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_sparse.py b/python/freetoken/attention/dsv4_sparse.py index 2c740e2..b944f5e 100644 --- a/python/freetoken/attention/dsv4_sparse.py +++ b/python/freetoken/attention/dsv4_sparse.py @@ -221,6 +221,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)] diff --git a/python/freetoken/core.py b/python/freetoken/core.py index ef0a539..cfb7366 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,23 @@ 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) + 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) + # Compressor carry saved before the block advanced it, restored on rejection. + carry_snapshot: object | 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..903ec41 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -72,8 +72,16 @@ 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 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/engine.py b/python/freetoken/engine/engine.py index ff3c985..0ac380b 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -2,6 +2,7 @@ import gc import math +import glob import os from datetime import timedelta from typing import Any, Dict, Iterable, NamedTuple, Tuple @@ -9,8 +10,17 @@ 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 +255,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). @@ -289,12 +309,98 @@ class ForwardOutput(NamedTuple): copy_done_event: torch.cuda.Event +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 +429,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 +450,11 @@ 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_drafted = 0 + self.spec_threshold = float(ENV.DSPARK_CONFIDENCE_THRESHOLD.value) if is_offload_moe_backend(config.moe_backend): self._init_offload_moe_cache(config) if hasattr(self.model, "prepare_for_runtime"): @@ -511,7 +635,44 @@ 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) + # The drafter's MoE layers are the tail of the sequence (see the model's + # _iter_offload_moe_layers). Keep them on the GPU whatever the CPU plan says: the + # draft is a hard serial dependency of every block -- the verify cannot start + # until it finishes -- so a draft layer on the CPU pool adds the engine's slowest + # path to each block, and it is 3 layers against the target's 43. + # Read the drafter's layer count from dsv4_args, NOT from extra_moe_layers. + # extra_moe_layers is patched onto model_config later than this (parse_config + # runs before the flag exists), so it is still 0 here even though + # num_moe_layers already counts the drafter's layers -- which made this whole + # path silently inert the first time it was written. + dsv4 = getattr(config.model_config, "dsv4_args", None) + n_draft = 0 + if getattr(dsv4, "dspark_enabled", False): + n_draft = int(getattr(dsv4, "n_draft_layers", 0) or 0) + if not n_draft: + n_draft = int(getattr(config.model_config, "extra_moe_layers", 0) or 0) + draft_layer_ids = frozenset(range(n_moe - n_draft, n_moe)) if n_draft else frozenset() + if draft_layer_ids: + # Always report it. The exclusion below is a no-op under --moe-backend + # hybrid (nothing is assigned to the CPU up front), so its absence from the + # log says nothing about whether the drafter's layers were identified at + # all -- and if extra_moe_layers is still 0 here, this whole path is inert + # and the drafter shares the capped fetch with the target. + logger.info_rank0( + f"dSpark: draft MoE layers {min(draft_layer_ids)}..{max(draft_layer_ids)} " + f"of {n_moe} fetch uncapped on the GPU (serial with every block)" + ) + if cpu_layer_ids & draft_layer_ids: + cpu_layer_ids = cpu_layer_ids - draft_layer_ids + elif getattr(dsv4, "dspark_enabled", False): + logger.warning_rank0( + "dSpark is on but no draft MoE layers were identified " + f"(extra_moe_layers={n_draft}): the drafter's layers will share the " + "target's capped fetch and can be computed on the CPU pool, in series " + "with every block" + ) if config.moe_backend == "hybrid": decode_target = "hybrid" elif cpu_layer_ids: @@ -575,6 +736,7 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: if decode_target == "hybrid": self._resolve_hybrid_fetch(config, cache) cache.cpu_layer_ids = cpu_layer_ids + cache.draft_layer_ids = draft_layer_ids # Must be set before CUDA graph capture so the (device-side) accumulation ops are # captured and re-run on every decode replay. cache.collect_stats = config.moe_collect_stats @@ -636,7 +798,18 @@ 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 verify carries block_size ROWS per request, not one: the last + # committed token plus the drafted block. max_tokens sizes the C++ pool's + # per-task scratch, so a pool built for one row per request would be handed + # block_size times that and overrun it. + rows_per_req = 1 + dsv4 = getattr(config.model_config, "dsv4_args", None) + if getattr(dsv4, "dspark_enabled", False): + rows_per_req = max(1, int(getattr(dsv4, "dspark_block_size", 1) or 1)) + max_tokens = max( + config.max_running_req * rows_per_req, config.cuda_graph_max_bs or 0, 1 + ) # gpt-oss mxfp4 carries clamped-swiglu scalars; other formats use the defaults. executor = CpuMoeExecutor( cache, @@ -860,6 +1033,201 @@ def rebuild_runtime_cache( moe_offload_cache=self.moe_offload_cache, ) + def _finish_speculative( + self, batch: Batch, logits: torch.Tensor, args: BatchSamplingArgs + ) -> 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, + draft_width, + rejection_accept, + sampling_probs, + ) + + k = batch.spec_block + sp = batch.reqs[0].sampling_params + greedy = sp.is_greedy + # 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." + ) + # Not self.sampler: a speculative batch has 1+k ROWS per request while the + # sampling args are sized per REQUEST. + target_cpu = logits.argmax(dim=-1).to("cpu", non_blocking=False).to(torch.int32) + p = None if greedy else sampling_probs( + logits.cpu(), sp.temperature, sp.top_p, sp.top_k + ) + q = None if greedy else batch.draft_probs.cpu() + confidence = batch.draft_confidence + + emitted: list[torch.Tensor] = [] + any_rejected = False + release_tail = getattr(batch, "release_tail", None) + off = 0 + for i, req in enumerate(batch.reqs): + span = 1 + k # this request's rows in the flat batch + start = req.device_len - k # where the block began + proposed = req.input_ids[start:start + k] + width = k + if confidence is not None: + width = draft_width( + confidence[off + 1:off + span], self.spec_threshold, 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: + 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. + n_acc, bonus = rejection_accept( + proposed[:width], + q[off:off + width], + p[off:off + width + 1], + ) + if n_acc < width: + any_rejected = True + keep = start + n_acc + req.input_ids = req._ids_buf[:keep] + 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 + 1) + req.cached_len, req.device_len = keep, keep + 1 + emitted.append( + torch.cat([proposed[:n_acc], torch.tensor([bonus], dtype=torch.int32)]) + ) + self._spec_accepted += n_acc + self._spec_drafted += width + off += span + + # Undo the carry the verify advanced across positions this block did not keep. + # The compressor's state is a reduction over the tokens seen, not a function of + # the KV, so a stranded carry cannot be rebuilt -- the request would simply + # continue from state that saw tokens it never emitted. + snap = getattr(batch, "carry_snapshot", None) + if any_rejected and snap is not None: + snap.restore() + + # 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 _snapshot_carry(self, batch: Batch): + """Save the compressor carry a rejected block would strand, or None. + + Only the page each request is currently in can be stranded: a rejection in a + LATER page leaves the surviving page untouched, and the abandoned pages take + their ring blocks with them. + """ + from freetoken.models.deepseek_v4.rollback import CarrySnapshot + + layers = getattr(self.model, "_transformer", None) + if layers is None or not batch.reqs: + return None + try: + slots = torch.stack([ + self.attn_backend.window_slots_of(r.table_idx, r.cached_len - 1, r.cached_len) + .reshape(()) + for r in batch.reqs + ]) + return CarrySnapshot( + self.attn_backend, [b.attn for b in layers.layers], slots + ) + except Exception as exc: # a snapshot is an optimization for correctness, not + # a correctness requirement on its own -- but say so, loudly, rather than + # silently running without the ability to undo a rejection. + logger.warning_rank0(f"dSpark: no carry snapshot this step ({exc})") + return None + + 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. + """ + from freetoken.models.deepseek_v4.dspark import sampling_probs + + drafter = getattr(self.model, "draft", None) + if drafter is None: + return None + # Snapshot the compressor carry BEFORE anything advances it. The verify pass + # walks the compressor across the whole block; if the block is then partly + # rejected, the carry the next real step must read has already been overwritten + # in place, and nothing else can rebuild it. + batch.carry_snapshot = self._snapshot_carry(batch) + with self.ctx.forward_batch(batch): + # Give the draft layers KV for the positions the target has committed since + # the last step. Without it they attend over a stale window and propose + # badly -- which reads as a weak drafter, not a missing call. + catch_up = getattr(self.model, "catch_up_draft_context", None) + if catch_up is not None: + catch_up(batch) + out = drafter() + if out is None: + return None + draft_logits, confidence = out + # q must be the distribution the sampler would DRAW from, not the raw softmax: + # the ratio test compares q against a target p shaped the same way. + sp = batch.reqs[0].sampling_params + q = sampling_probs(draft_logits, sp.temperature, sp.top_p, sp.top_k) + # Draft by SAMPLING from q, not by taking its argmax. Speculative sampling's + # guarantee assumes the proposal was drawn from q; an argmax proposal is not, + # and the acceptance probabilities would no longer preserve p. + proposed = ( + q.argmax(dim=-1) if sp.is_greedy + else torch.multinomial(q, 1).squeeze(-1) + ) + # Position i's logits predict i+1, so the proposal for the block's positions + # 1..k is proposed[0..k-1]; the last entry predicts past the block and is dropped. + batch.input_ids[1:] = proposed[:-1].to(batch.input_ids.dtype) + batch.draft_probs = q + return confidence + def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput: assert torch.cuda.current_stream() == self.stream with self.ctx.forward_batch(batch): @@ -872,6 +1240,9 @@ def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput: # -> stale expert outputs) as a loud error instead of silent corruption. self.cpu_moe_executor.raise_if_unhealthy() + if batch.speculative: + return self._finish_speculative(batch, logits, args) + for req in batch.reqs: req.complete_one() @@ -989,6 +1360,31 @@ def _adjust_dsv4_config(config: EngineConfig, override) -> None: 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 diff --git a/python/freetoken/env.py b/python/freetoken/env.py index 6878dc4..2df433d 100644 --- a/python/freetoken/env.py +++ b/python/freetoken/env.py @@ -71,6 +71,11 @@ class EnvClassSingleton: FLASHINFER_USE_TENSOR_CORES = EnvOption() DISABLE_OVERLAP_SCHEDULING = EnvBool(False) PYNCCL_MAX_BUFFER_SIZE = EnvMem(1024**3) + # dSpark adaptive verification: a drafted position scoring below this is not + # verified. Verification costs a full target pass over whatever width it covers, so + # carrying a tail the drafter itself doubts is work bought and thrown away. 0 keeps + # the whole block, 1 keeps only the first position. + DSPARK_CONFIDENCE_THRESHOLD = EnvFloat(0.5) # GatedDeltaNet recurrent (SSM) state dtype: float32 (default) | bfloat16 | float16. # fp32 matches the Qwen3.x configs (mamba_ssm_dtype); fp16/bf16 halves the GDN state # pool at some precision cost on the long recurrence (mirrors SGLang's mamba_ssm_dtype). 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/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..c3790f0 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,6 +146,13 @@ 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 @@ -128,9 +162,15 @@ def _prefill_segment(self, x_seg, qr_seg, kv_seg, ti: int, start_pos: int, n: in # prefix plus the new tokens, causal-masked to -1 past its own position. w_lo = max(0, 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 +188,21 @@ 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 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) + blocks = 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 +212,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)] diff --git a/python/freetoken/models/deepseek_v4/compress.py b/python/freetoken/models/deepseek_v4/compress.py index 4d578f3..43fe58a 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 @@ -358,7 +358,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)) 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..642d6d9 --- /dev/null +++ b/python/freetoken/models/deepseek_v4/dspark.py @@ -0,0 +1,562 @@ +"""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 + +from typing import Callable, NamedTuple + +import torch +import torch.nn.functional as F +from torch import nn + +from freetoken.kernel.triton.dsv4.bf16_linear import bf16_linear_fp32 +from freetoken.kernel.triton.dsv4.hc import hc_pre_combine + +from .args import DeepseekV4Args +from .layers import Linear, RMSNorm +from .parallel import div_tp +from .rollback import needs_rollback + + +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 is what makes verification ADAPTIVE: a block whose tail positions score + low can be verified short instead of paying for tokens that will be rejected. + Runs in fp32 -- it is one row of arithmetic per position, and the threshold it feeds + decides how many tokens the verify step covers. + """ + + 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, aux_hidden: torch.Tensor, input_ids: torch.Tensor, segments, + positions: torch.Tensor, target_logits_fn, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Run the draft stack over one block and propose its tokens. + + ``aux_hidden`` is the target's tap at ``dspark_target_layer_ids`` for these + positions, ``[1, T, dim * len(target_layer_ids)]``. ``input_ids`` is what the + block is fed -- the last committed token then noise -- and ``segments`` / + ``positions`` are the batch's, so the draft blocks address the same paged slots + the caller allocated. + + Returns ``(draft_logits [T, vocab], confidence [T])`` -- the LOGITS, not a token + choice. Speculative sampling needs the draft's whole distribution q, not just + its argmax: a sampled request accepts a proposal with probability + min(1, p(x)/q(x)) and, on rejection, resamples from the residual p - q. Reducing + to argmax here would throw q away and force the caller into greedy-only + acceptance. + """ + # The block runs on token embeddings ALONE. The projected target hidden + # (main_x) is not added here -- it is what the draft layers derive their CONTEXT + # KV from, over the positions the target has already committed. Adding it to the + # block would also mismatch: aux covers the previous forward's positions, which + # is the prompt on the first step and one token thereafter, never the block. + embeds = self.embed_block(input_ids) + head_hidden = self.head_hidden(embeds, input_ids, segments, positions) + return self.logits(head_hidden[0], target_logits_fn, input_ids.view(-1)) + + 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 + xf = x.flatten(2).float() + rsqrt = torch.rsqrt(xf.square().mean(-1, keepdim=True) + 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, 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 over one block and return the pre-logits 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.norm(self.hc_head(h)) + + def logits(self, head_hidden: torch.Tensor, target_logits_fn, + prev_token_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Draft logits and per-position acceptance confidence. + + ``target_logits_fn`` must be the TARGET's own logits method, not its head + weight. Under tensor parallelism that head is vocabulary-parallel: a bare + ``F.linear`` against it returns this rank's vocabulary SLICE, while the Markov + bias is full-vocabulary because the head is replicated. Adding the two would + line a [B, vocab/tp] tensor up against a [B, vocab] one -- broadcasting either + into nonsense or a shape error, depending on the TP size. Going through the + target's method keeps the all-gather in one place. + """ + markov = self.markov_head.embed(prev_token_ids) + logits = target_logits_fn(head_hidden) + self.markov_head.bias(markov) + return logits, self.confidence_head(head_hidden, markov) + + +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 + scaled = 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 draft_width(confidence: torch.Tensor, threshold: float, block_size: int) -> int: + """How many of the block's positions are worth verifying. + + The confidence head scores each drafted position. Verification costs a full target + pass over whatever width it covers, so carrying tail positions the drafter itself + doubts is paid-for work that will be thrown away. Cut the block at the first + position that scores below ``threshold``, keeping at least one -- that is the + "adaptive" in adaptive verification. + """ + below = (confidence < threshold).nonzero() + if below.numel() == 0: + return min(block_size, int(confidence.numel())) + return max(1, int(below[0])) + + +class StepResult(NamedTuple): + """What one speculative step produced.""" + + tokens: list[int] # what the request actually emits, in order + drafted: int # positions the drafter proposed + verified: int # positions verification covered (<= drafted) + accepted: int # positions the target agreed with + rolled_back: bool # whether the compressor carry had to be restored + + +class SpeculativeLoop: + """One dSpark step: draft a block, verify it, keep the prefix the target agrees with. + + This class exists for the ORDER, which is where speculative decoding goes wrong in + ways that do not crash. Three orderings matter and all three are enforced here: + + 1. Snapshot the compressor carry BEFORE drafting. The draft advances it in place; + once that has happened the state a rejection must return to no longer exists. + 2. Choose the verify width from the drafter's confidence BEFORE verifying, not + after. Verification costs a full target pass over whatever width it covers, so + deciding afterwards spends exactly what the confidence head exists to save. + 3. Accept a PREFIX, and take the target's own token at the first disagreement. A + block therefore always advances by at least one token and can never stall. + + The model-facing operations are injected, so the loop's control flow is testable + without a GPU -- and so the same logic serves both the eager and captured paths. + """ + + def __init__(self, block_size: int, confidence_threshold: float = 0.5, + page_size: int = 128): + if block_size < 1: + raise ValueError(f"block_size must be >= 1, got {block_size}") + self.block_size = block_size + self.confidence_threshold = confidence_threshold + self.page_size = page_size + + def step( + self, + *, + draft: "Callable[[], tuple[torch.Tensor, torch.Tensor]]", + verify: "Callable[[torch.Tensor], torch.Tensor]", + positions: torch.Tensor, + snapshot=None, + ) -> StepResult: + """Run one block. + + ``draft()`` returns ``(proposed_ids [k], confidence [k])``. + ``verify(tokens [w])`` returns the target's argmax at each of the ``w`` drafted + positions plus one more -- the token that follows a fully accepted block. + ``positions`` are the block's absolute positions, used only to decide whether a + rejection can strand the carry. + ``snapshot`` is called BEFORE drafting; its result is restored on rejection. + """ + saved = snapshot() if snapshot is not None else None + + proposed, confidence = draft() + width = draft_width(confidence, self.confidence_threshold, self.block_size) + proposed = proposed[:width] + + target = verify(proposed) + n_accepted, bonus = accepted_prefix(proposed, target) + + tokens = [int(t) for t in proposed[:n_accepted]] + if bonus >= 0: + tokens.append(bonus) + + rolled_back = False + if n_accepted < width and saved is not None: + if needs_rollback(n_accepted, positions[:width], self.page_size): + saved.restore() + rolled_back = True + return StepResult(tokens, self.block_size, width, n_accepted, rolled_back) + + +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, as ``[n, window]``. + + 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: EVERY row sees ``[last - window + 1, last]``, where ``last`` is the + block's final position, so each drafted query sees the whole block. + """ + end = start_pos + n + w_lo = max(0, start_pos - window + 1) + ar = torch.arange(window, device=device) + if non_causal: + last = end - 1 + lo = max(w_lo, last - window + 1) + cand = lo + ar + return torch.where(cand > last, -1, cand - w_lo).unsqueeze(0).expand(n, window) + 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", + "DSparkDrafter", + "MarkovHead", + "SpeculativeLoop", + "StepResult", + "accepted_prefix", + "draft_width", + "window_cols_for_block", +] diff --git a/python/freetoken/models/deepseek_v4/model.py b/python/freetoken/models/deepseek_v4/model.py index c00d20c..3223001 100644 --- a/python/freetoken/models/deepseek_v4/model.py +++ b/python/freetoken/models/deepseek_v4/model.py @@ -21,14 +21,18 @@ from __future__ import annotations +import os + 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.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 +49,20 @@ get_window_topk_idxs, ) from .moe import Expert, Gate # noqa: F401 +from .parallel import div_tp, tp_info, validate_tp + +logger = init_logger(__name__) 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) @@ -130,23 +137,142 @@ def decode_step(self, x, pos, rows, cmp_stage_cap, input_ids, wctx=None): 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. The checkpoint's + # dspark_target_layer_ids count layers from 1 (40, 41, 42 = the 40th, 41st and + # 42nd block), so the tap fires after 0-based 39, 40 and 41. + 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 + # Which base dspark_target_layer_ids uses is NOT documented, and the + # checkpoint does not settle it. It declares [40, 41, 42] with n_layers=43: + # read 0-based those are the last three layers; read 1-based they are the + # 2nd-to-4th from last. Both land inside the model, so no range check can + # tell them apart, and the only symptom of the wrong choice is a drafter + # whose proposals get rejected. + # + # Measured on this checkpoint at a fixed 800 drafted tokens: + # 1-based (ids-1 -> 39,40,41): 25% accepted + # 0-based (ids -> 40,41,42): 20% accepted + # so 1-based is what this checkpoint means, and is the default. The override + # exists because that is an empirical answer on one checkpoint, not a spec. + base = int(os.environ.get("FREETOKEN_DSPARK_LAYER_BASE", "1")) + if base not in (0, 1): + raise ValueError( + f"FREETOKEN_DSPARK_LAYER_BASE must be 0 or 1, got {base}" + ) + ids = tuple(i - base for i in 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"dspark_target_layer_ids {tuple(args.dspark_target_layer_ids)} read " + f"as {base}-based give {ids}, which fall outside the " + f"{args.n_layers} target layers" + ) + logger.info_rank0( + f"dSpark: tapping target layers {sorted(ids)} " + f"(config {list(args.dspark_target_layer_ids)} read as {base}-based)" + ) + # The concatenated tap from the last forward, [.., dim * len(target_layer_ids)], + # or None when dSpark is off. The drafter reads this and nothing else. + self._last_aux_hidden: torch.Tensor | None = None + # The positions and page-table rows the tap above was computed AT. The + # drafter writes context KV derived from the tap, so it must address the + # positions that produced it -- which are the PREVIOUS forward's, not the + # batch the drafter is about to fill. + self._last_aux_positions: torch.Tensor | None = None + self._last_aux_rows: torch.Tensor | 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 last_aux_hidden(self) -> torch.Tensor | None: + """The tap from the most recent forward, or None if nothing has run yet.""" + return self._last_aux_hidden + + def last_aux_addressing(self) -> tuple[torch.Tensor, torch.Tensor] | None: + """``(positions, table_rows)`` for the rows of :meth:`last_aux_hidden`. + + The tap and its addressing must travel together. A consumer that pairs the tap + with any OTHER batch's positions writes hidden states derived from one set of + tokens into the slots of a different set -- which raises nothing, and shows up + only as a drafter that proposes badly. + """ + if ( + self._last_aux_hidden is None + or self._last_aux_positions is None + or self._last_aux_rows is None + ): + return None + return self._last_aux_positions, self._last_aux_rows + + 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 @@ -160,7 +286,7 @@ def hc_head(self, x): 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 +298,30 @@ 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)) + self._last_aux_hidden = torch.cat(aux, dim=-1) if aux else None + if aux: + self._last_aux_positions = flat_positions + rows = torch.empty_like(flat_positions) + for off, n, ti, _start in segments: + rows[off:off + n] = ti + self._last_aux_rows = rows 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 +337,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 +345,18 @@ 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] + self._last_aux_hidden = torch.cat(aux, dim=-1) if aux else None + if aux: + self._last_aux_positions = pos.reshape(-1) + self._last_aux_rows = rows.reshape(-1) h = self.hc_head(h) h = self.norm(h) - return F.linear(h[:, -1], self.head) + return self.logits(h[:, -1]) class DeepseekV4ForCausalLM(BaseLLMModel): @@ -241,8 +391,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,6 +418,73 @@ 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 catch_up_draft_context(self, batch) -> None: + """Write the draft layers' KV for the positions the target has just committed. + + The drafter never runs the prompt; its sliding window fills in behind the target, + derived from the target's own hidden state at those positions. + + The addressing comes from the TAP, not from ``batch``. This runs before the + target's forward for the current step, so the newest tap belongs to the PREVIOUS + forward and covers that forward's positions. Pairing it with the current batch's + positions -- which is the block about to be drafted -- would store hidden states + derived from one set of tokens into the slots of a different set. That raises + nothing at all; it just makes the drafter propose badly, which reads as a weak + drafter rather than a wiring fault. + + ``batch`` is still taken so the caller need not know what the drafter addresses + by, and so a future batched form has it. + """ + del batch + drafter = self._transformer.drafter + if drafter is None: + return + aux = self._transformer.last_aux_hidden() + addressing = self._transformer.last_aux_addressing() + if aux is None or addressing is None: + return # nothing has run yet; the first block attends over a short window + positions, rows = addressing + flat = aux.view(-1, aux.shape[-1]) + if flat.shape[0] != positions.numel(): + # The tap and its addressing are written together, so this cannot drift -- + # unless a new forward path records one and not the other. + raise RuntimeError( + f"aux tap has {flat.shape[0]} rows but {positions.numel()} positions; " + "every forward that stores the tap must store its addressing too" + ) + backend = get_global_ctx().attn_backend + slots = backend.window_slots_at(rows, positions) + drafter.catch_up_context(flat, positions, slots) + + def draft(self) -> tuple[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 [T], confidence [T])``, or None when there is nothing to + draft from: the drafter needs the target's hidden state at its tap layers, and + that only exists once a forward has produced it. + """ + if self._transformer.drafter is None: + return None + aux = self._transformer.last_aux_hidden() + if aux is None: + return None + self._ensure_bound() + batch = get_global_ctx().batch + md = batch.attn_metadata + return self._transformer.drafter.propose( + aux, + batch.input_ids.long().view(1, -1), + md.segments, + batch.positions.long(), + self._transformer.logits, + ) + def forward(self) -> torch.Tensor: self._ensure_bound() batch = get_global_ctx().batch @@ -273,9 +497,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 diff --git a/python/freetoken/models/deepseek_v4/moe.py b/python/freetoken/models/deepseek_v4/moe.py index 9b28922..fa29154 100644 --- a/python/freetoken/models/deepseek_v4/moe.py +++ b/python/freetoken/models/deepseek_v4/moe.py @@ -7,12 +7,15 @@ 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 from .args import DeepseekV4Args from .layers import Linear +from .parallel import div_tp, tp_size class Gate(nn.Module): @@ -56,13 +59,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 +97,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 +123,25 @@ 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 +167,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 +182,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/rollback.py b/python/freetoken/models/deepseek_v4/rollback.py new file mode 100644 index 0000000..d90d762 --- /dev/null +++ b/python/freetoken/models/deepseek_v4/rollback.py @@ -0,0 +1,104 @@ +"""Undo a rejected speculative block's effect on DSV4's stateful compressor carry. + +Rejecting a speculated token is easy for a plain paged KV cache: stop advancing the +position and the slots are reused. DSV4 is not plain. Its compressors and Lightning +Indexer carry a rolling state per request, and a decode step READS that carry from the +previous token's window page, advances it, and WRITES it back to the current token's +page. Speculating N tokens therefore advances the carry N times. + +When the rejected tokens fall in a LATER window page than the last accepted token, no +undo is needed: the next step reads from the accepted token's page, which was never +touched, and the abandoned pages are freed with their ring blocks. + +The hazard is a rejection INSIDE the page the last accepted token lives in. There, +``prev`` and ``cur`` are the same ring block, so the speculative advance overwrote the +state the next real step must read. Nothing else in the engine can rebuild it: the +carry is a reduction over the tokens seen so far, not a function of the KV alone. + +So snapshot exactly those blocks before speculating, and restore on rejection. The +snapshot is bounded and small -- one ring block per (layer, tier) per request, which +for DSV4-Flash is about 5.5 MiB per row across the whole model -- and it is taken only +for the pages a rollback could not otherwise reconstruct. +""" + +from __future__ import annotations + +import torch + +# The carry rollback and the loop that drives it are imported together by the engine; +# keeping needs_rollback here (beside the snapshot it guards) means the decision and the +# mechanism cannot drift apart. + + +class CarrySnapshot: + """The compressor/indexer carry blocks a speculative block is about to overwrite. + + Take one before the draft advances the carry; ``restore()`` puts the engine back to + the state the last accepted token left behind. Restoring is idempotent, so it is + safe to call on the accept path too rather than branching at the call site. + """ + + __slots__ = ("_backend", "_saved", "_window_slots") + + def __init__(self, backend, layers, window_slots: torch.Tensor): + """``layers`` are the target's attention modules; ``window_slots`` is ``[B]``, + each row's CURRENT window slot -- the page a speculative step would advance in + place. Layers with no compressor contribute nothing.""" + self._backend = backend + self._saved: list[tuple[int, str, int, torch.Tensor]] = [] + for attn in layers: + for tier, module in (("attn", attn.compressor), ("idx", attn.indexer)): + comp = _compressor_of(module, tier) + if comp is None: + continue + block = backend.read_carry_blocks( + attn.layer_id, tier, window_slots, comp.ring_size + ) + # clone: read_carry_blocks may view the live ring, which is exactly the + # memory the speculative step is about to overwrite. + self._saved.append((attn.layer_id, tier, comp.ring_size, block.clone())) + self._window_slots = window_slots + + def restore(self) -> None: + """Put every saved carry block back. Call this when a block is rejected.""" + for layer_id, tier, ring_size, block in self._saved: + self._backend.write_carry_blocks( + layer_id, tier, self._window_slots, ring_size, block + ) + + def __len__(self) -> int: + return len(self._saved) + + @property + def bytes(self) -> int: + return sum(b.numel() * b.element_size() for _, _, _, b in self._saved) + + +def _compressor_of(module, tier: str): + """The Compressor that owns ``tier``'s ring, or None if this layer has none. + + A ratio-0 layer has neither; a ratio-128 layer has a compressor but no indexer; the + dSpark draft layers have neither, which is why they are not passed in at all. + """ + if module is None: + return None + return module if tier == "attn" else getattr(module, "compressor", None) + + +def needs_rollback(accepted_upto: int, block_positions: torch.Tensor, page_size: int) -> bool: + """Would a rejection at ``accepted_upto`` leave a stale carry behind? + + Only when a rejected position shares a window page with the last accepted one. If + the block crossed into a new page at or before the rejection, the abandoned pages + carry their own ring blocks away with them and the surviving page was never touched. + """ + if accepted_upto >= block_positions.numel(): + return False # nothing rejected + last_kept = block_positions[accepted_upto - 1] if accepted_upto else None + first_dropped = block_positions[accepted_upto] + if last_kept is None: + return False # the whole block was rejected; the carry never advanced past it + return int(last_kept) // page_size == int(first_dropped) // page_size + + +__all__ = ["CarrySnapshot", "needs_rollback"] 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..a699623 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -123,6 +123,9 @@ 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 + # MoE layer ids belonging to the dSpark drafter. They are kept off the CPU pool + # and fetched without a cap -- see _fetch_fraction_for. + draft_layer_ids: frozenset = frozenset() def __post_init__(self) -> None: policy_ids = {"lru": 0} @@ -786,9 +789,42 @@ 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: + """The fetch fraction, corrected for how many tokens this step carries. + + ``hybrid_fetch_fraction`` balances PCIe against the CPU pool for a ONE-token + decode: fetch pcie_bw/cpu_bw of the misses and both sides finish together. + + A speculative verify breaks that balance, because the two sides scale + differently with the token count T. Fetching an expert moves the same bytes + whether 1 or T tokens route to it -- the cost is flat in T. The CPU pool pays + per (expert, token) pair, so its side of the step grows with T. Holding the + 1-token fraction fixed therefore leaves the GPU idle while the CPU does T times + the work, which is what makes a wider block cost nearly T times a single step + instead of amortizing. + + Re-balancing moves the split point by exactly T. This is the difference between + speculation being free and speculation being pointless on an offload MoE. + """ + # The drafter's layers fetch everything. They are 3 of 46, they run on the + # critical path of EVERY block, and nothing overlaps them -- the target's + # verify cannot start until the draft is done. Leaving a draft expert to the + # CPU pool puts the slowest path in the engine in series with every block, + # to save a fetch that is a rounding error against the target's 43 layers. + if layer_id in self.draft_layer_ids: + return 1.0 + frac = self.hybrid_fetch_fraction + if frac <= 0.0: + return frac + n_tokens = int(expert_ids.shape[0]) if expert_ids.dim() > 1 else 1 + if n_tokens <= 1: + return frac + return min(1.0, frac * n_tokens) + def materialize_layer(self, layer_id: int) -> None: from freetoken.moe.offload_kernels import materialize_layer 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..e385df2 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 ) @@ -331,42 +333,68 @@ 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() - reply.append( - DetokenizeMsg( - uid=req.uid, - next_token=next_token, - finished=finished, - finish_reason=finish_reason, - matched_stop=matched_stop, - stop_strs=req.sampling_params.stop_strs or None, + finished = False + for pos, tok in enumerate(block): + 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) + 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 = req.input_ids.numel() + reply.append( + DetokenizeMsg( + uid=req.uid, + next_token=tok, + finished=finished, + finish_reason=finish_reason, + matched_stop=matched_stop, + stop_strs=req.sampling_params.stop_strs or None, + ) + ) + if finished: + break # drop the rest of the block: it is past the reply's end + next_token = block[-1] # NOTE: overlap scheduling may make the request freed twice, skip second free if finished and req not in self.finished_reqs: @@ -833,10 +861,52 @@ 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 + k = spec.block_size - 1 + if k < 1: + return + for req in batch.reqs: + # Room for the block AND its bonus token; a request near its budget just + # decodes normally rather than being truncated mid-block. + if req.remain_len <= k + 1: + return + noise = torch.full((k,), spec.noise_token_id, dtype=self.token_pool.dtype) + for req in batch.reqs: + req.append_host(noise) + req.device_len += k + batch.phase = "prefill" + batch.speculative = True + batch.spec_block = k + # 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 +935,42 @@ 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) forward_output = self.engine.forward_batch(batch, sample_args) 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 + confidence_threshold: float + + +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, + confidence_threshold=float(ENV.DSPARK_CONFIDENCE_THRESHOLD.value), + ) + + 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..5398d39 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -928,6 +928,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..8dda7de 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -221,6 +221,31 @@ 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( + "--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/launch.py b/python/freetoken/server/launch.py index a5f8438..c1f8cca 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,14 @@ def _run_scheduler(args: ServerArgs, ack_queue: mp.Queue[str]) -> None: try: scheduler = Scheduler(args) scheduler.sync_all_ranks() + 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 +130,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/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..1f773b2 --- /dev/null +++ b/scripts/freetoken-dsv4.sh @@ -0,0 +1,272 @@ +#!/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}" +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}" +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) + + 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" \ + --expert-load "$EXPERT_LOAD" \ + "${spec[@]}" \ + >> "$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_aux_addressing.py b/tests/dsv4/test_dsv4_aux_addressing.py new file mode 100644 index 0000000..4c4d4b2 --- /dev/null +++ b/tests/dsv4/test_dsv4_aux_addressing.py @@ -0,0 +1,118 @@ +"""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_every_site_that_stores_the_tap_stores_positions(self): + # The two forwards (ragged prefill, batched decode) each write the tap. A third + # that writes only the tensor would leave last_aux_addressing stale, and the + # drafter would silently address the wrong step. + sites = SRC.count("self._last_aux_hidden = torch.cat(aux, dim=-1)") + writes = SRC.count("self._last_aux_positions = ") + assert writes >= sites, ( + f"{sites} sites store the tap but only {writes} store its positions" + ) + + def test_addressing_is_exposed_as_one_unit(self): + from freetoken.models.deepseek_v4.model import Transformer + + assert hasattr(Transformer, "last_aux_addressing"), ( + "positions and rows must be readable together with the tap" + ) + + def test_addressing_returns_none_until_a_forward_has_run(self): + from freetoken.models.deepseek_v4.model import Transformer + + t = Transformer.__new__(Transformer) + t._last_aux_hidden = None + t._last_aux_positions = None + t._last_aux_rows = None + assert t.last_aux_addressing() is None + + +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_addressing_from_the_transformer(self): + assert "last_aux_addressing()" in _fn("catch_up_draft_context") + + 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..01cbf6a --- /dev/null +++ b/tests/dsv4/test_dsv4_aux_layer_ids.py @@ -0,0 +1,98 @@ +"""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 TestTheBaseIsSelectableAndMeasured: + """Which base the ids use is an empirical answer, not a spec. + + The checkpoint declares [40, 41, 42] with n_layers=43. Read 0-based those are the + last three layers; read 1-based they are the 2nd-to-4th from last. Both land inside + the model, so no assertion can decide it -- only a measurement can, and it did: + + 1-based (39, 40, 41): 25% accepted at 800 drafted tokens + 0-based (40, 41, 42): 20% accepted at the same point + + So 1-based is the default, and the override stays because that answer came from one + checkpoint rather than from documentation. + """ + + def test_the_default_is_one_based(self): + src = _model_src() + assert 'FREETOKEN_DSPARK_LAYER_BASE", "1"' in src, ( + "1-based measured better on this checkpoint; it is the default" + ) + + def test_the_base_is_an_override_not_a_hardcode(self): + assert "FREETOKEN_DSPARK_LAYER_BASE" in _model_src() + + def test_only_zero_or_one_is_accepted(self): + assert "base not in (0, 1)" in _model_src() + + def test_the_resolved_layers_are_logged(self): + # The wrong base costs acceptance and raises nothing, so the run must say which + # layers it actually tapped. + src = _model_src() + assert "dSpark: tapping target layers" in src + + @pytest.mark.parametrize("base,expected", [(1, {39, 40, 41}), (0, {40, 41, 42})]) + def test_both_readings_stay_inside_the_model(self, base, expected): + ids = tuple(i - base for i in (40, 41, 42)) + assert set(ids) == expected + assert all(0 <= i < 43 for i in ids), ( + "both readings are in range, which is why only a measurement can choose" + ) + + +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..50c0329 --- /dev/null +++ b/tests/dsv4/test_dsv4_dspark.py @@ -0,0 +1,65 @@ +"""dSpark acceptance and adaptive width. + +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 pytest +import torch + +from freetoken.models.deepseek_v4.dspark import accepted_prefix, draft_width + + +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) + + +class TestDraftWidth: + def test_full_width_when_every_position_is_confident(self): + assert draft_width(torch.tensor([0.9, 0.8, 0.7]), 0.5, 5) == 3 + + def test_cuts_at_the_first_doubtful_position(self): + # Verification costs a full target pass over whatever width it covers, so a + # tail the drafter itself doubts is work bought and thrown away. + assert draft_width(torch.tensor([0.9, 0.8, 0.2, 0.9]), 0.5, 5) == 2 + + def test_always_verifies_at_least_one_position(self): + assert draft_width(torch.tensor([0.1, 0.1]), 0.5, 5) == 1 + + def test_never_exceeds_the_block_size(self): + assert draft_width(torch.ones(9), 0.5, 5) == 5 + + @pytest.mark.parametrize("threshold", [0.0, 0.5, 1.0]) + def test_width_is_within_bounds_for_any_threshold(self, threshold): + w = draft_width(torch.rand(5), threshold, 5) + assert 1 <= w <= 5 diff --git a/tests/dsv4/test_dsv4_dspark_heads.py b/tests/dsv4/test_dsv4_dspark_heads.py new file mode 100644 index 0000000..0739189 --- /dev/null +++ b/tests/dsv4/test_dsv4_dspark_heads.py @@ -0,0 +1,169 @@ +"""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): + # draft_width compares them against a threshold, so anything outside [0, 1] + # would make the adaptive width meaningless. + 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 TestLogitsUnderTensorParallelism: + """The bug the review caught: the drafter must not use the head weight directly.""" + + def test_the_draft_logits_are_full_vocabulary_at_tp4(self): + from freetoken.models.deepseek_v4.dspark import DSparkDrafter + + info_mod._TP_INFO = DistributedInfo(0, 4) + args = _args(dim=512, moe_inter_dim=512, n_heads=8, o_groups=4, vocab_size=64) + with torch.device("meta"): + drafter = DSparkDrafter(args) + + # The target's logits() all-gathers the vocabulary; stand in for it here. + def target_logits(h): + return torch.zeros(h.shape[0], args.vocab_size) + + drafter.markov_head = MarkovHead(args) # real (small) weights, off meta + drafter.confidence_head = ConfidenceHead(args) + logits, conf = drafter.logits( + torch.zeros(3, args.dim), target_logits, torch.tensor([1, 2, 3]) + ) + assert logits.shape == (3, args.vocab_size), ( + "draft logits must span the full vocabulary; a per-rank slice cannot be " + "added to the replicated Markov bias" + ) + assert conf.shape == (3,) + + +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")) diff --git a/tests/dsv4/test_dsv4_dspark_masking.py b/tests/dsv4/test_dsv4_dspark_masking.py new file mode 100644 index 0000000..b683e51 --- /dev/null +++ b/tests/dsv4/test_dsv4_dspark_masking.py @@ -0,0 +1,118 @@ +"""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 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 + 1 + 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 + 1 + 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 + 1 + 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_the_window_for_any_block_size(self, n): + cols = window_cols_for_block(start_pos=900, n=n, window=WIN, non_causal=True) + assert cols.shape == (n, WIN) + + def test_a_single_token_block_matches_the_causal_mask(self): + # With n == 1 there is no future inside the block, so the two must agree -- + # a non-causal mask that differs here is reaching somewhere it should not. + 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 torch.equal(a, b) + + +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 + 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..d62b92c --- /dev/null +++ b/tests/dsv4/test_dsv4_dspark_wiring.py @@ -0,0 +1,139 @@ +"""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. +* ``CarrySnapshot`` was written and tested precisely because a mid-page rejection + corrupts the compressor carry in silence -- and was never invoked, so the one failure + it existed to prevent was live. +* ``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 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", "sampled requests would fall back to argmax bias"), + ("accepted_prefix", "greedy requests would have no acceptance rule"), + ("draft_width", "the confidence head would score nothing"), + ("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", "a rejected block would strand the compressor carry"), + ("_snapshot_carry", "there would be nothing to restore"), + ("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", "carry_snapshot"], + ) + 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.carry_snapshot is None + + +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" + ) diff --git a/tests/dsv4/test_dsv4_hybrid_balance.py b/tests/dsv4/test_dsv4_hybrid_balance.py new file mode 100644 index 0000000..0a2a743 --- /dev/null +++ b/tests/dsv4/test_dsv4_hybrid_balance.py @@ -0,0 +1,105 @@ +"""Where a speculative step's MoE work runs. + +Speculation is normally free because verifying K tokens costs what 1 costs: a GPU decode +is weight-bound and the batch dimension rides along. That is a property of where the +weights LIVE, and it does not survive an offload MoE. Fetching an expert over PCIe moves +the same bytes whether 1 or K tokens route to it, but the CPU pool pays per +(expert, token) pair. Hold the 1-token split fixed and a K-wide block costs nearly K +single steps -- at which point no acceptance rate can pay for it. + +These tests pin the two decisions that keep a block cheap, both of which are invisible +at runtime: a step that carries more tokens fetches a larger share, and the drafter's +own layers never touch the CPU pool. +""" + +from __future__ import annotations + +import pytest +import torch + + +class _Cache: + """The two fields _fetch_fraction_for reads, with the real method bound in.""" + + def __init__(self, fraction=0.25, draft_ids=frozenset()): + from freetoken.moe.offload_cache import OffloadMoeCache + + self.hybrid_fetch_fraction = fraction + self.draft_layer_ids = draft_ids + 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 TestTheSplitFollowsTheTokenCount: + def test_a_single_token_step_is_unchanged(self): + # The benched fraction is defined for exactly this case; changing it here would + # silently re-tune ordinary decode, which is not what any of this is for. + assert _Cache(0.28).frac(0, 1) == pytest.approx(0.28) + + def test_a_wider_step_fetches_proportionally_more(self): + c = _Cache(0.10) + assert c.frac(0, 2) == pytest.approx(0.20) + assert c.frac(0, 6) == pytest.approx(0.60) + + def test_the_fraction_never_exceeds_everything(self): + # 0.28 * 6 = 1.68. A fraction above 1 would ask for more fetches than there are + # misses; the kernel takes it as a count and would over-read the miss list. + assert _Cache(0.28).frac(0, 6) == pytest.approx(1.0) + assert _Cache(0.9).frac(0, 64) == pytest.approx(1.0) + + 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_a_one_dimensional_id_tensor_counts_as_one_token(self): + c = _Cache(0.3) + assert c._fn(c, 0, torch.zeros(6, dtype=torch.int32)) == pytest.approx(0.3) + + +class TestTheDrafterStaysOnTheGpu: + """The draft is serial with every block; the CPU pool is the slowest path there is.""" + + def test_a_draft_layer_fetches_everything(self): + c = _Cache(0.05, draft_ids=frozenset({43, 44, 45})) + assert c.frac(44, 1) == 1.0, "even a 1-token draft must not wait on the CPU" + assert c.frac(44, 6) == 1.0 + + def test_a_target_layer_is_unaffected(self): + c = _Cache(0.05, draft_ids=frozenset({43, 44, 45})) + assert c.frac(42, 1) == pytest.approx(0.05) + + +class TestDraftLayerIdsAreTheTail: + """The ids must match the order the model yields its MoE layers in. + + The model yields the target's layers, then the drafter's. Compute the tail wrongly + and the engine pins three of the TARGET's layers to the GPU while leaving the + drafter on the CPU -- the exact inverse, and nothing would report it. + """ + + @pytest.mark.parametrize( + "n_moe,n_draft,expected", + [(46, 3, {43, 44, 45}), (46, 0, set()), (10, 1, {9}), (4, 4, {0, 1, 2, 3})], + ) + def test_the_tail_is_taken(self, n_moe, n_draft, expected): + ids = frozenset(range(n_moe - n_draft, n_moe)) if n_draft else frozenset() + assert set(ids) == expected + + def test_the_model_yields_the_drafter_last(self): + # If this order ever flips, the id arithmetic above silently targets the wrong + # layers -- so assert the source still yields target-then-drafter. + import inspect + + from freetoken.models.deepseek_v4.model import DeepseekV4ForCausalLM + + src = inspect.getsource(DeepseekV4ForCausalLM._iter_offload_moe_layers) + target_at = src.index("self._transformer.layers") + draft_at = src.index("drafter.layers") + assert target_at < draft_at, ( + "the drafter's MoE layers must be yielded AFTER the target's; the engine " + "identifies them as the tail of the sequence" + ) 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_rejection_sampling.py b/tests/dsv4/test_dsv4_rejection_sampling.py new file mode 100644 index 0000000..865a8d9 --- /dev/null +++ b/tests/dsv4/test_dsv4_rejection_sampling.py @@ -0,0 +1,159 @@ +"""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 + +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" diff --git a/tests/dsv4/test_dsv4_rollback.py b/tests/dsv4/test_dsv4_rollback.py new file mode 100644 index 0000000..d28e19c --- /dev/null +++ b/tests/dsv4/test_dsv4_rollback.py @@ -0,0 +1,134 @@ +"""Undoing a rejected speculative block's effect on the DSV4 compressor carry. + +The carry is a reduction over the tokens seen so far, held per window page. Nothing +else in the engine can rebuild it from the KV, so if a speculative step advances it in +place and the block is then rejected, the state the next real step must read is simply +gone -- and the request continues from a carry that saw tokens it never emitted. That +is silent corruption, not a crash, so the contract is pinned here. + +CPU-only: a fake backend stands in for the paged ring. +""" + +from __future__ import annotations + +import torch + +from freetoken.models.deepseek_v4.rollback import CarrySnapshot, needs_rollback + +P = 128 # DSV4 window page + + +class FakeRing: + """The carry ring, keyed the way the real backend keys it: (layer, tier, slot).""" + + def __init__(self): + self.store: dict[tuple[int, str, int], torch.Tensor] = {} + self.writes = 0 + + def read_carry_blocks(self, layer_id, tier, window_slots, ring_size): + return torch.stack([ + self.store.setdefault( + (layer_id, tier, int(s)), torch.zeros(ring_size, 4) + ) + for s in window_slots + ]) + + def write_carry_blocks(self, layer_id, tier, window_slots, ring_size, blocks): + self.writes += 1 + for i, s in enumerate(window_slots): + self.store[(layer_id, tier, int(s))] = blocks[i].clone() + + +class FakeCompressor: + ring_size = 8 + + +class FakeAttn: + """A layer as the snapshot sees it: an id, and whichever tiers it owns.""" + + def __init__(self, layer_id, *, compressor=True, indexer=True): + self.layer_id = layer_id + self.compressor = FakeCompressor() if compressor else None + self.indexer = ( + type("Idx", (), {"compressor": FakeCompressor()})() if indexer else None + ) + + +def _advance(ring, layers, slots, value): + """Stand-in for a decode step advancing the carry in place.""" + for attn in layers: + for tier in ("attn", "idx"): + for s in slots: + ring.store[(attn.layer_id, tier, int(s))] = torch.full((8, 4), value) + + +class TestCarrySnapshot: + def test_restore_undoes_an_in_place_advance(self): + ring, layers, slots = FakeRing(), [FakeAttn(0), FakeAttn(1)], torch.tensor([3]) + _advance(ring, layers, slots, 1.0) # state the accepted token left + snap = CarrySnapshot(ring, layers, slots) + _advance(ring, layers, slots, 9.0) # speculative step overwrites it + assert ring.store[(0, "attn", 3)][0, 0] == 9.0 + + snap.restore() + for attn in layers: + for tier in ("attn", "idx"): + assert torch.equal( + ring.store[(attn.layer_id, tier, 3)], torch.full((8, 4), 1.0) + ), "the accepted token's carry must come back exactly" + + def test_the_snapshot_is_not_a_view_of_the_live_ring(self): + # If read_carry_blocks handed back a view, the speculative write would corrupt + # the snapshot itself and restore would be a no-op. + ring, layers, slots = FakeRing(), [FakeAttn(0)], torch.tensor([2]) + _advance(ring, layers, slots, 1.0) + snap = CarrySnapshot(ring, layers, slots) + _advance(ring, layers, slots, 7.0) + snap.restore() + assert ring.store[(0, "attn", 2)][0, 0] == 1.0 + + def test_restore_is_idempotent(self): + ring, layers, slots = FakeRing(), [FakeAttn(0)], torch.tensor([1]) + _advance(ring, layers, slots, 4.0) + snap = CarrySnapshot(ring, layers, slots) + _advance(ring, layers, slots, 5.0) + snap.restore() + snap.restore() + assert ring.store[(0, "attn", 1)][0, 0] == 4.0 + + def test_layers_without_a_compressor_are_skipped(self): + ring = FakeRing() + layers = [FakeAttn(0, compressor=False, indexer=False), FakeAttn(1)] + snap = CarrySnapshot(ring, layers, torch.tensor([0])) + assert len(snap) == 2, "only layer 1's two tiers" + + def test_a_ratio_128_layer_saves_its_compressor_but_no_indexer(self): + snap = CarrySnapshot(FakeRing(), [FakeAttn(0, indexer=False)], torch.tensor([0])) + assert len(snap) == 1 + + def test_every_row_of_a_batch_is_restored(self): + ring, layers = FakeRing(), [FakeAttn(0)] + slots = torch.tensor([4, 9]) + _advance(ring, layers, slots, 2.0) + snap = CarrySnapshot(ring, layers, slots) + _advance(ring, layers, slots, 8.0) + snap.restore() + assert ring.store[(0, "attn", 4)][0, 0] == 2.0 + assert ring.store[(0, "attn", 9)][0, 0] == 2.0 + + +class TestNeedsRollback: + def test_needed_when_the_rejection_shares_a_page_with_the_last_accepted_token(self): + assert needs_rollback(2, torch.tensor([10, 11, 12, 13]), P) is True + + def test_not_needed_when_the_block_crossed_into_a_new_page(self): + # The last accepted token is on page 0, the first rejected on page 1: the + # abandoned page carries its own ring block away, and page 0 was never touched. + assert needs_rollback(2, torch.tensor([126, 127, 128, 129]), P) is False + + def test_not_needed_when_nothing_was_rejected(self): + assert needs_rollback(4, torch.tensor([1, 2, 3, 4]), P) is False + + def test_not_needed_when_the_whole_block_was_rejected(self): + # No accepted token means the carry never advanced past the previous step. + assert needs_rollback(0, torch.tensor([5, 6, 7]), P) is False 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_speculative_loop.py b/tests/dsv4/test_dsv4_speculative_loop.py new file mode 100644 index 0000000..5f75c47 --- /dev/null +++ b/tests/dsv4/test_dsv4_speculative_loop.py @@ -0,0 +1,195 @@ +"""The dSpark speculative step, end to end over fakes. + +The loop's whole job is ORDER. Every bug available here is silent: + +* snapshot after drafting -> the carry to restore is already gone, rollback is a no-op, + and the request continues from state that saw tokens it never emitted; +* choose the verify width after verifying -> the confidence head saves nothing and the + adaptive feature is decorative; +* accept more than the agreed prefix -> the model emits text the target would not have, + which is the one thing speculative decoding must never do. + +None of those raise. So the loop is driven here against fakes that record what happened +and in which order. +""" + +from __future__ import annotations + +import pytest +import torch + +from freetoken.models.deepseek_v4.dspark import SpeculativeLoop + +BLOCK = 5 + + +class Recorder: + """Fakes for draft/verify/snapshot that log the order they were called in.""" + + def __init__(self, proposed, target, confidence=None): + self.proposed = torch.tensor(proposed, dtype=torch.long) + self.target = torch.tensor(target, dtype=torch.long) + self.confidence = ( + torch.ones(len(proposed)) if confidence is None + else torch.tensor(confidence, dtype=torch.float) + ) + self.calls: list[str] = [] + self.verified_with: torch.Tensor | None = None + self.restored = 0 + + def snapshot(self): + self.calls.append("snapshot") + return self + + def restore(self): + self.calls.append("restore") + self.restored += 1 + + def draft(self): + self.calls.append("draft") + return self.proposed, self.confidence + + def verify(self, tokens): + self.calls.append("verify") + self.verified_with = tokens.clone() + return self.target[: tokens.numel() + 1] + + +def _positions(start=500, n=BLOCK): + return torch.arange(start, start + n) + + +class TestOrdering: + def test_the_carry_is_snapshotted_before_the_draft_advances_it(self): + r = Recorder([1, 2, 3, 4, 5], [1, 2, 3, 4, 5, 6]) + SpeculativeLoop(BLOCK).step( + draft=r.draft, verify=r.verify, positions=_positions(), snapshot=r.snapshot + ) + assert r.calls.index("snapshot") < r.calls.index("draft"), ( + "drafting advances the carry in place; a snapshot taken afterwards " + "captures the state a rejection is supposed to undo" + ) + + def test_verification_happens_after_drafting(self): + r = Recorder([1, 2], [1, 2, 3]) + SpeculativeLoop(BLOCK).step(draft=r.draft, verify=r.verify, positions=_positions(2)) + assert r.calls == ["draft", "verify"] + + +class TestVerifyWidth: + def test_only_the_confident_prefix_is_verified(self): + # Verification costs a full target pass over whatever width it covers. + r = Recorder([1, 2, 3, 4, 5], [1, 2, 3, 4, 5, 6], confidence=[0.9, 0.9, 0.1, 0.9, 0.9]) + res = SpeculativeLoop(BLOCK, confidence_threshold=0.5).step( + draft=r.draft, verify=r.verify, positions=_positions() + ) + assert r.verified_with.numel() == 2, "the doubted tail must not be paid for" + assert res.verified == 2 + assert res.drafted == BLOCK + + def test_full_width_when_the_drafter_is_confident(self): + r = Recorder([1, 2, 3, 4, 5], [1, 2, 3, 4, 5, 6]) + res = SpeculativeLoop(BLOCK).step( + draft=r.draft, verify=r.verify, positions=_positions() + ) + assert res.verified == BLOCK + + +class TestEmission: + def test_a_fully_accepted_block_emits_every_token_plus_the_bonus(self): + r = Recorder([1, 2, 3], [1, 2, 3, 99]) + res = SpeculativeLoop(3).step(draft=r.draft, verify=r.verify, positions=_positions(500, 3)) + assert res.tokens == [1, 2, 3, 99] + assert res.accepted == 3 + + def test_a_partial_block_emits_the_prefix_then_the_targets_own_token(self): + r = Recorder([1, 2, 3], [1, 77, 3, 4]) + res = SpeculativeLoop(3).step(draft=r.draft, verify=r.verify, positions=_positions(500, 3)) + assert res.tokens == [1, 77], "the rejected position is replaced, not dropped" + assert res.accepted == 1 + + def test_a_fully_rejected_block_still_emits_one_token(self): + r = Recorder([1, 2, 3], [55, 2, 3, 4]) + res = SpeculativeLoop(3).step(draft=r.draft, verify=r.verify, positions=_positions(500, 3)) + assert res.tokens == [55], "a rejected block must never stall the request" + assert res.accepted == 0 + + def test_a_match_after_a_rejection_is_not_emitted(self): + # Position 2 agrees, but it was drafted from a token the target rejected. + r = Recorder([1, 2, 3], [1, 77, 3, 4]) + res = SpeculativeLoop(3).step(draft=r.draft, verify=r.verify, positions=_positions(500, 3)) + assert 3 not in res.tokens + + def test_every_step_emits_at_least_one_token(self): + # The liveness property: no combination of draft and target can produce a step + # that advances the request by nothing. + for target in ([9, 9, 9, 9], [1, 9, 9, 9], [1, 2, 9, 9], [1, 2, 3, 9]): + r = Recorder([1, 2, 3], target) + res = SpeculativeLoop(3).step( + draft=r.draft, verify=r.verify, positions=_positions(500, 3) + ) + assert len(res.tokens) >= 1, f"stalled on target {target}" + + +class TestRollback: + def test_rollback_runs_when_a_rejection_strands_the_carry(self): + r = Recorder([1, 2, 3], [1, 77, 3, 4]) + res = SpeculativeLoop(3).step( + draft=r.draft, verify=r.verify, + positions=torch.tensor([500, 501, 502]), # one page + snapshot=r.snapshot, + ) + assert res.rolled_back is True + assert r.restored == 1 + + def test_no_rollback_when_the_block_was_fully_accepted(self): + r = Recorder([1, 2, 3], [1, 2, 3, 4]) + res = SpeculativeLoop(3).step( + draft=r.draft, verify=r.verify, positions=_positions(500, 3), snapshot=r.snapshot + ) + assert res.rolled_back is False + assert r.restored == 0 + + def test_no_rollback_when_the_rejection_fell_in_a_later_page(self): + # Last accepted on page 0, first rejected on page 1: the abandoned page takes + # its ring block with it and the surviving page was never touched. + r = Recorder([1, 2, 3], [1, 2, 88, 4]) + res = SpeculativeLoop(3).step( + draft=r.draft, verify=r.verify, + positions=torch.tensor([126, 127, 128]), + snapshot=r.snapshot, + ) + assert res.rolled_back is False + assert r.restored == 0 + + def test_the_loop_works_without_a_snapshot_at_all(self): + r = Recorder([1, 2, 3], [1, 77, 3, 4]) + res = SpeculativeLoop(3).step(draft=r.draft, verify=r.verify, positions=_positions(500, 3)) + assert res.rolled_back is False + assert res.tokens == [1, 77] + + +class TestInvariants: + @pytest.mark.parametrize("seed", range(12)) + def test_random_blocks_never_break_the_contract(self, seed): + """Over random agreement patterns: emitted tokens must be exactly the agreed + prefix plus the target's own next token, and the step must always advance.""" + g = torch.Generator().manual_seed(seed) + proposed = torch.randint(0, 20, (BLOCK,), generator=g) + target = torch.randint(0, 20, (BLOCK + 1,), generator=g) + r = Recorder(proposed.tolist(), target.tolist()) + res = SpeculativeLoop(BLOCK).step( + draft=r.draft, verify=r.verify, positions=_positions() + ) + assert 1 <= len(res.tokens) <= BLOCK + 1 + assert res.accepted <= res.verified <= res.drafted + # Everything but the last emitted token must be an accepted proposal. + assert res.tokens[: res.accepted] == proposed[: res.accepted].tolist() + # And the emitted sequence must match what the target itself would produce. + assert res.tokens == target[: len(res.tokens)].tolist(), ( + "speculation emitted a sequence the target would not have" + ) + + def test_block_size_must_be_positive(self): + with pytest.raises(ValueError, match="block_size"): + SpeculativeLoop(0) 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/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]