From 2fc20205d3326213d95f2c06cf0503ea2aeafe08 Mon Sep 17 00:00:00 2001 From: chrisqian Date: Fri, 16 Oct 2026 12:46:10 +0800 Subject: [PATCH 01/14] feat(models): decode CSF slab scales for the FP4-Local NVFP4 experts The official FP4-Local checkpoint stores routed-expert block scales as lossless CSF slab pairs (weight_scale.nvfp4_csf_fixed / _exceptions) instead of the plain E4M3 matrix. Add decode_csf_scales and a scale_pair spec mechanism so the expert reader expands each pair into the 9 plain roles, and pick the CSF spec from the checkpoint's tensor names. --- python/freetoken/models/nvfp4_banks.py | 102 +++++++++++++++++-- python/freetoken/models/qwen4_exp/weight.py | 25 ++++- tests/models/qwen4_exp/test_weight.py | 107 ++++++++++++++++++++ 3 files changed, 225 insertions(+), 9 deletions(-) diff --git a/python/freetoken/models/nvfp4_banks.py b/python/freetoken/models/nvfp4_banks.py index 2dfdaaa35f..d8c1aa8149 100644 --- a/python/freetoken/models/nvfp4_banks.py +++ b/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,40 @@ 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(): + rec = exceptions.reshape(-1) + if rec.dtype == torch.uint8: + rec = rec.view(torch.uint32) + elif rec.dtype != torch.uint32: + rec = rec.to(torch.int32).view(torch.uint32) + scale = scale.to(torch.int32) + flat = scale.reshape(-1) + flat.scatter_(0, (rec & 0xFFFFFF).to(torch.int64), (rec >> 24).to(torch.int32)) + return scale.to(torch.uint8).view(torch.float8_e4m3fn) + + def iter_nvfp4_expert_pieces( model_path: str, config, @@ -90,6 +133,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 +145,19 @@ 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 not in pair_kinds: + raise ValueError(f"{spec.desc}: unknown NVFP4 expert tensor kind {kind!r}") + wanted[name] = (bank_layer, int(match.group("expert")), f"{role}~{kind}") + 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 +185,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"] diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index 691651d4cc..cb2ebbcdbb 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -24,9 +24,10 @@ from freetoken.models.qwen3_vl.weight import rename_vl_prefix from freetoken.models.config import VISION_KEY_PREFIXES -from freetoken.models.loader import drop_page_cache, iter_weight_files +from freetoken.models.loader import drop_page_cache, iter_weight_files, safetensors_weight_map from freetoken.models.nvfp4_banks import ( Nvfp4ExpertSourceSpec, + decode_csf_scales, ) from freetoken.layers.quantization import get_quant_config from freetoken.models.register import get_model_spec @@ -40,7 +41,8 @@ # stacked ``mtp.layers.N.mlp.experts.*`` tensors. _EXPERT_KEY_RE = re.compile( r"^model\.language_model\.layers\.(?P\d+)\.mlp\.experts\.(?P\d+)\." - r"(?Pgate_proj|up_proj|down_proj)\.(?Pweight|weight_scale|weight_scale_2)$" + r"(?Pgate_proj|up_proj|down_proj)\." + r"(?Pweight|weight_scale|weight_scale_2|input_scale|weight_scale\.nvfp4_csf_fixed|weight_scale\.nvfp4_csf_exceptions)$" ) _EXPERT_RE = re.compile(r"\.mlp\.experts\.\d+\.") _NVFP4_SOURCE_SPEC = Nvfp4ExpertSourceSpec( @@ -48,6 +50,20 @@ proj_to_role={"gate_proj": "gate", "up_proj": "up", "down_proj": "down"}, layer_to_bank=lambda layer, config: layer, # every layer is MoE desc="Qwen3.8-Flash-Next NVFP4 experts", + # per-expert input scales are runtime quantities, not stored by the banks + ignored_kinds=frozenset({"input_scale"}), +) +# FP4-Local stores the routed-expert block scales as lossless CSF slab pairs instead of the +# plain E4M3 matrix; decode_csf_scales expands each pair back into the matrix the banks read. +_CSF_KINDS = ("weight_scale.nvfp4_csf_fixed", "weight_scale.nvfp4_csf_exceptions") +_NVFP4_LOCAL_SOURCE_SPEC = Nvfp4ExpertSourceSpec( + key_pattern=_EXPERT_KEY_RE, + proj_to_role={"gate_proj": "gate", "up_proj": "up", "down_proj": "down"}, + layer_to_bank=lambda layer, config: layer, + desc="Qwen3.8-Flash-Next-FP4-Local CSF NVFP4 experts", + scale_pair=_CSF_KINDS, + scale_decoder=decode_csf_scales, + ignored_kinds=frozenset({"input_scale"}), ) # Per-tensor modelopt quant scales; consumed with their ``.weight`` (experts) or unused. _SCALE_SUFFIXES = (".weight_scale", ".weight_scale_2", ".input_scale") @@ -383,6 +399,11 @@ def load_ple_table(model_path: str, qwen4_args, *, pin: bool = True, def nvfp4_expert_spec(model_path: str, config): + """Pick the expert source spec for the checkpoint's scale encoding.""" + folder = download_hf_weight(model_path) + weight_map = safetensors_weight_map(folder) + if any(".nvfp4_csf_fixed" in name for name in weight_map): + return _NVFP4_LOCAL_SOURCE_SPEC return _NVFP4_SOURCE_SPEC diff --git a/tests/models/qwen4_exp/test_weight.py b/tests/models/qwen4_exp/test_weight.py index 1a9a100edb..e4c38cf9d1 100644 --- a/tests/models/qwen4_exp/test_weight.py +++ b/tests/models/qwen4_exp/test_weight.py @@ -16,11 +16,16 @@ from freetoken.distributed import set_tp_info, try_get_tp_info from freetoken.kernel.aot_models import SUPPORTED_MODELS, expert_bank_row_bytes +import freetoken.models.qwen4_exp.weight as qwen4_exp_weight +from freetoken.models.nvfp4_banks import decode_csf_scales, iter_nvfp4_expert_pieces from freetoken.models.qwen4_exp.weight import ( + _NVFP4_LOCAL_SOURCE_SPEC, + _NVFP4_SOURCE_SPEC, _ZERO_CENTERED_NORM_SUFFIXES, _DenseFuser, iter_weights, load_ple_table, + nvfp4_expert_spec, ) from freetoken.models.register import get_model_spec from freetoken.moe.host_banks import HostBank, read_range_into @@ -64,6 +69,108 @@ def _fp8_scale(weight: torch.Tensor) -> torch.Tensor: return torch.rand(weight.shape[0] // BLOCK, weight.shape[1] // BLOCK) + 0.5 +def _csf_target(rows: int, cols: int, seed: int = 0) -> torch.Tensor: + """An E4M3 scale matrix whose rows share a high nibble (so base + 4-bit offset is exact).""" + rng = random.Random(seed) + high = torch.tensor([[rng.randrange(16)] for _ in range(rows)], dtype=torch.uint8) + off = torch.tensor([[rng.randrange(16) for _ in range(cols)] for _ in range(rows)], dtype=torch.uint8) + return high << 4 | off + + +def _csf_encode(target: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """CSF slab pair of ``target``: fixed slabs (per-row high-nibble bases + 4-bit offsets) and a + few absolute exception records. The reference the decoder test compares against.""" + rows, cols = target.shape + assert rows % 16 == 0 and cols % 2 == 0 + row_bytes = cols // 2 + fixed = torch.zeros(rows // 16, 16 * (1 + row_bytes), dtype=torch.uint8) + exceptions: list[int] = [] + for r in range(rows): + slab, row = divmod(r, 16) + base = target[r, 0] & 0xF0 + fixed[slab, row] = base + for j in range(row_bytes): + fixed[slab, 16 + row * row_bytes + j] = (target[r, 2 * j + 1] & 0xF) << 4 | (target[r, 2 * j] & 0xF) + # absolute records: (value << 24) | row-major position + for pos in (3, rows * cols - 1): + value = 0xFF - (target.ravel()[pos] & 0xFF) + exceptions.append((value << 24) | pos) + exc = torch.tensor(exceptions, dtype=torch.int32) + ref = target.reshape(-1).clone() + for rec in exceptions: + ref[rec & 0xFFFFFF] = rec >> 24 + return fixed, exc, ref + + +# --------------------------------------------------------------------------- CSF slab scales + + +def test_csf_decode_roundtrip(): + for rows, cols in ((32, 8), (128, 160), (40 * 16, 40)): + target = _csf_target(rows, cols, seed=rows) + fixed, exc, ref = _csf_encode(target) + out = decode_csf_scales(fixed, exc) + assert out.shape == (rows, cols) and out.dtype is torch.float8_e4m3fn + assert torch.equal(out.view(torch.uint8), ref) + # the same records as a U32 byte stream + assert torch.equal(decode_csf_scales(fixed, exc.view(torch.uint8)).view(torch.uint8), ref) + # with no exceptions the slab decode is the target itself (rows share a high nibble) + no_exc = decode_csf_scales(fixed, torch.zeros(0, dtype=torch.int32)) + assert torch.equal(no_exc.view(torch.uint8), target.reshape(-1)) + + +def test_nvfp4_expert_spec_selects_the_csf_checkpoint(monkeypatch): + monkeypatch.setattr(qwen4_exp_weight, "download_hf_weight", lambda path: "") + weight_map = { + "model.language_model.layers.0.mlp.experts.0.gate_proj.weight": "m.safetensors", + "model.language_model.layers.0.mlp.experts.0.gate_proj.weight_scale.nvfp4_csf_fixed": "m.safetensors", + } + monkeypatch.setattr(qwen4_exp_weight, "safetensors_weight_map", lambda folder: weight_map) + spec = nvfp4_expert_spec("unused", SimpleNamespace()) + assert spec is _NVFP4_LOCAL_SOURCE_SPEC + assert spec.scale_pair == ("weight_scale.nvfp4_csf_fixed", "weight_scale.nvfp4_csf_exceptions") + plain = {k.replace(".weight_scale.nvfp4_csf_fixed", ".weight_scale"): v for k, v in weight_map.items()} + monkeypatch.setattr(qwen4_exp_weight, "safetensors_weight_map", lambda folder: plain) + assert nvfp4_expert_spec("unused", SimpleNamespace()) is _NVFP4_SOURCE_SPEC + assert "input_scale" in _NVFP4_SOURCE_SPEC.ignored_kinds + + +def test_csf_expert_reader_expands_the_scale_pieces(tmp_path): + """A CSF checkpoint streams slab pairs; the reader expands them into the 9 plain roles.""" + from freetoken.models.loader import drop_page_cache + + lm = "model.language_model" + H_, I_ = 128, 64 # hidden / expert width; K multiples of 16, rows multiples of 16 + raw: dict[str, torch.Tensor] = {} + refs: dict[tuple[int, str], torch.Tensor] = {} + for expert in range(2): + for proj, rows, k in (("gate_proj", I_, H_), ("up_proj", I_, H_), ("down_proj", H_, I_)): + base = f"{lm}.layers.0.mlp.experts.{expert}.{proj}" + raw[f"{base}.weight"] = torch.randint(0, 256, (rows, k // 2), dtype=torch.uint8) + raw[f"{base}.weight_scale_2"] = torch.rand(()) + 0.5 + raw[f"{base}.input_scale"] = torch.rand(()) + 0.5 # ignored by the banks + fixed, exc, ref = _csf_encode(_csf_target(rows, k // 16, seed=expert + k)) + raw[f"{base}.weight_scale.nvfp4_csf_fixed"] = fixed + raw[f"{base}.weight_scale.nvfp4_csf_exceptions"] = exc + refs[(expert, proj)] = ref + shard = "model-00001-of-00001.safetensors" + save_file(raw, tmp_path / shard) + (tmp_path / "model.safetensors.index.json").write_text(json.dumps({"metadata": None, "weight_map": {n: shard for n in raw}})) + config = SimpleNamespace(num_experts=2, num_layers=1, first_k_dense_replace=0) + pieces = list(iter_nvfp4_expert_pieces(str(tmp_path), config, _NVFP4_LOCAL_SOURCE_SPEC, drop_page_cache=drop_page_cache)) + assert sorted(p[1] for p in pieces) == [0, 1] + for layer, e0, e1, piece in pieces: + assert (layer, e0, e1) == (0, e0, e0 + 1) + for proj, role in (("gate_proj", "gate"), ("up_proj", "up"), ("down_proj", "down")): + base = f"{lm}.layers.0.mlp.experts.{e0}.{proj}" + assert torch.equal(piece[role][0], raw[f"{base}.weight"]) + assert piece[f"{role}_global"].dtype is torch.float16 and torch.equal( + piece[f"{role}_global"].reshape(-1), raw[f"{base}.weight_scale_2"].reshape(-1).to(torch.float16) + ) + assert piece[f"{role}_scale"].shape == (1, raw[f"{base}.weight"].shape[0], raw[f"{base}.weight"].shape[1] // 8) + assert torch.equal(piece[f"{role}_scale"].view(torch.uint8)[0].reshape(-1), refs[(e0, proj)]) + + def _raw_checkpoint(dense_fp8: bool = False) -> dict[str, torch.Tensor]: """Layer 0 = GDN + PLE, layer 1 = QSA; plus the mtp / visual / routed-expert noise. From 2db65bd7662d4e6d218ee61ab137ad330507e8df Mon Sep 17 00:00:00 2001 From: chrisqian Date: Fri, 16 Oct 2026 13:22:21 +0800 Subject: [PATCH 02/14] feat(qwen4_exp): load the FP4-Local PLE table as packed e2m1 --- python/freetoken/kernel/triton/ple.py | 48 +++++++++- python/freetoken/models/qwen4_exp/config.py | 3 + python/freetoken/models/qwen4_exp/model.py | 7 +- python/freetoken/models/qwen4_exp/ple.py | 68 +++++++++++--- python/freetoken/models/qwen4_exp/ple_disk.py | 5 + python/freetoken/models/qwen4_exp/weight.py | 94 +++++++++++++++---- tests/models/qwen4_exp/test_ple.py | 48 ++++++++++ tests/models/qwen4_exp/test_ple_disk.py | 16 ++++ tests/models/qwen4_exp/test_weight.py | 64 +++++++++++++ 9 files changed, 318 insertions(+), 35 deletions(-) diff --git a/python/freetoken/kernel/triton/ple.py b/python/freetoken/kernel/triton/ple.py index 2314e94e72..4f9861dff6 100644 --- a/python/freetoken/kernel/triton/ple.py +++ b/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,31 @@ _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))) + return tl.where((n >> 3) & 1, -mag, mag) + + @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 +67,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 +104,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 diff --git a/python/freetoken/models/qwen4_exp/config.py b/python/freetoken/models/qwen4_exp/config.py index 2bc96f67a1..45d0cfe95e 100644 --- a/python/freetoken/models/qwen4_exp/config.py +++ b/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( diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py index 4eb69165d5..d44456b568 100644 --- a/python/freetoken/models/qwen4_exp/model.py +++ b/python/freetoken/models/qwen4_exp/model.py @@ -201,12 +201,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 diff --git a/python/freetoken/models/qwen4_exp/ple.py b/python/freetoken/models/qwen4_exp/ple.py index c944c4946f..f7890c90b5 100644 --- a/python/freetoken/models/qwen4_exp/ple.py +++ b/python/freetoken/models/qwen4_exp/ple.py @@ -70,22 +70,54 @@ def prefetch(self, row_ids: torch.Tensor) -> None: ... class GpuResidentTable: - """PLE table held whole in GPU memory; ``index_select`` oracle for the pinned-host backend.""" + """PLE table held whole in GPU memory; ``index_select`` oracle for the pinned-host backend. + + Native fp8/bf16 tables read one row per value and apply the scalar ``scale``. Packed e2m1 + tables (FP4-Local) pass the per-16 e4m3 ``scale_tensor`` rows and the global ``scale``. + """ def __init__( - self, weight: torch.Tensor, scale: float = 1.0, dtype: torch.dtype | None = None + self, + weight: torch.Tensor, + scale: float = 1.0, + dtype: torch.dtype | None = None, + scale_tensor: torch.Tensor | None = None, ) -> None: self.weight = weight self.scale = float(scale) - self.num_rows, self.head_dim = weight.shape + self.scale_tensor = scale_tensor + if scale_tensor is None: + self.num_rows, self.head_dim = weight.shape + else: + assert weight.dtype == torch.uint8, "packed PLE tables hold e2m1, two per byte" + assert scale_tensor.shape[0] == weight.shape[0] and weight.shape[1] % 8 == 0 + self.num_rows = weight.shape[0] + self.head_dim = weight.shape[1] * 2 + # e2m1: 0, +/-0.5, +/-1, +/-1.5, +/-2, +/-3, +/-4, +/-6 (sign = bit 3) + self._lut = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, 0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], + dtype=torch.float32, + device=weight.device, + ) self.dtype = dtype if dtype is not None else ( torch.bfloat16 if weight.dtype.itemsize < 2 else weight.dtype ) def lookup(self, row_ids: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor: - rows = self.weight.index_select(0, row_ids.reshape(-1)).to(self.dtype) - if self.scale != 1.0: - rows = rows * self.scale + if self.scale_tensor is None: + rows = self.weight.index_select(0, row_ids.reshape(-1)).to(self.dtype) + if self.scale != 1.0: + rows = rows * self.scale + else: + flat = row_ids.reshape(-1) + packed = self.weight.index_select(0, flat) + scales = self.scale_tensor.index_select(0, flat).view(torch.float8_e4m3fn).float() + lo = self._lut[packed & 0x0F] + hi = self._lut[(packed >> 4) & 0x0F] + rows = torch.stack((lo, hi), dim=2).reshape(flat.shape[0], self.head_dim) + if self.scale != 1.0: + scales = scales * self.scale + rows = (rows * scales.repeat_interleave(16, dim=1)).to(self.dtype) rows = rows.view(*row_ids.shape[:-1], -1) if out is None: return rows @@ -121,10 +153,12 @@ class PinnedUVATable: """PLE table left in pinned host memory; rows are gathered over UVA by a Triton kernel. ``weight`` must be the filled and ``pin()``ed ``HostBank.tensor`` from - ``weight.load_ple_table`` (``[num_rows, head_dim]``, fp8-e4m3 or bf16); an unregistered host - buffer is not device-addressable and the kernel faults on it. ``scale`` is the checkpoint's - scalar ``weight_scale``. Gathers emit bf16 into a staging buffer, one per captured decode size - and one growable buffer for everything else. + ``weight.load_ple_table`` (``[num_rows, head_dim]`` fp8-e4m3 or bf16, or ``[num_rows, + head_dim // 2]`` packed e2m1 with the per-16 e4m3 ``scale_tensor`` rows); an unregistered + host buffer is not device-addressable and the kernel faults on it. ``scale`` is the + checkpoint's scalar scale (``weight_scale``, or ``weight_scale_2`` for packed tables). + Gathers emit bf16 into a staging buffer, one per captured decode size and one growable buffer + for everything else. ``prefetch`` runs the gather on a private stream and the next ``lookup`` joins it. ``lookup`` returns a view of that staging buffer, so the rows must be consumed before the next lookup. @@ -137,19 +171,26 @@ def __init__( *, device: torch.device | None = None, prefetch: bool = True, + scale_tensor: torch.Tensor | None = None, ) -> None: assert weight.device.type == "cpu" and weight.is_contiguous() - assert weight.dtype in (torch.float8_e4m3fn, torch.bfloat16), weight.dtype + assert weight.dtype in (torch.float8_e4m3fn, torch.bfloat16, torch.uint8), weight.dtype from freetoken.kernel.pinned import device_ptr self.weight = weight self.scale = float(scale) - self.num_rows, self.head_dim = weight.shape + self.scale_tensor = scale_tensor + self.num_rows = weight.shape[0] + self.head_dim = weight.shape[1] * 2 if weight.dtype == torch.uint8 else weight.shape[1] + if weight.dtype == torch.uint8: + assert scale_tensor is not None, "packed e2m1 tables need their per-16 scale rows" + assert scale_tensor.dtype == torch.float8_e4m3fn self.dtype = torch.bfloat16 self._is_fp8 = weight.dtype == torch.float8_e4m3fn self._device = device or torch.device("cuda", torch.cuda.current_device()) # WDDM maps registered host memory at a different device address; on Linux/UVA this is data_ptr self._table_ptr = device_ptr(weight) + self._scale_ptr = device_ptr(scale_tensor) if scale_tensor is not None else None self._stream = torch.cuda.Stream(device=self._device) if prefetch else None self._staging: torch.Tensor | None = None self._graph_staging: dict[int, torch.Tensor] = {} @@ -180,7 +221,8 @@ def _gather(self, row_ids: torch.Tensor, dst: torch.Tensor) -> torch.Tensor: row_ids.reshape(-1), dst, self.scale, - self._is_fp8, + self._is_fp8 and self._scale_ptr is None, + scale_ptr=self._scale_ptr, ) def prefetch(self, row_ids: torch.Tensor) -> None: diff --git a/python/freetoken/models/qwen4_exp/ple_disk.py b/python/freetoken/models/qwen4_exp/ple_disk.py index bb71da2ff7..f136794932 100644 --- a/python/freetoken/models/qwen4_exp/ple_disk.py +++ b/python/freetoken/models/qwen4_exp/ple_disk.py @@ -78,6 +78,11 @@ def source_from_safetensors(folder: str) -> PleRowSource: match = _PLE_SHARD_RE.search(key) if match is None: continue + if meta["dtype"] == "U8": + raise NotImplementedError( + f"the disk PLE backend reads native { _PLE_ST_DTYPE } rows; " + f"shard {key} holds packed e2m1 (NVFP4 table) -- use the default pinned backend" + ) if meta["dtype"] != _PLE_ST_DTYPE: raise ValueError(f"PLE shard {key} has dtype {meta['dtype']}, expected {_PLE_ST_DTYPE}") if rows and tuple(meta["shape"]) != (rows, cols): diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index cb2ebbcdbb..3383489d34 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -74,6 +74,12 @@ r"\.ple\.ple_embedding\.ngram_embedding\.shard_(?P\d+)\.weight$" ) _PLE_SCALE_SUFFIX = ".ple.ple_embedding.ngram_embedding.weight_scale" +# FP4-Local tables pack the rows as e2m1 (two per byte), keep a per-16 e4m3 scale row per shard +# and put the global scale in ``weight_scale_2``. +_PLE_SHARD_SCALE_RE = re.compile( + r"\.ple\.ple_embedding\.ngram_embedding\.shard_(?P\d+)\.weight_scale$" +) +_PLE_GLOBAL_SCALE_SUFFIX = ".ple.ple_embedding.ngram_embedding.weight_scale_2" _PLE_FILE_BYTES = 4 << 30 # ple-table-*.safetensors written by ftw_side_files # Zero-centered Qwen4ExpTextRMSNorm weights, loaded RAW: GroupedPlusOneRMSNorm / GemmaPlusOneRMSNorm @@ -267,14 +273,20 @@ def iter_vision_weights(model_path: str, device: torch.device) -> Iterator[tuple @dataclass(frozen=True) class PleTable: - """The filled n-gram table: one pinned host bank plus the checkpoint's per-tensor FP8 scale.""" + """The filled n-gram table: one pinned host bank plus the checkpoint's per-tensor scale. + + FP8 tables hold e4m3 values and a scalar ``weight_scale``; NVFP4 tables hold packed e2m1 + values in ``bank`` and the per-16 e4m3 scale rows in ``scale_bank``, with ``weight_scale`` + reading the checkpoint's global ``weight_scale_2``. + """ bank: HostBank - weight_scale: torch.Tensor # scalar, checkpoint dtype (bf16) + weight_scale: torch.Tensor # scalar, checkpoint dtype (bf16 or fp32) + scale_bank: HostBank | None = None # [total_rows, ngram_head_dim // 16] e4m3, NVFP4 tables @property def tensor(self) -> torch.Tensor: - """``[total_rows, ngram_head_dim]`` float8_e4m3fn view of the bank.""" + """``[total_rows, ngram_head_dim // (2 if packed else 1)]`` view of the value bank.""" return self.bank.tensor @@ -340,42 +352,70 @@ def load_ple_table(model_path: str, qwen4_args, *, pin: bool = True, bank is filled shard by shard at ``shard_index * rows_per_shard``. Each read is O_DIRECT: the table is ~47.7 GiB and must not also sit in the page cache while the bank holds the same bytes. """ + encoding = getattr(qwen4_args, "ple_embedding_dtype", "fp8") + if encoding not in ("fp8", "nvfp4"): + raise ValueError(f"unsupported ple_embedding_dtype {encoding!r}") folder = download_hf_weight(model_path) parts: dict[int, tuple[str, int, int]] = {} # shard index -> (path, file offset, bytes) + scale_parts: dict[int, tuple[str, int, int]] = {} scale: torch.Tensor | None = None rows = cols = 0 + st_dtype = _PLE_ST_DTYPE if encoding == "fp8" else "U8" for path in _ple_table_files(folder): header, base = _safetensors_header(path) for key, meta in header.items(): if key == "__metadata__": continue - if key.endswith(_PLE_SCALE_SUFFIX): + if key.endswith(_PLE_GLOBAL_SCALE_SUFFIX if encoding == "nvfp4" else _PLE_SCALE_SUFFIX): with safetensors.safe_open(path, framework="pt", device="cpu") as f: scale = f.get_tensor(key).reshape(()) continue match = _PLE_SHARD_RE.search(key) - if match is None: + if match is not None: + if meta["dtype"] != st_dtype: + raise ValueError( + f"PLE table shard {key} has dtype {meta['dtype']}, " + f"expected {st_dtype} for ple_embedding_dtype={encoding}" + ) + shape = meta["shape"] + if rows and tuple(shape) != (rows, cols): + raise ValueError(f"PLE table shard {key} is {shape}, expected {[rows, cols]}") + rows, cols = shape + begin, end = meta["data_offsets"] + parts[int(match.group("shard"))] = (path, base + begin, end - begin) continue - if meta["dtype"] != _PLE_ST_DTYPE: - raise ValueError(f"PLE table shard {key} has unsupported dtype {meta['dtype']}") - shape = meta["shape"] - if rows and tuple(shape) != (rows, cols): - raise ValueError(f"PLE table shard {key} is {shape}, expected {[rows, cols]}") - rows, cols = shape - begin, end = meta["data_offsets"] - parts[int(match.group("shard"))] = (path, base + begin, end - begin) + if encoding == "nvfp4": + match = _PLE_SHARD_SCALE_RE.search(key) + if match is not None: + if meta["dtype"] != _PLE_ST_DTYPE: + raise ValueError(f"PLE scale shard {key} has dtype {meta['dtype']}, expected {_PLE_ST_DTYPE}") + shape = meta["shape"] + if rows and tuple(shape) != (rows, cols // 8): + raise ValueError(f"PLE scale shard {key} is {shape}, expected {[rows, cols // 8]}") + begin, end = meta["data_offsets"] + scale_parts[int(match.group("shard"))] = (path, base + begin, end - begin) expected = int(qwen4_args.split_ngram_parts) if sorted(parts) != list(range(expected)): raise ValueError( f"PLE table needs shards 0..{expected - 1}, found {len(parts)}: {sorted(parts)[:8]}" ) - if cols != qwen4_args.ngram_head_dim: - raise ValueError(f"PLE table row is {cols} wide, config says {qwen4_args.ngram_head_dim}") + head_dim = qwen4_args.ngram_head_dim + if encoding == "nvfp4": + if cols != head_dim // 2: + raise ValueError( + f"PLE table row is {cols} packed bytes, config says {head_dim // 2} for a {head_dim}-dim head" + ) + if sorted(scale_parts) != list(range(expected)): + raise ValueError( + f"PLE table needs scale shards 0..{expected - 1}, found {len(scale_parts)}: {sorted(scale_parts)[:8]}" + ) + elif cols != head_dim: + raise ValueError(f"PLE table row is {cols} wide, config says {head_dim}") if scale is None: - raise ValueError("PLE table has no weight_scale") + raise ValueError(f"PLE table has no {'weight_scale_2' if encoding == 'nvfp4' else 'weight_scale'}") - bank = HostBank((expected * rows, cols), torch.float8_e4m3fn) + bank = HostBank((expected * rows, cols), torch.uint8 if encoding == "nvfp4" else torch.float8_e4m3fn) shard_bytes = rows * cols bar = byte_bar(expected * shard_bytes, "Loading PLE table") try: @@ -388,9 +428,27 @@ def load_ple_table(model_path: str, qwen4_args, *, pin: bool = True, bar.update(nbytes) finally: bar.close() + scale_bank = None + if encoding == "nvfp4": + scale_cols = cols // 8 + scale_bank = HostBank((expected * rows, scale_cols), torch.float8_e4m3fn) + scale_shard_bytes = rows * scale_cols + bar = byte_bar(expected * scale_shard_bytes, "Loading PLE table scales") + try: + sbuf = scale_bank.memoryview() + for shard in range(expected): + path, offset, nbytes = scale_parts[shard] + assert nbytes == scale_shard_bytes, f"PLE scale shard {shard} is {nbytes} B, expected {scale_shard_bytes}" + read_range_into(sbuf, path, file_offset=offset, nbytes=nbytes, + dest_offset=shard * scale_shard_bytes, workers=workers, chunk=chunk) + bar.update(nbytes) + finally: + bar.close() + if pin and torch.cuda.is_available(): + scale_bank.pin() if pin and torch.cuda.is_available(): bank.pin() - return PleTable(bank=bank, weight_scale=scale) + return PleTable(bank=bank, weight_scale=scale, scale_bank=scale_bank) # ====================================================================================== diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index 0d2461427e..70db4cbdea 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -224,6 +224,54 @@ def test_pinned_uva_zeroes_out_of_range_ids(): assert torch.equal(rows[0], bank.tensor[0].cuda().to(torch.bfloat16)) +@requires_cuda +def test_pinned_uva_fp4_matches_gpu_resident(): + """The FP4-Local table: packed e2m1 rows + per-16 e4m3 scale rows + global scale, UVA against the oracle.""" + from freetoken.moe.host_banks import HostBank + + gen = torch.Generator().manual_seed(9) + rows, dim, scale = 4096, 160, 0.03125 + value = HostBank((rows, dim // 2), torch.uint8) + value.tensor.copy_(torch.randint(0, 256, (rows, dim // 2), generator=gen)) + value.pin() + scale_bank = HostBank((rows, dim // 16), torch.float8_e4m3fn) + scale_bank.tensor.copy_((torch.rand(rows, dim // 16, generator=gen) * 8).to(torch.float8_e4m3fn)) + scale_bank.pin() + + oracle = GpuResidentTable( + value.tensor.cuda(), scale, dtype=torch.bfloat16, scale_tensor=scale_bank.tensor.cuda() + ) + pinned = PinnedUVATable(value.tensor, scale, scale_tensor=scale_bank.tensor) + + ids = torch.randint(0, rows, (37, 16), device="cuda") + want = oracle.lookup(ids) + assert torch.equal(pinned.lookup(ids), want) + + pinned.prefetch(ids) + assert torch.equal(pinned.lookup(ids), want) + + other = torch.randint(0, rows, (37, 16), device="cuda") + assert torch.equal(pinned.lookup(other), oracle.lookup(other)) + assert pinned.lookup(ids).shape == (37, 16, dim) + + +def test_gpu_resident_fp4_lookup_matches_the_lut(): + """CPU oracle: e2m1 nibble decode, per-16 e4m3 scale rows, global scale.""" + gen = torch.Generator().manual_seed(3) + rows, dim, scale = 64, 32, 0.5 + packed = torch.randint(0, 256, (rows, dim // 2), generator=gen) + scales = (torch.rand(rows, dim // 16, generator=gen) * 8).to(torch.float8_e4m3fn) + table = GpuResidentTable(packed, scale, dtype=torch.bfloat16, scale_tensor=scales) + ids = torch.randint(0, rows, (5, 16)) + + lut = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, 0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0] + ) + values = torch.stack((lut[packed & 0x0F], lut[(packed >> 4) & 0x0F]), dim=2).reshape(rows, dim) + want = (values * scales.float().repeat_interleave(16, dim=1) * scale).to(torch.bfloat16) + assert torch.equal(table.lookup(ids).view(-1, dim), want[ids.reshape(-1)]) + + @pytest.mark.skipif( not os.environ.get("FREETOKEN_QWEN4EXP_MODEL"), reason="needs FREETOKEN_QWEN4EXP_MODEL" ) diff --git a/tests/models/qwen4_exp/test_ple_disk.py b/tests/models/qwen4_exp/test_ple_disk.py index ba9cd98c63..b0eebb8611 100644 --- a/tests/models/qwen4_exp/test_ple_disk.py +++ b/tests/models/qwen4_exp/test_ple_disk.py @@ -79,6 +79,22 @@ def _write_checkpoint(tmp_path, table, n_shards): save_file(tensors, str(tmp_path / "model.safetensors")) +def test_source_rejects_a_packed_nvfp4_table(tmp_path): + """The disk store reads native FP8 rows; the packed e2m1 table needs the pinned backend.""" + save_file( + { + f"{_KEY_PREFIX}.shard_0.weight": torch.zeros(4, 8, dtype=torch.uint8), + f"{_KEY_PREFIX}.shard_0.weight_scale": torch.zeros(4, 1, dtype=torch.float8_e4m3fn), + f"{_KEY_PREFIX}.weight_scale_2": torch.tensor([1.0]), + }, + str(tmp_path / "model-plefp4-00000.safetensors"), + ) + from freetoken.models.qwen4_exp.ple_disk import source_from_safetensors + + with pytest.raises(NotImplementedError, match="packed e2m1"): + source_from_safetensors(str(tmp_path)) + + def _make_table(tmp_path): from freetoken.models.qwen4_exp.ple_disk import DiskRowTable, source_from_safetensors diff --git a/tests/models/qwen4_exp/test_weight.py b/tests/models/qwen4_exp/test_weight.py index e4c38cf9d1..1de581faaf 100644 --- a/tests/models/qwen4_exp/test_weight.py +++ b/tests/models/qwen4_exp/test_weight.py @@ -462,6 +462,70 @@ def test_load_ple_table_rejects_a_shard_count_mismatch(checkpoint): load_ple_table(folder, args, pin=False) +def _nvfp4_ngram_table() -> tuple[dict[str, torch.Tensor], torch.Tensor]: + """The FP4-Local table: packed e2m1 values (u8, two per byte) with per-16 e4m3 scale rows. + + head_dim is 16 so a packed row is 8 bytes with a single per-16 scale. + """ + prefix = "model.language_model.layers.0.ple.ple_embedding.ngram_embedding" + rows, parts, dim = 5, 3, 16 + tensors = { + f"{prefix}.shard_{i}.weight": ( + torch.arange(i * rows * dim // 2, (i + 1) * rows * dim // 2) + .remainder(255).to(torch.uint8).view(rows, dim // 2) + ) + for i in range(parts) + } + tensors.update( + { + f"{prefix}.shard_{i}.weight_scale": ( + torch.arange(i * rows * dim // 16, (i + 1) * rows * dim // 16) + .remainder(127).to(torch.uint8).view(rows, dim // 16).view(torch.float8_e4m3fn) + ) + for i in range(parts) + } + ) + scale = torch.tensor([0.25], dtype=torch.float32) + tensors[f"{prefix}.weight_scale_2"] = scale + return tensors, scale, rows, dim, parts + + +def test_load_ple_table_nvfp4_concatenates_packed_shards(tmp_path): + """The FP4-Local table: packed e2m1 value rows, per-16 e4m3 scale rows, global weight_scale_2.""" + tensors, scale, rows, dim, parts = _nvfp4_ngram_table() + names = sorted(tensors) + save_file({n: tensors[n] for n in names[:2]}, str(tmp_path / "model-plefp4-00000.safetensors")) + save_file({n: tensors[n] for n in names[2:]}, str(tmp_path / "model-plefp4-00001.safetensors")) + args = SimpleNamespace(split_ngram_parts=parts, ngram_head_dim=dim, ple_embedding_dtype="nvfp4") + table = load_ple_table(str(tmp_path), args, pin=False) + assert table.tensor.shape == (parts * rows, dim // 2) + assert table.tensor.dtype is torch.uint8 + assert table.scale_bank is not None + assert table.scale_bank.tensor.shape == (parts * rows, dim // 16) + assert table.scale_bank.tensor.dtype is torch.float8_e4m3fn + assert table.weight_scale.dtype is torch.float32 + assert float(table.weight_scale) == 0.25 + prefix = "model.language_model.layers.0.ple.ple_embedding.ngram_embedding" + for shard in range(parts): + assert torch.equal( + table.tensor[shard * rows : (shard + 1) * rows], tensors[f"{prefix}.shard_{shard}.weight"] + ) + assert torch.equal( + table.scale_bank.tensor[shard * rows : (shard + 1) * rows], + tensors[f"{prefix}.shard_{shard}.weight_scale"], + ) + + +def test_load_ple_table_rejects_a_nvfp4_table_under_fp8(tmp_path): + """A packed table's U8 value shards must not load as an FP8 table.""" + tensors, _scale, rows, dim, parts = _nvfp4_ngram_table() + for i, name in enumerate(sorted(tensors)): + save_file({name: tensors[name]}, str(tmp_path / f"model-plefp4-{i:05d}.safetensors")) + args = SimpleNamespace(split_ngram_parts=parts, ngram_head_dim=dim, ple_embedding_dtype="fp8") + with pytest.raises(ValueError, match="unsupported dtype"): + load_ple_table(str(tmp_path), args, pin=False) + + # ====================================================================================== # read_range_into: the O_DIRECT byte-range read the PLE table load is built on # ====================================================================================== From c624c02c5902b06a0438517c519c064f6976c2fc Mon Sep 17 00:00:00 2001 From: chrisqian Date: Fri, 16 Oct 2026 13:42:20 +0800 Subject: [PATCH 03/14] test(qwen4_exp): import save_file locally in the NVFP4 disk rejection test --- tests/models/qwen4_exp/test_ple_disk.py | 2 ++ tests/models/qwen4_exp/test_weight.py | 6 +++--- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/models/qwen4_exp/test_ple_disk.py b/tests/models/qwen4_exp/test_ple_disk.py index b0eebb8611..be310a8397 100644 --- a/tests/models/qwen4_exp/test_ple_disk.py +++ b/tests/models/qwen4_exp/test_ple_disk.py @@ -81,6 +81,8 @@ def _write_checkpoint(tmp_path, table, n_shards): def test_source_rejects_a_packed_nvfp4_table(tmp_path): """The disk store reads native FP8 rows; the packed e2m1 table needs the pinned backend.""" + from safetensors.torch import save_file + save_file( { f"{_KEY_PREFIX}.shard_0.weight": torch.zeros(4, 8, dtype=torch.uint8), diff --git a/tests/models/qwen4_exp/test_weight.py b/tests/models/qwen4_exp/test_weight.py index 1de581faaf..49124f4d41 100644 --- a/tests/models/qwen4_exp/test_weight.py +++ b/tests/models/qwen4_exp/test_weight.py @@ -462,7 +462,7 @@ def test_load_ple_table_rejects_a_shard_count_mismatch(checkpoint): load_ple_table(folder, args, pin=False) -def _nvfp4_ngram_table() -> tuple[dict[str, torch.Tensor], torch.Tensor]: +def _nvfp4_ngram_table() -> tuple[dict[str, torch.Tensor], torch.Tensor, int, int, int]: """The FP4-Local table: packed e2m1 values (u8, two per byte) with per-16 e4m3 scale rows. head_dim is 16 so a packed row is 8 bytes with a single per-16 scale. @@ -492,7 +492,7 @@ def _nvfp4_ngram_table() -> tuple[dict[str, torch.Tensor], torch.Tensor]: def test_load_ple_table_nvfp4_concatenates_packed_shards(tmp_path): """The FP4-Local table: packed e2m1 value rows, per-16 e4m3 scale rows, global weight_scale_2.""" - tensors, scale, rows, dim, parts = _nvfp4_ngram_table() + tensors, _scale, rows, dim, parts = _nvfp4_ngram_table() names = sorted(tensors) save_file({n: tensors[n] for n in names[:2]}, str(tmp_path / "model-plefp4-00000.safetensors")) save_file({n: tensors[n] for n in names[2:]}, str(tmp_path / "model-plefp4-00001.safetensors")) @@ -518,7 +518,7 @@ def test_load_ple_table_nvfp4_concatenates_packed_shards(tmp_path): def test_load_ple_table_rejects_a_nvfp4_table_under_fp8(tmp_path): """A packed table's U8 value shards must not load as an FP8 table.""" - tensors, _scale, rows, dim, parts = _nvfp4_ngram_table() + tensors, _scale, _rows, dim, parts = _nvfp4_ngram_table() for i, name in enumerate(sorted(tensors)): save_file({name: tensors[name]}, str(tmp_path / f"model-plefp4-{i:05d}.safetensors")) args = SimpleNamespace(split_ngram_parts=parts, ngram_head_dim=dim, ple_embedding_dtype="fp8") From b708c45af58d567e4caa2d35bedb6b9422900a10 Mon Sep 17 00:00:00 2001 From: chrisqian Date: Fri, 16 Oct 2026 14:03:54 +0800 Subject: [PATCH 04/14] fix(models): int32 CSF exception math and accept plain kinds in pair mode --- python/freetoken/models/nvfp4_banks.py | 17 +++++++++++------ tests/models/qwen4_exp/test_ple.py | 2 +- tests/models/qwen4_exp/test_weight.py | 2 +- 3 files changed, 13 insertions(+), 8 deletions(-) diff --git a/python/freetoken/models/nvfp4_banks.py b/python/freetoken/models/nvfp4_banks.py index d8c1aa8149..dd8a4cdd65 100644 --- a/python/freetoken/models/nvfp4_banks.py +++ b/python/freetoken/models/nvfp4_banks.py @@ -95,14 +95,16 @@ def decode_csf_scales(fixed: torch.Tensor, exceptions: torch.Tensor) -> torch.Te 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.uint32) - elif rec.dtype != torch.uint32: - rec = rec.to(torch.int32).view(torch.uint32) + 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).to(torch.int32)) + flat.scatter_(0, (rec & 0xFFFFFF).to(torch.int64), (rec >> 24) & 0xFF) return scale.to(torch.uint8).view(torch.float8_e4m3fn) @@ -149,9 +151,12 @@ def iter_nvfp4_expert_pieces( if kind in spec.ignored_kinds: continue if pair_kinds is not None: - if kind not in pair_kinds: + 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}") - wanted[name] = (bank_layer, int(match.group("expert")), f"{role}~{kind}") elif kind not in ("weight", "weight_scale", "weight_scale_2"): raise ValueError(f"{spec.desc}: unknown NVFP4 expert tensor kind {kind!r}") else: diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index 70db4cbdea..4fd6fb4879 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -259,7 +259,7 @@ def test_gpu_resident_fp4_lookup_matches_the_lut(): """CPU oracle: e2m1 nibble decode, per-16 e4m3 scale rows, global scale.""" gen = torch.Generator().manual_seed(3) rows, dim, scale = 64, 32, 0.5 - packed = torch.randint(0, 256, (rows, dim // 2), generator=gen) + packed = torch.randint(0, 256, (rows, dim // 2), generator=gen).to(torch.uint8) scales = (torch.rand(rows, dim // 16, generator=gen) * 8).to(torch.float8_e4m3fn) table = GpuResidentTable(packed, scale, dtype=torch.bfloat16, scale_tensor=scales) ids = torch.randint(0, rows, (5, 16)) diff --git a/tests/models/qwen4_exp/test_weight.py b/tests/models/qwen4_exp/test_weight.py index 49124f4d41..83f9039079 100644 --- a/tests/models/qwen4_exp/test_weight.py +++ b/tests/models/qwen4_exp/test_weight.py @@ -522,7 +522,7 @@ def test_load_ple_table_rejects_a_nvfp4_table_under_fp8(tmp_path): for i, name in enumerate(sorted(tensors)): save_file({name: tensors[name]}, str(tmp_path / f"model-plefp4-{i:05d}.safetensors")) args = SimpleNamespace(split_ngram_parts=parts, ngram_head_dim=dim, ple_embedding_dtype="fp8") - with pytest.raises(ValueError, match="unsupported dtype"): + with pytest.raises(ValueError, match="expected F8_E4M3"): load_ple_table(str(tmp_path), args, pin=False) From c5e01203261825c82ca9f5a43b9d399db0058585 Mon Sep 17 00:00:00 2001 From: chrisqian Date: Fri, 16 Oct 2026 14:30:03 +0800 Subject: [PATCH 05/14] test(qwen4_exp): build CSF test records as Python ints and compare slab shapes --- tests/models/qwen4_exp/test_weight.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/models/qwen4_exp/test_weight.py b/tests/models/qwen4_exp/test_weight.py index 83f9039079..844eee47b6 100644 --- a/tests/models/qwen4_exp/test_weight.py +++ b/tests/models/qwen4_exp/test_weight.py @@ -91,14 +91,14 @@ def _csf_encode(target: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch fixed[slab, row] = base for j in range(row_bytes): fixed[slab, 16 + row * row_bytes + j] = (target[r, 2 * j + 1] & 0xF) << 4 | (target[r, 2 * j] & 0xF) - # absolute records: (value << 24) | row-major position + # absolute records: (value << 24) | row-major position, as Python ints (uint8 would wrap) for pos in (3, rows * cols - 1): - value = 0xFF - (target.ravel()[pos] & 0xFF) + value = 0xFF - target.ravel()[pos].item() exceptions.append((value << 24) | pos) exc = torch.tensor(exceptions, dtype=torch.int32) - ref = target.reshape(-1).clone() + ref = target.clone() for rec in exceptions: - ref[rec & 0xFFFFFF] = rec >> 24 + ref.view(-1)[rec & 0xFFFFFF] = (rec >> 24) & 0xFF return fixed, exc, ref @@ -116,7 +116,7 @@ def test_csf_decode_roundtrip(): assert torch.equal(decode_csf_scales(fixed, exc.view(torch.uint8)).view(torch.uint8), ref) # with no exceptions the slab decode is the target itself (rows share a high nibble) no_exc = decode_csf_scales(fixed, torch.zeros(0, dtype=torch.int32)) - assert torch.equal(no_exc.view(torch.uint8), target.reshape(-1)) + assert torch.equal(no_exc.view(torch.uint8), target) def test_nvfp4_expert_spec_selects_the_csf_checkpoint(monkeypatch): @@ -168,7 +168,7 @@ def test_csf_expert_reader_expands_the_scale_pieces(tmp_path): piece[f"{role}_global"].reshape(-1), raw[f"{base}.weight_scale_2"].reshape(-1).to(torch.float16) ) assert piece[f"{role}_scale"].shape == (1, raw[f"{base}.weight"].shape[0], raw[f"{base}.weight"].shape[1] // 8) - assert torch.equal(piece[f"{role}_scale"].view(torch.uint8)[0].reshape(-1), refs[(e0, proj)]) + assert torch.equal(piece[f"{role}_scale"].view(torch.uint8)[0], refs[(e0, proj)]) def _raw_checkpoint(dense_fp8: bool = False) -> dict[str, torch.Tensor]: From 2d19499c286a4a8483bdb3594f9ce14de9861732 Mon Sep 17 00:00:00 2001 From: chrisqian Date: Fri, 16 Oct 2026 14:32:27 +0800 Subject: [PATCH 06/14] test(qwen4_exp): store CSF U32 exception records as signed int32 bit patterns --- tests/models/qwen4_exp/test_weight.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/models/qwen4_exp/test_weight.py b/tests/models/qwen4_exp/test_weight.py index 844eee47b6..568ba6d69f 100644 --- a/tests/models/qwen4_exp/test_weight.py +++ b/tests/models/qwen4_exp/test_weight.py @@ -91,10 +91,11 @@ def _csf_encode(target: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch fixed[slab, row] = base for j in range(row_bytes): fixed[slab, 16 + row * row_bytes + j] = (target[r, 2 * j + 1] & 0xF) << 4 | (target[r, 2 * j] & 0xF) - # absolute records: (value << 24) | row-major position, as Python ints (uint8 would wrap) + # absolute records: (value << 24) | row-major position, U32 bit patterns stored as int32 for pos in (3, rows * cols - 1): value = 0xFF - target.ravel()[pos].item() - exceptions.append((value << 24) | pos) + rec = (value << 24) | pos + exceptions.append(rec - 0x100000000 if rec >= 0x80000000 else rec) exc = torch.tensor(exceptions, dtype=torch.int32) ref = target.clone() for rec in exceptions: From 54b1dfbbb22ab8c8b6a0ad2966e749ac936e2ced Mon Sep 17 00:00:00 2001 From: chrisqian Date: Fri, 16 Oct 2026 14:34:01 +0800 Subject: [PATCH 07/14] test(qwen4_exp): index the e2m1 LUT with int64 nibbles --- tests/models/qwen4_exp/test_ple.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index 4fd6fb4879..ed916be88f 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -267,7 +267,8 @@ def test_gpu_resident_fp4_lookup_matches_the_lut(): lut = torch.tensor( [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, 0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0] ) - values = torch.stack((lut[packed & 0x0F], lut[(packed >> 4) & 0x0F]), dim=2).reshape(rows, dim) + nibbles = packed.to(torch.int64) + values = torch.stack((lut[nibbles & 0x0F], lut[(nibbles >> 4) & 0x0F]), dim=2).reshape(rows, dim) want = (values * scales.float().repeat_interleave(16, dim=1) * scale).to(torch.bfloat16) assert torch.equal(table.lookup(ids).view(-1, dim), want[ids.reshape(-1)]) From b3135b577ba304142be33cc1731d0fad5afb0fc7 Mon Sep 17 00:00:00 2001 From: chrisqian Date: Fri, 16 Oct 2026 14:34:46 +0800 Subject: [PATCH 08/14] fix(qwen4_exp): widen PLE nibbles to int64 before the e2m1 LUT lookup --- python/freetoken/models/qwen4_exp/ple.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/freetoken/models/qwen4_exp/ple.py b/python/freetoken/models/qwen4_exp/ple.py index f7890c90b5..c8cd172a63 100644 --- a/python/freetoken/models/qwen4_exp/ple.py +++ b/python/freetoken/models/qwen4_exp/ple.py @@ -112,8 +112,9 @@ def lookup(self, row_ids: torch.Tensor, out: torch.Tensor | None = None) -> torc flat = row_ids.reshape(-1) packed = self.weight.index_select(0, flat) scales = self.scale_tensor.index_select(0, flat).view(torch.float8_e4m3fn).float() - lo = self._lut[packed & 0x0F] - hi = self._lut[(packed >> 4) & 0x0F] + nibbles = packed.to(torch.int64) # this torch build rejects uint8 indices + lo = self._lut[nibbles & 0x0F] + hi = self._lut[(nibbles >> 4) & 0x0F] rows = torch.stack((lo, hi), dim=2).reshape(flat.shape[0], self.head_dim) if self.scale != 1.0: scales = scales * self.scale From 3ff0cf09e484770906fdfa9c3d653122b4b5198e Mon Sep 17 00:00:00 2001 From: chrisqian Date: Fri, 16 Oct 2026 15:06:32 +0800 Subject: [PATCH 09/14] fix(models): find index-only checkpoints that keep shards in a subdirectory --- python/freetoken/models/loader.py | 13 ++++++++- tests/models/test_models_loader.py | 45 ++++++++++++++++++++++++++++++ 2 files changed, 57 insertions(+), 1 deletion(-) diff --git a/python/freetoken/models/loader.py b/python/freetoken/models/loader.py index 0f310ccbb2..60a9e90200 100644 --- a/python/freetoken/models/loader.py +++ b/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]: diff --git a/tests/models/test_models_loader.py b/tests/models/test_models_loader.py index b0ac2b13f3..594a0a6560 100644 --- a/tests/models/test_models_loader.py +++ b/tests/models/test_models_loader.py @@ -31,6 +31,51 @@ def test_shard_tensor_splits_vocab_with_ceil_partition(): assert rank1.tolist() == value[3:5].tolist() +def test_iter_weight_files_reads_shards_from_an_indexed_subdirectory(tmp_path): + """FP4-Local releases keep shards in tensors/; the root glob alone finds nothing.""" + import json + import os + from freetoken.models.loader import iter_weight_files + + sub = tmp_path / "tensors" + sub.mkdir() + for i in (1, 2): + (sub / f"model-{i:05d}-of-00002.safetensors").write_bytes(b"") + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps( + { + "weight_map": { + "model.embed_tokens.weight": "tensors/model-00001-of-00002.safetensors", + "lm_head.weight": "tensors/model-00002-of-00002.safetensors", + } + } + ) + ) + + files = iter_weight_files(str(tmp_path)) + + assert len(files) == 2 + assert all(os.path.exists(f) for f in files) + + +def test_iter_weight_files_prefers_root_shards_over_the_index(tmp_path): + """Root-level shards keep the historical path; the index fallback is only for subdirectory layouts.""" + import json + from freetoken.models.loader import iter_weight_files + + (tmp_path / "model-00001-of-00001.safetensors").write_bytes(b"") + sub = tmp_path / "tensors" + sub.mkdir() + (sub / "model-00001-of-00002.safetensors").write_bytes(b"") + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": {"model.embed_tokens.weight": "tensors/model-00001-of-00002.safetensors"}}) + ) + + files = iter_weight_files(str(tmp_path)) + + assert files == [str(tmp_path / "model-00001-of-00001.safetensors")] + + def test_iter_root_safetensor_files_from_index_keeps_only_root_index_shards(tmp_path): import json from freetoken.models.loader import iter_root_safetensor_files_from_index From 545a7fe0af591cfce986a2577d591236f51fc325 Mon Sep 17 00:00:00 2001 From: chrisqian Date: Fri, 16 Oct 2026 15:40:51 +0800 Subject: [PATCH 10/14] fix(qwen4_exp): accept the transformers 5.16 full-attention layer alias --- python/freetoken/models/qwen4_exp/config.py | 6 ++++-- tests/models/qwen4_exp/test_config.py | 12 ++++++++++++ 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/python/freetoken/models/qwen4_exp/config.py b/python/freetoken/models/qwen4_exp/config.py index 45d0cfe95e..cbbb8237e5 100644 --- a/python/freetoken/models/qwen4_exp/config.py +++ b/python/freetoken/models/qwen4_exp/config.py @@ -103,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)) diff --git a/tests/models/qwen4_exp/test_config.py b/tests/models/qwen4_exp/test_config.py index f850927b99..83fe8ae6af 100644 --- a/tests/models/qwen4_exp/test_config.py +++ b/tests/models/qwen4_exp/test_config.py @@ -92,6 +92,18 @@ def test_groups_and_layer_split(): assert linear[0].num_key_heads == 16 and linear[0].num_value_heads == 48 +def test_layer_types_accept_both_hf_full_attention_aliases(): + """HF rewrites full_attention to a version-specific alias; both must map back.""" + from types import SimpleNamespace + from freetoken.models.qwen4_exp.config import _layer_types + + base = ["full_attention" if (i + 1) % 4 == 0 else "linear_attention" for i in range(48)] + for alias in ("indexed_attention", "qwen_sparse_attention"): + rewritten = [alias if t == "full_attention" else t for t in base] + text = SimpleNamespace(layer_types=rewritten, num_hidden_layers=48) + assert _layer_types(text) == base + + def test_kv_specs_resolve_qsa(): cfg = parse_config(_hf_config()) specs = {s.name: s for s in cfg.kv_cache_group_specs()} From 2ca3fce8b751675f32f92a411401fa1c5ece3469 Mon Sep 17 00:00:00 2001 From: chrisqian Date: Fri, 16 Oct 2026 17:30:13 +0800 Subject: [PATCH 11/14] feat(qwen4_exp): serve the FP4-Local dense projections as MXFP8 / W4A16 --- python/freetoken/models/qwen4_exp/weight.py | 109 +++++++++++++++++++- tests/models/qwen4_exp/test_weight.py | 64 ++++++++++++ 2 files changed, 168 insertions(+), 5 deletions(-) diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index 3383489d34..1bd7cc6162 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -29,7 +29,7 @@ Nvfp4ExpertSourceSpec, decode_csf_scales, ) -from freetoken.layers.quantization import get_quant_config +from freetoken.layers.quantization import QuantKind, get_quant_config from freetoken.models.register import get_model_spec from freetoken.moe.host_banks import HostBank, read_range_into from freetoken.utils import cached_load_hf_config, download_hf_weight @@ -67,6 +67,14 @@ ) # Per-tensor modelopt quant scales; consumed with their ``.weight`` (experts) or unused. _SCALE_SUFFIXES = (".weight_scale", ".weight_scale_2", ".input_scale") +# Scale tensor names the modelopt dialect stores per scheme role; nvidia's MIXED_PRECISION +# exports name them after the role, while the model's linears keep the role's FreeToken name +# (MXFP8 block scales arrive as weight_scale_inv, NVFP4 globals as weight_global). +_MODEL_SCALE_ROLE = ( + (QuantKind.MXFP8, ".weight_scale", "weight_scale_inv"), + (QuantKind.NVFP4, ".weight_scale", "weight_scale"), + (QuantKind.NVFP4, ".weight_scale_2", "weight_global"), +) # 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." @@ -115,10 +123,75 @@ def _rename(raw_name: str) -> str | None: if _EXPERT_RE.search(raw_name): return None # routed experts: offload source banks if raw_name.endswith(_SCALE_SUFFIXES): - return None + return _rename_scaled(raw_name) return rename_vl_prefix(raw_name) +def _rename_scaled(raw_name: str) -> str | None: + """Keep a checkpoint scale when the scheme that stores it names it after the model's role. + + nvidia's MIXED_PRECISION exports call the MXFP8 block scale ``weight_scale`` and the NVFP4 + per-tensor global ``weight_scale_2``; the model's linears keep the role names ``weight_scale_inv`` + / ``weight_global``, so rename along the module prefix. Routed experts read their scales + through the bank readers, so a scheme there keeps nothing. + """ + quant = get_quant_config() + if quant is None or _EXPERT_RE.search(raw_name): + return None + module, suffix = _split_scale_suffix(raw_name) + if suffix is None: + return None + scheme = quant.scheme_for_name(module) + if scheme is None: + return None + for kind, csuf, role in _MODEL_SCALE_ROLE: + if scheme.kind is kind and suffix == csuf: + return rename_vl_prefix(module) + "." + role + return None + + +def _split_scale_suffix(raw_name: str) -> tuple[str, str | None]: + for suffix in _SCALE_SUFFIXES: + if raw_name.endswith(suffix): + return raw_name[: -len(suffix)], suffix + return raw_name, None + + +class _GlobalExpander: + """Expand a scalar NVFP4 ``weight_global`` to the model's ``[out]`` fp16 buffer. + + nvidia's W4A16 exports store the per-tensor global as a scalar fp32, while the linear's + buffer holds one entry per output row (load_state_dict asserts the exact shape). The global + and its ``.weight`` arrive in either order, so the expansion is flushed once both are seen. + """ + + def __init__(self) -> None: + self._scalars: dict[str, torch.Tensor] = {} + self._rows: dict[str, int] = {} + + def on_weight(self, module: str, rows: int) -> tuple[str, torch.Tensor] | None: + self._rows[module] = rows + scalar = self._scalars.pop(module, None) + return None if scalar is None else self._expand(module, scalar, rows) + + def on_global(self, name: str, tensor: torch.Tensor) -> tuple[str, torch.Tensor] | None: + """A ready pair, or None for a scalar still waiting for its weight.""" + if tensor.numel() != 1: + return name, tensor + module = name[: -len(".weight_global")] + if module in self._rows: + return self._expand(module, tensor, self._rows[module]) + self._scalars[module] = tensor + return None + + def check_complete(self) -> None: + assert not self._scalars, f"scalar weight_global without a weight: {sorted(self._scalars)}" + + @staticmethod + def _expand(module: str, scalar: torch.Tensor, rows: int) -> tuple[str, torch.Tensor]: + return module + ".weight_global", scalar.reshape(-1).to(torch.float16).expand(rows).contiguous() + + def _split_kind(name: str) -> tuple[str, str]: """``name`` -> ``(module, kind)``; kind is "" for tensors that are neither a weight nor a block scale.""" for suffix in _KIND_SUFFIXES: @@ -221,7 +294,7 @@ def iter_weights( """Yield the dense (non-expert) weights, prefix-stripped and fused to the model's buffers. Keys keep the checkpoint's module names below the stripped prefix, so the emitted set is the model's state dict minus the routed experts. - A dense projection is bf16 or 128x128 block-fp8 (``.weight`` e4m3 + ``.weight_scale_inv``) as the checkpoint's QuantConfig says: the official releases skip everything but the routed experts, the community NVFP4-FP8 requants quantize the attention / GDN projections. + A dense projection is bf16, 128x128 block-fp8, or MXFP8 per-32 (``.weight`` e4m3 + ``.weight_scale``) as the checkpoint's QuantConfig says; the vision MLP is W4A16 NVFP4 in the FP4-Local release. Fusions, per kind: attention q|k|v -> ``qkv_proj``; GDN ``in_proj_{qkv,z,b,a}`` -> ``in_proj``, or ``in_proj_qkvz`` + bf16 ``in_proj_ba`` when qkv|z is quantized; shared-expert gate|up -> ``gate_up_proj``; each per-layer HC's ``input_mix_weight_down`` | ``block_inject_weight`` -> a zero-padded ``input_mix_weight_down_block_inject``. ``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. """ @@ -233,6 +306,7 @@ def iter_weights( hf_config = cached_load_hf_config(model_path) spec = get_model_spec(hf_config.architectures[0]) fuser = _DenseFuser(get_quant_config(), spec.packed_modules_mapping) + globals_ = _GlobalExpander() for file in tqdm( iter_weight_files(model_path), desc="Loading weights", @@ -246,24 +320,49 @@ def iter_weights( if not include_vision and name.startswith(VISION_KEY_PREFIXES): continue tensor = f.get_tensor(raw_name) + if name.endswith(".weight_global"): + pair = globals_.on_global(name, tensor) + if pair is None: + continue + yield pair + continue fused = fuser.fuse(name, tensor) if fused is None: fuser.check_unfused(name, tensor) yield name, tensor else: yield from fused + if name.endswith(".weight"): + pair = globals_.on_weight(name.rpartition(".")[0], tensor.shape[0]) + if pair is not None: + yield pair + globals_.check_complete() assert not fuser.buf, f"Incomplete projection fusions: {sorted(k[0] + k[1] for k in fuser.buf)}" def iter_vision_weights(model_path: str, device: torch.device) -> Iterator[tuple[str, torch.Tensor]]: """The vision tower alone, named as iter_weights names it.""" + globals_ = _GlobalExpander() for file in iter_weight_files(model_path): with safetensors.safe_open(file, framework="pt", device=str(device)) as f: for raw_name in f.keys(): name = _rename(raw_name) - if name is not None and name.startswith(VISION_KEY_PREFIXES): - yield name, f.get_tensor(raw_name) + if name is None or not name.startswith(VISION_KEY_PREFIXES): + continue + tensor = f.get_tensor(raw_name) + if name.endswith(".weight_global"): + pair = globals_.on_global(name, tensor) + if pair is None: + continue + yield pair + continue + yield name, tensor + if name.endswith(".weight"): + pair = globals_.on_weight(name.rpartition(".")[0], tensor.shape[0]) + if pair is not None: + yield pair + globals_.check_complete() # ====================================================================================== diff --git a/tests/models/qwen4_exp/test_weight.py b/tests/models/qwen4_exp/test_weight.py index 568ba6d69f..439f20c198 100644 --- a/tests/models/qwen4_exp/test_weight.py +++ b/tests/models/qwen4_exp/test_weight.py @@ -120,6 +120,70 @@ def test_csf_decode_roundtrip(): assert torch.equal(no_exc.view(torch.uint8), target) +def test_rename_scaled_maps_the_checkpoint_scale_names_to_the_model_roles(monkeypatch): + """FP4-Local stores dense projections as MXFP8 and the vision MLP as W4A16 NVFP4. + + The checkpoint names the block scale ``weight_scale`` while the engine's buffers are + ``weight_scale_inv`` (and the per-tensor global ``weight_scale_2``), so the reader must + rename along the module prefix -- and only for modules the quant config actually names. + """ + from freetoken.layers.quantization import ModelOptConfig + + quant = ModelOptConfig({ + "quant_algo": "MIXED_PRECISION", + "quantized_layers": { + "model.language_model.layers.0.linear_attn.in_proj_qkv": {"quant_algo": "MXFP8"}, + "model.visual.blocks.0.mlp.linear_fc1": {"quant_algo": "W4A16_NVFP4"}, + "model.language_model.layers.0.mlp.experts": {"quant_algo": "NVFP4"}, + }, + }) + monkeypatch.setattr(qwen4_exp_weight, "get_quant_config", lambda: quant) + + # MXFP8 block scale: checkpoint role name -> the engine's buffer name + got = qwen4_exp_weight._rename_scaled("model.language_model.layers.0.linear_attn.in_proj_qkv.weight_scale") + assert got == "model.layers.0.linear_attn.in_proj_qkv.weight_scale_inv" + # W4A16 per-tensor global: the vision prefix lands under visual., the role under weight_global + got = qwen4_exp_weight._rename_scaled("model.visual.blocks.0.mlp.linear_fc1.weight_scale_2") + assert got == "visual.blocks.0.mlp.linear_fc1.weight_global" + # a scale whose module the config leaves unquantized stays noise + assert qwen4_exp_weight._rename_scaled("model.language_model.layers.1.self_attn.q_proj.weight_scale") is None + # routed experts read their scales through the bank readers, never the dense path + assert qwen4_exp_weight._rename_scaled( + "model.language_model.layers.0.mlp.experts.0.gate_proj.weight_scale.nvfp4_csf_fixed" + ) is None + + +def test_global_expander_broadcasts_the_scalar_once_its_weight_is_seen(): + """The W4A16 global is a scalar fp32; the model buffer is one fp16 entry per output row. + + The pair may arrive in either order, so the expansion is flushed only once both are known. + """ + g = qwen4_exp_weight._GlobalExpander() + module = "visual.blocks.0.mlp.linear_fc1" + scalar = torch.tensor(0.0625) + # global first: buffered until the weight fixes the row count + assert g.on_global(module + ".weight_global", scalar) is None + name, tensor = g.on_weight(module, 256) + assert name == module + ".weight_global" + assert tensor.shape == (256,) and tensor.dtype is torch.float16 + assert torch.all(tensor == 0.0625) + # weight first: the global flushes on the spot + g2 = qwen4_exp_weight._GlobalExpander() + assert g2.on_weight(module, 64) is None + name, tensor = g2.on_global(module + ".weight_global", scalar) + assert name == module + ".weight_global" and tensor.shape == (64,) + # a non-scalar global passes through untouched + row = torch.arange(3, dtype=torch.float16) + assert g2.on_global(module + ".weight_global", row) == (module + ".weight_global", row) + g.check_complete() + g2.check_complete() + # a scalar that never met its weight fails the completeness assert + g3 = qwen4_exp_weight._GlobalExpander() + g3.on_global(module + ".weight_global", scalar) + with pytest.raises(AssertionError, match="without a weight"): + g3.check_complete() + + def test_nvfp4_expert_spec_selects_the_csf_checkpoint(monkeypatch): monkeypatch.setattr(qwen4_exp_weight, "download_hf_weight", lambda path: "") weight_map = { From 59953e31a5d09ed54b62ef1545fe6d6c7e49a953 Mon Sep 17 00:00:00 2001 From: chrisqian Date: Fri, 16 Oct 2026 17:55:48 +0800 Subject: [PATCH 12/14] fix(qwen4_exp): fall back to the pinned PLE backend for packed nvfp4 tables --- python/freetoken/engine/config.py | 2 +- python/freetoken/models/qwen4_exp/model.py | 14 ++++++++++++-- tests/models/qwen4_exp/test_skeleton.py | 10 ++++++++++ 3 files changed, 23 insertions(+), 3 deletions(-) diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index a6c47dd0fb..cfe0b4daa8 100644 --- a/python/freetoken/engine/config.py +++ b/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. diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py index d44456b568..921d9d4ccc 100644 --- a/python/freetoken/models/qwen4_exp/model.py +++ b/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 diff --git a/tests/models/qwen4_exp/test_skeleton.py b/tests/models/qwen4_exp/test_skeleton.py index 8fb9112df2..d183887d58 100644 --- a/tests/models/qwen4_exp/test_skeleton.py +++ b/tests/models/qwen4_exp/test_skeleton.py @@ -20,6 +20,7 @@ from freetoken.models.config import ModelConfig from freetoken.models.qwen4_exp.config import parse_config from freetoken.models.qwen4_exp.hc import GatedResidual +from freetoken.models.qwen4_exp.model import _resolve_ple_backend from freetoken.models.qwen4_exp.ple import GpuResidentTable, PLELayer, PLEMetadata from .common import EOS, hash_constants, requires_cuda, toy_hf_config @@ -507,3 +508,12 @@ def test_decoder_stack_prefill_and_decode(monkeypatch): decode_logits = model.forward() assert decode_logits.shape == (len(prompts), config.vocab_size) assert torch.isfinite(decode_logits.float()).all() + + +def test_ple_backend_falls_back_to_pinned_for_nvfp4_tables(): + """The disk row store streams native F8 rows; packed e2m1 (nvfp4) tables must go pinned.""" + assert _resolve_ple_backend("disk", "nvfp4") == "pinned" + # fp8 tables keep the requested backend either way + assert _resolve_ple_backend("disk", "fp8") == "disk" + assert _resolve_ple_backend("pinned", "nvfp4") == "pinned" + assert _resolve_ple_backend("pinned", "fp8") == "pinned" From f4a9df2681c11802501f8ef5b9a88d6e32920b8d Mon Sep 17 00:00:00 2001 From: chrisqian Date: Fri, 16 Oct 2026 18:02:27 +0800 Subject: [PATCH 13/14] fix(models): read the U32 CSF exceptions stream in the parallel expert loader --- python/freetoken/models/weight.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/freetoken/models/weight.py b/python/freetoken/models/weight.py index b03266fb67..405a7550ad 100644 --- a/python/freetoken/models/weight.py +++ b/python/freetoken/models/weight.py @@ -32,6 +32,8 @@ "I64": torch.int64, "I32": torch.int32, "I16": torch.int16, "I8": torch.int8, "U8": torch.uint8, "BOOL": torch.bool, "F8_E4M3": torch.float8_e4m3fn, "F8_E5M2": torch.float8_e5m2, "F8_E8M0": torch.float8_e8m0fnu, + # unsigned integer types (the CSF scale exceptions stream is U32) + "U16": torch.uint16, "U32": torch.uint32, "U64": torch.uint64, } _ODIRECT_BLK = 4096 From 994504613746c2ba74777a3cf4c8bea8ba1803de Mon Sep 17 00:00:00 2001 From: chrisqian Date: Fri, 16 Oct 2026 18:18:57 +0800 Subject: [PATCH 14/14] fix(kernels): decode the e2m1 sign without a non-boolean tl.where --- python/freetoken/kernel/triton/ple.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/freetoken/kernel/triton/ple.py b/python/freetoken/kernel/triton/ple.py index 4f9861dff6..94b68074ff 100644 --- a/python/freetoken/kernel/triton/ple.py +++ b/python/freetoken/kernel/triton/ple.py @@ -34,7 +34,8 @@ def _e2m1_f32(n: tl.uint8) -> 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))) - return tl.where((n >> 3) & 1, -mag, mag) + sign = ((n >> 3) & 1) * mag + return mag - 2 * sign @triton.jit