From 9e73f2cf126843de238aff50538d9bf575805262 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 4 Sep 2026 14:19:37 -0400 Subject: [PATCH 1/6] 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 +- python/freetoken/models/nvfp4_banks.py | 123 ++++++++---------- .../freetoken/models/qwen4_exp/attention.py | 39 ++++-- python/freetoken/models/qwen4_exp/gdn.py | 66 +++++++--- python/freetoken/models/qwen4_exp/moe.py | 29 ++++- 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 +++++++++++ 9 files changed, 388 insertions(+), 116 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 bcbe6bcf2..e190cac18 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -74,7 +74,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 f707e0429..15c973a98 100644 --- a/python/freetoken/layers/linear.py +++ b/python/freetoken/layers/linear.py @@ -59,10 +59,14 @@ def __init__( input_size: int, output_sizes: List[int], has_bias: bool, + local_output_sizes: List[int] | None = None, ): - # check that all output sizes are divisible by tp_size + # 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__(input_size, output_size, input_size, tp_output_size, has_bias) diff --git a/python/freetoken/models/nvfp4_banks.py b/python/freetoken/models/nvfp4_banks.py index 6b933ff1d..6980c9ed8 100644 --- a/python/freetoken/models/nvfp4_banks.py +++ b/python/freetoken/models/nvfp4_banks.py @@ -9,7 +9,8 @@ import safetensors import torch -from freetoken.utils import download_hf_weight +from freetoken.distributed import get_tp_info +from freetoken.utils import div_even, download_hf_weight from tqdm import tqdm LayerToBank = Callable[[int, object], int | None] @@ -78,6 +79,41 @@ def _alloc_nvfp4_host_banks(num_layers: int, E: int, H: int, I: int): }, num_layers) +def _tp_slice(inter: int) -> tuple[int, int]: + """``(i_local, i_lo)``: this rank's slice of the intermediate axis. TP shards every expert + along I (gate/up rows, down columns), the ``stream_moe_expert_sources`` convention, so the + routed output is a partial sum the MoE layer all-reduces.""" + tp = get_tp_info() + i_local = div_even(inter, tp.size) + assert i_local % 16 == 0, f"NVFP4 TP shard {i_local} must cover whole 16-wide scale blocks" + return i_local, tp.rank * i_local + + +class _Placer: + """Writes one checkpoint expert tensor into its bank slot (this rank's I slice only).""" + + def __init__(self, banks: dict, inter: int): + self.b = banks + self.i_local, self.i_lo = _tp_slice(inter) + + def put(self, layer: int, expert: int, role: str, kind: str, tensor, global_scale=None): + n, lo = self.i_local, self.i_lo + b = self.b + if role == "down": # [H, I/2] codes, [H, I/16] scales, [H] global + if kind == "weight": + b["down_packed"][layer][expert] = tensor[:, lo // 2 : (lo + n) // 2] + else: + b["down_scale"][layer][expert] = tensor[:, lo // 16 : (lo + n) // 16] + b["down_global"][layer][expert] = global_scale + return + rows = slice(0, n) if role == "gate" else slice(n, 2 * n) # gate | up on the row axis + if kind == "weight": + b["gate_up_packed"][layer][expert, rows] = tensor[lo : lo + n] + else: + b["gate_up_scale"][layer][expert, rows] = tensor[lo : lo + n] + b["gate_up_global"][layer][expert, rows] = global_scale # per-tensor scalar + + def load_nvfp4_expert_source_banks( model_path: str, config, @@ -149,13 +185,9 @@ def load_nvfp4_expert_source_banks( globals_map[key] = _ingest_global(spec, f.get_tensor(name)) drop_page_cache(path) - _hb = _alloc_nvfp4_host_banks(num_layers, E, H, I) # unpinned; pinned after fill - gate_up_packed = [b.tensor for b in _hb["gate_up_packed"]] - gate_up_scale = [b.tensor for b in _hb["gate_up_scale"]] - gate_up_global = [b.tensor for b in _hb["gate_up_global"]] - down_packed = [b.tensor for b in _hb["down_packed"]] - down_scale = [b.tensor for b in _hb["down_scale"]] - down_global = [b.tensor for b in _hb["down_global"]] + _hb = _alloc_nvfp4_host_banks(num_layers, E, H, _tp_slice(I)[0]) # unpinned; pinned after fill + banks = {name: [b.tensor for b in layers] for name, layers in _hb.items()} + place = _Placer(banks, I) from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline @@ -170,30 +202,11 @@ def _load(sink) -> int: expert = int(match.group("expert")) proj = match.group("proj") role = spec.proj_to_role[proj] + if role not in ("gate", "up", "down"): + raise ValueError(f"{spec.desc}: unknown projection role {role!r}") kind = _canon_kind(spec, match.group("kind")) - tensor = f.get_tensor(name) - if kind == "weight": - if role == "gate": - gate_up_packed[bank_layer_id][expert, :I] = tensor - elif role == "up": - gate_up_packed[bank_layer_id][expert, I:] = tensor - elif role == "down": - down_packed[bank_layer_id][expert] = tensor - else: - raise ValueError(f"{spec.desc}: unknown projection role {role!r}") - else: - global_scale = globals_map[(layer, expert, proj)] - if role == "gate": - gate_up_scale[bank_layer_id][expert, :I] = tensor - gate_up_global[bank_layer_id][expert, :I] = global_scale - elif role == "up": - gate_up_scale[bank_layer_id][expert, I:] = tensor - gate_up_global[bank_layer_id][expert, I:] = global_scale - elif role == "down": - down_scale[bank_layer_id][expert] = tensor - down_global[bank_layer_id][expert] = global_scale - else: - raise ValueError(f"{spec.desc}: unknown projection role {role!r}") + place.put(bank_layer_id, expert, role, kind, f.get_tensor(name), + None if kind == "weight" else globals_map[(layer, expert, proj)]) tracker.note(bank_layer_id) placed += 1 drop_page_cache(path) @@ -207,14 +220,7 @@ def _load(sink) -> int: expected = num_layers * E * 6 assert placed == expected, f"{spec.desc}: loaded {placed} expert tensors, expected {expected}" - return { - "gate_up_packed": gate_up_packed, - "gate_up_scale": gate_up_scale, - "gate_up_global": gate_up_global, - "down_packed": down_packed, - "down_scale": down_scale, - "down_global": down_global, - } + return banks def load_nvfp4_expert_source_banks_parallel( @@ -273,13 +279,9 @@ def load_nvfp4_expert_source_banks_parallel( ) drop_page_cache(path) - _hb = _alloc_nvfp4_host_banks(num_layers, E, H, I) # unpinned; pinned after fill - gate_up_packed = [b.tensor for b in _hb["gate_up_packed"]] - gate_up_scale = [b.tensor for b in _hb["gate_up_scale"]] - gate_up_global = [b.tensor for b in _hb["gate_up_global"]] - down_packed = [b.tensor for b in _hb["down_packed"]] - down_scale = [b.tensor for b in _hb["down_scale"]] - down_global = [b.tensor for b in _hb["down_global"]] + _hb = _alloc_nvfp4_host_banks(num_layers, E, H, _tp_slice(I)[0]) # unpinned; pinned after fill + banks = {name: [b.tensor for b in layers] for name, layers in _hb.items()} + place = _Placer(banks, I) from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline @@ -296,24 +298,8 @@ def _load(sink) -> int: proj = match.group("proj") role = spec.proj_to_role[proj] kind = _canon_kind(spec, match.group("kind")) - if kind == "weight": - if role == "gate": - gate_up_packed[bank_layer_id][expert, :I] = tensor - elif role == "up": - gate_up_packed[bank_layer_id][expert, I:] = tensor - else: - down_packed[bank_layer_id][expert] = tensor - else: - g = globals_map[(layer, expert, proj)] - if role == "gate": - gate_up_scale[bank_layer_id][expert, :I] = tensor - gate_up_global[bank_layer_id][expert, :I] = g - elif role == "up": - gate_up_scale[bank_layer_id][expert, I:] = tensor - gate_up_global[bank_layer_id][expert, I:] = g - else: - down_scale[bank_layer_id][expert] = tensor - down_global[bank_layer_id][expert] = g + place.put(bank_layer_id, expert, role, kind, tensor, + None if kind == "weight" else globals_map[(layer, expert, proj)]) tracker.note(bank_layer_id) placed += 1 return placed @@ -326,14 +312,7 @@ def _load(sink) -> int: expected = num_layers * E * 6 assert placed == expected, f"{spec.desc}: loaded {placed} expert tensors, expected {expected}" - return { - "gate_up_packed": gate_up_packed, - "gate_up_scale": gate_up_scale, - "gate_up_global": gate_up_global, - "down_packed": down_packed, - "down_scale": down_scale, - "down_global": down_global, - } + return banks __all__ = [ diff --git a/python/freetoken/models/qwen4_exp/attention.py b/python/freetoken/models/qwen4_exp/attention.py index d6ab2867a..db77d2b6e 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 @@ -120,11 +127,21 @@ def __init__(self, config: ModelConfig, layer_id: int) -> None: 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, ) - self.o_proj = LinearReplicated(self.qo_attn_dim, config.hidden_size, has_bias=False) + self.o_proj = LinearOProj(self.qo_attn_dim, config.hidden_size, has_bias=False) 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 @@ -140,21 +157,21 @@ def __init__(self, config: ModelConfig, layer_id: int) -> None: @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 69838153f..fc562e76c 100644 --- a/python/freetoken/models/qwen4_exp/gdn.py +++ b/python/freetoken/models/qwen4_exp/gdn.py @@ -3,8 +3,10 @@ import torch import torch.nn.functional as F from freetoken.core import get_global_ctx +from freetoken.distributed import get_tp_info from freetoken.kernel.causal_conv1d import causal_conv1d_decode, causal_conv1d_varlen -from freetoken.layers import BaseOP, LinearColParallelMerged +from freetoken.layers import BaseOP, LinearColParallelMerged, LinearOProj +from freetoken.utils import div_even from freetoken.kernel.triton.fp8_block_linear import Fp8BlockColMerged from freetoken.kernel.triton.fp8_pertensor_linear import Fp8PerTensorColMerged @@ -80,6 +82,15 @@ 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 # qkv|z carry a weight scale (block-fp8 weight_scale_inv, or per-tensor FP8 # weight_scale); b|a stay bf16. Both quant modes therefore split the four-way # fusion into an fp8 qkvz GEMM + a bf16 ba GEMM (matches sglang/vLLM). @@ -87,7 +98,14 @@ def __init__( self._pertensor_fp8 = attn_quant == "fp8_pertensor" self._fp8 = self._block_fp8 or self._pertensor_fp8 - self._in_proj_split = [self.conv_dim, self.value_dim, num_v_heads, num_v_heads] + self._in_proj_split = [ + self._local_conv_dim, self._local_value_dim, + self._local_num_v_heads, self._local_num_v_heads, + ] + if tp.size > 1: + assert not self._fp8 and attn_quant == "none", ( + "qwen4_exp TP shards bf16 GDN projections only" + ) if self._fp8: ColMerged = Fp8BlockColMerged if self._block_fp8 else Fp8PerTensorColMerged self.in_proj_qkvz = ColMerged( @@ -98,21 +116,27 @@ def __init__( ) else: # Fused input projection (one GEMM instead of four): qkv | z | b | a. - self.in_proj = LinearColParallelMerged(hidden_size, self._in_proj_split, has_bias=False) - self.conv1d = _DepthwiseConv1d(self.conv_dim, conv_kernel_size) + self.in_proj = LinearColParallelMerged( + hidden_size, [self.conv_dim, self.value_dim, num_v_heads, num_v_heads], + has_bias=False, local_output_sizes=self._in_proj_split, + ) + 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) # out_proj follows the checkpoint quant: block-fp8 / per-tensor-fp8 / compressed-tensors # NVFP4 (W4A16) / bf16. in_proj_* stay bf16 in every mode (above), so a compressed-tensors # NVFP4 checkpoint (attn_quant=="nvfp4") only makes out_proj native FP4. - self.out_proj = make_replicated_quant( - expert_quant, attn_quant, self.value_dim, hidden_size, has_bias=False - ) + if tp.size > 1: # row-parallel over the local v heads, all-reduce inside + self.out_proj = LinearOProj(self.value_dim, hidden_size, has_bias=False) + else: + self.out_proj = make_replicated_quant( + expert_quant, attn_quant, self.value_dim, hidden_size, has_bias=False + ) def _gate_params(self, a: torch.Tensor, b: torch.Tensor): beta = b.sigmoid() @@ -178,7 +202,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: @@ -187,10 +213,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, @@ -200,13 +226,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 9ef65c0a8..e743a3af0 100644 --- a/python/freetoken/models/qwen4_exp/moe.py +++ b/python/freetoken/models/qwen4_exp/moe.py @@ -4,8 +4,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.moe import make_moe_layer +from freetoken.layers import LinearRowParallel, silu_and_mul +from freetoken.layers.moe import OffloadMoELayer, make_moe_layer from freetoken.models.qwen3_5_moe.moe import Qwen3_5MoE if TYPE_CHECKING: @@ -16,9 +20,15 @@ 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: + self._comm = DistributedCommunicator() + self._tp_size = get_tp_info().size if getattr(config, "expert_quant", "none") != "fp8_block": super().__init__(config, layer_id=layer_id) return @@ -37,9 +47,22 @@ 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 f8d2a7494..b6868f6ba 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -21,13 +21,13 @@ 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, load_nvfp4_expert_source_banks, ) 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, @@ -157,16 +209,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 :func:`load_nvfp4_expert_sources`. """ - 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(): @@ -177,9 +231,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 d16d63492e969f9ae7f5efb497ea71f8a9505c39 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 4 Sep 2026 15:09:33 -0400 Subject: [PATCH 2/6] 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 b6868f6ba..5a0013bfb 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -215,7 +215,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 4d3518f3e89b6b101e85ab7b3c98f326940c8a73 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 4 Sep 2026 18:55:34 -0400 Subject: [PATCH 3/6] 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 | 19 ++- python/freetoken/models/qwen4_exp/config.py | 5 + python/freetoken/models/qwen4_exp/gdn.py | 29 +++- python/freetoken/models/qwen4_exp/weight.py | 67 +++++++- tests/models/qwen4_exp/test_fp8_dense.py | 121 ++++++++++++++ 7 files changed, 382 insertions(+), 20 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 1bce039cb..1a15f9cc8 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 db77d2b6e..1e684091c 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 @@ -135,13 +136,17 @@ def __init__(self, config: ModelConfig, layer_id: int) -> None: 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, - ) - self.o_proj = LinearOProj(self.qo_attn_dim, config.hidden_size, has_bias=False) + qkv_sizes = [self.qo_attn_dim * 2, self.kv_attn_dim, self.kv_attn_dim] + if config.attn_quant == "fp8_dynamic": # load-time per-tensor FP8, W8A8 GEMMs + 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 + ) + self.o_proj = LinearOProj(self.qo_attn_dim, config.hidden_size, has_bias=False) 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 bb5d1dff5..a32346b1b 100644 --- a/python/freetoken/models/qwen4_exp/config.py +++ b/python/freetoken/models/qwen4_exp/config.py @@ -7,6 +7,7 @@ import torch from freetoken.models.config import ( + fp8_dense_enabled, FullAttentionGroupConfig, LinearGatedDeltaGroupConfig, ModelConfig, @@ -179,6 +180,10 @@ def _quant(probe: str) -> str: attn_quant = _quant(f"{prefix}.self_attn.q_proj") lm_head_quant = _quant("lm_head") + if attn_quant == "none" and fp8_dense_enabled(): + # bf16 attention / GDN projections quantized at load: per-tensor e4m3 weights, + # W8A8 through cuBLASLt (layers/fp8_dynamic.py; the loader emits the fp8 tensors) + attn_quant = "fp8_dynamic" 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") diff --git a/python/freetoken/models/qwen4_exp/gdn.py b/python/freetoken/models/qwen4_exp/gdn.py index fc562e76c..a0ff870b3 100644 --- a/python/freetoken/models/qwen4_exp/gdn.py +++ b/python/freetoken/models/qwen4_exp/gdn.py @@ -9,6 +9,7 @@ from freetoken.utils import div_even from freetoken.kernel.triton.fp8_block_linear import Fp8BlockColMerged +from freetoken.layers.fp8_dynamic import Fp8DynamicColMerged, Fp8DynamicRowParallel from freetoken.kernel.triton.fp8_pertensor_linear import Fp8PerTensorColMerged from freetoken.models.qwen3_5_moe.gdn_kernels import gdn_decode_fla, gdn_prefill_chunk_fla from freetoken.models.quant_linear import make_replicated_quant @@ -96,17 +97,27 @@ def __init__( # fusion into an fp8 qkvz GEMM + a bf16 ba GEMM (matches sglang/vLLM). self._block_fp8 = expert_quant == "fp8_block" self._pertensor_fp8 = attn_quant == "fp8_pertensor" - self._fp8 = self._block_fp8 or self._pertensor_fp8 + self._dynamic_fp8 = attn_quant == "fp8_dynamic" # load-time per-tensor FP8, W8A8 + self._fp8 = self._block_fp8 or self._pertensor_fp8 or self._dynamic_fp8 self._in_proj_split = [ self._local_conv_dim, self._local_value_dim, self._local_num_v_heads, self._local_num_v_heads, ] if tp.size > 1: - assert not self._fp8 and attn_quant == "none", ( - "qwen4_exp TP shards bf16 GDN projections only" + assert self._dynamic_fp8 or (not self._fp8 and attn_quant == "none"), ( + "qwen4_exp TP shards bf16 or load-time fp8 GDN projections only" ) - if self._fp8: + 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._fp8: ColMerged = Fp8BlockColMerged if self._block_fp8 else Fp8PerTensorColMerged self.in_proj_qkvz = ColMerged( hidden_size, [self.conv_dim, self.value_dim], has_bias=False @@ -131,7 +142,9 @@ def __init__( # out_proj follows the checkpoint quant: block-fp8 / per-tensor-fp8 / compressed-tensors # NVFP4 (W4A16) / bf16. in_proj_* stay bf16 in every mode (above), so a compressed-tensors # NVFP4 checkpoint (attn_quant=="nvfp4") only makes out_proj native FP4. - if tp.size > 1: # row-parallel over the local v heads, all-reduce inside + if self._dynamic_fp8: + self.out_proj = Fp8DynamicRowParallel(self.value_dim, hidden_size) + elif tp.size > 1: # row-parallel over the local v heads, all-reduce inside self.out_proj = LinearOProj(self.value_dim, hidden_size, has_bias=False) else: self.out_proj = make_replicated_quant( @@ -194,15 +207,15 @@ 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._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/weight.py b/python/freetoken/models/qwen4_exp/weight.py index 5a0013bfb..82914c010 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, @@ -212,12 +254,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), @@ -234,9 +291,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 69e5a44556b496d996dfeb4d093a91a222c15fc3 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 4 Sep 2026 18:57:14 -0400 Subject: [PATCH 4/6] 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 82914c010..9799b974d 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -27,10 +27,12 @@ load_nvfp4_expert_source_banks, ) 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. @@ -267,6 +269,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 fec72ba28342e7f2bfe177e265973f1adad56d16 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 4 Sep 2026 19:09:48 -0400 Subject: [PATCH 5/6] 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 9799b974d..8c8e22c83 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -303,6 +303,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 3f8f2496972fd3b918cc78ba588304678116f7e4 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 4 Sep 2026 19:23:02 -0400 Subject: [PATCH 6/6] 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 8c8e22c83..be731b2d3 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():