From 189c8ad0c3350e30252209d40780274976cb7527 Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Sun, 4 Oct 2026 13:51:55 -0700 Subject: [PATCH 01/73] feat(moe): serve NVFP4 and MXFP4 experts fused through the expert banks (#601) * feat(moe): serve NVFP4 and MXFP4 experts fused through the expert banks Resident experts skip the state dict and load through the offload expert-bank loader (HF pieces or FTW banks) into per-layer GPU banks, so fused serves bf16, fp8_block, NVFP4 and MXFP4 alike. Unified-memory GPUs auto-select fused for these formats and ignore offload-only flags with a warning. BREAKING CHANGE: ft checkpoint always writes MoE experts as banks; an FTW converted with --moe-backend fused or triton (dense experts) no longer loads and must be reconverted. * docs(moe): name only GB10 / DGX Spark for unified memory --- docs/cli.md | 9 +- docs/models.md | 17 ++- python/freetoken/checkpoint/__main__.py | 2 +- python/freetoken/checkpoint/convert.py | 4 +- python/freetoken/checkpoint/ftw.py | 140 ++++++++++++++---- python/freetoken/engine/engine.py | 59 ++++++-- python/freetoken/layers/__init__.py | 3 +- python/freetoken/layers/moe.py | 17 ++- .../freetoken/layers/quantization/moe/base.py | 15 +- .../layers/quantization/moe/fp8_block.py | 16 +- .../layers/quantization/moe/mxfp4.py | 36 +---- .../layers/quantization/moe/nvfp4.py | 14 +- .../layers/quantization/moe/unquantized.py | 13 +- python/freetoken/models/qwen3_5_moe/weight.py | 39 +---- python/freetoken/moe/expert_banks.py | 71 +++++++-- python/freetoken/moe/offload_cache.py | 15 +- python/freetoken/server/args.py | 4 +- tests/engine/test_cache_budget.py | 7 +- tests/models/qwen4_exp/test_skeleton.py | 9 +- tests/models/test_quant_config.py | 20 +-- 20 files changed, 298 insertions(+), 212 deletions(-) diff --git a/docs/cli.md b/docs/cli.md index fb5e613f7b..deeda18419 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -80,7 +80,7 @@ See [models.md](models.md#moe-strategies) for what each strategy does. | Flag | Default | Meaning | |---|---|---| -| `--moe-strategy` | auto | `fused`/`offload`/`cpu`/`hybrid`; auto → offload, or hybrid with a `ft bench bw` profile. `--moe-backend` is the deprecated old spelling | +| `--moe-strategy` | auto | `fused`/`offload`/`cpu`/`hybrid`; auto → offload, or hybrid with a `ft bench bw` profile; fused on unified-memory GPUs (GB10), see [models.md](models.md#unified-memory-gpus-gb10--dgx-spark). `--moe-backend` is the deprecated old spelling | | `--quant-backend` | auto | Kernel per quantized layer type, `layer[.kind]=name` entries: `linear=marlin,moe=b12x` or `moe.nvfp4=triton`. A layer-level entry applies to every kind whose table lists the name | | `--nvfp4-backend` | — | Deprecated: stands in for `--quant-backend moe.nvfp4=` (`flashinfer` means b12x); cannot be combined with `--quant-backend` | | `--moe-cache-size` / `--moe-cache-rate` / `--moe-cache-auto` | auto | GPU expert-cache size as slots / fraction of all experts / sized from free VRAM (mutually exclusive; auto is enabled by default for offload-family strategies) | @@ -179,9 +179,10 @@ ft checkpoint --model --out [--dtype bfloat16] [--moe-backend ``` Converts an HF safetensors checkpoint to FTW, FreeToken's self-contained -fast-load format; point `ft serve --model` at the output dir. `--moe-backend -offload` (default) packs experts into offload banks; `--moe-backend triton` -keeps them dense for resident serving. See the FTW caveats in +fast-load format; point `ft serve --model` at the output dir. MoE experts are +always packed into expert banks, which serve the fused and the offload +strategies alike (GGUF and DeepSeek-V4 experts serve offload only); +`--moe-backend` is kept for compatibility. See the FTW caveats in [models.md](models.md#notes); FTW files from older builds can be repaired with [scripts/ftw_hotfix.py](ftw-hotfix.md) instead of reconverting. diff --git a/docs/models.md b/docs/models.md index 548a363c67..efd5268b37 100644 --- a/docs/models.md +++ b/docs/models.md @@ -38,7 +38,8 @@ These families accept image input by default; pass `--text-model-only` to skip t `ft serve --moe-strategy {auto,fused,offload,cpu,hybrid}` (`--moe-backend` is the deprecated old spelling): -- **fused** — experts resident on GPU (needs the VRAM); never auto-selected. +- **fused** — experts resident on GPU (needs the VRAM). Serves every expert + format except GGUF and DeepSeek-V4's. - **offload** — experts live in host RAM, an LRU cache of expert slots on GPU; misses stream over PCIe. - **cpu** — misses are computed on the CPU instead of fetched. @@ -46,12 +47,24 @@ These families accept image input by default; pass `--text-model-only` to skip t CPU, overlapped. Run `ft bench bw` once per machine to calibrate the split. - **auto** — dense models always resolve to `fused`; MoE models resolve to `offload`, upgraded to `hybrid` when a cached `ft bench bw` profile - recommends it. + recommends it. Unified-memory GPUs differ, see below. + +### Unified-memory GPUs (GB10 / DGX Spark) + +The GPU and the CPU share one DRAM, so offload only copies experts between two +names for the same memory. On these GPUs `auto` resolves every MoE model to +`fused`: the default is the resident path only. GGUF and DeepSeek-V4 experts +have no resident path and stay on `offload`. Offload-only flags +(`--moe-cpu-layers`, `--moe-cache-size` / `-rate` / `-auto`, +`--moe-prefill-hit-d2d`, `--disable-moe-prefill-overlap`) are ignored with a +warning; pass `--moe-strategy offload` to use them. ## Notes - `ft checkpoint` conversion is optional — it pre-converts a checkpoint into FreeToken's fast-load format, and `ft serve --model` auto-detects the result. +- An FTW converted by an older build with `--moe-backend fused` or `triton` keeps + its experts as dense weights and no longer loads; reconvert it with `ft checkpoint`. - FTW files converted by builds before the quantization refactor may fail to load; see [ftw-hotfix.md](ftw-hotfix.md) for the affected checkpoints and the repair tool. - An FTW converted before its family served images holds no vision encoder: `ft serve` diff --git a/python/freetoken/checkpoint/__main__.py b/python/freetoken/checkpoint/__main__.py index 2b7e35f1d1..0d7c2d7ed7 100644 --- a/python/freetoken/checkpoint/__main__.py +++ b/python/freetoken/checkpoint/__main__.py @@ -37,7 +37,7 @@ def main(argv: list[str] | None = None, prog: str = "freetoken.checkpoint") -> i p.add_argument("--out", required=True, help="output FTW checkpoint dir") p.add_argument("--dtype", choices=sorted(_DTYPES), default="bfloat16") p.add_argument("--moe-backend", default="offload", - help="offload (experts -> banks) or e.g. triton (experts stay dense)") + help="kept for compatibility: MoE experts always go to banks, which every strategy reads") p.add_argument("--quant-backend", type=_parse_quant_backend, default=None, help="kernel per quantized layer type, as for ft serve; the expert banks are packed for the " "chosen MoE kernel and the server picks that kernel back up from the FTW") diff --git a/python/freetoken/checkpoint/convert.py b/python/freetoken/checkpoint/convert.py index a875d77f34..152f3871c0 100644 --- a/python/freetoken/checkpoint/convert.py +++ b/python/freetoken/checkpoint/convert.py @@ -200,7 +200,8 @@ def convert_checkpoint( cfg = EngineConfig(model_path=model_path, tp_info=DistributedInfo(tp.rank, tp.size), dtype=dtype, moe_strategy=moe_backend, quant_backend=quant_backend) mc = cfg.model_config - offload = moe_backend == "offload" and getattr(mc, "is_moe", False) + # every MoE strategy reads the routed experts from the banks, fused included + offload = getattr(mc, "is_moe", False) include_moe_experts = not offload method = None if offload: @@ -208,6 +209,7 @@ def convert_checkpoint( from freetoken.layers import set_rope_device set_rope_device(dev) + object.__setattr__(cfg, "moe_strategy", "offload") method = offload_expert_method(cfg) from freetoken.utils.progress import byte_bar, count_bar diff --git a/python/freetoken/checkpoint/ftw.py b/python/freetoken/checkpoint/ftw.py index e2bd0de24e..e31e9d8278 100644 --- a/python/freetoken/checkpoint/ftw.py +++ b/python/freetoken/checkpoint/ftw.py @@ -432,6 +432,43 @@ def _producer(): raise err[0] +def _split_bank_entries(reader: FTWReader, path: str, num_layers: int): + """``(all, alphas, flat regions, {base: {layer: entry}})`` of the ``experts_bank`` entries; None when there are none.""" + bank_entries = reader.entries("experts_bank") + if not bank_entries: + return None + + alpha_entries = [e for e in bank_entries if e["name"] in _ALPHA_NAMES] + row_entries = [e for e in bank_entries if e["name"] not in _ALPHA_NAMES] + + meta_layers = reader.meta("expert_bank_num_layers") + if meta_layers is not None and meta_layers != num_layers: + raise RuntimeError( + f"{path!r} was converted with {meta_layers} expert-bank layers but the " + f"model config says num_moe_layers={num_layers}; the checkpoint does not " + "match its config" + ) + + # Split row entries into the two layouts by name. + flat_entries: list[dict] = [] + per_layer_groups: dict[str, dict[int, dict]] = {} + for e in row_entries: + m = _LAYER_ENTRY_RE.match(e["name"]) + if m is None: + flat_entries.append(e) + continue + per_layer_groups.setdefault(m.group("base"), {})[int(m.group("layer"))] = e + + mixed = {e["name"] for e in flat_entries} & per_layer_groups.keys() + assert not mixed, f"FTW bank(s) mix flat and per-layer row layouts: {sorted(mixed)}" + for base, by_layer in per_layer_groups.items(): + assert sorted(by_layer) == list(range(num_layers)), ( + f"FTW bank {base!r} has per-layer entries for layers {sorted(by_layer)}, " + f"expected exactly range({num_layers})" + ) + return bank_entries, alpha_entries, flat_entries, per_layer_groups + + def load_ftw_banks( path: str, *, num_layers: int, workers: int = 8, chunk: int = _DEFAULT_CHUNK, layer_residency: list[str] | None = None, @@ -482,40 +519,20 @@ def _backing(layer_id: int) -> str: return "mmap" reader = FTWReader(path) - bank_entries = reader.entries("experts_bank") - if not bank_entries: + try: + split = _split_bank_entries(reader, path, num_layers) + except BaseException: reader.close() - return None - - alpha_entries = [e for e in bank_entries if e["name"] in _ALPHA_NAMES] - row_entries = [e for e in bank_entries if e["name"] not in _ALPHA_NAMES] - - meta_layers = reader.meta("expert_bank_num_layers") - if meta_layers is not None and meta_layers != num_layers: + raise + if split is None: reader.close() - raise RuntimeError( - f"{path!r} was converted with {meta_layers} expert-bank layers but the " - f"model config says num_moe_layers={num_layers}; the checkpoint does not " - "match its config" - ) + return None + bank_entries, alpha_entries, flat_entries, per_layer_groups = split # Alphas: unchanged, one flat HostBank per entry. alpha_specs = {e["name"]: (tuple(e["shape"]), _dtype_of(e["dtype"])) for e in alpha_entries} alpha_hb = alloc_banks(alpha_specs) - # Split row entries into the two layouts by name. - flat_entries: list[dict] = [] - per_layer_groups: dict[str, dict[int, dict]] = {} - for e in row_entries: - m = _LAYER_ENTRY_RE.match(e["name"]) - if m is None: - flat_entries.append(e) - continue - per_layer_groups.setdefault(m.group("base"), {})[int(m.group("layer"))] = e - - mixed = {e["name"] for e in flat_entries} & per_layer_groups.keys() - assert not mixed, f"FTW bank(s) mix flat and per-layer row layouts: {sorted(mixed)}" - # Row banks: one padded-window HostBank per (name, layer_id) for the flat layout, plus # how to carve the real [num_experts, *row_shape] tensor out of its head; ``None`` marks # a per-layer entry (direct view, no carving needed). @@ -546,10 +563,6 @@ def _backing(layer_id: int) -> str: row_jobs.append((name, bank, win_off, win_end - win_off, layer_bytes, layer_id)) for base, by_layer in per_layer_groups.items(): - assert sorted(by_layer) == list(range(num_layers)), ( - f"FTW bank {base!r} has per-layer entries for layers {sorted(by_layer)}, " - f"expected exactly range({num_layers})" - ) row_hb[base] = [] row_view_args[base] = [] for layer_id in range(num_layers): @@ -659,8 +672,71 @@ def _read_layer(job): ) +def load_ftw_banks_to_device(path: str, *, num_layers: int, device: torch.device, workers: int = 8, chunk: int = _DEFAULT_CHUNK): + """:func:`load_ftw_banks` for resident experts: every bank lands on ``device``, one (bank, layer) at a time through a transient host buffer.""" + from freetoken.moe.expert_banks import ExpertBanks + from freetoken.moe.legacy_format import canonical_role, kind_kernel_for + from freetoken.utils.progress import byte_bar + + reader = FTWReader(path) + try: + split = _split_bank_entries(reader, path, num_layers) + if split is None: + return None + bank_entries, alpha_entries, flat_entries, per_layer_groups = split + + # (name, layer_id or None for an alpha, read offset, read length, head pad, shape, dtype) + jobs = [(e["name"], None, e["global_off"], e["nbytes"], 0, tuple(e["shape"]), _dtype_of(e["dtype"])) for e in alpha_entries] + for e in flat_entries: + total, *row_shape = e["shape"] + assert total % num_layers == 0, (e["name"], total, num_layers) + layer_bytes = total // num_layers * math.prod(row_shape) * _elsize(_dtype_of(e["dtype"])) + assert layer_bytes * num_layers == e["nbytes"], (e["name"], layer_bytes, num_layers, e["nbytes"]) + for layer_id in range(num_layers): + off = e["global_off"] + layer_id * layer_bytes + win_off = (off // ALIGN) * ALIGN + jobs.append((e["name"], layer_id, win_off, _align_up(off + layer_bytes) - win_off, off - win_off, (total // num_layers, *row_shape), _dtype_of(e["dtype"]))) + for base, by_layer in per_layer_groups.items(): + for layer_id in range(num_layers): + e = by_layer[layer_id] + jobs.append((base, layer_id, e["global_off"], e["nbytes"], 0, tuple(e["shape"]), _dtype_of(e["dtype"]))) + + def _read(job) -> torch.Tensor: + _name, _layer, off, nbytes, head_pad, shape, dtype = job + buf = _transient_buffer(nbytes) + reader.read_into(memoryview(buf), {"global_off": off, "nbytes": nbytes}, workers=workers, chunk=chunk) + count = math.prod(shape) * _elsize(dtype) + return torch.frombuffer(buf, dtype=torch.uint8, count=head_pad + count)[head_pad:].view(dtype).view(shape) + + sources: dict[str, list] = {} + alphas: dict[str, torch.Tensor] = {} + bar = byte_bar(sum(e["nbytes"] for e in bank_entries), "Loading expert banks (FTW)") + try: + # read the next (bank, layer) while the current one copies to the device + with ThreadPoolExecutor(1) as ex: + ahead = ex.submit(_read, jobs[0]) + for k, (name, layer_id, *_rest) in enumerate(jobs): + host = ahead.result() + if k + 1 < len(jobs): + ahead = ex.submit(_read, jobs[k + 1]) + if layer_id is None: + alphas[name] = host.to(device) + else: + sources.setdefault(canonical_role(name), [None] * num_layers)[layer_id] = host.to(device) + bar.update(host.numel() * host.element_size()) + del host + finally: + bar.close() + quant_format = reader.meta("quant_format") + finally: + reader.close() + + kind, kernel = kind_kernel_for(quant_format) if quant_format is not None else (None, None) + return ExpertBanks(quant_format, sources, **alphas, kind=kind, kernel=kernel) + + __all__ = [ "INDEX_NAME", "FORMAT_TAG", "FORMAT_VERSION", "ALIGN", "DEFAULT_SHARD_LIMIT", "is_ftw_checkpoint", "ftw_tensor_names", "FTWWriter", "FTWReader", - "iter_ftw_weights", "load_ftw_banks", "layer_bank_entry_name", + "iter_ftw_weights", "load_ftw_banks", "load_ftw_banks_to_device", "layer_bank_entry_name", ] diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index e855bcd32f..e5ccf71050 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -13,14 +13,15 @@ from freetoken.core import Batch, Context, Req, set_global_ctx from freetoken.distributed import destroy_distributed, enable_pynccl_distributed, set_tp_info from freetoken.gpu_select import gpu_identity -from freetoken.layers import set_rope_device +from freetoken.checkpoint.ftw import ftw_tensor_names, is_ftw_checkpoint +from freetoken.layers import OffloadMoELayer, iter_moe_layers, set_rope_device from freetoken.layers.quantization import LayerKind, QuantBackend, finalize_quant, set_quant_backend from freetoken.moe.offload_cache import iter_offload_moe_layers from freetoken.mm.config import ENCODER_SECTIONS from freetoken.models import create_model, load_weight from freetoken.models.weight import ftw_lacks_vision from freetoken.moe import is_offload_moe_strategy -from freetoken.moe.expert_banks import load_expert_banks +from freetoken.moe.expert_banks import attach_resident_banks, load_expert_banks from freetoken.moe.host_banks import PinFailed from freetoken.moe.offload_cache import OffloadMoeCache, attach_offload_moe_cache from freetoken.utils import align_ceil, init_logger, is_sm90_family, is_sm100_family, mem_GB, torch_dtype @@ -531,9 +532,14 @@ def _load_weights(self, config: EngineConfig) -> None: "family served images. Reconvert it with `ft checkpoint`, add the encoder in place with " "scripts/ftw_hotfix.py (docs/ftw-hotfix.md), or start with --text-model-only" ) + resident = [layer for layer in iter_moe_layers(self.model) if not isinstance(layer, OffloadMoELayer)] + if resident and not config.use_dummy_weight and is_ftw_checkpoint(config.model_path) and not ftw_tensor_names(config.model_path, "experts_bank"): + raise ValueError(f"{config.model_path} keeps its routed experts as dense weights (an older `ft checkpoint --moe-backend fused`); reconvert it with `ft checkpoint`") with _weight_load_context(): self.model.load_state_dict(self._load_weight_state_dict(config)) finalize_quant(self.model) + if resident: + self._load_resident_experts(config, resident) def _load_weight_state_dict(self, config: EngineConfig) -> Dict[str, torch.Tensor]: model_state = self.model.state_dict() @@ -541,18 +547,32 @@ def _load_weight_state_dict(self, config: EngineConfig) -> Dict[str, torch.Tenso return _make_dummy_weight_state_dict(model_state, device=self.device) # _materialize casts each loaded tensor to its model-param dtype (model_state), so # models declaring per-tensor dtypes (e.g. DSV4's mixed fp8/fp32/bf16) are preserved; - # offload models exclude experts (served from the offload cache, not dense weights). + # MoE models exclude routed experts (they load as expert banks, not dense weights). return _materialize_loaded_weight_state_dict( model_state, load_weight( config.model_path, self.device, - include_moe_experts=not is_offload_moe_strategy(config.moe_strategy), + include_moe_experts=not getattr(config.model_config, "is_moe", False), include_vision=bool(config.active_encoders), ), device=self.device, ) + def _load_resident_experts(self, config: EngineConfig, layers) -> None: + with _weight_load_context(): + banks = load_expert_banks( + config.model_path, + config.model_config, + method=layers[0].quant_method, + device=self.device, + dtype=self.dtype, + dummy=config.use_dummy_weight, + parallel={"serial": False, "parallel": True}.get(config.expert_load, None), + resident=True, + ) + attach_resident_banks(layers, banks) + @torch.inference_mode() def _warmup_encoders(self) -> None: for item in self.mm_processor.dummy_items(self.dtype, self.device): @@ -1135,15 +1155,14 @@ def _is_unified_memory_gpu(index: "int | None" = None) -> bool: return False +_RESIDENT_EXPERT_QUANTS = ("none", "fp8_block", "nvfp4") + + def _fused_resident_ok(model_config) -> bool: - """Whether the resident ('fused') MoE path can hold this model's experts. - - FIXME: auto resolves to fused only for bf16 and fp8_block experts; drop this gate once the other quant formats support fused. - """ - expert_quant = getattr(model_config, "expert_quant", "none") - if expert_quant not in ("none", "fp8_block"): + """Whether the resident ('fused') MoE path can hold this model's experts; GGUF q4_0 and DeepSeek-V4 experts stay offload.""" + if getattr(model_config, "expert_quant", "none") not in _RESIDENT_EXPERT_QUANTS: return False - return getattr(model_config, "moe_weight_format", None) in (None, "bf16") + return getattr(model_config, "moe_weight_format", None) in (None, "bf16", "mxfp4") def _ensure_expandable_segments() -> None: @@ -1594,8 +1613,8 @@ def override(attr: str, value: Any): # this is dangerous, use with caution if is_moe and config.moe_strategy == "auto": # A MoE model always defaults to the offload family: experts stream from pinned host # banks into an auto-sized GPU slot cache, which is the only default that serves a model - # bigger than the GPU. The resident 'fused' path (bf16 / block-fp8 experts, the two - # formats MoELayer can allocate) is still reachable, but only when asked for explicitly + # bigger than the GPU. The resident 'fused' path (every expert format but GGUF q4_0 and + # DeepSeek-V4's) is still reachable, but only when asked for explicitly # -- auto never picks it, because nothing here knows whether the experts would fit in # HBM and a wrong guess is a weight-load OOM rather than a slower-but-working run. default_backend = "offload" @@ -1614,6 +1633,18 @@ def override(attr: str, value: Any): # this is dangerous, use with caution "Unified-memory GPU detected; auto-selecting 'fused' MoE strategy " "(resident experts) instead of offload" ) + # the cache sizing flags get their own warning in the 'fused' block below + inert = [flag for flag, on in ( + (f"--moe-cpu-layers={config.moe_cpu_layers}", config.moe_cpu_layers), + ("--moe-prefill-hit-d2d", config.moe_prefill_hit_d2d), + ("--disable-moe-prefill-overlap", not config.moe_prefill_overlap), + ) if on] + if inert: + logger.warning_rank0( + f"Unified-memory GPU serves MoE experts resident; ignoring {', '.join(inert)} " + "(use --moe-strategy offload to keep the offload flags)" + ) + override("moe_cpu_layers", None) # Hardware-adaptive config: a cached `ft bench bw` profile can upgrade # the offload default to hybrid when this machine's CPU MoE bandwidth clears its PCIe # gather bandwidth by the bench threshold (default 2x). hybrid is VRAM-equivalent to @@ -1733,7 +1764,7 @@ def override(attr: str, value: Any): # this is dangerous, use with caution if ( is_moe - and expert_quant not in ("none", "fp8_block") + and expert_quant not in _RESIDENT_EXPERT_QUANTS and not is_offload_moe_strategy(config.moe_strategy) ): raise ValueError( diff --git a/python/freetoken/layers/__init__.py b/python/freetoken/layers/__init__.py index d868751a6b..a07484620e 100644 --- a/python/freetoken/layers/__init__.py +++ b/python/freetoken/layers/__init__.py @@ -16,7 +16,7 @@ LinearReplicated, LinearRowParallel, ) -from .moe import MoELayer, OffloadMoELayer, make_moe_layer +from .moe import MoELayer, OffloadMoELayer, iter_moe_layers, make_moe_layer from .norm import ( GatedRMSNorm, GemmaPlusOneRMSNorm, @@ -57,5 +57,6 @@ "LinearReplicated", "MoELayer", "OffloadMoELayer", + "iter_moe_layers", "make_moe_layer", ] diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index c209454996..996beaa796 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -1,5 +1,5 @@ import os -from typing import TYPE_CHECKING, Tuple +from typing import TYPE_CHECKING, Iterator, Tuple import torch from freetoken.core import get_global_ctx @@ -435,6 +435,19 @@ def _expert_gemm( raise AssertionError(f"offload experts without a quant method only serve q4_0 banks, got {fmt!r}") +def iter_moe_layers(model) -> Iterator[MoELayer]: + if isinstance(model, MoELayer): + yield model + if not isinstance(model, BaseOP): + return + for value in model.__dict__.values(): + if isinstance(value, BaseOP): + yield from iter_moe_layers(value) + elif isinstance(value, (list, tuple)): + for item in value: + yield from iter_moe_layers(item) + + def make_moe_layer( config: "ModelConfig", *, @@ -483,10 +496,10 @@ def make_moe_layer( has_bias=has_bias, quant_config=quant_config, prefix=prefix, + layer_id=layer_id, ) if offload: assert layer_id is not None, "offload MoE backends need the layer_id" - kwargs["layer_id"] = layer_id kwargs["strategy"] = config.moe_strategy kwargs["decode_target"] = config.decode_target return layer_cls(**kwargs) diff --git a/python/freetoken/layers/quantization/moe/base.py b/python/freetoken/layers/quantization/moe/base.py index 89f13ed09d..e4a428d3e4 100644 --- a/python/freetoken/layers/quantization/moe/base.py +++ b/python/freetoken/layers/quantization/moe/base.py @@ -32,6 +32,7 @@ class MoEConfig: apply_router_weight_on_input: bool = False strategy: str = "resident" decode_target: str = "gpu" + dtype: torch.dtype = torch.bfloat16 @classmethod def from_layer(cls, layer: Any, scheme: QuantScheme | None) -> "MoEConfig": @@ -52,6 +53,7 @@ def from_layer(cls, layer: Any, scheme: QuantScheme | None) -> "MoEConfig": apply_router_weight_on_input=bool(layer.apply_router_weight_on_input), strategy=layer.strategy, decode_target=layer.decode_target, + dtype=torch.get_default_dtype(), ) @property @@ -147,9 +149,7 @@ def slot_limit(self, cfg: MoEConfig) -> int | None: """Most GPU cache slots the kernel can address for ``cfg``; None for no limit.""" return self.max_slots - def _common_reject(self, cfg: MoEConfig, *, resident_ok: bool, tp_ok: bool, cpu_ok: bool, plain_silu_only: bool) -> str | None: - if not resident_ok and cfg.strategy == "resident": - return "not served resident; use --moe-strategy offload or cpu" + def _common_reject(self, cfg: MoEConfig, *, tp_ok: bool, cpu_ok: bool, plain_silu_only: bool) -> str | None: if not tp_ok and cfg.tp_size > 1: return "TP > 1 is not supported for this expert format" if not cpu_ok and cfg.decode_target != "gpu": @@ -190,12 +190,15 @@ def apply(self, x: torch.Tensor, topk_weights: torch.Tensor, topk_ids: torch.Ten def cpu_format(self) -> str | None: return self.kernel.cpu_format - @abstractmethod - def create_weights(self, layer: Any) -> None: ... + def create_weights(self, layer: Any) -> None: + # resident experts stay out of the state dict; attach_resident_banks fills these after the dense weights load + layer._expert_banks = None + layer._expert_alphas = None def finalize(self, layer: Any) -> None: pass - @abstractmethod def resident_view(self, layer: Any) -> ExpertView: """Layout roles -> the resident layer's tensors, for apply().""" + assert layer._expert_banks is not None, "resident expert banks were never loaded" + return ExpertView(layer._expert_banks, n=self.cfg.num_experts, alphas=layer._expert_alphas) diff --git a/python/freetoken/layers/quantization/moe/fp8_block.py b/python/freetoken/layers/quantization/moe/fp8_block.py index 6c0167b71b..20c3dde9dc 100644 --- a/python/freetoken/layers/quantization/moe/fp8_block.py +++ b/python/freetoken/layers/quantization/moe/fp8_block.py @@ -21,7 +21,7 @@ class TritonFp8BlockMoEKernel(MoEKernel): name = "triton" def unusable_reason(self, cfg: MoEConfig) -> str | None: - reason = self._common_reject(cfg, resident_ok=True, tp_ok=False, cpu_ok=False, plain_silu_only=False) + reason = self._common_reject(cfg, tp_ok=False, cpu_ok=False, plain_silu_only=False) if reason: return reason reason = gated_epilogue_reason(cfg) @@ -64,17 +64,3 @@ def apply(self, layer, x, topk_weights, topk_ids, view: ExpertView, *, is_prefil @register_method(QuantKind.FP8_BLOCK, LayerKind.MOE) class Fp8BlockMoEMethod(MoEMethod): candidates = (TritonFp8BlockMoEKernel,) - - def create_weights(self, layer) -> None: - g = self.cfg - e, i, h, b = g.num_experts, g.intermediate, g.hidden, BLOCK - layer.gate_up_proj = torch.empty(e, 2 * i, h, dtype=FP8) - layer.gate_up_scale_inv = torch.empty(e, 2 * i // b, h // b, dtype=torch.bfloat16) - layer.down_proj = torch.empty(e, h, i, dtype=FP8) - layer.down_scale_inv = torch.empty(e, h // b, i // b, dtype=torch.bfloat16) - - def resident_view(self, layer) -> ExpertView: - return ExpertView({ - "gate_up": layer.gate_up_proj, "gate_up_scale": layer.gate_up_scale_inv, - "down": layer.down_proj, "down_scale": layer.down_scale_inv, - }) diff --git a/python/freetoken/layers/quantization/moe/mxfp4.py b/python/freetoken/layers/quantization/moe/mxfp4.py index aa020ee33a..13cd97466c 100644 --- a/python/freetoken/layers/quantization/moe/mxfp4.py +++ b/python/freetoken/layers/quantization/moe/mxfp4.py @@ -22,7 +22,7 @@ def unusable_reason(self, cfg: MoEConfig) -> str | None: return "standard MXFP4 kernel reads the concatenated gate|up row order" if (cfg.alpha, cfg.beta) != (1.0, 0.0): return "standard MXFP4 kernel has no alpha / beta in its swiglu" - return self._common_reject(cfg, resident_ok=False, tp_ok=False, cpu_ok=True, plain_silu_only=False) + return self._common_reject(cfg, tp_ok=False, cpu_ok=True, plain_silu_only=False) def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]: i, h = cfg.intermediate, cfg.hidden @@ -106,37 +106,3 @@ def apply(self, layer, x, topk_weights, topk_ids, view: ExpertView, *, is_prefil @register_method(QuantKind.MXFP4, LayerKind.MOE) class Mxfp4MoEMethod(MoEMethod): candidates = (TritonMxfp4MoEKernel, TritonGptossMxfp4MoEKernel) - - def create_weights(self, layer) -> None: - if not self.cfg.has_bias: - raise NotImplementedError("standard MXFP4 experts are served from the offload cache, not resident") - g = self.cfg - e, i, h = g.num_experts, g.local_intermediate, g.hidden - if h % GROUP: - raise ValueError(f"MXFP4 hidden size must be divisible by {GROUP}") - layer.gate_up_proj_blocks = torch.empty(e, 2 * i, h // GROUP, 16, dtype=torch.uint8) - layer.gate_up_proj_scales = torch.empty(e, 2 * i, h // GROUP, dtype=torch.uint8) - layer.gate_up_proj_bias = torch.empty(e, 2 * i, dtype=torch.bfloat16) - layer.down_proj_blocks = torch.empty(e, h, i // GROUP, 16, dtype=torch.uint8) - layer.down_proj_scales = torch.empty(e, h, i // GROUP, dtype=torch.uint8) - layer.down_proj_bias = torch.empty(e, h, dtype=torch.bfloat16) - - def finalize(self, layer) -> None: - if getattr(layer, "gate_up_proj_blocks", None) is None: - return - from freetoken.moe.fused_mxfp4 import _transpose_mxfp4_for_decode - - # one transposed copy serves prefill and decode; the HF blocks are freed so 120B fits - layer._gu_blocks_t, layer._gu_scales_t = _transpose_mxfp4_for_decode(layer.gate_up_proj_blocks, layer.gate_up_proj_scales) - layer._dn_blocks_t, layer._dn_scales_t = _transpose_mxfp4_for_decode(layer.down_proj_blocks, layer.down_proj_scales) - layer.gate_up_proj_blocks = None - layer.gate_up_proj_scales = None - layer.down_proj_blocks = None - layer.down_proj_scales = None - torch.cuda.empty_cache() - - def resident_view(self, layer) -> ExpertView: - return ExpertView({ - "gate_up": layer._gu_blocks_t, "gate_up_scale": layer._gu_scales_t, "gate_up_bias": layer.gate_up_proj_bias, - "down": layer._dn_blocks_t, "down_scale": layer._dn_scales_t, "down_bias": layer.down_proj_bias, - }) diff --git a/python/freetoken/layers/quantization/moe/nvfp4.py b/python/freetoken/layers/quantization/moe/nvfp4.py index bb341cdfb7..cf66e2071b 100644 --- a/python/freetoken/layers/quantization/moe/nvfp4.py +++ b/python/freetoken/layers/quantization/moe/nvfp4.py @@ -1,4 +1,4 @@ -"""NVFP4 experts: served from the offload cache through the Triton inline-dequant kernels, vLLM's Marlin or flashinfer's b12x. +"""NVFP4 experts: served resident or from the offload cache through the Triton inline-dequant kernels, vLLM's Marlin or flashinfer's b12x. FreeToken stores MoE experts as ModelOpt NVFP4 (packed e2m1 codes + fp8-e4m3 per-16 block scales + per-tensor global scale) in pinned host banks and gathers the routed @@ -41,7 +41,7 @@ class TritonNvfp4MoEKernel(MoEKernel): cpu_format = "nvfp4" def unusable_reason(self, cfg: MoEConfig) -> str | None: - reason = self._common_reject(cfg, resident_ok=False, tp_ok=False, cpu_ok=True, plain_silu_only=False) + reason = self._common_reject(cfg, tp_ok=False, cpu_ok=True, plain_silu_only=False) if reason: return reason reason = gated_epilogue_reason(cfg) @@ -230,7 +230,7 @@ class MarlinNvfp4MoEKernel(MoEKernel): def unusable_reason(self, cfg: MoEConfig) -> str | None: if not backend.is_vllm_installed(): return "vLLM is not installed" - reason = self._common_reject(cfg, resident_ok=False, tp_ok=False, cpu_ok=False, plain_silu_only=True) + reason = self._common_reject(cfg, tp_ok=False, cpu_ok=False, plain_silu_only=True) if reason: return reason if not _marlin_symbols_ok(): @@ -496,7 +496,7 @@ def unusable_reason(self, cfg: MoEConfig) -> str | None: return f"b12x requires sm_120+, got sm_{cc[0]}{cc[1]}" if not backend.is_flashinfer_installed(): return "flashinfer is not installed" - reason = self._common_reject(cfg, resident_ok=False, tp_ok=False, cpu_ok=False, plain_silu_only=True) + reason = self._common_reject(cfg, tp_ok=False, cpu_ok=False, plain_silu_only=True) if reason: return reason return _b12x_unusable_reason(cc) @@ -567,9 +567,3 @@ def apply(self, layer, x, topk_weights, topk_ids, view: ExpertView, *, is_prefil @register_method(QuantKind.NVFP4, LayerKind.MOE) class Nvfp4MoEMethod(MoEMethod): candidates = (TritonNvfp4MoEKernel, MarlinNvfp4MoEKernel, B12xNvfp4MoEKernel) - - def create_weights(self, layer) -> None: - raise NotImplementedError("NVFP4 experts are served from the offload cache, not resident") - - def resident_view(self, layer) -> ExpertView: - raise NotImplementedError("NVFP4 experts are not resident") diff --git a/python/freetoken/layers/quantization/moe/unquantized.py b/python/freetoken/layers/quantization/moe/unquantized.py index ef435686c5..ded7d8135c 100644 --- a/python/freetoken/layers/quantization/moe/unquantized.py +++ b/python/freetoken/layers/quantization/moe/unquantized.py @@ -19,7 +19,9 @@ def unusable_reason(self, cfg: MoEConfig) -> str | None: def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]: i = cfg.local_intermediate - return {"gate_up": BankSpec((2 * i, cfg.hidden), torch.bfloat16), "down": BankSpec((cfg.hidden, i), torch.bfloat16)} + # resident experts compute in the model dtype; offload banks stay bf16, the FTW and CPU executor format + dt = cfg.dtype if cfg.strategy == "resident" else torch.bfloat16 + return {"gate_up": BankSpec((2 * i, cfg.hidden), dt), "down": BankSpec((cfg.hidden, i), dt)} def pack(self, pieces, cfg: MoEConfig, out): out["gate_up"].copy_(fused_piece(pieces, "gate_up")) @@ -37,12 +39,3 @@ def apply(self, layer, x, topk_weights, topk_ids, view: ExpertView, *, is_prefil @register_method(QuantKind.NONE, LayerKind.MOE) class UnquantizedMoEMethod(MoEMethod): candidates = (FusedMoEKernel,) - - def create_weights(self, layer) -> None: - g = self.cfg - i = g.local_intermediate - layer.gate_up_proj = torch.empty(g.num_experts, 2 * i, g.hidden) - layer.down_proj = torch.empty(g.num_experts, g.hidden, i) - - def resident_view(self, layer) -> ExpertView: - return ExpertView({"gate_up": layer.gate_up_proj, "down": layer.down_proj}) diff --git a/python/freetoken/models/qwen3_5_moe/weight.py b/python/freetoken/models/qwen3_5_moe/weight.py index 1bac8227b0..de5b9e8983 100644 --- a/python/freetoken/models/qwen3_5_moe/weight.py +++ b/python/freetoken/models/qwen3_5_moe/weight.py @@ -1,6 +1,6 @@ """Qwen3.5 / 3.6 / 3.8 checkpoint reader. -The dense pass reads every Linear module under the scheme the checkpoint's QuantConfig gives it, the same answer the model built its buffers from, so bf16, block-fp8, ModelOpt and llm-compressor exports in any mix all land as the model's state dict. Routed experts are read by the offload cache (``nvfp4_expert_spec`` / ``iter_expert_pieces``); only bf16 stacked experts and resident block-fp8 experts come from here. +The dense pass reads every Linear module under the scheme the checkpoint's QuantConfig gives it, the same answer the model built its buffers from, so bf16, block-fp8, ModelOpt and llm-compressor exports in any mix all land as the model's state dict. Routed experts are read by the expert-bank loader (``nvfp4_expert_spec`` / ``iter_expert_pieces``); only the bf16 stacked experts it reads come from here. """ from __future__ import annotations @@ -241,9 +241,9 @@ def iter_weights( include_non_moe: bool, include_vision: bool = True, ) -> Iterator[tuple[str, torch.Tensor]]: - """Yield the dense weights fused to the model's buffers, and the routed experts only where a resident path takes them from here: bf16 stacked experts as stored, block-fp8 experts restacked per layer. + """Yield the dense weights fused to the model's buffers, and with ``include_moe_experts`` the bf16 stacked experts as stored. - Per-expert NVFP4 experts always come from the offload cache's expert reader. + Block-fp8 and NVFP4 experts always come from the expert-bank reader. """ if get_tp_info().size > 1: raise NotImplementedError("qwen3_5_moe weight loading supports TP=1 only") @@ -253,8 +253,6 @@ def iter_weights( if include_non_moe or stacked: reader = _DenseReader(get_quant_config(), get_model_spec(hf_config.architectures[0])) if include_non_moe else None yield from _iter_shards(model_path, device, reader, stacked=stacked, include_vision=include_vision) - if include_moe_experts and config.is_moe and config.expert_quant == "fp8_block": - yield from _resident_fp8_experts(model_path, config) def _iter_shards(model_path: str, device: torch.device, reader: _DenseReader | None, *, stacked: bool, include_vision: bool): @@ -327,7 +325,7 @@ def _is_expert(raw_name: str) -> bool: # ====================================================================================== -# Block-FP8 routed experts (Qwen3.5-35B-A3B-FP8): offload expert pieces and resident stacks. +# Block-FP8 routed experts (Qwen3.5-35B-A3B-FP8): expert pieces. # ====================================================================================== # Routed-expert checkpoint key (per-expert, un-fused). ``mtp.layers...`` is excluded by the @@ -338,35 +336,6 @@ def _is_expert(raw_name: str) -> bool: ) -def _resident_fp8_experts(model_path, config): - from freetoken.kernel.triton.fp8_block_linear import FP8 - - B = 128 - L, E, H, I, dense = _moe_dims(config) - shapes = { - "gate_up_proj": ((E, 2 * I, H), FP8), - "gate_up_scale_inv": ((E, 2 * I // B, H // B), torch.bfloat16), - "down_proj": ((E, H, I), FP8), - "down_scale_inv": ((E, H // B, I // B), torch.bfloat16), - } - layers: dict[int, dict[str, torch.Tensor]] = {} - placed = [0] * L - for li, e0, e1, piece in iter_expert_pieces(model_path, config, QuantKind.FP8_BLOCK, parallel=None): - stack = layers.setdefault(li, {n: torch.empty(shape, dtype=dt) for n, (shape, dt) in shapes.items()}) - stack["gate_up_proj"][e0:e1, :I] = piece["gate"] - stack["gate_up_proj"][e0:e1, I:] = piece["up"] - stack["gate_up_scale_inv"][e0:e1, : I // B] = piece["gate_scale"] - stack["gate_up_scale_inv"][e0:e1, I // B :] = piece["up_scale"] - stack["down_proj"][e0:e1] = piece["down"] - stack["down_scale_inv"][e0:e1] = piece["down_scale"] - placed[li] += e1 - e0 - if placed[li] == E: - pre = f"model.layers.{dense + li}.mlp.experts" - for name, tensor in layers.pop(li).items(): - yield f"{pre}.{name}", tensor - assert not layers, f"incomplete resident fp8 experts for layers {sorted(layers)}" - - def _moe_dims(model_config): L = model_config.num_moe_layers return ( diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py index e3a768b6d0..7500174ae4 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/python/freetoken/moe/expert_banks.py @@ -75,6 +75,7 @@ def build_expert_banks( device: torch.device, layer_sink=None, dummy: bool = False, + resident: bool = False, ) -> ExpertBanks: """Fill host banks in the kernel's layout from a stream of expert pieces. @@ -83,6 +84,7 @@ def build_expert_banks( complete once its ``num_experts`` rows have arrived: with ``layer_sink=None`` its banks are pinned in the background, otherwise the sink receives them (converter). ``dummy`` skips the pieces and fills the banks with finite random contents. + ``resident`` packs into banks on ``device`` instead of host banks (the resident 'fused' experts). """ from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline, pin_banks from freetoken.moe.legacy_format import legacy_format_for @@ -91,8 +93,13 @@ def build_expert_banks( layout = method.layout() E = method.cfg.num_experts specs = {role: ((E, *spec.shape), spec.dtype) for role, spec in layout.items() if not spec.resident} - hb = alloc_layer_banks(specs, num_layers) - banks = {role: [b.tensor for b in hb[role]] for role in specs} + if resident: + assert layer_sink is None + hb = None + banks = {role: [torch.empty(shape, dtype=dtype, device=device) for _ in range(num_layers)] for role, (shape, dtype) in specs.items()} + else: + hb = alloc_layer_banks(specs, num_layers) + banks = {role: [b.tensor for b in hb[role]] for role in specs} alphas = { role: torch.empty(num_layers * E, dtype=spec.dtype, device=device) for role, spec in layout.items() if spec.resident @@ -104,7 +111,7 @@ def build_expert_banks( _dummy_fill(role, tensor) for alpha in alphas.values(): alpha.fill_(1.0) - if torch.cuda.is_available(): + if hb is not None and torch.cuda.is_available(): pin_banks(hb) return ExpertBanks( legacy_format_for(method.kind, kernel.name), banks, @@ -136,7 +143,7 @@ def _fill(sink) -> None: if layer_sink is not None: _fill(layer_sink) - elif torch.cuda.is_available(): + elif torch.cuda.is_available() and not resident: with PinPipeline() as pins: _fill(pins) else: @@ -192,16 +199,16 @@ def _legacy_expert_banks(model_path, model_config, device, dtype, dummy, paralle ) -def _method_expert_banks(model_path, model_config, method, device, dummy, parallel, workers, chunk, layer_sink=None) -> ExpertBanks: +def _method_expert_banks(model_path, model_config, method, device, dummy, parallel, workers, chunk, layer_sink=None, resident=False) -> ExpertBanks: from freetoken.moe.expert_pieces import iter_expert_pieces num_layers = model_config.num_moe_layers if dummy: - return build_expert_banks(method, num_layers, None, device=device, dummy=True) + return build_expert_banks(method, num_layers, None, device=device, dummy=True, resident=resident) pieces = iter_expert_pieces( model_path, model_config, method.kind, parallel=parallel, workers=workers, chunk=chunk ) - return build_expert_banks(method, num_layers, pieces, device=device, layer_sink=layer_sink) + return build_expert_banks(method, num_layers, pieces, device=device, layer_sink=layer_sink, resident=resident) def _host_ram_fits_parallel(model_path: str) -> bool: @@ -286,6 +293,7 @@ def load_expert_banks( decode_target: str = "gpu", layer_sink=None, layer_residency: list[str] | None = None, + resident: bool = False, ) -> ExpertBanks: """Load (or fabricate, with ``dummy=True``) the expert banks. Two paths, both returning the same normalized ``ExpertBanks`` and both pinning after fill: @@ -310,14 +318,21 @@ def load_expert_banks( ``layer_residency``: per-layer ``HostResidency`` labels applied at settle time -- explicitly on the FTW fast path, ambiently (``requested_residency``) in the slow-path providers. Applied labels are echoed on ``ExpertBanks.layer_residency``; a loader that settles some other way leaves it ``None`` (CPU-layer decode still works on pinned banks, it just saves no pin quota). + + ``resident`` (with ``method``): the banks land on ``device`` for the resident 'fused' experts; no host bank outlives the load. """ - from freetoken.checkpoint.ftw import is_ftw_checkpoint, load_ftw_banks + from freetoken.checkpoint.ftw import is_ftw_checkpoint, load_ftw_banks, load_ftw_banks_to_device + if resident: + assert method is not None and layer_sink is None and layer_residency is None if model_path and is_ftw_checkpoint(model_path) and not dummy: - banks = load_ftw_banks( - model_path, num_layers=model_config.num_moe_layers, workers=workers, chunk=chunk, - layer_residency=layer_residency, - ) + if resident: + banks = load_ftw_banks_to_device(model_path, num_layers=model_config.num_moe_layers, device=device, workers=workers, chunk=chunk) + else: + banks = load_ftw_banks( + model_path, num_layers=model_config.num_moe_layers, workers=workers, chunk=chunk, + layer_residency=layer_residency, + ) if banks is not None: logger.info_rank0(f"expert banks: FTW fast path (FTW checkpoint {model_path})") return banks @@ -354,7 +369,7 @@ def load_expert_banks( def _build(par: bool) -> ExpertBanks: if method is not None: - return _method_expert_banks(model_path, model_config, method, device, dummy, par, workers, chunk, layer_sink) + return _method_expert_banks(model_path, model_config, method, device, dummy, par, workers, chunk, layer_sink, resident) return _legacy_expert_banks(model_path, model_config, device, dtype, dummy, par, workers, chunk, decode_target, layer_sink) with requested_residency(layer_residency) as residency_plan: @@ -393,3 +408,33 @@ def _echo_residency(banks: ExpertBanks, requested, plan) -> ExpertBanks: "but saves no pinned quota" ) return banks + + +def attach_resident_banks(layers, banks: ExpertBanks) -> None: + """Give each resident MoE layer its rows of the device ``banks``; bank layer ``i`` goes to the layer whose ``layer_id`` is ``i``.""" + ids = sorted(layer.layer_id for layer in layers if layer.layer_id is not None) + num_layers = len(next(iter(banks.sources.values()))) + if ids != list(range(num_layers)) or len(ids) != len(layers): + raise ValueError(f"resident MoE layer ids {ids} do not cover the {num_layers} expert bank layers") + for layer in layers: + i, e = layer.layer_id, layer.num_experts + layout = layer.quant_method.layout() + rows = {role: (e, *spec.shape) for role, spec in layout.items() if not spec.resident} + if set(banks.sources) != set(rows): + raise ValueError(f"expert banks hold {sorted(banks.sources)}, the {layer.quant_method.kernel.name} kernel reads {sorted(rows)}") + for role, shape in rows.items(): + bank = banks.sources[role][i] + if tuple(bank.shape) != shape or bank.dtype != layout[role].dtype: + raise ValueError( + f"bank {role!r} of layer {i} is {tuple(bank.shape)} {bank.dtype} but the {layer.quant_method.kernel.name} kernel " + f"wants {shape} {layout[role].dtype}; the banks were packed for another kernel or model" + ) + alphas = [] + # apply() reads view.alphas as (gate_up, down) + for role in (role for role in ("gate_up_alpha", "down_alpha") if role in layout): + a = getattr(banks, role) + if a is None or a.numel() != num_layers * e or a.dtype != layout[role].dtype: + raise ValueError(f"the {layer.quant_method.kernel.name} kernel needs {num_layers * e} {layout[role].dtype} {role} values, the expert banks carry {None if a is None else (a.numel(), a.dtype)}") + alphas.append(a) + layer._expert_banks = {role: per_layer[i] for role, per_layer in banks.sources.items()} + layer._expert_alphas = tuple(a[i * e : (i + 1) * e] for a in alphas) or None diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 32e0f3cc1e..c6f975df7e 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -1054,20 +1054,9 @@ def copy_missing(self) -> None: def iter_offload_moe_layers(model) -> Iterator: - from freetoken.layers import BaseOP, OffloadMoELayer + from freetoken.layers import OffloadMoELayer, iter_moe_layers - if isinstance(model, OffloadMoELayer): - yield model - - if not isinstance(model, BaseOP): - return - - for value in model.__dict__.values(): - if isinstance(value, BaseOP): - yield from iter_offload_moe_layers(value) - elif isinstance(value, (list, tuple)): - for item in value: - yield from iter_offload_moe_layers(item) + return (layer for layer in iter_moe_layers(model) if isinstance(layer, OffloadMoELayer)) def attach_offload_moe_cache(model, cache: OffloadMoeCache) -> list: diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 23a62f8419..4dd57b1ae5 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -600,8 +600,8 @@ def _infer_reasoning_parser(model_path: str) -> str | None: choices=["auto", *MOE_STRATEGIES], help=( "How the routed experts are served. 'auto' resolves a MoE model to the offload family " - "(offload, or hybrid when a `ft bench bw` profile recommends it); resident " - "'fused' experts must be requested explicitly." + "(offload, or hybrid when a `ft bench bw` profile recommends it), and to resident " + "'fused' experts on unified-memory GPUs (GB10 / DGX Spark)." ), ) diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py index 073112bb5a..a77409f013 100644 --- a/tests/engine/test_cache_budget.py +++ b/tests/engine/test_cache_budget.py @@ -367,7 +367,7 @@ def _offload_engine_config(**overrides): num_layers=10, num_moe_layers=10, num_experts=8, - expert_quant="nvfp4", # quantized experts -> must resolve to an offload backend + expert_quant="nvfp4", moe_strategy="auto", ), ) @@ -389,7 +389,7 @@ def test_guard_raises_actionable_error_when_too_small(): assert "128" in msg and "moe-cache" in msg -def test_adjust_config_defaults_moe_cache_auto_for_auto_resolved_offload_backend(): +def test_adjust_config_defaults_moe_cache_auto_for_auto_resolved_offload_backend(monkeypatch): """Bare `ft serve `: no --moe-backend, no --moe-cache-* flags at all. args.py's parse-time default only fires when the backend is *already* @@ -400,9 +400,12 @@ def test_adjust_config_defaults_moe_cache_auto_for_auto_resolved_offload_backend moe_cache_auto=True, so _init_offload_moe_cache's _require_offload_cache_size guard is never reached with moe_cache_size still 0. """ + import freetoken.engine.engine as engine_module from freetoken.engine.engine import _adjust_config from freetoken.moe import is_offload_moe_strategy + # a unified-memory host (GB10) resolves auto to fused instead + monkeypatch.setattr(engine_module, "_is_unified_memory_gpu", lambda index=None: False) config = _offload_engine_config() _adjust_config(config) diff --git a/tests/models/qwen4_exp/test_skeleton.py b/tests/models/qwen4_exp/test_skeleton.py index ac598e65df..8fb9112df2 100644 --- a/tests/models/qwen4_exp/test_skeleton.py +++ b/tests/models/qwen4_exp/test_skeleton.py @@ -16,7 +16,7 @@ import torch import torch.nn.functional as F -from freetoken.layers import BaseOP, LinearReplicated +from freetoken.layers import BaseOP, LinearReplicated, iter_moe_layers from freetoken.models.config import ModelConfig from freetoken.models.qwen4_exp.config import parse_config from freetoken.models.qwen4_exp.hc import GatedResidual @@ -30,12 +30,17 @@ def _config(num_layers: int = 4) -> ModelConfig: def _fill(op, gen: torch.Generator, scale: float = 0.05) -> None: - """Random floats / zeroed ints for every state-dict tensor of an op tree.""" + """Random floats / zeroed ints for every state-dict tensor of an op tree, and random expert banks for its resident MoE layers.""" for tensor in op.state_dict().values(): if tensor.is_floating_point(): tensor.normal_(0.0, scale, generator=gen) else: tensor.zero_() + for layer in iter_moe_layers(op): + layer._expert_banks = { + role: torch.randn((layer.num_experts, *spec.shape), generator=gen, device=gen.device, dtype=spec.dtype) * scale + for role, spec in layer.quant_method.layout().items() + } def _group_norm(x, weight, eps, groups): diff --git a/tests/models/test_quant_config.py b/tests/models/test_quant_config.py index 6ac1b84b17..084a2e4aa7 100644 --- a/tests/models/test_quant_config.py +++ b/tests/models/test_quant_config.py @@ -121,14 +121,6 @@ class Case: check: Callable[[Any, dict], None] | None = None # family extras beyond the method table -def _installed(module: str) -> bool: - try: - __import__(module) - return True - except ImportError: - return False - - def _qwen35_fp8(model, idx): assert model.model.layers.op_list[idx["lin"]].linear_attn._split_in_proj @@ -162,13 +154,18 @@ def _dsv4(model, idx): assert "model.head.weight" in keys and not any(".ffn.experts." in k for k in keys) +def _banked_experts(model, idx): + experts = op_at(model, f"model.layers.{idx['lin']}.mlp.experts") + assert not experts.state_dict() + + def _gpt_oss(cls): def check(model, idx): experts = op_at(model, "model.layers.0.mlp.experts") assert type(experts) is cls assert experts.interleaved and experts.alpha == pytest.approx(1.702) and experts.limit == 7.0 if cls is GptOssMoELayer: - assert experts.gate_up_proj_blocks.dtype is torch.uint8 and experts.gate_up_proj_blocks.dim() == 4 + assert not experts.state_dict() return check @@ -198,9 +195,8 @@ def _gemma_tied_head(model, idx): "model.layers.{lin}.mlp.shared_expert.down_proj": NVFP4, "model.layers.{lin}.mlp.experts": Nvfp4MoEMethod, "lm_head": NVFP4, }, check=_qwen35_nvfp4), - # resident NVFP4 experts have no kernel without vLLM's Marlin - Case("Qwen3.6-35B-A3B-NVFP4", ModelOptConfig, {}, strategy="fused", raises=KernelSelectionError, - skip_if=lambda: not _installed("flashinfer") or _installed("vllm")), + # resident NVFP4 experts load through the expert-bank loader, not the state dict + Case("Qwen3.6-35B-A3B-NVFP4", ModelOptConfig, {"model.layers.{lin}.mlp.experts": (Nvfp4MoEMethod, "triton")}, strategy="fused", check=_banked_experts), Case("Qwen3.8-27B-NVFP4", ModelOptConfig, { "model.layers.{lin}.linear_attn.in_proj_qkvz": FP8T, "model.layers.{full}.self_attn.o_proj": FP8T, "model.layers.{lin}.mlp.gate_up_proj": NVFP4, "model.layers.{lin}.mlp.down_proj": NVFP4, "lm_head": NVFP4, From 581fdca310d429cca74ac9d94ed2b893b86d9320 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BA=91=E5=BE=AE?= <1067852565@qq.com> Date: Mon, 5 Oct 2026 09:24:33 +0200 Subject: [PATCH 02/73] fix(server): report the backend kv page size in stats (#599) * fix(server): report the backend kv page size in stats Use resolved backend cache metadata, retaining the frontend fallback. Refs FlashML-org/FreeToken#598 Assisted-by: Codex * test(server): cover resolved kv page size and fallback Refs FlashML-org/FreeToken#598 Assisted-by: Codex --- python/freetoken/server/stats.py | 3 ++- tests/server/test_generation_accounting.py | 13 +++++++++++++ 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/python/freetoken/server/stats.py b/python/freetoken/server/stats.py index 1857d3d7c0..d01828e21b 100644 --- a/python/freetoken/server/stats.py +++ b/python/freetoken/server/stats.py @@ -138,9 +138,10 @@ def build_stats(state: Any, p95_ms: int, ttft_mean_ms: int) -> dict: config = state.config ready_at = getattr(state, "ready_at", None) uptime_s = max(0, int(time.monotonic() - ready_at)) if ready_at is not None else 0 + pools = getattr(state, "cache_pools", None) or {} kv = ( {"used_pages": tr.kv_used_pages, "total_pages": tr.kv_total_pages, - "page_size": getattr(config, "page_size", 1)} + "page_size": pools.get("page_size") or getattr(config, "page_size", 1)} if tr.kv_total_pages > 0 else None ) mamba = ( diff --git a/tests/server/test_generation_accounting.py b/tests/server/test_generation_accounting.py index 8276163b0c..a6d2bdb916 100644 --- a/tests/server/test_generation_accounting.py +++ b/tests/server/test_generation_accounting.py @@ -91,6 +91,19 @@ def _last_row() -> dict: return rows[-1] +@pytest.mark.parametrize("pools, expected", [(None, 1), ({"page_size": 64}, 64)]) +def test_stats_reports_backend_page_size(pools, expected): + from freetoken.server.stats import StatsTracker, build_stats + + tracker = StatsTracker() + tracker.kv_total_pages = 130 + config = SimpleNamespace(served_model_name="unit-model", max_seq_len=262144, + page_size=1, served_modalities=(), model_config=SimpleNamespace()) + state = SimpleNamespace(stats=tracker, config=config, cache_pools=pools) + doc = build_stats(state, 0, 0) + assert doc["kv"] == {"used_pages": 0, "total_pages": 130, "page_size": expected} + + def test_non_stream_records_the_request_with_real_token_totals(): request_ring.reset() st = FakeState([_ack(prompt=5, completion=1, out="a"), _ack(completion=2, out="bc", finished=True)]) From 93b0a59164bb6d68749298f868ca3738fd464143 Mon Sep 17 00:00:00 2001 From: William Xinze Zheng Date: Mon, 5 Oct 2026 14:59:59 -0700 Subject: [PATCH 03/73] fix(server): keep Anthropic inline system messages in place for prefix reuse (#608) * Implement fix for Claude Code's inline system message * Fix for Muse --- docs/cli.md | 2 +- python/freetoken/message/tokenizer.py | 1 + python/freetoken/server/anthropic_api.py | 19 +- python/freetoken/server/args.py | 9 + python/freetoken/server/generation.py | 9 + python/freetoken/tokenizer/inline_system.py | 173 +++++++ python/freetoken/tokenizer/tokenize.py | 31 +- tests/README.md | 5 +- tests/server/test_anthropic_api.py | 103 ++++- tests/server/test_message_wire.py | 12 + tests/tokenizer/test_tokenize.py | 476 ++++++++++++++++++++ 11 files changed, 822 insertions(+), 18 deletions(-) create mode 100644 python/freetoken/tokenizer/inline_system.py diff --git a/docs/cli.md b/docs/cli.md index deeda18419..b3cec99826 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -99,6 +99,7 @@ See [models.md](models.md#moe-strategies) for what each strategy does. | `--tool-call-parser` | auto | Tool-call format; auto-inferred from the model family | | `--reasoning-parser` | auto | Splits chain-of-thought into `reasoning_content`; auto-inferred; `off` disables | | `--enable-cache-report` | off | Report prefix-cache hits in each response's usage block | +| `--anthropic-inline-system` | auto | Placement of late Anthropic system instructions: `auto`, `preserve`, or `fold` | ### Image input @@ -203,4 +204,3 @@ profile that `ft serve --moe-strategy auto` and `--moe-hybrid-max-fetch -1` then - What to measure: `--dtype`, `--model`, `--formats`, `--isa`. - `--threshold` (default 2.0) sets the call: recommend hybrid when CPU bandwidth beats PCIe by that factor. - diff --git a/python/freetoken/message/tokenizer.py b/python/freetoken/message/tokenizer.py index 9442ee059e..d52d658e63 100644 --- a/python/freetoken/message/tokenizer.py +++ b/python/freetoken/message/tokenizer.py @@ -73,6 +73,7 @@ class TokenizeMsg(BaseTokenizerMsg): chat_template_kwargs: Dict[str, Any] | None = None tools: List[Dict[str, Any]] | None = None images: List[bytes] | None = None + inline_system_policy: str | None = None @dataclass diff --git a/python/freetoken/server/anthropic_api.py b/python/freetoken/server/anthropic_api.py index 0d60940c97..d9a6972f09 100644 --- a/python/freetoken/server/anthropic_api.py +++ b/python/freetoken/server/anthropic_api.py @@ -120,6 +120,7 @@ async def handle_anthropic_messages( default_max_tokens=( getattr(state.config, "max_output_tokens", None) or DEFAULT_MAX_OUTPUT_TOKENS ), + inline_system_policy=getattr(state.config, "anthropic_inline_system", "auto"), ) uid = await submit_generation(spec, state) except ValueError as exc: @@ -164,7 +165,10 @@ async def handle_anthropic_count_tokens(req: AnthropicCountTokensRequest, state: 400, "invalid_request_error", "messages: no tokenizable content" ) try: - n_tokens = await count_prompt_tokens(messages, template_tools, ctk, state) + n_tokens = await count_prompt_tokens( + messages, template_tools, ctk, state, + inline_system_policy=getattr(state.config, "anthropic_inline_system", "auto"), + ) except GenerationError as exc: # The chat template could not render this conversation (bad role ordering, an unmatched # tool_result, ...) — a client error, exactly as /v1/messages classifies the same failure. @@ -186,9 +190,8 @@ def convert_anthropic_prompt( """(messages, template_tools, parser_tools, chat_template_kwargs) — the prompt side of the conversion, shared by /v1/messages and /v1/messages/count_tokens so a counted prompt is exactly the prompt a generation would tokenize.""" - # Collect all system content (top-level `system` + any system-role messages - # Claude Code interleaves in the array) and emit ONE system message at the - # front: strict chat templates (e.g. Qwen3.5) require system at the beginning. + # Hoisting turn-specific reminders rewrites the cached prompt prefix. + # Preserve their position; the tokenizer adapts them to the model's roles. system_texts: list[str] = [] if req.system: if isinstance(req.system, str): @@ -201,7 +204,11 @@ def convert_anthropic_prompt( other: list[dict[str, Any]] = [] for msg in req.messages: if msg.role == "system": - system_texts.append(_content_text(msg.content)) + text = _content_text(msg.content) + if not other: + system_texts.append(text) + elif text: + other.append({"role": "system", "content": text}) continue if isinstance(msg.content, str): @@ -308,6 +315,7 @@ def convert_anthropic_to_genspec( model_sampling: dict[str, Any], reasoning_parser: str | None = None, default_max_tokens: int = DEFAULT_MAX_OUTPUT_TOKENS, + inline_system_policy: str = "auto", ) -> GenSpec: messages, template_tools, parser_tools, ctk = convert_anthropic_prompt( req, reasoning_parser=reasoning_parser @@ -327,6 +335,7 @@ def convert_anthropic_to_genspec( chat_template_kwargs=ctk, template_tools=template_tools, parser_tools=parser_tools, + inline_system_policy=inline_system_policy, ) diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 4dd57b1ae5..2194b82e92 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -59,6 +59,7 @@ class ServerArgs(SchedulerConfig): # prompt_tokens_details.cached_tokens, Anthropic cache_read_input_tokens, Responses # input_tokens_details.cached_tokens). Mirrors sglang's --enable-cache-report. enable_cache_report: bool = False + anthropic_inline_system: str = "auto" # Comma-separated hostname allowlist for client-supplied image URLs; empty admits any domain. allowed_media_domains: str = "" # Directory file:// image refs may be read from; empty rejects local files. @@ -535,6 +536,14 @@ def _infer_reasoning_parser(model_path: str) -> str | None: ), ) + parser.add_argument( + "--anthropic-inline-system", + choices=("auto", "preserve", "fold"), + default=ServerArgs.anthropic_inline_system, + help="Preserve inline system instructions when supported by the renderer, " + "or fold them into nearby user/tool content without hoisting the prompt prefix.", + ) + parser.add_argument( "--sampling-defaults", type=str, diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index 91056e8d1f..ac988eddc7 100644 --- a/python/freetoken/server/generation.py +++ b/python/freetoken/server/generation.py @@ -25,6 +25,7 @@ from freetoken.message import TokenizeMsg from freetoken.mm.media import collect_image_refs, fetch_image_bytes, image_reject_reason from freetoken.tokenizer.tokenize import resolve_thinking_mode +from freetoken.tokenizer.inline_system import InlineSystemError try: # Chat templates render through jinja2 (a transformers dependency): a TemplateError means @@ -142,6 +143,7 @@ class GenSpec: chat_template_kwargs: dict[str, Any] = field(default_factory=dict) template_tools: list[dict[str, Any]] | None = None # tools the model sees (TokenizeMsg.tools) parser_tools: list[dict[str, Any]] | None = None # tools for FunctionCallParser; None disables parsing + inline_system_policy: str | None = None @property def parse_tools(self) -> bool: @@ -302,6 +304,7 @@ async def submit_generation(spec: GenSpec, state: Any) -> int: chat_template_kwargs=spec.chat_template_kwargs, tools=spec.template_tools, images=images, + inline_system_policy=spec.inline_system_policy, ) ) return uid @@ -312,6 +315,8 @@ async def count_prompt_tokens( tools: list[dict[str, Any]] | None, chat_template_kwargs: dict[str, Any], state: Any, + *, + inline_system_policy: str | None = None, ) -> int: """Token count of an already-converted (messages, tools, chat_template_kwargs) prompt, using the frontend's own tokenizer (``state.frontend_tokenizer()``) so the count equals the @@ -334,10 +339,13 @@ class the generation path maps to a 400. A tokenizer *initialization* failure (m chat_template_kwargs=chat_template_kwargs, tools=tools, images=images, + inline_system_policy=inline_system_policy, ) manager = await asyncio.to_thread(state.frontend_tokenizer) # init failure -> server fault try: (user_msg,) = await asyncio.to_thread(manager.tokenize, [msg]) + except InlineSystemError as exc: + raise GenerationError(str(exc)) from exc except _TemplateError as exc: raise GenerationError(str(exc)) from exc return int(user_msg.input_ids.numel()) @@ -361,6 +369,7 @@ async def prerender_error(spec: GenSpec, state: Any) -> GenerationError | None: sampling_params=SamplingParams(), chat_template_kwargs=spec.chat_template_kwargs, tools=spec.template_tools, + inline_system_policy=spec.inline_system_policy, ) try: manager = await asyncio.to_thread(build) diff --git a/python/freetoken/tokenizer/inline_system.py b/python/freetoken/tokenizer/inline_system.py new file mode 100644 index 0000000000..7e6f2c06f9 --- /dev/null +++ b/python/freetoken/tokenizer/inline_system.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: Apache-2.0 +# Renderer-probe approach adapted from vLLM PR #58772. +"""Preserve or fold late system instructions without hoisting the prompt head.""" +from __future__ import annotations + +from copy import deepcopy +from typing import Any, Callable + + +# Allow small role-framing differences, not a repeated environment header. +_INLINE_ROLE_SLACK_TOKENS = 8 + + +class InlineSystemError(ValueError): + """An input instruction cannot be placed without corrupting tool ordering.""" + + +def has_inline_system(messages: list[dict]) -> bool: + started = False + for message in messages: + if message.get("role") == "system": + if started: + return True + else: + started = True + return False + + +def _add_text(message: dict, text: str, *, prepend: bool = False) -> dict: + result = dict(message) + content = result.get("content") + def join(a, b): + return "\n\n".join(part for part in (a, b) if part) + if content is None or isinstance(content, str): + result["content"] = join(text, content) if prepend else join(content, text) + elif isinstance(content, list): + parts = deepcopy(content) + index = 0 if prepend else -1 + if parts and parts[index].get("type") == "text": + previous = parts[index].get("text") or "" + parts[index]["text"] = join(text, previous) if prepend else join(previous, text) + else: + parts.insert(0 if prepend else len(parts), {"type": "text", "text": text}) + result["content"] = parts + else: + raise InlineSystemError("inline system folding requires text or content blocks") + return result + + +def _after_tool_results(messages: list[dict]) -> list[dict]: + out, held = [], [] + pending_ids: set[str] = set() + for message in messages: + role = message.get("role") + if role == "system" and pending_ids: + held.append(message) + continue + if role == "tool": + out.append(message) + pending_ids.discard(message.get("tool_call_id", "")) + if not pending_ids: + out.extend(held) + held.clear() + continue + if pending_ids and held: + raise InlineSystemError("inline system instruction interrupts incomplete tool results") + out.append(message) + pending_ids = {call.get("id", "") for call in message.get("tool_calls", [])} if role == "assistant" else set() + if held: + raise InlineSystemError("inline system instruction is waiting for missing tool results") + return out + + +def normalize_inline_system(messages: list[dict], mode: str) -> list[dict]: + if mode not in ("preserve", "fold"): + raise InlineSystemError(f"invalid inline system mode: {mode}") + if not has_inline_system(messages): + return messages + ordered = _after_tool_results(deepcopy(messages)) + # Without a stable leading system, some renderers attach tools to the first + # late system message. Folding avoids moving tools when that message changes. + if mode == "preserve" and ordered[0].get("role") == "system": + return ordered + out, pending = [], [] + started = False + for message in ordered: + role = message.get("role") + if role == "system" and started: + text = message.get("content") or "" + if not isinstance(text, str): + raise InlineSystemError("inline system instructions must be text") + if text: + if out and out[-1].get("role") in ("user", "tool") and not pending: + out[-1] = _add_text(out[-1], text) + else: + pending.append(text) + continue + if pending: + text = "\n\n".join(pending) + pending.clear() + if role == "user": + message = _add_text(message, text, prepend=True) + else: + out.append({"role": "user", "content": text}) + out.append(message) + started = started or role != "system" + if pending: + out.append({"role": "user", "content": "\n\n".join(pending)}) + return out + + +def _insertion(base: list[int], new: list[int]) -> tuple[int, list[int]] | None: + start = 0 + while start < min(len(base), len(new)) and base[start] == new[start]: + start += 1 + suffix = 0 + while suffix < min(len(base), len(new)) - start and base[-suffix-1] == new[-suffix-1]: + suffix += 1 + if start + suffix != len(base): + return None + return start, new[start:len(new)-suffix] + + +def probe_inline_system(render: Callable[[list[dict]], str], tokenizer: Any) -> tuple[str, str]: + """Conservatively require a distinct, prefix-stable system segment.""" + note, reasoning = "Echo inline instruction", "Foxtrot previous reasoning" + system = {"role": "system", "content": "Alpha stable instructions"} + user = {"role": "user", "content": "Bravo first question"} + cases = [ + ("middle", [system, user, {"role": "assistant", "content": "Charlie first answer"}, + {"role": "user", "content": "Delta second question"}], 3), + ("trailing", [system, user], 2), + ("tool result", [system, user, + {"role": "assistant", "content": "", "reasoning_content": reasoning, "thinking": reasoning, + "reasoning": reasoning, "tool_calls": [{"id": "probe_call", "type": "function", + "function": {"name": "read", "arguments": {}}}]}, + {"role": "tool", "tool_call_id": "probe_call", "content": "Golf tool output"}], 4), + ] + max_system_tokens = None + try: + special_ids = set(tokenizer.all_special_ids) | set(tokenizer.get_added_vocab().values()) + if not special_ids: + return "fold", "renderer has no identifiable special tokens" + for name, messages, index in cases: + base_text = render(deepcopy(messages)) + base = tokenizer.encode(base_text, add_special_tokens=False) + with_note = messages[:index] + [{"role": "system", "content": note}] + messages[index:] + new_text = render(deepcopy(with_note)) + new = tokenizer.encode(new_text, add_special_tokens=False) + found = _insertion(base, new) + if found is None: + return "fold", f"{name}: other prompt tokens change" + position, segment = found + if index == len(messages) and position >= len(base): + return "fold", f"{name}: instruction follows the generation opening" + if tokenizer.decode(segment).count(note) != 1 or not special_ids.intersection(segment): + return "fold", f"{name}: no distinct role-marked instruction" + if (reasoning in base_text) != (reasoning in new_text): + return "fold", f"{name}: reasoning retention changes" + if name == "middle": + as_user = messages[:index] + [{"role": "user", "content": note}] + messages[index:] + user_tokens = tokenizer.encode(render(deepcopy(as_user)), add_special_tokens=False) + user_segment = _insertion(base, user_tokens) + if user_segment is None: + return "fold", "middle: cannot isolate a user instruction" + if user_segment[1] == segment: + return "fold", "system instructions render as ordinary user turns" + max_system_tokens = len(user_segment[1]) + _INLINE_ROLE_SLACK_TOKENS + if max_system_tokens is not None and len(segment) > max_system_tokens: + return "fold", f"{name}: system insertion overhead ({len(segment)} > {max_system_tokens} tokens)" + except Exception as exc: + return "fold", f"probe could not render: {type(exc).__name__}" + return "preserve", "renderer preserves distinct system turns and existing prompt tokens" diff --git a/python/freetoken/tokenizer/tokenize.py b/python/freetoken/tokenizer/tokenize.py index 758ffd4dcf..f473665267 100644 --- a/python/freetoken/tokenizer/tokenize.py +++ b/python/freetoken/tokenizer/tokenize.py @@ -4,6 +4,7 @@ import json import os import threading +from copy import deepcopy from types import ModuleType from typing import TYPE_CHECKING, Any, List @@ -22,6 +23,7 @@ probe_thinking_profile, quantize_effort, ) +from .inline_system import has_inline_system, normalize_inline_system, probe_inline_system logger = init_logger(__name__) @@ -58,6 +60,8 @@ def __init__(self, tokenizer: PreTrainedTokenizerBase, mm_processor: MMProcessor self._thinking_profile: ThinkingProfile | None = None self._effort_lock = threading.Lock() self._logged_effort_maps: set[tuple[Any, str | None]] = set() + self._inline_system_lock = threading.Lock() + self._inline_system_modes: dict[str, str] = {} def tokenize(self, msgs: List[TokenizeMsg]) -> List[UserMsg]: results: List[UserMsg] = [] @@ -103,9 +107,30 @@ def render_prompt(self, msg: TokenizeMsg) -> str: validation, count_tokens) must quantize identically.""" if not isinstance(msg.text, list): return msg.text - return self._render( - msg.text, msg.tools, self._sanitize_effort(msg.chat_template_kwargs or {}) - ) + kwargs = self._sanitize_effort(msg.chat_template_kwargs or {}) + messages = msg.text + if msg.inline_system_policy is not None: + if msg.inline_system_policy not in ("auto", "preserve", "fold"): + raise ValueError(f"invalid inline system policy: {msg.inline_system_policy}") + if has_inline_system(messages): + mode = msg.inline_system_policy + if mode == "auto": + mode = self.inline_system_mode(msg.tools, kwargs) + messages = normalize_inline_system(messages, mode) + return self._render(messages, msg.tools, kwargs) + + def inline_system_mode(self, tools: list[dict] | None, kwargs: dict[str, Any]) -> str: + key = json.dumps([tools, kwargs], sort_keys=True) + with self._inline_system_lock: + if key not in self._inline_system_modes: + mode, reason = probe_inline_system( + lambda messages: self._render(messages, deepcopy(tools), deepcopy(kwargs)), self.tokenizer + ) + if len(self._inline_system_modes) >= 128: + self._inline_system_modes.pop(next(iter(self._inline_system_modes))) + self._inline_system_modes[key] = mode + logger.info("Anthropic inline system policy: %s (%s)", mode, reason) + return self._inline_system_modes[key] def _render( self, diff --git a/tests/README.md b/tests/README.md index 26dba63391..be22c1fb20 100644 --- a/tests/README.md +++ b/tests/README.md @@ -53,7 +53,10 @@ checkpoint is set: | `FREETOKEN_GEMMA4_MODEL` | `models/test_gemma4_vision.py`, `tokenizer/test_mm_tokenize.py` — local Gemma-4 `gemma4` checkpoint (tower parity, boi/eoi expansion; the tokenizer test also takes the 12B `gemma4_unified` checkpoint) | | `FREETOKEN_GEMMA4_UNIFIED_MODEL` | `models/test_gemma4_vision.py` — local gemma-4-12B `gemma4_unified` checkpoint (embedder parity) | | `FREETOKEN_GLM53_MODEL` | `tokenizer/test_mm_tokenize.py` — local GLM-5.3-Flash checkpoint (image-token expansion) | -| `FREETOKEN_MUSE_MODEL` | `models/test_muse_glimmer_vision.py`, `tokenizer/test_mm_tokenize.py` — local Muse-Glimmer-30B checkpoint (tower parity, image_start/image_end expansion) | +| `FREETOKEN_DSV4_MODEL` | `tokenizer/test_tokenize.py` - local DSV4 checkpoint with its native `encoding/encoding_dsv4.py`; prefix stability and reminder retention, tokenizer only (no GPU or model weights loaded) | +| `FREETOKEN_INLINE_GLM_MODEL` | `tokenizer/test_tokenize.py` - local GLM-5.3 tokenizer/template; inline-system preserve policy, prefix stability and reasoning retention (no weights loaded) | +| `FREETOKEN_INLINE_QWEN_MODEL` | `tokenizer/test_tokenize.py` - local Qwen3.6 tokenizer/template; inline-system fold policy, prefix stability and reasoning retention (no weights loaded) | +| `FREETOKEN_MUSE_MODEL` | `models/test_muse_glimmer_vision.py`, `tokenizer/test_mm_tokenize.py`, `tokenizer/test_tokenize.py` - local Muse-Glimmer-30B checkpoint (vision tests; tokenizer-only inline-system regressions) | | `FREETOKEN_MINIMAX_M3_MODEL` | `tokenizer/test_mm_tokenize.py` — local MiniMax-M3 checkpoint (start/end expansion, token budget) | | `FREETOKEN_TEST_MOE_CACHE_SIZE` | `e2e/test_aime.py` — >0 switches to the offload MoE backend with this cache size | | `FREETOKEN_TEST_MEM_RATIO` | `e2e/test_aime.py` — offload-mode memory_ratio (default `0.9`) | diff --git a/tests/server/test_anthropic_api.py b/tests/server/test_anthropic_api.py index 327bfcd87d..f382093fc0 100644 --- a/tests/server/test_anthropic_api.py +++ b/tests/server/test_anthropic_api.py @@ -231,10 +231,7 @@ def test_convert_thinking_replay_in_tool_loop(): assert spec.messages[2]["role"] == "tool" -def test_convert_hoists_and_merges_system_messages(): - # Claude Code interleaves system messages mid-array; strict chat templates - # (e.g. Qwen3.5: "System message must be at the beginning") require ONE system - # message at the front. Merge top-level system + in-array system, hoist to front. +def test_convert_preserves_mid_conversation_system_reminders(): req = AnthropicMessagesRequest.model_validate( { "model": "claude-x", @@ -247,11 +244,69 @@ def test_convert_hoists_and_merges_system_messages(): ], } ) + original = req.model_dump() spec = A.convert_anthropic_to_genspec(req, {}) - assert [m["role"] for m in spec.messages] == ["system", "user", "assistant"] - assert sum(1 for m in spec.messages if m["role"] == "system") == 1 - assert "top-level sys" in spec.messages[0]["content"] - assert "mid-stream sys" in spec.messages[0]["content"] + assert spec.messages == [ + {"role": "system", "content": "top-level sys"}, + {"role": "user", "content": "hello"}, + {"role": "system", "content": "mid-stream sys"}, + {"role": "assistant", "content": "hi"}, + ] + assert req.model_dump() == original + + +def test_convert_merges_only_leading_system_messages(): + req = AnthropicMessagesRequest.model_validate({ + "model": "claude-x", "max_tokens": 64, + "system": [{"type": "text", "text": "top-level sys"}], + "messages": [ + {"role": "system", "content": "leading sys"}, + {"role": "system", "content": [{"type": "text", "text": "second sys"}]}, + {"role": "user", "content": "hello"}, + {"role": "system", "content": [{"type": "text", "text": "turn reminder"}]}, + {"role": "assistant", "content": "hi"}, + {"role": "system", "content": ""}, + {"role": "system", "content": "next reminder"}, + ], + }) + original = req.model_dump() + spec = A.convert_anthropic_to_genspec(req, {}) + assert spec.messages == [ + {"role": "system", "content": "top-level sys\n\nleading sys\n\nsecond sys"}, + {"role": "user", "content": "hello"}, + {"role": "system", "content": "turn reminder"}, + {"role": "assistant", "content": "hi"}, + {"role": "system", "content": "next reminder"}, + ] + assert req.model_dump() == original + + +@pytest.mark.parametrize("system", [None, "stable instructions"]) +@pytest.mark.parametrize("count_tokens", [False, True]) +def test_appending_system_reminder_preserves_converted_prefix(system, count_tokens): + from freetoken.server.anthropic_models import AnthropicCountTokensRequest + + data = { + "model": "claude-x", + "messages": [ + {"role": "user", "content": "inspect the source"}, + {"role": "assistant", "content": "I found the function"}, + ], + "tools": [{"name": "read", "input_schema": {"type": "object", "properties": {}}}], + } + if system is not None: + data["system"] = system + if not count_tokens: + data["max_tokens"] = 64 + request_type = AnthropicCountTokensRequest if count_tokens else AnthropicMessagesRequest + first, first_tools, _, _ = A.convert_anthropic_prompt(request_type.model_validate(data)) + data["messages"].append({"role": "system", "content": "1000 tokens left"}) + second, second_tools, _, _ = A.convert_anthropic_prompt(request_type.model_validate(data)) + assert second[:-1] == first + assert second[-1] == { + "role": "system", "content": "1000 tokens left" + } + assert second_tools == first_tools # --------------------------------------------------------------------------- # @@ -664,6 +719,7 @@ def test_count_tokens_route(): assert msg.text == spec.messages assert msg.tools == spec.template_tools assert msg.chat_template_kwargs == spec.chat_template_kwargs + assert msg.inline_system_policy == spec.inline_system_policy == "auto" def test_count_tokens_empty_messages_400(): @@ -673,6 +729,28 @@ def test_count_tokens_empty_messages_400(): assert r.json()["error"]["type"] == "invalid_request_error" +@pytest.mark.parametrize("policy", ["auto", "preserve", "fold"]) +def test_inline_system_override_reaches_counting_and_generation(policy): + manager = _FakeTokenizeManager() + fake = FakeState([("ok", True, 3, 1)]) + fake._count_manager = manager + fake.config.anthropic_inline_system = policy + submitted = [] + + async def record(msg): + submitted.append(msg) + + fake.send_one = record + client = _client(fake) + body = {**_COUNT_BODY, "messages": [ + {"role": "user", "content": "question"}, {"role": "system", "content": "instruction"}]} + assert client.post("/v1/messages/count_tokens", json=body).status_code == 200 + assert client.post("/v1/messages", json={**body, "max_tokens": 32}).status_code == 200 + assert manager.msgs[0].inline_system_policy == submitted[0].inline_system_policy == policy + assert manager.msgs[0].text == submitted[0].text + assert submitted[0].text[0]["content"] == body["system"] + + def test_count_tokens_works_while_not_serving(): # Counting never touches the engine, so it must not 503 during load/rebuild. client, _ = _count_client(_FakeTokenizeManager(), maintenance_state="rebuilding") @@ -719,6 +797,15 @@ def test_count_tokens_template_render_error_400(): assert r.json()["error"]["type"] == "invalid_request_error" +def test_count_tokens_inline_system_placement_error_is_400(): + from freetoken.tokenizer.inline_system import InlineSystemError + + client, _ = _count_client(_RaisingManager(InlineSystemError("missing tool results"))) + r = client.post("/v1/messages/count_tokens", json=_COUNT_BODY) + assert r.status_code == 400 + assert r.json()["error"]["type"] == "invalid_request_error" + + def test_count_tokens_excluded_from_request_ring(): # count_tokens never enters generation accounting; its latency (incl. first-touch tokenizer # load) must not land in the /v1/requests ring or the /v1/stats p95. diff --git a/tests/server/test_message_wire.py b/tests/server/test_message_wire.py index e228667154..fdc095921d 100644 --- a/tests/server/test_message_wire.py +++ b/tests/server/test_message_wire.py @@ -28,6 +28,18 @@ from freetoken.core import SamplingParams +def test_inline_system_policy_tokenizer_roundtrip(): + for policy in (None, "auto", "preserve", "fold"): + msg = TokenizeMsg(uid=1, text=[{"role": "system", "content": "instruction"}], + sampling_params=SamplingParams(), inline_system_policy=policy) + decoded = BaseTokenizerMsg.decoder(BaseTokenizerMsg.encoder(msg)) + assert decoded.inline_system_policy == policy + assert decoded.text == msg.text + legacy = BaseTokenizerMsg.encoder(msg) + legacy.pop("inline_system_policy") + assert BaseTokenizerMsg.decoder(legacy).inline_system_policy is None + + def test_cache_rebuild_msg_roundtrip(): msg = CacheRebuildMsg(request_id="abc", moe_cache_size=8, num_pages=1024, mode="if_idle") out = BaseTokenizerMsg.decoder(BaseTokenizerMsg.encoder(msg)) diff --git a/tests/tokenizer/test_tokenize.py b/tests/tokenizer/test_tokenize.py index 369576a18b..16ef6d7419 100644 --- a/tests/tokenizer/test_tokenize.py +++ b/tests/tokenizer/test_tokenize.py @@ -1,6 +1,8 @@ from __future__ import annotations import json +import os +from copy import deepcopy import pytest import torch @@ -8,14 +10,303 @@ from freetoken.core import SamplingParams from freetoken.message import TokenizeMsg from freetoken.tokenizer.tokenize import TokenizeManager, _dsv4_arguments_str +from freetoken.tokenizer.inline_system import InlineSystemError, normalize_inline_system, probe_inline_system + + +def test_fold_keeps_leading_system_and_appends_text_without_mutation(): + messages = [ + {"role": "system", "content": "stable"}, + {"role": "user", "content": [{"type": "image_url", "image_url": {"url": "image"}}, + {"type": "text", "text": "question"}]}, + {"role": "system", "content": "instruction 1"}, + {"role": "system", "content": "instruction 2"}, + ] + original = deepcopy(messages) + result = normalize_inline_system(messages, "fold") + assert result == [messages[0], {"role": "user", "content": [ + messages[1]["content"][0], {"type": "text", "text": "question\n\ninstruction 1\n\ninstruction 2"}]}] + assert messages == original + + +@pytest.mark.parametrize("mode", ["preserve", "fold"]) +def test_inline_system_waits_for_all_parallel_tool_results(mode): + messages = [ + {"role": "system", "content": "stable"}, {"role": "user", "content": "question"}, + {"role": "assistant", "reasoning_content": "reason", "tool_calls": [{"id": "a"}, {"id": "b"}]}, + {"role": "system", "content": "before results"}, + {"role": "tool", "tool_call_id": "a", "content": "result a"}, + {"role": "system", "content": "between results"}, + {"role": "tool", "tool_call_id": "b", "content": "result b"}, + {"role": "system", "content": "after results"}, + ] + original = deepcopy(messages) + result = normalize_inline_system(messages, mode) + assert [m["tool_call_id"] for m in result if m["role"] == "tool"] == ["a", "b"] + assert result[2] == messages[2] + assert [m["role"] for m in result[3:5]] == ["tool", "tool"] + if mode == "fold": + assert len(result) == 5 + assert result[-1]["content"] == "result b\n\nbefore results\n\nbetween results\n\nafter results" + else: + assert [m["content"] for m in result[5:]] == ["before results", "between results", "after results"] + assert messages == original + + +@pytest.mark.parametrize("mode", ["preserve", "fold"]) +def test_inline_system_does_not_hide_missing_tool_results(mode): + messages = [{"role": "assistant", "tool_calls": [{"id": "missing"}]}, + {"role": "system", "content": "instruction"}] + with pytest.raises(InlineSystemError, match="missing tool results"): + normalize_inline_system(messages, mode) + + +def test_fold_after_assistant_prepends_next_user_or_supplies_a_user_turn(): + messages = [{"role": "user", "content": "first"}, {"role": "assistant", "content": "answer"}, + {"role": "system", "content": "instruction"}] + assert normalize_inline_system(messages, "fold")[-1] == {"role": "user", "content": "instruction"} + result = normalize_inline_system(messages + [{"role": "user", "content": "next"}], "fold") + assert result == messages[:2] + [{"role": "user", "content": "instruction\n\nnext"}] + + +def test_preserve_keeps_instruction_role_and_only_folds_without_a_stable_system(): + messages = [{"role": "system", "content": "stable"}, {"role": "user", "content": "question"}, + {"role": "system", "content": "late"}] + assert normalize_inline_system(messages, "preserve") == messages + assert normalize_inline_system(messages[1:], "preserve") == [{"role": "user", "content": "question\n\nlate"}] + + +class ProbeTokenizer: + vocabulary = {text: 10000+i for i, text in enumerate( + ["", "", "", "", "", "", + "", "", "", ""])} + all_special_ids = list(vocabulary.values()) + + def get_added_vocab(self): + return self.vocabulary + + def encode(self, text, **kwargs): + import re + parts = re.split("(" + "|".join(re.escape(x) for x in self.vocabulary) + ")", text) + return [token for part in parts for token in ([self.vocabulary[part]] if part in self.vocabulary else map(ord, part))] + + def decode(self, ids): + inverse = {v: k for k, v in self.vocabulary.items()} + return "".join(inverse[x] if x in inverse else chr(x) for x in ids) + + +@pytest.mark.parametrize("behavior,expected", [("preserve", "preserve"), ("reject", "fold"), + ("drop", "fold"), ("bare", "fold"), ("user", "fold"), ("rewrite", "fold"), + ("after_generation", "fold"), ("drop_reasoning", "fold")]) +def test_renderer_probe_checks_more_than_template_acceptance(behavior, expected): + def render(messages): + inline = any(m["role"] == "system" for m in messages[1:]) + output, late = "", "" + for i, message in enumerate(messages): + role, content = message["role"], message.get("content", "") + if role == "system" and i: + if behavior == "reject": + raise ValueError("system first only") + if behavior == "drop": + continue + if behavior == "bare": + output += content + continue + if behavior == "user": + role = "user" + if behavior == "after_generation": + late += content + continue + if behavior == "rewrite" and inline and i == 0: + content = "changed head" + reason = message.get("reasoning_content") + if reason and not (behavior == "drop_reasoning" and inline): + content = "" + reason + "" + content + output += f"<{role}>{content}" + return output + "" + late + assert probe_inline_system(render, ProbeTokenizer())[0] == expected + + +@pytest.mark.parametrize("position", ["middle", "trailing", "tool result"]) +@pytest.mark.parametrize("metadata", ["Reasoning strength: high. Valid recipients: self, user.", "tool schema " * 1000], + ids=["metadata", "tools"]) +def test_inline_probe_rejects_repeated_system_metadata(position, metadata): + def render(messages): + output = "" + for i, message in enumerate(messages): + role = message["role"] + content = message.get("content", "") + if role == "system" and i: + if i < len(messages)-1: + location = "middle" + else: + location = "tool result" if messages[i-1]["role"] == "tool" else "trailing" + if location == position: + content += metadata + content = message.get("reasoning_content", "") + content + output += f"<{role}>{content}" + return output + "" + + mode, reason = probe_inline_system(render, ProbeTokenizer()) + assert mode == "fold" + assert "overhead" in reason + + +def test_inline_probe_allows_small_role_framing_overhead(): + def render(messages): + return "".join(f'<{m["role"]}>' + ("role: " if m["role"] == "system" else "") + + m.get("reasoning_content", "") + m.get("content", "") + + f'' for m in messages) + "" + + assert probe_inline_system(render, ProbeTokenizer())[0] == "preserve" + + +def test_inline_probe_folds_when_user_insertion_cannot_be_measured(): + def render(messages): + rewrite = any(m["role"] == "user" and m.get("content") == "Echo inline instruction" for m in messages) + return ("changed head" if rewrite else "stable head") + "".join( + f'<{m["role"]}>' + m.get("reasoning_content", "") + m.get("content", "") + + f'' for m in messages) + "" + + mode, reason = probe_inline_system(render, ProbeTokenizer()) + assert mode == "fold" + assert "cannot isolate" in reason + + +def test_inline_probe_cache_is_scoped_to_tools_and_render_kwargs(monkeypatch): + import freetoken.tokenizer.tokenize as module + calls = [] + monkeypatch.setattr(module, "probe_inline_system", lambda *_: (calls.append(1) or ("fold", "test"))) + manager = TokenizeManager(FakeTokenizer()) + for _ in range(3): + assert manager.inline_system_mode(None, {}) == "fold" + manager.inline_system_mode([], {}) + manager.inline_system_mode(None, {"enable_thinking": True}) + assert len(calls) == 3 + + +def test_inline_policy_is_opt_in_and_not_forwarded_to_templates(): + tokenizer = FakeTokenizer() + manager = TokenizeManager(tokenizer) + messages = [{"role": "user", "content": "question"}, {"role": "system", "content": "instruction"}] + msg = TokenizeMsg(uid=1, text=messages, sampling_params=SamplingParams()) + manager.render_prompt(msg) + assert tokenizer.chat_template_messages == messages + msg.inline_system_policy = "fold" + manager.render_prompt(msg) + assert tokenizer.chat_template_messages == [{"role": "user", "content": "question\n\ninstruction"}] + assert "inline_system_policy" not in tokenizer.chat_template_kwargs + + +@pytest.mark.needs_weights +@pytest.mark.parametrize("variable,expected_mode", [ + ("FREETOKEN_DSV4_MODEL", "fold"), + ("FREETOKEN_INLINE_GLM_MODEL", "preserve"), + ("FREETOKEN_INLINE_QWEN_MODEL", "fold"), +]) +def test_native_model_inline_system_tool_history(variable, expected_mode): + from transformers import AutoTokenizer + from freetoken.server.anthropic_api import convert_anthropic_to_genspec + from freetoken.server.anthropic_models import AnthropicMessagesRequest + + path = os.environ.get(variable) + if not path: + pytest.skip(f"set {variable} to a local checkpoint") + manager = TokenizeManager(AutoTokenizer.from_pretrained(path, local_files_only=True)) + data = { + "model": "local", "max_tokens": 64, "system": "Stable project instructions.", + "thinking": {"type": "enabled"}, + "tools": [{"name": "read", "input_schema": {"type": "object", "properties": {}}}], + "messages": [ + {"role": "user", "content": "Inspect source:\n" + "value = 1\n" * 512}, + {"role": "system", "content": "Keep all existing tests."}, + ], + } + ids, prompts = [], [] + for turn in (1, 2): + data["messages"].extend([ + {"role": "assistant", "content": [ + {"type": "thinking", "thinking": f"Retained thought {turn}.", "signature": ""}, + {"type": "tool_use", "id": f"read_{turn}", "name": "read", "input": {}}, + ]}, + {"role": "system", "content": f"During tool call {turn}."}, + {"role": "user", "content": [ + {"type": "tool_result", "tool_use_id": f"read_{turn}", "content": f"Source result {turn}."}, + ]}, + {"role": "system", "content": f"Remaining budget: {1000-turn}."}, + ]) + req = AnthropicMessagesRequest.model_validate(data) + spec = convert_anthropic_to_genspec(req, {}) + assert manager.inline_system_mode(spec.template_tools, spec.chat_template_kwargs) == expected_mode + original = deepcopy(spec.messages) + msg = TokenizeMsg(uid=turn, text=spec.messages, tools=spec.template_tools, + sampling_params=spec.sampling_params, chat_template_kwargs=spec.chat_template_kwargs, + inline_system_policy=spec.inline_system_policy) + prompts.append(manager.render_prompt(msg)) + ids.append(manager.tokenize([msg])[0].input_ids.tolist()) + assert spec.messages == original + assert prompts[-1].count("Keep all existing tests.") == 1 + assert "Retained thought 1." in prompts[-1] + for previous in range(1, turn+1): + assert prompts[-1].count(f"During tool call {previous}.") == 1 + assert prompts[-1].count(f"Remaining budget: {1000-previous}.") == 1 + common = next((i for i, pair in enumerate(zip(*ids)) if pair[0] != pair[1]), min(map(len, ids))) + assert common >= len(ids[0]) - 64 + assert common > 2048 + + +@pytest.mark.needs_weights +@pytest.mark.parametrize("tool_count", [0, 20]) +def test_native_muse_inline_system_does_not_repeat_environment(tool_count): + from transformers import AutoTokenizer + + path = os.environ.get("FREETOKEN_MUSE_MODEL") + if not path: + pytest.skip("set FREETOKEN_MUSE_MODEL to a local checkpoint") + manager = TokenizeManager(AutoTokenizer.from_pretrained(path, local_files_only=True)) + tools = [{"type": "function", "function": {"name": f"read_{i}", + "description": "Inspect the repository source. " * 100, + "parameters": {"type": "object", "properties": {}}}} for i in range(tool_count)] or None + messages = [{"role": "system", "content": "Stable instructions."}, + {"role": "user", "content": "Inspect this file.\n" + "value = 1\n" * 256}] + assert manager.inline_system_mode(tools, {}) == "fold" + for turn in range(3): + messages.extend([{"role": "assistant", "content": f"Answer {turn}."}, + {"role": "user", "content": f"Question {turn}."}, + {"role": "system", "content": f"Budget reminder {turn}."}]) + original = deepcopy(messages) + msg = TokenizeMsg(uid=turn, text=messages, tools=tools, sampling_params=SamplingParams(), + inline_system_policy="auto") + prompt = manager.render_prompt(msg) + assert prompt.count("Reasoning strength:") == 1 + assert prompt.count("# Valid recipients:") == 1 + assert prompt.count("// Function schemas") == bool(tools) + for previous in range(turn+1): + assert prompt.count(f"Budget reminder {previous}.") == 1 + ids = manager.tokenize([msg])[0].input_ids.tolist() + msg.inline_system_policy = "fold" + assert ids == manager.tokenize([msg])[0].input_ids.tolist() + assert messages == original + + latest = [m for i, m in enumerate(messages) if m["role"] != "system" or i in (0, len(messages)-1)] + updated = deepcopy(latest) + updated[-1]["content"] = "Budget reminder 900." + tokens = [] + for conversation in (latest, updated): + msg = TokenizeMsg(uid=0, text=conversation, tools=tools, sampling_params=SamplingParams(), + inline_system_policy="auto") + tokens.append(manager.tokenize([msg])[0].input_ids.tolist()) + common = next((i for i, pair in enumerate(zip(*tokens)) if pair[0] != pair[1]), min(map(len, tokens))) + assert common >= min(map(len, tokens)) - 16 class FakeTokenizer: def __init__(self) -> None: self.chat_template_kwargs = None + self.chat_template_messages = None def apply_chat_template(self, messages, **kwargs): self.chat_template_kwargs = kwargs + self.chat_template_messages = deepcopy(messages) return "rendered prompt" def encode(self, prompt, return_tensors=None, add_special_tokens=True): @@ -74,6 +365,31 @@ def test_tokenize_manager_passes_tools_to_chat_template(): assert input_ids.tolist() == [1, 2, 3] +@pytest.mark.parametrize("initial_system", [[], [{"role": "system", "content": "stable instructions"}]]) +def test_generic_template_folds_system_messages_without_mutating_input(initial_system): + tokenizer = FakeTokenizer() + manager = TokenizeManager(tokenizer) + messages = initial_system + [ + {"role": "user", "content": "inspect files"}, + {"role": "system", "content": "first reminder"}, + {"role": "assistant", "content": "found the bug"}, + {"role": "system", "content": "second reminder"}, + ] + original = deepcopy(messages) + tools = [{"type": "function", "function": {"name": "read", "parameters": {"type": "object"}}}] + original_tools = deepcopy(tools) + manager.tokenize([TokenizeMsg(uid=1, text=messages, tools=tools, sampling_params=SamplingParams(), + inline_system_policy="auto")]) + assert tokenizer.chat_template_messages == initial_system + [ + {"role": "user", "content": "inspect files\n\nfirst reminder"}, + {"role": "assistant", "content": "found the bug"}, + {"role": "user", "content": "second reminder"}, + ] + assert tokenizer.chat_template_kwargs["tools"] == original_tools + assert messages == original + assert tools == original_tools + + class FakeDsv4Tokenizer: chat_template = None @@ -126,6 +442,166 @@ def encode_messages(messages, thinking_mode, reasoning_effort=None): assert input_ids.tolist() == [4, 5, 6] +@pytest.mark.parametrize("initial_system", [[], [{"role": "system", "content": "stable instructions"}]]) +def test_dsv4_fold_keeps_instructions_and_tools_on_stable_prefix(tmp_path, initial_system): + encoding_dir = tmp_path / "encoding" + encoding_dir.mkdir() + (encoding_dir / "encoding_dsv4.py").write_text( + "import json\n" + "def encode_messages(messages, thinking_mode, reasoning_effort=None):\n" + " return json.dumps(messages)\n" + ) + tokenizer = FakeDsv4Tokenizer(tmp_path) + manager = TokenizeManager(tokenizer) + tools = [{"type": "function", "function": {"name": "read", "parameters": {"type": "object"}}}] + messages = initial_system + [ + {"role": "user", "content": "inspect files"}, + {"role": "system", "content": "first reminder"}, + {"role": "assistant", "content": "found the bug"}, + ] + original, original_tools = deepcopy(messages), deepcopy(tools) + manager.tokenize([TokenizeMsg(uid=1, text=messages, tools=tools, sampling_params=SamplingParams(), + inline_system_policy="fold")]) + first = json.loads(tokenizer.prompt) + extended = messages + [{"role": "system", "content": "next reminder"}] + original_extended = deepcopy(extended) + manager.tokenize([TokenizeMsg(uid=2, text=extended, tools=tools, sampling_params=SamplingParams(), + inline_system_policy="fold")]) + second = json.loads(tokenizer.prompt) + assert [m["role"] for m in first] == ["system", "user", "assistant"] + assert first[0]["content"] == (initial_system[0]["content"] if initial_system else "") + assert first[0]["tools"] == original_tools + assert first[1]["content"] == "inspect files\n\nfirst reminder" + assert second[:-1] == first + assert second[-1] == {"role": "user", "content": "next reminder"} + assert all("tools" not in m for m in second[1:]) + assert messages == original and extended == original_extended and tools == original_tools + + +@pytest.fixture(scope="module") +def native_dsv4_manager(): + model_path = os.environ.get("FREETOKEN_DSV4_MODEL") + if not model_path: + pytest.skip("set FREETOKEN_DSV4_MODEL to a local DSV4 checkpoint") + from transformers import AutoTokenizer + + manager = TokenizeManager(AutoTokenizer.from_pretrained(model_path, local_files_only=True)) + assert manager._dsv4_encoder is not None, "checkpoint must supply the native DSV4 encoder" + return manager + + +def test_dsv4_no_tools_retains_reminders_when_encoder_drops_developer_turns(tmp_path): + encoding_dir = tmp_path / "encoding" + encoding_dir.mkdir() + (encoding_dir / "encoding_dsv4.py").write_text( + "import json\n" + "def encode_messages(messages, thinking_mode, reasoning_effort=None):\n" + " if thinking_mode == 'thinking':\n" + " messages = [m for m in messages if m['role'] != 'developer']\n" + " return json.dumps(messages)\n" + ) + manager = TokenizeManager(FakeDsv4Tokenizer(tmp_path)) + messages = [ + {"role": "user", "content": "First task."}, + {"role": "system", "content": "Keep this earlier instruction."}, + {"role": "assistant", "content": "First answer."}, + {"role": "user", "content": "Second task."}, + ] + original = deepcopy(messages) + msg = TokenizeMsg(uid=1, text=messages, sampling_params=SamplingParams(), + chat_template_kwargs={"enable_thinking": True}, inline_system_policy="fold") + rendered = json.loads(manager.render_prompt(msg)) + assert all(m["role"] in ("system", "user", "assistant") for m in rendered) + assert sum(m["content"].count("Keep this earlier instruction.") for m in rendered) == 1 + assert messages == original + + +@pytest.mark.needs_weights +@pytest.mark.parametrize("system", [None, "You are a coding assistant."]) +@pytest.mark.parametrize("change", ["update", "append"]) +def test_native_dsv4_reminder_update_preserves_token_prefix(native_dsv4_manager, system, change): + from freetoken.server.anthropic_api import convert_anthropic_to_genspec + from freetoken.server.anthropic_models import AnthropicMessagesRequest + + data = { + "model": "dsv4", "max_tokens": 64, + "tools": [{"name": "read", "input_schema": {"type": "object", "properties": {}}}], + "messages": [ + {"role": "user", "content": "Inspect this file:\n" + "value = 1\n" * 512}, + {"role": "system", "content": "Keep the changes minimal."}, + {"role": "assistant", "content": [ + {"type": "thinking", "thinking": "I should read the source.", "signature": ""}, + {"type": "tool_use", "id": "call_read", "name": "read", "input": {}}, + ]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call_read", "content": "value = 1"}]}, + {"role": "system", "content": "Remaining token budget: 1000"}, + ], + } + if system is not None: + data["system"] = system + tokens = [] + for budget in (1000, 900): + if budget == 900 and change == "append": + data["messages"].extend([ + {"role": "assistant", "content": [ + {"type": "thinking", "thinking": "The file is simple.", "signature": ""}, + {"type": "text", "text": "I inspected the file."}, + ]}, + {"role": "system", "content": "Remaining token budget: 900"}, + ]) + else: + data["messages"][-1]["content"] = f"Remaining token budget: {budget}" + req = AnthropicMessagesRequest.model_validate(data) + original = req.model_dump() + spec = convert_anthropic_to_genspec(req, {}) + msg = TokenizeMsg(uid=budget, text=spec.messages, tools=spec.template_tools, + sampling_params=spec.sampling_params, chat_template_kwargs=spec.chat_template_kwargs, + inline_system_policy=spec.inline_system_policy) + prompt = native_dsv4_manager.render_prompt(msg) + assert prompt.count("Keep the changes minimal.") == 1 + assert prompt.index("Inspect this file:") < prompt.index(f"Remaining token budget: {budget}") + tokens.append(native_dsv4_manager.tokenize([msg])[0].input_ids.tolist()) + assert req.model_dump() == original + common = next((i for i, pair in enumerate(zip(*tokens)) if pair[0] != pair[1]), min(map(len, tokens))) + assert common > 2048 + assert common >= min(map(len, tokens)) - 64 + + +@pytest.mark.needs_weights +@pytest.mark.parametrize("disable_declared_tools", [False, True]) +@pytest.mark.parametrize("thinking_enabled", [False, True]) +def test_native_dsv4_no_tools_keeps_earlier_reminder( + native_dsv4_manager, disable_declared_tools, thinking_enabled +): + from freetoken.server.anthropic_api import convert_anthropic_to_genspec + from freetoken.server.anthropic_models import AnthropicMessagesRequest + + data = { + "model": "dsv4", "max_tokens": 64, "system": "Stable instructions.", + "thinking": {"type": "enabled" if thinking_enabled else "disabled"}, + "messages": [ + {"role": "user", "content": "First task."}, + {"role": "system", "content": "Keep this earlier instruction."}, + {"role": "assistant", "content": "First answer."}, + {"role": "user", "content": "Second task."}, + ], + } + if disable_declared_tools: + data["tools"] = [{"name": "read", "input_schema": {"type": "object", "properties": {}}}] + data["tool_choice"] = {"type": "none"} + req = AnthropicMessagesRequest.model_validate(data) + original = req.model_dump() + spec = convert_anthropic_to_genspec(req, {}) + assert not spec.template_tools + msg = TokenizeMsg(uid=1, text=spec.messages, tools=spec.template_tools, + sampling_params=spec.sampling_params, chat_template_kwargs=spec.chat_template_kwargs, + inline_system_policy=spec.inline_system_policy) + prompt = native_dsv4_manager.render_prompt(msg) + assert prompt.count("Keep this earlier instruction.") == 1 + assert "Second task." in prompt + assert req.model_dump() == original + + def test_dsv4_encoder_gets_tool_call_arguments_as_json_string(tmp_path): """Regression: render_messages hands the template dict arguments; the dsv4 encoder contract is a JSON-object STRING -- a dict trips its fallback that From 555efd89447232a555e05d187256b76a51c1aaa0 Mon Sep 17 00:00:00 2001 From: cherry77-cloud <1615405@qq.com> Date: Tue, 6 Oct 2026 06:29:44 +0800 Subject: [PATCH 04/73] fix(kvcache): clone inserted radix keys to release request buffers (#559) --- python/freetoken/kvcache/hybrid_radix_cache.py | 5 ++++- python/freetoken/kvcache/radix_cache.py | 5 ++++- python/freetoken/kvcache/swa_radix_cache.py | 3 ++- 3 files changed, 10 insertions(+), 3 deletions(-) diff --git a/python/freetoken/kvcache/hybrid_radix_cache.py b/python/freetoken/kvcache/hybrid_radix_cache.py index 1bd9e0b9cd..0b66253b17 100644 --- a/python/freetoken/kvcache/hybrid_radix_cache.py +++ b/python/freetoken/kvcache/hybrid_radix_cache.py @@ -98,7 +98,10 @@ def insert(self, input_ids: torch.Tensor, kv_indices: torch.Tensor, node, prefix_len = self._walk(input_ids) if prefix_len != insert_len: new_node = RadixTreeNode(self.key_fn) - new_node.set_key_value(input_ids[prefix_len:], kv_indices[prefix_len:].clone()) + # A view would keep the request's full output-budget token buffer alive. + new_node.set_key_value( + input_ids[prefix_len:].clone(), kv_indices[prefix_len:].clone() + ) new_node.set_parent(node) self.full_evictable += new_node.length node = new_node diff --git a/python/freetoken/kvcache/radix_cache.py b/python/freetoken/kvcache/radix_cache.py index 5437286d12..ec0bf75d35 100644 --- a/python/freetoken/kvcache/radix_cache.py +++ b/python/freetoken/kvcache/radix_cache.py @@ -167,7 +167,10 @@ def insert_prefix(self, input_ids: torch.Tensor, indices: torch.Tensor) -> Inser node, prefix_len = self._tree_walk(input_ids) if prefix_len != insert_len: # NOTE: prefix_len < insert_len new_node = RadixTreeNode(self.key_fn) - new_node.set_key_value(input_ids[prefix_len:], indices[prefix_len:].clone()) + # A view would keep the request's full output-budget token buffer alive. + new_node.set_key_value( + input_ids[prefix_len:].clone(), indices[prefix_len:].clone() + ) new_node.set_parent(node) self.evictable_size += new_node.length node = new_node diff --git a/python/freetoken/kvcache/swa_radix_cache.py b/python/freetoken/kvcache/swa_radix_cache.py index 260439ad66..d2cc89c26f 100644 --- a/python/freetoken/kvcache/swa_radix_cache.py +++ b/python/freetoken/kvcache/swa_radix_cache.py @@ -215,7 +215,8 @@ def insert(self, input_ids: torch.Tensor, kv_indices: torch.Tensor, # then a live (non-tombstone) leaf for the in-window remainder (never a tombstone leaf -- # the clamp + the free_swa -page_size margin guarantee a live tail). if total < insert_len: - suffix_ids = input_ids[total:] + # Both suffix nodes can share this copy without retaining the request buffer. + suffix_ids = input_ids[total:].clone() suffix_kv = kv_indices[total:].clone() boundary = max(0, min(swa_evicted_seqlen, insert_len) - total) boundary = min(boundary, max(0, len(suffix_ids) - self.page_size)) From c1a504a50b6ced2c28c5628d29df9885f0f3ae74 Mon Sep 17 00:00:00 2001 From: SuperHotDogCat <115141367+SuperHotDogCat@users.noreply.github.com> Date: Wed, 7 Oct 2026 11:54:38 +0900 Subject: [PATCH 05/73] perf(moe): Improved `fast_index_copy_multi_jit` for NVLink-C2C hosts (#610) * perf(moe): default fused gather to 32 blocks/bank on NVLink-C2C hosts * build(aot): prebuild fast_index_copy_multi (1024, 32) variant * perf(moe): make fused H2D gather grid configurable via env var * docs: document FREETOKEN_H2D_BLOCKS_PER_BANK for NVLink-C2C hosts * fix(moe): auto-detect H2D gather grid per device in the copy plan * chore(docs): revised docs for NVLinkC2C --- docs/cli.md | 3 +++ docs/models.md | 14 ++++++++++++++ python/freetoken/kernel/aot.py | 1 + python/freetoken/kernel/fast_index_copy.py | 17 +++++++++++++++++ python/freetoken/moe/offload_cache.py | 4 ++++ 5 files changed, 39 insertions(+) diff --git a/docs/cli.md b/docs/cli.md index b3cec99826..b9be878ce7 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -204,3 +204,6 @@ profile that `ft serve --moe-strategy auto` and `--moe-hybrid-max-fetch -1` then - What to measure: `--dtype`, `--model`, `--formats`, `--isa`. - `--threshold` (default 2.0) sets the call: recommend hybrid when CPU bandwidth beats PCIe by that factor. + +On NVLink-C2C hosts, see [models.md](models.md#nvlink-c2c-hosts-gh200--gb200) +for `FREETOKEN_H2D_BLOCKS_PER_BANK` if `PCIe-gather` sits well below the linear H2D ceiling. diff --git a/docs/models.md b/docs/models.md index efd5268b37..2cd3baa317 100644 --- a/docs/models.md +++ b/docs/models.md @@ -59,6 +59,20 @@ have no resident path and stay on `offload`. Offload-only flags `--moe-prefill-hit-d2d`, `--disable-moe-prefill-overlap`) are ignored with a warning; pass `--moe-strategy offload` to use them. +### NVLink-C2C hosts (GH200 / GB200) + +The host link is ~450 GB/s per direction instead of PCIe's ~32-64 GB/s. The +offload expert gather reads host memory zero-copy and is latency-bound, so it +needs a wider grid to keep enough loads in flight. Set +`FREETOKEN_H2D_BLOCKS_PER_BANK`: +``` + FREETOKEN_H2D_BLOCKS_PER_BANK=32 ft serve --model --moe-strategy offload +``` +On GH200 this raises the gather from ~220 to ~410 GB/s (`PCIe-gather` in +`ft bench bw`), matching the DMA ceiling. Run `ft bench bw` with the variable +set so the hybrid split is calibrated against the faster gather. Leave it +unset on PCIe GPUs: wider grids add no bandwidth there. + ## Notes - `ft checkpoint` conversion is optional — it pre-converts a checkpoint into diff --git a/python/freetoken/kernel/aot.py b/python/freetoken/kernel/aot.py index 5e87f8c923..bbc2d2c669 100644 --- a/python/freetoken/kernel/aot.py +++ b/python/freetoken/kernel/aot.py @@ -153,6 +153,7 @@ def default_kernel_specs() -> tuple[KernelSpec, ...]: for feature_size in DEFAULT_FAST_INDEX_COPY_FEATURE_SIZES ) specs.append(_fast_index_copy_multi_spec(num_threads=1024, blocks_per_bank=8)) + specs.append(_fast_index_copy_multi_spec(num_threads=1024, blocks_per_bank=32)) # prefill hit-D2D gather (HBM-bound: wide grid) + its miss-side batch H2D binding. specs.append(_fast_index_copy_multi_spec(num_threads=1024, blocks_per_bank=64)) specs.append(_batch_memcpy_spec()) diff --git a/python/freetoken/kernel/fast_index_copy.py b/python/freetoken/kernel/fast_index_copy.py index 1aaa1303d2..b2030620ee 100644 --- a/python/freetoken/kernel/fast_index_copy.py +++ b/python/freetoken/kernel/fast_index_copy.py @@ -149,6 +149,23 @@ def _jit_fast_index_copy_multi_module(*, num_threads: int, blocks_per_bank: int) cuda_wrappers=[("launch", f"&MultiIndexCopyKernel<{args}>::run")], ) +def default_blocks_per_bank(device: torch.device) -> int: + """H2D gather grid width per bank. The link's bandwidth-delay product sets how many + 16B loads must be in flight: PCIe (~31 GB/s) saturates at ~4K threads/bank, NVLink-C2C + (GH200, ~450 GB/s) needs ~32K (measured: 8x1024 -> 222, 32x1024 -> 412 GB/s).""" + env = os.environ.get("FREETOKEN_H2D_BLOCKS_PER_BANK") + if env is not None: + blocks_per_bank = int( + env + ) + assert blocks_per_bank > 0, ( + "FREETOKEN_H2D_BLOCKS_PER_BANK must be greater than 0, " + f"got {blocks_per_bank}" + ) + return blocks_per_bank # manual override (validate > 0) + if any(k in torch.cuda.get_device_name(device) for k in ("GH200", "GB200", "GB300")): + return 32 + return 8 def fast_index_copy_multi_jit( dst_ptrs: torch.Tensor, diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index c6f975df7e..0969f73829 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -25,6 +25,7 @@ # entry the batch sees is >= this size. _SMALL_BANK_FEAT_BYTES = 256 * 1024 +from freetoken.kernel.fast_index_copy import default_blocks_per_bank from freetoken.utils import init_logger logger = init_logger(__name__) @@ -257,6 +258,7 @@ def __post_init__(self) -> None: # Source pointers are per layer (_copy_src_ptrs[layer_id] -> [num_banks] device # tensor); dst/feat are layer-invariant. self._copy_fused_ok = False + self._copy_blocks_per_bank = 8 self._copy_dst_ptrs: torch.Tensor | None = None self._copy_src_ptrs: list[torch.Tensor] | None = None self._copy_feat_bytes: torch.Tensor | None = None @@ -438,6 +440,7 @@ def _build_fused_copy_plan(self) -> None: elif self._gather_bank_ids: self._gather_dst_ptrs = self._copy_dst_ptrs[self._gather_bank_ids].contiguous() self._gather_feat_bytes = self._copy_feat_bytes[self._gather_bank_ids].contiguous() + self._copy_blocks_per_bank = default_blocks_per_bank(self.device) self._copy_fused_ok = True def validate_rebuild(self, cache_size: int) -> None: @@ -1038,6 +1041,7 @@ def copy_missing(self) -> None: self.evict_slots, self.src_indices, self.num_indices, + blocks_per_bank=self._copy_blocks_per_bank, ) return From 1c183e74ae39b6f2708394056c00dda556f89bea Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Tue, 6 Oct 2026 21:38:10 -0700 Subject: [PATCH 06/73] perf(moe): spin-wait the cpu moe pool between layers (#619) Co-authored-by: shubham3-ucb <229030647+shubham3-ucb@users.noreply.github.com> --- .../kernel/csrc/cpu_moe/cpu_moe_ext.cpp | 71 +++++++++++++++---- python/freetoken/moe/cpu_executor.py | 18 ++++- 2 files changed, 75 insertions(+), 14 deletions(-) diff --git a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp index e2cb0d9c33..fb6a33f497 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp @@ -49,6 +49,14 @@ namespace { +inline void cpu_relax() { +#if CPU_MOE_X86 + _mm_pause(); +#elif defined(__aarch64__) + asm volatile("yield" ::: "memory"); +#endif +} + using bf16_t = uint16_t; inline float bf16_to_f32(bf16_t v) { @@ -1290,13 +1298,19 @@ struct CpuMoeExecutor { std::mutex sync_mtx; std::condition_variable sync_cv; - bool stop = false; + std::atomic stop{false}; // spinning workers read it without the lock uint64_t cur_gen = 0; MoeTask* cur_task = nullptr; - std::atomic submitted{0}; - std::atomic completed{0}; + // Separate cache lines: idle workers poll `submitted` while busy workers hit p1_next/p2_next. + alignas(64) std::atomic submitted{0}; + alignas(64) std::atomic completed{0}; + + // Spin-then-park waits. Off unless the Python wrapper finds spare CPUs for the spinners. + static constexpr auto kSpinWindow = std::chrono::milliseconds(50); + std::atomic spin_wait{false}; + bool coord_pinned = false; - std::atomic p1_next{0}; + alignas(64) std::atomic p1_next{0}; std::atomic p2_next{0}; std::atomic prt_next{0}; // ds_fp4 intermediate fp8 round-trip phase int64_t p1_total = 0, p2_total = 0, prt_total = 0; @@ -1866,19 +1880,34 @@ struct CpuMoeExecutor { } } + // Tasks arrive once per MoE layer, sub-ms apart. Spinning skips a futex wakeup per worker per layer; parking after kSpinWindow keeps an idle server off the CPU. void worker_loop(int tid) { pin_self(tid); uint64_t my_gen = 0; + using wclk = std::chrono::steady_clock; + auto last_active = wclk::now(); for (;;) { - MoeTask* t; - { + uint64_t gen = submitted.load(std::memory_order_acquire); + if (gen == my_gen && spin_wait.load(std::memory_order_relaxed)) { + unsigned polls = 0; + while (gen == my_gen) { + cpu_relax(); + if (stop.load(std::memory_order_relaxed)) return; + if ((++polls & 1023u) == 0 && wclk::now() - last_active > kSpinWindow) break; + gen = submitted.load(std::memory_order_acquire); + } + } + if (gen == my_gen) { std::unique_lock lk(task_mtx); - task_cv.wait(lk, [&] { return stop || cur_gen != my_gen; }); - if (stop) return; - my_gen = cur_gen; - t = cur_task; + task_cv.wait(lk, [&] { return stop.load() || cur_gen != my_gen; }); + if (stop.load()) return; + gen = cur_gen; } + // submit() writes cur_task before the release store to `submitted`. + my_gen = gen; + MoeTask* t = cur_task; run_task_body(t); + last_active = wclk::now(); if (done_count.fetch_add(1) + 1 == num_threads) { completed.store(my_gen, std::memory_order_release); { @@ -1961,8 +1990,19 @@ struct CpuMoeExecutor { task_cv.notify_all(); } - void sync() { + // Only a coordinator on its own core may spin: host-func callbacks and eager calls share cores with the workers. + void sync(bool spin = false) { const uint64_t target = submitted.load(std::memory_order_acquire); + if (spin) { + using sclk = std::chrono::steady_clock; + const auto t0 = sclk::now(); + unsigned polls = 0; + while (completed.load(std::memory_order_acquire) < target) { + cpu_relax(); + if ((++polls & 1023u) == 0 && sclk::now() - t0 > kSpinWindow) break; + } + if (completed.load(std::memory_order_acquire) >= target) return; + } std::unique_lock lk(sync_mtx); sync_cv.wait(lk, [&] { return completed.load(std::memory_order_acquire) >= target; }); } @@ -2003,6 +2043,7 @@ struct CpuMoeExecutor { } flag_served.assign(num_slots, 0); coord_stop.store(false); + coord_pinned = CPU_MOE_HAS_AFFINITY && pin_core >= 0; coord_thread = std::thread([this, pin_core] { #if CPU_MOE_HAS_AFFINITY if (pin_core >= 0) { @@ -2050,7 +2091,7 @@ struct CpuMoeExecutor { } if (t != nullptr) { submit(t); - sync(); + sync(coord_pinned && spin_wait.load(std::memory_order_relaxed)); } // Release: the workers' y stores are visible before the GPU sees done. flag_store_release(&done_flags[L], 1); @@ -2144,7 +2185,11 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { .def("set_input_prequant", [](CpuMoeExecutor& e, bool v) { e.input_prequant = v; }, py::arg("value")) - .def("isa_name", &CpuMoeExecutor::isa_name); + .def("isa_name", &CpuMoeExecutor::isa_name) + .def("set_spin_wait", + [](CpuMoeExecutor& e, bool v) { e.spin_wait.store(v, std::memory_order_relaxed); }, + py::arg("enabled")) + .def("get_spin_wait", [](const CpuMoeExecutor& e) { return e.spin_wait.load(); }); m.def("memops_probe", &cumemops_probe, py::arg("stream"), py::arg("scratch_addr")); m.def("memop_submit", &cumemop_submit, py::arg("stream"), py::arg("done_addr"), py::arg("ready_addr"), py::arg("slot")); diff --git a/python/freetoken/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py index e2e48d1461..a46806740d 100644 --- a/python/freetoken/moe/cpu_executor.py +++ b/python/freetoken/moe/cpu_executor.py @@ -53,6 +53,8 @@ # sizes plus any eager padded sizes); more than that is unheard of, and the overflow # just keeps the host-func path for the extra combos. _FLAG_SLOTS_PER_LAYER = 16 +# Pool threads spin up to 50 ms before they park. FREETOKEN_CPU_MOE_SPIN=0 opts out. +_SPIN_WAIT = os.getenv("FREETOKEN_CPU_MOE_SPIN", "1") != "0" # Activation ids must match ActKind in csrc/cpu_moe/cpu_moe_ext.cpp. Id 3 is the # clamped (up + 1) swiglu: "swigluoai" runs it in the generic GEMV epilogue, @@ -116,6 +118,17 @@ def physical_core_cpus() -> list[int]: return reps or allowed or [0] +def _spin_wait_ok(core_ids: list[int], coord_core: int) -> bool: + """Spin only if each pool thread owns a distinct CPU and two CPUs stay free for the main and CUDA callback threads.""" + if not _SPIN_WAIT or not hasattr(os, "sched_getaffinity"): + return False + pool = list(core_ids) + ([coord_core] if coord_core >= 0 else []) + if len(set(pool)) != len(pool): + return False + spare = set(os.sched_getaffinity(0)) - set(pool) + return len(spare) >= 2 + + def resolve_threads_and_affinity(requested: int) -> tuple[int, list[int]]: """Return (num_threads, core_ids) for the worker pool. @@ -246,6 +259,8 @@ def __init__( self.num_threads = nthreads self.core_ids = core_ids self.isa = self._ext.isa_name() + self.spin_wait = _spin_wait_ok(core_ids, coord_core) + self._ext.set_spin_wait(self.spin_wait) spare = len(physical_core_cpus()) - nthreads - (1 if coord_core >= 0 else 0) - 1 clamp = max(1, min(torch.get_num_threads(), spare)) @@ -319,7 +334,8 @@ def __init__( f"CPU MoE executor ready: threads={nthreads} (pinned to cores " f"{core_ids[0]}..{core_ids[-1]}) isa={self.isa} fmt={fmt} " f"H={self.H} I={self.I} experts={self.num_experts} layers={self.num_layers} " - f"top_k={self.top_k} act={activation} max_tokens={self.max_tokens}" + f"top_k={self.top_k} act={activation} max_tokens={self.max_tokens} " + f"spin_wait={self.spin_wait}" ) def _make_table(self, layers: list[torch.Tensor]) -> torch.Tensor: From 127c85665b49e39d78a5d09b95a92169b5794eb0 Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Wed, 7 Oct 2026 00:15:11 -0700 Subject: [PATCH 07/73] feat(deepseek-v41): support text and vision serving (#603) * test(deepseek-v41): add independent text and vision reference * feat(quant): support block-32 and unrounded inference outputs * feat(attention): add csa2 packed cache and sparse backend * feat(deepseek-v41): serve text and images through shared runtime * perf(mhc): tune batch-stable reduction tiles * perf(moe): autotune fp4 decode launch geometry * perf(moe): admit recurring misses against actual lru victims * revert: perf(moe): admit recurring misses against actual lru victims This reverts commit 1652d200f0ebb825379cf44c28baf52971e3a461. * perf(mhc): derive the reduction tile from the width, not a timing * perf(moe): replace the ds_fp4 decode autotune with a fixed two-warp launch * fix(deepseek-v41): fence engram staging against in-flight lookups * fix(moe): return the cpu routed sum in fp32 and round it in python * fix(deepseek-v41): keep the vision tower and aligner unquantized * fix(scheduler): admit bounded-replay hits by the history their resume reads * refactor(deepseek-v41): name the attention family after the model, not csa2 * refactor(kvcache): import the dsv41 pool only for dsv41 models * refactor(deepseek-v41): name the replay switch --swa-decoder-replay * fix(server): keep reasoning_effort a string on the wire * refactor(models): keep the prefill replay hooks off BaseLLMModel * refactor(deepseek-v41): drop the fp32 routed sum and batch-invariant prefill * fix(moe): weight ds_fp4 routes in the down epilogue again * docs(models): anchor the v4.1 replay paragraph on the tech report * refactor(deepseek-v41): run the decoder source's narrowed prefill on Block * refactor(kvcache): let the dsv41 pool own the history a resume keeps live * refactor(deepseek-v41): resume bounded-replay hits by capping the prefix match * refactor(deepseek-v41): describe the attention group as a dsv4 variant and derive its geometry from the args * refactor(deepseek-v41): keep the engram table locator in the model package * refactor(deepseek-v41): drop the fp32 lm head output - return the head's bf16 logits, as DeepSeek-V4 does - revert ParallelLMHead(out_dtype=) and LinearKernel.apply_out_dtype, which only this call used - remove test_head_preserves_fp32_logit_margin - cast test_vision's prefill logits to fp32 before comparing, as the harness does * refactor(quant): carry the w4a8 activation block on MoEConfig - drop ActDesc and QuantScheme.act; the scheme only describes stored tensors - Fp8BlockConfig sets MoEConfig.act_block from weight_block_size for its fp4 experts via a layer_config hook - the ds_fp4 GPU kernel and CPU executor read the block from the layer config and require it * refactor(deepseek-v41): size the kv pool from the args, not a geometry copy - delete kvcache/dsv41_geometry.py; the pool and cost model read dsv41_args by attribute, as the DSV4 pool reads dsv4_args - source / ratio aliasing reads args.roles; the duplicated source and ratio validation is gone - the cost model owns the tier row formats, row bytes and the private window layers - the pool keeps the format and window/ratio checks, and exposes the formats and is_private_window to the attention backend * refactor(kvcache): move the dsv4 pool into kvcache/dsv4 * refactor(kvcache): move the dsv41 pool into kvcache/dsv4 * chore(benchmarks): drop the dsv41 selection and fp8 linear benchmarks * test(deepseek-v41): drop tests that repeat other coverage * test(deepseek-v41): load the tiny model's resident experts as banks * fix(server): parse deepseek-v4.1's space-led dsml tool calls --------- Co-authored-by: andy-yang-1 --- docs/models.md | 32 + python/freetoken/attention/__init__.py | 10 + python/freetoken/attention/base.py | 3 + python/freetoken/attention/dsv41_sparse.py | 438 +++++ python/freetoken/core.py | 3 + python/freetoken/engine/cache_budget.py | 2 +- python/freetoken/engine/config.py | 3 + python/freetoken/engine/engine.py | 44 +- python/freetoken/kernel/aot_models.py | 11 + .../kernel/csrc/cpu_moe/cpu_moe_ext.cpp | 20 +- .../kernel/triton/dsv4/fp8_linear.py | 321 ++-- .../freetoken/kernel/triton/dsv4/fused_moe.py | 35 +- python/freetoken/kernel/triton/dsv4/router.py | 83 + .../freetoken/kernel/triton/dsv41/__init__.py | 9 + .../freetoken/kernel/triton/dsv41/indexer.py | 176 ++ python/freetoken/kernel/triton/dsv41/pack.py | 183 +++ .../kernel/triton/dsv41/row_format.py | 92 ++ .../kernel/triton/dsv41/sparse_attn.py | 323 ++++ python/freetoken/kernel/triton/dsv41/topk.py | 265 +++ python/freetoken/kernel/triton/e4m3_compat.py | 14 + python/freetoken/kernel/triton/mhc.py | 185 ++- python/freetoken/kvcache/__init__.py | 27 +- python/freetoken/kvcache/base.py | 17 + python/freetoken/kvcache/cache_status.py | 44 +- python/freetoken/kvcache/dsa_pool.py | 4 +- python/freetoken/kvcache/dsv4/__init__.py | 1 + .../freetoken/kvcache/dsv4/v41_cost_model.py | 216 +++ python/freetoken/kvcache/dsv4/v41_pool.py | 292 ++++ .../freetoken/kvcache/dsv4/v41_row_format.py | 69 + .../v4_cost_model.py} | 14 +- .../{dsv4_paged_pool.py => dsv4/v4_pool.py} | 313 +--- python/freetoken/kvcache/dsv4/window_tier.py | 331 ++++ python/freetoken/kvcache/hybrid_swa_pool.py | 12 +- python/freetoken/layers/mhc.py | 119 +- python/freetoken/layers/moe.py | 23 +- .../layers/quantization/configs/base.py | 6 +- .../layers/quantization/configs/fp8.py | 34 +- .../layers/quantization/linear/fp8_block.py | 24 +- .../freetoken/layers/quantization/moe/base.py | 6 + .../layers/quantization/moe/fp8_block.py | 4 +- .../layers/quantization/moe/mxfp4.py | 11 +- .../freetoken/layers/quantization/scheme.py | 16 +- .../freetoken/mm/processors/deepseek_v41.py | 88 + python/freetoken/models/config.py | 12 +- .../freetoken/models/deepseek_v4/__init__.py | 2 +- python/freetoken/models/deepseek_v4/model.py | 2 +- .../freetoken/models/deepseek_v41/__init__.py | 14 + python/freetoken/models/deepseek_v41/args.py | 225 +++ .../models/deepseek_v41/attention.py | 305 ++++ .../models/deepseek_v41/compressor.py | 131 ++ .../freetoken/models/deepseek_v41/config.py | 95 ++ .../freetoken/models/deepseek_v41/engram.py | 262 +++ .../models/deepseek_v41/engram_table.py | 296 ++++ .../freetoken/models/deepseek_v41/indexer.py | 127 ++ python/freetoken/models/deepseek_v41/model.py | 289 ++++ python/freetoken/models/deepseek_v41/moe.py | 90 + python/freetoken/models/deepseek_v41/rope.py | 58 + .../freetoken/models/deepseek_v41/vision.py | 139 ++ .../freetoken/models/deepseek_v41/weight.py | 242 +++ python/freetoken/models/register.py | 14 + python/freetoken/moe/benchbw.py | 8 +- python/freetoken/moe/cpu_executor.py | 12 +- python/freetoken/moe/fused_ds_fp4.py | 50 +- python/freetoken/scheduler/cache.py | 9 +- python/freetoken/scheduler/prefill.py | 6 +- python/freetoken/scheduler/scheduler.py | 17 +- python/freetoken/server/api_server.py | 27 +- python/freetoken/server/args.py | 14 + .../freetoken/server/function_call_parser.py | 49 +- python/freetoken/server/stats.py | 10 +- python/freetoken/tokenizer/effort.py | 35 +- python/freetoken/tokenizer/tokenize.py | 34 +- scripts/ftw_hotfix.py | 4 +- tests/README.md | 1 + tests/attention/test_dsv41_backend.py | 208 +++ tests/checkpoint/test_ftw_hotfix_vision.py | 6 +- tests/dsv4/test_dsv4_attn_backend.py | 4 +- tests/kernels/test_dsv41_indexer.py | 139 ++ tests/kernels/test_dsv41_pack.py | 112 ++ tests/kernels/test_dsv41_sparse_attn.py | 107 ++ tests/kernels/test_dsv41_topk.py | 159 ++ tests/kernels/test_dsv4_fp8_block_linear.py | 101 ++ tests/kernels/test_dsv4_router_select.py | 52 + tests/kvcache/test_cache_unit_bytes.py | 4 +- tests/kvcache/test_dsv41_pool.py | 205 +++ tests/kvcache/test_dsv4_free_list.py | 2 +- tests/kvcache/test_dsv4_pool.py | 4 +- tests/kvcache/test_kv_cache_rebuild.py | 12 +- tests/kvcache/test_pool_sizing_surface.py | 18 +- tests/layers/test_mhc.py | 97 ++ tests/models/deepseek_v41/__init__.py | 0 tests/models/deepseek_v41/common.py | 176 ++ tests/models/deepseek_v41/harness.py | 182 ++ tests/models/deepseek_v41/reference/LICENSE | 21 + .../models/deepseek_v41/reference/__init__.py | 2 + tests/models/deepseek_v41/reference/engram.py | 186 +++ tests/models/deepseek_v41/reference/kernel.py | 113 ++ tests/models/deepseek_v41/reference/model.py | 1458 +++++++++++++++++ tests/models/deepseek_v41/reference/vision.py | 121 ++ tests/models/deepseek_v41/test_config.py | 169 ++ .../models/deepseek_v41/test_engine_config.py | 81 + tests/models/deepseek_v41/test_engram.py | 281 ++++ .../models/deepseek_v41/test_forward_smoke.py | 184 +++ .../deepseek_v41/test_reference_parity.py | 219 +++ tests/models/deepseek_v41/test_vision.py | 214 +++ tests/models/test_quant_config.py | 26 +- tests/moe/test_cpu_moe.py | 3 +- tests/moe/test_dsfp4_prequant.py | 2 +- tests/moe/test_dsfp4_reference_semantics.py | 170 ++ tests/scheduler/test_dsv4_generic_manager.py | 31 +- tests/scheduler/test_scheduler_kv_usage.py | 11 +- tests/server/test_effort_dialect.py | 16 + tests/server/test_function_call_parser.py | 28 + tests/server/test_parser_auto_selection.py | 7 + tests/server/test_streaming_model_matrix.py | 7 + tests/tokenizer/test_effort.py | 36 + tests/tokenizer/test_tokenize.py | 55 + 117 files changed, 10881 insertions(+), 650 deletions(-) create mode 100644 python/freetoken/attention/dsv41_sparse.py create mode 100644 python/freetoken/kernel/triton/dsv4/router.py create mode 100644 python/freetoken/kernel/triton/dsv41/__init__.py create mode 100644 python/freetoken/kernel/triton/dsv41/indexer.py create mode 100644 python/freetoken/kernel/triton/dsv41/pack.py create mode 100644 python/freetoken/kernel/triton/dsv41/row_format.py create mode 100644 python/freetoken/kernel/triton/dsv41/sparse_attn.py create mode 100644 python/freetoken/kernel/triton/dsv41/topk.py create mode 100644 python/freetoken/kvcache/dsv4/__init__.py create mode 100644 python/freetoken/kvcache/dsv4/v41_cost_model.py create mode 100644 python/freetoken/kvcache/dsv4/v41_pool.py create mode 100644 python/freetoken/kvcache/dsv4/v41_row_format.py rename python/freetoken/kvcache/{dsv4_cost_model.py => dsv4/v4_cost_model.py} (96%) rename python/freetoken/kvcache/{dsv4_paged_pool.py => dsv4/v4_pool.py} (52%) create mode 100644 python/freetoken/kvcache/dsv4/window_tier.py create mode 100644 python/freetoken/mm/processors/deepseek_v41.py create mode 100644 python/freetoken/models/deepseek_v41/__init__.py create mode 100644 python/freetoken/models/deepseek_v41/args.py create mode 100644 python/freetoken/models/deepseek_v41/attention.py create mode 100644 python/freetoken/models/deepseek_v41/compressor.py create mode 100644 python/freetoken/models/deepseek_v41/config.py create mode 100644 python/freetoken/models/deepseek_v41/engram.py create mode 100644 python/freetoken/models/deepseek_v41/engram_table.py create mode 100644 python/freetoken/models/deepseek_v41/indexer.py create mode 100644 python/freetoken/models/deepseek_v41/model.py create mode 100644 python/freetoken/models/deepseek_v41/moe.py create mode 100644 python/freetoken/models/deepseek_v41/rope.py create mode 100644 python/freetoken/models/deepseek_v41/vision.py create mode 100644 python/freetoken/models/deepseek_v41/weight.py create mode 100644 tests/attention/test_dsv41_backend.py create mode 100644 tests/kernels/test_dsv41_indexer.py create mode 100644 tests/kernels/test_dsv41_pack.py create mode 100644 tests/kernels/test_dsv41_sparse_attn.py create mode 100644 tests/kernels/test_dsv41_topk.py create mode 100644 tests/kernels/test_dsv4_fp8_block_linear.py create mode 100644 tests/kernels/test_dsv4_router_select.py create mode 100644 tests/kvcache/test_dsv41_pool.py create mode 100644 tests/models/deepseek_v41/__init__.py create mode 100644 tests/models/deepseek_v41/common.py create mode 100644 tests/models/deepseek_v41/harness.py create mode 100644 tests/models/deepseek_v41/reference/LICENSE create mode 100644 tests/models/deepseek_v41/reference/__init__.py create mode 100644 tests/models/deepseek_v41/reference/engram.py create mode 100644 tests/models/deepseek_v41/reference/kernel.py create mode 100644 tests/models/deepseek_v41/reference/model.py create mode 100644 tests/models/deepseek_v41/reference/vision.py create mode 100644 tests/models/deepseek_v41/test_config.py create mode 100644 tests/models/deepseek_v41/test_engine_config.py create mode 100644 tests/models/deepseek_v41/test_engram.py create mode 100644 tests/models/deepseek_v41/test_forward_smoke.py create mode 100644 tests/models/deepseek_v41/test_reference_parity.py create mode 100644 tests/models/deepseek_v41/test_vision.py create mode 100644 tests/moe/test_dsfp4_reference_semantics.py diff --git a/docs/models.md b/docs/models.md index 2cd3baa317..b999d29e41 100644 --- a/docs/models.md +++ b/docs/models.md @@ -5,6 +5,7 @@ for them; other checkpoints of the same architectures work too. | Model | HF checkpoints | |---|---| +| DeepSeek-V4.1-Flash | [deepseek-ai/DeepSeek-V4.1-Flash](https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash) | | DeepSeek-V4 | [deepseek-ai/DeepSeek-V4-Flash-0731](https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash-0731) | | GLM-5.3-Flash | [RedHatAI/GLM-5.3-Flash-NVFP4](https://huggingface.co/RedHatAI/GLM-5.3-Flash-NVFP4) | | GLM-5.2 | [nvidia/GLM-5.2-NVFP4](https://huggingface.co/nvidia/GLM-5.2-NVFP4) | @@ -27,6 +28,7 @@ These families accept image input by default; pass `--text-model-only` to skip t | Family | Image tokens | `--image-min-tokens` / `--image-max-tokens` | `--mm-processor-kwargs` example | | --- | --- | --- | --- | +| DeepSeek-V4.1-Flash (ViT tower, streamed under `--mm-encoder-weights host`) | one feature per 42x42 pixels, plus a learned row break per row and two delimiters; checkpoint cap 1024 total span tokens | maximum caps the entire span; minimum maps to a pixel-area floor before the cap; checkpoint minimum area is 544x544 | `{"max_image_tokens": 512, "min_pixels": 295936}` | | Qwen3.6 (both variants, every listed weight format), Qwen3.8-Flash-Next, Qwen3-VL | one token per 32x32 pixels of the resized image, dynamic resolution | pixel areas in `size.shortest_edge` / `longest_edge`; checkpoint defaults 64 to 16384 tokens | `{"size": {"longest_edge": 1048576}}` | | Gemma-4 26B-A4B, 31B (`gemma4`: ViT tower, streamed under `--mm-encoder-weights host`) | one of the soft-token budgets 70 / 140 / 280 / 560 / 1120, every image scaled to its budget as far as the aspect ratio allows | the maximum picks the largest budget within it, below 70 is refused at start-up; the minimum has no effect | `{"max_soft_tokens": 1120}` | | Gemma-4 12B (`gemma4_unified`: linear patch embedder, resident under either placement) | same budgets, one 48x48 super-patch per soft token | same as the tower releases | same | @@ -86,4 +88,34 @@ unset on PCIe GPUs: wider grids add no bandwidth there. with `ft checkpoint`, or add the encoder in place with [scripts/ftw_hotfix.py](ftw-hotfix.md). - DeepSeek-V4 checkpoints must keep the `inference/config.json` subdir — the authoritative model args are read from there. +- DeepSeek-V4.1-Flash serves text and images from the HF checkpoint as shipped (fp8 block-32 + dense weights, fp4 experts on the offload cache, 890 B/token global KV in packed + fp4/fp8 pools). Its two 98 GiB Engram tables stream from the checkpoint shards on + demand (keep them on a fast NVMe; the 6 GiB of table scales stay in host RAM), using + bounded pinned staging buffers; an FTW conversion copies the shards that hold them next + to the checkpoint. + SWA bounded replay (DeepSeek_V41_Tech_Report.pdf, shipped in the checkpoint, §3.2.2): a replayed + segment recomputes only the SWA KV of its last `n_win` tokens and truncates each query's window to + the segment. *Encoder* replay rebuilds the encoder SWA KV behind a prefix hit from the global KV + alone (approximate; the report's fallback when the SWA KV of a hit has been evicted). *Decoder* + replay runs the decoder layers on each prompt's last `n_win` tokens only; their SWA KV is never + prefix-cached and post-training simulated it. FreeToken keeps encoder SWA KV in the radix cache (no + encoder replay); `--swa-decoder-replay bounded` (default) is the report's decoder replay, with a + prefix hit stopping at least `n_win` tokens before the prompt end so those tokens are prefilled; + `exact` runs the decoder on every prompt token (reference numerics). + Bounded output differs from exact by construction and does not depend on the prefill chunk size. + Under expert offload the prefill time is bounded by streaming each layer's experts, so decoder + replay saves decoder-layer compute, not prefill time. + Run `ft bench bw --model dsv4.1-flash` to measure local PCIe and CPU bandwidth before + choosing a MoE strategy; `--moe-strategy hybrid` splits expert misses using those + measurements. With `--moe-cache-auto`, `--kv-reserve-tokens` sets the token capacity reserved + before the expert cache is allocated, subject to the attention pool's structural minimum; size it + for the intended context and concurrency. `--max-extend-length` controls the prefill chunk size + and its workspace. Image input uses the shared multimodal processor and encoder cache, including + chunked prefill and prefix replay; `--text-model-only` skips the vision weights. DSpark + speculative decoding is not served. + `reasoning_effort` takes `low`, `high` or `max`; the 1-100 integer budget goes through the + template kwargs with thinking on, `"chat_template_kwargs": {"enable_thinking": true, + "reasoning_effort": 37}`, and the checkpoint's encoder validates it. Runtime window-cache + controls use 128-token pages. - Qwen3.8-Flash-Next keeps a 47.7 GiB PLE n-gram table pinned in host RAM. diff --git a/python/freetoken/attention/__init__.py b/python/freetoken/attention/__init__.py index 8a410ffe4a..eda3b61f55 100644 --- a/python/freetoken/attention/__init__.py +++ b/python/freetoken/attention/__init__.py @@ -102,6 +102,16 @@ def create_dsv4_sparse_backend(config: ModelConfig): return DSV4SparseAttnBackend(config) +@SUPPORTED_ATTENTION_BACKENDS.register( + "dsv41_sparse", + BackendInfo(supported_types=frozenset({AttnType.DSV41})), +) +def create_dsv41_sparse_backend(config: ModelConfig): + from .dsv41_sparse import DSV41SparseAttnBackend + + return DSV41SparseAttnBackend(config) + + @SUPPORTED_ATTENTION_BACKENDS.register( "dsa", BackendInfo(supported_types=frozenset({AttnType.MLA, AttnType.DSA})), diff --git a/python/freetoken/attention/base.py b/python/freetoken/attention/base.py index 9115b80929..846e99447a 100644 --- a/python/freetoken/attention/base.py +++ b/python/freetoken/attention/base.py @@ -20,6 +20,9 @@ class AttnType(str, Enum): MLA = "mla" # plain latent-KV MLA -> MLAKVCache DSA = "dsa" # latent-KV MLA + DSA sparse indexer -> DSAKVCache DSV4 = "dsv4" # DSV4 window+compressed sparse -> DSV4PagedKVCache + # DSV41 (DeepSeek-V4.1): window + cross-layer-shared compressed KV in packed fp8/fp4 + # rows, Full/Reindex/Reuse layer modes -> DSV41PagedKVCache + DSV41 = "dsv41" LINEAR = "linear" # GDN/mamba state layers -> LinearStatePool # GQA block-sparse (MiniMax-M3): paged GQA K/V + a per-sparse-layer index-key # slab; the indexer picks top-k 128-token blocks per query -> BSAKVCache diff --git a/python/freetoken/attention/dsv41_sparse.py b/python/freetoken/attention/dsv41_sparse.py new file mode 100644 index 0000000000..baaa2cf019 --- /dev/null +++ b/python/freetoken/attention/dsv41_sparse.py @@ -0,0 +1,438 @@ +"""DeepSeek-V4.1 sparse-attention backend. + +Same contract and division of labour as ``dsv4_sparse``: the backend owns the per-forward KV +ADDRESSING over the shared page table (window ring slots, compressed rows, compress-state carry, +the decode snapshot and the CUDA-graph staging buffers) plus the selection helpers that are pure +index arithmetic (causal top-k, the hierarchical candidate pool); the model computes projections +and hands the picks back to be resolved into pool rows. + +What DSV41 adds on top of DSV4's vocabulary: + +* **cross-layer sharing** -- the compressed tiers belong to the kv-source layer; a consumer layer + addresses ``source_of(layer)``'s pools. The Top-K rows and the candidate pool the Full / Reindex + layers produce ride on the per-forward ``SharedSelection`` for the Reuse layers to read. +* **packed rows** -- writes quantize into the pool (``pool.store_*``), reads dequantize in-kernel. +* **prefill passes with a window floor** -- a ``PrefillSegment`` carries ``window_floor``: the + absolute position a query's sliding window may not reach below (Decoder SWA Bounded Replay + truncates the decoder's window at ``L - W``, the start of the prompt's last window). The metadata + carries the encoder pass (every new token) and the decoder pass (the new tokens inside the + prompt's last window, or every new token in exact mode). +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, List + +import torch +from freetoken.core import Batch, get_global_ctx + +from .base import AttentionSpec, BaseAttnBackend, BaseAttnMetadata + +if TYPE_CHECKING: + from freetoken.models import ModelConfig + + +@dataclass(frozen=True) +class PrefillSegment: + """One request's slice of a flat prefill token stream. + + ``offset``/``n`` tile the stream; ``start_pos`` is the absolute position of the first token; + ``window_floor`` bounds every query's sliding window from below (0 = the exact window; bounded + replay sets it to the start of the prompt's last window). + """ + + offset: int + n: int + table_idx: int + start_pos: int + window_floor: int = 0 + + @property + def end(self) -> int: + return self.start_pos + self.n + + +@dataclass +class SharedSelection: + """What the index-producing layers hand down the stack within ONE forward. + + ``topk_rows`` are GLOBAL main-pool rows (``-1`` masked) for the current kv source, laid out + exactly as the attention kernel consumes them; ``cmp_counts`` is the decode-only per-row live + count that bounds the kernel's loop; ``candidates`` is the Hierarchical Sparse Indexer's + compressed-position pool (``-1`` = empty slot). Layers run in order and every producer writes + before its consumers read, so one slot each suffices. + """ + + source: int | None = None + topk_rows: torch.Tensor | None = None + cmp_counts: torch.Tensor | None = None + candidates: torch.Tensor | None = None + + +@dataclass +class DSV41AttnMetadata(BaseAttnMetadata): + last_indices: torch.Tensor + # Prefill: the encoder pass tiles the batch's new tokens; the decoder pass tiles the tokens the + # decoder layers process (``decoder_rows`` gathers them out of the encoder stream; None = all of + # it). ``decoder_pad``: some request has no decoder row this chunk, and its ``last_indices`` entry + # points at one placeholder row the model appends after the decoder stream (its logits are unused). + segments: List[PrefillSegment] | None = None + decoder_segments: List[PrefillSegment] | None = None + decoder_rows: torch.Tensor | None = None + decoder_pad: bool = False + # Decode: the whole-history full-loc snapshot (captured buffer under a replay, lazy eager copy). + full_snap: torch.Tensor | None = None + table_rows: torch.Tensor | None = None + window_ar: torch.Tensor | None = None + selection: SharedSelection = field(default_factory=SharedSelection) + + def get_last_indices(self, bs: int) -> torch.Tensor: + return self.last_indices[:bs] + + @property + def stage_width(self) -> int: + assert self.full_snap is not None, "stage_width is capture/replay-only" + return self.full_snap.shape[1] + + def full_snapshot(self) -> torch.Tensor: + if self.full_snap is None: + assert self.table_rows is not None, "snapshot is decode-only" + pool = get_global_ctx().kv_cache + self.full_snap = pool.full_loc_map.index_select(0, self.table_rows).to(torch.int64) + return self.full_snap + + def window_ctx(self, pos: torch.Tensor, rows: torch.Tensor): + """Layer-invariant decode ring context ``(window_slots, prev_window_slots, + window_slots_topk [B, 1, win])``; computed fresh on every call (never cached -- a capture + must record these gathers), off the snapshot so a concurrent allocation cannot redirect + an in-flight replay.""" + snap = self.full_snapshot() + translate = get_global_ctx().kv_cache.translate_full_to_window + bs = pos.shape[0] + j = self.window_ar + assert j is not None, "window_ctx is decode-only" + win = j.shape[0] + window_slots = translate(snap[rows, pos]) + prev_window_slots = translate(snap[rows, (pos - 1).clamp_min(0)]) + p = pos[:, None] - ((pos[:, None] - j[None, :]) % win) + ws = translate(snap[rows[:, None], p.clamp(min=0)]) + window_slots_topk = torch.where((p >= 0) & (j[None, :] <= pos[:, None]), ws, -1).view(bs, 1, win) + return window_slots, prev_window_slots, window_slots_topk + + def private_window_ctx(self, pos: torch.Tensor, rows: torch.Tensor): + """The private-ring counterpart of ``window_ctx`` for request-private window layers: + ``(window_slots [B], window_slots_topk [B, 1, win])`` in the layer's ring, where a request's + row holds positions by ``pos % win`` (candidate ``j`` is the position congruent to ``j``).""" + assert self.table_rows is not None and self.window_ar is not None, "private_window_ctx is decode-only" + pool = get_global_ctx().kv_cache + bs = pos.shape[0] + j = self.window_ar + win = j.shape[0] + table = self.table_rows[rows] + window_slots = pool.ring_slots(table, pos) + p = pos[:, None] - ((pos[:, None] - j[None, :]) % win) # the position at ring index j (< 0: none yet) + window_slots_topk = pool.ring_slots(table[:, None], p).view(bs, 1, win) + return window_slots, window_slots_topk + + +@dataclass +class DSV41CaptureData: + full_snap: torch.Tensor + last_indices: torch.Tensor + table_rows: torch.Tensor # each batch row's page-table row (the private window rings key on it) + + @classmethod + def create(cls, max_bs: int, width: int, device: torch.device) -> "DSV41CaptureData": + return cls( + full_snap=torch.full((max_bs, width), -1, dtype=torch.int64, device=device), + last_indices=torch.arange(max_bs, dtype=torch.int32, device=device), + table_rows=torch.zeros(max_bs, dtype=torch.int64, device=device), + ) + + +class DSV41SparseAttnBackend(BaseAttnBackend): + def __init__(self, config: ModelConfig): + self.config = config + self.device = get_global_ctx().kv_cache.device + args = config.dsv41_args + self.window_size = args.window_size + self.swa_decoder_replay: str = args.swa_decoder_replay + self.capture: DSV41CaptureData | None = None + self.capture_bs: List[int] = [] + self.max_graph_bs = 0 + self._window_ar = torch.arange(self.window_size, device=self.device) + + @property + def pool(self): + return get_global_ctx().kv_cache + + # ----- generic contract ------------------------------------------------------------- + def forward(self, q, k, v, layer_id, batch, attn_spec: AttentionSpec | None = None): + raise NotImplementedError("DSV41 attention is driven per-tier from the model module; use DSV41SparseAttnBackend.attend().") + + def prepare_metadata(self, batch: Batch) -> None: + if not batch.is_decode: + batch.attn_metadata = self._prefill_metadata(batch) + return + last = torch.tensor([r.extend_len for r in batch.padded_reqs], dtype=torch.int32, device=self.device).cumsum_(0) - 1 + batch.attn_metadata = DSV41AttnMetadata(last_indices=last, table_rows=self._table_rows(batch), window_ar=self._window_ar) + + def _prefill_metadata(self, batch: Batch) -> DSV41AttnMetadata: + """The encoder pass runs every request's new tokens ``[cached_len, device_len)``. Exact mode runs + the decoder on the same stream. Decoder SWA Bounded Replay runs it on the new tokens inside the + prompt's last window ``[L - W, L)`` (``L = prompt_len``) with every query's window floored at + ``L - W``: the chunks of a chunked prefill each take their part of that window (earlier rows + come from the request's private ring), and a chunk that ends before it runs no decoder row.""" + segments: List[PrefillSegment] = [] + off = 0 + for r in batch.reqs: + segments.append(PrefillSegment(off, r.extend_len, r.table_idx, r.cached_len)) + off += r.extend_len + if self.swa_decoder_replay == "exact": + last = torch.tensor([s.n for s in segments], dtype=torch.int32, device=self.device).cumsum_(0) - 1 + return DSV41AttnMetadata(last_indices=last, segments=segments, decoder_segments=segments) + win = self.window_size + dec: List[PrefillSegment] = [] + rows: list[torch.Tensor] = [] + last: list[int] = [] + doff = 0 + for r, s in zip(batch.reqs, segments): + floor = max(0, r.prompt_len - win) + lo = max(s.start_pos, floor) + n = max(0, s.end - lo) + if n: + # the decoder's window KV goes to the request's private rings, never to radix-shared pages + dec.append(PrefillSegment(doff, n, s.table_idx, lo, window_floor=floor)) + rows.append(torch.arange(s.offset + lo - s.start_pos, s.offset + s.n)) + doff += n + last.append(doff - 1 if n else -1) + pad = -1 in last + last_indices = torch.tensor([doff if i < 0 else i for i in last], dtype=torch.int32, device=self.device) + decoder_rows = torch.cat(rows).to(self.device, non_blocking=True) if rows else torch.empty(0, dtype=torch.int64, device=self.device) + return DSV41AttnMetadata(last_indices=last_indices, segments=segments, decoder_segments=dec, decoder_rows=decoder_rows, decoder_pad=pad) + + def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None: + assert self.capture is None, "Capture already initialized." + self.max_graph_bs = max(bs_list) + self.capture = DSV41CaptureData.create(self.max_graph_bs, max_seq_len, self.device) + self.capture_bs = sorted(bs_list) + + def prepare_for_capture(self, batch: Batch) -> None: + bs = batch.size + rows = torch.full((bs,), batch.padded_reqs[0].table_idx, dtype=torch.int64, device=self.device) + self._point_to_capture(batch, bs, rows) + + def prepare_for_replay(self, batch: Batch) -> None: + self._point_to_capture(batch, batch.padded_size, self._table_rows(batch)) + + def _table_rows(self, batch: Batch) -> torch.Tensor: + assert batch.active_table_idx is not None, "decode batch is missing its page-table rows" + return batch.active_table_idx.to(torch.int64) + + def _point_to_capture(self, batch: Batch, bs: int, rows_ti: torch.Tensor) -> None: + assert self.capture is not None and bs <= self.max_graph_bs + cap = self.capture + src = self.pool.full_loc_map.index_select(0, rows_ti) + w = min(src.shape[1], cap.full_snap.shape[1]) + cap.full_snap[:bs, :w].copy_(src[:bs, :w]) + cap.table_rows[:bs].copy_(rows_ti[:bs]) + batch.attn_metadata = DSV41AttnMetadata( + last_indices=cap.last_indices[:bs], full_snap=cap.full_snap[:bs], table_rows=cap.table_rows[:bs], window_ar=self._window_ar, + ) + + # ----- metadata / selection access ------------------------------------------------- + @property + def metadata(self) -> DSV41AttnMetadata: + md = get_global_ctx().batch.attn_metadata + assert isinstance(md, DSV41AttnMetadata) + return md + + @property + def selection(self) -> SharedSelection: + return self.metadata.selection + + def snapshot(self) -> torch.Tensor: + return self.metadata.full_snapshot() + + # ----- window tier ------------------------------------------------------------------- + def window_slots_of(self, ti: int, lo: int, hi: int) -> torch.Tensor: + """SHARED window slots of positions ``[lo, hi)`` off the request's LIVE full locs (slid-out -> -1).""" + return self.pool.translate_full_to_window(self.pool.full_loc_map[ti, lo:hi]) + + def ring_slots_of(self, ti: int, lo: int, hi: int) -> torch.Tensor: + """PRIVATE-ring slots of positions ``[lo, hi)`` of page-table row ``ti`` (``hi - lo <= win``).""" + assert hi - lo <= self.window_size, f"a private ring holds one window, not {hi - lo} positions" + return self.pool.ring_slots(torch.tensor(ti, device=self.device), torch.arange(lo, hi, device=self.device)) + + def layer_window_slots_of(self, layer_id: int, ti: int, lo: int, hi: int) -> torch.Tensor: + """Where ``layer_id`` keeps the window KV of positions ``[lo, hi)``: its private ring or the shared pool.""" + if self.pool.is_private_window(layer_id): + return self.ring_slots_of(ti, lo, hi) + return self.window_slots_of(ti, lo, hi) + + def store_window(self, kv: torch.Tensor, layer_id: int, window_slots: torch.Tensor) -> None: + self.pool.store_window(kv, layer_id, window_slots) + + def window_topk_prefill(self, seg: PrefillSegment, layer_id: int | None = None) -> torch.Tensor: + """Per-query window candidates for a prefill segment as GLOBAL window slots ``[1, n, win]`` + (``-1`` where empty) in the tier ``layer_id`` reads (shared pool by default). Query at + absolute ``p`` sees ``[max(floor, first_retained, p - win + 1), p]``.""" + win, device = self.window_size, self.device + lo = max(0, seg.start_pos - win + 1, seg.window_floor) + private = layer_id is not None and self.pool.is_private_window(layer_id) + ws_pool = (self.ring_slots_of if private else self.window_slots_of)(seg.table_idx, lo, seg.end) # [end - lo] + abs_p = seg.start_pos + torch.arange(seg.n, device=device).unsqueeze(1) + cand = (abs_p - win + 1).clamp(min=lo) + torch.arange(win, device=device) + cols = torch.where(cand > abs_p, -1, cand - lo) + g = ws_pool[cols.clamp_min(0)] + return torch.where(cols < 0, -1, g).unsqueeze(0) + + # ----- compressed tiers (per kv source) ---------------------------------------------- + def source_of(self, layer_id: int) -> int: + return self.pool.source_of(layer_id) + + def compressed_rows_of(self, ti: int, group_starts: torch.Tensor, ratio: int) -> torch.Tensor: + """Main / index rows of the compressed groups whose ABSOLUTE first positions are + ``group_starts``, off the request's LIVE full locs (a page is ratio-divisible, so every + position of a group shares one row).""" + return self.pool.cmp_rows(self.pool.full_loc_map[ti, group_starts], ratio) + + def locs_prefill(self, ti: int, end: int) -> torch.Tensor: + """``[1, end]`` int32: the request's live full locs for positions ``[0, end)`` -- what the indexer + derives compressed rows from (``loc // ratio``) for the positions its queries may see.""" + return self.pool.full_loc_map[ti : ti + 1, :end] + + def store_main(self, latent: torch.Tensor, source: int, rows: torch.Tensor) -> None: + self.pool.store_main(latent, source, rows) + + def store_index(self, k: torch.Tensor, source: int, rows: torch.Tensor) -> None: + self.pool.store_index(k, source, rows) + + def decode_store_rows(self, rows: torch.Tensor, pos: torch.Tensor, ratio: int, source: int, completed: torch.Tensor) -> torch.Tensor: + """Per-row decode destination: the completed group's arithmetic row, or the row's OWN scratch + row when this step did not complete a group (graph-safe, collision-free masked store).""" + row_of_group = self.pool.cmp_rows(self.snapshot()[rows, pos], ratio) + scratch = rows + self.pool.scratch_base[source] + return torch.where(completed, row_of_group, scratch) + + # ----- compress-state ring (ratio > 1 sources) --------------------------------------- + def ring_page_base(self, window_slots: torch.Tensor, ring_size: int) -> torch.Tensor: + return torch.div(window_slots, self.window_size, rounding_mode="floor") * ring_size + + def carry_state_loc(self, window_slot: int, ring_size: int) -> torch.Tensor: + base = (window_slot // self.window_size) * ring_size + return torch.arange(base, base + ring_size, device=self.device, dtype=torch.int64) + + def read_carry(self, source: int, window_slot: int) -> torch.Tensor: + """The ``[ring_size, 2 * head_dim]`` carry block at ``window_slot``'s page (kv | score).""" + ring = self.pool.state_ring[source] + return ring.get(self.carry_state_loc(window_slot, ring.ring_size)) + + def write_carry(self, source: int, window_slot: int, kv_score: torch.Tensor) -> None: + ring = self.pool.state_ring[source] + ring.set(self.carry_state_loc(window_slot, ring.ring_size), kv_score) + + def read_carry_blocks(self, source: int, window_slots: torch.Tensor) -> torch.Tensor: + ring = self.pool.state_ring[source] + return ring.get_blocks(self.ring_page_base(window_slots, ring.ring_size)) + + def write_carry_blocks(self, source: int, window_slots: torch.Tensor, blocks: torch.Tensor) -> None: + ring = self.pool.state_ring[source] + ring.set_blocks(self.ring_page_base(window_slots, ring.ring_size), blocks) + + def write_boundary_carries(self, source: int, *, lo: int, hi: int, window_slots: torch.Tensor) -> None: + """Persist the compressor carry at every window-page boundary ``B`` in ``(lo, hi]`` so a + page-aligned radix match can resume by value. A page holds whole groups (``P % ratio == 0``), + so the carry at a boundary is the empty group -- the reset block is written so a resume never + reads a stale one. ``window_slots`` is indexed by ``pos - lo``. One batched ring write, no host syncs.""" + ring = self.pool.state_ring.get(source) + if ring is None: + return + P = self.window_size + first = (lo // P + 1) * P + if first > hi: + return + bounds = torch.arange(first, hi + 1, P, device=self.device) + page_slots = window_slots[bounds - 1 - lo].to(torch.int64) # the last slot of each page + locs = (torch.div(page_slots, P, rounding_mode="floor") * ring.ring_size)[:, None] + torch.arange(ring.ring_size, device=self.device) + empty = torch.cat([ + torch.zeros(ring.ring_size, ring.item_size, dtype=torch.float32, device=self.device), + torch.full((ring.ring_size, ring.item_size), float("-inf"), dtype=torch.float32, device=self.device), + ], dim=-1) + ring.set(locs.flatten(), empty.repeat(bounds.numel(), 1)) + + # ----- indexer ---------------------------------------------------------------------- + def indexer_logits( + self, q: torch.Tensor, weights: torch.Tensor, source: int, locs: torch.Tensor, ratio: int, live: torch.Tensor, + T: int | None = None, candidates: torch.Tensor | None = None, + ) -> torch.Tensor: + """``[B, S, T]`` full-range scores (live columns written) or ``[B, S, NC]`` candidate scores; the + row of compressed position ``t`` is ``locs[b, t * ratio] // ratio``.""" + from freetoken.kernel.triton.dsv41.indexer import indexer_logits_packed + + return indexer_logits_packed(q, weights, self.pool.idx_pool[source], self.pool.idx_fmt, locs, ratio, live, T=T, candidates=candidates) + + @staticmethod + def select_topk(scores: torch.Tensor, live: torch.Tensor, topk: int) -> torch.Tensor: + """Exact causal top-k over ``[B, S, T]`` full-range scores (columns at or past ``live`` unread): + compressed positions ``[B, S, topk]`` int32 in ascending order, ``-1`` in the tail for picks + that do not exist (fewer live or finite positions than ``topk``). Ties: lowest position wins.""" + from freetoken.kernel.triton.dsv41.topk import dsv41_topk + + b, s, t = scores.shape + return dsv41_topk(scores.reshape(b * s, t), live.reshape(b * s), topk).view(b, s, topk) + + @staticmethod + def select_topk_in_candidates(scores: torch.Tensor, candidates: torch.Tensor, topk: int) -> torch.Tensor: + """Top-k over ``[B, S, NC]`` candidate-aligned scores -> the picked compressed positions + ``[B, S, topk]`` int32. The candidate list is a sorted valid prefix (``-1`` tail) and empty / + unreachable slots score ``-inf``, so the ascending winning slots map to ascending positions.""" + from freetoken.kernel.triton.dsv41.topk import dsv41_topk + + b, s, nc = scores.shape + live = torch.full((b * s,), nc, dtype=torch.int32, device=scores.device) + slots = dsv41_topk(scores.reshape(b * s, nc), live, topk).view(b, s, topk) + pos = candidates.gather(-1, slots.clamp_min(0).to(torch.int64)).to(torch.int32) + return torch.where(slots < 0, -1, pos) + + @staticmethod + def select_candidate_blocks(scores: torch.Tensor, live: torch.Tensor, topk_blocks: int, block_size: int) -> torch.Tensor: + """Level one of the Hierarchical Sparse Indexer: the ``topk_blocks`` best live ``block_size`` + blocks by their best position, the block holding the newest live position always kept, + expanded to the compact candidate list ``[B, S, topk_blocks * block_size]`` int32 of + compressed positions (ascending, ``-1`` tail). Mirrors the reference ``select_candidate_blocks`` + (which returns the equivalent boolean mask).""" + from freetoken.kernel.triton.dsv41.topk import dsv41_candidate_blocks + + b, s, t = scores.shape + return dsv41_candidate_blocks(scores.reshape(b * s, t), live.reshape(b * s), topk_blocks, block_size).view(b, s, -1) + + @staticmethod + def positions_to_rows(positions: torch.Tensor, locs: torch.Tensor, ratio: int) -> torch.Tensor: + """Compressed positions ``[B, S, K]`` (``-1`` masked) -> GLOBAL main rows ``locs[b, p * ratio] // ratio``.""" + b = positions.shape[0] + idx = (positions.to(torch.int64) * ratio).clamp_min(0).flatten(1) + full = locs.to(torch.int64).gather(1, idx.clamp(max=locs.shape[1] - 1)).view_as(positions) + rows = torch.div(full, ratio, rounding_mode="floor") + return torch.where((positions < 0) | (full < 0), -1, rows) + + # ----- attention ----------------------------------------------------------------------- + def attend( + self, q: torch.Tensor, layer_id: int, topk_idxs: torch.Tensor, n_window: int, attn_sink: torch.Tensor, + softmax_scale: float, cmp_counts: torch.Tensor | None = None, + ) -> torch.Tensor: + """Paged sparse attention over ``[window | compressed]`` global rows; the compressed half reads + the layer's kv source (window-only layers pass ``n_window == topk``).""" + from freetoken.kernel.triton.dsv41.sparse_attn import sparse_attn_packed + + pool, args = self.pool, self.config.dsv41_args + src = args.roles[layer_id].kv_source + cmp = pool.main_pool[src] if src is not None else pool.main_pool[args.backbone_kv_sources[0]] + return sparse_attn_packed( + q, pool.window_pool[layer_id], pool.win_fmt, cmp, pool.main_fmt, attn_sink, + topk_idxs, n_window, softmax_scale, cmp_counts=cmp_counts, + ) + + +__all__ = ["DSV41SparseAttnBackend", "DSV41AttnMetadata", "PrefillSegment", "SharedSelection"] diff --git a/python/freetoken/core.py b/python/freetoken/core.py index fd691d180c..45a763c0fe 100644 --- a/python/freetoken/core.py +++ b/python/freetoken/core.py @@ -43,6 +43,8 @@ class Req: mm_items: list | None = None mrope_positions_full: torch.Tensor | None = None # [3, prompt_len] int32, CPU mrope_delta: int = 0 + # The whole prompt's length; 0 = this request's input_ids (a ChunkedReq's end at its chunk). + prompt_len: int = 0 # --- hybrid-radix (GDN linear-state) per-request slots; None for non-hybrid models or # until allocated from LinearStatePool. Set by the scheduler (P2). --- @@ -69,6 +71,7 @@ def __post_init__(self) -> None: assert self.input_ids.is_cpu self.device_len = len(self.input_ids) self.max_device_len = len(self.input_ids) + self.output_len + self.prompt_len = self.prompt_len or self.device_len assert 0 <= self.cached_len < self.device_len <= self.max_device_len self._alloc_ids_buf() diff --git a/python/freetoken/engine/cache_budget.py b/python/freetoken/engine/cache_budget.py index d522be4991..b5f07ee8ab 100644 --- a/python/freetoken/engine/cache_budget.py +++ b/python/freetoken/engine/cache_budget.py @@ -24,7 +24,7 @@ def expert_bytes_per_slot(sources: dict[str, "list[torch.Tensor]"]) -> int: # marlin/b12x gate_up/down alpha scales are fixed [L*E] residency (do not scale # 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. + # slot-byte idiom in kvcache/linear_state_pool.py and kvcache/dsv4/v4_pool.py. return sum(t[0][0].numel() * t[0].element_size() for t in sources.values()) diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 3e38cb2165..b6fa8b8949 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -34,6 +34,9 @@ class EngineConfig: quant_backend: str | None = None # PLE table backend: "disk" (default) reads rows from the checkpoint files per fill, "pinned" preloads the table into page-locked host RAM. ple_backend: str = "disk" + # DeepSeek-V4.1 Decoder SWA Bounded Replay: "bounded" (default) runs the decoder layers on each + # prompt's last window; "exact" runs them on every token. + swa_decoder_replay: str = "bounded" # Expert-bank host load (--expert-load): auto|serial|parallel. "auto" reads scattered # experts in parallel but falls back to serial when free RAM can't cover the banks + the # parallel reader's extra (non-reclaimable) whole-shard buffer; "serial" forces the diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index e5ccf71050..bf3960ab8d 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -129,6 +129,8 @@ def _resolve_auto_attention_backend(required: frozenset[AttnType]) -> str: candidates: list[tuple[str, bool]] = [] if AttnType.DSV4 in required: candidates.append(("dsv4_sparse", True)) + if AttnType.DSV41 in required: + candidates.append(("dsv41_sparse", True)) if required & {AttnType.MLA, AttnType.DSA}: candidates.append(("dsa", True)) if AttnType.BSA in required: @@ -812,6 +814,7 @@ def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None: # Decode batches never exceed max_running_req, but CUDA-graph padding can # round a batch up to the largest captured size; cover both. max_tokens = max(config.max_running_req, config.cuda_graph_max_bs or 0, 1) + method = sample.quant_method executor = CpuMoeExecutor( cache, top_k=sample.top_k, @@ -823,7 +826,8 @@ def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None: swiglu_alpha=float(sample.alpha), swiglu_limit=sample.limit, # FIXME: the None branch serves GGUF q4_0 banks, which have no quant method yet; drop it once GGUF joins the quant path - fmt=sample.quant_method.cpu_format if sample.quant_method is not None else None, + fmt=method.cpu_format if method is not None else None, + act_block=method.cfg.act_block if method is not None else None, ) cache.set_cpu_executor(executor) self.cpu_moe_executor = executor @@ -1200,6 +1204,39 @@ def _resolve_cache_type(has_linear_attention: bool, requested: str) -> str: return requested +def _adjust_dsv41_config(config: EngineConfig, override) -> None: + """DeepSeek-V4.1 config reconciliation, the DSV4 policy with the replay knob: sync the + runtime into ``dsv41_args``, page_size = the window page P, radix -> swa_radix, decode graph + batches <= max_running_req. Unlike DSV4 the prefill chunk keeps ``max_extend_tokens`` (whole + window pages): at a 1M ceiling the window pool alone would admit ~40K-token chunks whose + activations (64 x 512 latent queries per token) do not fit next to the expert cache.""" + model_config = config.model_config + args = model_config.dsv41_args + args.max_seq_len = config.max_seq_len + args.max_batch_size = config.max_running_req + 1 # +1 dummy + args.swa_decoder_replay = config.swa_decoder_replay + P = args.window_size + override("page_size", P) + logger.info_rank0( + f"DSV41 KV pages are {P}-token window pages; page_size set to {P}; SWA decoder replay: {args.swa_decoder_replay}" + + ("" if args.swa_decoder_replay == "exact" else f" (prefix hits stop {P} tokens before the prompt end)") + ) + if getattr(config, "cache_type", "radix") != "naive": + override("cache_type", "swa_radix") + current = getattr(config, "max_extend_tokens", None) # a SchedulerConfig field + if current is not None and (extend := max(P, current // P * P)) != current: + override("max_extend_tokens", extend) + mr = config.max_running_req + if config.cuda_graph_max_bs is not None and config.cuda_graph_max_bs > mr: + logger.warning_rank0(f"cuda_graph_max_bs {config.cuda_graph_max_bs} exceeds max_running_req {mr}; clamping (larger decode batches never occur).") + override("cuda_graph_max_bs", mr) + if config.cuda_graph_bs is not None: + kept = [bs for bs in config.cuda_graph_bs if bs <= mr] + if kept != list(config.cuda_graph_bs): + logger.warning_rank0(f"dropping cuda_graph_bs entries above max_running_req {mr}: {[bs for bs in config.cuda_graph_bs if bs > mr]}") + override("cuda_graph_bs", kept) + + def _adjust_dsv4_config(config: EngineConfig, override) -> None: """DSV4 engine-config reconciliation at config-resolution time (before the pool exists). Syncs the resolved runtime config into the opaque ``dsv4_args`` payload, sets @@ -1477,6 +1514,7 @@ def override(attr: str, value: Any): # this is dangerous, use with caution model_config = config.model_config single_stream_only = getattr(model_config, "single_stream_only", False) is_dsv4 = getattr(model_config, "dsv4_args", None) is not None + is_dsv41 = getattr(model_config, "dsv41_args", None) is not None has_swa_attention = getattr(model_config, "has_swa_attention", False) has_linear_attention = getattr(model_config, "has_linear_attention", False) is_moe = getattr(model_config, "is_moe", False) @@ -1520,6 +1558,8 @@ def override(attr: str, value: Any): # this is dangerous, use with caution if is_dsv4: _adjust_dsv4_config(config, override) + if is_dsv41: + _adjust_dsv41_config(config, override) if has_swa_attention: # Both SWA cache paths use the global-paged swa pool (page_size==1 only for now). @@ -1806,7 +1846,7 @@ def override(attr: str, value: Any): # this is dangerous, use with caution # DSV4 is exempt: it sizes its own table from the resolved max_seq_len (_adjust_dsv4_config). rotary = getattr(model_config, "rotary_config", None) seq_override = getattr(config, "max_seq_len_override", None) - if seq_override is not None and rotary is not None and not is_dsv4: + if seq_override is not None and rotary is not None and not (is_dsv4 or is_dsv41): if seq_override > rotary.table_positions: raise ValueError( f"--max-seq-len-override {seq_override} exceeds the model's " diff --git a/python/freetoken/kernel/aot_models.py b/python/freetoken/kernel/aot_models.py index a8ef85d432..ecd1f8502b 100644 --- a/python/freetoken/kernel/aot_models.py +++ b/python/freetoken/kernel/aot_models.py @@ -332,6 +332,17 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int expert_formats=("ds_fp4",), embed_indexing=False, # plain nn.Embedding ), + AotModel( + # DSV41 pools are packed byte rows written by the dsv41 pack kernel, not store_cache; + # the routed experts share DSV4's ds_fp4 bank layout (hidden 5120, intermediate 2304). + name="deepseek-ai/DeepSeek-V4.1-Flash", + architecture="DeepseekV41ForCausalLM", + hidden_size=5120, + kv_groups=(), + top_k=6, + moe_intermediate_size=2304, + expert_formats=("ds_fp4",), + ), # ---- dense checkpoints (store/index only, no expert banks) ---- AotModel( name="Qwen/Qwen3.6-27B", diff --git a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp index fb6a33f497..86765013e3 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp @@ -1063,11 +1063,12 @@ inline void deinterleave_bf16_f32(const bf16_t* src, float* even, float* odd, in } } -// DeepSeek-V4 activation FP8 round-trip (bf16 in/out): per 128-block, -// s = 2^ceil(log2(max(|x|,1e-4)/448)); y = round_e4m3(clamp(x/s,+-448)) * s. -void fp8_roundtrip_bf16(const bf16_t* src, bf16_t* dst, int K) { - for (int b0 = 0; b0 < K; b0 += 128) { - const int b1 = std::min(K, b0 + 128); +// DeepSeek-V4 / V4.1 activation FP8 round-trip (bf16 in/out): per `block` (the checkpoint's +// fp8 block: 128 on V4, 32 on V4.1), s = 2^ceil(log2(max(|x|,1e-4)/448)); +// y = round_e4m3(clamp(x/s,+-448)) * s. +void fp8_roundtrip_bf16(const bf16_t* src, bf16_t* dst, int K, int block) { + for (int b0 = 0; b0 < K; b0 += block) { + const int b1 = std::min(K, b0 + block); float amax = 1e-4f; for (int i = b0; i < b1; ++i) amax = std::max(amax, std::fabs(bf16_to_f32(src[i]))); const float s = std::ldexp(1.0f, ceil_log2_pos(amax * (1.0f / 448.0f))); @@ -1275,6 +1276,8 @@ struct CpuMoeExecutor { // it to a captured GPU elementwise kernel removes it while keeping the official // W4A8 numerics bit-exact. Set via set_input_prequant (see cpu_executor.py). bool input_prequant = false; + // ds_fp4 activation round-trip block (the checkpoint's fp8 block); set via set_act_block. + int act_block = 128; // Q4_0 packed-row byte strides (H/32*18 for gate_up over K=H, I/32*18 for down over K=I). int q4_gu_row_bytes = 0, q4_dn_row_bytes = 0; float e2m1_lut[16]; @@ -1817,7 +1820,7 @@ struct CpuMoeExecutor { gas_scratch.data() + (size_t)r * (I / 32)); return; } - if (fmt == WF_DSFP4) fp8_roundtrip_bf16(g, g, I); + if (fmt == WF_DSFP4) fp8_roundtrip_bf16(g, g, I, act_block); float* ge = ge_scratch.data() + (size_t)r * (I / 2); float* go = go_scratch.data() + (size_t)r * (I / 2); deinterleave_bf16_f32(g, ge, go, I); @@ -1957,7 +1960,7 @@ struct CpuMoeExecutor { const bf16_t* src = t->x + (size_t)tok * H; if (ds) { // DSV4 FP8-round-trips the input before the gate_up GEMV bf16_t* xq = xq_scratch.data() + (size_t)tok * H; - fp8_roundtrip_bf16(src, xq, H); + fp8_roundtrip_bf16(src, xq, H, act_block); src = xq; } float* xe = xe_scratch.data() + (size_t)tok * (H / 2); @@ -2185,6 +2188,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { .def("set_input_prequant", [](CpuMoeExecutor& e, bool v) { e.input_prequant = v; }, py::arg("value")) + .def("set_act_block", + [](CpuMoeExecutor& e, int v) { e.act_block = v; }, + py::arg("value")) .def("isa_name", &CpuMoeExecutor::isa_name) .def("set_spin_wait", [](CpuMoeExecutor& e, bool v) { e.spin_wait.store(v, std::memory_order_relaxed); }, diff --git a/python/freetoken/kernel/triton/dsv4/fp8_linear.py b/python/freetoken/kernel/triton/dsv4/fp8_linear.py index 47e3c9f58a..fb17a35175 100644 --- a/python/freetoken/kernel/triton/dsv4/fp8_linear.py +++ b/python/freetoken/kernel/triton/dsv4/fp8_linear.py @@ -1,23 +1,25 @@ -"""Block-scaled FP8 (e4m3) linear for DeepSeek-V4, matching the reference numerics. +"""Block-scaled FP8 (e4m3) linear for DeepSeek-V4 / V4.1, matching the reference numerics. The reference (``inference/model.py`` ``linear`` + ``inference/kernel.py`` -``act_quant``/``fp8_gemm``) quantizes the *activation* to FP8 with a per-128 block +``act_quant``/``fp8_gemm``) quantizes the *activation* to FP8 with a per-``block`` power-of-two (ue8m0) scale, then runs an FP8xFP8 block-scaled GEMM against the FP8 -weight (which carries its own 128x128 ue8m0 block scale). Both operands' scales are -applied per 128-K block to a separate FP32 accumulator. This module reproduces that: +weight (which carries its own ``block x block`` ue8m0 scale). Both operands' scales are +applied per ``block``-K slab to a separate FP32 accumulator. This module reproduces that: - ``y = fp8_gemm(act_quant(x, 128, ue8m0), weight_fp8, weight_scale_e8m0)`` + ``y = fp8_gemm(act_quant(x, block, ue8m0), weight_fp8, weight_scale_e8m0)`` + +``block`` is the checkpoint's ``weight_block_size``: 128 for DeepSeek-V4, 32 for V4.1. ``act_quant`` (reference): per block ``s = 2**ceil(log2(max(|x|,1e-4)/448))`` (exact via IEEE bit ops -> matches ``fast_round_scale``), ``x_fp8 = round_e4m3(clamp(x/s, +-448))``, scale stored e8m0. The GEMM accumulates ``sum_k (A_fp8 @ B_fp8) * s_a * s_b`` -per 128-K block in FP32. +per ``block``-K slab in FP32. Also provides ``act_quant_fp8_inplace`` -- the fused FP8 quant+dequant round-trip the reference applies in-place to the window / compressor KV (``act_quant(..., 64, ..., inplace=True)``), returning BF16. -Assumes ``K % 128 == 0`` and ``N % 128 == 0`` (true for every DeepSeek-V4 projection). +Assumes ``K % block == 0`` and ``N % block == 0`` (true for every DeepSeek-V4 / V4.1 projection). """ from __future__ import annotations @@ -218,8 +220,8 @@ def fp4_act_quant_inplace(x: torch.Tensor, block: int = 32) -> torch.Tensor: def _fp8_act_gemm_kernel( a_ptr, # [M, K] float8_e4m3fn (quantized activation) w_ptr, # [N, K] float8_e4m3fn - sa_ptr, # [M, K//128] uint8 (e8m0 act codes) - sb_ptr, # [N//128, K//128] uint8 (e8m0 weight codes) + sa_ptr, # [M, K//BLOCK_K] uint8 (e8m0 act codes) + sb_ptr, # [N//BLOCK_K, K//BLOCK_K] uint8 (e8m0 weight codes; square BLOCK_K blocks) c_ptr, # [M, N] compute dtype M, N, K, stride_am, stride_ak, stride_wn, stride_wk, @@ -228,6 +230,11 @@ def _fp8_act_gemm_kernel( BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, compute_type: tl.constexpr, ): + """``BLOCK_K`` is the weight's square block edge, so each K step maps to exactly one scale + column. When the N tile IS one scale row (``BLOCK_N == BLOCK_K``, the DeepSeek-V4 block-128 + case) the weight scale is a scalar per step and the N masks vanish -- the original kernel; + otherwise (block 32) the tile spans ``BLOCK_N // BLOCK_K`` scale rows and N is masked.""" + SCALAR_SCALE: tl.constexpr = BLOCK_N == BLOCK_K pid_m = tl.program_id(0) pid_n = tl.program_id(1) offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) @@ -238,112 +245,217 @@ def _fp8_act_gemm_kernel( w_ptrs = w_ptr + offs_n[:, None] * stride_wn + offs_k[None, :] * stride_wk acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) num_k = tl.cdiv(K, BLOCK_K) - for k in range(num_k): - a = tl.load(a_ptrs, mask=m_mask[:, None], other=0.0) - w = tl.load(w_ptrs) - if e4m3_native_cx(): - p = tl.dot(a, tl.trans(w), out_dtype=tl.float32) - else: - # bf16 dot on the same e4m3 grid: operands exact in bf16, fp32 acc - p = tl.dot(a, tl.trans(e4m3_u8_to_f32(w).to(tl.bfloat16)), out_dtype=tl.float32) - sa_code = tl.load(sa_ptr + offs_m * stride_sam + k * stride_sak, mask=m_mask, other=0) - sca = tl.exp2(sa_code.to(tl.float32) - 127.0) # [BLOCK_M] - sb_code = tl.load(sb_ptr + pid_n * stride_sbn + k * stride_sbk) - scb = tl.exp2(sb_code.to(tl.float32) - 127.0) # scalar (one 128-N block) - acc += p * sca[:, None] * scb - a_ptrs += BLOCK_K * stride_ak - w_ptrs += BLOCK_K * stride_wk - c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn - tl.store(c_ptrs, acc.to(compute_type), mask=m_mask[:, None]) + if SCALAR_SCALE: + for k in range(num_k): + a = tl.load(a_ptrs, mask=m_mask[:, None], other=0.0) + w = tl.load(w_ptrs) + if e4m3_native_cx(): + p = tl.dot(a, tl.trans(w), out_dtype=tl.float32) + else: + # bf16 dot on the same e4m3 grid: operands exact in bf16, fp32 acc + p = tl.dot(a, tl.trans(e4m3_u8_to_f32(w).to(tl.bfloat16)), out_dtype=tl.float32) + sa_code = tl.load(sa_ptr + offs_m * stride_sam + k * stride_sak, mask=m_mask, other=0) + sca = tl.exp2(sa_code.to(tl.float32) - 127.0) # [BLOCK_M] + sb_code = tl.load(sb_ptr + pid_n * stride_sbn + k * stride_sbk) + scb = tl.exp2(sb_code.to(tl.float32) - 127.0) # scalar (one N block) + acc += p * sca[:, None] * scb + a_ptrs += BLOCK_K * stride_ak + w_ptrs += BLOCK_K * stride_wk + c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn + tl.store(c_ptrs, acc.to(compute_type), mask=m_mask[:, None]) + else: + n_mask = offs_n < N # N is a multiple of BLOCK_K, not necessarily of BLOCK_N + sn = offs_n // BLOCK_K + for k in range(num_k): + a = tl.load(a_ptrs, mask=m_mask[:, None], other=0.0) + w = tl.load(w_ptrs, mask=n_mask[:, None], other=0.0) + if e4m3_native_cx(): + p = tl.dot(a, tl.trans(w), out_dtype=tl.float32) + else: + p = tl.dot(a, tl.trans(e4m3_u8_to_f32(w).to(tl.bfloat16)), out_dtype=tl.float32) + sa_code = tl.load(sa_ptr + offs_m * stride_sam + k * stride_sak, mask=m_mask, other=0) + sca = tl.exp2(sa_code.to(tl.float32) - 127.0) # [BLOCK_M] + sb_code = tl.load(sb_ptr + sn * stride_sbn + k * stride_sbk, mask=n_mask, other=0) + scb = tl.exp2(sb_code.to(tl.float32) - 127.0) # [BLOCK_N] + acc += p * sca[:, None] * scb[None, :] + a_ptrs += BLOCK_K * stride_ak + w_ptrs += BLOCK_K * stride_wk + c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn + tl.store(c_ptrs, acc.to(compute_type), mask=m_mask[:, None] & n_mask[None, :]) @triton.jit -def _fp8_act_gemv_splitk_kernel( - a_ptr, # [K] float8_e4m3fn - sa_ptr, # [K//128] uint8 (e8m0 act codes) - w_ptr, # [N, K] float8_e4m3fn - sb_ptr, # [N//128, K//128] uint8 (e8m0 weight codes) - part_ptr, # [SPLIT_K, N] fp32 - N, K, - stride_ak, stride_wn, stride_wk, stride_sbn, stride_sbk, stride_pk, stride_pn, - BLOCK_N: tl.constexpr, SPLIT_K: tl.constexpr, +def _fp8_gemv_kernel( + x_ptr, # [M, G, K] activation in the compute dtype (bf16); QUANT_ACT: quantized in-kernel per SLAB + w_ptr, # [G * N, K] float8_e4m3fn (or its uint8 view on pre-Ada parts) + sb_ptr, # [G * N // SLAB, K // SLAB] uint8 (e8m0 weight codes) + out_ptr, # [M, G * N] OUT dtype (SPLIT_K == 1) + part_ptr, # [SPLIT_K, M, G * N] fp32 partials (SPLIT_K > 1) + M, N, K, + stride_xm, stride_xg, stride_wn, stride_wk, stride_sbn, stride_sbk, stride_om, stride_pk, stride_pm, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, SLAB: tl.constexpr, SPLIT_K: tl.constexpr, + QUANT_ACT: tl.constexpr, OUT: tl.constexpr, ): + """Small-M (decode) block-scaled fp8 GEMV/GEMM: one pass over the weight at near HBM bandwidth. + + Each program owns ``BLOCK_N`` rows and ``K // SPLIT_K`` contiguous K, walked ``BLOCK_K`` (several + ``SLAB``-wide scale slabs) at a time: a wide tile keeps enough bytes in flight per SM to saturate + the memory system, which a slab-at-a-time walk (512 B per load at the V4.1 slab of 32) cannot. + Every weight tile is applied to all ``M`` (<= ``BLOCK_M``) activation rows, so a decode batch + costs one weight read like a single token does (the tensor-core GEMM path is 3-6x slower at + M = 2..8: too few tiles, tiny K steps). ``QUANT_ACT`` quantizes the activation here rather than + by a separate ``act_quant`` launch: every program recomputes the ue8m0 scale and the e4m3 rounding + of the slabs it touches (the reference formula, bit-identical to ``act_quant_fp8``) from the + L2-resident bf16 rows, so the path is one kernel (plus a small reduce when ``SPLIT_K > 1``). + Without it the activation stays bf16 (W8A16: the reference's bf16 einsum over a dequantized + weight, e.g. ``wo_a``). The third grid axis is an independent group (block-diagonal weights: + group ``g`` reads ``x[:, g]`` and rows ``[g * N, (g + 1) * N)``).""" + tl.static_assert(BLOCK_K % SLAB == 0, "the K tile holds whole scale slabs") + NSLAB: tl.constexpr = BLOCK_K // SLAB pid_n = tl.program_id(0) pid_k = tl.program_id(1) + g = tl.program_id(2).to(tl.int64) + x_ptr += g * stride_xg offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) n_mask = offs_n < N - sn = offs_n // 128 + sn = (g * N + offs_n) // SLAB + offs_n = g * N + offs_n + rows = tl.arange(0, BLOCK_M) k_per = K // SPLIT_K k_start = pid_k * k_per - acc = tl.zeros((BLOCK_N,), dtype=tl.float32) - for k0 in range(0, k_per, 128): - offs_k = k_start + k0 + tl.arange(0, 128) - a = tl.load(a_ptr + offs_k * stride_ak).to(tl.float32) + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + for k0 in range(0, k_per, BLOCK_K): + offs_k = k_start + k0 + tl.arange(0, BLOCK_K) + # --- the weight tile, dequantized per slab (read once for every activation row) --- w_ptrs = w_ptr + offs_n[:, None] * stride_wn + offs_k[None, :] * stride_wk if e4m3_native_cx(): w = tl.load(w_ptrs, mask=n_mask[:, None], other=0.0).to(tl.float32) else: w = e4m3_u8_to_f32(tl.load(w_ptrs, mask=n_mask[:, None], other=0)) - kb = (k_start + k0) // 128 - sb_code = tl.load(sb_ptr + sn * stride_sbn + kb * stride_sbk, mask=n_mask, other=0) - scb = tl.exp2(sb_code.to(tl.float32) - 127.0) - sa_code = tl.load(sa_ptr + kb) - sca = tl.exp2(sa_code.to(tl.float32) - 127.0) - acc += tl.sum(w * a[None, :], axis=1) * scb * sca - tl.store(part_ptr + pid_k * stride_pk + offs_n * stride_pn, acc, mask=n_mask) + kb = (k_start + k0) // SLAB + tl.arange(0, NSLAB) + sb_code = tl.load(sb_ptr + sn[:, None] * stride_sbn + kb[None, :] * stride_sbk, mask=n_mask[:, None], other=0) + scb = tl.exp2(sb_code.to(tl.float32) - 127.0) # [BLOCK_N, NSLAB] + for m in tl.static_range(BLOCK_M): + if m < M: + xs = tl.reshape(tl.load(x_ptr + m * stride_xm + offs_k).to(tl.float32), (NSLAB, SLAB)) + if QUANT_ACT: + # --- reference act_quant on this row's slabs: s = 2**ceil(log2(max(|x|, 1e-4) / 448)) --- + amax = tl.maximum(tl.max(tl.abs(xs), axis=1), 1e-4) + sa = tl.exp2(_log2_ceil(amax * (1.0 / 448.0)).to(tl.float32)) # [NSLAB] + q = tl.clamp(xs / sa[:, None], -448.0, 448.0) + if e4m3_native_cx(): + q = q.to(tl.float8e4nv).to(tl.float32) + else: + q = round_e4m3(q).to(tl.float32) + else: + sa = tl.full((NSLAB,), 1.0, tl.float32) + q = xs + a = tl.reshape(q, (BLOCK_K,)) + slab = tl.sum(tl.reshape(w * a[None, :], (BLOCK_N, NSLAB, SLAB)), axis=2) # [BLOCK_N, NSLAB] + contrib = tl.sum(slab * scb * sa[None, :], axis=1) # [BLOCK_N] + acc = tl.where(rows[:, None] == m, acc + contrib[None, :], acc) + row_mask = rows < M + if SPLIT_K == 1: + tl.store(out_ptr + rows[:, None] * stride_om + offs_n[None, :], acc.to(OUT), mask=row_mask[:, None] & n_mask[None, :]) + else: + tl.store(part_ptr + pid_k * stride_pk + rows[:, None] * stride_pm + offs_n[None, :], acc, mask=row_mask[:, None] & n_mask[None, :]) @triton.jit -def _splitk_reduce_kernel(part_ptr, out_ptr, N, SPLIT_K: tl.constexpr, - stride_pk, stride_pn, BLOCK: tl.constexpr, OUT: tl.constexpr): +def _splitk_reduce_kernel(part_ptr, out_ptr, M, N, SPLIT_K: tl.constexpr, + stride_pk, stride_pm, stride_om, BLOCK: tl.constexpr, OUT: tl.constexpr): + m = tl.program_id(1) offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) mask = offs < N acc = tl.zeros((BLOCK,), dtype=tl.float32) for k in tl.static_range(SPLIT_K): - acc += tl.load(part_ptr + k * stride_pk + offs * stride_pn, mask=mask, other=0.0) - tl.store(out_ptr + offs, acc.to(OUT), mask=mask) - - -# Swept-best decode GEMV config per (N, K) -> (BLOCK_N, SPLIT_K, num_warps). -_DECODE_FP8_CFG = { - (1024, 4096): (16, 16, 1), # wq_a - (32768, 1024): (16, 1, 2), # wq_b - (512, 4096): (16, 16, 1), # attn wkv - (4096, 8192): (16, 8, 1), # wo_b - (2048, 4096): (16, 8, 1), # shared w1 / w3 - (4096, 2048): (16, 8, 1), # shared w2 - (8192, 1024): (16, 2, 4), # indexer wq_b -} - - -def _decode_cfg(N: int, K: int) -> tuple[int, int, int]: - cfg = _DECODE_FP8_CFG.get((N, K)) - if cfg is not None: - bn, sk, nw = cfg - return bn, max(1, min(sk, K // 128)), nw - bn = 16 - n_tiles = triton.cdiv(N, bn) - sk = max(1, 1536 // n_tiles) - sk = 1 << (sk.bit_length() - 1) - return bn, max(1, min(sk, K // 128)), 1 - - -def _fp8_act_gemv(a_fp8: torch.Tensor, sa: torch.Tensor, weight: torch.Tensor, - sb: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: - N, K = weight.shape - BLOCK_N, split_k, num_warps = _decode_cfg(N, K) - n_tiles = triton.cdiv(N, BLOCK_N) - part = torch.empty((split_k, N), dtype=torch.float32, device=a_fp8.device) - _fp8_act_gemv_splitk_kernel[(n_tiles, split_k)]( - a_fp8, sa, weight, sb, part, N, K, - a_fp8.stride(0), weight.stride(0), weight.stride(1), - sb.stride(0), sb.stride(1), part.stride(0), part.stride(1), - BLOCK_N=BLOCK_N, SPLIT_K=split_k, num_warps=num_warps, + acc += tl.load(part_ptr + k * stride_pk + m * stride_pm + offs, mask=mask, other=0.0) + tl.store(out_ptr + m * stride_om + offs, acc.to(OUT), mask=mask) + + +_GEMV_BLOCK_N = 32 +_GEMV_NARROW_BLOCK_N = 8 +def _gemv_cfg(N: int, K: int, slab: int, sm_count: int) -> tuple[int, int, int, int]: + """Choose whole scale tiles and enough independent programs to cover the device.""" + block_k = next((bk for bk in (512, 256, 128) if bk >= slab and K % bk == 0), slab) + if triton.cdiv(N, _GEMV_BLOCK_N) >= sm_count: + return _GEMV_BLOCK_N, block_k, 1, 4 + n_tiles = triton.cdiv(N, _GEMV_NARROW_BLOCK_N) + split = 1 + while n_tiles * split < sm_count and (K // block_k) % (split * 2) == 0: + split *= 2 + return _GEMV_NARROW_BLOCK_N, block_k, split, 4 + + +GEMV_MAX_M = 4 # bound the unrolled row loop's register footprint + + +def _fp8_gemv( + x: torch.Tensor, weight: torch.Tensor, sb: torch.Tensor, out_dtype: torch.dtype, block: int, *, groups: int = 1, quant_act: bool = True, +) -> torch.Tensor: + """``x [M, groups, K]`` bf16 x block-diagonal ``weight [groups * N, K]`` fp8 -> ``[M, groups * N]`` + in ``out_dtype`` (``M <= GEMV_MAX_M``). ``quant_act`` applies the reference fp8 activation quant + (W8A8); off, the activation stays bf16 (W8A16). ``groups == 1`` is the plain GEMV.""" + GN, K = weight.shape + N = GN // groups + x = x.reshape(-1, groups, K) + M = x.shape[0] + assert 1 <= M <= GEMV_MAX_M, M + block_m = triton.next_power_of_2(M) + block_n, block_k, split_k, num_warps = _gemv_cfg(GN, K, block, torch.cuda.get_device_properties(x.device).multi_processor_count) + out = torch.empty((M, GN), dtype=out_dtype, device=x.device) + part = out if split_k == 1 else torch.empty((split_k, M, GN), dtype=torch.float32, device=x.device) + _fp8_gemv_kernel[(triton.cdiv(N, block_n), split_k, groups)]( + x, weight, sb, out, part, M, N, K, + x.stride(0), x.stride(1), weight.stride(0), weight.stride(1), sb.stride(0), sb.stride(1), out.stride(0), + part.stride(0) if split_k > 1 else 0, part.stride(1) if split_k > 1 else 0, + BLOCK_M=block_m, BLOCK_N=block_n, BLOCK_K=block_k, SLAB=block, SPLIT_K=split_k, QUANT_ACT=quant_act, OUT=_TL_DTYPE[out_dtype], + num_warps=num_warps, ) - out = torch.empty(N, dtype=out_dtype, device=a_fp8.device) - _splitk_reduce_kernel[(triton.cdiv(N, 256),)]( - part, out, N, split_k, part.stride(0), part.stride(1), - BLOCK=256, OUT=_TL_DTYPE[out_dtype], num_warps=2, + if split_k > 1: + _splitk_reduce_kernel[(triton.cdiv(GN, 256), M)]( + part, out, M, GN, split_k, part.stride(0), part.stride(1), out.stride(0), + BLOCK=256, OUT=_TL_DTYPE[out_dtype], num_warps=2, + ) + return out + + +def grouped_w8a16_gemv(x: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor, *, block: int) -> torch.Tensor: + """Decode form of the reference's grouped bf16 einsum ``tgd,grd->tgr`` over a block-diagonal fp8 + weight: ``x [t, G, d]`` bf16 (``t <= GEMV_MAX_M``), ``weight [G * r, d]`` fp8 with ``[G * r // + block, d // block]`` e8m0 ``scale`` -> ``[t, G * r]`` bf16. The weight is dequantized in-kernel + (exact in bf16: e4m3 x 2^k), the activation is not quantized, so this is the einsum's math with a + different fp32 summation order.""" + assert x.ndim == 3 and weight.dtype == FP8 + G = x.shape[1] + sb = (scale.view(torch.uint8) if scale.dtype == torch.float8_e8m0fnu else scale).contiguous() + return _fp8_gemv(x.contiguous(), e4m3_kernel_view(weight), sb, x.dtype, block, groups=G, quant_act=False) + + +@triton.jit +def _dequant_block_fp8_kernel(w_ptr, sb_ptr, out_ptr, N, K, stride_wn, stride_sbn, stride_on, SLAB: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr): + pid_n = tl.program_id(0) + pid_k = tl.program_id(1) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + offs_k = pid_k * BLOCK_K + tl.arange(0, BLOCK_K) + mask = (offs_n < N)[:, None] & (offs_k < K)[None, :] + if e4m3_native_cx(): + w = tl.load(w_ptr + offs_n[:, None] * stride_wn + offs_k[None, :], mask=mask, other=0.0).to(tl.float32) + else: + w = e4m3_u8_to_f32(tl.load(w_ptr + offs_n[:, None] * stride_wn + offs_k[None, :], mask=mask, other=0)) + code = tl.load(sb_ptr + (offs_n // SLAB)[:, None] * stride_sbn + (offs_k // SLAB)[None, :], mask=mask, other=0) + tl.store(out_ptr + offs_n[:, None] * stride_on + offs_k[None, :], (w * tl.exp2(code.to(tl.float32) - 127.0)).to(tl.bfloat16), mask=mask) + + +def dequant_block_fp8(weight: torch.Tensor, scale: torch.Tensor, *, block: int) -> torch.Tensor: + """``[N, K]`` block-scaled e4m3 -> bf16 (exact), one pass; the prefill-side companion of + :func:`grouped_w8a16_gemv` for weights the reference consumes through a bf16 einsum.""" + N, K = weight.shape + sb = (scale.view(torch.uint8) if scale.dtype == torch.float8_e8m0fnu else scale).contiguous() + out = torch.empty((N, K), dtype=torch.bfloat16, device=weight.device) + _dequant_block_fp8_kernel[(triton.cdiv(N, 64), triton.cdiv(K, 256))]( + e4m3_kernel_view(weight), sb, out, N, K, weight.stride(0), sb.stride(0), out.stride(0), + SLAB=block, BLOCK_N=64, BLOCK_K=256, num_warps=4, ) return out @@ -353,36 +465,41 @@ def block_fp8_linear( weight: torch.Tensor, scale: torch.Tensor, bias: torch.Tensor | None = None, + *, + block: int = 128, ) -> torch.Tensor: """``y = act_quant(x) @ weight^T`` (reference FP8 path). ``x``: ``[..., K]`` bf16; ``weight``: ``[N, K]`` float8_e4m3fn; ``scale``: - ``[N//128, K//128]`` float8_e8m0fnu (weight block scale). Activation is quantized - to FP8 with a per-128 ue8m0 scale; the GEMM applies both scales per 128-K block. + ``[N//block, K//block]`` float8_e8m0fnu (weight block scale). Activation is quantized + to FP8 with a per-``block`` ue8m0 scale; the GEMM applies both scales per ``block``-K slab. """ assert weight.dtype == FP8 *lead, K = x.shape N = weight.shape[0] assert weight.shape[1] == K - assert K % 128 == 0 and N % 128 == 0, (N, K) + assert K % block == 0 and N % block == 0, (N, K, block) + assert tuple(scale.shape) == (N // block, K // block), (tuple(scale.shape), N, K, block) compute_dtype = x.dtype if x.dtype in _TL_DTYPE else torch.bfloat16 sb = scale.view(torch.uint8) if scale.dtype == torch.float8_e8m0fnu else scale sb = sb.contiguous() w = e4m3_kernel_view(weight) - a_fp8, sa = act_quant_fp8(x, 128) # [M,K] fp8, [M,K//128] e8m0 codes - M = a_fp8.shape[0] - - if M == 1: - out = _fp8_act_gemv(a_fp8[0], sa[0], w, sb, compute_dtype).reshape(*lead, N) + x2d = x.reshape(-1, K) + if x2d.shape[0] == 1: + # decode: one kernel quantizes the activation rows and streams the weight once + out = _fp8_gemv(x2d.contiguous(), w, sb, compute_dtype, block).reshape(*lead, N) if bias is not None: out = out + bias.to(out.dtype) return out + a_fp8, sa = act_quant_fp8(x, block) # [M,K] fp8, [M,K//block] e8m0 codes + M = a_fp8.shape[0] + out = torch.empty((M, N), dtype=compute_dtype, device=x.device) BLOCK_M = 32 BLOCK_N = 128 - BLOCK_K = 128 + BLOCK_K = block grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N)) _fp8_act_gemm_kernel[grid]( a_fp8, w, sa, sb, out, @@ -399,4 +516,4 @@ def block_fp8_linear( return out -__all__ = ["block_fp8_linear", "act_quant_fp8", "act_quant_fp8_inplace", "fp4_act_quant_inplace"] +__all__ = ["block_fp8_linear", "act_quant_fp8", "act_quant_fp8_inplace", "fp4_act_quant_inplace", "grouped_w8a16_gemv", "dequant_block_fp8", "GEMV_MAX_M"] diff --git a/python/freetoken/kernel/triton/dsv4/fused_moe.py b/python/freetoken/kernel/triton/dsv4/fused_moe.py index 6eea606c36..0adfbdd34b 100644 --- a/python/freetoken/kernel/triton/dsv4/fused_moe.py +++ b/python/freetoken/kernel/triton/dsv4/fused_moe.py @@ -24,6 +24,9 @@ import triton import triton.language as tl +from freetoken.kernel.triton.dsv4.fp8_linear import _log2_ceil +from freetoken.kernel.triton.e4m3_compat import e4m3_native_cx, round_e4m3 + _E2M1_VALUES = [ 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0, @@ -71,6 +74,12 @@ def _decode_dsfp4_moe_kernel( n_mask = offs_n < N slot = tl.load(topk_ids_ptr + token_id * stride_tid_m + route_k * stride_tid_k).to(tl.int64) + c_ptrs = c_ptr + token_id * stride_cm + route_k * stride_ck + offs_n * stride_cn + if slot < 0: + # an inactive route (hybrid decode computes it on the CPU): zero contribution, and no read of + # any slot -- a slot's storage may be uninitialized, and 0 * NaN is not 0 + tl.store(c_ptrs, tl.zeros((BLOCK_SIZE_N,), dtype=compute_type), mask=(route_id < total_routes) & n_mask) + return a_row = route_id if A_ROW_IS_ROUTE else token_id a_base = a_ptr + a_row * stride_am @@ -113,7 +122,6 @@ def _decode_dsfp4_moe_kernel( weight = tl.load(topk_weights_ptr + token_id * stride_tw_m + route_k * stride_tw_k) accumulator = accumulator * weight - c_ptrs = c_ptr + token_id * stride_cm + route_k * stride_ck + offs_n * stride_cn tl.store(c_ptrs, accumulator.to(compute_type), mask=(route_id < total_routes) & n_mask) @@ -124,11 +132,15 @@ def _swiglu_kernel( R, I, limit, stride_gr, stride_gi, stride_or, stride_oi, BLOCK: tl.constexpr, HAS_LIMIT: tl.constexpr, compute_type: tl.constexpr, + ACT_BLOCK: tl.constexpr, # > 0: fp8 (e4m3, ue8m0 per ACT_BLOCK) round-trip of the stored value, fused ): """Fused SwiGLU: ``out = silu(min(gate, limit)) * clamp(up, -limit, limit)``. Collapses the gate/up split + two clamps + silu + mul + the fp32 round-trip - (6 elementwise launches over [R, I]) between the two FP4 GEMVs into one pass.""" + (6 elementwise launches over [R, I]) between the two FP4 GEMVs into one pass. With + ``ACT_BLOCK`` the down projection's activation quant (``act_quant_fp8_inplace`` on the bf16 + value: ``s = 2**ceil(log2(max(|x|, 1e-4) / 448))``, e4m3 rounding, dequant) is applied to the + same tile before the store -- one launch and one memory pass fewer, bit-identical.""" row = tl.program_id(0) cb = tl.program_id(1) offs = cb * BLOCK + tl.arange(0, BLOCK) @@ -139,6 +151,17 @@ def _swiglu_kernel( g = tl.minimum(g, limit) u = tl.minimum(tl.maximum(u, -limit), limit) act = (g * tl.sigmoid(g)) * u + if ACT_BLOCK > 0: + tl.static_assert(BLOCK % ACT_BLOCK == 0, "the tile holds whole quant blocks") + xs = tl.reshape(act.to(compute_type).to(tl.float32), (BLOCK // ACT_BLOCK, ACT_BLOCK)) # the value the store would hold + amax = tl.maximum(tl.max(tl.abs(xs), axis=1), 1e-4) + sc = tl.exp2(_log2_ceil(amax * (1.0 / 448.0)).to(tl.float32)) + q = tl.clamp(xs / sc[:, None], -448.0, 448.0) + if e4m3_native_cx(): + q = q.to(tl.float8e4nv).to(tl.float32) + else: + q = round_e4m3(q) + act = tl.reshape(q * sc[:, None], (BLOCK,)) tl.store(out_ptr + row * stride_or + offs * stride_oi, act.to(compute_type), mask=mask) @@ -274,20 +297,22 @@ def _compute_type(dtype: torch.dtype): return {torch.bfloat16: tl.bfloat16, torch.float16: tl.float16, torch.float32: tl.float32}[dtype] -def fused_swiglu(gate_up: torch.Tensor, limit: float) -> torch.Tensor: - """``[..., 2I] -> [..., I]`` SwiGLU in a single kernel (see ``_swiglu_kernel``).""" +def fused_swiglu(gate_up: torch.Tensor, limit: float, *, act_block: int = 0) -> torch.Tensor: + """``[..., 2I] -> [..., I]`` SwiGLU in a single kernel (see ``_swiglu_kernel``). ``act_block > 0`` + also applies the down projection's fp8 activation round-trip (block ``act_block``) in the same pass.""" *lead, two_I = gate_up.shape I = two_I // 2 gu = gate_up.reshape(-1, two_I) R = gu.shape[0] out = torch.empty((R, I), dtype=gate_up.dtype, device=gate_up.device) BLOCK = 1024 + assert act_block == 0 or (I % act_block == 0 and BLOCK % act_block == 0), (I, act_block) grid = (R, triton.cdiv(I, BLOCK)) _swiglu_kernel[grid]( gu, out, R, I, float(limit), gu.stride(0), gu.stride(1), out.stride(0), out.stride(1), BLOCK=BLOCK, HAS_LIMIT=limit > 0, - compute_type=_compute_type(gate_up.dtype), num_warps=4, + compute_type=_compute_type(gate_up.dtype), ACT_BLOCK=act_block, num_warps=4, ) return out.reshape(*lead, I) diff --git a/python/freetoken/kernel/triton/dsv4/router.py b/python/freetoken/kernel/triton/dsv4/router.py new file mode 100644 index 0000000000..17e3506a19 --- /dev/null +++ b/python/freetoken/kernel/triton/dsv4/router.py @@ -0,0 +1,83 @@ +"""Fused MoE gate selection (DeepSeek-V4 / V4.1 ``Gate``): score activation, selection bias, top-k, +weight normalization and route scale in one program per token. + + s = act(scores) act: softmax | sigmoid | sqrt(softplus) + ids = topk(s + bias) selection is biased, the weights are not + w = s[ids] [/ (sum + 1e-20)] * route_scale + +One launch replaces the eight or so elementwise / top-k / gather launches of the torch expression -- +at decode each of those moves a few hundred floats and costs its launch floor. Ties in ``s + bias`` +go to the lowest expert id (torch.topk leaves them unspecified). +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl +from triton.language.extra import libdevice + +_SCORE = {"softmax": 0, "sigmoid": 1, "sqrtsoftplus": 2} + + +@triton.jit +def _router_select_kernel( + scores_ptr, bias_ptr, alternate_bias_ptr, alternate_mask_ptr, w_ptr, id_ptr, E, route_scale, + stride_s, stride_w, stride_id, + SCORE: tl.constexpr, TOPK: tl.constexpr, BLOCK_K: tl.constexpr, NORMALIZE: tl.constexpr, BLOCK_E: tl.constexpr, ALTERNATE: tl.constexpr, +): + t = tl.program_id(0).to(tl.int64) + offs = tl.arange(0, BLOCK_E) + mask = offs < E + s = tl.load(scores_ptr + t * stride_s + offs, mask=mask, other=float("-inf")).to(tl.float32) + if SCORE == 0: + s = tl.exp(s - tl.max(s)) + s = s / tl.sum(tl.where(mask, s, 0.0)) + elif SCORE == 1: + s = tl.sigmoid(s) + else: + s = tl.sqrt(tl.where(s > 20.0, s, libdevice.log1p(tl.exp(s)))) # torch softplus (threshold 20) + if ALTERNATE: + bias_ptr = tl.where(tl.load(alternate_mask_ptr + t), alternate_bias_ptr, bias_ptr) + bias = tl.load(bias_ptr + offs, mask=mask, other=0.0).to(tl.float32) + key = tl.where(mask, s + bias, float("-inf")) + slots = tl.arange(0, BLOCK_K) + picked_w = tl.zeros((BLOCK_K,), dtype=tl.float32) + picked_id = tl.zeros((BLOCK_K,), dtype=tl.int32) + for k in tl.static_range(TOPK): + best = tl.max(key) + idx = tl.min(tl.where(key == best, offs, BLOCK_E)) # lowest id among ties + w = tl.sum(tl.where(offs == idx, s, 0.0)) + picked_w = tl.where(slots == k, w, picked_w) + picked_id = tl.where(slots == k, idx.to(tl.int32), picked_id) + key = tl.where(offs == idx, float("-inf"), key) + if NORMALIZE: + picked_w = picked_w / (tl.sum(picked_w) + 1e-20) + picked_w = picked_w * route_scale + tl.store(w_ptr + t * stride_w + slots, picked_w, mask=slots < TOPK) + tl.store(id_ptr + t * stride_id + slots, picked_id, mask=slots < TOPK) + + +def router_select( + scores: torch.Tensor, bias: torch.Tensor, topk: int, *, score_func: str, normalize: bool, route_scale: float, + alternate_bias: torch.Tensor | None = None, alternate_mask: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """``scores [T, E]`` fp32 (pre-activation gate logits), ``bias [E]`` -> ``(weights [T, topk] fp32, + ids [T, topk] int32)`` in selection order (best first).""" + T, E = scores.shape + assert scores.dtype == torch.float32 and scores.stride(1) == 1, (scores.dtype, scores.stride()) + weights = torch.empty((T, topk), dtype=torch.float32, device=scores.device) + ids = torch.empty((T, topk), dtype=torch.int32, device=scores.device) + if T == 0: + return weights, ids + _router_select_kernel[(T,)]( + scores, bias, alternate_bias if alternate_bias is not None else bias, + alternate_mask if alternate_mask is not None else scores, weights, ids, E, float(route_scale), + scores.stride(0), weights.stride(0), ids.stride(0), + SCORE=_SCORE[score_func], TOPK=topk, BLOCK_K=triton.next_power_of_2(topk), NORMALIZE=normalize and topk > 1, + BLOCK_E=triton.next_power_of_2(E), ALTERNATE=alternate_mask is not None, num_warps=4, + ) + return weights, ids + + +__all__ = ["router_select"] diff --git a/python/freetoken/kernel/triton/dsv41/__init__.py b/python/freetoken/kernel/triton/dsv41/__init__.py new file mode 100644 index 0000000000..00ec72daba --- /dev/null +++ b/python/freetoken/kernel/triton/dsv41/__init__.py @@ -0,0 +1,9 @@ +"""Triton kernels for DeepSeek-V4.1 Compressed Sparse Attention 2 over packed KV pools. + +``row_format`` device helpers: load a tile of packed rows (bf16 / fp8+e8m0 / fp4+e4m3 / fp4+e8m0) + and dequantize it to fp32, encode / decode e2m1 nibbles. +``pack`` host wrappers: quantize rows into a packed pool (``pack_rows``) and read them back + (``unpack_rows``), matching the reference ``act_quant`` / ``fp4_act_quant`` numerics. +``sparse_attn`` paged sparse attention gathering window (fp8) and compressed (fp4) rows. +``indexer`` Lightning-indexer logits over fp4 index keys, full-range or candidate-restricted. +""" diff --git a/python/freetoken/kernel/triton/dsv41/indexer.py b/python/freetoken/kernel/triton/dsv41/indexer.py new file mode 100644 index 0000000000..3721315837 --- /dev/null +++ b/python/freetoken/kernel/triton/dsv41/indexer.py @@ -0,0 +1,176 @@ +"""Lightning-Indexer logits over PACKED index keys (DeepSeek-V4.1, Full and Reindex modes). + + logits[b, s, j] = sum_h relu(q[b, s, h, :] . k(b, pos_j)) * weights[b, s, h] + +``q`` carries the ``index_n_heads`` fp4-baked index queries, ``k`` is the single shared index key of +a compressed position, gathered from the kv-source's packed pool and dequantized in-kernel. The +row of compressed position ``t`` is derived in-kernel from the request's full-token locs (the page +table row at prefill, the decode snapshot): ``locs[b, t * ratio] // ratio`` (negative = no row) -- +no ``[B, T]`` row table is materialized per layer. ``weights`` folds ``softmax_scale * +n_heads**-0.5`` in. + +Two position sets, one kernel: + +* **full range** (Full mode): position ``j`` is the compressed position ``j``; the output is + ``[B, S, T]`` over the staged width ``T``, but only the live tiles (columns below ``live[b, s]`` + rounded up to ``BLOCK_T``, the rest of the last tile ``-inf``) are WRITTEN. A fixed grid of worker + programs strides over the live tiles, so both the work and the launch geometry follow the live + history rather than the staged width; consumers (``dsv41/topk.py``) read below ``live`` only, so + the dead area needs no fill. +* **candidate pool** (Reindex mode, Hierarchical Sparse Indexer): position ``j`` is + ``candidates[b, s, j]`` (a compressed position, ``-1`` = empty slot); the output is ``[B, S, NC]`` + aligned with the candidate list and fully written (``-inf`` for empty / unreachable slots), so + a top-k over it indexes straight back into the list. + +Causality lives in the kernel: ``live[b, s]`` is the number of compressed positions query ``s`` may +see, and every position at or past it (or with no row) scores ``-inf``. The dot product, weighted +scores and head reduction round to bf16 at the same boundaries as the reference expression. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from freetoken.kvcache.dsv4.v41_row_format import RowFormat + +from .row_format import load_rows + +BLOCK_T = 64 + + +@triton.jit +def _score_tile( + q, w, pool_ptr, locs_row, live, pos, store_mask, out_ptrs, + n_locs, RATIO: tl.constexpr, D: tl.constexpr, ROW_BYTES: tl.constexpr, FMT: tl.constexpr, +): + """Score one tile of compressed positions ``pos`` against the query (``q [BLOCK_H, D]`` bf16, + ``w [BLOCK_H]`` fp32) and store ``[BLOCK_T]`` logits (``-inf`` where unreachable).""" + pos_ok = store_mask & (pos >= 0) & (pos < live) + full_idx = pos * RATIO + full = tl.load(locs_row + tl.maximum(full_idx, 0), mask=pos_ok & (full_idx < n_locs), other=-1) + rows = tl.where(full >= 0, full // RATIO, -1).to(tl.int32) + valid = pos_ok & (rows >= 0) + offs_d = tl.arange(0, D) + k = load_rows(pool_ptr, rows, valid, offs_d, D, ROW_BYTES, FMT).to(tl.bfloat16) # [BLOCK_T, D] + score = tl.dot(q, tl.trans(k)).to(tl.bfloat16).to(tl.float32) + score = (tl.maximum(score, 0.0) * w[:, None]).to(tl.bfloat16).to(tl.float32) + logits = tl.sum(score, axis=0).to(tl.bfloat16).to(tl.float32) + tl.store(out_ptrs, tl.where(valid, logits, float("-inf")), mask=store_mask) + + +@triton.jit +def _indexer_logits_packed_kernel( + q_ptr, w_ptr, pool_ptr, locs_ptr, live_ptr, cand_ptr, out_ptr, + T, NC, n_locs, + stride_qb, stride_qs, stride_qh, stride_qd, + stride_wb, stride_ws, stride_wh, + stride_lb, stride_lt, + stride_vb, stride_vs, + stride_cb, stride_cs, stride_cn, + stride_ob, stride_os, stride_on, + RATIO: tl.constexpr, + H: tl.constexpr, + D: tl.constexpr, + ROW_BYTES: tl.constexpr, + FMT: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_T: tl.constexpr, + HAS_CAND: tl.constexpr, + N_WORKERS: tl.constexpr, +): + pid_s = tl.program_id(0) + pid_b = tl.program_id(1) + pid_t = tl.program_id(2) + + live = tl.load(live_ptr + pid_b * stride_vb + pid_s * stride_vs) + locs_row = locs_ptr + pid_b * stride_lb + out_row = out_ptr + pid_b * stride_ob + pid_s * stride_os + offs_d = tl.arange(0, D) + offs_h = tl.arange(0, BLOCK_H) + h_mask = offs_h < H + q = tl.load( + q_ptr + pid_b * stride_qb + pid_s * stride_qs + offs_h[:, None] * stride_qh + offs_d[None, :] * stride_qd, + mask=h_mask[:, None], other=0.0, + ).to(tl.bfloat16) + w = tl.load(w_ptr + pid_b * stride_wb + pid_s * stride_ws + offs_h * stride_wh, mask=h_mask, other=0.0).to(tl.bfloat16).to(tl.float32) + offs = tl.arange(0, BLOCK_T) + + if HAS_CAND: + # one program per candidate tile; the list is fixed-width and fully written + offs_t = pid_t * BLOCK_T + offs + store_mask = offs_t < NC + pos = tl.load(cand_ptr + pid_b * stride_cb + pid_s * stride_cs + offs_t * stride_cn, mask=store_mask, other=-1) + _score_tile(q, w, pool_ptr, locs_row, live, pos, store_mask, out_row + offs_t * stride_on, n_locs, RATIO, D, ROW_BYTES, FMT) + else: + # worker pid_t scores live tiles pid_t, pid_t + N_WORKERS, ...; nothing past live is written + for tile in range(pid_t, tl.cdiv(tl.minimum(live, T), BLOCK_T), N_WORKERS): + offs_t = tile * BLOCK_T + offs + store_mask = offs_t < T + _score_tile(q, w, pool_ptr, locs_row, live, offs_t, store_mask, out_row + offs_t * stride_on, n_locs, RATIO, D, ROW_BYTES, FMT) + + +def indexer_logits_packed( + q: torch.Tensor, # [B, S, H, D] bf16 (fp4 round-tripped index queries) + weights: torch.Tensor, # [B, S, H] (softmax_scale * n_heads**-0.5 folded in) + k_pool: torch.Tensor, # [R, fmt.row_bytes(D)] uint8, the kv-source's packed index keys + fmt: RowFormat, + locs: torch.Tensor, # [B, W] int32/int64 full-token locs of each request (page-table row / decode snapshot) + ratio: int, # compressed position t -> full loc locs[b, t * ratio]; its row is loc // ratio + live: torch.Tensor, # [B, S] int32: compressed positions query (b, s) may see + T: int | None = None, # staged width of the full-range output (default: W // ratio) + candidates: torch.Tensor | None = None, # [B, S, NC] int32 compressed positions (-1 = empty) + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Head-reduced indexer logits: ``[B, S, T]`` over the full staged range (columns below ``live`` + written, the rest untouched), or ``[B, S, NC]`` aligned with ``candidates`` (fully written). + Unreachable / empty positions are ``-inf``.""" + B, S, H, D = q.shape + assert locs.ndim == 2 and locs.shape[0] == B and live.shape == (B, S), (locs.shape, live.shape, (B, S)) + assert weights.shape == (B, S, H), (weights.shape, (B, S, H)) + assert D == triton.next_power_of_2(D), f"index_head_dim must be pow2, got {D}" + assert k_pool.dtype == torch.uint8 and k_pool.shape[1] == fmt.row_bytes(D) and k_pool.is_contiguous() + assert locs.dtype in (torch.int32, torch.int64) and locs.stride(1) == 1, (locs.dtype, locs.stride()) + n_locs = locs.shape[1] + if T is None: + T = n_locs // ratio + live = live.to(torch.int32) + has_cand = candidates is not None + if has_cand: + assert candidates.shape[:2] == (B, S), (candidates.shape, (B, S)) + cand = candidates.to(torch.int32) + NC = cand.shape[2] + cs = cand.stride() + else: + cand, NC, cs = live, 0, (0, 0, 0) + n_out = NC if has_cand else T + if out is None: + out = torch.empty((B, S, n_out), dtype=torch.float32, device=q.device) + assert out.shape == (B, S, n_out) and out.dtype == torch.float32 + if n_out == 0 or S == 0: + return out + + n_tiles = triton.cdiv(n_out, BLOCK_T) + # A device-sized worker grid strides over live history without launching a + # program for every reserved tile. Independent queries share the SM budget. + sm_count = torch.cuda.get_device_properties(q.device).multi_processor_count + workers = n_tiles if has_cand else min(n_tiles, max(1, triton.cdiv(sm_count, B * S))) + grid = (S, B, workers) + _indexer_logits_packed_kernel[grid]( + q, weights, k_pool, locs, live, cand, out, + T, NC, n_locs, + q.stride(0), q.stride(1), q.stride(2), q.stride(3), + weights.stride(0), weights.stride(1), weights.stride(2), + locs.stride(0), locs.stride(1), + live.stride(0), live.stride(1), + cs[0], cs[1], cs[2], + out.stride(0), out.stride(1), out.stride(2), + RATIO=ratio, H=H, D=D, ROW_BYTES=fmt.row_bytes(D), FMT=fmt.code, + BLOCK_H=triton.next_power_of_2(H), BLOCK_T=BLOCK_T, HAS_CAND=has_cand, N_WORKERS=workers, + num_warps=4, num_stages=2, + ) + return out + + +__all__ = ["indexer_logits_packed", "BLOCK_T"] diff --git a/python/freetoken/kernel/triton/dsv41/pack.py b/python/freetoken/kernel/triton/dsv41/pack.py new file mode 100644 index 0000000000..47a9164e6e --- /dev/null +++ b/python/freetoken/kernel/triton/dsv41/pack.py @@ -0,0 +1,183 @@ +"""Quantize rows into a packed pool and read them back (DeepSeek-V4.1 KV storage formats). + +``pack_rows`` reproduces the reference quantizers (``inference/kernel.py``) exactly: + +* fp8_e8m0_b32 ``act_quant(x, 32, "ue8m0", e8m0)``: per 32-group ``s = 2**ceil(log2(max(amax, 1e-4) / 448))``, + ``q = round_e4m3(clamp(x / s, +-448))`` +* fp4_e4m3_b16 ``fp4_act_quant(x, 16, scale_dtype=e4m3)``: ``s = e4m3(max(amax, 6 * 2^-9) / 6)`` (round to + nearest e4m3), ``q = round_e2m1(clamp(x / s, +-6))`` +* fp4_e8m0_b32 ``fp4_act_quant(x, 32)``: ``s = 2**ceil(log2(max(amax, 6 * 2^-126) / 6))``, same e2m1 rounding +* bf16 a plain byte copy + +Rows land at ``pool[row_ids[m]]`` (or ``pool[m]`` without ids), so the same launch serves the KV write +path (scatter into slots) and plain buffer conversion. ``unpack_rows`` is the inverse gather + dequant +(bf16-rounded like the reference cache); the attention / indexer kernels use the same device helper +(``row_format.load_rows``) in-line instead of materializing rows. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from freetoken.kernel.triton.dsv4.fp8_linear import _log2_ceil, _round_fp4 +from freetoken.kernel.triton.e4m3_compat import e4m3_f32_to_u8, e4m3_native_cx, round_e4m3 +from freetoken.kvcache.dsv4.v41_row_format import BF16, RowFormat + +from .row_format import fp4_f32_to_nibble, load_rows + + +@triton.jit +def _e4m3_codes(q): + """fp32 (|q| <= 448) -> e4m3 bit pattern, RNE.""" + if e4m3_native_cx(): + return q.to(tl.float8e4nv).to(tl.uint8, bitcast=True) + else: + return e4m3_f32_to_u8(round_e4m3(q)) + + +@triton.jit +def _e4m3_round(q): + """fp32 -> nearest e4m3 value, as fp32.""" + if e4m3_native_cx(): + return q.to(tl.float8e4nv).to(tl.float32) + else: + return round_e4m3(q) + + +@triton.jit +def _pack_rows_kernel( + x_ptr, out_ptr, idx_ptr, stride_xm, + D: tl.constexpr, GROUP: tl.constexpr, ROW_BYTES: tl.constexpr, FMT: tl.constexpr, HAS_IDX: tl.constexpr, +): + """One program per input row ``m``; writes packed bytes to pool row ``idx[m]`` (or ``m``).""" + m = tl.program_id(0).to(tl.int64) + if HAS_IDX: + row = tl.load(idx_ptr + m).to(tl.int64) + else: + row = m + dst = out_ptr + row * ROW_BYTES + src = x_ptr + m * stride_xm + + if FMT == 0: + offs_d = tl.arange(0, D) + bits = tl.load(src + offs_d).to(tl.bfloat16).to(tl.uint16, bitcast=True) + tl.store(dst + 2 * offs_d, (bits & 0xFF).to(tl.uint8)) + tl.store(dst + 2 * offs_d + 1, (bits >> 8).to(tl.uint8)) + elif FMT == 1: + NG: tl.constexpr = D // GROUP + offs_d = tl.arange(0, D) + offs_g = tl.arange(0, NG) + x = tl.reshape(tl.load(src + offs_d).to(tl.float32), [NG, GROUP]) + amax = tl.maximum(tl.max(tl.abs(x), axis=1), 1e-4) + e = _log2_ceil(amax * (1.0 / 448.0)) + s = tl.exp2(e.to(tl.float32)) + q = tl.clamp(tl.math.div_rn(x, s[:, None]), -448.0, 448.0) + tl.store(dst + offs_d, tl.reshape(_e4m3_codes(q), [D])) + tl.store(dst + D + offs_g, (e + 127).to(tl.uint8)) + else: + # e2m1 pairs: even channels from one strided load, odd from another, so a byte is (lo | hi << 4) + NG: tl.constexpr = D // GROUP + HALF: tl.constexpr = GROUP // 2 + offs_h = tl.arange(0, D // 2) + offs_g = tl.arange(0, NG) + xe = tl.reshape(tl.load(src + 2 * offs_h).to(tl.float32), [NG, HALF]) + xo = tl.reshape(tl.load(src + 2 * offs_h + 1).to(tl.float32), [NG, HALF]) + amax = tl.maximum(tl.max(tl.abs(xe), axis=1), tl.max(tl.abs(xo), axis=1)) + if FMT == 2: + amax = tl.maximum(amax, 6.0 * (2.0 ** -9)) + s = _e4m3_round(amax * (1.0 / 6.0)) + tl.store(dst + D // 2 + offs_g, _e4m3_codes(s)) + else: + amax = tl.maximum(amax, 6.0 * (2.0 ** -126)) + e = _log2_ceil(amax * (1.0 / 6.0)) + s = tl.exp2(e.to(tl.float32)) + tl.store(dst + D // 2 + offs_g, (e + 127).to(tl.uint8)) + # IEEE division: the reference's CUDA ``x / s`` is correctly rounded, and with an e4m3 + # (non power-of-two) scale triton's default approximate div would land a hair off the + # e2m1 midpoints and flip their round-to-even + qe = fp4_f32_to_nibble(_round_fp4(tl.clamp(tl.math.div_rn(xe, s[:, None]), -6.0, 6.0))) + qo = fp4_f32_to_nibble(_round_fp4(tl.clamp(tl.math.div_rn(xo, s[:, None]), -6.0, 6.0))) + tl.store(dst + offs_h, tl.reshape(qe | (qo << 4), [D // 2])) + + +@triton.jit +def _unpack_rows_kernel( + pool_ptr, idx_ptr, out_ptr, M, + D: tl.constexpr, ROW_BYTES: tl.constexpr, FMT: tl.constexpr, BLOCK_T: tl.constexpr, HAS_IDX: tl.constexpr, +): + pid = tl.program_id(0) + offs_t = pid * BLOCK_T + tl.arange(0, BLOCK_T) + valid = offs_t < M + if HAS_IDX: + rows = tl.load(idx_ptr + offs_t, mask=valid, other=0) + else: + rows = offs_t + offs_d = tl.arange(0, D) + vals = load_rows(pool_ptr, rows, valid, offs_d, D, ROW_BYTES, FMT) + tl.store( + out_ptr + offs_t[:, None].to(tl.int64) * D + offs_d[None, :], + vals.to(out_ptr.dtype.element_ty), mask=valid[:, None], + ) + + +def pack_rows( + x: torch.Tensor, fmt: RowFormat, pool: torch.Tensor | None = None, row_ids: torch.Tensor | None = None, +) -> torch.Tensor: + """Quantize ``x [M, D]`` (bf16 / fp16 / fp32) into packed rows. + + With ``pool [rows, row_bytes] uint8`` and ``row_ids [M]``, row ``m`` is written to ``pool[row_ids[m]]`` + (the KV write path). Without a pool a fresh ``[M, row_bytes]`` uint8 tensor is returned. + """ + assert x.dim() == 2, x.shape + m, d = x.shape + fmt.validate_dim(d) + row_bytes = fmt.row_bytes(d) + assert d & (d - 1) == 0 and d >= 16, f"pack_rows wants a power-of-two D >= 16, got {d}" + if pool is None: + assert row_ids is None + pool = torch.empty((m, row_bytes), dtype=torch.uint8, device=x.device) + else: + assert pool.dtype == torch.uint8 and pool.shape[1] == row_bytes and pool.is_contiguous(), (pool.shape, row_bytes) + assert row_ids is not None and row_ids.shape == (m,), (None if row_ids is None else row_ids.shape, m) + if m == 0: + return pool + if x.stride(1) != 1: + x = x.contiguous() + _pack_rows_kernel[(m,)]( + x, pool, row_ids if row_ids is not None else pool, x.stride(0), + D=d, GROUP=fmt.group or d, ROW_BYTES=row_bytes, FMT=fmt.code, HAS_IDX=row_ids is not None, + num_warps=4, + ) + return pool + + +def unpack_rows( + pool: torch.Tensor, fmt: RowFormat, dim: int, row_ids: torch.Tensor | None = None, + out_dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Gather ``pool[row_ids]`` (or every row) and dequantize to ``[M, dim]`` in ``out_dtype``.""" + fmt.validate_dim(dim) + assert pool.dtype == torch.uint8 and pool.shape[1] == fmt.row_bytes(dim) and pool.is_contiguous(), pool.shape + m = pool.shape[0] if row_ids is None else row_ids.numel() + out = torch.empty((m, dim), dtype=out_dtype, device=pool.device) + if m == 0: + return out + block_t = 16 + _unpack_rows_kernel[(triton.cdiv(m, block_t),)]( + pool, row_ids if row_ids is not None else pool, out, m, + D=dim, ROW_BYTES=fmt.row_bytes(dim), FMT=fmt.code, BLOCK_T=block_t, HAS_IDX=row_ids is not None, + num_warps=4, + ) + return out + + +def quant_roundtrip(x: torch.Tensor, fmt: RowFormat) -> torch.Tensor: + """``unpack(pack(x))`` -- the value the cache will hand back; the reference's ``inplace=True`` semantics.""" + if fmt is BF16: + return x.to(torch.bfloat16) + return unpack_rows(pack_rows(x, fmt), fmt, x.shape[1], out_dtype=torch.bfloat16) + + +__all__ = ["pack_rows", "unpack_rows", "quant_roundtrip"] diff --git a/python/freetoken/kernel/triton/dsv41/row_format.py b/python/freetoken/kernel/triton/dsv41/row_format.py new file mode 100644 index 0000000000..032d77e970 --- /dev/null +++ b/python/freetoken/kernel/triton/dsv41/row_format.py @@ -0,0 +1,92 @@ +"""Device-side row-format helpers shared by the DSV41 kernels. + +The Python-side descriptors live in ``freetoken.kvcache.dsv4.v41_row_format``; the ``FMT`` constexpr a +kernel receives is ``RowFormat.code``: + + 0 bf16 2 B / channel, no scale + 1 fp8_e8m0_b32 e4m3 byte / channel, ue8m0 scale byte per 32 channels + 2 fp4_e4m3_b16 e2m1 nibble / channel (even channel low nibble), e4m3 scale byte per 16 + 3 fp4_e8m0_b32 e2m1 nibble / channel, ue8m0 scale byte per 32 + +A row is ``[values | scales]``; ``ROW_BYTES`` is the pool's row stride. Every dequantized value +is rounded through bf16 before use -- the reference bakes its quant round-trip into a bf16 cache, +so this is what its attention / indexer kernels actually read. +""" + +from __future__ import annotations + +import triton +import triton.language as tl + +from freetoken.kernel.triton.e4m3_compat import e4m3_u8_to_f32 + +FMT_BF16 = 0 +FMT_FP8_E8M0_B32 = 1 +FMT_FP4_E4M3_B16 = 2 +FMT_FP4_E8M0_B32 = 3 + + +@triton.jit +def fp4_nibble_to_f32(nib): + """e2m1 code (uint8, 0..15) -> fp32: bit 3 sign, bits 2-1 exponent, bit 0 mantissa; + codes 0/1 are the subnormals 0 and 0.5, the rest ``(1 + m/2) * 2^(e-1)``.""" + c = (nib & 7).to(tl.int32) + neg = (nib & 8) != 0 + sub = c.to(tl.float32) * 0.5 + normal = (1.0 + 0.5 * (c & 1).to(tl.float32)) * tl.exp2(((c >> 1) - 1).to(tl.float32)) + mag = tl.where(c < 2, sub, normal) + return tl.where(neg, -mag, mag) + + +@triton.jit +def fp4_f32_to_nibble(q): + """fp32 value already on the e2m1 grid ({0, .5, 1, 1.5, 2, 3, 4, 6}, signed) -> e2m1 code. + Thresholds sit between grid points, so a value a hair off the grid still maps to its point.""" + a = tl.abs(q) + c = tl.where( + a < 1.75, + tl.where(a < 0.75, tl.where(a < 0.25, 0, 1), tl.where(a < 1.25, 2, 3)), + tl.where(a < 3.5, tl.where(a < 2.5, 4, 5), tl.where(a < 5.0, 6, 7)), + ) + return (c + tl.where(q < 0, 8, 0)).to(tl.uint8) + + +@triton.jit +def load_rows(pool_ptr, rows, valid, offs_d, D: tl.constexpr, ROW_BYTES: tl.constexpr, FMT: tl.constexpr): + """Gather + dequantize ``[T, D]`` fp32 from a packed pool. + + ``pool_ptr`` uint8 pool base; ``rows`` [T] row ids (any int width; masked-out rows may be -1); + ``valid`` [T] bool; ``offs_d`` [D] = ``tl.arange(0, D)``. Invalid rows read as 0. + """ + row_i64 = tl.maximum(rows, 0).to(tl.int64) + base = pool_ptr + row_i64[:, None] * ROW_BYTES + m = valid[:, None] & (offs_d >= 0)[None, :] + if FMT == 0: + lo = tl.load(base + 2 * offs_d[None, :], mask=m, other=0).to(tl.uint16) + hi = tl.load(base + 2 * offs_d[None, :] + 1, mask=m, other=0).to(tl.uint16) + val = (lo | (hi << 8)).to(tl.bfloat16, bitcast=True).to(tl.float32) + elif FMT == 1: + v = tl.load(base + offs_d[None, :], mask=m, other=0) + s = tl.load(base + D + (offs_d // 32)[None, :], mask=m, other=127) + val = e4m3_u8_to_f32(v) * tl.exp2(s.to(tl.float32) - 127.0) + else: + b = tl.load(base + (offs_d // 2)[None, :], mask=m, other=0) + nib = tl.where((offs_d % 2 == 0)[None, :], b & 0xF, b >> 4) + q = fp4_nibble_to_f32(nib) + if FMT == 2: + s = e4m3_u8_to_f32(tl.load(base + D // 2 + (offs_d // 16)[None, :], mask=m, other=0)) + else: + s = tl.exp2(tl.load(base + D // 2 + (offs_d // 32)[None, :], mask=m, other=127).to(tl.float32) - 127.0) + val = q * s + return val.to(tl.bfloat16).to(tl.float32) + + +__all__ = [ + "FMT_BF16", + "FMT_FP8_E8M0_B32", + "FMT_FP4_E4M3_B16", + "FMT_FP4_E8M0_B32", + "fp4_nibble_to_f32", + "fp4_f32_to_nibble", + "load_rows", +] diff --git a/python/freetoken/kernel/triton/dsv41/sparse_attn.py b/python/freetoken/kernel/triton/dsv41/sparse_attn.py new file mode 100644 index 0000000000..9d6b0816db --- /dev/null +++ b/python/freetoken/kernel/triton/dsv41/sparse_attn.py @@ -0,0 +1,323 @@ +"""Sparse gathered-KV flash attention over PACKED pools (DeepSeek-V4.1). + +Same contract as ``kernel/triton/dsv4/sparse_attn.py`` -- each query attends to a per-query list of +GLOBAL slots laid out ``[window part | compressed part]`` plus a per-head attention sink, K == V is one +shared latent head -- but the two pools are byte pools in (possibly different) ``RowFormat``s: the +window ring holds fp8 rows, the compressed KV holds fp4 rows, and every gathered row is dequantized +in-kernel (``row_format.load_rows``) to the bf16 value the reference cache would have held. + +Because the two pools dequantize differently, a tile never mixes them: the window part is +``n_window`` columns with ``n_window % BLOCK_T == 0`` (the caller pads it with ``-1``), so the kernel +walks the window tiles and then the compressed tiles in two loops, each with a single load site. +Both dots take bf16 operands with fp32 accumulation, like the reference ``sparse_attn`` kernel. + +Two implementations, picked from the launch shape (no knob): + * prefill (``m`` > 1): one program per (query, request, head block) + * decode (``m`` == 1): flash-decoding split over the candidate axis, merged through log-sum-exp + (the sink joins in the merge). Not bit-identical to the single-program path. + +``cmp_counts`` (device ``[b, m]`` int32) bounds the compressed columns a CUDA-graph replay visits; +without it the full ``topk_idxs`` width is walked. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from freetoken.kvcache.dsv4.v41_row_format import RowFormat + +from .row_format import load_rows + +BLOCK_H = 16 +BLOCK_T = 32 +MAX_SPLITS = 32 +# Decode splits at tile boundaries, bounded by the device's SM count below. +MIN_TILES_PER_SPLIT = 1 + + +@triton.jit +def _attend_tile(q, kv, valid, scale, m_i, l_i, acc): + """One online-softmax step over a ``[BLOCK_T, D]`` fp32 KV tile (bf16-exact values).""" + kv16 = kv.to(tl.bfloat16) + scores = tl.dot(q, tl.trans(kv16)) * scale # [BLOCK_H, BLOCK_T] fp32 + scores = tl.where(valid[None, :], scores, -float("inf")) + m_new = tl.maximum(m_i, tl.max(scores, axis=1)) + # an all-masked tile keeps m_new at -inf; short-circuit so -inf - -inf never forms + alpha = tl.where(m_new == -float("inf"), 1.0, tl.exp(m_i - m_new)) + p = tl.where(valid[None, :], tl.exp(scores - m_new[:, None]), 0.0) + l_i = l_i * alpha + tl.sum(p, axis=1) + acc = acc * alpha[:, None] + tl.dot(p.to(tl.bfloat16), kv16) + return m_new, l_i, acc + + +@triton.jit +def _sparse_attn_packed_kernel( + q_ptr, win_ptr, cmp_ptr, o_ptr, sink_ptr, idx_ptr, cnt_ptr, + scale, + H, TOPK, N_WINDOW, + stride_qb, stride_qm, stride_qh, stride_qd, + stride_ob, stride_om, stride_oh, stride_od, + stride_ib, stride_im, stride_it, + stride_nb, stride_nm, + D: tl.constexpr, + WIN_FMT: tl.constexpr, WIN_ROW_BYTES: tl.constexpr, + CMP_FMT: tl.constexpr, CMP_ROW_BYTES: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_T: tl.constexpr, + HAS_COUNTS: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_b = tl.program_id(1) + pid_h = tl.program_id(2) + + offs_h = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) + h_mask = offs_h < H + offs_d = tl.arange(0, D) + + q_ptrs = q_ptr + pid_b * stride_qb + pid_m * stride_qm + offs_h[:, None] * stride_qh + offs_d[None, :] * stride_qd + q = tl.load(q_ptrs, mask=h_mask[:, None], other=0.0).to(tl.bfloat16) # [BLOCK_H, D] + + m_i = tl.full((BLOCK_H,), -float("inf"), dtype=tl.float32) + l_i = tl.zeros((BLOCK_H,), dtype=tl.float32) + acc = tl.zeros((BLOCK_H, D), dtype=tl.float32) + + n_active = TOPK + if HAS_COUNTS: + n_active = N_WINDOW + tl.load(cnt_ptr + pid_b * stride_nb + pid_m * stride_nm) + + idx_base = idx_ptr + pid_b * stride_ib + pid_m * stride_im + for start in range(0, N_WINDOW, BLOCK_T): + offs_t = start + tl.arange(0, BLOCK_T) + idxs = tl.load(idx_base + offs_t * stride_it) + valid = idxs >= 0 + kv = load_rows(win_ptr, idxs, valid, offs_d, D, WIN_ROW_BYTES, WIN_FMT) + m_i, l_i, acc = _attend_tile(q, kv, valid, scale, m_i, l_i, acc) + for start in range(N_WINDOW, n_active, BLOCK_T): + offs_t = start + tl.arange(0, BLOCK_T) + t_mask = offs_t < n_active + idxs = tl.load(idx_base + offs_t * stride_it, mask=t_mask, other=-1) + valid = idxs >= 0 + kv = load_rows(cmp_ptr, idxs, valid, offs_d, D, CMP_ROW_BYTES, CMP_FMT) + m_i, l_i, acc = _attend_tile(q, kv, valid, scale, m_i, l_i, acc) + + sink = tl.load(sink_ptr + offs_h, mask=h_mask, other=0.0).to(tl.float32) + l_i = l_i + tl.exp(sink - m_i) + o = acc / l_i[:, None] + + o_ptrs = o_ptr + pid_b * stride_ob + pid_m * stride_om + offs_h[:, None] * stride_oh + offs_d[None, :] * stride_od + tl.store(o_ptrs, o.to(o_ptr.dtype.element_ty), mask=h_mask[:, None]) + + +@triton.jit +def _sparse_attn_packed_splitk_kernel( + q_ptr, win_ptr, cmp_ptr, mid_o_ptr, mid_lse_ptr, idx_ptr, cnt_ptr, + scale, + H, TOPK, N_WINDOW, + stride_qb, stride_qm, stride_qh, stride_qd, + stride_mb, stride_mm, stride_mh, stride_ms, stride_md, + stride_lb, stride_lm, stride_lh, stride_ls, + stride_ib, stride_im, stride_it, + stride_nb, stride_nm, + D: tl.constexpr, + WIN_FMT: tl.constexpr, WIN_ROW_BYTES: tl.constexpr, + CMP_FMT: tl.constexpr, CMP_ROW_BYTES: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_T: tl.constexpr, + HAS_COUNTS: tl.constexpr, + NUM_SPLITS: tl.constexpr, +): + """Stage 1: each program reduces one BLOCK_T-aligned slice of the candidate list and writes + its normalized partial output + log-sum-exp. A slice may straddle the window / compressed + boundary; it is walked as (window tiles, compressed tiles) since N_WINDOW is tile-aligned.""" + pid_ms = tl.program_id(0) + pid_m = pid_ms // NUM_SPLITS + split_id = pid_ms % NUM_SPLITS + pid_b = tl.program_id(1) + pid_h = tl.program_id(2) + + offs_h = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) + h_mask = offs_h < H + offs_d = tl.arange(0, D) + + n_active = TOPK + if HAS_COUNTS: + n_active = N_WINDOW + tl.load(cnt_ptr + pid_b * stride_nb + pid_m * stride_nm) + + per_split = tl.cdiv(tl.cdiv(n_active, NUM_SPLITS), BLOCK_T) * BLOCK_T + split_start = per_split * split_id + split_end = tl.minimum(split_start + per_split, n_active) + + m_i = tl.full((BLOCK_H,), -float("inf"), dtype=tl.float32) + l_i = tl.zeros((BLOCK_H,), dtype=tl.float32) + acc = tl.zeros((BLOCK_H, D), dtype=tl.float32) + + if split_end > split_start: + q_ptrs = ( + q_ptr + pid_b * stride_qb + pid_m * stride_qm + + offs_h[:, None] * stride_qh + offs_d[None, :] * stride_qd + ) + q = tl.load(q_ptrs, mask=h_mask[:, None], other=0.0).to(tl.bfloat16) + idx_base = idx_ptr + pid_b * stride_ib + pid_m * stride_im + + win_end = tl.minimum(split_end, N_WINDOW) + for start in range(split_start, win_end, BLOCK_T): + offs_t = start + tl.arange(0, BLOCK_T) + idxs = tl.load(idx_base + offs_t * stride_it) + valid = idxs >= 0 + kv = load_rows(win_ptr, idxs, valid, offs_d, D, WIN_ROW_BYTES, WIN_FMT) + m_i, l_i, acc = _attend_tile(q, kv, valid, scale, m_i, l_i, acc) + cmp_start = tl.maximum(split_start, N_WINDOW) + for start in range(cmp_start, split_end, BLOCK_T): + offs_t = start + tl.arange(0, BLOCK_T) + t_mask = offs_t < split_end + idxs = tl.load(idx_base + offs_t * stride_it, mask=t_mask, other=-1) + valid = idxs >= 0 + kv = load_rows(cmp_ptr, idxs, valid, offs_d, D, CMP_ROW_BYTES, CMP_FMT) + m_i, l_i, acc = _attend_tile(q, kv, valid, scale, m_i, l_i, acc) + + out = tl.where(l_i[:, None] == 0.0, 0.0, acc / l_i[:, None]) + lse = tl.where(l_i == 0.0, -float("inf"), m_i + tl.log(l_i)) + + mid_base = ( + mid_o_ptr + pid_b * stride_mb + pid_m * stride_mm + + offs_h[:, None] * stride_mh + split_id * stride_ms + offs_d[None, :] * stride_md + ) + tl.store(mid_base, out, mask=h_mask[:, None]) + lse_base = ( + mid_lse_ptr + pid_b * stride_lb + pid_m * stride_lm + + offs_h * stride_lh + split_id * stride_ls + ) + tl.store(lse_base, lse, mask=h_mask) + + +@triton.jit +def _sparse_attn_splitk_merge_kernel( + mid_o_ptr, mid_lse_ptr, o_ptr, sink_ptr, + stride_mb, stride_mm, stride_mh, stride_ms, stride_md, + stride_lb, stride_lm, stride_lh, stride_ls, + stride_ob, stride_om, stride_oh, stride_od, + D: tl.constexpr, + NUM_SPLITS: tl.constexpr, +): + """Stage 2: log-sum-exp merge over the splits; the attention sink joins here once, as a null + key with logit ``attn_sink[h]`` and zero value. One program per (query, request, head).""" + pid_m = tl.program_id(0) + pid_b = tl.program_id(1) + pid_h = tl.program_id(2) + offs_d = tl.arange(0, D) + + m_i = tl.load(sink_ptr + pid_h).to(tl.float32) + l_i = 1.0 + acc = tl.zeros((D,), dtype=tl.float32) + + mid_base = mid_o_ptr + pid_b * stride_mb + pid_m * stride_mm + pid_h * stride_mh + offs_d * stride_md + lse_base = mid_lse_ptr + pid_b * stride_lb + pid_m * stride_lm + pid_h * stride_lh + + for split_id in tl.range(0, NUM_SPLITS, num_stages=2): + partial = tl.load(mid_base + split_id * stride_ms) + lse = tl.load(lse_base + split_id * stride_ls) + m_new = tl.maximum(m_i, lse) + alpha = tl.exp(m_i - m_new) + beta = tl.where(lse == -float("inf"), 0.0, tl.exp(lse - m_new)) + acc = acc * alpha + partial * beta + l_i = l_i * alpha + beta + m_i = m_new + + o = acc / l_i + o_ptrs = o_ptr + pid_b * stride_ob + pid_m * stride_om + pid_h * stride_oh + offs_d * stride_od + tl.store(o_ptrs, o.to(o_ptr.dtype.element_ty)) + + +def split_count(b: int, m: int, h: int, topk: int, device) -> int: + """How many ways to split the candidate axis; 0 means run the single-program kernel. Decode + splits until every SM has a program or every split is down to one tile.""" + if m != 1: + return 0 + sm_count = torch.cuda.get_device_properties(device).multi_processor_count + n_splits = min( + MAX_SPLITS, + triton.cdiv(topk, MIN_TILES_PER_SPLIT * BLOCK_T), + max(1, sm_count // (b * triton.cdiv(h, BLOCK_H))), + ) + return n_splits if n_splits > 1 else 0 + + +def sparse_attn_packed( + q: torch.Tensor, # [b, m, h, d] bf16 + window_pool: torch.Tensor, # [n_win_slots, win_fmt.row_bytes(d)] uint8 + win_fmt: RowFormat, + cmp_pool: torch.Tensor, # [n_cmp, cmp_fmt.row_bytes(d)] uint8 + cmp_fmt: RowFormat, + attn_sink: torch.Tensor, # [h] + topk_idxs: torch.Tensor, # [b, m, topk] int32, GLOBAL rows, layout [window | compressed] + n_window: int, # window columns (a multiple of BLOCK_T; pad with -1) + softmax_scale: float, + cmp_counts: torch.Tensor | None = None, # [b, m] int32 live compressed columns per query +) -> torch.Tensor: + """Paged sparse MLA attention over packed KV pools; see the module docstring for the contract.""" + b, m, h, d = q.shape + topk = topk_idxs.shape[-1] + assert n_window % BLOCK_T == 0 and 0 <= n_window <= topk, (n_window, topk) + assert window_pool.dtype == torch.uint8 and window_pool.shape[1] == win_fmt.row_bytes(d), (window_pool.shape, win_fmt) + assert cmp_pool.dtype == torch.uint8 and cmp_pool.shape[1] == cmp_fmt.row_bytes(d), (cmp_pool.shape, cmp_fmt) + assert window_pool.is_contiguous() and cmp_pool.is_contiguous() + q = q.contiguous() + idx = topk_idxs.contiguous().to(torch.int32) + sink = attn_sink.contiguous().to(torch.float32) + o = torch.empty_like(q) + + has_counts = cmp_counts is not None + if has_counts: + cnt = cmp_counts.contiguous().to(torch.int32).view(b, m) + stride_nb, stride_nm = cnt.stride() + else: + cnt, stride_nb, stride_nm = idx, 0, 0 + + fmt_args = dict( + WIN_FMT=win_fmt.code, WIN_ROW_BYTES=win_fmt.row_bytes(d), + CMP_FMT=cmp_fmt.code, CMP_ROW_BYTES=cmp_fmt.row_bytes(d), + ) + n_splits = split_count(b, m, h, topk, q.device) + if n_splits: + head_blocks = triton.cdiv(h, BLOCK_H) + mid_o = torch.empty((b, m, h, n_splits, d), dtype=torch.float32, device=q.device) + mid_lse = torch.empty((b, m, h, n_splits), dtype=torch.float32, device=q.device) + _sparse_attn_packed_splitk_kernel[(m * n_splits, b, head_blocks)]( + q, window_pool, cmp_pool, mid_o, mid_lse, idx, cnt, + float(softmax_scale), + h, topk, int(n_window), + q.stride(0), q.stride(1), q.stride(2), q.stride(3), + mid_o.stride(0), mid_o.stride(1), mid_o.stride(2), mid_o.stride(3), mid_o.stride(4), + mid_lse.stride(0), mid_lse.stride(1), mid_lse.stride(2), mid_lse.stride(3), + idx.stride(0), idx.stride(1), idx.stride(2), + stride_nb, stride_nm, + D=d, BLOCK_H=BLOCK_H, BLOCK_T=BLOCK_T, HAS_COUNTS=has_counts, NUM_SPLITS=n_splits, + **fmt_args, num_warps=8, num_stages=2, + ) + _sparse_attn_splitk_merge_kernel[(m, b, h)]( + mid_o, mid_lse, o, sink, + mid_o.stride(0), mid_o.stride(1), mid_o.stride(2), mid_o.stride(3), mid_o.stride(4), + mid_lse.stride(0), mid_lse.stride(1), mid_lse.stride(2), mid_lse.stride(3), + o.stride(0), o.stride(1), o.stride(2), o.stride(3), + D=d, NUM_SPLITS=n_splits, num_warps=4, + ) + return o + + grid = (m, b, triton.cdiv(h, BLOCK_H)) + _sparse_attn_packed_kernel[grid]( + q, window_pool, cmp_pool, o, sink, idx, cnt, + float(softmax_scale), + h, topk, int(n_window), + q.stride(0), q.stride(1), q.stride(2), q.stride(3), + o.stride(0), o.stride(1), o.stride(2), o.stride(3), + idx.stride(0), idx.stride(1), idx.stride(2), + stride_nb, stride_nm, + D=d, BLOCK_H=BLOCK_H, BLOCK_T=BLOCK_T, HAS_COUNTS=has_counts, + **fmt_args, num_warps=8, num_stages=2, + ) + return o + + +__all__ = ["sparse_attn_packed", "split_count", "BLOCK_T"] diff --git a/python/freetoken/kernel/triton/dsv41/topk.py b/python/freetoken/kernel/triton/dsv41/topk.py new file mode 100644 index 0000000000..4d638e4d49 --- /dev/null +++ b/python/freetoken/kernel/triton/dsv41/topk.py @@ -0,0 +1,265 @@ +"""Exact top-k and block-candidate selection whose work follows the live history (DeepSeek-V4.1). + +The indexer scores a request's compressed positions into a buffer as wide as the staged history +(``max_seq_len / ratio``, fixed while a CUDA graph is captured), but only the first ``live[row]`` +columns are meaningful -- and at 1M staged positions a 4K-token request must not pay for the +other 99.6%. Every kernel here takes the live count per row on the device, keeps its output +addresses and shapes fixed, and restricts what it READS and reduces to the live prefix; nothing +past ``live`` is ever read, so the score buffer's dead area needs no initialization. + +Selection is a fixed pipeline of small kernels with explicit workspace contracts: + +* ``dsv41_topk``: exact top-k of ``[R, T]`` fp32 scores below ``live``. Level 0 gives every + ``CHUNK``-wide slice of a row its own program, which radix-selects the slice's top-k into a + ``(key, column)`` candidate workspace (a global winner beats at most k-1 columns, so it is + inside its own slice's top-k: the union is exact). Further levels apply the same program to the + candidates until one register-resident tile remains; the final program ranks it and emits the + winners in ascending column order, ``-1`` padded. Slices past ``live`` exit at once and a + slice with no more entries than ``k`` is copied through, so a short history costs a handful of + near-empty launches. The radix primitives are the QSA block top-k's (``kernel/triton/qsa/topk``). +* ``dsv41_block_scores``: per-``block_size`` block maxima over the live prefix (the Hierarchical + Sparse Indexer's level one), the block holding the newest live position pinned to ``+inf`` so it + is always kept; a fixed worker grid strides over the live blocks. +* ``dsv41_candidate_blocks``: the two above plus a fixed-width expansion into the compact candidate + list ``[R, topk_blocks * block_size]`` of compressed positions, ascending with ``-1`` in the tail. + +Tie rule (documented, deterministic): scores order descending, equal scores by ascending column, +so the lowest positions win among exact ties. ``torch.topk`` leaves ties unspecified. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from freetoken.kernel.triton.qsa.topk import _BINS, _PASSES, _RADIX, _monotone_key, _resident_prefix, _tile_ranks + +CHUNK = 4096 # register-resident slice per level program (fp32 keys + int32 columns) +FINAL_MAX = 4096 # the final program's resident candidate tile +BLOCK_SCORE_TILE = 512 # blocks per block-max program tile + + +@triton.jit +def _load_entries(src_ptr, col_ptr, base, offsets, limit, SRC_LOGITS: tl.constexpr): + """One slice of a level's input as ``(monotone key, column)``; key 0 = dead (past ``limit``, + ``-inf``, or a padding slot of the previous level).""" + live = offsets < limit + if SRC_LOGITS: + value = tl.load(src_ptr + base + offsets, mask=live, other=-float("inf")) + key = tl.where(live & (value > -float("inf")), _monotone_key(value), 0) + cols = (base + offsets).to(tl.int32) + else: + key = tl.load(src_ptr + base + offsets, mask=live, other=0).to(tl.uint32, bitcast=True) + cols = tl.load(col_ptr + base + offsets, mask=live, other=-1) + return key, cols + + +@triton.jit +def _dsv41_topk_level_kernel( + src_ptr, src_col_ptr, live_in_ptr, key_out_ptr, col_out_ptr, live_out_ptr, + stride_src_row, stride_out_row, num_columns, + TOP_K: tl.constexpr, PAD_K: tl.constexpr, CHUNK: tl.constexpr, N_SPLITS: tl.constexpr, + SRC_LOGITS: tl.constexpr, BINS: tl.constexpr, RADIX: tl.constexpr, PASSES: tl.constexpr, +): + """One program per (row, slice): the slice's top-k into candidate slot ``split * TOP_K``. Slot + ``split`` is read downstream only when ``split < cdiv(live, CHUNK)``; program 0 publishes that + live candidate count for the next level.""" + tl.static_assert(CHUNK <= 0xFFFF, "packed cumsum keeps 16 bits per half") + tl.static_assert(TOP_K <= CHUNK, "a slice must hold its own top-k") + row = tl.program_id(0) + split = tl.program_id(1) + visible = tl.maximum(tl.minimum(tl.load(live_in_ptr + row), num_columns), 0) + if split == 0: + tl.store(live_out_ptr + row, tl.minimum(tl.cdiv(visible, CHUNK), N_SPLITS) * TOP_K) + base = split * CHUNK + limit = tl.minimum(tl.maximum(visible - base, 0), CHUNK) + if limit <= 0: + return # a dead slice: nothing downstream reads its slot + slot = row.to(tl.int64) * stride_out_row + split * TOP_K + offsets = tl.arange(0, CHUNK) + key, cols = _load_entries(src_ptr + row.to(tl.int64) * stride_src_row, src_col_ptr + row.to(tl.int64) * stride_src_row, base, offsets, limit, SRC_LOGITS) + if limit <= TOP_K: + # every live entry is a winner: copy in order (dead keys stay 0 for the next level) + inside = offsets < TOP_K + tl.store(key_out_ptr + slot + offsets, key.to(tl.int32, bitcast=True), mask=inside) + tl.store(col_out_ptr + slot + offsets, cols, mask=inside) + return + prefix, ties = _resident_prefix(key, TOP_K, BINS, RADIX, PASSES) + rank, take, above, equal = _tile_ranks(key, prefix, ties, 0, 0) + write = take & (rank < TOP_K) + tl.store(key_out_ptr + slot + rank, key.to(tl.int32, bitcast=True), mask=write) + tl.store(col_out_ptr + slot + rank, cols, mask=write) + emitted = above + tl.minimum(equal, ties) + pad = tl.arange(0, PAD_K) + tl.store(key_out_ptr + slot + pad, 0, mask=(pad >= emitted) & (pad < TOP_K)) + + +@triton.jit +def _final_tile(src_row, col_row, out_row, limit, TOP_K: tl.constexpr, BLOCK: tl.constexpr, SRC_LOGITS: tl.constexpr, BINS: tl.constexpr, RADIX: tl.constexpr, PASSES: tl.constexpr): + tl.static_assert(BLOCK <= 0xFFFF, "packed cumsum keeps 16 bits per half") + offsets = tl.arange(0, BLOCK) + key, cols = _load_entries(src_row, col_row, 0, offsets, limit, SRC_LOGITS) + prefix, ties = _resident_prefix(key, tl.minimum(limit, TOP_K), BINS, RADIX, PASSES) + rank, take, above, equal = _tile_ranks(key, prefix, ties, 0, 0) + tl.store(out_row + rank, cols, mask=take & (rank < TOP_K)) + return above + tl.minimum(equal, ties) + + +@triton.jit +def _dsv41_topk_final_kernel( + src_ptr, src_col_ptr, live_ptr, out_ptr, + stride_src_row, stride_out_row, num_columns, + TOP_K: tl.constexpr, PAD_K: tl.constexpr, + BLOCK_SMALL: tl.constexpr, BLOCK_MID: tl.constexpr, BLOCK_FULL: tl.constexpr, + SRC_LOGITS: tl.constexpr, BINS: tl.constexpr, RADIX: tl.constexpr, PASSES: tl.constexpr, +): + """One program per row over the last level's live candidates (or the scores themselves when the + row fits one tile): winners at their ascending-column ranks, ``-1`` from ``emitted`` on.""" + row = tl.program_id(0) + limit = tl.maximum(tl.minimum(tl.load(live_ptr + row), num_columns), 0) + src_row = src_ptr + row.to(tl.int64) * stride_src_row + col_row = src_col_ptr + row.to(tl.int64) * stride_src_row + out_row = out_ptr + row.to(tl.int64) * stride_out_row + emitted = 0 + if limit > 0: + # register residency costs the whole tile, so a short row takes a narrower one + if limit <= BLOCK_SMALL: + emitted = _final_tile(src_row, col_row, out_row, limit, TOP_K, BLOCK_SMALL, SRC_LOGITS, BINS, RADIX, PASSES) + elif limit <= BLOCK_MID: + emitted = _final_tile(src_row, col_row, out_row, limit, TOP_K, BLOCK_MID, SRC_LOGITS, BINS, RADIX, PASSES) + else: + emitted = _final_tile(src_row, col_row, out_row, limit, TOP_K, BLOCK_FULL, SRC_LOGITS, BINS, RADIX, PASSES) + pad = tl.arange(0, PAD_K) + tl.store(out_row + pad, -1, mask=(pad >= emitted) & (pad < TOP_K)) + + +MAX_K = CHUNK // 2 # a slice must shed at least half its entries per level, or the plan cannot converge + + +def topk_plan(columns: int, k: int) -> list[tuple[int, int]]: + """``[(input width, n_splits), ...]`` of the reduction levels before the final tile; empty when + the row already fits the final tile. Derived from the buffer width alone (graph-static). Each + level maps ``width`` to ``cdiv(width, CHUNK) * k`` candidates, which shrinks the row by at least + ``CHUNK / k >= 2`` per level for ``k <= MAX_K``; larger ``k`` would stall (``k == CHUNK`` never + shrinks) and is rejected here, before any launch.""" + if not 0 < k <= MAX_K: + raise ValueError(f"dsv41_topk supports 1 <= k <= {MAX_K}, got {k}") + levels: list[tuple[int, int]] = [] + width = columns + while width > FINAL_MAX: + n_splits = -(-width // CHUNK) + levels.append((width, n_splits)) + nxt = n_splits * k + assert nxt < width, (width, k) # guaranteed by k <= CHUNK // 2 for width > FINAL_MAX == CHUNK + width = nxt + return levels + + +def dsv41_topk(scores: torch.Tensor, live: torch.Tensor, k: int, out: torch.Tensor | None = None) -> torch.Tensor: + """Exact top-``k`` columns of every ``scores`` row among its first ``live[row]`` columns. + + ``scores [R, T]`` fp32 row-contiguous (columns at or past ``live`` are never read; ``-inf`` + columns are dead), ``live [R]`` int32. Returns ``out [R, k]`` int32: the winners in ascending + column order packed at the front, ``-1`` from the ``emitted`` count on (fewer live or finite + columns than ``k`` leave a ``-1`` tail). Ties: lower column wins.""" + assert scores.ndim == 2 and scores.dtype == torch.float32 and scores.stride(1) == 1, (scores.shape, scores.dtype) + R, T = scores.shape + assert live.shape == (R,), (live.shape, R) + live = live.to(torch.int32).contiguous() + if out is None: + out = torch.empty((R, k), dtype=torch.int32, device=scores.device) + assert out.shape == (R, k) and out.dtype == torch.int32 and out.stride(1) == 1 + if R == 0 or k == 0: + return out + if T == 0: + out.fill_(-1) + return out + pad_k = triton.next_power_of_2(k) + src, src_col, src_live, width, src_logits = scores, scores, live, T, True + for level_width, n_splits in topk_plan(T, k): + keys = torch.empty((R, n_splits * k), dtype=torch.int32, device=scores.device) + cols = torch.empty((R, n_splits * k), dtype=torch.int32, device=scores.device) + live_out = torch.empty((R,), dtype=torch.int32, device=scores.device) + _dsv41_topk_level_kernel[(R, n_splits)]( + src, src_col, src_live, keys, cols, live_out, + src.stride(0), keys.stride(0), level_width, + TOP_K=k, PAD_K=pad_k, CHUNK=CHUNK, N_SPLITS=n_splits, SRC_LOGITS=src_logits, + BINS=_BINS, RADIX=_RADIX, PASSES=_PASSES, num_warps=8, num_stages=1, + ) + src, src_col, src_live, width, src_logits = keys, cols, live_out, n_splits * k, False + block = triton.next_power_of_2(max(width, 1)) + _dsv41_topk_final_kernel[(R,)]( + src, src_col, src_live, out, + src.stride(0), out.stride(0), width, + TOP_K=k, PAD_K=pad_k, + BLOCK_SMALL=min(1024, block), BLOCK_MID=min(4096, block), BLOCK_FULL=block, + SRC_LOGITS=src_logits, BINS=_BINS, RADIX=_RADIX, PASSES=_PASSES, num_warps=8, num_stages=1, + ) + return out + + +@triton.jit +def _dsv41_block_scores_kernel( + scores_ptr, live_ptr, blk_ptr, + stride_s_row, stride_b_row, num_blocks, + BLOCK_SIZE: tl.constexpr, TILE: tl.constexpr, N_WORKERS: tl.constexpr, +): + """Worker ``w`` of a row scores block tiles ``w, w + N_WORKERS, ...`` below the live block count: + the block's best live position, ``+inf`` for the block holding the newest live position.""" + row = tl.program_id(0) + worker = tl.program_id(1) + live = tl.maximum(tl.load(live_ptr + row), 0) + nb_live = tl.minimum(tl.cdiv(live, BLOCK_SIZE), num_blocks) + s_row = scores_ptr + row.to(tl.int64) * stride_s_row + b_row = blk_ptr + row.to(tl.int64) * stride_b_row + offs_b = tl.arange(0, TILE) + offs_p = tl.arange(0, BLOCK_SIZE) + for tile in range(worker, tl.cdiv(nb_live, TILE), N_WORKERS): + blocks = tile * TILE + offs_b + b_mask = blocks < nb_live + pos = blocks[:, None] * BLOCK_SIZE + offs_p[None, :] + v = tl.load(s_row + pos, mask=b_mask[:, None] & (pos < live), other=-float("inf")) + best = tl.max(v, axis=1) + best = tl.where(blocks == nb_live - 1, float("inf"), best) + tl.store(b_row + blocks, best, mask=b_mask) + + +def dsv41_block_scores(scores: torch.Tensor, live: torch.Tensor, block_size: int, out: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor]: + """``(block scores [R, cdiv(T, block_size)] fp32, live blocks [R] int32)``: each live block's best + live position (``+inf`` for the newest block); blocks at or past the live count are not written.""" + R, T = scores.shape + nb = -(-T // block_size) + live = live.to(torch.int32).contiguous() + if out is None: + out = torch.empty((R, nb), dtype=torch.float32, device=scores.device) + live_blocks = torch.div(live + block_size - 1, block_size, rounding_mode="floor").clamp_(min=0).to(torch.int32) + if R == 0 or nb == 0: + return out, live_blocks + assert block_size == triton.next_power_of_2(block_size), block_size + sm_count = torch.cuda.get_device_properties(scores.device).multi_processor_count + workers = min(triton.cdiv(nb, BLOCK_SCORE_TILE), max(1, triton.cdiv(sm_count, R))) + _dsv41_block_scores_kernel[(R, workers)]( + scores, live, out, scores.stride(0), out.stride(0), nb, + BLOCK_SIZE=block_size, TILE=BLOCK_SCORE_TILE, N_WORKERS=workers, num_warps=4, + ) + return out, live_blocks + + +def dsv41_candidate_blocks(scores: torch.Tensor, live: torch.Tensor, topk_blocks: int, block_size: int) -> torch.Tensor: + """Level one of the Hierarchical Sparse Indexer over ``[R, T]`` scores: the ``topk_blocks`` best + live blocks (the newest always among them) expanded to the compact candidate list ``[R, + topk_blocks * block_size]`` of compressed positions -- ascending, ``-1`` for positions at or past + ``live`` and for unfilled blocks, all in the tail (a sorted valid prefix).""" + R, T = scores.shape + blk, live_blocks = dsv41_block_scores(scores, live, block_size) + keep = dsv41_topk(blk, live_blocks, min(topk_blocks, blk.shape[1])) # [R, KB] ascending, -1 tail + pos = keep.to(torch.int64).unsqueeze(-1) * block_size + torch.arange(block_size, device=scores.device) + dead = (keep < 0).unsqueeze(-1) | (pos >= live.to(torch.int64).view(R, 1, 1)) + out = torch.where(dead, -1, pos).flatten(1).to(torch.int32) + if out.shape[1] < topk_blocks * block_size: + out = torch.nn.functional.pad(out, (0, topk_blocks * block_size - out.shape[1]), value=-1) + return out + + +__all__ = ["dsv41_topk", "dsv41_block_scores", "dsv41_candidate_blocks", "topk_plan", "CHUNK", "FINAL_MAX", "MAX_K"] diff --git a/python/freetoken/kernel/triton/e4m3_compat.py b/python/freetoken/kernel/triton/e4m3_compat.py index 6ea0c7537d..14a7acd835 100644 --- a/python/freetoken/kernel/triton/e4m3_compat.py +++ b/python/freetoken/kernel/triton/e4m3_compat.py @@ -109,6 +109,20 @@ def e4m3_u8_to_f16_x128(v): return h.to(tl.float16, bitcast=True) * 32768.0 +@jit +def e4m3_f32_to_u8(x): + """Encode an fp32 value that already sits on the e4m3 grid (see :func:`round_e4m3`, |x| <= 448) + into its e4m3 bit pattern (uint8). Normals: fp32 exponent rebased 127 -> 7 with the top three + mantissa bits; subnormals (|x| < 2^-6, grid 2^-9): mantissa = |x| / 2^-9, exponent field 0.""" + b = x.to(tl.uint32, bitcast=True) + sign = ((b >> 31) & 1).to(tl.uint8) << 7 + exp = ((b >> 23) & 0xFF).to(tl.int32) - 127 + 7 + man = ((b >> 20) & 0x7).to(tl.uint8) + normal = (exp.to(tl.uint8) << 3) | man + sub = (tl.abs(x) * 512.0).to(tl.uint8) # 0..7 on the 2^-9 grid + return sign | tl.where(tl.abs(x) >= 0.015625, normal, sub) + + @jit def round_e4m3(x): """Round fp32 onto the e4m3 value grid (RNE), fp32 -> fp32, in a SINGLE diff --git a/python/freetoken/kernel/triton/mhc.py b/python/freetoken/kernel/triton/mhc.py index 0782efe37d..876643c164 100644 --- a/python/freetoken/kernel/triton/mhc.py +++ b/python/freetoken/kernel/triton/mhc.py @@ -27,7 +27,7 @@ def _mhc_stage1_kernel( H: tl.constexpr, N: tl.constexpr, MIX: tl.constexpr, BLK_MIX: tl.constexpr, SPLIT: tl.constexpr, # hidden elems per split (multiple of BLOCK_H) BLOCK_H: tl.constexpr, - NS: tl.constexpr, + NS: tl.constexpr, # splits per token (the grid's second axis; need not be a power of two) HAS_POST: tl.constexpr, ): """Split-K stage of the fused mHC: each program owns one hidden slice of one @@ -79,24 +79,24 @@ def _mhc_stage1_kernel( @triton.jit -def _mhc_stage2_kernel( - sq_part_ptr, mix_part_ptr, scale_ptr, base_ptr, +def _mhc_gates( + t, sq_part_ptr, mix_part_ptr, scale_ptr, base_ptr, post_out_ptr, comb_out_ptr, pre_out_ptr, rms_eps, hc_eps, post_mult, SINKHORN: tl.constexpr, H: tl.constexpr, N: tl.constexpr, MIX: tl.constexpr, BLK_MIX: tl.constexpr, - NS: tl.constexpr, + NS: tl.constexpr, BLK_NS: tl.constexpr, ): - """Reduce the split partials and run the tiny gate math (sigmoid gates, - row-softmax + in-register 4x4 Sinkhorn); emits pre gates for stage 3.""" - t = tl.program_id(0).to(tl.int64) + """Reduce token ``t``'s split partials and run the tiny gate math (sigmoid gates, + row-softmax + in-register 4x4 Sinkhorn); writes post / comb / pre for the next sublayer.""" offs_mix = tl.arange(0, BLK_MIX) mix_mask = offs_mix < MIX - offs_s = tl.arange(0, NS) + offs_s = tl.arange(0, BLK_NS) + s_mask = offs_s < NS - sqsum = tl.sum(tl.load(sq_part_ptr + t * NS + offs_s)) + sqsum = tl.sum(tl.load(sq_part_ptr + t * NS + offs_s, mask=s_mask, other=0.0)) acc = tl.sum( - tl.load(mix_part_ptr + (t * NS + offs_s)[:, None] * BLK_MIX + offs_mix[None, :]), + tl.load(mix_part_ptr + (t * NS + offs_s)[:, None] * BLK_MIX + offs_mix[None, :], mask=s_mask[:, None], other=0.0), axis=0, ) inv_rms = tl.math.rsqrt(sqsum / (N * H) + rms_eps) @@ -142,6 +142,23 @@ def _mhc_stage2_kernel( tl.store(comb_out_ptr + t * N * N + offs_n2[:, None] * N + offs_n2[None, :], comb) + +@triton.jit +def _mhc_stage2_kernel( + sq_part_ptr, mix_part_ptr, scale_ptr, base_ptr, + post_out_ptr, comb_out_ptr, pre_out_ptr, + rms_eps, hc_eps, post_mult, + SINKHORN: tl.constexpr, + H: tl.constexpr, N: tl.constexpr, MIX: tl.constexpr, BLK_MIX: tl.constexpr, + NS: tl.constexpr, BLK_NS: tl.constexpr, +): + """Gate math alone (the two-pass variant, whose stage 3 consumes these pre gates).""" + _mhc_gates( + tl.program_id(0).to(tl.int64), sq_part_ptr, mix_part_ptr, scale_ptr, base_ptr, post_out_ptr, comb_out_ptr, pre_out_ptr, + rms_eps, hc_eps, post_mult, SINKHORN, H, N, MIX, BLK_MIX, NS, BLK_NS, + ) + + @triton.jit def _mhc_stage3_kernel( res_out_ptr, pre_ptr, li_out_ptr, @@ -161,7 +178,46 @@ def _mhc_stage3_kernel( tl.store(li_out_ptr + t * H + offs_h, li.to(li_out_ptr.dtype.element_ty), mask=h_mask) -def mhc_fused_post_pre_triton( +@triton.jit +def _mhc_stage23_kernel( + res_out_ptr, pre_ptr, li_out_ptr, + sq_part_ptr, mix_part_ptr, scale_ptr, base_ptr, + post_out_ptr, comb_out_ptr, pre_out_ptr, + rms_eps, hc_eps, post_mult, + SINKHORN: tl.constexpr, + H: tl.constexpr, N: tl.constexpr, MIX: tl.constexpr, BLK_MIX: tl.constexpr, + NS: tl.constexpr, BLK_NS: tl.constexpr, BLOCK_H: tl.constexpr, +): + """Single-pass mHC: the layer input mixes with the PREVIOUS sublayer's pre gates, so the + gate math for the next sublayer (stage 2) does not gate stage 3 -- one launch does both: + every program mixes one hidden chunk, program 0 of each token also reduces the split + partials and writes the next sublayer's gates.""" + t = tl.program_id(0).to(tl.int64) + hb = tl.program_id(1) + offs_h = hb * BLOCK_H + tl.arange(0, BLOCK_H) + h_mask = offs_h < H + offs_n = tl.arange(0, N) + pre = tl.load(pre_ptr + t * N + offs_n) + li = tl.zeros([BLOCK_H], dtype=tl.float32) + for n in tl.static_range(N): + r = tl.load(res_out_ptr + (t * N + n) * H + offs_h, mask=h_mask, other=0.0).to(tl.float32) + li += tl.sum(tl.where(offs_n == n, pre, 0.0)) * r + tl.store(li_out_ptr + t * H + offs_h, li.to(li_out_ptr.dtype.element_ty), mask=h_mask) + if hb == 0: + _mhc_gates( + t, sq_part_ptr, mix_part_ptr, scale_ptr, base_ptr, post_out_ptr, comb_out_ptr, pre_out_ptr, + rms_eps, hc_eps, post_mult, SINKHORN, H, N, MIX, BLK_MIX, NS, BLK_NS, + ) + + +def _stage1_hidden_tile(h: int) -> int: + """The stage-1 hidden split. It fixes the fp32 reduction order over the hidden width, so it + depends on the width alone and results reproduce across processes, machines and chunk sizes. + 128 gives 32 / 40 programs per token at widths 4096 / 5120.""" + return min(128, triton.next_power_of_2(h)) + + +def _mhc_fused_launch( x: torch.Tensor, residual: torch.Tensor, post_mix: torch.Tensor | None, @@ -173,11 +229,12 @@ def mhc_fused_post_pre_triton( hc_eps: float, post_mult: float, sinkhorn_repeat: int, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Fused hc_post (skipped when ``post_mix is None``) + hc_pre. Three-stage - split-K: the GEMV/sq-sum reduction fans out over NS hidden slices. Returns - (residual_new [T,N,H] bf16, post [T,N,1] fp32, comb [T,N,N] fp32, - layer_input [T,H] bf16).""" + pre_mix_in: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """The three stages behind both variants. ``pre_mix_in`` selects single-pass: the + layer input is mixed with it instead of with the pre gates stage 2 just computed + (those are still returned, for the next sublayer). Returns (residual_new, post, + comb, pre_new [T,N] fp32, layer_input).""" t, n, h = residual.shape mix = 2 * n + n * n assert fn.shape == (mix, n * h) and fn.dtype == torch.float32 @@ -190,12 +247,7 @@ def mhc_fused_post_pre_triton( comb_out = torch.empty(t, n, n, dtype=torch.float32, device=dev) li_out = torch.empty(t, h, dtype=residual.dtype, device=dev) - block_h = min(512, triton.next_power_of_2(h)) - # NS feeds a tl.arange in stage 2 -> keep it a power of two. - ns = 1 - while ns * 2 <= min(16, h // block_h): - ns *= 2 - split = triton.cdiv(triton.cdiv(h, ns), block_h) * block_h + block_h = split = _stage1_hidden_tile(h) ns = triton.cdiv(h, split) blk_mix = triton.next_power_of_2(mix) sq_part = torch.empty(t, ns, dtype=torch.float32, device=dev) @@ -213,20 +265,83 @@ def mhc_fused_post_pre_triton( HAS_POST=has_post, num_warps=4, num_stages=2, ) - _mhc_stage2_kernel[(t,)]( - sq_part, mix_part, hc_scale, hc_base, - post_out, comb_out, pre_out, - rms_eps, hc_eps, post_mult, - SINKHORN=sinkhorn_repeat, - H=h, N=n, MIX=mix, BLK_MIX=blk_mix, NS=ns, - num_warps=1, + block_h3 = min(1024, triton.next_power_of_2(h)) + if pre_mix_in is None: + # two-pass: stage 3 mixes with the gates stage 2 just produced + _mhc_stage2_kernel[(t,)]( + sq_part, mix_part, hc_scale, hc_base, + post_out, comb_out, pre_out, + rms_eps, hc_eps, post_mult, + SINKHORN=sinkhorn_repeat, + H=h, N=n, MIX=mix, BLK_MIX=blk_mix, NS=ns, BLK_NS=triton.next_power_of_2(ns), + num_warps=1, + ) + _mhc_stage3_kernel[(t, triton.cdiv(h, block_h3))]( + res_out, pre_out, li_out, + H=h, N=n, BLOCK_H=block_h3, + num_warps=4, + ) + else: + # single-pass: the mix uses the previous sublayer's gates, so stages 2 and 3 share a launch + assert pre_mix_in.shape == (t, n), (pre_mix_in.shape, (t, n)) + _mhc_stage23_kernel[(t, triton.cdiv(h, block_h3))]( + res_out, pre_mix_in.to(torch.float32).contiguous(), li_out, + sq_part, mix_part, hc_scale, hc_base, + post_out, comb_out, pre_out, + rms_eps, hc_eps, post_mult, + SINKHORN=sinkhorn_repeat, + H=h, N=n, MIX=mix, BLK_MIX=blk_mix, NS=ns, BLK_NS=triton.next_power_of_2(ns), BLOCK_H=block_h3, + num_warps=4, + ) + return res_out, post_out.view(t, n, 1), comb_out, pre_out, li_out + + +def mhc_fused_post_pre_triton( + x: torch.Tensor, + residual: torch.Tensor, + post_mix: torch.Tensor | None, + comb_mix: torch.Tensor | None, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_eps: float, + post_mult: float, + sinkhorn_repeat: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Fused hc_post (skipped when ``post_mix is None``) + hc_pre. Three-stage + split-K: the GEMV/sq-sum reduction fans out over NS hidden slices. Returns + (residual_new [T,N,H] bf16, post [T,N,1] fp32, comb [T,N,N] fp32, + layer_input [T,H] bf16).""" + res_out, post_out, comb_out, _pre_out, li_out = _mhc_fused_launch( + x, residual, post_mix, comb_mix, fn, hc_scale, hc_base, + rms_eps, hc_eps, post_mult, sinkhorn_repeat, pre_mix_in=None, ) - _mhc_stage3_kernel[(t, triton.cdiv(h, 1024))]( - res_out, pre_out, li_out, - H=h, N=n, BLOCK_H=min(1024, triton.next_power_of_2(h)), - num_warps=4, + return res_out, post_out, comb_out, li_out + + +def mhc_fused_post_pre_single_pass_triton( + x: torch.Tensor, + residual: torch.Tensor, + post_mix: torch.Tensor | None, + comb_mix: torch.Tensor | None, + pre_mix_in: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_eps: float, + post_mult: float, + sinkhorn_repeat: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Single-pass mHC (DeepSeek-V4.1): same fused stages, but the layer input is mixed + with the PREVIOUS sublayer's ``pre_mix_in`` and this sublayer's pre gates come back + as ``pre_new`` for the next one. Returns (residual_new, post [T,N,1], comb [T,N,N], + pre_new [T,N] fp32, layer_input).""" + return _mhc_fused_launch( + x, residual, post_mix, comb_mix, fn, hc_scale, hc_base, + rms_eps, hc_eps, post_mult, sinkhorn_repeat, pre_mix_in=pre_mix_in, ) - return res_out, post_out.view(t, n, 1), comb_out, li_out -__all__ = ["mhc_fused_post_pre_triton"] +__all__ = ["mhc_fused_post_pre_triton", "mhc_fused_post_pre_single_pass_triton"] diff --git a/python/freetoken/kvcache/__init__.py b/python/freetoken/kvcache/__init__.py index 1792da2b3c..e0cf606db1 100644 --- a/python/freetoken/kvcache/__init__.py +++ b/python/freetoken/kvcache/__init__.py @@ -34,7 +34,7 @@ def resolve_pool_class(model_config: ModelConfig) -> type[BaseKVCachePool]: specs_fn = getattr(model_config, "kv_cache_group_specs", None) if specs_fn is None: if getattr(model_config, "dsv4_args", None) is not None: - from .dsv4_paged_pool import DSV4PagedKVCache + from .dsv4.v4_pool import DSV4PagedKVCache return DSV4PagedKVCache from .mha_pool import MHAKVCache @@ -43,9 +43,13 @@ def resolve_pool_class(model_config: ModelConfig) -> type[BaseKVCachePool]: specs = list(specs_fn()) types = {spec.attn_type for spec in specs} if AttnType.DSV4 in types: - from .dsv4_paged_pool import DSV4PagedKVCache + from .dsv4.v4_pool import DSV4PagedKVCache return DSV4PagedKVCache + if AttnType.DSV41 in types: + from .dsv4.v41_pool import DSV41PagedKVCache + + return DSV41PagedKVCache if AttnType.SWA in types: from .hybrid_swa_pool import HybridSWAKVCache @@ -80,9 +84,9 @@ def create_kv_pool(config, num_pages: int, device: torch.device, dtype: torch.dt """Build the engine's KV pool for ``num_pages`` USABLE pages (the dummy page and every secondary tier -- window pool, index slab, state rings -- are derived here or inside the pool). Single factory entry for all pool families, DSV4 included.""" - from .dsv4_cost_model import _dsv4_pool_sizes + from .dsv4.v4_cost_model import _dsv4_pool_sizes from .hybrid_swa_pool import _naive_swa_num_tokens, _swa_paged_num_tokens - from .dsv4_paged_pool import DSV4PagedKVCache + from .dsv4.v4_pool import DSV4PagedKVCache model_config = config.model_config if resolve_pool_class(model_config) is DSV4PagedKVCache: @@ -99,6 +103,21 @@ def create_kv_pool(config, num_pages: int, device: torch.device, dtype: torch.dt ) pool._init_paged_state(config.max_running_req, config.cache_type != "naive") return pool + if getattr(model_config, "dsv41_args", None) is not None: + from .dsv4.v41_cost_model import _dsv41_pool_sizes + from .dsv4.v41_pool import DSV41PagedKVCache + + # Same route as DSV4: the generic CacheManager over the shared page table; the pool is the + # swa_pool plug-in (window tier + per-source packed main / index pools + state rings). + pool = DSV41PagedKVCache( + sizes=_dsv41_pool_sizes(config, num_pages + 1), # +1 for dummy page + args=model_config.dsv41_args, + device=device, + dtype=dtype, + n_scratch=config.max_running_req + 1, + ) + pool._init_paged_state(config.max_running_req, config.cache_type != "naive") + return pool num_swa_tokens = None # Both the naive and radix SWA paths share the global-paged swa pool; radix sizes it by diff --git a/python/freetoken/kvcache/base.py b/python/freetoken/kvcache/base.py index 95669e8c81..569da1cb6c 100644 --- a/python/freetoken/kvcache/base.py +++ b/python/freetoken/kvcache/base.py @@ -35,6 +35,14 @@ def spec_kv_bytes_per_token(spec, config) -> int: return per_token + spec.index_head_dim * spec.num_index_layers * 2 // spec.index_ratio +@dataclass(frozen=True) +class WindowPoolSpec: + """Window resize units and minimum usable capacity, excluding reserved pages.""" + + page_size: int + min_pages: int + + class BaseKVCachePool(ABC): """ Base class for key-value caches. @@ -45,6 +53,15 @@ class BaseKVCachePool(ABC): # model re-bound after a rebuild; the engine asks before it resizes. needs_rebind_on_rebuild: ClassVar[bool] = False + @classmethod + def window_spec(cls, config) -> WindowPoolSpec | None: + return None + + @property + def window_pages(self) -> int: + """Current usable window pages in the units declared by window_spec.""" + return 0 + # ---- sizing/cost classmethods: run BEFORE the pool exists (startup budget solve, # --moe-cache-auto). The engine measures memory and passes bytes in; each pool family # implements kv_cost for ITS OWN buffers only (the engine sums families, e.g. adds diff --git a/python/freetoken/kvcache/cache_status.py b/python/freetoken/kvcache/cache_status.py index 6aa13e7f75..5b531ce265 100644 --- a/python/freetoken/kvcache/cache_status.py +++ b/python/freetoken/kvcache/cache_status.py @@ -3,13 +3,14 @@ from typing import Any, Dict +def window_pool_spec(config): + from . import resolve_pool_class + + return resolve_pool_class(config.model_config).window_spec(config) + + def _supports_swa_ratio(config) -> bool: - """Whether ``swa_full_tokens_ratio`` sizes a separate window pool for this model -- DSV4 - (always) or a radix-SWA model (Gemma). Gates the ratio in telemetry and rebuild.""" - mc = config.model_config - if mc.dsv4_args is not None: - return True - return mc.has_swa_attention and config.cache_type == "swa_radix" + return window_pool_spec(config) is not None def compute_cache_unit_bytes(engine: "Engine") -> Dict[str, int]: @@ -120,18 +121,8 @@ def _mamba() -> int: return int(_linear_pool_min_slots(config) - 1) def _swa() -> int: - # Report usable tokens, matching num_swa_pages in rebuild requests. - # DSV4's physical floor includes a dummy page; the SWA floor already excludes slot 0. - from .dsv4_cost_model import _dsv4_window_floor_pages - from .hybrid_swa_pool import _swa_pool_floor - - mc = config.model_config - if mc.dsv4_args is not None: - P = mc.dsv4_args.window_size - return int((_dsv4_window_floor_pages(config, P) - 1) * P) - if not (mc.has_swa_attention and config.cache_type == "swa_radix"): - return 0 - return int(_swa_pool_floor(config)) + spec = window_pool_spec(config) + return spec.page_size * spec.min_pages if spec is not None else 0 for key, fn in (("kv_tokens", _kv), ("moe_experts", _moe), ("mamba_slots", _mamba), ("swa_tokens", _swa)): @@ -157,19 +148,10 @@ def compute_cache_pools(engine: "Engine") -> Dict[str, int]: pools["num_pages"] = int(engine.num_pages or 0) if config is not None: pools["page_size"] = int(config.page_size or 0) - # Window pool: its own page unit (swa_page_size -- DSV4 windows are P-token pages, - # radix-SWA is token-granular page_size 1) and the concrete current size in that unit - # (num_swa_pages, usable count). Same source as the scheduler's _current_cache_geometry. - # Both 0 for models without a window pool. Lets a client denominate the swa control. - mc = config.model_config - if mc.dsv4_args is not None: - pools["swa_page_size"] = int(mc.dsv4_args.window_size or 0) - sizes = getattr(engine.kv_cache, "sizes", None) # usable = physical minus dummy - if sizes is not None: - pools["num_swa_pages"] = max(0, int(sizes.n_win_pages) - 1) - elif mc.has_swa_attention and config.cache_type == "swa_radix": - pools["swa_page_size"] = 1 # usable = pool tokens minus the slot-0 sentinel - pools["num_swa_pages"] = max(0, int(getattr(engine.kv_cache, "swa_num_tokens", 0) or 0) - 1) + spec = window_pool_spec(config) + if spec is not None: + pools["swa_page_size"] = spec.page_size + pools["num_swa_pages"] = max(0, int(engine.kv_cache.window_pages)) moe = engine.moe_offload_cache if moe is not None: pools["moe_cache_size"] = int(moe.cache_size or 0) diff --git a/python/freetoken/kvcache/dsa_pool.py b/python/freetoken/kvcache/dsa_pool.py index 13225cec06..d7c36fd734 100644 --- a/python/freetoken/kvcache/dsa_pool.py +++ b/python/freetoken/kvcache/dsa_pool.py @@ -2,7 +2,7 @@ ``MLAKVCache`` is the MHA pool's sibling for latent-KV models: ONE slab holding the per-token latent ``ckv (kv_lora_rank) | kpe (qk_rope_head_dim)`` -- there is no -separate V (``v_cache`` aliases ``k_cache``, same convention as dsv4_paged_pool's +separate V (``v_cache`` aliases ``k_cache``, same convention as dsv4/v4_pool's single-latent tiers). ``DSAKVCache`` extends it with the DeepSeek-Sparse-Attention index-key slab: one ``index_head_dim``-wide bf16 row per token per full-indexer layer, addressed by the SAME physical rows as the latent slab (GLM-5.2, page 1; @@ -75,7 +75,7 @@ def k_cache(self, layer_id: int) -> torch.Tensor: ) def v_cache(self, layer_id: int) -> torch.Tensor: - # MLA: K == V (single latent); same buffer, dsv4_paged_pool precedent. + # MLA: K == V (single latent); same buffer, dsv4/v4_pool precedent. return self.k_cache(layer_id) def latent_rows(self, layer_id: int) -> torch.Tensor: diff --git a/python/freetoken/kvcache/dsv4/__init__.py b/python/freetoken/kvcache/dsv4/__init__.py new file mode 100644 index 0000000000..7531470f0f --- /dev/null +++ b/python/freetoken/kvcache/dsv4/__init__.py @@ -0,0 +1 @@ +"""Paged KV pools for the DeepSeek-V4 family.""" diff --git a/python/freetoken/kvcache/dsv4/v41_cost_model.py b/python/freetoken/kvcache/dsv4/v41_cost_model.py new file mode 100644 index 0000000000..4db1fb18c1 --- /dev/null +++ b/python/freetoken/kvcache/dsv4/v41_cost_model.py @@ -0,0 +1,216 @@ +"""Cost model + per-tier sizing for the DSV41 paged KV pool (DeepSeek-V4.1). + +Same shape as ``v4_cost_model``: one bytes-per-P-token number for the budget division, and +independent per-tier sizing from the anchor ``full_token = num_pages * P``. The tiers: + +* window pool, every layer -- ``swa_ratio`` of the full history, packed fp8 rows +* main KV pool, per kv source -- ``full_token // ratio`` packed fp4 rows +* index-key pool, per kv source -- ``full_token // ratio`` packed fp4 rows +* compress-state ring, ratio>1 sources -- ``ring_size`` fp32 ``kv|score`` slots per WINDOW page + +Sizing reads ``ModelConfig.dsv41_args`` by attribute, like ``v4_cost_model`` reads ``dsv4_args``: +the kvcache package never imports the model package. All byte widths come from the tier row formats +below, so there is exactly one place that knows what a row costs. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from .v41_row_format import FP4_E4M3_B16, FP4_E8M0_B32, FP8_E8M0_B32 +from .window_tier import reserved_window_pages + +_INT64_BYTES = 8 +_FP32_BYTES = 4 + +WIN_FMT = FP8_E8M0_B32 +MAIN_FMT = FP4_E4M3_B16 +IDX_FMT = FP4_E8M0_B32 + + +def _win_row_bytes(args) -> int: + return WIN_FMT.row_bytes(args.head_dim) + + +def _main_row_bytes(args) -> int: + return MAIN_FMT.row_bytes(args.head_dim) + + +def _idx_row_bytes(args) -> int: + return IDX_FMT.row_bytes(args.index_head_dim) + + +def _state_bytes(args) -> int: + """One compress-state ring slot: fp32 ``kv | score`` of the latent width.""" + return 2 * args.head_dim * _FP32_BYTES + + +def private_window_layer_ids(args) -> tuple[int, ...]: + """Layers whose window KV is REQUEST-PRIVATE: kept in a per-page-table-row ring of ``window`` + slots (slot = row * window + pos % window) instead of the shared, page-bound window pool, so + it never lands in radix-shared pages. The decoder layers under Decoder SWA Bounded Replay: + their KV is recomputed per request over the prompt's last window (positions before it are + never computed), so a shared page could hold rows another request never wrote. Exact mode + computes every position and keeps the decoder on the shared pool.""" + if args.swa_decoder_replay == "exact": + return () + return tuple(range(args.decoder_start_layer, args.n_layers)) + + +@dataclass +class DSV41PoolSizes: + """Per-tier row counts derived from the budget anchor ``full_token``; the global tiers are + keyed by kv-source layer id.""" + + P: int + swa_ratio: float + full_token: int + n_win_slots: int + n_win_pages: int + main_rows: dict[int, int] = field(default_factory=dict) + idx_rows: dict[int, int] = field(default_factory=dict) + state_slots: dict[int, int] = field(default_factory=dict) # ratio>1 sources only + + +def dsv41_pool_sizes( + num_pages: int, args, swa_ratio: float, P: int, n_win_pages: int | None = None +) -> DSV41PoolSizes: + full_token = num_pages * P + if n_win_pages is None: + n_win_pages = (round(swa_ratio * full_token) + P - 1) // P + n_win_pages = min(n_win_pages, num_pages) + sizes = DSV41PoolSizes(P=P, swa_ratio=swa_ratio, full_token=full_token, n_win_slots=n_win_pages * P, n_win_pages=n_win_pages) + for src in args.backbone_kv_sources: + ratio = args.compress_ratios[src] + sizes.main_rows[src] = full_token // ratio + sizes.idx_rows[src] = full_token // ratio + if ratio > 1: + sizes.state_slots[src] = n_win_pages * ratio # ring size == ratio + return sizes + + +def dsv41_pool_bytes(sizes: DSV41PoolSizes, args, n_scratch: int = 1) -> int: + """Exact bytes a ``DSV41PagedKVCache`` built from ``sizes`` allocates (mirror of ``total_bytes``).""" + n_private = len(private_window_layer_ids(args)) + total = (args.n_layers - n_private) * sizes.n_win_slots * _win_row_bytes(args) + total += n_private * n_scratch * sizes.P * _win_row_bytes(args) # per-request rings + total += (sizes.full_token + 1) * _INT64_BYTES # full_to_window (+ sentinel row) + for src in args.backbone_kv_sources: + total += (sizes.main_rows[src] + n_scratch) * _main_row_bytes(args) + total += (sizes.idx_rows[src] + n_scratch) * _idx_row_bytes(args) + if src in sizes.state_slots: + total += (sizes.state_slots[src] + 1) * _state_bytes(args) + return int(total) + + +def dsv41_cache_per_page(args, swa_ratio: float, P: int) -> int: + """Marginal bytes per P-token page across all tiers (window scaled by ``swa_ratio``).""" + n_shared = args.n_layers - len(private_window_layer_ids(args)) + total = n_shared * round(swa_ratio * P) * _win_row_bytes(args) + for src in args.backbone_kv_sources: + ratio = args.compress_ratios[src] + total += (P // ratio) * (_main_row_bytes(args) + _idx_row_bytes(args)) + if ratio > 1: + total += round(swa_ratio * ratio) * _state_bytes(args) + return int(total) + + +def dsv41_kv_unit_bytes(args, P: int) -> int: + """FULL-tier bytes per full-history token (main + index pools + the mapping); the slider's + ``kv_bytes_per_token``. 890 B/token for DeepSeek-V4.1-Flash before the mapping.""" + per_page = P * _INT64_BYTES + for src in args.backbone_kv_sources: + per_page += (P // args.compress_ratios[src]) * (_main_row_bytes(args) + _idx_row_bytes(args)) + return -(-per_page // P) + + +def dsv41_window_unit_bytes(args, P: int) -> int: + """WINDOW-tier bytes per window token (sliding KV on every shared-window layer + the state rings).""" + per_page = (args.n_layers - len(private_window_layer_ids(args))) * P * _win_row_bytes(args) + for src in args.backbone_kv_sources: + ratio = args.compress_ratios[src] + if ratio > 1: + per_page += ratio * _state_bytes(args) + return -(-per_page // P) + + +def dsv41_solve_num_pages( + available_bytes: int, args, swa_ratio: float, floor_win_pages: int, P: int, n_scratch: int = 1 +) -> DSV41PoolSizes: + """Largest budget-respecting pool; the window floor is honored in pages and the total is + byte-checked. Raises ``ValueError`` when even the minimal pool does not fit.""" + + def _sizes(num: int) -> DSV41PoolSizes: + win = max(floor_win_pages, (round(swa_ratio * num * P) + P - 1) // P) + return dsv41_pool_sizes(num, args, swa_ratio, P, n_win_pages=win) + + lo = max(floor_win_pages, 2) + if dsv41_pool_bytes(_sizes(lo), args, n_scratch) > available_bytes: + raise ValueError( + f"DSV41 KV budget {available_bytes} bytes cannot fit the minimal pool ({lo} pages incl. the " + f"window working-set floor {floor_win_pages}); raise memory_ratio or lower max_running_req/max_seq_len" + ) + hi = max(lo, available_bytes // max(1, dsv41_cache_per_page(args, 0.0, P))) + while dsv41_pool_bytes(_sizes(hi), args, n_scratch) <= available_bytes: + hi *= 2 + while lo < hi - 1: + mid = (lo + hi) // 2 + if dsv41_pool_bytes(_sizes(mid), args, n_scratch) <= available_bytes: + lo = mid + else: + hi = mid + return _sizes(lo) + + +def dsv41_auto_cost_model(args, swa_ratio: float, floor_win_pages: int, P: int, n_scratch: int = 1): + """Affine ``(cache_per_page, fixed_cache_size, min_reserve_tokens)`` for the MoE-first auto planner.""" + per_page = dsv41_cache_per_page(args, swa_ratio, P) + P * _INT64_BYTES + n0 = max(floor_win_pages, 2) + win0 = max(floor_win_pages, (round(swa_ratio * n0 * P) + P - 1) // P) + base = dsv41_pool_bytes(dsv41_pool_sizes(n0, args, swa_ratio, P, n_win_pages=win0), args, n_scratch) + # The engine combines this structural floor with the configured KV reserve. + return per_page, max(0, base - n0 * per_page), n0 * P + + +# ---- config-facing sizing (EngineConfig in, sizes out) ---- + + +def _dsv41_swa_ratio(config) -> float: + return float(config.swa_full_tokens_ratio) + + +def _dsv41_window_floor_pages(config, args) -> int: + """Minimum window pages: one prefill chunk's reach (capped at 8 pages) + the running requests' + working set (see ``reserved_window_pages``).""" + P = args.window_size + prefill_reach_pages = (config.max_seq_len + P - 1) // P + radix = config.cache_type != "naive" + return min(prefill_reach_pages, 8) + reserved_window_pages(config.max_running_req, radix) + + +def _dsv41_pool_sizes(config, num_pages: int, num_swa_pages: int | None = None) -> DSV41PoolSizes: + """Sizes for ``num_pages`` PHYSICAL pages (dummy included). Window precedence: an explicit + ``num_swa_pages`` > ``config.swa_num_pages_override`` > ``swa_ratio`` x full, floored ONCE in pages.""" + args = config.model_config.dsv41_args + P = args.window_size + swa_ratio = _dsv41_swa_ratio(config) + floor_pages = _dsv41_window_floor_pages(config, args) + target = num_swa_pages if num_swa_pages is not None else config.swa_num_pages_override + if target is not None: + win = min(num_pages, max(floor_pages, int(target) + 1)) + else: + win = max(floor_pages, (round(swa_ratio * num_pages * P) + P - 1) // P) + return dsv41_pool_sizes(num_pages, args, swa_ratio, P, n_win_pages=win) + + +__all__ = [ + "DSV41PoolSizes", + "dsv41_auto_cost_model", + "dsv41_cache_per_page", + "dsv41_kv_unit_bytes", + "dsv41_pool_bytes", + "dsv41_pool_sizes", + "dsv41_solve_num_pages", + "dsv41_window_unit_bytes", + "private_window_layer_ids", +] diff --git a/python/freetoken/kvcache/dsv4/v41_pool.py b/python/freetoken/kvcache/dsv4/v41_pool.py new file mode 100644 index 0000000000..36d3986ae5 --- /dev/null +++ b/python/freetoken/kvcache/dsv4/v41_pool.py @@ -0,0 +1,292 @@ +"""DSV41 paged KV pool (DeepSeek-V4.1): packed byte tiers over the shared page table. + +Tiers, sized from a budget (``v41_cost_model``) not from ``num_requests``: + +* ``window_pool[L]`` -- every layer; the P-sliding KV ring in ``win_fmt`` (fp8) rows, page-granular. +* ``main_pool[src]`` -- per kv-source layer; compressed KV latents in ``main_fmt`` (fp4) rows. +* ``idx_pool[src]`` -- per kv-source layer; indexer keys in ``idx_fmt`` (fp4) rows. +* ``state_ring[src]`` -- ratio>1 sources; the per-window-page compress-state ring (fp32 ``kv|score``). + +Consumer layers (Reindex / Reuse modes) alias their source's pools: ``main_pool_of(layer)``. + +Addressing is the shared ``WindowTierPagedPool``'s (``kvcache/dsv4/window_tier.py``, also under the DSV4 +pool): the page table maps ``(table_idx, pos)`` to a full-token loc; ``full_to_window`` maps that to a +window slot (page-bound by the free list); the main / index rows are pure arithmetic ``full_loc // +ratio``; the ring slot derives from the window slot. Scratch rows past every main / index pool's +arithmetic range give a graph-safe destination for the decode rows whose compressed group did not +complete this step. + +Rows are written through ``pack_rows`` (quantize + scatter) and read inside the attention / +indexer kernels; ``read_*`` helpers dequantize for tests and torch reference paths. +""" + +from __future__ import annotations + +import torch + +from freetoken.utils import init_logger + +from .v41_cost_model import ( + IDX_FMT, + MAIN_FMT, + WIN_FMT, + DSV41PoolSizes, + _idx_row_bytes, + _main_row_bytes, + _win_row_bytes, + dsv41_kv_unit_bytes, + dsv41_window_unit_bytes, + private_window_layer_ids, +) +from .window_tier import CompressStateRing, WindowTierPagedPool + +logger = init_logger(__name__) + + +class DSV41PagedKVCache(WindowTierPagedPool): + win_fmt = WIN_FMT + main_fmt = MAIN_FMT + idx_fmt = IDX_FMT + + def __init__( + self, + sizes: DSV41PoolSizes, + args, + device: torch.device, + dtype: torch.dtype = torch.bfloat16, + n_scratch: int = 1, + ) -> None: + for fmt, dim in ((self.win_fmt, args.head_dim), (self.main_fmt, args.head_dim), (self.idx_fmt, args.index_head_dim)): + fmt.validate_dim(dim) + for src in args.backbone_kv_sources: + if args.window_size % args.compress_ratios[src]: + raise ValueError(f"window {args.window_size} must be a multiple of compress ratio {args.compress_ratios[src]}") + self.args = args + self.private_window_layer_ids = private_window_layer_ids(args) + self.sizes = sizes + self._device = device + self._dtype = dtype # the model's compute dtype; the tiers themselves are packed bytes + self.P = args.window_size + self.head_dim = args.head_dim + self.index_head_dim = args.index_head_dim + self.n_scratch = int(n_scratch) + self._paged_params: tuple[int, bool] | None = None + self.full_loc_map: torch.Tensor | None = None + self._alloc_buffers() + + # ----- buffers ----- + def _alloc_buffers(self) -> None: + sizes, args, dev = self.sizes, self.args, self._device + self.full_to_window = torch.full((sizes.full_token + 1,), -1, dtype=torch.int64, device=dev) + # shared layers: the page-bound window pool; private layers: one ring of P slots per + # page-table row (n_scratch rows, the dummy included), addressed row * P + pos % P + self.window_pool: list[torch.Tensor] = [ + torch.zeros(self.n_scratch * self.P if self.is_private_window(l) else sizes.n_win_slots, _win_row_bytes(args), dtype=torch.uint8, device=dev) + for l in range(args.n_layers) + ] + self.main_pool: dict[int, torch.Tensor] = {} + self.idx_pool: dict[int, torch.Tensor] = {} + self.state_ring: dict[int, CompressStateRing] = {} + self.scratch_base: dict[int, int] = {} + for src in args.backbone_kv_sources: + rows = sizes.main_rows[src] + self.scratch_base[src] = rows + self.main_pool[src] = torch.zeros(rows + self.n_scratch, _main_row_bytes(args), dtype=torch.uint8, device=dev) + self.idx_pool[src] = torch.zeros(sizes.idx_rows[src] + self.n_scratch, _idx_row_bytes(args), dtype=torch.uint8, device=dev) + if src in sizes.state_slots: + self.state_ring[src] = CompressStateRing( + n_slots=sizes.state_slots[src], ring_size=args.compress_ratios[src], overlap=False, + head_dim=args.head_dim, device=dev, + ) + + def total_bytes(self) -> int: + n = self.full_to_window.numel() * self.full_to_window.element_size() + n += sum(t.numel() for t in self.window_pool) + n += sum(t.numel() for t in self.main_pool.values()) + n += sum(t.numel() for t in self.idx_pool.values()) + n += sum(r.buffer.numel() * r.buffer.element_size() for r in self.state_ring.values()) + return int(n) + + @property + def prefix_replay_tokens(self) -> int: + """Prompt-tail tokens a prefix hit must leave to the prefill: bounded replay runs the decoder on + the prompt's last window from that window's encoder outputs, which no cache keeps.""" + return self.P if self.private_window_layer_ids else 0 + + # ----- private window rings ----- + def is_private_window(self, layer_id: int) -> bool: + return layer_id in self.private_window_layer_ids + + def ring_slots(self, table_rows: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + """Ring slots of ``positions`` (int64, broadcastable with ``table_rows``) for private window layers; + negative positions map to ``-1``.""" + slots = table_rows.to(torch.int64) * self.P + positions % self.P + return torch.where(positions < 0, -1, slots) + + # ----- source aliasing ----- + def source_of(self, layer_id: int) -> int: + src = self.args.roles[layer_id].kv_source + assert src is not None, f"layer {layer_id} is window-only" + return src + + def main_pool_of(self, layer_id: int) -> torch.Tensor: + return self.main_pool[self.source_of(layer_id)] + + def idx_pool_of(self, layer_id: int) -> torch.Tensor: + return self.idx_pool[self.source_of(layer_id)] + + def ratio_of(self, layer_id: int) -> int: + return self.args.roles[layer_id].ratio + + # ----- engine-facing sizing / rebuild surface ----- + @classmethod + def kv_cost(cls, config) -> tuple[int, int, int, int]: + from .v41_cost_model import _dsv41_swa_ratio, _dsv41_window_floor_pages, dsv41_auto_cost_model + + args = config.model_config.dsv41_args + P = args.window_size + per_page, fixed, min_reserve = dsv41_auto_cost_model( + args, _dsv41_swa_ratio(config), _dsv41_window_floor_pages(config, args), P, n_scratch=config.max_running_req + 1 + ) + return per_page, fixed, config.page_size, min_reserve + + @classmethod + def solve_num_pages(cls, config, available_memory: int) -> int: + from freetoken.utils import mem_GB + + from .v41_cost_model import ( + _dsv41_pool_sizes, + _dsv41_swa_ratio, + _dsv41_window_floor_pages, + dsv41_pool_bytes, + dsv41_solve_num_pages, + ) + + args = config.model_config.dsv41_args + P = args.window_size + num_pages = config.num_page_override + if num_pages is None: + sizes = dsv41_solve_num_pages( + available_memory, args, _dsv41_swa_ratio(config), _dsv41_window_floor_pages(config, args), P, + n_scratch=config.max_running_req + 1, + ) + num_pages = sizes.full_token // P - 1 # one physical page is the dummy + else: + floor = _dsv41_window_floor_pages(config, args) + if num_pages < floor: + raise ValueError( + f"--num-pages {num_pages} ({num_pages * P} tokens) is below the DSV41 window working-set " + f"floor {floor} pages ({floor * P} tokens); raise --num-pages or lower max_running_req/max_seq_len" + ) + sizes = _dsv41_pool_sizes(config, num_pages + 1) + assert num_pages > 1, "Not enough memory for KV cache, try reducing --num-pages" + real = dsv41_pool_bytes(sizes, args, config.max_running_req + 1) + logger.info( + f"Allocating {num_pages * P} tokens for DSV41 KV cache ({sizes.n_win_pages} window pages), total = {mem_GB(real)}" + ) + return num_pages + + @classmethod + def window_spec(cls, config): + from ..base import WindowPoolSpec + from .v41_cost_model import _dsv41_window_floor_pages + + args = config.model_config.dsv41_args + return WindowPoolSpec(args.window_size, _dsv41_window_floor_pages(config, args) - 1) + + @classmethod + def min_kv_tokens(cls, config) -> int: + from .v41_cost_model import _dsv41_window_floor_pages + + args = config.model_config.dsv41_args + return _dsv41_window_floor_pages(config, args) * args.window_size + + def validate_rebuild( + self, config, *, num_pages: int | None, target_moe: int, per_expert_bytes: int, + baseline_free: int, weights_bytes: int, current_num_pages: int, + extra_fixed_bytes: int = 0, extra_note: str = "", + num_swa_pages: int | None = None, **targets, + ) -> None: + from freetoken.engine.cache_budget import net_cache_budget_bytes + from freetoken.utils import mem_GB + + from ..base import CacheRebuildRejected + from .v41_cost_model import _dsv41_pool_sizes, _dsv41_window_floor_pages, dsv41_pool_bytes + + if num_pages is not None: + floor = _dsv41_window_floor_pages(config, self.args) + if num_pages < floor: + raise CacheRebuildRejected( + f"num_pages {num_pages} is below the DSV41 window working-set floor {floor} " + f"(max_running_req={config.max_running_req}); admission would deadlock" + ) + if num_pages is not None or num_swa_pages is not None: + target_pages = num_pages if num_pages is not None else current_num_pages + kv_sizes = _dsv41_pool_sizes(config, target_pages + 1, num_swa_pages=num_swa_pages) + else: + kv_sizes = self.sizes + budget = net_cache_budget_bytes(config.memory_ratio, baseline_free, weights_bytes, 0) + need = target_moe * per_expert_bytes + dsv41_pool_bytes(kv_sizes, self.args, config.max_running_req + 1) + if need > budget: + kv_part = f"kv={num_pages} P-pages" if num_pages is not None else "kv=current pool" + raise CacheRebuildRejected( + f"requested cache (moe={target_moe} slots, {kv_part}) needs {mem_GB(need)} > budget {mem_GB(budget)}; " + "old cache kept, still serving" + ) + + def rebuild_from_config(self, config, num_pages: int, *, num_swa_pages: int | None = None) -> None: + from .v41_cost_model import _dsv41_pool_sizes + + self.rebuild(_dsv41_pool_sizes(config, num_pages + 1, num_swa_pages=num_swa_pages)) + + def _drop_buffers(self) -> None: + self.window_pool = self.main_pool = self.idx_pool = self.state_ring = None # type: ignore[assignment] + + def unit_bytes(self) -> tuple[int, int]: + return dsv41_kv_unit_bytes(self.args, self.P), dsv41_window_unit_bytes(self.args, self.P) + + # ----- writes (quantize + scatter) ----- + def store_window(self, kv: torch.Tensor, layer_id: int, window_slot: torch.Tensor) -> None: + from freetoken.kernel.triton.dsv41.pack import pack_rows + + pack_rows(kv, self.win_fmt, pool=self.window_pool[layer_id], row_ids=window_slot) + + def store_main(self, latent: torch.Tensor, source: int, rows: torch.Tensor) -> None: + from freetoken.kernel.triton.dsv41.pack import pack_rows + + pack_rows(latent, self.main_fmt, pool=self.main_pool[source], row_ids=rows) + + def store_index(self, k: torch.Tensor, source: int, rows: torch.Tensor) -> None: + from freetoken.kernel.triton.dsv41.pack import pack_rows + + pack_rows(k, self.idx_fmt, pool=self.idx_pool[source], row_ids=rows) + + # ----- reads (gather + dequantize; tests and torch reference paths) ----- + def read_window(self, layer_id: int, window_slot: torch.Tensor) -> torch.Tensor: + from freetoken.kernel.triton.dsv41.pack import unpack_rows + + return unpack_rows(self.window_pool[layer_id], self.win_fmt, self.head_dim, window_slot) + + def read_main(self, source: int, rows: torch.Tensor) -> torch.Tensor: + from freetoken.kernel.triton.dsv41.pack import unpack_rows + + return unpack_rows(self.main_pool[source], self.main_fmt, self.head_dim, rows) + + def read_index(self, source: int, rows: torch.Tensor) -> torch.Tensor: + from freetoken.kernel.triton.dsv41.pack import unpack_rows + + return unpack_rows(self.idx_pool[source], self.idx_fmt, self.index_head_dim, rows) + + # ----- compress-state ring accessors ----- + def get_state(self, source: int, state_loc: torch.Tensor) -> torch.Tensor: + return self.state_ring[source].get(state_loc) + + def set_state(self, source: int, state_loc: torch.Tensor, kv_score: torch.Tensor) -> None: + self.state_ring[source].set(state_loc, kv_score) + + @property + def num_layers(self) -> int: + return self.args.n_layers + + +__all__ = ["DSV41PagedKVCache"] diff --git a/python/freetoken/kvcache/dsv4/v41_row_format.py b/python/freetoken/kvcache/dsv4/v41_row_format.py new file mode 100644 index 0000000000..f273be4f56 --- /dev/null +++ b/python/freetoken/kvcache/dsv4/v41_row_format.py @@ -0,0 +1,69 @@ +"""RowFormat: how one DeepSeek-V4.1 KV row (a head_dim-wide vector) is stored in a packed byte pool. + +DeepSeek-V4.1 keeps three cache tiers in three storage formats (tech report sec. 2.4.4 and +``inference/model.py``): + +* SWA KV -- ``fp8_e8m0_b32``: e4m3 values, one ue8m0 (power-of-two) scale per 32 channels +* main KV -- ``fp4_e4m3_b16``: e2m1 values, one e4m3 scale per 16 channels (NVFP4 without the global scale) +* indexer K -- ``fp4_e8m0_b32``: e2m1 values, one ue8m0 scale per 32 channels (OCP MXFP4) + +A row is laid out as ``[values | scales]``: ``value_bytes(dim)`` bytes of (packed) values followed by +``scale_bytes(dim)`` scale bytes, so one gather fetches both. e2m1 pairs pack the even channel into the low +nibble and the odd channel into the high nibble (``torch.float4_e2m1fn_x2`` / ``convert.py`` order). + +The Triton side (``kernel/triton/dsv41/row_format.py``) selects the dequant path from ``RowFormat.code`` +(a constexpr); the pack / unpack wrappers in ``kernel/triton/dsv41/pack.py`` reproduce the reference +``act_quant`` / ``fp4_act_quant`` numerics exactly, including the bf16 rounding of the dequantized value +the reference bakes into its cache. +""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class RowFormat: + name: str + elem_bits: int # 16 (bf16), 8 (e4m3) or 4 (e2m1) + group: int # channels per scale; 0 = unscaled + scale: str | None # None | "e8m0" | "e4m3" + code: int # Triton constexpr selector + + def value_bytes(self, dim: int) -> int: + assert (dim * self.elem_bits) % 8 == 0, (dim, self.elem_bits) + return dim * self.elem_bits // 8 + + def scale_bytes(self, dim: int) -> int: + if not self.group: + return 0 + assert dim % self.group == 0, (dim, self.group) + return dim // self.group + + def row_bytes(self, dim: int) -> int: + return self.value_bytes(dim) + self.scale_bytes(dim) + + def num_groups(self, dim: int) -> int: + return dim // self.group if self.group else 0 + + def validate_dim(self, dim: int) -> None: + if dim <= 0 or (dim * self.elem_bits) % 8 or (self.group and dim % self.group): + raise ValueError(f"row format {self.name} cannot hold a {dim}-wide row") + + +BF16 = RowFormat("bf16", elem_bits=16, group=0, scale=None, code=0) +FP8_E8M0_B32 = RowFormat("fp8_e8m0_b32", elem_bits=8, group=32, scale="e8m0", code=1) +FP4_E4M3_B16 = RowFormat("fp4_e4m3_b16", elem_bits=4, group=16, scale="e4m3", code=2) +FP4_E8M0_B32 = RowFormat("fp4_e8m0_b32", elem_bits=4, group=32, scale="e8m0", code=3) + +ROW_FORMATS: dict[str, RowFormat] = {f.name: f for f in (BF16, FP8_E8M0_B32, FP4_E4M3_B16, FP4_E8M0_B32)} + + +def row_format(name: str) -> RowFormat: + try: + return ROW_FORMATS[name] + except KeyError: + raise ValueError(f"unknown row format {name!r}; one of {sorted(ROW_FORMATS)}") from None + + +__all__ = ["RowFormat", "BF16", "FP8_E8M0_B32", "FP4_E4M3_B16", "FP4_E8M0_B32", "ROW_FORMATS", "row_format"] diff --git a/python/freetoken/kvcache/dsv4_cost_model.py b/python/freetoken/kvcache/dsv4/v4_cost_model.py similarity index 96% rename from python/freetoken/kvcache/dsv4_cost_model.py rename to python/freetoken/kvcache/dsv4/v4_cost_model.py index 2c084aff84..c6f5b7ce41 100644 --- a/python/freetoken/kvcache/dsv4_cost_model.py +++ b/python/freetoken/kvcache/dsv4/v4_cost_model.py @@ -28,14 +28,12 @@ def dsv4_reserved_window_pages(max_running_req: int, radix: bool) -> int: - """Window pages the sliding pool must always keep for the concurrent working set: each - running request's decode transients (2 per req + dummy) plus, in radix mode, PER concurrent - request one locked live-tail page AND a retained (soft-pinned) prompt-end window -- the - window is 2 pages here because the retention gap page-aligns to a whole extra page at - P==window==128, so a distinct follow-up per running request can re-lock 2 pages each. - Shared by the engine's window-floor and the manager's prefill_chunk_budget so both - reserve the same set.""" - return 2 * (max_running_req + 1) + (3 * max_running_req if radix else 0) + 1 + """Window pages the sliding pool must always keep for the concurrent working set; the + formula is shared with every sliding-window paged pool (``window_tier.reserved_window_pages``) + and by the engine's window-floor and the manager's prefill_chunk_budget.""" + from .window_tier import reserved_window_pages + + return reserved_window_pages(max_running_req, radix) def ring_size_for_ratio(ratio: int) -> int: diff --git a/python/freetoken/kvcache/dsv4_paged_pool.py b/python/freetoken/kvcache/dsv4/v4_pool.py similarity index 52% rename from python/freetoken/kvcache/dsv4_paged_pool.py rename to python/freetoken/kvcache/dsv4/v4_pool.py index 3ce23eea7d..d5eae6d080 100644 --- a/python/freetoken/kvcache/dsv4_paged_pool.py +++ b/python/freetoken/kvcache/dsv4/v4_pool.py @@ -23,135 +23,22 @@ from freetoken.utils import init_logger -from .base import BaseKVCachePool -from .dsv4_cost_model import ( +from .v4_cost_model import ( DSV4PoolSizes, dsv4_kv_unit_bytes, dsv4_window_unit_bytes, ring_size_for_ratio, ) +# The window-tier building blocks and the pool base live in window_tier.py (shared with the DSV41 +# pool); the first two are re-exported here for the DSV4 callers and tests that import them from +# this module. +from .window_tier import CompressStateRing, FreeListAllocator, WindowTierPagedPool # noqa: F401 logger = init_logger(__name__) -# LIFO free-list allocator for the window tier: one allocated unit spans ``page_unit = P`` -# slots, so a unit base is always a multiple of ``P`` (the spec's ``G*P`` page-base -# invariant). The cmp/idx tiers have no allocator -- their rows are ``full_loc // ratio`` -# pure arithmetic. -class FreeListAllocator: - def __init__( - self, - capacity: int, - device: torch.device, - page_unit: int = 1, - ) -> None: - assert capacity % page_unit == 0, ( - f"capacity {capacity} must be a multiple of page_unit {page_unit}" - ) - self._capacity = int(capacity) - self._page_unit = int(page_unit) - self._device = device - # Unit base slots: [0, page_unit, 2*page_unit, ...]. LIFO -> pop the tail. - self._free = self._fresh_free() - - def _fresh_free(self) -> torch.Tensor: - n_units = self._capacity // self._page_unit - return torch.arange(n_units, dtype=torch.int64, device=self._device) * self._page_unit - - def alloc(self, n_units: int) -> torch.Tensor: - """Return ``n_units`` unit base slots (each a multiple of ``page_unit``).""" - if n_units < 0: - raise ValueError(f"n_units must be non-negative, got {n_units}") - if n_units > self._free.numel(): - raise RuntimeError( - f"FreeListAllocator out of slots: requested {n_units} units, " - f"have {self._free.numel()} (capacity {self._capacity}, unit {self._page_unit})" - ) - if n_units == 0: - return self._free[:0].clone() - taken = self._free[-n_units:].clone() - self._free = self._free[:-n_units] - return taken - - def free(self, units: torch.Tensor) -> None: - """Return previously-allocated unit base slots to the pool (LIFO recycle).""" - if units.numel() == 0: - return - units = units.to(device=self._device, dtype=torch.int64).reshape(-1) - self._free = torch.cat([self._free, units]) - - def available(self) -> int: - """Free capacity in slots (units * page_unit).""" - return int(self._free.numel()) * self._page_unit - - @property - def capacity(self) -> int: - return self._capacity - - def reset(self) -> None: - self._free = self._fresh_free() - - - -class CompressStateRing: - """Per-layer fp32 compress-state ring: ``[n_slots + 1, 2*(1+overlap)*head_dim]``. - - Last row (index ``-1``) is a permanent scratch slot, re-cleared on every - write. Last dim is split ``kv | score``; ``set_state`` writes both halves. - """ - - def __init__( - self, - n_slots: int, - ring_size: int, - overlap: bool, - head_dim: int, - device: torch.device, - dtype: torch.dtype = torch.float32, - ) -> None: - self.ring_size = ring_size - self.head_dim = head_dim - self._item_size = (1 + int(overlap)) * head_dim # width of kv (== width of score) - last_dim = 2 * self._item_size - # +1 trailing scratch row at index -1. - self.buffer = torch.zeros((n_slots + 1, last_dim), dtype=dtype, device=device) - self.n_slots = n_slots - self._clear_scratch() - - def _clear_scratch(self) -> None: - self.buffer[-1, : self._item_size].zero_() - self.buffer[-1, self._item_size :].fill_(float("-inf")) - - @property - def item_size(self) -> int: - return self._item_size - - def get(self, state_loc: torch.Tensor) -> torch.Tensor: - """Gather ``kv_score`` rows at ``state_loc`` (``-1`` -> scratch row).""" - return self.buffer[state_loc] - - def set(self, state_loc: torch.Tensor, kv_score: torch.Tensor) -> None: - """Scatter ``kv_score`` rows to ``state_loc`` then re-clear the scratch row.""" - self.buffer[state_loc] = kv_score - self._clear_scratch() - - def get_blocks(self, page_base: torch.Tensor) -> torch.Tensor: - """Batched per-row carry-block read. ``page_base`` is ``[B]`` (the ring block base - row per row, == ``(window_slot // P) * ring_size``). Returns ``[B, ring_size, 2*item]`` - — each row's whole rolling carry block. Distinct pages -> disjoint blocks (isolation).""" - rows = page_base[:, None] + torch.arange(self.ring_size, device=page_base.device) - return self.buffer[rows] - - def set_blocks(self, page_base: torch.Tensor, blocks: torch.Tensor) -> None: - """Batched per-row carry-block write: ``buffer[base+arange(ring_size)] = blocks[row]``. - ``page_base`` ``[B]``, ``blocks`` ``[B, ring_size, 2*item]``. Re-clears scratch.""" - rows = page_base[:, None] + torch.arange(self.ring_size, device=page_base.device) - self.buffer[rows] = blocks - self._clear_scratch() - - -class DSV4PagedKVCache(BaseKVCachePool): +class DSV4PagedKVCache(WindowTierPagedPool): def __init__( self, sizes: DSV4PoolSizes, @@ -273,50 +160,11 @@ def _alloc_buffers(self) -> None: ) ) - # ----- full-loc translation (gather-only -1 safety; see field comment) ----- - def translate_full_to_window(self, full_locs: torch.Tensor) -> torch.Tensor: - # int64 gather indices: the shared page_table stores full locs as int32. - return self.full_to_window[full_locs.to(dtype=torch.int64)] - - @staticmethod - def cmp_rows(full_locs: torch.Tensor, ratio: int) -> torch.Tensor: - return torch.div(full_locs.to(dtype=torch.int64), ratio, rounding_mode="floor") - - def bind_window_pages(self, full_page_base: int, window_page_base: int) -> None: - assert full_page_base % self.P == 0 and window_page_base % self.P == 0 - self.full_to_window[full_page_base : full_page_base + self.P] = torch.arange( - window_page_base, window_page_base + self.P, dtype=torch.int64, device=self._device - ) - - def unbind_window_pages(self, full_locs: torch.Tensor) -> None: - self.full_to_window[full_locs[full_locs >= 0]] = -1 - - # ----- generic swa_pool duck-type (the CacheManager plug-in surface) ----- - # ShadowRadix layering: the shared page_table is the virtual full-token coordinate; this pool - # projects it into the physical tiers. The window tier is the managed "second currency" -- - # token-face signatures (what the generic CacheManager speaks), PAGE-ATOMIC internals (window - # pages are 1:1 page-bound to full pages; the per-page state ring requires it). alloc_swa - # receives whole ascending pages (allocate_paged's _page_to_token expansion) and free paths - # are page-complete by construction (padded finish tails, align_down frontiers, page-aligned - # tree nodes) -- asserted here, not assumed. - swa_paged = True - - @property - def sliding_window_size(self) -> int: - return self.P - - @property - def prefill_chunk_budget(self) -> int: - return self._chunk_budget - # ----- engine-facing rebuild surface ----- - # The tier buffers are bound into per-forward model scratch, invalid after a realloc. - needs_rebind_on_rebuild = True - @classmethod def kv_cost(cls, config) -> tuple[int, int, int, int]: - from .dsv4_cost_model import _dsv4_swa_ratio, _dsv4_window_floor_pages - from .dsv4_cost_model import dsv4_auto_cost_model + from .v4_cost_model import _dsv4_swa_ratio, _dsv4_window_floor_pages + from .v4_cost_model import dsv4_auto_cost_model dsv4_args = config.model_config.dsv4_args P = dsv4_args.window_size @@ -335,8 +183,8 @@ def solve_num_pages(cls, config, available_memory: int) -> int: # graceful config error, not a late OOM. from freetoken.utils import mem_GB - from .dsv4_cost_model import _dsv4_pool_sizes, _dsv4_swa_ratio, _dsv4_window_floor_pages - from .dsv4_cost_model import dsv4_pool_bytes, dsv4_solve_num_pages + from .v4_cost_model import _dsv4_pool_sizes, _dsv4_swa_ratio, _dsv4_window_floor_pages + from .v4_cost_model import dsv4_pool_bytes, dsv4_solve_num_pages dsv4_args = config.model_config.dsv4_args P = dsv4_args.window_size @@ -369,11 +217,19 @@ def solve_num_pages(cls, config, available_memory: int) -> int: ) return num_pages + @classmethod + def window_spec(cls, config): + from ..base import WindowPoolSpec + from .v4_cost_model import _dsv4_window_floor_pages + + P = config.model_config.dsv4_args.window_size + return WindowPoolSpec(P, _dsv4_window_floor_pages(config, P) - 1) + @classmethod def min_kv_tokens(cls, config) -> int: # The full anchor must cover the window working-set floor (full >= window always), so # that floor -- the value validate_rebuild enforces -- is the pool's floor in tokens. - from .dsv4_cost_model import _dsv4_window_floor_pages + from .v4_cost_model import _dsv4_window_floor_pages P = config.model_config.dsv4_args.window_size return _dsv4_window_floor_pages(config, P) * P @@ -387,9 +243,9 @@ def validate_rebuild( from freetoken.engine.cache_budget import net_cache_budget_bytes from freetoken.utils import mem_GB - from .base import CacheRebuildRejected - from .dsv4_cost_model import _dsv4_pool_sizes, _dsv4_window_floor_pages - from .dsv4_cost_model import dsv4_pool_bytes + from ..base import CacheRebuildRejected + from .v4_cost_model import _dsv4_pool_sizes, _dsv4_window_floor_pages + from .v4_cost_model import dsv4_pool_bytes dsv4_args = config.model_config.dsv4_args if num_pages is not None: @@ -426,120 +282,19 @@ def validate_rebuild( def rebuild_from_config( self, config, num_pages: int, *, num_swa_pages: int | None = None ) -> None: - from .dsv4_cost_model import _dsv4_pool_sizes + from .v4_cost_model import _dsv4_pool_sizes # +1 for the dummy page self.rebuild(_dsv4_pool_sizes(config, num_pages + 1, num_swa_pages=num_swa_pages)) - def attach_page_table(self, page_table: torch.Tensor) -> None: - # The model reads full locs through full_loc_map; under the shared route that IS the - # page table. The attention backend re-allocates its decode snapshot on re-capture. - self.full_loc_map = page_table - def unit_bytes(self) -> tuple[int, int]: # No measurable flat buffer (owned paged pool): the full (cmp/idx + mapping) and window # (sliding KV + state rings) per-token costs come from the per-tier cost model. return dsv4_kv_unit_bytes(self.args, self.P), dsv4_window_unit_bytes(self.args, self.P) - def rebuild(self, sizes) -> None: - """In-place resize to ``sizes`` (identity-preserving; free-before-alloc). The manager's - tree/page bookkeeping reset is the scheduler's generic cache_manager.rebuild; the engine - re-attaches the page table via attach_page_table afterwards.""" - import gc - - assert self._paged_params is not None, "rebuild before _init_paged_state" - self.sizes = sizes + def _drop_buffers(self) -> None: self.window_pool = self.cmp_pool = self.idx_pool = None self.state_ring = self.indexer_state_ring = None - self.full_to_window = None - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - self._alloc_buffers() - self._init_paged_state(*self._paged_params) - - def _init_paged_state(self, max_running_req: int, radix: bool) -> None: - """Build the pool-owned window free-list + the tail dummy binding. The LAST full page and - LAST window page are the reserved dummy region: page_table's dummy row points at - ``full_token - P`` (== the generic ``fill_(num_tokens)`` convention with num_tokens = the - allocatable token count), permanently bound so graph-padded rows scatter to a real slot.""" - from .dsv4_cost_model import dsv4_reserved_window_pages - P = self.P - self._paged_params = (int(max_running_req), bool(radix)) - self.full_to_window.fill_(-1) - self._win_alloc = FreeListAllocator(self.sizes.n_win_slots - P, self._device, page_unit=P) - self.bind_window_pages(self.sizes.full_token - P, self.sizes.n_win_slots - P) - # Chunk cap: a batched prefill holds the whole chunk's window live at once (sliding frees - # only between chunks; peak ~2x the chunk), so reserve the concurrent working set and - # halve the rest -- the same formula the bespoke manager used. - n_win_pages = (self.sizes.n_win_slots // P) - 1 - reserved = dsv4_reserved_window_pages(max_running_req, radix) - self._chunk_budget = max(P, (n_win_pages - reserved) // 2 * P) - - @property - def swa_num_tokens(self) -> int: - # Allocatable window slots + 1: the generic capacity convention reserves slot 0 as a - # sentinel (cap == swa_num_tokens - 1); DSV4's reserved unit is the tail dummy page, - # already excluded from the free-list, so +1 re-encodes the same cap. - return (self.sizes.n_win_slots - self.P) + 1 - - def swa_available_size(self) -> int: - return int(self._win_alloc.available()) - - def alloc_swa(self, full_indices: torch.Tensor) -> None: - """Bind one window page per incoming FULL page. ``full_indices`` must be whole ascending - pages (the ``_page_to_token`` expansion); the in-page offsets are preserved - (``window_slot = wbase + pos % P``), which the state ring's page-block layout requires.""" - n = int(full_indices.numel()) - if n == 0: - return - P = self.P - assert n % P == 0, f"alloc_swa needs whole pages, got {n} slots" - fi = full_indices.to(device=self._device, dtype=torch.int64).view(-1, P) - fbases = fi[:, 0] - assert torch.equal(fi, fbases[:, None] + torch.arange(P, device=self._device)), ( - "alloc_swa pages must be contiguous ascending" - ) - wbases = self._win_alloc.alloc(fbases.numel()) # raises when exhausted (caller gated) - offsets = torch.arange(P, dtype=torch.int64, device=self._device) - self.full_to_window[(fbases[:, None] + offsets).flatten()] = ( - wbases[:, None] + offsets - ).flatten() - - def free_swa(self, full_indices: torch.Tensor) -> None: - """Return the window pages backing these FULL locs and unbind the mapping. Page-atomic: - the incoming locs must cover each touched page completely (guaranteed by the padded - finish tails / aligned frontiers / page-aligned tree values). Idempotent over already - unbound (slid/tombstoned) pages.""" - if full_indices.numel() == 0: - return - P = self.P - fi = full_indices.to(device=self._device, dtype=torch.int64) - fi = fi[fi >= 0] - if fi.numel() == 0: - return - fbases, counts = torch.unique( - torch.div(fi, P, rounding_mode="floor") * P, return_counts=True - ) - assert bool((counts == P).all()), ( - f"free_swa got partial pages (counts {counts[counts != P].tolist()[:4]})" - ) - ws = self.full_to_window[fbases] - live = ws[ws >= 0] - offsets = torch.arange(P, dtype=torch.int64, device=self._device) - self.full_to_window[(fbases[:, None] + offsets).flatten()] = -1 - if live.numel(): - self._win_alloc.free(torch.div(live, P, rounding_mode="floor") * P) - - def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor) -> torch.Tensor: - return self.full_to_window[kv_indices.to(dtype=torch.int64)] - - # ----- state_loc derivation (vectorized, LongTensor in/out) ----- - @staticmethod - def state_loc(window_slot: torch.Tensor, ring_size: int, P: int) -> torch.Tensor: - pages = torch.div(window_slot, P, rounding_mode="floor") - loc = pages * ring_size + (window_slot % ring_size) - return torch.where(window_slot < 0, torch.full_like(loc, -1), loc) def ring_size(self, layer_id: int) -> int: return ring_size_for_ratio(self.compress_ratios[layer_id]) @@ -586,26 +341,6 @@ def total_bytes(self) -> int: ) return int(n) - # ----- BaseKVCachePool interface ----- - def k_cache(self, index: int) -> torch.Tensor: - return self.window_pool[index] - - def v_cache(self, index: int) -> torch.Tensor: # MLA: K == V (single latent) - return self.window_pool[index] - - def store_kv(self, k, v, out_loc, layer_id) -> None: - # Thin window-write shim for ABC compat; DSV4 writes via the specialized - # setters above. K == V (single latent), so the window slot is out_loc. - self.store_window(k, layer_id, out_loc) - - @property - def device(self) -> torch.device: - return self._device - - @property - def dtype(self) -> torch.dtype: - return self._dtype - @property def num_layers(self) -> int: return self._n_layers diff --git a/python/freetoken/kvcache/dsv4/window_tier.py b/python/freetoken/kvcache/dsv4/window_tier.py new file mode 100644 index 0000000000..0a25c05872 --- /dev/null +++ b/python/freetoken/kvcache/dsv4/window_tier.py @@ -0,0 +1,331 @@ +"""Shared pieces of the sliding-window paged pools (DSV4, DSV41). + +* ``FreeListAllocator`` -- LIFO page allocator for the window tier (unit = one P-slot page). +* ``CompressStateRing`` -- per-window-page fp32 ``kv | score`` carry ring with a scratch row. +* ``reserved_window_pages`` -- the window working set every such pool keeps for its running requests. +* ``window_state_loc`` -- the derived (never stored) ring slot of a window slot. +* ``WindowTierPagedPool`` -- the pool base: the full-loc -> window-slot mapping over the shared page + table, the page-atomic window allocator behind the CacheManager's ``swa_pool`` duck-type, the + in-place rebuild, and the ABC glue every such pool shares. Subclasses allocate their own tiers + (``_alloc_buffers`` / ``_drop_buffers``), size themselves (the ``kv_cost`` family) and write rows. + +``DSV4PagedKVCache`` (the origin of this code) and ``DSV41PagedKVCache`` both build on these, so a +fix to allocation, freeing or the mapping lands in one place. +""" + +from __future__ import annotations + +import gc +from abc import abstractmethod + +import torch + +from ..base import BaseKVCachePool + + +class FreeListAllocator: + """LIFO free list over ``capacity // page_unit`` units; a unit's base slot is a multiple of + ``page_unit`` (the window tier's ``G*P`` page-base invariant).""" + + def __init__(self, capacity: int, device: torch.device, page_unit: int = 1) -> None: + assert capacity % page_unit == 0, f"capacity {capacity} must be a multiple of page_unit {page_unit}" + self._capacity = int(capacity) + self._page_unit = int(page_unit) + self._device = device + self._free = self._fresh_free() + + def _fresh_free(self) -> torch.Tensor: + n_units = self._capacity // self._page_unit + return torch.arange(n_units, dtype=torch.int64, device=self._device) * self._page_unit + + def alloc(self, n_units: int) -> torch.Tensor: + """Return ``n_units`` unit base slots (each a multiple of ``page_unit``).""" + if n_units < 0: + raise ValueError(f"n_units must be non-negative, got {n_units}") + if n_units > self._free.numel(): + raise RuntimeError( + f"FreeListAllocator out of slots: requested {n_units} units, " + f"have {self._free.numel()} (capacity {self._capacity}, unit {self._page_unit})" + ) + if n_units == 0: + return self._free[:0].clone() + taken = self._free[-n_units:].clone() + self._free = self._free[:-n_units] + return taken + + def free(self, units: torch.Tensor) -> None: + """Return previously-allocated unit base slots to the pool (LIFO recycle).""" + if units.numel() == 0: + return + units = units.to(device=self._device, dtype=torch.int64).reshape(-1) + self._free = torch.cat([self._free, units]) + + def available(self) -> int: + """Free capacity in slots (units * page_unit).""" + return int(self._free.numel()) * self._page_unit + + @property + def capacity(self) -> int: + return self._capacity + + def reset(self) -> None: + self._free = self._fresh_free() + + +class CompressStateRing: + """Per-layer fp32 compress-state ring: ``[n_slots + 1, 2*(1+overlap)*head_dim]``. + + Last row (index ``-1``) is a permanent scratch slot, re-cleared on every write. Last dim is + split ``kv | score``; ``set`` writes both halves. + """ + + def __init__( + self, + n_slots: int, + ring_size: int, + overlap: bool, + head_dim: int, + device: torch.device, + dtype: torch.dtype = torch.float32, + ) -> None: + self.ring_size = ring_size + self.head_dim = head_dim + self._item_size = (1 + int(overlap)) * head_dim # width of kv (== width of score) + last_dim = 2 * self._item_size + self.buffer = torch.zeros((n_slots + 1, last_dim), dtype=dtype, device=device) + self.n_slots = n_slots + self._clear_scratch() + + def _clear_scratch(self) -> None: + self.buffer[-1, : self._item_size].zero_() + self.buffer[-1, self._item_size :].fill_(float("-inf")) + + @property + def item_size(self) -> int: + return self._item_size + + def get(self, state_loc: torch.Tensor) -> torch.Tensor: + """Gather ``kv_score`` rows at ``state_loc`` (``-1`` -> scratch row).""" + return self.buffer[state_loc] + + def set(self, state_loc: torch.Tensor, kv_score: torch.Tensor) -> None: + """Scatter ``kv_score`` rows to ``state_loc`` then re-clear the scratch row.""" + self.buffer[state_loc] = kv_score + self._clear_scratch() + + def get_blocks(self, page_base: torch.Tensor) -> torch.Tensor: + """Batched per-row carry-block read: ``page_base`` ``[B]`` (``(window_slot // P) * ring_size``) + -> ``[B, ring_size, 2*item]``. Distinct pages -> disjoint blocks.""" + rows = page_base[:, None] + torch.arange(self.ring_size, device=page_base.device) + return self.buffer[rows] + + def set_blocks(self, page_base: torch.Tensor, blocks: torch.Tensor) -> None: + """Batched per-row carry-block write of ``[B, ring_size, 2*item]``; re-clears scratch.""" + rows = page_base[:, None] + torch.arange(self.ring_size, device=page_base.device) + self.buffer[rows] = blocks + self._clear_scratch() + + +def reserved_window_pages(max_running_req: int, radix: bool) -> int: + """Window pages the sliding pool must always keep for the concurrent working set: each + running request's decode transients (2 per req + dummy) plus, in radix mode, PER concurrent + request one locked live-tail page AND a retained (soft-pinned) prompt-end window -- the + window is 2 pages here because the retention gap page-aligns to a whole extra page at + P == window.""" + return 2 * (max_running_req + 1) + (3 * max_running_req if radix else 0) + 1 + + +def window_state_loc(window_slot: torch.Tensor, ring_size: int, P: int) -> torch.Tensor: + """Ring slot of a window slot: ``(ws // P) * ring_size + ws % ring_size``; ``-1`` stays ``-1`` + (the scratch row). ``ring_size | P`` keeps distinct pages on disjoint blocks.""" + pages = torch.div(window_slot, P, rounding_mode="floor") + loc = pages * ring_size + (window_slot % ring_size) + return torch.where(window_slot < 0, torch.full_like(loc, -1), loc) + + +class WindowTierPagedPool(BaseKVCachePool): + """A paged KV pool whose window tier is page-bound to the shared full-token page table. + + Subclasses set ``P`` (the window page), ``_device``, ``_dtype`` and ``sizes`` (with + ``full_token`` and ``n_win_slots``) before calling ``_alloc_buffers``, which must create + ``self.full_to_window`` (``[full_token + 1]`` int64, ``-1`` = unmapped, the trailing row a + permanent ``-1`` sentinel so a gather at ``-1`` is safe) and ``self.window_pool`` (one tensor + per layer, indexed by window slot). Scatter paths must never see a negative slot. + + ShadowRadix layering: the shared page table is the virtual full-token coordinate; the pool + projects it into physical tiers. The window tier is the managed "second currency" -- token-face + signatures (what the generic CacheManager speaks), PAGE-ATOMIC internals (window pages are 1:1 + page-bound to full pages; the per-page state ring requires it). ``alloc_swa`` receives whole + ascending pages and free paths are page-complete by construction -- asserted here, not assumed. + """ + + swa_paged = True + # the tier buffers are bound into per-forward model scratch, invalid after a realloc + needs_rebind_on_rebuild = True + + P: int + _device: torch.device + _dtype: torch.dtype + sizes: object + full_to_window: torch.Tensor + window_pool: list + + # ----- tiers (subclass) ----- + @abstractmethod + def _alloc_buffers(self) -> None: ... + + @abstractmethod + def _drop_buffers(self) -> None: + """Release every tier buffer before a rebuild reallocates them.""" + + # ----- full-loc translation (gather-only -1 safety; see the class docstring) ----- + def translate_full_to_window(self, full_locs: torch.Tensor) -> torch.Tensor: + # int64 gather indices: the shared page_table stores full locs as int32 + return self.full_to_window[full_locs.to(dtype=torch.int64)] + + @staticmethod + def cmp_rows(full_locs: torch.Tensor, ratio: int) -> torch.Tensor: + """The compressed row of a full loc: ``full_loc // ratio`` (negative stays negative).""" + return torch.div(full_locs.to(dtype=torch.int64), ratio, rounding_mode="floor") + + def bind_window_pages(self, full_page_base: int, window_page_base: int) -> None: + assert full_page_base % self.P == 0 and window_page_base % self.P == 0 + self.full_to_window[full_page_base : full_page_base + self.P] = torch.arange( + window_page_base, window_page_base + self.P, dtype=torch.int64, device=self._device + ) + + def unbind_window_pages(self, full_locs: torch.Tensor) -> None: + self.full_to_window[full_locs[full_locs >= 0]] = -1 + + @staticmethod + def state_loc(window_slot: torch.Tensor, ring_size: int, P: int) -> torch.Tensor: + return window_state_loc(window_slot, ring_size, P) + + # ----- generic swa_pool duck-type (the CacheManager plug-in surface) ----- + @property + def sliding_window_size(self) -> int: + return self.P + + @property + def prefill_chunk_budget(self) -> int: + return self._chunk_budget + + def _init_paged_state(self, max_running_req: int, radix: bool) -> None: + """Build the pool-owned window free list + the tail dummy binding. The LAST full page and + LAST window page are the reserved dummy region: page_table's dummy row points at + ``full_token - P`` (the generic ``fill_(num_tokens)`` convention), permanently bound so + graph-padded rows scatter to a real slot.""" + P = self.P + self._paged_params = (int(max_running_req), bool(radix)) + self.full_to_window.fill_(-1) + self._win_alloc = FreeListAllocator(self.sizes.n_win_slots - P, self._device, page_unit=P) + self.bind_window_pages(self.sizes.full_token - P, self.sizes.n_win_slots - P) + # Chunk cap: a batched prefill holds the whole chunk's window live at once (sliding frees + # only between chunks; peak ~2x the chunk), so reserve the concurrent working set and + # halve the rest. + n_win_pages = (self.sizes.n_win_slots // P) - 1 + reserved = reserved_window_pages(max_running_req, radix) + self._chunk_budget = max(P, (n_win_pages - reserved) // 2 * P) + + @property + def window_pages(self) -> int: + return self.sizes.n_win_pages - 1 + + @property + def swa_num_tokens(self) -> int: + # Allocatable window slots + 1: the generic capacity convention reserves slot 0 as a + # sentinel (cap == swa_num_tokens - 1); here the reserved unit is the tail dummy page, + # already excluded from the free list, so +1 re-encodes the same cap. + return (self.sizes.n_win_slots - self.P) + 1 + + def swa_available_size(self) -> int: + return int(self._win_alloc.available()) + + def alloc_swa(self, full_indices: torch.Tensor) -> None: + """Bind one window page per incoming FULL page. ``full_indices`` must be whole ascending + pages (the ``_page_to_token`` expansion); the in-page offsets are preserved + (``window_slot = wbase + pos % P``), which the state ring's page-block layout requires.""" + n = int(full_indices.numel()) + if n == 0: + return + P = self.P + assert n % P == 0, f"alloc_swa needs whole pages, got {n} slots" + fi = full_indices.to(device=self._device, dtype=torch.int64).view(-1, P) + fbases = fi[:, 0] + assert torch.equal(fi, fbases[:, None] + torch.arange(P, device=self._device)), "alloc_swa pages must be contiguous ascending" + wbases = self._win_alloc.alloc(fbases.numel()) # raises when exhausted (caller gated) + offsets = torch.arange(P, dtype=torch.int64, device=self._device) + self.full_to_window[(fbases[:, None] + offsets).flatten()] = (wbases[:, None] + offsets).flatten() + + def free_swa(self, full_indices: torch.Tensor) -> None: + """Return the window pages backing these FULL locs and unbind the mapping. Page-atomic: the + incoming locs must cover each touched page completely (guaranteed by the padded finish + tails / aligned frontiers / page-aligned tree values). Idempotent over already unbound + (slid / tombstoned) pages.""" + if full_indices.numel() == 0: + return + P = self.P + fi = full_indices.to(device=self._device, dtype=torch.int64) + fi = fi[fi >= 0] + if fi.numel() == 0: + return + fbases, counts = torch.unique(torch.div(fi, P, rounding_mode="floor") * P, return_counts=True) + assert bool((counts == P).all()), f"free_swa got partial pages (counts {counts[counts != P].tolist()[:4]})" + ws = self.full_to_window[fbases] + live = ws[ws >= 0] + offsets = torch.arange(P, dtype=torch.int64, device=self._device) + self.full_to_window[(fbases[:, None] + offsets).flatten()] = -1 + if live.numel(): + self._win_alloc.free(torch.div(live, P, rounding_mode="floor") * P) + + def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor) -> torch.Tensor: + return self.full_to_window[kv_indices.to(dtype=torch.int64)] + + # ----- rebuild / page table ----- + def rebuild(self, sizes) -> None: + """In-place resize to ``sizes`` (identity-preserving; free-before-alloc). The manager's + tree / page bookkeeping reset is the scheduler's generic cache_manager.rebuild; the engine + re-attaches the page table via attach_page_table afterwards.""" + assert self._paged_params is not None, "rebuild before _init_paged_state" + self.sizes = sizes + self._drop_buffers() + self.full_to_window = None # type: ignore[assignment] + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + self._alloc_buffers() + self._init_paged_state(*self._paged_params) + + def attach_page_table(self, page_table: torch.Tensor) -> None: + # the model reads full locs through full_loc_map; under the shared route that IS the page table + self.full_loc_map = page_table + + # ----- BaseKVCachePool glue: the window tier is the K == V latent cache ----- + def k_cache(self, index: int) -> torch.Tensor: + return self.window_pool[index] + + def v_cache(self, index: int) -> torch.Tensor: + return self.window_pool[index] + + def store_kv(self, k, v, out_loc, layer_id) -> None: + self.store_window(k, layer_id, out_loc) + + @abstractmethod + def store_window(self, kv: torch.Tensor, layer_id: int, window_slot: torch.Tensor) -> None: ... + + @property + def device(self) -> torch.device: + return self._device + + @property + def dtype(self) -> torch.dtype: + return self._dtype + + +__all__ = [ + "FreeListAllocator", + "CompressStateRing", + "WindowTierPagedPool", + "reserved_window_pages", + "window_state_loc", +] diff --git a/python/freetoken/kvcache/hybrid_swa_pool.py b/python/freetoken/kvcache/hybrid_swa_pool.py index 41c3880e34..022eb49918 100644 --- a/python/freetoken/kvcache/hybrid_swa_pool.py +++ b/python/freetoken/kvcache/hybrid_swa_pool.py @@ -8,7 +8,7 @@ from freetoken.models.config import KVCacheGroupSpec from freetoken.utils import align_ceil, div_even -from .base import BaseKVCachePool +from .base import BaseKVCachePool, WindowPoolSpec @dataclass(frozen=True) @@ -295,6 +295,16 @@ def rebuild(self, num_full_pages: int, num_swa_tokens: int | None = None) -> Non # by CacheManager.rebuild is then consistent by construction (idle-only). self._init_swa_paged_state() + @classmethod + def window_spec(cls, config) -> WindowPoolSpec | None: + if config.cache_type != "swa_radix": + return None + return WindowPoolSpec(1, _swa_pool_floor(config)) + + @property + def window_pages(self) -> int: + return self.swa_num_tokens - 1 + @classmethod def kv_cost(cls, config, *, num_swa_pages: int | None = None) -> tuple[int, int, int, int]: """Full groups ride cache_per_page; the window pool is either a pinned absolute diff --git a/python/freetoken/layers/mhc.py b/python/freetoken/layers/mhc.py index 1cd6459d7b..dcf4915140 100644 --- a/python/freetoken/layers/mhc.py +++ b/python/freetoken/layers/mhc.py @@ -19,6 +19,16 @@ correctness baseline; the fused triton kernel (``kernel/triton/mhc.py``, dispatched by ``mhc_fused_post_pre`` below) replaces it on CUDA and is validated against this file. + +Two variants share the gate math: + +* **multi-pass** (GLM-5.3, DeepSeek-V4): ``pre`` is computed from the residual the + sublayer is about to read and applied immediately -- ``mhc_pre`` / ``mhc_fused_post_pre``. +* **single-pass** (DeepSeek-V4.1 "Single-Pass mHC", tech report sec. 2.4.1): every + sublayer consumes the ``pre`` gates the *previous* sublayer produced, so the input + mix no longer waits on this sublayer's own reduction -- ``mhc_pre_single_pass`` / + ``mhc_fused_post_pre_single_pass``. The returned ``pre_new`` is what the next + sublayer applies; the model seeds the chain with a one-hot pre on stream 0. """ from __future__ import annotations @@ -36,8 +46,8 @@ def hc_contract(x: torch.Tensor) -> torch.Tensor: return x.mean(dim=1) -def mhc_pre( - residual: torch.Tensor, # [T, n, hidden] bf16 +def mhc_gates( + residual: torch.Tensor, # [T, n, hidden] fn: torch.Tensor, # [2n + n^2, n*hidden] fp32 hc_scale: torch.Tensor, # [3] fp32 hc_base: torch.Tensor, # [2n + n^2] fp32 @@ -46,8 +56,8 @@ def mhc_pre( post_mult: float, sinkhorn_repeat: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Returns (post_mix [T, n, 1] fp32, comb_mix [T, n, n] fp32, - layer_input [T, hidden] bf16).""" + """The coefficient predictor H(X): returns (pre_mix [T, n] fp32, post_mix [T, n, 1] + fp32, comb_mix [T, n, n] fp32) from one fp32 GEMM over the flattened streams.""" n, hidden = residual.shape[-2], residual.shape[-1] t = residual.shape[0] @@ -68,11 +78,46 @@ def mhc_pre( for _ in range(sinkhorn_repeat - 1): comb = comb / (comb.sum(dim=-1, keepdim=True) + hc_eps) comb = comb / (comb.sum(dim=-2, keepdim=True) + hc_eps) + return pre_mix, post_mix.view(t, n, 1), comb - layer_input = ( - (pre_mix.unsqueeze(-1) * residual.to(torch.float32)).sum(dim=1).to(residual.dtype) - ) - return post_mix.view(t, n, 1), comb, layer_input + +def mhc_mix_input(residual: torch.Tensor, pre_mix: torch.Tensor) -> torch.Tensor: + """The sublayer input ``sum_i pre_i * res_i``: [T, n, hidden] x [T, n] -> [T, hidden].""" + return (pre_mix.unsqueeze(-1) * residual.to(torch.float32)).sum(dim=1).to(residual.dtype) + + +def mhc_pre( + residual: torch.Tensor, # [T, n, hidden] bf16 + fn: torch.Tensor, # [2n + n^2, n*hidden] fp32 + hc_scale: torch.Tensor, # [3] fp32 + hc_base: torch.Tensor, # [2n + n^2] fp32 + rms_eps: float, + hc_eps: float, + post_mult: float, + sinkhorn_repeat: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Multi-pass: gates from ``residual``, pre applied to it right away. Returns + (post_mix [T, n, 1] fp32, comb_mix [T, n, n] fp32, layer_input [T, hidden] bf16).""" + pre_mix, post_mix, comb = mhc_gates(residual, fn, hc_scale, hc_base, rms_eps, hc_eps, post_mult, sinkhorn_repeat) + return post_mix, comb, mhc_mix_input(residual, pre_mix) + + +def mhc_pre_single_pass( + residual: torch.Tensor, # [T, n, hidden] + pre_mix_in: torch.Tensor, # [T, n] fp32 -- the previous sublayer's pre gates + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_eps: float, + post_mult: float, + sinkhorn_repeat: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Single-pass: gates from ``residual`` for THIS sublayer's post/comb and the NEXT + sublayer's pre; the input is mixed with ``pre_mix_in``. Returns (pre_new [T, n] fp32, + post_mix [T, n, 1] fp32, comb_mix [T, n, n] fp32, layer_input [T, hidden]).""" + pre_new, post_mix, comb = mhc_gates(residual, fn, hc_scale, hc_base, rms_eps, hc_eps, post_mult, sinkhorn_repeat) + return pre_new, post_mix, comb, mhc_mix_input(residual, pre_mix_in.to(torch.float32)) def mhc_post( @@ -140,11 +185,69 @@ def mhc_fused_post_pre( ) +def mhc_fused_post_pre_single_pass_torch( + x: torch.Tensor, + residual: torch.Tensor, + post_mix: torch.Tensor | None, + comb_mix: torch.Tensor | None, + pre_mix_in: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_eps: float, + post_mult: float, + sinkhorn_repeat: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Decomposed single-pass reference: the previous sublayer's hc_post (skipped when + ``post_mix is None`` -- the model's first sublayer), then this sublayer's gates with + the input mixed by ``pre_mix_in``. Returns (residual_new, post_new [T, n, 1], + comb_new [T, n, n], pre_new [T, n], layer_input).""" + residual_new = residual if post_mix is None else mhc_post(x, residual, post_mix, comb_mix) + pre_new, post_new, comb_new, layer_input = mhc_pre_single_pass( + residual_new, pre_mix_in, fn, hc_scale, hc_base, rms_eps, hc_eps, post_mult, sinkhorn_repeat + ) + return residual_new, post_new, comb_new, pre_new, layer_input + + +def mhc_fused_post_pre_single_pass( + x: torch.Tensor, + residual: torch.Tensor, + post_mix: torch.Tensor | None, + comb_mix: torch.Tensor | None, + pre_mix_in: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_eps: float, + post_mult: float, + sinkhorn_repeat: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Single-pass sublayer boundary (DeepSeek-V4.1): fused triton on CUDA, torch elsewhere.""" + if residual.is_cuda: + from freetoken.kernel.triton.mhc import mhc_fused_post_pre_single_pass_triton + + return mhc_fused_post_pre_single_pass_triton( + x, residual, post_mix, comb_mix, pre_mix_in, fn, hc_scale, hc_base, + rms_eps, hc_eps, post_mult, sinkhorn_repeat, + ) + return mhc_fused_post_pre_single_pass_torch( + x, residual, post_mix, comb_mix, pre_mix_in, fn, hc_scale, hc_base, + rms_eps, hc_eps, post_mult, sinkhorn_repeat, + ) + + __all__ = [ "hc_expand", "hc_contract", + "mhc_gates", + "mhc_mix_input", "mhc_pre", + "mhc_pre_single_pass", "mhc_post", "mhc_fused_post_pre", "mhc_fused_post_pre_torch", + "mhc_fused_post_pre_single_pass", + "mhc_fused_post_pre_single_pass_torch", ] diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index 996beaa796..4918a032fe 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -308,18 +308,21 @@ def _decode_hybrid( The CPU pool is kicked off (``decode_submit``) before the GPU PCIe fetch + GEMM so the CPU overflow GEMV runs concurrently with the GPU work. Capture-safe: the routing split is device-side elementwise and the CPU submit/sync are host nodes. - Each route is computed exactly once -- the GPU weights are zeroed for CPU-assigned - routes and the CPU ids are -1 for GPU-assigned routes (the C++ kernel skips id<0). + Each route is computed exactly once -- CPU-assigned routes reach the GPU as slot -1 with + weight 0 (a kernel with ``supports_inactive_slots`` skips them without a read; any other reads + slot 0, whose storage is a payload or the bank's zero fill) and the CPU ids are -1 for + GPU-assigned routes (the C++ kernel skips id<0). """ executor = cache.cpu_executor assert executor is not None, "CPU MoE executor was not initialized" - raw = topk_ids.clone() # raw expert ids for the CPU partial - cache.ensure_experts_hybrid(self.layer_id, topk_ids) # -> slot (hit/fetched) or -1 + raw_ids = topk_ids.clone() + cache.ensure_experts_hybrid(self.layer_id, topk_ids) + on_gpu = topk_ids >= 0 + cpu_ids = torch.where(on_gpu, -1, raw_ids).contiguous() + gpu_w = torch.where(on_gpu, topk_weights, 0.0).contiguous() + gpu_slots = topk_ids if cache.collect_stats: cache.record_decode_stats_hybrid(self.layer_id) - on_gpu = topk_ids >= 0 - - cpu_ids = torch.where(on_gpu, raw.new_full((), -1), raw).contiguous() pending = executor.decode_submit(self.layer_id, hidden_states, topk_weights, cpu_ids) # Measurement knob: FREETOKEN_HYBRID_OVERLAP=0 syncs the CPU pool *before* the @@ -329,8 +332,10 @@ def _decode_hybrid( ) cache.copy_missing() - gpu_slots = topk_ids.clamp_min(0) # -1 -> slot 0 (zero-weighted below) - gpu_w = torch.where(on_gpu, topk_weights, topk_weights.new_zeros(())).contiguous() + if not (self.quant_method is not None and self.quant_method.kernel.supports_inactive_slots): + # the kernel reads every route's slot: give the inactive ones slot 0, which holds a + # complete payload or the bank's zero fill (never uninitialized storage), times weight 0 + gpu_slots = gpu_slots.clamp_min(0) gpu_routed = self._expert_gemm( cache, hidden_states, diff --git a/python/freetoken/layers/quantization/configs/base.py b/python/freetoken/layers/quantization/configs/base.py index e84dcbcbba..a64136103b 100644 --- a/python/freetoken/layers/quantization/configs/base.py +++ b/python/freetoken/layers/quantization/configs/base.py @@ -89,9 +89,13 @@ def get_quant_method(self, layer: Any, prefix: str): layer_kind = layer.quant_layer_kind kind = scheme.kind if scheme else QuantKind.NONE cls = method_class(kind, layer_kind) - cfg = LAYER_CONFIGS[layer_kind].from_layer(layer, scheme) + cfg = self.layer_config(layer, layer_kind, scheme) return cls(cfg, get_quant_backend().select(layer_kind, kind)) + def layer_config(self, layer: Any, layer_kind: LayerKind, scheme: QuantScheme | None) -> LinearConfig | MoEConfig: + """The layer's Linear / MoE config; a dialect overrides it to add what only its checkpoint knows.""" + return LAYER_CONFIGS[layer_kind].from_layer(layer, scheme) + @staticmethod def from_hf( hf_config: Any, diff --git a/python/freetoken/layers/quantization/configs/fp8.py b/python/freetoken/layers/quantization/configs/fp8.py index c23cbfdee8..3e74f0c4d0 100644 --- a/python/freetoken/layers/quantization/configs/fp8.py +++ b/python/freetoken/layers/quantization/configs/fp8.py @@ -1,20 +1,25 @@ from __future__ import annotations +import dataclasses from typing import Any, ClassVar +from ..linear import LinearConfig +from ..moe import MoEConfig from ..names import is_routed_expert, name_set, substr_set -from ..registry import register_dialect +from ..registry import LayerKind, register_dialect from ..scheme import QuantKind, QuantScheme -from ..scheme import FP8_BLOCK, fp8_block_scheme, fp8_tensor_scheme, mxfp4_scheme +from ..scheme import FP8_BLOCK_SIZES, fp8_block_scheme, fp8_tensor_scheme, mxfp4_scheme from .base import QuantConfig, Stored, cfg_get @register_dialect class Fp8BlockConfig(QuantConfig): - """HF ``quant_method: fp8`` (DeepSeek-V3 style 128x128 block scales) plus the DeepSeek-V4 e8m0 / fp4-expert variant.""" + """HF ``quant_method: fp8`` (DeepSeek-V3 style 128x128 block scales) plus the DeepSeek-V4 e8m0 / fp4-expert + variant; DeepSeek-V4.1 keeps that dialect with 32x32 blocks (``weight_block_size: [32, 32]``).""" dialect = "fp8" + # the 128-block schemes every dialect table names; ``block_scheme`` carries the checkpoint's actual block edge SCHEMES: ClassVar[dict[str, QuantScheme]] = { "BLOCK": fp8_block_scheme("float"), "BLOCK_E8M0": fp8_block_scheme("e8m0"), @@ -32,15 +37,23 @@ class Fp8BlockConfig(QuantConfig): def __init__(self, q: dict[str, Any], hf_config: Any = None, *, name_map=None, unquantized=()): super().__init__(name_map, unquantized) block = tuple(int(x) for x in (q.get("weight_block_size") or ())) - if q.get("weight_per_tensor") or block != (FP8_BLOCK, FP8_BLOCK): - raise NotImplementedError(f"fp8 checkpoint with weight_block_size={block} per_tensor={q.get('weight_per_tensor')} is not supported; only 128x128 blocks are") + square = len(block) == 2 and block[0] == block[1] and block[0] in FP8_BLOCK_SIZES + if q.get("weight_per_tensor") or not square: + raise NotImplementedError( + f"fp8 checkpoint with weight_block_size={block} per_tensor={q.get('weight_per_tensor')} is not supported; " + f"only square blocks of {FP8_BLOCK_SIZES} are" + ) + self.block = block[0] # transformers skips lm_head when the checkpoint gives no list not_convert = tuple(q.get("modules_to_not_convert") or ("lm_head",)) self.not_convert = name_set(not_convert) self.not_convert_substr = substr_set(not_convert) self.convert_tables = name_set(tuple(q.get("modules_to_convert") or ())) self.e8m0 = str(q.get("scale_fmt") or "").lower() == "ue8m0" - self.expert_fp4 = str(cfg_get(hf_config, "expert_dtype") or "").lower() == "fp4" + # ``expert_dtype`` sits at the config top level (V4) or inside quantization_config (V4.1) + expert_dtype = q.get("expert_dtype") or cfg_get(hf_config, "expert_dtype") + self.expert_fp4 = str(expert_dtype or "").lower() == "fp4" + self.block_scheme = fp8_block_scheme("e8m0" if self.e8m0 else "float", self.block) def storage(self, scheme: QuantScheme) -> dict[str, Stored]: names = super().storage(scheme) @@ -55,4 +68,11 @@ def scheme_for_name(self, name: str) -> QuantScheme | None: return None if self.expert_fp4 and is_routed_expert(name): return self.SCHEMES["EXPERT_MXFP4"] - return self.SCHEMES["BLOCK_E8M0" if self.e8m0 else "BLOCK"] + return self.block_scheme + + def layer_config(self, layer: Any, layer_kind: LayerKind, scheme: QuantScheme | None) -> LinearConfig | MoEConfig: + cfg = super().layer_config(layer, layer_kind, scheme) + if layer_kind is LayerKind.MOE and scheme is not None and scheme.kind is QuantKind.MXFP4: + # the fp4 experts quantize their activations at the same block as the dense linears + cfg = dataclasses.replace(cfg, act_block=self.block) + return cfg diff --git a/python/freetoken/layers/quantization/linear/fp8_block.py b/python/freetoken/layers/quantization/linear/fp8_block.py index 3b78e9776d..3eafdfdf5b 100644 --- a/python/freetoken/layers/quantization/linear/fp8_block.py +++ b/python/freetoken/layers/quantization/linear/fp8_block.py @@ -1,4 +1,4 @@ -"""fp8 e4m3 weight with 128x128 block scales (DeepSeek-V3 style).""" +"""fp8 e4m3 weight with square block scales: 128x128 (DeepSeek-V3 style) or 32x32 (DeepSeek-V4.1).""" from __future__ import annotations @@ -7,7 +7,7 @@ import torch from ..registry import LayerKind, register_method -from ..scheme import FP8_BLOCK as BLOCK, QuantKind +from ..scheme import FP8_BLOCK, QuantKind, fp8_block_size from .base import LinearConfig, LinearKernel, LinearMethod FP8 = torch.float8_e4m3fn @@ -18,6 +18,11 @@ def _e8m0(cfg: LinearConfig) -> bool: return cfg.scheme is not None and cfg.scheme.weight.scale == "e8m0" +def layer_block_size(layer: Any) -> int: + """The block edge a finalized layer was declared with: ``K / (scale columns)``, exact by construction.""" + return layer.weight.shape[1] // layer.weight_scale_inv.shape[1] + + class Dsv4Fp8BlockLinearKernel(LinearKernel): """DeepSeek-V4's reference path: activations quantized to fp8 with power-of-two block scales, e8m0 weight scales read as codes.""" @@ -29,7 +34,7 @@ def unusable_reason(self, cfg: LinearConfig) -> str | None: def apply(self, layer: Any, x: torch.Tensor) -> torch.Tensor: from freetoken.kernel.triton.dsv4.fp8_linear import block_fp8_linear - return block_fp8_linear(x, layer.weight, layer.weight_scale_inv, layer.bias) + return block_fp8_linear(x, layer.weight, layer.weight_scale_inv, layer.bias, block=layer_block_size(layer)) class TritonFp8BlockLinearKernel(LinearKernel): @@ -38,7 +43,11 @@ class TritonFp8BlockLinearKernel(LinearKernel): name = "triton" def unusable_reason(self, cfg: LinearConfig) -> str | None: - return "reads float block scales; e8m0 codes go to the dsv4 kernel" if _e8m0(cfg) else None + if _e8m0(cfg): + return "reads float block scales; e8m0 codes go to the dsv4 kernel" + if fp8_block_size(cfg.scheme) != FP8_BLOCK: + return f"float-scale kernel serves {FP8_BLOCK}x{FP8_BLOCK} blocks only" + return None def apply(self, layer: Any, x: torch.Tensor) -> torch.Tensor: from freetoken.kernel.triton.fp8_block_linear import block_fp8_linear @@ -52,9 +61,10 @@ class Fp8BlockLinearMethod(LinearMethod): def create_weights(self, layer: Any) -> None: g = self.cfg - if g.in_features % BLOCK or any(o % BLOCK for o in g.output_sizes): - raise ValueError(f"block-fp8 needs in/out sizes divisible by {BLOCK}, got K={g.in_features} N={g.output_sizes}") + block = fp8_block_size(g.scheme) + if g.in_features % block or any(o % block for o in g.output_sizes): + raise ValueError(f"block-fp8 needs in/out sizes divisible by {block}, got K={g.in_features} N={g.output_sizes}") layer.weight = torch.empty(g.out_features, g.in_features, dtype=FP8) # e8m0 codes stay codes for the dsv4 kernel; float scales are bf16 as the readers push them today scale_dtype = E8M0 if _e8m0(g) else torch.bfloat16 - layer.weight_scale_inv = torch.empty(g.out_features // BLOCK, g.in_features // BLOCK, dtype=scale_dtype) + layer.weight_scale_inv = torch.empty(g.out_features // block, g.in_features // block, dtype=scale_dtype) diff --git a/python/freetoken/layers/quantization/moe/base.py b/python/freetoken/layers/quantization/moe/base.py index e4a428d3e4..de1ce004fe 100644 --- a/python/freetoken/layers/quantization/moe/base.py +++ b/python/freetoken/layers/quantization/moe/base.py @@ -28,6 +28,8 @@ class MoEConfig: beta: float = 0.0 limit: float | None = None interleaved: bool = False + # dynamic fp8 activation quant block of a W4A8 expert kernel, set by the quant dialect; None = no activation quant + act_block: int | None = None has_bias: bool = False apply_router_weight_on_input: bool = False strategy: str = "resident" @@ -138,6 +140,10 @@ class MoEKernel(ABC): name: ClassVar[str] cpu_format: ClassVar[str | None] = None max_slots: ClassVar[int | None] = None + # Whether ``apply`` treats a slot of -1 as an inactive route (zero output, no read of any slot). + # Hybrid decode hands the GPU a route split with such holes; kernels without this get the holes + # pointed at slot 0 with weight 0 (slot storage is zero-filled at allocation, so that is exact). + supports_inactive_slots: ClassVar[bool] = False def unusable_reason(self, cfg: MoEConfig) -> str | None: return None diff --git a/python/freetoken/layers/quantization/moe/fp8_block.py b/python/freetoken/layers/quantization/moe/fp8_block.py index 20c3dde9dc..7552978d3f 100644 --- a/python/freetoken/layers/quantization/moe/fp8_block.py +++ b/python/freetoken/layers/quantization/moe/fp8_block.py @@ -5,7 +5,7 @@ import torch from ..registry import LayerKind, register_method -from ..scheme import FP8_BLOCK as BLOCK, QuantKind +from ..scheme import FP8_BLOCK as BLOCK, QuantKind, fp8_block_size from .base import BankSpec, ExpertView, fused_piece, gated_epilogue_reason, limit_or_inf, MoEConfig, MoEKernel, MoEMethod FP8 = torch.float8_e4m3fn @@ -24,6 +24,8 @@ def unusable_reason(self, cfg: MoEConfig) -> str | None: reason = self._common_reject(cfg, tp_ok=False, cpu_ok=False, plain_silu_only=False) if reason: return reason + if cfg.scheme is not None and fp8_block_size(cfg.scheme) != BLOCK: + return f"fp8 block MoE kernel serves {BLOCK}x{BLOCK} expert blocks only" reason = gated_epilogue_reason(cfg) if reason: return f"fp8 block MoE kernel: {reason}" diff --git a/python/freetoken/layers/quantization/moe/mxfp4.py b/python/freetoken/layers/quantization/moe/mxfp4.py index 13cd97466c..1af6d9f43e 100644 --- a/python/freetoken/layers/quantization/moe/mxfp4.py +++ b/python/freetoken/layers/quantization/moe/mxfp4.py @@ -12,16 +12,20 @@ class TritonMxfp4MoEKernel(MoEKernel): - """Standard OCP MXFP4 experts (DeepSeek-V4 ds_fp4): e2m1 pairs + e8m0 scales, no bias.""" + """Standard OCP MXFP4 experts (DeepSeek-V4 / V4.1 ds_fp4): e2m1 pairs + e8m0 scales, no bias; + W4A8 with the fp8 activation round-trips at the checkpoint's block.""" name = "triton" cpu_format = "ds_fp4" + supports_inactive_slots = True def unusable_reason(self, cfg: MoEConfig) -> str | None: if cfg.interleaved: return "standard MXFP4 kernel reads the concatenated gate|up row order" if (cfg.alpha, cfg.beta) != (1.0, 0.0): return "standard MXFP4 kernel has no alpha / beta in its swiglu" + if cfg.act_block is None: + return "W4A8 ds_fp4 experts need the fp8 activation quant block" return self._common_reject(cfg, tp_ok=False, cpu_ok=True, plain_silu_only=False) def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]: @@ -44,13 +48,14 @@ def apply(self, layer, x, topk_weights, topk_ids, view: ExpertView, *, is_prefil t = view.tensors banks = (t["gate_up"], t["gate_up_scale"], t["down"], t["down_scale"]) limit = limit_or_inf(layer) + act_block = layer.quant_method.cfg.act_block if is_prefill and view.n is not None: from freetoken.moe.fused_ds_fp4 import routed_experts_fp4_prefill - return routed_experts_fp4_prefill(x, topk_ids, topk_weights, *banks, limit, view.n) + return routed_experts_fp4_prefill(x, topk_ids, topk_weights, *banks, limit, view.n, act_block=act_block) from freetoken.moe.fused_ds_fp4 import routed_experts_fp4 - return routed_experts_fp4(x, topk_ids, topk_weights, *banks, limit) + return routed_experts_fp4(x, topk_ids, topk_weights, *banks, limit, act_block=act_block) class TritonGptossMxfp4MoEKernel(MoEKernel): diff --git a/python/freetoken/layers/quantization/scheme.py b/python/freetoken/layers/quantization/scheme.py index 668eb6fb37..89f968afc8 100644 --- a/python/freetoken/layers/quantization/scheme.py +++ b/python/freetoken/layers/quantization/scheme.py @@ -25,7 +25,8 @@ def __str__(self) -> str: return self.value -FP8_BLOCK = 128 +FP8_BLOCK = 128 # DeepSeek-V3 style square block; V4.1 exports 32x32 +FP8_BLOCK_SIZES = (32, 128) NVFP4_GROUP = 16 MX_GROUP = 32 @@ -80,8 +81,17 @@ def fp8_tensor_scheme(scale: str, *, per_row: bool = False, input_scale: bool = return QuantScheme(QuantKind.FP8_TENSOR, WeightDesc("e4m3", (1, -1) if per_row else (-1, -1), scale), roles) -def fp8_block_scheme(scale: str) -> QuantScheme: - return QuantScheme(QuantKind.FP8_BLOCK, WeightDesc("e4m3", (FP8_BLOCK, FP8_BLOCK), scale), {"weight", "weight_scale_inv"}) +def fp8_block_scheme(scale: str, block: int = FP8_BLOCK) -> QuantScheme: + """Square ``block x block`` e4m3 weight blocks, one scale each; ``block`` is the checkpoint's ``weight_block_size``.""" + if block not in FP8_BLOCK_SIZES: + raise NotImplementedError(f"fp8 block size {block} is not supported; one of {FP8_BLOCK_SIZES}") + return QuantScheme(QuantKind.FP8_BLOCK, WeightDesc("e4m3", (block, block), scale), {"weight", "weight_scale_inv"}) + + +def fp8_block_size(scheme: QuantScheme) -> int: + """The square block edge of an FP8_BLOCK scheme (128 for DeepSeek-V3 style exports, 32 for DeepSeek-V4.1).""" + assert scheme.kind is QuantKind.FP8_BLOCK, scheme + return scheme.weight.group[0] def mxfp8_scheme() -> QuantScheme: diff --git a/python/freetoken/mm/processors/deepseek_v41.py b/python/freetoken/mm/processors/deepseek_v41.py new file mode 100644 index 0000000000..e40c961569 --- /dev/null +++ b/python/freetoken/mm/processors/deepseek_v41.py @@ -0,0 +1,88 @@ +"""DeepSeek-V4.1 image resize/patch layout and learned image-span delimiters.""" + +from __future__ import annotations + +import math +import struct + +import numpy as np +import torch +from PIL import ImageOps + +from freetoken.message import MMItem +from freetoken.mm import mm_pad_value +from freetoken.mm.processor import MMProcessor, PromptReplacement, content_hash + + +def image_grid(width, height, patch, ratio, min_pixels, max_tokens, max_wh_ratio=None): + if max_wh_ratio is not None: + width = min(width, height * max_wh_ratio) + if width * height < min_pixels: + scale = (min_pixels / (width * height)) ** 0.5 + width, height = int(width * scale), int(height * scale) + w, h = math.ceil(width / patch) * patch, math.ceil(height / patch) * patch + if math.ceil(h / (patch * ratio)) * (math.ceil(w / (patch * ratio)) + 1) + 2 > max_tokens: + aspect, cell = height / width, patch * ratio + wf = math.sqrt((max_tokens - 2) / aspect + 0.25) - 0.5 + hf = wf * aspect + if wf < 1: + h, w = (max_tokens - 2) // 2 * cell, cell + elif hf < 1: + h, w = cell, (max_tokens - 3) * cell + else: + scale = min(math.floor(wf) * cell / width, math.floor(hf) * cell / height) + h, w = math.floor(height * scale / patch) * patch, math.floor(width * scale / patch) * patch + return h, w + + +class DeepseekV41MMProcessor(MMProcessor): + def __init__(self, hf_config, model_path, mm): + super().__init__(model_path, mm) + self.vc = hf_config.vision_config + self.placeholder = [hf_config.image_token_id] + self.min_pixels = self.vc.min_pixels + self.max_tokens = mm.image_max_tokens if mm.image_max_tokens is not None else self.vc.max_image_tokens + if mm.image_min_tokens is not None: + # Delimiters and row breaks count against the total span budget. + self.min_pixels = max(1, mm.image_min_tokens - 2) * (self.vc.patch_size * self.vc.downsample_ratio)**2 + allowed = {"min_pixels", "max_image_tokens", "max_wh_ratio"} + unknown = set(mm.processor_kwargs) - allowed + if unknown: + raise ValueError(f"unsupported DeepSeek image processor options: {sorted(unknown)}") + self.min_pixels = mm.processor_kwargs.get("min_pixels", self.min_pixels) + self.max_tokens = mm.processor_kwargs.get("max_image_tokens", self.max_tokens) + self.max_wh_ratio = mm.processor_kwargs.get("max_wh_ratio", self.vc.max_wh_ratio) + if self.max_tokens < 4 or self.min_pixels < 0: + raise ValueError("image token budget must be >= 4 and min_pixels must be nonnegative") + if self.max_wh_ratio is not None and self.max_wh_ratio <= 0: + raise ValueError("max_wh_ratio must be positive") + + def process(self, images): + p, r = self.vc.patch_size, self.vc.downsample_ratio + items = [] + for image in images: + h, w = image_grid(image.width, image.height, p, r, self.min_pixels, self.max_tokens, self.max_wh_ratio) + if self.max_wh_ratio is not None and image.width >= self.max_wh_ratio * image.height: + image = image.resize((w, h)) + else: + image = ImageOps.pad(image, (w, h), color=(127, 127, 127)) + x = torch.from_numpy(np.asarray(image, dtype=np.float32)).permute(2, 0, 1) / 255 + x = ((x - 0.5) / 0.5).to(torch.bfloat16) + nh, nw = h // p, w // p + patches = x.reshape(3, nh, p, nw, p).permute(1, 3, 0, 2, 4).reshape(nh * nw, 3, p, p).contiguous() + item_hash = content_hash(patches, struct.pack("<3i", 1, nh, nw)) + items.append(MMItem(modality="image", hash=item_hash, pad_value=mm_pad_value(item_hash), offsets=[], feature=patches, + model_specific_data={"grid_thw": [1, nh, nw]})) + return items + + def prompt_replacement(self, item): + _, h, w = item.grid_thw + r = self.vc.downsample_ratio + n = ((h + r - 1) // r) * ((w + r - 1) // r + 1) + 2 + # Every span row, including delimiters/newlines, has a learned visual embedding. + return PromptReplacement([self.placeholder[0]] * n) + + def dummy_items(self, dtype, device): + p = self.vc.patch_size + return [MMItem(modality="image", hash=0, pad_value=0, offsets=[[0, 4]], + feature=torch.zeros(1, 3, p, p, dtype=dtype, device=device), model_specific_data={"grid_thw": [1, 1, 1]})] diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py index c5943dce1b..55b90bb58c 100644 --- a/python/freetoken/models/config.py +++ b/python/freetoken/models/config.py @@ -205,7 +205,8 @@ class DSV4AttentionGroupConfig(BaseAttentionGroupConfig): """DSV4 sparse attention (window + compressed tiers + Lightning Indexer). Standalone class on purpose: subclassing SWAAttentionGroupConfig would flip has_swa_attention and reroute DSV4 through the SWA gates. Geometry beyond the latent width lives in - dsv4_args; the pool prices itself from there, not from this spec.""" + dsv4_args (``dsv41_args`` for the ``v41`` variant); the pool prices itself from there, + not from this spec.""" kind: ClassVar[Literal["dsv4"]] = "dsv4" cache_kind: ClassVar[Literal["dsv4_paged"]] = "dsv4_paged" @@ -213,6 +214,8 @@ class on purpose: subclassing SWAAttentionGroupConfig would flip has_swa_attenti num_kv_heads: int head_dim: int sliding_window: int # the P-token window page + # "v41": DeepSeek-V4.1's cross-layer-shared packed tiers (AttnType.DSV41, its own pool and backend) + variant: Literal["v4", "v41"] = "v4" AttentionGroupConfig: TypeAlias = ( @@ -327,6 +330,9 @@ class ModelConfig: # CSA/HCA compressors, Lightning Indexer, manifold-constrained Hyper-Connections, # hash routing). Opaque to model-agnostic engine code; None for non-DSV4 models. dsv4_args: Any | None = None + # DeepSeek-V4.1 payload (DeepseekV41Args): DSV41 layer roles, Engram, replay policy. Opaque + # to model-agnostic engine code; None for every other model. + dsv41_args: Any | None = None # GLM-5.2 (glm_moe_dsa) MLA/DSA payload (GlmMoeDsaArgs): the MLA low-rank dims and the # DSA indexer geometry the model module needs. Opaque to model-agnostic engine code; # None for every other model. @@ -448,7 +454,7 @@ def attn_type_for_layer(self, layer_id: int) -> AttnType: if isinstance(group, SWAAttentionGroupConfig): return AttnType.SWA if isinstance(group, DSV4AttentionGroupConfig): - return AttnType.DSV4 + return AttnType.DSV41 if group.variant == "v41" else AttnType.DSV4 return _full_group_attn_type(group) def kv_cache_group_specs(self) -> Tuple[KVCacheGroupSpec, ...]: @@ -502,7 +508,7 @@ def kv_cache_group_specs(self) -> Tuple[KVCacheGroupSpec, ...]: num_kv_heads=group.num_kv_heads, head_dim=group.head_dim, sliding_window=group.sliding_window, - attn_type=AttnType.DSV4, + attn_type=AttnType.DSV41 if group.variant == "v41" else AttnType.DSV4, ) ) return tuple(specs) diff --git a/python/freetoken/models/deepseek_v4/__init__.py b/python/freetoken/models/deepseek_v4/__init__.py index 37f5076ba7..e19438d1ba 100644 --- a/python/freetoken/models/deepseek_v4/__init__.py +++ b/python/freetoken/models/deepseek_v4/__init__.py @@ -10,7 +10,7 @@ DeepSeek-V4-Flash is a first-class registered model on the shared paged-KV engine: its window / compressed-attention / compressed-index KV live in DSV4-owned pools addressed by page tables, and sparse attention is a physical-slot gather (see -:mod:`freetoken.attention.dsv4_sparse` and :mod:`freetoken.kvcache.dsv4_paged_pool`). +:mod:`freetoken.attention.dsv4_sparse` and :mod:`freetoken.kvcache.dsv4.v4_pool`). """ from .args import DeepseekV4Args, load_args diff --git a/python/freetoken/models/deepseek_v4/model.py b/python/freetoken/models/deepseek_v4/model.py index a2d4387ade..9faefb83fe 100644 --- a/python/freetoken/models/deepseek_v4/model.py +++ b/python/freetoken/models/deepseek_v4/model.py @@ -4,7 +4,7 @@ + stateful KV compressors / Lightning Indexer, manifold-constrained Hyper-Connections, sqrtsoftplus / hash MoE), wired onto FreeToken's shared paged engine: - - KV lives in DSV4-owned paged pools (:class:`~freetoken.kvcache.dsv4_paged_pool.DSV4PagedKVCache`): + - KV lives in DSV4-owned paged pools (:class:`~freetoken.kvcache.dsv4.v4_pool.DSV4PagedKVCache`): each layer's window-ring + compressed KV is a region of a shared global pool, addressed by per-layer slot maps. Sparse attention is a PAGED physical-slot gather: ``sparse_attn_paged`` reads KV directly from the two global pools (window / compressed) at diff --git a/python/freetoken/models/deepseek_v41/__init__.py b/python/freetoken/models/deepseek_v41/__init__.py new file mode 100644 index 0000000000..bcdef078c7 --- /dev/null +++ b/python/freetoken/models/deepseek_v41/__init__.py @@ -0,0 +1,14 @@ +"""DeepSeek-V4.1-Flash: Causal Encoder-Decoder MoE with Compressed Sparse Attention 2. + +40 layers (20 encoder + 20 decoder), every layer a sliding window (128) plus -- from layer 2 -- global +sparse attention over cross-layer-shared compressed KV in packed fp4 (Full / Reindex / Reuse +modes, a hierarchical candidate pool in the decoder), single-pass mHC residual streams, Engram n-gram +memory at layers 1 and 14 streamed from disk, one shared + 384 routed MXFP4 experts per layer served +from the offload cache. Images use the DeepSeek ViT; DSpark speculative decoding is not served. +""" + +from .config import parse_config +from .model import DeepseekV41ForCausalLM +from .weight import ftw_side_files, iter_expert_pieces, iter_weights, iter_vision_weights + +__all__ = ["DeepseekV41ForCausalLM", "parse_config", "iter_weights", "iter_vision_weights", "iter_expert_pieces", "ftw_side_files"] diff --git a/python/freetoken/models/deepseek_v41/args.py b/python/freetoken/models/deepseek_v41/args.py new file mode 100644 index 0000000000..867dcb21a9 --- /dev/null +++ b/python/freetoken/models/deepseek_v41/args.py @@ -0,0 +1,225 @@ +"""DeepSeek-V4.1-Flash hyperparameters and the static per-layer attention roles. + +``DeepseekV41Args`` mirrors the reference ``ModelArgs`` (``inference/model.py``) field for field, +read from the HF ``config.json``'s ``text_config`` (the ``inference/config.json`` dialect is also +accepted). Everything the engine reconciles at runtime (``max_seq_len``, the replay policy) is +attached by ``parse_config`` / the engine, not read from the checkpoint. + +``LayerRole`` is the attention mode table (tech report sec. 2.3.1): every compressing layer is Full +(owns the compressed KV + indexer K), Reindex (own indexer over a shared K) or Reuse (shared +Top-K); the candidate pool of the Hierarchical Sparse Indexer (sec. 2.3.2) is produced by the +first decoder Full layer and consumed by the decoder Reindex layers after it. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Literal + +# text_config key -> ModelArgs key, where the two dialects differ +_HF_TO_ARGS = { + "hidden_size": "dim", + "moe_intermediate_size": "moe_inter_dim", + "num_hidden_layers": "n_layers", + "num_nextn_predict_layers": "n_mtp_layers", + "num_attention_heads": "n_heads", + "num_experts_per_tok": "n_activated_experts", + "scoring_func": "score_func", + "routed_scaling_factor": "route_scale", + "qk_rope_head_dim": "rope_head_dim", + "rms_norm_eps": "norm_eps", + "sliding_window": "window_size", + "kv_source_layer_ids": "kv_source_layers", + "index_source_layer_ids": "index_source_layers", + "candidate_source_layer_id": "candidate_source_layer", + "engram_pad_token_id": "engram_pad_id", + "dspark_num_experts_per_tok": "dspark_n_activated_experts", +} + + +class Mode(str, Enum): + WINDOW = "window" # sliding window only (compress_ratio == 0) + FULL = "full" # computes main KV + indexer K, runs its indexer + REINDEX = "reindex" # reuses main KV + indexer K, runs its own indexer + REUSE = "reuse" # reuses main KV and the latest Top-K + + +@dataclass(frozen=True) +class LayerRole: + layer_id: int + ratio: int + mode: Mode + kv_source: int | None # the layer whose main KV / indexer K this layer reads (itself when Full) + is_candidate_source: bool # builds the hierarchical candidate pool + uses_candidates: bool # indexes only inside the candidate pool + + @property + def compresses(self) -> bool: + return self.ratio > 0 + + @property + def has_indexer(self) -> bool: + return self.mode in (Mode.FULL, Mode.REINDEX) + + +@dataclass +class DeepseekV41Args: + vocab_size: int = 129280 + dim: int = 5120 + moe_inter_dim: int = 2304 + n_layers: int = 40 + n_mtp_layers: int = 3 + n_heads: int = 64 + # moe + n_routed_experts: int = 384 + n_shared_experts: int = 1 + n_activated_experts: int = 6 + score_func: Literal["softmax", "sigmoid", "sqrtsoftplus"] = "sqrtsoftplus" + norm_topk_prob: bool = True + route_scale: float = 1.5 + swiglu_limit: float = 10.0 + # attention: latent q, one shared latent kv head, grouped low-rank output projection + q_lora_rank: int = 1280 + head_dim: int = 512 + rope_head_dim: int = 64 + norm_eps: float = 1e-20 + o_groups: int = 8 + o_lora_rank: int = 1024 + # sparse attention + window_size: int = 128 + compress_ratios: tuple[int, ...] = () # one per layer, MTP layers included + kv_source_layers: tuple[int, ...] = () + index_source_layers: tuple[int, ...] = () + compress_rope_theta: float = 160000.0 + original_seq_len: int = 65536 + rope_theta: float = 10000.0 + rope_factor: float = 16.0 + beta_fast: int = 32 + beta_slow: int = 1 + index_n_heads: int = 32 + index_head_dim: int = 128 + index_topk: int = 512 + candidate_source_layer: int = -1 + candidate_topk_blocks: int = 0 + candidate_block_size: int = 0 + # hyper-connections (single-pass mHC) + hc_mult: int = 4 + hc_sinkhorn_iters: int = 20 + hc_eps: float = 1e-6 + # engram + engram_layer_ids: tuple[int, ...] = () + engram_num_embeddings: tuple[int, ...] = () + engram_max_ngram_size: int = 4 + engram_vocab_size: int = 0 + engram_n_heads: int = 0 + engram_head_dim: int = 0 + engram_pad_id: int = 2 + engram_compressed_vocab_size: int = 0 + # dspark (weights are not served in v1; kept so the config round-trips) + dspark_block_size: int = 0 + dspark_target_layer_ids: tuple[int, ...] = () + # ---- runtime (engine-reconciled, never read from the checkpoint) ---- + max_seq_len: int = 4096 + max_batch_size: int = 1 + # "bounded": the decoder runs on the prompt's last window_size tokens with the sliding window + # truncated there (Decoder SWA Bounded Replay, tech report sec. 3.2.2); "exact": on every token. + swa_decoder_replay: Literal["bounded", "exact"] = "bounded" + roles: tuple[LayerRole, ...] = field(default_factory=tuple, repr=False) + + def __post_init__(self) -> None: + for name in ("compress_ratios", "kv_source_layers", "index_source_layers", "engram_layer_ids", + "engram_num_embeddings", "dspark_target_layer_ids"): + setattr(self, name, tuple(int(v) for v in getattr(self, name))) + if len(self.compress_ratios) < self.n_layers: + raise ValueError(f"compress_ratios has {len(self.compress_ratios)} entries for {self.n_layers} layers") + if self.n_shared_experts != 1: + raise ValueError("DeepSeek-V4.1 has exactly one shared expert") + if self.hc_mult < 1: + raise ValueError("hc_mult must be positive") + self.roles = tuple(self._role(layer) for layer in range(self.n_layers)) + if len(self.engram_layer_ids) != len(self.engram_num_embeddings): + raise ValueError("engram_layer_ids and engram_num_embeddings differ in length") + + # ---- derived ---- + def _role(self, layer: int) -> LayerRole: + ratio = self.compress_ratios[layer] + if ratio == 0: + return LayerRole(layer, 0, Mode.WINDOW, None, False, False) + sources = [s for s in self.kv_source_layers if s <= layer] + if not sources: + raise ValueError(f"layer {layer} compresses but no kv source precedes it") + source = max(sources) + if self.compress_ratios[source] != ratio: + raise ValueError(f"layer {layer} (ratio {ratio}) reads kv source {source} (ratio {self.compress_ratios[source]})") + if layer == source: + mode = Mode.FULL + elif layer in self.index_source_layers: + mode = Mode.REINDEX + else: + mode = Mode.REUSE + cand = self.candidate_source_layer + is_cand_src = cand >= 0 and layer == cand + uses = cand >= 0 and cand < layer and mode is Mode.REINDEX + if uses and self.kv_source_layers and max(s for s in self.kv_source_layers if s <= cand) != source: + raise ValueError(f"layer {layer} would index the candidate pool of layer {cand} over a different kv source") + return LayerRole(layer, ratio, mode, source, is_cand_src, uses) + + @property + def nope_head_dim(self) -> int: + return self.head_dim - self.rope_head_dim + + @property + def backbone_kv_sources(self) -> tuple[int, ...]: + return tuple(s for s in self.kv_source_layers if s < self.n_layers) + + @property + def decoder_start_layer(self) -> int: + """The first decoder layer: the last kv source (CED projects the decoder's global KV from + its input, the final encoder hidden state).""" + return max(self.backbone_kv_sources) if self.backbone_kv_sources else self.n_layers + + @property + def engram_hash_cols(self) -> int: + return (self.engram_max_ngram_size - 1) * self.engram_n_heads + + @property + def index_softmax_scale(self) -> float: + return self.index_head_dim**-0.5 + + def freqs_params(self, layer: int) -> tuple[int, int, float, float, int, int]: + """``(rotary_dim, original_seq_len, theta, factor, beta_fast, beta_slow)``: compressing layers + rotate at the compressed theta under YaRN, window-only layers at the base theta without it.""" + if self.compress_ratios[layer]: + return self.rope_head_dim, self.original_seq_len, self.compress_rope_theta, self.rope_factor, self.beta_fast, self.beta_slow + return self.rope_head_dim, 0, self.rope_theta, self.rope_factor, self.beta_fast, self.beta_slow + + @classmethod + def from_hf(cls, hf_config: Any) -> "DeepseekV41Args": + """From the HF config (``text_config`` of the multimodal wrapper, or the text config itself, + or the ``inference/config.json`` dialect).""" + text = getattr(hf_config, "text_config", None) or hf_config + if isinstance(text, dict): + raw = text + elif hasattr(text, "to_dict"): # PretrainedConfig / RawConfigShim + raw = text.to_dict() + else: + raw = vars(text) + names = {f.name for f in cls.__dataclass_fields__.values()} - {"roles", "max_seq_len", "max_batch_size", "swa_decoder_replay"} + kwargs: dict[str, Any] = {} + for key, value in raw.items(): + key = _HF_TO_ARGS.get(key, key) + if key in names: + kwargs[key] = value + # HF nests the YaRN parameters; the inference dialect keeps them flat + scaling = raw.get("rope_scaling") or raw.get("rope_parameters") + if isinstance(scaling, dict): + if "factor" in scaling: + kwargs["rope_factor"] = scaling["factor"] + for src, dst in (("beta_fast", "beta_fast"), ("beta_slow", "beta_slow"), ("original_max_position_embeddings", "original_seq_len")): + if src in scaling: + kwargs[dst] = scaling[src] + return cls(**kwargs) + + +__all__ = ["DeepseekV41Args", "LayerRole", "Mode"] diff --git a/python/freetoken/models/deepseek_v41/attention.py b/python/freetoken/models/deepseek_v41/attention.py new file mode 100644 index 0000000000..1643ddb6d1 --- /dev/null +++ b/python/freetoken/models/deepseek_v41/attention.py @@ -0,0 +1,305 @@ +"""DSV41 attention layer (reference ``Attention``): latent MLA over a sliding window plus, on +compressing layers, ``index_topk`` compressed positions -- one paged sparse-attention call. + +Per layer mode (``LayerRole``): + +* Full -- ``publish_*`` compresses its input into the kv source's main KV (fp4) and index keys + (fp4), then its indexer selects fresh Top-K rows and publishes them on the forward's + ``SharedSelection``; the decoder's first Full layer also builds the candidate pool. +* Reindex -- selects fresh Top-K rows with its own indexer over the shared keys (inside the + candidate pool in the decoder) and publishes them. +* Reuse -- attends over the latest published Top-K rows. +* Window -- the sliding window only. + +All modes compute their own query and sliding-window KV (fp8 into the layer's window ring). The +attention output is rotated back (K == V share one rotated latent) and projected through the +grouped low-rank ``wo_a`` (block-diagonal over ``o_groups``; fp8 payload kept, W8A16 at decode) and ``wo_b``. + +Prefill runs ragged over a flat token stream tiled by ``PrefillSegment``s; the compressor / indexer +state is per request, so those run per segment. Decode is batched (one token per row) and +CUDA-graph safe: every gather reads the backend's snapshot, never the live page table. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import List + +import torch +from freetoken.attention.dsv41_sparse import PrefillSegment +from freetoken.core import get_global_ctx +from freetoken.kernel.triton.dsv4.fp8_linear import GEMV_MAX_M, dequant_block_fp8, grouped_w8a16_gemv +from freetoken.layers import BaseOP, LinearReplicated, RMSNorm +from freetoken.layers.quantization import QuantKind +from freetoken.layers.quantization.scheme import fp8_block_size + +from .args import DeepseekV41Args, LayerRole, Mode +from .compressor import Compressor +from .indexer import Indexer +from .rope import apply_rotary_emb, apply_rotary_emb_decode, get_freqs_cis + + +class DSV41Attention(BaseOP): + def __init__(self, args: DeepseekV41Args, role: LayerRole, *, quant_config=None, prefix: str = ""): + self.args = args + self.role = role + self.layer_id = role.layer_id + self.n_heads = args.n_heads + self.head_dim = args.head_dim + self.rope_dim = args.rope_head_dim + self.n_groups = args.o_groups + self.o_lora_rank = args.o_lora_rank + self.window = args.window_size + self.softmax_scale = args.head_dim**-0.5 + + self.wq_a = LinearReplicated(args.dim, args.q_lora_rank, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.wq_a") + self.q_norm = RMSNorm(args.q_lora_rank, args.norm_eps) + self.wq_b = LinearReplicated(args.q_lora_rank, args.n_heads * args.head_dim, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.wq_b") + self.wkv = LinearReplicated(args.dim, args.head_dim, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.wkv") + self.kv_norm = RMSNorm(args.head_dim, args.norm_eps) + # Block-diagonal over groups. The reference dequantizes it to bf16 and runs a grouped einsum + # (the activation stays bf16); a quantized checkpoint keeps the e4m3 payload + e8m0 scale + # here and dequantizes in-kernel -- the same math, half the bytes per decode step, and no + # bf16 copy in VRAM. An unquantized checkpoint ships it bf16. + wo_a_scheme = quant_config.scheme_for(f"{prefix}.wo_a") if quant_config is not None else None + self._wo_a_block = fp8_block_size(wo_a_scheme) if wo_a_scheme is not None and wo_a_scheme.kind is QuantKind.FP8_BLOCK else None + rows, cols = args.o_groups * args.o_lora_rank, args.n_heads * args.head_dim // args.o_groups + if self._wo_a_block is not None: + self.wo_a = torch.empty(rows, cols, dtype=torch.float8_e4m3fn) + self.wo_a_scale = torch.empty(rows // self._wo_a_block, cols // self._wo_a_block, dtype=torch.float8_e8m0fnu) + else: + self.wo_a = torch.empty(rows, cols, dtype=torch.bfloat16) + self.wo_b = LinearReplicated(args.o_groups * args.o_lora_rank, args.dim, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.wo_b") + self.attn_sink = torch.empty(args.n_heads, dtype=torch.float32) + + self.compressor = Compressor(args, self.layer_id, role.ratio, quant_config=quant_config, prefix=f"{prefix}.compressor") if role.mode is Mode.FULL else None + self.indexer = Indexer(args, role, quant_config=quant_config, prefix=f"{prefix}.indexer") if role.has_indexer else None + self._freqs_params = args.freqs_params(self.layer_id) + self._freqs: torch.Tensor | None = None + + # addressing lives on the live backend, pool buffers on the live pool (both read per access) + @property + def attn(self): + return get_global_ctx().attn_backend + + @property + def freqs_key(self) -> tuple: + """The RoPE table this layer rotates with (window layers: base theta; compressing layers: the + compressed theta under YaRN). Layers with the same key share one table.""" + return self._freqs_params + + def bind(self, device: torch.device, tables: dict | None = None) -> None: + """Attach the RoPE table; ``tables`` (key -> table) lets the model share one table per key + instead of every layer holding its own ``[max_seq_len, rd // 2]`` copy.""" + if tables is not None and self.freqs_key in tables: + self._freqs = tables[self.freqs_key] + return + rd, original_seq_len, theta, factor, beta_fast, beta_slow = self._freqs_params + self._freqs = get_freqs_cis(rd, self.args.max_seq_len, original_seq_len, theta, factor, beta_fast, beta_slow, device) + if tables is not None: + tables[self.freqs_key] = self._freqs + + @property + def freqs(self) -> torch.Tensor: + assert self._freqs is not None, "attention layer is not bound (call bind first)" + return self._freqs + + # ----- shared pieces ------------------------------------------------------------------- + def _q(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + qr = self.q_norm.forward(self.wq_a.forward(x)) + q = self.wq_b.forward(qr).unflatten(-1, (self.n_heads, self.head_dim)) + return qr, q + + def _kv(self, x: torch.Tensor) -> torch.Tensor: + return self.kv_norm.forward(self.wkv.forward(x)) + + def _wo(self, o: torch.Tensor) -> torch.Tensor: + """``wo_b(wo_a(o))``: up to ``GEMV_MAX_M`` rows the grouped W8A16 GEMV over the fp8 payload, else the + reference's grouped einsum over the bf16 (or dequantized) ``wo_a``.""" + t = o.shape[0] + o = o.reshape(t, self.n_groups, -1) + if self._wo_a_block is not None and t <= GEMV_MAX_M: + return self.wo_b.forward(grouped_w8a16_gemv(o, self.wo_a, self.wo_a_scale, block=self._wo_a_block)) + wo_a = self.wo_a if self._wo_a_block is None else dequant_block_fp8(self.wo_a, self.wo_a_scale, block=self._wo_a_block) + return self.wo_b.forward(torch.einsum("tgd,grd->tgr", o, wo_a.view(self.n_groups, self.o_lora_rank, -1)).flatten(1)) + + def load_state_dict(self, state_dict, *, prefix: str = "", _internal: bool = False) -> None: + key = f"{prefix}.wo_a" if prefix else "wo_a" + scale_key = f"{key}_scale" + if self._wo_a_block is None and scale_key in state_dict: + # a bf16 layer (no fp8 scheme) fed a quantized checkpoint dequantizes wo_a on the way in + from .weight import dequant_fp8_block + + state_dict[key] = dequant_fp8_block(state_dict[key], state_dict.pop(scale_key)) + elif self._wo_a_block is not None and scale_key not in state_dict: + # an FTW conversion made before wo_a kept its fp8 payload carries a dequantized bf16 + # tensor and no scale: say so instead of failing on the missing key + raise RuntimeError( + f"{key}: the checkpoint has no {scale_key} for the fp8 wo_a layout; an FTW checkpoint " + "converted by an earlier build stores wo_a dequantized -- reconvert it with `ft checkpoint`" + ) + super().load_state_dict(state_dict, prefix=prefix, _internal=_internal) + + # ----- prefill ----------------------------------------------------------------------------- + def publish_prefill(self, x: torch.Tensor, segments: List[PrefillSegment]) -> None: + """Full mode: compress this layer's input for every segment into the source's main KV and + index keys. Runs on EVERY new token (under bounded replay the decoder source publishes for + the whole prompt even though its attention runs on the prompt's last window only).""" + assert self.compressor is not None and self.indexer is not None + attn, ratio, src = self.attn, self.role.ratio, self.layer_id + for seg in segments: + slots = attn.window_slots_of(seg.table_idx, seg.start_pos, seg.end) + tail = None + if self.compressor.needs_tail_carry(seg.start_pos): + tail = int(attn.window_slots_of(seg.table_idx, seg.start_pos - 1, seg.start_pos).item()) + latent, starts = self.compressor.forward_prefill(x[seg.offset : seg.offset + seg.n], seg.start_pos, slots, tail) + if not latent.shape[0]: + continue + rows = attn.compressed_rows_of(seg.table_idx, starts, ratio) + freqs = self.freqs.index_select(0, starts) + attn.store_index(self.indexer.index_keys(latent, freqs), src, rows) + apply_rotary_emb(latent[..., -self.rope_dim :], freqs) + attn.store_main(latent, src, rows) + + def forward_prefill(self, x: torch.Tensor, segments: List[PrefillSegment], positions: torch.Tensor) -> torch.Tensor: + """Ragged prefill over ``x [T, dim]`` tiled by ``segments`` (``positions [T]`` absolute).""" + attn, role = self.attn, self.role + freqs = self.freqs.index_select(0, positions) + qr, q = self._q(x) + apply_rotary_emb(q[..., -self.rope_dim :], freqs) + kv = self._kv(x) + apply_rotary_emb(kv[..., -self.rope_dim :], freqs) + + sel = attn.selection + win_parts, cmp_parts, pools = [], [], [] + for seg in segments: + lo, hi = seg.offset, seg.offset + seg.n + attn.store_window(kv[lo:hi], self.layer_id, attn.layer_window_slots_of(self.layer_id, seg.table_idx, seg.start_pos, seg.end)) + win_parts.append(attn.window_topk_prefill(seg, self.layer_id)) + if not role.compresses: + continue + if role.has_indexer: + locs = attn.locs_prefill(seg.table_idx, seg.end // role.ratio * role.ratio) + cand = sel.candidates[:, lo:hi] if role.uses_candidates else None + rows, pool = self.indexer.select_prefill( + x[lo:hi], qr[lo:hi], freqs[lo:hi], start_pos=seg.start_pos, ratio=role.ratio, locs=locs, candidates=cand, + ) + cmp_parts.append(rows) + if pool is not None: + pools.append(pool) + else: + assert sel.source == role.kv_source, "Reuse layer without a published selection for its source" + cmp_parts.append(sel.topk_rows[:, lo:hi]) + if role.has_indexer: + sel.source, sel.topk_rows = role.kv_source, torch.cat(cmp_parts, dim=1) + if pools: + sel.candidates = torch.cat(pools, dim=1) + win = torch.cat(win_parts, dim=1) + topk = torch.cat([win, torch.cat(cmp_parts, dim=1)], dim=-1) if role.compresses else win + + o = attn.attend(q.unsqueeze(0), self.layer_id, topk.to(torch.int32), self.window, self.attn_sink, self.softmax_scale)[0] + apply_rotary_emb(o[..., -self.rope_dim :], freqs, inverse=True) + return self._wo(o) + + # ----- decode ------------------------------------------------------------------------------ + def forward_decode(self, x: torch.Tensor, pos: torch.Tensor, rows: torch.Tensor, dctx: "DecodeStepContext", cmp_stage_cap: int) -> torch.Tensor: + """Batched single-token attention. ``pos [B]`` device positions, ``rows [B]`` local rows into + the decode snapshot, ``dctx`` the step's layer-invariant context (ring slots, window + candidates, RoPE rows), ``cmp_stage_cap`` the static compressed staging bound (max position + any row reaches; the capture width under a graph).""" + attn, role = self.attn, self.role + B = x.shape[0] + freqs_t = dctx.freqs[self.freqs_key] + qr, q = self._q(x) + apply_rotary_emb_decode(q[..., -self.rope_dim :].unsqueeze(1), freqs_t) + kv = self._kv(x) + apply_rotary_emb_decode(kv[..., -self.rope_dim :].unsqueeze(1), freqs_t) + # the shared ring context feeds the compressor's carry ring on every layer; a request-private + # window layer stores and reads its own KV through its ring context instead + window_slots, prev_window_slots = dctx.window_slots, dctx.prev_window_slots + private = dctx.private is not None and attn.pool.is_private_window(self.layer_id) + tier = dctx.private if private else dctx.shared + attn.store_window(kv, self.layer_id, tier.slots if private else window_slots) + + cmp_counts = None + if role.compresses: + ratio, src = role.ratio, role.kv_source + n_stage = (cmp_stage_cap + 1) // ratio + sel = attn.selection + if role.mode is Mode.FULL: + latent, completed = self.compressor.forward_decode(x, pos, prev_window_slots, window_slots) + dst = attn.decode_store_rows(rows, pos, ratio, src, completed) + freqs_g = self.freqs.index_select(0, (pos + 1 - ratio).clamp_min(0)) + k = self.indexer.k_norm.forward(self.indexer.wk.forward(latent)) + apply_rotary_emb_decode(k[..., -self.rope_dim :].unsqueeze(1), freqs_g) + attn.store_index(k, src, dst) + apply_rotary_emb_decode(latent[..., -self.rope_dim :].unsqueeze(1), freqs_g) + attn.store_main(latent, src, dst) + if role.has_indexer: + cand = sel.candidates if role.uses_candidates else None + cmp_rows, counts, pool = self.indexer.select_decode(x, qr, freqs_t, pos=pos, ratio=ratio, locs=attn.snapshot(), T=n_stage, candidates=cand) + sel.source, sel.topk_rows, sel.cmp_counts = src, cmp_rows, counts + if pool is not None: + sel.candidates = pool + # the step's [window | compressed] candidate buffer of this tier takes the new rows once; + # the Reuse layers that follow read them in place (no per-layer cat / cast) + tier.topk[..., self.window :].copy_(cmp_rows) + else: + assert sel.source == src, "Reuse layer without a published selection for its source" + counts = sel.cmp_counts + topk = tier.topk + cmp_counts = counts + else: + topk = tier.window_topk + + o = attn.attend(q.view(B, 1, self.n_heads, self.head_dim), self.layer_id, topk, self.window, self.attn_sink, self.softmax_scale, cmp_counts=cmp_counts) + apply_rotary_emb_decode(o[..., -self.rope_dim :], freqs_t, inverse=True) + return self._wo(o.view(B, self.n_heads, self.head_dim)) + + +@dataclass +class TierCandidates: + """One window tier's per-step decode candidates: the request's store slot and the int32 + ``[B, 1, win]`` window candidates, plus the ``[B, 1, win + topk]`` ``[window | compressed]`` + buffer whose compressed half the index layers of that tier fill as they publish.""" + + slots: torch.Tensor + window_topk: torch.Tensor + topk: torch.Tensor + + +@dataclass +class DecodeStepContext: + """Everything a decode step resolves once for all layers (a capture records it once): the shared + pool's ring slots (the compressor's carry ring keys on them), the shared / private tier + candidates, and each RoPE table gathered at the step's positions.""" + + window_slots: torch.Tensor + prev_window_slots: torch.Tensor + shared: TierCandidates + private: TierCandidates | None + freqs: dict + + @classmethod + def build(cls, layers, md, pos: torch.Tensor, rows: torch.Tensor, window: int, topk: int) -> "DecodeStepContext": + B = pos.shape[0] + + def tier(slots, window_topk): + win32 = window_topk.to(torch.int32) + buf = torch.empty((B, 1, window + topk), dtype=torch.int32, device=pos.device) + buf[..., :window].copy_(window_topk) + return TierCandidates(slots, win32, buf) + + window_slots, prev_window_slots, window_topk = md.window_ctx(pos, rows) + private = None + if get_global_ctx().kv_cache.private_window_layer_ids: + private = tier(*md.private_window_ctx(pos, rows)) + tables: dict = {} + for layer in layers: + key = layer.freqs_key + if key not in tables: + tables[key] = layer.freqs.index_select(0, pos) + return cls(window_slots, prev_window_slots, tier(window_slots, window_topk), private, tables) + + +__all__ = ["DSV41Attention", "DecodeStepContext", "TierCandidates"] diff --git a/python/freetoken/models/deepseek_v41/compressor.py b/python/freetoken/models/deepseek_v41/compressor.py new file mode 100644 index 0000000000..5d5d07fe5b --- /dev/null +++ b/python/freetoken/models/deepseek_v41/compressor.py @@ -0,0 +1,131 @@ +"""DSV41 KV compressor: ``compress_ratio`` consecutive tokens -> one latent (reference ``Compressor``). + +* ratio 1 (the CED decoder source): a plain bf16 projection + RMSNorm, one latent per token. +* ratio > 1 (the encoder sources): ``wkv`` / ``wgate`` promoted to fp32, a softmax over the group's + gate scores pools the group's ``wkv`` rows; the trailing partial group is carried across chunks + and decode steps in the pool's compress-state ring (per window page, so a page-aligned radix hit + resumes it by value -- see ``DSV41SparseAttnBackend`` for the addressing). + +Returns the latent BEFORE RoPE: the indexer key is projected from the unrotated latent, the +attention layer rotates and quantizes it afterwards. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F +from freetoken.core import get_global_ctx +from freetoken.layers import BaseOP, LinearReplicated, RMSNorm + +from .args import DeepseekV41Args + + +class Compressor(BaseOP): + def __init__(self, args: DeepseekV41Args, layer_id: int, ratio: int, *, quant_config=None, prefix: str = ""): + assert ratio >= 1 + self.layer_id = layer_id + self.ratio = ratio + self.dim = args.dim + self.head_dim = args.head_dim + self.window = args.window_size + if ratio > 1: + # fp32 like the reference (the softmax pooling runs in fp32); the checkpoint ships bf16 + self.wkv = torch.empty(args.head_dim, args.dim, dtype=torch.float32) + self.wgate = torch.empty(args.head_dim, args.dim, dtype=torch.float32) + else: + self.wkv = LinearReplicated(args.dim, args.head_dim, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.wkv") + self.norm = RMSNorm(args.head_dim, args.norm_eps) + + @property + def attn(self): + return get_global_ctx().attn_backend + + def _project(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """fp32 ``wkv`` / ``wgate`` projections.""" + xf = x.float() + return F.linear(xf, self.wkv), F.linear(xf, self.wgate) + + # ----- prefill ----------------------------------------------------------------------- + def needs_tail_carry(self, start_pos: int) -> bool: + """Whether a prefill from ``start_pos`` resumes a partial group (never at a page-aligned start).""" + return self.ratio > 1 and start_pos % self.ratio != 0 + + def forward_prefill( + self, x: torch.Tensor, start_pos: int, window_slots: torch.Tensor, tail_window_slot: int | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Compress the tokens ``[start_pos, start_pos + n)`` of one request. + + ``window_slots`` are the tokens' window slots (the ring block of each window page carries the + partial group); ``tail_window_slot`` is the previous token's slot, needed only when + ``needs_tail_carry(start_pos)``. Returns ``(latent [G, head_dim] bf16 pre-RoPE, group_starts [G] + int64)`` for every group that completed in this call. + """ + n = x.shape[0] + device = x.device + if self.ratio == 1: + latent = self.norm.forward(self.wkv.forward(x)) + return latent, torch.arange(start_pos, start_pos + n, device=device) + + ratio, d = self.ratio, self.head_dim + kv, score = self._project(x) + # the partial group in progress at start_pos: its rows live in the previous token's ring block + offset = start_pos % ratio + if offset: + assert tail_window_slot is not None, "resuming mid-group needs the previous token's window slot" + carry = self.attn.read_carry(self.layer_id, tail_window_slot) # [ratio, 2d] + kv = torch.cat([carry[:offset, :d], kv], dim=0) + score = torch.cat([carry[:offset, d:], score], dim=0) + total = kv.shape[0] + groups = total // ratio + remainder = total % ratio + cut = groups * ratio + if groups: + kv_g = kv[:cut].view(groups, ratio, d) + score_g = score[:cut].view(groups, ratio, d) + pooled = (kv_g * score_g.softmax(dim=1)).sum(dim=1) + latent = self.norm.forward(pooled.to(x.dtype)) + else: + latent = x.new_empty(0, d) + group_starts = (start_pos - offset) + ratio * torch.arange(groups, device=device) + # ring: the reset block at every page boundary crossed (a page holds whole groups), and the + # in-progress tail (the remainder rows) on the last page so decode / the next chunk resume + end = start_pos + n + self.attn.write_boundary_carries(self.layer_id, lo=start_pos, hi=end, window_slots=window_slots) + if end % self.window != 0: + block = self._carry_block(kv[cut:], score[cut:], remainder) + self.attn.write_carry(self.layer_id, int(window_slots[-1].item()), block) + return latent, group_starts + + def _carry_block(self, kv_tail: torch.Tensor, score_tail: torch.Tensor, filled: int) -> torch.Tensor: + ratio, d = self.ratio, self.head_dim + ks = kv_tail.new_zeros(ratio, d) + ss = kv_tail.new_full((ratio, d), float("-inf")) + if filled: + ks[:filled] = kv_tail + ss[:filled] = score_tail + return torch.cat([ks, ss], dim=-1) + + # ----- decode ------------------------------------------------------------------------ + def forward_decode(self, x: torch.Tensor, pos: torch.Tensor, prev_window_slots: torch.Tensor, window_slots: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """One token per row. Returns ``(latent [B, head_dim] bf16 pre-RoPE, completed [B] bool)``; + a row whose group did not complete carries a meaningless latent the caller discards.""" + B = x.shape[0] + if self.ratio == 1: + latent = self.norm.forward(self.wkv.forward(x)) + return latent, torch.ones(B, dtype=torch.bool, device=x.device) + + from freetoken.kernel.triton.dsv4.compress import gated_pool + + ratio, d = self.ratio, self.head_dim + kv, score = self._project(x) + idx = (pos % ratio).view(B, 1, 1).expand(B, 1, d) + block = self.attn.read_carry_blocks(self.layer_id, prev_window_slots) # [B, ratio, 2d] + ks = block[..., :d].clone().scatter_(1, idx, kv.unsqueeze(1)) + ss = block[..., d:].clone().scatter_(1, idx, score.unsqueeze(1)) + self.attn.write_carry_blocks(self.layer_id, window_slots, torch.cat([ks, ss], dim=-1)) + pooled = gated_pool(ks, ss, x.dtype).view(B, d) + completed = (pos + 1) % ratio == 0 + return self.norm.forward(pooled), completed + + +__all__ = ["Compressor"] diff --git a/python/freetoken/models/deepseek_v41/config.py b/python/freetoken/models/deepseek_v41/config.py new file mode 100644 index 0000000000..9e738939c8 --- /dev/null +++ b/python/freetoken/models/deepseek_v41/config.py @@ -0,0 +1,95 @@ +"""Engine-facing config for DeepSeek-V4.1-Flash. + +``parse_config`` maps the standard transformer fields into :class:`ModelConfig`, carries the full +:class:`DeepseekV41Args` in ``ModelConfig.dsv41_args`` for the model module and the KV pool (which +sizes its tiers from it), and declares the ``v41`` variant of the DSV4 attention group. +The engine reconciles the runtime knobs (``max_seq_len``, ``swa_decoder_replay``) onto ``dsv41_args`` +at config resolution. +""" + +from __future__ import annotations + +from dataclasses import dataclass, fields +from typing import Any + +from freetoken.models.config import DSV4AttentionGroupConfig, ModelConfig, RotaryConfig + +from .args import DeepseekV41Args + + +@dataclass(frozen=True) +class VisionConfig: + num_hidden_layers: int = 32 + hidden_size: int = 1024 + num_attention_heads: int = 16 + intermediate_size: int = 2816 + patch_size: int = 14 + rope_theta: float = 10000.0 + downsample_ratio: int = 3 + max_image_tokens: int = 1024 + min_pixels: int = 544 * 544 + max_wh_ratio: float | None = None + + +def parse_config(hf_config: Any) -> ModelConfig: + args = DeepseekV41Args.from_hf(hf_config) + text = getattr(hf_config, "text_config", None) or hf_config + max_position = int(getattr(text, "max_position_embeddings", 0) or getattr(hf_config, "max_position_embeddings", 0) or 0) + if max_position <= 0: + max_position = int(args.original_seq_len * args.rope_factor) + vc = getattr(hf_config, "vision_config", None) + vision = VisionConfig(**{f.name: getattr(vc, f.name, f.default) for f in fields(VisionConfig)}) if vc is not None else None + + return ModelConfig( + num_layers=args.n_layers, + num_qo_heads=args.n_heads, + num_kv_heads=1, # one shared latent KV head (K == V) + head_dim=args.head_dim, + hidden_size=args.dim, + vocab_size=args.vocab_size, + intermediate_size=args.moe_inter_dim, + hidden_act="silu", + rms_norm_eps=args.norm_eps, + tie_word_embeddings=False, + rotary_config=RotaryConfig( + head_dim=args.head_dim, + rotary_dim=args.rope_head_dim, + max_position=max_position, + base=args.rope_theta, + scaling={ + "rope_type": "yarn", + "factor": args.rope_factor, + "beta_fast": args.beta_fast, + "beta_slow": args.beta_slow, + "original_max_position_embeddings": args.original_seq_len, + }, + ), + num_experts=args.n_routed_experts, + num_experts_per_tok=args.n_activated_experts, + moe_intermediate_size=args.moe_inter_dim, + norm_topk_prob=args.norm_topk_prob, + n_shared_experts=args.n_shared_experts, + routed_scaling_factor=args.route_scale, + swiglu_limit=args.swiglu_limit, + model_type="deepseek_v41", + architectures=["DeepseekV41ForCausalLM"], + moe_enabled=True, + expert_quant="ds_fp4", + attn_sm_scale=args.head_dim**-0.5, + dsv41_args=args, + vision_config=vision, + image_token_id=getattr(hf_config, "image_token_id", None), + attention_groups=( + DSV4AttentionGroupConfig( + name="dsv41", + layer_ids=tuple(range(args.n_layers)), + num_kv_heads=1, + head_dim=args.head_dim, + sliding_window=args.window_size, + variant="v41", + ), + ), + ) + + +__all__ = ["parse_config"] diff --git a/python/freetoken/models/deepseek_v41/engram.py b/python/freetoken/models/deepseek_v41/engram.py new file mode 100644 index 0000000000..1c997a7cfd --- /dev/null +++ b/python/freetoken/models/deepseek_v41/engram.py @@ -0,0 +1,262 @@ +"""Engram conditional memory (reference ``inference/engram.py`` + ``Engram``): n-gram hash lookups +added into the residual streams before layers 1 and 14. + +Hashing (host side, ``EngramHash``): token ids are first mapped through a compressed vocabulary +(tokens that normalize alike collapse together), then each position is hashed with the +``max_ngram_size - 1`` tokens before it -- a rolling XOR of ``id * multiplier`` over the lookback, +one bucket range (a distinct prime modulus) per (n-gram order, head). Lookback stops at the sequence +start (pad id) and never crosses an image span. Addresses depend +only on the token ids, so the rows can be fetched before the forward starts. + +Lookup (``EngramTable`` protocol): the 384M-row fp8 tables live on disk; a backend gathers the +``n_hash_cols`` rows of every token into a device tensor ``[T, n_hash_cols * head_dim]`` in bf16. + +Gate (``EngramLayer``): ``wkv`` (fp8 linear) turns the gathered rows into one key per residual +stream plus a shared value; the gate is a normalized dot product of stream and key, signed-sqrt +then sigmoid; ``stream += gate * value``. +""" + +from __future__ import annotations + +from typing import Protocol + +import numpy as np +import torch +from freetoken.layers import BaseOP, LinearReplicated + +from .args import DeepseekV41Args + + +# ----- hashing --------------------------------------------------------------------------------- + + +def _is_prime(n: int) -> bool: + """Deterministic Miller-Rabin for n < 3.3e24 (the bases below are a proven set), matching + ``sympy.isprime`` on every 64-bit input.""" + if n < 2: + return False + small = (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41) + for p in small: + if n % p == 0: + return n == p + d, s = n - 1, 0 + while d % 2 == 0: + d //= 2 + s += 1 + for a in small: + x = pow(a, d, n) + if x in (1, n - 1): + continue + for _ in range(s - 1): + x = pow(x, 2, n) + if x == n - 1: + break + else: + return False + return True + + +def find_next_prime(start: int, seen: set[int]) -> int: + """The smallest prime above ``start`` not handed out yet (reference ``find_next_prime``).""" + candidate = start + 1 + while not _is_prime(candidate) or candidate in seen: + candidate += 1 + return candidate + + +def bucket_primes(args: DeepseekV41Args) -> list[list[list[int]]]: + """``[layer][ngram order - 2][head]`` bucket moduli: distinct primes drawn in order from + ``engram_vocab_size - 1`` upward, never reused across layers / orders / heads.""" + primes: list[list[list[int]]] = [] + seen: set[int] = set() + for _ in args.engram_layer_ids: + per_ngram = [] + for _ in range(args.engram_max_ngram_size - 1): + sizes, current = [], args.engram_vocab_size - 1 + for _ in range(args.engram_n_heads): + current = find_next_prime(current, seen) + seen.add(current) + sizes.append(current) + per_ngram.append(sizes) + primes.append(per_ngram) + return primes + + +def hash_multipliers(layer_ids: tuple[int, ...], max_ngram_size: int, compressed_vocab_size: int) -> torch.Tensor: + """``[n_layers, max_ngram_size]`` odd int64 multipliers from a per-layer RNG (reference + ``compute_hash_multipliers``), bounded so ``id * multiplier`` cannot overflow int64.""" + max_long = np.iinfo(np.int64).max + bound = max(1, (max_long // compressed_vocab_size) // 2) + rows = [] + for layer_id in layer_ids: + gen = np.random.default_rng(10007 * layer_id) + values = gen.integers(low=0, high=bound, size=(max_ngram_size,), dtype=np.int64) + rows.append(torch.tensor(values * 2 + 1)) + return torch.stack(rows) + + +def build_compressed_token_map(tokenizer) -> tuple[list[int], int]: + """Token id -> compressed id: tokens that normalize alike (NFKC, accents stripped, lowercased, + whitespace collapsed) share one id. Reference ``build_compressed_token_map``; the size feeds the + multipliers, so it is asserted against ``engram_compressed_vocab_size``.""" + from tokenizers import Regex, normalizers + + sentinel = "\ue000" + normalizer = normalizers.Sequence([ + normalizers.NFKC(), + normalizers.NFD(), + normalizers.StripAccents(), + normalizers.Lowercase(), + normalizers.Replace(Regex(r"[ \t\r\n]+"), " "), + normalizers.Replace(Regex(r"^ $"), sentinel), + normalizers.Strip(), + normalizers.Replace(sentinel, " "), + ]) + backend = tokenizer.backend_tokenizer + key_to_new: dict[str, int] = {} + lookup = [0] * len(tokenizer) + for token_id in range(len(tokenizer)): + text = backend.decode([token_id], skip_special_tokens=False) + if "\ufffd" in text: + key = backend.id_to_token(token_id) + else: + normalized = normalizer.normalize_str(text) + key = normalized if normalized else text + new_id = key_to_new.get(key) + if new_id is None: + new_id = len(key_to_new) + key_to_new[key] = new_id + lookup[token_id] = new_id + return lookup, len(key_to_new) + + +class EngramHash: + """Maps token runs to table rows: ``row_ids(tokens) -> [T, n_layers, n_hash_cols]`` int64. + + ``token_map`` is the compressed vocabulary (``build_compressed_token_map``); ``pad_id`` is the + compressed id of ``engram_pad_id``. Runs on the host (numpy / CPU torch) since the addresses + are needed before the forward launches. + """ + + def __init__(self, args: DeepseekV41Args, token_map: list[int] | torch.Tensor, compressed_vocab_size: int): + if compressed_vocab_size != args.engram_compressed_vocab_size: + raise ValueError( + f"compressed vocab has {compressed_vocab_size} ids, config says {args.engram_compressed_vocab_size}; " + "the hash multipliers derive from it, so the tokenizer does not match the checkpoint" + ) + self.args = args + self.max_ngram = args.engram_max_ngram_size + self.token_map = torch.as_tensor(token_map, dtype=torch.int64) + self.pad_id = int(self.token_map[args.engram_pad_id]) + primes = bucket_primes(args) + flat = [[p for per_ngram in layer for p in per_ngram] for layer in primes] + self.primes = torch.tensor(primes, dtype=torch.int64) # [L, ngram-1, heads] + self.offsets = torch.tensor(np.array([np.cumsum([0, *sizes[:-1]]) for sizes in flat]), dtype=torch.int64) # [L, cols] + self.multipliers = hash_multipliers(args.engram_layer_ids, self.max_ngram, compressed_vocab_size) # [L, ngram] + self.rows_per_layer = [int(o[-1] + s[-1]) for o, s in zip(self.offsets.tolist(), [[p for pn in layer for p in pn] for layer in primes])] + + @property + def n_cols(self) -> int: + return (self.max_ngram - 1) * self.args.engram_n_heads + + def row_ids(self, tokens: torch.Tensor, context: torch.Tensor | None = None) -> torch.Tensor: + """Rows for ``tokens [T]`` (raw ids, one contiguous run). ``context`` holds the + ``max_ngram_size - 1`` raw ids immediately before the run (fewer -> the run starts a + sequence and the missing lookback is the pad id). Returns ``[T, n_layers, n_hash_cols]``.""" + tokens = tokens.to(torch.int64).cpu() + ctx_len = self.max_ngram - 1 + if context is None: + context = tokens.new_empty(0) + context = context.to(torch.int64).cpu()[-ctx_len:] if ctx_len else tokens.new_empty(0) + raw = torch.cat([context, tokens]) + # Content-pad ids identify image spans; n-grams stop at them, including across chunks. + dead = (raw < 0) | (raw >= self.token_map.numel()) + compressed = self.token_map[raw.clamp(0, self.token_map.numel() - 1)] + n_ctx = context.numel() + T = tokens.numel() + positions = torch.arange(n_ctx, n_ctx + T) + cols = [] + blocked = torch.zeros(T, dtype=torch.bool) + for shift in range(self.max_ngram): + src = positions - shift + blocked |= (src < 0) | dead[src.clamp_min(0)] + cols.append(torch.where(blocked, torch.full_like(src, self.pad_id), compressed[src.clamp_min(0)])) + stack = torch.stack(cols, dim=-1) # [T, max_ngram]; column s is the token s places back + products = stack.unsqueeze(1) * self.multipliers # [T, L, max_ngram] + rolling, hashes = products[..., 0], [] + for i in range(1, self.max_ngram): + rolling = torch.bitwise_xor(rolling, products[..., i]) + hashes.append(rolling.unsqueeze(-1) % self.primes[:, i - 1]) # [T, L, heads] + return torch.cat(hashes, dim=-1) + self.offsets # [T, L, cols] + + +# ----- lookup ------------------------------------------------------------------------------------ + + +class EngramTable(Protocol): + """Row store of one Engram layer: ``lookup`` returns the dequantized rows of the current + forward's tokens as ``[T, n_hash_cols * head_dim]`` bf16 (staged by the host before launch).""" + + def lookup(self, num_tokens: int) -> torch.Tensor: ... + + +class ZeroEngramTable: + """Stand-in when no table is attached (dummy weights / unit tests): every lookup reads zeros.""" + + def __init__(self, width: int, device: torch.device) -> None: + self.width = width + self.device = device + + def lookup(self, num_tokens: int) -> torch.Tensor: + return torch.zeros(num_tokens, self.width, dtype=torch.bfloat16, device=self.device) + + +# ----- the layer ----------------------------------------------------------------------------------- + + +class EngramLayer(BaseOP): + """``streams [T, hc, dim] += gate * value`` with key / value from the gathered n-gram rows.""" + + def __init__(self, args: DeepseekV41Args, layer_id: int, *, quant_config=None, prefix: str = ""): + self.layer_id = layer_id + self.dim = args.dim + self.hc_mult = args.hc_mult + self.eps = args.norm_eps + self.clamp_value = 1e-6 + self.width = args.engram_hash_cols * args.engram_head_dim + self.wkv = LinearReplicated(self.width, args.dim * (args.hc_mult + 1), has_bias=False, quant_config=quant_config, prefix=f"{prefix}.wkv") + self.q_weight = torch.empty(args.hc_mult, args.dim, dtype=torch.float32) + self.k_weight = torch.empty(args.hc_mult, args.dim, dtype=torch.float32) + self._table: EngramTable | None = None + + def attach_table(self, table: EngramTable) -> None: + self._table = table + + def forward(self, streams: torch.Tensor, rows: torch.Tensor | None = None, *, image_mask: torch.Tensor | None = None) -> torch.Tensor: + """``streams [T, hc, dim]``; ``rows`` overrides the table lookup (tests / reference parity).""" + T = streams.shape[0] + if rows is None: + assert self._table is not None, "Engram table was never attached" + rows = self._table.lookup(T) + kv = self.wkv.forward(rows.to(streams.dtype)) + key, value = kv.split([self.hc_mult * self.dim, self.dim], dim=-1) + key = key.float().view(T, self.hc_mult, self.dim) + weight = self.q_weight * self.k_weight # only ever used as a product + h = streams.float() + # normalized per (token, stream) over dim, not jointly over the streams + rstd = torch.rsqrt(h.square().mean(-1) + self.eps) * torch.rsqrt(key.square().mean(-1) + self.eps) + dot = (h * weight * key).sum(-1) * rstd * self.dim**-0.5 + gate = torch.sigmoid(torch.copysign(dot.abs().clamp_min(self.clamp_value).sqrt(), dot)) + out = (h + gate.unsqueeze(-1) * value.float().unsqueeze(1)).to(streams.dtype) + return out if image_mask is None else torch.where(image_mask[:, None, None], streams, out) + + +__all__ = [ + "EngramHash", + "EngramLayer", + "EngramTable", + "ZeroEngramTable", + "build_compressed_token_map", + "bucket_primes", + "hash_multipliers", +] diff --git a/python/freetoken/models/deepseek_v41/engram_table.py b/python/freetoken/models/deepseek_v41/engram_table.py new file mode 100644 index 0000000000..55450bbd80 --- /dev/null +++ b/python/freetoken/models/deepseek_v41/engram_table.py @@ -0,0 +1,296 @@ +"""Engram tables streamed directly from checkpoint shards. + +Each Engram layer owns a ``[384M, 256]`` fp8 table (98 GB) plus a ``[384M, 8]`` ue8m0 scale +tensor (3 GB). The tables stay where the checkpoint put them: ``EngramDiskTable`` maps the shard +tensor as one ``RowStore`` extent and reads the 24 rows a token needs per fill (io_uring, O_DIRECT, +duplicates deduplicated), while the scales are host-resident and gathered next to the values, so +the staged row is exactly the pool row format ``fp8_e8m0_b32`` and the device dequantizes it with +the same kernel the KV pools use. + +Protocol (the Qwen3.8 PLE disk table's, ``models/qwen4_exp/ple_disk.py``): the engine wraps every +dispatch in ``EngramHost.forward_host_ctx``. Eager: hash + stage + read before the launch. CUDA-graph +decode: the sampled token lives on the device under overlap scheduling, so the host reads it back, +fills the graph's pinned staging and releases a per-table flag the captured ``lookup`` WAITs on (or, +without stream memops, fills before the replay). The fill runs on a worker thread that is started +BEFORE the graph is launched: with a graph this large (40 layers) the driver's launch call itself +blocks once the GPU sits on the WAIT with the rest of the graph queued behind it, so a fill scheduled +after the launch returns would deadlock. The worker works from an immutable dispatch snapshot plus +the readback event (which the GPU completes on its own), and the dispatch context waits for it after +the launch so its failure surfaces before the step's output is consumed. +""" + +from __future__ import annotations + +import json +import os +import struct +from concurrent.futures import Future, ThreadPoolExecutor +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Sequence + +import torch +from freetoken.core import Batch +from freetoken.kernel.pinned import alloc_pinned_tensor +from freetoken.kvcache.dsv4.v41_row_format import FP8_E8M0_B32 +from freetoken.utils import init_logger + +from .engram import EngramHash + +logger = init_logger(__name__) + + +@dataclass(frozen=True) +class _TensorExtent: + """Where a safetensors tensor's bytes sit in its shard: a contiguous ``[offset, offset + nbytes)``.""" + + path: str + offset: int + nbytes: int + shape: tuple[int, ...] + dtype: str # the safetensors dtype string (e.g. "F8_E4M3") + + +def _tensor_extent(folder: str, name: str) -> _TensorExtent: + """Locate ``name`` in the checkpoint directory without reading it (the table is streamed in place).""" + from freetoken.models.loader import safetensors_weight_map + + path = os.path.join(folder, safetensors_weight_map(folder)[name]) + with open(path, "rb") as fh: + n = struct.unpack(" EngramRowSource: + """Locate a layer's table in whatever safetensors files the checkpoint directory holds: the HF + shards (through the index) or the shards an FTW conversion copied next to it (through their headers).""" + from safetensors import safe_open + + weight = _tensor_extent(folder, f"layers.{layer_id}.engram.embed.weight") + if weight.dtype != "F8_E4M3": + raise ValueError(f"engram table of layer {layer_id} is {weight.dtype}, expected F8_E4M3") + rows, head_dim = weight.shape + scale = _tensor_extent(folder, f"layers.{layer_id}.engram.embed.scale") + with safe_open(scale.path, framework="pt", device="cpu") as f: + scales = f.get_tensor(f"layers.{layer_id}.engram.embed.scale").view(torch.uint8).contiguous() + if tuple(scales.shape) != (rows, head_dim // 32): + raise ValueError(f"engram scales of layer {layer_id} are {tuple(scales.shape)}, expected {(rows, head_dim // 32)}") + return EngramRowSource(weight.path, weight.offset, rows, head_dim, scales) + + +class EngramDiskTable: + """One layer's table: stages ``[T, n_cols]`` rows into pinned memory on the host, dequantizes on + the device. Implements the model's ``EngramTable`` protocol (``lookup``).""" + + def __init__( + self, source: EngramRowSource, n_cols: int, device: torch.device, *, max_graph_rows: int, max_extend_tokens: int, + use_io_uring: bool = True, + ) -> None: + from freetoken.kernel.row_store import RowStore + + self.head_dim = source.head_dim + self.n_cols = n_cols + self.num_rows = source.num_rows + self.row_bytes = FP8_E8M0_B32.row_bytes(source.head_dim) # values | scales + self.scales = source.scales + self.device = device + self.store = RowStore( + paths=[source.path], extent_file=[0], extent_base=[source.base], rows_per_extent=source.num_rows, + row_bytes=source.head_dim, row_stride=source.head_dim, use_io_uring=use_io_uring, + ) + # the graph staging is allocated up front (pinned alloc inside capture is illegal) and + # outlives any one capture so a re-capture keeps the same pointers; the eager staging + # grows to the largest prefill chunk seen (never inside a capture) + self.max_graph_rows = max_graph_rows + self._graph_pinned = alloc_pinned_tensor(max_graph_rows * n_cols * self.row_bytes, dtype=torch.uint8) + self._graph_pinned.zero_() + self._graph_dev = torch.empty(max_graph_rows * n_cols * self.row_bytes, dtype=torch.uint8, device=device) + # completes once the last graph launch that read ``_graph_pinned`` has run (recorded by the + # host coordinator after the launch: a record inside the capture would not be observable) + self.graph_consumed = torch.cuda.Event() + # the overlap scheduler stages batch k+1 while batch k's lookup copy may still be queued: the + # eager staging alternates two pinned buffers, and each is rewritten only once the H2D copy + # that last read it has run (``_eager_read[i]``, recorded right after that copy) + self._eager_pinned: list[torch.Tensor] = [] + self._eager_read = [torch.cuda.Event(), torch.cuda.Event()] + self._eager_slot = 0 + self._eager_dev = None + self._grow_eager(max_extend_tokens) + self.flag = alloc_pinned_tensor(1, dtype=torch.int64) + self.flag.zero_() + self.wait_sync = False # set by the host coordinator after its probe + + def _grow_eager(self, num_tokens: int) -> None: + nbytes = num_tokens * self.n_cols * self.row_bytes + if self._eager_pinned and self._eager_pinned[0].numel() >= nbytes: + return + for event in self._eager_read: + event.synchronize() + self._eager_pinned = [alloc_pinned_tensor(nbytes, dtype=torch.uint8) for _ in self._eager_read] + for buf in self._eager_pinned: + buf.zero_() # a warmup prefill stages nothing and reads whatever sits here + self._eager_dev = torch.empty(nbytes, dtype=torch.uint8, device=self.device) + + # ----- host ----- + def stage(self, row_ids: torch.Tensor, *, graph: bool) -> None: + """Queue ``row_ids [T, n_cols]`` (int64, host): values through the store, scales from the resident tensor.""" + n = row_ids.numel() + if graph: + self.graph_consumed.synchronize() + pinned = self._graph_pinned + assert n <= self.max_graph_rows * self.n_cols, f"graph staging holds {self.max_graph_rows} rows, batch has {n // self.n_cols}" + else: + self._grow_eager(n // self.n_cols) + self._eager_slot ^= 1 + self._eager_read[self._eager_slot].synchronize() + pinned = self._eager_pinned[self._eager_slot] + ids = row_ids.reshape(-1).to(torch.int64).contiguous() + rows = pinned[: n * self.row_bytes].view(n, self.row_bytes) + self.store.stage_rows(ids.data_ptr(), n, rows.data_ptr(), self.row_bytes) + rows[:, self.head_dim :].copy_(self.scales.index_select(0, ids)) + + def flush(self, *, signal: bool) -> None: + self.store.flush(self.flag.data_ptr() if signal else 0) + + # ----- device (EngramTable protocol) ----- + def lookup(self, num_tokens: int) -> torch.Tensor: + from freetoken.kernel.row_store import wait_reset + from freetoken.kernel.triton.dsv41.pack import unpack_rows + + stream = torch.cuda.current_stream(self.device) + capturing = torch.cuda.is_current_stream_capturing() + if capturing and self.wait_sync: + wait_reset(stream, self.flag) + pinned, dev = (self._graph_pinned, self._graph_dev) if capturing else (self._eager_pinned[self._eager_slot], self._eager_dev) + n = num_tokens * self.n_cols + nbytes = n * self.row_bytes + dev[:nbytes].copy_(pinned[:nbytes], non_blocking=True) + if not capturing: + self._eager_read[self._eager_slot].record(stream) + values = unpack_rows(dev[:nbytes].view(n, self.row_bytes), FP8_E8M0_B32, self.head_dim) + return values.view(num_tokens, self.n_cols * self.head_dim) + + +def _context(ids: torch.Tensor, position: int, width: int) -> torch.Tensor: + """The up-to-``width`` raw ids before ``position`` (fewer at the sequence start).""" + return ids[max(0, position - width) : position] + + +class EngramHost: + """The per-model coordinator: hashes each batch's token runs and fills every layer's table. + + The token range a prefill hashes is the attention metadata's encoder segments -- the one + execution plan the scheduler, the attention layers and the model already agreed on. A graph + decode's fill runs on a worker thread from an immutable dispatch snapshot (position + the + ``max_ngram - 1`` preceding ids per row) plus the readback of the sampled token; the request + objects themselves advance on the main thread right after dispatch and are never read late. + """ + + def __init__(self, hash: EngramHash, tables: Sequence[EngramDiskTable], device: torch.device, *, sync_mode: str = "auto") -> None: + from freetoken.kernel.row_store import probe_wait_sync + + self.hash = hash + self.tables = list(tables) + self.device = device + self.wait_sync = probe_wait_sync(sync_mode, device) + for t in self.tables: + t.wait_sync = self.wait_sync + self._token_readback = alloc_pinned_tensor(min(t.max_graph_rows for t in self.tables), dtype=torch.int32) + self._readback_event = torch.cuda.Event() + self._worker = ThreadPoolExecutor(max_workers=1, thread_name_prefix="engram-fill") + self._pending: Future | None = None + logger.info_rank0(f"Engram disk backend: {self.tables[0].store.io_backend()}, {'wait-sync' if self.wait_sync else 'launch-gating'}") + + def close(self) -> None: + self._worker.shutdown(wait=True) + + def __del__(self) -> None: + worker = getattr(self, "_worker", None) + if worker is not None: + worker.shutdown(wait=False) + + def fill(self, runs: Sequence[tuple[torch.Tensor, torch.Tensor]], *, graph: bool) -> None: + """``runs``: per request ``(context ids, token ids)`` in batch order; stages every layer's rows.""" + ids = torch.cat([self.hash.row_ids(tokens, ctx) for ctx, tokens in runs]) if runs else torch.empty(0, len(self.tables), self.hash.n_cols, dtype=torch.int64) + for l, table in enumerate(self.tables): + table.stage(ids[:, l], graph=graph) + table.flush(signal=graph and self.wait_sync) + + @torch.inference_mode() + def _deferred_fill(self, bs: int, histories: list[tuple[int, ...]]) -> None: + """Worker side of a graph decode: wait for the sampled tokens, hash them against the dispatch + snapshot, stage, and release the tables' flags. Reads nothing that the main thread mutates. + Inference mode is thread-local; staging buffers were created under it by the server.""" + try: + self._readback_event.synchronize() + tokens = self._token_readback[:bs].to(torch.int64) + runs = [(torch.tensor(hist, dtype=torch.int64), tokens[i : i + 1]) for i, hist in enumerate(histories)] + self.fill(runs, graph=True) + except BaseException: + from freetoken.kernel.row_store import signal + + for t in self.tables: # unblock the stream before surfacing; the step's output is discarded + signal(t.flag) + raise + + def _await_pending(self) -> None: + """Surface the deferred fill's outcome (it depends only on GPU work already enqueued, so this + cannot wait on the host).""" + if self._pending is not None: + pending, self._pending = self._pending, None + pending.result() + + def host_fill_batch(self, batch: Batch, use_graph: bool) -> None: + """Stage this batch's rows: eagerly, or -- a graph decode under flag-sync -- on the worker + thread once the sampled tokens have been read back.""" + self._await_pending() + width = self.hash.max_ngram - 1 + if batch.is_decode: + reqs = list(batch.padded_reqs) + # the dispatch snapshot: each row's position context, taken before anything advances + histories = [tuple(_context(r.input_ids, r.device_len - 1, width).tolist()) for r in reqs] + if use_graph and self.wait_sync: + bs = batch.padded_size + self._token_readback[:bs].copy_(batch.input_ids, non_blocking=True) + self._readback_event.record(torch.cuda.current_stream(self.device)) + self._pending = self._worker.submit(self._deferred_fill, bs, histories) + return + tokens = batch.input_ids.to("cpu").to(torch.int64) + runs = [(torch.tensor(hist, dtype=torch.int64), tokens[i : i + 1]) for i, hist in enumerate(histories)] + self.fill(runs, graph=use_graph) + return + # prefill: the encoder segments are the token range every consumer of this forward executes + segments = batch.attn_metadata.segments + assert segments is not None and len(segments) == len(batch.reqs) + runs = [(_context(r.input_ids, seg.start_pos, width), r.input_ids[seg.start_pos : seg.end]) for r, seg in zip(batch.reqs, segments)] + self.fill(runs, graph=False) + + @contextmanager + def forward_host_ctx(self, batch: Batch, use_graph: bool): + self.host_fill_batch(batch, use_graph) + yield + # the launch has been enqueued: wait for the deferred fill here so a failure surfaces before + # this step's output is consumed (the fill only waits on the GPU readback, never on the host) + self._await_pending() + if use_graph: + # only after this step's fill: the next fill waits on this launch, not the fill on it + stream = torch.cuda.current_stream(self.device) + for t in self.tables: + t.graph_consumed.record(stream) + + +__all__ = ["EngramDiskTable", "EngramHost", "EngramRowSource", "engram_row_source"] diff --git a/python/freetoken/models/deepseek_v41/indexer.py b/python/freetoken/models/deepseek_v41/indexer.py new file mode 100644 index 0000000000..4ab3da2bd4 --- /dev/null +++ b/python/freetoken/models/deepseek_v41/indexer.py @@ -0,0 +1,127 @@ +"""DSV41 Lightning Indexer (reference ``Indexer``): scores the compressed positions and keeps the +``index_topk`` best per query. + +A Full-mode layer owns the index keys: ``k = k_norm(wk(latent))`` from its compressor's unrotated +latent, RoPE'd at the group position, fp4-quantized (ue8m0 per 32) into the kv source's index pool. +Every Full / Reindex layer scores its own fp4-baked queries ``wq_b(qr)`` against those keys, +``sum_h relu(q_h . k) * w_h`` with ``w = weights_proj(x) * softmax_scale * n_heads**-0.5``. + +The Hierarchical Sparse Indexer (sec. 2.3.2): the candidate-source layer (the decoder's first Full +layer) additionally keeps the ``candidate_topk_blocks`` best ``candidate_block_size``-position blocks +as a candidate pool; the decoder's Reindex layers score only that pool. Logits are produced in query +blocks bounded by ``LOGITS_BUDGET_BYTES`` so a long-context prefill never materializes ``[T, 1M]``. +""" + +from __future__ import annotations + +import torch +from freetoken.core import get_global_ctx +from freetoken.kernel.triton.dsv4.fp8_linear import fp4_act_quant_inplace +from freetoken.layers import BaseOP, LinearReplicated, RMSNorm + +from .args import DeepseekV41Args, LayerRole, Mode +from .rope import apply_rotary_emb, apply_rotary_emb_decode + +LOGITS_BUDGET_BYTES = 512 << 20 + + +class Indexer(BaseOP): + def __init__(self, args: DeepseekV41Args, role: LayerRole, *, quant_config=None, prefix: str = ""): + self.role = role + self.n_heads = args.index_n_heads + self.head_dim = args.index_head_dim + self.rope_dim = args.rope_head_dim + self.topk = args.index_topk + self.weights_scale = args.index_softmax_scale * args.index_n_heads**-0.5 + self.candidate_topk_blocks = args.candidate_topk_blocks + self.candidate_block_size = args.candidate_block_size + self.wq_b = LinearReplicated(args.q_lora_rank, self.n_heads * self.head_dim, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.wq_b") + self.weights_proj = LinearReplicated(args.dim, self.n_heads, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.weights_proj") + if role.mode is Mode.FULL: + self.wk = LinearReplicated(args.head_dim, self.head_dim, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.wk") + self.k_norm = RMSNorm(self.head_dim, args.norm_eps) + + @property + def attn(self): + return get_global_ctx().attn_backend + + # ----- keys (Full mode) ---------------------------------------------------------------- + def index_keys(self, latent: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + """``[G, index_head_dim]`` keys from the unrotated latents, rotated at their group positions. The + fp4 quantization is the pool's packing.""" + k = self.k_norm.forward(self.wk.forward(latent)) + apply_rotary_emb(k[..., -self.rope_dim :], freqs) + return k + + # ----- queries --------------------------------------------------------------------------- + def _queries(self, x: torch.Tensor, qr: torch.Tensor, rotate) -> tuple[torch.Tensor, torch.Tensor]: + q = self.wq_b.forward(qr).unflatten(-1, (self.n_heads, self.head_dim)) + rotate(q[..., -self.rope_dim :]) + fp4_act_quant_inplace(q, 32) + weights = self.weights_proj.forward(x) * self.weights_scale + return q, weights + + # ----- selection ------------------------------------------------------------------------- + def select_prefill( + self, x: torch.Tensor, qr: torch.Tensor, freqs: torch.Tensor, *, start_pos: int, ratio: int, locs: torch.Tensor, + candidates: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Per-query picks for one request's ``n`` new tokens at ``[start_pos, start_pos + n)``. + + ``locs [1, W]`` are the request's live full locs for the positions its queries may see (the + source's row of compressed position ``t`` is ``locs[t * ratio] // ratio``); ``candidates [1, n, + NC]`` restricts a Reindex layer to the pool. Returns ``(rows [1, n, topk] int32 global main rows, + -1 tail; candidates [1, n, NC] | None -- the pool this layer built as candidate source)``. + """ + n = x.shape[0] + T = locs.shape[1] // ratio + if T == 0: # a prompt shorter than one compression block: no source row exists yet + rows = torch.full((1, n, self.topk), -1, dtype=torch.int32, device=x.device) + pool = None + if self.role.is_candidate_source: + pool = torch.full((1, n, self.candidate_topk_blocks * self.candidate_block_size), -1, dtype=torch.int32, device=x.device) + return rows, pool + q, w = self._queries(x, qr, lambda t: apply_rotary_emb(t, freqs)) + live = ((start_pos + torch.arange(1, n + 1, device=x.device)) // ratio).to(torch.int32).view(1, n) + width = candidates.shape[-1] if candidates is not None else T + qb = max(1, LOGITS_BUDGET_BYTES // max(1, width * 4)) + picks, pools = [], [] + for s0 in range(0, n, qb): + s1 = min(n, s0 + qb) + cand = candidates[:, s0:s1] if candidates is not None else None + logits = self.attn.indexer_logits(q[None, s0:s1], w[None, s0:s1], self.role.kv_source, locs, ratio, live[:, s0:s1], T=T, candidates=cand) + if cand is not None: + pos = self.attn.select_topk_in_candidates(logits, cand, self.topk) + else: + pos = self.attn.select_topk(logits, live[:, s0:s1], self.topk) + if self.role.is_candidate_source: + pools.append(self.attn.select_candidate_blocks(logits, live[:, s0:s1], self.candidate_topk_blocks, self.candidate_block_size)) + picks.append(self.attn.positions_to_rows(pos, locs, ratio)) + rows = torch.cat(picks, dim=1) + pool = torch.cat(pools, dim=1) if pools else None + return rows.to(torch.int32), pool + + def select_decode( + self, x: torch.Tensor, qr: torch.Tensor, freqs: torch.Tensor, *, pos: torch.Tensor, ratio: int, locs: torch.Tensor, + T: int, candidates: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Batched one-token selection. ``locs [B, W]`` is the decode snapshot (graph-static width), + ``T`` the staged compressed width, ``pos [B]``. Returns ``(rows [B, 1, topk] int32, counts [B, 1] + int32 valid picks, candidates [B, 1, NC] | None)``.""" + B = x.shape[0] + q, w = self._queries(x, qr, lambda t: apply_rotary_emb_decode(t.unsqueeze(1), freqs)) + live = ((pos + 1) // ratio).to(torch.int32).view(B, 1) + logits = self.attn.indexer_logits(q.view(B, 1, self.n_heads, self.head_dim), w.view(B, 1, -1), self.role.kv_source, locs, ratio, live, T=T, candidates=candidates) + pool = None + if candidates is not None: + picks = self.attn.select_topk_in_candidates(logits, candidates, self.topk) + else: + picks = self.attn.select_topk(logits, live, self.topk) + if self.role.is_candidate_source: + pool = self.attn.select_candidate_blocks(logits, live, self.candidate_topk_blocks, self.candidate_block_size) + rows = self.attn.positions_to_rows(picks, locs, ratio) + counts = (picks >= 0).sum(dim=-1, dtype=torch.int32) # valid picks form a prefix + return rows.to(torch.int32), counts, pool + + +__all__ = ["Indexer", "LOGITS_BUDGET_BYTES"] diff --git a/python/freetoken/models/deepseek_v41/model.py b/python/freetoken/models/deepseek_v41/model.py new file mode 100644 index 0000000000..8112022f76 --- /dev/null +++ b/python/freetoken/models/deepseek_v41/model.py @@ -0,0 +1,289 @@ +"""DeepSeek-V4.1-Flash model (engine-native port of ``inference/model.py``). + + embed -> hc_mult residual streams -> [Engram at 1, 14] -> 40 blocks (single-pass mHC around + DSV41 attention and MoE) -> collapse with the last block's pre gates -> norm -> head + +Single-pass mHC (``layers/mhc.py``): each sublayer boundary applies the PREVIOUS sublayer's post / +comb to the streams, predicts this sublayer's (pre, post, comb) from the updated streams, and mixes +the sublayer input with the pre gates the previous sublayer predicted. The chain state +``(y, residual, post, comb, pre)`` threads through the layer loop; layer 0's attention starts it +with a one-hot pre on stream 0 and no pending post. + +Causal Encoder-Decoder: the decoder's kv source (layer 20) compresses ITS INPUT -- the final encoder +hidden state -- into the global KV every decoder layer reads. Under Decoder SWA Bounded Replay the +decoder layers run on each request's last ``window_size`` prompt tokens only (their sliding window +truncated at the start of that window, see the backend; a prefix hit stops before it, so those tokens +are always prefilled); layer 20 still publishes global KV for every prompt token. +``--swa-decoder-replay exact`` runs the decoder on every token (the reference numerics). + +KV addressing is the attention backend's; pool buffers are read off the live pool per access, so a +runtime cache rebuild needs no unbind. Decode is batched and CUDA-graph safe (the DSV4 precedent). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, List, NamedTuple + +import torch +from freetoken.attention.dsv41_sparse import DSV41AttnMetadata, PrefillSegment +from freetoken.core import get_global_ctx +from freetoken.layers import BaseOP, OPList, ParallelLMHead, RMSNorm, VocabParallelEmbedding +from freetoken.layers.mhc import mhc_fused_post_pre_single_pass, mhc_mix_input, mhc_post +from freetoken.models.blocks import BaseLLMModel, embed_input_ids + +from .args import DeepseekV41Args +from .attention import DSV41Attention, DecodeStepContext +from .engram import EngramLayer +from .moe import MoE + +if TYPE_CHECKING: + from freetoken.models.config import ModelConfig + +POST_MULT = 2.0 # post gates are 2 * sigmoid (reference hc_split_sinkhorn) + + +class Streams(NamedTuple): + """The single-pass mHC chain state between sublayers: the last sublayer output ``y`` with its + not-yet-applied ``post`` / ``comb``, the residual streams ``[T, hc, dim]`` they apply to, and + the ``pre`` gates ``[T, hc]`` the next sublayer mixes its input with.""" + + y: torch.Tensor | None + residual: torch.Tensor + post: torch.Tensor | None + comb: torch.Tensor | None + pre: torch.Tensor + + def materialize(self) -> torch.Tensor: + """The streams with the pending post applied.""" + if self.post is None: + return self.residual + return mhc_post(self.y, self.residual, self.post, self.comb) + + +class Block(BaseOP): + def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""): + args: DeepseekV41Args = config.dsv41_args + self.layer_id = layer_id + self.role = args.roles[layer_id] + self.norm_eps = args.norm_eps + self.hc_eps = args.hc_eps + self.sinkhorn = args.hc_sinkhorn_iters + hc, dim = args.hc_mult, args.dim + mix = (2 + hc) * hc + self.attn = DSV41Attention(args, self.role, quant_config=config.quant, prefix=f"{prefix}.attn") + self.ffn = MoE(config, layer_id, prefix=f"{prefix}.ffn") + self.attn_norm = RMSNorm(dim, args.norm_eps) + self.ffn_norm = RMSNorm(dim, args.norm_eps) + self.engram = EngramLayer(args, layer_id, quant_config=config.quant, prefix=f"{prefix}.engram") if layer_id in args.engram_layer_ids else None + self.hc_attn_fn = torch.empty(mix, hc * dim, dtype=torch.float32) + self.hc_ffn_fn = torch.empty(mix, hc * dim, dtype=torch.float32) + self.hc_attn_base = torch.empty(mix, dtype=torch.float32) + self.hc_ffn_base = torch.empty(mix, dtype=torch.float32) + self.hc_attn_scale = torch.empty(3, dtype=torch.float32) + self.hc_ffn_scale = torch.empty(3, dtype=torch.float32) + + def _boundary(self, s: Streams, fn, scale, base): + residual, post, comb, pre_next, x = mhc_fused_post_pre_single_pass( + s.y if s.y is not None else s.residual[:, 0], s.residual, s.post, s.comb, s.pre, fn, scale, base, + self.norm_eps, self.hc_eps, POST_MULT, self.sinkhorn, + ) + return residual, post, comb, pre_next, x + + def attn_input(self, s: Streams, *, image_mask: torch.Tensor | None = None): + """Apply the pending post, predict the attention gates, mix and norm the attention input. + Returns ``(residual, post, comb, pre_next, x)``.""" + if self.engram is not None: + s = Streams(None, self.engram.forward(s.materialize(), image_mask=image_mask), None, None, s.pre) + residual, post, comb, pre_next, x = self._boundary(s, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base) + return residual, post, comb, pre_next, self.attn_norm.forward(x) + + def ffn_step(self, y: torch.Tensor, residual, post, comb, pre, *, image_mask: torch.Tensor | None = None) -> Streams: + """The FFN sublayer after the attention output ``y``: boundary, norm, MoE; returns the chain state.""" + residual, post, comb, pre_next, x = self._boundary(Streams(y, residual, post, comb, pre), self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base) + return Streams(self.ffn.forward(self.ffn_norm.forward(x), image_mask=image_mask), residual, post, comb, pre_next) + + def forward_prefill(self, s: Streams, segments: List[PrefillSegment], positions: torch.Tensor, image_mask: torch.Tensor | None = None) -> Streams: + residual, post, comb, pre_next, x = self.attn_input(s, image_mask=image_mask) + if self.attn.compressor is not None: + self.attn.publish_prefill(x, segments) + y = self.attn.forward_prefill(x, segments, positions) + return self.ffn_step(y, residual, post, comb, pre_next, image_mask=image_mask) + + def forward_prefill_narrowed( + self, s: Streams, segments: List[PrefillSegment], rows: torch.Tensor, query_segments: List[PrefillSegment], + query_positions: torch.Tensor, image_mask: torch.Tensor | None = None, query_image_mask: torch.Tensor | None = None, + ) -> Streams | None: + """``forward_prefill`` with a row cut after the global KV is published: the attention input and + ``publish_prefill`` cover every token of ``segments``, the attention and the FFN only ``rows`` + (tiled by ``query_segments``) -- the CED decoder's kv source under bounded replay. None when no + query segment is left (every chunk in the batch ends before its prompt's last window).""" + residual, post, comb, pre_next, x = self.attn_input(s, image_mask=image_mask) + if self.attn.compressor is not None: + self.attn.publish_prefill(x, segments) + if not query_segments: + return None + y = self.attn.forward_prefill(x[rows], query_segments, query_positions) + return self.ffn_step(y, residual[rows], post[rows], comb[rows], pre_next[rows], image_mask=query_image_mask) + + def forward_decode(self, s: Streams, pos: torch.Tensor, rows: torch.Tensor, dctx, cmp_stage_cap: int) -> Streams: + residual, post, comb, pre_next, x = self.attn_input(s) + y = self.attn.forward_decode(x, pos, rows, dctx, cmp_stage_cap) + return self.ffn_step(y, residual, post, comb, pre_next) + + +class Transformer(BaseOP): + def __init__(self, config: ModelConfig, *, prefix: str = "model"): + args: DeepseekV41Args = config.dsv41_args + self.args = args + self.embed = VocabParallelEmbedding(num_embeddings=args.vocab_size, embedding_dim=args.dim) + self.layers = OPList([Block(config, i, prefix=f"{prefix}.layers.{i}") for i in range(args.n_layers)]) + self.norm = RMSNorm(args.dim, args.norm_eps) + + def bind(self, device: torch.device) -> None: + tables: dict = {} # one RoPE table per distinct parameter set, shared by the layers using it + for block in self.layers.op_list: + block.attn.bind(device, tables) + + def _entry(self, input_ids: torch.Tensor) -> Streams: + h = embed_input_ids(self.embed, input_ids, get_global_ctx().batch) + residual = h.unsqueeze(1).expand(-1, self.args.hc_mult, -1).contiguous() + pre = torch.zeros(h.shape[0], self.args.hc_mult, dtype=torch.float32, device=h.device) + pre[:, 0] = 1.0 + return Streams(None, residual, None, None, pre) + + def _exit(self, s: Streams) -> torch.Tensor: + return self.norm.forward(mhc_mix_input(s.materialize(), s.pre)) + + def prefill(self, input_ids: torch.Tensor, positions: torch.Tensor, md: DSV41AttnMetadata) -> torch.Tensor: + """Ragged prefill over the flat new-token stream; returns the hidden states the head reads (the + decoder stream under bounded replay, where ``md.last_indices`` also point, plus the placeholder + row ``md.decoder_pad`` asks for).""" + s = self._entry(input_ids) + dtype = s.residual.dtype + image_mask = input_ids >= self.args.vocab_size if get_global_ctx().batch.mm_embeds is not None else None + segments, dec_segments, rows = md.segments, md.decoder_segments, md.decoder_rows + decoder_start = self.args.decoder_start_layer + if rows is not None: + dec_positions = positions[rows] + dec_image_mask = None if image_mask is None else image_mask[rows] + for block in self.layers.op_list: + if rows is None or block.layer_id < decoder_start: + s = block.forward_prefill(s, segments, positions, image_mask) + elif block.layer_id == decoder_start: + s = block.forward_prefill_narrowed(s, segments, rows, dec_segments, dec_positions, image_mask, dec_image_mask) + if s is None: + break + else: + s = block.forward_prefill(s, dec_segments, dec_positions, dec_image_mask) + pad = [torch.zeros(1, self.args.dim, dtype=dtype, device=input_ids.device)] if md.decoder_pad else [] + return torch.cat([self._exit(s)] + pad) if s is not None else pad[0] + + def decode(self, input_ids: torch.Tensor, pos: torch.Tensor, md: DSV41AttnMetadata, cmp_stage_cap: int) -> torch.Tensor: + B = input_ids.shape[0] + rows = torch.arange(B, device=input_ids.device) + s = self._entry(input_ids) + # the layer-invariant step context, resolved once (a capture records it once): ring slots and + # window candidates of the shared pool -- and of the private rings when bounded replay keeps + # the decoder's window KV per request -- plus each RoPE table at the step's positions + dctx = DecodeStepContext.build([b.attn for b in self.layers.op_list], md, pos, rows, self.args.window_size, self.args.index_topk) + for block in self.layers.op_list: + s = block.forward_decode(s, pos, rows, dctx, cmp_stage_cap) + return self._exit(s) + + +class DeepseekV41ForCausalLM(BaseLLMModel): + def __init__(self, config: ModelConfig): + self._config = config + self._args: DeepseekV41Args = config.dsv41_args + self.model = Transformer(config) + self.head = ParallelLMHead(num_embeddings=self._args.vocab_size, embedding_dim=self._args.dim, quant_config=config.quant, prefix="head") + self._bound = False + if config.is_multimodal: + from .vision import DeepseekV41Vision + + self.visual = DeepseekV41Vision(config.vision_config, config.hidden_size, + num_layers=self._args.n_layers, num_experts=self._args.n_routed_experts) + + def encode(self, item): + return self.visual.forward(item) + + def place_encoder_weights(self, mode: str) -> None: + self.visual.place_weights(mode) + + def _ensure_bound(self) -> None: + if not self._bound: + self.model.bind(get_global_ctx().kv_cache.device) + if self._config.is_multimodal: + for i, block in enumerate(self.model.layers.op_list): + block.ffn.gate._bias_vl = self.visual.routing_bias[i] + self._bound = True + + def mark_for_rebind(self) -> None: + """A runtime pool rebuild changes nothing the model caches (rope tables are pool-independent), + but the engine asks; re-bind is cheap and keeps the contract with the DSV4 sibling.""" + self._bound = False + + def engram_layers(self) -> List[EngramLayer]: + return [b.engram for b in self.model.layers.op_list if b.engram is not None] + + def load_host_tables(self, engine_config) -> int: + """Attach the Engram tables (streamed from the checkpoint shards; zeros for dummy weights) + and take over ``forward_host_ctx`` so every dispatch stages its rows. Returns the pinned + host bytes the engine reserves from its pin budget.""" + import os + + layers = self.engram_layers() + if not layers: + return 0 + device = torch.device("cuda", torch.cuda.current_device()) + if getattr(engine_config, "use_dummy_weight", False): + from .engram import ZeroEngramTable + + for layer in layers: + layer.attach_table(ZeroEngramTable(layer.width, device)) + return 0 + from freetoken.utils import download_hf_weight + from transformers import AutoTokenizer + + from .engram import EngramHash, build_compressed_token_map + from .engram_table import EngramDiskTable, EngramHost, engram_row_source + + folder = download_hf_weight(engine_config.model_path) + token_map, compressed_vocab = build_compressed_token_map(AutoTokenizer.from_pretrained(folder, trust_remote_code=True)) + hash = EngramHash(self._args, token_map, compressed_vocab) + graph_rows = max(engine_config.max_running_req, engine_config.cuda_graph_max_bs or 0, 1) + tables = [ + EngramDiskTable( + engram_row_source(folder, layer.layer_id), hash.n_cols, device, + max_graph_rows=graph_rows, + max_extend_tokens=engine_config.max_extend_tokens, + use_io_uring=os.getenv("FREETOKEN_ENGRAM_IO_URING", "1") != "0", + ) + for layer in layers + ] + for layer, table in zip(layers, tables): + layer.attach_table(table) + self._engram_host = EngramHost(hash, tables, device, sync_mode=os.getenv("FREETOKEN_ENGRAM_SYNC", "auto")) + self.forward_host_ctx = self._engram_host.forward_host_ctx + return sum(t._graph_pinned.numel() + sum(b.numel() for b in t._eager_pinned) for t in tables) + + def forward(self) -> torch.Tensor: + self._ensure_bound() + batch = get_global_ctx().batch + md = batch.attn_metadata + assert isinstance(md, DSV41AttnMetadata) + if batch.is_prefill: + hidden = self.model.prefill(batch.input_ids, batch.positions.long(), md) + else: + input_ids = batch.input_ids + B = batch.padded_size + pos = batch.positions.long().view(-1)[:B] + if torch.cuda.is_current_stream_capturing(): + cmp_stage_cap = md.stage_width - 1 + else: + cmp_stage_cap = int(pos.max().item()) + hidden = self.model.decode(input_ids.view(B), pos, md, cmp_stage_cap) + return self.head.forward(hidden) + + +__all__ = ["DeepseekV41ForCausalLM", "Transformer", "Block", "Streams"] diff --git a/python/freetoken/models/deepseek_v41/moe.py b/python/freetoken/models/deepseek_v41/moe.py new file mode 100644 index 0000000000..91c62ecb01 --- /dev/null +++ b/python/freetoken/models/deepseek_v41/moe.py @@ -0,0 +1,90 @@ +"""DeepSeek-V4.1 MoE block: sqrtsoftplus router with a selection-only correction bias, one shared +clamped-SwiGLU expert, and the routed MXFP4 experts on the shared offload cache. + +Router (reference ``Gate``): ``scores = sqrt(softplus(x @ W^T))`` in fp32; the bias picks the experts, +the unbiased scores weight them (renormalized, ``+1e-20`` as in training), times ``route_scale``. +Image-span tokens select experts using the checkpoint's separate ``bias_vl``. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Tuple + +import torch +from freetoken.kernel.triton.dsv4.bf16_linear import bf16_linear_fp32 +from freetoken.kernel.triton.dsv4.router import router_select +from freetoken.kernel.triton.dsv4.swiglu import fused_swiglu +from freetoken.layers import BaseOP, LinearReplicated, make_moe_layer + +from .args import DeepseekV41Args + +if TYPE_CHECKING: + from freetoken.models.config import ModelConfig + +TopK = Tuple[torch.Tensor, torch.Tensor] + + +class Router(BaseOP): + def __init__(self, args: DeepseekV41Args): + self.topk = args.n_activated_experts + self.score_func = args.score_func + self.norm_topk_prob = args.norm_topk_prob + self.route_scale = args.route_scale + self.weight = torch.empty(args.n_routed_experts, args.dim, dtype=torch.bfloat16) + self.bias = torch.empty(args.n_routed_experts, dtype=torch.float32) + self._bias_vl = None # bound to the visual stack's per-layer routing bias after loading + + def forward(self, x: torch.Tensor, *, image_mask: torch.Tensor | None = None) -> TopK: + if image_mask is not None: + assert self._bias_vl is not None, "image routing bias was not loaded" + scores = bf16_linear_fp32(x, self.weight) + return router_select( + scores.view(-1, scores.shape[-1]), self.bias, self.topk, + score_func=self.score_func, normalize=self.norm_topk_prob, route_scale=self.route_scale, + alternate_bias=self._bias_vl, alternate_mask=image_mask, + ) + + +class SharedExpert(BaseOP): + """The dense clamped-SwiGLU expert every token passes through (fp8 block-32 linears).""" + + def __init__(self, args: DeepseekV41Args, *, quant_config=None, prefix: str = ""): + self.w1 = LinearReplicated(args.dim, args.moe_inter_dim, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.w1") + self.w2 = LinearReplicated(args.moe_inter_dim, args.dim, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.w2") + self.w3 = LinearReplicated(args.dim, args.moe_inter_dim, has_bias=False, quant_config=quant_config, prefix=f"{prefix}.w3") + self.limit = args.swiglu_limit + + def forward(self, x: torch.Tensor) -> torch.Tensor: + h = fused_swiglu(self.w1.forward(x), self.w3.forward(x), self.limit, x.dtype) + return self.w2.forward(h) + + +class MoE(BaseOP): + def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""): + args: DeepseekV41Args = config.dsv41_args + self.dim = args.dim + self.gate = Router(args) + self.shared_experts = SharedExpert(args, quant_config=config.quant, prefix=f"{prefix}.shared_experts") + # every layer is MoE: the offload cache indexes experts by layer id directly + self.experts = make_moe_layer( + config, + layer_id=layer_id, + activation="swiglu_clamp", + renormalize=args.norm_topk_prob, + limit=args.swiglu_limit, + quant_config=config.quant, + prefix=f"{prefix}.experts", + ) + + def forward(self, x: torch.Tensor, *, image_mask: torch.Tensor | None = None) -> torch.Tensor: + shape = x.shape + x = x.view(-1, self.dim) + weights, indices = self.gate.forward(x, image_mask=image_mask) + # the shared expert runs before the routed experts (same stream, so it does not overlap the + # hybrid decode's CPU work, whose submit follows it -- reordering that is follow-up work) + shared = self.shared_experts.forward(x) + routed = self.experts.routed_forward(x, weights, indices) + return (routed + shared).view(shape) + + +__all__ = ["MoE", "Router", "SharedExpert"] diff --git a/python/freetoken/models/deepseek_v41/rope.py b/python/freetoken/models/deepseek_v41/rope.py new file mode 100644 index 0000000000..6d7df9f8d3 --- /dev/null +++ b/python/freetoken/models/deepseek_v41/rope.py @@ -0,0 +1,58 @@ +"""Rotary embeddings for DeepSeek-V4.1 (verbatim reference math, ``inference/model.py``). + +Interleaved complex rotation on the last ``rope_head_dim`` channels, with YaRN when +``original_seq_len > 0``. Two regimes share two cached tables: window-only layers at the base +theta without YaRN, compressing layers at ``compress_rope_theta`` under YaRN. A compressed latent +stands for the first token of its group, so group ``j`` rotates at position ``j * ratio``; the +attention output is rotated back (``inverse=True``) because K == V share one rotated latent. +""" + +from __future__ import annotations + +import math +from functools import lru_cache + +import torch + + +def precompute_freqs_cis(dim: int, seqlen: int, original_seq_len: int, base: float, factor: float, beta_fast: int, beta_slow: int) -> torch.Tensor: + freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + if original_seq_len > 0: + + def corrected_dim(rotations): + return dim * math.log(original_seq_len / (rotations * 2 * math.pi)) / (2 * math.log(base)) + + low = max(math.floor(corrected_dim(beta_fast)), 0) + high = min(math.ceil(corrected_dim(beta_slow)), dim - 1) + ramp = ((torch.arange(dim // 2, dtype=torch.float32) - low) / max(high - low, 1e-3)).clamp(0, 1) + smooth = 1 - ramp + freqs = freqs / factor * (1 - smooth) + freqs * smooth + freqs = torch.outer(torch.arange(seqlen), freqs) + return torch.polar(torch.ones_like(freqs), freqs) + + +@lru_cache(2) # the model's two rope regimes; a 1M table is ~256 MB, so never one per layer +def get_freqs_cis(dim: int, seqlen: int, original_seq_len: int, base: float, factor: float, beta_fast: int, beta_slow: int, device: torch.device) -> torch.Tensor: + return precompute_freqs_cis(dim, seqlen, original_seq_len, base, factor, beta_fast, beta_slow).to(device) + + +def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False) -> torch.Tensor: + """In-place rotation of ``x[..., :]`` (``[T, rd]`` or ``[T, H, rd]``) by per-token ``freqs_cis [T, rd//2]``.""" + xc = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) + if inverse: + freqs_cis = freqs_cis.conj() + if xc.ndim == 3: + freqs_cis = freqs_cis.unsqueeze(1) + x.copy_(torch.view_as_real(xc * freqs_cis).flatten(-2)) + return x + + +def apply_rotary_emb_decode(x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False) -> torch.Tensor: + """Batched decode rotation with PER-ROW freqs: ``x [B, 1, ..., rd]``, ``freqs_cis [B, rd//2]`` + (the fused dsv4 triton kernel; bit-identical to ``apply_rotary_emb``).""" + from freetoken.kernel.triton.dsv4.rope import rope_decode_inplace + + return rope_decode_inplace(x, freqs_cis, inverse) + + +__all__ = ["precompute_freqs_cis", "get_freqs_cis", "apply_rotary_emb", "apply_rotary_emb_decode"] diff --git a/python/freetoken/models/deepseek_v41/vision.py b/python/freetoken/models/deepseek_v41/vision.py new file mode 100644 index 0000000000..f5fea969ea --- /dev/null +++ b/python/freetoken/models/deepseek_v41/vision.py @@ -0,0 +1,139 @@ +"""DeepSeek ViT and 3x3 pixel-unshuffle aligner, with the shared encoder weight streamer.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F +from freetoken.layers import BaseOP, LinearReplicated, OPList +from freetoken.models.weight_stream import BlockWeightStreamer + + +class VisionNorm(BaseOP): + def __init__(self, dim: int): + self.weight = torch.empty(dim, dtype=torch.float32) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = x.float() + return (y * torch.rsqrt(y.square().mean(-1, keepdim=True) + 1e-6) * self.weight).to(x.dtype) + + +class VisionAttention(BaseOP): + def __init__(self, vc): + self.n_heads = vc.num_attention_heads + self.head_dim = vc.hidden_size // self.n_heads + self.wqkv = LinearReplicated(vc.hidden_size, 3 * vc.hidden_size, has_bias=True) + self.wo = LinearReplicated(vc.hidden_size, vc.hidden_size, has_bias=True) + + def forward(self, x, cos, sin): + n = x.shape[0] + q, k, v = (t.reshape(n, self.n_heads, self.head_dim) for t in self.wqkv.forward(x).chunk(3, -1)) + def rotary(t): + a, b = t.float().chunk(2, -1) + return torch.cat((a * cos - b * sin, b * cos + a * sin), -1).to(t.dtype) + q, k = rotary(q), rotary(k) + q, k, v = (t.transpose(0, 1).unsqueeze(0) for t in (q, k, v)) + y = F.scaled_dot_product_attention(q, k, v) + return self.wo.forward(y[0].transpose(0, 1).reshape(n, -1)) + + +class VisionMLP(BaseOP): + def __init__(self, vc): + self.w1 = LinearReplicated(vc.hidden_size, 2 * vc.intermediate_size, has_bias=False) + self.w2 = LinearReplicated(vc.intermediate_size, vc.hidden_size, has_bias=False) + + def forward(self, x): + gate, up = self.w1.forward(x).chunk(2, -1) + return self.w2.forward(F.silu(gate) * up) + + +class VisionBlock(BaseOP): + def __init__(self, vc): + self.norm1 = VisionNorm(vc.hidden_size) + self.attn = VisionAttention(vc) + self.norm2 = VisionNorm(vc.hidden_size) + self.mlp = VisionMLP(vc) + + def forward(self, x, cos, sin): + x = x + self.attn.forward(self.norm1.forward(x), cos, sin) + return x + self.mlp.forward(self.norm2.forward(x)) + + +class PatchEmbed(BaseOP): + def __init__(self, vc): + self.proj = LinearReplicated(3 * vc.patch_size**2, vc.hidden_size, has_bias=True) + + def forward(self, x): + return self.proj.forward(x.flatten(1)) + + +class VisionTransformer(BaseOP): + def __init__(self, vc): + self.patch_embed = PatchEmbed(vc) + self.blocks = OPList([VisionBlock(vc) for _ in range(vc.num_hidden_layers)]) + self.norm = VisionNorm(vc.hidden_size) + self._vc = vc + self._streamer = None + + def place_weights(self, mode: str) -> None: + if mode == "host" and self._streamer is None: + self._streamer = BlockWeightStreamer(self.blocks.op_list, self.patch_embed.proj.weight.device) + elif mode == "gpu" and self._streamer is not None: + self._streamer.unstream() + self._streamer = None + elif mode not in ("gpu", "host"): + raise ValueError(f"unknown vision weight placement {mode!r}") + + def forward(self, patches, h: int, w: int): + x = self.patch_embed.forward(patches) + dim = self._vc.hidden_size // self._vc.num_attention_heads // 2 + inv = 1.0 / (self._vc.rope_theta ** (torch.arange(0, dim, 2, dtype=torch.float32, device=x.device) / dim)) + rows = torch.arange(h, device=x.device)[:, None].expand(h, w) + cols = torch.arange(w, device=x.device)[None, :].expand(h, w) + angles = (torch.stack((rows, cols), -1).reshape(-1, 2, 1).float() * inv).flatten(1) + cos, sin = angles.cos().unsqueeze(1), angles.sin().unsqueeze(1) + blocks = enumerate(self.blocks.op_list) if self._streamer is None else self._streamer.blocks(self.blocks.op_list) + for _, block in blocks: + x = block.forward(x, cos, sin) + return self.norm.forward(x) + + +class VisionAligner(BaseOP): + def __init__(self, vc, dim: int): + self.ratio = vc.downsample_ratio + self.w1 = LinearReplicated(vc.hidden_size * self.ratio**2, dim, has_bias=True) + self.w2 = LinearReplicated(dim, dim, has_bias=True) + + def forward(self, x, h: int, w: int): + r = self.ratio + x = x.view(h, w, -1).permute(2, 0, 1) + x = F.pad(x, (0, -w % r, 0, -h % r)) + x = F.unfold(x.unsqueeze(0), r, stride=r).squeeze(0).T + return self.w2.forward(F.gelu(self.w1.forward(x))) + + +class DeepseekV41Vision(BaseOP): + def __init__(self, vc, dim: int, *, num_layers: int, num_experts: int): + self.vision = VisionTransformer(vc) + self.aligner = VisionAligner(vc, dim) + self.image_start = torch.empty(dim) + self.image_end = torch.empty(dim) + self.image_newline = torch.empty(dim) + # Keep all image-only parameters under visual, including FTW/text-only filtering. + self.routing_bias = torch.empty(num_layers, num_experts, dtype=torch.float32) + + def place_weights(self, mode: str) -> None: + self.vision.place_weights(mode) + + @torch.inference_mode() + def forward(self, item): + _, h, w = item.grid_thw + patches = item.feature.to(device=self.image_start.device, dtype=self.image_start.dtype) + x = self.aligner.forward(self.vision.forward(patches, h, w), h, w) + r = self.aligner.ratio + nh, nw = (h + r - 1) // r, (w + r - 1) // r + out = x.new_empty(nh * (nw + 1) + 2, x.shape[-1]) + out[0], out[-1] = self.image_start, self.image_end + rows = out[1:-1].view(nh, nw + 1, -1) + rows[:, :nw] = x.view(nh, nw, -1) + rows[:, nw] = self.image_newline + return out diff --git a/python/freetoken/models/deepseek_v41/weight.py b/python/freetoken/models/deepseek_v41/weight.py new file mode 100644 index 0000000000..1ccd23eaef --- /dev/null +++ b/python/freetoken/models/deepseek_v41/weight.py @@ -0,0 +1,242 @@ +"""Weight loading for DeepSeek-V4.1-Flash. + +* :func:`iter_weights` streams the resident (non-expert) tensors keyed by the model's attribute + paths: checkpoint ``layers.N.*`` -> ``model.layers.N.*``, ``embed`` / ``norm`` -> ``model.*``, + ``head`` -> ``head``. fp8 linears keep their e4m3 payload and e8m0 ``scale`` (declared by the quant + method as ``weight_scale_inv``); ``wo_a`` keeps its e4m3 payload and ``wo_a_scale`` too (the layer + dequantizes in-kernel for the reference's bf16 grouped einsum). The engine casts everything else to + the declared parameter dtype (fp32 compressor weights, hc gates, engram gate weights). +* :func:`iter_expert_pieces` streams the routed MXFP4 experts (e2m1 pairs + e8m0 per-32 scales) per + expert for the offload banks. + +Dropped: the DSpark ``mtp.*`` stage; vision weights and ``gate.bias_vl`` when serving text-only. +The Engram tables (``layers.{1,14}.engram.embed.*``) are +not state-dict tensors: the disk row store maps them in place (``engram_table.py``); an FTW conversion +carries them as side files (:func:`ftw_side_files`). +""" + +from __future__ import annotations + +from contextlib import ExitStack +import json +import os +import re + +import safetensors +import torch +from freetoken.distributed import get_tp_info +from freetoken.layers.quantization import QuantKind +from freetoken.models.loader import drop_page_cache +from tqdm import tqdm + +from .args import DeepseekV41Args, Mode + + +class _ShardReader: + def __init__(self, folder: str, weight_map: dict, device): + self._folder = folder + self._weight_map = weight_map + self._device = str(device) + self._handles: dict[str, object] = {} + self._stack = ExitStack() + + def has(self, name: str) -> bool: + return name in self._weight_map + + def get(self, name: str) -> torch.Tensor: + shard = self._weight_map[name] + handle = self._handles.get(shard) + if handle is None: + handle = self._stack.enter_context(safetensors.safe_open(os.path.join(self._folder, shard), framework="pt", device=self._device)) + self._handles[shard] = handle + return handle.get_tensor(name) + + def close(self) -> None: + try: + self._stack.close() + finally: + for shard in self._handles: + drop_page_cache(os.path.join(self._folder, shard)) + self._handles.clear() + + +def weight_map(model_path: str) -> dict: + with open(os.path.join(model_path, "model.safetensors.index.json")) as f: + return json.load(f)["weight_map"] + + +def load_args(model_path: str) -> DeepseekV41Args: + from freetoken.utils.hf import cached_load_hf_config + + return DeepseekV41Args.from_hf(cached_load_hf_config(model_path)) + + +def dequant_fp8_block(weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + """Square block-scaled e4m3 -> bf16; ``scale`` holds e8m0 codes (``2^(code-127)``), the block + edge follows from the shapes.""" + n, k = weight.shape + block = n // scale.shape[0] + assert (n // block, k // block) == tuple(scale.shape), (weight.shape, scale.shape) + s = torch.exp2(scale.view(torch.uint8).to(torch.float32) - 127.0) + s = s.repeat_interleave(block, dim=0).repeat_interleave(block, dim=1) + return (weight.to(torch.float32) * s).to(torch.bfloat16) + + +def iter_weights(model_path: str, device, *, include_moe_experts: bool = True, include_non_moe: bool = True, include_vision: bool = True): + args = load_args(model_path) + reader = _ShardReader(model_path, weight_map(model_path), device) + get = reader.get + def linear(src: str, dst: str): + yield f"{dst}.weight", get(f"{src}.weight") + if reader.has(f"{src}.scale"): + yield f"{dst}.weight_scale_inv", get(f"{src}.scale") + + try: + if include_moe_experts: + # Resident experts support unquantized checkpoints; FP4 uses the offload cache. + if reader.has("layers.0.ffn.experts.0.w1.scale"): + raise ValueError( + "DeepSeek-V4.1's fp4 routed experts are served from the offload cache; run with " + "--moe-strategy offload (include_moe_experts must be False)." + ) + for L in range(args.n_layers): + p = f"layers.{L}.ffn.experts" + gate_up = torch.stack([torch.cat([get(f"{p}.{e}.w1.weight"), get(f"{p}.{e}.w3.weight")]) for e in range(args.n_routed_experts)]) + down = torch.stack([get(f"{p}.{e}.w2.weight") for e in range(args.n_routed_experts)]) + yield f"model.layers.{L}.ffn.experts.gate_up_proj", gate_up + yield f"model.layers.{L}.ffn.experts.down_proj", down + if not include_non_moe: + return + if include_vision and reader.has("vision.patch_embed.proj.weight"): + yield from iter_vision_weights(model_path, device) + yield "model.embed.weight", get("embed.weight") + yield "model.norm.weight", get("norm.weight") + yield "head.weight", get("head.weight") + for L in range(args.n_layers): + role = args.roles[L] + a, m = f"layers.{L}.attn", f"model.layers.{L}.attn" + yield from linear(f"{a}.wq_a", f"{m}.wq_a") + yield f"{m}.q_norm.weight", get(f"{a}.q_norm.weight") + yield from linear(f"{a}.wq_b", f"{m}.wq_b") + yield from linear(f"{a}.wkv", f"{m}.wkv") + yield f"{m}.kv_norm.weight", get(f"{a}.kv_norm.weight") + yield f"{m}.wo_a", get(f"{a}.wo_a.weight") + if reader.has(f"{a}.wo_a.scale"): # the fp8 payload + scale; the layer dequantizes in-kernel + yield f"{m}.wo_a_scale", get(f"{a}.wo_a.scale") + yield from linear(f"{a}.wo_b", f"{m}.wo_b") + yield f"{m}.attn_sink", get(f"{a}.attn_sink") + if role.mode is Mode.FULL: + c = f"{a}.compressor" + yield f"{m}.compressor.norm.weight", get(f"{c}.norm.weight") + if role.ratio > 1: + yield f"{m}.compressor.wkv", get(f"{c}.wkv.weight") + yield f"{m}.compressor.wgate", get(f"{c}.wgate.weight") + else: + yield from linear(f"{c}.wkv", f"{m}.compressor.wkv") + yield from linear(f"{a}.indexer.wk", f"{m}.indexer.wk") + yield f"{m}.indexer.k_norm.weight", get(f"{a}.indexer.k_norm.weight") + if role.has_indexer: + yield from linear(f"{a}.indexer.wq_b", f"{m}.indexer.wq_b") + yield from linear(f"{a}.indexer.weights_proj", f"{m}.indexer.weights_proj") + + p = f"layers.{L}" + yield f"model.{p}.attn_norm.weight", get(f"{p}.attn_norm.weight") + yield f"model.{p}.ffn_norm.weight", get(f"{p}.ffn_norm.weight") + yield f"model.{p}.ffn.gate.weight", get(f"{p}.ffn.gate.weight") + yield f"model.{p}.ffn.gate.bias", get(f"{p}.ffn.gate.bias") + for proj in ("w1", "w2", "w3"): + yield from linear(f"{p}.ffn.shared_experts.{proj}", f"model.{p}.ffn.shared_experts.{proj}") + for nm in ("hc_attn_fn", "hc_ffn_fn", "hc_attn_base", "hc_ffn_base", "hc_attn_scale", "hc_ffn_scale"): + yield f"model.{p}.{nm}", get(f"{p}.{nm}") + if L in args.engram_layer_ids: + e = f"{p}.engram" + yield from linear(f"{e}.wkv", f"model.{e}.wkv") + yield f"model.{e}.q_weight", get(f"{e}.q_weight") + yield f"model.{e}.k_weight", get(f"{e}.k_weight") + finally: + reader.close() + + +def iter_vision_weights(model_path: str, device): + names = weight_map(model_path) + reader = _ShardReader(model_path, names, device) + try: + for name in names: + if name.startswith(("vision.", "aligner.")) or name in ("image_start", "image_end", "image_newline"): + yield f"visual.{name}", reader.get(name) + args = load_args(model_path) + yield "visual.routing_bias", torch.stack([reader.get(f"layers.{i}.ffn.gate.bias_vl") for i in range(args.n_layers)]) + finally: + reader.close() + + +# ----- FTW side files: the Engram tables ---------------------------------------------------------- +def engram_table_names(layer_id: int) -> tuple[str, str]: + return f"layers.{layer_id}.engram.embed.weight", f"layers.{layer_id}.engram.embed.scale" + + +def ftw_side_files(model_path: str, out_dir: str) -> list[str]: + """Copy the checkpoint shards holding the Engram tables (fp8 rows + ue8m0 scales) next to an FTW + checkpoint. The tables are never FTW entries: ``load_host_tables`` maps them in place from whatever + safetensors file holds them (``engram_table.engram_row_source`` reads the headers), and the FTW + loader never reads ``.safetensors`` as weights, so the shard's other tensors ride along unused.""" + import shutil + + from freetoken.models.loader import safetensors_weight_map + from freetoken.utils import download_hf_weight + + folder = download_hf_weight(model_path) + args = load_args(folder) + shard_of = safetensors_weight_map(folder) + shards = sorted({shard_of[n] for layer_id in args.engram_layer_ids if layer_id < args.n_layers for n in engram_table_names(layer_id)}) + for shard in shards: + shutil.copyfile(os.path.join(folder, shard), os.path.join(out_dir, shard)) + return shards + + +# ----- routed MXFP4 expert pieces --------------------------------------------------------------- +_EXPERT_RE = re.compile(r"^layers\.(?P\d+)\.ffn\.experts\.(?P\d+)\.(?Pw1|w2|w3)\.(?Pweight|scale)$") +_PROJ_ROLE = {"w1": "gate", "w3": "up", "w2": "down"} +_KIND_SUFFIX = {"weight": "", "scale": "_scale"} + + +def iter_expert_pieces(model_path: str, config, kind: QuantKind, *, parallel: bool | None = False, workers: int = 8, chunk: int = 8 << 20): + """One piece per expert: ``{gate, up, down}`` e2m1 pairs and their e8m0 ``_scale`` companions. + The DSpark stage's experts (``mtp.*``) never match.""" + if kind is not QuantKind.MXFP4: + return None + if get_tp_info().size > 1: + raise NotImplementedError("DeepSeek-V4.1 expert banks support TP=1 only") + from freetoken.models.weight import iter_expert_tensors_parallel + from freetoken.moe.expert_pieces import per_expert_pieces + + args = load_args(model_path) + L, E = args.n_layers, args.n_routed_experts + + def locate(raw_name: str): + m = _EXPERT_RE.match(raw_name) + if m is None or int(m["layer"]) >= L: + return None + return int(m["layer"]), int(m["expert"]), _PROJ_ROLE[m["proj"]] + _KIND_SUFFIX[m["kind"]] + + if parallel: + tensors = iter_expert_tensors_parallel(model_path, lambda n: locate(n) is not None, workers=workers, chunk=chunk) + return per_expert_pieces(tensors, locate, tensors_per_expert=6) + + def _serial(): + reader = _ShardReader(model_path, weight_map(model_path), torch.device("cpu")) + try: + for li in tqdm(range(L), desc="Loading DeepSeek-V4.1 experts (serial)", disable=not get_tp_info().is_primary()): + for e in range(E): + base = f"layers.{li}.ffn.experts.{e}" + for proj in ("w1", "w3", "w2"): + for kind_ in ("weight", "scale"): + name = f"{base}.{proj}.{kind_}" + yield name, reader.get(name) + finally: + reader.close() + + return per_expert_pieces(_serial(), locate, tensors_per_expert=6) + + +__all__ = ["iter_weights", "iter_vision_weights", "iter_expert_pieces", "ftw_side_files", "load_args", "weight_map", "dequant_fp8_block"] diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py index 2a4ced10af..fe150ddc29 100644 --- a/python/freetoken/models/register.py +++ b/python/freetoken/models/register.py @@ -156,6 +156,20 @@ class ModelSpec: # the head, the KV compressors and the indexer's scorer ship bf16; the fp8 config has no modules_to_not_convert unquantized_modules=("head", "*.compressor.wkv", "*.compressor.wgate", "*.indexer.weights_proj"), ), + # DeepSeek-V4.1-Flash (model_type deepseek_v41): multimodal wrapper config (text tower in + # text_config). The checkpoint has no ``model.`` root; the compressors, the indexer's key projection / + # scorer, the head and the vision tower + aligner ship bf16 (the fp8 config has no modules_to_not_convert). + "DeepseekV41ForCausalLM": ModelSpec( + "freetoken.models.deepseek_v41", + "DeepseekV41ForCausalLM", + mm_processor="freetoken.mm.processors.deepseek_v41:DeepseekV41MMProcessor", + encoders=(EncoderSpec("vision", "vision_config", ("image",)),), + checkpoint_roots=(("model.layers", "layers"), ("model.embed", "embed"), ("model.norm", "norm")), + packed_modules_mapping=_EXPERTS_W123_PACKED, + unquantized_modules=( + "head", "*.compressor.wkv", "*.compressor.wgate", "*.indexer.wk", "*.indexer.weights_proj", "vision.*", "aligner.*", + ), + ), "Qwen3_5MoeForConditionalGeneration": ModelSpec( "freetoken.models.qwen3_5_moe", "Qwen3_5MoeForConditionalGeneration", diff --git a/python/freetoken/moe/benchbw.py b/python/freetoken/moe/benchbw.py index 83b3d19665..e9d6ffe54e 100644 --- a/python/freetoken/moe/benchbw.py +++ b/python/freetoken/moe/benchbw.py @@ -108,6 +108,7 @@ class Workload: activation: str = "silu" swiglu_alpha: float = 1.702 swiglu_limit: float | None = None + act_block: int | None = None # ds_fp4 W4A8 activation quant block # Preset workloads. Dims from the model configs / benchmarks/bench_offload_cache_copy.py. @@ -120,7 +121,8 @@ class Workload: activation="gpt_oss_swiglu", swiglu_limit=7.0), "gpt-oss-20b": Workload("gpt-oss-20b", 2880, 2880, 32, 4, ("mxfp4_triton",), activation="gpt_oss_swiglu", swiglu_limit=7.0), - "dsv4": Workload("dsv4", 4096, 2048, 256, 6, ("ds_fp4",), swiglu_limit=7.0), + "dsv4": Workload("dsv4", 4096, 2048, 256, 6, ("ds_fp4",), swiglu_limit=7.0, act_block=128), + "dsv4.1-flash": Workload("dsv4.1-flash", 5120, 2304, 384, 6, ("ds_fp4",), swiglu_limit=10.0, act_block=32), "glm4.7-nvfp4": Workload("glm4.7-nvfp4", 5120, 1536, 160, 8, ("nvfp4",)), "glm5.3-flash-nvfp4": Workload( "glm5.3-flash-nvfp4", 4096, 2048, 288, 8, ("nvfp4",), @@ -142,7 +144,7 @@ class Workload: "fp8_block": Workload("dtype:fp8", 2048, 512, 128, 8, ("fp8_block",)), "mxfp4_triton": Workload("dtype:mxfp4", 2880, 2880, 128, 4, ("mxfp4_triton",), activation="gpt_oss_swiglu", swiglu_limit=7.0), - "ds_fp4": Workload("dtype:ds_fp4", 4096, 2048, 128, 6, ("ds_fp4",), swiglu_limit=7.0), + "ds_fp4": Workload("dtype:ds_fp4", 4096, 2048, 128, 6, ("ds_fp4",), swiglu_limit=7.0, act_block=128), } @@ -460,7 +462,7 @@ def _build_cpu_moe_executor(fmt: str, wl: Workload, banks: dict, num_threads: in cache, top_k=wl.top_k, activation=wl.activation, apply_router_weight_on_input=False, num_threads=num_threads, max_tokens=1, device=torch.device("cuda"), swiglu_alpha=wl.swiglu_alpha, - swiglu_limit=wl.swiglu_limit, + swiglu_limit=wl.swiglu_limit, act_block=wl.act_block, ) diff --git a/python/freetoken/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py index a46806740d..c4809075c8 100644 --- a/python/freetoken/moe/cpu_executor.py +++ b/python/freetoken/moe/cpu_executor.py @@ -171,6 +171,7 @@ def __init__( swiglu_alpha: float = 1.702, swiglu_limit: float | None = None, fmt: str | None = None, + act_block: int | None = None, ) -> None: from freetoken.kernel import _cpu_moe from freetoken.moe.legacy_format import canonical_role @@ -323,6 +324,11 @@ def __init__( # the C++ side to skip its own. Measured on DeepSeek-V4-Flash bs=1 decode: # 12.85 -> 15.65 tok/s, output bit-identical (tests/moe/test_dsfp4_prequant.py). self._gpu_prequant = fmt == "ds_fp4" and device.type == "cuda" + # the W4A8 activation round-trip block follows the checkpoint's fp8 block (128 on V4, 32 on V4.1) + self._act_block = act_block + if fmt == "ds_fp4": + assert act_block is not None, "ds_fp4 experts need act_block, the checkpoint's fp8 activation quant block" + self._ext.set_act_block(act_block) if self._gpu_prequant: self._ext.set_input_prequant(True) logger.info_rank0( @@ -482,7 +488,7 @@ def _resolve_dsfp4_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: """DeepSeek-V4 ``ds_fp4`` schema: row-major e2m1 (2/byte) + e8m0 per-32 block scales, no global, no bias. Layout matches nvfp4 (K contiguous per output row), so the C++ GEMV reads it in place. The kernel additionally FP8-round-trips the - activations (block 128) to match DSV4's W4A8 reference, hence the %128 dims.""" + activations (the checkpoint's fp8 block) to match the DeepSeek W4A8 reference.""" gup, gus = banks["gate_up"], banks["gate_up_scale"] dnp, dns = banks["down"], banks["down_scale"] assert gup[0].dtype == torch.uint8 and dnp[0].dtype == torch.uint8, (gup[0].dtype, dnp[0].dtype) @@ -490,7 +496,7 @@ def _resolve_dsfp4_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: I = int(gup[0].shape[1] // 2) H = int(gup[0].shape[2] * 2) assert gup[0].shape[1] == 2 * I - assert H % 128 == 0 and I % 128 == 0, (H, I) # FP8 activation round-trip block=128 + assert H % 32 == 0 and I % 32 == 0, (H, I) # e8m0 scales per 32 along K assert tuple(dnp[0].shape[1:]) == (H, I // 2), (dnp[0].shape, H, I) assert tuple(gus[0].shape[1:]) == (2 * I, H // 32), (gus[0].shape, I, H) assert tuple(dns[0].shape[1:]) == (H, I // 32), (dns[0].shape, H, I) @@ -580,7 +586,7 @@ def decode_submit( # pre-quantized activations and skips its serial scalar pass. from freetoken.kernel.triton.dsv4.fp8_linear import act_quant_fp8_roundtrip - hidden_states = act_quant_fp8_roundtrip(hidden_states, block=128) + hidden_states = act_quant_fp8_roundtrip(hidden_states, block=self._act_block) # D2H: ship this step's activations + routing to pinned host memory. io["x"].copy_(hidden_states, non_blocking=True) diff --git a/python/freetoken/moe/fused_ds_fp4.py b/python/freetoken/moe/fused_ds_fp4.py index 65fffb5cc9..b4dc5cd0d0 100644 --- a/python/freetoken/moe/fused_ds_fp4.py +++ b/python/freetoken/moe/fused_ds_fp4.py @@ -12,10 +12,7 @@ import torch import triton -from freetoken.kernel.triton.dsv4.fp8_linear import ( - act_quant_fp8_inplace, - act_quant_fp8_roundtrip, -) +from freetoken.kernel.triton.dsv4.fp8_linear import act_quant_fp8_roundtrip from freetoken.kernel.triton.dsv4.fused_moe import ( _decode_dsfp4_moe_kernel, _e2m1_lut, @@ -50,7 +47,8 @@ def _grouped_decode( """Grouped per-route GEMM ``c[t,r,:] = a[row] @ dequant(W[slot])^T``. Returns ``[T, top_k, N]``. Used for both gate_up (a_row=token) and down - (a_row=route). FP4 weights are dequantized inline from the slot cache. + (a_row=route). FP4 weights are dequantized inline from the slot cache. A slot + of ``-1`` is an inactive route: its output is zero and its storage is never read. """ T, top_k = slots.shape N = packed_cache.shape[1] @@ -62,14 +60,13 @@ def _grouped_decode( topk_weights = out.new_empty((1, 1), dtype=torch.float32) scale_u8 = scale_cache.view(torch.uint8) - # Decode is a per-route GEMV, bound by the inline FP4 dequant (LUT gather + block - # scale), not weight HBM bandwidth -- so it tops out ~29% of HBM peak. BN=16/ - # BKB=128/1 warp maximizes memory-level parallelism (many single-warp CTAs, no - # cross-warp reduction) -> ~950 GB/s on H100 (was ~830 at BN=8/BKB=256/2). K_BYTES - # must be a multiple of BLOCK_SIZE_KB (gate_up 2048, down 1024 -- both /128). + # One fixed launch geometry, never tuned per machine: the row / warp split changes the per-row + # K reduction layout (and so the output bits) on some Triton versions. 16 rows on 2 warps is at + # or near the fastest split for the V4 and V4.1 decode shapes. Each K tile contains whole + # packed FP4 scale groups. BLOCK_SIZE_N = 16 BLOCK_SIZE_KB = 128 - _NW = 1 + _NW = 2 assert (K // 2) % BLOCK_SIZE_KB == 0, (K, BLOCK_SIZE_KB) grid = (total_routes, triton.cdiv(N, BLOCK_SIZE_N)) _decode_dsfp4_moe_kernel[grid]( @@ -103,28 +100,28 @@ def routed_experts_fp4( down_packed: torch.Tensor, # [S, H, I//2] uint8 down_scale: torch.Tensor, # [S, H, I//32] e8m0 swiglu_limit: float, + act_block: int = 128, ) -> torch.Tensor: """Full routed-expert output (summed over the top-k routes), excludes shared expert. - Precision matches the reference ``fp4_gemm(act_quant(x, 128), W_fp4)``: the gate_up - and down activations are FP8-round-tripped (block 128, ue8m0) before each GEMM. Since - an fp8 value x pow2 scale is exact in bf16, the round-tripped activation entering the - bf16 decode kernel is bit-identical to the reference's dequantized FP8 activation - (validated max diff = 0 vs the tilelang ``fp4_gemm`` reference).""" + Precision matches the reference ``fp4_gemm(act_quant(x, act_block), W_fp4)``: the gate_up + and down activations are FP8-round-tripped (block ``act_block`` -- the checkpoint's fp8 block, + 128 on V4 and 32 on V4.1 -- ue8m0) before each GEMM. Since an fp8 value x pow2 scale is exact + in bf16, the round-tripped activation entering the bf16 decode kernel is bit-identical to the + reference's dequantized FP8 activation (validated max diff = 0 vs the tilelang ``fp4_gemm`` + reference). The routing weight scales each route's down output in the GEMM epilogue.""" T, top_k = slots.shape H = x.shape[1] two_I = gate_up_packed.shape[1] I = two_I // 2 - x = act_quant_fp8_roundtrip(x, 128) # gate_up activation -> FP8 round-trip (no clone) + x = act_quant_fp8_roundtrip(x, act_block) # gate_up activation -> FP8 round-trip (no clone) gate_up = _grouped_decode( x, gate_up_packed, gate_up_scale, slots, None, a_row_is_route=False, mul_routed_weight=False, ) # [T, top_k, 2I] - act = fused_swiglu(gate_up, swiglu_limit) # [T, top_k, I] - - act = act.reshape(T * top_k, I) - act_quant_fp8_inplace(act, 128) # down activation -> FP8 round-trip + # [T, top_k, I] with the down activation's FP8 round-trip in the same pass + act = fused_swiglu(gate_up, swiglu_limit, act_block=act_block).reshape(T * top_k, I) down = _grouped_decode( act, down_packed, down_scale, slots, topk_weights, a_row_is_route=True, mul_routed_weight=True, @@ -191,6 +188,7 @@ def routed_experts_fp4_prefill( down_scale: torch.Tensor, # [S, H, I//32] e8m0 swiglu_limit: float, num_rows: int, + act_block: int = 128, ) -> torch.Tensor: """Grouped-GEMM counterpart of :func:`routed_experts_fp4` for dense prefill chunks: one moe_align sort shared by both GEMMs, each expert's weights @@ -201,7 +199,7 @@ def routed_experts_fp4_prefill( if T * top_k < _GROUPED_MIN_ROUTES: return routed_experts_fp4( x, slots, topk_weights, - gate_up_packed, gate_up_scale, down_packed, down_scale, swiglu_limit, + gate_up_packed, gate_up_scale, down_packed, down_scale, swiglu_limit, act_block=act_block, ) H = x.shape[1] two_I = gate_up_packed.shape[1] @@ -217,16 +215,14 @@ def routed_experts_fp4_prefill( sorted_ids, expert_ids, ntpp = moe_align_block_size(slots, cfg["BLOCK_SIZE_M"], num_rows) tw = topk_weights.reshape(-1).contiguous() - x = act_quant_fp8_roundtrip(x, 128) # gate_up activation -> FP8 round-trip (no clone) + x = act_quant_fp8_roundtrip(x, act_block) # gate_up activation -> FP8 round-trip (no clone) gate_up = torch.empty((T, top_k, two_I), dtype=x.dtype, device=x.device) _grouped_prefill( x, gate_up_packed, gate_up_scale, gate_up, tw, sorted_ids, expert_ids, ntpp, routes, top_k, False, cfg, ) - act = fused_swiglu(gate_up, swiglu_limit) # [T, top_k, I] - - act = act.reshape(routes, I) - act_quant_fp8_inplace(act, 128) # down activation -> FP8 round-trip + # [T, top_k, I] with the down activation's FP8 round-trip in the same pass + act = fused_swiglu(gate_up, swiglu_limit, act_block=act_block).reshape(routes, I) down = torch.empty((T, top_k, H), dtype=x.dtype, device=x.device) _grouped_prefill( act, down_packed, down_scale, down, tw, diff --git a/python/freetoken/scheduler/cache.py b/python/freetoken/scheduler/cache.py index 8f77be9330..772297e069 100644 --- a/python/freetoken/scheduler/cache.py +++ b/python/freetoken/scheduler/cache.py @@ -66,6 +66,11 @@ def __init__(self, num_pages: int, page_size: int, page_table: torch.Tensor, typ supports_runtime_rebuild = True prefill_chunk_budget = None # generic shared page pool: no per-model prefill chunk cap + @property + def prefix_replay_tokens(self) -> int: + """Prompt-tail tokens a prefix hit must leave to the prefill (a pool whose model recomputes them).""" + return getattr(self.swa_pool, "prefix_replay_tokens", 0) if self.swa_pool is not None else 0 + @property def prefill_chunk_align(self) -> int: """Granularity a non-final prefill chunk should end on. A hybrid snapshot is donated only @@ -90,10 +95,10 @@ def _make_prefix_cache(self, device, page_size, type): return SWARadixCache(device, page_size, self.sliding_window_size) return create_prefix_cache(device=device, type=type, page_size=page_size) - def match_req(self, req: PendingReq) -> MatchResult: + def match_req(self, req: PendingReq, max_len: int | None = None) -> MatchResult: input_len = req.input_len assert input_len > 0, "Input length must be greater than 0." - ids = req.input_ids[: input_len - 1] + ids = req.input_ids[: input_len - 1 if max_len is None else max(0, min(input_len - 1, max_len))] if self.is_swa: from freetoken.kvcache.swa_radix_cache import SWACacheHandle m = self.prefix_cache.match_prefix(ids) diff --git a/python/freetoken/scheduler/prefill.py b/python/freetoken/scheduler/prefill.py index bd8129aae8..de4bddfb4a 100644 --- a/python/freetoken/scheduler/prefill.py +++ b/python/freetoken/scheduler/prefill.py @@ -71,8 +71,9 @@ def _try_allocate_one(self, req: PendingReq): if self.table_manager.available_size == 0: return None - # TODO: consider host cache match case - mr = self.cache_manager.match_req(req) + # TODO: consider host cache match case (it needs the same replay cap) + replay = self.cache_manager.prefix_replay_tokens + mr = self.cache_manager.match_req(req, max_len=max(0, req.input_len - replay) if replay else None) handle = mr.cuda_handle cached_len = handle.cached_len # TODO: better estimate policy @@ -214,6 +215,7 @@ def _add_one_req( uid=pending_req.uid, cache_handle=cache_handle, sampling_params=pending_req.sampling_params, + prompt_len=pending_req.input_len, ) req.mm_items = pending_req.mm_items req.mrope_positions_full = pending_req.mrope_positions_full diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index fa0d96a714..d057d3b908 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -349,7 +349,7 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: next_token = int(next_token.item()) # EOS / stop-string -> "stop", output budget exhausted -> "length"; # EOS and stop strings win over length. - # Overlap can advance device_len ahead of the token delivered to the host. + # Overlap can advance device_len one token ahead of this host reply. hit_length = req.input_ids.numel() >= req.max_device_len hit_eos = ( not req.sampling_params.ignore_eos and next_token in self.eos_token_ids @@ -737,17 +737,10 @@ def _current_cache_geometry(self) -> dict: the CONCRETE current window (usable pages) so a rollback restores it byte-for-byte, whether it was pinned or ratio-derived.""" eng = self.engine - config = self.config - mc = config.model_config - num_swa_pages = None - if getattr(mc, "dsv4_args", None) is not None: - sizes = getattr(eng.kv_cache, "sizes", None) - if sizes is not None: # usable window pages = physical n_win_pages minus the dummy page - num_swa_pages = max(0, sizes.n_win_pages - 1) - elif getattr(mc, "has_swa_attention", False) and ( - getattr(config, "cache_type", None) == "swa_radix" - ): # usable window tokens = pool tokens minus the slot-0 sentinel - num_swa_pages = max(0, int(getattr(eng.kv_cache, "swa_num_tokens", 0) or 0) - 1) + from freetoken.kvcache.cache_status import window_pool_spec + + spec = window_pool_spec(self.config) + num_swa_pages = eng.kv_cache.window_pages if spec is not None else None return dict( num_pages=eng.num_pages, moe_cache_size=eng.moe_offload_cache.cache_size if eng.moe_offload_cache is not None else None, diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index 5592a2b0d5..34221cfc70 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -563,8 +563,12 @@ def _resolve_num_swa_pages(state: FrontendManager, req: CacheRebuildRequest) -> or pools.get("num_pages", 0) or 0 ) page_size = int(pools.get("page_size", 0) or getattr(config, "page_size", 1) or 1) - is_dsv4 = getattr(getattr(config, "model_config", None), "dsv4_args", None) is not None - swa_page_size = page_size if is_dsv4 else 1 + from freetoken.kvcache.cache_status import window_pool_spec + + spec = window_pool_spec(config) + if spec is None: + raise ValueError("swa_full_tokens_ratio requested but this model has no resizable window pool") + swa_page_size = spec.page_size window_tokens = int(round(req.swa_full_tokens_ratio * num_pages * page_size)) return max(1, -(-window_tokens // swa_page_size)) # ceil-div to the pool's page unit @@ -604,12 +608,16 @@ async def cache_rebuild(req: CacheRebuildRequest): {"status": "failed", "error": "swa_full_tokens_ratio must be in (0, 1]"}, status_code=422, ) + try: + num_swa_pages = _resolve_num_swa_pages(state, req) + except ValueError as exc: + return JSONResponse({"status": "failed", "error": str(exc)}, status_code=422) result = await dispatch_rebuild( state, moe_cache_size=req.moe_cache_size, num_pages=req.num_pages, num_mamba_slots=req.num_mamba_slots, - num_swa_pages=_resolve_num_swa_pages(state, req), + num_swa_pages=num_swa_pages, mode=req.mode, timeout=req.timeout, ) @@ -739,11 +747,11 @@ def cache_geometry(state: Any) -> dict: # the internal currency), else the load-time meta value. 0.0 for models without a window pool. swa_full_tokens_ratio = float(getattr(state, "swa_full_tokens_ratio", 0.0) or 0.0) last_swa_pages = last.get("num_swa_pages") - if last_swa_pages: - is_dsv4 = getattr(getattr(config, "model_config", None), "dsv4_args", None) is not None - full = num_pages if is_dsv4 else num_pages * page_size + swa_page_size = int(pools.get("swa_page_size", 0) or 0) + if last_swa_pages and swa_page_size: + full = num_pages * page_size if full > 0: - swa_full_tokens_ratio = min(1.0, last_swa_pages / full) + swa_full_tokens_ratio = min(1.0, last_swa_pages * swa_page_size / full) try: model_config = config.model_config num_experts = int(getattr(model_config, "num_experts", 0) or 0) @@ -785,9 +793,8 @@ def cache_geometry(state: Any) -> dict: "moe_cache_policy": getattr(config, "moe_cache_policy", None), "unit_bytes": unit_bytes, "swa_full_tokens_ratio": swa_full_tokens_ratio, - # The window pool's page unit for num_swa_pages: DSV4 = P (== page_size), radix-SWA = 1 - # token, 0 for models without a window pool. Lets a client denominate the swa control. - "swa_page_size": int(pools.get("swa_page_size", 0) or 0), + # The pool's window page unit; 0 for models without a resizable window pool. + "swa_page_size": swa_page_size, # Concrete current window size in that unit (usable pages): the last rebuild's value if it # pinned/derived one, else the load-time pool size. 0 for models without a window pool. "num_swa_pages": int(last.get("num_swa_pages") or pools.get("num_swa_pages", 0) or 0), diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 2194b82e92..b8d06018d8 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -204,6 +204,8 @@ def _infer_tool_call_parser(model_path: str) -> str: return "qwen3_coder" if "qwen" in marker: return "qwen25" + if "deepseek_v41" in marker or "deepseekv41" in marker: + return "deepseekv41" if "deepseek" in marker and ("v4" in marker or "deepseek_v4" in marker): return "deepseekv32" if "deepseek" in marker and ("v3.2" in marker or "v32" in marker): @@ -576,6 +578,7 @@ def _infer_reasoning_parser(model_path: str) -> str | None: "qwen3_coder", "mistral", "deepseekv32", + "deepseekv41", "gemma4", "glm47", "minimax", @@ -645,6 +648,17 @@ def _infer_reasoning_parser(model_path: str) -> str | None: ), ) + parser.add_argument( + "--swa-decoder-replay", + default=ServerArgs.swa_decoder_replay, + choices=["bounded", "exact"], + help=( + "DeepSeek-V4.1 Decoder SWA Bounded Replay. 'bounded' (default) runs the 20 decoder layers " + "on each prompt's last 128 tokens with their sliding window truncated there, as in the " + "tech report; 'exact' runs them on every prompt token (the reference numerics)." + ), + ) + parser.add_argument( "--nvfp4-backend", action=_DeprecatedAlias, diff --git a/python/freetoken/server/function_call_parser.py b/python/freetoken/server/function_call_parser.py index 0804dcc3f7..296cd560dc 100644 --- a/python/freetoken/server/function_call_parser.py +++ b/python/freetoken/server/function_call_parser.py @@ -78,6 +78,8 @@ def _should_forward_unknown_tool(name: Any) -> bool: "<｜DSML｜function_calls>", "<｜DSML｜tool_calls>", "<｜DSML｜invoke", + "<｜DSML｜ calls>", + "<｜DSML｜ invoke", "", ] @@ -1563,37 +1565,45 @@ class DeepSeekV32Detector(BaseFormatDetector): Reference: https://huggingface.co/deepseek-ai/DeepSeek-V3.2 """ + # the DSML tag names: (block, alternative block), invoke, parameter + block_tags = ("function_calls", "tool_calls") + invoke_tag = "invoke" + param_tag = "parameter" + def __init__(self): super().__init__() self.dsml_token = "｜DSML｜" - self.bot_token = f"<{self.dsml_token}function_calls>" - self.eot_token = f"" - self.alt_bot_token = f"<{self.dsml_token}tool_calls>" - self.alt_eot_token = f"" - self.invoke_start_prefix = f"<{self.dsml_token}invoke" - self.invoke_end_token = f"" - self.param_end_token = f"" + block, alt_block = self.block_tags + self.bot_token = f"<{self.dsml_token}{block}>" + self.eot_token = f"" + self.alt_bot_token = f"<{self.dsml_token}{alt_block}>" + self.alt_eot_token = f"" + self.invoke_start_prefix = f"<{self.dsml_token}{self.invoke_tag}" + self.invoke_end_token = f"" + self.param_start_prefix = f"<{self.dsml_token}{self.param_tag}" + self.param_end_token = f"" # Regex for complete invoke extraction - _de = re.escape(self.dsml_token) + _inv = re.escape(self.dsml_token + self.invoke_tag) + _par = re.escape(self.dsml_token + self.param_tag) self.invoke_regex = re.compile( - rf'<{_de}invoke\s+name="([^"]+)"\s*>(.*?)', + rf'<{_inv}\s+name="([^"]+)"\s*>(.*?)', re.DOTALL, ) # Regex for parameter extraction self.param_regex = re.compile( - rf'<{_de}parameter\s+name="([^"]+)"(?:\s+string="(true|false)")?\s*>(.*?)', + rf'<{_par}\s+name="([^"]+)"(?:\s+string="(true|false)")?\s*>(.*?)', re.DOTALL, ) # Regex for partial invoke (name known, body still streaming) self.partial_invoke_regex = re.compile( - rf'<{_de}invoke\s+name="([^"]+)"\s*>(.*)', + rf'<{_inv}\s+name="([^"]+)"\s*>(.*)', re.DOTALL, ) # Streaming state machine tag regexes (anchored matches over the buffer). - self.invoke_open_regex = re.compile(rf'<{_de}invoke\s+name="([^"]+)"\s*>') + self.invoke_open_regex = re.compile(rf'<{_inv}\s+name="([^"]+)"\s*>') self.param_open_regex = re.compile( - rf'<{_de}parameter\s+name="([^"]+)"(?:\s+string="(true|false)")?\s*>' + rf'<{_par}\s+name="([^"]+)"(?:\s+string="(true|false)")?\s*>' ) self._last_arguments = "" @@ -1687,7 +1697,7 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami normal_parts: List[str] = [] calls: List[ToolCallItem] = [] - param_open_token = f"<{self.dsml_token}parameter" + param_open_token = self.param_start_prefix def _emit_args(fragment: str) -> None: if fragment: @@ -3512,6 +3522,16 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult return StreamingParseResult(normal_text="".join(normal_parts).strip(), calls=calls) +class DeepSeekV41Detector(DeepSeekV32Detector): + """DeepSeek-V4.1's DSML: the V3.2 grammar with every tag name led by a space and the block + named ``calls`` (``<｜DSML｜ calls>`` / ``<｜DSML｜ invoke name=...>`` / ``<｜DSML｜ parameter + ...>``, the checkpoint's ``encoding/encoding.py``).""" + + block_tags = (" calls", " calls") + invoke_tag = " invoke" + param_tag = " parameter" + + class FunctionCallParser: """ Parser for function/tool calls in model outputs. @@ -3523,6 +3543,7 @@ class FunctionCallParser: ToolCallParserEnum: Dict[str, Type[BaseFormatDetector]] = { "deepseekv32": DeepSeekV32Detector, + "deepseekv41": DeepSeekV41Detector, "gemma4": Gemma4Detector, "gpt-oss": GptOssDetector, "gpt_oss": GptOssDetector, diff --git a/python/freetoken/server/stats.py b/python/freetoken/server/stats.py index d01828e21b..3d48ee96a0 100644 --- a/python/freetoken/server/stats.py +++ b/python/freetoken/server/stats.py @@ -120,12 +120,10 @@ def derive_model_card(config: Any) -> dict: def _swa_page_size(config: Any) -> int: - """The window pool's own page unit: P (window_size) for DSV4, 1 token for radix-SWA. - Mirrors compute_cache_pools' swa_page_size.""" - dsv4 = getattr(getattr(config, "model_config", None), "dsv4_args", None) - if dsv4 is not None: - return int(getattr(dsv4, "window_size", 0) or 1) - return 1 + from freetoken.kvcache.cache_status import window_pool_spec + + spec = window_pool_spec(config) + return spec.page_size if spec is not None else 1 def build_stats(state: Any, p95_ms: int, ttft_mean_ms: int) -> dict: diff --git a/python/freetoken/tokenizer/effort.py b/python/freetoken/tokenizer/effort.py index c81cd7bef1..2e9e3cd21f 100644 --- a/python/freetoken/tokenizer/effort.py +++ b/python/freetoken/tokenizer/effort.py @@ -84,7 +84,8 @@ class EffortProfile: template interpolates anything". ``strength_dialect`` True means the template reads the ``reasoning_strength`` spelling (muse-glimmer's family, whose documented ladder tops out at xhigh) rather than only - ``reasoning_effort``. + ``reasoning_effort``. ``integer_range`` records a validated 1-100 budget + accepted independently of the named effort vocabulary. """ supported: frozenset[str] @@ -92,6 +93,7 @@ class EffortProfile: consumes_effort: bool validates: bool = False strength_dialect: bool = False + integer_range: tuple[int, int] | None = None @dataclass(frozen=True) @@ -118,7 +120,7 @@ class ThinkingProfile: _MAX_QUANTIZE_DISTANCE = 0.15 -def quantize_effort(value: Any, profile: EffortProfile) -> str | None: +def quantize_effort(value: Any, profile: EffortProfile) -> str | int | None: """Map a client's effort onto ``profile``; ``None`` means "send nothing". In-vocabulary values pass through untouched. Other named levels land on the @@ -128,6 +130,11 @@ def quantize_effort(value: Any, profile: EffortProfile) -> str | None: the template default. With "max" excluded the remaining scale values are unique, so quantization is deterministic across processes. """ + if profile.integer_range is not None and isinstance(value, (int, float)): + low, high = profile.integer_range + if type(value) is not int or not low <= value <= high: + raise ValueError(f"reasoning_effort must be an integer in [{low}, {high}]") + return value supported = effective_efforts(profile) if not supported: return None @@ -181,6 +188,7 @@ def probe_effort_profile( diverged: set[str] = set() matches_baseline: dict[str, bool] = {name: True for name in KNOWN_REASONING_EFFORTS} ran_rounds = 0 + integer_range = None for base_kwargs, tools in _PROBE_ROUNDS: try: @@ -198,6 +206,10 @@ def probe_effort_profile( if _renderings_differ(rendering, baseline): diverged.add(name) matches_baseline[name] = False + # Only a validated budget range is a numeric capability; arbitrary Jinja + # interpolation must not advertise support for numeric reasoning effort. + if _accepts_integer_budget(render, base_kwargs, tools): + integer_range = (1, 100) if ran_rounds == 0: # Nothing learnable: sending no effort is the only safe rendering. @@ -214,7 +226,7 @@ def probe_effort_profile( strength_dialect = False supported = frozenset(name for name in KNOWN_REASONING_EFFORTS if name not in rejected) - consumes = bool(rejected or diverged) + consumes = bool(rejected or diverged or integer_range) default = None if consumes: defaults = [name for name in supported if matches_baseline[name]] @@ -226,9 +238,26 @@ def probe_effort_profile( consumes_effort=consumes, validates=bool(rejected), strength_dialect=strength_dialect, + integer_range=integer_range, ) +def _accepts_integer_budget(render, base_kwargs, tools) -> bool: + for value in (0, 101): + try: + render({**base_kwargs, "reasoning_effort": value}, tools) + except Exception: # noqa: BLE001 -- the encoder owns validation + continue + return False + outputs = [] + for value in range(1, 101): + try: + outputs.append(render({**base_kwargs, "reasoning_effort": value}, tools)) + except Exception: # noqa: BLE001 + return False + return _renderings_differ(outputs[0], outputs[-1]) + + def probe_thinking_profile( render: Callable[[dict[str, Any], list[dict[str, Any]] | None], Any], efforts: EffortProfile, diff --git a/python/freetoken/tokenizer/tokenize.py b/python/freetoken/tokenizer/tokenize.py index f473665267..772472573a 100644 --- a/python/freetoken/tokenizer/tokenize.py +++ b/python/freetoken/tokenizer/tokenize.py @@ -59,7 +59,7 @@ def __init__(self, tokenizer: PreTrainedTokenizerBase, mm_processor: MMProcessor self._effort_profile: EffortProfile | None = None self._thinking_profile: ThinkingProfile | None = None self._effort_lock = threading.Lock() - self._logged_effort_maps: set[tuple[Any, str | None]] = set() + self._logged_effort_maps: set[tuple[Any, str | int | None]] = set() self._inline_system_lock = threading.Lock() self._inline_system_modes: dict[str, str] = {} @@ -117,7 +117,7 @@ def render_prompt(self, msg: TokenizeMsg) -> str: if mode == "auto": mode = self.inline_system_mode(msg.tools, kwargs) messages = normalize_inline_system(messages, mode) - return self._render(messages, msg.tools, kwargs) + return self._render(messages, msg.tools, kwargs, images=msg.images) def inline_system_mode(self, tools: list[dict] | None, kwargs: dict[str, Any]) -> str: key = json.dumps([tools, kwargs], sort_keys=True) @@ -137,12 +137,13 @@ def _render( messages: list[dict[str, Any]], tools: list[dict[str, Any]] | None, chat_template_kwargs: dict[str, Any], + *, images: list[bytes] | None = None, ) -> str: """Raw render, no effort sanitation — the probe needs unsupported values to actually reach the template so rejection is observable.""" if self._dsv4_encoder is not None: return _apply_dsv4_chat_encoder( - self._dsv4_encoder, messages, tools, chat_template_kwargs + self._dsv4_encoder, messages, tools, chat_template_kwargs, images=images ) # Broadcast the effort in every spelling the ecosystem's templates read # (muse-glimmer grades ``reasoning_strength``; Jinja ignores undeclared @@ -222,8 +223,12 @@ def _load_dsv4_encoder_if_needed(tokenizer: PreTrainedTokenizerBase) -> ModuleTy model_path = getattr(tokenizer, "name_or_path", None) or getattr(tokenizer, "_name_or_path", "") if not model_path: return None - encoder_path = os.path.join(str(model_path), "encoding", "encoding_dsv4.py") - if not os.path.isfile(encoder_path): + # DeepSeek-V4 ships encoding/encoding_dsv4.py, V4.1 encoding/encoding.py; same encode_messages contract + encoder_path = next( + (p for p in (os.path.join(str(model_path), "encoding", name) for name in ("encoding_dsv4.py", "encoding.py")) if os.path.isfile(p)), + None, + ) + if encoder_path is None: return None spec = importlib.util.spec_from_file_location("encoding_dsv4", encoder_path) if spec is None or spec.loader is None: @@ -240,9 +245,28 @@ def _apply_dsv4_chat_encoder( messages: list[dict], tools: list[dict] | None, chat_template_kwargs: dict, + *, images: list[bytes] | None = None, ) -> str: rendered_messages = [dict(message) for message in messages] + image_data = iter(images) if images is not None else None for message in rendered_messages: + if isinstance(message.get("content"), list): + content = [] + for part in message["content"]: + if isinstance(part, dict) and part.get("type") == "image": + # Prerender still has refs; the tokenizer worker receives fetched bytes. + # The checkpoint encoder validates sources even when rendering only text. + ref = part.get("freetoken_ref") + if image_data is not None: + try: + part = {"type": "image", "data": next(image_data)} + except StopIteration: + raise ValueError("image blocks outnumber fetched images") from None + elif ref is not None: + key = "data" if ref["kind"] == "b64" else "url" + part = {"type": "image", key: ref["data"]} + content.append(part) + message["content"] = content if message.get("tool_calls"): message["tool_calls"] = _dsv4_tool_calls(message["tool_calls"]) if tools: diff --git a/scripts/ftw_hotfix.py b/scripts/ftw_hotfix.py index f9cb62d1a8..9640c7b7f7 100644 --- a/scripts/ftw_hotfix.py +++ b/scripts/ftw_hotfix.py @@ -263,7 +263,9 @@ def is_checkpoint_tower_name(name: str) -> bool: head = name.split(".") if head[0] == "model" and len(head) > 1: head = head[1:] - return "vision" in head[0] or "visual" in head[0] or head[0] in _TOWER_SEGMENTS + return ("vision" in head[0] or "visual" in head[0] or head[0] in _TOWER_SEGMENTS + or head[0] in ("aligner", "image_start", "image_end", "image_newline") + or name.endswith(".ffn.gate.bias_vl")) def read_tower(source: TensorSource, ftw_dir: str, checkpoint_names: list[str]) -> dict[str, torch.Tensor]: diff --git a/tests/README.md b/tests/README.md index be22c1fb20..76b2b1fb03 100644 --- a/tests/README.md +++ b/tests/README.md @@ -62,6 +62,7 @@ checkpoint is set: | `FREETOKEN_TEST_MEM_RATIO` | `e2e/test_aime.py` — offload-mode memory_ratio (default `0.9`) | | `FREETOKEN_REBUILD_TEST_MODEL` | `e2e/test_cache_rebuild.py` — a SMALL local model dir; boots a real server (falls back to `FREETOKEN_TEST_MODEL`) | | `FREETOKEN_GEMMA4_GGUF_GLOB` | `models/test_gemma4_gguf_rope.py` — glob matching a local gemma-4 GGUF file | +| `FREETOKEN_TEST_MODELS_ROOT` | `models/test_quant_config.py` — directory whose checkpoints the stored-tensor scheme scan checks (default `/mnt/nvme/models`) | `test_aime.py` takes its sampling protocol from the checkpoint's own `generation_config.json` (pass@N at the recommended temperature, or a single greedy diff --git a/tests/attention/test_dsv41_backend.py b/tests/attention/test_dsv41_backend.py new file mode 100644 index 0000000000..387742eb63 --- /dev/null +++ b/tests/attention/test_dsv41_backend.py @@ -0,0 +1,208 @@ +"""DSV41 attention backend: addressing and selection contracts, CPU-only (kernels live in tests/kernels). + +* prefill metadata: encoder segments, the decoder pass under exact vs bounded replay; +* window candidates with a floor (bounded replay truncates the window at the prompt's last window); +* decode snapshot staging and the layer-invariant ring context; +* the selection helpers' reshape, candidate gather and position-to-row mapping. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.core import Batch, Context, Req, SamplingParams, get_global_ctx, set_global_ctx +from freetoken.kvcache.dsv4.v41_cost_model import dsv41_pool_sizes +from freetoken.kvcache.dsv4.v41_pool import DSV41PagedKVCache +from freetoken.models.deepseek_v41.args import DeepseekV41Args + +P, MRR, DEVICE = 128, 4, torch.device("cpu") +RATIOS = (0, 2, 2, 1, 1) +SOURCES = (1, 3) + + +def _ctx(pool): + try: + ctx = get_global_ctx() + except AssertionError: + ctx = Context(page_size=P) + set_global_ctx(ctx) + ctx.kv_cache = pool + return ctx + + +def _stack(swa_decoder_replay="exact", num_pages=32, max_seq_len=8192): + from freetoken.attention.dsv41_sparse import DSV41SparseAttnBackend + + args = DeepseekV41Args(n_layers=len(RATIOS), head_dim=512, index_head_dim=128, window_size=P, compress_ratios=RATIOS, + kv_source_layers=SOURCES, swa_decoder_replay=swa_decoder_replay) + pool = DSV41PagedKVCache(dsv41_pool_sizes(num_pages + 1, args, 1.0, P), args, DEVICE, n_scratch=MRR + 1) + pool._init_paged_state(MRR, True) + pt = torch.zeros(MRR + 1, max_seq_len, dtype=torch.int32) + pt[MRR].fill_(num_pages * P) + pt[2, :300] = torch.arange(300, dtype=torch.int32) + for page in range(3): + pool.bind_window_pages(page * P, page * P) + pool.full_loc_map = pt + _ctx(pool) + backend = DSV41SparseAttnBackend(SimpleNamespace(dsv41_args=args)) + return backend, pool, pt + + +def _req(table_idx, cached_len, n_new, output_len=16, uid=0, prompt_len=0): + return Req(input_ids=torch.zeros(cached_len + n_new, dtype=torch.int32), table_idx=table_idx, cached_len=cached_len, + output_len=output_len, uid=uid, sampling_params=SamplingParams(), cache_handle=None, prompt_len=prompt_len) + + +def _prefill_batch(reqs): + batch = Batch(reqs=reqs, phase="prefill") + batch.padded_reqs = reqs + return batch + + +def test_prefill_metadata_exact_mode_runs_the_decoder_on_every_token(): + backend, pool, _ = _stack("exact") + assert pool.prefix_replay_tokens == 0 + batch = _prefill_batch([_req(0, 0, 300), _req(2, 256, 44, uid=1)]) + backend.prepare_metadata(batch) + md = batch.attn_metadata + assert [(s.offset, s.n, s.table_idx, s.start_pos, s.window_floor) for s in md.segments] == [(0, 300, 0, 0, 0), (300, 44, 2, 256, 0)] + assert md.decoder_segments is md.segments and md.decoder_rows is None + assert md.last_indices.tolist() == [299, 343] + + +def test_prefill_metadata_bounded_mode_runs_the_decoder_on_the_last_window(): + """The encoder runs every new token; the decoder runs each request's part of its prompt's last + window ``[L - P, L)``, floored at ``L - P``. A chunk ending before that window runs no decoder row + and its ``last_indices`` entry points at the placeholder row the model appends.""" + backend, pool, _ = _stack("bounded") + assert pool.prefix_replay_tokens == P and pool.sliding_window_size == P + cold = _req(0, 0, 300) # a whole prompt: L - P = 172 + early = _req(1, 256, 128, uid=1, prompt_len=600) # chunk [256, 384) of a 600-token prompt: before 472 + straddle = _req(2, 384, 128, uid=2, prompt_len=600) # chunk [384, 512): its rows [472, 512) are in the window + batch = _prefill_batch([cold, early, straddle]) + backend.prepare_metadata(batch) + md = batch.attn_metadata + assert [(s.offset, s.n, s.start_pos, s.window_floor) for s in md.segments] == [(0, 300, 0, 0), (300, 128, 256, 0), (428, 128, 384, 0)] + assert [(d.offset, d.n, d.table_idx, d.start_pos, d.window_floor) for d in md.decoder_segments] == [(0, 128, 0, 172, 172), (128, 40, 2, 472, 472)] + assert md.decoder_rows.tolist() == list(range(172, 300)) + list(range(428 + 88, 556)) + assert md.decoder_pad and md.last_indices.tolist() == [127, 168, 167] # 168: the placeholder after 168 decoder rows + + +def test_prefill_metadata_bounded_mode_without_any_decoder_row(): + backend, _, _ = _stack("bounded") + batch = _prefill_batch([_req(0, 0, 256, prompt_len=1024)]) + backend.prepare_metadata(batch) + md = batch.attn_metadata + assert md.decoder_segments == [] and md.decoder_rows.numel() == 0 + assert md.decoder_pad and md.last_indices.tolist() == [0] + + +def test_prefill_metadata_short_prompt_runs_the_decoder_everywhere(): + """A prompt shorter than the window floors the decoder's window at 0: it sees what exact mode sees.""" + backend, _, _ = _stack("bounded") + batch = _prefill_batch([_req(0, 0, 50)]) + backend.prepare_metadata(batch) + md = batch.attn_metadata + assert [(d.offset, d.n, d.start_pos, d.window_floor) for d in md.decoder_segments] == [(0, 50, 0, 0)] + assert not md.decoder_pad and md.last_indices.tolist() == [49] + + +def test_private_window_layers_address_per_request_rings(): + """Under bounded replay the decoder layers keep their window KV in per-request rings (slot = + table row * P + pos % P), never in the radix-shared window pages; the encoder layers stay on + the shared pool.""" + from freetoken.attention.dsv41_sparse import PrefillSegment + + backend, pool, _ = _stack("bounded") + assert pool.private_window_layer_ids == (3, 4) and not any(pool.is_private_window(l) for l in (0, 1, 2)) + assert pool.window_pool[3].shape[0] == (MRR + 1) * P and pool.window_pool[0].shape[0] == pool.sizes.n_win_slots + # prefill: a decoder segment's slots are its ring row; an encoder segment's are page-bound + seg = PrefillSegment(0, 44, 2, 256, window_floor=256) + assert backend.layer_window_slots_of(3, 2, 256, 300).tolist() == [2 * P + p % P for p in range(256, 300)] + assert backend.layer_window_slots_of(0, 2, 256, 300).tolist() == list(range(256, 300)) + g = backend.window_topk_prefill(seg, layer_id=3)[0] + assert g.shape == (44, P) and g[0, 0].item() == 2 * P + 0 and (g[0, 1:] == -1).all() # position 256 sees itself only (floor) + assert set(g[43].tolist()) - {-1} == {2 * P + p % P for p in range(256, 300)} + # decode: the private ring context keys on the batch row's page-table row + batch = _decode_batch([2], [300]) + backend.prepare_metadata(batch) + md = batch.attn_metadata + slots, topk = md.private_window_ctx(torch.tensor([300]), torch.arange(1)) + assert slots.tolist() == [2 * P + 300 % P] and topk.shape == (1, 1, P) + assert set(topk[0, 0].tolist()) == {2 * P + j for j in range(P)} # every ring index holds a position <= 300 + + +def test_window_candidates_honor_the_floor_and_the_retained_prefix(): + from freetoken.attention.dsv41_sparse import PrefillSegment + + backend, pool, _ = _stack() + # a cold segment at positions [0, 10): query p sees [0, p] + g = backend.window_topk_prefill(PrefillSegment(0, 10, 2, 0)) + assert g.shape == (1, 10, P) + assert g[0, 3, :4].tolist() == [0, 1, 2, 3] and (g[0, 3, 4:] == -1).all() + # an extend at [256, 300) sees back into the retained prefix, but not below the floor + g = backend.window_topk_prefill(PrefillSegment(0, 44, 2, 256, window_floor=200)) + row = g[0, 0] # query at 256 -> [200, 256] + assert row[0].item() == 200 and row[56].item() == 256 and (row[57:] == -1).all() + # without a floor the query at 256 sees [129, 256] + g = backend.window_topk_prefill(PrefillSegment(0, 44, 2, 256)) + assert g[0, 0, 0].item() == 129 and g[0, 0, 127].item() == 256 + + +def _decode_batch(rows, positions): + reqs = [_req(int(t), 0, 1, uid=i) for i, t in enumerate(rows)] + batch = Batch(reqs=reqs, phase="decode") + batch.padded_reqs = reqs + batch.active_table_idx = torch.tensor(rows, dtype=torch.int64) + batch.positions = torch.tensor(positions, dtype=torch.int64) + return batch + + +def test_decode_snapshot_and_ring_context(): + backend, pool, pt = _stack() + batch = _decode_batch([2, MRR], [259, 0]) + backend.prepare_metadata(batch) + md = batch.attn_metadata + assert md.full_snap is None + snap = md.full_snapshot() + assert torch.equal(snap[0, :300], torch.arange(300)) + pt[2, :10] = 7 # a later mutation must not reach the snapshot + assert torch.equal(snap[0, :300], torch.arange(300)) + pos, rows = batch.positions, torch.arange(2) + ws, prev, topk = md.window_ctx(pos, rows) + dummy_ws = pool.sizes.n_win_slots - P # the dummy row's full loc binds to the dummy window page + assert ws.tolist() == [259, dummy_ws] and prev.tolist() == [258, dummy_ws] + assert topk.shape == (2, 1, P) and (topk[0] >= 0).all() + assert set(topk[0, 0].tolist()) == set(range(259 - 127, 260)) + assert topk[1, 0, 0].item() == dummy_ws and (topk[1, 0, 1:] == -1).all() + # the indexer derives compressed rows from the live locs: prefill off the page table row, + # decode off the snapshot (position t -> full_loc(t * ratio) // ratio) + assert torch.equal(backend.locs_prefill(2, 20), pt[2:3, :20]) + assert snap[0, :16].tolist() == list(range(16)) # taken before the mutation above + # capture staging reuses one buffer + backend.init_capture_graph(512, [1, 2]) + backend.prepare_for_replay(batch) + assert batch.attn_metadata.full_snap.data_ptr() == backend.capture.full_snap.data_ptr() + assert batch.attn_metadata.stage_width == 512 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="the selection kernels need CUDA") +def test_topk_helpers(): + from freetoken.attention.dsv41_sparse import DSV41SparseAttnBackend as B + + scores = torch.tensor([[[0.1, 5.0, 3.0, 7.0, 9.0]]], device="cuda") # columns past live are never read + live = torch.tensor([[3]], device="cuda", dtype=torch.int32) + assert B.select_topk(scores, live, 4).tolist() == [[[0, 1, 2, -1]]] + assert B.select_topk(scores, live, 2).tolist() == [[[1, 2]]] + # a candidate list is a sorted valid prefix; empty slots score -inf + cand = torch.tensor([[[2, 7, 9, -1]]], device="cuda", dtype=torch.int32) + cscores = torch.tensor([[[1.0, 4.0, -torch.inf, -torch.inf]]], device="cuda") + assert B.select_topk_in_candidates(cscores, cand, 3).tolist() == [[[2, 7, -1]]] + # rows come from the live locs: compressed position p -> locs[p * ratio] // ratio + locs = torch.arange(20, 40, device="cuda", dtype=torch.int32).view(1, 20) + locs[0, 14:16] = -1 # position 7 (ratio 2) has slid out + rows = B.positions_to_rows(torch.tensor([[[2, 7, 9, -1]]], device="cuda", dtype=torch.int32), locs, 2) + assert rows.tolist() == [[[12, -1, 19, -1]]] diff --git a/tests/checkpoint/test_ftw_hotfix_vision.py b/tests/checkpoint/test_ftw_hotfix_vision.py index d4a3eac68e..eb0c592f35 100644 --- a/tests/checkpoint/test_ftw_hotfix_vision.py +++ b/tests/checkpoint/test_ftw_hotfix_vision.py @@ -14,7 +14,7 @@ # env var -> local checkpoint, as tests/README.md lists them _CHECKPOINTS = ("FREETOKEN_QWEN3VL_MODEL", "FREETOKEN_GEMMA4_MODEL", "FREETOKEN_GEMMA4_UNIFIED_MODEL", - "FREETOKEN_GLM53_MODEL", "FREETOKEN_MUSE_MODEL", "FREETOKEN_MINIMAX_M3_MODEL") + "FREETOKEN_GLM53_MODEL", "FREETOKEN_MUSE_MODEL", "FREETOKEN_MINIMAX_M3_MODEL", "FREETOKEN_DSV41_MODEL") @pytest.fixture(scope="module") @@ -34,6 +34,10 @@ def hotfix(): ("vision_tower.vision_model.embeddings.patch_embedding.weight", True), ("multi_modal_projector.linear_1.weight", True), ("patch_merge_mlp.linear_2.bias", True), + ("vision.blocks.0.norm1.weight", True), + ("aligner.w1.weight", True), + ("image_newline", True), + ("layers.3.ffn.gate.bias_vl", True), ("model.embed_audio.embedding_projection.weight", False), ("model.language_model.layers.0.self_attn.q_proj.weight", False), ("model.layers.3.mlp.experts.0.gate_proj.weight", False), diff --git a/tests/dsv4/test_dsv4_attn_backend.py b/tests/dsv4/test_dsv4_attn_backend.py index 8ffc6fd4d9..7bc72210f2 100644 --- a/tests/dsv4/test_dsv4_attn_backend.py +++ b/tests/dsv4/test_dsv4_attn_backend.py @@ -11,8 +11,8 @@ import torch from freetoken.core import Batch, Context, Req, SamplingParams, get_global_ctx, set_global_ctx -from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache -from freetoken.kvcache.dsv4_cost_model import dsv4_pool_sizes +from freetoken.kvcache.dsv4.v4_pool import DSV4PagedKVCache +from freetoken.kvcache.dsv4.v4_cost_model import dsv4_pool_sizes from freetoken.models.deepseek_v4.args import DeepseekV4Args P, MRR, DEVICE = 128, 4, torch.device("cpu") diff --git a/tests/kernels/test_dsv41_indexer.py b/tests/kernels/test_dsv41_indexer.py new file mode 100644 index 0000000000..19b45ab8b9 --- /dev/null +++ b/tests/kernels/test_dsv41_indexer.py @@ -0,0 +1,139 @@ +"""DSV41 indexer logits over packed fp4 index keys: full-range (Full mode) and candidate-restricted +(Reindex mode) against a torch reference on the dequantized keys, with causality from ``live``. The +row of a compressed position comes from the request's full-token locs (``locs[t * ratio] // ratio``); +in Full mode only the live columns are written.""" + +from __future__ import annotations + +import pytest +import torch + +from freetoken.kvcache.dsv4.v41_row_format import FP4_E8M0_B32 + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + +D, H = 128, 32 +ROWS = 4200 + + +RATIO = 2 + + +def _setup(b, s, t, seed): + from freetoken.kernel.triton.dsv41.pack import pack_rows, unpack_rows + + torch.manual_seed(seed) + keys = torch.randn(ROWS, D, device="cuda", dtype=torch.bfloat16) + pool = pack_rows(keys, FP4_E8M0_B32) + deq = unpack_rows(pool, FP4_E8M0_B32, D) + q = torch.randn(b, s, H, D, device="cuda", dtype=torch.bfloat16) + w = (torch.randn(b, s, H, device="cuda") * 0.3).bfloat16() + # each request owns a random permutation of pool rows for its compressed positions; a few holes. + # The kernel derives rows from full-token locs: compressed position p -> loc = row * RATIO (+ any + # in-group offset) at locs[p * RATIO]; a hole is a negative loc. + k_rows = torch.stack([torch.randperm(ROWS, device="cuda")[:t] for _ in range(b)]).to(torch.int32) + k_rows[:, 5::11] = -1 + locs = torch.full((b, t * RATIO), -1, device="cuda", dtype=torch.int32) + for i in range(b): + for p in range(t): + r = int(k_rows[i, p]) + if r >= 0: + locs[i, p * RATIO : (p + 1) * RATIO] = torch.arange(r * RATIO, (r + 1) * RATIO, device="cuda", dtype=torch.int32) + return pool, deq, q, w, k_rows, locs + + +def _reference(q, w, deq, k_rows, live, positions): + """positions: [B, S, N] compressed positions (-1 empty) -> logits [B, S, N].""" + b, s, n = positions.shape + out = torch.full((b, s, n), float("-inf"), device="cuda") + for i in range(b): + for j in range(s): + for c in range(n): + p = int(positions[i, j, c]) + if p < 0 or p >= int(live[i, j]): + continue + r = int(k_rows[i, p]) + if r < 0: + continue + score = (q[i, j] @ deq[r]).relu() * w[i, j] + out[i, j, c] = score.sum() + return out + + +def test_full_range_matches_reference_with_causal_live(): + from freetoken.kernel.triton.dsv41.indexer import indexer_logits_packed + + b, s, t = 2, 5, 300 + pool, deq, q, w, k_rows, locs = _setup(b, s, t, seed=1) + live = torch.tensor([[10, 64, 65, 200, 300], [0, 1, 128, 129, 250]], device="cuda", dtype=torch.int32) + out = torch.full((b, s, t), 12345.0, device="cuda") # a sentinel: the dead area must stay untouched + got = indexer_logits_packed(q, w, pool, FP4_E8M0_B32, locs, RATIO, live, out=out) + positions = torch.arange(t, device="cuda").view(1, 1, t).expand(b, s, t) + want = _reference(q, w, deq, k_rows, live, positions) + assert got.shape == (b, s, t) + written = positions < live.unsqueeze(-1) + touched = positions < (torch.div(live + 63, 64, rounding_mode="floor") * 64).unsqueeze(-1) # whole live tiles + assert (got[~touched] == 12345.0).all() and torch.isneginf(got[touched & ~written]).all() + assert torch.equal(torch.isinf(got[written]), torch.isinf(want[written])) + finite = written & ~torch.isinf(want) + torch.testing.assert_close(got[finite], want[finite], atol=2e-2, rtol=2e-2) + + +def test_quantized_queries_preserve_reference_score_rounding(): + from freetoken.kernel.triton.dsv41.indexer import indexer_logits_packed + from freetoken.kernel.triton.dsv41.pack import pack_rows, unpack_rows + + torch.manual_seed(23) + t, s, topk = 4096, 32, 512 + pool = pack_rows(torch.randn(t, D, device="cuda", dtype=torch.bfloat16), FP4_E8M0_B32) + keys = unpack_rows(pool, FP4_E8M0_B32, D) + q = torch.randn(s * H, D, device="cuda", dtype=torch.bfloat16) + q = unpack_rows(pack_rows(q, FP4_E8M0_B32), FP4_E8M0_B32, D).view(1, s, H, D) + w = torch.randn(1, s, H, device="cuda", dtype=torch.bfloat16) * (D * H) ** -0.5 + locs = torch.arange(t, device="cuda", dtype=torch.int32).view(1, t) + live = torch.full((1, s), t, device="cuda", dtype=torch.int32) + got = indexer_logits_packed(q, w, pool, FP4_E8M0_B32, locs, 1, live) + reference = torch.einsum("bshd,btd->bsht", q, keys.unsqueeze(0)) + reference = (reference.relu() * w.unsqueeze(-1)).sum(dim=2).float() + torch.testing.assert_close(got, reference, atol=0, rtol=0) + # BF16 scores can tie at the cutoff; compare selected scores, not arbitrary tie order. + torch.testing.assert_close(got.topk(topk).values, reference.topk(topk).values, atol=0, rtol=0) + + +def test_full_range_worker_grid_covers_wide_histories(monkeypatch): + """More live tiles than worker programs: each worker strides over its share.""" + from freetoken.kernel.triton.dsv41 import indexer as mod + + b, s, t = 1, 2, 4096 # 64 tiles + pool, deq, q, w, k_rows, locs = _setup(b, s, t, seed=4) + live = torch.tensor([[4096, 4000]], device="cuda", dtype=torch.int32) + from types import SimpleNamespace + + properties = torch.cuda.get_device_properties(q.device) + with monkeypatch.context() as patch: + patch.setattr(torch.cuda, "get_device_properties", lambda _: SimpleNamespace( + multi_processor_count=5, major=properties.major, minor=properties.minor, + )) + got = mod.indexer_logits_packed(q, w, pool, FP4_E8M0_B32, locs, RATIO, live) + positions = torch.arange(t, device="cuda").view(1, 1, t).expand(b, s, t) + want = _reference(q, w, deq, k_rows, live, positions) + written = positions < live.unsqueeze(-1) + assert torch.equal(torch.isinf(got[written]), torch.isinf(want[written])) + finite = written & ~torch.isinf(want) + torch.testing.assert_close(got[finite], want[finite], atol=2e-2, rtol=2e-2) + + +def test_candidate_pool_matches_reference(): + from freetoken.kernel.triton.dsv41.indexer import indexer_logits_packed + + b, s, t, nc = 2, 3, 400, 96 + pool, deq, q, w, k_rows, locs = _setup(b, s, t, seed=2) + live = torch.tensor([[400, 350, 33], [1, 400, 200]], device="cuda", dtype=torch.int32) + cand = torch.randint(0, t, (b, s, nc), device="cuda", dtype=torch.int32) + cand[..., -10:] = -1 # empty tail of the pool + got = indexer_logits_packed(q, w, pool, FP4_E8M0_B32, locs, RATIO, live, candidates=cand) + want = _reference(q, w, deq, k_rows, live, cand) + assert got.shape == (b, s, nc) + assert torch.equal(torch.isinf(got), torch.isinf(want)) + finite = ~torch.isinf(want) + torch.testing.assert_close(got[finite], want[finite], atol=2e-2, rtol=2e-2) diff --git a/tests/kernels/test_dsv41_pack.py b/tests/kernels/test_dsv41_pack.py new file mode 100644 index 0000000000..176889dc7f --- /dev/null +++ b/tests/kernels/test_dsv41_pack.py @@ -0,0 +1,112 @@ +"""Packed KV row formats (``kvcache/dsv4/v41_row_format.py`` + ``kernel/triton/dsv41/pack.py``) against a torch +transcription of the reference quantizers in DeepSeek-V4.1's ``inference/kernel.py``: the round-tripped +value the pool hands back must equal what the reference bakes into its bf16 cache, bit for bit.""" + +from __future__ import annotations + +import pytest +import torch + +from freetoken.kvcache.dsv4.v41_row_format import BF16, FP4_E4M3_B16, FP4_E8M0_B32, FP8_E8M0_B32, RowFormat + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + +FP4_GRID = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]) + + +def _round_fp4(x: torch.Tensor) -> torch.Tensor: + """Round-to-nearest-even onto the signed e2m1 grid (the ``float4_e2m1fn`` cast).""" + a = x.abs() + grid = FP4_GRID.to(x.device) + # nearest grid point; ties go to the even-mantissa neighbour (0, 1, 2, 4 are even) + dist = (a[..., None] - grid).abs() + lo = torch.searchsorted(grid, a.reshape(-1).contiguous(), right=True).reshape(a.shape).clamp(1, 7) + below, above = grid[lo - 1], grid[lo] + mid = (below + above) / 2 + even_above = torch.isin(above, torch.tensor([0.0, 1.0, 2.0, 4.0], device=x.device)) + pick_above = (a > mid) | ((a == mid) & even_above) + r = torch.where(pick_above, above, below) + r = torch.where(a >= 6.0, torch.full_like(r, 6.0), r) + del dist + return torch.copysign(r, x) + + +def reference_roundtrip(x: torch.Tensor, fmt: RowFormat) -> torch.Tensor: + """``act_quant`` / ``fp4_act_quant`` with ``inplace=True``: quant + dequant back to bf16.""" + m, d = x.shape + if fmt is BF16: + return x.to(torch.bfloat16) + g = x.float().view(m, d // fmt.group, fmt.group) + amax = g.abs().amax(dim=-1) + if fmt is FP8_E8M0_B32: + s = torch.exp2(torch.ceil(torch.log2(amax.clamp_min(1e-4) / 448.0))) + q = (g / s[..., None]).clamp(-448.0, 448.0).to(torch.float8_e4m3fn).float() + elif fmt is FP4_E8M0_B32: + s = torch.exp2(torch.ceil(torch.log2(amax.clamp_min(6.0 * 2.0**-126) / 6.0))) + q = _round_fp4((g / s[..., None]).clamp(-6.0, 6.0)) + elif fmt is FP4_E4M3_B16: + s = (amax.clamp_min(6.0 * 2.0**-9) / 6.0).to(torch.float8_e4m3fn).float() + q = _round_fp4((g / s[..., None]).clamp(-6.0, 6.0)) + else: + raise AssertionError(fmt) + return (q * s[..., None]).view(m, d).to(torch.bfloat16) + + +@pytest.mark.parametrize("fmt", [BF16, FP8_E8M0_B32, FP4_E4M3_B16, FP4_E8M0_B32], ids=lambda f: f.name) +@pytest.mark.parametrize("dim", [128, 512]) +def test_roundtrip_matches_the_reference_quantizer(fmt: RowFormat, dim: int): + from freetoken.kernel.triton.dsv41.pack import pack_rows, unpack_rows + + torch.manual_seed(0) + m = 300 + # mixed magnitudes, exact zeros, a tiny group and a huge one + x = torch.randn(m, dim, device="cuda") * torch.exp2(torch.randint(-8, 6, (m, 1), device="cuda").float()) + x[7] = 0 + x[11, : fmt.group or 32] *= 1e-6 + x[13, -32:] = 400.0 + x = x.to(torch.bfloat16) + + packed = pack_rows(x, fmt) + assert packed.shape == (m, fmt.row_bytes(dim)) and packed.dtype == torch.uint8 + got = unpack_rows(packed, fmt, dim) + want = reference_roundtrip(x, fmt) + assert torch.equal(got, want), f"{fmt.name}: {(got.float() - want.float()).abs().max().item()}" + + +def test_scatter_into_pool_rows_and_gather_back(): + from freetoken.kernel.triton.dsv41.pack import pack_rows, unpack_rows + + torch.manual_seed(1) + fmt, dim, rows = FP4_E4M3_B16, 512, 64 + pool = torch.zeros((rows, fmt.row_bytes(dim)), dtype=torch.uint8, device="cuda") + x = torch.randn(5, dim, device="cuda", dtype=torch.bfloat16) + ids = torch.tensor([3, 60, 0, 17, 42], device="cuda", dtype=torch.int64) + assert pack_rows(x, fmt, pool=pool, row_ids=ids) is pool + got = unpack_rows(pool, fmt, dim, row_ids=ids.to(torch.int32)) + assert torch.equal(got, reference_roundtrip(x, fmt)) + untouched = torch.ones(rows, dtype=torch.bool, device="cuda") + untouched[ids] = False + assert int(pool[untouched].sum()) == 0 + + +def test_e4m3_encode_helper_matches_torch_bits(): + """``e4m3_f32_to_u8`` (the pre-sm_89 scale encoder) on every finite e4m3 code.""" + import triton + import triton.language as tl + + from freetoken.kernel.triton.e4m3_compat import e4m3_f32_to_u8 + + @triton.jit + def _enc(x_ptr, o_ptr, N: tl.constexpr): + offs = tl.arange(0, N) + tl.store(o_ptr + offs, e4m3_f32_to_u8(tl.load(x_ptr + offs))) + + codes = torch.arange(256, dtype=torch.uint8, device="cuda") + finite = (codes & 0x7F) != 0x7F # skip the NaN codes + vals = codes.view(torch.float8_e4m3fn).float() + vals = torch.where(finite, vals, torch.zeros_like(vals)) + out = torch.empty(256, dtype=torch.uint8, device="cuda") + _enc[(1,)](vals, out, N=256) + want = torch.where(finite, codes, torch.zeros_like(codes)) + # -0.0 encodes to 0x80 in torch and 0x80 here; +0.0 -> 0x00 + assert torch.equal(out, want) diff --git a/tests/kernels/test_dsv41_sparse_attn.py b/tests/kernels/test_dsv41_sparse_attn.py new file mode 100644 index 0000000000..8a4f5f2a4f --- /dev/null +++ b/tests/kernels/test_dsv41_sparse_attn.py @@ -0,0 +1,107 @@ +"""DSV41 sparse attention over packed pools against a torch reference on the DEQUANTIZED pools. + +``m == 1`` takes the split-k decode path, ``m > 1`` the single-program prefill one. Tolerance follows +the dsv4 sparse-attention tests (5e-2 on fp32-cast bf16): the online-softmax accumulation order differs +from the reference and split-k merges through log-sum-exp on top of that.""" + +from __future__ import annotations + +import pytest +import torch + +from freetoken.kvcache.dsv4.v41_row_format import BF16, FP4_E4M3_B16, FP8_E8M0_B32 + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + +D, H = 128, 8 +N_WINDOW = 128 # a multiple of BLOCK_T +N_WIN_SLOTS, N_CMP = 512, 1024 +TOL = dict(atol=5e-2, rtol=5e-2) + + +def _reference(q, win, cmp, sink, idx, n_window, scale, counts): + """Per (request, query, head) softmax over the live columns; ``sink`` is a null key.""" + b, m, h, d = q.shape + out = torch.zeros_like(q, dtype=torch.float32) + for i in range(b): + for j in range(m): + width = idx.shape[-1] if counts is None else n_window + int(counts[i, j]) + kv = [] + for c in range(min(width, idx.shape[-1])): + slot = int(idx[i, j, c]) + if slot >= 0: + kv.append((win if c < n_window else cmp)[slot].float()) + if not kv: + continue + kv_t = torch.stack(kv) + logits = q[i, j].float() @ kv_t.T * scale + mx = torch.maximum(logits.max(dim=-1).values, sink.float()) + probs = torch.exp(logits - mx[:, None]) + denom = probs.sum(dim=-1) + torch.exp(sink.float() - mx) + out[i, j] = (probs @ kv_t) / denom[:, None] + return out + + +def _pools(win_fmt, cmp_fmt, seed): + from freetoken.kernel.triton.dsv41.pack import pack_rows, unpack_rows + + torch.manual_seed(seed) + win_vals = torch.randn(N_WIN_SLOTS, D, device="cuda", dtype=torch.bfloat16) + cmp_vals = torch.randn(N_CMP, D, device="cuda", dtype=torch.bfloat16) + win_pool = pack_rows(win_vals, win_fmt) + cmp_pool = pack_rows(cmp_vals, cmp_fmt) + return win_pool, cmp_pool, unpack_rows(win_pool, win_fmt, D), unpack_rows(cmp_pool, cmp_fmt, D) + + +def _topk(b, m, n_cmp_cols, seed, ring_fill=None): + torch.manual_seed(seed) + win = torch.randint(0, N_WIN_SLOTS, (b, m, N_WINDOW), device="cuda", dtype=torch.int32) + if ring_fill is not None: # a ring still filling: trailing window columns empty + win[..., ring_fill:] = -1 + cmp = torch.randint(0, N_CMP, (b, m, n_cmp_cols), device="cuda", dtype=torch.int32) + cmp[..., ::7] = -1 # scattered unreachable picks + return torch.cat([win, cmp], dim=-1) + + +@pytest.mark.parametrize("m", [1, 3]) +@pytest.mark.parametrize("fmts", [(FP8_E8M0_B32, FP4_E4M3_B16), (BF16, BF16)], ids=["fp8+fp4", "bf16"]) +def test_matches_reference_on_dequantized_pools(m, fmts): + from freetoken.kernel.triton.dsv41.sparse_attn import sparse_attn_packed + + win_fmt, cmp_fmt = fmts + b, n_cmp_cols = 2, 256 + win_pool, cmp_pool, win_deq, cmp_deq = _pools(win_fmt, cmp_fmt, seed=1) + q = torch.randn(b, m, H, D, device="cuda", dtype=torch.bfloat16) + sink = torch.randn(H, device="cuda") * 0.5 + idx = _topk(b, m, n_cmp_cols, seed=2, ring_fill=100) + scale = D**-0.5 + + got = sparse_attn_packed(q, win_pool, win_fmt, cmp_pool, cmp_fmt, sink, idx, N_WINDOW, scale) + want = _reference(q, win_deq, cmp_deq, sink, idx, N_WINDOW, scale, None) + torch.testing.assert_close(got.float(), want, **TOL) + + +def test_cmp_counts_bound_the_visited_columns(): + from freetoken.kernel.triton.dsv41.sparse_attn import sparse_attn_packed + + b, m, n_cmp_cols = 3, 1, 512 + win_pool, cmp_pool, win_deq, cmp_deq = _pools(FP8_E8M0_B32, FP4_E4M3_B16, seed=3) + q = torch.randn(b, m, H, D, device="cuda", dtype=torch.bfloat16) + sink = torch.zeros(H, device="cuda") + idx = _topk(b, m, n_cmp_cols, seed=4) + counts = torch.tensor([[0], [37], [512]], device="cuda", dtype=torch.int32) + scale = D**-0.5 + + got = sparse_attn_packed(q, win_pool, FP8_E8M0_B32, cmp_pool, FP4_E4M3_B16, sink, idx, N_WINDOW, scale, cmp_counts=counts) + want = _reference(q, win_deq, cmp_deq, sink, idx, N_WINDOW, scale, counts) + torch.testing.assert_close(got.float(), want, **TOL) + + +def test_all_masked_query_yields_zeros(): + from freetoken.kernel.triton.dsv41.sparse_attn import sparse_attn_packed + + win_pool, cmp_pool, _, _ = _pools(FP8_E8M0_B32, FP4_E4M3_B16, seed=5) + q = torch.randn(1, 2, H, D, device="cuda", dtype=torch.bfloat16) + idx = torch.full((1, 2, N_WINDOW + 64), -1, device="cuda", dtype=torch.int32) + out = sparse_attn_packed(q, win_pool, FP8_E8M0_B32, cmp_pool, FP4_E4M3_B16, torch.zeros(H, device="cuda"), idx, N_WINDOW, 0.1) + assert torch.equal(out, torch.zeros_like(out)) diff --git a/tests/kernels/test_dsv41_topk.py b/tests/kernels/test_dsv41_topk.py new file mode 100644 index 0000000000..e7fc66f8d8 --- /dev/null +++ b/tests/kernels/test_dsv41_topk.py @@ -0,0 +1,159 @@ +"""dsv41_topk / dsv41_candidate_blocks against an independent torch reference with the documented tie +rule, over the contract's edge cases: empty and short histories, unequal live counts in one batch, +-inf holes, widths crossing the slice and level boundaries, exact ties, newest-block inclusion, +the sorted-valid-prefix / -1-tail output layout, and CUDA-graph replay across changing live lengths +(stale workspace reads would surface as wrong picks).""" + +from __future__ import annotations + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + + +def ref_topk(scores: torch.Tensor, live: torch.Tensor, k: int) -> torch.Tensor: + """(score desc, column asc) over the live, finite columns; ascending columns, -1 padded.""" + R, T = scores.shape + out = torch.full((R, k), -1, dtype=torch.int32) + for r in range(R): + n = int(live[r]) + s = scores[r, :n].double().cpu() + cols = torch.arange(n) + alive = torch.isfinite(s) | (s == float("inf")) + cand = [(float(-s[c]), int(c)) for c in cols[alive]] + cand.sort() + picked = sorted(c for _, c in cand[:k]) + out[r, : len(picked)] = torch.tensor(picked, dtype=torch.int32) + return out.to(scores.device) + + +def ref_candidates(scores, live, topk_blocks, block_size): + R, T = scores.shape + nb = -(-T // block_size) + out = torch.full((R, topk_blocks * block_size), -1, dtype=torch.int32) + for r in range(R): + n = int(live[r]) + if n == 0: + continue + nbl = -(-n // block_size) + s = torch.nn.functional.pad(scores[r, :n].float().cpu(), (0, nbl * block_size - n), value=float("-inf")).view(nbl, block_size).amax(-1) + s[nbl - 1] = float("inf") + cand = sorted((float(-s[b]), b) for b in range(nbl) if s[b] > float("-inf")) + keep = sorted(b for _, b in cand[:topk_blocks]) + pos = [p for b in keep for p in range(b * block_size, (b + 1) * block_size) if p < n] + out[r, : len(pos)] = torch.tensor(pos, dtype=torch.int32) + return out.to(scores.device) + + +def _scores(R, T, seed=0, ties=False): + g = torch.Generator(device="cpu").manual_seed(seed) + s = torch.randn(R, T, generator=g) + if ties: + s = torch.randint(0, 5, (R, T), generator=g).float() # heavy exact ties + return s.cuda() + + +@pytest.mark.parametrize("T,k", [(64, 8), (4096, 512), (4097, 512), (8192, 512), (100_000, 512), (262_144, 512), (1_048_576, 512), (131_072, 2048)]) +def test_topk_matches_reference_across_widths(T, k): + from freetoken.kernel.triton.dsv41.topk import dsv41_topk + + R = 3 + scores = _scores(R, T, seed=T) + live = torch.tensor([0, min(T, 5000), T], dtype=torch.int32, device="cuda")[:R] + # holes inside the live range + scores[2, ::7] = float("-inf") + got = dsv41_topk(scores, live, k) + want = ref_topk(scores, live, k) + assert torch.equal(got, want), (got[:, :8], want[:, :8]) + + +def test_topk_short_and_unequal_rows_and_exact_ties(): + from freetoken.kernel.triton.dsv41.topk import dsv41_topk + + T, k = 20_000, 512 + scores = _scores(5, T, seed=3, ties=True) + live = torch.tensor([0, 1, 300, 4096, 20_000], dtype=torch.int32, device="cuda") + got = dsv41_topk(scores, live, k) + want = ref_topk(scores, live, k) + assert torch.equal(got, want) + # the layout contract: a sorted valid prefix, then -1 + for r in range(5): + row = got[r].tolist() + n = sum(1 for v in row if v >= 0) + assert row[n:] == [-1] * (k - n) and row[:n] == sorted(row[:n]) and n == min(int(live[r]), k) + + +def test_topk_all_dead_row_and_fewer_finite_than_k(): + from freetoken.kernel.triton.dsv41.topk import dsv41_topk + + scores = torch.full((2, 10_000), float("-inf"), device="cuda") + scores[1, [10, 4095, 4096, 9000]] = torch.tensor([1.0, 2.0, 3.0, 0.5], device="cuda") + live = torch.tensor([10_000, 9001], dtype=torch.int32, device="cuda") + got = dsv41_topk(scores, live, 8) + assert (got[0] == -1).all() + assert got[1].tolist() == [10, 4095, 4096, 9000, -1, -1, -1, -1] + + +@pytest.mark.parametrize("T", [4000, 200_000]) +def test_candidate_blocks_match_reference_and_pin_the_newest_block(T): + from freetoken.kernel.triton.dsv41.topk import dsv41_candidate_blocks + + R, KB, BS = 4, 2048, 8 + scores = _scores(R, T, seed=T + 1) + scores[0, :] = -1000.0 # the newest block is kept even when every score is poor + live = torch.tensor([13, 8, min(T, 16_000), T], dtype=torch.int32, device="cuda") + got = dsv41_candidate_blocks(scores, live, KB, BS) + want = ref_candidates(scores, live, KB, BS) + assert got.shape == (R, KB * BS) + assert torch.equal(got, want) + # short histories: every live position is a candidate (the implicit fast path) + assert got[0].tolist()[:13] == list(range(13)) and (got[0, 13:] == -1).all() + assert got[1].tolist()[:8] == list(range(8)) and (got[1, 8:] == -1).all() + + +def test_topk_under_graph_replay_across_live_lengths(): + """Capture once at the staged width, replay with live lengths that move up and down: every replay + must equal the reference (a stale candidate slot or live count from a previous replay would not).""" + from freetoken.kernel.triton.dsv41.topk import dsv41_candidate_blocks, dsv41_topk + + T, k, R = 300_000, 512, 2 + scores = torch.empty((R, T), device="cuda") + live = torch.empty((R,), dtype=torch.int32, device="cuda") + out = torch.empty((R, k), dtype=torch.int32, device="cuda") + cand = torch.empty((R, 2048 * 8), dtype=torch.int32, device="cuda") + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + scores.copy_(_scores(R, T, seed=9)) + live.fill_(T) + dsv41_topk(scores, live, k, out) # warm-up compiles + cand.copy_(dsv41_candidate_blocks(scores, live, 2048, 8)) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + dsv41_topk(scores, live, k, out) + cand.copy_(dsv41_candidate_blocks(scores, live, 2048, 8)) + torch.cuda.synchronize() + for step, (l0, l1) in enumerate([(T, 7), (4096, 4097), (0, 1), (250_000, 12_000), (3, T), (16_384, 16_385)]): + scores.copy_(_scores(R, T, seed=100 + step)) + live.copy_(torch.tensor([l0, l1], dtype=torch.int32)) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, ref_topk(scores, live, k)), (step, l0, l1) + assert torch.equal(cand, ref_candidates(scores, live, 2048, 8)), (step, l0, l1) + + +def test_topk_plan_terminates_or_rejects(): + """Every accepted k converges (each level at least halves the row); k past CHUNK // 2 -- which + would stall the planner (k == CHUNK never shrinks, 3072 stalls at 6144) -- is rejected up front.""" + from freetoken.kernel.triton.dsv41.topk import CHUNK, FINAL_MAX, MAX_K, topk_plan + + for k in (1, 8, 512, 2048, MAX_K): + for columns in (1, 4096, 4097, 1 << 20, (1 << 20) + 3): + levels = topk_plan(columns, k) + widths = [w for w, _ in levels] + [(levels[-1][1] * k) if levels else columns] + assert all(a > b for a, b in zip(widths, widths[1:])) and widths[-1] <= FINAL_MAX + assert len(levels) <= 12 + for k in (MAX_K + 1, 2049, 3072, CHUNK, 0, -1): + with pytest.raises(ValueError): + topk_plan(1 << 20, k) diff --git a/tests/kernels/test_dsv4_fp8_block_linear.py b/tests/kernels/test_dsv4_fp8_block_linear.py new file mode 100644 index 0000000000..2194a65e1b --- /dev/null +++ b/tests/kernels/test_dsv4_fp8_block_linear.py @@ -0,0 +1,101 @@ +"""``kernel/triton/dsv4/fp8_linear.block_fp8_linear`` against a torch transcription of the reference +``act_quant`` + ``fp8_gemm`` (ue8m0 activation scales, ``block x block`` e8m0 weight scales), for the +DeepSeek-V4 block (128) and the DeepSeek-V4.1 block (32), on the GEMV (M=1) and GEMM (M>1) paths.""" + +from __future__ import annotations + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + +FP8 = torch.float8_e4m3fn +E8M0 = torch.float8_e8m0fnu + + +def _pow2_ceil_scale(amax: torch.Tensor, qmax: float) -> torch.Tensor: + """``2 ** ceil(log2(amax / qmax))`` -- the reference ``fast_round_scale``.""" + return torch.exp2(torch.ceil(torch.log2(amax / qmax))) + + +def quantize_weight(w: torch.Tensor, block: int) -> tuple[torch.Tensor, torch.Tensor]: + """Per ``block x block`` ue8m0-scaled e4m3 weight, scale as e8m0 codes.""" + n, k = w.shape + blocks = w.float().view(n // block, block, k // block, block).permute(0, 2, 1, 3) + amax = blocks.abs().amax(dim=(-1, -2)).clamp_min(1e-4) + scale = _pow2_ceil_scale(amax, 448.0) # [n//block, k//block] + q = (blocks / scale[..., None, None]).clamp(-448.0, 448.0).to(FP8) + w_fp8 = q.permute(0, 2, 1, 3).reshape(n, k) + return w_fp8, scale.to(E8M0) + + +def act_quant_reference(x: torch.Tensor, block: int) -> tuple[torch.Tensor, torch.Tensor]: + m, k = x.shape + groups = x.float().view(m, k // block, block) + amax = groups.abs().amax(dim=-1).clamp_min(1e-4) + scale = _pow2_ceil_scale(amax, 448.0) + q = (groups / scale[..., None]).clamp(-448.0, 448.0).to(FP8) + return q.view(m, k), scale + + +def linear_reference(x: torch.Tensor, w_fp8: torch.Tensor, w_scale: torch.Tensor, block: int) -> torch.Tensor: + m, k = x.shape + n = w_fp8.shape[0] + a_fp8, a_scale = act_quant_reference(x, block) + a = a_fp8.float().view(m, k // block, block) + w = w_fp8.float().view(n // block, block, k // block, block) + ws = w_scale.float() # [n//block, k//block] + out = torch.zeros(m, n, dtype=torch.float32, device=x.device) + for kb in range(k // block): + wk = w[:, :, kb, :].reshape(n, block) # [n, block] + partial = a[:, kb, :] @ wk.t() # [m, n] + out += partial * a_scale[:, kb, None] * ws[:, kb].repeat_interleave(block)[None, :] + return out + + +@pytest.mark.parametrize("block", [32, 128]) +@pytest.mark.parametrize("m", [1, 2, 5, 8, 9, 64]) +def test_block_fp8_linear_matches_reference(block: int, m: int): + from freetoken.kernel.triton.dsv4.fp8_linear import block_fp8_linear + + torch.manual_seed(0) + # K // block = 5 slabs exercises the GEMV split that must divide the slab count; N spans two GEMM N-tiles + k, n = 5 * block, 256 if block == 128 else 224 + x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) * 3 + w_fp8, w_scale = quantize_weight(torch.randn(n, k, device="cuda", dtype=torch.bfloat16), block) + + got = block_fp8_linear(x, w_fp8, w_scale, block=block).float() + want = linear_reference(x, w_fp8, w_scale, block) + # bf16 output; the fp32 reduction order inside the MMA may differ from the slab loop + torch.testing.assert_close(got, want, rtol=2e-2, atol=2e-2 * want.abs().max().item()) + + +def test_block_argument_must_match_the_scale_shape(): + from freetoken.kernel.triton.dsv4.fp8_linear import block_fp8_linear + + x = torch.randn(4, 256, device="cuda", dtype=torch.bfloat16) + w_fp8, w_scale = quantize_weight(torch.randn(256, 256, device="cuda", dtype=torch.bfloat16), 32) + with pytest.raises(AssertionError): + block_fp8_linear(x, w_fp8, w_scale, block=128) + + +@pytest.mark.parametrize("block", [32, 128]) +def test_grouped_w8a16_gemv_matches_the_dequantized_einsum(block): + """``wo_a`` at decode: the fp8 payload dequantized in-kernel against a bf16 activation must equal + the reference's bf16 einsum over the dequantized weight (fp32 accumulation, order aside).""" + from freetoken.kernel.triton.dsv4.fp8_linear import dequant_block_fp8, grouped_w8a16_gemv + + G, r, d = 4, 256, 1024 + torch.manual_seed(block) + w = torch.randn(G * r, d, device="cuda") * 0.05 + w_fp8, scale = quantize_weight(w, block) + x = torch.randn(G, d, device="cuda", dtype=torch.bfloat16) + deq = dequant_block_fp8(w_fp8, scale, block=block) + ref_deq = (w_fp8.float() * torch.exp2(scale.view(torch.uint8).float() - 127.0).repeat_interleave(block, 0).repeat_interleave(block, 1)).to(torch.bfloat16) + assert torch.equal(deq, ref_deq) + for t in (1, 3): + xt = torch.randn(t, G, d, device="cuda", dtype=torch.bfloat16) if t > 1 else x.unsqueeze(0) + want = torch.einsum("tgd,grd->tgr", xt, ref_deq.view(G, r, d)).flatten(1) + got = grouped_w8a16_gemv(xt, w_fp8, scale, block=block) + assert got.shape == (t, G * r) and got.dtype == torch.bfloat16 + torch.testing.assert_close(got.float(), want.float(), atol=2e-2, rtol=2e-2) diff --git a/tests/kernels/test_dsv4_router_select.py b/tests/kernels/test_dsv4_router_select.py new file mode 100644 index 0000000000..9829ed038a --- /dev/null +++ b/tests/kernels/test_dsv4_router_select.py @@ -0,0 +1,52 @@ +"""Fused gate selection (``kernel/triton/dsv4/router``) against the torch expression it replaces, for +every score function, with and without top-k weight normalization, at decode and prefill widths.""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + + +def reference(scores, bias, topk, score_func, normalize, route_scale): + if score_func == "softmax": + s = scores.softmax(dim=-1) + elif score_func == "sigmoid": + s = scores.sigmoid() + else: + s = F.softplus(scores).sqrt() + ids = (s + bias).topk(topk, dim=-1)[1] + w = s.gather(1, ids) + if normalize and topk > 1: + w = w / (w.sum(dim=-1, keepdim=True) + 1e-20) + return w * route_scale, ids + + +@pytest.mark.parametrize("score_func", ["sqrtsoftplus", "sigmoid", "softmax"]) +@pytest.mark.parametrize("normalize", [True, False]) +@pytest.mark.parametrize("T", [1, 3, 300]) +def test_router_select_matches_torch(score_func, normalize, T): + from freetoken.kernel.triton.dsv4.router import router_select + + g = torch.Generator(device="cpu").manual_seed(T) + E, K = 384, 6 + scores = (torch.randn(T, E, generator=g) * 3).cuda() + bias = (torch.randn(E, generator=g) * 0.1).cuda() + w, ids = router_select(scores, bias, K, score_func=score_func, normalize=normalize, route_scale=2.5) + w_ref, ids_ref = reference(scores, bias, K, score_func, normalize, 2.5) + assert torch.equal(ids.sort(dim=-1).values, ids_ref.to(torch.int32).sort(dim=-1).values) + # same expert set; compare the weights per expert (selection order may differ only on exact ties) + order = ids.argsort(dim=-1) + order_ref = ids_ref.argsort(dim=-1) + torch.testing.assert_close(w.gather(1, order), w_ref.gather(1, order_ref), atol=1e-6, rtol=1e-5) + + +def test_router_select_breaks_ties_toward_the_lowest_expert(): + from freetoken.kernel.triton.dsv4.router import router_select + + scores = torch.zeros(1, 384, device="cuda") # every expert ties + bias = torch.zeros(384, device="cuda") + _, ids = router_select(scores, bias, 6, score_func="sigmoid", normalize=True, route_scale=1.0) + assert ids.tolist() == [[0, 1, 2, 3, 4, 5]] diff --git a/tests/kvcache/test_cache_unit_bytes.py b/tests/kvcache/test_cache_unit_bytes.py index 496622f985..cce5df6246 100644 --- a/tests/kvcache/test_cache_unit_bytes.py +++ b/tests/kvcache/test_cache_unit_bytes.py @@ -16,7 +16,7 @@ compute_cache_status_meta, compute_cache_unit_bytes, ) -from freetoken.kvcache.dsv4_cost_model import _dsv4_window_floor_pages +from freetoken.kvcache.dsv4.v4_cost_model import _dsv4_window_floor_pages def _init_tp() -> None: @@ -204,7 +204,7 @@ def test_floors_dsv4_reports_real_window_floor(): dsv4_args=SimpleNamespace(window_size=P), has_swa_attention=False ), ) - from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache + from freetoken.kvcache.dsv4.v4_pool import DSV4PagedKVCache eng = SimpleNamespace(config=cfg, moe_offload_cache=None, linear_state_pool=None) eng.kv_cache = object.__new__(DSV4PagedKVCache) # min_kv_tokens is a classmethod diff --git a/tests/kvcache/test_dsv41_pool.py b/tests/kvcache/test_dsv41_pool.py new file mode 100644 index 0000000000..f41df85174 --- /dev/null +++ b/tests/kvcache/test_dsv41_pool.py @@ -0,0 +1,205 @@ +"""DSV41 paged KV pool + cost model (CPU, no model): args validation, per-source tiers, sizing / +byte accounting, the window free-list duck-type, full-loc translation; packed writes on CUDA.""" + +from __future__ import annotations + +import pytest +import torch + +from freetoken.kvcache.dsv4.v41_cost_model import ( + dsv41_cache_per_page, + dsv41_kv_unit_bytes, + dsv41_pool_bytes, + dsv41_pool_sizes, + dsv41_solve_num_pages, + dsv41_window_unit_bytes, +) +from freetoken.kvcache.dsv4.v41_pool import DSV41PagedKVCache +from freetoken.kvcache.dsv4.v41_row_format import FP4_E4M3_B16, FP4_E8M0_B32, FP8_E8M0_B32 +from freetoken.models.deepseek_v41.args import DeepseekV41Args + +DEVICE = torch.device("cpu") +P = 128 +# a 10-layer miniature of V4.1's layout: 2 window-only, ratio-2 encoder (source 2, source 5), ratio-1 decoder (source 7) +RATIOS = (0, 0, 2, 2, 2, 2, 2, 1, 1, 1) +SOURCES = (2, 5, 7) + + +def _args(**over) -> DeepseekV41Args: + base = dict(n_layers=10, head_dim=512, index_head_dim=128, window_size=P, compress_ratios=RATIOS, + kv_source_layers=SOURCES, swa_decoder_replay="exact") + base.update(over) + return DeepseekV41Args(**base) + + +def _pool(num_pages=8, swa_ratio=0.5, n_scratch=1, **over): + args = _args(**over) + sizes = dsv41_pool_sizes(num_pages, args, swa_ratio, P) + return DSV41PagedKVCache(sizes, args, DEVICE, n_scratch=n_scratch), sizes, args + + +def test_window_control_uses_pool_units_and_restores_concrete_capacity(): + from types import SimpleNamespace + + from freetoken.attention import AttnType + from freetoken.kvcache.cache_status import ( + _supports_swa_ratio, compute_cache_floors, compute_cache_pools, + ) + from freetoken.kvcache.dsv4.v41_cost_model import _dsv41_pool_sizes + from freetoken.scheduler.scheduler import Scheduler + from freetoken.server.api_server import CacheRebuildRequest, _resolve_num_swa_pages, cache_geometry + from freetoken.server.stats import _swa_page_size + + pool, _, args = _pool(num_pages=64) + config = SimpleNamespace( + page_size=P, max_running_req=1, max_seq_len=1024, cache_type="swa_radix", + swa_full_tokens_ratio=0.5, swa_num_pages_override=None, + model_config=SimpleNamespace( + dsv4_args=None, dsv41_args=_args(), has_swa_attention=False, + kv_cache_group_specs=lambda: [SimpleNamespace(attn_type=AttnType.DSV41)], + ), + ) + engine = SimpleNamespace( + config=config, kv_cache=pool, num_pages=63, + moe_offload_cache=None, linear_state_pool=None, + ) + assert _supports_swa_ratio(config) + assert compute_cache_pools(engine)["swa_page_size"] == P + assert compute_cache_pools(engine)["num_swa_pages"] == 31 + floor = _dsv41_pool_sizes(config, 64, num_swa_pages=1).n_win_pages - 1 + assert compute_cache_floors(engine)["swa_tokens"] == floor * P + prior = Scheduler._current_cache_geometry(SimpleNamespace(engine=engine, config=config)) + assert prior["num_swa_pages"] == 31 + state = SimpleNamespace(config=config, cache_pools=compute_cache_pools(engine)) + req = CacheRebuildRequest(num_pages=64, swa_full_tokens_ratio=0.5) + assert _resolve_num_swa_pages(state, req) == 32 + assert _swa_page_size(config) == P + state.stats = SimpleNamespace(kv_total_pages=63, mamba_total_slots=0) + state.last_rebuild = {"num_pages": 64, "num_swa_pages": 32} + assert cache_geometry(state)["swa_full_tokens_ratio"] == 0.5 + pool._init_paged_state(1, True) + pool.rebuild_from_config(config, 63, num_swa_pages=24) + assert pool.window_pages == 24 + pool.rebuild_from_config(config, prior["num_pages"], num_swa_pages=prior["num_swa_pages"]) + assert pool.window_pages == 31 + + +def test_pool_reads_sources_and_rings_from_args(): + pool, _, _ = _pool() + assert [pool.source_of(l) for l in range(2, 10)] == [2, 2, 2, 5, 5, 7, 7, 7] + assert set(pool.state_ring) == {2, 5} and pool.state_ring[2].ring_size == 2 + with pytest.raises(ValueError): # the window is not a multiple of a compress ratio + _pool(window_size=P - 1) + with pytest.raises(ValueError): # an indexer key the fp4 row format cannot hold + _pool(index_head_dim=48) + + +@pytest.mark.parametrize("replay", ["bounded", "exact"]) +def test_v41_flash_row_bytes_and_private_window_layers(replay): + """The tech report's headline: 3 ratio-2 encoder sources + 1 ratio-1 decoder source = 890 B/token.""" + ratios = (0, 0) + (2,) * 18 + (1,) * 20 + args = DeepseekV41Args(n_layers=40, head_dim=512, index_head_dim=128, window_size=128, compress_ratios=ratios, + kv_source_layers=(2, 8, 14, 20), swa_decoder_replay=replay) + pool = DSV41PagedKVCache(dsv41_pool_sizes(4, args, 1.0, 128), args, DEVICE) + win, main, idx = pool.window_pool[0].shape[1], pool.main_pool[2].shape[1], pool.idx_pool[2].shape[1] + assert (win, main, idx) == (528, 288, 68) + assert pool.state_ring[2].buffer.shape[1] * 4 == 4096 + assert sum((main + idx) // pool.ratio_of(s) for s in args.backbone_kv_sources) == 890 + assert dsv41_kv_unit_bytes(args, 128) == 890 + 8 # + the int64 full->window map slot + assert pool.private_window_layer_ids == (tuple(range(20, 40)) if replay == "bounded" else ()) + + +def test_pool_tiers_per_source_and_aliasing(): + pool, sizes, args = _pool() + assert len(pool.window_pool) == 10 and pool.window_pool[0].shape == (sizes.n_win_slots, 528) + assert set(pool.main_pool) == set(SOURCES) == set(pool.idx_pool) + assert set(pool.state_ring) == {2, 5} # the ratio-1 decoder source carries no partial group + assert pool.main_pool[2].shape == (sizes.full_token // 2 + 1, 288) + assert pool.main_pool[7].shape == (sizes.full_token + 1, 288) + assert pool.idx_pool[5].shape == (sizes.full_token // 2 + 1, 68) + assert pool.main_pool_of(4) is pool.main_pool[2] and pool.idx_pool_of(9) is pool.idx_pool[7] + assert pool.scratch_base[2] == sizes.full_token // 2 + with pytest.raises(AssertionError): + pool.main_pool_of(0) + ring = pool.state_ring[2] + assert ring.ring_size == 2 and ring.item_size == 512 and ring.buffer.shape == (sizes.state_slots[2] + 1, 1024) + assert torch.all(ring.buffer[-1, :512] == 0) and torch.all(torch.isneginf(ring.buffer[-1, 512:])) + + +def test_pool_bytes_match_allocation_and_solver_respects_budget(): + pool, sizes, args = _pool(num_pages=16, swa_ratio=0.25) + assert pool.total_bytes() == dsv41_pool_bytes(sizes, args, n_scratch=1) + assert dsv41_cache_per_page(args, 0.25, P) > 0 + assert dsv41_window_unit_bytes(args, P) == -(-(10 * P * 528 + 2 * 2 * 4096) // P) + budget = 64 << 20 + solved = dsv41_solve_num_pages(budget, args, 0.25, floor_win_pages=4, P=P, n_scratch=3) + assert dsv41_pool_bytes(solved, args, 3) <= budget + # one more page, sized the way the solver sizes (window = max(floor, ceil(ratio * pages))), overflows + more = solved.full_token // P + 1 + bigger = dsv41_pool_sizes(more, args, 0.25, P, n_win_pages=max(4, (round(0.25 * more * P) + P - 1) // P)) + assert dsv41_pool_bytes(bigger, args, 3) > budget + assert solved.n_win_pages >= 4 + with pytest.raises(ValueError): + dsv41_solve_num_pages(1 << 10, args, 0.25, floor_win_pages=4, P=P) + + +def test_translation_state_loc_and_cmp_rows(): + pool, sizes, args = _pool(num_pages=16) + pool.bind_window_pages(full_page_base=0, window_page_base=2 * P) + pool.bind_window_pages(full_page_base=3 * P, window_page_base=0) + assert pool.translate_full_to_window(torch.tensor([0, 1, 127])).tolist() == [2 * P, 2 * P + 1, 2 * P + 127] + assert pool.translate_full_to_window(torch.tensor([3 * P + 5])).item() == 5 + assert pool.translate_full_to_window(torch.tensor([P + 7, -1])).tolist() == [-1, -1] + full = torch.tensor([0, 1, 2, 127, 128, 4 * P - 1]) + assert pool.cmp_rows(full, 2).tolist() == [0, 0, 1, 63, 64, 2 * P - 1] + assert pool.cmp_rows(full, 1).tolist() == full.tolist() + assert pool.cmp_rows(torch.tensor([-1]), 2).item() < 0 + top = pool.cmp_rows(torch.tensor([sizes.full_token - 1]), 2).item() + assert top < pool.scratch_base[2] + ws = pool.translate_full_to_window(torch.tensor([127, 3 * P])) + assert DSV41PagedKVCache.state_loc(ws, 2, P).tolist() == [2 * 2 + 1, 0] + assert DSV41PagedKVCache.state_loc(torch.tensor([-1]), 2, P).item() == -1 + + +def _expand(bases): + return (torch.tensor(bases, dtype=torch.int64)[:, None] + torch.arange(P)).flatten() + + +def test_swa_duck_type_alloc_free_and_dummy(): + pool, sizes, _ = _pool(num_pages=8) + pool._init_paged_state(max_running_req=2, radix=True) + cap = sizes.n_win_slots - P + assert pool.swa_available_size() == cap and pool.swa_num_tokens - 1 == cap + assert pool.sliding_window_size == P and pool.prefill_chunk_budget >= P + pool.alloc_swa(_expand([0, 2 * P])) + assert pool.swa_available_size() == cap - 2 * P + ws = pool.translate_loc_from_full_to_swa(torch.arange(2 * P, 3 * P)) + assert (ws >= 0).all() and int(ws[0]) % P == 0 and torch.equal(ws - ws[0], torch.arange(P)) + pool.free_swa(_expand([0])) + pool.free_swa(_expand([0])) # idempotent + assert pool.swa_available_size() == cap - P + with pytest.raises(AssertionError): + pool.free_swa(torch.arange(2 * P, 2 * P + 5)) + dummy = pool.translate_loc_from_full_to_swa(torch.arange(sizes.full_token - P, sizes.full_token)) + assert int(dummy[0]) == sizes.n_win_slots - P + assert pool.k_cache(3) is pool.window_pool[3] and pool.num_layers == 10 and pool.unit_bytes()[0] > 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="packed writes run through triton") +def test_packed_writes_round_trip_on_cuda(): + from kernels.test_dsv41_pack import reference_roundtrip # tests/ is on sys.path under pytest + + args = _args() + sizes = dsv41_pool_sizes(4, args, 0.5, P) + pool = DSV41PagedKVCache(sizes, args, torch.device("cuda"), n_scratch=2) + kv = torch.randn(3, 512, device="cuda", dtype=torch.bfloat16) + slots = torch.tensor([0, 130, 5], device="cuda") + pool.store_window(kv, 4, slots) + assert torch.equal(pool.read_window(4, slots), reference_roundtrip(kv, FP8_E8M0_B32)) + pool.store_main(kv, 2, slots) + assert torch.equal(pool.read_main(2, slots), reference_roundtrip(kv, FP4_E4M3_B16)) + k = torch.randn(3, 128, device="cuda", dtype=torch.bfloat16) + pool.store_index(k, 7, slots) + assert torch.equal(pool.read_index(7, slots), reference_roundtrip(k, FP4_E8M0_B32)) + pool.store_kv(kv, kv, slots, 0) + assert torch.equal(pool.read_window(0, slots), reference_roundtrip(kv, FP8_E8M0_B32)) diff --git a/tests/kvcache/test_dsv4_free_list.py b/tests/kvcache/test_dsv4_free_list.py index a3ae8d0cd7..7b0dfb5566 100644 --- a/tests/kvcache/test_dsv4_free_list.py +++ b/tests/kvcache/test_dsv4_free_list.py @@ -5,7 +5,7 @@ import pytest import torch -from freetoken.kvcache.dsv4_paged_pool import FreeListAllocator +from freetoken.kvcache.dsv4.v4_pool import FreeListAllocator DEVICE = torch.device("cpu") P = 128 diff --git a/tests/kvcache/test_dsv4_pool.py b/tests/kvcache/test_dsv4_pool.py index c345b67c82..c8e21997ed 100644 --- a/tests/kvcache/test_dsv4_pool.py +++ b/tests/kvcache/test_dsv4_pool.py @@ -10,8 +10,8 @@ import pytest import torch -from freetoken.kvcache.dsv4_cost_model import dsv4_pool_sizes -from freetoken.kvcache.dsv4_paged_pool import CompressStateRing, DSV4PagedKVCache +from freetoken.kvcache.dsv4.v4_cost_model import dsv4_pool_sizes +from freetoken.kvcache.dsv4.v4_pool import CompressStateRing, DSV4PagedKVCache from freetoken.models.deepseek_v4.args import DeepseekV4Args DEVICE = torch.device("cpu") diff --git a/tests/kvcache/test_kv_cache_rebuild.py b/tests/kvcache/test_kv_cache_rebuild.py index a5ec94a084..0e3b7050e8 100644 --- a/tests/kvcache/test_kv_cache_rebuild.py +++ b/tests/kvcache/test_kv_cache_rebuild.py @@ -178,13 +178,13 @@ def test_dsv4_rebuild_from_config_builds_the_pool_sizes_and_attaches_the_table() the engine used to do -- and re-points full_loc_map at the shared page table.""" from types import SimpleNamespace - from freetoken.kvcache.dsv4_cost_model import _dsv4_pool_sizes - from freetoken.kvcache.dsv4_cost_model import ( + from freetoken.kvcache.dsv4.v4_cost_model import _dsv4_pool_sizes + from freetoken.kvcache.dsv4.v4_cost_model import ( dsv4_kv_unit_bytes, dsv4_pool_sizes, dsv4_window_unit_bytes, ) - from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache + from freetoken.kvcache.dsv4.v4_pool import DSV4PagedKVCache from freetoken.models.deepseek_v4.args import DeepseekV4Args P, mrr = 128, 1 @@ -222,7 +222,7 @@ def test_dsv4_refresh_seq_state_tracks_page_table_width(): from freetoken.engine.engine import Engine from freetoken.utils import align_ceil - from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache + from freetoken.kvcache.dsv4.v4_pool import DSV4PagedKVCache from freetoken.scheduler.table import TableManager P, mrr = 128, 4 @@ -264,7 +264,7 @@ def test_every_kv_pool_answers_the_sizing_surface(): the model resolved to. DSV4 must carry its OWN overrides (getattr would resolve to the base class and silently misprice the window floor if an override vanished).""" from freetoken.kvcache.dsa_pool import DSAKVCache, MLAKVCache - from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache + from freetoken.kvcache.dsv4.v4_pool import DSV4PagedKVCache from freetoken.kvcache.hybrid_swa_pool import HybridSWAKVCache from freetoken.kvcache.mha_pool import MHAKVCache @@ -279,7 +279,7 @@ def test_every_kv_pool_answers_the_rebuild_surface(): """The engine drives every pool through this one surface; a pool that skips an implementation must fail loudly at import/instantiation, not at rebuild time.""" from freetoken.kvcache.dsa_pool import DSAKVCache, MLAKVCache - from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache + from freetoken.kvcache.dsv4.v4_pool import DSV4PagedKVCache from freetoken.kvcache.hybrid_swa_pool import HybridSWAKVCache from freetoken.kvcache.mha_pool import MHAKVCache diff --git a/tests/kvcache/test_pool_sizing_surface.py b/tests/kvcache/test_pool_sizing_surface.py index 516090b099..33f27983de 100644 --- a/tests/kvcache/test_pool_sizing_surface.py +++ b/tests/kvcache/test_pool_sizing_surface.py @@ -30,7 +30,7 @@ def _model_config(specs, **attrs): def test_resolve_pool_class_follows_attn_type(): from freetoken.kvcache.dsa_pool import DSAKVCache, MLAKVCache - from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache + from freetoken.kvcache.dsv4.v4_pool import DSV4PagedKVCache from freetoken.kvcache.hybrid_swa_pool import HybridSWAKVCache from freetoken.kvcache.mha_pool import MHAKVCache @@ -91,11 +91,11 @@ def _dsv4_config(num_page_override=None): def test_dsv4_kv_cost_and_floor_parity(): - from freetoken.kvcache.dsv4_cost_model import _dsv4_swa_ratio, _dsv4_window_floor_pages - from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache + from freetoken.kvcache.dsv4.v4_cost_model import _dsv4_swa_ratio, _dsv4_window_floor_pages + from freetoken.kvcache.dsv4.v4_pool import DSV4PagedKVCache - pytest.importorskip("freetoken.kvcache.dsv4_cost_model") - from freetoken.kvcache.dsv4_cost_model import dsv4_auto_cost_model + pytest.importorskip("freetoken.kvcache.dsv4.v4_cost_model") + from freetoken.kvcache.dsv4.v4_cost_model import dsv4_auto_cost_model config = _dsv4_config() P = 128 @@ -147,8 +147,8 @@ def check(pages, baseline): def test_dsv4_validate_rebuild_floor(): from freetoken.kvcache.base import CacheRebuildRejected - from freetoken.kvcache.dsv4_cost_model import _dsv4_window_floor_pages - from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache + from freetoken.kvcache.dsv4.v4_cost_model import _dsv4_window_floor_pages + from freetoken.kvcache.dsv4.v4_pool import DSV4PagedKVCache config = _dsv4_config() pool = object.__new__(DSV4PagedKVCache) # floor path reads no instance state @@ -165,7 +165,7 @@ def test_rebuild_from_config_explicit_num_swa_pages_wins(): # The explicit parameter must take precedence over config.swa_num_pages_override -- # the engine's override write would otherwise mask a dropped parameter forever. from freetoken.distributed import set_tp_info, try_get_tp_info - from freetoken.kvcache.dsv4_cost_model import _dsv4_pool_sizes + from freetoken.kvcache.dsv4.v4_cost_model import _dsv4_pool_sizes from freetoken.kvcache.hybrid_swa_pool import _swa_paged_num_tokens, _swa_pool_floor if try_get_tp_info() is None: @@ -183,7 +183,7 @@ def test_rebuild_from_config_explicit_num_swa_pages_wins(): assert _swa_paged_num_tokens(swa_cfg, 11, num_swa_pages=explicit) == explicit + 1 assert _swa_paged_num_tokens(swa_cfg, 11) == 999_999 + 1 # override path unchanged - from freetoken.kvcache.dsv4_cost_model import _dsv4_window_floor_pages + from freetoken.kvcache.dsv4.v4_cost_model import _dsv4_window_floor_pages d_cfg = _dsv4_config() d_floor = _dsv4_window_floor_pages(d_cfg, d_cfg.page_size) diff --git a/tests/layers/test_mhc.py b/tests/layers/test_mhc.py index 4cf144e720..0359bed9c9 100644 --- a/tests/layers/test_mhc.py +++ b/tests/layers/test_mhc.py @@ -152,6 +152,75 @@ def test_triton_fused_respects_input_dtype(dtype, tol): assert err < tol, f"{name} [{dtype}]: max abs err {err}" +def test_single_pass_applies_the_previous_pre_and_hands_on_its_own(): + """Single-pass mHC (DeepSeek-V4.1): the input mix uses ``pre_mix_in``, the gates + computed here are the next sublayer's pre; post/comb are unchanged from multi-pass.""" + from freetoken.layers.mhc import ( + mhc_fused_post_pre_single_pass, + mhc_gates, + mhc_mix_input, + mhc_pre_single_pass, + ) + + fn, scale, base = _weights(seed=21) + res = _residual(seed=22) + x = torch.randn(T, HIDDEN, dtype=torch.bfloat16) + pre_in = torch.rand(T, N, dtype=torch.float32) + post0, comb0, _ = mhc_pre(res, fn, scale, base, RMS_EPS, EPS, POST_MULT, SINKHORN) + + r1, p1, c1, pre1, li1 = mhc_fused_post_pre_single_pass( + x, res, post0, comb0, pre_in, fn, scale, base, RMS_EPS, EPS, POST_MULT, SINKHORN + ) + r_ref = mhc_post(x, res, post0, comb0) + pre_ref, p_ref, c_ref = mhc_gates(r_ref, fn, scale, base, RMS_EPS, EPS, POST_MULT, SINKHORN) + assert torch.equal(r1, r_ref) and torch.equal(p1, p_ref) and torch.equal(c1, c_ref) + assert torch.equal(pre1, pre_ref) + assert torch.equal(li1, mhc_mix_input(r_ref, pre_in)) + # the multi-pass input would have used pre_ref, not pre_in + assert not torch.equal(li1, mhc_mix_input(r_ref, pre_ref)) + # a one-hot pre picks one stream (the model's seed for the chain) + one_hot = torch.zeros(T, N) + one_hot[:, 0] = 1.0 + _, _, _, li0 = mhc_pre_single_pass(res, one_hot, fn, scale, base, RMS_EPS, EPS, POST_MULT, SINKHORN) + assert torch.equal(li0, res[:, 0]) + # no post on the first sublayer: the streams pass through untouched + r0, *_ = mhc_fused_post_pre_single_pass( + x, res, None, None, one_hot, fn, scale, base, RMS_EPS, EPS, POST_MULT, SINKHORN + ) + assert torch.equal(r0, res) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +@pytest.mark.parametrize("t,hidden,with_post", [(1, 64, True), (9, 64, True), (3, 5120, True), (3, 5120, False)]) +def test_triton_single_pass_matches_torch(t, hidden, with_post): + from freetoken.layers.mhc import mhc_fused_post_pre_single_pass_torch + from freetoken.kernel.triton.mhc import mhc_fused_post_pre_single_pass_triton + + torch.manual_seed(23) + mix = 2 * N + N * N + fn = torch.randn(mix, N * hidden, dtype=torch.float32, device="cuda") * 0.05 + scale = torch.rand(3, dtype=torch.float32, device="cuda") + 0.5 + base = torch.randn(mix, dtype=torch.float32, device="cuda") * 0.3 + res = torch.randn(t, N, hidden, dtype=torch.bfloat16, device="cuda") + x = torch.randn(t, hidden, dtype=torch.bfloat16, device="cuda") + pre_in = torch.rand(t, N, dtype=torch.float32, device="cuda") + post0 = torch.rand(t, N, 1, dtype=torch.float32, device="cuda") * POST_MULT if with_post else None + comb0 = torch.softmax(torch.randn(t, N, N, device="cuda"), dim=-1) if with_post else None + + ref = mhc_fused_post_pre_single_pass_torch( + x, res, post0, comb0, pre_in, fn, scale, base, RMS_EPS, EPS, POST_MULT, SINKHORN + ) + got = mhc_fused_post_pre_single_pass_triton( + x, res, post0, comb0, pre_in, fn, scale, base, RMS_EPS, EPS, POST_MULT, SINKHORN + ) + names = ("residual", "post", "comb", "pre_new", "layer_input") + tols = (2e-2, 2e-3, 2e-3, 2e-3, 2e-2) + for name, r, g, tol in zip(names, ref, got, tols): + assert g.shape == r.shape, (name, g.shape, r.shape) + err = (g.float() - r.float()).abs().max().item() + assert err < tol, f"{name}: max abs err {err}" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") def test_triton_pre_only_matches_torch(): """HAS_POST=False path (layer 0's standalone hc_pre through the fused kernel).""" @@ -176,3 +245,31 @@ def test_triton_pre_only_matches_torch(): assert (got_post.float() - ref_post.float()).abs().max().item() < 2e-3 assert (got_comb.float() - ref_comb.float()).abs().max().item() < 2e-3 assert (got_li.float() - ref_li.float()).abs().max().item() < 2e-2 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_mhc_split_preserves_chunked_results_and_graph_replay(): + from freetoken.kernel.triton.mhc import mhc_fused_post_pre_single_pass_triton as fused + + torch.manual_seed(26) + hidden, tokens = 5120, 7 + res = torch.randn(tokens, N, hidden, device="cuda", dtype=torch.bfloat16) + x = torch.randn(tokens, hidden, device="cuda", dtype=torch.bfloat16) + post = torch.rand(tokens, N, 1, device="cuda") + comb = torch.softmax(torch.randn(tokens, N, N, device="cuda"), -1) + pre = torch.rand(tokens, N, device="cuda") + fn = torch.randn(2 * N + N * N, N * hidden, device="cuda") * 0.05 + scale, base = torch.ones(3, device="cuda"), torch.zeros(2 * N + N * N, device="cuda") + def run(lo, hi): + return fused(x[lo:hi], res[lo:hi], post[lo:hi], comb[lo:hi], pre[lo:hi], fn, scale, base, + RMS_EPS, EPS, POST_MULT, SINKHORN) + whole = run(0, tokens) + chunks = [run(0, 2), run(2, tokens)] + for i, expected in enumerate(whole): + assert torch.equal(expected, torch.cat([chunk[i] for chunk in chunks])) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = run(0, tokens) + graph.replay() + for got, expected in zip(captured, whole): + assert torch.equal(got, expected) diff --git a/tests/models/deepseek_v41/__init__.py b/tests/models/deepseek_v41/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/models/deepseek_v41/common.py b/tests/models/deepseek_v41/common.py new file mode 100644 index 0000000000..d0f8efeef9 --- /dev/null +++ b/tests/models/deepseek_v41/common.py @@ -0,0 +1,176 @@ +"""Shared fixtures for the deepseek_v41 tests: a tiny V4.1-shaped config (2 window layers, a ratio-2 +encoder group with one Full + one Reuse layer, a ratio-1 decoder group with a Full candidate source and +a Reindex layer) and synthetic bf16 and FP8/MXFP4 checkpoints in the real tensor naming. Holds no tests itself.""" + +from __future__ import annotations + +import json +import os +from types import SimpleNamespace + +import pytest +import torch + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") + +VOCAB = 512 +DIM = 256 +N_LAYERS = 6 +RATIOS = [0, 0, 2, 2, 1, 1, 0] # + one trailing MTP layer the model ignores +KV_SOURCES = [2, 4] +INDEX_SOURCES = [2, 4, 5] +ENGRAM_LAYERS = [1] +ENGRAM_VOCAB = 500 # bucket prime start; 3 orders x 2 heads ~= 3000 rows +ENGRAM_ROWS = 4096 + + +def tiny_text_config(**over) -> dict: + text = dict( + model_type="deepseek_v41_text", + vocab_size=VOCAB, hidden_size=DIM, moe_intermediate_size=64, num_hidden_layers=N_LAYERS, num_attention_heads=4, + num_key_value_heads=1, head_dim=64, qk_rope_head_dim=16, q_lora_rank=64, o_lora_rank=32, o_groups=2, + hidden_act="silu", swiglu_limit=10.0, rms_norm_eps=1e-20, max_position_embeddings=16384, + rope_theta=10000, rope_scaling=dict(rope_type="yarn", factor=16, beta_fast=32, beta_slow=1, original_max_position_embeddings=1024), + n_routed_experts=8, n_shared_experts=1, num_experts_per_tok=2, scoring_func="sqrtsoftplus", topk_method="noaux_tc", + norm_topk_prob=True, routed_scaling_factor=1.5, sliding_window=128, + compress_ratios=list(RATIOS), compress_rope_theta=160000, kv_source_layer_ids=list(KV_SOURCES), + index_source_layer_ids=list(INDEX_SOURCES), index_n_heads=2, index_head_dim=32, index_topk=8, + candidate_source_layer_id=4, candidate_topk_blocks=2, candidate_block_size=4, + hc_mult=4, hc_sinkhorn_iters=20, hc_eps=1e-6, + engram_layer_ids=list(ENGRAM_LAYERS), engram_num_embeddings=[ENGRAM_ROWS], engram_max_ngram_size=4, + engram_vocab_size=ENGRAM_VOCAB, engram_n_heads=2, engram_head_dim=32, engram_pad_token_id=2, + engram_compressed_vocab_size=VOCAB, # the synthetic tokenizer maps every id to itself + num_nextn_predict_layers=1, dspark_block_size=0, dspark_target_layer_ids=[], + ) + text.update(over) + return text + + +def tiny_hf_config(**over) -> SimpleNamespace: + return SimpleNamespace(architectures=["DeepseekV41ForCausalLM"], model_type="deepseek_v41", text_config=tiny_text_config(**over)) + + +def checkpoint_tensors(text: dict, seed: int = 0) -> dict[str, torch.Tensor]: + """Random bf16 weights under the checkpoint's names (the reference ``convert.py`` dialect without + quantization: no ``.scale`` tensors). hc / sink / router-bias / engram gates in fp32 as shipped.""" + g = torch.Generator().manual_seed(seed) + + def rnd(*shape, scale=0.02, dtype=torch.bfloat16): + return (torch.randn(*shape, generator=g) * scale).to(dtype) + + dim, heads, hd, rd = text["hidden_size"], text["num_attention_heads"], text["head_dim"], text["qk_rope_head_dim"] + qr, ol, og = text["q_lora_rank"], text["o_lora_rank"], text["o_groups"] + inter, E, hc = text["moe_intermediate_size"], text["n_routed_experts"], text["hc_mult"] + ih, ihd = text["index_n_heads"], text["index_head_dim"] + mix = (2 + hc) * hc + t: dict[str, torch.Tensor] = { + "embed.weight": rnd(text["vocab_size"], dim, scale=1.0), + "norm.weight": torch.ones(dim, dtype=torch.bfloat16), + "head.weight": rnd(text["vocab_size"], dim), + } + for L in range(text["num_hidden_layers"]): + a = f"layers.{L}.attn" + t[f"{a}.wq_a.weight"] = rnd(qr, dim) + t[f"{a}.q_norm.weight"] = torch.ones(qr, dtype=torch.bfloat16) + t[f"{a}.wq_b.weight"] = rnd(heads * hd, qr) + t[f"{a}.wkv.weight"] = rnd(hd, dim) + t[f"{a}.kv_norm.weight"] = torch.ones(hd, dtype=torch.bfloat16) + t[f"{a}.wo_a.weight"] = rnd(og * ol, heads * hd // og) + t[f"{a}.wo_b.weight"] = rnd(dim, og * ol) + t[f"{a}.attn_sink"] = rnd(heads, scale=0.5, dtype=torch.float32) + ratio = text["compress_ratios"][L] + if L in text["kv_source_layer_ids"]: + t[f"{a}.compressor.wkv.weight"] = rnd(hd, dim) + if ratio > 1: + t[f"{a}.compressor.wgate.weight"] = rnd(hd, dim) + t[f"{a}.compressor.norm.weight"] = torch.ones(hd, dtype=torch.bfloat16) + t[f"{a}.indexer.wk.weight"] = rnd(ihd, hd) + t[f"{a}.indexer.k_norm.weight"] = torch.ones(ihd, dtype=torch.bfloat16) + if L in text["index_source_layer_ids"]: + t[f"{a}.indexer.wq_b.weight"] = rnd(ih * ihd, qr) + t[f"{a}.indexer.weights_proj.weight"] = rnd(ih, dim) + p = f"layers.{L}" + t[f"{p}.attn_norm.weight"] = torch.ones(dim, dtype=torch.bfloat16) + t[f"{p}.ffn_norm.weight"] = torch.ones(dim, dtype=torch.bfloat16) + t[f"{p}.ffn.gate.weight"] = rnd(E, dim) + t[f"{p}.ffn.gate.bias"] = rnd(E, scale=0.1, dtype=torch.float32) + t[f"{p}.ffn.gate.bias_vl"] = rnd(E, scale=0.1, dtype=torch.float32) + for w, shape in (("w1", (inter, dim)), ("w2", (dim, inter)), ("w3", (inter, dim))): + t[f"{p}.ffn.shared_experts.{w}.weight"] = rnd(*shape) + for e in range(E): + t[f"{p}.ffn.experts.{e}.{w}.weight"] = rnd(*shape) + for nm in ("hc_attn_fn", "hc_ffn_fn"): + t[f"{p}.{nm}"] = rnd(mix, hc * dim, scale=0.05, dtype=torch.float32) + for nm in ("hc_attn_base", "hc_ffn_base"): + t[f"{p}.{nm}"] = rnd(mix, scale=0.3, dtype=torch.float32) + for nm in ("hc_attn_scale", "hc_ffn_scale"): + t[f"{p}.{nm}"] = (torch.rand(3, generator=g) + 0.5).to(torch.float32) + if L in text["engram_layer_ids"]: + e = f"{p}.engram" + cols = (text["engram_max_ngram_size"] - 1) * text["engram_n_heads"] + t[f"{e}.wkv.weight"] = rnd(dim * (hc + 1), cols * text["engram_head_dim"]) + t[f"{e}.q_weight"] = torch.ones(hc, dim, dtype=torch.bfloat16) + t[f"{e}.k_weight"] = torch.ones(hc, dim, dtype=torch.bfloat16) + rows = text["engram_num_embeddings"][text["engram_layer_ids"].index(L)] + t[f"{e}.embed.weight"] = rnd(rows, text["engram_head_dim"], scale=1.0).to(torch.float8_e4m3fn) + t[f"{e}.embed.scale"] = torch.full((rows, text["engram_head_dim"] // 32), 127, dtype=torch.uint8).view(torch.float8_e8m0fnu) + # the DSpark stage and the vision tower: present in a real checkpoint, dropped by the loader + t["mtp.0.attn_norm.weight"] = torch.ones(dim, dtype=torch.bfloat16) + t["vision.norm.weight"] = torch.ones(8, dtype=torch.bfloat16) + return t + + +def write_tiny_checkpoint(folder: str, seed: int = 0, *, quantized: bool = False, **over) -> dict[str, torch.Tensor]: + """A loadable checkpoint directory: HF-style config.json + one safetensors shard + index.""" + from safetensors.torch import save_file + + if quantized: + over.setdefault("moe_intermediate_size", 256) + text = tiny_text_config(**over) + os.makedirs(folder, exist_ok=True) + with open(os.path.join(folder, "config.json"), "w") as f: + config = {"architectures": ["DeepseekV41ForCausalLM"], "model_type": "deepseek_v41", "text_config": text} + if quantized: + config["quantization_config"] = { + "quant_method": "fp8", "activation_scheme": "dynamic", + "weight_block_size": [32, 32], "scale_fmt": "ue8m0", "expert_dtype": "fp4", + } + json.dump(config, f) + tensors = checkpoint_tensors(text, seed) + if quantized: + tensors = _quantize_checkpoint(tensors) + save_file(tensors, os.path.join(folder, "model-00001-of-00001.safetensors")) + with open(os.path.join(folder, "model.safetensors.index.json"), "w") as f: + json.dump({"metadata": {}, "weight_map": {k: "model-00001-of-00001.safetensors" for k in tensors}}, f) + return tensors + + +def _quantize_checkpoint(tensors): + from .reference.kernel import _round_fp4 + + out = dict(tensors) + fp8_suffixes = ( + ".attn.wq_a", ".attn.wq_b", ".attn.wkv", ".attn.wo_a", ".attn.wo_b", + ".indexer.wq_b", ".shared_experts.w1", ".shared_experts.w2", ".shared_experts.w3", ".engram.wkv", + ) + grid = torch.tensor([0., .5, 1., 1.5, 2., 3., 4., 6.]) + for name, weight in tensors.items(): + if not name.endswith(".weight"): + continue + stem = name.removesuffix(".weight") + if ".ffn.experts." in name: + groups = weight.float().unflatten(-1, (-1, 32)) + scale = torch.exp2(torch.ceil(torch.log2(groups.abs().amax(-1).clamp_min(1e-12) / 6))) + q = _round_fp4(groups / scale.unsqueeze(-1)).flatten(-2) + codes = (q.abs().unsqueeze(-1) == grid).to(torch.int32).argmax(-1).to(torch.uint8) + codes |= torch.signbit(q).to(torch.uint8) * 8 + out[name] = codes[:, 0::2] | (codes[:, 1::2] << 4) + out[stem + ".scale"] = scale.to(torch.float8_e8m0fnu) + elif stem.endswith(fp8_suffixes): + n, k = weight.shape + blocks = weight.float().reshape(n // 32, 32, k // 32, 32).permute(0, 2, 1, 3) + scale = torch.exp2(torch.ceil(torch.log2(blocks.abs().amax((-1, -2)).clamp_min(1e-4) / 448))) + expanded = scale.repeat_interleave(32, 0).repeat_interleave(32, 1) + out[name] = (weight.float() / expanded).to(torch.float8_e4m3fn) + out[stem + ".scale"] = scale.to(torch.float8_e8m0fnu) + return out diff --git a/tests/models/deepseek_v41/harness.py b/tests/models/deepseek_v41/harness.py new file mode 100644 index 0000000000..98b19e2de4 --- /dev/null +++ b/tests/models/deepseek_v41/harness.py @@ -0,0 +1,182 @@ +"""A minimal in-process engine for the tiny DeepSeek-V4.1 model: builds the DSV41 pool, backend and +context by hand (no scheduler), loads a synthetic bf16 or FP8/MXFP4 checkpoint with resident or offloaded experts, and drives +ragged prefill / batched decode batches through ``model.forward()``. Shared by the forward smoke test +and the reference-parity test.""" + +from __future__ import annotations + +from dataclasses import replace + +import torch + +from freetoken.core import Batch, Context, Req, SamplingParams, get_global_ctx, set_global_ctx +from freetoken.distributed.info import set_tp_info, try_get_tp_info +from freetoken.engine.engine import _materialize_loaded_weight_state_dict +from freetoken.kvcache.dsv4.v41_cost_model import dsv41_pool_sizes +from freetoken.kvcache.dsv4.v41_pool import DSV41PagedKVCache +from freetoken.layers import set_rope_device +from freetoken.layers.quantization import NoQuantConfig +from freetoken.layers.quantization.method import finalize_quant +from freetoken.models import create_model +from freetoken.models.deepseek_v41.config import parse_config +from freetoken.models.deepseek_v41.engram import ZeroEngramTable +from freetoken.models.deepseek_v41.weight import iter_weights +from freetoken.utils.hf import cached_load_hf_config +from freetoken.utils.torch_utils import torch_dtype + +P = 128 + + +class TinyEngine: + def __init__(self, checkpoint: str, *, max_seq_len: int = 1024, max_running_req: int = 2, swa_decoder_replay: str = "exact", engram_table=None, quantized: bool = False): + self.device = torch.device("cuda") + self.checkpoint = checkpoint + if try_get_tp_info() is None: + set_tp_info(0, 1) + set_rope_device(self.device) + hf_config = cached_load_hf_config(checkpoint) + mc = parse_config(hf_config) + args = mc.dsv41_args + args.max_seq_len = max_seq_len + args.max_batch_size = max_running_req + 1 + args.swa_decoder_replay = swa_decoder_replay + self.args = args + quant = NoQuantConfig() + if quantized: + from freetoken.models.register import checkpoint_quant_config, get_model_spec + + quant = checkpoint_quant_config(checkpoint, hf_config, get_model_spec("DeepseekV41ForCausalLM")) + self.config = replace(mc, moe_strategy="offload" if quantized else "fused", decode_target="gpu", quant=quant) + mc = self.config + with torch.device("meta"), torch_dtype(torch.bfloat16): + self.model = create_model(self.config) + state = _materialize_loaded_weight_state_dict( + self.model.state_dict(), iter_weights(checkpoint, self.device, include_moe_experts=False), device=self.device, + ) + self.model.load_state_dict(state) + finalize_quant(self.model) + from freetoken.moe.expert_banks import load_expert_banks + + method = self.model.model.layers.op_list[0].ffn.experts.quant_method + if quantized: + from freetoken.moe.offload_cache import OffloadMoeCache, attach_offload_moe_cache + + self.banks = load_expert_banks( + checkpoint, mc, method=method, device=self.device, dtype=torch.bfloat16, parallel=False, + ) + self.expert_cache = OffloadMoeCache( + num_layers=mc.num_moe_layers, num_experts=mc.num_experts, + cache_size=mc.num_moe_layers * mc.num_experts, device=self.device, + quant_format=self.banks.quant_format, layout=method.layout(), + ) + self.expert_cache.set_bank_sources(self.banks.sources, layer_residency=self.banks.layer_residency) + attach_offload_moe_cache(self.model, self.expert_cache) + else: # resident experts load as banks, not through the state dict (as in Engine._load_resident_experts) + from freetoken.layers import iter_moe_layers + from freetoken.moe.expert_banks import attach_resident_banks + + attach_resident_banks(list(iter_moe_layers(self.model)), load_expert_banks( + checkpoint, mc, method=method, device=self.device, dtype=torch.bfloat16, parallel=False, resident=True, + )) + for layer in self.model.engram_layers(): + layer.attach_table(engram_table if engram_table is not None else ZeroEngramTable(layer.width, self.device)) + + self.max_running_req = max_running_req + num_pages = max_seq_len // P + self.pool = DSV41PagedKVCache(dsv41_pool_sizes(num_pages + 1, args, 1.0, P), args, self.device, n_scratch=max_running_req + 1) + self.pool._init_paged_state(max_running_req, radix=False) + self.page_table = torch.zeros(max_running_req + 1, max_seq_len, dtype=torch.int32, device=self.device) + self.page_table[max_running_req].fill_(num_pages * P) # the dummy row -> the reserved tail page + self.pool.attach_page_table(self.page_table) + try: + ctx = get_global_ctx() + except AssertionError: + ctx = Context(page_size=P) + set_global_ctx(ctx) + from freetoken.attention.dsv41_sparse import DSV41SparseAttnBackend + + self.ctx = ctx + ctx.page_table = self.page_table + ctx.kv_cache = self.pool # the backend reads the pool's device from the context + self.backend = DSV41SparseAttnBackend(self.config) + self._next_page = 0 + self._bind() + + def _bind(self) -> None: + """Point the process-wide ``Context`` at this engine's page table, pool and backend. The + context is a singleton the model reads through ``get_global_ctx()``, so a test that holds two + engines must rebind before each forward; ``prefill`` / ``decode`` do this, and the model + re-resolves its backend when the binding changes.""" + ctx = self.ctx + if ( + getattr(ctx, "attn_backend", None) is self.backend + and getattr(ctx, "kv_cache", None) is self.pool + and getattr(ctx, "page_table", None) is self.page_table + ): + return + ctx.page_table = self.page_table + ctx.kv_cache = self.pool + ctx.attn_backend = self.backend + self.model.mark_for_rebind() + + # ----- requests ----------------------------------------------------------------------- + def new_request(self, table_idx: int, tokens: list[int], max_new: int = 64) -> Req: + """A request with pages for ``len(tokens) + max_new`` positions bound on row ``table_idx``.""" + n_pages = -(-(len(tokens) + max_new) // P) + first = self._next_page + self._next_page += n_pages + locs = torch.arange(first * P, (first + n_pages) * P, device=self.device) + self.page_table[table_idx, : n_pages * P] = locs.to(torch.int32) + self.pool.alloc_swa(locs) + return Req(input_ids=torch.tensor(tokens, dtype=torch.int32), table_idx=table_idx, cached_len=0, output_len=max_new, + uid=table_idx, sampling_params=SamplingParams(), cache_handle=None) + + def new_request_on_prefix(self, table_idx: int, donor: Req, tokens: list[int], max_new: int = 64) -> Req: + """A request whose first ``hit`` positions alias ``donor``'s pages (what a radix prefix hit does + through the page table) and whose remaining positions get pages of their own. ``hit`` is what the + scheduler admits: the donor's prompt, capped by ``match_req`` (the last token, and the pool's + ``prefix_replay_tokens`` under bounded replay) and page-aligned. The caller prefills the rest.""" + limit = len(tokens) - 1 + if self.pool.prefix_replay_tokens: + limit = min(limit, max(0, len(tokens) - self.pool.prefix_replay_tokens)) + hit = min(donor.device_len, limit) // P * P + assert tokens[:hit] == donor.input_ids[:hit].tolist(), "the prefix must match the donor's prompt" + self.page_table[table_idx, :hit] = self.page_table[donor.table_idx, :hit] + n_pages = -(-(len(tokens) + max_new - hit) // P) + first = self._next_page + self._next_page += n_pages + locs = torch.arange(first * P, (first + n_pages) * P, device=self.device) + self.page_table[table_idx, hit : hit + n_pages * P] = locs.to(torch.int32) + self.pool.alloc_swa(locs) + return Req(input_ids=torch.tensor(tokens, dtype=torch.int32), table_idx=table_idx, cached_len=hit, output_len=max_new, + uid=table_idx, sampling_params=SamplingParams(), cache_handle=None) + + def prefill(self, reqs: list[Req]) -> torch.Tensor: + """Run each request's ``[cached_len, device_len)`` tokens; returns ``[B, vocab]`` logits.""" + batch = Batch(reqs=reqs, phase="prefill") + batch.padded_reqs = list(reqs) + batch.input_ids = torch.cat([r.input_ids[r.cached_len : r.device_len] for r in reqs]).to(self.device) + batch.positions = torch.cat([torch.arange(r.cached_len, r.device_len) for r in reqs]).to(self.device) + self._bind() + self.backend.prepare_metadata(batch) + with self.ctx.forward_batch(batch), self.model.forward_host_ctx(batch, False): + return self.model.forward().float() + + def decode(self, reqs: list[Req], tokens: list[int]) -> torch.Tensor: + """Feed one new token per request (appended at ``device_len - 1``); returns ``[B, vocab]``.""" + for r, t in zip(reqs, tokens): + r.append_host(torch.tensor([t], dtype=torch.int32)) + r.complete_one() + batch = Batch(reqs=reqs, phase="decode") + batch.padded_reqs = list(reqs) + batch.input_ids = torch.tensor(tokens, dtype=torch.int32, device=self.device) + batch.positions = torch.tensor([r.device_len - 1 for r in reqs], device=self.device) + batch.active_table_idx = torch.tensor([r.table_idx for r in reqs], device=self.device) + self._bind() + self.backend.prepare_metadata(batch) + with self.ctx.forward_batch(batch), self.model.forward_host_ctx(batch, False): + return self.model.forward().float() + + def finish_prefill(self, reqs: list[Req]) -> None: + for r in reqs: + r.cached_len = r.device_len diff --git a/tests/models/deepseek_v41/reference/LICENSE b/tests/models/deepseek_v41/reference/LICENSE new file mode 100644 index 0000000000..d42fae903e --- /dev/null +++ b/tests/models/deepseek_v41/reference/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2023 DeepSeek + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. \ No newline at end of file diff --git a/tests/models/deepseek_v41/reference/__init__.py b/tests/models/deepseek_v41/reference/__init__.py new file mode 100644 index 0000000000..79d0c33108 --- /dev/null +++ b/tests/models/deepseek_v41/reference/__init__.py @@ -0,0 +1,2 @@ +"""The DeepSeek-V4.1 reference implementation (MIT, vendored from the HF repo's ``inference/``) with its +tilelang kernels transcribed to torch -- the numerical oracle FreeToken's port is tested against.""" diff --git a/tests/models/deepseek_v41/reference/engram.py b/tests/models/deepseek_v41/reference/engram.py new file mode 100644 index 0000000000..625553113f --- /dev/null +++ b/tests/models/deepseek_v41/reference/engram.py @@ -0,0 +1,186 @@ +# Vendored from https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash (inference/engram.py), MIT License +# (see LICENSE next to this file); unmodified. +from dataclasses import dataclass + +import numpy as np +import torch +from sympy import isprime +from torch import nn + + +def find_next_prime(start: int, seen_primes: set[int]) -> int: + """The smallest prime above `start` that has not been handed out yet.""" + candidate = start + 1 + while not isprime(candidate) or candidate in seen_primes: + candidate += 1 + return candidate + + +def build_compressed_token_map(tokenizer) -> tuple[list[int], int]: + """Map every token id onto a smaller id space where tokens that normalize alike collapse together. + + N-grams are hashed over these compressed ids, so " The", "the" and "THE" all hash the same way. + Returns the lookup plus the size of the compressed vocab -- and that size matters beyond bounds + checking, because every hash multiplier is derived from it. + """ + from tokenizers import Regex, normalizers + + # a private-use char, so a token that is exactly one space survives Strip() instead of + # collapsing to the empty string and merging with unrelated tokens + sentinel = "\ue000" + normalizer = normalizers.Sequence( + [ + normalizers.NFKC(), + normalizers.NFD(), + normalizers.StripAccents(), + normalizers.Lowercase(), + normalizers.Replace(Regex(r"[ \t\r\n]+"), " "), + normalizers.Replace(Regex(r"^ $"), sentinel), + normalizers.Strip(), + normalizers.Replace(sentinel, " "), + ] + ) + + # the raw Rust tokenizer, matching what training decodes with (no clean_up_tokenization_spaces) + backend = tokenizer.backend_tokenizer + key_to_new: dict[str, int] = {} + lookup = [0] * len(tokenizer) + for token_id in range(len(tokenizer)): + text = backend.decode([token_id], skip_special_tokens=False) + if "\ufffd" in text: + # a partial UTF-8 byte token: nothing to normalize, so key it by its raw form + key = backend.id_to_token(token_id) + else: + normalized = normalizer.normalize_str(text) + key = normalized if normalized else text + + new_id = key_to_new.get(key) + if new_id is None: + new_id = len(key_to_new) + key_to_new[key] = new_id + lookup[token_id] = new_id + + return lookup, len(key_to_new) + + +def compute_hash_multipliers( + layer_ids: tuple[int, ...], max_ngram_size: int, tokenizer_vocab_size: int +) -> torch.Tensor: + """One multiplier per (layer, lookback), from a per-layer RNG so layers hash differently. + + Kept odd, and bounded so that `token_id * multiplier` cannot overflow int64. + """ + max_long = np.iinfo(np.int64).max + multiplier_bound = max(1, (max_long // tokenizer_vocab_size) // 2) + rows = [] + for layer_id in layer_ids: + generator = np.random.default_rng(10007 * layer_id) + values = generator.integers( + low=0, + high=multiplier_bound, + size=(max_ngram_size,), + dtype=np.int64, + ) + rows.append(torch.tensor(values * 2 + 1)) + return torch.stack(rows) + + +@dataclass(frozen=True) +class EngramLayout: + """Bucket layout of the n-gram hash tables. + + A position is hashed as `max_ngram_size - 1` n-grams (2-gram .. max_ngram_size-gram), each split + over `n_heads` heads. Every (n-gram size, head) pair owns its own prime-sized bucket range in the + layer's table; the primes are drawn in order and never reused, which keeps the ranges disjoint. + """ + + max_ngram_size: int + layer_ids: tuple[int, ...] + num_embeddings: tuple[int, ...] # table rows, per engram layer + primes: tuple[tuple[tuple[int, ...], ...], ...] # [layer][n-gram size][head] bucket modulus + n_heads: int + head_dim: int + + @classmethod + def from_args(cls, args) -> "EngramLayout | None": + layer_ids = tuple(args.engram_layer_ids) + if not layer_ids: + return None + max_ngram_size, n_heads = args.engram_max_ngram_size, args.engram_n_heads + primes, seen = [], set() + for _ in layer_ids: + per_ngram = [] + for _ in range(max_ngram_size - 1): + sizes, current = [], args.engram_vocab_size - 1 + for _ in range(n_heads): + current = find_next_prime(current, seen) + seen.add(current) + sizes.append(current) + per_ngram.append(tuple(sizes)) + primes.append(tuple(per_ngram)) + return cls( + max_ngram_size=max_ngram_size, + layer_ids=layer_ids, + num_embeddings=tuple(args.engram_num_embeddings), + primes=tuple(primes), + n_heads=n_heads, + head_dim=args.engram_head_dim, + ) + + +class NgramHashState(nn.Module): + """Maps each position to the hash ids of the n-grams ending there. + + Ids go through the compressed table, then each position is hashed with the `max_ngram_size - 1` + tokens before it. Look-back stops at the start of the sequence and at any dead token (an image + span, cached as DEAD), so an n-gram never spans one. The cache carries all of this across the + prefill/decode split. + """ + + DEAD = -1 + + def __init__(self, args, layout: EngramLayout, tokenizer): + super().__init__() + self.layout = layout + # every hash multiplier derives from the compressed vocab size, so a mismatch there would + # silently rehash the whole table + token_map, vocab_size = build_compressed_token_map(tokenizer) + assert vocab_size == args.engram_compressed_vocab_size, (vocab_size, args.engram_compressed_vocab_size) + self.pad_id = token_map[args.engram_pad_id] + flat = [[p for per_ngram in layer for p in per_ngram] for layer in layout.primes] + offsets = [np.cumsum([0, *sizes[:-1]]) for sizes in flat] + multipliers = compute_hash_multipliers(layout.layer_ids, layout.max_ngram_size, vocab_size) + self.register_buffer("primes", torch.tensor(layout.primes), persistent=False) + self.register_buffer("offsets", torch.tensor(np.array(offsets)), persistent=False) + self.register_buffer("multipliers", multipliers, persistent=False) + self.register_buffer("token_map", torch.tensor(token_map), persistent=False) + self.register_buffer( + "cache", torch.empty(args.max_batch_size, args.max_seq_len, dtype=torch.int64), persistent=False + ) + + @torch.inference_mode() + def forward(self, input_ids: torch.Tensor, start_pos: int, token_mask: torch.Tensor | None = None) -> torch.Tensor: + """token_mask: [B, L], False for tokens that take no part in an n-gram (image spans). + Returns the hash ids, shaped [B, L, n_engram_layers, n_hash_cols].""" + batch, seqlen = input_ids.shape + compressed = self.token_map[input_ids] + if token_mask is not None: + compressed = torch.where(token_mask, compressed, self.DEAD) + self.cache[:batch, start_pos : start_pos + seqlen] = compressed + + positions = torch.arange(start_pos, start_pos + seqlen, device=input_ids.device).expand(batch, seqlen) + tokens, blocked = [], torch.zeros_like(positions, dtype=torch.bool) + for shift in range(self.layout.max_ngram_size): + source = self.cache[:batch].gather(1, (positions - shift).clamp_min(0)) + blocked = blocked | (positions < shift) | (source == self.DEAD) + tokens.append(torch.where(blocked, self.pad_id, source)) + tokens = torch.stack(tokens, dim=-1) # [B, L, max_ngram_size] + + # XOR the multiplied ids together one lookback at a time, so the running value after step i + # is the hash of the (i+1)-gram; each lands in its own prime-sized bucket range + products = tokens.unsqueeze(2) * self.multipliers # [B, L, n_engram_layers, max_ngram_size] + rolling, hashes = products[..., 0], [] + for i in range(1, self.layout.max_ngram_size): + rolling = torch.bitwise_xor(rolling, products[..., i]) + hashes.append(rolling.unsqueeze(-1) % self.primes[:, i - 1]) + return torch.cat(hashes, dim=-1) + self.offsets diff --git a/tests/models/deepseek_v41/reference/kernel.py b/tests/models/deepseek_v41/reference/kernel.py new file mode 100644 index 0000000000..fab0ae08e2 --- /dev/null +++ b/tests/models/deepseek_v41/reference/kernel.py @@ -0,0 +1,113 @@ +"""Pure-torch transcriptions of the reference tilelang kernels (inference/kernel.py), for the test oracle. + +Same numerics as the tilelang code: fp32 math, power-of-two (ue8m0) or e4m3 scales rounded the way +the kernels round them, the e2m1 grid with the hardware's round-to-nearest-even, and the fp8/fp4 +in-place round trips that write the dequantized value back in the input dtype. Quantized GEMMs +dequantize independently in torch and round each projection's output to bf16. +""" + +from __future__ import annotations + +import torch + +_FP4_GRID = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]) + + +def _pow2_ceil(x: torch.Tensor) -> torch.Tensor: + """``fast_round_scale``: 2 ** ceil(log2(x)) for x > 0.""" + return torch.exp2(torch.ceil(torch.log2(x))) + + +def _round_fp4(x: torch.Tensor) -> torch.Tensor: + """Round-to-nearest-even onto the signed e2m1 grid (the ``float4_e2m1fn`` cast).""" + a = x.abs() + grid = _FP4_GRID.to(x.device) + hi = torch.searchsorted(grid, a.reshape(-1).contiguous(), right=True).reshape(a.shape).clamp(1, 7) + below, above = grid[hi - 1], grid[hi] + mid = (below + above) / 2 + even_above = torch.isin(above, torch.tensor([0.0, 1.0, 2.0, 4.0], device=x.device)) + r = torch.where((a > mid) | ((a == mid) & even_above), above, below) + r = torch.where(a >= 6.0, torch.full_like(r, 6.0), r) + return torch.copysign(r, x) + + +def act_quant(x: torch.Tensor, block_size: int = 128, scale_fmt=None, scale_dtype=torch.float32, inplace: bool = False): + """Block-wise fp8 quantization; ``inplace=True`` writes the quant+dequant round trip back into ``x``.""" + n = x.size(-1) + assert n % block_size == 0 + g = x.float().unflatten(-1, (-1, block_size)) + amax = g.abs().amax(dim=-1).clamp_min(1e-4) + s = _pow2_ceil(amax / 448.0) if scale_fmt is not None else amax / 448.0 + q = (g / s.unsqueeze(-1)).clamp(-448.0, 448.0).to(torch.float8_e4m3fn) + if inplace: + x.copy_((q.float() * s.unsqueeze(-1)).flatten(-2).to(x.dtype)) + return x + return q.flatten(-2), s.to(scale_dtype) + + +def fp4_act_quant(x: torch.Tensor, block_size: int = 32, inplace: bool = False, scale_dtype=torch.float8_e8m0fnu): + """Block-wise fp4 with ue8m0 (pow2) or e4m3 scales; ``inplace=True`` writes the round trip back.""" + n = x.size(-1) + assert n % block_size == 0 + g = x.float().unflatten(-1, (-1, block_size)) + amax = g.abs().amax(dim=-1) + if scale_dtype == torch.float8_e4m3fn: + amax = amax.clamp_min(6 * 2**-9) + s = (amax / 6.0).to(torch.float8_e4m3fn).float() + else: + amax = amax.clamp_min(6 * 2**-126) + s = _pow2_ceil(amax / 6.0) + q = _round_fp4((g / s.unsqueeze(-1)).clamp(-6.0, 6.0)) + if inplace: + x.copy_((q * s.unsqueeze(-1)).flatten(-2).to(x.dtype)) + return x + return q.flatten(-2), s.to(scale_dtype) + + +def fp8_gemm(a, a_s, b, b_s, scale_dtype=torch.float32, block_size=128): + """``C = (A * s_a) @ (B * s_b)^T`` over dequantized operands.""" + k = a.size(-1) + a_deq = (a.float().unflatten(-1, (-1, block_size)) * a_s.float().unsqueeze(-1)).flatten(-2) + n = b.size(0) + bs = b_s.float().repeat_interleave(block_size, 0)[:n].repeat_interleave(block_size, 1)[:, :k] + return (a_deq @ (b.float() * bs).t()).bfloat16() + + +def fp4_gemm(a, a_s, b, b_s, scale_dtype=torch.float8_e8m0fnu, act_block_size=32): + packed = b.view(torch.uint8) + codes = torch.stack((packed & 15, packed >> 4), dim=-1).flatten(-2).long() + grid = torch.tensor([0., .5, 1., 1.5, 2., 3., 4., 6., -0., -.5, -1., -1.5, -2., -3., -4., -6.], device=b.device) + weight = grid[codes] * b_s.float().repeat_interleave(32, -1) + x = (a.float().unflatten(-1, (-1, act_block_size)) * a_s.float().unsqueeze(-1)).flatten(-2) + return (x @ weight.t()).bfloat16() + + +def sparse_attn(q: torch.Tensor, kv: torch.Tensor, attn_sink: torch.Tensor, topk_idxs: torch.Tensor, softmax_scale: float) -> torch.Tensor: + """Per (batch, query): softmax over the gathered ``topk_idxs`` (-1 = absent) plus the sink null key. + + ``q [b, m, h, d]``, ``kv [b, n, d]``, ``topk_idxs [b, m, topk]`` -> ``o [b, m, h, d]`` in q's dtype. + A row with no valid index yields zeros (the finite ``-1e30`` running max of the kernel).""" + b, m, h, d = q.shape + valid = topk_idxs >= 0 + gathered = kv[torch.arange(b, device=q.device)[:, None, None], topk_idxs.clamp_min(0).long()] # [b, m, topk, d] + scores = torch.einsum("bmhd,bmtd->bmht", q.float(), gathered.float()) * softmax_scale + scores = scores.masked_fill(~valid[:, :, None, :], float("-inf")) + mx = scores.amax(dim=-1).clamp_min(-1e30) + p = torch.exp(scores - mx.unsqueeze(-1)) + denom = p.sum(-1) + torch.exp(attn_sink.float()[None, None, :] - mx) + o = torch.einsum("bmht,bmtd->bmhd", p, gathered.float()) / denom.unsqueeze(-1) + return o.to(q.dtype) + + +def hc_split_sinkhorn(mixes: torch.Tensor, hc_scale: torch.Tensor, hc_base: torch.Tensor, hc_mult: int = 4, sinkhorn_iters: int = 20, eps: float = 1e-6): + hc = hc_mult + mixes, sc, base = mixes.float(), hc_scale.float(), hc_base.float() + pre = torch.sigmoid(mixes[..., :hc] * sc[0] + base[:hc]) + eps + post = 2 * torch.sigmoid(mixes[..., hc : 2 * hc] * sc[1] + base[hc : 2 * hc]) + comb = (mixes[..., 2 * hc :] * sc[2] + base[2 * hc :]).unflatten(-1, (hc, hc)) + comb = comb.softmax(dim=-1) + eps + comb = comb / (comb.sum(dim=-2, keepdim=True) + eps) + for _ in range(sinkhorn_iters - 1): + comb = comb / (comb.sum(dim=-1, keepdim=True) + eps) + comb = comb / (comb.sum(dim=-2, keepdim=True) + eps) + return pre, post, comb diff --git a/tests/models/deepseek_v41/reference/model.py b/tests/models/deepseek_v41/reference/model.py new file mode 100644 index 0000000000..593c7518c4 --- /dev/null +++ b/tests/models/deepseek_v41/reference/model.py @@ -0,0 +1,1458 @@ +# Vendored from https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash (inference/model.py), MIT License +# (see LICENSE next to this file). FreeToken keeps this copy as the numerical ORACLE its port is tested +# against. Edits, all marked "FreeToken:", are limited to (1) local imports with the tilelang kernels +# replaced by the pure-torch transcriptions in kernel.py, and (2) forward_bounded(), a Decoder SWA +# Bounded Replay oracle built from the unmodified layer math. +import math +from contextlib import contextmanager +from dataclasses import dataclass +from functools import lru_cache +from typing import Literal + +import torch +import torch.distributed as dist +import torch.nn.functional as F +from torch import nn + +# FreeToken: local imports. +from .engram import EngramLayout, NgramHashState +from .kernel import ( + act_quant, + fp4_act_quant, + fp4_gemm, + fp8_gemm, + hc_split_sinkhorn, + sparse_attn, +) + +from .vision import ViT, Aligner + +IMAGE_START, IMAGE, IMAGE_NEW_LINE, IMAGE_END = range(4) # FreeToken: image_processor constants + +# Set once by Transformer.__init__; one model per process, so layers just read them. +world_size = 1 +rank = 0 +default_dtype = torch.float8_e4m3fn # storage dtype for Linear weights, from ModelArgs.dtype +fp8_block_size = 32 # one fp8 scale per 32x32 weight block / 32 activations +fp4_block_size = 32 # one fp4 scale per 32 elements along K +scale_fmt = "ue8m0" +scale_dtype = torch.float8_e8m0fnu + + +@contextmanager +def set_dtype(dtype): + """Temporarily override torch's default dtype, restoring it even if the body raises.""" + prev = torch.get_default_dtype() + torch.set_default_dtype(dtype) + try: + yield + finally: + torch.set_default_dtype(prev) + + +@dataclass +class ModelArgs: + """Field names are exactly the config JSON keys. The defaults are a small model that + `python model.py` can run, not the released shapes -- though the scale-independent + values (norm_eps, score_func, hc_*, engram_*) do match it.""" + + # runtime limits rather than model shape: they size the KV caches + max_batch_size: int = 4 + max_seq_len: int = 4096 + temperature: float = 1 + dtype: Literal["bf16", "fp8"] = "fp8" + expert_dtype: Literal["fp4"] | None = "fp4" + vocab_size: int = 129280 + dim: int = 1024 + moe_inter_dim: int = 1024 + n_layers: int = 5 + n_mtp_layers: int = 1 # extra draft layers appended after the backbone, indices n_layers.. + n_heads: int = 16 + # moe + n_routed_experts: int = 8 + n_shared_experts: int = 1 + n_activated_experts: int = 2 + score_func: Literal["softmax", "sigmoid", "sqrtsoftplus"] = "sqrtsoftplus" + gate_temp: float = 1.0 + norm_topk_prob: bool = True + route_scale: float = 1.0 + swiglu_limit: float = 0.0 + # attention: latent q/kv projections, plus a LoRA-factorised output projection over o_groups + q_lora_rank: int = 256 + head_dim: int = 128 + rope_head_dim: int = 32 + norm_eps: float = 1e-20 + o_groups: int = 8 + o_lora_rank: int = 256 + # sparse attention: every layer attends over a sliding window, and may add compressed KV on top + window_size: int = 128 + # one entry per layer, MTP layers included: 0 = sliding window only, r = KV compressed r-to-1 + compress_ratios: tuple[int, ...] = (0, 2, 2, 1, 1, 0) + # layers sharing a ratio also share one compressed KV and one indexer, produced by the first + kv_source_layers: tuple[int, ...] = (1, 3) + index_source_layers: tuple[int, ...] = (1, 3) + # rope, with YaRN extrapolation when original_seq_len > 0. Compressed KV rotates at its own + # theta because one latent stands for compress_ratio tokens, so its positions are further apart. + compress_rope_theta: float = 40000.0 + original_seq_len: int = 0 + rope_theta: float = 10000.0 + rope_factor: float = 40 + beta_fast: int = 32 + beta_slow: int = 1 + # the indexer: a small extra attention that scores compressed positions, so each query can keep + # just `index_topk` of them. Names match DeepSeek-V3.2-Exp, where this mechanism first appeared. + index_n_heads: int = 16 + index_head_dim: int = 64 + index_topk: int = 64 + # candidate pre-filtering: candidate_source_layer < 0 turns it off and the other two are unused + candidate_source_layer: int = -1 + candidate_topk_blocks: int = 0 + candidate_block_size: int = 0 + # hyper-connections: the residual stream is carried as hc_mult parallel copies + hc_mult: int = 4 + hc_sinkhorn_iters: int = 20 + hc_eps: float = 1e-6 + # engram: n-gram hash lookups added into the residual stream at a few layers + engram_layer_ids: tuple[int, ...] = () + engram_num_embeddings: tuple[int, ...] = () # unpadded table rows; each rank allocates ceil(rows / world_size) + engram_max_ngram_size: int = 1 + engram_vocab_size: int = 0 # bucket size each (n-gram size, head) starts searching primes from + engram_n_heads: int = 0 + engram_head_dim: int = 0 + engram_pad_id: int = 2 # token that fills n-gram slots with no history; matches training + # size of the compressed tokenizer vocab; every hash multiplier is derived from it + engram_compressed_vocab_size: int = 0 + # vision (VL); vision_n_layers == 0 disables the vision path + vision_n_layers: int = 0 + vision_dim: int = 1024 + vision_n_heads: int = 16 + vision_inter_dim: int = 2816 + vision_patch_size: int = 14 + vision_rope_theta: float = 10000.0 + vision_downsample_ratio: int = 3 + vision_max_n_token: int = 1024 + vision_min_pixels: int = 544 * 544 + vision_max_wh_ratio: int | None = None + # raw id of <｜deepseek_image｜>; every position of an image span carries this id in input_ids + image_token_id: int = 129264 + # dspark draft head. Only the forward pass is implemented here -- nothing calls forward_spec, + # so these are read but the speculative-decoding loop itself is out of scope for this repo. + dspark_block_size: int = 0 + dspark_noise_token_id: int = 0 + dspark_target_layer_ids: tuple[int, ...] = () + dspark_markov_rank: int = 256 + dspark_n_routed_experts: int = 0 + dspark_n_activated_experts: int = 0 + + @property + def vision_enabled(self) -> bool: + return self.vision_n_layers > 0 + + def get_moe_config(self, layer_id: int) -> tuple[int, int]: + """Return the routed/activated expert counts for a given layer.""" + if layer_id < self.n_layers: + return self.n_routed_experts, self.n_activated_experts + return ( + self.dspark_n_routed_experts or self.n_routed_experts, + self.dspark_n_activated_experts or self.n_activated_experts, + ) + + +class ParallelEmbedding(nn.Module): + """Embedding sharded along the vocab dimension. Each rank holds vocab_size // world_size rows. + Out-of-range indices are zero-masked before all_reduce to combine partial embeddings.""" + + def __init__(self, vocab_size: int, dim: int): + super().__init__() + self.vocab_size = vocab_size + self.dim = dim + assert vocab_size % world_size == 0, ( + f"Vocabulary size must be divisible by world size (world_size={world_size})" + ) + self.part_vocab_size = vocab_size // world_size + self.vocab_start_idx = rank * self.part_vocab_size + self.vocab_end_idx = self.vocab_start_idx + self.part_vocab_size + self.weight = nn.Parameter(torch.empty(self.part_vocab_size, self.dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if world_size > 1: + # ids off this rank read row 0 then get zeroed, so the all_reduce sums one real row + mask = (x < self.vocab_start_idx) | (x >= self.vocab_end_idx) + x = x - self.vocab_start_idx + x[mask] = 0 + y = F.embedding(x, self.weight) + if world_size > 1: + y[mask] = 0 + dist.all_reduce(y) + return y + + +def linear(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor: + """Pick a GEMM from the weight dtype. Quantized weights need a quantized activation, and both + fp4 and fp8 weights take an fp8 one -- for fp4 the kernel handles the mixed precision.""" + assert bias is None + + if weight.dtype == torch.float4_e2m1fn_x2: + x, s = act_quant(x, fp8_block_size, scale_fmt, scale_dtype) + return fp4_gemm( + x, + s, + weight, + weight.scale, + scale_dtype, + act_block_size=fp8_block_size, + ) + elif weight.dtype == torch.float8_e4m3fn: + x, s = act_quant(x, fp8_block_size, scale_fmt, scale_dtype) + return fp8_gemm( + x, + s, + weight, + weight.scale, + scale_dtype, + block_size=fp8_block_size, + ) + else: + return F.linear(x, weight) + + +class Linear(nn.Module): + """bf16, fp8 or fp4 weights. Quantized ones get a `scale`, also attached to `.weight` so that + `linear()` can reach it.""" + + def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype=None): + super().__init__() + self.in_features = in_features + self.out_features = out_features + dtype = dtype or default_dtype + if dtype == torch.float4_e2m1fn_x2: + # two values per byte: [out, in] logically, [out, in//2] stored + self.weight = nn.Parameter(torch.empty(out_features, in_features // 2, dtype=torch.float4_e2m1fn_x2)) + self.weight.scale = self.scale = nn.Parameter( + torch.empty(out_features, in_features // fp4_block_size, dtype=torch.float8_e8m0fnu) + ) + elif dtype == torch.float8_e4m3fn: + self.weight = nn.Parameter(torch.empty(out_features, in_features, dtype=dtype)) + self.weight.scale = self.scale = nn.Parameter( + torch.empty( + (out_features + fp8_block_size - 1) // fp8_block_size, + (in_features + fp8_block_size - 1) // fp8_block_size, + dtype=torch.float8_e8m0fnu, + ) + ) + else: + self.weight = nn.Parameter(torch.empty(out_features, in_features, dtype=dtype)) + self.register_parameter("scale", None) + if bias: + self.bias = nn.Parameter(torch.empty(out_features)) + else: + self.register_parameter("bias", None) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return linear(x, self.weight, self.bias) + + +class ColumnParallelLinear(Linear): + """Splits the output dim across ranks; each rank's slice of the output is already complete.""" + + def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype=None): + assert out_features % world_size == 0, ( + f"Output features must be divisible by world size (world_size={world_size})" + ) + self.part_out_features = out_features // world_size + super().__init__(in_features, self.part_out_features, bias, dtype) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return linear(x, self.weight, self.bias) + + +class RowParallelLinear(Linear): + """Splits the reduction dim, so each rank holds a partial sum: hence the fp32 all_reduce, with + the bias added only after it.""" + + def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype=None): + assert in_features % world_size == 0, ( + f"Input features must be divisible by world size (world_size={world_size})" + ) + self.part_in_features = in_features // world_size + super().__init__(self.part_in_features, out_features, bias, dtype) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = linear(x, self.weight, None) + if world_size > 1: + y = y.float() + dist.all_reduce(y) + if self.bias is not None: + y += self.bias + return y.type_as(x) + + +class RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.dim = dim + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x: torch.Tensor): + dtype = x.dtype + x = x.float() + var = x.square().mean(-1, keepdim=True) + x = x * torch.rsqrt(var + self.eps) + return (self.weight * x).to(dtype) + + +class ParallelEngramEmbedding(nn.Module): + """The n-gram hash table, sharded over its rows. Stays fp8: rows are dequantized on lookup.""" + + def __init__(self, num_embeddings: int, dim: int): + super().__init__() + self.num_embeddings = num_embeddings + self.dim = dim + self.part_num_embeddings = (num_embeddings + world_size - 1) // world_size + self.vocab_start_idx = rank * self.part_num_embeddings + self.vocab_end_idx = self.vocab_start_idx + self.part_num_embeddings + + self.block_size = fp8_block_size + # the table stays fp8 as stored: rows are dequantized with `scale` on lookup + self.weight = nn.Parameter(torch.empty(self.part_num_embeddings, dim, dtype=torch.float8_e4m3fn)) + self.scale = nn.Parameter(torch.empty(self.part_num_embeddings, dim // self.block_size, dtype=scale_dtype)) + + def forward(self, indices: torch.Tensor) -> torch.Tensor: + mask = (indices < self.vocab_start_idx) | (indices >= self.vocab_end_idx) + local_indices = indices - self.vocab_start_idx + local_indices = local_indices.masked_fill(mask, 0) + + values = F.embedding(local_indices, self.weight) + scales = F.embedding(local_indices, self.scale) + values = values.float().unflatten(-1, (-1, self.block_size)) * scales.float().unsqueeze(-1) + values = values.flatten(-2).to(torch.bfloat16) + values = values.masked_fill(mask.unsqueeze(-1), 0) + + if world_size > 1: + dist.all_reduce(values) + return values + + +class Engram(nn.Module): + """Writes an n-gram lookup into the residual stream, gated by how well it matches that stream. + + The hash ids fetch `n_hash_cols` rows; `wkv` turns them into one key per hc copy plus a shared + value. The gate is a normalized dot product of stream against key. + """ + + def __init__(self, args: ModelArgs, layer_id: int, layout: EngramLayout): + super().__init__() + self.layer_id = layer_id + self.layer_hash_index = layout.layer_ids.index(layer_id) + self.dim = args.dim + self.hc_mult = args.hc_mult + self.clamp_value = 1e-6 + + self.embed = ParallelEngramEmbedding(layout.num_embeddings[self.layer_hash_index], layout.head_dim) + n_hash_cols = (layout.max_ngram_size - 1) * layout.n_heads + self.wkv = Linear(n_hash_cols * layout.head_dim, args.dim * (args.hc_mult + 1)) + self.eps = args.norm_eps + self.q_weight = nn.Parameter(torch.ones(args.hc_mult, args.dim)) + self.k_weight = nn.Parameter(torch.ones(args.hc_mult, args.dim)) + + def forward(self, x: torch.Tensor, hash_ids: torch.Tensor, token_mask: torch.Tensor | None = None) -> torch.Tensor: + """x: [B, L, hc_mult, dim]; hash_ids: [B, L, n_hash_cols]; token_mask: [B, L], False shuts + the gate so those positions pass through untouched.""" + kv = self.wkv(self.embed(hash_ids).flatten(-2)) + key, value = kv.split([self.hc_mult * self.dim, self.dim], dim=-1) + key = key.float().unflatten(-1, (self.hc_mult, self.dim)) + weight = self.q_weight.float() * self.k_weight.float() # only ever used as a product + h, eps = x.float(), self.eps + # normalized per (token, hc copy) over `dim`, NOT jointly over the copies + rstd = torch.rsqrt(h.square().mean(-1) + eps) * torch.rsqrt(key.square().mean(-1) + eps) + dot = (h * weight * key).sum(-1) * rstd * self.dim**-0.5 + # signed sqrt before the sigmoid, matching the training kernel + gate = torch.sigmoid(torch.copysign(dot.abs().clamp_min(self.clamp_value).sqrt(), dot)) + if token_mask is not None: + gate = gate.masked_fill(~token_mask.unsqueeze(-1), 0) + return (h + gate.unsqueeze(-1) * value.float().unsqueeze(-2)).to(x.dtype) + + +@lru_cache(2) +def precompute_freqs_cis(dim, seqlen, original_seq_len, base, factor, beta_fast, beta_slow) -> torch.Tensor: + """Rotary frequencies as complex exponentials, one row per position. + + With original_seq_len > 0 this applies YaRN: dimensions whose wavelength already fits inside the + training context keep their frequency, those far beyond it are divided by `factor`, and the + `beta_fast`..`beta_slow` band in between is faded across with a linear ramp. + """ + freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + if original_seq_len > 0: + # the dim whose wavelength completes `rotations` turns over the training context + def corrected_dim(rotations): + return dim * math.log(original_seq_len / (rotations * 2 * math.pi)) / (2 * math.log(base)) + + low = max(math.floor(corrected_dim(beta_fast)), 0) + high = min(math.ceil(corrected_dim(beta_slow)), dim - 1) + ramp = ((torch.arange(dim // 2, dtype=torch.float32) - low) / max(high - low, 1e-3)).clamp(0, 1) + smooth = 1 - ramp + freqs = freqs / factor * (1 - smooth) + freqs * smooth + + freqs = torch.outer(torch.arange(seqlen), freqs) + return torch.polar(torch.ones_like(freqs), freqs) + + +def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False) -> torch.Tensor: + """Rotate `x` in place, taking adjacent element pairs as complex numbers. Accepts [b, s, d] and + [b, s, h, d]; `inverse` conjugates the rotation, which is how the attention output gets the + query's rotation removed again so the cache can stay in one shared rotated form.""" + y = x + x = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) + if inverse: + freqs_cis = freqs_cis.conj() + if x.ndim == 3: + freqs_cis = freqs_cis.view(1, x.size(1), x.size(-1)) + else: + freqs_cis = freqs_cis.view(1, x.size(1), 1, x.size(-1)) + x = torch.view_as_real(x * freqs_cis).flatten(-2) + y.copy_(x) + return y + + +@lru_cache(1) +def get_window_topk_idxs(window_size: int, bsz: int, seqlen: int, start_pos: int): + """Which sliding-window cache slots each query attends to; -1 marks a slot holding nothing. + + The cache is a ring of `window_size` slots. Prefill needs one row per query, each seeing its own + causal window. A decode step has a single query that sees the whole ring, listed oldest first. + Order within a row does not matter to `sparse_attn`, which handles every slot independently. + """ + if start_pos == 0: + end = torch.arange(seqlen).unsqueeze(1) + idxs = (end - window_size + 1).clamp(0) + torch.arange(min(seqlen, window_size)) + idxs = torch.where(idxs > end, -1, idxs) # before the sequence started + else: + oldest = start_pos % window_size + 1 + idxs = torch.cat([torch.arange(oldest, window_size), torch.arange(oldest)]) + idxs = torch.where(idxs > start_pos, -1, idxs) # ring still filling + # sparse_attn needs real [b, m, topk] int32 memory, hence the materializing expand + return idxs.int().unsqueeze(0).expand(bsz, -1, -1).contiguous() + + +class Compressor(nn.Module): + """Pools `compress_ratio` consecutive tokens into one KV latent with a learned softmax gate. + + Returns the latent before RoPE, or None while a group is still filling up -- so during decode it + only yields every `compress_ratio` steps, holding the partial group in `kv_state`/`score_state`. + Pre-RoPE is deliberate: the indexer needs the unrotated form, so Attention rotates afterwards. + """ + + def __init__(self, args: ModelArgs, layer_id: int): + super().__init__() + compress_ratio = args.compress_ratios[layer_id] + head_dim = args.head_dim + self.compress_ratio = compress_ratio + self.head_dim = head_dim + self.norm = RMSNorm(head_dim, args.norm_eps) + # ratio 1 is a plain projection, so it stays in the checkpoint's bf16; the softmax pooling + # above ratio 1 runs in fp32, so those weights are promoted to fp32 to match + self.wkv = Linear(args.dim, head_dim, dtype=torch.float32 if compress_ratio > 1 else torch.bfloat16) + if compress_ratio > 1: + self.wgate = Linear(args.dim, head_dim, dtype=torch.float32) + # tail of an incomplete group, carried across decode steps + self.kv_state: torch.Tensor + self.score_state: torch.Tensor + state_shape = (args.max_batch_size, compress_ratio, head_dim) + self.register_buffer("kv_state", torch.zeros(state_shape, dtype=torch.float32), persistent=False) + self.register_buffer( + "score_state", torch.full(state_shape, -torch.inf, dtype=torch.float32), persistent=False + ) + + def forward(self, x: torch.Tensor, start_pos: int) -> torch.Tensor | None: + bsz, seqlen, _ = x.size() + ratio, dtype = self.compress_ratio, x.dtype + if ratio == 1: # one token per group: nothing to pool, so no gate and no fp32 + return self.norm(self.wkv(x)) + + x = x.float() + kv, score = self.wkv(x), self.wgate(x) + if start_pos == 0: + should_compress = seqlen >= ratio + remainder = seqlen % ratio + cutoff = seqlen - remainder + if remainder: # trailing partial group waits in the state + kv, self.kv_state[:bsz, :remainder] = kv.split([cutoff, remainder], dim=1) + score, self.score_state[:bsz, :remainder] = score.split([cutoff, remainder], dim=1) + kv = kv.unflatten(1, (-1, ratio)) + score = score.unflatten(1, (-1, ratio)) + kv = (kv * score.softmax(dim=2)).sum(dim=2) + else: # one token per step: fill a slot, and pool only when the group just completed + should_compress = (start_pos + 1) % ratio == 0 + slot = start_pos % ratio + self.kv_state[:bsz, slot] = kv.squeeze(1) + self.score_state[:bsz, slot] = score.squeeze(1) + if should_compress: + kv = (self.kv_state[:bsz] * self.score_state[:bsz].softmax(dim=1)).sum(dim=1, keepdim=True) + if not should_compress: + return None + return self.norm(kv.to(dtype)) + + +class Indexer(torch.nn.Module): + """Keeps the `index_topk` best compressed positions per query. + + A small side attention: fp4 query heads against one shared key per compressed position, scores + rectified then combined by `weights_proj`. With a candidate source this is the second of two + levels; `select_candidate_blocks` is the first. + """ + + def __init__(self, args: ModelArgs, layer_id: int): + super().__init__() + # the index keys are derived from the compressor's latent, so only a layer that compresses + # its own KV can produce them; every other indexer reads them from that layer's cache + self.owns_k = layer_id in args.kv_source_layers + self.compress_ratio = args.compress_ratios[layer_id] + self.is_candidate_source = layer_id == args.candidate_source_layer + self.uses_candidates = 0 <= args.candidate_source_layer < layer_id + self.candidate_topk_blocks = args.candidate_topk_blocks + self.candidate_block_size = args.candidate_block_size + self.dim = args.dim + self.n_heads = args.index_n_heads + self.n_local_heads = args.index_n_heads // world_size + self.index_head_dim = args.index_head_dim + self.rope_head_dim = args.rope_head_dim + self.index_topk = args.index_topk + self.q_lora_rank = args.q_lora_rank + self.softmax_scale = self.index_head_dim**-0.5 + self.wq_b = ColumnParallelLinear(self.q_lora_rank, self.n_heads * self.index_head_dim) + self.weights_proj = ColumnParallelLinear(self.dim, self.n_heads, dtype=torch.bfloat16) + self.freqs_cis: torch.Tensor | None = None + if self.owns_k: + self.wk = Linear(args.head_dim, self.index_head_dim, dtype=torch.bfloat16) + self.k_norm = RMSNorm(self.index_head_dim, args.norm_eps) + self.k_cache: torch.Tensor + self.register_buffer( + "k_cache", + torch.zeros(args.max_batch_size, args.max_seq_len // self.compress_ratio, args.index_head_dim), + persistent=False, + ) + + def forward(self, x: torch.Tensor, qr: torch.Tensor, latent: torch.Tensor, start_pos: int, offset: int): + """`latent` is this layer's RoPE-free compressed latent, None when this layer does not + compress or when its current group is still incomplete. An index-key owner turns it into + index keys here, which has to happen before Attention overwrites that same storage with + the RoPE'd, quantized values.""" + assert self.freqs_cis is not None + bsz, seqlen, _ = x.size() + ratio, rd, end_pos = self.compress_ratio, self.rope_head_dim, start_pos + seqlen + + # latent is None while a group is still filling up, so there is nothing to publish yet + if self.owns_k and latent is not None: + # a latent stands for the first token of its group, so group j takes position j * ratio + freqs = ( + self.freqs_cis[: seqlen - seqlen % ratio : ratio] + if start_pos == 0 + else self.freqs_cis[start_pos + 1 - ratio].unsqueeze(0) + ) + k = self.k_norm(self.wk(latent)) + apply_rotary_emb(k[..., -rd:], freqs) + fp4_act_quant(k, fp4_block_size, True) + self.k_cache[:bsz, start_pos // ratio : start_pos // ratio + k.size(1)] = k + shared_attn.index_k = self.k_cache + + q = self.wq_b(qr).unflatten(-1, (self.n_local_heads, self.index_head_dim)) + apply_rotary_emb(q[..., -rd:], self.freqs_cis[start_pos:end_pos]) + fp4_act_quant(q, fp4_block_size, True) + + index_k = shared_attn.index_k[:bsz, : end_pos // ratio] + weights = self.weights_proj(x) * (self.softmax_scale * self.n_heads**-0.5) + index_score = torch.einsum("bshd,btd->bsht", q, index_k) + index_score = (index_score.relu_() * weights.unsqueeze(-1)).sum(dim=2) + if world_size > 1: + dist.all_reduce(index_score) + + # how many compressed positions each query can see: a block becomes visible once the query + # has passed its last token. One query per decode step, so there it is just a number. + if start_pos == 0: + compress_lens = (torch.arange(1, seqlen + 1, device=x.device) // ratio).unsqueeze(-1) + index_score.masked_fill_(torch.arange(seqlen // ratio, device=x.device) >= compress_lens, -torch.inf) + else: + compress_lens = end_pos // ratio + + if self.is_candidate_source: + shared_attn.candidates = select_candidate_blocks( + index_score, compress_lens, self.candidate_topk_blocks, self.candidate_block_size + ) + elif self.uses_candidates: + # level two: score with our own weights, but only inside the source's candidate blocks + index_score = index_score.masked_fill(~shared_attn.candidates, -torch.inf) + + # top-k by score, re-sorted into position order; unreachable -> -1, rest shifted by offset + topk = min(self.index_topk, end_pos // ratio) + idxs = index_score.topk(topk, dim=-1, sorted=False).indices.sort(dim=-1).values + return torch.where(idxs < compress_lens, idxs + offset, -1).int() + + # FreeToken: Decoder SWA Bounded Replay oracle -- the indexer for a TAIL prefill: queries are the + # positions [s, s + n) of a prompt of L_all tokens; the keys (a source's) were published for all + # L_all tokens by forward()'s start_pos == 0 path. Same scoring / masking / top-k as forward(). + def forward_tail(self, x, qr, s: int, n: int, offset: int): + assert self.freqs_cis is not None + bsz = x.size(0) + ratio, rd = self.compress_ratio, self.rope_head_dim + q = self.wq_b(qr).unflatten(-1, (self.n_local_heads, self.index_head_dim)) + apply_rotary_emb(q[..., -rd:], self.freqs_cis[s : s + n]) + fp4_act_quant(q, fp4_block_size, True) + end_pos = s + n + index_k = shared_attn.index_k[:bsz, : end_pos // ratio] + weights = self.weights_proj(x) * (self.softmax_scale * self.n_heads**-0.5) + index_score = torch.einsum("bshd,btd->bsht", q, index_k) + index_score = (index_score.relu_() * weights.unsqueeze(-1)).sum(dim=2) + compress_lens = ((s + torch.arange(1, n + 1, device=x.device)) // ratio).unsqueeze(-1) + index_score.masked_fill_(torch.arange(end_pos // ratio, device=x.device) >= compress_lens, -torch.inf) + if self.is_candidate_source: + shared_attn.candidates = select_candidate_blocks( + index_score, compress_lens, self.candidate_topk_blocks, self.candidate_block_size + ) + elif self.uses_candidates: + index_score = index_score.masked_fill(~shared_attn.candidates, -torch.inf) + topk = min(self.index_topk, end_pos // ratio) + idxs = index_score.topk(topk, dim=-1, sorted=False).indices.sort(dim=-1).values + return torch.where(idxs < compress_lens, idxs + offset, -1).int() + + +def select_candidate_blocks( + logits: torch.Tensor, + compress_lens: torch.Tensor | int, + topk_blocks: int, + block_size: int, +) -> torch.Tensor: + """Level one of the two-level top-k: keep the `topk_blocks` highest-scoring blocks per query. + + `logits` is [..., n_positions] with positions the query cannot reach already at -inf, which is + what makes a block score of -inf mean "not reachable yet". `compress_lens` is a plain int during + decode, or broadcasts against logits' leading dims during prefill. Returns a bool mask shaped + like `logits`, so the layers consuming it just mask and never think about blocks again. + """ + width = logits.size(-1) + # score each block by its best position; -inf pads the last one out to block_size + scores = F.pad(logits, (0, -width % block_size), value=-torch.inf) + scores = scores.unflatten(-1, (-1, block_size)).amax(dim=-1) + num_blocks = scores.size(-1) + + # the block with this query's newest position is only partly filled, so pin it in: it holds the + # most recent tokens but could otherwise be outscored by an older, full block + last = (compress_lens - 1) // block_size + scores = scores.masked_fill(torch.arange(num_blocks, device=logits.device) == last, torch.inf) + + top = scores.topk(min(topk_blocks, num_blocks), dim=-1) + # fewer reachable blocks than topk_blocks means leftover picks came back -inf: drop them + keep = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, top.indices, top.values > -torch.inf) + return keep.repeat_interleave(block_size, dim=-1)[..., :width] + + +class Attention(nn.Module): + """Latent attention over two KV sources at once, concatenated into one `sparse_attn` call: a + sliding window of raw KV, plus -- when compress_ratio > 0 -- `index_topk` compressed positions + reaching further back. Q and the output projection are both low-rank, the latter grouped. + + compress_ratio > 0 does not mean the layer compresses its own KV: only kv_source_layers do, + the rest read that same cache. + """ + + def __init__(self, layer_id: int, args: ModelArgs): + super().__init__() + self.layer_id = layer_id + self.dim = args.dim + self.n_heads = args.n_heads + self.n_local_heads = args.n_heads // world_size + self.q_lora_rank = args.q_lora_rank + self.o_lora_rank = args.o_lora_rank + self.head_dim = args.head_dim + self.rope_head_dim = args.rope_head_dim + self.nope_head_dim = args.head_dim - args.rope_head_dim + self.n_groups = args.o_groups + self.n_local_groups = self.n_groups // world_size + self.window_size = args.window_size + self.compress_ratio = args.compress_ratios[layer_id] + self.eps = args.norm_eps + + self.attn_sink = nn.Parameter(torch.empty(self.n_local_heads, dtype=torch.float32)) + self.wq_a = Linear(self.dim, self.q_lora_rank) + self.q_norm = RMSNorm(self.q_lora_rank, self.eps) + self.wq_b = ColumnParallelLinear(self.q_lora_rank, self.n_heads * self.head_dim) + self.wkv = Linear(self.dim, self.head_dim) + self.kv_norm = RMSNorm(self.head_dim, self.eps) + self.wo_a = ColumnParallelLinear( + self.n_heads * self.head_dim // self.n_groups, + self.n_groups * args.o_lora_rank, + dtype=torch.bfloat16, + ) + self.wo_b = RowParallelLinear(self.n_groups * args.o_lora_rank, self.dim) + self.softmax_scale = self.head_dim**-0.5 + + is_backbone = layer_id < args.n_layers + self.is_kv_source = is_backbone and layer_id in args.kv_source_layers + self.is_index_source = is_backbone and layer_id in args.index_source_layers + self.compressor: Compressor | None = None + self.indexer: Indexer | None = None + if self.is_kv_source: + self.compressor = Compressor(args, layer_id) + if self.is_index_source: + self.indexer = Indexer(args, layer_id) + + self.window_kv_cache: torch.Tensor + self.register_buffer( + "window_kv_cache", + torch.zeros(args.max_batch_size, args.window_size, self.head_dim), + persistent=False, + ) + if self.is_kv_source: + self.compress_kv_cache: torch.Tensor + self.register_buffer( + "compress_kv_cache", + torch.zeros( + args.max_batch_size, + args.max_seq_len // self.compress_ratio, + self.head_dim, + ), + persistent=False, + ) + if self.compress_ratio: + original_seq_len, rope_theta = ( + args.original_seq_len, + args.compress_rope_theta, + ) + else: + # disable YaRN and use base rope_theta in pure sliding-window attention + original_seq_len, rope_theta = 0, args.rope_theta + freqs_cis = precompute_freqs_cis( + self.rope_head_dim, + args.max_seq_len, + original_seq_len, + rope_theta, + args.rope_factor, + args.beta_fast, + args.beta_slow, + ) + self.freqs_cis: torch.Tensor + self.register_buffer("freqs_cis", freqs_cis, persistent=False) + + def _window_kv(self, x, freqs_cis, start_pos): + """This layer's sliding-window K and the window positions every query may attend to. The K + stays fp8, quantized over the whole post-RoPE vector, RoPE tail included.""" + bsz, seqlen, _ = x.size() + win = self.window_size + kv = self.kv_norm(self.wkv(x)) + apply_rotary_emb(kv[..., -self.rope_head_dim :], freqs_cis) + act_quant(kv, fp8_block_size, scale_fmt, scale_dtype, True) + if start_pos == 0: # prefill: attend over this chunk, seeding the ring buffer for decode + if seqlen <= win: + self.window_kv_cache[:bsz, :seqlen] = kv + else: + cutoff = seqlen % win + self.window_kv_cache[:bsz, cutoff:win], self.window_kv_cache[:bsz, :cutoff] = kv[:, -win:].split( + [win - cutoff, cutoff], dim=1 + ) + window_kv = kv + else: # decode: one token into the ring buffer, attend over the whole window + self.window_kv_cache[:bsz, start_pos % win] = kv.squeeze(1) + window_kv = self.window_kv_cache[:bsz] + return window_kv, get_window_topk_idxs(win, bsz, seqlen, start_pos) + + def _compress_topk_idxs(self, x, qr, latent, start_pos, offset, compress_len): + """Which compressed positions each query attends to. Index sources run their own indexer; + the layers in between reuse the result their source published.""" + if not self.is_index_source: + return shared_attn.topk_idxs + + bsz, seqlen, _ = x.size() + if compress_len == 0: + idxs = torch.empty(bsz, seqlen, 0, dtype=torch.int32, device=x.device) + else: + assert self.indexer is not None + if self.indexer.freqs_cis is None: + self.indexer.freqs_cis = self.freqs_cis + idxs = self.indexer(x, qr, latent, start_pos, offset) + shared_attn.topk_idxs = idxs + return idxs + + def _compress_kv(self, x, qr, start_pos, offset): + """The shared compressed KV and the compressed positions every query may attend to. This + layer compresses its own KV only when it is a source; otherwise it just reads the cache.""" + bsz, seqlen, _ = x.size() + ratio = self.compress_ratio + compress_len = (start_pos + seqlen) // ratio + latent = None + if self.is_kv_source: + latent = self.compressor(x, start_pos) + shared_attn.compress_kv = self.compress_kv_cache + # the indexer needs the latent before RoPE, so it runs before the cache is written + idxs = self._compress_topk_idxs(x, qr, latent, start_pos, offset, compress_len) + if latent is not None: + # a latent stands for the first token of its group, so group j takes position j * ratio + freqs = ( + self.freqs_cis[: seqlen - seqlen % ratio : ratio] + if start_pos == 0 + else self.freqs_cis[start_pos + 1 - ratio].unsqueeze(0) + ) + apply_rotary_emb(latent[..., -self.rope_head_dim :], freqs) + # Compressed KV uses groups of 16 with E4M3 scales; the indexer uses 32 with E8M0. + fp4_act_quant(latent, 16, True, scale_dtype=torch.float8_e4m3fn) + self.compress_kv_cache[:bsz, start_pos // ratio : start_pos // ratio + latent.size(1)] = latent + # read after the write, so this does not depend on the slice aliasing the cache + return shared_attn.compress_kv[:bsz, :compress_len], idxs + + def forward(self, x: torch.Tensor, start_pos: int): + bsz, seqlen, _ = x.size() + freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen] + rd = self.rope_head_dim + + qr = self.q_norm(self.wq_a(x)) + q = self.wq_b(qr).unflatten(-1, (self.n_local_heads, self.head_dim)) + apply_rotary_emb(q[..., -rd:], freqs_cis) + + kv, topk_idxs = self._window_kv(x, freqs_cis, start_pos) + if self.compress_ratio: + compress_kv, compress_idxs = self._compress_kv(x, qr, start_pos, kv.size(1)) + kv = torch.cat([kv, compress_kv], dim=1) + topk_idxs = torch.cat([topk_idxs, compress_idxs], dim=-1) + + o = sparse_attn(q, kv, self.attn_sink, topk_idxs, self.softmax_scale) + apply_rotary_emb(o[..., -rd:], freqs_cis, True) + + # wo_a is block-diagonal over groups (each projects only its own heads), hence einsum not + # Linear. convert.py dequantizes it to bf16; an fp8 grouped GEMM would halve the memory. + o = o.view(bsz, seqlen, self.n_local_groups, -1) + wo_a = self.wo_a.weight.view(self.n_local_groups, self.o_lora_rank, -1) + o = torch.einsum("bsgd,grd->bsgr", o, wo_a) + x = self.wo_b(o.flatten(2)) + return x + + # FreeToken: Decoder SWA Bounded Replay oracle. Prefill of the prompt TAIL [s, s + n) as its own + # window segment (a query sees SWA keys in [max(s, i - win + 1), i]) while the global branch + # covers the whole prompt: a kv source compresses ``x_all`` (every token, forward()'s start_pos + # == 0 path), the indexer scores the tail queries at their absolute positions. + def forward_tail(self, x: torch.Tensor, s: int, x_all: torch.Tensor | None = None): + bsz, n, _ = x.size() + win, rd, ratio = self.window_size, self.rope_head_dim, self.compress_ratio + freqs_cis = self.freqs_cis[s : s + n] + + qr = self.q_norm(self.wq_a(x)) + q = self.wq_b(qr).unflatten(-1, (self.n_local_heads, self.head_dim)) + apply_rotary_emb(q[..., -rd:], freqs_cis) + + kv = self.kv_norm(self.wkv(x)) + apply_rotary_emb(kv[..., -rd:], freqs_cis) + act_quant(kv, fp8_block_size, scale_fmt, scale_dtype, True) + # seed the decode ring with the tail's keys at their absolute slots + for i in range(max(0, n - win), n): + self.window_kv_cache[:bsz, (s + i) % win] = kv[:, i] + topk_idxs = get_window_topk_idxs(win, bsz, n, 0).to(x.device) + + if ratio: + end = s + n + compress_len = end // ratio + latent = None + if self.is_kv_source: + assert x_all is not None, "the kv source compresses every prompt token" + latent = self.compressor(x_all, 0) + shared_attn.compress_kv = self.compress_kv_cache + if self.indexer is not None: + if self.indexer.freqs_cis is None: + self.indexer.freqs_cis = self.freqs_cis + # publish the index keys for every group (forward()'s start_pos == 0 path) + L = x_all.size(1) + k = self.indexer.k_norm(self.indexer.wk(latent)) + apply_rotary_emb(k[..., -rd:], self.freqs_cis[: L - L % ratio : ratio]) + fp4_act_quant(k, fp4_block_size, True) + self.indexer.k_cache[:bsz, : k.size(1)] = k + shared_attn.index_k = self.indexer.k_cache + if not self.is_index_source: + idxs = shared_attn.topk_idxs + elif compress_len == 0: + idxs = torch.empty(bsz, n, 0, dtype=torch.int32, device=x.device) + else: + if self.indexer.freqs_cis is None: + self.indexer.freqs_cis = self.freqs_cis + idxs = self.indexer.forward_tail(x, qr, s, n, kv.size(1)) + shared_attn.topk_idxs = idxs + if latent is not None: + L = x_all.size(1) + apply_rotary_emb(latent[..., -rd:], self.freqs_cis[: L - L % ratio : ratio]) + fp4_act_quant(latent, 16, True, scale_dtype=torch.float8_e4m3fn) + self.compress_kv_cache[:bsz, : latent.size(1)] = latent + kv = torch.cat([kv, shared_attn.compress_kv[:bsz, :compress_len]], dim=1) + topk_idxs = torch.cat([topk_idxs, idxs], dim=-1) + + o = sparse_attn(q, kv, self.attn_sink, topk_idxs, self.softmax_scale) + apply_rotary_emb(o[..., -rd:], freqs_cis, True) + o = o.view(bsz, n, self.n_local_groups, -1) + wo_a = self.wo_a.weight.view(self.n_local_groups, self.o_lora_rank, -1) + o = torch.einsum("bsgd,grd->bsgr", o, wo_a) + return self.wo_b(o.flatten(2)) + + +class Gate(nn.Module): + """MoE gating. The correction bias steers expert selection only; the routing weights come from the + unbiased scores. Image-span tokens use a separate bias (training `noaux_tc_for_vl`).""" + + def __init__(self, layer_id: int, args: ModelArgs): + super().__init__() + n_routed_experts, n_activated_experts = args.get_moe_config(layer_id) + self.dim = args.dim + self.topk = n_activated_experts + self.score_func = args.score_func + self.gate_temp = args.gate_temp + self.norm_topk_prob = args.norm_topk_prob + self.route_scale = args.route_scale + self.weight = nn.Parameter(torch.empty(n_routed_experts, args.dim)) + self.bias = nn.Parameter(torch.empty(n_routed_experts, dtype=torch.float32)) + self.bias_vl = nn.Parameter(torch.empty(n_routed_experts, dtype=torch.float32)) if args.vision_enabled else None + + def forward(self, x: torch.Tensor, image_mask: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor]: + """x: [n, dim]; image_mask: [n] bool, True for tokens inside an image span.""" + scores = linear(x.float(), self.weight.float()) / self.gate_temp + if self.score_func == "softmax": + scores = scores.softmax(dim=-1) + elif self.score_func == "sigmoid": + scores = scores.sigmoid() + else: + scores = F.softplus(scores).sqrt() + bias = self.bias + if image_mask is not None and self.bias_vl is not None: + bias = torch.where(image_mask.unsqueeze(-1), self.bias_vl, bias) + # the bias picks experts but does not scale them: weights come from the raw scores + indices = (scores + bias).topk(self.topk, dim=-1)[1] + weights = scores.gather(1, indices) + if self.norm_topk_prob and self.topk > 1: + weights /= weights.sum(dim=-1, keepdim=True) + 1e-20 # not norm_eps, matches training + weights *= self.route_scale + return weights, indices + + +class Expert(nn.Module): + """One SwiGLU FFN. The clamps come straight from training, where they keep fp8/fp4 activations in + range: the up branch is clamped on both sides, the gate branch only from above.""" + + def __init__(self, dim: int, inter_dim: int, dtype=None, swiglu_limit=0.0): + super().__init__() + self.w1 = Linear(dim, inter_dim, dtype=dtype) + self.w2 = Linear(inter_dim, dim, dtype=dtype) + self.w3 = Linear(dim, inter_dim, dtype=dtype) + self.swiglu_limit = swiglu_limit + + def forward(self, x: torch.Tensor, weights: torch.Tensor | None = None) -> torch.Tensor: + dtype = x.dtype + gate = self.w1(x).float() + up = self.w3(x).float() + if self.swiglu_limit > 0: + up = torch.clamp(up, min=-self.swiglu_limit, max=self.swiglu_limit) + gate = torch.clamp(gate, max=self.swiglu_limit) + x = F.silu(gate) * up + if weights is not None: + x = weights * x + return self.w2(x.to(dtype)) + + +class MoE(nn.Module): + """Top-k routed experts plus one shared expert every token goes through. Experts are split + across ranks, so `self.experts` is None for those another rank owns.""" + + def __init__(self, layer_id: int, args: ModelArgs): + super().__init__() + n_routed_experts, n_activated_experts = args.get_moe_config(layer_id) + self.layer_id = layer_id + self.dim = args.dim + assert n_routed_experts % world_size == 0, ( + f"Number of experts must be divisible by world size (world_size={world_size})" + ) + self.n_routed_experts = n_routed_experts + self.n_local_experts = n_routed_experts // world_size + self.n_activated_experts = n_activated_experts + self.experts_start_idx = rank * self.n_local_experts + self.experts_end_idx = self.experts_start_idx + self.n_local_experts + self.gate = Gate(layer_id, args) + expert_dtype = torch.float4_e2m1fn_x2 if args.expert_dtype == "fp4" else None + self.experts = nn.ModuleList( + [ + Expert( + args.dim, + args.moe_inter_dim, + dtype=expert_dtype, + swiglu_limit=args.swiglu_limit, + ) + if self.experts_start_idx <= i < self.experts_end_idx + else None + for i in range(self.n_routed_experts) + ] + ) + assert args.n_shared_experts == 1 + self.shared_experts = Expert(args.dim, args.moe_inter_dim, swiglu_limit=args.swiglu_limit) + + def forward(self, x: torch.Tensor, image_mask: torch.Tensor | None = None) -> torch.Tensor: + shape = x.size() + x = x.view(-1, self.dim) + weights, indices = self.gate(x, None if image_mask is None else image_mask.flatten()) + y = torch.zeros_like(x, dtype=torch.float32) + counts = torch.bincount(indices.flatten(), minlength=self.n_routed_experts).tolist() + for i in range(self.experts_start_idx, self.experts_end_idx): + if counts[i] == 0: + continue + expert = self.experts[i] + idx, top = torch.where(indices == i) + y[idx] += expert(x[idx], weights[idx, top, None]) + if world_size > 1: + dist.all_reduce(y) + y += self.shared_experts(x) + return y.type_as(x).view(shape) + + +class Block(nn.Module): + """A block whose residual stream is `hc_mult` parallel copies (Hyper-Connections). + + Attention and FFN each sit between `hc_pre` (collapse the copies into one sublayer input) and + `hc_post` (expand back out, mixing the residual in through `comb`). `hc_mixes` derives all three + coefficient sets from the stream itself, `comb` made doubly stochastic by Sinkhorn. + + The coefficients a sublayer computes are used by the *next* one -- see `forward`. + """ + + attention_cls = Attention + + def __init__( + self, + layer_id: int, + args: ModelArgs, + engram_layout: EngramLayout | None = None, + ): + super().__init__() + self.layer_id = layer_id + self.norm_eps = args.norm_eps + self.attn = self.attention_cls(layer_id, args) + self.ffn = MoE(layer_id, args) + self.engram = None + if engram_layout is not None and layer_id in engram_layout.layer_ids: + self.engram = Engram(args, layer_id, engram_layout) + self.attn_norm = RMSNorm(args.dim, self.norm_eps) + self.ffn_norm = RMSNorm(args.dim, self.norm_eps) + self.hc_mult = hc_mult = args.hc_mult + self.hc_sinkhorn_iters = args.hc_sinkhorn_iters + self.hc_eps = args.hc_eps + mix_hc = (2 + hc_mult) * hc_mult + hc_dim = hc_mult * args.dim + with set_dtype(torch.float32): + self.hc_attn_fn = nn.Parameter(torch.empty(mix_hc, hc_dim)) + self.hc_ffn_fn = nn.Parameter(torch.empty(mix_hc, hc_dim)) + self.hc_attn_base = nn.Parameter(torch.empty(mix_hc)) + self.hc_ffn_base = nn.Parameter(torch.empty(mix_hc)) + self.hc_attn_scale = nn.Parameter(torch.empty(3)) + self.hc_ffn_scale = nn.Parameter(torch.empty(3)) + + def hc_mixes(self, x: torch.Tensor, hc_fn: torch.Tensor, hc_scale: torch.Tensor, hc_base: torch.Tensor): + """x: [b,s,hc,d], hc_fn: [mix_hc, hc*d], hc_scale: [3], hc_base: [mix_hc]. Returns the + pre / post / comb coefficients, split out of one projection of the flattened stream.""" + # normalized over the whole flattened hc*d stream, one statistic per token + x = x.flatten(2).float() + rsqrt = torch.rsqrt(x.square().mean(-1, keepdim=True) + self.norm_eps) + mixes = F.linear(x, hc_fn) * rsqrt + return hc_split_sinkhorn(mixes, hc_scale, hc_base, self.hc_mult, self.hc_sinkhorn_iters, self.hc_eps) + + def hc_pre(self, x: torch.Tensor, pre_mix: torch.Tensor): + """Collapse the hc copies into one, weighted by pre_mix. [b,s,hc,d] x [b,s,hc] -> [b,s,d]""" + y = torch.sum(pre_mix.unsqueeze(-1) * x.float(), dim=2) + return y.to(x.dtype) + + def hc_post(self, x: torch.Tensor, residual: torch.Tensor, post: torch.Tensor, comb: torch.Tensor): + """Expand the sublayer output back to hc copies and mix the residual in through `comb`. + x: [b,s,d], residual: [b,s,hc,d], post: [b,s,hc], comb: [b,s,hc,hc] -> [b,s,hc,d]""" + y = post.unsqueeze(-1) * x.unsqueeze(-2) + torch.sum(comb.unsqueeze(-1) * residual.unsqueeze(-2), dim=2) + return y.type_as(x) + + def forward( + self, + x: torch.Tensor, + start_pos: int, + pre_mix: torch.Tensor, + image_mask: torch.Tensor | None, + *attn_args, + ) -> tuple[torch.Tensor, torch.Tensor]: + """`pre_mix` collapses the hc_mult copies down to one input for this block's attention. Each + sub-block's own `hc_mixes` produces the mix for the *next* one, so attention uses what the + previous layer's FFN produced and the FFN uses what this attention produced. + + image_mask: [b, s] bool, True inside image spans (selects the VL routing bias).""" + residual = x + attn_pre, attn_post, attn_comb = self.hc_mixes(x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base) + x = self.hc_pre(x, pre_mix) + x = self.attn_norm(x) + x = self.attn(x, start_pos, *attn_args) + x = self.hc_post(x, residual, attn_post, attn_comb) + + residual = x + ffn_pre, ffn_post, ffn_comb = self.hc_mixes(x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base) + x = self.hc_pre(x, attn_pre) + x = self.ffn_norm(x) + x = self.ffn(x, image_mask) + x = self.hc_post(x, residual, ffn_post, ffn_comb) + return x, ffn_pre + + # FreeToken: Decoder SWA Bounded Replay oracle. ``h`` / ``pre_mix`` are the tail [s, N) streams, + # or -- at the decoder's kv source -- ``h_all`` / ``pre_mix_all`` cover every token so the + # source compresses them all before the block narrows to the tail. + def forward_tail(self, h, s: int, pre_mix, h_all=None, pre_mix_all=None): + if h_all is not None: + attn_pre, attn_post, attn_comb = self.hc_mixes(h_all, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base) + x_all = self.attn_norm(self.hc_pre(h_all, pre_mix_all)) + residual = h_all[:, s:] + attn_pre, attn_post, attn_comb = attn_pre[:, s:], attn_post[:, s:], attn_comb[:, s:] + x = self.attn.forward_tail(x_all[:, s:], s, x_all) + else: + residual = h + attn_pre, attn_post, attn_comb = self.hc_mixes(h, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base) + x = self.attn_norm(self.hc_pre(h, pre_mix)) + x = self.attn.forward_tail(x, s) + x = self.hc_post(x, residual, attn_post, attn_comb) + + residual = x + ffn_pre, ffn_post, ffn_comb = self.hc_mixes(x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base) + x = self.hc_pre(x, attn_pre) + x = self.ffn_norm(x) + x = self.ffn(x, None) + x = self.hc_post(x, residual, ffn_post, ffn_comb) + return x, ffn_pre + + +class ParallelHead(nn.Module): + def __init__(self, vocab_size: int, dim: int, norm_eps: float = 1e-6, hc_eps: float = 1e-6): + super().__init__() + self.vocab_size = vocab_size + self.dim = dim + self.norm_eps = norm_eps + self.hc_eps = hc_eps + self.part_vocab_size = vocab_size // world_size + # bf16 in the checkpoint, kept as fp32 here so the logits come out in fp32 directly + self.weight = nn.Parameter(torch.empty(self.part_vocab_size, self.dim, dtype=torch.float32)) + + def forward(self, x: torch.Tensor, full_logits=False): + """x: [b, s, d]. Generation only needs the last position, so that is the default.""" + if not full_logits: + x = x[:, -1] + logits = F.linear(x.float(), self.weight) + if world_size > 1: + all_logits = [torch.empty_like(logits) for _ in range(world_size)] + dist.all_gather(all_logits, logits) + logits = torch.cat(all_logits, dim=-1) + return logits + + +@lru_cache(1) +def get_dspark_topk_idxs(window_size: int, bsz: int, block_size: int, start_pos: int): + assert start_pos > 0 + matrix = torch.cat( + [ + torch.arange(min(window_size, start_pos + 1)), + window_size + torch.arange(block_size), + ] + ) + return matrix.int().view(1, 1, -1).expand(bsz, block_size, -1).contiguous() + + +class DSparkAttention(Attention): + def forward(self, x: torch.Tensor, start_pos: int, main_x: torch.Tensor): + assert self.compress_ratio == 0 + bsz, seqlen, _ = main_x.size() + win = self.window_size + rd = self.rope_head_dim + + main_freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen] + main_kv = self.kv_norm(self.wkv(main_x)) + apply_rotary_emb(main_kv[..., -rd:], main_freqs_cis) + act_quant(main_kv, fp8_block_size, scale_fmt, scale_dtype, True) + + if start_pos == 0: + if seqlen <= win: + self.window_kv_cache[:bsz, :seqlen] = main_kv + else: + cutoff = seqlen % win + self.window_kv_cache[:bsz, cutoff:win], self.window_kv_cache[:bsz, :cutoff] = main_kv[:, -win:].split( + [win - cutoff, cutoff], dim=1 + ) + return x + + bsz, block_size, _ = x.size() + freqs_cis = self.freqs_cis[start_pos + seqlen : start_pos + seqlen + block_size] + + qr = self.q_norm(self.wq_a(x)) + q = self.wq_b(qr).unflatten(-1, (self.n_local_heads, self.head_dim)) + apply_rotary_emb(q[..., -rd:], freqs_cis) + kv = self.kv_norm(self.wkv(x)) + apply_rotary_emb(kv[..., -rd:], freqs_cis) + act_quant(kv, fp8_block_size, scale_fmt, scale_dtype, True) + + topk_idxs = get_dspark_topk_idxs(win, bsz, block_size, start_pos) + self.window_kv_cache[:bsz, start_pos % win] = main_kv.squeeze(1) + kv = torch.cat([self.window_kv_cache[:bsz], kv], dim=1) + o = sparse_attn(q, kv, self.attn_sink, topk_idxs, self.softmax_scale) + apply_rotary_emb(o[..., -rd:], freqs_cis, True) + + o = o.view(bsz, block_size, self.n_local_groups, -1) + wo_a = self.wo_a.weight.view(self.n_local_groups, self.o_lora_rank, -1) + o = torch.einsum("bsgd,grd->bsgr", o, wo_a) + x = self.wo_b(o.flatten(2)) + return x + + +class DSparkMarkovHead(nn.Module): + def __init__(self, vocab_size: int, dspark_markov_rank: int): + super().__init__() + self.embed = ParallelEmbedding(vocab_size, dspark_markov_rank) + self.head = ParallelHead(vocab_size, dspark_markov_rank) + + def forward(self, token_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + embed = self.embed(token_ids) + logits = self.head(embed, full_logits=True) + return logits, embed + + +class DSparkConfidenceHead(nn.Module): + def __init__(self, input_dim: int): + super().__init__() + # proj in the checkpoint is stored in bf16, while the parameter here is stored in fp32 for fp32 confidence score. + self.proj = Linear(input_dim, 1, dtype=torch.float32) + + def forward(self, hidden: torch.Tensor, markov_embed: torch.Tensor): + hidden = torch.cat([hidden, markov_embed], dim=-1) + return self.proj(hidden.float()).squeeze(-1) + + +class DSparkBlock(Block): + """DSpark stage stored under the mtp.* checkpoint namespace.""" + + attention_cls = DSparkAttention + + def __init__(self, layer_id: int, args: ModelArgs): + super().__init__(layer_id, args) + stage_id = layer_id - args.n_layers + self.block_size = args.dspark_block_size + self.noise_token_id = args.dspark_noise_token_id + self.temperature = args.temperature + if stage_id == 0: + assert len(args.dspark_target_layer_ids) > 0, "DSpark needs target layers" + self.main_proj = Linear(args.dim * len(args.dspark_target_layer_ids), args.dim) + self.main_norm = RMSNorm(args.dim, args.norm_eps) + if stage_id == args.n_mtp_layers - 1: + self.norm = RMSNorm(args.dim, args.norm_eps) + self.markov_head = DSparkMarkovHead(args.vocab_size, args.dspark_markov_rank) + self.confidence_head = DSparkConfidenceHead(args.dim + args.dspark_markov_rank) + self.embed: ParallelEmbedding | None = None + self.head: ParallelHead | None = None + + def forward(self, x: torch.Tensor, start_pos: int, pre_mix: torch.Tensor, main_x: torch.Tensor): + if start_pos == 0: + self.attn(x, start_pos, main_x) # prefill only seeds the window KV cache + return x, pre_mix + return super().forward(x, start_pos, pre_mix, None, main_x) # drafts are text: no VL bias + + def forward_embed(self, main_hidden: torch.Tensor, input_ids: torch.Tensor): + assert self.embed is not None + main_x = self.main_norm(self.main_proj(main_hidden)) + draft_input_ids = input_ids.new_full([input_ids.size(0), self.block_size], self.noise_token_id) + draft_input_ids[:, 0] = input_ids + x = self.embed(draft_input_ids) + x = x.unsqueeze(2).repeat(1, 1, self.hc_mult, 1) + return x, main_x + + def forward_head( + self, + x: torch.Tensor, + pre_mix: torch.Tensor, + input_ids: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + assert self.head is not None + x = self.hc_pre(x, pre_mix) + logits = self.head(self.norm(x), full_logits=True) + output_ids = input_ids.new_empty(input_ids.size(0), self.block_size + 1) + output_ids[:, 0] = input_ids + markov_embeds = [] + for i in range(self.block_size): + logits_bias, markov_embed = self.markov_head(output_ids[:, i]) + logits[:, i].add_(logits_bias) + markov_embeds.append(markov_embed) + output_ids[:, i + 1] = sample(logits[:, i], self.temperature) + markov_embed = torch.stack(markov_embeds, dim=1) + confidence = self.confidence_head(x, markov_embed) + return output_ids, logits, confidence + + +def make_identity_pre_mix(x: torch.Tensor, hc_mult: int) -> torch.Tensor: + """initial one-hot mix""" + pre_mix = x.new_zeros(x.size(0), x.size(1), hc_mult, dtype=torch.float32) + pre_mix[:, :, 0] = 1.0 + return pre_mix + + +class SharedAttentionRuntime: + """What attention layers hand down the stack instead of recomputing. Layers run in order and + every source writes before its consumers read, so one slot each is enough and nothing needs + resetting between forwards. Sources: compress_kv and index_k from kv_source_layers, topk_idxs + from index_source_layers, candidates from candidate_source_layer.""" + + def __init__(self): + self.compress_kv: torch.Tensor | None = None + self.index_k: torch.Tensor | None = None + self.topk_idxs: torch.Tensor | None = None + self.candidates: torch.Tensor | None = None + + +# Only ever one model per process, same as world_size / rank / default_dtype above. +shared_attn = SharedAttentionRuntime() + + +class Transformer(nn.Module): + """DeepSeek-V4.1: embed -> expand to hc_mult copies -> blocks -> collapse -> logits. Building + this sets the globals at the top of the file. The tokenizer only feeds the engram token map.""" + + def __init__(self, args: ModelArgs, tokenizer=None): + global world_size, rank, default_dtype + world_size = dist.get_world_size() if dist.is_initialized() else 1 + rank = dist.get_rank() if dist.is_initialized() else 0 + default_dtype = torch.float8_e4m3fn if args.dtype == "fp8" else torch.bfloat16 + super().__init__() + self.max_seq_len = args.max_seq_len + self.temperature = args.temperature + self.norm_eps = args.norm_eps + self.hc_eps = args.hc_eps + self.engram_layout = EngramLayout.from_args(args) + self.engram_hash = ( + NgramHashState(args, self.engram_layout, tokenizer) if self.engram_layout is not None else None + ) + self.embed = ParallelEmbedding(args.vocab_size, args.dim) + self.layers = torch.nn.ModuleList() + for layer_id in range(args.n_layers): + self.layers.append(Block(layer_id, args, self.engram_layout)) + self.norm = RMSNorm(args.dim, self.norm_eps) + self.head = ParallelHead(args.vocab_size, args.dim, self.norm_eps, self.hc_eps) + self.mtp = torch.nn.ModuleList() + self.target_layer_ids = args.dspark_target_layer_ids + if args.dspark_block_size: + for layer_id in range(args.n_mtp_layers): + self.mtp.append(DSparkBlock(args.n_layers + layer_id, args)) + self.mtp[-1].embed = self.embed + self.mtp[-1].head = self.head + self.hc_mult = args.hc_mult + self.vision = None + if args.vision_enabled: + self.vision = ViT(args) + self.aligner = Aligner(args) + # learned embeddings for the image span delimiters + self.image_start = nn.Parameter(torch.empty(args.dim)) + self.image_end = nn.Parameter(torch.empty(args.dim)) + self.image_newline = nn.Parameter(torch.empty(args.dim)) + + @torch.inference_mode() + def encode_image(self, patches: torch.Tensor, n_vit_h: int, n_vit_w: int) -> torch.Tensor: + return self.aligner(self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w) + + def merge_image_embeddings(self, images, h: torch.Tensor): + """Overwrite each image's token span in h with its ViT/aligner features. The IMAGE slots take + the aligner rows in row-major order; the span delimiters take learned embeddings.""" + for i, sample in enumerate(images): + for img in sample or (): + types = img.types.to(h.device) + span = h[i, img.start : img.start + types.numel()] + span[types == IMAGE_START] = self.image_start.to(h.dtype) + span[types == IMAGE_END] = self.image_end.to(h.dtype) + span[types == IMAGE_NEW_LINE] = self.image_newline.to(h.dtype) + embeds = self.encode_image(img.patches.to(h.device), img.n_vit_h, img.n_vit_w) + span[types == IMAGE] = embeds.to(h.dtype) + + @torch.inference_mode() + def forward( + self, input_ids: torch.Tensor, start_pos: int = 0, images=None, token_types: torch.Tensor | None = None + ): + """input_ids: [b, s], every entry a real token id -- generate.py only ever passes positions it + has already filled, so the padding it uses internally never reaches here. token_types / + images carry the VL inputs built by image_processor.prepare_vl_inputs; image spans must lie + inside the first (start_pos 0) chunk.""" + image_mask = None if token_types is None else token_types >= 0 # TEXT is -1 + # image tokens take no part in an n-gram and get no engram contribution; text-only needs no mask + engram_mask = None if image_mask is None else ~image_mask + engram_hashes = self.engram_hash(input_ids, start_pos, engram_mask) if self.engram_hash is not None else None + h = self.embed(input_ids) + if images is not None: + assert start_pos == 0, "image spans must be prefilled in a single chunk" + self.merge_image_embeddings(images, h) + # Expand to hc_mult copies for Hyper-Connections + h = h.unsqueeze(2).repeat(1, 1, self.hc_mult, 1) + main_hiddens = [] + pre_mix = make_identity_pre_mix(h, self.hc_mult) + for i, layer in enumerate(self.layers): + if layer.engram is not None: + h = layer.engram(h, engram_hashes[:, :, layer.engram.layer_hash_index, :], engram_mask) + # the MTP head reads the attention input of its target layers, not their output + if i in self.target_layer_ids: + main_hiddens.append(h.mean(dim=2)) + h, pre_mix = layer(h, start_pos, pre_mix, image_mask) + h = layer.hc_pre(h, pre_mix) + logits = self.head(self.norm(h)) + output_ids = sample(logits, self.temperature) + main_hidden = torch.cat(main_hiddens, dim=-1) if main_hiddens else None + return output_ids, logits, main_hidden + + # FreeToken: Decoder SWA Bounded Replay oracle (tech report sec. 3.2.2): the encoder runs on + # every prompt token, the decoder on the last ``n_win`` with its sliding window truncated + # there; the decoder's kv source (the last kv_source layer) compresses every token. Returns + # the last position's logits and leaves the caches ready for decode from position seqlen. + @torch.inference_mode() + def forward_bounded(self, input_ids: torch.Tensor, n_win: int): + engram_hashes = self.engram_hash(input_ids, 0, None) if self.engram_hash is not None else None + h = self.embed(input_ids).unsqueeze(2).repeat(1, 1, self.hc_mult, 1) + pre_mix = make_identity_pre_mix(h, self.hc_mult) + N = input_ids.size(1) + s = max(0, N - n_win) + dec_start = max(i for i, layer in enumerate(self.layers) if layer.attn.is_kv_source) + for i, layer in enumerate(self.layers): + if layer.engram is not None: + assert i < dec_start, "Engram sits in the encoder" + h = layer.engram(h, engram_hashes[:, :, layer.engram.layer_hash_index, :], None) + if i < dec_start: + h, pre_mix = layer(h, 0, pre_mix, None) + elif i == dec_start: + h, pre_mix = layer.forward_tail(None, s, None, h_all=h, pre_mix_all=pre_mix) + else: + h, pre_mix = layer.forward_tail(h, s, pre_mix) + h = layer.hc_pre(h, pre_mix) + return self.head(self.norm(h)) + + @torch.inference_mode() + def forward_spec(self, input_ids: torch.Tensor, main_hidden: torch.Tensor, start_pos: int = 0): + h, main_x = self.mtp[0].forward_embed(main_hidden, input_ids) + pre_mix = make_identity_pre_mix(h, self.hc_mult) + for layer in self.mtp: + h, pre_mix = layer(h, start_pos, pre_mix, main_x) + if start_pos == 0: + return None + return self.mtp[-1].forward_head(h, pre_mix, input_ids) + + +def sample(logits, temperature: float = 1.0): + """Gumbel-max trick: equivalent to multinomial sampling but faster on GPU, + since it avoids the GPU-to-CPU sync in torch.multinomial.""" + if temperature == 0: + return logits.argmax(dim=-1) + logits = logits / max(temperature, 1e-5) + probs = torch.softmax(logits, dim=-1, dtype=torch.float32) + return probs.div_(torch.empty_like(probs).exponential_(1)).argmax(dim=-1) + + +if __name__ == "__main__": + torch.set_default_dtype(torch.bfloat16) + torch.set_default_device("cuda") + torch.manual_seed(0) + args = ModelArgs(dspark_block_size=6, dspark_target_layer_ids=(3, 4)) + x = torch.randint(0, args.vocab_size, (2, 150)) + model = Transformer(args) + + output_ids, logits, main_hidden = model(x[:, :128]) + model.forward_spec(output_ids, main_hidden) + for i in range(128, 150): + output_ids, logits, main_hidden = model(x[:, i : i + 1], i) + result = model.forward_spec(output_ids, main_hidden, i) + assert result is not None + output_ids, logits, confidence = result diff --git a/tests/models/deepseek_v41/reference/vision.py b/tests/models/deepseek_v41/reference/vision.py new file mode 100644 index 0000000000..2fd86160f7 --- /dev/null +++ b/tests/models/deepseek_v41/reference/vision.py @@ -0,0 +1,121 @@ +# Vendored from deepseek-ai/DeepSeek-V4.1-Flash, revision dba1be0a, inference/vision.py. +# MIT License; see LICENSE in this directory. +from functools import lru_cache + +import torch +import torch.nn.functional as F +from torch import nn + + +@lru_cache(8) +def get_vision_cos_sin(n_h: int, n_w: int, dim: int, theta: float): + inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + hpos = torch.arange(n_h).unsqueeze(1).expand(n_h, n_w) + wpos = torch.arange(n_w).unsqueeze(0).expand(n_h, n_w) + freqs = torch.stack([hpos, wpos], dim=-1).reshape(-1, 2, 1).float() * inv_freq + freqs = freqs.flatten(1) + return freqs.cos().unsqueeze(1), freqs.sin().unsqueeze(1) + + +def apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + dtype = x.dtype + x1, x2 = x.float().chunk(2, dim=-1) + return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1).to(dtype) + + +class RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + dtype = x.dtype + x = x.float() + x = x * torch.rsqrt(x.square().mean(-1, keepdim=True) + self.eps) + return (self.weight * x).to(dtype) + + +class PatchEmbed(nn.Module): + def __init__(self, args): + super().__init__() + self.proj = nn.Linear(3 * args.vision_patch_size**2, args.vision_dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.proj(x.flatten(1)) + + +class Attention(nn.Module): + def __init__(self, args): + super().__init__() + self.n_heads = args.vision_n_heads + self.head_dim = args.vision_dim // args.vision_n_heads + self.wqkv = nn.Linear(args.vision_dim, 3 * args.vision_dim) + self.wo = nn.Linear(args.vision_dim, args.vision_dim) + + def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + n = x.size(0) + q, k, v = (t.view(n, self.n_heads, self.head_dim) for t in self.wqkv(x).chunk(3, dim=-1)) + q = apply_rotary(q, cos, sin) + k = apply_rotary(k, cos, sin) + o = F.scaled_dot_product_attention(q.transpose(0, 1), k.transpose(0, 1), v.transpose(0, 1)) + return self.wo(o.transpose(0, 1).reshape(n, -1)) + + +class MLP(nn.Module): + def __init__(self, args): + super().__init__() + self.w1 = nn.Linear(args.vision_dim, 2 * args.vision_inter_dim, bias=False) + self.w2 = nn.Linear(args.vision_inter_dim, args.vision_dim, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate, up = self.w1(x).chunk(2, dim=-1) + return self.w2(F.silu(gate) * up) + + +class Block(nn.Module): + def __init__(self, args): + super().__init__() + self.norm1 = RMSNorm(args.vision_dim) + self.attn = Attention(args) + self.norm2 = RMSNorm(args.vision_dim) + self.mlp = MLP(args) + + def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + x = x + self.attn(self.norm1(x), cos, sin) + return x + self.mlp(self.norm2(x)) + + +class ViT(nn.Module): + """DeepSeek ViT: full bidirectional attention over one image with 2D RoPE.""" + + def __init__(self, args): + super().__init__() + self.rope_dim = args.vision_dim // args.vision_n_heads // 2 + self.rope_theta = args.vision_rope_theta + self.patch_embed = PatchEmbed(args) + self.blocks = nn.ModuleList([Block(args) for _ in range(args.vision_n_layers)]) + self.norm = RMSNorm(args.vision_dim) + + def forward(self, patches: torch.Tensor, n_h: int, n_w: int) -> torch.Tensor: + x = self.patch_embed(patches) + cos, sin = get_vision_cos_sin(n_h, n_w, self.rope_dim, self.rope_theta) + for block in self.blocks: + x = block(x, cos, sin) + return self.norm(x) + + +class Aligner(nn.Module): + def __init__(self, args): + super().__init__() + self.downsample_ratio = args.vision_downsample_ratio + in_dim = args.vision_dim * self.downsample_ratio**2 + self.w1 = nn.Linear(in_dim, args.dim) + self.w2 = nn.Linear(args.dim, args.dim) + + def forward(self, x: torch.Tensor, n_h: int, n_w: int) -> torch.Tensor: + r = self.downsample_ratio + x = x.view(n_h, n_w, -1).permute(2, 0, 1) + x = F.pad(x, (0, -n_w % r, 0, -n_h % r)) + x = F.unfold(x.unsqueeze(0), r, stride=r).squeeze(0).transpose(0, 1) + return self.w2(F.gelu(self.w1(x))) diff --git a/tests/models/deepseek_v41/test_config.py b/tests/models/deepseek_v41/test_config.py new file mode 100644 index 0000000000..fab385fb93 --- /dev/null +++ b/tests/models/deepseek_v41/test_config.py @@ -0,0 +1,169 @@ +"""Config parsing, the DSV41 layer-role table, registry dispatch and loader / model key parity.""" + +from __future__ import annotations + +from dataclasses import replace + +import pytest +import torch + +from freetoken.layers import set_rope_device +from freetoken.layers.quantization import NoQuantConfig +from freetoken.distributed.info import set_tp_info, try_get_tp_info +from freetoken.models import create_model +from freetoken.models.deepseek_v41.args import DeepseekV41Args, Mode +from freetoken.models.deepseek_v41.config import parse_config +from freetoken.models.register import get_model_spec +from freetoken.utils.torch_utils import torch_dtype + +from .common import ENGRAM_LAYERS, INDEX_SOURCES, KV_SOURCES, N_LAYERS, RATIOS, tiny_hf_config, write_tiny_checkpoint + + +def test_roles_follow_the_dsv41_mode_table(): + args = parse_config(tiny_hf_config()).dsv41_args + modes = [r.mode for r in args.roles] + assert modes == [Mode.WINDOW, Mode.WINDOW, Mode.FULL, Mode.REUSE, Mode.FULL, Mode.REINDEX] + assert [r.kv_source for r in args.roles] == [None, None, 2, 2, 4, 4] + assert args.roles[4].is_candidate_source and args.roles[5].uses_candidates and not args.roles[2].uses_candidates + assert args.decoder_start_layer == 4 and args.backbone_kv_sources == (2, 4) + assert args.freqs_params(0)[1] == 0 and args.freqs_params(2)[1] == 1024 # yarn only on compressing layers + assert args.swa_decoder_replay == "bounded" + + +def test_shipping_layout_roles(): + """The real 40-layer layout: 3 encoder Full layers (ratio 2), the decoder Full at 20 building the + candidate pool, Reindex at 24/28/32/36 inside it, Reuse elsewhere.""" + text = tiny_hf_config().text_config + text.update( + num_hidden_layers=40, compress_ratios=[0, 0] + [2] * 18 + [1] * 20 + [0, 0, 0], + kv_source_layer_ids=[2, 8, 14, 20], index_source_layer_ids=[2, 8, 14, 20, 24, 28, 32, 36], + candidate_source_layer_id=20, num_nextn_predict_layers=3, engram_layer_ids=[1, 14], engram_num_embeddings=[10, 10], + ) + args = DeepseekV41Args.from_hf(tiny_hf_config(**text)) + full = [r.layer_id for r in args.roles if r.mode is Mode.FULL] + reindex = [r.layer_id for r in args.roles if r.mode is Mode.REINDEX] + assert full == [2, 8, 14, 20] and reindex == [24, 28, 32, 36] + assert sum(r.mode is Mode.REUSE for r in args.roles) == 40 - 2 - 4 - 4 + assert [r.layer_id for r in args.roles if r.uses_candidates] == reindex + assert args.decoder_start_layer == 20 + + +def test_bad_layouts_are_rejected(): + with pytest.raises(ValueError, match="no kv source"): + DeepseekV41Args.from_hf(tiny_hf_config(kv_source_layer_ids=[4])) + with pytest.raises(ValueError, match="reads kv source"): + DeepseekV41Args.from_hf(tiny_hf_config(compress_ratios=[0, 0, 2, 1, 1, 1, 0])) + with pytest.raises(ValueError, match="different kv source"): + DeepseekV41Args.from_hf(tiny_hf_config(compress_ratios=[0, 0, 2, 2, 2, 2, 0], kv_source_layer_ids=[2, 4], candidate_source_layer_id=2, index_source_layer_ids=[2, 4, 5])) + + +def test_registry_and_kv_sources(): + spec = get_model_spec("DeepseekV41ForCausalLM") + assert spec.module == "freetoken.models.deepseek_v41" + from freetoken.attention import AttnType + from freetoken.models.config import DSV4AttentionGroupConfig + + mc = parse_config(tiny_hf_config()) + args = mc.dsv41_args + assert args.backbone_kv_sources == (2, 4) and tuple(r.ratio for r in args.roles) == tuple(RATIOS[:N_LAYERS]) + (group,) = mc.attention_groups + assert type(group) is DSV4AttentionGroupConfig and group.variant == "v41" + assert mc.kv_cache_group_specs()[0].attn_type is AttnType.DSV41 and mc.attn_type_for_layer(0) is AttnType.DSV41 + assert mc.num_layers == N_LAYERS and mc.dsv41_args is not None + + +def _meta_model(mc): + if try_get_tp_info() is None: + set_tp_info(0, 1) + set_rope_device(torch.device("cpu")) + mc = replace(mc, moe_strategy="offload", decode_target="gpu", quant=NoQuantConfig()) + with torch.device("meta"), torch_dtype(torch.bfloat16): + return create_model(mc) + + +def test_loader_names_every_model_parameter(tmp_path): + from freetoken.models.deepseek_v41.weight import iter_weights + + write_tiny_checkpoint(str(tmp_path)) + model = _meta_model(parse_config(tiny_hf_config())) + params = model.state_dict() + loaded = dict(iter_weights(str(tmp_path), torch.device("cpu"), include_moe_experts=False)) + assert set(loaded) == set(params), (set(loaded) ^ set(params)) + for name, tensor in loaded.items(): + assert tuple(tensor.shape) == tuple(params[name].shape), name + # routed experts, the DSpark stage, the vision tower and the image routing bias are not resident weights + assert not any(".experts." in k or k.startswith(("mtp.", "vision.")) or k.endswith("bias_vl") for k in loaded) + # per-role structure + assert model.model.layers.op_list[3].attn.compressor is None and model.model.layers.op_list[3].attn.indexer is None + assert model.model.layers.op_list[4].attn.compressor.ratio == 1 and model.model.layers.op_list[2].attn.compressor.ratio == 2 + assert [b.engram is not None for b in model.model.layers.op_list] == [L in ENGRAM_LAYERS for L in range(N_LAYERS)] + assert all(L in INDEX_SOURCES for L in range(N_LAYERS) if model.model.layers.op_list[L].attn.indexer is not None) + assert all(L in KV_SOURCES for L in range(N_LAYERS) if model.model.layers.op_list[L].attn.compressor is not None) + + +@pytest.mark.parametrize("missing_weight", [False, True]) +def test_loader_closes_shards_when_expert_loading_stops(tmp_path, monkeypatch, missing_weight): + from freetoken.models.deepseek_v41 import weight + + write_tiny_checkpoint(str(tmp_path)) + closed = [] + close, get = weight._ShardReader.close, weight._ShardReader.get + + def close_reader(reader): + close(reader) + closed.append(True) + + def read_weight(reader, name): + if missing_weight and name.endswith(".w3.weight"): + raise KeyError(name) + return get(reader, name) + + monkeypatch.setattr(weight._ShardReader, "close", close_reader) + monkeypatch.setattr(weight._ShardReader, "get", read_weight) + tensors = weight.iter_weights(str(tmp_path), torch.device("cpu")) + if missing_weight: + with pytest.raises(KeyError): + next(tensors) + else: + next(tensors) + tensors.close() + assert closed == [True] + + +def test_fp8_wo_a_layer_rejects_a_dequantized_ftw_layout_clearly(): + """An fp8-declared wo_a fed a state dict without ``wo_a_scale`` (an FTW converted before wo_a kept + its payload) fails with a reconvert message, not a raw missing-key error; a bf16-declared layer fed + the fp8 payload + scale dequantizes it on the way in.""" + from types import SimpleNamespace + + from freetoken.layers.quantization.scheme import fp8_block_scheme + from freetoken.models.deepseek_v41.attention import DSV41Attention + + if try_get_tp_info() is None: + set_tp_info(0, 1) + args = parse_config(tiny_hf_config()).dsv41_args + args.max_seq_len = 512 + # wo_a is fp8; the LinearReplicated modules stay unquantized (the bf16 quant method) + plain = NoQuantConfig() + quant = SimpleNamespace( + scheme_for=lambda name: fp8_block_scheme("ue8m0", 32) if name.endswith(".wo_a") else None, + get_quant_method=lambda layer, prefix: plain.get_quant_method(layer, prefix), + ) + with torch.device("meta"), torch_dtype(torch.bfloat16): + fp8_layer = DSV41Attention(args, args.roles[2], quant_config=quant, prefix="attn") + bf16_layer = DSV41Attention(args, args.roles[2], quant_config=None, prefix="attn") + assert fp8_layer.wo_a.dtype == torch.float8_e4m3fn and bf16_layer.wo_a.dtype == torch.bfloat16 + rows, cols = fp8_layer.wo_a.shape + with torch.device("meta"): + state = {k: torch.empty_like(v) for k, v in fp8_layer.state_dict(prefix="attn").items()} + del state["attn.wo_a_scale"] + state["attn.wo_a"] = torch.empty(rows, cols, dtype=torch.bfloat16, device="meta") + with pytest.raises(RuntimeError, match="reconvert"): + fp8_layer.load_state_dict(dict(state), prefix="attn") + # the bf16 layer given the fp8 payload + scale dequantizes (real tensors: the dequant runs) + w = (torch.randn(rows, cols) * 0.05).to(torch.float8_e4m3fn) + scale = torch.full((rows // 32, cols // 32), 127, dtype=torch.uint8).view(torch.float8_e8m0fnu) + state2 = {k: torch.zeros(v.shape, dtype=v.dtype) for k, v in bf16_layer.state_dict(prefix="attn").items()} + state2["attn.wo_a"], state2["attn.wo_a_scale"] = w, scale + bf16_layer.load_state_dict(state2, prefix="attn") + assert torch.equal(bf16_layer.wo_a, w.to(torch.bfloat16)) diff --git a/tests/models/deepseek_v41/test_engine_config.py b/tests/models/deepseek_v41/test_engine_config.py new file mode 100644 index 0000000000..116209ea51 --- /dev/null +++ b/tests/models/deepseek_v41/test_engine_config.py @@ -0,0 +1,81 @@ +"""Engine config resolution for DeepSeek-V4.1 on the tiny checkpoint: the DSV41 backend and pool are +picked, the page size is the window page, the replay knob lands on the model args, and the fp8 +dialect (block 32, fp4 experts) reaches the quant layer with the matching activation block.""" + +from __future__ import annotations + +import json +import os + +import pytest +import torch + +from freetoken.distributed import DistributedInfo +from freetoken.scheduler.config import SchedulerConfig +from freetoken.engine.engine import _adjust_config +from freetoken.kvcache import resolve_pool_class +from freetoken.kvcache.dsv4.v41_pool import DSV41PagedKVCache + +from .common import write_tiny_checkpoint + + +def _engine_config(path, **over): + return SchedulerConfig(model_path=path, tp_info=DistributedInfo(rank=0, size=1), dtype=torch.bfloat16, **over) + + +@pytest.mark.parametrize("replay", ["bounded", "exact"]) +def test_resolution_picks_dsv41_and_the_replay_knob(tmp_path, monkeypatch, replay): + from freetoken.engine import engine + + monkeypatch.setattr(engine, "is_sm100_family", lambda: False) + monkeypatch.setattr(engine, "is_sm90_family", lambda: True) + write_tiny_checkpoint(str(tmp_path)) + config = _engine_config(str(tmp_path), attention_backend="auto", moe_strategy="offload", swa_decoder_replay=replay, max_seq_len_override=2048) + _adjust_config(config) + assert config.attention_backend == "dsv41_sparse" + assert config.page_size == 128 and config.cache_type == "swa_radix" + assert resolve_pool_class(config.model_config) is DSV41PagedKVCache + args = config.model_config.dsv41_args + assert args.swa_decoder_replay == replay and args.max_seq_len == 2048 and args.max_batch_size == config.max_running_req + 1 + assert config.max_extend_tokens == 8192 # the prefill chunk stays bounded (whole window pages) + # the cache contract follows the replay mode: a bounded-mode prefix hit must leave the prompt's last + # window to the prefill (the pool's prefix_replay_tokens); the history a resume reads stays one window + from freetoken.kvcache.dsv4.v41_cost_model import dsv41_pool_sizes + + spec = config.model_config.kv_cache_group_specs()[0] + pool = DSV41PagedKVCache(dsv41_pool_sizes(16, args, 1.0, 128), args, torch.device("cpu")) + assert pool.prefix_replay_tokens == (128 if replay == "bounded" else 0) + assert pool.sliding_window_size == 128 and spec.sliding_window == 128 and not spec.is_swa + # bounded replay keeps the decoder's per-request window KV in private rings, off the shared pages + assert pool.private_window_layer_ids == (tuple(range(args.decoder_start_layer, args.n_layers)) if replay == "bounded" else ()) + assert DSV41PagedKVCache.min_kv_tokens(config) // 128 == 8 + 3 * config.max_running_req + 2 * (config.max_running_req + 1) + 1 + + +def test_fp8_block32_dialect_reaches_the_quant_layer(tmp_path): + """A V4.1-style quantization_config on the tiny checkpoint: dense linears get the 32-block + e8m0 scheme, routed experts the MXFP4 scheme with a 32-wide activation block.""" + from freetoken.distributed.info import set_tp_info, try_get_tp_info + from freetoken.layers import OffloadMoELayer + from freetoken.layers.quantization import Fp8BlockConfig, QuantKind + from freetoken.layers.quantization.scheme import fp8_block_size + + write_tiny_checkpoint(str(tmp_path)) + cfg_path = os.path.join(str(tmp_path), "config.json") + with open(cfg_path) as f: + cfg = json.load(f) + cfg["quantization_config"] = {"quant_method": "fp8", "activation_scheme": "dynamic", "weight_block_size": [32, 32], "scale_fmt": "ue8m0", "expert_dtype": "fp4"} + with open(cfg_path, "w") as f: + json.dump(cfg, f) + config = _engine_config(str(tmp_path), moe_strategy="offload") + quant = config.model_config.quant + assert type(quant) is Fp8BlockConfig and quant.block == 32 + dense = quant.scheme_for("model.layers.3.attn.wq_a") + assert dense.kind is QuantKind.FP8_BLOCK and fp8_block_size(dense) == 32 + if try_get_tp_info() is None: + set_tp_info(0, 1) + experts = OffloadMoELayer(3, num_experts=4, top_k=2, hidden_size=64, intermediate_size=64, limit=10.0, quant_config=quant, prefix="model.layers.3.ffn.experts") + assert experts.quant_method.scheme.kind is QuantKind.MXFP4 and experts.quant_method.cfg.act_block == 32 + # the family's bf16 modules stay unquantized although the fp8 config lists no exceptions + for name in ("model.layers.2.attn.compressor.wkv", "model.layers.2.attn.indexer.wk", "model.layers.5.attn.indexer.weights_proj", "head"): + assert quant.scheme_for(name) is None, name + assert quant.scheme_for("model.layers.1.engram.wkv").kind is QuantKind.FP8_BLOCK diff --git a/tests/models/deepseek_v41/test_engram.py b/tests/models/deepseek_v41/test_engram.py new file mode 100644 index 0000000000..dec852ec03 --- /dev/null +++ b/tests/models/deepseek_v41/test_engram.py @@ -0,0 +1,281 @@ +"""Engram: the hash against a transcription of the reference ``NgramHashState``, the prime buckets, +the disk table against a torch oracle over the synthetic checkpoint (eager, decode with context, +CUDA-graph flag sync), and the layer's gate math against the reference ``Engram.forward``.""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +from freetoken.models.deepseek_v41.args import DeepseekV41Args +from freetoken.models.deepseek_v41.engram import EngramHash, bucket_primes, hash_multipliers + +from .common import VOCAB, requires_cuda, tiny_hf_config, write_tiny_checkpoint + + +def _args(**over) -> DeepseekV41Args: + return DeepseekV41Args.from_hf(tiny_hf_config(**over)) + + +def _reference_hashes(args: DeepseekV41Args, ids: torch.Tensor, start_pos: int, cache: torch.Tensor) -> torch.Tensor: + """``NgramHashState.forward`` for one sequence with an identity token map; ``cache`` is the + sequence's compressed-id history buffer (mutated like the reference's).""" + layout_primes = bucket_primes(args) + flat = [[p for per_ngram in layer for p in per_ngram] for layer in layout_primes] + offsets = torch.tensor(np.array([np.cumsum([0, *sizes[:-1]]) for sizes in flat])) + primes = torch.tensor(layout_primes) + multipliers = hash_multipliers(args.engram_layer_ids, args.engram_max_ngram_size, args.engram_compressed_vocab_size) + pad_id = args.engram_pad_id + seqlen = ids.numel() + cache[start_pos : start_pos + seqlen] = ids + positions = torch.arange(start_pos, start_pos + seqlen) + tokens, blocked = [], torch.zeros_like(positions, dtype=torch.bool) + for shift in range(args.engram_max_ngram_size): + source = cache[(positions - shift).clamp_min(0)] + blocked = blocked | (positions < shift) + tokens.append(torch.where(blocked, pad_id, source)) + tokens = torch.stack(tokens, dim=-1) + products = tokens.unsqueeze(1) * multipliers + rolling, hashes = products[..., 0], [] + for i in range(1, args.engram_max_ngram_size): + rolling = torch.bitwise_xor(rolling, products[..., i]) + hashes.append(rolling.unsqueeze(-1) % primes[:, i - 1]) + return torch.cat(hashes, dim=-1) + offsets + + +def test_primes_are_distinct_and_ordered(): + args = _args() + primes = bucket_primes(args) + flat = [p for layer in primes for per in layer for p in per] + assert len(flat) == len(set(flat)) and flat == sorted(flat) and min(flat) > args.engram_vocab_size - 1 + from sympy import isprime + + assert all(isprime(p) for p in flat) + + +def test_hash_matches_the_reference_across_a_prefill_decode_split(): + args = _args() + hash = EngramHash(args, list(range(VOCAB)), VOCAB) + ids = torch.randint(3, VOCAB, (37,), generator=torch.Generator().manual_seed(0)) + cache = torch.empty(64, dtype=torch.int64) + want = _reference_hashes(args, ids[:30], 0, cache) + got = hash.row_ids(ids[:30]) + assert torch.equal(got, want) + # decode continues with the three preceding tokens as context + for step in range(30, 37): + want = _reference_hashes(args, ids[step : step + 1], step, cache) + got = hash.row_ids(ids[step : step + 1], context=ids[:step]) + assert torch.equal(got, want), step + # a run resumed mid-sequence (chunked prefill) hashes like the uninterrupted one + whole = hash.row_ids(ids) + assert torch.equal(torch.cat([hash.row_ids(ids[:20]), hash.row_ids(ids[20:], context=ids[:20])]), whole) + assert whole.shape == (37, 1, hash.n_cols) and int(whole.max()) < hash.rows_per_layer[0] + with pytest.raises(ValueError, match="compressed vocab"): + EngramHash(args, list(range(VOCAB)), VOCAB + 1) + + +@requires_cuda +def test_disk_table_matches_the_dequantized_shard(tmp_path): + from freetoken.models.deepseek_v41.engram_table import EngramDiskTable, engram_row_source + + tensors = write_tiny_checkpoint(str(tmp_path)) + args = _args() + hash = EngramHash(args, list(range(VOCAB)), VOCAB) + src = engram_row_source(str(tmp_path), 1) + assert src.num_rows == tensors["layers.1.engram.embed.weight"].shape[0] and src.head_dim == 32 + table = EngramDiskTable(src, hash.n_cols, torch.device("cuda"), max_graph_rows=4, max_extend_tokens=16) + fp8 = tensors["layers.1.engram.embed.weight"] + scale = torch.exp2(tensors["layers.1.engram.embed.scale"].view(torch.uint8).float() - 127.0) + + def oracle(rows: torch.Tensor) -> torch.Tensor: + vals = fp8[rows.flatten()].float().view(-1, 32 // 32, 32) * scale[rows.flatten()].unsqueeze(-1) + return vals.view(rows.shape[0], -1).to(torch.bfloat16).cuda() + + ids = torch.randint(3, VOCAB, (40,), generator=torch.Generator().manual_seed(1)) + rows = hash.row_ids(ids)[:, 0] + table.stage(rows, graph=False) # larger than the initial eager staging: grows + table.flush(signal=False) + assert torch.equal(table.lookup(40), oracle(rows)) + # the graph staging + flag protocol: capture a lookup, fill after the launch, replay + from freetoken.kernel.row_store import probe_wait_sync + + table.wait_sync = probe_wait_sync("auto", torch.device("cuda")) + dec_rows = hash.row_ids(ids[40 - 1 :], context=ids[: 40 - 1])[:, 0] + out = torch.empty(1, hash.n_cols * 32, dtype=torch.bfloat16, device="cuda") + graph = torch.cuda.CUDAGraph() + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + if not table.wait_sync: + table.stage(dec_rows, graph=True) + table.flush(signal=False) + with torch.cuda.graph(graph, stream=stream): + out.copy_(table.lookup(1)) + torch.cuda.synchronize() + graph.replay() # under wait-sync the replay blocks until the host signals + table.stage(dec_rows, graph=True) + table.flush(signal=table.wait_sync) + torch.cuda.synchronize() + assert torch.equal(out, oracle(dec_rows)) + + +@requires_cuda +def test_eager_staging_does_not_overwrite_an_in_flight_lookup(tmp_path): + """The overlap scheduler stages batch k+1 while batch k's lookup copy is still queued behind + the layers before the Engram layer (a sleep kernel here): batch k must read its own rows.""" + from freetoken.models.deepseek_v41.engram_table import EngramDiskTable, engram_row_source + + tensors = write_tiny_checkpoint(str(tmp_path)) + hash = EngramHash(_args(), list(range(VOCAB)), VOCAB) + table = EngramDiskTable(engram_row_source(str(tmp_path), 1), hash.n_cols, torch.device("cuda"), max_graph_rows=4, max_extend_tokens=64) + fp8 = tensors["layers.1.engram.embed.weight"] + scale = torch.exp2(tensors["layers.1.engram.embed.scale"].view(torch.uint8).float() - 127.0) + + def oracle(rows): + vals = fp8[rows.flatten()].float().view(-1, 1, 32) * scale[rows.flatten()].unsqueeze(-1) + return vals.view(rows.shape[0], -1).to(torch.bfloat16).cuda() + + gen = torch.Generator().manual_seed(0) + batches = [hash.row_ids(torch.randint(3, VOCAB, (40,), generator=gen))[:, 0] for _ in range(4)] + got = [] + for rows in batches: + table.stage(rows, graph=False) + table.flush(signal=False) + torch.cuda._sleep(50_000_000) + got.append(table.lookup(rows.shape[0])) + torch.cuda.synchronize() + for i, (rows, out) in enumerate(zip(batches, got)): + assert torch.equal(out, oracle(rows)), f"batch {i} read another batch's rows" + + +@requires_cuda +@torch.inference_mode() +def test_host_fills_a_graph_decode_step_from_a_worker_thread(tmp_path): + """The deferred graph-decode fill must not wait for the launch call to return: a 40-layer graph + blocks cuGraphLaunch on the host while the GPU sits on the WAIT, so the host coordinator reads the + tokens back and signals from a worker thread. The worker hashes against the DISPATCH snapshot of + the request state: advancing the request right after dispatch (what the engine does) must not + change which rows are fetched.""" + from types import SimpleNamespace + + from freetoken.models.deepseek_v41.engram_table import EngramDiskTable, EngramHost, engram_row_source + + tensors = write_tiny_checkpoint(str(tmp_path)) + args = _args() + hash = EngramHash(args, list(range(VOCAB)), VOCAB) + table = EngramDiskTable(engram_row_source(str(tmp_path), 1), hash.n_cols, torch.device("cuda"), max_graph_rows=4, max_extend_tokens=16) + host = EngramHost(hash, [table], torch.device("cuda")) + if not host.wait_sync: + pytest.skip("stream memops unavailable: launch-gating mode needs no worker") + ids = torch.randint(3, VOCAB, (20,), generator=torch.Generator().manual_seed(7)) + req = SimpleNamespace(input_ids=ids.clone(), device_len=20, cached_len=19) + batch = SimpleNamespace(is_decode=True, is_prefill=False, padded_reqs=[req], padded_size=1, input_ids=ids[19:20].to(torch.int32).cuda()) + out = torch.empty(1, hash.n_cols * 32, dtype=torch.bfloat16, device="cuda") + stream = torch.cuda.Stream() + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream): + with torch.cuda.graph(graph, stream=stream): + out.copy_(table.lookup(1)) + torch.cuda.synchronize() + with torch.cuda.stream(stream): + # slow the readback so the request advances (complete_one + append) before the worker runs + torch.cuda._sleep(200_000_000) + with host.forward_host_ctx(batch, use_graph=True): + graph.replay() + req.device_len += 1 + req.input_ids = torch.cat([req.input_ids, torch.tensor([5], dtype=ids.dtype)]) + torch.cuda.synchronize() + fp8 = tensors["layers.1.engram.embed.weight"] + rows = hash.row_ids(ids[19:20], context=ids[:19])[:, 0].flatten() + want = fp8[rows].float().view(1, -1).to(torch.bfloat16).cuda() # scales are all 1.0 in the synthetic table + assert torch.equal(out, want) + host.close() + + +@requires_cuda +def test_host_surfaces_a_failed_deferred_fill_at_dispatch_exit(tmp_path): + from types import SimpleNamespace + + from freetoken.models.deepseek_v41.engram_table import EngramDiskTable, EngramHost, engram_row_source + + write_tiny_checkpoint(str(tmp_path)) + args = _args() + hash = EngramHash(args, list(range(VOCAB)), VOCAB) + table = EngramDiskTable(engram_row_source(str(tmp_path), 1), hash.n_cols, torch.device("cuda"), max_graph_rows=4, max_extend_tokens=16) + host = EngramHost(hash, [table], torch.device("cuda")) + if not host.wait_sync: + pytest.skip("stream memops unavailable") + ids = torch.randint(3, VOCAB, (8,), generator=torch.Generator().manual_seed(8)) + req = SimpleNamespace(input_ids=ids, device_len=8, cached_len=7) + batch = SimpleNamespace(is_decode=True, is_prefill=False, padded_reqs=[req], padded_size=1, input_ids=ids[7:8].to(torch.int32).cuda()) + stream = torch.cuda.Stream() + graph = torch.cuda.CUDAGraph() + out = torch.empty(1, hash.n_cols * 32, dtype=torch.bfloat16, device="cuda") + with torch.cuda.stream(stream): + with torch.cuda.graph(graph, stream=stream): + out.copy_(table.lookup(1)) + torch.cuda.synchronize() + table.stage = lambda *a, **k: (_ for _ in ()).throw(RuntimeError("disk gone")) # the fill fails + with pytest.raises(RuntimeError, match="disk gone"): + with torch.cuda.stream(stream): + with host.forward_host_ctx(batch, use_graph=True): + graph.replay() + torch.cuda.synchronize() # the failing worker released the flag: the stream is not stuck + host.close() + + +def test_ftw_side_files_carry_the_tables_and_reopen(tmp_path): + """An FTW conversion copies the shards holding the Engram tables next to it (the converter's + ``ftw_side_files`` hook, which skips the HF index); the row source resolves the copies through + their headers exactly like the HF shards, byte for byte.""" + from freetoken.kernel.row_store import RowStore + from freetoken.models.deepseek_v41.engram_table import engram_row_source + from freetoken.models.deepseek_v41.weight import ftw_side_files + + src, out = tmp_path / "hf", tmp_path / "ftw" + out.mkdir() + tensors = write_tiny_checkpoint(str(src)) + (out / "config.json").write_bytes((src / "config.json").read_bytes()) # the converter copies the metadata + shard = "model-00001-of-00001.safetensors" # the tiny checkpoint's single shard holds the table + assert ftw_side_files(str(src), str(out)) == [shard] + assert (out / shard).read_bytes() == (src / shard).read_bytes() and not (out / "model.safetensors.index.json").exists() + source = engram_row_source(str(out), 1) + assert source.path == str(out / shard) and source.num_rows == 4096 and source.head_dim == 32 + assert torch.equal(source.scales, tensors["layers.1.engram.embed.scale"].view(torch.uint8)) + store = RowStore(paths=[source.path], extent_file=[0], extent_base=[source.base], rows_per_extent=source.num_rows, + row_bytes=source.head_dim, row_stride=source.head_dim, use_io_uring=False) + ids = torch.tensor([0, 7, 4095], dtype=torch.int64) + buf = torch.zeros(3 * 32, dtype=torch.uint8) + store.stage_rows(ids.data_ptr(), 3, buf.data_ptr(), 0) + store.flush(0) + assert torch.equal(buf.view(3, 32), tensors["layers.1.engram.embed.weight"][ids].view(torch.uint8)) + + +def test_layer_gate_matches_the_reference(): + """``EngramLayer.forward`` against the reference ``Engram.forward`` (bf16 wkv, fp32 gate math).""" + from freetoken.distributed.info import set_tp_info, try_get_tp_info + from freetoken.models.deepseek_v41.engram import EngramLayer + + if try_get_tp_info() is None: + set_tp_info(0, 1) + args = _args() + torch.manual_seed(0) + layer = EngramLayer(args, 1) + T, hc, dim = 5, args.hc_mult, args.dim + wkv = torch.randn(dim * (hc + 1), layer.width, dtype=torch.bfloat16) * 0.05 + layer.wkv.weight = wkv + layer.q_weight = torch.rand(hc, dim) + 0.5 + layer.k_weight = torch.rand(hc, dim) + 0.5 + x = torch.randn(T, hc, dim, dtype=torch.bfloat16) + rows = torch.randn(T, layer.width, dtype=torch.bfloat16) + + kv = torch.nn.functional.linear(rows, wkv) + key, value = kv.split([hc * dim, dim], dim=-1) + key = key.float().view(T, hc, dim) + weight = layer.q_weight * layer.k_weight + h = x.float() + rstd = torch.rsqrt(h.square().mean(-1) + args.norm_eps) * torch.rsqrt(key.square().mean(-1) + args.norm_eps) + dot = (h * weight * key).sum(-1) * rstd * dim**-0.5 + gate = torch.sigmoid(torch.copysign(dot.abs().clamp_min(1e-6).sqrt(), dot)) + want = (h + gate.unsqueeze(-1) * value.float().unsqueeze(-2)).to(x.dtype) + assert torch.equal(layer.forward(x, rows), want) diff --git a/tests/models/deepseek_v41/test_forward_smoke.py b/tests/models/deepseek_v41/test_forward_smoke.py new file mode 100644 index 0000000000..e0861e8874 --- /dev/null +++ b/tests/models/deepseek_v41/test_forward_smoke.py @@ -0,0 +1,184 @@ +"""End-to-end forward smoke on the tiny model: cold prefill, batched decode, a second request, chunked +prefill and bounded replay all produce finite logits of the right shape (numerics: test_reference_parity).""" + +from __future__ import annotations + +import pytest +import torch + +from .common import VOCAB, requires_cuda, write_tiny_checkpoint + +pytestmark = requires_cuda + + +@pytest.fixture(scope="module") +def checkpoint(tmp_path_factory): + folder = tmp_path_factory.mktemp("dsv41-tiny") + write_tiny_checkpoint(str(folder)) + return str(folder) + + +def _tokens(n: int, seed: int) -> list[int]: + g = torch.Generator().manual_seed(seed) + return torch.randint(3, VOCAB, (n,), generator=g).tolist() + + +def test_prefill_and_decode(checkpoint): + from .harness import TinyEngine + + eng = TinyEngine(checkpoint, max_seq_len=1024, max_running_req=2, swa_decoder_replay="exact") + r0 = eng.new_request(0, _tokens(300, 1)) + logits = eng.prefill([r0]) + assert logits.shape == (1, VOCAB) and torch.isfinite(logits).all() + eng.finish_prefill([r0]) + r1 = eng.new_request(1, _tokens(50, 2)) + assert torch.isfinite(eng.prefill([r1])).all() + eng.finish_prefill([r1]) + for step in range(5): # batched decode over two requests at different positions + out = eng.decode([r0, r1], [10 + step, 20 + step]) + assert out.shape == (2, VOCAB) and torch.isfinite(out).all() + + +def test_chunked_prefill_matches_single_shot(checkpoint): + """Two 128-aligned chunks must reproduce the single-shot prefill's last-token logits: the + compressor carry, the window ring and the compressed rows are all resumed through the pool. The + chunks run their projections at other row counts, so the match is to rounding, not bitwise.""" + from .harness import TinyEngine + from .test_reference_parity import _compare + + eng = TinyEngine(checkpoint, max_seq_len=1024, max_running_req=2, swa_decoder_replay="exact") + toks = _tokens(300, 3) + whole = eng.new_request(0, toks) + ref = eng.prefill([whole]) + chunked = eng.new_request(1, toks) + chunked.device_len = 256 # first chunk [0, 256) + eng.prefill([chunked]) + chunked.cached_len, chunked.device_len = 256, 300 + got = eng.prefill([chunked]) + _compare("chunked prefill", got, ref) + + +def test_bounded_replay_is_exact_within_one_window(checkpoint): + """A prompt no longer than the window replays every token with no floor: bounded == exact.""" + from .harness import TinyEngine + from .test_reference_parity import _compare + + toks = _tokens(100, 5) + exact = TinyEngine(checkpoint, max_seq_len=1024, max_running_req=2, swa_decoder_replay="exact") + want = exact.prefill([exact.new_request(0, toks)]) + bounded = TinyEngine(checkpoint, max_seq_len=1024, max_running_req=2, swa_decoder_replay="bounded") + got = bounded.prefill([bounded.new_request(0, toks)]) + _compare("bounded within one window", got, want) + + +@pytest.mark.parametrize("cut", [128, 256]) +def test_bounded_chunked_prefill_matches_unchunked(checkpoint, cut): + """A chunked bounded prefill gives each chunk its part of the prompt's last window ``[172, 300)``: + a first chunk that ends before it (cut 128) runs no decoder row and hands the head a placeholder; + one that reaches into it (cut 256) leaves rows in the private ring for the last chunk, which is + shorter than the window (44 tokens). Both match the unchunked prefill.""" + from .harness import TinyEngine + from .test_reference_parity import _compare + + eng = TinyEngine(checkpoint, max_seq_len=1024, max_running_req=2, swa_decoder_replay="bounded") + toks = _tokens(300, 6) + ref = eng.prefill([eng.new_request(0, toks)]) + chunked = eng.new_request(1, toks) + chunked.device_len = cut + first = eng.prefill([chunked]) + assert first.shape == (1, VOCAB) and torch.isfinite(first).all() + chunked.cached_len, chunked.device_len = cut, 300 + _compare(f"bounded chunks cut at {cut}", eng.prefill([chunked]), ref) + + +def test_bounded_batch_mixes_a_decoderless_chunk_with_a_whole_prompt(checkpoint): + """One batch: a chunk that ends before its prompt's last window (no decoder row, a placeholder for + the head) next to a whole prompt; the whole prompt's logits are those of its own prefill.""" + from .harness import TinyEngine + from .test_reference_parity import _compare + + eng = TinyEngine(checkpoint, max_seq_len=2048, max_running_req=2, swa_decoder_replay="bounded") + whole = _tokens(300, 8) + want = eng.prefill([eng.new_request(0, whole)]) + chunk = eng.new_request(0, _tokens(600, 9)) + chunk.device_len = 256 + other = eng.new_request(1, whole) + got = eng.prefill([chunk, other]) + assert got.shape == (2, VOCAB) and torch.isfinite(got).all() + _compare("whole prompt next to a decoderless chunk", got[1:], want) + + +def test_bounded_prefix_hit_matches_cold_prefill(checkpoint): + """A prefix hit that would leave fewer than a window of new tokens (768 cached + 20 new) is capped + one window before the prompt end, so the request prefills ``[640, 788)`` -- the decoder's whole + window among it -- and produces the cold prefill's logits.""" + from .harness import TinyEngine + from .test_reference_parity import _compare + + eng = TinyEngine(checkpoint, max_seq_len=2048, max_running_req=2, swa_decoder_replay="bounded") + prefix = _tokens(768, 11) + donor = eng.new_request(0, prefix + _tokens(20, 12)) + cold = eng.prefill([donor]) + eng.finish_prefill([donor]) + hit = eng.new_request_on_prefix(1, donor, prefix + _tokens(20, 12)) # the same prompt over the shared pages + assert hit.cached_len == 640 and hit.extend_len == 148 + got = eng.prefill([hit]) + _compare("bounded prefix hit", got, cold) + + +def test_shared_prefix_replays_do_not_disturb_each_other(checkpoint): + """Two requests sharing a 768-token prefix with different suffix lengths: B's hit is capped at 640 + (a window before its prompt end), its decoder writes only B's private ring and it reads the shared + encoder history, so A's pending decode sees the logits it would have seen without B, and neither + the shared pages nor A's ring change.""" + from .harness import TinyEngine + + prefix = _tokens(768, 13) + a_toks, b_toks = prefix + _tokens(20, 14), prefix + _tokens(40, 15) + + def run(interleave: bool): + eng = TinyEngine(checkpoint, max_seq_len=2048, max_running_req=2, swa_decoder_replay="bounded") + a = eng.new_request(0, a_toks) + eng.prefill([a]) + eng.finish_prefill([a]) + dec_layer, enc_layer = eng.args.decoder_start_layer, 2 + shared_slots = eng.pool.translate_full_to_window(eng.page_table[0, 640:768].long()) + enc_before = eng.pool.window_pool[enc_layer][shared_slots].clone() + P = eng.args.window_size + ring_before = eng.pool.window_pool[dec_layer][0 * P : 1 * P].clone() # A's ring row + if interleave: + b = eng.new_request_on_prefix(1, a, b_toks) + eng.prefill([b]) + eng.finish_prefill([b]) + assert torch.equal(eng.pool.window_pool[enc_layer][shared_slots], enc_before) + assert torch.equal(eng.pool.window_pool[dec_layer][0 * P : 1 * P], ring_before) + return eng.decode([a], [7]) + + assert torch.equal(run(True), run(False)) + + +def test_commit_dedup_onto_a_longer_prompts_pages_keeps_the_decoder_state(checkpoint): + """The radix commit may replace a request's freshly written pages with an existing node's (the + reviewer's donor case): D prefilled 256 tokens, so under bounded replay its decoder never computed + positions [0, 128); A (exactly D's first 128 tokens) is repointed onto D's pages at commit. A's + decoder KV lives in A's own ring, so its next decode is unchanged by the repoint (D's shared rows + came from a longer prefill, so to rounding).""" + from .harness import TinyEngine + from .test_reference_parity import _compare + + donor_toks = _tokens(256, 21) + a_toks = donor_toks[:128] + + def run(repoint: bool): + eng = TinyEngine(checkpoint, max_seq_len=1024, max_running_req=2, swa_decoder_replay="bounded") + d = eng.new_request(0, donor_toks) + eng.prefill([d]) + eng.finish_prefill([d]) + a = eng.new_request(1, a_toks) + eng.prefill([a]) + eng.finish_prefill([a]) + if repoint: # what CacheManager.cache_req does when the tree already holds the prefix + eng.page_table[1, :128] = eng.page_table[0, :128] + return eng.decode([a], [3]) + + _compare("decode after a commit repoint", run(True), run(False)) diff --git a/tests/models/deepseek_v41/test_reference_parity.py b/tests/models/deepseek_v41/test_reference_parity.py new file mode 100644 index 0000000000..839cdafb82 --- /dev/null +++ b/tests/models/deepseek_v41/test_reference_parity.py @@ -0,0 +1,219 @@ +"""Logits parity against the vendored reference implementation on the tiny synthetic checkpoint. + +The reference runs bf16 or FP8/FP4 weights with its own torch quantizers, sparse attention, engram +hash and single-pass mHC; FreeToken runs the same checkpoint through its kernels and packed pools. +Compared: exact-mode prefill, greedy decode steps, a two-request batch, and the bounded-replay path +against the reference's ``forward_bounded`` oracle. +""" + +from __future__ import annotations + +from dataclasses import asdict, fields + +import pytest +import torch + +from .common import VOCAB, requires_cuda, tiny_text_config, write_tiny_checkpoint + +pytestmark = requires_cuda + +# The sparse-attention implementations accumulate in different orders. +ATOL = RTOL = 1e-2 +# The fp4 experts also round the routed sum to bf16 before the shared-expert add and weight each +# expert's down output; the reference keeps that sum fp32 and weights the intermediate before its fp8 +# round-trip (10 seeds x 16 checks: max |err| 0.019, against 0.007 in the reference's order). +QUANT_TOL = 2e-2 + + +@pytest.fixture(scope="module") +def checkpoint(tmp_path_factory): + folder = tmp_path_factory.mktemp("dsv41-parity") + tensors = write_tiny_checkpoint(str(folder), seed=11) + return str(folder), tensors + + +def _tokens(n: int, seed: int) -> torch.Tensor: + return torch.randint(3, VOCAB, (n,), generator=torch.Generator().manual_seed(seed)) + + +class Reference: + """The vendored ``Transformer`` on CUDA, loaded from the synthetic checkpoint (identity token map).""" + + def __init__(self, tensors: dict, text: dict, *, max_seq_len: int, max_batch_size: int, quantized: bool = False, vision_config: dict | None = None): + from freetoken.models.deepseek_v41.args import DeepseekV41Args + from types import SimpleNamespace + + from . import reference + from .reference import engram as ref_engram + from .reference import model as ref_model + + args = DeepseekV41Args.from_hf(SimpleNamespace(text_config=text)) + names = {f.name for f in fields(ref_model.ModelArgs)} + kwargs = {k: v for k, v in asdict(args).items() if k in names} + kwargs.update(dtype="fp8" if quantized else "bf16", expert_dtype="fp4" if quantized else None, + max_seq_len=max_seq_len, max_batch_size=max_batch_size, vision_n_layers=0, temperature=0.0) + if vision_config: + kwargs.update(vision_n_layers=vision_config["num_hidden_layers"], vision_dim=vision_config["hidden_size"], + vision_n_heads=vision_config["num_attention_heads"], vision_inter_dim=vision_config["intermediate_size"], + vision_patch_size=vision_config["patch_size"], vision_downsample_ratio=vision_config["downsample_ratio"]) + self.margs = ref_model.ModelArgs(**kwargs) + # the synthetic checkpoint's tokenizer is the identity over VOCAB ids + ref_engram.build_compressed_token_map = lambda tokenizer: (list(range(VOCAB)), VOCAB) + torch.set_default_dtype(torch.bfloat16) + try: + with torch.device("cuda"): + self.model = ref_model.Transformer(self.margs, tokenizer=object()) + finally: + torch.set_default_dtype(torch.float32) + state = {k: v for k, v in tensors.items() if not k.startswith("mtp.") + and (vision_config or (not k.startswith("vision.") and not k.endswith("bias_vl")))} + if quantized: + state = {k: v.view(torch.float4_e2m1fn_x2) if ".ffn.experts." in k and k.endswith(".weight") else v + for k, v in state.items()} + # The reference's grouped output projection holds dequantized bf16 weights. + for name in list(state): + if name.endswith(".wo_a.scale"): + scale = state.pop(name).float().repeat_interleave(32, 0).repeat_interleave(32, 1) + weight_name = name.removesuffix(".scale") + ".weight" + state[weight_name] = (state[weight_name].float() * scale).bfloat16() + missing, unexpected = self.model.load_state_dict({k: v.cuda() for k, v in state.items()}, strict=False) + assert not unexpected, unexpected + assert not [m for m in missing if "freqs_cis" not in m], missing + self.reference = reference + + # the reference builds its index tensors on the default device (generate.py sets it to cuda) + def prefill(self, ids: torch.Tensor) -> torch.Tensor: + with torch.device("cuda"): + return self.model(ids.view(1, -1).cuda(), 0)[1].float() + + def prefill_batch(self, ids: torch.Tensor) -> torch.Tensor: + with torch.device("cuda"): + return self.model(ids.cuda(), 0)[1].float() + + def prefill_bounded(self, ids: torch.Tensor) -> torch.Tensor: + with torch.device("cuda"): + return self.model.forward_bounded(ids.reshape(-1, ids.shape[-1]).cuda(), self.margs.window_size).float() + + def decode(self, tokens: torch.Tensor, pos: int) -> torch.Tensor: + with torch.device("cuda"): + return self.model(tokens.view(-1, 1).cuda(), pos)[1].float() + + +def _engram_table_for(engine, tensors, text): + """Our engine reads the Engram rows straight from the synthetic shard (the disk table), hashed + with the same identity token map the reference oracle uses.""" + from freetoken.models.deepseek_v41.engram import EngramHash + from freetoken.models.deepseek_v41.engram_table import EngramDiskTable, EngramHost, engram_row_source + + hash = EngramHash(engine.args, list(range(VOCAB)), VOCAB) + tables = [ + EngramDiskTable(engram_row_source(engine.checkpoint, layer.layer_id), hash.n_cols, engine.device, max_graph_rows=8, max_extend_tokens=64) + for layer in engine.model.engram_layers() + ] + for layer, table in zip(engine.model.engram_layers(), tables): + layer.attach_table(table) + host = EngramHost(hash, tables, engine.device) + engine.model.forward_host_ctx = host.forward_host_ctx + engine.engram_host = host + + +def _compare(name, got, want, *, tol=ATOL): + err = (got - want).abs().max().item() + assert torch.equal(got.argmax(-1), want.argmax(-1)), f"{name}: argmax differs (max abs err {err:.4f})" + torch.testing.assert_close(got, want, atol=tol, rtol=tol, msg=lambda m: f"{name}: {m}") + + +@pytest.mark.parametrize("mode", ["exact", "bounded"]) +@pytest.mark.parametrize("batch_size", [1, 2]) +def test_quantized_offload_prefill_and_decode_match_reference(tmp_path, mode, batch_size): + from .harness import TinyEngine + + tensors = write_tiny_checkpoint(str(tmp_path), seed=11, quantized=True) + text = tiny_text_config(moe_intermediate_size=256) + ref = Reference(tensors, text, max_seq_len=1024, max_batch_size=3, quantized=True) + eng = TinyEngine(str(tmp_path), swa_decoder_replay=mode, quantized=True) + _engram_table_for(eng, tensors, text) + ids = torch.stack([_tokens(300, 1 + i) for i in range(batch_size)]) + want = ref.prefill_batch(ids) if mode == "exact" else ref.prefill_bounded(ids) + reqs = [eng.new_request(i, row.tolist()) for i, row in enumerate(ids)] + got = eng.prefill(reqs) + _compare(f"quantized {mode} prefill", got, want, tol=QUANT_TOL) + eng.finish_prefill(reqs) + for step in range(3): + token = want.argmax(-1) + want = ref.decode(token, ids.shape[-1] + step) + got = eng.decode(reqs, token.tolist()) + _compare(f"quantized {mode} decode {step}", got, want, tol=QUANT_TOL) + + +def test_exact_prefill_and_greedy_decode_match_the_reference(checkpoint): + from .harness import TinyEngine + + folder, tensors = checkpoint + text = tiny_text_config() + ref = Reference(tensors, text, max_seq_len=1024, max_batch_size=3) + eng = TinyEngine(folder, max_seq_len=1024, max_running_req=2, swa_decoder_replay="exact") + _engram_table_for(eng, tensors, text) + + ids = _tokens(300, 1) + want = ref.prefill(ids) + req = eng.new_request(0, ids.tolist()) + got = eng.prefill([req]) + _compare("prefill", got, want) + eng.finish_prefill([req]) + # greedy decode from the reference's own picks keeps both sides on one trajectory + pos = 300 + for step in range(12): + nxt = int(want.argmax(-1).item()) + want = ref.decode(torch.tensor([nxt]), pos) + got = eng.decode([req], [nxt]) + _compare(f"decode step {step}", got, want) + pos += 1 + + +def test_batched_prefill_and_decode_match_the_reference(checkpoint): + from .harness import TinyEngine + + folder, tensors = checkpoint + text = tiny_text_config() + ref = Reference(tensors, text, max_seq_len=1024, max_batch_size=3) + eng = TinyEngine(folder, max_seq_len=1024, max_running_req=2, swa_decoder_replay="exact") + _engram_table_for(eng, tensors, text) + + ids = torch.stack([_tokens(200, 2), _tokens(200, 3)]) + want = ref.prefill_batch(ids) + reqs = [eng.new_request(i, ids[i].tolist()) for i in range(2)] + got = eng.prefill(reqs) + _compare("batched prefill", got, want) + eng.finish_prefill(reqs) + pos = 200 + for step in range(6): + nxt = want.argmax(-1).cpu() + want = ref.decode(nxt, pos) + got = eng.decode(reqs, nxt.tolist()) + _compare(f"batched decode step {step}", got, want) + pos += 1 + + +def test_bounded_replay_matches_the_reference_oracle(checkpoint): + from .harness import TinyEngine + + folder, tensors = checkpoint + text = tiny_text_config() + ref = Reference(tensors, text, max_seq_len=1024, max_batch_size=3) + eng = TinyEngine(folder, max_seq_len=1024, max_running_req=2, swa_decoder_replay="bounded") + _engram_table_for(eng, tensors, text) + + ids = _tokens(300, 4) + want = ref.prefill_bounded(ids) + req = eng.new_request(0, ids.tolist()) + got = eng.prefill([req]) + _compare("bounded prefill", got, want) + eng.finish_prefill([req]) + pos = 300 + for step in range(6): + nxt = int(want.argmax(-1).item()) + want = ref.decode(torch.tensor([nxt]), pos) + got = eng.decode([req], [nxt]) + _compare(f"decode after bounded prefill, step {step}", got, want) + pos += 1 diff --git a/tests/models/deepseek_v41/test_vision.py b/tests/models/deepseek_v41/test_vision.py new file mode 100644 index 0000000000..0e1d4a4830 --- /dev/null +++ b/tests/models/deepseek_v41/test_vision.py @@ -0,0 +1,214 @@ +"""Image layout, official ViT/logits parity and image replay through the shared encoder cache.""" + +from __future__ import annotations + +import io +import json +from dataclasses import asdict +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from PIL import Image +from safetensors.torch import save_file + +from freetoken.core import Batch +from freetoken.engine.engine import Engine +from freetoken.mm.config import MultimodalConfig +from freetoken.mm.encoder_cache import EncoderCache +from freetoken.mm.processors.deepseek_v41 import DeepseekV41MMProcessor +from freetoken.models.deepseek_v41.config import VisionConfig +from freetoken.scheduler.mm import mm_rows_after, plan_mm_batch +from freetoken.utils.torch_utils import torch_dtype + +from .common import DIM, VOCAB, requires_cuda, tiny_text_config, write_tiny_checkpoint +from .reference.vision import Aligner, ViT + +VC = asdict(VisionConfig(num_hidden_layers=2, hidden_size=32, num_attention_heads=4, intermediate_size=48, + patch_size=2, min_pixels=0, max_image_tokens=64)) + + +def _args(): + return SimpleNamespace(vision_n_layers=2, vision_dim=32, vision_n_heads=4, vision_inter_dim=48, + vision_patch_size=2, vision_downsample_ratio=3, vision_rope_theta=10000.0, dim=DIM) + + +def _checkpoint(path, quantized=False): + tensors = write_tiny_checkpoint(str(path), seed=11, quantized=quantized) + torch.manual_seed(12) + with torch_dtype(torch.bfloat16): + vit, aligner = ViT(_args()), Aligner(_args()) + # Replace the unused dummy vision weight from the text-only fixture. + tensors = {k: v for k, v in tensors.items() if not k.startswith("vision.")} + for prefix, module in (("vision", vit), ("aligner", aligner)): + tensors.update({f"{prefix}.{k}": v.detach().bfloat16().contiguous() for k, v in module.state_dict().items()}) + tensors.update({k: torch.randn(DIM, dtype=torch.bfloat16) for k in ("image_start", "image_end", "image_newline")}) + save_file(tensors, path / 'model-00001-of-00001.safetensors') + (path / 'model.safetensors.index.json').write_text(json.dumps({'weight_map': {k: 'model-00001-of-00001.safetensors' for k in tensors}})) + config = json.loads((path / 'config.json').read_text()) + config.update(vision_config=VC, image_token_id=VOCAB - 1) + (path / 'config.json').write_text(json.dumps(config)) + return tensors + + +def _processor(): + cfg = SimpleNamespace(vision_config=SimpleNamespace(**VC), image_token_id=VOCAB - 1) + return DeepseekV41MMProcessor(cfg, '.', MultimodalConfig()) + + +def _image(color): + buf = io.BytesIO() + Image.new('RGB', (14, 10), color).save(buf, format='PNG') + return buf.getvalue() + + +def _inputs(length=151): + ids = [3] * 17 + [VOCAB - 1] + [4] * (length - 19) + [VOCAB - 1] + return _processor().apply(torch.tensor(ids, dtype=torch.int32), [_image('red'), _image('blue')]) + + +def _prefill(eng, reqs): + if not hasattr(eng, 'encoder_cache'): + eng.encoder_cache = EncoderCache() + eng.dtype = torch.bfloat16 + eng._mm_registered = set() + for req in reqs: + if id(req) not in eng._mm_registered: + for item in req.mm_items or (): # what the scheduler claims at admission + eng.encoder_cache.register(item.hash, req.uid, mm_rows_after(item, req.cached_len)) + eng._mm_registered.add(id(req)) + batch = Batch(reqs=reqs, phase='prefill') + batch.padded_reqs = reqs + batch.input_ids = torch.cat([r.input_ids[r.cached_len:r.device_len] for r in reqs]).cuda() + batch.positions = torch.cat([torch.arange(r.cached_len, r.device_len) for r in reqs]).cuda() + jobs, plan, rows, _ = plan_mm_batch(reqs, eng.encoder_cache) + batch.mm_encoder_jobs, batch.mm_gather_plan = jobs, plan + if plan: + batch.mm_rows = torch.tensor(rows, dtype=torch.long, device=eng.device) + Engine._run_mm_encoder(eng, batch) + eng._bind() + eng.backend.prepare_metadata(batch) + with eng.ctx.forward_batch(batch), eng.model.forward_host_ctx(batch, False): + return eng.model.forward().float() + + +def test_processor_patch_order_and_all_span_rows_are_embeddings(): + pixels = np.arange(10 * 14 * 3, dtype=np.uint8).reshape(10, 14, 3) + item = _processor().process([Image.fromarray(pixels)])[0] + assert item.grid_thw == [1, 5, 7] + expected = ((torch.tensor(pixels[:2, :2]).permute(2, 0, 1).float() / 255 - .5) / .5).bfloat16() + assert torch.equal(item.feature[0], expected) + assert len(_processor().prompt_replacement(item).full) == 10 + result = _inputs() + assert [i.offsets for i in result.mm_items] == [[[17, 27]], [[159, 169]]] + assert result.mrope_positions is None + assert result.mm_items[0].hash != result.mm_items[1].hash + + +def test_text_only_skips_every_visual_parameter(tmp_path): + from freetoken.models.deepseek_v41.weight import iter_weights + from freetoken.mm.processor import get_mm_processor + from .test_engine_config import _engine_config + _checkpoint(tmp_path) + enabled = _engine_config(str(tmp_path), moe_strategy='offload') + assert enabled.model_config.is_multimodal + disabled = MultimodalConfig(disabled_encoders=frozenset({'vision', 'audio'})) + text = _engine_config(str(tmp_path), moe_strategy='offload', mm=disabled) + assert not text.model_config.is_multimodal and not text.active_encoders + assert get_mm_processor(str(tmp_path), disabled) is None + loaded = dict(iter_weights(str(tmp_path), 'cpu', include_moe_experts=False, include_vision=False)) + assert not any(k.startswith('visual.') for k in loaded) + assert get_mm_processor(str(tmp_path)).placeholder == [VOCAB - 1] + + +@requires_cuda +@pytest.mark.parametrize('placement', ['gpu', 'host']) +def test_vision_encoder_matches_official_with_padded_grid(tmp_path, placement): + from .harness import TinyEngine + tensors = _checkpoint(tmp_path) + eng = TinyEngine(str(tmp_path)) + eng.model.place_encoder_weights(placement) + item = _processor().process([Image.new('RGB', (14, 10), 'red')])[0] + with torch_dtype(torch.bfloat16), torch.device('cuda'): + vit, aligner = ViT(_args()), Aligner(_args()) + for prefix, module in (('vision', vit), ('aligner', aligner)): + module.load_state_dict({k.removeprefix(prefix + '.'): v.cuda() for k, v in tensors.items() if k.startswith(prefix + '.')}) + with torch.inference_mode(), torch.device('cuda'): + x = aligner(vit(item.feature.cuda(), 5, 7), 5, 7) + got = eng.model.encode(item) + torch.testing.assert_close(got[1:-1].view(2, 4, DIM)[:, :3].reshape(-1, DIM), x, atol=1e-3, rtol=1e-3) + assert torch.equal(got[0], tensors['image_start'].cuda()) + assert torch.equal(got[-1], tensors['image_end'].cuda()) + assert torch.equal(got[[4, 8]], tensors['image_newline'].cuda().expand(2, -1)) + eng.model.place_encoder_weights('gpu') + + +@requires_cuda +@pytest.mark.parametrize('quantized', [False, True]) +def test_image_logits_and_decode_match_reference(tmp_path, quantized): + from .harness import TinyEngine + from .test_reference_parity import ATOL, QUANT_TOL, Reference, _compare, _engram_table_for + tol = QUANT_TOL if quantized else ATOL + tensors = _checkpoint(tmp_path, quantized) + text = tiny_text_config(moe_intermediate_size=256) if quantized else tiny_text_config() + ref = Reference(tensors, text, max_seq_len=512, max_batch_size=2, quantized=quantized, vision_config=VC) + eng = TinyEngine(str(tmp_path), max_seq_len=512, swa_decoder_replay='exact', quantized=quantized) + _engram_table_for(eng, tensors, text) + result = _inputs() + ids = result.input_ids.clone() + types = torch.full_like(ids, -1) + images = [] + for item in result.mm_items: + lo, hi = item.offsets[0] + ids[lo:hi] = VOCAB - 1 + typ = torch.tensor([0] + ([1] * 3 + [2]) * 2 + [3]) + types[lo:hi] = typ + images.append(SimpleNamespace(start=lo, patches=item.feature, n_vit_h=5, n_vit_w=7, types=typ)) + with torch.device('cuda'): + want = ref.model(ids.cuda()[None], images=[images], token_types=types.cuda()[None])[1] + req = eng.new_request(0, result.input_ids.tolist()) + req.mm_items = result.mm_items + got = _prefill(eng, [req]) + _compare('image prefill', got, want, tol=tol) + eng.finish_prefill([req]) + for step in range(3): + token = want.argmax(-1) + want = ref.decode(token, len(ids) + step) + got = eng.decode([req], token.tolist()) + _compare(f'image decode {step}', got, want, tol=tol) + + +@requires_cuda +@pytest.mark.parametrize('mode', ['exact', 'bounded']) +def test_images_survive_chunking_and_prefix_replay(tmp_path, mode): + from .harness import TinyEngine + from .test_reference_parity import _compare, _engram_table_for + tensors = _checkpoint(tmp_path) + eng = TinyEngine(str(tmp_path), max_seq_len=2048, swa_decoder_replay=mode) + _engram_table_for(eng, tensors, tiny_text_config()) + # bounded mode caps a hit a window before the prompt end: trailing text keeps that cap at 128 + tail = 150 if mode == 'bounded' else 0 + def inputs(): + ids = [3] * 80 + [VOCAB - 1] + [4] * 37 + [VOCAB - 1] + [4] * 31 + [VOCAB - 1] + [4] * tail + return _processor().apply(torch.tensor(ids, dtype=torch.int32), [_image(c) for c in ('red', 'blue', 'green')]) + whole = inputs() + assert [i.offsets for i in whole.mm_items] == [[[80, 90]], [[127, 137]], [[168, 178]]] + a = eng.new_request(0, whole.input_ids.tolist()) + a.mm_items = whole.mm_items + want = _prefill(eng, [a]) + chunk = inputs() + b = eng.new_request(1, chunk.input_ids.tolist()) + b.mm_items = chunk.mm_items + b.device_len = 128 + _prefill(eng, [b]) + b.cached_len, b.device_len = 128, len(chunk.input_ids) + got = _prefill(eng, [b]) + _compare('chunked image prefill', got, want) + prefix = inputs() + # The second image straddles the chunk / prefix boundary at 128 (in bounded mode the first chunk + # also ends before the prompt's last window and runs no decoder row). + hit = eng.new_request_on_prefix(1, a, prefix.input_ids.tolist()) + assert hit.cached_len == 128 + hit.mm_items = prefix.mm_items + _compare('image prefix hit', _prefill(eng, [hit]), want) diff --git a/tests/models/test_quant_config.py b/tests/models/test_quant_config.py index 084a2e4aa7..bfcfb87664 100644 --- a/tests/models/test_quant_config.py +++ b/tests/models/test_quant_config.py @@ -1,6 +1,7 @@ """What each checkpoint's quant config resolves to: the dialect, the method behind every layer, and agreement with the stored tensors. -Rows need the checkpoint's config.json locally (/mnt/nvme/models or the HF cache); absent ones skip. +Rows need the checkpoint's config.json locally (/mnt/nvme/models or the HF cache); absent ones skip. The +stored-tensor scan also reads $FREETOKEN_TEST_MODELS_ROOT. """ from __future__ import annotations @@ -47,6 +48,8 @@ from freetoken.utils.torch_utils import torch_dtype MODELS = "/mnt/nvme/models" +# the stored-tensor scan takes every checkpoint under this root (the case table above matches by name) +SCAN_ROOT = os.environ.get("FREETOKEN_TEST_MODELS_ROOT", MODELS) HF_CACHE = os.path.expanduser("~/.cache/huggingface/hub") BF16, FP8B, FP8T, MXFP8, NVFP4 = UnquantizedLinearMethod, Fp8BlockLinearMethod, Fp8TensorLinearMethod, Mxfp8LinearMethod, Nvfp4LinearMethod @@ -435,6 +438,25 @@ def test_modelopt_nvfp4_reads_the_activation_quantizer_from_config_groups(extra, assert scheme.kind is QuantKind.NVFP4 and scheme.has("input_scale") is has_input_scale +def test_fp8_dialect_reads_the_block_size_and_the_nested_expert_dtype(): + """DeepSeek-V4.1 keeps the ``fp8`` dialect with 32x32 blocks and puts ``expert_dtype`` inside quantization_config.""" + from freetoken.layers.quantization.scheme import fp8_block_size + + q = {"quant_method": "fp8", "activation_scheme": "dynamic", "weight_block_size": [32, 32], "scale_fmt": "ue8m0", "expert_dtype": "fp4"} + quant = QuantConfig.from_hf(SimpleNamespace(quantization_config=q)) + assert type(quant) is Fp8BlockConfig and quant.block == 32 + dense = quant.scheme_for("layers.3.attn.wq_a") + assert dense.kind is QuantKind.FP8_BLOCK and fp8_block_size(dense) == 32 and dense.weight.scale == "e8m0" + assert quant.storage(dense)["weight_scale_inv"].name == "scale" + assert quant.scheme_for("layers.3.ffn.experts.7.w1").kind is QuantKind.MXFP4 + # the V4 128-block export still resolves to the 128 scheme + q128 = {"quant_method": "fp8", "activation_scheme": "dynamic", "weight_block_size": [128, 128], "scale_fmt": "ue8m0"} + assert fp8_block_size(QuantConfig.from_hf(SimpleNamespace(quantization_config=q128)).scheme_for("layers.3.attn.wq_a")) == 128 + for bad in ([64, 64], [32, 128], [128]): + with pytest.raises(NotImplementedError): + QuantConfig.from_hf(SimpleNamespace(quantization_config={**q, "weight_block_size": bad})) + + def test_every_dialect_names_the_tensors_behind_its_schemes(): from freetoken.layers.quantization.registry import dialects @@ -491,7 +513,7 @@ def _weight_map(p: Path) -> dict[str, str] | None: def _candidate_dirs() -> list[Path]: - roots = glob.glob(os.path.join(HF_CACHE, "models--*/snapshots/*/")) + glob.glob(os.path.join(MODELS, "*/")) + roots = glob.glob(os.path.join(HF_CACHE, "models--*/snapshots/*/")) + glob.glob(os.path.join(SCAN_ROOT, "*/")) return [Path(d) for d in roots if (Path(d) / "config.json").exists() and _weight_map(Path(d)) is not None] diff --git a/tests/moe/test_cpu_moe.py b/tests/moe/test_cpu_moe.py index f496de3142..547c94c60a 100644 --- a/tests/moe/test_cpu_moe.py +++ b/tests/moe/test_cpu_moe.py @@ -391,6 +391,7 @@ def test_cpu_decode_dsfp4_matches_gpu(bs): max_tokens=bs, device=dev, swiglu_limit=swiglu_limit, + act_block=128, ) hidden = torch.randn(bs, H, device=dev, dtype=torch.bfloat16) @@ -585,7 +586,7 @@ def test_cpu_moe_decode_cuda_graph_replay_dsfp4(): ex = CpuMoeExecutor( cache, top_k=top_k, activation="silu", apply_router_weight_on_input=False, num_threads=8, max_tokens=bs, device=dev, - swiglu_limit=limit, + swiglu_limit=limit, act_block=128, ) b = cache.bank_sources layer_banks = [b[n][layer].to(dev) for n in ( diff --git a/tests/moe/test_dsfp4_prequant.py b/tests/moe/test_dsfp4_prequant.py index 6d6b49d395..11519502e2 100644 --- a/tests/moe/test_dsfp4_prequant.py +++ b/tests/moe/test_dsfp4_prequant.py @@ -46,7 +46,7 @@ def test_dsfp4_gpu_prequant_bit_parity(): def run(cpu_roundtrip: bool) -> torch.Tensor: ex = CpuMoeExecutor( cache, top_k=TOPK, activation="silu", apply_router_weight_on_input=False, - num_threads=4, max_tokens=2, device=dev, swiglu_limit=10.0, + num_threads=4, max_tokens=2, device=dev, swiglu_limit=10.0, act_block=128, ) assert ex._gpu_prequant, "ds_fp4 on CUDA should enable GPU prequant" if cpu_roundtrip: # control: the original CPU-side scalar round-trip diff --git a/tests/moe/test_dsfp4_reference_semantics.py b/tests/moe/test_dsfp4_reference_semantics.py new file mode 100644 index 0000000000..2f89cf334d --- /dev/null +++ b/tests/moe/test_dsfp4_reference_semantics.py @@ -0,0 +1,170 @@ +"""ds_fp4 (DeepSeek W4A8) experts against an INDEPENDENT torch transcription over +``fp4_gemm(act_quant(x, block), W)`` at both fp8 activation blocks (128 on V4, 32 on V4.1): + + x_q = fp8_roundtrip(x, block) # act_quant(..., inplace) + gate = bf16(x_q @ deq(W1)^T), up = bf16(x_q @ deq(W3)^T) # fp4_gemm outputs bf16 + h = bf16(silu(min(gate, L)) * clamp(up, -L, L)) + y_r = bf16(w_r * (fp8_roundtrip(h, block) @ deq(W2)^T)) # routing weight in the down epilogue + y = sum_r y_r + +The routing weight scales the down output (the placement main serves V4 with); the reference +``Expert.forward`` scales the intermediate before its fp8 quant instead. The GEMV (decode), the +grouped GEMM (prefill) and the CPU executor all follow this transcription. +""" + +from __future__ import annotations + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + +E, H, I, TOP_K, LIMIT = 16, 512, 256, 4, 7.0 +E2M1 = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0]) + + +def _banks(device, seed=0): + g = torch.Generator(device="cpu").manual_seed(seed) + + def u8(*shape, low=0, high=256): + return torch.randint(low, high, shape, dtype=torch.uint8, generator=g).to(device) + + return (u8(E, 2 * I, H // 2), u8(E, 2 * I, H // 32, low=120, high=130), u8(E, H, I // 2), u8(E, H, I // 32, low=120, high=130)) + + +def _dequant(packed: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + """[E, N, K//2] e2m1 pairs (even channel low nibble) + [E, N, K//32] e8m0 -> [E, N, K] fp32.""" + lut = E2M1.to(packed.device) + lo, hi = lut[(packed & 0xF).long()], lut[(packed >> 4).long()] + vals = torch.stack([lo, hi], dim=-1).flatten(-2) # [E, N, K] + s = torch.exp2(scale.view(torch.uint8).float() - 127.0).repeat_interleave(32, dim=-1) + return vals * s + + +def _fp8_roundtrip(x: torch.Tensor, block: int) -> torch.Tensor: + g = x.float().unflatten(-1, (-1, block)) + s = torch.exp2(torch.ceil(torch.log2(g.abs().amax(-1).clamp_min(1e-4) / 448.0))) + q = (g / s.unsqueeze(-1)).clamp(-448.0, 448.0).to(torch.float8_e4m3fn).float() + return (q * s.unsqueeze(-1)).flatten(-2).to(torch.bfloat16) + + +def reference(x, slots, weights, banks, block): + """The MoE over dequantized banks: per-route bf16 expert outputs accumulated into an fp32 ``y``, + returned in fp32 (the kernels round that sum to bf16).""" + gu_p, gu_s, dn_p, dn_s = banks + W13 = _dequant(gu_p, gu_s) # [E, 2I, H] + W2 = _dequant(dn_p, dn_s) # [E, H, I] + xq = _fp8_roundtrip(x, block).float() + T = x.shape[0] + y = torch.zeros(T, H, dtype=torch.float32, device=x.device) + for t in range(T): + for r in range(TOP_K): + e = int(slots[t, r]) + gu = (xq[t] @ W13[e].T).to(torch.bfloat16).float() + gate, up = gu[:I].clamp(max=LIMIT), gu[I:].clamp(-LIMIT, LIMIT) + h = (torch.nn.functional.silu(gate) * up).to(torch.bfloat16) + hq = _fp8_roundtrip(h.view(1, -1), block).float().view(-1) + y[t] += (weights[t, r].float() * (hq @ W2[e].T)).to(torch.bfloat16).float() + return y + + +def _inputs(T, device, seed=1): + g = torch.Generator(device="cpu").manual_seed(seed) + x = (torch.randn(T, H, dtype=torch.bfloat16, generator=g) * 0.5).to(device) + slots = torch.stack([torch.randperm(E, generator=g)[:TOP_K] for _ in range(T)]).to(device=device, dtype=torch.int32).contiguous() + w = torch.rand(T, TOP_K, generator=g) + w = (1.5 * w / w.sum(-1, keepdim=True)).float().to(device).contiguous() # route_scale-style weights > 1 too + return x, slots, w + + +@pytest.mark.parametrize("block", [32, 128]) +def test_gemv_path_matches_the_transcription(block): + from freetoken.moe.fused_ds_fp4 import routed_experts_fp4 + + banks = _banks("cuda") + x, slots, w = _inputs(6, "cuda") + got = routed_experts_fp4(x, slots, w, *banks, LIMIT, act_block=block).float() + want = reference(x, slots, w, banks, block) + torch.testing.assert_close(got, want, atol=2e-2, rtol=2e-2) + + +@pytest.mark.parametrize("block", [32, 128]) +def test_cpu_executor_matches_the_transcription(block): + """The ``ds_fp4`` CPU executor (cpu / hybrid decode) implements the same contract as the GPU path.""" + from types import SimpleNamespace + + from freetoken.kernel.pinned import alloc_pinned_tensor + from freetoken.moe.cpu_executor import CpuMoeExecutor + + gu_p, gu_s, dn_p, dn_s = _banks("cuda") + pinned = {} + for name, t in (("gate_up_packed", gu_p), ("gate_up_scale", gu_s), ("down_packed", dn_p), ("down_scale", dn_s)): + buf = alloc_pinned_tensor(*t.shape, dtype=torch.uint8) + buf.copy_(t.cpu()) + pinned[name] = buf + cache = SimpleNamespace(quant_format="ds_fp4", bank_sources={k: [v] for k, v in pinned.items()}, num_layers=1, num_experts=E) + dev = torch.device("cuda", 0) + x, slots, w = _inputs(2, dev) + ex = CpuMoeExecutor(cache, top_k=TOP_K, activation="silu", apply_router_weight_on_input=False, num_threads=4, max_tokens=2, device=dev, swiglu_limit=LIMIT, act_block=block) + got = ex.decode(0, x, w, slots.clone()).clone().float() + torch.cuda.synchronize() + del ex + want = reference(x, slots, w, (gu_p, gu_s, dn_p, dn_s), block) + torch.testing.assert_close(got, want, atol=2e-2, rtol=2e-2) + + +def test_grouped_prefill_path_matches_the_transcription(): + from freetoken.moe import fused_ds_fp4 + + banks = _banks("cuda") + x, slots, w = _inputs(64, "cuda") + got = fused_ds_fp4.routed_experts_fp4_prefill(x, slots, w, *banks, LIMIT, E, act_block=32).float() + want = reference(x, slots, w, banks, 32) + torch.testing.assert_close(got, want, atol=2e-2, rtol=2e-2) + + +@pytest.mark.parametrize("block", [32, 128]) +def test_inactive_routes_contribute_zero_without_reading_their_slots(block): + """Hybrid decode hands the GPU kernel slot -1 (weight 0) for the routes the CPU computes. The kernel + must produce exactly zero for them without touching any slot: every slot the active routes do not + use is filled with 0xFF (an e8m0 scale of NaN), and the result still equals the reference over the + active routes and the CPU executor's output for the same split.""" + from types import SimpleNamespace + + from freetoken.kernel.pinned import alloc_pinned_tensor + from freetoken.moe.cpu_executor import CpuMoeExecutor + from freetoken.moe.fused_ds_fp4 import routed_experts_fp4 + + gu_p, gu_s, dn_p, dn_s = (t.clone() for t in _banks("cuda")) + x, slots, w = _inputs(2, "cuda") + # routes: token 0 keeps routes 0..1 on the GPU, token 1 keeps route 3; the rest go to the CPU + on_gpu = torch.zeros_like(slots, dtype=torch.bool) + on_gpu[0, :2] = True + on_gpu[1, 3] = True + gpu_slots = torch.where(on_gpu, slots, -1) + gpu_w = torch.where(on_gpu, w, 0.0) + used = set(slots[on_gpu].tolist()) + poison = [e for e in range(E) if e not in used] + for t in (gu_p, gu_s, dn_p, dn_s): + t[poison] = 0xFF + got = routed_experts_fp4(x, gpu_slots, gpu_w, gu_p, gu_s, dn_p, dn_s, LIMIT, act_block=block) + assert torch.isfinite(got).all() + clean = _banks("cuda") + # the poisoned-bank result equals the clean-bank result with the same holes (neither reads them), + # and each share matches the reference over its own routes (a zero weight drops a route) + got_clean = routed_experts_fp4(x, gpu_slots, gpu_w, *clean, LIMIT, act_block=block) + assert torch.equal(got, got_clean) + torch.testing.assert_close(got.float(), reference(x, slots, gpu_w, clean, block), atol=2e-2, rtol=2e-2) + # the CPU executor computes the complementary split from the same raw ids (ids < 0 skipped) + pinned = {} + for name, t in (("gate_up_packed", clean[0]), ("gate_up_scale", clean[1]), ("down_packed", clean[2]), ("down_scale", clean[3])): + buf = alloc_pinned_tensor(*t.shape, dtype=torch.uint8) + buf.copy_(t.cpu()) + pinned[name] = buf + cache = SimpleNamespace(quant_format="ds_fp4", bank_sources={k: [v] for k, v in pinned.items()}, num_layers=1, num_experts=E) + ex = CpuMoeExecutor(cache, top_k=TOP_K, activation="silu", apply_router_weight_on_input=False, num_threads=4, max_tokens=2, device=torch.device("cuda", 0), swiglu_limit=LIMIT, act_block=block) + cpu_ids = torch.where(on_gpu, -1, slots) + cpu_out = ex.decode(0, x, w, cpu_ids.clone()).clone() + torch.cuda.synchronize() + del ex + torch.testing.assert_close(cpu_out.float(), reference(x, slots, torch.where(on_gpu, 0.0, w), clean, block), atol=2e-2, rtol=2e-2) diff --git a/tests/scheduler/test_dsv4_generic_manager.py b/tests/scheduler/test_dsv4_generic_manager.py index ad69b89ec5..1bca10bc8a 100644 --- a/tests/scheduler/test_dsv4_generic_manager.py +++ b/tests/scheduler/test_dsv4_generic_manager.py @@ -14,8 +14,8 @@ import torch from freetoken.core import Req, SamplingParams -from freetoken.kvcache.dsv4_cost_model import dsv4_pool_sizes -from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache +from freetoken.kvcache.dsv4.v4_cost_model import dsv4_pool_sizes +from freetoken.kvcache.dsv4.v4_pool import DSV4PagedKVCache from freetoken.models.deepseek_v4.args import DeepseekV4Args from freetoken.scheduler.cache import CacheManager @@ -191,6 +191,33 @@ def _pending(uid, n): cm.check_integrity() +@pytest.mark.parametrize("replay, new_tokens, cached", [(0, 2, 256), (P, 2, 128), (P, 200, 256)]) +def test_admission_caps_the_hit_by_the_pools_prefix_replay(replay, new_tokens, cached): + """A pool whose model recomputes the prompt's tail (DeepSeek-V4.1's bounded decoder replay) holds + back ``prefix_replay_tokens`` at admission: the match ends at or before ``L - replay`` (page-aligned + by the radix), a hit already that far from the prompt end is untouched, and the request's + ``prompt_len`` is the whole prompt.""" + from freetoken.scheduler.decode import DecodeManager + from freetoken.scheduler.prefill import PrefillManager + from freetoken.scheduler.table import TableManager + from freetoken.scheduler.utils import PendingReq + + cm, pool, pt = _stack(num_pages=48) + pool.prefix_replay_tokens = replay + prompt = torch.arange(1, 301, dtype=torch.int32) + _lifecycle(cm, _req(MRR - 1, prompt, n_decode=310), total_len=310) + pm = PrefillManager(cm, TableManager(max_running_reqs=MRR, page_table=pt), DecodeManager(page_size=P)) + tokens = torch.cat([prompt, torch.arange(1000, 1000 + new_tokens, dtype=torch.int32)]) + pm.pending_list = [PendingReq(uid=1, input_ids=tokens, sampling_params=SamplingParams(max_tokens=1))] + batch = pm.schedule_next_batch(1024) + req = batch.reqs[0] + assert req.cached_len == cached and req.extend_len == len(tokens) - cached and req.prompt_len == len(tokens) + cm.allocate_paged(batch.reqs) + req.complete_one() + cm.cache_req(req, finished=True) + cm.check_integrity() + + def test_abort_anywhere_fuzz_conserves_and_isolates(): """admit / extend / finish-or-abort at arbitrary points with shared prefixes + eviction pressure, on a deliberately small window tier. The generic conservation check (free + tree diff --git a/tests/scheduler/test_scheduler_kv_usage.py b/tests/scheduler/test_scheduler_kv_usage.py index 6ab949705e..bf0e825c31 100644 --- a/tests/scheduler/test_scheduler_kv_usage.py +++ b/tests/scheduler/test_scheduler_kv_usage.py @@ -2,6 +2,8 @@ from types import SimpleNamespace +from freetoken.attention import AttnType + def test_kv_usage_pages_excludes_evictable_prefix_cache(): # page_usage now lives on the CacheManagerLike interface (polymorphic vs DSV4); test the @@ -53,10 +55,15 @@ def _fake_engine(swa=True, moe=True, mamba=True): config=SimpleNamespace( page_size=16, cache_type="swa_radix", - model_config=SimpleNamespace(dsv4_args=None, has_swa_attention=swa), + max_running_req=1, + model_config=SimpleNamespace( + dsv4_args=None, has_swa_attention=swa, + kv_cache_group_specs=lambda: [SimpleNamespace(attn_type=AttnType.SWA, is_swa=True, sliding_window=128)] if swa else [], + swa_attention_group=lambda: SimpleNamespace(sliding_window=128), + ), ), kv_cache=SimpleNamespace( - swa_num_tokens=513, unit_bytes=lambda: (1 << 20, 1 << 21) + window_pages=512, swa_num_tokens=513, unit_bytes=lambda: (1 << 20, 1 << 21) ), moe_offload_cache=SimpleNamespace( cache_size=24, num_layers=8, num_experts=16, diff --git a/tests/server/test_effort_dialect.py b/tests/server/test_effort_dialect.py index 6a68a4ece3..a25de8e728 100644 --- a/tests/server/test_effort_dialect.py +++ b/tests/server/test_effort_dialect.py @@ -159,6 +159,22 @@ def test_empty_effort_is_treated_as_absent(): assert state.sent.chat_template_kwargs == {} +def test_numeric_effort_rides_template_kwargs_not_the_wire_field(): + # The wire field takes named tiers only (as vLLM and SGLang type it); a numeric budget goes + # through chat_template_kwargs, where the checkpoint's encoder validates it at render time. + from pydantic import ValidationError + import pytest + + for effort in (37, True, 25.5): + with pytest.raises(ValidationError): + chat_request(reasoning_effort=effort) + ctk = {"enable_thinking": True, "reasoning_effort": 37} + state = FakeState(reasoning_parser="qwen3") + response = run(handle_chat_completion(chat_request(chat_template_kwargs=ctk), None, state, {})) + assert not isinstance(response, JSONResponse) + assert state.sent.chat_template_kwargs == ctk + + def test_foreign_thinking_shapes_stay_ignored(): # extra="allow" swallowed any thinking shape before the field existed; # a bare string, a bool, or a typeless dict must keep working unchanged. diff --git a/tests/server/test_function_call_parser.py b/tests/server/test_function_call_parser.py index 34c1c04d10..bc22296760 100644 --- a/tests/server/test_function_call_parser.py +++ b/tests/server/test_function_call_parser.py @@ -337,6 +337,34 @@ def test_dsv32_streaming_multi_param_args_prefix_stable(): } +V41_BLOCK = ( + "<｜DSML｜ calls>\n" + '<｜DSML｜ invoke name="glob">\n' + '<｜DSML｜ parameter name="pattern" string="true">*.py\n' + '<｜DSML｜ parameter name="path" string="true">/src\n' + "\n" + "" +) + + +def test_dsv41_dsml_tags_lead_with_a_space(): + """DeepSeek-V4.1 leads every DSML tag name with a space (its encoder's ``<｜DSML｜ calls>``): + the V3.2 detector sees no tool call there, the V4.1 one parses it whole and streamed.""" + text = "Let me look.\n\n" + V41_BLOCK + assert not FunctionCallParser(OPENCODE_TOOLS, tool_call_parser="deepseekv32").has_tool_call(text) + parser = FunctionCallParser(OPENCODE_TOOLS, tool_call_parser="deepseekv41") + assert parser.has_tool_call(text) + result = parser.parse_non_stream(text) + assert result.normal_text == "Let me look." + assert [(c.name, json.loads(c.parameters)) for c in result.calls] == [("glob", {"pattern": "*.py", "path": "/src"})] + + streamed = FunctionCallParser(OPENCODE_TOOLS, tool_call_parser="deepseekv41") + texts, calls = _feed(streamed, [text[i : i + 5] for i in range(0, len(text), 5)]) + assert "".join(texts).strip() + streamed.finish_stream() == "Let me look." + assert [c.name for c in calls if c.name] == ["glob"] + assert json.loads("".join(c.parameters for c in calls if c.name is None)) == {"pattern": "*.py", "path": "/src"} + + def test_streaming_support_flags(): # Every registered detector is incremental-safe (buffered fallback remains as # the escape hatch for future formats, exercised via monkeypatch in diff --git a/tests/server/test_parser_auto_selection.py b/tests/server/test_parser_auto_selection.py index 78b1663e8d..e1eb87982b 100644 --- a/tests/server/test_parser_auto_selection.py +++ b/tests/server/test_parser_auto_selection.py @@ -87,6 +87,13 @@ def test_qwen3_5_is_not_shadowed_by_the_generic_qwen_branch(): assert _inferred("Qwen3MoeForCausalLM")[0] == "qwen25" +def test_deepseek_v41_is_not_shadowed_by_the_v4_branch(): + """V4.1's DSML leads every tag name with a space; the V4 arm (``"v4"`` is a substring of + ``"v41"``) would hand it a tool parser that never sees a call.""" + assert _inferred("DeepseekV41ForCausalLM") == ("deepseekv41", "deepseekv32") + assert _inferred("DeepseekV4ForCausalLM")[0] == "deepseekv32" + + 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): diff --git a/tests/server/test_streaming_model_matrix.py b/tests/server/test_streaming_model_matrix.py index 43922021a2..eb6e3a2ed2 100644 --- a/tests/server/test_streaming_model_matrix.py +++ b/tests/server/test_streaming_model_matrix.py @@ -67,6 +67,11 @@ '<｜DSML｜parameter name="filePath" string="true">/tmp/test_calc.py' "" ), + "deepseekv41": ( + '<｜DSML｜ calls>\n<｜DSML｜ invoke name="read">\n' + '<｜DSML｜ parameter name="filePath" string="true">/tmp/test_calc.py\n' + "\n" + ), "qwen25": '\n{"name": "read", "arguments": {"filePath": "/tmp/test_calc.py"}}\n', "qwen3_coder": ( "/tmp/test_calc.py" @@ -104,6 +109,7 @@ # Substrings that must never leak into user-visible content. MARKUP_MARKERS = { "deepseekv32": ["｜DSML｜"], + "deepseekv41": ["｜DSML｜"], "qwen25": ["", ""], "qwen3_coder": ["", "", "", ""], @@ -121,6 +127,7 @@ # opens the block implicitly (model emits only the closing marker). REASONING_FAMILIES = { "dsv4": ("deepseekv32", "deepseekv32", "", ""), + "dsv4.1": ("deepseekv41", "deepseekv32", "", ""), "qwen3.5": ("qwen3_coder", "qwen3", "", ""), "qwen": ("qwen25", "qwen3", "", ""), "glm4.7": ("glm47", "glm", "", ""), diff --git a/tests/tokenizer/test_effort.py b/tests/tokenizer/test_effort.py index e56df5fe68..e0cf7a81f6 100644 --- a/tests/tokenizer/test_effort.py +++ b/tests/tokenizer/test_effort.py @@ -1,6 +1,8 @@ """Unit tests for the reasoning-effort dialect layer (tokenizer/effort.py).""" from __future__ import annotations +import pytest + from freetoken.tokenizer.effort import ( EFFORT_SCALE, EffortProfile, @@ -131,3 +133,37 @@ def render(kwargs, tools): profile = probe_effort_profile(render) assert profile.supported == frozenset({"xhigh", "medium", "low"}) assert profile.consumes_effort + + +def _numeric_render(kwargs, tools): + if not tools and not kwargs.get("enable_thinking"): + return "chat" + effort = kwargs.get("reasoning_effort", "high") + if isinstance(effort, str): + effort = {"low": 50, "high": 75, "max": 100}[effort] + if type(effort) is not int or not 1 <= effort <= 100: + raise ValueError("invalid budget") + return f"Reasoning Effort: {effort}" + + +def test_probe_preserves_validated_numeric_budgets(): + profile = probe_effort_profile(_numeric_render) + assert profile.integer_range == (1, 100) + assert profile.default == "high" + for budget in range(1, 101): + assert quantize_effort(budget, profile) == budget + for bad in (0, 101, -1, True, False, 1.0, 1.5): + with pytest.raises(ValueError, match="integer"): + quantize_effort(bad, profile) + + +def test_numeric_capability_requires_validation_and_the_whole_range(): + assert probe_effort_profile(_dsv4_render).integer_range is None + assert probe_effort_profile(lambda kw, tools: str(kw)).integer_range is None + + def holes(kwargs, tools): + if kwargs.get("reasoning_effort") == 42: + raise ValueError("unsupported") + return _numeric_render(kwargs, tools) + + assert probe_effort_profile(holes).integer_range is None diff --git a/tests/tokenizer/test_tokenize.py b/tests/tokenizer/test_tokenize.py index 16ef6d7419..fb73fd317a 100644 --- a/tests/tokenizer/test_tokenize.py +++ b/tests/tokenizer/test_tokenize.py @@ -405,6 +405,36 @@ def encode(self, prompt, return_tensors=None, add_special_tokens=True): return torch.tensor([[4, 5, 6]], dtype=torch.long) +def test_dsv41_encoder_gets_image_sources_before_and_after_fetch(tmp_path): + import copy + from freetoken.mm.media import collect_image_refs + encoding_dir = tmp_path / "encoding" + encoding_dir.mkdir() + (encoding_dir / "encoding_dsv4.py").write_text(''' +def encode_messages(messages, thinking_mode, reasoning_effort=None): + parts = messages[0]["content"] + assert parts[1] == {"type": "text", "text": "compare"} + assert parts[0].get("url") == "https://example.com/image.png" or parts[0].get("data") == b"first image bytes" + assert parts[2].get("data") in ("aW1hZ2U=", b"second image bytes") + return "two images" +''') + messages = [{"role": "user", "content": [ + {"type": "image", "freetoken_ref": {"kind": "url", "data": "https://example.com/image.png"}}, + {"type": "text", "text": "compare"}, + {"type": "image", "freetoken_ref": {"kind": "b64", "data": "aW1hZ2U="}}, + ]}] + manager = TokenizeManager(FakeDsv4Tokenizer(tmp_path)) + msg = TokenizeMsg(uid=1, text=messages, sampling_params=SamplingParams()) + original = copy.deepcopy(messages) + assert manager.render_prompt(msg) == "two images" + assert messages == original + assert len(collect_image_refs(messages)) == 2 + msg.images = [b"first image bytes", b"second image bytes"] + stripped = copy.deepcopy(messages) + assert manager.render_prompt(msg) == "two images" + assert messages == stripped + + def test_tokenize_manager_uses_dsv4_encoder_when_chat_template_is_missing(tmp_path): encoding_dir = tmp_path / "encoding" encoding_dir.mkdir() @@ -761,3 +791,28 @@ def test_tokenize_survives_an_unhashable_effort(): manager.tokenize([msg]) assert "reasoning_effort" not in tokenizer.chat_template_kwargs + + +@pytest.mark.parametrize("budget", [1, 25, 50, 75, 100]) +def test_numeric_encoder_effort_survives_prompt_rendering(tmp_path, budget): + encoding_dir = tmp_path / "encoding" + encoding_dir.mkdir() + (encoding_dir / "encoding.py").write_text(''' +def encode_messages(messages, thinking_mode, reasoning_effort=None): + if thinking_mode != "thinking": + return "chat" + effort = "high" if reasoning_effort is None else reasoning_effort + if isinstance(effort, str): + effort = {"low": 50, "high": 75, "max": 100}[effort] + if type(effort) is not int or not 1 <= effort <= 100: + raise ValueError("invalid budget") + return f"Reasoning Effort: {effort}" +''') + manager = TokenizeManager(FakeDsv4Tokenizer(tmp_path)) + msg = TokenizeMsg( + uid=1, text=[{"role": "user", "content": "hello"}], + sampling_params=SamplingParams(), + chat_template_kwargs={"enable_thinking": True, "reasoning_effort": budget}, + ) + assert manager.render_prompt(msg) == f"Reasoning Effort: {budget}" + assert msg.chat_template_kwargs["reasoning_effort"] == budget From 377a3bc4284f032e9d30777ce099828a158e7e4c Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Wed, 7 Oct 2026 13:24:45 -0700 Subject: [PATCH 08/73] ci: auto-label amd issues and prs by title (#620) --- .github/issue-title-labeler.yml | 2 ++ .github/workflows/issue-labels.yml | 18 +++++++++++++++--- 2 files changed, 17 insertions(+), 3 deletions(-) create mode 100644 .github/issue-title-labeler.yml diff --git a/.github/issue-title-labeler.yml b/.github/issue-title-labeler.yml new file mode 100644 index 0000000000..48826301d5 --- /dev/null +++ b/.github/issue-title-labeler.yml @@ -0,0 +1,2 @@ +amd: + - '/\b(ROCm|HIP|RCCL|RDNA\s?\d?|CDNA\s?\d?|gfx\d{3,4}[a-z]?|MI\d{3}X?|RX\s?\d{4}|Radeon(?!\s*(\d{3}M|Graphics|Vega))|AMD[\s-]*(GPUs?|graphics|cards?|Instinct))\b|\bAMD[\s-]*显卡|\(amd\)|(? Date: Mon, 14 Sep 2026 18:23:04 -0400 Subject: [PATCH 09/73] test(kvcache): give the qsa fp8 pool fakes the mrope flag main now reads Assisted-by: Claude Opus 5 (cherry picked from commit 150c95d6f39b83e0516f5c9c42fef54c679eae09) (cherry picked from commit 4db8b858375b01ff7edf37031dc81671b0181853) (cherry picked from commit 495770e33dcb5c3aa6979c6515b61dfade5e29b3) --- tests/kvcache/test_qsa_pool_fp8.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/kvcache/test_qsa_pool_fp8.py b/tests/kvcache/test_qsa_pool_fp8.py index c0a91f862a..02b525f679 100644 --- a/tests/kvcache/test_qsa_pool_fp8.py +++ b/tests/kvcache/test_qsa_pool_fp8.py @@ -74,7 +74,7 @@ def _spec(): def _config(kv_quant, *, page_size=PAGE_SIZE, max_running_req=3): - mc = SimpleNamespace(num_layers=8, has_swa_attention=False, has_linear_attention=True) + mc = SimpleNamespace(num_layers=8, has_swa_attention=False, has_linear_attention=True, model_is_mrope=False) mc.kv_cache_group_specs = lambda: (_spec(),) return SimpleNamespace( model_config=mc, @@ -200,7 +200,7 @@ def test_factory_threads_kv_quant_into_the_qsa_pool(): mc = SimpleNamespace( num_layers=8, has_swa_attention=False, has_linear_attention=True, - num_kv_heads=HEADS, head_dim=DIM, dsv4_args=None, + num_kv_heads=HEADS, head_dim=DIM, dsv4_args=None, model_is_mrope=False, ) mc.kv_cache_group_specs = lambda: (_spec(),) pool = create_kvcache_pool( From fda3efd3e637e2e2e84e9b9fb979b6d0679d01ce Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Mon, 14 Sep 2026 23:02:28 -0400 Subject: [PATCH 10/73] fix(models): restore the os import and truthy set the fp8 env helpers read Assisted-by: Claude Opus 5 (cherry picked from commit 82a21ae2432fb891175c8c881f52cb4b710241bc) (cherry picked from commit 0369ed94d7d085ece40ffdcc0adf93a5c26a6b4e) (cherry picked from commit b5c7812fceb68b4b7835a80ab7c533e0385d9577) --- python/freetoken/models/config.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py index fd0802b256..2370d7b82c 100644 --- a/python/freetoken/models/config.py +++ b/python/freetoken/models/config.py @@ -1,4 +1,6 @@ from __future__ import annotations + +import os from dataclasses import dataclass from typing import Any, ClassVar, Dict, List, Literal, Tuple, TypeAlias @@ -7,12 +9,14 @@ # State-dict key prefixes of the vision stack; load_weight drops them when the engine serves text-only. VISION_KEY_PREFIXES = ("vision_tower.", "embed_vision.", "vision_embedder.", "visual.") +_ENV_TRUE = {"1", "true", "yes", "on"} + def fp8_dense_enabled() -> bool: """Load-time FP8 for a model's bf16 attention / GDN projections (opt-in, default OFF): per-tensor e4m3 weights, per-tensor dynamic activation scale, cuBLASLt W8A8 GEMMs (``torch._scaled_mm``, sm_89+). ``FREETOKEN_FP8_DENSE=1``.""" - return os.getenv("FREETOKEN_FP8_DENSE", "0").strip().lower() in _VISION_TRUE + return os.getenv("FREETOKEN_FP8_DENSE", "0").strip().lower() in _ENV_TRUE def fp8_lmhead_enabled() -> bool: @@ -20,7 +24,7 @@ def fp8_lmhead_enabled() -> bool: for the rest). Separate from :func:`fp8_dense_enabled` because this one moves the logits: a per-tensor e4m3 vocab matrix changes every sampled token's score, so it carries its own quality gate. ``FREETOKEN_FP8_LMHEAD=1``.""" - return os.getenv("FREETOKEN_FP8_LMHEAD", "0").strip().lower() in _VISION_TRUE + return os.getenv("FREETOKEN_FP8_LMHEAD", "0").strip().lower() in _ENV_TRUE def detect_expert_quant(hf_config: Any) -> str: From 35d2e25c9b4ae6e6259603d4e03043e6b088bae3 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Mon, 14 Sep 2026 23:00:44 -0400 Subject: [PATCH 11/73] test(qwen4_exp): restore the rope device and clear get_rope after the cpu build test Assisted-by: Claude Opus 5 (cherry picked from commit 27ad92e0a554be016d6e533683c40a2eb67c6164) (cherry picked from commit e152f665cc1f35f073fbb82daedece77d02f91e1) (cherry picked from commit 333f7f9f0501116694ede067202070353faf341b) --- tests/models/qwen4_exp/test_build_cpu.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/tests/models/qwen4_exp/test_build_cpu.py b/tests/models/qwen4_exp/test_build_cpu.py index 0cf2027d14..58ad0c4dc3 100644 --- a/tests/models/qwen4_exp/test_build_cpu.py +++ b/tests/models/qwen4_exp/test_build_cpu.py @@ -10,7 +10,7 @@ import torch from freetoken.distributed import set_tp_info, try_get_tp_info -from freetoken.layers import set_rope_device +from freetoken.layers import rotary, set_rope_device from freetoken.models.qwen4_exp.config import parse_config from freetoken.models.qwen4_exp.model import Qwen4ExpForCausalLM @@ -20,9 +20,15 @@ def _build(): if try_get_tp_info() is None: set_tp_info(rank=0, size=1) + previous = rotary._ROPE_DEVICE set_rope_device(torch.device("cpu")) - with torch.device("meta"): - return Qwen4ExpForCausalLM(parse_config(toy_hf_config())) + try: + with torch.device("meta"): + return Qwen4ExpForCausalLM(parse_config(toy_hf_config())) + finally: + # get_rope is cached per geometry: a CPU cos/sin table left behind feeds later CUDA tests + rotary.get_rope.cache_clear() + set_rope_device(previous) def test_model_constructs_on_meta_without_cuda(): From 7dc5db7c85823a9cc648c69eaea0ce24e70ae899 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Mon, 14 Sep 2026 23:19:07 -0400 Subject: [PATCH 12/73] test(kernels): give the block_ends backend fake the kv scale accessors the quantized-kv backend reads Assisted-by: Claude Opus 5 (cherry picked from commit 600819fd7ce18c813cd4b8400cf534bff38f6f7f) (cherry picked from commit dd2ebdfda3e1c53e5269bd9adfbee7291cc8ffc2) (cherry picked from commit 12754d41c5fe646ccc6081ce26f05fdf7b0a57b9) --- tests/kernels/test_triton_attention.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/kernels/test_triton_attention.py b/tests/kernels/test_triton_attention.py index a5abb3d7b3..3176ad71ce 100644 --- a/tests/kernels/test_triton_attention.py +++ b/tests/kernels/test_triton_attention.py @@ -804,6 +804,12 @@ def k_cache(self, layer_id): def v_cache(self, layer_id): return self.v + def k_scale(self, layer_id): # 16-bit pool: no fp8 scales + return None + + def v_scale(self, layer_id): + return None + device = torch.device("cuda") head_dim = 256 kv_cache = FakeKVCache(device, head_dim) From fa609ed60f7ba88f7520fa37b7c592709b02ba73 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 23 Sep 2026 19:17:58 -0400 Subject: [PATCH 13/73] merge fork #18: fix(moe): report the routing top-k figure as a static bound, not an oracle --- python/freetoken/engine/engine.py | 19 +++++++---------- python/freetoken/moe/offload_cache.py | 12 +++++------ tests/moe/test_moe_collect_stats.py | 30 +++++++++++++++++++++++---- 3 files changed, 40 insertions(+), 21 deletions(-) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 6c2e50b933..9740954dbd 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -820,9 +820,8 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: cache.collect_stats = config.moe_collect_stats # The routing histogram rides the same switch: on its own the miss rate says how # often we fetch, but not whether a smarter policy could have avoided the fetch. - # decode_routing_stats turns it into an oracle hit rate -- the ceiling any policy - # holding this many slots could reach on the observed routing -- which is the number - # worth having before anyone rewrites eviction. + # decode_routing_stats turns it into routing-skew numbers, including the hit rate of + # the best fixed expert set per layer (not a ceiling: LRU's temporal locality can beat it). cache.collect_decode_freq = config.moe_collect_stats # attach_offload_moe_cache walks for OffloadMoELayers, or defers to a model's # _iter_offload_moe_layers() hook when its MoE blocks are bespoke nn.Modules (DSV4). @@ -1179,8 +1178,8 @@ def _emit_moe_stats(self) -> None: Accumulation is device-side and captured into the decode graph, so it is free to leave running; reading it is not (the counters have to come back to the host), which is why this only fires every MOE_STATS_INTERVAL decode steps. The routing histogram - is deliberately *not* reset -- the oracle bound wants the whole run's distribution, - not one window's. + is deliberately *not* reset -- it describes the whole run's distribution, not one + window's, so the static top-k figure flattens as a session covers more topics. """ cache = self.moe_offload_cache agg = cache.decode_miss_stats() @@ -1206,14 +1205,12 @@ def _emit_moe_stats(self) -> None: ) routing = cache.decode_routing_stats() if routing: - # oracle_hit_at_slots is the upper bound on hit rate for *any* policy with this - # many slots per layer. If it sits near the realized hit rate, the cache is - # already doing as well as the routing allows and the win has to come from - # somewhere else (more slots, more bandwidth); if it sits far above, eviction - # policy is leaving something on the table. + # static_topk_hit is what pinning each layer's most frequent experts would score over + # the whole run. A realized hit rate above it means the cache is winning on temporal + # locality; it says nothing about how close to optimal eviction is. logger.info_rank0( "MoE routing: " - f"oracle_hit={routing['oracle_hit_at_slots']:.3f} " + f"static_topk_hit={routing['static_topk_hit_at_slots']:.3f} " f"(realized {1.0 - agg['miss_rate']:.3f}), " f"slots/layer={routing['slots_per_layer']:.1f}, " f"working_set={routing['working_set_mean']:.1f}" diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 0969f73829..b2659d41b6 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -982,10 +982,10 @@ def decode_miss_stats_per_layer(self) -> dict: def decode_routing_stats(self) -> dict: """Per-layer decode routing concentration, for cache-skew analysis. - Uses the histogram from ``collect_decode_freq``. The ``oracle_hit`` is the best a - per-layer LRU holding ``cache_size/num_layers`` slots could achieve on the observed - (stationary) routing distribution -- i.e. an upper bound on hit rate that depends - purely on how skewed routing is, independent of any LRU/LFU dynamics. + Uses the histogram from ``collect_decode_freq``. ``static_topk_hit_at_slots`` is the + hit rate of the best *fixed* expert set: each layer's ``cache_size/num_layers`` (rounded) + most frequent experts over the whole histogram. It measures routing skew only. It is + not an upper bound on a dynamic cache: LRU exploits temporal locality and can beat it. """ freq = self.decode_freq.float() total = freq.sum(dim=1) @@ -995,7 +995,7 @@ def decode_routing_stats(self) -> dict: slots_per_layer = self.cache_size / self.num_layers C = max(1, int(round(slots_per_layer))) sorted_f, _ = torch.sort(freq, dim=1, descending=True) - oracle_hit = (sorted_f[:, :C].sum(dim=1)[valid] / total[valid]).mean().item() + static_hit = (sorted_f[:, :C].sum(dim=1)[valid] / total[valid]).mean().item() ws = (freq > 0).sum(dim=1).float() cdf = torch.cumsum(sorted_f, dim=1) / total.clamp(min=1).unsqueeze(1) cover90 = ((cdf < 0.9).sum(dim=1).float() + 1)[valid] @@ -1007,7 +1007,7 @@ def decode_routing_stats(self) -> dict: "working_set_mean": ws[valid].mean().item(), "working_set_max": int(ws[valid].max().item()), "experts_for_90pct": cover90.mean().item(), - "oracle_hit_at_slots": oracle_hit, + "static_topk_hit_at_slots": static_hit, "norm_entropy": norm_ent, } diff --git a/tests/moe/test_moe_collect_stats.py b/tests/moe/test_moe_collect_stats.py index f960abbdc0..bd1c4e9b3f 100644 --- a/tests/moe/test_moe_collect_stats.py +++ b/tests/moe/test_moe_collect_stats.py @@ -10,7 +10,10 @@ import io from types import SimpleNamespace +import torch + from freetoken.engine.engine import MOE_STATS_INTERVAL, Engine +from freetoken.moe.offload_cache import OffloadMoeCache from freetoken.server.args import ServerArgs, parse_args @@ -63,7 +66,7 @@ def decode_routing_stats(self): "working_set_mean": 173.1, "working_set_max": 243, "experts_for_90pct": 92.3, - "oracle_hit_at_slots": 0.764, + "static_topk_hit_at_slots": 0.764, "norm_entropy": 0.813, } @@ -82,9 +85,9 @@ def test_emit_reports_and_resets_the_window(caplog): cache = _StubCache() out = _emit(cache, caplog) assert "miss_rate=0.250" in out - # The oracle bound is the whole point of the report: it says how much room a different - # eviction policy could possibly have. - assert "oracle_hit=0.764" in out + # Labelled as the fixed-set figure it is; "oracle" implied an upper bound it is not. + assert "static_topk_hit=0.764" in out + assert "oracle" not in out assert "(realized 0.750)" in out # Ranked worst-first, and the layer that never ran is left out entirely. assert "L0=0.500, L1=0.100" in out @@ -107,3 +110,22 @@ def test_idle_window_emits_nothing_and_keeps_counters(caplog): def test_interval_is_a_sane_window(): assert MOE_STATS_INTERVAL >= 1 + + +def test_static_topk_hit_is_not_an_upper_bound_on_lru(): + """One slot, four experts each routed in a run of four tokens: A A A A B B B B ... + + The best fixed expert catches 4 of 16 tokens; an LRU holding one slot misses only on + each run's first token. The figure is routing skew, not a ceiling on a dynamic cache. + """ + trace = [e for e in range(4) for _ in range(4)] + cached, lru_hits = None, 0 + for e in trace: + lru_hits += e == cached + cached = e + freq = torch.bincount(torch.tensor(trace), minlength=4).unsqueeze(0) + cache = SimpleNamespace(decode_freq=freq, cache_size=1, num_layers=1, num_experts=4) + stats = OffloadMoeCache.decode_routing_stats(cache) + assert stats["static_topk_hit_at_slots"] == 0.25 + assert "oracle_hit_at_slots" not in stats + assert lru_hits / len(trace) == 0.75 From 11338882bc77105a373702107239dd69782981dc Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 23 Sep 2026 19:17:58 -0400 Subject: [PATCH 14/73] merge fork #19: fix(qwen4_exp): refuse an ftw checkpoint whose fp8 split differs from the fp8 env flags --- python/freetoken/models/qwen4_exp/__init__.py | 2 + python/freetoken/models/qwen4_exp/weight.py | 26 +++++++ python/freetoken/models/weight.py | 10 ++- tests/models/qwen4_exp/test_fp8_dense.py | 69 +++++++++++++++++++ 4 files changed, 106 insertions(+), 1 deletion(-) diff --git a/python/freetoken/models/qwen4_exp/__init__.py b/python/freetoken/models/qwen4_exp/__init__.py index 0fe4a4ffa8..84154b26f4 100644 --- a/python/freetoken/models/qwen4_exp/__init__.py +++ b/python/freetoken/models/qwen4_exp/__init__.py @@ -12,6 +12,7 @@ from .config import parse_config from .model import Qwen4ExpForCausalLM, Qwen4ExpForConditionalGeneration from .weight import ( + check_ftw_weights, ftw_side_files, nvfp4_expert_spec, iter_vision_weights, @@ -24,6 +25,7 @@ from freetoken.models.qwen3_5_moe.weight import iter_expert_pieces __all__ = [ + "check_ftw_weights", "ftw_side_files", "nvfp4_expert_spec", "Qwen4ExpForCausalLM", diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index 30a3d1f797..91c8c1bd17 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -417,6 +417,32 @@ def _ple_table_files(folder: str) -> list[str]: return sorted(os.path.join(folder, shard) for shard in files) +def check_ftw_weights(config, model_path: str, names: set[str]) -> None: + """Refuse an FTW checkpoint whose stored FP8 split disagrees with this process's FP8 flags. + + The FTW holds the tensors ``iter_weights`` emitted at conversion, so a mismatch otherwise + surfaces only after the dense read, as a KeyError on the first fused GDN projection.""" + from .config import use_fp8_lmhead + + stored = { + "FREETOKEN_FP8_DENSE": any(n.endswith(".in_proj_qkvz.weight_scale") for n in names), + "FREETOKEN_FP8_LMHEAD": "lm_head.weight_scale" in names, + } + wanted = { + "FREETOKEN_FP8_DENSE": config.attn_quant == "fp8_dynamic", + "FREETOKEN_FP8_LMHEAD": use_fp8_lmhead(config), + } + wrong = [var for var in stored if stored[var] != wanted[var]] + if wrong: + was = ", ".join(f"{var}={int(stored[var])}" for var in wrong) + now = ", ".join(f"{var}={int(wanted[var])}" for var in wrong) + raise ValueError( + f"FTW checkpoint {model_path} stores its dense weights as converted with {was}, " + f"but this process has {now}; serve it with the conversion-time setting, or " + "reconvert with ft checkpoint under the setting you want" + ) + + def ftw_side_files(model_path: str, out_dir: str) -> list[str]: """Write the PLE n-gram table tensors, and only those, into ``ple-table-*.safetensors`` next to an FTW checkpoint. diff --git a/python/freetoken/models/weight.py b/python/freetoken/models/weight.py index e3382b1b59..9016b9477e 100644 --- a/python/freetoken/models/weight.py +++ b/python/freetoken/models/weight.py @@ -218,12 +218,20 @@ def load_weight( # decided at conversion (offload -> experts live in banks, not here); a backend mismatch # 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.ftw import FTWReader, is_ftw_checkpoint, iter_ftw_weights from freetoken.models.config import VISION_KEY_PREFIXES # a text-only engine never built the tower, so its tensors are not even read keep = None if include_vision else (lambda name: not name.startswith(VISION_KEY_PREFIXES)) if is_ftw_checkpoint(model_path): + spec = get_model_spec(cached_load_hf_config(model_path).architectures[0]) + # a family whose stored tensors depend on load-time settings checks them against the index before the read + check = _model_override(spec, "check_ftw_weights") + if check is not None: + reader = FTWReader(model_path) + names = {e["name"] for e in reader.entries("weight")} + reader.close() + check(_spec_for_model_path(model_path)[0], model_path, names) weights = iter_ftw_weights(model_path, keep=keep) else: _config, spec = _spec_for_model_path(model_path) diff --git a/tests/models/qwen4_exp/test_fp8_dense.py b/tests/models/qwen4_exp/test_fp8_dense.py index 058515b183..74da973341 100644 --- a/tests/models/qwen4_exp/test_fp8_dense.py +++ b/tests/models/qwen4_exp/test_fp8_dense.py @@ -151,3 +151,72 @@ def test_reader_gates_lm_head_behind_its_own_flag(): passthrough = dict(_fp8_dense("x.self_attn.qkv_proj.weight", t, cfg, 1, dense=False, lm_head=True)) assert list(passthrough) == ["x.self_attn.qkv_proj.weight"] + + +def _ftw_with(tmp_path, names): + from freetoken.checkpoint.ftw import FTWWriter + + writer = FTWWriter(str(tmp_path)) + for name in names: + writer.add_tensor(name, torch.zeros(4, dtype=torch.bfloat16)) + writer.finalize({}) + return str(tmp_path) + + +_BF16_NAMES = ["model.layers.0.linear_attn.in_proj.weight", "lm_head.weight"] +_FP8_NAMES = [ + "model.layers.0.linear_attn.in_proj_qkvz.weight", + "model.layers.0.linear_attn.in_proj_qkvz.weight_scale", + "model.layers.0.linear_attn.in_proj_ba.weight", + "lm_head.weight", + "lm_head.weight_scale", +] + + +@pytest.mark.parametrize( + "names, env, bad", + [ + (_BF16_NAMES, {"FREETOKEN_FP8_DENSE": "1", "FREETOKEN_FP8_LMHEAD": "1"}, + "FREETOKEN_FP8_DENSE=0, FREETOKEN_FP8_LMHEAD=0"), + (_FP8_NAMES, {}, "FREETOKEN_FP8_DENSE=1, FREETOKEN_FP8_LMHEAD=1"), + (_FP8_NAMES, {"FREETOKEN_FP8_DENSE": "1"}, "converted with FREETOKEN_FP8_LMHEAD=1, but"), + ], +) +def test_ftw_load_refuses_fp8_flags_that_differ_from_conversion(tmp_path, monkeypatch, names, env, bad): + """#8: the FTW replays conversion-time tensors, so the FP8 flags must match before any byte is read.""" + from freetoken.models import weight as model_weight + + for var in ("FREETOKEN_FP8_DENSE", "FREETOKEN_FP8_LMHEAD"): + monkeypatch.delenv(var, raising=False) + for var, value in env.items(): + monkeypatch.setenv(var, value) + path = _ftw_with(tmp_path, names) + spec = SimpleNamespace(module="freetoken.models.qwen4_exp") + config = SimpleNamespace( + attn_quant="fp8_dynamic" if env.get("FREETOKEN_FP8_DENSE") else "none", + tie_word_embeddings=False, + quant=None, + ) + monkeypatch.setattr(model_weight, "cached_load_hf_config", lambda _: SimpleNamespace(architectures=["x"])) + monkeypatch.setattr(model_weight, "get_model_spec", lambda _: spec) + monkeypatch.setattr(model_weight, "_spec_for_model_path", lambda _: (config, spec)) + with pytest.raises(ValueError, match=bad): + next(model_weight.load_weight(path, torch.device("cpu"))) + + +@pytest.mark.parametrize("names, fp8", [(_BF16_NAMES, False), (_FP8_NAMES, True)]) +def test_ftw_load_accepts_matching_fp8_flags(tmp_path, monkeypatch, names, fp8): + from freetoken.models import weight as model_weight + + for var in ("FREETOKEN_FP8_DENSE", "FREETOKEN_FP8_LMHEAD"): + if fp8: + monkeypatch.setenv(var, "1") + else: + monkeypatch.delenv(var, raising=False) + path = _ftw_with(tmp_path, names) + spec = SimpleNamespace(module="freetoken.models.qwen4_exp") + config = SimpleNamespace(attn_quant="fp8_dynamic" if fp8 else "none", tie_word_embeddings=False, quant=None) + monkeypatch.setattr(model_weight, "cached_load_hf_config", lambda _: SimpleNamespace(architectures=["x"])) + monkeypatch.setattr(model_weight, "get_model_spec", lambda _: spec) + monkeypatch.setattr(model_weight, "_spec_for_model_path", lambda _: (config, spec)) + assert [name for name, _ in model_weight.load_weight(path, torch.device("cpu"))] == names From ed8de88e58ec3635b52d40dc5d6d92930e9c9f86 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 23 Sep 2026 19:17:58 -0400 Subject: [PATCH 15/73] merge fork #20: fix(engine): let moe-cache-auto drop prefill overlap when its slot floor does not fit --- python/freetoken/engine/cache_budget.py | 16 +++++++++-- python/freetoken/engine/engine.py | 5 ++++ tests/engine/test_cache_budget.py | 37 +++++++++++++++++++++++++ 3 files changed, 55 insertions(+), 3 deletions(-) diff --git a/python/freetoken/engine/cache_budget.py b/python/freetoken/engine/cache_budget.py index b5f07ee8ab..f297b9002d 100644 --- a/python/freetoken/engine/cache_budget.py +++ b/python/freetoken/engine/cache_budget.py @@ -54,6 +54,7 @@ def plan_cache_budget( prefill_overlap: bool, kv_reserve_pages: int, max_slots: int, + page_size: int = 1, ) -> tuple[int, int, bool]: """Split ``budget_bytes`` MoE-first into (moe_cache_size, num_pages, prefill_overlap). @@ -61,7 +62,8 @@ def plan_cache_budget( weights + fixed_cache_size; the (1-memory_ratio) remainder is the graph headroom). Experts greedily fill the budget after reserving ``kv_reserve_pages`` for KV, clamped to ``[floor, min(total_experts, max_slots)]`` (floor is ``2*num_experts`` when prefill - overlap is feasible else ``num_experts``); KV pages take whatever remains. + overlap is feasible, in slots and in bytes, else ``num_experts``); KV pages take whatever + remains. ``page_size`` (tokens per page) only feeds the error message. """ assert per_expert_bytes > 0, "per_expert_bytes must be positive" assert cache_per_page > 0, "cache_per_page must be positive (owned-KV models unsupported here)" @@ -74,6 +76,10 @@ def plan_cache_budget( assert hi >= lo, f"slot cap {hi} below the minimum {lo} slots" kv_reserve_bytes = kv_reserve_pages * cache_per_page + # Overlap is a prefill speedup, not a requirement: when its 2*num_experts floor leaves no + # room for the KV reserve (and the 2 pages asserted below), plan without it. + if overlap and 2 * num_experts * per_expert_bytes + max(kv_reserve_pages, 2) * cache_per_page > budget_bytes: + overlap, lo = False, num_experts # MoE-priority: reserve KV first, then experts greedily take the remaining budget. raw = (budget_bytes - kv_reserve_bytes) // per_expert_bytes moe_cache_size = max(lo, min(raw, hi)) @@ -86,10 +92,13 @@ def plan_cache_budget( # the reserve (or negative), yielding a plan that exceeds budget_bytes. Reject here so # --moe-cache-auto fails in arithmetic instead of OOMing in a later CUDA allocation. total = moe_cache_size * per_expert_bytes + num_pages * cache_per_page + fit_pages = max(0, (budget_bytes - moe_cache_size * per_expert_bytes) // cache_per_page) assert total <= budget_bytes, ( f"cache budget too small: minimum plan (moe={moe_cache_size} slots, " - f"kv={num_pages} pages) needs {total} B > budget {budget_bytes} B " - "(raise memory_ratio, lower kv_reserve_tokens, or free GPU memory)" + f"kv={num_pages} pages) needs {total / 2**30:.2f} GiB > budget {budget_bytes / 2**30:.2f} GiB; " + f"beside {moe_cache_size} slots at most {fit_pages} KV pages ({fit_pages * page_size} tokens) fit. " + "Lower --kv-reserve-tokens, use a smaller --kv-cache-dtype, raise --memory-ratio, " + "or set --moe-cache-size and --max-seq-len-override explicitly" ) assert num_pages > 1, "not enough memory for KV cache after MoE allocation" return moe_cache_size, num_pages, overlap @@ -130,4 +139,5 @@ def resolve_moe_cache_auto( prefill_overlap=prefill_overlap, kv_reserve_pages=kv_reserve_pages, max_slots=max_slots, + page_size=page_size, ) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 9740954dbd..b01599291d 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -768,6 +768,11 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: raise RuntimeError(f"{exc}; {_pin_hint(self._host_tables_bytes)}") from exc if config.moe_cache_auto: size, pages, overlap = self._resolve_auto_moe_cache_size(config, banks, method) + if config.moe_prefill_overlap and not overlap: + logger.info_rank0( + f"--moe-cache-auto: prefill overlap disabled, its {2 * config.model_config.num_experts} " + "slot floor does not fit beside the KV reserve" + ) object.__setattr__(config, "moe_cache_size", size) object.__setattr__(config, "moe_prefill_overlap", overlap) if config.num_page_override is None: diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py index a77409f013..aa0332dec4 100644 --- a/tests/engine/test_cache_budget.py +++ b/tests/engine/test_cache_budget.py @@ -80,6 +80,43 @@ def test_budget_too_small_for_min_moe_plus_reserve_raises(): ) # min moe = 4 slots (400 B) + reserve (10 pages = 100 B) = 500 B > 300 B budget +def test_overlap_floor_that_does_not_fit_falls_back_to_num_experts(): + # #4: 2*num_experts slots (800 B) + the KV reserve (100 B) exceed the 700 B budget, but + # num_experts slots do fit. Plan without overlap instead of refusing to start. + size, pages, overlap = plan_cache_budget( + budget_bytes=700, per_expert_bytes=100, cache_per_page=10, + num_experts=4, total_experts=50, prefill_overlap=True, + kv_reserve_pages=10, max_slots=50, + ) + assert overlap is False + assert size == 6 # (700 - 100) // 100, above the num_experts floor + assert pages == 10 + assert size * 100 + pages * 10 <= 700 + + +def test_overlap_kept_when_its_floor_fits_the_budget(): + size, pages, overlap = plan_cache_budget( + budget_bytes=900, per_expert_bytes=100, cache_per_page=10, + num_experts=4, total_experts=50, prefill_overlap=True, + kv_reserve_pages=10, max_slots=50, + ) + assert (size, pages, overlap) == (8, 10, True) + + +def test_budget_too_small_message_names_what_fits_and_the_flags(): + # num_experts slots (400 B) + 10 reserved pages (100 B) > 450 B: 5 pages fit beside them. + with pytest.raises(AssertionError) as err: + plan_cache_budget( + budget_bytes=450, per_expert_bytes=100, cache_per_page=10, + num_experts=4, total_experts=50, prefill_overlap=True, + kv_reserve_pages=10, max_slots=50, page_size=64, + ) + msg = str(err.value) + assert "beside 4 slots at most 5 KV pages (320 tokens) fit" in msg + for flag in ("--kv-reserve-tokens", "--kv-cache-dtype", "--memory-ratio", "--moe-cache-size"): + assert flag in msg + + def test_prefill_overlap_false_is_honored(): # Even when the cache could fit 2*num_experts, an explicit False stays False. size, pages, overlap = plan_cache_budget( From 02e9f04b6922502768e7039f6d5037ac45fbb929 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 23 Sep 2026 19:17:59 -0400 Subject: [PATCH 16/73] merge fork #21: fix(triton): keep the float glm dsa sparse launches two-stage to fit sm_120 shared memory --- python/freetoken/kernel/triton/glm_dsa_sparse.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/python/freetoken/kernel/triton/glm_dsa_sparse.py b/python/freetoken/kernel/triton/glm_dsa_sparse.py index d73ece8cf9..bac2cd7eed 100644 --- a/python/freetoken/kernel/triton/glm_dsa_sparse.py +++ b/python/freetoken/kernel/triton/glm_dsa_sparse.py @@ -498,9 +498,9 @@ def glm_dsa_sparse_attn( D_V=d_v, D_R=d_r, BLOCK_H=BLOCK_H, BLOCK_T=block_t, HAS_COUNTS=has_counts, HAS_ROPE=d_r > 0, HAS_FP8=has_fp8, HAS_NVFP4=has_nvfp4, NUM_SPLITS=n_splits, - # The 512-wide latent accumulator exceeds the 99 KiB shared-memory - # limit of consumer Blackwell GPUs with a two-stage pipeline. - num_warps=4, num_stages=1, + # Consumer Blackwell caps kernel shared memory at 99 KiB. Single-stage, the float path + # needs 100 KiB; two stages (its setting before quantized KV) fit. + num_warps=4, num_stages=1 if has_fp8 or has_nvfp4 else 2, ) grid2 = (m, b, h) _glm_dsa_merge_kernel[grid2]( @@ -527,7 +527,8 @@ def glm_dsa_sparse_attn( D_V=d_v, D_R=d_r, BLOCK_H=BLOCK_H, BLOCK_T=block_t, HAS_COUNTS=has_counts, HAS_ROPE=d_r > 0, HAS_FP8=has_fp8, HAS_NVFP4=has_nvfp4, - num_warps=4, num_stages=1, + # same shared-memory split as the split-k launch above + num_warps=4, num_stages=1 if has_fp8 or has_nvfp4 else 2, ) return o From 460c011c4a42029da913d63a45b1c1fc36c6ad60 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 23 Sep 2026 19:17:59 -0400 Subject: [PATCH 17/73] merge fork #22: fix(engine): resize shrinking cache pools before growing ones in a runtime rebuild --- python/freetoken/engine/engine.py | 51 ++++++++++++++++++--------- tests/scheduler/test_cache_rebuild.py | 38 ++++++++++++++++++++ 2 files changed, 73 insertions(+), 16 deletions(-) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index b01599291d..dcf6db72ac 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -969,6 +969,40 @@ def _refresh_seq_state(self, config) -> None: self.page_table[self.dummy_req.table_idx].fill_(num_tokens) self.kv_cache.attach_page_table(self.page_table) + def _resize_pools( + self, config, moe_cache_size: int | None, num_pages: int | None, + num_mamba_slots: int | None, num_swa_pages: int | None, + ) -> None: + """Resize the requested pools, shrinking ones first. + + Each pool frees its own tensors before allocating, but a pool that grows ahead of one + that shrinks needs both geometries resident at once: on a nearly full card that OOMs + a resize whose final total fits (MoE 512 -> 1024 slots with KV 262K -> 64K tokens).""" + steps = [] # (shrinks, resize) + if moe_cache_size is not None: + assert self.moe_offload_cache is not None, "no MoE offload cache to resize" + steps.append(( + moe_cache_size < self.moe_offload_cache.cache_size, + lambda: self.moe_offload_cache.rebuild(moe_cache_size), + )) + if num_pages is not None or num_swa_pages is not None: + # A window-only change (num_pages None) re-derives the window pool at the new pin + # against the CURRENT page count; _resize_kv_pool sets self.num_pages either way. + pages = num_pages if num_pages is not None else self.num_pages + steps.append(( + pages < self.num_pages, + lambda: self._resize_kv_pool(config, pages, num_swa_pages), + )) + if num_mamba_slots is not None: + # Must sit between graph teardown and re-capture so the recaptured graphs bind the new + # state tensors. +1 for the reserved padding sink: num_mamba_slots is the usable count. + steps.append(( + num_mamba_slots + 1 < self.linear_state_pool.num_slots, + lambda: self.linear_state_pool.rebuild(num_mamba_slots + 1), + )) + for _, resize in sorted(steps, key=lambda step: not step[0]): + resize() + @torch.inference_mode() def rebuild_runtime_cache( self, @@ -1078,22 +1112,7 @@ def rebuild_runtime_cache( # FrozenInstanceError, which here aborts the rebuild after the CUDA graphs are gone (→ 503). if num_swa_pages is not None: object.__setattr__(config, "swa_num_pages_override", num_swa_pages) - if moe_cache_size is not None: - assert self.moe_offload_cache is not None, "no MoE offload cache to resize" - self.moe_offload_cache.rebuild(moe_cache_size) - if num_pages is not None: - # sets self.num_pages (rebuilds KV + window) - self._resize_kv_pool(config, num_pages, num_swa_pages) - elif num_swa_pages is not None: - # Window-only change: no page-count change, but re-derive the window pool at the new - # pin against the CURRENT page count. This re-allocs the same-size full pool and - # the resized window, both inside the pool's own rebuild_from_config. - self._resize_kv_pool(config, self.num_pages, num_swa_pages) - if num_mamba_slots is not None: - # Reallocate the GDN state pool (frees old tensors first). Must sit between graph - # teardown and re-capture so the recaptured graphs bind the new state tensors. - # +1 for the reserved padding sink: num_mamba_slots is the usable count. - self.linear_state_pool.rebuild(num_mamba_slots + 1) + self._resize_pools(config, moe_cache_size, num_pages, num_mamba_slots, num_swa_pages) # 3. Refresh max_seq_len (+ generic page table) for the new token budget. self._refresh_seq_state(config) aligned_max_seq_len = _page_table_width(self.max_seq_len, config.page_size) diff --git a/tests/scheduler/test_cache_rebuild.py b/tests/scheduler/test_cache_rebuild.py index 0c0580708f..5e280ae7fd 100644 --- a/tests/scheduler/test_cache_rebuild.py +++ b/tests/scheduler/test_cache_rebuild.py @@ -168,3 +168,41 @@ def test_rebuild_cache_refreshes_prefill_budget(monkeypatch): cache_manager.prefill_chunk_budget = 1000 # the (stubbed) engine rebuild shrank the pool Scheduler.rebuild_cache(sched, num_pages=16) assert sched.prefill_budget == 1000 # tracks the shrunk cap, not the stale 5000 + + +def _pool_engine(calls, *, moe_slots=512, num_pages=4096, mamba_slots=9): + """An engine carrying only the three pools _resize_pools touches, each logging its resize.""" + from types import SimpleNamespace + + engine = SimpleNamespace(num_pages=num_pages) + engine.moe_offload_cache = SimpleNamespace( + cache_size=moe_slots, rebuild=lambda n: calls.append(("moe", n))) + engine.linear_state_pool = SimpleNamespace( + num_slots=mamba_slots, rebuild=lambda n: calls.append(("mamba", n))) + engine._resize_kv_pool = lambda config, pages, swa: calls.append(("kv", pages, swa)) + return engine + + +def test_resize_pools_frees_shrinking_pools_before_growing_ones(): + """#5: MoE 512 -> 1024 slots with KV 4096 -> 1024 pages OOMed on a full card because the MoE + pool was reallocated while the KV pool still held its old pages.""" + from freetoken.engine.engine import Engine + + calls = [] + Engine._resize_pools(_pool_engine(calls), None, moe_cache_size=1024, num_pages=1024, + num_mamba_slots=4, num_swa_pages=None) + assert calls == [("kv", 1024, None), ("mamba", 5), ("moe", 1024)] + + calls.clear() # the rollback of that request: MoE shrinks back first, then KV and mamba grow + Engine._resize_pools(_pool_engine(calls, moe_slots=1024, num_pages=1024, mamba_slots=5), None, + moe_cache_size=512, num_pages=4096, num_mamba_slots=8, num_swa_pages=None) + assert calls == [("moe", 512), ("kv", 4096, None), ("mamba", 9)] + + +def test_resize_pools_keeps_request_order_among_growers_and_window_only_uses_current_pages(): + from freetoken.engine.engine import Engine + + calls = [] + Engine._resize_pools(_pool_engine(calls), None, moe_cache_size=600, num_pages=None, + num_mamba_slots=None, num_swa_pages=7) + assert calls == [("moe", 600), ("kv", 4096, 7)] From d85529832368e20a2ddb534d26945ce2b92ae56d Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 23 Sep 2026 19:17:59 -0400 Subject: [PATCH 18/73] merge fork #23 (upstream FlashML-org/FreeToken#17 rebased): fix(server): resolve Hugging Face repo ids to a local checkpoint dir # Conflicts: # python/freetoken/utils/hf.py --- python/freetoken/server/args.py | 20 +++-- python/freetoken/utils/__init__.py | 2 + python/freetoken/utils/hf.py | 45 ++++++++++- tests/server/test_model_path_resolution.py | 87 ++++++++++++++++++++++ 4 files changed, 148 insertions(+), 6 deletions(-) create mode 100644 tests/server/test_model_path_resolution.py diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 52e04c7edb..528ea9bdae 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -920,16 +920,26 @@ def _infer_reasoning_parser(model_path: str) -> str | None: if is_offload_moe_strategy(kwargs["moe_strategy"]) and _no_cache_flag: kwargs["moe_cache_auto"] = True - if kwargs["model_source"] == "modelscope": - model_path = kwargs["model_path"] - if not os.path.isdir(model_path): + # Resolve a hub repo id ("org/model") to a local checkpoint directory. This has to + # happen here, after served_model_name and the parser cascade have read the repo id + # (they want "DeepSeek-V4-Flash-0731", not a snapshot hash) and before anything opens + # the checkpoint: config parsing reads model-specific files straight off disk, so a + # bare repo id reaches `open()` as a relative path and dies with a FileNotFoundError. + model_path = kwargs["model_path"] + if not os.path.isdir(model_path): + if kwargs["model_source"] == "modelscope": from modelscope import snapshot_download ignore_patterns = [] if kwargs["use_dummy_weight"]: ignore_patterns = ["*.bin", "*.safetensors", "*.pt", "*.ckpt"] - model_path = snapshot_download(model_path, ignore_patterns=ignore_patterns) - kwargs["model_path"] = model_path + kwargs["model_path"] = snapshot_download(model_path, ignore_patterns=ignore_patterns) + else: + from freetoken.utils import download_hf_checkpoint + + kwargs["model_path"] = download_hf_checkpoint( + model_path, dummy_weight=kwargs["use_dummy_weight"] + ) del kwargs["model_source"] # "auto" (or an unspecified dtype) resolves to the checkpoint's dtype. Multimodal / diff --git a/python/freetoken/utils/__init__.py b/python/freetoken/utils/__init__.py index af68bf5816..a1dd10b7c3 100644 --- a/python/freetoken/utils/__init__.py +++ b/python/freetoken/utils/__init__.py @@ -9,6 +9,7 @@ ) from .hf import ( cached_load_hf_config, + download_hf_checkpoint, download_hf_weight, load_eos_token_ids, load_generation_sampling, @@ -30,6 +31,7 @@ __all__ = [ "cached_load_hf_config", + "download_hf_checkpoint", "download_hf_weight", "load_eos_token_ids", "load_generation_sampling", diff --git a/python/freetoken/utils/hf.py b/python/freetoken/utils/hf.py index 489585e8c2..6c6a5f7a00 100644 --- a/python/freetoken/utils/hf.py +++ b/python/freetoken/utils/hf.py @@ -2,6 +2,7 @@ import functools import json import os +import re from typing import Any, FrozenSet, Mapping from huggingface_hub import hf_hub_download, snapshot_download @@ -296,4 +297,46 @@ def download_hf_weight(model_path: str) -> str: except Exception as e: raise ValueError( f"Model path '{model_path}' is neither a local directory nor a valid model ID: {e}" - ) \ No newline at end of file + ) + + +# Weight formats FreeToken never loads -- it reads safetensors (plus GGUF for Gemma-4). +# Plenty of repos still ship a duplicate .bin/.pth copy of the same tensors, so skipping +# them can halve a first-time pull. +_UNLOADABLE_WEIGHTS = ("*.bin", "*.pt", "*.pth", "*.h5", "*.msgpack", "*.onnx") + + +# A hub repo id is "org/name" and nothing else. Anything with a leading separator, a +# "~", or more than one "/" is a filesystem path the user got wrong, and turning that +# into a hub lookup replaces "no such directory" with a confusing download error. +_HUB_REPO_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*/[A-Za-z0-9][A-Za-z0-9._-]*$") + + +def looks_like_hub_repo_id(model_path: str) -> bool: + """Whether ``model_path`` should be resolved against the Hub rather than the disk.""" + return bool(_HUB_REPO_ID.match(model_path)) + + +def download_hf_checkpoint(model_path: str, *, dummy_weight: bool = False) -> str: + """Resolve a Hugging Face repo id to a complete local checkpoint directory. + + ``download_hf_weight`` fetches only ``*.safetensors``, which is all the weight + loaders need once the engine is running. Config parsing happens earlier and wants + more than that: the tokenizer, and the model-specific subdirs some architectures + read directly off disk -- DeepSeek-V4 fails without ``inference/config.json`` + (see ``models/deepseek_v4/args.py``). So this pulls the whole snapshot. + + A path that is already a local directory is returned untouched. + """ + if os.path.isdir(model_path) or not looks_like_hub_repo_id(model_path): + return model_path + ignore_patterns = list(_UNLOADABLE_WEIGHTS) + if dummy_weight: + # Config and tokenizer are still needed; the tensors are not. + ignore_patterns += ["*.safetensors", "*.gguf"] + try: + return snapshot_download(model_path, ignore_patterns=ignore_patterns) + except Exception as e: + raise ValueError( + f"Model path '{model_path}' is neither a local directory nor a valid model ID: {e}" + ) diff --git a/tests/server/test_model_path_resolution.py b/tests/server/test_model_path_resolution.py new file mode 100644 index 0000000000..76016e6869 --- /dev/null +++ b/tests/server/test_model_path_resolution.py @@ -0,0 +1,87 @@ +"""``--model org/name``: a hub repo id resolves to a local checkpoint directory. + +``ft serve --model deepseek-ai/DeepSeek-V4-Flash-0731`` used to pass the repo id through +untouched as a filesystem path. Config parsing then opened it relative to cwd: + + FileNotFoundError: No DeepSeek-V4 ModelArgs JSON found under + deepseek-ai/DeepSeek-V4-Flash-0731 (looked for inference/config.json) + +Only the modelscope branch ever downloaded, and ``download_hf_weight`` (used later, by the +weight loaders) fetches ``*.safetensors`` alone -- so it would never have produced +``inference/config.json`` either. + +These pin the resolve for both sources, and pin the two things that must keep reading the +*repo id* rather than the resolved snapshot path: the served model name and the parser +cascade. Resolving too early renames the served model to a snapshot hash. +""" + +from __future__ import annotations + +import sys +import types +from unittest.mock import patch + +import pytest + +from freetoken.server.args import parse_args + +REPO_ID = "deepseek-ai/DeepSeek-V4-Flash-0731" +SNAPSHOT = "/cache/hub/models--deepseek-ai--DeepSeek-V4-Flash-0731/snapshots/7872f01" + +# Explicit dtype and parsers keep parse_args off the network: all three otherwise read the +# checkpoint config, which is exactly what cannot be loaded until the path resolves. +BASE_ARGS = ["--dtype", "bfloat16", "--tool-call-parser", "llama3", "--reasoning-parser", "off"] + + +def _parse(model_path: str, *extra: str): + return parse_args([*BASE_ARGS, "--model", model_path, *extra])[0] + + +@pytest.fixture +def hf_download(): + with patch("freetoken.utils.hf.snapshot_download", return_value=SNAPSHOT) as m: + yield m + + +def test_repo_id_resolves_to_local_snapshot(hf_download): + assert _parse(REPO_ID).model_path == SNAPSHOT + assert hf_download.call_args.args[0] == REPO_ID + + +def test_served_model_name_keeps_the_repo_id(hf_download): + """Not the snapshot hash -- this is the id clients send in the `model` field.""" + assert _parse(REPO_ID).served_model_name == "DeepSeek-V4-Flash-0731" + + +def test_local_dir_is_never_downloaded(tmp_path, hf_download): + assert _parse(str(tmp_path)).model_path == str(tmp_path) + hf_download.assert_not_called() + + +def test_formats_freetoken_cannot_load_are_skipped(hf_download): + """Many repos ship a duplicate .bin copy; the loaders only read safetensors/gguf.""" + _parse(REPO_ID) + ignored = hf_download.call_args.kwargs["ignore_patterns"] + assert "*.bin" in ignored + assert "*.safetensors" not in ignored + + +def test_dummy_weight_skips_the_tensors(hf_download): + """--dummy-weight still needs config + tokenizer, but not 156 GiB of weights.""" + _parse(REPO_ID, "--dummy-weight") + ignored = hf_download.call_args.kwargs["ignore_patterns"] + assert "*.safetensors" in ignored + assert "*.gguf" in ignored + + +def test_modelscope_source_still_uses_modelscope(monkeypatch): + """The isdir check moved out of the modelscope branch; it must still short-circuit.""" + fake = types.ModuleType("modelscope") + calls = [] + fake.snapshot_download = lambda repo, **kw: calls.append((repo, kw)) or SNAPSHOT + monkeypatch.setitem(sys.modules, "modelscope", fake) + + with patch("freetoken.utils.hf.snapshot_download") as hf: + assert _parse(REPO_ID, "--model-source", "modelscope").model_path == SNAPSHOT + hf.assert_not_called() + assert calls[0][0] == REPO_ID From 4f8b0e9e45741333980999525314827aca91c03c Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 23 Sep 2026 19:17:59 -0400 Subject: [PATCH 19/73] merge fork #47 (upstream FlashML-org/FreeToken#393 rebased): OpenAI protocol coverage # Conflicts: # python/freetoken/tokenizer/tokenize.py # Conflicts: # python/freetoken/engine/sample.py # python/freetoken/tokenizer/tokenize.py # Conflicts: # python/freetoken/engine/sample.py # python/freetoken/tokenizer/tokenize.py # Conflicts: # python/freetoken/engine/sample.py # python/freetoken/tokenizer/tokenize.py --- docs/openai_api.md | 101 +++++ python/freetoken/core.py | 34 ++ python/freetoken/engine/sample.py | 181 ++++++++- python/freetoken/message/frontend.py | 3 + python/freetoken/message/tokenizer.py | 4 + python/freetoken/scheduler/scheduler.py | 9 +- python/freetoken/server/api_models.py | 58 +++ python/freetoken/server/api_server.py | 10 +- python/freetoken/server/control_api.py | 18 + python/freetoken/server/generation.py | 58 ++- python/freetoken/server/openai_api.py | 471 ++++++++++++++++------ python/freetoken/server/stats.py | 41 ++ python/freetoken/tokenizer/detokenize.py | 54 ++- python/freetoken/tokenizer/server.py | 24 +- python/freetoken/tokenizer/tokenize.py | 8 +- tests/engine/test_logits_processors.py | 155 +++++++ tests/server/test_openai_extras.py | 383 ++++++++++++++++++ tests/tokenizer/test_detokenize_extras.py | 81 ++++ 18 files changed, 1545 insertions(+), 148 deletions(-) create mode 100644 docs/openai_api.md create mode 100644 tests/engine/test_logits_processors.py create mode 100644 tests/server/test_openai_extras.py create mode 100644 tests/tokenizer/test_detokenize_extras.py diff --git a/docs/openai_api.md b/docs/openai_api.md new file mode 100644 index 0000000000..0b1cc9d93c --- /dev/null +++ b/docs/openai_api.md @@ -0,0 +1,101 @@ +# OpenAI-compatible API: endpoints and parameters + +What `ft serve` accepts on its OpenAI-shaped routes, what it ignores, and what it refuses, +next to the same surface in vLLM, SGLang and llama.cpp's server. "honoured" means the value +changes the generation; "accepted" means the request is not rejected but the value has no +effect; "400" means the server refuses the request with a clear message. + +## Endpoints + +| route | FreeToken | vLLM | SGLang | llama.cpp | +|---|---|---|---|---| +| `POST /v1/chat/completions` | yes | yes | yes | yes | +| `POST /v1/completions` | yes | yes | yes | yes | +| `GET /v1/models`, `GET /v1/models/{id}` | yes | yes | yes | yes | +| `POST /v1/responses` (+ get, cancel) | yes | yes | yes | yes | +| `POST /v1/messages` (Anthropic), `count_tokens` | yes | yes | no | yes | +| `POST /tokenize`, `POST /detokenize` (also under `/v1/`) | yes | yes | yes | yes (`/tokenize`, `/detokenize`) | +| `GET /metrics` (Prometheus) | yes | yes | yes | yes | +| `GET /version` | yes | yes | yes | no (`/props`) | +| `GET /health` | yes | yes | yes | yes | +| `GET /v1/stats`, `/v1/cache/*`, `/v1/requests`, `/v1/admin/*` | yes (FreeToken control plane) | `/metrics` only | `/get_server_info`, `/flush_cache` | `/slots`, `/props`, `/metrics` | +| `POST /v1/embeddings`, `/rerank`, `/score`, `/pooling` | no (generative models only) | embedding / reranker models | embedding / reranker models | yes, any model with pooling | +| `POST /v1/audio/*`, `/v1/files`, `/v1/batches` | no | transcriptions, batch runner (offline) | transcriptions, files, batches | no | +| `POST /infill`, `/apply-template`, `/completion` (non-OpenAI) | `suffix` on `/v1/completions`; `/tokenize` with `messages` renders the template | no | no | yes | + +## Chat completions: request + +| parameter | FreeToken | vLLM | SGLang | llama.cpp | +|---|---|---|---|---| +| `model`, `messages` (text + `image_url` parts) | honoured | honoured | honoured | honoured | +| `max_tokens`, `max_completion_tokens` | honoured (default 32k) | honoured | honoured | `max_tokens` | +| `temperature`, `top_p`, `top_k` | honoured | honoured | honoured | honoured | +| `min_p` | honoured (prob. below `min_p` x top prob dropped) | honoured | honoured | honoured | +| `presence_penalty`, `frequency_penalty` | honoured, over generated tokens | honoured | honoured | honoured | +| `repetition_penalty` | honoured, over prompt + generated tokens (HF semantics) | honoured | honoured | `repeat_penalty` | +| `logit_bias` | honoured (token id -> bias, clamped to [-100, 100]) | honoured | honoured | honoured | +| `min_tokens` | honoured: no EOS / stop token before this many generated tokens | honoured | `min_new_tokens` | no | +| `stop` (strings), `stop_token_ids` | honoured | honoured | honoured | `stop` only | +| `include_stop_str_in_output` | honoured | honoured | `no_stop_trim` | no | +| `skip_special_tokens` | honoured; default off here, the reasoning and tool parsers read the tags | honoured (default on) | honoured (default on) | n/a | +| `n` | honoured, 1..16, one generation per choice (the prefix cache shares the prompt) | honoured | honoured | `n_cmpl` | +| `stream`, `stream_options.include_usage` | honoured | honoured | honoured | honoured | +| `logprobs`, `top_logprobs` | 400 on this route (see #224); `/v1/completions` `logprobs` 400 | honoured | honoured | `n_probs` | +| `tools`, `tool_choice` (`auto` / `none` / `required` / named), `parallel_tool_calls` | honoured | honoured | honoured | honoured | +| `reasoning_effort`, `thinking`, `chat_template_kwargs` | honoured | `chat_template_kwargs` | honoured | honoured | +| `continue_final_message` | honoured (the last assistant message is continued, no generation prompt) | honoured | honoured | no | +| `response_format` `json_object` / `json_schema`, guided decoding | 400 (no constrained decoding) | honoured (xgrammar / guidance) | honoured (xgrammar / outlines / llguidance) | `grammar`, `json_schema` | +| `seed` | accepted, not honoured (one batched sampling kernel, no per-request generator) | honoured | `sampling_seed` | honoured | +| `user`, `metadata`, `store`, `service_tier` | accepted, ignored | accepted, ignored | accepted, ignored | accepted, ignored | +| `request_id` | honoured: becomes the response id | honoured | `rid` | `id_slot` (different meaning) | +| `function_call` (legacy) | 400 | 400 | 400 | no | +| `best_of`, `use_beam_search`, `length_penalty`, `prompt_logprobs`, `echo` (chat), `truncate_prompt_tokens`, `priority`, `cache_salt`, `lora` | no | honoured | partly | partly | +| `ignore_eos` | honoured (benchmarking) | honoured | honoured | honoured | + +## Legacy completions: request + +`prompt` as a string or a list of strings (token-id prompts are refused); the same sampling +parameters as above; plus: + +| parameter | FreeToken | vLLM | SGLang | llama.cpp | +|---|---|---|---|---| +| `echo` | honoured (the prompt leads the text; not with `logprobs`) | honoured | honoured | no | +| `suffix` | honoured on models with fill-in-the-middle tokens (`<|fim_prefix|>` ... `<|fim_middle|>`): the prompt becomes prefix-suffix-middle | 400 | 400 | `/infill` | +| `n` | honoured, prompt-major choice order (`prompt i`, sample `k` -> `index i*n+k`) | honoured | honoured | no | +| `logprobs` | 400 (see #224 for the sampled-token logprobs) | honoured | honoured | `n_probs` | + +## Responses + +| field | FreeToken | notes | +|---|---|---| +| `id` | `chatcmpl-` / `cmpl-`, or `-` when the client sent one | | +| `created`, `model`, `object` | yes | | +| `system_fingerprint` | `null` | vLLM and SGLang also send `null` | +| `choices[].message.reasoning_content` | yes, when the model reasoned | vLLM `reasoning_content`, llama.cpp `reasoning_content` | +| `choices[].message.tool_calls` | yes | | +| `choices[].finish_reason` | `stop`, `length`, `tool_calls` | | +| `usage.prompt_tokens`, `completion_tokens`, `total_tokens` | yes | for `n > 1` the prompt is counted once, completions summed | +| `usage.prompt_tokens_details.cached_tokens` | with `--enable-cache-report` | vLLM, SGLang: always | +| `usage.completion_tokens_details.reasoning_tokens` | yes, when the model emitted a reasoning end tag | vLLM, SGLang | +| streaming: one `[DONE]`, an optional final usage chunk | yes | | +| llama.cpp `timings` | no; `/v1/stats` and `/metrics` carry the server-side rates | | + +## Tokenize / detokenize + +``` +POST /tokenize {"prompt": "text", "add_special_tokens": true, "return_token_strs": false} +POST /tokenize {"messages": [...], "add_generation_prompt": true, "chat_template_kwargs": {}} + -> {"count": N, "max_model_len": L, "tokens": [...], "token_strs": [...]?} +POST /detokenize {"tokens": [...], "skip_special_tokens": false} -> {"prompt": "text"} +``` + +`messages` are rendered through the model's chat template exactly as a generation would +render them, so `count` equals the `prompt_tokens` that generation reports. + +## Not implemented, on purpose + +- Constrained decoding (`response_format` json / json_schema, regex, grammar): needs a + grammar engine on the sampling path. +- Prompt logprobs and `echo` with `logprobs`: need the prefill logits. +- Per-request `seed`: the batched sampling kernel draws from one generator. +- Embeddings, reranking, scoring, audio: not generation. diff --git a/python/freetoken/core.py b/python/freetoken/core.py index 45a763c0fe..a160beaf81 100644 --- a/python/freetoken/core.py +++ b/python/freetoken/core.py @@ -24,11 +24,45 @@ class SamplingParams: # Stop strings (OpenAI `stop` / Anthropic `stop_sequences`). Generation finishes when one # appears in the decoded output; the matched substring (and anything after) is trimmed. stop_strs: list[str] = field(default_factory=list) + # ---- logits processors (engine/sample.py), all off by default ---- + # min_p: drop tokens whose probability is below min_p x the top probability. + min_p: float = 0.0 + # OpenAI presence/frequency penalties over the generated tokens; repetition_penalty is the + # HF/vLLM multiplicative penalty over prompt + generated tokens (1.0 = off). + presence_penalty: float = 0.0 + frequency_penalty: float = 0.0 + repetition_penalty: float = 1.0 + # OpenAI logit_bias as [[token id, bias], ...] (clamped to [-100, 100] by the API). A list, + # not a dict: the messages cross process boundaries as msgpack, which refuses int map keys. + logit_bias: list[list[float]] | None = None + # No EOS / stop token before this many generated tokens (vLLM min_tokens, SGLang + # min_new_tokens). The scheduler fills min_tokens_stop_ids with the ids to mask. + min_tokens: int = 0 + min_tokens_stop_ids: list[int] = field(default_factory=list) + # Extra token ids that end generation with finish_reason "stop" (vLLM stop_token_ids). + stop_token_ids: list[int] = field(default_factory=list) + # Keep the matched stop string in the output instead of trimming it (vLLM + # include_stop_str_in_output, SGLang no_stop_trim). + include_stop_str_in_output: bool = False + # Decode with skip_special_tokens. Off by default here: the reasoning and tool parsers + # consume the /tool tags from the decoded text. + skip_special_tokens: bool = False @property def is_greedy(self) -> bool: return self.temperature <= 0.0 or self.top_k == 1 + @property + def needs_logits_processing(self) -> bool: + return bool( + self.min_p > 0.0 + or self.presence_penalty != 0.0 + or self.frequency_penalty != 0.0 + or self.repetition_penalty != 1.0 + or self.logit_bias + or self.min_tokens > 0 + ) + @dataclass(eq=False) class Req: diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py index 3bdae76957..24a677a5f7 100644 --- a/python/freetoken/engine/sample.py +++ b/python/freetoken/engine/sample.py @@ -10,18 +10,114 @@ from freetoken.core import Batch +@dataclass +class LogitsPlan: + """Inputs of the logits processors for the rows of one batch that asked for any + (presence / frequency / repetition penalties, logit_bias, min_tokens, min_p). + + Built on the host in ``Sampler.prepare`` from the requests' SamplingParams and + token histories, applied on the device by ``apply_logits_processors`` before the + sampling kernel. ``None`` when no request in the batch uses a processor, so the + default path stays exactly as it was. Row indices inside the plan are LOCAL (0..r-1, + in the order of ``rows``); ``rows`` maps them back to the batch. + """ + + rows: torch.Tensor # [r] int64 batch rows + temps: torch.Tensor # [r, 1] float32 temperatures (min_p works on probabilities) + presence: torch.Tensor | None = None # [r, 1] float32 + frequency: torch.Tensor | None = None # [r, 1] float32 + repetition: torch.Tensor | None = None # [r, 1] float32 (1.0 = off) + # Generated ids per row, right-padded with vocab_size (a scratch column that is + # dropped), for the presence / frequency counts and the repetition mask. + hist_ids: torch.Tensor | None = None # [r, L] int64 + # Prompt ids per row, same padding: repetition_penalty also covers the prompt. + prompt_ids: torch.Tensor | None = None # [r, P] int64 + bias_rows: torch.Tensor | None = None # [b] int64 local rows + bias_ids: torch.Tensor | None = None # [b] int64 + bias_vals: torch.Tensor | None = None # [b] float32 + # Rows still under min_tokens and the ids (EOS + stop_token_ids, padded) they must + # not sample yet. + min_rows: torch.Tensor | None = None # [m] int64 local rows + min_ids: torch.Tensor | None = None # [m, S] int64 + min_p: torch.Tensor | None = None # [r, 1] float32 (0 = off for that row) + + @dataclass class BatchSamplingArgs: temperatures: torch.Tensor | None top_k: torch.Tensor | None = None top_p: torch.Tensor | None = None greedy_mask: torch.Tensor | None = None + plan: LogitsPlan | None = None def make_device_tensor(data: List, dtype: torch.dtype, device: torch.device) -> torch.Tensor: + if device.type != "cuda": # the unit tests build plans on the CPU + return torch.tensor(data, dtype=dtype, device=device) return torch.tensor(data, dtype=dtype, pin_memory=True).to(device, non_blocking=True) +def _padded_rows(rows: list[list[int]], pad: int, device: torch.device) -> torch.Tensor: + """[len(rows), max_len] int64 with ``pad`` on the right. A synchronous copy: the + histories can be long (prompt + output) and a pinned buffer must not be freed under + an in-flight non_blocking copy.""" + width = max(1, max((len(r) for r in rows), default=1)) + out = torch.full((len(rows), width), pad, dtype=torch.int64) + for i, r in enumerate(rows): + if r: + out[i, : len(r)] = torch.as_tensor(r, dtype=torch.int64) + return out.to(device) + + +def apply_logits_processors( + logits: torch.Tensor, plan: LogitsPlan, vocab_size: int +) -> torch.Tensor: + """A float32 copy of ``logits`` with the plan applied to its rows, in vLLM's order: + repetition -> presence -> frequency penalties, logit_bias, the min_tokens mask, then + min_p. Rows outside the plan are copied unchanged.""" + out = logits.to(torch.float32, copy=True) + sub = out[plan.rows] # [r, V] (advanced indexing copies) + r = sub.shape[0] + device = sub.device + neg_inf = float("-inf") + + if plan.hist_ids is not None: + counts = torch.zeros((r, vocab_size + 1), dtype=torch.float32, device=device) + counts.scatter_add_(1, plan.hist_ids, torch.ones_like(plan.hist_ids, dtype=torch.float32)) + counts = counts[:, :vocab_size] + if plan.repetition is not None: + seen = counts > 0 + if plan.prompt_ids is not None: + in_prompt = torch.zeros((r, vocab_size + 1), dtype=torch.bool, device=device) + in_prompt.scatter_(1, plan.prompt_ids, torch.ones_like(plan.prompt_ids, dtype=torch.bool)) + seen = seen | in_prompt[:, :vocab_size] + # HF / vLLM semantics: divide positive logits, multiply negative ones. + penalized = torch.where(sub > 0, sub / plan.repetition, sub * plan.repetition) + sub = torch.where(seen, penalized, sub) + if plan.presence is not None: + sub = sub - plan.presence * (counts > 0).to(sub.dtype) + if plan.frequency is not None: + sub = sub - plan.frequency * counts + + if plan.bias_ids is not None: + sub.index_put_((plan.bias_rows, plan.bias_ids), plan.bias_vals, accumulate=True) + + if plan.min_ids is not None: + block = torch.zeros((plan.min_ids.shape[0], vocab_size + 1), dtype=torch.bool, device=device) + block.scatter_(1, plan.min_ids, torch.ones_like(plan.min_ids, dtype=torch.bool)) + rows = sub[plan.min_rows] + sub[plan.min_rows] = torch.where(block[:, :vocab_size], torch.full_like(rows, neg_inf), rows) + + if plan.min_p is not None: + probs = torch.softmax(sub / plan.temps, dim=-1) + keep = probs >= plan.min_p * probs.amax(dim=-1, keepdim=True) + keep = keep | (plan.min_p <= 0.0) # rows without min_p stay untouched + sub = torch.where(keep, sub, torch.full_like(sub, neg_inf)) + + out[plan.rows] = sub + return out + + def sample_impl( logits: torch.Tensor, temperatures: torch.Tensor, @@ -56,11 +152,88 @@ class Sampler: device: torch.device vocab_size: int + def _plan(self, batch: Batch, params) -> LogitsPlan | None: + """The LogitsPlan for this batch, or None when no request asks for a processor. + Token histories come from the host-side ``req.input_ids``; under overlap + scheduling the previous step's token may not be appended yet, so a penalty sees + the history one token late. That is the accepted cost of keeping the counts off + the device.""" + MIN_T = 1e-6 + rows = [i for i, p in enumerate(params) if p.needs_logits_processing] + if not rows: + return None + picked = [(batch.reqs[i], params[i]) for i in rows] + temps = [[max(0.0 if p.is_greedy else p.temperature, MIN_T)] for _, p in picked] + plan = LogitsPlan( + rows=make_device_tensor(rows, torch.int64, self.device), + temps=make_device_tensor(temps, torch.float32, self.device), + ) + pad = self.vocab_size + + want_hist = any( + p.presence_penalty != 0.0 or p.frequency_penalty != 0.0 or p.repetition_penalty != 1.0 + for _, p in picked + ) + if want_hist: + hist: list[list[int]] = [] + prompts: list[list[int]] = [] + want_prompt = False + for req, p in picked: + prompt_len = req.max_device_len - req.output_len + ids = req.input_ids + penalized = p.presence_penalty != 0.0 or p.frequency_penalty != 0.0 or p.repetition_penalty != 1.0 + hist.append(ids[prompt_len:].tolist() if penalized else []) + if p.repetition_penalty != 1.0: + want_prompt = True + prompts.append(ids[:prompt_len].tolist()) + else: + prompts.append([]) + plan.hist_ids = _padded_rows(hist, pad, self.device) + if want_prompt: + plan.prompt_ids = _padded_rows(prompts, pad, self.device) + if any(p.presence_penalty != 0.0 for _, p in picked): + plan.presence = make_device_tensor([[p.presence_penalty] for _, p in picked], torch.float32, self.device) + if any(p.frequency_penalty != 0.0 for _, p in picked): + plan.frequency = make_device_tensor([[p.frequency_penalty] for _, p in picked], torch.float32, self.device) + if any(p.repetition_penalty != 1.0 for _, p in picked): + plan.repetition = make_device_tensor([[p.repetition_penalty] for _, p in picked], torch.float32, self.device) + + bias_rows: list[int] = [] + bias_ids: list[int] = [] + bias_vals: list[float] = [] + for local, (_, p) in enumerate(picked): + for tid, val in p.logit_bias or (): + if 0 <= int(tid) < self.vocab_size: + bias_rows.append(local) + bias_ids.append(int(tid)) + bias_vals.append(float(val)) + if bias_ids: + plan.bias_rows = make_device_tensor(bias_rows, torch.int64, self.device) + plan.bias_ids = make_device_tensor(bias_ids, torch.int64, self.device) + plan.bias_vals = make_device_tensor(bias_vals, torch.float32, self.device) + + min_rows: list[int] = [] + min_ids: list[list[int]] = [] + for local, (req, p) in enumerate(picked): + if p.min_tokens > 0 and p.min_tokens_stop_ids: + generated = len(req.input_ids) - (req.max_device_len - req.output_len) + if generated < p.min_tokens: + min_rows.append(local) + min_ids.append([t for t in p.min_tokens_stop_ids if 0 <= t < self.vocab_size]) + if min_rows: + plan.min_rows = make_device_tensor(min_rows, torch.int64, self.device) + plan.min_ids = _padded_rows(min_ids, pad, self.device) + + if any(p.min_p > 0.0 for _, p in picked): + plan.min_p = make_device_tensor([[p.min_p] for _, p in picked], torch.float32, self.device) + return plan + def prepare(self, batch: Batch) -> BatchSamplingArgs: params = [r.sampling_params for r in batch.reqs] is_greedy = [p.is_greedy for p in params] + plan = self._plan(batch, params) if all(is_greedy): - return BatchSamplingArgs(temperatures=None) + return BatchSamplingArgs(temperatures=None, plan=plan) MIN_P = MIN_T = 1e-6 # Greedy outputs are selected explicitly in sample(); use neutral sampling @@ -83,11 +256,15 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs: greedy_mask = ( make_device_tensor(is_greedy, torch.bool, self.device) if any(is_greedy) else None ) - return BatchSamplingArgs(temperatures, top_k=top_k, top_p=top_p, greedy_mask=greedy_mask) + return BatchSamplingArgs( + temperatures, top_k=top_k, top_p=top_p, greedy_mask=greedy_mask, plan=plan + ) @nvtx_annotate("Sampler") def sample(self, logits: torch.Tensor, args: BatchSamplingArgs) -> torch.Tensor: with torch.cuda.nvtx.range("Sampler"): + if args.plan is not None: + logits = apply_logits_processors(logits, args.plan, self.vocab_size) if args.temperatures is None: # greedy sampling return torch.argmax(logits, dim=-1) tokens = sample_impl(logits.float(), args.temperatures, args.top_k, args.top_p) diff --git a/python/freetoken/message/frontend.py b/python/freetoken/message/frontend.py index 24725567b9..f3200650d8 100644 --- a/python/freetoken/message/frontend.py +++ b/python/freetoken/message/frontend.py @@ -56,6 +56,9 @@ class UserReply(BaseFrontendMsg): finish_reason: str | None = None # The stop string that ended generation (Anthropic reports it as stop_reason='stop_sequence'). matched_stop: str | None = None + # On the finished reply: tokens generated up to and including the reasoning end tag + # (0 when the model emitted none) -> usage.completion_tokens_details.reasoning_tokens. + reasoning_tokens: int = 0 @dataclass diff --git a/python/freetoken/message/tokenizer.py b/python/freetoken/message/tokenizer.py index d52d658e63..0d2116038f 100644 --- a/python/freetoken/message/tokenizer.py +++ b/python/freetoken/message/tokenizer.py @@ -36,6 +36,10 @@ class DetokenizeMsg(BaseTokenizerMsg): # The request's stop strings (None when it has none), so the detokenizer can hold back # a trailing partial-stop prefix instead of streaming it and then needing to retract. stop_strs: list[str] | None = None + # Keep the matched stop string in the output (SamplingParams.include_stop_str_in_output). + keep_stop_str: bool = False + # Decode this request with skip_special_tokens (SamplingParams.skip_special_tokens). + skip_special_tokens: bool = False # KV page-pool usage snapshot at this step (not-evictable used/total), passed # through to the frontend for the shell status bar. 0/0 for owned-KV models. kv_used_pages: int = 0 diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index d057d3b908..f508a9dc89 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -353,7 +353,7 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: hit_length = req.input_ids.numel() >= req.max_device_len hit_eos = ( not req.sampling_params.ignore_eos and next_token in self.eos_token_ids - ) + ) or next_token in req.sampling_params.stop_token_ids matched_stop = ( self._match_stop_str(req) if not hit_eos and req.sampling_params.stop_strs @@ -379,6 +379,8 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: finish_reason=finish_reason, matched_stop=matched_stop, stop_strs=req.sampling_params.stop_strs or None, + keep_stop_str=req.sampling_params.include_stop_str_in_output, + skip_special_tokens=req.sampling_params.skip_special_tokens, ) ) @@ -542,6 +544,11 @@ def _process_one_msg(self, msg: BaseBackendMsg) -> None: logger.warning_rank0( f"Adjust max_tokens to {max_output_len} for request {msg.uid}." ) + sp = msg.sampling_params + if sp.min_tokens > 0: + # min_tokens: the sampler masks these ids while the output is shorter. + sp.min_tokens = min(sp.min_tokens, sp.max_tokens) + sp.min_tokens_stop_ids = sorted(set(self.eos_token_ids) | set(sp.stop_token_ids)) self.prefill_manager.add_one_req(msg) elif isinstance(msg, AbortBackendMsg): logger.debug_rank0("Aborting request %d", msg.uid) diff --git a/python/freetoken/server/api_models.py b/python/freetoken/server/api_models.py index ffd7172802..579443313d 100644 --- a/python/freetoken/server/api_models.py +++ b/python/freetoken/server/api_models.py @@ -90,11 +90,30 @@ class ChatCompletionRequest(BaseModel): function_call: Any | None = None logit_bias: dict[str, float] | None = None response_format: dict[str, Any] | None = None + # The last message is an assistant prefix the model continues (no generation prompt). + continue_final_message: bool = False + # ---- sampling extras shared with vLLM / SGLang / llama.cpp (all optional) ---- + min_p: float | None = None + repetition_penalty: float | None = None + min_tokens: int = 0 + stop_token_ids: list[int] | None = None + include_stop_str_in_output: bool = False + skip_special_tokens: bool | None = None + # Accepted for OpenAI compatibility, not honoured: the batched sampling kernel has no + # per-request generator (seed); user / metadata / store are opaque to the server. + seed: int | None = None + user: str | None = None + # A client-chosen id echoed as the response id (vLLM request_id / SGLang rid). + request_id: str | None = None @model_validator(mode="after") def _sync_max_completion_tokens(self) -> "ChatCompletionRequest": if self.max_completion_tokens is not None: self.max_tokens = self.max_completion_tokens + if self.continue_final_message and ( + not self.messages or self.messages[-1].role != "assistant" + ): + raise ValueError("continue_final_message needs a final assistant message") return self @@ -120,6 +139,19 @@ class CompletionRequest(BaseModel): suffix: str | None = None logit_bias: dict[str, float] | None = None response_format: dict[str, Any] | None = None + # ---- sampling extras shared with vLLM / SGLang / llama.cpp (all optional) ---- + min_p: float | None = None + repetition_penalty: float | None = None + min_tokens: int = 0 + stop_token_ids: list[int] | None = None + include_stop_str_in_output: bool = False + skip_special_tokens: bool | None = None + # Accepted for OpenAI compatibility, not honoured: the batched sampling kernel has no + # per-request generator (seed); user / metadata / store are opaque to the server. + seed: int | None = None + user: str | None = None + # A client-chosen id echoed as the response id (vLLM request_id / SGLang rid). + request_id: str | None = None @model_validator(mode="after") def _sync_max_completion_tokens(self) -> "CompletionRequest": @@ -128,6 +160,32 @@ def _sync_max_completion_tokens(self) -> "CompletionRequest": return self +MAX_N = 16 # choices per request; each is a full generation, so keep the fan-out bounded + + +class TokenizeRequest(BaseModel): + """``POST /tokenize`` (the vLLM / SGLang shape): a raw prompt, or chat messages that + are rendered through the model's chat template first.""" + + model_config = ConfigDict(extra="allow") + + model: str | None = None + prompt: str | None = None + messages: list[Message] | None = None + add_special_tokens: bool = True + add_generation_prompt: bool = True + return_token_strs: bool = False + chat_template_kwargs: dict[str, Any] = Field(default_factory=dict) + + +class DetokenizeRequest(BaseModel): + model_config = ConfigDict(extra="allow") + + model: str | None = None + tokens: list[int] + skip_special_tokens: bool = False + + class ModelCard(BaseModel): id: str object: str = "model" diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index 34221cfc70..fea281a815 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -358,16 +358,20 @@ async def stream_generate(self, uid: int): yield b"data: [DONE]\n\n" logger.debug("Finished streaming response for user %s", uid) - async def stream_with_cancellation(self, generator, request: Request, uid: int): + async def stream_with_cancellation(self, generator, request: Request, uid): + """``uid`` is one request id or a sequence of them (an ``n > 1`` fan-out streams + several generations into one response): a disconnect aborts every one.""" + uids = list(uid) if isinstance(uid, (list, tuple)) else [uid] try: async for chunk in generator: # detect if the client has disconnected if await request.is_disconnected(): - logger.info("Client disconnected for user %s", uid) + logger.info("Client disconnected for user %s", uids) raise asyncio.CancelledError yield chunk except asyncio.CancelledError: - asyncio.create_task(self.abort_user(uid)) + for one in uids: + asyncio.create_task(self.abort_user(one)) raise async def abort_user(self, uid: int): diff --git a/python/freetoken/server/control_api.py b/python/freetoken/server/control_api.py index 7158e4e5fa..cef632306e 100644 --- a/python/freetoken/server/control_api.py +++ b/python/freetoken/server/control_api.py @@ -68,6 +68,24 @@ async def list_requests(since: int = 0, limit: int = 100): from .stats import build_stats + @app.get("/metrics") + async def metrics(): + """Prometheus text exposition of /v1/stats (the llama.cpp / vLLM convention).""" + from fastapi.responses import PlainTextResponse + + from .stats import prometheus_text + + doc = build_stats( + get_state(), request_ring.requests_p95_ms(), request_ring.requests_ttft_mean_ms() + ) + return PlainTextResponse(prometheus_text(doc), media_type="text/plain; version=0.0.4") + + @app.get("/version") + async def version(): + from freetoken import __version__ + + return {"version": __version__} + @app.get("/v1/stats") async def stats(): doc = build_stats( diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index ac988eddc7..c7d658688c 100644 --- a/python/freetoken/server/generation.py +++ b/python/freetoken/server/generation.py @@ -115,6 +115,7 @@ class GenDone: completion_tokens: int matched_stop: str | None = None cached_tokens: int = 0 + reasoning_tokens: int = 0 GenEvent = ReasoningDelta | ContentDelta | ToolCallStart | ToolCallArgsDelta | ToolCallsDelta | GenDone @@ -130,6 +131,7 @@ class GenResult: completion_tokens: int matched_stop: str | None = None cached_tokens: int = 0 + reasoning_tokens: int = 0 @dataclass @@ -169,9 +171,19 @@ def resolve_sampling( model_sampling: dict[str, Any], stop: str | list[str] | None = None, default_max_tokens: int = DEFAULT_MAX_OUTPUT_TOKENS, + min_p: float | None = None, + presence_penalty: float = 0.0, + frequency_penalty: float = 0.0, + repetition_penalty: float | None = None, + logit_bias: dict[str, float] | dict[int, float] | None = None, + min_tokens: int = 0, + stop_token_ids: list[int] | None = None, + include_stop_str_in_output: bool = False, + skip_special_tokens: bool | None = None, ) -> SamplingParams: """Map a protocol's sampling fields onto the engine's neutral SamplingParams, - filling unspecified fields from the checkpoint's recommended defaults.""" + filling unspecified fields from the checkpoint's recommended defaults. Range errors + raise ValueError (the adapters answer 400).""" def pick(value, key, framework): return value if value is not None else model_sampling.get(key, framework) @@ -183,14 +195,50 @@ def pick(value, key, framework): # non-positive value is a client error. if max_tokens is not None and max_tokens < 1: raise ValueError(f"max_tokens must be at least 1, got {max_tokens}") - return SamplingParams( + resolved_min_p = float(pick(min_p, "min_p", 0.0)) + resolved_rep = float(pick(repetition_penalty, "repetition_penalty", 1.0)) + if not 0.0 <= resolved_min_p <= 1.0: + raise ValueError(f"min_p must be in [0, 1], got {resolved_min_p}") + for name, value in (("presence_penalty", presence_penalty), ("frequency_penalty", frequency_penalty)): + if not -2.0 <= float(value) <= 2.0: + raise ValueError(f"{name} must be in [-2, 2], got {value}") + if resolved_rep <= 0.0: + raise ValueError(f"repetition_penalty must be positive, got {resolved_rep}") + if min_tokens < 0: + raise ValueError(f"min_tokens must be >= 0, got {min_tokens}") + bias: list[list[float]] | None = None + if logit_bias: + bias = [] + for key, value in logit_bias.items(): + try: + tid = int(key) + except (TypeError, ValueError): + raise ValueError(f"logit_bias keys must be token ids, got {key!r}") from None + if tid < 0: + raise ValueError(f"logit_bias token id must be >= 0, got {tid}") + bias.append([tid, max(-100.0, min(100.0, float(value)))]) # the OpenAI range + ids = [int(t) for t in (stop_token_ids or [])] + if any(t < 0 for t in ids): + raise ValueError("stop_token_ids must be non-negative token ids") + params = SamplingParams( ignore_eos=ignore_eos, max_tokens=default_max_tokens if max_tokens is None else max_tokens, temperature=pick(temperature, "temperature", 0.0), top_k=pick(top_k, "top_k", -1), top_p=pick(top_p, "top_p", 1.0), stop_strs=[s for s in stop_list if s], # drop empty strings (would match everything) + min_p=resolved_min_p, + presence_penalty=float(presence_penalty), + frequency_penalty=float(frequency_penalty), + repetition_penalty=resolved_rep, + logit_bias=bias, + min_tokens=int(min_tokens), + stop_token_ids=ids, + include_stop_str_in_output=bool(include_stop_str_in_output), ) + if skip_special_tokens is not None: + params.skip_special_tokens = bool(skip_special_tokens) + return params def render_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: @@ -603,6 +651,7 @@ async def _generate_events_impl(uid: int, spec: GenSpec, state: Any) -> AsyncIte prompt_tokens = 0 completion_tokens = 0 cached_tokens = 0 + reasoning_tokens = 0 pending = "" parse_tools = spec.parse_tools reasoning_parser = _make_reasoning_parser(spec, state) @@ -698,6 +747,7 @@ def _route_tool_text(piece: str) -> list[GenEvent]: prompt_tokens += ack.prompt_tokens_delta completion_tokens += ack.completion_tokens_delta cached_tokens += ack.cached_tokens + reasoning_tokens = max(reasoning_tokens, getattr(ack, "reasoning_tokens", 0)) content_delta = ack.incremental_output if reasoning_parser is not None and content_delta: reasoning_delta, content_delta = reasoning_parser.parse_stream_chunk(content_delta) @@ -780,6 +830,7 @@ def _route_tool_text(piece: str) -> list[GenEvent]: yield GenDone( finish_reason, prompt_tokens, completion_tokens, matched_stop=engine_matched_stop, cached_tokens=cached_tokens, + reasoning_tokens=reasoning_tokens, ) @@ -790,6 +841,7 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any) -> GenResult: prompt_tokens = 0 completion_tokens = 0 cached_tokens = 0 + reasoning_tokens = 0 engine_finish_reason: str | None = None engine_matched_stop: str | None = None async for ack in state.wait_for_ack(uid): @@ -798,6 +850,7 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any) -> GenResult: prompt_tokens += ack.prompt_tokens_delta completion_tokens += ack.completion_tokens_delta cached_tokens += ack.cached_tokens + reasoning_tokens = max(reasoning_tokens, getattr(ack, "reasoning_tokens", 0)) full_content += ack.incremental_output if ack.finished: engine_finish_reason = getattr(ack, "finish_reason", None) @@ -824,4 +877,5 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any) -> GenResult: completion_tokens=completion_tokens, matched_stop=engine_matched_stop, cached_tokens=cached_tokens, + reasoning_tokens=reasoning_tokens, ) diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index dc2f73a979..f1bb0888f0 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -14,6 +14,9 @@ from freetoken.tokenizer.effort import EFFORT_SCALE, KNOWN_REASONING_EFFORTS from .api_models import ( + MAX_N, + DetokenizeRequest, + TokenizeRequest, ChatCompletionRequest, CompletionRequest, ModelCard, @@ -68,18 +71,11 @@ def chat_request_to_genspec( thinking_type = _thinking_type(req) if req.reasoning_effort or thinking_type: ctk = effort_toggle_kwargs(req.reasoning_effort, ctk, thinking_type=thinking_type) + if req.continue_final_message: + ctk = {**ctk, "continue_final_message": True} return GenSpec( messages=render_messages([m.model_dump(exclude_none=True) for m in req.messages]), - sampling_params=resolve_sampling( - temperature=req.temperature, - top_k=req.top_k, - top_p=req.top_p, - max_tokens=req.max_tokens, - ignore_eos=req.ignore_eos, - model_sampling=model_sampling, - stop=req.stop, - default_max_tokens=default_max_tokens, - ), + sampling_params=_resolve_sampling(req, model_sampling, default_max_tokens=default_max_tokens), chat_template_kwargs=ctk, template_tools=_tools_for_template(req), parser_tools=(_all_tool_dicts(req.tools) if _should_parse_tools(req) else None), @@ -131,6 +127,16 @@ async def v1_completions(req: CompletionRequest, request: Request): return gate return await handle_completion(req, request, state, get_model_sampling()) + @app.post("/tokenize") + @app.post("/v1/tokenize") + async def tokenize(req: TokenizeRequest): + return await handle_tokenize(req, get_state()) + + @app.post("/detokenize") + @app.post("/v1/detokenize") + async def detokenize(req: DetokenizeRequest): + return await handle_detokenize(req, get_state()) + @app.get("/v1/models") async def v1_models(): state = get_state() @@ -155,15 +161,13 @@ async def handle_chat_completion( ): if req.function_call is not None: return create_error_response("function_call is not supported; use tools/tool_choice instead") - if req.logit_bias is not None: - return create_error_response("logit_bias is not supported") if _response_format_unsupported(req.response_format): return create_error_response( "response_format json_object/json_schema is not supported (no constrained decoding)", param="response_format", ) - if req.n != 1: - return create_error_response("Only n=1 is supported", param="n") + if not 1 <= req.n <= MAX_N: + return create_error_response(f"n must be between 1 and {MAX_N}", param="n") # Case/whitespace and the "off" disable synonym stay accepted here because # effort_toggle_kwargs normalizes and honors them downstream. effort = req.reasoning_effort.strip().lower() if isinstance(req.reasoning_effort, str) else None @@ -196,43 +200,65 @@ async def handle_chat_completion( if err is not None: return create_error_response(str(err), code=err.code) + # n > 1: one generation per choice, submitted together (the prefix cache serves the + # shared prompt after the first prefill). + uids: list[int] = [] try: - uid = await submit_generation(spec, state) + for _ in range(req.n): + uids.append(await submit_generation(spec, state)) except GenerationError as exc: + for u in uids: # choices already queued before a later one failed + await state.abort_user(u) return create_error_response(str(exc), code=exc.code) + uid = uids[0] if req.stream: - chunks = stream_chat_completion_chunks(uid, req, state, spec) + if req.n == 1: + chunks = stream_chat_completion_chunks(uid, req, state, spec) + else: + chunks = _merge_streams( + [ + stream_chat_completion_chunks(u, req, state, spec, index=i, terminal=False) + for i, u in enumerate(uids) + ], + lambda: _chat_chunk(req, uid, []), + include_usage=bool(req.stream_options and req.stream_options.include_usage), + state=state, + ) if request is not None: - chunks = state.stream_with_cancellation(chunks, request, uid) + chunks = state.stream_with_cancellation(chunks, request, uids) return StreamingResponse(chunks, media_type="text/event-stream") try: - result = await generate_full(uid, spec, state, source="/v1/chat/completions") + results = await asyncio.gather( + *(generate_full(u, spec, state, source="/v1/chat/completions") for u in uids) + ) except GenerationError as exc: + for u in uids: + await state.abort_user(u) return create_error_response(str(exc), code=exc.code) - message: dict[str, Any] = {"role": "assistant", "content": result.content} - if result.reasoning: - message["reasoning_content"] = result.reasoning - if result.tool_calls: - message["tool_calls"] = _tool_calls_to_openai(result.tool_calls) - + choices: list[dict[str, Any]] = [] + for index, result in enumerate(results): + message: dict[str, Any] = {"role": "assistant", "content": result.content} + if result.reasoning: + message["reasoning_content"] = result.reasoning + if result.tool_calls: + message["tool_calls"] = _tool_calls_to_openai(result.tool_calls) + choices.append({"index": index, "message": message, "finish_reason": result.finish_reason}) + + first = results[0] return { - "id": f"chatcmpl-{uid}", + "id": _response_id("chatcmpl", req, uid), "object": "chat.completion", "created": int(time.time()), "model": req.model, - "choices": [ - { - "index": 0, - "message": message, - "finish_reason": result.finish_reason, - } - ], + "system_fingerprint": None, + "choices": choices, "usage": _usage( - result.prompt_tokens, - result.completion_tokens, - _reported_cached(state, result.cached_tokens), + first.prompt_tokens, + sum(r.completion_tokens for r in results), + _reported_cached(state, max(r.cached_tokens for r in results)), + reasoning_tokens=sum(r.reasoning_tokens for r in results), ), } @@ -242,21 +268,28 @@ async def stream_chat_completion_chunks( req: ChatCompletionRequest, state: Any, spec: GenSpec | None = None, + index: int = 0, + terminal: bool = True, ) -> AsyncIterator[bytes]: - """Format generate_events() into the OpenAI chat.completion.chunk SSE stream.""" + """Format generate_events() into the OpenAI chat.completion.chunk SSE stream. + + ``index`` is the choice index (``n > 1`` runs one of these per choice); with + ``terminal=False`` the usage chunk and ``[DONE]`` are left to the merger, which reads + this stream's usage from the ``_StreamUsage`` yielded last.""" if spec is None: spec = chat_request_to_genspec(req, {}) yield _sse( _chat_chunk( req, uid, - [{"delta": {"role": "assistant", "content": ""}, "index": 0, "finish_reason": None}], + [{"delta": {"role": "assistant", "content": ""}, "index": index, "finish_reason": None}], ) ) prompt_tokens = 0 completion_tokens = 0 cached_tokens = 0 + reasoning_tokens = 0 tool_calls_sent = 0 open_tool: dict[str, Any] | None = None events = generate_events(uid, spec, state, source="/v1/chat/completions") @@ -277,7 +310,7 @@ async def stream_chat_completion_chunks( _chat_chunk( req, uid, - [{"delta": {"reasoning_content": ev.text}, "index": 0, "finish_reason": None}], + [{"delta": {"reasoning_content": ev.text}, "index": index, "finish_reason": None}], ) ) elif isinstance(ev, ContentDelta): @@ -285,7 +318,7 @@ async def stream_chat_completion_chunks( _chat_chunk( req, uid, - [{"delta": {"content": ev.text}, "index": 0, "finish_reason": None}], + [{"delta": {"content": ev.text}, "index": index, "finish_reason": None}], ) ) elif isinstance(ev, ToolCallStart): @@ -305,7 +338,7 @@ async def stream_chat_completion_chunks( "type": "function", "function": {"name": ev.name, "arguments": ""}, }]}, - "index": 0, "finish_reason": None, + "index": index, "finish_reason": None, }], ) ) @@ -322,7 +355,7 @@ async def stream_chat_completion_chunks( "index": open_tool["index"], "function": {"arguments": ev.fragment}, }]}, - "index": 0, "finish_reason": None, + "index": index, "finish_reason": None, }], ) ) @@ -346,7 +379,7 @@ async def stream_chat_completion_chunks( "index": open_tool["index"], "function": {"arguments": remainder}, }]}, - "index": 0, "finish_reason": None, + "index": index, "finish_reason": None, }], ) ) @@ -358,7 +391,7 @@ async def stream_chat_completion_chunks( yield _sse( _chat_chunk( req, uid, - [{"delta": {"tool_calls": [delta]}, "index": 0, "finish_reason": None}], + [{"delta": {"tool_calls": [delta]}, "index": index, "finish_reason": None}], ) ) tool_calls_sent += 1 @@ -366,21 +399,18 @@ async def stream_chat_completion_chunks( prompt_tokens = ev.prompt_tokens completion_tokens = ev.completion_tokens cached_tokens = ev.cached_tokens - yield _sse(_chat_chunk(req, uid, [{"delta": {}, "index": 0, "finish_reason": ev.finish_reason}])) + reasoning_tokens = ev.reasoning_tokens + yield _sse(_chat_chunk(req, uid, [{"delta": {}, "index": index, "finish_reason": ev.finish_reason}])) + usage = _usage( + prompt_tokens, completion_tokens, _reported_cached(state, cached_tokens), + reasoning_tokens=reasoning_tokens, + ) + if not terminal: + yield _StreamUsage(usage) + return if req.stream_options and req.stream_options.include_usage: - yield _sse( - { - "id": f"chatcmpl-{uid}", - "object": "chat.completion.chunk", - "created": int(time.time()), - "model": req.model, - "choices": [], - "usage": _usage( - prompt_tokens, completion_tokens, _reported_cached(state, cached_tokens) - ), - } - ) + yield _sse({**_chat_chunk(req, uid, []), "usage": usage}) yield b"data: [DONE]\n\n" @@ -394,47 +424,63 @@ async def handle_completion( unsupported = _completion_unsupported_reason(req) if unsupported is not None: return create_error_response(unsupported) - try: # surfaces an out-of-range max_tokens as a 400 rather than a 500 from the worker + if not 1 <= req.n <= MAX_N: + return create_error_response(f"n must be between 1 and {MAX_N}", param="n") + try: # surfaces an out-of-range value as a 400 rather than a 500 from the worker default_max_tokens = ( getattr(state.config, "max_output_tokens", None) or DEFAULT_MAX_OUTPUT_TOKENS ) - _resolve_sampling(req, model_sampling, default_max_tokens=default_max_tokens) + sampling = _resolve_sampling(req, model_sampling, default_max_tokens=default_max_tokens) except ValueError as exc: - return create_error_response(str(exc), param="max_tokens") + return create_error_response(str(exc)) prompts = [req.prompt] if isinstance(req.prompt, str) else req.prompt assert isinstance(prompts, list) + if req.suffix is not None: + markers = await _fim_markers(state) + if markers is None: + return create_error_response( + "suffix needs a model with fill-in-the-middle tokens " + "(<|fim_prefix|> / <|fim_suffix|> / <|fim_middle|>)", + param="suffix", + ) + prompts = [_fim_prompt(p, req.suffix, markers) for p in prompts] + echo_texts = list(prompts) if req.echo else [""] * len(prompts) + if req.suffix is not None and req.echo: + # echo returns what the client sent, not the FIM-wrapped prompt + echo_texts = [req.prompt] if isinstance(req.prompt, str) else list(req.prompt) + + # choices are ordered prompt-major: prompt i, sample k -> index i * n + k (OpenAI) + uids: list[int] = [] + for prompt in prompts: + for _ in range(req.n): + uid = state.new_user() + await state.send_one(TokenizeMsg(uid=uid, text=prompt, sampling_params=sampling)) + uids.append(uid) + if req.stream: - if len(prompts) != 1: - return create_error_response("Streaming completions only support a single text prompt") - uid = state.new_user() - await state.send_one( - TokenizeMsg(uid=uid, text=prompts[0], sampling_params=_resolve_sampling( - req, model_sampling, default_max_tokens=default_max_tokens - )) - ) - chunks = stream_completion_chunks(uid, req, state) + if len(uids) == 1: + chunks = stream_completion_chunks(uids[0], req, state, echo_text=echo_texts[0]) + else: + chunks = _merge_streams( + [ + stream_completion_chunks( + u, req, state, index=i, terminal=False, echo_text=echo_texts[i // req.n] + ) + for i, u in enumerate(uids) + ], + lambda: _completion_chunk(req, uids[0], []), + include_usage=bool(req.stream_options and req.stream_options.include_usage), + state=state, + ) if request is not None: - chunks = state.stream_with_cancellation(chunks, request, uid) + chunks = state.stream_with_cancellation(chunks, request, uids) return StreamingResponse(chunks, media_type="text/event-stream") - choices: list[dict[str, Any]] = [] - prompt_tokens = 0 - completion_tokens = 0 - cached_tokens = 0 - for index, prompt in enumerate(prompts): - uid = state.new_user() - await state.send_one( - TokenizeMsg( - uid=uid, - text=prompt, - sampling_params=_resolve_sampling( - req, model_sampling, default_max_tokens=default_max_tokens - ), - ) - ) + async def _drain_one(uid: int, index: int) -> dict[str, Any] | JSONResponse: text = "" finish_reason = "stop" + prompt_tokens = completion_tokens = cached_tokens = 0 async for ack in state.wait_for_ack(uid): if getattr(ack, "error", None): return create_error_response(ack.error) @@ -445,77 +491,225 @@ async def handle_completion( if ack.finished: finish_reason = getattr(ack, "finish_reason", None) or "stop" break - choices.append({"index": index, "text": text, "finish_reason": finish_reason, "logprobs": None}) + return { + "index": index, + "text": echo_texts[index // req.n] + text, + "finish_reason": finish_reason, + "logprobs": None, + "_usage": (prompt_tokens, completion_tokens, cached_tokens), + } + + drained = await asyncio.gather(*(_drain_one(u, i) for i, u in enumerate(uids))) + choices: list[dict[str, Any]] = [] + prompt_tokens = completion_tokens = cached_tokens = 0 + for choice in drained: + if isinstance(choice, JSONResponse): + return choice + p, c, cached = choice.pop("_usage") + # the prompt is counted once per prompt, not once per sample of it + if choice["index"] % req.n == 0: + prompt_tokens += p + completion_tokens += c + cached_tokens = max(cached_tokens, cached) + choices.append(choice) return { - "id": f"cmpl-{uuid.uuid4().hex}", + "id": _response_id("cmpl", req, uuid.uuid4().hex), "object": "text_completion", "created": int(time.time()), "model": req.model, + "system_fingerprint": None, "choices": choices, "usage": _usage(prompt_tokens, completion_tokens, _reported_cached(state, cached_tokens)), } -async def stream_completion_chunks(uid: int, req: CompletionRequest, state: Any) -> AsyncIterator[bytes]: +def _completion_chunk(req: CompletionRequest, uid: Any, choices: list[dict[str, Any]]) -> dict[str, Any]: + return { + "id": _response_id("cmpl", req, uid), + "object": "text_completion.chunk", + "created": int(time.time()), + "model": req.model, + "choices": choices, + } + + +async def stream_completion_chunks( + uid: int, + req: CompletionRequest, + state: Any, + index: int = 0, + terminal: bool = True, + echo_text: str = "", +) -> AsyncIterator[bytes]: prompt_tokens = 0 completion_tokens = 0 cached_tokens = 0 finish_reason = "stop" + if echo_text: + # echo: the prompt leads the stream (vLLM sends it in the first chunk) + yield _sse(_completion_chunk(req, uid, [{"text": echo_text, "index": index, "finish_reason": None, "logprobs": None}])) async for ack in state.wait_for_ack(uid): if getattr(ack, "error", None): yield _sse({"error": {"message": ack.error, "type": "invalid_request_error", "code": None}}) - yield b"data: [DONE]\n\n" + if terminal: + yield b"data: [DONE]\n\n" return prompt_tokens += ack.prompt_tokens_delta completion_tokens += ack.completion_tokens_delta cached_tokens += ack.cached_tokens if ack.incremental_output: yield _sse( - { - "id": f"cmpl-{uid}", - "object": "text_completion.chunk", - "created": int(time.time()), - "model": req.model, - "choices": [ - { - "text": ack.incremental_output, - "index": 0, - "finish_reason": None, - "logprobs": None, - } - ], - } + _completion_chunk( + req, uid, + [{"text": ack.incremental_output, "index": index, "finish_reason": None, "logprobs": None}], + ) ) if ack.finished: finish_reason = getattr(ack, "finish_reason", None) or "stop" break - yield _sse( - { - "id": f"cmpl-{uid}", - "object": "text_completion.chunk", - "created": int(time.time()), - "model": req.model, - "choices": [{"text": "", "index": 0, "finish_reason": finish_reason, "logprobs": None}], - } - ) + yield _sse(_completion_chunk(req, uid, [{"text": "", "index": index, "finish_reason": finish_reason, "logprobs": None}])) + usage = _usage(prompt_tokens, completion_tokens, _reported_cached(state, cached_tokens)) + if not terminal: + yield _StreamUsage(usage) + return if req.stream_options and req.stream_options.include_usage: - yield _sse( - { - "id": f"cmpl-{uid}", - "object": "text_completion.chunk", - "created": int(time.time()), - "model": req.model, - "choices": [], - "usage": _usage( - prompt_tokens, completion_tokens, _reported_cached(state, cached_tokens) - ), - } + yield _sse({**_completion_chunk(req, uid, []), "usage": usage}) + yield b"data: [DONE]\n\n" + + +class _StreamUsage: + """Yielded last by a non-terminal sub-stream: its usage dict for the merger.""" + + __slots__ = ("usage",) + + def __init__(self, usage: dict[str, Any]) -> None: + self.usage = usage + + +async def _merge_streams( + streams: list[AsyncIterator[Any]], + envelope: Callable[[], dict[str, Any]], + *, + include_usage: bool, + state: Any, +) -> AsyncIterator[bytes]: + """Interleave the SSE chunks of several choice streams into one response (``n > 1``), + then one usage chunk (prompt counted once, completions summed) and one ``[DONE]``.""" + queue: asyncio.Queue[tuple[int, Any]] = asyncio.Queue() + usages: dict[int, dict[str, Any]] = {} + + async def pump(slot: int, stream: AsyncIterator[Any]) -> None: + try: + async for item in stream: + if isinstance(item, _StreamUsage): + usages[slot] = item.usage + else: + await queue.put((slot, item)) + finally: + await queue.put((slot, None)) + + tasks = [asyncio.create_task(pump(i, st)) for i, st in enumerate(streams)] + try: + open_slots = len(streams) + while open_slots: + slot, item = await queue.get() + if item is None: + open_slots -= 1 + continue + yield item + finally: + for task in tasks: + task.cancel() + if include_usage and usages: + first = usages[min(usages)] + merged = _usage( + first["prompt_tokens"], + sum(u["completion_tokens"] for u in usages.values()), + max(u.get("prompt_tokens_details", {}).get("cached_tokens", 0) for u in usages.values()), + reasoning_tokens=sum( + u.get("completion_tokens_details", {}).get("reasoning_tokens", 0) for u in usages.values() + ), ) + yield _sse({**envelope(), "usage": merged}) yield b"data: [DONE]\n\n" +_FIM_TOKENS = ("<|fim_prefix|>", "<|fim_suffix|>", "<|fim_middle|>") + + +async def _fim_markers(state: Any) -> tuple[str, str, str] | None: + """The model's fill-in-the-middle marker strings when its vocabulary has them + (Qwen / DeepSeek-Coder style), else None.""" + build = getattr(state, "frontend_tokenizer", None) + if build is None: + return None + try: + manager = await asyncio.to_thread(build) + tok = manager.tokenizer + unk = getattr(tok, "unk_token_id", None) + for name in _FIM_TOKENS: + tid = tok.convert_tokens_to_ids(name) + if not isinstance(tid, int) or tid < 0 or tid == unk: + return None + except Exception: # noqa: BLE001 -- no tokenizer here means no FIM + return None + return _FIM_TOKENS + + +def _fim_prompt(prompt: str, suffix: str, markers: tuple[str, str, str]) -> str: + """The prefix-suffix-middle prompt of OpenAI's ``suffix`` (llama.cpp's /infill).""" + prefix_tok, suffix_tok, middle_tok = markers + return f"{prefix_tok}{prompt}{suffix_tok}{suffix}{middle_tok}" + + +async def handle_tokenize(req: TokenizeRequest, state: Any): + """``POST /tokenize``: the vLLM / SGLang shape. A raw ``prompt`` is encoded as is; chat + ``messages`` are rendered through the model's chat template first (the exact prompt a + generation would tokenize).""" + if (req.prompt is None) == (req.messages is None): + return create_error_response("pass exactly one of prompt or messages") + build = getattr(state, "frontend_tokenizer", None) + if build is None: + return create_error_response("no tokenizer on this server", status_code=503, err_type="api_error") + try: + manager = await asyncio.to_thread(build) + if req.prompt is not None: + ids = await asyncio.to_thread( + manager.tokenizer.encode, req.prompt, add_special_tokens=req.add_special_tokens + ) + else: + messages = render_messages([m.model_dump(exclude_none=True) for m in req.messages]) + ctk = dict(req.chat_template_kwargs) + if not req.add_generation_prompt: + ctk["continue_final_message"] = True + msg = TokenizeMsg(uid=-1, text=messages, sampling_params=SamplingParams(), chat_template_kwargs=ctk) + text = await asyncio.to_thread(manager.render_prompt, msg) + ids = await asyncio.to_thread(manager.tokenizer.encode, text, add_special_tokens=False) + except ValueError as exc: + return create_error_response(str(exc)) + body: dict[str, Any] = { + "count": len(ids), + "max_model_len": _model_context_length(state), + "tokens": [int(t) for t in ids], + } + if req.return_token_strs: + body["token_strs"] = manager.tokenizer.convert_ids_to_tokens(ids) + return body + + +async def handle_detokenize(req: DetokenizeRequest, state: Any): + build = getattr(state, "frontend_tokenizer", None) + if build is None: + return create_error_response("no tokenizer on this server", status_code=503, err_type="api_error") + manager = await asyncio.to_thread(build) + text = await asyncio.to_thread( + manager.tokenizer.decode, req.tokens, skip_special_tokens=req.skip_special_tokens + ) + return {"prompt": text} + + def create_error_response( message: str, status_code: int = 400, @@ -550,9 +744,25 @@ def _resolve_sampling( model_sampling=model_sampling, stop=req.stop, default_max_tokens=default_max_tokens, + min_p=req.min_p, + presence_penalty=req.presence_penalty, + frequency_penalty=req.frequency_penalty, + repetition_penalty=req.repetition_penalty, + logit_bias=req.logit_bias, + min_tokens=req.min_tokens, + stop_token_ids=req.stop_token_ids, + include_stop_str_in_output=req.include_stop_str_in_output, + skip_special_tokens=req.skip_special_tokens, ) +def _response_id(prefix: str, req: Any, fallback: Any) -> str: + """``request_id`` from the client when given (vLLM request_id / SGLang rid), else the + server's own id.""" + rid = getattr(req, "request_id", None) + return f"{prefix}-{rid}" if rid else f"{prefix}-{fallback}" + + def _tools_for_template(req: ChatCompletionRequest) -> list[dict[str, Any]] | None: if not _should_parse_tools(req): return None @@ -613,10 +823,11 @@ def _tool_call_id(name: str | None, index: int) -> str: def _chat_chunk(req: ChatCompletionRequest, uid: int, choices: list[dict[str, Any]]) -> dict[str, Any]: return { - "id": f"chatcmpl-{uid}", + "id": _response_id("chatcmpl", req, uid), "object": "chat.completion.chunk", "created": int(time.time()), "model": req.model, + "system_fingerprint": None, "choices": choices, } @@ -630,7 +841,9 @@ def _reported_cached(state: Any, cached_tokens: int) -> int: return cached_tokens if getattr(state.config, "enable_cache_report", False) else 0 -def _usage(prompt_tokens: int, completion_tokens: int, cached_tokens: int = 0) -> dict[str, Any]: +def _usage( + prompt_tokens: int, completion_tokens: int, cached_tokens: int = 0, reasoning_tokens: int = 0 +) -> dict[str, Any]: usage: dict[str, Any] = { "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, @@ -640,6 +853,10 @@ def _usage(prompt_tokens: int, completion_tokens: int, cached_tokens: int = 0) - # disabled report and a 0-token hit serialize identically. if cached_tokens > 0: usage["prompt_tokens_details"] = {"cached_tokens": cached_tokens} + # OpenAI's completion_tokens_details.reasoning_tokens: tokens up to and including the + # reasoning end tag, only when the model emitted one. + if reasoning_tokens > 0: + usage["completion_tokens_details"] = {"reasoning_tokens": reasoning_tokens} return usage @@ -653,12 +870,6 @@ def _completion_unsupported_reason(req: CompletionRequest) -> str | None: return "OpenAI token-id prompt inputs are not supported; pass text prompt strings instead" if req.logprobs is not None: return "logprobs is not supported" - if req.echo: - return "echo is not supported" - if req.suffix is not None: - return "suffix is not supported" - if req.logit_bias is not None: - return "logit_bias is not supported" if _response_format_unsupported(req.response_format): return "response_format json_object/json_schema is not supported (no constrained decoding)" return None diff --git a/python/freetoken/server/stats.py b/python/freetoken/server/stats.py index 3d48ee96a0..c9d0212c32 100644 --- a/python/freetoken/server/stats.py +++ b/python/freetoken/server/stats.py @@ -174,3 +174,44 @@ def build_stats(state: Any, p95_ms: int, ttft_mean_ms: int) -> dict: "completion_tokens_total": tr.completion_tokens_total, }, } + + + +def prometheus_text(doc: dict) -> str: + """The /v1/stats document as Prometheus text exposition (gauges and counters).""" + lines: list[str] = [] + + def gauge(name: str, value, help_text: str, kind: str = "gauge") -> None: + if value is None: + return + lines.append(f"# HELP {name} {help_text}") + lines.append(f"# TYPE {name} {kind}") + lines.append(f"{name} {value}") + + model = doc.get("model") or {} + model_id = str(model.get("id") or model.get("name") or "").replace('"', "'") + lines.append("# HELP freetoken_info Server identity.") + lines.append("# TYPE freetoken_info gauge") + lines.append(f'freetoken_info{{instance_id="{doc.get("instance_id")}",model="{model_id}"}} 1') + gauge("freetoken_uptime_seconds", doc.get("uptime_s"), "Seconds since the server became ready.") + req = doc.get("requests") or {} + gauge("freetoken_requests_active", req.get("active"), "Requests in flight.") + gauge("freetoken_requests_completed_total", req.get("completed"), "Requests completed.", "counter") + gauge("freetoken_request_latency_p95_milliseconds", req.get("p95_ms"), "p95 request latency over the recent window.") + gauge("freetoken_ttft_mean_milliseconds", req.get("ttft_mean_ms"), "Mean time to first token over the recent window.") + gauge("freetoken_prompt_tokens_total", req.get("prompt_tokens_total"), "Prompt tokens processed.", "counter") + gauge("freetoken_completion_tokens_total", req.get("completion_tokens_total"), "Tokens generated.", "counter") + tp = doc.get("throughput") or {} + gauge("freetoken_decode_tokens_per_second", tp.get("decode_tps"), "Decode throughput, sliding window.") + gauge("freetoken_prefill_tokens_per_second", tp.get("prefill_tps"), "Prefill throughput, sliding window.") + for pool in ("kv", "swa"): + p = doc.get(pool) + if p: + gauge(f"freetoken_{pool}_pages_used", p.get("used_pages"), f"{pool} pages in use.") + gauge(f"freetoken_{pool}_pages_total", p.get("total_pages"), f"{pool} pages in the pool.") + mamba = doc.get("mamba") + if mamba: + gauge("freetoken_mamba_slots_used", mamba.get("used_slots"), "GDN state slots in use.") + gauge("freetoken_mamba_slots_total", mamba.get("total_slots"), "GDN state slots in the pool.") + gauge("freetoken_vram_bytes", doc.get("vram_bytes"), "GPU memory in use by the engine.") + return "\n".join(lines) + "\n" diff --git a/python/freetoken/tokenizer/detokenize.py b/python/freetoken/tokenizer/detokenize.py index d138ac1a68..27dd3b8034 100644 --- a/python/freetoken/tokenizer/detokenize.py +++ b/python/freetoken/tokenizer/detokenize.py @@ -71,11 +71,16 @@ class DecodeStatus: read_offset: int # length of read ids surr_offset: int # length of surr ids sent_offset: int # length of sent out string + skip_special: bool = False # decode this request with skip_special_tokens + think_end_at: int | None = None # generated tokens up to and incl. the reasoning end tag class DetokenizeManager: def __init__( - self, tokenizer: PreTrainedTokenizerBase, eos_token_ids: FrozenSet[int] | None = None + self, + tokenizer: PreTrainedTokenizerBase, + eos_token_ids: FrozenSet[int] | None = None, + think_end_id: int | None = None, ) -> None: # uid -> DecodeStatus self.decode_map: Dict[int, DecodeStatus] = {} @@ -83,6 +88,9 @@ def __init__( self.eos_token_ids = ( eos_token_ids if eos_token_ids is not None else frozenset({tokenizer.eos_token_id}) ) + # The reasoning end tag's id (e.g. ````): tokens generated up to and + # including it are reported as usage.completion_tokens_details.reasoning_tokens. + self.think_end_id = think_end_id if think_end_id is None else int(think_end_id) def discard(self, uid: int) -> None: """Drop a uid's decode state without a finished DetokenizeMsg. An aborted or errored @@ -91,8 +99,26 @@ def discard(self, uid: int) -> None: self.decode_map.pop(uid, None) def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]: + return self.detokenize_with_meta(msgs)[0] + + def _batch_decode(self, ids: List[List[int]], skip: List[bool]) -> List[str]: + """batch_decode honouring a per-request skip_special_tokens flag (one call per + distinct flag value, results back in order).""" + out: List[str | None] = [None] * len(ids) + for flag in (False, True): + idx = [i for i, f in enumerate(skip) if f == flag] + if idx: + texts = self.tokenizer.batch_decode([ids[i] for i in idx], skip_special_tokens=flag) + for i, t in zip(idx, texts, strict=True): + out[i] = t + return [t if t is not None else "" for t in out] + + def detokenize_with_meta(self, msgs: List[DetokenizeMsg]) -> tuple[List[str], List[int]]: + """Incremental text per message plus, for a finished message, its reasoning token + count (0 otherwise or when the model emitted no reasoning end tag).""" read_ids: List[List[int]] = [] surr_ids: List[List[int]] = [] + skip: List[bool] = [] for msg in msgs: if msg.uid not in self.decode_map: self.decode_map[msg.uid] = DecodeStatus( @@ -101,17 +127,26 @@ def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]: read_offset=0, surr_offset=0, sent_offset=0, + skip_special=bool(getattr(msg, "skip_special_tokens", False)), ) s = self.decode_map[msg.uid] if not (msg.finished and msg.next_token in self.eos_token_ids): s.decoded_ids.append(msg.next_token) + if ( + self.think_end_id is not None + and s.think_end_at is None + and msg.next_token == self.think_end_id + ): + s.think_end_at = len(s.decoded_ids) read_ids.append(s.decoded_ids[s.surr_offset :]) surr_ids.append(s.decoded_ids[s.surr_offset : s.read_offset]) + skip.append(s.skip_special) - read_texts = self.tokenizer.batch_decode(read_ids) - surr_texts = self.tokenizer.batch_decode(surr_ids) + read_texts = self._batch_decode(read_ids, skip) + surr_texts = self._batch_decode(surr_ids, skip) incremental_strs: List[str] = [] + reasoning_tokens: List[int] = [] for msg, read_str, surr_str in zip(msgs, read_texts, surr_texts, strict=True): s = self.decode_map[msg.uid] new_text = read_str[len(surr_str) :] @@ -127,10 +162,16 @@ def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]: prev_sent = s.sent_offset if msg.finished: - # Generation is over: flush everything, trimming at the matched stop string. + # Generation is over: flush everything, trimming at the matched stop string + # (or keeping it, for include_stop_str_in_output). if msg.matched_stop: cut = output_str.find(msg.matched_stop) - emit_end = cut if cut >= 0 else len(output_str) + if cut < 0: + emit_end = len(output_str) + elif getattr(msg, "keep_stop_str", False): + emit_end = cut + len(msg.matched_stop) + else: + emit_end = cut else: emit_end = len(output_str) elif msg.stop_strs: @@ -141,7 +182,8 @@ def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]: incremental_output = output_str[prev_sent:emit_end] if emit_end > prev_sent else "" s.sent_offset = max(prev_sent, emit_end) incremental_strs.append(incremental_output) + reasoning_tokens.append((s.think_end_at or 0) if msg.finished else 0) if msg.finished: del self.decode_map[msg.uid] - return incremental_strs + return incremental_strs, reasoning_tokens diff --git a/python/freetoken/tokenizer/server.py b/python/freetoken/tokenizer/server.py index d3365bfecd..2050fea3f6 100644 --- a/python/freetoken/tokenizer/server.py +++ b/python/freetoken/tokenizer/server.py @@ -81,6 +81,19 @@ def _send_generation_replies( _put_user_replies(send_frontend, terminal_errors) +def _reasoning_end_id(tokenizer) -> int | None: + """The id of the reasoning end tag when the vocabulary has one (```` for the + Qwen / DeepSeek / GLM families), else None. Only used to count reasoning tokens.""" + for tag in ("", "<|end_of_thinking|>"): + try: + tid = tokenizer.convert_tokens_to_ids(tag) + except Exception: # noqa: BLE001 -- a tokenizer without the tag answers unk or raises + continue + if isinstance(tid, int) and tid >= 0 and tid != getattr(tokenizer, "unk_token_id", -1): + return tid + return None + + def _tokenize_requests( tokenize_manager: Any, messages: List[TokenizeMsg], @@ -151,7 +164,9 @@ def tokenize_worker( tokenize_manager = TokenizeManager(tokenizer, get_mm_processor(tokenizer_path, mm)) detokenize_manager = DetokenizeManager( - tokenizer, load_eos_token_ids(tokenizer_path, tokenizer) + tokenizer, + load_eos_token_ids(tokenizer_path, tokenizer), + think_end_id=_reasoning_end_id(tokenizer), ) if ack_queue is not None: @@ -209,7 +224,7 @@ def tokenize_worker( ) sampled_replies: List[UserReply] = [] if len(detokenize_msg) > 0: - replies = detokenize_manager.detokenize(detokenize_msg) + replies, reasoning_counts = detokenize_manager.detokenize_with_meta(detokenize_msg) sampled_replies = [ UserReply( uid=msg.uid, @@ -217,6 +232,7 @@ def tokenize_worker( finished=msg.finished, finish_reason=msg.finish_reason, matched_stop=msg.matched_stop, + reasoning_tokens=reasoning_tokens, completion_tokens_delta=1, kv_used_pages=msg.kv_used_pages, kv_total_pages=msg.kv_total_pages, @@ -226,7 +242,9 @@ def tokenize_worker( swa_total_tokens=msg.swa_total_tokens, gpu_mem_bytes=msg.gpu_mem_bytes, ) - for msg, reply in zip(detokenize_msg, replies, strict=True) + for msg, reply, reasoning_tokens in zip( + detokenize_msg, replies, reasoning_counts, strict=True + ) ] # An error reply and a client abort are both terminal for their uid, and neither diff --git a/python/freetoken/tokenizer/tokenize.py b/python/freetoken/tokenizer/tokenize.py index fc66391cf2..f512930dee 100644 --- a/python/freetoken/tokenizer/tokenize.py +++ b/python/freetoken/tokenizer/tokenize.py @@ -176,10 +176,16 @@ def _render( if tools is not None: chat_template_kwargs = {**chat_template_kwargs, "tools": tools} messages = _map_developer_role(messages, getattr(self.tokenizer, "chat_template", None)) + # continue_final_message (an OpenAI-compatible extra, as in vLLM / SGLang): the last + # message is an assistant prefix the model must continue, so no generation prompt. + chat_template_kwargs = dict(chat_template_kwargs) + continue_final = bool(chat_template_kwargs.pop("continue_final_message", False)) + if continue_final: + chat_template_kwargs["continue_final_message"] = True prompt = self.tokenizer.apply_chat_template( messages, tokenize=False, - add_generation_prompt=True, + add_generation_prompt=not continue_final, **chat_template_kwargs, ) assert isinstance(prompt, str) diff --git a/tests/engine/test_logits_processors.py b/tests/engine/test_logits_processors.py new file mode 100644 index 0000000000..f051dec51f --- /dev/null +++ b/tests/engine/test_logits_processors.py @@ -0,0 +1,155 @@ +"""The logits processors behind min_p, the penalties, logit_bias and min_tokens: pure +torch on the CPU, so the math is checked without a GPU or the sampling kernels.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import torch +from freetoken.core import SamplingParams +from freetoken.engine.sample import LogitsPlan, Sampler, apply_logits_processors + +V = 8 +CPU = torch.device("cpu") + + +def _plan(rows, **fields) -> LogitsPlan: + plan = LogitsPlan( + rows=torch.tensor(rows, dtype=torch.int64), + temps=torch.ones((len(rows), 1), dtype=torch.float32), + ) + for k, v in fields.items(): + setattr(plan, k, v) + return plan + + +def _ids(rows: list[list[int]]) -> torch.Tensor: + width = max(len(r) for r in rows) + out = torch.full((len(rows), width), V, dtype=torch.int64) + for i, r in enumerate(rows): + out[i, : len(r)] = torch.tensor(r) + return out + + +def test_presence_and_frequency_penalties_count_generated_tokens(): + logits = torch.zeros((2, V)) + plan = _plan( + [1], + hist_ids=_ids([[3, 3, 5]]), + presence=torch.tensor([[0.5]]), + frequency=torch.tensor([[0.25]]), + ) + out = apply_logits_processors(logits, plan, V) + assert torch.equal(out[0], logits[0]) # row outside the plan untouched + assert out[1, 3].item() == -(0.5 + 2 * 0.25) + assert out[1, 5].item() == -(0.5 + 0.25) + assert out[1, 0].item() == 0.0 + + +def test_repetition_penalty_covers_prompt_and_output_with_hf_semantics(): + logits = torch.tensor([[2.0, -2.0, 2.0, -2.0, 1.0, 1.0, 1.0, 1.0]]) + plan = _plan( + [0], + hist_ids=_ids([[0, 1]]), # generated: 0 (positive), 1 (negative) + prompt_ids=_ids([[2, 3]]), # prompt: 2 (positive), 3 (negative) + repetition=torch.tensor([[2.0]]), + ) + out = apply_logits_processors(logits, plan, V) + assert out[0, 0].item() == 1.0 and out[0, 2].item() == 1.0 # divided + assert out[0, 1].item() == -4.0 and out[0, 3].item() == -4.0 # multiplied + assert out[0, 4].item() == 1.0 # unseen + + +def test_logit_bias_adds_per_row(): + logits = torch.zeros((2, V)) + plan = _plan( + [0, 1], + bias_rows=torch.tensor([0, 1, 1]), + bias_ids=torch.tensor([2, 2, 6]), + bias_vals=torch.tensor([5.0, -100.0, 1.5]), + ) + out = apply_logits_processors(logits, plan, V) + assert ( + out[0, 2].item() == 5.0 + and out[1, 2].item() == -100.0 + and out[1, 6].item() == 1.5 + ) + + +def test_min_tokens_masks_the_stop_ids_of_the_row_only(): + logits = torch.zeros((2, V)) + plan = _plan([0, 1], min_rows=torch.tensor([1]), min_ids=_ids([[0, 7]])) + out = apply_logits_processors(logits, plan, V) + assert out[1, 0].item() == float("-inf") and out[1, 7].item() == float("-inf") + assert out[1, 3].item() == 0.0 + assert torch.isfinite(out[0]).all() + + +def test_min_p_drops_tokens_below_the_fraction_of_the_top_probability(): + logits = torch.log(torch.tensor([[0.5, 0.3, 0.1, 0.05, 0.03, 0.01, 0.005, 0.005]])) + plan = _plan( + [0], min_p=torch.tensor([[0.15]]) + ) # threshold 0.075, between 0.1 and 0.05 + out = apply_logits_processors(logits, plan, V) + kept = torch.isfinite(out[0]) + assert kept.tolist() == [True, True, True, False, False, False, False, False] + # a row with min_p 0 is left alone + plan0 = _plan([0], min_p=torch.tensor([[0.0]])) + assert torch.isfinite(apply_logits_processors(logits, plan0, V)[0]).all() + + +def _req( + prompt: list[int], generated: list[int], max_tokens: int, **sp +) -> SimpleNamespace: + ids = torch.tensor(prompt + generated, dtype=torch.int32) + return SimpleNamespace( + input_ids=ids, + output_len=max_tokens, + max_device_len=len(prompt) + max_tokens, + sampling_params=SamplingParams(**sp), + ) + + +def test_prepare_builds_no_plan_without_processors_and_a_plan_with_them(): + sampler = Sampler(CPU, V) + plain = SimpleNamespace(reqs=[_req([1, 2], [3], 4)]) + assert sampler.prepare(plain).plan is None + + # all greedy: the argmax path runs on the CPU, the sampling kernels need a GPU + batch = SimpleNamespace( + reqs=[ + _req([1, 2], [3], 4), + _req( + [1, 2], + [3, 3], + 4, + presence_penalty=1.0, + repetition_penalty=1.5, + logit_bias=[[5, 2.0]], + ), + _req([4], [], 4, min_tokens=2, min_tokens_stop_ids=[0, 7]), + _req( + [4], [6, 6], 4, min_tokens=2, min_tokens_stop_ids=[0] + ), # already past min_tokens + ] + ) + args = sampler.prepare(batch) + plan = args.plan + assert plan is not None + assert plan.rows.tolist() == [1, 2, 3] + assert plan.hist_ids[0].tolist()[:2] == [3, 3] # row 1's generated tokens + assert plan.prompt_ids[0].tolist()[:2] == [1, 2] + assert plan.bias_rows.tolist() == [0] and plan.bias_ids.tolist() == [5] + assert plan.min_rows.tolist() == [1] # local row of the third request only + assert plan.min_ids[0].tolist() == [0, 7] + + logits = torch.zeros((4, V)) + out = apply_logits_processors(logits, plan, V) + assert ( + out[1, 3].item() == -1.0 + ) # presence on the repeated 3 (repetition on 0 is a no-op) + assert out[1, 5].item() == 2.0 + assert out[2, 0].item() == float("-inf") and out[2, 7].item() == float("-inf") + assert torch.isfinite(out[3]).all() + # greedy path: argmax over the processed logits + assert sampler.sample(logits, args)[1].item() == 5 diff --git a/tests/server/test_openai_extras.py b/tests/server/test_openai_extras.py new file mode 100644 index 0000000000..115112b9d0 --- /dev/null +++ b/tests/server/test_openai_extras.py @@ -0,0 +1,383 @@ +"""OpenAI protocol coverage: n > 1 fan-out, echo and suffix (FIM) on completions, the +sampling extras reaching SamplingParams, request_id, usage details, and the tokenize / +detokenize routes.""" + +from __future__ import annotations + +import asyncio +import json +from types import SimpleNamespace + +import pytest +from freetoken.message import TokenizeMsg, UserReply +from freetoken.server.api_models import ( + ChatCompletionRequest, + CompletionRequest, + DetokenizeRequest, + TokenizeRequest, +) +from freetoken.server.openai_api import ( + _usage, + chat_request_to_genspec, + handle_chat_completion, + handle_completion, + handle_detokenize, + handle_tokenize, +) +from freetoken.server.stats import prometheus_text + + +def run(coro): + return asyncio.run(coro) + + +class FakeState: + """Distinct uids per generation; every uid gets the same two-ack reply.""" + + def __init__(self, text=("hel", "lo"), reasoning_tokens=0, tokenizer=None): + self.config = SimpleNamespace( + model_path="/models/unit-model", + served_model_name="unit-model", + tool_call_parser="llama3", + reasoning_parser=None, + max_seq_len=4096, + ) + self.text = text + self.reasoning_tokens = reasoning_tokens + self._next = 100 + self.sent: list[TokenizeMsg] = [] + self.aborted: list[int] = [] + if tokenizer is not None: + manager = SimpleNamespace( + tokenizer=tokenizer, + render_prompt=lambda msg: "RENDERED:" + json.dumps(msg.text), + ) + self.frontend_tokenizer = lambda: manager + + def new_user(self) -> int: + self._next += 1 + return self._next + + async def send_one(self, msg): + self.sent.append(msg) + + async def abort_user(self, uid): + self.aborted.append(uid) + + async def wait_for_ack(self, uid: int): + yield UserReply( + uid=uid, incremental_output="", finished=False, prompt_tokens_delta=5 + ) + for i, piece in enumerate(self.text): + last = i == len(self.text) - 1 + yield UserReply( + uid=uid, + incremental_output=piece, + finished=last, + completion_tokens_delta=1, + finish_reason="stop" if last else None, + reasoning_tokens=self.reasoning_tokens if last else 0, + ) + + +def chat(**kwargs) -> ChatCompletionRequest: + body = {"model": "unit-model", "messages": [{"role": "user", "content": "hi"}]} + body.update(kwargs) + return ChatCompletionRequest(**body) + + +def parse_sse(chunks: list[bytes]) -> list: + out = [] + for chunk in chunks: + for line in chunk.decode().split("\n"): + if not line.startswith("data: "): + continue + payload = line[len("data: ") :] + out.append("[DONE]" if payload == "[DONE]" else json.loads(payload)) + return out + + +async def _collect(generator): + return [item async for item in generator] + + +# ---- n > 1 ------------------------------------------------------------------------------ + + +def test_chat_n_returns_one_choice_per_generation_and_counts_the_prompt_once(): + state = FakeState(reasoning_tokens=2) + body = run(handle_chat_completion(chat(n=2), None, state, {})) + assert [c["index"] for c in body["choices"]] == [0, 1] + assert all(c["message"]["content"] == "hello" for c in body["choices"]) + assert len(state.sent) == 2 and len({m.uid for m in state.sent}) == 2 + usage = body["usage"] + assert ( + usage["prompt_tokens"] == 5 + and usage["completion_tokens"] == 4 + and usage["total_tokens"] == 9 + ) + assert usage["completion_tokens_details"] == {"reasoning_tokens": 4} + assert body["system_fingerprint"] is None + + +def test_chat_n_stream_interleaves_choice_indices_and_ends_once(): + state = FakeState() + response = run( + handle_chat_completion( + chat(n=2, stream=True, stream_options={"include_usage": True}), + None, + state, + {}, + ) + ) + events = parse_sse(run(_collect(response.body_iterator))) + assert events.count("[DONE]") == 1 and events[-1] == "[DONE]" + indices = {c["index"] for ev in events[:-1] for c in ev.get("choices", [])} + assert indices == {0, 1} + usage_chunks = [ev for ev in events[:-1] if ev.get("usage")] + assert len(usage_chunks) == 1 + assert usage_chunks[0]["usage"]["prompt_tokens"] == 5 + assert usage_chunks[0]["usage"]["completion_tokens"] == 4 + finals = [ + c for ev in events[:-1] for c in ev.get("choices", []) if c.get("finish_reason") + ] + assert len(finals) == 2 + + +def test_n_out_of_range_is_400(): + for n in (0, 17): + body = run(handle_chat_completion(chat(n=n), None, FakeState(), {})) + assert body.status_code == 400 + assert json.loads(body.body)["error"]["param"] == "n" + + +# ---- completions: echo, suffix, n --------------------------------------------------------- + + +def test_completion_echo_prepends_the_prompt(): + state = FakeState() + req = CompletionRequest(model="unit-model", prompt="Once upon", echo=True) + body = run(handle_completion(req, None, state, {})) + assert body["choices"][0]["text"] == "Once uponhello" + + +def test_completion_echo_streams_the_prompt_first(): + state = FakeState() + req = CompletionRequest( + model="unit-model", prompt="Once upon", echo=True, stream=True + ) + response = run(handle_completion(req, None, state, {})) + events = parse_sse(run(_collect(response.body_iterator))) + assert events[0]["choices"][0]["text"] == "Once upon" + assert events[-1] == "[DONE]" + + +class _FimTokenizer: + unk_token_id = -1 + + def convert_tokens_to_ids(self, name): + return {"<|fim_prefix|>": 10, "<|fim_suffix|>": 11, "<|fim_middle|>": 12}.get( + name, -1 + ) + + +class _NoFimTokenizer: + unk_token_id = 3 + + def convert_tokens_to_ids(self, name): + return 3 + + +def test_completion_suffix_builds_a_fill_in_the_middle_prompt(): + state = FakeState(tokenizer=_FimTokenizer()) + req = CompletionRequest(model="unit-model", prompt="def f(", suffix=" return 1") + run(handle_completion(req, None, state, {})) + assert ( + state.sent[0].text + == "<|fim_prefix|>def f(<|fim_suffix|> return 1<|fim_middle|>" + ) + + +def test_completion_suffix_needs_fim_tokens(): + state = FakeState(tokenizer=_NoFimTokenizer()) + req = CompletionRequest(model="unit-model", prompt="def f(", suffix="x") + body = run(handle_completion(req, None, state, {})) + assert ( + body.status_code == 400 and json.loads(body.body)["error"]["param"] == "suffix" + ) + + +def test_completion_n_orders_choices_prompt_major(): + state = FakeState() + req = CompletionRequest(model="unit-model", prompt=["a", "b"], n=2) + body = run(handle_completion(req, None, state, {})) + assert [c["index"] for c in body["choices"]] == [0, 1, 2, 3] + assert [m.text for m in state.sent] == ["a", "a", "b", "b"] + assert body["usage"]["prompt_tokens"] == 10 # two prompts, each counted once + assert body["usage"]["completion_tokens"] == 8 + + +# ---- sampling extras -------------------------------------------------------------------- + + +def test_sampling_extras_reach_sampling_params(): + spec = chat_request_to_genspec( + chat( + min_p=0.1, + repetition_penalty=1.2, + presence_penalty=0.5, + frequency_penalty=0.3, + logit_bias={"5": 150, "7": -3}, + min_tokens=3, + stop_token_ids=[9], + include_stop_str_in_output=True, + skip_special_tokens=True, + seed=7, + user="u1", + ), + {}, + ) + sp = spec.sampling_params + assert sp.min_p == 0.1 and sp.repetition_penalty == 1.2 + assert sp.presence_penalty == 0.5 and sp.frequency_penalty == 0.3 + assert sp.logit_bias == [[5, 100.0], [7, -3.0]] # clamped to the OpenAI range + assert sp.min_tokens == 3 and sp.stop_token_ids == [9] + assert sp.include_stop_str_in_output is True and sp.skip_special_tokens is True + assert sp.needs_logits_processing + + +@pytest.mark.parametrize( + "field", + [ + {"min_p": 1.5}, + {"presence_penalty": 3.0}, + {"frequency_penalty": -2.5}, + {"repetition_penalty": 0.0}, + {"logit_bias": {"x": 1.0}}, + {"min_tokens": -1}, + {"stop_token_ids": [-1]}, + ], +) +def test_bad_sampling_extras_are_400(field): + body = run(handle_chat_completion(chat(**field), None, FakeState(), {})) + assert body.status_code == 400, field + + +def test_continue_final_message_needs_an_assistant_tail_and_reaches_the_template(): + spec = chat_request_to_genspec( + chat( + messages=[ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "Sure,"}, + ], + continue_final_message=True, + ), + {}, + ) + assert spec.chat_template_kwargs["continue_final_message"] is True + with pytest.raises(ValueError): + chat(continue_final_message=True) + + +# ---- ids and usage ------------------------------------------------------------------------ + + +def test_request_id_becomes_the_response_id(): + state = FakeState() + body = run(handle_chat_completion(chat(request_id="abc-1"), None, state, {})) + assert body["id"] == "chatcmpl-abc-1" + response = run( + handle_chat_completion(chat(request_id="abc-2", stream=True), None, state, {}) + ) + events = parse_sse(run(_collect(response.body_iterator))) + assert events[0]["id"] == "chatcmpl-abc-2" + + +def test_usage_details_appear_only_when_nonzero(): + assert "completion_tokens_details" not in _usage(10, 5) + assert _usage(10, 5, reasoning_tokens=3)["completion_tokens_details"] == { + "reasoning_tokens": 3 + } + assert _usage(10, 5, cached_tokens=4)["prompt_tokens_details"] == { + "cached_tokens": 4 + } + + +# ---- tokenize / detokenize / metrics ----------------------------------------------------- + + +class _Tokenizer: + def encode(self, text, add_special_tokens=True): + ids = [ord(c) for c in text] + return ([1] + ids) if add_special_tokens else ids + + def decode(self, ids, skip_special_tokens=False): + return "".join(chr(i) for i in ids if not (skip_special_tokens and i == 1)) + + def convert_ids_to_tokens(self, ids): + return [chr(i) for i in ids] + + +def test_tokenize_prompt_and_messages_and_detokenize(): + state = FakeState(tokenizer=_Tokenizer()) + body = run(handle_tokenize(TokenizeRequest(prompt="ab"), state)) + assert body == {"count": 3, "max_model_len": 4096, "tokens": [1, 97, 98]} + body = run( + handle_tokenize( + TokenizeRequest( + prompt="ab", add_special_tokens=False, return_token_strs=True + ), + state, + ) + ) + assert body["tokens"] == [97, 98] and body["token_strs"] == ["a", "b"] + body = run( + handle_tokenize( + TokenizeRequest(messages=[{"role": "user", "content": "x"}]), state + ) + ) + assert body["count"] == len( + "RENDERED:" + json.dumps([{"role": "user", "content": "x"}]) + ) + bad = run( + handle_tokenize( + TokenizeRequest(prompt="a", messages=[{"role": "user", "content": "x"}]), + state, + ) + ) + assert bad.status_code == 400 + body = run( + handle_detokenize( + DetokenizeRequest(tokens=[1, 104, 105], skip_special_tokens=True), state + ) + ) + assert body == {"prompt": "hi"} + + +def test_metrics_exposition_lists_the_stats_document(): + doc = { + "instance_id": "abc", + "model": {"id": "unit-model"}, + "uptime_s": 12, + "kv": {"used_pages": 3, "total_pages": 10, "page_size": 64}, + "mamba": {"used_slots": 1, "total_slots": 4}, + "swa": None, + "vram_bytes": 123, + "throughput": {"decode_tps": 99.5, "prefill_tps": 1000.0}, + "requests": { + "active": 1, + "completed": 7, + "p95_ms": 50, + "ttft_mean_ms": 20, + "prompt_tokens_total": 100, + "completion_tokens_total": 40, + }, + } + text = prometheus_text(doc) + assert 'freetoken_info{instance_id="abc",model="unit-model"} 1' in text + assert "freetoken_requests_active 1" in text + assert "freetoken_requests_completed_total 7" in text + assert "freetoken_kv_pages_total 10" in text + assert "freetoken_mamba_slots_used 1" in text + assert "freetoken_decode_tokens_per_second 99.5" in text + assert "# TYPE freetoken_prompt_tokens_total counter" in text diff --git a/tests/tokenizer/test_detokenize_extras.py b/tests/tokenizer/test_detokenize_extras.py new file mode 100644 index 0000000000..5617a10537 --- /dev/null +++ b/tests/tokenizer/test_detokenize_extras.py @@ -0,0 +1,81 @@ +"""Per-request detokenizer options: include_stop_str_in_output, skip_special_tokens, and +the reasoning token count reported on the finished reply.""" + +from __future__ import annotations + +from freetoken.message import DetokenizeMsg +from freetoken.tokenizer.detokenize import DetokenizeManager + +EOS, SPECIAL, THINK_END = 0, 99, 7 + + +class _Tok: + """Token ids decode to letters; 99 is a special marker, 7 the reasoning end tag.""" + + eos_token_id = EOS + unk_token_id = -1 + + def decode(self, ids, skip_special_tokens=False): + out = [] + for t in ids: + if t == SPECIAL: + if not skip_special_tokens: + out.append("") + elif t == THINK_END: + out.append("") + else: + out.append(chr(ord("a") + t)) + return "".join(out) + + def batch_decode(self, ids_list, skip_special_tokens=False): + return [self.decode(ids, skip_special_tokens) for ids in ids_list] + + +def _run(manager, uid, tokens, *, finished_kwargs=None, **msg_kwargs): + text = "" + meta = 0 + for i, t in enumerate(tokens): + last = i == len(tokens) - 1 + kwargs = dict(msg_kwargs) + if last and finished_kwargs: + kwargs.update(finished_kwargs) + strs, counts = manager.detokenize_with_meta( + [DetokenizeMsg(uid=uid, next_token=t, finished=last, **kwargs)] + ) + text += strs[0] + meta = counts[0] + return text, meta + + +def test_stop_string_is_trimmed_by_default_and_kept_on_request(): + manager = DetokenizeManager(_Tok(), frozenset({EOS})) + trimmed, _ = _run( + manager, 1, [1, 2, 3], stop_strs=["cd"], finished_kwargs={"matched_stop": "cd"} + ) + assert trimmed == "b" + kept, _ = _run( + manager, + 2, + [1, 2, 3], + stop_strs=["cd"], + keep_stop_str=True, + finished_kwargs={"matched_stop": "cd"}, + ) + assert kept == "bcd" + + +def test_skip_special_tokens_is_per_request(): + manager = DetokenizeManager(_Tok(), frozenset({EOS})) + shown, _ = _run(manager, 1, [1, SPECIAL, 2, EOS]) + hidden, _ = _run(manager, 2, [1, SPECIAL, 2, EOS], skip_special_tokens=True) + assert shown == "bc" + assert hidden == "bc" + + +def test_reasoning_tokens_count_up_to_the_end_tag(): + manager = DetokenizeManager(_Tok(), frozenset({EOS}), think_end_id=THINK_END) + text, reasoning = _run(manager, 1, [1, 2, THINK_END, 3, 4, EOS]) + assert text == "bcde" + assert reasoning == 3 # b, c and the tag itself + _, none = _run(manager, 2, [1, 2, EOS]) + assert none == 0 From 718e7036a154fdb390249d570963f5af3f6e3642 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 17 Sep 2026 22:53:22 -0400 Subject: [PATCH 20/73] fix(models): skip the ftw layout check for a bare ftw without an hf config Fork #19 asks the model family to check the stored fp8 split against the env flags; a synthetic FTW written straight from tensors (upstream's test_ftw_weights) carries no config.json, so there is no family to ask and the read proceeds as before. Assisted-by: Claude Fable 5.1 (cherry picked from commit cab9be35b8e56fa3cf08a109017183691e2e69c5) (cherry picked from commit e59e5441e44dac96d530f4e2cac34d2c6273411f) --- python/freetoken/models/weight.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/python/freetoken/models/weight.py b/python/freetoken/models/weight.py index 9016b9477e..29cce4eefa 100644 --- a/python/freetoken/models/weight.py +++ b/python/freetoken/models/weight.py @@ -224,9 +224,13 @@ def load_weight( # a text-only engine never built the tower, so its tensors are not even read keep = None if include_vision else (lambda name: not name.startswith(VISION_KEY_PREFIXES)) if is_ftw_checkpoint(model_path): - spec = get_model_spec(cached_load_hf_config(model_path).architectures[0]) - # a family whose stored tensors depend on load-time settings checks them against the index before the read - check = _model_override(spec, "check_ftw_weights") + # a family whose stored tensors depend on load-time settings checks them against the index before the + # read; a bare FTW without a readable HF config (synthetic test checkpoints) has no family to ask + try: + spec = get_model_spec(cached_load_hf_config(model_path).architectures[0]) + except (OSError, ValueError): + spec = None + check = _model_override(spec, "check_ftw_weights") if spec is not None else None if check is not None: reader = FTWReader(model_path) names = {e["name"] for e in reader.entries("weight")} From 79e086cb80a601c1136c78274635ac3af45e4fdf Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 17 Sep 2026 22:53:25 -0400 Subject: [PATCH 21/73] test(scheduler): give the rejection test's engine fake the pool size the admission clamp reads Assisted-by: Claude Fable 5.1 (cherry picked from commit 4584957f3d63449a153097f380ffdac035a245dd) (cherry picked from commit 0ab1764343228845db270505185dd2be7090086d) --- tests/scheduler/test_cost_accounting_core.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/scheduler/test_cost_accounting_core.py b/tests/scheduler/test_cost_accounting_core.py index 045010c6b4..116d4b88c7 100644 --- a/tests/scheduler/test_cost_accounting_core.py +++ b/tests/scheduler/test_cost_accounting_core.py @@ -132,7 +132,8 @@ def fail_prepare(_batch): def test_scheduler_rejection_emits_error_but_no_admission(): scheduler = Scheduler.__new__(Scheduler) - scheduler.engine = SimpleNamespace(max_seq_len=4) + scheduler.engine = SimpleNamespace(max_seq_len=4, num_pages=4) + scheduler.config = SimpleNamespace(page_size=1) added = [] scheduler.prefill_manager = SimpleNamespace(add_one_req=added.append) sent = [] From 6dd2dc0f9f7e505c5050ea8670be82d7416a34e8 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 17 Sep 2026 22:41:19 -0400 Subject: [PATCH 22/73] test(server): abort every sample of an n>1 non-streaming request on disconnect Guards the #222 resolution against #393's fan-out: the disconnect watcher takes the whole uid list, not the first sample's uid. Assisted-by: Claude Fable 5.1 (cherry picked from commit 93ff96349235ea3ca1504733a2abbd6e6bcccf94) (cherry picked from commit 69a5efa3300e93bd331b4394e54eea8cef82ad3f) (cherry picked from commit 40e994121fb79a21ba0d7ffff70e4394a4498899) --- tests/server/test_stream_cancellation.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/tests/server/test_stream_cancellation.py b/tests/server/test_stream_cancellation.py index e38fa56bfb..ba7da7c17b 100644 --- a/tests/server/test_stream_cancellation.py +++ b/tests/server/test_stream_cancellation.py @@ -294,3 +294,25 @@ def test_non_stream_connected_client_gets_result_without_abort(monkeypatch): assert result["choices"][0]["message"]["content"] == "Hi" assert state.aborted == [] + + +def test_chat_non_stream_disconnect_aborts_every_sample(monkeypatch): + # n > 1 fans out into one engine request per sample; a disconnect must stop all of them + monkeypatch.setattr(openai_api, "_DISCONNECT_POLL_SECONDS", 0.01) + + class _State(_ApiState): + def __init__(self): + super().__init__(acks=None) + self._next = 7 + + def new_user(self): + self._next += 1 + return self._next + + state = _State() + resp = asyncio.run( + openai_api.handle_chat_completion(_chat_req(n=2), _Request(disconnected=True), state, {}) + ) + + assert resp.status_code == 499 + assert sorted(state.aborted) == [8, 9] From 02d463c20f85e4b1d986e5daa12c4f8e1d5131fe Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 17 Sep 2026 23:04:45 -0400 Subject: [PATCH 23/73] fix(server): record requests with a pure asgi middleware so handlers can see a client disconnect Starlette's @app.middleware("http") (BaseHTTPMiddleware) wraps the handler's receive channel and never yields the client's http.disconnect while a non-streaming request is still computing, so request.is_disconnected() stays False and the disconnect watcher from FlashML-org/FreeToken#222 could not stop an abandoned request. Streaming requests only worked because the socket write failed. The recorder is the one such middleware; as pure ASGI it records the same row at response start. Assisted-by: Claude Fable 5.1 (cherry picked from commit 8adde913c94e2f76b3962b8d10f10ded4b1cccc6) (cherry picked from commit 202d7fec30b29a9a4320370b0c62426422422339) --- python/freetoken/server/api_server.py | 79 ++++++++++++------- .../test_request_recorder_disconnect.py | 75 ++++++++++++++++++ 2 files changed, 124 insertions(+), 30 deletions(-) create mode 100644 tests/server/test_request_recorder_disconnect.py diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index 7acb3fc133..84f1085cff 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -470,38 +470,57 @@ def _served_model_name() -> str | None: return getattr(cfg, "served_model_name", None) -@app.middleware("http") -async def _record_request_middleware(request: Request, call_next): +class _RecordRequestMiddleware: """Time every generation request into the ring for /v1/requests + /v1/stats p95. Single- model server, so model = served_model_name; stream is inferred from the response media - type. Token counts are P3 (SSE usage arrives after the handler returns) — kept as None.""" - path = request.url.path - if path.startswith(_UNTRACKED_REQUEST_PREFIXES) or not path.startswith( - _TRACKED_REQUEST_PREFIXES - ): - return await call_next(request) - import time as _time - - start = _time.monotonic() - response = await call_next(request) - duration_ms = int((_time.monotonic() - start) * 1000) - ctype = response.headers.get("content-type", "") - request_ring.record_request( - request_ring.RequestRecord( - ts=_time.strftime("%Y-%m-%dT%H:%M:%SZ", _time.gmtime()), - method=request.method, - path=path, - status=response.status_code, - model=_served_model_name(), - duration_ms=duration_ms, - ttft_ms=None, - prompt_tokens=None, - completion_tokens=None, - stream=ctype.startswith("text/event-stream"), - error=None, - ) - ) - return response + type. Token counts are P3 (SSE usage arrives after the handler returns) — kept as None. + + Pure ASGI on purpose: Starlette's ``@app.middleware("http")`` (BaseHTTPMiddleware) hands the + handler a wrapped receive channel that never yields the client's ``http.disconnect`` while a + non-streaming request is still computing, so ``request.is_disconnected()`` stays False and an + abandoned request runs to max_tokens. The recorder is the only such middleware here.""" + + def __init__(self, app): + self.app = app + + async def __call__(self, scope, receive, send): + path = scope.get("path", "") if scope["type"] == "http" else "" + if not path.startswith(_TRACKED_REQUEST_PREFIXES) or path.startswith(_UNTRACKED_REQUEST_PREFIXES): + await self.app(scope, receive, send) + return + import time as _time + + start = _time.monotonic() + + async def send_recording(message): + # the base-class version recorded once the handler returned its response object, + # i.e. at response start, before a streaming body -- keep that timing + if message["type"] == "http.response.start": + ctype = "" + for name, value in message.get("headers", ()): + if name.lower() == b"content-type": + ctype = value.decode("latin-1") + request_ring.record_request( + request_ring.RequestRecord( + ts=_time.strftime("%Y-%m-%dT%H:%M:%SZ", _time.gmtime()), + method=scope["method"], + path=path, + status=message["status"], + model=_served_model_name(), + duration_ms=int((_time.monotonic() - start) * 1000), + ttft_ms=None, + prompt_tokens=None, + completion_tokens=None, + stream=ctype.startswith("text/event-stream"), + error=None, + ) + ) + await send(message) + + await self.app(scope, receive, send_recording) + + +app.add_middleware(_RecordRequestMiddleware) class CacheRebuildRequest(BaseModel): diff --git a/tests/server/test_request_recorder_disconnect.py b/tests/server/test_request_recorder_disconnect.py new file mode 100644 index 0000000000..06ff17f99b --- /dev/null +++ b/tests/server/test_request_recorder_disconnect.py @@ -0,0 +1,75 @@ +"""The request recorder must not hide a client disconnect from a non-streaming handler. + +Real sockets on purpose: the failure only shows through uvicorn's receive channel (a +BaseHTTPMiddleware wrapper never yields ``http.disconnect`` while the handler computes), which +TestClient does not model. +""" + +from __future__ import annotations + +import asyncio +import socket +import threading +import time + +from fastapi import FastAPI, Request + +from freetoken.server.api_server import _RecordRequestMiddleware + + +def _free_port() -> int: + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def _serve(app: FastAPI, port: int): + import uvicorn + + server = uvicorn.Server(uvicorn.Config(app, host="127.0.0.1", port=port, log_level="error")) + thread = threading.Thread(target=server.run, daemon=True) + thread.start() + deadline = time.monotonic() + 10 + while not server.started: + assert time.monotonic() < deadline, "uvicorn did not start" + time.sleep(0.02) + return server + + +def _post_and_hang_up(port: int) -> None: + body = b'{"prompt": "x"}' + with socket.create_connection(("127.0.0.1", port)) as sock: + sock.sendall( + b"POST /v1/completions HTTP/1.1\r\nHost: t\r\nContent-Type: application/json\r\n" + b"Content-Length: " + str(len(body)).encode() + b"\r\n\r\n" + body + ) + time.sleep(0.2) # the handler is now polling + + +def test_non_streaming_handler_sees_the_client_leave(): + app = FastAPI() + app.add_middleware(_RecordRequestMiddleware) + seen: dict[str, float | None] = {} + + @app.post("/v1/completions") + async def slow(request: Request): + t0 = time.monotonic() + for _ in range(60): + await asyncio.sleep(0.05) + if await request.is_disconnected(): + seen["after"] = time.monotonic() - t0 + return {"detected": True} + seen["after"] = None + return {"detected": False} + + port = _free_port() + server = _serve(app, port) + try: + _post_and_hang_up(port) + deadline = time.monotonic() + 5 + while "after" not in seen and time.monotonic() < deadline: + time.sleep(0.05) + finally: + server.should_exit = True + assert seen.get("after") is not None, "handler never saw the disconnect" + assert seen["after"] < 2.0 From 2253e4bfe8ef77e585844ae255a1627aaf7cb647 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Sat, 19 Sep 2026 17:21:12 -0400 Subject: [PATCH 24/73] fix(kvcache): refuse the host kv tiers on an nvfp4 pool whose block scales they cannot mirror Upstream #499/#525 mirror KV codes plus the fp8 pool's per-token _scale_buffer; the nvfp4 pool (#408) keeps per-block scales in _block_scale_buffer, which neither the offloader nor the prefix tier copies, so --kv-host-pages or FT_PREFIX_HOST would rehydrate codes without scales. Both constructors now raise a clear error naming the knob; the GDN-snapshot-only tier (FT_GDN_HOST_TIER) never reads the KV pool and stays usable. Assisted-by: Claude Fable 5.1 (cherry picked from commit 4bb35fc34b0cb65b41584c65d7390b6c5f4f854a) --- python/freetoken/kvcache/host_prefix_tier.py | 3 ++ python/freetoken/kvcache/kv_host_offload.py | 14 +++++++ tests/kvcache/test_host_tier_nvfp4_guard.py | 44 ++++++++++++++++++++ 3 files changed, 61 insertions(+) create mode 100644 tests/kvcache/test_host_tier_nvfp4_guard.py diff --git a/python/freetoken/kvcache/host_prefix_tier.py b/python/freetoken/kvcache/host_prefix_tier.py index 0265b126ef..3b2d2d5dc3 100644 --- a/python/freetoken/kvcache/host_prefix_tier.py +++ b/python/freetoken/kvcache/host_prefix_tier.py @@ -79,6 +79,9 @@ def __init__(self, state_pool, kv_pool=None, *, gdn_slots_host: int = 64, self._kv_scales: HostBank | None = None self._kv_free: list[int] = [] if kv_pool is not None and kv_budget_bytes > 0: + from freetoken.kvcache.kv_host_offload import _reject_block_scaled_pool + + _reject_block_scaled_pool(kv_pool, "FT_PREFIX_HOST") codes = kv_pool._kv_buffer # [2, L, P, page, H, D] scales = kv_pool._scale_buffer # [2, L, P*page, H] | None (fp8) _, layers, _, page_size, heads, dim = codes.shape diff --git a/python/freetoken/kvcache/kv_host_offload.py b/python/freetoken/kvcache/kv_host_offload.py index dbb49f4814..1f8c358426 100644 --- a/python/freetoken/kvcache/kv_host_offload.py +++ b/python/freetoken/kvcache/kv_host_offload.py @@ -33,10 +33,24 @@ logger = init_logger(__name__) +def _reject_block_scaled_pool(pool, knob: str) -> None: + """The host mirrors carry codes plus the fp8 pool's per-token scales; an nvfp4 pool + (--kv-cache-dtype nvfp4) keeps per-block scales in _block_scale_buffer, which no mirror + copies, so a rehydrated page would be codes without scales. Refuse instead of decoding + garbage.""" + if getattr(pool, "_block_scale_buffer", None) is not None: + raise ValueError( + f"{knob} needs a KV pool with per-token scales (bf16 or fp8); the nvfp4 pool keeps " + "block scales the host mirror does not carry -- leave the tier off or serve with " + "--kv-cache-dtype fp8" + ) + + class KVHostOffloader: """GPU-slot LRU cache over the pinned host mirror of the logical KV page space.""" def __init__(self, pool, num_logical_pages: int, device: torch.device) -> None: + _reject_block_scaled_pool(pool, "--kv-host-pages") buf = pool._kv_buffer # [2, L, P_gpu + 1, page_size, kv_heads, head_dim] assert buf.shape[0] == 2 and buf.is_contiguous() _, num_layers, num_slots, page_size, kv_heads, head_dim = buf.shape diff --git a/tests/kvcache/test_host_tier_nvfp4_guard.py b/tests/kvcache/test_host_tier_nvfp4_guard.py new file mode 100644 index 0000000000..a793fb294a --- /dev/null +++ b/tests/kvcache/test_host_tier_nvfp4_guard.py @@ -0,0 +1,44 @@ +"""The host KV tiers (upstream #499 / #525) mirror codes plus the fp8 pool's per-token scales; +an nvfp4 pool keeps block scales they do not copy, so both must refuse it up front.""" + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.kvcache.host_prefix_tier import HostPrefixTier +from freetoken.kvcache.kv_host_offload import KVHostOffloader + + +def _kv_pool(block_scaled: bool): + return SimpleNamespace( + _kv_buffer=torch.zeros(2, 1, 3, 4, 1, 8, dtype=torch.uint8), # [2, L, P+1, page, H, D] + _scale_buffer=None, + _block_scale_buffer=torch.zeros(1) if block_scaled else None, + ) + + +def _state_pool(): + return SimpleNamespace( + conv_states=torch.zeros(2, 2, 4, 3), + recurrent_states=torch.zeros(2, 2, 2, 4), + slot_states={}, + ) + + +def test_kv_host_offloader_refuses_an_nvfp4_pool(): + with pytest.raises(ValueError, match="--kv-host-pages .* nvfp4"): + KVHostOffloader(_kv_pool(block_scaled=True), num_logical_pages=4, device=torch.device("cpu")) + + +def test_prefix_tier_refuses_kv_pages_on_an_nvfp4_pool(): + with pytest.raises(ValueError, match="FT_PREFIX_HOST .* nvfp4"): + HostPrefixTier(_state_pool(), kv_pool=_kv_pool(block_scaled=True), gdn_slots_host=2, + kv_budget_bytes=1 << 16, device="cpu") + + +def test_prefix_tier_gdn_snapshots_do_not_touch_the_kv_pool(): + # the snapshot-only tier (no KV budget) never reads the pool, nvfp4 or not + tier = HostPrefixTier(_state_pool(), kv_pool=_kv_pool(block_scaled=True), gdn_slots_host=2, + kv_budget_bytes=0, device="cpu") + assert tier.num_slots == 2 and not tier.has_kv From 797cdb8a8fbbceafebbed9b5b26086ae2b523d40 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 23 Sep 2026 19:57:03 -0400 Subject: [PATCH 25/73] test(scheduler): give the scheduler stubs the prefill-start map the metrics timing reads Upstream #504 times the prefill span in Scheduler._prefill_start; the stubs from #222 and #505's tests build the scheduler without it. Assisted-by: Claude Opus 5.5 --- tests/scheduler/test_abort_inflight_prefill.py | 1 + tests/scheduler/test_hybrid_cache_manager.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/scheduler/test_abort_inflight_prefill.py b/tests/scheduler/test_abort_inflight_prefill.py index 85c217859f..f2cbd6e0d2 100644 --- a/tests/scheduler/test_abort_inflight_prefill.py +++ b/tests/scheduler/test_abort_inflight_prefill.py @@ -63,6 +63,7 @@ def _setup(): config=SimpleNamespace(page_size=1), status_reporter=SimpleNamespace(report_batch=lambda *_, **__: None), send_result=sent.extend, + _prefill_start={}, _kv_usage_pages=cm.page_usage, _mamba_slot_usage=lambda: None, _swa_token_usage=lambda: None, diff --git a/tests/scheduler/test_hybrid_cache_manager.py b/tests/scheduler/test_hybrid_cache_manager.py index 282ab2ca43..07385d7046 100644 --- a/tests/scheduler/test_hybrid_cache_manager.py +++ b/tests/scheduler/test_hybrid_cache_manager.py @@ -213,7 +213,7 @@ def test_chunked_prefill_retains_resumable_snapshot(monkeypatch, full_chunks, ta assert torch.all(pool.conv_states[:, match.mamba_value] == states[expected]) # Finish/abort may occur before the prefill commit; use the real idempotent cleanup path. - stub = SimpleNamespace(cache_manager=cm, table_manager=tm) + stub = SimpleNamespace(cache_manager=cm, table_manager=tm, _prefill_start={}) Scheduler._free_req_resources(stub, final) Scheduler._free_req_resources(stub, final) match = cm.match_req(_pend(prompt.tolist() + [999])) From 3224ef8de28e3ccdac7e52b8da007ead15ce43f6 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 23 Sep 2026 19:59:17 -0400 Subject: [PATCH 26/73] test(engine): give the logprobs sampler fakes the logits-processing flag the sampler plan reads Fork #47's sampler plan asks each request whether it needs logits processing; the upstream #224 test fakes predate it. Assisted-by: Claude Opus 5.5 --- tests/engine/test_sample_logprobs.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/engine/test_sample_logprobs.py b/tests/engine/test_sample_logprobs.py index aabbdc5fe7..6bd204604d 100644 --- a/tests/engine/test_sample_logprobs.py +++ b/tests/engine/test_sample_logprobs.py @@ -137,8 +137,8 @@ def test_compute_logprobs_with_prepare_keeps_max_top_logprobs_clamped(monkeypatc batch = SimpleNamespace( reqs=[ - SimpleNamespace(sampling_params=SimpleNamespace(logprobs=True, top_logprobs=17, is_greedy=True)), - SimpleNamespace(sampling_params=SimpleNamespace(logprobs=False, top_logprobs=0, is_greedy=True)), + SimpleNamespace(sampling_params=SimpleNamespace(needs_logits_processing=False, logprobs=True, top_logprobs=17, is_greedy=True)), + SimpleNamespace(sampling_params=SimpleNamespace(needs_logits_processing=False, logprobs=False, top_logprobs=0, is_greedy=True)), ] ) args = sampler.prepare(batch) From b458a1c144100a0fc5b358653a8b495c00d6bc0b Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 23 Sep 2026 20:48:29 -0400 Subject: [PATCH 27/73] merge fork review fixes: 11 bugs found re-reviewing the 5080deploy stack Assisted-by: Claude Opus 5.5 # Conflicts: # python/freetoken/server/openai_api.py # tests/server/test_stream_cancellation.py --- python/freetoken/engine/sample.py | 4 +- python/freetoken/kvcache/host_prefix_tier.py | 53 +++++++++-------- python/freetoken/models/qwen4_exp/ple.py | 13 +++-- python/freetoken/scheduler/cache.py | 7 ++- python/freetoken/server/anthropic_api.py | 8 ++- python/freetoken/server/generation.py | 13 ++++- python/freetoken/server/openai_api.py | 15 +++-- python/freetoken/server/responses_api.py | 8 ++- python/freetoken/tokenizer/tokenize.py | 10 +++- python/freetoken/utils/mp.py | 4 +- tests/engine/test_logits_processors.py | 10 ++++ tests/kvcache/test_host_tier_live_pool.py | 53 +++++++++++++++++ .../kvcache/test_host_tier_rehydrate_exact.py | 39 +++++++++++++ tests/models/qwen4_exp/test_ple.py | 13 +++++ tests/server/test_openai_api.py | 20 +++++++ tests/server/test_sampling_validation.py | 23 ++++++++ tests/server/test_stream_cancellation.py | 58 +++++++++++++++++++ tests/tokenizer/test_tokenize.py | 18 ++++++ tests/utils/test_mp_large_frame.py | 13 +++++ 19 files changed, 338 insertions(+), 44 deletions(-) create mode 100644 tests/kvcache/test_host_tier_live_pool.py create mode 100644 tests/kvcache/test_host_tier_rehydrate_exact.py create mode 100644 tests/server/test_sampling_validation.py create mode 100644 tests/utils/test_mp_large_frame.py diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py index f1984f7576..bca85c56fa 100644 --- a/python/freetoken/engine/sample.py +++ b/python/freetoken/engine/sample.py @@ -187,7 +187,9 @@ def _plan(self, batch: Batch, params) -> LogitsPlan | None: hist.append(ids[prompt_len:].tolist() if penalized else []) if p.repetition_penalty != 1.0: want_prompt = True - prompts.append(ids[:prompt_len].tolist()) + # multimodal placeholders are pseudo-ids >= vocab_size (MM_PAD_SHIFT_VALUE); + # send them to the scratch column, not past the scatter target + prompts.append(ids[:prompt_len].clamp(max=pad).tolist()) else: prompts.append([]) plan.hist_ids = _padded_rows(hist, pad, self.device) diff --git a/python/freetoken/kvcache/host_prefix_tier.py b/python/freetoken/kvcache/host_prefix_tier.py index 3b2d2d5dc3..39c96dfd6d 100644 --- a/python/freetoken/kvcache/host_prefix_tier.py +++ b/python/freetoken/kvcache/host_prefix_tier.py @@ -59,13 +59,15 @@ def __init__(self, state_pool, kv_pool=None, *, gdn_slots_host: int = 64, tensors = {"conv": state_pool.conv_states, "rec": state_pool.recurrent_states} tensors.update(state_pool.slot_states) self._order = list(tensors) - self._banks: dict[str, tuple[torch.Tensor, HostBank]] = {} + # banks only: the pool tensors are re-read on every copy, since + # LinearStatePool.rebuild replaces them (a runtime cache rebuild) + self._banks: dict[str, HostBank] = {} for name, t in tensors.items(): # pool: [layers, slots, *rest] -> bank: [host_slots, layers, *rest] bank = HostBank((gdn_slots_host, t.shape[0], *t.shape[2:]), t.dtype) if self.device.type == "cuda" and torch.cuda.is_available(): bank.pin() - self._banks[name] = (t, bank) + self._banks[name] = bank self.num_slots = gdn_slots_host self._free = list(range(gdn_slots_host)) self._in_use: set[int] = set() @@ -102,40 +104,41 @@ def __init__(self, state_pool, kv_pool=None, *, gdn_slots_host: int = 64, torch.cuda.Stream(self.device) if self.device.type == "cuda" else None) # ---------------- GDN snapshot storage ---------------- + def _pairs(self): + """(live pool tensor [layers, slots, ...], host bank) per state tensor.""" + pool = self.pool + live = {"conv": pool.conv_states, "rec": pool.recurrent_states, **pool.slot_states} + return [(live[name], self._banks[name]) for name in self._order] + + def _copy(self, copies) -> None: + if self.copy_stream is None: + for dst, src in copies: + dst.copy_(src) + return + # Order the copies after the caller's stream: the scheduler stream (itself made to wait + # on the engine stream) or the engine stream. Overlap launches the next batch before + # the previous one is drained, so a just-freed slot may still be written by it. + self.copy_stream.wait_stream(torch.cuda.current_stream(self.device)) + with torch.cuda.stream(self.copy_stream): + for dst, src in copies: + dst.copy_(src, non_blocking=True) + self.copy_stream.synchronize() + def write_gdn(self, gpu_slot: int) -> int: """D2H do snapshot do gpu_slot pra um slot host livre. Retorna o host_slot. - Deve ser enfileirado ANTES de state_pool.free(gpu_slot): a ordem de stream do - engine garante leitura antes de qualquer reescrita do slot (ver spec).""" + Deve ser enfileirado ANTES de state_pool.free(gpu_slot).""" if not self._free: raise RuntimeError("HostPrefixTier: sem slots host GDN livres (chame evict_host)") hs = self._free.pop() self._in_use.add(hs) - if self.copy_stream is not None: - with torch.cuda.stream(self.copy_stream): - for name in self._order: - t, bank = self._banks[name] - bank.tensor[hs].copy_(t[:, gpu_slot], non_blocking=True) - self.copy_stream.synchronize() - else: - for name in self._order: - t, bank = self._banks[name] - bank.tensor[hs].copy_(t[:, gpu_slot]) + self._copy([(bank.tensor[hs], t[:, gpu_slot]) for t, bank in self._pairs()]) return hs def read_gdn(self, host_slot: int, dst_gpu_slot: int) -> None: """H2D host_slot -> dst_gpu_slot do pool (já alocado pelo caller).""" if host_slot not in self._in_use: raise KeyError(f"HostPrefixTier: slot host {host_slot} não está em uso") - if self.copy_stream is not None: - with torch.cuda.stream(self.copy_stream): - for name in self._order: - t, bank = self._banks[name] - t[:, dst_gpu_slot].copy_(bank.tensor[host_slot], non_blocking=True) - self.copy_stream.synchronize() - else: - for name in self._order: - t, bank = self._banks[name] - t[:, dst_gpu_slot].copy_(bank.tensor[host_slot]) + self._copy([(t[:, dst_gpu_slot], bank.tensor[host_slot]) for t, bank in self._pairs()]) def free_gdn(self, host_slot: int) -> None: self._in_use.discard(host_slot) @@ -198,6 +201,7 @@ def _go(): del staged_s if self.copy_stream is not None: + self.copy_stream.wait_stream(torch.cuda.current_stream(self.device)) with torch.cuda.stream(self.copy_stream): _go() self.copy_stream.synchronize() @@ -235,6 +239,7 @@ def _go(): del staged_s if self.copy_stream is not None: + self.copy_stream.wait_stream(torch.cuda.current_stream(self.device)) with torch.cuda.stream(self.copy_stream): _go() self.copy_stream.synchronize() diff --git a/python/freetoken/models/qwen4_exp/ple.py b/python/freetoken/models/qwen4_exp/ple.py index 284e023bdd..44973bc0a6 100644 --- a/python/freetoken/models/qwen4_exp/ple.py +++ b/python/freetoken/models/qwen4_exp/ple.py @@ -49,6 +49,10 @@ _FUSED_HASH_ENV = "FREETOKEN_PLE_FUSED_HASH" +def _capturing(device: torch.device) -> bool: + return device.type == "cuda" and torch.cuda.is_current_stream_capturing() + + def _fused_row_ids_enabled() -> bool: """The fused hash kernel, on unless ``FREETOKEN_PLE_FUSED_HASH=0`` takes it back to torch ops.""" return (os.getenv(_FUSED_HASH_ENV) or "1").strip() not in ("0", "false", "False") @@ -482,12 +486,13 @@ def _token_index(self, meta: PLEMetadata) -> Tuple[torch.Tensor, torch.Tensor]: The fused kernel addresses the hash window through these instead of materializing the ``[B, ctx+max_len]`` packed window. Memoized on the shape (which is all they depend on) - so a captured replay reads a stable address instead of re-running the build; a build - that happens DURING capture is not cached, since its buffers live in the graph pool. + for eager forwards only: a capture always builds its own, which the graph recomputes + into its pool on every replay. A cached pair (the eager warmup forward right before + capture leaves one) would be baked in by address, then evicted and reused. """ device = meta.input_ids.device num_tokens = meta.input_ids.numel() - capturing = device.type == "cuda" and torch.cuda.is_current_stream_capturing() + capturing = _capturing(device) # is_decode is part of the key, not just the shape: a decode of B requests and a # prefill of ONE B-token request are the same [T] and mean opposite things (one token # per request at offset 0 vs B offsets inside one request). @@ -496,7 +501,7 @@ def _token_index(self, meta: PLEMetadata) -> Tuple[torch.Tensor, torch.Tensor]: (num_tokens,) if meta.is_decode else tuple(meta.seq_lens), str(device), ) - cached = self._token_index_cache.get(key) + cached = None if capturing else self._token_index_cache.get(key) if cached is not None: return cached if meta.is_decode: # one token per request, each at offset 0 diff --git a/python/freetoken/scheduler/cache.py b/python/freetoken/scheduler/cache.py index 7e4aa8e695..4e6b91df88 100644 --- a/python/freetoken/scheduler/cache.py +++ b/python/freetoken/scheduler/cache.py @@ -624,8 +624,11 @@ def _rehydrate_snapshot(self, node) -> int | None: prefix, copy the snapshot back into a fresh pool slot and return the slot; None = keep walking (the match truncates as before).""" tier = self.host_tier - hit = tier.lookup(self.prefix_cache._collect_key(node), self.page_size) - if hit is None or hit.gdn_slot is None: + key = self.prefix_cache._collect_key(node) + hit = tier.lookup(key, self.page_size) + # lookup is longest-prefix; a shorter entry's state would resume this node's full + # depth with the tokens in between missing from the recurrence + if hit is None or hit.gdn_slot is None or hit.host_len != len(key): return None pool = self.linear_state_pool if pool.num_free_slots < 1: diff --git a/python/freetoken/server/anthropic_api.py b/python/freetoken/server/anthropic_api.py index 385c87f3f1..4cb4b34822 100644 --- a/python/freetoken/server/anthropic_api.py +++ b/python/freetoken/server/anthropic_api.py @@ -55,6 +55,7 @@ submit_generation, with_keepalive, ) +from .openai_api import _await_watching_disconnect from .request_logger import log_request # Emit a protocol-native `ping` event after this many seconds of stream silence, @@ -141,9 +142,14 @@ async def handle_anthropic_messages( return StreamingResponse(events, media_type="text/event-stream") try: - result = await generate_full(uid, spec, state, source="/v1/messages") + # an abandoned request must not keep decoding to max_tokens (#222's watcher) + result = await _await_watching_disconnect( + generate_full(uid, spec, state, source="/v1/messages"), request, state, [uid] + ) except GenerationError as exc: return _anthropic_error_response(400, "invalid_request_error", str(exc)) + if result is None: + return _anthropic_error_response(499, "api_error", "client disconnected before the response was ready") response = anthropic_full_response( result, req.model, uid, cache_report=cache_report, metrics=metrics ) diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index e7d2b970a7..3db3913909 100644 --- a/python/freetoken/server/generation.py +++ b/python/freetoken/server/generation.py @@ -14,6 +14,7 @@ from __future__ import annotations import asyncio +import math import json import time from collections.abc import AsyncIterator @@ -231,8 +232,9 @@ def pick(value, key, framework): for name, value in (("presence_penalty", presence_penalty), ("frequency_penalty", frequency_penalty)): if not -2.0 <= float(value) <= 2.0: raise ValueError(f"{name} must be in [-2, 2], got {value}") - if resolved_rep <= 0.0: - raise ValueError(f"repetition_penalty must be positive, got {resolved_rep}") + # not `<= 0`: NaN compares False and would poison every seen token's logit + if not (math.isfinite(resolved_rep) and resolved_rep > 0.0): + raise ValueError(f"repetition_penalty must be a positive finite number, got {resolved_rep}") if min_tokens < 0: raise ValueError(f"min_tokens must be >= 0, got {min_tokens}") bias: list[list[float]] | None = None @@ -487,7 +489,12 @@ def _make_reasoning_parser(spec: GenSpec, state: Any) -> ReasoningParser | None: # The qwen3 chat template opens an implicit (thinking on) unless # enable_thinking is explicitly false, so the model emits only the closing # . Mirror that default here, else the chain-of-thought leaks into content. - force_reasoning = (spec.chat_template_kwargs or {}).get("enable_thinking") is not False + # continue_final_message is the exception: the template renders the final assistant + # turn as ... + content, so the continuation starts in content. + ctk = spec.chat_template_kwargs or {} + force_reasoning = ctk.get("enable_thinking") is not False and not ctk.get( + "continue_final_message" + ) elif parser_name == "glm": # GLM's template honors enable_thinking (default on) even with tools; the # generic fallback would force thinking and mislabel disabled output as reasoning. diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index b67825bfaf..9ce63060ab 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -401,14 +401,17 @@ async def handle_chat_completion( chunks = state.stream_with_cancellation(chunks, request, uids) return StreamingResponse(chunks, media_type="text/event-stream") + tasks = [ + asyncio.ensure_future(generate_full(u, spec, state, source="/v1/chat/completions")) + for u in uids + ] try: - results = await _await_watching_disconnect( - asyncio.gather(*(generate_full(u, spec, state, source="/v1/chat/completions") for u in uids)), - request, - state, - uids, - ) + results = await _await_watching_disconnect(asyncio.gather(*tasks), request, state, uids) except GenerationError as exc: + # gather does not cancel the other choices; the aborts below drop the events they + # wait on, so without the cancel they would stay pending forever + for task in tasks: + task.cancel() for u in uids: await state.abort_user(u) return create_error_response(str(exc), code=exc.code) diff --git a/python/freetoken/server/responses_api.py b/python/freetoken/server/responses_api.py index eea41722bb..9fd063a992 100644 --- a/python/freetoken/server/responses_api.py +++ b/python/freetoken/server/responses_api.py @@ -79,6 +79,7 @@ submit_generation, with_keepalive, ) +from .openai_api import _await_watching_disconnect from .request_logger import log_request # Seconds of event silence before a keep-alive frame is emitted on the stream. @@ -177,9 +178,14 @@ async def handle_responses( return StreamingResponse(events, media_type="text/event-stream") try: - result = await generate_full(uid, spec, state, source="/v1/responses") + # an abandoned request must not keep decoding to max_tokens (#222's watcher) + result = await _await_watching_disconnect( + generate_full(uid, spec, state, source="/v1/responses"), request, state, [uid] + ) except GenerationError as exc: return _error_response(400, str(exc), exc.code) + if result is None: + return _error_response(499, "client disconnected before the response was ready") response = build_responses_response( result, req, response_id, created, cache_report=cache_report, metrics=metrics ) diff --git a/python/freetoken/tokenizer/tokenize.py b/python/freetoken/tokenizer/tokenize.py index f512930dee..9648974dab 100644 --- a/python/freetoken/tokenizer/tokenize.py +++ b/python/freetoken/tokenizer/tokenize.py @@ -67,7 +67,15 @@ def _map_developer_role( m.get("role") == "developer" for m in messages ): return messages - return [{**m, "role": "system"} if m.get("role") == "developer" else m for m in messages] + mapped = [{**m, "role": "system"} if m.get("role") == "developer" else m for m in messages] + # The frontend hoisted and merged system messages before this mapping ran (templates such + # as Qwen's raise on a system message that is not first), so do it again for the new ones. + system = [m for m in mapped if m.get("role") == "system"] + if len(system) == 1 and mapped[0] is system[0]: + return mapped + rest = [m for m in mapped if m.get("role") != "system"] + content = "\n\n".join(str(m.get("content") or "") for m in system if m.get("content")) + return [{"role": "system", "content": content}] + rest class TokenizeManager: diff --git a/python/freetoken/utils/mp.py b/python/freetoken/utils/mp.py index f91498fbb2..ee35c057c1 100644 --- a/python/freetoken/utils/mp.py +++ b/python/freetoken/utils/mp.py @@ -22,7 +22,9 @@ class _CoalescedUnpacker: """ def __init__(self) -> None: - self._unpacker = msgpack.Unpacker(raw=False) + # 0 = effectively uncapped (2**31-1), as unpackb had; the Unpacker default of 100 MiB would + # kill the worker on a large message (e.g. an image's pixel values) + self._unpacker = msgpack.Unpacker(raw=False, max_buffer_size=0) self._pending: list[Any] = [] def feed(self, frame: bytes) -> None: diff --git a/tests/engine/test_logits_processors.py b/tests/engine/test_logits_processors.py index f051dec51f..a9645f4210 100644 --- a/tests/engine/test_logits_processors.py +++ b/tests/engine/test_logits_processors.py @@ -153,3 +153,13 @@ def test_prepare_builds_no_plan_without_processors_and_a_plan_with_them(): assert torch.isfinite(out[3]).all() # greedy path: argmax over the processed logits assert sampler.sample(logits, args)[1].item() == 5 + + +def test_repetition_penalty_ignores_multimodal_placeholder_ids_in_the_prompt(): + # image tokens carry pseudo-ids >= MM_PAD_SHIFT_VALUE; scattering them into the + # [rows, V+1] prompt mask indexed out of bounds (a device-side assert on CUDA) + sampler = Sampler(CPU, V) + batch = SimpleNamespace(reqs=[_req([1, 1_000_000, 2], [3], 4, repetition_penalty=2.0)]) + args = sampler.prepare(batch) + out = apply_logits_processors(torch.ones(1, V), args.plan, V) + assert out[0, 1] == 0.5 and out[0, 2] == 0.5 and out[0, 4] == 1.0 diff --git a/tests/kvcache/test_host_tier_live_pool.py b/tests/kvcache/test_host_tier_live_pool.py new file mode 100644 index 0000000000..bd12836893 --- /dev/null +++ b/tests/kvcache/test_host_tier_live_pool.py @@ -0,0 +1,53 @@ +"""The GDN snapshot tier (upstream #525) must copy against the pool's CURRENT tensors, which +LinearStatePool.rebuild replaces, and must order its copies after the engine stream.""" + +import contextlib +from types import SimpleNamespace + +import torch + +import freetoken.kvcache.host_prefix_tier as tier_module +from freetoken.kvcache.host_prefix_tier import HostPrefixTier + + +def _pool(slots: int): + return SimpleNamespace( + conv_states=torch.zeros(2, slots, 4, 3), + recurrent_states=torch.zeros(2, slots, 2, 4), + slot_states={}, + ) + + +def test_read_after_a_pool_rebuild_lands_in_the_new_tensors(): + pool = _pool(3) + tier = HostPrefixTier(pool, gdn_slots_host=2, device="cpu") + pool.conv_states[:, 1] = 5.0 + pool.recurrent_states[:, 1] = 7.0 + host = tier.write_gdn(1) + # a runtime rebuild swaps in fresh (larger) state tensors + pool.conv_states = torch.zeros(2, 6, 4, 3) + pool.recurrent_states = torch.zeros(2, 6, 2, 4) + tier.read_gdn(host, 5) + assert torch.all(pool.conv_states[:, 5] == 5.0) + assert torch.all(pool.recurrent_states[:, 5] == 7.0) + + +def test_copies_wait_for_the_engine_stream(monkeypatch): + events = [] + engine = object() + + class FakeStream: + def wait_stream(self, other): + events.append(("wait", other)) + + def synchronize(self): + events.append(("sync",)) + + monkeypatch.setattr(tier_module.torch.cuda, "current_stream", lambda device=None: engine) + monkeypatch.setattr(tier_module.torch.cuda, "stream", lambda s: contextlib.nullcontext()) + pool = _pool(3) + tier = HostPrefixTier(pool, gdn_slots_host=2, device="cpu") + tier.copy_stream = FakeStream() + host = tier.write_gdn(1) + tier.read_gdn(host, 2) + assert events == [("wait", engine), ("sync",), ("wait", engine), ("sync",)] diff --git a/tests/kvcache/test_host_tier_rehydrate_exact.py b/tests/kvcache/test_host_tier_rehydrate_exact.py new file mode 100644 index 0000000000..f67402041f --- /dev/null +++ b/tests/kvcache/test_host_tier_rehydrate_exact.py @@ -0,0 +1,39 @@ +"""The GDN snapshot tier (upstream #525) may only revive a tombstone from a snapshot of +exactly that node's prefix: the host index lookup is longest-prefix, and a shorter +snapshot would resume the recurrence with the tokens in between silently missing.""" + +from types import SimpleNamespace + +import torch + +from freetoken.kvcache.host_prefix_tier import HostMatch +from freetoken.scheduler.cache import CacheManager + + +def _manager(host_len: int, key_len: int = 128): + reads = [] + tier = SimpleNamespace( + lookup=lambda key, page_size: HostMatch(host_len, None, 7), + read_gdn=lambda host_slot, gpu_slot: reads.append((host_slot, gpu_slot)), + ) + pool = SimpleNamespace(num_free_slots=4, alloc=lambda n: [3]) + manager = SimpleNamespace( + host_tier=tier, + prefix_cache=SimpleNamespace(_collect_key=lambda node: torch.arange(key_len)), + page_size=64, + linear_state_pool=pool, + ensure_mamba_slots=lambda n: None, + ) + return manager, reads + + +def test_a_shorter_host_snapshot_does_not_revive_a_deeper_node(): + manager, reads = _manager(host_len=64) + assert CacheManager._rehydrate_snapshot(manager, node=object()) is None + assert reads == [] + + +def test_an_exact_host_snapshot_revives_the_node(): + manager, reads = _manager(host_len=128) + assert CacheManager._rehydrate_snapshot(manager, node=object()) == 3 + assert reads == [(7, 3)] diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index 74d4fd614e..02abaee1f4 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -201,6 +201,19 @@ def test_token_index_cache_is_keyed_by_shape(): assert embedding._token_index(other)[0] is not first[0] +def test_token_index_capture_never_reuses_an_evictable_cache_entry(monkeypatch): + """The eager warmup forward caches the decode pair right before capture; a capture that + reused it would bake an address the FIFO later evicts while the graph still reads it.""" + embedding = _make_layer(_config()).ple_embedding + decode = _meta([[5]], [[EOS, EOS]], decode=True) + eager = embedding._token_index(decode) # the warmup forward + monkeypatch.setattr(ple_module, "_capturing", lambda device: True) + captured = embedding._token_index(decode) + assert captured[0] is not eager[0] and captured[1] is not eager[1] + assert all(v[0] is not captured[0] for v in embedding._token_index_cache.values()) + assert captured[0].tolist() == [0] and captured[1].tolist() == [0] + + def test_token_index_does_not_confuse_a_decode_with_a_one_request_prefill(): """Same [T], opposite meaning: B decode rows at offset 0 vs B offsets in one request.""" layer = _make_layer(_config()) diff --git a/tests/server/test_openai_api.py b/tests/server/test_openai_api.py index ce675a2508..7d5b3cd114 100644 --- a/tests/server/test_openai_api.py +++ b/tests/server/test_openai_api.py @@ -223,6 +223,26 @@ def test_glm_reasoning_parser_honors_disabled_thinking_with_tools(): assert parser is not None and parser.detector.force_reasoning is True +def test_qwen3_continuation_of_a_final_assistant_message_is_content_not_reasoning(): + # The Qwen template closes the think block before the assistant's content, so a + # continue_final_message continuation never emits ; starting the parser inside + # reasoning filed the whole continuation under reasoning_content. + from freetoken.server.generation import _make_reasoning_parser + + state = FakeState([], reasoning_parser="qwen3") + req = chat_request( + messages=[ + {"role": "user", "content": "Count to five."}, + {"role": "assistant", "content": "One, two,"}, + ], + continue_final_message=True, + ) + parser = _make_reasoning_parser(chat_request_to_genspec(req, {}), state) + assert parser is not None and parser.detector.force_reasoning is False + plain = _make_reasoning_parser(chat_request_to_genspec(chat_request(), {}), state) + assert plain.detector.force_reasoning is True + + def test_non_stream_chat_completion_returns_openai_tool_calls_and_sends_tools(): output = '[TOOL_CALLS] [{"name":"get_weather","arguments":{"city":"Paris"}}]' state = FakeState( diff --git a/tests/server/test_sampling_validation.py b/tests/server/test_sampling_validation.py new file mode 100644 index 0000000000..3b2e1d6473 --- /dev/null +++ b/tests/server/test_sampling_validation.py @@ -0,0 +1,23 @@ +"""resolve_sampling rejects non-finite penalties: NaN passes a `<= 0` check and would +turn every already-seen token's logit into NaN.""" + +import pytest + +from freetoken.server.generation import resolve_sampling + + +def _resolve(**extra): + return resolve_sampling( + temperature=None, top_k=None, top_p=None, max_tokens=8, ignore_eos=False, + model_sampling={}, **extra, + ) + + +@pytest.mark.parametrize("value", [float("nan"), float("inf"), 0.0, -1.0]) +def test_repetition_penalty_must_be_positive_and_finite(value): + with pytest.raises(ValueError, match="repetition_penalty"): + _resolve(repetition_penalty=value) + + +def test_a_normal_repetition_penalty_is_kept(): + assert _resolve(repetition_penalty=1.1).repetition_penalty == pytest.approx(1.1) diff --git a/tests/server/test_stream_cancellation.py b/tests/server/test_stream_cancellation.py index ba7da7c17b..71d5ac8b85 100644 --- a/tests/server/test_stream_cancellation.py +++ b/tests/server/test_stream_cancellation.py @@ -316,3 +316,61 @@ def new_user(self): assert resp.status_code == 499 assert sorted(state.aborted) == [8, 9] + + +def test_a_failed_choice_cancels_the_other_choices(monkeypatch): + # n > 1: one choice errors while another is still generating. gather passes the error up + # without cancelling the sibling, and the aborts drop the event it waits on. + monkeypatch.setattr(openai_api, "_DISCONNECT_POLL_SECONDS", 0.01) + + class _State(_ApiState): + def __init__(self): + super().__init__(acks=None) + self._next = 7 + + def new_user(self): + self._next += 1 + return self._next + + async def wait_for_ack(self, uid): + if uid == 8: + yield SimpleNamespace(**{**vars(_ack("")), "error": "boom"}) + return + await asyncio.sleep(3600) + yield _ack("never") + + async def scenario(): + state = _State() + resp = await openai_api.handle_chat_completion( + _chat_req(n=2), _Request(disconnected=False), state, {} + ) + await asyncio.sleep(0) + pending = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()] + return resp, state, pending + + resp, state, pending = asyncio.run(scenario()) + assert resp.status_code == 400 + assert sorted(state.aborted) == [8, 9] + assert pending == [] + + +def test_messages_and_responses_non_stream_disconnect_delivers_abort(monkeypatch): + # #222 only wrapped the OpenAI handlers; the Anthropic and Responses endpoints kept + # decoding an abandoned non-streaming request to max_tokens. + from freetoken.server import anthropic_api, responses_api + from freetoken.server.anthropic_models import AnthropicMessagesRequest + from freetoken.server.responses_api import ResponsesRequest + + monkeypatch.setattr(openai_api, "_DISCONNECT_POLL_SECONDS", 0.01) + + state = _ApiState(acks=None) + req = AnthropicMessagesRequest(model="m", max_tokens=8, messages=[{"role": "user", "content": "hi"}]) + handler = anthropic_api.handle_anthropic_messages(req, _Request(disconnected=True), state, {}) + resp = asyncio.run(asyncio.wait_for(handler, timeout=5)) # unwatched: hangs, then times out + assert resp.status_code == 499 and state.aborted == [7] + + state = _ApiState(acks=None) + req = ResponsesRequest(model="m", input="hi", max_output_tokens=8) + handler = responses_api.handle_responses(req, _Request(disconnected=True), state, {}) + resp = asyncio.run(asyncio.wait_for(handler, timeout=5)) + assert resp.status_code == 499 and state.aborted == [7] diff --git a/tests/tokenizer/test_tokenize.py b/tests/tokenizer/test_tokenize.py index c82a76861c..390fcd4cce 100644 --- a/tests/tokenizer/test_tokenize.py +++ b/tests/tokenizer/test_tokenize.py @@ -827,3 +827,21 @@ def test_developer_role_maps_to_system_unless_the_template_knows_it(): assert mapped[0]["content"] == "be brief" and msgs[0]["role"] == "developer" # input untouched assert _map_developer_role(msgs, "{% if message.role == 'developer' %}") is msgs assert _map_developer_role(msgs[1:], None) == msgs[1:] # nothing to map + + +def test_mapped_developer_messages_are_merged_into_one_leading_system_message(): + """The frontend hoists system messages before the tokenizer maps developer -> system; a + template that rejects a non-first system message must still get exactly one, first.""" + from freetoken.tokenizer.tokenize import _map_developer_role + + template = "{% if message.role == 'system' %}" + msgs = [ + {"role": "system", "content": "be brief"}, + {"role": "developer", "content": "answer in German"}, + {"role": "user", "content": "hi"}, + ] + mapped = _map_developer_role(msgs, template) + assert [m["role"] for m in mapped] == ["system", "user"] + assert mapped[0]["content"] == "be brief\n\nanswer in German" + late = [{"role": "user", "content": "hi"}, {"role": "developer", "content": "be brief"}] + assert [m["role"] for m in _map_developer_role(late, template)] == ["system", "user"] diff --git a/tests/utils/test_mp_large_frame.py b/tests/utils/test_mp_large_frame.py new file mode 100644 index 0000000000..f91ee5604d --- /dev/null +++ b/tests/utils/test_mp_large_frame.py @@ -0,0 +1,13 @@ +"""The coalesced-frame decoder (#466) must accept a message larger than msgpack's +Unpacker default buffer (100 MiB), which unpackb never limited.""" + +import msgpack + +from freetoken.utils.mp import _CoalescedUnpacker + + +def test_a_message_over_100_mib_decodes(): + payload = {"pixels": b"x" * (101 * 1024 * 1024)} + unpacker = _CoalescedUnpacker() + unpacker.feed(msgpack.packb(payload)) + assert len(unpacker) == 1 and len(unpacker.take()["pixels"]) == len(payload["pixels"]) From 22851b9c84d09418e7d96cb78a9cbbf796a3b1bb Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 23 Sep 2026 23:52:44 -0400 Subject: [PATCH 28/73] merge fork fix (issue #65): run the PLE graph fill on its own thread, submitted before the launch # Conflicts: # python/freetoken/models/qwen4_exp/ple_disk.py --- python/freetoken/models/qwen4_exp/ple_disk.py | 126 ++++++++++++---- .../qwen4_exp/test_ple_disk_fill_order.py | 136 ++++++++++++++++++ 2 files changed, 237 insertions(+), 25 deletions(-) create mode 100644 tests/models/qwen4_exp/test_ple_disk_fill_order.py diff --git a/python/freetoken/models/qwen4_exp/ple_disk.py b/python/freetoken/models/qwen4_exp/ple_disk.py index 69ee9fa93a..d2eaa99df8 100644 --- a/python/freetoken/models/qwen4_exp/ple_disk.py +++ b/python/freetoken/models/qwen4_exp/ple_disk.py @@ -6,6 +6,9 @@ from __future__ import annotations import os +import queue +import threading +from concurrent.futures import Future from contextlib import contextmanager from dataclasses import dataclass from typing import Sequence @@ -100,6 +103,52 @@ def resolve_row_source(folder: str) -> PleRowSource: return source_from_safetensors(folder) +class _Filler: + """The one thread that drives the PleStore (engine-thread-only C++, no locks), FIFO. + + Graph fills run here, submitted before the launch: the fill releases the captured decode's + WAIT, and cuGraphLaunch may block until the GPU drains, so it must not wait for the launch + to return (a large graph deadlocked that way; fork issue #65).""" + + def __init__(self) -> None: + self._jobs: queue.SimpleQueue = queue.SimpleQueue() + self._thread = threading.Thread(target=self._run, name="ple-filler", daemon=True) + self._thread.start() + + def submit(self, fn) -> Future: + fut: Future = Future() + self._jobs.put((fn, fut)) + return fut + + def _run(self) -> None: + while (item := self._jobs.get()) is not None: + fn, fut = item + try: + fut.set_result(fn()) + except BaseException as exc: # noqa: BLE001 -- surfaced to the engine thread + fut.set_exception(exc) + + def close(self) -> None: + self._jobs.put(None) + self._thread.join(timeout=5) + + +class _PendingFill: + """A graph fill submitted before its launch; ``cancel`` after a failed launch.""" + + def __init__(self, table: "DiskRowTable", future: Future, cancelled: threading.Event) -> None: + self._table, self._future, self._cancelled = table, future, cancelled + + def cancel(self) -> None: + # no WAIT consumed the signal: clear it, or the next step's WAIT passes on stale rows + self._cancelled.set() + try: + self._future.result() + except BaseException: # noqa: BLE001 -- the launch error is the one to surface + pass + self._table._flag.zero_() + + class DiskRowTable: """``PLETableBackend`` whose rows are read from disk per fill (--ple-backend disk).""" @@ -150,9 +199,6 @@ def __init__( self._graph_dev = torch.empty( max_graph_rows * self._token_bytes, dtype=torch.uint8, device=self._device ) - # completes once the last graph launch that read ``_graph_pinned`` has run (recorded by the - # dispatch context after the step's fill: a record inside the capture would not be observable) - self._graph_consumed = torch.cuda.Event() eager_bytes = max_extend_tokens * self._token_bytes # the overlap scheduler fills batch k+1 while batch k's lookup copy may still be queued: the # eager staging alternates two pinned buffers, and each is rewritten only once the H2D copy @@ -177,10 +223,36 @@ def __init__( # ---------------- host side (engine thread, before the forward launches) ---------------- + def _filler(self) -> _Filler: + filler = self.__dict__.get("_filler_thread") + if filler is None: + filler = self._filler_thread = _Filler() + return filler + + def close(self) -> None: + filler = self.__dict__.pop("_filler_thread", None) + if filler is not None: + filler.close() + + def _current_stream(self): + return torch.cuda.current_stream(self._device) + + def _raise_failed_fill(self) -> None: + pending = self.__dict__.pop("_pending_fill", None) + if pending is not None and pending.done() and pending.exception() is not None: + raise pending.exception() + if pending is not None and not pending.done(): + self._pending_fill = pending + + def _fill_now(self, runs: Sequence[torch.Tensor], *, graph: bool) -> None: + """A fill that completes before its launch, on the store's thread (after any queued graph fill).""" + self._filler().submit(lambda: self.fill(runs, graph=graph)).result() + def fill(self, runs: Sequence[torch.Tensor], *, graph: bool) -> None: """Stage per-request token runs (two context ids, then the new tokens) in batch order.""" if graph: - self._graph_consumed.synchronize() + # no launch fence: the job already waited on the readback recorded after the last graph, + # and a fence recorded after this step's launch would wait on the WAIT this fill releases pinned = self._graph_pinned else: self._eager_slot ^= 1 @@ -202,22 +274,27 @@ def _ple_context(self, ids: torch.Tensor, position: int) -> list[int]: return _context(ids, position, self.eos_token_id) return [self.image_token_id if t >= MM_PAD_SHIFT_VALUE else t for t in _context(ids, position, self.eos_token_id)] - def host_fill_batch(self, batch: Batch, use_graph: bool): - """Stage this batch's rows; returns the post-dispatch fill callable under flag-sync, else None.""" - eos = self.eos_token_id + def host_fill_batch(self, batch: Batch, use_graph: bool) -> _PendingFill | None: + """Stage this batch's rows before the dispatch. Under flag-sync a graph fill is only submitted: + it completes on the filler thread, independently of the launch; returns its handle, else None.""" + self._raise_failed_fill() if batch.is_decode: reqs = list(batch.reqs) + # the context comes from the request as it is NOW: the engine mutates it once the launch returns + contexts = [self._ple_context(r.input_ids, r.device_len - 1) for r in reqs] if use_graph and self._wait_sync: bs = batch.padded_size self._token_readback[:bs].copy_(batch.input_ids, non_blocking=True) - self._readback_event.record(torch.cuda.current_stream(self._device)) + self._readback_event.record(self._current_stream()) + cancelled = threading.Event() - def _complete() -> None: + def _job() -> None: try: self._readback_event.synchronize() tokens = self._token_readback[:bs].to(torch.int64).tolist() - runs = [torch.tensor([*self._ple_context(r.input_ids, r.device_len - 1), t], dtype=torch.int64) - for r, t in zip(reqs, tokens)] + if cancelled.is_set(): + return + runs = [torch.tensor([*ctx, t], dtype=torch.int64) for ctx, t in zip(contexts, tokens)] self.fill(runs, graph=True) except BaseException: from freetoken.kernel.row_store import signal @@ -226,12 +303,12 @@ def _complete() -> None: signal(self._flag) raise - return _complete + self._pending_fill = self._filler().submit(_job) + return _PendingFill(self, self._pending_fill, cancelled) # launch-gating: this D2H is the step's readback and orders the fill after sampling tokens = batch.input_ids.to("cpu").to(torch.int64).tolist() - runs = [torch.tensor([*self._ple_context(r.input_ids, r.device_len - 1), t], dtype=torch.int64) - for r, t in zip(reqs, tokens)] - self.fill(runs, graph=use_graph) + runs = [torch.tensor([*ctx, t], dtype=torch.int64) for ctx, t in zip(contexts, tokens)] + self._fill_now(runs, graph=use_graph) return None runs = [ torch.cat(( @@ -240,20 +317,19 @@ def _complete() -> None: )) for req in batch.padded_reqs ] - self.fill(runs, graph=False) + self._fill_now(runs, graph=False) return None @contextmanager def forward_host_ctx(self, batch: Batch, use_graph: bool): - """Around one dispatch: stage on enter, run the deferred fill+signal on exit.""" - deferred = self.host_fill_batch(batch, use_graph) - yield - # no try/finally: a failed launch leaves no WAIT pending, so the fill must not run - if deferred is not None: - deferred() - if use_graph: - # only after this step's fill: the next fill waits on this launch, not the fill on it - self._graph_consumed.record(torch.cuda.current_stream(self._device)) + """Around one dispatch: stage (or submit the graph fill) on enter; cancel a submitted fill if the launch fails.""" + pending = self.host_fill_batch(batch, use_graph) + try: + yield + except BaseException: + if pending is not None: + pending.cancel() + raise # ---------------- device side (PLETableBackend protocol) ---------------- diff --git a/tests/models/qwen4_exp/test_ple_disk_fill_order.py b/tests/models/qwen4_exp/test_ple_disk_fill_order.py new file mode 100644 index 0000000000..e7fe58f7e0 --- /dev/null +++ b/tests/models/qwen4_exp/test_ple_disk_fill_order.py @@ -0,0 +1,136 @@ +"""The PLE disk backend's wait-sync fill must not depend on the graph launch returning. + +The captured decode graph WAITs on ``_flag``; ``cuGraphLaunch`` may block until the GPU drains +(a large graph fills the launch queue), and the GPU is parked on that WAIT. So the fill that sets +the flag has to run on its own thread, started before the launch (fork issue #65). +These tests stand in for the GPU with host fakes: the "launch" blocks until the flag is set. +""" + +from __future__ import annotations + +import threading +import time +from types import SimpleNamespace + +import pytest +import torch + +import freetoken.models.qwen4_exp.ple_disk as ple_disk +from freetoken.models.qwen4_exp.ple_disk import DiskRowTable + +EOS = 9 + + +class _Store: + """Stands in for the C++ PleStore: records calls and the thread that made them.""" + + def __init__(self, flag: torch.Tensor): + self.flag = flag + self.calls: list[tuple[str, list[int] | None, str]] = [] + + def stage(self, tokens_addr, n, staging_addr): + tokens = ctypes_read(tokens_addr, n + 2) + self.calls.append(("stage", tokens, threading.current_thread().name)) + + def flush(self, signal_addr): + self.calls.append(("flush", None, threading.current_thread().name)) + if signal_addr: + self.flag[0] = 1 + + +def ctypes_read(addr: int, n: int) -> list[int]: + import ctypes + + return list((ctypes.c_int64 * n).from_address(addr)) + + +class _Event: + def record(self, stream=None): + pass + + def synchronize(self): + pass + + +def _table(monkeypatch) -> DiskRowTable: + # host-only: no CUDA stream behind the (faked) readback event + monkeypatch.setattr(ple_disk.torch.cuda, "current_stream", lambda *a, **k: None) + t = DiskRowTable.__new__(DiskRowTable) + t._device = None + t._wait_sync = True + t._flag = torch.zeros(1, dtype=torch.int64) + t._store = _Store(t._flag) + t._token_readback = torch.zeros(8, dtype=torch.int32) + t._readback_event = _Event() + t._graph_pinned = torch.zeros(64, dtype=torch.uint8) + t._eager_pinned = [torch.zeros(64, dtype=torch.uint8), torch.zeros(64, dtype=torch.uint8)] + t._eager_read = [_Event(), _Event()] + t._eager_slot = 0 + t._token_bytes = 1 + t.eos_token_id = EOS + t.image_token_id = None + return t + + +def _close(t: DiskRowTable) -> None: + close = getattr(t, "close", None) + if close is not None: + close() + + +def _decode_batch(prev_tokens: list[int], new_token: int): + req = SimpleNamespace(input_ids=torch.tensor(prev_tokens + [new_token]), device_len=len(prev_tokens) + 1) + return SimpleNamespace(is_decode=True, reqs=[req], padded_size=1, + input_ids=torch.tensor([new_token], dtype=torch.int32)) + + +def _blocking_launch(flag: torch.Tensor, timeout: float = 5.0) -> bool: + """A cuGraphLaunch that cannot return until the GPU passes the graph's WAIT(flag).""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if int(flag[0]) == 1: + flag[0] = 0 # the graph's WAIT-then-RESET + return True + time.sleep(0.001) + return False + + +def test_graph_fill_does_not_wait_for_the_launch_to_return(monkeypatch): + t = _table(monkeypatch) + with t.forward_host_ctx(_decode_batch([1, 2, 3], 7), use_graph=True): + assert _blocking_launch(t._flag), "deadlock: the fill only runs after the launch returns" + _close(t) + + +def test_fills_run_in_step_order_on_one_thread_with_contexts_taken_before_the_launch(monkeypatch): + t = _table(monkeypatch) + for prev, new in (([1, 2, 3], 7), ([1, 2, 3, 7], 8), ([1, 2, 3, 7, 8], 4)): + batch = _decode_batch(prev, new) + with t.forward_host_ctx(batch, use_graph=True): + # the engine mutates the request as soon as the launch returns (overlap scheduling) + batch.reqs[0].input_ids = torch.tensor([0, 0, 0, 0, 0, 0]) + batch.reqs[0].device_len += 1 + assert _blocking_launch(t._flag) + _close(t) + staged = [c for c in t._store.calls if c[0] == "stage"] + assert [c[1] for c in staged] == [[2, 3, 7], [3, 7, 8], [7, 8, 4]] + assert len({c[2] for c in t._store.calls}) == 1 # the store stays single-threaded + + +def test_eager_fill_completes_before_the_launch(monkeypatch): + t = _table(monkeypatch) + req = SimpleNamespace(input_ids=torch.tensor([1, 2, 3, 4]), cached_len=1, device_len=4) + batch = SimpleNamespace(is_decode=False, padded_reqs=[req]) + with t.forward_host_ctx(batch, use_graph=False): + assert [c[0] for c in t._store.calls] == ["stage", "flush"] + _close(t) + + +def test_a_failed_launch_leaves_no_stale_signal(monkeypatch): + t = _table(monkeypatch) + with pytest.raises(RuntimeError, match="launch failed"): + with t.forward_host_ctx(_decode_batch([1, 2, 3], 7), use_graph=True): + raise RuntimeError("launch failed") + # no WAIT consumed it, so a leftover 1 would let the next step read stale rows + assert int(t._flag[0]) == 0 + _close(t) From a3120104662d8def972b0ed2297dbdcbc15acdbc Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 23 Sep 2026 23:52:44 -0400 Subject: [PATCH 29/73] merge fork perf: load only the routed experts for short prefills (FREETOKEN_MOE_SMALL_PREFILL_TOKENS, off by default) --- python/freetoken/layers/moe.py | 19 +++++++ tests/moe/test_small_prefill_on_demand.py | 69 +++++++++++++++++++++++ 2 files changed, 88 insertions(+) create mode 100644 tests/moe/test_small_prefill_on_demand.py diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index 4918a032fe..22beb4fdd8 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -24,6 +24,9 @@ # default. Set FREETOKEN_HYBRID_OVERLAP=0 to force the serial path (CPU sync before the # GPU work) -- a measurement-only escape hatch to A/B the overlap benefit. _HYBRID_OVERLAP = os.getenv("FREETOKEN_HYBRID_OVERLAP", "1") != "0" +# Prefills of at most this many tokens load only their routed experts (the decode path's +# on-demand LRU) instead of streaming every expert of every layer; 0 = off. +_SMALL_PREFILL_TOKENS = int(os.getenv("FREETOKEN_MOE_SMALL_PREFILL_TOKENS", "0")) class MoELayer(BaseOP): @@ -361,6 +364,22 @@ def _prefill_routed( pass through unmapped.""" cache = self.offload_cache assert cache is not None + if 0 < hidden_states.shape[0] <= _SMALL_PREFILL_TOKENS and cache.decode_target in ("gpu", "hybrid"): + # A short extension (an agent turn over a cached prefix) touches a fraction of each + # layer's experts; fetching those beats streaming all of them (~68 GB here). The same + # token count holds in every layer, so the overlap double buffer is never begun. + cache.ensure_experts(self.layer_id, topk_ids) + cache.copy_missing() + return self._expert_gemm( + cache, + hidden_states, + topk_weights, + topk_ids, + views=cache.bank_views(), + n=None, + alphas=cache.alphas_for_slots(self.layer_id), + is_prefill=False, + ) if cache.prefill_overlap: views = self._wait_prefill_overlap(cache) out = self._expert_gemm( diff --git a/tests/moe/test_small_prefill_on_demand.py b/tests/moe/test_small_prefill_on_demand.py new file mode 100644 index 0000000000..5d984a0dec --- /dev/null +++ b/tests/moe/test_small_prefill_on_demand.py @@ -0,0 +1,69 @@ +"""Short prefills load only their routed experts (FREETOKEN_MOE_SMALL_PREFILL_TOKENS) instead of +streaming whole layers; above the threshold, or with it off, the whole-layer path is unchanged.""" + +import pytest +import torch + +import freetoken.layers.moe as moe_mod +from freetoken.distributed import set_tp_info, try_get_tp_info + + +def _layer_and_cache(device, num_experts=8, top_k=2, hidden=16, inter=32): + from freetoken.layers.moe import OffloadMoELayer + from freetoken.layers.quantization import NoQuantConfig + from freetoken.moe.offload_cache import OffloadMoeCache + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + layer = OffloadMoELayer(0, num_experts, top_k, hidden, inter, quant_config=NoQuantConfig(), + prefix="model.layers.0.mlp.experts") + cache = OffloadMoeCache(num_layers=1, num_experts=num_experts, cache_size=num_experts, device=device) + g = torch.Generator().manual_seed(0) + cache.set_bank_sources({ + "gate_up": [torch.randn(num_experts, 2 * inter, hidden, generator=g, dtype=torch.bfloat16) * 0.1], + "down": [torch.randn(num_experts, hidden, inter, generator=g, dtype=torch.bfloat16) * 0.1], + }) + layer.offload_cache = cache + return layer, cache + + +@pytest.mark.parametrize("tokens, threshold, on_demand", [ + (3, 8, True), (8, 8, True), (9, 8, False), (3, 0, False), +]) +def test_small_prefill_dispatch(monkeypatch, tokens, threshold, on_demand): + layer, cache = _layer_and_cache(torch.device("cpu")) + monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", threshold) + calls = [] + monkeypatch.setattr(cache, "ensure_experts", lambda layer_id, ids: calls.append("ensure")) + monkeypatch.setattr(cache, "materialize_layer", lambda layer_id: calls.append("materialize")) + monkeypatch.setattr(cache, "copy_missing", lambda: calls.append("copy")) + monkeypatch.setattr(cache, "begin_prefill", lambda: calls.append("begin_prefill"), raising=False) + got = {} + + def fake_gemm(cache_, hs, w, ids, *, views, n, alphas, is_prefill): + got.update(n=n, is_prefill=is_prefill) + return hs + + monkeypatch.setattr(layer, "_expert_gemm", fake_gemm) + hs = torch.randn(tokens, 16) + layer._prefill_routed(hs, torch.full((tokens, 2), 0.5), torch.zeros(tokens, 2, dtype=torch.int32)) + if on_demand: + assert calls == ["ensure", "copy"] and got == {"n": None, "is_prefill": False} + else: + assert "ensure" not in calls and "begin_prefill" not in calls[1:] and got["n"] == layer.num_experts + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_small_prefill_matches_the_whole_layer_path(monkeypatch): + from freetoken.core import get_global_ctx # noqa: F401 (layer reads the ctx only in forward()) + + dev = torch.device("cuda") + hs = torch.randn(5, 16, dtype=torch.bfloat16, device=dev) * 0.5 + ids = torch.tensor([[0, 3], [1, 3], [7, 2], [3, 4], [5, 0]], dtype=torch.int32, device=dev) + w = torch.tensor([[0.6, 0.4], [0.5, 0.5], [0.9, 0.1], [0.3, 0.7], [0.2, 0.8]], device=dev) + outs = {} + for threshold in (0, 8): + monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", threshold) + layer, cache = _layer_and_cache(dev) + outs[threshold] = layer._prefill_routed(hs.clone(), w.clone(), ids.clone()).float() + torch.testing.assert_close(outs[8], outs[0], rtol=2e-2, atol=2e-2) From 531244a384a800ad7968cb826083595a5738a73b Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 24 Sep 2026 00:27:13 -0400 Subject: [PATCH 30/73] merge fork fix: small-prefill on-demand path hands the LRU kernel unique ids under Triton's element limit --- python/freetoken/layers/moe.py | 49 +++++++++++++++++------ tests/moe/test_small_prefill_on_demand.py | 37 +++++++++++++++-- 2 files changed, 70 insertions(+), 16 deletions(-) diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index 22beb4fdd8..8161621775 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -27,6 +27,11 @@ # Prefills of at most this many tokens load only their routed experts (the decode path's # on-demand LRU) instead of streaming every expert of every layer; 0 = off. _SMALL_PREFILL_TOKENS = int(os.getenv("FREETOKEN_MOE_SMALL_PREFILL_TOKENS", "0")) +_TRITON_MAX_NUMEL = 1 << 20 # Triton's per-tensor element limit (the LRU kernel's block) + + +def triton_next_power_of_2(n: int) -> int: + return 1 << max(0, int(n) - 1).bit_length() class MoELayer(BaseOP): @@ -368,18 +373,18 @@ def _prefill_routed( # A short extension (an agent turn over a cached prefix) touches a fraction of each # layer's experts; fetching those beats streaming all of them (~68 GB here). The same # token count holds in every layer, so the overlap double buffer is never begun. - cache.ensure_experts(self.layer_id, topk_ids) - cache.copy_missing() - return self._expert_gemm( - cache, - hidden_states, - topk_weights, - topk_ids, - views=cache.bank_views(), - n=None, - alphas=cache.alphas_for_slots(self.layer_id), - is_prefill=False, - ) + slots = self._ensure_unique(cache, topk_ids) + if slots is not None: + return self._expert_gemm( + cache, + hidden_states, + topk_weights, + slots, + views=cache.bank_views(), + n=None, + alphas=cache.alphas_for_slots(self.layer_id), + is_prefill=False, + ) if cache.prefill_overlap: views = self._wait_prefill_overlap(cache) out = self._expert_gemm( @@ -407,6 +412,26 @@ def _prefill_routed( is_prefill=True, ) + def _ensure_unique(self, cache: OffloadMoeCache, topk_ids: torch.Tensor) -> torch.Tensor | None: + """Make this layer's routed experts resident and return ``topk_ids`` mapped to slots, or + None (take the whole-layer path) when they cannot all be resident at once. + + The LRU kernel is sized for decode: its block is next_pow2(ids) x next_pow2(slots), so a + prefill's tokens x top_k ids overflow Triton's element limit. It gets the unique ids + instead (at most num_experts), chunked under that limit; each chunk's misses are copied + before the next ensure reuses the copy plan.""" + uniq, inverse = torch.unique(topk_ids.reshape(-1), return_inverse=True) + n = uniq.numel() + if n > cache.cache_size // 2: # leave room so a chunk never evicts the one before it + return None + uniq = uniq.to(torch.int32) + per_call = max(1, _TRITON_MAX_NUMEL // triton_next_power_of_2(cache.cache_size)) + for start in range(0, n, per_call): + part = uniq[start : start + per_call] # a contiguous view: ensure rewrites it in place + cache.ensure_experts(self.layer_id, part) + cache.copy_missing() + return uniq[inverse].view_as(topk_ids) + def _wait_prefill_overlap(self, cache: OffloadMoeCache) -> tuple[torch.Tensor, ...]: """Double-buffer choreography for this layer's overlap prefill: kick off the next layer's full-layer H2D copy, then return this layer's bank views (in diff --git a/tests/moe/test_small_prefill_on_demand.py b/tests/moe/test_small_prefill_on_demand.py index 5d984a0dec..e7fb70721a 100644 --- a/tests/moe/test_small_prefill_on_demand.py +++ b/tests/moe/test_small_prefill_on_demand.py @@ -8,7 +8,7 @@ from freetoken.distributed import set_tp_info, try_get_tp_info -def _layer_and_cache(device, num_experts=8, top_k=2, hidden=16, inter=32): +def _layer_and_cache(device, num_experts=8, top_k=2, hidden=16, inter=32, cache_size=None): from freetoken.layers.moe import OffloadMoELayer from freetoken.layers.quantization import NoQuantConfig from freetoken.moe.offload_cache import OffloadMoeCache @@ -17,7 +17,7 @@ def _layer_and_cache(device, num_experts=8, top_k=2, hidden=16, inter=32): set_tp_info(rank=0, size=1) layer = OffloadMoELayer(0, num_experts, top_k, hidden, inter, quant_config=NoQuantConfig(), prefix="model.layers.0.mlp.experts") - cache = OffloadMoeCache(num_layers=1, num_experts=num_experts, cache_size=num_experts, device=device) + cache = OffloadMoeCache(num_layers=1, num_experts=num_experts, cache_size=cache_size or num_experts, device=device) g = torch.Generator().manual_seed(0) cache.set_bank_sources({ "gate_up": [torch.randn(num_experts, 2 * inter, hidden, generator=g, dtype=torch.bfloat16) * 0.1], @@ -64,6 +64,35 @@ def test_small_prefill_matches_the_whole_layer_path(monkeypatch): outs = {} for threshold in (0, 8): monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", threshold) - layer, cache = _layer_and_cache(dev) - outs[threshold] = layer._prefill_routed(hs.clone(), w.clone(), ids.clone()).float() + layer, cache = _layer_and_cache(dev, cache_size=32) + outs[threshold] = _run_and_check_path(layer, hs, w, ids, on_demand=threshold > 0) torch.testing.assert_close(outs[8], outs[0], rtol=2e-2, atol=2e-2) + + +def _run_and_check_path(layer, hs, w, ids, *, on_demand): + taken = [] + real = getattr(layer, "_ensure_unique", None) + if real is not None: + layer._ensure_unique = lambda cache, t: taken.append(r := real(cache, t)) or r + out = layer._prefill_routed(hs.clone(), w.clone(), ids.clone()).float() + assert (bool(taken) and taken[0] is not None) == on_demand + return out + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_a_real_sized_small_prefill_fits_the_lru_kernel(monkeypatch): + # 200 tokens x top-10 = 2000 routed ids against a cache that rounds up to 2048 slots: the + # LRU kernel's block is next_pow2(ids) x next_pow2(slots), so the raw ids overflow Triton's + # 1M-element limit; the path must hand it the (at most num_experts) unique ids instead. + dev = torch.device("cuda") + tokens, experts, top_k = 200, 512, 10 + g = torch.Generator(device=dev).manual_seed(1) + hs = torch.randn(tokens, 16, dtype=torch.bfloat16, device=dev, generator=g) * 0.5 + ids = torch.randint(0, experts, (tokens, top_k), dtype=torch.int32, device=dev, generator=g) + w = torch.softmax(torch.randn(tokens, top_k, device=dev, generator=g), dim=-1) + outs = {} + for threshold in (0, 1024): + monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", threshold) + layer, cache = _layer_and_cache(dev, num_experts=experts, top_k=top_k, cache_size=1100) + outs[threshold] = _run_and_check_path(layer, hs, w, ids, on_demand=threshold > 0) + torch.testing.assert_close(outs[1024], outs[0], rtol=2e-2, atol=2e-2) From 41b4444f8d2ba0b8e0f7cc9447594067132ab7cf Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 24 Sep 2026 00:50:11 -0400 Subject: [PATCH 31/73] merge fork fix: small-prefill ensures its experts in decode-width chunks --- python/freetoken/layers/moe.py | 19 ++++++++----------- tests/moe/test_small_prefill_on_demand.py | 6 +++++- 2 files changed, 13 insertions(+), 12 deletions(-) diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index 8161621775..799458dc06 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -27,11 +27,9 @@ # Prefills of at most this many tokens load only their routed experts (the decode path's # on-demand LRU) instead of streaming every expert of every layer; 0 = off. _SMALL_PREFILL_TOKENS = int(os.getenv("FREETOKEN_MOE_SMALL_PREFILL_TOKENS", "0")) -_TRITON_MAX_NUMEL = 1 << 20 # Triton's per-tensor element limit (the LRU kernel's block) - - -def triton_next_power_of_2(n: int) -> int: - return 1 << max(0, int(n) - 1).bit_length() +# ids per LRU ensure in that path: the kernel's block is next_pow2(ids) x next_pow2(slots), and a +# block much wider than decode's spills to local memory whose launch reservation can OOM +_ENSURE_CHUNK_IDS = 32 class MoELayer(BaseOP): @@ -417,17 +415,16 @@ def _ensure_unique(self, cache: OffloadMoeCache, topk_ids: torch.Tensor) -> torc None (take the whole-layer path) when they cannot all be resident at once. The LRU kernel is sized for decode: its block is next_pow2(ids) x next_pow2(slots), so a - prefill's tokens x top_k ids overflow Triton's element limit. It gets the unique ids - instead (at most num_experts), chunked under that limit; each chunk's misses are copied - before the next ensure reuses the copy plan.""" + prefill's tokens x top_k ids overflow Triton's element limit, and even a block under it + can fail to launch on a full GPU. It gets the unique ids (at most num_experts) in + decode-width chunks; each chunk's misses are copied before the next ensure reuses the plan.""" uniq, inverse = torch.unique(topk_ids.reshape(-1), return_inverse=True) n = uniq.numel() if n > cache.cache_size // 2: # leave room so a chunk never evicts the one before it return None uniq = uniq.to(torch.int32) - per_call = max(1, _TRITON_MAX_NUMEL // triton_next_power_of_2(cache.cache_size)) - for start in range(0, n, per_call): - part = uniq[start : start + per_call] # a contiguous view: ensure rewrites it in place + for start in range(0, n, _ENSURE_CHUNK_IDS): + part = uniq[start : start + _ENSURE_CHUNK_IDS] # a contiguous view: ensure rewrites it in place cache.ensure_experts(self.layer_id, part) cache.copy_missing() return uniq[inverse].view_as(topk_ids) diff --git a/tests/moe/test_small_prefill_on_demand.py b/tests/moe/test_small_prefill_on_demand.py index e7fb70721a..eea5842d23 100644 --- a/tests/moe/test_small_prefill_on_demand.py +++ b/tests/moe/test_small_prefill_on_demand.py @@ -83,7 +83,8 @@ def _run_and_check_path(layer, hs, w, ids, *, on_demand): def test_a_real_sized_small_prefill_fits_the_lru_kernel(monkeypatch): # 200 tokens x top-10 = 2000 routed ids against a cache that rounds up to 2048 slots: the # LRU kernel's block is next_pow2(ids) x next_pow2(slots), so the raw ids overflow Triton's - # 1M-element limit; the path must hand it the (at most num_experts) unique ids instead. + # 1M-element limit, and a 512-id block failed to launch on a full GPU (CUDA OOM). The path + # must hand it the unique ids in decode-width chunks. dev = torch.device("cuda") tokens, experts, top_k = 200, 512, 10 g = torch.Generator(device=dev).manual_seed(1) @@ -94,5 +95,8 @@ def test_a_real_sized_small_prefill_fits_the_lru_kernel(monkeypatch): for threshold in (0, 1024): monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", threshold) layer, cache = _layer_and_cache(dev, num_experts=experts, top_k=top_k, cache_size=1100) + widths, ensure = [], cache.ensure_experts + monkeypatch.setattr(cache, "ensure_experts", lambda lid, q: (widths.append(q.numel()), ensure(lid, q))[1]) outs[threshold] = _run_and_check_path(layer, hs, w, ids, on_demand=threshold > 0) + assert widths and max(widths) <= 32, f"LRU ensure widths {sorted(set(widths))}: wider than decode's" torch.testing.assert_close(outs[1024], outs[0], rtol=2e-2, atol=2e-2) From 3c73cea709fb947228e11770c8496974d8f85319 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Sat, 26 Sep 2026 16:14:49 -0400 Subject: [PATCH 32/73] merge fork perf: scored expert eviction (--moe-cache-policy), shared expert beside the miss copy (FREETOKEN_MOE_COPY_OVERLAP), opt-in slim copy --- benchmarks/README.md | 30 + benchmarks/bench_expert_copy_corun.py | 518 +++++++++++++++ benchmarks/bench_moe_copy_overlap.py | 233 +++++++ benchmarks/bench_scored_ensure.py | 158 +++++ python/freetoken/env.py | 2 + python/freetoken/kernel/aot.py | 19 + .../kernel/csrc/jit/fast_index_copy.cuh | 109 +++ python/freetoken/kernel/fast_index_copy.py | 42 ++ python/freetoken/layers/moe.py | 21 +- python/freetoken/models/deepseek_v4/moe.py | 2 +- python/freetoken/models/qwen4_exp/moe.py | 12 +- python/freetoken/moe/offload_cache.py | 171 ++++- python/freetoken/moe/offload_kernels.py | 45 ++ python/freetoken/moe/scored_ensure.py | 395 +++++++++++ python/freetoken/server/args.py | 9 +- tests/README.md | 1 + tests/models/qwen4_exp/test_skeleton.py | 124 ++++ tests/moe/ref_scored_cache.py | 159 +++++ tests/moe/test_fused_copy.py | 173 ++++- tests/moe/test_offload.py | 117 +++- tests/moe/test_scored_ensure.py | 622 ++++++++++++++++++ tests/moe/test_small_prefill_on_demand.py | 43 +- 22 files changed, 2972 insertions(+), 33 deletions(-) create mode 100644 benchmarks/bench_expert_copy_corun.py create mode 100644 benchmarks/bench_moe_copy_overlap.py create mode 100644 benchmarks/bench_scored_ensure.py create mode 100644 python/freetoken/moe/scored_ensure.py create mode 100644 tests/moe/ref_scored_cache.py create mode 100644 tests/moe/test_scored_ensure.py diff --git a/benchmarks/README.md b/benchmarks/README.md index 5cbb0b34e7..3f0df72f81 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -27,6 +27,36 @@ batch size x miss rate. python benchmarks/bench_offload_cache_copy.py ``` +**`bench_scored_ensure.py`** — synthetic (no checkpoint): us per decode `ensure` call at a +fixed miss count, flashlib `lru_ensure` vs the `--moe-cache-policy` kernels (lru, kd, kdfb, +rule), from a CUDA graph of two decode steps on a full cache. `--stats` times it with +`--moe-collect-stats` accumulation on. + +```bash +python benchmarks/bench_scored_ensure.py --slots 1650 --k 10,20 +``` + +**`bench_moe_copy_overlap.py`** — synthetic (no checkpoint): `FREETOKEN_MOE_COPY_OVERLAP` +off vs on for Qwen4ExpMoE decode layers at the Qwen3.8-Flash-Next expert geometry, as CUDA +graphs. Checks the outputs are bitwise identical, times replays ABAB, and profiles the copy's +start after the ensure and how much of the side-stream shared expert runs inside it. + +```bash +python benchmarks/bench_moe_copy_overlap.py --bs 1 --slots 160 +``` + +**`bench_expert_copy_corun.py`** — synthetic (no checkpoint): the fused expert copy kernel +alone and beside compute, on registered host banks with the qwen4_exp NVFP4 row sizes. Sweeps +`fast_index_copy_multi` (threads x blocks/bank) and `fast_index_copy_multi_slim` +(blocks x threads x unroll) over 1..10-expert plans (GB/s, "hot" back-to-back and "cold" +spacer-separated copies), then reports the stretch sigma of bf16/fp8 GEMVs, a small matmul and +the NVFP4 routed decode GEMV while a copy loop runs on another stream. `--json` keeps the raw +per-trial times. + +```bash +python benchmarks/bench_expert_copy_corun.py --trials 5 --json copy_corun.json +``` + For host RAM vs PCIe bandwidth and the offload/hybrid backend pick, use `ft bench bw` instead — it writes the JSON profile the engine reads. diff --git a/benchmarks/bench_expert_copy_corun.py b/benchmarks/bench_expert_copy_corun.py new file mode 100644 index 0000000000..b76b82ee6f --- /dev/null +++ b/benchmarks/bench_expert_copy_corun.py @@ -0,0 +1,518 @@ +"""Expert-copy kernel sweep and co-run stretch on real-size NVFP4 expert banks. + +Two measurements, no checkpoint needed: + +* sweep: GB/s of one fused expert copy (all 6 banks, 1..10 experts per plan) for each + ``fast_index_copy_multi`` (num_threads, blocks_per_bank) config, and for each + ``fast_index_copy_multi_slim`` (blocks, threads, unroll) config when that kernel exists; +* co-run: the stretch ``sigma = t_corun / t_alone - 1`` of decode compute kernels (bf16 and + fp8 GEMV [1,H]x[H,N], a small bf16 matmul, the NVFP4 routed decode GEMV) while a copy loop + runs on another stream, and the copy's GB/s while each compute loop runs. + +The host banks are registered anonymous mmaps (``HostBank``, like the serving banks) with +the qwen4_exp triton-NVFP4 row sizes. Source rows are not reused within a loop, the GEMV +weights are cycled over copies that together exceed L2, and every loop is a CUDA graph. +Alone and co-run trials are interleaved; each cell reports min and median over trials. + +Copy loops run in two modes. "hot": copy nodes back to back. "cold": a one-element spacer +kernel after every copy, as the ensure kernel precedes every decode copy. On the RTX 5080, in +graph replay, a copy node that follows another kernel's node runs ~8 us longer than one that +follows a copy node, for every variant (CUPTI puts it in the copy's own duration; eager +launches do not show it; cause not identified). Cold GB/s subtracts the spacer's alone time. + +Run: CUDA_VISIBLE_DEVICES=0 PYTHONPATH=python python benchmarks/bench_expert_copy_corun.py +""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import statistics +import tempfile +from dataclasses import dataclass +from typing import Callable + +import torch +import torch.nn.functional as F + +from freetoken.gpu_select import assign_gpu, bind_assigned_gpu, single_gpu_arg +from freetoken.kernel import fast_index_copy as fic +from freetoken.kernel.pinned import device_ptr +from freetoken.moe.host_banks import HostBank +from freetoken.moe.offload_cache import _BANK_BYTES_PER_EXPERT + +L2_BYTES = 64 << 20 +SLEEP_CYCLES = 2_000_000 # ~1 ms head start so the host enqueues every graph before the GPU starts +LMAX = 16 # plan buffer length; the kernels read the live count from num_indices + + +def nvfp4_bank_bytes(h: int, i: int) -> list[int]: + # triton NVFP4 layout() order: gate_up, gate_up_scale, gate_up_global, down, down_scale, down_global + return [2 * i * (h // 2), 2 * i * (h // 16), 2 * i * 2, h * (i // 2), h * (i // 16), h * 2] + + +@dataclass(frozen=True) +class CopyVariant: + name: str + launch: Callable[..., None] + + +def multi_variant(threads: int, blocks_per_bank: int) -> CopyVariant: + def launch(dst, src, feat, di, si, num): + fic.fast_index_copy_multi_jit( + dst, src, feat, di, si, num, num_threads=threads, blocks_per_bank=blocks_per_bank + ) + + return CopyVariant(f"multi {threads}x{blocks_per_bank}", launch) + + +def slim_variant(blocks: int, threads: int, unroll: int) -> CopyVariant: + def launch(dst, src, feat, di, si, num): + fic.fast_index_copy_multi_slim_jit( + dst, src, feat, di, si, num, num_blocks=blocks, num_threads=threads, unroll=unroll + ) + + return CopyVariant(f"slim {blocks}x{threads}u{unroll}", launch) + + +class CopyRig: + """Registered host banks, GPU slot caches and the fused-copy descriptor tensors.""" + + def __init__(self, feats: list[int], n_src: int, n_slots: int, device: torch.device, seed: int): + self.feats = feats + self.expert_bytes = sum(feats) + self.n_src = n_src + self.n_slots = n_slots + self.device = device + gen = torch.Generator().manual_seed(seed) + self.host = [] + for feat in feats: + bank = HostBank((n_src, feat), torch.uint8) + bank.tensor.copy_(torch.randint(0, 256, (n_src, feat), dtype=torch.uint8, generator=gen)) + bank.pin() + self.host.append(bank) + self.slots = [torch.empty((n_slots, feat), dtype=torch.uint8, device=device) for feat in feats] + self.dst_ptrs = torch.tensor([c.data_ptr() for c in self.slots], dtype=torch.int64, device=device) + self.src_ptrs = torch.tensor([device_ptr(b.tensor) for b in self.host], dtype=torch.int64, device=device) + self.feat_bytes = torch.tensor(feats, dtype=torch.int64, device=device) + self.spacer_buf = torch.zeros(1, dtype=torch.int32, device=device) + self.gen = gen + self._cursor = 0 + self._perm = torch.randperm(n_src, generator=gen) + + def plans(self, n_launch: int, experts: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """[n_launch, LMAX] dst/src plans: source rows walk a permutation of the host rows + (no reuse until all are used), destination slots are distinct within a plan.""" + src = torch.zeros((n_launch, LMAX), dtype=torch.int32) + dst = torch.zeros((n_launch, LMAX), dtype=torch.int32) + for r in range(n_launch): + for j in range(experts): + if self._cursor == self.n_src: + self._perm = torch.randperm(self.n_src, generator=self.gen) + self._cursor = 0 + src[r, j] = self._perm[self._cursor] + self._cursor += 1 + dst[r, :experts] = torch.randperm(self.n_slots, generator=self.gen)[:experts] + num = torch.full((1,), experts, dtype=torch.int64, device=self.device) + return dst.to(self.device), src.to(self.device), num + + def copy_loop(self, variant: CopyVariant, n_launch: int, experts: int, cold: bool) -> Callable[[], None]: + dst, src, num = self.plans(n_launch, experts) + + def run(): + for r in range(n_launch): + variant.launch(self.dst_ptrs, self.src_ptrs, self.feat_bytes, dst[r], src[r], num) + if cold: + self.spacer_buf.add_(1) + + return run + + def spacer_us(self, stream: torch.cuda.Stream, trials: int) -> float: + n = 200 + + def run(): + for _ in range(n): + self.spacer_buf.add_(1) + + g = capture(run, stream) + return statistics.median(time_alone(g, n, stream) for _ in range(trials)) + + +@dataclass +class ComputeOp: + name: str + make: Callable[[int], Callable[[], None]] # n_launch -> loop body + single: Callable[[], None] + + +def _weight_copies(nbytes: int) -> int: + return max(2, math.ceil(2 * L2_BYTES / nbytes)) + + +def gemv_bf16(h: int, n: int, device) -> ComputeOp: + x = torch.randn(1, h, dtype=torch.bfloat16, device=device) + ws = [torch.randn(n, h, dtype=torch.bfloat16, device=device) for _ in range(_weight_copies(n * h * 2))] + + def make(n_launch): + def run(): + for r in range(n_launch): + F.linear(x, ws[r % len(ws)]) + + return run + + return ComputeOp(f"gemv_bf16 1x{h}x{n}", make, lambda: F.linear(x, ws[0])) + + +def gemv_fp8(h: int, n: int, device) -> ComputeOp: + fp8 = torch.float8_e4m3fn + x = torch.randn(1, h, device=device).to(fp8) + ws = [torch.randn(n, h, device=device).to(fp8) for _ in range(_weight_copies(n * h))] + one = torch.ones((), dtype=torch.float32, device=device) + + def mm(w): + torch._scaled_mm(x, w.t(), scale_a=one, scale_b=one, out_dtype=torch.bfloat16) + + def make(n_launch): + def run(): + for r in range(n_launch): + mm(ws[r % len(ws)]) + + return run + + return ComputeOp(f"gemv_fp8 1x{h}x{n}", make, lambda: mm(ws[0])) + + +def matmul_bf16(m: int, h: int, n: int, device) -> ComputeOp: + a = torch.randn(m, h, dtype=torch.bfloat16, device=device) + ws = [torch.randn(n, h, dtype=torch.bfloat16, device=device) for _ in range(_weight_copies(n * h * 2))] + + def make(n_launch): + def run(): + for r in range(n_launch): + F.linear(a, ws[r % len(ws)]) + + return run + + return ComputeOp(f"mm_bf16 {m}x{h}x{n}", make, lambda: F.linear(a, ws[0])) + + +def nvfp4_routed(h: int, i: int, top_k: int, n_slots: int, device) -> ComputeOp: + """The decode routed-expert GEMV pair (gate_up -> act -> down) over its own slot banks.""" + from freetoken.moe.fused_nvfp4 import fused_experts_decode_nvfp4_marlin + + fp8 = torch.float8_e4m3fn + gu = torch.randint(0, 256, (n_slots, 2 * i, h // 2), dtype=torch.uint8, device=device) + gus = torch.full((n_slots, 2 * i, h // 16), 1.0, device=device).to(fp8) + gug = torch.full((n_slots, 2 * i), 1.0, dtype=torch.float16, device=device) + dn = torch.randint(0, 256, (n_slots, h, i // 2), dtype=torch.uint8, device=device) + dns = torch.full((n_slots, h, i // 16), 1.0, device=device).to(fp8) + dng = torch.full((n_slots, h), 1.0, dtype=torch.float16, device=device) + x = torch.randn(1, h, dtype=torch.bfloat16, device=device) * 0.01 + w = torch.full((1, top_k), 1.0 / top_k, dtype=torch.float32, device=device) + groups = max(1, n_slots // top_k) + ids = [ + (torch.arange(top_k, dtype=torch.int32, device=device) + g * top_k).view(1, top_k) % n_slots + for g in range(groups) + ] + + def call(t): + fused_experts_decode_nvfp4_marlin(x, gu, gus, gug, dn, dns, dng, w, t, "silu", False) + + def make(n_launch): + def run(): + for r in range(n_launch): + call(ids[r % groups]) + + return run + + return ComputeOp(f"nvfp4_routed k={top_k}", make, lambda: call(ids[0])) + + +def capture(fn: Callable[[], None], stream: torch.cuda.Stream) -> torch.cuda.CUDAGraph: + with torch.cuda.stream(stream): + fn() # warm-up (JIT / Triton compile, cuBLAS heuristics) outside capture + stream.synchronize() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g, stream=stream): + fn() + torch.cuda.synchronize() + g.keepalive = fn # the closure owns the plan tensors the graph reads + return g + + +def _event() -> torch.cuda.Event: + return torch.cuda.Event(enable_timing=True) + + +def time_alone(g: torch.cuda.CUDAGraph, n: int, stream: torch.cuda.Stream) -> float: + """us per launch of a graph holding ``n`` launches, replayed alone.""" + start, end = _event(), _event() + with torch.cuda.stream(stream): + torch.cuda._sleep(SLEEP_CYCLES) + start.record(stream) + g.replay() + end.record(stream) + end.synchronize() + return start.elapsed_time(end) * 1e3 / n + + +def time_corun( + g_fg: torch.cuda.CUDAGraph, n_fg: int, s_fg: torch.cuda.Stream, + g_bg: torch.cuda.CUDAGraph, s_bg: torch.cuda.Stream, +) -> tuple[float, bool]: + """us per launch of the foreground graph while the background graph runs on another + stream, and whether the background outlasted the foreground (a valid co-run).""" + ev0, bg_end, fg_start, fg_end = _event(), _event(), _event(), _event() + s0 = torch.cuda.current_stream() + torch.cuda._sleep(SLEEP_CYCLES) + ev0.record(s0) + s_bg.wait_event(ev0) + s_fg.wait_event(ev0) + with torch.cuda.stream(s_bg): + g_bg.replay() + bg_end.record(s_bg) + with torch.cuda.stream(s_fg): + torch.cuda._sleep(SLEEP_CYCLES // 20) # let the background reach steady state + fg_start.record(s_fg) + g_fg.replay() + fg_end.record(s_fg) + torch.cuda.synchronize() + covered = ev0.elapsed_time(bg_end) >= ev0.elapsed_time(fg_end) + return fg_start.elapsed_time(fg_end) * 1e3 / n_fg, covered + + +def kernel_info(fn: Callable[[], None]) -> list[dict]: + """(name, grid, block, regs/thread, smem) of every kernel one call launches.""" + from torch.profiler import ProfilerActivity, profile + + fn() + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CUDA]) as prof: + fn() + torch.cuda.synchronize() + with tempfile.TemporaryDirectory() as d: + path = os.path.join(d, "trace.json") + prof.export_chrome_trace(path) + with open(path) as f: + events = json.load(f)["traceEvents"] + out = [] + for e in events: + if e.get("cat") != "kernel": + continue + a = e.get("args", {}) + out.append({ + "name": e["name"][:80], + "grid": a.get("grid"), + "block": a.get("block"), + "regs": a.get("registers per thread"), + "smem": a.get("shared memory"), + }) + return out + + +def _fmt_info(info: list[dict]) -> str: + return "; ".join( + f"{k['name'][:48]} grid={k['grid']} block={k['block']} regs={k['regs']} smem={k['smem']}" for k in info + ) + + +def _stats(xs: list[float]) -> tuple[float, float]: + return min(xs), statistics.median(xs) + + +def parse_pairs(text: str, n: int) -> list[tuple[int, ...]]: + out = [] + for item in text.split(","): + item = item.strip() + if not item: + continue + parts = tuple(int(p) for p in item.lower().replace("u", "x").split("x")) + if len(parts) != n: + raise argparse.ArgumentTypeError(f"{item!r}: expected {n} 'x'-separated ints") + out.append(parts) + return out + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--gpu", type=single_gpu_arg, default=None, + help="GPU UUID or nvidia-smi index (default: the first visible GPU)") + p.add_argument("--hidden", type=int, default=2560) + p.add_argument("--inter", type=int, default=640) + p.add_argument("--src-experts", type=int, default=256, help="host rows per bank (pinned)") + p.add_argument("--slots", type=int, default=64, help="GPU slot-cache rows per bank") + p.add_argument("--multi", default="1024x8,1024x4,512x8,512x4,256x16,256x8,128x32", + help="fast_index_copy_multi configs, THREADSxBLOCKS_PER_BANK") + p.add_argument("--slim", default="8x256x4,16x256x2,12x256x4,16x256x4,8x256x2,4x256x4", + help="fast_index_copy_multi_slim configs, BLOCKSxTHREADSxUNROLL " + "(skipped if the kernel is absent)") + p.add_argument("--experts", default="1,2,3,4,5,6,7,8,9,10", help="plan sizes for the sweep") + p.add_argument("--corun-experts", default="4", help="plan sizes for the co-run copy loop") + p.add_argument("--link", choices=["hot", "cold", "both"], default="both", help="sweep link mode(s)") + p.add_argument("--corun-link", choices=["hot", "cold"], default="cold", help="co-run copy loop link mode") + p.add_argument("--trials", type=int, default=5) + p.add_argument("--copy-ms", type=float, default=6.0, help="target length of one timed copy loop") + p.add_argument("--compute-ms", type=float, default=3.0, help="target length of one timed compute loop") + p.add_argument("--cover", type=float, default=3.0, + help="background loop length as a multiple of the foreground's alone time") + p.add_argument("--skip-sweep", action="store_true") + p.add_argument("--skip-corun", action="store_true") + p.add_argument("--ops", default="all", + help="comma list of co-run ops by prefix: gemv_bf16,gemv_fp8,mm_bf16,nvfp4_routed") + p.add_argument("--json", default=None, help="write all results to this JSON path") + p.add_argument("--seed", type=int, default=0) + return p.parse_args() + + +def build_variants(args) -> list[CopyVariant]: + variants = [multi_variant(t, b) for t, b in parse_pairs(args.multi, 2)] + if hasattr(fic, "fast_index_copy_multi_slim_jit"): + variants += [slim_variant(g, t, u) for g, t, u in parse_pairs(args.slim, 3)] + else: + print("fast_index_copy_multi_slim_jit not available: slim configs skipped", flush=True) + return variants + + +def run_sweep(rig: CopyRig, variants, experts_list, args, stream, spacer_us, results) -> None: + modes = ["hot", "cold"] if args.link == "both" else [args.link] + for mode in modes: + cold = mode == "cold" + sub = spacer_us if cold else 0.0 + print(f"\n== sweep ({mode} link{f', spacer {sub:.2f} us subtracted' if cold else ''}): copy alone, " + f"expert={rig.expert_bytes} B, trials={args.trials}; cells are GB/s at min-time / " + f"median-time (us/copy median)", flush=True) + print("variant".ljust(18) + "".join(f"{e:>20d}" for e in experts_list), flush=True) + for v in variants: + row = v.name.ljust(18) + for e in experts_list: + n_launch = max(8, math.ceil(args.copy_ms * 1e3 / (e * 55.0))) + g = capture(rig.copy_loop(v, n_launch, e, cold), stream) + ts = [time_alone(g, n_launch, stream) - sub for _ in range(args.trials)] + t_min, t_med = _stats(ts) + nbytes = e * rig.expert_bytes + results["sweep"].append({ + "link": mode, "variant": v.name, "experts": e, "us_min": t_min, "us_median": t_med, + "gbps_best": nbytes / (t_min * 1e3), "gbps_median": nbytes / (t_med * 1e3), "us_all": ts, + }) + row += f"{nbytes / (t_min * 1e3):>7.1f}/{nbytes / (t_med * 1e3):5.1f} ({t_med:5.0f})" + del g + print(row, flush=True) + + +def run_corun(rig: CopyRig, variants, ops, corun_experts, args, s_fg, s_bg, spacer_us, results) -> None: + cold = args.corun_link == "cold" + sub = spacer_us if cold else 0.0 + print(f"\n== co-run ({args.corun_link} link copy loop): compute on one stream, copy loop on another; " + f"trials={args.trials} (interleaved alone/co-run); sigma = t_corun / t_alone - 1", flush=True) + for op in ops: + n_fg_probe = 20 + g_probe = capture(op.make(n_fg_probe), s_fg) + t_probe = statistics.median(time_alone(g_probe, n_fg_probe, s_fg) for _ in range(3)) + del g_probe + n_op = max(20, math.ceil(args.compute_ms * 1e3 / t_probe)) + g_op = capture(op.make(n_op), s_fg) + print(f"\n-- {op.name}: ~{t_probe:.1f} us/launch alone, loop of {n_op}", flush=True) + results["kernels"][op.name] = kernel_info(op.single) + print(f" kernels: {_fmt_info(results['kernels'][op.name])}", flush=True) + print(" " + "copy variant".ljust(18) + "E".rjust(3) + + " op_alone_us(min/med) op_corun_us(min/med) sigma(min/med)" + + " copy_GBps alone(med) -> under op(med) valid (! = alone median >10% over the probe:" + + " other GPU load, trust the min column)", flush=True) + for v in variants: + for e in corun_experts: + copy_launch_us = e * 55.0 + n_copy_long = max(4, math.ceil(args.cover * n_op * t_probe / copy_launch_us)) + g_copy_long = capture(rig.copy_loop(v, n_copy_long, e, cold), s_bg) + n_copy = max(8, math.ceil(args.copy_ms * 1e3 / copy_launch_us)) + g_copy = capture(rig.copy_loop(v, n_copy, e, cold), s_bg) + n_op_long = max(n_op, math.ceil(args.cover * n_copy * copy_launch_us / t_probe)) + g_op_long = capture(op.make(n_op_long), s_fg) if n_op_long > n_op else g_op + op_alone, op_co, cp_alone, cp_co = [], [], [], [] + valid_op = valid_cp = 0 + for _ in range(args.trials): + op_alone.append(time_alone(g_op, n_op, s_fg)) + t, ok = time_corun(g_op, n_op, s_fg, g_copy_long, s_bg) + op_co.append(t) + valid_op += ok + cp_alone.append(time_alone(g_copy, n_copy, s_bg) - sub) + t, ok = time_corun(g_copy, n_copy, s_bg, g_op_long, s_fg) + cp_co.append(t - sub) + valid_cp += ok + a_min, a_med = _stats(op_alone) + c_min, c_med = _stats(op_co) + sig_min, sig_med = c_min / a_min - 1, c_med / a_med - 1 + nbytes = e * rig.expert_bytes + gb_alone = nbytes / (statistics.median(cp_alone) * 1e3) + gb_co = nbytes / (statistics.median(cp_co) * 1e3) + suspect = a_med > 1.1 * t_probe + results["corun"].append({ + "link": args.corun_link, "op": op.name, "variant": v.name, "experts": e, "suspect": suspect, + "op_alone_us": op_alone, "op_corun_us": op_co, + "sigma_min": sig_min, "sigma_median": sig_med, + "copy_alone_us": cp_alone, "copy_corun_us": cp_co, + "copy_gbps_alone_median": gb_alone, "copy_gbps_corun_median": gb_co, + "valid_op_trials": valid_op, "valid_copy_trials": valid_cp, + }) + print(f" {v.name.ljust(18)}{e:>3d} {a_min:8.2f}/{a_med:8.2f} {c_min:8.2f}/{c_med:8.2f}" + f" {sig_min:+6.3f}/{sig_med:+6.3f} {gb_alone:6.1f} -> {gb_co:6.1f}" + f" {valid_op}/{valid_cp}{' !' if suspect else ''}", flush=True) + del g_copy_long, g_copy, g_op_long + del g_op + + +def main() -> None: + args = parse_args() + assert torch.cuda.is_available(), "CUDA is required" + try: + assign_gpu(args.gpu) + device = bind_assigned_gpu() + except (ValueError, RuntimeError) as e: + raise SystemExit(f"error: {e}") from e + + feats = nvfp4_bank_bytes(args.hidden, args.inter) + assert sum(feats) == _BANK_BYTES_PER_EXPERT["nvfp4"](args.hidden, args.inter) + if (args.hidden, args.inter) == (2560, 640): + assert sum(feats) == 2_772_480 + print("gpu", torch.cuda.get_device_name(device), "| banks", feats, "| expert", sum(feats), "B", flush=True) + rig = CopyRig(feats, args.src_experts, args.slots, device, args.seed) + variants = build_variants(args) + results = {"feats": feats, "sweep": [], "corun": [], "kernels": {}} + + s_fg = torch.cuda.Stream(device=device) + s_bg = torch.cuda.Stream(device=device) + spacer_us = rig.spacer_us(s_bg, args.trials) + results["spacer_us"] = spacer_us + print(f"spacer kernel alone: {spacer_us:.2f} us", flush=True) + for v in variants: + dst, src, num = rig.plans(1, 4) + info = kernel_info(lambda: v.launch(rig.dst_ptrs, rig.src_ptrs, rig.feat_bytes, dst[0], src[0], num)) + results["kernels"][v.name] = info + print(f"{v.name.ljust(18)} {_fmt_info(info)}", flush=True) + + if not args.skip_sweep: + experts_list = [int(x) for x in args.experts.split(",") if x] + run_sweep(rig, variants, experts_list, args, s_bg, spacer_us, results) + + if not args.skip_corun: + h, i = args.hidden, args.inter + all_ops = [ + gemv_bf16(h, 5120, device), gemv_bf16(h, 12288, device), + gemv_fp8(h, 5120, device), gemv_fp8(h, 12288, device), + matmul_bf16(256, h, h, device), + nvfp4_routed(h, i, 10, args.slots, device), + ] + wanted = [w.strip() for w in args.ops.split(",") if w.strip()] + ops = all_ops if wanted == ["all"] else [o for o in all_ops if any(o.name.startswith(w) for w in wanted)] + corun_experts = [int(x) for x in args.corun_experts.split(",") if x] + run_corun(rig, variants, ops, corun_experts, args, s_fg, s_bg, spacer_us, results) + + if args.json: + with open(args.json, "w") as f: + json.dump(results, f, indent=1) + print(f"\nwrote {args.json}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_moe_copy_overlap.py b/benchmarks/bench_moe_copy_overlap.py new file mode 100644 index 0000000000..535f7b1bd2 --- /dev/null +++ b/benchmarks/bench_moe_copy_overlap.py @@ -0,0 +1,233 @@ +"""Decode MoE copy/compute overlap bench (FREETOKEN_MOE_COPY_OVERLAP). + +Qwen4ExpMoE layers at the Qwen3.8-Flash-Next expert geometry (hidden 2560, moe and shared +intermediate 640, top-10, NVFP4 triton banks: 2,772,480 B per expert) over one offload cache, +with fewer experts per layer than the model so host and GPU memory stay small. Each arm +(overlap off / on) captures one decode step through every layer as a CUDA graph, then: + 1. bitwise check of the two arms over the same fresh inputs, + 2. replay time per step with CUDA events, arms interleaved ABAB, fresh inputs every round, + 3. a torch.profiler (CUPTI) trace of graph replays: per layer, the copy kernel's start slip + after the ensure, the shared expert's span (side stream) and whether it ends inside the copy, the join + gap, and the copy's duration against the off arm. +No attention or hyper-connection work runs, so the percentages overstate the model's. + +Run: CUDA_VISIBLE_DEVICES=0 PYTHONPATH=python python benchmarks/bench_moe_copy_overlap.py +""" + +from __future__ import annotations + +import argparse +import statistics +from types import SimpleNamespace + +import torch + +H, I, K = 2560, 640, 10 + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--layers", type=int, default=4) + parser.add_argument("--experts", type=int, default=64, help="experts per layer") + parser.add_argument("--slots", type=int, default=160, help="cache slots (160 of 4 x 64 gives ~4.4 misses/layer)") + parser.add_argument("--bs", type=int, default=1) + parser.add_argument("--steps", type=int, default=40, help="decode steps per timing round") + parser.add_argument("--rounds", type=int, default=6) + return parser.parse_args() + + +def nvfp4_quant(): + from freetoken.layers.quantization import QuantBackend, QuantConfig, set_quant_backend + + set_quant_backend(QuantBackend.parse("moe.nvfp4=triton")) + return QuantConfig.from_hf( + {"quantization_config": {"quant_method": "modelopt", "quant_algo": "NVFP4", "ignore": ["lm_head"]}} + ) + + +def build_arm(args, overlap: bool, banks) -> SimpleNamespace: + import freetoken.core as core + from freetoken.core import Context, set_global_ctx + from freetoken.layers.moe import OffloadMoELayer + from freetoken.models.qwen4_exp.moe import Qwen4ExpMoE + from freetoken.moe.offload_cache import OffloadMoeCache + from freetoken.utils.torch_utils import torch_dtype + + E, L = args.experts, args.layers + config = SimpleNamespace( + hidden_size=H, num_experts=E, num_experts_per_tok=K, moe_intermediate_size=I, + shared_expert_intermediate_size=I, norm_topk_prob=True, quant=None, + moe_strategy="offload", decode_target="gpu", + ) + device = torch.device("cuda") + with torch.device(device), torch_dtype(torch.bfloat16): + moes = [Qwen4ExpMoE(config, l, prefix=f"model.layers.{l}.mlp") for l in range(L)] + gen = torch.Generator(device=device).manual_seed(11) + for moe in moes: + for t in moe.state_dict().values(): + if t.is_floating_point(): + t.normal_(0.0, 0.02, generator=gen) + quant = nvfp4_quant() + experts = [ + OffloadMoELayer(l, E, K, H, I, quant_config=quant, prefix=f"model.layers.{l}.mlp.experts") + for l in range(L) + ] + cache = OffloadMoeCache( + num_layers=L, num_experts=E, cache_size=args.slots, device=device, + quant_format=banks.quant_format, layout=banks.layout, + max_slots=experts[0].quant_method.slot_limit(), decode_copy_overlap=overlap, + ) + cache.set_bank_sources(banks.sources) + cache.set_alphas(banks.gate_up_alpha, banks.down_alpha) + cache.collect_stats = True + for moe, ex in zip(moes, experts): + ex.offload_cache = cache + moe.experts = ex + cache.reset() + core._GLOBAL_CTX = None + ctx = Context(page_size=1) + set_global_ctx(ctx) + ctx._batch = SimpleNamespace(is_prefill=False) + + static = torch.randn(args.bs, H, device=device, dtype=torch.bfloat16) + out = torch.empty(L, args.bs, H, device=device, dtype=torch.bfloat16) + + def step(): + for l, moe in enumerate(moes): + out[l].copy_(moe.forward(static.clone())) + + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + step() # the eager warm-up the graph runner also does + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + step() + cache.reset() + cache.reset_stats() + # the graph holds raw pointers only: the modules must outlive it or replays read freed weights + return SimpleNamespace(graph=graph, static=static, out=out, cache=cache, moes=moes) + + +def fresh_inputs(args, n: int, seed: int) -> list[torch.Tensor]: + gen = torch.Generator(device="cuda").manual_seed(seed) + return [torch.randn(args.bs, H, device="cuda", dtype=torch.bfloat16, generator=gen) for _ in range(n)] + + +def replay(arm, inputs, keep: bool = False) -> list[torch.Tensor]: + outs = [] + for x in inputs: + arm.static.copy_(x) + arm.graph.replay() + if keep: + outs.append(arm.out.clone()) + return outs + + +def time_arm(arm, inputs) -> float: + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + replay(arm, inputs) + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / len(inputs) + + +def trace(arm, inputs) -> list[tuple[float, float, str]]: + from torch.profiler import ProfilerActivity, profile + + with profile(activities=[ProfilerActivity.CUDA]) as prof: + replay(arm, inputs) + torch.cuda.synchronize() + return sorted( + (e.time_range.start, e.time_range.end, e.name) for e in prof.events() if e.device_type.name == "CUDA" + ) + + +def layer_calls(ev, overlap: bool) -> list[dict]: + """Per MoE layer call, located by kernel name (graph replay spreads branches over streams).""" + rows = [] + for i, (_, ens_end, name) in enumerate(ev): + if "_lru_ensure_kernel" not in name: + continue + copy = next((k for k in ev[i + 1:] if "fast_index_copy" in k[2]), None) + gemm = next((k for k in ev[i + 1:] if "_decode_nvfp4" in k[2]), None) + if copy is None or gemm is None: + continue + # the shared expert chain starts at the gate kernel before the ensure; with the overlap it runs on + # the side stream beside the routed path, without it serially ahead of that path's topk + j = max(j for j in range(i) if "_gate_sigmoid_kernel" in ev[j][2]) + routed_path = ("router", "topk", "ensure", "fast_index_copy") + chain = [k for k in (ev[j:] if overlap else ev[j:i]) if k[1] <= gemm[0] + and not any(m in k[2] for m in routed_path)] + rows.append({ + "copy_us": copy[1] - copy[0], + "slip_us": copy[0] - ens_end, + "join_gap_us": gemm[0] - copy[1], + "shared_span_us": max(k[1] for k in chain) - chain[0][0], + "shared_ends_in_copy": max(k[1] for k in chain) <= copy[1], + }) + return rows + + +def pct(xs, q: float) -> float: + xs = sorted(xs) + return xs[min(len(xs) - 1, int(q * len(xs)))] + + +def main() -> None: + args = parse_args() + from freetoken.distributed import set_tp_info + from freetoken.layers.moe import OffloadMoELayer + from freetoken.moe.expert_banks import build_expert_banks + + set_tp_info(rank=0, size=1) + torch.manual_seed(0) + method = OffloadMoELayer(0, args.experts, K, H, I, quant_config=nvfp4_quant(), prefix="model.layers.0.mlp.experts").quant_method + banks = build_expert_banks(method, args.layers, None, device=torch.device("cuda"), dummy=True) + expert_bytes = sum(layers[0][0].numel() * layers[0].element_size() for layers in banks.sources.values()) + print(f"expert {expert_bytes} B; {args.layers} layers x {args.experts} experts, {args.slots} slots, bs {args.bs}") + arms = {False: build_arm(args, False, banks), True: build_arm(args, True, banks)} + + seed = 100 + inputs = fresh_inputs(args, args.steps, seed) + outs = {k: replay(arm, inputs, keep=True) for k, arm in arms.items()} + same = all(torch.equal(a, b) for a, b in zip(outs[True], outs[False])) + active, missed, calls = arms[False].cache.lru_stats.sum(0).tolist() + print(f"bitwise identical on/off over {args.steps} steps: {same}; misses per layer call {missed / calls:.2f}") + + times = {False: [], True: []} + for r in range(args.rounds): + seed += 1 + inputs = fresh_inputs(args, args.steps, seed) + for k in (False, True) if r % 2 == 0 else (True, False): + times[k].append(time_arm(arms[k], inputs)) + for k in (False, True): + t = times[k] + print(f"overlap {'on ' if k else 'off'}: {statistics.median(t):.4f} ms/step median (min {min(t):.4f}, max {max(t):.4f}, n={len(t)})") + per_layer = [1000 * (a - b) / args.layers for a, b in zip(times[False], times[True])] + print(f"saving per layer, paired rounds (us): {[round(x, 1) for x in per_layer]}, median {statistics.median(per_layer):.1f}") + + inputs = fresh_inputs(args, 12, seed + 1) + rows = {k: layer_calls(trace(arms[k], inputs), k) for k in (False, True)} + for k in (False, True): + r = [x for x in rows[k] if x["copy_us"] > 5.0] # layer calls with at least one miss + if not r: + continue + slips = [x["slip_us"] for x in r] + print( + f"overlap {'on ' if k else 'off'}: {len(r)} layer calls with misses; copy {statistics.median(x['copy_us'] for x in r):.1f} us; " + f"copy start - ensure end median {statistics.median(slips):.2f} us (p90 {pct(slips, 0.9):.2f}, <=5 us in {100 * sum(s <= 5 for s in slips) / len(slips):.0f}%); " + f"GEMM start - copy end median {statistics.median(x['join_gap_us'] for x in r):.2f} us; " + f"shared expert span median {statistics.median(x['shared_span_us'] for x in r):.1f} us" + + (f", ends inside the copy in {100 * sum(x['shared_ends_in_copy'] for x in r) / len(r):.0f}%" if k else "") + ) + pairs = [(a["copy_us"], b["copy_us"]) for a, b in zip(rows[False], rows[True]) if a["copy_us"] > 5.0] + if pairs: + ratio = [b / a for a, b in pairs] + print(f"copy duration on/off: median {statistics.median(ratio):.3f} (p10 {pct(ratio, 0.1):.3f}, p90 {pct(ratio, 0.9):.3f})") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_scored_ensure.py b/benchmarks/bench_scored_ensure.py new file mode 100644 index 0000000000..5aa852716e --- /dev/null +++ b/benchmarks/bench_scored_ensure.py @@ -0,0 +1,158 @@ +"""Decode ensure cost per call: flashlib lru_ensure vs the scored --moe-cache-policy kernels. + +One ensure call per MoE layer per decode step, on a full cache (S slots over L x E experts). +Each timed call routes K ids of which exactly m miss: the hits are hot, recently used experts +and the victims are cold filler, under every policy. A CUDA graph of 2 decode steps (2 * L +calls; 1 step when that many fresh ids would not fit) is replayed from a restored snapshot, +so the number is the graph-node cost the decode graph pays. The miss count is verified +through the kernels' stats before timing. + +Run: CUDA_VISIBLE_DEVICES=0 PYTHONPATH=python python benchmarks/bench_scored_ensure.py +""" + +from __future__ import annotations + +import argparse +import statistics + +import numpy as np +import torch +from flashlib.kernels.slot_cache import N_STATS, Stat, lru_ensure + +from freetoken.moe import scored_ensure as se + +L, E, TOP_K = 48, 512, 10 +VARIANTS = ("flashlib-lru", "lru", "kd", "kdfb", "rule") + + +def build(variant: str, slots: int, k: int, m: int, dev: torch.device): + """Full-cache state plus 1-2 steps of L queries of k ids, the first k - m hot, the rest fresh misses.""" + policy = 0 if variant.endswith("lru") else se.POLICY_IDS[variant] + hot = k - m + calls = 2 * L if hot * L + 2 * L * m <= slots else L + step0, tok0 = 1_000_000, 1000 + slot_of_id = np.full(L * E, -1, np.int32) + id_of_slot = np.full(slots, -1, np.int32) + usage = np.zeros(slots, np.int64) + last_tok = np.full(L * E, se.LAST_TOK_NEVER, np.int64) + lc = np.full(L * E, se.LC_NEVER, np.int32) + ct = np.zeros(L * E, np.int64) + ids = [layer * E + e for layer in range(L) for e in range(hot)] + assert len(ids) + calls * m <= slots, "fresh ids would evict hot ones" + filler = (layer * E + e for e in range(k + m, E) for layer in range(L)) + ids += [next(filler) for _ in range(slots - len(ids))] + for s, key in enumerate(ids): + id_of_slot[s], slot_of_id[key] = key, s + is_hot = key % E < hot + usage[s] = step0 - 1 if is_hot else 1 + s + last_tok[key] = tok0 if is_hot else tok0 - 500 - s % 7 + if is_hot: + lc[key], ct[key] = 5 << se.Q, tok0 * L + key // E + queries = np.empty((calls, k), np.int32) + for i in range(calls): + visit = i // L + queries[i] = np.r_[np.arange(hot), hot + visit * m + np.arange(m)] + rng = np.random.default_rng(0) + logits = rng.normal(size=(calls, k // TOP_K, E)).astype(np.float32) + rows = queries.reshape(calls, k // TOP_K, -1).astype(np.int64) + np.put_along_axis(logits, rows, np.take_along_axis(logits, rows, -1) + 3.0, -1) + held = id_of_slot.astype(np.int64) + t = lambda a, dtype=None: torch.from_numpy(a).to(dev, dtype) # noqa: E731 + state = dict( + owner=t(id_of_slot.copy()), slot_last=t(last_tok[held]), slot_lc=t(lc[held]), slot_ct=t(ct[held]), + slot_of_id=t(slot_of_id), id_of_slot=t(id_of_slot), usage=t(usage), + step=torch.tensor(step0, dtype=torch.int64, device=dev), + tok=torch.tensor(tok0, dtype=torch.int64, device=dev), + last_tok=t(last_tok), lc=t(lc), ct=t(ct), queries=t(queries), + src=torch.zeros(slots, dtype=torch.int32, device=dev), + dst=torch.zeros(slots, dtype=torch.int32, device=dev), + num=torch.zeros((), dtype=torch.int64, device=dev), + stats=torch.zeros(L, N_STATS, dtype=torch.int64, device=dev), + ) + return policy, state, t(logits, torch.bfloat16), se.softplus2_table(dev) + + +def step_fn(variant, policy, st, logits, g_table, stats: bool): + def run(): + for i in range(st["queries"].shape[0]): + layer = i % L + q = st["queries"][i] + s = st["stats"][layer] if stats else None + if variant == "flashlib-lru": + lru_ensure(q, st["slot_of_id"], st["id_of_slot"], st["usage"], st["step"], q, + st["src"], st["dst"], st["num"], stats=s, id_base=layer * E) + continue + se.scored_ensure( + q, st["slot_of_id"], st["id_of_slot"], st["usage"], st["step"], q, + st["src"], st["dst"], st["num"], policy=policy, num_layers=L, num_experts=E, + tok=st["tok"], last_tok=st["last_tok"], lc=st["lc"], ct=st["ct"], g_table=g_table, + slot_owner=st["owner"], slot_last_tok=st["slot_last"], slot_lc=st["slot_lc"], slot_ct=st["slot_ct"], + router_logits=logits[i] if policy == 3 else None, near_miss_thr=0.25, + stats=s, id_base=layer * E, bump_tok=layer == 0, + ) + return run + + +def bench(variant, slots, k, m, dev, reps, stats): + policy, st, logits, g_table = build(variant, slots, k, m, dev) + snap = {name: v.clone() for name, v in st.items()} + + def restore(): + for name, v in snap.items(): + st[name].copy_(v) + + # verify the miss count (a stats build of the same calls), then time the requested build + restore() + step_fn(variant, policy, st, logits, g_table, True)() + got = int(st["stats"][:, Stat.MISS].sum()) + calls = st["queries"].shape[0] + assert got == calls * m, f"{variant} K={k} m={m}: {got} misses, expected {calls * m}" + restore() + run = step_fn(variant, policy, st, logits, g_table, stats) + run() # compile outside capture + restore() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + times = [] + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + for rep in range(reps + 3): + restore() + torch.cuda.synchronize() + start.record() + graph.replay() + end.record() + end.synchronize() + if rep >= 3: + times.append(start.elapsed_time(end) * 1000.0 / calls) + return statistics.median(times) + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--slots", type=int, default=1650) + ap.add_argument("--k", default="10,20", help="routed ids per call (10 = bs 1, 20 = bs 2)") + ap.add_argument("--misses", default="0-10", help="list or range; capped at K per call") + ap.add_argument("--variants", default=",".join(VARIANTS)) + ap.add_argument("--reps", type=int, default=30) + ap.add_argument("--stats", action="store_true", help="time with --moe-collect-stats accumulation on") + args = ap.parse_args() + dev = torch.device("cuda") + variants = args.variants.split(",") + misses = [] + for part in args.misses.split(","): + lo, _, hi = part.partition("-") + misses += range(int(lo), int(hi or lo) + 1) + print(f"# {torch.cuda.get_device_name(dev)}, S={args.slots}, us per ensure call " + f"(median of {args.reps} graph replays), stats={'on' if args.stats else 'off'}") + print(f"{'K':>3} {'m':>3} " + " ".join(f"{v:>13}" for v in variants) + " max delta vs flashlib-lru") + for k in (int(x) for x in args.k.split(",")): + for m in (x for x in misses if x <= k): + us = {v: bench(v, args.slots, k, m, dev, args.reps, args.stats) for v in variants} + base = us.get("flashlib-lru") + delta = "" if base is None else f" {max(us[v] - base for v in variants):+.2f}" + print(f"{k:>3} {m:>3} " + " ".join(f"{us[v]:>13.2f}" for v in variants) + delta, flush=True) + + +if __name__ == "__main__": + main() diff --git a/python/freetoken/env.py b/python/freetoken/env.py index c8eb1e7513..0e56633f48 100644 --- a/python/freetoken/env.py +++ b/python/freetoken/env.py @@ -77,6 +77,8 @@ class EnvClassSingleton: MAMBA_SSM_DTYPE = EnvStr("float32") # cap the longest prefill warmup forward; 0 = the engine's own chunk cap WARMUP_MAX_LEN = EnvInt(0) + # GPU-decode MoE: run a layer's expert miss copy on a side stream under the shared expert + MOE_COPY_OVERLAP = EnvBool(False) def __new__(cls): # single instance diff --git a/python/freetoken/kernel/aot.py b/python/freetoken/kernel/aot.py index bbc2d2c669..54539d952c 100644 --- a/python/freetoken/kernel/aot.py +++ b/python/freetoken/kernel/aot.py @@ -124,6 +124,23 @@ def build(build_directory: pathlib.Path) -> object: return KernelSpec(name=_make_name("fast_index_copy_multi", *args), build=build) +def _fast_index_copy_multi_slim_spec() -> KernelSpec: + from .fast_index_copy import SLIM_COPY_BLOCKS, SLIM_COPY_THREADS, SLIM_COPY_UNROLL + + args = make_cpp_args(SLIM_COPY_BLOCKS, SLIM_COPY_THREADS, SLIM_COPY_UNROLL) + + def build(build_directory: pathlib.Path) -> object: + return load_jit( + "fast_index_copy_multi_slim", + *args, + cuda_files=["fast_index_copy.cuh"], + cuda_wrappers=[("launch", f"&MultiIndexCopySlimKernel<{args}>::run")], + build_directory=str(build_directory), + ) + + return KernelSpec(name=_make_name("fast_index_copy_multi_slim", *args), build=build) + + def _batch_memcpy_spec() -> KernelSpec: def build(build_directory: pathlib.Path) -> object: return load_jit( @@ -156,6 +173,8 @@ def default_kernel_specs() -> tuple[KernelSpec, ...]: specs.append(_fast_index_copy_multi_spec(num_threads=1024, blocks_per_bank=32)) # prefill hit-D2D gather (HBM-bound: wide grid) + its miss-side batch H2D binding. specs.append(_fast_index_copy_multi_spec(num_threads=1024, blocks_per_bank=64)) + # opt-in decode copy (FREETOKEN_MOE_SLIM_COPY=1) + specs.append(_fast_index_copy_multi_slim_spec()) specs.append(_batch_memcpy_spec()) specs.append(_radix_spec()) diff --git a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh index 43f3fc1ba5..e37e3d48ad 100644 --- a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh +++ b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh @@ -595,3 +595,112 @@ struct MultiIndexCopyKernel { device.unwrap())(kernel, params); } }; + + +// --------------------------------------------------------------------------- +// Slim multi-bank copy: the fused copy from a small grid, so it can share the GPU with compute. The +// grid strides each (row, bank) segment in turn, carrying the phase across, so no bank runs alone at the tail. +template +__always_inline __device__ void slim_copy_segment( + uint4* __restrict__ dst, const uint4* __restrict__ src, int32_t count, int32_t i, int32_t stride +) { + constexpr int32_t kU = static_cast(kUnroll); + // kUnroll independent loads in flight per thread: host-link latency, not issue rate, bounds a small grid + for (; i + (kU - 1) * stride < count; i += kU * stride) { + uint4 v[kUnroll]; +#pragma unroll + for (int32_t k = 0; k < kU; ++k) { + v[k] = details::load_nc(src + i + k * stride); + } +#pragma unroll + for (int32_t k = 0; k < kU; ++k) { + dst[i + k * stride] = v[k]; + } + } + uint4 v[kUnroll]; +#pragma unroll + for (int32_t k = 0; k < kU; ++k) { + if (i + k * stride < count) { + v[k] = details::load_nc(src + i + k * stride); + } + } +#pragma unroll + for (int32_t k = 0; k < kU; ++k) { + if (i + k * stride < count) { + dst[i + k * stride] = v[k]; + } + } +} + +template +__global__ __launch_bounds__(kNumThreads) void fast_index_copy_multi_slim( + const __grid_constant__ MultiIndexCopyParams p +) { + const int64_t n = p.valid_length ? p.valid_length[0] : p.length; + const auto stride = static_cast(gridDim.x * kNumThreads); + const auto tid = static_cast(blockIdx.x * kNumThreads + threadIdx.x); + const auto* di = static_cast(p.dst_indices); + const auto* si = static_cast(p.src_indices); + int32_t phase = 0; // grid thread that takes the next segment's first chunk + for (int64_t row = 0; row < n; ++row) { + const auto ps = static_cast(si[row]); + const auto pd = static_cast(di[row]); + for (int b = 0; b < p.num_banks; ++b) { + const int64_t feat = p.feat_bytes[b]; + const auto units = static_cast(feat >> 4); // feat % 16 == 0; int32 chunk index: rows < 32 GiB + const auto* src = reinterpret_cast(reinterpret_cast(p.src_ptrs[b]) + ps * feat); + auto* dst = reinterpret_cast(reinterpret_cast(p.dst_ptrs[b]) + pd * feat); + const int32_t first = tid >= phase ? tid - phase : tid - phase + stride; + slim_copy_segment(dst, src, units, first, stride); + phase = static_cast((static_cast(phase) + units) % stride); + } + } +} + +template +struct MultiIndexCopySlimKernel { + static void run( + tvm::ffi::TensorView dst_ptrs, + tvm::ffi::TensorView src_ptrs, + tvm::ffi::TensorView feat_bytes, + tvm::ffi::TensorView dst_indices, + tvm::ffi::TensorView src_indices, + tvm::ffi::Optional num_indices + ) { + using namespace host; + auto device = SymbolicDevice{}; + auto B = SymbolicSize{"num_banks"}; + auto L = SymbolicSize{"indices length"}; + auto ptr_dtype = SymbolicDType{}; + auto indices_dtype = SymbolicDType{}; + auto num_indices_dtype = SymbolicDType{}; + + TensorMatcher({B}).with_dtype(ptr_dtype).with_device(device) + .verify(dst_ptrs).verify(src_ptrs).verify(feat_bytes); + TensorMatcher({L}).with_dtype(indices_dtype).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) + .verify(num_indices.value()); + valid_length = static_cast(num_indices.value().data_ptr()); + } + + const auto params = MultiIndexCopyParams{ + static_cast(dst_ptrs.data_ptr()), + static_cast(src_ptrs.data_ptr()), + static_cast(feat_bytes.data_ptr()), + dst_indices.data_ptr(), + src_indices.data_ptr(), + valid_length, + static_cast(L.unwrap()), + static_cast(B.unwrap()), + }; + const auto use_int32 = indices_dtype.unwrap().bits == 32; + const auto kernel = use_int32 + ? fast_index_copy_multi_slim + : fast_index_copy_multi_slim; + LaunchKernel(kNumBlocks, kNumThreads, device.unwrap())(kernel, params); + } +}; diff --git a/python/freetoken/kernel/fast_index_copy.py b/python/freetoken/kernel/fast_index_copy.py index b2030620ee..3047947c78 100644 --- a/python/freetoken/kernel/fast_index_copy.py +++ b/python/freetoken/kernel/fast_index_copy.py @@ -14,6 +14,11 @@ DEFAULT_NUM_BLOCKS = 4 +# fast_index_copy_multi_slim grid (8K loads in flight, the least that holds ~51 GB/s on an +# RTX 5080 at 4 experts); the AOT spec (kernel/aot.py) builds exactly this config +SLIM_COPY_BLOCKS = 8 +SLIM_COPY_THREADS = 256 +SLIM_COPY_UNROLL = 4 SKIP_FAST_INDEX_COPY_ENV = "FREETOKEN_SKIP_FAST_INDEX_COPY" _TRUE_VALUES = {"1", "true", "yes", "on"} @@ -202,6 +207,43 @@ def fast_index_copy_multi_jit( module.launch(dst_ptrs, src_ptrs, feat_bytes, dst_indices, src_indices, num_indices) +@lru_cache(maxsize=None) +def _jit_fast_index_copy_multi_slim_module(*, num_blocks: int, num_threads: int, unroll: int) -> Module: + args = make_cpp_args(num_blocks, num_threads, unroll) + return load_jit( + "fast_index_copy_multi_slim", + *args, + cuda_files=["fast_index_copy.cuh"], + cuda_wrappers=[("launch", f"&MultiIndexCopySlimKernel<{args}>::run")], + ) + + +def fast_index_copy_multi_slim_jit( + dst_ptrs: torch.Tensor, + src_ptrs: torch.Tensor, + feat_bytes: torch.Tensor, + dst_indices: torch.Tensor, + src_indices: torch.Tensor, + num_indices: torch.Tensor | None = None, + *, + num_blocks: int = SLIM_COPY_BLOCKS, + num_threads: int = SLIM_COPY_THREADS, + unroll: int = SLIM_COPY_UNROLL, +) -> None: + """:func:`fast_index_copy_multi_jit` from a small fixed grid, for copies that run beside compute. + + ``num_blocks`` CTAs split the concatenated (row, bank) 16-byte chunk range evenly and each + thread keeps ``unroll`` loads in flight, so a few SMs still saturate the host link while + the rest stay free for concurrent kernels. Same arguments and alignment contract. + """ + if _skip_fast_index_copy_enabled(): + return + module = _jit_fast_index_copy_multi_slim_module( + num_blocks=num_blocks, num_threads=num_threads, unroll=unroll + ) + module.launch(dst_ptrs, src_ptrs, feat_bytes, dst_indices, src_indices, num_indices) + + def update_copy_flag_jit(sync_flag: torch.Tensor, delta: int) -> None: assert sync_flag.is_cuda assert sync_flag.numel() == 1 diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index 799458dc06..918a3f9731 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -241,7 +241,19 @@ def decode_forward( topk=self.top_k, renormalize=self.renormalize, ) - return self._decode_routed(hidden_states, topk_weights, topk_ids) + return self._decode_routed(hidden_states, topk_weights, topk_ids, router_logits=router_logits) + + def decode_side_applies(self) -> bool: + """Whether FREETOKEN_MOE_COPY_OVERLAP applies to this layer's next forward: a GPU decode on a + cache with a side stream (not prefill, CPU/hybrid decode or TP > 1); see decode_side.""" + cache = self.offload_cache + return not ( + cache is None + or cache.decode_copy_stream is None + or self.tp_size > 1 + or get_global_ctx().batch.is_prefill + or cache.is_cpu_layer(self.layer_id) + ) def prefill_forward( self, @@ -269,6 +281,7 @@ def _decode_routed( hidden_states: torch.Tensor, topk_weights: torch.Tensor, topk_ids: torch.Tensor, + router_logits: torch.Tensor | None = None, ) -> torch.Tensor: """On-demand load: ``ensure_experts`` rewrites ``topk_ids`` into cache slot ids in place (loading missing experts), then the GEMM reads the full slot @@ -288,7 +301,7 @@ def _decode_routed( return executor.decode(self.layer_id, hidden_states, topk_weights, topk_ids) if cache.decode_target == "hybrid": return self._decode_hybrid(cache, hidden_states, topk_weights, topk_ids) - cache.ensure_experts(self.layer_id, topk_ids) + cache.ensure_experts(self.layer_id, topk_ids, router_logits=router_logits) cache.copy_missing() return self._expert_gemm( cache, @@ -423,9 +436,11 @@ def _ensure_unique(self, cache: OffloadMoeCache, topk_ids: torch.Tensor) -> torc if n > cache.cache_size // 2: # leave room so a chunk never evicts the one before it return None uniq = uniq.to(torch.int32) + # scored eviction ranks an earlier chunk's fresh installs coldest, so pin them explicitly + pin_since = cache.step.clone() if cache.cache_policy_id else None for start in range(0, n, _ENSURE_CHUNK_IDS): part = uniq[start : start + _ENSURE_CHUNK_IDS] # a contiguous view: ensure rewrites it in place - cache.ensure_experts(self.layer_id, part) + cache.ensure_experts(self.layer_id, part, update_state=False, pin_since=pin_since) # not a decode step cache.copy_missing() return uniq[inverse].view_as(topk_ids) diff --git a/python/freetoken/models/deepseek_v4/moe.py b/python/freetoken/models/deepseek_v4/moe.py index 6b6245674d..83e6585403 100644 --- a/python/freetoken/models/deepseek_v4/moe.py +++ b/python/freetoken/models/deepseek_v4/moe.py @@ -106,7 +106,7 @@ def _prefill_routed( or cache.is_unpinned_layer(self.layer_id) ): return super()._prefill_routed(hidden_states, topk_weights, topk_ids) - cache.ensure_experts(self.layer_id, topk_ids) # in-place expert-id -> slot + cache.ensure_experts(self.layer_id, topk_ids, update_state=False) # in-place expert-id -> slot; not a decode step cache.copy_missing() if cache.collect_stats: cache.record_decode_stats(self.layer_id) diff --git a/python/freetoken/models/qwen4_exp/moe.py b/python/freetoken/models/qwen4_exp/moe.py index c10407861f..6f729231f1 100644 --- a/python/freetoken/models/qwen4_exp/moe.py +++ b/python/freetoken/models/qwen4_exp/moe.py @@ -35,9 +35,19 @@ def __init__( def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: num_tokens, hidden_dim = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_dim) + se, ex = self.shared_expert, self.experts + if isinstance(ex, OffloadMoELayer) and ex.decode_side_applies(): + # FREETOKEN_MOE_COPY_OVERLAP: the shared expert runs on a side stream while the router, + # ensure and miss copy stay back to back on this one; joined before the combine + cache = ex.offload_cache + with cache.decode_side(ex.layer_id): + gate = shared_gate_sigmoid(hidden_states, self.shared_expert_gate.weight.view(-1)) + shared = se.forward(hidden_states) + routed = ex.forward(hidden_states=hidden_states, router_logits=self.gate.forward(hidden_states)) + cache.join_decode_side(ex.layer_id) + return shared_gate_mul_add(routed, shared, gate).view(num_tokens, hidden_dim) router_logits = self.gate.forward(hidden_states) gate = shared_gate_sigmoid(hidden_states, self.shared_expert_gate.weight.view(-1)) - se, ex = self.shared_expert, self.experts if ( self._tp_size > 1 and isinstance(se.down_proj, LinearRowParallel) diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 8730c8d734..477504bcf3 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -2,6 +2,7 @@ import math import os +from contextlib import contextmanager from dataclasses import dataclass from typing import Iterator @@ -13,6 +14,9 @@ # (kept for A/B profiling). Falls back to per-bank automatically if a bank's row bytes or # base address are not 16-byte aligned. _FUSED_COPY = os.getenv("FREETOKEN_FUSED_COPY", "1").strip().lower() not in {"0", "false", "no", "off"} +# Opt-in: run the fused copy_missing on the slim kernel (a few small CTAs, leaves most SMs to +# concurrent compute) instead of the 8-CTAs-per-bank one. Off by default. +_SLIM_COPY = os.getenv("FREETOKEN_MOE_SLIM_COPY", "0").strip().lower() in {"1", "true", "yes", "on"} # cudaMemcpyBatchAsync silently degrades to a SYNCHRONOUS copy when a batch mixes # large entries with sub-~256KB entries on registered host memory (H100 + CUDA 13.0, @@ -26,6 +30,7 @@ _SMALL_BANK_FEAT_BYTES = 256 * 1024 from freetoken.kernel.fast_index_copy import default_blocks_per_bank +from freetoken.env import ENV from freetoken.utils import init_logger logger = init_logger(__name__) @@ -145,11 +150,20 @@ class OffloadMoeCache: # bank layout from the expert kernel (a BankSpec per role); when given it replaces the _BANK_SCHEMAS lookup and the slot cap comes from max_slots layout: dict | None = None max_slots: int | None = None + # GPU decode only: the caller's work that is independent of the routed experts (qwen4_exp: the + # shared expert) runs on decode_copy_stream beside the miss copy. None = FREETOKEN_MOE_COPY_OVERLAP. + decode_copy_overlap: bool | None = None def __post_init__(self) -> None: - policy_ids = {"lru": 0} - assert self.cache_policy in policy_ids + from freetoken.moe.scored_ensure import POLICY_IDS + + assert self.cache_policy in POLICY_IDS, self.cache_policy assert self.decode_target in ("gpu", "cpu", "hybrid"), self.decode_target + if self.cache_policy != "lru" and self.decode_target != "gpu": + # the hybrid and CPU decode paths run their own kernels and would silently stay LRU + raise ValueError( + f"--moe-cache-policy {self.cache_policy} needs GPU decode, not {self.decode_target!r}" + ) if self.layout is None: assert self.quant_format in _BANK_SCHEMAS, f"unknown quant_format {self.quant_format!r}" # Attached by the engine for decode_target == "cpu" (CpuMoeExecutor); None @@ -166,7 +180,11 @@ def __post_init__(self) -> None: "cache, so cache_size must be at least 2 * num_experts " "(raise moe_cache_size or disable moe_prefill_overlap)" ) - self.cache_policy_id = policy_ids[self.cache_policy] + self.cache_policy_id = POLICY_IDS[self.cache_policy] + # rule: an unrouted expert within this many logits of the top-k-th is a near miss + self.near_miss_thr = ( + float(os.getenv("FREETOKEN_MOE_NEAR_MISS_THR", "0.25")) if self.cache_policy == "rule" else None + ) self.slot_for_id = torch.full( (self.num_layers, self.num_experts), -1, @@ -200,6 +218,7 @@ def __post_init__(self) -> None: self.expert_recency = torch.full( (self.num_layers, self.num_experts), -1, dtype=torch.int64, device=self.device ) + self._init_evict_state() # Host source banks (one [num_experts, ...] tensor per layer, so layers can # carry independent host attributes -- see layer_residency) and their GPU # slot caches, keyed by the format's bank schema (attached by @@ -290,6 +309,87 @@ def __post_init__(self) -> None: self._batch_memcpy = None self.prefill_hit_rows = 0 self.prefill_total_rows = 0 + # Decode copy overlap: one side stream plus a (forked, done) event pair per layer. + # Neither is cache_size-shaped, so rebuild() keeps them. + if self.decode_copy_overlap is None: + self.decode_copy_overlap = bool(ENV.MOE_COPY_OVERLAP) + self.decode_copy_stream: torch.cuda.Stream | None = None + self._decode_copy_events: list[tuple[torch.cuda.Event, torch.cuda.Event]] = [] + if self.decode_copy_overlap and self.device.type == "cuda" and self.decode_target == "gpu": + self._init_decode_copy_overlap() + + def _init_decode_copy_overlap(self) -> None: + # a dedicated stream, not one of torch's 32 pooled ones: a pooled stream can alias the engine + # or capture stream, and the side work would then run serially (correct, but no overlap) + from cuda.bindings import runtime as cudart + + with torch.cuda.device(self.device): + err, handle = cudart.cudaStreamCreateWithFlags(cudart.cudaStreamNonBlocking) + if err != cudart.cudaError_t.cudaSuccess: + raise RuntimeError(f"cudaStreamCreateWithFlags failed: {err}") + self.decode_copy_stream = torch.cuda.ExternalStream(int(handle), device=self.device) + self._decode_copy_events = [ + (torch.cuda.Event(), torch.cuda.Event()) for _ in range(self.num_layers) + ] + # torch creates the CUDA event on its first record: do that now, never inside a capture + for pair in self._decode_copy_events: + for event in pair: + event.record(self.decode_copy_stream) + logger.info("MoE decode copy overlap on: expert misses copy on a side stream") + + def _init_evict_state(self) -> None: + """Scored policies' per-(layer, expert) state (moe/scored_ensure.py); None under LRU. + + It is not slot-shaped: it survives eviction and prefill invalidation, and a rebuild + only resets it. ``evict_tok`` is the decode-step clock the first ensure of each step bumps. + """ + from freetoken.moe import scored_ensure as se + + self.evict_tok = self.evict_last_tok = self.evict_lc = self.evict_ct = None + self.evict_g_table = None + n, dev = self.num_layers * self.num_experts, self.device + if self.cache_policy_id >= 1: + self.evict_tok = torch.zeros((), dtype=torch.int64, device=dev) + self.evict_last_tok = torch.full((n,), se.LAST_TOK_NEVER, dtype=torch.int64, device=dev) + if self.cache_policy_id >= 2: + self.evict_lc = torch.full((n,), se.LC_NEVER, dtype=torch.int32, device=dev) + self.evict_ct = torch.zeros((n,), dtype=torch.int64, device=dev) + self.evict_g_table = se.softplus2_table(dev) + self._alloc_evict_mirrors() + + def _alloc_evict_mirrors(self) -> None: + """Slot-shaped copies of each slot owner's state, read coalesced by the victim scan. + + A mirror is valid while ``evict_slot_owner`` equals ``id_of_slot``; the kernel keeps it + so, and anything else that installs an id into a slot must set the owner to -1. + """ + self.evict_slot_owner = self.evict_slot_last_tok = self.evict_slot_lc = self.evict_slot_ct = None + if self.evict_tok is None: + return + size, dev = self.cache_size, self.device + self.evict_slot_owner = torch.full((size,), -1, dtype=torch.int32, device=dev) + self.evict_slot_last_tok = torch.zeros((size,), dtype=torch.int64, device=dev) + if self.evict_lc is not None: + self.evict_slot_lc = torch.zeros((size,), dtype=torch.int32, device=dev) + self.evict_slot_ct = torch.zeros((size,), dtype=torch.int64, device=dev) + + def reset_evict_state(self) -> None: + if self.evict_tok is None: + return + from freetoken.moe import scored_ensure as se + + self.evict_tok.zero_() + self.evict_last_tok.fill_(se.LAST_TOK_NEVER) + self.evict_slot_owner.fill_(-1) + if self.evict_lc is not None: + self.evict_lc.fill_(se.LC_NEVER) + self.evict_ct.zero_() + + def evict_clock_layer(self) -> int: + """The first layer that ensures on the GPU each decode step; its call bumps ``evict_tok``.""" + if not self.cpu_layer_ids: + return 0 + return min(set(range(self.num_layers)) - self.cpu_layer_ids, default=0) def set_bank_sources( self, @@ -512,6 +612,8 @@ def rebuild(self, cache_size: int) -> None: self.num_indices.zero_() self.num_missing_full.zero_() self.expert_recency.fill_(-1) + self._alloc_evict_mirrors() + self.reset_evict_state() self.stat_missing.zero_() self.stat_active.zero_() self.stat_calls.zero_() @@ -846,8 +948,21 @@ def release_prefill_layer(self, layer_id: int) -> None: self._prefill_buffer_has_release_event[buffer_id] = True self._prefill_buffer_released[buffer_id] = True - def ensure_experts(self, layer_id: int, expert_ids: torch.Tensor) -> None: - from freetoken.moe.offload_kernels import ensure_experts + def ensure_experts( + self, + layer_id: int, + expert_ids: torch.Tensor, + *, + update_state: bool = True, + router_logits: torch.Tensor | None = None, + pin_since: torch.Tensor | None = None, + ) -> None: + """``update_state=False`` (small-prefill ensures) keeps a scored policy's state and + decode-step clock still; ``router_logits`` (``[rows, num_experts]``) feed the ``rule`` + policy's near misses; ``pin_since`` (a ``step`` value) keeps every slot touched after it + (a small prefill's earlier chunks). All are ignored under LRU, whose recency order + already keeps those slots.""" + from freetoken.moe.offload_kernels import ensure_experts, ensure_experts_scored if self.collect_decode_freq: # ``expert_ids`` still holds raw expert ids here (the kernel rewrites them to @@ -856,7 +971,18 @@ def ensure_experts(self, layer_id: int, expert_ids: torch.Tensor) -> None: self.decode_freq[layer_id].scatter_add_(0, ids, torch.ones_like(ids)) self._pending_src_layer = layer_id self._pending_whole_layer = False - ensure_experts(self, layer_id, expert_ids) + if self.cache_policy_id == 0: + ensure_experts(self, layer_id, expert_ids) + return + ensure_experts_scored( + self, + layer_id, + expert_ids, + bump_tok=update_state and layer_id == self.evict_clock_layer(), + update_state=update_state, + router_logits=router_logits, + pin_since=pin_since, + ) def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None: """Capped-fetch LRU for the hybrid backend. @@ -885,6 +1011,8 @@ def materialize_layer(self, layer_id: int) -> None: self._pending_src_layer = layer_id self._pending_whole_layer = True materialize_layer(self, layer_id) + if self.evict_slot_owner is not None: + self.evict_slot_owner[: self.num_experts].fill_(-1) # it installs the layer into these slots def reset(self) -> None: from freetoken.moe.offload_kernels import reset_cache @@ -893,6 +1021,7 @@ def reset(self) -> None: # 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.reset_evict_state() def reset_stats(self) -> None: self.prefill_hit_rows = 0 @@ -1014,9 +1143,12 @@ def decode_routing_stats(self) -> dict: "norm_entropy": norm_ent, } - def copy_missing(self) -> None: + def copy_missing(self, layer_id: int | None = None) -> None: + """Copy the misses the last ensure_experts/materialize_layer staged, from ``layer_id``'s + host banks (default: the layer that staged them).""" assert self.banks, "set_bank_sources must register the banks first" - layer_id = self._pending_src_layer + if layer_id is None: + 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._unpinned_layers: if not self._pending_whole_layer: @@ -1031,13 +1163,14 @@ def copy_missing(self) -> None: cache[: self.num_experts].copy_(per_layer[layer_id]) return if self._copy_fused_ok: - from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit + from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit, fast_index_copy_multi_slim_jit # One launch copies the missing rows for every bank (instead of one launch per # bank). evict_slots/src_indices/num_indices are shared across banks; # src_indices holds layer-local expert rows, resolved against this layer's # source pointers (layer_id is a static int per captured graph node). - fast_index_copy_multi_jit( + copy = fast_index_copy_multi_slim_jit if _SLIM_COPY else fast_index_copy_multi_jit + copy( self._copy_dst_ptrs, self._copy_src_ptrs[layer_id], self._copy_feat_bytes, @@ -1059,6 +1192,24 @@ def copy_missing(self) -> None: self.num_indices, ) + @contextmanager + def decode_side(self, layer_id: int) -> Iterator[None]: + """Run the body on ``decode_copy_stream``, ordered after the work already on the current + stream; ``join_decode_side`` must follow before its results are read. + + The copy stays on the current stream right after its ensure, so it never waits on a + cross-stream edge; the side work (slack-rich) takes that latency and any SM contention. + Under capture the event pair becomes a fork and a join edge of the graph.""" + forked, done = self._decode_copy_events[layer_id] + forked.record(torch.cuda.current_stream(self.device)) + self.decode_copy_stream.wait_event(forked) + with torch.cuda.stream(self.decode_copy_stream): + yield + done.record(self.decode_copy_stream) + + def join_decode_side(self, layer_id: int) -> None: + torch.cuda.current_stream(self.device).wait_event(self._decode_copy_events[layer_id][1]) + def iter_offload_moe_layers(model) -> Iterator: from freetoken.layers import OffloadMoELayer, iter_moe_layers diff --git a/python/freetoken/moe/offload_kernels.py b/python/freetoken/moe/offload_kernels.py index cf513f52d1..bb03aae14c 100644 --- a/python/freetoken/moe/offload_kernels.py +++ b/python/freetoken/moe/offload_kernels.py @@ -40,6 +40,51 @@ def ensure_experts(cache, layer_id: int, expert_ids: torch.Tensor) -> None: ) +def ensure_experts_scored( + cache, + layer_id: int, + expert_ids: torch.Tensor, + *, + bump_tok: bool, + update_state: bool, + router_logits: torch.Tensor | None = None, + pin_since: torch.Tensor | None = None, +) -> None: + """``ensure_experts`` for the scored ``--moe-cache-policy`` kinds (kd, kdfb, rule).""" + from freetoken.moe.scored_ensure import scored_ensure + + scored_ensure( + expert_ids, + cache.slot_for_id.view(-1), + cache.id_of_slot, + cache.usage, + cache.step, + expert_ids, + cache.src_indices, + cache.evict_slots, + cache.num_indices, + policy=cache.cache_policy_id, + num_layers=cache.num_layers, + num_experts=cache.num_experts, + tok=cache.evict_tok, + last_tok=cache.evict_last_tok, + lc=cache.evict_lc, + ct=cache.evict_ct, + g_table=cache.evict_g_table, + slot_owner=cache.evict_slot_owner, + slot_last_tok=cache.evict_slot_last_tok, + slot_lc=cache.evict_slot_lc, + slot_ct=cache.evict_slot_ct, + router_logits=router_logits if cache.near_miss_thr is not None else None, + near_miss_thr=cache.near_miss_thr or 0.0, + stats=cache.lru_stats[layer_id] if cache.collect_stats else None, + id_base=layer_id * cache.num_experts, + bump_tok=bump_tok, + update_state=update_state, + pin_since=pin_since, + ) + + def ensure_experts_hybrid( cache, layer_id: int, expert_ids: torch.Tensor, max_fetch: int, fetch_fraction: float = 0.0 ) -> None: diff --git a/python/freetoken/moe/scored_ensure.py b/python/freetoken/moe/scored_ensure.py new file mode 100644 index 0000000000..123086ed51 --- /dev/null +++ b/python/freetoken/moe/scored_ensure.py @@ -0,0 +1,395 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright 2026 The FlashLib Authors +# Adapted from flashlib 0.3.0 (kernels/slot_cache/triton/lru_ensure.py: _phase1, _stats, +# _lru_ensure_kernel, _num_warps_for), https://github.com/FlashML-org/flashlib +"""Scored-victim slot-cache admission for the MoE expert cache (``--moe-cache-policy``). + +flashlib's sequential-argmin ``lru_ensure`` with the victim key made pluggable. POLICY 0 is +LRU and matches flashlib bit for bit; the others evict the lowest packed score key: + +- 1 ``kd``: ``-(k + beta * d / L)``. ``k`` is the tokens the expert has been passed over + since its last use, ``d`` the layer-steps until its layer comes up again. +- 2 ``kdfb``: kd plus ``w * log2(decayed use count)`` with a ``halflife``-token decay. +- 3 ``rule``: kdfb plus the near-miss refresh: this layer's experts whose router logit is + within a margin of the top-k-th count as used one token ago (``last_tok``), unpinned. + +The per-(layer, expert) state survives eviction; slot-indexed mirrors of it (tagged with the +owning id) let the victim scan read it coalesced instead of gathering by id. Scores and +counts are int32 Q16 fixed point and the count update ``log2(2^x + 1)`` reads a host-built +table, so a CPU reference can reproduce every victim exactly (``tl.exp2``/``tl.log2`` are +approximate). +""" +from __future__ import annotations + +import math + +import torch +import triton +import triton.language as tl + +POLICY_IDS = {"lru": 0, "kd": 1, "kdfb": 2, "rule": 3} + +# Replay-tuned on sep5_s01 at 1650 slots; the plateau is flat for beta 1-8, w 2-4, halflife 32-128. +BETA = 4.0 +W = 3.0 +HALFLIFE = 64 + +Q = 16 # fraction bits of scores and counts +LAST_TOK_NEVER = -1_000_000 +LC_NEVER = -(1 << 30) +LC_FLOOR = -(30 << Q) +K_MAX = 1 << 14 # tokens; with the caps on beta and w below, a score fits int32 +MAX_BETA, MAX_W = 64.0, 16.0 +DECAY_SHIFT = 8 +SCORE_BIAS = 1 << 31 +MAX_SLOT_BITS = 30 +# g(u) = log2(1 + 2^-u) sampled every 2^-G_STEP_BITS for u in [0, G_RANGE]; zero beyond +G_STEP_BITS = 6 +G_RANGE = 24 +G_LAST = G_RANGE << G_STEP_BITS +# jit code may only read globals that are constexpr +_Q = tl.constexpr(Q) +_LC_NEVER = tl.constexpr(LC_NEVER) +_LC_FLOOR = tl.constexpr(LC_FLOOR) +_K_MAX = tl.constexpr(K_MAX) +_DECAY_SHIFT = tl.constexpr(DECAY_SHIFT) +_SCORE_BIAS = tl.constexpr(SCORE_BIAS) +_G_SHIFT = tl.constexpr(Q - G_STEP_BITS) +_G_LAST = tl.constexpr(G_LAST) + + +def softplus2_table(device=None) -> torch.Tensor: + """Q16 ``log2(1 + 2^-u)`` at ``u = i / 64``, two entries past ``G_LAST`` for the interpolation.""" + step = 1 << G_STEP_BITS + vals = [round(math.log2(1.0 + 2.0 ** (-i / step)) * (1 << Q)) for i in range(G_LAST + 2)] + return torch.tensor(vals, dtype=torch.int32, device=device) + + +def score_params(beta: float, w: float, halflife: int, num_layers: int) -> tuple[int, int, int, int]: + """``(BETA_STEP, W_Q4, DECAY_MUL, DT_MAX)`` shared by the kernel and the CPU reference. + + ``beta * d / L`` is ``d * BETA_STEP`` in Q16; the decay over ``dt`` layer-steps is + ``(dt * DECAY_MUL) >> DECAY_SHIFT``, with ``dt`` capped at 64 half-lives (past the -30 floor). + """ + assert 0 <= beta <= MAX_BETA and 0 <= w <= MAX_W and halflife >= 1 + span = num_layers * int(halflife) + decay_mul = round((1 << (Q + DECAY_SHIFT)) / span) + dt_max = 64 * span + assert dt_max * decay_mul < (1 << 31) + return round(beta * (1 << Q) / num_layers), round(w * 16), decay_mul, dt_max + + +@triton.jit +def _phase1(query_ptr, slot_of_id_ptr, lru_usage_ptr, num_copy_ptr, step, K, + BLOCK_K: tl.constexpr, id_base): + """Vendored unchanged: dedup the query, split hit/miss, rank the misses, bump the hits.""" + k = tl.arange(0, BLOCK_K) + kmask = k < K + q = tl.load(query_ptr + k, mask=kmask, other=-1) + id_base + s = tl.load(slot_of_id_ptr + q, mask=kmask, other=-1) + hit = kmask & (s >= 0) + miss = kmask & (s == -1) + same = ( + (q[:, None] == q[None, :]) + & (k[:, None] > k[None, :]) + & kmask[:, None] + & kmask[None, :] + ) + first = kmask & (tl.sum(same.to(tl.int32), axis=1) == 0) + first_miss = miss & first + smaller = (q[None, :] < q[:, None]) & first_miss[None, :] + rank = tl.sum(smaller.to(tl.int32), axis=1) + num_missing = tl.sum(first_miss.to(tl.int32)) + tl.store(num_copy_ptr, num_missing.to(tl.int64)) + # Duplicated hits write the same value to the same slot -- idempotent. + tl.store(lru_usage_ptr + s, step, mask=hit) + return q, kmask, miss, first, first_miss, rank, num_missing, tl.where(hit, s, -1) + + +@triton.jit +def _stats(stats_ptr, first, num_missing): + # One vectorized atomic over 3 lanes, not three scalar ones: a scalar atomic serializes the CTA. + si = tl.arange(0, 4) + v = tl.where(si == 0, tl.sum(first.to(tl.int32)), tl.where(si == 1, num_missing, 1)) + # relaxed: only later kernels read the counters; flashlib's acq_rel waits on this call's stores (~0.3 us) + tl.atomic_add(stats_ptr + si, v.to(tl.int64), mask=si < 3, sem="relaxed") + + +@triton.jit +def _decayed(lc, ct, now, DECAY_MUL: tl.constexpr, DT_MAX: tl.constexpr): + """int32 Q16 log2 count decayed from ``ct`` to ``now`` (int64 layer-steps).""" + dt = tl.minimum(tl.maximum(now - ct, 0), DT_MAX).to(tl.int32) + return lc - ((dt * DECAY_MUL) >> _DECAY_SHIFT) + + +@triton.jit +def _softplus2(x, g_ptr, mask): + """Q16 ``log2(2^x + 1)`` = ``max(x, 0) + g(|x|)``, g linearly interpolated from the table.""" + u = tl.abs(x) + i = tl.minimum(u >> _G_SHIFT, _G_LAST) + frac = u & ((1 << _G_SHIFT) - 1) + g0 = tl.load(g_ptr + i, mask=mask, other=0) + g1 = tl.load(g_ptr + i + 1, mask=mask, other=0) + g = tl.where(i < _G_LAST, g0 + (((g1 - g0) * frac) >> _G_SHIFT), 0) + return tl.maximum(x, 0) + g + + +@triton.jit +def _owner_state(mirror_ptr, state_ptr, c, oid, mirrored, stale, other): + """The slot owner's state: its coalesced mirror, or a gather by id where the mirror is stale.""" + return tl.where( + mirrored, + tl.load(mirror_ptr + c, mask=mirrored, other=other), + tl.load(state_ptr + oid, mask=stale, other=other), + ) + + +@triton.jit(do_not_specialize=["K", "num_cached", "id_base", "nm_topk", "nm_stride"]) +def _scored_ensure_kernel( + query_ptr, slot_of_id_ptr, id_of_slot_ptr, lru_usage_ptr, lru_step_ptr, + out_ptr, src_ptr, dst_ptr, num_copy_ptr, stats_ptr, + tok_ptr, last_tok_ptr, lc_ptr, ct_ptr, g_ptr, logits_ptr, + owner_ptr, slot_last_ptr, slot_lc_ptr, slot_ct_ptr, pin_ptr, + K, num_cached, id_base, nm_topk, nm_stride, nm_thr, + BLOCK_K: tl.constexpr, BLOCK_C: tl.constexpr, BLOCK_E: tl.constexpr, BLOCK_TOPK: tl.constexpr, + NM_ROWS: tl.constexpr, + USAGE_MAX: tl.constexpr, COLLECT_STATS: tl.constexpr, + POLICY: tl.constexpr, BUMP_TOK: tl.constexpr, UPDATE_STATE: tl.constexpr, PIN_SINCE: tl.constexpr, + NUM_LAYERS: tl.constexpr, NUM_EXPERTS: tl.constexpr, + BETA_STEP: tl.constexpr, W_Q4: tl.constexpr, DECAY_MUL: tl.constexpr, DT_MAX: tl.constexpr, + SLOT_BITS: tl.constexpr, +): + step = tl.load(lru_step_ptr) + 1 + tl.store(lru_step_ptr, step) + if POLICY != 0: + tok = tl.load(tok_ptr) + if BUMP_TOK: + tok = tok + 1 + tl.store(tok_ptr, tok) + layer = id_base // NUM_EXPERTS + now = tok * NUM_LAYERS + layer + if POLICY == 3: + if UPDATE_STATE: + # Routed ids qualify too; they are pinned now and set to tok below, so that is harmless. + ex = tl.arange(0, BLOCK_E) + j = tl.arange(0, BLOCK_TOPK) + prev = tl.load(last_tok_ptr + id_base + ex, mask=ex < NUM_EXPERTS, other=0) + held_at = tl.load(slot_of_id_ptr + id_base + ex, mask=ex < NUM_EXPERTS, other=-1) + near = ex < 0 + for r in tl.static_range(NM_ROWS): + routed = tl.load(query_ptr + r * nm_topk + j, mask=j < nm_topk, other=-1) + row_ptr = logits_ptr + r * nm_stride + kth = tl.min(tl.load(row_ptr + routed, mask=routed >= 0, other=float("inf")).to(tl.float32), axis=0) + row = tl.load(row_ptr + ex, mask=ex < NUM_EXPERTS, other=float("-inf")).to(tl.float32) + near = near | (row >= kth - nm_thr) + refreshed = tl.maximum(prev, tok - 1) + tl.store(last_tok_ptr + id_base + ex, refreshed, mask=near) + tl.store(slot_last_ptr + held_at, refreshed, mask=near & (held_at >= 0)) + # the victim scan and the routed update below read and rewrite what this wrote + tl.debug_barrier() + q, kmask, miss, first, first_miss, rank, num_missing, out = _phase1( + query_ptr, slot_of_id_ptr, lru_usage_ptr, num_copy_ptr, step, K, BLOCK_K, + id_base) + + if num_missing > 0: + # REQUIRED (flashlib): the reload below must see _phase1's hit bump, or a hit is evicted. + tl.debug_barrier() + c = tl.arange(0, BLOCK_C) + cmask = c < num_cached + u = tl.load(lru_usage_ptr + c, mask=cmask, other=USAGE_MAX) + if POLICY == 0: + umax = tl.full([BLOCK_C], USAGE_MAX, u.dtype) + u = tl.where((u == step) | (~cmask), umax, u) + else: + oid = tl.load(id_of_slot_ptr + c, mask=cmask, other=-1) + held = cmask & (oid >= 0) + owner = tl.load(owner_ptr + c, mask=cmask, other=-1) + mirrored = held & (owner == oid) + stale = held & (owner != oid) + lk = oid // NUM_EXPERTS + ahead = lk > layer + d = tl.where(ahead, lk - layer, NUM_LAYERS - layer + lk) + last = _owner_state(slot_last_ptr, last_tok_ptr, c, oid, mirrored, stale, 0) + k = tl.minimum(tl.maximum(tl.where(ahead, tok - 1, tok) - last, 0), _K_MAX).to(tl.int32) + # int32 throughout: |score| < 2^31 by the caps on k, beta and w + score = -((k << _Q) + d * BETA_STEP) + if POLICY >= 2: + lc = _owner_state(slot_lc_ptr, lc_ptr, c, oid, mirrored, stale, _LC_NEVER) + ct = _owner_state(slot_ct_ptr, ct_ptr, c, oid, mirrored, stale, 0) + lcv = _decayed(lc, ct, now, DECAY_MUL, DT_MAX) + lcv = tl.where(lc == _LC_NEVER, _LC_FLOOR, tl.maximum(lcv, _LC_FLOOR)) + score += (W_Q4 * lcv) >> 4 + # Every key is distinct and ties go to the lowest slot; an empty slot keys below any held one. + key = tl.where(held, ((score.to(tl.int64) + _SCORE_BIAS) << SLOT_BITS) | c, c.to(tl.int64)) + evictable = cmask & (u != step) + if PIN_SINCE: + # the score ignores recency, so slots an earlier call touched since *pin_ptr stay pinned too + evictable = evictable & (u <= tl.load(pin_ptr)) + # the pinned key's slot bits are 0, so an all-pinned scan falls back to slot 0 like flashlib's argmin + pinned_key = 0x7FFFFFFFFFFFFFFF ^ ((1 << SLOT_BITS) - 1) + key = tl.where(evictable, key, pinned_key) + for i in tl.range(num_missing): + if POLICY == 0: + victim = tl.argmin(u, axis=0).to(tl.int32) + else: + victim = (tl.min(key, axis=0) & ((1 << SLOT_BITS) - 1)).to(tl.int32) + # Scalar load: victims are distinct, so no earlier iteration wrote this slot. + old = tl.load(id_of_slot_ptr + victim) + if old >= 0: + tl.store(slot_of_id_ptr + old, -1) + e = tl.sum(tl.where((rank == i) & first_miss, q, 0)) + tl.store(id_of_slot_ptr + victim, e) + tl.store(slot_of_id_ptr + e, victim) + tl.store(lru_usage_ptr + victim, step) + tl.store(dst_ptr + i, victim) + tl.store(src_ptr + i, e - id_base) # back to the caller's id space + out = tl.where((rank == i) & miss, victim, out) + if POLICY == 0: + u = tl.where(c == victim, umax, u) # claim in-register + else: + key = tl.where(c == victim, pinned_key, key) + + # Written from registers, never re-read from slot_of_id, so out_ptr may alias query_ptr. + tl.store(out_ptr + tl.arange(0, BLOCK_K), out, mask=kmask) + if POLICY != 0: + if UPDATE_STATE: + # distinct ids only, so a duplicate at bs > 1 counts once per call; out is each id's final slot + tl.store(last_tok_ptr + q, tok, mask=first) + tl.store(owner_ptr + out, q, mask=first) + tl.store(slot_last_ptr + out, tok, mask=first) + if POLICY >= 2: + lc = tl.load(lc_ptr + q, mask=first, other=_LC_NEVER) + x = _decayed(lc, tl.load(ct_ptr + q, mask=first, other=0), now, DECAY_MUL, DT_MAX) + lc_new = tl.where(lc == _LC_NEVER, 0, _softplus2(x, g_ptr, first)) + tl.store(lc_ptr + q, lc_new, mask=first) + tl.store(ct_ptr + q, now, mask=first) + tl.store(slot_lc_ptr + out, lc_new, mask=first) + tl.store(slot_ct_ptr + out, now, mask=first) + else: + # a slot filled without a state update reads the per-id state until a decode hit re-mirrors it + tl.store(owner_ptr + out, -1, mask=first_miss) + if COLLECT_STATS: + _stats(stats_ptr, first, num_missing) + + +def _num_warps_for(block_c: int) -> int: + # flashlib's rule: 8 warps from 2048 slots, more only past ~32 elements per thread. + warps = 8 if block_c >= 2048 else 4 + return max(warps, min(triton.next_power_of_2(max(block_c // 1024, 1)), 32)) + + +def scored_ensure( + query: torch.Tensor, + slot_of_id: torch.Tensor, + id_of_slot: torch.Tensor, + lru_usage: torch.Tensor, + lru_step: torch.Tensor, + out_indices: torch.Tensor, + src_indices: torch.Tensor, + dst_indices: torch.Tensor, + num_copy: torch.Tensor, + *, + policy: int, + num_layers: int, + num_experts: int, + tok: torch.Tensor | None = None, + last_tok: torch.Tensor | None = None, + lc: torch.Tensor | None = None, + ct: torch.Tensor | None = None, + g_table: torch.Tensor | None = None, + slot_owner: torch.Tensor | None = None, + slot_last_tok: torch.Tensor | None = None, + slot_lc: torch.Tensor | None = None, + slot_ct: torch.Tensor | None = None, + router_logits: torch.Tensor | None = None, + near_miss_thr: float = 0.0, + stats: torch.Tensor | None = None, + id_base: int = 0, + bump_tok: bool = False, + update_state: bool = True, + pin_since: torch.Tensor | None = None, + beta: float = BETA, + w: float = W, + halflife: int = HALFLIFE, +) -> None: + """``flashlib.lru_ensure`` (sequential strategy) with the victim picked by ``policy``. + + Same contract as ``lru_ensure``; ``policy`` > 0 also reads and updates the per-(layer, + expert) state ``last_tok``/``lc``/``ct`` (flat ``layer * num_experts + expert``) against + the device token counter ``tok``, and keeps its ``[num_cached]`` mirrors ``slot_last_tok``/ + ``slot_lc``/``slot_ct`` valid where ``slot_owner`` equals ``id_of_slot`` (-1 marks a stale + mirror). ``bump_tok`` advances the counter first (once per decode step); + ``update_state=False`` leaves the state untouched (small-prefill ensures). + ``pin_since`` (a ``lru_step`` value) also pins every slot touched after it, so a small + prefill's chunked calls cannot evict each other's experts; policy > 0 only. + Policy 3 refreshes the experts whose ``router_logits`` (``[rows, num_experts]``, rows of + ``query``) are within ``near_miss_thr`` of their row's lowest routed logit. + """ + k = query.numel() + num_cached = id_of_slot.numel() + assert query.dtype == torch.int32 and query.is_contiguous() + assert slot_of_id.dtype == torch.int32 and id_of_slot.dtype == torch.int32 + assert out_indices.dtype == torch.int32 and out_indices.numel() == k + assert lru_usage.dtype == lru_step.dtype and lru_usage.numel() == num_cached + plan = min(k, num_cached) + assert src_indices.numel() >= plan and dst_indices.numel() >= plan + block_c = triton.next_power_of_2(num_cached) + slot_bits = max(block_c.bit_length() - 1, 1) + beta_step, w_q4, decay_mul, dt_max = score_params(beta, w, halflife, num_layers) + if policy != 0: + assert slot_bits <= MAX_SLOT_BITS, f"{num_cached} slots overflow the packed score key" + assert tok is not None and last_tok is not None + assert last_tok.dtype == torch.int64 and last_tok.numel() == num_layers * num_experts + assert slot_owner is not None and slot_last_tok is not None + assert slot_owner.dtype == torch.int32 and slot_owner.numel() == num_cached + assert slot_last_tok.dtype == torch.int64 and slot_last_tok.numel() == num_cached + if pin_since is not None: + assert policy != 0 and pin_since.dtype == lru_step.dtype and pin_since.numel() == 1 + if policy >= 2: + assert lc is not None and ct is not None and g_table is not None + assert lc.dtype == torch.int32 and ct.dtype == torch.int64 + assert slot_lc is not None and slot_lc.dtype == torch.int32 and slot_lc.numel() == num_cached + assert slot_ct is not None and slot_ct.dtype == torch.int64 and slot_ct.numel() == num_cached + nm_rows = nm_topk = nm_stride = 0 + if policy == 3 and router_logits is not None: + assert router_logits.ndim == 2 and router_logits.shape[1] == num_experts + assert router_logits.stride(1) == 1 and k % router_logits.shape[0] == 0 + nm_rows, nm_stride = router_logits.shape[0], router_logits.stride(0) + nm_topk = k // nm_rows + # unused pointers still need a tensor argument + dummy = lru_step + _scored_ensure_kernel[(1,)]( + query, slot_of_id, id_of_slot, lru_usage, lru_step, + out_indices, src_indices, dst_indices, num_copy, stats, + dummy if tok is None else tok, + dummy if last_tok is None else last_tok, + dummy if lc is None else lc, + dummy if ct is None else ct, + dummy if g_table is None else g_table, + dummy if nm_rows == 0 else router_logits, + dummy if slot_owner is None else slot_owner, + dummy if slot_last_tok is None else slot_last_tok, + dummy if slot_lc is None else slot_lc, + dummy if slot_ct is None else slot_ct, + dummy if pin_since is None else pin_since, + k, num_cached, id_base, nm_topk, nm_stride, float(near_miss_thr), + BLOCK_K=triton.next_power_of_2(k), + BLOCK_C=block_c, + BLOCK_E=triton.next_power_of_2(num_experts), + BLOCK_TOPK=triton.next_power_of_2(max(nm_topk, 1)), + NM_ROWS=nm_rows, + USAGE_MAX=torch.iinfo(lru_usage.dtype).max, + COLLECT_STATS=stats is not None, + POLICY=policy, + BUMP_TOK=bool(bump_tok) and policy != 0, + UPDATE_STATE=bool(update_state) and policy != 0, + PIN_SINCE=pin_since is not None, + NUM_LAYERS=num_layers, + NUM_EXPERTS=num_experts, + BETA_STEP=beta_step, + W_Q4=w_q4, + DECAY_MUL=decay_mul, + DT_MAX=dt_max, + SLOT_BITS=slot_bits, + num_warps=_num_warps_for(block_c), + ) diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 99172e1ec8..00209604bc 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -785,8 +785,13 @@ def _infer_reasoning_parser(model_path: str) -> str | None: parser.add_argument( "--moe-cache-policy", default=ServerArgs.moe_cache_policy, - choices=["lru"], - help="The unified MoE cache eviction policy.", + choices=["lru", "kd", "kdfb", "rule"], + help=( + "The unified MoE cache eviction policy. lru (default); kd: evict the expert passed " + "over for the most tokens, then the one whose layer comes up last; kdfb: kd plus a " + "decayed per-expert use count; rule: kdfb plus router near misses counted as recent " + "(margin FREETOKEN_MOE_NEAR_MISS_THR logits). GPU decode only." + ), ) parser.add_argument( diff --git a/tests/README.md b/tests/README.md index 76b2b1fb03..47c7455c35 100644 --- a/tests/README.md +++ b/tests/README.md @@ -63,6 +63,7 @@ checkpoint is set: | `FREETOKEN_REBUILD_TEST_MODEL` | `e2e/test_cache_rebuild.py` — a SMALL local model dir; boots a real server (falls back to `FREETOKEN_TEST_MODEL`) | | `FREETOKEN_GEMMA4_GGUF_GLOB` | `models/test_gemma4_gguf_rope.py` — glob matching a local gemma-4 GGUF file | | `FREETOKEN_TEST_MODELS_ROOT` | `models/test_quant_config.py` — directory whose checkpoints the stored-tensor scheme scan checks (default `/mnt/nvme/models`) | +| `FREETOKEN_MOE_ROUTING_TRACE` | `moe/test_scored_ensure.py` — a compact decode routing trace directory (`experts.npy`, `logit_idx.npy`, `logit_val.npy`); its rows are replayed through the scored eviction kernel and the CPU reference | `test_aime.py` takes its sampling protocol from the checkpoint's own `generation_config.json` (pass@N at the recommended temperature, or a single greedy diff --git a/tests/models/qwen4_exp/test_skeleton.py b/tests/models/qwen4_exp/test_skeleton.py index 8fb9112df2..a2412f8346 100644 --- a/tests/models/qwen4_exp/test_skeleton.py +++ b/tests/models/qwen4_exp/test_skeleton.py @@ -432,6 +432,130 @@ def test_shared_gate_kernels_match_torch(num_tokens, hidden, dtype): assert (fused.float() - ref).abs().max() <= (eager.float() - ref).abs().max() + 1e-6 +def _offload_experts(kind: str, layer_id: int, num_experts: int, top_k: int, hidden: int, inter: int): + from freetoken.layers.moe import OffloadMoELayer + from freetoken.layers.quantization import NoQuantConfig, QuantBackend, QuantConfig, set_quant_backend + + if kind == "bf16": + quant = NoQuantConfig() + else: + set_quant_backend(QuantBackend.parse("moe.nvfp4=triton")) # the sm_120 decode kernel + quant = QuantConfig.from_hf({"quantization_config": {"quant_method": "modelopt", "quant_algo": "NVFP4", "ignore": ["lm_head"]}}) + return OffloadMoELayer( + layer_id, num_experts, top_k, hidden, inter, quant_config=quant, + prefix=f"model.layers.{layer_id}.mlp.experts", + ) + + +def _copy_overlap_run(kind, banks, overlap, bs, graph, inputs, monkeypatch): + """Two Qwen4ExpMoE layers sharing one offload cache as small as a layer, so every step evicts + and copies: warm up eagerly (and capture), reset the cache, then decode ``inputs``.""" + import freetoken.kernel.fast_index_copy as fic + from flashlib.kernels.slot_cache import Stat + from freetoken.models.qwen4_exp.moe import Qwen4ExpMoE + from freetoken.moe.offload_cache import OffloadMoeCache + from freetoken.utils.torch_utils import torch_dtype + + num_experts, top_k, hidden, inter = 8, 2, 256, 128 + config = parse_config(toy_hf_config( + hidden_size=hidden, num_experts=num_experts, num_experts_per_tok=top_k, + moe_intermediate_size=inter, shared_expert_intermediate_size=inter, + )) + device = torch.device("cuda") + with torch.device(device), torch_dtype(torch.bfloat16): + moes = [Qwen4ExpMoE(config, layer_id) for layer_id in range(2)] + gen = torch.Generator(device=device).manual_seed(41) + for moe in moes: + _fill(moe, gen, scale=0.2) + experts = [_offload_experts(kind, i, num_experts, top_k, hidden, inter) for i in range(2)] + cache = OffloadMoeCache( + num_layers=2, num_experts=num_experts, cache_size=num_experts, device=device, + quant_format=banks.quant_format, layout=banks.layout, + max_slots=experts[0].quant_method.slot_limit(), decode_copy_overlap=overlap, + ) + cache.set_bank_sources(banks.sources) + cache.set_alphas(banks.gate_up_alpha, banks.down_alpha) + cache.collect_stats = True + for moe, ex in zip(moes, experts): + ex.offload_cache = cache + moe.experts = ex + cache.reset() + + import freetoken.models.qwen4_exp.moe as qmoe + + copy_streams, gate_streams = [], [] + real_gate = qmoe.shared_gate_sigmoid + for name in ("fast_index_copy_multi_jit", "fast_index_copy_multi_slim_jit"): # FREETOKEN_MOE_SLIM_COPY + real_copy = getattr(fic, name) + monkeypatch.setattr( + fic, name, + lambda *a, _real=real_copy, **k: (copy_streams.append(torch.cuda.current_stream().cuda_stream), _real(*a, **k))[1], + ) + monkeypatch.setattr( + qmoe, "shared_gate_sigmoid", + lambda *a, **k: (gate_streams.append(torch.cuda.current_stream().cuda_stream), real_gate(*a, **k))[1], + ) + + def step(x0, x1): + return moes[0].forward(x0.clone()), moes[1].forward(x1.clone()) + + _fresh_ctx(_batch=SimpleNamespace(is_prefill=False)) + static = [torch.randn(bs, hidden, device=device, dtype=torch.bfloat16) for _ in range(2)] + step(*static) # the eager warm-up the graph runner also does + if graph: + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + outs = step(*static) + cache.reset() + cache.reset_stats() + results = [] + for x0, x1 in inputs: + if graph: + static[0].copy_(x0) + static[1].copy_(x1) + g.replay() + else: + outs = step(x0, x1) + results.append(tuple(o.clone() for o in outs)) + torch.cuda.synchronize() + state = [t.clone() for t in (cache.slot_for_id, cache.id_of_slot, cache.usage, *cache.bank_views())] + misses = int(cache.lru_stats[:, Stat.MISS].sum()) + side = getattr(cache.decode_copy_stream, "cuda_stream", None) + return results, state, misses, [s == side for s in copy_streams], [s == side for s in gate_streams] + + +@requires_cuda +@pytest.mark.parametrize("kind", ["bf16", "nvfp4"]) +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +def test_moe_copy_overlap_is_bitwise_identical(kind, graph, bs, monkeypatch): + """FREETOKEN_MOE_COPY_OVERLAP moves the shared expert onto a side stream beside the decode miss + copy; the MoE outputs, the slot map and the slot contents must not change by a single bit.""" + from freetoken.moe.expert_banks import build_expert_banks + + torch.manual_seed(7) + method = _offload_experts(kind, 0, 8, 2, 256, 128).quant_method + banks = build_expert_banks(method, 2, None, device=torch.device("cuda"), dummy=True) + gen = torch.Generator(device="cuda").manual_seed(bs) + inputs = [ + tuple(torch.randn(bs, 256, device="cuda", dtype=torch.bfloat16, generator=gen) * 0.5 for _ in range(2)) + for _ in range(6) + ] + off = _copy_overlap_run(kind, banks, False, bs, graph, inputs, monkeypatch) + on = _copy_overlap_run(kind, banks, True, bs, graph, inputs, monkeypatch) + + assert off[2] > 0 and on[2] == off[2], "the sequence must miss, and copy, in both runs" + assert off[3] and not any(off[3]) and not any(off[4]), "flag off: everything stays on the compute stream" + assert on[3] and not any(on[3]), "flag on: the miss copy stays right behind its ensure" + assert on[4] and all(on[4]), "flag on: the shared expert runs on the cache's side stream" + for got, want in zip(on[0], off[0]): + for a, b in zip(got, want): + assert torch.isfinite(b.float()).all() + assert torch.equal(a, b) + for a, b in zip(on[1], off[1]): + assert torch.equal(a, b) + + @requires_cuda def test_decoder_stack_prefill_and_decode(monkeypatch): """Ragged bs=3 prefill then a bs=3 decode step through the whole model with dummy weights.""" diff --git a/tests/moe/ref_scored_cache.py b/tests/moe/ref_scored_cache.py new file mode 100644 index 0000000000..d7d8b0f2fb --- /dev/null +++ b/tests/moe/ref_scored_cache.py @@ -0,0 +1,159 @@ +"""CPU reference for the scored slot-cache ensure (freetoken.moe.scored_ensure). + +A port of the adversarial-synthesis replay (global slot cache with lru_ensure semantics: the +call's ids pinned, empty slots first, ties to the lowest slot, missing ids ascending paired with +ascending victims) onto the kernel's Q16 integer arithmetic, so victims match the GPU exactly. +""" +from __future__ import annotations + +import numpy as np + +from freetoken.moe import scored_ensure as se + +_KEY_MAX = np.iinfo(np.int64).max + + +class RefScoredCache: + def __init__(self, num_layers, num_experts, cache_size, policy, + beta=se.BETA, w=se.W, halflife=se.HALFLIFE): + self.L, self.E, self.S, self.policy = num_layers, num_experts, cache_size, policy + self.beta_step, self.w_q4, self.decay_mul, self.dt_max = se.score_params(beta, w, halflife, num_layers) + block_c = 1 << max(cache_size - 1, 1).bit_length() + self.slot_bits = max(block_c.bit_length() - 1, 1) + self.g = se.softplus2_table().numpy().astype(np.int64) + self.reset() + + def reset(self): + n = self.L * self.E + self.slot_of_id = np.full(n, -1, np.int64) + self.id_of_slot = np.full(self.S, -1, np.int64) + self.usage = np.zeros(self.S, np.int64) + self.step = 0 + self.tok = 0 + self.last_tok = np.full(n, se.LAST_TOK_NEVER, np.int64) + self.lc = np.full(n, se.LC_NEVER, np.int64) + self.ct = np.zeros(n, np.int64) + + def _decayed(self, ids, now): + dt = np.minimum(np.maximum(now - self.ct[ids], 0), self.dt_max) + return self.lc[ids] - ((dt * self.decay_mul) >> se.DECAY_SHIFT) + + def _softplus2(self, x): + shift = se.Q - se.G_STEP_BITS + u = np.abs(x) + i = np.minimum(u >> shift, se.G_LAST) + frac = u & ((1 << shift) - 1) + g0, g1 = self.g[i], self.g[i + 1] + g = np.where(i < se.G_LAST, g0 + (((g1 - g0) * frac) >> shift), 0) + return np.maximum(x, 0) + g + + def _keys(self, layer, tok, now, pin_since=None): + c = np.arange(self.S, dtype=np.int64) + oid = self.id_of_slot + held = oid >= 0 + if self.policy == 0: + key = (self.usage << self.slot_bits) | c + else: + ids = np.where(held, oid, 0) + lk = ids // self.E + ahead = lk > layer + d = np.where(ahead, lk - layer, self.L - layer + lk) + k = np.where(ahead, tok - 1, tok) - self.last_tok[ids] + k = np.minimum(np.maximum(k, 0), se.K_MAX) + score = -((k << se.Q) + d * self.beta_step) + if self.policy >= 2: + lc = self.lc[ids] + lcv = self._decayed(ids, now) + lcv = np.where(lc == se.LC_NEVER, se.LC_FLOOR, np.maximum(lcv, se.LC_FLOOR)) + score = score + ((self.w_q4 * lcv) >> 4) + assert np.abs(score).max() < 2**31 # the kernel computes it in int32 + key = np.where(held, ((score + se.SCORE_BIAS) << self.slot_bits) | c, c) + evictable = self.usage != self.step + if pin_since is not None: + evictable &= self.usage <= pin_since + return np.where(evictable, key, _KEY_MAX) + + def ensure(self, layer, ids, *, bump_tok=False, update_state=True, logits=None, thr=0.0, near_miss=None, + pin_since=None): + """One ensure call; returns ``(out_slots, src_ids, dst_slots)`` like the kernel's plan. + + ``pin_since`` (a ``step`` value) also pins every slot touched after it. + + Policy 3 refreshes the ids whose ``logits`` (``[rows, E]``) are within ``thr`` of their + row's lowest routed logit, in fp32 as the kernel does, or the ids of a given ``near_miss`` mask. + """ + self.step += 1 + base = layer * self.E + tok = now = 0 + if self.policy != 0: + if bump_tok: + self.tok += 1 + tok = self.tok + now = tok * self.L + layer + if self.policy == 3 and update_state and logits is not None: + lg = np.asarray(logits, dtype=np.float32).reshape(-1, self.E) + routed = np.asarray(ids, dtype=np.int64).reshape(lg.shape[0], -1) + kth = lg[np.arange(lg.shape[0])[:, None], routed].min(axis=1) + near_miss = lg >= (kth - np.float32(thr))[:, None] + if self.policy == 3 and update_state and near_miss is not None: + near = base + np.flatnonzero(np.asarray(near_miss).reshape(-1, self.E).any(0)) + self.last_tok[near] = np.maximum(self.last_tok[near], tok - 1) + q = np.asarray(ids, dtype=np.int64).reshape(-1) + base + slots = self.slot_of_id[q] + self.usage[slots[slots >= 0]] = self.step + uniq = np.unique(q) + missing = uniq[self.slot_of_id[uniq] < 0] # ascending, as the kernel ranks them + src = missing - base + dst = np.empty(missing.size, np.int64) + if missing.size: + keys = self._keys(layer, tok, now, pin_since) + part = np.argpartition(keys, missing.size - 1)[: missing.size] + victims = part[np.argsort(keys[part])] + for i, (e, v) in enumerate(zip(missing, victims)): + old = self.id_of_slot[v] + if old >= 0: + self.slot_of_id[old] = -1 + self.id_of_slot[v] = e + self.slot_of_id[e] = v + self.usage[v] = self.step + dst[i] = v + out = self.slot_of_id[q] + if self.policy != 0 and update_state: + self.last_tok[uniq] = tok + if self.policy >= 2: + lc = self.lc[uniq] + x = self._decayed(uniq, now) + self.lc[uniq] = np.where(lc == se.LC_NEVER, 0, self._softplus2(x)) + self.ct[uniq] = now + return out, src, dst + + + def materialize(self, layer): + """The prefill materialize kernel: the whole layer into slots [0, E), position == expert id.""" + base = layer * self.E + same = (self.id_of_slot >= base) & (self.id_of_slot < base + self.E) + self.id_of_slot[same] = -1 + self.usage[same] = 0 + old = self.id_of_slot[: self.E].copy() + self.slot_of_id[old[old >= 0]] = -1 + self.step += 1 + self.id_of_slot[: self.E] = base + np.arange(self.E) + self.slot_of_id[base : base + self.E] = np.arange(self.E) + self.usage[: self.E] = self.step + + +def replay(cache: RefScoredCache, rows, logits=None, thr=0.0, near_miss=None): + """Decode ``rows`` ``[N, L, K]`` (one step per row, all layers); per-row miss counts. + + ``rows`` may be ``[N, L, B, K]`` for a batch of B rows per step, one K*B call per layer. + Policy 3 takes the matching ``[N, L, B, E]`` router ``logits`` (or a ``near_miss`` mask). + """ + rows = np.asarray(rows) + misses = np.zeros(rows.shape[0], np.int64) + for r in range(rows.shape[0]): + for layer in range(rows.shape[1]): + lg = None if logits is None else logits[r, layer] + nm = None if near_miss is None else near_miss[r, layer] + _, src, _ = cache.ensure(layer, rows[r, layer], bump_tok=layer == 0, logits=lg, thr=thr, near_miss=nm) + misses[r] += src.size + return misses diff --git a/tests/moe/test_fused_copy.py b/tests/moe/test_fused_copy.py index 7a0c086444..c5336ea117 100644 --- a/tests/moe/test_fused_copy.py +++ b/tests/moe/test_fused_copy.py @@ -1,6 +1,8 @@ """The fused multi-bank ``copy_missing`` path must move exactly the same bytes as the legacy per-bank ``fast_index_copy_jit`` loop, for every miss count (including the -zero-copy case), across banks of differing per-row sizes. +zero-copy case), across banks of differing per-row sizes. The opt-in slim kernel +(``FREETOKEN_MOE_SLIM_COPY``) must move exactly the same bytes as the fused one, also when +replayed from a CUDA graph with a device-side plan. """ from __future__ import annotations @@ -8,6 +10,7 @@ import pytest import torch +from freetoken.moe import offload_cache from freetoken.moe.offload_cache import _BANK_SCHEMAS, OffloadMoeCache CUDA = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") @@ -17,22 +20,30 @@ # kernel/aot_models.py must be listed in its TEST_FEATURE_SIZES so the per-bank kernels # stay prebuilt under FREETOKEN_DISABLE_JIT=1. FEATS = [8192, 512, 256, 4096, 512, 256] +# qwen4_exp (Qwen3.8-Flash-Next) triton-NVFP4 rows in _BANK_SCHEMAS["nvfp4"] order: gate_up +# packed/scale/global, down packed/scale/global -- 2,772,480 B per expert. +NVFP4_FEATS = [1638400, 204800, 2560, 819200, 102400, 5120] +# 16 * odd, several rows shorter than one grid pass, so segments end mid-round and the +# stride phase wraps across banks and rows. +ODD_FEATS = [16 * 1031, 16 * 7, 16 * 250, 16, 16 * 4099, 16 * 33] -def _build_cache(num_layers, num_experts, cache_size): +def _build_cache(num_layers, num_experts, cache_size, quant_format="mxfp4_triton", feats=FEATS, pinned=False): dev = torch.device("cuda") cache = OffloadMoeCache( num_layers=num_layers, num_experts=num_experts, cache_size=cache_size, - device=dev, cache_policy="lru", prefill_overlap=False, quant_format="mxfp4_triton", + device=dev, cache_policy="lru", prefill_overlap=False, quant_format=quant_format, ) - schema = _BANK_SCHEMAS["mxfp4_triton"] + schema = _BANK_SCHEMAS[quant_format] # Views into one flat tensor: only per-layer addressing matters here, not # independent allocations. - sources = { - name: list(torch.randint(0, 256, (num_layers * num_experts, feat), dtype=torch.uint8, device=dev) - .split(num_experts)) - for name, feat in zip(schema, FEATS) - } + sources = {} + for name, feat in zip(schema, feats): + if pinned: + flat = torch.randint(0, 256, (num_layers * num_experts, feat), dtype=torch.uint8).pin_memory() + else: + flat = torch.randint(0, 256, (num_layers * num_experts, feat), dtype=torch.uint8, device=dev) + sources[name] = list(flat.split(num_experts)) cache.set_bank_sources(sources) # also builds the fused-copy descriptor return cache @@ -74,3 +85,147 @@ def test_fused_copy_matches_per_bank(num_indices): for b, (r, (_, c)) in enumerate(zip(ref, cache.banks)): assert torch.equal(r, c), f"bank {b} (feat={FEATS[b]}) fused != per-bank at num_indices={num_indices}" + + +def _host_banks(feats, n_src, seed): + from freetoken.kernel.pinned import device_ptr + + gen = torch.Generator().manual_seed(seed) + src = [torch.randint(0, 256, (n_src, f), dtype=torch.uint8, generator=gen).pin_memory() for f in feats] + dev = torch.device("cuda") + src_ptrs = torch.tensor([device_ptr(t) for t in src], dtype=torch.int64, device=dev) + feat_bytes = torch.tensor(feats, dtype=torch.int64, device=dev) + return src, src_ptrs, feat_bytes + + +def _slot_banks(feats, n_slots): + dst = [torch.full((n_slots, f), 0xAB, dtype=torch.uint8, device="cuda") for f in feats] + return dst, torch.tensor([t.data_ptr() for t in dst], dtype=torch.int64, device="cuda") + + +def _expected(src, feats, n_slots, di, si, num): + # untouched slots keep the 0xAB fill; plan rows land at their dst slots + out = [] + for b, f in enumerate(feats): + e = torch.full((n_slots, f), 0xAB, dtype=torch.uint8) + e[di[:num].long().cpu()] = src[b][si[:num].long().cpu()] + out.append(e.cuda()) + return out + + +def _random_plan(gen, n_src, n_slots, plan, dtype): + di = torch.randperm(n_slots, generator=gen)[:plan].to(dtype).cuda() # distinct: no write races + si = torch.randint(0, n_src, (plan,), generator=gen).to(dtype).cuda() + return di, si + + +def _check_slim_matches_multi(feats, idx_dtype, **grid): + from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit, fast_index_copy_multi_slim_jit + + n_src, n_slots, plan = 24, 16, 16 + src, src_ptrs, feat_bytes = _host_banks(feats, n_src, seed=0) + gen = torch.Generator().manual_seed(1) + for num in [0, 1, 3, 10, plan, None]: # None: no num_indices, the whole index length + di, si = _random_plan(gen, n_src, n_slots, plan, idx_dtype) + count = None if num is None else torch.tensor([num], dtype=torch.int64, device="cuda") + ref, ref_ptrs = _slot_banks(feats, n_slots) + fast_index_copy_multi_jit(ref_ptrs, src_ptrs, feat_bytes, di, si, count) + out, out_ptrs = _slot_banks(feats, n_slots) + fast_index_copy_multi_slim_jit(out_ptrs, src_ptrs, feat_bytes, di, si, count, **grid) + torch.cuda.synchronize() + exp = _expected(src, feats, n_slots, di, si, plan if num is None else num) + for b, f in enumerate(feats): + assert torch.equal(ref[b], exp[b]), f"multi reference wrong: bank {b} (feat={f}) num={num}" + assert torch.equal(out[b], ref[b]), f"slim != multi: bank {b} (feat={f}) num={num} grid={grid}" + + +@CUDA +@pytest.mark.parametrize("idx_dtype", [torch.int32, torch.int64]) +@pytest.mark.parametrize("feats", [NVFP4_FEATS, ODD_FEATS], ids=["nvfp4", "odd"]) +def test_slim_copy_matches_multi(feats, idx_dtype): + _check_slim_matches_multi(feats, idx_dtype) + + +@CUDA +@pytest.mark.slow +@pytest.mark.parametrize("grid", [(3, 128, 2), (16, 256, 2), (5, 64, 8)]) +@pytest.mark.parametrize("feats", [NVFP4_FEATS, ODD_FEATS], ids=["nvfp4", "odd"]) +def test_slim_copy_matches_multi_other_grids(feats, grid): + blocks, threads, unroll = grid + _check_slim_matches_multi(feats, torch.int32, num_blocks=blocks, num_threads=threads, unroll=unroll) + + +@CUDA +def test_slim_copy_graph_replay_follows_device_plan(): + from freetoken.kernel.fast_index_copy import fast_index_copy_multi_slim_jit + + feats, n_src, n_slots, plan = NVFP4_FEATS, 24, 16, 16 + src, src_ptrs, feat_bytes = _host_banks(feats, n_src, seed=2) + dst, dst_ptrs = _slot_banks(feats, n_slots) + di = torch.zeros(plan, dtype=torch.int32, device="cuda") + si = torch.zeros(plan, dtype=torch.int32, device="cuda") + count = torch.zeros(1, dtype=torch.int64, device="cuda") + fast_index_copy_multi_slim_jit(dst_ptrs, src_ptrs, feat_bytes, di, si, count) # JIT outside capture + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + fast_index_copy_multi_slim_jit(dst_ptrs, src_ptrs, feat_bytes, di, si, count) + + gen = torch.Generator().manual_seed(3) + for num in [4, 0, plan, 1, 10]: + new_di, new_si = _random_plan(gen, n_src, n_slots, plan, torch.int32) + di.copy_(new_di) + si.copy_(new_si) + count.fill_(num) + for d in dst: + d.fill_(0xAB) + graph.replay() + torch.cuda.synchronize() + exp = _expected(src, feats, n_slots, di, si, num) + for b, f in enumerate(feats): + assert torch.equal(dst[b], exp[b]), f"graph replay: bank {b} (feat={f}) num={num}" + + +def test_slim_copy_flag_defaults_off(monkeypatch): + import importlib.util + + monkeypatch.delenv("FREETOKEN_MOE_SLIM_COPY", raising=False) + spec = importlib.util.find_spec("freetoken.moe.offload_cache") + fresh = importlib.util.module_from_spec(spec) + spec.loader.exec_module(fresh) # a private copy: the imported module and its classes stay untouched + assert fresh._SLIM_COPY is False + monkeypatch.setenv("FREETOKEN_MOE_SLIM_COPY", "1") + spec.loader.exec_module(fresh) + assert fresh._SLIM_COPY is True + + +@CUDA +@pytest.mark.parametrize("slim", [False, True], ids=["multi", "slim"]) +def test_copy_missing_slim_switch(monkeypatch, slim): + from freetoken.kernel import fast_index_copy as fic + + calls = [] + for name in ("fast_index_copy_multi_jit", "fast_index_copy_multi_slim_jit"): + real = getattr(fic, name) + monkeypatch.setattr(fic, name, lambda *a, _n=name, _r=real, **k: (calls.append(_n), _r(*a, **k))) + monkeypatch.setattr(offload_cache, "_SLIM_COPY", slim) + + num_layers, num_experts, cache_size, layer_id, num = 2, 8, 16, 1, 5 + cache = _build_cache(num_layers, num_experts, cache_size, "nvfp4", NVFP4_FEATS, pinned=True) + assert cache._copy_fused_ok + cache._pending_src_layer = layer_id + dev = torch.device("cuda") + cache.num_indices.fill_(num) + cache.evict_slots[:num] = torch.tensor([9, 2, 15, 0, 7], dtype=torch.int32, device=dev) + cache.src_indices[:num] = torch.tensor([3, 3, 0, 7, 5], dtype=torch.int32, device=dev) + for _, c in cache.banks: + c.fill_(0xAB) + cache.copy_missing() + torch.cuda.synchronize() + + assert calls == ["fast_index_copy_multi_slim_jit" if slim else "fast_index_copy_multi_jit"] + slots, rows = cache.evict_slots[:num].long().cpu(), cache.src_indices[:num].long().cpu() + for b, (per_layer, c) in enumerate(cache.banks): + exp = torch.full(c.shape, 0xAB, dtype=torch.uint8) + exp[slots] = per_layer[layer_id][rows] + assert torch.equal(c.cpu(), exp), f"bank {b} (feat={NVFP4_FEATS[b]})" diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index aaf3a4c097..19e25eafe3 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -387,7 +387,7 @@ def test_offload_moe_layer_decode_forward_uses_remapped_slot_ids(monkeypatch): lambda *, hidden_states, gating_output, topk, renormalize: (topk_weights, topk_ids), ) - def fake_ensure(layer_id, expert_ids): + def fake_ensure(layer_id, expert_ids, **kw): calls["ensure_layer_id"] = layer_id calls["ensure_expert_ids"] = expert_ids.clone() expert_ids.copy_(torch.tensor([[5, 0]], dtype=torch.int32)) @@ -427,6 +427,121 @@ def fake_fused_decode( assert calls["topk_ids"].tolist() == [[5, 0]] +def test_decode_copy_overlap_follows_the_flag_and_the_decode_target(monkeypatch): + from freetoken.env import ENV + from freetoken.moe.offload_cache import OffloadMoeCache + + def cache(device, **kw): + return OffloadMoeCache(num_layers=3, num_experts=4, cache_size=8, device=torch.device(device), **kw) + + monkeypatch.setattr(ENV.MOE_COPY_OVERLAP, "value", True) + assert cache("cpu").decode_copy_overlap is True # unset -> FREETOKEN_MOE_COPY_OVERLAP + assert cache("cpu").decode_copy_stream is None + monkeypatch.setattr(ENV.MOE_COPY_OVERLAP, "value", False) + assert cache("cpu").decode_copy_overlap is False + if not torch.cuda.is_available(): + return + assert cache("cuda").decode_copy_stream is None + for target in ("cpu", "hybrid"): + assert cache("cuda", decode_copy_overlap=True, decode_target=target).decode_copy_stream is None + on = cache("cuda", decode_copy_overlap=True) + assert on.decode_copy_stream is not None and len(on._decode_copy_events) == 3 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_copy_missing_takes_an_explicit_layer_id(): + # a second copy path (the overlap fork, later prefetch) names its source layer instead of trusting whichever call staged last + from freetoken.moe.offload_cache import OffloadMoeCache + + num_layers, num_experts = 2, 4 + fingerprint = torch.arange(num_layers * num_experts, dtype=torch.float32).view(num_layers, num_experts, 1, 1) + sources = { + "gate_up": [fingerprint[l].expand(num_experts, 32, 8).contiguous().pin_memory() for l in range(num_layers)], + "down": [fingerprint[l].expand(num_experts, 8, 16).contiguous().pin_memory() for l in range(num_layers)], + } + cache = OffloadMoeCache( + num_layers=num_layers, num_experts=num_experts, cache_size=num_experts, device=torch.device("cuda") + ) + cache.set_bank_sources(sources) + cache.reset() + ids = torch.tensor([2, 3], dtype=torch.int32, device="cuda") + cache.ensure_experts(1, ids) + cache._pending_src_layer = 0 + cache.copy_missing(1) + torch.cuda.synchronize() + for bank in cache.bank_views(): + assert [int(bank[s].flatten()[0]) for s in ids.tolist()] == [num_experts + 2, num_experts + 3] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_decode_side_applies_only_to_gpu_decode(monkeypatch): + # prefill, CPU-routed layers, TP > 1 and the flag off keep the shared expert on the compute stream + from types import SimpleNamespace + + import freetoken.core as core + from freetoken.core import Context + from freetoken.moe.offload_cache import OffloadMoeCache + + _init_tp() + layer = _bf16_offload_layer(0, 4, 2, 8, 16) + ctx = Context(page_size=1) + monkeypatch.setattr(core, "_GLOBAL_CTX", ctx) + + def applies(is_prefill=False, cpu_layers=(), **cache_kw): + layer.offload_cache = OffloadMoeCache( + num_layers=1, num_experts=4, cache_size=4, device=torch.device("cuda"), **cache_kw + ) + layer.offload_cache.cpu_layer_ids = frozenset(cpu_layers) + ctx._batch = SimpleNamespace(is_prefill=is_prefill) + return layer.decode_side_applies() + + assert not applies(decode_copy_overlap=False) + assert not applies(decode_copy_overlap=True, decode_target="hybrid") + assert not applies(is_prefill=True, decode_copy_overlap=True) + assert not applies(cpu_layers=(0,), decode_copy_overlap=True) + monkeypatch.setattr(layer, "tp_size", 2) + assert not applies(decode_copy_overlap=True) + monkeypatch.setattr(layer, "tp_size", 1) + assert applies(decode_copy_overlap=True) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_decode_side_stream_never_aliases_a_pooled_stream(): + # torch.cuda.Stream() hands out 32 pooled streams per priority round robin; a pooled side stream + # can come back as the engine or capture stream, and then the side work silently runs serially + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache( + num_layers=1, num_experts=4, cache_size=4, device=torch.device("cuda"), decode_copy_overlap=True + ) + side = cache.decode_copy_stream.cuda_stream + pooled = {torch.cuda.Stream().cuda_stream for _ in range(64)} + assert side not in pooled and side != torch.cuda.current_stream().cuda_stream + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_decode_side_survives_a_rebuild(): + # the side stream and events are not cache_size-shaped, so a runtime rebuild keeps them working + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache( + num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"), decode_copy_overlap=True + ) + cache.set_bank_sources({ + "gate_up": [torch.zeros(4, 32, 8, dtype=torch.bfloat16).pin_memory() for _ in range(2)], + "down": [torch.zeros(4, 8, 16, dtype=torch.bfloat16).pin_memory() for _ in range(2)], + }) + stream = cache.decode_copy_stream + cache.rebuild(6) + assert cache.decode_copy_stream is stream is not None + x = torch.arange(1024, dtype=torch.float32, device="cuda") + for layer_id in range(2): + with cache.decode_side(layer_id): + assert torch.cuda.current_stream() == stream + y = x * 2 + cache.join_decode_side(layer_id) + assert torch.equal(y + 0, x * 2) + def test_lru_gpu_cache_assigns_unique_slots_for_large_miss_batch(): import pytest diff --git a/tests/moe/test_scored_ensure.py b/tests/moe/test_scored_ensure.py new file mode 100644 index 0000000000..db734180b6 --- /dev/null +++ b/tests/moe/test_scored_ensure.py @@ -0,0 +1,622 @@ +"""--moe-cache-policy kd / kdfb / rule: the vendored scored-victim ensure kernel. + +The kernel must reproduce flashlib's LRU bit for bit under POLICY=lru, and match the CPU +reference (ref_scored_cache.py, a port of the eviction replay) victim for victim under the +scored policies. Set FREETOKEN_MOE_ROUTING_TRACE to a compact routing trace directory +(experts.npy, logit_idx.npy, logit_val.npy) to also replay real decode rows. +""" +import os +from pathlib import Path + +import numpy as np +import pytest +import torch + +L, E, S, TOP_K = 48, 512, 1650, 10 +NEAR_MISS_THR = 0.25 +_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +_TRACE = os.environ.get("FREETOKEN_MOE_ROUTING_TRACE", "") + + +def _trace(): + if not _TRACE or not (Path(_TRACE) / "experts.npy").exists(): + pytest.skip("FREETOKEN_MOE_ROUTING_TRACE is not set to a routing trace directory") + load = lambda name: np.load(Path(_TRACE) / name, mmap_mode="r") # noqa: E731 + return load("experts.npy"), load("logit_idx.npy"), load("logit_val.npy") + + +def _trace_logits(li, lv): + """Router logits rebuilt from the trace's recorded top-32 (the rest far below any margin).""" + li, lv = np.asarray(li).astype(np.int64), np.asarray(lv).astype(np.float32) + logits = np.full(li.shape[:-1] + (E,), -1e4, np.float32) + np.put_along_axis(logits, li, lv, axis=-1) + return logits + + +def _zipf_rows(n, rows_per_step=1, seed=0, num_layers=L, num_experts=E, top_k=TOP_K, skew=0.9): + """Skewed decode routing ``[n, L, rows_per_step, top_k]``: distinct ids within a row, duplicates across rows.""" + rng = np.random.default_rng(seed) + p = 1.0 / np.arange(1, num_experts + 1) ** skew + perm = [rng.permutation(num_experts) for _ in range(num_layers)] + out = np.empty((n, num_layers, rows_per_step, top_k), np.int64) + for r in range(n): + for layer in range(num_layers): + for b in range(rows_per_step): + pick = rng.choice(num_experts, top_k, replace=False, p=p / p.sum()) + out[r, layer, b] = perm[layer][pick] + return out + + +def _random_logits(rows, seed=1, num_experts=E): + """bf16-exact router logits shaped like ``rows`` with the routed ids on top (~10% near misses at 0.25).""" + rng = np.random.default_rng(seed) + logits = rng.normal(size=rows.shape[:-1] + (num_experts,)).astype(np.float32) + np.put_along_axis(logits, rows, np.take_along_axis(logits, rows, axis=-1) + 3.0, axis=-1) + return torch.from_numpy(logits).bfloat16().float().numpy() + + +def _cache(policy, cache_size=S, num_layers=L, num_experts=E, device="cuda", **kw): + from freetoken.moe.offload_cache import OffloadMoeCache + + return OffloadMoeCache( + num_layers=num_layers, num_experts=num_experts, cache_size=cache_size, + device=torch.device(device), cache_policy=policy, **kw, + ) + + +def _ref(policy, cache_size=S, num_layers=L, num_experts=E): + from freetoken.moe.scored_ensure import POLICY_IDS + + from .ref_scored_cache import RefScoredCache + + return RefScoredCache(num_layers, num_experts, cache_size, POLICY_IDS[policy]) + + +def _run_gpu(cache, rows, logits=None, dtype=torch.bfloat16): + """Decode ``rows`` ``[N, L, B, K]`` through ``cache.ensure_experts``; per-row miss counts.""" + dev = cache.device + ids = torch.from_numpy(np.ascontiguousarray(rows, dtype=np.int32)).to(dev) + lg = None if logits is None else torch.from_numpy(np.ascontiguousarray(logits)).to(dev, dtype) + cache.collect_stats = True + cache.lru_stats.zero_() + marks = torch.zeros(rows.shape[0], dtype=torch.int64, device=dev) + for r in range(rows.shape[0]): + for layer in range(rows.shape[1]): + cache.ensure_experts(layer, ids[r, layer], router_logits=None if lg is None else lg[r, layer]) + marks[r] = cache.lru_stats[:, 1].sum() + return np.diff(marks.cpu().numpy(), prepend=0) + + +def _run_ref(ref, rows, logits=None): + from .ref_scored_cache import replay + + return replay(ref, rows, logits, NEAR_MISS_THR) + + +def _assert_same_tables(cache, ref): + got = lambda t: t.detach().cpu().numpy().astype(np.int64) # noqa: E731 + np.testing.assert_array_equal(got(cache.slot_for_id.view(-1)), ref.slot_of_id) + np.testing.assert_array_equal(got(cache.id_of_slot), ref.id_of_slot) + np.testing.assert_array_equal(got(cache.usage), ref.usage) + assert int(cache.step) == ref.step + if ref.policy: + assert int(cache.evict_tok) == ref.tok + np.testing.assert_array_equal(got(cache.evict_last_tok), ref.last_tok) + if ref.policy >= 2: + np.testing.assert_array_equal(got(cache.evict_lc), ref.lc) + np.testing.assert_array_equal(got(cache.evict_ct), ref.ct) + if ref.policy: + _assert_mirrors_hold(cache) + + +def _assert_mirrors_hold(cache): + """Every slot whose owner tag matches its id mirrors that id's state.""" + got = lambda t: t.detach().cpu().numpy().astype(np.int64) # noqa: E731 + ids, owner = got(cache.id_of_slot), got(cache.evict_slot_owner) + tagged = (ids >= 0) & (owner == ids) + np.testing.assert_array_equal(got(cache.evict_slot_last_tok)[tagged], got(cache.evict_last_tok)[ids[tagged]]) + if cache.evict_lc is not None: + np.testing.assert_array_equal(got(cache.evict_slot_lc)[tagged], got(cache.evict_lc)[ids[tagged]]) + np.testing.assert_array_equal(got(cache.evict_slot_ct)[tagged], got(cache.evict_ct)[ids[tagged]]) + return tagged + + +def _check_against_ref(policy, rows, logits=None, dtype=torch.bfloat16): + cache, ref = _cache(policy), _ref(policy) + assert policy != "rule" or cache.near_miss_thr == NEAR_MISS_THR + got = _run_gpu(cache, rows, logits, dtype) + want = _run_ref(ref, rows, logits) + np.testing.assert_array_equal(got, want) + _assert_same_tables(cache, ref) + return got + + +# (a) POLICY=lru is flashlib's kernel ------------------------------------------------------- + + +@_cuda +@pytest.mark.parametrize("rows_per_step", [1, 2]) +def test_vendored_lru_is_bit_identical_to_flashlib(rows_per_step): + from flashlib.kernels.slot_cache import N_STATS, lru_ensure + + from freetoken.moe.scored_ensure import scored_ensure + + dev = torch.device("cuda") + k = rows_per_step * TOP_K + # skewed enough that some calls hit every id + stream = _zipf_rows(63, rows_per_step, seed=rows_per_step, skew=1.6).reshape(-1, k)[:3000] + layers = np.arange(stream.shape[0]) % L + + def state(): + return dict( + slot_of_id=torch.full((L * E,), -1, dtype=torch.int32, device=dev), + id_of_slot=torch.full((S,), -1, dtype=torch.int32, device=dev), + usage=torch.zeros(S, dtype=torch.int64, device=dev), + step=torch.zeros((), dtype=torch.int64, device=dev), + src=torch.zeros(S, dtype=torch.int32, device=dev), + dst=torch.zeros(S, dtype=torch.int32, device=dev), + num=torch.zeros((), dtype=torch.int64, device=dev), + stats=torch.zeros(N_STATS, dtype=torch.int64, device=dev), + ) + + a, b = state(), state() + queries = torch.from_numpy(stream.astype(np.int32)).to(dev) + qa, qb = queries.clone(), queries.clone() + plans = ([], []) + for i in range(stream.shape[0]): + base = int(layers[i]) * E + lru_ensure(qa[i], a["slot_of_id"], a["id_of_slot"], a["usage"], a["step"], qa[i], + a["src"], a["dst"], a["num"], stats=a["stats"], id_base=base) + scored_ensure(qb[i], b["slot_of_id"], b["id_of_slot"], b["usage"], b["step"], qb[i], + b["src"], b["dst"], b["num"], policy=0, num_layers=L, num_experts=E, + stats=b["stats"], id_base=base) + for st, plan in zip((a, b), plans): + plan.append(torch.cat([st["num"].view(1), st["src"][:k].long(), st["dst"][:k].long()])) + assert torch.equal(qa, qb), "out slots differ" + pa, pb = (torch.stack(plan) for plan in plans) + n = pa[:, 0] + assert int(n.sum()) > 1000 and int((n == 0).sum()) > 0 # the stream exercises misses and all-hit calls + for i in range(pa.shape[0]): + m = int(n[i]) + assert torch.equal(pa[i, : 1 + m], pb[i, : 1 + m]) and torch.equal(pa[i, 1 + k : 1 + k + m], pb[i, 1 + k : 1 + k + m]) + for key in ("slot_of_id", "id_of_slot", "usage", "step", "stats"): + assert torch.equal(a[key], b[key]), key + + +def test_lru_policy_stays_on_flashlib(monkeypatch): + import freetoken.moe.offload_kernels as ok + + cache = _cache("lru", device="cpu") + calls = [] + monkeypatch.setattr(ok, "lru_ensure", lambda *a, **kw: calls.append("flashlib")) + monkeypatch.setattr(ok, "ensure_experts_scored", lambda *a, **kw: calls.append("scored")) + cache.ensure_experts(3, torch.zeros(TOP_K, dtype=torch.int32)) + cache.ensure_experts(3, torch.zeros(TOP_K, dtype=torch.int32), update_state=False) + assert calls == ["flashlib", "flashlib"] + assert cache.evict_tok is None and cache.evict_last_tok is None and cache.near_miss_thr is None + + +# (b) scored policies == CPU reference ------------------------------------------------------ + + +@_cuda +@pytest.mark.parametrize("policy", ["lru", "kd", "kdfb", "rule"]) +def test_scored_kernel_matches_cpu_reference_on_synthetic_routing(policy): + rows = _zipf_rows(40, seed=3) + got = _check_against_ref(policy, rows, _random_logits(rows) if policy == "rule" else None) + assert got.sum() > 0 + + +@_cuda +@pytest.mark.slow +@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"]) +def test_scored_kernel_matches_cpu_reference_on_trace_rows(policy): + experts, li, lv = _trace() + rows = np.asarray(experts[:2000]).astype(np.int64)[:, :, None, :] + logits = _trace_logits(li[:2000], lv[:2000])[:, :, None, :] if policy == "rule" else None + got = _check_against_ref(policy, rows, logits, torch.float32) + lru = _run_ref(_ref("lru"), rows) + assert got.sum() < lru.sum() # the point of the policy + + +# (c) victims: never the call's own ids, empty slots first ---------------------------------- + + +@_cuda +@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"]) +@pytest.mark.parametrize("rows_per_step", [1, 2]) +def test_victims_skip_the_call_and_take_empty_slots_first(policy, rows_per_step): + from freetoken.moe import scored_ensure as se + + num_layers, num_experts, cache_size = 4, 64, 100 + cache = _cache(policy, cache_size, num_layers, num_experts) + ref = _ref(policy, cache_size, num_layers, num_experts) + warm = _zipf_rows(6, rows_per_step, seed=5, num_layers=num_layers, num_experts=num_experts) + _run_gpu(cache, warm) + _run_ref(ref, warm) + assert int((cache.id_of_slot >= 0).sum()) == cache_size + + layer = 2 + held = cache.id_of_slot.cpu().numpy() + resident = sorted(int(i) - layer * num_experts for i in held if layer * num_experts <= i < (layer + 1) * num_experts) + absent = sorted(set(range(num_experts)) - set(resident)) + k = rows_per_step * TOP_K + hits = resident[: k // 2] + ids = np.array((hits + absent)[:k], np.int64) + if rows_per_step == 2: + ids[-1] = ids[0] # a duplicate hit + # make this call's hits the stalest ids in the cache, so only the pin keeps them + hit_keys = [layer * num_experts + e for e in hits] + for key in hit_keys: + cache.evict_last_tok[key] = se.LAST_TOK_NEVER + ref.last_tok[key] = se.LAST_TOK_NEVER + if cache.evict_lc is not None: + cache.evict_lc[key] = se.LC_NEVER + ref.lc[key] = se.LC_NEVER + hit_slots = {int(cache.slot_for_id[layer, e]) for e in hits} + # invalidate three slots (as the prefill overlap does): they must be the first victims + empty = sorted([s for s in (7, 40, 91, 13, 55, 77) if s not in hit_slots][:3]) + for slot in empty: + for t, v in ((cache.id_of_slot, held[slot]), (ref.id_of_slot, held[slot])): + t[slot] = -1 + cache.slot_for_id.view(-1)[int(held[slot])] = -1 + ref.slot_of_id[int(held[slot])] = -1 + cache.usage[slot] = 0 + ref.usage[slot] = 0 + + q = torch.tensor(ids, dtype=torch.int32, device="cuda") + lg = _random_logits(ids.reshape(rows_per_step, TOP_K), seed=9, num_experts=num_experts) + cache.ensure_experts(layer, q, router_logits=torch.from_numpy(lg).cuda()) + out, src, dst = ref.ensure(layer, ids, logits=lg, thr=NEAR_MISS_THR) + + n = int(cache.num_indices) + got_dst = cache.evict_slots[:n].tolist() + assert n == len(set(ids.tolist()) - set(hits)) == src.size + assert got_dst[: len(empty)] == empty + assert not hit_slots & set(got_dst) + assert q.tolist() == out.tolist() + assert all(int(cache.slot_for_id[layer, e]) >= 0 for e in ids.tolist()) + assert got_dst == dst.tolist() + _assert_same_tables(cache, ref) + + +# (d) the decode-step clock ----------------------------------------------------------------- + + +@_cuda +@pytest.mark.parametrize("rows_per_step", [1, 2]) +def test_token_clock_moves_once_per_decode_step(rows_per_step): + cache = _cache("kdfb") + rows = _zipf_rows(5, rows_per_step, seed=7) + ids = torch.from_numpy(rows.astype(np.int32)).cuda() + for r in range(rows.shape[0]): + for layer in range(L): + cache.ensure_experts(layer, ids[r, layer]) + assert int(cache.evict_tok) == r + 1 + # a CPU-decoded first layer moves the clock to the first GPU layer + cache.reset() + cache.cpu_layer_ids = frozenset({0, 1}) + ids = torch.from_numpy(rows.astype(np.int32)).cuda() + for layer in range(2, L): + cache.ensure_experts(layer, ids[0, layer]) + assert int(cache.evict_tok) == 1 + + +@_cuda +def test_small_prefill_ensures_leave_the_state_alone(monkeypatch): + import freetoken.layers.moe as moe_mod + from freetoken.distributed import set_tp_info, try_get_tp_info + from freetoken.layers.moe import OffloadMoELayer + from freetoken.layers.quantization import NoQuantConfig + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + dev = torch.device("cuda") + experts, top_k, hidden, inter = 64, 4, 16, 32 + layer = OffloadMoELayer(0, experts, top_k, hidden, inter, quant_config=NoQuantConfig(), + prefix="model.layers.0.mlp.experts") + cache = _cache("kdfb", 96, 1, experts) + g = torch.Generator().manual_seed(0) + cache.set_bank_sources({ + "gate_up": [torch.randn(experts, 2 * inter, hidden, generator=g, dtype=torch.bfloat16) * 0.1], + "down": [torch.randn(experts, hidden, inter, generator=g, dtype=torch.bfloat16) * 0.1], + }) + layer.offload_cache = cache + decode = torch.tensor([[1, 2, 3, 4]], dtype=torch.int32, device=dev) + cache.ensure_experts(0, decode) + before = [t.clone() for t in (cache.evict_tok, cache.evict_last_tok, cache.evict_lc, cache.evict_ct)] + + monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", 64) + tokens = 40 + ids = torch.randint(0, experts // 2, (tokens, top_k), dtype=torch.int32, device=dev) + w = torch.full((tokens, top_k), 1.0 / top_k, device=dev) + layer._prefill_routed(torch.randn(tokens, hidden, dtype=torch.bfloat16, device=dev), w, ids.clone()) + assert bool((cache.slot_for_id[0, ids.long().unique()] >= 0).all()) # it did make them resident + assert int((cache.slot_for_id[0] >= 0).sum()) < experts # through the on-demand path, not the whole layer + after = (cache.evict_tok, cache.evict_last_tok, cache.evict_lc, cache.evict_ct) + for b, a in zip(before, after): + assert torch.equal(b, a) + + +@_cuda +def test_dsv4_on_demand_prefill_leaves_the_state_alone(): + from types import SimpleNamespace + + from freetoken.models.deepseek_v4.moe import DSV4OffloadMoELayer + + experts, top_k = 64, 4 + cache = _cache("kdfb", 96, 1, experts) + cache.set_bank_sources({"gate_up": [torch.zeros(experts, 2, 4)], "down": [torch.zeros(experts, 4, 2)]}) + cache.ensure_experts(0, torch.tensor([1, 2, 3, 4], dtype=torch.int32, device="cuda")) + before = [t.clone() for t in (cache.evict_tok, cache.evict_last_tok, cache.evict_lc, cache.evict_ct)] + # few enough routes (3 x 4 < 64) for the slot path rather than whole-layer streaming + layer = SimpleNamespace(offload_cache=cache, num_experts=experts, top_k=top_k, layer_id=0, + _expert_gemm=lambda *a, **kw: None) + ids = torch.tensor([[5, 6, 7, 8], [9, 10, 11, 12], [5, 13, 14, 15]], dtype=torch.int32, device="cuda") + DSV4OffloadMoELayer._prefill_routed(layer, torch.zeros(3, 8, device="cuda"), torch.full((3, top_k), 0.25), ids) + assert bool((cache.slot_for_id[0, 5:16] >= 0).all()) + for b, a in zip(before, (cache.evict_tok, cache.evict_last_tok, cache.evict_lc, cache.evict_ct)): + assert torch.equal(b, a) + + +# (e) reset() and rebuild() ----------------------------------------------------------------- + + +@_cuda +@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"]) +def test_reset_and_rebuild_clear_the_state(policy): + from freetoken.moe import scored_ensure as se + + cache = _cache(policy, 1480) + cache.set_bank_sources({"gate_up": [torch.zeros(E, 2, 4)] * L, "down": [torch.zeros(E, 4, 2)] * L}) + rows = _zipf_rows(3, seed=11) + fresh = _ref(policy) + + def assert_cold(): + assert int(cache.evict_tok) == 0 + assert bool((cache.evict_last_tok == se.LAST_TOK_NEVER).all()) + assert bool((cache.evict_slot_owner == -1).all()) and cache.evict_slot_owner.numel() == cache.cache_size + if cache.evict_lc is not None: + assert bool((cache.evict_lc == se.LC_NEVER).all()) and bool((cache.evict_ct == 0).all()) + + for clear in (cache.reset, lambda: cache.rebuild(S)): + _run_gpu(cache, rows) + assert int(cache.evict_tok) == rows.shape[0] + clear() + assert_cold() + assert cache.cache_size == S + # a rebuilt cache replays exactly like a fresh one + np.testing.assert_array_equal(_run_gpu(cache, rows), _run_ref(fresh, rows)) + _assert_same_tables(cache, fresh) + + +# (f) CUDA graph replay == eager ------------------------------------------------------------ + + +@_cuda +@pytest.mark.parametrize("policy", ["kdfb", "rule"]) +@pytest.mark.parametrize("rows_per_step", [1, 2]) +def test_graph_replay_matches_eager(policy, rows_per_step): + steps = 50 + rows = _zipf_rows(steps + 1, rows_per_step, seed=13) + ids = torch.from_numpy(rows.astype(np.int32)).cuda() + logits = torch.from_numpy(_random_logits(rows, seed=14)).cuda().bfloat16() + + eager = _cache(policy) + for r in range(1, steps + 1): + for layer in range(L): + eager.ensure_experts(layer, ids[r, layer].clone(), router_logits=logits[r, layer]) + + graphed = _cache(policy) + static_ids = ids[0].clone() + static_logits = logits[0].clone() + + def step(): + for layer in range(L): + graphed.ensure_experts(layer, static_ids[layer], router_logits=static_logits[layer]) + + step() # warm-up compiles the kernel variants outside capture + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + step() + graphed.reset() # as graph.py does after capture + for r in range(1, steps + 1): + static_ids.copy_(ids[r]) + static_logits.copy_(logits[r]) + graph.replay() + torch.cuda.synchronize() + + assert int(graphed.evict_tok) == steps + for name in ("slot_for_id", "id_of_slot", "usage", "step", "evict_tok", "evict_last_tok", "evict_lc", "evict_ct", + "evict_slot_owner", "evict_slot_last_tok", "evict_slot_lc", "evict_slot_ct"): + assert torch.equal(getattr(eager, name), getattr(graphed, name)), name + + +# (g) bs=2: two rows per call --------------------------------------------------------------- + + +@_cuda +@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"]) +def test_two_rows_per_call_match_cpu_reference(policy): + rows = _zipf_rows(30, 2, seed=17) + assert any(len(set(rows[r, layer].ravel())) < 2 * TOP_K for r in range(30) for layer in range(L)) + _check_against_ref(policy, rows, _random_logits(rows, seed=18) if policy == "rule" else None) + + +@_cuda +@pytest.mark.slow +@pytest.mark.parametrize("policy", ["kdfb", "rule"]) +def test_two_trace_streams_per_call_match_cpu_reference(policy): + experts, li, lv = _trace() + a, b = slice(0, 600), slice(45000, 45600) # two different requests decoding together + rows = np.stack([np.asarray(experts[a]), np.asarray(experts[b])], axis=2).astype(np.int64) + logits = None + if policy == "rule": + logits = np.stack([_trace_logits(li[a], lv[a]), _trace_logits(li[b], lv[b])], axis=2) + _check_against_ref(policy, rows, logits, torch.float32) + + +# the slot mirrors: foreign installs and state-free ensures fall back to the per-id state ----- + + +@_cuda +@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"]) +def test_materialized_prefill_between_decodes_matches_cpu_reference(policy): + cache, ref = _cache(policy), _ref(policy) + rows = _zipf_rows(24, seed=21) + logits = _random_logits(rows, seed=22) if policy == "rule" else None + # layer 5 comes back to slots it lost to layer 9: its tags there would otherwise look current + for part, layer in ((slice(0, 6), 5), (slice(6, 12), 9), (slice(12, 18), 5), (slice(18, 24), None)): + lg = None if logits is None else logits[part] + np.testing.assert_array_equal(_run_gpu(cache, rows[part], lg), _run_ref(ref, rows[part], lg)) + _assert_same_tables(cache, ref) + if layer is not None: + # the whole-layer prefill path re-seats this layer's experts at slots [0, E) + cache.materialize_layer(layer) + ref.materialize(layer) + assert not bool((cache.evict_slot_owner[:E] >= 0).any()) + _assert_same_tables(cache, ref) + + +@_cuda +@pytest.mark.parametrize("policy", ["kd", "kdfb"]) +def test_materialize_untags_the_slots_it_fills(policy): + # expert 3 of layer 0 leaves slot 3, is used elsewhere, then materialize puts it back at slot 3 + cache = _cache(policy, 24, 2, 8) + dev = torch.device("cuda") + ids = lambda *e: torch.tensor(e, dtype=torch.int32, device=dev) # noqa: E731 + cache.materialize_layer(0) + cache.ensure_experts(0, ids(3, 4)) # a hit at slot 3 tags it with expert 3 + assert int(cache.evict_slot_owner[3]) == 3 + cache.materialize_layer(1) # slot 3 now holds layer 1's expert 3 + cache.ensure_experts(1, ids(0, 1)) + cache.ensure_experts(0, ids(3, 5)) # expert 3 comes back elsewhere, with new state + assert int(cache.slot_for_id[0, 3]) != 3 + cache.materialize_layer(0) + assert int(cache.slot_for_id[0, 3]) == 3 + _assert_mirrors_hold(cache) + + +@_cuda +@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"]) +@pytest.mark.parametrize("chunks", [1, 2]) +def test_state_free_ensures_between_decodes_match_cpu_reference(policy, chunks): + cache, ref = _cache(policy), _ref(policy) + rows = _zipf_rows(12, seed=23) + extra = _zipf_rows(4, 3, seed=24, skew=0.3) + logits = _random_logits(rows, seed=25) if policy == "rule" else None + for r in range(rows.shape[0]): + lg = None if logits is None else logits[r : r + 1] + np.testing.assert_array_equal(_run_gpu(cache, rows[r : r + 1], lg), _run_ref(ref, rows[r : r + 1], lg)) + if r % 3 == 2: + for layer in range(0, L, 7): # a small prefill's chunked ensures + ids = np.unique(extra[r // 3, layer]) + # two chunks pin each other's slots as _ensure_unique does; one chunk keeps the plain pin + pin = cache.step.clone() if chunks > 1 else None + ref_pin = ref.step if chunks > 1 else None + for part in np.array_split(ids, chunks): + q = torch.from_numpy(part.astype(np.int32)).cuda() + cache.ensure_experts(layer, q, update_state=False, pin_since=pin) + out, _, dst = ref.ensure(layer, part, update_state=False, pin_since=ref_pin) + assert q.tolist() == out.tolist() + assert bool((cache.evict_slot_owner[torch.from_numpy(dst).cuda()] == -1).all()) + assert bool((cache.slot_for_id[layer, torch.from_numpy(ids).cuda()] >= 0).all()) + _assert_same_tables(cache, ref) + assert not _assert_mirrors_hold(cache).all() # some slots still read the per-id state + + +# rule: the near misses come from the router logits ---------------------------------------- + + +@_cuda +@pytest.mark.parametrize("rows_per_step", [1, 2]) +def test_rule_refreshes_the_router_near_misses(rows_per_step): + from freetoken.moe import scored_ensure as se + + cache = _cache("rule") + layer = 5 + g = torch.Generator(device="cuda").manual_seed(rows_per_step) + logits = (torch.randn(rows_per_step, E, device="cuda", generator=g) * 1.5).bfloat16() + ids = torch.topk(logits.float(), TOP_K, dim=-1).indices.to(torch.int32) + cache.ensure_experts(0, torch.arange(TOP_K, dtype=torch.int32, device="cuda")) # step 1 starts + cache.ensure_experts(layer, ids.clone(), router_logits=logits) + kth = logits.float().gather(1, ids.long()).min(dim=1, keepdim=True).values + near = (logits.float() >= kth - NEAR_MISS_THR).any(0) + routed = torch.zeros(E, dtype=torch.bool, device="cuda") + routed[ids.long().view(-1)] = True + last = cache.evict_last_tok[layer * E : (layer + 1) * E] + assert bool((last[routed] == 1).all()) + assert bool((last[near & ~routed] == 0).all()) + assert bool((last[~near & ~routed] == se.LAST_TOK_NEVER).all()) + assert 0 < int((near & ~routed).sum()) < E - TOP_K * rows_per_step + + +@pytest.mark.parametrize("policy", ["lru", "kdfb", "rule"]) +def test_decode_forward_hands_the_router_logits_to_ensure(monkeypatch, policy): + from freetoken.distributed import set_tp_info, try_get_tp_info + from freetoken.layers.moe import OffloadMoELayer + from freetoken.layers.quantization import NoQuantConfig + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + layer = OffloadMoELayer(0, 8, 2, 8, 16, quant_config=NoQuantConfig(), prefix="model.layers.0.mlp.experts") + layer.offload_cache = cache = _cache(policy, 12, 1, 8, device="cpu") + seen = {} + monkeypatch.setattr( + "freetoken.layers.moe.fused_topk", + lambda **kw: (torch.full((2, 2), 0.5), torch.zeros((2, 2), dtype=torch.int32)), + ) + monkeypatch.setattr(cache, "ensure_experts", lambda lid, ids, **kw: seen.update(kw)) + monkeypatch.setattr(cache, "copy_missing", lambda: None) + monkeypatch.setattr(cache, "bank_views", lambda: ()) + monkeypatch.setattr(layer, "_expert_gemm", lambda *a, **kw: None) + router_logits = torch.randn(2, 8) + layer.decode_forward(torch.randn(2, 8), router_logits) + assert seen["router_logits"] is router_logits + assert cache.near_miss_thr == (NEAR_MISS_THR if policy == "rule" else None) + + +def test_policy_flag_and_gpu_only(): + import contextlib + import io + + from freetoken.server.args import parse_args + + buf = io.StringIO() + with contextlib.redirect_stdout(buf), contextlib.suppress(SystemExit): + parse_args(["--help"]) + assert "{lru,kd,kdfb,rule}" in buf.getvalue() + for policy in ("kd", "kdfb", "rule"): + _cache(policy, 12, 1, 8, device="cpu") + with pytest.raises(ValueError, match="GPU decode"): + _cache(policy, 12, 1, 8, device="cpu", decode_target="hybrid") + + +@_cuda +@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"]) +def test_an_all_pinned_call_never_stores_past_the_slot_arrays(policy): + # 5 slots pad to an 8-lane block; with every slot pinned the sentinel key must not decode to lane 7. + # The cache enforces cache_size >= num_experts, so the kernel gets 5-slot views of a 16-slot cache. + from freetoken.moe.offload_kernels import ensure_experts_scored + + s, pad = 5, 3 + cache = _cache(policy, cache_size=16, num_layers=1, num_experts=16) + guarded = {} + for name in ("id_of_slot", "usage", "evict_slot_owner", "evict_slot_last_tok", "evict_slot_lc", "evict_slot_ct"): + t = getattr(cache, name) + if t is None: # kd keeps no lc/ct mirrors + continue + buf = torch.full((s + pad,), 7, dtype=t.dtype, device=t.device) + buf[:s] = t[:s] + setattr(cache, name, buf[:s]) + guarded[name] = buf + dev = cache.id_of_slot.device + ensure_experts_scored(cache, 0, torch.arange(s, dtype=torch.int32, device=dev), bump_tok=True, update_state=True) + # 6 distinct ids against 5 slots breaks the caller contract; the kernel must still stay in bounds + q = torch.arange(s + 1, dtype=torch.int32, device=dev) + ensure_experts_scored(cache, 0, q, bump_tok=True, update_state=True) + torch.cuda.synchronize() + for name, buf in guarded.items(): + assert torch.all(buf[s:] == 7), f"{name} written past slot {s - 1}" + assert int(cache.num_indices) == 1 and 0 <= int(cache.evict_slots[0]) < s diff --git a/tests/moe/test_small_prefill_on_demand.py b/tests/moe/test_small_prefill_on_demand.py index eea5842d23..70cec00019 100644 --- a/tests/moe/test_small_prefill_on_demand.py +++ b/tests/moe/test_small_prefill_on_demand.py @@ -8,7 +8,7 @@ from freetoken.distributed import set_tp_info, try_get_tp_info -def _layer_and_cache(device, num_experts=8, top_k=2, hidden=16, inter=32, cache_size=None): +def _layer_and_cache(device, num_experts=8, top_k=2, hidden=16, inter=32, cache_size=None, policy="lru", num_layers=1): from freetoken.layers.moe import OffloadMoELayer from freetoken.layers.quantization import NoQuantConfig from freetoken.moe.offload_cache import OffloadMoeCache @@ -17,11 +17,14 @@ def _layer_and_cache(device, num_experts=8, top_k=2, hidden=16, inter=32, cache_ set_tp_info(rank=0, size=1) layer = OffloadMoELayer(0, num_experts, top_k, hidden, inter, quant_config=NoQuantConfig(), prefix="model.layers.0.mlp.experts") - cache = OffloadMoeCache(num_layers=1, num_experts=num_experts, cache_size=cache_size or num_experts, device=device) + cache = OffloadMoeCache(num_layers=num_layers, num_experts=num_experts, cache_size=cache_size or num_experts, + device=device, cache_policy=policy) g = torch.Generator().manual_seed(0) cache.set_bank_sources({ - "gate_up": [torch.randn(num_experts, 2 * inter, hidden, generator=g, dtype=torch.bfloat16) * 0.1], - "down": [torch.randn(num_experts, hidden, inter, generator=g, dtype=torch.bfloat16) * 0.1], + "gate_up": [torch.randn(num_experts, 2 * inter, hidden, generator=g, dtype=torch.bfloat16) * 0.1 + for _ in range(num_layers)], + "down": [torch.randn(num_experts, hidden, inter, generator=g, dtype=torch.bfloat16) * 0.1 + for _ in range(num_layers)], }) layer.offload_cache = cache return layer, cache @@ -34,7 +37,7 @@ def test_small_prefill_dispatch(monkeypatch, tokens, threshold, on_demand): layer, cache = _layer_and_cache(torch.device("cpu")) monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", threshold) calls = [] - monkeypatch.setattr(cache, "ensure_experts", lambda layer_id, ids: calls.append("ensure")) + monkeypatch.setattr(cache, "ensure_experts", lambda layer_id, ids, **kw: calls.append("ensure")) monkeypatch.setattr(cache, "materialize_layer", lambda layer_id: calls.append("materialize")) monkeypatch.setattr(cache, "copy_missing", lambda: calls.append("copy")) monkeypatch.setattr(cache, "begin_prefill", lambda: calls.append("begin_prefill"), raising=False) @@ -96,7 +99,35 @@ def test_a_real_sized_small_prefill_fits_the_lru_kernel(monkeypatch): monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", threshold) layer, cache = _layer_and_cache(dev, num_experts=experts, top_k=top_k, cache_size=1100) widths, ensure = [], cache.ensure_experts - monkeypatch.setattr(cache, "ensure_experts", lambda lid, q: (widths.append(q.numel()), ensure(lid, q))[1]) + monkeypatch.setattr(cache, "ensure_experts", lambda lid, q, **kw: (widths.append(q.numel()), ensure(lid, q, **kw))[1]) outs[threshold] = _run_and_check_path(layer, hs, w, ids, on_demand=threshold > 0) assert widths and max(widths) <= 32, f"LRU ensure widths {sorted(set(widths))}: wider than decode's" torch.testing.assert_close(outs[1024], outs[0], rtol=2e-2, atol=2e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +@pytest.mark.parametrize("policy", ["lru", "kd", "kdfb", "rule"]) +def test_a_later_chunk_keeps_the_experts_an_earlier_chunk_loaded(monkeypatch, policy): + # Decode fills every slot, then 40 never-decoded experts take two ensure chunks. A scored + # policy ranks the first chunk's fresh installs as its coldest slots, so they need a pin. + dev = torch.device("cuda") + experts, top_k, slots, tokens = 64, 4, 80, 10 + g = torch.Generator(device=dev).manual_seed(2) + hs = torch.randn(tokens, 16, dtype=torch.bfloat16, device=dev, generator=g) * 4 # a wrong expert is ~1.0 off + ids = torch.arange(16, 16 + tokens * top_k, dtype=torch.int32, device=dev).view(tokens, top_k) + w = torch.softmax(torch.randn(tokens, top_k, device=dev, generator=g), dim=-1) + outs = {} + for threshold in (0, 64): + monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", threshold) + layer, cache = _layer_and_cache(dev, num_experts=experts, top_k=top_k, cache_size=slots, policy=policy, + num_layers=2) + for layer_id, first, last in ((1, 0, experts), (0, 0, 16)): + for start in range(first, last, 8): + cache.ensure_experts(layer_id, torch.arange(start, start + 8, dtype=torch.int32, device=dev)) + cache.copy_missing() + assert bool((cache.id_of_slot >= 0).all()) + outs[threshold] = _run_and_check_path(layer, hs, w, ids, on_demand=threshold > 0) + if threshold: + slot_of = cache.slot_for_id[0, ids.long()] + assert bool((slot_of >= 0).all()), f"{int((slot_of < 0).sum())} routed experts were evicted" + torch.testing.assert_close(outs[64], outs[0], rtol=2e-2, atol=0.1) From 4b4f9f8407a7bc8900372261cc13fa80c3ed299a Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Sun, 27 Sep 2026 03:33:34 -0400 Subject: [PATCH 33/73] merge fork perf: expert prefetch into low-priority slots with single-writer slot maps (FREETOKEN_MOE_PREFETCH), plus decode verify and slot audit debug modes --- benchmarks/README.md | 23 + benchmarks/bench_moe_prefetch.py | 580 ++++++ python/freetoken/engine/engine.py | 71 +- python/freetoken/env.py | 14 + python/freetoken/layers/moe.py | 48 +- python/freetoken/models/qwen4_exp/model.py | 3 +- python/freetoken/models/qwen4_exp/moe.py | 17 +- python/freetoken/moe/offload_cache.py | 236 ++- python/freetoken/moe/offload_kernels.py | 44 +- python/freetoken/moe/prefetch.py | 438 +++++ python/freetoken/moe/scored_ensure.py | 168 +- python/freetoken/moe/slot_audit.py | 688 +++++++ python/freetoken/moe/verify.py | 516 ++++++ tests/models/qwen4_exp/conftest.py | 13 + tests/models/qwen4_exp/test_build_cpu.py | 29 + tests/models/qwen4_exp/test_skeleton.py | 1911 ++++++++++++++++++++ tests/moe/ref_prefetch.py | 22 + tests/moe/ref_scored_cache.py | 44 +- tests/moe/test_prefetch.py | 424 +++++ tests/moe/test_scored_ensure.py | 319 ++++ 20 files changed, 5541 insertions(+), 67 deletions(-) create mode 100644 benchmarks/bench_moe_prefetch.py create mode 100644 python/freetoken/moe/prefetch.py create mode 100644 python/freetoken/moe/slot_audit.py create mode 100644 python/freetoken/moe/verify.py create mode 100644 tests/moe/ref_prefetch.py create mode 100644 tests/moe/test_prefetch.py diff --git a/benchmarks/README.md b/benchmarks/README.md index 3f0df72f81..817f99da18 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -45,6 +45,29 @@ start after the ensure and how much of the side-stream shared expert runs inside python benchmarks/bench_moe_copy_overlap.py --bs 1 --slots 160 ``` +**`bench_moe_prefetch.py`** — synthetic (no checkpoint): `FREETOKEN_MOE_PREFETCH=measure` +cost. The router-lookahead predictor alone per layer (GEMV `[bs, 2560] x [2560, 512]`, top-K +select, and the count kernel), then prefetch off vs measure on a real-geometry Qwen4ExpMoE decode +stack (512 experts, rule eviction, copy overlap on) with forced routing that misses exactly m +experts per layer, as CUDA graphs. `--on` times whole 48-layer decode steps with prefetch off vs +on instead: a GEMV per layer stands in for attention (`--filler-us`), each layer prefetches 3 or +4 candidates of which `--useful` are its misses, and `--delay-us` holds every prefetch copy back +to price a late one. It reports ms/step, demand misses/step and the late fraction. `--steady +--trace DIR` runs steady-state decode on a recorded routing trace instead (experts.npy, +logit_idx.npy, logit_val.npy, req_index.npy): no cache restore, so wrong prefetches and evictions +carry over, and a noisy copy of each layer's traced logits (`--sigma`) stands in for the +lookahead. Arms such as `off,measure,on,on:4/4` (budgets before GDN / QSA layers) replay step by +step in turn, each on its own copy of the cache state; it adds pollution (non-resident demands +per layer against off's misses) and the late fraction per layer kind. Slots and sigma were +calibrated to production's measure line (misses/layer ~4.3, precision ~0.56) at 24 layers. + +```bash +python benchmarks/bench_moe_prefetch.py --bs 1,2 --misses 0,4 +python benchmarks/bench_moe_prefetch.py --on --bs 1,2 --misses 4 --delay-us 0,50,150,300 +python benchmarks/bench_moe_prefetch.py --steady --trace DIR --bs 1,2 --steady-slots 740 --sigma 0.55 \ + --arms off,measure,on,on:empty,on:3/3,on:4/4,on:4/5 +``` + **`bench_expert_copy_corun.py`** — synthetic (no checkpoint): the fused expert copy kernel alone and beside compute, on registered host banks with the qwen4_exp NVFP4 row sizes. Sweeps `fast_index_copy_multi` (threads x blocks/bank) and `fast_index_copy_multi_slim` diff --git a/benchmarks/bench_moe_prefetch.py b/benchmarks/bench_moe_prefetch.py new file mode 100644 index 0000000000..be5d38ebd7 --- /dev/null +++ b/benchmarks/bench_moe_prefetch.py @@ -0,0 +1,580 @@ +"""FREETOKEN_MOE_PREFETCH cost and gain: the router-lookahead predictor alone, its effect on the +decode compute stream in measure mode, and (--on) whole decode steps with prefetch off vs on. + +1. Predictor alone, per layer: the lookahead GEMV ([bs, 2560] x [2560, 512] bf16, a distinct router + weight per layer as in the model), the top-K + select kernel, and both; one stream, CUDA graph. +2. The count kernel alone (it runs on the compute stream in measure mode). +3. A real-geometry decode stack (Qwen3.8-Flash-Next: hidden 2560, 512 experts, top-10, moe and + shared intermediate 640, NVFP4 triton banks) of N Qwen4ExpMoE blocks: router GEMV, shared + expert (on the FREETOKEN_MOE_COPY_OVERLAP side stream by default), top-k, rule ensure, miss + copy of exactly m experts from pinned host, grouped GEMM. The routing is forced (10 - m hot + experts + m cold ones per row) and the cache restored before every replay, so both arms copy + the same bytes. One CUDA graph per step, prefetch off vs measure. No attention runs, so the + percentages overstate the model's. +4. --on: the same stack at the model's depth (48 layers), with a bf16 GEMV "filler" before every + block standing in for attention and the rest of the layer (--filler-us, ~225 us between one + demand copy and the next in production). Prefetch on installs a fixed candidate list per layer: + 3 before GDN / 4 before attention layers, of which --useful are the layer's cold experts and the + rest never-routed ids (production: ~3.24 issued, ~1.75 useful, ~4.2 misses per layer at bs 1). + Reports ms/step, demand misses/step and the late fraction; --delay-us adds a sleep ahead of every + prefetch copy to price the compute stream's wait on a late copy. +5. --steady --trace DIR: steady-state decode on a recorded routing trace (no restore, so wrong + prefetches and evictions carry over), --steady-layers layers of it with --steady-slots slots, + a GEMV filler of --filler-gdn-us / --filler-qsa-us before GDN / QSA layers, and a lookahead + that is each layer's traced top-32 logits plus N(0, --sigma) noise. Every arm (off, measure, + on[:gdn/qsa budgets | :empty]) keeps its own cache state and all arms replay step by step in + turn, so drift hits them alike. Reports ms/step (mean, p50, p90), demand misses per layer, + issued/useful/precision, non-resident demands per layer (misses + prefetch hits: above off's + misses is pollution) and the late fraction before GDN and QSA layers. + +Run: CUDA_VISIBLE_DEVICES=0 PYTHONPATH=python python benchmarks/bench_moe_prefetch.py + ... benchmarks/bench_moe_prefetch.py --on --bs 1,2 --misses 4 + ... benchmarks/bench_moe_prefetch.py --steady --trace DIR --steady-slots 740 --sigma 0.55 +""" + +from __future__ import annotations + +import argparse +import gc +import statistics + +import torch +import torch.nn.functional as F + +H, E, TOP_K, INTER = 2560, 512, 10, 640 + + +def time_graph(run, reps: int, before=None) -> float: + """Median ms of one replay of ``run`` captured as a CUDA graph (after an eager warm-up).""" + if before: + before() + run() + torch.cuda.synchronize() + if before: + before() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + times = [] + for rep in range(reps + 3): + if before: + before() + torch.cuda.synchronize() + start.record() + graph.replay() + end.record() + end.synchronize() + if rep >= 3: + times.append(start.elapsed_time(end)) + return statistics.median(times) + + +def bench_predictor(bs: int, layers: int, k: int, budget: int, reps: int) -> dict: + from freetoken.moe.prefetch import lookahead_select, prefetch_count + + dev = torch.device("cuda") + gen = torch.Generator(device=dev).manual_seed(bs) + weights = [torch.randn(E, H, device=dev, dtype=torch.bfloat16, generator=gen) * 0.02 for _ in range(layers)] + x = torch.randn(bs, H, device=dev, dtype=torch.bfloat16, generator=gen) + resident = torch.where(torch.rand(E, device=dev, generator=gen) < 0.4, 1, -1).to(torch.int32) + sel = torch.full((layers, 2 * k), -1, dtype=torch.int32, device=dev) + res = sel.clone() + logits = [F.linear(x, w) for w in weights] + + def gemv(): + for w in weights: + F.linear(x, w) + + def select(): + for layer in range(layers): + lookahead_select(logits[layer], resident, sel[layer], res[layer], k=k, budget=budget) + + def both(): + for layer, w in enumerate(weights): + lookahead_select(F.linear(x, w), resident, sel[layer], res[layer], k=k, budget=budget) + + slots = torch.randint(0, 1024, (bs * TOP_K,), device=dev, dtype=torch.int32, generator=gen) + id_of_slot = torch.randint(0, E, (1024,), device=dev, dtype=torch.int32, generator=gen) + misses = torch.ones(1, dtype=torch.int64, device=dev) + counters = torch.zeros((layers, 6), dtype=torch.int64, device=dev) + + def count(): + for layer in range(layers): + prefetch_count(slots, id_of_slot, misses, sel[layer], res[layer], counters[layer], id_base=0, rows=bs) + + per = 1000.0 / layers + return {name: time_graph(fn, reps) * per for name, fn in + (("gemv", gemv), ("select", select), ("predictor", both), ("count", count))} + + +class ForcedGate: + """A layer's real router GEMV whose logits are then replaced, so every replay routes (and + misses) the same experts; the next layer's lookahead runs through it too.""" + + def __init__(self, gate, logits: torch.Tensor): + self.gate, self.logits = gate, logits + + def forward(self, x: torch.Tensor) -> torch.Tensor: + self.gate.forward(x) + return self.logits + + +class DecodeStack: + """N Qwen4ExpMoE blocks (router, shared expert, nvfp4 offload experts) over one shared cache, + wired like qwen4_exp for the lookahead.""" + + def __init__(self, layers: int, slots: int, bs: int, misses: int, mode: str, overlap: bool, filler=None, + routing=None, budgets=(3, 4)): + import freetoken.core as core + from types import SimpleNamespace + + from freetoken.core import Context, set_global_ctx + from freetoken.distributed import set_tp_info, try_get_tp_info + from freetoken.layers.moe import OffloadMoELayer + from freetoken.layers.quantization import QuantBackend, QuantConfig, set_quant_backend + from freetoken.models.qwen4_exp.moe import Qwen4ExpMoE + from freetoken.moe.expert_banks import build_expert_banks + from freetoken.moe.offload_cache import OffloadMoeCache + from freetoken.utils.torch_utils import torch_dtype + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + set_quant_backend(QuantBackend.parse("moe.nvfp4=triton")) + quant = QuantConfig.from_hf({"quantization_config": {"quant_method": "modelopt", "quant_algo": "NVFP4", "ignore": ["lm_head"]}}) + dev = torch.device("cuda") + config = SimpleNamespace( + hidden_size=H, num_experts=E, num_experts_per_tok=TOP_K, moe_intermediate_size=INTER, + shared_expert_intermediate_size=INTER, norm_topk_prob=True, quant=None, + moe_strategy="offload", decode_target="gpu", + ) + with torch.device(dev), torch_dtype(torch.bfloat16): + self.moes = [Qwen4ExpMoE(config, i, prefix=f"model.layers.{i}.mlp") for i in range(layers)] + gen = torch.Generator(device=dev).manual_seed(0) + for moe in self.moes: + for t in moe.state_dict().values(): + if t.is_floating_point(): + t.normal_(0.0, 0.02, generator=gen) + experts = [ + OffloadMoELayer(i, E, TOP_K, H, INTER, quant_config=quant, prefix=f"model.layers.{i}.mlp.experts") + for i in range(layers) + ] + method = experts[0].quant_method + if not hasattr(DecodeStack, "_banks"): + DecodeStack._banks = build_expert_banks(method, 1, None, device=dev, dummy=True) # one layer, shared + banks = DecodeStack._banks + self.cache = cache = OffloadMoeCache( + num_layers=layers, num_experts=E, cache_size=slots, device=dev, cache_policy="rule", + quant_format=banks.quant_format, layout=banks.layout, max_slots=method.slot_limit(), + decode_copy_overlap=overlap, prefetch_mode=mode, + ) + cache.set_bank_sources({role: per_layer * layers for role, per_layer in banks.sources.items()}) + cache.set_alphas(banks.gate_up_alpha, banks.down_alpha) + core._GLOBAL_CTX = None + ctx = Context(page_size=1) + set_global_ctx(ctx) + ctx._batch = SimpleNamespace(is_prefill=False) + self.x = [torch.randn(bs, H, device=dev, dtype=torch.bfloat16, generator=gen) * 0.5 for _ in range(layers)] + self.filler = filler + if routing is not None: + # --steady: every layer routes this step's row of route_buf and layer L's lookahead (from + # layer L-1) returns look_buf[L]; both run the real router GEMV first + route_buf, look_buf = routing + gates = [moe.gate for moe in self.moes] + for i, (moe, ex) in enumerate(zip(self.moes, experts)): + ex.offload_cache = cache + moe.experts = ex + moe.gate = ForcedGate(gates[i], route_buf[i]) + for i in range(1, layers): + experts[i - 1].set_lookahead(ForcedGate(gates[i], look_buf[i]), i, budgets[(i + 1) % 4 == 0]) + self.bs = bs + cache.reset() + torch.cuda.synchronize() + return + # forced routing: row r routes hot experts [r*hot, (r+1)*hot) and m fresh ones past all rows' hot ones + hot = TOP_K - misses + for moe, ex in zip(self.moes, experts): + ex.offload_cache = cache + moe.experts = ex + lg = torch.zeros(bs, E, device=dev, dtype=torch.bfloat16) + for r in range(bs): + lg[r, list(range(r * hot, (r + 1) * hot))] = 10.0 + lg[r, list(range(bs * hot + r * misses, bs * hot + (r + 1) * misses))] = 10.0 + moe.gate = ForcedGate(moe.gate, lg) + for i in range(1, layers): + experts[i - 1].set_lookahead(self.moes[i].gate, i, budgets[(i + 1) % 4 == 0]) + self.bs, self.hot, self.misses = bs, hot, misses + # warm the hot experts in, then drop the fresh ones: every replay then misses exactly those + cache.reset() + self.step() + cache.slot_for_id[:, bs * hot:].fill_(-1) + fresh = (cache.id_of_slot >= 0) & (cache.id_of_slot % E >= bs * hot) + cache.id_of_slot[fresh] = -1 + cache.usage[fresh] = 0 + torch.cuda.synchronize() + self.state = [cache.slot_for_id, cache.id_of_slot, cache.usage, cache.step] + for name in ("evict_tok", "evict_last_tok", "evict_lc", "evict_ct", "evict_slot_owner", + "evict_slot_last_tok", "evict_slot_lc", "evict_slot_ct"): + if getattr(cache, name) is not None: + self.state.append(getattr(cache, name)) + self.snap = [t.clone() for t in self.state] + + def restore(self): + for dst, src in zip(self.state, self.snap): + dst.copy_(src) + + def step(self): + for i, (moe, x) in enumerate(zip(self.moes, self.x)): + filler = self.filler[i] if isinstance(self.filler, list) else self.filler + if filler is not None: + F.linear(x, filler) + moe.forward(x.clone()) + + def cold(self, layer: int) -> list[int]: + """The experts ``layer`` misses every replay, merged rank by rank across rows like the select.""" + per_row = [list(range(self.bs * self.hot + r * self.misses, self.bs * self.hot + (r + 1) * self.misses)) + for r in range(self.bs)] + return [row[i] for i in range(self.misses) for row in per_row] + + +def bench_stack(layers, slots, bs, misses, reps, overlap) -> dict: + from flashlib.kernels.slot_cache import Stat + + out = {} + for mode in ("off", "measure"): + stack = DecodeStack(layers, slots, bs, misses, mode, overlap) + stack.cache.collect_stats = True + stack.restore() + stack.cache.lru_stats.zero_() + stack.step() + got = stack.cache.lru_stats[:, Stat.MISS].tolist() + assert got == [bs * misses] * layers, f"expected {bs * misses} misses per layer, got {got}" + stack.cache.collect_stats = False + out[mode] = time_graph(stack.step, reps, before=stack.restore) * 1000.0 / layers + if stack.cache.prefetch is not None: + stack.cache.prefetch.counters.zero_() + stack.restore() + stack.step() + c = stack.cache.prefetch.counters.sum(0).tolist() + out["issued/layer"] = c[0] / max(c[3], 1) + del stack + torch.cuda.empty_cache() + return out + + +def budget_of(layer: int) -> int: + return 3 if (layer + 1) % 4 else 4 + + +def bench_on(layers, slots, bs, misses, useful, reps, filler_us, delays_us) -> list[dict]: + """ms per decode step, prefetch off vs on (and on with every prefetch copy delayed).""" + from flashlib.kernels.slot_cache import Stat + + from freetoken.moe.prefetch import COPIED, ISSUED, LATE, USEFUL + + dev = torch.device("cuda") + x = torch.randn(1, H, device=dev, dtype=torch.bfloat16) + probe = torch.randn(32768, H, device=dev, dtype=torch.bfloat16) + per_row_us = time_graph(lambda: F.linear(x, probe), reps) * 1000.0 / 32768 + del probe + filler = torch.randn(max(int(filler_us / per_row_us), 1), H, device=dev, dtype=torch.bfloat16) * 0.01 + cycles_per_us = 1_000_000 / (time_graph(lambda: torch.cuda._sleep(1_000_000), 5) * 1000.0) + out = [] + for arm in ["off", "on", *[f"on+{d}us" for d in delays_us if d]]: + stack = DecodeStack(layers, slots, bs, misses, "off" if arm == "off" else "on", True, filler=filler) + cache = stack.cache + want = [bs * misses] * layers + if arm != "off": + pf = cache.prefetch + pf.sel_override = torch.full_like(pf.sel, -1) + for layer in range(1, layers): + u = min(useful[layer % len(useful)], bs * misses) + wrong = list(range(E - 1, E - 1 - (budget_of(layer) - u), -1)) + ids = stack.cold(layer)[:u] + wrong + pf.sel_override[layer, : len(ids)] = torch.tensor(ids, dtype=torch.int32) + want[layer] -= u + if arm.startswith("on+"): + pf.delay_copy_cycles = int(float(arm[3:-2]) * cycles_per_us) + cache.collect_stats = True + stack.restore() + cache.lru_stats.zero_() + stack.step() + got = cache.lru_stats[:, Stat.MISS].tolist() + assert got == want, f"{arm}: expected demand misses {want}, got {got}" + cache.collect_stats = False + stack.restore() + stack.step() + torch.cuda.synchronize() + stack.restore() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + stack.step() + if cache.prefetch is not None: + cache.prefetch.stats.zero_() + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + times = [] + for rep in range(reps + 3): + stack.restore() + torch.cuda.synchronize() + start.record() + graph.replay() + end.record() + end.synchronize() + if rep >= 3: + times.append(start.elapsed_time(end)) + row = {"arm": arm, "ms": statistics.median(times), "misses": sum(want)} + if cache.prefetch is not None: + st = cache.prefetch.stats.sum(0).tolist() + calls = max(reps + 3, 1) * (layers - 1) + row.update(issued=st[ISSUED] / calls, useful=st[USEFUL] / calls, late=st[LATE] / max(st[COPIED], 1)) + out.append(row) + del graph, stack, cache + gc.collect() # the module graph holds cycles: free this arm's slot banks before the next one + torch.cuda.empty_cache() + return out + + +def load_trace(path: str, layers: int, offset: int, bs: int, steps: int, skip: int, sigma: float, seed: int): + """Per-step routing and lookahead logits ``[steps, layers, bs, E]`` from a routing trace. + + ``path`` holds experts.npy ``[N, 48, 10]`` (routed ids), logit_idx.npy / logit_val.npy + ``[N, 48, 32]`` (each layer's own top-32 router logits) and req_index.npy ``[N]``. Row r of the + batch follows the r-th longest request from its ``skip``-th token on. The routing logits pick + exactly the traced top-10; the lookahead of layer L is layer L's traced top-32 logits plus + N(0, sigma) noise (everything else far below), a stand-in for W_gate[L] . x_{L-1}.""" + import os + from collections import Counter + + import numpy as np + + ex = np.load(os.path.join(path, "experts.npy"), mmap_mode="r") + li = np.load(os.path.join(path, "logit_idx.npy"), mmap_mode="r") + lv = np.load(os.path.join(path, "logit_val.npy"), mmap_mode="r") + req = np.load(os.path.join(path, "req_index.npy")) + reqs = [r for r, n in Counter(req.tolist()).most_common(bs)] + rng = np.random.default_rng(seed) + route = np.full((steps, layers, bs, E), -10.0, np.float32) + look = np.full((steps, layers, bs, E), -30.0, np.float32) + lsl = slice(offset, offset + layers) + for r, rid in enumerate(reqs): + rows = np.flatnonzero(req == rid)[skip: skip + steps] + assert len(rows) == steps, f"request {rid} has {len(rows) + skip} tokens, fewer than skip + steps" + routed = np.asarray(ex[rows, lsl]).astype(np.int64) # [steps, layers, 10] + np.put_along_axis(route[:, :, r], routed, 10.0 - 0.1 * np.arange(TOP_K), -1) + idx = np.asarray(li[rows, lsl]).astype(np.int64) + val = np.asarray(lv[rows, lsl]).astype(np.float32) + rng.normal(0.0, sigma, idx.shape).astype(np.float32) + np.put_along_axis(look[:, :, r], idx, val, -1) + dev = torch.device("cuda") + return (torch.from_numpy(route).to(dev, torch.bfloat16), torch.from_numpy(look).to(dev, torch.bfloat16), reqs) + + +def bench_steady(args, bs: int, arms: list[str]) -> list[dict]: + """ms per decode step in steady state on traced routing: no restore, the cache evolves. + + One stack, one graph per arm, each arm with its own copy of the cache's slot maps, eviction + state and counters: every step replays all arms back to back (rotating their order), so + clock or link drift hits every arm alike. Slot bytes are shared (timing does not read them).""" + from flashlib.kernels.slot_cache import Stat + + from freetoken.moe.prefetch import COPIED, ISSUED, LATE, MISSES, USEFUL + + dev = torch.device("cuda") + steps = args.warm + args.timed + route_all, look_all, reqs = load_trace(args.trace, args.steady_layers, args.layer_offset, bs, steps, + args.skip, args.sigma, args.seed) + layers = args.steady_layers + x = torch.randn(1, H, device=dev, dtype=torch.bfloat16) + probe = torch.randn(32768, H, device=dev, dtype=torch.bfloat16) + per_row_us = time_graph(lambda: F.linear(x, probe), args.reps) * 1000.0 / 32768 + del probe + fill = {us: torch.randn(max(int(us / per_row_us), 1), H, device=dev, dtype=torch.bfloat16) * 0.01 + for us in (args.filler_gdn_us, args.filler_qsa_us)} + # the filler of layer i stands in for its attention: QSA layers are 3, 7, ... (i % 4 == 3) + fillers = [fill[args.filler_qsa_us if i % 4 == 3 else args.filler_gdn_us] for i in range(layers)] + qsa = torch.tensor([i % 4 == 3 for i in range(layers)]) + route_buf = torch.empty_like(route_all[0]) + look_buf = torch.empty_like(look_all[0]) + stack = DecodeStack(layers, args.steady_slots, bs, 0, "on", True, filler=fillers, routing=(route_buf, look_buf)) + cache, pf = stack.cache, stack.cache.prefetch + cache.collect_stats = True + experts = [moe.experts for moe in stack.moes] + state = [cache.slot_for_id, cache.id_of_slot, cache.usage, cache.step, cache.lru_stats, pf.stats] + for name in ("evict_tok", "evict_last_tok", "evict_lc", "evict_ct", "evict_slot_owner", + "evict_slot_last_tok", "evict_slot_lc", "evict_slot_ct"): + if getattr(cache, name) is not None: + state.append(getattr(cache, name)) + init = [t.clone() for t in state] + + def load(snap): + for dst, src in zip(state, snap): + dst.copy_(src) + + def save(snap): + for dst, src in zip(snap, state): + dst.copy_(src) + + route_buf.copy_(route_all[0]) + look_buf.copy_(look_all[0]) + graphs, snaps, keep = {}, {}, [] + for arm in dict.fromkeys(arms): + mode, _, spec = arm.partition(":") + budgets = tuple(int(b) for b in spec.split("/")) if spec and spec != "empty" else (3, 4) + for i in range(1, layers): + experts[i - 1].set_lookahead(experts[i - 1]._lookahead[0], i, budgets[(i + 1) % 4 == 0]) + # every mode switch is read at enqueue time, so each arm's graph bakes its own path in + pf.mode = "on" if mode == "off" else mode + pf.sel_override = torch.full_like(pf.sel, -1) if spec == "empty" else None + keep.append(pf.sel_override) + cache.prefetch = None if mode == "off" else pf + load(init) + stack.step() # compile outside capture + torch.cuda.synchronize() + load(init) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + stack.step() + graphs[arm], snaps[arm] = graph, [t.clone() for t in init] + cache.prefetch, pf.mode, pf.sel_override = pf, "on", None + names = list(graphs) + events = {arm: [] for arm in names} + for t in range(steps): + if t == args.warm: + for arm in names: + snaps[arm][4].zero_() + snaps[arm][5].zero_() + route_buf.copy_(route_all[t]) + look_buf.copy_(look_all[t]) + for j in range(len(names)): + arm = names[(t + j) % len(names)] + load(snaps[arm]) + ev = (torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)) + ev[0].record() + graphs[arm].replay() + ev[1].record() + save(snaps[arm]) + if t >= args.warm: + events[arm].append(ev) + if t % 50 == 49: + torch.cuda.synchronize() # bound the event backlog + torch.cuda.synchronize() + out = [] + n = args.timed + for arm in names: + mode = arm.partition(":")[0] + times = sorted(a.elapsed_time(b) for a, b in events[arm]) + miss = snaps[arm][4][:, Stat.MISS].double().cpu() + row = {"arm": arm, "mean": statistics.fmean(times), "p50": times[len(times) // 2], + "p90": times[int(len(times) * 0.9)], "miss": miss.sum().item() / (n * layers), + "miss1": miss[1:].sum().item() / (n * (layers - 1))} + if mode != "off": + st = snaps[arm][5].double().cpu() + calls = n * (layers - 1) + row.update(issued=st[:, ISSUED].sum().item() / calls, useful=st[:, USEFUL].sum().item() / calls) + row["precision"] = row["useful"] / row["issued"] if row["issued"] else 0.0 + if mode == "measure": + row["coverage"] = st[:, USEFUL].sum().item() / max(st[:, MISSES].sum().item(), 1) + if mode == "on": + # every routed id that was not a normal resident: residual misses plus prefetch hits + row["nonres1"] = row["miss1"] + row["useful"] + for name, mask in (("late_gdn", ~qsa), ("late_qsa", qsa)): + m = mask.clone() + m[0] = False + row[name] = st[m, LATE].sum().item() / max(st[m, COPIED].sum().item(), 1) + row["late"] = st[:, LATE].sum().item() / max(st[:, COPIED].sum().item(), 1) + out.append(row) + print(f"# peak allocated {torch.cuda.max_memory_allocated(dev) / 2**30:.2f} GiB, requests {reqs}", flush=True) + del graphs, stack, cache, pf, events + gc.collect() + torch.cuda.empty_cache() + return out + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--bs", default="1,2") + ap.add_argument("--k", type=int, default=16) + ap.add_argument("--budget", type=int, default=4) + ap.add_argument("--layers", type=int, default=8, help="MoE layers in the decode stack") + ap.add_argument("--slots", type=int, default=512, help="expert slots (2.77 MB each)") + ap.add_argument("--misses", default="0,4", help="demand misses per layer per row") + ap.add_argument("--reps", type=int, default=30) + ap.add_argument("--overlap", type=int, default=1, help="FREETOKEN_MOE_COPY_OVERLAP in the stack (production: 1)") + ap.add_argument("--on", action="store_true", help="time whole decode steps with prefetch off vs on instead") + ap.add_argument("--on-layers", type=int, default=48) + ap.add_argument("--on-slots", type=int, default=0, help="0 = every layer's routed experts + 64 (the warm-up must not evict)") + ap.add_argument("--filler-us", type=float, default=160.0, help="per-layer compute standing in for attention") + ap.add_argument("--useful", default="2,1,2,2", help="useful prefetches per layer, cycled over layers") + ap.add_argument("--delay-us", default="0,150,300", help="sleep ahead of every prefetch copy (0 = none)") + ap.add_argument("--steady", action="store_true", help="steady-state decode on a routing trace instead") + ap.add_argument("--trace", default=None, help="--steady: dir with experts/logit_idx/logit_val/req_index .npy") + ap.add_argument("--steady-layers", type=int, default=24, help="--steady: layers (VRAM: slots * 2.77 MB)") + ap.add_argument("--layer-offset", type=int, default=0, help="--steady: first traced layer") + ap.add_argument("--steady-slots", type=int, default=825, help="--steady: slots (1650 * 24 / 48 by default)") + ap.add_argument("--warm", type=int, default=300, help="--steady: untimed steps before timing") + ap.add_argument("--timed", type=int, default=500) + ap.add_argument("--skip", type=int, default=200, help="--steady: tokens skipped at each request's start") + ap.add_argument("--sigma", type=float, default=0.3, help="--steady: lookahead noise on the traced logits") + ap.add_argument("--seed", type=int, default=0) + ap.add_argument("--filler-gdn-us", type=float, default=171.0, help="--steady: compute before a GDN layer's MoE") + ap.add_argument("--filler-qsa-us", type=float, default=212.0, help="--steady: compute before a QSA layer's MoE") + ap.add_argument("--arms", default="off,measure,on,on:empty,on:2/3,on:3/3,on:4/4", + help="--steady: mode[:gdn/qsa budgets | :empty (on, nothing selected)]") + args = ap.parse_args() + from freetoken.env import ENV + + ENV.MOE_PREFETCH_K.value, ENV.MOE_PREFETCH_BUDGET.value = args.k, 0 if args.steady else args.budget + dev = torch.device("cuda") + if args.steady: + print(f"# {torch.cuda.get_device_name(dev)}, steady state on {args.trace}: {args.steady_layers} layers from " + f"{args.layer_offset}, {args.steady_slots} slots, rule, copy overlap on, filler {args.filler_gdn_us:.0f}/" + f"{args.filler_qsa_us:.0f} us (GDN/QSA), sigma {args.sigma}, {args.warm} warm + {args.timed} timed steps") + print(f"{'bs':>3} {'arm':>10} {'mean ms':>8} {'vs off':>7} {'p50':>7} {'p90':>7} {'miss/L':>7} {'miss/L1+':>8} " + f"{'iss/L':>6} {'use/L':>6} {'prec':>6} {'nonres/L1+':>10} {'late':>6} {'lateGDN':>7} {'lateQSA':>7}") + for bs in (int(b) for b in args.bs.split(",")): + rows = bench_steady(args, bs, args.arms.split(",")) + base = rows[0]["mean"] + for r in rows: + g = r.get + print(f"{bs:>3} {r['arm']:>10} {r['mean']:>8.3f} {r['mean'] / base - 1:>+7.1%} {r['p50']:>7.3f} " + f"{r['p90']:>7.3f} {r['miss']:>7.3f} {r['miss1']:>8.3f} {g('issued', 0):>6.2f} {g('useful', 0):>6.2f} " + f"{g('precision', 0):>6.3f} {g('nonres1', 0):>10.3f} {g('late', 0):>6.3f} {g('late_gdn', 0):>7.3f} " + f"{g('late_qsa', 0):>7.3f}", flush=True) + return + if args.on: + useful = [int(u) for u in args.useful.split(",")] + delays = [float(d) for d in args.delay_us.split(",")] + print(f"# {torch.cuda.get_device_name(dev)}, {args.on_layers} Qwen4ExpMoE nvfp4 layers, " + f"{args.on_slots or 'layers * bs * top-k + 64'} slots, " + f"rule, copy overlap on, filler {args.filler_us:.0f} us/layer, useful {useful}, median of {args.reps} replays") + print(f"{'bs':>3} {'m':>3} {'arm':>10} {'ms/step':>9} {'vs off':>8} {'misses/step':>12} " + f"{'issued/layer':>13} {'useful/layer':>13} {'late':>6}") + for bs in (int(b) for b in args.bs.split(",")): + for m in (int(x) for x in args.misses.split(",")): + slots = args.on_slots or args.on_layers * bs * TOP_K + 64 + rows = bench_on(args.on_layers, slots, bs, m, useful, args.reps, args.filler_us, delays) + base = rows[0]["ms"] + for r in rows: + extra = (f"{r['issued']:>13.2f} {r['useful']:>13.2f} {r['late']:>6.3f}" if "late" in r + else f"{'':>13} {'':>13} {'':>6}") + print(f"{bs:>3} {m:>3} {r['arm']:>10} {r['ms']:>9.3f} {r['ms'] / base - 1:>+8.1%} " + f"{r['misses']:>12d} {extra}", flush=True) + print(f"# peak allocated {torch.cuda.max_memory_allocated(dev) / 2**30:.2f} GiB") + return + print(f"# {torch.cuda.get_device_name(dev)}, K={args.k}, budget={args.budget}, median of {args.reps} graph replays") + print("# predictor alone, us per layer (47 layers, one stream)") + print(f"{'bs':>3} {'gemv':>8} {'select':>8} {'predictor':>10} {'count':>8}") + for bs in (int(b) for b in args.bs.split(",")): + r = bench_predictor(bs, 47, args.k, args.budget, args.reps) + print(f"{bs:>3} {r['gemv']:>8.2f} {r['select']:>8.2f} {r['predictor']:>10.2f} {r['count']:>8.2f}", flush=True) + print(f"# decode stack ({args.layers} Qwen4ExpMoE nvfp4 layers, {args.slots} slots, rule, overlap={args.overlap}), " + "us per layer, prefetch off vs measure") + print(f"{'bs':>3} {'m':>3} {'off':>9} {'measure':>9} {'delta':>8} {'issued/layer':>13}") + for bs in (int(b) for b in args.bs.split(",")): + for m in (int(x) for x in args.misses.split(",")): + r = bench_stack(args.layers, args.slots, bs, m, args.reps, bool(args.overlap)) + print(f"{bs:>3} {m:>3} {r['off']:>9.2f} {r['measure']:>9.2f} {r['measure'] - r['off']:>+8.2f} " + f"{r.get('issued/layer', 0.0):>13.2f}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 9e7bfe7c8d..1227d60c04 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -1205,6 +1205,9 @@ def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput: use_graph = self.graph_runner.can_use_cuda_graph(batch) with self.ctx.forward_batch(batch), self.model.forward_host_ctx(batch, use_graph): logits = self.graph_runner.replay(batch) if use_graph else self.model.forward() + audit = self.moe_offload_cache.audit if self.moe_offload_cache is not None else None + if audit is not None and batch.is_decode: + audit.end_decode_step(self.moe_offload_cache) # every side stream has joined the compute stream if self.cpu_moe_executor is not None: # One pinned read: surfaces a fired flag-handshake watchdog (dead coordinator # -> stale expert outputs) as a loud error instead of silent corruption. @@ -1219,11 +1222,21 @@ def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput: logprobs_out = self.sampler.compute_logprobs(batch_logits, next_tokens_gpu, args) copy_done_event = torch.cuda.Event() copy_done_event.record(self.stream) - if self.moe_offload_cache is not None and self.moe_offload_cache.collect_stats and batch.is_decode: + cache = self.moe_offload_cache + if cache is not None and batch.is_decode and ( + cache.collect_stats or cache.prefetch is not None or cache.verify is not None or cache.audit is not None + ): self._moe_stats_step += 1 if self._moe_stats_step >= MOE_STATS_INTERVAL: self._moe_stats_step = 0 - self._emit_moe_stats() + if cache.prefetch is not None: + self._emit_prefetch_stats() + if cache.verify is not None: + self._emit_verify_stats() + if cache.audit is not None: + self._emit_audit_stats() + if cache.collect_stats: + self._emit_moe_stats() if logprobs_out is None: return ForwardOutput(next_tokens_gpu, next_tokens_cpu, copy_done_event) chosen_logprobs, top_ids, top_logprobs = logprobs_out @@ -1267,6 +1280,39 @@ def _warmup_prefill_lens(self) -> list[int]: lens.add(8193) # chunk_delta_h NT bucket 2 (NT > 128) return sorted(n for n in lens if 2 <= n <= cap) + def _emit_prefetch_stats(self) -> None: + """Report one window of the FREETOKEN_MOE_PREFETCH counters (one host sync), then restart them.""" + from freetoken.moe.prefetch import format_per_layer, format_summary, summarize + + prefetch = self.moe_offload_cache.prefetch + window = prefetch.take_window() + stats = summarize(window, prefetch.mode) + if not stats["layer_calls"]: + return + logger.info_rank0(f"MoE prefetch {prefetch.mode} ({MOE_STATS_INTERVAL} decode steps): {format_summary(stats)}") + if ENV.MOE_PREFETCH_DEBUG: + logger.info_rank0(f"MoE prefetch per layer (useful/issued/misses per call): {format_per_layer(window)}") + + def _emit_verify_stats(self) -> None: + """Report one window of the FREETOKEN_MOE_PREFETCH_VERIFY counts and the new bad records (host syncs).""" + _log_lines(self.moe_offload_cache.verify.report_window(MOE_STATS_INTERVAL)) + + def _emit_audit_stats(self) -> None: + """Report one window of FREETOKEN_MOE_SLOT_AUDIT totals and the new bad slots (host syncs).""" + _log_lines(self.moe_offload_cache.audit.report_window(MOE_STATS_INTERVAL)) + + def _log_prefetch_totals(self) -> None: + from freetoken.moe.prefetch import format_per_layer, format_summary, summarize + + prefetch = self.moe_offload_cache.prefetch if self.moe_offload_cache is not None else None + if prefetch is None: + return + totals = prefetch.totals + prefetch.stats.cpu() + stats = summarize(totals, prefetch.mode) + if stats["layer_calls"]: + logger.info_rank0(f"MoE prefetch {prefetch.mode} (session, {stats['tokens']} tokens): {format_summary(stats)}") + logger.info_rank0(f"MoE prefetch per layer (useful/issued/misses per call): {format_per_layer(totals)}") + def _emit_moe_stats(self) -> None: """Report one window of expert-cache behaviour, then reset the miss counters. @@ -1389,11 +1435,32 @@ def _warmup_prefill(self) -> None: ) def shutdown(self) -> None: + try: + self._log_prefetch_totals() + except Exception as exc: # noqa: BLE001 -- a diagnostic must never block shutdown + logger.warning(f"MoE prefetch totals unavailable at shutdown: {exc}") + try: + verify = self.moe_offload_cache.verify if self.moe_offload_cache is not None else None + if verify is not None: + _log_lines(verify.report_session()) + except Exception as exc: # noqa: BLE001 + logger.warning(f"MoE verify totals unavailable at shutdown: {exc}") + try: + audit = self.moe_offload_cache.audit if self.moe_offload_cache is not None else None + if audit is not None: + _log_lines(audit.report_session()) + except Exception as exc: # noqa: BLE001 + logger.warning(f"MoE slot audit totals unavailable at shutdown: {exc}") self.graph_runner.destroy_cuda_graphs() torch.distributed.destroy_process_group() destroy_distributed() +def _log_lines(lines: list[tuple[str, str]]) -> None: + for level, line in lines: + (logger.warning_rank0 if level == "warning" else logger.info_rank0)(line) + + def _profile_gpu(index: "int | None" = None) -> Tuple[str | None, str | None]: """(name, uuid) of visible device ``index`` (default: the current, i.e. bound, device); (None, None) without CUDA.""" if not torch.cuda.is_available(): diff --git a/python/freetoken/env.py b/python/freetoken/env.py index 0e56633f48..9a06b1c029 100644 --- a/python/freetoken/env.py +++ b/python/freetoken/env.py @@ -79,6 +79,20 @@ class EnvClassSingleton: WARMUP_MAX_LEN = EnvInt(0) # GPU-decode MoE: run a layer's expert miss copy on a side stream under the shared expert MOE_COPY_OVERLAP = EnvBool(False) + # GPU-decode MoE cross-layer prefetch: off | measure (predict and count, no copies) | on (predict and copy) + MOE_PREFETCH = EnvStr("off") + # router-lookahead candidates per layer, and the non-resident ones kept per target layer + # (0 = 3 before GDN layers, 4 before full-attention layers) + MOE_PREFETCH_K = EnvInt(16) + MOE_PREFETCH_BUDGET = EnvInt(0) + # measure mode: also log the per-layer breakdown with every summary + MOE_PREFETCH_DEBUG = EnvBool(False) + # debug: check every GPU decode layer call's slot maps (meta) and slot bytes before/after the GEMM (1) + MOE_PREFETCH_VERIFY = EnvStr("0") + # routed experts a call byte-checks (scratch rows, ~2.8 MB each on Qwen3.8-Flash) + MOE_PREFETCH_VERIFY_ROWS = EnvInt(20) + # debug: record every expert slot writer and audit every held slot each N decode steps (0 = off) + MOE_SLOT_AUDIT = EnvInt(0) def __new__(cls): # single instance diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index 918a3f9731..56dcd712ab 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -194,6 +194,13 @@ def __init__( prefix=prefix, ) self.offload_cache: OffloadMoeCache | None = None + # FREETOKEN_MOE_PREFETCH: (next MoE layer's router op, its layer id, its budget); the + # underscore keeps the borrowed router out of this layer's state dict + self._lookahead: tuple[BaseOP, int, int] | None = None + + def set_lookahead(self, gate: BaseOP, target_layer: int, budget: int) -> None: + """Let this layer's GPU decode predict ``target_layer``'s experts with that layer's router.""" + self._lookahead = (gate, target_layer, budget) def forward( self, @@ -301,9 +308,25 @@ def _decode_routed( return executor.decode(self.layer_id, hidden_states, topk_weights, topk_ids) if cache.decode_target == "hybrid": return self._decode_hybrid(cache, hidden_states, topk_weights, topk_ids) + prefetch, verify = cache.prefetch, cache.verify + if verify is not None: + verify.begin(cache, self.layer_id, topk_ids, prefetch and prefetch.plan_of(self.layer_id)) cache.ensure_experts(self.layer_id, topk_ids, router_logits=router_logits) + target = None + if prefetch is not None: + # on: install the next layer's prediction (forked at this block's start) on this stream, the slot + # maps' only writer; before this layer's copy join, so a late prefetch copy never delays it + target = prefetch.install(self._lookahead and self._lookahead[1]) + prefetch.join_copy(self.layer_id) cache.copy_missing() - return self._expert_gemm( + if prefetch is not None: + prefetch.join_and_count(self.layer_id, topk_ids, cache.id_of_slot, cache.num_indices) + if target is not None: + prefetch.issue_copy(target) + if verify is not None: + # after issue_copy: the next layer's prefetch copy starts when it would without the checks + verify.before_gemm(cache, topk_ids, None if target is None else prefetch.plan_of(target)) + out = self._expert_gemm( cache, hidden_states, topk_weights, @@ -313,6 +336,29 @@ def _decode_routed( alphas=cache.alphas_for_slots(self.layer_id), is_prefill=False, ) + if verify is not None: + verify.after_gemm(cache) + return out + + def fork_lookahead(self, hidden_states: torch.Tensor) -> None: + """FREETOKEN_MOE_PREFETCH: at the start of this layer's MoE block, predict the next layer's + experts from this block's router input on the predictor stream, beside this layer's router + and ensure. Only on the GPU decode path that ``_decode_routed`` takes and not under TP, and + never toward a CPU-decoded layer, whose ensure would not use it.""" + cache = self.offload_cache + if ( + self._lookahead is None + or cache is None + or cache.prefetch is None + or self.tp_size > 1 + or get_global_ctx().batch.is_prefill + or cache.is_cpu_layer(self.layer_id) + or cache.decode_target == "hybrid" + ): + return + gate, target, budget = self._lookahead + if not cache.is_cpu_layer(target): + cache.prefetch.fork(target, hidden_states, gate, cache.slot_for_id[target], budget) def _decode_hybrid( self, diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py index 5d4c6f9d2f..7f4f826807 100644 --- a/python/freetoken/models/qwen4_exp/model.py +++ b/python/freetoken/models/qwen4_exp/model.py @@ -26,7 +26,7 @@ from .attention import Qwen4ExpAttention from .hc import GatedResidual -from .moe import Qwen4ExpMoE +from .moe import Qwen4ExpMoE, wire_router_lookahead from .ple import PLELayer from freetoken.models.blocks import embed_input_ids from freetoken.models.qwen3_vl.vision import Qwen3VLVisionModel, QwenVLVisionMixin @@ -101,6 +101,7 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model") -> None: for layer_id in range(config.num_layers) ] ) + wire_router_lookahead([layer.mlp for layer in self.layers.op_list], config) self.hyper_connection_mixer = GatedResidual(config, use_combine=False, prefix=f"{prefix}.hyper_connection_mixer") # plain tuple (not an OP child), so it never shows up in the state dict self._ple = tuple(layer.ple for layer in self.layers.op_list if layer.ple is not None) diff --git a/python/freetoken/models/qwen4_exp/moe.py b/python/freetoken/models/qwen4_exp/moe.py index 6f729231f1..4a1907d2ce 100644 --- a/python/freetoken/models/qwen4_exp/moe.py +++ b/python/freetoken/models/qwen4_exp/moe.py @@ -36,6 +36,10 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: num_tokens, hidden_dim = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_dim) se, ex = self.shared_expert, self.experts + if isinstance(ex, OffloadMoELayer): + # FREETOKEN_MOE_PREFETCH: the next layer's prediction needs only this block's input, so it + # starts here and runs beside this layer's router and ensure + ex.fork_lookahead(hidden_states) if isinstance(ex, OffloadMoELayer) and ex.decode_side_applies(): # FREETOKEN_MOE_COPY_OVERLAP: the shared expert runs on a side stream while the router, # ensure and miss copy stay back to back on this one; joined before the combine @@ -65,4 +69,15 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return shared_gate_mul_add(routed, shared, gate).view(num_tokens, hidden_dim) -__all__ = ["Qwen4ExpMoE"] +def wire_router_lookahead(moes: list[Qwen4ExpMoE], config: ModelConfig) -> None: + """Give each offloaded MoE block the next layer's router, for FREETOKEN_MOE_PREFETCH's + lookahead (inert when it is off). ``moes[i]`` is decoder layer ``i``'s block.""" + from freetoken.moe.prefetch import default_budget + + for target in range(1, len(moes)): + experts = moes[target - 1].experts + if isinstance(experts, OffloadMoELayer): + experts.set_lookahead(moes[target].gate, target, default_budget(config.is_linear_layer(target))) + + +__all__ = ["Qwen4ExpMoE", "wire_router_lookahead"] diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 477504bcf3..246b6b248f 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -2,6 +2,7 @@ import math import os +import weakref from contextlib import contextmanager from dataclasses import dataclass from typing import Iterator @@ -153,6 +154,13 @@ class OffloadMoeCache: # GPU decode only: the caller's work that is independent of the routed experts (qwen4_exp: the # shared expert) runs on decode_copy_stream beside the miss copy. None = FREETOKEN_MOE_COPY_OVERLAP. decode_copy_overlap: bool | None = None + # GPU decode only: cross-layer expert prefetch (moe/prefetch.py). None = FREETOKEN_MOE_PREFETCH. + prefetch_mode: str | None = None + # GPU decode only: debug slot checks (moe/verify.py). None = FREETOKEN_MOE_PREFETCH_VERIFY. + verify_mode: str | None = None + # GPU decode only: debug slot writer history + audit every N decode steps (moe/slot_audit.py). + # None = FREETOKEN_MOE_SLOT_AUDIT, 0 = off. + slot_audit: int | None = None def __post_init__(self) -> None: from freetoken.moe.scored_ensure import POLICY_IDS @@ -287,6 +295,9 @@ def __post_init__(self) -> None: # _pending_whole_layer records WHICH staged it: the pageable branch is only sound after materialize_layer self._pending_src_layer: int | None = None self._pending_whole_layer = False + # the slot audit's label for the staged copy, and the most entries its plan can have + self._pending_kind = 0 + self._pending_width: int | None = None # Per-bank [2, num_experts, ...] double-buffer views over the slot cache's # first 2 * num_experts slots (set up when prefill_overlap is enabled). self.prefill_bank_buffers: list[torch.Tensor] = [] @@ -317,6 +328,70 @@ def __post_init__(self) -> None: self._decode_copy_events: list[tuple[torch.cuda.Event, torch.cuda.Event]] = [] if self.decode_copy_overlap and self.device.type == "cuda" and self.decode_target == "gpu": self._init_decode_copy_overlap() + self.prefetch = self._init_prefetch() + self.verify = self._init_verify() + self.audit = self._init_audit() + + def _init_audit(self): + from freetoken.kernel.fast_index_copy import _skip_fast_index_copy_enabled + from freetoken.moe.slot_audit import SlotAuditor, resolve_interval + + from_env = self.slot_audit is None + self.slot_audit = resolve_interval(self.slot_audit) + if not self.slot_audit or self.device.type != "cuda": + return None + if self.decode_target != "gpu": + if not from_env: + raise ValueError(f"FREETOKEN_MOE_SLOT_AUDIT={self.slot_audit} needs GPU decode, not {self.decode_target!r}") + logger.warning(f"FREETOKEN_MOE_SLOT_AUDIT={self.slot_audit} ignored: decode target is {self.decode_target!r}") + self.slot_audit = 0 + return None + if _skip_fast_index_copy_enabled(): + logger.warning("FREETOKEN_MOE_SLOT_AUDIT ignored: FREETOKEN_SKIP_FAST_INDEX_COPY makes every copy a no-op") + self.slot_audit = 0 + return None + return SlotAuditor(self.num_layers, self.num_experts, self.device, interval=self.slot_audit) + + def _init_verify(self): + from freetoken.moe.verify import ExpertVerifier, resolve_mode + + from_env = self.verify_mode is None + self.verify_mode = resolve_mode(self.verify_mode) + if self.verify_mode == "off" or self.device.type != "cuda": + return None + if self.decode_target != "gpu": + if not from_env: + raise ValueError(f"FREETOKEN_MOE_PREFETCH_VERIFY={self.verify_mode} needs GPU decode, not {self.decode_target!r}") + logger.warning(f"FREETOKEN_MOE_PREFETCH_VERIFY={self.verify_mode} ignored: decode target is {self.decode_target!r}") + self.verify_mode = "off" + return None + return ExpertVerifier(self.num_layers, self.num_experts, self.device, mode=self.verify_mode) + + def _init_prefetch(self): + from freetoken.moe.prefetch import ExpertPrefetcher, resolve_mode + + from_env = self.prefetch_mode is None + self.prefetch_mode = resolve_mode(self.prefetch_mode) + if self.prefetch_mode == "off": + return None + if self.decode_target != "gpu": + if not from_env: + raise ValueError(f"FREETOKEN_MOE_PREFETCH={self.prefetch_mode} needs GPU decode, not {self.decode_target!r}") + # the process-wide flag skips hybrid/CPU-decoded caches, like FREETOKEN_MOE_COPY_OVERLAP + logger.warning(f"FREETOKEN_MOE_PREFETCH={self.prefetch_mode} ignored: decode target is {self.decode_target!r}") + self.prefetch_mode = "off" + return None + if self.device.type != "cuda": + return None + prefetch = ExpertPrefetcher(self.num_layers, self.num_experts, self.device, mode=self.prefetch_mode) + # a proxy, not a cycle: dropping the cache must free its slot banks without waiting for gc + prefetch.cache = weakref.proxy(self) + return prefetch + + @property + def prefetch_on(self) -> bool: + """FREETOKEN_MOE_PREFETCH=on: every ensure keys prefetched slots low priority.""" + return self.prefetch is not None and self.prefetch.mode == "on" def _init_decode_copy_overlap(self) -> None: # a dedicated stream, not one of torch's 32 pooled ones: a pooled stream can alias the engine @@ -463,6 +538,10 @@ def set_bank_sources( ) self.banks = [(self.bank_sources[n], self.bank_caches[n]) for n in self.bank_schema] self._build_copy_plan() + if self.verify is not None: + self.verify.bind(self) + if self.audit is not None: + self.audit.bind(self) if self.prefill_overlap: self._init_prefill_overlap_buffers() @@ -600,6 +679,8 @@ def rebuild(self, cache_size: int) -> None: ) 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 + if self.verify is not None: + self.verify.bind(self) # 4. Reallocate cache_size-shaped bookkeeping; reset the slot map (cold start). self.slot_for_id.fill_(-1) self.id_of_slot = torch.full((cache_size,), -1, dtype=torch.int32, device=self.device) @@ -625,6 +706,10 @@ def rebuild(self, cache_size: int) -> None: self.stat_fetched_layer.zero_() self.stat_steps_layer.zero_() self.decode_freq.zero_() + if self.prefetch is not None: + self.prefetch.reset() + if self.audit is not None: + self.audit.bind(self) # new slot arrays: a new history self.prefill_hit_rows = 0 self.prefill_total_rows = 0 self._hit_d2d_fallback_logged = False # geometry changed; re-log if still unusable @@ -737,8 +822,12 @@ def _init_prefill_overlap_buffers(self) -> None: ) self._prefill_hit_num = torch.zeros((1,), dtype=torch.int64, device=self.device) - def _invalidate_prefill_buffer(self, buffer_id: int) -> None: + def _invalidate_prefill_buffer(self, buffer_id: int, layer_id: int = -1) -> None: slot_start = buffer_id * self.num_experts + if self.audit is not None: + from freetoken.moe.slot_audit import INVALIDATE + + self.audit.record_clear(self, INVALIDATE, layer_id, slot_start, self.num_experts) # One fixed-shape launch. The previous boolean-mask index # (slot_for_id[old_ids[old_ids >= 0]] = -1) has a data-dependent shape, so every # call hid a device-to-host sync -- and with two buffer reuses per chunk x 48 @@ -790,9 +879,13 @@ def prefetch_prefill_layer(self, layer_id: int) -> None: ) def copy() -> None: - self._invalidate_prefill_buffer(buffer_id) + self._invalidate_prefill_buffer(buffer_id, layer_id) for (per_layer, _), buffer in zip(self.banks, self.prefill_bank_buffers): buffer[buffer_id].copy_(per_layer[layer_id], non_blocking=True) + if self.audit is not None: + from freetoken.moe.slot_audit import PREFILL_BUFFER + + self.audit.record_range(self, PREFILL_BUFFER, layer_id, buffer_id * self.num_experts, self.num_experts, meta=False) if self._prefill_hit_d2d_active: self._prefetch_split(layer_id, buffer_id) @@ -892,11 +985,16 @@ def _prefetch_split(self, layer_id: int, buffer_id: int) -> None: self._prefill_hit_num, blocks_per_bank=64, ) + if self.audit is not None: + from freetoken.moe.slot_audit import PREFILL_HIT_D2D + + self.audit.record_plan(self, PREFILL_HIT_D2D, layer_id, self._prefill_hit_dst, None, self._prefill_hit_num, + meta=False, aux=self._prefill_hit_src, row_base=buffer_id * E) miss = np.nonzero(~hit_mask)[0] with torch.cuda.stream(self.prefill_copy_stream): if self._prefill_buffer_has_release_event[buffer_id]: self.prefill_copy_stream.wait_event(self.prefill_release_events[buffer_id]) - self._invalidate_prefill_buffer(buffer_id) + self._invalidate_prefill_buffer(buffer_id, layer_id) if miss.size: run_starts = np.concatenate(([0], np.nonzero(np.diff(miss) != 1)[0] + 1)) starts = miss[run_starts] @@ -921,6 +1019,11 @@ def _prefetch_split(self, layer_id: int, buffer_id: int) -> None: torch.tensor(nbytes, dtype=torch.int64), torch.cuda.current_stream(self.device).cuda_stream, ) + if self.audit is not None: + from freetoken.moe.slot_audit import PREFILL_SPLIT_H2D + + # every row gets at least its small banks from the host, the misses all of them + self.audit.record_range(self, PREFILL_SPLIT_H2D, layer_id, buffer_id * E, E, meta=False) self.prefill_ready_events[buffer_id].record(self.prefill_copy_stream) def wait_prefill_layer(self, layer_id: int) -> tuple[torch.Tensor, ...]: @@ -971,18 +1074,49 @@ def ensure_experts( self.decode_freq[layer_id].scatter_add_(0, ids, torch.ones_like(ids)) self._pending_src_layer = layer_id self._pending_whole_layer = False - if self.cache_policy_id == 0: + audit = self.audit + if audit is not None and update_state and layer_id == self.evict_clock_layer(): + audit.begin_step() + lowpri = self.prefetch_on + width = expert_ids.numel() + if self.cache_policy_id == 0 and not lowpri: ensure_experts(self, layer_id, expert_ids) - return - ensure_experts_scored( - self, - layer_id, - expert_ids, - bump_tok=update_state and layer_id == self.evict_clock_layer(), - update_state=update_state, - router_logits=router_logits, - pin_since=pin_since, - ) + else: + # prefetch on: LRU runs the vendored kernel (flashlib's cannot key low-priority slots) + ensure_experts_scored( + self, + layer_id, + expert_ids, + bump_tok=update_state and layer_id == self.evict_clock_layer(), + update_state=update_state, + router_logits=router_logits, + pin_since=pin_since, + lowpri=lowpri, + pf_count=self.prefetch.count_args(layer_id) if lowpri and update_state else None, + ) + if audit is not None: + from freetoken.moe import slot_audit as sa + + kind = sa.DEMAND_INSTALL if update_state else sa.PREFILL_INSTALL + self._pending_kind = sa.DEMAND_COPY if update_state else sa.PREFILL_COPY + self._pending_width = width + audit.record_plan(self, kind, layer_id, self.evict_slots, self.src_indices, self.num_indices, meta=True, width=width) + + def prefetch_ensure( + self, layer_id: int, query: torch.Tensor, dst_slots: torch.Tensor, src_rows: torch.Tensor, + num: torch.Tensor, stats_row: torch.Tensor, ready: torch.Tensor, budget: int | None = None, + ) -> None: + """Install the first ``budget`` of ``query``'s non-resident experts of ``layer_id`` as + low-priority slots and plan their copy into ``dst_slots``/``src_rows``/``num`` (never the + demand plan, which copy(L-1) still reads). Runs on the compute stream like every other slot + map writer.""" + from freetoken.moe.offload_kernels import prefetch_ensure_experts + + prefetch_ensure_experts(self, layer_id, query, dst_slots, src_rows, num, stats_row, ready, budget=budget) + if self.audit is not None: + from freetoken.moe.slot_audit import PREFETCH_INSTALL + + self.audit.record_plan(self, PREFETCH_INSTALL, layer_id, dst_slots, src_rows, num, meta=True) def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None: """Capped-fetch LRU for the hybrid backend. @@ -1010,18 +1144,36 @@ def materialize_layer(self, layer_id: int) -> None: self._pending_src_layer = layer_id self._pending_whole_layer = True + audit = self.audit + if audit is not None: + from freetoken.moe import slot_audit as sa + + # the kernel also frees this layer's slots above E; record them while they still name it + audit.record_clear(self, sa.MATERIALIZE_CLEAR, layer_id, 0, self.cache_size, layer_above_e=True) materialize_layer(self, layer_id) if self.evict_slot_owner is not None: self.evict_slot_owner[: self.num_experts].fill_(-1) # it installs the layer into these slots + if audit is not None: + self._pending_kind = sa.MATERIALIZE_COPY + self._pending_width = self.num_experts + audit.record_plan(self, sa.MATERIALIZE_INSTALL, layer_id, self.evict_slots, self.src_indices, self.num_indices, + meta=True, width=self.num_experts) def reset(self) -> None: from freetoken.moe.offload_kernels import reset_cache + if self.audit is not None: + from freetoken.moe.slot_audit import RESET + + self.audit.record_clear(self, RESET, -1, 0, self.cache_size) reset_cache(self) # 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.reset_evict_state() + if self.prefetch is not None: + # the graph runner resets after each capture: drop the capture-time counts too + self.prefetch.reset_counters() def reset_stats(self) -> None: self.prefill_hit_rows = 0 @@ -1161,36 +1313,52 @@ def copy_missing(self, layer_id: int | None = None) -> None: # never CUDA-graph captured: prefill is not captured, and decode never reaches this branch (it routes to the CPU executor) for per_layer, cache in self.banks: cache[: self.num_experts].copy_(per_layer[layer_id]) + if self.audit is not None: + from freetoken.moe.slot_audit import PAGEABLE_COPY + + self.audit.record_range(self, PAGEABLE_COPY, layer_id, 0, self.num_experts, meta=False) return + self.copy_rows(layer_id, self.evict_slots, self.src_indices, self.num_indices, slim=_SLIM_COPY, + kind=self._pending_kind, width=self._pending_width) + + def copy_rows( + self, layer_id: int, dst_slots: torch.Tensor, src_rows: torch.Tensor, num: torch.Tensor, *, slim: bool, + kind: int = 0, width: int | None = None, + ) -> None: + """Copy host rows ``src_rows[:num]`` of ``layer_id``'s banks into slots ``dst_slots[:num]`` + on the current stream (the demand plan, or a prefetch plan on the prefetch copy stream); + ``kind`` labels the copy for the slot audit and ``width`` bounds ``num`` for it.""" + self._copy_rows(layer_id, dst_slots, src_rows, num, slim=slim) + if self.audit is not None: + from freetoken.moe.slot_audit import OTHER_COPY + + self.audit.record_plan(self, kind or OTHER_COPY, layer_id, dst_slots, src_rows, num, meta=False, width=width) + + def _copy_rows( + self, layer_id: int, dst_slots: torch.Tensor, src_rows: torch.Tensor, num: torch.Tensor, *, slim: bool + ) -> None: if self._copy_fused_ok: from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit, fast_index_copy_multi_slim_jit # One launch copies the missing rows for every bank (instead of one launch per - # bank). evict_slots/src_indices/num_indices are shared across banks; - # src_indices holds layer-local expert rows, resolved against this layer's - # source pointers (layer_id is a static int per captured graph node). - copy = fast_index_copy_multi_slim_jit if _SLIM_COPY else fast_index_copy_multi_jit - copy( - self._copy_dst_ptrs, - self._copy_src_ptrs[layer_id], - self._copy_feat_bytes, - self.evict_slots, - self.src_indices, - self.num_indices, - blocks_per_bank=self._copy_blocks_per_bank, - ) + # bank). The slot/row plan is shared across banks; src_rows holds layer-local + # expert rows, resolved against this layer's source pointers (layer_id is a + # static int per captured graph node). + if slim: + fast_index_copy_multi_slim_jit( + self._copy_dst_ptrs, self._copy_src_ptrs[layer_id], self._copy_feat_bytes, dst_slots, src_rows, num, + ) + else: + fast_index_copy_multi_jit( + self._copy_dst_ptrs, self._copy_src_ptrs[layer_id], self._copy_feat_bytes, dst_slots, src_rows, num, + blocks_per_bank=self._copy_blocks_per_bank, + ) return from freetoken.kernel import fast_index_copy_jit for per_layer, cache in self.banks: - fast_index_copy_jit( - cache, - self.evict_slots, - per_layer[layer_id], - self.src_indices, - self.num_indices, - ) + fast_index_copy_jit(cache, dst_slots, per_layer[layer_id], src_rows, num) @contextmanager def decode_side(self, layer_id: int) -> Iterator[None]: diff --git a/python/freetoken/moe/offload_kernels.py b/python/freetoken/moe/offload_kernels.py index bb03aae14c..5993738d10 100644 --- a/python/freetoken/moe/offload_kernels.py +++ b/python/freetoken/moe/offload_kernels.py @@ -49,10 +49,15 @@ def ensure_experts_scored( update_state: bool, router_logits: torch.Tensor | None = None, pin_since: torch.Tensor | None = None, + lowpri: bool = False, + pf_count: tuple | None = None, ) -> None: - """``ensure_experts`` for the scored ``--moe-cache-policy`` kinds (kd, kdfb, rule).""" + """``ensure_experts`` for the scored ``--moe-cache-policy`` kinds (kd, kdfb, rule), and for + every policy under FREETOKEN_MOE_PREFETCH=on (``lowpri``; ``pf_count`` is the prefetcher's + ``(stats_row, ready, num, rows)`` when this layer was prefetched this step).""" from freetoken.moe.scored_ensure import scored_ensure + stats_row, ready, num, rows = pf_count if pf_count is not None else (None, None, None, 0) scored_ensure( expert_ids, cache.slot_for_id.view(-1), @@ -63,6 +68,23 @@ def ensure_experts_scored( cache.src_indices, cache.evict_slots, cache.num_indices, + router_logits=router_logits if cache.near_miss_thr is not None else None, + near_miss_thr=cache.near_miss_thr or 0.0, + stats=cache.lru_stats[layer_id] if cache.collect_stats else None, + bump_tok=bump_tok, + update_state=update_state, + pin_since=pin_since, + lowpri=lowpri, + pf_stats=stats_row, + pf_ready=ready, + pf_num=num, + pf_rows=rows, + **_evict_state(cache, layer_id), + ) + + +def _evict_state(cache, layer_id: int) -> dict: + return dict( policy=cache.cache_policy_id, num_layers=cache.num_layers, num_experts=cache.num_experts, @@ -75,13 +97,21 @@ def ensure_experts_scored( slot_last_tok=cache.evict_slot_last_tok, slot_lc=cache.evict_slot_lc, slot_ct=cache.evict_slot_ct, - router_logits=router_logits if cache.near_miss_thr is not None else None, - near_miss_thr=cache.near_miss_thr or 0.0, - stats=cache.lru_stats[layer_id] if cache.collect_stats else None, id_base=layer_id * cache.num_experts, - bump_tok=bump_tok, - update_state=update_state, - pin_since=pin_since, + ) + + +def prefetch_ensure_experts( + cache, layer_id: int, query: torch.Tensor, dst_slots: torch.Tensor, src_rows: torch.Tensor, + num: torch.Tensor, stats_row: torch.Tensor, ready: torch.Tensor, budget: int | None = None, +) -> None: + """FREETOKEN_MOE_PREFETCH=on: install the first ``budget`` of ``query``'s non-resident experts + of ``layer_id`` as low-priority slots and plan their copy (scored_ensure.prefetch_ensure).""" + from freetoken.moe.scored_ensure import prefetch_ensure + + prefetch_ensure( + query, cache.slot_for_id.view(-1), cache.id_of_slot, cache.usage, cache.step, + dst_slots, src_rows, num, stats_row, ready, budget=budget, **_evict_state(cache, layer_id), ) diff --git a/python/freetoken/moe/prefetch.py b/python/freetoken/moe/prefetch.py new file mode 100644 index 0000000000..01129d3997 --- /dev/null +++ b/python/freetoken/moe/prefetch.py @@ -0,0 +1,438 @@ +"""Cross-layer expert prefetch for GPU decode (``FREETOKEN_MOE_PREFETCH``). + +A training-free router lookahead predicts layer L's experts from the tensor layer L-1's router +consumed (``W_gate[L] . x_{L-1}``, top-K) and keeps the first BUDGET candidates that are not +resident. At bs 2 both rows' top-K lists merge rank by rank (row 0's first, then row 1's +first, ...), duplicates dropped, under one budget. + +``measure``: after layer L's ensure, count how many of them layer L actually routed. Nothing is +copied and the cache is never touched. The predictor reads ``slot_for_id[L]`` without ordering +against the ensures; a late read only skews the counters, never the model. + +The prediction forks at the start of layer L-1's MoE block, where x_{L-1} first exists, so it runs +beside layer L-1's own router and ensure: + + compute: -fork(L)-> router_{L-1} ensure(L-1) copy(L-1) GEMM(L-1) ... ensure(L) copy(L) -join-> count(L) GEMM(L) + predict: \\-> router_L(x_{L-1}) -> top-K -> select(L) -rec sel(L) -/ + +``on``: the predictor only merges the candidates (it reads no cache state); the compute stream +installs the first BUDGET of them that are not resident as low-priority slots (prefetch_ensure: +held, usage 0, keyed between the empty slots and every resident) and a third stream copies them +from layer L's host banks with the slim copy kernel: + + compute: -fork(L)-> ... ensure(L-1) [wait sel(L)] prefetch_ensure(L) [wait pfcopy(L-1)] copy(L-1) -rec dcopy(L)-> GEMM(L-1) ... ensure(L) [wait pfcopy(L)] copy(L) GEMM(L) + predict: \\-> router_L -> top-K -> merge(L) -rec sel(L) + copy: [wait dcopy(L)] slim copy(L) -> ready(L) -rec pfcopy(L) + +Every writer of the slot maps and their mirrors runs on the compute stream, so the maps have one +writer in stream order and no cross-stream edge guards them, and the installs follow layer L's +residency right after ensure(L-1) however late the predictor runs. dcopy(L) follows +prefetch_ensure(L) and the demand copy(L-1) on the compute stream, so the copy stream reads a +finished plan and owns the host link second. prefetch_ensure(L) keeps usage == step, i.e. layer +L-1's routed slots, pinned while GEMM(L-1) may read them; the compute stream joins the prefetch copy +before its own copy(L), so ensure(L) may reuse a prefetched slot as a demand victim and GEMM(L) only +ever reads landed bytes. A demand hit on a prefetched slot sets its usage (a normal resident) and +counts as useful. Layer 0 is never predicted. +""" +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from freetoken.env import ENV +from freetoken.moe.slot_audit import PREFETCH_COPY +from freetoken.utils import init_logger + +logger = init_logger(__name__) + +MODES = ("off", "measure", "on") +# non-resident candidates per target layer: GDN layers leave a shorter copy window than QSA ones +GDN_BUDGET = 3 +ATTN_BUDGET = 4 +# counters[layer, col]; on mode adds LATE and COPIED (calls whose plan had a copy) in a wider stats row +ISSUED, USEFUL, RESIDENT_HITS, CALLS, ROWS, MISSES, LATE, COPIED = range(8) +NUM_COLS = 6 +NUM_STAT_COLS = 8 +# the select kernel's dedup tile is (rows * k)^2; larger (eager-only) decode batches skip the prediction +MAX_ROWS = 8 + + +def resolve_mode(mode: str | None = None) -> str: + mode = (ENV.MOE_PREFETCH.value if mode is None else mode).strip().lower() + if mode not in MODES: + raise ValueError(f"FREETOKEN_MOE_PREFETCH={mode!r}: expected one of {', '.join(MODES)}") + return mode + + +def default_budget(before_linear: bool) -> int: + """Non-resident candidates kept for a target layer, by the target's attention kind.""" + return GDN_BUDGET if before_linear else ATTN_BUDGET + + +def dedicated_stream(device: torch.device, *, highest_priority: bool = False) -> torch.cuda.ExternalStream: + # not one of torch's 32 pooled streams: a pooled one can alias the engine or capture stream + from cuda.bindings import runtime as cudart + + with torch.cuda.device(device): + if highest_priority: + err, _, greatest = cudart.cudaDeviceGetStreamPriorityRange() + if err == cudart.cudaError_t.cudaSuccess: + err, handle = cudart.cudaStreamCreateWithPriority(cudart.cudaStreamNonBlocking, greatest) + else: + err, handle = cudart.cudaStreamCreateWithFlags(cudart.cudaStreamNonBlocking) + if err != cudart.cudaError_t.cudaSuccess: + raise RuntimeError(f"creating a dedicated CUDA stream failed: {err}") + return torch.cuda.ExternalStream(int(handle), device=device) + + +@triton.jit(do_not_specialize=["stride", "budget"]) +def _lookahead_select_kernel( + logits_ptr, stride, resident_ptr, sel_ptr, res_ptr, budget, + ROWS: tl.constexpr, ROWS_P: tl.constexpr, E: tl.constexpr, E_P: tl.constexpr, + K: tl.constexpr, N_P: tl.constexpr, W: tl.constexpr, W_P: tl.constexpr, +): + b = tl.arange(0, ROWS_P) + e = tl.arange(0, E_P) + live = (b < ROWS)[:, None] & (e < E)[None, :] + v = tl.load(logits_ptr + b[:, None] * stride + e[None, :], mask=live, other=float("-inf")).to(tl.float32) + # candidates in merged order: position r * ROWS + b holds row b's rank-r id + j = tl.arange(0, N_P) + cand = tl.full([N_P], -1, tl.int32) + for r in tl.static_range(K): + top = tl.argmax(v, axis=1).to(tl.int32) # ties go to the lowest id + at = (j[:, None] == r * ROWS + b[None, :]) & (b < ROWS)[None, :] + cand = tl.where(tl.sum(at.to(tl.int32), axis=1) > 0, tl.sum(tl.where(at, top[None, :], 0), axis=1), cand) + v = tl.where(e[None, :] == top[:, None], float("-inf"), v) + valid = j < ROWS * K + earlier = (cand[:, None] == cand[None, :]) & (j[None, :] < j[:, None]) & valid[None, :] + first = valid & (tl.sum(earlier.to(tl.int32), axis=1) == 0) + held = tl.load(resident_ptr + cand, mask=first, other=-1) >= 0 + pick = first & (~held) + rank = tl.cumsum(pick.to(tl.int32), axis=0) - 1 + kept = first & held & (tl.cumsum(first.to(tl.int32), axis=0) <= budget) + krank = tl.cumsum(kept.to(tl.int32), axis=0) - 1 + p = tl.arange(0, W_P) + m_sel = pick[None, :] & (rank[None, :] == p[:, None]) & (p < budget)[:, None] + m_res = kept[None, :] & (krank[None, :] == p[:, None]) + # at most one candidate matches each output position + sel = tl.sum(tl.where(m_sel, cand[None, :] + 1, 0), axis=1) - 1 + res = tl.sum(tl.where(m_res, cand[None, :] + 1, 0), axis=1) - 1 + tl.store(sel_ptr + p, sel, mask=p < W) + tl.store(res_ptr + p, res, mask=p < W) + + +@triton.jit(do_not_specialize=["id_base"]) +def _prefetch_count_kernel( + slots_ptr, id_of_slot_ptr, misses_ptr, sel_ptr, res_ptr, counters_ptr, id_base, + N: tl.constexpr, N_P: tl.constexpr, W: tl.constexpr, W_P: tl.constexpr, ROWS: tl.constexpr, +): + i = tl.arange(0, N_P) + slot = tl.load(slots_ptr + i, mask=i < N, other=-1) + # ensure rewrote the routed ids to slots, all pinned to this layer's ids until the next ensure + routed = tl.where(slot >= 0, tl.load(id_of_slot_ptr + slot, mask=slot >= 0, other=0) - id_base, -1) + p = tl.arange(0, W_P) + sel = tl.load(sel_ptr + p, mask=p < W, other=-1) + res = tl.load(res_ptr + p, mask=p < W, other=-1) + sel_hit = tl.sum(((sel[:, None] == routed[None, :]) & (i < N)[None, :]).to(tl.int32), axis=1) > 0 + res_hit = tl.sum(((res[:, None] == routed[None, :]) & (i < N)[None, :]).to(tl.int32), axis=1) > 0 + issued = tl.sum((sel >= 0).to(tl.int32)) + useful = tl.sum((sel_hit & (sel >= 0)).to(tl.int32)) + res_hits = tl.sum((res_hit & (res >= 0)).to(tl.int32)) + misses = tl.load(misses_ptr) + c = tl.arange(0, 8) + add = tl.where(c == 0, issued, tl.where(c == 1, useful, tl.where(c == 2, res_hits, tl.where(c == 3, 1, ROWS)))) + add = tl.where(c == 5, misses, add.to(tl.int64)) + # one writer per layer row, stream-ordered: no atomics needed + old = tl.load(counters_ptr + c, mask=c < 6, other=0) + tl.store(counters_ptr + c, old + add, mask=c < 6) + + +def lookahead_select( + logits: torch.Tensor, resident: torch.Tensor, sel: torch.Tensor, res: torch.Tensor, *, k: int, budget: int +) -> None: + """Top-``k`` of each ``logits`` row, merged rank by rank; ``sel`` gets the first ``budget`` + distinct ids with ``resident[id] < 0`` and ``res`` the resident ids among the first ``budget`` + distinct ones, both in merge order and padded with -1. One CTA, no host sync.""" + rows, num_experts = logits.shape + width = sel.numel() + assert logits.stride(1) == 1 and resident.numel() == num_experts and resident.dtype == torch.int32 + assert sel.dtype == res.dtype == torch.int32 and res.numel() == width and 1 <= k <= num_experts + _lookahead_select_kernel[(1,)]( + logits, logits.stride(0), resident, sel, res, min(budget, width), + ROWS=rows, ROWS_P=triton.next_power_of_2(rows), + E=num_experts, E_P=triton.next_power_of_2(num_experts), + K=k, N_P=triton.next_power_of_2(rows * k), + W=width, W_P=triton.next_power_of_2(width), + num_warps=4, + ) + + +def prefetch_count( + slots: torch.Tensor, id_of_slot: torch.Tensor, misses: torch.Tensor, + sel: torch.Tensor, res: torch.Tensor, counters: torch.Tensor, *, id_base: int, rows: int, +) -> None: + """Add one call's issued / useful / resident-hit / call / row / miss counts to ``counters``.""" + assert slots.dtype == torch.int32 and slots.is_contiguous() and counters.dtype == torch.int64 + n, width = slots.numel(), sel.numel() + _prefetch_count_kernel[(1,)]( + slots, id_of_slot, misses, sel, res, counters, id_base, + N=n, N_P=triton.next_power_of_2(n), W=width, W_P=triton.next_power_of_2(width), ROWS=rows, + num_warps=1, + ) + + +@triton.jit +def _mark_ready_kernel(ready_ptr): + tl.store(ready_ptr, 1) + + +class ExpertPrefetcher: + """Predictor stream, per-layer buffers and device counters for the router lookahead. + + ``fork`` runs layer L's prediction on the dedicated stream at the start of layer L-1's MoE + block. ``measure``: ``join_and_count`` waits for it after layer L's copy is enqueued and counts + on the compute stream. ``on``: ``install`` (right after layer L-1's ensure) waits for it and + runs prefetch_ensure(L) on the compute stream; ``issue_copy`` puts its copy on a second + dedicated stream after layer L-1's demand copy; the compute stream joins the copy before its + copy(L) (``join_copy``) and ensure(L) counts the call (``count_args``). The events are created + up front, the shapes are fixed and nothing syncs the host, so every side captures into the + decode graph. Nothing here is cache_size-shaped: a rebuild keeps it and only resets the counters. + """ + + def __init__( + self, num_layers: int, num_experts: int, device: torch.device, *, + mode: str = "measure", k: int | None = None, budget: int | None = None, + ) -> None: + self.mode = mode + self.num_layers = num_layers + self.num_experts = num_experts + self.device = device + k = ENV.MOE_PREFETCH_K.value if k is None else int(k) + if k < 1: + raise ValueError(f"FREETOKEN_MOE_PREFETCH_K={k} must be >= 1") + self.k = min(k, num_experts) + # 0 = the per-layer default the model wires in + self.budget_override = ENV.MOE_PREFETCH_BUDGET.value if budget is None else int(budget) + if self.budget_override < 0: + raise ValueError(f"FREETOKEN_MOE_PREFETCH_BUDGET={self.budget_override} must be >= 0") + self.stream = dedicated_stream(device) + # (forked, selected): the predictor starts after forked and records selected after the select + self._events = [(torch.cuda.Event(), torch.cuda.Event()) for _ in range(num_layers)] + # torch creates the CUDA event on its first record: do that now, never inside a capture + for pair in self._events: + for event in pair: + event.record(self.stream) + # on mode's select sees nothing resident: prefetch_ensure applies residency and the budget + self._none_resident = torch.full((num_experts,), -1, dtype=torch.int32, device=device) + self._budget = [0] * num_layers + width = 2 * self.k # room for the whole bs-2 union + self.sel = torch.full((num_layers, width), -1, dtype=torch.int32, device=device) + self.res = torch.full((num_layers, width), -1, dtype=torch.int32, device=device) + self.stats = torch.zeros((num_layers, NUM_STAT_COLS), dtype=torch.int64, device=device) + # measure mode's six columns, a view: its count kernel and readers keep the stage-A layout + self.counters = self.stats[:, :NUM_COLS] + self.totals = torch.zeros((num_layers, NUM_STAT_COLS), dtype=torch.int64) + # the router input of each forked prediction, alive until the compute stream has waited for + # the select so the allocator cannot hand its block to later work the predictor still reads + self._inflight: list[torch.Tensor | None] = [None] * num_layers + # on mode: whether install ran prefetch_ensure for the forked prediction (its copy is then due) + self._installed = [False] * num_layers + # on mode: the cache whose slot maps prefetch_ensure writes and whose banks the copy reads + self.cache = None + self.copy_stream = None + if mode == "on": + # the slim copy is an SM kernel: at default priority a late copy queues for SMs behind the next + # compute grid while the compute stream waits on it (bench --on: a 50 us slip cost 14% per step) + self.copy_stream = dedicated_stream(device, highest_priority=True) + # (dcopy: after prefetch_ensure(L) and layer L-1's demand copy on the compute stream, + # pfcopy: after layer L's prefetch copy) + self._copy_events = [(torch.cuda.Event(), torch.cuda.Event()) for _ in range(num_layers)] + for pair in self._copy_events: + for event in pair: + event.record(self.copy_stream) + # one plan per target layer: prefetch_ensure(L+1) must not overwrite the plan copy(L) reads + self.pf_slots = torch.full((num_layers, width), -1, dtype=torch.int32, device=device) + self.pf_src = torch.full((num_layers, width), -1, dtype=torch.int32, device=device) + self.pf_num = torch.zeros((num_layers, 1), dtype=torch.int64, device=device) + self.pf_ready = torch.zeros((num_layers, 1), dtype=torch.int32, device=device) + # test hooks, read at enqueue (capture) time: a [num_layers, width] device buffer replacing + # sel (on mode still installs at most the budget), and sleeps put on the predictor / copy + # stream ahead of their work + self.sel_override: torch.Tensor | None = None + self.delay_predict_cycles = 0 + self.delay_copy_cycles = 0 + logger.info( + f"MoE prefetch {mode}: router lookahead top-{self.k}, budget " + + (str(self.budget_override) if self.budget_override else f"{GDN_BUDGET} before GDN / {ATTN_BUDGET} before attention layers") + + (", slim copies on a dedicated high-priority stream" if mode == "on" else "") + ) + + def budget_for(self, wired: int) -> int: + return min(self.budget_override or wired, self.sel.shape[1]) + + def fork(self, target: int, x: torch.Tensor, gate, resident: torch.Tensor, budget: int) -> bool: + """Predict ``target``'s experts from ``x`` (the previous layer's router input) on the + predictor stream, ordered after the work already on the current stream. Measure mode skips + the ids ``resident`` holds (a racy read of the cache); on mode reads no cache state and + leaves residency and the budget to ``install``. Returns whether it forked.""" + if x.shape[0] > MAX_ROWS: + return False + budget = self.budget_for(budget) + if self.mode == "on": + self._budget[target] = budget + resident, budget = self._none_resident, self.sel.shape[1] + forked, selected = self._events[target] + forked.record(torch.cuda.current_stream(self.device)) + self.stream.wait_event(forked) + with torch.cuda.stream(self.stream): + if self.delay_predict_cycles: + torch.cuda._sleep(self.delay_predict_cycles) + logits = gate.forward(x) + lookahead_select(logits, resident, self.sel[target], self.res[target], k=self.k, budget=budget) + if self.sel_override is not None: + self.sel[target].copy_(self.sel_override[target]) + selected.record(self.stream) + self._inflight[target] = x + self._installed[target] = False + return True + + def install(self, target: int | None) -> int | None: + """On mode, right after the previous layer's ensure: wait for ``target``'s prediction and + install its first budget non-resident candidates with prefetch_ensure on the current + (compute) stream. Returns ``target`` when its copy is due (``issue_copy``), else None.""" + if target is None or self.mode != "on" or self._inflight[target] is None: + return None + torch.cuda.current_stream(self.device).wait_event(self._events[target][1]) + self.cache.prefetch_ensure( + target, self.sel[target], self.pf_slots[target], self.pf_src[target], + self.pf_num[target], self.stats[target], self.pf_ready[target], budget=self._budget[target], + ) + self._installed[target] = True + return target + + def _pending(self, layer_id: int) -> bool: + """On mode: ``layer_id``'s prediction was installed and its copy is not joined yet.""" + return self.mode == "on" and self._installed[layer_id] and self._inflight[layer_id] is not None + + def count_args(self, layer_id: int) -> tuple | None: + """On mode: what ``layer_id``'s demand ensure counts into, or None when it was not prefetched.""" + if not self._pending(layer_id): + return None + return self.stats[layer_id], self.pf_ready[layer_id], self.pf_num[layer_id], self._inflight[layer_id].shape[0] + + def plan_of(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor] | None: + """On mode: ``layer_id``'s prefetch copy plan (slots, count) while its prefetch is in flight.""" + if not self._pending(layer_id): + return None + return self.pf_slots[layer_id], self.pf_num[layer_id] + + def join_copy(self, layer_id: int) -> None: + """On mode: the compute stream waits for ``layer_id``'s prefetch copy before its own copy, so + a prefetched slot ensure(L) picked as a demand victim is never written by two copies.""" + if self.mode != "on" or self._inflight[layer_id] is None: + return + # a prediction that was never installed still has to rejoin (a capture needs every fork joined) + done = self._copy_events[layer_id][1] if self._pending(layer_id) else self._events[layer_id][1] + torch.cuda.current_stream(self.device).wait_event(done) + self._inflight[layer_id] = None + self._installed[layer_id] = False + + def issue_copy(self, target: int) -> None: + """On mode: copy ``target``'s prefetch plan on the copy stream after the work already on the + current stream: its prefetch_ensure and the demand copy (which owns the host link first).""" + if not self._pending(target): + return + dcopy, pfcopy = self._copy_events[target] + dcopy.record(torch.cuda.current_stream(self.device)) + self.copy_stream.wait_event(dcopy) + with torch.cuda.stream(self.copy_stream): + if self.delay_copy_cycles: + torch.cuda._sleep(self.delay_copy_cycles) + self.cache.copy_rows(target, self.pf_slots[target], self.pf_src[target], self.pf_num[target], slim=True, + kind=PREFETCH_COPY) + _mark_ready_kernel[(1,)](self.pf_ready[target]) + pfcopy.record(self.copy_stream) + + def join_and_count(self, layer_id: int, slots: torch.Tensor, id_of_slot: torch.Tensor, misses: torch.Tensor) -> None: + """Measure mode: join ``layer_id``'s prediction (if one was forked this step) and count it + against the routed slots ``slots`` of this layer's ensure.""" + x = self._inflight[layer_id] if self.mode == "measure" else None + if x is None: + return + torch.cuda.current_stream(self.device).wait_event(self._events[layer_id][1]) + prefetch_count( + slots.view(-1), id_of_slot, misses, self.sel[layer_id], self.res[layer_id], self.counters[layer_id], + id_base=layer_id * self.num_experts, rows=x.shape[0], + ) + self._inflight[layer_id] = None + + def reset_counters(self) -> None: + self.stats.zero_() + + def reset(self) -> None: + """Counters and in-flight state; the streams, events and buffers stay.""" + self.stats.zero_() + self.totals.zero_() + self._inflight = [None] * self.num_layers + self._installed = [False] * self.num_layers + + def take_window(self) -> torch.Tensor: + """This window's ``[num_layers, NUM_STAT_COLS]`` stats on the host (one sync), added to + ``totals``; the device ones restart.""" + window = self.stats.cpu() + self.stats.zero_() + self.totals += window + return window + + +def summarize(counters: torch.Tensor, mode: str = "measure") -> dict: + """Aggregate ``[num_layers, NUM_COLS or NUM_STAT_COLS]`` counters into per-layer-call and + per-token rates. In on mode the misses are what was left after the prefetch, so coverage is + useful / (useful + misses), and late is the share of copied calls whose copy had not landed + when the compute stream reached it.""" + total = counters.sum(0).tolist() + calls, issued, useful = total[CALLS], total[ISSUED], total[USEFUL] + tokens = int(counters[:, ROWS].max()) if counters.numel() else 0 + wanted = useful + total[MISSES] if mode == "on" else total[MISSES] + stats = { + "layer_calls": calls, + "tokens": tokens, + "issued_per_layer": issued / calls if calls else 0.0, + "useful_per_layer": useful / calls if calls else 0.0, + "resident_hits_per_layer": total[RESIDENT_HITS] / calls if calls else 0.0, + "misses_per_layer": total[MISSES] / calls if calls else 0.0, + "precision": useful / issued if issued else 0.0, + "coverage": useful / wanted if wanted else 0.0, + "useful_per_token": useful / tokens if tokens else 0.0, + } + if mode == "on": + copied = total[COPIED] if len(total) > COPIED else 0 + stats["copied_calls"] = copied + stats["late_fraction"] = total[LATE] / copied if copied else 0.0 + return stats + + +def format_summary(stats: dict) -> str: + if "late_fraction" in stats: + return ( + f"issued/layer={stats['issued_per_layer']:.2f}, useful/layer={stats['useful_per_layer']:.2f}, " + f"precision={stats['precision']:.3f}, useful/token={stats['useful_per_token']:.1f}, " + f"misses/layer={stats['misses_per_layer']:.2f}, coverage={stats['coverage']:.3f}, " + f"late={stats['late_fraction']:.3f} of {stats['copied_calls']} copies" + ) + return ( + f"issued/layer={stats['issued_per_layer']:.2f}, useful/layer={stats['useful_per_layer']:.2f}, " + f"precision={stats['precision']:.3f}, useful/token={stats['useful_per_token']:.1f}, " + f"misses/layer={stats['misses_per_layer']:.2f}, coverage={stats['coverage']:.3f}, " + f"resident_hits/layer={stats['resident_hits_per_layer']:.2f}" + ) + + +def format_per_layer(counters: torch.Tensor) -> str: + parts = [] + for layer, row in enumerate(counters.tolist()): + if row[CALLS]: + c = row[CALLS] + parts.append(f"L{layer}={row[USEFUL] / c:.2f}/{row[ISSUED] / c:.2f}/{row[MISSES] / c:.2f}") + return ", ".join(parts) diff --git a/python/freetoken/moe/scored_ensure.py b/python/freetoken/moe/scored_ensure.py index 123086ed51..ba570b5169 100644 --- a/python/freetoken/moe/scored_ensure.py +++ b/python/freetoken/moe/scored_ensure.py @@ -106,6 +106,22 @@ def _phase1(query_ptr, slot_of_id_ptr, lru_usage_ptr, num_copy_ptr, step, K, return q, kmask, miss, first, first_miss, rank, num_missing, tl.where(hit, s, -1) +@triton.jit +def _prefetch_phase1(query_ptr, slot_of_id_ptr, K, BLOCK_K: tl.constexpr, id_base): + """_phase1 for a prefetch: -1 entries are padding, hits are left alone (a prefetch is not an + access), and the misses rank in query order so the select's best candidates install first.""" + k = tl.arange(0, BLOCK_K) + raw = tl.load(query_ptr + k, mask=k < K, other=-1) + kmask = raw >= 0 + q = tl.where(kmask, raw + id_base, 0) + s = tl.load(slot_of_id_ptr + q, mask=kmask, other=0) + miss = kmask & (s == -1) + same = (q[:, None] == q[None, :]) & (k[:, None] > k[None, :]) & kmask[:, None] & kmask[None, :] + first_miss = miss & (tl.sum(same.to(tl.int32), axis=1) == 0) + rank = tl.cumsum(first_miss.to(tl.int32), axis=0) - 1 + return q, first_miss, rank, tl.sum(first_miss.to(tl.int32)) + + @triton.jit def _stats(stats_ptr, first, num_missing): # One vectorized atomic over 3 lanes, not three scalar ones: a scalar atomic serializes the CTA. @@ -144,13 +160,14 @@ def _owner_state(mirror_ptr, state_ptr, c, oid, mirrored, stale, other): ) -@triton.jit(do_not_specialize=["K", "num_cached", "id_base", "nm_topk", "nm_stride"]) +@triton.jit(do_not_specialize=["K", "num_cached", "id_base", "nm_topk", "nm_stride", "pf_rows", "pf_budget"]) def _scored_ensure_kernel( query_ptr, slot_of_id_ptr, id_of_slot_ptr, lru_usage_ptr, lru_step_ptr, out_ptr, src_ptr, dst_ptr, num_copy_ptr, stats_ptr, tok_ptr, last_tok_ptr, lc_ptr, ct_ptr, g_ptr, logits_ptr, owner_ptr, slot_last_ptr, slot_lc_ptr, slot_ct_ptr, pin_ptr, - K, num_cached, id_base, nm_topk, nm_stride, nm_thr, + pf_stats_ptr, pf_ready_ptr, pf_num_ptr, + K, num_cached, id_base, nm_topk, nm_stride, nm_thr, pf_rows, pf_budget, BLOCK_K: tl.constexpr, BLOCK_C: tl.constexpr, BLOCK_E: tl.constexpr, BLOCK_TOPK: tl.constexpr, NM_ROWS: tl.constexpr, USAGE_MAX: tl.constexpr, COLLECT_STATS: tl.constexpr, @@ -158,15 +175,27 @@ def _scored_ensure_kernel( NUM_LAYERS: tl.constexpr, NUM_EXPERTS: tl.constexpr, BETA_STEP: tl.constexpr, W_Q4: tl.constexpr, DECAY_MUL: tl.constexpr, DT_MAX: tl.constexpr, SLOT_BITS: tl.constexpr, + LOWPRI: tl.constexpr, PF_COUNT: tl.constexpr, PREFETCH: tl.constexpr, ): - step = tl.load(lru_step_ptr) + 1 - tl.store(lru_step_ptr, step) + if PREFETCH: + # not an access: the clock stays, so usage == step still pins the last demand call's slots + # (at step 0, right after a reset, every slot reads pinned: nothing installs before a demand call) + step = tl.load(lru_step_ptr) + tl.store(num_copy_ptr, tl.zeros([], tl.int64)) + tl.store(pf_ready_ptr, 0) + else: + step = tl.load(lru_step_ptr) + 1 + tl.store(lru_step_ptr, step) if POLICY != 0: tok = tl.load(tok_ptr) if BUMP_TOK: tok = tok + 1 tl.store(tok_ptr, tok) layer = id_base // NUM_EXPERTS + if PREFETCH: + # score from where the pass is, right after layer - 1's demand call: the target layer's + # residents are next up, not a whole pass away (else it evicts what the layer is about to route) + layer = layer - 1 now = tok * NUM_LAYERS + layer if POLICY == 3: if UPDATE_STATE: @@ -187,9 +216,18 @@ def _scored_ensure_kernel( tl.store(slot_last_ptr + held_at, refreshed, mask=near & (held_at >= 0)) # the victim scan and the routed update below read and rewrite what this wrote tl.debug_barrier() - q, kmask, miss, first, first_miss, rank, num_missing, out = _phase1( - query_ptr, slot_of_id_ptr, lru_usage_ptr, num_copy_ptr, step, K, BLOCK_K, - id_base) + if PREFETCH: + q, first_miss, rank, num_missing = _prefetch_phase1(query_ptr, slot_of_id_ptr, K, BLOCK_K, id_base) + else: + if PF_COUNT: + kk = tl.arange(0, BLOCK_K) + sk = tl.load(slot_of_id_ptr + tl.load(query_ptr + kk, mask=kk < K, other=0) + id_base, mask=kk < K, other=-1) + # a hit on a held slot with usage 0 consumes a prefetch; read it before _phase1 bumps it + was_pf = (sk >= 0) & (tl.load(lru_usage_ptr + sk, mask=sk >= 0, other=1) == 0) + tl.debug_barrier() + q, kmask, miss, first, first_miss, rank, num_missing, out = _phase1( + query_ptr, slot_of_id_ptr, lru_usage_ptr, num_copy_ptr, step, K, BLOCK_K, + id_base) if num_missing > 0: # REQUIRED (flashlib): the reload below must see _phase1's hit bump, or a hit is evicted. @@ -199,7 +237,13 @@ def _scored_ensure_kernel( u = tl.load(lru_usage_ptr + c, mask=cmask, other=USAGE_MAX) if POLICY == 0: umax = tl.full([BLOCK_C], USAGE_MAX, u.dtype) - u = tl.where((u == step) | (~cmask), umax, u) + if LOWPRI: + held0 = cmask & (tl.load(id_of_slot_ptr + c, mask=cmask, other=-1) >= 0) + # empty slots, then unconsumed prefetches (held, usage 0), then residents by recency + u = tl.where((u == step) | (~cmask), umax, tl.where(held0, u + 1, 0)) + else: + u = tl.where((u == step) | (~cmask), umax, u) + n_free = tl.sum((u != umax).to(tl.int32)) else: oid = tl.load(id_of_slot_ptr + c, mask=cmask, other=-1) held = cmask & (oid >= 0) @@ -221,6 +265,10 @@ def _scored_ensure_kernel( score += (W_Q4 * lcv) >> 4 # Every key is distinct and ties go to the lowest slot; an empty slot keys below any held one. key = tl.where(held, ((score.to(tl.int64) + _SCORE_BIAS) << SLOT_BITS) | c, c.to(tl.int64)) + if LOWPRI: + # an unconsumed prefetch (held, usage 0) keys above every empty slot and below every + # resident: the caps keep score + 2^31 > 2^30 - 2^25 - 2^22, far above this band's 1 + key = tl.where(held & (u == 0), (1 << SLOT_BITS) | c.to(tl.int64), key) evictable = cmask & (u != step) if PIN_SINCE: # the score ignores recency, so slots an earlier call touched since *pin_ptr stay pinned too @@ -228,7 +276,16 @@ def _scored_ensure_kernel( # the pinned key's slot bits are 0, so an all-pinned scan falls back to slot 0 like flashlib's argmin pinned_key = 0x7FFFFFFFFFFFFFFF ^ ((1 << SLOT_BITS) - 1) key = tl.where(evictable, key, pinned_key) - for i in tl.range(num_missing): + n_free = tl.sum(evictable.to(tl.int32)) + if PREFETCH: + # never the all-pinned fallback: that slot may be under the running GEMM; misses rank in + # query order, so the budget keeps the first non-resident candidates + n_iter = tl.minimum(tl.minimum(num_missing, n_free), pf_budget) + tl.store(num_copy_ptr, n_iter.to(tl.int64)) + tl.store(pf_stats_ptr, tl.load(pf_stats_ptr) + n_iter) + else: + n_iter = num_missing + for i in tl.range(n_iter): if POLICY == 0: victim = tl.argmin(u, axis=0).to(tl.int32) else: @@ -240,17 +297,38 @@ def _scored_ensure_kernel( e = tl.sum(tl.where((rank == i) & first_miss, q, 0)) tl.store(id_of_slot_ptr + victim, e) tl.store(slot_of_id_ptr + e, victim) - tl.store(lru_usage_ptr + victim, step) + if PREFETCH: + # usage 0 on a held slot marks it low priority until a demand call touches it + tl.store(lru_usage_ptr + victim, 0) + if POLICY != 0: + tl.store(owner_ptr + victim, -1) + else: + tl.store(lru_usage_ptr + victim, step) tl.store(dst_ptr + i, victim) tl.store(src_ptr + i, e - id_base) # back to the caller's id space - out = tl.where((rank == i) & miss, victim, out) + if not PREFETCH: + out = tl.where((rank == i) & miss, victim, out) if POLICY == 0: u = tl.where(c == victim, umax, u) # claim in-register else: key = tl.where(c == victim, pinned_key, key) - # Written from registers, never re-read from slot_of_id, so out_ptr may alias query_ptr. - tl.store(out_ptr + tl.arange(0, BLOCK_K), out, mask=kmask) + if not PREFETCH: + # Written from registers, never re-read from slot_of_id, so out_ptr may alias query_ptr. + tl.store(out_ptr + tl.arange(0, BLOCK_K), out, mask=kmask) + if PF_COUNT: + # one writer per column: pf_ensure adds ISSUED (column 0), earlier on this stream + pn = tl.load(pf_num_ptr) + late = (pn > 0) & (tl.load(pf_ready_ptr, volatile=True) == 0) + col = tl.arange(0, 8) + add = tl.where(col == 1, tl.sum((first & was_pf).to(tl.int32)).to(tl.int64), 0) + add = tl.where(col == 3, 1, add) + add = tl.where(col == 4, pf_rows, add) + add = tl.where(col == 5, num_missing.to(tl.int64), add) + add = tl.where(col == 6, late.to(tl.int64), add) + add = tl.where(col == 7, (pn > 0).to(tl.int64), add) + cols = (col == 1) | (col >= 3) + tl.store(pf_stats_ptr + col, tl.load(pf_stats_ptr + col, mask=cols, other=0) + add, mask=cols) if POLICY != 0: if UPDATE_STATE: # distinct ids only, so a duplicate at bs > 1 counts once per call; out is each id's final slot @@ -265,7 +343,7 @@ def _scored_ensure_kernel( tl.store(ct_ptr + q, now, mask=first) tl.store(slot_lc_ptr + out, lc_new, mask=first) tl.store(slot_ct_ptr + out, now, mask=first) - else: + elif not PREFETCH: # a slot filled without a state update reads the per-id state until a decode hit re-mirrors it tl.store(owner_ptr + out, -1, mask=first_miss) if COLLECT_STATS: @@ -311,6 +389,13 @@ def scored_ensure( beta: float = BETA, w: float = W, halflife: int = HALFLIFE, + lowpri: bool = False, + pf_stats: torch.Tensor | None = None, + pf_ready: torch.Tensor | None = None, + pf_num: torch.Tensor | None = None, + pf_rows: int = 0, + pf_budget: int | None = None, + prefetch: bool = False, ) -> None: """``flashlib.lru_ensure`` (sequential strategy) with the victim picked by ``policy``. @@ -324,7 +409,21 @@ def scored_ensure( prefill's chunked calls cannot evict each other's experts; policy > 0 only. Policy 3 refreshes the experts whose ``router_logits`` (``[rows, num_experts]``, rows of ``query``) are within ``near_miss_thr`` of their row's lowest routed logit. + + FREETOKEN_MOE_PREFETCH=on: ``lowpri`` keys held slots with usage 0 (unconsumed prefetches) + between the empty slots and every resident. ``pf_stats`` (a ``[8]`` int64 prefetch stats row) + also counts this demand call: hits on such slots, the call, ``pf_rows``, the misses, and whether + the prefetch plan ``pf_num`` had rows whose copy had not set ``pf_ready`` yet. + ``prefetch=True`` is :func:`prefetch_ensure`'s install mode instead, installing at most + ``pf_budget`` ids (default: all). """ + if prefetch: + assert lowpri and pf_stats is not None and pf_ready is not None and stats is None + assert not update_state and not bump_tok and pin_since is None and router_logits is None + if pf_stats is not None: + assert lowpri and pf_stats.dtype == torch.int64 and pf_stats.numel() == 8 and pf_stats.is_contiguous() + assert pf_ready is not None and pf_ready.dtype == torch.int32 and pf_ready.numel() == 1 + assert prefetch or (pf_num is not None and pf_num.dtype == torch.int64 and pf_num.numel() == 1) k = query.numel() num_cached = id_of_slot.numel() assert query.dtype == torch.int32 and query.is_contiguous() @@ -372,7 +471,11 @@ def scored_ensure( dummy if slot_lc is None else slot_lc, dummy if slot_ct is None else slot_ct, dummy if pin_since is None else pin_since, - k, num_cached, id_base, nm_topk, nm_stride, float(near_miss_thr), + dummy if pf_stats is None else pf_stats, + dummy if pf_ready is None else pf_ready, + dummy if pf_num is None else pf_num, + k, num_cached, id_base, nm_topk, nm_stride, float(near_miss_thr), int(pf_rows), + k if pf_budget is None else int(pf_budget), BLOCK_K=triton.next_power_of_2(k), BLOCK_C=block_c, BLOCK_E=triton.next_power_of_2(num_experts), @@ -391,5 +494,40 @@ def scored_ensure( DECAY_MUL=decay_mul, DT_MAX=dt_max, SLOT_BITS=slot_bits, + LOWPRI=bool(lowpri), + PF_COUNT=pf_stats is not None and not prefetch, + PREFETCH=bool(prefetch), num_warps=_num_warps_for(block_c), ) + + +def prefetch_ensure( + query: torch.Tensor, + slot_of_id: torch.Tensor, + id_of_slot: torch.Tensor, + lru_usage: torch.Tensor, + lru_step: torch.Tensor, + dst_slots: torch.Tensor, + src_rows: torch.Tensor, + num_copy: torch.Tensor, + pf_stats: torch.Tensor, + pf_ready: torch.Tensor, + budget: int | None = None, + **state, +) -> None: + """Install the ids of ``query`` (-1 = padding) that are not resident as low-priority slots + (held, usage 0), at most the first ``budget`` of them in query order (default: all), and + write their copy plan to ``dst_slots``/``src_rows``/``num_copy``. + + Unlike a demand call it bumps no clock, touches no hit and no per-id state, and so keeps + every slot of the last demand call (usage == step) pinned; victims follow the demand key + (empty, then earlier unconsumed prefetches, then the coldest residents), scored as of the + previous layer's call, the position the prefetch runs at. With fewer + evictable slots than misses it installs only the first ones, in query order. Adds the + installs to ``pf_stats[0]`` and clears ``pf_ready``. + """ + scored_ensure( + query, slot_of_id, id_of_slot, lru_usage, lru_step, query, src_rows, dst_slots, num_copy, + update_state=False, lowpri=True, pf_stats=pf_stats, pf_ready=pf_ready, pf_budget=budget, prefetch=True, + **state, + ) diff --git a/python/freetoken/moe/slot_audit.py b/python/freetoken/moe/slot_audit.py new file mode 100644 index 0000000000..981239758f --- /dev/null +++ b/python/freetoken/moe/slot_audit.py @@ -0,0 +1,688 @@ +"""Expert slot audit (``FREETOKEN_MOE_SLOT_AUDIT=N``): a debug detector for PERSISTENT slot corruption. + +Two parts, both on the device and off by default: + +- writer history: every writer of the slot maps (demand and prefetch ensures, prefill ensures, the + whole-layer materialize, prefill buffer invalidation, reset) and every writer of the slot banks + (demand, prefetch and prefill copies, the prefill double buffers) is followed on its writer's + stream by a tiny kernel keyed on the same plan (or slot range) that appends ``(seq, kind, decode + step, lru step, layer, expert row, ...)`` to a per-slot ring of the last ``DEPTH`` metadata events + and one of the last ``DEPTH`` byte events. The map writers all run on the compute stream except the + prefill buffer invalidation (on the prefill copy stream, fenced behind it); the byte writers also + run on the prefetch and prefill copy streams. ``seq`` is one global atomic counter, so the two rings + merge into one cross-stream order. Clears are recorded just before their writer, so they carry the + old owner. + Fixed shapes and no host sync: the recorders capture into the decode graph. +- audit: every N decode steps, after the forward on the compute stream, every held slot is checked: + its bytes in every bank against the host row of the (layer, expert) its ``id_of_slot`` names + (host rows gathered with the fused copy into a scratch of fewer than 64 MiB, one layer chunk at a + time), ``slot_for_id[id_of_slot[s]] == s``, and, over every id, ``id_of_slot[slot_for_id[id]] == id`` + (a stale entry routes that id to another expert's bytes). The first ``RING`` bad (slot, owner) + pairs are kept with both history rings; the engine logs them and the totals every + MOE_STATS_INTERVAL decode steps and at shutdown, with a diagnosis read from the history. +""" +from __future__ import annotations + +import math + +import torch +import triton +import triton.language as tl + +from freetoken.env import ENV +from freetoken.utils import init_logger + +logger = init_logger(__name__) + +# events kept per slot and ring (metadata / bytes) +DEPTH = 4 +# event fields +(E_SEQ, E_KIND, E_DSTEP, E_LSTEP, E_LAYER, E_ROW, E_AUX, E_OWNER, E_USAGE, E_POS) = range(10) +EVENT_FIELDS = 10 +# metadata kinds (installs name the new owner; clears the old one, recorded before the writer) +DEMAND_INSTALL, PREFETCH_INSTALL, PREFILL_INSTALL, MATERIALIZE_INSTALL = 1, 2, 3, 4 +MATERIALIZE_CLEAR, INVALIDATE, RESET = 5, 6, 7 +# byte kinds +DEMAND_COPY, PREFETCH_COPY, PREFILL_COPY, MATERIALIZE_COPY, PAGEABLE_COPY = 11, 12, 13, 14, 15 +PREFILL_BUFFER, PREFILL_SPLIT_H2D, PREFILL_HIT_D2D, OTHER_COPY = 16, 17, 18, 19 +KIND_NAMES = { + DEMAND_INSTALL: "demand install", PREFETCH_INSTALL: "prefetch install", PREFILL_INSTALL: "prefill ensure install", + MATERIALIZE_INSTALL: "materialize install", MATERIALIZE_CLEAR: "materialize clear", + INVALIDATE: "prefill buffer invalidate", RESET: "reset", + DEMAND_COPY: "demand copy", PREFETCH_COPY: "prefetch copy", PREFILL_COPY: "prefill ensure copy", + MATERIALIZE_COPY: "materialize copy", PAGEABLE_COPY: "pageable materialize copy", + PREFILL_BUFFER: "prefill buffer copy", PREFILL_SPLIT_H2D: "prefill buffer H2D (misses and small banks)", + PREFILL_HIT_D2D: "prefill buffer D2D", OTHER_COPY: "copy", +} +INSTALLS = (DEMAND_INSTALL, PREFETCH_INSTALL, PREFILL_INSTALL, MATERIALIZE_INSTALL) +# the copy each install expects next +COPY_FOR = { + DEMAND_INSTALL: DEMAND_COPY, PREFETCH_INSTALL: PREFETCH_COPY, PREFILL_INSTALL: PREFILL_COPY, + MATERIALIZE_INSTALL: MATERIALIZE_COPY, +} +# recorder row sources: the plan's src rows, slot - row_base, the range index, the old owner +ROW_PLAN, ROW_FROM_SLOT, ROW_INDEX, ROW_OLD_OWNER = range(4) +# recorder predicates: every slot, held slots, slots >= E held by the layer +PRED_ALL, PRED_HELD, PRED_LAYER_ABOVE_E = range(3) +# totals: audits, held slots seen, slots byte-checked, held slots not byte-checked, bad slots found +# (summed over audits), of them with bad bytes / a forward map miss / named by a stale slot_for_id / +# an owner out of range, slot_for_id entries past the cache, new records, repeat findings, events +# recorded, plan entries naming no slot +(T_AUDITS, T_HELD, T_BYTES_CHECKED, T_UNCHECKED, T_BAD, T_BYTES, T_FWD, T_STALE, T_OWNER_RANGE, T_MAP_RANGE, + T_NEW, T_STILL, T_EVENTS, T_BAD_PLAN) = range(14) +NUM_TOTALS = 16 +# per-slot audit flags +F_BYTES, F_FWD, F_STALE, F_OWNER_RANGE = 1, 2, 4, 8 +# record header, then the metadata ring and the byte ring as stored ([DEPTH, EVENT_FIELDS] each) +(H_AUDIT, H_DSTEP, H_LSTEP, H_SLOT, H_OWNER, H_MAPPED, H_USAGE, H_FLAGS, H_BANKS, H_FIRST_BANK, H_FIRST_BYTE, + H_BAD_WORDS, H_STALE_ID, H_STALE_N, H_META_N, H_BYTES_N) = range(16) +HEADER = 16 +RING_EVENTS = DEPTH * EVENT_FIELDS +RECORD_FIELDS = HEADER + 2 * RING_EVENTS +RING = 32 +SCRATCH_BYTES = 64 << 20 +MAX_CHUNK_ROWS = 64 +COMPARE_WORDS = 4096 +_NO_BAD = tl.constexpr(0x7FFFFFFF) +_D = tl.constexpr(DEPTH) +_F = tl.constexpr(EVENT_FIELDS) + + +def resolve_interval(interval: int | None = None) -> int: + value = ENV.MOE_SLOT_AUDIT.value if interval is None else interval + value = int(value or 0) + if value < 0: + raise ValueError(f"FREETOKEN_MOE_SLOT_AUDIT={value}: expected 0 (off) or the decode steps between audits") + return value + + +@triton.jit +def _put(vals, f, idx: tl.constexpr, v): + return tl.where(f[None, :] == idx, v.to(tl.int64)[:, None], vals) + + +@triton.jit +def _bump_step_kernel(dstep_ptr): + tl.store(dstep_ptr, tl.load(dstep_ptr) + 1) + + +@triton.jit(do_not_specialize=["kind", "layer", "start", "row_base", "width", "num_slots"]) +def _record_kernel( + hist_ptr, cur_ptr, seq_ptr, dstep_ptr, lstep_ptr, id_of_slot_ptr, usage_ptr, totals_ptr, + dst_ptr, src_ptr, aux_ptr, num_ptr, + kind, layer, start, row_base, width, num_slots, + PLAN: tl.constexpr, ROW: tl.constexpr, HAS_AUX: tl.constexpr, HAS_NUM: tl.constexpr, PRED: tl.constexpr, + E: tl.constexpr, BLOCK: tl.constexpr, +): + """Append one event per written slot to that slot's ring: plan entries ``i < min(*num, width)`` + (``PLAN``) or the slots ``start + i`` for ``i < width``.""" + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + n = width + if HAS_NUM: + n = tl.minimum(tl.load(num_ptr).to(tl.int32), width) + valid = i < n + if PLAN: + slot = tl.load(dst_ptr + i, mask=valid, other=-1) + else: + slot = start + i + ok = valid & (slot >= 0) & (slot < num_slots) + bad_plan = valid & ~ok + owner = tl.load(id_of_slot_ptr + slot, mask=ok, other=-3) + usage = tl.load(usage_ptr + slot, mask=ok, other=-3) + if PRED == 1: + ok = ok & (owner >= 0) + if PRED == 2: + ok = ok & (owner >= layer * E) & (owner < layer * E + E) & (slot >= E) + lay = tl.zeros([BLOCK], tl.int32) + layer + aux = tl.zeros([BLOCK], tl.int32) - 1 + if ROW == 0: + row = tl.load(src_ptr + i, mask=ok, other=-1) + elif ROW == 1: + row = slot - row_base + elif ROW == 2: + row = i + else: + row = tl.where(owner >= 0, owner % E, -1) + lay = tl.where(owner >= 0, owner // E, -1) + aux = aux + 1 + layer + if HAS_AUX: + aux = tl.load(aux_ptr + i, mask=ok, other=-1) + cnt = tl.sum(ok.to(tl.int64), axis=0) + seq = tl.atomic_add(seq_ptr, cnt) + tl.cumsum(ok.to(tl.int64), axis=0) - 1 + pos = tl.atomic_add(cur_ptr + slot, 1, mask=ok) + f = tl.arange(0, 16) + z = tl.zeros([BLOCK], tl.int64) + vals = tl.zeros([BLOCK, 16], tl.int64) - 1 + vals = _put(vals, f, 0, seq) + vals = _put(vals, f, 1, z + kind) + vals = _put(vals, f, 2, z + tl.load(dstep_ptr)) + vals = _put(vals, f, 3, z + tl.load(lstep_ptr)) + vals = _put(vals, f, 4, lay) + vals = _put(vals, f, 5, row) + vals = _put(vals, f, 6, aux) + vals = _put(vals, f, 7, owner) + vals = _put(vals, f, 8, usage) + vals = _put(vals, f, 9, i) + at = (slot.to(tl.int64) * _D + (pos % _D)) * _F + tl.store(hist_ptr + at[:, None] + f[None, :], vals, mask=ok[:, None] & (f[None, :] < _F)) + t = tl.arange(0, 2) + add = tl.where(t == 0, cnt, tl.sum(bad_plan.to(tl.int64), axis=0)) + tl.atomic_add(totals_ptr + 12 + t, add, sem="relaxed") + + +@triton.jit(do_not_specialize=["num_slots", "num_ids"]) +def _audit_meta_kernel( + id_of_slot_ptr, slot_for_id_ptr, layer_ok_ptr, count_ptr, rows_ptr, slots_ptr, + flags_ptr, stale_id_ptr, stale_n_ptr, totals_ptr, num_slots, num_ids, + E: tl.constexpr, CAP: tl.constexpr, BLOCK: tl.constexpr, +): + """Both map checks, and each byte-checkable held slot appended to its layer's list.""" + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + ms = i < num_slots + o = tl.load(id_of_slot_ptr + i, mask=ms, other=-1) + held = ms & (o >= 0) + in_range = held & (o < num_ids) + mapped = tl.load(slot_for_id_ptr + o, mask=in_range, other=-2) + fwd = in_range & (mapped != i) + out_of_range = held & ~in_range + fl = tl.where(fwd, 2, 0) | tl.where(out_of_range, 8, 0) + tl.atomic_or(flags_ptr + i, fl, mask=ms & (fl != 0)) + layer = tl.where(in_range, o // E, 0) + ok = in_range & (tl.load(layer_ok_ptr + layer, mask=in_range, other=0) != 0) + p = tl.atomic_add(count_ptr + layer, 1, mask=ok) + take = ok & (p < CAP) + tl.store(rows_ptr + layer * CAP + p, o - layer * E, mask=take) + tl.store(slots_ptr + layer * CAP + p, i, mask=take) + mi = i < num_ids + s = tl.load(slot_for_id_ptr + i, mask=mi, other=-1) + live = mi & (s >= 0) + s_ok = live & (s < num_slots) + owner = tl.load(id_of_slot_ptr + s, mask=s_ok, other=-1) + stale = s_ok & (owner != i) + tl.atomic_or(flags_ptr + s, 4, mask=stale) + tl.store(stale_id_ptr + s, i, mask=stale) + tl.atomic_add(stale_n_ptr + s, 1, mask=stale) + t = tl.arange(0, 16) + add = tl.where(t == 1, tl.sum(held.to(tl.int64), axis=0), 0) + add = tl.where(t == 2, tl.sum(take.to(tl.int64), axis=0), add) + add = tl.where(t == 3, tl.sum((held & ~take).to(tl.int64), axis=0), add) + add = tl.where(t == 6, tl.sum(fwd.to(tl.int64), axis=0), add) + add = tl.where(t == 7, tl.sum(stale.to(tl.int64), axis=0), add) + add = tl.where(t == 8, tl.sum(out_of_range.to(tl.int64), axis=0), add) + add = tl.where(t == 9, tl.sum((live & ~s_ok).to(tl.int64), axis=0), add) + add = tl.where((t == 0) & (tl.program_id(0) == 0), 1, add) + tl.atomic_add(totals_ptr + t, add, mask=add != 0, sem="relaxed") + + +@triton.jit +def _audit_nums_kernel(count_ptr, nums_ptr, L: tl.constexpr, C: tl.constexpr, R: tl.constexpr, BLOCK: tl.constexpr): + j = tl.arange(0, BLOCK) + m = j < L * C + cnt = tl.load(count_ptr + j // C, mask=m, other=0) + tl.store(nums_ptr + j, tl.minimum(tl.maximum(cnt - (j % C) * R, 0), R).to(tl.int64), mask=m) + + +@triton.jit +def _audit_compare_kernel(cache_ptrs, ref_ptrs, words_ptr, slots_ptr, num_ptr, bad_cnt_ptr, bad_first_ptr, + NB: tl.constexpr, BLOCK: tl.constexpr): + """Program (row, bank, chunk): one chunk of a listed slot's bank row against its gathered host row.""" + r = tl.program_id(0) + if r < tl.load(num_ptr): + b = tl.program_id(1) + words = tl.load(words_ptr + b) + off = tl.program_id(2).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + m = off < words + slot = tl.load(slots_ptr + r).to(tl.int64) + got = tl.load(tl.load(cache_ptrs + b).to(tl.pointer_type(tl.int32)) + slot * words + off, mask=m, other=0) + want = tl.load(tl.load(ref_ptrs + b).to(tl.pointer_type(tl.int32)) + r * words + off, mask=m, other=0) + bad = m & (got != want) + nbad = tl.sum(bad.to(tl.int32), axis=0) + if nbad > 0: + tl.atomic_add(bad_cnt_ptr + slot * NB + b, nbad) + tl.atomic_min(bad_first_ptr + slot * NB + b, tl.min(tl.where(bad, off, _NO_BAD), axis=0).to(tl.int32)) + + +@triton.jit(do_not_specialize=["num_slots", "num_ids"]) +def _audit_report_kernel( + id_of_slot_ptr, slot_for_id_ptr, usage_ptr, lstep_ptr, dstep_ptr, + flags_ptr, stale_id_ptr, stale_n_ptr, bad_cnt_ptr, bad_first_ptr, reported_ptr, + meta_hist_ptr, meta_cur_ptr, bytes_hist_ptr, bytes_cur_ptr, + ring_ptr, cursor_ptr, totals_ptr, count_ptr, + num_slots, num_ids, + NB: tl.constexpr, NB_P: tl.constexpr, L: tl.constexpr, L_P: tl.constexpr, RING_N: tl.constexpr, + BLOCK: tl.constexpr, +): + """Keep each new bad (slot, owner) with its history, count the findings and clear the work arrays.""" + s = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + ms = s < num_slots + b = tl.arange(0, NB_P) + cell = ms[:, None] & (b < NB)[None, :] + cnt = tl.load(bad_cnt_ptr + s[:, None] * NB + b[None, :], mask=cell, other=0) + first = tl.load(bad_first_ptr + s[:, None] * NB + b[None, :], mask=cell, other=_NO_BAD) + bank_bad = cnt > 0 + bytes_bad = tl.sum(bank_bad.to(tl.int32), axis=1) > 0 + bank_mask = tl.sum(tl.where(bank_bad, 1 << b[None, :], 0), axis=1) + first_bank = tl.min(tl.where(bank_bad, b[None, :], NB_P), axis=1) + first_word = tl.sum(tl.where(b[None, :] == first_bank[:, None], first, 0), axis=1) + fl = tl.load(flags_ptr + s, mask=ms, other=0) | tl.where(bytes_bad, 1, 0) + bad = ms & (fl != 0) + owner = tl.load(id_of_slot_ptr + s, mask=ms, other=-1) + sid = tl.load(stale_id_ptr + s, mask=ms, other=-1) + key = (owner.to(tl.int64) + 1) * (num_ids + 2) + sid + 1 + new = bad & (tl.load(reported_ptr + s, mask=ms, other=-1) != key) + tl.store(reported_ptr + s, tl.where(bad, key, -1), mask=ms) + pos = tl.atomic_add(cursor_ptr + s * 0, 1, mask=new) + keep = new & (pos < RING_N) + in_range = ms & (owner >= 0) & (owner < num_ids) + f = tl.arange(0, 16) + z = tl.zeros([BLOCK], tl.int64) + vals = tl.zeros([BLOCK, 16], tl.int64) - 1 + vals = _put(vals, f, 0, z + tl.load(totals_ptr) - 1) + vals = _put(vals, f, 1, z + tl.load(dstep_ptr)) + vals = _put(vals, f, 2, z + tl.load(lstep_ptr)) + vals = _put(vals, f, 3, s) + vals = _put(vals, f, 4, owner) + vals = _put(vals, f, 5, tl.load(slot_for_id_ptr + owner, mask=in_range, other=-2)) + vals = _put(vals, f, 6, tl.load(usage_ptr + s, mask=ms, other=-1)) + vals = _put(vals, f, 7, fl) + vals = _put(vals, f, 8, bank_mask) + vals = _put(vals, f, 9, tl.where(bytes_bad, first_bank, -1)) + vals = _put(vals, f, 10, tl.where(bytes_bad, first_word.to(tl.int64) * 4, -1)) + vals = _put(vals, f, 11, tl.sum(cnt, axis=1)) + vals = _put(vals, f, 12, sid) + vals = _put(vals, f, 13, tl.load(stale_n_ptr + s, mask=ms, other=0)) + vals = _put(vals, f, 14, tl.load(meta_cur_ptr + s, mask=ms, other=0)) + vals = _put(vals, f, 15, tl.load(bytes_cur_ptr + s, mask=ms, other=0)) + rec = ring_ptr + pos.to(tl.int64) * (16 + 2 * _D * _F) + tl.store(rec[:, None] + f[None, :], vals, mask=keep[:, None]) + j = tl.arange(0, 64) + jm = keep[:, None] & (j < _D * _F)[None, :] + src = s.to(tl.int64)[:, None] * (_D * _F) + j[None, :] + tl.store(rec[:, None] + 16 + j[None, :], tl.load(meta_hist_ptr + src, mask=jm, other=-1), mask=jm) + tl.store(rec[:, None] + 16 + _D * _F + j[None, :], tl.load(bytes_hist_ptr + src, mask=jm, other=-1), mask=jm) + t = tl.arange(0, 16) + add = tl.where(t == 4, tl.sum(bad.to(tl.int64), axis=0), 0) + add = tl.where(t == 5, tl.sum((ms & bytes_bad).to(tl.int64), axis=0), add) + add = tl.where(t == 10, tl.sum(new.to(tl.int64), axis=0), add) + add = tl.where(t == 11, tl.sum((bad & ~new).to(tl.int64), axis=0), add) + tl.atomic_add(totals_ptr + t, add, mask=add != 0, sem="relaxed") + tl.store(flags_ptr + s, tl.zeros([BLOCK], tl.int32), mask=ms) + tl.store(stale_id_ptr + s, tl.zeros([BLOCK], tl.int32) - 1, mask=ms) + tl.store(stale_n_ptr + s, tl.zeros([BLOCK], tl.int32), mask=ms) + tl.store(bad_cnt_ptr + s[:, None] * NB + b[None, :], tl.zeros([BLOCK, NB_P], tl.int32), mask=cell) + tl.store(bad_first_ptr + s[:, None] * NB + b[None, :], tl.zeros([BLOCK, NB_P], tl.int32) + _NO_BAD, mask=cell) + if tl.program_id(0) == 0: + ll = tl.arange(0, L_P) + tl.store(count_ptr + ll, tl.zeros([L_P], tl.int32), mask=ll < L) + + +class SlotAuditor: + """Per-slot writer history (recorded by ``record_*`` on the writer's stream) and the periodic + audit (``end_decode_step`` -> ``scan`` on the compute stream). ``bind`` sizes everything for a + cache's slots and banks; a rebuild rebinds and starts a new history, the totals and records stay. + Holds no reference to its cache: every call takes it.""" + + def __init__(self, num_layers: int, num_experts: int, device: torch.device, *, interval: int, + ring: int | None = None, scratch_bytes: int = SCRATCH_BYTES) -> None: + self.num_layers = num_layers + self.num_experts = num_experts + self.device = device + self.interval = interval + self.ring_n = RING if ring is None else ring + self.scratch_bytes = scratch_bytes + ring = self.ring_n + self.dstep = torch.zeros((1,), dtype=torch.int64, device=device) + self.seq = torch.zeros((1,), dtype=torch.int64, device=device) + self.totals = torch.zeros((NUM_TOTALS,), dtype=torch.int64, device=device) + self.ring = torch.full((ring, RECORD_FIELDS), -1, dtype=torch.int64, device=device) + self.cursor = torch.zeros((1,), dtype=torch.int64, device=device) + self.count = torch.zeros((num_layers,), dtype=torch.int32, device=device) + self.num_slots = 0 + self.bank_names: list[str] = [] + self._feats: list[int] = [] + self._steps = 0 + self._printed = 0 + self._last = torch.zeros((NUM_TOTALS,), dtype=torch.int64) + self._dummy = torch.zeros((1,), dtype=torch.int32, device=device) + logger.info(f"MoE slot audit: recording every slot writer, auditing every held slot each {interval} decode steps") + + # ------------------------------------------------------------------ setup + + def bind(self, cache) -> None: + """Size the history and audit arrays for ``cache``'s slots and banks (a rebuild starts a new + history) and point the byte checks at its banks.""" + caches = [bank for _, bank in cache.banks] + feats = [math.prod(bank.shape[1:]) * bank.element_size() for bank in caches] + assert all(f % 4 == 0 for f in feats), feats + if self.bank_names: + assert list(cache.bank_schema) == self.bank_names and feats == self._feats, "the bank layout changed" + else: + self.bank_names = list(cache.bank_schema) + self._feats = feats + self._alloc_scratch(cache, feats) + self.cache_ptrs = torch.tensor([bank.data_ptr() for bank in caches], dtype=torch.int64, device=self.device) + S, dev = cache.cache_size, self.device + self.num_slots = S + self.meta_hist = torch.full((S, DEPTH, EVENT_FIELDS), -1, dtype=torch.int64, device=dev) + self.bytes_hist = torch.full((S, DEPTH, EVENT_FIELDS), -1, dtype=torch.int64, device=dev) + self.meta_cur = torch.zeros((S,), dtype=torch.int32, device=dev) + self.bytes_cur = torch.zeros((S,), dtype=torch.int32, device=dev) + self.flags = torch.zeros((S,), dtype=torch.int32, device=dev) + self.stale_id = torch.full((S,), -1, dtype=torch.int32, device=dev) + self.stale_n = torch.zeros((S,), dtype=torch.int32, device=dev) + self.reported = torch.full((S,), -1, dtype=torch.int64, device=dev) + nb = len(feats) + self.bad_cnt = torch.zeros((S * nb,), dtype=torch.int32, device=dev) + self.bad_first = torch.full((S * nb,), 0x7FFFFFFF, dtype=torch.int32, device=dev) + # a correct cache holds at most num_experts slots of a layer; more are counted unchecked + per_layer = min(self.num_experts, S) + self.chunks = triton.cdiv(per_layer, self.rows) + self.cap = self.chunks * self.rows + self.list_rows = torch.zeros((self.num_layers, self.cap), dtype=torch.int32, device=dev) + self.list_slots = torch.zeros((self.num_layers, self.cap), dtype=torch.int32, device=dev) + self.nums = torch.zeros((self.num_layers * self.chunks,), dtype=torch.int64, device=dev) + # layers without a device alias (LOCKED/PAGEABLE host banks) cannot be gathered + self.layer_ok_host = [layer not in cache._unpinned_layers for layer in range(self.num_layers)] + self.layer_ok = torch.tensor(self.layer_ok_host, dtype=torch.int32, device=dev) + self.count.zero_() + + def _alloc_scratch(self, cache, feats: list[int]) -> None: + row = sum(feats) + self.rows = max(1, min(MAX_CHUNK_ROWS, (self.scratch_bytes - 1) // row)) + rows = self.rows + self.scratch = torch.empty((rows * row,), dtype=torch.uint8, device=self.device) + offsets = [rows * sum(feats[:b]) for b in range(len(feats))] + self.scratch_views = [ + self.scratch[o : o + rows * f].view(bank.dtype).view(rows, *bank.shape[1:]) + for o, f, (_, bank) in zip(offsets, feats, cache.banks) + ] + self.scratch_ptrs = torch.tensor([self.scratch.data_ptr() + o for o in offsets], dtype=torch.int64, device=self.device) + self.words = torch.tensor([f // 4 for f in feats], dtype=torch.int64, device=self.device) + self.iota = torch.arange(rows, dtype=torch.int32, device=self.device) + self._word_chunks = triton.cdiv(max(feats) // 4, COMPARE_WORDS) + logger.info(f"MoE slot audit scratch: {self.scratch.numel() / 2**20:.1f} MiB, {rows} experts a chunk") + + # ------------------------------------------------------------------ recording (graph-capturable) + + def begin_step(self) -> None: + """Advance the decode-step clock (the first GPU layer's decode ensure calls this).""" + _bump_step_kernel[(1,)](self.dstep) + + def _launch(self, cache, kind: int, layer: int, *, meta: bool, width: int, dst=None, src=None, aux=None, + num=None, start: int = 0, row: int = ROW_PLAN, row_base: int = 0, pred: int = PRED_ALL) -> None: + if width <= 0 or not self.num_slots: + return + block = min(triton.next_power_of_2(width), 128) + d = self._dummy + _record_kernel[(triton.cdiv(width, block),)]( + self.meta_hist if meta else self.bytes_hist, self.meta_cur if meta else self.bytes_cur, + self.seq, self.dstep, cache.step, cache.id_of_slot, cache.usage, self.totals, + d if dst is None else dst, d if src is None else src, d if aux is None else aux, d if num is None else num, + kind, layer, start, row_base, width, self.num_slots, + PLAN=dst is not None, ROW=row, HAS_AUX=aux is not None, HAS_NUM=num is not None, PRED=pred, + E=self.num_experts, BLOCK=block, num_warps=4, + ) + + def record_plan(self, cache, kind: int, layer: int, dst: torch.Tensor, src: torch.Tensor | None, num: torch.Tensor, + *, meta: bool, width: int | None = None, aux: torch.Tensor | None = None, row_base: int = 0) -> None: + """The writer just wrote slots ``dst[:num]`` (rows ``src[:num]``, else ``dst - row_base``).""" + width = dst.numel() if width is None else min(width, dst.numel()) + self._launch(cache, kind, layer, meta=meta, width=width, dst=dst, src=src, aux=aux, num=num, + row=ROW_PLAN if src is not None else ROW_FROM_SLOT, row_base=row_base) + + def record_range(self, cache, kind: int, layer: int, start: int, n: int, *, meta: bool) -> None: + """The writer just wrote expert ``i`` of ``layer`` into slot ``start + i`` for ``i < n``.""" + self._launch(cache, kind, layer, meta=meta, width=n, start=start, row=ROW_INDEX) + + def record_clear(self, cache, kind: int, context_layer: int, start: int, n: int, *, layer_above_e: bool = False) -> None: + """Before a writer clears held slots in ``[start, start + n)`` (with ``layer_above_e``: only + those >= E held by ``context_layer``); the event names the old owner.""" + self._launch(cache, kind, context_layer, meta=True, width=n, start=start, row=ROW_OLD_OWNER, + pred=PRED_LAYER_ABOVE_E if layer_above_e else PRED_HELD) + + # ------------------------------------------------------------------ audit + + def end_decode_step(self, cache) -> bool: + """After a decode forward, on the compute stream: audit every ``interval`` steps.""" + self._steps += 1 + if self._steps % self.interval: + return False + self.scan(cache) + return True + + def scan(self, cache) -> None: + """Audit every held slot now (enqueued on the current stream; no host sync).""" + S, num_ids = self.num_slots, self.num_layers * self.num_experts + assert S == cache.cache_size and cache.id_of_slot.numel() == S, "bind after a rebuild" + block = 1024 + _audit_meta_kernel[(triton.cdiv(max(S, num_ids), block),)]( + cache.id_of_slot, cache.slot_for_id.view(-1), self.layer_ok, self.count, self.list_rows, self.list_slots, + self.flags, self.stale_id, self.stale_n, self.totals, S, num_ids, + E=self.num_experts, CAP=self.cap, BLOCK=block, num_warps=4, + ) + n_lists = self.num_layers * self.chunks + _audit_nums_kernel[(1,)](self.count, self.nums, L=self.num_layers, C=self.chunks, R=self.rows, + BLOCK=triton.next_power_of_2(n_lists)) + nb = len(self.bank_names) + R = self.rows + for layer in range(self.num_layers): + if not self.layer_ok_host[layer]: + continue + for c in range(self.chunks): + rows = self.list_rows[layer, c * R : (c + 1) * R] + slots = self.list_slots[layer, c * R : (c + 1) * R] + num = self.nums[layer * self.chunks + c : layer * self.chunks + c + 1] + self._gather(cache, layer, rows, num) + _audit_compare_kernel[(R, nb, self._word_chunks)]( + self.cache_ptrs, self.scratch_ptrs, self.words, slots, num, self.bad_cnt, self.bad_first, + NB=nb, BLOCK=COMPARE_WORDS, num_warps=8, + ) + rblock = 256 + _audit_report_kernel[(triton.cdiv(S, rblock),)]( + cache.id_of_slot, cache.slot_for_id.view(-1), cache.usage, cache.step, self.dstep, + self.flags, self.stale_id, self.stale_n, self.bad_cnt, self.bad_first, self.reported, + self.meta_hist, self.meta_cur, self.bytes_hist, self.bytes_cur, + self.ring, self.cursor, self.totals, self.count, S, num_ids, + NB=nb, NB_P=triton.next_power_of_2(nb), L=self.num_layers, L_P=triton.next_power_of_2(self.num_layers), + RING_N=self.ring_n, BLOCK=rblock, num_warps=4, + ) + + def _gather(self, cache, layer: int, rows: torch.Tensor, num: torch.Tensor) -> None: + # the plain fused copy, not the slim one prefetch copies use: a slim-kernel bug cannot hide itself + if cache._copy_fused_ok: + from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit + + fast_index_copy_multi_jit( + self.scratch_ptrs, cache._copy_src_ptrs[layer], cache._copy_feat_bytes, self.iota, rows, num, + ) + return + from freetoken.kernel import fast_index_copy_jit + + for (per_layer, _), view in zip(cache.banks, self.scratch_views): + fast_index_copy_jit(view, self.iota, per_layer[layer], rows, num) + + # ------------------------------------------------------------------ reporting (host syncs) + + def take_window(self) -> tuple[list[int], list[int], list[list[int]], int]: + """(totals now, totals of this window, records kept since the last call, records in all).""" + totals = self.totals.cpu() + window = (totals - self._last).tolist() + self._last = totals + cursor = int(self.cursor.item()) + upto = min(cursor, self.ring_n) + records = self.ring[self._printed : upto].cpu().tolist() if upto > self._printed else [] + self._printed = max(self._printed, upto) + return totals.tolist(), window, records, cursor + + def report_window(self, steps: int) -> list[tuple[str, str]]: + """(level, line) pairs for one MOE_STATS_INTERVAL window; nothing when no audit ran.""" + totals, window, records, cursor = self.take_window() + if not window[T_AUDITS] and not records: + return [] + bad = bool(window[T_BAD] or window[T_MAP_RANGE] or window[T_BAD_PLAN] or records) + lines = [("warning" if bad else "info", + f"MoE slot audit ({steps} decode steps, every {self.interval}): {self.format_counts(window)}; " + f"session: {totals[T_AUDITS]} audits, {totals[T_NEW]} bad (slot, owner) pairs")] + lines += [("warning", "MoE slot audit record: " + self.format_record(r)) for r in records] + if records and cursor > self.ring_n: + lines.append(("warning", f"MoE slot audit: {cursor} bad (slot, owner) pairs in all, the first {self.ring_n} are kept")) + return lines + + def report_session(self) -> list[tuple[str, str]]: + """Session totals and every kept record (shutdown).""" + totals = self.totals.cpu().tolist() + cursor = int(self.cursor.item()) + if not totals[T_AUDITS]: + return [] + bad = bool(totals[T_BAD] or totals[T_MAP_RANGE] or totals[T_BAD_PLAN]) + lines = [("warning" if bad else "info", f"MoE slot audit (session, every {self.interval} decode steps): " + f"{self.format_counts(totals)}")] + records = self.ring[: min(cursor, self.ring_n)].cpu().tolist() + lines += [("warning", "MoE slot audit record: " + self.format_record(r)) for r in records] + if cursor > self.ring_n: + lines.append(("warning", f"MoE slot audit: {cursor} bad (slot, owner) pairs in all, the first {self.ring_n} are kept")) + return lines + + @staticmethod + def format_counts(t: list[int]) -> str: + audits = t[T_AUDITS] + per = (lambda v: f"{v / audits:.0f}") if audits else str + text = ( + f"{audits} audits, {per(t[T_HELD])} held slots an audit ({per(t[T_BYTES_CHECKED])} byte-checked), " + f"bad slots={t[T_BAD]} (bytes={t[T_BYTES]}, map={t[T_FWD]}, stale slot_for_id={t[T_STALE]}, " + f"owner out of range={t[T_OWNER_RANGE]}), new={t[T_NEW]}, repeat={t[T_STILL]}, events={t[T_EVENTS]}" + ) + if t[T_UNCHECKED]: + text += f", held slots not byte-checked={t[T_UNCHECKED]}" + if t[T_MAP_RANGE]: + text += f", slot_for_id entries past the cache={t[T_MAP_RANGE]}" + if t[T_BAD_PLAN]: + text += f", plan entries naming no slot={t[T_BAD_PLAN]}" + return text + + # ------------------------------------------------------------------ record decoding + + def _name(self, flat: int) -> str: + e = self.num_experts + return f"L{flat // e}/e{flat % e}" if flat >= 0 else str(flat) + + @staticmethod + def events(rec: list[int]) -> list[dict]: + """Both rings of a record as event dicts in ``seq`` order (``ring`` 'meta' or 'bytes').""" + out = [] + for ring, base in (("meta", HEADER), ("bytes", HEADER + RING_EVENTS)): + for d in range(DEPTH): + ev = rec[base + d * EVENT_FIELDS : base + (d + 1) * EVENT_FIELDS] + if ev[E_SEQ] >= 0: + out.append({ + "ring": ring, "seq": ev[E_SEQ], "kind": ev[E_KIND], "dstep": ev[E_DSTEP], "lstep": ev[E_LSTEP], + "layer": ev[E_LAYER], "row": ev[E_ROW], "aux": ev[E_AUX], "owner": ev[E_OWNER], + "usage": ev[E_USAGE], "pos": ev[E_POS], + }) + return sorted(out, key=lambda ev: ev["seq"]) + + def _flat(self, ev: dict) -> int: + return ev["layer"] * self.num_experts + ev["row"] if ev["layer"] >= 0 and ev["row"] >= 0 else -1 + + def format_event(self, ev: dict) -> str: + e = self.num_experts + text = f"#{ev['seq']} {KIND_NAMES.get(ev['kind'], ev['kind'])} L{ev['layer']}/e{ev['row']}" + if ev["kind"] == PREFILL_HIT_D2D: + text += f" from slot {ev['aux']}" + elif ev["kind"] in (MATERIALIZE_CLEAR, INVALIDATE) and ev["aux"] >= 0: + text += f" (for layer {ev['aux']})" + seen = "slot held" if ev["ring"] == "bytes" else "then held" + if ev["kind"] in (MATERIALIZE_CLEAR, INVALIDATE, RESET): + seen = "held" + owner = f"L{ev['owner'] // e}/e{ev['owner'] % e}" if ev["owner"] >= 0 else "nothing" + return f"{text} @step {ev['dstep']}/lru {ev['lstep']} ({seen} {owner}, usage {ev['usage']})" + + def _kind(self, ev: dict) -> str: + return KIND_NAMES.get(ev["kind"], str(ev["kind"])) + + def _pending_install(self, copy: dict, meta: list[dict], byte: list[dict], before: dict) -> dict | None: + """The earlier install (before ``before``) that ``copy`` completes: an install of the id it + copied with no copy of that id between them.""" + flat = self._flat(copy) + for inst in reversed([ev for ev in meta if ev["kind"] in INSTALLS and ev["seq"] < before["seq"]]): + if self._flat(inst) == flat: + done = any(self._flat(ev) == flat and inst["seq"] < ev["seq"] < copy["seq"] for ev in byte) + return None if done else inst + return None + + def _bytes_story(self, owner: int, meta: list[dict], byte: list[dict]) -> str: + insts = [ev for ev in meta if ev["kind"] in INSTALLS and self._flat(ev) == owner] + if not insts: + return (f"no recorded install of {self._name(owner)} in the kept history " + f"(an unrecorded map writer, or it is older than the last {DEPTH} events)") + inst = insts[-1] + name = self._kind(inst) + after = [ev for ev in byte if ev["seq"] > inst["seq"]] + own = [ev for ev in after if self._flat(ev) == owner] + if own: + last = after[-1] + if self._flat(last) == owner: + return (f"{name} and {self._kind(last)} of {self._name(owner)} (#{last['seq']}), then the bytes changed " + f"with no recorded write: overwritten after the last recorded write, or that copy did not land") + text = (f"{name} and {self._kind(own[-1])} of {self._name(owner)} (#{own[-1]['seq']}), then " + f"{self._kind(last)} of {self._name(self._flat(last))} (#{last['seq']}) overwrote it") + prior = self._pending_install(last, meta, byte, inst) + if prior is not None: + text += f": the late copy of the earlier {self._kind(prior)} of {self._name(self._flat(prior))} (#{prior['seq']})" + return text + text = f"{name}, no {KIND_NAMES[COPY_FOR[inst['kind']]]}" + if not after: + return text + last = after[-1] + prior = self._pending_install(last, meta, byte, inst) + if prior is not None: + return (f"{text}; the last write, {self._kind(last)} of {self._name(self._flat(last))} (#{last['seq']}), " + f"is the copy of the earlier {self._kind(prior)} (#{prior['seq']}) and landed after this install") + where = (f"from layer {last['layer']}'s source" if last["row"] == owner % self.num_experts + else f"of {self._name(self._flat(last))}") + return f"{name} of {self._name(owner)}, then {self._kind(last)} {where} ({self._name(self._flat(last))}, #{last['seq']})" + + def diagnose(self, rec: list[int]) -> str: + """What the history says about a bad slot, in words.""" + owner, fl = rec[H_OWNER], rec[H_FLAGS] + evs = self.events(rec) + meta = [ev for ev in evs if ev["ring"] == "meta"] + byte = [ev for ev in evs if ev["ring"] == "bytes"] + parts = [] + if fl & F_OWNER_RANGE: + parts.append(f"id_of_slot={owner} names no expert") + if fl & F_BYTES and owner >= 0: + parts.append(self._bytes_story(owner, meta, byte)) + if fl & F_FWD: + parts.append(f"slot_for_id[{self._name(owner)}]={rec[H_MAPPED]}, not this slot") + if fl & F_STALE: + sid = rec[H_STALE_ID] + text = (f"stale slot_for_id: {self._name(sid)} -> this slot, which holds {self._name(owner)}" + + (f" ({rec[H_STALE_N]} such ids)" if rec[H_STALE_N] > 1 else "")) + gone = [ev for ev in meta if ev["kind"] in INSTALLS and self._flat(ev) == sid] + if gone: + later = [ev for ev in meta if ev["seq"] > gone[-1]["seq"]] + text += (f"; {self._name(sid)} was installed here by {self._kind(gone[-1])} #{gone[-1]['seq']}" + + (f" and replaced by {self._kind(later[0])} #{later[0]['seq']} without clearing " + f"its map entry" if later else "")) + parts.append(text) + # an install's recorder reads the slot right after the install on the same stream + for ev in meta: + if ev["kind"] in INSTALLS and ev["owner"] != self._flat(ev): + parts.append(f"right after the {self._kind(ev)} of {self._name(self._flat(ev))} (#{ev['seq']}) the slot " + f"named {self._name(ev['owner'])}: another map writer ran concurrently") + return "; ".join(parts) or "clean" + + def format_record(self, rec: list[int]) -> str: + fl = rec[H_FLAGS] + text = (f"audit={rec[H_AUDIT]} step={rec[H_DSTEP]} lru={rec[H_LSTEP]} slot={rec[H_SLOT]} " + f"holds {self._name(rec[H_OWNER])} (slot_for_id={rec[H_MAPPED]}, usage={rec[H_USAGE]})") + if fl & F_BYTES: + banks = [self.bank_names[b] if b < len(self.bank_names) else str(b) + for b in range(max(len(self.bank_names), 1)) if rec[H_BANKS] >> b & 1] + first = self.bank_names[rec[H_FIRST_BANK]] if 0 <= rec[H_FIRST_BANK] < len(self.bank_names) else rec[H_FIRST_BANK] + text += (f" | bytes differ in {', '.join(banks)} (first at {first}+{rec[H_FIRST_BYTE]}, " + f"{rec[H_BAD_WORDS]} bad words)") + text += f" | diagnosis: {self.diagnose(rec)}" + history = "; ".join(self.format_event(ev) for ev in self.events(rec)) + return text + f" | history ({rec[H_META_N]} map / {rec[H_BYTES_N]} byte events in all): {history or 'none'}" diff --git a/python/freetoken/moe/verify.py b/python/freetoken/moe/verify.py new file mode 100644 index 0000000000..1f096c9c3a --- /dev/null +++ b/python/freetoken/moe/verify.py @@ -0,0 +1,516 @@ +"""Decode expert-cache verifier (``FREETOKEN_MOE_PREFETCH_VERIFY``): a debug detector for slot races. + +For every GPU decode layer call it checks, on the compute stream and in its order, what the +grouped GEMM is about to read and has just read: + +- meta (before and after the GEMM): each routed entry's slot maps to the expert the router chose, + in both slot maps (``id_of_slot[s] == layer * E + id`` and ``slot_for_id[layer, id] == s``); the + router's ids are copied before ensure rewrites them in place; +- bytes before the GEMM (``full``): every routed slot's row in every bank equals the host bank row + of that expert for that layer, gathered into a scratch buffer with the fused copy; +- bytes after the GEMM (``full``): the same slots against the same scratch, so a slot overwritten + while the GEMM ran shows up even when a later copy restores it. + +Nothing here records or waits on an event: the prefetch streams keep exactly the edges they have +without it. The checks do delay the compute stream (``full`` gathers every routed expert over +PCIe), which moves where the other streams land relative to it. ``meta`` skips the byte checks and +costs a few one-CTA launches per layer. Counts accumulate per layer on the device and the first +``RING`` offending records are kept with the slot state seen at check time; the engine reports +them every MOE_STATS_INTERVAL decode steps and at shutdown. +""" +from __future__ import annotations + +import math + +import torch +import triton +import triton.language as tl + +from freetoken.env import ENV +from freetoken.utils import init_logger + +logger = init_logger(__name__) + +MODES = ("off", "meta", "full") +# counters[layer, col]: checked calls, offending records per check, routed experts left unchecked +CHECKS, META_BAD, PRE_BAD, POST_BAD, POST_META_BAD, UNCHECKED = range(6) +NUM_COLS = 6 +COL_NAMES = ("checks", "meta_bad", "pre_bad", "post_bad", "post_meta_bad", "unchecked") +META_PRE, BYTES_PRE, META_POST, BYTES_POST = range(4) +KINDS = ("meta_pre", "bytes_pre", "meta_post", "bytes_post") +# record fields; owner / mapped / usage / step are read when the check runs +(R_CALL, R_LAYER, R_ROW, R_EXPERT, R_SLOT, R_BANK, R_KIND, R_OWNER, R_MAPPED, R_USAGE, R_STEP, + R_SOURCE, R_FIRST_BYTE, R_BAD_WORDS) = range(14) +NUM_FIELDS = 14 +# R_SOURCE bits: the slot is in this call's demand copy plan, this layer's prefetch copy plan, or +# the next layer's prefetch plan (installed earlier on the compute stream) +SRC_DEMAND, SRC_PREFETCH, SRC_NEXT_PREFETCH = 1, 2, 4 +RING = 64 +# routed entries a call may have (bs * top_k); wider eager batches are counted unchecked +MAX_ENTRIES = 64 +CHUNK_WORDS = 4096 +_NF = tl.constexpr(NUM_FIELDS) +_NC = tl.constexpr(NUM_COLS) +_NO_BAD = tl.constexpr(0x7FFFFFFF) + + +def resolve_mode(mode: str | None = None) -> str: + raw = (ENV.MOE_PREFETCH_VERIFY.value if mode is None else mode).strip().lower() + if raw in ("", "0", "off", "false", "no"): + return "off" + if raw in ("1", "on", "true", "yes", "full"): + return "full" + if raw == "meta": + return "meta" + raise ValueError(f"FREETOKEN_MOE_PREFETCH_VERIFY={raw!r}: expected 0, 1 (full) or meta") + + +@triton.jit +def _put(vals, f, idx: tl.constexpr, v): + return tl.where(f[None, :] == idx, v.to(tl.int64)[:, None], vals) + + +@triton.jit +def _in_plan(s, plan_ptr, num_ptr, W: tl.constexpr, W_P: tl.constexpr): + j = tl.arange(0, W_P) + plan = tl.load(plan_ptr + j, mask=j < tl.minimum(tl.load(num_ptr), W), other=-3) + return tl.sum((s[:, None] == plan[None, :]).to(tl.int32), axis=1) > 0 + + +@triton.jit +def _source(s, dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, + DEM_W: tl.constexpr, DEM_P: tl.constexpr, PF_W: tl.constexpr, PF_P: tl.constexpr, + HAS_PF: tl.constexpr, HAS_NX: tl.constexpr): + src = tl.where(_in_plan(s, dem_ptr, dem_num_ptr, DEM_W, DEM_P), 1, 0) + if HAS_PF: + src += tl.where(_in_plan(s, pf_ptr, pf_num_ptr, PF_W, PF_P), 2, 0) + if HAS_NX: + src += tl.where(_in_plan(s, nx_ptr, nx_num_ptr, PF_W, PF_P), 4, 0) + return src + + +@triton.jit +def _write(ring_ptr, cursor, bad, vals, RING_N: tl.constexpr): + """Append the ``bad`` rows of ``vals`` after ``cursor`` while the ring has room; the new cursor.""" + f = tl.arange(0, 16) + pos = cursor + tl.cumsum(bad.to(tl.int64), axis=0) - 1 + keep = bad & (pos < RING_N) + tl.store(ring_ptr + pos[:, None] * _NF + f[None, :], vals, mask=keep[:, None] & (f[None, :] < _NF)) + return cursor + tl.sum(bad.to(tl.int64), axis=0) + + +@triton.jit +def _meta(ids_ptr, slots_ptr, id_of_slot_ptr, slot_for_id_ptr, usage_ptr, step_ptr, + dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, + ring_ptr, cursor, counters_ptr, col: tl.constexpr, kind: tl.constexpr, call, layer, num_slots, + N: tl.constexpr, N_P: tl.constexpr, E: tl.constexpr, + DEM_W: tl.constexpr, DEM_P: tl.constexpr, PF_W: tl.constexpr, PF_P: tl.constexpr, + HAS_PF: tl.constexpr, HAS_NX: tl.constexpr, RING_N: tl.constexpr): + """Check every routed entry's slot against both maps; returns the new ring cursor.""" + i = tl.arange(0, N_P) + m = i < N + ids = tl.load(ids_ptr + i, mask=m, other=0) + s = tl.load(slots_ptr + i, mask=m, other=-1) + ok_s = m & (s >= 0) & (s < num_slots) + ok_id = m & (ids >= 0) & (ids < E) + fid = layer * E + ids + owner = tl.load(id_of_slot_ptr + s, mask=ok_s, other=-2) + mapped = tl.load(slot_for_id_ptr + fid, mask=ok_id, other=-2) + bad = m & ~(ok_s & ok_id & (owner == fid) & (mapped == s)) + usage = tl.load(usage_ptr + s, mask=ok_s, other=-1) + src = _source(s, dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, DEM_W, DEM_P, PF_W, PF_P, HAS_PF, HAS_NX) + f = tl.arange(0, 16) + z = tl.zeros([N_P], tl.int64) + vals = tl.zeros([N_P, 16], tl.int64) - 1 + vals = _put(vals, f, 0, z + call) + vals = _put(vals, f, 1, z + layer) + vals = _put(vals, f, 2, i) + vals = _put(vals, f, 3, ids) + vals = _put(vals, f, 4, s) + vals = _put(vals, f, 6, z + kind) + vals = _put(vals, f, 7, owner) + vals = _put(vals, f, 8, mapped) + vals = _put(vals, f, 9, usage) + vals = _put(vals, f, 10, z + tl.load(step_ptr)) + vals = _put(vals, f, 11, src) + nbad = tl.sum(bad.to(tl.int64), axis=0) + tl.store(counters_ptr + col, tl.load(counters_ptr + col) + nbad) + return _write(ring_ptr, cursor, bad, vals, RING_N) + + +@triton.jit(do_not_specialize=["layer", "num_slots"]) +def _plan_kernel( + ids_ptr, slots_ptr, id_of_slot_ptr, slot_for_id_ptr, usage_ptr, step_ptr, + dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, + ring_ptr, cursor_ptr, counters_ptr, calls_ptr, + u_ids_ptr, u_slots_ptr, u_rows_ptr, u_num_ptr, slots_copy_ptr, + layer, num_slots, + N: tl.constexpr, N_P: tl.constexpr, E: tl.constexpr, CAP: tl.constexpr, + DEM_W: tl.constexpr, DEM_P: tl.constexpr, PF_W: tl.constexpr, PF_P: tl.constexpr, + HAS_PF: tl.constexpr, HAS_NX: tl.constexpr, BYTES: tl.constexpr, RING_N: tl.constexpr, +): + """Before the GEMM: count the call, check the maps, keep the slots for the after-GEMM checks and + list the distinct (expert, slot) pairs to byte-check (first CAP of them).""" + call = tl.load(calls_ptr + layer) + tl.store(calls_ptr + layer, call + 1) + row = counters_ptr + layer * _NC + tl.store(row + 0, tl.load(row + 0) + 1) + cursor = tl.load(cursor_ptr) + cursor = _meta(ids_ptr, slots_ptr, id_of_slot_ptr, slot_for_id_ptr, usage_ptr, step_ptr, + dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, + ring_ptr, cursor, row, 1, 0, call, layer, num_slots, + N, N_P, E, DEM_W, DEM_P, PF_W, PF_P, HAS_PF, HAS_NX, RING_N) + tl.store(cursor_ptr, cursor) + i = tl.arange(0, N_P) + m = i < N + ids = tl.load(ids_ptr + i, mask=m, other=-1) + s = tl.load(slots_ptr + i, mask=m, other=-1) + tl.store(slots_copy_ptr + i, s, mask=m) + if BYTES: + # an unreadable slot or id is already a meta record; a duplicate pair shares its first row + ok = m & (s >= 0) & (s < num_slots) & (ids >= 0) & (ids < E) + same = (ids[:, None] == ids[None, :]) & (s[:, None] == s[None, :]) & (i[None, :] < i[:, None]) & ok[None, :] + first = ok & (tl.sum(same.to(tl.int32), axis=1) == 0) + rank = tl.cumsum(first.to(tl.int32), axis=0) - 1 + take = first & (rank < CAP) + tl.store(u_ids_ptr + rank, ids, mask=take) + tl.store(u_slots_ptr + rank, s, mask=take) + tl.store(u_rows_ptr + rank, i, mask=take) + total = tl.sum(first.to(tl.int64), axis=0) + kept = tl.minimum(total, CAP) + tl.store(u_num_ptr, kept) + tl.store(row + 5, tl.load(row + 5) + total - kept) + + +@triton.jit +def _compare_kernel(cache_ptrs, ref_ptrs, words_ptr, u_slots_ptr, u_num_ptr, bad_cnt_ptr, bad_first_ptr, + NB: tl.constexpr, BLOCK: tl.constexpr): + """Program (row, bank, chunk): compare one chunk of a routed slot's row with its host copy.""" + r = tl.program_id(0) + b = tl.program_id(1) + words = tl.load(words_ptr + b) + off = tl.program_id(2).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + m = (r < tl.load(u_num_ptr)) & (off < words) + slot = tl.load(u_slots_ptr + r).to(tl.int64) + got = tl.load(tl.load(cache_ptrs + b).to(tl.pointer_type(tl.int32)) + slot * words + off, mask=m, other=0) + want = tl.load(tl.load(ref_ptrs + b).to(tl.pointer_type(tl.int32)) + r * words + off, mask=m, other=0) + bad = m & (got != want) + nbad = tl.sum(bad.to(tl.int32), axis=0) + if nbad > 0: + tl.atomic_add(bad_cnt_ptr + r * NB + b, nbad) + tl.atomic_min(bad_first_ptr + r * NB + b, tl.min(tl.where(bad, off, _NO_BAD), axis=0).to(tl.int32)) + + +@triton.jit(do_not_specialize=["layer", "num_slots"]) +def _tally_kernel( + ids_ptr, slots_copy_ptr, id_of_slot_ptr, slot_for_id_ptr, usage_ptr, step_ptr, + dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, + ring_ptr, cursor_ptr, counters_ptr, calls_ptr, + u_ids_ptr, u_slots_ptr, u_rows_ptr, u_num_ptr, bad_cnt_ptr, bad_first_ptr, + layer, num_slots, + N: tl.constexpr, N_P: tl.constexpr, E: tl.constexpr, CAP: tl.constexpr, CAP_P: tl.constexpr, + NB: tl.constexpr, NB_P: tl.constexpr, + DEM_W: tl.constexpr, DEM_P: tl.constexpr, PF_W: tl.constexpr, PF_P: tl.constexpr, + HAS_PF: tl.constexpr, HAS_NX: tl.constexpr, BYTES: tl.constexpr, POST: tl.constexpr, + BYTES_KIND: tl.constexpr, BYTES_COL: tl.constexpr, RING_N: tl.constexpr, +): + """Turn the compare flags into counts and records (and clear them); after the GEMM also recheck + the maps of the slots it read.""" + call = tl.load(calls_ptr + layer) - 1 + row = counters_ptr + layer * _NC + cursor = tl.load(cursor_ptr) + if POST: + cursor = _meta(ids_ptr, slots_copy_ptr, id_of_slot_ptr, slot_for_id_ptr, usage_ptr, step_ptr, + dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, + ring_ptr, cursor, row, 4, 2, call, layer, num_slots, + N, N_P, E, DEM_W, DEM_P, PF_W, PF_P, HAS_PF, HAS_NX, RING_N) + if BYTES: + j = tl.arange(0, CAP_P * NB_P) + r = j // NB_P + b = j % NB_P + cell = (r < CAP) & (b < NB) + live = cell & (r < tl.load(u_num_ptr)) + cnt = tl.load(bad_cnt_ptr + r * NB + b, mask=live, other=0) + first = tl.load(bad_first_ptr + r * NB + b, mask=live, other=_NO_BAD) + tl.store(bad_cnt_ptr + r * NB + b, tl.zeros([CAP_P * NB_P], tl.int32), mask=cell) + tl.store(bad_first_ptr + r * NB + b, tl.zeros([CAP_P * NB_P], tl.int32) + _NO_BAD, mask=cell) + bad = live & (cnt > 0) + ids = tl.load(u_ids_ptr + r, mask=live, other=0) + s = tl.load(u_slots_ptr + r, mask=live, other=-1) + ok_s = live & (s >= 0) & (s < num_slots) + src = _source(s, dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, DEM_W, DEM_P, PF_W, PF_P, HAS_PF, HAS_NX) + f = tl.arange(0, 16) + z = tl.zeros([CAP_P * NB_P], tl.int64) + vals = tl.zeros([CAP_P * NB_P, 16], tl.int64) - 1 + vals = _put(vals, f, 0, z + call) + vals = _put(vals, f, 1, z + layer) + vals = _put(vals, f, 2, tl.load(u_rows_ptr + r, mask=live, other=-1)) + vals = _put(vals, f, 3, ids) + vals = _put(vals, f, 4, s) + vals = _put(vals, f, 5, b) + vals = _put(vals, f, 6, z + BYTES_KIND) + vals = _put(vals, f, 7, tl.load(id_of_slot_ptr + s, mask=ok_s, other=-2)) + vals = _put(vals, f, 8, tl.load(slot_for_id_ptr + layer * E + ids, mask=live, other=-2)) + vals = _put(vals, f, 9, tl.load(usage_ptr + s, mask=ok_s, other=-1)) + vals = _put(vals, f, 10, z + tl.load(step_ptr)) + vals = _put(vals, f, 11, src) + vals = _put(vals, f, 12, first.to(tl.int64) * 4) + vals = _put(vals, f, 13, cnt) + tl.store(row + BYTES_COL, tl.load(row + BYTES_COL) + tl.sum(bad.to(tl.int64), axis=0)) + cursor = _write(ring_ptr, cursor, bad, vals, RING_N) + tl.store(cursor_ptr, cursor) + + +class ExpertVerifier: + """Device counters, the record ring and (``full``) the scratch the byte checks compare against. + + ``begin`` copies the router's ids before ensure rewrites them, ``before_gemm`` runs after the + joins and the demand copy, ``after_gemm`` right after the GEMM; all three enqueue on the current + (compute) stream with fixed shapes and no host sync, so they capture into the decode graph. + ``bind`` sizes the scratch from the banks once and refreshes the slot bank addresses in place + (a rebuild reallocates the banks); counters survive rebuilds and resets. + """ + + def __init__(self, num_layers: int, num_experts: int, device: torch.device, *, mode: str = "full", + rows: int | None = None) -> None: + from freetoken.kernel.fast_index_copy import _skip_fast_index_copy_enabled + + if mode == "full" and _skip_fast_index_copy_enabled(): + logger.warning( + "FREETOKEN_MOE_PREFETCH_VERIFY=1 with FREETOKEN_SKIP_FAST_INDEX_COPY: the host rows cannot be " + "gathered, so only the slot maps are checked (meta)" + ) + mode = "meta" + self.mode = mode + self.num_layers = num_layers + self.num_experts = num_experts + self.device = device + self.rows = ENV.MOE_PREFETCH_VERIFY_ROWS.value if rows is None else int(rows) + if self.rows < 1: + raise ValueError(f"FREETOKEN_MOE_PREFETCH_VERIFY_ROWS={self.rows} must be >= 1") + self.counters = torch.zeros((num_layers, NUM_COLS), dtype=torch.int64, device=device) + self.totals = torch.zeros((num_layers, NUM_COLS), dtype=torch.int64) + self.calls = torch.zeros((num_layers,), dtype=torch.int64, device=device) + self.ring = torch.full((RING, NUM_FIELDS), -1, dtype=torch.int64, device=device) + self.cursor = torch.zeros((1,), dtype=torch.int64, device=device) + self.ids = torch.full((MAX_ENTRIES,), -1, dtype=torch.int32, device=device) + self.slots = torch.full((MAX_ENTRIES,), -1, dtype=torch.int32, device=device) + self.u_ids = torch.zeros((self.rows,), dtype=torch.int32, device=device) + self.u_slots = torch.zeros((self.rows,), dtype=torch.int32, device=device) + self.u_rows = torch.zeros((self.rows,), dtype=torch.int32, device=device) + self.u_num = torch.zeros((1,), dtype=torch.int64, device=device) + self.iota = torch.arange(self.rows, dtype=torch.int32, device=device) + self.bank_names: list[str] = [] + self.scratch: torch.Tensor | None = None + self._printed = 0 + self._top_k = 1 + self._call = None # (layer, entries, this layer's prefetch plan or None) between begin and after_gemm + self._next = None + logger.info( + f"MoE verify {self.mode}: checking every GPU decode layer call's slot maps" + + (f" and slot bytes before and after the GEMM (up to {self.rows} experts a call)" if self.mode == "full" else "") + ) + + def bind(self, cache) -> None: + """Size the scratch for ``cache``'s banks (first call) or point the checks at its new banks.""" + caches = [bank for _, bank in cache.banks] + feats = [math.prod(bank.shape[1:]) * bank.element_size() for bank in caches] + if self.bank_names: + assert list(cache.bank_schema) == self.bank_names and feats == self._feats, "the bank layout changed" + else: + assert all(f % 4 == 0 for f in feats), feats + self.bank_names = list(cache.bank_schema) + self._feats = feats + self.cache_ptrs = torch.zeros((len(feats),), dtype=torch.int64, device=self.device) + if self.mode == "full": + self._alloc_scratch(cache, feats) + self.cache_ptrs.copy_(torch.tensor([bank.data_ptr() for bank in caches], dtype=torch.int64)) + + def _alloc_scratch(self, cache, feats: list[int]) -> None: + rows = self.rows + self.scratch = torch.empty((rows * sum(feats),), dtype=torch.uint8, device=self.device) + offsets = [rows * sum(feats[:b]) for b in range(len(feats))] + self.scratch_views = [ + self.scratch[o : o + rows * f].view(bank.dtype).view(rows, *bank.shape[1:]) + for o, f, (_, bank) in zip(offsets, feats, cache.banks) + ] + self.scratch_ptrs = torch.tensor([self.scratch.data_ptr() + o for o in offsets], dtype=torch.int64, device=self.device) + self.words = torch.tensor([f // 4 for f in feats], dtype=torch.int64, device=self.device) + self.bad_cnt = torch.zeros((rows * len(feats),), dtype=torch.int32, device=self.device) + self.bad_first = torch.full((rows * len(feats),), 0x7FFFFFFF, dtype=torch.int32, device=self.device) + self._chunks = triton.cdiv(max(feats) // 4, CHUNK_WORDS) + logger.info(f"MoE verify scratch: {self.scratch.numel() / 2**20:.1f} MiB for {rows} experts") + + def begin(self, cache, layer_id: int, topk_ids: torch.Tensor, plan: tuple | None) -> None: + """Copy the router's ids of ``layer_id`` (``plan``: this layer's prefetch copy plan, if forked).""" + assert topk_ids.dtype == torch.int32 and topk_ids.is_contiguous() + n = topk_ids.numel() + self._top_k = topk_ids.shape[-1] + if n > MAX_ENTRIES: + self.counters[layer_id, UNCHECKED] += n + self._call = None + return + self.ids[:n].copy_(topk_ids.view(-1)) + self._call = (layer_id, n, plan) + + def _plans(self, cache) -> tuple: + _, n, plan = self._call + pf, nx = plan, self._next + dummy = (self.u_slots, self.u_num) + width = (pf or nx or dummy)[0].numel() + return ( + cache.evict_slots, cache.num_indices, *(pf or dummy), *(nx or dummy), + dict(DEM_W=n, DEM_P=triton.next_power_of_2(n), PF_W=width, PF_P=triton.next_power_of_2(width), + HAS_PF=pf is not None, HAS_NX=nx is not None), + ) + + def before_gemm(self, cache, slots: torch.Tensor, next_plan: tuple | None) -> None: + """Check the maps and (``full``) the bytes of the slots ``slots`` the GEMM is about to read + (``next_plan``: the next layer's prefetch plan, if one was issued).""" + if self._call is None: + return + layer, n, _ = self._call + self._next = next_plan + *plans, plan_kw = self._plans(cache) + full = self.mode == "full" + _plan_kernel[(1,)]( + self.ids, slots, cache.id_of_slot, cache.slot_for_id.view(-1), cache.usage, cache.step, *plans, + self.ring, self.cursor, self.counters, self.calls, + self.u_ids, self.u_slots, self.u_rows, self.u_num, self.slots, + layer, cache.cache_size, + N=n, N_P=triton.next_power_of_2(n), E=self.num_experts, CAP=self.rows, + BYTES=full, RING_N=RING, num_warps=4, **plan_kw, + ) + if not full: + return + self._gather(cache, layer) + self._compare_and_tally(cache, post=False) + + def after_gemm(self, cache) -> None: + """Recheck the maps and (``full``) the bytes of the slots the GEMM just read.""" + if self._call is None: + return + if self.mode == "full": + self._compare_and_tally(cache, post=True) + else: + self._tally(cache, post=True, bytes_=False) + self._call = self._next = None + + def _gather(self, cache, layer: int) -> None: + if cache._copy_fused_ok: + from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit + + fast_index_copy_multi_jit( + self.scratch_ptrs, cache._copy_src_ptrs[layer], cache._copy_feat_bytes, self.iota, self.u_ids, self.u_num, + ) + return + from freetoken.kernel import fast_index_copy_jit + + for (per_layer, _), view in zip(cache.banks, self.scratch_views): + fast_index_copy_jit(view, self.iota, per_layer[layer], self.u_ids, self.u_num) + + def _compare_and_tally(self, cache, *, post: bool) -> None: + nb = len(self.bank_names) + _compare_kernel[(self.rows, nb, self._chunks)]( + self.cache_ptrs, self.scratch_ptrs, self.words, self.u_slots, self.u_num, self.bad_cnt, self.bad_first, + NB=nb, BLOCK=CHUNK_WORDS, num_warps=8, + ) + self._tally(cache, post=post, bytes_=True) + + def _tally(self, cache, *, post: bool, bytes_: bool) -> None: + layer, n, _ = self._call + *plans, plan_kw = self._plans(cache) + nb = max(len(self.bank_names), 1) + _tally_kernel[(1,)]( + self.ids, self.slots, cache.id_of_slot, cache.slot_for_id.view(-1), cache.usage, cache.step, *plans, + self.ring, self.cursor, self.counters, self.calls, + self.u_ids, self.u_slots, self.u_rows, self.u_num, + self.bad_cnt if bytes_ else self.u_rows, self.bad_first if bytes_ else self.u_rows, + layer, cache.cache_size, + N=n, N_P=triton.next_power_of_2(n), E=self.num_experts, CAP=self.rows, + CAP_P=triton.next_power_of_2(self.rows), NB=nb, NB_P=triton.next_power_of_2(nb), + BYTES=bytes_, POST=post, BYTES_KIND=BYTES_POST if post else BYTES_PRE, BYTES_COL=POST_BAD if post else PRE_BAD, + RING_N=RING, num_warps=4, **plan_kw, + ) + + # ------------------------------------------------------------------ reporting (host syncs) + + def take_window(self) -> tuple[torch.Tensor, list[list[int]], int]: + """This window's ``[num_layers, NUM_COLS]`` counts (added to ``totals``; the device ones + restart), the records kept since the last call, and how many records there were in all.""" + window = self.counters.cpu() + self.counters.zero_() + self.totals += window + cursor = int(self.cursor.item()) + upto = min(cursor, RING) + records = self.ring[self._printed : upto].cpu().tolist() if upto > self._printed else [] + self._printed = max(self._printed, upto) + return window, records, cursor + + def report_window(self, steps: int) -> list[tuple[str, str]]: + """(level, line) pairs for one MOE_STATS_INTERVAL window; nothing when no call was checked.""" + window, records, cursor = self.take_window() + if not int(window[:, CHECKS].sum()) and not records: + return [] + bad = self._bad(window) or bool(records) + session = self._bad(self.totals) + lines = [("warning" if bad else "info", + f"MoE verify {self.mode} ({steps} decode steps): {self.format_counts(window)}; session bad={session}")] + lines += [("warning", "MoE verify record: " + self.format_record(r)) for r in records] + if bad and cursor > RING: + lines.append(("warning", f"MoE verify: {cursor} bad records in all, the first {RING} are kept")) + return lines + + def report_session(self) -> list[tuple[str, str]]: + """Session counts and every kept record (shutdown).""" + totals = self.totals + self.counters.cpu() + cursor = int(self.cursor.item()) + if not int(totals[:, CHECKS].sum()): + return [] + bad = self._bad(totals) + lines = [("warning" if bad else "info", f"MoE verify {self.mode} (session): {self.format_counts(totals)}")] + records = self.ring[: min(cursor, RING)].cpu().tolist() + lines += [("warning", "MoE verify record: " + self.format_record(r)) for r in records] + if cursor > RING: + lines.append(("warning", f"MoE verify: {cursor} bad records in all, the first {RING} are kept")) + return lines + + @staticmethod + def _bad(counts: torch.Tensor) -> int: + return int(counts[:, [META_BAD, PRE_BAD, POST_BAD, POST_META_BAD]].sum()) + + def format_counts(self, counts: torch.Tensor) -> str: + total = counts.sum(0).tolist() + text = f"{total[CHECKS]} layer calls checked, " + ", ".join( + f"{COL_NAMES[c]}={total[c]}" for c in (META_BAD, PRE_BAD, POST_BAD, POST_META_BAD) + ) + if total[UNCHECKED]: + text += f", unchecked experts={total[UNCHECKED]}" + per_layer = [ + f"L{layer}=" + "/".join(str(row[c]) for c in (META_BAD, PRE_BAD, POST_BAD, POST_META_BAD)) + for layer, row in enumerate(counts.tolist()) + if any(row[c] for c in (META_BAD, PRE_BAD, POST_BAD, POST_META_BAD)) + ] + if per_layer: + text += " (meta/pre/post/post_meta by layer: " + ", ".join(per_layer) + ")" + return text + + def format_record(self, rec: list[int]) -> str: + e = self.num_experts + kind = KINDS[rec[R_KIND]] if 0 <= rec[R_KIND] < len(KINDS) else str(rec[R_KIND]) + owner = rec[R_OWNER] + owner_text = f"L{owner // e}/e{owner % e}" if owner >= 0 else str(owner) + src = [name for bit, name in ((SRC_DEMAND, "demand copy"), (SRC_PREFETCH, "prefetch copy"), + (SRC_NEXT_PREFETCH, "next layer's prefetch plan")) if rec[R_SOURCE] & bit] + row = rec[R_ROW] + text = ( + f"{kind} call={rec[R_CALL]} layer={rec[R_LAYER]} row={row} (token {row // self._top_k}, rank {row % self._top_k}) " + f"expert={rec[R_EXPERT]} slot={rec[R_SLOT]}" + ) + if rec[R_BANK] >= 0: + bank = self.bank_names[rec[R_BANK]] if rec[R_BANK] < len(self.bank_names) else rec[R_BANK] + text += f" bank={bank} first_bad_byte={rec[R_FIRST_BYTE]} bad_words={rec[R_BAD_WORDS]}" + return text + ( + f" | slot now holds {owner_text}, slot_for_id={rec[R_MAPPED]}, usage={rec[R_USAGE]}, lru_step={rec[R_STEP]}, " + f"slot in: {', '.join(src) or 'no copy plan (a resident hit)'}" + ) diff --git a/tests/models/qwen4_exp/conftest.py b/tests/models/qwen4_exp/conftest.py index 2b9e80214f..4e2bb67bb8 100644 --- a/tests/models/qwen4_exp/conftest.py +++ b/tests/models/qwen4_exp/conftest.py @@ -13,3 +13,16 @@ def _runtime(): core._GLOBAL_CTX = None yield core._GLOBAL_CTX = None + import torch + + if torch.cuda.is_available() and torch.cuda.is_initialized(): + import gc + + # torch keeps a cuBLAS workspace per stream for good: each offload cache's dedicated + # streams would otherwise pin ~16 MiB for the rest of the session + gc.collect() + torch.cuda.synchronize() + clear = getattr(torch._C, "_cuda_clearCublasWorkspaces", None) # private torch API + if clear is not None: + clear() + torch.cuda.empty_cache() diff --git a/tests/models/qwen4_exp/test_build_cpu.py b/tests/models/qwen4_exp/test_build_cpu.py index 58ad0c4dc3..85366638a5 100644 --- a/tests/models/qwen4_exp/test_build_cpu.py +++ b/tests/models/qwen4_exp/test_build_cpu.py @@ -47,3 +47,32 @@ def test_every_layer_gets_its_own_prefixed_weights(): mlp_keys = {k for k in sd if ".mlp." in k} layers = {k.split(".layers.")[1].split(".")[0] for k in mlp_keys if ".layers." in k} assert len(layers) > 1, f"MoE weights landed on a single layer: {sorted(layers)}" + + +def test_offloaded_layers_borrow_the_next_router_for_the_lookahead(): + # FREETOKEN_MOE_PREFETCH: layer L-1's experts hold layer L's router, outside the state dict + # and without the engine's offload-layer walk finding extra layers through it + import dataclasses + + from freetoken.moe.offload_cache import iter_offload_moe_layers + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + previous = rotary._ROPE_DEVICE + set_rope_device(torch.device("cpu")) + try: + config = dataclasses.replace(parse_config(toy_hf_config()), moe_strategy="offload") + with torch.device("meta"): + model = Qwen4ExpForCausalLM(config) + finally: + rotary.get_rope.cache_clear() + set_rope_device(previous) + mlps = [layer.mlp for layer in model.model.layers.op_list] + assert len(list(iter_offload_moe_layers(model))) == len(mlps) == 4 + wired = [mlp.experts._lookahead for mlp in mlps] + assert wired[-1] is None + assert [(gate is mlps[t].gate, t, budget) for gate, t, budget in wired[:-1]] == [ + (True, 1, 3), (True, 2, 3), (True, 3, 4), # layer 3 is the full-attention one + ] + assert not any("lookahead" in k for k in model.state_dict()) + assert sum(k.endswith(".mlp.gate.weight") for k in model.state_dict()) == 4 diff --git a/tests/models/qwen4_exp/test_skeleton.py b/tests/models/qwen4_exp/test_skeleton.py index a2412f8346..2b47748f53 100644 --- a/tests/models/qwen4_exp/test_skeleton.py +++ b/tests/models/qwen4_exp/test_skeleton.py @@ -556,6 +556,1917 @@ def test_moe_copy_overlap_is_bitwise_identical(kind, graph, bs, monkeypatch): assert torch.equal(a, b) +# FREETOKEN_MOE_PREFETCH stack: 4 layers (full attention on layer 3), 16 experts, top-2, and one +# shared cache as small as a layer so every step evicts +_PF_LAYERS, _PF_EXPERTS, _PF_TOPK, _PF_HIDDEN, _PF_INTER = 4, 16, 2, 256, 128 +# ~50-100 us between layers: the predictor must finish before the next ensure even at toy size +_PF_SLEEP_CYCLES = 200_000 + + +def _prefetch_banks(kind): + from freetoken.moe.expert_banks import build_expert_banks + + torch.manual_seed(7) + method = _offload_experts(kind, 0, _PF_EXPERTS, _PF_TOPK, _PF_HIDDEN, _PF_INTER).quant_method + return build_expert_banks(method, _PF_LAYERS, None, device=torch.device("cuda"), dummy=True) + + +def _prefetch_stack(kind, banks, mode, monkeypatch, *, k=6, budget=0, wire=True, overlap=False, policy="rule", + cache_size=_PF_EXPERTS, **cache_kw): + from freetoken.env import ENV + from freetoken.models.qwen4_exp.moe import Qwen4ExpMoE, wire_router_lookahead + from freetoken.moe.offload_cache import OffloadMoeCache + from freetoken.utils.torch_utils import torch_dtype + + monkeypatch.setattr(ENV.MOE_PREFETCH_K, "value", k) + monkeypatch.setattr(ENV.MOE_PREFETCH_BUDGET, "value", budget) + config = parse_config(toy_hf_config( + _PF_LAYERS, hidden_size=_PF_HIDDEN, num_experts=_PF_EXPERTS, num_experts_per_tok=_PF_TOPK, + moe_intermediate_size=_PF_INTER, shared_expert_intermediate_size=_PF_INTER, + )) + device = torch.device("cuda") + with torch.device(device), torch_dtype(torch.bfloat16): + moes = [Qwen4ExpMoE(config, layer_id) for layer_id in range(_PF_LAYERS)] + gen = torch.Generator(device=device).manual_seed(43) + for moe in moes: + _fill(moe, gen, scale=0.2) + experts = [_offload_experts(kind, i, _PF_EXPERTS, _PF_TOPK, _PF_HIDDEN, _PF_INTER) for i in range(_PF_LAYERS)] + cache = OffloadMoeCache( + num_layers=_PF_LAYERS, num_experts=_PF_EXPERTS, cache_size=cache_size, device=device, + cache_policy=policy, quant_format=banks.quant_format, layout=banks.layout, + max_slots=experts[0].quant_method.slot_limit(), prefetch_mode=mode, decode_copy_overlap=overlap, + **cache_kw, + ) + cache.set_bank_sources(banks.sources) + cache.set_alphas(banks.gate_up_alpha, banks.down_alpha) + cache.collect_stats = True + for moe, ex in zip(moes, experts): + ex.offload_cache = cache + moe.experts = ex + if wire: + wire_router_lookahead(moes, config) + cache.reset() + return moes, cache + + +def _prefetch_run(moes, cache, bs, graph, inputs, before_decode=None, sleep=_PF_SLEEP_CYCLES, before_step=None): + """Warm up eagerly (and capture), reset the cache, then decode ``inputs`` (per step, one + ``[bs, hidden]`` input per layer); ``sleep`` cycles between layers, ``before_step(i)`` before step i.""" + from flashlib.kernels.slot_cache import Stat + + _fresh_ctx(_batch=SimpleNamespace(is_prefill=False)) + + def step(xs): + outs = [] + for layer, (moe, x) in enumerate(zip(moes, xs)): + if layer and sleep: + torch.cuda._sleep(sleep) + outs.append(moe.forward(x.clone())) + return outs + + static = [torch.randn(bs, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16) for _ in moes] + step(static) # the eager warm-up the graph runner also does + if graph: + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + outs = step(static) + cache.reset() + cache.reset_stats() + if before_decode is not None: + before_decode() + results = [] + for i, xs in enumerate(inputs): + if before_step is not None: + before_step(i) + if graph: + for s, x in zip(static, xs): + s.copy_(x) + g.replay() + else: + outs = step(xs) + if cache.audit is not None: + cache.audit.end_decode_step(cache) # where the engine calls it: after the forward + results.append([o.clone() for o in outs]) + torch.cuda.synchronize() + state = [t.clone() for t in (cache.slot_for_id, cache.id_of_slot, cache.usage, *cache.bank_views())] + counters = None if cache.prefetch is None else cache.prefetch.counters.cpu() + return results, state, counters, cache.lru_stats[:, Stat.MISS].cpu() + + +def _prefetch_inputs(bs, steps=8, chain=False, seed=0): + gen = torch.Generator(device="cuda").manual_seed(seed * 10 + bs) + + def x(): + return torch.randn(bs, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16, generator=gen) * 0.5 + + if chain: # every layer sees the same input, so the lookahead equals the next router + return [[x()] * _PF_LAYERS for _ in range(steps)] + return [[x() for _ in range(_PF_LAYERS)] for _ in range(steps)] + + +@requires_cuda +@pytest.mark.parametrize("kind", ["bf16", "nvfp4"]) +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("overlap,policy", [(False, "rule"), (True, "rule"), (True, "lru")], ids=["serial", "overlap", "overlap-lru"]) +def test_moe_prefetch_measure_is_bitwise_identical(kind, graph, bs, overlap, policy, monkeypatch): + """FREETOKEN_MOE_PREFETCH=measure predicts and counts beside the decode (with and without + FREETOKEN_MOE_COPY_OVERLAP's side stream); the MoE outputs, the slot map, the usage clock and + the slot contents must not change by a single bit.""" + banks = _prefetch_banks(kind) + inputs = _prefetch_inputs(bs) + off = _prefetch_run(*_prefetch_stack(kind, banks, "off", monkeypatch, overlap=overlap, policy=policy), bs, graph, inputs) + on = _prefetch_run(*_prefetch_stack(kind, banks, "measure", monkeypatch, overlap=overlap, policy=policy), bs, graph, inputs) + + assert off[2] is None and int(off[3].sum()) > 0 and torch.equal(on[3], off[3]) + for got, want in zip(on[0], off[0]): + for a, b in zip(got, want): + assert torch.isfinite(b.float()).all() + assert torch.equal(a, b) + for a, b in zip(on[1], off[1]): + assert torch.equal(a, b) + from freetoken.moe.prefetch import CALLS, ISSUED, ROWS + + counters = on[2] + assert counters[0].tolist() == [0] * counters.shape[1], "layer 0 has no predecessor" + assert counters[1:, CALLS].tolist() == [len(inputs)] * (_PF_LAYERS - 1) + assert counters[1:, ROWS].tolist() == [len(inputs) * bs] * (_PF_LAYERS - 1) + assert int(counters[1:, ISSUED].sum()) > 0 + + +def _reference_counters(selects, ensures, budget_override, k): + """CPU counters from the spied predictor logits and the residency its select read, and each + target layer's routing and residency at its ensure.""" + from freetoken.moe.prefetch import NUM_COLS, default_budget + from tests.moe.ref_prefetch import ref_count, ref_select + + config = parse_config(toy_hf_config(_PF_LAYERS)) + counters = torch.zeros((_PF_LAYERS, NUM_COLS), dtype=torch.int64) + for layer in range(1, _PF_LAYERS): + assert len(selects[layer]) == len(ensures[layer]) > 0 + want_budget = budget_override or default_budget(config.is_linear_layer(layer)) + for (logits, budget, seen), (raw, resident) in zip(selects[layer], ensures[layer]): + assert budget == want_budget + sel, res = ref_select(logits.tolist(), seen.tolist(), k, budget) + routed = raw.view(-1).tolist() + misses = len({e for e in routed if resident[e] < 0}) + counters[layer] += torch.tensor(ref_count(sel, res, routed, misses, logits.shape[0])) + return counters + + +@requires_cuda +@pytest.mark.parametrize("kind", ["bf16", "nvfp4"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("budget", [0, 5]) +def test_moe_prefetch_counters_match_reference(kind, bs, budget, monkeypatch): + """Eager counters equal a CPU reference built from the predictor's own logits and the residency + its select read, and each layer's routing and pre-ensure residency; the captured graph counts + the same. The predictor runs on its dedicated stream, the count on the compute stream.""" + import freetoken.moe.prefetch as pf_mod + + banks = _prefetch_banks(kind) + inputs = _prefetch_inputs(bs, seed=1) + k = 6 + moes, cache = _prefetch_stack(kind, banks, "measure", monkeypatch, k=k, budget=budget) + selects = {layer: [] for layer in range(_PF_LAYERS)} + ensures = {layer: [] for layer in range(_PF_LAYERS)} + streams = {"select": set(), "count": set()} + real_select, real_count, real_ensure = pf_mod.lookahead_select, pf_mod.prefetch_count, cache.ensure_experts + + def spy_select(logits, resident, sel, res, *, k, budget): + # the select reads residency beside layer L-1's ensure: syncing on both sides pins what it saw + torch.cuda.synchronize() + layer = (sel.data_ptr() - cache.prefetch.sel.data_ptr()) // cache.prefetch.sel[0].nbytes + selects[layer].append((logits.float().cpu(), budget, resident.cpu())) + streams["select"].add(torch.cuda.current_stream().cuda_stream) + real_select(logits, resident, sel, res, k=k, budget=budget) + torch.cuda.synchronize() + + def spy_count(*args, **kwargs): + streams["count"].add(torch.cuda.current_stream().cuda_stream) + return real_count(*args, **kwargs) + + def spy_ensure(layer_id, expert_ids, **kwargs): + torch.cuda.synchronize() # the prediction forked after the previous ensure has landed + ensures[layer_id].append((expert_ids.cpu(), cache.slot_for_id[layer_id].cpu())) + return real_ensure(layer_id, expert_ids, **kwargs) + + def install(): + monkeypatch.setattr(pf_mod, "lookahead_select", spy_select) + monkeypatch.setattr(pf_mod, "prefetch_count", spy_count) + monkeypatch.setattr(cache, "ensure_experts", spy_ensure) + + eager = _prefetch_run(moes, cache, bs, False, inputs, before_decode=install) + want = _reference_counters(selects, ensures, budget, k) + assert torch.equal(eager[2], want), (eager[2], want) + assert int(want[:, 1].sum()) > 0, "the sequence must have useful predictions" + assert streams["select"] == {cache.prefetch.stream.cuda_stream} + assert streams["count"] == {torch.cuda.current_stream().cuda_stream} + + monkeypatch.undo() + moes, cache = _prefetch_stack(kind, banks, "measure", monkeypatch, k=k, budget=budget) + graph = _prefetch_run(moes, cache, bs, True, inputs) + assert torch.equal(graph[2], eager[2]) + for got, want_out in zip(graph[0], eager[0]): + for a, b in zip(got, want_out): + assert torch.equal(a, b) + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +def test_moe_prefetch_perfect_lookahead_counts_every_miss(graph, bs, monkeypatch): + """Known answer: when layer L-1's router input equals layer L's, the lookahead is layer L's own + router, so with the budget covering every candidate each demand miss is a useful prediction.""" + from freetoken.moe.prefetch import ISSUED, MISSES, RESIDENT_HITS, USEFUL + + banks = _prefetch_banks("nvfp4") + moes, cache = _prefetch_stack("nvfp4", banks, "measure", monkeypatch, k=4, budget=8) + _, _, counters, lru_misses = _prefetch_run(moes, cache, bs, graph, _prefetch_inputs(bs, chain=True, seed=2)) + assert int(counters[1:, MISSES].sum()) > 0 + assert counters[1:, USEFUL].tolist() == counters[1:, MISSES].tolist() == lru_misses[1:].tolist() + assert (counters[1:, ISSUED] >= counters[1:, USEFUL]).all() + assert int(counters[1:, RESIDENT_HITS].sum()) > 0, "routed hits among the top candidates" + + +@requires_cuda +def test_moe_prefetch_off_launches_nothing_new(monkeypatch): + """With the flag off the wired lookahead is inert: no prefetcher, and the decode launches + exactly the kernels of an unwired stack (measure adds some, so the check can see them).""" + from torch.profiler import ProfilerActivity, profile + + banks = _prefetch_banks("bf16") + inputs = _prefetch_inputs(1, steps=2) + + def kernels(mode, wire): + moes, cache = _prefetch_stack("bf16", banks, mode, monkeypatch, wire=wire) + _prefetch_run(moes, cache, 1, False, inputs[:1]) # warm up and compile outside the trace + with profile(activities=[ProfilerActivity.CUDA]) as prof: + _prefetch_run(moes, cache, 1, False, inputs) + events = [e for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA] + return cache, [e.name for e in sorted(events, key=lambda e: e.time_range.start)] + + cache, wired_off = kernels("off", True) + assert cache.prefetch is None + _, unwired_off = kernels("off", False) + _, measured = kernels("measure", True) + assert wired_off == unwired_off + assert len(measured) > len(wired_off) + + +@requires_cuda +def test_moe_prefetch_survives_a_rebuild(monkeypatch): + """A runtime rebuild resets the counters and keeps the predictor stream; a recaptured graph + counts like an eager run on the rebuilt cache.""" + banks = _prefetch_banks("bf16") + inputs = _prefetch_inputs(1, seed=3) + moes, cache = _prefetch_stack("bf16", banks, "measure", monkeypatch) + before = _prefetch_run(moes, cache, 1, True, inputs)[2] + stream = cache.prefetch.stream + assert int(before.sum()) > 0 + cache.rebuild(_PF_EXPERTS + 4) + assert int(cache.prefetch.counters.abs().sum()) == 0 and cache.prefetch.stream is stream + graph = _prefetch_run(moes, cache, 1, True, inputs) + eager = _prefetch_run(moes, cache, 1, False, inputs) + assert int(graph[2].sum()) > 0 and torch.equal(graph[2], eager[2]) + for a, b in zip(graph[1], eager[1]): + assert torch.equal(a, b) + + +@requires_cuda +@pytest.mark.parametrize("kind", ["bf16", "nvfp4"]) +@pytest.mark.parametrize("bs", [1, 2]) +def test_moe_prefetch_predicts_layer_l_from_layer_l_minus_1_input(kind, bs, monkeypatch): + """With a distinct input per layer, the prediction for layer L reads exactly layer L-1's router + input through layer L's router, and the select sees those logits and layer L's residency.""" + import freetoken.moe.prefetch as pf_mod + + banks = _prefetch_banks(kind) + inputs = _prefetch_inputs(bs, steps=3, seed=4) + moes, cache = _prefetch_stack(kind, banks, "measure", monkeypatch) + forks, selects = [], [] + real_fork, real_select = cache.prefetch.fork, pf_mod.lookahead_select + + def spy_fork(target, x, gate, resident, budget): + forks.append((target, x.clone(), gate, resident.data_ptr())) + return real_fork(target, x, gate, resident, budget) + + def spy_select(logits, resident, sel, res, *, k, budget): + selects.append(logits.clone()) + return real_select(logits, resident, sel, res, k=k, budget=budget) + + def install(): + monkeypatch.setattr(cache.prefetch, "fork", spy_fork) + monkeypatch.setattr(pf_mod, "lookahead_select", spy_select) + + _prefetch_run(moes, cache, bs, False, inputs, before_decode=install) + assert [f[0] for f in forks] == list(range(1, _PF_LAYERS)) * len(inputs) and len(selects) == len(forks) + for call, ((target, x, gate, resident), logits) in enumerate(zip(forks, selects)): + assert torch.equal(x, inputs[call // (_PF_LAYERS - 1)][target - 1]) + assert gate is moes[target].gate and resident == cache.slot_for_id[target].data_ptr() + assert torch.equal(logits, F.linear(x, moes[target].gate.weight)) + + +# FREETOKEN_MOE_PREFETCH=on: the same stack, with copies. No sleep between layers, so the prefetch +# streams race the compute stream at toy size and only the joins keep the GEMMs on landed bytes. +_PF_ON_STEPS = 200 + + +def _assert_slots_hold_their_experts(cache): + """Lossless: every held slot's bytes equal the host rows of the expert its id names.""" + held = (cache.id_of_slot >= 0).nonzero().view(-1).cpu() + ids = cache.id_of_slot[held.cuda()].long().cpu() + assert held.numel() > 0 + for (per_layer, slots), view in zip(cache.banks, cache.bank_views()): + host = torch.stack(per_layer).flatten(0, 1) + want = host[ids].contiguous().view(torch.uint8) + got = view[held.cuda()].cpu().contiguous().view(torch.uint8) + assert torch.equal(got, want) + + +def _routed_ids(moes, inputs): + """[step][layer] -> the experts each layer routes (the router alone decides; no cache input).""" + from freetoken.moe.fused import fused_topk + + out = [] + for xs in inputs: + row = [] + for moe, x in zip(moes, xs): + _, ids = fused_topk(hidden_states=x, gating_output=moe.gate.forward(x), topk=_PF_TOPK, renormalize=True) + row.append(ids.view(-1).tolist()) + out.append(row) + return out + + +def _garbage(routed, width, kind, seed): + """Per step a [layers, width] candidate override: 'random' ids (and -1), 'adversarial' ones + that mix layer L's routed ids with ids it does not route, duplicates and padding, or 'flood': + every expert of layer L, so a prefetch wants more slots than the unpinned ones.""" + gen = torch.Generator().manual_seed(seed) + steps = [] + for row in routed: + sel = torch.full((_PF_LAYERS, width), -1, dtype=torch.int32) + for layer in range(1, _PF_LAYERS): + if kind == "random": + ids = torch.randint(-1, _PF_EXPERTS, (width,), generator=gen) + elif kind == "flood": + ids = torch.randperm(_PF_EXPERTS, generator=gen)[:width] + else: + hit = list(dict.fromkeys(row[layer])) + miss = [e for e in torch.randperm(_PF_EXPERTS, generator=gen).tolist() if e not in hit][:3] + ids = torch.tensor((hit[:3] + miss + hit[:1] + [-1] + miss[:1]) * 2)[:width] + ids = ids[torch.randperm(ids.numel(), generator=gen)] + sel[layer, : ids.numel()] = ids.to(torch.int32) + steps.append(sel) + return steps + + +def _on_stack(kind, banks, monkeypatch, *, policy, overlap, k=6, budget=0): + return _prefetch_stack(kind, banks, "on", monkeypatch, k=k, budget=budget, overlap=overlap, policy=policy) + + +def _with_override(cache, steps, **delays): + """before_step hook that loads step i's candidate override into the prefetcher's static buffer.""" + pf = cache.prefetch + pf.sel_override = torch.full_like(pf.sel, -1) + for name, cycles in delays.items(): + setattr(pf, name, cycles) + + def before_step(i): + pf.sel_override.copy_(steps[i][:, : pf.sel.shape[1]]) + + return before_step + + +def _assert_same_outputs(got, want): + for a_step, b_step in zip(got, want): + for a, b in zip(a_step, b_step): + assert torch.isfinite(b.float()).all() + assert torch.equal(a, b) + + +@requires_cuda +@pytest.mark.parametrize("kind", ["bf16", "nvfp4"]) +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("policy", ["rule", "lru"]) +@pytest.mark.parametrize("overlap", [False, True], ids=["serial", "overlap"]) +def test_moe_prefetch_on_is_bitwise_identical(kind, graph, bs, policy, overlap, monkeypatch): + """FREETOKEN_MOE_PREFETCH=on copies predicted experts into low-priority slots beside the decode: + over 200 steps of a cache that evicts every step the MoE outputs are those of prefetch off to + the bit, every held slot holds its expert's bytes, and the prefetches were used.""" + from freetoken.moe.prefetch import CALLS, COPIED, ISSUED, USEFUL + + banks = _prefetch_banks(kind) + inputs = _prefetch_inputs(bs, steps=_PF_ON_STEPS, seed=5) + moes, cache = _prefetch_stack(kind, banks, "off", monkeypatch, overlap=overlap, policy=policy) + off = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0) + _assert_slots_hold_their_experts(cache) + moes, cache = _on_stack(kind, banks, monkeypatch, policy=policy, overlap=overlap) + on = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0) + _assert_same_outputs(on[0], off[0]) + _assert_slots_hold_their_experts(cache) + stats = cache.prefetch.stats.cpu() + assert stats[0].tolist() == [0] * stats.shape[1], "layer 0 is never prefetched" + assert stats[1:, CALLS].tolist() == [_PF_ON_STEPS] * (_PF_LAYERS - 1) + assert int(stats[:, ISSUED].sum()) > 0 and int(stats[:, USEFUL].sum()) > 0 and int(stats[:, COPIED].sum()) > 0 + assert int(on[3].sum()) < int(off[3].sum()), "useful prefetches remove demand misses" + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("policy", ["rule", "lru"]) +@pytest.mark.parametrize("garbage", ["random", "adversarial"]) +def test_moe_prefetch_on_survives_a_garbage_predictor(graph, bs, policy, garbage, monkeypatch): + """Whatever the candidates (random ids and padding, or the next layer's routed ids mixed with ids + it will not route, duplicates and -1), prefetch on changes no output bit and no slot's bytes.""" + from freetoken.moe.prefetch import ISSUED, USEFUL + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(bs, steps=_PF_ON_STEPS, seed=6) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, policy=policy) + off = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0) + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True, budget=8) + steps = _garbage(_routed_ids(moes, inputs), 12, garbage, seed=bs) + on = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0, before_step=_with_override(cache, steps)) + _assert_same_outputs(on[0], off[0]) + _assert_slots_hold_their_experts(cache) + stats = cache.prefetch.stats.cpu() + assert int(stats[:, ISSUED].sum()) > 0 and int(stats[:, USEFUL].sum()) > 0 + assert int(stats[:, USEFUL].sum()) < int(stats[:, ISSUED].sum()), "some prefetches must be wrong" + + +# ~0.5 ms at the RTX 5080's clock: far longer than a toy layer, so every delayed stage is late +_PF_DELAY_CYCLES = 1_000_000 + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("policy", ["rule", "lru"]) +@pytest.mark.parametrize("delay", ["copy", "predict", "both"]) +def test_moe_prefetch_on_joins_hold_when_the_prefetch_streams_run_late(graph, bs, policy, delay, monkeypatch): + """Race test: a sleep on the prefetch copy stream (every copy lands after its layer's ensure) or on + the predictor stream (prefetch_ensure runs long after the fork) must not change a bit. The + candidates include each layer's routed ids, so GEMMs do read prefetched slots. The joins also + make the slot maps timing-free: the delayed run ends with the undelayed run's maps and counts.""" + from freetoken.moe.prefetch import COPIED, LATE, NUM_STAT_COLS, USEFUL + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(bs, steps=_PF_ON_STEPS, seed=7) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, policy=policy) + off = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0) + runs = [] + for cycles in (0, _PF_DELAY_CYCLES): + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True, budget=8) + steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=bs + 10) + delays = {"delay_copy_cycles": cycles * (delay in ("copy", "both")), + "delay_predict_cycles": cycles * (delay in ("predict", "both"))} + on = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0, before_step=_with_override(cache, steps, **delays)) + _assert_same_outputs(on[0], off[0]) + _assert_slots_hold_their_experts(cache) + runs.append((on[1][:3], cache.prefetch.stats.cpu())) + (maps, stats), (delayed_maps, delayed_stats) = runs + assert int(stats[:, USEFUL].sum()) > 0 + for a, b in zip(maps, delayed_maps): + assert torch.equal(a, b) + cols = [c for c in range(NUM_STAT_COLS) if c != LATE] + assert torch.equal(stats[:, cols], delayed_stats[:, cols]) + if delay == "copy": + # the late counter sees what the join waited for: the copies were still sleeping at ensure(L) + # (all of them under replay; in eager a loaded host can take longer than the sleep per layer) + late, copied = int(delayed_stats[:, LATE].sum()), int(delayed_stats[:, COPIED].sum()) + assert copied > 0 and (late == copied if graph else late > 0) + + +@requires_cuda +def test_moe_prefetch_race_test_catches_a_missing_copy_join(monkeypatch): + """Negative control for the race test: without the compute stream's wait on the prefetch copy, + GEMMs read prefetched slots before their bytes land and the outputs change.""" + from freetoken.moe.prefetch import ExpertPrefetcher + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(1, steps=40, seed=7) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True) + off = _prefetch_run(moes, cache, 1, False, inputs, sleep=0) + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy="rule", overlap=True, budget=8) + + def no_wait(self, layer_id): + self._inflight[layer_id] = None + + monkeypatch.setattr(ExpertPrefetcher, "join_copy", no_wait) + steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=11) + on = _prefetch_run(moes, cache, 1, False, inputs, sleep=0, + before_step=_with_override(cache, steps, delay_copy_cycles=_PF_DELAY_CYCLES)) + assert any(not torch.equal(a, b) for got, want in zip(on[0], off[0]) for a, b in zip(got, want)) + + +def _replay_on_reference(policy, calls, rows): + """Replay the spied prefetch_ensure / ensure sequence on the CPU reference; per-layer stats.""" + from freetoken.moe.prefetch import CALLS, COPIED, ISSUED, MISSES, ROWS, USEFUL + from freetoken.moe.scored_ensure import POLICY_IDS + from tests.moe.ref_scored_cache import RefScoredCache + + ref = RefScoredCache(_PF_LAYERS, _PF_EXPERTS, _PF_EXPERTS, POLICY_IDS[policy]) + stats = torch.zeros((_PF_LAYERS, 8), dtype=torch.int64) + planned = {} + for call in calls: + if call[0] == "prefetch": + _, layer, sel, budget = call + src, _ = ref.prefetch(layer, sel.tolist(), budget) + stats[layer, ISSUED] += src.size + planned[layer] = src.size + else: + _, layer, ids, logits = call + _, src, _ = ref.ensure(layer, ids.view(-1).numpy(), bump_tok=layer == 0, lowpri=True, + logits=None if logits is None else logits.float().numpy(), thr=0.25) + if layer in planned: + stats[layer, [USEFUL, CALLS, ROWS, MISSES, COPIED]] += torch.tensor( + [ref.useful, 1, rows, src.size, int(planned.pop(layer) > 0)]) + return ref, stats + + +@requires_cuda +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("policy", ["rule", "lru"]) +def test_moe_prefetch_on_counters_match_reference(bs, policy, monkeypatch): + """issued / useful / calls / rows / misses / copied equal a CPU replay of the spied candidates + and routings, the slot maps end identical, late is 0 when each ensure waits for the copies + (known answer), and with the candidates fixed a captured graph counts the same as eager. + prefetch_ensure and the demand copy run on the compute stream, the slim prefetch copy on the + copy stream.""" + import freetoken.kernel.fast_index_copy as fic + from freetoken.moe.prefetch import LATE, NUM_STAT_COLS, RESIDENT_HITS + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(bs, steps=60, seed=8) + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True) + calls, streams = [], {"pf_ensure": set(), "slim": set(), "demand": set()} + real_pf, real_ensure = cache.prefetch_ensure, cache.ensure_experts + real_slim, real_multi = fic.fast_index_copy_multi_slim_jit, fic.fast_index_copy_multi_jit + + def spy_pf(layer, query, *args, **kw): + streams["pf_ensure"].add(torch.cuda.current_stream().cuda_stream) + calls.append(("prefetch", layer, query.clone(), kw.get("budget"))) # on the compute stream after its wait on the select + return real_pf(layer, query, *args, **kw) + + def spy_ensure(layer, ids, **kw): + torch.cuda.synchronize() # the copies forked before this ensure have landed: late must be 0 + logits = kw.get("router_logits") + calls.append(("ensure", layer, ids.clone(), None if logits is None or policy != "rule" else logits.clone())) + return real_ensure(layer, ids, **kw) + + def spy_slim(*args): + streams["slim"].add(torch.cuda.current_stream().cuda_stream) + return real_slim(*args) + + def spy_multi(*args): + streams["demand"].add(torch.cuda.current_stream().cuda_stream) + return real_multi(*args) + + def install(): + monkeypatch.setattr(cache, "prefetch_ensure", spy_pf) + monkeypatch.setattr(cache, "ensure_experts", spy_ensure) + monkeypatch.setattr(fic, "fast_index_copy_multi_slim_jit", spy_slim) + monkeypatch.setattr(fic, "fast_index_copy_multi_jit", spy_multi) + + eager = _prefetch_run(moes, cache, bs, False, inputs, before_decode=install, sleep=0) + torch.cuda.synchronize() + calls = [c if c[0] == "prefetch" else (c[0], c[1], c[2].cpu(), None if c[3] is None else c[3].cpu()) for c in calls] + calls = [(c[0], c[1], c[2].cpu(), c[3]) if c[0] == "prefetch" else c for c in calls] + ref, want = _replay_on_reference(policy, calls, bs) + got = cache.prefetch.stats.cpu() + assert got.shape == (_PF_LAYERS, NUM_STAT_COLS) + assert int(got[:, LATE].sum()) == 0 and int(got[:, RESIDENT_HITS].sum()) == 0 + assert torch.equal(got, want), (got, want) + assert int(want[:, 1].sum()) > 0, "the sequence must have useful prefetches" + assert torch.equal(cache.slot_for_id.view(-1).cpu().long(), torch.from_numpy(ref.slot_of_id)) + assert torch.equal(cache.id_of_slot.cpu().long(), torch.from_numpy(ref.id_of_slot)) + assert torch.equal(cache.usage.cpu(), torch.from_numpy(ref.usage)) + assert streams["pf_ensure"] == {torch.cuda.current_stream().cuda_stream} + assert streams["slim"] == {cache.prefetch.copy_stream.cuda_stream} + assert streams["demand"] == {torch.cuda.current_stream().cuda_stream} + + # the select reads residency beside layer L-1's ensure, so two runs may pick different candidates: + # with the candidates fixed, a captured graph must install, count and end like eager + monkeypatch.undo() + runs = [] + for graph in (False, True): + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True) + steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=bs + 8) + run = _prefetch_run(moes, cache, bs, graph, inputs, sleep=_PF_SLEEP_CYCLES, before_step=_with_override(cache, steps)) + _assert_same_outputs(run[0], eager[0]) + runs.append((run[1][:3], cache.prefetch.stats.cpu())) + cols = [c for c in range(NUM_STAT_COLS) if c != LATE] + (maps, stats), (graph_maps, graph_stats) = runs + for a, b in zip(maps, graph_maps): + assert torch.equal(a, b) + assert torch.equal(graph_stats[:, cols], stats[:, cols]) and int(stats[:, 1].sum()) > 0 + + +@requires_cuda +def test_moe_prefetch_on_rebuild_and_reset_clear_the_prefetch_state(monkeypatch): + """A rebuild keeps both prefetch streams, the events and plan buffers, and drops the counters, + the in-flight state and every low-priority slot; a recaptured graph then matches eager and off. + A reset (as after a capture) also drops the counters and low-priority slots.""" + banks = _prefetch_banks("bf16") + inputs = _prefetch_inputs(1, steps=40, seed=9) + moes, cache = _on_stack("bf16", banks, monkeypatch, policy="rule", overlap=True) + _prefetch_run(moes, cache, 1, True, inputs, sleep=0) + pf = cache.prefetch + streams, plan = (pf.stream, pf.copy_stream), (pf.pf_slots, pf.pf_src, pf.pf_num, pf.pf_ready) + assert int(pf.stats.sum()) > 0 and bool(((cache.id_of_slot >= 0) & (cache.usage == 0)).any()) + pf._inflight[2] = torch.zeros(1) + cache.rebuild(_PF_EXPERTS + 4) + assert (pf.stream, pf.copy_stream) == streams and (pf.pf_slots, pf.pf_src, pf.pf_num, pf.pf_ready) == plan + assert int(pf.stats.abs().sum()) == 0 and int(pf.totals.abs().sum()) == 0 and pf._inflight == [None] * _PF_LAYERS + assert not bool(((cache.id_of_slot >= 0) & (cache.usage == 0)).any()) + graph = _prefetch_run(moes, cache, 1, True, inputs, sleep=0) + eager = _prefetch_run(moes, cache, 1, False, inputs, sleep=0) + _assert_same_outputs(graph[0], eager[0]) + _assert_slots_hold_their_experts(cache) + assert int(pf.stats.sum()) > 0 + cache.reset() + assert int(pf.stats.abs().sum()) == 0 and not bool(((cache.id_of_slot >= 0) & (cache.usage == 0)).any()) + moes_off, cache_off = _prefetch_stack("bf16", banks, "off", monkeypatch, overlap=True) + cache_off.rebuild(_PF_EXPERTS + 4) + off = _prefetch_run(moes_off, cache_off, 1, False, inputs, sleep=0) + _assert_same_outputs(eager[0], off[0]) + + +@requires_cuda +def test_moe_prefetch_on_skips_batches_past_the_select_tile(monkeypatch): + """An eager decode wider than MAX_ROWS forks nothing: no join, no count, the same outputs.""" + from freetoken.moe.prefetch import MAX_ROWS + + banks = _prefetch_banks("bf16") + bs = MAX_ROWS + 1 + inputs = _prefetch_inputs(bs, steps=6, seed=10) + off = _prefetch_run(*_prefetch_stack("bf16", banks, "off", monkeypatch), bs, False, inputs, sleep=0) + moes, cache = _on_stack("bf16", banks, monkeypatch, policy="rule", overlap=True) + on = _prefetch_run(moes, cache, bs, False, inputs, sleep=0) + _assert_same_outputs(on[0], off[0]) + assert int(cache.prefetch.stats.abs().sum()) == 0 and cache.prefetch._inflight == [None] * _PF_LAYERS + + +def _stall_hooks(monkeypatch): + """Sleeps (cycles, read at enqueue time) on the compute stream right before each decode GEMM, i.e. + after issue_copy(L+1), and on the stream that runs the shared expert gate (the side stream).""" + import freetoken.models.qwen4_exp.moe as q4moe + from freetoken.layers.moe import OffloadMoELayer + + stall = {"gemm": 0, "shared": 0} + real_gemm, real_gate = OffloadMoELayer._expert_gemm, q4moe.shared_gate_sigmoid + + def gemm(self, cache, *args, is_prefill, **kw): + if stall["gemm"] and not is_prefill: + torch.cuda._sleep(stall["gemm"]) + return real_gemm(self, cache, *args, is_prefill=is_prefill, **kw) + + def gate(*args, **kw): + if stall["shared"]: + torch.cuda._sleep(stall["shared"]) + return real_gate(*args, **kw) + + monkeypatch.setattr(OffloadMoELayer, "_expert_gemm", gemm) + monkeypatch.setattr(q4moe, "shared_gate_sigmoid", gate) + return stall + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("policy", ["rule", "lru"]) +@pytest.mark.parametrize("where", ["gemm", "shared"]) +@pytest.mark.parametrize("cands", ["adversarial", "flood"]) +def test_moe_prefetch_on_holds_when_the_compute_side_stalls(graph, bs, policy, where, cands, monkeypatch): + """Race test from the other side: a sleep on the compute stream between issue_copy(L+1) and + GEMM(L) lets copy(L+1) land before GEMM(L) reads, so a prefetch into any of layer L's routed + slots changes outputs deterministically; a sleep on the shared expert's side stream must not + matter either. 'flood' asks for all 16 experts a layer, more than the unpinned slots, so even + LRU's recency order would reach the pinned ones. Bit-identical to off, and the same final maps + as the unstalled run.""" + k, width = (8, 16) if cands == "flood" else (6, 12) + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(bs, steps=80, seed=21) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, policy=policy) + off = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0) + stall = _stall_hooks(monkeypatch) + runs = [] + for cycles in (0, _PF_DELAY_CYCLES): + stall[where] = cycles + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True, k=k, budget=width) + steps = _garbage(_routed_ids(moes, inputs), width, cands, seed=bs + 20) + on = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0, before_step=_with_override(cache, steps)) + _assert_same_outputs(on[0], off[0]) + _assert_slots_hold_their_experts(cache) + runs.append(on[1][:3]) + for a, b in zip(*runs): + assert torch.equal(a, b) + + +def _assert_held_slots_hold_their_experts(cache): + """_assert_slots_hold_their_experts, but a prefill may leave no slot held.""" + if bool((cache.id_of_slot >= 0).any()): + _assert_slots_hold_their_experts(cache) + + +def _mixed_run(moes, cache, bs, graph, schedule, before_step=None, late=None, settle=True): + """Warm up (and capture) the decode like _prefetch_run, reset, then run ``schedule``: ("decode", xs) + is one decode step (a replay under ``graph``), ("prefill", xs) an eager prefill of xs's tokens. + ``late=(cycles, steps)`` sleeps that long on the predictor stream in those decode steps (a second + graph captured with the sleep, in the same pool, replays them); ``settle`` syncs and checks the + slots after each prefill, else the next decode is enqueued right behind it.""" + ctx = _fresh_ctx(_batch=SimpleNamespace(is_prefill=False)) + + def step(xs): + return [moe.forward(x.clone()) for moe, x in zip(moes, xs)] + + static = [torch.randn(bs, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16) for _ in moes] + step(static) + if graph: + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + graphs = {False: (g, step(static))} + if late is not None: + cache.prefetch.delay_predict_cycles = late[0] + g_late = torch.cuda.CUDAGraph() + with torch.cuda.graph(g_late, pool=g.pool()): + graphs[True] = (g_late, step(static)) + cache.prefetch.delay_predict_cycles = 0 + cache.reset() + results = [] + for i, (phase, xs) in enumerate(schedule): + if before_step is not None: + before_step(i) + ctx._batch.is_prefill = phase == "prefill" + slow = late is not None and i in late[1] + if phase == "decode" and graph: + for s, x in zip(static, xs): + s.copy_(x) + g, outs = graphs[slow] + g.replay() + else: + if late is not None: + cache.prefetch.delay_predict_cycles = late[0] if slow else 0 + outs = step(xs) + if phase == "decode" and cache.audit is not None: + cache.audit.end_decode_step(cache) + results.append([o.clone() for o in outs]) + if phase == "prefill" and settle: + torch.cuda.synchronize() + _assert_held_slots_hold_their_experts(cache) + torch.cuda.synchronize() + _assert_held_slots_hold_their_experts(cache) + return results + + +_PREFILL_PATHS = { + # (cache_size, prefill tokens, cache kwargs, FREETOKEN_MOE_SMALL_PREFILL_TOKENS) + "materialize": (_PF_EXPERTS + 4, 24, {}, 0), + "overlap": (2 * _PF_EXPERTS + 4, 24, {"prefill_overlap": True}, 0), + "overlap_d2d": (2 * _PF_EXPERTS + 4, 24, {"prefill_overlap": True, "prefill_hit_d2d": True}, 0), + "small": (_PF_EXPERTS, 3, {}, 8), +} + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("policy", ["rule", "lru"]) +@pytest.mark.parametrize("path", list(_PREFILL_PATHS)) +def test_moe_prefetch_on_survives_prefills_between_decodes(graph, bs, policy, path, monkeypatch): + """Prefills between decode steps rewrite slots behind the prefetch's back: the whole-layer + materialize installs slots [0, E), the overlap double buffers invalidate and stream [0, 2E) + (with the hit-D2D gather reading resident slots), and small prefills ensure their experts + under the low-priority keys. Decode, prefill, decode, prefill, decode: every output of prefetch + on equals off to the bit, and after each prefill every held slot holds its expert.""" + import freetoken.layers.moe as layers_moe + + cache_size, tokens, cache_kw, small = _PREFILL_PATHS[path] + monkeypatch.setattr(layers_moe, "_SMALL_PREFILL_TOKENS", small) + banks = _prefetch_banks("nvfp4") + decode = _prefetch_inputs(bs, steps=18, seed=31) + gen = torch.Generator(device="cuda").manual_seed(32) + prefills = [[torch.randn(tokens, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16, generator=gen) * 0.5 + for _ in range(_PF_LAYERS)] for _ in range(2)] + schedule = ([("decode", xs) for xs in decode[:6]] + [("prefill", prefills[0])] + + [("decode", xs) for xs in decode[6:12]] + [("prefill", prefills[1])] + + [("decode", xs) for xs in decode[12:]]) + kw = dict(overlap=True, policy=policy, cache_size=cache_size, **cache_kw) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, **kw) + off = _mixed_run(moes, cache, bs, graph, schedule) + moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, budget=8, **kw) + steps = _garbage(_routed_ids(moes, [xs for _, xs in schedule]), 12, "adversarial", seed=bs + 30) + on = _mixed_run(moes, cache, bs, graph, schedule, before_step=_with_override(cache, steps)) + _assert_same_outputs(on, off) + if path == "overlap_d2d": + assert cache.prefill_hit_rows > 0, "the hit-D2D gather must have read resident slots" + assert int(cache.prefetch.stats[:, 0].sum()) > 0 + + +def _record_compute_streams(monkeypatch, moes): + """The streams the MoE blocks are entered on (the default stream, and the capture stream under a + graph): the compute streams of the decode.""" + compute = set() + for moe in moes: + real = moe.forward + + def forward(x, real=real): + compute.add(torch.cuda.current_stream().cuda_stream) + return real(x) + + monkeypatch.setattr(moe, "forward", forward) + return compute + + +def _spy_map_writers(monkeypatch): + """(writer, stream) of every launch that writes the slot maps or their mirrors, at enqueue time.""" + import freetoken.kernel.triton.moe as triton_moe + import freetoken.moe.offload_kernels as ok + import freetoken.moe.scored_ensure as se + + seen = [] + + def spy(name, real): + def launch(*args, **kwargs): + seen.append((name(kwargs) if callable(name) else name, torch.cuda.current_stream().cuda_stream)) + return real(*args, **kwargs) + + return launch + + # scored_ensure is the kernel of both the demand ensures and prefetch_ensure (prefetch=True) + monkeypatch.setattr(se, "scored_ensure", spy(lambda kw: "prefetch_ensure" if kw.get("prefetch") else "ensure", + se.scored_ensure)) + monkeypatch.setattr(ok, "lru_ensure", spy("ensure", ok.lru_ensure)) + monkeypatch.setattr(ok, "materialize_layer", spy("materialize", ok.materialize_layer)) + monkeypatch.setattr(ok, "reset_cache", spy("reset", ok.reset_cache)) + monkeypatch.setattr(triton_moe, "invalidate_prefill_slots", spy("invalidate", triton_moe.invalidate_prefill_slots)) + return seen + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("path", list(_PREFILL_PATHS)) +def test_moe_prefetch_on_writes_the_slot_maps_only_on_the_compute_stream(graph, bs, path, monkeypatch): + """Single writer: every launch that writes the slot maps (the demand and prefetch ensures, prefill + ensures, the whole-layer materialize, reset) runs on the stream the MoE blocks run on (the capture + stream under a graph), never on the predictor, prefetch copy or shared-expert side stream. The only + other writer is the overlap prefill's buffer invalidation, on the prefill copy stream it fences. + The predictor's select reads no cache state.""" + import freetoken.layers.moe as layers_moe + import freetoken.moe.prefetch as pf_mod + from freetoken.moe.prefetch import ISSUED + + cache_size, tokens, cache_kw, small = _PREFILL_PATHS[path] + monkeypatch.setattr(layers_moe, "_SMALL_PREFILL_TOKENS", small) + banks = _prefetch_banks("nvfp4") + decode = _prefetch_inputs(bs, steps=6, seed=95) + gen = torch.Generator(device="cuda").manual_seed(96) + prefill = [torch.randn(tokens, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16, generator=gen) * 0.5 + for _ in range(_PF_LAYERS)] + schedule = [("decode", xs) for xs in decode[:3]] + [("prefill", prefill)] + [("decode", xs) for xs in decode[3:]] + moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, budget=8, overlap=True, cache_size=cache_size, + **cache_kw) + steps = _garbage(_routed_ids(moes, [xs for _, xs in schedule]), 12, "adversarial", seed=bs + 97) + seen = _spy_map_writers(monkeypatch) + compute = _record_compute_streams(monkeypatch, moes) + selects, real_select = [], pf_mod.lookahead_select + + def spy_select(logits, resident, *args, **kwargs): + selects.append((resident.data_ptr(), torch.cuda.current_stream().cuda_stream)) + return real_select(logits, resident, *args, **kwargs) + + monkeypatch.setattr(pf_mod, "lookahead_select", spy_select) + _mixed_run(moes, cache, bs, graph, schedule, before_step=_with_override(cache, steps)) + pf = cache.prefetch + assert not compute & {pf.stream.cuda_stream, pf.copy_stream.cuda_stream, cache.decode_copy_stream.cuda_stream} + assert len(compute) == 1 + graph + writers = {name for name, _ in seen} + want = {"ensure", "prefetch_ensure", "reset"} | {"materialize": {"materialize"}, "overlap": {"invalidate"}, + "overlap_d2d": {"invalidate"}, "small": set()}[path] + assert writers == want, writers + off_compute = [(name, stream) for name, stream in seen if stream not in compute] + fence = {cache.prefill_copy_stream.cuda_stream} if cache.prefill_copy_stream is not None else set() + assert all(name == "invalidate" and stream in fence for name, stream in off_compute), off_compute + assert {stream for name, stream in seen if name == "prefetch_ensure"} == compute # eager, and the capture + assert int(pf.stats[:, ISSUED].sum()) > 0 + # after the writer checks, which are what a predictor-side install fails + assert selects and set(selects) == {(pf._none_resident.data_ptr(), pf.stream.cuda_stream)} + + +# ~4 ms at the RTX 5080's clock: a select far later than its whole layer +_PF_LATE_CYCLES = 8 * _PF_DELAY_CYCLES + + +def _step_starts(times): + """before_step hook recording a timed event on the compute stream at the start of each step.""" + def before_step(i): + times.append(torch.cuda.Event(enable_timing=True)) + times[-1].record() + + return before_step + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("policy", ["rule", "lru"]) +def test_moe_prefetch_on_waits_for_a_late_select(graph, bs, policy, monkeypatch): + """The compute stream waits for the select before it installs, however late the predictor runs: + with every select ~4 ms late (far past its whole layer) a step lasts at least its three late + selects, and the outputs equal off's and the maps and counts (all but LATE) those of the run + without the delay.""" + from freetoken.moe.prefetch import LATE, NUM_STAT_COLS, USEFUL + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(bs, steps=24, seed=100) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, policy=policy) + off = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0) + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + torch.cuda._sleep(_PF_LATE_CYCLES) + end.record() + end.synchronize() + late_ms = start.elapsed_time(end) + runs = [] + for cycles in (0, _PF_LATE_CYCLES): + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True, budget=8) + steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=bs + 101) + load, times = _with_override(cache, steps, delay_predict_cycles=cycles), [] + mark = _step_starts(times) + + def before_step(i, load=load, mark=mark): + mark(i) + load(i) + + on = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0, before_step=before_step) + _assert_same_outputs(on[0], off[0]) + _assert_slots_hold_their_experts(cache) + step_ms = sorted(a.elapsed_time(b) for a, b in zip(times, times[1:])) + runs.append((on[1][:3], cache.prefetch.stats.cpu(), step_ms[len(step_ms) // 2])) + (maps, stats, fast_ms), (late_maps, late_stats, slow_ms) = runs + assert int(stats[:, USEFUL].sum()) > 0 + for a, b in zip(maps, late_maps): + assert torch.equal(a, b) + cols = [c for c in range(NUM_STAT_COLS) if c != LATE] + assert torch.equal(stats[:, cols], late_stats[:, cols]) + # layers 1-3 are predicted; each prediction forks after the previous one was waited for + assert slow_ms >= 0.9 * (_PF_LAYERS - 1) * late_ms > 4 * fast_ms, (slow_ms, late_ms, fast_ms) + + +@requires_cuda +def test_moe_prefetch_on_late_select_without_the_wait_changes_the_maps_but_no_output(monkeypatch): + """Negative control for the test above (eager: a capture needs the predictor joined): without the + compute stream's wait, prefetch_ensure installs stale candidates, so the maps differ from the + undelayed run's, yet every output still equals off's: the slot maps have one writer, so a wrong + prediction only costs a copy.""" + from freetoken.moe.prefetch import ExpertPrefetcher + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(1, steps=24, seed=100) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True) + off = _prefetch_run(moes, cache, 1, False, inputs, sleep=0) + + def no_wait(self, target): + """install without the compute stream's wait on the select.""" + if target is None or self._inflight[target] is None: + return None + self.cache.prefetch_ensure( + target, self.sel[target], self.pf_slots[target], self.pf_src[target], + self.pf_num[target], self.stats[target], self.pf_ready[target], budget=self._budget[target], + ) + self._installed[target] = True + return target + + runs = [] + for skip in (False, True): + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy="rule", overlap=True, budget=8) + if skip: + monkeypatch.setattr(ExpertPrefetcher, "install", no_wait) + steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=102) + on = _prefetch_run(moes, cache, 1, False, inputs, sleep=0, + before_step=_with_override(cache, steps, delay_predict_cycles=_PF_LATE_CYCLES * skip)) + _assert_same_outputs(on[0], off[0]) + _assert_slots_hold_their_experts(cache) + runs.append(on[1][:3]) + assert any(not torch.equal(a, b) for a, b in zip(*runs)) + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("path", list(_PREFILL_PATHS)) +def test_moe_prefetch_on_late_select_right_after_a_prefill(graph, bs, path, monkeypatch): + """The suspected production trigger: a decode step enqueued right behind a prefill (which + rewrote slots behind the prefetch's back, the overlap buffers' [0, 2E) included) with the + predictor ~4 ms late. Every output equals off's to the bit, the maps equal those of the same run + with no late step, every held slot holds its expert and the decode verifier finds nothing.""" + import freetoken.layers.moe as layers_moe + from freetoken.moe.verify import CHECKS + + cache_size, tokens, cache_kw, small = _PREFILL_PATHS[path] + monkeypatch.setattr(layers_moe, "_SMALL_PREFILL_TOKENS", small) + banks = _prefetch_banks("nvfp4") + decode = _prefetch_inputs(bs, steps=15, seed=105) + gen = torch.Generator(device="cuda").manual_seed(106) + prefills = [[torch.randn(tokens, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16, generator=gen) * 0.5 + for _ in range(_PF_LAYERS)] for _ in range(2)] + schedule = ([("decode", xs) for xs in decode[:5]] + [("prefill", prefills[0])] + + [("decode", xs) for xs in decode[5:10]] + [("prefill", prefills[1])] + + [("decode", xs) for xs in decode[10:]]) + after_prefill = {i + 1 for i, (phase, _) in enumerate(schedule) if phase == "prefill"} + kw = dict(overlap=True, cache_size=cache_size, **cache_kw) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, **kw) + off = _mixed_run(moes, cache, bs, graph, schedule) + runs = [] + for late in (None, (_PF_LATE_CYCLES, after_prefill)): + moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, budget=8, verify_mode="meta", **kw) + steps = _garbage(_routed_ids(moes, [xs for _, xs in schedule]), 12, "adversarial", seed=bs + 107) + on = _mixed_run(moes, cache, bs, graph, schedule, before_step=_with_override(cache, steps), late=late, + settle=False) + _assert_same_outputs(on, off) + counts, records = _verify_state(cache) + assert counts[CHECKS] > 0 and counts[CHECKS + 1 :] == [0] * (len(counts) - 1) and records == [], records + runs.append([t.clone() for t in (cache.slot_for_id, cache.id_of_slot, cache.usage)]) + for a, b in zip(*runs): + assert torch.equal(a, b) + if path == "overlap_d2d": + assert cache.prefill_hit_rows > 0, "the hit-D2D gather must have read resident slots" + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("policy", ["rule", "lru"]) +def test_moe_prefetch_on_install_may_evict_a_slot_whose_prefetch_copy_is_in_flight(graph, bs, policy, monkeypatch): + """prefetch_ensure(L) runs on the compute stream before it joins copy(L-1), so it may evict a slot + that plan L-1 installed and whose slim copy is still writing it. The next writer of that slot's + bytes (copy(L) on the same copy stream, or a demand copy after join_copy) is ordered after copy(L-1): + with every copy late and flood candidates the eviction happens (counted on the device, so it holds + under replay too), and still no output bit and no slot's bytes change.""" + from freetoken.moe.prefetch import ExpertPrefetcher + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(bs, steps=40, seed=110) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, policy=policy) + off = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0) + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True, budget=8) + pf = cache.prefetch + evicted_in_flight = torch.zeros((), dtype=torch.int64, device="cuda") + col = torch.arange(pf.pf_slots.shape[1], device="cuda") + real_install = ExpertPrefetcher.install + + def install(self, target): + got = real_install(self, target) + if got is not None and got >= 2: + # on the compute stream right after prefetch_ensure(got), before join_copy(got - 1) + new = torch.where(col < self.pf_num[got], self.pf_slots[got], -2) + old = torch.where(col < self.pf_num[got - 1], self.pf_slots[got - 1], -3) + shared = (new[:, None] == old[None, :]).any() + evicted_in_flight.add_((shared & (self.pf_ready[got - 1][0] == 0)).long()) + return got + + monkeypatch.setattr(ExpertPrefetcher, "install", install) + steps = _garbage(_routed_ids(moes, inputs), 12, "flood", seed=bs + 111) + on = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0, + before_step=_with_override(cache, steps, delay_copy_cycles=_PF_DELAY_CYCLES)) + _assert_same_outputs(on[0], off[0]) + _assert_slots_hold_their_experts(cache) + assert int(evicted_in_flight) > 0, "no install evicted a slot of the previous layer's in-flight copy" + + +@requires_cuda +def test_moe_prefetch_in_flight_eviction_test_catches_an_unordered_next_copy(monkeypatch): + """Negative control for the test above: order each prefetch copy only after its own install (not + after copy(L-1)'s join) and put odd and even layers' copies on two streams, only the even ones + late. A late copy then lands on a slot the next install took over, and the outputs or the held + slots' bytes go wrong.""" + from freetoken.moe.prefetch import ExpertPrefetcher, _mark_ready_kernel, dedicated_stream + from freetoken.moe.slot_audit import PREFETCH_COPY + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(1, steps=40, seed=110) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True) + off = _prefetch_run(moes, cache, 1, False, inputs, sleep=0) + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy="rule", overlap=True, budget=8) + pf = cache.prefetch + streams = [pf.copy_stream, dedicated_stream(torch.device("cuda"), highest_priority=True)] + installed = [torch.cuda.Event() for _ in range(_PF_LAYERS)] + real_install = ExpertPrefetcher.install + + def install(self, target): + got = real_install(self, target) + if got is not None: + installed[got].record(torch.cuda.current_stream()) + return got + + def issue_copy(self, target): + if not self._pending(target): + return + stream = streams[target % 2] + stream.wait_event(installed[target]) + with torch.cuda.stream(stream): + if target % 2 == 0: + torch.cuda._sleep(_PF_DELAY_CYCLES) + self.cache.copy_rows(target, self.pf_slots[target], self.pf_src[target], self.pf_num[target], slim=True, + kind=PREFETCH_COPY) + _mark_ready_kernel[(1,)](self.pf_ready[target]) + self._copy_events[target][1].record(stream) + + monkeypatch.setattr(ExpertPrefetcher, "install", install) + monkeypatch.setattr(ExpertPrefetcher, "issue_copy", issue_copy) + steps = _garbage(_routed_ids(moes, inputs), 12, "flood", seed=112) + on = _prefetch_run(moes, cache, 1, False, inputs, sleep=0, before_step=_with_override(cache, steps)) + try: + _assert_same_outputs(on[0], off[0]) + _assert_slots_hold_their_experts(cache) + except AssertionError: + return + pytest.fail("a prefetch copy landing on a slot the next install took over went unnoticed") + + +@requires_cuda +@pytest.mark.parametrize("kind", ["bf16", "nvfp4"]) +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +def test_moe_prefetch_on_per_bank_copy_fallback_is_bitwise_identical(kind, graph, bs, monkeypatch): + """Without the fused copy plan (FREETOKEN_FUSED_COPY=0 or misaligned banks) the prefetch copy + takes copy_rows' per-bank launches on the copy stream; late copies still change no bit.""" + import freetoken.moe.offload_cache as offload_cache + + monkeypatch.setattr(offload_cache, "_FUSED_COPY", False) + banks = _prefetch_banks(kind) + inputs = _prefetch_inputs(bs, steps=60, seed=5) + moes, cache = _prefetch_stack(kind, banks, "off", monkeypatch, overlap=True) + off = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0) + moes, cache = _on_stack(kind, banks, monkeypatch, policy="rule", overlap=True, budget=8) + assert not cache._copy_fused_ok + steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=bs) + on = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0, + before_step=_with_override(cache, steps, delay_copy_cycles=_PF_DELAY_CYCLES // 2)) + _assert_same_outputs(on[0], off[0]) + _assert_slots_hold_their_experts(cache) + assert int(cache.prefetch.stats[:, 0].sum()) > 0 + + +def _multi_graph_run(moes, cache, schedule, before_step=None): + """The graph runner's sequence: per graph size (largest first) an eager warm-up, a capture into + one shared pool and a reset; then ``schedule`` of (bs, xs, replay) steps, replays and eager + steps interleaved.""" + _fresh_ctx(_batch=SimpleNamespace(is_prefill=False)) + + def step(xs): + return [moe.forward(x.clone()) for moe, x in zip(moes, xs)] + + graphs, pool = {}, None + for bs in (2, 1): + static = [torch.randn(bs, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16) for _ in moes] + step(static) + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g, pool=pool): + outs = step(static) + cache.reset() + pool = g.pool() if pool is None else pool + graphs[bs] = (g, static, outs) + results = [] + for i, (bs, xs, replay) in enumerate(schedule): + if before_step is not None: + before_step(i) + if replay: + g, static, outs = graphs[bs] + for s, x in zip(static, xs): + s.copy_(x) + g.replay() + else: + outs = step(xs) + if cache.audit is not None: + cache.audit.end_decode_step(cache) + results.append([o.clone() for o in outs]) + torch.cuda.synchronize() + return results + + +@requires_cuda +@pytest.mark.parametrize("policy", ["rule", "lru"]) +def test_moe_prefetch_on_interleaves_graph_sizes_and_eager_steps(policy, monkeypatch): + """bs 1 and bs 2 graphs captured into one pool (as the graph runner does), replayed in an + arbitrary order with eager steps between them, including an eager bs 9 step that forks nothing: + every output equals off's to the bit and every held slot holds its expert.""" + from freetoken.moe.prefetch import MAX_ROWS + + banks = _prefetch_banks("nvfp4") + gen = torch.Generator().manual_seed(41) + order = [(int(b), bool(r)) for b, r in zip(torch.randint(0, 4, (60,), generator=gen), torch.randint(0, 2, (60,), generator=gen))] + sizes = {0: 1, 1: 2, 2: 3, 3: MAX_ROWS + 1} + xs_by_bs = {bs: iter(_prefetch_inputs(bs, steps=60, seed=42 + bs)) for bs in sizes.values()} + schedule = [] + for b, r in order: + bs = sizes[b] + schedule.append((bs, next(xs_by_bs[bs]), r and bs <= 2)) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, policy=policy) + off = _multi_graph_run(moes, cache, schedule) + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True, budget=8) + steps = _garbage(_routed_ids(moes, [xs for _, xs, _ in schedule]), 12, "adversarial", seed=43) + on = _multi_graph_run(moes, cache, schedule, before_step=_with_override(cache, steps)) + _assert_same_outputs(on, off) + _assert_slots_hold_their_experts(cache) + assert cache.prefetch._inflight == [None] * _PF_LAYERS + assert int(cache.prefetch.stats[:, 0].sum()) > 0 and int(cache.prefetch.stats[:, 1].sum()) > 0 + + +@requires_cuda +@pytest.mark.parametrize("seed", range(4)) +@pytest.mark.parametrize("rows", ["bs1", "bs2", "bs2_same"]) +@pytest.mark.parametrize("policy", ["rule", "lru"]) +def test_moe_prefetch_on_fuzzes_per_layer_stream_delays(seed, rows, policy, monkeypatch): + """Eager fuzz: before every layer a random choice of sleeps on the predictor stream, the prefetch + copy stream, the compute stream before the GEMM and the shared expert's side stream, with up to + 12 candidates a layer (most of the 16 slots re-installed every call) and, for bs2_same, two + identical rows. Outputs equal off's, held slots hold their experts, and the final maps and + stats (all but LATE) equal those of the same run without delays.""" + from freetoken.moe.prefetch import LATE, NUM_STAT_COLS + + bs = 1 if rows == "bs1" else 2 + inputs = _prefetch_inputs(bs, steps=60, seed=50 + seed) + if rows == "bs2_same": + inputs = [[x[:1].expand(2, -1).contiguous() for x in xs] for xs in inputs] + banks = _prefetch_banks("nvfp4") + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, policy=policy) + off = _prefetch_run(moes, cache, bs, False, inputs, sleep=0) + stall = _stall_hooks(monkeypatch) + gen = torch.Generator().manual_seed(seed) + choices = torch.randint(0, 16, (len(inputs), _PF_LAYERS), generator=gen).tolist() + runs = [] + for delayed in (False, True): + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True, budget=12) + steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=60 + seed) + pf = cache.prefetch + load = _with_override(cache, steps) + at = {"step": -1} + + def before_step(i, load=load): + at["step"] = i + load(i) + + for layer, moe in enumerate(moes): + real = moe.forward + + def forward(x, layer=layer, real=real, pf=pf): + bits = choices[at["step"]][layer] if delayed and at["step"] >= 0 else 0 + # layer L's fork runs layer L+1's prediction and copy: these delays hit them + pf.delay_predict_cycles = _PF_DELAY_CYCLES // 2 if bits & 1 else 0 + pf.delay_copy_cycles = _PF_DELAY_CYCLES // 2 if bits & 2 else 0 + stall["gemm"] = _PF_DELAY_CYCLES // 2 if bits & 4 else 0 + stall["shared"] = _PF_DELAY_CYCLES // 2 if bits & 8 else 0 + return real(x) + + monkeypatch.setattr(moe, "forward", forward) + on = _prefetch_run(moes, cache, bs, False, inputs, sleep=0, before_step=before_step) + _assert_same_outputs(on[0], off[0]) + _assert_slots_hold_their_experts(cache) + runs.append((on[1][:3], pf.stats.cpu())) + (maps, stats), (delayed_maps, delayed_stats) = runs + for a, b in zip(maps, delayed_maps): + assert torch.equal(a, b) + cols = [c for c in range(NUM_STAT_COLS) if c != LATE] + assert torch.equal(stats[:, cols], delayed_stats[:, cols]) + + +# FREETOKEN_MOE_PREFETCH_VERIFY: the debug slot checks on the same stacks, quiet on a correct decode +# and loud on each kind of fault they exist for +_VERIFY_STEPS = 60 + + +def _verify_state(cache): + """(counters summed over layers, kept records) of the stack's verifier.""" + from freetoken.moe.verify import RING + + v = cache.verify + torch.cuda.synchronize() + return v.counters.cpu().sum(0).tolist(), v.ring[: min(int(v.cursor), RING)].cpu().tolist() + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("prefetch", ["off", "on"]) +@pytest.mark.parametrize("verify", ["full", "meta"]) +def test_moe_prefetch_verify_is_quiet_on_a_correct_decode(graph, bs, prefetch, verify, monkeypatch): + """FREETOKEN_MOE_PREFETCH_VERIFY checks every decode layer call and finds nothing on a correct + decode: prefetch off, and prefetch on with candidates the GEMMs read, late copies and a late + predictor. The checks only read, so the outputs equal those of verify off to the bit.""" + from freetoken.moe.prefetch import USEFUL + from freetoken.moe.verify import CHECKS + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(bs, steps=_VERIFY_STEPS, seed=70) + runs = [] + for mode in ("off", verify): + moes, cache = _prefetch_stack("nvfp4", banks, prefetch, monkeypatch, overlap=True, budget=8, verify_mode=mode) + before_step = None + if prefetch == "on": + steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=bs + 70) + before_step = _with_override(cache, steps, delay_copy_cycles=_PF_DELAY_CYCLES // 2, + delay_predict_cycles=_PF_DELAY_CYCLES // 4) + runs.append(_prefetch_run(moes, cache, bs, graph, inputs, sleep=0, before_step=before_step)) + _assert_same_outputs(runs[1][0], runs[0][0]) + counts, records = _verify_state(cache) + assert counts[CHECKS] == (_VERIFY_STEPS + 1) * _PF_LAYERS # the eager warm-up is checked too + assert counts[CHECKS + 1 :] == [0] * (len(counts) - 1) and records == [], records + assert (cache.verify.scratch is None) == (verify == "meta") + if prefetch == "on": + assert int(cache.prefetch.stats[:, USEFUL].sum()) > 0, "the GEMMs must have read prefetched slots" + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("prefetch", ["off", "on"]) +@pytest.mark.parametrize("path", list(_PREFILL_PATHS)) +def test_moe_prefetch_verify_is_quiet_across_prefills(graph, prefetch, path, monkeypatch): + """Every request prefills before it decodes: the whole-layer materialize, the overlap double + buffers (with the hit-D2D gather) and small prefills rewrite slots between decode steps, and the + decode checks after each of them still find nothing.""" + import freetoken.layers.moe as layers_moe + from freetoken.moe.verify import CHECKS + + cache_size, tokens, cache_kw, small = _PREFILL_PATHS[path] + monkeypatch.setattr(layers_moe, "_SMALL_PREFILL_TOKENS", small) + banks = _prefetch_banks("nvfp4") + decode = _prefetch_inputs(2, steps=18, seed=75) + gen = torch.Generator(device="cuda").manual_seed(76) + prefills = [[torch.randn(tokens, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16, generator=gen) * 0.5 + for _ in range(_PF_LAYERS)] for _ in range(2)] + schedule = ([("decode", xs) for xs in decode[:6]] + [("prefill", prefills[0])] + + [("decode", xs) for xs in decode[6:12]] + [("prefill", prefills[1])] + + [("decode", xs) for xs in decode[12:]]) + moes, cache = _prefetch_stack("nvfp4", banks, prefetch, monkeypatch, budget=8, overlap=True, cache_size=cache_size, + verify_mode="full", **cache_kw) + before_step = None + if prefetch == "on": + steps = _garbage(_routed_ids(moes, [xs for _, xs in schedule]), 12, "adversarial", seed=77) + before_step = _with_override(cache, steps) + _mixed_run(moes, cache, 2, graph, schedule, before_step=before_step) + counts, records = _verify_state(cache) + if path == "overlap_d2d": + assert cache.prefill_hit_rows > 0, "the hit-D2D gather must have read resident slots" + assert counts[CHECKS] == (len(decode) + 1) * _PF_LAYERS + assert counts[CHECKS + 1 :] == [0] * (len(counts) - 1) and records == [], records + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +def test_moe_prefetch_verify_catches_a_missing_copy_join(graph, monkeypatch): + """Negative control: without the compute stream's wait on the prefetch copy (and with that copy + late), GEMMs read prefetched slots before their bytes land. The byte check before the GEMM flags + them, and the records place the slot in the layer's prefetch copy plan.""" + from freetoken.moe.prefetch import ExpertPrefetcher + from freetoken.moe.verify import BYTES_PRE, META_BAD, PRE_BAD, R_KIND, R_SOURCE, SRC_PREFETCH + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(1, steps=40, seed=7) + moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, overlap=True, budget=8, verify_mode="full") + + def no_wait(self, layer_id): + self._inflight[layer_id] = None + + monkeypatch.setattr(ExpertPrefetcher, "join_copy", no_wait) + if graph: + # a capture needs the copy stream joined back: after the last layer, past every GEMM it races + last, copy_stream = moes[-1].forward, cache.prefetch.copy_stream + + def forward(x): + out = last(x) + torch.cuda.current_stream().wait_stream(copy_stream) + return out + + monkeypatch.setattr(moes[-1], "forward", forward) + steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=11) + # ~2.5 ms a copy: the copy stream falls ever further behind even an eager host's slow enqueue + _prefetch_run(moes, cache, 1, graph, inputs, sleep=0, + before_step=_with_override(cache, steps, delay_copy_cycles=5 * _PF_DELAY_CYCLES)) + counts, records = _verify_state(cache) + assert counts[PRE_BAD] > 0 and counts[META_BAD] == 0 + assert any(r[R_KIND] == BYTES_PRE and r[R_SOURCE] & SRC_PREFETCH for r in records) + + +def _on_call(layer, call, action): + """Wrap ExpertVerifier.before_gemm so ``action(verifier, cache, slots, real)`` replaces it on the + ``call``-th decode call of ``layer`` (call 0 is the eager warm-up).""" + from freetoken.moe.verify import ExpertVerifier + + real, seen = ExpertVerifier.before_gemm, {} + + def before_gemm(self, cache, slots, next_plan): + at = self._call[0] + seen[at] = seen.get(at, -1) + 1 + if (at, seen[at]) == (layer, call): + return action(self, cache, slots, lambda: real(self, cache, slots, next_plan)) + return real(self, cache, slots, next_plan) + + return before_gemm + + +@requires_cuda +@pytest.mark.parametrize("prefetch", ["off", "on"]) +def test_moe_prefetch_verify_catches_a_corrupt_slot(prefetch, monkeypatch): + """Negative control: bytes flipped in one bank of one routed slot right before the checks are + flagged before the GEMM, with the call, layer, slot, expert, bank and byte offset of the flip.""" + from freetoken.moe.verify import (BYTES_PRE, META_BAD, PRE_BAD, R_BANK, R_CALL, R_EXPERT, R_FIRST_BYTE, + R_KIND, R_LAYER, R_SLOT, ExpertVerifier) + + banks = _prefetch_banks("nvfp4") + moes, cache = _prefetch_stack("nvfp4", banks, prefetch, monkeypatch, overlap=True, verify_mode="full") + hit = {} + + def corrupt(verifier, cache, slots, real): + torch.cuda.synchronize() + hit.update(slot=int(slots.view(-1)[0]), expert=int(verifier.ids[0])) + cache.bank_views()[3][hit["slot"]].view(torch.uint8).view(-1)[100:108].bitwise_not_() + real() + + monkeypatch.setattr(ExpertVerifier, "before_gemm", _on_call(2, 10, corrupt)) + _prefetch_run(moes, cache, 1, False, _prefetch_inputs(1, steps=20, seed=71), sleep=0) + counts, records = _verify_state(cache) + assert counts[PRE_BAD] >= 1 and counts[META_BAD] == 0 + first = records[0] + assert [first[f] for f in (R_KIND, R_CALL, R_LAYER, R_SLOT, R_EXPERT, R_BANK, R_FIRST_BYTE)] == [ + BYTES_PRE, 10, 2, hit["slot"], hit["expert"], 3, 100] + + +@requires_cuda +@pytest.mark.parametrize("verify", ["full", "meta"]) +def test_moe_prefetch_verify_catches_a_wrong_slot_owner(verify, monkeypatch): + """Negative control: a routed slot's id_of_slot names another layer's expert while the checks + before the GEMM run (restored right after them): exactly that entry is flagged, with the owner it + saw, and the checks after the GEMM find the maps whole again.""" + from freetoken.moe.verify import (CHECKS, META_BAD, META_PRE, R_CALL, R_EXPERT, R_KIND, R_LAYER, R_MAPPED, + R_OWNER, R_ROW, R_SLOT, ExpertVerifier) + + banks = _prefetch_banks("nvfp4") + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, verify_mode=verify) + hit = {} + + def corrupt(verifier, cache, slots, real): + torch.cuda.synchronize() + s, e = int(slots.view(-1)[1]), int(verifier.ids[1]) + hit.update(slot=s, expert=e, saved=int(cache.id_of_slot[s])) + cache.id_of_slot[s] = 3 * _PF_EXPERTS + e # in range: an ensure reading it stays in bounds + real() + cache.id_of_slot[s] = hit["saved"] + + monkeypatch.setattr(ExpertVerifier, "before_gemm", _on_call(1, 7, corrupt)) + _prefetch_run(moes, cache, 1, False, _prefetch_inputs(1, steps=12, seed=72), sleep=0) + counts, records = _verify_state(cache) + assert hit["saved"] == _PF_EXPERTS + hit["expert"] + assert counts[META_BAD] == 1 and sum(counts[CHECKS + 1 :]) == 1 and len(records) == 1 + rec = records[0] + assert [rec[f] for f in (R_KIND, R_CALL, R_LAYER, R_ROW, R_SLOT, R_EXPERT, R_OWNER, R_MAPPED)] == [ + META_PRE, 7, 1, 1, hit["slot"], hit["expert"], 3 * _PF_EXPERTS + hit["expert"], hit["slot"]] + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("prefetch", ["off", "on"]) +def test_moe_prefetch_verify_catches_a_write_during_the_gemm(graph, prefetch, monkeypatch): + """Negative control: another stream rewrites a routed slot while the GEMM runs, after the checks + before it (a sleep on the compute stream holds the GEMM until the write lands). The byte check + after the GEMM flags it; the check before it saw nothing.""" + from freetoken.layers.moe import OffloadMoELayer + from freetoken.moe.verify import (BYTES_POST, BYTES_PRE, POST_BAD, R_BANK, R_CALL, R_FIRST_BYTE, R_KIND, R_LAYER, + R_SLOT, ExpertVerifier) + + banks = _prefetch_banks("nvfp4") + moes, cache = _prefetch_stack("nvfp4", banks, prefetch, monkeypatch, overlap=True, verify_mode="full") + side, real, armed = torch.cuda.Stream(), OffloadMoELayer._expert_gemm, {"on": False} + slot = torch.zeros((1,), dtype=torch.int64, device="cuda") + offset = torch.arange(16, device="cuda") + + def gemm(self, cache, hidden, weights, topk_ids, *, is_prefill, **kw): + if armed["on"] and self.layer_id == 3 and not is_prefill: + slot.copy_(topk_ids.view(-1)[:1]) # device side, so a captured graph rewrites this step's slot + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + torch.cuda._sleep(_PF_DELAY_CYCLES // 20) + row = cache.bank_views()[0].view(cache.cache_size, -1) + row.index_put_((slot, offset), row[slot, offset].bitwise_not()) + torch.cuda._sleep(_PF_DELAY_CYCLES) + out = real(self, cache, hidden, weights, topk_ids, is_prefill=is_prefill, **kw) + torch.cuda.current_stream().wait_stream(side) # a capture must join the side stream back + return out + return real(self, cache, hidden, weights, topk_ids, is_prefill=is_prefill, **kw) + + monkeypatch.setattr(OffloadMoELayer, "_expert_gemm", gemm) + real_run = ExpertVerifier.before_gemm + + def arm_at_call_5(self, cache, slots, next_plan): + # graph: the arm must be set while capturing, so the rewrite replays every step; eager: once at call 5 + armed["on"] = self._call[0] == 3 and (graph or int(self.calls[3]) == 5) + return real_run(self, cache, slots, next_plan) + + monkeypatch.setattr(ExpertVerifier, "before_gemm", arm_at_call_5) + _prefetch_run(moes, cache, 1, graph, _prefetch_inputs(1, steps=8, seed=73), sleep=0) + torch.cuda.synchronize() + counts, records = _verify_state(cache) + assert counts[POST_BAD] >= 1 + first = records[0] + assert [first[f] for f in (R_KIND, R_LAYER, R_BANK, R_FIRST_BYTE)] == [BYTES_POST, 3, 0, 0] + assert not any(r[R_KIND] == BYTES_PRE and (r[R_CALL], r[R_SLOT]) == (first[R_CALL], first[R_SLOT]) for r in records) + + +@requires_cuda +def test_moe_prefetch_verify_off_builds_and_launches_nothing(monkeypatch): + """With the flag off there is no verifier and a decode launches only what it launches without + one; meta and full only add launches (their kernels, the id copy and the host-row gather).""" + from collections import Counter + + from torch.profiler import ProfilerActivity, profile + + from freetoken.env import ENV + + monkeypatch.setattr(ENV.MOE_PREFETCH_VERIFY, "value", "0") + banks = _prefetch_banks("bf16") + inputs = _prefetch_inputs(1, steps=2) + + def kernels(verify): + moes, cache = _prefetch_stack("bf16", banks, "off", monkeypatch, verify_mode=verify) + _prefetch_run(moes, cache, 1, False, inputs[:1]) # warm up and compile outside the trace + with profile(activities=[ProfilerActivity.CUDA]) as prof: + _prefetch_run(moes, cache, 1, False, inputs) + return cache, Counter(e.name for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA) + + cache, off = kernels(None) + assert cache.verify is None and cache.verify_mode == "off" + ours = ("_plan_kernel", "_compare_kernel", "_tally_kernel") + assert not any(k in name for name in off for k in ours) + for mode in ("meta", "full"): + _, on = kernels(mode) + assert not off - on, f"{mode} dropped a launch" + added = on - off + assert all(any(k in name for k in (*ours, "Memcpy", "fast_index_copy_multi")) for name in added), added + assert any("_compare_kernel" in name for name in added) == (mode == "full") + + +@requires_cuda +def test_moe_prefetch_verify_follows_a_rebuild(monkeypatch): + """A rebuild reallocates the slot banks (production rebuilds to its final slot count after start): + the verifier keeps its scratch and counts, points its byte checks at the new banks, and a + recaptured graph checks clean.""" + from freetoken.moe.verify import CHECKS + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(2, steps=30, seed=74) + moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, overlap=True, verify_mode="full") + v = cache.verify + scratch = v.scratch + _prefetch_run(moes, cache, 2, True, inputs, sleep=0) + cache.rebuild(_PF_EXPERTS + 4) + assert cache.verify is v and v.scratch is scratch + assert v.cache_ptrs.tolist() == [bank.data_ptr() for bank in cache.bank_views()] + _prefetch_run(moes, cache, 2, True, inputs, sleep=0) + counts, records = _verify_state(cache) + assert counts[CHECKS] == 2 * (30 + 1) * _PF_LAYERS + assert counts[CHECKS + 1 :] == [0] * (len(counts) - 1) and records == [] + + +# FREETOKEN_MOE_SLOT_AUDIT (moe/slot_audit.py): quiet and output-neutral on a correct decode with +# prefills between the steps, and each injected persistent fault found with the history that explains it + + +def _audit_records(cache): + """(totals, kept records) of the stack's slot audit.""" + a = cache.audit + torch.cuda.synchronize() + return a.totals.cpu().tolist(), a.ring[: min(int(a.cursor), a.ring_n)].cpu().tolist() + + +def _bad_slots(cache): + """Ground truth: the (slot, owner) pairs whose bytes differ from the owner's host rows.""" + torch.cuda.synchronize() + held = (cache.id_of_slot >= 0).nonzero().view(-1).cpu() + ids = cache.id_of_slot[held.cuda()].long().cpu() + bad = torch.zeros(held.numel(), dtype=torch.bool) + for (per_layer, _), view in zip(cache.banks, cache.bank_views()): + want = torch.stack(per_layer).flatten(0, 1)[ids].contiguous().view(torch.uint8).view(held.numel(), -1) + got = view[held.cuda()].cpu().contiguous().view(torch.uint8).view(held.numel(), -1) + bad |= (got != want).any(dim=1) + return {(int(s), int(o)) for s, o, b in zip(held, ids, bad) if b} + + +def _spy_audit_kinds(monkeypatch): + """The event kinds the recorders were asked for, at enqueue time, each with the streams its + recorders ran on (a recorder follows its writer on the writer's stream).""" + from freetoken.moe.slot_audit import SlotAuditor + + kinds, real = {}, SlotAuditor._launch + + def launch(self, cache, kind, *args, **kwargs): + kinds.setdefault(kind, set()).add(torch.cuda.current_stream().cuda_stream) + return real(self, cache, kind, *args, **kwargs) + + monkeypatch.setattr(SlotAuditor, "_launch", launch) + return kinds + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("prefetch", ["off", "on"]) +@pytest.mark.parametrize("path", list(_PREFILL_PATHS)) +def test_moe_slot_audit_is_quiet_and_changes_no_bit(graph, bs, prefetch, path, monkeypatch): + """FREETOKEN_MOE_SLOT_AUDIT=1 records every slot writer (the decode and prefetch ensures and copies, + and between the decode steps the whole-layer materialize, the overlap double buffers with the + hit-D2D gather, or small prefills) and audits every held slot after each decode step: it finds + nothing, and the outputs, slot maps and slot bytes equal those of the audit off to the bit. Every + map writer it records ran on the compute stream but the prefill buffer invalidation (on the + prefill copy stream), and the prefetch copy on the prefetch copy stream.""" + import freetoken.layers.moe as layers_moe + import freetoken.moe.offload_cache as offload_cache + from freetoken.moe import slot_audit as sa + + cache_size, tokens, cache_kw, small = _PREFILL_PATHS[path] + monkeypatch.setattr(layers_moe, "_SMALL_PREFILL_TOKENS", small) + if path == "overlap_d2d": + # every toy bank is under the 256 KiB small-bank floor, which would leave the hit gather idle + monkeypatch.setattr(offload_cache, "_SMALL_BANK_FEAT_BYTES", 8192) + banks = _prefetch_banks("nvfp4") + decode = _prefetch_inputs(bs, steps=18, seed=81) + gen = torch.Generator(device="cuda").manual_seed(82) + prefills = [[torch.randn(tokens, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16, generator=gen) * 0.5 + for _ in range(_PF_LAYERS)] for _ in range(2)] + schedule = ([("decode", xs) for xs in decode[:6]] + [("prefill", prefills[0])] + + [("decode", xs) for xs in decode[6:12]] + [("prefill", prefills[1])] + + [("decode", xs) for xs in decode[12:]]) + kinds = _spy_audit_kinds(monkeypatch) + runs = [] + for interval in (0, 1): + moes, cache = _prefetch_stack("nvfp4", banks, prefetch, monkeypatch, budget=8, overlap=True, + cache_size=cache_size, slot_audit=interval, **cache_kw) + assert (cache.audit is None) == (interval == 0) + before_step = None + if prefetch == "on": + steps = _garbage(_routed_ids(moes, [xs for _, xs in schedule]), 12, "adversarial", seed=bs + 83) + before_step = _with_override(cache, steps) + compute = _record_compute_streams(monkeypatch, moes) + out = _mixed_run(moes, cache, bs, graph, schedule, before_step=before_step) + runs.append((out, [t.clone() for t in (cache.slot_for_id, cache.id_of_slot, cache.usage, *cache.bank_views())])) + _assert_same_outputs(runs[1][0], runs[0][0]) + for a, b in zip(runs[1][1], runs[0][1]): + assert torch.equal(a, b) + totals, records = _audit_records(cache) + assert records == [] and totals[sa.T_BAD] == 0, [cache.audit.format_record(r) for r in records] + assert totals[sa.T_MAP_RANGE] == totals[sa.T_BAD_PLAN] == totals[sa.T_UNCHECKED] == 0 + assert totals[sa.T_AUDITS] == len(decode) + assert totals[sa.T_HELD] > 0 and totals[sa.T_BYTES_CHECKED] == totals[sa.T_HELD] and totals[sa.T_EVENTS] > 0 + expected = {sa.DEMAND_INSTALL, sa.DEMAND_COPY, sa.RESET} + if prefetch == "on": + expected |= {sa.PREFETCH_INSTALL, sa.PREFETCH_COPY} + expected |= { + "materialize": {sa.MATERIALIZE_CLEAR, sa.MATERIALIZE_INSTALL, sa.MATERIALIZE_COPY}, + "overlap": {sa.INVALIDATE, sa.PREFILL_BUFFER}, + "overlap_d2d": {sa.INVALIDATE, sa.PREFILL_SPLIT_H2D, sa.PREFILL_HIT_D2D}, + "small": {sa.PREFILL_INSTALL, sa.PREFILL_COPY}, + }[path] + assert set(kinds) == expected, (set(kinds) - expected, expected - set(kinds)) + fence = set() if cache.prefill_copy_stream is None else {cache.prefill_copy_stream.cuda_stream} + for kind in (sa.DEMAND_INSTALL, sa.PREFETCH_INSTALL, sa.PREFILL_INSTALL, sa.MATERIALIZE_INSTALL, + sa.MATERIALIZE_CLEAR, sa.RESET, sa.INVALIDATE): + assert kinds.get(kind, set()) <= (fence if kind == sa.INVALIDATE else compute), kind + if prefetch == "on": + assert kinds[sa.PREFETCH_COPY] == {cache.prefetch.copy_stream.cuda_stream} + + +@requires_cuda +@pytest.mark.parametrize("policy", ["rule", "lru"]) +def test_moe_slot_audit_is_quiet_across_graph_sizes_and_eager_steps(policy, monkeypatch): + """The graph runner's shape: bs 1 and bs 2 graphs in one pool, replayed in any order with eager + steps between them (a bs 9 one forks no prefetch); audited after every step, with the decode + verifier on beside it, neither finds anything.""" + from freetoken.moe import slot_audit as sa + from freetoken.moe.prefetch import MAX_ROWS + from freetoken.moe.verify import CHECKS + + banks = _prefetch_banks("nvfp4") + gen = torch.Generator().manual_seed(90) + order = [(int(b), bool(r)) for b, r in zip(torch.randint(0, 4, (40,), generator=gen), torch.randint(0, 2, (40,), generator=gen))] + sizes = {0: 1, 1: 2, 2: 3, 3: MAX_ROWS + 1} + xs_by_bs = {bs: iter(_prefetch_inputs(bs, steps=40, seed=91 + bs)) for bs in sizes.values()} + schedule = [(sizes[b], next(xs_by_bs[sizes[b]]), r and sizes[b] <= 2) for b, r in order] + moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, policy=policy, overlap=True, budget=8, slot_audit=1, + verify_mode="meta") + steps = _garbage(_routed_ids(moes, [xs for _, xs, _ in schedule]), 12, "adversarial", seed=92) + _multi_graph_run(moes, cache, schedule, before_step=_with_override(cache, steps)) + totals, records = _audit_records(cache) + assert records == [] and totals[sa.T_BAD] == 0 and totals[sa.T_AUDITS] == len(schedule) + assert totals[sa.T_BYTES_CHECKED] == totals[sa.T_HELD] > 0 + counts, verify_records = _verify_state(cache) + assert counts[CHECKS] > 0 and counts[CHECKS + 1 :] == [0] * (len(counts) - 1) and verify_records == [] + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("policy", ["rule", "lru"]) +def test_moe_slot_audit_explains_a_skipped_prefetch_copy(graph, policy, monkeypatch): + """Fault injection: layer 2's prefetch copies copy nothing. Every slot the audit finds bad holds a + layer-2 expert whose history reads 'prefetch install, no prefetch copy', and every bad slot still + held at the end was found.""" + from freetoken.moe import slot_audit as sa + + monkeypatch.setattr(sa, "RING", 1024) + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(1, steps=40, seed=84) + moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, policy=policy, overlap=True, budget=8, slot_audit=1) + real = cache.copy_rows + + def copy_rows(layer_id, dst, src, num, *, slim, **kw): + if kw.get("kind") == sa.PREFETCH_COPY and layer_id == 2: + return None # the fault: no bytes move (and so nothing is recorded) + return real(layer_id, dst, src, num, slim=slim, **kw) + + monkeypatch.setattr(cache, "copy_rows", copy_rows) + steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=85) + _prefetch_run(moes, cache, 1, graph, inputs, sleep=0, before_step=_with_override(cache, steps)) + truth = _bad_slots(cache) + totals, records = _audit_records(cache) + assert records and totals[sa.T_BAD] > 0 and totals[sa.T_NEW] == len(records) + a = cache.audit + for rec in records: + assert rec[sa.H_FLAGS] == sa.F_BYTES and rec[sa.H_OWNER] // _PF_EXPERTS == 2, a.format_record(rec) + # a layer-1 prefetch copy still in flight may land after the install and is named as such + assert a.diagnose(rec).startswith("prefetch install, no prefetch copy"), a.format_record(rec) + assert truth <= {(r[sa.H_SLOT], r[sa.H_OWNER]) for r in records} + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("prefetch", ["off", "on"]) +def test_moe_slot_audit_explains_a_demand_copy_from_the_wrong_layer(graph, prefetch, monkeypatch): + """Fault injection: layer 2's demand copies read layer 1's host banks (the right rows). The bad + slots are layer-2 demand installs whose last copy came from layer 1's source.""" + from freetoken.moe import slot_audit as sa + from freetoken.moe.offload_cache import OffloadMoeCache + + monkeypatch.setattr(sa, "RING", 1024) + real = OffloadMoeCache.copy_missing + + def copy_missing(self, layer_id=None): + if self._pending_src_layer == 2 and not self._pending_whole_layer: + return self.copy_rows(1, self.evict_slots, self.src_indices, self.num_indices, slim=False, kind=self._pending_kind) + return real(self, layer_id) + + monkeypatch.setattr(OffloadMoeCache, "copy_missing", copy_missing) + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(1, steps=30, seed=86) + moes, cache = _prefetch_stack("nvfp4", banks, prefetch, monkeypatch, overlap=True, budget=8, slot_audit=1) + _prefetch_run(moes, cache, 1, graph, inputs, sleep=0) + truth = _bad_slots(cache) + totals, records = _audit_records(cache) + assert records and truth <= {(r[sa.H_SLOT], r[sa.H_OWNER]) for r in records} + for rec in records: + owner = rec[sa.H_OWNER] + assert owner // _PF_EXPERTS == 2 and rec[sa.H_FLAGS] == sa.F_BYTES + e = owner % _PF_EXPERTS + assert cache.audit.diagnose(rec).startswith( + f"demand install of L2/e{e}, then demand copy from layer 1's source (L1/e{e}, #"), cache.audit.format_record(rec) + + +@requires_cuda +@pytest.mark.parametrize("prefetch", ["off", "on"]) +def test_moe_slot_audit_explains_bytes_overwritten_after_install(prefetch, monkeypatch): + """Fault injection: bytes of a held slot flipped behind every recorder's back. The next audit keeps + one record (bank, byte offset, bad words) whose history ends in the owner's install and copy, a + later audit counts it again without a second record, and a clean slot is forgotten.""" + from freetoken.moe import slot_audit as sa + + banks = _prefetch_banks("nvfp4") + moes, cache = _prefetch_stack("nvfp4", banks, prefetch, monkeypatch, overlap=True, budget=8, slot_audit=1000) + _prefetch_run(moes, cache, 1, False, _prefetch_inputs(1, steps=12, seed=87), sleep=0) + a = cache.audit + a.scan(cache) + assert _audit_records(cache)[1] == [] + slot = int((cache.id_of_slot >= 0).nonzero()[0]) + owner = int(cache.id_of_slot[slot]) + row = cache.bank_views()[3][slot].view(torch.uint8).view(-1) + row[100:108].bitwise_not_() + a.scan(cache) + a.scan(cache) + totals, records = _audit_records(cache) + assert len(records) == 1 and [totals[t] for t in (sa.T_AUDITS, sa.T_BAD, sa.T_NEW, sa.T_STILL)] == [3, 2, 1, 1] + rec = records[0] + assert [rec[f] for f in (sa.H_AUDIT, sa.H_SLOT, sa.H_OWNER, sa.H_FLAGS, sa.H_BANKS, sa.H_FIRST_BANK, + sa.H_FIRST_BYTE, sa.H_BAD_WORDS)] == [1, slot, owner, sa.F_BYTES, 1 << 3, 3, 100, 2] + events = a.events(rec) + install = [ev for ev in events if ev["kind"] in sa.INSTALLS][-1] + copy = [ev for ev in events if ev["ring"] == "bytes"][-1] + assert a._flat(install) == owner == a._flat(copy) and copy["seq"] > install["seq"] + assert "then the bytes changed with no recorded write" in a.diagnose(rec) + row[100:108].bitwise_not_() + a.scan(cache) + row[100:108].bitwise_not_() + a.scan(cache) + totals, records = _audit_records(cache) + assert len(records) == 2 and records[1][sa.H_AUDIT] == 4 and totals[sa.T_BAD] == 3 + + +@requires_cuda +def test_moe_slot_audit_explains_stale_and_missing_map_entries(monkeypatch): + """Fault injection on the maps: slot_for_id still names a slot for its previous owner (an eviction + that forgot to clear it: routing that expert would read another's bytes), and a held slot's owner + has lost its entry. Both are found and the stale one is traced to the install it outlived.""" + from freetoken.moe import slot_audit as sa + + banks = _prefetch_banks("nvfp4") + moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, overlap=True, budget=8, slot_audit=1000) + _prefetch_run(moes, cache, 1, False, _prefetch_inputs(1, steps=12, seed=88), sleep=0) + a = cache.audit + flat = cache.slot_for_id.view(-1) + held = (cache.id_of_slot >= 0).nonzero().view(-1).tolist() + found = None + for s in held: + for ev in sorted(a.meta_hist[s].tolist()): + fid = ev[sa.E_LAYER] * _PF_EXPERTS + ev[sa.E_ROW] + if ev[sa.E_KIND] in sa.INSTALLS and int(flat[fid]) == -1: + found = (s, fid) + if found: + break + assert found, "some slot must have outlived an earlier owner" + stale_slot, stale_id = found + fwd_slot = next(t for t in held if t != stale_slot) + fwd_owner = int(cache.id_of_slot[fwd_slot]) + flat[stale_id] = stale_slot + flat[fwd_owner] = -1 + a.scan(cache) + totals, records = _audit_records(cache) + assert [totals[t] for t in (sa.T_BAD, sa.T_STALE, sa.T_FWD, sa.T_BYTES)] == [2, 1, 1, 0] and len(records) == 2 + by_slot = {r[sa.H_SLOT]: r for r in records} + rec = by_slot[stale_slot] + assert rec[sa.H_FLAGS] == sa.F_STALE and rec[sa.H_STALE_ID] == stale_id + holder = a._name(int(cache.id_of_slot[stale_slot])) + text = a.diagnose(rec) + assert text.startswith(f"stale slot_for_id: {a._name(stale_id)} -> this slot, which holds {holder}; " + f"{a._name(stale_id)} was installed here by "), text + assert text.endswith("without clearing its map entry"), text + rec = by_slot[fwd_slot] + assert rec[sa.H_FLAGS] == sa.F_FWD and rec[sa.H_MAPPED] == -1 + assert a.diagnose(rec) == f"slot_for_id[{a._name(fwd_owner)}]=-1, not this slot" + + +@requires_cuda +def test_moe_slot_audit_off_builds_and_launches_nothing(monkeypatch): + """With the flag off there is no auditor and a decode launches only what it launches without one; + on, a decode only adds the recorders (and the audit its own kernels and host-row gathers).""" + from collections import Counter + + from torch.profiler import ProfilerActivity, profile + + from freetoken.env import ENV + + monkeypatch.setattr(ENV.MOE_SLOT_AUDIT, "value", 0) + banks = _prefetch_banks("bf16") + inputs = _prefetch_inputs(1, steps=2) + + def kernels(interval): + moes, cache = _prefetch_stack("bf16", banks, "on", monkeypatch, slot_audit=interval) + _prefetch_run(moes, cache, 1, False, inputs[:1]) # warm up and compile outside the trace + with profile(activities=[ProfilerActivity.CUDA]) as prof: + _prefetch_run(moes, cache, 1, False, inputs) + return cache, Counter(e.name for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA) + + cache, off = kernels(None) + assert cache.audit is None and cache.slot_audit == 0 + ours = ("_record_kernel", "_bump_step_kernel", "_audit_") + assert not any(k in name for name in off for k in ours) + _, on = kernels(1) + assert not off - on, "the audit dropped a launch" + added = on - off + assert all(any(k in name for k in (*ours, "fast_index_copy_multi")) for name in added), added + assert any("_audit_report_kernel" in name for name in added) + + +@requires_cuda +def test_moe_slot_audit_follows_a_rebuild(monkeypatch): + """A rebuild reallocates the slot banks and maps: the auditor keeps its scratch and totals, starts a + new history sized for the new slot count, points its byte checks at the new banks, and a recaptured + graph audits clean.""" + from freetoken.moe import slot_audit as sa + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(2, steps=20, seed=89) + moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, overlap=True, slot_audit=1) + a = cache.audit + scratch = a.scratch + _prefetch_run(moes, cache, 2, True, inputs, sleep=0) + cache.rebuild(_PF_EXPERTS + 4) + assert cache.audit is a and a.scratch is scratch and a.num_slots == _PF_EXPERTS + 4 + assert a.meta_hist.shape[0] == a.bytes_cur.shape[0] == _PF_EXPERTS + 4 and int(a.meta_cur.sum()) == 0 + assert a.cache_ptrs.tolist() == [bank.data_ptr() for bank in cache.bank_views()] + _prefetch_run(moes, cache, 2, True, inputs, sleep=0) + totals, records = _audit_records(cache) + assert records == [] and totals[sa.T_BAD] == 0 and totals[sa.T_AUDITS] == 2 * len(inputs) + + @requires_cuda def test_decoder_stack_prefill_and_decode(monkeypatch): """Ragged bs=3 prefill then a bs=3 decode step through the whole model with dummy weights.""" diff --git a/tests/moe/ref_prefetch.py b/tests/moe/ref_prefetch.py new file mode 100644 index 0000000000..c40d9351a8 --- /dev/null +++ b/tests/moe/ref_prefetch.py @@ -0,0 +1,22 @@ +"""CPU reference for the router-lookahead select and count (freetoken.moe.prefetch).""" +from __future__ import annotations + + +def ref_select(logits, resident, k: int, budget: int): + """Each row's top-k (ties to the lowest id), merged rank by rank, duplicates dropped; the first + ``budget`` non-resident ids, and the resident ids among the first ``budget`` distinct ones.""" + rows = [sorted(range(len(row)), key=lambda e: (-row[e], e))[:k] for row in logits] + merged = [] + for r in range(k): + for row in rows: + if row[r] not in merged: + merged.append(row[r]) + sel = [e for e in merged if resident[e] < 0][:budget] + res = [e for e in merged[:budget] if resident[e] >= 0] + return sel, res + + +def ref_count(sel, res, routed, misses: int, rows: int): + """``[issued, useful, resident_hits, calls, rows, misses]`` for one layer call.""" + routed = set(routed) + return [len(sel), sum(e in routed for e in sel), sum(e in routed for e in res), 1, rows, misses] diff --git a/tests/moe/ref_scored_cache.py b/tests/moe/ref_scored_cache.py index d7d8b0f2fb..15153c86d5 100644 --- a/tests/moe/ref_scored_cache.py +++ b/tests/moe/ref_scored_cache.py @@ -47,12 +47,13 @@ def _softplus2(self, x): g = np.where(i < se.G_LAST, g0 + (((g1 - g0) * frac) >> shift), 0) return np.maximum(x, 0) + g - def _keys(self, layer, tok, now, pin_since=None): + def _keys(self, layer, tok, now, pin_since=None, lowpri=False): c = np.arange(self.S, dtype=np.int64) oid = self.id_of_slot held = oid >= 0 if self.policy == 0: - key = (self.usage << self.slot_bits) | c + # lowpri: empty (0) < unconsumed prefetch (held, usage 0 -> 1) < resident (usage + 1) + key = (np.where(held, self.usage + 1, 0) << self.slot_bits) | c if lowpri else (self.usage << self.slot_bits) | c else: ids = np.where(held, oid, 0) lk = ids // self.E @@ -68,16 +69,20 @@ def _keys(self, layer, tok, now, pin_since=None): score = score + ((self.w_q4 * lcv) >> 4) assert np.abs(score).max() < 2**31 # the kernel computes it in int32 key = np.where(held, ((score + se.SCORE_BIAS) << self.slot_bits) | c, c) + if lowpri: + key = np.where(held & (self.usage == 0), (1 << self.slot_bits) | c, key) evictable = self.usage != self.step if pin_since is not None: evictable &= self.usage <= pin_since return np.where(evictable, key, _KEY_MAX) def ensure(self, layer, ids, *, bump_tok=False, update_state=True, logits=None, thr=0.0, near_miss=None, - pin_since=None): + pin_since=None, lowpri=False): """One ensure call; returns ``(out_slots, src_ids, dst_slots)`` like the kernel's plan. - ``pin_since`` (a ``step`` value) also pins every slot touched after it. + ``pin_since`` (a ``step`` value) also pins every slot touched after it. ``lowpri`` + (FREETOKEN_MOE_PREFETCH=on) keys unconsumed prefetches between the empty slots and the + residents; ``self.useful`` is then the call's distinct hits on such slots. Policy 3 refreshes the ids whose ``logits`` (``[rows, E]``) are within ``thr`` of their row's lowest routed logit, in fp32 as the kernel does, or the ids of a given ``near_miss`` mask. @@ -100,13 +105,15 @@ def ensure(self, layer, ids, *, bump_tok=False, update_state=True, logits=None, self.last_tok[near] = np.maximum(self.last_tok[near], tok - 1) q = np.asarray(ids, dtype=np.int64).reshape(-1) + base slots = self.slot_of_id[q] - self.usage[slots[slots >= 0]] = self.step uniq = np.unique(q) + held_at = self.slot_of_id[uniq] + self.useful = int(((held_at >= 0) & (self.usage[np.maximum(held_at, 0)] == 0)).sum()) + self.usage[slots[slots >= 0]] = self.step missing = uniq[self.slot_of_id[uniq] < 0] # ascending, as the kernel ranks them src = missing - base dst = np.empty(missing.size, np.int64) if missing.size: - keys = self._keys(layer, tok, now, pin_since) + keys = self._keys(layer, tok, now, pin_since, lowpri) part = np.argpartition(keys, missing.size - 1)[: missing.size] victims = part[np.argsort(keys[part])] for i, (e, v) in enumerate(zip(missing, victims)): @@ -127,6 +134,31 @@ def ensure(self, layer, ids, *, bump_tok=False, update_state=True, logits=None, self.ct[uniq] = now return out, src, dst + def prefetch(self, layer, sel, budget=None): + """``prefetch_ensure``: install ``sel``'s non-resident ids (-1 = padding, query order), at most + the first ``budget``, as held slots with usage 0, without bumping a clock or touching a hit; + the last demand call's slots (usage == step) stay pinned. Returns ``(src_ids, dst_slots)``, + the copy plan.""" + base = layer * self.E + tok = self.tok if self.policy else 0 + pos = layer - 1 # the kernel scores from the previous layer's call, where the prefetch runs + now = tok * self.L + pos + order = [] + for e in sel: + if e >= 0 and e + base not in order: + order.append(int(e) + base) + missing = [i for i in order if self.slot_of_id[i] < 0] + keys = self._keys(pos, tok, now, lowpri=True) + n = min(len(missing), int((keys != _KEY_MAX).sum()), len(missing) if budget is None else budget) + victims = np.argsort(keys, kind="stable")[:n] + for e, v in zip(missing[:n], victims): + old = self.id_of_slot[v] + if old >= 0: + self.slot_of_id[old] = -1 + self.id_of_slot[v] = e + self.slot_of_id[e] = v + self.usage[v] = 0 + return np.asarray(missing[:n], np.int64) - base, victims.astype(np.int64) def materialize(self, layer): """The prefill materialize kernel: the whole layer into slots [0, E), position == expert id.""" diff --git a/tests/moe/test_prefetch.py b/tests/moe/test_prefetch.py new file mode 100644 index 0000000000..47c26066a4 --- /dev/null +++ b/tests/moe/test_prefetch.py @@ -0,0 +1,424 @@ +"""FREETOKEN_MOE_PREFETCH stage A: the router-lookahead select / count kernels against a Python +reference, the mode switch, and the predictor stream's lifetime across resets and rebuilds. +The qwen4_exp stack tests (bit-identical outputs, counters over a real decode) live in +tests/models/qwen4_exp/test_skeleton.py.""" + +from __future__ import annotations + +import weakref + +import pytest +import torch + +from .ref_prefetch import ref_count, ref_select + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + + +def _padded(ids, width): + return ids + [-1] * (width - len(ids)) + + +def _select(logits, resident, k, budget, width=16): + from freetoken.moe.prefetch import lookahead_select + + sel = torch.full((width,), -7, dtype=torch.int32, device="cuda") + res = torch.full((width,), -7, dtype=torch.int32, device="cuda") + lookahead_select( + torch.tensor(logits, dtype=torch.bfloat16, device="cuda"), + torch.tensor(resident, dtype=torch.int32, device="cuda"), + sel, res, k=k, budget=budget, + ) + return sel.tolist(), res.tolist() + + +def test_mode_switch(): + from freetoken.moe.prefetch import default_budget, resolve_mode + + assert resolve_mode("off") == "off" and resolve_mode(" Measure ") == "measure" and resolve_mode("ON") == "on" + with pytest.raises(ValueError): + resolve_mode("lookahead") + assert (default_budget(True), default_budget(False)) == (3, 4) + + +@requires_cuda +def test_cache_builds_the_prefetcher_only_when_asked(): + from freetoken.moe.offload_cache import OffloadMoeCache + + def cache(**kw): + return OffloadMoeCache(num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"), **kw) + + assert cache(prefetch_mode="off").prefetch is None + assert cache(prefetch_mode="measure").prefetch.copy_stream is None + on = cache(prefetch_mode="on") + assert on.prefetch.mode == "on" and on.prefetch.cache.prefetch is on.prefetch and on.prefetch_on + # no reference cycle: dropping the cache frees its tensors at once + ref = weakref.ref(on) + del on + assert ref() is None + for mode in ("measure", "on"): + with pytest.raises(ValueError): + cache(prefetch_mode=mode, decode_target="hybrid") + + +@requires_cuda +def test_the_prefetch_env_flag_skips_caches_that_do_not_decode_on_the_gpu(monkeypatch): + # a process-wide flag must not break hybrid or CPU-decoded caches, as FREETOKEN_MOE_COPY_OVERLAP does not + from freetoken.env import ENV + from freetoken.moe.offload_cache import OffloadMoeCache + + monkeypatch.setattr(ENV.MOE_PREFETCH, "value", "on") + for target in ("hybrid", "cpu"): + c = OffloadMoeCache(num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"), decode_target=target) + assert c.prefetch is None and not c.prefetch_on + assert OffloadMoeCache(num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda")).prefetch_on + + +@requires_cuda +def test_select_known_answers(): + # order by logit: 1, 4, 6, 2, 3, 7, 0, 5; experts 4 and 2 are resident + logits = [[0.1, 0.9, 0.5, 0.3, 0.8, 0.0, 0.7, 0.2]] + resident = [-1, -1, 7, -1, 3, -1, -1, -1] + assert _select(logits, resident, 5, 2) == (_padded([1, 6], 16), _padded([4], 16)) + assert _select(logits, resident, 5, 4) == (_padded([1, 6, 3], 16), _padded([4, 2], 16)) + # the top-k cap: 3 and 7 are out of a top-4 list + assert _select(logits, resident, 4, 4) == (_padded([1, 6], 16), _padded([4, 2], 16)) + # all-equal logits rank by id + assert _select([[0.5] * 8], [-1] * 8, 8, 3) == (_padded([0, 1, 2], 16), _padded([], 16)) + # bs 2: row 0 ranks 1, 4, 6 and row 1 ranks 4, 0, 5 -> merged 1, 4, 0, 6, 5; expert 0 resident + rows = [[0.0, 0.9, 0.1, 0.1, 0.8, 0.1, 0.7, 0.1], [0.8, 0.1, 0.1, 0.1, 0.9, 0.7, 0.1, 0.1]] + resident = [2, -1, -1, -1, -1, -1, -1, -1] + assert _select(rows, resident, 3, 3) == (_padded([1, 4, 6], 16), _padded([0], 16)) + assert _select(rows, resident, 3, 16) == (_padded([1, 4, 6, 5], 16), _padded([0], 16)) + + +@requires_cuda +@pytest.mark.parametrize("rows", [1, 2]) +@pytest.mark.parametrize("k,budget", [(3, 3), (16, 4), (16, 3), (8, 16)]) +def test_select_matches_reference(rows, k, budget): + gen = torch.Generator().manual_seed(rows * 100 + k + budget) + for trial in range(20): + # coarse values make bf16 ties common, so the tie rule is exercised + logits = (torch.randint(-24, 24, (rows, 512), generator=gen).float() / 8).tolist() + resident = torch.where(torch.rand(512, generator=gen) < 0.4, 1, -1).tolist() + want = ref_select(logits, resident, k, budget) + width = 2 * k + got = _select(logits, resident, k, budget, width=width) + assert got == (_padded(want[0], width), _padded(want[1], width)), trial + + +@requires_cuda +def test_count_known_answers(): + from freetoken.moe.prefetch import prefetch_count + + num_experts, layer = 8, 2 + base = layer * num_experts + # slots 0, 1, 3 hold this layer's experts 3, 1, 6; slot 4 another layer's expert + id_of_slot = torch.tensor([base + 3, base + 1, -1, base + 6, 5, -1], dtype=torch.int32, device="cuda") + slots = torch.tensor([0, 1, 3], dtype=torch.int32, device="cuda") + sel = torch.tensor(_padded([3, 5], 6), dtype=torch.int32, device="cuda") + res = torch.tensor(_padded([1, 4], 6), dtype=torch.int32, device="cuda") + misses = torch.tensor([2], dtype=torch.int64, device="cuda") + counters = torch.zeros(6, dtype=torch.int64, device="cuda") + for _ in range(2): + prefetch_count(slots, id_of_slot, misses, sel, res, counters, id_base=base, rows=1) + want = ref_count([3, 5], [1, 4], [3, 1, 6], 2, 1) + assert counters.tolist() == [2 * v for v in want] == [4, 2, 2, 2, 2, 4] + # bs 2 with a duplicate routed expert counts it once + counters.zero_() + slots2 = torch.tensor([0, 1, 0, 3], dtype=torch.int32, device="cuda") + prefetch_count(slots2, id_of_slot, misses, sel, res, counters, id_base=base, rows=2) + assert counters.tolist() == ref_count([3, 5], [1, 4], [3, 1, 3, 6], 2, 2) + + +@requires_cuda +def test_prefetch_stream_never_aliases_a_pooled_stream(): + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache( + num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"), + decode_copy_overlap=True, prefetch_mode="measure", + ) + pred = cache.prefetch.stream.cuda_stream + pooled = {torch.cuda.Stream().cuda_stream for _ in range(64)} + assert pred not in pooled and pred != torch.cuda.current_stream().cuda_stream + assert pred != cache.decode_copy_stream.cuda_stream + + +@requires_cuda +def test_prefetch_copy_stream_is_a_third_dedicated_stream(): + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache( + num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"), + decode_copy_overlap=True, prefetch_mode="on", + ) + pf = cache.prefetch + copy = pf.copy_stream.cuda_stream + pooled = {torch.cuda.Stream().cuda_stream for _ in range(64)} + assert copy not in pooled and pf.stream.cuda_stream not in pooled + assert len({copy, pf.stream.cuda_stream, cache.decode_copy_stream.cuda_stream, torch.cuda.current_stream().cuda_stream}) == 4 + # its CTAs go ahead of pending compute ones, so a late copy does not also wait for SMs + from cuda.bindings import runtime as cudart + + _, _, greatest = cudart.cudaDeviceGetStreamPriorityRange() + assert cudart.cudaStreamGetPriority(copy)[1] == greatest < 0 + assert cudart.cudaStreamGetPriority(pf.stream.cuda_stream)[1] == 0 + + +@requires_cuda +def test_rebuild_and_reset_clear_the_counters(): + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache( + num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"), prefetch_mode="measure" + ) + cache.set_bank_sources({ + "gate_up": [torch.zeros(4, 32, 8, dtype=torch.bfloat16).pin_memory() for _ in range(2)], + "down": [torch.zeros(4, 8, 16, dtype=torch.bfloat16).pin_memory() for _ in range(2)], + }) + pf = cache.prefetch + stream, sel = pf.stream, pf.sel + pf.counters.fill_(3) + cache.reset() + assert int(pf.counters.abs().sum()) == 0 + pf.counters.fill_(5) + pf.take_window() + assert int(pf.totals.sum()) > 0 + pf.counters.fill_(2) + pf._inflight[1] = torch.zeros(1) + cache.rebuild(6) + assert cache.prefetch is pf and pf.stream is stream and pf.sel is sel + assert int(pf.counters.abs().sum()) == 0 and int(pf.totals.abs().sum()) == 0 + assert pf._inflight == [None, None] + + +def _emit_prefetch(window, monkeypatch, *, debug=False, mode="measure"): + from types import SimpleNamespace + + from freetoken.engine import engine as engine_mod + from freetoken.env import ENV + + taken, lines = [], [] + prefetch = SimpleNamespace(mode=mode, take_window=lambda: taken.append(1) or window) + engine = SimpleNamespace(moe_offload_cache=SimpleNamespace(prefetch=prefetch)) + monkeypatch.setattr(ENV.MOE_PREFETCH_DEBUG, "value", debug) + # the freetoken loggers do not propagate, so caplog only sees them in some import orders + monkeypatch.setattr(engine_mod.logger, "info_rank0", lambda msg, *a, **k: lines.append(msg)) + engine_mod.Engine._emit_prefetch_stats(engine) + return "\n".join(lines), taken + + +def test_engine_reports_one_prefetch_window(monkeypatch): + from freetoken.engine.engine import MOE_STATS_INTERVAL + + window = torch.tensor([[0] * 6, [30, 12, 4, 10, 10, 40], [40, 20, 6, 10, 10, 30]]) + out, taken = _emit_prefetch(window, monkeypatch) + assert taken == [1] + assert f"MoE prefetch measure ({MOE_STATS_INTERVAL} decode steps): issued/layer=3.50, useful/layer=1.60" in out + assert "per layer" not in out + out, _ = _emit_prefetch(window, monkeypatch, debug=True) + assert "L1=1.20/3.00/4.00, L2=2.00/4.00/3.00" in out and "L0" not in out + assert _emit_prefetch(torch.zeros((3, 6), dtype=torch.int64), monkeypatch)[0] == "" + + +def test_engine_reports_the_on_mode_window(monkeypatch): + from freetoken.engine.engine import MOE_STATS_INTERVAL + + # issued, useful, resident_hits, calls, rows, misses, late, copied + window = torch.tensor([[0] * 8, [30, 12, 0, 10, 10, 18, 2, 10], [40, 20, 0, 10, 10, 20, 0, 6]]) + out, _ = _emit_prefetch(window, monkeypatch, mode="on") + assert f"MoE prefetch on ({MOE_STATS_INTERVAL} decode steps): issued/layer=3.50, useful/layer=1.60" in out + assert "misses/layer=1.90, coverage=0.457, late=0.125 of 16 copies" in out and "resident_hits" not in out + + +def test_summary_rates(): + from freetoken.moe.prefetch import format_per_layer, format_summary, summarize + + # layer 0 never predicts; layers 1 and 2 over 10 bs-1 steps + counters = torch.tensor([[0] * 6, [30, 12, 4, 10, 10, 40], [40, 20, 6, 10, 10, 30]]) + stats = summarize(counters) + assert stats["layer_calls"] == 20 and stats["tokens"] == 10 + assert stats["issued_per_layer"] == 3.5 and stats["useful_per_layer"] == 1.6 + assert stats["precision"] == 32 / 70 and stats["coverage"] == 32 / 70 + assert stats["useful_per_token"] == 3.2 and stats["resident_hits_per_layer"] == 0.5 + assert "useful/token=3.2" in format_summary(stats) + assert format_per_layer(counters) == "L1=1.20/3.00/4.00, L2=2.00/4.00/3.00" + + +# FREETOKEN_MOE_PREFETCH_VERIFY (moe/verify.py): the mode switch, where the verifier is built, and +# the engine's report; the checks themselves run on the qwen4_exp stacks in test_skeleton.py + + +def test_verify_mode_switch(): + from freetoken.moe.verify import resolve_mode + + assert [resolve_mode(m) for m in ("0", "", "off", "1", " True ", "full", "META")] == [ + "off", "off", "off", "full", "full", "full", "meta"] + with pytest.raises(ValueError): + resolve_mode("bytes") + + +@requires_cuda +def test_cache_builds_the_verifier_only_for_gpu_decode(monkeypatch): + from freetoken.env import ENV + from freetoken.moe.offload_cache import OffloadMoeCache + + def cache(**kw): + return OffloadMoeCache(num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"), **kw) + + monkeypatch.setattr(ENV.MOE_PREFETCH_VERIFY, "value", "0") + assert cache().verify is None and cache(verify_mode="off").verify is None + full = cache(verify_mode="full", prefetch_mode="on") + assert full.verify.mode == "full" and cache(verify_mode="meta").verify.mode == "meta" + # the verifier holds no reference to its cache + ref = weakref.ref(full) + del full + assert ref() is None + with pytest.raises(ValueError): + cache(verify_mode="full", decode_target="hybrid") + monkeypatch.setattr(ENV.MOE_PREFETCH_VERIFY, "value", "1") + for target in ("hybrid", "cpu"): + assert cache(decode_target=target).verify is None + assert cache().verify.mode == "full" + + +@requires_cuda +def test_engine_reports_the_verify_window_and_session(monkeypatch): + from types import SimpleNamespace + + from freetoken.engine import engine as engine_mod + from freetoken.engine.engine import MOE_STATS_INTERVAL + from freetoken.moe.offload_cache import OffloadMoeCache + from freetoken.moe.verify import BYTES_PRE, CHECKS, NUM_FIELDS, PRE_BAD + + cache = OffloadMoeCache(num_layers=3, num_experts=4, cache_size=4, device=torch.device("cuda"), verify_mode="full") + cache.set_bank_sources({ + "gate_up": [torch.randn(4, 32, 8, dtype=torch.bfloat16).pin_memory() for _ in range(3)], + "down": [torch.randn(4, 8, 16, dtype=torch.bfloat16).pin_memory() for _ in range(3)], + }) + v = cache.verify + lines = [] + for level in ("info_rank0", "warning_rank0"): + monkeypatch.setattr(engine_mod.logger, level, lambda msg, *a, level=level, **k: lines.append((level, msg))) + engine = SimpleNamespace(moe_offload_cache=cache) + + v.counters[:, CHECKS] = 256 + engine_mod.Engine._emit_verify_stats(engine) + assert lines == [("info_rank0", f"MoE verify full ({MOE_STATS_INTERVAL} decode steps): 768 layer calls checked, " + "meta_bad=0, pre_bad=0, post_bad=0, post_meta_bad=0; session bad=0")] + lines.clear() + v.counters[:, CHECKS] = 256 + v.counters[1, PRE_BAD] = 1 + # call 7, layer 1, row 3 of a top-2 call, expert 2, slot 1, bank down, before the GEMM; slot 1 holds + # layer 1's expert 2, usage 9 at lru step 9, demand-copied; bytes 64.. differ in 4 words + record = [7, 1, 3, 2, 1, 1, BYTES_PRE, 6, 1, 9, 9, 1, 64, 4] + assert len(record) == NUM_FIELDS + v.ring[0] = torch.tensor(record) + v.cursor.fill_(1) + v._top_k = 2 + engine_mod.Engine._emit_verify_stats(engine) + assert [level for level, _ in lines] == ["warning_rank0", "warning_rank0"] + assert "pre_bad=1" in lines[0][1] and "(meta/pre/post/post_meta by layer: L1=0/1/0/0); session bad=1" in lines[0][1] + assert lines[1][1] == ( + "MoE verify record: bytes_pre call=7 layer=1 row=3 (token 1, rank 1) expert=2 slot=1 bank=down " + "first_bad_byte=64 bad_words=4 | slot now holds L1/e2, slot_for_id=1, usage=9, lru_step=9, slot in: demand copy") + lines.clear() + engine_mod.Engine._emit_verify_stats(engine) # nothing checked and no new record: silent + assert lines == [] + session = v.report_session() + assert session[0] == ("warning", "MoE verify full (session): 1536 layer calls checked, meta_bad=0, pre_bad=1, " + "post_bad=0, post_meta_bad=0 (meta/pre/post/post_meta by layer: L1=0/1/0/0)") + assert session[1:] == [("warning", "MoE verify record: " + v.format_record(record))] + + +# FREETOKEN_MOE_SLOT_AUDIT (moe/slot_audit.py): the switch, where the auditor is built, and the +# engine's report; the recorders and the audit run on the qwen4_exp stacks in test_skeleton.py + + +def test_slot_audit_interval_switch(monkeypatch): + from freetoken.env import ENV + from freetoken.moe.slot_audit import resolve_interval + + monkeypatch.setattr(ENV.MOE_SLOT_AUDIT, "value", 64) + assert [resolve_interval(v) for v in (None, 0, 1, 8)] == [64, 0, 1, 8] + with pytest.raises(ValueError): + resolve_interval(-1) + + +@requires_cuda +def test_cache_builds_the_auditor_only_for_gpu_decode(monkeypatch): + from freetoken.env import ENV + from freetoken.moe.offload_cache import OffloadMoeCache + + def cache(**kw): + return OffloadMoeCache(num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"), **kw) + + monkeypatch.setattr(ENV.MOE_SLOT_AUDIT, "value", 0) + assert cache().audit is None and cache(slot_audit=0).audit is None + on = cache(slot_audit=8, prefetch_mode="on") + assert on.audit.interval == 8 and on.slot_audit == 8 + ref = weakref.ref(on) # the auditor holds no reference to its cache + del on + assert ref() is None + with pytest.raises(ValueError): + cache(slot_audit=8, decode_target="hybrid") + monkeypatch.setattr(ENV.MOE_SLOT_AUDIT, "value", 16) + for target in ("hybrid", "cpu"): + assert cache(decode_target=target).audit is None + assert cache().audit.interval == 16 + monkeypatch.setenv("FREETOKEN_SKIP_FAST_INDEX_COPY", "1") + assert cache().audit is None, "copies that move nothing cannot be audited" + + +@requires_cuda +def test_engine_reports_the_audit_window_and_session(monkeypatch): + from types import SimpleNamespace + + from freetoken.engine import engine as engine_mod + from freetoken.engine.engine import MOE_STATS_INTERVAL + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache(num_layers=3, num_experts=4, cache_size=6, device=torch.device("cuda"), slot_audit=2) + cache.set_bank_sources({ + "gate_up": [torch.randn(4, 32, 8, dtype=torch.bfloat16).pin_memory() for _ in range(3)], + "down": [torch.randn(4, 8, 16, dtype=torch.bfloat16).pin_memory() for _ in range(3)], + }) + cache.reset() + lines = [] + for level in ("info_rank0", "warning_rank0"): + monkeypatch.setattr(engine_mod.logger, level, lambda msg, *a, level=level, **k: lines.append((level, msg))) + engine = SimpleNamespace(moe_offload_cache=cache) + engine_mod.Engine._emit_audit_stats(engine) + assert lines == [], "no audit ran: nothing to say" + + # one decode step of layer 1 routing experts 2 and 0: two demand installs and their copy + ids = torch.tensor([[2, 0]], dtype=torch.int32, device="cuda") + cache.ensure_experts(1, ids) + cache.copy_missing() + audit = cache.audit + assert not audit.end_decode_step(cache) and audit.end_decode_step(cache), "every 2 decode steps" + engine_mod.Engine._emit_audit_stats(engine) + assert lines == [("info_rank0", f"MoE slot audit ({MOE_STATS_INTERVAL} decode steps, every 2): 1 audits, 2 held slots " + "an audit (2 byte-checked), bad slots=0 (bytes=0, map=0, stale slot_for_id=0, owner out " + "of range=0), new=0, repeat=0, events=4; session: 1 audits, 0 bad (slot, owner) pairs")] + lines.clear() + slot = int(ids[0, 0]) + cache.bank_views()[1][slot].view(torch.uint8).view(-1)[8:12].bitwise_not_() + audit.scan(cache) + engine_mod.Engine._emit_audit_stats(engine) + assert [level for level, _ in lines] == ["warning_rank0", "warning_rank0"] + assert "1 audits, 2 held slots an audit (2 byte-checked), bad slots=1 (bytes=1," in lines[0][1] + assert lines[1][1].startswith( + f"MoE slot audit record: audit=1 step=0 lru=1 slot={slot} holds L1/e2 (slot_for_id={slot}, usage=1) | bytes differ in " + "down (first at down+8, 1 bad words) | diagnosis: demand install and demand copy of L1/e2 (#"), lines[1][1] + assert "then the bytes changed with no recorded write" in lines[1][1] + # the kernel ranks misses by id: e0 installs first (#0), e2 second (#1); the copy records after both + assert lines[1][1].endswith("history (1 map / 1 byte events in all): #1 demand install L1/e2 @step 0/lru 1 (then held " + "L1/e2, usage 1); #3 demand copy L1/e2 @step 0/lru 1 (slot held L1/e2, usage 1)"), lines[1][1] + lines.clear() + engine_mod.Engine._emit_audit_stats(engine) + assert lines == [] + session = audit.report_session() + assert session[0][0] == "warning" and session[0][1].startswith("MoE slot audit (session, every 2 decode steps): 2 audits") + assert session[1:] == [("warning", "MoE slot audit record: " + audit.format_record(audit.ring[0].tolist()))] diff --git a/tests/moe/test_scored_ensure.py b/tests/moe/test_scored_ensure.py index db734180b6..6bfa84a5ed 100644 --- a/tests/moe/test_scored_ensure.py +++ b/tests/moe/test_scored_ensure.py @@ -6,6 +6,7 @@ (experts.npy, logit_idx.npy, logit_val.npy) to also replay real decode rows. """ import os +from types import SimpleNamespace from pathlib import Path import numpy as np @@ -620,3 +621,321 @@ def test_an_all_pinned_call_never_stores_past_the_slot_arrays(policy): for name, buf in guarded.items(): assert torch.all(buf[s:] == 7), f"{name} written past slot {s - 1}" assert int(cache.num_indices) == 1 and 0 <= int(cache.evict_slots[0]) < s + + +# (h) FREETOKEN_MOE_PREFETCH=on: prefetch installs and low-priority keys ---------------------- + +_PF_L, _PF_E, _PF_S, _PF_K, _PF_W = 4, 64, 100, 6, 16 + + +def _pf_cache(policy, cache_size=_PF_S, num_layers=_PF_L, num_experts=_PF_E): + return _cache(policy, cache_size, num_layers, num_experts, prefetch_mode="on") + + +def _gpu_prefetch(cache, layer, sel, stats_row, width=_PF_W, budget=None): + """prefetch_ensure on the current stream; returns (src_ids, dst_slots), the plan it wrote.""" + from freetoken.moe.offload_kernels import prefetch_ensure_experts + + dev = cache.id_of_slot.device + query = torch.tensor(list(sel) + [-1] * (width - len(sel)), dtype=torch.int32, device=dev) + dst = torch.full((width,), -7, dtype=torch.int32, device=dev) + src = torch.full((width,), -7, dtype=torch.int32, device=dev) + num = torch.full((1,), 99, dtype=torch.int64, device=dev) + ready = torch.ones(1, dtype=torch.int32, device=dev) + prefetch_ensure_experts(cache, layer, query, dst, src, num, stats_row, ready, budget=budget) + n = int(num) + assert int(ready) == 0, "prefetch_ensure clears the copy's ready flag" + assert dst[n:].tolist() == [-7] * (width - n) + return src[:n].cpu().numpy().astype(np.int64), dst[:n].cpu().numpy().astype(np.int64), num + + +def _lowpri_slots(cache): + return set(((cache.id_of_slot >= 0) & (cache.usage == 0)).nonzero().view(-1).tolist()) + + +@_cuda +@pytest.mark.parametrize("policy", ["lru", "kd", "kdfb", "rule"]) +@pytest.mark.parametrize("rows_per_step", [1, 2]) +def test_prefetch_mode_matches_cpu_reference(policy, rows_per_step): + """prefetch_ensure and the low-priority demand ensure against the CPU reference victim for victim, + with adversarial candidates (next routed ids, resident ids, duplicates, padding, junk) and a + materialized prefill in between; the per-layer prefetch stats match the reference's counts.""" + from freetoken.moe.offload_kernels import ensure_experts_scored + from freetoken.moe.prefetch import CALLS, COPIED, ISSUED, LATE, MISSES, RESIDENT_HITS, ROWS, USEFUL + + rng = np.random.default_rng(31 + rows_per_step) + rows = _zipf_rows(40, rows_per_step, seed=33, num_layers=_PF_L, num_experts=_PF_E, top_k=_PF_K) + logits = _random_logits(rows, seed=34, num_experts=_PF_E) if policy == "rule" else None + cache, ref = _pf_cache(policy), _ref(policy, _PF_S, _PF_L, _PF_E) + dev = cache.id_of_slot.device + stats = torch.zeros((_PF_L, 8), dtype=torch.int64, device=dev) + want = np.zeros((_PF_L, 8), np.int64) + ready = torch.zeros(1, dtype=torch.int32, device=dev) + for r in range(rows.shape[0]): + if r == 20: + cache.materialize_layer(2) + ref.materialize(2) + assert not {s for s in _lowpri_slots(cache) if s < _PF_E}, "materialize installs at usage step" + for layer in range(_PF_L): + ids = rows[r, layer] + plan = None + if layer and rng.random() < 0.9: + routed = list(dict.fromkeys(ids.ravel().tolist())) + held = [e for e in range(_PF_E) if ref.slot_of_id[layer * _PF_E + e] >= 0] + sel = rng.choice(routed, size=min(3, len(routed)), replace=False).tolist() + sel += rng.choice(_PF_E, size=4).tolist() + held[:2] + sel[:1] + [-1] + rng.shuffle(sel) + pinned = set(np.flatnonzero((ref.usage == ref.step) & (ref.id_of_slot >= 0)).tolist()) + budget = [None, 0, 1, 2, 4][int(rng.integers(5))] + src, dst, num = _gpu_prefetch(cache, layer, sel, stats[layer], budget=budget) + want_src, want_dst = ref.prefetch(layer, sel, budget) + np.testing.assert_array_equal(src, want_src) + np.testing.assert_array_equal(dst, want_dst) + assert not pinned & set(dst.tolist()), "a prefetch evicted a slot of the last demand call" + assert set(dst.tolist()) <= _lowpri_slots(cache) + want[layer, ISSUED] += len(want_src) + _assert_same_tables(cache, ref) + plan = num + ready.fill_(int(rng.random() < 0.5)) # a copy that has (1) or has not (0) landed + q = torch.from_numpy(np.ascontiguousarray(ids, dtype=np.int32)).to(dev).view(-1) + lg = None if logits is None else torch.from_numpy(np.ascontiguousarray(logits[r, layer])).to(dev, torch.bfloat16) + lowpri_before = _lowpri_slots(cache) + ensure_experts_scored( + cache, layer, q, bump_tok=layer == 0, update_state=True, router_logits=lg, lowpri=True, + pf_count=None if plan is None else (stats[layer], ready, plan, rows_per_step), + ) + out, src, _ = ref.ensure(layer, ids, bump_tok=layer == 0, logits=None if lg is None else lg.float().cpu().numpy(), + thr=NEAR_MISS_THR, lowpri=True) + np.testing.assert_array_equal(q.cpu().numpy(), out) + _assert_same_tables(cache, ref) + # a demand hit makes a prefetched slot an ordinary resident, tagged with its owner's state + converted = lowpri_before & set(out.tolist()) + assert all(int(cache.usage[s]) == int(cache.step) for s in converted) + if policy != "lru": + assert all(int(cache.evict_slot_owner[s]) == int(cache.id_of_slot[s]) for s in converted) + if plan is not None: + n = int(plan) + want[layer, [USEFUL, CALLS, ROWS, MISSES]] += [ref.useful, 1, rows_per_step, src.size] + want[layer, LATE] += int(n > 0 and int(ready) == 0) + want[layer, COPIED] += int(n > 0) + got = stats.cpu().numpy() + np.testing.assert_array_equal(got, want) + assert got[:, USEFUL].sum() > 0 and got[:, LATE].sum() > 0 and got[:, RESIDENT_HITS].sum() == 0 + assert got[:, COPIED].sum() > got[:, LATE].sum() + + +@_cuda +@pytest.mark.parametrize("policy", ["lru", "rule"]) +@pytest.mark.parametrize("rows_per_step", [1, 2, 8]) +def test_prefetch_mode_matches_cpu_reference_at_production_geometry(policy, rows_per_step): + """The same victim-for-victim check at the deployment's shape: 48 layers x 512 experts, 1480 + slots (a 2048-wide, 8-warp victim scan), top-10, up to 8 rows (every graph size that forks) + and 32-wide prefetch queries, with small-prefill ensures (pinned since their first chunk) and a + materialized layer in between. No prefetch ever evicts a slot of the last demand call.""" + from freetoken.moe.offload_kernels import ensure_experts_scored + + num_layers, num_experts, size, width = L, E, 1480, 32 + rng = np.random.default_rng(71 + rows_per_step) + rows = _zipf_rows(4, rows_per_step, seed=72, num_layers=num_layers, num_experts=num_experts, top_k=TOP_K) + logits = _random_logits(rows, seed=73, num_experts=num_experts) if policy == "rule" else None + cache, ref = _pf_cache(policy, size, num_layers, num_experts), _ref(policy, size, num_layers, num_experts) + dev = cache.id_of_slot.device + stats = torch.zeros((num_layers, 8), dtype=torch.int64, device=dev) + ready = torch.zeros(1, dtype=torch.int32, device=dev) + installs = 0 + for r in range(rows.shape[0]): + if r == 2: + cache.materialize_layer(5) + ref.materialize(5) + for layer in range(num_layers): + ids = rows[r, layer] + plan = None + if layer: + routed = list(dict.fromkeys(ids.ravel().tolist())) + held = [e for e in range(num_experts) if ref.slot_of_id[layer * num_experts + e] >= 0] + sel = rng.choice(routed, size=min(6, len(routed)), replace=False).tolist() + sel += rng.choice(num_experts, size=16).tolist() + held[:4] + sel[:2] + [-1, -1] + sel = sel[:width] + rng.shuffle(sel) + pinned = set(np.flatnonzero((ref.usage == ref.step) & (ref.id_of_slot >= 0)).tolist()) + src, dst, plan = _gpu_prefetch(cache, layer, sel, stats[layer], width=width) + want_src, want_dst = ref.prefetch(layer, sel) + np.testing.assert_array_equal(src, want_src) + np.testing.assert_array_equal(dst, want_dst) + assert not pinned & set(dst.tolist()), "a prefetch evicted a slot of the last demand call" + installs += len(dst) + _assert_same_tables(cache, ref) + q = torch.from_numpy(np.ascontiguousarray(ids, dtype=np.int32)).to(dev).view(-1) + lg = None if logits is None else torch.from_numpy(np.ascontiguousarray(logits[r, layer])).to(dev, torch.bfloat16) + ensure_experts_scored( + cache, layer, q, bump_tok=layer == 0, update_state=True, router_logits=lg, lowpri=True, + pf_count=None if plan is None else (stats[layer], ready, plan, rows_per_step), + ) + out, _, _ = ref.ensure(layer, ids, bump_tok=layer == 0, logits=None if lg is None else lg.float().cpu().numpy(), + thr=NEAR_MISS_THR, lowpri=True) + np.testing.assert_array_equal(q.cpu().numpy(), out) + _assert_same_tables(cache, ref) + # a short prefill of one layer between decode steps: chunked, pinned since its first chunk + layer = int(rng.integers(num_layers)) + uniq = np.unique(rng.choice(num_experts, size=60, replace=False)) + pin = cache.step.clone() if policy != "lru" else None + ref_pin = ref.step if policy != "lru" else None + for start in range(0, uniq.size, 32): + part = torch.from_numpy(uniq[start : start + 32].astype(np.int32)).to(dev) + cache.ensure_experts(layer, part, update_state=False, pin_since=pin) + out, _, _ = ref.ensure(layer, uniq[start : start + 32], update_state=False, pin_since=ref_pin, lowpri=True) + np.testing.assert_array_equal(part.cpu().numpy(), out) + _assert_same_tables(cache, ref) + assert installs > 0 and int((ref.id_of_slot >= 0).sum()) == size, "the cache must be full and evicting" + + +def _plant(cache, ref, slot, flat_id, usage): + for t_id, t_slot in ((cache.id_of_slot, cache.slot_for_id.view(-1)), (ref.id_of_slot, ref.slot_of_id)): + t_id[slot] = flat_id + t_slot[flat_id] = slot + cache.usage[slot] = usage + ref.usage[slot] = usage + + +@_cuda +@pytest.mark.parametrize("policy", ["lru", "kd", "kdfb", "rule"]) +def test_lowpri_band_sits_between_empty_slots_and_every_resident(policy): + """An unconsumed prefetch is evicted after every empty slot and before any resident, even the + coldest resident the score caps allow, whatever the slot order.""" + from freetoken.moe import scored_ensure as se + from freetoken.moe.offload_kernels import ensure_experts_scored + + num_layers, num_experts, size = 2, 16, 16 + cache, ref = _pf_cache(policy, size, num_layers, num_experts), _ref(policy, size, num_layers, num_experts) + dev = cache.id_of_slot.device + cache.step.fill_(5) + ref.step = 5 + # slots 0-12: layer-0 residents at the oldest usage 1 with never-seen state (the lowest possible + # score); slot 13 empty; slots 14, 15 unconsumed prefetches + for slot in range(13): + _plant(cache, ref, slot, slot, 1) + _plant(cache, ref, 14, 13, 0) + _plant(cache, ref, 15, 14, 0) + if policy != "lru": + cache.evict_slot_owner.fill_(-1) + victims = [] + for e in (15, 16 + 0, 16 + 1, 16 + 2): # four misses, one per call + layer, expert = divmod(e, num_experts) + q = torch.tensor([expert], dtype=torch.int32, device=dev) + ensure_experts_scored(cache, layer, q, bump_tok=False, update_state=True, lowpri=True) + ref.ensure(layer, [expert], lowpri=True) + victims.append(int(cache.evict_slots[0])) + _assert_same_tables(cache, ref) + assert victims[:3] == [13, 14, 15] and victims[3] < 13 + # the band's keys are below 2 << SLOT_BITS; the caps' worst resident score keeps score + 2^31 >= 2 + beta_step, w_q4, _, _ = se.score_params(se.MAX_BETA, se.MAX_W, 1, 1) + worst = -((se.K_MAX << se.Q) + 1 * beta_step) + ((w_q4 * se.LC_FLOOR) >> 4) + assert worst + se.SCORE_BIAS >= 2 + + +@_cuda +@pytest.mark.parametrize("policy", ["lru", "rule"]) +def test_prefetch_budget_keeps_the_first_non_resident_candidates(policy): + """The install budget counts only non-resident candidates, in query order: resident ids, padding + and duplicates ahead of them take none of it, and ISSUED counts the installs.""" + num_experts = 16 + cache, ref = _pf_cache(policy, num_experts, 2, num_experts), _ref(policy, num_experts, 2, num_experts) + stats = torch.zeros(8, dtype=torch.int64, device="cuda") + cache.ensure_experts(0, torch.tensor([0, 1], dtype=torch.int32, device="cuda")) # step 0 pins every slot + ref.ensure(0, np.arange(2), bump_tok=True, lowpri=True) + src, _, _ = _gpu_prefetch(cache, 1, [4, 6], stats) + ref.prefetch(1, [4, 6]) + assert src.tolist() == [4, 6] + sel = [4, -1, 6, 9, 9, 3, 7, 1] + src, dst, _ = _gpu_prefetch(cache, 1, sel, stats, budget=2) + assert src.tolist() == [9, 3] and int(stats[0]) == 4 + np.testing.assert_array_equal(np.stack(ref.prefetch(1, sel, 2)), np.stack([src, dst])) + _assert_same_tables(cache, ref) + src, _, _ = _gpu_prefetch(cache, 1, [5, 8], stats, budget=0) + assert src.size == 0 and int(stats[0]) == 4 + + +@_cuda +@pytest.mark.parametrize("policy", ["lru", "rule"]) +def test_prefetch_installs_only_into_evictable_slots(policy): + """With fewer evictable slots than candidates a prefetch installs the first candidates only, and + with every slot pinned by the last demand call it installs nothing (no all-pinned fallback).""" + num_experts = 16 + cache, ref = _pf_cache(policy, num_experts, 2, num_experts), _ref(policy, num_experts, 2, num_experts) + stats = torch.zeros(8, dtype=torch.int64, device="cuda") + ids = np.arange(12) + cache.ensure_experts(0, torch.from_numpy(ids.astype(np.int32)).cuda()) + ref.ensure(0, ids, bump_tok=True, lowpri=True) + pinned = cache.id_of_slot.clone() + src, dst, _ = _gpu_prefetch(cache, 1, [9, 3, 7, 1, 5, 11, 2], stats) + assert src.tolist() == [9, 3, 7, 1] and sorted(dst.tolist()) == [12, 13, 14, 15] + np.testing.assert_array_equal(np.stack(ref.prefetch(1, [9, 3, 7, 1, 5, 11, 2])), np.stack([src, dst])) + assert torch.equal(cache.id_of_slot[:12], pinned[:12]) and int(stats[0]) == 4 + all16 = np.arange(16) + cache.ensure_experts(0, torch.from_numpy(all16.astype(np.int32)).cuda()) + before = [t.clone() for t in (cache.id_of_slot, cache.slot_for_id, cache.usage)] + src, dst, num = _gpu_prefetch(cache, 1, [0, 1, 2], stats) + assert int(num) == 0 and src.size == 0 and int(stats[0]) == 4 + for a, b in zip(before, (cache.id_of_slot, cache.slot_for_id, cache.usage)): + assert torch.equal(a, b) + + +@_cuda +@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"]) +def test_prefetch_scores_victims_from_the_previous_layers_position(policy): + """prefetch_ensure(L) runs right after layer L-1's demand call and scores victims from there: + layer L's residents are next up (d = 1), so one used last token outlives an older resident of + a later layer instead of looking a whole pass away.""" + num_layers, num_experts, size = 8, 16, 16 + cache, ref = _pf_cache(policy, size, num_layers, num_experts), _ref(policy, size, num_layers, num_experts) + cache.step.fill_(7) + ref.step = 7 + cache.evict_tok.fill_(10) + ref.tok = 10 + cache.evict_slot_owner.fill_(-1) + # slots 0-13: layer 1's routed experts (pinned), slot 14: a layer-2 expert used last token, + # slot 15: a layer-5 expert last used three tokens ago; all with the same (never-counted) use count + planted = [(slot, 16 + slot, 7, 10) for slot in range(14)] + [(14, 2 * 16 + 3, 6, 9), (15, 5 * 16 + 4, 5, 7)] + for slot, flat, usage, last in planted: + _plant(cache, ref, slot, flat, usage) + cache.evict_last_tok[flat] = last + ref.last_tok[flat] = last + stats = torch.zeros(8, dtype=torch.int64, device="cuda") + src, dst, _ = _gpu_prefetch(cache, 2, [9], stats) + assert src.tolist() == [9] and dst.tolist() == [15], "the prefetch evicted the target layer's fresh resident" + np.testing.assert_array_equal(np.stack(ref.prefetch(2, [9])), np.stack([src, dst])) + _assert_same_tables(cache, ref) + + +@_cuda +@pytest.mark.parametrize("rows_per_step", [1, 2]) +def test_lowpri_lru_without_prefetches_is_flashlib(rows_per_step): + """Prefetch on routes LRU through the vendored kernel with the low-priority key; while no slot is + a prefetch it evicts exactly what flashlib's lru_ensure does.""" + num_layers, num_experts, size = 8, 64, 200 + on = _pf_cache("lru", size, num_layers, num_experts) + off = _cache("lru", size, num_layers, num_experts) + rows = _zipf_rows(30, rows_per_step, seed=41, num_layers=num_layers, num_experts=num_experts) + ids = torch.from_numpy(rows.astype(np.int32)).cuda() + for r in range(rows.shape[0]): + for layer in range(num_layers): + a, b = ids[r, layer].clone(), ids[r, layer].clone() + on.ensure_experts(layer, a) + off.ensure_experts(layer, b) + assert torch.equal(a, b) + assert torch.equal(on.evict_slots[: int(on.num_indices)], off.evict_slots[: int(off.num_indices)]) + for name in ("slot_for_id", "id_of_slot", "usage", "step"): + assert torch.equal(getattr(on, name), getattr(off, name)), name + + +def test_lru_policy_leaves_flashlib_only_when_prefetch_is_on(monkeypatch): + import freetoken.moe.offload_kernels as ok + + calls = [] + monkeypatch.setattr(ok, "lru_ensure", lambda *a, **kw: calls.append("flashlib")) + monkeypatch.setattr(ok, "ensure_experts_scored", lambda *a, **kw: calls.append(("scored", kw["lowpri"]))) + for mode in ("off", "measure", "on"): + cache = _cache("lru", 12, 2, 8, device="cpu", prefetch_mode=mode) + cache.prefetch = SimpleNamespace(mode=mode, count_args=lambda layer: None) if mode != "off" else None + cache.ensure_experts(1, torch.zeros(2, dtype=torch.int32)) + assert calls == ["flashlib", "flashlib", ("scored", True)] From 05cb9f4de803785c806d0b3038dfbfd5515adff9 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 1 Oct 2026 23:18:22 -0400 Subject: [PATCH 34/73] test(engine): give the kv quant config fake the rope table size #573's max_seq_len reads Assisted-by: Claude Opus 5.5 --- tests/engine/test_kv_quant_config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/engine/test_kv_quant_config.py b/tests/engine/test_kv_quant_config.py index ba4196be09..1f460166d1 100644 --- a/tests/engine/test_kv_quant_config.py +++ b/tests/engine/test_kv_quant_config.py @@ -39,7 +39,7 @@ def _model_config(kind): has_swa_attention=False, has_linear_attention=False, num_layers=4, - rotary_config=SimpleNamespace(max_position=1024), + rotary_config=SimpleNamespace(max_position=1024, table_positions=1024), ) specs = { "full": (_spec("full", AttnType.FULL),), From a8aa8125aaa73e2fe6e670154fe6560216d680ae Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 06:55:23 -0400 Subject: [PATCH 35/73] fix(engine): hand the rendezvous listen fd to TCPStore instead of keeping its socket TCPStore owns master_listen_fd and closes it with the store; #461 also kept the Python socket on the engine, whose finalizer then closed whatever file reused that fd number after destroy_process_group at shutdown. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit b54cced6af48bffcce88c280124692d54c69f792) --- python/freetoken/engine/engine.py | 8 ++--- tests/engine/test_distributed_store.py | 49 +++++++++++++++++++++++++- 2 files changed, 52 insertions(+), 5 deletions(-) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 9d268a66ee..57519c4eca 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -640,12 +640,12 @@ def _make_distributed_store(self, config: EngineConfig) -> torch.distributed.Sto sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) sock.bind(("127.0.0.1", config.distributed_port)) sock.listen(1024) - # TCPStore uses this fd directly rather than duplicating it, so it must outlive the - # store; kept alive on self, closed only when the process (and the store) exits. - self._distributed_listen_socket = sock + self._distributed_listen_addr = sock.getsockname() + # TCPStore takes ownership of the fd and closes it with the store; detach so the + # Python socket never closes that number again after it has been reused. return torch.distributed.TCPStore( "127.0.0.1", config.distributed_port, config.tp_info.size, - is_master=True, timeout=timeout, multi_tenant=True, master_listen_fd=sock.fileno(), + is_master=True, timeout=timeout, multi_tenant=True, master_listen_fd=sock.detach(), ) def _init_communication(self, config: EngineConfig) -> torch.distributed.ProcessGroup: diff --git a/tests/engine/test_distributed_store.py b/tests/engine/test_distributed_store.py index f70be12668..4fc2473762 100644 --- a/tests/engine/test_distributed_store.py +++ b/tests/engine/test_distributed_store.py @@ -8,7 +8,9 @@ from __future__ import annotations +import gc import multiprocessing as mp +import os from datetime import timedelta from types import SimpleNamespace @@ -38,7 +40,7 @@ def _run_rank(rank: int, port: int, result_q: mp.Queue) -> None: ) total = torch.tensor([float(rank)]) dist.all_reduce(total) - bind_host = engine._distributed_listen_socket.getsockname()[0] if rank == 0 else None + bind_host = engine._distributed_listen_addr[0] if rank == 0 else None result_q.put((rank, total.item(), bind_host)) dist.destroy_process_group() @@ -64,3 +66,48 @@ def test_two_ranks_rendezvous_with_loopback_only_store(): # the master's listening socket was pre-bound to loopback, not the wildcard # address the C10d TCPStore server binds to on its own assert results[0][1] == "127.0.0.1" + + +def _teardown_rank(port: int, result_q: mp.Queue) -> None: + listen_fds = [] + real_store = dist.TCPStore + + def recording_store(*args, **kwargs): + listen_fds.append(kwargs.get("master_listen_fd")) + return real_store(*args, **kwargs) + + dist.TCPStore = recording_store + engine = Engine.__new__(Engine) + config = SimpleNamespace( + tp_info=DistributedInfo(0, 1), distributed_timeout=TIMEOUT.total_seconds(), + distributed_port=port, + ) + store = engine._make_distributed_store(config) + dist.init_process_group(backend="gloo", rank=0, world_size=1, timeout=TIMEOUT, store=store) + dist.destroy_process_group() + del store + gc.collect() + # the listening fd number is reused by an unrelated file, as a later open() would + fd = listen_fds[0] + r, w = os.pipe() + os.dup2(r, fd) + del engine + gc.collect() + try: + os.fstat(fd) + result_q.put("open") + except OSError as exc: + result_q.put(f"closed: {exc}") + + +def test_the_engine_does_not_close_the_listen_fd_the_store_owns(): + # TCPStore takes ownership of master_listen_fd and closes it itself; an engine that also + # kept the Python socket closed whatever file reused that number when it was freed + ctx = mp.get_context("spawn") + result_q = ctx.Queue() + proc = ctx.Process(target=_teardown_rank, args=(29513, result_q)) + proc.start() + outcome = result_q.get(timeout=30) + proc.join(timeout=30) + assert proc.exitcode == 0 + assert outcome == "open" From 5f777ecded94f53ff36f20aba679dcc60a1a5799 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 06:52:14 -0400 Subject: [PATCH 36/73] fix(scheduler): trim the stop-match EOS only at the delivered output limit #464 used `not req.can_decode`, but since #546 overlap runs device_len a step ahead of the host, so a non-terminal EOS was dropped from the matched text one step early. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 1814a3acc1848366a59d16e40d25c3d66a3be301) --- python/freetoken/scheduler/scheduler.py | 4 +- tests/scheduler/test_match_stop_str.py | 71 +++++++++++++++++++++++++ 2 files changed, 74 insertions(+), 1 deletion(-) create mode 100644 tests/scheduler/test_match_stop_str.py diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index 01f1678761..895f012c75 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -478,9 +478,11 @@ def _match_stop_str(self, req: Req) -> str | None: prompt_len = req.max_device_len - req.output_len end = len(req.input_ids) # The frontend omits a terminal EOS, including at the output limit with ignore_eos. + # The limit is the delivered length, as in hit_length: under overlap device_len (and + # so can_decode) runs a step ahead of the host. if ( end > prompt_len - and not req.can_decode + and end >= req.max_device_len and int(req.input_ids[-1]) in self.eos_token_ids ): end -= 1 diff --git a/tests/scheduler/test_match_stop_str.py b/tests/scheduler/test_match_stop_str.py new file mode 100644 index 0000000000..351509b759 --- /dev/null +++ b/tests/scheduler/test_match_stop_str.py @@ -0,0 +1,71 @@ +"""Scheduler-side stop-string matching (Scheduler._match_stop_str) against the same +incrementally decoded text the frontend streams.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import torch + +from freetoken.core import Req, SamplingParams +from freetoken.scheduler.scheduler import Scheduler + +EOS = 0 + + +class _ByteTok: + """Byte-level ids like Qwen's BPE fallback: a 4-byte emoji is four tokens. 0 is EOS.""" + + def decode(self, ids, skip_special_tokens=False): + return b"".join(b"" if t == EOS else bytes([t]) for t in ids).decode( + "utf-8", errors="replace" + ) + + +def _req(prompt: list[int], output_len: int, **sp) -> Req: + return Req( + input_ids=torch.tensor(prompt, dtype=torch.int32), + table_idx=0, + cached_len=0, + output_len=output_len, + uid=1, + sampling_params=SamplingParams(**sp), + cache_handle=None, + ) + + +def _match(req: Req) -> str | None: + stub = SimpleNamespace(tokenizer=_ByteTok(), eos_token_ids={EOS}) + return Scheduler._match_stop_str(stub, req) + + +def _feed(req: Req, ids: list[int]) -> list[str | None]: + seen = [] + for t in ids: + req.append_host(torch.tensor([t], dtype=torch.int32)) + req.complete_one() + seen.append(_match(req)) + return seen + + +def test_a_stop_spanning_more_tokens_than_characters_is_matched(): + dna = list("\N{DNA DOUBLE HELIX}".encode()) # one character, four byte tokens + req = _req([65, 66], 32, stop_strs=["\N{DNA DOUBLE HELIX}"]) + text = list(b"Here is DNA ") + dna + list(b" and more") + seen = _feed(req, text) + first = next(i for i, m in enumerate(seen) if m is not None) + assert first == len(b"Here is DNA ") + len(dna) - 1 + + +def test_an_eos_short_of_the_delivered_limit_stays_in_the_matched_text(): + # Overlap scheduling advances device_len ahead of the host: the budget is spent on the + # device (can_decode is False) while the host holds one token less than the limit. That + # EOS is not terminal, so the frontend decodes it and the scheduler must too. + req = _req([65], 3, stop_strs=["zz"], ignore_eos=True) + for t in (65, EOS): + req.append_host(torch.tensor([t], dtype=torch.int32)) + while req.can_decode: + req.complete_one() + assert req.input_ids.numel() < req.max_device_len + assert _match(req) is None + assert req.stop_decode_status.decoded_ids == [65, EOS] From 25d1ed0bcf7dc275ccae7e494e27de34827820ad Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 06:52:27 -0400 Subject: [PATCH 37/73] test(tokenizer): cover a detokenize batch that carries one uid more than once Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 577427a2cecf047ac64616854a51aa338197336b) --- tests/tokenizer/test_detokenize_extras.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/tokenizer/test_detokenize_extras.py b/tests/tokenizer/test_detokenize_extras.py index 5617a10537..4717617b9a 100644 --- a/tests/tokenizer/test_detokenize_extras.py +++ b/tests/tokenizer/test_detokenize_extras.py @@ -79,3 +79,19 @@ def test_reasoning_tokens_count_up_to_the_end_tag(): assert reasoning == 3 # b, c and the tag itself _, none = _run(manager, 2, [1, 2, EOS]) assert none == 0 + + +def test_a_batch_with_one_uid_twice_advances_its_state_in_order(): + # each message must see the decode state the previous one left, and the reasoning + # counts must stay aligned with the texts + manager = DetokenizeManager(_Tok(), frozenset({EOS}), think_end_id=THINK_END) + msgs = [ + DetokenizeMsg(uid=1, next_token=1, finished=False), + DetokenizeMsg(uid=1, next_token=THINK_END, finished=False), + DetokenizeMsg(uid=2, next_token=3, finished=False), + DetokenizeMsg(uid=1, next_token=2, finished=True), + ] + texts, counts = manager.detokenize_with_meta(msgs) + assert texts == ["b", "", "d", "c"] + assert counts == [0, 0, 0, 2] + assert manager.detokenize([DetokenizeMsg(uid=2, next_token=4, finished=True)]) == ["e"] From 5d76e35ee3d75532a505402de40c70d3ce334d25 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 06:53:13 -0400 Subject: [PATCH 38/73] fix(scheduler): match stop strings in the text the request's skip_special_tokens decodes #464's scheduler-side DecodeStatus decoded with special tokens while #393 lets a request skip them on the frontend, so a stop spanning a special token matched different text. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit fc2392b5367e841b329454bddf33514f808417e1) --- python/freetoken/scheduler/scheduler.py | 7 ++++++- python/freetoken/tokenizer/detokenize.py | 7 +++++-- tests/scheduler/test_match_stop_str.py | 25 +++++++++++++++++++----- 3 files changed, 31 insertions(+), 8 deletions(-) diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index 895f012c75..cb6085d344 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -490,7 +490,12 @@ def _match_stop_str(self, req: Req) -> str | None: return None if req.stop_decode_status is None: req.stop_decode_status = DecodeStatus( - decoded_ids=[], decoded_str="", read_offset=0, surr_offset=0, sent_offset=0 + decoded_ids=[], + decoded_str="", + read_offset=0, + surr_offset=0, + sent_offset=0, + skip_special=req.sampling_params.skip_special_tokens, ) state = req.stop_decode_status state.decoded_ids.extend(req.input_ids[prompt_len + len(state.decoded_ids) : end].tolist()) diff --git a/python/freetoken/tokenizer/detokenize.py b/python/freetoken/tokenizer/detokenize.py index b4c0b04733..2eb3029df0 100644 --- a/python/freetoken/tokenizer/detokenize.py +++ b/python/freetoken/tokenizer/detokenize.py @@ -84,9 +84,12 @@ def update(self, read_str: str, surr_str: str) -> str: return self.decoded_str + find_printable_text(new_text) def decode(self, tokenizer: PreTrainedTokenizerBase) -> str: + skip = self.skip_special return self.update( - tokenizer.decode(self.decoded_ids[self.surr_offset :]), - tokenizer.decode(self.decoded_ids[self.surr_offset : self.read_offset]), + tokenizer.decode(self.decoded_ids[self.surr_offset :], skip_special_tokens=skip), + tokenizer.decode( + self.decoded_ids[self.surr_offset : self.read_offset], skip_special_tokens=skip + ), ) diff --git a/tests/scheduler/test_match_stop_str.py b/tests/scheduler/test_match_stop_str.py index 351509b759..e9a5c07da5 100644 --- a/tests/scheduler/test_match_stop_str.py +++ b/tests/scheduler/test_match_stop_str.py @@ -10,16 +10,22 @@ from freetoken.core import Req, SamplingParams from freetoken.scheduler.scheduler import Scheduler -EOS = 0 +EOS, SPECIAL = 0, 1 class _ByteTok: - """Byte-level ids like Qwen's BPE fallback: a 4-byte emoji is four tokens. 0 is EOS.""" + """Byte-level ids like Qwen's BPE fallback: a 4-byte emoji is four tokens. 0 is EOS, + 1 a special marker that skip_special_tokens drops.""" def decode(self, ids, skip_special_tokens=False): - return b"".join(b"" if t == EOS else bytes([t]) for t in ids).decode( - "utf-8", errors="replace" - ) + def piece(t): + if t == EOS: + return b"" + if t == SPECIAL: + return b"" if skip_special_tokens else b"" + return bytes([t]) + + return b"".join(piece(t) for t in ids).decode("utf-8", errors="replace") def _req(prompt: list[int], output_len: int, **sp) -> Req: @@ -69,3 +75,12 @@ def test_an_eos_short_of_the_delivered_limit_stays_in_the_matched_text(): assert req.input_ids.numel() < req.max_device_len assert _match(req) is None assert req.stop_decode_status.decoded_ids == [65, EOS] + + +def test_stops_match_the_text_decoded_with_the_requests_skip_special_tokens(): + # the frontend streams "ab" for a request that skips special tokens; the scheduler has + # to match the stop in that same text, not in "ab" + req = _req([65], 8, stop_strs=["ab"], skip_special_tokens=True) + assert _feed(req, [ord("a"), SPECIAL, ord("b")]) == [None, None, "ab"] + req = _req([65], 8, stop_strs=["ab"]) + assert _feed(req, [ord("a"), SPECIAL, ord("b")]) == [None, None, None] From dd9cf451a424ef0f186a8be4c573833dd7a5296f Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 05:43:14 -0400 Subject: [PATCH 39/73] perf(ple): skip the device n-gram hash for the disk table, which reads only its shape DiskRowTable hashes on the host in its fill; its lookup and prefetch read only the row_ids shape, so every prefill chunk and captured decode step ran the hash for nothing. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 52a18015ffff759a9f2b7733954bd9998b900f32) --- python/freetoken/models/qwen4_exp/ple.py | 17 ++++++-- python/freetoken/models/qwen4_exp/ple_disk.py | 3 ++ tests/models/qwen4_exp/test_ple.py | 40 +++++++++++++++++++ 3 files changed, 57 insertions(+), 3 deletions(-) diff --git a/python/freetoken/models/qwen4_exp/ple.py b/python/freetoken/models/qwen4_exp/ple.py index 44973bc0a6..22bda9e6c8 100644 --- a/python/freetoken/models/qwen4_exp/ple.py +++ b/python/freetoken/models/qwen4_exp/ple.py @@ -581,8 +581,19 @@ def row_ids_reference(self, meta: PLEMetadata) -> torch.Tensor: blocks.append(head_ids + self.ngram_heads_offsets[start:end]) return torch.cat(blocks, dim=-1) + def lookup_row_ids(self, meta: PLEMetadata) -> torch.Tensor: + """``row_ids`` as the attached table's ``prefetch``/``lookup`` take them. A table that + hashes on the host (``reads_row_ids = False``) reads only their shape: skip the hash.""" + if not getattr(self.table, "reads_row_ids", True): + return torch.empty( + (meta.input_ids.numel(), self.num_heads), + dtype=torch.int64, + device=meta.input_ids.device, + ) + return self.row_ids(meta) + def forward(self, meta: PLEMetadata, out: torch.Tensor | None = None) -> torch.Tensor: - return self.table.lookup(self.row_ids(meta), out) + return self.table.lookup(self.lookup_row_ids(meta), out) class _DepthwiseConv1d(BaseOP): @@ -680,7 +691,7 @@ def start_prefetch(self, batch: Batch, meta: PLEMetadata | None = None) -> None: """Hash this forward's n-grams and start the table gather on the side stream.""" if meta is None: meta = build_ple_metadata(batch, self.args, batch.input_ids.device) - row_ids = self.ple_embedding.row_ids(meta) + row_ids = self.ple_embedding.lookup_row_ids(meta) self._pending = (meta, row_ids) self.ple_embedding.table.prefetch(row_ids) @@ -701,7 +712,7 @@ def forward( elif pending is not None and pending[0] is meta: row_ids = pending[1] if row_ids is None: - row_ids = self.ple_embedding.row_ids(meta) + row_ids = self.ple_embedding.lookup_row_ids(meta) embeddings = self.ple_embedding.table.lookup(row_ids).to(R.dtype) key = self.norm_key.forward(self.key_proj.forward(embeddings)) diff --git a/python/freetoken/models/qwen4_exp/ple_disk.py b/python/freetoken/models/qwen4_exp/ple_disk.py index d2eaa99df8..56cfcb922b 100644 --- a/python/freetoken/models/qwen4_exp/ple_disk.py +++ b/python/freetoken/models/qwen4_exp/ple_disk.py @@ -152,6 +152,9 @@ def cancel(self) -> None: class DiskRowTable: """``PLETableBackend`` whose rows are read from disk per fill (--ple-backend disk).""" + # the fill hashes on the host; lookup/prefetch read only the row_ids shape + reads_row_ids = False + def __init__( self, source: PleRowSource, diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index 02abaee1f4..954d898945 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -339,6 +339,46 @@ def test_the_fused_hash_refuses_a_geometry_it_cannot_address(ctx_len, heads_per_ ) + +def test_a_host_hashing_table_skips_the_device_hash(monkeypatch): + """The disk table hashes on the host and reads only the row_ids shape, so neither the + prefetch nor the lookup path runs the device hash for it.""" + from freetoken.models.qwen4_exp.ple import ZeroTable + from freetoken.models.qwen4_exp.ple_disk import DiskRowTable + + assert DiskRowTable.reads_row_ids is False + config = _config() + args = config.qwen4_args + seen = [] + + class HostHashedTable(ZeroTable): + reads_row_ids = False + + def lookup(self, row_ids, out=None): + seen.append(("lookup", tuple(row_ids.shape), row_ids.dtype)) + return super().lookup(row_ids, out) + + def prefetch(self, row_ids): + seen.append(("prefetch", tuple(row_ids.shape), row_ids.dtype)) + + layer = _make_layer(config, table=HostHashedTable(_padded_vocab(args), args.ngram_head_dim)) + embedding = layer.ple_embedding + + def no_hash(*a, **k): + raise AssertionError("the device hash ran for a table that never reads it") + + monkeypatch.setattr(embedding, "row_ids", no_hash) + sequences, contexts = [[3, 4, EOS, 5], [2, 7]], [[EOS, EOS], [21, 22]] + meta = _meta(sequences, contexts) + total = sum(len(s) for s in sequences) + states = torch.zeros(len(sequences), args.ple_state_width, args.ple_conv_state_len) + + layer.start_prefetch(None, meta) + _forward(layer, torch.randn(total, layer.hc_count * layer.hidden_size), meta, states) + embedding.forward(meta) + rows = (total, embedding.num_heads) + assert seen == [("prefetch", rows, torch.int64), ("lookup", rows, torch.int64), ("lookup", rows, torch.int64)] + # -------------------------------------------------------------------------------------- # table backends # -------------------------------------------------------------------------------------- From 36c9eadfc842a2acb4c1d3bdf4433215d6723cb8 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 10:24:20 -0400 Subject: [PATCH 40/73] fix(ple): give each graph fill its own readback event and keep every pending fill A shared event re-recorded by the next step before the filler reached synchronize() made fill N wait on graph N, which waits on fill N; and a fill that failed after the next step submitted was dropped, so its graph ran on stale rows without an error. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 79988ccdf596e54b37817971628de84c802f3a58) --- python/freetoken/models/qwen4_exp/ple_disk.py | 34 ++++--- .../qwen4_exp/test_ple_disk_fill_order.py | 95 ++++++++++++++++++- 2 files changed, 117 insertions(+), 12 deletions(-) diff --git a/python/freetoken/models/qwen4_exp/ple_disk.py b/python/freetoken/models/qwen4_exp/ple_disk.py index 56cfcb922b..05b18d5ce7 100644 --- a/python/freetoken/models/qwen4_exp/ple_disk.py +++ b/python/freetoken/models/qwen4_exp/ple_disk.py @@ -8,6 +8,7 @@ import os import queue import threading +from collections import deque from concurrent.futures import Future from contextlib import contextmanager from dataclasses import dataclass @@ -146,6 +147,9 @@ def cancel(self) -> None: self._future.result() except BaseException: # noqa: BLE001 -- the launch error is the one to surface pass + pending = self._table._pending_fills() + if self._future in pending: + pending.remove(self._future) # its outcome was consumed here self._table._flag.zero_() @@ -216,11 +220,10 @@ def __init__( from freetoken.kernel.row_store import probe_wait_sync self._wait_sync = probe_wait_sync(os.getenv(_SYNC_ENV, "auto"), self._device) - # one flag for all graphs: the readback event orders a fill after the previous graph, so signals never overlap + # one flag for all graphs: each fill's readback event orders it after the previous graph, so signals never overlap self._flag = alloc_pinned_tensor(1, dtype=torch.int64) self._flag.zero_() self._token_readback = alloc_pinned_tensor(max_graph_rows, dtype=torch.int32) - self._readback_event = torch.cuda.Event() sync = "wait-sync" if self._wait_sync else "launch-gating" logger.info_rank0(f"PLE disk backend: {self._store.io_backend()}, {sync}") @@ -240,12 +243,17 @@ def close(self) -> None: def _current_stream(self): return torch.cuda.current_stream(self._device) + def _pending_fills(self) -> deque: + return self.__dict__.setdefault("_pending_fill_queue", deque()) + def _raise_failed_fill(self) -> None: - pending = self.__dict__.pop("_pending_fill", None) - if pending is not None and pending.done() and pending.exception() is not None: - raise pending.exception() - if pending is not None and not pending.done(): - self._pending_fill = pending + """Surface a finished graph fill's error on the engine thread, oldest first; a failed + fill signalled its WAIT, so that graph ran on stale rows.""" + pending = self._pending_fills() + while pending and pending[0].done(): + exc = pending.popleft().exception() + if exc is not None: + raise exc def _fill_now(self, runs: Sequence[torch.Tensor], *, graph: bool) -> None: """A fill that completes before its launch, on the store's thread (after any queued graph fill).""" @@ -288,12 +296,15 @@ def host_fill_batch(self, batch: Batch, use_graph: bool) -> _PendingFill | None: if use_graph and self._wait_sync: bs = batch.padded_size self._token_readback[:bs].copy_(batch.input_ids, non_blocking=True) - self._readback_event.record(self._current_stream()) + # one event per fill: a shared one re-recorded by the next step before this job + # reached synchronize() would wait on a graph parked on this job's own WAIT + readback = torch.cuda.Event() + readback.record(self._current_stream()) cancelled = threading.Event() def _job() -> None: try: - self._readback_event.synchronize() + readback.synchronize() tokens = self._token_readback[:bs].to(torch.int64).tolist() if cancelled.is_set(): return @@ -306,8 +317,9 @@ def _job() -> None: signal(self._flag) raise - self._pending_fill = self._filler().submit(_job) - return _PendingFill(self, self._pending_fill, cancelled) + future = self._filler().submit(_job) + self._pending_fills().append(future) + return _PendingFill(self, future, cancelled) # launch-gating: this D2H is the step's readback and orders the fill after sampling tokens = batch.input_ids.to("cpu").to(torch.int64).tolist() runs = [torch.tensor([*ctx, t], dtype=torch.int64) for ctx, t in zip(contexts, tokens)] diff --git a/tests/models/qwen4_exp/test_ple_disk_fill_order.py b/tests/models/qwen4_exp/test_ple_disk_fill_order.py index e7fe58f7e0..4d2bbfd71b 100644 --- a/tests/models/qwen4_exp/test_ple_disk_fill_order.py +++ b/tests/models/qwen4_exp/test_ple_disk_fill_order.py @@ -8,6 +8,7 @@ from __future__ import annotations +import queue import threading import time from types import SimpleNamespace @@ -55,13 +56,13 @@ def synchronize(self): def _table(monkeypatch) -> DiskRowTable: # host-only: no CUDA stream behind the (faked) readback event monkeypatch.setattr(ple_disk.torch.cuda, "current_stream", lambda *a, **k: None) + monkeypatch.setattr(ple_disk.torch.cuda, "Event", _Event) t = DiskRowTable.__new__(DiskRowTable) t._device = None t._wait_sync = True t._flag = torch.zeros(1, dtype=torch.int64) t._store = _Store(t._flag) t._token_readback = torch.zeros(8, dtype=torch.int32) - t._readback_event = _Event() t._graph_pinned = torch.zeros(64, dtype=torch.uint8) t._eager_pinned = [torch.zeros(64, dtype=torch.uint8), torch.zeros(64, dtype=torch.uint8)] t._eager_read = [_Event(), _Event()] @@ -134,3 +135,95 @@ def test_a_failed_launch_leaves_no_stale_signal(monkeypatch): # no WAIT consumed it, so a leftover 1 would let the next step read stale rows assert int(t._flag[0]) == 0 _close(t) + + +class _Gpu: + """One in-order stream: each launched graph parks on WAIT(flag), then RESETs it.""" + + def __init__(self, flag: torch.Tensor): + self.flag = flag + self.enqueued = self.completed = 0 + self.stalled = False + self.cv = threading.Condition() + self._launches: queue.SimpleQueue = queue.SimpleQueue() + threading.Thread(target=self._run, daemon=True).start() + + def launch(self) -> None: + self.enqueued += 1 + self._launches.put(True) + + def close(self) -> None: + self._launches.put(False) + + def _run(self) -> None: + while self._launches.get(): + self.stalled |= not _blocking_launch(self.flag) + with self.cv: + self.completed += 1 + self.cv.notify_all() + + def wait_completed(self, n: int, timeout: float = 5.0) -> bool: + with self.cv: + return self.cv.wait_for(lambda: self.completed >= n, timeout) + + +class _StreamEvent: + """cudaEvent semantics: synchronize() waits for the work enqueued before the LAST record.""" + + def __init__(self, gpu: _Gpu): + self.gpu, self.target = gpu, 0 + + def record(self, stream=None): + self.target = self.gpu.enqueued + + def synchronize(self): + if not self.gpu.wait_completed(self.target, timeout=1.0): + raise TimeoutError("the readback waits on a graph parked on this fill's own WAIT") + + +def _two_steps_behind_a_slow_filler(t: DiskRowTable, gpu: _Gpu) -> None: + """Steps N and N+1 both launch before the filler thread reaches fill N.""" + gate = threading.Event() + t._filler().submit(gate.wait) + for prev, new in (([1, 2, 3], 7), ([1, 2, 3, 7], 8)): + with t.forward_host_ctx(_decode_batch(prev, new), use_graph=True): + gpu.launch() + gate.set() + assert gpu.wait_completed(2), "decode deadlocked" + t._filler().submit(lambda: None).result() # every submitted fill has finished + gpu.close() + + +def test_each_fill_waits_on_its_own_steps_readback(monkeypatch): + """A readback event shared across steps is re-recorded by step N+1 behind graph N, so + fill N would wait on the graph that is parked on fill N's own WAIT.""" + t = _table(monkeypatch) + gpu = _Gpu(t._flag) + monkeypatch.setattr(ple_disk.torch.cuda, "Event", lambda: _StreamEvent(gpu)) + _two_steps_behind_a_slow_filler(t, gpu) + t._raise_failed_fill() + _close(t) + assert not gpu.stalled + # the fake D2H readback is not stream-ordered, so compare only the contexts + assert [c[1][:2] for c in t._store.calls if c[0] == "stage"] == [[2, 3], [3, 7]] + + +def test_a_failed_fill_is_raised_after_the_next_step_was_submitted(monkeypatch): + """Fill N fails (its graph runs on stale rows) after step N+1 already submitted: the error + must still reach the engine thread at the next launch.""" + t = _table(monkeypatch) + gpu = _Gpu(t._flag) + monkeypatch.setattr(ple_disk.torch.cuda, "Event", lambda: _StreamEvent(gpu)) + flush = t._store.flush + + def fail_step_n(signal_addr): + if t._store.calls[-1][1][:2] == [2, 3]: # the run staged just before is step N's + raise OSError("PLE row read failed") + flush(signal_addr) + + t._store.flush = fail_step_n + _two_steps_behind_a_slow_filler(t, gpu) + req = SimpleNamespace(input_ids=torch.tensor([1, 2, 3, 4]), cached_len=1, device_len=4) + with pytest.raises(OSError, match="PLE row read failed"): + t.host_fill_batch(SimpleNamespace(is_decode=False, padded_reqs=[req]), use_graph=False) + _close(t) From 86a06f16a938c71786fe09e2d644f63cfd496c94 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 15:01:03 -0400 Subject: [PATCH 41/73] fix(ple): recycle the fill readback events instead of freeing them on the filler The filler thread dropped each fill's last event reference, and cuEventDestroy there blocks on the driver lock that a blocked cuGraphLaunch holds (host-func nodes in the decode graph make the launch wait for the GPU, parked on the next fill's WAIT), so decode deadlocked with FREETOKEN_CPU_MOE_FLAG_SYNC=0. Events now return to an engine-thread pool once their fill is done. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 198574d266608fd6744f0ad90f01461ecc6b1c2d) --- python/freetoken/models/qwen4_exp/ple_disk.py | 21 ++++++++++++++-- .../qwen4_exp/test_ple_disk_fill_order.py | 25 +++++++++++++++++++ 2 files changed, 44 insertions(+), 2 deletions(-) diff --git a/python/freetoken/models/qwen4_exp/ple_disk.py b/python/freetoken/models/qwen4_exp/ple_disk.py index 05b18d5ce7..a0c71f7046 100644 --- a/python/freetoken/models/qwen4_exp/ple_disk.py +++ b/python/freetoken/models/qwen4_exp/ple_disk.py @@ -150,6 +150,7 @@ def cancel(self) -> None: pending = self._table._pending_fills() if self._future in pending: pending.remove(self._future) # its outcome was consumed here + self._table._retire(self._future) self._table._flag.zero_() @@ -246,12 +247,24 @@ def _current_stream(self): def _pending_fills(self) -> deque: return self.__dict__.setdefault("_pending_fill_queue", deque()) + def _readback_pool(self) -> list: + return self.__dict__.setdefault("_readback_events", []) + + def _retire(self, future: Future) -> None: + """Return a finished fill's readback event for reuse (engine thread only).""" + event = getattr(future, "readback", None) + if event is not None: + self._readback_pool().append(event) + future.readback = None + def _raise_failed_fill(self) -> None: """Surface a finished graph fill's error on the engine thread, oldest first; a failed fill signalled its WAIT, so that graph ran on stale rows.""" pending = self._pending_fills() while pending and pending[0].done(): - exc = pending.popleft().exception() + future = pending.popleft() + self._retire(future) + exc = future.exception() if exc is not None: raise exc @@ -298,7 +311,10 @@ def host_fill_batch(self, batch: Batch, use_graph: bool) -> _PendingFill | None: self._token_readback[:bs].copy_(batch.input_ids, non_blocking=True) # one event per fill: a shared one re-recorded by the next step before this job # reached synchronize() would wait on a graph parked on this job's own WAIT - readback = torch.cuda.Event() + # recycled, never destroyed while serving: cuEventDestroy from the filler blocks on + # the driver lock a blocked cuGraphLaunch holds, and the fill it waits for never runs + pool = self._readback_pool() + readback = pool.pop() if pool else torch.cuda.Event() readback.record(self._current_stream()) cancelled = threading.Event() @@ -318,6 +334,7 @@ def _job() -> None: raise future = self._filler().submit(_job) + future.readback = readback self._pending_fills().append(future) return _PendingFill(self, future, cancelled) # launch-gating: this D2H is the step's readback and orders the fill after sampling diff --git a/tests/models/qwen4_exp/test_ple_disk_fill_order.py b/tests/models/qwen4_exp/test_ple_disk_fill_order.py index 4d2bbfd71b..9d979b64eb 100644 --- a/tests/models/qwen4_exp/test_ple_disk_fill_order.py +++ b/tests/models/qwen4_exp/test_ple_disk_fill_order.py @@ -227,3 +227,28 @@ def fail_step_n(signal_addr): with pytest.raises(OSError, match="PLE row read failed"): t.host_fill_batch(SimpleNamespace(is_decode=False, padded_reqs=[req]), use_graph=False) _close(t) + + +def test_the_readback_event_is_never_freed_on_the_filler_thread(monkeypatch): + """cuEventDestroy on the filler blocks on the driver lock a blocked cuGraphLaunch holds (host + nodes in the graph make the launch wait for the GPU, which is parked on the next fill's WAIT).""" + import gc + + freed_on: list[str] = [] + + class _TrackedEvent(_Event): + def __del__(self): + freed_on.append(threading.current_thread().name) + + t = _table(monkeypatch) + monkeypatch.setattr(ple_disk.torch.cuda, "Event", _TrackedEvent) + for prev, new in (([1, 2, 3], 7), ([1, 2, 3, 7], 8)): + with t.forward_host_ctx(_decode_batch(prev, new), use_graph=True): + assert _blocking_launch(t._flag) + time.sleep(0.05) # let the filler drop its work items + gc.collect() + t._raise_failed_fill() + gc.collect() + assert "ple-filler" not in freed_on, freed_on + assert len(t._readback_pool()) >= 1 # finished fills hand their events back for reuse + _close(t) From b610dbe31bef232cbc29828aa8af2c9e7fd08670 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 05:25:06 -0400 Subject: [PATCH 42/73] fix(kvcache): refuse the host kv-page tier on a pool whose index slab it cannot restore Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit eb462bd73c2fa5eb96b38abe56fb841d82d68a32) --- python/freetoken/kvcache/host_prefix_tier.py | 8 +++++++- tests/kvcache/test_host_tier_nvfp4_guard.py | 7 +++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/python/freetoken/kvcache/host_prefix_tier.py b/python/freetoken/kvcache/host_prefix_tier.py index 39c96dfd6d..5cee24058a 100644 --- a/python/freetoken/kvcache/host_prefix_tier.py +++ b/python/freetoken/kvcache/host_prefix_tier.py @@ -84,8 +84,14 @@ def __init__(self, state_pool, kv_pool=None, *, gdn_slots_host: int = 64, from freetoken.kvcache.kv_host_offload import _reject_block_scaled_pool _reject_block_scaled_pool(kv_pool, "FT_PREFIX_HOST") + if getattr(kv_pool, "_cmp_k_buffer", None) is not None: + # a QSA pool's indexer slab is keyed by the same pages and is not mirrored here + raise ValueError( + "FT_PREFIX_HOST cannot restore a sparse-attention pool's compressed index " + "slab; keep FT_GDN_HOST_TIER alone for this model" + ) codes = kv_pool._kv_buffer # [2, L, P, page, H, D] - scales = kv_pool._scale_buffer # [2, L, P*page, H] | None (fp8) + scales = getattr(kv_pool, "_scale_buffer", None) # [2, L, P*page, H] | None (fp8) _, layers, _, page_size, heads, dim = codes.shape page_bytes = 2 * layers * page_size * heads * dim * codes.element_size() if scales is not None: diff --git a/tests/kvcache/test_host_tier_nvfp4_guard.py b/tests/kvcache/test_host_tier_nvfp4_guard.py index a793fb294a..50227267a6 100644 --- a/tests/kvcache/test_host_tier_nvfp4_guard.py +++ b/tests/kvcache/test_host_tier_nvfp4_guard.py @@ -42,3 +42,10 @@ def test_prefix_tier_gdn_snapshots_do_not_touch_the_kv_pool(): tier = HostPrefixTier(_state_pool(), kv_pool=_kv_pool(block_scaled=True), gdn_slots_host=2, kv_budget_bytes=0, device="cpu") assert tier.num_slots == 2 and not tier.has_kv + + +def test_prefix_tier_refuses_kv_pages_on_a_qsa_pool(): + # the KV-page store mirrors codes and scales, not the QSA indexer's compressed-key slab + pool = SimpleNamespace(**vars(_kv_pool(block_scaled=False)), _cmp_k_buffer=torch.zeros(1)) + with pytest.raises(ValueError, match="FT_PREFIX_HOST .* index slab"): + HostPrefixTier(_state_pool(), kv_pool=pool, gdn_slots_host=2, kv_budget_bytes=1 << 16, device="cpu") From c312d3807f7a032374b226b50e3338b77b5e8210 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 05:51:57 -0400 Subject: [PATCH 43/73] perf(kvcache): keep one gpu gdn snapshot slot when the host snapshot tier is on Evicted snapshots spill to RAM and return on a hit, so the floor of four resident ones only costs VRAM: three 110 MiB Qwen3.8-Flash-Next slots go back to the expert cache. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 6751cf86b302dd539fdf81c1b32aaa8820650a5b) --- python/freetoken/kvcache/linear_state_pool.py | 6 +++++- tests/scheduler/test_hybrid_cache_manager.py | 11 +++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/python/freetoken/kvcache/linear_state_pool.py b/python/freetoken/kvcache/linear_state_pool.py index ff7539410c..d89556899d 100644 --- a/python/freetoken/kvcache/linear_state_pool.py +++ b/python/freetoken/kvcache/linear_state_pool.py @@ -1,6 +1,7 @@ from __future__ import annotations import math +import os import torch from freetoken.distributed import get_tp_info @@ -285,7 +286,10 @@ def _linear_pool_num_slots(config) -> int: if config.cache_type != "hybrid_radix": return mr + 1 # live + dummy/padding ratio = config.linear_state_cache_ratio - n_cache = max(4, int(ratio * mr)) + # With the host snapshot tier (FT_GDN_HOST_TIER=1) evicted snapshots spill to RAM and come + # back on a hit, so the GPU cache need not keep four of them resident. + floor = 1 if os.environ.get("FT_GDN_HOST_TIER", "0") == "1" else 4 + n_cache = max(floor, int(ratio * mr)) return 4 * mr + n_cache + 1 # live + 2 ping-pong + locked committed snapshot + cache + padding diff --git a/tests/scheduler/test_hybrid_cache_manager.py b/tests/scheduler/test_hybrid_cache_manager.py index 07385d7046..165689ab3b 100644 --- a/tests/scheduler/test_hybrid_cache_manager.py +++ b/tests/scheduler/test_hybrid_cache_manager.py @@ -241,5 +241,16 @@ def test_pool_sizing_covers_4mr_floor(): assert _linear_pool_num_slots(c) >= 4 * mr + 1, (mr, _linear_pool_num_slots(c)) +def test_host_snapshot_tier_lowers_the_gpu_snapshot_floor(monkeypatch): + """With FT_GDN_HOST_TIER evicted snapshots live in RAM, so one GPU snapshot slot is enough.""" + from types import SimpleNamespace + from freetoken.kvcache.linear_state_pool import _linear_pool_num_slots + c = SimpleNamespace(max_running_req=1, cache_type="hybrid_radix", linear_state_cache_ratio=0.0) + monkeypatch.delenv("FT_GDN_HOST_TIER", raising=False) + assert _linear_pool_num_slots(c) == 4 + 4 + 1 + monkeypatch.setenv("FT_GDN_HOST_TIER", "1") + assert _linear_pool_num_slots(c) == 4 + 1 + 1 + + if __name__ == "__main__": raise SystemExit(pytest.main([__file__])) From 4ea7d4a369a0c8fc53ef3bbcb3aeeb6426134f93 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 10:43:02 -0400 Subject: [PATCH 44/73] perf(scheduler): probe the host tier for the exact tombstone prefix Revival only accepts an exact entry, but the longest-prefix lookup re-hashed every shorter page-aligned prefix on a miss and refreshed the LRU age of a shorter entry it then rejected. _rehydrate_snapshot on a 262144-token key, CPU, median of 3: no host entry 1852 ms -> 4.0 ms entry 4096 tokens in 1494 ms -> 2.9 ms entry 8192 below 91 ms -> 3.5 ms exact hit 1.4 ms -> 1.6 ms Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit b349af0cac8ee9334b49023460bc5d7afa9e92d1) --- python/freetoken/kvcache/host_prefix_tier.py | 10 +++++ python/freetoken/scheduler/cache.py | 8 ++-- .../kvcache/test_host_tier_rehydrate_exact.py | 39 ++++++++++++------- 3 files changed, 38 insertions(+), 19 deletions(-) diff --git a/python/freetoken/kvcache/host_prefix_tier.py b/python/freetoken/kvcache/host_prefix_tier.py index 5cee24058a..4295aa8e4b 100644 --- a/python/freetoken/kvcache/host_prefix_tier.py +++ b/python/freetoken/kvcache/host_prefix_tier.py @@ -318,6 +318,16 @@ def lookup(self, input_ids: torch.Tensor, page_size: int = 64) -> HostMatch | No n -= page_size return None + def lookup_exact(self, key_tokens: torch.Tensor) -> HostMatch | None: + """The entry for exactly this prefix: one hash, where lookup re-hashes every shorter + page-aligned prefix on a miss.""" + key = self._key_of(key_tokens) + hit = self._index.get(key) + if hit is None: + return None + self._index.move_to_end(key) + return HostMatch(key_tokens.numel(), hit.kv_slots, hit.gdn_slot) + def evict_host(self, key_tokens: torch.Tensor) -> None: """Remove uma entrada do índice e libera os slots host (KV + GDN).""" ent = self._index.pop(self._key_of(key_tokens), None) diff --git a/python/freetoken/scheduler/cache.py b/python/freetoken/scheduler/cache.py index 4e6b91df88..715dfdc21a 100644 --- a/python/freetoken/scheduler/cache.py +++ b/python/freetoken/scheduler/cache.py @@ -625,10 +625,10 @@ def _rehydrate_snapshot(self, node) -> int | None: keep walking (the match truncates as before).""" tier = self.host_tier key = self.prefix_cache._collect_key(node) - hit = tier.lookup(key, self.page_size) - # lookup is longest-prefix; a shorter entry's state would resume this node's full - # depth with the tokens in between missing from the recurrence - if hit is None or hit.gdn_slot is None or hit.host_len != len(key): + # exact only: a shorter entry's state would resume this node's full depth with the + # tokens in between missing from the recurrence + hit = tier.lookup_exact(key) + if hit is None or hit.gdn_slot is None: return None pool = self.linear_state_pool if pool.num_free_slots < 1: diff --git a/tests/kvcache/test_host_tier_rehydrate_exact.py b/tests/kvcache/test_host_tier_rehydrate_exact.py index f67402041f..d6f4d09215 100644 --- a/tests/kvcache/test_host_tier_rehydrate_exact.py +++ b/tests/kvcache/test_host_tier_rehydrate_exact.py @@ -1,21 +1,19 @@ """The GDN snapshot tier (upstream #525) may only revive a tombstone from a snapshot of -exactly that node's prefix: the host index lookup is longest-prefix, and a shorter -snapshot would resume the recurrence with the tokens in between silently missing.""" +exactly that node's prefix: a shorter snapshot would resume the recurrence with the tokens +in between silently missing.""" from types import SimpleNamespace import torch -from freetoken.kvcache.host_prefix_tier import HostMatch +from freetoken.kvcache.host_prefix_tier import HostPrefixTier from freetoken.scheduler.cache import CacheManager -def _manager(host_len: int, key_len: int = 128): - reads = [] - tier = SimpleNamespace( - lookup=lambda key, page_size: HostMatch(host_len, None, 7), - read_gdn=lambda host_slot, gpu_slot: reads.append((host_slot, gpu_slot)), - ) +def _manager(key_len: int = 128): + state = SimpleNamespace(conv_states=torch.zeros(1, 4, 1, 1), + recurrent_states=torch.zeros(1, 4, 1, 1), slot_states={}) + tier = HostPrefixTier(state, gdn_slots_host=4, device="cpu") pool = SimpleNamespace(num_free_slots=4, alloc=lambda n: [3]) manager = SimpleNamespace( host_tier=tier, @@ -24,16 +22,27 @@ def _manager(host_len: int, key_len: int = 128): linear_state_pool=pool, ensure_mamba_slots=lambda n: None, ) - return manager, reads + return manager, tier, state + + +def _snapshot(tier, state, key, value): + state.conv_states[:, 1] = value + tier.put_prefix(key, None, tier.write_gdn(1)) -def test_a_shorter_host_snapshot_does_not_revive_a_deeper_node(): - manager, reads = _manager(host_len=64) +def test_a_shorter_host_snapshot_neither_revives_nor_refreshes_a_deeper_node(): + manager, tier, state = _manager(key_len=128) + _snapshot(tier, state, torch.arange(64), 1.0) + _snapshot(tier, state, torch.arange(1000, 1064), 2.0) assert CacheManager._rehydrate_snapshot(manager, node=object()) is None - assert reads == [] + # the rejected entry keeps its LRU age + assert [len(e.key_tokens) for e in tier._index.values()] == [64, 64] + assert next(iter(tier._index.values())).key_tokens[0] == 0 + assert torch.all(state.conv_states[:, 3] == 0) def test_an_exact_host_snapshot_revives_the_node(): - manager, reads = _manager(host_len=128) + manager, tier, state = _manager(key_len=128) + _snapshot(tier, state, torch.arange(128), 5.0) assert CacheManager._rehydrate_snapshot(manager, node=object()) == 3 - assert reads == [(7, 3)] + assert torch.all(state.conv_states[:, 3] == 5.0) From 844bc4972ac7e5174ff689c7e91e788a1b195730 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 10:46:01 -0400 Subject: [PATCH 45/73] fix(scheduler): re-probe the host entry after making room to rehydrate it The evictions that make room for a revived snapshot or a host prefix write their victims to the host tier, which can drop the hit's own entry and hand its GDN and KV slots to the victim, so the stale hit restored another prefix's state. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 44f25067541a6fb84161effe4dd2d7ee8e3a0306) --- python/freetoken/scheduler/cache.py | 12 ++- .../kvcache/test_host_tier_rehydrate_exact.py | 79 ++++++++++++++++++- 2 files changed, 87 insertions(+), 4 deletions(-) diff --git a/python/freetoken/scheduler/cache.py b/python/freetoken/scheduler/cache.py index 715dfdc21a..ff4467328e 100644 --- a/python/freetoken/scheduler/cache.py +++ b/python/freetoken/scheduler/cache.py @@ -632,9 +632,11 @@ def _rehydrate_snapshot(self, node) -> int | None: return None pool = self.linear_state_pool if pool.num_free_slots < 1: - # make room: evicting another snapshot writes IT to the host tier first + # make room: evicting another snapshot writes IT to the host tier first, which may + # drop this entry and hand its host slot to the victim, so probe again self.ensure_mamba_slots(1) - if pool.num_free_slots < 1: + hit = tier.lookup_exact(key) + if pool.num_free_slots < 1 or hit is None or hit.gdn_slot is None: return None slot = pool.alloc(1)[0] tier.read_gdn(hit.gdn_slot, slot) @@ -651,7 +653,11 @@ def _rehydrate_host_prefix(self, ids: torch.Tensor, hit) -> None: pool = self.linear_state_pool if hit.gdn_slot is not None: self.ensure_mamba_slots(1) - if pool.num_free_slots < 1: + # _allocate and ensure_mamba_slots write their victims to the tier, which may drop + # this entry and hand its host slots to them: use what the key maps to now + hit = tier.lookup_exact(ids[: hit.host_len]) + if (pool.num_free_slots < 1 or hit is None or hit.kv_pages is None + or hit.gdn_slot is None): self._free(kv_new) # sem slot GDN não dá pra reidratar; devolve as páginas return slot = pool.alloc(1)[0] diff --git a/tests/kvcache/test_host_tier_rehydrate_exact.py b/tests/kvcache/test_host_tier_rehydrate_exact.py index d6f4d09215..77dae0fd08 100644 --- a/tests/kvcache/test_host_tier_rehydrate_exact.py +++ b/tests/kvcache/test_host_tier_rehydrate_exact.py @@ -1,12 +1,16 @@ """The GDN snapshot tier (upstream #525) may only revive a tombstone from a snapshot of exactly that node's prefix: a shorter snapshot would resume the recurrence with the tokens -in between silently missing.""" +in between silently missing. Making room for the revival evicts other prefixes into the +tier, which may recycle the entry's own host slots, so it must not read stale ones.""" from types import SimpleNamespace import torch +import freetoken.core as core from freetoken.kvcache.host_prefix_tier import HostPrefixTier +from freetoken.kvcache.linear_state_pool import LinearStatePool +from freetoken.models.config import LinearGatedDeltaGroupConfig from freetoken.scheduler.cache import CacheManager @@ -46,3 +50,76 @@ def test_an_exact_host_snapshot_revives_the_node(): _snapshot(tier, state, torch.arange(128), 5.0) assert CacheManager._rehydrate_snapshot(manager, node=object()) == 3 assert torch.all(state.conv_states[:, 3] == 5.0) + + +def _cache_manager(monkeypatch, num_pages, page_size, state_slots, **env): + for k, v in env.items(): + monkeypatch.setenv(k, v) + g = LinearGatedDeltaGroupConfig( + name="linear", layer_ids=(0,), num_key_heads=2, num_value_heads=4, + key_head_dim=16, value_head_dim=16, conv_kernel_dim=4, output_gate="silu", + ) + pool = LinearStatePool(group=g, num_slots=state_slots, dtype=torch.bfloat16, + device=torch.device("cpu"), tp_size=1) + page_table = torch.zeros(2, 64, dtype=torch.int32) + cm = CacheManager(num_pages, page_size, page_table, "hybrid_radix", linear_state_pool=pool) + return cm, pool + + +def _put(cm, pool, ids, state, stamp, kv_pool=None, kv_value=None): + kv = cm._page_to_token(cm._allocate(len(ids) // cm.page_size)) + if kv_pool is not None: + kv_pool._kv_buffer[:, :, kv[:: cm.page_size] // cm.page_size] = kv_value + slot = pool.alloc(1)[0] + pool.conv_states[:, slot] = state + cm.prefix_cache.insert(ids, kv, slot) + cm.prefix_cache.match_prefix(ids).node.timestamp = stamp + + +def _match(cm, ids): + tail = torch.tensor([7, 9], dtype=torch.int32) + return cm.match_req(SimpleNamespace(input_ids=torch.cat([ids, tail]), input_len=len(ids) + 2)) + + +def test_making_room_to_revive_does_not_restore_the_victim_into_the_host_slot(monkeypatch): + # one host slot: evicting B for the revived slot drops A's entry and writes B there + cm, pool = _cache_manager(monkeypatch, 8, 64, 3, FT_GDN_HOST_TIER="1", FT_GDN_HOST_SLOTS="1") + a = torch.arange(128, dtype=torch.int32) + b = torch.arange(1000, 1128, dtype=torch.int32) + _put(cm, pool, a, 1.0, stamp=1) + _put(cm, pool, b, 2.0, stamp=2) + cm.ensure_mamba_slots(1) # A -> host tier, A's node becomes a tombstone + pool.conv_states[:, pool.alloc(1)[0]] = 9.0 # someone else takes the freed slot + + mr = _match(cm, a) + restored = None if mr.mamba_value is None else pool.conv_states[:, mr.mamba_value].unique().tolist() + assert restored in (None, [1.0]), (mr.cuda_handle.cached_len, restored) + assert cm.host_tier.lookup_exact(b) is not None + + +def test_making_room_to_rehydrate_a_host_prefix_does_not_restore_the_victim(monkeypatch): + # a 4-page KV store: writing the 4-page victim V drops the 2-page hit H and reuses its slots + page, n_pages = 4, 6 + kv_pool = SimpleNamespace(_kv_buffer=torch.zeros(2, 1, n_pages, page, 1, 2), + _scale_buffer=None, _block_scale_buffer=None) + monkeypatch.setattr(core, "_GLOBAL_CTX", SimpleNamespace(kv_cache=kv_pool, kv_offloader=None)) + page_bytes = kv_pool._kv_buffer[:, :, 0].numel() * kv_pool._kv_buffer.element_size() + cm, pool = _cache_manager(monkeypatch, n_pages, page, 4, FT_PREFIX_HOST="1", + FT_PREFIX_HOST_GB=repr(4 * page_bytes / (1 << 30)), + FT_GDN_HOST_SLOTS="4") + h = torch.arange(2 * page, dtype=torch.int32) + v = torch.arange(100, 100 + 4 * page, dtype=torch.int32) + _put(cm, pool, h, 1.0, stamp=1, kv_pool=kv_pool, kv_value=1.0) + er = cm.prefix_cache.evict_full(len(h)) # H -> host tier (KV + snapshot), out of the tree + pool.free(er.mamba_slots) + cm._free(er.kv_indices) + _put(cm, pool, v, 2.0, stamp=2, kv_pool=kv_pool, kv_value=2.0) + cm._allocate(len(cm.free_slots)) # a running request holds the rest of the GPU pages + + mr = _match(cm, h) + handle = mr.cuda_handle + if handle.cached_len: + pages = handle.kv_indices[::page] // page + assert pool.conv_states[:, mr.mamba_value].unique().tolist() == [1.0] + assert kv_pool._kv_buffer[:, :, pages].unique().tolist() == [1.0] + assert cm.host_tier.lookup_exact(v) is not None From 8905d0af030beb5ccdb2307f7b8896d8c03bbbc3 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 10:49:19 -0400 Subject: [PATCH 46/73] test(kvcache): cover the engine-stream waits of the prefix tier's kv page copies Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit af5ca25838e9c690ba32a6cae20c72692e61477c) --- tests/kvcache/test_host_tier_live_pool.py | 23 +++++++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/tests/kvcache/test_host_tier_live_pool.py b/tests/kvcache/test_host_tier_live_pool.py index bd12836893..0eea1990e3 100644 --- a/tests/kvcache/test_host_tier_live_pool.py +++ b/tests/kvcache/test_host_tier_live_pool.py @@ -32,7 +32,7 @@ def test_read_after_a_pool_rebuild_lands_in_the_new_tensors(): assert torch.all(pool.recurrent_states[:, 5] == 7.0) -def test_copies_wait_for_the_engine_stream(monkeypatch): +def _fake_copy_stream(monkeypatch): events = [] engine = object() @@ -45,9 +45,28 @@ def synchronize(self): monkeypatch.setattr(tier_module.torch.cuda, "current_stream", lambda device=None: engine) monkeypatch.setattr(tier_module.torch.cuda, "stream", lambda s: contextlib.nullcontext()) + return FakeStream(), events, engine + + +def test_copies_wait_for_the_engine_stream(monkeypatch): + stream, events, engine = _fake_copy_stream(monkeypatch) pool = _pool(3) tier = HostPrefixTier(pool, gdn_slots_host=2, device="cpu") - tier.copy_stream = FakeStream() + tier.copy_stream = stream host = tier.write_gdn(1) tier.read_gdn(host, 2) assert events == [("wait", engine), ("sync",), ("wait", engine), ("sync",)] + + +def test_kv_page_copies_wait_for_the_engine_stream(monkeypatch): + stream, events, engine = _fake_copy_stream(monkeypatch) + kv_pool = SimpleNamespace( + _kv_buffer=torch.zeros(2, 1, 4, 4, 1, 2), # [2, L, P, page, H, D] + _scale_buffer=torch.zeros(2, 1, 16, 1), # [2, L, P*page, H] + _block_scale_buffer=None, + ) + tier = HostPrefixTier(_pool(3), kv_pool, gdn_slots_host=2, kv_budget_bytes=1 << 12, device="cpu") + tier.copy_stream = stream + slots = tier.write_kv(torch.arange(4, 8, dtype=torch.int32)) + tier.read_kv(slots, torch.tensor([2])) + assert events == [("wait", engine), ("sync",), ("wait", engine), ("sync",)] From 2cc45a49c917419eb4c59891d8393e0dd81735d4 Mon Sep 17 00:00:00 2001 From: fangyuan-3149 <314945674@qq.com> Date: Sun, 30 Aug 2026 09:35:36 +0800 Subject: [PATCH 47/73] fix(engine): reserve the GDN prefill workspace in the cache budget chunk_gated_delta_rule_fwd_h materializes the per-chunk recurrent states for the whole extend at once -- [NT, H, V, K] with NT = ceil(extend/64) -- plus a transient v_new buffer. Subtract it at the startup sizing and the --moe-cache-auto plan for hybrid-GDN models. (cherry picked from commit 5c3b9141b6b658f4d4ebe59f2a68b2208c6f6771) Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit ec8731a65ff5ca2cce0ff073d273f482a8faa5e6) --- python/freetoken/engine/engine.py | 7 ++++- python/freetoken/kvcache/linear_state_pool.py | 27 +++++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 57519c4eca..eab1d64828 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -35,7 +35,8 @@ from freetoken.kvcache.base import CacheRebuildRejected from freetoken.kvcache.cache_status import _supports_swa_ratio from freetoken.kvcache.linear_state_pool import ( - _linear_pool_min_slots, _linear_pool_num_slots, state_pool_bytes, + _linear_pool_min_slots, _linear_pool_num_slots, gdn_prefill_workspace_bytes, + state_pool_bytes, ) logger = init_logger(__name__) @@ -506,6 +507,9 @@ def __init__(self, config: EngineConfig): # off it; the KV pool family owns every geometry-specific formula behind the rest. available_memory = _startup_kv_budget(config.memory_ratio, init_free_memory, new_free) available_memory -= state_pool_bytes(config) + # reserve the GDN prefill transient (chunk recurrent states + v_new); without + # it long prefills OOM inside the GDN kernel when memory_ratio packs the pools + available_memory -= gdn_prefill_workspace_bytes(config) num_gpu_pages = self._pool_cls.solve_num_pages(config, available_memory) if config.kv_host_pages > 0: # KV host offload: the scheduler, page table and radix tree run in a LOGICAL page @@ -761,6 +765,7 @@ def _resolve_auto_moe_cache_size(self, config: EngineConfig, banks, method=None) cache_per_page, fixed_cache_size, page_tokens, min_reserve = self._pool_cls.kv_cost(config) fixed_cache_size += state_pool_bytes(config) # sibling GDN state pool, engine-summed + fixed_cache_size += gdn_prefill_workspace_bytes(config) # GDN prefill transient num_experts = config.model_config.num_experts total_experts = config.model_config.num_moe_layers * num_experts return resolve_moe_cache_auto( diff --git a/python/freetoken/kvcache/linear_state_pool.py b/python/freetoken/kvcache/linear_state_pool.py index d89556899d..16d2fb537e 100644 --- a/python/freetoken/kvcache/linear_state_pool.py +++ b/python/freetoken/kvcache/linear_state_pool.py @@ -261,6 +261,33 @@ def linear_state_bytes_per_req( __all__ = ["LinearStatePool", "linear_state_bytes_per_req"] +_GDN_CHUNK_SIZE = 64 # keep in lockstep with kernel/fla/chunk_delta_h.py CHUNK_SIZE + + +def gdn_prefill_workspace_bytes(config) -> int: + """Peak transient VRAM of one GDN prefill forward. + + ``chunk_gated_delta_rule_fwd_h`` materializes the per-chunk recurrent states for + the WHOLE extend at once -- ``[NT, H, V, K]`` with ``NT = ceil(extend / 64)`` -- + plus the ``v_new`` value buffer; ~192 MiB at an 8192-token extent for + ornith-class geometry. The generic ``(1 - memory_ratio)`` headroom is sized for + CUDA graphs; without reserving this term the GDN kernel OOMs on long prefills + at aggressive memory ratios. 0 for non-GDN models.""" + linear_group = config.model_config.linear_attention_group() + if linear_group is None: + return 0 + extent = min(config.max_extend_tokens, config.max_seq_len) + nt = -(-extent // _GDN_CHUNK_SIZE) + states = ( + nt + * linear_group.num_value_heads + * linear_group.value_head_dim + * linear_group.key_head_dim + ) + v_new = extent * linear_group.num_value_heads * linear_group.value_head_dim + return (states + v_new) * config.dtype.itemsize + + def state_pool_bytes(config, num_slots: int | None = None) -> int: """Total GDN state-pool bytes at ``num_slots`` PHYSICAL slots (default: the startup slot count). The engine adds this to the KV family's fixed cost when budgeting -- From 1c7dc2c6ed6783c87449154e76b38025701e2446 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 05:29:45 -0400 Subject: [PATCH 48/73] fix(engine): reserve the gdn prefill workspace in the rebuild fit-check too The cherry-picked reservation covered startup sizing and the auto plan but not validate_rebuild, so a runtime rebuild could grow the pools back into it. Also price the scheduler's default extend when the config is a bare EngineConfig, and document that h + v_new is a lower bound on the GDN prefill peak. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 9d836ed0b0b6c448975b26b4d86368e087dfe0af) --- python/freetoken/engine/engine.py | 7 ++-- python/freetoken/kvcache/linear_state_pool.py | 13 +++--- tests/engine/test_cache_budget.py | 38 +++++++++++++++++ tests/kvcache/test_linear_state_pool_alloc.py | 42 +++++++++++++++++++ 4 files changed, 91 insertions(+), 9 deletions(-) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index eab1d64828..9ccc44763c 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -507,8 +507,8 @@ def __init__(self, config: EngineConfig): # off it; the KV pool family owns every geometry-specific formula behind the rest. available_memory = _startup_kv_budget(config.memory_ratio, init_free_memory, new_free) available_memory -= state_pool_bytes(config) - # reserve the GDN prefill transient (chunk recurrent states + v_new); without - # it long prefills OOM inside the GDN kernel when memory_ratio packs the pools + # Reserve part of the GDN prefill transient (chunk states + v_new); the rest of the + # prefill activations still live in the (1 - memory_ratio) headroom. available_memory -= gdn_prefill_workspace_bytes(config) num_gpu_pages = self._pool_cls.solve_num_pages(config, available_memory) if config.kv_host_pages > 0: @@ -1181,7 +1181,8 @@ def rebuild_runtime_cache( per_expert_bytes=per_expert_bytes, baseline_free=self._baseline_free, weights_bytes=self._weights_bytes, current_num_pages=self.num_pages, extra_fixed_bytes=( - state_pool_bytes(config, target_mamba) if target_mamba is not None else 0 + (state_pool_bytes(config, target_mamba) if target_mamba is not None else 0) + + gdn_prefill_workspace_bytes(config) ), extra_note=( f", mamba={target_mamba - 1} slots" if target_mamba is not None else "" diff --git a/python/freetoken/kvcache/linear_state_pool.py b/python/freetoken/kvcache/linear_state_pool.py index 16d2fb537e..df5a3396da 100644 --- a/python/freetoken/kvcache/linear_state_pool.py +++ b/python/freetoken/kvcache/linear_state_pool.py @@ -265,18 +265,19 @@ def linear_state_bytes_per_req( def gdn_prefill_workspace_bytes(config) -> int: - """Peak transient VRAM of one GDN prefill forward. + """Bytes of the GDN chunked-prefill buffers that scale with the extend length. ``chunk_gated_delta_rule_fwd_h`` materializes the per-chunk recurrent states for the WHOLE extend at once -- ``[NT, H, V, K]`` with ``NT = ceil(extend / 64)`` -- - plus the ``v_new`` value buffer; ~192 MiB at an 8192-token extent for - ornith-class geometry. The generic ``(1 - memory_ratio)`` headroom is sized for - CUDA graphs; without reserving this term the GDN kernel OOMs on long prefills - at aggressive memory ratios. 0 for non-GDN models.""" + plus the ``v_new`` value buffer (288 MiB at an 8192-token extend for qwen4_exp). + This is a lower bound on the GDN prefill peak, not the peak: w/u/A/o and the + in_proj/conv activations still come out of the ``(1 - memory_ratio)`` headroom. + 0 for non-GDN models.""" linear_group = config.model_config.linear_attention_group() if linear_group is None: return 0 - extent = min(config.max_extend_tokens, config.max_seq_len) + # EngineConfig has no max_extend_tokens (the served config is a SchedulerConfig subclass) + extent = min(getattr(config, "max_extend_tokens", 8192), config.max_seq_len) nt = -(-extent // _GDN_CHUNK_SIZE) states = ( nt diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py index 7d5b7ec5cd..da242a2cc4 100644 --- a/tests/engine/test_cache_budget.py +++ b/tests/engine/test_cache_budget.py @@ -398,6 +398,44 @@ def slot_limit(self): assert captured["max_slots"] == 5 + +def test_rebuild_fit_check_reserves_the_gdn_prefill_workspace(): + """The runtime-rebuild fit-check prices the GDN prefill workspace like the startup sizing, + so a rebuild cannot grow the pools back into the space reserved for it.""" + from freetoken.engine.engine import Engine + from freetoken.kvcache.base import CacheRebuildRejected + from freetoken.kvcache.linear_state_pool import gdn_prefill_workspace_bytes, state_pool_bytes + from freetoken.models.config import LinearGatedDeltaGroupConfig + + group = LinearGatedDeltaGroupConfig( + name="linear", layer_ids=(0, 1), num_key_heads=2, num_value_heads=4, + key_head_dim=16, value_head_dim=16, conv_kernel_dim=4, output_gate="silu", + ) + config = SimpleNamespace( + model_config=SimpleNamespace(linear_attention_group=lambda: group), + dtype=torch.bfloat16, tp_info=SimpleNamespace(size=1), + max_extend_tokens=1024, max_seq_len=4096, + ) + captured = {} + + def validate_rebuild(config, **kwargs): + captured.update(kwargs) + raise CacheRebuildRejected("stop before teardown") + + engine = SimpleNamespace( + config=config, moe_offload_cache=None, kv_offloader=None, + linear_state_pool=SimpleNamespace(num_slots=5), + kv_cache=SimpleNamespace(validate_rebuild=validate_rebuild), + _baseline_free=1 << 30, _weights_bytes=0, num_pages=64, + ) + engine._target_moe_and_expert_bytes = lambda size: (0, 0) + with pytest.raises(CacheRebuildRejected): + Engine.rebuild_runtime_cache(engine, num_pages=32) + + workspace = gdn_prefill_workspace_bytes(config) + assert workspace > 0 + assert captured["extra_fixed_bytes"] == state_pool_bytes(config, 5) + workspace + # --------------------------------------------------------------------------- # offload-cache sizing guard + auto-resolution (_require_offload_cache_size / _adjust_config), # the floor rule compute_cache_floors documents above. diff --git a/tests/kvcache/test_linear_state_pool_alloc.py b/tests/kvcache/test_linear_state_pool_alloc.py index e9fe405426..e0edd0438e 100644 --- a/tests/kvcache/test_linear_state_pool_alloc.py +++ b/tests/kvcache/test_linear_state_pool_alloc.py @@ -9,6 +9,7 @@ from freetoken.kvcache.linear_state_pool import ( LinearStatePool, + gdn_prefill_workspace_bytes, linear_state_bytes_per_req, state_pool_bytes, ) @@ -189,3 +190,44 @@ def test_slot_state_bytes_for_the_real_geometry(): delta = linear_state_bytes_per_req(group, 1, torch.bfloat16, (spec,)) - \ linear_state_bytes_per_req(group, 1, torch.bfloat16) assert delta == 4 * 2560 * 9 * 2 == 180 * 1024 + + +def _gdn_config(**fields): + # qwen4_exp GDN geometry: 48 value heads of 128x128 + group = LinearGatedDeltaGroupConfig( + name="linear", layer_ids=(0, 1), + num_key_heads=16, num_value_heads=48, + key_head_dim=128, value_head_dim=128, conv_kernel_dim=4, output_gate="silu", + ) + mc = SimpleNamespace(linear_attention_group=lambda: group) + return SimpleNamespace(model_config=mc, dtype=torch.bfloat16, max_seq_len=262144, **fields) + + +def test_gdn_prefill_workspace_for_the_real_geometry(): + # h [8192/64 chunks, 48, 128, 128] bf16 = 192 MiB, plus v_new [8192, 48, 128] = 96 MiB + assert gdn_prefill_workspace_bytes(_gdn_config(max_extend_tokens=8192)) == 288 * 2**20 + assert gdn_prefill_workspace_bytes(_gdn_config(max_extend_tokens=2048)) == 72 * 2**20 + + +def test_gdn_prefill_workspace_without_a_scheduler_extend_budget(): + # a bare EngineConfig has no max_extend_tokens: price the scheduler's 8192 default + assert gdn_prefill_workspace_bytes(_gdn_config()) == 288 * 2**20 + + +def test_gdn_prefill_workspace_is_capped_by_the_sequence_length(): + config = _gdn_config(max_extend_tokens=8192) + config.max_seq_len = 100 # 2 chunks, 100 value rows + assert gdn_prefill_workspace_bytes(config) == (2 * 48 * 128 * 128 + 100 * 48 * 128) * 2 + + +def test_gdn_prefill_workspace_is_zero_without_a_linear_group(): + config = _gdn_config(max_extend_tokens=8192) + config.model_config.linear_attention_group = lambda: None + assert gdn_prefill_workspace_bytes(config) == 0 + + +def test_gdn_prefill_workspace_chunk_matches_the_kernel(): + from freetoken.kernel.fla.chunk_delta_h import CHUNK_SIZE + from freetoken.kvcache.linear_state_pool import _GDN_CHUNK_SIZE + + assert _GDN_CHUNK_SIZE == CHUNK_SIZE From c54922afc6e6ea58601d77eb31c29684412195de Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 10:25:31 -0400 Subject: [PATCH 49/73] fix(engine): count min_tokens progress from the device length Under overlap scheduling the sampler prepares batch N before batch N-1's token lands in req.input_ids, so the host count lagged by one and kept the stop ids masked for one token past min_tokens. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 8e4891fd54037f19aa42d43422a50cafc38513d9) --- python/freetoken/engine/sample.py | 4 +++- tests/engine/test_logits_processors.py | 26 ++++++++++++++++++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py index bca85c56fa..89d2a6fb88 100644 --- a/python/freetoken/engine/sample.py +++ b/python/freetoken/engine/sample.py @@ -220,7 +220,9 @@ def _plan(self, batch: Batch, params) -> LogitsPlan | None: min_ids: list[list[int]] = [] for local, (req, p) in enumerate(picked): if p.min_tokens > 0 and p.min_tokens_stop_ids: - generated = len(req.input_ids) - (req.max_device_len - req.output_len) + # device_len, not the host ids: under overlap scheduling the previous token + # reaches req.input_ids only after this batch is prepared + generated = req.device_len - (req.max_device_len - req.output_len) if generated < p.min_tokens: min_rows.append(local) min_ids.append([t for t in p.min_tokens_stop_ids if 0 <= t < self.vocab_size]) diff --git a/tests/engine/test_logits_processors.py b/tests/engine/test_logits_processors.py index a9645f4210..c9f498183f 100644 --- a/tests/engine/test_logits_processors.py +++ b/tests/engine/test_logits_processors.py @@ -85,6 +85,31 @@ def test_min_tokens_masks_the_stop_ids_of_the_row_only(): assert torch.isfinite(out[0]).all() +@pytest.mark.parametrize("overlap", [True, False]) +def test_min_tokens_releases_the_stop_ids_at_output_index_min_tokens(overlap): + """Overlap scheduling prepares batch N before batch N-1's token reaches req.input_ids on + the host; the count must not lag by that token and hold EOS one step too long.""" + from freetoken.core import Req + + sp = SamplingParams(min_tokens=3, max_tokens=10) + sp.min_tokens_stop_ids = [0] + req = Req(input_ids=torch.tensor([5, 6], dtype=torch.int32), table_idx=0, cached_len=0, + output_len=10, uid=1, sampling_params=sp, cache_handle=None) + sampler = Sampler(CPU, V) + masked, pending = [], None + for _ in range(6): + plan = sampler.prepare(SimpleNamespace(reqs=[req])).plan + masked.append(plan is not None and plan.min_rows is not None) + req.complete_one() # the forward launch advances the device length + if overlap: + if pending is not None: # batch N-1 drains after batch N launched + req.append_host(pending) + pending = torch.tensor([7], dtype=torch.int32) + else: + req.append_host(torch.tensor([7], dtype=torch.int32)) + assert [i for i, m in enumerate(masked) if m] == [0, 1, 2] + + def test_min_p_drops_tokens_below_the_fraction_of_the_top_probability(): logits = torch.log(torch.tensor([[0.5, 0.3, 0.1, 0.05, 0.03, 0.01, 0.005, 0.005]])) plan = _plan( @@ -104,6 +129,7 @@ def _req( ids = torch.tensor(prompt + generated, dtype=torch.int32) return SimpleNamespace( input_ids=ids, + device_len=len(ids), output_len=max_tokens, max_device_len=len(prompt) + max_tokens, sampling_params=SamplingParams(**sp), From 36a8fb5898d4ed0d12b0e9da8bfd716e5ed50274 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 1 Oct 2026 23:23:00 -0400 Subject: [PATCH 50/73] test(engine): import pytest for the parametrized min_tokens test Assisted-by: Claude Opus 5.5 --- tests/engine/test_logits_processors.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/engine/test_logits_processors.py b/tests/engine/test_logits_processors.py index c9f498183f..abd2e1f939 100644 --- a/tests/engine/test_logits_processors.py +++ b/tests/engine/test_logits_processors.py @@ -5,6 +5,7 @@ from types import SimpleNamespace +import pytest import torch from freetoken.core import SamplingParams from freetoken.engine.sample import LogitsPlan, Sampler, apply_logits_processors From 4f1ecf002e0b128eda719593abda64dca08c02df Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 06:01:42 -0400 Subject: [PATCH 51/73] fix(server): let protocol-level thinking requests win over --default-thinking-mode returned early, so reasoning_effort / reasoning.effort / thinking.type were silently turned off under --default-thinking-mode chat. count_tokens also skipped the default. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 8276ccd49a8fcb9123ea28f113ac39684b725665) --- python/freetoken/server/anthropic_api.py | 16 ++-- python/freetoken/server/args.py | 17 ++-- python/freetoken/server/model_meta.py | 19 +++++ python/freetoken/server/openai_api.py | 33 +------- python/freetoken/server/responses_api.py | 9 +-- tests/server/test_default_thinking_mode.py | 94 +++++++++++++++++++--- 6 files changed, 127 insertions(+), 61 deletions(-) diff --git a/python/freetoken/server/anthropic_api.py b/python/freetoken/server/anthropic_api.py index 4cb4b34822..a7c88b141a 100644 --- a/python/freetoken/server/anthropic_api.py +++ b/python/freetoken/server/anthropic_api.py @@ -166,7 +166,9 @@ async def handle_anthropic_count_tokens(req: AnthropicCountTokensRequest, state: # so it must not fall into the convert/empty-prompt ValueError branch. try: messages, template_tools, _, ctk = convert_anthropic_prompt( - req, reasoning_parser=getattr(state.config, "reasoning_parser", None) + req, + reasoning_parser=getattr(state.config, "reasoning_parser", None), + default_thinking_mode=getattr(state.config, "default_thinking_mode", "auto"), ) except ValueError as exc: return _anthropic_error_response(400, "invalid_request_error", str(exc)) @@ -198,6 +200,7 @@ async def handle_anthropic_count_tokens(req: AnthropicCountTokensRequest, state: def convert_anthropic_prompt( req: AnthropicMessagesRequest | AnthropicCountTokensRequest, reasoning_parser: str | None = None, + default_thinking_mode: str | None = None, ) -> tuple[list[dict[str, Any]], list[dict[str, Any]] | None, list[dict[str, Any]] | None, dict[str, Any]]: """(messages, template_tools, parser_tools, chat_template_kwargs) — the prompt side of the conversion, shared by /v1/messages and /v1/messages/count_tokens so @@ -310,7 +313,7 @@ def convert_anthropic_prompt( # Native extended-thinking toggle -> template kwargs, broadcast in every # spelling the ecosystem's templates read (a bare enable_thinking bool is # inert for templates that read a different knob, e.g. M3's thinking_mode). - from .model_meta import thinking_toggle_kwargs + from .model_meta import apply_default_thinking_mode, thinking_toggle_kwargs ctk: dict[str, Any] = {} if req.thinking: @@ -318,6 +321,10 @@ def convert_anthropic_prompt( ctk = thinking_toggle_kwargs(True) elif req.thinking.get("type") == "disabled": ctk = thinking_toggle_kwargs(False) + else: + # Any thinking block (e.g. "adaptive") is the client's explicit choice; the server + # default only covers requests that leave thinking out. + ctk = apply_default_thinking_mode(ctk, default_thinking_mode) return render_messages(messages), template_tools, parser_tools, ctk @@ -330,12 +337,9 @@ def convert_anthropic_to_genspec( inline_system_policy: str = "auto", default_thinking_mode: str | None = None, ) -> GenSpec: - from .openai_api import apply_default_thinking_mode - messages, template_tools, parser_tools, ctk = convert_anthropic_prompt( - req, reasoning_parser=reasoning_parser + req, reasoning_parser=reasoning_parser, default_thinking_mode=default_thinking_mode ) - ctk = apply_default_thinking_mode(ctk, default_thinking_mode) return GenSpec( messages=messages, sampling_params=resolve_sampling( diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 6d06cfcbf4..aaba2fc557 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -50,10 +50,10 @@ class ServerArgs(SchedulerConfig): # responses. None disables it (default for models without a reasoning protocol). reasoning_parser: str | None = None # Server-wide default thinking mode for reasoning-capable models: "auto" keeps the - # current per-request behavior; "chat" forces enable_thinking=False for every request - # that does not explicitly set it in chat_template_kwargs (for OpenAI-compatible - # clients that never send template kwargs, like Vercel AI SDK or llama-swap); - # "thinking" forces thinking on the same way. + # current per-request behavior; "chat" turns thinking off for every request that does + # not choose itself (template kwargs or a protocol-level effort/thinking field), for + # OpenAI-compatible clients that never send template kwargs, like Vercel AI SDK or + # llama-swap; "thinking" turns it on the same way. default_thinking_mode: str = "auto" # "model": fill unspecified request sampling params from generation_config.json # (temperature/top_k/top_p), like sglang. "none": use framework defaults only. @@ -687,10 +687,11 @@ def _infer_reasoning_parser(model_path: str) -> str | None: choices=["auto", "chat", "thinking"], help=( "Server-wide default thinking mode for reasoning-capable models. 'auto' keeps " - "the current per-request behavior; 'chat' forces enable_thinking=False for " - "every request that does not explicitly set it in chat_template_kwargs (for " - "OpenAI-compatible clients that never send template kwargs, like Vercel AI SDK " - "or llama-swap); 'thinking' forces thinking on the same way." + "the current per-request behavior; 'chat' turns thinking off for every request " + "that does not choose itself (chat_template_kwargs, reasoning_effort, " + "reasoning.effort, thinking) -- for OpenAI-compatible clients that never send " + "template kwargs, like Vercel AI SDK or llama-swap; 'thinking' turns it on the " + "same way." ), ) diff --git a/python/freetoken/server/model_meta.py b/python/freetoken/server/model_meta.py index 8ce8009a05..c9187ad629 100644 --- a/python/freetoken/server/model_meta.py +++ b/python/freetoken/server/model_meta.py @@ -119,6 +119,25 @@ def effort_toggle_kwargs( return mapped +def apply_default_thinking_mode( + ctk: dict[str, Any] | None, + default_mode: str | None, +) -> dict[str, Any] | None: + """Merge the server's --default-thinking-mode into a request's template kwargs. + + An explicit per-request value always wins; the server default only fills + in what the request left unset. "auto" (the default) is a no-op. Adapters + apply it after folding their protocol-level toggles (effort_toggle_kwargs), + so those count as explicit too. + """ + if default_mode not in ("chat", "thinking"): + return ctk + ctk = dict(ctk or {}) + if not any(key in ctk for key in _THINKING_KWARG_KEYS): + ctk.update(thinking_toggle_kwargs(default_mode == "thinking")) + return ctk + + def moe_total_experts(config: Any) -> int: """Total routed-expert slots the model has: experts per layer x MoE layers. Matches the engine's own basis (``Engine._resolve_auto_moe_cache_size``), so a residency rate derived diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index b09c480e1d..54c10ac2e6 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -71,34 +71,6 @@ def _thinking_type(req: Any) -> str | None: -def apply_default_thinking_mode( - ctk: dict[str, Any] | None, - default_mode: str | None, -) -> dict[str, Any] | None: - """Merge the server's --default-thinking-mode into a request's template kwargs. - - An explicit per-request value always wins; the server default only fills - in what the request left unset. "auto" (the default) is a no-op. - """ - if not default_mode or default_mode == "auto": - return ctk - if ctk is None: - ctk = {} - else: - ctk = dict(ctk) - has_explicit = ( - "enable_thinking" in ctk - or "thinking" in ctk - or "thinking_mode" in ctk - ) - if not has_explicit: - if default_mode == "chat": - ctk["enable_thinking"] = False - elif default_mode == "thinking": - ctk["enable_thinking"] = True - return ctk - - def chat_request_to_genspec( req: ChatCompletionRequest, model_sampling: dict[str, Any], @@ -106,12 +78,13 @@ def chat_request_to_genspec( default_thinking_mode: str | None = None, ) -> GenSpec: """OpenAI ChatCompletionRequest -> GenSpec (the OpenAI 'to_sampling_params').""" - from .model_meta import effort_toggle_kwargs + from .model_meta import apply_default_thinking_mode, effort_toggle_kwargs - ctk = apply_default_thinking_mode(req.chat_template_kwargs, default_thinking_mode) + ctk = req.chat_template_kwargs thinking_type = _thinking_type(req) if req.reasoning_effort or thinking_type: ctk = effort_toggle_kwargs(req.reasoning_effort, ctk, thinking_type=thinking_type) + ctk = apply_default_thinking_mode(ctk, default_thinking_mode) if req.continue_final_message: ctk = {**ctk, "continue_final_message": True} sampling_params = _resolve_sampling(req, model_sampling, default_max_tokens=default_max_tokens) diff --git a/python/freetoken/server/responses_api.py b/python/freetoken/server/responses_api.py index 9fd063a992..a84c68cb6d 100644 --- a/python/freetoken/server/responses_api.py +++ b/python/freetoken/server/responses_api.py @@ -245,15 +245,12 @@ def convert_responses_to_genspec( else: template_tools, parser_tools = split_tool_lists(raw_tools, selected) - from .model_meta import effort_toggle_kwargs + from .model_meta import apply_default_thinking_mode, effort_toggle_kwargs - from .openai_api import apply_default_thinking_mode - - ctk = apply_default_thinking_mode( - dict(getattr(req, "chat_template_kwargs", None) or {}), default_thinking_mode - ) + ctk = dict(getattr(req, "chat_template_kwargs", None) or {}) if req.reasoning: ctk = effort_toggle_kwargs(req.reasoning.get("effort"), ctk) + ctk = apply_default_thinking_mode(ctk, default_thinking_mode) return GenSpec( messages=render_messages(messages), diff --git a/tests/server/test_default_thinking_mode.py b/tests/server/test_default_thinking_mode.py index ac979e625d..7f9938f60c 100644 --- a/tests/server/test_default_thinking_mode.py +++ b/tests/server/test_default_thinking_mode.py @@ -1,7 +1,18 @@ """Unit tests for the --default-thinking-mode server flag merge logic.""" from __future__ import annotations -from freetoken.server.openai_api import apply_default_thinking_mode +import asyncio +from types import SimpleNamespace + +from freetoken.server import anthropic_api as A +from freetoken.server.anthropic_models import AnthropicCountTokensRequest, AnthropicMessagesRequest +from freetoken.server.api_models import ChatCompletionRequest +from freetoken.server.model_meta import apply_default_thinking_mode +from freetoken.server.openai_api import chat_request_to_genspec +from freetoken.server.responses_api import ResponsesRequest, convert_responses_to_genspec + +OFF = {"enable_thinking": False, "thinking_mode": "disabled"} +ON = {"enable_thinking": True, "thinking_mode": "enabled"} def test_auto_is_noop(): @@ -13,17 +24,14 @@ def test_auto_is_noop(): def test_chat_fills_unset_kwargs(): - assert apply_default_thinking_mode(None, "chat") == {"enable_thinking": False} - assert apply_default_thinking_mode({}, "chat") == {"enable_thinking": False} - assert apply_default_thinking_mode({"other": 1}, "chat") == { - "other": 1, - "enable_thinking": False, - } + assert apply_default_thinking_mode(None, "chat") == OFF + assert apply_default_thinking_mode({}, "chat") == OFF + assert apply_default_thinking_mode({"other": 1}, "chat") == {"other": 1, **OFF} def test_thinking_fills_unset_kwargs(): - assert apply_default_thinking_mode(None, "thinking") == {"enable_thinking": True} - assert apply_default_thinking_mode({}, "thinking") == {"enable_thinking": True} + assert apply_default_thinking_mode(None, "thinking") == ON + assert apply_default_thinking_mode({}, "thinking") == ON def test_explicit_request_value_wins(): @@ -38,14 +46,78 @@ def test_explicit_request_value_wins(): assert apply_default_thinking_mode({"thinking_mode": "thinking"}, "chat") == { "thinking_mode": "thinking" } + # A raw graded-effort request asks for thinking too. + assert apply_default_thinking_mode({"reasoning_effort": "high"}, "chat") == { + "reasoning_effort": "high" + } def test_original_dict_not_mutated(): original = {"other": 1} result = apply_default_thinking_mode(original, "chat") - assert result == {"other": 1, "enable_thinking": False} + assert result == {"other": 1, **OFF} assert original == {"other": 1} def test_unknown_mode_is_noop(): - assert apply_default_thinking_mode({"other": 1}, "bogus") == {"other": 1} \ No newline at end of file + assert apply_default_thinking_mode({"other": 1}, "bogus") == {"other": 1} + + +def _chat(**extra): + body = {"model": "m", "messages": [{"role": "user", "content": "hi"}], **extra} + return chat_request_to_genspec( + ChatCompletionRequest.model_validate(body), {}, default_thinking_mode="chat" + ).chat_template_kwargs + + +def test_chat_default_fills_a_request_without_thinking_controls(): + assert _chat() == OFF + + +def test_chat_default_keeps_protocol_level_thinking_requests(): + assert _chat(reasoning_effort="high")["enable_thinking"] is True + assert _chat(thinking={"type": "enabled"})["enable_thinking"] is True + assert _chat(reasoning_effort="none")["enable_thinking"] is False + + +def test_responses_chat_default_keeps_reasoning_effort(): + def ctk(**extra): + req = ResponsesRequest.model_validate({"model": "m", "input": "hi", **extra}) + return convert_responses_to_genspec(req, {}, default_thinking_mode="chat").chat_template_kwargs + + assert ctk() == OFF + on = ctk(reasoning={"effort": "medium"}) + assert on["enable_thinking"] is True and on["reasoning_effort"] == "medium" + + +def _anthropic(**extra): + return AnthropicMessagesRequest.model_validate( + {"model": "claude-x", "max_tokens": 64, "messages": [{"role": "user", "content": "hi"}], **extra} + ) + + +def test_anthropic_chat_default_applies_only_without_a_thinking_block(): + def ctk(**extra): + return A.convert_anthropic_to_genspec( + _anthropic(**extra), {}, default_thinking_mode="chat" + ).chat_template_kwargs + + assert ctk() == OFF + assert ctk(thinking={"type": "enabled", "budget_tokens": 1024}) == ON + assert ctk(thinking={"type": "adaptive"}) == {} + + +def test_anthropic_count_tokens_renders_the_generation_prompt_under_the_default(monkeypatch): + captured = {} + + async def fake_count(messages, template_tools, ctk, state): + captured["ctk"] = ctk + return 1 + + monkeypatch.setattr(A, "count_prompt_tokens", fake_count) + state = SimpleNamespace(config=SimpleNamespace(reasoning_parser=None, default_thinking_mode="chat")) + body = {"model": "claude-x", "messages": [{"role": "user", "content": "hi"}]} + asyncio.run(A.handle_anthropic_count_tokens(AnthropicCountTokensRequest.model_validate(body), state)) + + generation = A.convert_anthropic_to_genspec(_anthropic(), {}, default_thinking_mode="chat") + assert captured["ctk"] == generation.chat_template_kwargs == OFF From 58bd0ae33eea0ab1a6f867362999612ab27d6a68 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 05:57:18 -0400 Subject: [PATCH 52/73] fix(server): type allOf-wrapped refs and untyped enum/const tool parameters #435 sends any parameter it cannot type to the loose JSON parser, so an untyped string enum member such as "1" came back as the integer 1, outside the enum. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit ab2693c977c64e425d428468364e2a2743f170e4) --- .../freetoken/server/function_call_parser.py | 38 +++++++++- ..._function_call_parser_schema_resolution.py | 70 +++++++++++++++++++ 2 files changed, 106 insertions(+), 2 deletions(-) diff --git a/python/freetoken/server/function_call_parser.py b/python/freetoken/server/function_call_parser.py index 08c60a2993..88cc65b367 100644 --- a/python/freetoken/server/function_call_parser.py +++ b/python/freetoken/server/function_call_parser.py @@ -382,7 +382,27 @@ def _resolve_local_ref(self, schema: Dict, root_schema: Dict, _seen: Optional[se return self._resolve_local_ref(merged, root_schema, _seen) return merged - def _normalize_param_schema(self, schema: Dict, root_schema: Dict) -> Dict: + @staticmethod + def _literal_schema_type(schema: Dict) -> Optional[str]: + """Type of an untyped ``enum``/``const`` schema, inferred from its literals.""" + if "const" in schema: + values = [schema["const"]] + elif isinstance(schema.get("enum"), list): + values = schema["enum"] + else: + return None + values = [v for v in values if v is not None] + if not values: + return None + if all(isinstance(v, str) for v in values): + return "string" + if all(isinstance(v, bool) for v in values): + return "boolean" + if all(isinstance(v, int) and not isinstance(v, bool) for v in values): + return "integer" + return "loose" + + def _normalize_param_schema(self, schema: Dict, root_schema: Dict, _depth: int = 0) -> Dict: """Resolve the schema parts needed for tool argument typing.""" if not isinstance(schema, dict): return {"type": "string"} @@ -395,6 +415,18 @@ def _normalize_param_schema(self, schema: Dict, root_schema: Dict) -> Dict: return result if isinstance(param_type, str): return schema + all_of = schema.get("allOf") + if ( + isinstance(all_of, list) + and len(all_of) == 1 + and isinstance(all_of[0], dict) + and _depth < 16 # bounds a self-referential allOf chain + ): + # Pydantic v1 / OpenAPI 3.0 wrap a $ref in a one-element allOf to attach + # siblings such as default/description; type it like the bare $ref. + merged = dict(all_of[0]) + merged.update({k: v for k, v in schema.items() if k != "allOf"}) + return self._normalize_param_schema(merged, root_schema, _depth + 1) for keyword in ("oneOf", "anyOf"): branches = schema.get(keyword) if not isinstance(branches, list): @@ -405,6 +437,8 @@ def _normalize_param_schema(self, schema: Dict, root_schema: Dict) -> Dict: continue branch = self._resolve_local_ref(branch, root_schema) branch_type = branch.get("type") + if not isinstance(branch_type, str): + branch_type = self._literal_schema_type(branch) if isinstance(branch_type, str) and branch_type != "null": types.append(branch_type) types = list(dict.fromkeys(types)) @@ -420,7 +454,7 @@ def _normalize_param_schema(self, schema: Dict, root_schema: Dict) -> Dict: result["type"] = "array" return result result = dict(schema) - result["type"] = "loose" + result["type"] = self._literal_schema_type(schema) or "loose" return result def _get_tool_schema(self, func_name: str, tools: List[Tool]) -> Dict: diff --git a/tests/server/test_function_call_parser_schema_resolution.py b/tests/server/test_function_call_parser_schema_resolution.py index fbc0fe270b..468acc3bd2 100644 --- a/tests/server/test_function_call_parser_schema_resolution.py +++ b/tests/server/test_function_call_parser_schema_resolution.py @@ -375,6 +375,67 @@ def _parse_args(parser_name: str, prop_schema: dict, root_extra: dict | None, ra 5, id="ambiguous-union-loose-json", ), + + # An enum/const with no "type" takes its type from the literals, so a + # numeric-looking string member is not coerced out of the enum. + pytest.param( + {"enum": ["1", "2"]}, + None, + "1", + "1", + id="untyped-string-enum", + ), + pytest.param( + {"enum": [1, 2]}, + None, + "2", + 2, + id="untyped-integer-enum", + ), + pytest.param( + {"const": "10"}, + None, + "10", + "10", + id="untyped-string-const", + ), + pytest.param( + { + "anyOf": [ + {"enum": ["1", "2"]}, + {"type": "null"}, + ] + }, + None, + "1", + "1", + id="anyof-untyped-string-enum-nullable", + ), + + # Pydantic v1 / OpenAPI 3.0 wrap a $ref in a one-element allOf to attach + # siblings such as default/description. + pytest.param( + {"allOf": [{"$ref": "#/$defs/Mode"}], "default": "a"}, + { + "$defs": { + "Mode": {"type": "string", "enum": ["a", "true"]}, + } + }, + "true", + "true", + id="allof-ref-string-enum", + ), + pytest.param( + {"allOf": [{"$ref": "#/$defs/Limit"}], "description": "max rows"}, + { + "$defs": { + "Limit": {"type": "integer"}, + } + }, + "5", + 5, + id="allof-ref-integer", + ), ] @@ -454,6 +515,15 @@ def test_schema_aware_parsers_resolve_indirect_schema_types( }, id="mutually-recursive-ref", ), + pytest.param( + {"allOf": [{"$ref": "#/$defs/A"}]}, + { + "$defs": { + "A": {"allOf": [{"$ref": "#/$defs/A"}]}, + } + }, + id="self-referential-allof", + ), ] From 90a61759415534cc4b704019b7e84eb1a7c29e83 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 06:03:44 -0400 Subject: [PATCH 53/73] fix(server): keep /v1/models/ listing the model after the retrieve route Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 7ccc433dcc8608b23f33c1a4be6313b5f9307c47) --- python/freetoken/server/openai_api.py | 3 +++ tests/server/test_openai_api.py | 26 ++++++++++++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index 54c10ac2e6..7f94efc250 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -176,6 +176,9 @@ async def v1_model(model_id: str): OpenAI-compatible clients may use this endpoint to validate a model name before sending completion requests. """ + if not model_id: + # The :path converter also matches "/v1/models/", which used to redirect to the list. + return await v1_models() state = get_state() served_model_id = _served_model_name(state) diff --git a/tests/server/test_openai_api.py b/tests/server/test_openai_api.py index 376cc2a7ba..7e6d4ed30e 100644 --- a/tests/server/test_openai_api.py +++ b/tests/server/test_openai_api.py @@ -523,6 +523,32 @@ def test_models_route_prefers_the_last_rebuild_over_the_load_time_pool(): assert card["max_model_len"] == 40000 +def test_model_retrieve_route_returns_the_served_card_or_404(): + state = FakeState([]) + state.config.max_seq_len = 262144 + app = FastAPI() + register_openai_routes(app, lambda: state, lambda: {}) + client = TestClient(app) + + card = client.get("/v1/models/unit-model").json() + assert card["id"] == "unit-model" and card["context_length"] == 262144 + + missing = client.get("/v1/models/other-model") + assert missing.status_code == 404 + assert missing.json()["error"]["code"] == "model_not_found" + + +def test_models_route_with_a_trailing_slash_still_lists(): + state = FakeState([]) + app = FastAPI() + register_openai_routes(app, lambda: state, lambda: {}) + + response = TestClient(app).get("/v1/models/") + + assert response.status_code == 200 + assert response.json()["data"][0]["id"] == "unit-model" + + async def _collect(generator): return [chunk async for chunk in generator] From a33da2f47af07345bae4adb444c29d6bc9dd6059 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 06:22:47 -0400 Subject: [PATCH 54/73] fix(server): open the streamed thinking block with the signature field The full response and the closing signature_delta already send signature="", but exclude_none dropped it from content_block_start, where the SDK types it as required. Taken from upstream #517 (b0442764b014bd32a68ba29055eb20eda9bfb026). Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 55a8e72f21cb679b585cfe6f7f738aafedeb7369) --- python/freetoken/server/anthropic_api.py | 4 +++- tests/server/test_anthropic_api.py | 3 ++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/python/freetoken/server/anthropic_api.py b/python/freetoken/server/anthropic_api.py index a7c88b141a..545cf4e7cc 100644 --- a/python/freetoken/server/anthropic_api.py +++ b/python/freetoken/server/anthropic_api.py @@ -480,7 +480,9 @@ def _open_thinking() -> str: block_open = "thinking" return _event(AnthropicStreamEvent( type="content_block_start", index=block_index, - content_block=AnthropicContentBlock(type="thinking", thinking=""), + # The SDK types signature as a required string on the opening block, before any + # signature_delta. Local reasoning is unsigned, like the full-response path. + content_block=AnthropicContentBlock(type="thinking", thinking="", signature=""), )) def _stop_block() -> list[str]: diff --git a/tests/server/test_anthropic_api.py b/tests/server/test_anthropic_api.py index 1a09b83a1c..890339522c 100644 --- a/tests/server/test_anthropic_api.py +++ b/tests/server/test_anthropic_api.py @@ -431,7 +431,8 @@ def test_stream_reasoning_as_thinking_block(): "message_stop", ] blocks = [e[1] for e in collected if e[0] == "content_block_start"] - assert blocks[0]["content_block"]["type"] == "thinking" + # signature is a required string in the SDK's thinking block, present from the start. + assert blocks[0]["content_block"] == {"type": "thinking", "thinking": "", "signature": ""} assert blocks[0]["index"] == 0 assert blocks[1]["content_block"]["type"] == "text" assert blocks[1]["index"] == 1 From 481738d28d8ec1afc51aedbf85fd99e4ef223cb3 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 06:22:47 -0400 Subject: [PATCH 55/73] test(server): exercise detector SDK and lifecycle contracts Upstream #517 tests without its docs and pyproject dev-extra: the SDK tests importorskip anthropic/openai/httpx so the server suite still collects without them. (cherry picked from commit b0442764b014bd32a68ba29055eb20eda9bfb026, tests only) Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 50f99d167304858d9eaebfbcd3d00aa3c1f06430) --- tests/server/test_accounting.py | 121 ++++++++++++ tests/server/test_function_call_parser.py | 129 +++++++++++++ tests/server/test_protocol_compat.py | 213 ++++++++++++++++++++++ tests/server/test_reasoning_parser.py | 108 +++++++++++ tests/server/test_supervisor.py | 37 ++++ 5 files changed, 608 insertions(+) create mode 100644 tests/server/test_accounting.py create mode 100644 tests/server/test_protocol_compat.py create mode 100644 tests/server/test_reasoning_parser.py diff --git a/tests/server/test_accounting.py b/tests/server/test_accounting.py new file mode 100644 index 0000000000..fe29d414d8 --- /dev/null +++ b/tests/server/test_accounting.py @@ -0,0 +1,121 @@ +"""Failure/retry and concurrency contracts for the terminal accounting barrier.""" +import asyncio +from types import SimpleNamespace + +import pytest + +from freetoken.message import UserReply +from freetoken.server.accounting import AccountingDrainError, prepare_stop_accounting +from freetoken.server.stats import StatsTracker + + +def state_with_request(): + stats = StatsTracker() + stats.on_new_user(7) + return SimpleNamespace(stats=stats, maintenance_state="serving", instance_id="instance", + ready_at=None, config=SimpleNamespace(served_model_name="test")) + + +def terminal(state, tokens=3): + state.stats.observe(UserReply(uid=7, incremental_output="", finished=True, + completion_tokens_delta=tokens, error="aborted")) + + +def test_drain_timeout_then_late_terminal_can_retry_without_losing_tokens(): + async def run(): + state = state_with_request() + aborted = [] + + async def abort(uid): + aborted.append(uid) + state.stats.on_abort(uid) # enqueueing abort is NOT a terminal reply + + state.abort_user = abort + with pytest.raises(AccountingDrainError, match="abort barrier timed out"): + await prepare_stop_accounting(state, drain_timeout_s=0, abort_timeout_s=0) + assert state.maintenance_state == "stopping" + assert state.stats.active == 1 and not hasattr(state, "_sealed_accounting") + terminal(state, tokens=11) # last sampled-token delta arrives AFTER the timeout + sealed = await prepare_stop_accounting(state, drain_timeout_s=0, abort_timeout_s=0) + assert sealed["completion_tokens_total"] == 11 + assert sealed["drain_complete"] is True + assert aborted == [7] + sealed["completion_tokens_total"] = -1 + assert (await prepare_stop_accounting(state))["completion_tokens_total"] == 11 + asyncio.run(run()) + + +def test_concurrent_prepare_stop_serializes_abort_and_returns_independent_snapshots(): + async def run(): + state = state_with_request() + entered, release = asyncio.Event(), asyncio.Event() + aborted = [] + + async def abort(uid): + assert state.maintenance_state == "stopping" + aborted.append(uid) + entered.set() + await release.wait() + terminal(state) + + state.abort_user = abort + first = asyncio.create_task(prepare_stop_accounting(state, drain_timeout_s=0)) + await entered.wait() + second = asyncio.create_task(prepare_stop_accounting(state, drain_timeout_s=0)) + await asyncio.sleep(0) # second caller reaches the shared lock, no wall-clock race + assert not second.done() + release.set() + a, b = await asyncio.gather(first, second) + assert a == b and a is not b + assert aborted == [7] + assert a["completion_tokens_total"] == 3 + asyncio.run(run()) + + +def test_abort_transport_failure_preserves_engine_and_closes_gate(): + state = state_with_request() + + async def abort(uid): + raise ConnectionError("scheduler disconnected") + + state.abort_user = abort + with pytest.raises(AccountingDrainError, match="failed to abort active requests") as error: + asyncio.run(prepare_stop_accounting(state, drain_timeout_s=0, abort_timeout_s=0)) + assert isinstance(error.value.__cause__, ConnectionError) + assert state.maintenance_state == "stopping" and state.stats.active == 1 + assert not hasattr(state, "_sealed_accounting") + + +def test_unidentified_active_request_cannot_be_sealed(): + state = state_with_request() + state.stats = SimpleNamespace(active=1, inflight_uids=()) + with pytest.raises(AccountingDrainError, match="unidentified"): + asyncio.run(prepare_stop_accounting(state, drain_timeout_s=0)) + assert state.maintenance_state == "stopping" + assert not hasattr(state, "_sealed_accounting") + + +@pytest.mark.parametrize("maintenance", ["rebuilding", "unknown"]) +def test_invalid_prepare_transition_preserves_previous_state(maintenance): + state = state_with_request() + state.maintenance_state = maintenance + with pytest.raises(AccountingDrainError): + asyncio.run(prepare_stop_accounting(state, drain_timeout_s=0)) + assert state.maintenance_state == maintenance + assert not hasattr(state, "_sealed_accounting") + + +def test_cancelled_drain_releases_lock_without_reopening_admission(): + async def run(): + state = state_with_request() + task = asyncio.create_task(prepare_stop_accounting(state, drain_timeout_s=10)) + await asyncio.sleep(0) # admission closes before the first drain wait + assert state.maintenance_state == "stopping" + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert not state._accounting_prepare_lock.locked() + assert not hasattr(state, "_sealed_accounting") + terminal(state) + assert (await prepare_stop_accounting(state))["drain_complete"] is True + asyncio.run(run()) diff --git a/tests/server/test_function_call_parser.py b/tests/server/test_function_call_parser.py index bc22296760..23cd12272e 100644 --- a/tests/server/test_function_call_parser.py +++ b/tests/server/test_function_call_parser.py @@ -371,3 +371,132 @@ def test_streaming_support_flags(): # test_streaming_model_matrix.py::test_non_streaming_detector_falls_back_to_buffered_parse). for name in SUPPORTED_TOOL_CALL_PARSERS: assert FunctionCallParser(TOOLS, tool_call_parser=name).supports_streaming() is True + + +# One class per concrete detector makes format coverage visible in pytest output. +# A shared contract avoids accidentally giving a newly added format weaker checks. +class _DetectorContract: + parser_name: str + block: str + truncate_before: str + buffered_recovery = False + + def parser(self): + return FunctionCallParser(OPENCODE_TOOLS, self.parser_name) + + def test_complete_call(self): + result = self.parser().parse_non_stream(self.block) + assert result.normal_text == "" + assert len(result.calls) == 1 + assert result.calls[0].name == "read" + assert json.loads(result.calls[0].parameters) == {"filePath": "/tmp/test_calc.py"} + + def test_every_character_boundary(self): + parser = self.parser() + texts, calls = _feed(parser, self.block) + assert "".join(texts) + parser.finish_stream() == "" + assert [call.name for call in calls if call.name] == ["read"] + assert json.loads("".join(call.parameters for call in calls)) == { + "filePath": "/tmp/test_calc.py" + } + # A completed call must never be duplicated by the end-of-stream recovery. + assert parser.recover_truncated_call() == [] + + def test_truncated_mid_call_recovery(self): + parser = self.parser() + cut = self.block.rindex(self.truncate_before) + texts, calls = _feed(parser, self.block[:cut]) + recovered = parser.recover_truncated_call() + assert bool(recovered) is self.buffered_recovery + # Incremental detectors have already consumed the opener and emitted the + # call identity. Their recovery method MUST NOT invent a duplicate call: + # generation.py recovers missing arguments from their partial-parse ledger. + # Buffered M3 instead recovers a complete call from the retained markup. + all_calls = calls + recovered + assert [call.name for call in all_calls if call.name] == ["read"] + arguments = "".join(call.parameters for call in all_calls) + try: + parsed = json.loads(arguments) + except json.JSONDecodeError: + parsed = json.loads(parser.unstreamed_arguments(0)) + assert parsed == {"filePath": "/tmp/test_calc.py"} + assert parser.recover_truncated_call() == [] + assert "".join(texts) + parser.finish_stream() == "" + + +class TestQwen25Detector(_DetectorContract): + parser_name = 'qwen25' + block = '{"name": "read", "arguments": {"filePath": "/tmp/test_calc.py"}}' + truncate_before = '' + + +class TestMistralDetector(_DetectorContract): + parser_name = 'mistral' + block = '[TOOL_CALLS] [{"name": "read", "arguments": {"filePath": "/tmp/test_calc.py"}}]' + truncate_before = '}]' + + +class TestLlama32Detector(_DetectorContract): + parser_name = 'llama3' + block = '<|python_tag|>{"name": "read", "arguments": {"filePath": "/tmp/test_calc.py"}}' + truncate_before = '}' + + +class TestGlm47Detector(_DetectorContract): + parser_name = 'glm47' + block = 'readfilePath/tmp/test_calc.py' + truncate_before = '' + + +class TestDeepSeekV32Detector(_DetectorContract): + parser_name = 'deepseekv32' + block = '<｜DSML｜function_calls><｜DSML｜invoke name="read"><｜DSML｜parameter name="filePath" string="true">/tmp/test_calc.py' + truncate_before = '' + + +class TestQwen3CoderDetector(_DetectorContract): + parser_name = 'qwen3_coder' + block = '/tmp/test_calc.py' + truncate_before = '' + + +class TestGemma4Detector(_DetectorContract): + parser_name = 'gemma4' + block = '<|tool_call>call:read{filePath:<|"|>/tmp/test_calc.py<|"|>}' + truncate_before = '' + + +class TestMiniMaxDetector(_DetectorContract): + parser_name = 'minimax' + block = '/tmp/test_calc.py' + truncate_before = '' + + +class TestMiniMaxM3Detector(_DetectorContract): + parser_name = 'minimax_m3' + block = ']<]minimax[>[\n]<]minimax[>[]<]minimax[>[/tmp/test_calc.py]<]minimax[>[]<]minimax[>[\n]<]minimax[>[' + truncate_before = ']<]minimax[>[' + buffered_recovery = True + + +class TestGptOssDetector(_DetectorContract): + parser_name = 'gpt_oss' + block = '<|start|>assistant<|channel|>commentary to=functions.read <|constrain|>json<|message|>{"filePath": "/tmp/test_calc.py"}<|end|>' + truncate_before = '<|end|>' + + +class TestMuseGlimmerDetector(_DetectorContract): + parser_name = 'muse_glimmer' + block = '<|start|>assistant to=read<|message|>\n\n/tmp/test_calc.py\n\n<|eot|>' + truncate_before = '' + + +def test_contract_classes_cover_every_concrete_detector(): + from freetoken.server.function_call_parser import BaseFormatDetector + + covered = { + type(cls().parser().detector) + for cls in _DetectorContract.__subclasses__() + } + assert covered == set(BaseFormatDetector.__subclasses__()) + assert covered == set(FunctionCallParser.ToolCallParserEnum.values()) diff --git a/tests/server/test_protocol_compat.py b/tests/server/test_protocol_compat.py new file mode 100644 index 0000000000..8dc22908d6 --- /dev/null +++ b/tests/server/test_protocol_compat.py @@ -0,0 +1,213 @@ +"""Real SDK model/stream validation over the real HTTP adapters and GenSpec path. + +Only the scheduler/tokenizer transport is scripted. In-process ASGI transport +keeps these deterministic and CPU-only; it does not prove TCP chunk timing, +disconnect behavior, tokenizer templates, or a real model's output quality. +""" +import asyncio +import json +from types import SimpleNamespace + +import pytest +from fastapi import FastAPI + +from freetoken.message import TokenizeMsg, UserReply +from freetoken.server.anthropic_api import register_anthropic_routes +from freetoken.server.openai_api import register_openai_routes +from freetoken.server.responses_api import register_responses_routes + +# The SDKs are optional test dependencies: skip what needs a missing one instead of +# failing collection of the whole server suite. The Anthropic half skips per test. +httpx = pytest.importorskip("httpx") +openai = pytest.importorskip("openai") + + +SCHEMA = {"type": "object", "properties": {"city": {"type": "string"}}} +TOOL = {"type": "function", "function": {"name": "get_weather", "parameters": SCHEMA}} +CALL = '{"name":"get_weather","arguments":{"city":"Paris"}}' + + +class ScriptedBackend: + def __init__(self, text, *, reasoning=False): + self.config = SimpleNamespace( + model_path="/models/test", served_model_name="test", max_seq_len=4096, + tool_call_parser="qwen25", reasoning_parser="qwen3" if reasoning else None, + ) + self.text = text + self.sent = [] + self.maintenance_state = "serving" + + def new_user(self): + return 42 + + async def send_one(self, message): + assert isinstance(message, TokenizeMsg) + self.sent.append(message) + + async def wait_for_ack(self, uid): + # One character per backend reply exercises adapter stream aggregation + # as well as the SDK SSE decoder. Usage deltas remain authoritative. + for i, char in enumerate(self.text): + yield UserReply(uid=uid, incremental_output=char, + finished=i == len(self.text) - 1, + prompt_tokens_delta=5 if i == 0 else 0, + completion_tokens_delta=1) + + async def stream_with_cancellation(self, events, request, uid): + async for event in events: + yield event + + +def client(state, vendor="openai"): + app = FastAPI() + for register in (register_openai_routes, register_anthropic_routes, register_responses_routes): + register(app, lambda: state, dict) + # Recent Anthropic SDKs use httpx2 and reject httpx client instances. + # Select the transport matching its public default client class; older SDKs + # in the supported range still use httpx. Neither path bypasses SDK validation. + http = httpx + if vendor == "anthropic": + anthropic = pytest.importorskip("anthropic") + if not issubclass(anthropic.DefaultAsyncHttpxClient, httpx.AsyncClient): + http = pytest.importorskip("httpx2") + transport = http.AsyncClient(transport=http.ASGITransport(app=app)) + cls = openai.AsyncOpenAI if vendor == "openai" else anthropic.AsyncAnthropic + return cls(base_url="http://test/v1" if vendor == "openai" else "http://test", + api_key="test", http_client=transport, max_retries=0, + _strict_response_validation=True) + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("tool", [False, True]) +def test_openai_chat_sdk_parses_reasoning_text_tools_and_usage(stream, tool): + async def run(): + state = ScriptedBackend("consider" + (CALL if tool else "Hello"), reasoning=True) + async with client(state) as sdk: + response = await sdk.chat.completions.create( + model="test", messages=[{"role": "user", "content": "Hi"}], + tools=[TOOL], stream=stream, stream_options={"include_usage": True}, + ) + if stream: + chunks = [chunk async for chunk in response] + choices = [choice for chunk in chunks for choice in chunk.choices] + assert "".join(getattr(c.delta, "reasoning_content", "") or "" for c in choices) == "consider" + assert [c.finish_reason for c in choices if c.finish_reason] == ["tool_calls" if tool else "stop"] + if tool: + calls = [tc for c in choices for tc in c.delta.tool_calls or []] + assert [tc.function.name for tc in calls if tc.function.name] == ["get_weather"] + assert json.loads("".join(tc.function.arguments or "" for tc in calls)) == {"city": "Paris"} + assert all(tc.index == 0 for tc in calls) + else: + assert "".join(c.delta.content or "" for c in choices) == "Hello" + usage = chunks[-1].usage + else: + assert response.choices[0].message.reasoning_content == "consider" + assert response.choices[0].finish_reason == ("tool_calls" if tool else "stop") + if tool: + call = response.choices[0].message.tool_calls[0] + assert call.id and call.function.name == "get_weather" + assert json.loads(call.function.arguments) == {"city": "Paris"} + else: + assert response.choices[0].message.content == "Hello" + usage = response.usage + assert usage.prompt_tokens == 5 + assert usage.completion_tokens == len(state.text) + assert usage.total_tokens == 5 + len(state.text) + assert len(state.sent) == 1 + asyncio.run(run()) + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("tool", [False, True]) +def test_anthropic_sdk_assembles_typed_message(stream, tool): + pytest.importorskip("anthropic") + + async def run(): + state = ScriptedBackend("consider" + (CALL if tool else "Hello"), reasoning=True) + async with client(state, "anthropic") as sdk: + kwargs = dict(model="test", max_tokens=256, + messages=[{"role": "user", "content": "Hi"}], + tools=[{"name": "get_weather", "input_schema": SCHEMA}]) + if stream: + async with sdk.messages.stream(**kwargs) as response: + events = [event async for event in response] + result = await response.get_final_message() + assert events[0].type == "message_start" + assert events[-1].type == "message_stop" + else: + result = await sdk.messages.create(**kwargs) + assert result.role == "assistant" and result.model == "test" + assert result.stop_reason == ("tool_use" if tool else "end_turn") + assert result.content[0].type == "thinking" + assert result.content[0].thinking == "consider" + if tool: + # The established full-response shape includes an empty text + # block before tools; preserving that is part of compatibility. + assert all(b.text == "" for b in result.content if b.type == "text") + block = next(b for b in result.content if b.type == "tool_use") + assert block.id + assert block.name == "get_weather" and block.input == {"city": "Paris"} + else: + block = result.content[1] + assert block.type == "text" and block.text == "Hello" + assert result.usage.input_tokens == 5 + assert result.usage.output_tokens == len(state.text) + assert len(state.sent) == 1 + asyncio.run(run()) + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("tool", [False, True]) +def test_responses_sdk_parses_output_and_events(stream, tool): + async def run(): + state = ScriptedBackend(CALL if tool else "Hello") + async with client(state) as sdk: + response = await sdk.responses.create( + model="test", input="Hi", stream=stream, + tools=[{"type": "function", "name": "get_weather", "parameters": SCHEMA}], + ) + if stream: + events = [event async for event in response] + assert events[0].type == "response.created" + assert events[-1].type == "response.completed" + assert [e.sequence_number for e in events] == list(range(len(events))) + result = events[-1].response + deltas = [e for e in events if e.type == ("response.function_call_arguments.delta" if tool else "response.output_text.delta")] + text = "".join(e.delta for e in deltas) + assert json.loads(text) == {"city": "Paris"} if tool else text == "Hello" + else: + result = response + assert result.status == "completed" + if tool: + call = next(item for item in result.output if item.type == "function_call") + assert call.call_id and call.name == "get_weather" + assert json.loads(call.arguments) == {"city": "Paris"} + else: + assert result.output_text == "Hello" + assert result.usage.input_tokens == 5 + assert result.usage.output_tokens == len(state.text) + assert result.usage.total_tokens == 5 + len(state.text) + assert len(state.sent) == 1 + asyncio.run(run()) + + +@pytest.mark.parametrize("protocol", ["chat", "responses", "anthropic"]) +def test_sdk_classifies_invalid_request_as_400(protocol): + vendor = "anthropic" if protocol == "anthropic" else "openai" + sdk_module = pytest.importorskip("anthropic") if vendor == "anthropic" else openai + + async def run(): + state = ScriptedBackend("unused") + error = sdk_module.BadRequestError + async with client(state, vendor) as sdk: + with pytest.raises(error) as caught: + if protocol == "chat": + await sdk.chat.completions.create(model="test", messages=[{"role": "user", "content": "Hi"}], n=2) + elif protocol == "responses": + await sdk.responses.create(model="test", input="Hi", background=True) + else: + await sdk.messages.create(model="test", max_tokens=0, messages=[{"role": "user", "content": "Hi"}]) + assert caught.value.status_code == 400 + assert "invalid_request_error" in str(caught.value) + assert state.sent == [] + asyncio.run(run()) diff --git a/tests/server/test_reasoning_parser.py b/tests/server/test_reasoning_parser.py new file mode 100644 index 0000000000..05c207d0e1 --- /dev/null +++ b/tests/server/test_reasoning_parser.py @@ -0,0 +1,108 @@ +"""Character-boundary contracts for every reasoning detector and template mode.""" +from types import SimpleNamespace + +import pytest + +from freetoken.server.generation import GenSpec, _make_reasoning_parser +from freetoken.server.reasoning_parser import BaseReasoningParser, ReasoningParser + + +class _ReasoningContract: + name: str + opener: str + closer: str + content_opener = "" + + def parser(self): + return ReasoningParser(self.name, force_reasoning=False) + + def stream(self, text): + parser = self.parser() + parts = [parser.parse_stream_chunk(char) for char in text] + parts.append(parser.flush()) + assert parser.flush() == ("", "") + return tuple("".join(pair[i] for pair in parts) for i in (0, 1)) + + def test_complete_reasoning(self): + text = self.opener + "consider" + self.closer + self.content_opener + "answer" + assert self.parser().parse_non_stream(text) == ("consider", "answer") + + def test_every_character_boundary(self): + text = self.opener + "consider" + self.closer + self.content_opener + "answer" + assert self.stream(text) == ("consider", "answer") + + def test_truncated_reasoning_flush(self): + # Reasoning detectors expose flush(), not recover_truncated_call(). EOS + # inside an open reasoning body must retain it in the reasoning channel. + assert self.stream(self.opener + "partial thought") == ("partial thought", "") + + +class TestDeepSeekV32ReasoningParser(_ReasoningContract): + name = "deepseekv32" + opener = "" + closer = "" + + +class TestThinkReasoningParser(_ReasoningContract): + name = "qwen3" + opener = "" + closer = "" + + +class TestMiniMaxM3ReasoningParser(_ReasoningContract): + name = "minimax_m3" + opener = "" + closer = "" + + +class TestGemmaThoughtReasoningParser(_ReasoningContract): + name = "gemma4" + opener = "<|channel>thought\n" + closer = "" + + +class TestGptOssHarmonyReasoningParser(_ReasoningContract): + name = "gpt_oss" + opener = "<|channel|>analysis<|message|>" + closer = "<|end|>" + content_opener = "<|start|>assistant<|channel|>final<|message|>" + + +class TestMuseGlimmerReasoningParser(_ReasoningContract): + name = "muse_glimmer" + opener = "<|start|>assistant to=self<|message|>" + closer = "<|eom|>" + content_opener = "<|start|>assistant to=user<|message|>" + + +def test_contract_classes_cover_every_reasoning_detector(): + covered = {type(cls().parser().detector) for cls in _ReasoningContract.__subclasses__()} + assert covered == set(BaseReasoningParser.__subclasses__()) + assert covered == set(ReasoningParser.ReasoningParserEnum.values()) + + +@pytest.mark.parametrize("tools", [None, [{"type": "function", "function": {"name": "read"}}]]) +@pytest.mark.parametrize(("name", "kwargs", "forced"), [ + ("qwen3", {}, True), + ("qwen3", {"enable_thinking": True}, True), + ("qwen3", {"enable_thinking": False}, False), + ("glm", {}, True), + ("glm", {"enable_thinking": True}, True), + ("glm", {"enable_thinking": False}, False), + ("gemma4", {}, False), + ("gemma4", {"thinking_mode": "thinking"}, True), + ("gemma4", {"enable_thinking": True}, True), + ("gemma4", {"thinking": True}, True), + ("gemma4", {"thinking_mode": "chat", "enable_thinking": False}, False), + ("minimax_m3", {}, False), + ("minimax_m3", {"thinking_mode": "adaptive"}, False), + ("minimax_m3", {"thinking_mode": "disabled"}, False), + ("minimax_m3", {"thinking_mode": "enabled"}, True), +]) +def test_generation_force_reasoning_matches_template(name, kwargs, forced, tools): + spec = GenSpec(messages=[], sampling_params=None, chat_template_kwargs=kwargs, + template_tools=tools, parser_tools=tools) + parser = _make_reasoning_parser(spec, SimpleNamespace(config=SimpleNamespace(reasoning_parser=name))) + # Observe routing, not an internal flag: tools must not turn a disabled + # template's ordinary response into invisible reasoning. + assert parser.parse_non_stream("probe") == (("probe", "") if forced else ("", "probe")) diff --git a/tests/server/test_supervisor.py b/tests/server/test_supervisor.py index 3568d90505..c4772603b0 100644 --- a/tests/server/test_supervisor.py +++ b/tests/server/test_supervisor.py @@ -1,5 +1,7 @@ from __future__ import annotations +import pytest + import queue from queue import Empty as _Empty @@ -270,3 +272,38 @@ def test_byte_bar_emits_to_installed_sink_then_stops_after_clear(): bar.update(100) bar.close() assert seen == [] + + +@pytest.mark.parametrize("shutdown", [False, True]) +def test_shutdown_flag_is_read_at_later_worker_death(monkeypatch, shutdown): + """Readiness must not snapshot shutdown state: the flag changes on a later poll. + + Keep TP0 alive and kill a different worker, so the first process cannot mask + a dead tokenizer. No real processes or sleep timing enter the assertion. + """ + import queue + from types import SimpleNamespace + from freetoken.server import supervisor + + alive = {"scheduler": True, "tokenizer": True} + processes = [SimpleNamespace(name=name, is_alive=lambda name=name: alive[name]) + for name in alive] + acks = queue.Queue() + acks.put("scheduler ready") + acks.put("tokenizer ready") + handle = BackendHandle(ack_queue=acks, processes=processes, expected_acks=2) + stopping = False + seen = [] + + def next_poll(_): + nonlocal stopping + stopping = shutdown + alive["tokenizer"] = False + + monkeypatch.setattr(supervisor.time, "sleep", next_poll) + supervisor.run_backend_supervisor( + handle, LoadProgress(), on_ready=lambda: seen.append("ready"), + on_failure=lambda message: seen.append(message), + is_shutting_down=lambda: stopping, + ) + assert seen == (["ready"] if shutdown else ["ready", "backend worker tokenizer exited"]) From 4d99e7f3a0e0df86a8b989f71806135c6b7ee026 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 06:49:51 -0400 Subject: [PATCH 56/73] fix(server): drop /v1/completions from the --enable-metrics-report help, which it does not serve Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit c6b7ce7cc462c99c561a08071ec9cd308fdcc26e) --- python/freetoken/server/args.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index aaba2fc557..ce307bf375 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -610,7 +610,7 @@ def _infer_reasoning_parser(model_path: str) -> str | None: default=ServerArgs.enable_metrics_report, help=( "Serve a per-request `metrics` object next to usage on /v1/chat/completions, " - "/v1/completions, /v1/messages and /v1/responses: ttft_ms, prefill_time_ms and " + "/v1/messages and /v1/responses: ttft_ms, prefill_time_ms and " "prefill_tokens_per_second (the prefill span measured by the scheduler itself), " "decode_time_ms and decode_tokens_per_second, cached_prompt_tokens and " "total_time_ms. Streaming responses carry it on the same final chunk as usage, so " From 42a053b8c7df2418e6fef99ea3cb5ed1d92f327e Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 10:36:38 -0400 Subject: [PATCH 57/73] fix(server): keep developer ahead of system and reject images when merging system messages render_messages hoisted system ahead of a leading developer message, which the tokenizer then merged as "S\n\nD"; it now leaves developer requests to the tokenizer's per-template mapping. Both merge sites joined str() of an image part list into the prompt, data URL included; they now raise the template's 400 instead. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit e60b2c39d78626e8ecc5c21c7f5b1d77664698ba) --- python/freetoken/server/generation.py | 13 ++++++----- python/freetoken/tokenizer/tokenize.py | 13 +++++++++-- tests/server/test_mm_input.py | 32 ++++++++++++++++++++++++++ tests/tokenizer/test_tokenize.py | 3 +++ 4 files changed, 53 insertions(+), 8 deletions(-) diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index b86bcab0f0..0ed933feb8 100644 --- a/python/freetoken/server/generation.py +++ b/python/freetoken/server/generation.py @@ -25,7 +25,7 @@ from freetoken.core import SamplingParams from freetoken.message import TokenizeMsg from freetoken.mm.media import collect_image_refs, fetch_image_bytes, image_reject_reason -from freetoken.tokenizer.tokenize import resolve_thinking_mode +from freetoken.tokenizer.tokenize import join_system_contents, resolve_thinking_mode from freetoken.tokenizer.inline_system import InlineSystemError try: @@ -525,6 +525,11 @@ def render_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: hoist system messages to the front and merge multiples into one to satisfy that constraint.""" rendered = [_render_message(m) for m in messages] + if any(m.get("role") == "developer" for m in rendered): + # The tokenizer maps developer per template and then merges in order + # (_map_developer_role); hoisting system alone here would put it ahead of + # a developer message that came first. + return rendered system_msgs = [m for m in rendered if m.get("role") == "system"] if system_msgs and rendered[0].get("role") != "system": non_system = [m for m in rendered if m.get("role") != "system"] @@ -533,12 +538,8 @@ def render_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: # system messages that are not at loop.first, i.e. after the first). system_msgs = [m for m in rendered if m.get("role") == "system"] if len(system_msgs) > 1: - sep = chr(10) + chr(10) - merged_content = sep.join( - str(m.get("content") or "") for m in system_msgs if m.get("content") - ) non_system = [m for m in rendered if m.get("role") != "system"] - rendered = [{"role": "system", "content": merged_content}] + non_system + rendered = [{"role": "system", "content": join_system_contents(system_msgs)}] + non_system return rendered diff --git a/python/freetoken/tokenizer/tokenize.py b/python/freetoken/tokenizer/tokenize.py index 9648974dab..31b74bdac2 100644 --- a/python/freetoken/tokenizer/tokenize.py +++ b/python/freetoken/tokenizer/tokenize.py @@ -57,6 +57,16 @@ def resolve_thinking_mode(chat_template_kwargs: dict[str, Any] | None, tools: An _EFFORT_PROBE_MESSAGES = [{"role": "user", "content": "ping"}] +def join_system_contents(messages: list[dict[str, Any]]) -> str: + """The contents of several system/developer messages as one system prompt, in order. + Image content (a part list left by text flattening) cannot be joined into text: raise + the error templates give for an image in the system turn.""" + contents = [m.get("content") for m in messages if m.get("content")] + if any(not isinstance(c, str) for c in contents): + raise ValueError("System message cannot contain images") + return "\n\n".join(contents) + + def _map_developer_role( messages: list[dict[str, Any]], chat_template: str | None ) -> list[dict[str, Any]]: @@ -74,8 +84,7 @@ def _map_developer_role( if len(system) == 1 and mapped[0] is system[0]: return mapped rest = [m for m in mapped if m.get("role") != "system"] - content = "\n\n".join(str(m.get("content") or "") for m in system if m.get("content")) - return [{"role": "system", "content": content}] + rest + return [{"role": "system", "content": join_system_contents(system)}] + rest class TokenizeManager: diff --git a/tests/server/test_mm_input.py b/tests/server/test_mm_input.py index 5ac5834bf5..a694a50a84 100644 --- a/tests/server/test_mm_input.py +++ b/tests/server/test_mm_input.py @@ -173,3 +173,35 @@ def test_render_messages_no_system_unchanged(): {"role": "user", "content": "Bye"}, ]) assert [m["role"] for m in msgs] == ["user", "assistant", "user"] + + +def test_render_messages_keeps_a_leading_developer_message_ahead_of_system(): + """[developer, system, user]: the tokenizer maps developer -> system and merges in order, so + render_messages must not hoist system ahead of it first. A template that knows the + developer role gets both messages as sent.""" + from freetoken.tokenizer.tokenize import _map_developer_role + + msgs = [ + {"role": "developer", "content": "D"}, + {"role": "system", "content": "S"}, + {"role": "user", "content": "hi"}, + ] + rendered = render_messages(msgs) + assert [m["role"] for m in rendered] == ["developer", "system", "user"] + mapped = _map_developer_role(rendered, "{% if message.role == 'system' %}") + assert [m["role"] for m in mapped] == ["system", "user"] and mapped[0]["content"] == "D\n\nS" + assert _map_developer_role(rendered, "{% if message.role == 'developer' %}") is rendered + + +def test_render_messages_rejects_an_image_in_a_merged_system_message(): + """Merging system messages joins text; an image part must not be pasted into the prompt as + the repr of its part list (with the whole data URL).""" + with pytest.raises(ValueError, match="System message cannot contain images"): + render_messages([ + {"role": "system", "content": "A"}, + {"role": "system", "content": [ + {"type": "text", "text": "B"}, + {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{PNG}"}}, + ]}, + {"role": "user", "content": "hi"}, + ]) diff --git a/tests/tokenizer/test_tokenize.py b/tests/tokenizer/test_tokenize.py index 390fcd4cce..9e441886b2 100644 --- a/tests/tokenizer/test_tokenize.py +++ b/tests/tokenizer/test_tokenize.py @@ -845,3 +845,6 @@ def test_mapped_developer_messages_are_merged_into_one_leading_system_message(): assert mapped[0]["content"] == "be brief\n\nanswer in German" late = [{"role": "user", "content": "hi"}, {"role": "developer", "content": "be brief"}] assert [m["role"] for m in _map_developer_role(late, template)] == ["system", "user"] + image = [{"type": "text", "text": "B"}, {"type": "image", "freetoken_ref": {"kind": "url", "data": "x"}}] + with pytest.raises(ValueError, match="System message cannot contain images"): + _map_developer_role([{"role": "developer", "content": "A"}, {"role": "system", "content": image}], template) From ce31cd4edb3457eda3f87d52f278776f36e4275a Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 1 Oct 2026 23:24:07 -0400 Subject: [PATCH 58/73] test(server): ask for n=0 in the invalid-request sdk test, since n=2 is served here Assisted-by: Claude Opus 5.5 --- tests/server/test_protocol_compat.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/server/test_protocol_compat.py b/tests/server/test_protocol_compat.py index 8dc22908d6..d91582e423 100644 --- a/tests/server/test_protocol_compat.py +++ b/tests/server/test_protocol_compat.py @@ -202,7 +202,7 @@ async def run(): async with client(state, vendor) as sdk: with pytest.raises(error) as caught: if protocol == "chat": - await sdk.chat.completions.create(model="test", messages=[{"role": "user", "content": "Hi"}], n=2) + await sdk.chat.completions.create(model="test", messages=[{"role": "user", "content": "Hi"}], n=0) elif protocol == "responses": await sdk.responses.create(model="test", input="Hi", background=True) else: From 14d1b92308d19cea511e3642b8812907f21b3509 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 05:40:55 -0400 Subject: [PATCH 59/73] fix(moe): keep unpinned layers off the small-prefill path and take it in every layer An unpinned layer's copy_missing is the whole-layer pageable copy, which cannot serve the LRU slot remap, so it keeps streaming. The per-layer fallback when a layer routes to more than half the cache is gone: the engine's num_experts slot floor already keeps every chunk resident, and a uniform path cannot leave a later layer in an overlap buffer that was never begun. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 37c2479212b8b7028827218b23d1cee68e2c274c) --- python/freetoken/layers/moe.py | 44 +++++++++--------- tests/moe/test_small_prefill_on_demand.py | 54 +++++++++++++++++++++++ 2 files changed, 78 insertions(+), 20 deletions(-) diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index 56dcd712ab..71858b54d0 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -426,22 +426,26 @@ def _prefill_routed( pass through unmapped.""" cache = self.offload_cache assert cache is not None - if 0 < hidden_states.shape[0] <= _SMALL_PREFILL_TOKENS and cache.decode_target in ("gpu", "hybrid"): + if ( + 0 < hidden_states.shape[0] <= _SMALL_PREFILL_TOKENS + and cache.decode_target in ("gpu", "hybrid") + and not cache.is_unpinned_layer(self.layer_id) + ): # A short extension (an agent turn over a cached prefix) touches a fraction of each - # layer's experts; fetching those beats streaming all of them (~68 GB here). The same - # token count holds in every layer, so the overlap double buffer is never begun. - slots = self._ensure_unique(cache, topk_ids) - if slots is not None: - return self._expert_gemm( - cache, - hidden_states, - topk_weights, - slots, - views=cache.bank_views(), - n=None, - alphas=cache.alphas_for_slots(self.layer_id), - is_prefill=False, - ) + # layer's experts; fetching those beats streaming all of them (~68 GB here). An + # unpinned layer's copy_missing presumes position == expert id, so it streams below; + # unpinned layers only exist with the prefill overlap off, so no layer can enter + # the overlap double buffer after an earlier layer skipped begin_prefill. + return self._expert_gemm( + cache, + hidden_states, + topk_weights, + self._ensure_unique(cache, topk_ids), + views=cache.bank_views(), + n=None, + alphas=cache.alphas_for_slots(self.layer_id), + is_prefill=False, + ) if cache.prefill_overlap: views = self._wait_prefill_overlap(cache) out = self._expert_gemm( @@ -469,9 +473,8 @@ def _prefill_routed( is_prefill=True, ) - def _ensure_unique(self, cache: OffloadMoeCache, topk_ids: torch.Tensor) -> torch.Tensor | None: - """Make this layer's routed experts resident and return ``topk_ids`` mapped to slots, or - None (take the whole-layer path) when they cannot all be resident at once. + def _ensure_unique(self, cache: OffloadMoeCache, topk_ids: torch.Tensor) -> torch.Tensor: + """Make this layer's routed experts resident and return ``topk_ids`` mapped to slots. The LRU kernel is sized for decode: its block is next_pow2(ids) x next_pow2(slots), so a prefill's tokens x top_k ids overflow Triton's element limit, and even a block under it @@ -479,8 +482,9 @@ def _ensure_unique(self, cache: OffloadMoeCache, topk_ids: torch.Tensor) -> torc decode-width chunks; each chunk's misses are copied before the next ensure reuses the plan.""" uniq, inverse = torch.unique(topk_ids.reshape(-1), return_inverse=True) n = uniq.numel() - if n > cache.cache_size // 2: # leave room so a chunk never evicts the one before it - return None + # Each ensure stamps a newer LRU step, so a later chunk evicts the earlier chunks' slots + # only once every other slot is gone; the engine's num_experts slot floor rules that out. + assert n <= cache.cache_size, f"{n} routed experts exceed the {cache.cache_size}-slot cache" uniq = uniq.to(torch.int32) # scored eviction ranks an earlier chunk's fresh installs coldest, so pin them explicitly pin_since = cache.step.clone() if cache.cache_policy_id else None diff --git a/tests/moe/test_small_prefill_on_demand.py b/tests/moe/test_small_prefill_on_demand.py index 70cec00019..717f835d9c 100644 --- a/tests/moe/test_small_prefill_on_demand.py +++ b/tests/moe/test_small_prefill_on_demand.py @@ -56,6 +56,60 @@ def fake_gemm(cache_, hs, w, ids, *, views, n, alphas, is_prefill): assert "ensure" not in calls and "begin_prefill" not in calls[1:] and got["n"] == layer.num_experts +def _record_paths(monkeypatch, layer, cache): + """Stub the cache movement and the GEMM; log each call and what the GEMM was handed.""" + calls, got = [], {} + + def ensure(layer_id, ids): + calls.append(("ensure", ids.tolist())) + ids.add_(100) # the LRU kernel rewrites its query to slot ids in place + + monkeypatch.setattr(cache, "ensure_experts", ensure) + monkeypatch.setattr(cache, "materialize_layer", lambda layer_id: calls.append(("materialize",))) + monkeypatch.setattr(cache, "copy_missing", lambda: calls.append(("copy",))) + + def fake_gemm(cache_, hs, w, ids, *, views, n, alphas, is_prefill): + got.update(ids=ids.tolist(), n=n, is_prefill=is_prefill) + return hs + + monkeypatch.setattr(layer, "_expert_gemm", fake_gemm) + return calls, got + + +def test_small_prefill_leaves_an_unpinned_layer_on_the_whole_layer_path(monkeypatch): + # an unpinned (locked) layer's copy_missing is the whole-layer pageable copy, which presumes + # position == expert id and cannot serve an LRU slot remap + layer, cache = _layer_and_cache(torch.device("cpu")) + monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", 8) + cache._unpinned_layers = frozenset({layer.layer_id}) + calls, got = _record_paths(monkeypatch, layer, cache) + layer._prefill_routed(torch.randn(3, 16), torch.full((3, 2), 0.5), torch.zeros(3, 2, dtype=torch.int32)) + assert calls == [("materialize",), ("copy",)] and got["n"] == layer.num_experts + + +def test_small_prefill_path_does_not_depend_on_how_many_experts_a_layer_routes_to(monkeypatch): + # 6 distinct experts in an 8-slot cache: every layer of a short prefill still takes the + # on-demand path, so no layer can fall back into an overlap buffer that was never begun + layer, cache = _layer_and_cache(torch.device("cpu")) + monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", 8) + monkeypatch.setattr(moe_mod, "_ENSURE_CHUNK_IDS", 4) + calls, got = _record_paths(monkeypatch, layer, cache) + ids = torch.tensor([[5, 0], [1, 7], [2, 5], [4, 0]], dtype=torch.int32) + layer._prefill_routed(torch.randn(4, 16), torch.full((4, 2), 0.5), ids.clone()) + assert calls == [("ensure", [0, 1, 2, 4]), ("copy",), ("ensure", [5, 7]), ("copy",)] + assert got == {"ids": (ids + 100).tolist(), "n": None, "is_prefill": False} + + +def test_small_prefill_refuses_more_experts_than_the_cache_holds(monkeypatch): + layer, cache = _layer_and_cache(torch.device("cpu")) + monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", 8) + _record_paths(monkeypatch, layer, cache) + cache.cache_size = 4 # below the engine's num_experts floor + ids = torch.tensor([[5, 0], [1, 7], [2, 5], [4, 0]], dtype=torch.int32) + with pytest.raises(AssertionError, match="exceed"): + layer._prefill_routed(torch.randn(4, 16), torch.full((4, 2), 0.5), ids) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") def test_small_prefill_matches_the_whole_layer_path(monkeypatch): from freetoken.core import get_global_ctx # noqa: F401 (layer reads the ctx only in forward()) From b0c020dd8f039a7885c6e8e0ea7cc20b4d9368da Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 1 Oct 2026 23:24:22 -0400 Subject: [PATCH 60/73] test(moe): let the small-prefill ensure fake take the eviction policy kwargs Assisted-by: Claude Opus 5.5 --- tests/moe/test_small_prefill_on_demand.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/moe/test_small_prefill_on_demand.py b/tests/moe/test_small_prefill_on_demand.py index 717f835d9c..8b7b9d802c 100644 --- a/tests/moe/test_small_prefill_on_demand.py +++ b/tests/moe/test_small_prefill_on_demand.py @@ -60,7 +60,7 @@ def _record_paths(monkeypatch, layer, cache): """Stub the cache movement and the GEMM; log each call and what the GEMM was handed.""" calls, got = [], {} - def ensure(layer_id, ids): + def ensure(layer_id, ids, **_): calls.append(("ensure", ids.tolist())) ids.add_(100) # the LRU kernel rewrites its query to slot ids in place From 71a92dbd4181f00311e42965534a904a202f59e4 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 4 Sep 2026 17:31:34 -0400 Subject: [PATCH 61/73] test(qwen4_exp): compare chunked prefill to one-shot at bf16 tolerance `torch.equal` fails for all three cut points on torch 2.11.0+cu130 / flashinfer 0.6.18 / triton 3.6.0 (sm_89): cuBLAS picks a different kernel for M=5000 and M=5000-cut, so the indexer's q and k round differently, a few rows flip one block at the top-k margin, and the layer output differs by up to 2.6e-3 (a few bf16 ulps). Both paths are deterministic run to run. Use the tolerance of the neighbouring oracle tests. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt (cherry picked from commit 3849e3caeebe4f1376708fe677cc0ecb106fb78c) Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit e700e8fb7339eb958507401c39d81b2623f56ca7) --- tests/models/qwen4_exp/test_qsa_backend.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/models/qwen4_exp/test_qsa_backend.py b/tests/models/qwen4_exp/test_qsa_backend.py index 9302fad172..3827b62416 100644 --- a/tests/models/qwen4_exp/test_qsa_backend.py +++ b/tests/models/qwen4_exp/test_qsa_backend.py @@ -4,7 +4,8 @@ tokens every complete block is selected, so QSA IS dense attention: the selection must be exactly the causal prefix and the layer output must match ``TorchDenseQSAReference`` (fp32) and a flashinfer dense run over the same pool; -(b) chunked prefill at unaligned cut points equals one-shot prefill (the dual-source compress); +(b) chunked prefill at unaligned cut points matches one-shot prefill to bf16 tolerance (the + dual-source compress); (c) a captured decode replay equals the eager decode step; (d) an fp8 KV pool (``--kv-cache-dtype fp8``) keeps block selection bit-identical to the 16-bit run -- only the selected K/V rows are read back as e4m3 codes -- and the layer @@ -136,7 +137,11 @@ def test_flashinfer_dense_matches_the_sparse_path(): @requires_cuda @pytest.mark.parametrize("cut", [1001, 4096, 4097], ids=["unaligned", "page-boundary", "boundary+1"]) def test_chunked_prefill_matches_one_shot(cut: int): - """Cut points that are not multiples of index_ratio exercise the dual-source compress.""" + """Cut points that are not multiples of index_ratio exercise the dual-source compress. + + Not bit for bit: the projections see a different M (5000 vs 5000 - cut rows) and cuBLAS + may pick a different kernel per shape, so the indexer's q and k round differently and a + few rows flip one block at the top-k margin. Same tolerance as the oracle tests above.""" config = parsed_config() fixture = Fixture(config, num_pages=512) attn = fixture.layer(QSA_LAYER) @@ -148,7 +153,7 @@ def test_chunked_prefill_matches_one_shot(cut: int): attn.forward(x[:cut], fixture.batch([head], "prefill")) tail = fixture.req(1, cut, length) got = attn.forward(x[cut:], fixture.batch([tail], "prefill")) - assert torch.equal(got, one_shot[cut:]) + torch.testing.assert_close(got.float(), one_shot[cut:].float(), rtol=2e-2, atol=2e-2) @requires_cuda From f8e43969a699ca7206d0f3ab7533582a5c20c01b Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 09:23:20 -0400 Subject: [PATCH 62/73] test(qwen4_exp): compare the chunked prefill's compressed slab rows to the one-shot run At bf16 output tolerance a broken dual-source compress only swaps a block in a few rows and passes; the slab keys are unit-scale, so a wrong straddling-group key stands out. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit cedf179739d2f17e8dc8ee27152f35c5584c24c4) --- tests/models/qwen4_exp/test_qsa_backend.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/tests/models/qwen4_exp/test_qsa_backend.py b/tests/models/qwen4_exp/test_qsa_backend.py index 3827b62416..8e836f5632 100644 --- a/tests/models/qwen4_exp/test_qsa_backend.py +++ b/tests/models/qwen4_exp/test_qsa_backend.py @@ -155,6 +155,18 @@ def test_chunked_prefill_matches_one_shot(cut: int): got = attn.forward(x[cut:], fixture.batch([tail], "prefill")) torch.testing.assert_close(got.float(), one_shot[cut:].float(), rtol=2e-2, atol=2e-2) + # A wrong key for the group straddling the cut only swaps one block in a few rows, which + # the output tolerance hides; the unit-scale slab rows show it as an O(1) error. + ratio = config.qwen4_args.index_ratio + slab = fixture.pool.cmp_k_cache(0) + group_starts = slice(0, length - length % ratio, ratio) + one_shot_rows, chunked_rows = ( + (fixture.page_table[table_idx, group_starts] // ratio).long() for table_idx in (0, 1) + ) + torch.testing.assert_close( + slab[chunked_rows].float(), slab[one_shot_rows].float(), rtol=2e-2, atol=2e-2 + ) + @requires_cuda def test_decode_graph_replay_matches_eager(): From e06b2b24048c8dd183cfa3a67681223d1f9f12aa Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 06:40:34 -0400 Subject: [PATCH 63/73] fix(kernel): keep the triton sampling softmax finite over a fully masked prefix min_p (#393) writes -inf over most of the vocabulary; a column chunk or leading block that is -inf throughout computed exp(-inf - -inf) = NaN and poisoned the whole row. Pascal samples through this softmax (no flashinfer below sm_75), so min_p > 0 drew uniform noise. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit dbe82ad95f97dc0d88e647e87a7cb90f8d7783cf) --- python/freetoken/kernel/triton/sampling.py | 8 ++++++-- tests/engine/test_logits_processors.py | 15 ++++++++++++++ tests/kernels/test_triton_sampling.py | 23 ++++++++++++++++++++++ 3 files changed, 44 insertions(+), 2 deletions(-) create mode 100644 tests/kernels/test_triton_sampling.py diff --git a/python/freetoken/kernel/triton/sampling.py b/python/freetoken/kernel/triton/sampling.py index a970714406..869046aa1c 100644 --- a/python/freetoken/kernel/triton/sampling.py +++ b/python/freetoken/kernel/triton/sampling.py @@ -92,7 +92,10 @@ def _sm_partial( x = tl.load(logits_ptr + base + offs, mask=mask, other=-float("inf")).to(tl.float32) * inv_t blk_max = tl.max(x, 0) new_m = tl.maximum(m, blk_max) - d = d * tl.exp(m - new_m) + tl.sum(tl.exp(x - new_m), 0) + # A fully masked prefix (min_p / min_tokens write -inf) leaves new_m at -inf, and + # exp(-inf - -inf) is NaN; shift by 0 there instead, every term is exp(-inf) = 0. + safe_m = tl.where(new_m == -float("inf"), 0.0, new_m) + d = d * tl.exp(m - safe_m) + tl.sum(tl.exp(x - safe_m), 0) m = new_m tl.store(pm_ptr + pid, m) tl.store(pl_ptr + pid, d) @@ -117,7 +120,8 @@ def _sm_finalize( pm = tl.load(pm_ptr + row * G + goff, mask=gmask, other=-float("inf")) pl = tl.load(pl_ptr + row * G + goff, mask=gmask, other=0.0) gm = tl.max(pm, 0) - gl = tl.sum(pl * tl.exp(pm - gm), 0) + # a chunk that is -inf throughout holds no mass; keep its NaN-prone exp out of the sum + gl = tl.sum(tl.where(pm == -float("inf"), 0.0, pl * tl.exp(pm - gm)), 0) inv_gl = 1.0 / gl base = row * row_stride diff --git a/tests/engine/test_logits_processors.py b/tests/engine/test_logits_processors.py index abd2e1f939..a215961ef7 100644 --- a/tests/engine/test_logits_processors.py +++ b/tests/engine/test_logits_processors.py @@ -190,3 +190,18 @@ def test_repetition_penalty_ignores_multimodal_placeholder_ids_in_the_prompt(): args = sampler.prepare(batch) out = apply_logits_processors(torch.ones(1, V), args.plan, V) assert out[0, 1] == 0.5 and out[0, 2] == 0.5 and out[0, 4] == 1.0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_min_p_with_temperature_draws_only_the_kept_tokens(): + # min_p leaves -inf over nearly the whole row; the sampling softmax must not turn that + # into NaN (which the Pascal fallback then samples as a uniform row) + V = 32768 + sampler = Sampler(torch.device("cuda"), V) + req = _req([1, 2], [], 4, temperature=0.8, min_p=0.1) + args = sampler.prepare(SimpleNamespace(reqs=[req])) + logits = torch.full((1, V), -10.0, device="cuda") + kept = [100, 20000, 31000] + logits[0, kept] = torch.tensor([5.0, 4.5, 4.0], device="cuda") + for _ in range(32): + assert sampler.sample(logits, args).item() in kept diff --git a/tests/kernels/test_triton_sampling.py b/tests/kernels/test_triton_sampling.py new file mode 100644 index 0000000000..0766ea9343 --- /dev/null +++ b/tests/kernels/test_triton_sampling.py @@ -0,0 +1,23 @@ +"""The Triton sampling softmax, which samples when flashinfer is not installed.""" + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + + +def test_softmax_survives_a_fully_masked_prefix(): + """min_p / min_tokens write -inf over most of the vocabulary. A column chunk or leading + block that is -inf throughout gave exp(-inf - -inf) = NaN and a NaN row.""" + from freetoken.kernel.triton.sampling import softmax + + V = 32768 # several column chunks per row + torch.manual_seed(0) + logits = torch.full((2, V), float("-inf"), device="cuda") + logits[0, [20000, 30000, 32000]] = torch.tensor([1.0, 2.0, 0.5], device="cuda") + logits[1] = torch.randn(V, device="cuda") + logits[1, :8192] = float("-inf") + temps = torch.tensor([0.7, 1.3], device="cuda") + probs = softmax(logits, temps) + assert not probs.isnan().any() + torch.testing.assert_close(probs, torch.softmax(logits / temps[:, None], dim=-1)) From 588677eab7966c651f8e382a74305b3521896dff Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 05:20:22 -0400 Subject: [PATCH 64/73] feat(engine): keep the token embedding in pinned host ram (--embed-weights host) The embedding is a row lookup, so a single-GPU engine can leave the table in page-locked host RAM and gather each step's rows over UVA with the PLE row kernel (5 KiB per decode token, graph-safe): 1.2 GiB of VRAM back for Qwen3.8-Flash-Next. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 9f2b7a2e33424e2c0e06b240f1fb7a408c850212) --- python/freetoken/engine/config.py | 2 + python/freetoken/engine/engine.py | 35 +++++++++++++++++ python/freetoken/kernel/triton/ple.py | 8 +++- python/freetoken/layers/embedding.py | 31 ++++++++++++--- python/freetoken/server/args.py | 10 +++++ tests/kernels/test_host_embedding.py | 55 +++++++++++++++++++++++++++ 6 files changed, 134 insertions(+), 7 deletions(-) create mode 100644 tests/kernels/test_host_embedding.py diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 81cd2dc4d8..a54dd105b6 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -38,6 +38,8 @@ class EngineConfig: # DeepSeek-V4.1 Decoder SWA Bounded Replay: "bounded" (default) runs the decoder layers on each # prompt's last window; "exact" runs them on every token. swa_decoder_replay: str = "bounded" + # --embed-weights: "host" keeps the token embedding in pinned host RAM, gathered over PCIe. + embed_weights: str = "gpu" # Expert-bank host load (--expert-load): auto|serial|parallel. "auto" reads scattered # experts in parallel but falls back to serial when free RAM can't cover the banks + the # parallel reader's extra (non-reclaimable) whole-shard buffer; "serial" forces the diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 9ccc44763c..c7a4fefcca 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -85,6 +85,39 @@ def _sgl_flash_attn_available() -> bool: return True +def _place_embeddings_on_host(model) -> None: + """--embed-weights host: every input embedding table moves to pinned host RAM, except one an lm_head is tied to.""" + from freetoken.layers.base import BaseOP + from freetoken.layers.embedding import ParallelLMHead, VocabParallelEmbedding + + found: list[VocabParallelEmbedding] = [] + tied: set[int] = set() + stack, seen = [model], set() + while stack: + op = stack.pop() + if id(op) in seen: + continue + seen.add(id(op)) + if isinstance(op, ParallelLMHead): + if op.tied_embedding is not None: + tied.add(id(op.tied_embedding)) + continue + if isinstance(op, VocabParallelEmbedding): + found.append(op) + continue + for value in vars(op).values(): + items = value if isinstance(value, (list, tuple)) else (value,) + stack.extend(v for v in items if isinstance(v, BaseOP)) + moved = [e for e in found if id(e) not in tied] + if not moved: + raise ValueError("--embed-weights host: the model has no untied input embedding to move") + for emb in moved: + emb.place_on_host() + torch.cuda.empty_cache() + size = sum(e.weight.numel() * e.weight.element_size() for e in moved) + logger.info_rank0(f"Token embedding in pinned host RAM ({mem_GB(size)}), gathered over PCIe") + + def _startup_kv_budget(memory_ratio: float, init_free_memory: int, new_free_memory: int) -> int: """Bytes available to the KV pool at startup: ratio-scaled pre-load free memory minus what the resident model consumed. Kept as a pure function so the composition with the @@ -457,6 +490,8 @@ def __init__(self, config: EngineConfig): ) # before the residency snapshot, so streamed blocks are not charged as resident weights self.model.place_encoder_weights(config.mm.encoder_weights) + if config.embed_weights == "host": + _place_embeddings_on_host(self.model) post_weights_free = self._sync_get_memory()[0] self._weights_bytes = self._baseline_free - post_weights_free # Pool-budget baseline for the desktop cache sliders: free VRAM after the weights are diff --git a/python/freetoken/kernel/triton/ple.py b/python/freetoken/kernel/triton/ple.py index 2314e94e72..cded101a2e 100644 --- a/python/freetoken/kernel/triton/ple.py +++ b/python/freetoken/kernel/triton/ple.py @@ -21,6 +21,7 @@ # Latency-bound over PCIe, so keep the block small and let many of them be in flight. _NUM_WARPS = 1 +_TL_TABLE_DTYPES = {torch.bfloat16: tl.bfloat16, torch.float16: tl.float16} @triton.jit @@ -33,6 +34,7 @@ def _ple_gather_kernel( EMB_DIM: tl.constexpr, IS_FP8: tl.constexpr, BLOCK_D: tl.constexpr, + TABLE_DTYPE: tl.constexpr, ): row = tl.program_id(0) idx = tl.load(ids_ptr + row).to(tl.int64) @@ -50,7 +52,7 @@ def _ple_gather_kernel( base = table_ptr.to(tl.int64).to(tl.pointer_type(tl.uint8)) values = e4m3_u8_to_f32(tl.load(base + idx * EMB_DIM + offsets, mask=mask, other=0)) else: - base = table_ptr.to(tl.int64).to(tl.pointer_type(tl.bfloat16)) + base = table_ptr.to(tl.int64).to(tl.pointer_type(TABLE_DTYPE)) values = tl.load(base + idx * EMB_DIM + offsets, mask=mask, other=0.0).to(tl.float32) values = tl.where(in_range, values * scale, 0.0) tl.store( @@ -68,11 +70,12 @@ def ple_gather_rows( out: torch.Tensor, scale: float = 1.0, is_fp8: bool = True, + table_dtype: torch.dtype = torch.bfloat16, ) -> torch.Tensor: """Gather ``row_ids`` from the host-resident table at ``table_ptr`` into ``out``. ``row_ids`` is a flat device int tensor; ``out`` is ``[row_ids.numel(), embed_dim]`` - bf16 on the same device. ``table_ptr`` is the address the GPU must dereference + on the same device. A non-fp8 table holds ``table_dtype`` elements. ``table_ptr`` is the address the GPU must dereference (``kernel/pinned.device_ptr``), not necessarily the host ``data_ptr``. """ n = row_ids.numel() @@ -87,6 +90,7 @@ def ple_gather_rows( EMB_DIM=embed_dim, IS_FP8=is_fp8, BLOCK_D=triton.next_power_of_2(embed_dim), + TABLE_DTYPE=_TL_TABLE_DTYPES[table_dtype], num_warps=_NUM_WARPS, ) return out diff --git a/python/freetoken/layers/embedding.py b/python/freetoken/layers/embedding.py index 6939156cd6..fc5deaac58 100644 --- a/python/freetoken/layers/embedding.py +++ b/python/freetoken/layers/embedding.py @@ -36,16 +36,37 @@ def __init__( self._embed_scale = embed_scale self._embed_scale_t: torch.Tensor | None = None self._comm = DistributedCommunicator() + self._host_ptr: int | None = None + + def place_on_host(self) -> None: + """Move the table into pinned host RAM; lookups then read their rows over PCIe (UVA).""" + from freetoken.kernel.pinned import alloc_pinned_tensor, device_ptr + + if self.tp_size > 1: + raise ValueError("a host-resident embedding serves a single GPU only") + host = alloc_pinned_tensor(*self.weight.shape, dtype=self.weight.dtype) + host.copy_(self.weight) + self.weight = host + self._host_ptr = device_ptr(host) @nvtx_annotate("Embedding") def forward(self, x: torch.Tensor) -> torch.Tensor: from freetoken.kernel import indexing - y = indexing( - weights=self.weight, - indices=x, - vocab_range=self.vocab_range if self.tp_size > 1 else None, - ) + if self._host_ptr is not None: + from freetoken.kernel.triton.ple import ple_gather_rows + + rows, dim = self.weight.shape + y = torch.empty(x.numel(), dim, dtype=self.weight.dtype, device=x.device) + ple_gather_rows(self._host_ptr, rows, dim, x.reshape(-1), y, is_fp8=False, + table_dtype=self.weight.dtype) + y = y.view(*x.shape, dim) + else: + y = indexing( + weights=self.weight, + indices=x, + vocab_range=self.vocab_range if self.tp_size > 1 else None, + ) if self.tp_size > 1: y = self._comm.all_reduce(y) diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index ce307bf375..3e693bbe78 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -748,6 +748,16 @@ def _infer_reasoning_parser(model_path: str) -> str | None: ), ) + parser.add_argument( + "--embed-weights", + default=ServerArgs.embed_weights, + choices=["gpu", "host"], + help=( + "Where the token embedding table lives. 'host' keeps it in pinned host RAM and gathers " + "the rows each step needs over PCIe, freeing its VRAM (single GPU only)." + ), + ) + parser.add_argument( "--nvfp4-backend", action=_DeprecatedAlias, diff --git a/tests/kernels/test_host_embedding.py b/tests/kernels/test_host_embedding.py new file mode 100644 index 0000000000..76c8c612c3 --- /dev/null +++ b/tests/kernels/test_host_embedding.py @@ -0,0 +1,55 @@ +"""--embed-weights host: the UVA row gather from pinned host RAM must equal the GPU table lookup.""" + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + + +@pytest.fixture(autouse=True) +def _tp(): + 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) + + +def _embedding(dtype: torch.dtype, rows: int = 1000, dim: int = 2560): + from freetoken.layers.embedding import VocabParallelEmbedding + + torch.manual_seed(0) + emb = VocabParallelEmbedding(rows, dim) + emb.weight = torch.randn(rows, dim, device="cuda", dtype=dtype) + return emb + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_host_rows_equal_the_gpu_lookup(dtype): + emb = _embedding(dtype) + ids = torch.randint(0, 1000, (37,), device="cuda", dtype=torch.int32) + want = emb.forward(ids) + emb.place_on_host() + assert not emb.weight.is_cuda and emb.weight.is_pinned() + got = emb.forward(ids) + assert got.is_cuda and got.dtype is dtype + assert torch.equal(got, want) + + +def test_host_lookup_replays_in_a_cuda_graph_with_new_ids(): + emb = _embedding(torch.bfloat16) + table = emb.weight.clone() + emb.place_on_host() + ids = torch.zeros(8, device="cuda", dtype=torch.int64) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + emb.forward(ids) # compile outside the capture + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = emb.forward(ids) + for seed in range(3): + ids.copy_(torch.randint(0, 1000, (8,), generator=torch.Generator().manual_seed(seed))) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, table[ids]) From bf064fdcfd04f74e4c25608eddd7cdc7a3428afb Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 10:51:34 -0400 Subject: [PATCH 65/73] fix(layers): gather fp32 host embeddings and refuse other dtypes at placement Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit ccfb2b8148c8c108897f37b7f337e99cbb5a0134) --- python/freetoken/kernel/triton/ple.py | 2 +- python/freetoken/layers/embedding.py | 3 +++ tests/kernels/test_host_embedding.py | 17 +++++++++++++---- 3 files changed, 17 insertions(+), 5 deletions(-) diff --git a/python/freetoken/kernel/triton/ple.py b/python/freetoken/kernel/triton/ple.py index cded101a2e..8fa25acdfd 100644 --- a/python/freetoken/kernel/triton/ple.py +++ b/python/freetoken/kernel/triton/ple.py @@ -21,7 +21,7 @@ # Latency-bound over PCIe, so keep the block small and let many of them be in flight. _NUM_WARPS = 1 -_TL_TABLE_DTYPES = {torch.bfloat16: tl.bfloat16, torch.float16: tl.float16} +_TL_TABLE_DTYPES = {torch.bfloat16: tl.bfloat16, torch.float16: tl.float16, torch.float32: tl.float32} @triton.jit diff --git a/python/freetoken/layers/embedding.py b/python/freetoken/layers/embedding.py index fc5deaac58..860c220010 100644 --- a/python/freetoken/layers/embedding.py +++ b/python/freetoken/layers/embedding.py @@ -41,9 +41,12 @@ def __init__( def place_on_host(self) -> None: """Move the table into pinned host RAM; lookups then read their rows over PCIe (UVA).""" from freetoken.kernel.pinned import alloc_pinned_tensor, device_ptr + from freetoken.kernel.triton.ple import _TL_TABLE_DTYPES if self.tp_size > 1: raise ValueError("a host-resident embedding serves a single GPU only") + if self.weight.dtype not in _TL_TABLE_DTYPES: + raise ValueError(f"a host-resident embedding cannot gather a {self.weight.dtype} table") host = alloc_pinned_tensor(*self.weight.shape, dtype=self.weight.dtype) host.copy_(self.weight) self.weight = host diff --git a/tests/kernels/test_host_embedding.py b/tests/kernels/test_host_embedding.py index 76c8c612c3..866687a9f1 100644 --- a/tests/kernels/test_host_embedding.py +++ b/tests/kernels/test_host_embedding.py @@ -3,7 +3,7 @@ import pytest import torch -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +needs_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") @pytest.fixture(autouse=True) @@ -14,16 +14,17 @@ def _tp(): set_tp_info(rank=0, size=1) -def _embedding(dtype: torch.dtype, rows: int = 1000, dim: int = 2560): +def _embedding(dtype: torch.dtype, rows: int = 1000, dim: int = 2560, device: str = "cuda"): from freetoken.layers.embedding import VocabParallelEmbedding torch.manual_seed(0) emb = VocabParallelEmbedding(rows, dim) - emb.weight = torch.randn(rows, dim, device="cuda", dtype=dtype) + emb.weight = torch.randn(rows, dim, device=device, dtype=dtype) return emb -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@needs_cuda +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) def test_host_rows_equal_the_gpu_lookup(dtype): emb = _embedding(dtype) ids = torch.randint(0, 1000, (37,), device="cuda", dtype=torch.int32) @@ -35,6 +36,14 @@ def test_host_rows_equal_the_gpu_lookup(dtype): assert torch.equal(got, want) +def test_a_table_the_gather_cannot_read_is_refused_before_pinning(): + emb = _embedding(torch.float64, rows=4, dim=8, device="cpu") + with pytest.raises(ValueError, match="cannot gather a torch.float64 table"): + emb.place_on_host() + assert emb._host_ptr is None + + +@needs_cuda def test_host_lookup_replays_in_a_cuda_graph_with_new_ids(): emb = _embedding(torch.bfloat16) table = emb.weight.clone() From b7676a3e1102e4223f84527c27105fe1d85d2bbf Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 10:54:43 -0400 Subject: [PATCH 66/73] fix(kernel): compute the host gather output offset in int64 No regression test: overflowing it needs an output of 2^31+ elements (4 GiB+ in bf16). Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit cc5af7adc5d99c17c0edc3f7947040ac35a48445) --- python/freetoken/kernel/triton/ple.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/freetoken/kernel/triton/ple.py b/python/freetoken/kernel/triton/ple.py index 8fa25acdfd..466fce11a0 100644 --- a/python/freetoken/kernel/triton/ple.py +++ b/python/freetoken/kernel/triton/ple.py @@ -36,7 +36,8 @@ def _ple_gather_kernel( BLOCK_D: tl.constexpr, TABLE_DTYPE: tl.constexpr, ): - row = tl.program_id(0) + # the output offset row * EMB_DIM can pass 2^31 for a long prefill at a wide hidden size + row = tl.program_id(0).to(tl.int64) idx = tl.load(ids_ptr + row).to(tl.int64) in_range = (idx >= 0) & (idx < num_rows) idx = tl.where(in_range, idx, 0) From bc1d696b37ccd7d2c2f5adf6eba1536c5e417525 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 10:56:09 -0400 Subject: [PATCH 67/73] fix(engine): charge the host embedding's pinned bytes to the pin budget Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit a8b2f0a6924c9362b014c200dd702115b0c33d89) --- python/freetoken/engine/engine.py | 16 +++++++++------- tests/engine/test_cache_budget.py | 16 ++++++++++++++++ 2 files changed, 25 insertions(+), 7 deletions(-) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index c7a4fefcca..5a93e139f5 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -85,8 +85,10 @@ def _sgl_flash_attn_available() -> bool: return True -def _place_embeddings_on_host(model) -> None: - """--embed-weights host: every input embedding table moves to pinned host RAM, except one an lm_head is tied to.""" +def _place_embeddings_on_host(model) -> int: + """--embed-weights host: every input embedding table moves to pinned host RAM, except one an lm_head is tied to. + + Returns the pinned bytes.""" from freetoken.layers.base import BaseOP from freetoken.layers.embedding import ParallelLMHead, VocabParallelEmbedding @@ -116,6 +118,7 @@ def _place_embeddings_on_host(model) -> None: torch.cuda.empty_cache() size = sum(e.weight.numel() * e.weight.element_size() for e in moved) logger.info_rank0(f"Token embedding in pinned host RAM ({mem_GB(size)}), gathered over PCIe") + return size def _startup_kv_budget(memory_ratio: float, init_free_memory: int, new_free_memory: int) -> int: @@ -490,8 +493,7 @@ def __init__(self, config: EngineConfig): ) # before the residency snapshot, so streamed blocks are not charged as resident weights self.model.place_encoder_weights(config.mm.encoder_weights) - if config.embed_weights == "host": - _place_embeddings_on_host(self.model) + embed_host_bytes = _place_embeddings_on_host(self.model) if config.embed_weights == "host" else 0 post_weights_free = self._sync_get_memory()[0] self._weights_bytes = self._baseline_free - post_weights_free # Pool-budget baseline for the desktop cache sliders: free VRAM after the weights are @@ -505,11 +507,11 @@ def __init__(self, config: EngineConfig): self.cpu_moe_executor = None # Host-side auxiliary stores (qwen4_exp's pinned PLE table): after the weights so a # load failure is not masked, before the MoE offload cache so the bank residency - # planning sees the pin quota the table already spent. - self._host_tables_bytes = 0 + # planning sees the pin quota the table (and a host embedding) already spent. + self._host_tables_bytes = embed_host_bytes if hasattr(self.model, "load_host_tables"): with _weight_load_context(): - self._host_tables_bytes = int(self.model.load_host_tables(config) or 0) + self._host_tables_bytes += int(self.model.load_host_tables(config) or 0) if is_offload_moe_strategy(config.moe_strategy): self._init_offload_moe_cache(config) if hasattr(self.model, "prepare_for_runtime"): diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py index da242a2cc4..69ab3f2c40 100644 --- a/tests/engine/test_cache_budget.py +++ b/tests/engine/test_cache_budget.py @@ -621,3 +621,19 @@ def test_uncapped_platform_stays_uncapped(monkeypatch): if hasattr(os, "uname") and "microsoft" in os.uname().release.lower(): pytest.skip("WSL caps pinning") assert _pin_budget_bytes(reserved=2**30) is None + + +def test_host_embedding_placement_reports_the_bytes_it_pinned(monkeypatch): + # the engine charges these bytes to the pin budget before it plans the expert banks + from freetoken.distributed import set_tp_info, try_get_tp_info + from freetoken.engine.engine import _place_embeddings_on_host + from freetoken.layers.embedding import ParallelLMHead, VocabParallelEmbedding + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + monkeypatch.setattr(VocabParallelEmbedding, "place_on_host", lambda self: None) + embed, tied = VocabParallelEmbedding(100, 64), VocabParallelEmbedding(50, 64) + embed.weight = torch.empty(100, 64, dtype=torch.bfloat16) + head = ParallelLMHead(50, 64, tie_word_embeddings=True, tied_embedding=tied) + model = SimpleNamespace(embed_tokens=embed, per_layer_embed=tied, lm_head=head) + assert _place_embeddings_on_host(model) == 100 * 64 * 2 From 2b711346dcc506b2378684cb7250c72a889eba0d Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 10:57:23 -0400 Subject: [PATCH 68/73] fix(engine): refuse --embed-weights host for a tied embedding or tp > 1 at config time Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 579e03a7aae077ac209cbf22ed4de6cda5bb390f) --- python/freetoken/engine/config.py | 10 ++++++++ tests/engine/test_embed_weights.py | 39 ++++++++++++++++++++++++++++++ 2 files changed, 49 insertions(+) create mode 100644 tests/engine/test_embed_weights.py diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index a54dd105b6..5567164573 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -163,8 +163,18 @@ def model_config(self) -> ModelConfig: quant = checkpoint_quant_config(self.model_path, hf_config, spec) set_quant_config(quant) model_config = _load_attr(spec.module, spec.parse_config)(hf_config) + self._check_embed_weights(model_config) return replace(model_config, quant=quant) + def _check_embed_weights(self, model_config) -> None: + # the parsed flag, not the hf key: gemma4 ties by default and gguf infers it from the tensors + if self.embed_weights != "host": + return + if model_config.tie_word_embeddings: + raise ValueError("--embed-weights host needs an untied embedding (a tied lm_head reads the table on the GPU)") + if self.tp_info.size > 1: + raise ValueError("--embed-weights host serves a single GPU only (--tensor-parallel-size 1)") + @property def max_seq_len(self) -> int: if self.max_seq_len_override is not None: diff --git a/tests/engine/test_embed_weights.py b/tests/engine/test_embed_weights.py new file mode 100644 index 0000000000..b8cb0fdc26 --- /dev/null +++ b/tests/engine/test_embed_weights.py @@ -0,0 +1,39 @@ +"""--embed-weights host: a setup the host gather cannot serve is refused while the config is still a dataclass, not after the weight load.""" + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.distributed import DistributedInfo +from freetoken.engine.config import EngineConfig + + +def _config(monkeypatch, *, tie: bool, tp: int = 1, embed_weights: str = "host"): + import freetoken.engine.config as engine_config + + hf = SimpleNamespace( + architectures=["Qwen3ForCausalLM"], model_type="qwen3", num_hidden_layers=2, num_attention_heads=4, + num_key_value_heads=2, hidden_size=64, vocab_size=128, intermediate_size=128, rms_norm_eps=1e-6, + max_position_embeddings=1024, hidden_act="silu", tie_word_embeddings=tie, + ) + monkeypatch.setattr(engine_config, "cached_load_hf_config", lambda path: hf) + monkeypatch.setattr(engine_config, "checkpoint_quant_config", lambda *args: None) + return EngineConfig( + model_path="/fake", tp_info=DistributedInfo(rank=0, size=tp), dtype=torch.bfloat16, embed_weights=embed_weights + ) + + +def test_a_tied_embedding_is_refused(monkeypatch): + with pytest.raises(ValueError, match="needs an untied embedding"): + _config(monkeypatch, tie=True).model_config + + +def test_tensor_parallel_is_refused(monkeypatch): + with pytest.raises(ValueError, match="single GPU only"): + _config(monkeypatch, tie=False, tp=2).model_config + + +def test_an_untied_single_gpu_model_and_the_gpu_default_pass(monkeypatch): + assert not _config(monkeypatch, tie=False).model_config.tie_word_embeddings + assert _config(monkeypatch, tie=True, tp=2, embed_weights="gpu").model_config.tie_word_embeddings From ee3f94fb9fe001042b738869632ce7c92b6ef74b Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 1 Oct 2026 23:26:39 -0400 Subject: [PATCH 69/73] test(engine): fake cached_load_hf_config with the overrides argument #573 passes Assisted-by: Claude Opus 5.5 --- tests/engine/test_embed_weights.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/engine/test_embed_weights.py b/tests/engine/test_embed_weights.py index b8cb0fdc26..d3aa102f94 100644 --- a/tests/engine/test_embed_weights.py +++ b/tests/engine/test_embed_weights.py @@ -17,7 +17,7 @@ def _config(monkeypatch, *, tie: bool, tp: int = 1, embed_weights: str = "host") num_key_value_heads=2, hidden_size=64, vocab_size=128, intermediate_size=128, rms_norm_eps=1e-6, max_position_embeddings=1024, hidden_act="silu", tie_word_embeddings=tie, ) - monkeypatch.setattr(engine_config, "cached_load_hf_config", lambda path: hf) + monkeypatch.setattr(engine_config, "cached_load_hf_config", lambda path, *_: hf) monkeypatch.setattr(engine_config, "checkpoint_quant_config", lambda *args: None) return EngineConfig( model_path="/fake", tp_info=DistributedInfo(rank=0, size=tp), dtype=torch.bfloat16, embed_weights=embed_weights From 423f2053be86b2b193dc2592e11b2ea66dd52e31 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 1 Oct 2026 23:27:21 -0400 Subject: [PATCH 70/73] test(ple): drive overlapped graph fills against an in-order fake stream so a post-launch fence fails as a hang Upstream #588's _graph_consumed fence is left out of the merge: our graph fill runs on the filler thread before its own launch passes its WAIT, so that fence deadlocks the second step. Assisted-by: Claude Opus 5.5 --- .../qwen4_exp/test_ple_disk_no_graph_fence.py | 136 ++++++++++++++++++ 1 file changed, 136 insertions(+) create mode 100644 tests/models/qwen4_exp/test_ple_disk_no_graph_fence.py diff --git a/tests/models/qwen4_exp/test_ple_disk_no_graph_fence.py b/tests/models/qwen4_exp/test_ple_disk_no_graph_fence.py new file mode 100644 index 0000000000..cb9bc32356 --- /dev/null +++ b/tests/models/qwen4_exp/test_ple_disk_no_graph_fence.py @@ -0,0 +1,136 @@ +"""The wait-sync graph fill must never wait on an event recorded after its own launch. + +Upstream #588 fences ``_graph_pinned`` with ``_graph_consumed``, recorded once the launch returns and +synchronized by the next graph fill. Upstream runs the fill after the launch returns; here (fork +#65/#67) the fill runs on the filler thread BEFORE its own graph can pass its WAIT, so that fence +makes fill k wait on graph k, which waits on fill k: decode hangs on the second overlapped step. +The previous graph's read of ``_graph_pinned`` is already ordered before the next fill by the +per-fill readback event. This test drives the overlapped decode loop against an asynchronous, +in-order fake stream (launch returns at once; events record a stream position) and fails if any +step stalls, including when the fence is reintroduced. +""" + +from __future__ import annotations + +import threading +import time +from types import SimpleNamespace + +import torch + +import freetoken.models.qwen4_exp.ple_disk as ple_disk +from freetoken.models.qwen4_exp.ple_disk import DiskRowTable + +EOS = 9 + + +class _Stream: + """One in-order stream: a graph item WAITs on the flag, then RESETs it; launch is async.""" + + def __init__(self, flag: torch.Tensor): + self.flag = flag + self.cv = threading.Condition() + self.queue: list[str] = [] + self.enqueued = self.done = 0 + self.stop = False + threading.Thread(target=self._run, daemon=True).start() + + def enqueue(self, kind: str) -> None: + with self.cv: + self.queue.append(kind) + self.enqueued += 1 + self.cv.notify_all() + + def _run(self) -> None: + while True: + with self.cv: + self.cv.wait_for(lambda: self.queue or self.stop) + if self.stop: + return + kind = self.queue.pop(0) + if kind == "graph": + deadline = time.monotonic() + 10.0 + while int(self.flag[0]) != 1 and time.monotonic() < deadline: + time.sleep(0.0005) + self.flag[0] = 0 + with self.cv: + self.done += 1 + self.cv.notify_all() + + def wait_done(self, n: int, timeout: float) -> bool: + with self.cv: + return self.cv.wait_for(lambda: self.done >= n, timeout) + + def close(self) -> None: + with self.cv: + self.stop = True + self.cv.notify_all() + + +class _Event: + """cudaEvent semantics: synchronize() waits for the work enqueued before the LAST record().""" + + def __init__(self, stream: _Stream): + self.stream, self.pos = stream, 0 + + def record(self, _stream=None) -> None: + self.pos = self.stream.enqueued + + def synchronize(self) -> None: + if not self.stream.wait_done(self.pos, 2.0): + raise TimeoutError(f"a fill waited on stream position {self.pos}, which never completed") + + +class _Store: + def __init__(self, flag: torch.Tensor): + self.flag = flag + + def stage(self, tokens_addr, n, staging_addr): + pass + + def flush(self, signal_addr): + if signal_addr: + self.flag[0] = 1 + + +def test_overlapped_graph_fills_never_wait_on_their_own_launch(monkeypatch): + flag = torch.zeros(1, dtype=torch.int64) + stream = _Stream(flag) + monkeypatch.setattr(ple_disk.torch.cuda, "current_stream", lambda *a, **k: None) + monkeypatch.setattr(ple_disk.torch.cuda, "Event", lambda *a, **k: _Event(stream)) + t = DiskRowTable.__new__(DiskRowTable) + t._device = None + t._wait_sync = True + t._flag = flag + t._store = _Store(flag) + t._token_readback = torch.zeros(8, dtype=torch.int32) + t._graph_pinned = torch.zeros(64, dtype=torch.uint8) + t._eager_pinned = [torch.zeros(64, dtype=torch.uint8), torch.zeros(64, dtype=torch.uint8)] + t._eager_read = [_Event(stream), _Event(stream)] + t._eager_slot = 0 + # if upstream's post-launch fence is ever reintroduced, model it with real stream semantics + # so this test fails by hanging the way production would, not on a missing attribute + t._graph_consumed = _Event(stream) + t._token_bytes = 1 + t.eos_token_id = EOS + t.image_token_id = None + + tokens = [1, 2, 3] + steps = 12 + try: + for step in range(steps): + new = 10 + step + req = SimpleNamespace(input_ids=torch.tensor(tokens + [new]), device_len=len(tokens) + 1) + batch = SimpleNamespace(is_decode=True, reqs=[req], padded_size=1, + input_ids=torch.tensor([new], dtype=torch.int32)) + with t.forward_host_ctx(batch, use_graph=True): + stream.enqueue("graph") # cuGraphLaunch returns before the GPU passes the WAIT + tokens.append(new) + # the overlap loop keeps the engine at most one launch ahead of the GPU + assert stream.wait_done(step, 4.0), f"decode hung: graph {step - 1} never completed" + t._raise_failed_fill() + assert stream.wait_done(steps, 4.0), "decode hung on the last step" + t._raise_failed_fill() + finally: + t.close() + stream.close() From 28ba6af86820a8d4718bec6f0990fb11237b0c24 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 7 Oct 2026 17:33:31 -0400 Subject: [PATCH 71/73] fix(kernel): keep one e4m3_f32_to_u8, #408's copy duplicates the one #603 added to main Assisted-by: Claude Opus 5.5 --- python/freetoken/kernel/triton/e4m3_compat.py | 27 ------------------- 1 file changed, 27 deletions(-) diff --git a/python/freetoken/kernel/triton/e4m3_compat.py b/python/freetoken/kernel/triton/e4m3_compat.py index 6c8e568b25..f0a4464a5a 100644 --- a/python/freetoken/kernel/triton/e4m3_compat.py +++ b/python/freetoken/kernel/triton/e4m3_compat.py @@ -211,33 +211,6 @@ def round_e4m3(x): return tl.where(tl.abs(x) >= 0.015625, y_norm, y_sub) -@jit -def e4m3_f32_to_u8(x): - """Encode an fp32 value that ALREADY lies on the e4m3 grid -- the output of - :func:`round_e4m3`, clamped to +-448 -- into its e4m3 byte code. This is the - encoder the pre-sm_89 emulated path needs to STORE fp8-sized data (the fp8 - type itself is unavailable there, so the bytes live in a uint8 buffer that - :func:`e4m3_u8_to_f32` decodes back). - - Normal range: read the (unbiased) exponent and the now-zero-padded fp32 - mantissa back out of the fp32 header. Subnormal range (|x| < 2^-6, grid step - 2^-9): the value is an exact multiple of 2^-9, so ``|x| * 512`` IS the mantissa - field -- the sign bit has to be carried in by hand, since that branch never - looks at the header. The same 0.015625 boundary as :func:`round_e4m3` keeps the - two consistent: ``e4m3_u8_to_f32(e4m3_f32_to_u8(round_e4m3(x)))`` is x's - single-rounded value for every input, and no code it emits is a NaN pattern (the - caller's +-448 clamp caps the code at 0x7E). ``-0.0`` encodes as 0x00 after - round_e4m3 (which documents returning +0.0 for it). - """ - u = x.to(tl.uint32, bitcast=True) - sign = ((u >> 31) & 1).to(tl.int32) - exp = ((u >> 23) & 0xFF).to(tl.int32) - 127 - mant = ((u >> 20) & 7).to(tl.int32) - normal = (sign << 7) | ((exp + 7) << 3) | mant - sub = (sign << 7) | (tl.abs(x) * 512.0).to(tl.int32) - return tl.where(tl.abs(x) >= 0.015625, normal, sub).to(tl.uint8) - - @jit def kv_load_e4m3_tile_f32(ptrs, mask): """Load a tile of KV e4m3 codes and widen it to fp32. From 99c0413d6051678077fe2c43df0792b8bcb0477b Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 7 Oct 2026 17:38:06 -0400 Subject: [PATCH 72/73] fix(server): leave late Anthropic system messages to the inline-system policy instead of hoisting them #487's hoist ran before #608's fold on /v1/messages and moved a late system reminder into the system head, so every reminder missed the prefix cache; the OpenAI and Responses paths keep the hoist the qwen templates need. Assisted-by: Claude Opus 5.5 --- python/freetoken/server/anthropic_api.py | 3 ++- python/freetoken/server/generation.py | 7 ++++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/python/freetoken/server/anthropic_api.py b/python/freetoken/server/anthropic_api.py index 545cf4e7cc..35d902b9d6 100644 --- a/python/freetoken/server/anthropic_api.py +++ b/python/freetoken/server/anthropic_api.py @@ -326,7 +326,8 @@ def convert_anthropic_prompt( # default only covers requests that leave thinking out. ctk = apply_default_thinking_mode(ctk, default_thinking_mode) - return render_messages(messages), template_tools, parser_tools, ctk + # late system reminders stay in place: the tokenizer folds them instead (prefix reuse) + return render_messages(messages, hoist_system=False), template_tools, parser_tools, ctk def convert_anthropic_to_genspec( diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index 0ed933feb8..e5a215658b 100644 --- a/python/freetoken/server/generation.py +++ b/python/freetoken/server/generation.py @@ -516,16 +516,17 @@ def pick(value, key, framework): return params -def render_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: +def render_messages(messages: list[dict[str, Any]], *, hoist_system: bool = True) -> list[dict[str, Any]]: """Normalize OpenAI-shaped message dicts for the chat template: flatten text content parts to a string and decode tool-call arguments from JSON. Raises ValueError on a non-text content part (text-only server). Shared by all adapters. Some chat templates (e.g. Qwen3.6) require the system message at index 0; hoist system messages to the front and merge multiples into one to satisfy - that constraint.""" + that constraint. ``hoist_system=False`` leaves late system messages in place + for the tokenizer's inline-system policy (the Anthropic path).""" rendered = [_render_message(m) for m in messages] - if any(m.get("role") == "developer" for m in rendered): + if not hoist_system or any(m.get("role") == "developer" for m in rendered): # The tokenizer maps developer per template and then merges in order # (_map_developer_role); hoisting system alone here would put it ahead of # a developer message that came first. From 1b6e0f0e4334b246956cef9dc1aeba8bda70667f Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 7 Oct 2026 17:39:01 -0400 Subject: [PATCH 73/73] test(server): give main's DeepSeek-V4.1 detector its contract class and the count fake the inline-system kwarg Assisted-by: Claude Opus 5.5 --- tests/server/test_default_thinking_mode.py | 2 +- tests/server/test_function_call_parser.py | 13 ++++++++++++- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/tests/server/test_default_thinking_mode.py b/tests/server/test_default_thinking_mode.py index 7f9938f60c..56854641f2 100644 --- a/tests/server/test_default_thinking_mode.py +++ b/tests/server/test_default_thinking_mode.py @@ -110,7 +110,7 @@ def ctk(**extra): def test_anthropic_count_tokens_renders_the_generation_prompt_under_the_default(monkeypatch): captured = {} - async def fake_count(messages, template_tools, ctk, state): + async def fake_count(messages, template_tools, ctk, state, **_): captured["ctk"] = ctk return 1 diff --git a/tests/server/test_function_call_parser.py b/tests/server/test_function_call_parser.py index 23cd12272e..8af5bc88b6 100644 --- a/tests/server/test_function_call_parser.py +++ b/tests/server/test_function_call_parser.py @@ -454,6 +454,12 @@ class TestDeepSeekV32Detector(_DetectorContract): truncate_before = '' +class TestDeepSeekV41Detector(_DetectorContract): + parser_name = 'deepseekv41' + block = '<｜DSML｜ calls><｜DSML｜ invoke name="read"><｜DSML｜ parameter name="filePath" string="true">/tmp/test_calc.py' + truncate_before = '' + + class TestQwen3CoderDetector(_DetectorContract): parser_name = 'qwen3_coder' block = '/tmp/test_calc.py' @@ -498,5 +504,10 @@ def test_contract_classes_cover_every_concrete_detector(): type(cls().parser().detector) for cls in _DetectorContract.__subclasses__() } - assert covered == set(BaseFormatDetector.__subclasses__()) + concrete, todo = set(), list(BaseFormatDetector.__subclasses__()) + while todo: # a detector may extend another (DeepSeekV41Detector < DeepSeekV32Detector) + cls = todo.pop() + concrete.add(cls) + todo.extend(cls.__subclasses__()) + assert covered == concrete assert covered == set(FunctionCallParser.ToolCallParserEnum.values())