diff --git a/python/freetoken/checkpoint/__init__.py b/python/freetoken/checkpoint/__init__.py index f3ce70fc..f78d90ca 100644 --- a/python/freetoken/checkpoint/__init__.py +++ b/python/freetoken/checkpoint/__init__.py @@ -12,8 +12,16 @@ load_ftw_banks, ) from .convert import convert_checkpoint +from .q3_ple import ( + Q3PLEReader, + Q3PLESegment, + write_q3_ple_from_safetensors, + write_q3_ple_sidecar, +) __all__ = [ "FTWReader", "FTWWriter", "is_ftw_checkpoint", "iter_ftw_weights", "load_ftw_banks", "convert_checkpoint", + "Q3PLEReader", "Q3PLESegment", "write_q3_ple_sidecar", + "write_q3_ple_from_safetensors", ] diff --git a/python/freetoken/checkpoint/convert.py b/python/freetoken/checkpoint/convert.py index 2f643bca..106cbdb2 100644 --- a/python/freetoken/checkpoint/convert.py +++ b/python/freetoken/checkpoint/convert.py @@ -114,7 +114,9 @@ def _copy_host_mapped_weights(model_path: str, out_dir: str) -> list[str]: return copied -def _copy_metadata(model_path: str, out_dir: str) -> list[str]: +def _copy_metadata( + model_path: str, out_dir: str, *, include_host_mapped_weights: bool = True +) -> list[str]: """Copy all non-weight files (config, tokenizer, remote-code, nested model configs) preserving directory structure, so the FTW dir is a self-contained checkpoint.""" if os.path.isfile(model_path): @@ -148,10 +150,22 @@ def _copy_metadata(model_path: str, out_dir: str) -> list[str]: os.makedirs(os.path.dirname(dst), exist_ok=True) shutil.copy2(src, dst) copied.append(rel) - copied.extend(_copy_host_mapped_weights(model_path, out_dir)) + if include_host_mapped_weights: + copied.extend(_copy_host_mapped_weights(model_path, out_dir)) return copied +def _iter_qwen4_modular_dense_entries(entries): + """Apply the frozen active map and text-only filtering to a source stream.""" + from freetoken.models.config import VISION_KEY_PREFIXES + from freetoken.models.qwen4_exp.weight import iter_active_nvfp4_runtime_entries + + for name, tensor in iter_active_nvfp4_runtime_entries(entries): + if name.startswith(("visual.",) + VISION_KEY_PREFIXES): + continue + yield name, tensor + + class _ConvertSink: """Layer-completion sink for ``load_expert_banks(layer_sink=...)``: writes each completed layer's banks as their own FTW entries immediately (name @@ -217,6 +231,8 @@ def convert_checkpoint( moe_backend: str = "offload", shard_limit: int = DEFAULT_SHARD_LIMIT, device: str | None = None, + artifact_format: str | None = None, + source_inventory_sha256: str | None = None, ) -> dict: """Write ``model_path`` as an FTW checkpoint at ``out_dir``. Returns the index dict. @@ -239,26 +255,62 @@ def convert_checkpoint( f"FTW conversion runs single-process and the format records no TP layout, " f"but TP is already set to size={tp.size}" ) - dev = torch.device(device or "cuda:0") - torch.cuda.set_device(dev) - torch.zeros(1, device=dev) # init CUDA context (needed by nvfp4 backend pick / pinning) + if artifact_format not in (None, "qwen4_modular_v1"): + raise ValueError( + f"unsupported artifact_format {artifact_format!r}; expected None or 'qwen4_modular_v1'" + ) + # The modular target is pre-encoded entirely on CPU. It deliberately omits + # expert banks from this active FTW component, so initializing CUDA here would + # add an unnecessary conversion dependency and obscure the zero-VRAM envelope. + dev = torch.device("cpu" if artifact_format == "qwen4_modular_v1" else (device or "cuda:0")) + if dev.type == "cuda": + torch.cuda.set_device(dev) + torch.zeros(1, device=dev) # needed by legacy expert backend selection / pinning cfg = EngineConfig(model_path=model_path, tp_info=DistributedInfo(tp.rank, tp.size), dtype=dtype, moe_backend=moe_backend) mc = cfg.model_config - offload = moe_backend == "offload" and getattr(mc, "is_moe", False) - include_moe_experts = not offload + is_qwen4 = any("Qwen4" in str(arch) for arch in getattr(mc, "architectures", ())) + if artifact_format is not None and not is_qwen4: + raise ValueError("artifact_format='qwen4_modular_v1' requires a Qwen4 checkpoint") + modular = artifact_format == "qwen4_modular_v1" + if modular: + source_inventory_sha256 = str(source_inventory_sha256 or "").lower() + if len(source_inventory_sha256) != 64 or any( + char not in "0123456789abcdef" for char in source_inventory_sha256 + ): + raise ValueError( + "qwen4_modular_v1 conversion requires source_inventory_sha256" + ) + offload = not modular and moe_backend == "offload" and getattr(mc, "is_moe", False) + include_moe_experts = False if modular else not offload from freetoken.utils.progress import byte_bar, count_bar - writer = FTWWriter(out_dir, shard_limit=shard_limit) + # For the modular target ``out_dir`` is the artifact root; active FTW bytes + # live in their own component directory while config/tokenizer metadata stays + # at the root used by normal Engine startup. + active_out_dir = ( + os.path.join(out_dir, "qwen4-active-v1.ftw") if modular else out_dir + ) + writer = FTWWriter(active_out_dir, shard_limit=shard_limit) n_weight = n_bank = n_alpha = 0 # 1) dense weights (host tensors; load straight to CPU to avoid GPU pressure) _progress("dense", 0, 0) # phase start; per-tensor cumulative bytes follow (total unknown) dense_bytes = 0 - for name, tensor in count_bar(load_weight(model_path, torch.device("cpu"), - include_moe_experts=include_moe_experts), + dense_entries = load_weight( + model_path, + torch.device("cpu"), + include_moe_experts=include_moe_experts, + ) + if artifact_format == "qwen4_modular_v1": + # Quantization is an explicit artifact-build policy, never a generic runtime + # fallback. The Qwen4 iterator has already fused canonical projections; this + # wrapper only converts the frozen active map while leaving routers, PLE, and + # all non-active entries untouched. + dense_entries = _iter_qwen4_modular_dense_entries(dense_entries) + for name, tensor in count_bar(dense_entries, "Converting dense weights"): writer.add_tensor(name, tensor, kind="weight") n_weight += 1 @@ -324,7 +376,24 @@ def convert_checkpoint( bar.close() _progress("finalize") # writing shard index + copying config/tokenizer - copied = _copy_metadata(model_path, out_dir) + copied = _copy_metadata( + model_path, + out_dir, + include_host_mapped_weights=artifact_format != "qwen4_modular_v1", + ) + if artifact_format == "qwen4_modular_v1": + config_path = os.path.join(out_dir, "config.json") + if not os.path.isfile(config_path): + raise ValueError("Qwen4 modular conversion requires a copied config.json") + with open(config_path, "r", encoding="utf-8") as handle: + config_data = json.load(handle) + config_data["freetoken_text_only"] = "qwen4_text_only_v1" + config_data["freetoken_active_quant"] = "nvfp4_w4a16_v1" + tmp_config = config_path + ".tmp" + with open(tmp_config, "w", encoding="utf-8") as handle: + json.dump(config_data, handle, indent=2, sort_keys=True) + handle.write("\n") + os.replace(tmp_config, config_path) try: fingerprint = _source_fingerprint(model_path, mc, device=dev) @@ -334,6 +403,7 @@ def convert_checkpoint( index = writer.finalize({ "source_model_path": os.path.abspath(model_path), "fingerprint": fingerprint, + "source_inventory_sha256": source_inventory_sha256 if modular else None, # quant_format records the actual on-disk bank layout (e.g. nvfp4_marlin vs # nvfp4_b12x): the suffix is a runtime backend pick (GPU capability / env), NOT in # config, and the stored bytes are physically repacked into it -- so it's kept and diff --git a/python/freetoken/checkpoint/nvfp4.py b/python/freetoken/checkpoint/nvfp4.py new file mode 100644 index 00000000..906a912c --- /dev/null +++ b/python/freetoken/checkpoint/nvfp4.py @@ -0,0 +1,173 @@ +"""Deterministic host-side NVFP4 W4A16 encoding helpers. + +The native FreeToken dense NVFP4 operators consume three row-major tensors: + +* packed E2M1 codes (two low-bit-first nibbles per byte), +* one positive E4M3 scale for every 16 input values, and +* one FP16 positive global scale per output row. + +This module is intentionally CPU-safe and does not retain a BF16 copy. It is +used by metadata/conversion code and by synthetic component tests; runtime +operators continue to live in :mod:`freetoken.kernel.triton.nvfp4_linear`. +""" + +from __future__ import annotations + +import torch + + +# Keep this table in lock-step with the Triton/native dequant implementations. +# The unsigned codes are magnitudes; bit 3 is the sign bit. +E2M1_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) +E2M1_MAGNITUDES = torch.tensor(E2M1_VALUES, dtype=torch.float32) +E2M1_SIGNED = torch.tensor( + E2M1_VALUES + tuple(-v for v in E2M1_VALUES), dtype=torch.float32 +) +_FP16_MAX = float(torch.finfo(torch.float16).max) +_FP16_MIN_SUBNORMAL = 2.0 ** -24 +_E4M3_MAX = 448.0 +_E4M3_MIN_SUBNORMAL = 2.0 ** -9 + + +def _round_e2m1_rne(magnitude: torch.Tensor) -> torch.Tensor: + """Round non-negative values to E2M1 using the shared tie-to-even rule. + + Ties are resolved by the parity of the integer E2M1 code (for example, + 0.5/1.0 resolves to code 2, while 1.0/1.5 resolves to code 2). The + comparison is carried out in float64 so all BF16 inputs have deterministic + behavior at exact midpoints. + """ + + if torch.any(~torch.isfinite(magnitude)) or torch.any(magnitude < 0): + raise ValueError("E2M1 rounding expects finite non-negative values") + grid = E2M1_MAGNITUDES.to(device=magnitude.device, dtype=torch.float64) + x = magnitude.to(torch.float64).unsqueeze(-1) + distance = (x - grid).abs() + minimum = distance.min(dim=-1, keepdim=True).values + candidates = distance == minimum + # Prefer the even code among exact ties. Since candidates are at most two + # adjacent codes, selecting the last even candidate gives the desired rule. + codes = torch.arange(8, device=magnitude.device).expand_as(distance) + even = candidates & ((codes & 1) == 0) + picked = torch.where(even, codes, torch.full_like(codes, -1)).amax(dim=-1) + # Non-ties have no even candidate only for an impossible malformed grid; + # retain the nearest code as a defensive total fallback. + nearest = distance.argmin(dim=-1) + return torch.where(picked >= 0, picked, nearest).to(torch.uint8) + + +def _round_e4m3_positive(values: torch.Tensor) -> torch.Tensor: + """Encode finite non-negative values to E4M3 bytes with explicit bounds. + + E4M3's finite range is [0, 448]. Values above the finite range saturate + to 448 before the PyTorch cast (which otherwise produces the NaN sentinel), + and values below the representable subnormal range round to zero. This is + the documented, deterministic total policy for synthetic conversion. + """ + + if torch.any(~torch.isfinite(values)) or torch.any(values < 0): + raise ValueError("E4M3 scale encoding expects finite non-negative values") + bounded = values.to(torch.float32).clamp_(0.0, _E4M3_MAX) + # torch's CPU float8 conversion is round-to-nearest-even on the E4M3 grid. + return bounded.to(torch.float8_e4m3fn) + + +def encode_bf16_nvfp4(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Encode a BF16/FP16/FP32 matrix into the native row-major NVFP4 triple. + + Args: + weight: ``[out_features, in_features]`` finite real matrix. The input + width must be divisible by 16, matching ``Nvfp4DenseLinear``. + + Returns: + ``(packed, block_scale, global_scale)`` where packed is uint8 + ``[N,K//2]``, block_scale is native ``torch.float8_e4m3fn`` + ``[N,K//16]`` and + global_scale is FP16 ``[N]``. The returned tensors are newly allocated + and no copy of ``weight`` is retained. + + Scale rule: + ``global = round_fp16(max_abs / 6)`` (clamped to the finite FP16 range, + with the smallest FP16 subnormal used when a positive value would round + to zero); ``block = round_e4m3(max_abs_block / (6*global))``. A zero + row uses global=1 and zero block scales. Quantization then rounds each + value to E2M1 after dividing by ``global*block``. Zero block scales + produce zero codes. These explicit bounds make conversion total for + every finite input, including under/overflow extrema. + """ + + if weight.ndim != 2: + raise ValueError(f"NVFP4 encoder expects a rank-2 matrix, got {tuple(weight.shape)}") + if weight.shape[1] % 16: + raise ValueError(f"NVFP4 input width must be divisible by 16, got {weight.shape[1]}") + if not weight.dtype.is_floating_point: + raise TypeError(f"NVFP4 encoder expects a floating tensor, got {weight.dtype}") + if not torch.isfinite(weight).all(): + raise ValueError("NVFP4 encoder rejects NaN and infinity inputs") + + source = weight.to(dtype=torch.float32) + n_rows, width = source.shape + abs_source = source.abs() + row_max = abs_source.amax(dim=1) + nonzero = row_max > 0 + + # Rounding through one float16 conversion is intentional. Explicitly clamp + # the target first because a large BF16 row otherwise converts to inf. + global_target = (row_max / 6.0).clamp(_FP16_MIN_SUBNORMAL, _FP16_MAX) + global_target = torch.where(nonzero, global_target, torch.ones_like(global_target)) + global_scale = global_target.to(torch.float16) + # A positive target below the FP16 subnormal can still become zero on some + # CPU implementations; repair it explicitly and deterministically. + global_scale = torch.where( + nonzero & (global_scale == 0), + torch.full_like(global_scale, _FP16_MIN_SUBNORMAL, dtype=torch.float16), + global_scale, + ) + + blocks = source.view(n_rows, width // 16, 16) + block_max = blocks.abs().amax(dim=-1) + denom = global_scale.float().unsqueeze(-1) * 6.0 + block_target = torch.where(block_max > 0, block_max / denom, torch.zeros_like(block_max)) + block_scale = _round_e4m3_positive(block_target) + block_real = block_scale.view(torch.float8_e4m3fn).float() + + # Quantize against the *rounded* scales consumed by the kernel. Saturating + # normalized values to +/-6 is the finite E2M1 endpoint policy. + scale_real = global_scale.float().unsqueeze(-1).unsqueeze(-1) * block_real.unsqueeze(-1) + normalized = torch.where(scale_real > 0, blocks / scale_real, torch.zeros_like(blocks)) + magnitude = normalized.abs().clamp_(0.0, 6.0) + mag_code = _round_e2m1_rne(magnitude.reshape(-1)).view(n_rows, width // 16, 16) + sign = (normalized < 0).to(torch.uint8) + code = mag_code | (sign << 3) + # Two values per byte, low nibble first, as required by the native kernels. + packed = code.reshape(n_rows, width // 2, 2) + packed = packed[..., 0] | (packed[..., 1] << 4) + return packed.contiguous(), block_scale.contiguous(), global_scale.contiguous() + + +def decode_nvfp4( + packed: torch.Tensor, + block_scale: torch.Tensor, + global_scale: torch.Tensor, + *, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Reference dequantization for the native row-major NVFP4 triple.""" + + if packed.dtype != torch.uint8 or block_scale.dtype not in (torch.uint8, torch.float8_e4m3fn): + raise TypeError("packed must be uint8 and block_scale must be uint8-view or float8_e4m3fn") + if packed.ndim != 2 or block_scale.ndim != 2 or global_scale.ndim != 1: + raise ValueError("NVFP4 tensors must be packed[N,K/2], scale[N,K/16], global[N]") + rows, packed_width = packed.shape + width = packed_width * 2 + if block_scale.shape != (rows, width // 16) or global_scale.shape != (rows,): + raise ValueError("NVFP4 tensor shapes do not agree") + lo = packed & 0x0F + hi = packed >> 4 + codes = torch.stack((lo, hi), dim=-1).reshape(rows, width).to(torch.long) + values = E2M1_SIGNED.to(device=packed.device)[codes] + scales = block_scale.view(torch.float8_e4m3fn).float().repeat_interleave(16, dim=1) + return (values * scales * global_scale.float().unsqueeze(-1)).to(dtype) + + +__all__ = ["E2M1_VALUES", "encode_bf16_nvfp4", "decode_nvfp4"] diff --git a/python/freetoken/checkpoint/q3_ple.py b/python/freetoken/checkpoint/q3_ple.py new file mode 100644 index 00000000..cd62cc5b --- /dev/null +++ b/python/freetoken/checkpoint/q3_ple.py @@ -0,0 +1,718 @@ +"""Native reader for the Qwen4 ``Q3_PLE_32`` lookup-table sidecar. + +The reader deliberately knows only the PLE GET_ROWS format. It does not expose a +matmul/dequant path and it never maps or allocates the complete table. A small JSON +directory describes the logical rows and the byte ranges containing each segment; +the payload remains an ordinary read-only file on the project's required ``Z:`` +volume. +""" + +from __future__ import annotations + +import hashlib +import json +import math +import os +import operator +import struct +import threading +import uuid +from dataclasses import dataclass +from pathlib import Path +from typing import Iterable, Sequence + +import torch + + +BLOCK_VALUES = 32 +BLOCK_BYTES = 14 +ROW_VALUES = 160 +BLOCKS_PER_ROW = 5 +ROW_BYTES = BLOCKS_PER_ROW * BLOCK_BYTES +FORMAT = "q3_ple_32" +VERSION = 1 +ALIGN = 4096 +REFINEMENT_PASSES = 2 +DEFAULT_SEGMENT_ROWS = 128 + + +def _z_path(path: str | os.PathLike[str]) -> Path: + """Resolve *path* and fail closed unless its physical drive is ``Z:``.""" + + resolved = Path(path).expanduser().resolve(strict=True) + drive, _ = os.path.splitdrive(str(resolved)) + # On Windows splitdrive is authoritative. The second clause keeps the + # check useful in a POSIX test harness mounted as ``/z`` while still never + # accepting an unqualified relative path. + if drive.upper() != "Z:" and not str(resolved).lower().startswith("/z/"): + raise ValueError(f"Q3_PLE_32 backing must resolve to Z:, got {resolved}") + return resolved + + +def _z_output_path(path: str | os.PathLike[str]) -> Path: + """Resolve an output path and require an existing ``Z:`` parent directory. + + Unlike :func:`_z_path`, this helper permits the leaf file not to exist. The + writer intentionally does not create arbitrary parent directories: callers + must choose an already-created, Z-backed fixture or checkpoint directory. + """ + + candidate = Path(path).expanduser() + if not candidate.is_absolute(): + raise ValueError(f"Q3_PLE_32 output path must be absolute: {path}") + lexical_drive, _ = os.path.splitdrive(str(candidate)) + if lexical_drive.upper() != "Z:" and not str(candidate).lower().startswith("/z/"): + raise ValueError(f"Q3_PLE_32 output must resolve to Z:, got {candidate}") + parent = candidate.parent.resolve(strict=True) + resolved = parent / candidate.name + drive, _ = os.path.splitdrive(str(resolved)) + if drive.upper() != "Z:" and not str(resolved).lower().startswith("/z/"): + raise ValueError(f"Q3_PLE_32 output must resolve to Z:, got {resolved}") + return resolved + + +def _pack_codes(codes: Sequence[int]) -> bytes: + if len(codes) != BLOCK_VALUES: + raise ValueError(f"expected {BLOCK_VALUES} codes, got {len(codes)}") + packed = 0 + for index, code in enumerate(codes): + if not 0 <= code <= 7: + raise ValueError(f"code {index} is outside 0..7: {code}") + packed |= int(code) << (3 * index) + return packed.to_bytes(12, "little") + + +def _bf16_bits(value: float) -> int: + """Round a finite Python float to an IEEE BF16 bit pattern.""" + + try: + bits = struct.unpack("> 16) & 1) + return (rounded >> 16) & 0xFFFF + + +def _bf16_from_bits(bits: int) -> float: + return struct.unpack(" tuple[bytes, float]: + if not math.isfinite(value) or value < 0.0: + raise ValueError(f"Q3_PLE_32 scale must be finite and non-negative: {value!r}") + bits = _bf16_bits(value) + # A positive source scale must remain representable after BF16 storage. A + # zero source scale is reserved for an all-zero block. + if value > 0.0 and bits == 0: + bits = 1 + stored = _bf16_from_bits(bits) + if not math.isfinite(stored): + raise ValueError("Q3_PLE_32 stored BF16 scale is not finite") + return struct.pack(" list[int]: + if scale == 0.0: + return [4] * BLOCK_VALUES + return [max(-4, min(3, int(round(value / scale)))) + 4 for value in values] + + +def quantize_block(values: Sequence[float], *, refinement_passes: int = REFINEMENT_PASSES) -> bytes: + """Encode one 32-value block using the canonical Q3_PLE_32 recipe. + + This mirrors ``scripts/q3_ple_32_reference.py`` while keeping production + conversion independent of the repository's executable reference script. + """ + + if len(values) != BLOCK_VALUES: + raise ValueError(f"expected {BLOCK_VALUES} values, got {len(values)}") + if refinement_passes < 0: + raise ValueError("refinement_passes must be non-negative") + try: + source = [float(value) for value in values] + except (TypeError, ValueError) as exc: + raise ValueError("Q3_PLE_32 values must be numeric") from exc + if not all(math.isfinite(value) for value in source): + raise ValueError("Q3_PLE_32 cannot encode non-finite values") + + minimum = min(source) + maximum = max(source) + scale = max(-minimum / 4.0, maximum / 3.0) + if scale == 0.0: + scale_bytes, _ = _store_bf16_scale(0.0) + return scale_bytes + _pack_codes([4] * BLOCK_VALUES) + + codes = _codes_for_scale(source, scale) + for _ in range(refinement_passes): + quants = [code - 4 for code in codes] + denominator = sum(quant * quant for quant in quants) + if denominator == 0: + break + refined = sum(value * quant for value, quant in zip(source, quants)) / denominator + if refined <= 0.0 or not math.isfinite(refined): + break + new_codes = _codes_for_scale(source, refined) + scale = refined + if new_codes == codes: + codes = new_codes + break + codes = new_codes + + scale_bytes, stored_scale = _store_bf16_scale(scale) + # Requantize once against the stored BF16 value so decoding exactly follows + # runtime behavior rather than the pre-rounded Python scale. + codes = _codes_for_scale(source, stored_scale) + block = scale_bytes + _pack_codes(codes) + if len(block) != BLOCK_BYTES: + raise AssertionError(f"Q3_PLE_32 block has wrong size: {len(block)}") + return block + + +def quantize_row(values: Sequence[float], *, refinement_passes: int = REFINEMENT_PASSES) -> bytes: + """Encode a 160-value row as five canonical Q3_PLE_32 blocks.""" + + if len(values) != ROW_VALUES: + raise ValueError(f"expected {ROW_VALUES} values, got {len(values)}") + return b"".join( + quantize_block(values[offset : offset + BLOCK_VALUES], refinement_passes=refinement_passes) + for offset in range(0, ROW_VALUES, BLOCK_VALUES) + ) + + +def _unpack_codes(payload: bytes) -> list[int]: + if len(payload) != 12: + raise ValueError(f"Q3_PLE_32 code payload must be 12 bytes, got {len(payload)}") + bits = int.from_bytes(payload, "little") + return [(bits >> (3 * i)) & 0x7 for i in range(BLOCK_VALUES)] + + +def _decode_row(row: bytes) -> torch.Tensor: + if len(row) != ROW_BYTES: + raise ValueError(f"Q3_PLE_32 row must be {ROW_BYTES} bytes, got {len(row)}") + values: list[float] = [] + for block_start in range(0, ROW_BYTES, BLOCK_BYTES): + block = row[block_start : block_start + BLOCK_BYTES] + # The format authority specifies a little-endian BF16 scalar. Decode + # through an integer bit pattern so host endianness cannot leak in. + scale_bits = int.from_bytes(block[:2], "little") + scale = struct.unpack(" int: + return self.end_row - self.first_row + + +class Q3PLEReader: + """Bounded random-row reader for a validated Q3_PLE_32 sidecar. + + ``manifest_path`` points to ``ple-q3.json`` and ``data_path`` may override + its ``data_file``. Opening validates the JSON schema, segment coverage, + file length, and whole-file/segment hashes in bounded chunks. ``gather`` + then reads only the requested 70-byte rows and returns BF16 values. + """ + + def __init__(self, manifest_path: str | os.PathLike[str], *, data_path: str | os.PathLike[str] | None = None): + self.manifest_path = _z_path(manifest_path) + with self.manifest_path.open("r", encoding="utf-8") as handle: + manifest = json.load(handle) + self.manifest = manifest + if manifest.get("format") != FORMAT or int(manifest.get("version", -1)) != VERSION: + raise ValueError("unsupported Q3_PLE_32 format/version") + if manifest.get("endianness", "little") != "little": + raise ValueError("Q3_PLE_32 requires little-endian metadata") + if int(manifest.get("block_values", BLOCK_VALUES)) != BLOCK_VALUES: + raise ValueError("Q3_PLE_32 block_values mismatch") + if int(manifest.get("block_bytes", BLOCK_BYTES)) != BLOCK_BYTES: + raise ValueError("Q3_PLE_32 block_bytes mismatch") + if int(manifest.get("row_values", ROW_VALUES)) != ROW_VALUES: + raise ValueError("Q3_PLE_32 row_values mismatch") + if int(manifest.get("row_bytes", ROW_BYTES)) != ROW_BYTES: + raise ValueError("Q3_PLE_32 row_bytes mismatch") + # Older Stage 6 fixtures predate this field, so absence remains + # readable. A present fingerprint is always canonical SHA-256 hex; + # malformed provenance must fail closed rather than being ignored. + if "source_fingerprint" in manifest: + _validate_source_fingerprint(manifest["source_fingerprint"]) + + candidate = data_path + if candidate is None: + candidate = self.manifest.get("data_file") + if not candidate: + raise ValueError("Q3_PLE_32 manifest has no data_file") + data_candidate = Path(candidate) + if not data_candidate.is_absolute(): + data_candidate = self.manifest_path.parent / data_candidate + self.data_path = _z_path(data_candidate) + self.row_count = int(manifest.get("rows", 0)) + self.total_payload_bytes = self.row_count * ROW_BYTES + if self.row_count <= 0: + raise ValueError("Q3_PLE_32 rows must be positive") + declared_payload = int(manifest.get("payload_bytes", self.total_payload_bytes)) + if declared_payload != self.total_payload_bytes: + raise ValueError("Q3_PLE_32 payload_bytes does not equal rows * row_bytes") + + raw_segments = manifest.get("segments") + if not isinstance(raw_segments, list) or not raw_segments: + raise ValueError("Q3_PLE_32 segment directory is empty") + self.segments: tuple[Q3PLESegment, ...] = tuple(self._parse_segment(item) for item in raw_segments) + self._validate_segments() + stat = self.data_path.stat() + expected_file_bytes = int(manifest.get("file_bytes", stat.st_size)) + if stat.st_size != expected_file_bytes: + raise ValueError(f"Q3_PLE_32 file length mismatch: {stat.st_size} != {expected_file_bytes}") + self._handle = self.data_path.open("rb") + self._io_lock = threading.Lock() + try: + self._verify_hashes() + except Exception: + self._handle.close() + raise + + self.weight_scale = float(manifest.get("weight_scale", 1.0)) + if not math.isfinite(self.weight_scale): + raise ValueError("Q3_PLE_32 weight_scale must be finite") + + def _parse_segment(self, item: object) -> Q3PLESegment: + if not isinstance(item, dict): + raise ValueError("Q3_PLE_32 segment must be an object") + try: + data_offset = item.get("data_offset") + if data_offset is None: + data_offset = item["offset"] + byte_length = item.get("byte_length") + if byte_length is None: + byte_length = item["length"] + segment = Q3PLESegment( + first_row=int(item["first_row"]), + end_row=int(item["end_row"]), + data_offset=int(data_offset), + byte_length=int(byte_length), + sha256=str(item["sha256"]).lower(), + ) + except (KeyError, TypeError, ValueError) as exc: + raise ValueError("malformed Q3_PLE_32 segment directory") from exc + if len(segment.sha256) != 64 or any(c not in "0123456789abcdef" for c in segment.sha256): + raise ValueError("malformed Q3_PLE_32 segment hash") + return segment + + def _validate_segments(self) -> None: + expected_row = 0 + previous_end = 0 + contiguous = self.manifest.get("storage_layout") == "contiguous_rows_v1" + for segment in self.segments: + if segment.first_row != expected_row or segment.end_row <= segment.first_row: + raise ValueError("Q3_PLE_32 segment rows have a gap, overlap, or bad order") + if segment.data_offset < 0: + raise ValueError("Q3_PLE_32 segment data offset is negative") + if contiguous and segment.data_offset != previous_end: + raise ValueError("Q3_PLE_32 contiguous segment directory has a gap or overlap") + if not contiguous and segment.data_offset % ALIGN: + raise ValueError("Q3_PLE_32 segment data offset is not 4 KiB aligned") + if segment.byte_length != segment.rows * ROW_BYTES: + raise ValueError("Q3_PLE_32 segment byte length does not match rows") + if segment.data_offset < previous_end: + raise ValueError("Q3_PLE_32 segment byte ranges overlap") + expected_row = segment.end_row + previous_end = segment.data_offset + segment.byte_length + if expected_row != self.row_count: + raise ValueError("Q3_PLE_32 segment rows do not cover the table") + if previous_end > int(self.manifest.get("file_bytes", previous_end)): + raise ValueError("Q3_PLE_32 segment exceeds file length") + + def _read_exact(self, offset: int, length: int) -> bytes: + with self._io_lock: + self._handle.seek(offset) + payload = self._handle.read(length) + if len(payload) != length: + raise OSError(f"short Q3_PLE_32 read at {offset}: {len(payload)}/{length}") + return payload + + def _verify_hashes(self) -> None: + # Hashes are checked incrementally; this never allocates the table. + whole = hashlib.sha256() + with self.data_path.open("rb") as handle: + while True: + chunk = handle.read(8 << 20) + if not chunk: + break + whole.update(chunk) + expected_whole = str(self.manifest.get("sha256", "")).lower() + if len(expected_whole) != 64 or whole.hexdigest() != expected_whole: + raise ValueError("Q3_PLE_32 whole-file hash mismatch") + for segment in self.segments: + digest_ctx = hashlib.sha256() + remaining = segment.byte_length + offset = segment.data_offset + while remaining: + take = min(8 << 20, remaining) + digest_ctx.update(self._read_exact(offset, take)) + offset += take + remaining -= take + digest = digest_ctx.hexdigest() + if digest != segment.sha256: + raise ValueError(f"Q3_PLE_32 segment hash mismatch at row {segment.first_row}") + + def _segment_for_row(self, row: int) -> Q3PLESegment: + if not 0 <= row < self.row_count: + raise IndexError(f"Q3_PLE_32 row {row} outside 0..{self.row_count - 1}") + # Segment count is small (128 in the production sidecar); linear search + # keeps the directory representation transparent and deterministic. + for segment in self.segments: + if segment.first_row <= row < segment.end_row: + return segment + raise AssertionError("validated Q3_PLE_32 directory did not locate row") + + def read_row(self, row: int) -> torch.Tensor: + segment = self._segment_for_row(int(row)) + offset = segment.data_offset + (int(row) - segment.first_row) * ROW_BYTES + return _decode_row(self._read_exact(offset, ROW_BYTES)) + + def gather(self, row_indices: Sequence[int], *, apply_weight_scale: bool = False) -> torch.Tensor: + """Gather rows in the caller's order without deduplication or reordering.""" + + rows = [self.read_row(int(index)) for index in row_indices] + if not rows: + output = torch.empty((0, ROW_VALUES), dtype=torch.bfloat16) + else: + output = torch.stack(rows, dim=0) + if apply_weight_scale: + output = output * self.weight_scale + return output + + def gather16(self, row_indices: Sequence[int], *, apply_weight_scale: bool = False) -> torch.Tensor: + if len(row_indices) != 16: + raise ValueError(f"Qwen4 PLE requires exactly 16 logical rows, got {len(row_indices)}") + return self.gather(row_indices, apply_weight_scale=apply_weight_scale) + + def close(self) -> None: + handle, self._handle = getattr(self, "_handle", None), None + if handle is not None: + handle.close() + + def __enter__(self) -> "Q3PLEReader": + return self + + def __exit__(self, exc_type, exc_value, traceback) -> None: + self.close() + + +def _align_up(value: int, alignment: int = ALIGN) -> int: + return (value + alignment - 1) // alignment * alignment + + +def _validate_source_fingerprint(source_fingerprint: str) -> str: + if not isinstance(source_fingerprint, str): + raise ValueError("source_fingerprint must be a 64-character SHA-256 hex string") + fingerprint = source_fingerprint.lower() + if len(fingerprint) != 64 or any(char not in "0123456789abcdef" for char in fingerprint): + raise ValueError("source_fingerprint must be a 64-character SHA-256 hex string") + return fingerprint + + +def _materialize_row(row: object) -> list[float]: + """Materialize one bounded row without retaining any other source rows.""" + + if isinstance(row, torch.Tensor): + if row.ndim != 1 or row.numel() != ROW_VALUES: + raise ValueError(f"Q3_PLE_32 row must contain exactly {ROW_VALUES} values") + try: + values = row.detach().cpu().tolist() + except Exception as exc: + raise ValueError("Q3_PLE_32 row tensor could not be copied to CPU") from exc + else: + try: + values = list(row) # type: ignore[arg-type] + except (TypeError, ValueError) as exc: + raise ValueError(f"Q3_PLE_32 row must contain exactly {ROW_VALUES} values") from exc + if len(values) != ROW_VALUES: + raise ValueError(f"Q3_PLE_32 row must contain exactly {ROW_VALUES} values, got {len(values)}") + return values + + +def _partial_path(path: Path, token: str) -> Path: + return path.with_name(f".{path.name}.partial-{os.getpid()}-{threading.get_ident()}-{token}") + + +def _validate_segment_directory( + segments: Sequence[dict[str, int | str]], row_count: int, file_bytes: int +) -> None: + expected_row = 0 + previous_end = 0 + for segment in segments: + first_row = int(segment["first_row"]) + end_row = int(segment["end_row"]) + data_offset = int(segment["data_offset"]) + byte_length = int(segment["byte_length"]) + digest = str(segment["sha256"]) + if first_row != expected_row or end_row <= first_row: + raise ValueError("Q3_PLE_32 writer generated a malformed segment directory") + if data_offset != previous_end: + raise ValueError("Q3_PLE_32 writer generated a non-contiguous segment") + if byte_length != (end_row - first_row) * ROW_BYTES: + raise ValueError("Q3_PLE_32 writer generated a segment length mismatch") + if data_offset + byte_length > file_bytes: + raise ValueError("Q3_PLE_32 writer generated overlapping/out-of-range segments") + if len(digest) != 64 or any(char not in "0123456789abcdef" for char in digest): + raise ValueError("Q3_PLE_32 writer generated a malformed segment hash") + expected_row = end_row + previous_end = data_offset + byte_length + if not segments or expected_row != row_count: + raise ValueError("Q3_PLE_32 writer generated incomplete segment coverage") + + +def _fsync(handle: object) -> None: + # This helper exists to keep the finalize path explicit and easy to audit; + # the writer only passes ordinary binary file handles here. + file_handle = handle # type narrowing for type checkers without a runtime dependency + file_handle.flush() # type: ignore[attr-defined] + os.fsync(file_handle.fileno()) # type: ignore[attr-defined] + + +def write_q3_ple_sidecar( + rows: Iterable[Sequence[float] | torch.Tensor], + data_path: str | os.PathLike[str], + manifest_path: str | os.PathLike[str], + *, + source_fingerprint: str, + weight_scale: float, + segment_rows: int = DEFAULT_SEGMENT_ROWS, +) -> dict: + """Stream rows into an atomic, reader-compatible Q3_PLE_32 sidecar. + + ``rows`` is consumed exactly once and only the current 160-value row is held + in memory. Segments are logical hash/addressing ranges over one contiguous + row stream. There is no inter-row or inter-segment padding: the production + table is exactly ``rows * 70`` bytes. Both files are written under unique partial + names, fsynced, and atomically renamed into place on successful completion. + + The converter rule is intentionally fixed at two least-squares refinement + passes (the provisional Q3_PLE_32 recipe), and ``weight_scale`` is metadata + applied by the runtime after row dequantization rather than folded into the + per-block scales. + """ + + data_final = _z_output_path(data_path) + manifest_final = _z_output_path(manifest_path) + if data_final == manifest_final: + raise ValueError("Q3_PLE_32 data_path and manifest_path must differ") + source_digest = _validate_source_fingerprint(source_fingerprint) + if isinstance(segment_rows, bool): + raise ValueError("segment_rows must be a positive integer") + try: + segment_size = operator.index(segment_rows) + except TypeError as exc: + raise ValueError("segment_rows must be a positive integer") from exc + if segment_size <= 0: + raise ValueError("segment_rows must be a positive integer") + try: + global_scale = float(weight_scale) + except (TypeError, ValueError) as exc: + raise ValueError("weight_scale must be finite") from exc + if not math.isfinite(global_scale): + raise ValueError("weight_scale must be finite") + + # Unique tokens make concurrent conversion attempts independent and avoid + # ever truncating a stale partial file left by an interrupted process. + token = uuid.uuid4().hex + data_partial = _partial_path(data_final, token) + manifest_partial = _partial_path(manifest_final, token) + segments: list[dict[str, int | str]] = [] + whole_digest = hashlib.sha256() + payload_digest = hashlib.sha256() + rows_written = 0 + file_offset = 0 + current_segment: dict[str, int | str] | None = None + segment_digest: hashlib._Hash | None = None + + try: + with data_partial.open("wb") as output: + for source_row in rows: + row_values = _materialize_row(source_row) + encoded_row = quantize_row(row_values, refinement_passes=REFINEMENT_PASSES) + if len(encoded_row) != ROW_BYTES: + raise AssertionError(f"Q3_PLE_32 row has wrong size: {len(encoded_row)}") + + if rows_written % segment_size == 0: + if current_segment is not None: + assert segment_digest is not None + current_segment["end_row"] = rows_written + current_segment["byte_length"] = ( + rows_written * ROW_BYTES - int(current_segment["first_row"]) * ROW_BYTES + ) + current_segment["sha256"] = segment_digest.hexdigest() + segments.append(current_segment) + current_segment = { + "first_row": rows_written, + "end_row": rows_written, + "data_offset": file_offset, + "byte_length": 0, + "sha256": "", + } + segment_digest = hashlib.sha256() + + assert current_segment is not None and segment_digest is not None + output.write(encoded_row) + whole_digest.update(encoded_row) + payload_digest.update(encoded_row) + segment_digest.update(encoded_row) + file_offset += len(encoded_row) + rows_written += 1 + + if current_segment is None: + raise ValueError("Q3_PLE_32 rows must contain at least one row") + assert segment_digest is not None + current_segment["end_row"] = rows_written + current_segment["byte_length"] = ( + rows_written * ROW_BYTES - int(current_segment["first_row"]) * ROW_BYTES + ) + current_segment["sha256"] = segment_digest.hexdigest() + segments.append(current_segment) + _fsync(output) + except Exception: + data_partial.unlink(missing_ok=True) + manifest_partial.unlink(missing_ok=True) + raise + + try: + file_bytes = data_partial.stat().st_size + except Exception: + data_partial.unlink(missing_ok=True) + raise + if file_bytes != file_offset: + data_partial.unlink(missing_ok=True) + raise OSError(f"Q3_PLE_32 partial length mismatch: {file_bytes} != {file_offset}") + try: + _validate_segment_directory(segments, rows_written, file_bytes) + except Exception: + data_partial.unlink(missing_ok=True) + raise + manifest = { + "format": FORMAT, + "version": VERSION, + "endianness": "little", + "block_values": BLOCK_VALUES, + "block_bytes": BLOCK_BYTES, + "row_values": ROW_VALUES, + "row_bytes": ROW_BYTES, + "rows": rows_written, + "payload_bytes": rows_written * ROW_BYTES, + "file_bytes": file_bytes, + "storage_layout": "contiguous_rows_v1", + "data_file": os.path.relpath(data_final, manifest_final.parent), + "weight_scale": global_scale, + "source_fingerprint": source_digest, + "sha256": whole_digest.hexdigest(), + "payload_sha256": payload_digest.hexdigest(), + "segments": segments, + } + try: + with manifest_partial.open("w", encoding="utf-8", newline="\n") as manifest_handle: + json.dump(manifest, manifest_handle, ensure_ascii=False, indent=2, sort_keys=True) + manifest_handle.write("\n") + _fsync(manifest_handle) + os.replace(data_partial, data_final) + os.replace(manifest_partial, manifest_final) + except Exception: + data_partial.unlink(missing_ok=True) + manifest_partial.unlink(missing_ok=True) + raise + return manifest + + +def write_q3_ple_from_safetensors( + model_path: str | os.PathLike[str], + data_path: str | os.PathLike[str], + manifest_path: str | os.PathLike[str], + *, + layer_id: int, + split_parts: int, + source_fingerprint: str, + rows_per_chunk: int = 8192, + segment_rows: int = DEFAULT_SEGMENT_ROWS, +) -> dict: + """Stream the official FP8 PLE shards into the native Q3 sidecar. + + Shards and rows are consumed in exact ``shard_0..shard_N`` order. A + Safetensors slice is read in bounded row chunks; the full 51.2-GiB table is + never materialized. The source per-model ``weight_scale`` remains a separate + scalar in the Q3 manifest and is not folded into block scales. + """ + + folder = Path(model_path).expanduser().resolve() + if not folder.is_dir(): + raise ValueError(f"Q3 PLE source must be a local checkpoint directory: {folder}") + if rows_per_chunk <= 0 or split_parts <= 0: + raise ValueError("rows_per_chunk and split_parts must be positive") + index_path = folder / "model.safetensors.index.json" + with index_path.open("r", encoding="utf-8") as handle: + weight_map = json.load(handle)["weight_map"] + prefix = ( + f"model.language_model.layers.{int(layer_id)}.ple.ple_embedding." + "ngram_embedding" + ) + shard_keys = [f"{prefix}.shard_{part}.weight" for part in range(int(split_parts))] + missing = [key for key in shard_keys if key not in weight_map] + scale_key = prefix + ".weight_scale" + if missing or scale_key not in weight_map: + raise ValueError(f"incomplete PLE source mapping under {prefix}") + + import safetensors + + scale_file = folder / weight_map[scale_key] + with safetensors.safe_open(scale_file, framework="pt", device="cpu") as handle: + scale = handle.get_tensor(scale_key).reshape(()) + weight_scale = float(scale.float().item()) + + def iter_rows(): + for key in shard_keys: + source_file = folder / weight_map[key] + with safetensors.safe_open(source_file, framework="pt", device="cpu") as handle: + sliced = handle.get_slice(key) + shape = tuple(int(value) for value in sliced.get_shape()) + if len(shape) != 2 or shape[1] != ROW_VALUES: + raise ValueError(f"unexpected PLE source shape for {key}: {shape}") + for start in range(0, shape[0], int(rows_per_chunk)): + chunk = sliced[start : min(start + int(rows_per_chunk), shape[0])] + if chunk.dtype != torch.float8_e4m3fn: + raise ValueError(f"unexpected PLE source dtype for {key}: {chunk.dtype}") + for row in chunk.float(): + yield row + + return write_q3_ple_sidecar( + iter_rows(), + data_path, + manifest_path, + source_fingerprint=source_fingerprint, + weight_scale=weight_scale, + segment_rows=segment_rows, + ) + + +__all__ = [ + "ALIGN", + "BLOCK_BYTES", + "BLOCK_VALUES", + "DEFAULT_SEGMENT_ROWS", + "FORMAT", + "REFINEMENT_PASSES", + "Q3PLEReader", + "Q3PLESegment", + "ROW_BYTES", + "ROW_VALUES", + "VERSION", + "quantize_block", + "quantize_row", + "write_q3_ple_sidecar", + "write_q3_ple_from_safetensors", +] diff --git a/python/freetoken/checkpoint/qwen4_artifact.py b/python/freetoken/checkpoint/qwen4_artifact.py new file mode 100644 index 00000000..4e209628 --- /dev/null +++ b/python/freetoken/checkpoint/qwen4_artifact.py @@ -0,0 +1,900 @@ +"""Validation and runtime wiring for the Qwen4 modular artifact. + +The modular artifact is intentionally a small manifest around three independently +validated pieces: native active weights, a Q3 PLE sidecar, and a mixed resident/file +expert tier. This module owns only the manifest contract and the wiring seam. The +large writers and the byte-level readers live in their existing modules. + +An absent ``manifest.json`` is not an error and keeps the normal Qwen4 checkpoint path +unchanged. Once the marker is present, malformed or unknown values fail closed rather +than silently falling back to a different representation. +""" + +from __future__ import annotations + +import json +import hashlib +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Mapping + + +FORMAT = "freetoken-qwen4-modular-v1" +VERSION = 1 +TEXT_ONLY_MARKER = "qwen4_text_only_v1" +ARTIFACT_FORMAT = "qwen4_modular_v1" +ACTIVE_FORMAT = "nvfp4_w4a16_v1" +PLE_FORMAT = "q3_ple_32" +EXPERT_FORMAT = "ftexpert1_nvfp4_v1" +REQUIRED_VOLUME = "Z:" +MANIFEST_NAME = "manifest.json" +ACTIVE_TARGET_BYTES = 4_804_403_200 +PLE_TARGET_BYTES = 22_400_107_520 +EXPERT_FILE_BYTES = 1_419_776_000 +EXPERT_LAYERS = 48 +EXPERT_NUM_EXPERTS = 512 +KNOWN_TARGET_BYTES = 95_353_758_720 +FILE_TIER_LAYERS = (0, 1, 2, 3, 4, 5, 42, 43, 44, 45, 46, 47) +PINNED_SOURCE_REPOSITORY = "RadixArk/Qwen3.8-Flash-Next-NVFP4" +PINNED_SOURCE_REVISION = "7b719225242aacd3dbd3f9407468c2ee9a9d2594" +TVM_FFI_PATCH_SHA256 = "889310b8152a147a6552a3e451b3251a7df70cdc8e6e4c1c87c7adf3854182ec" + + +class Qwen4ArtifactError(ValueError): + """Raised when a Qwen4 modular manifest cannot be trusted.""" + + +def _resolve_z(path: str | os.PathLike[str], *, label: str) -> Path: + resolved = Path(path).expanduser().resolve() + drive = (resolved.drive or os.path.splitdrive(str(resolved))[0]).upper() + if drive != REQUIRED_VOLUME: + raise Qwen4ArtifactError(f"{label} must resolve to Z:, got {resolved}") + return resolved + + +def _path_from(root: Path, value: object, *, label: str) -> Path: + if not isinstance(value, str) or not value.strip(): + raise Qwen4ArtifactError(f"{label} must be a non-empty path") + candidate = Path(value) + if not candidate.is_absolute(): + candidate = root / candidate + return _resolve_z(candidate, label=label) + + +def _path_within_root(root: Path, value: object, *, label: str) -> Path: + path = _path_from(root, value, label=label) + try: + path.relative_to(root) + except ValueError as exc: + raise Qwen4ArtifactError(f"{label} resolves outside artifact root") from exc + return path + + +def _require_mapping(value: object, *, label: str) -> Mapping[str, Any]: + if not isinstance(value, Mapping): + raise Qwen4ArtifactError(f"{label} must be an object") + return value + + +def _require_sha(value: object, *, label: str) -> str: + text = str(value).lower() + if len(text) != 64 or any(char not in "0123456789abcdef" for char in text): + raise Qwen4ArtifactError(f"{label} must be a SHA-256 hex digest") + return text + + +@dataclass(frozen=True) +class ExpertFile: + layer: int + path: Path + bytes: int + sha256: str + source_fingerprint: str + + +@dataclass(frozen=True) +class ComponentFile: + path: Path + bytes: int + sha256: str + + +def _sha256_file(path: Path, *, chunk_bytes: int = 8 << 20) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(chunk_bytes): + digest.update(chunk) + return digest.hexdigest() + + +def _verify_component_file(entry: ComponentFile, *, label: str) -> None: + if not entry.path.is_file(): + raise Qwen4ArtifactError(f"{label} is missing: {entry.path}") + actual_bytes = entry.path.stat().st_size + if actual_bytes != entry.bytes: + raise Qwen4ArtifactError( + f"{label} length mismatch: {actual_bytes} != {entry.bytes}" + ) + actual_sha = _sha256_file(entry.path) + if actual_sha != entry.sha256: + raise Qwen4ArtifactError(f"{label} SHA-256 mismatch") + + +def _component_file(root: Path, path: Path) -> dict[str, Any]: + resolved = _resolve_z(path, label="artifact component") + try: + relative = resolved.relative_to(root) + except ValueError as exc: + raise Qwen4ArtifactError(f"artifact component resolves outside root: {resolved}") from exc + return { + "path": relative.as_posix(), + "bytes": resolved.stat().st_size, + "sha256": _sha256_file(resolved), + } + + +@dataclass(frozen=True) +class Qwen4ArtifactManifest: + """Validated view of ``manifest.json``. + + ``raw`` is retained for provenance and future additive fields. Paths are absolute + Z: paths so callers never accidentally resolve a relative sidecar against cwd. + """ + + path: Path + raw: Mapping[str, Any] + active_path: Path + active_files: tuple[ComponentFile, ...] + ple_manifest_path: Path + ple_data_bytes: int + ple_sha256: str + expert_files: tuple[ExpertFile, ...] + metadata_files: tuple[ComponentFile, ...] + file_tier_layers: tuple[int, ...] + resident_layers: tuple[int, ...] + production_geometry: bool + + @property + def root(self) -> Path: + return self.path.parent + + @property + def manifest_path(self) -> Path: + return self.path + + def __getitem__(self, key: str): + return self.raw[key] + + def get(self, key: str, default=None): + return self.raw.get(key, default) + + @property + def text_only(self) -> bool: + return bool(self.raw.get("text_only", False)) + + @property + def active_format(self) -> str: + return str(_require_mapping(self.raw.get("active"), label="active")["format"]) + + @property + def source(self) -> Mapping[str, Any]: + return _require_mapping(self.raw.get("source", {}), label="source") + + @property + def active(self) -> Mapping[str, Any]: + return _require_mapping(self.raw.get("active"), label="active") + + @property + def ple(self) -> Mapping[str, Any]: + return _require_mapping(self.raw.get("ple"), label="ple") + + @property + def experts(self) -> Mapping[str, Any]: + return _require_mapping(self.raw.get("experts"), label="experts") + + @property + def expert_format(self) -> str: + return str(_require_mapping(self.raw.get("experts"), label="experts")["format"]) + + @property + def num_layers(self) -> int: + layers = set(self.file_tier_layers) | set(self.resident_layers) + return max(layers) + 1 if layers else 0 + + def file_for_layer(self, layer_id: int) -> ExpertFile: + for entry in self.expert_files: + if entry.layer == int(layer_id): + return entry + raise Qwen4ArtifactError(f"manifest has no expert file for layer {layer_id}") + + def verify_active(self) -> None: + for index, entry in enumerate(self.active_files): + _verify_component_file(entry, label=f"active.files[{index}]") + from freetoken.checkpoint.ftw import INDEX_NAME + + with (self.active_path / INDEX_NAME).open("r", encoding="utf-8") as handle: + index = json.load(handle) + expected = str(self.source["inventory_sha256"]).lower() + if str(index.get("source_inventory_sha256", "")).lower() != expected: + raise Qwen4ArtifactError("active FTW source inventory fingerprint mismatch") + + +def _validate_layers(value: object, *, label: str) -> tuple[int, ...]: + if not isinstance(value, list): + raise Qwen4ArtifactError(f"experts.{label} must be a list") + result = [] + for item in value: + if isinstance(item, bool): + raise Qwen4ArtifactError(f"experts.{label} contains a non-integer layer") + try: + layer = int(item) + except (TypeError, ValueError) as exc: + raise Qwen4ArtifactError(f"experts.{label} contains a non-integer layer") from exc + if layer < 0 or layer in result: + raise Qwen4ArtifactError(f"experts.{label} contains an invalid/duplicate layer") + result.append(layer) + return tuple(result) + + +def _read_manifest_path(model_path: str | os.PathLike[str]) -> Path | None: + candidate = Path(model_path) + if candidate.name == MANIFEST_NAME and candidate.is_file(): + return _resolve_z(candidate, label="Qwen4 modular manifest") + if candidate.is_dir(): + path = candidate / MANIFEST_NAME + if path.is_file(): + return _resolve_z(path, label="Qwen4 modular manifest") + return None + + +def load_qwen4_artifact_manifest( + model_path: str | os.PathLike[str], *, require: bool = False, + allow_synthetic_geometry: bool = False, +) -> Qwen4ArtifactManifest | None: + """Load and validate a Qwen4 modular manifest. + + ``None`` means no manifest is present. This is the compatibility path for all + unmarked checkpoints. ``require=True`` is useful at a marked call site where a + missing manifest must not silently fall back to source weights. + """ + + path = _read_manifest_path(model_path) + if path is None: + if require: + raise Qwen4ArtifactError(f"Qwen4 modular manifest missing under {model_path}") + return None + try: + with path.open("r", encoding="utf-8") as handle: + raw = json.load(handle) + except (OSError, json.JSONDecodeError) as exc: + raise Qwen4ArtifactError(f"cannot read Qwen4 modular manifest {path}") from exc + root = _require_mapping(raw, label="Qwen4 modular manifest") + if root.get("format") != FORMAT or int(root.get("version", -1)) != VERSION: + raise Qwen4ArtifactError("unsupported Qwen4 modular manifest format/version") + if root.get("artifact_schema") != FORMAT: + raise Qwen4ArtifactError("unsupported Qwen4 modular artifact_schema") + if root.get("text_only") is not True: + raise Qwen4ArtifactError("Qwen4 modular manifest must declare text_only=true") + source = _require_mapping(root.get("source"), label="source") + if not str(source.get("repository", "")).strip() or not str(source.get("revision", "")).strip(): + raise Qwen4ArtifactError("source repository and revision are required") + source_inventory_sha256 = _require_sha( + source.get("inventory_sha256"), label="source.inventory_sha256" + ) + minimum_commit = str(root.get("minimum_freetoken_commit", "")).lower() + if len(minimum_commit) != 40 or any(char not in "0123456789abcdef" for char in minimum_commit): + raise Qwen4ArtifactError("minimum_freetoken_commit must be a 40-character Git OID") + _require_sha(root.get("tvm_ffi_patch_sha256"), label="tvm_ffi_patch_sha256") + if not allow_synthetic_geometry: + if source.get("repository") != PINNED_SOURCE_REPOSITORY: + raise Qwen4ArtifactError("source repository does not match the pinned production source") + if source.get("revision") != PINNED_SOURCE_REVISION: + raise Qwen4ArtifactError("source revision does not match the pinned production revision") + if str(root.get("tvm_ffi_patch_sha256", "")).lower() != TVM_FFI_PATCH_SHA256: + raise Qwen4ArtifactError("TVM-FFI patch does not match the frozen contract") + declared_fingerprint = _require_sha( + root.get("complete_artifact_fingerprint"), label="complete_artifact_fingerprint" + ) + unsigned = dict(root) + unsigned.pop("complete_artifact_fingerprint", None) + actual_fingerprint = hashlib.sha256( + json.dumps(unsigned, sort_keys=True, separators=(",", ":")).encode("utf-8") + ).hexdigest() + if actual_fingerprint != declared_fingerprint: + raise Qwen4ArtifactError("complete artifact fingerprint mismatch") + metadata = _require_mapping(root.get("metadata"), label="metadata") + if not isinstance(metadata.get("files"), list) or not metadata["files"]: + raise Qwen4ArtifactError("metadata.files must be a non-empty list") + metadata_files: list[ComponentFile] = [] + for index, item in enumerate(metadata["files"]): + entry = _require_mapping(item, label=f"metadata.files[{index}]") + component = ComponentFile( + path=_path_within_root( + path.parent, entry.get("path"), label=f"metadata.files[{index}].path" + ), + bytes=int(entry.get("bytes", -1)), + sha256=_require_sha( + entry.get("sha256"), label=f"metadata.files[{index}].sha256" + ), + ) + if component.bytes < 0: + raise Qwen4ArtifactError("metadata file bytes must be non-negative") + _verify_component_file(component, label=f"metadata.files[{index}]") + metadata_files.append(component) + + active = _require_mapping(root.get("active"), label="active") + if active.get("format") != ACTIVE_FORMAT: + raise Qwen4ArtifactError(f"unsupported active format {active.get('format')!r}") + active_path = _path_within_root(path.parent, active.get("path"), label="active.path") + if "bytes" in active and int(active["bytes"]) < 0: + raise Qwen4ArtifactError("active.bytes must be non-negative") + raw_active_files = active.get("files") + if not isinstance(raw_active_files, list) or not raw_active_files: + raise Qwen4ArtifactError("active.files must be a non-empty list") + active_files: list[ComponentFile] = [] + for index, item in enumerate(raw_active_files): + entry = _require_mapping(item, label=f"active.files[{index}]") + try: + size = int(entry["bytes"]) + except (KeyError, TypeError, ValueError) as exc: + raise Qwen4ArtifactError(f"active.files[{index}].bytes must be an integer") from exc + if size < 0: + raise Qwen4ArtifactError(f"active.files[{index}].bytes must be non-negative") + component = ComponentFile( + path=_path_within_root(path.parent, entry.get("path"), label=f"active.files[{index}].path"), + bytes=size, + sha256=_require_sha(entry.get("sha256"), label=f"active.files[{index}].sha256"), + ) + try: + component.path.relative_to(active_path) + except ValueError as exc: + raise Qwen4ArtifactError("active file resolves outside active.path") from exc + active_files.append(component) + + ple = _require_mapping(root.get("ple"), label="ple") + if ple.get("format") != PLE_FORMAT: + raise Qwen4ArtifactError(f"unsupported PLE format {ple.get('format')!r}") + if str(ple.get("required_volume", REQUIRED_VOLUME)).upper() != REQUIRED_VOLUME: + raise Qwen4ArtifactError("Qwen4 PLE sidecar must reside on Z:") + ple_manifest_path = _path_within_root(path.parent, ple.get("manifest"), label="ple.manifest") + ple_data_bytes = int(ple.get("data_bytes", 0)) + if ple_data_bytes < 0: + raise Qwen4ArtifactError("ple.data_bytes must be non-negative") + ple_sha256 = _require_sha(ple.get("sha256"), label="ple.sha256") + if not ple_manifest_path.is_file(): + raise Qwen4ArtifactError(f"PLE manifest is missing: {ple_manifest_path}") + try: + with ple_manifest_path.open("r", encoding="utf-8") as handle: + ple_sidecar_manifest = json.load(handle) + except (OSError, json.JSONDecodeError) as exc: + raise Qwen4ArtifactError("cannot read Q3 PLE manifest") from exc + if str(ple_sidecar_manifest.get("source_fingerprint", "")).lower() != source_inventory_sha256: + raise Qwen4ArtifactError("Q3 PLE source fingerprint mismatch") + + experts = _require_mapping(root.get("experts"), label="experts") + if experts.get("format") != EXPERT_FORMAT: + raise Qwen4ArtifactError(f"unsupported expert format {experts.get('format')!r}") + if str(experts.get("required_volume", REQUIRED_VOLUME)).upper() != REQUIRED_VOLUME: + raise Qwen4ArtifactError("Qwen4 expert sidecars must reside on Z:") + file_layers = _validate_layers(experts.get("file_tier_layers"), label="file_tier_layers") + resident_layers = _validate_layers(experts.get("resident_layers"), label="resident_layers") + if set(file_layers) & set(resident_layers): + raise Qwen4ArtifactError("experts file_tier_layers and resident_layers overlap") + files_raw = experts.get("files") + if not isinstance(files_raw, list) or not files_raw: + raise Qwen4ArtifactError("experts.files must be a non-empty list") + files: list[ExpertFile] = [] + seen: set[int] = set() + for item in files_raw: + entry = _require_mapping(item, label="experts.files[]") + try: + layer = int(entry["layer"]) + size = int(entry["bytes"]) + except (KeyError, TypeError, ValueError) as exc: + raise Qwen4ArtifactError("expert file requires integer layer and bytes") from exc + if layer < 0 or layer in seen or size < 0: + raise Qwen4ArtifactError("expert file has invalid/duplicate layer or bytes") + seen.add(layer) + files.append( + ExpertFile( + layer=layer, + path=_path_within_root(path.parent, entry.get("path"), label=f"experts.files[{layer}].path"), + bytes=size, + sha256=_require_sha(entry.get("sha256"), label=f"experts.files[{layer}].sha256"), + source_fingerprint=_require_sha( + entry.get("source_fingerprint"), + label=f"experts.files[{layer}].source_fingerprint", + ), + ) + ) + if any(item.source_fingerprint != source_inventory_sha256 for item in files): + raise Qwen4ArtifactError("expert sidecar source fingerprint mismatch") + declared_layers = set(file_layers) | set(resident_layers) + if declared_layers != seen: + raise Qwen4ArtifactError( + "experts.files layers must exactly match file_tier_layers + resident_layers" + ) + production_geometry = not allow_synthetic_geometry + if production_geometry: + active_payload = int(active.get("payload_bytes", active.get("bytes", -1))) + if active_payload != ACTIVE_TARGET_BYTES: + raise Qwen4ArtifactError("active payload does not match frozen production bytes") + if ple_data_bytes != PLE_TARGET_BYTES: + raise Qwen4ArtifactError("Q3 PLE extent does not match frozen production bytes") + if len(files) != EXPERT_LAYERS or any(item.bytes != EXPERT_FILE_BYTES for item in files): + raise Qwen4ArtifactError("expert sidecars do not match frozen production geometry") + if file_layers != FILE_TIER_LAYERS: + raise Qwen4ArtifactError("file-tier layers do not match the frozen production policy") + expected_resident = tuple(layer for layer in range(EXPERT_LAYERS) if layer not in FILE_TIER_LAYERS) + if resident_layers != expected_resident: + raise Qwen4ArtifactError("resident layers do not match the frozen production policy") + if active_payload + ple_data_bytes + sum(item.bytes for item in files) != KNOWN_TARGET_BYTES: + raise Qwen4ArtifactError("known target component byte reconciliation failed") + return Qwen4ArtifactManifest( + path=path, + raw=root, + active_path=active_path, + active_files=tuple(active_files), + ple_manifest_path=ple_manifest_path, + ple_data_bytes=ple_data_bytes, + ple_sha256=ple_sha256, + expert_files=tuple(sorted(files, key=lambda item: item.layer)), + metadata_files=tuple(metadata_files), + file_tier_layers=file_layers, + resident_layers=resident_layers, + production_geometry=production_geometry, + ) + + +def qwen4_text_only_marker(config: Any) -> bool: + """Validate and return the explicit target config marker. + + The marker is deliberately target-specific. A typo or future marker is an error, + not permission to disable vision under an unknown policy. + """ + + marker = getattr(config, "freetoken_text_only", None) + if marker is None: + return False + if marker != TEXT_ONLY_MARKER: + raise Qwen4ArtifactError(f"unsupported freetoken_text_only marker {marker!r}") + return True + + +def build_mixed_expert_sources( + manifest: Qwen4ArtifactManifest, + *, + num_experts: int = 512, + resident_residency: list[str] | None = None, + allocator=None, + verify_hash: bool = True, +): + """Materialize only resident layers from bounded expert sidecars. + + Each resident layer is allocated as six independent :class:`HostBank` buffers and + filled one record at a time through ``FileExpertSource.read_record``. File-tier + layers remain ``None`` in the returned bank lists and retain an open + ``FileExpertSource`` for demand paging. ``allocator`` is an injectable + ``(shape, dtype) -> buffer`` callback for CPU-only tests; a production call uses + the normal HostBank allocator and settles each completed layer to the requested + residency class. + """ + + if not isinstance(manifest, Qwen4ArtifactManifest): + raise TypeError("manifest must be a validated Qwen4ArtifactManifest") + num_experts = int(num_experts) + from freetoken.moe.expert_source import FileExpertSource + from freetoken.moe.host_banks import HostBank, HostResidency + + if not 1 <= num_experts <= FileExpertSource.num_experts: + raise ValueError( + f"Qwen4 modular expert sidecars support num_experts in [1, {FileExpertSource.num_experts}]" + ) + + layers = manifest.num_layers + residency = resident_residency or [HostResidency.PINNED.value] * layers + if len(residency) != layers: + raise ValueError(f"resident_residency has {len(residency)} layers, expected {layers}") + unknown_residency = set(residency) - {item.value for item in HostResidency} + if unknown_residency: + raise ValueError(f"unknown host residency values: {sorted(unknown_residency)}") + if allocator is None: + allocator = lambda shape, dtype: HostBank(shape, dtype) + + sources = {name: [None] * layers for name in FileExpertSource.bank_schema} + file_sources = {} + for layer in sorted(manifest.file_tier_layers): + entry = manifest.file_for_layer(layer) + file_sources[layer] = FileExpertSource( + entry.path, + expected_sha256=entry.sha256, + expected_source_fingerprint=entry.source_fingerprint, + expected_layer_id=layer, + num_experts=num_experts, + verify_hash=verify_hash, + ) + + # Resident layers are bounded by one six-plane record at a time. We do not + # retain a second full-layer staging tensor and close each source after fill. + # If construction fails, close already-open tier sources so partial startup + # cannot retain file handles or make fixture cleanup impossible. + try: + for layer in sorted(manifest.resident_layers): + entry = manifest.file_for_layer(layer) + with FileExpertSource( + entry.path, + expected_sha256=entry.sha256, + expected_source_fingerprint=entry.source_fingerprint, + expected_layer_id=layer, + num_experts=num_experts, + verify_hash=verify_hash, + ) as source: + buffers = { + name: allocator(shape=(num_experts, *shape), dtype=dtype) + for name, (shape, dtype) in source.plane_specs.items() + } + for expert_id in range(num_experts): + row = source.read_record(expert_id) + for name, buffer in buffers.items(): + destination = getattr(buffer, "tensor", buffer) + destination[expert_id].copy_(row[name]) + settle = residency[layer] + for buffer in buffers.values(): + if settle == HostResidency.PINNED.value and hasattr(buffer, "pin"): + buffer.pin() + elif settle == HostResidency.LOCKED.value and hasattr(buffer, "lock"): + buffer.lock() + for name, buffer in buffers.items(): + sources[name][layer] = getattr(buffer, "tensor", buffer) + except Exception: + for source in file_sources.values(): + source.close() + raise + return sources, file_sources + + +def configure_mixed_expert_sources( + cache, + manifest: Qwen4ArtifactManifest | Mapping[str, Any], + resident_sources, + *, + file_sources: Mapping[int, object] | None = None, + layer_residency: list[str] | None = None, +): + """Wire resident bank entries and file-backed layers into an offload cache. + + ``resident_sources`` is injected by the caller (normally the existing FTW bank + loader). This keeps the helper unit-testable with tiny synthetic tensors and avoids + implementing another expert writer here. File tiers use only the public + :class:`FileExpertSource` reader API. + """ + + if not isinstance(manifest, Qwen4ArtifactManifest): + raise TypeError("manifest must be a validated Qwen4ArtifactManifest") + if set(manifest.file_tier_layers) & set(cache.cpu_layer_ids): + raise ValueError("Qwen4 modular file-tier layers are GPU-only") + if cache.prefill_overlap: + raise ValueError("Qwen4 modular file tiers require prefill_overlap=False") + # The artifact's native expert geometry is 512 slots. Enforce this independently + # of a malformed/synthetic model config so rebuilds cannot create an undersized cache. + if int(cache.cache_size) < 512: + raise ValueError("Qwen4 modular expert cache requires at least 512 slots") + if not isinstance(resident_sources, Mapping): + raise TypeError("resident_sources must be a bank-name -> per-layer mapping") + layers = manifest.num_layers + if layers <= 0 or int(cache.num_layers) != layers: + raise ValueError( + f"manifest layer geometry ({layers}) does not match cache ({cache.num_layers})" + ) + file_layers = set(manifest.file_tier_layers) + resident_layers = set(manifest.resident_layers) + if file_layers | resident_layers != set(range(layers)): + raise ValueError("manifest expert layer sets must cover a contiguous model") + if set(resident_sources) != set(cache.bank_schema): + raise ValueError("resident bank schema does not match cache quant_format") + per_layer: dict[str, list[Any]] = {} + for name in cache.bank_schema: + values = list(resident_sources[name]) + if len(values) != layers: + raise ValueError(f"resident bank {name!r} has {len(values)} layers, expected {layers}") + for layer in file_layers: + if values[layer] is not None: + # A file tier must be represented explicitly as a None resident source. + values[layer] = None + for layer in resident_layers: + if values[layer] is None: + raise ValueError(f"resident layer {layer} has no resident source for {name}") + per_layer[name] = values + cache.set_bank_sources(per_layer, layer_residency=layer_residency) + from freetoken.moe.expert_source import FileExpertSource + + if file_sources is None: + file_sources = {} + for layer in sorted(file_layers): + entry = manifest.file_for_layer(layer) + source = FileExpertSource( + entry.path, + expected_sha256=entry.sha256, + expected_source_fingerprint=entry.source_fingerprint, + expected_layer_id=layer, + num_experts=cache.num_experts, + ) + file_sources[layer] = source + else: + file_sources = {int(layer): source for layer, source in file_sources.items()} + if set(file_sources) != file_layers: + raise ValueError("file_sources keys do not match manifest file_tier_layers") + cache.set_file_sources(dict(file_sources)) + return dict(file_sources) + + +def build_qwen4_modular_artifact( + source_path: str | os.PathLike[str], + artifact_root: str | os.PathLike[str], + *, + source_repository: str = PINNED_SOURCE_REPOSITORY, + source_revision: str = PINNED_SOURCE_REVISION, + source_inventory_sha256: str, + minimum_freetoken_commit: str, + ple_layer_id: int = 2, + ple_split_parts: int = 128, + expert_layers: tuple[int, ...] = tuple(range(EXPERT_LAYERS)), + file_tier_layers: tuple[int, ...] = FILE_TIER_LAYERS, + expert_num_experts: int = EXPERT_NUM_EXPERTS, + expert_geometry=None, + allow_synthetic_geometry: bool = False, +) -> dict[str, Any]: + """Run the canonical C1-C5 modular conversion sequence. + + C0 (full source-file identity verification) remains a mandatory caller gate. + This function performs no network I/O and accepts only an already-local Z:-backed + source snapshot. Each component writer owns its bounded streaming/atomic contract; + the final manifest is published only after all components reopen successfully. + """ + + source = _resolve_z(source_path, label="source checkpoint") + root = _resolve_z(artifact_root, label="artifact root") + if not source.is_dir(): + raise Qwen4ArtifactError(f"source checkpoint is not a directory: {source}") + root.mkdir(parents=True, exist_ok=True) + inventory = _require_sha(source_inventory_sha256, label="source_inventory_sha256") + + from freetoken.checkpoint.convert import convert_checkpoint + from freetoken.checkpoint.q3_ple import write_q3_ple_from_safetensors + from freetoken.moe.expert_source import write_expert_sidecar_from_safetensors + + active_index = convert_checkpoint( + str(source), + str(root), + artifact_format=ARTIFACT_FORMAT, + source_inventory_sha256=inventory, + ) + write_q3_ple_from_safetensors( + source, + root / "ple-q3-000.bin", + root / "ple-q3.json", + layer_id=int(ple_layer_id), + split_parts=int(ple_split_parts), + source_fingerprint=inventory, + ) + expert_paths: dict[int, str] = {} + for layer in expert_layers: + name = f"experts-L{int(layer):02d}.nvfp4" + write_expert_sidecar_from_safetensors( + source, + root / name, + layer_id=int(layer), + source_fingerprint=inventory, + num_experts=int(expert_num_experts), + geometry=expert_geometry, + ) + expert_paths[int(layer)] = name + metadata_paths = list(active_index.get("copied_metadata", ())) + if not metadata_paths: + raise Qwen4ArtifactError("active conversion copied no target metadata") + return finalize_qwen4_modular_manifest( + root, + source_repository=source_repository, + source_revision=source_revision, + source_inventory_sha256=inventory, + minimum_freetoken_commit=minimum_freetoken_commit, + tvm_ffi_patch_sha256=TVM_FFI_PATCH_SHA256, + expert_paths=expert_paths, + file_tier_layers=file_tier_layers, + metadata_paths=metadata_paths, + expert_num_experts=expert_num_experts, + allow_synthetic_geometry=allow_synthetic_geometry, + ) + + +def finalize_qwen4_modular_manifest( + artifact_root: str | os.PathLike[str], + *, + source_repository: str, + source_revision: str, + source_inventory_sha256: str, + minimum_freetoken_commit: str, + tvm_ffi_patch_sha256: str, + active_dir: str | os.PathLike[str] = "qwen4-active-v1.ftw", + ple_manifest: str | os.PathLike[str] = "ple-q3.json", + expert_paths: Mapping[int, str | os.PathLike[str]], + file_tier_layers: list[int] | tuple[int, ...], + metadata_paths: list[str | os.PathLike[str]], + expert_num_experts: int = 512, + allow_synthetic_geometry: bool = False, +) -> dict[str, Any]: + """Validate completed components and atomically publish ``manifest.json``. + + This is the final C5 orchestration seam. It never creates weight payloads; + the C2/C3/C4 writers must already have atomically finalized their components. + Every file is length/hash inventoried here, the Q3 and FTEXPERT1 readers reopen + their formats, and only then is the complete manifest promoted. + """ + + root = _resolve_z(artifact_root, label="artifact root") + root.mkdir(parents=True, exist_ok=True) + active_root = _path_from(root, active_dir, label="active_dir") + from freetoken.checkpoint.ftw import INDEX_NAME, is_ftw_checkpoint + + if not active_root.is_dir() or not is_ftw_checkpoint(str(active_root)): + raise Qwen4ArtifactError(f"active component is not an FTW checkpoint: {active_root}") + with (active_root / INDEX_NAME).open("r", encoding="utf-8") as handle: + active_index = json.load(handle) + inventory_digest = _require_sha( + source_inventory_sha256, label="source_inventory_sha256" + ) + if str(active_index.get("source_inventory_sha256", "")).lower() != inventory_digest: + raise Qwen4ArtifactError("active FTW source inventory fingerprint mismatch") + active_payload_bytes = int(active_index.get("total_bytes", -1)) + if active_payload_bytes < 0: + raise Qwen4ArtifactError("active FTW index has no valid total_bytes") + active_files = [ + _component_file(root, item) + for item in sorted(path for path in active_root.rglob("*") if path.is_file()) + ] + + ple_path = _path_within_root(root, ple_manifest, label="ple_manifest") + from freetoken.checkpoint.q3_ple import Q3PLEReader + + with Q3PLEReader(ple_path) as ple_reader: + ple_data = ple_reader.data_path + ple_data_bytes = ple_data.stat().st_size + ple_sha256 = _sha256_file(ple_data) + ple_source_fingerprint = str(ple_reader.manifest.get("source_fingerprint", "")) + if ple_source_fingerprint.lower() != inventory_digest: + raise Qwen4ArtifactError("Q3 PLE source fingerprint mismatch") + + tiered = tuple(sorted(_validate_layers(list(file_tier_layers), label="file_tier_layers"))) + expert_items: list[dict[str, Any]] = [] + from freetoken.moe.expert_source import FileExpertSource + + for layer, value in sorted((int(layer), path) for layer, path in expert_paths.items()): + expert_path = _path_within_root(root, value, label=f"expert layer {layer}") + with FileExpertSource( + expert_path, + expected_source_fingerprint=inventory_digest, + expected_layer_id=layer, + num_experts=int(expert_num_experts), + verify_hash=True, + ) as source: + if source.layer_id != layer: + raise Qwen4ArtifactError(f"expert sidecar layer mismatch for {expert_path}") + expert_items.append( + { + "layer": layer, + **_component_file(root, expert_path), + "source_fingerprint": source.source_fingerprint, + } + ) + all_layers = tuple(item["layer"] for item in expert_items) + if all_layers != tuple(range(len(all_layers))): + raise Qwen4ArtifactError("expert sidecars must cover contiguous layers from zero") + if not set(tiered) <= set(all_layers): + raise Qwen4ArtifactError("file tier contains a layer without an expert sidecar") + resident = sorted(set(all_layers) - set(tiered)) + if not allow_synthetic_geometry: + if source_repository != PINNED_SOURCE_REPOSITORY or source_revision != PINNED_SOURCE_REVISION: + raise Qwen4ArtifactError("production artifact source pin mismatch") + commit = str(minimum_freetoken_commit).lower() + if len(commit) != 40 or any(char not in "0123456789abcdef" for char in commit): + raise Qwen4ArtifactError("production artifact requires a concrete FreeToken commit") + if str(tvm_ffi_patch_sha256).lower() != TVM_FFI_PATCH_SHA256: + raise Qwen4ArtifactError("production artifact TVM-FFI patch mismatch") + if int(expert_num_experts) != EXPERT_NUM_EXPERTS: + raise Qwen4ArtifactError("production artifact requires 512 experts per layer") + if active_payload_bytes != ACTIVE_TARGET_BYTES: + raise Qwen4ArtifactError("active FTW does not match frozen production bytes") + if ple_data_bytes != PLE_TARGET_BYTES: + raise Qwen4ArtifactError("Q3 PLE does not match frozen production extent") + if tuple(all_layers) != tuple(range(EXPERT_LAYERS)): + raise Qwen4ArtifactError("production artifact requires 48 expert sidecars") + if any(item["bytes"] != EXPERT_FILE_BYTES for item in expert_items): + raise Qwen4ArtifactError("expert sidecar does not match frozen production bytes") + if tiered != FILE_TIER_LAYERS: + raise Qwen4ArtifactError("file tier does not match the frozen production policy") + if active_payload_bytes + ple_data_bytes + sum(item["bytes"] for item in expert_items) != KNOWN_TARGET_BYTES: + raise Qwen4ArtifactError("known target component byte reconciliation failed") + + metadata_files = [ + _component_file(root, _path_within_root(root, value, label="metadata file")) + for value in metadata_paths + ] + if not any(item["path"] == "config.json" for item in metadata_files): + raise Qwen4ArtifactError("modular artifact metadata must include config.json") + config_path = root / "config.json" + with config_path.open("r", encoding="utf-8") as handle: + config = json.load(handle) + if config.get("freetoken_text_only") != TEXT_ONLY_MARKER: + raise Qwen4ArtifactError("config.json lacks the accepted text-only marker") + if config.get("freetoken_active_quant") != ACTIVE_FORMAT: + raise Qwen4ArtifactError("config.json lacks the accepted active-quant marker") + + manifest: dict[str, Any] = { + "format": FORMAT, + "version": VERSION, + "artifact_schema": FORMAT, + "text_only": True, + "source": { + "repository": str(source_repository), + "revision": str(source_revision), + "inventory_sha256": inventory_digest, + }, + "minimum_freetoken_commit": str(minimum_freetoken_commit), + "tvm_ffi_patch_sha256": _require_sha( + tvm_ffi_patch_sha256, label="tvm_ffi_patch_sha256" + ), + "active": { + "format": ACTIVE_FORMAT, + "path": active_root.relative_to(root).as_posix(), + "payload_bytes": active_payload_bytes, + "physical_file_bytes": sum(item["bytes"] for item in active_files), + "files": active_files, + }, + "ple": { + "format": PLE_FORMAT, + "manifest": ple_path.relative_to(root).as_posix(), + "data_bytes": ple_data_bytes, + "sha256": ple_sha256, + "source_fingerprint": ple_source_fingerprint, + "required_volume": REQUIRED_VOLUME, + }, + "experts": { + "format": EXPERT_FORMAT, + "files": expert_items, + "file_tier_layers": list(tiered), + "resident_layers": resident, + "required_volume": REQUIRED_VOLUME, + }, + "metadata": {"files": metadata_files}, + } + canonical = json.dumps(manifest, sort_keys=True, separators=(",", ":")).encode("utf-8") + manifest["complete_artifact_fingerprint"] = hashlib.sha256(canonical).hexdigest() + partial = root / f".{MANIFEST_NAME}.partial-{os.getpid()}" + with partial.open("w", encoding="utf-8", newline="\n") as handle: + json.dump(manifest, handle, indent=2, sort_keys=True) + handle.write("\n") + handle.flush() + os.fsync(handle.fileno()) + os.replace(partial, root / MANIFEST_NAME) + return manifest + + +__all__ = [ + "ACTIVE_FORMAT", + "ARTIFACT_FORMAT", + "EXPERT_FORMAT", + "FORMAT", + "MANIFEST_NAME", + "PLE_FORMAT", + "Qwen4ArtifactError", + "Qwen4ArtifactManifest", + "TEXT_ONLY_MARKER", + "VERSION", + "build_qwen4_modular_artifact", + "configure_mixed_expert_sources", + "build_mixed_expert_sources", + "finalize_qwen4_modular_manifest", + "load_qwen4_artifact_manifest", + "qwen4_text_only_marker", +] diff --git a/python/freetoken/engine/cache_budget.py b/python/freetoken/engine/cache_budget.py index ab7c0a9f..5214619f 100644 --- a/python/freetoken/engine/cache_budget.py +++ b/python/freetoken/engine/cache_budget.py @@ -25,7 +25,13 @@ def expert_bytes_per_slot(sources: dict[str, "list[torch.Tensor]"]) -> int: # with cache_size), so they are intentionally excluded from the per-slot growth term. # tensor[0].numel() is the per-row element count (one expert slot); see the matching # slot-byte idiom in kvcache/linear_state_pool.py and kvcache/dsv4_paged_pool.py. - return sum(t[0][0].numel() * t[0].element_size() for t in sources.values()) + total = 0 + for per_layer in sources.values(): + representative = next((tensor for tensor in per_layer if tensor is not None), None) + if representative is None: + raise ValueError("expert bank has no resident shape for cache sizing") + total += representative[0].numel() * representative.element_size() + return total def net_cache_budget_bytes( diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index dd9a8499..a4b5122e 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -506,6 +506,107 @@ def _resolve_auto_moe_cache_size(self, config: EngineConfig, banks) -> tuple[int ) def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: + # Qwen4 modular artifacts carry their own mixed expert tier. The ordinary + # loader cannot infer file-backed layers from a generic checkpoint path, so + # resolve the manifest first and route through the explicit wiring seam. + from freetoken.checkpoint.qwen4_artifact import ( + build_mixed_expert_sources, + configure_mixed_expert_sources, + load_qwen4_artifact_manifest, + ) + + artifact = load_qwen4_artifact_manifest(config.model_path) + if artifact is not None: + if config.moe_backend not in ("offload",): + raise ValueError( + "Qwen4 modular file tiers require --moe-backend offload; " + "CPU execution may be selected only for resident layers" + ) + if config.moe_prefill_overlap: + object.__setattr__(config, "moe_prefill_overlap", False) + # Expert sidecars are authoritative for both tiers. Resident layers are + # streamed one record at a time into HostBanks; file layers remain open + # FileExpertSource readers and never materialize a full layer. + from freetoken.moe.host_banks import HostResidency + from freetoken.moe.offload_cache import _BANK_BYTES_PER_EXPERT + + cpu_layer_ids = _resolve_cpu_layers(config, artifact.num_layers) + tiered = set(artifact.file_tier_layers) + if cpu_layer_ids & tiered: + raise ValueError( + "--moe-cpu-layers selects a file-tier layer; only resident layers " + "are CPU/hybrid eligible" + ) + resident = set(artifact.resident_layers) + pin_budget = _pin_budget_bytes() + if pin_budget is not None: + row_bytes = _BANK_BYTES_PER_EXPERT["nvfp4"]( + config.model_config.hidden_size, + config.model_config.moe_intermediate_size, + ) + max_pinned = pin_budget // (row_bytes * config.model_config.num_experts) + required_locked = max(0, len(resident - set(cpu_layer_ids)) - max_pinned) + if required_locked: + if not _cpu_moe_executor_viable(config.model_config): + raise ValueError( + "resident expert banks exceed the Windows pin budget and the " + "CPU executor is unavailable; refusing pageable GPU sources" + ) + # Deterministic outer-to-inner choice within the resident set. + candidates = sorted( + resident - set(cpu_layer_ids), + key=lambda layer: (min(layer, artifact.num_layers - 1 - layer), layer), + ) + cpu_layer_ids = frozenset(set(cpu_layer_ids) | set(candidates[:required_locked])) + residency = [HostResidency.PINNED.value] * artifact.num_layers + for layer in cpu_layer_ids: + residency[layer] = HostResidency.LOCKED.value + + resident_sources, file_sources = build_mixed_expert_sources( + artifact, + num_experts=config.model_config.num_experts, + resident_residency=residency, + verify_hash=not config.use_dummy_weight, + ) + from freetoken.moe.expert_banks import ExpertBanks + + banks = ExpertBanks("nvfp4", resident_sources, layer_residency=residency) + if config.moe_cache_auto: + size, pages, _overlap = self._resolve_auto_moe_cache_size(config, banks) + object.__setattr__(config, "moe_cache_size", max(size, 512)) + if config.num_page_override is None: + object.__setattr__(config, "num_page_override", pages) + _require_offload_cache_size(config.moe_cache_size, 512) + cache = OffloadMoeCache( + num_layers=config.model_config.num_moe_layers, + num_experts=config.model_config.num_experts, + cache_size=config.moe_cache_size, + device=self.device, + cache_policy=config.moe_cache_policy, + prefill_overlap=False, + prefill_hit_d2d=False, + quant_format=banks.quant_format, + decode_target="cpu" if cpu_layer_ids else "gpu", + hybrid_max_fetch=config.moe_hybrid_max_fetch, + ) + cache.cpu_layer_ids = cpu_layer_ids + configure_mixed_expert_sources( + cache, + artifact, + banks.sources, + file_sources=file_sources, + layer_residency=residency, + ) + cache.set_alphas(banks.gate_up_alpha, banks.down_alpha) + cache.collect_stats = config.moe_collect_stats + layers = attach_offload_moe_cache(self.model, cache) + assert len(layers) == config.model_config.num_moe_layers + if cache.decode_target == "cpu": + self._init_cpu_moe_executor(config, cache, layers) + self.ctx.moe_offload_cache = cache + self.moe_offload_cache = cache + return cache + # A model may fully own cache construction via make_offload_moe_cache. # Otherwise load_expert_banks gives the model module a setup hook first, then # falls back to per-quant providers, and the engine wires the banks into cache. diff --git a/python/freetoken/kernel/csrc/include/freetoken/tensor.h b/python/freetoken/kernel/csrc/include/freetoken/tensor.h index 9b591a87..2d0af903 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/tensor.h +++ b/python/freetoken/kernel/csrc/include/freetoken/tensor.h @@ -397,6 +397,15 @@ struct TensorMatcher { } template + auto with_device(SymbolicDevice &device) && -> TensorMatcher && { + m_init_device(); + if constexpr (sizeof...(Codes) > 0) { + device.set_options(); + } + m_device.rebind(device); + return std::move(*this); + } + auto with_device(DeviceRef &&device) && -> TensorMatcher && { m_init_device(); m_device.rebind(*device); diff --git a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh index bb83c23e..c6a47136 100644 --- a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh +++ b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh @@ -354,7 +354,7 @@ struct FastIndexCopyKernel { TensorMatcher({L}) .with_dtype(indices_dtype) - .with_device(device) + .template with_device(device) .verify(src_indices) .verify(dst_indices); @@ -363,7 +363,7 @@ struct FastIndexCopyKernel { const auto num_indices_tensor = num_indices.value(); TensorMatcher({1}) .with_dtype(num_indices_dtype) - .with_device(device) + .template with_device(device) .verify(num_indices_tensor); num_indices_data_ptr = static_cast(num_indices_tensor.data_ptr()); @@ -529,14 +529,14 @@ struct MultiIndexCopyKernel { auto indices_dtype = SymbolicDType{}; auto num_indices_dtype = SymbolicDType{}; - TensorMatcher({B}).with_dtype(ptr_dtype).with_device(device) + TensorMatcher({B}).with_dtype(ptr_dtype).template with_device(device) .verify(dst_ptrs).verify(src_ptrs).verify(feat_bytes); - TensorMatcher({L}).with_dtype(indices_dtype).with_device(device) + TensorMatcher({L}).with_dtype(indices_dtype).template with_device(device) .verify(dst_indices).verify(src_indices); const int64_t* valid_length = nullptr; if (num_indices.has_value()) { - TensorMatcher({1}).with_dtype(num_indices_dtype).with_device(device) + TensorMatcher({1}).with_dtype(num_indices_dtype).template with_device(device) .verify(num_indices.value()); valid_length = static_cast(num_indices.value().data_ptr()); } diff --git a/python/freetoken/models/qwen3_5_moe/gdn.py b/python/freetoken/models/qwen3_5_moe/gdn.py index 2e732005..b370f3cb 100644 --- a/python/freetoken/models/qwen3_5_moe/gdn.py +++ b/python/freetoken/models/qwen3_5_moe/gdn.py @@ -53,7 +53,7 @@ class Qwen3_5GatedDeltaNet(BaseOP): def __init__( self, hidden_size, num_k_heads, num_v_heads, head_k_dim, head_v_dim, conv_kernel_size, rms_norm_eps, layer_id, expert_quant: str = "none", - attn_quant: str = "none", + attn_quant: str = "none", *, nvfp4_qkvz: bool = False, ): self.layer_id = layer_id # The fla chunk/decode kernels read+write the recurrent state and the per-chunk h as @@ -71,16 +71,25 @@ def __init__( self.value_dim = num_v_heads * head_v_dim self.conv_dim = 2 * self.key_dim + self.value_dim self.conv_kernel_size = conv_kernel_size - # qkv|z carry a weight scale (block-fp8 weight_scale_inv, or per-tensor FP8 - # weight_scale); b|a stay bf16. Both quant modes therefore split the four-way - # fusion into an fp8 qkvz GEMM + a bf16 ba GEMM (matches sglang/vLLM). + # qkv|z carry a weight scale (block-fp8 weight_scale_inv, per-tensor FP8 + # weight_scale, or native NVFP4 scales); b|a stay bf16. The explicit + # ``nvfp4_qkvz`` opt-in is used by Qwen4 only. Keeping it separate from + # ``attn_quant`` preserves Qwen3.5's historical NVFP4 behavior (where only + # out_proj is native FP4). self._block_fp8 = expert_quant == "fp8_block" self._pertensor_fp8 = attn_quant == "fp8_pertensor" + self._nvfp4_qkvz = bool(nvfp4_qkvz) self._fp8 = self._block_fp8 or self._pertensor_fp8 + self._split_input = self._fp8 or self._nvfp4_qkvz self._in_proj_split = [self.conv_dim, self.value_dim, num_v_heads, num_v_heads] - if self._fp8: - ColMerged = Fp8BlockColMerged if self._block_fp8 else Fp8PerTensorColMerged + if self._split_input: + if self._nvfp4_qkvz: + from freetoken.kernel.triton.nvfp4_linear import Nvfp4DenseColMerged + + ColMerged = Nvfp4DenseColMerged + else: + ColMerged = Fp8BlockColMerged if self._block_fp8 else Fp8PerTensorColMerged self.in_proj_qkvz = ColMerged( hidden_size, [self.conv_dim, self.value_dim], has_bias=False ) @@ -98,9 +107,9 @@ def __init__( self.dt_bias = torch.empty(num_v_heads, dtype=torch.float32) self.A_log = torch.empty(num_v_heads, dtype=torch.float32) self.norm = _GatedRMSNorm(head_v_dim, eps=rms_norm_eps) - # out_proj follows the checkpoint quant: block-fp8 / per-tensor-fp8 / compressed-tensors - # NVFP4 (W4A16) / bf16. in_proj_* stay bf16 in every mode (above), so a compressed-tensors - # NVFP4 checkpoint (attn_quant=="nvfp4") only makes out_proj native FP4. + # out_proj follows the checkpoint quant: block-fp8 / per-tensor-fp8 / NVFP4 (W4A16) / + # bf16. In Qwen4's explicit ``nvfp4_qkvz`` mode, qkv|z is native NVFP4 while b|a + # remains BF16; Qwen3.5 callers retain the historical fused-BF16 input path. self.out_proj = make_replicated_quant( expert_quant, attn_quant, self.value_dim, hidden_size, has_bias=False ) @@ -161,7 +170,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: fla = build_fla_metadata(batch, hidden_states.device) batch.fla_metadata = fla - if self._fp8: + if self._split_input: qkvz = self.in_proj_qkvz.forward(hidden_states) conv_in, z = torch.split(qkvz, [self.conv_dim, self.value_dim], dim=-1) ba = self.in_proj_ba.forward(hidden_states) diff --git a/python/freetoken/models/qwen4_exp/config.py b/python/freetoken/models/qwen4_exp/config.py index b53811a9..ccd0a319 100644 --- a/python/freetoken/models/qwen4_exp/config.py +++ b/python/freetoken/models/qwen4_exp/config.py @@ -9,6 +9,7 @@ RotaryConfig, detect_expert_quant, ) +from freetoken.checkpoint.qwen4_artifact import qwen4_text_only_marker from .args import Qwen4ExpArgs, Qwen4VisionConfig @@ -94,7 +95,11 @@ def parse_config(hf_config: Any) -> ModelConfig: "Qwen4-Exp mrope_section must cover the rotary dimension: " f"{qwen4_args.mrope_section} vs {rotary_dim}" ) - raw_vision = getattr(hf_config, "vision_config", None) + # A modular Qwen4 artifact is explicitly text-only. Keep the ordinary source + # checkpoint behavior untouched (vision remains part of the parsed model) and + # fail closed on an unknown target marker in qwen4_text_only_marker(). + text_only = qwen4_text_only_marker(hf_config) + raw_vision = None if text_only else getattr(hf_config, "vision_config", None) vision_config = None if raw_vision is not None: vision_config = Qwen4VisionConfig( @@ -143,6 +148,11 @@ def parse_config(hf_config: Any) -> ModelConfig: f"detected {detected_quant!r}" ) + active_quant = getattr(hf_config, "freetoken_active_quant", None) + if active_quant not in (None, "nvfp4_w4a16_v1"): + raise ValueError(f"unsupported Qwen4 active-weight format: {active_quant}") + active_linear_quant = "nvfp4" if active_quant == "nvfp4_w4a16_v1" else "none" + return ModelConfig( num_layers=int(text.num_hidden_layers), num_qo_heads=int(text.num_attention_heads), @@ -159,24 +169,25 @@ def parse_config(hf_config: Any) -> ModelConfig: num_experts_per_tok=int(text.num_experts_per_tok), moe_intermediate_size=int(text.moe_intermediate_size), shared_expert_intermediate_size=int(text.shared_expert_intermediate_size), - # The released Qwen3.8-Flash-Next configs omit this older Qwen MoE - # field. Omission means that the router weights are not renormalized. - norm_topk_prob=bool(getattr(text, "norm_topk_prob", False)), + # The official Qwen4-Exp config defaults this field to True. Released + # checkpoints may omit it, while an explicit False must remain False. + norm_topk_prob=bool(getattr(text, "norm_topk_prob", True)), model_type=str(hf_config.model_type), architectures=list(hf_config.architectures), moe_enabled=True, expert_quant=expert_quant, weight_block_size=block_size, - # Only routed experts and PLE are FP8 in the official checkpoint. All - # attention, hyper-connection, and shared-expert projections stay BF16. - attn_quant="none", - dense_quant="none", + # The published experts-only checkpoint has no active-weight marker and + # therefore retains BF16 operators. Only the canonical preconverted + # FTW artifact may opt into the frozen native W4A16 map. + attn_quant=active_linear_quant, + dense_quant=active_linear_quant, lm_head_quant="none", use_qk_norm=True, # Qwen3.8-Flash-Next is a VL checkpoint. Vision is part of this model, # not an optional text-only add-on. vision_config=vision_config, - image_token_id=getattr(hf_config, "image_token_id", None), + image_token_id=None if text_only else getattr(hf_config, "image_token_id", None), attention_groups=groups, qwen4_args=qwen4_args, # PLE keeps per-request dilated-convolution state outside the generic diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py index d0ba393a..c1e50724 100644 --- a/python/freetoken/models/qwen4_exp/model.py +++ b/python/freetoken/models/qwen4_exp/model.py @@ -25,13 +25,14 @@ get_rope, ) from freetoken.models.blocks import BaseLLMModel +from freetoken.models.config import ModelConfig from freetoken.models.qwen3_5_moe.attention import Qwen3_5Attention +from freetoken.checkpoint.q3_ple import Q3PLEReader from freetoken.models.qwen3_5_moe.gdn import Qwen3_5GatedDeltaNet +from freetoken.models.quant_linear import make_col_merged_quant, make_replicated_quant from freetoken.utils import download_hf_weight, nvtx_annotate if TYPE_CHECKING: - from freetoken.models.config import ModelConfig - from .args import Qwen4ExpArgs @@ -187,8 +188,16 @@ def __init__(self, config: ModelConfig, combine: bool = True): self.hidden_size = config.hidden_size hc_size = self.hc_count * self.hidden_size self.hc_norm = _GroupedRMSNorm(hc_size, self.hidden_size, config.rms_norm_eps) - self.input_mix_weight_down = LinearReplicated(hc_size, args.hc_lowrank, has_bias=False) - self.input_mix_weight_up = LinearReplicated(args.hc_lowrank, hc_size, has_bias=False) + # The frozen Qwen4 active map keeps mHC input-mix down/up native NVFP4 when + # ``dense_quant=nvfp4`` is explicitly selected. Block injection remains BF16. + self.input_mix_weight_down = make_replicated_quant( + "none", getattr(config, "dense_quant", "none"), hc_size, args.hc_lowrank, + has_bias=False, + ) + self.input_mix_weight_up = make_replicated_quant( + "none", getattr(config, "dense_quant", "none"), args.hc_lowrank, hc_size, + has_bias=False, + ) self.block_inject_weight = ( LinearReplicated(hc_size, self.hc_count, has_bias=False) if combine else None ) @@ -208,10 +217,14 @@ def forward(self, hyper_input: torch.Tensor): class _SharedExpert(BaseOP): def __init__(self, config: ModelConfig): width = config.shared_expert_intermediate_size - self.gate_up_proj = LinearColParallelMerged( - config.hidden_size, [width, width], has_bias=False + self.gate_up_proj = make_col_merged_quant( + "none", getattr(config, "dense_quant", "none"), config.hidden_size, + [width, width], has_bias=False, + ) + self.down_proj = make_replicated_quant( + "none", getattr(config, "dense_quant", "none"), width, config.hidden_size, + has_bias=False, ) - self.down_proj = LinearRowParallel(width, config.hidden_size, has_bias=False) def forward(self, hidden: torch.Tensor) -> torch.Tensor: return self.down_proj.forward(silu_and_mul(self.gate_up_proj.forward(hidden))) @@ -278,17 +291,28 @@ def build_ngram_ids( return torch.cat(blocks, dim=-1) -def _ple_request_tokens(req, forwarded_ids: torch.Tensor | None = None) -> torch.Tensor: - """Return the complete host token history visible to this forward. +def _ple_request_tokens( + req, + forwarded_ids: torch.Tensor | None = None, + *, + start: int = 0, +) -> torch.Tensor: + """Return host-visible request tokens from ``start`` through ``device_len``. The overlap scheduler advances ``device_len`` before it drains the prior sampled token to ``req.input_ids``. During decode, that one current token is already present in ``batch.input_ids``. Join it to the committed host prefix - so PLE hashes the same history as a non-overlapped forward. + so PLE hashes the same history as a non-overlapped forward. The default + preserves the complete-history contract; callers may select a validated + suffix when only the N-gram dependency prefix is needed. """ + if not 0 <= start <= req.device_len: + raise ValueError( + f"Qwen4-Exp PLE history start {start} is outside [0, {req.device_len}]" + ) host_len = req.input_ids.numel() if host_len >= req.device_len: - return req.input_ids[: req.device_len] + return req.input_ids[start : req.device_len] if host_len != req.cached_len: raise RuntimeError( "Qwen4-Exp PLE host history has an unexpected gap: " @@ -300,7 +324,14 @@ def _ple_request_tokens(req, forwarded_ids: torch.Tensor | None = None) -> torch "Qwen4-Exp PLE needs the current forwarded tokens: " f"got {actual}, expected {req.extend_len}" ) - return torch.cat((req.input_ids[: req.cached_len], forwarded_ids.to(device="cpu"))) + host_start = min(start, req.cached_len) + forwarded_start = max(0, start - req.cached_len) + return torch.cat( + ( + req.input_ids[host_start : req.cached_len], + forwarded_ids[forwarded_start:].to(device="cpu"), + ) + ) class _HostNGramEmbedding(BaseOP): @@ -323,8 +354,20 @@ def __init__(self, config: ModelConfig, layer_id: int): self._scale = torch.tensor(1.0, dtype=torch.bfloat16) self._host_constants: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None self._dummy = False + self._q3_reader: Q3PLEReader | None = None - def load_host_weights(self, model_path: str, *, dummy: bool = False) -> None: + def load_host_weights( + self, + model_path: str, + *, + dummy: bool = False, + ple_format: str = "fp8_safetensors", + ) -> None: + if ple_format == "q3_ple_32": + self.load_q3_ple_weights(model_path) + return + if ple_format != "fp8_safetensors": + raise ValueError(f"unsupported Qwen4 PLE format: {ple_format}") if dummy: self._dummy = True return @@ -382,6 +425,28 @@ def load_host_weights(self, model_path: str, *, dummy: bool = False) -> None: f"PLE table has {int(self._shard_ends[-1])} rows, needs {expected_rows}" ) + def load_q3_ple_weights(self, manifest_path: str) -> None: + """Opt into the native Q3_PLE_32 sidecar; FP8 Safetensors stays default.""" + + reader = Q3PLEReader(manifest_path) + if self._host_constants is None: + self._host_constants = ( + self.layer_multipliers.cpu(), + self.ngram_heads_vocab_sizes.cpu(), + self.ngram_heads_offsets.cpu(), + ) + expected_rows = int(self._host_constants[1][-1] + self._host_constants[2][-1]) + if reader.row_count < expected_rows: + reader.close() + raise RuntimeError( + f"Q3_PLE_32 table has {reader.row_count} rows, needs {expected_rows}" + ) + self._q3_reader = reader + self._shards = [] + self._handles = [] + self._shard_ends = torch.empty(0, dtype=torch.long) + self._scale = torch.tensor(reader.weight_scale, dtype=torch.bfloat16) + def _current_ngram_ids(self) -> torch.Tensor: if self._host_constants is None: raise RuntimeError("Qwen4-Exp PLE host weights are not loaded") @@ -400,7 +465,8 @@ def _current_ngram_ids(self) -> torch.Tensor: forwarded = forwarded_host[ forwarded_offset : forwarded_offset + extend_len ] - tokens = _ple_request_tokens(req, forwarded) + history_start = max(0, req.cached_len - (self.ngram_size - 1)) + tokens = _ple_request_tokens(req, forwarded, start=history_start) all_ids = build_ngram_ids( tokens, ngram_size=self.ngram_size, @@ -410,8 +476,17 @@ def _current_ngram_ids(self) -> torch.Tensor: vocab_sizes=vocab_sizes, offsets=offsets, ) - pieces.append(all_ids[req.cached_len : req.device_len]) + pieces.append( + all_ids[ + req.cached_len - history_start : req.device_len - history_start + ] + ) forwarded_offset += extend_len + if forwarded_offset != batch.input_ids.numel(): + raise RuntimeError( + f"Qwen4-Exp PLE consumed {forwarded_offset} forwarded tokens, " + f"but the batch carries {batch.input_ids.numel()}" + ) result = torch.cat(pieces, dim=0) if result.shape[0] != batch.input_ids.numel(): raise RuntimeError( @@ -424,6 +499,10 @@ def forward(self, device: torch.device, dtype: torch.dtype) -> torch.Tensor: token_count = get_global_ctx().batch.input_ids.numel() return torch.zeros(token_count, self.embedding_dim, device=device, dtype=dtype) ngram_ids = self._current_ngram_ids().reshape(-1) + if self._q3_reader is not None: + rows = self._q3_reader.gather(ngram_ids.tolist()) + embedded = rows.to(device=device, dtype=dtype) * self._scale.to(device=device, dtype=dtype) + return embedded.view(-1, self.embedding_dim) shard_ids = torch.bucketize(ngram_ids, self._shard_ends, right=True) output = torch.empty( ngram_ids.numel(), @@ -468,6 +547,10 @@ def __init__(self, config: ModelConfig, layer_id: int): def load_host_weights(self, model_path: str, *, dummy: bool = False) -> None: self.ple_embedding.load_host_weights(model_path, dummy=dummy) + def load_q3_ple_weights(self, manifest_path: str) -> None: + """Load an explicit Q3_PLE_32 sidecar for this layer.""" + self.ple_embedding.load_q3_ple_weights(manifest_path) + def _short_conv(self, hidden: torch.Tensor) -> torch.Tensor: batch = get_global_ctx().batch reqs = batch.padded_reqs if batch.is_decode else batch.reqs @@ -600,7 +683,9 @@ class Qwen4ExpDecoderLayer(BaseOP): def __init__(self, config: ModelConfig, layer_id: int): self._layer_id = layer_id self._is_linear = config.is_linear_layer(layer_id) - dense_config = replace(config, expert_quant="none", attn_quant="none") + # Strip routed-expert quantization from the dense attention constructor, but + # preserve the explicit Qwen4 active ``attn_quant`` selection. + dense_config = replace(config, expert_quant="none") if self._is_linear: group = config.linear_attention_group() assert group is not None @@ -614,7 +699,8 @@ def __init__(self, config: ModelConfig, layer_id: int): rms_norm_eps=config.rms_norm_eps, layer_id=layer_id, expert_quant="none", - attn_quant="none", + attn_quant=config.attn_quant, + nvfp4_qkvz=(config.attn_quant == "nvfp4"), ) self.linear_attn.norm = _GatedRMSNorm( group.value_head_dim, @@ -659,14 +745,56 @@ def __init__(self, config: ModelConfig): self._image_token_id = config.image_token_id def load_host_weights(self, model_path: str, *, dummy: bool = False) -> None: + # Modular Qwen4 artifacts carry an explicit Q3_PLE sidecar. Once the + # target marker is present, selecting that sidecar is mandatory: silently + # falling back to the 51-GiB FP8 safetensors table would defeat the artifact's + # bounded host-memory contract. Unmarked source checkpoints keep the original + # FP8 path unchanged. + if not dummy: + from freetoken.checkpoint.qwen4_artifact import ( + load_qwen4_artifact_manifest, + qwen4_text_only_marker, + ) + + artifact = load_qwen4_artifact_manifest(model_path) + if artifact is not None: + self.load_q3_ple_weights(str(artifact.ple_manifest_path)) + return + # A known text-only target marker without its modular manifest is + # incomplete. Do not silently reopen the source FP8 PLE table. + try: + from freetoken.utils import cached_load_hf_config + + if qwen4_text_only_marker(cached_load_hf_config(model_path)): + raise ValueError( + "Qwen4 text-only target is marked but manifest.json is missing" + ) + except ValueError: + raise + except Exception: + # Unmarked hub/source paths retain the historical loader behavior; + # parse_config remains the authoritative marker validator. + pass for layer in self.layers.op_list: if layer.ple is not None: layer.ple.load_host_weights(model_path, dummy=dummy) + def load_q3_ple_weights(self, manifest_paths: str | dict[int, str]) -> None: + """Opt into Q3_PLE_32 using one manifest or a layer-id manifest map.""" + for layer_id, layer in enumerate(self.layers.op_list): + if layer.ple is None: + continue + manifest = manifest_paths[layer_id] if isinstance(manifest_paths, dict) else manifest_paths + layer.ple.load_q3_ple_weights(manifest) + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: hidden = self.embed_tokens.forward(input_ids) mm_embeds = getattr(get_global_ctx().batch, "mm_embeds", None) - if mm_embeds is not None and self._image_token_id is not None: + if mm_embeds is not None and self._image_token_id is None: + raise RuntimeError( + "image inputs are not supported by this text-only Qwen4 modular artifact" + ) + if mm_embeds is not None: mask = input_ids == self._image_token_id slots = int(mask.sum().item()) if slots != mm_embeds.shape[0]: @@ -707,6 +835,9 @@ def encode_images( def load_host_weights(self, model_path: str, *, dummy: bool = False) -> None: self.model.load_host_weights(model_path, dummy=dummy) + def load_q3_ple_weights(self, manifest_paths: str | dict[int, str]) -> None: + self.model.load_q3_ple_weights(manifest_paths) + def forward(self) -> torch.Tensor: hidden = self.model.forward(get_global_ctx().batch.input_ids) return self.lm_head.forward(hidden) diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index 65793983..7f754eb7 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -1,10 +1,12 @@ from __future__ import annotations +from collections.abc import Iterable from typing import Iterator import safetensors import torch from freetoken.distributed import get_tp_info +from freetoken.checkpoint.nvfp4 import encode_bf16_nvfp4 from freetoken.models.loader import iter_weight_files from tqdm import tqdm @@ -22,9 +24,14 @@ ".self_attn.k_proj.weight", ".self_attn.v_proj.weight", ), - ".linear_attn.in_proj.weight": ( + # Canonical runtime-state names match the explicit Qwen4 GDN split: native + # NVFP4 qkv|z and a separate BF16 b|a projection. Keeping these as two + # entries avoids a load-time dequant/re-fusion ambiguity. + ".linear_attn.in_proj_qkvz.weight": ( ".linear_attn.in_proj_qkv.weight", ".linear_attn.in_proj_z.weight", + ), + ".linear_attn.in_proj_ba.weight": ( ".linear_attn.in_proj_b.weight", ".linear_attn.in_proj_a.weight", ), @@ -34,6 +41,44 @@ ), } +ACTIVE_NVFP4_FORMAT = "nvfp4_w4a16_v1" +_ACTIVE_NVFP4_WEIGHT_SUFFIXES = ( + ".self_attn.qkv_proj.weight", + ".self_attn.o_proj.weight", + ".linear_attn.in_proj_qkvz.weight", + ".linear_attn.out_proj.weight", + ".attn_hyper_connection.input_mix_weight_down.weight", + ".attn_hyper_connection.input_mix_weight_up.weight", + ".mlp_hyper_connection.input_mix_weight_down.weight", + ".mlp_hyper_connection.input_mix_weight_up.weight", + ".hyper_connection_mixer.input_mix_weight_down.weight", + ".hyper_connection_mixer.input_mix_weight_up.weight", + ".mlp.shared_expert.gate_up_proj.weight", + ".mlp.shared_expert.down_proj.weight", +) + + +def is_active_nvfp4_weight(name: str) -> bool: + """Whether a fused runtime-state weight belongs to the frozen Qwen4 map.""" + + return name.endswith(_ACTIVE_NVFP4_WEIGHT_SUFFIXES) + + +def iter_active_nvfp4_runtime_entries( + entries: Iterable[tuple[str, torch.Tensor]], +) -> Iterator[tuple[str, torch.Tensor]]: + """Stream fused BF16 state into canonical native NVFP4 FTW entries.""" + + for name, tensor in entries: + if not is_active_nvfp4_weight(name): + yield name, tensor + continue + packed, scale, global_scale = encode_bf16_nvfp4(tensor) + prefix = name.removesuffix(".weight") + yield name, packed + yield prefix + ".weight_scale", scale + yield prefix + ".weight_global", global_scale + def _rename(raw_name: str) -> str | None: if raw_name.startswith("mtp."): @@ -107,6 +152,10 @@ def iter_weights( __all__ = [ "iter_weights", + "encode_bf16_nvfp4", + "ACTIVE_NVFP4_FORMAT", + "is_active_nvfp4_weight", + "iter_active_nvfp4_runtime_entries", "iter_weights_parallel", "load_nvfp4_expert_sources", "load_nvfp4_expert_sources_parallel", diff --git a/python/freetoken/models/weight.py b/python/freetoken/models/weight.py index 6a34f3b9..52bf3ef5 100644 --- a/python/freetoken/models/weight.py +++ b/python/freetoken/models/weight.py @@ -232,8 +232,29 @@ def load_weight( # fails loudly in load_state_dict (strict missing/unexpected expert keys), so the reader # just yields the stored weight tensors regardless of the include_moe_experts flag. from freetoken.checkpoint.ftw import is_ftw_checkpoint, iter_ftw_weights + from freetoken.checkpoint.qwen4_artifact import load_qwen4_artifact_manifest from freetoken.models.config import VISION_KEY_PREFIXES, vision_load_enabled + # A modular Qwen4 artifact keeps active weights in a nested FTW target while + # the root directory owns the manifest, Q3 PLE sidecar, and mixed expert tiers. + # Resolve that target before the ordinary FTW/source branches. The resolver is + # deliberately optional so every unmarked checkpoint follows the historical path. + artifact = load_qwen4_artifact_manifest(model_path) + if artifact is not None: + artifact.verify_active() + active_path = artifact.active_path + if active_path.is_file() and active_path.name == "freetoken_weight.json": + active_path = active_path.parent + if not is_ftw_checkpoint(str(active_path)): + raise ValueError(f"Qwen4 modular active target is not an FTW checkpoint: {active_path}") + for name, tensor in iter_ftw_weights(str(active_path)): + # Qwen4's loader uses ``visual.*`` after canonical renaming; the generic + # prefixes cover model variants that retain ``vision_tower``/``embed_vision``. + if name.startswith(("visual.",) + VISION_KEY_PREFIXES): + continue + yield name, tensor + return + if is_ftw_checkpoint(model_path): # The FTW dense shard stores whatever existed at conversion, including the vision # stack. Vision is opt-in (default OFF, see vision_load_enabled): when it is off the diff --git a/python/freetoken/moe/__init__.py b/python/freetoken/moe/__init__.py index 5fc41911..d18f2cf6 100644 --- a/python/freetoken/moe/__init__.py +++ b/python/freetoken/moe/__init__.py @@ -71,4 +71,19 @@ def create_moe_backend(backend: str) -> BaseMoeBackend: "SUPPORTED_MOE_BACKENDS", "OFFLOAD_MOE_BACKENDS", "is_offload_moe_backend", + "FileExpertSource", + "ExpertSourceError", + "write_expert_sidecar", + "write_expert_sidecar_from_safetensors", + "adapt_expert_tensor_record", ] + +# Kept at module bottom to avoid importing torch/file-I/O helpers while the +# backend registry is initialized by lightweight callers. +from .expert_source import ( # noqa: E402 + ExpertSourceError, + FileExpertSource, + adapt_expert_tensor_record, + write_expert_sidecar, + write_expert_sidecar_from_safetensors, +) diff --git a/python/freetoken/moe/expert_source.py b/python/freetoken/moe/expert_source.py new file mode 100644 index 00000000..42401cb5 --- /dev/null +++ b/python/freetoken/moe/expert_source.py @@ -0,0 +1,986 @@ +"""Bounded, file-backed routed-expert sources. + +The normal :class:`~freetoken.moe.offload_cache.OffloadMoeCache` source is a +resident HostBank. ``FileExpertSource`` is the explicit alternative used by +the Qwen4 text-only tier: one aligned record per expert and no materialised +full-layer tensor. The format is deliberately small and boring so corruption +is detected before the first request is served. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import struct +import threading +from pathlib import Path +from collections.abc import Iterable, Mapping, Sequence +from typing import Any + +import torch + + +MAGIC = b"FTEXPERT1" +# ``FTEXNV4`` was the private Stage 6 fixture format. Keep the reader able to +# reopen those fixtures while making every new artifact unambiguously FTEXPERT1. +LEGACY_MAGIC = b"FTEXNV4\0" +VERSION = 1 +HEADER_BYTES = 4096 +RECORD_BYTES = 2_772_992 +RAW_RECORD_BYTES = 2_772_480 +NUM_EXPERTS = 512 +ALIGNMENT = 4096 + +# Native ModelOpt NVFP4 planes for Qwen4's local expert geometry (H=2560, +# I=640). The six planes are kept in this order in every record. +PLANE_LAYOUT: tuple[tuple[str, int, str], ...] = ( + ("gate_up_packed", 1_638_400, "uint8"), + ("gate_up_scale", 204_800, "float8_e4m3fn"), + ("gate_up_global", 2_560, "float16"), + ("down_packed", 819_200, "uint8"), + ("down_scale", 102_400, "float8_e4m3fn"), + ("down_global", 5_120, "float16"), +) +_PLANE_OFFSETS = {} +_cursor = 0 +for _name, _size, _dtype in PLANE_LAYOUT: + _PLANE_OFFSETS[_name] = _cursor + _cursor += _size +assert _cursor == RAW_RECORD_BYTES + +_HEADER_STRUCT = struct.Struct("<9sIIIIIQQ16s32s32s32s") +_LEGACY_HEADER_STRUCT = struct.Struct("<8sIIIIIQQ16s32s32s") + +# A fixed descriptor in the otherwise-reserved header makes reduced synthetic +# geometry self-describing. Native artifacts retain the exact six-plane +# ModelOpt layout; tests can use tiny tensors without teaching the reader a +# second out-of-band schema. The descriptor is deliberately binary and fixed +# width so two writes of the same inputs are byte-for-byte identical. +_GEOMETRY_MAGIC = b"GEO1" +_GEOMETRY_ENTRY = struct.Struct(" str: + return { + torch.uint8: "uint8", + torch.float8_e4m3fn: "float8_e4m3fn", + torch.float16: "float16", + torch.float32: "float32", + torch.bfloat16: "bfloat16", + }[dtype] + + +class ExpertSourceError(RuntimeError): + """Raised when a tier cannot be trusted or read exactly.""" + + +def _z_path(path: str | os.PathLike[str]) -> Path: + resolved = Path(path).resolve() + drive, _ = os.path.splitdrive(str(resolved)) + # ``Path.drive`` is reliable on Windows; splitdrive also keeps tests clear + # on environments where pathlib's Windows flavour is not selected. + drive = (resolved.drive or drive).upper() + if drive != "Z:": + raise ExpertSourceError(f"file-backed experts must reside on Z:, got {resolved}") + return resolved + + +def _dtype(name: str) -> torch.dtype: + return { + "uint8": torch.uint8, + "float16": torch.float16, + "float8_e4m3fn": torch.float8_e4m3fn, + }[name] + + +def _plane_shape(name: str) -> tuple[int, ...]: + return { + "gate_up_packed": (1280, 1280), + "gate_up_scale": (1280, 160), + "gate_up_global": (1280,), + "down_packed": (2560, 320), + "down_scale": (2560, 40), + "down_global": (2560,), + }[name] + + +def _align_up(value: int, alignment: int = ALIGNMENT) -> int: + return ((int(value) + alignment - 1) // alignment) * alignment + + +def _normalise_fingerprint(value: str | bytes | bytearray | None) -> bytes: + """Return the fixed 32-byte source fingerprint stored in the header. + + Fingerprints are normally SHA-256 bytes or their 64-character hexadecimal + spelling. Short byte strings are accepted for deterministic synthetic + fixtures and zero padded; this mirrors the Stage 6 fixture contract while + still rejecting an accidentally over-wide digest. + """ + + if value is None: + return b"\0" * 32 + if isinstance(value, str): + try: + value = bytes.fromhex(value) + except ValueError: + # Human-readable fixture labels are useful in bounded tests; keep + # them deterministic while documenting that production callers + # should pass SHA-256 bytes/hex. + value = value.encode("utf-8") + raw = bytes(value) + if len(raw) > 32: + raise ValueError("source_fingerprint must be at most 32 bytes") + return raw.ljust(32, b"\0") + + +def _normalise_geometry( + geometry: Mapping[str, Any] | Sequence[tuple[str, Any, Any]] | None, +) -> tuple[tuple[str, tuple[int, ...], torch.dtype], ...]: + """Validate/normalise a six-plane geometry declaration. + + ``geometry`` may map plane names to ``(shape, dtype)`` pairs, or be a + sequence of ``(name, shape, dtype)`` entries. Names must appear exactly in + :data:`PLANE_LAYOUT` order; accepting a mapping is convenient for callers, + but serialization remains ordered and deterministic. + """ + + if geometry is None: + return tuple((name, _plane_shape(name), _dtype(dtype_name)) for name, _size, dtype_name in PLANE_LAYOUT) + if isinstance(geometry, Mapping): + unknown = set(geometry) - set(name for name, _size, _dtype_name in PLANE_LAYOUT) + missing = set(name for name, _size, _dtype_name in PLANE_LAYOUT) - set(geometry) + if unknown or missing: + raise ValueError(f"geometry must contain exactly six planes; missing={sorted(missing)}, unknown={sorted(unknown)}") + entries = [] + for name, _size, _dtype_name in PLANE_LAYOUT: + value = geometry[name] + if not isinstance(value, (tuple, list)) or len(value) != 2: + raise TypeError(f"geometry[{name!r}] must be (shape, dtype)") + shape, dtype = value + entries.append((name, tuple(int(dim) for dim in shape), dtype)) + else: + entries = [] + for item in geometry: + if not isinstance(item, (tuple, list)) or len(item) != 3: + raise TypeError("geometry entries must be (name, shape, dtype)") + name, shape, dtype = item + entries.append((str(name), tuple(int(dim) for dim in shape), dtype)) + expected_names = tuple(name for name, _size, _dtype_name in PLANE_LAYOUT) + actual_names = tuple(name for name, _shape, _dtype in entries) + if actual_names != expected_names: + raise ValueError(f"geometry plane order must be {expected_names}, got {actual_names}") + normalised = [] + for name, shape, dtype in entries: + if not shape or any(dim <= 0 for dim in shape) or len(shape) > 4: + raise ValueError(f"{name} shape must have 1-4 positive dimensions") + if not isinstance(dtype, torch.dtype) or dtype not in _DTYPE_CODES: + raise TypeError(f"unsupported dtype for {name}: {dtype!r}") + normalised.append((name, shape, dtype)) + return tuple(normalised) + + +def _geometry_descriptor( + specs: tuple[tuple[str, tuple[int, ...], torch.dtype], ...], +) -> bytes: + descriptor = bytearray(_GEOMETRY_BYTES) + descriptor[:4] = _GEOMETRY_MAGIC + cursor = 4 + for _name, shape, dtype in specs: + _GEOMETRY_ENTRY.pack_into( + descriptor, + cursor, + _DTYPE_CODES[dtype], + len(shape), + 0, + *(tuple(shape) + (0,) * (4 - len(shape))), + ) + cursor += _GEOMETRY_ENTRY.size + return bytes(descriptor) + + +def _parse_geometry_descriptor(header: bytes) -> tuple[tuple[str, tuple[int, ...], torch.dtype], ...] | None: + if len(header) < _GEOMETRY_BYTES: + return None + offset = _HEADER_STRUCT.size + if header[offset : offset + 4] != _GEOMETRY_MAGIC: + return None + cursor = offset + 4 + specs = [] + try: + for name, _size, _dtype_name in PLANE_LAYOUT: + code, rank, _reserved, d0, d1, d2, d3 = _GEOMETRY_ENTRY.unpack_from(header, cursor) + dtype = _CODE_DTYPES[code] + if not 1 <= rank <= 4: + return None + dims = (d0, d1, d2, d3)[:rank] + if any(dim <= 0 for dim in dims): + return None + specs.append((name, dims, dtype)) + cursor += _GEOMETRY_ENTRY.size + except (KeyError, struct.error): + return None + return tuple(specs) + + +def _canonical_whole_sha256(path: Path) -> str: + """Hash a finalized sidecar with the stored whole-hash field zeroed.""" + + digest = hashlib.sha256() + with path.open("rb") as handle: + offset = 0 + while chunk := handle.read(8 << 20): + if offset <= _WHOLE_HASH_OFFSET < offset + len(chunk): + begin = _WHOLE_HASH_OFFSET - offset + chunk = chunk[:begin] + b"\0" * 32 + chunk[begin + 32 :] + digest.update(chunk) + offset += len(chunk) + return digest.hexdigest() + + +def _sha256_path(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(8 << 20): + digest.update(chunk) + return digest.hexdigest() + + +def _tensor_bytes(value: Any, *, name: str, shape: tuple[int, ...], dtype: torch.dtype) -> bytes: + """Validate one plane and return its native little-endian bytes.""" + + if isinstance(value, (bytes, bytearray, memoryview)): + raw = bytes(value) + expected = int(torch.empty(shape, dtype=dtype).numel()) * torch.empty((), dtype=dtype).element_size() + if len(raw) != expected: + raise ValueError(f"{name} bytes length {len(raw)} != expected {expected}") + return raw + try: + tensor = value if isinstance(value, torch.Tensor) else torch.as_tensor(value) + except Exception as exc: + raise TypeError(f"{name} must be a torch tensor or bytes-like value") from exc + # ModelOpt emits a scalar ``weight_scale_2`` for each source projection. + # Adapters may pass that scalar directly for a global plane; expand only in + # this explicit case and require the declared FP16 dtype afterward. + if name.endswith("_global") and tensor.numel() == 1 and tuple(tensor.shape) != shape: + if not tensor.dtype.is_floating_point: + raise TypeError(f"{name} scalar expansion requires a floating source") + tensor = tensor.to(dtype=dtype).expand(shape) + if tuple(tensor.shape) != shape: + raise ValueError(f"{name} shape {tuple(tensor.shape)} != expected {shape}") + if tensor.dtype != dtype: + raise TypeError(f"{name} dtype {tensor.dtype} != expected {dtype}") + tensor = tensor.detach().to(device="cpu").contiguous() + return tensor.view(torch.uint8).numpy().tobytes() + + +def _record_from_planes( + record: Mapping[str, Any], + specs: tuple[tuple[str, tuple[int, ...], torch.dtype], ...], +) -> bytes: + expected = tuple(name for name, _shape, _dtype in specs) + keys = tuple(key for key in record if key not in {"expert_id", "id"}) + if set(keys) != set(expected): + missing = sorted(set(expected) - set(keys)) + unknown = sorted(set(keys) - set(expected)) + raise ValueError(f"expert record planes mismatch; missing={missing}, unknown={unknown}") + return b"".join( + _tensor_bytes(record[name], name=name, shape=shape, dtype=dtype) + for name, shape, dtype in specs + ) + + +_SOURCE_TENSOR_NAMES = tuple( + f"{projection}.{suffix}" + for projection in ("gate_proj", "up_proj", "down_proj") + for suffix in ("weight", "weight_scale", "weight_scale_2", "input_scale") +) + + +def _source_scalar(value: Any, *, name: str) -> torch.Tensor: + """Validate one ModelOpt source scalar (F32, rank-zero or one element).""" + + tensor = value if isinstance(value, torch.Tensor) else torch.as_tensor(value) + if tensor.dtype != torch.float32: + raise TypeError(f"{name} dtype {tensor.dtype} != expected torch.float32") + if tensor.numel() != 1: + raise ValueError(f"{name} must be a scalar, got shape {tuple(tensor.shape)}") + return tensor.detach().reshape(()) + + +def _as_named_mapping(record: Any) -> dict[str, Any]: + if isinstance(record, Mapping): + return dict(record) + try: + items = list(record) + except TypeError as exc: + raise TypeError("expert tensor record must be a mapping or (name, tensor) iterable") from exc + mapped: dict[str, Any] = {} + for item in items: + if not isinstance(item, (tuple, list)) or len(item) != 2: + raise TypeError("expert tensor entries must be (name, tensor) pairs") + name, value = item + if name in mapped: + raise ValueError(f"duplicate expert tensor name {name!r}") + mapped[str(name)] = value + return mapped + + +def adapt_expert_tensor_record(record: Mapping[str, Any] | Iterable[tuple[str, Any]]) -> dict[str, Any]: + """Adapt one ModelOpt expert's twelve source tensors to six native planes. + + ModelOpt stores ``gate_proj``, ``up_proj`` and ``down_proj`` independently, + each with ``weight``, ``weight_scale``, ``weight_scale_2`` and + ``input_scale``. The native runtime sidecar fuses gate/up along the output + axis, ignores the activation ``input_scale`` (W4A16), and expands each F32 + ``weight_scale_2`` scalar to the FP16 per-output-row global plane. + + A six-plane mapping is returned unchanged (but copied), allowing callers + that have already performed the adaptation to use the same writer API. + """ + + # ``_iter_record_items`` permits an explicit ID alongside the planes; the + # ID is routing metadata, not one of the twelve source tensors. + named_record = _as_named_mapping(record) + clean_record = {key: value for key, value in named_record.items() if key not in {"expert_id", "id"}} + keys = tuple(clean_record) + native_names = tuple(name for name, _shape, _dtype in _normalise_geometry(None)) + if set(keys) == set(native_names): + return {name: clean_record[name] for name in native_names} + if set(keys) != set(_SOURCE_TENSOR_NAMES): + missing = sorted(set(_SOURCE_TENSOR_NAMES) - set(keys)) + unknown = sorted(set(keys) - set(_SOURCE_TENSOR_NAMES)) + raise ValueError(f"expert source tensors mismatch; missing={missing}, unknown={unknown}") + # Validate all twelve source names, including both metadata scalar kinds. + for projection in ("gate_proj", "up_proj", "down_proj"): + _source_scalar(clean_record[f"{projection}.input_scale"], name=f"{projection}.input_scale") + _source_scalar(clean_record[f"{projection}.weight_scale_2"], name=f"{projection}.weight_scale_2") + gate_weight = clean_record["gate_proj.weight"] + up_weight = clean_record["up_proj.weight"] + gate_scale = clean_record["gate_proj.weight_scale"] + up_scale = clean_record["up_proj.weight_scale"] + down_weight = clean_record["down_proj.weight"] + down_scale = clean_record["down_proj.weight_scale"] + # Let the regular six-plane validator provide exact dtype/shape diagnostics + # after these bounded concatenations. Concatenation happens per expert and + # therefore never materialises a layer or model-sized tensor. + gate_weight_t = gate_weight if isinstance(gate_weight, torch.Tensor) else torch.as_tensor(gate_weight) + up_weight_t = up_weight if isinstance(up_weight, torch.Tensor) else torch.as_tensor(up_weight) + gate_scale_t = gate_scale if isinstance(gate_scale, torch.Tensor) else torch.as_tensor(gate_scale) + up_scale_t = up_scale if isinstance(up_scale, torch.Tensor) else torch.as_tensor(up_scale) + down_weight_t = down_weight if isinstance(down_weight, torch.Tensor) else torch.as_tensor(down_weight) + down_scale_t = down_scale if isinstance(down_scale, torch.Tensor) else torch.as_tensor(down_scale) + if gate_weight_t.ndim != up_weight_t.ndim or gate_weight_t.ndim < 1: + raise ValueError("gate_proj.weight and up_proj.weight must have matching rank") + if gate_scale_t.ndim != up_scale_t.ndim or gate_scale_t.ndim < 1: + raise ValueError("gate_proj.weight_scale and up_proj.weight_scale must have matching rank") + gate_rows = int(gate_weight_t.shape[0]) + up_rows = int(up_weight_t.shape[0]) + gate_global = _source_scalar(clean_record["gate_proj.weight_scale_2"], name="gate_proj.weight_scale_2").to(torch.float16).expand(gate_rows) + up_global = _source_scalar(clean_record["up_proj.weight_scale_2"], name="up_proj.weight_scale_2").to(torch.float16).expand(up_rows) + down_global = _source_scalar(clean_record["down_proj.weight_scale_2"], name="down_proj.weight_scale_2").to(torch.float16).expand(int(down_weight_t.shape[0])) + return { + "gate_up_packed": torch.cat((gate_weight_t, up_weight_t), dim=0), + "gate_up_scale": torch.cat((gate_scale_t, up_scale_t), dim=0), + "gate_up_global": torch.cat((gate_global, up_global), dim=0), + "down_packed": down_weight_t, + "down_scale": down_scale_t, + "down_global": down_global, + } + + +def _iter_record_items( + records_or_planes: Any, + *, + num_experts: int, +) -> Iterable[tuple[int, Any]]: + """Yield ``(expert_id, record)`` without materialising the expert bank.""" + + expected_names = {name for name, _size, _dtype_name in PLANE_LAYOUT} + if isinstance(records_or_planes, Mapping): + keys = set(records_or_planes) + source_names = set(_SOURCE_TENSOR_NAMES) + if keys & (expected_names | source_names): + if keys.issubset(source_names): + bank_names = tuple(_SOURCE_TENSOR_NAMES) + elif keys == expected_names: + bank_names = tuple(name for name, _s, _d in PLANE_LAYOUT) + else: + raise ValueError("plane-bank input must contain exactly the six native or twelve source tensor names") + # Stacked tensors/sequences are indexed lazily one expert at a time. + banks = records_or_planes + for eid in range(num_experts): + row = {} + for name in bank_names: + bank = banks[name] + if isinstance(bank, Mapping): + bank_ids = set(bank) + expected_ids = set(range(num_experts)) + if bank_ids != expected_ids: + raise ValueError( + f"plane bank {name} IDs mismatch; missing={sorted(expected_ids - bank_ids)}, " + f"unknown={sorted(bank_ids - expected_ids)}" + ) + if eid not in bank: + raise ValueError(f"missing expert id {eid} in plane bank {name}") + row[name] = bank[eid] + else: + try: + bank_len = len(bank) + except TypeError: + bank_len = None + if bank_len is not None and bank_len != num_experts: + raise ValueError(f"plane bank {name} length {bank_len} != num_experts {num_experts}") + try: + row[name] = bank[eid] + except (IndexError, KeyError, TypeError) as exc: + raise ValueError(f"missing expert id {eid} in plane bank {name}") from exc + yield eid, row + return + for key in sorted(records_or_planes): + if not isinstance(key, int): + raise TypeError("expert mapping keys must be integer expert IDs") + yield key, records_or_planes[key] + return + + for index, item in enumerate(records_or_planes): + expert_id = index + record = item + if isinstance(item, tuple) and len(item) == 2 and isinstance(item[0], int): + expert_id, record = item + elif isinstance(item, Mapping): + explicit = item.get("expert_id", item.get("id", None)) + if explicit is not None: + expert_id = explicit + yield int(expert_id), record + + +def write_expert_sidecar( + path: str | os.PathLike[str], + records_or_planes: Any, + *, + layer_id: int, + source_fingerprint: str | bytes, + num_experts: int = NUM_EXPERTS, + geometry: Mapping[str, Any] | Sequence[tuple[str, Any, Any]] | None = None, + overwrite: bool = True, +) -> dict[str, Any]: + """Stream a deterministic FTEXPERT1 routed-expert sidecar to ``path``. + + ``records_or_planes`` can be an iterable of six-plane mappings (or the + twelve ModelOpt source-tensor mappings, optionally ``(expert_id, mapping)``), + an ``{expert_id: mapping}`` mapping, or a mapping of six/twelve plane names + to stacked tensors/sequences. Exactly one record for every ID + ``0..num_experts-1`` is required. The destination is written to a + ``.partial`` sibling and atomically replaced only after all validation, + payload hashing, and header hashes complete. + """ + + destination = _z_path(path) + if int(num_experts) < 1 or int(num_experts) > NUM_EXPERTS: + raise ValueError(f"num_experts must be in [1, {NUM_EXPERTS}]") + num_experts = int(num_experts) + if not 0 <= int(layer_id) <= 0xFFFFFFFF: + raise ValueError("layer_id must fit an unsigned 32-bit field") + layer_id = int(layer_id) + fingerprint = _normalise_fingerprint(source_fingerprint) + specs = _normalise_geometry(geometry) + sizes = tuple(int(torch.empty(shape, dtype=dtype).numel()) * torch.empty((), dtype=dtype).element_size() for _name, shape, dtype in specs) + raw_record_bytes = sum(sizes) + record_bytes = _align_up(raw_record_bytes) + descriptor = _geometry_descriptor(specs) + destination.parent.mkdir(parents=True, exist_ok=True) + partial = Path(str(destination) + ".partial") + if partial.exists(): + partial.unlink() + payload_hash = hashlib.sha256() + seen: set[int] = set() + sample_ids = tuple(sorted({0, num_experts - 1, num_experts // 4, num_experts // 2, (3 * num_experts) // 4})) + sample_raw: dict[int, bytes] = {} + try: + with partial.open("wb") as handle: + header = bytearray(HEADER_BYTES) + _HEADER_STRUCT.pack_into( + header, + 0, + MAGIC, + VERSION, + HEADER_BYTES, + num_experts, + len(specs), + layer_id, + raw_record_bytes, + record_bytes, + b"nvfp4-qwen4-v1\0\0", + fingerprint, + b"\0" * 32, + b"\0" * 32, + ) + header[_HEADER_STRUCT.size : _HEADER_STRUCT.size + len(descriptor)] = descriptor + handle.write(header) + for expert_id, record in _iter_record_items(records_or_planes, num_experts=num_experts): + if not 0 <= expert_id < num_experts: + raise ValueError(f"expert_id {expert_id} outside [0, {num_experts})") + if expert_id in seen: + raise ValueError(f"duplicate expert_id {expert_id}") + seen.add(expert_id) + if isinstance(record, (bytes, bytearray, memoryview)): + if len(record) != raw_record_bytes: + raise ValueError(f"expert {expert_id} raw bytes length {len(record)} != {raw_record_bytes}") + raw = bytes(record) + elif isinstance(record, Mapping) or isinstance(record, Iterable): + raw = _record_from_planes(adapt_expert_tensor_record(record), specs) + else: + raise TypeError(f"expert {expert_id} must be a plane mapping or raw bytes") + if len(raw) != raw_record_bytes: + raise ValueError(f"expert {expert_id} serialized length {len(raw)} != {raw_record_bytes}") + if expert_id in sample_ids: + sample_raw[expert_id] = raw + padded = raw + b"\0" * (record_bytes - raw_record_bytes) + payload_hash.update(padded) + handle.write(padded) + missing = sorted(set(range(num_experts)) - seen) + if missing: + raise ValueError(f"missing expert IDs: {missing}") + handle.flush() + os.fsync(handle.fileno()) + digest = payload_hash.digest() + # Patch payload hash first, then derive a canonical whole hash over the + # finalized header with only the whole-hash field zeroed. + with partial.open("r+b") as handle: + handle.seek(_PAYLOAD_HASH_OFFSET) + handle.write(digest) + handle.flush() + os.fsync(handle.fileno()) + canonical = _canonical_whole_sha256(partial) + with partial.open("r+b") as handle: + handle.seek(_WHOLE_HASH_OFFSET) + handle.write(bytes.fromhex(canonical)) + handle.flush() + os.fsync(handle.fileno()) + if not overwrite and destination.exists(): + raise FileExistsError(destination) + os.replace(partial, destination) + except Exception: + try: + partial.unlink() + except FileNotFoundError: + pass + raise + whole = _sha256_path(destination) + # Reopen through the production reader and byte-compare deterministic sample + # records before reporting the sidecar complete. The bounded sample set is + # at most five native records (~13.3 MiB at real geometry). + with FileExpertSource( + destination, + expected_sha256=whole, + expected_source_fingerprint=fingerprint, + expected_layer_id=layer_id, + num_experts=num_experts, + verify_hash=True, + ) as source: + for expert_id in sample_ids: + actual = _record_from_planes(source.read_record(expert_id), specs) + if actual != sample_raw[expert_id]: + raise ExpertSourceError(f"expert {expert_id} failed writer reopen verification") + return { + "path": str(destination), + "format": "FTEXPERT1", + "version": VERSION, + "layer_id": layer_id, + "num_experts": num_experts, + "planes": tuple(name for name, _shape, _dtype in specs), + "raw_record_bytes": raw_record_bytes, + "record_bytes": record_bytes, + "source_fingerprint": fingerprint.hex(), + "payload_sha256": digest.hex(), + "canonical_sha256": canonical, + "whole_sha256": whole, + "sha256": whole, + "sample_ids": sample_ids, + } + + +def write_expert_sidecar_from_safetensors( + model_path: str | os.PathLike[str], + path: str | os.PathLike[str], + *, + layer_id: int, + source_fingerprint: str | bytes, + num_experts: int = NUM_EXPERTS, + geometry: Mapping[str, Any] | Sequence[tuple[str, Any, Any]] | None = None, +) -> dict[str, Any]: + """Stream one official ModelOpt NVFP4 layer into ``FTEXPERT1``. + + The source index must contain exactly twelve tensors for every expert. Only + one expert's tensors and one Safetensors mapping remain live at a time; no + full layer or bank is materialized in Python memory. + """ + + folder = Path(model_path).expanduser().resolve() + if not folder.is_dir(): + raise ValueError(f"expert source must be a local checkpoint directory: {folder}") + index_path = folder / "model.safetensors.index.json" + with index_path.open("r", encoding="utf-8") as handle: + weight_map = json.load(handle)["weight_map"] + prefix = f"model.language_model.layers.{int(layer_id)}.mlp.experts." + suffixes = tuple( + f"{projection}.{field}" + for projection in ("gate_proj", "up_proj", "down_proj") + for field in ("weight", "weight_scale", "weight_scale_2", "input_scale") + ) + expected = { + f"{prefix}{expert_id}.{suffix}" + for expert_id in range(int(num_experts)) + for suffix in suffixes + } + actual = {name for name in weight_map if name.startswith(prefix)} + missing = expected - actual + unexpected = actual - expected + if missing or unexpected: + raise ValueError( + f"expert layer {layer_id} tensor set mismatch: " + f"missing={len(missing)} unexpected={len(unexpected)}" + ) + + import safetensors + + def iter_records(): + current_file: str | None = None + current_context = None + current_handle = None + try: + for expert_id in range(int(num_experts)): + record = {} + for suffix in suffixes: + full_name = f"{prefix}{expert_id}.{suffix}" + filename = weight_map[full_name] + if filename != current_file: + if current_context is not None: + current_context.__exit__(None, None, None) + current_context = safetensors.safe_open( + folder / filename, framework="pt", device="cpu" + ) + current_handle = current_context.__enter__() + current_file = filename + record[suffix] = current_handle.get_tensor(full_name) + yield expert_id, record + finally: + if current_context is not None: + current_context.__exit__(None, None, None) + + return write_expert_sidecar( + path, + iter_records(), + layer_id=layer_id, + source_fingerprint=source_fingerprint, + num_experts=num_experts, + geometry=geometry, + ) + + +class FileExpertSource: + """Read fixed NVFP4 expert records from one layer sidecar. + + ``read_record`` returns independent CPU tensors, while ``read_into`` copies + directly into the destination slot planes. Calls are synchronous and each + call owns at most one bounded record buffer. ``read_records`` exposes an + explicit queue-depth guard for future asynchronous readers; this reference + implementation is intentionally serial (depth one) and therefore graph + capture incompatible. + """ + + bank_schema = tuple(name for name, _, _ in PLANE_LAYOUT) + record_bytes = RECORD_BYTES + raw_record_bytes = RAW_RECORD_BYTES + num_experts = NUM_EXPERTS + max_queue_depth = 16 + plane_specs = { + name: (_plane_shape(name), _dtype(dtype_name)) + for name, _size, dtype_name in PLANE_LAYOUT + } + + def __init__( + self, + path: str | os.PathLike[str], + *, + expected_sha256: str | None = None, + expected_source_fingerprint: str | bytes | None = None, + expected_layer_id: int | None = None, + num_experts: int = NUM_EXPERTS, + max_queue_depth: int = 1, + verify_hash: bool = True, + ) -> None: + self.path = _z_path(path) + if not self.path.is_file(): + raise ExpertSourceError(f"missing expert tier: {self.path}") + if not 1 <= int(max_queue_depth) <= 16: + raise ValueError("max_queue_depth must be in [1, 16]") + if not 1 <= int(num_experts) <= NUM_EXPERTS: + raise ValueError(f"num_experts must be in [1, {NUM_EXPERTS}]") + self.num_experts = int(num_experts) + self.requested_queue_depth = int(max_queue_depth) + self.staging_record_bytes = RECORD_BYTES + self.max_staging_records = self.requested_queue_depth + self._lock = threading.Lock() + self._closed = False + self._fd = os.open(str(self.path), os.O_RDONLY | getattr(os, "O_BINARY", 0)) + try: + self._validate_file(expected_source_fingerprint, expected_layer_id) + self.staging_record_bytes = self.record_bytes + self.sha256 = self._hash_file() if verify_hash else None + if expected_sha256 is not None: + expected_sha256 = expected_sha256.lower() + if self.sha256 is None: + self.sha256 = self._hash_file() + if self.sha256 != expected_sha256: + raise ExpertSourceError( + f"expert tier hash mismatch: expected {expected_sha256}, got {self.sha256}" + ) + except Exception: + os.close(self._fd) + raise + self.read_count = 0 + self.bytes_read = 0 + self.max_inflight = 0 + self._inflight = 0 + + @classmethod + def create_synthetic( + cls, + path: str | os.PathLike[str], + *, + num_experts: int = NUM_EXPERTS, + records: Iterable[bytes] | None = None, + source_fingerprint: bytes | None = None, + layer_id: int = 0, + ) -> str: + """Create a tiny deterministic sidecar for tests (never model data).""" + path = _z_path(path) + if num_experts < 1: + raise ValueError("num_experts must be positive") + source_fingerprint = source_fingerprint or hashlib.sha256(b"synthetic").digest() + rows = iter(records) if records is not None else None + + def _rows() -> Iterable[bytes]: + for expert_id in range(num_experts): + raw = next(rows) if rows is not None else bytes([expert_id & 0xFF]) * RAW_RECORD_BYTES + if len(raw) != RAW_RECORD_BYTES: + raise ValueError("synthetic record must contain exactly RAW_RECORD_BYTES") + yield raw + + result = write_expert_sidecar( + path, + _rows(), + layer_id=layer_id, + source_fingerprint=source_fingerprint, + num_experts=num_experts, + ) + # Keep synthetic native-geometry fixtures bounded: the real 512-record + # shape is roughly 1.42 GB and must never be read into one Python bytes + # object merely to produce its verification hash. + return str(result["sha256"]) + + def _validate_file( + self, expected_source_fingerprint: str | bytes | None, expected_layer_id: int | None + ) -> None: + size = self.path.stat().st_size + header = self._read_exact(HEADER_BYTES, 0) + if len(header) != HEADER_BYTES: + raise ExpertSourceError("truncated expert tier header") + try: + if header[: len(MAGIC)] == MAGIC: + ( + magic, + version, + hbytes, + experts, + planes, + layer_id, + raw_bytes, + rec_bytes, + tag, + fingerprint, + payload_hash, + whole_hash, + ) = _HEADER_STRUCT.unpack_from(header) + whole_offset = _WHOLE_HASH_OFFSET + specs = _parse_geometry_descriptor(header) + if specs is None: + specs = tuple((name, _plane_shape(name), _dtype(dtype_name)) for name, _size, dtype_name in PLANE_LAYOUT) + elif header[: len(LEGACY_MAGIC)] == LEGACY_MAGIC: + ( + magic, + version, + hbytes, + experts, + planes, + layer_id, + raw_bytes, + rec_bytes, + tag, + fingerprint, + payload_hash, + ) = _LEGACY_HEADER_STRUCT.unpack_from(header) + whole_hash = b"\0" * 32 + whole_offset = None + specs = tuple((name, _plane_shape(name), _dtype(dtype_name)) for name, _size, dtype_name in PLANE_LAYOUT) + else: + raise ExpertSourceError("unsupported expert tier magic/version/header") + except struct.error as exc: + raise ExpertSourceError("malformed expert tier header") from exc + if magic not in (MAGIC, LEGACY_MAGIC) or version != VERSION or hbytes != HEADER_BYTES: + raise ExpertSourceError("unsupported expert tier magic/version/header") + if experts != self.num_experts or planes != len(PLANE_LAYOUT): + raise ExpertSourceError("expert tier geometry mismatch") + if tag.rstrip(b"\0") != b"nvfp4-qwen4-v1": + raise ExpertSourceError("expert tier layout mismatch") + self.layer_id = int(layer_id) + if expected_layer_id is not None and self.layer_id != int(expected_layer_id): + raise ExpertSourceError(f"expert tier layer mismatch: {self.layer_id} != {int(expected_layer_id)}") + self.plane_layout = tuple( + (name, int(torch.empty(shape, dtype=dtype).numel()) * torch.empty((), dtype=dtype).element_size(), _dtype_name(dtype)) + for name, shape, dtype in specs + ) + self.plane_specs = {name: (shape, dtype) for name, shape, dtype in specs} + self._plane_offsets = {} + cursor = 0 + for name, shape, dtype in specs: + self._plane_offsets[name] = cursor + cursor += int(torch.empty(shape, dtype=dtype).numel()) * torch.empty((), dtype=dtype).element_size() + if cursor != int(raw_bytes) or int(rec_bytes) != _align_up(cursor): + raise ExpertSourceError("expert tier layout mismatch") + self.raw_record_bytes = int(raw_bytes) + self.record_bytes = int(rec_bytes) + expected_size = HEADER_BYTES + self.num_experts * self.record_bytes + if size != expected_size: + raise ExpertSourceError(f"expert tier length mismatch: {size} != {expected_size}") + if expected_source_fingerprint is not None: + try: + expected = _normalise_fingerprint(expected_source_fingerprint) + except ValueError as exc: + raise ExpertSourceError("invalid expected source fingerprint") from exc + if fingerprint != expected: + raise ExpertSourceError("expert tier source fingerprint mismatch") + self.source_fingerprint = fingerprint.hex() + self.payload_sha256 = payload_hash.hex() + if magic == MAGIC and payload_hash == b"\0" * 32: + raise ExpertSourceError("expert tier missing payload hash") + if self.payload_sha256 != "00" * 32: + h = hashlib.sha256() + with self.path.open("rb") as fh: + fh.seek(HEADER_BYTES) + while block := fh.read(8 << 20): + h.update(block) + if h.digest() != payload_hash: + raise ExpertSourceError("expert tier payload hash mismatch") + self.whole_sha256 = _sha256_path(self.path) + self.canonical_sha256 = None + if magic == MAGIC and whole_hash == b"\0" * 32: + raise ExpertSourceError("expert tier missing whole hash") + if whole_offset is not None and whole_hash != b"\0" * 32: + self.canonical_sha256 = _canonical_whole_sha256(self.path) + if self.canonical_sha256 != whole_hash.hex(): + raise ExpertSourceError("expert tier whole hash mismatch") + + def _read_exact(self, size: int, offset: int) -> bytes: + if self._closed: + raise ExpertSourceError("expert tier is closed") + with self._lock: + if hasattr(os, "pread"): + data = os.pread(self._fd, size, offset) + else: # pragma: no cover - Windows Python fallback + os.lseek(self._fd, offset, os.SEEK_SET) + data = os.read(self._fd, size) + if len(data) != size: + raise ExpertSourceError(f"short expert tier read at offset {offset}: {len(data)} != {size}") + return data + + def _hash_file(self) -> str: + return _sha256_path(self.path) + + def _record_bytes(self, expert_id: int) -> bytes: + if self._closed: + raise ExpertSourceError("expert tier is closed") + expert_id = int(expert_id) + if not 0 <= expert_id < self.num_experts: + raise IndexError(f"expert_id {expert_id} outside [0, {self.num_experts})") + self._inflight += 1 + self.max_inflight = max(self.max_inflight, self._inflight) + try: + offset = HEADER_BYTES + expert_id * self.record_bytes + if offset % ALIGNMENT != 0 or self.record_bytes % ALIGNMENT != 0: + raise ExpertSourceError("expert tier record is not aligned to 4096 bytes") + data = self._read_exact(self.record_bytes, offset) + self.read_count += 1 + self.bytes_read += len(data) + return data + finally: + self._inflight -= 1 + + def read_record(self, expert_id: int) -> dict[str, torch.Tensor]: + raw = self._record_bytes(expert_id) + out: dict[str, torch.Tensor] = {} + for name, (shape, dtype) in self.plane_specs.items(): + size = int(torch.empty(shape, dtype=dtype).numel()) * torch.empty((), dtype=dtype).element_size() + offset = self._plane_offsets[name] + # Clone the record-local staging slice so returned tensors remain + # independent after this bounded read buffer is released. + staging = raw[offset : offset + size] + out[name] = torch.frombuffer(bytearray(staging), dtype=dtype).clone().reshape(shape) + return out + + def read_records(self, expert_ids: Iterable[int], *, max_concurrency: int = 1) -> list[dict[str, torch.Tensor]]: + if not 1 <= int(max_concurrency) <= self.requested_queue_depth: + raise ValueError(f"max_concurrency must be in [1, {self.requested_queue_depth}]") + # Serial is deliberate: it keeps staging bounded and is graph-safe only + # outside CUDA capture. A future async implementation may use up to 16. + return [self.read_record(eid) for eid in expert_ids] + + def read_into(self, expert_id: int, destinations: dict[str, torch.Tensor], slot: int) -> int: + """Read one record and copy its six planes into a GPU/CPU cache slot.""" + rows = self.read_record(expert_id) + for name in self.bank_schema: + dst = destinations[name] + if dst.ndim < 1 or not 0 <= slot < dst.shape[0]: + raise ValueError(f"destination slot {slot} invalid for {name}") + if tuple(dst.shape[1:]) != tuple(rows[name].shape): + raise ValueError(f"destination shape mismatch for {name}") + dst[slot].copy_(rows[name], non_blocking=False) + return self.record_bytes + + def close(self) -> None: + if not self._closed: + os.close(self._fd) + self._closed = True + + def __enter__(self) -> "FileExpertSource": + return self + + def __exit__(self, *_exc) -> None: + self.close() + + +__all__ = [ + "FileExpertSource", + "ExpertSourceError", + "PLANE_LAYOUT", + "HEADER_BYTES", + "RECORD_BYTES", + "RAW_RECORD_BYTES", + "MAGIC", + "write_expert_sidecar", + "adapt_expert_tensor_record", +] diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 6ee76406..78c372e9 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -195,6 +195,15 @@ def __post_init__(self) -> None: self.bank_schema = _BANK_SCHEMAS[self.quant_format] self.bank_sources: dict[str, list[torch.Tensor]] = {} self.bank_caches: dict[str, torch.Tensor] = {} + # Optional true file-backed tier. Unlike ``bank_sources`` these entries + # never point at a full-layer HostBank; one aligned record is fetched per + # miss into the existing slot cache. The source object is deliberately + # kept separate so CPU/hybrid callers cannot accidentally treat a tier as + # pageable resident memory. + self.file_sources = {} + self._pending_file_fetches: list[tuple[int, int]] = [] + self._pending_file_materialize = False + self._pending_file_rollback = None # per-layer host residency: the GPU movement paths require "pinned"; LOCKED/PAGEABLE layers decode on the CPU executor and prefill via copy_missing's pageable branch # _unpinned_layers is the derived id set the hot paths test against self.layer_residency: list[str] = [] @@ -277,7 +286,7 @@ def __post_init__(self) -> None: def set_bank_sources( self, - sources: dict[str, list[torch.Tensor]], + sources: dict[str, list[torch.Tensor | None]], layer_residency: list[str] | None = None, ) -> None: """Attach the host (CPU pinned) expert source banks and allocate a GPU slot @@ -321,8 +330,16 @@ def set_bank_sources( for name in self.bank_schema: per_layer = sources[name] assert len(per_layer) == self.num_layers, (name, len(per_layer)) - head = per_layer[0] + heads = [source for source in per_layer if source is not None] + if not heads: + raise ValueError(f"bank {name!r} has no resident shape; attach a file source") + head = heads[0] for layer_id, source in enumerate(per_layer): + if source is None: + # A None row is an explicit file-tier placeholder. The + # corresponding layer must be registered with set_file_sources + # before any request reaches it. + continue assert source.is_contiguous(), f"bank {name!r} layer {layer_id} must be contiguous" assert source.size(0) == self.num_experts, (name, layer_id, source.shape) assert source.shape == head.shape and source.dtype == head.dtype, ( @@ -339,6 +356,65 @@ def set_bank_sources( if self.prefill_overlap: self._init_prefill_overlap_buffers() + def set_file_sources(self, sources: dict[int, object]) -> None: + """Register GPU-only, fixed-record expert tiers by MoE layer. + + ``set_bank_sources`` must contain ``None`` placeholders for these layers, + which makes the absence of a HostBank explicit. File tiers are never + valid for CPU/hybrid decode or prefill overlap because both paths require + resident host tensors and CUDA graph capture cannot contain synchronous + file I/O. + """ + # A mixed cache may use the CPU executor for selected resident layers, + # but a file-backed layer itself is always GPU-only. The per-layer set is + # the precise invariant; rejecting the entire cache would unnecessarily + # disable existing CPU/hybrid support for resident HostBanks. + requested = {int(layer) for layer in sources} + if requested & set(self.cpu_layer_ids): + raise ValueError( + "file-backed expert tiers are GPU-only; CPU/hybrid layers must be resident" + ) + if self.prefill_overlap: + raise ValueError("file-backed expert tiers require prefill_overlap=False") + if not sources: + return + if not self.bank_sources: + # A pure file-tier fixture (and a future all-file policy) has no + # resident HostBank from which to infer cache-plane geometry. The + # source contract supplies the exact per-expert native plane specs. + representative = next(iter(sources.values())) + if tuple(getattr(representative, "bank_schema", ())) != tuple(self.bank_schema): + raise ValueError("file tier bank schema does not match cache quant_format") + specs = getattr(representative, "plane_specs", None) + if not isinstance(specs, dict): + raise ValueError("file source does not declare native plane_specs") + for name in self.bank_schema: + shape, dtype = specs[name] + self.bank_sources[name] = [None] * self.num_layers + self.bank_caches[name] = torch.empty( + (self.cache_size, *shape), dtype=dtype, device=self.device + ) + self.banks = [ + (self.bank_sources[name], self.bank_caches[name]) for name in self.bank_schema + ] + self._build_copy_plan() + for layer_id, source in sources.items(): + layer_id = int(layer_id) + if not 0 <= layer_id < self.num_layers: + raise ValueError(f"file tier layer {layer_id} outside cache geometry") + if not hasattr(source, "bank_schema") or tuple(source.bank_schema) != tuple(self.bank_schema): + raise ValueError("file tier bank schema does not match cache quant_format") + if int(source.num_experts) != self.num_experts: + raise ValueError("file tier expert count does not match cache geometry") + if int(getattr(source, "layer_id", -1)) != layer_id: + raise ValueError( + f"file tier declares layer {getattr(source, 'layer_id', None)}, registered as {layer_id}" + ) + for name in self.bank_schema: + if self.bank_sources[name][layer_id] is not None: + raise ValueError(f"layer {layer_id} already has a resident HostBank for {name}") + self.file_sources[layer_id] = source + def _build_copy_plan(self) -> None: """Precompute the fused multi-bank copy descriptor (base addrs + per-row bytes). @@ -364,7 +440,10 @@ def _build_copy_plan(self) -> None: dst_ptrs, feats = [], [] layer_src_ptrs = [[] for _ in range(self.num_layers)] for per_layer, cache in self.banks: - feat = math.prod(per_layer[0].shape[1:]) * per_layer[0].element_size() + head = next((source for source in per_layer if source is not None), None) + if head is None: + return + feat = math.prod(head.shape[1:]) * head.element_size() if feat % 16 != 0 or cache.data_ptr() % 16 != 0: return # leave fused disabled; copy_missing uses the per-bank path for layer_id, source in enumerate(per_layer): @@ -376,6 +455,12 @@ def _build_copy_plan(self) -> None: # The kernel dereferences these on the GPU, so store each host bank's # device alias (== data_ptr() under UVA identity; differs on # Windows/WDDM). + if source is None: + # A file-tier layer has no host pointer and is never passed + # through the CUDA copy kernel; copy_missing dispatches its + # synchronous record reader below. + layer_src_ptrs[layer_id].append(0) + continue src_dev = device_ptr(source) if src_dev % 16 != 0: return @@ -448,9 +533,18 @@ def rebuild(self, cache_size: int) -> None: torch.cuda.empty_cache() # 3. Reallocate the slot cache from the retained host sources. for name in self.bank_schema: - head = self.bank_sources[name][0] + head = next( + (source for source in self.bank_sources[name] if source is not None), None + ) + if head is None: + representative = next(iter(self.file_sources.values()), None) + if representative is None: + raise ValueError(f"bank {name!r} has no source shape for rebuild") + shape, dtype = representative.plane_specs[name] + else: + shape, dtype = head.shape[1:], head.dtype self.bank_caches[name] = torch.empty( - (cache_size, *head.shape[1:]), dtype=head.dtype, device=self.device + (cache_size, *shape), dtype=dtype, device=self.device ) self.banks = [(self.bank_sources[n], self.bank_caches[n]) for n in self.bank_schema] self._build_copy_plan() # slot caches were reallocated -> refresh fused-copy addrs @@ -798,6 +892,9 @@ def release_prefill_layer(self, layer_id: int) -> None: self._prefill_buffer_released[buffer_id] = True def ensure_experts(self, layer_id: int, expert_ids: torch.Tensor) -> None: + if layer_id in self.file_sources: + self.ensure_file_experts(layer_id, expert_ids) + return from freetoken.moe.offload_kernels import ensure_experts if self.collect_decode_freq: @@ -809,6 +906,62 @@ def ensure_experts(self, layer_id: int, expert_ids: torch.Tensor) -> None: self._pending_whole_layer = False ensure_experts(self, layer_id, expert_ids) + def ensure_file_experts(self, layer_id: int, expert_ids: torch.Tensor) -> None: + """Host-bookkeep file-tier misses and rewrite IDs to GPU cache slots. + + This path intentionally performs a bounded host sync and is not CUDA-graph + capturable. Qwen4 disables graph capture when a file source is attached; + callers that cannot make that guarantee must reject the configuration. + """ + if layer_id not in self.file_sources: + raise ValueError(f"layer {layer_id} has no file source") + if self.decode_target != "gpu": + raise RuntimeError("file-backed experts cannot serve CPU/hybrid decode") + original_shape = tuple(expert_ids.shape) + # File reads happen after the host-side slot plan is installed. Retain + # the small bookkeeping tensors so an I/O error cannot expose a cache + # slot whose six planes were never completely populated. + rollback = ( + self.slot_for_id.clone(), + self.id_of_slot.clone(), + self.usage.clone(), + self.step.clone(), + ) + raw_ids = [int(value) for value in expert_ids.detach().cpu().reshape(-1).tolist()] + self.step += 1 + step = int(self.step.item()) + mapped: list[int] = [] + pending: list[tuple[int, int]] = [] + for expert_id in raw_ids: + if not 0 <= expert_id < self.num_experts: + raise IndexError(f"expert_id {expert_id} outside cache geometry") + slot = int(self.slot_for_id[layer_id, expert_id].item()) + if slot < 0: + free = torch.nonzero(self.id_of_slot < 0, as_tuple=False).reshape(-1) + if free.numel(): + slot = int(free[0].item()) + else: + slot = int(torch.argmin(self.usage).item()) + old = int(self.id_of_slot[slot].item()) + if old >= 0: + self.slot_for_id.view(-1)[old] = -1 + flat_id = layer_id * self.num_experts + expert_id + self.slot_for_id[layer_id, expert_id] = slot + self.id_of_slot[slot] = flat_id + pending.append((slot, expert_id)) + self.usage[slot] = step + mapped.append(slot) + self._pending_src_layer = layer_id + self._pending_whole_layer = False + self._pending_file_materialize = False + # ``copy_missing`` has no use for the LRU scratch arrays on this path, but + # keep num_indices truthful for diagnostics and callers that inspect it. + self.num_indices[0] = len(pending) + self._pending_file_fetches = pending + self._pending_file_rollback = rollback if pending else None + replacement = torch.tensor(mapped, dtype=expert_ids.dtype, device=expert_ids.device).reshape(original_shape) + expert_ids.copy_(replacement) + def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None: """Capped-fetch LRU for the hybrid backend. @@ -831,6 +984,31 @@ def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None ) def materialize_layer(self, layer_id: int) -> None: + if layer_id in self.file_sources: + if self.cache_size < self.num_experts: + raise RuntimeError("file-tier prefill requires one slot per expert") + self._pending_file_rollback = ( + self.slot_for_id.clone(), + self.id_of_slot.clone(), + self.usage.clone(), + self.step.clone(), + ) + self.step += 1 + step = int(self.step.item()) + self.slot_for_id[layer_id].fill_(-1) + for expert_id in range(self.num_experts): + slot = expert_id + old = int(self.id_of_slot[slot].item()) + if old >= 0: + self.slot_for_id.view(-1)[old] = -1 + self.slot_for_id[layer_id, expert_id] = slot + self.id_of_slot[slot] = layer_id * self.num_experts + expert_id + self.usage[slot] = step + self._pending_src_layer = layer_id + self._pending_whole_layer = True + self._pending_file_materialize = True + self._pending_file_fetches = [] + return from freetoken.moe.offload_kernels import materialize_layer self._pending_src_layer = layer_id @@ -838,12 +1016,25 @@ def materialize_layer(self, layer_id: int) -> None: materialize_layer(self, layer_id) def reset(self) -> None: - from freetoken.moe.offload_kernels import reset_cache + if self.device.type == "cuda": + from freetoken.moe.offload_kernels import reset_cache - reset_cache(self) + reset_cache(self) + else: + # The production path is CUDA-only, but a CPU synthetic FileExpertSource + # fixture still needs deterministic reset/reuse semantics without + # attempting to launch a Triton kernel on a CPU tensor. + self.slot_for_id.fill_(-1) + self.id_of_slot.fill_(-1) + self.usage.zero_() + self.step.zero_() + self.num_indices.zero_() # Per-expert recency is not cache_size-shaped, so reset_cache leaves it alone; wipe # it here so a new sequence starts with cold hybrid fetch priorities. self.expert_recency.fill_(-1) + self._pending_file_fetches = [] + self._pending_file_materialize = False + self._pending_file_rollback = None def reset_stats(self) -> None: self.prefill_hit_rows = 0 @@ -969,6 +1160,33 @@ def copy_missing(self) -> None: assert self.banks, "set_bank_sources must register the banks first" layer_id = self._pending_src_layer assert layer_id is not None, "no staged misses (ensure_experts/materialize_layer first)" + if layer_id in self.file_sources: + source = self.file_sources[layer_id] + if self._pending_file_materialize: + pairs = [(expert_id, expert_id) for expert_id in range(self.num_experts)] + else: + pairs = list(self._pending_file_fetches) + destinations = { + name: cache for name, (_, cache) in zip(self.bank_schema, self.banks) + } + try: + for slot, expert_id in pairs: + source.read_into(expert_id, destinations, slot) + except Exception: + if self._pending_file_rollback is not None: + slot_for_id, id_of_slot, usage, step = self._pending_file_rollback + self.slot_for_id.copy_(slot_for_id) + self.id_of_slot.copy_(id_of_slot) + self.usage.copy_(usage) + self.step.copy_(step) + self._pending_file_fetches = [] + self._pending_file_materialize = False + self._pending_file_rollback = None + raise + self._pending_file_fetches = [] + self._pending_file_materialize = False + self._pending_file_rollback = None + return if layer_id in self._unpinned_layers: if not self._pending_whole_layer: raise RuntimeError( diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index adef3e39..22567e43 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -174,6 +174,8 @@ def _infer_tool_call_parser(model_path: str) -> str: if ( "qwen3_5" in marker or "qwen3.5" in marker + or "qwen4_exp" in marker + or "qwen4exp" in marker or ("qwen3" in marker and "coder" in marker) ): return "qwen3_coder" diff --git a/tests/checkpoint/test_convert_metadata.py b/tests/checkpoint/test_convert_metadata.py index a58b00b7..124d6bf9 100644 --- a/tests/checkpoint/test_convert_metadata.py +++ b/tests/checkpoint/test_convert_metadata.py @@ -5,7 +5,11 @@ import torch -from freetoken.checkpoint.convert import _copy_metadata +from freetoken.checkpoint.convert import ( + _copy_metadata, + _iter_qwen4_modular_dense_entries, + convert_checkpoint, +) from freetoken.checkpoint.ftw import FTWReader, FTWWriter, iter_ftw_weights @@ -46,6 +50,108 @@ def test_copy_metadata_keeps_only_qwen4_host_mapped_shards(tmp_path: Path) -> No ] +def test_modular_metadata_does_not_copy_source_ple_payload(tmp_path: Path) -> None: + source = tmp_path / "source" + output = tmp_path / "output" + source.mkdir() + (source / "config.json").write_text("{}", encoding="utf-8") + (source / "model-plefp8-00000.safetensors").write_bytes(b"ple") + (source / "model.safetensors.index.json").write_text( + json.dumps({ + "weight_map": { + "model.language_model.layers.0.ple.ple_embedding.ngram_embedding.shard_0.weight": "model-plefp8-00000.safetensors" + } + }), + encoding="utf-8", + ) + copied = _copy_metadata(str(source), str(output), include_host_mapped_weights=False) + assert copied == ["config.json"] + assert not (output / "model-plefp8-00000.safetensors").exists() + + +def test_modular_dense_stream_quantizes_map_and_excludes_vision() -> None: + active = torch.ones((2, 16), dtype=torch.bfloat16) + protected = torch.ones((4,), dtype=torch.bfloat16) + names = [ + name + for name, _tensor in _iter_qwen4_modular_dense_entries( + [ + ("model.layers.0.self_attn.o_proj.weight", active), + ("model.layers.0.input_layernorm.weight", protected), + ("visual.blocks.0.weight", active), + ] + ) + ] + assert names == [ + "model.layers.0.self_attn.o_proj.weight", + "model.layers.0.self_attn.o_proj.weight_scale", + "model.layers.0.self_attn.o_proj.weight_global", + "model.layers.0.input_layernorm.weight", + ] + + +def test_convert_checkpoint_builds_marked_modular_active_ftw(tmp_path: Path, monkeypatch) -> None: + source = tmp_path / "source" + output = tmp_path / "target" + source.mkdir() + (source / "config.json").write_text( + json.dumps({"architectures": ["Qwen4ExpForCausalLM"], "vision_config": {}}), + encoding="utf-8", + ) + + class FakeEngineConfig: + def __init__(self, **_kwargs): + self.model_config = type( + "ModelConfig", + (), + { + "architectures": ("Qwen4ExpForCausalLM",), + "expert_quant": "nvfp4", + "is_moe": True, + }, + )() + + active = torch.ones((2, 16), dtype=torch.bfloat16) + protected = torch.arange(4, dtype=torch.bfloat16) + + def fake_load_weight(_path, device, *, include_moe_experts): + assert device.type == "cpu" + assert include_moe_experts is False + return iter( + [ + ("model.layers.0.self_attn.o_proj.weight", active), + ("model.layers.0.input_layernorm.weight", protected), + ] + ) + + import freetoken.engine.config as engine_config + import freetoken.models.weight as weight_module + + monkeypatch.setattr(engine_config, "EngineConfig", FakeEngineConfig) + monkeypatch.setattr(weight_module, "load_weight", fake_load_weight) + inventory = "a" * 64 + index = convert_checkpoint( + str(source), + str(output), + artifact_format="qwen4_modular_v1", + source_inventory_sha256=inventory, + shard_limit=4096 * 16, + ) + + assert index["source_inventory_sha256"] == inventory + assert index["counts"] == {"weight": 4, "experts_bank": 0} + config = json.loads((output / "config.json").read_text(encoding="utf-8")) + assert config["freetoken_text_only"] == "qwen4_text_only_v1" + assert config["freetoken_active_quant"] == "nvfp4_w4a16_v1" + loaded = dict(iter_ftw_weights(str(output / "qwen4-active-v1.ftw"), workers=1)) + assert set(loaded) == { + "model.layers.0.self_attn.o_proj.weight", + "model.layers.0.self_attn.o_proj.weight_scale", + "model.layers.0.self_attn.o_proj.weight_global", + "model.layers.0.input_layernorm.weight", + } + + def test_ftw_buffered_reader_works_on_the_current_platform(tmp_path: Path) -> None: output = tmp_path / "ftw" writer = FTWWriter(str(output), shard_limit=4096) diff --git a/tests/checkpoint/test_q3_ple.py b/tests/checkpoint/test_q3_ple.py new file mode 100644 index 00000000..004bd94b --- /dev/null +++ b/tests/checkpoint/test_q3_ple.py @@ -0,0 +1,185 @@ +from __future__ import annotations + +import hashlib +import json +import os +import shutil +import sys +from pathlib import Path +from uuid import uuid4 + +import pytest +import torch + +from freetoken.checkpoint.q3_ple import ( + ALIGN, + BLOCK_BYTES, + BLOCK_VALUES, + Q3PLEReader, + ROW_BYTES, + ROW_VALUES, +) +ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(ROOT.parent.parent / "scripts")) +from q3_ple_32_reference import dequantize_row, encode_table + + +def _write_fixture(fixture_root: Path) -> tuple[Path, Path, int]: + fixture_root.mkdir(parents=True, exist_ok=True) + manifest_path = fixture_root / "ple-q3.json" + data_path = fixture_root / "ple-q3-000.bin" + rows = [] + for row in range(20): + rows.append([((row + 1) * 0.125) * ((i % 17) - 8) for i in range(ROW_VALUES)]) + # Explicit zero/extreme rows exercise every one of the five block decoders. + rows[0] = [0.0] * ROW_VALUES + rows[-1] = [(-1.0 if i & 1 else 1.0) * 448.0 for i in range(ROW_VALUES)] + encoded = encode_table(rows, refinement_passes=2, scale_dtype="bf16") + split = 9 * ROW_BYTES + second_offset = ALIGN + payload = encoded[:split] + bytes(second_offset - split) + encoded[split:] + data_path.write_bytes(payload) + segments = [] + for first, end, offset in ((0, 9, 0), (9, 20, second_offset)): + segment_bytes = encoded[first * ROW_BYTES : end * ROW_BYTES] + segments.append( + { + "first_row": first, + "end_row": end, + "data_offset": offset, + "byte_length": len(segment_bytes), + "sha256": hashlib.sha256(segment_bytes).hexdigest(), + } + ) + manifest = { + "format": "q3_ple_32", + "version": 1, + "endianness": "little", + "block_values": BLOCK_VALUES, + "block_bytes": BLOCK_BYTES, + "row_values": ROW_VALUES, + "row_bytes": ROW_BYTES, + "rows": len(rows), + "payload_bytes": len(encoded), + "file_bytes": len(payload), + "data_file": data_path.name, + "weight_scale": 1.25, + "sha256": hashlib.sha256(payload).hexdigest(), + "segments": segments, + } + manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8") + return manifest_path, data_path, len(rows) + + +@pytest.fixture(scope="module") +def fixture_paths(): + root = ROOT / ".stage6-test-fixtures" / uuid4().hex + assert root.drive.upper() == "Z:" + try: + yield _write_fixture(root) + finally: + shutil.rmtree(root, ignore_errors=True) + + +def test_q3_reader_constants_and_ordered_gather(fixture_paths): + manifest, _, row_count = fixture_paths + with Q3PLEReader(manifest) as reader: + assert reader.row_count == row_count + assert reader.total_payload_bytes == row_count * 70 + assert reader.gather([11, 0, 11]).shape == (3, ROW_VALUES) + scaled = reader.gather16(list(range(16)), apply_weight_scale=True) + assert scaled.dtype == torch.bfloat16 + assert scaled.shape == (16, ROW_VALUES) + assert torch.equal(scaled[0], reader.gather([0], apply_weight_scale=True)[0]) + + +def test_q3_reader_matches_authoritative_codec(fixture_paths): + manifest, data, _ = fixture_paths + # The fixture uses two segments and includes alignment padding between them. + encoded = data.read_bytes() + with Q3PLEReader(manifest) as reader: + for row in (0, 1, 8, 9, 10, 19): + raw_row = (encoded[row * ROW_BYTES : (row + 1) * ROW_BYTES] + if row < 9 else encoded[ALIGN + (row - 9) * ROW_BYTES : ALIGN + (row - 8) * ROW_BYTES]) + expected = torch.tensor( + dequantize_row(raw_row), + dtype=torch.bfloat16, + ) + assert torch.equal(reader.gather([row])[0], expected) + + +def test_q3_reader_all_block_boundaries_and_random_order(fixture_paths): + manifest, data, _ = fixture_paths + encoded = data.read_bytes() + with Q3PLEReader(manifest) as reader: + rows = reader.gather([19, 9, 0, 19, 8]) + assert rows.shape == (5, ROW_VALUES) + assert torch.equal(rows[0], rows[3]) + for row_index, row in zip((19, 9, 0, 19, 8), rows): + offset = row_index * ROW_BYTES if row_index < 9 else ALIGN + (row_index - 9) * ROW_BYTES + expected = torch.tensor(dequantize_row(encoded[offset : offset + ROW_BYTES]), dtype=torch.bfloat16) + assert torch.equal(row, expected) + for boundary in (0, 31, 32, 63, 64, 95, 96, 127, 128, 159): + assert row[boundary] == expected[boundary] + + +@pytest.mark.parametrize( + "field,value", + [("version", 2), ("endianness", "big"), ("row_bytes", 69), ("sha256", "0" * 64)], +) +def test_q3_reader_rejects_bad_manifest(fixture_paths, field, value): + manifest, _, _ = fixture_paths + raw = json.loads(manifest.read_text(encoding="utf-8")) + raw[field] = value + bad = manifest.with_name(f"bad-{field}.json") + bad.write_text(json.dumps(raw), encoding="utf-8") + try: + with pytest.raises(ValueError): + Q3PLEReader(bad) + finally: + bad.unlink(missing_ok=True) + + +def test_q3_reader_rejects_gap_overlap_and_truncation(fixture_paths): + manifest, data, _ = fixture_paths + raw = json.loads(manifest.read_text(encoding="utf-8")) + raw["segments"][1]["first_row"] = 6 + bad = manifest.with_name("bad-segments.json") + bad.write_text(json.dumps(raw), encoding="utf-8") + try: + with pytest.raises(ValueError): + Q3PLEReader(bad) + finally: + bad.unlink(missing_ok=True) + + truncated = data.with_name("truncated.bin") + truncated.write_bytes(data.read_bytes()[:-1]) + raw = json.loads(manifest.read_text(encoding="utf-8")) + raw["data_file"] = truncated.name + raw["file_bytes"] -= 1 + bad = manifest.with_name("bad-truncated.json") + bad.write_text(json.dumps(raw), encoding="utf-8") + try: + with pytest.raises(ValueError): + Q3PLEReader(bad) + finally: + bad.unlink(missing_ok=True) + truncated.unlink(missing_ok=True) + + +def test_q3_reader_rejects_corrupt_segment_hash(fixture_paths): + manifest, _, _ = fixture_paths + raw = json.loads(manifest.read_text(encoding="utf-8")) + raw["segments"][1]["sha256"] = "f" * 64 + bad = manifest.with_name("bad-segment-hash.json") + bad.write_text(json.dumps(raw), encoding="utf-8") + try: + with pytest.raises(ValueError, match="segment hash mismatch"): + Q3PLEReader(bad) + finally: + bad.unlink(missing_ok=True) + + +def test_q3_reader_requires_z_backing(): + with pytest.raises((ValueError, FileNotFoundError)): + Q3PLEReader("C:\\q3-ple\\ple-q3.json") diff --git a/tests/checkpoint/test_q3_ple_writer.py b/tests/checkpoint/test_q3_ple_writer.py new file mode 100644 index 00000000..f3d3c22b --- /dev/null +++ b/tests/checkpoint/test_q3_ple_writer.py @@ -0,0 +1,210 @@ +from __future__ import annotations + +import hashlib +import importlib.util +import json +import shutil +import sys +from uuid import uuid4 +from pathlib import Path + +import pytest +import torch + +from freetoken.checkpoint.q3_ple import ( + ALIGN, + BLOCK_BYTES, + ROW_BYTES, + ROW_VALUES, + Q3PLEReader, + quantize_block, + write_q3_ple_from_safetensors, + write_q3_ple_sidecar, +) + + +def _rows(count: int): + for row_index in range(count): + yield [((row_index + 1) * 0.125) * ((column % 17) - 8) for column in range(ROW_VALUES)] + + +def _load_authoritative_reference(): + reference_path = Path(__file__).resolve().parents[4] / "scripts" / "q3_ple_32_reference.py" + if not reference_path.exists(): + pytest.skip("authoritative Q3_PLE_32 reference script is not present in this checkout") + spec = importlib.util.spec_from_file_location("q3_ple_32_authoritative_reference", reference_path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +@pytest.fixture() +def z_fixture_dir() -> Path: + root = Path(__file__).resolve().parents[2] / ".stage7-q3-writer-fixtures" / uuid4().hex + assert root.drive.upper() == "Z:" + root.mkdir(parents=True) + try: + yield root + finally: + shutil.rmtree(root, ignore_errors=True) + + +def test_writer_matches_reference_vectors_and_reader(z_fixture_dir: Path) -> None: + # These vectors are the byte-for-byte values from the authoritative + # scripts/q3_ple_32_reference.py codec (two refinement passes, BF16 scale). + assert quantize_block([0.0] * 32).hex() == "0000244992244992244992244992" + assert quantize_block([(index - 16) / 4.0 for index in range(32)]).hex() == ( + "9c3f492249dab69124dbb6b6edff" + ) + + data_path = z_fixture_dir / "ple-q3.bin" + manifest_path = z_fixture_dir / "ple-q3.json" + manifest = write_q3_ple_sidecar( + _rows(5), + data_path, + manifest_path, + source_fingerprint="A" * 64, + weight_scale=1.25, + segment_rows=2, + ) + assert manifest["source_fingerprint"] == "a" * 64 + assert manifest["weight_scale"] == 1.25 + assert manifest["payload_bytes"] == 5 * ROW_BYTES + assert manifest["file_bytes"] == data_path.stat().st_size + assert [segment["first_row"] for segment in manifest["segments"]] == [0, 2, 4] + assert manifest["storage_layout"] == "contiguous_rows_v1" + assert manifest["file_bytes"] == manifest["payload_bytes"] == 5 * ROW_BYTES + assert [segment["data_offset"] for segment in manifest["segments"]] == [0, 2 * ROW_BYTES, 4 * ROW_BYTES] + assert all( + segment["byte_length"] == (2 if segment["first_row"] < 4 else 1) * ROW_BYTES + for segment in manifest["segments"] + ) + + raw = data_path.read_bytes() + assert hashlib.sha256(raw).hexdigest() == manifest["sha256"] + loaded = json.loads(manifest_path.read_text(encoding="utf-8")) + assert loaded == manifest + with Q3PLEReader(manifest_path) as reader: + gathered = reader.gather([4, 0, 4]) + assert gathered.shape == (3, ROW_VALUES) + assert gathered[0].equal(gathered[2]) + scaled = reader.gather([0], apply_weight_scale=True) + assert scaled.equal(reader.gather([0]) * 1.25) + + +def test_writer_payload_is_byte_identical_to_authoritative_reference(z_fixture_dir: Path) -> None: + reference = _load_authoritative_reference() + rows = list(_rows(3)) + expected = reference.encode_table(rows, refinement_passes=2, scale_dtype="bf16") + manifest = write_q3_ple_sidecar( + iter(rows), + z_fixture_dir / "reference.bin", + z_fixture_dir / "reference.json", + source_fingerprint="e" * 64, + weight_scale=1.0, + segment_rows=128, + ) + assert (z_fixture_dir / "reference.bin").read_bytes() == expected + assert manifest["file_bytes"] == len(expected) + + +def test_writer_consumes_rows_once_and_rejects_nonfinite_without_finalizing(z_fixture_dir: Path) -> None: + data_path = z_fixture_dir / "ple-q3.bin" + manifest_path = z_fixture_dir / "ple-q3.json" + consumed = 0 + + def source_rows(): + nonlocal consumed + consumed += 1 + yield [0.0] * ROW_VALUES + consumed += 1 + bad = [0.0] * ROW_VALUES + bad[31] = float("nan") + yield bad + + with pytest.raises(ValueError, match="non-finite"): + write_q3_ple_sidecar( + source_rows(), + data_path, + manifest_path, + source_fingerprint="b" * 64, + weight_scale=1.0, + segment_rows=1, + ) + assert consumed == 2 + assert not data_path.exists() + assert not manifest_path.exists() + assert not list(z_fixture_dir.glob(".*.partial-*")) + + +@pytest.mark.parametrize( + "kwargs", + [ + {"source_fingerprint": "short", "weight_scale": 1.0}, + {"source_fingerprint": "c" * 64, "weight_scale": float("inf")}, + {"source_fingerprint": "c" * 64, "weight_scale": 1.0, "segment_rows": 0}, + ], +) +def test_writer_rejects_bad_integrity_metadata(z_fixture_dir: Path, kwargs: dict) -> None: + with pytest.raises(ValueError): + write_q3_ple_sidecar( + _rows(1), + z_fixture_dir / "ple-q3.bin", + z_fixture_dir / "ple-q3.json", + **kwargs, + ) + + +def test_writer_requires_z_backing() -> None: + forbidden = Path("C:/stage7-q3-writer-must-not-create") + with pytest.raises(ValueError, match="Z:"): + write_q3_ple_sidecar( + _rows(1), + forbidden / "ple-q3.bin", + forbidden / "ple-q3.json", + source_fingerprint="d" * 64, + weight_scale=1.0, + ) + + +def test_production_writer_streams_safetensor_shards_in_source_order(z_fixture_dir: Path) -> None: + from safetensors.torch import save_file + + prefix = "model.language_model.layers.2.ple.ple_embedding.ngram_embedding" + weight_map = {} + source_rows = [] + for part in range(2): + key = f"{prefix}.shard_{part}.weight" + filename = f"model-plefp8-{part:05d}.safetensors" + rows = ( + torch.tensor(list(_rows(2)), dtype=torch.float32) + float(part) + ).to(torch.float8_e4m3fn).contiguous() + tensors = {key: rows} + if part == 0: + tensors[prefix + ".weight_scale"] = torch.tensor(0.5, dtype=torch.bfloat16) + weight_map[prefix + ".weight_scale"] = filename + save_file(tensors, z_fixture_dir / filename) + weight_map[key] = filename + source_rows.extend(rows.float().tolist()) + (z_fixture_dir / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": weight_map}), encoding="utf-8" + ) + manifest = write_q3_ple_from_safetensors( + z_fixture_dir, + z_fixture_dir / "ple-q3.bin", + z_fixture_dir / "ple-q3.json", + layer_id=2, + split_parts=2, + source_fingerprint="f" * 64, + rows_per_chunk=1, + segment_rows=3, + ) + reference = _load_authoritative_reference() + assert (z_fixture_dir / "ple-q3.bin").read_bytes() == reference.encode_table( + source_rows, refinement_passes=2, scale_dtype="bf16" + ) + assert manifest["rows"] == 4 + assert manifest["file_bytes"] == 4 * ROW_BYTES + assert manifest["weight_scale"] == 0.5 diff --git a/tests/checkpoint/test_qwen4_artifact.py b/tests/checkpoint/test_qwen4_artifact.py new file mode 100644 index 00000000..096d7cff --- /dev/null +++ b/tests/checkpoint/test_qwen4_artifact.py @@ -0,0 +1,490 @@ +from __future__ import annotations + +import json +import hashlib +import shutil +from pathlib import Path +from uuid import uuid4 + +import pytest +import torch +from types import SimpleNamespace + +from freetoken.checkpoint.qwen4_artifact import ( + FORMAT, + TEXT_ONLY_MARKER, + build_qwen4_modular_artifact, + configure_mixed_expert_sources, + build_mixed_expert_sources, + finalize_qwen4_modular_manifest, + load_qwen4_artifact_manifest, + qwen4_text_only_marker, +) +from freetoken.checkpoint.ftw import FTWWriter +from freetoken.checkpoint.q3_ple import ROW_VALUES, write_q3_ple_sidecar +from freetoken.models.qwen4_exp.weight import iter_active_nvfp4_runtime_entries +from freetoken.models.weight import load_weight +from freetoken.moe.expert_source import ( + FileExpertSource, + RAW_RECORD_BYTES, + write_expert_sidecar, +) + + +@pytest.fixture +def z_fixture_dir(): + root = Path.cwd() / ".stage7-test-fixtures" / uuid4().hex + root.mkdir(parents=True, exist_ok=False) + try: + assert (root.drive or "").upper() == "Z:", root + yield root + finally: + shutil.rmtree(root, ignore_errors=True) + + +class _FakeCache: + bank_schema = FileExpertSource.bank_schema + decode_target = "gpu" + prefill_overlap = False + cache_size = 512 + num_layers = 2 + num_experts = 1 + cpu_layer_ids = frozenset() + + def set_bank_sources(self, sources, layer_residency=None): + self.bank_sources = sources + self.layer_residency = layer_residency + + def set_file_sources(self, sources): + self.file_sources = sources + + +def _manifest(root: Path, sidecar: Path, digest: str) -> Path: + source_fingerprint = hashlib.sha256(b"synthetic").hexdigest() + active_root = root / "qwen4-active-v1.ftw" + active_root.mkdir(exist_ok=True) + active_index = active_root / "freetoken_weight.json" + active_index.write_text( + json.dumps({"source_inventory_sha256": source_fingerprint}), encoding="utf-8" + ) + active_digest = __import__("hashlib").sha256(active_index.read_bytes()).hexdigest() + ple_manifest = root / "ple-q3.json" + ple_manifest.write_text( + json.dumps({"source_fingerprint": source_fingerprint}), encoding="utf-8" + ) + resident_sidecar = root / "experts-L01.nvfp4" + resident_digest = FileExpertSource.create_synthetic( + resident_sidecar, num_experts=1, records=[bytes([8]) * RAW_RECORD_BYTES], layer_id=1 + ) + config_path = root / "config.json" + config_path.write_bytes(b"") + config_digest = hashlib.sha256(b"").hexdigest() + data = { + "format": FORMAT, + "version": 1, + "artifact_schema": FORMAT, + "text_only": True, + "source": {"repository": "synthetic", "revision": "test", "inventory_sha256": source_fingerprint}, + "minimum_freetoken_commit": "0" * 40, + "tvm_ffi_patch_sha256": "0" * 64, + "active": { + "format": "nvfp4_w4a16_v1", + "path": active_root.name, + "bytes": active_index.stat().st_size, + "files": [{"path": str(active_index.relative_to(root)), "bytes": active_index.stat().st_size, "sha256": active_digest}], + }, + "ple": { + "format": "q3_ple_32", + "manifest": "ple-q3.json", + "data_bytes": 0, + "sha256": "0" * 64, + "required_volume": "Z:", + }, + "experts": { + "format": "ftexpert1_nvfp4_v1", + "files": [ + {"layer": 0, "path": sidecar.name, "bytes": sidecar.stat().st_size, "sha256": digest, "source_fingerprint": source_fingerprint}, + {"layer": 1, "path": resident_sidecar.name, "bytes": resident_sidecar.stat().st_size, "sha256": resident_digest, "source_fingerprint": source_fingerprint}, + ], + "file_tier_layers": [0], + "resident_layers": [1], + "required_volume": "Z:", + }, + "metadata": {"files": [{"path": "config.json", "bytes": 0, "sha256": config_digest}]}, + } + data["complete_artifact_fingerprint"] = hashlib.sha256( + json.dumps(data, sort_keys=True, separators=(",", ":")).encode("utf-8") + ).hexdigest() + path = root / "manifest.json" + path.write_text(json.dumps(data), encoding="utf-8") + return path + + +def test_qwen4_marker_is_explicit_and_unknown_fails_closed(): + class Config: + freetoken_text_only = TEXT_ONLY_MARKER + + assert qwen4_text_only_marker(Config()) is True + Config.freetoken_text_only = "future_policy" + with pytest.raises(ValueError, match="unsupported freetoken_text_only"): + qwen4_text_only_marker(Config()) + + +def test_active_component_hash_fails_closed(z_fixture_dir): + sidecar = z_fixture_dir / "experts-L00.nvfp4" + digest = FileExpertSource.create_synthetic( + sidecar, num_experts=1, records=[bytes([7]) * RAW_RECORD_BYTES], layer_id=0 + ) + manifest = load_qwen4_artifact_manifest(_manifest(z_fixture_dir, sidecar, digest), require=True, allow_synthetic_geometry=True) + assert manifest is not None + manifest.verify_active() + manifest.active_files[0].path.write_bytes(b"tampered") + with pytest.raises(ValueError, match="length mismatch|SHA-256 mismatch"): + manifest.verify_active() + + +def test_manifest_unknown_schema_and_fingerprint_fail_closed(z_fixture_dir): + sidecar = z_fixture_dir / "experts-L00.nvfp4" + digest = FileExpertSource.create_synthetic( + sidecar, num_experts=1, records=[bytes([7]) * RAW_RECORD_BYTES], layer_id=0 + ) + path = _manifest(z_fixture_dir, sidecar, digest) + data = json.loads(path.read_text(encoding="utf-8")) + data["artifact_schema"] = "future-v99" + path.write_text(json.dumps(data), encoding="utf-8") + with pytest.raises(ValueError, match="artifact_schema"): + load_qwen4_artifact_manifest(path, require=True, allow_synthetic_geometry=True) + + data["artifact_schema"] = FORMAT + data["complete_artifact_fingerprint"] = "f" * 64 + path.write_text(json.dumps(data), encoding="utf-8") + with pytest.raises(ValueError, match="fingerprint mismatch"): + load_qwen4_artifact_manifest(path, require=True, allow_synthetic_geometry=True) + + +def test_manifest_rejects_component_source_fingerprint_drift(z_fixture_dir): + sidecar = z_fixture_dir / "experts-L00.nvfp4" + digest = FileExpertSource.create_synthetic( + sidecar, num_experts=1, records=[bytes([7]) * RAW_RECORD_BYTES], layer_id=0 + ) + path = _manifest(z_fixture_dir, sidecar, digest) + data = json.loads(path.read_text(encoding="utf-8")) + data["experts"]["files"][0]["source_fingerprint"] = "f" * 64 + data.pop("complete_artifact_fingerprint") + data["complete_artifact_fingerprint"] = hashlib.sha256( + json.dumps(data, sort_keys=True, separators=(",", ":")).encode("utf-8") + ).hexdigest() + path.write_text(json.dumps(data), encoding="utf-8") + with pytest.raises(ValueError, match="expert sidecar source fingerprint mismatch"): + load_qwen4_artifact_manifest(path, require=True, allow_synthetic_geometry=True) + + +def test_production_loader_rejects_reduced_geometry(z_fixture_dir): + sidecar = z_fixture_dir / "experts-L00.nvfp4" + digest = FileExpertSource.create_synthetic( + sidecar, num_experts=1, records=[bytes([7]) * RAW_RECORD_BYTES], layer_id=0 + ) + path = _manifest(z_fixture_dir, sidecar, digest) + data = json.loads(path.read_text(encoding="utf-8")) + data["source"]["repository"] = "RadixArk/Qwen3.8-Flash-Next-NVFP4" + data["source"]["revision"] = "7b719225242aacd3dbd3f9407468c2ee9a9d2594" + data["minimum_freetoken_commit"] = "846504bf9d81119cb72400e6c5a3cc860f2b1dd8" + data["tvm_ffi_patch_sha256"] = "889310b8152a147a6552a3e451b3251a7df70cdc8e6e4c1c87c7adf3854182ec" + data.pop("complete_artifact_fingerprint") + data["complete_artifact_fingerprint"] = hashlib.sha256( + json.dumps(data, sort_keys=True, separators=(",", ":")).encode("utf-8") + ).hexdigest() + path.write_text(json.dumps(data), encoding="utf-8") + with pytest.raises(ValueError, match="active payload does not match"): + load_qwen4_artifact_manifest(path, require=True) + + +def test_manifest_rejects_metadata_tamper_and_out_of_root_component(z_fixture_dir): + sidecar = z_fixture_dir / "experts-L00.nvfp4" + digest = FileExpertSource.create_synthetic( + sidecar, num_experts=1, records=[bytes([7]) * RAW_RECORD_BYTES], layer_id=0 + ) + path = _manifest(z_fixture_dir, sidecar, digest) + (z_fixture_dir / "config.json").write_bytes(b"tampered") + with pytest.raises(ValueError, match="length mismatch|SHA-256 mismatch"): + load_qwen4_artifact_manifest(path, require=True, allow_synthetic_geometry=True) + + path = _manifest(z_fixture_dir, sidecar, digest) + data = json.loads(path.read_text(encoding="utf-8")) + data["experts"]["files"][0]["path"] = str(sidecar.parent.parent / sidecar.name) + data.pop("complete_artifact_fingerprint") + data["complete_artifact_fingerprint"] = hashlib.sha256( + json.dumps(data, sort_keys=True, separators=(",", ":")).encode("utf-8") + ).hexdigest() + path.write_text(json.dumps(data), encoding="utf-8") + with pytest.raises(ValueError, match="outside artifact root"): + load_qwen4_artifact_manifest(path, require=True, allow_synthetic_geometry=True) + + +def test_manifest_reopens_and_wires_mixed_sources(z_fixture_dir): + sidecar = z_fixture_dir / "experts-L00.nvfp4" + digest = FileExpertSource.create_synthetic( + sidecar, num_experts=1, records=[bytes([7]) * RAW_RECORD_BYTES], layer_id=0 + ) + manifest_path = _manifest(z_fixture_dir, sidecar, digest) + manifest = load_qwen4_artifact_manifest(manifest_path, require=True, allow_synthetic_geometry=True) + assert manifest is not None + assert manifest.file_tier_layers == (0,) + assert manifest.resident_layers == (1,) + + resident = {name: [None, object()] for name in FileExpertSource.bank_schema} + cache = _FakeCache() + sources = configure_mixed_expert_sources(cache, manifest, resident) + assert sorted(sources) == [0] + assert all(cache.bank_sources[name][0] is None for name in cache.bank_schema) + assert all(cache.bank_sources[name][1] is not None for name in cache.bank_schema) + assert cache.file_sources[0].layer_id == 0 + cache.file_sources[0].close() + + +def test_resident_builder_streams_one_record_without_full_layer(z_fixture_dir): + source_fingerprint = hashlib.sha256(b"synthetic").hexdigest() + first_path = z_fixture_dir / "experts-L00.nvfp4" + first_digest = FileExpertSource.create_synthetic( + first_path, num_experts=1, records=[bytes([8]) * RAW_RECORD_BYTES], layer_id=0 + ) + resident_path = z_fixture_dir / "experts-L01.nvfp4" + digest = FileExpertSource.create_synthetic( + resident_path, num_experts=1, records=[bytes([9]) * RAW_RECORD_BYTES], layer_id=1 + ) + manifest_data = { + "format": FORMAT, + "version": 1, + "artifact_schema": FORMAT, + "text_only": True, + "source": {"repository": "synthetic", "revision": "test", "inventory_sha256": source_fingerprint}, + "minimum_freetoken_commit": "0" * 40, + "tvm_ffi_patch_sha256": "0" * 64, + "active": {"format": "nvfp4_w4a16_v1", "path": ".", "files": [{"path": "manifest.json", "bytes": 0, "sha256": "0" * 64}]}, + "ple": {"format": "q3_ple_32", "manifest": "ple-q3.json", "sha256": "0" * 64}, + "experts": { + "format": "ftexpert1_nvfp4_v1", + "files": [{"layer": 0, "path": first_path.name, "bytes": first_path.stat().st_size, "sha256": first_digest, "source_fingerprint": source_fingerprint}, + {"layer": 1, "path": resident_path.name, "bytes": resident_path.stat().st_size, "sha256": digest, "source_fingerprint": source_fingerprint}], + "file_tier_layers": [], + "resident_layers": [0, 1], + "required_volume": "Z:", + }, + "metadata": {"files": [{"path": "config.json", "bytes": 0, "sha256": hashlib.sha256(b"").hexdigest()}]}, + } + manifest_data["complete_artifact_fingerprint"] = hashlib.sha256( + json.dumps(manifest_data, sort_keys=True, separators=(",", ":")).encode("utf-8") + ).hexdigest() + path = z_fixture_dir / "manifest.json" + (z_fixture_dir / "config.json").write_bytes(b"") + (z_fixture_dir / "ple-q3.json").write_text( + json.dumps({"source_fingerprint": source_fingerprint}), encoding="utf-8" + ) + path.write_text(json.dumps(manifest_data), encoding="utf-8") + manifest = load_qwen4_artifact_manifest(path, require=True, allow_synthetic_geometry=True) + resident_sources, file_sources = build_mixed_expert_sources( + manifest, + num_experts=1, + resident_residency=["pageable", "pageable"], + allocator=lambda shape, dtype: torch.empty(shape, dtype=dtype), + ) + assert file_sources == {} + assert resident_sources["gate_up_packed"][0].shape == (1, 1280, 1280) + assert int(resident_sources["gate_up_packed"][1][0, 0, 0]) == 9 + + +def _tiny_geometry(): + return { + "gate_up_packed": ((4, 4), torch.uint8), + "gate_up_scale": ((4, 1), torch.float8_e4m3fn), + "gate_up_global": ((4,), torch.float16), + "down_packed": ((2, 16), torch.uint8), + "down_scale": ((2, 1), torch.float8_e4m3fn), + "down_global": ((2,), torch.float16), + } + + +def _tiny_planes(value: int): + return { + name: torch.full(shape, 1 if dtype == torch.float8_e4m3fn else value, dtype=dtype) + for name, (shape, dtype) in _tiny_geometry().items() + } + + +def test_end_to_end_synthetic_modular_artifact_reopens_normal_paths(z_fixture_dir, monkeypatch): + source_fingerprint = "4" * 64 + config = { + "architectures": ["Qwen4ExpForConditionalGeneration"], + "freetoken_text_only": TEXT_ONLY_MARKER, + "freetoken_active_quant": "nvfp4_w4a16_v1", + } + (z_fixture_dir / "config.json").write_text(json.dumps(config), encoding="utf-8") + (z_fixture_dir / "tokenizer_config.json").write_text("{}", encoding="utf-8") + + active = z_fixture_dir / "qwen4-active-v1.ftw" + writer = FTWWriter(str(active), shard_limit=4096 * 16) + fused = torch.arange(64, dtype=torch.float32).reshape(4, 16).to(torch.bfloat16) + protected = torch.arange(16, dtype=torch.bfloat16) + entries = iter_active_nvfp4_runtime_entries( + [ + ("model.layers.0.self_attn.qkv_proj.weight", fused), + ("model.layers.0.self_attn.index_qk_proj.weight", protected), + ] + ) + emitted_names = [] + for name, tensor in entries: + emitted_names.append(name) + writer.add_tensor(name, tensor) + writer.finalize({ + "artifact_format": "qwen4_modular_v1", + "source_inventory_sha256": source_fingerprint, + }) + assert emitted_names == [ + "model.layers.0.self_attn.qkv_proj.weight", + "model.layers.0.self_attn.qkv_proj.weight_scale", + "model.layers.0.self_attn.qkv_proj.weight_global", + "model.layers.0.self_attn.index_qk_proj.weight", + ] + + write_q3_ple_sidecar( + ([float((row + column) % 9 - 4) for column in range(ROW_VALUES)] for row in range(4)), + z_fixture_dir / "ple-q3-000.bin", + z_fixture_dir / "ple-q3.json", + source_fingerprint=source_fingerprint, + weight_scale=1.0, + segment_rows=2, + ) + experts = {} + for layer in range(2): + path = z_fixture_dir / f"experts-L{layer:02d}.nvfp4" + write_expert_sidecar( + path, + [_tiny_planes(layer + expert + 1) for expert in range(2)], + layer_id=layer, + source_fingerprint=source_fingerprint, + num_experts=2, + geometry=_tiny_geometry(), + ) + experts[layer] = path.name + + finalized = finalize_qwen4_modular_manifest( + z_fixture_dir, + source_repository="synthetic/qwen4", + source_revision="3" * 40, + source_inventory_sha256=source_fingerprint, + minimum_freetoken_commit="846504bf9d81119cb72400e6c5a3cc860f2b1dd8", + tvm_ffi_patch_sha256="889310b8152a147a6552a3e451b3251a7df70cdc8e6e4c1c87c7adf3854182ec", + expert_paths=experts, + file_tier_layers=[0], + metadata_paths=["config.json", "tokenizer_config.json"], + expert_num_experts=2, + allow_synthetic_geometry=True, + ) + assert finalized["experts"]["resident_layers"] == [1] + manifest = load_qwen4_artifact_manifest(z_fixture_dir, require=True, allow_synthetic_geometry=True) + assert manifest is not None + manifest.verify_active() + import freetoken.checkpoint.qwen4_artifact as artifact_module + + original_loader = artifact_module.load_qwen4_artifact_manifest + monkeypatch.setattr( + artifact_module, + "load_qwen4_artifact_manifest", + lambda _path, **_kwargs: manifest, + ) + loaded = dict(load_weight(str(z_fixture_dir), torch.device("cpu"), include_moe_experts=False)) + assert set(loaded) == set(emitted_names) + resident, file_sources = build_mixed_expert_sources( + manifest, + num_experts=2, + resident_residency=["pageable", "pageable"], + allocator=lambda shape, dtype: torch.empty(shape, dtype=dtype), + ) + assert all(resident[name][0] is None for name in FileExpertSource.bank_schema) + assert all(resident[name][1] is not None for name in FileExpertSource.bank_schema) + assert set(file_sources) == {0} + file_sources[0].close() + + # Normal Qwen4 host-load dispatch selects Q3 from the manifest; no manual + # load_q3_ple_weights injection is involved. + calls = [] + fake = SimpleNamespace(load_q3_ple_weights=lambda path: calls.append(path)) + from freetoken.models.qwen4_exp.model import Qwen4ExpModel + + Qwen4ExpModel.load_host_weights(fake, str(z_fixture_dir), dummy=False) + assert calls == [str(z_fixture_dir / "ple-q3.json")] + monkeypatch.setattr(artifact_module, "load_qwen4_artifact_manifest", original_loader) + + +def test_production_orchestrator_sequences_all_modular_components(z_fixture_dir, monkeypatch): + source = z_fixture_dir / "source" + target = z_fixture_dir / "target" + source.mkdir() + inventory = "5" * 64 + + import freetoken.checkpoint.convert as convert_module + import freetoken.checkpoint.q3_ple as q3_module + import freetoken.moe.expert_source as expert_module + + def fake_convert(_source, out_dir, **kwargs): + assert kwargs["artifact_format"] == "qwen4_modular_v1" + assert kwargs["source_inventory_sha256"] == inventory + out = Path(out_dir) + out.mkdir(parents=True, exist_ok=True) + (out / "config.json").write_text( + json.dumps({ + "freetoken_text_only": TEXT_ONLY_MARKER, + "freetoken_active_quant": "nvfp4_w4a16_v1", + }), + encoding="utf-8", + ) + active = out / "qwen4-active-v1.ftw" + writer = FTWWriter(str(active), shard_limit=4096 * 4) + writer.add_tensor("protected.weight", torch.ones(1, dtype=torch.bfloat16)) + return writer.finalize({ + "source_inventory_sha256": inventory, + "copied_metadata": ["config.json"], + }) + + def fake_q3(_source, data_path, manifest_path, **kwargs): + assert kwargs["source_fingerprint"] == inventory + return write_q3_ple_sidecar( + ([0.0] * ROW_VALUES for _ in range(2)), + data_path, + manifest_path, + source_fingerprint=inventory, + weight_scale=1.0, + segment_rows=1, + ) + + def fake_expert(_source, path, **kwargs): + return write_expert_sidecar( + path, + [_tiny_planes(kwargs["layer_id"] + 1) for _ in range(kwargs["num_experts"])], + layer_id=kwargs["layer_id"], + source_fingerprint=kwargs["source_fingerprint"], + num_experts=kwargs["num_experts"], + geometry=kwargs["geometry"], + ) + + monkeypatch.setattr(convert_module, "convert_checkpoint", fake_convert) + monkeypatch.setattr(q3_module, "write_q3_ple_from_safetensors", fake_q3) + monkeypatch.setattr(expert_module, "write_expert_sidecar_from_safetensors", fake_expert) + manifest = build_qwen4_modular_artifact( + source, + target, + source_repository="synthetic/qwen4", + source_revision="6" * 40, + source_inventory_sha256=inventory, + minimum_freetoken_commit="7" * 40, + ple_split_parts=1, + expert_layers=(0, 1), + file_tier_layers=(0,), + expert_num_experts=2, + expert_geometry=_tiny_geometry(), + allow_synthetic_geometry=True, + ) + assert manifest["active"]["format"] == "nvfp4_w4a16_v1" + assert manifest["ple"]["source_fingerprint"] == inventory + assert [item["layer"] for item in manifest["experts"]["files"]] == [0, 1] + assert (target / "manifest.json").is_file() diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py index a164f0b4..5e28d40d 100644 --- a/tests/engine/test_cache_budget.py +++ b/tests/engine/test_cache_budget.py @@ -97,6 +97,14 @@ def test_expert_bytes_per_slot_sums_row_bytes_over_banks(): assert expert_bytes_per_slot(sources) == 512 + 256 +def test_expert_bytes_per_slot_accepts_file_tier_placeholders(): + sources = { + "packed": [None, torch.empty((2, 16), dtype=torch.uint8)], + "scale": [None, torch.empty((2, 4), dtype=torch.float16)], + } + assert expert_bytes_per_slot(sources) == 16 + 8 + + def test_resolve_auto_applies_ratio_once_and_marlin_cap(): # baseline 1000, weights 100, ratio 0.9 -> budget = 900 - 100 - 0(fixed) = 800 size, pages, overlap = resolve_moe_cache_auto( diff --git a/tests/kernels/test_qsa_differential.py b/tests/kernels/test_qsa_differential.py new file mode 100644 index 00000000..20d06caa --- /dev/null +++ b/tests/kernels/test_qsa_differential.py @@ -0,0 +1,314 @@ +from __future__ import annotations + +import json +import math +import time + +import pytest +import torch + +from freetoken.attention.qsa import select_qsa_logical_rows +from freetoken.kernel.triton.qsa import qsa_sparse_gqa + + +SEED = 38038 +COMPRESS_RATIO = 4 +TOKEN_BUDGET = 2048 +OUTPUT_WIDTH = TOKEN_BUDGET + COMPRESS_RATIO - 1 +QUERY_POSITIONS = ( + 0, + 1, + 2, + 3, + 4, + 2047, + 2048, + 2049, + 2050, + 2051, + 2052, + 4095, + 8191, + 65535, +) + + +def _official_contiguous_oracle( + index_q: torch.Tensor, + compressed_keys: torch.Tensor, + query_positions: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Independent eager transcription of the pinned official indexer semantics. + + This oracle intentionally does not call any FreeToken selection or compaction + helper. Its scope is the serving topology supported by FreeToken: one + contiguous causal history per request. + """ + rows = torch.full( + (query_positions.numel(), OUTPUT_WIDTH), -1, dtype=torch.int32 + ) + counts = torch.empty(query_positions.numel(), dtype=torch.int32) + offsets = torch.arange(COMPRESS_RATIO, dtype=torch.long) + for query_row, position_tensor in enumerate(query_positions.cpu()): + position = int(position_tensor) + visible = position + 1 + complete_groups = visible // COMPRESS_RATIO + width = min(TOKEN_BUDGET // COMPRESS_RATIO, complete_groups) + if width: + query = index_q[query_row].cpu().float() + keys = compressed_keys[:complete_groups, 0].cpu().float() + scores = torch.relu(query @ keys.transpose(0, 1)).sum(dim=0) + scores /= math.sqrt(query.shape[-1]) + groups = torch.topk(scores, width, sorted=True).indices + selected = (groups[:, None] * COMPRESS_RATIO + offsets).flatten() + else: + selected = torch.empty(0, dtype=torch.long) + tail = torch.arange( + complete_groups * COMPRESS_RATIO, visible, dtype=torch.long + ) + selected = torch.cat((selected, tail)) + rows[query_row, : selected.numel()] = selected.to(torch.int32) + counts[query_row] = selected.numel() + return rows, counts + + +def _score_separated_fixture( + positions: tuple[int, ...] = QUERY_POSITIONS, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + query_positions = torch.tensor(positions, dtype=torch.int64) + complete_groups = (max(positions) + 1) // COMPRESS_RATIO + dim = 8 + index_q = torch.zeros(len(positions), 4, dim, dtype=torch.float32) + index_q[:, :, 0] = torch.tensor([1.0, 1.5, 2.0, 2.5]) + compressed_keys = torch.zeros(complete_groups, 1, dim, dtype=torch.float32) + # Strictly increasing positive scores make exact top-k order authoritative. + compressed_keys[:, 0, 0] = torch.arange( + 1, complete_groups + 1, dtype=torch.float32 + ) + return index_q, compressed_keys, query_positions + + +def _live_rows(rows: torch.Tensor, counts: torch.Tensor, index: int) -> torch.Tensor: + return rows[index, : int(counts[index])].long() + + +def _assert_group_and_tail_invariants( + selected: torch.Tensor, position: int, expected_count: int +) -> None: + assert selected.numel() == expected_count + assert torch.all((selected >= 0) & (selected <= position)) + assert torch.unique(selected).numel() == selected.numel() + visible = position + 1 + complete_groups = visible // COMPRESS_RATIO + tail = torch.arange(complete_groups * COMPRESS_RATIO, visible) + if tail.numel(): + assert torch.equal(selected[-tail.numel() :].cpu(), tail) + selected = selected[: -tail.numel()] + assert selected.numel() % COMPRESS_RATIO == 0 + if selected.numel(): + groups = selected.view(-1, COMPRESS_RATIO) + assert torch.equal( + groups, + groups[:, :1] + torch.arange(COMPRESS_RATIO, device=groups.device), + ) + assert torch.all(groups[:, 0] % COMPRESS_RATIO == 0) + + +def _independent_sparse_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + rows: torch.Tensor, + counts: torch.Tensor, + scale: float, +) -> torch.Tensor: + output = torch.zeros_like(q, dtype=torch.float32) + gqa = q.shape[1] // k.shape[1] + for query_row in range(q.shape[0]): + selected = rows[query_row, : int(counts[query_row])].long().cpu() + for kv_head in range(k.shape[1]): + heads = slice(kv_head * gqa, (kv_head + 1) * gqa) + scores = torch.einsum( + "hd,td->ht", + q[query_row, heads].cpu().float(), + k[selected, kv_head].cpu().float(), + ) * scale + probabilities = torch.softmax(scores, dim=-1) + output[query_row, heads] = torch.einsum( + "ht,td->hd", probabilities, v[selected, kv_head].cpu().float() + ) + return output + + +def test_qsa_unique_score_selection_matches_pinned_official_oracle_cpu(): + torch.manual_seed(SEED) + index_q, compressed_keys, positions = _score_separated_fixture() + expected_rows, expected_counts = _official_contiguous_oracle( + index_q, compressed_keys, positions + ) + actual_rows, actual_counts = select_qsa_logical_rows( + index_q, + compressed_keys, + positions, + compress_ratio=COMPRESS_RATIO, + token_budget=TOKEN_BUDGET, + ) + assert torch.equal(actual_counts.cpu(), expected_counts) + assert torch.equal(actual_rows.cpu(), expected_rows) + + for row, position in enumerate(QUERY_POSITIONS): + selected = _live_rows(actual_rows, actual_counts, row) + assert selected.numel() <= OUTPUT_WIDTH + assert torch.all(selected <= position) + if position <= 2050: + assert set(selected.tolist()) == set(range(position + 1)) + + boundary = QUERY_POSITIONS.index(2051) + boundary_rows = _live_rows(actual_rows, actual_counts, boundary) + assert actual_counts[boundary].item() == TOKEN_BUDGET + assert set(boundary_rows.tolist()) == set(range(4, 2052)) + assert not any(token in boundary_rows.tolist() for token in range(4)) + + after = QUERY_POSITIONS.index(2052) + assert actual_counts[after].item() == TOKEN_BUDGET + 1 + assert _live_rows(actual_rows, actual_counts, after)[-1].item() == 2052 + + +@pytest.mark.parametrize("key_value", [0.0, -1.0]) +def test_qsa_tied_scores_obey_order_insensitive_official_invariants_cpu(key_value): + positions = torch.tensor([2051, 2052, 4095], dtype=torch.int64) + index_q = torch.ones(positions.numel(), 4, 8) + compressed_keys = torch.full((1024, 1, 8), key_value) + rows, counts = select_qsa_logical_rows( + index_q, + compressed_keys, + positions, + compress_ratio=COMPRESS_RATIO, + token_budget=TOKEN_BUDGET, + ) + for row, position in enumerate(positions.tolist()): + complete = (position + 1) // COMPRESS_RATIO + tail = (position + 1) % COMPRESS_RATIO + expected_count = min(complete, TOKEN_BUDGET // COMPRESS_RATIO) * 4 + tail + _assert_group_and_tail_invariants( + _live_rows(rows, counts, row), position, expected_count + ) + + +def test_qsa_float32_sparse_attention_matches_independent_eager_oracle_cpu(): + torch.manual_seed(SEED) + positions_tuple = (2051, 2052, 4095) + index_q, compressed_keys, positions = _score_separated_fixture(positions_tuple) + selected, counts = select_qsa_logical_rows( + index_q, + compressed_keys, + positions, + compress_ratio=COMPRESS_RATIO, + token_budget=TOKEN_BUDGET, + ) + q = torch.randn(len(positions_tuple), 4, 16, dtype=torch.float32) + k = torch.randn(max(positions_tuple) + 1, 2, 16, dtype=torch.float32) + v = torch.randn_like(k) + scale = 16**-0.5 + actual = qsa_sparse_gqa(q, k, v, selected, counts, scale) + expected = _independent_sparse_attention(q, k, v, selected, counts, scale) + error = (actual.float() - expected).abs() + print( + "STAGE3_QSA_METRIC", + json.dumps( + { + "device": "cpu", + "dtype": "float32", + "max_abs_error": error.max().item(), + "mean_abs_error": error.mean().item(), + }, + sort_keys=True, + ), + ) + torch.testing.assert_close(actual.float(), expected, rtol=1e-5, atol=1e-6) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_qsa_cuda_selection_and_compaction_match_independent_oracle(): + torch.manual_seed(SEED) + index_q, compressed_keys, positions = _score_separated_fixture() + expected_rows, expected_counts = _official_contiguous_oracle( + index_q, compressed_keys, positions + ) + torch.cuda.reset_peak_memory_stats() + started = time.perf_counter() + actual_rows, actual_counts = select_qsa_logical_rows( + index_q.cuda(), + compressed_keys.cuda(), + positions.cuda(), + compress_ratio=COMPRESS_RATIO, + token_budget=TOKEN_BUDGET, + ) + torch.cuda.synchronize() + elapsed = time.perf_counter() - started + assert torch.equal(actual_counts.cpu(), expected_counts) + assert torch.equal(actual_rows.cpu(), expected_rows) + peak = torch.cuda.max_memory_allocated() + print( + "STAGE3_QSA_METRIC", + json.dumps( + { + "case": "cuda_selection_compaction", + "device": torch.cuda.get_device_name(), + "driver_runtime": torch.version.cuda, + "elapsed_seconds": elapsed, + "peak_memory_bytes": peak, + }, + sort_keys=True, + ), + ) + assert peak < 1 << 30 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_qsa_cuda_bf16_sparse_gqa_matches_independent_eager_oracle(): + torch.manual_seed(SEED) + positions_tuple = (2051, 65535) + index_q, compressed_keys, positions = _score_separated_fixture(positions_tuple) + selected, counts = _official_contiguous_oracle(index_q, compressed_keys, positions) + dim = 64 + q_cpu = torch.randn(len(positions_tuple), 8, dim, dtype=torch.bfloat16) + k_cpu = torch.randn(max(positions_tuple) + 1, 2, dim, dtype=torch.bfloat16) + v_cpu = torch.randn_like(k_cpu) + scale = dim**-0.5 + expected = _independent_sparse_attention( + q_cpu, k_cpu, v_cpu, selected, counts, scale + ) + + torch.cuda.reset_peak_memory_stats() + started = time.perf_counter() + actual = qsa_sparse_gqa( + q_cpu.cuda(), + k_cpu.cuda(), + v_cpu.cuda(), + selected.cuda(), + counts.cuda(), + scale, + ) + torch.cuda.synchronize() + elapsed = time.perf_counter() - started + actual_cpu = actual.cpu().float() + error = (actual_cpu - expected).abs() + peak = torch.cuda.max_memory_allocated() + print( + "STAGE3_QSA_METRIC", + json.dumps( + { + "case": "cuda_bf16_sparse_gqa", + "device": torch.cuda.get_device_name(), + "driver_runtime": torch.version.cuda, + "elapsed_seconds": elapsed, + "peak_memory_bytes": peak, + "max_abs_error": error.max().item(), + "mean_abs_error": error.mean().item(), + }, + sort_keys=True, + ), + ) + assert peak < 1 << 30 + torch.testing.assert_close(actual_cpu, expected, rtol=2e-2, atol=2e-2) diff --git a/tests/kernels/test_qwen4_nvfp4_active.py b/tests/kernels/test_qwen4_nvfp4_active.py new file mode 100644 index 00000000..350481e8 --- /dev/null +++ b/tests/kernels/test_qwen4_nvfp4_active.py @@ -0,0 +1,59 @@ +"""Bounded synthetic NVFP4 W4A16 differential tests (no model payloads).""" + +import pytest +import torch + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable") +@pytest.mark.parametrize( + "out_features,in_features", + [ + (320, 10240), # mHC down + (10240, 320), # mHC up + (1280, 2560), # shared gate|up + (2560, 640), # shared down + (13312, 2560), # QSA q|k|v + (2560, 6144), # QSA/GDN output + (16384, 2560), # GDN qkv|z + ], +) +@pytest.mark.parametrize("rows", [1, 2, 64, 65]) +def test_native_nvfp4_dense_matches_dequantized_reference(rows, out_features, in_features): + from freetoken.checkpoint.nvfp4 import encode_bf16_nvfp4 + from freetoken.kernel.triton.nvfp4_linear import nvfp4_dense_linear + from freetoken.kernel.triton.nvfp4_dequant import dequant_nvfp4 + + torch.manual_seed(38038 + rows + out_features) + source = torch.randn(out_features, in_features, dtype=torch.bfloat16) + packed, scales, globals_ = encode_bf16_nvfp4(source) + x = torch.randn(rows, in_features, dtype=torch.bfloat16, device="cuda") + packed_cuda, scales_cuda, globals_cuda = packed.cuda(), scales.cuda(), globals_.cuda() + out = nvfp4_dense_linear(x, packed_cuda, scales_cuda, globals_cuda) + weight = dequant_nvfp4( + packed_cuda.unsqueeze(0), scales_cuda.unsqueeze(0), globals_cuda.unsqueeze(0), + torch.zeros(1, dtype=torch.int32, device="cuda"), dtype=torch.bfloat16, + )[0] + reference = x @ weight.t() + # Reuse the pre-existing native NVFP4 backend tolerance verbatim: BF16 + # grouped/dense GEMMs accumulate large reductions, so the absolute bound is + # relative to this fixture's output magnitude. + atol = 0.03 * float(reference.abs().max()) + torch.testing.assert_close(out.float(), reference.float(), rtol=3e-2, atol=atol) + assert torch.cuda.max_memory_allocated() < 6 * (1 << 30) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable") +def test_native_nvfp4_state_dict_repackages_without_bf16_weight_copy(): + from freetoken.checkpoint.nvfp4 import encode_bf16_nvfp4 + from freetoken.kernel.triton.nvfp4_linear import Nvfp4DenseLinear + + torch.manual_seed(38039) + source = torch.randn(17, 32, dtype=torch.bfloat16) + packed, scales, globals_ = encode_bf16_nvfp4(source) + op = Nvfp4DenseLinear(32, 17) + state = {"weight": packed, "weight_scale": scales, "weight_global": globals_} + op.load_state_dict(state) + assert not state + assert op.weight.dtype == torch.int32 and op.weight.shape == (4, 17) + assert op.weight_scale.shape == (2, 17) + assert op.weight_global.dtype == torch.float16 diff --git a/tests/kernels/test_tensor_matcher.py b/tests/kernels/test_tensor_matcher.py new file mode 100644 index 00000000..a245aa0d --- /dev/null +++ b/tests/kernels/test_tensor_matcher.py @@ -0,0 +1,150 @@ +from __future__ import annotations + +from functools import lru_cache +import os +from pathlib import Path +import sys + +import pytest +import torch + + +_SOURCE = r""" +#include +#include + +int symbolic_cuda_same(tvm::ffi::TensorView first, tvm::ffi::TensorView second) { + auto device = host::SymbolicDevice{}; + host::TensorMatcher({-1}) + .with_device(device) + .verify(first) + .verify(second); + return device.unwrap().device_id; +} + +int symbolic_cuda_same_typed(tvm::ffi::TensorView first, + tvm::ffi::TensorView second) { + auto length = host::SymbolicSize{"length"}; + auto dtype = host::SymbolicDType{}; + auto device = host::SymbolicDevice{}; + host::TensorMatcher({length}) + .with_dtype(dtype) + .with_device(device) + .verify(first) + .verify(second); + return device.unwrap().device_id; +} + +template struct SymbolicCudaTemplate { + static int run(tvm::ffi::TensorView first, tvm::ffi::TensorView second) { + auto length = host::SymbolicSize{"length"}; + auto dtype = host::SymbolicDType{}; + auto device = host::SymbolicDevice{}; + host::TensorMatcher({length}) + .with_dtype(dtype) + .template with_device(device) + .verify(first) + .verify(second); + return device.unwrap().device_id + Tag - Tag; + } +}; + +int symbolic_cuda_same_templated(tvm::ffi::TensorView first, + tvm::ffi::TensorView second) { + return SymbolicCudaTemplate<1>::run(first, second); +} + +int symbolic_unrestricted(tvm::ffi::TensorView value) { + auto device = host::SymbolicDevice{}; + host::TensorMatcher({-1}).with_device(device).verify(value); + return static_cast(device.unwrap().device_type); +} + +void fixed_cpu(tvm::ffi::TensorView value) { + host::TensorMatcher({-1}).with_device().verify(value); +} + +void explicit_cpu(tvm::ffi::TensorView value) { + host::TensorMatcher({-1}).with_device({{kDLCPU, 0}}).verify(value); +} + +void reject_different_cuda_device_ids() { + auto device = host::SymbolicDevice{}; + device.set_options(); + device.verify({kDLCUDA, 0}); + device.verify({kDLCUDA, 1}); +} +""" + +_FUNCTIONS = [ + "symbolic_cuda_same", + "symbolic_cuda_same_typed", + "symbolic_cuda_same_templated", + "symbolic_unrestricted", + "fixed_cpu", + "explicit_cpu", + "reject_different_cuda_device_ids", +] + + +@lru_cache(maxsize=1) +def _cpu_module(): + from freetoken.kernel.utils import DEFAULT_CFLAGS, DEFAULT_INCLUDE + from tvm_ffi.cpp import load_inline + + return load_inline( + "freetoken_tensor_matcher_cpu_test_v7", + cpp_sources=_SOURCE, + functions=_FUNCTIONS, + extra_cflags=DEFAULT_CFLAGS, + extra_include_paths=DEFAULT_INCLUDE, + ) + + +@lru_cache(maxsize=1) +def _cuda_module(): + from freetoken.kernel.utils import DEFAULT_INCLUDE, _cuda_cflags + from tvm_ffi.cpp import load_inline + + extra_ldflags = [] + if sys.platform == "win32": + cuda_home = Path(os.environ["CUDA_HOME"]) + extra_ldflags = [f"/LIBPATH:{cuda_home / 'lib' / 'x64'}", "cudart.lib"] + + return load_inline( + "freetoken_tensor_matcher_cuda_test_v7", + cuda_sources=_SOURCE, + functions=_FUNCTIONS, + extra_cuda_cflags=_cuda_cflags([]), + extra_ldflags=extra_ldflags, + extra_include_paths=DEFAULT_INCLUDE, + backend="cuda", + ) + + +def test_symbolic_device_restrictions_and_existing_cpu_paths(): + module = _cpu_module() + value = torch.ones(4) + + assert module.symbolic_unrestricted(value) == 1 # DLPack kDLCPU + module.fixed_cpu(value) + module.explicit_cpu(value) + with pytest.raises(Exception, match="Device"): + module.symbolic_cuda_same(value, value) + with pytest.raises(Exception, match="Device"): + module.reject_different_cuda_device_ids() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_symbolic_cuda_binding_compiles_and_checks_same_device(): + module = _cuda_module() + first = torch.ones(4, device="cuda") + second = torch.zeros(4, device="cuda") + + assert module.symbolic_cuda_same(first, second) == torch.cuda.current_device() + first_int = torch.ones(4, dtype=torch.int32, device="cuda") + second_int = torch.zeros(4, dtype=torch.int32, device="cuda") + assert module.symbolic_cuda_same_typed(first_int, second_int) == torch.cuda.current_device() + assert module.symbolic_cuda_same_templated(first_int, second_int) == torch.cuda.current_device() + with pytest.raises(Exception, match="Device"): + module.symbolic_cuda_same(torch.ones(4), torch.zeros(4)) diff --git a/tests/kvcache/test_qsa_incremental_differential.py b/tests/kvcache/test_qsa_incremental_differential.py new file mode 100644 index 00000000..85f5aca1 --- /dev/null +++ b/tests/kvcache/test_qsa_incremental_differential.py @@ -0,0 +1,369 @@ +from __future__ import annotations + +import json +import math +import time +from types import SimpleNamespace + +import pytest +import torch + +import freetoken.attention.qsa as qsa_module +from freetoken.attention.qsa import QSAAttnBackend, select_qsa_logical_rows +from freetoken.distributed.info import DistributedInfo +from freetoken.kvcache.qsa_pool import QSAKVCache + + +SEED = 38038 +RATIO = 4 +BUDGET = 2048 +LAYER_ID = 1 +INDEX_DIM = 8 + + +class _RecordingSyntheticIndexer: + """Observable stand-in for official K RMSNorm + block-start RoPE. + + The backend owns pooling and position routing, while the model indexer owns + normalization/RoPE. This deterministic transform makes both inputs visible + without importing model weights or calling a production selection helper. + """ + + def __init__(self): + self.calls: list[tuple[torch.Tensor, torch.Tensor]] = [] + + def normalize_compressed_keys( + self, keys: torch.Tensor, positions: torch.Tensor + ) -> torch.Tensor: + self.calls.append((keys.detach().cpu().clone(), positions.detach().cpu().clone())) + values = keys.float() + normalized = values * torch.rsqrt(values.square().mean(-1, keepdim=True) + 1e-6) + # A small block-start marker makes an incorrect current/end position observable. + marker = positions[0].float().view(-1, 1, 1) / 4096.0 + return (normalized + marker).to(keys.dtype) + + +def _oracle_normalize_and_position( + pooled: torch.Tensor, block_start_rope: torch.Tensor, dtype: torch.dtype +) -> torch.Tensor: + values = pooled.float() + normalized = values * torch.rsqrt(values.square().mean(-1, keepdim=True) + 1e-6) + marker = block_start_rope[0].float().view(-1, 1, 1) / 4096.0 + return (normalized + marker).to(dtype) + + +def _pool(monkeypatch, device: torch.device, num_pages: int = 70) -> QSAKVCache: + monkeypatch.setattr( + "freetoken.kvcache.mha_pool.get_tp_info", + lambda: DistributedInfo(rank=0, size=1), + ) + return QSAKVCache( + num_kv_heads=2, + num_layers=2, + head_dim=16, + num_pages=num_pages, + page_size=64, + dtype=torch.bfloat16, + device=device, + index_num_kv_heads=1, + index_head_dim=INDEX_DIM, + compress_ratio=RATIO, + layer_ids=(LAYER_ID,), + ) + + +def _backend_and_context(monkeypatch, pool: QSAKVCache, max_tokens: int = 2112): + stride = max_tokens + page_table = torch.stack( + ( + torch.arange(max_tokens, device=pool.device), + torch.arange(stride, stride + max_tokens, device=pool.device), + ) + ).long() + assert int(page_table.max()) < pool.k_cache(LAYER_ID).numel() // (2 * 16) + context = SimpleNamespace(kv_cache=pool, page_table=page_table) + monkeypatch.setattr(qsa_module, "get_global_ctx", lambda: context) + config = SimpleNamespace( + qwen4_args=SimpleNamespace( + indexer_compress_ratio=RATIO, + indexer_budget=BUDGET, + ) + ) + return QSAAttnBackend(config), context + + +def _raw_keys(length: int, request_id: int, device: torch.device) -> torch.Tensor: + base = torch.arange(length * INDEX_DIM, dtype=torch.float32).view(length, 1, INDEX_DIM) + values = base / 97.0 + 1.0 + request_id * 100.0 + return values.to(device=device, dtype=torch.bfloat16) + + +def _rope_positions(length: int, request_id: int, device: torch.device) -> torch.Tensor: + logical = torch.arange(length, device=device, dtype=torch.int64) + return torch.stack( + (logical, logical + 1000 * request_id, logical + 2000 * request_id) + ) + + +def _request(start: int, end: int, table_idx: int): + return SimpleNamespace( + cached_len=start, + device_len=end, + extend_len=end - start, + table_idx=table_idx, + ) + + +def _batch(reqs, rope_chunks): + lengths = [req.extend_len for req in reqs] + return SimpleNamespace( + reqs=reqs, + padded_reqs=reqs, + rope_positions=torch.cat(rope_chunks, dim=1), + input_ids=torch.empty(sum(lengths), dtype=torch.long), + ) + + +def _oracle_completed( + full_keys: torch.Tensor, full_rope: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + groups = full_keys.shape[0] // RATIO + if not groups: + return ( + full_keys.new_empty((0, 1, INDEX_DIM)), + full_rope.new_empty((3, 0)), + ) + members = full_keys[: groups * RATIO].view(groups, RATIO, 1, INDEX_DIM) + pooled = members.float().mean(dim=1).to(full_keys.dtype) + starts = torch.arange(0, groups * RATIO, RATIO, device=full_keys.device) + rope = full_rope.index_select(1, starts) + return _oracle_normalize_and_position(pooled, rope, full_keys.dtype), rope + + +def _oracle_selection( + query: torch.Tensor, keys: torch.Tensor, position: int +) -> tuple[torch.Tensor, int]: + complete = (position + 1) // RATIO + width = min(BUDGET // RATIO, complete) + if width: + score = torch.relu( + query.float() @ keys[:complete, 0].float().transpose(0, 1) + ).sum(dim=0) / math.sqrt(query.shape[-1]) + groups = torch.topk(score, width, sorted=True).indices + rows = ( + groups[:, None] * RATIO + torch.arange(RATIO, device=groups.device) + ).flatten() + else: + rows = torch.empty(0, dtype=torch.long, device=query.device) + tail = torch.arange(complete * RATIO, position + 1, device=query.device) + rows = torch.cat((rows, tail)) + return rows, rows.numel() + + +def _assert_incremental_state( + backend: QSAAttnBackend, + context, + full_keys: torch.Tensor, + full_rope: torch.Tensor, + end: int, + table_idx: int, +) -> None: + expected, _ = _oracle_completed(full_keys[:end], full_rope[:, :end]) + starts = torch.arange(0, expected.shape[0] * RATIO, RATIO, device=full_keys.device) + physical = context.page_table[table_idx].index_select(0, starts) + compressed_rows = torch.div(physical, RATIO, rounding_mode="floor") + actual = backend.kvcache.compressed_k_cache(LAYER_ID).index_select( + 0, compressed_rows.long() + ) + assert torch.equal(actual, expected) + + latest_start = max(0, end - RATIO) + latest = torch.arange(latest_start, end, device=full_keys.device) + assert torch.equal( + backend.kvcache.pending_group(LAYER_ID, table_idx, latest), + full_keys.index_select(0, latest), + ) + assert torch.equal( + backend.kvcache.pending_rope_group(LAYER_ID, table_idx, latest), + full_rope.index_select(1, latest).transpose(0, 1), + ) + + query = torch.ones(1, 4, INDEX_DIM, device=full_keys.device) + selected, counts = select_qsa_logical_rows( + query, + actual, + torch.tensor([end - 1], device=full_keys.device), + compress_ratio=RATIO, + token_budget=BUDGET, + ) + expected_rows, expected_count = _oracle_selection(query[0], actual, end - 1) + assert int(counts[0]) == expected_count + assert torch.equal(selected[0, :expected_count].long(), expected_rows.long()) + + +def _run_pattern(monkeypatch, pattern: list[int], device: torch.device): + total = sum(pattern) + pool = _pool(monkeypatch, device) + backend, context = _backend_and_context(monkeypatch, pool) + indexer = _RecordingSyntheticIndexer() + full_keys = _raw_keys(total, 0, device) + full_rope = _rope_positions(total, 0, device) + start = 0 + for length in pattern: + end = start + length + req = _request(start, end, 0) + batch = _batch([req], [full_rope[:, start:end]]) + backend.prepare_metadata(batch) + backend._compress_current_keys( + indexer, full_keys[start:end], LAYER_ID, batch + ) + _assert_incremental_state( + backend, context, full_keys, full_rope, end, table_idx=0 + ) + start = end + return backend, context, indexer, full_keys, full_rope + + +@pytest.mark.parametrize( + "pattern", + ([1, 3], [3, 1], [5, 2, 1], [1, 1, 1, 1, 1, 1, 1, 1]), +) +def test_qsa_incremental_compression_matches_full_history_cpu(monkeypatch, pattern): + torch.manual_seed(SEED) + backend, context, indexer, full_keys, full_rope = _run_pattern( + monkeypatch, list(pattern), torch.device("cpu") + ) + expected, expected_rope = _oracle_completed(full_keys, full_rope) + assert indexer.calls + # Across all calls, every emitted completed group uses its block-start RoPE. + recorded_rope = torch.cat([positions for _, positions in indexer.calls], dim=1) + assert torch.equal(recorded_rope, expected_rope.cpu()) + starts = torch.arange(0, expected.shape[0] * RATIO, RATIO) + rows = torch.div(context.page_table[0].cpu().index_select(0, starts), RATIO, rounding_mode="floor") + assert torch.equal( + backend.kvcache.compressed_k_cache(LAYER_ID).cpu().index_select(0, rows), + expected.cpu(), + ) + + +def test_qsa_incremental_crosses_first_sparse_boundary_cpu(monkeypatch): + torch.manual_seed(SEED) + pattern = [2047, 1, 1, 1, 1, 1] + backend, context, _, full_keys, full_rope = _run_pattern( + monkeypatch, pattern, torch.device("cpu") + ) + assert full_keys.shape[0] == 2052 + expected, _ = _oracle_completed(full_keys, full_rope) + assert expected.shape[0] == 513 + query = torch.ones(1, 4, INDEX_DIM) + selected, counts = select_qsa_logical_rows( + query, + expected, + torch.tensor([2051]), + compress_ratio=RATIO, + token_budget=BUDGET, + ) + oracle_rows, oracle_count = _oracle_selection(query[0], expected, 2051) + assert oracle_count == BUDGET + assert int(counts[0]) == BUDGET + assert torch.equal(selected[0, :BUDGET].long(), oracle_rows.long()) + # The least-scoring complete group is the only omitted group. + omitted = set(range(2052)) - set(oracle_rows.tolist()) + assert omitted == set(range(4)) + assert backend.kvcache is context.kv_cache + + +def test_qsa_multi_request_cache_and_pending_state_are_isolated_cpu(monkeypatch): + torch.manual_seed(SEED) + pool = _pool(monkeypatch, torch.device("cpu")) + backend, context = _backend_and_context(monkeypatch, pool) + indexer = _RecordingSyntheticIndexer() + lengths = (5, 7) + keys = [_raw_keys(length, req_id, torch.device("cpu")) for req_id, length in enumerate(lengths)] + rope = [_rope_positions(length, req_id, torch.device("cpu")) for req_id, length in enumerate(lengths)] + reqs = [_request(0, length, req_id) for req_id, length in enumerate(lengths)] + batch = _batch(reqs, rope) + backend.prepare_metadata(batch) + backend._compress_current_keys(indexer, torch.cat(keys), LAYER_ID, batch) + + for req_id, length in enumerate(lengths): + _assert_incremental_state( + backend, context, keys[req_id], rope[req_id], length, req_id + ) + first_expected, _ = _oracle_completed(keys[0], rope[0]) + second_expected, _ = _oracle_completed(keys[1], rope[1]) + assert not torch.equal(first_expected[0], second_expected[0]) + + index_q = torch.ones(sum(lengths), 4, INDEX_DIM) + physical, counts = backend._select_physical_rows(index_q, LAYER_ID, batch) + offset = 0 + for req_id, length in enumerate(lengths): + for local_position in range(length): + row = offset + local_position + assert int(counts[row]) == local_position + 1 + actual = set(physical[row, : int(counts[row])].tolist()) + expected = set( + context.page_table[req_id, : local_position + 1].tolist() + ) + assert actual == expected + offset += length + + +def test_qsa_request_slot_reset_invalidates_pending_state_cpu(monkeypatch): + torch.manual_seed(SEED) + backend, context, _, old_keys, old_rope = _run_pattern( + monkeypatch, [5, 2, 1], torch.device("cpu") + ) + new_keys = _raw_keys(3, 1, torch.device("cpu")) + new_rope = _rope_positions(3, 1, torch.device("cpu")) + req = _request(0, 3, 0) + batch = _batch([req], [new_rope]) + backend.prepare_metadata(batch) + backend._compress_current_keys( + _RecordingSyntheticIndexer(), new_keys, LAYER_ID, batch + ) + assert torch.equal( + backend.kvcache.pending_group(LAYER_ID, 0, torch.arange(3)), new_keys + ) + with pytest.raises(RuntimeError, match="pending-key state is missing"): + backend.kvcache.pending_group(LAYER_ID, 0, torch.tensor([4])) + selected, counts = select_qsa_logical_rows( + torch.ones(1, 4, INDEX_DIM), + backend.kvcache.compressed_k_cache(LAYER_ID)[:0], + torch.tensor([2]), + compress_ratio=RATIO, + token_budget=BUDGET, + ) + assert int(counts[0]) == 3 + assert torch.equal(selected[0, :3], torch.arange(3, dtype=torch.int32)) + assert old_keys.shape[0] == old_rope.shape[1] == 8 + assert backend.kvcache is context.kv_cache + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_qsa_incremental_cache_public_surface_matches_cpu_oracle_cuda(monkeypatch): + torch.manual_seed(SEED) + torch.cuda.reset_peak_memory_stats() + started = time.perf_counter() + _run_pattern( + monkeypatch, + [1, 1, 1, 1, 1, 1, 1, 1], + torch.device("cuda"), + ) + torch.cuda.synchronize() + elapsed = time.perf_counter() - started + peak = torch.cuda.max_memory_allocated() + print( + "STAGE3_QSA_METRIC", + json.dumps( + { + "case": "cuda_incremental_cache", + "device": torch.cuda.get_device_name(), + "driver_runtime": torch.version.cuda, + "elapsed_seconds": elapsed, + "peak_memory_bytes": peak, + }, + sort_keys=True, + ), + ) + assert peak < 1 << 30 diff --git a/tests/models/test_qwen4_exp.py b/tests/models/test_qwen4_exp.py index 7d9912e4..ce902997 100644 --- a/tests/models/test_qwen4_exp.py +++ b/tests/models/test_qwen4_exp.py @@ -6,8 +6,13 @@ import torch import freetoken.models.qwen4_exp as qwen4_exp +import freetoken.models.qwen4_exp.model as qwen4_model from freetoken.models.qwen4_exp.config import parse_config -from freetoken.models.qwen4_exp.model import _ple_request_tokens, build_ngram_ids +from freetoken.models.qwen4_exp.model import ( + _HostNGramEmbedding, + _ple_request_tokens, + build_ngram_ids, +) from freetoken.models.qwen4_exp.weight import _rename, _try_fuse from freetoken.models.register import get_model_spec @@ -106,6 +111,13 @@ def test_qwen4_config_accepts_missing_norm_topk_prob(): hf_config = _config() del hf_config.text_config.norm_topk_prob config = parse_config(hf_config) + assert config.norm_topk_prob + + +def test_qwen4_config_preserves_explicit_false_norm_topk_prob(): + hf_config = _config() + hf_config.text_config.norm_topk_prob = False + config = parse_config(hf_config) assert not config.norm_topk_prob @@ -211,6 +223,19 @@ def test_ple_request_tokens_uses_complete_prefill_history(): assert _ple_request_tokens(req).tolist() == [11, 12, 13] +def test_ple_request_tokens_uses_bounded_non_overlap_decode_suffix(): + history = torch.tensor([10, 11, 12, 13, 14, 15]) + req = SimpleNamespace( + input_ids=history, + cached_len=5, + device_len=6, + extend_len=1, + ) + + assert _ple_request_tokens(req, start=3).tolist() == [13, 14, 15] + assert torch.equal(req.input_ids, history) + + def test_ple_request_tokens_joins_overlap_decode_token(): req = SimpleNamespace( input_ids=torch.tensor([11, 12]), @@ -221,6 +246,33 @@ def test_ple_request_tokens_joins_overlap_decode_token(): assert _ple_request_tokens(req, torch.tensor([13])).tolist() == [11, 12, 13] +def test_ple_request_tokens_bounds_overlap_multi_token_extension(): + history = torch.tensor([10, 11, 12, 13, 14, 15, 16, 17]) + req = SimpleNamespace( + input_ids=history, + cached_len=8, + device_len=11, + extend_len=3, + ) + + tokens = _ple_request_tokens(req, torch.tensor([18, 19, 20]), start=6) + + assert tokens.tolist() == [16, 17, 18, 19, 20] + assert tokens.numel() == 2 + req.extend_len + assert torch.equal(req.input_ids, history) + + +def test_ple_request_tokens_validates_overlap_forwarded_token_count(): + req = SimpleNamespace( + input_ids=torch.tensor([11, 12]), + cached_len=2, + device_len=4, + extend_len=2, + ) + with pytest.raises(RuntimeError, match="needs the current forwarded tokens"): + _ple_request_tokens(req, torch.tensor([13])) + + def test_ple_request_tokens_rejects_noncontiguous_host_history(): req = SimpleNamespace( input_ids=torch.tensor([11]), @@ -230,3 +282,94 @@ def test_ple_request_tokens_rejects_noncontiguous_host_history(): ) with pytest.raises(RuntimeError, match="unexpected gap"): _ple_request_tokens(req, torch.tensor([13])) + + +@pytest.mark.parametrize("start", [-1, 4]) +def test_ple_request_tokens_rejects_invalid_suffix_start(start): + req = SimpleNamespace( + input_ids=torch.tensor([11, 12, 13]), + cached_len=2, + device_len=3, + extend_len=1, + ) + with pytest.raises(ValueError, match="history start"): + _ple_request_tokens(req, start=start) + + +def test_incremental_ngram_ids_match_full_history_across_eos_boundary(): + tokens = torch.tensor([10, 99, 4, 5, 6]) + kwargs = { + "ngram_size": 3, + "heads_per_ngram": 1, + "eos_token_id": 99, + "multipliers": torch.tensor([3, 5, 7]), + "vocab_sizes": torch.tensor([101, 103]), + "offsets": torch.tensor([0, 101]), + } + cached_len = 3 + history_start = cached_len - (kwargs["ngram_size"] - 1) + + full = build_ngram_ids(tokens, **kwargs) + incremental = build_ngram_ids(tokens[history_start:], **kwargs) + + assert torch.equal( + incremental[cached_len - history_start :], + full[cached_len:], + ) + + +def test_current_ngram_ids_keeps_forwarded_offsets_request_local(monkeypatch): + req_a_history = torch.tensor([1, 2]) + req_b_history = torch.tensor([10, 11, 12]) + req_a = SimpleNamespace( + input_ids=req_a_history, + cached_len=2, + device_len=4, + extend_len=2, + ) + req_b = SimpleNamespace( + input_ids=req_b_history, + cached_len=3, + device_len=4, + extend_len=1, + ) + batch = SimpleNamespace( + is_decode=True, + padded_reqs=[req_a, req_b], + input_ids=torch.tensor([3, 4, 13]), + ) + monkeypatch.setattr(qwen4_model, "get_global_ctx", lambda: SimpleNamespace(batch=batch)) + + multipliers = torch.tensor([3, 5, 7]) + vocab_sizes = torch.tensor([101, 103]) + offsets = torch.tensor([0, 101]) + embedding = SimpleNamespace( + _host_constants=(multipliers, vocab_sizes, offsets), + ngram_size=3, + heads_per_ngram=1, + eos_token_id=99, + ) + + actual = _HostNGramEmbedding._current_ngram_ids(embedding) + expected_a = build_ngram_ids( + torch.tensor([1, 2, 3, 4]), + ngram_size=3, + heads_per_ngram=1, + eos_token_id=99, + multipliers=multipliers, + vocab_sizes=vocab_sizes, + offsets=offsets, + )[2:] + expected_b = build_ngram_ids( + torch.tensor([10, 11, 12, 13]), + ngram_size=3, + heads_per_ngram=1, + eos_token_id=99, + multipliers=multipliers, + vocab_sizes=vocab_sizes, + offsets=offsets, + )[3:] + + assert torch.equal(actual, torch.cat((expected_a, expected_b))) + assert torch.equal(req_a.input_ids, req_a_history) + assert torch.equal(req_b.input_ids, req_b_history) diff --git a/tests/models/test_qwen4_exp_nvfp4_components.py b/tests/models/test_qwen4_exp_nvfp4_components.py new file mode 100644 index 00000000..3efc69d9 --- /dev/null +++ b/tests/models/test_qwen4_exp_nvfp4_components.py @@ -0,0 +1,241 @@ +"""Synthetic Qwen4 active-weight NVFP4 component coverage. + +No model files are used. These tests exercise the deterministic host encoder, +canonical runtime fusion names, and the explicit GDN split without constructing +the full Qwen4 model. +""" + +from types import SimpleNamespace +import math + +import pytest +import torch + + +def _independent_nvfp4_reference(source: torch.Tensor): + """Tiny scalar oracle implementing the documented format, not FreeToken helpers.""" + + grid = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) + packed_rows, scale_rows, globals_out = [], [], [] + for row in source.float().tolist(): + row_max = max(abs(value) for value in row) + if row_max == 0: + global_scale = 1.0 + else: + target = min(max(row_max / 6.0, 2.0**-24), 65504.0) + global_scale = float(torch.tensor(target, dtype=torch.float16)) + if global_scale == 0: + global_scale = 2.0**-24 + globals_out.append(global_scale) + codes, block_scales = [], [] + for start in range(0, len(row), 16): + block = row[start : start + 16] + block_max = max(abs(value) for value in block) + target = 0.0 if block_max == 0 else min(block_max / (6.0 * global_scale), 448.0) + scale = float(torch.tensor(target, dtype=torch.float8_e4m3fn)) + block_scales.append(scale) + for value in block: + normalized = 0.0 if scale == 0 else max(-6.0, min(6.0, value / (scale * global_scale))) + magnitude = abs(normalized) + code = min(range(8), key=lambda item: (abs(grid[item] - magnitude), item & 1)) + codes.append(code | (8 if normalized < 0 else 0)) + packed_rows.append([codes[i] | (codes[i + 1] << 4) for i in range(0, len(codes), 2)]) + scale_rows.append(block_scales) + return ( + torch.tensor(packed_rows, dtype=torch.uint8), + torch.tensor(scale_rows, dtype=torch.float8_e4m3fn), + torch.tensor(globals_out, dtype=torch.float16), + ) + + +def test_encoder_is_deterministic_and_round_trips_layout(): + from freetoken.checkpoint.nvfp4 import decode_nvfp4, encode_bf16_nvfp4 + + torch.manual_seed(38038) + source = torch.randn(3, 32, dtype=torch.bfloat16) + first = encode_bf16_nvfp4(source) + second = encode_bf16_nvfp4(source.clone()) + assert all(torch.equal(a, b) for a, b in zip(first, second)) + packed, scales, globals_ = first + assert packed.shape == (3, 16) and packed.dtype == torch.uint8 + assert scales.shape == (3, 2) and scales.dtype == torch.float8_e4m3fn + assert globals_.shape == (3,) and globals_.dtype == torch.float16 + assert decode_nvfp4(packed, scales, globals_).shape == source.shape + + +def test_encoder_bytes_match_independent_scalar_oracle(): + from freetoken.checkpoint.nvfp4 import encode_bf16_nvfp4 + + torch.manual_seed(38038) + source = torch.randn(4, 32, dtype=torch.bfloat16) + actual = encode_bf16_nvfp4(source) + expected = _independent_nvfp4_reference(source) + assert all(torch.equal(a, b) for a, b in zip(actual, expected)) + + +@pytest.mark.parametrize( + "value", + [ + 0.0, + 2.0**-133, # smallest finite BF16 subnormal + 2.0**-126, # smallest finite BF16 normal + 448.0, + -448.0, + 3.38953139e38, # largest finite BF16 (saturating policy is defined) + ], +) +def test_encoder_total_for_finite_extrema(value): + from freetoken.checkpoint.nvfp4 import encode_bf16_nvfp4 + + source = torch.tensor([[value] * 16], dtype=torch.bfloat16) + packed, scales, globals_ = encode_bf16_nvfp4(source) + assert torch.isfinite(globals_).all() + assert torch.isfinite(scales.view(torch.float8_e4m3fn).float()).all() + assert packed.numel() == 8 + + +@pytest.mark.parametrize("bad", [float("nan"), float("inf"), -float("inf")]) +def test_encoder_rejects_nonfinite(bad): + from freetoken.checkpoint.nvfp4 import encode_bf16_nvfp4 + + with pytest.raises(ValueError, match="rejects NaN"): + encode_bf16_nvfp4(torch.full((1, 16), bad, dtype=torch.float32)) + + +def test_encoder_e2m1_midpoint_tie_to_even(): + from freetoken.checkpoint.nvfp4 import encode_bf16_nvfp4, decode_nvfp4 + + # Select a single block with global=1 and block=1; row max=6 establishes + # that scale, then the midpoint pairs exercise each E2M1 tie. + values = [0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0] + [6.0] * 9 + source = torch.tensor([values], dtype=torch.float32).to(torch.bfloat16) + packed, scales, globals_ = encode_bf16_nvfp4(source) + decoded = decode_nvfp4(packed, scales, globals_)[0] + # E2M1 even-code tie choices: 0, 1, 1, 2, 2, 4, 4. + assert decoded[:7].tolist() == pytest.approx([0.0, 1.0, 1.0, 2.0, 2.0, 4.0, 4.0]) + + +def test_qwen4_weight_fusions_use_runtime_state_names(): + from freetoken.models.qwen4_exp.weight import _try_fuse + + base = "model.layers.0.linear_attn." + buf = {} + assert _try_fuse(base + "in_proj_qkv.weight", torch.ones(2, 16), buf) == () + name, merged = _try_fuse(base + "in_proj_z.weight", torch.full((1, 16), 2.0), buf) + assert name == base + "in_proj_qkvz.weight" + assert merged.shape == (3, 16) + + buf = {} + assert _try_fuse(base + "in_proj_b.weight", torch.ones(1, 16), buf) == () + name, merged = _try_fuse(base + "in_proj_a.weight", torch.full((1, 16), 3.0), buf) + assert name == base + "in_proj_ba.weight" + assert merged[:, 0].tolist() == [1.0, 3.0] + + +def test_active_converter_emits_canonical_runtime_triple_and_preserves_slices(): + from freetoken.checkpoint.nvfp4 import encode_bf16_nvfp4 + from freetoken.models.qwen4_exp.weight import iter_active_nvfp4_runtime_entries + + q = torch.full((2, 16), 1.0, dtype=torch.bfloat16) + k = torch.full((1, 16), 2.0, dtype=torch.bfloat16) + v = torch.full((1, 16), -3.0, dtype=torch.bfloat16) + name = "model.layers.0.self_attn.qkv_proj.weight" + emitted = dict(iter_active_nvfp4_runtime_entries([(name, torch.cat((q, k, v), dim=0))])) + prefix = name.removesuffix(".weight") + assert list(emitted) == [name, prefix + ".weight_scale", prefix + ".weight_global"] + packed, scales, globals_ = emitted[name], emitted[prefix + ".weight_scale"], emitted[prefix + ".weight_global"] + cursor = 0 + for constituent in (q, k, v): + expected = encode_bf16_nvfp4(constituent) + end = cursor + constituent.shape[0] + assert torch.equal(packed[cursor:end], expected[0]) + assert torch.equal(scales[cursor:end], expected[1]) + assert torch.equal(globals_[cursor:end], expected[2]) + cursor = end + + +def test_active_converter_protects_non_map_tensors(): + from freetoken.models.qwen4_exp.weight import iter_active_nvfp4_runtime_entries + + protected = torch.randn(3, 16, dtype=torch.bfloat16) + rows = list(iter_active_nvfp4_runtime_entries([ + ("model.layers.0.self_attn.index_qk_proj.weight", protected), + ])) + assert len(rows) == 1 and rows[0][0].endswith("index_qk_proj.weight") + assert rows[0][1] is protected + + +def test_gdn_nvfp4_qkvz_is_explicit_opt_in(): + from freetoken.kernel.triton.nvfp4_linear import Nvfp4DenseColMerged, Nvfp4DenseLinear + from freetoken.layers import LinearColParallelMerged + from freetoken.models.qwen3_5_moe.gdn import Qwen3_5GatedDeltaNet + from freetoken.distributed import set_tp_info, try_get_tp_info + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + + kwargs = dict( + hidden_size=32, + num_k_heads=2, + num_v_heads=2, + head_k_dim=8, + head_v_dim=8, + conv_kernel_size=4, + rms_norm_eps=1e-6, + layer_id=0, + expert_quant="none", + attn_quant="nvfp4", + ) + legacy = Qwen3_5GatedDeltaNet(**kwargs) + assert isinstance(legacy.in_proj, LinearColParallelMerged) + explicit = Qwen3_5GatedDeltaNet(**kwargs, nvfp4_qkvz=True) + assert isinstance(explicit.in_proj_qkvz, Nvfp4DenseColMerged) + assert isinstance(explicit.in_proj_ba, LinearColParallelMerged) + assert isinstance(explicit.out_proj, Nvfp4DenseLinear) + + +def test_qwen4_frozen_operator_map_and_protected_linears(): + from freetoken.kernel.triton.nvfp4_linear import Nvfp4DenseColMerged, Nvfp4DenseLinear + from freetoken.layers import LinearReplicated + from freetoken.models.qwen3_5_moe.attention import Qwen3_5Attention + from freetoken.models.qwen4_exp.model import _GatedResidual, _SharedExpert + + rotary = SimpleNamespace(rotary_dim=128, max_position=128, base=10_000.0, scaling=None) + attention_config = SimpleNamespace( + head_dim=128, num_qo_heads=2, num_kv_heads=1, hidden_size=256, + rms_norm_eps=1e-6, rotary_config=rotary, expert_quant="none", attn_quant="nvfp4", + ) + attention = Qwen3_5Attention(attention_config, 0) + assert isinstance(attention.qkv_proj, Nvfp4DenseColMerged) + assert isinstance(attention.o_proj, Nvfp4DenseLinear) + + qwen_config = SimpleNamespace( + hidden_size=32, rms_norm_eps=1e-6, dense_quant="nvfp4", + shared_expert_intermediate_size=16, + qwen4_args=SimpleNamespace(hc_count=4, hc_lowrank=16), + ) + residual = _GatedResidual(qwen_config, combine=True) + assert isinstance(residual.input_mix_weight_down, Nvfp4DenseLinear) + assert isinstance(residual.input_mix_weight_up, Nvfp4DenseLinear) + assert isinstance(residual.block_inject_weight, LinearReplicated) + shared = _SharedExpert(qwen_config) + assert isinstance(shared.gate_up_proj, Nvfp4DenseColMerged) + assert isinstance(shared.down_proj, Nvfp4DenseLinear) + + +def test_legacy_gdn_quant_modes_keep_their_original_dispatch(): + from freetoken.kernel.triton.fp8_block_linear import Fp8BlockColMerged + from freetoken.kernel.triton.fp8_pertensor_linear import Fp8PerTensorColMerged + from freetoken.layers import LinearColParallelMerged + from freetoken.models.qwen3_5_moe.gdn import Qwen3_5GatedDeltaNet + + base = dict( + hidden_size=256, num_k_heads=1, num_v_heads=1, head_k_dim=128, + head_v_dim=128, conv_kernel_size=4, rms_norm_eps=1e-6, layer_id=0, + ) + bf16 = Qwen3_5GatedDeltaNet(**base, expert_quant="none", attn_quant="none") + assert isinstance(bf16.in_proj, LinearColParallelMerged) + block = Qwen3_5GatedDeltaNet(**base, expert_quant="fp8_block", attn_quant="none") + assert isinstance(block.in_proj_qkvz, Fp8BlockColMerged) + pertensor = Qwen3_5GatedDeltaNet(**base, expert_quant="none", attn_quant="fp8_pertensor") + assert isinstance(pertensor.in_proj_qkvz, Fp8PerTensorColMerged) diff --git a/tests/models/test_qwen4_exp_raw_config.py b/tests/models/test_qwen4_exp_raw_config.py new file mode 100644 index 00000000..e9d17d53 --- /dev/null +++ b/tests/models/test_qwen4_exp_raw_config.py @@ -0,0 +1,147 @@ +from __future__ import annotations + +from freetoken.models.qwen4_exp.config import parse_config +from freetoken.utils.hf import RawConfigShim + + +def _raw_checkpoint_config() -> RawConfigShim: + """Raw config shape used when installed Transformers predates Qwen4-Exp.""" + return RawConfigShim( + { + "architectures": ["Qwen4ExpForConditionalGeneration"], + "model_type": "qwen4_exp", + "image_token_id": 248056, + "quantization_config": { + "quant_method": "fp8", + "weight_block_size": [128, 128], + }, + "text_config": { + "model_type": "qwen4_exp_text", + "layer_types": [ + "linear_attention", + "linear_attention", + "linear_attention", + "qwen_sparse_attention", + ], + "head_dim": 256, + "rope_parameters": { + "partial_rotary_factor": 0.25, + "rope_theta": 10_000_000, + "rope_type": "default", + "mrope_interleaved": True, + "mrope_section": [11, 11, 10], + }, + "indexer_budget": 2048, + "indexer_n_heads": 4, + "indexer_kv_heads": 1, + "indexer_head_dim": 128, + "indexer_compress_ratio": 4, + "max_position_embeddings": 262_144, + "num_key_value_heads": 2, + "linear_num_key_heads": 16, + "linear_num_value_heads": 48, + "linear_key_head_dim": 128, + "linear_value_head_dim": 128, + "linear_conv_kernel_dim": 4, + "eos_token_id": 248044, + "hc_count": 4, + "hc_lowrank": 320, + "ple_layer_ids": [2], + "ple_embed_dim": 2560, + "ple_conv_kernel_size": 4, + "ngram_size": 3, + "heads_per_ngram": 8, + "ngram_vocab_size_base": 20_000_000, + "split_ngram_parts": 128, + "output_gate_type": "sigmoid", + "hidden_act": "silu", + "num_hidden_layers": 4, + "num_attention_heads": 24, + "hidden_size": 2560, + "vocab_size": 248320, + "rms_norm_eps": 1e-6, + "num_experts": 512, + "num_experts_per_tok": 10, + "moe_intermediate_size": 640, + "shared_expert_intermediate_size": 640, + "tie_word_embeddings": False, + }, + } + ) + + +def test_qwen4_raw_config_uses_official_topk_normalization_default(): + config = parse_config(_raw_checkpoint_config()) + + assert config.norm_topk_prob is True + assert config.rotary_config.max_position == 262_144 + assert config.attn_type_for_layer(3).value == "qsa" + + +def test_qwen4_active_nvfp4_requires_explicit_artifact_marker(): + published = _raw_checkpoint_config() + config = parse_config(published) + assert config.attn_quant == "none" + assert config.dense_quant == "none" + + converted = _raw_checkpoint_config() + converted.freetoken_active_quant = "nvfp4_w4a16_v1" + config = parse_config(converted) + assert config.attn_quant == "nvfp4" + assert config.dense_quant == "nvfp4" + + +def test_qwen4_rejects_unknown_active_weight_marker(): + config = _raw_checkpoint_config() + config.freetoken_active_quant = "ambiguous-q8" + try: + parse_config(config) + except ValueError as exc: + assert "unsupported Qwen4 active-weight format" in str(exc) + else: + raise AssertionError("unknown active-weight marker was accepted") + + +def _add_minimal_vision(config: RawConfigShim) -> None: + config.vision_config = RawConfigShim( + { + "depth": 2, + "hidden_size": 32, + "intermediate_size": 64, + "num_heads": 4, + "num_position_embeddings": 16, + "out_hidden_size": 2560, + "patch_size": 14, + "spatial_merge_size": 2, + "temporal_patch_size": 2, + "in_channels": 3, + "hidden_act": "gelu", + "deepstack_visual_indexes": [0], + } + ) + + +def test_qwen4_text_only_marker_disables_only_marked_artifact_vision(): + published = _raw_checkpoint_config() + _add_minimal_vision(published) + parsed = parse_config(published) + assert parsed.vision_config is not None + assert parsed.image_token_id == 248056 + + target = _raw_checkpoint_config() + _add_minimal_vision(target) + target.freetoken_text_only = "qwen4_text_only_v1" + parsed = parse_config(target) + assert parsed.vision_config is None + assert parsed.image_token_id is None + + +def test_qwen4_rejects_unknown_text_only_marker(): + config = _raw_checkpoint_config() + config.freetoken_text_only = "textish-v0" + try: + parse_config(config) + except ValueError as exc: + assert "unsupported freetoken_text_only marker" in str(exc) + else: + raise AssertionError("unknown text-only marker was accepted") diff --git a/tests/moe/test_expert_sidecar_writer.py b/tests/moe/test_expert_sidecar_writer.py new file mode 100644 index 00000000..604ab21c --- /dev/null +++ b/tests/moe/test_expert_sidecar_writer.py @@ -0,0 +1,174 @@ +from __future__ import annotations + +import hashlib +import shutil +from pathlib import Path +from uuid import uuid4 + +import pytest +import torch + +from freetoken.moe.expert_source import ( + ExpertSourceError, + FileExpertSource, + MAGIC, + adapt_expert_tensor_record, + write_expert_sidecar, + write_expert_sidecar_from_safetensors, +) + + +@pytest.fixture +def z_dir(): + root = Path.cwd() / ".stage7-test-fixtures" / uuid4().hex + root.mkdir(parents=True) + assert (root.drive or "").upper() == "Z:" + try: + yield root + finally: + shutil.rmtree(root, ignore_errors=True) + + +def _geometry(): + return { + "gate_up_packed": ((4, 4), torch.uint8), + "gate_up_scale": ((4, 1), torch.float8_e4m3fn), + "gate_up_global": ((4,), torch.float16), + "down_packed": ((2, 16), torch.uint8), + "down_scale": ((2, 1), torch.float8_e4m3fn), + "down_global": ((2,), torch.float16), + } + + +def _planes(value: int): + out = {} + for name, (shape, dtype) in _geometry().items(): + if dtype == torch.float8_e4m3fn: + out[name] = torch.full(shape, 1, dtype=dtype) + else: + out[name] = torch.full(shape, value, dtype=dtype) + return out + + +def _source_record(value: int): + out = {} + for projection, rows, width in (("gate_proj", 2, 4), ("up_proj", 2, 4), ("down_proj", 2, 16)): + out[f"{projection}.weight"] = torch.full((rows, width), value, dtype=torch.uint8) + out[f"{projection}.weight_scale"] = torch.full((rows, max(1, width // 16)), 1, dtype=torch.float8_e4m3fn) + out[f"{projection}.weight_scale_2"] = torch.tensor(1.5, dtype=torch.float32) + out[f"{projection}.input_scale"] = torch.tensor(1.0, dtype=torch.float32) + return out + + +def test_writer_reduced_geometry_reopens_and_is_deterministic(z_dir): + first = z_dir / "layer-00.ftex" + second = z_dir / "layer-00-copy.ftex" + kwargs = dict(layer_id=7, source_fingerprint=b"source", num_experts=3, geometry=_geometry()) + result = write_expert_sidecar(first, [_planes(i) for i in range(3)], **kwargs) + result_copy = write_expert_sidecar(second, [_planes(i) for i in range(3)], **kwargs) + assert result["format"] == "FTEXPERT1" + assert result["raw_record_bytes"] == 66 + assert result["record_bytes"] == 4096 + assert result["sample_ids"] == (0, 1, 2) + assert first.read_bytes() == second.read_bytes() + assert result["sha256"] == hashlib.sha256(first.read_bytes()).hexdigest() + assert result["sha256"] == result_copy["sha256"] + with FileExpertSource(first, num_experts=3, expected_sha256=result["sha256"], expected_layer_id=7) as source: + assert source.read_record(0)["gate_up_packed"].flatten()[0].item() == 0 + assert source.read_record(2)["gate_up_packed"].flatten()[0].item() == 2 + assert source.record_bytes == 4096 + + +def test_source_tensor_adapter_validates_twelve_names_and_expands_globals(z_dir): + source_record = _source_record(9) + adapted = adapt_expert_tensor_record(source_record) + assert adapted["gate_up_packed"].shape == (4, 4) + assert adapted["gate_up_scale"].shape == (4, 1) + assert adapted["gate_up_global"].dtype == torch.float16 + assert adapted["gate_up_global"].tolist() == [1.5] * 4 + assert adapted["down_global"].tolist() == [1.5] * 2 + result = write_expert_sidecar( + z_dir / "source.ftex", [source_record], layer_id=0, source_fingerprint=b"source", num_experts=1, geometry=_geometry() + ) + with FileExpertSource(result["path"], num_experts=1) as source: + assert source.read_record(0)["gate_up_packed"].flatten()[0].item() == 9 + + +def test_production_writer_streams_indexed_safetensor_experts(z_dir): + import json + from safetensors.torch import save_file + + prefix = "model.language_model.layers.3.mlp.experts" + tensors = {} + weight_map = {} + for expert_id in range(2): + for suffix, tensor in _source_record(5 + expert_id).items(): + name = f"{prefix}.{expert_id}.{suffix}" + tensors[name] = tensor + weight_map[name] = "layer-3.safetensors" + save_file(tensors, z_dir / "layer-3.safetensors") + (z_dir / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": weight_map}), encoding="utf-8" + ) + output = z_dir / "experts-L03.nvfp4" + result = write_expert_sidecar_from_safetensors( + z_dir, + output, + layer_id=3, + source_fingerprint="a" * 64, + num_experts=2, + geometry=_geometry(), + ) + assert result["sample_ids"] == (0, 1) + with FileExpertSource(output, num_experts=2, expected_layer_id=3) as source: + assert int(source.read_record(0)["gate_up_packed"][0, 0]) == 5 + assert int(source.read_record(1)["down_packed"][0, 0]) == 6 + + +def test_writer_accepts_explicit_id_and_named_pairs(z_dir): + named = [(name, value) for name, value in _planes(4).items()] + path = z_dir / "explicit.ftex" + write_expert_sidecar( + path, + [(0, named)], + layer_id=0, + source_fingerprint="fixture", + num_experts=1, + geometry=_geometry(), + ) + with FileExpertSource(path, num_experts=1) as source: + assert source.read_record(0)["down_global"].flatten()[0].item() == 4 + + +@pytest.mark.parametrize("records", [ + [(0, _planes(1)), (0, _planes(2))], + [(0, _planes(1))], + [(0, _planes(1)), (2, _planes(2))], +]) +def test_writer_rejects_duplicate_or_missing_ids_without_publishing(z_dir, records): + path = z_dir / "bad.ftex" + with pytest.raises(ValueError, match="(duplicate|missing|outside)"): + write_expert_sidecar(path, records, layer_id=0, source_fingerprint=b"x", num_experts=2, geometry=_geometry()) + assert not path.exists() + assert not Path(str(path) + ".partial").exists() + + +def test_writer_partial_and_payload_corruption_fail_closed(z_dir): + path = z_dir / "corrupt.ftex" + result = write_expert_sidecar(path, [_planes(1)], layer_id=0, source_fingerprint=b"x", num_experts=1, geometry=_geometry()) + with path.open("r+b") as handle: + handle.seek(4096 + 1) + handle.write(b"x") + with pytest.raises(ExpertSourceError, match="payload hash mismatch"): + FileExpertSource(path, num_experts=1) + path.unlink() + partial = Path(str(path) + ".partial") + partial.write_bytes(b"partial") + with pytest.raises(ExpertSourceError): + FileExpertSource(partial, num_experts=1) + + +def test_writer_emits_ft_expert_magic(z_dir): + path = z_dir / "magic.ftex" + write_expert_sidecar(path, [_planes(1)], layer_id=0, source_fingerprint=b"x", num_experts=1, geometry=_geometry()) + assert path.read_bytes()[: len(MAGIC)] == MAGIC diff --git a/tests/moe/test_file_expert_source.py b/tests/moe/test_file_expert_source.py new file mode 100644 index 00000000..996003aa --- /dev/null +++ b/tests/moe/test_file_expert_source.py @@ -0,0 +1,231 @@ +from __future__ import annotations + +import hashlib +import shutil +from pathlib import Path +from uuid import uuid4 + +import pytest +import torch + +from freetoken.moe.expert_source import ( + ExpertSourceError, + FileExpertSource, + PLANE_LAYOUT, + RAW_RECORD_BYTES, +) + + +@pytest.fixture +def z_fixture_dir(): + root = Path.cwd() / ".stage6-test-fixtures" / uuid4().hex + root.mkdir(parents=True, exist_ok=False) + try: + assert (root.drive or "").upper() == "Z:", root + yield root + finally: + shutil.rmtree(root, ignore_errors=True) + + +def _record(byte: int) -> bytes: + return bytes([byte]) * RAW_RECORD_BYTES + + +def _resident_banks(num_layers: int, experts: int): + shapes = { + "gate_up_packed": ((experts, 1280, 1280), torch.uint8), + "gate_up_scale": ((experts, 1280, 160), torch.float8_e4m3fn), + "gate_up_global": ((experts, 1280), torch.float16), + "down_packed": ((experts, 2560, 320), torch.uint8), + "down_scale": ((experts, 2560, 40), torch.float8_e4m3fn), + "down_global": ((experts, 2560), torch.float16), + } + return { + name: [torch.zeros(shape, dtype=dtype) if layer == 0 else None for layer in range(num_layers)] + for name, (shape, dtype) in shapes.items() + } + + +def test_file_expert_source_reads_exact_planes_and_rejects_tamper(z_fixture_dir): + root = z_fixture_dir + path = root / "experts-L00.nvfp4" + digest = FileExpertSource.create_synthetic(path, num_experts=3, records=[_record(i) for i in range(3)]) + with FileExpertSource(path, num_experts=3, expected_sha256=digest) as src: + rows = src.read_record(2) + assert set(rows) == {name for name, _, _ in PLANE_LAYOUT} + assert rows["gate_up_packed"].shape == (1280, 1280) + assert rows["gate_up_scale"].shape == (1280, 160) + assert int(rows["gate_up_packed"].flatten()[0]) == 2 + assert src.read_count == 1 + assert src.max_inflight <= 1 <= 16 + with pytest.raises(ExpertSourceError, match="closed"): + src.read_record(0) + + +def test_file_expert_source_cache_miss_fills_slots_without_host_layer(z_fixture_dir): + path = z_fixture_dir / "experts-L01.nvfp4" + digest = FileExpertSource.create_synthetic(path, num_experts=3, records=[_record(i + 11) for i in range(3)], layer_id=1) + src = FileExpertSource(path, num_experts=3, expected_sha256=digest, expected_layer_id=1) + try: + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache( + num_layers=2, + num_experts=3, + cache_size=3, + device=torch.device("cpu"), + quant_format="nvfp4", + decode_target="gpu", + ) + cache.set_bank_sources(_resident_banks(2, 3)) + cache.set_file_sources({1: src}) + assert all(cache.bank_sources[name][1] is None for name in cache.bank_schema) + ids = torch.tensor([2, 0], dtype=torch.int32) + cache.ensure_experts(1, ids) + assert ids.tolist() == [0, 1] + cache.copy_missing() + assert cache.bank_caches["gate_up_packed"][0, 0, 0].item() == 13 + assert cache.bank_caches["gate_up_packed"][1, 0, 0].item() == 11 + assert src.bytes_read == 2 * src.record_bytes + assert cache._pending_file_fetches == [] + cache.reset() + assert cache.id_of_slot.tolist() == [-1, -1, -1] + assert cache.slot_for_id.tolist() == [[-1, -1, -1], [-1, -1, -1]] + finally: + src.close() + + +def test_file_tier_materialize_streams_complete_layer(z_fixture_dir): + path = z_fixture_dir / "experts-L01.nvfp4" + digest = FileExpertSource.create_synthetic(path, num_experts=3, records=[_record(i + 21) for i in range(3)], layer_id=1) + src = FileExpertSource(path, num_experts=3, expected_sha256=digest, expected_layer_id=1) + try: + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache(2, 3, 3, torch.device("cpu"), quant_format="nvfp4") + cache.set_bank_sources(_resident_banks(2, 3)) + cache.set_file_sources({1: src}) + cache.materialize_layer(1) + cache.copy_missing() + values = cache.bank_caches["gate_up_packed"][:, 0, 0].tolist() + assert values == [21, 22, 23] + assert src.read_count == 3 + assert cache.slot_for_id[1].tolist() == [0, 1, 2] + assert cache.id_of_slot.tolist() == [3, 4, 5] + finally: + src.close() + + +def test_file_only_cache_derives_planes_without_any_hostbank(z_fixture_dir): + path = z_fixture_dir / "experts-L00.nvfp4" + digest = FileExpertSource.create_synthetic(path, num_experts=2, records=[_record(41), _record(42)]) + src = FileExpertSource(path, num_experts=2, expected_sha256=digest) + try: + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache(1, 2, 2, torch.device("cpu"), quant_format="nvfp4") + cache.set_file_sources({0: src}) + assert all(all(item is None for item in layers) for layers in cache.bank_sources.values()) + ids = torch.tensor([1], dtype=torch.int32) + cache.ensure_experts(0, ids) + cache.copy_missing() + assert cache.bank_caches["gate_up_packed"][0, 0, 0].item() == 42 + finally: + src.close() + + +def test_file_only_multilayer_identity_eviction_and_reset(z_fixture_dir): + path0 = z_fixture_dir / "experts-L00.nvfp4" + path1 = z_fixture_dir / "experts-L01.nvfp4" + digest0 = FileExpertSource.create_synthetic(path0, num_experts=2, records=[_record(51), _record(52)], layer_id=0) + digest1 = FileExpertSource.create_synthetic(path1, num_experts=2, records=[_record(61), _record(62)], layer_id=1) + src0 = FileExpertSource(path0, num_experts=2, expected_sha256=digest0, expected_layer_id=0) + src1 = FileExpertSource(path1, num_experts=2, expected_sha256=digest1, expected_layer_id=1) + try: + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache(2, 2, 2, torch.device("cpu"), quant_format="nvfp4") + cache.set_file_sources({0: src0, 1: src1}) + ids0 = torch.tensor([0, 1], dtype=torch.int32) + cache.ensure_experts(0, ids0) + cache.copy_missing() + ids1 = torch.tensor([0], dtype=torch.int32) + cache.ensure_experts(1, ids1) + cache.copy_missing() + slot = int(ids1.item()) + assert cache.slot_for_id[1, 0].item() == slot + assert cache.slot_for_id[0, slot].item() == -1 + assert cache.bank_caches["gate_up_packed"][slot, 0, 0].item() == 61 + cache.reset() + assert (cache.slot_for_id == -1).all() + assert (cache.id_of_slot == -1).all() + finally: + src0.close() + src1.close() + + +def test_file_tier_read_failure_rolls_back_slot_bookkeeping(z_fixture_dir, monkeypatch): + path = z_fixture_dir / "experts-L01.nvfp4" + digest = FileExpertSource.create_synthetic(path, num_experts=2, records=[_record(31), _record(32)], layer_id=1) + src = FileExpertSource(path, num_experts=2, expected_sha256=digest, expected_layer_id=1) + try: + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache(2, 2, 2, torch.device("cpu"), quant_format="nvfp4") + cache.set_bank_sources(_resident_banks(2, 2)) + cache.set_file_sources({1: src}) + before_slots = cache.slot_for_id.clone() + before_ids = cache.id_of_slot.clone() + ids = torch.tensor([1], dtype=torch.int32) + cache.ensure_experts(1, ids) + + def fail_read(*_args, **_kwargs): + raise ExpertSourceError("synthetic short read") + + monkeypatch.setattr(src, "read_into", fail_read) + with pytest.raises(ExpertSourceError, match="short read"): + cache.copy_missing() + assert torch.equal(cache.slot_for_id, before_slots) + assert torch.equal(cache.id_of_slot, before_ids) + assert cache._pending_file_fetches == [] + finally: + src.close() + + +def test_file_tier_payload_hash_fails_closed(z_fixture_dir): + path = z_fixture_dir / "experts-L00.nvfp4" + FileExpertSource.create_synthetic(path, num_experts=1, records=[_record(9)]) + with path.open("r+b") as fh: + fh.seek(4096 + 17) + fh.write(b"x") + with pytest.raises(ExpertSourceError, match="payload hash mismatch"): + FileExpertSource(path, num_experts=1) + + +def test_file_tier_rejects_only_cpu_selected_file_layers(z_fixture_dir): + path = z_fixture_dir / "experts-L00.nvfp4" + digest = FileExpertSource.create_synthetic(path, num_experts=1, records=[_record(7)]) + src = FileExpertSource(path, num_experts=1, expected_sha256=digest) + try: + from freetoken.moe.offload_cache import OffloadMoeCache + + for target in ("cpu", "hybrid"): + cache = OffloadMoeCache(1, 1, 1, torch.device("cpu"), decode_target=target, quant_format="nvfp4") + cache.cpu_layer_ids = frozenset({0}) + # Source registration itself fails before any source can be used; + # a direct call is sufficient to prove the policy and avoids giant + # synthetic resident allocations in this negative test. + with pytest.raises(ValueError, match="GPU-only"): + cache.set_file_sources({0: src}) + + mixed = OffloadMoeCache(2, 1, 1, torch.device("cpu"), decode_target="cpu", quant_format="nvfp4") + mixed.cpu_layer_ids = frozenset({1}) + mixed.set_file_sources({0: src}) + assert set(mixed.file_sources) == {0} + finally: + src.close() + + +def test_file_expert_source_rejects_wrong_volume(monkeypatch): + with pytest.raises(ExpertSourceError, match="Z:"): + FileExpertSource(r"C:\\not-a-tier.nvfp4") diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index 8b62a907..e573c929 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -415,7 +415,9 @@ def test_adjust_config_converts_moe_cache_rate_to_cache_size(): model_path="/tmp/freetoken-test-model", tp_info=DistributedInfo(rank=0, size=1), dtype=torch.float16, - attention_backend="fi", + # This test exercises MoE cache-rate conversion, not FlashInfer package + # discovery; keep it runnable in the project-local component-test env. + attention_backend="triton", moe_cache_rate=0.3, ) object.__setattr__( diff --git a/tests/server/test_parser_auto_selection.py b/tests/server/test_parser_auto_selection.py index 78b1663e..78ceaaff 100644 --- a/tests/server/test_parser_auto_selection.py +++ b/tests/server/test_parser_auto_selection.py @@ -87,6 +87,14 @@ def test_qwen3_5_is_not_shadowed_by_the_generic_qwen_branch(): assert _inferred("Qwen3MoeForCausalLM")[0] == "qwen25" +@pytest.mark.parametrize( + "architecture", + ["Qwen4ExpForConditionalGeneration", "qwen4_exp"], +) +def test_qwen4_exp_uses_qwen3_coder_tool_parser(architecture): + assert _inferred(architecture)[0] == "qwen3_coder" + + def test_an_explicit_choice_beats_inference(): config = _Config({"architectures": ["DeepseekV4ForCausalLM"], "torch_dtype": "bfloat16"}) with patch("freetoken.utils.cached_load_hf_config", lambda _path: config):