From 7f0d8fc80ce129772d136243e7b1755493469e0d Mon Sep 17 00:00:00 2001 From: Carlos Alvarado Date: Sat, 22 Aug 2026 21:10:38 -0400 Subject: [PATCH 1/3] feat(dsv4): add tensor-parallel DeepSeek-V4 runtime Shard MLA heads, vocabulary, shared experts, and routed FP4 expert banks across TP ranks. Give every shard independent storage, size collectives for skewed expert loading, partition CPU workers, and avoid redundant checkpoint reads. This is the foundation that makes DeepSeek-V4-Flash fit and start reliably on 4x RTX A4000. --- docs/cli.md | 2 + docs/models.md | 21 ++ python/freetoken/engine/config.py | 7 +- python/freetoken/engine/engine.py | 51 ++++- python/freetoken/models/deepseek_v4/args.py | 41 ++++ .../freetoken/models/deepseek_v4/attention.py | 23 +- python/freetoken/models/deepseek_v4/model.py | 56 ++++- python/freetoken/models/deepseek_v4/moe.py | 34 ++- .../freetoken/models/deepseek_v4/parallel.py | 114 ++++++++++ python/freetoken/models/deepseek_v4/weight.py | 207 ++++++++++++------ python/freetoken/moe/cpu_executor.py | 57 ++++- python/freetoken/server/args.py | 12 + tests/dsv4/test_dsv4_tensor_parallel.py | 150 +++++++++++++ 13 files changed, 680 insertions(+), 95 deletions(-) create mode 100644 python/freetoken/models/deepseek_v4/parallel.py create mode 100644 tests/dsv4/test_dsv4_tensor_parallel.py 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/engine/config.py b/python/freetoken/engine/config.py index 543012f..0bb4582 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -72,7 +72,12 @@ 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 use_pynccl: bool = True max_seq_len_override: int | None = None diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index ff3c985..4a021fa 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -245,8 +245,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 +299,38 @@ class ForwardOutput(NamedTuple): copy_done_event: torch.cuda.Event +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 + total = os.cpu_count() or tp_size + per_rank = max(1, total // tp_size) + if per_rank < torch.get_num_threads(): + torch.set_num_threads(per_rank) + logger.info_rank0( + f"torch intra-op threads: {total} -> {per_rank} per rank " + f"({tp_size} ranks share this host)" + ) + + 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) + _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 +359,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 — diff --git a/python/freetoken/models/deepseek_v4/args.py b/python/freetoken/models/deepseek_v4/args.py index b75d280..6e5239f 100644 --- a/python/freetoken/models/deepseek_v4/args.py +++ b/python/freetoken/models/deepseek_v4/args.py @@ -69,15 +69,56 @@ 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 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.""" diff --git a/python/freetoken/models/deepseek_v4/attention.py b/python/freetoken/models/deepseek_v4/attention.py index 357585e..b266b4e 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): @@ -30,17 +32,23 @@ def __init__(self, layer_id: int, args: DeepseekV4Args): 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.eps = args.norm_eps + # 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 +56,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 +108,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, diff --git a/python/freetoken/models/deepseek_v4/model.py b/python/freetoken/models/deepseek_v4/model.py index c00d20c..df75e17 100644 --- a/python/freetoken/models/deepseek_v4/model.py +++ b/python/freetoken/models/deepseek_v4/model.py @@ -26,6 +26,7 @@ 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 @@ -45,6 +46,7 @@ get_window_topk_idxs, ) from .moe import Expert, Gate # noqa: F401 +from .parallel import div_tp, tp_info, validate_tp class Block(nn.Module): @@ -130,15 +132,30 @@ 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) @@ -148,6 +165,33 @@ def bind(self, pool, device: torch.device) -> None: for layer in self.layers: layer.attn.bind(pool, device) + def embed_tokens(self, input_ids: torch.Tensor) -> torch.Tensor: + """Vocabulary-parallel lookup. Rows outside this rank's block contribute zero, so + the all-reduce sums exactly one real row per token. Index arithmetic only -- no + host sync, so a captured decode graph replays it.""" + if self._comm is None: + return self.embed(input_ids) + local = input_ids - self.vocab_start + outside = (local < 0) | (local >= self.vocab_local) + y = self.embed(local.masked_fill(outside, 0)) + y = y.masked_fill(outside.unsqueeze(-1), 0) + return self._comm.all_reduce(y) + + def logits(self, h: torch.Tensor) -> torch.Tensor: + """Head projection over this rank's vocabulary block, then an all-gather back to + the full [rows, vocab]. all_gather concatenates on dim 0, so the gathered tensor + is rank-major: one row block per rank, which is the vocabulary order.""" + local = F.linear(h, self.head) # [rows, vocab_local] + if self._comm is None: + return local + rows = local.shape[0] + gathered = self._comm.all_gather(local) # [tp*rows, vocab_local] + if rows == 1: + return gathered.view(1, -1)[:, : self.vocab_size] + gathered = gathered.view(self.tp_size, rows, self.vocab_local) + gathered = gathered.permute(1, 0, 2).reshape(rows, self.tp_size * self.vocab_local) + return gathered[:, : self.vocab_size] + def hc_head(self, x): shape, dtype = x.size(), x.dtype dim = self.args.dim @@ -172,13 +216,13 @@ 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: h = layer.prefill_batched(h, input_ids, segments, flat_positions) h = self.hc_head(h) h = self.norm(h) - return F.linear(h[0, last_indices], self.head) # [B, vocab] + return self.logits(h[0, last_indices]) # [B, vocab] def decode( self, input_ids: torch.Tensor, pos: torch.Tensor, cmp_stage_cap: int @@ -194,7 +238,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 @@ -206,7 +250,7 @@ def decode( h = layer.decode_step(h, pos, rows, cmp_stage_cap, input_ids, wctx) h = self.hc_head(h) h = self.norm(h) - return F.linear(h[:, -1], self.head) + return self.logits(h[:, -1]) class DeepseekV4ForCausalLM(BaseLLMModel): diff --git a/python/freetoken/models/deepseek_v4/moe.py b/python/freetoken/models/deepseek_v4/moe.py index 9b28922..31caeca 100644 --- a/python/freetoken/models/deepseek_v4/moe.py +++ b/python/freetoken/models/deepseek_v4/moe.py @@ -7,12 +7,14 @@ import torch.nn.functional as F from torch import nn +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 +58,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 +96,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, @@ -129,6 +147,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 +162,9 @@ def forward(self, x: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor: routed = self.experts.routed_forward( x, weights.float().contiguous(), indices.to(torch.int32).contiguous() ) - return (routed + shared).view(shape) + out = routed + shared + # Both halves are partial sums over the split intermediate dim; one collective + # completes the layer (see DSV4OffloadMoELayer._maybe_all_reduce). + if self._comm is not None: + out = self._comm.all_reduce(out) + return out.view(shape) diff --git a/python/freetoken/models/deepseek_v4/parallel.py b/python/freetoken/models/deepseek_v4/parallel.py new file mode 100644 index 0000000..4459e0b --- /dev/null +++ b/python/freetoken/models/deepseek_v4/parallel.py @@ -0,0 +1,114 @@ +"""Tensor-parallel helpers for DeepSeek-V4-Flash. + +What DSV4 shards, and why: + +* **MLA query/output heads.** ``wq_b`` is column-parallel over ``n_heads``, + ``wo_a`` is sharded over the ``o_groups`` output groups, and ``wo_b`` is + row-parallel with one all-reduce at the end of the block. ``o_groups`` is the + binding constraint: a group owns ``n_heads // o_groups`` heads, so a rank must + own whole groups. +* **MoE intermediate dimension.** The routed FP4 expert banks and the shared + SwiGLU expert both split ``moe_inter_dim``. This is the memory win that makes + the model fit: the host expert banks divide by the TP size, so 4 ranks hold + the same 143 GB in total that 1 rank holds today. + +What stays REPLICATED, and why: + +* **The latent KV path** (``wkv``, ``kv_norm``, the paged pools). MLA keeps one + latent KV per token that every head reads, so there is nothing to split; the + KV pool cost per rank is unchanged. +* **The compressors and the Lightning Indexer.** They select which blocks the + attention reads. Every rank must select the SAME blocks. Replicating them + guarantees that without a collective on the selection path. +* **The router (``Gate``) and the hyper-connection mixers.** Small, and the + routing decision must agree across ranks. +""" + +from __future__ import annotations + +import torch + +from freetoken.distributed import try_get_tp_info + + +def tp_info() -> tuple[int, int]: + """``(rank, size)`` for this process; ``(0, 1)`` before the engine sets it.""" + info = try_get_tp_info() + return (0, 1) if info is None else (info.rank, info.size) + + +def tp_size() -> int: + return tp_info()[1] + + +def tp_rank() -> int: + return tp_info()[0] + + +def div_tp(total: int, what: str, *, multiple_of: int = 1) -> int: + """Per-rank size of ``total``. Fails loudly when the split is not exact. + + ``multiple_of`` guards the block-quantized weights: a 128x128 FP8 scale block + and a 32-element FP4 scale block cannot be cut in half. + """ + size = tp_size() + if total % size != 0: + raise ValueError( + f"DeepSeek-V4 tensor parallelism: {what}={total} is not divisible by " + f"--tensor-parallel-size {size}" + ) + local = total // size + if local % multiple_of != 0: + raise ValueError( + f"DeepSeek-V4 tensor parallelism: {what}={total} split over {size} ranks " + f"gives {local}, which is not a multiple of {multiple_of} (quantization " + f"block size)" + ) + return local + + +def shard(t: torch.Tensor, dim: int) -> torch.Tensor: + """This rank's slice of ``t`` along ``dim``, in its OWN storage. + + The copy is the point. ``narrow`` returns a view, and on dim 0 that view is already + contiguous, so ``.contiguous()`` hands it straight back -- keeping the whole parent + tensor alive behind a 1/N-sized view. Every rank then pays for the full weight it + just sharded: measured at 5.1 GiB per GPU on DeepSeek-V4-Flash at TP=4, silently + charged to the model and taken out of the KV and expert-cache budget. + """ + rank, size = tp_info() + if size == 1: + return t + total = t.shape[dim] + if total % size != 0: + raise ValueError( + f"cannot shard a tensor of shape {tuple(t.shape)} on dim {dim} over {size} ranks" + ) + step = total // size + return t.narrow(dim, rank * step, step).clone(memory_format=torch.contiguous_format) + + +def validate_tp(args) -> None: + """Check every DSV4 split up front, so a bad ``--tensor-parallel-size`` fails + at config time with one clear message instead of at the first bad reshape.""" + size = tp_size() + if size == 1: + return + div_tp(args.o_groups, "o_groups") + div_tp(args.n_heads, "n_heads") + if args.n_heads % args.o_groups != 0: + raise ValueError( + f"DeepSeek-V4 expects n_heads ({args.n_heads}) to be a multiple of " + f"o_groups ({args.o_groups})" + ) + # wq_b output rows and the FP8 128x128 scale grid. + div_tp(args.n_heads * args.head_dim, "n_heads*head_dim", multiple_of=128) + # wo_b input columns. + div_tp(args.o_groups * args.o_lora_rank, "o_groups*o_lora_rank", multiple_of=128) + # Shared expert (FP8 128-blocks) and routed FP4 experts (32-element scale blocks). + div_tp(args.moe_inter_dim, "moe_inter_dim", multiple_of=128) + # Vocabulary-parallel embedding + head. + div_tp(args.vocab_size, "vocab_size") + + +__all__ = ["div_tp", "shard", "tp_info", "tp_rank", "tp_size", "validate_tp"] diff --git a/python/freetoken/models/deepseek_v4/weight.py b/python/freetoken/models/deepseek_v4/weight.py index f5a3fd1..eb8c9c6 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,8 +165,10 @@ 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", @@ -171,6 +187,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 +236,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 +289,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 +303,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 _place_dsfp4(banks: dict, name: str, t: torch.Tensor, I: int) -> int: +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, 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 +376,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..a193663 100644 --- a/python/freetoken/moe/cpu_executor.py +++ b/python/freetoken/moe/cpu_executor.py @@ -88,8 +88,51 @@ 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 + per = len(cores) // info.size + if per < 1: # fewer cores than ranks: let the OS schedule them + return cores + return cores[info.rank * per: (info.rank + 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 +155,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 +170,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)] diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index b2857b7..0948a9e 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -221,6 +221,18 @@ def _infer_reasoning_parser(model_path: str) -> str | None: help="The tensor parallelism size.", ) + 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/tests/dsv4/test_dsv4_tensor_parallel.py b/tests/dsv4/test_dsv4_tensor_parallel.py new file mode 100644 index 0000000..dae1237 --- /dev/null +++ b/tests/dsv4/test_dsv4_tensor_parallel.py @@ -0,0 +1,150 @@ +"""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) From 423edd11d7ebffb5f7a8a61dd7c226b0d5e93509 Mon Sep 17 00:00:00 2001 From: Carlos Alvarado Date: Sat, 22 Aug 2026 21:10:38 -0400 Subject: [PATCH 2/3] feat(dsv4): implement exact DSpark speculative decoding Load the checkpoint DSpark blocks and heads, capture target features, maintain draft KV, run non-causal block drafting, verify candidate prefixes, and integrate multi-token request advancement. Use standard p/q rejection sampling for temperature-1 requests and preserve DSV4 compressor state across rejected tails. --- python/freetoken/attention/dsv4_compress.py | 13 +- python/freetoken/core.py | 30 + python/freetoken/engine/config.py | 3 + python/freetoken/engine/engine.py | 294 +++++++++- python/freetoken/env.py | 5 + python/freetoken/kvcache/dsv4_cost_model.py | 10 +- python/freetoken/kvcache/dsv4_paged_pool.py | 10 +- python/freetoken/models/config.py | 10 +- python/freetoken/models/deepseek_v4/args.py | 33 ++ .../freetoken/models/deepseek_v4/attention.py | 71 ++- .../freetoken/models/deepseek_v4/compress.py | 4 +- python/freetoken/models/deepseek_v4/config.py | 9 +- python/freetoken/models/deepseek_v4/dspark.py | 553 ++++++++++++++++++ python/freetoken/models/deepseek_v4/model.py | 131 ++++- .../freetoken/models/deepseek_v4/rollback.py | 104 ++++ python/freetoken/models/deepseek_v4/weight.py | 61 ++ python/freetoken/moe/cpu_executor.py | 28 +- python/freetoken/scheduler/scheduler.py | 161 ++++- python/freetoken/scheduler/status.py | 35 +- python/freetoken/server/args.py | 13 + python/freetoken/utils/numa.py | 182 ++++++ scripts/freetoken-dsv4.sh | 160 +++++ tests/dsv4/test_dsv4_dspark.py | 65 ++ tests/dsv4/test_dsv4_dspark_heads.py | 169 ++++++ tests/dsv4/test_dsv4_dspark_masking.py | 118 ++++ tests/dsv4/test_dsv4_dspark_wiring.py | 139 +++++ tests/dsv4/test_dsv4_multi_token_advance.py | 70 +++ tests/dsv4/test_dsv4_rejection_sampling.py | 159 +++++ tests/dsv4/test_dsv4_rollback.py | 134 +++++ tests/dsv4/test_dsv4_speculative_loop.py | 195 ++++++ tests/dsv4/test_dsv4_tensor_parallel.py | 92 +++ tests/engine/test_numa_placement.py | 211 +++++++ 32 files changed, 3198 insertions(+), 74 deletions(-) create mode 100644 python/freetoken/models/deepseek_v4/dspark.py create mode 100644 python/freetoken/models/deepseek_v4/rollback.py create mode 100644 python/freetoken/utils/numa.py create mode 100644 scripts/freetoken-dsv4.sh create mode 100644 tests/dsv4/test_dsv4_dspark.py create mode 100644 tests/dsv4/test_dsv4_dspark_heads.py create mode 100644 tests/dsv4/test_dsv4_dspark_masking.py create mode 100644 tests/dsv4/test_dsv4_dspark_wiring.py create mode 100644 tests/dsv4/test_dsv4_multi_token_advance.py create mode 100644 tests/dsv4/test_dsv4_rejection_sampling.py create mode 100644 tests/dsv4/test_dsv4_rollback.py create mode 100644 tests/dsv4/test_dsv4_speculative_loop.py create mode 100644 tests/engine/test_numa_placement.py 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/core.py b/python/freetoken/core.py index ef0a539..2206b8e 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,20 @@ 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) # 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 0bb4582..903ec41 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -79,6 +79,9 @@ class EngineConfig: # 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 4a021fa..b11f928 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 @@ -299,6 +309,52 @@ 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. @@ -314,13 +370,24 @@ def _share_cpu_threads_across_ranks(tp_size: int) -> None: """ if tp_size <= 1 or os.environ.get("OMP_NUM_THREADS"): return - total = os.cpu_count() or tp_size - per_rank = max(1, total // tp_size) + # 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) - logger.info_rank0( - f"torch intra-op threads: {total} -> {per_rank} per rank " - f"({tp_size} ranks share this host)" + 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" ) @@ -330,6 +397,9 @@ def __init__(self, config: EngineConfig): 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}") @@ -380,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"): @@ -909,6 +984,185 @@ 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 + 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 + ) + 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)) + 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): @@ -921,6 +1175,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() @@ -1038,6 +1295,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 6e5239f..0b15eca 100644 --- a/python/freetoken/models/deepseek_v4/args.py +++ b/python/freetoken/models/deepseek_v4/args.py @@ -109,6 +109,18 @@ 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. @@ -135,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``. @@ -145,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 b266b4e..c3790f0 100644 --- a/python/freetoken/models/deepseek_v4/attention.py +++ b/python/freetoken/models/deepseek_v4/attention.py @@ -28,7 +28,10 @@ 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 @@ -37,8 +40,17 @@ def __init__(self, layer_id: int, args: DeepseekV4Args): self.head_dim = args.head_dim self.rope_head_dim = args.rope_head_dim 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. @@ -134,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 @@ -143,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: @@ -163,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) @@ -172,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..bba90a6 --- /dev/null +++ b/python/freetoken/models/deepseek_v4/dspark.py @@ -0,0 +1,553 @@ +"""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 + + from .ops import apply_rotary_emb + + 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)) + apply_rotary_emb(kv[..., -rd:], 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 df75e17..e8cd727 100644 --- a/python/freetoken/models/deepseek_v4/model.py +++ b/python/freetoken/models/deepseek_v4/model.py @@ -52,12 +52,12 @@ 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) @@ -160,10 +160,42 @@ def __init__(self, args: DeepseekV4Args): 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 + self._aux_layer_ids = frozenset(i - 1 for i in args.dspark_target_layer_ids) + bad = [i for i in self._aux_layer_ids if not 0 <= i < args.n_layers] + if bad: + raise ValueError( + f"dspark_target_layer_ids {args.dspark_target_layer_ids} are 1-based " + f"and must land inside the {args.n_layers} target layers" + ) + # 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 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 embed_tokens(self, input_ids: torch.Tensor) -> torch.Tensor: """Vocabulary-parallel lookup. Rows outside this rank's block contribute zero, so @@ -204,7 +236,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 @@ -218,11 +250,22 @@ def prefill_batched( # final token -> its next-token logits row. 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 h = self.hc_head(h) h = self.norm(h) - return self.logits(h[0, last_indices]) # [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 @@ -246,8 +289,12 @@ 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 h = self.hc_head(h) h = self.norm(h) return self.logits(h[:, -1]) @@ -285,8 +332,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 @@ -305,6 +359,61 @@ 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 + one step at a time, derived from the target's own hidden state. Skipped when + there is no tap yet -- the first draft then attends over a short window, which + costs acceptance on the opening block and nothing after. + """ + drafter = self._transformer.drafter + if drafter is None: + return + aux = self._transformer.last_aux_hidden() + if aux is None: + return + md = batch.attn_metadata + if not md.segments: + return + backend = get_global_ctx().attn_backend + for off, n, ti, start in md.segments: + positions = torch.arange(start, start + n, device=aux.device) + slots = backend.window_slots_of(ti, start, start + n) + flat = aux.view(-1, aux.shape[-1]) + if flat.shape[0] < off + n: + return # the tap covers a different span; skip rather than misalign + drafter.catch_up_context(flat[off:off + n], 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 @@ -317,9 +426,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/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 eb8c9c6..aecd389 100644 --- a/python/freetoken/models/deepseek_v4/weight.py +++ b/python/freetoken/models/deepseek_v4/weight.py @@ -175,10 +175,71 @@ def linear(prefix: str, split: int | None = None): "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. # -------------------------------------------------------------------------------------- diff --git a/python/freetoken/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py index a193663..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 @@ -105,10 +106,20 @@ def _tp_core_share(cores: list[int]) -> list[int]: info = try_get_tp_info() if info is None or info.size <= 1: return cores - per = len(cores) // info.size + # ``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[info.rank * per: (info.rank + 1) * per] + return cores[index * per: (index + 1) * per] def _smt_siblings_of(cores: list[int]) -> list[int]: @@ -351,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/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index 3554116..d9a5c59 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,49 @@ 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 + def _report_prompt_admissions(self, batch: Batch) -> None: """Publish first-prefill accounting only after batch preparation succeeded. @@ -865,12 +932,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/args.py b/python/freetoken/server/args.py index 0948a9e..8dda7de 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -221,6 +221,19 @@ 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, 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..8a5c1bc --- /dev/null +++ b/scripts/freetoken-dsv4.sh @@ -0,0 +1,160 @@ +#!/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 +# ./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=1 builds and loads the checkpoint's dSpark drafter (the mtp.* stack). +# It costs 9.5 GiB of host expert banks and a slice of the GPU slot cache. Leave it off +# until the draft/verify loop is wired -- today it would load the drafter and never call it. +SPECULATIVE_DSPARK="${SPECULATIVE_DSPARK:-0}" +LOG="${LOG:-/tmp/freetoken-dsv4.log}" +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" +} + +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 + + 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 pid; pid=$(running_pid) || pid="" + if [ -z "$pid" ]; then + pkill -TERM -f "ft serve --model $MODEL" 2>/dev/null + else + kill -TERM "$pid" 2>/dev/null + fi + echo -n "stopping" + for _ in $(seq 1 30); do + pgrep -f "ft serve --model $MODEL" >/dev/null 2>&1 || break + sleep 2; printf '.' + done + echo + rm -f "$PIDFILE" + if pgrep -f "ft serve --model $MODEL" >/dev/null 2>&1; then + echo "warning: something is still running. Give it longer before forcing it --" >&2 + echo "a SIGKILL here leaks the pinned host banks." >&2 + return 1 + fi + echo "stopped" +} + +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 ;; + restart) cmd_stop; cmd_start ;; + status) cmd_status ;; + logs) cmd_logs ;; + test) cmd_test ;; + *) echo "usage: $0 {start|stop|restart|status|logs|test}"; exit 2 ;; +esac 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_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_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 index dae1237..1f38b2a 100644 --- a/tests/dsv4/test_dsv4_tensor_parallel.py +++ b/tests/dsv4/test_dsv4_tensor_parallel.py @@ -148,3 +148,95 @@ def test_a_split_that_does_not_divide_o_groups_fails_loudly(args): _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] From d42be9b32144146da6bfcc6fcf165ac8b8937184 Mon Sep 17 00:00:00 2001 From: Carlos Alvarado Date: Sat, 22 Aug 2026 21:10:38 -0400 Subject: [PATCH 3/3] fix(dsv4): make hybrid DSpark serving operational Correct target-feature addressing, per-row rotary context writes, target-layer taps, GPU affinity and NUMA placement, speculative page release, and drafter residency. Route short verification blocks through FreeToken hybrid CPU/GPU MoE execution, and harden the TP service lifecycle and startup diagnostics. --- assets/ascii-art.txt | 18 ++ python/freetoken/attention/dsv4_sparse.py | 12 + python/freetoken/core.py | 3 + python/freetoken/engine/engine.py | 69 ++++- python/freetoken/models/deepseek_v4/dspark.py | 13 +- python/freetoken/models/deepseek_v4/model.py | 111 +++++-- python/freetoken/models/deepseek_v4/moe.py | 20 ++ python/freetoken/moe/offload_cache.py | 38 ++- python/freetoken/scheduler/cache.py | 35 +++ python/freetoken/scheduler/scheduler.py | 3 + python/freetoken/server/api_server.py | 17 ++ python/freetoken/server/launch.py | 30 +- python/freetoken/utils/banner.py | 282 ++++++++++++++++++ scripts/freetoken-dsv4.sh | 144 ++++++++- tests/dsv4/test_dsv4_aux_addressing.py | 118 ++++++++ tests/dsv4/test_dsv4_aux_layer_ids.py | 98 ++++++ tests/dsv4/test_dsv4_hybrid_balance.py | 105 +++++++ tests/dsv4/test_dsv4_rotary_contract.py | 79 +++++ tests/dsv4/test_dsv4_slot_release.py | 154 ++++++++++ 19 files changed, 1306 insertions(+), 43 deletions(-) create mode 100644 assets/ascii-art.txt create mode 100644 python/freetoken/utils/banner.py create mode 100644 tests/dsv4/test_dsv4_aux_addressing.py create mode 100644 tests/dsv4/test_dsv4_aux_layer_ids.py create mode 100644 tests/dsv4/test_dsv4_hybrid_balance.py create mode 100644 tests/dsv4/test_dsv4_rotary_contract.py create mode 100644 tests/dsv4/test_dsv4_slot_release.py 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/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 2206b8e..cfb7366 100644 --- a/python/freetoken/core.py +++ b/python/freetoken/core.py @@ -163,6 +163,9 @@ class Batch: 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/engine.py b/python/freetoken/engine/engine.py index b11f928..0ac380b 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -635,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: @@ -699,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 @@ -760,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, @@ -1031,6 +1080,7 @@ def _finish_speculative( 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 @@ -1041,6 +1091,14 @@ def _finish_speculative( 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] @@ -1059,6 +1117,13 @@ def _finish_speculative( 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)]) diff --git a/python/freetoken/models/deepseek_v4/dspark.py b/python/freetoken/models/deepseek_v4/dspark.py index bba90a6..642d6d9 100644 --- a/python/freetoken/models/deepseek_v4/dspark.py +++ b/python/freetoken/models/deepseek_v4/dspark.py @@ -194,7 +194,12 @@ def store_context_kv( """ from freetoken.kernel.triton.dsv4.fp8_linear import act_quant_fp8_inplace - from .ops import apply_rotary_emb + # 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( @@ -206,7 +211,11 @@ def store_context_kv( attn = layer.attn freqs = attn.freqs_cis.index_select(0, positions) kv = attn.kv_norm(attn.wkv(main_x)) - apply_rotary_emb(kv[..., -rd:], freqs) + # [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) diff --git a/python/freetoken/models/deepseek_v4/model.py b/python/freetoken/models/deepseek_v4/model.py index e8cd727..3223001 100644 --- a/python/freetoken/models/deepseek_v4/model.py +++ b/python/freetoken/models/deepseek_v4/model.py @@ -21,6 +21,8 @@ from __future__ import annotations +import os + import torch import torch.nn.functional as F from torch import nn @@ -30,6 +32,7 @@ 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 @@ -48,6 +51,8 @@ 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).""" @@ -176,16 +181,45 @@ def __init__(self, args: DeepseekV4Args): # 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 - self._aux_layer_ids = frozenset(i - 1 for i in args.dspark_target_layer_ids) - bad = [i for i in self._aux_layer_ids if not 0 <= i < args.n_layers] + # 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 {args.dspark_target_layer_ids} are 1-based " - f"and must land inside the {args.n_layers} target layers" + 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: @@ -197,6 +231,22 @@ 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 @@ -258,6 +308,12 @@ def prefill_batched( # 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) # Normally only each request's final token needs logits. A speculative VERIFY @@ -295,6 +351,9 @@ def decode( 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 self.logits(h[:, -1]) @@ -362,28 +421,40 @@ def load_state_dict(self, state_dict, *, prefix: str = "", _internal: bool = Fal 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 - one step at a time, derived from the target's own hidden state. Skipped when - there is no tap yet -- the first draft then attends over a short window, which - costs acceptance on the opening block and nothing after. + 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() - if aux is None: - return - md = batch.attn_metadata - if not md.segments: - return + 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 - for off, n, ti, start in md.segments: - positions = torch.arange(start, start + n, device=aux.device) - slots = backend.window_slots_of(ti, start, start + n) - flat = aux.view(-1, aux.shape[-1]) - if flat.shape[0] < off + n: - return # the tap covers a different span; skip rather than misalign - drafter.catch_up_context(flat[off:off + n], positions, slots) + 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. diff --git a/python/freetoken/models/deepseek_v4/moe.py b/python/freetoken/models/deepseek_v4/moe.py index 31caeca..fa29154 100644 --- a/python/freetoken/models/deepseek_v4/moe.py +++ b/python/freetoken/models/deepseek_v4/moe.py @@ -7,6 +7,7 @@ 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 @@ -122,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: 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 d9a5c59..e385df2 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -903,6 +903,9 @@ def _maybe_make_speculative(self, batch: Batch) -> None: 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. 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/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/scripts/freetoken-dsv4.sh b/scripts/freetoken-dsv4.sh index 8a5c1bc..1f773b2 100644 --- a/scripts/freetoken-dsv4.sh +++ b/scripts/freetoken-dsv4.sh @@ -2,7 +2,8 @@ # 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 +# ./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 @@ -20,11 +21,17 @@ 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=1 builds and loads the checkpoint's dSpark drafter (the mtp.* stack). -# It costs 9.5 GiB of host expert banks and a slice of the GPU slot cache. Leave it off -# until the draft/verify loop is wired -- today it would load the drafter and never call it. -SPECULATIVE_DSPARK="${SPECULATIVE_DSPARK:-0}" +# 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 @@ -46,6 +53,28 @@ running_pid() { 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" @@ -56,6 +85,21 @@ cmd_start() { 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) @@ -102,27 +146,94 @@ 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 pid; pid=$(running_pid) || pid="" - if [ -z "$pid" ]; then - pkill -TERM -f "ft serve --model $MODEL" 2>/dev/null - else - kill -TERM "$pid" 2>/dev/null - fi + local pids; pids=$(service_pids) + [ -n "$pids" ] && kill -TERM $pids 2>/dev/null echo -n "stopping" for _ in $(seq 1 30); do - pgrep -f "ft serve --model $MODEL" >/dev/null 2>&1 || break + [ -z "$(service_pids)" ] && break sleep 2; printf '.' done echo rm -f "$PIDFILE" - if pgrep -f "ft serve --model $MODEL" >/dev/null 2>&1; then + 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." >&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" @@ -151,10 +262,11 @@ cmd_test() { case "${1:-}" in start) cmd_start ;; - stop) cmd_stop ;; + 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|restart|status|logs|test}"; exit 2 ;; + *) 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_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_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