feat(models): add nvidia/Qwen3.8-Flash-Next-FP4-Local checkpoint support

docs/models.md

@@ -10,7 +10,7 @@ for them; other checkpoints of the same architectures work too.
 | GLM-5.3-Flash | [RedHatAI/GLM-5.3-Flash-NVFP4](https://huggingface.co/RedHatAI/GLM-5.3-Flash-NVFP4) |
 | GLM-5.2 | [nvidia/GLM-5.2-NVFP4](https://huggingface.co/nvidia/GLM-5.2-NVFP4) |
 | GLM-4.7 | [nvidia/GLM-4.7-NVFP4](https://huggingface.co/nvidia/GLM-4.7-NVFP4) |
-| Qwen3.8-Flash-Next | [Qwen/Qwen3.8-Flash-Next-FP8](https://huggingface.co/Qwen/Qwen3.8-Flash-Next-FP8), [RadixArk/Qwen3.8-Flash-Next-NVFP4](https://huggingface.co/RadixArk/Qwen3.8-Flash-Next-NVFP4), [nvidia/Qwen3.8-Flash-Next-NVFP4](https://huggingface.co/nvidia/Qwen3.8-Flash-Next-NVFP4) |
+| Qwen3.8-Flash-Next | [Qwen/Qwen3.8-Flash-Next-FP8](https://huggingface.co/Qwen/Qwen3.8-Flash-Next-FP8), [RadixArk/Qwen3.8-Flash-Next-NVFP4](https://huggingface.co/RadixArk/Qwen3.8-Flash-Next-NVFP4), [nvidia/Qwen3.8-Flash-Next-NVFP4](https://huggingface.co/nvidia/Qwen3.8-Flash-Next-NVFP4), [nvidia/Qwen3.8-Flash-Next-FP4-Local](https://huggingface.co/nvidia/Qwen3.8-Flash-Next-FP4-Local) |
 | Qwen3.6 / Qwen3.5 MoE | [Qwen/Qwen3.6-35B-A3B](https://huggingface.co/Qwen/Qwen3.6-35B-A3B) ([-FP8](https://huggingface.co/Qwen/Qwen3.6-35B-A3B-FP8)), [nvidia/Qwen3.6-35B-A3B-NVFP4](https://huggingface.co/nvidia/Qwen3.6-35B-A3B-NVFP4), [Qwen/Qwen3.5-35B-A3B](https://huggingface.co/Qwen/Qwen3.5-35B-A3B) ([-FP8](https://huggingface.co/Qwen/Qwen3.5-35B-A3B-FP8)) |
 | Qwen3.8 / Qwen3.6 dense | [Qwen/Qwen3.8-27B](https://huggingface.co/Qwen/Qwen3.8-27B) ([-FP8](https://huggingface.co/Qwen/Qwen3.8-27B-FP8)), [RadixArk/Qwen3.8-27B-NVFP4](https://huggingface.co/RadixArk/Qwen3.8-27B-NVFP4), [Qwen/Qwen3.6-27B](https://huggingface.co/Qwen/Qwen3.6-27B) ([-FP8](https://huggingface.co/Qwen/Qwen3.6-27B-FP8)), [nvidia/Qwen3.6-27B-NVFP4](https://huggingface.co/nvidia/Qwen3.6-27B-NVFP4) |
 | Qwen3-MoE | [Qwen/Qwen3-30B-A3B](https://huggingface.co/Qwen/Qwen3-30B-A3B) |

python/freetoken/engine/config.py

@@ -32,7 +32,7 @@ class EngineConfig:
     moe_backend: str | None = field(default=None, repr=False)
     # --quant-backend: layer[.kind]=kernel entries, comma separated
     quant_backend: str | None = None
-    # PLE table backend: "disk" (default) reads rows from the checkpoint files per fill, "pinned" preloads the table into page-locked host RAM.
+    # PLE table backend: "disk" (default) reads rows from the checkpoint files per fill, "pinned" preloads the table into page-locked host RAM. NVFP4 (packed e2m1) tables fall back to "pinned" automatically; the disk row store streams native rows only.
     ple_backend: str = "disk"
     # DeepSeek-V4.1 Decoder SWA Bounded Replay: "bounded" (default) runs the decoder layers on each
     # prompt's last window; "exact" runs them on every token.

python/freetoken/kernel/triton/ple.py

@@ -8,6 +8,9 @@
 the mapped device address on WDDM (``kernel/pinned.device_ptr``). One program per requested
 row: read the row, widen to fp32, apply the per-tensor scale, store bf16.
 
+NVFP4 checkpoints pack the rows as e2m1 (two per byte) plus a per-16 e4m3 scale row and one
+global scale: the gather reads both rows, decodes the nibbles, and multiplies in place.
+
 Ids outside the table store zeros.
 """
 
@@ -23,16 +26,32 @@
 _NUM_WARPS = 1
 
 
+@triton.jit
+def _e2m1_f32(n: tl.uint8) -> tl.float32:
+    """The 16 e2m1 nibbles: 0, +/-0.5, +/-1, +/-1.5, +/-2, +/-3, +/-4, +/-6 (sign = bit 3)."""
+    n = n.to(tl.int32)
+    mant = (n & 1).to(tl.float32)
+    exp2 = (n >> 1) & 3
+    # exponent 0 is subnormal (0.mant); otherwise (1.mant) * 2 ** (exp2 - 1)
+    mag = tl.where(exp2 == 0, mant * 0.5, (1.0 + mant * 0.5) * tl.where(exp2 == 1, 1.0, tl.where(exp2 == 2, 2.0, 4.0)))
+    sign = ((n >> 3) & 1) * mag
+    return mag - 2 * sign
+
+
 @triton.jit
 def _ple_gather_kernel(
     table_ptr,
+    scale_ptr,
     ids_ptr,
     out_ptr,
     scale,
     num_rows,
     EMB_DIM: tl.constexpr,
     IS_FP8: tl.constexpr,
+    IS_FP4: tl.constexpr,
     BLOCK_D: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    BLOCK_S: tl.constexpr,
 ):
     row = tl.program_id(0)
     idx = tl.load(ids_ptr + row).to(tl.int64)
@@ -49,6 +68,24 @@ def _ple_gather_kernel(
             # pre-sm_89 has no fp8e4nv type: load raw bytes and decode in software
             base = table_ptr.to(tl.int64).to(tl.pointer_type(tl.uint8))
             values = e4m3_u8_to_f32(tl.load(base + idx * EMB_DIM + offsets, mask=mask, other=0))
+    elif IS_FP4:
+        # packed e2m1: two values per byte (low nibble first), per-16 e4m3 scale row, global scale
+        pack_offsets = tl.arange(0, BLOCK_H)
+        vbytes = tl.load(table_ptr.to(tl.int64).to(tl.pointer_type(tl.uint8)) + idx * (EMB_DIM // 2) + pack_offsets,
+                         mask=pack_offsets < EMB_DIM // 2, other=0)
+        lo = _e2m1_f32(vbytes & 0x0F)
+        hi = _e2m1_f32((vbytes >> 4) & 0x0F)
+        values = tl.join(lo, hi).reshape([BLOCK_D])
+        scale_offsets = tl.arange(0, BLOCK_S)
+        scale_mask = scale_offsets < EMB_DIM // 16
+        if e4m3_native_cx():
+            scale_base = scale_ptr.to(tl.int64).to(tl.pointer_type(tl.float8e4nv))
+            svals = tl.load(scale_base + idx * (EMB_DIM // 16) + scale_offsets, mask=scale_mask, other=0.0).to(tl.float32)
+        else:
+            scale_base = scale_ptr.to(tl.int64).to(tl.pointer_type(tl.uint8))
+            svals = e4m3_u8_to_f32(tl.load(scale_base + idx * (EMB_DIM // 16) + scale_offsets, mask=scale_mask, other=0))
+        svals = tl.broadcast_to(svals.reshape([BLOCK_S, 1]), [BLOCK_S, 16]).reshape([BLOCK_D])
+        values = values * svals
     else:
         base = table_ptr.to(tl.int64).to(tl.pointer_type(tl.bfloat16))
         values = tl.load(base + idx * EMB_DIM + offsets, mask=mask, other=0.0).to(tl.float32)
@@ -68,25 +105,35 @@ def ple_gather_rows(
     out: torch.Tensor,
     scale: float = 1.0,
     is_fp8: bool = True,
+    scale_ptr: int | None = None,
 ) -> torch.Tensor:
     """Gather ``row_ids`` from the host-resident table at ``table_ptr`` into ``out``.
 
     ``row_ids`` is a flat device int tensor; ``out`` is ``[row_ids.numel(), embed_dim]``
     bf16 on the same device. ``table_ptr`` is the address the GPU must dereference
-    (``kernel/pinned.device_ptr``), not necessarily the host ``data_ptr``.
+    (``kernel/pinned.device_ptr``), not necessarily the host ``data_ptr``. Packed e2m1 tables
+    (``is_fp4``) pass the per-16 e4m3 scale row bank as ``scale_ptr`` and the global scale as
+    ``scale``.
     """
+    assert not (is_fp8 and scale_ptr is not None), "a packed table is not also a native fp8 table"
+    if scale_ptr is not None:
+        is_fp8 = False
     n = row_ids.numel()
     assert out.shape == (n, embed_dim) and out.is_contiguous(), out.shape
     if n:
         _ple_gather_kernel[(n,)](
             table_ptr,
+            scale_ptr if scale_ptr is not None else table_ptr,
             row_ids,
             out,
             float(scale),
             num_rows,
             EMB_DIM=embed_dim,
             IS_FP8=is_fp8,
+            IS_FP4=scale_ptr is not None,
             BLOCK_D=triton.next_power_of_2(embed_dim),
+            BLOCK_H=triton.next_power_of_2(embed_dim // 2),
+            BLOCK_S=triton.next_power_of_2(embed_dim // 16),
             num_warps=_NUM_WARPS,
         )
     return out

python/freetoken/models/nvfp4_banks.py

@@ -27,6 +27,12 @@ class Nvfp4ExpertSourceSpec:
     # The checkpoint stores the QUANT-side global scale (local fp8 scales were
     # multiplied by it before the cast); the banks keep its reciprocal.
     global_reciprocal: bool = False
+    # The checkpoint stores the block scales as a lossless CSF slab pair instead of
+    # the plain E4M3 matrix (e.g. "weight_scale.nvfp4_csf_fixed" /
+    # "weight_scale.nvfp4_csf_exceptions"); scale_decoder expands the parts into the
+    # plain matrix the banks consume.
+    scale_pair: tuple[str, ...] | None = None
+    scale_decoder: Callable[..., torch.Tensor] | None = None
 
 
 def _canon_kind(spec: "Nvfp4ExpertSourceSpec", kind: str) -> str:
@@ -85,6 +91,43 @@ def _kind_suffix(kind: str) -> str:
     return {"weight": "", "weight_scale": "_scale", "weight_scale_2": "_global"}[kind]
 
 
+def decode_csf_scales(fixed: torch.Tensor, exceptions: torch.Tensor) -> torch.Tensor:
+    """Expand the CSF slab encoding of an NVFP4 per-16 scale into the plain E4M3 matrix (CPU).
+
+    ``fixed`` is U8 ``[slabs, 16 * (1 + row_bytes)]``: the first 16 bytes of each slab hold the
+    per-row scale bases (one byte per row of the slab's 16 rows), the rest are row-major code
+    bytes with column c at bit 4c. ``exceptions`` is U32 records, ``pos = rec & 0xFFFFFF``
+    (row-major scale index, below 2**24) and ``value = rec >> 24`` (the exact E4M3 byte),
+    applied after base + offset.
+    """
+    if fixed.ndim != 2 or fixed.dtype != torch.uint8:
+        raise ValueError(f"CSF fixed slab tensor is {tuple(fixed.shape)} {fixed.dtype}, expected U8 [slabs, slab_bytes]")
+    slabs, slab_bytes = fixed.shape
+    row_bytes = slab_bytes // 16 - 1
+    if row_bytes < 1 or slab_bytes != 16 * (1 + row_bytes):
+        raise ValueError(f"CSF slab width {slab_bytes} does not hold 16 bases plus 16 * row_bytes code bytes")
+    rows = slabs * 16
+    cols = row_bytes * 2
+    data = fixed.to(torch.int16)
+    base = data[:, :16].reshape(rows, 1)
+    codes = data[:, 16:].reshape(rows, row_bytes)
+    bits = 4 * torch.arange(cols)
+    scale = (base + ((codes[:, bits // 8] >> bits % 8) & 0x0F)) & 0xFF
+    if exceptions.numel():
+        # Records are U32 on disk or raw U8 bytes; the bit math runs in int64 because CPU torch
+        # has no uint32 shifts and uint32 -> int32 casts wrap records with the sign bit set.
+        rec = exceptions.reshape(-1)
+        if rec.dtype == torch.uint8:
+            rec = rec.view(torch.uint32)
+        if rec.dtype not in (torch.uint32, torch.int32):
+            raise ValueError(f"CSF exceptions are U32 records (or U8 bytes), got {rec.dtype}")
+        rec = rec.to(torch.int64)
+        scale = scale.to(torch.int32)
+        flat = scale.view(-1)
+        flat.scatter_(0, rec & 0xFFFFFF, ((rec >> 24) & 0xFF).to(torch.int32))
+    return scale.to(torch.uint8).view(torch.float8_e4m3fn)
+
+
 def iter_nvfp4_expert_pieces(
     model_path: str,
     config,
@@ -117,6 +160,7 @@ def iter_nvfp4_expert_pieces(
     keep = None if layers is None else set(layers)
     act = _activation_scale()
     wanted: dict[str, tuple[int, int, str]] = {}
+    pair_kinds = spec.scale_pair
     for name in weight_map:
         match = spec.key_pattern.match(name)
         if match is None:
@@ -128,23 +172,41 @@ def iter_nvfp4_expert_pieces(
         if proj not in spec.proj_to_role:
             raise ValueError(f"{spec.desc}: unknown NVFP4 expert projection {proj!r}")
         kind = _canon_kind(spec, match.group("kind"))
-        if kind not in ("weight", "weight_scale", "weight_scale_2"):
+        role = spec.proj_to_role[proj]
+        if pair_kinds is not None:
+            if kind in pair_kinds:
+                wanted[name] = (bank_layer, int(match.group("expert")), f"{role}~{kind}")
+            elif kind in ("weight", "weight_scale_2"):
+                wanted[name] = (bank_layer, int(match.group("expert")), role + _kind_suffix(kind))
+            else:
+                raise ValueError(f"{spec.desc}: unknown NVFP4 expert tensor kind {kind!r}")
+        elif kind not in ("weight", "weight_scale", "weight_scale_2"):
             raise ValueError(f"{spec.desc}: unknown NVFP4 expert tensor kind {kind!r}")
-        wanted[name] = (bank_layer, int(match.group("expert")), spec.proj_to_role[proj] + _kind_suffix(kind))
+        else:
+            wanted[name] = (bank_layer, int(match.group("expert")), role + _kind_suffix(kind))
         if kind == "weight" and act is not None:
             act_name = name[: match.start("kind")] + act.name
             if act_name in weight_map:
-                wanted[act_name] = (bank_layer, int(match.group("expert")), spec.proj_to_role[proj] + "_input")
+                wanted[act_name] = (bank_layer, int(match.group("expert")), role + "_input")
     experts = (_num_moe_layers(config) if keep is None else len(keep)) * config.num_experts
-    acts = [name for name, (_, _, role) in wanted.items() if role.endswith("_input")]
+    acts = [name for name, (_, _, slot) in wanted.items() if slot.endswith("_input")]
     # a partial set serves no W4A4 kernel; the pieces then carry none
     if len(acts) != experts * 3:
         for name in acts:
             del wanted[name]
         acts = []
-    per_expert = 12 if acts else 9
-    if len(wanted) != experts * per_expert:
-        raise ValueError(f"{spec.desc}: found {len(wanted)} expert tensors, expected {experts * per_expert}")
+    # A CSF checkpoint keeps every input_scale in its own trailing shard; streaming them
+    # with the weights makes the first piece wait on that shard while every other tensor
+    # piles up unread in the queue. The CSF path reads them out of band instead.
+    act_keys: dict[str, tuple[int, int, str]] = {}
+    if pair_kinds is not None:
+        for name in acts:
+            act_keys[name] = wanted.pop(name)
+        acts = list(act_keys)
+    per_expert = (3 * (2 + len(pair_kinds)) + (3 if acts else 0)) if pair_kinds is not None else (12 if acts else 9)
+    streamed = per_expert - (3 if pair_kinds is not None and acts else 0)
+    if len(wanted) != experts * streamed:
+        raise ValueError(f"{spec.desc}: found {len(wanted)} expert tensors, expected {experts * streamed}")
 
     def _serial():
         by_shard: dict[str, list[str]] = collections.defaultdict(list)
@@ -165,7 +227,57 @@ def _parallel():
         for name, tensor in iter_expert_tensors_parallel(folder, lambda n: n in wanted, workers=workers, chunk=chunk):
             yield name, _ingest(spec, act, wanted[name][2], tensor)
 
-    return per_expert_pieces(_parallel() if parallel else _serial(), wanted.get, tensors_per_expert=per_expert)
+    if pair_kinds is None:
+        return per_expert_pieces(_parallel() if parallel else _serial(), wanted.get, tensors_per_expert=per_expert)
+    if spec.scale_decoder is None:
+        raise ValueError(f"{spec.desc}: scale_pair requires a scale_decoder")
+    decoder = spec.scale_decoder
+
+    # CSF checkpoints stream the slab parts as separate tensors; group them back into the
+    # same logical roles (plus the W4A4 input scales when the set is complete), decoding
+    # the scale pair on the fly.
+    roles = list(spec.proj_to_role.values())
+    slots_per_expert = {
+        slot
+        for role in roles
+        for slot in [f"{role}", f"{role}_global", *(f"{role}~{kind}" for kind in pair_kinds)]
+    }
+    act_index = {key: name for name, key in act_keys.items()}
+
+    def _csf_pieces():
+        acts_in: dict[str, torch.Tensor] = {}
+        if act_keys:
+            act_shards: dict[str, list[str]] = collections.defaultdict(list)
+            for name in act_keys:
+                act_shards[weight_map[name]].append(name)
+            for shard in tqdm(sorted(act_shards), desc=f"Loading {spec.desc} scales", disable=not primary):
+                path = os.path.join(folder, shard)
+                _drop(path)
+                with safetensors.safe_open(path, framework="pt", device="cpu") as f:
+                    for name in act_shards[shard]:
+                        acts_in[name] = _ingest(spec, act, act_keys[name][2], f.get_tensor(name))
+                _drop(path)
+
+        def _assemble(layer: int, expert: int, parts: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+            grouped: dict[str, torch.Tensor] = {}
+            for role in roles:
+                grouped[role] = parts[role].unsqueeze(0)
+                grouped[f"{role}_global"] = parts[f"{role}_global"].reshape(1, -1)
+                if acts:
+                    grouped[f"{role}_input"] = acts_in[act_index[(layer, expert, f"{role}_input")]].reshape(1, -1)
+                grouped[f"{role}_scale"] = decoder(
+                    *[parts[f"{role}~{kind}"] for kind in pair_kinds]
+                ).unsqueeze(0)
+            return grouped
+
+        yield from per_expert_pieces(
+            _parallel() if parallel else _serial(),
+            wanted.get,
+            tensors_per_expert=len(slots_per_expert),
+            assemble=_assemble,
+        )
+
+    return _csf_pieces()
 
 
-__all__ = ["Nvfp4ExpertSourceSpec", "iter_nvfp4_expert_pieces"]
+__all__ = ["Nvfp4ExpertSourceSpec", "decode_csf_scales", "iter_nvfp4_expert_pieces"]

python/freetoken/models/qwen4_exp/config.py

@@ -45,6 +45,8 @@ class Qwen4ExpArgs:
     image_token_id: int | None = None
     # GGUF stores ssm_out with its value-head columns in llama.cpp's tiled order
     gdn_out_heads_tiled: bool = False
+    # PLE n-gram table encoding: native FP8-e4m3 (default) or packed e2m1 NVFP4 (FP4-Local)
+    ple_embedding_dtype: str = "fp8"
 
     @property
     def index_topk_blocks(self) -> int:
@@ -227,6 +229,7 @@ def parse_config(hf_config: Any) -> ModelConfig:
         index_budget=int(text.indexer_budget),
         index_ratio=int(text.indexer_compress_ratio),
         image_token_id=getattr(hf_config, "image_token_id", None),
+        ple_embedding_dtype=str(getattr(text, "ple_embedding_dtype", "fp8")),
     )
 
     return ModelConfig(

python/freetoken/models/qwen4_exp/model.py

@@ -21,7 +21,7 @@
 from freetoken.core import get_global_ctx
 from freetoken.layers import BaseOP, OPList, ParallelLMHead, VocabParallelEmbedding
 from freetoken.models.blocks import BaseLLMModel
-from freetoken.utils import nvtx_annotate
+from freetoken.utils import init_logger, nvtx_annotate
 
 from .attention import Qwen4ExpAttention
 from .hc import GatedResidual
@@ -34,6 +34,16 @@
     from freetoken.core import Batch
     from freetoken.models.config import ModelConfig
 
+logger = init_logger(__name__)
+
+
+def _resolve_ple_backend(requested: str, ple_embedding_dtype: str) -> str:
+    """"disk" streams raw F8 rows through the C++ row store, which cannot unpack packed e2m1 rows; the NVFP4 PLE tables are always pinned."""
+    if requested == "disk" and ple_embedding_dtype == "nvfp4":
+        logger.warning_rank0("PLE table is packed e2m1 (nvfp4); the disk backend streams native rows only -- falling back to the pinned backend")
+        return "pinned"
+    return requested
+
 
 def build_linear_mixer(config: ModelConfig, layer_id: int, prefix: str) -> BaseOP:
     """GDN mixer of a linear_attention layer (Qwen3.5's GDN with a configurable output gate)."""
@@ -170,7 +180,7 @@ def load_host_tables(self, engine_config) -> int:
                 emb.attach_table(ZeroTable(offsets[-1] + sizes[-1], args.ngram_head_dim))
             return 0
 
-        if engine_config.ple_backend == "disk":
+        if _resolve_ple_backend(engine_config.ple_backend, getattr(self._config.qwen4_args, "ple_embedding_dtype", "fp8")) == "disk":
             from .ple_disk import DiskRowTable, resolve_row_source
 
             # one WAIT node per captured graph: the flag protocol supports a single consume
@@ -201,12 +211,13 @@ def load_host_tables(self, engine_config) -> int:
         from .weight import load_ple_table
 
         table = load_ple_table(engine_config.model_path, self._config.qwen4_args)
-        self._ple_table = table  # owns the pinned HostBank; keep it alive
+        self._ple_table = table  # owns the pinned HostBank(s); keep it alive
+        kwargs = {"scale_tensor": table.scale_bank.tensor} if table.scale_bank is not None else {}
         for ple in ple_layers:
             ple.ple_embedding.attach_table(
-                PinnedUVATable(table.bank.tensor, float(table.weight_scale))
+                PinnedUVATable(table.bank.tensor, float(table.weight_scale), **kwargs)
             )
-        return table.bank.nbytes
+        return table.bank.nbytes + (table.scale_bank.nbytes if table.scale_bank is not None else 0)
 
     def forward(self) -> torch.Tensor:
         batch = get_global_ctx().batch