qwen4_exp: serve the block-FP8 dense projections natively (+25% decode)

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)

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

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 [])

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
 

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,19 +248,27 @@ 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),
         desc="Loading 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