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/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 new file mode 100644 index 000000000..caeebb322 --- /dev/null +++ b/python/freetoken/layers/fp8_dynamic.py @@ -0,0 +1,211 @@ +"""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 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. +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 +from .embedding import ParallelLMHead + +FP8 = torch.float8_e4m3fn +E4M3_MAX = 448.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. 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) + 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 _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 + 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) + + +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) + 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 + 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 + + +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 + ) + + +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", + "fp8_dynamic_linear", + "quant_per_tensor", +] 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..e2dc1478c 100644 --- a/python/freetoken/layers/quantization/moe/nvfp4.py +++ b/python/freetoken/layers/quantization/moe/nvfp4.py @@ -41,14 +41,20 @@ class TritonNvfp4MoEKernel(MoEKernel): cpu_format = "nvfp4" def unusable_reason(self, cfg: MoEConfig) -> str | None: - reason = self._common_reject(cfg, resident_ok=False, tp_ok=False, cpu_ok=True, plain_silu_only=False) + # tp_ok: the source stream is sliced along the intermediate axis per rank + # (nvfp4_banks._tp_shard) and this kernel sizes its banks from + # cfg.local_intermediate, so a rank holds and reads only its own half. The routed + # output is then a partial sum, which the MoE layer reduces (_maybe_all_reduce, or + # one combined all-reduce in qwen4_exp's block). marlin and b12x stay off: their + # pack() repacks the rows and neither has been checked against a sharded bank. + reason = self._common_reject(cfg, resident_ok=False, tp_ok=True, cpu_ok=True, plain_silu_only=False) if reason: return reason reason = gated_epilogue_reason(cfg) return f"triton nvfp4 MoE kernel: {reason}" if reason else None def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]: - i, h = cfg.intermediate, cfg.hidden + i, h = cfg.local_intermediate, cfg.hidden return { "gate_up": BankSpec((2 * i, h // 2), torch.uint8), "gate_up_scale": BankSpec((2 * i, h // GROUP), FP8), @@ -61,7 +67,7 @@ def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]: def pack(self, pieces, cfg: MoEConfig, out): out["gate_up"].copy_(fused_piece(pieces, "gate_up")) out["gate_up_scale"].copy_(fused_piece(pieces, "gate_up_scale")) - out["gate_up_global"].copy_(fused_global(pieces, cfg.intermediate)) + out["gate_up_global"].copy_(fused_global(pieces, cfg.local_intermediate)) out["down"].copy_(pieces["down"]) out["down_scale"].copy_(pieces["down_scale"]) out["down_global"].copy_(global_rows(pieces["down_global"], cfg.hidden)) @@ -241,7 +247,7 @@ def worth_it(self, cfg: MoEConfig) -> bool: return (8, 0) <= backend.device_capability() < (10, 0) def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]: - i, h = cfg.intermediate, cfg.hidden + i, h = cfg.local_intermediate, cfg.hidden return { "gate_up": BankSpec((h // GROUP, 4 * i), torch.int32), "gate_up_scale": BankSpec((h // GROUP, 2 * i), FP8), @@ -252,7 +258,7 @@ def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]: } def pack(self, pieces, cfg: MoEConfig, out): - i, h = cfg.intermediate, cfg.hidden + i, h = cfg.local_intermediate, cfg.hidden device = torch.device("cuda") gu, gus, gug = fused_piece(pieces, "gate_up"), fused_piece(pieces, "gate_up_scale"), fused_global(pieces, i) dn, dns, dng = pieces["down"], pieces["down_scale"], global_rows(pieces["down_global"], h) @@ -503,7 +509,7 @@ def unusable_reason(self, cfg: MoEConfig) -> str | None: def worth_it(self, cfg: MoEConfig) -> bool: # NOTE: never auto-selected. flashinfer's cute launcher indexes each bank with int32 element offsets, so the GPU slot cache is capped at (2^31 - 1) / elements-per-slot (about 1000 slots for GLM-5.3-Flash's 12960 experts); until flashinfer lifts that, b12x stays behind triton in the table; the selection does not weigh the cap, cache-auto clamps to slot_limit() and OffloadMoeCache refuses a larger cache. - return cfg.intermediate >= B12X_MIN_INTERMEDIATE + return cfg.local_intermediate >= B12X_MIN_INTERMEDIATE def slot_limit(self, cfg: MoEConfig) -> int | None: # the cute launcher indexes each bank with int32 element offsets @@ -512,7 +518,7 @@ def slot_limit(self, cfg: MoEConfig) -> int | None: def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]: # flashinfer's prepared tiles: the Marlin shapes, byte-identical to the native rows - i, h = cfg.intermediate, cfg.hidden + i, h = cfg.local_intermediate, cfg.hidden return { "gate_up": BankSpec((h // GROUP, 4 * i), torch.int32), "gate_up_scale": BankSpec((h // GROUP, 2 * i), FP8), @@ -525,7 +531,7 @@ def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]: def pack(self, pieces, cfg: MoEConfig, out): from flashinfer.fused_moe.cute_dsl.blackwell_sm12x.moe_w4a16_prepare import prepare_w4a16_packed_weights - i, h = cfg.intermediate, cfg.hidden + i, h = cfg.local_intermediate, cfg.hidden device = torch.device("cuda") gu = fused_piece(pieces, "gate_up").to(device) gug = fused_global(pieces, i).to(device).float() diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py index 07a16df5c..62340ab75 100644 --- a/python/freetoken/models/config.py +++ b/python/freetoken/models/config.py @@ -18,6 +18,21 @@ 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 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/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..5b7e90397 100644 --- a/python/freetoken/models/qwen4_exp/attention.py +++ b/python/freetoken/models/qwen4_exp/attention.py @@ -19,9 +19,17 @@ 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.fp8_dynamic import Fp8DynamicColMerged, Fp8DynamicRowParallel 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,15 +131,38 @@ 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] - self.qkv_proj = LinearColParallelMerged( - config.hidden_size, self._qkv_split, has_bias=False, - quant_config=config.quant, prefix=f"{prefix}.qkv_proj", - ) - self.o_proj = LinearReplicated( - self.qo_attn_dim, config.hidden_size, has_bias=False, - quant_config=config.quant, prefix=f"{prefix}.o_proj", + # 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] + 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 @@ -147,21 +178,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/config.py b/python/freetoken/models/qwen4_exp/config.py index 4bcd28de3..bb9966987 100644 --- a/python/freetoken/models/qwen4_exp/config.py +++ b/python/freetoken/models/qwen4_exp/config.py @@ -7,6 +7,8 @@ from freetoken.layers.quantization import QuantConfig from freetoken.models.config import ( + fp8_dense_enabled, + fp8_lmhead_enabled, FullAttentionGroupConfig, LinearGatedDeltaGroupConfig, ModelConfig, @@ -109,6 +111,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) @@ -138,6 +157,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 +270,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 ed49dd42f..763cb78ca 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.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 @@ -33,6 +36,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 @@ -50,13 +54,42 @@ 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] - if self._split_in_proj: + # 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, + ] + # 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 or load-time fp8 GDN projections only" + ) + 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", @@ -68,21 +101,29 @@ 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( - 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() @@ -140,15 +181,17 @@ 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: + 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.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) - z = z.reshape(total, self.num_v_heads, self.head_v_dim) + 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 +200,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 +213,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/model.py b/python/freetoken/models/qwen4_exp/model.py index a4b63362d..3238459a6 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 .config import use_fp8_lmhead from freetoken.utils import nvtx_annotate from .attention import Qwen4ExpAttention @@ -50,6 +51,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, ) @@ -128,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 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 + + 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/moe.py b/python/freetoken/models/qwen4_exp/moe.py index 3687f2803..c10407861 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,39 @@ 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, *, prefix: str = "" + ) -> None: + super().__init__(config, layer_id=layer_id, prefix=prefix) + 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..10f960852 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -21,15 +21,17 @@ 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, 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. @@ -137,6 +139,107 @@ 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) + + +# 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", +) +# 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 + + +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, *, + 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 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) + 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 + # 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 (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 + else: + yield name, t + + def iter_weights( model_path: str, device: torch.device, @@ -158,16 +261,44 @@ 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 freetoken.models.config import fp8_dense_enabled, fp8_lmhead_enabled + + 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). + config = ( + parse_config(cached_load_hf_config(model_path)) + 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( + "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 or lmhead: + yield from _fp8_dense(name, tensor, config, tp.size, dense=fp8, lm_head=lmhead) + else: + yield name, tensor + 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,11 +309,17 @@ 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 from emit(name, tensor) continue - yield name, tensor + yield from emit(name, tensor) assert not fuse_buf, f"Incomplete projection fusions: {sorted(fuse_buf)}" + 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. + torch.cuda.empty_cache() # ====================================================================================== 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)}" 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..058515b18 --- /dev/null +++ b/tests/models/qwen4_exp/test_fp8_dense.py @@ -0,0 +1,153 @@ +"""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:]) + # 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(): + 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() + + +@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)) + + +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"] 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..b943cff1c --- /dev/null +++ b/tests/models/test_nvfp4_banks_tp.py @@ -0,0 +1,64 @@ +"""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.layers.quantization.moe.base import MoEConfig +from freetoken.models.nvfp4_banks import _tp_shard + +INTER, HIDDEN = 640, 2560 + + +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 + ) + + +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