diff --git a/python/freetoken/models/gemma4/attention.py b/python/freetoken/models/gemma4/attention.py index 4f85255f0..abfbd1b5d 100644 --- a/python/freetoken/models/gemma4/attention.py +++ b/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", ) diff --git a/python/freetoken/models/muse_glimmer/attention.py b/python/freetoken/models/muse_glimmer/attention.py index 7c814d4c1..aedef62b5 100644 --- a/python/freetoken/models/muse_glimmer/attention.py +++ b/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", ) diff --git a/python/freetoken/models/qwen3_5_moe/attention.py b/python/freetoken/models/qwen3_5_moe/attention.py index 12a91bb3e..fde694489 100644 --- a/python/freetoken/models/qwen3_5_moe/attention.py +++ b/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", )