diff --git a/assets/ascii-art.txt b/assets/ascii-art.txt new file mode 100644 index 0000000..9aa5e82 --- /dev/null +++ b/assets/ascii-art.txt @@ -0,0 +1,18 @@ + + + MMM MMMMMMMMMMMMMMMMMMM + M MMMMMMMMMMMMMMMMMMM MMMMMMMMMMMMMMMMMMM MMMMMM + MMMMMMMMMMMMMMMMMMMMMM MMMMMMMMMMMMMMMMMMM MMMMMM + MMMMMMMMMMMM MMMM MMMP MMMMMMMMMMMMMMMMMMM MMM MMMMMM MMM MM + MMM MMMMMMMMMMMMMMMM MMMMMMMMMMMMM MMMMMMPMMMMM MMMMMMPMMMMM MMMMMMM MMMMMMMMMMMM MMMMMMM MMMMMMMM MMMMMMMMMMMM MMMMMMMMMMMMMMMM + MMMM MMMMMMMMMMMMMMMMM MMMMMMMMMMMPM MPMMMPMMMMMPMMM MMPMMPMMMMPMMMM MMMMMM MMMMMMMMMMMMMMM MMMMMM MMMMMMMM MMMMMMMMMMMMMM MMMMMMMMMMMMMMMMM + MMMMMMMMMMMMMMM MPMPMPMPMPMMPMMMMPM PMMMMP MPMMMM MMMPMM MMMMMMM MMMMMMM MMMMMMM MMMMMMMMMMMMM MMMMMMM MMMMMM MMMMMMMM MMMMMMMM + MMMMMMMMMMMMMMMMMMMM MMMMMMM MMMPMMMMMMMMMPMMM MMMMPMMMMMMMPMMMM MMMMMMM MMMMMMM MMMMMM MMMMMMMMMMMM MMMMMMMMMMMMMMMMM MMMMMMM MMMMMMM + MMMMMM MMMMMMM MPMMMMMMMMMPMMMPM MMPMMMMMMMPMMMPMM MMMMMM MMMMMMM MMMMMMM MMMMMMMMMMMM MMMMMMMMMMMMMMMMM MMMMMM MMMMMM + MMM MMMMMMMMMM MPMPMP MMMPMP MMMMPM MMMMMMM MMMMMM MMMMMM MMMMMMMMMMMMM MMMMMMM MMMMMM MMMMMM + M MMMMMMMMMM MMMMMMM MMMMMMMMMMMMMMMM MMPMMMMMMMMMMMMM MMMMMMM MMMMMMMMMMMMMMMM MMMMMM MMMMMMM MMMMMMMMMMMMMMM MMMMMMM MMMMMMM + MMMMMMMMM MMMMMM PMMMMMMMMMMMMM MMMMMMMMMMMMPM MMMMMM MMMMMMMMMMMMMM MMMMMM MMMMMMMM MMMMMMMMMMMMM MMMMMM MMMMMM + MMMMMMMMM MPMPMP PMMMPMPMM MPMPMPMPM MMMMMM MMMMMMMMMM MMMMMM MMMMMMM MMMMMMMM MMMMMM MMMMMM + + + diff --git a/docs/cli.md b/docs/cli.md index cf4b27a..c04b0ab 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -40,6 +40,8 @@ parsers all resolve automatically from the checkpoint and the GPU. |---|---|---| | `--host` | 127.0.0.1 | Bind address | | `--port` | 1919 | Bind port | +| `--tensor-parallel-size`, `--tp-size` | 1 | GPUs to shard the model over; one scheduler process per rank, rank `i` on `cuda:i` | +| `--distributed-timeout` | 1800 | Seconds a TP collective waits before failing; sized for the per-rank skew in weight-load time | | `--max-running-requests` | 4 | Max concurrently running requests | | `--max-output-tokens` | 32768 | Default output budget for requests that omit one | | `--max-seq-len-override` | from checkpoint | Max sequence length | diff --git a/docs/models.md b/docs/models.md index e4850a1..4381a69 100644 --- a/docs/models.md +++ b/docs/models.md @@ -31,6 +31,27 @@ for them; other checkpoints of the same architectures work too. `offload`, upgraded to `hybrid` when a cached `ft bench bw` profile recommends it. +## Tensor parallelism + +`ft serve --tensor-parallel-size N` shards the model over N GPUs on one host. One +scheduler process runs per rank, and rank `i` uses `cuda:i`. + +DeepSeek-V4 shards: + +- MLA query and output heads (`wq_b` column-parallel, `wo_a` by output group, + `wo_b` row-parallel with one all-reduce per block) +- the MoE intermediate dimension, for both the shared expert and the routed FP4 + experts, so the host expert banks divide by N instead of being replicated +- the vocabulary, for the embedding table and the output head + +It keeps replicated: the MLA latent KV path (`wkv` and the paged KV pools — every +head reads the same latent KV), the compressors and the Lightning Indexer (every +rank must select the same blocks), and the router. + +Constraints for DeepSeek-V4: N must divide `o_groups` (8), so N is 1, 2, 4 or 8. +The KV pool is replicated, so its cost per GPU does not fall with N; the weights +and the expert banks do. + ## Notes - `ft checkpoint` conversion is optional — it pre-converts a checkpoint into diff --git a/python/freetoken/attention/dsv4_compress.py b/python/freetoken/attention/dsv4_compress.py index c034b1b..b889fed 100644 --- a/python/freetoken/attention/dsv4_compress.py +++ b/python/freetoken/attention/dsv4_compress.py @@ -51,7 +51,7 @@ def compress_rows_of(self, ti: int, block_starts: torch.Tensor, ratio: int) -> t def decode_compress_rows( self, rows: torch.Tensor, pos: torch.Tensor, ratio: int, layer_id: int, tier: str, - completed: torch.Tensor, + completed: torch.Tensor, ti: int | None = None, ) -> torch.Tensor: """Per-row decode store destination: the completed block's arithmetic row, or the row's OWN scratch row when this step did not finish a block. @@ -59,8 +59,17 @@ def decode_compress_rows( Reads the decode SNAPSHOT (not the live map) so a concurrent allocate_paged cannot redirect the write; scratch keeps the masked store free of negative indices (which ``index_copy_`` would treat as out of bounds) and collision-free across rows. + + ``ti`` addresses the LIVE map by table index instead, for a caller that is not a + decode batch. A speculative block walks the compressor per token from inside a + prefill, where no decode snapshot exists -- and it is safe there for the same + reason the rest of prefill is: the batch owns its rows for the whole forward. """ - row_of_block = self.pool.cmp_rows(self.snapshot()[rows, pos], ratio) + source = ( + self.pool.full_loc_map[ti, pos] if ti is not None + else self.snapshot()[rows, pos] + ) + row_of_block = self.pool.cmp_rows(source, ratio) scratch = rows + self.compress_scratch_base(layer_id, tier) return torch.where(completed, row_of_block, scratch) diff --git a/python/freetoken/attention/dsv4_indexer.py b/python/freetoken/attention/dsv4_indexer.py index aba5f63..f3cf595 100644 --- a/python/freetoken/attention/dsv4_indexer.py +++ b/python/freetoken/attention/dsv4_indexer.py @@ -43,7 +43,7 @@ def indexer_prefill_logits( def indexer_decode_scores( self, q: torch.Tensor, weights: torch.Tensor, valid: torch.Tensor, n_stage: int, - ratio: int, layer_id: int, + ratio: int, layer_id: int, rows: torch.Tensor, ) -> torch.Tensor: """Head-reduced scores ``[B, n_stage]`` for a decode step, gathering each block's key off the decode SNAPSHOT and bounding the work by the live block count read from device @@ -52,7 +52,7 @@ def indexer_decode_scores( return indexer_decode_logits( q, weights, self.compress_pool(layer_id, "idx"), self.snapshot(), - valid, n_stage, ratio, + rows, valid, n_stage, ratio, ) def indexer_select_prefill( diff --git a/python/freetoken/attention/dsv4_sparse.py b/python/freetoken/attention/dsv4_sparse.py index 2c740e2..606e0c5 100644 --- a/python/freetoken/attention/dsv4_sparse.py +++ b/python/freetoken/attention/dsv4_sparse.py @@ -208,6 +208,18 @@ def prepare_for_capture(self, batch: Batch) -> None: def prepare_for_replay(self, batch: Batch) -> None: self._point_to_capture(batch, batch.padded_size, self._table_rows(batch)) + def prepare_for_spec_capture(self, batch: Batch) -> None: + """Point a fixed-width DSpark target capture at the reserved dummy row.""" + bs = batch.padded_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_spec_replay(self, batch: Batch) -> None: + """Stage one snapshot row per request for a DSpark target replay.""" + self._point_to_capture(batch, batch.padded_size, self._table_rows(batch)) + # ----- DSV4 addressing vocabulary --------------------------------------------------- def snapshot(self) -> torch.Tensor: """The current decode batch's snapshot (see ``DSV4AttnMetadata.full_snapshot``); @@ -221,6 +233,18 @@ def window_slots_of(self, ti: int, lo: int, hi: int) -> torch.Tensor: that have slid out of the window translate to -1).""" return self.pool.translate_full_to_window(self.pool.full_loc_map[ti, lo:hi]) + def window_slots_at(self, rows: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + """Window slots for SCATTERED ``(row, position)`` pairs, one per token. + + ``window_slots_of`` covers one request's contiguous range. A flat batch spans + several requests, and its tokens arrive as parallel row/position vectors, so the + same lookup is a gather rather than a slice. Positions that have slid out of the + window translate to -1, exactly as they do there. + """ + return self.pool.translate_full_to_window( + self.pool.full_loc_map[rows.long(), positions.long()] + ) + def win_cols_to_global(self, win_cols: torch.Tensor, slot_lut: torch.Tensor) -> torch.Tensor: """Per-query window columns (or -1) -> GLOBAL window-pool slots via ``slot_lut``.""" g = slot_lut[win_cols.clamp_min(0)] @@ -279,11 +303,14 @@ def _point_to_capture(self, batch: Batch, bs: int, rows_ti: torch.Tensor) -> Non """ assert self.capture is not None and bs <= self.max_graph_bs cap = self.capture + prior = getattr(batch, "attn_metadata", None) + segments = getattr(prior, "segments", None) 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]) batch.attn_metadata = DSV4AttnMetadata( - last_indices=cap.last_indices[:bs], full_snap=cap.full_snap[:bs], + last_indices=cap.last_indices[:bs], segments=segments, + full_snap=cap.full_snap[:bs], window_ar=self._window_ar, ) diff --git a/python/freetoken/core.py b/python/freetoken/core.py index ef0a539..dafeb98 100644 --- a/python/freetoken/core.py +++ b/python/freetoken/core.py @@ -89,6 +89,22 @@ def complete_one(self) -> None: self.cached_len = self.device_len self.device_len += 1 + def complete_n(self, n: int) -> None: + """Advance by ``n`` accepted tokens instead of one. + + A speculative step emits a whole block, and each request in a batch keeps a + DIFFERENT amount of it -- whatever prefix the target agreed with. So the advance + is per request, not per batch. + + ``cached_len`` moves to where the step started, not to ``device_len - 1``: every + token of the accepted prefix is now cached history, and the next step extends + from the end of it. + """ + if n < 1: + raise ValueError(f"a step must advance by at least one token, got {n}") + self.cached_len = self.device_len + self.device_len += n + def append_host(self, next_token: torch.Tensor) -> None: n = self.input_ids.numel() m = n + next_token.numel() @@ -133,6 +149,33 @@ class Batch: active_table_idx: "torch.Tensor | None" = None # this field should be set by attention backend attn_metadata: BaseAttnMetadata = field(init=False) + # Speculative block state. A verify step rides the PREFILL path (1+k tokens per + # request resuming from its own position), so the phase alone cannot distinguish + # it from a real prefill -- these say so explicitly. + speculative: bool = field(default=False, init=False) + spec_block: int = field(default=0, init=False) + # GraphRunner sets this only while the target verify graph is capturing/replaying. + # It selects DSV4's fixed-shape, device-addressed verify implementation; the eager + # prefill implementation remains the fallback when no verify graph is available. + spec_verify_decode: bool = field(default=False, init=False) + draft_confidence: torch.Tensor | None = field(default=None, init=False) + # Per-request accepted tokens from a speculative step: the reply path emits one + # token per request, so the rest of each block is read from here. + spec_emitted: List[torch.Tensor] | None = field(default=None, init=False) + # The draft distribution q for this block, kept so acceptance can run the + # p/q ratio test rather than an argmax comparison. + draft_probs: torch.Tensor | None = field(default=None, init=False) + # Tokens produced by the paper's sequential Markov stage, [requests * gamma]. + # They stay on GPU through verification; only the accepted prefix is copied back. + draft_tokens: torch.Tensor | None = field(default=None, init=False) + # Per-token compressor partial states produced by a speculative target verify. + # FreeToken's compressor ring is page-addressed, so later tokens in the same page + # overwrite earlier states. This journal is the engine-native equivalent of + # vLLM's save_partial_states: acceptance selects the state at the last valid row. + spec_carry_states: dict | None = field(default=None, init=False) + # Returns a rejected block's unused pages/SWA slots. Supplied by the scheduler, + # which owns the cache manager and is what inflated device_len to the block width. + release_tail: object | None = field(default=None, init=False) # concatenated multimodal soft-token embeddings for a prefill batch (or None) mm_embeds: torch.Tensor | None = field(default=None, init=False) # Prefill log stats snapshotted at schedule time (before forward's complete_one() diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 543012f..7fddd5a 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -72,8 +72,23 @@ class EngineConfig: # ratio default above. A runtime cache rebuild sets this (num_swa_pages) to pin the window # regardless of the full anchor; the ratio is the startup default and the fallback. swa_num_pages_override: int | None = None - distributed_timeout: float = 60.0 + # Collective timeout for the TP process group. The FIRST collective is the post-load + # free-memory all-reduce, and ranks reach it minutes apart on a large MoE checkpoint + # (each reads its own expert shards, at its own speed), so a short timeout kills the + # server during a healthy startup rather than protecting it. Sized for that load skew; + # lower it with --distributed-timeout if a hung rank should fail faster. + distributed_timeout: float = 1800.0 use_dummy_weight: bool = False + # DeepSeek-V4 only: build the checkpoint's dSpark drafter for block speculative + # decoding. Reaches the model through dsv4_args.dspark_enabled (_adjust_dsv4_config). + speculative_dspark: bool = False + # Experimental, request-local DSpark circuit breaker. 0 disables it. Once at least + # dspark_fallback_min_drafted proposals have been measured below this acceptance + # rate, use ordinary target decode for dspark_fallback_steps steps, then probe the + # drafter again. This is based on observed acceptance, not prompt classification. + dspark_fallback_acceptance: float = 0.0 + dspark_fallback_min_drafted: int = 32 + dspark_fallback_steps: int = 64 use_pynccl: bool = True max_seq_len_override: int | None = None num_page_override: int | None = None # if not None, will override the number of pages diff --git a/python/freetoken/engine/distribution.py b/python/freetoken/engine/distribution.py new file mode 100644 index 0000000..93ba77b --- /dev/null +++ b/python/freetoken/engine/distribution.py @@ -0,0 +1,155 @@ +"""Final CPU/GPU placement summary for a loaded engine. + +The startup banner describes hardware. This module describes where the loaded model +actually landed after host expert banks, GPU caches, and CUDA graphs have reached their +serving sizes. The data is gathered from every TP rank immediately before readiness. +""" + +from __future__ import annotations + +import os +from collections import defaultdict +from typing import Any + +import torch + +from freetoken.utils import numa + + +def _gib(n_bytes: int) -> str: + return f"{n_bytes / (1 << 30):.2f} GiB" + + +def _percent(part: int, whole: int) -> float: + return 100.0 * part / whole if whole > 0 else 0.0 + + +def host_expert_bytes(cache: Any) -> int: + """Pinned host expert bytes owned by one TP rank.""" + if cache is None: + return 0 + return sum( + tensor.numel() * tensor.element_size() + for layers in getattr(cache, "bank_sources", {}).values() + for tensor in layers + ) + + +def _thread_siblings(cpu: int) -> str | None: + """Return the kernel topology key without depending on optional NUMA helpers.""" + try: + with open( + f"/sys/devices/system/cpu/cpu{cpu}/topology/thread_siblings_list" + ) as f: + return f.read().strip() + except OSError: + return None + + +def _physical_core_count(cpus: list[int]) -> int: + siblings = {_thread_siblings(cpu) or f"cpu:{cpu}" for cpu in cpus} + return len(siblings) + + +def _node_memory_total(node: int | None) -> int: + if node is None: + return 0 + try: + with open(f"/sys/devices/system/node/node{node}/meminfo") as f: + for line in f: + if " MemTotal:" in line: + return int(line.split()[-2]) * 1024 + except (OSError, ValueError, IndexError): + pass + return 0 + + +def _system_memory_total() -> int: + try: + return int(os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES")) + except (OSError, ValueError): + return 0 + + +def local_distribution_snapshot(engine: Any) -> dict[str, Any]: + """Collect one rank's final serving placement without synchronizing a hot path.""" + rank = int(engine.config.tp_info.rank) + free, total = torch.cuda.mem_get_info(engine.device) + placed = numa.placement() + node = placed[0] if placed is not None else numa.device_numa_node(rank) + node_cpus = ( + list(placed[1]) if placed is not None else sorted(numa.allowed_cpus()) + ) + executor = getattr(engine, "cpu_moe_executor", None) + workers = int(getattr(executor, "num_threads", 0) or 0) + coordinator = int(getattr(executor, "_coord_core", -1) >= 0) + return { + "rank": rank, + "gpu": int(engine.device.index if engine.device.index is not None else rank), + "gpu_name": torch.cuda.get_device_name(engine.device), + "gpu_used": int(total - free), + "gpu_total": int(total), + "node": node, + "node_memory_total": _node_memory_total(node), + "node_physical_cores": _physical_core_count(node_cpus), + "cpu_workers": workers, + "cpu_coordinator": coordinator, + "host_experts": host_expert_bytes(getattr(engine, "moe_offload_cache", None)), + "system_memory_total": _system_memory_total(), + } + + +def format_distribution_summary(rows: list[dict[str, Any]]) -> list[str]: + """Format gathered TP snapshots as a compact NUMA-aware workstation map.""" + if not rows: + return [] + rows = sorted(rows, key=lambda row: int(row["rank"])) + by_node: dict[int | None, list[dict[str, Any]]] = defaultdict(list) + for row in rows: + by_node[row.get("node")].append(row) + + lines = ["Model distribution (final, NUMA-aware):"] + for node, node_rows in sorted( + by_node.items(), key=lambda item: (-1 if item[0] is None else int(item[0])) + ): + host = sum(int(row["host_experts"]) for row in node_rows) + memory_total = max(int(row["node_memory_total"]) for row in node_rows) + cores = max(int(row["node_physical_cores"]) for row in node_rows) + assigned = sum( + int(row["cpu_workers"]) + int(row["cpu_coordinator"]) + for row in node_rows + ) + label = "NUMA unknown" if node is None else f"NUMA node {node}" + lines.append( + f" {label}: host experts {_gib(host)}/{_gib(memory_total)} " + f"({_percent(host, memory_total):.1f}% RAM), CPU cores {assigned}/{cores} assigned" + ) + for row in node_rows: + coord = " + 1 coordinator" if row["cpu_coordinator"] else "" + lines.append( + f" rank {row['rank']} -> GPU {row['gpu']} " + f"{_gib(int(row['gpu_used']))}/{_gib(int(row['gpu_total']))} " + f"({_percent(int(row['gpu_used']), int(row['gpu_total'])):.1f}% VRAM), " + f"CPU-MoE {row['cpu_workers']} workers{coord}, " + f"host experts {_gib(int(row['host_experts']))}" + ) + + host = sum(int(row["host_experts"]) for row in rows) + gpu_used = sum(int(row["gpu_used"]) for row in rows) + gpu_total = sum(int(row["gpu_total"]) for row in rows) + system_memory = max(int(row["system_memory_total"]) for row in rows) + counted = host + gpu_used + lines.append( + f" total: CPU host experts {_gib(host)}/{_gib(system_memory)} " + f"({_percent(host, system_memory):.1f}% RAM), GPU {_gib(gpu_used)}/{_gib(gpu_total)} " + f"({_percent(gpu_used, gpu_total):.1f}% VRAM); counted placement " + f"CPU {_percent(host, counted):.1f}% / GPU {_percent(gpu_used, counted):.1f}%" + ) + return lines + + +__all__ = [ + "format_distribution_summary", + "host_expert_bytes", + "local_distribution_snapshot", +] diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index ff3c985..78d6edc 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -2,15 +2,27 @@ import gc import math +import glob import os +import time +from contextlib import nullcontext from datetime import timedelta from typing import Any, Dict, Iterable, NamedTuple, Tuple import torch from freetoken.attention import AttnType, attention_backend_info, create_attention_backend from freetoken.core import Batch, Context, Req, set_global_ctx -from freetoken.distributed import destroy_distributed, enable_pynccl_distributed, set_tp_info +from freetoken.distributed import ( + destroy_distributed, + enable_pynccl_distributed, + get_tp_info, + set_tp_info, +) +from freetoken.env import ENV from freetoken.layers import set_rope_device +from freetoken.utils.numa import numa_nodes +from freetoken.utils.numa import placement as numa_placement +from freetoken.utils.numa import resolve_placement from freetoken.models import create_model, load_weight from freetoken.moe import create_moe_backend, is_offload_moe_backend from freetoken.moe.expert_banks import load_expert_banks @@ -245,8 +257,18 @@ def _make_dummy_weight_state_dict( ) -> Dict[str, torch.Tensor]: state_dict: Dict[str, torch.Tensor] = {} fp8_dtypes = (torch.float8_e4m3fn, torch.float8_e5m2) + e8m0 = getattr(torch, "float8_e8m0fnu", None) for key, param in model_state.items(): - if param.dtype in fp8_dtypes: + if e8m0 is not None and param.dtype == e8m0: + # e8m0 is a bare exponent code (value = 2^(code-127)), so 127 is scale 1.0. + # It reports is_floating_point, but there is no normal_ kernel for it -- and a + # RANDOM exponent would scale a block by up to 2^127 anyway, so 1.0 is both the + # only fillable value and the only sane one. Without this a --dummy-weight run + # of any block-scaled model dies with "normal_kernel_cuda not implemented". + t = torch.empty(param.shape, dtype=param.dtype, device=device) + t.view(torch.uint8).fill_(127) + state_dict[key] = t + elif param.dtype in fp8_dtypes: # torch.randn is not implemented for fp8; fill via a uint8 view with small # codes (avoid NaN/inf fp8 encodings). Lets dummy-weight startup work for # block-fp8 models (the dense fp8 linears are fp8 regardless of moe_backend). @@ -283,18 +305,130 @@ def _materialize_loaded_weight_state_dict( return state_dict +# FREETOKEN_SPEC_DEBUG=1 dumps the first few blocks' token flow. +_SPEC_DEBUG = os.environ.get("FREETOKEN_SPEC_DEBUG", "0") == "1" +# Diagnostic only: synchronize and report the first speculative cycles by stage. +# Disabled by default because the per-cycle synchronization deliberately removes overlap. +_SPEC_TIMING = os.environ.get("FREETOKEN_SPEC_TIMING", "0") == "1" +# One-shot operator profile of the target verify. Explicit opt-in: Kineto adds +# substantial CPU overhead and synchronizes CUDA when the table is emitted. +_SPEC_PROFILE = os.environ.get("FREETOKEN_SPEC_PROFILE", "0") == "1" + class ForwardOutput(NamedTuple): next_tokens_gpu: torch.Tensor next_tokens_cpu: torch.Tensor copy_done_event: torch.cuda.Event +def _cpu_moe_max_tokens(config: EngineConfig) -> int: + """Largest row count one CPU-expert task can receive. + + Ordinary decode has one row per request. DSpark target verification has the + anchor plus ``gamma`` proposal rows per request, while the draft itself has only + ``gamma``; therefore ``1 + gamma`` is the required upper bound. + """ + rows_per_req = 1 + dsv4 = getattr(config.model_config, "dsv4_args", None) + if getattr(dsv4, "dspark_enabled", False): + rows_per_req = 1 + max( + 1, int(getattr(dsv4, "dspark_block_size", 1) or 1) + ) + max_reqs = max(config.max_running_req, config.cuda_graph_max_bs or 0, 1) + return max_reqs * rows_per_req + + +def _bind_rank_to_numa_node(tp_info) -> None: + """Pin this rank to one NUMA node BEFORE it allocates anything. + + The host expert banks are anonymous mmaps, so their pages land on whichever node + first TOUCHES them -- the loader's threads. Unbound, those threads scatter across + every socket, so a rank's banks end up interleaved and the CPU MoE pool, pinned to + one socket's cores, reads much of its experts across the interconnect. That is the + critical path: the hybrid backend computes most of each decode step's expert misses + there. + + Binding first makes first-touch place a rank's banks on the node that will read + them, which also puts every socket's memory controller to work instead of one. + No-op on a single-node host. + """ + # Resolve BEFORE the bind: binding narrows the affinity mask to one node, after + # which the topology looks single-node and the answer would be lost. + nodes = numa_nodes() + placement = resolve_placement(tp_info.rank, tp_info.size) + if placement is None: + # Say WHY nothing happened, so "no NUMA lines" is never ambiguous between + # "single socket" and "the binding silently did not run". + logger.info_rank0( + f"NUMA: not placing ranks ({len(nodes)} usable node(s), " + f"tp_size={tp_info.size}) -- nothing to spread" + ) + return + node, cpus, siblings, index = placement + try: + os.sched_setaffinity(0, cpus) + except (AttributeError, OSError) as exc: # not Linux, or a restricted cpuset + logger.warning( + f"NUMA rank{tp_info.rank}: could not pin to node {node} ({exc}); " + "expert banks may land on a remote node and decode will be slower" + ) + return + # EVERY rank reports, not just rank 0. A placement bug shows up as ranks disagreeing + # with each other -- and that is invisible if only one of them speaks. This is + # exactly how a real bug hid here: the bind said "2 ranks on this node" while the + # thread split said 4, and only rank 0 was logging. + logger.info( + f"NUMA rank{tp_info.rank}/{tp_info.size}: node {node} of {len(nodes)} | " + f"cpus {cpus[0]}..{cpus[-1]} ({len(cpus)}) | " + f"{siblings} rank(s) here, I am #{index} | banks first-touch local" + ) + + +def _share_cpu_threads_across_ranks(tp_size: int) -> None: + """Split the machine's CPU threads across the TP ranks, before any tensor work. + + torch sizes its intra-op pool from the WHOLE machine, and every rank does the same, + so N ranks each start a machine-sized pool and oversubscribe every core N times. + Measured on a 40-core / 80-thread host at TP=4: 117 threads per rank, 468 in total, + ~6x the hardware, all of it contending through the weight load -- which is thread- + bound (host-side copies into the expert banks), not disk-bound. The NVMe sat at + 187 MB/s while the CPU ran at 92% user. + + The engine already narrows threads once the pinned CPU MoE pool exists; this covers + everything BEFORE that, the weight load included. Honour an explicit OMP_NUM_THREADS. + """ + if tp_size <= 1 or os.environ.get("OMP_NUM_THREADS"): + return + # Count what this rank may actually run on, not what the machine has. After the NUMA + # bind above, the affinity mask is one node -- dividing the whole machine by tp_size + # would under-count every rank by the number of nodes. + try: + allowed = len(os.sched_getaffinity(0)) + except AttributeError: + allowed = os.cpu_count() or tp_size + placed = numa_placement() + siblings = placed[2] if placed is not None else tp_size + per_rank = max(1, allowed // siblings) + if per_rank < torch.get_num_threads(): + torch.set_num_threads(per_rank) + from freetoken.distributed import get_tp_info + + logger.info( + f"NUMA rank{get_tp_info().rank}/{tp_size}: torch intra-op threads " + f"{torch.get_num_threads()} | {allowed} cpus visible | " + f"{siblings} rank(s) share them" + ) + + class Engine: def __init__(self, config: EngineConfig): assert not torch.cuda.is_initialized() set_tp_info(rank=config.tp_info.rank, size=config.tp_info.size) _ensure_expandable_segments() # before the first CUDA allocation below _adjust_config(config) + # Bind BEFORE any allocation: the expert banks land on the node that + # first touches them, and that must be the node whose cores will read them. + _bind_rank_to_numa_node(config.tp_info) + _share_cpu_threads_across_ranks(config.tp_info.size) self.device = torch.device(f"cuda:{config.tp_info.rank}") torch.cuda.set_device(self.device) @@ -323,6 +457,19 @@ def __init__(self, config: EngineConfig): self.model.load_state_dict(self._load_weight_state_dict(config)) post_weights_free = self._sync_get_memory()[0] self._weights_bytes = self._baseline_free - post_weights_free + # What the parameters declare vs what the load actually cost. Every byte of the gap + # is resident-but-unaccounted (staging buffers, per-layer scratch) and comes straight + # out of the MoE-cache + KV budget, so surface it instead of leaving it to arithmetic + # on the cache-sizing line. + declared = sum(t.numel() * t.element_size() for t in self.model.state_dict().values()) + logger.info_rank0( + f"Weights: {mem_GB(declared)} declared by parameters, " + f"{mem_GB(self._weights_bytes)} measured on device " + f"(torch allocator holds {mem_GB(torch.cuda.memory_allocated(self.device))} live, " + f"{mem_GB(torch.cuda.memory_reserved(self.device))} reserved; " + f"peak {mem_GB(torch.cuda.max_memory_reserved(self.device))}). " + "A gap between 'reserved' and 'measured' is held OUTSIDE the allocator." + ) # Pool-budget baseline for the desktop cache sliders: free VRAM after the weights are # resident but before ANY runtime cache pool (MoE expert cache below, KV pages, GDN # state) is allocated. This is the stable "if all free VRAM went to one pool" budget — @@ -331,6 +478,13 @@ def __init__(self, config: EngineConfig): self._post_weights_free = post_weights_free self.moe_offload_cache = None self.cpu_moe_executor = None + # Speculative accounting: acceptance rate is the one number that says whether + # speculation is paying for itself, and it cannot be inferred from tokens/s. + self._spec_accepted = 0 + self._spec_debug_left = 6 + self._spec_drafted = 0 + self._spec_timing_left = 12 + self._spec_profile_left = 1 if is_offload_moe_backend(config.moe_backend): self._init_offload_moe_cache(config) if hasattr(self.model, "prepare_for_runtime"): @@ -418,10 +572,29 @@ def __init__(self, config: EngineConfig): dummy_req=self.dummy_req, moe_offload_cache=self.moe_offload_cache, ) + self._init_dspark_adaptive_verification() if config.attention_backend.split(",")[0] == "triton": # Prefill runs on the first comma part; warm its autotune cache. self._warmup_prefill() + def log_distribution_summary(self) -> None: + """Gather final TP placement and print it immediately before readiness.""" + from freetoken.engine.distribution import ( + format_distribution_summary, + local_distribution_snapshot, + ) + + tp = get_tp_info() + rows: list[dict[str, Any] | None] = [None] * tp.size + torch.distributed.all_gather_object( + rows, + local_distribution_snapshot(self), + group=self.tp_cpu_group, + ) + if tp.is_primary(): + for line in format_distribution_summary([row for row in rows if row is not None]): + logger.info(line) + def _init_communication(self, config: EngineConfig) -> torch.distributed.ProcessGroup: if config.tp_info.size == 1 or config.use_pynccl: torch.distributed.init_process_group( @@ -511,7 +684,8 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: # cpu/hybrid both read experts on the CPU, so banks load in the native (CPU-readable) # layout; the GPU slot-cache GEMM reads those same native rows. decode_target also # gates the CPU executor build below. - cpu_layer_ids = _resolve_cpu_layers(config, config.model_config.num_moe_layers) + n_moe = config.model_config.num_moe_layers + cpu_layer_ids = _resolve_cpu_layers(config, n_moe) if config.moe_backend == "hybrid": decode_target = "hybrid" elif cpu_layer_ids: @@ -636,7 +810,12 @@ 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) + # + # A dSpark target verify carries 1 + block_size ROWS per request: the anchor + # plus every drafted token. max_tokens sizes the C++ pool's + # per-task scratch, so a pool built for one row per request would be handed + # more rows than it owns and overrun it. + max_tokens = _cpu_moe_max_tokens(config) # gpt-oss mxfp4 carries clamped-swiglu scalars; other formats use the defaults. executor = CpuMoeExecutor( cache, @@ -859,19 +1038,539 @@ def rebuild_runtime_cache( dummy_req=self.dummy_req, moe_offload_cache=self.moe_offload_cache, ) + self._init_dspark_adaptive_verification() + + def _init_dspark_adaptive_verification(self) -> None: + """Install the paper's measured-cost scheduler when every prefix is captured.""" + self._adaptive_verification = None + curve = list(getattr(self.graph_runner, "spec_verify_cost_curve", ()) or ()) + block_size = int(getattr(self.graph_runner, "spec_block_size", 0) or 0) + if not curve or block_size < 1: + return + if self.config.max_running_req != 1: + logger.warning_rank0( + "DSpark adaptive verification needs the paper's marker-tensor varlen " + "layout for request batches >1; using fixed gamma for this run" + ) + return + + # Every TP rank must select the same graph shape. vLLM broadcasts its + # profiled cost curves from rank 0; do exactly that rather than letting + # small per-GPU timing noise choose different collective graphs. + tp = get_tp_info() + if tp.size > 1: + payload = [curve if tp.is_primary() else None] + torch.distributed.broadcast_object_list( + payload, src=0, group=self.tp_cpu_group + ) + curve = payload[0] + from freetoken.models.deepseek_v4.dspark import DSparkAdaptiveVerification + + self._adaptive_verification = DSparkAdaptiveVerification( + block_size, + curve, + self.device, + fallback_acceptance=self.config.dspark_fallback_acceptance, + fallback_min_drafted=self.config.dspark_fallback_min_drafted, + fallback_steps=self.config.dspark_fallback_steps, + ) + if self.config.dspark_fallback_acceptance > 0: + logger.info_rank0( + "DSpark acceptance fallback enabled: threshold=%.1f%%, " + "min-drafted=%d, target-only-steps=%d", + 100.0 * self.config.dspark_fallback_acceptance, + self.config.dspark_fallback_min_drafted, + self.config.dspark_fallback_steps, + ) + + def should_speculate(self, req: Req) -> bool: + """Return whether this request should pay for a DSpark draft this step.""" + manager = self._adaptive_verification + if manager is None: + return True + return manager.should_speculate(req.uid) + + def _record_dspark_acceptance( + self, req: Req, accepted: int, drafted: int + ) -> None: + manager = self._adaptive_verification + if manager is not None: + manager.record_acceptance(req.uid, accepted, drafted) + + def adapt_speculative_batch(self, batch: Batch) -> None: + """Compact one prepared DSpark block to the paper-selected prefix width. + + Allocation deliberately remains at gamma until acceptance, so abandoned + page/window slots are still returned by ``release_speculative_tail``. Only + the target's input views and metadata shrink before graph replay. + """ + manager = self._adaptive_verification + if manager is None: + return + if len(batch.reqs) != 1 or batch.padded_size != 1: + raise RuntimeError( + "adaptive DSpark compaction currently requires one unpadded request" + ) + confidence = batch.draft_confidence + if confidence is None: + raise RuntimeError("adaptive DSpark received no confidence probabilities") + + max_width = int(batch.spec_block) + if max_width != manager.block_size: + raise RuntimeError( + f"DSpark prepared width {max_width}, expected checkpoint gamma " + f"{manager.block_size}" + ) + req = batch.reqs[0] + width = manager.record_and_choose(confidence, req.uid) + if width == max_width: + return + if not 0 <= width < max_width: + raise RuntimeError(f"adaptive DSpark selected invalid width {width}") + + base = req.input_ids.numel() - max_width + span = width + 1 + req.input_ids = req._ids_buf[: base + width] + batch.input_ids = batch.input_ids[:span] + batch.positions = batch.positions[:span] + if batch.out_loc is not None: + batch.out_loc = batch.out_loc[:span] + if batch.draft_tokens is None or batch.draft_probs is None: + raise RuntimeError("adaptive DSpark cannot compact a missing draft") + batch.draft_tokens = batch.draft_tokens[:width] + batch.draft_probs = batch.draft_probs[:width] + batch.draft_confidence = confidence[:width] + segments = getattr(batch.attn_metadata, "segments", None) + if segments is None or len(segments) != 1: + raise RuntimeError("adaptive DSpark needs one target metadata segment") + _off, _old_n, table_idx, start_pos = segments[0] + batch.attn_metadata.segments = [(0, span, table_idx, start_pos)] + batch.spec_block = width + + def _record_adaptive_draft_cost(self, batch: Batch) -> None: + """Publish rank-0's real drafter time to the shared five-sample profile.""" + manager = self._adaptive_verification + if manager is None or not manager.needs_draft_profile: + return + start = getattr(batch, "_spec_draft_start", None) + end = getattr(batch, "_spec_draft_end", None) + if start is None or end is None: + raise RuntimeError("adaptive DSpark draft timing events are missing") + tp = get_tp_info() + cost = start.elapsed_time(end) if tp.is_primary() else 0.0 + if tp.size > 1: + shared = torch.tensor([cost], dtype=torch.float64) + torch.distributed.broadcast(shared, src=0, group=self.tp_cpu_group) + cost = float(shared.item()) + manager.record_draft_cost(cost) + + def _finish_speculative( + self, batch: Batch, logits: torch.Tensor, args: BatchSamplingArgs, target_features + ) -> ForwardOutput: + """Keep the prefix of each block the target agrees with, and collapse the rest. + + The verify pass scored every drafted position, so ``logits[j]`` is the target's + own prediction for the token after position j. Acceptance is a PREFIX: scan in + order, stop at the first disagreement, and take the target's token there. That + is what makes speculation invisible in the output -- the emitted sequence is + exactly what the target alone would have produced. + + Each request keeps a different amount, so the state is fixed up per request: + ``input_ids`` truncates to the accepted prefix, the bonus token is appended, and + ``cached_len`` / ``device_len`` move to match the KV the verify actually wrote. + """ + from freetoken.models.deepseek_v4.dspark import ( + accepted_prefix, + rejection_accept_device, + sampling_probs, + ) + + k = batch.spec_block + # Acceptance reads one logits row per position of every block. Getting fewer + # means the verify scored only each request's LAST token -- the default for a + # prefill -- and the failure would otherwise surface as an opaque IndexError + # inside the sampler, several frames from the cause. + expected = (1 + k) * len(batch.reqs) + if logits.shape[0] != expected: + raise RuntimeError( + f"speculative verify produced {logits.shape[0]} logits rows, expected " + f"{expected} ({1 + k} per request x {len(batch.reqs)}). The forward must " + "pass logit_indices covering every drafted position." + ) + any_sampled = any(not req.sampling_params.is_greedy for req in batch.reqs) + any_greedy = any(req.sampling_params.is_greedy for req in batch.reqs) + # Not self.sampler: a speculative batch has 1+k rows per request while the + # sampling args are sized per request. Greedy verification needs target argmax + # ids. Probabilistic rejection consumes p directly; reducing and copying argmax + # before its p/q test inserted an otherwise unused synchronization every cycle. + target_cpu = ( + logits.argmax(dim=-1).to(torch.int32).to("cpu", non_blocking=False) + if any_greedy + else None + ) + if any_sampled and batch.draft_probs is None: + raise RuntimeError("sampled DSpark verification is missing draft probabilities") + confidence = batch.draft_confidence + proposed_gpu = batch.draft_tokens + if proposed_gpu is None or proposed_gpu.numel() != k * len(batch.reqs): + raise RuntimeError( + "DSpark verify is missing the gamma proposals produced by its sequential stage" + ) + # Greedy acceptance consumes proposal ids immediately. Sampled acceptance keeps + # them on device until rejection has queued directly behind target verification. + proposed_cpu = ( + proposed_gpu.to("cpu", non_blocking=False).to(torch.int32) + if any_greedy + else None + ) + draft_cost_recorded = False + if any_greedy: + # The blocking id copies prove the draft events have completed. + self._record_adaptive_draft_cost(batch) + draft_cost_recorded = True + + emitted: list[torch.Tensor] = [] + release_tail = getattr(batch, "release_tail", None) + selected_rows: list[int] = [] + accepted_counts: list[int] = [] + off = 0 + for i, req in enumerate(batch.reqs): + sp = req.sampling_params + greedy = sp.is_greedy + span = 1 + k # this request's rows in the flat batch + # Where the block begins IN THE BUFFER. Not device_len - k: device_len runs + # one ahead of the buffer during decode (see the write-back above), so that + # form silently yielded k-1 proposals -- a short slice, not an error. + start = req.input_ids.numel() - k + # Fixed-width verification is the paper/vLLM fallback when the profiled + # hardware-aware scheduler is disabled. A static per-token threshold is + # deliberately not used: DSpark section 3.2 schedules a GLOBAL token budget + # from cumulative survival probabilities and the measured draft/verify cost + # curves. Until that manager is ported, verify the full gamma without + # claiming that the confidence head is itself a threshold rule. + width = k + # A block emits n_acc accepted tokens PLUS the bonus token, so it can carry + # up to width+1 past `start`. Nothing upstream clamps that to the request's + # output budget: a block that starts with 2 tokens left would write 6, run + # off the end of _ids_buf, and leave device_len past max_device_len -- where + # remain_len goes negative, can_decode never turns False, and the request + # decodes forever instead of finishing on "length". + budget = req.max_device_len - start - 1 + width = max(0, min(width, budget)) + if greedy: + assert proposed_cpu is not None and target_cpu is not None + proposed = proposed_cpu[i * k:(i + 1) * k] + n_acc, bonus = accepted_prefix( + proposed[:width], target_cpu[off:off + width + 1] + ) + else: + # Sampled: accept with probability min(1, p(x)/q(x)) and resample from + # the residual on rejection, so the emitted token stays exactly + # p-distributed. An argmax comparison here would bias towards the mode. + p_req = sampling_probs( + logits[off:off + width + 1], + sp.temperature, + sp.top_p, + sp.top_k, + ) + assert batch.draft_probs is not None + n_acc, bonus = rejection_accept_device( + proposed_gpu[i * k:i * k + width], + batch.draft_probs[i * k:i * k + width], + p_req, + ) + if not draft_cost_recorded: + # The rejection sampler's one small D2H result proves the draft + # timing events have completed; do not add another synchronization. + self._record_adaptive_draft_cost(batch) + draft_cost_recorded = True + proposed = proposed_gpu[i * k:(i + 1) * k].to( + "cpu", non_blocking=False + ).to(torch.int32) + # Snapshot before req._ids_buf is rewritten below. Without it, debug logs + # display post-mutation tokens and invent proposal/target agreements. + proposed_snapshot = proposed.clone() if _SPEC_DEBUG else None + if proposed.numel() != k: + raise RuntimeError( + f"block has {proposed.numel()} proposals, expected {k}; the " + "request's buffer and the batch's block width disagree" + ) + keep = start + n_acc + req.input_ids = req._ids_buf[:start] + if n_acc: + req.append_host(proposed[:n_acc].to(req.input_ids.dtype)) + req.append_host(torch.tensor([bonus], dtype=req.input_ids.dtype)) + # Hand back the pages and SWA slots of the positions this block did not + # keep, BEFORE device_len drops past them. allocate_paged sized itself from + # the full block width, and nothing else walks a range above the request's + # current length -- so what is not released here is leaked until the next + # idle integrity check fails, far from the cause. + if release_tail is not None: + release_tail(req, keep) + req.cached_len, req.device_len = keep, keep + 1 + emitted.append( + torch.cat([proposed[:n_acc], torch.tensor([bonus], dtype=torch.int32)]) + ) + if _SPEC_DEBUG and self._spec_debug_left > 0: + self._spec_debug_left -= 1 + # One block's whole token flow, in the order acceptance sees it. Reading + # this beats reasoning about the indices: the greedy output showed every + # word duplicated ("are are", "first first"), which is an emission fault, + # and only the actual tokens say where. + logger.info_rank0( + "spec block: greedy=%s temp=%s topp=%s conf=%s " + "committed=%s proposed=%s target=%s width=%d n_acc=%d bonus=%s " + "emitted=%s", + greedy, sp.temperature, sp.top_p, + None if confidence is None else + [round(float(c), 3) for c in confidence[i * k:(i + 1) * k]], + int(req._ids_buf[start - 1]) if start > 0 else None, + proposed_snapshot[:width].tolist(), + None if target_cpu is None else + target_cpu[off:off + width + 1].tolist(), + width, n_acc, bonus, + emitted[-1].tolist() if emitted else None, + ) + self._spec_accepted += n_acc + self._spec_drafted += width + self._record_dspark_acceptance(req, n_acc, width) + selected_rows.append(off + n_acc) + accepted_counts.append(n_acc) + off += span + + # Select the target's saved compressor state after anchor + accepted prefix. + # Later rejected rows may share its 128-token page and overwrite the live ring. + self._restore_speculative_carry(batch, selected_rows) + committed_features = self._trim_dspark_target_features( + target_features, batch, accepted_counts + ) + self._commit_dspark_target_features(committed_features) + + # The reply path reads one token per request; hand it the LAST emitted token and + # let the scheduler read the rest off req.input_ids, which already holds them. + last = torch.tensor([int(e[-1]) for e in emitted], dtype=torch.int32) + batch.spec_emitted = emitted + gpu = last.to(self.device, non_blocking=True) + done = torch.cuda.Event() + done.record(self.stream) + return ForwardOutput(gpu, last, done) + + def _restore_speculative_carry(self, batch: Batch, selected_rows: list[int]) -> None: + """Commit the per-token partial state selected by DSpark acceptance.""" + journal = batch.spec_carry_states + if not journal: + raise RuntimeError("a DSpark target verify produced no compressor partial states") + device_rows = torch.tensor(selected_rows, dtype=torch.long, device=self.device) + md = batch.attn_metadata + table_rows = torch.tensor( + [md.segments[i][2] for i in range(len(selected_rows))], + dtype=torch.long, + device=self.device, + ) + positions = batch.positions.index_select(0, device_rows).long() + slots = self.attn_backend.window_slots_at(table_rows, positions) + for (layer_id, tier, ring_size), pieces in journal.items(): + if len(pieces) != batch.input_ids.numel(): + raise RuntimeError( + f"DSpark partial-state journal for layer {layer_id}/{tier} has " + f"{len(pieces)} rows, expected {batch.input_ids.numel()}" + ) + states = torch.cat(pieces, dim=0).index_select(0, device_rows) + self.attn_backend.write_carry_blocks( + layer_id, tier, slots, ring_size, states + ) + + def _trim_dspark_target_features(self, features, batch: Batch, accepted: list[int]): + """Keep target features through each request's accepted predecessor row.""" + if features is None: + raise RuntimeError("a DSpark target verify returned no target features") + from freetoken.models.deepseek_v4.model import DSparkTargetFeatures + + segments = batch.attn_metadata.segments + if segments is None or len(segments) != len(accepted): + raise RuntimeError( + "DSpark feature trimming needs one target segment per accepted count" + ) + indices = [] + for (off, _n, _ti, _start), n_acc in zip(segments, accepted, strict=True): + indices.append( + torch.arange(off, off + n_acc + 1, dtype=torch.long, device=self.device) + ) + idx = torch.cat(indices) + if features.hidden.shape[0] != batch.input_ids.numel(): + raise RuntimeError( + f"DSpark target returned {features.hidden.shape[0]} feature rows for " + f"{batch.input_ids.numel()} verify inputs" + ) + return DSparkTargetFeatures( + features.hidden.index_select(0, idx), + features.positions.index_select(0, idx), + features.table_rows.index_select(0, idx), + ) + + def _real_dspark_target_features(self, features, batch: Batch): + """Drop CUDA-graph padding rows before context KV is committed.""" + if features is None or batch.is_prefill or features.hidden.shape[0] == batch.size: + return features + from freetoken.models.deepseek_v4.model import DSparkTargetFeatures + + if features.hidden.shape[0] < batch.size: + raise RuntimeError( + f"target produced {features.hidden.shape[0]} feature rows for " + f"{batch.size} real decode requests" + ) + idx = torch.arange(batch.size, dtype=torch.long, device=self.device) + return DSparkTargetFeatures( + features.hidden.index_select(0, idx), + features.positions.index_select(0, idx), + features.table_rows.index_select(0, idx), + ) + + def _commit_dspark_target_features(self, features) -> None: + """Precompute draft context KV as soon as target rows become committed. + + vLLM performs this precompute immediately before its next proposal. FreeToken + cannot safely retain one model-global bundle across unrelated prefill/decode + batches (or across several captured batch sizes), so it performs the same write + once the target rows are known to be valid. The next proposal observes the + identical draft-layer KV without stale cross-request ownership. + """ + if features is None: + return + catch_up = getattr(self.model, "catch_up_draft_context", None) + if catch_up is not None: + catch_up(features) + + def draft_into_batch(self, batch: Batch) -> torch.Tensor | None: + """Fill a speculative batch's placeholder positions with the drafter's proposal. + + Called on a batch the scheduler has already prepared and extended: its segments + cover ``1 + k`` positions per request, ``allocate_paged`` has given every layer + -- target and draft alike -- slots at those positions, and the token ids past the + first are the checkpoint's noise placeholder. + + The draft runs first and writes the DRAFT layers' KV; the verify forward that + follows runs over the same positions and writes the TARGET layers' KV. The two + never collide, because the draft layers' ids continue past the target's. + + Returns the confidence per position, or None when the model cannot draft yet. + """ + drafter = getattr(self.model, "draft", None) + if drafter is None: + return None + measure_draft = ( + self._adaptive_verification is not None + and self._adaptive_verification.needs_draft_profile + ) or (_SPEC_TIMING and self._spec_timing_left > 0) + if measure_draft: + batch._spec_wall_start = time.perf_counter() + batch._spec_draft_start = torch.cuda.Event(enable_timing=True) + batch._spec_draft_end = torch.cuda.Event(enable_timing=True) + batch._spec_draft_start.record(self.stream) + with self.ctx.forward_batch(batch): + out = drafter( + [req.sampling_params for req in batch.reqs], + ) + if measure_draft: + batch._spec_draft_end.record(self.stream) + if out is None: + raise RuntimeError("DSpark was scheduled without a loaded drafter") + proposed, q, confidence = out + k = batch.spec_block + span = 1 + k + if proposed.numel() != k * len(batch.reqs): + raise RuntimeError( + f"DSpark proposed {proposed.numel()} tokens, expected " + f"{k} x {len(batch.reqs)}" + ) + for i in range(len(batch.reqs)): + batch.input_ids[i * span + 1:(i + 1) * span].copy_( + proposed[i * k:(i + 1) * k].to(batch.input_ids.dtype) + ) + batch.draft_tokens = proposed + batch.draft_probs = q + batch.spec_carry_states = {} + return confidence def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput: assert torch.cuda.current_stream() == self.stream - with self.ctx.forward_batch(batch): - if self.graph_runner.can_use_cuda_graph(batch): - logits = self.graph_runner.replay(batch) - else: - logits = self.model.forward() + target_features = None + timing = bool( + batch.speculative and _SPEC_TIMING and self._spec_timing_left > 0 + ) + profiling = bool( + batch.speculative and _SPEC_PROFILE and self._spec_profile_left > 0 + ) + if timing: + target_start = torch.cuda.Event(enable_timing=True) + target_end = torch.cuda.Event(enable_timing=True) + target_start.record(self.stream) + profile_ctx = ( + torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + record_shapes=True, + ) + if profiling else nullcontext() + ) + with profile_ctx as prof: + with self.ctx.forward_batch(batch): + if self.graph_runner.can_use_spec_cuda_graph(batch): + logits = self.graph_runner.replay_spec(batch) + target_features = self.graph_runner.dspark_spec_target_features(batch) + elif self.graph_runner.can_use_cuda_graph(batch): + logits = self.graph_runner.replay(batch) + target_features = self.graph_runner.dspark_target_features(batch) + else: + logits = self.model.forward() + get_features = getattr(self.model, "dspark_target_features", None) + target_features = get_features() if get_features is not None else None + if profiling: + torch.cuda.synchronize(self.device) + assert prof is not None + logger.info( + "speculative verify profile (by CUDA time):\n%s", + prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=30), + ) + logger.info( + "speculative verify profile (by CPU time):\n%s", + prof.key_averages().table(sort_by="self_cpu_time_total", row_limit=30), + ) + self._spec_profile_left -= 1 + if timing: + target_end.record(self.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. self.cpu_moe_executor.raise_if_unhealthy() + if batch.speculative: + finish_wall_start = time.perf_counter() + output = self._finish_speculative(batch, logits, args, target_features) + if timing: + cycle_end = torch.cuda.Event(enable_timing=True) + cycle_end.record(self.stream) + cycle_end.synchronize() + draft_start = batch._spec_draft_start + draft_end = batch._spec_draft_end + logger.info_rank0( + "spec timing: draft_cuda=%.2fms target_cuda=%.2fms " + "post_cuda=%.2fms finish_wall=%.2fms cycle_wall=%.2fms", + draft_start.elapsed_time(draft_end), + target_start.elapsed_time(target_end), + target_end.elapsed_time(cycle_end), + (time.perf_counter() - finish_wall_start) * 1000, + (time.perf_counter() - batch._spec_wall_start) * 1000, + ) + self._spec_timing_left -= 1 + return output + + self._commit_dspark_target_features( + self._real_dspark_target_features(target_features, batch) + ) + for req in batch.reqs: req.complete_one() @@ -983,12 +1682,37 @@ def _resolve_cache_type(has_linear_attention: bool, requested: str) -> str: 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 - page_size to the window page P, forces single-chunk prefill, and clamps cuda_graph_bs/max_bs to - the DSV4 decode batch size. + page_size to the window page P, and clamps cuda_graph_bs/max_bs to the DSV4 + decode batch size. """ model_config = config.model_config model_config.dsv4_args.max_seq_len = config.max_seq_len model_config.dsv4_args.max_batch_size = config.max_running_req + 1 # +1 dummy + # dSpark is opt-in: the drafter's routed experts enlarge both the host expert banks + # and the GPU slot cache, so a run that will not speculate must not pay for them. + if getattr(config, "speculative_dspark", False): + if not model_config.dsv4_args.has_dspark: + raise ValueError( + "--speculative-dspark: this checkpoint ships no dSpark drafter " + "(needs mtp.* weights with dspark_block_size > 1 in inference/config.json)" + ) + from freetoken.models.deepseek_v4.args import set_dspark_enabled + + model_config.dsv4_args.dspark_enabled = True + set_dspark_enabled(True) # so the weight reader builds the same model + # parse_config already ran, before the flag existed, so its extra_moe_layers is + # still 0. The expert banks would then build n_layers + n_draft entries while + # the offload cache was sized for n_layers, and the two assert against each + # other AFTER the full expert load -- five minutes to learn it. + # frozen dataclass -- same in-place idiom the offload-cache sizing uses below. + object.__setattr__( + model_config, "extra_moe_layers", model_config.dsv4_args.n_draft_layers + ) + logger.info_rank0( + f"dSpark drafter enabled: {model_config.dsv4_args.n_draft_layers} draft layers, " + f"block size {model_config.dsv4_args.dspark_block_size}, " + f"target layers {model_config.dsv4_args.dspark_target_layer_ids}" + ) # config.swa_full_tokens_ratio is the DSV4 window/full ratio directly (default sizing); # a runtime rebuild pins an absolute window via swa_num_pages_override instead. # DSV4's KV page IS the P-token window page (window == radix reuse granularity == lcm of @@ -1000,12 +1724,12 @@ def _adjust_dsv4_config(config: EngineConfig, override) -> None: # 'naive' stays naive with the pool's swa currency riding swa_paged. if getattr(config, "cache_type", "radix") != "naive": override("cache_type", "swa_radix") - # 'radix' (SWARadixCache on the full-loc currency, carry-aware re-prefill) is the default and is - # honored, as is an explicit 'naive'. Don't let max_extend_tokens force a second chunk within - # one prompt (the pool's prefill_chunk_budget still chunks prompts larger than the window - # pool); prefill batches ragged (bs>=1), each segment resuming from its own cached_len. - if getattr(config, "max_extend_tokens", 0) < config.max_seq_len: - override("max_extend_tokens", config.max_seq_len) + # Honor --max-prefill-length. DSV4's compressor carry and dSpark context + # catch-up both support continuation chunks; overriding the user's limit to + # max_seq_len made a long prompt materialize the full mHC activation at once. + # On memory-tight deployments that turns one oversized client history into a + # process-fatal prefill OOM. The pool's prefill_chunk_budget remains a second, + # capacity-derived upper bound in Scheduler. # DSV4 decode batches at most max_running_req rows; its full-loc snapshot is sized to that, # so a graph bs above it would exceed the backend's captured snapshot rows. Clamp any diff --git a/python/freetoken/engine/graph.py b/python/freetoken/engine/graph.py index 4f20250..bfe1dd5 100644 --- a/python/freetoken/engine/graph.py +++ b/python/freetoken/engine/graph.py @@ -1,8 +1,9 @@ from __future__ import annotations import gc +import statistics from dataclasses import dataclass -from typing import TYPE_CHECKING, Dict, List +from typing import TYPE_CHECKING, Dict, List, Tuple import torch from freetoken.core import Batch, Req, get_global_ctx @@ -26,6 +27,7 @@ class GraphCaptureBuffer: positions: torch.Tensor logits: torch.Tensor table_idx: torch.Tensor # per-request slot id for GatedDeltaNet state gather/scatter + request_table_idx: torch.Tensor # per-request paged-KV table row # Decode GDN query indptr = arange(bs+1); a constant per captured bs, filled once. fla_cu_seqlens: torch.Tensor @@ -37,6 +39,7 @@ def init(cls, bs: int, vocab_size: int, device: torch.device) -> GraphCaptureBuf positions=torch.zeros(bs, dtype=torch.int32, device=device), logits=torch.empty(bs, vocab_size, dtype=torch.float32, device=device), table_idx=torch.zeros(bs, dtype=torch.int32, device=device), + request_table_idx=torch.zeros(bs, dtype=torch.int64, device=device), fla_cu_seqlens=torch.arange(bs + 1, dtype=torch.int32, device=device), ) @@ -49,6 +52,7 @@ def set_batch(self, batch: Batch) -> None: batch.out_loc = self.out_loc[_slice] batch.positions = self.positions[_slice] batch.linear_table_idx = self.table_idx[_slice] + batch.active_table_idx = self.request_table_idx[_slice] # Decode GDN metadata reads the persistent cu_seqlens (constant arange) and the # persistent table_idx slot map, so the captured kernels see stable addresses. batch.fla_metadata = FLAMetadata( @@ -63,6 +67,113 @@ def copy_from(self, batch: Batch) -> None: self.positions[_slice] = batch.positions if batch.linear_table_idx is not None: self.table_idx[_slice] = batch.linear_table_idx + if batch.active_table_idx is not None: + self.request_table_idx[_slice] = batch.active_table_idx + + +@dataclass +class SpecGraphCaptureBuffer: + """Persistent inputs/outputs for ``anchor + prefix`` target graphs. + + The allocation is sized for the checkpoint's maximum ``anchor + gamma`` span, + while each captured graph takes a prefix of it. Keeping one backing allocation + gives every replay a stable address without paying one logits buffer per prefix. + """ + + input_ids: torch.Tensor + out_loc: torch.Tensor + positions: torch.Tensor + logits: torch.Tensor + request_table_idx: torch.Tensor + span: int + + @classmethod + def init( + cls, max_reqs: int, span: int, vocab_size: int, device: torch.device + ) -> "SpecGraphCaptureBuffer": + max_tokens = max_reqs * span + return cls( + input_ids=torch.zeros(max_tokens, dtype=torch.int32, device=device), + out_loc=torch.zeros(max_tokens, dtype=torch.int32, device=device), + positions=torch.zeros(max_tokens, dtype=torch.int32, device=device), + logits=torch.empty(max_tokens, vocab_size, dtype=torch.float32, device=device), + request_table_idx=torch.zeros(max_reqs, dtype=torch.int64, device=device), + span=span, + ) + + def set_batch(self, batch: Batch) -> None: + reqs = batch.padded_size + span = int(batch.spec_block) + 1 + if not 1 <= span <= self.span: + raise RuntimeError( + f"DSpark graph span {span} is outside captured range 1..{self.span}" + ) + tokens = reqs * span + batch.input_ids = self.input_ids[:tokens] + batch.out_loc = self.out_loc[:tokens] + batch.positions = self.positions[:tokens] + batch.active_table_idx = self.request_table_idx[:reqs] + + def copy_from(self, batch: Batch) -> None: + reqs = batch.padded_size + span = int(batch.spec_block) + 1 + if not 1 <= span <= self.span: + raise RuntimeError( + f"DSpark graph span {span} is outside captured range 1..{self.span}" + ) + tokens = reqs * span + if batch.input_ids.numel() != tokens or batch.positions.numel() != tokens: + raise RuntimeError( + f"DSpark graph expected {tokens} target rows, got " + f"{batch.input_ids.numel()}/{batch.positions.numel()}" + ) + self.input_ids[:tokens].copy_(batch.input_ids) + if batch.out_loc is not None: + self.out_loc[:tokens].copy_(batch.out_loc) + self.positions[:tokens].copy_(batch.positions) + if batch.active_table_idx is None or batch.active_table_idx.numel() != reqs: + raise RuntimeError("DSpark graph needs one active table row per request") + self.request_table_idx[:reqs].copy_(batch.active_table_idx) + + +class SharedSpecCarryJournal: + """Maximum-span carry outputs shared by every adaptive verify graph. + + The largest graph owns the actual tensors produced by each compressor step. + Smaller graph captures copy their prefix states into those same addresses. This + is safe because verify graphs never replay concurrently, and it avoids retaining + one very large compressor journal per CUDA-graph bucket. + """ + + def __init__(self, storage: dict) -> None: + if not storage: + raise RuntimeError("DSpark maximum verify graph produced no carry states") + self.storage = storage + self._active_rows = 0 + self._cursors: dict = {} + + def reset(self, active_rows: int) -> None: + if active_rows < 1: + raise ValueError(f"DSpark carry journal needs at least one row, got {active_rows}") + self._active_rows = active_rows + self._cursors = {key: 0 for key in self.storage} + + def record(self, key, state: torch.Tensor) -> None: + pieces = self.storage.get(key) + if pieces is None: + raise RuntimeError(f"DSpark carry journal saw unknown compressor {key}") + cursor = self._cursors[key] + if cursor >= self._active_rows or cursor >= len(pieces): + raise RuntimeError( + f"DSpark carry journal overflow for {key}: row {cursor}, " + f"active={self._active_rows}, capacity={len(pieces)}" + ) + pieces[cursor].copy_(state) + self._cursors[key] = cursor + 1 + + def items(self): + for key, pieces in self.storage.items(): + yield key, pieces[: self._active_rows] def _determine_cuda_graph_bs( @@ -118,12 +229,45 @@ def __init__( self.moe_offload_cache = moe_offload_cache self.stream = stream self.device = device + # A captured model forward does not execute Python on replay, so a model-global + # "last features" attribute remains pointed at whichever graph was captured + # last. Keep the actual output bundle owned by each graph size instead. + self._dspark_feature_map: Dict[int, object] = {} + self._spec_feature_map: Dict[Tuple[int, int], object] = {} + self._spec_carry_map: Dict[Tuple[int, int], dict] = {} + # Measured (span, milliseconds) curve for the paper's hardware-aware + # verification scheduler. It is populated only for the single-request + # graphs that FreeToken's current DSV4 serving mode can compact safely. + self.spec_verify_cost_curve: List[Tuple[int, float]] = [] + self.spec_block_size = int( + getattr(model, "speculative_verify_block_size", 0) or 0 + ) + self.spec_span = self.spec_block_size + 1 if self.spec_block_size else 0 self._capture_graphs(max_seq_len, vocab_size, model) def _reset_moe_offload_cache(self) -> None: if self.moe_offload_cache is not None: self.moe_offload_cache.reset() + def _profile_graph_ms(self, graph: torch.cuda.CUDAGraph, replays: int = 5) -> float: + """Median replay cost for DSpark's hardware-capacity curve. + + Begin the series cold, then let ordinary LRU fills and hits participate in + the curve. This prices the cache reuse present during normal engine execution + instead of forcing an artificial cold miss before every timed replay. + """ + samples: list[float] = [] + self._reset_moe_offload_cache() + for _ in range(replays): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record(self.stream) + graph.replay() + end.record(self.stream) + end.synchronize() + samples.append(start.elapsed_time(end)) + return float(statistics.median(samples)) + def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel): # Mark the post-weights "warmup" phase for /health: this stretch (graph capture — or the # remaining readiness work when graphs are disabled) moves no bytes, so without this the @@ -132,6 +276,7 @@ def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel # graphs-disabled early return so that config gets the phase too. emit_progress("Capturing CUDA graphs / warming up", 0, 0) self.graph_map: Dict[int, torch.cuda.CUDAGraph] = {} + self.spec_graph_map: Dict[Tuple[int, int], torch.cuda.CUDAGraph] = {} if self.max_graph_bs == 0: return logger.info_rank0("CUDA graph is disabled.") @@ -178,10 +323,104 @@ def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel with torch.cuda.graph(graph, pool=pool, stream=self.stream): self.buffer.logits[:bs] = model.forward() self._reset_moe_offload_cache() + get_features = getattr(model, "dspark_target_features", None) + if get_features is not None: + features = get_features() + if features is not None: + # These tensors are graph outputs: replay of this exact graph writes + # them in place. Holding the bundle preserves the matching addresses. + self._dspark_feature_map[bs] = features if pool is None: pool = graph.pool() # reuse cuda graph handle to reduce memory self.graph_map[bs] = graph + spec_capture = getattr(self.attn_backend, "prepare_for_spec_capture", None) + if self.spec_span and spec_capture is not None: + self.spec_buffer = SpecGraphCaptureBuffer.init( + self.max_graph_bs, self.spec_span, vocab_size, self.device + ) + # Section 5.2 chooses among the actual CUDA-graph cost cliffs. DSV4 + # currently serves one request at a time, so capture every legal prefix + # there. Keep the prior fixed-width behavior for larger request batches; + # those need the paper's marker-tensor varlen layout before independent + # per-request widths are safe. + adaptive_single_req = self.graph_bs_list == [1] + spans = ( + list(range(1, self.spec_span + 1)) + if adaptive_single_req + else [self.spec_span] + ) + logger.info_rank0( + f"Capturing DSpark target verify graphs: requests={self.graph_bs_list}, " + f"rows/request={spans}" + ) + shared_carry: SharedSpecCarryJournal | None = None + for span in sorted(spans, reverse=True): + for bs in sorted(self.graph_bs_list, reverse=True): + tokens = bs * span + graph = torch.cuda.CUDAGraph() + batch = Batch(reqs=[self.dummy_req] * bs, phase="prefill") + batch.padded_reqs = batch.reqs + batch.speculative = True + batch.spec_block = span - 1 + batch.spec_verify_decode = True + batch.spec_carry_states = {} + self.spec_buffer.set_batch(batch) + # Valid device positions keep capture-time writes within the reserved + # dummy row; replay overwrites this same stable buffer. + self.spec_buffer.positions[:tokens].copy_( + torch.arange(span, device=self.device, dtype=torch.int32).repeat(bs) + ) + self.spec_buffer.request_table_idx[:bs].fill_(self.dummy_req.table_idx) + spec_capture(batch) + active_rows = tokens + with get_global_ctx().forward_batch(batch): + if shared_carry is not None: + shared_carry.reset(active_rows) + batch.spec_carry_states = shared_carry + self.spec_buffer.logits[:tokens] = model.forward() + if shared_carry is None: + batch.spec_carry_states = {} + else: + shared_carry.reset(active_rows) + batch.spec_carry_states = shared_carry + with torch.cuda.graph(graph, pool=pool, stream=self.stream): + self.spec_buffer.logits[:tokens] = model.forward() + self._reset_moe_offload_cache() + if shared_carry is None and adaptive_single_req: + # Spans are descending, so this is the maximum graph. Its + # captured outputs become the one persistent journal backing + # all remaining prefix graphs. + shared_carry = SharedSpecCarryJournal(batch.spec_carry_states) + shared_carry.reset(active_rows) + get_features = getattr(model, "dspark_target_features", None) + key = (bs, span) + if get_features is not None: + features = get_features() + if features is not None: + self._spec_feature_map[key] = features + self._spec_carry_map[key] = ( + shared_carry if shared_carry is not None else batch.spec_carry_states + ) + self.spec_graph_map[key] = graph + if adaptive_single_req and bs == 1: + cost_ms = self._profile_graph_ms(graph) + self.spec_verify_cost_curve.append((span, cost_ms)) + + if self.spec_verify_cost_curve: + # vLLM makes profiled costs nondecreasing before scheduling. Sampling + # noise must not make a wider target batch appear artificially cheaper. + high = 0.0 + monotonic = [] + for span, cost in sorted(self.spec_verify_cost_curve): + high = max(high, cost) + monotonic.append((span, high)) + self.spec_verify_cost_curve = monotonic + logger.info_rank0( + "DSpark profiled target verify curve: %s", + ", ".join(f"{span} rows={cost:.2f}ms" for span, cost in monotonic), + ) + self._reset_moe_offload_cache() free_memory = get_free_memory(self.device) logger.info_rank0(f"Free GPU memory after capturing CUDA graphs: {mem_GB(free_memory)}") @@ -189,6 +428,14 @@ def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel def can_use_cuda_graph(self, batch: Batch) -> bool: return batch.is_decode and batch.size <= self.max_graph_bs + def can_use_spec_cuda_graph(self, batch: Batch) -> bool: + key = (batch.padded_size, int(batch.spec_block) + 1) + return bool( + batch.speculative + and key in self.spec_graph_map + and batch.padded_size == batch.size + ) + def replay(self, batch: Batch) -> torch.Tensor: assert self.can_use_cuda_graph(batch) self.buffer.copy_from(batch) @@ -197,6 +444,33 @@ def replay(self, batch: Batch) -> torch.Tensor: g.replay() return self.buffer.logits[: batch.size] + def replay_spec(self, batch: Batch) -> torch.Tensor: + assert self.can_use_spec_cuda_graph(batch) + span = int(batch.spec_block) + 1 + tokens = batch.padded_size * span + key = (batch.padded_size, span) + self.spec_buffer.copy_from(batch) + self.spec_buffer.set_batch(batch) + batch.spec_verify_decode = True + journal = self._spec_carry_map[key] + reset_journal = getattr(journal, "reset", None) + if reset_journal is not None: + reset_journal(tokens) + batch.spec_carry_states = journal + prepare = getattr(self.attn_backend, "prepare_for_spec_replay") + prepare(batch) + self.spec_graph_map[key].replay() + return self.spec_buffer.logits[:tokens] + + def dspark_target_features(self, batch: Batch): + """The feature output owned by the graph replayed for ``batch``.""" + return self._dspark_feature_map.get(batch.padded_size) + + def dspark_spec_target_features(self, batch: Batch): + return self._spec_feature_map.get( + (batch.padded_size, int(batch.spec_block) + 1) + ) + def pad_batch(self, batch: Batch) -> None: padded_size = ( # choose the first available batch size next(bs for bs in self.graph_bs_list if bs >= batch.size) @@ -213,5 +487,11 @@ def destroy_cuda_graphs(self) -> None: # free-before-alloc cannot reclaim this GPU memory. empty_cache() is left to the # caller / next capture (GraphRunner._capture_graphs already runs it). self.graph_map = {} + self.spec_graph_map = {} + self._dspark_feature_map = {} + self._spec_feature_map = {} + self._spec_carry_map = {} + self.spec_verify_cost_curve = [] self.buffer = None + self.spec_buffer = None gc.collect() diff --git a/python/freetoken/kernel/csrc/src/pynccl.cu b/python/freetoken/kernel/csrc/src/pynccl.cu index 9d41e00..33ffcc0 100644 --- a/python/freetoken/kernel/csrc/src/pynccl.cu +++ b/python/freetoken/kernel/csrc/src/pynccl.cu @@ -72,22 +72,12 @@ inline constexpr auto template_fn = struct NCCLWrapper : public tvm::ffi::Object { public: NCCLWrapper(int rank, int world_size, const size_t max_bytes, NCCLIDList uid) - : m_rank(rank), m_world_size(world_size), m_max_bytes(max_bytes) { + : m_rank(rank), m_world_size(world_size) { + (void)max_bytes; ncclUniqueId id = get_uid(uid); ncclComm_t comm; NCCL_CHECK(::ncclCommInitRank(&comm, m_world_size, id, m_rank)); m_comm = {comm, template_fn<::ncclCommDestroy>}; - - void *buf; - NCCL_CHECK(::ncclMemAlloc(&buf, max_bytes)); - m_sym_mem = {buf, template_fn<::ncclMemFree>}; - - ncclWindow_t win; - NCCL_CHECK(::ncclCommWindowRegister(comm, buf, max_bytes, &win, - NCCL_WIN_COLL_SYMMETRIC)); - m_win = {win, [comm = m_comm](ncclWindow_t w) { - return NCCL_CHECK(::ncclCommWindowDeregister(comm.get(), w)); - }}; } auto all_reduce(tvm::ffi::TensorView t, std::string op) const -> void { @@ -97,40 +87,17 @@ public: RuntimeCheck(t.is_contiguous(), "Tensor must be contiguous"); const auto size_dim = static_cast(t.shape().Product()); const auto dtype = kNCCLDtypeMap.at(t.dtype()); - const auto size_bytes = size_dim * (t.dtype().bits / 8); const auto data_ptr = t.data_ptr(); const auto reduce_op = kNCCLReduceOPMap.at(op); const auto stream = LaunchKernel::resolve_device(t.device()); - - if (size_bytes <= m_max_bytes) { // use internal buffer - const auto buf_ptr = m_sym_mem.get(); - const auto need_memcpy = (buf_ptr != data_ptr); - if (need_memcpy) { - CUDA_CHECK(::cudaMemcpyAsync(buf_ptr, data_ptr, size_bytes, - ::cudaMemcpyDeviceToDevice, stream)); - } - NCCL_CHECK(::ncclAllReduce( - /*sendbuff=*/buf_ptr, - /*recvbuff=*/buf_ptr, - /*count=*/size_dim, - /*datatype=*/dtype, - /*op=*/reduce_op, - /*comm=*/m_comm.get(), - /*stream=*/stream)); - if (need_memcpy) { - CUDA_CHECK(::cudaMemcpyAsync(data_ptr, buf_ptr, size_bytes, - ::cudaMemcpyDeviceToDevice, stream)); - } - } else { - NCCL_CHECK(::ncclAllReduce( - /*sendbuff=*/data_ptr, - /*recvbuff=*/data_ptr, - /*count=*/size_dim, - /*datatype=*/dtype, - /*op=*/reduce_op, - /*comm=*/m_comm.get(), - /*stream=*/stream)); - } + NCCL_CHECK(::ncclAllReduce( + /*sendbuff=*/data_ptr, + /*recvbuff=*/data_ptr, + /*count=*/size_dim, + /*datatype=*/dtype, + /*op=*/reduce_op, + /*comm=*/m_comm.get(), + /*stream=*/stream)); } auto all_gather(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) const @@ -160,18 +127,13 @@ public: /*stream=*/stream)); } - auto get_buffer() const -> void * { return m_sym_mem.get(); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("freetoken.NCCLWrapper", NCCLWrapper, tvm::ffi::Object); private: int m_rank; int m_world_size; - size_t m_max_bytes; shared_obj m_comm; - shared_ptr m_sym_mem; - shared_obj m_win; }; TVM_FFI_STATIC_INIT_BLOCK() { @@ -179,8 +141,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::ObjectDef() .def(refl::init(), "__init__") .def("all_reduce", &NCCLWrapper::all_reduce) - .def("all_gather", &NCCLWrapper::all_gather) - .def("get_buffer", &NCCLWrapper::get_buffer); + .def("all_gather", &NCCLWrapper::all_gather); } TVM_FFI_DLL_EXPORT_TYPED_FUNC(create_nccl_uid, &create_uid); diff --git a/python/freetoken/kernel/pynccl.py b/python/freetoken/kernel/pynccl.py index 23ea573..d2fe424 100644 --- a/python/freetoken/kernel/pynccl.py +++ b/python/freetoken/kernel/pynccl.py @@ -18,8 +18,6 @@ class PyNCCLCommunicator: def all_reduce(self, input: torch.Tensor, op: Literal["sum"]) -> None: ... @abstractmethod def all_gather(self, output: torch.Tensor, input: torch.Tensor) -> None: ... - @abstractmethod - def get_buffer(self) -> int: ... else: PyNCCLCommunicator = Any diff --git a/python/freetoken/kernel/triton/dsv4/indexer.py b/python/freetoken/kernel/triton/dsv4/indexer.py index 3d724e7..86040b0 100644 --- a/python/freetoken/kernel/triton/dsv4/indexer.py +++ b/python/freetoken/kernel/triton/dsv4/indexer.py @@ -117,7 +117,7 @@ def indexer_logits( @triton.jit def _indexer_decode_logits_kernel( - q_ptr, w_ptr, pool_ptr, snap_ptr, valid_ptr, out_ptr, + q_ptr, w_ptr, pool_ptr, snap_ptr, rows_ptr, valid_ptr, out_ptr, N_STAGE, RATIO, stride_qb, stride_qh, stride_qd, stride_wb, stride_wh, @@ -156,8 +156,11 @@ def _indexer_decode_logits_kernel( h_mask = offs_h < H # block b -> its compressed row: full_loc(b * RATIO) // RATIO off the decode snapshot. + # There may be several query rows per request (DSpark target verification), so + # pid_b is a QUERY row and rows_ptr maps it to the corresponding SNAPSHOT row. # Masked-off lanes read snapshot column 0 and are dropped by ``t_mask`` below. - snap = tl.load(snap_ptr + pid_b * stride_sb + (offs_t * RATIO) * stride_sw, + snap_row = tl.load(rows_ptr + pid_b) + snap = tl.load(snap_ptr + snap_row * stride_sb + (offs_t * RATIO) * stride_sw, mask=t_mask, other=0) rows = tl.maximum(snap, 0) // RATIO @@ -179,8 +182,9 @@ def indexer_decode_logits( q: torch.Tensor, # [B, H, D] bf16, one query per request weights: torch.Tensor, # [B, H] idx_pool: torch.Tensor, # [R, D] bf16, this layer's paged indexer keys - full_snap: torch.Tensor, # [B, W] int64, the decode full-loc snapshot - valid: torch.Tensor, # [B] live compressed block count per row + full_snap: torch.Tensor, # [R, W] int64, one snapshot row per request + rows: torch.Tensor, # [Q] query row -> snapshot row + valid: torch.Tensor, # [Q] live compressed block count per query n_stage: int, # static staged width (the buffer the top-k scans) ratio: int, out: torch.Tensor | None = None, @@ -199,6 +203,7 @@ def indexer_decode_logits( """ B, H, D = q.shape assert weights.shape == (B, H), (weights.shape, (B, H)) + assert rows.shape == (B,), (rows.shape, (B,)) assert idx_pool.shape[1] == D, (idx_pool.shape, D) assert D == triton.next_power_of_2(D), f"index_head_dim must be pow2, got {D}" @@ -209,7 +214,7 @@ def indexer_decode_logits( BLOCK_T = 64 _indexer_decode_logits_kernel[(B, triton.cdiv(n_stage, BLOCK_T))]( - q, weights, idx_pool, full_snap, valid, out, + q, weights, idx_pool, full_snap, rows, valid, out, n_stage, ratio, q.stride(0), q.stride(1), q.stride(2), weights.stride(0), weights.stride(1), diff --git a/python/freetoken/kernel/triton/dsv4/norm.py b/python/freetoken/kernel/triton/dsv4/norm.py index b1c4399..523448a 100644 --- a/python/freetoken/kernel/triton/dsv4/norm.py +++ b/python/freetoken/kernel/triton/dsv4/norm.py @@ -52,4 +52,48 @@ def rms_norm(x: torch.Tensor, weight: torch.Tensor | None, eps: float) -> torch. return out.reshape(x.shape) -__all__ = ["rms_norm"] +@triton.jit +def _inv_rms_kernel(x_ptr, o_ptr, D, eps, stride_xm, BLOCK_D: tl.constexpr): + row = tl.program_id(0) + acc = tl.zeros((BLOCK_D,), dtype=tl.float32) + for off in range(0, D, BLOCK_D): + offs = off + tl.arange(0, BLOCK_D) + v = tl.load(x_ptr + row * stride_xm + offs, mask=offs < D, other=0.0).to(tl.float32) + acc += v * v + tl.store(o_ptr + row, tl.rsqrt(tl.sum(acc, axis=0) / D + eps)) + + +def inv_rms(x: torch.Tensor, eps: float) -> torch.Tensor: + """``rsqrt(mean(x^2, -1) + eps)`` -> ``[..., 1]`` fp32, one pass, no fp32 temp. + + :func:`rms_norm` applies the scale and writes a full-size output; callers that need + the *scalar* (DSV4's hyper-connection pre-norm multiplies it into a [.., mix_hc] + tensor, not into x) were spelling it ``x.float().square().mean(-1)``, which + materialises a second full fp32 copy of the hidden state purely to reduce it away. + At hc_dim 16384, T=8192 that is a 537 MB write plus a 537 MB read for a [T, 1] + result. Reading the bf16 source instead of its fp32 upcast halves the bytes again; + the upcast is exact, so the sum of squares is unchanged. + + RTX 6000 Ada, hc_dim 16384: 1.944 ms -> 0.304 ms at T=8192 (6.4x), 0.020 -> 0.013 ms + at T=64. The large end lands at ~880 GB/s against a 960 GB/s peak -- the memory roof. + + NOT bit-identical to ``square().mean()``: the reduction order differs, and ATen's + order is not reproducible anyway (it varies with M -- a tree/2 fold matches at M=4 + and nothing matches at M=64, because the block/grid split is chosen per shape). The + accuracy is equivalent, not merely close: against an fp64 reference the mean relative + error is 3.65e-08 here vs 3.60e-08 for ATen at M=8192, and 3.60e-08 vs 3.63e-08 at + M=4096 -- i.e. it wins at some shapes and loses at others, within 8% of ATen's own + distance from the truth. (``linalg.vector_norm`` is the same speed but consistently + ~23% worse, at 4.47e-08, because sqrt-then-square rounds twice.) + """ + D = x.shape[-1] + x2d = x.reshape(-1, D) + out = torch.empty(x2d.shape[0], dtype=torch.float32, device=x.device) + _inv_rms_kernel[(x2d.shape[0],)]( + x2d, out, D, eps, x2d.stride(0), + BLOCK_D=min(2048, triton.next_power_of_2(D)), num_warps=8, + ) + return out.view(*x.shape[:-1], 1) + + +__all__ = ["rms_norm", "inv_rms"] diff --git a/python/freetoken/kvcache/dsv4_cost_model.py b/python/freetoken/kvcache/dsv4_cost_model.py index 2c084af..58f144a 100644 --- a/python/freetoken/kvcache/dsv4_cost_model.py +++ b/python/freetoken/kvcache/dsv4_cost_model.py @@ -80,7 +80,7 @@ def dsv4_cache_per_page(args, swa_ratio: float, P: int = 128) -> int: idx_b = _index_bytes(args) total = 0 - for ratio in tuple(args.compress_ratios)[: args.n_layers]: + for ratio in args.layer_compress_ratios: # Window tier exists on EVERY layer (all-sliding), scaled by swa_ratio. total += round(swa_ratio * P) * kv_b if ratio == 0: @@ -107,7 +107,7 @@ def dsv4_kv_unit_bytes(args, P: int = 128) -> int: kv_b = _kv_bytes(args) idx_b = _index_bytes(args) per_page = P * _INT64_BYTES # full_to_window map: one int64 slot per full token - for ratio in tuple(args.compress_ratios)[: args.n_layers]: + for ratio in args.layer_compress_ratios: if ratio == 0: continue per_page += (P // ratio) * kv_b # compressed KV @@ -122,7 +122,7 @@ def dsv4_window_unit_bytes(args, P: int = 128) -> int: ``state_slots = n_win_pages * ring_size``). Independent of ``swa_ratio`` -- the ratio only sets how many window tokens exist, not the per-token cost. This is ``swa_bytes_per_token``.""" kv_b = _kv_bytes(args) - ratios = tuple(args.compress_ratios)[: args.n_layers] + ratios = args.layer_compress_ratios per_page = len(ratios) * P * kv_b # window KV: P slots per page, every layer for ratio in ratios: if ratio == 0: @@ -174,7 +174,7 @@ def dsv4_pool_sizes( state_slots: list[int | None] = [] ring_sizes: list[int | None] = [] idx_state_slots: list[int | None] = [] - for ratio in tuple(args.compress_ratios)[: args.n_layers]: + for ratio in args.layer_compress_ratios: if ratio == 0: cmp_blocks.append(None) idx_blocks.append(None) @@ -213,7 +213,7 @@ def dsv4_pool_bytes(sizes: DSV4PoolSizes, args, n_scratch: int = 1) -> int: +1 ring scratch row, +1 mapping sentinel row).""" kv_b = _kv_bytes(args) idx_b = _index_bytes(args) - ratios = tuple(args.compress_ratios)[: args.n_layers] + ratios = args.layer_compress_ratios total = len(ratios) * sizes.n_win_slots * kv_b # window pool, every layer total += (sizes.full_token + 1) * _INT64_BYTES # full_to_window (+ sentinel row) diff --git a/python/freetoken/kvcache/dsv4_paged_pool.py b/python/freetoken/kvcache/dsv4_paged_pool.py index 3ce23ee..0e46b03 100644 --- a/python/freetoken/kvcache/dsv4_paged_pool.py +++ b/python/freetoken/kvcache/dsv4_paged_pool.py @@ -167,13 +167,13 @@ def __init__( self._device = device self._dtype = dtype self.P = P - self._n_layers = args.n_layers + # Every layer that OWNS KV, target and dSpark draft alike. The draft layers' + # ids continue past the target's, so a pool sized to n_layers alone leaves them + # indexing past the end of window_pool the moment one stores its KV. + self.compress_ratios = args.layer_compress_ratios + self._n_layers = len(self.compress_ratios) self.head_dim = args.head_dim self.index_head_dim = args.index_head_dim - # The checkpoint can ship one extra trailing ratio (44 for 43 layers); the model only uses - # the first n_layers, so truncate to match. - self.compress_ratios = tuple(args.compress_ratios)[: self._n_layers] - assert len(self.compress_ratios) == self._n_layers # Scratch rows appended to each cmp/idx pool tensor BEYOND the allocator's capacity # (never handed out): batched decode routes each row whose compressed block did NOT # complete this step to its own scratch row ``cmp_scratch_base + row`` (a discarded diff --git a/python/freetoken/llm/llm.py b/python/freetoken/llm/llm.py index 1d4aaef..92b215c 100644 --- a/python/freetoken/llm/llm.py +++ b/python/freetoken/llm/llm.py @@ -99,8 +99,10 @@ def offline_send_result(self, reply: List[BaseTokenizerMsg]) -> None: continue assert isinstance(msg, DetokenizeMsg) status = self.status_map[msg.uid] - if not (msg.finished and msg.next_token in self.eos_token_ids): - status.output_ids.append(msg.next_token) + token_ids = msg.token_ids if msg.token_ids is not None else [msg.next_token] + if msg.finished and token_ids[-1] in self.eos_token_ids: + token_ids = token_ids[:-1] + status.output_ids.extend(token_ids) def generate( self, diff --git a/python/freetoken/message/tokenizer.py b/python/freetoken/message/tokenizer.py index 33b75c7..273a229 100644 --- a/python/freetoken/message/tokenizer.py +++ b/python/freetoken/message/tokenizer.py @@ -29,6 +29,12 @@ class DetokenizeMsg(BaseTokenizerMsg): uid: int next_token: int finished: bool + # A speculative verify can accept several tokens in one engine step. Keep + # ``next_token`` as the final token for wire/backward compatibility, while + # carrying the complete ordered block here. This mirrors vLLM's + # EngineCoreOutput.new_token_ids contract and lets the detokenizer advance + # through the block exactly once. + token_ids: list[int] | None = None finish_reason: str | None = None # The stop string that ended generation (if any), so the detokenizer can trim it # and everything after it from the final output. diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py index f6105e1..3998ae0 100644 --- a/python/freetoken/models/config.py +++ b/python/freetoken/models/config.py @@ -250,6 +250,9 @@ class ModelConfig: # sparse MoE block (GLM-4: 3). Experts (and the offload cache) therefore only exist # for layers ``[first_k_dense_replace, num_layers)``. first_k_dense_replace: int = 0 + # MoE blocks beyond the target stack that still need expert banks and offload-cache + # slots -- a speculative drafter's layers, indexed after the target's. + extra_moe_layers: int = 0 # Always-on shared expert(s) added to every MoE layer's output (GLM-4: 1). n_shared_experts: int = 0 # Selected-expert weights are multiplied by this after renormalization (GLM-4: 2.5). @@ -302,8 +305,13 @@ def num_moe_layers(self) -> int: Models with leading dense layers (``first_k_dense_replace`` > 0, e.g. GLM-4) only store experts for the trailing layers; everything else has all layers MoE. + + ``extra_moe_layers`` covers MoE blocks that are NOT part of the target stack -- + today a speculative drafter's layers, which own routed experts of their own and + are indexed after the target's. The expert banks and the offload cache must agree + on this count, so it lives beside ``num_layers`` rather than being re-derived. """ - return self.num_layers - self.first_k_dense_replace + return self.num_layers - self.first_k_dense_replace + self.extra_moe_layers @property def is_multimodal(self) -> bool: diff --git a/python/freetoken/models/deepseek_v4/args.py b/python/freetoken/models/deepseek_v4/args.py index b75d280..0b15eca 100644 --- a/python/freetoken/models/deepseek_v4/args.py +++ b/python/freetoken/models/deepseek_v4/args.py @@ -69,15 +69,68 @@ class DeepseekV4Args: hc_sinkhorn_iters: int = 20 hc_eps: float = 1e-6 + # ----- dSpark (semi-autoregressive speculative decoding) ----- + # The checkpoint's ``mtp.{0..n_mtp_layers-1}`` stack. Unlike a one-token-per-step MTP + # head, dSpark proposes a whole block: ``dspark_block_size`` tokens per draft pass, + # scored by a Markov head and gated by a confidence head. Each mtp layer is a FULL + # DSV4 block, so enabling it costs n_mtp_layers x n_routed_experts more expert banks. + dspark_block_size: int = 0 + dspark_markov_rank: int = 0 + dspark_noise_token_id: int = -1 + # Which target layers feed the drafter its hidden state. + dspark_target_layer_ids: Tuple[int, ...] = () + # Runtime opt-in, set by the server from --speculative-dspark. The fields above only + # describe what the CHECKPOINT ships; this says whether to pay for it. Off by default + # because the drafter's own routed experts enlarge the host banks and the GPU cache. + dspark_enabled: bool = False + def __post_init__(self) -> None: # JSON lists -> tuple so the dataclass stays hashable / immutable-ish. if isinstance(self.compress_ratios, list): self.compress_ratios = tuple(self.compress_ratios) + if isinstance(self.dspark_target_layer_ids, list): + self.dspark_target_layer_ids = tuple(self.dspark_target_layer_ids) @property def nope_head_dim(self) -> int: return self.head_dim - self.rope_head_dim + @property + def has_dspark(self) -> bool: + """Does this checkpoint ship a usable dSpark drafter?""" + return ( + self.dspark_block_size > 1 + and self.n_mtp_layers > 0 + and self.dspark_markov_rank > 0 + ) + + @property + def n_draft_layers(self) -> int: + """Drafter layers actually built. Zero unless dSpark is both shipped and enabled.""" + return self.n_mtp_layers if (self.has_dspark and self.dspark_enabled) else 0 + + @property + def layer_compress_ratios(self) -> Tuple[int, ...]: + """Compression ratio per KV-owning layer, target layers then draft layers. + + ``compress_ratios`` in the checkpoint is longer than ``n_layers`` and describes + only the target. The dSpark draft layers ship no compressor or indexer, so they + append zeros: they own a sliding-window tier and nothing else. Every per-layer + KV structure -- pool sizing, the window/compressed/indexer regions -- should walk + THIS list rather than slicing ``compress_ratios``, so draft layers get their KV. + """ + return tuple(self.compress_ratios)[: self.n_layers] + (0,) * self.n_draft_layers + + @property + def n_moe_layers(self) -> int: + """MoE layers the expert banks and the offload cache must cover. + + The drafter's layers are appended AFTER the target's, so a draft layer keeps the + cache-facing id ``n_layers + k``. Every layer-indexed structure (host banks, GPU + slot cache, KV pools) then addresses target and draft layers the same way. + """ + return self.n_layers + self.n_draft_layers + def _config_path(model_path: str) -> str: """Locate the authors' ModelArgs JSON inside the checkpoint directory.""" @@ -94,6 +147,26 @@ def _config_path(model_path: str) -> str: ) +_DSPARK_ENABLED = False + + +def set_dspark_enabled(enabled: bool) -> None: + """Record the run's dSpark choice for every later ``load_args``. + + ``load_args`` reads the checkpoint, which only says what dSpark weights EXIST -- + never whether this run wants them. Several call sites (config resolution, the weight + reader, the expert-bank builder) each build their own args from that file, so the + runtime choice has to live beside the file rather than in one of those instances. + Without this the weight reader silently skips the drafter the model just built. + """ + global _DSPARK_ENABLED + _DSPARK_ENABLED = enabled + + +def dspark_enabled() -> bool: + return _DSPARK_ENABLED + + def load_args(model_path: str, **overrides) -> DeepseekV4Args: """Build :class:`DeepseekV4Args` from the checkpoint's ``inference/config.json``. @@ -104,6 +177,7 @@ def load_args(model_path: str, **overrides) -> DeepseekV4Args: raw = json.load(f) valid = {f.name for f in fields(DeepseekV4Args)} kwargs = {k: v for k, v in raw.items() if k in valid} + kwargs.setdefault("dspark_enabled", _DSPARK_ENABLED) kwargs.update(overrides) return DeepseekV4Args(**kwargs) diff --git a/python/freetoken/models/deepseek_v4/attention.py b/python/freetoken/models/deepseek_v4/attention.py index 357585e..21fda2d 100644 --- a/python/freetoken/models/deepseek_v4/attention.py +++ b/python/freetoken/models/deepseek_v4/attention.py @@ -8,6 +8,7 @@ from torch import nn from freetoken.core import get_global_ctx +from freetoken.distributed import DistributedCommunicator from freetoken.kernel.triton.dsv4.fp8_linear import act_quant_fp8_inplace from freetoken.kernel.triton.dsv4.norm import rms_norm @@ -15,6 +16,7 @@ from .compress import Compressor, Indexer from .layers import Linear, RMSNorm, get_compress_topk_idxs, get_window_topk_idxs from .ops import apply_rotary_emb, apply_rotary_emb_decode, get_freqs_cis +from .parallel import div_tp, tp_size class Attention(nn.Module): @@ -26,21 +28,39 @@ class Attention(nn.Module): ``window_pool[L]`` / ``cmp_pool[L]`` inside the kernel -- no per-forward staging slab. Both prefill and decode use the same paged kernel.""" - def __init__(self, layer_id: int, args: DeepseekV4Args): + def __init__(self, layer_id: int, args: DeepseekV4Args, compress_ratio: int | None = None): + # ``compress_ratio`` overrides the per-layer table. The dSpark draft layers pass 0: + # their checkpoint ships no compressor or indexer weights, because a draft block + # attends only over the sliding window the target already filled. super().__init__() self.layer_id = layer_id self.dim = args.dim - self.n_heads = args.n_heads 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.n_groups = args.o_groups self.window_size = args.window_size - self.compress_ratio = args.compress_ratios[layer_id] + self.compress_ratio = ( + args.compress_ratios[layer_id] if compress_ratio is None else compress_ratio + ) self.eps = args.norm_eps + # dSpark draft layers set this. A drafted BLOCK is proposed at once, so each of + # its queries may see the whole block, not just the positions before it -- that + # is what makes the proposal semi-autoregressive instead of five serial steps. + # Expressed through the sparse top-k: the causal cutoff moves from "my own + # position" to "the end of my block". Context positions precede the block, so + # they stay visible either way and need no special case. + self.non_causal = False + + # Head parallelism: a rank owns whole o_groups, and therefore whole heads. The + # per-group head width (wo_a_k) is a property of one group, so it does NOT shard. + self.tp_size = tp_size() + self.n_heads = div_tp(args.n_heads, "n_heads") + self.n_groups = div_tp(args.o_groups, "o_groups") + self._comm = DistributedCommunicator() if self.tp_size > 1 else None self.attn_sink = nn.Parameter(torch.empty(self.n_heads, dtype=torch.float32), requires_grad=False) + # wq_a / wkv / kv_norm stay replicated: the MLA latent KV is shared by every head. self.wq_a = Linear(self.dim, self.q_lora_rank, kind="fp8") self.q_norm = RMSNorm(self.q_lora_rank, self.eps) self.wq_b = Linear(self.q_lora_rank, self.n_heads * self.head_dim, kind="fp8") @@ -48,8 +68,9 @@ def __init__(self, layer_id: int, args: DeepseekV4Args): self.kv_norm = RMSNorm(self.head_dim, self.eps) # wo_a: dequantized to bf16 (reference runs a bf16 grouped-output einsum). wo_a_rows = self.n_groups * args.o_lora_rank - wo_a_k = self.n_heads * self.head_dim // self.n_groups + wo_a_k = args.n_heads * self.head_dim // args.o_groups self.wo_a = nn.Parameter(torch.empty(wo_a_rows, wo_a_k, dtype=torch.bfloat16), requires_grad=False) + # Row-parallel: each rank consumes its own groups and contributes a partial sum. self.wo_b = Linear(self.n_groups * args.o_lora_rank, self.dim, kind="fp8") self.softmax_scale = self.head_dim ** -0.5 @@ -99,10 +120,16 @@ def reset(self) -> None: def _wo(self, o: torch.Tensor, bsz: int, seqlen: int) -> torch.Tensor: + # Under TP, ``o`` holds only this rank's heads, so the grouped einsum and wo_b + # produce a PARTIAL sum over the full output dim. One all-reduce completes it -- + # the single attention-side collective per block. o = o.reshape(bsz, seqlen, self.n_groups, -1) wo_a = self.wo_a.view(self.n_groups, self.o_lora_rank, -1) o = torch.einsum("bsgd,grd->bsgr", o, wo_a).flatten(2) - return self.wo_b(o) + o = self.wo_b(o) + if self._comm is not None: + o = self._comm.all_reduce(o) + return o def _prefill_segment(self, x_seg, qr_seg, kv_seg, ti: int, start_pos: int, n: int): """One request's prefill work inside a (possibly ragged) batch: persist its window KV, @@ -119,18 +146,35 @@ def _prefill_segment(self, x_seg, qr_seg, kv_seg, ti: int, start_pos: int, n: in slots = self.attn.window_slots_of(ti, start_pos, end) self.attn.store_window(kv_seg, self.layer_id, slots) + if self.non_causal and start_pos == 0: + # A draft block always continues an existing context, so it never takes the + # cold path. Refuse rather than silently emit a causal grid. + raise ValueError( + "dSpark non-causal attention needs a resumed segment (start_pos > 0); " + "the drafter is not meant to run a cold prompt" + ) if start_pos == 0: win_cols = get_window_topk_idxs(win, 1, n, 0).to(device) # natural width min(n, win); the caller pads to the batch-uniform width win_global = self.attn.win_cols_to_global(win_cols, slots) else: - # Candidates span [w_lo, end): each new query's 128-sliding window over the retained - # prefix plus the new tokens, causal-masked to -1 past its own position. - w_lo = max(0, start_pos - win + 1) + # Target attention uses its ordinary 128-token causal window. A DSpark + # draft query instead sees 128 CONTEXT tokens plus every token in its + # parallel query block (paper section 3.1; same sparse list as vLLM). + w_lo = max( + 0, + start_pos - win if self.non_causal else start_pos - win + 1, + ) ws_pool = self.attn.window_slots_of(ti, w_lo, end) - abs_p = start_pos + torch.arange(n, device=device).unsqueeze(1) - cand = (abs_p - win + 1).clamp(min=w_lo) + torch.arange(win, device=device) - win_cols = torch.where(cand > abs_p, -1, cand - w_lo).unsqueeze(0) + # One implementation of the masking rule, shared with the tests that pin it. + # Non-causal (dSpark draft layers) gives every query in the block the SAME + # window, ending at the block's last position, so a drafted query sees the + # whole block plus the context before it. + from .dspark import window_cols_for_block + + win_cols = window_cols_for_block( + start_pos, n, win, self.non_causal, device=device + ).unsqueeze(0) win_global = self.attn.win_cols_to_global(win_cols, ws_pool) if not ratio: @@ -148,6 +192,36 @@ def _prefill_segment(self, x_seg, qr_seg, kv_seg, ti: int, start_pos: int, n: in else get_compress_topk_idxs(ratio, 1, n, 0, 0).to(device) ) self.compressor(x_seg, 0, slots, ti=ti) + elif ( + self.non_causal + or getattr(get_global_ctx().batch, "speculative", False) + or start_pos % self.P != 0 + ): + # A speculative block starts wherever generation reached, which is mid-page + # 127 times out of 128. The compressor's EXTEND path cannot serve that: it + # reduces tokens with ``start_pos // ratio`` arithmetic and asserts + # 128-alignment, because it seeds the carry from a matched tail PAGE. + # + # Its per-token path has no such assumption -- it reads the carry from the + # previous token's slot, advances it, and writes it to this token's slot, + # which is exactly what vLLM's save_partial_states kernel does (one program + # per token, ape chosen by position % compress_ratio, no alignment anywhere). + # So drive the compressor per token and leave attention and MoE batched: + # the compressor is cheap next to the expert fetch this block exists to + # amortize. + self._advance_compressor_per_token(x_seg, start_pos, n, slots, tail_ws, ti) + # Run the INDEXER here too. Falling back to _compress_topk_extend chooses + # compressed blocks positionally, which silently replaces DeepSeek-V4's + # learned sparse selection with a fixed one -- so the target attends to the + # wrong blocks and its own logits are wrong. Since a speculative block is + # unaligned 127 times out of 128, that was every verify. + blocks = ( + self.indexer.extend_unaligned( + x_seg, qr_seg, start_pos, 0, slots, tail_ws, ti + ) + if self.indexer is not None + else self._compress_topk_extend(n, start_pos, end, 0, device, 1) + ) else: blocks = ( self.indexer.extend(x_seg, qr_seg, start_pos, 0, slots, tail_ws, ti) @@ -157,6 +231,27 @@ def _prefill_segment(self, x_seg, qr_seg, kv_seg, ti: int, start_pos: int, n: in self.compressor(x_seg, start_pos, slots, tail_window_slot=tail_ws, ti=ti) return win_global, self.attn.blocks_to_global(blocks, ratio, ti=ti) + def _advance_compressor_per_token( + self, x_seg, start_pos: int, n: int, slots: torch.Tensor, tail_ws: int | None, + ti: int, + ) -> None: + """Walk the compressor one token at a time across an unaligned segment. + + Each step reads the carry from the PREVIOUS token's window slot and writes the + advanced carry to this token's -- the same movement a decode step makes, so no + block-boundary arithmetic is involved and any ``start_pos`` is valid. + """ + device = x_seg.device + rows = torch.zeros(1, dtype=torch.int64, device=device) + prev = torch.full((1,), tail_ws if tail_ws is not None else int(slots[0]), + dtype=slots.dtype, device=device) + for t in range(n): + cur = slots[t:t + 1] + pos_t = torch.full((1,), start_pos + t, dtype=torch.int64, device=device) + # ti, not the decode snapshot: this runs inside a prefill, where none exists. + self.compressor.decode_step(x_seg[:, t], pos_t, prev, cur, rows, ti=ti) + prev = cur + def forward_ragged(self, x, segments, flat_positions): """Ragged batched prefill (cu_seqlens). ``x`` is [1, T, dim] -- the requests' NEW token streams concatenated (NO padding); ``segments`` is [(offset, n, table_idx, start_pos)] @@ -212,7 +307,16 @@ def forward_ragged(self, x, segments, flat_positions): # they shift the kernel's window/compressed tile split, and that rounding-order change # is enough to flip a greedy near-tie. bs>1 pads every part to win (the historical # batched behavior). - n_window = win if len(segments) > 1 else win_parts[0].shape[-1] + # Ordinary target segments keep the historical batch-uniform ``win`` width. + # DSpark's non-causal query block is wider: trailing ``win`` context columns + # PLUS all gamma query columns. Labeling only the first 128 columns as window + # in a multi-request draft misclassifies/drops the query-query suffix even + # though the sparse indices themselves correctly contain all 133 columns. + n_window = ( + max(part.shape[-1] for part in win_parts) + if self.non_causal + else (win if len(segments) > 1 else win_parts[0].shape[-1]) + ) flat = [] for i, (off, n, ti, _start) in enumerate(segments): wg = win_parts[i] @@ -322,3 +426,77 @@ def decode_step( ) apply_rotary_emb_decode(o[..., -rd:], freqs_t, True) # per-row position freqs (inverse) return self._wo(o, B, 1) + + def verify_block( + self, + x: torch.Tensor, + pos: torch.Tensor, + rows: torch.Tensor, + cmp_stage_cap: int, + num_reqs: int, + span: int, + wctx, + ) -> torch.Tensor: + """Fixed-shape causal target verification using decode-addressed kernels. + + Large projections, MoE work and sparse attention remain batched over all + verification rows. Stateful compressors advance in request-major token order + so each token reads its predecessor's new carry, with no host-derived position + or address embedded in the CUDA graph. + """ + n = x.shape[1] + if n != num_reqs * span: + raise ValueError(f"verify block has {n} rows, expected {num_reqs} * {span}") + win, ratio, rd = self.window_size, self.compress_ratio, self.rope_head_dim + flat = x.reshape(n, self.dim) + xq = flat.unsqueeze(1) + freqs_t = self.freqs_cis.index_select(0, pos) + + qr = q = self.q_norm(self.wq_a(xq)) + q = self.wq_b(q).unflatten(-1, (self.n_heads, self.head_dim)) + q = rms_norm(q, None, self.eps) + apply_rotary_emb_decode(q[..., -rd:], freqs_t) + + kv = self.kv_norm(self.wkv(xq)) + apply_rotary_emb_decode(kv[..., -rd:], freqs_t) + act_quant_fp8_inplace(kv[..., :-rd], 64) + + window_slots, prev_window_slots, window_slots_topk = wctx + self.attn.store_window(kv.view(n, -1), self.layer_id, window_slots) + n_cmp_stage = (cmp_stage_cap + 1) // ratio if ratio else 0 + cmp_counts = None + if ratio: + # This order is also the carry-journal order consumed after acceptance. + for r in range(num_reqs): + base = r * span + row = rows[base:base + 1] + for s in range(span): + i = base + s + self.compressor.decode_step( + flat[i:i + 1], pos[i:i + 1], + prev_window_slots[i:i + 1], window_slots[i:i + 1], row, + ) + valid = (pos + 1) // ratio + if self.indexer is not None: + blocks = self.indexer.verify_block( + flat, qr.reshape(n, -1), pos, + prev_window_slots, window_slots, rows, n_cmp_stage, + num_reqs, span, + ) + else: + blk = torch.arange(n_cmp_stage, device=x.device) + blocks = torch.where( + blk[None, :] < valid[:, None], blk[None, :], -1 + ).view(n, 1, n_cmp_stage) + compress_idxs = self.attn.blocks_to_global(blocks, ratio, rows=rows) + topk_idxs = torch.cat([window_slots_topk, compress_idxs], dim=-1) + cmp_counts = valid.clamp(max=compress_idxs.shape[-1]).to(torch.int32).view(n, 1) + else: + topk_idxs = window_slots_topk + + o = self.attn.attend( + q, self.layer_id, topk_idxs.int(), win, self.attn_sink, + self.softmax_scale, cmp_counts=cmp_counts, has_compression=bool(ratio), + ) + apply_rotary_emb_decode(o[..., -rd:], freqs_t, True) + return self._wo(o, n, 1).view(1, n, self.dim) diff --git a/python/freetoken/models/deepseek_v4/compress.py b/python/freetoken/models/deepseek_v4/compress.py index 4d578f3..ac2cf3f 100644 --- a/python/freetoken/models/deepseek_v4/compress.py +++ b/python/freetoken/models/deepseek_v4/compress.py @@ -293,7 +293,7 @@ def extend(self, x, start_pos: int, window_slots: torch.Tensor, tail_window_slot def decode_step( self, x: torch.Tensor, pos: torch.Tensor, prev_window_slots: torch.Tensor, - window_slots: torch.Tensor, rows: torch.Tensor, + window_slots: torch.Tensor, rows: torch.Tensor, ti: int | None = None, ) -> None: """Batched single-token compressor update (EAGER/graph; ``pos``/``prev_window_slots``/ ``window_slots`` are GPU int tensors ``[B]``; ``rows`` is the LOCAL row index [B] into the @@ -338,9 +338,21 @@ def decode_step( compressed = gated_pool(ks, ss, dtype) # Persist the advanced carry block to THIS token's page (per row, disjoint blocks); the # next step reads it from here. Within a page prev==cur, so this is a self-consistent roll. + carry_state = torch.cat([ks, ss], dim=-1) self.attn.write_carry_blocks( - self.layer_id, self.tier, window_slots, self.ring_size, torch.cat([ks, ss], dim=-1) + self.layer_id, self.tier, window_slots, self.ring_size, carry_state ) + batch = get_global_ctx().batch + if getattr(batch, "speculative", False): + journal = batch.spec_carry_states + if journal is None: + raise RuntimeError("a speculative verify did not initialize its carry journal") + key = (self.layer_id, self.tier, self.ring_size) + record = getattr(journal, "record", None) + if record is None: + journal.setdefault(key, []).append(carry_state) + else: + record(key, carry_state) compressed = self.norm(compressed) freqs_t = self.freqs_cis.index_select(0, (pos + 1 - ratio).clamp_min(0)) # [B, rd//2] @@ -358,7 +370,7 @@ def decode_step( # it). Read off the full-loc SNAPSHOT so a concurrent next-batch allocate cannot redirect # this write (overlap safety); real rows resolve to a live row, dummy rows to reserved 0. cmp_dst = self.attn.decode_compress_rows( - rows, pos, ratio, self.layer_id, self.tier, should.view(B) + rows, pos, ratio, self.layer_id, self.tier, should.view(B), ti=ti ) self.attn.scatter_compressed(self.layer_id, self.tier, cmp_dst, compressed.view(B, -1)) @@ -446,6 +458,57 @@ def extend(self, x, qr, start_pos, offset, window_slots, tail_window_slot, ti: i topk=self.index_topk, offset=offset, ) + def extend_unaligned(self, x, qr, start_pos, offset, window_slots, tail_window_slot, ti: int = 0): + """``extend``, for a segment that does NOT start on a compress-ratio boundary. + + ``extend`` cannot serve one: it advances its compressor through + ``Compressor.forward``, which for ``start_pos > 0`` delegates to + ``Compressor.extend`` and asserts 128-alignment, because that path seeds the + carry from a matched tail PAGE. + + A speculative block starts wherever generation reached, so it is unaligned 127 + times out of 128 -- and without this the caller had no indexer to fall back on + and chose compressed blocks POSITIONALLY instead. That silently replaces + DeepSeek-V4's learned sparse selection with a fixed one, so the TARGET attends + to the wrong blocks and its own logits are wrong. It never raises; it just + degenerates the output. + + Everything below the compressor advance is identical to ``extend``: only the way + the compressor is walked differs, one token at a time off the previous token's + window slot -- the same movement a decode step makes, with no block-boundary + arithmetic, so any ``start_pos`` is valid. + """ + bsz, seqlen, _ = x.size() + end = start_pos + seqlen + freqs_cis = self.freqs_cis[start_pos:end] + ratio, rd = self.compress_ratio, self.rope_head_dim + q = self.wq_b(qr).unflatten(-1, (self.n_heads, self.head_dim)) + apply_rotary_emb(q[..., -rd:], freqs_cis) + q = hadamard_transform(q) + fp4_act_quant_inplace(q, 32) + + # The one difference from extend(): walk the indexer's own compressor per token. + device = x.device + rows = torch.zeros(1, dtype=torch.int64, device=device) + prev = torch.full( + (1,), + tail_window_slot if tail_window_slot is not None else int(window_slots[0]), + dtype=window_slots.dtype, device=device, + ) + for t in range(seqlen): + cur = window_slots[t:t + 1] + pos_t = torch.full((1,), start_pos + t, dtype=torch.int64, device=device) + self.compressor.decode_step(x[:, t], pos_t, prev, cur, rows, ti=ti) + prev = cur + + weights = self.weights_proj(x) * (self.softmax_scale * self.n_heads ** -0.5) + keys = self.attn.indexer_keys(ti, end // ratio, ratio, self.layer_id, bsz) + scores = self.attn.indexer_prefill_logits(q, keys, weights) + return self.attn.indexer_select_prefill( + scores, start_pos=start_pos, seqlen=seqlen, ratio=ratio, + topk=self.index_topk, offset=offset, + ) + def decode_step( self, x: torch.Tensor, qr: torch.Tensor, pos: torch.Tensor, offset: int, prev_window_slots: torch.Tensor, window_slots: torch.Tensor, rows: torch.Tensor, @@ -475,8 +538,56 @@ def decode_step( index_score = self.attn.indexer_decode_scores( q.reshape(B, self.n_heads, self.head_dim), weights.reshape(B, self.n_heads), - valid, n_stage, ratio, self.layer_id, + valid, n_stage, ratio, self.layer_id, rows, ).view(B, 1, n_stage) return self.attn.indexer_select_decode( index_score, valid=valid, topk=self.index_topk, offset=offset ) + + def verify_block( + self, + x: torch.Tensor, + qr: torch.Tensor, + pos: torch.Tensor, + prev_window_slots: torch.Tensor, + window_slots: torch.Tensor, + rows: torch.Tensor, + n_stage: int, + num_reqs: int, + span: int, + ) -> torch.Tensor: + """Graph-safe Lightning Indexer for a fixed DSpark target block. + + Projection and scoring stay batched over every verification row. Only the + indexer's recurrent compressor walks in request-major token order, matching + ``extend_unaligned`` while positions and addresses remain device-resident. + """ + n, _ = x.shape + ratio, rd = self.compress_ratio, self.rope_head_dim + q = self.wq_b(qr).unflatten(-1, (self.n_heads, self.head_dim)) + apply_rotary_emb_decode( + q[..., -rd:].unsqueeze(1), self.freqs_cis.index_select(0, pos) + ) + q = hadamard_transform(q) + fp4_act_quant_inplace(q, 32) + + for r in range(num_reqs): + base = r * span + row = rows[base:base + 1] + for s in range(span): + i = base + s + self.compressor.decode_step( + x[i:i + 1], pos[i:i + 1], + prev_window_slots[i:i + 1], window_slots[i:i + 1], row, + ) + + weights = self.weights_proj(x) * (self.softmax_scale * self.n_heads ** -0.5) + valid = (pos + 1) // ratio + scores = self.attn.indexer_decode_scores( + q.reshape(n, self.n_heads, self.head_dim), + weights.reshape(n, self.n_heads), + valid, n_stage, ratio, self.layer_id, rows, + ).view(n, 1, n_stage) + return self.attn.indexer_select_decode( + scores, valid=valid, topk=self.index_topk, offset=0 + ) diff --git a/python/freetoken/models/deepseek_v4/config.py b/python/freetoken/models/deepseek_v4/config.py index 1005b8a..199082f 100644 --- a/python/freetoken/models/deepseek_v4/config.py +++ b/python/freetoken/models/deepseek_v4/config.py @@ -64,6 +64,11 @@ def parse_config(hf_config: Any) -> ModelConfig: base=args.rope_theta, scaling=rope_scaling, ), + # The dSpark drafter's layers own routed experts too, and are indexed after the + # target's. Without this the expert banks hold n_layers + n_draft_layers entries + # while the offload cache is built for n_layers, and the two assert against each + # other at startup. + extra_moe_layers=args.n_draft_layers, num_experts=args.n_routed_experts, num_experts_per_tok=args.n_activated_experts, moe_intermediate_size=args.moe_inter_dim, @@ -84,7 +89,9 @@ def parse_config(hf_config: Any) -> ModelConfig: attention_groups=( DSV4AttentionGroupConfig( name="dsv4", - layer_ids=tuple(range(args.n_layers)), + # Draft layers continue the target's ids and own a window tier each, so + # the group must cover them or they would have nowhere to store KV. + layer_ids=tuple(range(len(args.layer_compress_ratios))), num_kv_heads=1, # MLA-style shared latent (K == V) head_dim=args.head_dim, sliding_window=args.window_size, diff --git a/python/freetoken/models/deepseek_v4/dspark.py b/python/freetoken/models/deepseek_v4/dspark.py new file mode 100644 index 0000000..a2df857 --- /dev/null +++ b/python/freetoken/models/deepseek_v4/dspark.py @@ -0,0 +1,887 @@ +"""dSpark: DeepSeek-V4-Flash's semi-autoregressive draft model. + +DSV4-Flash ships a drafter under ``mtp.{0..n_mtp_layers-1}`` in the checkpoint. It is +NOT a one-token-per-step MTP head: it proposes a whole BLOCK of ``dspark_block_size`` +tokens per draft pass, scores them with a low-rank Markov transition head, and gates +acceptance with a confidence head. That block structure is what makes it worth having +on an offload MoE -- one pass over the experts yields five candidate tokens instead of +one, and expert traffic is what bounds decode here. + +Shape of a draft pass: + +1. The target hands over its hidden state at ``dspark_target_layer_ids`` (40, 41, 42), + concatenated. ``main_norm(main_proj(...))`` projects that back to one hidden width. +2. Each draft layer derives its CONTEXT KV from that same projected tensor, through its + own ``wkv``/``kv_norm``, and writes it at the context slots -- the drafter never runs + the prompt itself. +3. The block's tokens run through ``n_mtp_layers`` full DSV4 blocks (MLA attention, + routed FP4 experts, hyper-connections), then ``hc_head`` and ``norm``. +4. ``markov_head`` adds a transition bias keyed on the previously sampled token, and + ``confidence_head`` scores how likely each position is to survive verification. + +Draft layers carry no compressor and no Lightning Indexer -- the checkpoint has no such +weights for them -- so they attend over the sliding window only. Their layer ids +continue the target's (``n_layers + k``), which is what lets the expert banks, the GPU +slot cache and the KV pools address draft and target layers with one index space. +""" + +from __future__ import annotations + +import math +import statistics +from collections import deque +from typing import Sequence + +import torch +import torch.nn.functional as F +from torch import nn + +from freetoken.kernel.triton.dsv4.hc import hc_pre_combine +from freetoken.kernel.triton.dsv4.norm import inv_rms +from freetoken.utils import init_logger + +from .args import DeepseekV4Args +from .layers import Linear, RMSNorm +from .parallel import div_tp + +logger = init_logger(__name__) + + +class DSparkAcceptanceFallback: + """Request-local measured-acceptance circuit breaker. + + DSpark's confidence scheduler chooses how many proposals to verify after the + drafter has run. It cannot recover that draft cost when a high-entropy request + repeatedly rejects the proposals. This controller observes actual target + acceptance over a bounded window, temporarily selects ordinary target decoding, + and then probes DSpark again. It deliberately knows nothing about the prompt. + + ``take_decision`` returns ``(should_speculate, is_resume_probe)``. A tripped + controller skips exactly ``cooldown_steps`` calls before returning one probe. + """ + + def __init__( + self, + threshold: float, + min_drafted: int, + cooldown_steps: int, + ) -> None: + if not math.isfinite(threshold) or not 0 < threshold <= 1: + raise ValueError("DSpark fallback threshold must be in (0, 1]") + if min_drafted < 1: + raise ValueError("DSpark fallback min_drafted must be >= 1") + if cooldown_steps < 1: + raise ValueError("DSpark fallback cooldown_steps must be >= 1") + self.threshold = float(threshold) + self.min_drafted = int(min_drafted) + self.cooldown_steps = int(cooldown_steps) + self.reset() + + def reset(self) -> None: + self.accepted = 0 + self.drafted = 0 + self.cooldown_left = 0 + self._resume_probe = False + + def take_decision(self) -> tuple[bool, bool]: + if self.cooldown_left > 0: + self.cooldown_left -= 1 + if self.cooldown_left == 0: + self._resume_probe = True + return False, False + resumed = self._resume_probe + self._resume_probe = False + return True, resumed + + def record(self, accepted: int, drafted: int) -> tuple[float, int] | None: + if drafted < 0 or accepted < 0 or accepted > drafted: + raise ValueError( + f"invalid DSpark acceptance sample {accepted}/{drafted}" + ) + self.accepted += int(accepted) + self.drafted += int(drafted) + if self.drafted < self.min_drafted: + return None + + measured_drafted = self.drafted + rate = self.accepted / measured_drafted + self.accepted = 0 + self.drafted = 0 + if rate >= self.threshold: + return None + self.cooldown_left = self.cooldown_steps + self._resume_probe = False + return rate, measured_drafted + + +def choose_adaptive_draft_width( + stale_confidence: Sequence[float] | torch.Tensor, + draft_cost_ms: float, + verify_cost_ms: Sequence[float], +) -> int: + """Choose the DSpark prefix that maximizes expected tokens per millisecond. + + This is Algorithm 1 / Section 5.2 specialized to FreeToken's present DSV4 + serving shape of one active request. ``verify_cost_ms[k]`` is the measured + target cost for anchor + ``k`` draft rows. The drafter has already produced + its whole block, so its measured cost is common to every candidate width. + + For one request, globally sorting cumulative survival probabilities is exactly + the same as considering its prefixes in order: the survival product cannot + increase as the prefix grows. Ties retain the smaller prefix, matching + ``argmax`` in the vLLM implementation. + """ + confidence = [float(c) for c in stale_confidence] + if len(verify_cost_ms) != len(confidence) + 1: + raise ValueError( + "adaptive DSpark needs one verify cost for every width 0..gamma; " + f"got {len(verify_cost_ms)} costs for gamma={len(confidence)}" + ) + if not math.isfinite(draft_cost_ms) or draft_cost_ms < 0: + raise ValueError(f"invalid DSpark draft cost {draft_cost_ms}") + + survival = 1.0 + expected = 1.0 # target bonus token; a verify always advances at least once + best_width = 0 + best_throughput = expected / max( + draft_cost_ms + float(verify_cost_ms[0]), 1e-6 + ) + for width, token_confidence in enumerate(confidence, start=1): + survival *= token_confidence + expected += survival + throughput = expected / max( + draft_cost_ms + float(verify_cost_ms[width]), 1e-6 + ) + if throughput > best_throughput: + best_width = width + best_throughput = throughput + return best_width + + +class DSparkAdaptiveVerification: + """Two-buffer stale-confidence scheduler from DSpark Section 5.2. + + Current confidences are copied asynchronously to pinned host memory. Capacity + selection reads the other buffer, creating the paper's causal barrier across + jagged CUDA-graph buckets. The target curve is profiled at graph capture; the + fixed-cost draft pass is priced from five real executions, like vLLM's five + startup profiling replays. + """ + + _DRAFT_PROFILE_SAMPLES = 5 + + def __init__( + self, + block_size: int, + verify_curve: Sequence[tuple[int, float]], + device: torch.device, + fallback_acceptance: float = 0.0, + fallback_min_drafted: int = 32, + fallback_steps: int = 64, + ) -> None: + expected_spans = list(range(1, block_size + 2)) + curve = sorted((int(span), float(cost)) for span, cost in verify_curve) + if [span for span, _ in curve] != expected_spans: + raise ValueError( + "adaptive DSpark needs profiled graphs for every anchor+prefix span " + f"{expected_spans}, got {[span for span, _ in curve]}" + ) + high = 0.0 + self.verify_cost_ms: list[float] = [] + for _span, cost in curve: + high = max(high, cost, 1e-6) + self.verify_cost_ms.append(high) + + self.block_size = block_size + self.device = device + self._copy_stream = torch.cuda.Stream(device=device) + self._stale = [ + torch.ones(block_size, dtype=torch.float32, pin_memory=True) + for _ in range(2) + ] + self._events: list[torch.cuda.Event | None] = [None, None] + self._stale_idx = 0 + self._active_uid: int | None = None + self._draft_samples: deque[float] = deque( + maxlen=self._DRAFT_PROFILE_SAMPLES + ) + self._debug_left = 12 + self._acceptance_fallback = ( + DSparkAcceptanceFallback( + fallback_acceptance, fallback_min_drafted, fallback_steps + ) + if fallback_acceptance > 0 + else None + ) + + @property + def needs_draft_profile(self) -> bool: + return len(self._draft_samples) < self._DRAFT_PROFILE_SAMPLES + + @property + def draft_cost_ms(self) -> float | None: + if self.needs_draft_profile: + return None + return float(statistics.median(self._draft_samples)) + + def record_draft_cost(self, cost_ms: float) -> None: + if self.needs_draft_profile and math.isfinite(cost_ms) and cost_ms >= 0: + self._draft_samples.append(float(cost_ms)) + if not self.needs_draft_profile: + logger.info_rank0( + "DSpark profiled draft cost: %.2fms (median of %d real steps)", + self.draft_cost_ms, + self._DRAFT_PROFILE_SAMPLES, + ) + + def _reset_request(self, uid: int) -> None: + for event in self._events: + if event is not None: + event.synchronize() + for slot in self._stale: + slot.fill_(1.0) + self._events = [None, None] + self._stale_idx = 0 + self._active_uid = uid + if self._acceptance_fallback is not None: + self._acceptance_fallback.reset() + + def should_speculate(self, uid: int) -> bool: + if self._active_uid != uid: + self._reset_request(uid) + fallback = self._acceptance_fallback + if fallback is None: + return True + speculate, resumed = fallback.take_decision() + if resumed: + logger.info_rank0( + "DSpark fallback: request %d probing speculation after %d " + "target-only steps", + uid, + fallback.cooldown_steps, + ) + return speculate + + def record_acceptance(self, uid: int, accepted: int, drafted: int) -> None: + if self._active_uid != uid: + self._reset_request(uid) + fallback = self._acceptance_fallback + if fallback is None: + return + trip = fallback.record(accepted, drafted) + if trip is None: + return + rate, measured_drafted = trip + logger.info_rank0( + "DSpark fallback: request %d accepted %.1f%% of %d proposals; " + "target-only for %d steps", + uid, + 100.0 * rate, + measured_drafted, + fallback.cooldown_steps, + ) + + def record_and_choose(self, confidence: torch.Tensor, uid: int) -> int: + """Publish this block's confidence and choose from the stale buffer.""" + flat = confidence.detach().float().reshape(-1) + if flat.numel() != self.block_size: + raise RuntimeError( + f"DSpark confidence has {flat.numel()} rows, expected {self.block_size}" + ) + if self._active_uid != uid: + self._reset_request(uid) + + ready_idx = self._stale_idx ^ 1 + ready = self._events[ready_idx] + if ready is not None: + ready.synchronize() + self._stale_idx, write_idx = ready_idx, self._stale_idx + stale = self._stale[self._stale_idx] + + current_stream = torch.cuda.current_stream(self.device) + self._copy_stream.wait_stream(current_stream) + with torch.cuda.stream(self._copy_stream): + self._stale[write_idx].copy_(flat, non_blocking=True) + event = torch.cuda.Event(blocking=True) + event.record(self._copy_stream) + self._events[write_idx] = event + + draft_cost = self.draft_cost_ms + # Collect five real drafter samples before changing shape. This is a + # measured warmup, not a guessed initial cost. + if draft_cost is None: + return self.block_size + width = choose_adaptive_draft_width( + stale, draft_cost, self.verify_cost_ms + ) + if self._debug_left > 0: + self._debug_left -= 1 + logger.info_rank0( + "DSpark adaptive verify: stale=%s draft=%.2fms width=%d/%d", + [round(float(c), 3) for c in stale], + draft_cost, + width, + self.block_size, + ) + return width + + +class MarkovHead(nn.Module): + """Low-rank transition bias over the vocabulary (V x r, then r x V). + + ``markov_w1[token]`` embeds the previously sampled token; ``markov_w2`` turns that + into a per-vocabulary bias added to the draft logits. Both stay REPLICATED under + tensor parallelism: the head runs once per draft position, so sharding it would buy + a little memory and cost an all-reduce plus a full-vocabulary gather every position. + """ + + def __init__(self, args: DeepseekV4Args): + super().__init__() + self.rank = args.dspark_markov_rank + self.markov_w1 = nn.Embedding(args.vocab_size, self.rank) + self.markov_w1.weight.requires_grad_(False) + self.markov_w2 = nn.Parameter( + torch.empty(args.vocab_size, self.rank, dtype=torch.bfloat16), + requires_grad=False, + ) + + def embed(self, token_ids: torch.Tensor) -> torch.Tensor: + """``[B]`` token ids -> ``[B, r]`` Markov embedding.""" + return self.markov_w1(token_ids) + + def bias(self, markov_embed: torch.Tensor) -> torch.Tensor: + """``[B, r]`` -> ``[B, vocab]`` transition bias.""" + return F.linear(markov_embed.to(self.markov_w2.dtype), self.markov_w2) + + +class ConfidenceHead(nn.Module): + """Per-position acceptance confidence, from the head hidden plus the Markov embedding. + + Its output feeds DSpark's hardware-aware prefix scheduler: cumulative survival + probabilities are ranked against profiled draft/verify costs. Runs in fp32 because + the checkpoint defines this as a calibrated scalar probability per position. + """ + + def __init__(self, args: DeepseekV4Args): + super().__init__() + self.proj = nn.Parameter( + torch.empty(1, args.dim + args.dspark_markov_rank, dtype=torch.float32), + requires_grad=False, + ) + + def forward(self, hidden: torch.Tensor, markov_embed: torch.Tensor) -> torch.Tensor: + x = torch.cat([hidden, markov_embed], dim=-1).float() + return torch.sigmoid(F.linear(x, self.proj).squeeze(-1)) + + +class DSparkDrafter(nn.Module): + """The ``mtp.*`` stack: block proposal for DeepSeek-V4-Flash. + + Built only when dSpark is both shipped by the checkpoint and enabled at runtime; + see ``DeepseekV4Args.n_draft_layers``. Holds no KV of its own -- its layers read and + write the same paged pools as the target, at layer ids ``n_layers + k``. + """ + + def __init__(self, args: DeepseekV4Args): + super().__init__() + # Imported here: model.py imports this module, so a top-level import would cycle. + from .model import Block + + self.args = args + self.dim = args.dim + self.hc_mult = args.hc_mult + self.hc_eps = args.hc_eps + self.norm_eps = args.norm_eps + self.block_size = args.dspark_block_size + self.noise_token_id = args.dspark_noise_token_id + self.target_layer_ids = tuple(args.dspark_target_layer_ids) + self.n_layers = n_draft = args.n_draft_layers + + # The target's hidden state at three layers, concatenated, projected back to one + # hidden width. Replicated: it is the drafter's single input, and sharding it + # would need a collective before the first draft layer. + self.main_proj = Linear(self.dim * len(self.target_layer_ids), self.dim, kind="fp8") + self.main_norm = RMSNorm(self.dim, self.norm_eps) + + # Full DSV4 blocks, continuing the target's layer ids. compress_ratio=0: the + # checkpoint ships no compressor or indexer for these layers. + self.layers = nn.ModuleList( + [Block(args.n_layers + k, args, compress_ratio=0) for k in range(n_draft)] + ) + # The block is proposed as a unit, so a drafted query may see the whole block + # rather than only what precedes it. Without this the five positions would be + # five serial steps wearing a block's clothing. + for layer in self.layers: + layer.attn.non_causal = True + + self.norm = RMSNorm(self.dim, self.norm_eps) + hc_dim = self.hc_mult * self.dim + self.hc_head_fn = nn.Parameter( + torch.empty(self.hc_mult, hc_dim, dtype=torch.float32), requires_grad=False + ) + self.hc_head_base = nn.Parameter( + torch.empty(self.hc_mult, dtype=torch.float32), requires_grad=False + ) + self.hc_head_scale = nn.Parameter( + torch.empty(1, dtype=torch.float32), requires_grad=False + ) + + self.markov_head = MarkovHead(args) + self.confidence_head = ConfidenceHead(args) + # Bound by the Transformer: the drafter shares the target's embedding table and + # output head, both vocabulary-parallel under TP, so it must reach them through + # the target's methods rather than holding tensors of its own. + self._embed_tokens = None + + # Sanity: the draft layers shard exactly like the target's, so any TP split that + # the target accepts must work here too. + div_tp(args.moe_inter_dim, "moe_inter_dim", multiple_of=128) + + def bind(self, pool, device: torch.device) -> None: + for layer in self.layers: + layer.attn.bind(pool, device) + + def combine_target_hidden(self, aux_hidden: torch.Tensor) -> torch.Tensor: + """``[T, dim * len(target_layer_ids)]`` -> ``[T, dim]``. + + The drafter's whole view of the context arrives through this projection, which + is why the target must tap exactly ``dspark_target_layer_ids`` and in that order. + """ + expect = self.dim * len(self.target_layer_ids) + if aux_hidden.shape[-1] != expect: + raise ValueError( + f"dSpark expects the target hidden at layers {self.target_layer_ids} " + f"concatenated ({expect} wide), got {aux_hidden.shape[-1]}" + ) + return self.main_norm(self.main_proj(aux_hidden)) + + def store_context_kv( + self, main_x: torch.Tensor, positions: torch.Tensor, window_slots: torch.Tensor + ) -> None: + """Give every draft layer the context's KV, without running the context. + + The drafter never sees the prompt. Each draft layer instead derives its own KV + from ``main_x`` -- the projected target hidden -- through that layer's ``wkv``, + ``kv_norm``, RoPE and FP8 quant, and writes it at the same window slots the + target used. So a draft block attends over the real context while costing one + projection per layer instead of a full prefill. + + ``main_x`` is ``[T, dim]``, ``positions`` the tokens' ABSOLUTE positions ``[T]``, + and ``window_slots`` their global window slots ``[T]``. + + RoPE keys on the absolute position, so the frequencies must be gathered by + ``positions`` rather than sliced from the front. A context that does not start + at zero -- which is every context after the first block -- would otherwise be + rotated as though it did, and the drafter would attend against phases the target + never used. That degrades acceptance without ever raising. + """ + from freetoken.kernel.triton.dsv4.fp8_linear import act_quant_fp8_inplace + + # apply_rotary_emb_decode, not apply_rotary_emb: the frequencies here are + # PER ROW (gathered by absolute position), which is the decode variant's + # contract. apply_rotary_emb broadcasts ONE row of frequencies across the + # sequence and reshapes freqs_cis to [1, T, 1, D] -- with T rows of frequencies + # it fails on the view rather than rotating anything wrongly. + from .ops import apply_rotary_emb_decode + + if positions.shape[0] != main_x.shape[0]: + raise ValueError( + f"positions ({positions.shape[0]}) must cover every context token " + f"({main_x.shape[0]})" + ) + rd = self.args.rope_head_dim + for layer in self.layers: + attn = layer.attn + freqs = attn.freqs_cis.index_select(0, positions) + kv = attn.kv_norm(attn.wkv(main_x)) + # [T, rd] -> [T, 1, rd]: the decode rotary wants a head dim to broadcast + # each row's frequencies over. The view aliases kv, so the in-place write + # lands in the tensor that gets stored. + rope_view = kv[..., -rd:].unsqueeze(1) + apply_rotary_emb_decode(rope_view, freqs) + act_quant_fp8_inplace(kv[..., :-rd], 64) + attn.attn.store_window(kv, attn.layer_id, window_slots) + + def catch_up_context( + self, aux_hidden: torch.Tensor, positions: torch.Tensor, window_slots: torch.Tensor + ) -> None: + """Give the draft layers KV for the positions the target just committed. + + The drafter never runs the prompt. Instead each draft layer derives KV for + already-committed positions from the target's projected hidden, and writes it at + those positions' window slots -- so the draft layers' sliding window fills in + behind the target, one step at a time, for the price of one projection per layer. + + Skipped silently when the tap is unavailable, because a drafter attending to a + stale window is merely a bad drafter, not a broken one. + """ + if aux_hidden is None or positions.numel() == 0: + return + main_x = self.combine_target_hidden(aux_hidden.view(-1, aux_hidden.shape[-1])) + self.store_context_kv(main_x, positions, window_slots) + + def propose( + self, + input_ids: torch.Tensor, + segments, + positions: torch.Tensor, + target_logits_fn, + sampling_params: Sequence, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Run the paper's parallel backbone and sequential Markov sampler. + + The prepared target VERIFY span has ``1 + gamma`` rows per request: the anchor + followed by ``gamma`` proposal slots. The DSpark backbone takes only ``gamma`` + inputs -- anchor plus ``gamma - 1`` masks -- and produces ``gamma`` base-logit + rows. Equation 4 then samples those rows left to right, conditioning row ``k`` + on the token sampled at ``k - 1`` (the anchor for the first row). + + Returns the proposed tokens, the exact shaped distribution each proposal was + drawn from, and the confidence from equation 7, all flattened request-major. + """ + gamma = self.block_size + draft_segments = [] + rows = [] + for i, (off, n, ti, start) in enumerate(segments): + if n != gamma + 1: + raise ValueError( + f"a DSpark verify span has {n} rows, expected gamma + 1 = " + f"{gamma + 1}" + ) + rows.append(torch.arange(off, off + gamma, device=input_ids.device)) + draft_segments.append((i * gamma, gamma, ti, start)) + row_idx = torch.cat(rows) + # The scheduler extends the CPU Req with noise placeholders, but the model's + # flat input is gathered from the GPU token pool; those newly allocated pool + # positions have never been initialized. Construct the reference layout here + # from device-resident data: one live anchor followed by gamma-1 checkpoint + # noise tokens per request (vLLM's sample_from_anchor path). + if self.noise_token_id < 0: + raise ValueError("DSpark requires the checkpoint's noise token id") + draft_ids = torch.full( + (len(segments), gamma), + self.noise_token_id, + dtype=input_ids.dtype, + device=input_ids.device, + ) + draft_ids[:, 0] = input_ids[0, row_idx.view(len(segments), gamma)[:, 0]] + draft_ids = draft_ids.view(1, -1) + draft_pos = positions[row_idx] + head_hidden = self.head_hidden( + self.embed_block(draft_ids), draft_ids, draft_segments, draft_pos + )[0] + # The checkpoint's two heads consume different forms of this tensor. vLLM's + # DSV4 DSpark model returns the PRE-norm hc_head output to the speculator: + # base logits use norm(head_hidden), while equation 7's confidence projection + # consumes head_hidden itself. Normalizing before returning silently changes + # every confidence score even though the proposal logits still look plausible. + base_logits = target_logits_fn(self.norm(head_hidden)).view( + len(segments), gamma, -1 + ) + return self.sample_block( + base_logits, + head_hidden.view(len(segments), gamma, -1), + draft_ids.view(len(segments), gamma)[:, 0], + sampling_params, + ) + + def embed_block(self, input_ids: torch.Tensor) -> torch.Tensor: + """Token embeddings for the block, ``[1, T, dim]``. + + The drafter has no embedding table of its own -- it shares the target's, which is + vocabulary-parallel under TP, so the lookup has to go through the target's method + rather than a bare index. + """ + return self._embed_tokens(input_ids) + + def block_input_ids(self, last_token: int, device: torch.device) -> torch.Tensor: + """The token ids a draft block is fed: the real last token, then noise. + + dSpark is a BLOCK predictor, not an autoregressive one. It is handed the last + committed token followed by ``block_size - 1`` copies of the checkpoint's noise + token, and fills every position in a single pass -- which is the whole reason + one pass yields several candidates. Feeding it anything else at those positions + (zeros, the last token repeated) puts it off its training distribution and the + proposals stop being accepted. + """ + if self.noise_token_id < 0: + raise ValueError( + "this checkpoint declares no dspark_noise_token_id, so a draft block " + "has nothing to place at its unknown positions" + ) + ids = torch.full( + (self.block_size,), self.noise_token_id, dtype=torch.long, device=device + ) + ids[0] = last_token + return ids + + def hc_head(self, x: torch.Tensor) -> torch.Tensor: + """Reduce the hyper-connection copies to one hidden state (mirrors Transformer).""" + shape, dtype = x.size(), x.dtype + x2d = x.flatten(2) + xf = x2d.float() + mixes = F.linear(xf, self.hc_head_fn) * inv_rms(x2d, self.norm_eps) + pre = torch.sigmoid(mixes * self.hc_head_scale + self.hc_head_base) + self.hc_eps + M = shape[0] * shape[1] + return hc_pre_combine( + x.view(M, self.hc_mult, self.dim), pre.view(M, self.hc_mult), dtype + ).view(*shape[:2], self.dim) + + def head_hidden( + self, embeds: torch.Tensor, input_ids: torch.Tensor, segments, positions: torch.Tensor + ) -> torch.Tensor: + """Run the draft stack and return the pre-norm ``hc_head`` hidden. + + ``embeds`` is ``[1, T, dim]`` -- the block's token embeddings, already summed with + the projected target hidden by the caller. Returns ``[1, T, dim]``. + """ + h = embeds.unsqueeze(2).repeat(1, 1, self.hc_mult, 1) + for layer in self.layers: + h = layer.prefill_batched(h, input_ids, segments, positions) + return self.hc_head(h) + + def sample_block( + self, + base_logits: torch.Tensor, + head_hidden: torch.Tensor, + anchor: torch.Tensor, + sampling_params: Sequence, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Equation 4/5/7 over ``[request, gamma, ...]`` backbone outputs.""" + requests, gamma, vocab = base_logits.shape + if len(sampling_params) != requests: + raise ValueError( + f"got {len(sampling_params)} sampling parameter sets for {requests} requests" + ) + prev = anchor.long() + proposed = torch.empty( + (requests, gamma), dtype=torch.long, device=base_logits.device + ) + q = torch.empty( + (requests, gamma, vocab), dtype=torch.float32, device=base_logits.device + ) + confidence = torch.empty( + (requests, gamma), dtype=torch.float32, device=base_logits.device + ) + for k in range(gamma): + markov = self.markov_head.embed(prev) + logits_k = base_logits[:, k].float() + self.markov_head.bias(markov).float() + step_tokens = [] + for r, params in enumerate(sampling_params): + q_r = sampling_probs( + logits_k[r : r + 1], + params.temperature, + params.top_p, + params.top_k, + ) + q[r, k].copy_(q_r[0]) + step_tokens.append( + q_r.argmax(dim=-1) + if params.is_greedy + else torch.multinomial(q_r, 1).squeeze(-1) + ) + next_token = torch.cat(step_tokens) + proposed[:, k].copy_(next_token) + confidence[:, k].copy_( + self.confidence_head(head_hidden[:, k], markov) + ) + prev = next_token + return proposed.flatten(), q.flatten(0, 1), confidence.flatten() + + +def accepted_prefix( + proposed: torch.Tensor, target_argmax: torch.Tensor +) -> tuple[int, int]: + """How much of a proposed block survives verification, and what follows it. + + Speculative decoding is only correct if it emits exactly what the target would + have emitted. So acceptance is a PREFIX: scan the block in order and stop at the + first position where the target disagrees. Everything before it is what the target + would have produced anyway; everything after it was conditioned on a token the + target rejected, and is worthless -- even if it happens to match. + + Returns ``(n_accepted, bonus_token)``. The bonus is the target's own token at the + first rejected position, which verification already computed: a block always + advances by at least one token, so a fully-rejected block is not a wasted step. + """ + n = int(proposed.numel()) + for i in range(n): + if int(proposed[i]) != int(target_argmax[i]): + return i, int(target_argmax[i]) + # Whole block accepted; the target's extra position supplies the next token. + return n, int(target_argmax[n]) if target_argmax.numel() > n else -1 + + +def sampling_probs( + logits: torch.Tensor, temperature: float, top_p: float, top_k: int = -1 +) -> torch.Tensor: + """Turn logits into the distribution the sampler would actually draw from. + + Speculative sampling's guarantee is that the emitted token matches what the TARGET + would have produced -- which means the target's distribution AFTER temperature and + top-p, not the raw softmax. Both the draft q and the target p have to be shaped the + same way, or the ratio test compares two different things and the guarantee is void. + """ + if temperature <= 0: + # Greedy: a point mass on the argmax. The ratio test then degenerates to an + # equality check, which is exactly the greedy acceptance rule. + out = torch.zeros_like(logits) + out.scatter_(-1, logits.argmax(dim=-1, keepdim=True), 1.0) + return out + # DSpark evaluates standard speculative sampling at temperature 1.0. Dividing by + # exactly one is numerically redundant, but it still launches an elementwise kernel + # for every draft and target row in every cycle. + scaled = logits.float() if temperature == 1.0 else logits.float() / temperature + if top_k and top_k > 0: + kth = scaled.topk(min(top_k, scaled.shape[-1]), dim=-1).values[..., -1:] + scaled = scaled.masked_fill(scaled < kth, float("-inf")) + probs = torch.softmax(scaled, dim=-1) + if 0 < top_p < 1: + ordered, idx = probs.sort(dim=-1, descending=True) + cum = ordered.cumsum(dim=-1) + # Keep the smallest prefix whose mass reaches top_p; the shift keeps the first + # token even when it alone already exceeds the threshold. + drop = cum - ordered > top_p + ordered = ordered.masked_fill(drop, 0.0) + probs = torch.zeros_like(probs).scatter_(-1, idx, ordered) + probs = probs / probs.sum(dim=-1, keepdim=True) + return probs + + +def rejection_accept( + proposed: torch.Tensor, q: torch.Tensor, p: torch.Tensor, + generator: torch.Generator | None = None, +) -> tuple[int, int]: + """Speculative sampling: accept a prefix, then resample from the residual. + + This is what lets a SAMPLED request speculate without changing its distribution. + Comparing against the target's argmax (what greedy acceptance does) would bias the + output towards the mode -- faster, and not what the caller asked for. + + For each drafted position the proposal x is accepted with probability + ``min(1, p(x) / q(x))``. On rejection the replacement is drawn from the normalized + residual ``max(0, p - q)``, and the block stops there. Together those two steps make + the emitted token exactly p-distributed -- the standard result: speculation becomes + invisible in the output, not merely close. + + ``q`` and ``p`` are the draft and target probabilities, ``[k, vocab]`` and + ``[k+1, vocab]``. Returns ``(n_accepted, next_token)``. + """ + n = int(proposed.numel()) + for i in range(n): + x = int(proposed[i]) + px, qx = float(p[i, x]), float(q[i, x]) + u = float(torch.rand((), generator=generator)) + # The test is u < p(x)/q(x), but evaluated as p(x) > u * q(x) so nothing is + # divided by a probability. q(x) is denormal for a token the drafter would + # essentially never pick, and the ratio there is meaningless while the product + # is exact. (vLLM's kernel does the same thing in log space.) + if px > u * qx: + continue + residual = torch.clamp(p[i] - q[i], min=0.0) + total = float(residual.sum()) + # A degenerate residual (p entirely inside q) means there is nothing left to + # prefer; fall back to p itself rather than dividing by zero. + dist = residual / total if total > 0 else p[i] + return i, int(torch.multinomial(dist, 1, generator=generator)) + # Whole block accepted: the verify's extra position supplies the next token, drawn + # from the target's own distribution so the bonus is p-distributed too. + return n, int(torch.multinomial(p[n], 1, generator=generator)) + + +def rejection_accept_device( + proposed: torch.Tensor, + q: torch.Tensor, + p: torch.Tensor, + generator: torch.Generator | None = None, +) -> tuple[int, int]: + """Device-resident form of vLLM's standard rejection sampler. + + The acceptance ratios, prefix break, residual distribution, and recovered-token + draw stay on the logits device. Only the accepted length and recovered token + cross to the host, which the scheduler needs to resize the request. This has the + same p/q rule as :func:`rejection_accept`; it removes the full-vocabulary D2H copy, + not any part of the distribution-correct algorithm. + """ + n = int(proposed.numel()) + if q.shape[0] != n or p.shape[0] < n + 1: + raise ValueError( + f"rejection sampler shape mismatch: proposed={n}, q={q.shape}, p={p.shape}" + ) + if n: + token_ids = proposed.long() + rows = torch.arange(n, device=proposed.device) + p_token = p[:n][rows, token_ids] + q_token = q[:n][rows, token_ids] + uniforms = torch.rand( + n, device=p.device, dtype=torch.float32, generator=generator + ) + # Keep the prefix reduction on device. Converting the first rejection to a + # Python int here would synchronize once for the length and again for the + # recovered token below. vLLM publishes both results together as well. + positions = torch.arange(n, device=p.device, dtype=torch.int64) + n_acc_device = torch.where( + p_token <= uniforms * q_token, + positions, + torch.full_like(positions, n), + ).amin() + else: + n_acc_device = torch.zeros((), dtype=torch.int64, device=p.device) + + if n: + rejected_row = n_acc_device.clamp_max(n - 1) + residual = torch.clamp(p[rejected_row] - q[rejected_row], min=0.0) + total = residual.sum() + # clamp_min protects the unselected division branch from producing NaNs. + residual_dist = torch.where( + total > 0, + residual / total.clamp_min(torch.finfo(residual.dtype).tiny), + p[rejected_row], + ) + dist = torch.where(n_acc_device < n, residual_dist, p[n]) + else: + dist = p[n] + bonus_device = torch.multinomial(dist, 1, generator=generator).squeeze(0) + result = torch.stack((n_acc_device, bonus_device.to(torch.int64))).to( + "cpu", non_blocking=False + ) + return int(result[0]), int(result[1]) + + +def window_cols_for_block( + start_pos: int, n: int, window: int, non_causal: bool, + device: torch.device | None = None, +) -> torch.Tensor: + """The window candidate columns a segment's queries may read. + + Extracted from ``Attention._prefill_segment`` so the masking rule -- the one thing + that decides whether a block is genuinely semi-autoregressive -- can be checked + without a GPU or a KV pool. ``-1`` marks a column no query may read. + + Causal: row i sees ``[pos_i - window + 1, pos_i]``. + Non-causal DSpark: every row sees the trailing ``window`` CONTEXT tokens plus all + ``n`` block tokens, matching the paper and vLLM's sparse-index construction. + """ + end = start_pos + n + if non_causal: + w_lo = max(0, start_pos - window) + width = end - w_lo + return torch.arange(width, device=device).unsqueeze(0).expand(n, width) + w_lo = max(0, start_pos - window + 1) + ar = torch.arange(window, device=device) + abs_p = start_pos + torch.arange(n, device=device).unsqueeze(1) + cand = (abs_p - window + 1).clamp(min=w_lo) + ar + return torch.where(cand > abs_p, -1, cand - w_lo) + + +__all__ = [ + "ConfidenceHead", + "DSparkAcceptanceFallback", + "DSparkAdaptiveVerification", + "DSparkDrafter", + "MarkovHead", + "accepted_prefix", + "choose_adaptive_draft_width", + "rejection_accept_device", + "window_cols_for_block", +] diff --git a/python/freetoken/models/deepseek_v4/model.py b/python/freetoken/models/deepseek_v4/model.py index c00d20c..7da008a 100644 --- a/python/freetoken/models/deepseek_v4/model.py +++ b/python/freetoken/models/deepseek_v4/model.py @@ -21,14 +21,20 @@ from __future__ import annotations +import os +from dataclasses import dataclass + import torch import torch.nn.functional as F from torch import nn from freetoken.core import get_global_ctx +from freetoken.distributed import DistributedCommunicator from freetoken.kernel.triton.dsv4.hc import hc_post_combine, hc_pre_combine +from freetoken.kernel.triton.dsv4.norm import inv_rms from freetoken.kernel.triton.dsv4.sinkhorn import hc_split_sinkhorn from freetoken.models.blocks import BaseLLMModel +from freetoken.utils import init_logger from .args import DeepseekV4Args from .attention import Attention @@ -45,17 +51,38 @@ get_window_topk_idxs, ) from .moe import Expert, Gate # noqa: F401 +from .parallel import div_tp, tp_info, validate_tp + +logger = init_logger(__name__) + +# FREETOKEN_SPEC_DEBUG=1 also reports where the draft context KV lands. +_DRAFT_CTX_DEBUG = os.environ.get("FREETOKEN_SPEC_DEBUG", "0") == "1" + + +@dataclass(frozen=True) +class DSparkTargetFeatures: + """The target outputs consumed by the next DSpark proposal. + + vLLM passes auxiliary hidden states, target positions, and context slot mappings + together. FreeToken stores request table rows instead of physical slots so the + slots are resolved from the live page map when the feature bundle is consumed. + Every tensor is produced by the same target forward. + """ + + hidden: torch.Tensor + positions: torch.Tensor + table_rows: torch.Tensor class Block(nn.Module): """Decoder block with manifold-constrained Hyper-Connections (4 residual streams).""" - def __init__(self, layer_id: int, args: DeepseekV4Args): + def __init__(self, layer_id: int, args: DeepseekV4Args, compress_ratio: int | None = None): super().__init__() self.layer_id = layer_id self.norm_eps = args.norm_eps self.dim = args.dim - self.attn = Attention(layer_id, args) + self.attn = Attention(layer_id, args, compress_ratio=compress_ratio) self.ffn = MoE(layer_id, args) self.attn_norm = RMSNorm(args.dim, self.norm_eps) self.ffn_norm = RMSNorm(args.dim, self.norm_eps) @@ -73,14 +100,15 @@ def __init__(self, layer_id: int, args: DeepseekV4Args): def hc_pre(self, x, hc_fn, hc_scale, hc_base): shape, dtype = x.size(), x.dtype - xf = x.flatten(2).float() - rsqrt = torch.rsqrt(xf.square().mean(-1, keepdim=True) + self.norm_eps) + x2d = x.flatten(2) + xf = x2d.float() + rsqrt = inv_rms(x2d, self.norm_eps) mixes = F.linear(xf, hc_fn) * rsqrt pre, post, comb = hc_split_sinkhorn( mixes.view(-1, mixes.size(-1)), hc_scale, hc_base, self.hc_mult, self.hc_sinkhorn_iters, self.hc_eps ) M = shape[0] * shape[1] - y = hc_pre_combine(xf.view(M, self.hc_mult, self.dim), pre, dtype).view(*shape[:2], self.dim) + y = hc_pre_combine(x.view(M, self.hc_mult, self.dim), pre, dtype).view(*shape[:2], self.dim) return y, post.view(M, self.hc_mult), comb.view(M, self.hc_mult, self.hc_mult) def hc_post(self, x, residual, post, comb): @@ -126,41 +154,141 @@ def decode_step(self, x, pos, rows, cmp_stage_cap, input_ids, wctx=None): x = self.hc_post(x, residual, post, comb) return x + def verify_block( + self, x, pos, rows, cmp_stage_cap, input_ids, num_reqs, span, wctx + ): + """Fixed-width DSpark target block; algebra matches ``decode_step``.""" + residual = x + x, post, comb = self.hc_pre( + x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base + ) + x = self.attn_norm(x) + x = self.attn.verify_block( + x, pos, rows, cmp_stage_cap, num_reqs, span, wctx + ) + x = self.hc_post(x, residual, post, comb) + + residual = x + x, post, comb = self.hc_pre( + x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base + ) + x = self.ffn_norm(x) + x = self.ffn(x, input_ids) + return self.hc_post(x, residual, post, comb) + class Transformer(nn.Module): def __init__(self, args: DeepseekV4Args): super().__init__() + # Check every tensor-parallel split before building a single layer, so a bad + # --tensor-parallel-size fails with one clear message, not a reshape deep in a + # forward. embed / head / norm / hyper-connections stay replicated. + validate_tp(args) self.args = args self.norm_eps = args.norm_eps self.hc_eps = args.hc_eps self.hc_mult = hc_mult = args.hc_mult - self.embed = nn.Embedding(args.vocab_size, args.dim) + # Vocabulary parallelism. The embedding table and the output head are the two + # largest replicated tensors (3.0 GiB together), so each rank keeps one + # contiguous vocabulary block: the embedding all-reduces its masked lookup, the + # head all-gathers its logit slice. + rank, tp = tp_info() + self.tp_size = tp + self.vocab_size = args.vocab_size + vocab_local = div_tp(args.vocab_size, "vocab_size") + self.vocab_start = rank * vocab_local + self.vocab_local = vocab_local + self._comm = DistributedCommunicator() if tp > 1 else None + self.embed = nn.Embedding(vocab_local, args.dim) self.embed.weight.requires_grad_(False) self.layers = nn.ModuleList([Block(i, args) for i in range(args.n_layers)]) self.norm = RMSNorm(args.dim, self.norm_eps) - self.head = nn.Parameter(torch.empty(args.vocab_size, args.dim, dtype=torch.bfloat16), requires_grad=False) + self.head = nn.Parameter(torch.empty(vocab_local, args.dim, dtype=torch.bfloat16), requires_grad=False) hc_dim = hc_mult * args.dim self.hc_head_fn = nn.Parameter(torch.empty(hc_mult, hc_dim, dtype=torch.float32), requires_grad=False) self.hc_head_base = nn.Parameter(torch.empty(hc_mult, dtype=torch.float32), requires_grad=False) self.hc_head_scale = nn.Parameter(torch.empty(1, dtype=torch.float32), requires_grad=False) + # dSpark drafter, only when the checkpoint ships it AND the run asked for it. + # Its blocks continue this model's layer ids, so they share the expert banks, + # the GPU slot cache and the KV pools with no separate index space. + self.drafter = None + # Layers whose output the drafter reads, as 0-based indices. vLLM converts the + # checkpoint ids [40, 41, 42] to aux ids [41, 42, 43], then captures when + # ``idx + 1`` is in that set: target outputs 40, 41 and 42 exactly. + self._aux_layer_ids: frozenset[int] = frozenset() + if args.n_draft_layers: + from .dspark import DSparkDrafter + + self.drafter = DSparkDrafter(args) + # The drafter shares this model's embedding table and output head; both are + # vocabulary-parallel under TP, so it reaches them through these methods + # rather than holding tensors that would only cover one rank's slice. + self.drafter._embed_tokens = self.embed_tokens + ids = tuple(args.dspark_target_layer_ids) + self._aux_layer_ids = frozenset(ids) + bad = [i for i in ids if not 0 <= i < args.n_layers] + if bad: + raise ValueError( + f"0-based dspark_target_layer_ids {ids} fall outside the " + f"{args.n_layers} target layers" + ) + logger.info_rank0( + f"dSpark: tapping target layer outputs {sorted(ids)}" + ) + self._target_features: DSparkTargetFeatures | None = None def bind(self, pool, device: torch.device) -> None: for layer in self.layers: layer.attn.bind(pool, device) + if self.drafter is not None: + self.drafter.bind(pool, device) + + def target_features(self) -> DSparkTargetFeatures | None: + return self._target_features + + def embed_tokens(self, input_ids: torch.Tensor) -> torch.Tensor: + """Vocabulary-parallel lookup. Rows outside this rank's block contribute zero, so + the all-reduce sums exactly one real row per token. Index arithmetic only -- no + host sync, so a captured decode graph replays it.""" + if self._comm is None: + return self.embed(input_ids) + local = input_ids - self.vocab_start + outside = (local < 0) | (local >= self.vocab_local) + y = self.embed(local.masked_fill(outside, 0)) + y = y.masked_fill(outside.unsqueeze(-1), 0) + return self._comm.all_reduce(y) + + def logits(self, h: torch.Tensor) -> torch.Tensor: + """Head projection over this rank's vocabulary block, then an all-gather back to + the full [rows, vocab]. all_gather concatenates on dim 0, so the gathered tensor + is rank-major: one row block per rank, which is the vocabulary order.""" + local = F.linear(h, self.head) # [rows, vocab_local] + if self._comm is None: + return local + rows = local.shape[0] + gathered = self._comm.all_gather(local) # [tp*rows, vocab_local] + if rows == 1: + return gathered.view(1, -1)[:, : self.vocab_size] + gathered = gathered.view(self.tp_size, rows, self.vocab_local) + gathered = gathered.permute(1, 0, 2).reshape(rows, self.tp_size * self.vocab_local) + return gathered[:, : self.vocab_size] def hc_head(self, x): shape, dtype = x.size(), x.dtype dim = self.args.dim - xf = x.flatten(2).float() - rsqrt = torch.rsqrt(xf.square().mean(-1, keepdim=True) + self.norm_eps) + x2d = x.flatten(2) + xf = x2d.float() + rsqrt = inv_rms(x2d, self.norm_eps) mixes = F.linear(xf, self.hc_head_fn) * rsqrt pre = torch.sigmoid(mixes * self.hc_head_scale + self.hc_head_base) + self.hc_eps M = shape[0] * shape[1] - return hc_pre_combine(xf.view(M, self.hc_mult, dim), pre.view(M, self.hc_mult), dtype).view(*shape[:2], dim) + return hc_pre_combine( + x.view(M, self.hc_mult, dim), pre.view(M, self.hc_mult), dtype + ).view(*shape[:2], dim) def prefill_batched( self, input_ids: torch.Tensor, segments, flat_positions: torch.Tensor, - last_indices: torch.Tensor, + last_indices: torch.Tensor, logit_indices: torch.Tensor | None = None, ) -> torch.Tensor: # Ragged batched prefill (bs >= 1). ``input_ids`` is [1, T] -- the requests' NEW tokens # concatenated (cu_seqlens, no padding); each request starts at its own cached_len @@ -172,13 +300,38 @@ def prefill_batched( # metadata; ``flat_positions`` [T] is the scheduler-staged batch.positions (per-token # ABSOLUTE position); ``last_indices`` [B] is the flattened index of each request's # final token -> its next-token logits row. - h = self.embed(input_ids) + h = self.embed_tokens(input_ids) h = h.unsqueeze(2).repeat(1, 1, self.hc_mult, 1) - for layer in self.layers: + aux: list[torch.Tensor] = [] + for i, layer in enumerate(self.layers): h = layer.prefill_batched(h, input_ids, segments, flat_positions) + if i in self._aux_layer_ids: + # The drafter's whole view of the context: this block's output with the + # hyper-connection copies averaged away, [1, T, dim]. + aux.append(h.mean(dim=2)) + aux_hidden = ( + torch.cat(aux, dim=-1).view(-1, self.args.dim * len(aux)) + if aux else None + ) + if aux: + rows = torch.empty_like(flat_positions) + for off, n, ti, _start in segments: + rows[off:off + n] = ti + self._target_features = DSparkTargetFeatures( + aux_hidden, + flat_positions.clone(), + rows.clone(), + ) + else: + self._target_features = None h = self.hc_head(h) h = self.norm(h) - return F.linear(h[0, last_indices], self.head) # [B, vocab] + # Normally only each request's final token needs logits. A speculative VERIFY + # pass is the same ragged prefill -- each request resuming from its own + # start_pos -- but it needs the logits at EVERY drafted position, to compare + # them against what the drafter proposed. logit_indices selects which rows. + idx = last_indices if logit_indices is None else logit_indices + return self.logits(h[0, idx]) # [len(idx), vocab] def decode( self, input_ids: torch.Tensor, pos: torch.Tensor, cmp_stage_cap: int @@ -194,7 +347,7 @@ def decode( # in-flight replay (it mutates only the live map). B = input_ids.size(0) rows = torch.arange(B, device=input_ids.device) - h = self.embed(input_ids) + h = self.embed_tokens(input_ids) h = h.unsqueeze(2).repeat(1, 1, self.hc_mult, 1) # Hoist the layer-invariant per-step decode tensors (shared window-ring global slots): # resolved ONCE off the attention metadata (recomputed per call, so a capture records @@ -202,11 +355,78 @@ def decode( # into every layer. They read only the shared snapshot / positions, so they are identical # across layers. wctx = get_global_ctx().batch.attn_metadata.window_ctx(pos, rows) - for layer in self.layers: + aux: list[torch.Tensor] = [] + for i, layer in enumerate(self.layers): h = layer.decode_step(h, pos, rows, cmp_stage_cap, input_ids, wctx) + if i in self._aux_layer_ids: + aux.append(h.mean(dim=2)) # [B, 1, dim] + aux_hidden = ( + torch.cat(aux, dim=-1).view(-1, self.args.dim * len(aux)) + if aux else None + ) + if aux: + table_rows = get_global_ctx().batch.active_table_idx + if table_rows is None: + raise RuntimeError("DSpark decode features require request table rows") + # clone() is captured: replay updates these output buffers without leaving + # them aliased to the graph input buffers that the next batch overwrites. + self._target_features = DSparkTargetFeatures( + aux_hidden, + pos.reshape(-1).clone(), + table_rows[:B].reshape(-1).clone(), + ) + else: + self._target_features = None h = self.hc_head(h) h = self.norm(h) - return F.linear(h[:, -1], self.head) + return self.logits(h[:, -1]) + + def verify_block( + self, input_ids: torch.Tensor, pos: torch.Tensor, span: int + ) -> torch.Tensor: + """Run one fixed-shape DSpark target verification inside a CUDA graph. + + ``input_ids`` is request-major ``[1, R * span]`` and every request contributes + exactly anchor + gamma rows. The attention backend supplies one persistent + whole-history snapshot per request; ``rows`` maps every token to that request. + """ + total = input_ids.numel() + if span < 1 or total % span: + raise ValueError(f"invalid DSpark verify shape: {total} rows, span={span}") + num_reqs = total // span + rows = torch.arange(num_reqs, device=input_ids.device).repeat_interleave(span) + md = get_global_ctx().batch.attn_metadata + cmp_stage_cap = md.stage_width - 1 + wctx = md.window_ctx(pos, rows) + + h = self.embed_tokens(input_ids) + h = h.unsqueeze(2).repeat(1, 1, self.hc_mult, 1) + aux: list[torch.Tensor] = [] + for i, layer in enumerate(self.layers): + h = layer.verify_block( + h, pos, rows, cmp_stage_cap, input_ids, + num_reqs, span, wctx, + ) + if i in self._aux_layer_ids: + aux.append(h.mean(dim=2)) + + aux_hidden = ( + torch.cat(aux, dim=-1).view(-1, self.args.dim * len(aux)) + if aux else None + ) + if aux: + table_rows = get_global_ctx().batch.active_table_idx + if table_rows is None or table_rows.numel() < num_reqs: + raise RuntimeError("DSpark verify features require one table row per request") + self._target_features = DSparkTargetFeatures( + aux_hidden, + pos.reshape(-1).clone(), + table_rows[:num_reqs].repeat_interleave(span).clone(), + ) + else: + self._target_features = None + h = self.norm(self.hc_head(h)) + return self.logits(h.view(-1, self.args.dim)) class DeepseekV4ForCausalLM(BaseLLMModel): @@ -222,6 +442,9 @@ def __init__(self, config): self._args: DeepseekV4Args = config.dsv4_args self._transformer = Transformer(self._args) self._bound = False + self.speculative_verify_block_size = ( + self._args.dspark_block_size if self._args.n_draft_layers else 0 + ) def _ensure_bound(self) -> None: if self._bound: @@ -241,8 +464,15 @@ def _iter_offload_moe_layers(self): # Hook for moe.offload_cache.iter_offload_moe_layers: DSV4's MoE block (block.ffn) # is a bespoke nn.Module (not OffloadMoELayer), so yield its inner routed-expert # layer instead, whose .offload_cache the generic attach sets. + # + # The dSpark draft blocks own routed experts too, and their layer ids continue the + # target's, so they come AFTER and in order -- the cache indexes by position in + # this sequence, and it asserts the count against model_config.num_moe_layers. for block in self._transformer.layers: yield block.ffn.experts + if self._transformer.drafter is not None: + for block in self._transformer.drafter.layers: + yield block.ffn.experts def state_dict(self, *, prefix: str = "", result=None): result = {} if result is None else result @@ -261,11 +491,75 @@ def load_state_dict(self, state_dict, *, prefix: str = "", _internal: bool = Fal raise RuntimeError(f"Unexpected keys in state_dict: {list(state_dict.keys())}") self._transformer.load_state_dict(casted, assign=True, strict=False) + def dspark_target_features(self) -> DSparkTargetFeatures | None: + """Outputs of the most recent target forward, as one addressed bundle.""" + return self._transformer.target_features() + + def catch_up_draft_context(self, features: DSparkTargetFeatures) -> None: + """Precompute draft-layer context KV from one target feature bundle.""" + drafter = self._transformer.drafter + if drafter is None: + return + flat = features.hidden + positions = features.positions + rows = features.table_rows + if flat.shape[0] != positions.numel() or rows.numel() != positions.numel(): + raise RuntimeError( + "DSpark target features must have one hidden, position, and table row " + f"per token; got {flat.shape[0]}, {positions.numel()}, {rows.numel()}" + ) + backend = get_global_ctx().attn_backend + slots = backend.window_slots_at(rows, positions) + valid = slots >= 0 + flat = flat[valid] + positions = positions[valid] + slots = slots[valid] + if positions.numel() == 0: + return + if _DRAFT_CTX_DEBUG: + logger.debug_rank0( + "draft ctx: %d rows, positions %s..%s, slots %s..%s", + flat.shape[0], int(positions.min()), int(positions.max()), + int(slots.min()), int(slots.max()), + ) + drafter.catch_up_context(flat, positions, slots) + + def draft( + self, sampling_params + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None: + """Propose the current batch's block with the dSpark drafter. + + Runs over the SAME prepared batch as the verify pass that follows -- the same + positions, the same segments, the same paged slots. The draft layers own layer + ids after the target's, so ``allocate_paged`` has already given them their own + window slots at these positions and they write KV there without disturbing the + target's. + + Returns ``(proposed [R*gamma], q [R*gamma,V], confidence [R*gamma])``. + """ + if self._transformer.drafter is None: + return None + self._ensure_bound() + batch = get_global_ctx().batch + md = batch.attn_metadata + return self._transformer.drafter.propose( + batch.input_ids.long().view(1, -1), + md.segments, + batch.positions.long(), + self._transformer.logits, + sampling_params, + ) + def forward(self) -> torch.Tensor: self._ensure_bound() batch = get_global_ctx().batch input_ids = batch.input_ids.long() md = batch.attn_metadata + if getattr(batch, "spec_verify_decode", False): + span = int(batch.spec_block) + 1 + return self._transformer.verify_block( + input_ids.view(1, -1), batch.positions.long(), span + ) if batch.is_prefill: # Ragged batched prefill (bs >= 1): each request starts from its own cached_len. # A cold segment (start_pos == 0) re-seeds the compressor carry register inside its @@ -273,9 +567,17 @@ def forward(self) -> torch.Tensor: # FROM THE RING. Per-token ops (embed / HC / norm / MoE) run batched over the # concatenated tokens; attention runs per segment so the carry / slot maps never # cross requests. + # A speculative verify needs logits at EVERY drafted position, not just each + # request's last one: acceptance compares the target's own prediction at each + # position against what the drafter proposed there. + logit_indices = None + if getattr(batch, "speculative", False): + logit_indices = torch.arange( + input_ids.numel(), device=input_ids.device, dtype=torch.long + ) return self._transformer.prefill_batched( input_ids.view(1, -1), md.segments, batch.positions.long(), - md.last_indices.long(), + md.last_indices.long(), logit_indices=logit_indices, ) # DECODE (bs>=1): per-row position (GPU int tensor -> no host syncs / graph safe). The # compressed staging cap is the max position any row reaches (eager); a static max_seq-1 @@ -293,4 +595,4 @@ def forward(self) -> torch.Tensor: return self._transformer.decode(input_ids.view(B, 1), pos, cmp_stage_cap) -__all__ = ["Transformer", "DeepseekV4ForCausalLM"] +__all__ = ["DSparkTargetFeatures", "Transformer", "DeepseekV4ForCausalLM"] diff --git a/python/freetoken/models/deepseek_v4/moe.py b/python/freetoken/models/deepseek_v4/moe.py index 9b28922..bb4617c 100644 --- a/python/freetoken/models/deepseek_v4/moe.py +++ b/python/freetoken/models/deepseek_v4/moe.py @@ -7,6 +7,8 @@ import torch.nn.functional as F from torch import nn +from freetoken.core import get_global_ctx +from freetoken.distributed import DistributedCommunicator from freetoken.kernel.triton.dsv4.bf16_linear import bf16_linear_fp32 from freetoken.kernel.triton.dsv4.swiglu import fused_swiglu from freetoken.layers import OffloadMoELayer @@ -14,6 +16,8 @@ from .args import DeepseekV4Args from .layers import Linear +from .parallel import div_tp, tp_size + class Gate(nn.Module): """MoE router: sqrtsoftplus scoring + hash routing (first ``n_hash_layers``).""" @@ -56,13 +60,19 @@ def forward(self, x: torch.Tensor, input_ids: torch.Tensor): class Expert(nn.Module): - """Dense SwiGLU expert (the shared expert; routed experts are offloaded FP4).""" + """Dense SwiGLU expert (the shared expert; routed experts are offloaded FP4). + + Under TP the intermediate dimension splits: ``w1``/``w3`` are column-parallel and + ``w2`` is row-parallel, so the output is a partial sum. ``MoE.forward`` owns the + single all-reduce that completes it together with the routed half. + """ def __init__(self, dim: int, inter_dim: int, swiglu_limit: float): super().__init__() - self.w1 = Linear(dim, inter_dim, kind="fp8") - self.w2 = Linear(inter_dim, dim, kind="fp8") - self.w3 = Linear(dim, inter_dim, kind="fp8") + inter_local = div_tp(inter_dim, "moe_inter_dim", multiple_of=128) + self.w1 = Linear(dim, inter_local, kind="fp8") + self.w2 = Linear(inter_local, dim, kind="fp8") + self.w3 = Linear(dim, inter_local, kind="fp8") self.swiglu_limit = swiglu_limit def forward(self, x: torch.Tensor) -> torch.Tensor: @@ -88,6 +98,16 @@ def __init__(self, layer_id: int, args: DeepseekV4Args): ) self.swiglu_limit = args.swiglu_limit + def _maybe_all_reduce(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Suppress the base class's per-call collective. + + The routed and the shared expert are both partial sums over the same output + dim, so DSV4 adds them first and all-reduces ONCE in ``MoE.forward``. Letting + the base reduce here would cost a second collective per layer and would also + double-count the shared expert. + """ + return hidden_states + def _prefill_routed( self, hidden_states: torch.Tensor, @@ -104,6 +124,26 @@ def _prefill_routed( return super()._prefill_routed(hidden_states, topk_weights, topk_ids) cache = self.offload_cache assert cache is not None + + # A speculative verify is a decode wearing a prefill's clothes. The scheduler + # marks the batch "prefill" because the block is an extend over several + # positions, but it carries block_size rows per request, not a prompt -- and it + # runs on the critical path of every decode step. + # + # The path below fetches EVERY missing expert over PCIe, uncapped, with no CPU + # overlap. That is the right trade for a prompt, where the fetch amortizes over + # hundreds of tokens. For a 5-row block it moves up to T*top_k experts per layer + # with nothing hiding the latency: measured at ~0.4s per block against a 0.06s + # single-token step, which is the whole reason speculation lost to plain decode + # here rather than a low acceptance rate. + # + # Hybrid decode caps the fetch and overlaps the overflow on the CPU pool, which + # is what a handful of rows wants. + if ( + getattr(get_global_ctx().batch, "speculative", False) + and cache.decode_target == "hybrid" + ): + return self._decode_routed(hidden_states, topk_weights, topk_ids) cache.ensure_experts(self.layer_id, topk_ids) # in-place expert-id -> slot cache.copy_missing() if cache.collect_stats: @@ -129,6 +169,7 @@ def __init__(self, layer_id: int, args: DeepseekV4Args): self.gate = Gate(layer_id, args) self.shared_experts = Expert(args.dim, args.moe_inter_dim, args.swiglu_limit) self.experts = DSV4OffloadMoELayer(layer_id, args) + self._comm = DistributedCommunicator() if tp_size() > 1 else None def forward(self, x: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor: shape = x.size() @@ -143,4 +184,9 @@ def forward(self, x: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor: routed = self.experts.routed_forward( x, weights.float().contiguous(), indices.to(torch.int32).contiguous() ) - return (routed + shared).view(shape) + out = routed + shared + # Both halves are partial sums over the split intermediate dim; one collective + # completes the layer (see DSV4OffloadMoELayer._maybe_all_reduce). + if self._comm is not None: + out = self._comm.all_reduce(out) + return out.view(shape) diff --git a/python/freetoken/models/deepseek_v4/parallel.py b/python/freetoken/models/deepseek_v4/parallel.py new file mode 100644 index 0000000..4459e0b --- /dev/null +++ b/python/freetoken/models/deepseek_v4/parallel.py @@ -0,0 +1,114 @@ +"""Tensor-parallel helpers for DeepSeek-V4-Flash. + +What DSV4 shards, and why: + +* **MLA query/output heads.** ``wq_b`` is column-parallel over ``n_heads``, + ``wo_a`` is sharded over the ``o_groups`` output groups, and ``wo_b`` is + row-parallel with one all-reduce at the end of the block. ``o_groups`` is the + binding constraint: a group owns ``n_heads // o_groups`` heads, so a rank must + own whole groups. +* **MoE intermediate dimension.** The routed FP4 expert banks and the shared + SwiGLU expert both split ``moe_inter_dim``. This is the memory win that makes + the model fit: the host expert banks divide by the TP size, so 4 ranks hold + the same 143 GB in total that 1 rank holds today. + +What stays REPLICATED, and why: + +* **The latent KV path** (``wkv``, ``kv_norm``, the paged pools). MLA keeps one + latent KV per token that every head reads, so there is nothing to split; the + KV pool cost per rank is unchanged. +* **The compressors and the Lightning Indexer.** They select which blocks the + attention reads. Every rank must select the SAME blocks. Replicating them + guarantees that without a collective on the selection path. +* **The router (``Gate``) and the hyper-connection mixers.** Small, and the + routing decision must agree across ranks. +""" + +from __future__ import annotations + +import torch + +from freetoken.distributed import try_get_tp_info + + +def tp_info() -> tuple[int, int]: + """``(rank, size)`` for this process; ``(0, 1)`` before the engine sets it.""" + info = try_get_tp_info() + return (0, 1) if info is None else (info.rank, info.size) + + +def tp_size() -> int: + return tp_info()[1] + + +def tp_rank() -> int: + return tp_info()[0] + + +def div_tp(total: int, what: str, *, multiple_of: int = 1) -> int: + """Per-rank size of ``total``. Fails loudly when the split is not exact. + + ``multiple_of`` guards the block-quantized weights: a 128x128 FP8 scale block + and a 32-element FP4 scale block cannot be cut in half. + """ + size = tp_size() + if total % size != 0: + raise ValueError( + f"DeepSeek-V4 tensor parallelism: {what}={total} is not divisible by " + f"--tensor-parallel-size {size}" + ) + local = total // size + if local % multiple_of != 0: + raise ValueError( + f"DeepSeek-V4 tensor parallelism: {what}={total} split over {size} ranks " + f"gives {local}, which is not a multiple of {multiple_of} (quantization " + f"block size)" + ) + return local + + +def shard(t: torch.Tensor, dim: int) -> torch.Tensor: + """This rank's slice of ``t`` along ``dim``, in its OWN storage. + + The copy is the point. ``narrow`` returns a view, and on dim 0 that view is already + contiguous, so ``.contiguous()`` hands it straight back -- keeping the whole parent + tensor alive behind a 1/N-sized view. Every rank then pays for the full weight it + just sharded: measured at 5.1 GiB per GPU on DeepSeek-V4-Flash at TP=4, silently + charged to the model and taken out of the KV and expert-cache budget. + """ + rank, size = tp_info() + if size == 1: + return t + total = t.shape[dim] + if total % size != 0: + raise ValueError( + f"cannot shard a tensor of shape {tuple(t.shape)} on dim {dim} over {size} ranks" + ) + step = total // size + return t.narrow(dim, rank * step, step).clone(memory_format=torch.contiguous_format) + + +def validate_tp(args) -> None: + """Check every DSV4 split up front, so a bad ``--tensor-parallel-size`` fails + at config time with one clear message instead of at the first bad reshape.""" + size = tp_size() + if size == 1: + return + div_tp(args.o_groups, "o_groups") + div_tp(args.n_heads, "n_heads") + if args.n_heads % args.o_groups != 0: + raise ValueError( + f"DeepSeek-V4 expects n_heads ({args.n_heads}) to be a multiple of " + f"o_groups ({args.o_groups})" + ) + # wq_b output rows and the FP8 128x128 scale grid. + div_tp(args.n_heads * args.head_dim, "n_heads*head_dim", multiple_of=128) + # wo_b input columns. + div_tp(args.o_groups * args.o_lora_rank, "o_groups*o_lora_rank", multiple_of=128) + # Shared expert (FP8 128-blocks) and routed FP4 experts (32-element scale blocks). + div_tp(args.moe_inter_dim, "moe_inter_dim", multiple_of=128) + # Vocabulary-parallel embedding + head. + div_tp(args.vocab_size, "vocab_size") + + +__all__ = ["div_tp", "shard", "tp_info", "tp_rank", "tp_size", "validate_tp"] diff --git a/python/freetoken/models/deepseek_v4/weight.py b/python/freetoken/models/deepseek_v4/weight.py index f5a3fd1..aecd389 100644 --- a/python/freetoken/models/deepseek_v4/weight.py +++ b/python/freetoken/models/deepseek_v4/weight.py @@ -22,6 +22,7 @@ from freetoken.models.loader import drop_page_cache from .args import DeepseekV4Args, load_args +from .parallel import div_tp, shard, tp_info class _ShardReader: @@ -99,31 +100,44 @@ def iter_weights( def get(name: str) -> torch.Tensor: return reader.get(name) - def linear(prefix: str): - yield f"{prefix}.weight", get(f"{prefix}.weight") + def linear(prefix: str, split: int | None = None): + """Yield one linear's tensors, sharded for this rank. + + ``split=0`` is column-parallel (split the output rows), ``split=1`` is + row-parallel (split the input columns), ``None`` replicates. The 128x128 FP8 + ``scale`` grid splits on the SAME axis as its weight, so the two stay aligned. + """ + w = get(f"{prefix}.weight") + yield f"{prefix}.weight", w if split is None else shard(w, split) if reader.has(f"{prefix}.scale"): - yield f"{prefix}.scale", get(f"{prefix}.scale") + s = get(f"{prefix}.scale") + yield f"{prefix}.scale", s if split is None else shard(s, split) try: - yield "embed.weight", get("embed.weight") + # Vocabulary-parallel: each rank keeps one contiguous block of rows. + yield "embed.weight", shard(get("embed.weight"), 0) yield "norm.weight", get("norm.weight") - yield "head", get("head.weight") + yield "head", shard(get("head.weight"), 0) for nm in ("hc_head_fn", "hc_head_base", "hc_head_scale"): yield nm, get(nm) for L in range(args.n_layers): a = f"layers.{L}.attn" + # wq_a / wkv / the norms stay replicated: MLA keeps ONE latent KV per token + # that every head reads, so there is nothing to split on that path. yield from linear(f"{a}.wq_a") yield f"{a}.q_norm.weight", get(f"{a}.q_norm.weight") - yield from linear(f"{a}.wq_b") + yield from linear(f"{a}.wq_b", split=0) # column-parallel over heads yield from linear(f"{a}.wkv") yield f"{a}.kv_norm.weight", get(f"{a}.kv_norm.weight") # wo_a: FP8 in the checkpoint, dequantized to bf16 (reference bf16 einsum). - yield f"{a}.wo_a", _dequant_fp8_block( + # Rows are o_groups blocks of o_lora_rank, so a dim-0 split hands each rank + # whole groups -- matching the heads its wq_b shard produced. + yield f"{a}.wo_a", shard(_dequant_fp8_block( get(f"{a}.wo_a.weight"), get(f"{a}.wo_a.scale") - ) - yield from linear(f"{a}.wo_b") - yield f"{a}.attn_sink", get(f"{a}.attn_sink") + ), 0) + yield from linear(f"{a}.wo_b", split=1) # row-parallel; all-reduced in _wo + yield f"{a}.attn_sink", shard(get(f"{a}.attn_sink"), 0) ratio = args.compress_ratios[L] if ratio: @@ -151,18 +165,81 @@ def linear(prefix: str): yield f"{g}.tid2eid", get(f"{g}.tid2eid") else: yield f"{g}.bias", get(f"{g}.bias") - for proj in ("w1", "w2", "w3"): - yield from linear(f"layers.{L}.ffn.shared_experts.{proj}") + # Shared expert: the intermediate dim splits (w1/w3 column, w2 row); the + # partial sum is all-reduced together with the routed half in MoE.forward. + for proj, split in (("w1", 0), ("w2", 1), ("w3", 0)): + yield from linear(f"layers.{L}.ffn.shared_experts.{proj}", split=split) for nm in ( "hc_attn_fn", "hc_ffn_fn", "hc_attn_base", "hc_ffn_base", "hc_attn_scale", "hc_ffn_scale", ): yield f"layers.{L}.{nm}", get(f"layers.{L}.{nm}") + + yield from _iter_dspark_weights(args, reader, linear) finally: reader.close() +def _iter_dspark_weights(args: DeepseekV4Args, reader: "_ShardReader", linear): + """The ``mtp.*`` drafter, renamed onto the ``drafter.*`` module tree. + + Each draft block shards exactly like a target block, so the same split rules apply. + The per-block weights live under ``mtp.k``; the shared head pieces (main_proj, + main_norm, norm, hc_head, markov, confidence) live under ``mtp.0`` only. + """ + if args.n_draft_layers == 0: + return + get = reader.get + + for k in range(args.n_draft_layers): + src, dst = f"mtp.{k}", f"drafter.layers.{k}" + a_src, a_dst = f"{src}.attn", f"{dst}.attn" + # Same split map as a target block: heads column-parallel, groups on wo_a, + # wo_b row-parallel, shared expert on the intermediate dim. + for name, tensor in linear(f"{a_src}.wq_a"): + yield name.replace(a_src, a_dst), tensor + yield f"{a_dst}.q_norm.weight", get(f"{a_src}.q_norm.weight") + for name, tensor in linear(f"{a_src}.wq_b", split=0): + yield name.replace(a_src, a_dst), tensor + for name, tensor in linear(f"{a_src}.wkv"): + yield name.replace(a_src, a_dst), tensor + yield f"{a_dst}.kv_norm.weight", get(f"{a_src}.kv_norm.weight") + yield f"{a_dst}.wo_a", shard( + _dequant_fp8_block(get(f"{a_src}.wo_a.weight"), get(f"{a_src}.wo_a.scale")), 0 + ) + for name, tensor in linear(f"{a_src}.wo_b", split=1): + yield name.replace(a_src, a_dst), tensor + yield f"{a_dst}.attn_sink", shard(get(f"{a_src}.attn_sink"), 0) + + yield f"{dst}.attn_norm.weight", get(f"{src}.attn_norm.weight") + yield f"{dst}.ffn_norm.weight", get(f"{src}.ffn_norm.weight") + yield f"{dst}.ffn.gate.weight", get(f"{src}.ffn.gate.weight") + yield f"{dst}.ffn.gate.bias", get(f"{src}.ffn.gate.bias") + for proj, split in (("w1", 0), ("w2", 1), ("w3", 0)): + for name, tensor in linear(f"{src}.ffn.shared_experts.{proj}", split=split): + yield name.replace(src, dst), tensor + for nm in ( + "hc_attn_fn", "hc_ffn_fn", "hc_attn_base", + "hc_ffn_base", "hc_attn_scale", "hc_ffn_scale", + ): + yield f"{dst}.{nm}", get(f"{src}.{nm}") + + # The shared pieces are stored once each, on the layer whose role they serve: + # the INPUT projection on the first draft layer, and the output norm / hc_head / + # Markov / confidence heads on the LAST one. + first, last = "mtp.0", f"mtp.{args.n_draft_layers - 1}" + for name, tensor in linear(f"{first}.main_proj"): + yield name.replace(first, "drafter"), tensor + yield "drafter.main_norm.weight", get(f"{first}.main_norm.weight") + yield "drafter.norm.weight", get(f"{last}.norm.weight") + for nm in ("hc_head_fn", "hc_head_base", "hc_head_scale"): + yield f"drafter.{nm}", get(f"{last}.{nm}") + yield "drafter.markov_head.markov_w1.weight", get(f"{last}.markov_head.markov_w1.weight") + yield "drafter.markov_head.markov_w2", get(f"{last}.markov_head.markov_w2.weight") + yield "drafter.confidence_head.proj", get(f"{last}.confidence_head.proj.weight") + + # -------------------------------------------------------------------------------------- # Routed FP4 expert pinned banks. # -------------------------------------------------------------------------------------- @@ -171,6 +248,35 @@ def linear(prefix: str): r"(?Pw1|w2|w3)\.(?Pweight|scale)$" ) +# The dSpark drafter's own routed experts. Its layers are appended after the target's, +# so ``mtp.k`` becomes bank layer ``n_layers + k`` and every layer-indexed structure +# (host banks, GPU slot cache) addresses draft and target layers identically. +_MTP_EXPERT_RE = re.compile( + r"^mtp\.(?P\d+)\.ffn\.experts\.(?P\d+)\." + r"(?Pw1|w2|w3)\.(?Pweight|scale)$" +) + + +def _expert_key(name: str, args: DeepseekV4Args): + """``(bank_layer, expert, proj, kind)`` for a routed-expert tensor, else None. + + Returns None for anything this run does not serve: a target layer past ``n_layers`` + and, when dSpark is off, every ``mtp.*`` expert. + """ + m = _EXPERT_RE.match(name) + if m is not None: + layer = int(m.group("layer")) + if layer >= args.n_layers: # a trailing MTP layer in the target namespace + return None + return layer, int(m.group("expert")), m.group("proj"), m.group("kind") + m = _MTP_EXPERT_RE.match(name) + if m is None: + return None + k = int(m.group("mtp")) + if k >= args.n_draft_layers: # dSpark off, or a layer beyond n_mtp_layers + return None + return args.n_layers + k, int(m.group("expert")), m.group("proj"), m.group("kind") + def load_dsfp4_expert_sources( model_path: str, args: DeepseekV4Args, *, layer_sink=None @@ -191,39 +297,39 @@ def load_dsfp4_expert_sources( folder = model_path weight_map = _weight_map(folder) - L, E = args.n_layers, args.n_routed_experts - H, I = args.dim, args.moe_inter_dim + # L covers the target layers plus, when dSpark is enabled, the drafter's own. + L, E = args.n_moe_layers, args.n_routed_experts - for shard in sorted(set(weight_map.values())): - drop_page_cache(os.path.join(folder, shard)) + for shard_file in sorted(set(weight_map.values())): + drop_page_cache(os.path.join(folder, shard_file)) - shards: dict[str, list[tuple[str, re.Match]]] = collections.defaultdict(list) - for name, shard in weight_map.items(): - m = _EXPERT_RE.match(name) - if m is None: - continue - if int(m.group("layer")) >= L: # skip the MTP layer (index L) - continue - shards[shard].append((name, m)) + shards: dict[str, list[tuple[str, tuple]]] = collections.defaultdict(list) + for name, shard_file in weight_map.items(): + key = _expert_key(name, args) + if key is not None: + shards[shard_file].append((name, key)) - e8m0 = torch.float8_e8m0fnu - specs = { # alloc UNPINNED, fill, then pin-after-fill (skips slow cudaHostAlloc zero-fill) - "gate_up_packed": ((E, 2 * I, H // 2), torch.uint8), - "gate_up_scale": ((E, 2 * I, H // 32), e8m0), - "down_packed": ((E, H, I // 2), torch.uint8), - "down_scale": ((E, H, I // 32), e8m0), - } + specs, I, i_lo = _expert_bank_specs(args) hb = alloc_layer_banks(specs, L) banks = {name: [b.tensor for b in hb[name]] for name in specs} + # Under TP a rank keeps only rows [i_lo, i_lo+I) of w1/w3, and those rows are + # CONTIGUOUS in the file -- so read just them instead of reading the whole tensor and + # throwing (N-1)/N of it away. w2 needs a COLUMN slice, whose rows are strided, so it + # is still read whole. w1+w3 are two thirds of the expert bytes, so at TP=4 this cuts + # a rank's read (and its copy work) roughly in half. + sliced_read = I != args.moe_inter_dim + def _load(sink) -> int: tracker = LayerCompletionTracker(E * 6, hb, sink) # {w1,w2,w3} x {weight,scale} x experts placed = 0 - for shard in tqdm(sorted(shards), desc="Loading DSV4 FP4 experts"): - path = os.path.join(folder, shard) + for shard_file in tqdm(sorted(shards), desc="Loading DSV4 FP4 experts"): + path = os.path.join(folder, shard_file) with safetensors.safe_open(path, framework="pt", device="cpu") as f: - for name, m in shards[shard]: - layer = _place_dsfp4(banks, name, f.get_tensor(name), I) + for name, key in shards[shard_file]: + rows_ready = sliced_read and key[2] in ("w1", "w3") + t = f.get_slice(name)[i_lo:i_lo + I] if rows_ready else f.get_tensor(name) + layer = _place_dsfp4(banks, key, t, I, i_lo, rows_ready) tracker.note(layer) placed += 1 drop_page_cache(path) @@ -244,15 +350,8 @@ def dummy_dsfp4_expert_sources(args: DeepseekV4Args) -> dict[str, list[torch.Ten """Fabricate the 4 ds_fp4 banks for --dummy-weight (no checkpoint on disk).""" from freetoken.moe.host_banks import alloc_layer_banks, pin_banks - L, E = args.n_layers, args.n_routed_experts - H, I = args.dim, args.moe_inter_dim - e8m0 = torch.float8_e8m0fnu - specs = { - "gate_up_packed": ((E, 2 * I, H // 2), torch.uint8), - "gate_up_scale": ((E, 2 * I, H // 32), e8m0), - "down_packed": ((E, H, I // 2), torch.uint8), - "down_scale": ((E, H, I // 32), e8m0), - } + L = args.n_moe_layers + specs, _I, _i_lo = _expert_bank_specs(args) hb = alloc_layer_banks(specs, L) banks = {name: [b.tensor for b in hb[name]] for name in specs} for t in banks["gate_up_packed"]: # packed e2m1; scales stay 0 (valid e8m0) @@ -265,30 +364,66 @@ def dummy_dsfp4_expert_sources(args: DeepseekV4Args) -> dict[str, list[torch.Ten def is_expert_tensor(name: str) -> bool: """Predicate for the common parallel reader: is this a routed-expert tensor?""" - return _EXPERT_RE.match(name) is not None + return _EXPERT_RE.match(name) is not None or _MTP_EXPERT_RE.match(name) is not None + + +def _expert_bank_specs(args: DeepseekV4Args) -> tuple[dict, int, int]: + """The 4 routed-expert bank specs for THIS rank, plus its slice of ``moe_inter_dim``. + + Tensor parallelism splits the intermediate dim, so each rank allocates and reads + only its own ``I // tp`` rows. This is what divides the 143 GB of host expert banks + across the ranks instead of replicating them. + """ + E, H = args.n_routed_experts, args.dim + # FP4 packs 2 values per byte and carries one e8m0 scale per 32 values, so a rank's + # slice must stay a multiple of 32 on the intermediate axis. + i_local = div_tp(args.moe_inter_dim, "moe_inter_dim", multiple_of=32) + i_lo = tp_info()[0] * i_local + e8m0 = torch.float8_e8m0fnu + specs = { # alloc UNPINNED, fill, then pin-after-fill (skips slow cudaHostAlloc zero-fill) + "gate_up_packed": ((E, 2 * i_local, H // 2), torch.uint8), + "gate_up_scale": ((E, 2 * i_local, H // 32), e8m0), + "down_packed": ((E, H, i_local // 2), torch.uint8), + "down_scale": ((E, H, i_local // 32), e8m0), + } + return specs, i_local, i_lo -def _place_dsfp4(banks: dict, name: str, t: torch.Tensor, I: int) -> int: +def _place_dsfp4( + banks: dict, key: tuple, t: torch.Tensor, I: int, i_lo: int = 0, + rows_ready: bool = False, +) -> int: """Copy one expert tensor into its layer/expert slot (shared by serial + parallel - readers): w1->gate_up[:,:I], w3->gate_up[:,I:], w2->down. Returns the layer index.""" - m = _EXPERT_RE.match(name) - layer, expert = int(m.group("layer")), int(m.group("expert")) - proj, kind = m.group("proj"), m.group("kind") + readers): w1->gate_up[:,:I], w3->gate_up[:,I:], w2->down. Returns the layer index. + + ``key`` is ``_expert_key``'s ``(bank_layer, expert, proj, kind)``, so a drafter + expert lands in bank layer ``n_layers + k`` with no special case here. + + ``I`` is this rank's intermediate width and ``i_lo`` its offset into the checkpoint's + full width, so each rank keeps only its own rows of w1/w3 and its own columns of w2. + ``rows_ready`` says the caller already read JUST those rows (the serial reader slices + w1/w3 at the file), so they must not be sliced a second time. + """ + layer, expert, proj, kind = key + + def rows(x: torch.Tensor) -> torch.Tensor: + return x if rows_ready else x[i_lo:i_lo + I] + if kind == "weight": t = t.view(torch.uint8) if proj == "w1": - banks["gate_up_packed"][layer][expert, :I] = t + banks["gate_up_packed"][layer][expert, :I] = rows(t) elif proj == "w3": - banks["gate_up_packed"][layer][expert, I:] = t - else: # w2 -> down - banks["down_packed"][layer][expert] = t - else: # scale (e8m0) + banks["gate_up_packed"][layer][expert, I:] = rows(t) + else: # w2 -> down; the intermediate axis is packed 2-per-byte + banks["down_packed"][layer][expert] = t[:, i_lo // 2:(i_lo + I) // 2] + else: # scale (e8m0), one per 32 values on the intermediate axis if proj == "w1": - banks["gate_up_scale"][layer][expert, :I] = t + banks["gate_up_scale"][layer][expert, :I] = rows(t) elif proj == "w3": - banks["gate_up_scale"][layer][expert, I:] = t + banks["gate_up_scale"][layer][expert, I:] = rows(t) else: - banks["down_scale"][layer][expert] = t + banks["down_scale"][layer][expert] = t[:, i_lo // 32:(i_lo + I) // 32] return layer @@ -302,27 +437,20 @@ def load_dsfp4_expert_sources_parallel( from freetoken.models.weight import iter_expert_tensors_parallel from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline, alloc_layer_banks - L, E = args.n_layers, args.n_routed_experts - H, I = args.dim, args.moe_inter_dim - e8m0 = torch.float8_e8m0fnu - specs = { - "gate_up_packed": ((E, 2 * I, H // 2), torch.uint8), - "gate_up_scale": ((E, 2 * I, H // 32), e8m0), - "down_packed": ((E, H, I // 2), torch.uint8), - "down_scale": ((E, H, I // 32), e8m0), - } + # L covers the target layers plus, when dSpark is enabled, the drafter's own. + L, E = args.n_moe_layers, args.n_routed_experts + specs, I, i_lo = _expert_bank_specs(args) hb = alloc_layer_banks(specs, L) # lazy host banks (unpinned) banks = {name: [b.tensor for b in hb[name]] for name in specs} def _is_expert(name: str) -> bool: - m = _EXPERT_RE.match(name) - return m is not None and int(m.group("layer")) < L # skip the MTP layer (index L) + return _expert_key(name, args) is not None def _load(sink) -> int: tracker = LayerCompletionTracker(E * 6, hb, sink) placed = 0 for name, t in iter_expert_tensors_parallel(model_path, _is_expert, workers=workers, chunk=chunk): - layer = _place_dsfp4(banks, name, t, I) + layer = _place_dsfp4(banks, _expert_key(name, args), t, I, i_lo) tracker.note(layer) placed += 1 return placed diff --git a/python/freetoken/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py index b96205a..f1ddebd 100644 --- a/python/freetoken/moe/cpu_executor.py +++ b/python/freetoken/moe/cpu_executor.py @@ -17,6 +17,7 @@ from __future__ import annotations +import glob import os import threading import time @@ -88,8 +89,61 @@ def compiled_extension_supports(activation: str) -> bool: return _ACT_IDS[activation] <= getattr(_cpu_moe, "max_generic_act_id", lambda: 2)() +def _tp_core_share(cores: list[int]) -> list[int]: + """This rank's slice of ``cores`` when several TP ranks share one host. + + Every rank sees the same affinity mask, so without this each rank pins its whole + worker pool to the SAME physical cores: N ranks then oversubscribe those cores N + times and the spin-barrier collapses. Ranks take contiguous blocks rather than + striding, which keeps one rank's threads on one socket. + + A leftover from an uneven split stays unused: giving it to some ranks and not + others would desynchronize the per-layer CPU expert compute that every rank + finishes before the same collective. + """ + from freetoken.distributed import try_get_tp_info + + info = try_get_tp_info() + if info is None or info.size <= 1: + return cores + # ``cores`` is already filtered by this process's affinity. When the engine has bound + # each rank to a NUMA node, that mask is ONE node's cpus and only the ranks on that + # node compete for them -- dividing by the full TP size would hand each rank a + # fraction of a fraction and leave most of the socket idle. + from freetoken.utils.numa import placement as numa_placement + + placed = numa_placement() + siblings, index = ( + (placed[2], placed[3]) if placed is not None else (info.size, info.rank) + ) + per = len(cores) // siblings + if per < 1: # fewer cores than ranks: let the OS schedule them + return cores + return cores[index * per: (index + 1) * per] + + +def _smt_siblings_of(cores: list[int]) -> list[int]: + """Every logical CPU that shares a physical core with one of ``cores``.""" + out: list[int] = [] + for cpu in cores: + try: + with open(f"/sys/devices/system/cpu/cpu{cpu}/topology/thread_siblings_list") as f: + spec = f.read().strip() + except OSError: + out.append(cpu) + continue + for part in spec.split(","): + if "-" in part: + lo, hi = part.split("-") + out.extend(range(int(lo), int(hi) + 1)) + elif part: + out.append(int(part)) + return sorted(set(out)) + + def physical_core_cpus() -> list[int]: - """One logical CPU per physical core, restricted to this process's affinity. + """One logical CPU per physical core, restricted to this process's affinity and, + under tensor parallelism, to this rank's share of them. MoE decode is memory-bandwidth-bound, so SMT siblings only contend for the same core's load ports without adding bandwidth. Picking one logical CPU per @@ -112,7 +166,7 @@ def physical_core_cpus() -> list[int]: if key not in seen: seen.add(key) reps.append(cpu) - return reps or allowed or [0] + return _tp_core_share(reps or allowed or [0]) def resolve_threads_and_affinity(requested: int) -> tuple[int, list[int]]: @@ -127,12 +181,10 @@ def resolve_threads_and_affinity(requested: int) -> tuple[int, list[int]]: reps = physical_core_cpus() if requested and requested > 0: n = int(requested) - try: - allowed = sorted(os.sched_getaffinity(0)) - except AttributeError: - allowed = list(range(os.cpu_count() or 1)) - # physical-core reps first, then the rest of the logical CPUs. - order = reps + [c for c in allowed if c not in set(reps)] + # physical-core reps first, then the SMT siblings OF THOSE CORES. Taking the + # siblings of this rank's own cores (rather than the rest of the affinity mask) + # keeps an explicit thread count from spilling onto another TP rank's cores. + order = reps + [c for c in _smt_siblings_of(reps) if c not in set(reps)] if not order: order = [0] core_ids = [order[i % len(order)] for i in range(n)] @@ -310,9 +362,16 @@ def __init__( "(bit-identical grid; the CPU-side scalar round-trip is skipped)" ) - logger.info_rank0( - f"CPU MoE executor ready: threads={nthreads} (pinned to cores " - f"{core_ids[0]}..{core_ids[-1]}) isa={self.isa} fmt={fmt} " + # Per-rank, not rank0-only: this pool sits on the decode critical path, and the + # way its core split goes wrong is ranks landing on each other's cores. That is + # only visible when every rank says where it went. + from freetoken.distributed import try_get_tp_info + + info = try_get_tp_info() + who = f"rank{info.rank}/{info.size}" if info is not None else "single" + logger.info( + f"NUMA {who}: CPU MoE pool threads={nthreads} on 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}" ) diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index debea28..a28b561 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -123,7 +123,6 @@ class OffloadMoeCache: # pcie_bw / cpu_bw ratio so the PCIe fetch and the CPU overflow GEMV take equal # time (perfect overlap): fetched : cpu = pcie : cpu - pcie. hybrid_fetch_fraction: float = 0.0 - def __post_init__(self) -> None: policy_ids = {"lru": 0} assert self.cache_policy in policy_ids @@ -786,9 +785,24 @@ def ensure_experts_hybrid(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 ensure_experts_hybrid( - self, layer_id, expert_ids, self.hybrid_max_fetch, self.hybrid_fetch_fraction + self, layer_id, expert_ids, self.hybrid_max_fetch, + self._fetch_fraction_for(layer_id, expert_ids), ) + def _fetch_fraction_for(self, layer_id: int, expert_ids: torch.Tensor) -> float: + """Return Equation 4's measured split, unchanged for every exact MoE layer. + + FreeToken defines ``m`` as the number of missing experts in the current layer + and chooses ``q* = m * B_PCIe / B_Host``. ``ensure_experts_hybrid`` already + deduplicates the routes of the whole token batch to obtain that ``m`` before it + applies this fraction. Scaling the ratio again by the token count double-counts + the wider batch and can collapse a speculative verify into pure PCIe fetching. + + ``layer_id`` and ``expert_ids`` remain arguments because the cache-control call + site has this interface; neither changes the hardware bandwidth ratio. + """ + return self.hybrid_fetch_fraction + def materialize_layer(self, layer_id: int) -> None: from freetoken.moe.offload_kernels import materialize_layer diff --git a/python/freetoken/moe/offload_kernels.py b/python/freetoken/moe/offload_kernels.py index cf513f5..e1f2265 100644 --- a/python/freetoken/moe/offload_kernels.py +++ b/python/freetoken/moe/offload_kernels.py @@ -8,8 +8,11 @@ from flashlib.kernels.slot_cache import lru_ensure # Hybrid backend: which of a step's missing experts to fetch (when capped below the miss -# count). "recency" (default) fetches the experts most-recently active before this step -# (LRU on the expert -> prioritizes recurring misses, lowering the steady miss rate); +# count). "recency" (default) first fetches misses used by the most routes in THIS +# batch, then the experts most-recently active before this step (LRU on the expert). +# The route-count priority spends each paper-derived q* PCIe copy where it removes the +# most repeated CPU GEMVs during a multi-row speculative verify; recency breaks ties and +# preserves the steady-state hit-rate policy. # "lowest_id" fetches the smallest expert ids (the original, routing-blind heuristic). _HYBRID_FETCH_BY_RECENCY = ( os.getenv("FREETOKEN_HYBRID_FETCH", "recency").strip().lower() != "lowest_id" @@ -132,9 +135,11 @@ def _ensure_experts_hybrid_cpu( the GPU path; see tests/test_offload_lru_kernels.py). Fetches at most ``max_fetch`` (or the bandwidth-matched ``~frac_q16/2^16 * misses`` when ``frac_q16`` > 0) of the missing experts; overflow misses are rewritten to -1. With ``BY_RECENCY`` the fetch set is the - most-recently-active misses (ties -> lower id); else the lowest ids.""" + highest-current-route-count, most-recently-active misses (ties -> lower id); + else the lowest ids.""" + flat_experts = expert_ids.view(-1).tolist() seen = [] - for expert in expert_ids.view(-1).tolist(): + for expert in flat_experts: if expert not in seen: seen.append(expert) @@ -150,9 +155,10 @@ def _ensure_experts_hybrid_cpu( cache.usage[slot] = step missing = [e for e in seen if int(cache.slot_for_id[layer_id, e].item()) == -1] + route_counts = {e: flat_experts.count(e) for e in seen} if _HYBRID_FETCH_BY_RECENCY: rec = cache.expert_recency[layer_id].tolist() - missing.sort(key=lambda e: (-rec[e], e)) + missing.sort(key=lambda e: (-route_counts[e], -rec[e], e)) else: missing.sort() if frac_q16 > 0: @@ -321,10 +327,11 @@ def _ensure_experts_hybrid_kernel( = the capped fetch count (copy_missing), ``num_missing_full`` = the pre-cap miss count (stats). - Which misses to fetch is the cap policy. ``BY_RECENCY`` (default) fetches the experts - most-recently active before this step (LRU on the expert, via ``expert_recency``), - breaking ties toward the lower expert id -- this prioritizes *recurring* misses for - caching, lowering the steady miss rate. Otherwise the lowest expert ids are fetched + Which misses to fetch is the cap policy. ``BY_RECENCY`` (default) first prioritizes + the experts used by the most routes in the current batch, so one PCIe copy removes + as many repeated CPU GEMVs as possible during a speculative verify. It then uses + most-recent activity (LRU on the expert, via ``expert_recency``), breaking final ties + toward the lower expert id. Otherwise the lowest expert ids are fetched (``missing_rank``), the original routing-blind heuristic.""" step = tl.load(step_ptr) + 1 tl.store(step_ptr, step) @@ -333,10 +340,11 @@ def _ensure_experts_hybrid_kernel( # ---- Phase 1: active + missing over experts ---- off_e = tl.arange(0, BLOCK_E) e_mask = off_e < num_experts - is_active = tl.zeros((BLOCK_E,), dtype=tl.int1) + route_count = tl.zeros((BLOCK_E,), dtype=tl.int32) for i in tl.range(num_active): e = tl.load(expert_ids_ptr + i) - is_active = is_active | (off_e == e) + route_count += (off_e == e).to(tl.int32) + is_active = route_count > 0 tl.store(active_mask_ptr + off_e, is_active.to(tl.int32), mask=e_mask) slot = tl.load(slot_for_id_ptr + base + off_e, mask=e_mask, other=-1) is_missing = is_active & (slot == -1) & e_mask @@ -358,13 +366,17 @@ def _ensure_experts_hybrid_kernel( is_hit = is_active & (slot >= 0) tl.store(usage_ptr + slot, step, mask=is_hit) - # Fetch-selection priority: encode (recency desc, id asc) into one strictly-ordered - # score so argmax has no ties (rec deltas are multiples of num_experts; the id term - # spans only [0, num_experts), so it can only break exact-recency ties). + # Fetch-selection priority: encode (current route count desc, recency desc, id asc) + # into one strictly-ordered score. Previous recency is <= step-1, so multiplying + # route_count by step+1 makes one additional current route dominate the full recency + # range. The id term spans only [0, num_experts), breaking exact ties only. if BY_RECENCY: rec = tl.load(expert_recency_ptr + base + off_e, mask=e_mask, other=-1).to(tl.int64) score = tl.where( - is_missing, rec * num_experts + (num_experts - 1 - off_e), -1152921504606846976 + is_missing, + (route_count.to(tl.int64) * (step + 1) + rec) * num_experts + + (num_experts - 1 - off_e), + -1152921504606846976, ).to(tl.int64) else: missing_rank = tl.cumsum(is_missing.to(tl.int32)) - 1 diff --git a/python/freetoken/scheduler/cache.py b/python/freetoken/scheduler/cache.py index 9be235b..588cc6b 100644 --- a/python/freetoken/scheduler/cache.py +++ b/python/freetoken/scheduler/cache.py @@ -253,6 +253,41 @@ def _free_swa(self, indices: torch.Tensor) -> None: if self.swa_pool is not None and len(indices) > 0: self.swa_pool.free_swa(indices) + def release_speculative_tail(self, req: Req, new_device_len: int) -> None: + """Give back the pages and SWA slots of a rejected block's abandoned positions. + + allocate_paged sizes itself from ``req.device_len``, so a speculative step + allocates for the whole block up front. Acceptance then LOWERS device_len to the + prefix it kept -- and the pages above it are simply forgotten. Nothing else + reclaims them: the per-step window free walks the live range, and the + end-of-request free walks the page table from the request's CURRENT length, so + the tail is outside both. + + The result is a slow leak that never fails a decode. It surfaces at the next + idle integrity check, as far from the cause as it is possible to get: + + SWA-slot leak/double-free: free(11520) + tree(256) != capacity(12160) + + Call this BEFORE lowering device_len -- it reads the current value to find the + range to return. + """ + first_page = div_ceil(new_device_len, self.page_size) + last_page = div_ceil(req.device_len, self.page_size) + if last_page <= first_page or req.table_idx < 0: + return + # The page table is indexed by TOKEN POSITION and holds token slots (see + # _write_page_table, which fills one entry per position over the page's span) -- + # so the slice runs over positions, not pages, and needs no page->token mapping. + # clone: this is a view into the row, and the next allocation rewrites it -- the + # same aliasing lazy_free_region guards against. + slots = self.page_table[ + req.table_idx, first_page * self.page_size : last_page * self.page_size + ].clone() + if self.swa_paged: + self._free_swa(slots) + # free_slots holds page-START slots, one per page. + self.free_slots = torch.cat([self.free_slots, slots[:: self.page_size]]) + def allocate_paged(self, reqs: List[Req]) -> None: needed_pages = 0 allocation_info: List[Tuple[int, int, int]] = [] diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index 3554116..8e8d3df 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -85,6 +85,8 @@ def __init__(self, config: SchedulerConfig): ) or getattr(self.engine.kv_cache, "sliding_window_size", None), ) self.decode_manager = DecodeManager(config.page_size) + self._speculative = _speculative_config(config) + self._warned_non_greedy = False self.prefill_manager = PrefillManager( self.cache_manager, self.table_manager, self.decode_manager ) @@ -225,6 +227,33 @@ def overlap_loop(self, last_data: ForwardData | None) -> ForwardData | None: # table_idx can have its freshly copied prompt clobbered by the prior occupant's # still-pending output write -- corrupting tokens (e.g. dropping an image # placeholder, which the multimodal merge then rejects). + # Speculation cannot be overlapped as written. The loop below schedules and + # LAUNCHES batch N+1 before _process_last_data applies batch N's tokens. Plain + # decode survives that because complete_one() pre-advances device_len by exactly + # one -- the token count is known before the forward runs, which is why + # device_len legitimately leads the buffer by one. + # + # A speculative block contributes 1..k+1 tokens and the count is only known + # AFTER its verify. Overlapped, N+1 is scheduled as though nothing was accepted, + # so it re-drafts positions N already consumed. Measured directly: + # + # draft ctx: 5 rows, positions 18..22 <- block N + # draft ctx: 5 rows, positions 18..22 <- N+1, same positions, same proposals + # draft ctx: 5 rows, positions 20..24 <- then a jump of two + # + # which duplicates tokens in the reply and makes the acceptance figure + # meaningless. + # + # vLLM overlaps and corrects for this by threading num_rejected into the next + # step's input preparation (valid_ctx_end = ctx_end - num_rejected in its + # dflash kernel). Without that correction the only correct option is to drain + # first, which costs the overlap but not the answer. + if self._speculative is not None and last_data is not None: + self.stream.wait_stream(self.engine.stream) + self._process_last_data(last_data) + self._flush_abort_acks() + last_data = None + self.stream.wait_stream(self.engine.stream) forward_input = self._schedule_next_batch() ongoing_data = None @@ -243,8 +272,10 @@ def overlap_loop(self, last_data: ForwardData | None) -> ForwardData | None: # full_to_window INSIDE the captured graph, so an unordered drain can redirect an # in-flight forward. copy_done only covers batch N; order against N+1 explicitly. self.stream.wait_stream(self.engine.stream) - self._process_last_data(last_data) - self._flush_abort_acks() + # None when the speculative drain above already handled it. + if last_data is not None: + self._process_last_data(last_data) + self._flush_abort_acks() return ongoing_data def normal_loop(self) -> None: @@ -331,37 +362,75 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: # are freed below/already; shipping this token would append past the # client's terminal reply. continue - next_token = next_tokens_cpu[i] - req.append_host(next_token.unsqueeze(0)) - next_token = int(next_token.item()) + spec_emitted = getattr(batch, "spec_emitted", None) + if spec_emitted is not None: + # A speculative step already wrote the accepted block into + # req.input_ids and moved the counters, so this must NOT append + # again. The block is reported to the client as a unit; the EOS and + # stop checks below run on its final token, and any earlier EOS is + # caught by the trimming the detokenizer already does. + block = spec_emitted[i] + reply_tokens = block + next_token = int(block[-1].item()) + else: + next_token = next_tokens_cpu[i] + req.append_host(next_token.unsqueeze(0)) + next_token = int(next_token.item()) + reply_tokens = None # EOS / stop-string -> "stop", output budget exhausted -> "length"; # EOS and stop strings win over length. - hit_length = not req.can_decode - hit_eos = ( - not req.sampling_params.ignore_eos and next_token in self.eos_token_ids - ) - matched_stop = ( - self._match_stop_str(req) - if not hit_eos and req.sampling_params.stop_strs - else None - ) - finished = hit_length or hit_eos or matched_stop is not None - finish_reason = ( - ("stop" if (hit_eos or matched_stop is not None) else "length") - if finished - else None + # A speculative step emits a whole accepted block. Every token needs its + # own stop check: an EOS at block position 2 must end the request there, + # not after position 5 -- otherwise the client receives text generated + # past the end of the reply. + block = ( + [int(t) for t in reply_tokens] + if reply_tokens is not None + else [next_token] ) - if ( - next_token == self.toolcall_anchor_id - and req.toolcall_anchor_len is None - and not finished - ): - req.toolcall_anchor_len = req.input_ids.numel() + block_start = req.input_ids.numel() - len(block) + finished = False + visible_tokens: list[int] = [] + finish_reason = None + matched_stop = None + for pos, tok in enumerate(block): + visible_len = block_start + pos + 1 + hit_length = not req.can_decode and pos == len(block) - 1 + hit_eos = ( + not req.sampling_params.ignore_eos and tok in self.eos_token_ids + ) + matched_stop = ( + self._match_stop_str(req, end_len=visible_len) + if not hit_eos and req.sampling_params.stop_strs + else None + ) + finished = hit_length or hit_eos or matched_stop is not None + finish_reason = ( + ("stop" if (hit_eos or matched_stop is not None) else "length") + if finished + else None + ) + if ( + tok == self.toolcall_anchor_id + and req.toolcall_anchor_len is None + and not finished + ): + req.toolcall_anchor_len = visible_len + visible_tokens.append(tok) + if finished: + break # drop the rest of the block: it is past the reply's end + # vLLM carries all newly accepted speculative ids in one engine + # output and advances its incremental detokenizer sequentially. Do + # the same here. Sending one DetokenizeMsg per token in a single + # BatchTokenizerMsg made the old batched decoder append cumulative + # prefixes ("pres" + "presidencia" + ...). + next_token = visible_tokens[-1] reply.append( DetokenizeMsg( uid=req.uid, next_token=next_token, finished=finished, + token_ids=visible_tokens if len(visible_tokens) > 1 else None, finish_reason=finish_reason, matched_stop=matched_stop, stop_strs=req.sampling_params.stop_strs or None, @@ -414,17 +483,20 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: ) self.send_result(reply) - def _match_stop_str(self, req: Req) -> str | None: + def _match_stop_str(self, req: Req, end_len: int | None = None) -> str | None: """First stop string present in this request's generated tail, else None. Decodes only a short suffix (bounded by the longest stop string's char length, so a stop of - N chars spans at most N tokens) to keep the per-step cost small.""" + N chars spans at most N tokens) to keep the per-step cost small. ``end_len`` + makes a speculative block visible one token at a time: looking at the already- + appended whole block would let a stop in token 5 terminate the reply at token 1.""" stop_strs = req.sampling_params.stop_strs prompt_len = req.max_device_len - req.output_len - if len(req.input_ids) <= prompt_len: + end_len = len(req.input_ids) if end_len is None else end_len + if end_len <= prompt_len: return None max_chars = max(len(s) for s in stop_strs) - tail_start = max(prompt_len, len(req.input_ids) - (max_chars + 1)) - tail = self.tokenizer.decode(req.input_ids[tail_start:].tolist()) + tail_start = max(prompt_len, end_len - (max_chars + 1)) + tail = self.tokenizer.decode(req.input_ids[tail_start:end_len].tolist()) for s in stop_strs: if s in tail: return s @@ -803,10 +875,19 @@ def _prepare_batch(self, batch: Batch) -> ForwardInput: # built once here instead of rebuilt in each of the 30 GDN layers. For decode # under CUDA graph the persistent cu_seqlens buffer is supplied by set_batch. batch.fla_metadata = build_fla_metadata(batch, self.device) - if batch.is_decode: + if batch.is_decode or getattr(batch, "speculative", False): # This batch's padded per-row page-table rows. Backends that snapshot the table for # a captured replay (DSV4) read them in prepare_metadata / prepare_for_replay. - batch.active_table_idx = input_mapping[0].view(-1) + # A DSpark verify is phase="prefill" but its fixed-shape target graph uses + # the same device-addressed snapshot contract. + if getattr(batch, "speculative", False): + batch.active_table_idx = torch.tensor( + [r.table_idx for r in batch.padded_reqs], + dtype=torch.int64, + device=self.device, + ) + else: + batch.active_table_idx = input_mapping[0].view(-1) self.engine.attn_backend.prepare_metadata(batch) return ForwardInput( batch=batch, @@ -833,10 +914,62 @@ def _schedule_next_batch(self) -> ForwardInput | None: ) if batch is None: return None + self._maybe_make_speculative(batch) forward_input = self._prepare_batch(batch) self._report_prompt_admissions(batch) return forward_input + def _maybe_make_speculative(self, batch: Batch) -> None: + """Turn a decode batch into a speculative block, in place. + + A verify step is not a special kind of forward -- it is a PREFILL of ``1 + k`` + tokens per request, each resuming from its own position. The engine already + builds exactly that from ``extend_len`` and ``cached_len``, so extending the + requests and flipping the phase is the whole change: allocation, positions, the + compressor carry and the per-position logits all follow. + + The extra ``k`` tokens are the checkpoint's noise placeholder. Slots depend on + POSITIONS, not on token values, so ``allocate_paged`` can hand out every layer's + slots before the drafter has decided what the tokens are; ``_forward`` fills + them in immediately afterwards. + """ + # getattr, not attribute access: the accounting tests drive this method on a + # partially built Scheduler, and speculation must be an opt-in that a stripped + # object simply does not have rather than something that breaks it. + spec = getattr(self, "_speculative", None) + if not hasattr(self, "_warned_non_greedy"): + self._warned_non_greedy = False + if spec is None or not batch.is_decode: + return + should_speculate = getattr( + getattr(self, "engine", None), "should_speculate", None + ) + if should_speculate is not None and any( + not should_speculate(req) for req in batch.reqs + ): + return + gamma = spec.block_size + if gamma < 1: + return + for req in batch.reqs: + # Room for gamma proposals AND the bonus token; equality fits exactly. + # A request with less room decodes normally rather than being truncated. + if req.remain_len < gamma + 1: + return + # The target verify has anchor + gamma proposal rows. The drafter takes the + # first gamma of those rows -- anchor + gamma-1 masks -- exactly as DSpark's + # sample-from-anchor layout specifies. + noise = torch.full((gamma,), spec.noise_token_id, dtype=self.token_pool.dtype) + for req in batch.reqs: + req.append_host(noise) + req.device_len += gamma + batch.phase = "prefill" + batch.speculative = True + batch.spec_block = gamma + # device_len now covers the whole block, so allocate_paged will size for it. + # Acceptance keeps only a prefix and must give the rest back through this. + batch.release_tail = self.cache_manager.release_speculative_tail + def _report_prompt_admissions(self, batch: Batch) -> None: """Publish first-prefill accounting only after batch preparation succeeded. @@ -865,12 +998,59 @@ def _forward(self, forward_input: ForwardInput) -> ForwardOutput: batch.input_ids = self.token_pool[input_mapping] if self.toolcall_anchor_id is not None and not batch.is_prefill: self.cache_manager.snapshot_toolcall_anchor(batch.reqs) + if getattr(batch, "speculative", False): + # Draft BEFORE the verify forward, over the same prepared batch: the draft + # replaces this batch's placeholder token ids in place, so the verify that + # follows scores the drafter's actual proposal. + batch.draft_confidence = self.engine.draft_into_batch(batch) + self.engine.adapt_speculative_batch(batch) forward_output = self.engine.forward_batch(batch, sample_args) - self.token_pool[output_mapping] = forward_output.next_tokens_gpu + if getattr(batch, "speculative", False): + # output_mapping was prepared before acceptance, when device_len still + # pointed past the whole gamma-wide verify span. A rejection moves each + # request's live bonus position independently; writing through the stale + # mapping leaves noise at the next cycle's anchor. Rebuild the tiny mapping + # from the post-acceptance frontiers. + rows = torch.tensor( + [req.table_idx for req in batch.reqs], + dtype=torch.int64, + device=self.device, + ) + positions = torch.tensor( + [req.device_len - 1 for req in batch.reqs], + dtype=torch.int64, + device=self.device, + ) + self.token_pool[rows, positions] = forward_output.next_tokens_gpu + else: + self.token_pool[output_mapping] = forward_output.next_tokens_gpu self.decode_manager.filter_reqs(forward_input.batch.reqs) return forward_output +class _SpeculativeConfig(NamedTuple): + block_size: int + noise_token_id: int + + +def _speculative_config(config) -> "_SpeculativeConfig | None": + """The run's speculative settings, or None when it will not speculate. + + Reads the MODEL's own numbers rather than taking a block size on the command line: + dSpark's block width and noise token are properties of the checkpoint, and a run + that used different ones would put the drafter off its training distribution. + """ + args = getattr(config.model_config, "dsv4_args", None) + if args is None or not getattr(args, "dspark_enabled", False): + return None + if not args.has_dspark: + return None + return _SpeculativeConfig( + block_size=args.dspark_block_size, + noise_token_id=args.dspark_noise_token_id, + ) + + def _make_positions(batch: Batch, device: torch.device) -> torch.Tensor: needed_size = sum(r.extend_len for r in batch.padded_reqs) indices_host = torch.empty(needed_size, dtype=torch.int32, pin_memory=True) diff --git a/python/freetoken/scheduler/status.py b/python/freetoken/scheduler/status.py index ca706c5..1bede18 100644 --- a/python/freetoken/scheduler/status.py +++ b/python/freetoken/scheduler/status.py @@ -16,6 +16,11 @@ class SchedulerStatusReporter: _last_decode_time: float = field(init=False) _decode_forward_count: int = field(default=0, init=False) _decode_generated_tokens: int = field(default=0, init=False) + # Speculative accounting. Acceptance rate is the one number that says whether + # speculation is paying for itself; tokens/s alone cannot distinguish a good + # drafter from a cheap one that is always rejected. + _spec_accepted: int = field(default=0, init=False) + _spec_drafted: int = field(default=0, init=False) def __post_init__(self) -> None: now = self.clock() @@ -35,7 +40,10 @@ def report_batch( mamba_slots: tuple[int, int] | None = None, swa_tokens: tuple[int, int] | None = None, ) -> None: - if batch.is_prefill: + # A speculative step rides the PREFILL path but IS a decode step: reporting + # it as a prefill would hide the tokens it produced and make decode + # throughput read as zero while generation is clearly happening. + if batch.is_prefill and not getattr(batch, "speculative", False): self._report_prefill( batch, running_reqs=running_reqs, @@ -45,7 +53,7 @@ def report_batch( mamba_slots=mamba_slots, swa_tokens=swa_tokens, ) - elif batch.is_decode: + else: self._report_decode( batch, running_reqs=running_reqs, @@ -103,7 +111,16 @@ def _report_decode( swa_tokens: tuple[int, int] | None = None, ) -> None: self._decode_forward_count += 1 - self._decode_generated_tokens += len(batch.reqs) + emitted = getattr(batch, "spec_emitted", None) + if emitted is not None: + # A speculative step emits a whole accepted block per request, not one + # token -- counting requests would understate throughput by the very + # factor speculation is supposed to deliver. + self._decode_generated_tokens += sum(int(e.numel()) for e in emitted) + self._spec_accepted += sum(int(e.numel()) - 1 for e in emitted) + self._spec_drafted += batch.spec_block * len(batch.reqs) + else: + self._decode_generated_tokens += len(batch.reqs) if self._decode_forward_count % self.decode_log_interval != 0: return @@ -120,6 +137,7 @@ def _report_decode( f"{_swa_msg(swa_tokens)}" f"{_mamba_msg(mamba_slots)}" f"gen throughput (token/s): {gen_throughput:.2f}, " + f"{_spec_msg(self._spec_accepted, self._spec_drafted)}" f"#queue-req: {queue_reqs}" ) @@ -142,3 +160,14 @@ def _swa_msg(swa_tokens: tuple[int, int] | None) -> str: return "" used, total = swa_tokens return f"#swa-token: {used}/{total}, swa usage: {_usage_ratio(used, total):.2f}, " + + +def _spec_msg(accepted: int, drafted: int) -> str: + """Speculative acceptance, or nothing at all when not speculating. + + accepted/drafted is what decides whether the drafter earns its cost: a block that + is usually rejected still pays for the draft pass and the wider verify. + """ + if drafted <= 0: + return "" + return f"spec: {accepted}/{drafted} accepted ({100 * accepted / drafted:.0f}%), " diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index 80dbf6a..42d04e8 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -365,6 +365,30 @@ async def stream_with_cancellation(self, generator, request: Request, uid: int): asyncio.create_task(self.abort_user(uid)) raise + async def wait_for_ack_watching_client(self, uid: int, request): + """``wait_for_ack``, but abort the generation when the client goes away. + + Streaming responses already notice a disconnect, because yielding a chunk into a + dead socket is observable. A non-streaming request awaits the whole completion + and never looks, so a client that hangs up leaves the GPUs generating to + max_tokens for nobody -- and the queue behind it waiting on work whose answer + will be discarded. + + ``request`` may be None for callers with no HTTP request in scope (offline, + internal); then this is exactly ``wait_for_ack``. + """ + try: + async for ack in self.wait_for_ack(uid): + if request is not None and await request.is_disconnected(): + logger.info("Client disconnected for user %s; aborting", uid) + raise asyncio.CancelledError + yield ack + except asyncio.CancelledError: + # create_task, not await: the caller is unwinding, and abort_user sleeps + # before it sends. Mirrors stream_with_cancellation. + asyncio.create_task(self.abort_user(uid)) + raise + async def abort_user(self, uid: int): await asyncio.sleep(0.1) if uid in self.ack_map: @@ -928,6 +952,23 @@ def run_api_server(config: ServerArgs, start_backend: Callable[[], "Any"], run_s global _GLOBAL_STATE, _MODEL_SAMPLING + # Before anything slow. Loading takes ~90s, and a log that opens part-way through it + # cannot answer what the run was configured as -- which is the first question asked + # of every one of these logs. + from freetoken.utils.banner import print_banner + + # Every field via getattr: these kwargs are evaluated BEFORE print_banner's own + # try/except, so one wrong attribute name here takes the whole server down at + # startup. A banner must not be able to do that. + print_banner( + model_path=getattr(config, "model_path", None), + tp_size=getattr(getattr(config, "tp_info", None), "size", 1) or 1, + dspark=bool(getattr(config, "speculative_dspark", False)), + moe_backend=getattr(config, "moe_backend", None), + host=getattr(config, "server_host", None), + port=getattr(config, "server_port", None), + ) + if config.sampling_defaults == "model" and not config.use_dummy_weight: _MODEL_SAMPLING = load_generation_sampling(config.model_path) # Always surface the effective default sampling (model-recommended where available, diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index b2857b7..c3c67f1 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -221,6 +221,56 @@ def _infer_reasoning_parser(model_path: str) -> str | None: help="The tensor parallelism size.", ) + parser.add_argument( + "--speculative-dspark", + action="store_true", + dest="speculative_dspark", + default=False, + help=( + "DeepSeek-V4 only: build the checkpoint's dSpark drafter (the mtp.* stack) " + "for block speculative decoding. Off by default because the drafter's own " + "routed experts enlarge the host expert banks and the GPU slot cache. Fails " + "at config time if the checkpoint ships no dSpark weights." + ), + ) + + parser.add_argument( + "--dspark-fallback-acceptance", + type=_parse_moe_cache_rate, + default=ServerArgs.dspark_fallback_acceptance, + help=( + "Experimental request-local DSpark fallback threshold in [0, 1]. " + "After measured proposal acceptance falls below this value, temporarily " + "use ordinary target decode; 0 disables the fallback." + ), + ) + + parser.add_argument( + "--dspark-fallback-min-drafted", + type=_positive_int, + default=ServerArgs.dspark_fallback_min_drafted, + help="Minimum measured DSpark proposals before evaluating the fallback.", + ) + + parser.add_argument( + "--dspark-fallback-steps", + type=_positive_int, + default=ServerArgs.dspark_fallback_steps, + help="Ordinary target-decode steps before probing DSpark again.", + ) + + parser.add_argument( + "--distributed-timeout", + type=float, + default=ServerArgs.distributed_timeout, + help=( + "Seconds a tensor-parallel collective waits before failing. The first one is " + "the post-load memory all-reduce, which ranks reach minutes apart on a large " + "MoE checkpoint, so the default is sized for that load skew; lower it to make " + "a hung rank fail faster." + ), + ) + parser.add_argument( "--max-running-requests", type=int, diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index be05d90..a55ad22 100644 --- a/python/freetoken/server/generation.py +++ b/python/freetoken/server/generation.py @@ -490,6 +490,19 @@ def _record_generation( ) +def _acks(state, uid, request): + """The ack stream, watching the client when both sides support it. + + getattr, not a bare call: a caller may pass a state that predates the watcher (the + test fakes do), and a missing disconnect check should cost the abort optimisation, + not the request. + """ + watch = getattr(state, "wait_for_ack_watching_client", None) + if request is not None and watch is not None: + return watch(uid, request) + return state.wait_for_ack(uid) + + async def generate_events( uid: int, spec: GenSpec, state: Any, *, source: str | None = None ) -> AsyncIterator[GenEvent]: @@ -521,7 +534,7 @@ async def generate_events( async def generate_full( - uid: int, spec: GenSpec, state: Any, *, source: str | None = None + uid: int, spec: GenSpec, state: Any, *, source: str | None = None, request: Any = None ) -> GenResult: """Wraps `_generate_full_impl` to log the request with its totals; the `finally` also records a `GenerationError` as a failed row.""" @@ -529,7 +542,7 @@ async def generate_full( result: GenResult | None = None error: str | None = None try: - result = await _generate_full_impl(uid, spec, state) + result = await _generate_full_impl(uid, spec, state, request) return result except GenerationError as exc: error = str(exc) @@ -647,6 +660,8 @@ def _route_tool_text(piece: str) -> list[GenEvent]: engine_finish_reason: str | None = None engine_matched_stop: str | None = None + # Streaming: a disconnect is already caught at the transport layer by + # stream_with_cancellation, which sees the write into a dead socket. async for ack in state.wait_for_ack(uid): if getattr(ack, "error", None): raise GenerationError(ack.error, getattr(ack, "error_code", None)) @@ -738,7 +753,9 @@ def _route_tool_text(piece: str) -> list[GenEvent]: ) -async def _generate_full_impl(uid: int, spec: GenSpec, state: Any) -> GenResult: +async def _generate_full_impl( + uid: int, spec: GenSpec, state: Any, request: Any = None +) -> GenResult: """Protocol-neutral non-streaming generation: accumulate, split reasoning, parse tool calls, strip special tokens. The adapters format the GenResult into their wire.""" full_content = "" @@ -747,7 +764,10 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any) -> GenResult: cached_tokens = 0 engine_finish_reason: str | None = None engine_matched_stop: str | None = None - async for ack in state.wait_for_ack(uid): + # Watch the client. A non-streaming request awaits the whole completion, so a + # caller that hangs up otherwise leaves the GPUs generating to max_tokens for an + # answer nobody will read -- with the queue behind it waiting on that work. + async for ack in _acks(state, uid, request): if getattr(ack, "error", None): raise GenerationError(ack.error, getattr(ack, "error_code", None)) prompt_tokens += ack.prompt_tokens_delta diff --git a/python/freetoken/server/launch.py b/python/freetoken/server/launch.py index a5f8438..ef29d95 100644 --- a/python/freetoken/server/launch.py +++ b/python/freetoken/server/launch.py @@ -3,6 +3,7 @@ import logging import multiprocessing as mp import os +import signal import sys from dataclasses import replace from typing import TYPE_CHECKING @@ -62,6 +63,23 @@ def _run_scheduler(args: ServerArgs, ack_queue: mp.Queue[str]) -> None: import torch from freetoken.scheduler import Scheduler + # The parent stops workers with p.terminate(), i.e. SIGTERM, whose default action is + # to kill the process outright. scheduler.shutdown() then never runs: the rank's + # collectives are never torn down, its multiprocessing primitives are never released + # -- the parent's "leaked semaphore objects" warning -- and a rank caught mid-forward + # dies loudly enough to look like a crash rather than a requested stop. + # + # Turning it into KeyboardInterrupt routes SIGTERM into the graceful path that ^C + # already uses, so there is one shutdown path rather than two. + def _terminate(signum, _frame): + raise KeyboardInterrupt(f"signal {signum}") + + for _sig in (signal.SIGTERM, signal.SIGINT): + try: + signal.signal(_sig, _terminate) + except (ValueError, OSError): + pass # not the main thread of this process; the parent's SIGKILL backstop covers it + if args.tp_info.is_primary(): from freetoken.utils.progress import set_progress_sink @@ -73,6 +91,19 @@ def _run_scheduler(args: ServerArgs, ack_queue: mp.Queue[str]) -> None: try: scheduler = Scheduler(args) scheduler.sync_all_ranks() + # All final serving allocations now exist (host expert banks, GPU caches, + # CUDA graphs). Every rank contributes its real placement before rank 0 + # sends the ready ack, so the log maps the loaded model rather than merely + # describing the workstation's hardware. + scheduler.engine.log_distribution_summary() + except KeyboardInterrupt: + # A stop during the ~90s load. KeyboardInterrupt is a BaseException, so it + # slips past the handler below and prints a traceback from wherever the + # loader happened to be -- which reads as a crash, and is what a stopped + # load looked like before this. There is nothing built yet to shut down. + if args.tp_info.is_primary(): + init_logger(__name__).info("Stopped during load; nothing to shut down.") + return except Exception as exc: # noqa: BLE001 -- surface the reason, then let it propagate # A startup failure (bad config, OOM, corrupt weights) would otherwise reach the # parent only as a dead process -> a generic "exited during load". Push the real @@ -104,11 +135,11 @@ def _run_scheduler(args: ServerArgs, ack_queue: mp.Queue[str]) -> None: try: scheduler.run_forever() - except KeyboardInterrupt: + except KeyboardInterrupt as stop: logger = init_logger(__name__) if args.tp_info.is_primary(): print() # for a clean newline after ^C - logger.info("Scheduler exiting gracefully...") + logger.info("Scheduler exiting gracefully (%s)...", stop or "interrupt") scheduler.shutdown() diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index b4becd2..a7dbc40 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -23,6 +23,7 @@ from .function_call_parser import ToolCallItem from .request_logger import log_request from .generation import ( + _acks, ContentDelta, GenDone, GenerationError, @@ -199,7 +200,9 @@ async def handle_chat_completion( return StreamingResponse(chunks, media_type="text/event-stream") try: - result = await generate_full(uid, spec, state, source="/v1/chat/completions") + result = await generate_full( + uid, spec, state, source="/v1/chat/completions", request=request + ) except GenerationError as exc: return create_error_response(str(exc), code=exc.code) message: dict[str, Any] = {"role": "assistant", "content": result.content} @@ -413,7 +416,9 @@ async def handle_completion( await state.send_one(TokenizeMsg(uid=uid, text=prompt, sampling_params=_resolve_sampling(req, model_sampling))) text = "" finish_reason = "stop" - async for ack in state.wait_for_ack(uid): + # Watch the client: a non-streaming request otherwise generates to max_tokens + # for a caller that has already hung up. + async for ack in _acks(state, uid, request): if getattr(ack, "error", None): return create_error_response(ack.error) prompt_tokens += ack.prompt_tokens_delta diff --git a/python/freetoken/tokenizer/detokenize.py b/python/freetoken/tokenizer/detokenize.py index d138ac1..975b00a 100644 --- a/python/freetoken/tokenizer/detokenize.py +++ b/python/freetoken/tokenizer/detokenize.py @@ -90,9 +90,37 @@ def discard(self, uid: int) -> None: accumulated ids and text stay in ``decode_map`` for the life of the worker.""" self.decode_map.pop(uid, None) + def _append_token(self, s: DecodeStatus, token: int, *, skip: bool = False) -> str: + """Advance one request by one token, preserving tokenizer boundary context.""" + if not skip: + s.decoded_ids.append(token) + read_str, surr_str = self.tokenizer.batch_decode( + [ + s.decoded_ids[s.surr_offset :], + s.decoded_ids[s.surr_offset : s.read_offset], + ] + ) + new_text = read_str[len(surr_str) :] + if len(new_text) > 0 and not new_text.endswith("�"): + output_str = s.decoded_str + new_text + s.decoded_str = output_str + s.surr_offset = s.read_offset + s.read_offset = len(s.decoded_ids) + else: + new_text = find_printable_text(new_text) + output_str = s.decoded_str + new_text + return output_str + def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]: - read_ids: List[List[int]] = [] - surr_ids: List[List[int]] = [] + """Detokenize each request's newly emitted token block in order. + + A speculative step can deliver multiple accepted tokens for one uid. The + original two-pass batch implementation appended every token before updating + the uid's offsets, so each decoded prefix was appended cumulatively. vLLM's + speculative output processor advances tokens sequentially; doing so here is + both exact and safe when multiple messages for one uid are coalesced. + """ + incremental_strs: List[str] = [] for msg in msgs: if msg.uid not in self.decode_map: self.decode_map[msg.uid] = DecodeStatus( @@ -103,27 +131,15 @@ def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]: sent_offset=0, ) 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) - read_ids.append(s.decoded_ids[s.surr_offset :]) - surr_ids.append(s.decoded_ids[s.surr_offset : s.read_offset]) - - read_texts = self.tokenizer.batch_decode(read_ids) - surr_texts = self.tokenizer.batch_decode(surr_ids) - - incremental_strs: List[str] = [] - 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) :] - # Streaming chunk: update the decode status - if len(new_text) > 0 and not new_text.endswith("�"): - output_str = s.decoded_str + new_text - s.decoded_str = output_str - s.surr_offset = s.read_offset - s.read_offset = len(s.decoded_ids) - else: - new_text = find_printable_text(new_text) - output_str = s.decoded_str + new_text + tokens = msg.token_ids if msg.token_ids is not None else [msg.next_token] + output_str = s.decoded_str + for pos, token in enumerate(tokens): + final_eos = ( + msg.finished + and pos == len(tokens) - 1 + and token in self.eos_token_ids + ) + output_str = self._append_token(s, token, skip=final_eos) prev_sent = s.sent_offset if msg.finished: diff --git a/python/freetoken/tokenizer/server.py b/python/freetoken/tokenizer/server.py index 530e862..0f69cba 100644 --- a/python/freetoken/tokenizer/server.py +++ b/python/freetoken/tokenizer/server.py @@ -215,7 +215,9 @@ def tokenize_worker( finished=msg.finished, finish_reason=msg.finish_reason, matched_stop=msg.matched_stop, - completion_tokens_delta=1, + completion_tokens_delta=( + len(msg.token_ids) if msg.token_ids is not None else 1 + ), kv_used_pages=msg.kv_used_pages, kv_total_pages=msg.kv_total_pages, mamba_used_slots=msg.mamba_used_slots, diff --git a/python/freetoken/utils/banner.py b/python/freetoken/utils/banner.py new file mode 100644 index 0000000..14b4a3b --- /dev/null +++ b/python/freetoken/utils/banner.py @@ -0,0 +1,282 @@ +"""The startup banner: what this process is, printed before it takes 90 seconds to load. + +A log that opens mid-way through loading experts answers none of the questions asked of +it afterwards. Nearly every confusing run in this engine's history was a configuration +question -- was dSpark on, how many ranks, which NUMA layout, which checkpoint -- and the +answer was derivable from the log only by knowing which incidental line to read. A run +with speculation off is distinguishable from one with it on by the expert-loader's +``layers=43`` versus ``layers=46``, which is not a thing anyone should have to know. + +So the banner states the configuration outright, at the top, where a pasted log excerpt +will usually include it. +""" + +from __future__ import annotations + +import os +import platform +import sys + +# The wordmark, from assets/freetoken-logo-dark.svg. 82 columns. +# +# Halved from a 163-column rendering by folding each 2x2 cell into one quadrant block +# glyph, which keeps the letterforms legible at half the width -- a plain every-other- +# column drop would have shredded the thin strokes. 163 columns overflowed narrower +# monitors, and a banner that wraps is worse than no banner. +# +# It prints in one colour. The glyphs interlock along the logo's slant, so the SVG's +# two-tone split ("ree" in off-white, the rest blue) cannot be reproduced by slicing +# columns without cutting through letters. Blue is the mark's dominant colour. +_ART = [ + " ▜▘█████████▌ ▗▄▄▄▄▄▄▄▄▄ ▄▄▄", + " ██████▀▀▀▀▀ ▄▄ ▄▄ █████████▌ ▄▖ ███ ▗▄ ▄", + " ▗█▌▟███████▌ ██████▌▟█████▙▖ ▟█████▙▖ ███▘ ▗▟█████▙ ▐██▛▄███▀▗██████▖ ████████▖", + " ▗▄▄███████▌ ▗███▀▀▀███▙▄▟██▌▟██▙▄▟██▌ ▐███ ▗███▘ ▜██▌▟█████▀ ▐███▄▄███▗███▀▝███▌", + " ▄▖▄▄███ ▐██▛ ███▀▀▀▀▀▘███▀▀▀▀▀▘ ▟██▌ ▐██▛ ███▘██████▖ ▟██▛▀▀▀▀▀▐██▌ ███", + " ▝ ████▛ ███▘ ▜██████▛ ▜██████▛ ███▘ ▝███████▘▗██▛▝███▄▝██████▛ ███▘ ▐██▛", + " ▀▀▀▀▘ ▀▀▀ ▝▀▀▀▀ ▝▀▀▀▀ ▀▀▀ ▀▀▀▀▀ ▝▀▀▘ ▝▀▀▀ ▀▀▀▀ ▀▀▀ ▝▀▀▘", +] + +# Shown instead when the terminal is too narrow for _ART. A 163-column banner that wraps +# is worse than no banner: every row folds and the logo becomes noise. +_ART_NARROW = [ + " ══ ______ ______ __ ", + " ════ / ____/_______ ___ /_ __/___ / /_____ ____ ", + " ═══ / /_ / ___/ _ \\/ _ \\ / / / __ \\/ //_/ _ \\/ __ \\ ", + " ══ / __/ / / / __/ __// / / /_/ / ,< / __/ / / / ", + " ═ /_/ /_/ \\___/\\___//_/ \\____/_/|_|\\___/_/ /_/ ", +] + +_BLUE = "\033[38;2;88;166;240m" # #58a6f0, the logo's blue +_WHITE = "\033[38;2;242;245;248m" # #f2f5f8, the logo's off-white +_DIM = "\033[2m" +_OFF = "\033[0m" + + +def _use_colour(stream) -> bool: + """Colour only where it will not become literal escape codes in a file. + + A banner is for humans, and this one is usually read in a redirected log. NO_COLOR + is honoured because a startup banner is exactly the sort of thing people redirect + into files and paste into issues. + """ + if os.environ.get("NO_COLOR") is not None: + return False + if os.environ.get("FORCE_COLOR"): + return True + try: + return bool(stream.isatty()) + except Exception: + return False + + +def _art(colour: bool) -> list[str]: + """The wordmark, or the compact one when it will not fit. + + Width is only knowable for a terminal. A redirected log has no width, and the file + holds long lines fine, so the full mark is the default there. + """ + rows = _ART + try: + cols = os.get_terminal_size().columns + if cols and cols < max(len(r) for r in _ART): + rows = _ART_NARROW + except OSError: + pass # not a terminal: a file takes the full mark + if not colour: + return list(rows) + return [f"{_BLUE}{r}{_OFF}" for r in rows] + + +def _gpu_line() -> str | None: + try: + import torch + + if not torch.cuda.is_available(): + return None + n = torch.cuda.device_count() + name = torch.cuda.get_device_name(0) + gib = torch.cuda.get_device_properties(0).total_memory / 1024**3 + # One decimal on both, or a 15.6 GiB card reads as "16 GiB each, 62 GiB + # total" and the arithmetic looks broken. + return f"{n} x {name} ({gib:.1f} GiB each, {n * gib:.1f} GiB total)" + except Exception: + return None + + +def _cpu_line() -> str | None: + """Physical cores, logical CPUs, and how many this process may actually use. + + The last number is the one that matters and the one nobody has: a cpuset or a + taskset makes os.cpu_count() a lie, and the CPU MoE pool sizes itself from the + permitted set, not the machine. + """ + try: + logical = os.cpu_count() or 0 + allowed = len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else logical + phys = None + try: + from freetoken.moe.cpu_executor import physical_core_cpus + + phys = len(physical_core_cpus()) + except Exception: + pass + desc = f"{phys} physical cores, {logical} logical" if phys else f"{logical} logical cpus" + if allowed and allowed != logical: + desc += f" ({allowed} permitted to this process)" + return desc + except Exception: + return None + + +def _ram_line() -> str | None: + """Total and available RAM. + + An offload MoE keeps its expert banks in host memory, so "available" is the number + that decides whether a run fits -- and a SIGKILL'd rank can leave tens of GiB of + pinned Shmem behind, which shows up here as a shortfall with no process to blame. + """ + try: + info = {} + with open("/proc/meminfo") as f: + for line in f: + k, _, rest = line.partition(":") + info[k] = int(rest.split()[0]) # kB + total = info.get("MemTotal", 0) / 1024**2 + avail = info.get("MemAvailable", 0) / 1024**2 + shmem = info.get("Shmem", 0) / 1024**2 + desc = f"{total:.0f} GiB total, {avail:.0f} GiB available" + if shmem >= 1: + desc += f", {shmem:.0f} GiB shmem" + return desc + except Exception: + return None + + +def _numa_line(world_size: int) -> str | None: + """The rank -> node map, which is otherwise only inferable from per-rank log lines.""" + try: + from freetoken.utils import numa + + nodes = numa.numa_nodes() + if len(nodes) < 2: + return f"1 node ({len(nodes[0]) if nodes else '?'} cpus) -- nothing to spread" + placed = [] + for r in range(world_size): + p = numa.rank_placement(r, world_size) + placed.append(str(p[0]) if p else "?") + cpus = sum(len(n) for n in nodes) + return f"{len(nodes)} nodes, {cpus} cpus | rank->node: {' '.join(placed)}" + except Exception: + return None + + +def format_banner( + *, + model_path: str | None = None, + quant: str | None = None, + tp_size: int = 1, + dspark: bool = False, + dspark_block: int | None = None, + dspark_layers: int | None = None, + moe_backend: str | None = None, + host: str | None = None, + port: int | None = None, + colour: bool = False, +) -> str: + from freetoken.version import __version__ + + rows = _art(colour) + d = _DIM if colour else "" + off = _OFF if colour else "" + + # Facts worth a line each. Everything here has cost real debugging time to recover + # from a log that did not say it. + facts: list[tuple[str, str]] = [("version", __version__)] + if model_path: + facts.append(("model", os.path.basename(model_path.rstrip("/")))) + if quant: + facts.append(("experts", quant)) + if moe_backend: + facts.append(("moe", moe_backend)) + facts.append(("parallel", f"TP={tp_size}")) + gpu = _gpu_line() + if gpu: + facts.append(("gpu", gpu)) + cpu = _cpu_line() + if cpu: + facts.append(("cpu", cpu)) + ram = _ram_line() + if ram: + facts.append(("ram", ram)) + numa_desc = _numa_line(tp_size) + if numa_desc: + facts.append(("numa", numa_desc)) + if dspark: + detail = "on" + if dspark_block: + detail += f" (block={dspark_block}" + detail += f", {dspark_layers} draft layers)" if dspark_layers else ")" + facts.append(("dspark", detail)) + else: + facts.append(("dspark", "off -- plain autoregressive decode")) + if host is not None and port is not None: + facts.append(("serving", f"{host}:{port}")) + facts.append( + ("runtime", f"python {platform.python_version()} | {_torch_desc()}") + ) + + width = max(len(k) for k, _ in facts) + lines = list(rows) + lines.append("") + for k, v in facts: + lines.append(f" {d}{k.rjust(width)}{off} {v}") + lines.append("") + return "\n".join(lines) + + +def _torch_desc() -> str: + """torch's build, and BOTH CUDA versions, because they legitimately differ. + + torch.version.cuda is the toolkit the wheel was COMPILED against, not the CUDA the + installed driver provides. A cu130 wheel runs on a 13.3 driver through minor-version + compatibility. Printing only the build version reads as "this box is on the wrong + CUDA" to anyone who knows what their driver reports -- so print both, labelled. + """ + try: + import torch + except Exception: + return "torch unavailable" + + built = torch.version.cuda or "none" + desc = f"torch {torch.__version__} | cuda {built} build" + try: + # Ask libcuda directly. torch._C._cuda_getDriverVersion does not exist in every + # build (it is absent in 2.11.0+cu130), and cuDriverGetVersion is the same + # number nvidia-smi reports, with no subprocess. + import ctypes + + raw = ctypes.c_int(0) + lib = ctypes.CDLL("libcuda.so.1") + if lib.cuDriverGetVersion(ctypes.byref(raw)) != 0: + raise OSError("cuDriverGetVersion failed") + raw = raw.value # packed 1000*major + 10*minor, e.g. 13030 -> 13.3 + if raw: + desc += f" / {raw // 1000}.{(raw % 1000) // 10} driver" + except Exception: + pass + return desc + + +def print_banner(stream=None, **kw) -> None: + """Write the banner. Rank 0 only -- the caller decides that, not this.""" + stream = stream or sys.stderr + try: + stream.write(format_banner(colour=_use_colour(stream), **kw) + "\n") + stream.flush() + except Exception: + # A banner must never be the reason a server fails to start. + pass + + +__all__ = ["format_banner", "print_banner"] diff --git a/python/freetoken/utils/numa.py b/python/freetoken/utils/numa.py new file mode 100644 index 0000000..5a4f55c --- /dev/null +++ b/python/freetoken/utils/numa.py @@ -0,0 +1,182 @@ +"""NUMA topology, and how tensor-parallel ranks should be spread across it. + +Why an engine cares. A rank's host-side working set -- for an offload MoE, tens of GiB +of expert banks -- is anonymous memory, so its pages land on whichever node FIRST +TOUCHES them. If the loader's threads are unbound they scatter across every socket, and +the CPU expert pool, pinned to one socket's cores, then reads much of its data across +the interconnect. On a two-socket Xeon that is the difference between one memory +controller and two: each socket has its own channels, so binding ranks to nodes turns +one socket's bandwidth into the machine's. + +Everything here is read from sysfs and intersected with the process's affinity mask, so +it is correct on a single-socket laptop (no-op), a 2-socket server, a 4- or 8-node +system, and inside a container whose cpuset covers part of a node. +""" + +from __future__ import annotations + +import glob +import os + + +def _parse_cpulist(spec: str) -> set[int]: + """``"0-3,8,12-13"`` -> ``{0,1,2,3,8,12,13}``.""" + cpus: set[int] = set() + for part in spec.split(","): + part = part.strip() + if not part: + continue + if "-" in part: + lo, hi = part.split("-", 1) + cpus.update(range(int(lo), int(hi) + 1)) + else: + cpus.add(int(part)) + return cpus + + +def allowed_cpus() -> set[int]: + """CPUs this process may run on. Falls back to every CPU where unsupported.""" + try: + return set(os.sched_getaffinity(0)) + except AttributeError: # not Linux + return set(range(os.cpu_count() or 1)) + + +def numa_nodes(allowed: set[int] | None = None) -> list[list[int]]: + """CPUs per NUMA node, restricted to what this process may use. + + Nodes with no usable CPU are dropped -- a memory-only node (CXL, or a node whose + cores a cpuset excludes) is not somewhere a rank can run, and treating it as one + would strand a rank with an empty affinity mask. + """ + if allowed is None: + allowed = allowed_cpus() + nodes: list[list[int]] = [] + for path in sorted( + glob.glob("/sys/devices/system/node/node[0-9]*/cpulist"), + key=lambda p: int(p.rsplit("node", 1)[1].split("/", 1)[0]), + ): + try: + with open(path) as f: + cpus = _parse_cpulist(f.read()) + except OSError: + return [] + usable = sorted(cpus & allowed) + if usable: + nodes.append(usable) + return nodes + + +def device_numa_node(device_index: int) -> int | None: + """The NUMA node a CUDA device hangs off, or None if it cannot be determined. + + A GPU sits behind a PCIe root complex owned by ONE socket. Host-to-device traffic + from the other socket crosses the interconnect -- which, for an offload MoE, is the + per-step expert fetch. So a rank's CPU should live on its GPU's node, and that node + is a property of the hardware, not of the rank's index. + """ + try: + import torch + + props = torch.cuda.get_device_properties(device_index) + bdf = f"{props.pci_domain_id:04x}:{props.pci_bus_id:02x}:{props.pci_device_id:02x}.0" + except Exception: + return None + try: + with open(f"/sys/bus/pci/devices/{bdf}/numa_node") as f: + node = int(f.read().strip()) + except (OSError, ValueError): + return None + return node if node >= 0 else None + + +def rank_placement(rank: int, world_size: int) -> tuple[int, list[int], int, int] | None: + """Where rank ``rank`` of ``world_size`` should run. + + Returns ``(node_index, cpus, siblings, index_among_siblings)``, or None when there + is nothing to decide -- one node, one rank, or no readable topology. + + Ranks are dealt to nodes in contiguous blocks, so consecutive ranks share a node and + a node's ranks are adjacent. ``siblings`` is the EXACT number of ranks placed on this + rank's node, not an average: with 3 ranks over 2 nodes one node holds two and the + other one, and a rank that assumed the average would either oversubscribe its cores + or leave half of them idle. + + With more nodes than ranks each rank still gets a whole node, and the surplus nodes + go unused -- spreading one rank's threads over two nodes would recreate the remote + access this exists to avoid. + """ + nodes = numa_nodes() + if len(nodes) < 2 or world_size < 2: + return None + if not 0 <= rank < world_size: + raise ValueError(f"rank {rank} outside world size {world_size}") + + n_nodes = len(nodes) + + # Prefer the node this rank's GPU actually hangs off. Deriving it from the rank + # index instead only works when CUDA enumerates devices in the same order as the + # sockets own them -- true on some machines, and silently wrong on others, where + # every rank would then reach its own card across the interconnect. + gpu_node = device_numa_node(rank) + if gpu_node is not None and gpu_node < n_nodes: + mates = [r for r in range(world_size) if device_numa_node(r) == gpu_node] + if mates: + return gpu_node, nodes[gpu_node], len(mates), mates.index(rank) + if world_size <= n_nodes: + # Spread out: every rank owns a node, evenly spaced so a partly used machine + # still uses distinct memory controllers. + node = (rank * n_nodes) // world_size + return node, nodes[node], 1, 0 + + # More ranks than nodes: contiguous blocks, remainder spread over the first nodes. + base, extra = divmod(world_size, n_nodes) + node, first = 0, 0 + for i in range(n_nodes): + count = base + (1 if i < extra else 0) + if rank < first + count: + node = i + return node, nodes[i], count, rank - first + first += count + raise AssertionError("unreachable: ranks did not cover the node blocks") + + +_PLACEMENT: tuple[int, list[int], int, int] | None = None + + +def resolve_placement(rank: int, world_size: int) -> tuple[int, list[int], int, int] | None: + """``rank_placement``, computed once and remembered for the rest of the process. + + Binding a rank ERASES the evidence the placement was derived from: afterwards the + affinity mask is one node, so ``numa_nodes`` sees a single node and + ``rank_placement`` correctly answers "nothing to spread". Every later consumer -- + the torch thread split, the CPU MoE core split -- would then fall back to dividing + by the whole world size and hand each rank a fraction of a fraction. + + So the answer is computed BEFORE the bind and reused. Call this once, early. + """ + global _PLACEMENT + if _PLACEMENT is None: + _PLACEMENT = rank_placement(rank, world_size) + return _PLACEMENT + + +def placement() -> tuple[int, list[int], int, int] | None: + """The remembered placement, or None if this process was never placed.""" + return _PLACEMENT + + +def reset_placement() -> None: + """Forget the remembered placement (tests).""" + global _PLACEMENT + _PLACEMENT = None + + +__all__ = [ + "allowed_cpus", + "numa_nodes", + "placement", + "rank_placement", + "reset_placement", + "resolve_placement", +] diff --git a/scripts/freetoken-dsv4.sh b/scripts/freetoken-dsv4.sh new file mode 100644 index 0000000..cb25348 --- /dev/null +++ b/scripts/freetoken-dsv4.sh @@ -0,0 +1,303 @@ +#!/usr/bin/env bash +# Start / stop DeepSeek-V4-Flash-0731 on the 4x RTX A4000, served on 0.0.0.0:8081. +# +# ./freetoken-dsv4.sh start start it (returns when it is really ready) +# ./freetoken-dsv4.sh stop stop it cleanly (add --force to escalate) +# ./freetoken-dsv4.sh kill SIGKILL a wedged rank, then reclaim /dev/shm +# ./freetoken-dsv4.sh status is it up, and what is it holding +# ./freetoken-dsv4.sh logs follow the log +# ./freetoken-dsv4.sh test send one request +# +# Override any setting from the environment, e.g. +# MEMORY_RATIO=0.88 ./freetoken-dsv4.sh start + +set -uo pipefail + +FT_DIR="${FT_DIR:-$HOME/FreeToken}" +MODEL="${MODEL:-$HOME/models/DeepSeek-V4-Flash-0731}" +HOST="${HOST:-0.0.0.0}" +PORT="${PORT:-8081}" +TP_SIZE="${TP_SIZE:-4}" +MEMORY_RATIO="${MEMORY_RATIO:-0.90}" +MAX_RUNNING_REQUESTS="${MAX_RUNNING_REQUESTS:-1}" +MAX_PREFILL_LENGTH="${MAX_PREFILL_LENGTH:-1024}" +EXPERT_LOAD="${EXPERT_LOAD:-serial}" +# SPECULATIVE_DSPARK=0 turns OFF the checkpoint's dSpark drafter (the mtp.* stack). +# On by default: the draft/verify loop is wired, and dSpark is what the checkpoint was +# trained to decode with. It costs 9.5 GiB of host expert banks and a slice of the GPU +# slot cache. Set 0 to measure a plain-decode baseline -- the drafter's 3 MoE layers then +# disappear from the log's `layers=` count, which is the quickest way to tell which +# mode a run was in. +SPECULATIVE_DSPARK="${SPECULATIVE_DSPARK:-1}" +# Experimental request-local circuit breaker. A zero acceptance threshold leaves +# production behavior unchanged. The measured Lenovo A/B starts at 0.60/32/64. +DSPARK_FALLBACK_ACCEPTANCE="${DSPARK_FALLBACK_ACCEPTANCE:-0}" +DSPARK_FALLBACK_MIN_DRAFTED="${DSPARK_FALLBACK_MIN_DRAFTED:-32}" +DSPARK_FALLBACK_STEPS="${DSPARK_FALLBACK_STEPS:-64}" +# Which base dspark_target_layer_ids is read with (0 or 1). Exported so it reaches the +# rank subprocesses; unset means the model's own default. See the note in model.py -- +# this is settled by measurement, not by the config. +[ -n "${FREETOKEN_DSPARK_LAYER_BASE:-}" ] && export FREETOKEN_DSPARK_LAYER_BASE +# FREETOKEN_DSPARK_MARKOV=0 drops the Markov bias from the draft logits. +[ -n "${FREETOKEN_DSPARK_MARKOV:-}" ] && export FREETOKEN_DSPARK_MARKOV +# FREETOKEN_SPEC_HYBRID_MOE=0 sends speculative verifies down the prefill expert path. +[ -n "${FREETOKEN_SPEC_HYBRID_MOE:-}" ] && export FREETOKEN_SPEC_HYBRID_MOE +# FREETOKEN_SPEC_DEBUG=1 dumps the first few blocks' token flow. +[ -n "${FREETOKEN_SPEC_DEBUG:-}" ] && export FREETOKEN_SPEC_DEBUG +[ -n "${FREETOKEN_SPEC_TIMING:-}" ] && export FREETOKEN_SPEC_TIMING +[ -n "${FREETOKEN_SPEC_PROFILE:-}" ] && export FREETOKEN_SPEC_PROFILE +[ -n "${FREETOKEN_SPEC_FORCE_REJECT:-}" ] && export FREETOKEN_SPEC_FORCE_REJECT +[ -n "${FREETOKEN_CMP_FIRST_ONLY:-}" ] && export FREETOKEN_CMP_FIRST_ONLY +[ -n "${FREETOKEN_SPEC_WINDOW_ONLY:-}" ] && export FREETOKEN_SPEC_WINDOW_ONLY +LOG="${LOG:-/tmp/freetoken-dsv4.log}" +# The TP ranks' torch.distributed rendezvous. Held by every rank, not just the +# frontend, so it is the one that lingers after a stop. +RDZV_PORT="${RDZV_PORT:-8082}" +PIDFILE="${PIDFILE:-/tmp/freetoken-dsv4.pid}" +START_TIMEOUT="${START_TIMEOUT:-2400}" # seconds; the expert banks take a while + +# nvcc rejects gcc newer than 15, and this host runs gcc 16. Point the JIT host pass +# at gcc-15, which is installed alongside it. Without this every kernel build fails +# with "unsupported GNU version". +export NVCC_APPEND_FLAGS="${NVCC_APPEND_FLAGS:--ccbin /usr/bin/g++-15}" +export CXX="${CXX:-/usr/bin/g++-15}" +export CC="${CC:-/usr/bin/gcc-15}" + +FT="$FT_DIR/.venv/bin/ft" + +die() { echo "error: $*" >&2; exit 1; } + +running_pid() { + [ -f "$PIDFILE" ] || return 1 + local pid; pid=$(cat "$PIDFILE" 2>/dev/null) || return 1 + [ -n "$pid" ] && kill -0 "$pid" 2>/dev/null || return 1 + echo "$pid" +} + + +# Every process belonging to this service, as PIDs. +# +# The TP ranks are multiprocessing.spawn children: their cmdline is a generic +# python3 -c from multiprocessing.spawn import spawn_main; spawn_main(...) +# with no "ft serve" and no model path in it, and launch.py detaches them into their own +# process group. So neither a pattern match on the serve command nor a process-group kill +# reaches them -- which is why a "stopped" service could still be holding its ports, and +# why the ranks had to be found by hand more than once. +# +# Two handles that do work: the venv interpreter they were spawned from (specific to this +# checkout), and whoever is listening on the service's ports. +service_pids() { + { + pgrep -u "$(id -u)" -f "ft serve --model $MODEL" 2>/dev/null + pgrep -u "$(id -u)" -f "^$FT_DIR/.venv/bin/python3 -c from multiprocessing.spawn" 2>/dev/null + ss -ltnp 2>/dev/null \ + | grep -E ":($PORT|$RDZV_PORT) " \ + | grep -oE 'pid=[0-9]+' | cut -d= -f2 + } | sort -un +} + +cmd_start() { + [ -x "$FT" ] || die "no ft binary at $FT (run: cd $FT_DIR && uv venv && uv pip install -e '.[accel]')" + [ -d "$MODEL" ] || die "no model directory at $MODEL" + if running_pid >/dev/null; then + echo "already running (pid $(running_pid)) on $HOST:$PORT"; return 0 + fi + if ss -ltn 2>/dev/null | grep -q ":$PORT "; then + die "port $PORT is already in use -- every inference service on this host shares it, so stop the other one first" + fi + + # The TP ranks rendezvous over a second port, and it outlives the API port on a + # restart: the previous run's ranks hold it while they tear down. Starting into that + # window fails every rank at once with a message about the wrong thing -- + # DistNetworkError: ... port: 8082 ... EADDRINUSE + # -- after the model has already begun loading. So wait for it, and say what for. + local waited=0 + while ss -ltn 2>/dev/null | grep -q ":$RDZV_PORT "; do + if [ "$waited" -eq 0 ]; then + echo -n "waiting for the TP rendezvous port $RDZV_PORT to be released" + fi + [ "$waited" -ge 60 ] && { echo; die "port $RDZV_PORT still held after 60s -- a previous rank is stuck; try '$0 kill'"; } + sleep 2; waited=$((waited + 2)); printf '.' + done + [ "$waited" -gt 0 ] && echo " ok" + + local spec=() + [ "$SPECULATIVE_DSPARK" = "1" ] && spec=(--speculative-dspark) + local fallback=() + if [ "$DSPARK_FALLBACK_ACCEPTANCE" != "0" ] && [ "$DSPARK_FALLBACK_ACCEPTANCE" != "0.0" ]; then + fallback=( + --dspark-fallback-acceptance "$DSPARK_FALLBACK_ACCEPTANCE" + --dspark-fallback-min-drafted "$DSPARK_FALLBACK_MIN_DRAFTED" + --dspark-fallback-steps "$DSPARK_FALLBACK_STEPS" + ) + fi + + echo "starting DeepSeek-V4-Flash on $HOST:$PORT (TP=$TP_SIZE, memory-ratio $MEMORY_RATIO${spec:+, dSpark drafter})" + : > "$LOG" + cd "$FT_DIR" || die "cannot cd to $FT_DIR" + setsid "$FT" serve \ + --model "$MODEL" \ + --host "$HOST" --port "$PORT" \ + --tensor-parallel-size "$TP_SIZE" \ + --memory-ratio "$MEMORY_RATIO" \ + --max-running-requests "$MAX_RUNNING_REQUESTS" \ + --max-prefill-length "$MAX_PREFILL_LENGTH" \ + --expert-load "$EXPERT_LOAD" \ + "${spec[@]}" \ + "${fallback[@]}" \ + >> "$LOG" 2>&1 < /dev/null & + echo $! > "$PIDFILE" + + # The HTTP frontend binds BEFORE the model loads, so /v1/models answers 200 while + # generation would still fail. The log line below is the only honest readiness signal. + echo "loading the expert banks -- this is the slow part; follow it with: $0 logs" + local waited=0 + while [ "$waited" -lt "$START_TIMEOUT" ]; do + if grep -aq "ready to serve" "$LOG"; then + echo + echo "ready on $HOST:$PORT after ${waited}s" + grep -a "moe_cache_size\|Allocating .* tokens\|Weights:" "$LOG" \ + | sed "s/.*INFO *//" | sed "s/\x1b\[[0-9;]*m//g" | cut -c1-120 | tail -3 + return 0 + fi + if grep -aqE "Traceback|OutOfMemoryError|cannot be restarted" "$LOG"; then + echo; echo "startup FAILED -- last lines:" >&2 + grep -aE "Error|OutOfMemory|assert" "$LOG" | tail -3 >&2 + echo "hint: an OOM during CUDA-graph capture means MEMORY_RATIO is too high." >&2 + cmd_stop >/dev/null 2>&1 + return 1 + fi + sleep 10; waited=$((waited + 10)) + printf '.' + done + echo; die "still not ready after ${START_TIMEOUT}s -- see $LOG" +} + +cmd_stop() { + # SIGTERM, never SIGKILL. The host expert banks are cudaHostRegister'd shared + # mappings; a SIGKILL'd rank leaves them behind (145 GB of ownerless Shmem, seen + # in practice) and only a reboot gets it back. + local pids; pids=$(service_pids) + [ -n "$pids" ] && kill -TERM $pids 2>/dev/null + echo -n "stopping" + for _ in $(seq 1 30); do + [ -z "$(service_pids)" ] && break + sleep 2; printf '.' + done + echo + rm -f "$PIDFILE" + if [ -n "$(service_pids)" ]; then + if [ "${1:-}" = "--force" ]; then + echo "still running after 60s -- forcing" >&2 + cmd_kill + return $? + fi + echo "warning: something is still running. Give it longer before forcing it --" >&2 + echo "a SIGKILL here leaks the pinned host banks. Use '$0 kill' to force." >&2 + return 1 + fi + # "stopped" must mean "ready to start again". The ranks release the rendezvous port + # slightly after the processes go, and a start inside that window fails every rank. + for _ in $(seq 1 15); do + ss -ltn 2>/dev/null | grep -q ":$RDZV_PORT " || break + sleep 1 + done + echo "stopped" +} + +# SIGKILL, plus the cleanup that makes it survivable. +# +# A killed rank never runs its atexit, so the host expert banks stay cudaHostRegister'd: +# ownerless Shmem that no process owns and no reboot-free path reclaims -- 145 GB of it, +# seen here. The banks are backed by files under /dev/shm, so once every rank is gone +# those files ARE the leak, and deleting them is what a clean shutdown would have done. +# +# Use this when a rank is wedged (a stuck decode, a hung collective) and SIGTERM has not +# landed. Prefer 'stop': a graceful exit needs no cleanup at all. +cmd_kill() { + echo "SIGKILL -- the pinned host banks will be reclaimed by hand below" >&2 + local pids; pids=$(service_pids) + [ -n "$pids" ] && kill -KILL $pids 2>/dev/null + for _ in $(seq 1 15); do + [ -z "$(service_pids)" ] && break + sleep 1 + done + rm -f "$PIDFILE" + if [ -n "$(service_pids)" ]; then + echo "error: processes survived SIGKILL -- they are in uninterruptible sleep (D)," >&2 + echo "usually a stuck GPU or NFS call. Check 'ps -o stat,pid,cmd -p \$(pgrep -f ft)'." >&2 + return 1 + fi + + # Reclaim only what this model's ranks left behind, and only with no rank alive to + # be using it -- deleting a live rank's segment corrupts a running model. + local before after + before=$(awk '/^Shmem:/{print $2}' /proc/meminfo) + find /dev/shm -maxdepth 1 \( -name 'freetoken*' -o -name 'ft_*' -o -name 'torch_*' \) \ + -user "$(id -un)" -delete 2>/dev/null + ipcs -m | awk -v me="$(id -u)" '$3==me && $6==0 {print $2}' \ + | xargs -r -n1 ipcrm -m 2>/dev/null + # The kernel returns the pages asynchronously, so a reading taken straight after the + # kill still shows the whole model resident and the check below cries leak on a + # perfectly clean shutdown. Wait for it to stop moving. + # `$_` is a bash special variable (the previous command's last argument), not the + # loop counter, so reading it here compared an empty string as an integer. + after=$before + local tick=0 now + while [ "$tick" -lt 10 ]; do + sleep 1 + tick=$((tick + 1)) + now=$(awk '/^Shmem:/{print $2}' /proc/meminfo) + # Stop once the figure stops falling, after giving it a couple of ticks to start. + if [ "$tick" -gt 2 ] && [ "$now" -ge "$after" ]; then + after=$now + break + fi + after=$now + done + echo "killed; Shmem $((before/1024/1024)) GiB -> $((after/1024/1024)) GiB" + + # Only a shortfall RELATIVE to what was released is evidence of a leak. An absolute + # threshold flags every host that legitimately uses shared memory for something else. + if [ "$after" -gt $((before / 2)) ] && [ "$before" -gt $((32*1024*1024)) ]; then + echo "note: less than half the shared memory came back. If nothing else on this" >&2 + echo "host uses it, some pinned banks were not reclaimed." >&2 + fi +} + +cmd_status() { + if grep -aq "ready to serve" "$LOG" 2>/dev/null && curl -sf -m 5 -o /dev/null "http://127.0.0.1:$PORT/v1/models"; then + echo "UP on $HOST:$PORT" + elif running_pid >/dev/null; then + echo "LOADING (pid $(running_pid)) -- not ready yet" + tail -c 300 "$LOG" 2>/dev/null | tr '\r' '\n' | grep -a . | tail -1 + else + echo "DOWN" + fi + echo + nvidia-smi --query-gpu=index,memory.used,memory.total,utilization.gpu \ + --format=csv,noheader 2>/dev/null | sed 's/^/ gpu /' + free -g | sed -n '2p' | awk '{printf " host ram: %s GiB used, %s GiB available of %s GiB\n", $3, $7, $2}' +} + +cmd_logs() { tail -f "$LOG"; } + +cmd_test() { + curl -sf -m 300 "http://127.0.0.1:$PORT/v1/chat/completions" \ + -H 'Content-Type: application/json' \ + -d '{"model":"DeepSeek-V4-Flash-0731", + "messages":[{"role":"user","content":"Say hello in one short sentence."}], + "max_tokens":32}' \ + && echo +} + +case "${1:-}" in + start) cmd_start ;; + stop) cmd_stop "${2:-}" ;; + kill) cmd_kill ;; + restart) cmd_stop; cmd_start ;; + status) cmd_status ;; + logs) cmd_logs ;; + test) cmd_test ;; + *) echo "usage: $0 {start|stop [--force]|kill|restart|status|logs|test}"; exit 2 ;; +esac diff --git a/tests/dsv4/test_dsv4_adaptive_verification.py b/tests/dsv4/test_dsv4_adaptive_verification.py new file mode 100644 index 0000000..55e61ef --- /dev/null +++ b/tests/dsv4/test_dsv4_adaptive_verification.py @@ -0,0 +1,197 @@ +"""Paper-faithful adaptive verification tests that do not require CUDA.""" + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.engine.engine import Engine +from freetoken.engine.graph import GraphRunner, SharedSpecCarryJournal +from freetoken.models.deepseek_v4.dspark import ( + DSparkAcceptanceFallback, + choose_adaptive_draft_width, +) +from freetoken.scheduler.scheduler import Scheduler + + +def test_acceptance_fallback_trips_only_after_minimum_sample(): + fallback = DSparkAcceptanceFallback(0.60, min_drafted=8, cooldown_steps=3) + + assert fallback.record(2, 5) is None + trip = fallback.record(1, 3) + + assert trip == pytest.approx((3 / 8, 8)) + + +def test_acceptance_fallback_skips_exact_cooldown_then_probes(): + fallback = DSparkAcceptanceFallback(0.60, min_drafted=4, cooldown_steps=3) + assert fallback.record(1, 4) is not None + + assert [fallback.take_decision() for _ in range(3)] == [ + (False, False), + (False, False), + (False, False), + ] + assert fallback.take_decision() == (True, True) + assert fallback.take_decision() == (True, False) + + +def test_acceptance_fallback_keeps_high_acceptance_speculating(): + fallback = DSparkAcceptanceFallback(0.60, min_drafted=8, cooldown_steps=3) + + assert fallback.record(5, 8) is None + assert fallback.take_decision() == (True, False) + + +def test_acceptance_fallback_reset_clears_request_state(): + fallback = DSparkAcceptanceFallback(0.60, min_drafted=4, cooldown_steps=3) + fallback.record(1, 4) + + fallback.reset() + + assert fallback.take_decision() == (True, False) + assert fallback.record(1, 3) is None + + +def test_scheduler_fallback_bypasses_block_before_appending_noise(): + scheduler = Scheduler.__new__(Scheduler) + scheduler._speculative = SimpleNamespace(block_size=5, noise_token_id=17) + scheduler.engine = SimpleNamespace(should_speculate=lambda _req: False) + scheduler._warned_non_greedy = False + req = SimpleNamespace(remain_len=100, input_ids=torch.tensor([1, 2])) + batch = SimpleNamespace(reqs=[req], is_decode=True, phase="decode") + + scheduler._maybe_make_speculative(batch) + + assert batch.phase == "decode" + assert req.input_ids.tolist() == [1, 2] + + +def test_profile_cliff_stops_before_expensive_graph_bucket(): + width = choose_adaptive_draft_width( + [0.9, 0.9], + draft_cost_ms=0.0, + verify_cost_ms=[1.0, 1.0, 100.0], + ) + assert width == 1 + + +def test_flat_verify_cost_admits_the_whole_prefix(): + width = choose_adaptive_draft_width( + [0.8, 0.7, 0.6], + draft_cost_ms=2.0, + verify_cost_ms=[10.0, 10.0, 10.0, 10.0], + ) + assert width == 3 + + +def test_selector_uses_cumulative_survival_not_raw_confidence(): + # The second token's contribution is 0.5 * 0.9, not 0.9 by itself. + width = choose_adaptive_draft_width( + [0.5, 0.9], + draft_cost_ms=0.0, + verify_cost_ms=[1.0, 1.0, 1.4], + ) + assert width == 1 + + +def test_tie_keeps_the_smaller_prefix_like_argmax(): + assert choose_adaptive_draft_width( + [1.0], draft_cost_ms=0.0, verify_cost_ms=[1.0, 2.0] + ) == 0 + + +def test_cost_curve_must_cover_every_prefix(): + with pytest.raises(ValueError, match="every width"): + choose_adaptive_draft_width( + [0.9, 0.9], draft_cost_ms=1.0, verify_cost_ms=[1.0, 2.0] + ) + + +def test_prefix_graphs_share_the_maximum_carry_journal_storage(): + key = (3, "compress", 128) + storage = {key: [torch.zeros(1) for _ in range(6)]} + journal = SharedSpecCarryJournal(storage) + journal.reset(2) + journal.record(key, torch.tensor([4.0])) + journal.record(key, torch.tensor([5.0])) + + [(seen_key, pieces)] = list(journal.items()) + assert seen_key == key + assert len(pieces) == 2 + assert torch.cat(pieces).tolist() == [4.0, 5.0] + assert storage[key][2].item() == 0.0 + + +def test_shared_carry_journal_rejects_prefix_overflow(): + key = (3, "compress", 128) + journal = SharedSpecCarryJournal({key: [torch.zeros(1)]}) + journal.reset(1) + journal.record(key, torch.ones(1)) + with pytest.raises(RuntimeError, match="overflow"): + journal.record(key, torch.ones(1)) + + +class _ChooseTwo: + block_size = 5 + + def record_and_choose(self, confidence, uid): + assert confidence.tolist() == pytest.approx([0.9] * 5) + assert uid == 7 + return 2 + + +def test_engine_compacts_only_target_views_and_keeps_full_allocation_frontier(): + engine = Engine.__new__(Engine) + engine._adaptive_verification = _ChooseTwo() + ids = torch.arange(20, dtype=torch.int32) + req = SimpleNamespace(input_ids=ids[:15], _ids_buf=ids, uid=7, device_len=15) + metadata = SimpleNamespace(segments=[(0, 6, 3, 9)]) + batch = SimpleNamespace( + reqs=[req], + padded_size=1, + spec_block=5, + draft_confidence=torch.full((5,), 0.9), + draft_tokens=torch.arange(5, dtype=torch.int32), + draft_probs=torch.empty(5, 4), + input_ids=torch.arange(6, dtype=torch.int32), + positions=torch.arange(9, 15, dtype=torch.int32), + out_loc=torch.arange(6, dtype=torch.int32), + attn_metadata=metadata, + ) + + engine.adapt_speculative_batch(batch) + + assert batch.spec_block == 2 + assert batch.input_ids.tolist() == [0, 1, 2] + assert batch.positions.tolist() == [9, 10, 11] + assert batch.out_loc.tolist() == [0, 1, 2] + assert batch.draft_tokens.tolist() == [0, 1] + assert batch.draft_confidence.tolist() == pytest.approx([0.9, 0.9]) + assert metadata.segments == [(0, 3, 3, 9)] + assert req.input_ids.numel() == 12 + assert req.device_len == 15, "tail release still needs the gamma-wide frontier" + + +def test_greedy_verify_reduces_logits_before_crossing_pcie(): + import inspect + + body = inspect.getsource(Engine._finish_speculative) + assert 'logits.argmax(dim=-1).to(torch.int32).to("cpu"' in body + assert 'logits.to("cpu"' not in body + + +def test_probabilistic_verify_does_not_compute_unused_target_argmax(): + import inspect + + body = inspect.getsource(Engine._finish_speculative) + assert "any_greedy = any(req.sampling_params.is_greedy" in body + assert "if any_greedy" in body + assert "else None" in body + + +def test_capacity_profile_flushes_cache_once_not_before_every_replay(): + import inspect + + body = inspect.getsource(GraphRunner._profile_graph_ms) + assert body.count("self._reset_moe_offload_cache()") == 1 diff --git a/tests/dsv4/test_dsv4_aux_addressing.py b/tests/dsv4/test_dsv4_aux_addressing.py new file mode 100644 index 0000000..5073829 --- /dev/null +++ b/tests/dsv4/test_dsv4_aux_addressing.py @@ -0,0 +1,128 @@ +"""The aux tap and the positions it was computed at must travel together. + +The drafter derives context KV from the target's hidden state and writes it at those +tokens' window slots. Which tokens is not a free choice: the tap belongs to a specific +forward, and the KV must land where THAT forward's tokens live. + +The catch-up runs before the target's forward for the current step, so the newest tap is +the PREVIOUS forward's. Addressing it with the current batch's positions -- the block +about to be drafted -- stores hidden states derived from one set of tokens into the slots +of another. Nothing raises. Attention still runs, the shapes all match, and the only +symptom is a drafter whose proposals are rejected, which reads as a weak drafter rather +than a wiring fault. That is what these tests exist to prevent. +""" + +from __future__ import annotations + +import inspect +import pathlib + +import pytest + +MODEL = ( + pathlib.Path(__file__).resolve().parents[2] + / "python" / "freetoken" / "models" / "deepseek_v4" / "model.py" +) +SRC = MODEL.read_text() + + +def _fn(name: str) -> str: + start = SRC.index(f" def {name}(") + nxt = SRC.index("\n def ", start + 10) + return SRC[start:nxt] + + +class TestTheTapCarriesItsAddressing: + def test_prefill_and_decode_build_one_feature_bundle(self): + assert SRC.count("DSparkTargetFeatures(") >= 2 + + def test_addressing_is_exposed_as_one_unit(self): + from freetoken.models.deepseek_v4.model import DSparkTargetFeatures + + assert set(DSparkTargetFeatures.__dataclass_fields__) == { + "hidden", "positions", "table_rows" + } + + def test_addressing_returns_none_until_a_forward_has_run(self): + from freetoken.models.deepseek_v4.model import Transformer + + t = Transformer.__new__(Transformer) + t._target_features = None + assert t.target_features() is None + + def test_tap_reshape_reads_the_dimension_from_model_args(self): + # Transformer has no direct ``dim`` field. This typo survives source-only + # wiring tests and crashes only when the first real aux tap runs. + assert "self.dim * len(aux)" not in SRC + assert SRC.count("self.args.dim * len(aux)") == 3 + + def test_cuda_graph_features_are_owned_per_captured_batch_size(self): + graph_src = ( + pathlib.Path(__file__).resolve().parents[2] + / "python" / "freetoken" / "engine" / "graph.py" + ).read_text() + assert "self._dspark_feature_map[bs] = features" in graph_src + assert "self._dspark_feature_map.get(batch.padded_size)" in graph_src + assert "self._spec_feature_map[key] = features" in graph_src + assert "self._spec_carry_map[key]" in graph_src + assert "shared_carry if shared_carry is not None" in graph_src + + +class TestTheCatchUpUsesTheTapsOwnPositions: + def test_it_does_not_read_the_batch_metadata(self): + body = _fn("catch_up_draft_context") + for forbidden in ("batch.attn_metadata", "md.segments", "batch.positions"): + assert forbidden not in body, ( + f"catch_up_draft_context reads {forbidden}: the tap is the PREVIOUS " + "forward's, so the current batch's positions address the wrong tokens" + ) + + def test_it_takes_one_explicit_feature_bundle(self): + body = _fn("catch_up_draft_context") + for field in ("features.hidden", "features.positions", "features.table_rows"): + assert field in body + + def test_a_row_count_mismatch_is_an_error_not_a_silent_skip(self): + # Truncating to the shorter of the two would write SOME rows at wrong positions, + # which is the failure mode this whole file is about. + body = _fn("catch_up_draft_context") + assert "raise RuntimeError" in body + + +class TestScatteredSlotLookup: + """A flat batch spans requests, so slots come from a gather, not a slice.""" + + def test_the_backend_exposes_the_gather_form(self): + from freetoken.attention.dsv4_sparse import DSV4SparseAttnBackend as M + + assert hasattr(M, "window_slots_at") + params = list(inspect.signature(M.window_slots_at).parameters) + assert params[1:] == ["rows", "positions"] + + def test_it_gathers_per_token_rather_than_slicing_one_request(self): + torch = pytest.importorskip("torch") + + from freetoken.attention.dsv4_sparse import DSV4SparseAttnBackend as M + + class _Pool: + def __init__(self): + # full_loc_map[row, position] -> a distinct full slot per pair. + self.full_loc_map = torch.arange(4 * 8).reshape(4, 8) + + def translate_full_to_window(self, x): + return x # identity: the gather itself is under test + + class _B(M): + # `pool` is a read-only property on the real backend, so shadow it rather + # than constructing one (which would need a live KV pool and a CUDA device). + pool = property(lambda self: self._pool) + + def __init__(self): + self._pool = _Pool() + + rows = torch.tensor([0, 2, 2, 3]) + pos = torch.tensor([1, 0, 5, 7]) + got = _B().window_slots_at(rows, pos) + assert got.tolist() == [1, 16, 21, 31], ( + "each token's slot must come from its OWN (row, position) pair" + ) diff --git a/tests/dsv4/test_dsv4_aux_layer_ids.py b/tests/dsv4/test_dsv4_aux_layer_ids.py new file mode 100644 index 0000000..cd9b469 --- /dev/null +++ b/tests/dsv4/test_dsv4_aux_layer_ids.py @@ -0,0 +1,72 @@ +"""Which target layers the drafter taps. + +dspark_target_layer_ids is 0-based. This checkpoint declares [40, 41, 42] against +n_layers=43 -- exactly the last three layers, and already valid indices into 0..42. +Reading them as 1-based and subtracting one taps 39, 40, 41: the drafter reads the wrong +three layers on every single draft, and nothing anywhere raises. + +A range check cannot separate the two readings, because both land inside the model. The +guard that does is structural: a draft tap reads the END of the target stack, so the ids +must reach the last layer. Under the wrong base they stop one short, every time. +""" + +from __future__ import annotations + +import pytest + + +def _args(ids, n_layers=43): + from freetoken.models.deepseek_v4.args import DeepseekV4Args + + a = DeepseekV4Args.__new__(DeepseekV4Args) + object.__setattr__(a, "dspark_target_layer_ids", tuple(ids)) + object.__setattr__(a, "n_layers", n_layers) + return a + + +def _resolve(args): + """The id resolution as the Transformer performs it, without building a model.""" + ids = tuple(args.dspark_target_layer_ids) + bad = [i for i in ids if not 0 <= i < args.n_layers] + if bad: + raise ValueError(f"{ids} are 0-based and must land inside {args.n_layers}; {bad} do not") + return frozenset(ids) + + +class TestIdsMustLandInsideTheModel: + def test_in_range_ids_resolve(self): + assert _resolve(_args([40, 41, 42], 43)) == {40, 41, 42} + + +class TestTheWrongBaseIsRejected: + def test_out_of_range_ids_raise(self): + with pytest.raises(ValueError, match="0-based"): + _resolve(_args([41, 42, 43], 43)) + + def test_no_ids_is_not_an_error(self): + # dSpark off, or a checkpoint that declares none. + assert _resolve(_args([], 43)) == frozenset() + + +class TestTheReferenceMappingIsFixed: + def test_the_resolved_layers_are_logged(self): + assert "dSpark: tapping target layer outputs" in _model_src() + + def test_no_runtime_base_override_remains(self): + assert "FREETOKEN_DSPARK_LAYER_BASE" not in _model_src() + + def test_vllm_mapping_selects_the_last_three_outputs(self): + # vLLM adds one to the config ids, then captures when idx+1 is selected. + config = (40, 41, 42) + aux_ids = tuple(i + 1 for i in config) + captured = tuple(i for i in range(43) if i + 1 in aux_ids) + assert captured == config + + +def _model_src() -> str: + import pathlib as _p + + return ( + _p.Path(__file__).resolve().parents[2] + / "python" / "freetoken" / "models" / "deepseek_v4" / "model.py" + ).read_text() diff --git a/tests/dsv4/test_dsv4_dspark.py b/tests/dsv4/test_dsv4_dspark.py new file mode 100644 index 0000000..0767bec --- /dev/null +++ b/tests/dsv4/test_dsv4_dspark.py @@ -0,0 +1,43 @@ +"""DSpark greedy prefix acceptance. + +These two functions decide what the model EMITS. A wrong acceptance rule does not +crash or slow anything down -- it silently changes the output distribution, which is +the one failure mode speculative decoding must never have. So the rule is pinned here +rather than left to the loop that calls it. +""" + +from __future__ import annotations + +import torch + +from freetoken.models.deepseek_v4.dspark import accepted_prefix + + +def _t(*xs: int) -> torch.Tensor: + return torch.tensor(xs, dtype=torch.long) + + +class TestAcceptedPrefix: + def test_a_fully_matching_block_is_wholly_accepted(self): + n, bonus = accepted_prefix(_t(5, 6, 7), _t(5, 6, 7, 8)) + assert n == 3 + assert bonus == 8, "verification's extra position supplies the next token free" + + def test_acceptance_stops_at_the_first_disagreement(self): + n, bonus = accepted_prefix(_t(5, 6, 7), _t(5, 99, 7)) + assert n == 1 + assert bonus == 99, "the target's own token replaces the rejected one" + + def test_a_later_match_after_a_rejection_does_not_count(self): + # Position 2 agrees, but it was drafted from a token the target rejected, so + # keeping it would emit a sequence the target never would have produced. + n, _ = accepted_prefix(_t(1, 2, 3), _t(1, 99, 3)) + assert n == 1 + + def test_a_wholly_rejected_block_still_advances_by_one(self): + n, bonus = accepted_prefix(_t(4, 5), _t(9, 9)) + assert (n, bonus) == (0, 9), "a rejected block must not stall the request" + + def test_no_bonus_when_verification_covered_no_extra_position(self): + n, bonus = accepted_prefix(_t(1, 2), _t(1, 2)) + assert (n, bonus) == (2, -1) diff --git a/tests/dsv4/test_dsv4_dspark_heads.py b/tests/dsv4/test_dsv4_dspark_heads.py new file mode 100644 index 0000000..16a5dc5 --- /dev/null +++ b/tests/dsv4/test_dsv4_dspark_heads.py @@ -0,0 +1,235 @@ +"""The dSpark heads, and the two contracts a review pass caught after they were written. + +Both bugs fixed here were silent, and both were invisible to the shape checks that +already existed: + +* The drafter reused the target's head WEIGHT. Under tensor parallelism that head is + vocabulary-parallel, so a bare F.linear returns this rank's vocabulary slice, while + the Markov bias is full-vocabulary because the Markov head is replicated. Adding them + broadcasts into nonsense or raises, depending on the TP size -- and at TP=1, where + most testing happens, it is perfectly correct. +* store_context_kv sliced the RoPE table from the front. Every context after the first + block starts at a non-zero position, so it would rotate the drafter's keys as though + they began at zero. Attention still runs; acceptance just quietly drops. + +CPU-only. +""" + +from __future__ import annotations + +import pytest +import torch + +import freetoken.distributed.info as info_mod +from freetoken.distributed import DistributedInfo +from freetoken.models.deepseek_v4.args import DeepseekV4Args +from freetoken.models.deepseek_v4.dspark import ConfidenceHead, MarkovHead + +VOCAB, DIM, RANK = 64, 32, 8 + + +def _args(**kw) -> DeepseekV4Args: + base = dict( + vocab_size=VOCAB, dim=DIM, dspark_markov_rank=RANK, dspark_block_size=5, + n_mtp_layers=3, dspark_enabled=True, n_layers=4, n_hash_layers=1, + compress_ratios=(0, 4, 128, 4), max_seq_len=4096, max_batch_size=1, + ) + base.update(kw) + return DeepseekV4Args(**base) + + +@pytest.fixture(autouse=True) +def _tp1(): + info_mod._TP_INFO = DistributedInfo(0, 1) + yield + info_mod._TP_INFO = DistributedInfo(0, 1) + + +class TestMarkovHead: + def test_the_bias_covers_the_whole_vocabulary(self): + # It must, because the Markov head is replicated under TP while the target's + # output head is sharded. A per-rank bias could not be added to gathered logits. + head = MarkovHead(_args()) + bias = head.bias(head.embed(torch.tensor([3, 7]))) + assert bias.shape == (2, VOCAB) + + def test_the_embedding_is_the_low_rank_width(self): + head = MarkovHead(_args()) + assert head.embed(torch.tensor([1, 2, 3])).shape == (3, RANK) + + def test_different_previous_tokens_give_different_biases(self): + # The head's entire purpose is to condition on the previous token; a bias that + # ignores it would leave the drafter proposing the same continuation everywhere. + head = MarkovHead(_args()) + torch.nn.init.normal_(head.markov_w1.weight) + torch.nn.init.normal_(head.markov_w2) + a = head.bias(head.embed(torch.tensor([1]))) + b = head.bias(head.embed(torch.tensor([2]))) + assert not torch.allclose(a, b) + + +class TestConfidenceHead: + def test_it_scores_one_value_per_position(self): + head = ConfidenceHead(_args()) + out = head(torch.randn(4, DIM), torch.randn(4, RANK)) + assert out.shape == (4,) + + def test_scores_are_probabilities(self): + # Equation 7 defines a probability consumed by the load-aware scheduler. + head = ConfidenceHead(_args()) + torch.nn.init.normal_(head.proj, std=10.0) # push the pre-sigmoid range wide + out = head(torch.randn(16, DIM), torch.randn(16, RANK)) + assert bool((out >= 0).all() and (out <= 1).all()) + + def test_it_consumes_both_the_hidden_and_the_markov_embedding(self): + head = ConfidenceHead(_args()) + assert head.proj.shape == (1, DIM + RANK), ( + "confidence reads the hidden state AND the previous-token embedding" + ) + + +class TestSequentialMarkovStage: + def test_each_step_reads_the_token_sampled_by_the_previous_step(self): + from freetoken.models.deepseek_v4.dspark import DSparkDrafter + + seen = [] + + class _Markov(torch.nn.Module): + def embed(self, ids): + seen.append(ids.clone()) + return torch.zeros(ids.shape[0], RANK) + + def bias(self, embeds): + return torch.zeros(embeds.shape[0], VOCAB) + + class _Confidence(torch.nn.Module): + def forward(self, hidden, embeds): + return torch.ones(hidden.shape[0]) + + class _Params: + temperature = 0.0 + top_p = 1.0 + top_k = -1 + is_greedy = True + + drafter = DSparkDrafter.__new__(DSparkDrafter) + torch.nn.Module.__init__(drafter) + drafter.markov_head = _Markov() + drafter.confidence_head = _Confidence() + base = torch.full((1, 3, VOCAB), -10.0) + base[0, 0, 2] = 10.0 + base[0, 1, 3] = 10.0 + base[0, 2, 4] = 10.0 + proposed, q, confidence = drafter.sample_block( + base, torch.zeros(1, 3, DIM), torch.tensor([1]), [_Params()] + ) + assert proposed.tolist() == [2, 3, 4] + assert [x.tolist() for x in seen] == [[1], [2], [3]] + assert q.shape == (3, VOCAB) + assert confidence.shape == (3,) + + +class TestDraftHeadNormalization: + def test_backbone_returns_pre_norm_hidden_for_the_confidence_head(self): + import inspect + + from freetoken.models.deepseek_v4.dspark import DSparkDrafter + + body = inspect.getsource(DSparkDrafter.head_hidden) + assert "return self.hc_head(h)" in body + assert "return self.norm(self.hc_head(h))" not in body + + def test_only_base_logits_apply_the_final_norm(self): + import inspect + + from freetoken.models.deepseek_v4.dspark import DSparkDrafter + + body = inspect.getsource(DSparkDrafter.propose) + assert "target_logits_fn(self.norm(head_hidden))" in body + + +class TestContextKvPositions: + """The other bug: RoPE must key on absolute positions, not a front slice.""" + + def test_a_position_length_mismatch_is_refused(self): + from freetoken.models.deepseek_v4.dspark import DSparkDrafter + + args = _args(dim=512, moe_inter_dim=512, n_heads=8, o_groups=4) + with torch.device("meta"): + drafter = DSparkDrafter(args) + with pytest.raises(ValueError, match="positions"): + drafter.store_context_kv( + torch.zeros(5, args.dim), torch.arange(3), torch.arange(5) + ) + + +class TestBlockInputIds: + """What a draft block is actually fed. + + dSpark is a BLOCK predictor: it gets the last committed token, then noise at every + unknown position, and fills them in one pass. Feeding zeros or repeating the last + token puts it off its training distribution -- proposals stop being accepted, which + reads as "the drafter is weak" rather than "we fed it the wrong thing". + """ + + def _drafter(self, **kw): + from freetoken.models.deepseek_v4.dspark import DSparkDrafter + + opts = dict(dim=512, moe_inter_dim=512, n_heads=8, o_groups=4, + dspark_noise_token_id=128799) + opts.update(kw) + args = _args(**opts) + with torch.device("meta"): + return DSparkDrafter(args), args + + def test_the_first_position_is_the_real_last_token(self): + d, _ = self._drafter() + ids = d.block_input_ids(4242, torch.device("cpu")) + assert int(ids[0]) == 4242 + + def test_every_other_position_is_the_noise_token(self): + d, _ = self._drafter() + ids = d.block_input_ids(7, torch.device("cpu")) + assert ids.numel() == d.block_size + assert all(int(t) == 128799 for t in ids[1:]), ( + "unknown positions must carry the checkpoint's noise token" + ) + + def test_a_checkpoint_without_a_noise_token_is_refused(self): + d, _ = self._drafter(dspark_noise_token_id=-1) + with pytest.raises(ValueError, match="noise"): + d.block_input_ids(1, torch.device("cpu")) + + def test_propose_does_not_trust_uninitialized_gpu_placeholder_rows(self): + from freetoken.models.deepseek_v4.dspark import DSparkDrafter + + d = DSparkDrafter.__new__(DSparkDrafter) + torch.nn.Module.__init__(d) + d.block_size = 3 + d.noise_token_id = 99 + d.norm = torch.nn.Identity() + seen = {} + d.embed_block = lambda ids: torch.zeros(1, ids.numel(), DIM) + + def head_hidden(_embeds, ids, _segments, _positions): + seen["ids"] = ids.clone() + return torch.zeros(1, ids.numel(), DIM) + + d.head_hidden = head_hidden + d.sample_block = lambda *_args: ( + torch.zeros(6, dtype=torch.long), + torch.zeros(6, VOCAB), + torch.zeros(6), + ) + # Each target span is anchor + gamma proposal slots. The proposal slots here + # deliberately contain stale values, exactly like an unwritten GPU token-pool row. + target_ids = torch.tensor([[11, 7, 8, 9, 22, 70, 80, 90]]) + segments = [(0, 4, 1, 10), (4, 4, 2, 20)] + d.propose( + target_ids, + segments, + torch.arange(8), + lambda h: torch.zeros(h.shape[0], VOCAB), + [object(), object()], + ) + assert seen["ids"].tolist() == [[11, 99, 99, 22, 99, 99]] diff --git a/tests/dsv4/test_dsv4_dspark_masking.py b/tests/dsv4/test_dsv4_dspark_masking.py new file mode 100644 index 0000000..ab0e424 --- /dev/null +++ b/tests/dsv4/test_dsv4_dspark_masking.py @@ -0,0 +1,131 @@ +"""The dSpark block mask, and the invariants that keep speculation honest. + +Two failure modes drive these tests, and neither one crashes: + +* A block mask that is secretly causal. The drafter still runs, still produces five + tokens, and acceptance still works -- it is just five serial steps wearing a block's + clothing, so the whole feature buys nothing and nothing reports that. +* A mask that leaks. If a query can read a slot outside its own request or beyond the + block, the drafter proposes from context the target never had, acceptance collapses, + and the only symptom is a poor acceptance rate that looks like a bad model. + +CPU-only: the masking rule is pure index arithmetic, extracted so it can be checked +without a GPU or a KV pool. +""" + +from __future__ import annotations + +import inspect + +import pytest +import torch + +from freetoken.models.deepseek_v4.dspark import window_cols_for_block + +WIN = 128 + + +def _visible(cols: torch.Tensor, row: int) -> set[int]: + """Columns row ``row`` may actually read (``-1`` is masked out).""" + return {int(c) for c in cols[row] if int(c) >= 0} + + +class TestCausalBaseline: + """The unchanged path must stay exactly as it was.""" + + def test_a_row_never_sees_past_its_own_position(self): + cols = window_cols_for_block(start_pos=500, n=4, window=WIN, non_causal=False) + w_lo = 500 - WIN + 1 + for row in range(4): + highest = max(_visible(cols, row)) + w_lo + assert highest == 500 + row, "causal rows end at their own position" + + def test_each_row_sees_exactly_one_window(self): + cols = window_cols_for_block(start_pos=500, n=3, window=WIN, non_causal=False) + for row in range(3): + assert len(_visible(cols, row)) == WIN + + +class TestNonCausalBlock: + """The whole point: every drafted query sees the whole block.""" + + def test_every_row_sees_the_blocks_last_position(self): + n = 5 + start = 500 + cols = window_cols_for_block(start_pos=start, n=n, window=WIN, non_causal=True) + w_lo = start - WIN + last_col = (start + n - 1) - w_lo + for row in range(n): + assert last_col in _visible(cols, row), ( + f"row {row} cannot see the block's final position -- the mask is " + "still causal, so the block is five serial steps in disguise" + ) + + def test_all_rows_see_an_identical_candidate_set(self): + cols = window_cols_for_block(start_pos=800, n=5, window=WIN, non_causal=True) + first = _visible(cols, 0) + for row in range(1, 5): + assert _visible(cols, row) == first + + def test_it_differs_from_the_causal_mask(self): + # Guards against a refactor that quietly drops the non_causal branch. + a = window_cols_for_block(start_pos=500, n=5, window=WIN, non_causal=True) + b = window_cols_for_block(start_pos=500, n=5, window=WIN, non_causal=False) + assert not torch.equal(a, b) + + def test_it_never_reaches_beyond_the_block(self): + n, start = 5, 500 + cols = window_cols_for_block(start_pos=start, n=n, window=WIN, non_causal=True) + w_lo = start - WIN + highest = max(max(_visible(cols, r)) for r in range(n)) + w_lo + assert highest == start + n - 1, "a query must not read past its own block" + + def test_the_shared_context_stays_visible(self): + n, start = 5, 500 + cols = window_cols_for_block(start_pos=start, n=n, window=WIN, non_causal=True) + w_lo = start - WIN + lowest = min(min(_visible(cols, r)) for r in range(n)) + w_lo + assert lowest < start, "the block must still attend over the context before it" + + @pytest.mark.parametrize("n", [1, 2, 5, 8]) + def test_the_candidate_width_is_context_window_plus_block(self, n): + cols = window_cols_for_block(start_pos=900, n=n, window=WIN, non_causal=True) + assert cols.shape == (n, WIN + n) + + def test_a_single_token_still_keeps_all_context_tokens(self): + a = window_cols_for_block(start_pos=700, n=1, window=WIN, non_causal=True) + b = window_cols_for_block(start_pos=700, n=1, window=WIN, non_causal=False) + assert a.shape[-1] == WIN + 1 + assert b.shape[-1] == WIN + + def test_ragged_launcher_labels_the_full_non_causal_width(self): + # The index helper can correctly build 133 columns while the attention call + # still labels only 128 of them as window columns. That only fails at bs > 1. + from freetoken.models.deepseek_v4.attention import Attention + + body = inspect.getsource(Attention.forward_ragged) + assert "max(part.shape[-1] for part in win_parts)" in body + assert "if self.non_causal" in body + + +class TestMaskBounds: + """A column index that escapes its range would read another request's KV.""" + + @pytest.mark.parametrize("non_causal", [True, False]) + @pytest.mark.parametrize("start_pos,n", [(1, 5), (127, 5), (128, 5), (5000, 5)]) + def test_columns_stay_inside_the_gathered_slot_range(self, non_causal, start_pos, n): + cols = window_cols_for_block(start_pos, n, WIN, non_causal) + w_lo = max( + 0, + start_pos - WIN if non_causal else start_pos - WIN + 1, + ) + width = (start_pos + n) - w_lo # slots the caller gathered for this segment + assert int(cols.max()) < width, "a column past the gathered slots reads foreign KV" + assert int(cols.min()) >= -1, "only -1 encodes 'masked'" + + @pytest.mark.parametrize("non_causal", [True, False]) + def test_no_row_is_entirely_masked(self, non_causal): + # A fully masked row would make its query attend to nothing at all. + cols = window_cols_for_block(300, 5, WIN, non_causal) + for row in range(5): + assert _visible(cols, row), f"row {row} has no visible column" diff --git a/tests/dsv4/test_dsv4_dspark_wiring.py b/tests/dsv4/test_dsv4_dspark_wiring.py new file mode 100644 index 0000000..e48526d --- /dev/null +++ b/tests/dsv4/test_dsv4_dspark_wiring.py @@ -0,0 +1,174 @@ +"""Every dSpark piece must be REACHED, not merely defined. + +The costliest failures in this feature were not wrong code -- they were correct code +nobody called: + +* ``logit_indices`` was added so a verify could score every drafted position, then never + passed. The verify returned one row per request and acceptance died on an IndexError + four frames away from the cause. +* Per-token compressor states must be selected after acceptance; restoring a single + pre-block snapshot loses any accepted prefix in the same page. +* ``catch_up_context`` was written to keep the draft layers' window fresh, and never + called, which reads as a weak drafter rather than a missing call. + +None of those show up in a unit test of the piece itself, because the piece works. They +show up here: an import-graph check that each public entry point has a caller in the +package, and behavioural checks that the wiring carries data through. +""" + +from __future__ import annotations + +import ast +import inspect +import pathlib + +import pytest + +PKG = pathlib.Path(__file__).resolve().parents[2] / "python" / "freetoken" + + +def _calls_in_package() -> set[str]: + """Every attribute and plain name that appears in a call position, package-wide.""" + called: set[str] = set() + for path in PKG.rglob("*.py"): + try: + tree = ast.parse(path.read_text()) + except SyntaxError: # pragma: no cover - a broken file is its own failure + continue + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + fn = node.func + if isinstance(fn, ast.Name): + called.add(fn.id) + elif isinstance(fn, ast.Attribute): + called.add(fn.attr) + # getattr(obj, "name") is how the engine reaches optional model hooks. + for node in ast.walk(tree): + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "getattr" + and node.args + and len(node.args) > 1 + and isinstance(node.args[1], ast.Constant) + and isinstance(node.args[1].value, str) + ): + called.add(node.args[1].value) + return called + + +@pytest.mark.parametrize( + "symbol,why", + [ + ("catch_up_context", "draft layers would attend over a stale window"), + ("store_context_kv", "draft layers would have no context KV at all"), + ("propose", "nothing would produce a draft"), + ("rejection_accept_device", "sampled requests would fall back to argmax bias"), + ("accepted_prefix", "greedy requests would have no acceptance rule"), + ("sampling_probs", "q and p would not be shaped like the sampler's draws"), + ("window_cols_for_block", "the block mask would silently stay causal"), + ("_restore_speculative_carry", "acceptance would leave the rejected carry live"), + ("_trim_dspark_target_features", "rejected target rows would become draft context"), + ("draft_into_batch", "the block would keep its noise placeholders"), + ("_finish_speculative", "nothing would apply acceptance"), + ("_maybe_make_speculative", "no batch would ever become speculative"), + ("catch_up_draft_context", "the engine's hook into the drafter's context"), + ], +) +def test_the_piece_has_a_caller(symbol, why): + assert symbol in _calls_in_package(), ( + f"{symbol}() is defined but never called anywhere in freetoken -- {why}" + ) + + +class TestSpeculativeStateIsCarried: + """The batch fields the speculative path hands between its stages must exist. + + Each of these is written in one place and read in another; a rename on one side is + invisible until a block runs. + """ + + @pytest.mark.parametrize( + "field", + ["speculative", "spec_block", "spec_emitted", "draft_probs", + "draft_confidence", "draft_tokens", "spec_carry_states", + "spec_verify_decode"], + ) + def test_batch_declares_the_field(self, field): + import torch + + from freetoken.core import Batch + + assert field in Batch.__dataclass_fields__, ( + f"Batch.{field} is read by the speculative path but not declared" + ) + assert not Batch.__dataclass_fields__[field].init, ( + f"Batch.{field} must not be a constructor argument" + ) + del torch + + def test_a_fresh_batch_is_not_speculative(self): + from freetoken.core import Batch + + b = Batch(reqs=[], phase="decode") + assert b.speculative is False + assert b.spec_block == 0 + assert b.spec_emitted is None + assert b.draft_tokens is None + assert b.spec_carry_states is None + assert b.spec_verify_decode is False + + +class TestFixedVerifyGraphWiring: + def test_target_graph_is_selected_before_ordinary_decode_graph(self): + engine = (PKG / "engine" / "engine.py").read_text() + spec_at = engine.index("can_use_spec_cuda_graph") + decode_at = engine.index("can_use_cuda_graph", spec_at) + assert spec_at < decode_at + + def test_graph_replays_graph_owned_partial_states(self): + graph = (PKG / "engine" / "graph.py").read_text() + assert "journal = self._spec_carry_map[key]" in graph + assert "batch.spec_carry_states = journal" in graph + assert "prepare_for_spec_replay" in graph + + def test_verify_uses_device_positions_not_a_host_start_position(self): + from freetoken.models.deepseek_v4.model import Transformer + + body = inspect.getsource(Transformer.verify_block) + assert "pos: torch.Tensor" in body + assert "start_pos" not in body + + +class TestVerifyReturnsEveryPosition: + """The bug that cost a restart: a verify must score all 1+k positions.""" + + def test_prefill_batched_accepts_logit_indices(self): + import inspect + + from freetoken.models.deepseek_v4.model import Transformer + + sig = inspect.signature(Transformer.prefill_batched) + assert "logit_indices" in sig.parameters, ( + "the verify pass needs to select every drafted position" + ) + + def test_the_forward_passes_them_for_a_speculative_batch(self): + src = (PKG / "models" / "deepseek_v4" / "model.py").read_text() + assert "logit_indices=logit_indices" in src, ( + "prefill_batched takes logit_indices but the forward never passes it, so a " + "speculative verify scores only each request's last token" + ) + + +class TestMHCBoundsLargePrefillResiduals: + def test_target_and_drafter_use_the_fused_inv_rms(self): + model = (PKG / "models" / "deepseek_v4" / "model.py").read_text() + dspark = (PKG / "models" / "deepseek_v4" / "dspark.py").read_text() + assert "inv_rms(x2d, self.norm_eps)" in model + assert "inv_rms(x2d, self.norm_eps)" in dspark + norm = (PKG / "kernel" / "triton" / "dsv4" / "norm.py").read_text() + assert "def _inv_rms_kernel" in norm + assert "x.float().square()" not in model + assert "x.float().square()" not in dspark diff --git a/tests/dsv4/test_dsv4_hybrid_balance.py b/tests/dsv4/test_dsv4_hybrid_balance.py new file mode 100644 index 0000000..4b6d130 --- /dev/null +++ b/tests/dsv4/test_dsv4_hybrid_balance.py @@ -0,0 +1,65 @@ +"""The hybrid MoE split follows Equation 4 of the FreeToken paper.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + + +class _Cache: + """Bind the real split method without constructing a GPU cache.""" + + def __init__(self, fraction=0.25): + from freetoken.moe.offload_cache import OffloadMoeCache + + self.hybrid_fetch_fraction = fraction + self._fn = OffloadMoeCache._fetch_fraction_for + + def frac(self, layer_id, n_tokens, top_k=6): + ids = torch.zeros(n_tokens, top_k, dtype=torch.int32) + return self._fn(self, layer_id, ids) + + +class TestTheSplitIsTheProfiledBandwidthRatio: + @pytest.mark.parametrize("n_tokens", [1, 2, 5, 6, 64]) + def test_token_width_does_not_rescale_equation_four(self, n_tokens): + # The cache kernel deduplicates all routes first, so its miss count m already + # reflects the wider batch. Equation 4 applies B_P/B_H to that m exactly once. + assert _Cache(0.28).frac(0, n_tokens) == pytest.approx(0.28) + + def test_a_disabled_fraction_stays_disabled(self): + # 0.0 means "use the fixed hybrid_max_fetch cap instead". Scaling it would turn + # the cap path on for speculative steps only. + assert _Cache(0.0).frac(0, 6) == 0.0 + + def test_tensor_shape_does_not_change_the_ratio(self): + c = _Cache(0.3) + assert c._fn(c, 0, torch.zeros(6, dtype=torch.int32)) == pytest.approx(0.3) + + +class TestCpuExecutorScratchWidth: + def test_anchor_plus_five_drafts_needs_six_rows(self): + from freetoken.engine.engine import _cpu_moe_max_tokens + + config = SimpleNamespace( + model_config=SimpleNamespace( + dsv4_args=SimpleNamespace(dspark_enabled=True, dspark_block_size=5) + ), + max_running_req=1, + cuda_graph_max_bs=None, + ) + assert _cpu_moe_max_tokens(config) == 6 + + def test_bound_scales_with_requests_and_graph_capture(self): + from freetoken.engine.engine import _cpu_moe_max_tokens + + config = SimpleNamespace( + model_config=SimpleNamespace( + dsv4_args=SimpleNamespace(dspark_enabled=True, dspark_block_size=5) + ), + max_running_req=3, + cuda_graph_max_bs=32, + ) + assert _cpu_moe_max_tokens(config) == 32 * 6 diff --git a/tests/dsv4/test_dsv4_multi_token_advance.py b/tests/dsv4/test_dsv4_multi_token_advance.py new file mode 100644 index 0000000..b3faac4 --- /dev/null +++ b/tests/dsv4/test_dsv4_multi_token_advance.py @@ -0,0 +1,70 @@ +"""Advancing a request by more than one token. + +A speculative step emits a block, and each request keeps a different prefix of it. The +engine's decode step advances every request by exactly one, so this is the first thing +that has to change -- and getting it wrong desynchronises a request's position from its +KV, which produces wrong text rather than an error. +""" + +from __future__ import annotations + +import pytest +import torch + +from freetoken.core import Req, SamplingParams + + +def _req(prompt_len: int = 8, budget: int = 64) -> Req: + return Req( + input_ids=torch.arange(prompt_len, dtype=torch.int32), + table_idx=0, + cached_len=0, + output_len=budget, + uid=1, + sampling_params=SamplingParams(), + cache_handle=None, + ) + + +class TestCompleteN: + def test_it_advances_by_exactly_n(self): + r = _req() + before = r.device_len + r.complete_n(3) + assert r.device_len == before + 3 + + def test_cached_len_marks_where_the_step_started(self): + # The accepted prefix becomes cached history, and the next step extends from + # the end of it. Pointing cached_len anywhere else re-runs or skips tokens. + r = _req() + start = r.device_len + r.complete_n(4) + assert r.cached_len == start + assert r.extend_len == 4 + + def test_n_equals_one_matches_complete_one(self): + a, b = _req(), _req() + a.complete_one() + b.complete_n(1) + assert (a.device_len, a.cached_len) == (b.device_len, b.cached_len) + + def test_a_step_that_advances_by_nothing_is_refused(self): + # A speculative step always emits at least one token (the target's own token at + # the first rejection). Zero means the caller lost that guarantee. + with pytest.raises(ValueError, match="at least one"): + _req().complete_n(0) + + def test_repeated_steps_accumulate(self): + r = _req() + start = r.device_len + for n in (3, 1, 5, 2): + r.complete_n(n) + assert r.device_len == start + 11 + + def test_append_host_and_complete_n_stay_consistent(self): + # The token buffer and the position counter must agree: device_len is what the + # KV and page table are sized against, input_ids is what the client receives. + r = _req(prompt_len=8) + r.append_host(torch.tensor([101, 102, 103], dtype=torch.int32)) + r.complete_n(3) + assert r.input_ids.numel() == r.device_len == 11 diff --git a/tests/dsv4/test_dsv4_proposal_writeback.py b/tests/dsv4/test_dsv4_proposal_writeback.py new file mode 100644 index 0000000..d3ea46d --- /dev/null +++ b/tests/dsv4/test_dsv4_proposal_writeback.py @@ -0,0 +1,71 @@ +"""DSpark proposal layout from the paper's gamma queries to target verification.""" + +from __future__ import annotations + +import pathlib + + +def _method(name: str) -> str: + src = ( + pathlib.Path(__file__).resolve().parents[2] + / "python" / "freetoken" / "engine" / "engine.py" + ).read_text() + i = src.index(f" def {name}(") + end = src.find("\n def ", i + 10) + return src[i : len(src) if end < 0 else end] + + +def _scheduler_method(name: str) -> str: + src = ( + pathlib.Path(__file__).resolve().parents[2] + / "python" / "freetoken" / "scheduler" / "scheduler.py" + ).read_text() + i = src.index(f" def {name}(") + end = src.find("\n def ", i + 10) + return src[i : len(src) if end < 0 else end] + + +class TestProposalLayout: + def test_all_gamma_proposals_fill_target_rows_one_through_gamma(self): + body = _method("draft_into_batch") + assert "i * span + 1:(i + 1) * span" in body + assert "proposed[i * k:(i + 1) * k]" in body + + def test_proposals_stay_on_gpu_until_acceptance(self): + body = _method("draft_into_batch") + assert "batch.draft_tokens = proposed" in body + assert "req.input_ids[-k:]" not in body + + def test_acceptance_reads_the_saved_proposals_not_noise_placeholders(self): + body = _method("_finish_speculative") + assert "proposed_gpu = batch.draft_tokens" in body + assert "proposed_cpu[i * k:(i + 1) * k]" in body + + def test_only_the_accepted_prefix_reaches_the_request_buffer(self): + body = _method("_finish_speculative") + assert "req.append_host(proposed[:n_acc]" in body + + def test_target_span_is_anchor_plus_gamma(self): + assert "span = 1 + k" in _method("draft_into_batch") + + def test_draft_masks_are_built_on_gpu_from_the_checkpoint_noise_id(self): + dspark = ( + pathlib.Path(__file__).resolve().parents[2] + / "python" / "freetoken" / "models" / "deepseek_v4" / "dspark.py" + ).read_text() + start = dspark.index(" def propose(") + end = dspark.index("\n def embed_block", start) + body = dspark[start:end] + assert "torch.full(" in body + assert "self.noise_token_id" in body + assert "draft_ids[:, 0] = input_ids" in body + + def test_bonus_writeback_uses_the_post_acceptance_frontier(self): + body = _scheduler_method("_forward") + assert "[req.device_len - 1 for req in batch.reqs]" in body + assert "self.token_pool[rows, positions]" in body + + def test_acceptance_uses_each_requests_sampling_params(self): + body = _method("_finish_speculative") + assert "sp = req.sampling_params" in body + assert "sp = batch.reqs[0].sampling_params" not in body diff --git a/tests/dsv4/test_dsv4_rejection_sampling.py b/tests/dsv4/test_dsv4_rejection_sampling.py new file mode 100644 index 0000000..fe86c10 --- /dev/null +++ b/tests/dsv4/test_dsv4_rejection_sampling.py @@ -0,0 +1,213 @@ +"""Speculative sampling must not change what the model would have written. + +This is the only property that matters here, and the only one that cannot be checked by +inspection: over many draws, the emitted token has to be distributed exactly as the +TARGET's own distribution p, whatever the drafter q proposed. If it is not, speculation +is a silent quality change -- output that reads fine, is faster, and is not the model +the caller asked for. + +So these tests are statistical. They drive the acceptance rule thousands of times with +a fixed seed and compare the empirical distribution against p. +""" + +from __future__ import annotations + +import pytest +import torch + +from freetoken.models.deepseek_v4.dspark import ( + rejection_accept, + rejection_accept_device, + sampling_probs, +) + +VOCAB = 6 + + +def _probs(*weights: float) -> torch.Tensor: + t = torch.tensor(weights, dtype=torch.float) + return t / t.sum() + + +def _empirical(q_row, p_rows, draws=20000, seed=0): + """Run one drafted position many times; return the emitted-token distribution. + + The draft samples from q, exactly as the real drafter does, so this exercises the + accept/reject path over the whole support rather than one fixed proposal. + """ + g = torch.Generator().manual_seed(seed) + counts = torch.zeros(VOCAB) + for _ in range(draws): + x = torch.multinomial(q_row, 1, generator=g) + n_acc, tok = rejection_accept(x, q_row.unsqueeze(0), p_rows, generator=g) + counts[int(x) if n_acc == 1 else tok] += 1 + return counts / counts.sum() + + +class TestDistributionIsPreserved: + def test_emitted_tokens_follow_the_target_not_the_draft(self): + q = _probs(5, 1, 1, 1, 1, 1) # drafter loves token 0 + p = _probs(1, 1, 3, 1, 1, 1) # target prefers token 2 + p_rows = torch.stack([p, p]) + got = _empirical(q, p_rows) + assert torch.allclose(got, p, atol=0.02), ( + f"emitted {got.tolist()} but the target's distribution is {p.tolist()} -- " + "speculation changed the output distribution" + ) + + def test_it_holds_when_draft_and_target_agree(self): + p = _probs(3, 2, 1, 1, 1, 1) + got = _empirical(p, torch.stack([p, p])) + assert torch.allclose(got, p, atol=0.02) + + def test_it_holds_when_they_barely_overlap(self): + # The hard case: almost every proposal is rejected, so nearly every emitted + # token comes from the residual. If the residual is wrong, this is where it shows. + q = _probs(10, 10, 1, 1, 1, 1) + p = _probs(1, 1, 1, 1, 10, 10) + got = _empirical(q, torch.stack([p, p])) + assert torch.allclose(got, p, atol=0.03) + + @pytest.mark.parametrize("seed", [1, 2, 3]) + def test_random_pairs_preserve_the_target(self, seed): + g = torch.Generator().manual_seed(seed) + q = torch.softmax(torch.randn(VOCAB, generator=g), dim=-1) + p = torch.softmax(torch.randn(VOCAB, generator=g), dim=-1) + got = _empirical(q, torch.stack([p, p]), draws=20000, seed=seed) + assert torch.allclose(got, p, atol=0.03) + + +class TestAcceptanceBehaviour: + def test_a_proposal_the_target_is_certain_of_is_always_accepted(self): + q = _probs(1, 1, 1, 1, 1, 1) + p = torch.zeros(VOCAB) + p[3] = 1.0 + g = torch.Generator().manual_seed(0) + n_acc, _ = rejection_accept( + torch.tensor([3]), q.unsqueeze(0), torch.stack([p, p]), generator=g + ) + assert n_acc == 1, "p(x)/q(x) >= 1 must accept outright" + + def test_a_proposal_the_target_rules_out_is_always_rejected(self): + q = _probs(1, 1, 1, 1, 1, 1) + p = torch.zeros(VOCAB) + p[3] = 1.0 + g = torch.Generator().manual_seed(0) + n_acc, tok = rejection_accept( + torch.tensor([0]), q.unsqueeze(0), torch.stack([p, p]), generator=g + ) + assert (n_acc, tok) == (0, 3), "p(x) == 0 must reject and resample to p's mass" + + def test_acceptance_stops_at_the_first_rejection(self): + # Position 1 is impossible under p, so position 2 must never be reached even + # though the target would have agreed with it. + p_bad = torch.zeros(VOCAB) + p_bad[5] = 1.0 + p_ok = torch.zeros(VOCAB) + p_ok[1] = 1.0 + q = _probs(1, 1, 1, 1, 1, 1) + g = torch.Generator().manual_seed(0) + n_acc, _ = rejection_accept( + torch.tensor([1, 0, 1]), + torch.stack([q, q, q]), + torch.stack([p_ok, p_bad, p_ok, p_ok]), + generator=g, + ) + assert n_acc == 1 + + def test_a_fully_accepted_block_draws_its_bonus_from_the_target(self): + p = torch.zeros(VOCAB) + p[4] = 1.0 + q = _probs(1, 1, 1, 1, 9, 1) + g = torch.Generator().manual_seed(0) + n_acc, tok = rejection_accept( + torch.tensor([4, 4]), torch.stack([q, q]), torch.stack([p, p, p]), generator=g + ) + assert (n_acc, tok) == (2, 4) + + def test_a_step_always_emits_at_least_one_token(self): + g = torch.Generator().manual_seed(0) + for seed in range(20): + gg = torch.Generator().manual_seed(seed) + q = torch.softmax(torch.randn(VOCAB, generator=gg), dim=-1) + p = torch.softmax(torch.randn(VOCAB, generator=gg), dim=-1) + x = torch.multinomial(q, 2, replacement=True, generator=gg) + n_acc, tok = rejection_accept( + x, torch.stack([q, q]), torch.stack([p, p, p]), generator=g + ) + assert 0 <= n_acc <= 2 + assert 0 <= tok < VOCAB + + +class TestDegenerateInputs: + def test_a_zero_probability_draft_token_is_accepted_rather_than_dividing_by_zero(self): + q = torch.zeros(VOCAB) + q[0] = 1.0 + p = _probs(1, 1, 1, 1, 1, 1) + g = torch.Generator().manual_seed(0) + n_acc, _ = rejection_accept( + torch.tensor([3]), q.unsqueeze(0), torch.stack([p, p]), generator=g + ) + assert n_acc == 1 + + def test_a_residual_of_zero_falls_back_to_the_target(self): + # p entirely inside q: the residual is all zeros, so there is nothing to prefer + # and normalizing it would divide by zero. + p = _probs(1, 1, 0, 0, 0, 0) + q = _probs(2, 2, 1, 1, 1, 1) + g = torch.Generator().manual_seed(0) + n_acc, tok = rejection_accept( + torch.tensor([2]), q.unsqueeze(0), torch.stack([p, p]), generator=g + ) + assert n_acc == 0 + assert tok in (0, 1), "the fallback must still draw from the target's support" + + +class TestDeviceResidentAcceptance: + def test_full_acceptance_and_bonus_stay_exact(self): + q = torch.zeros(2, VOCAB) + q[0, 1] = 1.0 + q[1, 2] = 1.0 + p = torch.zeros(3, VOCAB) + p[0, 1] = 1.0 + p[1, 2] = 1.0 + p[2, 4] = 1.0 + got = rejection_accept_device( + torch.tensor([1, 2]), q, p, generator=torch.Generator().manual_seed(0) + ) + assert got == (2, 4) + + def test_rejection_recovers_from_the_target_residual(self): + q = torch.zeros(1, VOCAB) + q[0, 0] = 1.0 + p = torch.zeros(2, VOCAB) + p[0, 3] = 1.0 + p[1, 4] = 1.0 + got = rejection_accept_device( + torch.tensor([0]), q, p, generator=torch.Generator().manual_seed(0) + ) + assert got == (0, 3) + + def test_zero_width_draws_the_target_bonus(self): + p = torch.zeros(1, VOCAB) + p[0, 5] = 1.0 + got = rejection_accept_device( + torch.empty(0, dtype=torch.long), + torch.empty(0, VOCAB), + p, + generator=torch.Generator().manual_seed(0), + ) + assert got == (0, 5) + + def test_accepted_length_and_bonus_share_one_host_transfer(self): + import inspect + + body = inspect.getsource(rejection_accept_device) + assert 'torch.stack((n_acc_device, bonus_device.to(torch.int64))).to(' in body + assert ".item()" not in body + + +def test_temperature_one_probabilities_equal_the_direct_softmax(): + logits = torch.tensor([[1.0, -2.0, 3.0, 0.5]]) + got = sampling_probs(logits, temperature=1.0, top_p=1.0, top_k=-1) + assert torch.equal(got, torch.softmax(logits.float(), dim=-1)) diff --git a/tests/dsv4/test_dsv4_rotary_contract.py b/tests/dsv4/test_dsv4_rotary_contract.py new file mode 100644 index 0000000..1c28553 --- /dev/null +++ b/tests/dsv4/test_dsv4_rotary_contract.py @@ -0,0 +1,79 @@ +"""Which rotary variant the draft-context write uses. + +There are two, and they are not interchangeable: + +* ``apply_rotary_emb`` takes ONE row of frequencies and broadcasts it across the + sequence. It reshapes freqs_cis to ``[1, T, 1, D]``. +* ``apply_rotary_emb_decode`` takes one row of frequencies PER ROW of x. + +``store_context_kv`` gathers frequencies by absolute position -- ``freqs_cis.index_select +(0, positions)`` -- which is T rows, so it needs the decode variant. Handing those to +``apply_rotary_emb`` fails on the view, which is the lucky outcome: at T == 1 the two +shapes coincide, so a single-token catch-up would pass and every longer one would crash. +Worse, had the view happened to be valid, every context token would have been rotated to +the phase of position zero, and the drafter would attend against phases the target never +used -- degrading acceptance with nothing raised. +""" + +from __future__ import annotations + +import pathlib + +import pytest + +SRC = ( + pathlib.Path(__file__).resolve().parents[2] + / "python" / "freetoken" / "models" / "deepseek_v4" / "dspark.py" +).read_text() + + +class TestTheDecodeVariantIsUsed: + def test_store_context_kv_uses_the_per_row_rotary(self): + body = SRC[SRC.index("def store_context_kv") : SRC.index("def catch_up_context")] + assert "apply_rotary_emb_decode(" in body, ( + "store_context_kv gathers per-position frequencies, so it must use the " + "per-row rotary" + ) + + def test_it_does_not_use_the_broadcast_rotary(self): + body = SRC[SRC.index("def store_context_kv") : SRC.index("def catch_up_context")] + calls = [ + ln for ln in body.splitlines() + if "apply_rotary_emb(" in ln and not ln.strip().startswith("#") + ] + assert not calls, ( + f"the broadcast rotary cannot take T rows of frequencies: {calls}" + ) + + def test_a_head_dim_is_added_before_the_call(self): + # The decode rotary broadcasts each row's frequencies over an inner head dim, so + # a [T, rd] slice has to become [T, 1, rd] first. + body = SRC[SRC.index("def store_context_kv") : SRC.index("def catch_up_context")] + assert "unsqueeze(1)" in body, ( + "the decode rotary needs a head dim to broadcast each row's freqs over" + ) + + +class TestTheContractsDiffer: + """Pin the difference itself, so neither helper can quietly grow the other's shape.""" + + def test_the_broadcast_variant_rejects_per_row_frequencies(self): + torch = pytest.importorskip("torch") + + from freetoken.models.deepseek_v4.ops import apply_rotary_emb + + T, rd = 5, 64 + x = torch.zeros(T, 1, rd) + per_row = torch.ones(T, rd // 2, dtype=torch.complex64) + with pytest.raises(RuntimeError): + apply_rotary_emb(x, per_row) + + def test_a_single_token_hides_the_difference(self): + # Why this was not caught earlier: at T == 1 the two shapes coincide, so the + # wrong variant works exactly until a block is wider than one token. + torch = pytest.importorskip("torch") + + from freetoken.models.deepseek_v4.ops import apply_rotary_emb + + x = torch.zeros(1, 1, 64) + apply_rotary_emb(x, torch.ones(1, 32, dtype=torch.complex64)) # no raise diff --git a/tests/dsv4/test_dsv4_slot_release.py b/tests/dsv4/test_dsv4_slot_release.py new file mode 100644 index 0000000..800cc66 --- /dev/null +++ b/tests/dsv4/test_dsv4_slot_release.py @@ -0,0 +1,154 @@ +"""A rejected block must give its pages and SWA slots back. + +allocate_paged sizes itself from ``req.device_len``, so a speculative step allocates for +the whole block before it knows how much of it survives. Acceptance then lowers +device_len to the prefix it kept. Nothing else reclaims the difference -- the per-step +window free walks the live range, and the end-of-request free walks the page table from +the request's CURRENT length, so the abandoned tail is outside both. + +The leak never fails a decode. It surfaces at the next idle integrity check, with a +message that names neither speculation nor the request that leaked: + + AssertionError: SWA-slot leak/double-free: free(11520) + tree(256) != capacity(12160) + +These tests fail at the cause instead: conservation across an accept, a partial reject, +and a total reject. +""" + +from __future__ import annotations + +import pathlib + +import pytest + +torch = pytest.importorskip("torch") + + +class _Req: + def __init__(self, table_idx=0, device_len=0): + self.table_idx = table_idx + self.device_len = device_len + + +class _Pool: + """Counts slots in and out. Conservation is the whole property under test.""" + + def __init__(self): + self.live: set[int] = set() + + def alloc_swa(self, idx): + for i in idx.tolist(): + assert i not in self.live, f"double-alloc of swa slot {i}" + self.live.add(i) + + def free_swa(self, idx): + for i in idx.tolist(): + assert i in self.live, f"double-free of swa slot {i}" + self.live.discard(i) + + +def _manager(page_size=1, capacity=64): + """A CacheManager with only the fields release_speculative_tail touches.""" + from freetoken.scheduler.cache import CacheManager + + m = CacheManager.__new__(CacheManager) + m.page_size = page_size + m.swa_paged = True + m.is_swa = True + m.swa_pool = _Pool() + m.device = torch.device("cpu") + m.free_slots = torch.empty(0, dtype=torch.int32) + m.page_table = torch.zeros(1, capacity, dtype=torch.int32) + return m + + +def _occupy(m, req, lo, hi): + """Pretend allocate_paged handed positions [lo, hi) the slots [lo, hi).""" + slots = torch.arange(lo, hi, dtype=torch.int32) + m.page_table[req.table_idx, lo:hi] = slots + m.swa_pool.alloc_swa(slots) + req.device_len = hi + + +class TestConservation: + def test_a_fully_rejected_block_returns_every_slot(self): + m, req = _manager(), _Req() + _occupy(m, req, 0, 1) # the committed token + _occupy(m, req, 1, 6) # a 5-wide block + m.release_speculative_tail(req, 1) + assert m.swa_pool.live == {0}, "only the committed token should still hold a slot" + assert m.free_slots.tolist() == [1, 2, 3, 4, 5] + + def test_a_partly_accepted_block_returns_only_the_tail(self): + m, req = _manager(), _Req() + _occupy(m, req, 0, 1) + _occupy(m, req, 1, 6) + m.release_speculative_tail(req, 3) # kept 2 positions plus the bonus + assert sorted(m.swa_pool.live) == [0, 1, 2] + assert m.free_slots.tolist() == [3, 4, 5] + + def test_a_fully_accepted_block_returns_nothing(self): + m, req = _manager(), _Req() + _occupy(m, req, 0, 6) + m.release_speculative_tail(req, 6) + assert len(m.swa_pool.live) == 6 + assert m.free_slots.numel() == 0 + + def test_release_is_a_no_op_above_the_current_length(self): + # Defensive: a caller that passes a length beyond what was allocated must not + # free slots the request never held. + m, req = _manager(), _Req() + _occupy(m, req, 0, 4) + m.release_speculative_tail(req, 10) + assert len(m.swa_pool.live) == 4 + assert m.free_slots.numel() == 0 + + def test_an_unallocated_request_is_skipped(self): + m, req = _manager(), _Req(table_idx=-1, device_len=5) + m.release_speculative_tail(req, 1) # must not index page_table[-1] + assert m.free_slots.numel() == 0 + + @pytest.mark.parametrize("keep", [1, 2, 3, 4, 5, 6]) + def test_every_prefix_conserves_the_total(self, keep): + m, req = _manager(), _Req() + _occupy(m, req, 0, 6) + m.release_speculative_tail(req, keep) + assert len(m.swa_pool.live) + m.free_slots.numel() == 6, ( + "every slot must be either still held or back on the free list" + ) + + +class TestPagedGeometry: + def test_free_slots_receives_one_entry_per_page(self): + # free_slots holds page-START slots. Pushing every token slot would inflate the + # free list and hand out overlapping pages on the next allocation. + m, req = _manager(page_size=4), _Req() + _occupy(m, req, 0, 16) + m.release_speculative_tail(req, 4) + assert m.free_slots.tolist() == [4, 8, 12] + assert len(m.swa_pool.live) == 4, "all 12 freed token slots leave the swa pool" + + +class TestItIsActuallyCalled: + """The release exists only if the acceptance path invokes it.""" + + def test_the_engine_releases_before_lowering_device_len(self): + src = ( + pathlib.Path(__file__).resolve().parents[2] + / "python" / "freetoken" / "engine" / "engine.py" + ).read_text() + body = src[src.index("def _finish_speculative") :] + body = body[: body.index("\n def ", 10)] + rel = body.index("release_tail(req") + lower = body.index("req.cached_len, req.device_len = keep") + assert rel < lower, ( + "release_speculative_tail reads req.device_len to find the range to return, " + "so it must run BEFORE acceptance lowers it" + ) + + def test_the_scheduler_supplies_it(self): + src = ( + pathlib.Path(__file__).resolve().parents[2] + / "python" / "freetoken" / "scheduler" / "scheduler.py" + ).read_text() + assert "batch.release_tail = self.cache_manager.release_speculative_tail" in src diff --git a/tests/dsv4/test_dsv4_tensor_parallel.py b/tests/dsv4/test_dsv4_tensor_parallel.py new file mode 100644 index 0000000..1f38b2a --- /dev/null +++ b/tests/dsv4/test_dsv4_tensor_parallel.py @@ -0,0 +1,242 @@ +"""DSV4 tensor parallelism: the shard contract, on the meta device. + +Every sharded parameter must be exactly ``1/tp`` of its TP=1 shape on exactly ONE axis, +and the routed FP4 expert banks must divide the same way -- that division is what lets +N ranks hold the same host bank bytes that one rank holds today, instead of N copies. +Replicated tensors (the MLA latent KV path, the compressors, the Lightning Indexer, the +router) must keep their full shape, because every rank reads the same latent KV and must +select the same blocks. + +CPU-only: shapes are read off a meta-device build, so no weights and no GPU. +""" + +from __future__ import annotations + +import pytest +import torch + +import freetoken.distributed.info as info_mod +from freetoken.distributed import DistributedInfo +from freetoken.models.deepseek_v4.args import DeepseekV4Args +from freetoken.models.deepseek_v4.weight import _expert_bank_specs + +# o_groups=8 bounds the split: a rank must own whole output groups. +TP_SIZES = (2, 4, 8) + +REPLICATED = (".wkv", ".kv_norm", ".q_norm", ".wq_a", ".compressor.", ".indexer.", "hc_") + + +@pytest.fixture +def args() -> DeepseekV4Args: + # Shipped DSV4-Flash widths (every divisibility rule keys off these), on a 4-layer + # stack so the meta build stays cheap. The ratio pattern mirrors tests/dsv4's: + # one uncompressed layer, one indexed (ratio 4) layer, one ratio-128 layer. + return DeepseekV4Args( + max_batch_size=1, + max_seq_len=4096, + n_layers=4, + n_hash_layers=1, + compress_ratios=(0, 4, 128, 4), + ) + + +def _set_tp(size: int, rank: int = 0) -> None: + # set_tp_info() is write-once by design; these tests walk several sizes in one process. + info_mod._TP_INFO = DistributedInfo(rank, size) + + +def _shapes(args: DeepseekV4Args, tp: int, rank: int = 0) -> dict[str, tuple[int, ...]]: + _set_tp(tp, rank) + from freetoken.models.deepseek_v4.model import Transformer + + with torch.device("meta"): + model = Transformer(args) + return {n: tuple(p.shape) for n, p in model.named_parameters()} + + +@pytest.fixture(autouse=True) +def _restore_tp(): + yield + _set_tp(1) + + +@pytest.mark.parametrize("tp", TP_SIZES) +def test_every_shard_is_a_clean_split_on_one_axis(args, tp): + base = _shapes(args, 1) + got = _shapes(args, tp) + assert got.keys() == base.keys(), "TP must not add or drop parameters" + + sharded = [n for n in got if got[n] != base[n]] + assert sharded, f"tp={tp} sharded nothing" + for name in sharded: + a, b = base[name], got[name] + axes = [i for i in range(len(a)) if a[i] != b[i]] + assert len(axes) == 1, f"{name}: {a} -> {b} splits on {len(axes)} axes, want 1" + assert a[axes[0]] == b[axes[0]] * tp, f"{name}: {a} -> {b} is not a 1/{tp} split" + + +@pytest.mark.parametrize("tp", TP_SIZES) +def test_replicated_tensors_keep_their_full_shape(args, tp): + base = _shapes(args, 1) + got = _shapes(args, tp) + for name, shape in got.items(): + if any(marker in name for marker in REPLICATED): + assert shape == base[name], f"{name} must stay replicated under TP" + + +@pytest.mark.parametrize("tp", TP_SIZES) +def test_ranks_cover_the_vocabulary_exactly_once(args, tp): + _set_tp(tp) + from freetoken.models.deepseek_v4.model import Transformer + + covered = 0 + for rank in range(tp): + _set_tp(tp, rank) + with torch.device("meta"): + model = Transformer(args) + assert model.vocab_start == covered + covered += model.vocab_local + assert covered == args.vocab_size + + +@pytest.mark.parametrize("tp", TP_SIZES) +def test_expert_banks_divide_and_tile_the_intermediate_dim(args, tp): + _set_tp(1) + full, full_i, _ = _expert_bank_specs(args) + + covered = 0 + for rank in range(tp): + _set_tp(tp, rank) + specs, i_local, i_lo = _expert_bank_specs(args) + assert i_lo == covered, "rank slices must tile the intermediate dim with no gap" + covered += i_local + assert i_local * tp == full_i + # A rank's bank bytes are exactly 1/tp of the whole -- the memory win. + for name, (shape, dtype) in specs.items(): + whole = torch.Size(full[name][0]).numel() + part = torch.Size(shape).numel() + assert part * tp == whole, f"{name}: {shape} is not 1/{tp} of {full[name][0]}" + assert dtype == full[name][1] + assert covered == full_i + + +@pytest.mark.parametrize("dim", [0, 1]) +def test_a_shard_does_not_keep_its_parent_alive(dim): + """A shard must own its storage. + + ``narrow`` returns a view, and a dim-0 view is already contiguous, so a + ``.contiguous()`` there hands the view straight back and pins the whole parent. + Every rank then pays for the full tensor it just sharded -- 5.1 GiB per GPU on + DSV4-Flash at TP=4, charged to the model and taken out of the cache budget. + """ + from freetoken.models.deepseek_v4.parallel import shard + + _set_tp(4, rank=1) + parent = torch.zeros(64, 64) + piece = shard(parent, dim) + assert piece.shape[dim] == 16 + assert piece.is_contiguous() + assert piece.untyped_storage().nbytes() == piece.numel() * piece.element_size(), ( + "shard is a view into the parent's storage" + ) + assert piece.data_ptr() != parent.data_ptr() + + +def test_a_split_that_does_not_divide_o_groups_fails_loudly(args): + from freetoken.models.deepseek_v4.parallel import validate_tp + + _set_tp(16) # o_groups == 8, so a rank cannot own a whole group + with pytest.raises(ValueError, match="o_groups"): + validate_tp(args) + + +def test_the_expert_banks_and_the_offload_cache_agree_on_layer_count(args): + """The banks and the cache must be built for the SAME number of MoE layers. + + They are derived independently -- the banks from the checkpoint's mtp.* keys, the + cache from ModelConfig.num_moe_layers -- so a drafter that adds layers to one and + not the other asserts at startup, after the full expert load has already run: + + AssertionError: ('gate_up_packed', 46) + + which costs a five-minute load to discover. Check it in a millisecond instead. + """ + import dataclasses + + from freetoken.models.deepseek_v4.config import parse_config + from freetoken.models.deepseek_v4.weight import _expert_bank_specs + + class _HF: # parse_config only needs the checkpoint path off the hf config + _name_or_path = "/home/carlos/models/DeepSeek-V4-Flash-0731" + + for enabled in (False, True): + from freetoken.models.deepseek_v4.args import set_dspark_enabled + + set_dspark_enabled(enabled) + try: + cfg = parse_config(_HF()) + except Exception: # no checkpoint on this host -- skip rather than fail + return + a = dataclasses.replace(cfg.dsv4_args, dspark_enabled=enabled) + _specs, _i, _lo = _expert_bank_specs(a) + assert cfg.num_moe_layers == a.n_moe_layers, ( + f"dspark_enabled={enabled}: offload cache expects {cfg.num_moe_layers} " + f"layers, expert banks build {a.n_moe_layers}" + ) + set_dspark_enabled(False) + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_every_consumer_of_the_moe_layer_count_agrees(enabled): + """Four places derive "how many MoE layers are there" independently. + + Enabling the dSpark drafter adds three, and each consumer learned about them in a + separate commit -- every miss cost a full expert load to discover, because the + assertions fire only after the banks are built: + + AssertionError: ('gate_up_packed', 46) # the offload cache's banks + assert len(layers) == num_moe_layers # the model's MoE layer iterator + + So check all of them together, offline, instead of one per five-minute run. + """ + import dataclasses + + import torch + + from freetoken.models.deepseek_v4.args import load_args, set_dspark_enabled + from freetoken.models.deepseek_v4.model import DeepseekV4ForCausalLM + from freetoken.models.deepseek_v4.weight import _expert_bank_specs + + model_path = "/home/carlos/models/DeepSeek-V4-Flash-0731" + _set_tp(4) + set_dspark_enabled(enabled) + try: + args = load_args(model_path, max_seq_len=4096) + except Exception: # no checkpoint on this host + set_dspark_enabled(False) + return + + expected = args.n_moe_layers + assert expected == args.n_layers + (3 if enabled else 0) + + # 1. the KV-owning layer list + assert len(args.layer_compress_ratios) == expected + + # 2. the host expert banks (one entry per layer, per bank) + specs, _i, _lo = _expert_bank_specs(args) + assert all(len(s) == 3 for s, _d in specs.values()), "bank specs are per-layer shapes" + + # 3. the DSV4 KV pool, which indexes window_pool by layer id -- a draft layer + # storing its KV past the end of that list is an IndexError mid-generation, + # not at startup. + from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache + assert len(args.layer_compress_ratios) == expected + + # 4. the model's offload-MoE layer iterator, which the cache counts + class _Cfg: + dsv4_args = args + with torch.device("meta"): + model = DeepseekV4ForCausalLM(_Cfg()) + assert len(list(model._iter_offload_moe_layers())) == expected + + set_dspark_enabled(False) diff --git a/tests/dsv4/test_hc_inv_rms.py b/tests/dsv4/test_hc_inv_rms.py new file mode 100644 index 0000000..61eca39 --- /dev/null +++ b/tests/dsv4/test_hc_inv_rms.py @@ -0,0 +1,67 @@ +"""``inv_rms`` must agree with the expression it replaced, and be no less accurate. + +The old form was ``torch.rsqrt(xf.square().mean(-1, keepdim=True) + eps)`` over the +fp32 upcast. ``inv_rms`` reduces the bf16 source directly: the upcast is exact, so the +per-element squares are identical and only the reduction order differs. Bit-equality is +not asserted -- ATen's order varies with M and is not reproducible -- so the contract +is (a) agreement to fp32 rounding and (b) accuracy no worse than ATen's against fp64. +""" +import pytest +import torch + +from freetoken.kernel.triton.dsv4.norm import inv_rms + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="triton kernel") + +HC_DIM = 4 * 4096 # hc_mult * dim at DSV4-Flash +EPS = 1e-6 + + +def _old(xf, eps): + return torch.rsqrt(xf.square().mean(-1, keepdim=True) + eps) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("rows", [1, 7, 512]) +def test_matches_square_mean(dtype, rows): + dev = "cuda" + x = torch.randn(rows, HC_DIM, device=dev, dtype=dtype) + got = inv_rms(x, EPS) + ref = _old(x.float(), EPS) + assert got.shape == ref.shape == (rows, 1) + assert got.dtype == torch.float32 + torch.testing.assert_close(got, ref, rtol=1e-5, atol=0) + + +def test_keeps_leading_dims(): + """hc_pre feeds a [B, T, hc_dim] view and broadcasts the result against [B, T, mix].""" + dev = "cuda" + x = torch.randn(2, 3, HC_DIM, device=dev, dtype=torch.bfloat16) + got = inv_rms(x, EPS) + assert got.shape == (2, 3, 1) + torch.testing.assert_close(got, _old(x.float(), EPS), rtol=1e-5, atol=0) + + +def test_tiny_and_large_magnitudes(): + """Check the scale range a hidden state can reach.""" + dev = "cuda" + for scale in (1e-3, 1.0, 1e3): + x = torch.randn(4, HC_DIM, device=dev, dtype=torch.bfloat16) * scale + torch.testing.assert_close(inv_rms(x, EPS), _old(x.float(), EPS), + rtol=1e-5, atol=0) + + +def test_accuracy_not_worse_than_aten(): + """The point of the Triton reduce over ``linalg.vector_norm``: no sqrt round-trip, + so it stays level with ATen against an fp64 reference instead of ~23% behind.""" + torch.manual_seed(0) + for rows in (64, 1024): + x = torch.randn(rows, HC_DIM, device="cuda", dtype=torch.bfloat16) + xf = x.float() + xd = xf.double() + truth = torch.rsqrt((xd * xd).mean(-1, keepdim=True) + EPS) + err_new = ((inv_rms(x, EPS).double() - truth).abs() / truth.abs()).mean() + err_aten = ((_old(xf, EPS).double() - truth).abs() / truth.abs()).mean() + # same quality of fp32 answer -- allow a small margin either way, but catch a + # real regression such as accumulating in bf16 or re-squaring a norm. + assert err_new < err_aten * 1.25, (rows, err_new.item(), err_aten.item()) diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py index a164f0b..7cfbc18 100644 --- a/tests/engine/test_cache_budget.py +++ b/tests/engine/test_cache_budget.py @@ -168,6 +168,15 @@ def test_adjust_config_allows_auto_for_dsv4(): assert cfg.page_size == 128 # DSV4's KV page is the P-token window page +def test_adjust_config_honors_dsv4_prefill_chunk_limit(): + """DSV4 continuation prefill is stateful, so a safety cap must survive resolution.""" + from freetoken.engine.engine import _adjust_config + + cfg = _dsv4_adjust_cfg(max_seq_len=8192, max_extend_tokens=1024) + _adjust_config(cfg) + assert cfg.max_extend_tokens == 1024 + + def test_adjust_config_resolves_num_tokens_for_dsv4(): # --num-tokens resolves AFTER every page_size override, so DSV4's P=128 page divides it. from freetoken.engine.engine import _adjust_config diff --git a/tests/engine/test_distribution_summary.py b/tests/engine/test_distribution_summary.py new file mode 100644 index 0000000..f472145 --- /dev/null +++ b/tests/engine/test_distribution_summary.py @@ -0,0 +1,67 @@ +from types import SimpleNamespace + +import torch + +from freetoken.engine import distribution +from freetoken.engine.distribution import format_distribution_summary, host_expert_bytes + + +def test_host_expert_bytes_sums_every_rank_local_bank(): + cache = SimpleNamespace( + bank_sources={ + "gate": [torch.empty(2, 8, dtype=torch.uint8), torch.empty(2, 8, dtype=torch.uint8)], + "down": [torch.empty(2, 4, dtype=torch.bfloat16)], + } + ) + assert host_expert_bytes(cache) == 48 + assert host_expert_bytes(None) == 0 + + +def test_physical_core_count_uses_sysfs_sibling_groups(monkeypatch): + siblings = {0: "0,40", 40: "0,40", 1: "1,41", 41: "1,41"} + monkeypatch.setattr(distribution, "_thread_siblings", siblings.get) + + assert distribution._physical_core_count([0, 1, 40, 41]) == 2 + + +def test_summary_groups_ranks_by_numa_and_reports_cpu_gpu_split(): + gib = 1 << 30 + rows = [ + { + "rank": 0, + "gpu": 0, + "gpu_name": "GPU", + "gpu_used": 12 * gib, + "gpu_total": 16 * gib, + "node": 0, + "node_memory_total": 128 * gib, + "node_physical_cores": 20, + "cpu_workers": 9, + "cpu_coordinator": 1, + "host_experts": 32 * gib, + "system_memory_total": 256 * gib, + }, + { + "rank": 1, + "gpu": 1, + "gpu_name": "GPU", + "gpu_used": 12 * gib, + "gpu_total": 16 * gib, + "node": 0, + "node_memory_total": 128 * gib, + "node_physical_cores": 20, + "cpu_workers": 9, + "cpu_coordinator": 1, + "host_experts": 32 * gib, + "system_memory_total": 256 * gib, + }, + ] + + lines = format_distribution_summary(rows) + + assert lines[0] == "Model distribution (final, NUMA-aware):" + assert "host experts 64.00 GiB/128.00 GiB (50.0% RAM)" in lines[1] + assert "CPU cores 20/20 assigned" in lines[1] + assert "GPU 0 12.00 GiB/16.00 GiB (75.0% VRAM)" in lines[2] + assert "CPU-MoE 9 workers + 1 coordinator" in lines[2] + assert "counted placement CPU 72.7% / GPU 27.3%" in lines[-1] diff --git a/tests/engine/test_numa_placement.py b/tests/engine/test_numa_placement.py new file mode 100644 index 0000000..0c69d12 --- /dev/null +++ b/tests/engine/test_numa_placement.py @@ -0,0 +1,211 @@ +"""Spreading tensor-parallel ranks over NUMA nodes. + +A rank's expert banks are anonymous memory: their pages land on whichever node first +TOUCHES them. Get the placement wrong and nothing fails -- the banks simply sit on the +far side of the interconnect from the cores that read them, and decode is slower for +reasons no log mentions. + +So the placement is checked against topologies this host does not have: one node, four, +eight, uneven rank counts, and a cpuset that hides part of a node. +""" + +from __future__ import annotations + +import pytest + +from freetoken.utils import numa + + +@pytest.fixture +def fake_topology(monkeypatch): + """Install a synthetic NUMA layout: a list of CPU lists, one per node.""" + + def install(nodes: list[list[int]], allowed: set[int] | None = None): + mask = allowed if allowed is not None else {c for n in nodes for c in n} + monkeypatch.setattr(numa, "allowed_cpus", lambda: set(mask)) + monkeypatch.setattr( + numa, "numa_nodes", + lambda a=None: [sorted(set(n) & set(mask)) for n in nodes if set(n) & set(mask)], + ) + # Synthetic topology means synthetic GPUs too. Without this the host's REAL + # GPU affinity leaks in and silently overrides the layout under test -- which + # is how these tests first broke when GPU-aware placement landed. + monkeypatch.setattr(numa, "device_numa_node", lambda i: None) + + return install + + +class TestNoPlacementNeeded: + def test_single_node_is_a_no_op(self, fake_topology): + fake_topology([list(range(16))]) + assert numa.rank_placement(0, 4) is None, "one node: nothing to spread" + + def test_single_rank_is_a_no_op(self, fake_topology): + fake_topology([[0, 1], [2, 3]]) + assert numa.rank_placement(0, 1) is None + + +class TestOneRankPerNode: + def test_four_ranks_over_four_nodes_get_one_each(self, fake_topology): + fake_topology([[0, 1], [2, 3], [4, 5], [6, 7]]) + seen = [numa.rank_placement(r, 4) for r in range(4)] + assert [p[0] for p in seen] == [0, 1, 2, 3] + assert all(p[2] == 1 for p in seen), "each rank owns its node alone" + + def test_two_ranks_over_four_nodes_spread_out(self, fake_topology): + # Half a machine should still use distinct memory controllers, not crowd node 0. + fake_topology([[0], [1], [2], [3]]) + assert numa.rank_placement(0, 2)[0] == 0 + assert numa.rank_placement(1, 2)[0] == 2 + + def test_a_rank_never_straddles_two_nodes(self, fake_topology): + fake_topology([[0, 1], [2, 3], [4, 5], [6, 7]]) + for r in range(3): + _node, cpus, _s, _i = numa.rank_placement(r, 3) + assert cpus in ([0, 1], [2, 3], [4, 5], [6, 7]), ( + "splitting one rank over nodes recreates the remote access this avoids" + ) + + +class TestMoreRanksThanNodes: + def test_four_ranks_over_two_nodes_pair_up(self, fake_topology): + fake_topology([list(range(20)), list(range(20, 40))]) + placements = [numa.rank_placement(r, 4) for r in range(4)] + assert [p[0] for p in placements] == [0, 0, 1, 1], "contiguous blocks" + assert [p[2] for p in placements] == [2, 2, 2, 2], "two ranks per node" + assert [p[3] for p in placements] == [0, 1, 0, 1], "index within the node" + + def test_an_uneven_split_reports_the_exact_count_per_node(self, fake_topology): + # 3 ranks over 2 nodes: one node holds two, the other one. A rank that assumed + # the average would either oversubscribe its cores or leave half of them idle. + fake_topology([[0, 1, 2, 3], [4, 5, 6, 7]]) + p0, p1, p2 = (numa.rank_placement(r, 3) for r in range(3)) + assert (p0[0], p0[2], p0[3]) == (0, 2, 0) + assert (p1[0], p1[2], p1[3]) == (0, 2, 1) + assert (p2[0], p2[2], p2[3]) == (1, 1, 0), "the lone rank has its node to itself" + + def test_eight_ranks_over_two_nodes(self, fake_topology): + fake_topology([list(range(20)), list(range(20, 40))]) + assert [numa.rank_placement(r, 8)[0] for r in range(8)] == [0, 0, 0, 0, 1, 1, 1, 1] + + def test_every_rank_is_placed_exactly_once(self, fake_topology): + fake_topology([[0, 1], [2, 3], [4, 5]]) + for world in (2, 3, 4, 5, 6, 7): + per_node: dict[int, list[int]] = {} + for r in range(world): + node, _cpus, siblings, index = numa.rank_placement(r, world) + per_node.setdefault(node, []).append(index) + assert 0 <= index < siblings + for node, indices in per_node.items(): + assert sorted(indices) == list(range(len(indices))), ( + f"world={world} node={node}: indices must be 0..n-1 with no gaps" + ) + + +class TestRestrictedCpusets: + def test_a_node_with_no_usable_cpu_is_dropped(self, fake_topology): + # A memory-only node, or one whose cores a container excludes, is not somewhere + # a rank can run -- placing one there would leave it with an empty mask. + fake_topology([[0, 1], [2, 3]], allowed={0, 1}) + assert numa.rank_placement(0, 2) is None, "only one usable node remains" + + def test_placement_uses_only_permitted_cpus(self, fake_topology): + fake_topology([[0, 1, 2, 3], [4, 5, 6, 7]], allowed={0, 1, 4, 5}) + for r in range(2): + _n, cpus, _s, _i = numa.rank_placement(r, 2) + assert set(cpus) <= {0, 1, 4, 5} + + +class TestParsing: + @pytest.mark.parametrize( + "spec,expected", + [ + ("0-3", {0, 1, 2, 3}), + ("0,2,4", {0, 2, 4}), + ("0-1,8-9", {0, 1, 8, 9}), + ("5", {5}), + ("", set()), + ], + ) + def test_cpulist_forms(self, spec, expected): + assert numa._parse_cpulist(spec) == expected + + +def test_a_rank_outside_the_world_is_refused(fake_topology): + fake_topology([[0], [1]]) + with pytest.raises(ValueError, match="outside world size"): + numa.rank_placement(5, 4) + + +class TestPlacementSurvivesBinding: + """Binding erases the evidence the placement was derived from. + + Once a rank's affinity is one node, the topology LOOKS single-node, so + rank_placement correctly answers "nothing to spread" -- and every later consumer + (the torch thread split, the CPU MoE core split) falls back to dividing by the whole + world size. Each rank then gets a fraction of a fraction and most of the socket + sits idle, with nothing in the logs but two lines that disagree: + + rank 0 -> NUMA node 0 (40 cpus, 2 rank(s) on this node) + torch intra-op threads: 10 per rank (shared with 3 sibling rank(s)) + """ + + def test_the_answer_is_remembered_across_the_bind(self, fake_topology): + numa.reset_placement() + fake_topology([list(range(20)), list(range(20, 40))]) + before = numa.resolve_placement(0, 4) + assert before[2] == 2, "two ranks share node 0" + + # Simulate the bind: affinity is now one node, so the topology looks single. + fake_topology([list(range(20)), list(range(20, 40))], allowed=set(range(20))) + assert numa.rank_placement(0, 4) is None, "a fresh derivation loses it" + assert numa.placement() == before, "the remembered answer must survive" + numa.reset_placement() + + def test_resolve_is_idempotent(self, fake_topology): + numa.reset_placement() + fake_topology([[0, 1], [2, 3]]) + first = numa.resolve_placement(1, 4) + second = numa.resolve_placement(1, 4) + assert first is second + numa.reset_placement() + + def test_an_unplaced_process_reports_none(self): + numa.reset_placement() + assert numa.placement() is None + + +class TestGpuAffinityDrivesPlacement: + """A rank should follow its GPU, not its index. + + Each GPU hangs off a PCIe root complex owned by one socket, and for an offload MoE + the per-step expert fetch is host-to-device traffic. Placing a rank by index happens + to be right when CUDA enumerates devices in socket order -- and is silently wrong + otherwise, sending every fetch across the interconnect with nothing in the logs. + """ + + def test_placement_follows_the_gpu_not_the_rank_index(self, fake_topology, monkeypatch): + numa.reset_placement() + fake_topology([list(range(20)), list(range(20, 40))]) + # A machine that enumerates GPUs across sockets: 0,2 -> node 0; 1,3 -> node 1. + # Index-derived placement would put ranks 0,1 on node 0 and 2,3 on node 1, so + # ranks 1 and 2 would each be on the wrong socket for their card. + monkeypatch.setattr(numa, "device_numa_node", lambda i: {0: 0, 1: 1, 2: 0, 3: 1}[i]) + assert numa.rank_placement(0, 4)[0] == 0 + assert numa.rank_placement(1, 4)[0] == 1 + assert numa.rank_placement(2, 4)[0] == 0 + assert numa.rank_placement(3, 4)[0] == 1 + + def test_ranks_on_a_node_are_indexed_without_gaps(self, fake_topology, monkeypatch): + fake_topology([list(range(20)), list(range(20, 40))]) + monkeypatch.setattr(numa, "device_numa_node", lambda i: {0: 0, 1: 1, 2: 0, 3: 1}[i]) + # Ranks 0 and 2 share node 0; they must take index 0 and 1 there, so the CPU + # MoE pool hands them different cores. + assert numa.rank_placement(0, 4)[2:] == (2, 0) + assert numa.rank_placement(2, 4)[2:] == (2, 1) + + def test_it_falls_back_to_the_index_when_affinity_is_unknown(self, fake_topology, monkeypatch): + # No sysfs, a virtualized GPU, or a driver that does not report it. + fake_topology([list(range(20)), list(range(20, 40))]) + monkeypatch.setattr(numa, "device_numa_node", lambda i: None) + assert [numa.rank_placement(r, 4)[0] for r in range(4)] == [0, 0, 1, 1] diff --git a/tests/kernels/test_dsv4_indexer_decode.py b/tests/kernels/test_dsv4_indexer_decode.py index a3dc63b..d5480e4 100644 --- a/tests/kernels/test_dsv4_indexer_decode.py +++ b/tests/kernels/test_dsv4_indexer_decode.py @@ -41,10 +41,10 @@ def _build(positions, n_stage, seed=0): return q, w, snap, valid -def _reference(q, w, pool, snap, valid, n_stage): +def _reference(q, w, pool, snap, valid, n_stage, row_map=None): """fp32 ground truth: sum_h relu(q_h . k_t) * w_h over the live blocks, -inf past them.""" b = q.shape[0] - rows = torch.arange(b, device=q.device) + rows = torch.arange(b, device=q.device) if row_map is None else row_map blk = torch.arange(n_stage, device=q.device) live = blk[None, :] < valid[:, None] at = snap[rows[:, None], (blk * RATIO)[None, :]] @@ -55,7 +55,8 @@ def _reference(q, w, pool, snap, valid, n_stage): def _check(pool, positions, n_stage, seed=0): q, w, snap, valid = _build(positions, n_stage, seed) - got = indexer_decode_logits(q, w, pool, snap, valid, n_stage, RATIO).float() + rows = torch.arange(q.shape[0], dtype=torch.int64, device=q.device) + got = indexer_decode_logits(q, w, pool, snap, rows, valid, n_stage, RATIO).float() ref = _reference(q, w, pool, snap, valid, n_stage) assert torch.equal(torch.isfinite(got), torch.isfinite(ref)), "live/-inf column set differs" live = torch.isfinite(ref) @@ -98,12 +99,27 @@ def test_per_row_valid(pool): assert torch.isinf(got[i, v:]).all() +def test_multiple_queries_share_one_request_snapshot(pool): + """DSpark verifies anchor+gamma query rows against one snapshot row per request.""" + n_stage = 256 + q, w, _, _ = _build([400] * 6, n_stage, seed=17) + _, _, snap, _ = _build([400], n_stage, seed=19) + rows = torch.zeros(6, dtype=torch.int64, device=q.device) + valid = torch.tensor([1, 7, 19, 41, 73, 100], dtype=torch.int64, device=q.device) + got = indexer_decode_logits(q, w, pool, snap, rows, valid, n_stage, RATIO).float() + ref = _reference(q, w, pool, snap, valid, n_stage, row_map=rows) + assert torch.equal(torch.isfinite(got), torch.isfinite(ref)) + live = torch.isfinite(ref) + torch.testing.assert_close(got[live], ref[live], **TOL) + + def test_cuda_graph_follows_valid(pool): """A captured graph must read the bound from device memory, not from capture time.""" n_stage = 2048 q, w, snap, valid = _build([800], n_stage) out = torch.empty((1, n_stage), dtype=pool.dtype, device="cuda") - call = lambda: indexer_decode_logits(q, w, pool, snap, valid, n_stage, RATIO) + rows = torch.arange(q.shape[0], dtype=torch.int64, device=q.device) + call = lambda: indexer_decode_logits(q, w, pool, snap, rows, valid, n_stage, RATIO) side = torch.cuda.Stream() side.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(side): diff --git a/tests/moe/test_hybrid_fetch.py b/tests/moe/test_hybrid_fetch.py index e080e16..c072b2f 100644 --- a/tests/moe/test_hybrid_fetch.py +++ b/tests/moe/test_hybrid_fetch.py @@ -106,3 +106,21 @@ def test_hybrid_fixed_cap_unchanged(): cache.ensure_experts_hybrid(0, ids) assert int(cache.num_missing_full.item()) == 8 assert int(cache.num_indices.item()) == 1 + + +def test_hybrid_fetch_spends_cap_on_most_repeated_current_miss(): + """q* still counts unique misses, but its copies should remove the most CPU routes.""" + cache = OffloadMoeCache( + num_layers=1, num_experts=16, cache_size=20, device=torch.device("cpu"), + quant_format="bf16", decode_target="hybrid", hybrid_max_fetch=1, + ) + # Expert 2 was historically recent, but expert 7 occurs three times in this + # multi-row verify. Fetching 7 avoids three CPU GEMVs for the price of one copy. + cache.step.fill_(11) + cache.expert_recency[0, 2] = 10 + ids = torch.tensor([7, 2, 7, 3, 7], dtype=torch.int32) + cache.ensure_experts_hybrid(0, ids) + + assert ids[0] >= 0 and ids[2] == ids[0] and ids[4] == ids[0] + assert ids[1] == -1 and ids[3] == -1 + assert int(cache.src_indices[0]) == 7 diff --git a/tests/server/test_client_disconnect_aborts.py b/tests/server/test_client_disconnect_aborts.py new file mode 100644 index 0000000..0769615 --- /dev/null +++ b/tests/server/test_client_disconnect_aborts.py @@ -0,0 +1,136 @@ +"""A client that hangs up must stop the generation. + +Streaming responses already notice: yielding a chunk into a dead socket is observable, +and stream_with_cancellation acts on it. A NON-streaming request awaits the whole +completion and never looks -- so a caller that disconnects leaves the GPUs generating to +max_tokens for an answer nobody will read, with the queue behind it waiting on that work. + +On an offload MoE that is minutes of the whole machine per abandoned request. +""" + +from __future__ import annotations + +import asyncio +import pathlib + +import pytest + + +class _Request: + """Stand-in for a Starlette Request: disconnects after `after` polls.""" + + def __init__(self, after: int | None = None): + self.after = after + self.polls = 0 + + async def is_disconnected(self) -> bool: + self.polls += 1 + return self.after is not None and self.polls > self.after + + +class _Ack: + def __init__(self, finished=False): + self.finished = finished + self.incremental_output = "" + self.prompt_tokens_delta = 0 + self.completion_tokens_delta = 1 + self.cached_tokens = 0 + self.error = None + + +class _State: + """Only the pieces wait_for_ack_watching_client touches.""" + + def __init__(self, n_acks=100): + from freetoken.server.api_server import FrontendManager + + self.n_acks = n_acks + self.yielded = 0 + self.aborted: list[int] = [] + self._fn = FrontendManager.wait_for_ack_watching_client + + async def wait_for_ack(self, uid): + for i in range(self.n_acks): + self.yielded += 1 + yield _Ack(finished=(i == self.n_acks - 1)) + await asyncio.sleep(0) + + async def abort_user(self, uid): + self.aborted.append(uid) + + def watch(self, uid, request): + return self._fn(self, uid, request) + + +async def _drain(state, request): + got = 0 + try: + async for _ack in state.watch(7, request): + got += 1 + except asyncio.CancelledError: + pass + # abort_user is scheduled, not awaited, so let the loop run it. + await asyncio.sleep(0) + await asyncio.sleep(0) + return got + + +class TestADisconnectStopsTheGeneration: + def test_it_stops_early_instead_of_running_to_completion(self): + state = _State(n_acks=100) + got = asyncio.run(_drain(state, _Request(after=3))) + assert got < 100, "the loop ran to completion for a client that had hung up" + assert got <= 5 + + def test_it_aborts_the_request(self): + state = _State(n_acks=100) + asyncio.run(_drain(state, _Request(after=2))) + assert state.aborted == [7], "the engine must be told to drop the request" + + def test_a_connected_client_receives_everything(self): + state = _State(n_acks=10) + got = asyncio.run(_drain(state, _Request(after=None))) + assert got == 10 + assert state.aborted == [] + + def test_no_request_means_no_watching(self): + # Offline / internal callers have no HTTP request; they must still work. + state = _State(n_acks=10) + got = asyncio.run(_drain(state, None)) + assert got == 10 + assert state.aborted == [] + + +class TestTheNonStreamingPathsUseIt: + """The wrapper is worthless if the handlers keep calling the bare loop.""" + + @staticmethod + def _src(name: str) -> str: + return ( + pathlib.Path(__file__).resolve().parents[2] + / "python" / "freetoken" / "server" / name + ).read_text() + + def test_generate_full_watches_the_client(self): + src = self._src("generation.py") + body = src[src.index("async def _generate_full_impl") :] + assert "_acks(state, uid, request)" in body, ( + "a non-streaming ack loop that does not watch the client will generate to " + "max_tokens for a caller that has gone" + ) + + def test_the_helper_tolerates_a_state_without_the_watcher(self): + # Not every state implementation has it (the test fakes do not). A missing + # disconnect check must cost the abort, not the request. + src = self._src("generation.py") + assert 'getattr(state, "wait_for_ack_watching_client", None)' in src + + def test_chat_completions_passes_its_request_down(self): + src = self._src("openai_api.py") + assert "request=request" in src, ( + "generate_full cannot watch a client it was never given" + ) + + def test_completions_watches_too(self): + src = self._src("openai_api.py") + assert "_acks(state, uid, request)" in src diff --git a/tests/server/test_message_wire.py b/tests/server/test_message_wire.py index 3bd7cc6..8c6e652 100644 --- a/tests/server/test_message_wire.py +++ b/tests/server/test_message_wire.py @@ -89,13 +89,14 @@ def test_user_reply_token_deltas_round_trip(): def test_detokenize_msg_carries_kv_usage_round_trip(): msg = DetokenizeMsg( - uid=3, next_token=42, finished=True, + uid=3, next_token=42, finished=True, token_ids=[40, 41, 42], kv_used_pages=10, kv_total_pages=256, gpu_mem_bytes=1 << 30, mamba_used_slots=7, mamba_total_slots=64, swa_used_tokens=8448, swa_total_tokens=76800, ) decoded = BaseTokenizerMsg.decoder(BaseTokenizerMsg.encoder(msg)) assert isinstance(decoded, DetokenizeMsg) + assert decoded.token_ids == [40, 41, 42] assert (decoded.kv_used_pages, decoded.kv_total_pages, decoded.gpu_mem_bytes) == (10, 256, 1 << 30) assert (decoded.mamba_used_slots, decoded.mamba_total_slots) == (7, 64) assert (decoded.swa_used_tokens, decoded.swa_total_tokens) == (8448, 76800) diff --git a/tests/tokenizer/test_speculative_detokenize.py b/tests/tokenizer/test_speculative_detokenize.py new file mode 100644 index 0000000..5a0a41a --- /dev/null +++ b/tests/tokenizer/test_speculative_detokenize.py @@ -0,0 +1,74 @@ +from freetoken.message import DetokenizeMsg +from freetoken.tokenizer.detokenize import DetokenizeManager + + +class _PieceTokenizer: + eos_token_id = 0 + + _pieces = { + 0: "", + 1: "pres", + 2: "idencia", + 3: " de", + 4: " López", + 5: " Obrador", + 6: ".", + } + + def batch_decode(self, rows): + return ["".join(self._pieces[token] for token in row) for row in rows] + + +def _manager() -> DetokenizeManager: + return DetokenizeManager(_PieceTokenizer(), frozenset({0})) + + +def test_speculative_block_emits_each_token_once(): + manager = _manager() + + chunks = manager.detokenize( + [ + DetokenizeMsg( + uid=7, + next_token=5, + token_ids=[1, 2, 3, 4, 5], + finished=False, + ) + ] + ) + + assert chunks == ["presidencia de López Obrador"] + + +def test_coalesced_messages_for_one_request_do_not_repeat_prefixes(): + manager = _manager() + + chunks = manager.detokenize( + [ + DetokenizeMsg(uid=9, next_token=2, token_ids=[1, 2], finished=False), + DetokenizeMsg(uid=9, next_token=5, token_ids=[3, 4, 5], finished=False), + DetokenizeMsg(uid=9, next_token=0, finished=True), + ] + ) + + assert chunks == ["presidencia", " de López Obrador", ""] + assert 9 not in manager.decode_map + + +def test_finished_block_counts_text_before_eos_once(): + manager = _manager() + + chunks = manager.detokenize( + [ + DetokenizeMsg( + uid=11, + next_token=0, + token_ids=[1, 2, 6, 0], + finished=True, + finish_reason="stop", + ) + ] + ) + + assert chunks == ["presidencia."] + assert 11 not in manager.decode_map