fix(models): build o_proj row-parallel in the remaining column-parallel families

python/freetoken/models/gemma4/attention.py

@@ -5,7 +5,7 @@
 import torch
 from freetoken.attention import AttentionSpec
 from freetoken.core import get_global_ctx
-from freetoken.layers import BaseOP, GemmaRMSNorm, LinearQKVMerged, LinearReplicated
+from freetoken.layers import BaseOP, GemmaRMSNorm, LinearOProj, LinearQKVMerged
 from freetoken.layers.rotary import get_rope
 from freetoken.models.config import FullAttentionGroupConfig, SWAAttentionGroupConfig
 from freetoken.utils import nvtx_annotate
@@ -40,7 +40,12 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""):
             quant_config=config.quant,
             prefix=f"{prefix}.qkv_proj",
         )
-        self.o_proj = LinearReplicated(
+        # Row-parallel, as every other family with a column-parallel qkv builds o_proj:
+        # each rank's attention output is its local head slice, so o_proj takes the sharded
+        # input dim and all-reduces the partial sums. At TP=1 this degenerates to the previous
+        # replicated behaviour exactly (div_even(x, 1) == x, all-reduce skipped), so it is a
+        # no-op today and correct whenever this family gains tensor parallelism.
+        self.o_proj = LinearOProj(
             self.q_dim, config.hidden_size, has_bias=False,
             quant_config=config.quant, prefix=f"{prefix}.o_proj",
         )

python/freetoken/models/muse_glimmer/attention.py

@@ -5,7 +5,7 @@
 import torch
 from freetoken.attention import AttentionSpec
 from freetoken.core import get_global_ctx
-from freetoken.layers import BaseOP, GemmaRMSNorm, LinearColParallelMerged, LinearReplicated
+from freetoken.layers import BaseOP, GemmaRMSNorm, LinearColParallelMerged, LinearOProj
 from freetoken.layers.rotary import get_rope
 from freetoken.models.config import SWAAttentionGroupConfig
 from freetoken.utils import nvtx_annotate
@@ -46,7 +46,12 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""):
             config.hidden_size, self._qkvg_split, has_bias=False,
             quant_config=config.quant, prefix=f"{prefix}.qkvg_proj",
         )
-        self.o_proj = LinearReplicated(
+        # Row-parallel, as every other family with a column-parallel qkv builds o_proj:
+        # each rank's attention output is its local head slice, so o_proj takes the sharded
+        # input dim and all-reduces the partial sums. At TP=1 this degenerates to the previous
+        # replicated behaviour exactly (div_even(x, 1) == x, all-reduce skipped), so it is a
+        # no-op today and correct whenever this family gains tensor parallelism.
+        self.o_proj = LinearOProj(
             self.qo_attn_dim, config.hidden_size, has_bias=False,
             quant_config=config.quant, prefix=f"{prefix}.o_proj",
         )

python/freetoken/models/qwen3_5_moe/attention.py

@@ -4,7 +4,7 @@
 
 import torch
 from freetoken.core import get_global_ctx
-from freetoken.layers import BaseOP, GemmaRMSNorm, LinearColParallelMerged, LinearReplicated
+from freetoken.layers import BaseOP, GemmaRMSNorm, LinearColParallelMerged, LinearOProj
 from freetoken.layers.rotary import get_rope
 from freetoken.utils import nvtx_annotate
 
@@ -56,7 +56,12 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""):
                 else None
             ),
         )
-        self.o_proj = LinearReplicated(
+        # Row-parallel, as every other family with a column-parallel qkv builds o_proj:
+        # each rank's attention output is its local head slice, so o_proj takes the sharded
+        # input dim and all-reduces the partial sums. At TP=1 this degenerates to the previous
+        # replicated behaviour exactly (div_even(x, 1) == x, all-reduce skipped), so it is a
+        # no-op today and correct whenever this family gains tensor parallelism.
+        self.o_proj = LinearOProj(
             self.qo_attn_dim, config.hidden_size, has_bias=False,
             quant_config=config.quant, prefix=f"{prefix}.o_proj",
         )