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

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) * 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/loader.py

@@ -51,7 +51,18 @@ def shard_tensor(
 def iter_weight_files(model_path: str) -> list[str]:
     model_folder = download_hf_weight(model_path)
     files = glob.glob(f"{model_folder}/*.safetensors")
-    return [f for f in files if not f.endswith("consolidated.safetensors")] or files
+    files = [f for f in files if not f.endswith("consolidated.safetensors")] or files
+    if files:
+        return files
+    # Root-glob misses checkpoints that keep their shards in a subdirectory (nvidia's FP4-Local
+    # releases ship tensors/model-NNNNN-of-NNNNN.safetensors): the index's weight_map holds the
+    # shard paths relative to the model folder, so resolve them from there.
+    index = os.path.join(model_folder, "model.safetensors.index.json")
+    if not os.path.exists(index):
+        return files
+    with open(index, encoding="utf-8") as f:
+        weight_map = json.load(f)["weight_map"]
+    return [os.path.join(model_folder, shard) for shard in dict.fromkeys(weight_map.values())]
 
 
 def safetensors_weight_map(folder: str) -> dict[str, str]:

python/freetoken/models/nvfp4_banks.py

@@ -27,6 +27,15 @@ 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
+    # Tensor kinds present in the checkpoint but never fetched (per-expert input
+    # scales are runtime quantities, not stored here).
+    ignored_kinds: frozenset[str] = frozenset()
+    # 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:
@@ -63,6 +72,42 @@ 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():
+        # CPU torch has no uint32 shifts; the same bit math in int32 is well defined,
+        # including (rec >> 24) & 0xFF for records with the sign bit set.
+        rec = exceptions.reshape(-1)
+        if rec.dtype == torch.uint8:
+            rec = rec.view(torch.int32)
+        elif rec.dtype != torch.int32:
+            rec = rec.to(torch.int32)
+        scale = scale.to(torch.int32)
+        flat = scale.reshape(-1)
+        flat.scatter_(0, (rec & 0xFFFFFF).to(torch.int64), (rec >> 24) & 0xFF)
+    return scale.to(torch.uint8).view(torch.float8_e4m3fn)
+
+
 def iter_nvfp4_expert_pieces(
     model_path: str,
     config,
@@ -90,6 +135,7 @@ def iter_nvfp4_expert_pieces(
     weight_map = safetensors_weight_map(folder)
 
     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:
@@ -101,10 +147,22 @@ 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 kind in spec.ignored_kinds:
+            continue
+        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))
-    expected = _num_moe_layers(config) * config.num_experts * 9
+        else:
+            wanted[name] = (bank_layer, int(match.group("expert")), role + _kind_suffix(kind))
+    per_expert = 9 if pair_kinds is None else 3 * (2 + len(pair_kinds))
+    expected = _num_moe_layers(config) * config.num_experts * per_expert
     if len(wanted) != expected:
         raise ValueError(f"{spec.desc}: found {len(wanted)} expert tensors, expected {expected}")
 
@@ -132,7 +190,42 @@ 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)
-
-
-__all__ = ["Nvfp4ExpertSourceSpec", "iter_nvfp4_expert_pieces"]
+    if pair_kinds is None:
+        return per_expert_pieces(_parallel() if parallel else _serial(), wanted.get, tensors_per_expert=9)
+    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 9 logical roles, 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)]
+    }
+
+    def _csf_pieces():
+        pending: dict[tuple[int, int], dict[str, torch.Tensor]] = {}
+        for name, tensor in _parallel() if parallel else _serial():
+            layer, expert, slot = wanted[name]
+            piece = pending.setdefault((layer, expert), {})
+            piece[slot] = tensor
+            if len(piece) != len(slots_per_expert):
+                continue
+            del pending[(layer, expert)]
+            grouped: dict[str, torch.Tensor] = {}
+            for role in roles:
+                grouped[role] = piece[role].unsqueeze(0)
+                grouped[f"{role}_global"] = piece[f"{role}_global"].reshape(1, -1)
+                grouped[f"{role}_scale"] = decoder(
+                    *[piece[f"{role}~{kind}"] for kind in pair_kinds]
+                ).unsqueeze(0)
+            yield layer, expert, expert + 1, grouped
+        if pending:
+            raise ValueError(f"{spec.desc}: incomplete expert tensors for (layer, expert) {sorted(pending)[:8]}")
+
+    return _csf_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:
@@ -101,9 +103,11 @@ def ple_slot_states(args: Qwen4ExpArgs) -> Tuple[SlotStateSpec, ...]:
 def _layer_types(text: Any) -> list[str]:
     layer_types = getattr(text, "layer_types", None)
     if layer_types is not None:
-        # HF Qwen4ExpTextConfig rewrites full_attention to indexed_attention in __post_init__.
+        # HF Qwen4ExpTextConfig rewrites full_attention to a non-"full" token in
+        # __post_init__: indexed_attention in transformers 5.18+, qwen_sparse_attention in 5.16.
         return [
-            "full_attention" if t == "indexed_attention" else t for t in layer_types
+            "full_attention" if t in ("indexed_attention", "qwen_sparse_attention") else t
+            for t in layer_types
         ]
     # Fall back to full_attention_interval: every Nth layer (1-indexed) is full.
     interval = int(getattr(text, "full_attention_interval", 4))
@@ -227,6 +231,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.info_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