diff --git a/python/freetoken/models/quant_linear.py b/python/freetoken/models/quant_linear.py index 6ba67a9ee..a937541bd 100644 --- a/python/freetoken/models/quant_linear.py +++ b/python/freetoken/models/quant_linear.py @@ -2,7 +2,8 @@ dense projections (qwen3_5_moe, muse_glimmer). Maps the model's quant config (``expert_quant`` for the dense MLP / shared-expert path, -``attn_quant`` for attention + GatedDeltaNet projections) to the right ``BaseOP`` linear: +``attn_quant`` for attention + GatedDeltaNet projections; block-fp8 may be declared by +either, since a modelopt MIXED_PRECISION checkpoint block-quantizes only the dense side) to the right ``BaseOP`` linear: block-FP8, per-tensor-FP8 and NVFP4 implementations live under ``freetoken.kernel.triton``; the bf16 fallback is the framework's TP-aware ``freetoken.layers``. Only the *dispatch* (config -> layer class) lives here. @@ -14,7 +15,7 @@ def make_col_merged_quant(expert_quant: str, attn_quant: str, in_f: int, output_sizes: list[int], has_bias: bool = False): """Column-merged linear for a dense projection: block-fp8 / per-tensor-fp8 / nvfp4 / bf16.""" - if expert_quant == "fp8_block": + if "fp8_block" in (expert_quant, attn_quant): from freetoken.kernel.triton.fp8_block_linear import Fp8BlockColMerged return Fp8BlockColMerged(in_f, output_sizes, has_bias) @@ -34,7 +35,7 @@ def make_col_merged_quant(expert_quant: str, attn_quant: str, in_f: int, def make_replicated_quant(expert_quant: str, attn_quant: str, in_f: int, out_f: int, has_bias: bool = False): """Replicated linear for a dense projection: block-fp8 / per-tensor-fp8 / nvfp4 / bf16.""" - if expert_quant == "fp8_block": + if "fp8_block" in (expert_quant, attn_quant): from freetoken.kernel.triton.fp8_block_linear import Fp8BlockLinear return Fp8BlockLinear(in_f, out_f, has_bias) diff --git a/python/freetoken/models/qwen4_exp/attention.py b/python/freetoken/models/qwen4_exp/attention.py index d6ab2867a..9bd5a154c 100644 --- a/python/freetoken/models/qwen4_exp/attention.py +++ b/python/freetoken/models/qwen4_exp/attention.py @@ -20,6 +20,7 @@ import torch from freetoken.core import get_global_ctx from freetoken.layers import BaseOP, GemmaPlusOneRMSNorm, LinearColParallelMerged, LinearReplicated +from freetoken.models.quant_linear import make_col_merged, make_replicated from freetoken.layers.rotary import get_rope from freetoken.utils import nvtx_annotate @@ -121,10 +122,22 @@ def __init__(self, config: ModelConfig, layer_id: int) -> None: 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 - ) - self.o_proj = LinearReplicated(self.qo_attn_dim, config.hidden_size, has_bias=False) + # q|k|v are all quantized together (or all bf16), so the merged GEMM stays a + # single kernel; a modelopt MIXED_PRECISION checkpoint declares them FP8_PB_WO. + # + # The quantized factories are used ONLY on the block-FP8 path. Routing the bf16 + # case through them too would swap in the factory's generic fallback and drop the + # tensor-parallel classes this model needs (per-rank ``local_output_sizes``, a + # row-parallel ``o_proj``), which is exactly the bf16 path a rank falls back to + # under TP>1 - where the block-FP8 linears have no parallel variant. + if getattr(config, "attn_quant", "none") == "fp8_block": + self.qkv_proj = make_col_merged(config, config.hidden_size, self._qkv_split) + self.o_proj = make_replicated(config, self.qo_attn_dim, config.hidden_size) + else: + self.qkv_proj = LinearColParallelMerged( + config.hidden_size, self._qkv_split, has_bias=False + ) + self.o_proj = LinearReplicated(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..c24712bc2 100644 --- a/python/freetoken/models/qwen4_exp/config.py +++ b/python/freetoken/models/qwen4_exp/config.py @@ -6,6 +6,7 @@ import torch +from freetoken.distributed import try_get_tp_info from freetoken.models.config import ( FullAttentionGroupConfig, LinearGatedDeltaGroupConfig, @@ -120,6 +121,45 @@ def _layer_types(text: Any) -> list[str]: ] +# modelopt spellings for 128x128 per-block, weight-only FP8 on the dense modules. +_FP8_BLOCK_ALGOS = frozenset({"FP8_PB_WO", "FP8_BLOCK"}) + + +def dense_quant_mode(algo: str, quantized_layers: Any) -> str: + """The quantization mode the dense (non-expert) projections will actually be SERVED in. + + Single source of truth for the two sides that must agree: :func:`parse_config`, which + decides the modules the model BUILDS, and the weight loader, which decides the buffers + it EMITS. They previously derived this independently from the same declaration - safe + only while they cannot disagree, and they can: the block-FP8 linears have no + tensor-parallel variant, so a rank running under TP>1 has to fall back to bf16. If only + one side knew that, the loaded buffers would not match the built modules. + + Returns ``"fp8_block"`` only when the checkpoint declares per-block weight-only FP8 on a + non-expert module AND this rank can serve it; ``"none"`` otherwise (bf16, via the + dequantize-at-load path). A checkpoint carrying ``weight_scale_inv`` without declaring + the algo is "none" here and keeps the pre-existing dequant behaviour. + """ + if str(algo or "").lower() != "mixed_precision": + return "none" + declared = any( + ".mlp.experts" not in str(module) + and str((spec or {}).get("quant_algo", "")).upper() in _FP8_BLOCK_ALGOS + for module, spec in (quantized_layers or {}).items() + ) + if not declared: + return "none" + # Resolved here, once, so both sides downgrade together. ``Engine.__init__`` sets TP + # info as its very first statement, before the model config or any weight is built, so + # a rank always knows its size by the time this matters. try_get_tp_info is used rather + # than get_tp_info because config parsing also happens with no engine at all (checkpoint + # conversion, tooling, tests), where get_tp_info raises; unset means a single rank. + tp = try_get_tp_info() + if tp is not None and tp.size > 1: + return "none" + return "fp8_block" + + def parse_config(hf_config: Any) -> ModelConfig: text = getattr(hf_config, "text_config", hf_config) @@ -163,6 +203,28 @@ def parse_config(hf_config: Any) -> ModelConfig: assert bs == (128, 128), f"only 128x128 block-fp8 is supported, got {bs}" expert_quant = "fp8_block" attn_quant = dense_quant = lm_head_quant = "none" + elif algo == "mixed_precision": + # modelopt MIXED_PRECISION: the quant algo is declared per module in + # ``quantized_layers`` rather than once at the top level. The community + # NVFP4-FP8 build of Qwen3.8-Flash-Next quantizes the routed experts to NVFP4 + # (read natively by the offload cache) and the dense attn/GDN projections to + # 128x128 block-FP8, declared per module as ``FP8_PB_WO``. + quantized = get("quantized_layers") or {} + experts_nvfp4 = any( + ".mlp.experts" in str(module) + and str((spec or {}).get("quant_algo", "")).upper() == "NVFP4" + for module, spec in quantized.items() + ) + # The same map declares the dense attn/GDN projections as FP8_PB_WO + # (per-block, weight-only FP8 with a ``weight_scale_inv`` sibling). Serve + # them natively instead of dequantizing at load: the four-way in_proj fusion + # splits into an fp8 qkv|z GEMM plus a small bf16 b|a GEMM (see gdn.py), which + # halves the dense bytes read on every decode step. Resolved through + # dense_quant_mode so the loader reaches the same answer, TP downgrade + # included - see that function. + expert_quant = "nvfp4" if experts_nvfp4 else "none" + attn_quant = dense_quant_mode(algo, quantized) + dense_quant = lm_head_quant = "none" else: is_fp4 = "fp4" in algo ignore = list(get("ignore") or []) diff --git a/python/freetoken/models/qwen4_exp/gdn.py b/python/freetoken/models/qwen4_exp/gdn.py index 69838153f..1f3e6e44b 100644 --- a/python/freetoken/models/qwen4_exp/gdn.py +++ b/python/freetoken/models/qwen4_exp/gdn.py @@ -83,7 +83,10 @@ def __init__( # qkv|z carry a weight scale (block-fp8 weight_scale_inv, or per-tensor FP8 # weight_scale); b|a stay bf16. Both quant modes therefore split the four-way # fusion into an fp8 qkvz GEMM + a bf16 ba GEMM (matches sglang/vLLM). - self._block_fp8 = expert_quant == "fp8_block" + # Block-fp8 dense can come from either a wholly block-fp8 checkpoint + # (expert_quant) or a modelopt MIXED_PRECISION one where only the dense + # attn/GDN projections are FP8_PB_WO while the experts are NVFP4 (attn_quant). + self._block_fp8 = "fp8_block" in (expert_quant, attn_quant) self._pertensor_fp8 = attn_quant == "fp8_pertensor" self._fp8 = self._block_fp8 or self._pertensor_fp8 diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index f8d2a7494..e1f2b22ed 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -22,6 +22,7 @@ import torch from freetoken.distributed import get_tp_info from freetoken.models.loader import drop_page_cache, iter_weight_files +from freetoken.models.qwen4_exp.config import dense_quant_mode from freetoken.models.nvfp4_banks import ( Nvfp4ExpertSourceSpec, load_nvfp4_expert_source_banks, @@ -46,7 +47,8 @@ desc="Qwen3.8-Flash-Next NVFP4 experts", ) # Per-tensor modelopt quant scales; consumed with their ``.weight`` (experts) or unused. -_SCALE_SUFFIXES = (".weight_scale", ".weight_scale_2", ".input_scale") +# ``.weight_scale_inv`` is the 128x128 block-FP8 reciprocal scale (see _load_maybe_block_fp8). +_SCALE_SUFFIXES = (".weight_scale", ".weight_scale_2", ".weight_scale_inv", ".input_scale") # The n-gram table itself: too big for the dense state dict, loaded by load_ple_table. _PLE_TABLE_INFIX = ".ple.ple_embedding.ngram_embedding." @@ -98,15 +100,20 @@ } -def _rename(raw_name: str) -> str | None: - """Checkpoint key -> FreeToken state-dict key, or None to skip.""" +def _rename(raw_name: str, keep_scale_inv: bool = False) -> str | None: + """Checkpoint key -> FreeToken state-dict key, or None to skip. + + ``keep_scale_inv`` retains the block-FP8 ``weight_scale_inv`` tensors, which the + fp8 linears need alongside their weight; they are dropped otherwise.""" if raw_name.startswith(("mtp.", "model.visual.", "visual.")): return None if _PLE_TABLE_INFIX in raw_name: return None # n-gram table + its scale: load_ple_table if _EXPERT_RE.search(raw_name): return None # routed experts: offload source banks - if raw_name.endswith(_SCALE_SUFFIXES): + if raw_name.endswith(_SCALE_SUFFIXES) and not ( + keep_scale_inv and raw_name.endswith(".weight_scale_inv") + ): return None if raw_name.startswith("model.language_model."): return "model." + raw_name[len("model.language_model.") :] @@ -116,10 +123,11 @@ def _rename(raw_name: str) -> str | None: def _try_fuse( - name: str, tensor: torch.Tensor, buf: dict[str, dict[int, torch.Tensor]] + name: str, tensor: torch.Tensor, buf: dict[str, dict[int, torch.Tensor]], + table: dict[str, tuple[tuple[str, ...], int]] | None = None, ) -> tuple[str, torch.Tensor] | tuple[()] | None: """Buffer a fusion part; return the merged ``(name, tensor)`` once all parts arrive, ``()`` while incomplete, ``None`` if ``name`` is not a fusion part.""" - for fused_suffix, (parts, pad_to) in _FUSIONS.items(): + for fused_suffix, (parts, pad_to) in (table or _FUSIONS).items(): for idx, part in enumerate(parts): if not name.endswith(part): continue @@ -137,6 +145,84 @@ def _try_fuse( return None +def _load_maybe_block_fp8(f, raw_name: str, keyset: set[str]) -> torch.Tensor: + """Load ``raw_name``, dequantizing 128x128 block-FP8 to bf16 when a sibling + ``.weight_scale_inv`` is present in the same shard; pass plain bf16 through unchanged. + + The official modelopt checkpoint keeps the dense attn/GDN/HC/PLE projections bf16 (they are + on the quant ``ignore`` list), but some community requants -- e.g. the lovedheart NVFP4-FP8 + build that fits Qwen3.8-Flash-Next on a 24 GB card -- store those dense weights as block-FP8. + Without this they reach ``_try_fuse`` as fp8 and crash on the fp8+bf16 ``torch.cat``.""" + tensor = f.get_tensor(raw_name) + if raw_name.endswith(".weight"): + base = raw_name[: -len(".weight")] + if base + ".weight_scale_inv" in keyset: + from freetoken.kernel.triton.fp8_block_linear import dequant_block_fp8 + + return dequant_block_fp8( + tensor, f.get_tensor(base + ".weight_scale_inv") + ).to(torch.bfloat16) + return tensor + + +# Serving the dense side natively as block-FP8 changes which buffers the model expects: +# the four-way in_proj fusion splits into an fp8 qkv|z GEMM plus a small bf16 b|a GEMM +# (see gdn.py), and each fp8 group fuses its ``weight_scale_inv`` on the same axis as its +# ``weight``. Every fp8 part is a whole number of 128-row blocks, so the per-block scales +# concatenate exactly alongside the rows they describe. +_BLOCK_FP8_FUSE: dict[str, tuple[str, ...]] = { + ".self_attn.qkv_proj": ( + ".self_attn.q_proj", ".self_attn.k_proj", ".self_attn.v_proj", + ), + ".linear_attn.in_proj_qkvz": ( + ".linear_attn.in_proj_qkv", ".linear_attn.in_proj_z", + ), +} +_BLOCK_BF16_FUSE: dict[str, tuple[str, ...]] = { + ".linear_attn.in_proj_ba": (".linear_attn.in_proj_b", ".linear_attn.in_proj_a"), +} +_BLOCK_FP8_KINDS = (".weight", ".weight_scale_inv") + + +def _block_fp8_fusions() -> dict[str, tuple[tuple[str, ...], int]]: + """``_FUSIONS`` with the two attention groups replaced by their block-FP8 split.""" + table = { + key: val + for key, val in _FUSIONS.items() + if key not in (".self_attn.qkv_proj.weight", ".linear_attn.in_proj.weight") + } + for fused, parts in _BLOCK_FP8_FUSE.items(): + for kind in _BLOCK_FP8_KINDS: + table[fused + kind] = (tuple(part + kind for part in parts), 0) + for fused, parts in _BLOCK_BF16_FUSE.items(): + table[fused + ".weight"] = (tuple(part + ".weight" for part in parts), 0) + return table + + +_FUSIONS_BLOCK_FP8 = _block_fp8_fusions() + + +def _dense_is_block_fp8(model_path: str) -> bool: + """True when the dense (non-expert) modules will be SERVED as per-block weight-only + FP8 -- exactly when config.py sets ``attn_quant="fp8_block"``. + + The decision itself lives in :func:`~freetoken.models.qwen4_exp.config.dense_quant_mode` + and is only read here, so the buffers this loader emits cannot disagree with the modules + the model built. That matters under TP: the block-FP8 linears have no parallel variant, + so a rank at TP>1 downgrades to bf16, and both sides have to downgrade together. + + A checkpoint carrying ``weight_scale_inv`` WITHOUT declaring the algo is not block-FP8 + here either; it falls through to ``_load_maybe_block_fp8`` and is dequantized as before. + """ + try: + with open(os.path.join(model_path, "config.json"), encoding="utf-8") as fh: + quant = json.load(fh).get("quantization_config") or {} + except (OSError, ValueError): + return False + algo = quant.get("quant_algo") or quant.get("quant_method") or "" + return dense_quant_mode(algo, quant.get("quantized_layers")) == "fp8_block" + + def iter_weights( model_path: str, device: torch.device, @@ -162,6 +248,9 @@ def iter_weights( if not include_non_moe: return + # Declared block-FP8 dense is served natively; anything else keeps the dequant path. + block_fp8 = _dense_is_block_fp8(model_path) + fusions = _FUSIONS_BLOCK_FP8 if block_fp8 else _FUSIONS fuse_buf: dict[str, dict[int, torch.Tensor]] = {} for file in tqdm( iter_weight_files(model_path), @@ -169,12 +258,17 @@ def iter_weights( disable=not get_tp_info().is_primary(), ): with safetensors.safe_open(file, framework="pt", device=str(device)) as f: + keyset = set(f.keys()) for raw_name in f.keys(): - name = _rename(raw_name) + name = _rename(raw_name, keep_scale_inv=block_fp8) if name is None: continue - tensor = f.get_tensor(raw_name) - fused = _try_fuse(name, tensor, fuse_buf) + tensor = ( + f.get_tensor(raw_name) + if block_fp8 + else _load_maybe_block_fp8(f, raw_name, keyset) + ) + fused = _try_fuse(name, tensor, fuse_buf, fusions) if fused is not None: if fused != (): # () means buffered, not yet complete yield fused