From 2d980071bb56c0c6e3d3acf94774c8b5ba26b7ff Mon Sep 17 00:00:00 2001 From: chrisqian Date: Sat, 24 Oct 2026 10:16:38 +0800 Subject: [PATCH] feat(models): add nvidia/Qwen3.8-Flash-Next-FP4-Local checkpoint support --- docs/models.md | 2 +- python/freetoken/engine/config.py | 2 +- python/freetoken/kernel/triton/ple.py | 49 ++- python/freetoken/models/nvfp4_banks.py | 130 ++++++- python/freetoken/models/qwen4_exp/config.py | 3 + python/freetoken/models/qwen4_exp/model.py | 21 +- python/freetoken/models/qwen4_exp/ple.py | 70 +++- python/freetoken/models/qwen4_exp/ple_disk.py | 10 + python/freetoken/models/qwen4_exp/weight.py | 250 +++++++++++-- python/freetoken/models/weight.py | 2 + python/freetoken/moe/expert_pieces.py | 15 +- tests/models/qwen4_exp/test_ple.py | 49 +++ tests/models/qwen4_exp/test_ple_disk.py | 18 + tests/models/qwen4_exp/test_skeleton.py | 10 + tests/models/qwen4_exp/test_weight.py | 343 ++++++++++++++++++ 15 files changed, 915 insertions(+), 59 deletions(-) diff --git a/docs/models.md b/docs/models.md index a9b66a2480..9ad91ffa10 100644 --- a/docs/models.md +++ b/docs/models.md @@ -10,7 +10,7 @@ for them; other checkpoints of the same architectures work too. | GLM-5.3-Flash | [RedHatAI/GLM-5.3-Flash-NVFP4](https://huggingface.co/RedHatAI/GLM-5.3-Flash-NVFP4) | | GLM-5.2 | [nvidia/GLM-5.2-NVFP4](https://huggingface.co/nvidia/GLM-5.2-NVFP4) | | GLM-4.7 | [nvidia/GLM-4.7-NVFP4](https://huggingface.co/nvidia/GLM-4.7-NVFP4) | -| Qwen3.8-Flash-Next | [Qwen/Qwen3.8-Flash-Next-FP8](https://huggingface.co/Qwen/Qwen3.8-Flash-Next-FP8), [RadixArk/Qwen3.8-Flash-Next-NVFP4](https://huggingface.co/RadixArk/Qwen3.8-Flash-Next-NVFP4), [nvidia/Qwen3.8-Flash-Next-NVFP4](https://huggingface.co/nvidia/Qwen3.8-Flash-Next-NVFP4) | +| Qwen3.8-Flash-Next | [Qwen/Qwen3.8-Flash-Next-FP8](https://huggingface.co/Qwen/Qwen3.8-Flash-Next-FP8), [RadixArk/Qwen3.8-Flash-Next-NVFP4](https://huggingface.co/RadixArk/Qwen3.8-Flash-Next-NVFP4), [nvidia/Qwen3.8-Flash-Next-NVFP4](https://huggingface.co/nvidia/Qwen3.8-Flash-Next-NVFP4), [nvidia/Qwen3.8-Flash-Next-FP4-Local](https://huggingface.co/nvidia/Qwen3.8-Flash-Next-FP4-Local) | | Qwen3.6 / Qwen3.5 MoE | [Qwen/Qwen3.6-35B-A3B](https://huggingface.co/Qwen/Qwen3.6-35B-A3B) ([-FP8](https://huggingface.co/Qwen/Qwen3.6-35B-A3B-FP8)), [nvidia/Qwen3.6-35B-A3B-NVFP4](https://huggingface.co/nvidia/Qwen3.6-35B-A3B-NVFP4), [Qwen/Qwen3.5-35B-A3B](https://huggingface.co/Qwen/Qwen3.5-35B-A3B) ([-FP8](https://huggingface.co/Qwen/Qwen3.5-35B-A3B-FP8)) | | Qwen3.8 / Qwen3.6 dense | [Qwen/Qwen3.8-27B](https://huggingface.co/Qwen/Qwen3.8-27B) ([-FP8](https://huggingface.co/Qwen/Qwen3.8-27B-FP8)), [RadixArk/Qwen3.8-27B-NVFP4](https://huggingface.co/RadixArk/Qwen3.8-27B-NVFP4), [Qwen/Qwen3.6-27B](https://huggingface.co/Qwen/Qwen3.6-27B) ([-FP8](https://huggingface.co/Qwen/Qwen3.6-27B-FP8)), [nvidia/Qwen3.6-27B-NVFP4](https://huggingface.co/nvidia/Qwen3.6-27B-NVFP4) | | Qwen3-MoE | [Qwen/Qwen3-30B-A3B](https://huggingface.co/Qwen/Qwen3-30B-A3B) | diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 9593667e2c..d89156fc98 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/kernel/triton/ple.py b/python/freetoken/kernel/triton/ple.py index 2314e94e72..ec232f175b 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,32 @@ _NUM_WARPS = 1 +@triton.jit +def _e2m1_f32(n: tl.uint8) -> tl.float32: + """The 16 e2m1 nibbles: 0, +/-0.5, +/-1, +/-1.5, +/-2, +/-3, +/-4, +/-6 (sign = bit 3).""" + n = n.to(tl.int32) + mant = (n & 1).to(tl.float32) + exp2 = (n >> 1) & 3 + # exponent 0 is subnormal (0.mant); otherwise (1.mant) * 2 ** (exp2 - 1) + mag = tl.where(exp2 == 0, mant * 0.5, (1.0 + mant * 0.5) * tl.where(exp2 == 1, 1.0, tl.where(exp2 == 2, 2.0, 4.0))) + sign = ((n >> 3) & 1) * mag + return mag - 2 * sign + + @triton.jit def _ple_gather_kernel( table_ptr, + scale_ptr, ids_ptr, out_ptr, scale, num_rows, EMB_DIM: tl.constexpr, IS_FP8: tl.constexpr, + IS_FP4: tl.constexpr, BLOCK_D: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_S: tl.constexpr, ): row = tl.program_id(0) idx = tl.load(ids_ptr + row).to(tl.int64) @@ -49,6 +68,24 @@ def _ple_gather_kernel( # pre-sm_89 has no fp8e4nv type: load raw bytes and decode in software base = table_ptr.to(tl.int64).to(tl.pointer_type(tl.uint8)) values = e4m3_u8_to_f32(tl.load(base + idx * EMB_DIM + offsets, mask=mask, other=0)) + elif IS_FP4: + # packed e2m1: two values per byte (low nibble first), per-16 e4m3 scale row, global scale + pack_offsets = tl.arange(0, BLOCK_H) + vbytes = tl.load(table_ptr.to(tl.int64).to(tl.pointer_type(tl.uint8)) + idx * (EMB_DIM // 2) + pack_offsets, + mask=pack_offsets < EMB_DIM // 2, other=0) + lo = _e2m1_f32(vbytes & 0x0F) + hi = _e2m1_f32((vbytes >> 4) & 0x0F) + values = tl.join(lo, hi).reshape([BLOCK_D]) + scale_offsets = tl.arange(0, BLOCK_S) + scale_mask = scale_offsets < EMB_DIM // 16 + if e4m3_native_cx(): + scale_base = scale_ptr.to(tl.int64).to(tl.pointer_type(tl.float8e4nv)) + svals = tl.load(scale_base + idx * (EMB_DIM // 16) + scale_offsets, mask=scale_mask, other=0.0).to(tl.float32) + else: + scale_base = scale_ptr.to(tl.int64).to(tl.pointer_type(tl.uint8)) + svals = e4m3_u8_to_f32(tl.load(scale_base + idx * (EMB_DIM // 16) + scale_offsets, mask=scale_mask, other=0)) + svals = tl.broadcast_to(svals.reshape([BLOCK_S, 1]), [BLOCK_S, 16]).reshape([BLOCK_D]) + values = values * svals else: base = table_ptr.to(tl.int64).to(tl.pointer_type(tl.bfloat16)) values = tl.load(base + idx * EMB_DIM + offsets, mask=mask, other=0.0).to(tl.float32) @@ -68,25 +105,35 @@ def ple_gather_rows( out: torch.Tensor, scale: float = 1.0, is_fp8: bool = True, + scale_ptr: int | None = None, ) -> torch.Tensor: """Gather ``row_ids`` from the host-resident table at ``table_ptr`` into ``out``. ``row_ids`` is a flat device int tensor; ``out`` is ``[row_ids.numel(), embed_dim]`` bf16 on the same device. ``table_ptr`` is the address the GPU must dereference - (``kernel/pinned.device_ptr``), not necessarily the host ``data_ptr``. + (``kernel/pinned.device_ptr``), not necessarily the host ``data_ptr``. Packed e2m1 tables + (``is_fp4``) pass the per-16 e4m3 scale row bank as ``scale_ptr`` and the global scale as + ``scale``. """ + assert not (is_fp8 and scale_ptr is not None), "a packed table is not also a native fp8 table" + if scale_ptr is not None: + is_fp8 = False n = row_ids.numel() assert out.shape == (n, embed_dim) and out.is_contiguous(), out.shape if n: _ple_gather_kernel[(n,)]( table_ptr, + scale_ptr if scale_ptr is not None else table_ptr, row_ids, out, float(scale), num_rows, EMB_DIM=embed_dim, IS_FP8=is_fp8, + IS_FP4=scale_ptr is not None, BLOCK_D=triton.next_power_of_2(embed_dim), + BLOCK_H=triton.next_power_of_2(embed_dim // 2), + BLOCK_S=triton.next_power_of_2(embed_dim // 16), num_warps=_NUM_WARPS, ) return out diff --git a/python/freetoken/models/nvfp4_banks.py b/python/freetoken/models/nvfp4_banks.py index 5382294573..c329aa49ab 100644 --- a/python/freetoken/models/nvfp4_banks.py +++ b/python/freetoken/models/nvfp4_banks.py @@ -27,6 +27,12 @@ class Nvfp4ExpertSourceSpec: # The checkpoint stores the QUANT-side global scale (local fp8 scales were # multiplied by it before the cast); the banks keep its reciprocal. global_reciprocal: bool = False + # The checkpoint stores the block scales as a lossless CSF slab pair instead of + # the plain E4M3 matrix (e.g. "weight_scale.nvfp4_csf_fixed" / + # "weight_scale.nvfp4_csf_exceptions"); scale_decoder expands the parts into the + # plain matrix the banks consume. + scale_pair: tuple[str, ...] | None = None + scale_decoder: Callable[..., torch.Tensor] | None = None def _canon_kind(spec: "Nvfp4ExpertSourceSpec", kind: str) -> str: @@ -85,6 +91,43 @@ def _kind_suffix(kind: str) -> str: return {"weight": "", "weight_scale": "_scale", "weight_scale_2": "_global"}[kind] +def decode_csf_scales(fixed: torch.Tensor, exceptions: torch.Tensor) -> torch.Tensor: + """Expand the CSF slab encoding of an NVFP4 per-16 scale into the plain E4M3 matrix (CPU). + + ``fixed`` is U8 ``[slabs, 16 * (1 + row_bytes)]``: the first 16 bytes of each slab hold the + per-row scale bases (one byte per row of the slab's 16 rows), the rest are row-major code + bytes with column c at bit 4c. ``exceptions`` is U32 records, ``pos = rec & 0xFFFFFF`` + (row-major scale index, below 2**24) and ``value = rec >> 24`` (the exact E4M3 byte), + applied after base + offset. + """ + if fixed.ndim != 2 or fixed.dtype != torch.uint8: + raise ValueError(f"CSF fixed slab tensor is {tuple(fixed.shape)} {fixed.dtype}, expected U8 [slabs, slab_bytes]") + slabs, slab_bytes = fixed.shape + row_bytes = slab_bytes // 16 - 1 + if row_bytes < 1 or slab_bytes != 16 * (1 + row_bytes): + raise ValueError(f"CSF slab width {slab_bytes} does not hold 16 bases plus 16 * row_bytes code bytes") + rows = slabs * 16 + cols = row_bytes * 2 + data = fixed.to(torch.int16) + base = data[:, :16].reshape(rows, 1) + codes = data[:, 16:].reshape(rows, row_bytes) + bits = 4 * torch.arange(cols) + scale = (base + ((codes[:, bits // 8] >> bits % 8) & 0x0F)) & 0xFF + if exceptions.numel(): + # Records are U32 on disk or raw U8 bytes; the bit math runs in int64 because CPU torch + # has no uint32 shifts and uint32 -> int32 casts wrap records with the sign bit set. + rec = exceptions.reshape(-1) + if rec.dtype == torch.uint8: + rec = rec.view(torch.uint32) + if rec.dtype not in (torch.uint32, torch.int32): + raise ValueError(f"CSF exceptions are U32 records (or U8 bytes), got {rec.dtype}") + rec = rec.to(torch.int64) + scale = scale.to(torch.int32) + flat = scale.view(-1) + flat.scatter_(0, rec & 0xFFFFFF, ((rec >> 24) & 0xFF).to(torch.int32)) + return scale.to(torch.uint8).view(torch.float8_e4m3fn) + + def iter_nvfp4_expert_pieces( model_path: str, config, @@ -117,6 +160,7 @@ def iter_nvfp4_expert_pieces( keep = None if layers is None else set(layers) act = _activation_scale() wanted: dict[str, tuple[int, int, str]] = {} + pair_kinds = spec.scale_pair for name in weight_map: match = spec.key_pattern.match(name) if match is None: @@ -128,23 +172,41 @@ def iter_nvfp4_expert_pieces( if proj not in spec.proj_to_role: raise ValueError(f"{spec.desc}: unknown NVFP4 expert projection {proj!r}") kind = _canon_kind(spec, match.group("kind")) - if kind not in ("weight", "weight_scale", "weight_scale_2"): + role = spec.proj_to_role[proj] + if pair_kinds is not None: + if kind in pair_kinds: + wanted[name] = (bank_layer, int(match.group("expert")), f"{role}~{kind}") + elif kind in ("weight", "weight_scale_2"): + wanted[name] = (bank_layer, int(match.group("expert")), role + _kind_suffix(kind)) + else: + raise ValueError(f"{spec.desc}: unknown NVFP4 expert tensor kind {kind!r}") + elif kind not in ("weight", "weight_scale", "weight_scale_2"): raise ValueError(f"{spec.desc}: unknown NVFP4 expert tensor kind {kind!r}") - wanted[name] = (bank_layer, int(match.group("expert")), spec.proj_to_role[proj] + _kind_suffix(kind)) + else: + wanted[name] = (bank_layer, int(match.group("expert")), role + _kind_suffix(kind)) if kind == "weight" and act is not None: act_name = name[: match.start("kind")] + act.name if act_name in weight_map: - wanted[act_name] = (bank_layer, int(match.group("expert")), spec.proj_to_role[proj] + "_input") + wanted[act_name] = (bank_layer, int(match.group("expert")), role + "_input") experts = (_num_moe_layers(config) if keep is None else len(keep)) * config.num_experts - acts = [name for name, (_, _, role) in wanted.items() if role.endswith("_input")] + acts = [name for name, (_, _, slot) in wanted.items() if slot.endswith("_input")] # a partial set serves no W4A4 kernel; the pieces then carry none if len(acts) != experts * 3: for name in acts: del wanted[name] acts = [] - per_expert = 12 if acts else 9 - if len(wanted) != experts * per_expert: - raise ValueError(f"{spec.desc}: found {len(wanted)} expert tensors, expected {experts * per_expert}") + # A CSF checkpoint keeps every input_scale in its own trailing shard; streaming them + # with the weights makes the first piece wait on that shard while every other tensor + # piles up unread in the queue. The CSF path reads them out of band instead. + act_keys: dict[str, tuple[int, int, str]] = {} + if pair_kinds is not None: + for name in acts: + act_keys[name] = wanted.pop(name) + acts = list(act_keys) + per_expert = (3 * (2 + len(pair_kinds)) + (3 if acts else 0)) if pair_kinds is not None else (12 if acts else 9) + streamed = per_expert - (3 if pair_kinds is not None and acts else 0) + if len(wanted) != experts * streamed: + raise ValueError(f"{spec.desc}: found {len(wanted)} expert tensors, expected {experts * streamed}") def _serial(): by_shard: dict[str, list[str]] = collections.defaultdict(list) @@ -165,7 +227,57 @@ def _parallel(): for name, tensor in iter_expert_tensors_parallel(folder, lambda n: n in wanted, workers=workers, chunk=chunk): yield name, _ingest(spec, act, wanted[name][2], tensor) - return per_expert_pieces(_parallel() if parallel else _serial(), wanted.get, tensors_per_expert=per_expert) + if pair_kinds is None: + return per_expert_pieces(_parallel() if parallel else _serial(), wanted.get, tensors_per_expert=per_expert) + if spec.scale_decoder is None: + raise ValueError(f"{spec.desc}: scale_pair requires a scale_decoder") + decoder = spec.scale_decoder + + # CSF checkpoints stream the slab parts as separate tensors; group them back into the + # same logical roles (plus the W4A4 input scales when the set is complete), decoding + # the scale pair on the fly. + roles = list(spec.proj_to_role.values()) + slots_per_expert = { + slot + for role in roles + for slot in [f"{role}", f"{role}_global", *(f"{role}~{kind}" for kind in pair_kinds)] + } + act_index = {key: name for name, key in act_keys.items()} + + def _csf_pieces(): + acts_in: dict[str, torch.Tensor] = {} + if act_keys: + act_shards: dict[str, list[str]] = collections.defaultdict(list) + for name in act_keys: + act_shards[weight_map[name]].append(name) + for shard in tqdm(sorted(act_shards), desc=f"Loading {spec.desc} scales", disable=not primary): + path = os.path.join(folder, shard) + _drop(path) + with safetensors.safe_open(path, framework="pt", device="cpu") as f: + for name in act_shards[shard]: + acts_in[name] = _ingest(spec, act, act_keys[name][2], f.get_tensor(name)) + _drop(path) + + def _assemble(layer: int, expert: int, parts: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + grouped: dict[str, torch.Tensor] = {} + for role in roles: + grouped[role] = parts[role].unsqueeze(0) + grouped[f"{role}_global"] = parts[f"{role}_global"].reshape(1, -1) + if acts: + grouped[f"{role}_input"] = acts_in[act_index[(layer, expert, f"{role}_input")]].reshape(1, -1) + grouped[f"{role}_scale"] = decoder( + *[parts[f"{role}~{kind}"] for kind in pair_kinds] + ).unsqueeze(0) + return grouped + + yield from per_expert_pieces( + _parallel() if parallel else _serial(), + wanted.get, + tensors_per_expert=len(slots_per_expert), + assemble=_assemble, + ) + + return _csf_pieces() -__all__ = ["Nvfp4ExpertSourceSpec", "iter_nvfp4_expert_pieces"] +__all__ = ["Nvfp4ExpertSourceSpec", "decode_csf_scales", "iter_nvfp4_expert_pieces"] 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..2351560b42 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.warning_rank0("PLE table is packed e2m1 (nvfp4); the disk backend streams native rows only -- falling back to the pinned backend") + return "pinned" + return requested + def build_linear_mixer(config: ModelConfig, layer_id: int, prefix: str) -> BaseOP: """GDN mixer of a linear_attention layer (Qwen3.5's GDN with a configurable output gate).""" @@ -170,7 +180,7 @@ def load_host_tables(self, engine_config) -> int: emb.attach_table(ZeroTable(offsets[-1] + sizes[-1], args.ngram_head_dim)) return 0 - if engine_config.ple_backend == "disk": + if _resolve_ple_backend(engine_config.ple_backend, getattr(self._config.qwen4_args, "ple_embedding_dtype", "fp8")) == "disk": from .ple_disk import DiskRowTable, resolve_row_source # one WAIT node per captured graph: the flag protocol supports a single consume @@ -201,12 +211,13 @@ def load_host_tables(self, engine_config) -> int: from .weight import load_ple_table table = load_ple_table(engine_config.model_path, self._config.qwen4_args) - self._ple_table = table # owns the pinned HostBank; keep it alive + self._ple_table = table # owns the pinned HostBank(s); keep it alive + kwargs = {"scale_tensor": table.scale_bank.tensor} if table.scale_bank is not None else {} for ple in ple_layers: ple.ple_embedding.attach_table( - PinnedUVATable(table.bank.tensor, float(table.weight_scale)) + PinnedUVATable(table.bank.tensor, float(table.weight_scale), **kwargs) ) - return table.bank.nbytes + return table.bank.nbytes + (table.scale_bank.nbytes if table.scale_bank is not None else 0) def forward(self) -> torch.Tensor: batch = get_global_ctx().batch diff --git a/python/freetoken/models/qwen4_exp/ple.py b/python/freetoken/models/qwen4_exp/ple.py index c944c4946f..2d0b3f4241 100644 --- a/python/freetoken/models/qwen4_exp/ple.py +++ b/python/freetoken/models/qwen4_exp/ple.py @@ -70,22 +70,56 @@ 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.dtype == torch.float8_e4m3fn, "per-16 PLE scales are e4m3 bytes" + assert scale_tensor.shape == (weight.shape[0], weight.shape[1] // 8), "one scale per 16 values" + 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() + 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 + 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 +155,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 +173,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 +223,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 536853be3b..defd5a439a 100644 --- a/python/freetoken/models/qwen4_exp/ple_disk.py +++ b/python/freetoken/models/qwen4_exp/ple_disk.py @@ -21,6 +21,7 @@ from .weight import ( _PLE_SCALE_SUFFIX, + _PLE_SHARD_SCALE_RE, _PLE_SHARD_RE, _PLE_ST_DTYPES, _ple_table_files, @@ -70,6 +71,15 @@ def source_from_safetensors(folder: str) -> PleRowSource: shards: dict[int, tuple[int, int]] = {} for path in _ple_table_files(folder): header, base = _safetensors_header(path) + # a packed e2m1 table pairs U8 shard rows with per-shard e4m3 scale rows; plain + # checkpoints never carry those scale keys, so U8 there is just an unsupported dtype + if any(_PLE_SHARD_SCALE_RE.search(k) is not None for k in header) and any( + header[k]["dtype"] == "U8" for k in header if _PLE_SHARD_RE.search(k) is not None + ): + raise NotImplementedError( + f"the disk PLE backend reads native rows only; {path} holds a packed e2m1 " + "(NVFP4) table -- use the pinned backend (--ple-backend pinned)" + ) for key, meta in header.items(): if key == "__metadata__": continue diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index 6c62d141cb..9877b759b9 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 QuantConfig, QuantKind, QuantScheme, get_quant_config from freetoken.models.register import get_model_spec @@ -58,12 +59,33 @@ _ELEM_DTYPES = {"e4m3": torch.float8_e4m3fn, "e2m1": torch.uint8} _QUANT_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2, torch.uint8, torch.int8) +# 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") + +# The vision tower bypasses the scheme-driven reader (its leaves are not in _LINEAR_LEAVES), +# so its role-named ModelOpt scales are renamed here instead; text dense roles go through +# QuantConfig storage aliases in _DenseReader. +_VISION_PREFIXES = ("model.visual.", "visual.") +_SCALE_SUFFIXES = (".weight_scale", ".weight_scale_2", ".input_scale") +_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." _PLE_SHARD_RE = re.compile( 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 @@ -98,9 +120,96 @@ def _rename(raw_name: str) -> str | None: # static KV-cache scales of the quantizers; the KV cache runs in the engine's dtype if raw_name.endswith((".k_scale", ".v_scale", ".q_scale", ".prob_scale")): return None + if raw_name.startswith(_VISION_PREFIXES) and raw_name.endswith(_SCALE_SUFFIXES): + return _rename_scaled(raw_name) return rename_vl_prefix(raw_name) +def _rename_scaled(raw_name: str) -> str | None: + """Keep a vision scale when the scheme that stores it names it after the model's role. + + The vision tower bypasses _DenseReader, so nvidia's MIXED_PRECISION role names are mapped + here: the MXFP8 block scale ``weight_scale`` -> ``weight_scale_inv``, the NVFP4 per-tensor + global ``weight_scale_2`` -> ``weight_global``. 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 vision NVFP4 ``weight_global`` to the model's ``[out]`` fp32 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: + rows = self._rows.get(name[: -len(".weight_global")]) + if rows is not None and tensor.numel() != rows: + raise ValueError( + f"{name} holds {tensor.numel()} elements for a weight with {rows} output rows" + ) + 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: + if self._scalars: + raise ValueError(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", _per_row_scale(scalar, rows) + + +def _emit_vision(globals_: _GlobalExpander, name: str, tensor: torch.Tensor) -> Iterator[tuple[str, torch.Tensor]]: + """Yield one vision tensor, flushing a scalar ``weight_global`` once its weight is seen.""" + if name.endswith(".weight_global"): + pair = globals_.on_global(name, tensor) + if pair is not None: + yield pair + return + yield name, tensor + if name.endswith(".weight"): + pair = globals_.on_weight(name.rpartition(".")[0], tensor.shape[0]) + if pair is not None: + yield pair + + def _per_row_scale(scale: torch.Tensor, rows: int) -> torch.Tensor: """Per-tensor scalar or per-channel ``[rows, 1]`` scale -> fp32 ``[rows]``; any other count is refused rather than broadcast onto the wrong rows.""" flat = scale.reshape(-1).to(torch.float32) @@ -271,6 +380,7 @@ def iter_weights( hf_config = cached_load_hf_config(model_path) spec = get_model_spec(hf_config.architectures[0]) reader = _DenseReader(get_quant_config(), spec.packed_modules_mapping) + globals_ = _GlobalExpander() # the vision tower's scalar weight_globals bypass the reader for file in tqdm( iter_weight_files(model_path), desc="Loading weights", @@ -284,12 +394,16 @@ def iter_weights( if not include_vision and name.startswith(VISION_KEY_PREFIXES): continue tensor = f.get_tensor(raw_name) - emitted = None if name.startswith(VISION_KEY_PREFIXES) else reader.add(name, tensor) + if name.startswith(VISION_KEY_PREFIXES): + yield from _emit_vision(globals_, name, tensor) + continue + emitted = reader.add(name, tensor) if emitted is None: yield name, tensor else: yield from emitted + globals_.check_complete() if reader.pending: lines = reader.missing() shown = "\n ".join(lines[:8]) + (f"\n ... {len(lines) - 8} more" if len(lines) > 8 else "") @@ -298,12 +412,16 @@ def iter_weights( 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) + yield from _emit_vision(globals_, name, tensor) + globals_.check_complete() # ====================================================================================== @@ -313,14 +431,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 scale (1 for a bf16 table).""" + """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``; bf16 tables hold bf16 values and + a scale of 1; 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]`` view of the bank in the stored dtype.""" + """``[total_rows, ngram_head_dim // (2 if packed else 1)]`` view of the value bank in the stored dtype.""" return self.bank.tensor @@ -387,42 +511,82 @@ def load_ple_table(model_path: str, qwen4_args, *, pin: bool = True, table is ~47.7 GiB (fp8) and must not also sit in the page cache while the bank holds the same bytes. An fp8 table needs the checkpoint's scalar ``weight_scale``; a bf16 table without one is unscaled. """ + encoding = getattr(qwen4_args, "ple_embedding_dtype", "fp8") + if encoding not in ("fp8", "bf16", "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]] = {} + shapes: dict[int, tuple[int, ...]] = {} + scale_shapes: dict[int, tuple[int, ...]] = {} scale: torch.Tensor | None = None - rows = cols = 0 dtype: torch.dtype | None = None + # fp8 stays dtype-driven (F8_E4M3 or a bf16 table); nvfp4 packs e2m1, bf16 demands bf16 + _want_st = {"nvfp4": ("U8",), "bf16": ("BF16",), "fp8": ("F8_E4M3", "BF16")}[encoding] 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"] not in _want_st: + raise ValueError( + f"PLE table shard {key} has dtype {meta['dtype']}, " + f"expected {' or '.join(_want_st)} for ple_embedding_dtype={encoding}" + ) + shard_dtype = {**_PLE_ST_DTYPES, "U8": torch.uint8}[meta["dtype"]] + if dtype not in (None, shard_dtype): + raise ValueError(f"PLE table shard {key} mixes dtypes: {meta['dtype']} after {dtype}") + dtype = shard_dtype + index = int(match.group("shard")) + begin, end = meta["data_offsets"] + parts[index] = (path, base + begin, end - begin) + shapes[index] = tuple(meta["shape"]) continue - shard_dtype = _PLE_ST_DTYPES.get(meta["dtype"]) - if shard_dtype is None or dtype not in (None, shard_dtype): - raise ValueError(f"PLE table shard {key} has dtype {meta['dtype']}, expected one of {sorted(_PLE_ST_DTYPES)} for every shard") - dtype = shard_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"] != "F8_E4M3": + raise ValueError(f"PLE scale shard {key} has dtype {meta['dtype']}, expected F8_E4M3") + index = int(match.group("shard")) + begin, end = meta["data_offsets"] + scale_parts[index] = (path, base + begin, end - begin) + scale_shapes[index] = tuple(meta["shape"]) 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}") + if len(set(shapes.values())) > 1: + raise ValueError(f"PLE table shards disagree on shape: {sorted(set(shapes.values()))[:8]}") + rows, cols = shapes[min(shapes)] + head_dim = qwen4_args.ngram_head_dim + if encoding == "nvfp4": + if head_dim % 16: + raise ValueError(f"nvfp4 PLE heads must be a multiple of 16, got {head_dim}") + 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]}" + ) + want = (rows, cols // 8) + for shard in sorted(scale_shapes): + if scale_shapes[shard] != want: + raise ValueError(f"PLE scale shard {shard} is {list(scale_shapes[shard])}, expected {list(want)}") + elif cols != head_dim: + raise ValueError(f"PLE table row is {cols} wide, config says {head_dim}") if scale is None: + if encoding == "nvfp4": + raise ValueError("PLE table has no weight_scale_2") if dtype is torch.float8_e4m3fn: raise ValueError("PLE table has no weight_scale") scale = torch.tensor(1.0, dtype=torch.bfloat16) @@ -434,15 +598,35 @@ def load_ple_table(model_path: str, qwen4_args, *, pin: bool = True, buf = bank.memoryview() for shard in range(expected): path, offset, nbytes = parts[shard] - assert nbytes == shard_bytes, f"PLE shard {shard} is {nbytes} B, expected {shard_bytes}" + if nbytes != shard_bytes: + raise ValueError(f"PLE shard {shard} is {nbytes} B, expected {shard_bytes}") read_range_into(buf, path, file_offset=offset, nbytes=nbytes, dest_offset=shard * shard_bytes, workers=workers, chunk=chunk) 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] + if nbytes != scale_shard_bytes: + raise ValueError(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) # ====================================================================================== @@ -451,17 +635,29 @@ def load_ple_table(model_path: str, qwen4_args, *, pin: bool = True, def nvfp4_expert_spec(model_path: str, config) -> Nvfp4ExpertSourceSpec: - """The per-expert NVFP4 layout under the checkpoint's dialect names (ModelOpt or llm-compressor).""" + """The per-expert NVFP4 layout under the checkpoint's dialect names (ModelOpt or llm-compressor). + + FP4-Local exports the block scales as CSF slab pairs; scan the weight map and decode them + through the shared decoder so the banks still see plain e4m3 matrices. + """ quant = get_quant_config() stored = quant.stored_tensors(QuantKind.NVFP4) kind_map = {stored[role].name: kind for role, kind in _BANK_KINDS.items()} + weight_map = safetensors_weight_map(download_hf_weight(model_path)) + csf = any(name.endswith(".nvfp4_csf_fixed") for name in weight_map) + kinds = list(kind_map) + if csf: + kinds += list(_CSF_KINDS) + kind_map = {**kind_map, **{kind: kind for kind in _CSF_KINDS}} return Nvfp4ExpertSourceSpec( - key_pattern=re.compile(_EXPERT_KEY_RE.format(kinds="|".join(map(re.escape, kind_map)))), + key_pattern=re.compile(_EXPERT_KEY_RE.format(kinds="|".join(map(re.escape, kinds)))), proj_to_role={"gate_proj": "gate", "up_proj": "up", "down_proj": "down"}, layer_to_bank=lambda layer, config: layer, # every layer is MoE - desc=f"Qwen3.8-Flash-Next NVFP4 experts ({quant.dialect})", + desc=f"Qwen3.8-Flash-Next NVFP4 experts ({quant.dialect}{' CSF' if csf else ''})", kind_map=kind_map, global_reciprocal=stored["weight_global"].reciprocal, + scale_pair=_CSF_KINDS if csf else None, + scale_decoder=decode_csf_scales if csf else None, ) 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 diff --git a/python/freetoken/moe/expert_pieces.py b/python/freetoken/moe/expert_pieces.py index 74c4a41b0f..70a28ea3d9 100644 --- a/python/freetoken/moe/expert_pieces.py +++ b/python/freetoken/moe/expert_pieces.py @@ -142,16 +142,26 @@ def _bf16_pieces(model_path: str, config, spec, *, layers: Collection[int] | Non return stacked_expert_pieces(tensors, config, layers=layers) +def _shape_piece(piece: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + return { + role: tensor.reshape(1, -1) if role.endswith(("_global", "_input")) else tensor.unsqueeze(0) + for role, tensor in piece.items() + } + + def per_expert_pieces( tensors: Iterable[tuple[str, torch.Tensor]], locate: Callable[[str], tuple[int, int, str] | None], *, tensors_per_expert: int, + assemble: Callable[[int, int, dict[str, torch.Tensor]], dict[str, torch.Tensor]] | None = None, ) -> Iterator[Piece]: """Group per-expert tensors into one piece per expert, in whatever order they arrive. ``locate(name)`` -> ``(bank_layer, expert, role)`` or None to skip; a role tensor is one expert's ``[...]`` row and becomes the piece's ``[1, ...]`` entry. Globals and activation scales may be scalars. + ``assemble(layer, expert, raw)`` replaces the default shaping for sources whose roles need + more than an unsqueeze (e.g. decoding a stored scale pair back into the plain matrix). """ pending: dict[tuple[int, int], dict[str, torch.Tensor]] = {} for name, tensor in tensors: @@ -160,10 +170,11 @@ def per_expert_pieces( continue bank_layer, expert, role = where piece = pending.setdefault((bank_layer, expert), {}) - piece[role] = tensor.reshape(1, -1) if role.endswith(("_global", "_input")) else tensor.unsqueeze(0) + piece[role] = tensor if len(piece) == tensors_per_expert: del pending[(bank_layer, expert)] - yield bank_layer, expert, expert + 1, piece + shaped = _shape_piece(piece) if assemble is None else assemble(bank_layer, expert, piece) + yield bank_layer, expert, expert + 1, shaped if pending: missing = sorted(pending)[:8] raise ValueError(f"Incomplete expert tensors in checkpoint (layer, expert): {missing}") diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index 0d2461427e..5b83d47e03 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -224,6 +224,55 @@ 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).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)) + + 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] + ) + 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)]) + + @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 8381720740..091d7a2496 100644 --- a/tests/models/qwen4_exp/test_ple_disk.py +++ b/tests/models/qwen4_exp/test_ple_disk.py @@ -82,6 +82,24 @@ 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.""" + from safetensors.torch import save_file + + 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, dtype=torch.float8_e4m3fn): from freetoken.models.qwen4_exp.ple_disk import DiskRowTable, source_from_safetensors 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" diff --git a/tests/models/qwen4_exp/test_weight.py b/tests/models/qwen4_exp/test_weight.py index e066acbc2c..785ddbe18b 100644 --- a/tests/models/qwen4_exp/test_weight.py +++ b/tests/models/qwen4_exp/test_weight.py @@ -6,6 +6,7 @@ from __future__ import annotations +import dataclasses import json import random from types import SimpleNamespace @@ -16,6 +17,8 @@ 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 ( _ZERO_CENTERED_NORM_SUFFIXES, _DenseReader, @@ -65,6 +68,282 @@ 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, U32 bit patterns stored as int32 + for pos in (3, rows * cols - 1): + value = 0xFF - target.ravel()[pos].item() + 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: + ref.view(-1)[rec & 0xFFFFFF] = (rec >> 24) & 0xFF + 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) + # the same records as their true uint32 dtype (sign-bit records included) + assert torch.equal(decode_csf_scales(fixed, exc.to(torch.uint32)).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) + + +def test_csf_decode_rejects_streams_of_the_wrong_dtype(): + target = _csf_target(32, 8, seed=1) + fixed, exc, _ref = _csf_encode(target) + with pytest.raises(ValueError, match="U32 records"): + decode_csf_scales(fixed, exc.to(torch.float32)) + with pytest.raises(ValueError, match="expected U8"): + decode_csf_scales(fixed.to(torch.int32), exc) + + +def test_rename_scaled_maps_the_checkpoint_scale_names_to_the_model_roles(monkeypatch): + """FP4-Local quantises the vision tower: attention as MXFP8, the 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.visual.blocks.0.attn.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.visual.blocks.0.attn.qkv.weight_scale") + assert got == "visual.blocks.0.attn.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.visual.blocks.1.attn.qkv.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 fp32 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.float32 + 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 full-width global passes through untouched; a miscounted one raises + full = torch.ones(64, dtype=torch.float16) + assert g2.on_global(module + ".weight_global", full) == (module + ".weight_global", full) + row = torch.arange(3, dtype=torch.float16) + with pytest.raises(ValueError, match="holds 3 elements"): + g2.on_global(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(ValueError, match="without a weight"): + g3.check_complete() + + +def _csf_expert_raw() -> tuple[dict[str, torch.Tensor], dict[tuple[int, str], torch.Tensor]]: + """One CSF-quantised MoE layer: packed codes, slab scale pairs, globals and input scales.""" + 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 + 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 + return raw, refs + + +def test_nvfp4_expert_spec_selects_the_csf_checkpoint(tmp_path): + """A weight map carrying CSF slab pairs switches the spec onto the shared decoder.""" + quant = mixed_precision_quant(gdn_layers=(), attn_layers=(), moe_layers=(0,)) + raw, _refs = _csf_expert_raw() + (tmp_path / "csf").mkdir() + folder, _full = _write_checkpoint(tmp_path / "csf", raw, quant) + install_quant_config(folder) + spec = nvfp4_expert_spec(folder, None) + assert spec.scale_pair == qwen4_exp_weight._CSF_KINDS + assert spec.scale_decoder is qwen4_exp_weight.decode_csf_scales + assert "CSF" in spec.desc + plain = {k.replace(".weight_scale.nvfp4_csf_fixed", ".weight_scale"): v for k, v in raw.items()} + (tmp_path / "plain").mkdir() + folder, _full = _write_checkpoint(tmp_path / "plain", plain, quant) + install_quant_config(folder) + spec = nvfp4_expert_spec(folder, None) + assert spec.scale_pair is None and spec.scale_decoder is None + + +def test_csf_expert_reader_expands_the_scale_pieces(tmp_path): + """A CSF checkpoint streams slab pairs; the reader expands them into the plain bank roles.""" + from freetoken.models.loader import drop_page_cache + + lm = "model.language_model" + raw, refs = _csf_expert_raw() + quant = mixed_precision_quant(gdn_layers=(), attn_layers=(), moe_layers=(0,)) + (tmp_path / "reader").mkdir() + folder, _full = _write_checkpoint(tmp_path / "reader", raw, quant) + install_quant_config(folder) + spec = nvfp4_expert_spec(folder, None) + assert spec.scale_pair is not None + config = SimpleNamespace(num_experts=2, num_layers=1, first_k_dense_replace=0) + pieces = list(iter_nvfp4_expert_pieces(folder, config, 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.float32 and torch.equal( + piece[f"{role}_global"].reshape(-1), raw[f"{base}.weight_scale_2"].reshape(-1).float() + ) + assert piece[f"{role}_input"].dtype is torch.float32 and torch.equal( + piece[f"{role}_input"].reshape(-1), raw[f"{base}.input_scale"].reshape(-1).float() + ) + 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], refs[(e0, proj)]) + + +def test_csf_pieces_decoder_errors_surface_to_the_consumer(tmp_path): + """A raising scale decoder must fail the stream, not silently drop pieces.""" + raw, _refs = _csf_expert_raw() + quant = mixed_precision_quant(gdn_layers=(), attn_layers=(), moe_layers=(0,)) + folder, _full = _write_checkpoint(tmp_path, raw, quant) + install_quant_config(folder) + spec = nvfp4_expert_spec(folder, None) + calls = iter(range(99)) + + def boom(fixed, exceptions): + if next(calls) == 4: + raise RuntimeError("decoder boom") + return decode_csf_scales(fixed, exceptions) + + spec = dataclasses.replace(spec, scale_decoder=boom) + config = SimpleNamespace(num_experts=2, num_layers=1, first_k_dense_replace=0) + with pytest.raises(RuntimeError, match="decoder boom"): + list(iter_nvfp4_expert_pieces(folder, config, spec)) + + +def test_csf_expert_reader_parallel_streams_csf_pieces(tmp_path): + """The O_DIRECT parallel reader (the production path) expands CSF pairs identically.""" + raw, refs = _csf_expert_raw() + quant = mixed_precision_quant(gdn_layers=(), attn_layers=(), moe_layers=(0,)) + folder, _full = _write_checkpoint(tmp_path, raw, quant) + install_quant_config(folder) + spec = nvfp4_expert_spec(folder, None) + config = SimpleNamespace(num_experts=2, num_layers=1, first_k_dense_replace=0) + pieces = list(iter_nvfp4_expert_pieces(folder, config, spec, parallel=True, workers=2, chunk=1 << 20)) + assert sorted(p[1] for p in pieces) == [0, 1] + for _layer, e0, e1, piece in pieces: + assert (e0, e1) == (0, 1) or (e0, e1) == (1, 2) + for proj, role in (("gate_proj", "gate"), ("up_proj", "up"), ("down_proj", "down")): + base = f"model.language_model.layers.0.mlp.experts.{e0}.{proj}" + 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], refs[(e0, proj)]) + + +def test_csf_pieces_read_activation_scales_out_of_band(tmp_path, monkeypatch): + """input_scale rides no part of the expert stream. + + The FP4-Local checkpoint keeps every input_scale in one trailing shard; streaming + them made the first piece wait on that shard with all weights already in hand. + """ + from freetoken.models import weight as models_weight + + probed: list = [] + real = models_weight.iter_expert_tensors_parallel + + def spy(folder, is_expert, **kwargs): + probed.append(is_expert) + yield from real(folder, is_expert, **kwargs) + + monkeypatch.setattr(models_weight, "iter_expert_tensors_parallel", spy) + quant = mixed_precision_quant(gdn_layers=(), attn_layers=(), moe_layers=(0,)) + raw, _refs = _csf_expert_raw() + # mimic the real checkpoint: the whole input_scale set lives in its own shard + act_tensors = {name: raw.pop(name) for name in list(raw) if name.endswith(".input_scale")} + assert len(act_tensors) == 6 + folder, _full = _write_checkpoint(tmp_path, raw, quant) + save_file(act_tensors, str(tmp_path / "model-act-00000.safetensors")) + assert folder == str(tmp_path) + install_quant_config(folder) + spec = nvfp4_expert_spec(folder, None) + config = SimpleNamespace(num_experts=2, num_layers=1, first_k_dense_replace=0) + pieces = list(iter_nvfp4_expert_pieces(folder, config, spec, parallel=True, workers=2, chunk=1 << 20)) + assert len(pieces) == 2 + for _layer, _e0, _e1, piece in pieces: + for role in ("gate", "up", "down"): + assert piece[f"{role}_input"].dtype is torch.float32 + assert probed + want = probed[0] + for expert in range(2): + for proj in ("gate_proj", "up_proj", "down_proj"): + base = f"model.language_model.layers.0.mlp.experts.{expert}.{proj}" + assert not want(f"{base}.input_scale") + assert want(f"{base}.weight") + + def _raw_checkpoint(dense_fp8: bool = False, channel_fp8: bool = False, mxfp8_scale: str | None = None) -> dict[str, torch.Tensor]: """Layer 0 = GDN + PLE, layer 1 = QSA; plus the mtp / visual / routed-expert noise. @@ -426,6 +705,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, 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. + """ + 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="expected F8_E4M3"): + 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 # ======================================================================================