feat(models): support TP for qwen3_5_moe

python/freetoken/layers/linear.py

@@ -59,10 +59,14 @@ def __init__(
         input_size: int,
         output_sizes: List[int],
         has_bias: bool,
+        local_output_sizes: List[int] | None = None,
     ):
         # check that all output sizes are divisible by tp_size
         tp_info = get_tp_info()
-        tp_output_sizes = [div_even(size, tp_info.size) for size in output_sizes]
+        if local_output_sizes is not None:
+            tp_output_sizes = local_output_sizes
+        else:
+            tp_output_sizes = [div_even(size, tp_info.size) for size in output_sizes]
         output_size = sum(output_sizes)
         tp_output_size = sum(tp_output_sizes)
         super().__init__(input_size, output_size, input_size, tp_output_size, has_bias)

python/freetoken/layers/moe.py

@@ -219,6 +219,11 @@ def __init__(
         self.layer_id = layer_id
         self.offload_cache: OffloadMoeCache | None = None
 
+    def _maybe_all_reduce(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Offload MoE: each rank computes the full expert output (banks are not TP-sharded),
+        so no all-reduce is needed — the result is already the complete output."""
+        return hidden_states
+
     def forward(
         self,
         hidden_states: torch.Tensor,

python/freetoken/models/quant_linear.py

@@ -12,8 +12,11 @@
 
 
 def make_col_merged_quant(expert_quant: str, attn_quant: str, in_f: int,
-                          output_sizes: list[int], has_bias: bool = False):
+                          output_sizes: list[int], has_bias: bool = False,
+                          local_output_sizes: list[int] | None = None):
     """Column-merged linear for a dense projection: block-fp8 / per-tensor-fp8 / nvfp4 / bf16."""
+    if local_output_sizes is not None and (expert_quant != "none" or attn_quant != "none"):
+        raise NotImplementedError("local_output_sizes (GQA KV replication) not yet supported for quantized checkpoints")
     if expert_quant == "fp8_block":
         from freetoken.kernel.triton.fp8_block_linear import Fp8BlockColMerged
 
@@ -28,7 +31,8 @@ def make_col_merged_quant(expert_quant: str, attn_quant: str, in_f: int,
         return Nvfp4DenseColMerged(in_f, output_sizes, has_bias)
     from freetoken.layers import LinearColParallelMerged
 
-    return LinearColParallelMerged(in_f, output_sizes, has_bias=has_bias)
+    return LinearColParallelMerged(in_f, output_sizes, has_bias=has_bias,
+                                   local_output_sizes=local_output_sizes)
 
 
 def make_replicated_quant(expert_quant: str, attn_quant: str, in_f: int, out_f: int,
@@ -51,6 +55,21 @@ def make_replicated_quant(expert_quant: str, attn_quant: str, in_f: int, out_f:
     return LinearReplicated(in_f, out_f, has_bias=has_bias)
 
 
+def make_row_parallel_quant(expert_quant: str, attn_quant: str, in_f: int, out_f: int,
+                            has_bias: bool = False):
+    """Row-parallel linear for a dense projection: block-fp8 / per-tensor-fp8 / nvfp4 / bf16.
+    Shards the input dimension by tp_size and all-reduces on forward."""
+    if expert_quant == "fp8_block":
+        raise NotImplementedError("row-parallel Fp8BlockLinear not yet implemented")
+    if attn_quant == "fp8_pertensor":
+        raise NotImplementedError("row-parallel Fp8PerTensorLinear not yet implemented")
+    if attn_quant == "nvfp4":
+        raise NotImplementedError("row-parallel Nvfp4DenseLinear not yet implemented")
+    from freetoken.layers import LinearOProj
+
+    return LinearOProj(in_f, out_f, has_bias=has_bias)
+
+
 def make_replicated(config, in_f: int, out_f: int, has_bias: bool = False):
     """Config-driven replicated linear: ``Fp8BlockLinear`` under block-fp8, ``Fp8PerTensorLinear``
     under per-tensor-fp8 attention, ``Nvfp4DenseLinear`` under nvfp4, else ``LinearReplicated``."""
@@ -60,19 +79,30 @@ def make_replicated(config, in_f: int, out_f: int, has_bias: bool = False):
     )
 
 
-def make_col_merged(config, in_f: int, output_sizes: list[int], has_bias: bool = False):
+def make_col_merged(config, in_f: int, output_sizes: list[int], has_bias: bool = False,
+                     local_output_sizes: list[int] | None = None):
     """Config-driven column-merged linear: ``Fp8BlockColMerged`` under block-fp8,
     ``Fp8PerTensorColMerged`` under per-tensor-fp8 attention, ``Nvfp4DenseColMerged`` under
     nvfp4, else ``LinearColParallelMerged``."""
     return make_col_merged_quant(
         getattr(config, "expert_quant", "none"), getattr(config, "attn_quant", "none"),
-        in_f, output_sizes, has_bias,
+        in_f, output_sizes, has_bias, local_output_sizes=local_output_sizes,
+    )
+
+
+def make_row_parallel(config, in_f: int, out_f: int, has_bias: bool = False):
+    """Config-driven row-parallel linear: ``LinearOProj`` (bf16); quant variants TBD."""
+    return make_row_parallel_quant(
+        getattr(config, "expert_quant", "none"), getattr(config, "attn_quant", "none"),
+        in_f, out_f, has_bias,
     )
 
 
 __all__ = [
     "make_col_merged_quant",
     "make_replicated_quant",
+    "make_row_parallel_quant",
     "make_replicated",
     "make_col_merged",
+    "make_row_parallel",
 ]

python/freetoken/models/qwen3_5_moe/attention.py

@@ -4,11 +4,12 @@
 
 import torch
 from freetoken.core import get_global_ctx
+from freetoken.distributed import get_tp_info
 from freetoken.layers import BaseOP, GemmaRMSNorm
 from freetoken.layers.rotary import get_rope
-from freetoken.utils import nvtx_annotate
+from freetoken.utils import div_even, nvtx_annotate
 
-from .quant_linear import make_col_merged, make_replicated
+from .quant_linear import make_col_merged, make_row_parallel
 
 if TYPE_CHECKING:
     from freetoken.models.config import ModelConfig
@@ -22,9 +23,6 @@ class Qwen3_5Attention(BaseOP):
         q, k = rope(q, k)                       # first rotary_dim dims
         attn = paged_attention(q, k, v)
         out = o_proj(attn * sigmoid(gate))
-
-    TP note: uses replicated linears (tp=1 correctness milestone); swap to
-    column/row-parallel for tensor parallelism later.
     """
 
     def __init__(self, config: ModelConfig, layer_id: int):
@@ -36,13 +34,21 @@ def __init__(self, config: ModelConfig, layer_id: int):
         self.qo_attn_dim = self.num_q * head_dim
         self.kv_attn_dim = self.num_kv * head_dim
 
+        tp = get_tp_info()
         # Fused q/k/v projection (one GEMM instead of three); q half is 2x for the
         # output gate. Split sizes: [num_q*head_dim*2, num_kv*head_dim, num_kv*head_dim].
-        self._qkv_split = [self.num_q * head_dim * 2, self.kv_attn_dim, self.kv_attn_dim]
+        # _qkv_split is TP-local (for torch.split in forward); make_col_merged gets full sizes.
         # Block-fp8 (Fp8BlockColMerged) when the checkpoint is quantized, else bf16
         # LinearColParallelMerged. q/k/v out dims are all /128, so the merged fp8 weight +
         # weight_scale_inv concatenate cleanly along the output dim.
-        self.qkv_proj = make_col_merged(config, config.hidden_size, self._qkv_split, has_bias=False)
+        full_split = [self.num_q * head_dim * 2, self.kv_attn_dim, self.kv_attn_dim]
+        self._local_num_q = div_even(self.num_q, tp.size)
+        self._local_num_kv = div_even(self.num_kv, tp.size, allow_replicate=True)
+        self._local_qo_attn_dim = self._local_num_q * head_dim
+        self._local_kv_attn_dim = self._local_num_kv * head_dim
+        self._qkv_split = [self._local_num_q * head_dim * 2, self._local_kv_attn_dim, self._local_kv_attn_dim]
+        local_qkv_sizes = [self._local_num_q * head_dim * 2, self._local_kv_attn_dim, self._local_kv_attn_dim]
+        self.qkv_proj = make_col_merged(config, config.hidden_size, full_split, has_bias=False, local_output_sizes=local_qkv_sizes)
         # Qwen3.5 uses Gemma-style (1+weight) RMSNorm; the weight loader bakes the +1
         # into the stored weight (GemmaRMSNorm scales by the raw weight).
         self.q_norm = GemmaRMSNorm(head_dim, eps=config.rms_norm_eps)
@@ -58,26 +64,26 @@ def __init__(self, config: ModelConfig, layer_id: int):
                 else None
             ),
         )
-        self.o_proj = make_replicated(config, self.qo_attn_dim, config.hidden_size, has_bias=False)
+        self.o_proj = make_row_parallel(config, self.qo_attn_dim, config.hidden_size, has_bias=False)
 
     def _project(self, x: torch.Tensor):
         """Returns (q, k, v, gate): q [N, num_q, head_dim] post qk-norm+rope,
         k [N, num_kv*head_dim] post norm+rope, v [N, num_kv*head_dim], gate [N, num_q*head_dim]."""
         positions = get_global_ctx().batch.positions
         qkv = self.qkv_proj.forward(x)
         qg, k, v = torch.split(qkv, self._qkv_split, dim=-1)
-        qg = qg.view(-1, self.num_q, self.head_dim * 2)
+        qg = qg.view(-1, self._local_num_q, self.head_dim * 2)
         q = qg[..., : self.head_dim].contiguous()  # [N, num_q, head_dim]
-        gate = qg[..., self.head_dim :].reshape(-1, self.qo_attn_dim)
-        k = k.view(-1, self.num_kv, self.head_dim).contiguous()
+        gate = qg[..., self.head_dim :].reshape(-1, self._local_qo_attn_dim)
+        k = k.reshape(-1, self._local_num_kv, self.head_dim).contiguous()
         v = v.contiguous()  # split view has the qkv row stride; the KV store needs contiguous
-        q = self.q_norm.forward(q).reshape(-1, self.qo_attn_dim)
-        k = self.k_norm.forward(k).reshape(-1, self.kv_attn_dim)
+        q = self.q_norm.forward(q).reshape(-1, self._local_qo_attn_dim)
+        k = self.k_norm.forward(k).reshape(-1, self._local_kv_attn_dim)
         q, k = self.rotary.forward(positions, q, k)
-        return q.view(-1, self.num_q, self.head_dim), k, v, gate
+        return q.view(-1, self._local_num_q, self.head_dim), k, v, gate
 
     def _combine(self, attn_out: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
-        gated = attn_out.reshape(-1, self.qo_attn_dim) * torch.sigmoid(gate)
+        gated = attn_out.reshape(-1, self._local_qo_attn_dim) * torch.sigmoid(gate)
         return self.o_proj.forward(gated)
 
     @nvtx_annotate("MHA")

python/freetoken/models/qwen3_5_moe/gdn.py

@@ -3,14 +3,16 @@
 import torch
 import torch.nn.functional as F
 from freetoken.core import get_global_ctx
+from freetoken.distributed import get_tp_info
 from freetoken.kernel.causal_conv1d import causal_conv1d_decode, causal_conv1d_varlen
 from freetoken.layers import BaseOP, LinearColParallelMerged
+from freetoken.utils import div_even
 
 from freetoken.kernel.triton.fp8_block_linear import Fp8BlockColMerged
 from freetoken.kernel.triton.fp8_pertensor_linear import Fp8PerTensorColMerged
 
 from .gdn_kernels import gdn_decode_fla, gdn_prefill_chunk_fla
-from .quant_linear import make_replicated_quant
+from .quant_linear import make_row_parallel_quant
 
 
 class _DepthwiseConv1d(BaseOP):
@@ -56,6 +58,7 @@ def __init__(
         attn_quant: str = "none",
     ):
         self.layer_id = layer_id
+        tp = get_tp_info()
         # The fla chunk/decode kernels read+write the recurrent state and the per-chunk h as
         # [V, K] while the LinearStatePool declares it [K, V]; these coincide (and the
         # hybrid-radix snapshot scatter h[h_row]->slot is a plain copy) only when the two head
@@ -71,14 +74,21 @@ def __init__(
         self.value_dim = num_v_heads * head_v_dim
         self.conv_dim = 2 * self.key_dim + self.value_dim
         self.conv_kernel_size = conv_kernel_size
+        # TP-local head counts for forward (reshape/split); LinearColParallelMerged gets full sizes
+        self._local_num_k_heads = div_even(num_k_heads, tp.size)
+        self._local_num_v_heads = div_even(num_v_heads, tp.size, allow_replicate=True)
+        self._local_key_dim = self._local_num_k_heads * head_k_dim
+        self._local_value_dim = self._local_num_v_heads * head_v_dim
+        self._local_conv_dim = 2 * self._local_key_dim + self._local_value_dim
         # qkv|z carry a weight scale (block-fp8 weight_scale_inv, or per-tensor FP8
         # weight_scale); b|a stay bf16. Both quant modes therefore split the four-way
         # fusion into an fp8 qkvz GEMM + a bf16 ba GEMM (matches sglang/vLLM).
         self._block_fp8 = expert_quant == "fp8_block"
         self._pertensor_fp8 = attn_quant == "fp8_pertensor"
         self._fp8 = self._block_fp8 or self._pertensor_fp8
 
-        self._in_proj_split = [self.conv_dim, self.value_dim, num_v_heads, num_v_heads]
+        full_split = [self.conv_dim, self.value_dim, num_v_heads, num_v_heads]
+        self._in_proj_split = [div_even(s, tp.size) for s in full_split]  # TP-local for torch.split
         if self._fp8:
             ColMerged = Fp8BlockColMerged if self._block_fp8 else Fp8PerTensorColMerged
             self.in_proj_qkvz = ColMerged(
@@ -89,19 +99,19 @@ def __init__(
             )
         else:
             # Fused input projection (one GEMM instead of four): qkv | z | b | a.
-            self.in_proj = LinearColParallelMerged(hidden_size, self._in_proj_split, has_bias=False)
-        self.conv1d = _DepthwiseConv1d(self.conv_dim, conv_kernel_size)
+            self.in_proj = LinearColParallelMerged(hidden_size, full_split, has_bias=False)
+        self.conv1d = _DepthwiseConv1d(self._local_conv_dim, conv_kernel_size)
         # Recurrence-gating params kept in fp32 (exp/softplus is precision-sensitive,
         # and the fla kernel reads them as fp32) -- matches HF/sglang, and avoids a
         # per-call .float() upcast in the decode wrapper. The weight loader exempts
         # *.A_log / *.dt_bias from the model-dtype downcast.
-        self.dt_bias = torch.empty(num_v_heads, dtype=torch.float32)
-        self.A_log = torch.empty(num_v_heads, dtype=torch.float32)
+        self.dt_bias = torch.empty(self._local_num_v_heads, dtype=torch.float32)
+        self.A_log = torch.empty(self._local_num_v_heads, dtype=torch.float32)
         self.norm = _GatedRMSNorm(head_v_dim, eps=rms_norm_eps)
         # out_proj follows the checkpoint quant: block-fp8 / per-tensor-fp8 / compressed-tensors
         # NVFP4 (W4A16) / bf16. in_proj_* stay bf16 in every mode (above), so a compressed-tensors
         # NVFP4 checkpoint (attn_quant=="nvfp4") only makes out_proj native FP4.
-        self.out_proj = make_replicated_quant(
+        self.out_proj = make_row_parallel_quant(
             expert_quant, attn_quant, self.value_dim, hidden_size, has_bias=False
         )
 
@@ -163,13 +173,13 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 
         if self._fp8:
             qkvz = self.in_proj_qkvz.forward(hidden_states)
-            conv_in, z = torch.split(qkvz, [self.conv_dim, self.value_dim], dim=-1)
+            conv_in, z = torch.split(qkvz, [self._local_conv_dim, self._local_value_dim], dim=-1)
             ba = self.in_proj_ba.forward(hidden_states)
-            b, a = torch.split(ba, [self.num_v_heads, self.num_v_heads], dim=-1)
+            b, a = torch.split(ba, [self._local_num_v_heads, self._local_num_v_heads], dim=-1)
         else:
             proj = self.in_proj.forward(hidden_states)
             conv_in, z, b, a = torch.split(proj, self._in_proj_split, dim=-1)
-        z = z.reshape(total, self.num_v_heads, self.head_v_dim)
+        z = z.reshape(total, self._local_num_v_heads, self.head_v_dim)
         li = pool.local_index(self.layer_id)
 
         if batch.is_decode:
@@ -178,10 +188,10 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
             # no clone, no external l2norm). q/k stay at num_k_heads (kernel handles GQA).
             mixed = self._conv_decode(conv_in, fla.cache_indices, pool)  # [B, conv_dim]
             B = mixed.shape[0]
-            qf, kf, vf = torch.split(mixed, [self.key_dim, self.key_dim, self.value_dim], dim=-1)
-            q = qf.reshape(1, B, self.num_k_heads, self.head_k_dim).to(dtype)
-            k = kf.reshape(1, B, self.num_k_heads, self.head_k_dim).to(dtype)
-            v = vf.reshape(1, B, self.num_v_heads, self.head_v_dim).to(dtype)
+            qf, kf, vf = torch.split(mixed, [self._local_key_dim, self._local_key_dim, self._local_value_dim], dim=-1)
+            q = qf.reshape(1, B, self._local_num_k_heads, self.head_k_dim).to(dtype)
+            k = kf.reshape(1, B, self._local_num_k_heads, self.head_k_dim).to(dtype)
+            v = vf.reshape(1, B, self._local_num_v_heads, self.head_v_dim).to(dtype)
             core_out = gdn_decode_fla(
                 q, k, v, a, b, A_log=self.A_log, dt_bias=self.dt_bias,
                 state_source=pool.recurrent_states[li], indices=fla.cache_indices,
@@ -191,13 +201,13 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
             mixed = self._conv_prefill(
                 conv_in, pool, fla.cu_seqlens, fla.cache_indices, fla.has_initial_state)
             # fla chunk handles GQA in-kernel: q/k stay at num_k_heads, v at num_v_heads.
-            qf, kf, vf = torch.split(mixed, [self.key_dim, self.key_dim, self.value_dim], dim=-1)
-            q = qf.reshape(1, total, self.num_k_heads, self.head_k_dim).to(dtype)
-            k = kf.reshape(1, total, self.num_k_heads, self.head_k_dim).to(dtype)
-            v = vf.reshape(1, total, self.num_v_heads, self.head_v_dim).to(dtype)
+            qf, kf, vf = torch.split(mixed, [self._local_key_dim, self._local_key_dim, self._local_value_dim], dim=-1)
+            q = qf.reshape(1, total, self._local_num_k_heads, self.head_k_dim).to(dtype)
+            k = kf.reshape(1, total, self._local_num_k_heads, self.head_k_dim).to(dtype)
+            v = vf.reshape(1, total, self._local_num_v_heads, self.head_v_dim).to(dtype)
             g, beta = self._gate_params(a, b)
-            g = g.reshape(1, total, self.num_v_heads)
-            beta = beta.float().reshape(1, total, self.num_v_heads)
+            g = g.reshape(1, total, self._local_num_v_heads)
+            beta = beta.float().reshape(1, total, self._local_num_v_heads)
             # The chunk kernel reads + writes back initial_state[cache_indices] in place;
             # fresh sequences (cached_len==0) must start from a zeroed slot.
             if fla.fresh_state_indices is not None:

python/freetoken/models/qwen3_5_moe/quant_linear.py

@@ -9,11 +9,15 @@
     make_col_merged_quant,
     make_replicated,
     make_replicated_quant,
+    make_row_parallel,
+    make_row_parallel_quant,
 )
 
 __all__ = [
     "make_col_merged_quant",
     "make_replicated_quant",
+    "make_row_parallel_quant",
     "make_replicated",
     "make_col_merged",
+    "make_row_parallel",
 ]