From c3d0f975bba332f364c695f4ac146ef46a2fa072 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 4 Sep 2026 14:19:37 -0400 Subject: [PATCH 01/14] feat(qwen4_exp): tensor parallelism for Qwen3.8-Flash-Next (offload backend) Shard the dense weights per rank at load (attention qkv by head, GDN in_proj as its six parts with the matching conv1d channels and A_log/dt_bias, shared-expert gate_up per part; o_proj/out_proj/down_proj row-parallel; embed/lm_head by vocab rows) and the NVFP4 expert banks along the intermediate axis, so every rank holds half the experts and each MoE layer needs one all-reduce (routed + gate * shared are combined before the reduce). Router, QSA indexer, norms, hyper-connections and PLE stay replicated so all ranks select the same blocks and n-gram rows. Also: LinearColParallelMerged(local_output_sizes=) for the kv-replicated case and distributed_timeout 60 -> 1800 s (ranks reach their first collective minutes apart behind a 100+ GiB load). Limits: offload backend with bf16 dense projections; fp8_block / nvfp4 dense checkpoints raise under TP. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt --- python/freetoken/engine/config.py | 2 +- python/freetoken/layers/linear.py | 8 +- .../layers/quantization/moe/nvfp4.py | 14 +-- python/freetoken/models/nvfp4_banks.py | 44 ++++++++- .../freetoken/models/qwen4_exp/attention.py | 42 ++++++--- python/freetoken/models/qwen4_exp/gdn.py | 63 +++++++++---- python/freetoken/models/qwen4_exp/moe.py | 31 ++++++- python/freetoken/models/qwen4_exp/weight.py | 69 ++++++++++++-- tests/models/qwen4_exp/test_tp_shard.py | 92 +++++++++++++++++++ tests/models/test_nvfp4_banks_tp.py | 76 +++++++++++++++ 10 files changed, 392 insertions(+), 49 deletions(-) create mode 100644 tests/models/qwen4_exp/test_tp_shard.py create mode 100644 tests/models/test_nvfp4_banks_tp.py diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 7ab792f9e..e14fd3118 100644 --- a/python/freetoken/engine/config.py +++ b/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 diff --git a/python/freetoken/layers/linear.py b/python/freetoken/layers/linear.py index d8a3e8619..40ebe934c 100644 --- a/python/freetoken/layers/linear.py +++ b/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__( diff --git a/python/freetoken/layers/quantization/moe/nvfp4.py b/python/freetoken/layers/quantization/moe/nvfp4.py index bb341cdfb..dfe35604b 100644 --- a/python/freetoken/layers/quantization/moe/nvfp4.py +++ b/python/freetoken/layers/quantization/moe/nvfp4.py @@ -48,7 +48,7 @@ def unusable_reason(self, cfg: MoEConfig) -> str | None: 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 +61,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 +241,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 +252,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 +503,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 +512,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 +525,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() diff --git a/python/freetoken/models/nvfp4_banks.py b/python/freetoken/models/nvfp4_banks.py index 1d1b42859..0feddaf2a 100644 --- a/python/freetoken/models/nvfp4_banks.py +++ b/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"] diff --git a/python/freetoken/models/qwen4_exp/attention.py b/python/freetoken/models/qwen4_exp/attention.py index 66211c7c2..651da87d9 100644 --- a/python/freetoken/models/qwen4_exp/attention.py +++ b/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) diff --git a/python/freetoken/models/qwen4_exp/gdn.py b/python/freetoken/models/qwen4_exp/gdn.py index ed49dd42f..3b2e02934 100644 --- a/python/freetoken/models/qwen4_exp/gdn.py +++ b/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: diff --git a/python/freetoken/models/qwen4_exp/moe.py b/python/freetoken/models/qwen4_exp/moe.py index 3687f2803..7d9f6adf7 100644 --- a/python/freetoken/models/qwen4_exp/moe.py +++ b/python/freetoken/models/qwen4_exp/moe.py @@ -3,7 +3,12 @@ from typing import TYPE_CHECKING import torch +import torch.nn.functional as F +from freetoken.core import get_global_ctx +from freetoken.distributed import DistributedCommunicator, get_tp_info from freetoken.kernel.triton.moe_shared_gate import shared_gate_mul_add, shared_gate_sigmoid +from freetoken.layers import LinearRowParallel, silu_and_mul +from freetoken.layers.moe import OffloadMoELayer from freetoken.models.qwen3_5_moe.moe import Qwen3_5MoE if TYPE_CHECKING: @@ -14,15 +19,37 @@ class Qwen4ExpMoE(Qwen3_5MoE): """Qwen3_5MoE with the shared-expert gate on triton instead of gemv + sigmoid + mul + add. Same weights, same state dict. The gate reduction stays ahead of the routed experts, which may write into ``hidden_states`` in place. + + TP: the offload experts are sharded along the intermediate axis and the bf16 shared expert is + row-parallel, so both produce partial sums; ``routed + gate * shared`` is linear in them and + is reduced once (one all-reduce per MoE layer instead of two). """ + def __init__(self, config: ModelConfig, layer_id: int | None = None) -> None: + super().__init__(config, layer_id=layer_id) + self._comm = DistributedCommunicator() + self._tp_size = get_tp_info().size + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: num_tokens, hidden_dim = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_dim) router_logits = self.gate.forward(hidden_states) - shared = self.shared_expert.forward(hidden_states) gate = shared_gate_sigmoid(hidden_states, self.shared_expert_gate.weight.view(-1)) - routed = self.experts.forward(hidden_states=hidden_states, router_logits=router_logits) + se, ex = self.shared_expert, self.experts + if ( + self._tp_size > 1 + and isinstance(se.down_proj, LinearRowParallel) + and isinstance(ex, OffloadMoELayer) + ): + shared = F.linear(silu_and_mul(se.gate_up_proj.forward(hidden_states)), se.down_proj.weight) + if get_global_ctx().batch.is_prefill: + routed = ex.prefill_forward(hidden_states, router_logits) + else: + routed = ex.decode_forward(hidden_states, router_logits) + out = self._comm.all_reduce(shared_gate_mul_add(routed, shared, gate)) + return out.view(num_tokens, hidden_dim) + shared = se.forward(hidden_states) + routed = ex.forward(hidden_states=hidden_states, router_logits=router_logits) return shared_gate_mul_add(routed, shared, gate).view(num_tokens, hidden_dim) diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index 2ce7462b8..97198f408 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -21,12 +21,12 @@ import safetensors import torch from freetoken.distributed import get_tp_info -from freetoken.models.loader import drop_page_cache, iter_weight_files +from freetoken.models.loader import drop_page_cache, iter_weight_files, shard_tensor from freetoken.models.nvfp4_banks import ( Nvfp4ExpertSourceSpec, ) from freetoken.moe.host_banks import HostBank, read_range_into -from freetoken.utils import download_hf_weight +from freetoken.utils import cached_load_hf_config, div_even, download_hf_weight from freetoken.utils.progress import byte_bar from tqdm import tqdm @@ -137,6 +137,58 @@ def _try_fuse( return None +def _shard_rows( + t: torch.Tensor, parts: list[tuple[int, int]], rank: int, world: int +) -> torch.Tensor: + """Column-parallel slice (dim 0) of a ``[part0 | part1 | ...]`` fusion; ``parts`` gives each + part as ``(heads, rows_per_head)``. Heads split evenly across ranks; a part with fewer heads + than ranks (GQA kv) replicates head ``rank * heads // world``, the ``div_even(..., + allow_replicate=True)`` convention of the TP-aware layers.""" + out, off = [], 0 + for heads, rows in parts: + local = div_even(heads, world, allow_replicate=True) + first = rank * heads // world + out.append(t[off + first * rows : off + (first + local) * rows]) + off += heads * rows + assert off == t.shape[0], f"fusion parts {parts} cover {off} rows, tensor has {t.shape[0]}" + return torch.cat(out, dim=0) + + +def _shard(name: str, t: torch.Tensor, config, rank: int, world: int) -> torch.Tensor: + """TP shard of one state-dict tensor (fused projections included); identity at TP=1. + + Column-parallel (dim 0, by head): attention ``qkv_proj`` [q|gate per head | k | v], GDN + ``in_proj`` [q | k | v | z | b | a] and the matching ``conv1d`` channels, ``A_log`` / + ``dt_bias``, shared-expert ``gate_up_proj``. Row-parallel (dim 1): ``o_proj``, + ``out_proj``, shared-expert ``down_proj``. Vocab rows: ``embed_tokens`` / ``lm_head``. + Everything else (router, indexer, norms, HC, PLE, shared-expert gate) is replicated. + """ + if world == 1: + return t + if name.endswith(".self_attn.qkv_proj.weight"): + q = (config.num_qo_heads, 2 * config.head_dim) + kv = (config.num_kv_heads, config.head_dim) + return _shard_rows(t, [q, kv, kv], rank, world) + if ".linear_attn." in name: + g = config.linear_attention_group() + k = (g.num_key_heads, g.key_head_dim) + v = (g.num_value_heads, g.value_head_dim) + if name.endswith(".in_proj.weight"): + return _shard_rows(t, [k, k, v, v, (v[0], 1), (v[0], 1)], rank, world) + if name.endswith(".conv1d.weight"): + return _shard_rows(t, [k, k, v], rank, world) + if name.endswith((".A_log", ".dt_bias")): + return _shard_rows(t, [(v[0], 1)], rank, world) + if name.endswith(".out_proj.weight"): + return t.chunk(world, dim=1)[rank].clone() + return t + if name.endswith(".shared_expert.gate_up_proj.weight"): + half = t.shape[0] // 2 + return _shard_rows(t, [(half, 1), (half, 1)], rank, world) + # o_proj / down_proj: dim 1; embed_tokens / lm_head: vocab rows; others unchanged. + return shard_tensor(name, t, rank=rank, world_size=world, num_kv_heads=None) + + def iter_weights( model_path: str, device: torch.device, @@ -158,16 +210,18 @@ def iter_weights( ``include_moe_experts`` is accepted for the loader contract but never yields anything: the routed experts are NVFP4 and always come from the offload cache's expert reader. """ - if get_tp_info().size > 1: - raise NotImplementedError("qwen4_exp weight loading supports TP=1 only") if not include_non_moe: return + from .config import parse_config + + tp = get_tp_info() + config = parse_config(cached_load_hf_config(model_path)) fuse_buf: dict[str, dict[int, torch.Tensor]] = {} for file in tqdm( iter_weight_files(model_path), desc="Loading weights", - disable=not get_tp_info().is_primary(), + disable=not tp.is_primary(), ): with safetensors.safe_open(file, framework="pt", device=str(device)) as f: for raw_name in f.keys(): @@ -178,9 +232,10 @@ def iter_weights( fused = _try_fuse(name, tensor, fuse_buf) if fused is not None: if fused != (): # () means buffered, not yet complete - yield fused + name, tensor = fused + yield name, _shard(name, tensor, config, tp.rank, tp.size) continue - yield name, tensor + yield name, _shard(name, tensor, config, tp.rank, tp.size) assert not fuse_buf, f"Incomplete projection fusions: {sorted(fuse_buf)}" diff --git a/tests/models/qwen4_exp/test_tp_shard.py b/tests/models/qwen4_exp/test_tp_shard.py new file mode 100644 index 000000000..87a305ef0 --- /dev/null +++ b/tests/models/qwen4_exp/test_tp_shard.py @@ -0,0 +1,92 @@ +"""TP sharding of the qwen4_exp dense weights (pure tensor math, no TP runtime needed).""" + +from types import SimpleNamespace + +import torch +from freetoken.models.qwen4_exp.weight import _shard, _shard_rows + + +def _cfg(): + g = SimpleNamespace( + num_key_heads=4, key_head_dim=8, num_value_heads=12, value_head_dim=8 + ) + return SimpleNamespace( + num_qo_heads=6, num_kv_heads=2, head_dim=8, linear_attention_group=lambda: g + ) + + +def _gather(name, t, cfg, world, dim=0): + return torch.cat([_shard(name, t, cfg, r, world) for r in range(world)], dim=dim) + + +def test_shard_rows_splits_each_part_by_head(): + t = torch.arange(4 * 4 + 2 * 4).reshape( + -1, 1 + ) # part A: 4 heads x 4 rows, part B: 2 heads x 4 + s = [_shard_rows(t, [(4, 4), (2, 4)], r, 4) for r in range(4)] + assert all(x.shape[0] == 8 for x in s) + assert s[0][:4].flatten().tolist() == [0, 1, 2, 3] + assert s[3][:4].flatten().tolist() == [12, 13, 14, 15] + # 2 kv heads over 4 ranks replicate: rank * heads // world -> heads 0, 0, 1, 1 + assert s[0][4:].equal(s[1][4:]) and s[2][4:].flatten().tolist() == [20, 21, 22, 23] + + +def test_shard_round_trips_the_fused_projections(): + cfg, world = _cfg(), 2 + qkv = torch.randn(6 * 16 + 2 * 8 + 2 * 8, 5) + assert ( + _gather("model.layers.0.self_attn.qkv_proj.weight", qkv, cfg, world).shape + == qkv.shape + ) + q0 = _shard("model.layers.0.self_attn.qkv_proj.weight", qkv, cfg, 0, world) + assert q0.shape[0] == 3 * 16 + 8 + 8 + torch.testing.assert_close(q0[:48], qkv[:48]) # q heads 0-2 (16 rows each: q|gate) + torch.testing.assert_close(q0[48:56], qkv[96:104]) # kv head 0 of k + torch.testing.assert_close(q0[56:64], qkv[112:120]) # kv head 0 of v + kd, vd, nv = 4 * 8, 12 * 8, 12 + in_proj = torch.randn(kd + kd + vd + vd + nv + nv, 5) + s = _shard("model.layers.1.linear_attn.in_proj.weight", in_proj, cfg, 1, world) + assert s.shape[0] == (kd + kd + vd + vd + nv + nv) // 2 + torch.testing.assert_close(s[:16], in_proj[16:32]) # q: k heads 2,3 + torch.testing.assert_close(s[-6:], in_proj[-6:]) # a: v heads 6-11 + conv = torch.randn(kd + kd + vd, 1, 4) + assert _shard( + "model.layers.1.linear_attn.conv1d.weight", conv, cfg, 0, world + ).shape == ((kd + kd + vd) // 2, 1, 4) + a_log = torch.randn(nv) + torch.testing.assert_close( + _shard("model.layers.1.linear_attn.A_log", a_log, cfg, 1, world), a_log[6:] + ) + + +def test_shard_row_parallel_and_vocab_and_replicated(): + cfg, world = _cfg(), 2 + for name in ( + "model.layers.0.self_attn.o_proj.weight", + "model.layers.1.linear_attn.out_proj.weight", + "model.layers.0.mlp.shared_expert.down_proj.weight", + ): + t = torch.randn(3, 8) + torch.testing.assert_close(_gather(name, t, cfg, world, dim=1), t) + gate_up = torch.randn(2 * 6, 3) + s1 = _shard( + "model.layers.0.mlp.shared_expert.gate_up_proj.weight", gate_up, cfg, 1, world + ) + torch.testing.assert_close(s1, torch.cat([gate_up[3:6], gate_up[9:12]])) + emb = torch.randn(10, 3) + torch.testing.assert_close( + _gather("model.embed_tokens.weight", emb, cfg, world), emb + ) + torch.testing.assert_close(_gather("lm_head.weight", emb, cfg, world), emb) + for name in ( + "model.layers.0.mlp.gate.weight", + "model.layers.0.self_attn.indexer.index_qk_proj.weight", + "model.layers.0.attn_hyper_connection.input_mix_weight_down_block_inject.weight", + "model.layers.1.ple.value_proj.weight", + "model.layers.0.mlp.shared_expert_gate.weight", + ): + t = torch.randn(4, 6) + assert _shard(name, t, cfg, 1, world).equal(t) + assert _shard("model.layers.0.self_attn.qkv_proj.weight", gate_up, cfg, 0, 1).equal( + gate_up + ) diff --git a/tests/models/test_nvfp4_banks_tp.py b/tests/models/test_nvfp4_banks_tp.py new file mode 100644 index 000000000..13bf2a2c7 --- /dev/null +++ b/tests/models/test_nvfp4_banks_tp.py @@ -0,0 +1,76 @@ +"""TP sharding of the NVFP4 expert source banks: the two ranks' banks concatenated along the +intermediate axis must equal the unsharded placement.""" + +import torch +from freetoken.distributed.info import DistributedInfo +from freetoken.models import nvfp4_banks +from freetoken.models.nvfp4_banks import _alloc_nvfp4_host_banks, _Placer + +E, H, INTER = 2, 64, 32 + + +def _place(monkeypatch, rank, world): + monkeypatch.setattr( + nvfp4_banks, "get_tp_info", lambda: DistributedInfo(rank, world) + ) + hb = _alloc_nvfp4_host_banks(1, E, H, INTER // world) + banks = {name: [b.tensor for b in layers] for name, layers in hb.items()} + placer = _Placer(banks, INTER) + torch.manual_seed(0) + for e in range(E): + for role in ("gate", "up"): + placer.put( + 0, + e, + role, + "weight", + torch.randint(0, 255, (INTER, H // 2), dtype=torch.uint8), + ) + placer.put( + 0, + e, + role, + "weight_scale", + torch.randn(INTER, H // 16).to(torch.float8_e4m3fn), + torch.tensor(0.5 + e, dtype=torch.float16), + ) + placer.put( + 0, + e, + "down", + "weight", + torch.randint(0, 255, (H, INTER // 2), dtype=torch.uint8), + ) + placer.put( + 0, + e, + "down", + "weight_scale", + torch.randn(H, INTER // 16).to(torch.float8_e4m3fn), + torch.tensor(2.0 + e, dtype=torch.float16), + ) + return {k: v[0] for k, v in banks.items()} + + +def test_rank_banks_concatenate_to_the_full_placement(monkeypatch): + full = _place(monkeypatch, 0, 1) + r0, r1 = _place(monkeypatch, 0, 2), _place(monkeypatch, 1, 2) + n = INTER // 2 + for name in ( + "gate_up_packed", + "gate_up_scale", + "gate_up_global", + ): # rows: [gate I | up I] + gate = torch.cat([r0[name][:, :n], r1[name][:, :n]], dim=1) + up = torch.cat([r0[name][:, n:], r1[name][:, n:]], dim=1) + assert torch.equal( + torch.cat([gate, up], dim=1).view(torch.uint8), full[name].view(torch.uint8) + ), name + for name in ("down_packed", "down_scale"): # columns + assert torch.equal( + torch.cat([r0[name], r1[name]], dim=2).view(torch.uint8), + full[name].view(torch.uint8), + ), name + assert torch.equal(r0["down_global"], full["down_global"]) and torch.equal( + r1["down_global"], full["down_global"] + ) From 96e190c3685b5a5fc1f3ac6c75983cb1c5921598 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 4 Sep 2026 15:09:33 -0400 Subject: [PATCH 02/14] fix(qwen4_exp): load the HF config for sharding only when TP > 1 tests/models/qwen4_exp/test_weight.py feeds iter_weights a synthetic checkpoint whose config.json has no model_type; at TP=1 nothing is sharded, so do not touch the config. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt --- python/freetoken/models/qwen4_exp/weight.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index 97198f408..905d32d89 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -216,7 +216,9 @@ def iter_weights( from .config import parse_config tp = get_tp_info() - config = parse_config(cached_load_hf_config(model_path)) + # The sharding geometry needs the HF config; TP=1 never shards, so keep the plain path free + # of a config load (synthetic test checkpoints carry no model_type). + config = parse_config(cached_load_hf_config(model_path)) if tp.size > 1 else None fuse_buf: dict[str, dict[int, torch.Tensor]] = {} for file in tqdm( iter_weight_files(model_path), From a99433c683a3ed214d5baabc1d9b041799d17295 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 10 Sep 2026 14:55:22 -0400 Subject: [PATCH 03/14] test(qwen4_exp): TP shard test against the new expert-piece stream --- tests/models/test_nvfp4_banks_tp.py | 124 +++++++++++++--------------- 1 file changed, 56 insertions(+), 68 deletions(-) diff --git a/tests/models/test_nvfp4_banks_tp.py b/tests/models/test_nvfp4_banks_tp.py index 13bf2a2c7..b943cff1c 100644 --- a/tests/models/test_nvfp4_banks_tp.py +++ b/tests/models/test_nvfp4_banks_tp.py @@ -1,76 +1,64 @@ -"""TP sharding of the NVFP4 expert source banks: the two ranks' banks concatenated along the -intermediate axis must equal the unsharded placement.""" +"""TP sharding of the NVFP4 routed-expert stream. +Every rank reads the same checkpoint tensors and keeps only its own slice of the expert +intermediate axis, so the ranks' slices must partition the original exactly -- the MoE layer +all-reduces the partial sums, which is only correct if nothing is dropped or double-counted. +The expert kernel sizes its banks from ``MoEConfig.local_intermediate``, the same split. +""" + +import pytest import torch -from freetoken.distributed.info import DistributedInfo -from freetoken.models import nvfp4_banks -from freetoken.models.nvfp4_banks import _alloc_nvfp4_host_banks, _Placer +from freetoken.layers.quantization.moe.base import MoEConfig +from freetoken.models.nvfp4_banks import _tp_shard -E, H, INTER = 2, 64, 32 +INTER, HIDDEN = 640, 2560 -def _place(monkeypatch, rank, world): - monkeypatch.setattr( - nvfp4_banks, "get_tp_info", lambda: DistributedInfo(rank, world) - ) - hb = _alloc_nvfp4_host_banks(1, E, H, INTER // world) - banks = {name: [b.tensor for b in layers] for name, layers in hb.items()} - placer = _Placer(banks, INTER) - torch.manual_seed(0) - for e in range(E): - for role in ("gate", "up"): - placer.put( - 0, - e, - role, - "weight", - torch.randint(0, 255, (INTER, H // 2), dtype=torch.uint8), - ) - placer.put( - 0, - e, - role, - "weight_scale", - torch.randn(INTER, H // 16).to(torch.float8_e4m3fn), - torch.tensor(0.5 + e, dtype=torch.float16), - ) - placer.put( - 0, - e, - "down", - "weight", - torch.randint(0, 255, (H, INTER // 2), dtype=torch.uint8), - ) - placer.put( - 0, - e, - "down", - "weight_scale", - torch.randn(H, INTER // 16).to(torch.float8_e4m3fn), - torch.tensor(2.0 + e, dtype=torch.float16), +def _roles(): + """One expert's checkpoint tensors, in the modelopt layout.""" + return { + "gate": torch.arange(INTER * (HIDDEN // 2), dtype=torch.uint8).reshape(INTER, HIDDEN // 2), + "gate_scale": torch.arange(INTER * (HIDDEN // 16), dtype=torch.uint8).reshape(INTER, HIDDEN // 16), + "gate_global": torch.tensor([2.5]), + "up": torch.arange(INTER * (HIDDEN // 2), dtype=torch.uint8).reshape(INTER, HIDDEN // 2), + "up_scale": torch.arange(INTER * (HIDDEN // 16), dtype=torch.uint8).reshape(INTER, HIDDEN // 16), + "up_global": torch.tensor([2.5]), + "down": torch.arange(HIDDEN * (INTER // 2), dtype=torch.uint8).reshape(HIDDEN, INTER // 2), + "down_scale": torch.arange(HIDDEN * (INTER // 16), dtype=torch.uint8).reshape(HIDDEN, INTER // 16), + "down_global": torch.tensor([3.5]), + } + + +@pytest.mark.parametrize("tp_size", [1, 2, 4]) +def test_rank_slices_partition_the_original_tensor(tp_size): + full = _roles() + for role, tensor in full.items(): + parts = [_tp_shard(role, tensor, INTER, tp_size, r) for r in range(tp_size)] + if role.endswith("_global"): + # per-tensor scalar: no I axis, every rank keeps it whole + assert all(torch.equal(p, tensor) for p in parts) + continue + axis = 1 if role.startswith("down") else 0 + assert torch.equal(torch.cat(parts, dim=axis), tensor), role + assert all(p.shape[axis] == tensor.shape[axis] // tp_size for p in parts), role + + +def test_tp1_is_the_identity(): + for role, tensor in _roles().items(): + assert _tp_shard(role, tensor, INTER, 1, 0) is tensor or torch.equal( + _tp_shard(role, tensor, INTER, 1, 0), tensor ) - return {k: v[0] for k, v in banks.items()} -def test_rank_banks_concatenate_to_the_full_placement(monkeypatch): - full = _place(monkeypatch, 0, 1) - r0, r1 = _place(monkeypatch, 0, 2), _place(monkeypatch, 1, 2) - n = INTER // 2 - for name in ( - "gate_up_packed", - "gate_up_scale", - "gate_up_global", - ): # rows: [gate I | up I] - gate = torch.cat([r0[name][:, :n], r1[name][:, :n]], dim=1) - up = torch.cat([r0[name][:, n:], r1[name][:, n:]], dim=1) - assert torch.equal( - torch.cat([gate, up], dim=1).view(torch.uint8), full[name].view(torch.uint8) - ), name - for name in ("down_packed", "down_scale"): # columns - assert torch.equal( - torch.cat([r0[name], r1[name]], dim=2).view(torch.uint8), - full[name].view(torch.uint8), - ), name - assert torch.equal(r0["down_global"], full["down_global"]) and torch.equal( - r1["down_global"], full["down_global"] - ) +def test_a_shard_that_splits_a_scale_block_is_rejected(): + # scales cover 16 values; a rank slice that is not a whole number of blocks cannot be + # sliced consistently across the weight and its scale + with pytest.raises(AssertionError, match="16-wide scale blocks"): + _tp_shard("gate", torch.zeros(24, 8, dtype=torch.uint8), 24, 3, 0) + + +def test_moe_config_local_intermediate_matches_the_shard(): + for tp_size in (1, 2, 4): + cfg = MoEConfig(num_experts=8, hidden=HIDDEN, intermediate=INTER, top_k=2, tp_size=tp_size) + assert cfg.local_intermediate == INTER // tp_size + assert _tp_shard("gate", _roles()["gate"], INTER, tp_size, 0).shape[0] == cfg.local_intermediate From be41deae6bff788f3392e45477452e450720c834 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 4 Sep 2026 18:55:34 -0400 Subject: [PATCH 04/14] feat(qwen4_exp): load-time per-tensor FP8 dense projections (W8A8 via _scaled_mm) Opt-in with FREETOKEN_FP8_DENSE=1 on a bf16-dense checkpoint (e.g. the RadixArk NVFP4 build): the weight reader quantizes qkv_proj / o_proj, GDN in_proj (q|k|v|z; the b|a gate rows stay bf16 as in_proj_ba) and out_proj to per-tensor e4m3 after TP sharding, and layers/fp8_dynamic.py runs them as cuBLASLt W8A8 GEMMs with a dynamic per-tensor activation scale (one fused Triton launch at decode sizes; no host sync, CUDA-graph safe). Column-merged and row-parallel variants, so it works at TP>1. Why: on an RTX 6000 Ada (sm_89, torch 2.11.0+cu130) these projections are 2.67 GB of the ~4 GB a TP=2 rank reads per token; bf16 cuBLAS takes 3.2-3.4 ms per step per rank, raw _scaled_mm 1.9 ms, while the existing Triton FP8 kernels are slower than bf16 there (measured, weights rotated past the L2). Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt --- python/freetoken/layers/fp8_dynamic.py | 154 ++++++++++++++++++ python/freetoken/models/config.py | 7 + .../freetoken/models/qwen4_exp/attention.py | 37 +++-- python/freetoken/models/qwen4_exp/config.py | 8 + python/freetoken/models/qwen4_exp/gdn.py | 48 ++++-- python/freetoken/models/qwen4_exp/model.py | 1 + python/freetoken/models/qwen4_exp/weight.py | 67 +++++++- tests/models/qwen4_exp/test_fp8_dense.py | 121 ++++++++++++++ 8 files changed, 409 insertions(+), 34 deletions(-) create mode 100644 python/freetoken/layers/fp8_dynamic.py create mode 100644 tests/models/qwen4_exp/test_fp8_dense.py diff --git a/python/freetoken/layers/fp8_dynamic.py b/python/freetoken/layers/fp8_dynamic.py new file mode 100644 index 000000000..1a956fbab --- /dev/null +++ b/python/freetoken/layers/fp8_dynamic.py @@ -0,0 +1,154 @@ +"""Load-time per-tensor FP8 dense linear (W8A8 through cuBLASLt ``torch._scaled_mm``). + +The weight is e4m3 ``[out_local, in_local]`` with one fp32 ``weight_scale`` (shape ``()``), +produced by the model's weight reader from the bf16 checkpoint tensor (after TP sharding). +The activation is quantized per call with a dynamic per-tensor scale: one fused Triton launch +(amax pass, then cast) at decode sizes, a torch reduction plus a cast kernel above that. No +host sync anywhere (the scales stay on the device), so the decode path is CUDA-graph safe; +the branch between the two paths is on the tensor *shape*, never on its values. + +Measured on an RTX 6000 Ada (sm_89, torch 2.11.0+cu130) at the qwen4_exp TP=2 shapes, per +decode step per rank over 48 layers: bf16 cuBLAS 3.2-3.4 ms, this path 1.9 ms at M=1/8/16. +Requires sm_89+ (``_scaled_mm``'s floor; Ampere has no FP8 tensor cores). +""" + +from __future__ import annotations + +from typing import List + +import torch +import triton +import triton.language as tl +from freetoken.distributed import DistributedCommunicator, get_tp_info +from freetoken.utils import div_even + +from .base import BaseOP + +FP8 = torch.float8_e4m3fn +E4M3_MAX = 448.0 +_FUSED_MAX_ELEMENTS = ( + 65536 # one program handles the whole tensor below this (decode sizes) +) +_SCALE_FLOOR = 1e-12 # an all-zero activation (graph warmup buffers) must not give 1/0 + + +@triton.jit +def _quant_fused_kernel(x_ptr, out_ptr, scale_ptr, n, BLOCK: tl.constexpr): + """One program: max|x| over the tensor, then the cast. Decode-sized inputs only.""" + acc = tl.zeros([BLOCK], dtype=tl.float32) + for start in range(0, n, BLOCK): + offs = start + tl.arange(0, BLOCK) + v = tl.load(x_ptr + offs, mask=offs < n, other=0.0).to(tl.float32) + acc = tl.maximum(acc, tl.abs(v)) + amax = tl.maximum(tl.max(acc, axis=0), 1e-12) + tl.store(scale_ptr, amax / 448.0) + inv = 448.0 / amax + for start in range(0, n, BLOCK): + offs = start + tl.arange(0, BLOCK) + mask = offs < n + v = tl.load(x_ptr + offs, mask=mask, other=0.0).to(tl.float32) * inv + v = tl.minimum(tl.maximum(v, -448.0), 448.0) + tl.store(out_ptr + offs, v.to(tl.float8e4nv), mask=mask) + + +@triton.jit +def _quant_cast_kernel(x_ptr, out_ptr, scale_ptr, n, BLOCK: tl.constexpr): + """Cast under a scale already on the device (prefill sizes; the amax is a torch reduction).""" + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = offs < n + inv = 1.0 / tl.load(scale_ptr) + v = tl.load(x_ptr + offs, mask=mask, other=0.0).to(tl.float32) * inv + v = tl.minimum(tl.maximum(v, -448.0), 448.0) + tl.store(out_ptr + offs, v.to(tl.float8e4nv), mask=mask) + + +def quant_per_tensor(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """``(x_fp8, scale)`` with ``x ~= x_fp8 * scale``; ``x`` contiguous, ``scale`` fp32 ``()``.""" + n = x.numel() + out = torch.empty_like(x, dtype=FP8) + if n <= _FUSED_MAX_ELEMENTS: + scale = torch.empty((), dtype=torch.float32, device=x.device) + _quant_fused_kernel[(1,)](x, out, scale, n, BLOCK=4096, num_warps=8) + return out, scale + amax = torch.linalg.vector_norm(x, ord=float("inf")).float() + scale = amax.clamp_min_(_SCALE_FLOOR).div_(E4M3_MAX) + _quant_cast_kernel[(triton.cdiv(n, 4096),)]( + x, out, scale, n, BLOCK=4096, num_warps=4 + ) + return out, scale + + +def fp8_dynamic_linear( + x: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor +) -> torch.Tensor: + """``x @ (weight * weight_scale)^T`` in W8A8; ``weight`` [N, K] e4m3 row-major, whose ``.t()`` + is the column-major operand cuBLASLt wants (a stride change, never a copy).""" + *lead, k = x.shape + x2 = x.reshape(-1, k).contiguous() + x8, scale = quant_per_tensor(x2) + y = torch._scaled_mm( + x8, weight.t(), scale_a=scale, scale_b=weight_scale, out_dtype=x.dtype + ) + return y.reshape(*lead, weight.shape[0]) + + +class Fp8DynamicLinear(BaseOP): + """Per-tensor FP8 linear over the local shard; ``all_reduce`` adds the TP sum (row-parallel).""" + + def __init__(self, local_isize: int, local_osize: int, *, all_reduce: bool = False): + assert local_isize % 16 == 0 and local_osize % 16 == 0, ( + local_isize, + local_osize, + ) + self.local_input_size = local_isize + self.local_output_size = local_osize + self.weight = torch.empty(local_osize, local_isize, dtype=FP8) + self.weight_scale = torch.empty((), dtype=torch.float32) + self._comm = ( + DistributedCommunicator() if all_reduce and get_tp_info().size > 1 else None + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = fp8_dynamic_linear(x, self.weight, self.weight_scale) + if self._comm is not None: + y = self._comm.all_reduce(y) + return y + + +class Fp8DynamicColMerged(Fp8DynamicLinear): + """Drop-in for ``LinearColParallelMerged``: one weight concatenating several projections + along the output dim; the caller splits the output by the local sizes as before.""" + + def __init__( + self, + input_size: int, + output_sizes: List[int], + local_output_sizes: List[int] | None = None, + ): + tp = get_tp_info() + if local_output_sizes is None: + local_output_sizes = [div_even(size, tp.size) for size in output_sizes] + self.output_sizes = list(output_sizes) + self.local_output_sizes = list(local_output_sizes) + super().__init__(input_size, sum(local_output_sizes)) + + +class Fp8DynamicRowParallel(Fp8DynamicLinear): + """Drop-in for ``LinearOProj`` / ``LinearRowParallel``: the input dim is sharded, the + all-reduce runs after the local GEMM (each rank scales its own shard).""" + + def __init__(self, input_size: int, output_size: int): + super().__init__( + div_even(input_size, get_tp_info().size), output_size, all_reduce=True + ) + + +__all__ = [ + "FP8", + "E4M3_MAX", + "Fp8DynamicColMerged", + "Fp8DynamicLinear", + "Fp8DynamicRowParallel", + "fp8_dynamic_linear", + "quant_per_tensor", +] diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py index 07a16df5c..0ad73b752 100644 --- a/python/freetoken/models/config.py +++ b/python/freetoken/models/config.py @@ -18,6 +18,13 @@ def vision_load_enabled() -> bool: return os.getenv("FREETOKEN_LOAD_VISION", "0").strip().lower() in _VISION_TRUE +def fp8_dense_enabled() -> bool: + """Load-time FP8 for a model's bf16 attention / GDN projections (opt-in, default OFF): + per-tensor e4m3 weights, per-tensor dynamic activation scale, cuBLASLt W8A8 GEMMs + (``torch._scaled_mm``, sm_89+). ``FREETOKEN_FP8_DENSE=1``.""" + return os.getenv("FREETOKEN_FP8_DENSE", "0").strip().lower() in _VISION_TRUE + + def detect_expert_quant(hf_config: Any) -> str: """Routed-expert quantization from a checkpoint's ``quantization_config``: ``"nvfp4"`` for a ModelOpt FP4 build (``quant_algo: NVFP4``) OR an llm-compressor NVFP4 export diff --git a/python/freetoken/models/qwen4_exp/attention.py b/python/freetoken/models/qwen4_exp/attention.py index 651da87d9..5b7e90397 100644 --- a/python/freetoken/models/qwen4_exp/attention.py +++ b/python/freetoken/models/qwen4_exp/attention.py @@ -27,6 +27,7 @@ LinearOProj, LinearReplicated, ) +from freetoken.layers.fp8_dynamic import Fp8DynamicColMerged, Fp8DynamicRowParallel from freetoken.layers.rotary import get_rope from freetoken.utils import div_even, nvtx_annotate @@ -138,20 +139,30 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = "") -> N 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.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", - ) - # 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", + qkv_sizes = [self.qo_attn_dim * 2, self.kv_attn_dim, self.kv_attn_dim] + # Load-time per-tensor FP8 applies only where the checkpoint declares no scheme of + # its own; a quantized checkpoint always wins, and its modules go through QuantConfig. + fp8_dyn = config.attn_quant == "fp8_dynamic" and ( + config.quant is None or config.quant.scheme_for(f"{prefix}.qkv_proj") is None ) + if fp8_dyn: # W8A8 through _scaled_mm, weights quantized by the loader + self.qkv_proj = Fp8DynamicColMerged( + config.hidden_size, qkv_sizes, local_output_sizes=self._qkv_split + ) + self.o_proj = Fp8DynamicRowParallel(self.qo_attn_dim, config.hidden_size) + else: + self.qkv_proj = LinearColParallelMerged( + config.hidden_size, qkv_sizes, has_bias=False, + local_output_sizes=self._qkv_split, + quant_config=config.quant, prefix=f"{prefix}.qkv_proj", + ) + # 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. 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", + ) self.q_norm = GemmaPlusOneRMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = GemmaPlusOneRMSNorm(self.head_dim, eps=config.rms_norm_eps) rotary = config.rotary_config diff --git a/python/freetoken/models/qwen4_exp/config.py b/python/freetoken/models/qwen4_exp/config.py index 4bcd28de3..e8b62766a 100644 --- a/python/freetoken/models/qwen4_exp/config.py +++ b/python/freetoken/models/qwen4_exp/config.py @@ -7,6 +7,7 @@ from freetoken.layers.quantization import QuantConfig from freetoken.models.config import ( + fp8_dense_enabled, FullAttentionGroupConfig, LinearGatedDeltaGroupConfig, ModelConfig, @@ -138,6 +139,12 @@ def parse_config(hf_config: Any) -> ModelConfig: else {k: v for k, v in rope_params.items() if not isinstance(v, (list, dict))} ) + # FREETOKEN_FP8_DENSE=1: the bf16 attention / GDN projections are quantized at LOAD to + # per-tensor e4m3 and served W8A8 through cuBLASLt (layers/fp8_dynamic.py). This is a + # synthetic scheme for a checkpoint that ships those modules unquantized -- distinct from + # the QuantConfig path, which serves what the checkpoint declares. The module builders + # apply it only where the checkpoint has no scheme of its own, so the two never collide. + attn_quant = "fp8_dynamic" if fp8_dense_enabled() else "none" layer_types = _layer_types(text) full_ids = tuple(i for i, t in enumerate(layer_types) if t == "full_attention") linear_ids = tuple(i for i, t in enumerate(layer_types) if t == "linear_attention") @@ -245,6 +252,7 @@ def parse_config(hf_config: Any) -> ModelConfig: image_token_id=getattr(hf_config, "image_token_id", None), attention_groups=groups, expert_quant=expert_quant, + attn_quant=attn_quant, qwen4_args=qwen4_args, slot_states=ple_slot_states(qwen4_args), ) diff --git a/python/freetoken/models/qwen4_exp/gdn.py b/python/freetoken/models/qwen4_exp/gdn.py index 3b2e02934..f0d95c3d7 100644 --- a/python/freetoken/models/qwen4_exp/gdn.py +++ b/python/freetoken/models/qwen4_exp/gdn.py @@ -9,6 +9,7 @@ from freetoken.layers import BaseOP, GatedRMSNorm, LinearColParallelMerged, LinearOProj from freetoken.layers.quantization import QuantConfig from freetoken.utils import div_even +from freetoken.layers.fp8_dynamic import Fp8DynamicColMerged, Fp8DynamicRowParallel from freetoken.models.qwen3_5_moe.gdn_kernels import gdn_decode_fla, gdn_prefill_chunk_fla @@ -36,6 +37,7 @@ def __init__( self, hidden_size, num_k_heads, num_v_heads, head_k_dim, head_v_dim, conv_kernel_size, rms_norm_eps, layer_id, output_gate: str = "sigmoid", *, quant_config: QuantConfig | None = None, prefix: str = "", + attn_quant: str = "none", ): self.layer_id = layer_id # The fla chunk/decode kernels read+write the recurrent state and the per-chunk h as @@ -66,17 +68,29 @@ def __init__( self._split_in_proj = ( quant_config is not None and quant_config.scheme_for(f"{prefix}.in_proj_qkvz") is not None ) + # Load-time per-tensor FP8 splits the fusion the same way, but it is synthetic: it + # applies only where the checkpoint declares no scheme of its own. + self._dynamic_fp8 = attn_quant == "fp8_dynamic" and not self._split_in_proj 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. + # A checkpoint-declared scheme has no TP-aware variant here: those two linears are + # built from the checkpoint's own scheme and are not sharded. Load-time fp8 is. assert not (tp.size > 1 and self._split_in_proj), ( - "qwen4_exp TP shards bf16 GDN projections only" + "qwen4_exp TP shards bf16 or load-time fp8 GDN projections only" ) - if self._split_in_proj: + if self._dynamic_fp8: + self.in_proj_qkvz = Fp8DynamicColMerged( + hidden_size, [self.conv_dim, self.value_dim], + local_output_sizes=[self._local_conv_dim, self._local_value_dim], + ) + self.in_proj_ba = LinearColParallelMerged( + hidden_size, [num_v_heads, num_v_heads], has_bias=False, + local_output_sizes=[self._local_num_v_heads, self._local_num_v_heads], + ) + elif self._split_in_proj: self.in_proj_qkvz = LinearColParallelMerged( hidden_size, [self.conv_dim, self.value_dim], has_bias=False, quant_config=quant_config, prefix=f"{prefix}.in_proj_qkvz", @@ -100,14 +114,17 @@ def __init__( 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) - # 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", - ) + # Row-parallel over the local v heads, all-reduce inside. At TP=1 LinearOProj 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. + if self._dynamic_fp8: + self.out_proj = Fp8DynamicRowParallel(self.value_dim, hidden_size) + else: + self.out_proj = LinearOProj( + self.value_dim, hidden_size, has_bias=False, + quant_config=quant_config, prefix=f"{prefix}.out_proj", + ) def _gate_params(self, a: torch.Tensor, b: torch.Tensor): beta = b.sigmoid() @@ -165,15 +182,14 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: fla = build_fla_metadata(batch, hidden_states.device) batch.fla_metadata = fla - if self._split_in_proj: + if self._split_in_proj or self._dynamic_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._in_proj_split[:2], 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, [nv, nv], dim=-1) else: proj = self.in_proj.forward(hidden_states) conv_in, z, b, a = torch.split(proj, self._in_proj_split, dim=-1) - 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) diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py index a4b63362d..c9cf9c131 100644 --- a/python/freetoken/models/qwen4_exp/model.py +++ b/python/freetoken/models/qwen4_exp/model.py @@ -50,6 +50,7 @@ def build_linear_mixer(config: ModelConfig, layer_id: int, prefix: str) -> BaseO output_gate=g.output_gate, quant_config=config.quant, prefix=prefix, + attn_quant=config.attn_quant, ) diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index 905d32d89..71aad6154 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -189,6 +189,48 @@ def _shard(name: str, t: torch.Tensor, config, rank: int, world: int) -> torch.T return shard_tensor(name, t, rank=rank, world_size=world, num_kv_heads=None) +# Load-time per-tensor FP8 (attn_quant == "fp8_dynamic", layers/fp8_dynamic.py): these keep +# their name and gain a sibling ``weight_scale``; GDN ``in_proj`` splits into the fp8 +# ``in_proj_qkvz`` and the bf16 ``in_proj_ba`` (the gate projections stay bf16, as in the +# block-fp8 checkpoints and in sglang / vLLM). +_FP8_DENSE_SUFFIXES = ( + ".self_attn.qkv_proj.weight", + ".self_attn.o_proj.weight", + ".linear_attn.out_proj.weight", +) +_E4M3_MAX = 448.0 + + +def _quantize_per_tensor(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """``(e4m3 weight, fp32 scale ())`` with ``w ~= weight * scale``.""" + w = w.float() + scale = (w.abs().amax() / _E4M3_MAX).clamp_min(1e-12) + return (w / scale).clamp_(-_E4M3_MAX, _E4M3_MAX).to(torch.float8_e4m3fn), scale.reshape(()) + + +def _fp8_dense( + name: str, t: torch.Tensor, config, world: int +) -> Iterator[tuple[str, torch.Tensor]]: + """The (already sharded) dense tensor as the fp8_dynamic model expects it.""" + if name.endswith(".linear_attn.in_proj.weight"): + g = config.linear_attention_group() + nk = div_even(g.num_key_heads, world, allow_replicate=True) + nv = div_even(g.num_value_heads, world, allow_replicate=True) + qkvz = 2 * nk * g.key_head_dim + 2 * nv * g.value_head_dim # [q | k | v | z] local rows + assert t.shape[0] == qkvz + 2 * nv, (name, t.shape, qkvz, nv) + base = name[: -len("in_proj.weight")] + w8, scale = _quantize_per_tensor(t[:qkvz]) + yield base + "in_proj_qkvz.weight", w8 + yield base + "in_proj_qkvz.weight_scale", scale + yield base + "in_proj_ba.weight", t[qkvz:].contiguous() + elif name.endswith(_FP8_DENSE_SUFFIXES): + w8, scale = _quantize_per_tensor(t) + yield name, w8 + yield name[: -len("weight")] + "weight_scale", scale + else: + yield name, t + + def iter_weights( model_path: str, device: torch.device, @@ -213,12 +255,27 @@ def iter_weights( if not include_non_moe: return + from freetoken.models.config import fp8_dense_enabled + from .config import parse_config tp = get_tp_info() - # The sharding geometry needs the HF config; TP=1 never shards, so keep the plain path free - # of a config load (synthetic test checkpoints carry no model_type). - config = parse_config(cached_load_hf_config(model_path)) if tp.size > 1 else None + # The sharding geometry (and the fp8 split) need the HF config; TP=1 bf16 never does, so + # keep that path free of a config load (synthetic test checkpoints carry no model_type). + config = ( + parse_config(cached_load_hf_config(model_path)) + if tp.size > 1 or fp8_dense_enabled() + else None + ) + fp8 = config is not None and config.attn_quant == "fp8_dynamic" + + def emit(name: str, tensor: torch.Tensor): + tensor = _shard(name, tensor, config, tp.rank, tp.size) + if fp8: + yield from _fp8_dense(name, tensor, config, tp.size) + else: + yield name, tensor + fuse_buf: dict[str, dict[int, torch.Tensor]] = {} for file in tqdm( iter_weight_files(model_path), @@ -235,9 +292,9 @@ def iter_weights( if fused is not None: if fused != (): # () means buffered, not yet complete name, tensor = fused - yield name, _shard(name, tensor, config, tp.rank, tp.size) + yield from emit(name, tensor) continue - yield name, _shard(name, tensor, config, tp.rank, tp.size) + yield from emit(name, tensor) assert not fuse_buf, f"Incomplete projection fusions: {sorted(fuse_buf)}" diff --git a/tests/models/qwen4_exp/test_fp8_dense.py b/tests/models/qwen4_exp/test_fp8_dense.py new file mode 100644 index 000000000..058562e5c --- /dev/null +++ b/tests/models/qwen4_exp/test_fp8_dense.py @@ -0,0 +1,121 @@ +"""Load-time per-tensor FP8 for the qwen4_exp dense projections (FREETOKEN_FP8_DENSE=1).""" + +from types import SimpleNamespace + +import pytest +import torch +from freetoken.models.qwen4_exp.weight import _fp8_dense, _quantize_per_tensor + +from .common import requires_cuda + +E4M3_MAX = 448.0 + + +def _cfg(): + g = SimpleNamespace( + num_key_heads=4, key_head_dim=8, num_value_heads=12, value_head_dim=8 + ) + return SimpleNamespace(linear_attention_group=lambda: g) + + +def _assert_e4m3_close( + deq: torch.Tensor, ref: torch.Tensor, scale: torch.Tensor +) -> None: + # e4m3 keeps 3 mantissa bits (rel 2^-4); below scale * 2^-6 it is subnormal (abs 2^-9 steps) + tol = ref.abs() * 2**-4 + float(scale) * 2**-9 + 1e-7 + assert ((deq - ref).abs() <= tol).all() + + +def test_quantize_per_tensor_round_trips_within_e4m3(): + w = torch.randn(64, 32) * 0.02 + w8, scale = _quantize_per_tensor(w) + assert ( + w8.dtype == torch.float8_e4m3fn + and scale.dtype == torch.float32 + and scale.shape == () + ) + assert torch.isclose(scale * E4M3_MAX, w.abs().max()) + _assert_e4m3_close(w8.float() * scale, w, scale) + + +def test_in_proj_splits_into_fp8_qkvz_and_bf16_ba_per_rank(): + cfg, world = ( + _cfg(), + 2, + ) # local: 2 k heads, 6 v heads -> qkvz = 2*2*8 + 2*6*8 = 128, ba = 12 + t = torch.randn(128 + 12, 16, dtype=torch.bfloat16) + out = dict(_fp8_dense("model.layers.3.linear_attn.in_proj.weight", t, cfg, world)) + assert sorted(out) == [ + "model.layers.3.linear_attn.in_proj_ba.weight", + "model.layers.3.linear_attn.in_proj_qkvz.weight", + "model.layers.3.linear_attn.in_proj_qkvz.weight_scale", + ] + w8 = out["model.layers.3.linear_attn.in_proj_qkvz.weight"] + scale = out["model.layers.3.linear_attn.in_proj_qkvz.weight_scale"] + assert w8.shape == (128, 16) and w8.dtype == torch.float8_e4m3fn + _assert_e4m3_close(w8.float() * scale, t[:128].float(), scale) + ba = out["model.layers.3.linear_attn.in_proj_ba.weight"] + assert ba.dtype == torch.bfloat16 and torch.equal(ba, t[128:]) + + +def test_other_projections_gain_a_scale_and_the_rest_pass_through(): + cfg = _cfg() + t = torch.randn(32, 16, dtype=torch.bfloat16) + for name in ( + "x.self_attn.qkv_proj.weight", + "x.self_attn.o_proj.weight", + "x.linear_attn.out_proj.weight", + ): + out = dict(_fp8_dense(name, t, cfg, 1)) + assert sorted(out) == [name, name[: -len("weight")] + "weight_scale"] + assert out[name].dtype == torch.float8_e4m3fn + out = dict(_fp8_dense("x.mlp.shared_expert.gate_up_proj.weight", t, cfg, 1)) + assert ( + list(out) == ["x.mlp.shared_expert.gate_up_proj.weight"] + and out.popitem()[1] is t + ) + + +def test_ops_declare_fp8_weight_and_scalar_scale(): + from freetoken.distributed import set_tp_info, try_get_tp_info + from freetoken.layers.fp8_dynamic import Fp8DynamicColMerged, Fp8DynamicRowParallel + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + col = Fp8DynamicColMerged(32, [64, 16, 16]) + row = Fp8DynamicRowParallel(64, 32) + for op in (col, row): + sd = op.state_dict() + assert set(sd) == {"weight", "weight_scale"} + assert ( + sd["weight"].dtype == torch.float8_e4m3fn and sd["weight_scale"].shape == () + ) + assert col.weight.shape == (96, 32) and row.weight.shape == (32, 64) + + +@requires_cuda +@pytest.mark.parametrize("rows", [1, 16, 300], ids=["decode-1", "decode-16", "prefill"]) +def test_fp8_linear_matches_bf16_on_the_dequantized_weight(rows: int): + from freetoken.layers.fp8_dynamic import fp8_dynamic_linear + + torch.manual_seed(0) + w = (torch.randn(256, 128, device="cuda") * 0.02).to(torch.bfloat16) + w8, scale = _quantize_per_tensor(w) + x = torch.randn(rows, 128, device="cuda", dtype=torch.bfloat16) + ref = torch.nn.functional.linear(x, (w8.float() * scale).to(torch.bfloat16)) + got = fp8_dynamic_linear(x, w8, scale.to("cuda")) + assert got.shape == ref.shape and got.dtype == torch.bfloat16 + # the per-tensor activation cast is the only extra rounding: 2^-4 relative on the inputs + torch.testing.assert_close( + got.float(), ref.float(), rtol=0.1, atol=0.08 * ref.abs().max().item() + ) + + +@requires_cuda +def test_quant_per_tensor_zero_input_is_finite(): + from freetoken.layers.fp8_dynamic import quant_per_tensor + + x8, scale = quant_per_tensor( + torch.zeros(16, 128, device="cuda", dtype=torch.bfloat16) + ) + assert torch.isfinite(scale).item() and (x8.float() == 0).all() From 0968c24011915d10682c20326f5cabec20261c77 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 4 Sep 2026 18:57:14 -0400 Subject: [PATCH 05/14] chore(qwen4_exp): log the load-time FP8 dense mode once at weight load Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt --- python/freetoken/models/qwen4_exp/weight.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index 71aad6154..e7c67658f 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -26,10 +26,12 @@ Nvfp4ExpertSourceSpec, ) from freetoken.moe.host_banks import HostBank, read_range_into -from freetoken.utils import cached_load_hf_config, div_even, download_hf_weight +from freetoken.utils import cached_load_hf_config, div_even, download_hf_weight, init_logger from freetoken.utils.progress import byte_bar from tqdm import tqdm +logger = init_logger(__name__) + # Routed NVFP4 experts (nvidia modelopt layout): per-expert, un-fused. Matched against the RAW # weight_map key in nvfp4_banks. The ``model.language_model.`` anchor excludes the MTP head's # stacked ``mtp.layers.N.mlp.experts.*`` tensors. @@ -268,6 +270,11 @@ def iter_weights( else None ) fp8 = config is not None and config.attn_quant == "fp8_dynamic" + if fp8: + logger.info( + "qwen4_exp dense projections: load-time per-tensor FP8 (W8A8 via _scaled_mm), " + "FREETOKEN_FP8_DENSE=1" + ) def emit(name: str, tensor: torch.Tensor): tensor = _shard(name, tensor, config, tp.rank, tp.size) From 8d37231d3793499bec36c53d6cc939641f1b61f9 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 4 Sep 2026 19:09:48 -0400 Subject: [PATCH 06/14] fix(qwen4_exp): return the FP8 quantization slack to the allocator before the cache planner runs Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt --- python/freetoken/models/qwen4_exp/weight.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index e7c67658f..7f347cec7 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -304,6 +304,11 @@ def emit(name: str, tensor: torch.Tensor): yield from emit(name, tensor) assert not fuse_buf, f"Incomplete projection fusions: {sorted(fuse_buf)}" + if fp8 and device.type == "cuda": + # The bf16 originals and fp32 temporaries of the quantization sit in the caching + # allocator; hand them back so the expert-cache planner (free VRAM after load) sees + # the halved dense footprint instead of the slack. + torch.cuda.empty_cache() # ====================================================================================== From 1a6f504c05d061a6af577c6ac23103897136e92c Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 4 Sep 2026 19:23:02 -0400 Subject: [PATCH 07/14] fix(qwen4_exp): give in_proj_ba its own storage under FP8 dense t[qkvz:].contiguous() on a contiguous row slice returns a view, so every GDN layer's bf16 gate rows kept the whole sharded bf16 in_proj resident next to the fp8 copy: 36 x 42 MB = 1.5 GiB per TP=2 rank, which is why the expert cache planner saw no saving (22,594 -> 22,458 slots) after the FP8 switch. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt --- python/freetoken/models/qwen4_exp/weight.py | 4 +++- tests/models/qwen4_exp/test_fp8_dense.py | 2 ++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index 7f347cec7..57f9d5f4f 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -224,7 +224,9 @@ def _fp8_dense( w8, scale = _quantize_per_tensor(t[:qkvz]) yield base + "in_proj_qkvz.weight", w8 yield base + "in_proj_qkvz.weight_scale", scale - yield base + "in_proj_ba.weight", t[qkvz:].contiguous() + # clone, not contiguous(): a contiguous row slice IS contiguous, so .contiguous() would + # hand back a view that keeps the whole bf16 in_proj (36 x 42 MB per rank) resident + yield base + "in_proj_ba.weight", t[qkvz:].clone() elif name.endswith(_FP8_DENSE_SUFFIXES): w8, scale = _quantize_per_tensor(t) yield name, w8 diff --git a/tests/models/qwen4_exp/test_fp8_dense.py b/tests/models/qwen4_exp/test_fp8_dense.py index 058562e5c..fac51bdb5 100644 --- a/tests/models/qwen4_exp/test_fp8_dense.py +++ b/tests/models/qwen4_exp/test_fp8_dense.py @@ -56,6 +56,8 @@ def test_in_proj_splits_into_fp8_qkvz_and_bf16_ba_per_rank(): _assert_e4m3_close(w8.float() * scale, t[:128].float(), scale) ba = out["model.layers.3.linear_attn.in_proj_ba.weight"] assert ba.dtype == torch.bfloat16 and torch.equal(ba, t[128:]) + # its own storage: a view would keep the whole bf16 in_proj alive next to the fp8 copy + assert ba.untyped_storage().data_ptr() != t.untyped_storage().data_ptr() def test_other_projections_gain_a_scale_and_the_rest_pass_through(): From 74ccbea7a3cc60bdb90f0e28109d6af862913bc8 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Sat, 5 Sep 2026 17:30:28 -0400 Subject: [PATCH 08/14] perf(fp8): quantize activations with a block-parallel amax quant_per_tensor ran a single Triton program over the whole tensor, so its cost grew linearly with the input: 9.8 us for an [8, 10240] activation and 68.6 us at [64, 10240], against ~1.2 us of useful work. Above 16384 elements, split it into a block-parallel partial amax and a reduce+cast, which flattens the cost to ~2.4 us. The arithmetic is the one the single-program kernel already used, so a given tensor quantizes bit for bit as before; the old three-launch torch reduction path above 65536 elements goes away with it. The partial count is a runtime argument and the pass-1 grid is strided and capped at _MAX_PARTS, so both kernels have exactly ONE compiled Triton variant. Letting the partial count reach the kernel as a constexpr instead costs a fresh compilation for every distinct input length -- unbounded variant growth in a server that sees arbitrary prompt lengths, and a compile stall mid-generation. Measured on an RTX 6000 Ada (sm_89, torch 2.11) under CUDA-graph capture, one-program -> split: 20480 elts 2.59 -> 2.22 us, 81920 elts 9.83 -> 2.44 us, 655360 elts 68.60 -> 3.06 us. Below the threshold the single program still wins (2560 elts: 1.25 vs 2.07 us) and is kept. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt (cherry picked from commit d254729340b4ac87d02d6956b57069904cb0e73d) --- python/freetoken/layers/fp8_dynamic.py | 79 +++++++++++++++++------- tests/models/qwen4_exp/test_fp8_dense.py | 17 +++++ 2 files changed, 75 insertions(+), 21 deletions(-) diff --git a/python/freetoken/layers/fp8_dynamic.py b/python/freetoken/layers/fp8_dynamic.py index 1a956fbab..5e84a929c 100644 --- a/python/freetoken/layers/fp8_dynamic.py +++ b/python/freetoken/layers/fp8_dynamic.py @@ -2,10 +2,12 @@ The weight is e4m3 ``[out_local, in_local]`` with one fp32 ``weight_scale`` (shape ``()``), produced by the model's weight reader from the bf16 checkpoint tensor (after TP sharding). -The activation is quantized per call with a dynamic per-tensor scale: one fused Triton launch -(amax pass, then cast) at decode sizes, a torch reduction plus a cast kernel above that. No -host sync anywhere (the scales stay on the device), so the decode path is CUDA-graph safe; -the branch between the two paths is on the tensor *shape*, never on its values. +The activation is quantized per call with a dynamic per-tensor scale: one Triton launch (a +single program walking the tensor) for tiny inputs, two above ``_SPLIT_MIN_ELEMENTS`` (a +block-parallel amax, then a reduce+cast). Both compute the same arithmetic, so the split is +a pure speed choice. No host sync anywhere (the scales stay on the device), so the decode +path is CUDA-graph safe; the branch between the two paths is on the tensor *shape*, never +on its values. Measured on an RTX 6000 Ada (sm_89, torch 2.11.0+cu130) at the qwen4_exp TP=2 shapes, per decode step per rank over 48 layers: bf16 cuBLAS 3.2-3.4 ms, this path 1.9 ms at M=1/8/16. @@ -26,15 +28,21 @@ FP8 = torch.float8_e4m3fn E4M3_MAX = 448.0 -_FUSED_MAX_ELEMENTS = ( - 65536 # one program handles the whole tensor below this (decode sizes) -) -_SCALE_FLOOR = 1e-12 # an all-zero activation (graph warmup buffers) must not give 1/0 +_BLOCK = 4096 +# One program per block above this, two launches (partial amax, then reduce+cast); below it +# a single program walks the whole tensor. The one-program kernel is serial over the tensor, +# so it must not be given the [T, hc_count*hidden] hyper-connection activations. +_SPLIT_MIN_ELEMENTS = 16384 +# Partial-amax programs, and so the constexpr width of the fold in pass 2. Capping it (rather +# than letting it follow the input) keeps BOTH split kernels to ONE compiled Triton variant: +# the partial count is a runtime argument, so a server that sees a new prompt length does not +# compile a new kernel mid-generation. +_MAX_PARTS = 512 @triton.jit def _quant_fused_kernel(x_ptr, out_ptr, scale_ptr, n, BLOCK: tl.constexpr): - """One program: max|x| over the tensor, then the cast. Decode-sized inputs only.""" + """One program: max|x| over the tensor, then the cast. Tiny inputs only (see SPLIT_MIN).""" acc = tl.zeros([BLOCK], dtype=tl.float32) for start in range(0, n, BLOCK): offs = start + tl.arange(0, BLOCK) @@ -52,12 +60,38 @@ def _quant_fused_kernel(x_ptr, out_ptr, scale_ptr, n, BLOCK: tl.constexpr): @triton.jit -def _quant_cast_kernel(x_ptr, out_ptr, scale_ptr, n, BLOCK: tl.constexpr): - """Cast under a scale already on the device (prefill sizes; the amax is a torch reduction).""" - offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) +def _amax_partial_kernel(x_ptr, part_ptr, n, nprog, BLOCK: tl.constexpr): + """max|x| -> one fp32 partial per program (pass 1). Grid-strided, so ``nprog`` bounds the + partial count however large the tensor is.""" + pid = tl.program_id(0) + acc = tl.zeros([BLOCK], dtype=tl.float32) + for start in range(pid * BLOCK, n, nprog * BLOCK): + offs = start + tl.arange(0, BLOCK) + v = tl.load(x_ptr + offs, mask=offs < n, other=0.0).to(tl.float32) + acc = tl.maximum(acc, tl.abs(v)) + tl.store(part_ptr + pid, tl.max(acc, axis=0)) + + +@triton.jit +def _reduce_cast_kernel( + x_ptr, out_ptr, part_ptr, scale_ptr, n, nprog, + MAX_PARTS: tl.constexpr, BLOCK: tl.constexpr, +): + """Pass 2: fold the partials to the tensor amax, then cast this program's block. + + Every program repeats the (nprog-element, L2-resident) fold instead of paying a third + launch for it; program 0 also publishes the scale. The arithmetic is the one + ``_quant_fused_kernel`` uses, so both paths quantize a given tensor identically. + """ + poffs = tl.arange(0, MAX_PARTS) + amax = tl.max(tl.load(part_ptr + poffs, mask=poffs < nprog, other=0.0), axis=0) + amax = tl.maximum(amax, 1e-12) + pid = tl.program_id(0) + if pid == 0: + tl.store(scale_ptr, amax / 448.0) + offs = pid * BLOCK + tl.arange(0, BLOCK) mask = offs < n - inv = 1.0 / tl.load(scale_ptr) - v = tl.load(x_ptr + offs, mask=mask, other=0.0).to(tl.float32) * inv + v = tl.load(x_ptr + offs, mask=mask, other=0.0).to(tl.float32) * (448.0 / amax) v = tl.minimum(tl.maximum(v, -448.0), 448.0) tl.store(out_ptr + offs, v.to(tl.float8e4nv), mask=mask) @@ -66,14 +100,17 @@ def quant_per_tensor(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """``(x_fp8, scale)`` with ``x ~= x_fp8 * scale``; ``x`` contiguous, ``scale`` fp32 ``()``.""" n = x.numel() out = torch.empty_like(x, dtype=FP8) - if n <= _FUSED_MAX_ELEMENTS: - scale = torch.empty((), dtype=torch.float32, device=x.device) - _quant_fused_kernel[(1,)](x, out, scale, n, BLOCK=4096, num_warps=8) + scale = torch.empty((), dtype=torch.float32, device=x.device) + if n <= _SPLIT_MIN_ELEMENTS: + # below this, the second launch costs more than the parallelism buys + _quant_fused_kernel[(1,)](x, out, scale, n, BLOCK=_BLOCK, num_warps=8) return out, scale - amax = torch.linalg.vector_norm(x, ord=float("inf")).float() - scale = amax.clamp_min_(_SCALE_FLOOR).div_(E4M3_MAX) - _quant_cast_kernel[(triton.cdiv(n, 4096),)]( - x, out, scale, n, BLOCK=4096, num_warps=4 + nblock = triton.cdiv(n, _BLOCK) + nprog = min(nblock, _MAX_PARTS) + part = torch.empty(nprog, dtype=torch.float32, device=x.device) + _amax_partial_kernel[(nprog,)](x, part, n, nprog, BLOCK=_BLOCK, num_warps=8) + _reduce_cast_kernel[(nblock,)]( + x, out, part, scale, n, nprog, MAX_PARTS=_MAX_PARTS, BLOCK=_BLOCK, num_warps=8, ) return out, scale diff --git a/tests/models/qwen4_exp/test_fp8_dense.py b/tests/models/qwen4_exp/test_fp8_dense.py index fac51bdb5..4b12598a1 100644 --- a/tests/models/qwen4_exp/test_fp8_dense.py +++ b/tests/models/qwen4_exp/test_fp8_dense.py @@ -121,3 +121,20 @@ def test_quant_per_tensor_zero_input_is_finite(): torch.zeros(16, 128, device="cuda", dtype=torch.bfloat16) ) assert torch.isfinite(scale).item() and (x8.float() == 0).all() + + +@requires_cuda +@pytest.mark.parametrize("n", [4096, 16384, 20480, 81920, 200000]) +def test_split_quant_path_agrees_with_the_single_program_one(n: int): + """The two-launch path above _SPLIT_MIN_ELEMENTS must quantize bit for bit like the + one-program kernel: same amax (max is exact under any grouping), same arithmetic.""" + import freetoken.layers.fp8_dynamic as m + + torch.manual_seed(0) + x = (torch.randn(n, device="cuda") * 3.0).to(torch.bfloat16) + got8, got_scale = m.quant_per_tensor(x) + ref8 = torch.empty_like(x, dtype=m.FP8) + ref_scale = torch.empty((), dtype=torch.float32, device="cuda") + m._quant_fused_kernel[(1,)](x, ref8, ref_scale, n, BLOCK=m._BLOCK, num_warps=8) + assert torch.equal(got_scale, ref_scale) + assert torch.equal(got8.view(torch.uint8), ref8.view(torch.uint8)) From 85c4b6042a55211d6983a839bcbc2dc2b9d457b0 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Sat, 5 Sep 2026 17:10:21 -0400 Subject: [PATCH 09/14] feat(layers): optional per-tensor FP8 lm_head (FREETOKEN_FP8_LMHEAD=1) The vocab-parallel head is the last large bf16 read on the decode path: 0.64 GB per step per rank at the Qwen3.8-Flash-Next geometry ([124160, 2560] after the TP=2 vocab split), 0.80 ms of the 9.9 ms step in an nsys trace of production. Its own flag, not FREETOKEN_FP8_DENSE, because this one moves the logits: every other quantized module feeds a norm or a sigmoid downstream, while a per-tensor e4m3 vocab matrix changes each sampled token's score directly, so it carries its own quality gate rather than riding along with the pure-throughput changes. ParallelLMHead.forward grows a _logits() seam (the local vocab-shard GEMM); Fp8ParallelLMHead overrides only that, leaving the all_gather of the logits above it untouched. Untied embeddings only -- a tied head shares the bf16 embedding table, which the lookup side still reads as bf16. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt (cherry picked from commit 35e80b30e30052fbac360baa167f0f12a791ed34) (cherry picked from commit 0693735a60efd17fded6f9b686f3b6d298c5bea0) (cherry picked from commit 1797ab56e8ee9cd7138111770896bca54b7d9ce7) --- python/freetoken/layers/embedding.py | 16 +++++++++---- python/freetoken/layers/fp8_dynamic.py | 20 ++++++++++++++++ python/freetoken/models/config.py | 8 +++++++ python/freetoken/models/qwen4_exp/model.py | 26 ++++++++++++++------- python/freetoken/models/qwen4_exp/weight.py | 24 ++++++++++++------- tests/models/qwen4_exp/test_fp8_dense.py | 13 +++++++++++ 6 files changed, 87 insertions(+), 20 deletions(-) diff --git a/python/freetoken/layers/embedding.py b/python/freetoken/layers/embedding.py index a22216548..6939156cd 100644 --- a/python/freetoken/layers/embedding.py +++ b/python/freetoken/layers/embedding.py @@ -121,6 +121,17 @@ def state_dict( return super().state_dict(prefix=prefix, result=result) return {} if result is None else result + def _logits(self, x: torch.Tensor) -> torch.Tensor: + """The local vocab-shard GEMM; the seam a quantized head overrides.""" + module = self.tied_embedding or self + return F.linear(x, module.weight, self.bias) + + def _logits(self, x: torch.Tensor) -> torch.Tensor: + """The local vocab-shard GEMM; the seam a quantized head overrides.""" + if self.tied_embedding is not None: + return F.linear(x, self.tied_embedding.weight, self.bias) + return self.quant_method.apply(self, x) + @nvtx_annotate("LMHead") def forward(self, x: torch.Tensor) -> torch.Tensor: ctx = get_global_ctx() @@ -131,10 +142,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: x = x[indices].contiguous() del indices - if self.tied_embedding is not None: - logits = F.linear(x, self.tied_embedding.weight, self.bias) - else: - logits = self.quant_method.apply(self, x) + logits = self._logits(x) if self.tp_size == 1: return logits input_shape = logits.shape diff --git a/python/freetoken/layers/fp8_dynamic.py b/python/freetoken/layers/fp8_dynamic.py index 5e84a929c..caeebb322 100644 --- a/python/freetoken/layers/fp8_dynamic.py +++ b/python/freetoken/layers/fp8_dynamic.py @@ -25,6 +25,7 @@ from freetoken.utils import div_even from .base import BaseOP +from .embedding import ParallelLMHead FP8 = torch.float8_e4m3fn E4M3_MAX = 448.0 @@ -180,9 +181,28 @@ def __init__(self, input_size: int, output_size: int): ) +class Fp8ParallelLMHead(ParallelLMHead): + """``ParallelLMHead`` whose vocab shard is a per-tensor e4m3 weight (FREETOKEN_FP8_LMHEAD=1). + + Only the GEMM changes; the vocab-parallel all_gather of the logits above it is untouched. + Untied embeddings only -- a tied head shares the bf16 embedding table, which the lookup + side still reads as bf16. + """ + + def __init__(self, num_embeddings: int, embedding_dim: int): + super().__init__(num_embeddings, embedding_dim) + self.weight = torch.empty(self.num_embeddings_tp, embedding_dim, dtype=FP8) + self.weight_scale = torch.empty((), dtype=torch.float32) + + def _logits(self, x: torch.Tensor) -> torch.Tensor: + y = fp8_dynamic_linear(x, self.weight, self.weight_scale) + return y if self.bias is None else y + self.bias + + __all__ = [ "FP8", "E4M3_MAX", + "Fp8ParallelLMHead", "Fp8DynamicColMerged", "Fp8DynamicLinear", "Fp8DynamicRowParallel", diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py index 0ad73b752..62340ab75 100644 --- a/python/freetoken/models/config.py +++ b/python/freetoken/models/config.py @@ -25,6 +25,14 @@ def fp8_dense_enabled() -> bool: return os.getenv("FREETOKEN_FP8_DENSE", "0").strip().lower() in _VISION_TRUE +def fp8_lmhead_enabled() -> bool: + """Load-time FP8 for the lm_head as well (opt-in, default OFF, needs FREETOKEN_FP8_DENSE=1 + for the rest). Separate from :func:`fp8_dense_enabled` because this one moves the logits: + a per-tensor e4m3 vocab matrix changes every sampled token's score, so it carries its own + quality gate. ``FREETOKEN_FP8_LMHEAD=1``.""" + return os.getenv("FREETOKEN_FP8_LMHEAD", "0").strip().lower() in _VISION_TRUE + + def detect_expert_quant(hf_config: Any) -> str: """Routed-expert quantization from a checkpoint's ``quantization_config``: ``"nvfp4"`` for a ModelOpt FP4 build (``quant_algo: NVFP4``) OR an llm-compressor NVFP4 export diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py index c9cf9c131..848b4d645 100644 --- a/python/freetoken/models/qwen4_exp/model.py +++ b/python/freetoken/models/qwen4_exp/model.py @@ -21,6 +21,7 @@ from freetoken.core import get_global_ctx from freetoken.layers import BaseOP, OPList, ParallelLMHead, VocabParallelEmbedding from freetoken.models.blocks import BaseLLMModel +from freetoken.models.config import fp8_lmhead_enabled from freetoken.utils import nvtx_annotate from .attention import Qwen4ExpAttention @@ -129,14 +130,23 @@ class Qwen4ExpForCausalLM(BaseLLMModel): def __init__(self, config: ModelConfig) -> None: self._config = config self.model = Qwen4ExpModel(config) - self.lm_head = ParallelLMHead( - num_embeddings=config.vocab_size, - embedding_dim=config.hidden_size, - tie_word_embeddings=config.tie_word_embeddings, - tied_embedding=self.model.embed_tokens if config.tie_word_embeddings else None, - quant_config=config.quant, - prefix="lm_head", - ) + if fp8_lmhead_enabled() and not config.tie_word_embeddings and config.quant is None: + # Load-time per-tensor FP8 head, for a checkpoint that ships lm_head unquantized. + # A checkpoint that declares its own scheme goes through QuantConfig instead. + from freetoken.layers.fp8_dynamic import Fp8ParallelLMHead + + self.lm_head = Fp8ParallelLMHead( + num_embeddings=config.vocab_size, embedding_dim=config.hidden_size + ) + else: + self.lm_head = ParallelLMHead( + num_embeddings=config.vocab_size, + embedding_dim=config.hidden_size, + tie_word_embeddings=config.tie_word_embeddings, + tied_embedding=self.model.embed_tokens if config.tie_word_embeddings else None, + quant_config=config.quant, + prefix="lm_head", + ) super().__init__() def load_host_tables(self, engine_config) -> int: diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index 57f9d5f4f..125cea2c0 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -200,6 +200,8 @@ def _shard(name: str, t: torch.Tensor, config, rank: int, world: int) -> torch.T ".self_attn.o_proj.weight", ".linear_attn.out_proj.weight", ) +# Behind its own flag: this one moves the logits (see models.config.fp8_lmhead_enabled). +_FP8_LMHEAD_SUFFIX = "lm_head.weight" _E4M3_MAX = 448.0 @@ -211,10 +213,11 @@ def _quantize_per_tensor(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: def _fp8_dense( - name: str, t: torch.Tensor, config, world: int + name: str, t: torch.Tensor, config, world: int, *, + dense: bool = True, lm_head: bool = False, ) -> Iterator[tuple[str, torch.Tensor]]: """The (already sharded) dense tensor as the fp8_dynamic model expects it.""" - if name.endswith(".linear_attn.in_proj.weight"): + if dense and name.endswith(".linear_attn.in_proj.weight"): g = config.linear_attention_group() nk = div_even(g.num_key_heads, world, allow_replicate=True) nv = div_even(g.num_value_heads, world, allow_replicate=True) @@ -227,7 +230,9 @@ def _fp8_dense( # clone, not contiguous(): a contiguous row slice IS contiguous, so .contiguous() would # hand back a view that keeps the whole bf16 in_proj (36 x 42 MB per rank) resident yield base + "in_proj_ba.weight", t[qkvz:].clone() - elif name.endswith(_FP8_DENSE_SUFFIXES): + elif (dense and name.endswith(_FP8_DENSE_SUFFIXES)) or ( + lm_head and name.endswith(_FP8_LMHEAD_SUFFIX) + ): w8, scale = _quantize_per_tensor(t) yield name, w8 yield name[: -len("weight")] + "weight_scale", scale @@ -259,16 +264,17 @@ def iter_weights( if not include_non_moe: return - from freetoken.models.config import fp8_dense_enabled + from freetoken.models.config import fp8_dense_enabled, fp8_lmhead_enabled from .config import parse_config tp = get_tp_info() # The sharding geometry (and the fp8 split) need the HF config; TP=1 bf16 never does, so # keep that path free of a config load (synthetic test checkpoints carry no model_type). + lmhead = fp8_lmhead_enabled() config = ( parse_config(cached_load_hf_config(model_path)) - if tp.size > 1 or fp8_dense_enabled() + if tp.size > 1 or fp8_dense_enabled() or lmhead else None ) fp8 = config is not None and config.attn_quant == "fp8_dynamic" @@ -277,11 +283,13 @@ def iter_weights( "qwen4_exp dense projections: load-time per-tensor FP8 (W8A8 via _scaled_mm), " "FREETOKEN_FP8_DENSE=1" ) + if lmhead: + logger.info("qwen4_exp lm_head: load-time per-tensor FP8, FREETOKEN_FP8_LMHEAD=1") def emit(name: str, tensor: torch.Tensor): tensor = _shard(name, tensor, config, tp.rank, tp.size) - if fp8: - yield from _fp8_dense(name, tensor, config, tp.size) + if fp8 or lmhead: + yield from _fp8_dense(name, tensor, config, tp.size, dense=fp8, lm_head=lmhead) else: yield name, tensor @@ -306,7 +314,7 @@ def emit(name: str, tensor: torch.Tensor): yield from emit(name, tensor) assert not fuse_buf, f"Incomplete projection fusions: {sorted(fuse_buf)}" - if fp8 and device.type == "cuda": + if (fp8 or lmhead) and device.type == "cuda": # The bf16 originals and fp32 temporaries of the quantization sit in the caching # allocator; hand them back so the expert-cache planner (free VRAM after load) sees # the halved dense footprint instead of the slack. diff --git a/tests/models/qwen4_exp/test_fp8_dense.py b/tests/models/qwen4_exp/test_fp8_dense.py index 4b12598a1..058515b18 100644 --- a/tests/models/qwen4_exp/test_fp8_dense.py +++ b/tests/models/qwen4_exp/test_fp8_dense.py @@ -138,3 +138,16 @@ def test_split_quant_path_agrees_with_the_single_program_one(n: int): m._quant_fused_kernel[(1,)](x, ref8, ref_scale, n, BLOCK=m._BLOCK, num_warps=8) assert torch.equal(got_scale, ref_scale) assert torch.equal(got8.view(torch.uint8), ref8.view(torch.uint8)) + + +def test_reader_gates_lm_head_behind_its_own_flag(): + cfg = _cfg() + t = torch.randn(64, 32, dtype=torch.bfloat16) + assert list(dict(_fp8_dense("lm_head.weight", t, cfg, 1))) == ["lm_head.weight"] + out = dict(_fp8_dense("lm_head.weight", t, cfg, 1, dense=False, lm_head=True)) + assert sorted(out) == ["lm_head.weight", "lm_head.weight_scale"] + assert out["lm_head.weight"].dtype == torch.float8_e4m3fn + # the lm_head flag alone must not pull in the dense rewrites + passthrough = dict(_fp8_dense("x.self_attn.qkv_proj.weight", t, cfg, 1, + dense=False, lm_head=True)) + assert list(passthrough) == ["x.self_attn.qkv_proj.weight"] From 62a9a8c1031fa779c674d18b03a86d37fdeeed5f Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 10 Sep 2026 18:02:07 -0400 Subject: [PATCH 10/14] fix(qwen4_exp): forward prefix through Qwen4ExpMoE, and test the CPU build Qwen4ExpDecoderLayer builds its MoE as `Qwen4ExpMoE(config, layer_id, prefix=...)`, but this PR's override of __init__ (added to hold the TP communicator) took only (config, layer_id), so a server boot died with TypeError: Qwen4ExpMoE.__init__() got an unexpected keyword argument 'prefix' The whole CPU test suite was green with that bug in place, because every test that builds a decoder layer is behind requires_cuda -- nothing without a GPU ever constructed the model. tests/models/qwen4_exp/test_build_cpu.py closes that: it builds the full model on the meta device (no GPU, no memory) and asserts the state dict has both layer families, an lm_head, and MoE weights on more than one layer, so a dropped or shared prefix fails too. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt --- python/freetoken/models/qwen4_exp/moe.py | 6 ++-- tests/models/qwen4_exp/test_build_cpu.py | 43 ++++++++++++++++++++++++ 2 files changed, 47 insertions(+), 2 deletions(-) create mode 100644 tests/models/qwen4_exp/test_build_cpu.py diff --git a/python/freetoken/models/qwen4_exp/moe.py b/python/freetoken/models/qwen4_exp/moe.py index 7d9f6adf7..c10407861 100644 --- a/python/freetoken/models/qwen4_exp/moe.py +++ b/python/freetoken/models/qwen4_exp/moe.py @@ -25,8 +25,10 @@ class Qwen4ExpMoE(Qwen3_5MoE): is reduced once (one all-reduce per MoE layer instead of two). """ - def __init__(self, config: ModelConfig, layer_id: int | None = None) -> None: - super().__init__(config, layer_id=layer_id) + def __init__( + self, config: ModelConfig, layer_id: int | None = None, *, prefix: str = "" + ) -> None: + super().__init__(config, layer_id=layer_id, prefix=prefix) self._comm = DistributedCommunicator() self._tp_size = get_tp_info().size diff --git a/tests/models/qwen4_exp/test_build_cpu.py b/tests/models/qwen4_exp/test_build_cpu.py new file mode 100644 index 000000000..0cf2027d1 --- /dev/null +++ b/tests/models/qwen4_exp/test_build_cpu.py @@ -0,0 +1,43 @@ +"""The whole model must CONSTRUCT without a GPU. + +Every other model test that builds a decoder layer is behind requires_cuda, so a constructor +signature mismatch between Qwen4ExpDecoderLayer and the ops it builds is invisible to a +CPU-only run -- it only shows up when a server boots. That happened: Qwen4ExpMoE did not +accept the `prefix` kwarg model.py passes, and a full CPU suite still went green. + +Building on the meta device costs no memory and no GPU, so there is no reason not to. +""" + +import torch +from freetoken.distributed import set_tp_info, try_get_tp_info +from freetoken.layers import set_rope_device +from freetoken.models.qwen4_exp.config import parse_config +from freetoken.models.qwen4_exp.model import Qwen4ExpForCausalLM + +from .common import toy_hf_config + + +def _build(): + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + set_rope_device(torch.device("cpu")) + with torch.device("meta"): + return Qwen4ExpForCausalLM(parse_config(toy_hf_config())) + + +def test_model_constructs_on_meta_without_cuda(): + model = _build() + sd = model.state_dict() + assert sd, "state dict is empty" + # both layer families and the head must be present + assert any(".self_attn." in k for k in sd), "no full-attention layer built" + assert any(".linear_attn." in k for k in sd), "no GDN layer built" + assert any(k.startswith("lm_head") for k in sd), "no lm_head built" + + +def test_every_layer_gets_its_own_prefixed_weights(): + # a prefix that is dropped or shared silently collapses layers onto one another + sd = _build().state_dict() + mlp_keys = {k for k in sd if ".mlp." in k} + layers = {k.split(".layers.")[1].split(".")[0] for k in mlp_keys if ".layers." in k} + assert len(layers) > 1, f"MoE weights landed on a single layer: {sorted(layers)}" From d2ad39b8d55550938741f7788183557fd2524591 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 10 Sep 2026 18:40:25 -0400 Subject: [PATCH 11/14] feat(quant): let the triton NVFP4 expert kernel serve a tensor-parallel shard With the expert piece stream sliced per rank and the banks sized from MoEConfig.local_intermediate, a rank holds exactly its half of every expert, so this kernel can serve TP>1 -- the routed output is a partial sum and the MoE layer already reduces it (_maybe_all_reduce, or the single combined all-reduce in qwen4_exp's block). Without this the whole selection table is empty under TP=2 on sm_89 and the server refuses to start: KernelSelectionError: no usable kernel in table; triton: TP > 1 is not supported for this expert format; marlin: vLLM is not installed; b12x: b12x requires sm_120+, got sm_89 marlin and b12x keep tp_ok=False deliberately: their pack() repacks the native rows and neither has been verified against a per-rank bank. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt --- python/freetoken/layers/quantization/moe/nvfp4.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/python/freetoken/layers/quantization/moe/nvfp4.py b/python/freetoken/layers/quantization/moe/nvfp4.py index dfe35604b..e2dc1478c 100644 --- a/python/freetoken/layers/quantization/moe/nvfp4.py +++ b/python/freetoken/layers/quantization/moe/nvfp4.py @@ -41,7 +41,13 @@ 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) From 1da3c6c7353b9ee304a8c797bee42b8cd54629fa Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 10 Sep 2026 19:00:04 -0400 Subject: [PATCH 12/14] fix(qwen4_exp): one predicate decides the FP8 lm_head, for builder and loader alike The model builder and the weight reader tested this separately and drifted, so on the shipping checkpoint the reader emitted lm_head.weight_scale while the builder made a plain ParallelLMHead, and startup died with RuntimeError: Unexpected keys in state_dict: ['lm_head.weight_scale'] The builder's extra condition was `config.quant is None`, which is wrong twice over: what matters is whether the checkpoint declares a scheme for **lm_head**, not whether it has a QuantConfig at all. This model ships NVFP4 routed experts (so quant is not None) with lm_head in the modelopt ignore list (so it has no scheme and the synthetic FP8 head is exactly what is wanted). Both sides now call config.use_fp8_lmhead(), which owns the flag, the tied-embedding exclusion and the scheme test in one place. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt --- python/freetoken/models/qwen4_exp/config.py | 17 +++++++++++++++++ python/freetoken/models/qwen4_exp/model.py | 4 ++-- python/freetoken/models/qwen4_exp/weight.py | 7 ++++--- 3 files changed, 23 insertions(+), 5 deletions(-) diff --git a/python/freetoken/models/qwen4_exp/config.py b/python/freetoken/models/qwen4_exp/config.py index e8b62766a..6f1939eb3 100644 --- a/python/freetoken/models/qwen4_exp/config.py +++ b/python/freetoken/models/qwen4_exp/config.py @@ -110,6 +110,23 @@ def _layer_types(text: Any) -> list[str]: ] +def use_fp8_lmhead(config: ModelConfig) -> bool: + """Whether to build/load the lm_head as load-time per-tensor FP8. + + The MODEL BUILDER and the WEIGHT READER must agree exactly: if one says yes and the + other no, the state dict gains or loses ``lm_head.weight_scale`` and load fails. They + used to test this separately and drifted. Conditions: the operator asked + (FREETOKEN_FP8_LMHEAD=1), the head is untied (a tied head shares the bf16 embedding + table), and the checkpoint declares no scheme OF ITS OWN for lm_head -- a quantized head + goes through QuantConfig instead. Note it is the lm_head scheme that matters, not whether + the checkpoint has a QuantConfig at all: this model ships NVFP4 experts with lm_head in + the modelopt ignore list.""" + if not fp8_lmhead_enabled() or config.tie_word_embeddings: + return False + quant = getattr(config, "quant", None) + return quant is None or quant.scheme_for("lm_head") is None + + def parse_config(hf_config: Any) -> ModelConfig: text = getattr(hf_config, "text_config", hf_config) diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py index 848b4d645..3238459a6 100644 --- a/python/freetoken/models/qwen4_exp/model.py +++ b/python/freetoken/models/qwen4_exp/model.py @@ -21,7 +21,7 @@ from freetoken.core import get_global_ctx from freetoken.layers import BaseOP, OPList, ParallelLMHead, VocabParallelEmbedding from freetoken.models.blocks import BaseLLMModel -from freetoken.models.config import fp8_lmhead_enabled +from .config import use_fp8_lmhead from freetoken.utils import nvtx_annotate from .attention import Qwen4ExpAttention @@ -130,7 +130,7 @@ class Qwen4ExpForCausalLM(BaseLLMModel): def __init__(self, config: ModelConfig) -> None: self._config = config self.model = Qwen4ExpModel(config) - if fp8_lmhead_enabled() and not config.tie_word_embeddings and config.quant is None: + if use_fp8_lmhead(config): # Load-time per-tensor FP8 head, for a checkpoint that ships lm_head unquantized. # A checkpoint that declares its own scheme goes through QuantConfig instead. from freetoken.layers.fp8_dynamic import Fp8ParallelLMHead diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index 125cea2c0..10f960852 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -266,17 +266,18 @@ def iter_weights( from freetoken.models.config import fp8_dense_enabled, fp8_lmhead_enabled - from .config import parse_config + from .config import parse_config, use_fp8_lmhead tp = get_tp_info() # The sharding geometry (and the fp8 split) need the HF config; TP=1 bf16 never does, so # keep that path free of a config load (synthetic test checkpoints carry no model_type). - lmhead = fp8_lmhead_enabled() config = ( parse_config(cached_load_hf_config(model_path)) - if tp.size > 1 or fp8_dense_enabled() or lmhead + if tp.size > 1 or fp8_dense_enabled() or fp8_lmhead_enabled() else None ) + # the SAME predicate the model builder uses, so the emitted keys always match its slots + lmhead = config is not None and use_fp8_lmhead(config) fp8 = config is not None and config.attn_quant == "fp8_dynamic" if fp8: logger.info( From d377a20c09fe90b148c03f66d72aad22e4d81e15 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 10 Sep 2026 19:12:27 -0400 Subject: [PATCH 13/14] fix(qwen4_exp): import fp8_lmhead_enabled, which use_fp8_lmhead calls The predicate moved into config.py but its import did not come with it, so building the model raised NameError on the very first line of use_fp8_lmhead -- under every flag setting, including the default one that wants no FP8 head at all. Found by the CPU key check (build the model on meta, diff its slots against the names the loader emits) before it could cost a GPU window. No behaviour change beyond the module now importing. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt --- python/freetoken/models/qwen4_exp/config.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/freetoken/models/qwen4_exp/config.py b/python/freetoken/models/qwen4_exp/config.py index 6f1939eb3..bb9966987 100644 --- a/python/freetoken/models/qwen4_exp/config.py +++ b/python/freetoken/models/qwen4_exp/config.py @@ -8,6 +8,7 @@ from freetoken.layers.quantization import QuantConfig from freetoken.models.config import ( fp8_dense_enabled, + fp8_lmhead_enabled, FullAttentionGroupConfig, LinearGatedDeltaGroupConfig, ModelConfig, From bd106c9aa5963f3d19c543dc557d7649f2fb61a6 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 10 Sep 2026 19:53:05 -0400 Subject: [PATCH 14/14] fix(qwen4_exp): restore the local head counts the GDN forward reads The rebase onto the quantization refactor dropped nk, nv = self._local_num_k_heads, self._local_num_v_heads from the top of GatedDeltaNet.forward while keeping all eleven uses of nk and nv below it. Every decode died at the first CUDA-graph capture: File "python/freetoken/models/qwen4_exp/gdn.py", line 189, in forward b, a = torch.split(ba, [nv, nv], dim=-1) NameError: name 'nv' is not defined Also drops a duplicate `from freetoken.distributed import get_tp_info` two lines under the first, left by the same merge. `ruff check --select F821` reports both, and would have reported the missing fp8_lmhead_enabled import a commit earlier. Lint the tree before asking for the GPU, not after: this file's eleven undefined names cost a window that a sub-second check would have saved. The eight F821 hits outside qwen4_exp are pre-existing on main. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt --- python/freetoken/models/qwen4_exp/gdn.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/freetoken/models/qwen4_exp/gdn.py b/python/freetoken/models/qwen4_exp/gdn.py index f0d95c3d7..763cb78ca 100644 --- a/python/freetoken/models/qwen4_exp/gdn.py +++ b/python/freetoken/models/qwen4_exp/gdn.py @@ -5,7 +5,6 @@ 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.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 @@ -182,6 +181,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: fla = build_fla_metadata(batch, hidden_states.device) batch.fla_metadata = fla + nk, nv = self._local_num_k_heads, self._local_num_v_heads if self._split_in_proj or self._dynamic_fp8: qkvz = self.in_proj_qkvz.forward(hidden_states) conv_in, z = torch.split(qkvz, self._in_proj_split[:2], dim=-1)