feat: remote code support for Hybrid FP8 Attn/GDN + INT4-g16 MLP

modeling_wemm_embedding.py

@@ -34,7 +34,7 @@ class FlatQuantFP8Linear(nn.Module):
         return F.linear(x, w_deq, self.bias.to(x.dtype) if self.bias is not None else None)
 
 class FlatQuantW4A8Linear(nn.Module):
-    def __init__(self, in_features, out_features, bias=False, group_size=32):
+    def __init__(self, in_features, out_features, bias=False, group_size=16):
         super().__init__()
         self.in_features = in_features
         self.out_features = out_features
@@ -60,8 +60,6 @@ class FlatQuantW4A8Linear(nn.Module):
 class WeMMEmbedding(Qwen3_5ForConditionalGeneration):
     def __init__(self, config):
         super().__init__(config)
-        full_attn_layers = {3, 7, 11, 15, 19, 23}
-        pas_layers = {2, 6, 10, 14, 18, 22}
         for name, mod in list(self.model.named_modules()):
             if name.endswith("embed_tokens") and isinstance(mod, nn.Embedding):
                 parent = self.model.get_submodule(name.rsplit(".", 1)[0]) if "." in name else self.model
@@ -70,13 +68,12 @@ class WeMMEmbedding(Qwen3_5ForConditionalGeneration):
             elif isinstance(mod, nn.Linear):
                 parent = self.model.get_submodule(name.rsplit(".", 1)[0]) if "." in name else self.model
                 child = name.rsplit(".", 1)[-1]
-                is_full = any(f"layers.{l}.self_attn" in name for l in full_attn_layers)
-                is_pas_kv = any(f"layers.{l}.linear_attn" in name for l in pas_layers) and (child in ["k_proj", "v_proj"])
-                is_down_proj_outer = ("down_proj" in name) and not (4 <= int(name.split(".layers.")[1].split(".")[0]) <= 18 if ".layers." in name else False)
-                if is_full or is_pas_kv or is_down_proj_outer:
+                is_attn = ("self_attn" in name) or ("linear_attn" in name)
+                is_down = "down_proj" in name
+                if is_attn or is_down:
                     setattr(parent, child, FlatQuantFP8Linear(mod.in_features, mod.out_features, bias=mod.bias is not None))
                 else:
-                    setattr(parent, child, FlatQuantW4A8Linear(mod.in_features, mod.out_features, bias=mod.bias is not None, group_size=32))
+                    setattr(parent, child, FlatQuantW4A8Linear(mod.in_features, mod.out_features, bias=mod.bias is not None, group_size=16))
 
     def embedding(self, input_ids=None, attention_mask=None, **kwargs):
         self.model.rope_deltas = None