feat(qwen4_exp): tensor parallelism for Qwen3.8-Flash-Next (offload backend)

python/freetoken/engine/config.py

@@ -79,7 +79,7 @@ class EngineConfig:
     # ratio default above. A runtime cache rebuild sets this (num_swa_pages) to pin the window
     # regardless of the full anchor; the ratio is the startup default and the fallback.
     swa_num_pages_override: int | None = None
-    distributed_timeout: float = 60.0
+    distributed_timeout: float = 1800.0  # ranks reach the first collective minutes apart on a 100+ GiB offload load
     use_dummy_weight: bool = False
     use_pynccl: bool = True
     max_seq_len_override: int | None = None

python/freetoken/layers/linear.py

@@ -83,13 +83,17 @@ def __init__(
         input_size: int,
         output_sizes: List[int],
         has_bias: bool,
+        local_output_sizes: List[int] | None = None,
         *,
         quant_config: QuantConfig | None = None,
         prefix: str = "",
     ):
-        # check that all output sizes are divisible by tp_size
+        # check that all output sizes are divisible by tp_size (a caller that replicates
+        # GQA kv heads across ranks passes the per-rank sizes explicitly)
         tp_info = get_tp_info()
-        tp_output_sizes = [div_even(size, tp_info.size) for size in output_sizes]
+        if local_output_sizes is None:
+            local_output_sizes = [div_even(size, tp_info.size) for size in output_sizes]
+        tp_output_sizes = local_output_sizes
         output_size = sum(output_sizes)
         tp_output_size = sum(tp_output_sizes)
         super().__init__(

python/freetoken/layers/quantization/moe/nvfp4.py

@@ -41,14 +41,20 @@ class TritonNvfp4MoEKernel(MoEKernel):
     cpu_format = "nvfp4"
 
     def unusable_reason(self, cfg: MoEConfig) -> str | None:
-        reason = self._common_reject(cfg, resident_ok=False, tp_ok=False, cpu_ok=True, plain_silu_only=False)
+        # tp_ok: the source stream is sliced along the intermediate axis per rank
+        # (nvfp4_banks._tp_shard) and this kernel sizes its banks from
+        # cfg.local_intermediate, so a rank holds and reads only its own half. The routed
+        # output is then a partial sum, which the MoE layer reduces (_maybe_all_reduce, or
+        # one combined all-reduce in qwen4_exp's block). marlin and b12x stay off: their
+        # pack() repacks the rows and neither has been checked against a sharded bank.
+        reason = self._common_reject(cfg, resident_ok=False, tp_ok=True, cpu_ok=True, plain_silu_only=False)
         if reason:
             return reason
         reason = gated_epilogue_reason(cfg)
         return f"triton nvfp4 MoE kernel: {reason}" if reason else None
 
     def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]:
-        i, h = cfg.intermediate, cfg.hidden
+        i, h = cfg.local_intermediate, cfg.hidden
         return {
             "gate_up": BankSpec((2 * i, h // 2), torch.uint8),
             "gate_up_scale": BankSpec((2 * i, h // GROUP), FP8),
@@ -61,7 +67,7 @@ def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]:
     def pack(self, pieces, cfg: MoEConfig, out):
         out["gate_up"].copy_(fused_piece(pieces, "gate_up"))
         out["gate_up_scale"].copy_(fused_piece(pieces, "gate_up_scale"))
-        out["gate_up_global"].copy_(fused_global(pieces, cfg.intermediate))
+        out["gate_up_global"].copy_(fused_global(pieces, cfg.local_intermediate))
         out["down"].copy_(pieces["down"])
         out["down_scale"].copy_(pieces["down_scale"])
         out["down_global"].copy_(global_rows(pieces["down_global"], cfg.hidden))
@@ -241,7 +247,7 @@ def worth_it(self, cfg: MoEConfig) -> bool:
         return (8, 0) <= backend.device_capability() < (10, 0)
 
     def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]:
-        i, h = cfg.intermediate, cfg.hidden
+        i, h = cfg.local_intermediate, cfg.hidden
         return {
             "gate_up": BankSpec((h // GROUP, 4 * i), torch.int32),
             "gate_up_scale": BankSpec((h // GROUP, 2 * i), FP8),
@@ -252,7 +258,7 @@ def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]:
         }
 
     def pack(self, pieces, cfg: MoEConfig, out):
-        i, h = cfg.intermediate, cfg.hidden
+        i, h = cfg.local_intermediate, cfg.hidden
         device = torch.device("cuda")
         gu, gus, gug = fused_piece(pieces, "gate_up"), fused_piece(pieces, "gate_up_scale"), fused_global(pieces, i)
         dn, dns, dng = pieces["down"], pieces["down_scale"], global_rows(pieces["down_global"], h)
@@ -503,7 +509,7 @@ def unusable_reason(self, cfg: MoEConfig) -> str | None:
 
     def worth_it(self, cfg: MoEConfig) -> bool:
         # NOTE: never auto-selected. flashinfer's cute launcher indexes each bank with int32 element offsets, so the GPU slot cache is capped at (2^31 - 1) / elements-per-slot (about 1000 slots for GLM-5.3-Flash's 12960 experts); until flashinfer lifts that, b12x stays behind triton in the table; the selection does not weigh the cap, cache-auto clamps to slot_limit() and OffloadMoeCache refuses a larger cache.
-        return cfg.intermediate >= B12X_MIN_INTERMEDIATE
+        return cfg.local_intermediate >= B12X_MIN_INTERMEDIATE
 
     def slot_limit(self, cfg: MoEConfig) -> int | None:
         # the cute launcher indexes each bank with int32 element offsets
@@ -512,7 +518,7 @@ def slot_limit(self, cfg: MoEConfig) -> int | None:
 
     def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]:
         # flashinfer's prepared tiles: the Marlin shapes, byte-identical to the native rows
-        i, h = cfg.intermediate, cfg.hidden
+        i, h = cfg.local_intermediate, cfg.hidden
         return {
             "gate_up": BankSpec((h // GROUP, 4 * i), torch.int32),
             "gate_up_scale": BankSpec((h // GROUP, 2 * i), FP8),
@@ -525,7 +531,7 @@ def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]:
     def pack(self, pieces, cfg: MoEConfig, out):
         from flashinfer.fused_moe.cute_dsl.blackwell_sm12x.moe_w4a16_prepare import prepare_w4a16_packed_weights
 
-        i, h = cfg.intermediate, cfg.hidden
+        i, h = cfg.local_intermediate, cfg.hidden
         device = torch.device("cuda")
         gu = fused_piece(pieces, "gate_up").to(device)
         gug = fused_global(pieces, i).to(device).float()

python/freetoken/models/nvfp4_banks.py

@@ -60,6 +60,36 @@ def _bank_layer(spec: Nvfp4ExpertSourceSpec, layer: int, config) -> int | None:
     return bank_layer
 
 
+def _tp_slice() -> tuple[int, int] | None:
+    """``(i_local, i_lo)``: this rank's slice of the expert intermediate axis, or None at TP=1.
+
+    TP shards every routed expert along I (gate/up output rows, down input columns), so each
+    rank caches only its own half and the MoE layer all-reduces the partial sum. The expert
+    kernel sizes its banks from ``MoEConfig.local_intermediate``, which is the same split.
+    """
+    from freetoken.distributed import try_get_tp_info
+
+    tp = try_get_tp_info()
+    if tp is None or tp.size == 1:
+        return None
+    return tp.size, tp.rank
+
+
+def _tp_shard(role: str, tensor: torch.Tensor, inter: int, tp_size: int, tp_rank: int):
+    """This rank's I-slice of one checkpoint expert tensor. ``_global`` is a per-tensor scalar
+    and has no I axis; gate/up carry I on the row axis, down on the column axis (halved for the
+    packed FP4 codes, sixteenthed for the fp8 block scales)."""
+    if role.endswith("_global"):
+        return tensor
+    n = inter // tp_size
+    assert n % 16 == 0, f"NVFP4 TP shard {n} must cover whole 16-wide scale blocks"
+    lo = tp_rank * n
+    if role.startswith("down"):
+        d = 2 if role == "down" else 16
+        return tensor[:, lo // d : (lo + n) // d]
+    return tensor[lo : lo + n]
+
+
 def _kind_suffix(kind: str) -> str:
     return {"weight": "", "weight_scale": "_scale", "weight_scale_2": "_global"}[kind]
 
@@ -133,7 +163,19 @@ def _parallel():
                 tensor = _ingest_global(spec, tensor)
             yield name, tensor
 
-    return per_expert_pieces(_parallel() if parallel else _serial(), wanted.get, tensors_per_expert=9)
+    stream = _parallel() if parallel else _serial()
+    shard = _tp_slice()
+    if shard is not None:
+        tp_size, tp_rank = shard
+        inter = config.moe_intermediate_size
+        base = stream
+
+        def _sharded():
+            for name, tensor in base:
+                yield name, _tp_shard(wanted[name][2], tensor, inter, tp_size, tp_rank)
+
+        stream = _sharded()
+    return per_expert_pieces(stream, wanted.get, tensors_per_expert=9)
 
 
 __all__ = ["Nvfp4ExpertSourceSpec", "iter_nvfp4_expert_pieces"]

python/freetoken/models/qwen4_exp/attention.py

@@ -19,9 +19,16 @@
 
 import torch
 from freetoken.core import get_global_ctx
-from freetoken.layers import BaseOP, GemmaPlusOneRMSNorm, LinearColParallelMerged, LinearReplicated
+from freetoken.distributed import get_tp_info
+from freetoken.layers import (
+    BaseOP,
+    GemmaPlusOneRMSNorm,
+    LinearColParallelMerged,
+    LinearOProj,
+    LinearReplicated,
+)
 from freetoken.layers.rotary import get_rope
-from freetoken.utils import nvtx_annotate
+from freetoken.utils import div_even, nvtx_annotate
 
 if TYPE_CHECKING:
     from freetoken.core import Batch
@@ -123,12 +130,25 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = "") -> N
         self.head_dim = config.head_dim
         self.qo_attn_dim = self.num_q * self.head_dim
         self.kv_attn_dim = self.num_kv * self.head_dim
-        self._qkv_split = [self.qo_attn_dim * 2, self.kv_attn_dim, self.kv_attn_dim]
+        # TP: q heads split across ranks, kv heads split or (num_kv < tp) replicated; the
+        # indexer stays replicated so every rank selects the same blocks.
+        tp = get_tp_info()
+        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_dim = self._local_num_q * self.head_dim
+        self._local_kv_dim = self._local_num_kv * self.head_dim
+        self._qkv_split = [self._local_qo_dim * 2, self._local_kv_dim, self._local_kv_dim]
         self.qkv_proj = LinearColParallelMerged(
-            config.hidden_size, self._qkv_split, has_bias=False,
+            config.hidden_size,
+            [self.qo_attn_dim * 2, self.kv_attn_dim, self.kv_attn_dim],
+            has_bias=False,
+            local_output_sizes=self._qkv_split,
             quant_config=config.quant, prefix=f"{prefix}.qkv_proj",
         )
-        self.o_proj = LinearReplicated(
+        # row-parallel, NOT LinearReplicated: under TP>1 qkv_proj is column-parallel, so this
+        # rank's attention output is its own head slice and o_proj must consume the sharded
+        # input dim and all-reduce the partial sums. At TP=1 the two are equivalent.
+        self.o_proj = LinearOProj(
             self.qo_attn_dim, config.hidden_size, has_bias=False,
             quant_config=config.quant, prefix=f"{prefix}.o_proj",
         )
@@ -147,21 +167,21 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = "") -> N
     @nvtx_annotate("QSA")
     def forward(self, x: torch.Tensor, batch: Batch) -> torch.Tensor:
         qg, k, v = self.qkv_proj.forward(x).split(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()
-        gate = qg[..., self.head_dim :].reshape(-1, self.qo_attn_dim)
-        k = k.contiguous().view(-1, self.num_kv, self.head_dim)
+        gate = qg[..., self.head_dim :].reshape(-1, self._local_qo_dim)
+        k = k.contiguous().view(-1, self._local_num_kv, self.head_dim)
         v = v.contiguous()
         self.q_norm.forward_inplace(q)
         self.k_norm.forward_inplace(k)
         q, k = self.rotary.forward(
-            batch.positions, q.view(-1, self.qo_attn_dim), k.view(-1, self.kv_attn_dim)
+            batch.positions, q.view(-1, self._local_qo_dim), k.view(-1, self._local_kv_dim)
         )
         index = self.indexer.forward(x)
         o = get_global_ctx().attn_backend.qsa_forward(
-            q.view(-1, self.num_q, self.head_dim), k, v, index, self.layer_id, batch
+            q.view(-1, self._local_num_q, self.head_dim), k, v, index, self.layer_id, batch
         )
-        gated = o.reshape(-1, self.qo_attn_dim) * torch.sigmoid(gate)
+        gated = o.reshape(-1, self._local_qo_dim) * torch.sigmoid(gate)
         return self.o_proj.forward(gated)
 
 

python/freetoken/models/qwen4_exp/gdn.py

@@ -3,9 +3,12 @@
 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, GatedRMSNorm, LinearColParallelMerged, LinearReplicated
+from freetoken.distributed import get_tp_info
+from freetoken.layers import BaseOP, GatedRMSNorm, LinearColParallelMerged, LinearOProj
 from freetoken.layers.quantization import QuantConfig
+from freetoken.utils import div_even
 from freetoken.models.qwen3_5_moe.gdn_kernels import gdn_decode_fla, gdn_prefill_chunk_fla
 
 
@@ -50,12 +53,29 @@ 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: k heads and their v-head groups split evenly across ranks
+        # (the fla kernels take the GQA ratio from the shapes); the state pool is sharded the
+        # same way (kvcache.linear_state_pool._linear_local_dims).
+        tp = get_tp_info()
+        self._local_num_k_heads = div_even(num_k_heads, tp.size, allow_replicate=True)
+        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
         # quantized checkpoints quantize qkv|z but not b|a, so the fusion splits into a qkvz GEMM and a ba GEMM with their own schemes (matches sglang / vLLM)
         self._split_in_proj = (
             quant_config is not None and quant_config.scheme_for(f"{prefix}.in_proj_qkvz") is not None
         )
 
-        self._in_proj_split = [self.conv_dim, self.value_dim, num_v_heads, num_v_heads]
+        self._in_proj_split = [
+            self._local_conv_dim, self._local_value_dim,
+            self._local_num_v_heads, self._local_num_v_heads,
+        ]
+        # The split (quantized) branch has no TP-aware variant: its two linears are built from
+        # the checkpoint's own scheme and are not sharded here.
+        assert not (tp.size > 1 and self._split_in_proj), (
+            "qwen4_exp TP shards bf16 GDN projections only"
+        )
         if self._split_in_proj:
             self.in_proj_qkvz = LinearColParallelMerged(
                 hidden_size, [self.conv_dim, self.value_dim], has_bias=False,
@@ -68,18 +88,23 @@ 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,
+                hidden_size, [self.conv_dim, self.value_dim, num_v_heads, num_v_heads],
+                has_bias=False, local_output_sizes=self._in_proj_split,
                 quant_config=quant_config, prefix=f"{prefix}.in_proj",
             )
-        self.conv1d = _DepthwiseConv1d(self.conv_dim, conv_kernel_size)
+        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, activation=output_gate)
-        self.out_proj = LinearReplicated(
+        # Row-parallel over the local v heads, all-reduce inside. At TP=1 this is exactly
+        # LinearReplicated (div_even(x, 1) == x and the reduction is skipped), so one class
+        # covers both -- and under TP>1 a replicated o_proj would silently take the wrong
+        # input width and skip the reduction.
+        self.out_proj = LinearOProj(
             self.value_dim, hidden_size, has_bias=False,
             quant_config=quant_config, prefix=f"{prefix}.out_proj",
         )
@@ -148,7 +173,9 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         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)
+        nk, nv = self._local_num_k_heads, self._local_num_v_heads
+        kd, vd = self._local_key_dim, self._local_value_dim
+        z = z.reshape(total, nv, self.head_v_dim)
         li = pool.local_index(self.layer_id)
 
         if batch.is_decode:
@@ -157,10 +184,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, [kd, kd, vd], dim=-1)
+            q = qf.reshape(1, B, nk, self.head_k_dim).to(dtype)
+            k = kf.reshape(1, B, nk, self.head_k_dim).to(dtype)
+            v = vf.reshape(1, B, nv, 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,
@@ -170,13 +197,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, [kd, kd, vd], dim=-1)
+            q = qf.reshape(1, total, nk, self.head_k_dim).to(dtype)
+            k = kf.reshape(1, total, nk, self.head_k_dim).to(dtype)
+            v = vf.reshape(1, total, nv, 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, nv)
+            beta = beta.float().reshape(1, total, nv)
             # 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: