From 7f0d8fc80ce129772d136243e7b1755493469e0d Mon Sep 17 00:00:00 2001 From: Carlos Alvarado Date: Sat, 22 Aug 2026 21:10:38 -0400 Subject: [PATCH 01/13] 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 02/13] 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 03/13] 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 From 13d4560901d13c0a78d3a30fd8aa15a7801dcdc6 Mon Sep 17 00:00:00 2001 From: Carlos Alvarado Date: Sat, 22 Aug 2026 21:10:38 -0400 Subject: [PATCH 04/13] perf(dsv4): add TP verification and adaptive DSpark scheduling Fix TP>1 Lightning Indexer row addressing, capture all legal verification widths, share compressor carry storage across graphs, and select widths from cumulative survival probabilities plus the measured draft/verify cost curve. Keep sampled verification exact and device-resident, preserve request and cache state through rejection, and cover the production path with focused DSV4 tests. --- python/freetoken/attention/dsv4_indexer.py | 4 +- python/freetoken/attention/dsv4_sparse.py | 17 +- python/freetoken/core.py | 14 +- python/freetoken/engine/engine.py | 503 +++++++++++++----- python/freetoken/engine/graph.py | 282 +++++++++- python/freetoken/env.py | 5 - .../freetoken/kernel/triton/dsv4/indexer.py | 15 +- .../freetoken/models/deepseek_v4/attention.py | 114 +++- .../freetoken/models/deepseek_v4/compress.py | 115 +++- python/freetoken/models/deepseek_v4/dspark.py | 487 ++++++++++++----- python/freetoken/models/deepseek_v4/model.py | 257 +++++---- python/freetoken/models/deepseek_v4/moe.py | 6 +- .../freetoken/models/deepseek_v4/rollback.py | 104 ---- python/freetoken/moe/offload_cache.py | 40 +- python/freetoken/scheduler/scheduler.py | 107 +++- python/freetoken/server/api_server.py | 24 + python/freetoken/server/generation.py | 28 +- python/freetoken/server/openai_api.py | 9 +- scripts/freetoken-dsv4.sh | 13 + tests/dsv4/test_dsv4_adaptive_verification.py | 124 +++++ tests/dsv4/test_dsv4_aux_addressing.py | 48 +- tests/dsv4/test_dsv4_aux_layer_ids.py | 48 +- tests/dsv4/test_dsv4_dspark.py | 26 +- tests/dsv4/test_dsv4_dspark_heads.py | 110 +++- tests/dsv4/test_dsv4_dspark_masking.py | 33 +- tests/dsv4/test_dsv4_dspark_wiring.py | 41 +- tests/dsv4/test_dsv4_hybrid_balance.py | 104 ++-- tests/dsv4/test_dsv4_proposal_writeback.py | 71 +++ tests/dsv4/test_dsv4_rejection_sampling.py | 31 +- tests/dsv4/test_dsv4_rollback.py | 134 ----- tests/dsv4/test_dsv4_speculative_loop.py | 195 ------- tests/kernels/test_dsv4_indexer_decode.py | 24 +- tests/server/test_client_disconnect_aborts.py | 136 +++++ 33 files changed, 2175 insertions(+), 1094 deletions(-) delete mode 100644 python/freetoken/models/deepseek_v4/rollback.py create mode 100644 tests/dsv4/test_dsv4_adaptive_verification.py create mode 100644 tests/dsv4/test_dsv4_proposal_writeback.py delete mode 100644 tests/dsv4/test_dsv4_rollback.py delete mode 100644 tests/dsv4/test_dsv4_speculative_loop.py create mode 100644 tests/server/test_client_disconnect_aborts.py diff --git a/python/freetoken/attention/dsv4_indexer.py b/python/freetoken/attention/dsv4_indexer.py index aba5f63..f3cf595 100644 --- a/python/freetoken/attention/dsv4_indexer.py +++ b/python/freetoken/attention/dsv4_indexer.py @@ -43,7 +43,7 @@ def indexer_prefill_logits( def indexer_decode_scores( self, q: torch.Tensor, weights: torch.Tensor, valid: torch.Tensor, n_stage: int, - ratio: int, layer_id: int, + ratio: int, layer_id: int, rows: torch.Tensor, ) -> torch.Tensor: """Head-reduced scores ``[B, n_stage]`` for a decode step, gathering each block's key off the decode SNAPSHOT and bounding the work by the live block count read from device @@ -52,7 +52,7 @@ def indexer_decode_scores( return indexer_decode_logits( q, weights, self.compress_pool(layer_id, "idx"), self.snapshot(), - valid, n_stage, ratio, + rows, valid, n_stage, ratio, ) def indexer_select_prefill( diff --git a/python/freetoken/attention/dsv4_sparse.py b/python/freetoken/attention/dsv4_sparse.py index b944f5e..606e0c5 100644 --- a/python/freetoken/attention/dsv4_sparse.py +++ b/python/freetoken/attention/dsv4_sparse.py @@ -208,6 +208,18 @@ def prepare_for_capture(self, batch: Batch) -> None: def prepare_for_replay(self, batch: Batch) -> None: self._point_to_capture(batch, batch.padded_size, self._table_rows(batch)) + def prepare_for_spec_capture(self, batch: Batch) -> None: + """Point a fixed-width DSpark target capture at the reserved dummy row.""" + bs = batch.padded_size + rows = torch.full( + (bs,), batch.padded_reqs[0].table_idx, dtype=torch.int64, device=self.device + ) + self._point_to_capture(batch, bs, rows) + + def prepare_for_spec_replay(self, batch: Batch) -> None: + """Stage one snapshot row per request for a DSpark target replay.""" + self._point_to_capture(batch, batch.padded_size, self._table_rows(batch)) + # ----- DSV4 addressing vocabulary --------------------------------------------------- def snapshot(self) -> torch.Tensor: """The current decode batch's snapshot (see ``DSV4AttnMetadata.full_snapshot``); @@ -291,11 +303,14 @@ def _point_to_capture(self, batch: Batch, bs: int, rows_ti: torch.Tensor) -> Non """ assert self.capture is not None and bs <= self.max_graph_bs cap = self.capture + prior = getattr(batch, "attn_metadata", None) + segments = getattr(prior, "segments", None) src = self.pool.full_loc_map.index_select(0, rows_ti) w = min(src.shape[1], cap.full_snap.shape[1]) cap.full_snap[:bs, :w].copy_(src[:bs, :w]) batch.attn_metadata = DSV4AttnMetadata( - last_indices=cap.last_indices[:bs], full_snap=cap.full_snap[:bs], + last_indices=cap.last_indices[:bs], segments=segments, + full_snap=cap.full_snap[:bs], window_ar=self._window_ar, ) diff --git a/python/freetoken/core.py b/python/freetoken/core.py index cfb7366..dafeb98 100644 --- a/python/freetoken/core.py +++ b/python/freetoken/core.py @@ -154,6 +154,10 @@ class Batch: # it from a real prefill -- these say so explicitly. speculative: bool = field(default=False, init=False) spec_block: int = field(default=0, init=False) + # GraphRunner sets this only while the target verify graph is capturing/replaying. + # It selects DSV4's fixed-shape, device-addressed verify implementation; the eager + # prefill implementation remains the fallback when no verify graph is available. + spec_verify_decode: bool = field(default=False, init=False) draft_confidence: torch.Tensor | None = field(default=None, init=False) # Per-request accepted tokens from a speculative step: the reply path emits one # token per request, so the rest of each block is read from here. @@ -161,8 +165,14 @@ class Batch: # 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) + # Tokens produced by the paper's sequential Markov stage, [requests * gamma]. + # They stay on GPU through verification; only the accepted prefix is copied back. + draft_tokens: torch.Tensor | None = field(default=None, init=False) + # Per-token compressor partial states produced by a speculative target verify. + # FreeToken's compressor ring is page-addressed, so later tokens in the same page + # overwrite earlier states. This journal is the engine-native equivalent of + # vLLM's save_partial_states: acceptance selects the state at the last valid row. + spec_carry_states: dict | None = field(default=None, init=False) # Returns a rejected block's unused pages/SWA slots. Supplied by the scheduler, # which owns the cache manager and is what inflated device_len to the block width. release_tail: object | None = field(default=None, init=False) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 0ac380b..7afe8b2 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -4,6 +4,7 @@ import math import glob import os +import time from datetime import timedelta from typing import Any, Dict, Iterable, NamedTuple, Tuple @@ -303,12 +304,35 @@ def _materialize_loaded_weight_state_dict( return state_dict +# FREETOKEN_SPEC_DEBUG=1 dumps the first few blocks' token flow. +_SPEC_DEBUG = os.environ.get("FREETOKEN_SPEC_DEBUG", "0") == "1" +# Diagnostic only: synchronize and report the first speculative cycles by stage. +# Disabled by default because the per-cycle synchronization deliberately removes overlap. +_SPEC_TIMING = os.environ.get("FREETOKEN_SPEC_TIMING", "0") == "1" + class ForwardOutput(NamedTuple): next_tokens_gpu: torch.Tensor next_tokens_cpu: torch.Tensor copy_done_event: torch.cuda.Event +def _cpu_moe_max_tokens(config: EngineConfig) -> int: + """Largest row count one CPU-expert task can receive. + + Ordinary decode has one row per request. DSpark target verification has the + anchor plus ``gamma`` proposal rows per request, while the draft itself has only + ``gamma``; therefore ``1 + gamma`` is the required upper bound. + """ + rows_per_req = 1 + dsv4 = getattr(config.model_config, "dsv4_args", None) + if getattr(dsv4, "dspark_enabled", False): + rows_per_req = 1 + max( + 1, int(getattr(dsv4, "dspark_block_size", 1) or 1) + ) + max_reqs = max(config.max_running_req, config.cuda_graph_max_bs or 0, 1) + return max_reqs * rows_per_req + + def _bind_rank_to_numa_node(tp_info) -> None: """Pin this rank to one NUMA node BEFORE it allocates anything. @@ -453,8 +477,9 @@ def __init__(self, config: EngineConfig): # Speculative accounting: acceptance rate is the one number that says whether # speculation is paying for itself, and it cannot be inferred from tokens/s. self._spec_accepted = 0 + self._spec_debug_left = 6 self._spec_drafted = 0 - self.spec_threshold = float(ENV.DSPARK_CONFIDENCE_THRESHOLD.value) + self._spec_timing_left = 12 if is_offload_moe_backend(config.moe_backend): self._init_offload_moe_cache(config) if hasattr(self.model, "prepare_for_runtime"): @@ -542,6 +567,7 @@ def __init__(self, config: EngineConfig): dummy_req=self.dummy_req, moe_offload_cache=self.moe_offload_cache, ) + self._init_dspark_adaptive_verification() if config.attention_backend.split(",")[0] == "triton": # Prefill runs on the first comma part; warm its autotune cache. self._warmup_prefill() @@ -637,42 +663,6 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: # gates the CPU executor build below. 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: @@ -736,7 +726,6 @@ 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 @@ -799,17 +788,11 @@ 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. # - # 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 + # A dSpark target verify carries 1 + block_size ROWS per request: the anchor + # plus every drafted token. max_tokens sizes the C++ pool's # per-task scratch, so a pool built for one row per request would be handed - # 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 - ) + # more rows than it owns and overrun it. + max_tokens = _cpu_moe_max_tokens(config) # gpt-oss mxfp4 carries clamped-swiglu scalars; other formats use the defaults. executor = CpuMoeExecutor( cache, @@ -1032,9 +1015,107 @@ def rebuild_runtime_cache( dummy_req=self.dummy_req, moe_offload_cache=self.moe_offload_cache, ) + self._init_dspark_adaptive_verification() + + def _init_dspark_adaptive_verification(self) -> None: + """Install the paper's measured-cost scheduler when every prefix is captured.""" + self._adaptive_verification = None + curve = list(getattr(self.graph_runner, "spec_verify_cost_curve", ()) or ()) + block_size = int(getattr(self.graph_runner, "spec_block_size", 0) or 0) + if not curve or block_size < 1: + return + if self.config.max_running_req != 1: + logger.warning_rank0( + "DSpark adaptive verification needs the paper's marker-tensor varlen " + "layout for request batches >1; using fixed gamma for this run" + ) + return + + # Every TP rank must select the same graph shape. vLLM broadcasts its + # profiled cost curves from rank 0; do exactly that rather than letting + # small per-GPU timing noise choose different collective graphs. + tp = get_tp_info() + if tp.size > 1: + payload = [curve if tp.is_primary() else None] + torch.distributed.broadcast_object_list( + payload, src=0, group=self.tp_cpu_group + ) + curve = payload[0] + from freetoken.models.deepseek_v4.dspark import DSparkAdaptiveVerification + + self._adaptive_verification = DSparkAdaptiveVerification( + block_size, curve, self.device + ) + + def adapt_speculative_batch(self, batch: Batch) -> None: + """Compact one prepared DSpark block to the paper-selected prefix width. + + Allocation deliberately remains at gamma until acceptance, so abandoned + page/window slots are still returned by ``release_speculative_tail``. Only + the target's input views and metadata shrink before graph replay. + """ + manager = self._adaptive_verification + if manager is None: + return + if len(batch.reqs) != 1 or batch.padded_size != 1: + raise RuntimeError( + "adaptive DSpark compaction currently requires one unpadded request" + ) + confidence = batch.draft_confidence + if confidence is None: + raise RuntimeError("adaptive DSpark received no confidence probabilities") + + max_width = int(batch.spec_block) + if max_width != manager.block_size: + raise RuntimeError( + f"DSpark prepared width {max_width}, expected checkpoint gamma " + f"{manager.block_size}" + ) + req = batch.reqs[0] + width = manager.record_and_choose(confidence, req.uid) + if width == max_width: + return + if not 0 <= width < max_width: + raise RuntimeError(f"adaptive DSpark selected invalid width {width}") + + base = req.input_ids.numel() - max_width + span = width + 1 + req.input_ids = req._ids_buf[: base + width] + batch.input_ids = batch.input_ids[:span] + batch.positions = batch.positions[:span] + if batch.out_loc is not None: + batch.out_loc = batch.out_loc[:span] + if batch.draft_tokens is None or batch.draft_probs is None: + raise RuntimeError("adaptive DSpark cannot compact a missing draft") + batch.draft_tokens = batch.draft_tokens[:width] + batch.draft_probs = batch.draft_probs[:width] + batch.draft_confidence = confidence[:width] + segments = getattr(batch.attn_metadata, "segments", None) + if segments is None or len(segments) != 1: + raise RuntimeError("adaptive DSpark needs one target metadata segment") + _off, _old_n, table_idx, start_pos = segments[0] + batch.attn_metadata.segments = [(0, span, table_idx, start_pos)] + batch.spec_block = width + + def _record_adaptive_draft_cost(self, batch: Batch) -> None: + """Publish rank-0's real drafter time to the shared five-sample profile.""" + manager = self._adaptive_verification + if manager is None or not manager.needs_draft_profile: + return + start = getattr(batch, "_spec_draft_start", None) + end = getattr(batch, "_spec_draft_end", None) + if start is None or end is None: + raise RuntimeError("adaptive DSpark draft timing events are missing") + tp = get_tp_info() + cost = start.elapsed_time(end) if tp.is_primary() else 0.0 + if tp.size > 1: + shared = torch.tensor([cost], dtype=torch.float64) + torch.distributed.broadcast(shared, src=0, group=self.tp_cpu_group) + cost = float(shared.item()) + manager.record_draft_cost(cost) def _finish_speculative( - self, batch: Batch, logits: torch.Tensor, args: BatchSamplingArgs + self, batch: Batch, logits: torch.Tensor, args: BatchSamplingArgs, target_features ) -> ForwardOutput: """Keep the prefix of each block the target agrees with, and collapse the rest. @@ -1050,14 +1131,11 @@ def _finish_speculative( """ from freetoken.models.deepseek_v4.dspark import ( accepted_prefix, - draft_width, - rejection_accept, + rejection_accept_device, 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 @@ -1070,27 +1148,56 @@ def _finish_speculative( "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() + # sampling args are sized per REQUEST. Greedy verification reduces the + # [rows, vocab] logits on the GPU and transfers only one int per row. Copying + # the whole matrix to the CPU was ~3 MiB per six-row cycle and then repeated + # the reduction on the slower side of PCIe. + target_cpu = logits.argmax(dim=-1).to(torch.int32).to("cpu", non_blocking=False) + # The blocking target copy above proves the draft events have completed. + # Record the five startup samples here, after target launch, so timing the + # drafter never introduces a synchronization between draft and verify. + self._record_adaptive_draft_cost(batch) + any_sampled = any(not req.sampling_params.is_greedy for req in batch.reqs) + if any_sampled and batch.draft_probs is None: + raise RuntimeError("sampled DSpark verification is missing draft probabilities") confidence = batch.draft_confidence + proposed_gpu = batch.draft_tokens + if proposed_gpu is None or proposed_gpu.numel() != k * len(batch.reqs): + raise RuntimeError( + "DSpark verify is missing the gamma proposals produced by its sequential stage" + ) + proposed_cpu = proposed_gpu.to("cpu", non_blocking=False).to(torch.int32) emitted: list[torch.Tensor] = [] - any_rejected = False release_tail = getattr(batch, "release_tail", None) + selected_rows: list[int] = [] + accepted_counts: list[int] = [] off = 0 for i, req in enumerate(batch.reqs): + sp = req.sampling_params + greedy = sp.is_greedy span = 1 + k # this request's rows in the flat batch - 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 + # Where the block begins IN THE BUFFER. Not device_len - k: device_len runs + # one ahead of the buffer during decode (see the write-back above), so that + # form silently yielded k-1 proposals -- a short slice, not an error. + start = req.input_ids.numel() - k + proposed = proposed_cpu[i * k:(i + 1) * k] + # A VIEW of _ids_buf, which the accept path overwrites below. Snapshot it + # for the debug line, or the log shows post-mutation values and invents + # agreements that never happened. + proposed_snapshot = proposed.clone() if _SPEC_DEBUG else None + if proposed.numel() != k: + raise RuntimeError( + f"block has {proposed.numel()} proposals, expected {k}; the " + "request's buffer and the batch's block width disagree" ) + # Fixed-width verification is the paper/vLLM fallback when the profiled + # hardware-aware scheduler is disabled. A static per-token threshold is + # deliberately not used: DSpark section 3.2 schedules a GLOBAL token budget + # from cumulative survival probabilities and the measured draft/verify cost + # curves. Until that manager is ported, verify the full gamma without + # claiming that the confidence head is itself a threshold rule. + width = k # A block emits n_acc accepted tokens PLUS the bonus token, so it can carry # up to width+1 past `start`. Nothing upstream clamps that to the request's # output budget: a block that starts with 2 tokens left would write 6, run @@ -1107,15 +1214,22 @@ def _finish_speculative( # 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], + p_req = sampling_probs( + logits[off:off + width + 1], + sp.temperature, + sp.top_p, + sp.top_k, + ) + assert batch.draft_probs is not None + n_acc, bonus = rejection_accept_device( + proposed_gpu[i * k:i * k + width], + batch.draft_probs[i * k:i * k + width], + p_req, ) - if n_acc < width: - any_rejected = True keep = start + n_acc - req.input_ids = req._ids_buf[:keep] + req.input_ids = req._ids_buf[:start] + if n_acc: + req.append_host(proposed[:n_acc].to(req.input_ids.dtype)) req.append_host(torch.tensor([bonus], dtype=req.input_ids.dtype)) # Hand back the pages and SWA slots of the positions this block did not # keep, BEFORE device_len drops past them. allocate_paged sized itself from @@ -1123,22 +1237,43 @@ def _finish_speculative( # 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) + release_tail(req, keep) req.cached_len, req.device_len = keep, keep + 1 emitted.append( torch.cat([proposed[:n_acc], torch.tensor([bonus], dtype=torch.int32)]) ) + if _SPEC_DEBUG and self._spec_debug_left > 0: + self._spec_debug_left -= 1 + # One block's whole token flow, in the order acceptance sees it. Reading + # this beats reasoning about the indices: the greedy output showed every + # word duplicated ("are are", "first first"), which is an emission fault, + # and only the actual tokens say where. + logger.info_rank0( + "spec block: greedy=%s temp=%s topp=%s conf=%s " + "committed=%s proposed=%s target=%s width=%d n_acc=%d bonus=%s " + "emitted=%s", + greedy, sp.temperature, sp.top_p, + None if confidence is None else + [round(float(c), 3) for c in confidence[i * k:(i + 1) * k]], + int(req._ids_buf[start - 1]) if start > 0 else None, + proposed_snapshot[:width].tolist(), + target_cpu[off:off + width + 1].tolist(), + width, n_acc, bonus, + emitted[-1].tolist() if emitted else None, + ) self._spec_accepted += n_acc self._spec_drafted += width + selected_rows.append(off + n_acc) + accepted_counts.append(n_acc) 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() + # Select the target's saved compressor state after anchor + accepted prefix. + # Later rejected rows may share its 128-token page and overwrite the live ring. + self._restore_speculative_carry(batch, selected_rows) + committed_features = self._trim_dspark_target_features( + target_features, batch, accepted_counts + ) + self._commit_dspark_target_features(committed_features) # The reply path reads one token per request; hand it the LAST emitted token and # let the scheduler read the rest off req.input_ids, which already holds them. @@ -1149,32 +1284,91 @@ def _finish_speculative( 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. + def _restore_speculative_carry(self, batch: Batch, selected_rows: list[int]) -> None: + """Commit the per-token partial state selected by DSpark acceptance.""" + journal = batch.spec_carry_states + if not journal: + raise RuntimeError("a DSpark target verify produced no compressor partial states") + device_rows = torch.tensor(selected_rows, dtype=torch.long, device=self.device) + md = batch.attn_metadata + table_rows = torch.tensor( + [md.segments[i][2] for i in range(len(selected_rows))], + dtype=torch.long, + device=self.device, + ) + positions = batch.positions.index_select(0, device_rows).long() + slots = self.attn_backend.window_slots_at(table_rows, positions) + for (layer_id, tier, ring_size), pieces in journal.items(): + if len(pieces) != batch.input_ids.numel(): + raise RuntimeError( + f"DSpark partial-state journal for layer {layer_id}/{tier} has " + f"{len(pieces)} rows, expected {batch.input_ids.numel()}" + ) + states = torch.cat(pieces, dim=0).index_select(0, device_rows) + self.attn_backend.write_carry_blocks( + layer_id, tier, slots, ring_size, states + ) - 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 + def _trim_dspark_target_features(self, features, batch: Batch, accepted: list[int]): + """Keep target features through each request's accepted predecessor row.""" + if features is None: + raise RuntimeError("a DSpark target verify returned no target features") + from freetoken.models.deepseek_v4.model import DSparkTargetFeatures - 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 + segments = batch.attn_metadata.segments + if segments is None or len(segments) != len(accepted): + raise RuntimeError( + "DSpark feature trimming needs one target segment per accepted count" ) - 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 + indices = [] + for (off, _n, _ti, _start), n_acc in zip(segments, accepted, strict=True): + indices.append( + torch.arange(off, off + n_acc + 1, dtype=torch.long, device=self.device) + ) + idx = torch.cat(indices) + if features.hidden.shape[0] != batch.input_ids.numel(): + raise RuntimeError( + f"DSpark target returned {features.hidden.shape[0]} feature rows for " + f"{batch.input_ids.numel()} verify inputs" + ) + return DSparkTargetFeatures( + features.hidden.index_select(0, idx), + features.positions.index_select(0, idx), + features.table_rows.index_select(0, idx), + ) + + def _real_dspark_target_features(self, features, batch: Batch): + """Drop CUDA-graph padding rows before context KV is committed.""" + if features is None or batch.is_prefill or features.hidden.shape[0] == batch.size: + return features + from freetoken.models.deepseek_v4.model import DSparkTargetFeatures + + if features.hidden.shape[0] < batch.size: + raise RuntimeError( + f"target produced {features.hidden.shape[0]} feature rows for " + f"{batch.size} real decode requests" + ) + idx = torch.arange(batch.size, dtype=torch.long, device=self.device) + return DSparkTargetFeatures( + features.hidden.index_select(0, idx), + features.positions.index_select(0, idx), + features.table_rows.index_select(0, idx), + ) + + def _commit_dspark_target_features(self, features) -> None: + """Precompute draft context KV as soon as target rows become committed. + + vLLM performs this precompute immediately before its next proposal. FreeToken + cannot safely retain one model-global bundle across unrelated prefill/decode + batches (or across several captured batch sizes), so it performs the same write + once the target rows are known to be valid. The next proposal observes the + identical draft-layer KV without stale cross-request ownership. + """ + if features is None: + return + catch_up = getattr(self.model, "catch_up_draft_context", None) + if catch_up is not None: + catch_up(features) def draft_into_batch(self, batch: Batch) -> torch.Tensor | None: """Fill a speculative batch's placeholder positions with the drafter's proposal. @@ -1190,58 +1384,95 @@ def draft_into_batch(self, batch: Batch) -> torch.Tensor | None: 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) + measure_draft = ( + self._adaptive_verification is not None + and self._adaptive_verification.needs_draft_profile + ) or (_SPEC_TIMING and self._spec_timing_left > 0) + if measure_draft: + batch._spec_wall_start = time.perf_counter() + batch._spec_draft_start = torch.cuda.Event(enable_timing=True) + batch._spec_draft_end = torch.cuda.Event(enable_timing=True) + batch._spec_draft_start.record(self.stream) with self.ctx.forward_batch(batch): - # 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() + out = drafter( + [req.sampling_params for req in batch.reqs], + ) + if measure_draft: + batch._spec_draft_end.record(self.stream) 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) + raise RuntimeError("DSpark was scheduled without a loaded drafter") + proposed, q, confidence = out + k = batch.spec_block + span = 1 + k + if proposed.numel() != k * len(batch.reqs): + raise RuntimeError( + f"DSpark proposed {proposed.numel()} tokens, expected " + f"{k} x {len(batch.reqs)}" + ) + for i in range(len(batch.reqs)): + batch.input_ids[i * span + 1:(i + 1) * span].copy_( + proposed[i * k:(i + 1) * k].to(batch.input_ids.dtype) + ) + batch.draft_tokens = proposed batch.draft_probs = q + batch.spec_carry_states = {} return confidence def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput: assert torch.cuda.current_stream() == self.stream + target_features = None + timing = bool( + batch.speculative and _SPEC_TIMING and self._spec_timing_left > 0 + ) + if timing: + target_start = torch.cuda.Event(enable_timing=True) + target_end = torch.cuda.Event(enable_timing=True) + target_start.record(self.stream) with self.ctx.forward_batch(batch): - if self.graph_runner.can_use_cuda_graph(batch): + if self.graph_runner.can_use_spec_cuda_graph(batch): + logits = self.graph_runner.replay_spec(batch) + target_features = self.graph_runner.dspark_spec_target_features(batch) + elif self.graph_runner.can_use_cuda_graph(batch): logits = self.graph_runner.replay(batch) + target_features = self.graph_runner.dspark_target_features(batch) else: logits = self.model.forward() + get_features = getattr(self.model, "dspark_target_features", None) + target_features = get_features() if get_features is not None else None + if timing: + target_end.record(self.stream) if self.cpu_moe_executor is not None: # One pinned read: surfaces a fired flag-handshake watchdog (dead coordinator # -> stale expert outputs) as a loud error instead of silent corruption. self.cpu_moe_executor.raise_if_unhealthy() if batch.speculative: - return self._finish_speculative(batch, logits, args) + finish_wall_start = time.perf_counter() + output = self._finish_speculative(batch, logits, args, target_features) + if timing: + cycle_end = torch.cuda.Event(enable_timing=True) + cycle_end.record(self.stream) + cycle_end.synchronize() + draft_start = batch._spec_draft_start + draft_end = batch._spec_draft_end + logger.info_rank0( + "spec timing: draft_cuda=%.2fms target_cuda=%.2fms " + "post_cuda=%.2fms finish_wall=%.2fms cycle_wall=%.2fms", + draft_start.elapsed_time(draft_end), + target_start.elapsed_time(target_end), + target_end.elapsed_time(cycle_end), + (time.perf_counter() - finish_wall_start) * 1000, + (time.perf_counter() - batch._spec_wall_start) * 1000, + ) + self._spec_timing_left -= 1 + return output + + self._commit_dspark_target_features( + self._real_dspark_target_features(target_features, batch) + ) for req in batch.reqs: req.complete_one() diff --git a/python/freetoken/engine/graph.py b/python/freetoken/engine/graph.py index 4f20250..29ff569 100644 --- a/python/freetoken/engine/graph.py +++ b/python/freetoken/engine/graph.py @@ -1,8 +1,9 @@ from __future__ import annotations import gc +import statistics from dataclasses import dataclass -from typing import TYPE_CHECKING, Dict, List +from typing import TYPE_CHECKING, Dict, List, Tuple import torch from freetoken.core import Batch, Req, get_global_ctx @@ -26,6 +27,7 @@ class GraphCaptureBuffer: positions: torch.Tensor logits: torch.Tensor table_idx: torch.Tensor # per-request slot id for GatedDeltaNet state gather/scatter + request_table_idx: torch.Tensor # per-request paged-KV table row # Decode GDN query indptr = arange(bs+1); a constant per captured bs, filled once. fla_cu_seqlens: torch.Tensor @@ -37,6 +39,7 @@ def init(cls, bs: int, vocab_size: int, device: torch.device) -> GraphCaptureBuf positions=torch.zeros(bs, dtype=torch.int32, device=device), logits=torch.empty(bs, vocab_size, dtype=torch.float32, device=device), table_idx=torch.zeros(bs, dtype=torch.int32, device=device), + request_table_idx=torch.zeros(bs, dtype=torch.int64, device=device), fla_cu_seqlens=torch.arange(bs + 1, dtype=torch.int32, device=device), ) @@ -49,6 +52,7 @@ def set_batch(self, batch: Batch) -> None: batch.out_loc = self.out_loc[_slice] batch.positions = self.positions[_slice] batch.linear_table_idx = self.table_idx[_slice] + batch.active_table_idx = self.request_table_idx[_slice] # Decode GDN metadata reads the persistent cu_seqlens (constant arange) and the # persistent table_idx slot map, so the captured kernels see stable addresses. batch.fla_metadata = FLAMetadata( @@ -63,6 +67,113 @@ def copy_from(self, batch: Batch) -> None: self.positions[_slice] = batch.positions if batch.linear_table_idx is not None: self.table_idx[_slice] = batch.linear_table_idx + if batch.active_table_idx is not None: + self.request_table_idx[_slice] = batch.active_table_idx + + +@dataclass +class SpecGraphCaptureBuffer: + """Persistent inputs/outputs for ``anchor + prefix`` target graphs. + + The allocation is sized for the checkpoint's maximum ``anchor + gamma`` span, + while each captured graph takes a prefix of it. Keeping one backing allocation + gives every replay a stable address without paying one logits buffer per prefix. + """ + + input_ids: torch.Tensor + out_loc: torch.Tensor + positions: torch.Tensor + logits: torch.Tensor + request_table_idx: torch.Tensor + span: int + + @classmethod + def init( + cls, max_reqs: int, span: int, vocab_size: int, device: torch.device + ) -> "SpecGraphCaptureBuffer": + max_tokens = max_reqs * span + return cls( + input_ids=torch.zeros(max_tokens, dtype=torch.int32, device=device), + out_loc=torch.zeros(max_tokens, dtype=torch.int32, device=device), + positions=torch.zeros(max_tokens, dtype=torch.int32, device=device), + logits=torch.empty(max_tokens, vocab_size, dtype=torch.float32, device=device), + request_table_idx=torch.zeros(max_reqs, dtype=torch.int64, device=device), + span=span, + ) + + def set_batch(self, batch: Batch) -> None: + reqs = batch.padded_size + span = int(batch.spec_block) + 1 + if not 1 <= span <= self.span: + raise RuntimeError( + f"DSpark graph span {span} is outside captured range 1..{self.span}" + ) + tokens = reqs * span + batch.input_ids = self.input_ids[:tokens] + batch.out_loc = self.out_loc[:tokens] + batch.positions = self.positions[:tokens] + batch.active_table_idx = self.request_table_idx[:reqs] + + def copy_from(self, batch: Batch) -> None: + reqs = batch.padded_size + span = int(batch.spec_block) + 1 + if not 1 <= span <= self.span: + raise RuntimeError( + f"DSpark graph span {span} is outside captured range 1..{self.span}" + ) + tokens = reqs * span + if batch.input_ids.numel() != tokens or batch.positions.numel() != tokens: + raise RuntimeError( + f"DSpark graph expected {tokens} target rows, got " + f"{batch.input_ids.numel()}/{batch.positions.numel()}" + ) + self.input_ids[:tokens].copy_(batch.input_ids) + if batch.out_loc is not None: + self.out_loc[:tokens].copy_(batch.out_loc) + self.positions[:tokens].copy_(batch.positions) + if batch.active_table_idx is None or batch.active_table_idx.numel() != reqs: + raise RuntimeError("DSpark graph needs one active table row per request") + self.request_table_idx[:reqs].copy_(batch.active_table_idx) + + +class SharedSpecCarryJournal: + """Maximum-span carry outputs shared by every adaptive verify graph. + + The largest graph owns the actual tensors produced by each compressor step. + Smaller graph captures copy their prefix states into those same addresses. This + is safe because verify graphs never replay concurrently, and it avoids retaining + one very large compressor journal per CUDA-graph bucket. + """ + + def __init__(self, storage: dict) -> None: + if not storage: + raise RuntimeError("DSpark maximum verify graph produced no carry states") + self.storage = storage + self._active_rows = 0 + self._cursors: dict = {} + + def reset(self, active_rows: int) -> None: + if active_rows < 1: + raise ValueError(f"DSpark carry journal needs at least one row, got {active_rows}") + self._active_rows = active_rows + self._cursors = {key: 0 for key in self.storage} + + def record(self, key, state: torch.Tensor) -> None: + pieces = self.storage.get(key) + if pieces is None: + raise RuntimeError(f"DSpark carry journal saw unknown compressor {key}") + cursor = self._cursors[key] + if cursor >= self._active_rows or cursor >= len(pieces): + raise RuntimeError( + f"DSpark carry journal overflow for {key}: row {cursor}, " + f"active={self._active_rows}, capacity={len(pieces)}" + ) + pieces[cursor].copy_(state) + self._cursors[key] = cursor + 1 + + def items(self): + for key, pieces in self.storage.items(): + yield key, pieces[: self._active_rows] def _determine_cuda_graph_bs( @@ -118,12 +229,45 @@ def __init__( self.moe_offload_cache = moe_offload_cache self.stream = stream self.device = device + # A captured model forward does not execute Python on replay, so a model-global + # "last features" attribute remains pointed at whichever graph was captured + # last. Keep the actual output bundle owned by each graph size instead. + self._dspark_feature_map: Dict[int, object] = {} + self._spec_feature_map: Dict[Tuple[int, int], object] = {} + self._spec_carry_map: Dict[Tuple[int, int], dict] = {} + # Measured (span, milliseconds) curve for the paper's hardware-aware + # verification scheduler. It is populated only for the single-request + # graphs that FreeToken's current DSV4 serving mode can compact safely. + self.spec_verify_cost_curve: List[Tuple[int, float]] = [] + self.spec_block_size = int( + getattr(model, "speculative_verify_block_size", 0) or 0 + ) + self.spec_span = self.spec_block_size + 1 if self.spec_block_size else 0 self._capture_graphs(max_seq_len, vocab_size, model) def _reset_moe_offload_cache(self) -> None: if self.moe_offload_cache is not None: self.moe_offload_cache.reset() + def _profile_graph_ms(self, graph: torch.cuda.CUDAGraph, replays: int = 5) -> float: + """Median replay cost with a cold expert cache, matching vLLM's profiler. + + A reset precedes every sample because the real hybrid workload continually + changes routed experts. Pricing a repeated dummy route after its first replay + would measure an all-hot cache and systematically underprice verification. + """ + samples: list[float] = [] + for _ in range(replays): + self._reset_moe_offload_cache() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record(self.stream) + graph.replay() + end.record(self.stream) + end.synchronize() + samples.append(start.elapsed_time(end)) + return float(statistics.median(samples)) + def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel): # Mark the post-weights "warmup" phase for /health: this stretch (graph capture — or the # remaining readiness work when graphs are disabled) moves no bytes, so without this the @@ -132,6 +276,7 @@ def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel # graphs-disabled early return so that config gets the phase too. emit_progress("Capturing CUDA graphs / warming up", 0, 0) self.graph_map: Dict[int, torch.cuda.CUDAGraph] = {} + self.spec_graph_map: Dict[Tuple[int, int], torch.cuda.CUDAGraph] = {} if self.max_graph_bs == 0: return logger.info_rank0("CUDA graph is disabled.") @@ -178,10 +323,104 @@ def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel with torch.cuda.graph(graph, pool=pool, stream=self.stream): self.buffer.logits[:bs] = model.forward() self._reset_moe_offload_cache() + get_features = getattr(model, "dspark_target_features", None) + if get_features is not None: + features = get_features() + if features is not None: + # These tensors are graph outputs: replay of this exact graph writes + # them in place. Holding the bundle preserves the matching addresses. + self._dspark_feature_map[bs] = features if pool is None: pool = graph.pool() # reuse cuda graph handle to reduce memory self.graph_map[bs] = graph + spec_capture = getattr(self.attn_backend, "prepare_for_spec_capture", None) + if self.spec_span and spec_capture is not None: + self.spec_buffer = SpecGraphCaptureBuffer.init( + self.max_graph_bs, self.spec_span, vocab_size, self.device + ) + # Section 5.2 chooses among the actual CUDA-graph cost cliffs. DSV4 + # currently serves one request at a time, so capture every legal prefix + # there. Keep the prior fixed-width behavior for larger request batches; + # those need the paper's marker-tensor varlen layout before independent + # per-request widths are safe. + adaptive_single_req = self.graph_bs_list == [1] + spans = ( + list(range(1, self.spec_span + 1)) + if adaptive_single_req + else [self.spec_span] + ) + logger.info_rank0( + f"Capturing DSpark target verify graphs: requests={self.graph_bs_list}, " + f"rows/request={spans}" + ) + shared_carry: SharedSpecCarryJournal | None = None + for span in sorted(spans, reverse=True): + for bs in sorted(self.graph_bs_list, reverse=True): + tokens = bs * span + graph = torch.cuda.CUDAGraph() + batch = Batch(reqs=[self.dummy_req] * bs, phase="prefill") + batch.padded_reqs = batch.reqs + batch.speculative = True + batch.spec_block = span - 1 + batch.spec_verify_decode = True + batch.spec_carry_states = {} + self.spec_buffer.set_batch(batch) + # Valid device positions keep capture-time writes within the reserved + # dummy row; replay overwrites this same stable buffer. + self.spec_buffer.positions[:tokens].copy_( + torch.arange(span, device=self.device, dtype=torch.int32).repeat(bs) + ) + self.spec_buffer.request_table_idx[:bs].fill_(self.dummy_req.table_idx) + spec_capture(batch) + active_rows = tokens + with get_global_ctx().forward_batch(batch): + if shared_carry is not None: + shared_carry.reset(active_rows) + batch.spec_carry_states = shared_carry + self.spec_buffer.logits[:tokens] = model.forward() + if shared_carry is None: + batch.spec_carry_states = {} + else: + shared_carry.reset(active_rows) + batch.spec_carry_states = shared_carry + with torch.cuda.graph(graph, pool=pool, stream=self.stream): + self.spec_buffer.logits[:tokens] = model.forward() + self._reset_moe_offload_cache() + if shared_carry is None and adaptive_single_req: + # Spans are descending, so this is the maximum graph. Its + # captured outputs become the one persistent journal backing + # all remaining prefix graphs. + shared_carry = SharedSpecCarryJournal(batch.spec_carry_states) + shared_carry.reset(active_rows) + get_features = getattr(model, "dspark_target_features", None) + key = (bs, span) + if get_features is not None: + features = get_features() + if features is not None: + self._spec_feature_map[key] = features + self._spec_carry_map[key] = ( + shared_carry if shared_carry is not None else batch.spec_carry_states + ) + self.spec_graph_map[key] = graph + if adaptive_single_req and bs == 1: + cost_ms = self._profile_graph_ms(graph) + self.spec_verify_cost_curve.append((span, cost_ms)) + + if self.spec_verify_cost_curve: + # vLLM makes profiled costs nondecreasing before scheduling. Sampling + # noise must not make a wider target batch appear artificially cheaper. + high = 0.0 + monotonic = [] + for span, cost in sorted(self.spec_verify_cost_curve): + high = max(high, cost) + monotonic.append((span, high)) + self.spec_verify_cost_curve = monotonic + logger.info_rank0( + "DSpark profiled target verify curve: %s", + ", ".join(f"{span} rows={cost:.2f}ms" for span, cost in monotonic), + ) + self._reset_moe_offload_cache() free_memory = get_free_memory(self.device) logger.info_rank0(f"Free GPU memory after capturing CUDA graphs: {mem_GB(free_memory)}") @@ -189,6 +428,14 @@ def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel def can_use_cuda_graph(self, batch: Batch) -> bool: return batch.is_decode and batch.size <= self.max_graph_bs + def can_use_spec_cuda_graph(self, batch: Batch) -> bool: + key = (batch.padded_size, int(batch.spec_block) + 1) + return bool( + batch.speculative + and key in self.spec_graph_map + and batch.padded_size == batch.size + ) + def replay(self, batch: Batch) -> torch.Tensor: assert self.can_use_cuda_graph(batch) self.buffer.copy_from(batch) @@ -197,6 +444,33 @@ def replay(self, batch: Batch) -> torch.Tensor: g.replay() return self.buffer.logits[: batch.size] + def replay_spec(self, batch: Batch) -> torch.Tensor: + assert self.can_use_spec_cuda_graph(batch) + span = int(batch.spec_block) + 1 + tokens = batch.padded_size * span + key = (batch.padded_size, span) + self.spec_buffer.copy_from(batch) + self.spec_buffer.set_batch(batch) + batch.spec_verify_decode = True + journal = self._spec_carry_map[key] + reset_journal = getattr(journal, "reset", None) + if reset_journal is not None: + reset_journal(tokens) + batch.spec_carry_states = journal + prepare = getattr(self.attn_backend, "prepare_for_spec_replay") + prepare(batch) + self.spec_graph_map[key].replay() + return self.spec_buffer.logits[:tokens] + + def dspark_target_features(self, batch: Batch): + """The feature output owned by the graph replayed for ``batch``.""" + return self._dspark_feature_map.get(batch.padded_size) + + def dspark_spec_target_features(self, batch: Batch): + return self._spec_feature_map.get( + (batch.padded_size, int(batch.spec_block) + 1) + ) + def pad_batch(self, batch: Batch) -> None: padded_size = ( # choose the first available batch size next(bs for bs in self.graph_bs_list if bs >= batch.size) @@ -213,5 +487,11 @@ def destroy_cuda_graphs(self) -> None: # free-before-alloc cannot reclaim this GPU memory. empty_cache() is left to the # caller / next capture (GraphRunner._capture_graphs already runs it). self.graph_map = {} + self.spec_graph_map = {} + self._dspark_feature_map = {} + self._spec_feature_map = {} + self._spec_carry_map = {} + self.spec_verify_cost_curve = [] self.buffer = None + self.spec_buffer = None gc.collect() diff --git a/python/freetoken/env.py b/python/freetoken/env.py index 2df433d..6878dc4 100644 --- a/python/freetoken/env.py +++ b/python/freetoken/env.py @@ -71,11 +71,6 @@ 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/kernel/triton/dsv4/indexer.py b/python/freetoken/kernel/triton/dsv4/indexer.py index 3d724e7..86040b0 100644 --- a/python/freetoken/kernel/triton/dsv4/indexer.py +++ b/python/freetoken/kernel/triton/dsv4/indexer.py @@ -117,7 +117,7 @@ def indexer_logits( @triton.jit def _indexer_decode_logits_kernel( - q_ptr, w_ptr, pool_ptr, snap_ptr, valid_ptr, out_ptr, + q_ptr, w_ptr, pool_ptr, snap_ptr, rows_ptr, valid_ptr, out_ptr, N_STAGE, RATIO, stride_qb, stride_qh, stride_qd, stride_wb, stride_wh, @@ -156,8 +156,11 @@ def _indexer_decode_logits_kernel( h_mask = offs_h < H # block b -> its compressed row: full_loc(b * RATIO) // RATIO off the decode snapshot. + # There may be several query rows per request (DSpark target verification), so + # pid_b is a QUERY row and rows_ptr maps it to the corresponding SNAPSHOT row. # Masked-off lanes read snapshot column 0 and are dropped by ``t_mask`` below. - snap = tl.load(snap_ptr + pid_b * stride_sb + (offs_t * RATIO) * stride_sw, + snap_row = tl.load(rows_ptr + pid_b) + snap = tl.load(snap_ptr + snap_row * stride_sb + (offs_t * RATIO) * stride_sw, mask=t_mask, other=0) rows = tl.maximum(snap, 0) // RATIO @@ -179,8 +182,9 @@ def indexer_decode_logits( q: torch.Tensor, # [B, H, D] bf16, one query per request weights: torch.Tensor, # [B, H] idx_pool: torch.Tensor, # [R, D] bf16, this layer's paged indexer keys - full_snap: torch.Tensor, # [B, W] int64, the decode full-loc snapshot - valid: torch.Tensor, # [B] live compressed block count per row + full_snap: torch.Tensor, # [R, W] int64, one snapshot row per request + rows: torch.Tensor, # [Q] query row -> snapshot row + valid: torch.Tensor, # [Q] live compressed block count per query n_stage: int, # static staged width (the buffer the top-k scans) ratio: int, out: torch.Tensor | None = None, @@ -199,6 +203,7 @@ def indexer_decode_logits( """ B, H, D = q.shape assert weights.shape == (B, H), (weights.shape, (B, H)) + assert rows.shape == (B,), (rows.shape, (B,)) assert idx_pool.shape[1] == D, (idx_pool.shape, D) assert D == triton.next_power_of_2(D), f"index_head_dim must be pow2, got {D}" @@ -209,7 +214,7 @@ def indexer_decode_logits( BLOCK_T = 64 _indexer_decode_logits_kernel[(B, triton.cdiv(n_stage, BLOCK_T))]( - q, weights, idx_pool, full_snap, valid, out, + q, weights, idx_pool, full_snap, rows, valid, out, n_stage, ratio, q.stride(0), q.stride(1), q.stride(2), weights.stride(0), weights.stride(1), diff --git a/python/freetoken/models/deepseek_v4/attention.py b/python/freetoken/models/deepseek_v4/attention.py index c3790f0..21fda2d 100644 --- a/python/freetoken/models/deepseek_v4/attention.py +++ b/python/freetoken/models/deepseek_v4/attention.py @@ -158,9 +158,13 @@ def _prefill_segment(self, x_seg, qr_seg, kv_seg, ti: int, start_pos: int, n: in # natural width min(n, win); the caller pads to the batch-uniform width win_global = self.attn.win_cols_to_global(win_cols, slots) else: - # Candidates span [w_lo, end): each new query's 128-sliding window over the retained - # prefix plus the new tokens, causal-masked to -1 past its own position. - w_lo = max(0, start_pos - win + 1) + # Target attention uses its ordinary 128-token causal window. A DSpark + # draft query instead sees 128 CONTEXT tokens plus every token in its + # parallel query block (paper section 3.1; same sparse list as vLLM). + w_lo = max( + 0, + start_pos - win if self.non_causal else start_pos - win + 1, + ) ws_pool = self.attn.window_slots_of(ti, w_lo, end) # 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 @@ -188,7 +192,11 @@ 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: + elif ( + self.non_causal + or getattr(get_global_ctx().batch, "speculative", False) + or start_pos % self.P != 0 + ): # A speculative block starts wherever generation reached, which is mid-page # 127 times out of 128. The compressor's EXTEND path cannot serve that: it # reduces tokens with ``start_pos // ratio`` arithmetic and asserts @@ -202,7 +210,18 @@ def _prefill_segment(self, x_seg, qr_seg, kv_seg, ti: int, start_pos: int, n: in # 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) + # Run the INDEXER here too. Falling back to _compress_topk_extend chooses + # compressed blocks positionally, which silently replaces DeepSeek-V4's + # learned sparse selection with a fixed one -- so the target attends to the + # wrong blocks and its own logits are wrong. Since a speculative block is + # unaligned 127 times out of 128, that was every verify. + blocks = ( + self.indexer.extend_unaligned( + x_seg, qr_seg, start_pos, 0, slots, tail_ws, ti + ) + if self.indexer is not None + else self._compress_topk_extend(n, start_pos, end, 0, device, 1) + ) else: blocks = ( self.indexer.extend(x_seg, qr_seg, start_pos, 0, slots, tail_ws, ti) @@ -288,7 +307,16 @@ def forward_ragged(self, x, segments, flat_positions): # they shift the kernel's window/compressed tile split, and that rounding-order change # is enough to flip a greedy near-tie. bs>1 pads every part to win (the historical # batched behavior). - n_window = win if len(segments) > 1 else win_parts[0].shape[-1] + # Ordinary target segments keep the historical batch-uniform ``win`` width. + # DSpark's non-causal query block is wider: trailing ``win`` context columns + # PLUS all gamma query columns. Labeling only the first 128 columns as window + # in a multi-request draft misclassifies/drops the query-query suffix even + # though the sparse indices themselves correctly contain all 133 columns. + n_window = ( + max(part.shape[-1] for part in win_parts) + if self.non_causal + else (win if len(segments) > 1 else win_parts[0].shape[-1]) + ) flat = [] for i, (off, n, ti, _start) in enumerate(segments): wg = win_parts[i] @@ -398,3 +426,77 @@ def decode_step( ) apply_rotary_emb_decode(o[..., -rd:], freqs_t, True) # per-row position freqs (inverse) return self._wo(o, B, 1) + + def verify_block( + self, + x: torch.Tensor, + pos: torch.Tensor, + rows: torch.Tensor, + cmp_stage_cap: int, + num_reqs: int, + span: int, + wctx, + ) -> torch.Tensor: + """Fixed-shape causal target verification using decode-addressed kernels. + + Large projections, MoE work and sparse attention remain batched over all + verification rows. Stateful compressors advance in request-major token order + so each token reads its predecessor's new carry, with no host-derived position + or address embedded in the CUDA graph. + """ + n = x.shape[1] + if n != num_reqs * span: + raise ValueError(f"verify block has {n} rows, expected {num_reqs} * {span}") + win, ratio, rd = self.window_size, self.compress_ratio, self.rope_head_dim + flat = x.reshape(n, self.dim) + xq = flat.unsqueeze(1) + freqs_t = self.freqs_cis.index_select(0, pos) + + qr = q = self.q_norm(self.wq_a(xq)) + q = self.wq_b(q).unflatten(-1, (self.n_heads, self.head_dim)) + q = rms_norm(q, None, self.eps) + apply_rotary_emb_decode(q[..., -rd:], freqs_t) + + kv = self.kv_norm(self.wkv(xq)) + apply_rotary_emb_decode(kv[..., -rd:], freqs_t) + act_quant_fp8_inplace(kv[..., :-rd], 64) + + window_slots, prev_window_slots, window_slots_topk = wctx + self.attn.store_window(kv.view(n, -1), self.layer_id, window_slots) + n_cmp_stage = (cmp_stage_cap + 1) // ratio if ratio else 0 + cmp_counts = None + if ratio: + # This order is also the carry-journal order consumed after acceptance. + for r in range(num_reqs): + base = r * span + row = rows[base:base + 1] + for s in range(span): + i = base + s + self.compressor.decode_step( + flat[i:i + 1], pos[i:i + 1], + prev_window_slots[i:i + 1], window_slots[i:i + 1], row, + ) + valid = (pos + 1) // ratio + if self.indexer is not None: + blocks = self.indexer.verify_block( + flat, qr.reshape(n, -1), pos, + prev_window_slots, window_slots, rows, n_cmp_stage, + num_reqs, span, + ) + else: + blk = torch.arange(n_cmp_stage, device=x.device) + blocks = torch.where( + blk[None, :] < valid[:, None], blk[None, :], -1 + ).view(n, 1, n_cmp_stage) + compress_idxs = self.attn.blocks_to_global(blocks, ratio, rows=rows) + topk_idxs = torch.cat([window_slots_topk, compress_idxs], dim=-1) + cmp_counts = valid.clamp(max=compress_idxs.shape[-1]).to(torch.int32).view(n, 1) + else: + topk_idxs = window_slots_topk + + o = self.attn.attend( + q, self.layer_id, topk_idxs.int(), win, self.attn_sink, + self.softmax_scale, cmp_counts=cmp_counts, has_compression=bool(ratio), + ) + apply_rotary_emb_decode(o[..., -rd:], freqs_t, True) + return self._wo(o, n, 1).view(1, n, self.dim) diff --git a/python/freetoken/models/deepseek_v4/compress.py b/python/freetoken/models/deepseek_v4/compress.py index 43fe58a..ac2cf3f 100644 --- a/python/freetoken/models/deepseek_v4/compress.py +++ b/python/freetoken/models/deepseek_v4/compress.py @@ -338,9 +338,21 @@ def decode_step( compressed = gated_pool(ks, ss, dtype) # Persist the advanced carry block to THIS token's page (per row, disjoint blocks); the # next step reads it from here. Within a page prev==cur, so this is a self-consistent roll. + carry_state = torch.cat([ks, ss], dim=-1) self.attn.write_carry_blocks( - self.layer_id, self.tier, window_slots, self.ring_size, torch.cat([ks, ss], dim=-1) + self.layer_id, self.tier, window_slots, self.ring_size, carry_state ) + batch = get_global_ctx().batch + if getattr(batch, "speculative", False): + journal = batch.spec_carry_states + if journal is None: + raise RuntimeError("a speculative verify did not initialize its carry journal") + key = (self.layer_id, self.tier, self.ring_size) + record = getattr(journal, "record", None) + if record is None: + journal.setdefault(key, []).append(carry_state) + else: + record(key, carry_state) compressed = self.norm(compressed) freqs_t = self.freqs_cis.index_select(0, (pos + 1 - ratio).clamp_min(0)) # [B, rd//2] @@ -446,6 +458,57 @@ def extend(self, x, qr, start_pos, offset, window_slots, tail_window_slot, ti: i topk=self.index_topk, offset=offset, ) + def extend_unaligned(self, x, qr, start_pos, offset, window_slots, tail_window_slot, ti: int = 0): + """``extend``, for a segment that does NOT start on a compress-ratio boundary. + + ``extend`` cannot serve one: it advances its compressor through + ``Compressor.forward``, which for ``start_pos > 0`` delegates to + ``Compressor.extend`` and asserts 128-alignment, because that path seeds the + carry from a matched tail PAGE. + + A speculative block starts wherever generation reached, so it is unaligned 127 + times out of 128 -- and without this the caller had no indexer to fall back on + and chose compressed blocks POSITIONALLY instead. That silently replaces + DeepSeek-V4's learned sparse selection with a fixed one, so the TARGET attends + to the wrong blocks and its own logits are wrong. It never raises; it just + degenerates the output. + + Everything below the compressor advance is identical to ``extend``: only the way + the compressor is walked differs, one token at a time off the previous token's + window slot -- the same movement a decode step makes, with no block-boundary + arithmetic, so any ``start_pos`` is valid. + """ + bsz, seqlen, _ = x.size() + end = start_pos + seqlen + freqs_cis = self.freqs_cis[start_pos:end] + ratio, rd = self.compress_ratio, self.rope_head_dim + q = self.wq_b(qr).unflatten(-1, (self.n_heads, self.head_dim)) + apply_rotary_emb(q[..., -rd:], freqs_cis) + q = hadamard_transform(q) + fp4_act_quant_inplace(q, 32) + + # The one difference from extend(): walk the indexer's own compressor per token. + device = x.device + rows = torch.zeros(1, dtype=torch.int64, device=device) + prev = torch.full( + (1,), + tail_window_slot if tail_window_slot is not None else int(window_slots[0]), + dtype=window_slots.dtype, device=device, + ) + for t in range(seqlen): + cur = window_slots[t:t + 1] + pos_t = torch.full((1,), start_pos + t, dtype=torch.int64, device=device) + self.compressor.decode_step(x[:, t], pos_t, prev, cur, rows, ti=ti) + prev = cur + + weights = self.weights_proj(x) * (self.softmax_scale * self.n_heads ** -0.5) + keys = self.attn.indexer_keys(ti, end // ratio, ratio, self.layer_id, bsz) + scores = self.attn.indexer_prefill_logits(q, keys, weights) + return self.attn.indexer_select_prefill( + scores, start_pos=start_pos, seqlen=seqlen, ratio=ratio, + topk=self.index_topk, offset=offset, + ) + def decode_step( self, x: torch.Tensor, qr: torch.Tensor, pos: torch.Tensor, offset: int, prev_window_slots: torch.Tensor, window_slots: torch.Tensor, rows: torch.Tensor, @@ -475,8 +538,56 @@ def decode_step( index_score = self.attn.indexer_decode_scores( q.reshape(B, self.n_heads, self.head_dim), weights.reshape(B, self.n_heads), - valid, n_stage, ratio, self.layer_id, + valid, n_stage, ratio, self.layer_id, rows, ).view(B, 1, n_stage) return self.attn.indexer_select_decode( index_score, valid=valid, topk=self.index_topk, offset=offset ) + + def verify_block( + self, + x: torch.Tensor, + qr: torch.Tensor, + pos: torch.Tensor, + prev_window_slots: torch.Tensor, + window_slots: torch.Tensor, + rows: torch.Tensor, + n_stage: int, + num_reqs: int, + span: int, + ) -> torch.Tensor: + """Graph-safe Lightning Indexer for a fixed DSpark target block. + + Projection and scoring stay batched over every verification row. Only the + indexer's recurrent compressor walks in request-major token order, matching + ``extend_unaligned`` while positions and addresses remain device-resident. + """ + n, _ = x.shape + ratio, rd = self.compress_ratio, self.rope_head_dim + q = self.wq_b(qr).unflatten(-1, (self.n_heads, self.head_dim)) + apply_rotary_emb_decode( + q[..., -rd:].unsqueeze(1), self.freqs_cis.index_select(0, pos) + ) + q = hadamard_transform(q) + fp4_act_quant_inplace(q, 32) + + for r in range(num_reqs): + base = r * span + row = rows[base:base + 1] + for s in range(span): + i = base + s + self.compressor.decode_step( + x[i:i + 1], pos[i:i + 1], + prev_window_slots[i:i + 1], window_slots[i:i + 1], row, + ) + + weights = self.weights_proj(x) * (self.softmax_scale * self.n_heads ** -0.5) + valid = (pos + 1) // ratio + scores = self.attn.indexer_decode_scores( + q.reshape(n, self.n_heads, self.head_dim), + weights.reshape(n, self.n_heads), + valid, n_stage, ratio, self.layer_id, rows, + ).view(n, 1, n_stage) + return self.attn.indexer_select_decode( + scores, valid=valid, topk=self.index_topk, offset=0 + ) diff --git a/python/freetoken/models/deepseek_v4/dspark.py b/python/freetoken/models/deepseek_v4/dspark.py index 642d6d9..942efde 100644 --- a/python/freetoken/models/deepseek_v4/dspark.py +++ b/python/freetoken/models/deepseek_v4/dspark.py @@ -27,19 +27,188 @@ from __future__ import annotations -from typing import Callable, NamedTuple +import math +import statistics +from collections import deque +from typing import Sequence import torch import torch.nn.functional as F from torch import nn -from freetoken.kernel.triton.dsv4.bf16_linear import bf16_linear_fp32 from freetoken.kernel.triton.dsv4.hc import hc_pre_combine +from freetoken.utils import init_logger from .args import DeepseekV4Args from .layers import Linear, RMSNorm from .parallel import div_tp -from .rollback import needs_rollback + +logger = init_logger(__name__) + + +def choose_adaptive_draft_width( + stale_confidence: Sequence[float] | torch.Tensor, + draft_cost_ms: float, + verify_cost_ms: Sequence[float], +) -> int: + """Choose the DSpark prefix that maximizes expected tokens per millisecond. + + This is Algorithm 1 / Section 5.2 specialized to FreeToken's present DSV4 + serving shape of one active request. ``verify_cost_ms[k]`` is the measured + target cost for anchor + ``k`` draft rows. The drafter has already produced + its whole block, so its measured cost is common to every candidate width. + + For one request, globally sorting cumulative survival probabilities is exactly + the same as considering its prefixes in order: the survival product cannot + increase as the prefix grows. Ties retain the smaller prefix, matching + ``argmax`` in the vLLM implementation. + """ + confidence = [float(c) for c in stale_confidence] + if len(verify_cost_ms) != len(confidence) + 1: + raise ValueError( + "adaptive DSpark needs one verify cost for every width 0..gamma; " + f"got {len(verify_cost_ms)} costs for gamma={len(confidence)}" + ) + if not math.isfinite(draft_cost_ms) or draft_cost_ms < 0: + raise ValueError(f"invalid DSpark draft cost {draft_cost_ms}") + + survival = 1.0 + expected = 1.0 # target bonus token; a verify always advances at least once + best_width = 0 + best_throughput = expected / max( + draft_cost_ms + float(verify_cost_ms[0]), 1e-6 + ) + for width, token_confidence in enumerate(confidence, start=1): + survival *= token_confidence + expected += survival + throughput = expected / max( + draft_cost_ms + float(verify_cost_ms[width]), 1e-6 + ) + if throughput > best_throughput: + best_width = width + best_throughput = throughput + return best_width + + +class DSparkAdaptiveVerification: + """Two-buffer stale-confidence scheduler from DSpark Section 5.2. + + Current confidences are copied asynchronously to pinned host memory. Capacity + selection reads the other buffer, creating the paper's causal barrier across + jagged CUDA-graph buckets. The target curve is profiled at graph capture; the + fixed-cost draft pass is priced from five real executions, like vLLM's five + startup profiling replays. + """ + + _DRAFT_PROFILE_SAMPLES = 5 + + def __init__( + self, + block_size: int, + verify_curve: Sequence[tuple[int, float]], + device: torch.device, + ) -> None: + expected_spans = list(range(1, block_size + 2)) + curve = sorted((int(span), float(cost)) for span, cost in verify_curve) + if [span for span, _ in curve] != expected_spans: + raise ValueError( + "adaptive DSpark needs profiled graphs for every anchor+prefix span " + f"{expected_spans}, got {[span for span, _ in curve]}" + ) + high = 0.0 + self.verify_cost_ms: list[float] = [] + for _span, cost in curve: + high = max(high, cost, 1e-6) + self.verify_cost_ms.append(high) + + self.block_size = block_size + self.device = device + self._copy_stream = torch.cuda.Stream(device=device) + self._stale = [ + torch.ones(block_size, dtype=torch.float32, pin_memory=True) + for _ in range(2) + ] + self._events: list[torch.cuda.Event | None] = [None, None] + self._stale_idx = 0 + self._active_uid: int | None = None + self._draft_samples: deque[float] = deque( + maxlen=self._DRAFT_PROFILE_SAMPLES + ) + self._debug_left = 12 + + @property + def needs_draft_profile(self) -> bool: + return len(self._draft_samples) < self._DRAFT_PROFILE_SAMPLES + + @property + def draft_cost_ms(self) -> float | None: + if self.needs_draft_profile: + return None + return float(statistics.median(self._draft_samples)) + + def record_draft_cost(self, cost_ms: float) -> None: + if self.needs_draft_profile and math.isfinite(cost_ms) and cost_ms >= 0: + self._draft_samples.append(float(cost_ms)) + if not self.needs_draft_profile: + logger.info_rank0( + "DSpark profiled draft cost: %.2fms (median of %d real steps)", + self.draft_cost_ms, + self._DRAFT_PROFILE_SAMPLES, + ) + + def _reset_request(self, uid: int) -> None: + for event in self._events: + if event is not None: + event.synchronize() + for slot in self._stale: + slot.fill_(1.0) + self._events = [None, None] + self._stale_idx = 0 + self._active_uid = uid + + def record_and_choose(self, confidence: torch.Tensor, uid: int) -> int: + """Publish this block's confidence and choose from the stale buffer.""" + flat = confidence.detach().float().reshape(-1) + if flat.numel() != self.block_size: + raise RuntimeError( + f"DSpark confidence has {flat.numel()} rows, expected {self.block_size}" + ) + if self._active_uid != uid: + self._reset_request(uid) + + ready_idx = self._stale_idx ^ 1 + ready = self._events[ready_idx] + if ready is not None: + ready.synchronize() + self._stale_idx, write_idx = ready_idx, self._stale_idx + stale = self._stale[self._stale_idx] + + current_stream = torch.cuda.current_stream(self.device) + self._copy_stream.wait_stream(current_stream) + with torch.cuda.stream(self._copy_stream): + self._stale[write_idx].copy_(flat, non_blocking=True) + event = torch.cuda.Event(blocking=True) + event.record(self._copy_stream) + self._events[write_idx] = event + + draft_cost = self.draft_cost_ms + # Collect five real drafter samples before changing shape. This is a + # measured warmup, not a guessed initial cost. + if draft_cost is None: + return self.block_size + width = choose_adaptive_draft_width( + stale, draft_cost, self.verify_cost_ms + ) + if self._debug_left > 0: + self._debug_left -= 1 + logger.info_rank0( + "DSpark adaptive verify: stale=%s draft=%.2fms width=%d/%d", + [round(float(c), 3) for c in stale], + draft_cost, + width, + self.block_size, + ) + return width class MarkovHead(nn.Module): @@ -73,10 +242,9 @@ def bias(self, markov_embed: torch.Tensor) -> torch.Tensor: 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. + Its output feeds DSpark's hardware-aware prefix scheduler: cumulative survival + probabilities are ranked against profiled draft/verify costs. Runs in fp32 because + the checkpoint defines this as a calibrated scalar probability per position. """ def __init__(self, args: DeepseekV4Args): @@ -238,32 +406,69 @@ def catch_up_context( 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. + self, + input_ids: torch.Tensor, + segments, + positions: torch.Tensor, + target_logits_fn, + sampling_params: Sequence, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Run the paper's parallel backbone and sequential Markov sampler. + + The prepared target VERIFY span has ``1 + gamma`` rows per request: the anchor + followed by ``gamma`` proposal slots. The DSpark backbone takes only ``gamma`` + inputs -- anchor plus ``gamma - 1`` masks -- and produces ``gamma`` base-logit + rows. Equation 4 then samples those rows left to right, conditioning row ``k`` + on the token sampled at ``k - 1`` (the anchor for the first row). + + Returns the proposed tokens, the exact shaped distribution each proposal was + drawn from, and the confidence from equation 7, all flattened request-major. """ - # 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)) + gamma = self.block_size + draft_segments = [] + rows = [] + for i, (off, n, ti, start) in enumerate(segments): + if n != gamma + 1: + raise ValueError( + f"a DSpark verify span has {n} rows, expected gamma + 1 = " + f"{gamma + 1}" + ) + rows.append(torch.arange(off, off + gamma, device=input_ids.device)) + draft_segments.append((i * gamma, gamma, ti, start)) + row_idx = torch.cat(rows) + # The scheduler extends the CPU Req with noise placeholders, but the model's + # flat input is gathered from the GPU token pool; those newly allocated pool + # positions have never been initialized. Construct the reference layout here + # from device-resident data: one live anchor followed by gamma-1 checkpoint + # noise tokens per request (vLLM's sample_from_anchor path). + if self.noise_token_id < 0: + raise ValueError("DSpark requires the checkpoint's noise token id") + draft_ids = torch.full( + (len(segments), gamma), + self.noise_token_id, + dtype=input_ids.dtype, + device=input_ids.device, + ) + draft_ids[:, 0] = input_ids[0, row_idx.view(len(segments), gamma)[:, 0]] + draft_ids = draft_ids.view(1, -1) + draft_pos = positions[row_idx] + head_hidden = self.head_hidden( + self.embed_block(draft_ids), draft_ids, draft_segments, draft_pos + )[0] + # The checkpoint's two heads consume different forms of this tensor. vLLM's + # DSV4 DSpark model returns the PRE-norm hc_head output to the speculator: + # base logits use norm(head_hidden), while equation 7's confidence projection + # consumes head_hidden itself. Normalizing before returning silently changes + # every confidence score even though the proposal logits still look plausible. + base_logits = target_logits_fn(self.norm(head_hidden)).view( + len(segments), gamma, -1 + ) + return self.sample_block( + base_logits, + head_hidden.view(len(segments), gamma, -1), + draft_ids.view(len(segments), gamma)[:, 0], + sampling_params, + ) def embed_block(self, input_ids: torch.Tensor) -> torch.Tensor: """Token embeddings for the block, ``[1, T, dim]``. @@ -310,7 +515,7 @@ def hc_head(self, x: torch.Tensor) -> torch.Tensor: 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. + """Run the draft stack and return the pre-norm ``hc_head`` hidden. ``embeds`` is ``[1, T, dim]`` -- the block's token embeddings, already summed with the projected target hidden by the caller. Returns ``[1, T, dim]``. @@ -318,23 +523,55 @@ def head_hidden( 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) + return self.hc_head(h) + + def sample_block( + self, + base_logits: torch.Tensor, + head_hidden: torch.Tensor, + anchor: torch.Tensor, + sampling_params: Sequence, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Equation 4/5/7 over ``[request, gamma, ...]`` backbone outputs.""" + requests, gamma, vocab = base_logits.shape + if len(sampling_params) != requests: + raise ValueError( + f"got {len(sampling_params)} sampling parameter sets for {requests} requests" + ) + prev = anchor.long() + proposed = torch.empty( + (requests, gamma), dtype=torch.long, device=base_logits.device + ) + q = torch.empty( + (requests, gamma, vocab), dtype=torch.float32, device=base_logits.device + ) + confidence = torch.empty( + (requests, gamma), dtype=torch.float32, device=base_logits.device + ) + for k in range(gamma): + markov = self.markov_head.embed(prev) + logits_k = base_logits[:, k].float() + self.markov_head.bias(markov).float() + step_tokens = [] + for r, params in enumerate(sampling_params): + q_r = sampling_probs( + logits_k[r : r + 1], + params.temperature, + params.top_p, + params.top_k, + ) + q[r, k].copy_(q_r[0]) + step_tokens.append( + q_r.argmax(dim=-1) + if params.is_greedy + else torch.multinomial(q_r, 1).squeeze(-1) + ) + next_token = torch.cat(step_tokens) + proposed[:, k].copy_(next_token) + confidence[:, k].copy_( + self.confidence_head(head_hidden[:, k], markov) + ) + prev = next_token + return proposed.flatten(), q.flatten(0, 1), confidence.flatten() def accepted_prefix( @@ -434,117 +671,75 @@ def rejection_accept( 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. +def rejection_accept_device( + proposed: torch.Tensor, + q: torch.Tensor, + p: torch.Tensor, + generator: torch.Generator | None = None, +) -> tuple[int, int]: + """Device-resident form of vLLM's standard rejection sampler. - The 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. + The acceptance ratios, prefix break, residual distribution, and recovered-token + draw stay on the logits device. Only the accepted length and recovered token + cross to the host, which the scheduler needs to resize the request. This has the + same p/q rule as :func:`rejection_accept`; it removes the full-vocabulary D2H copy, + not any part of the distribution-correct algorithm. """ - - 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) + n = int(proposed.numel()) + if q.shape[0] != n or p.shape[0] < n + 1: + raise ValueError( + f"rejection sampler shape mismatch: proposed={n}, q={q.shape}, p={p.shape}" + ) + if n: + token_ids = proposed.long() + rows = torch.arange(n, device=proposed.device) + p_token = p[:n][rows, token_ids] + q_token = q[:n][rows, token_ids] + uniforms = torch.rand( + n, device=p.device, dtype=torch.float32, generator=generator + ) + rejected = torch.nonzero(p_token <= uniforms * q_token, as_tuple=False) + n_acc = int(rejected[0, 0]) if rejected.numel() else n + else: + n_acc = 0 + + if n_acc < n: + residual = torch.clamp(p[n_acc] - q[n_acc], min=0.0) + total = residual.sum() + # torch.where keeps the degenerate fallback on device; clamp_min only + # protects the unselected division branch from producing NaNs. + dist = torch.where( + total > 0, + residual / total.clamp_min(torch.finfo(residual.dtype).tiny), + p[n_acc], + ) + else: + dist = p[n] + bonus = int(torch.multinomial(dist, 1, generator=generator).item()) + return n_acc, bonus 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]``. + """The window candidate columns a segment's queries may read. Extracted from ``Attention._prefill_segment`` so the masking rule -- the one thing that decides whether a block is genuinely semi-autoregressive -- can be checked without a GPU or a KV pool. ``-1`` marks a column no query may read. Causal: row i sees ``[pos_i - window + 1, pos_i]``. - Non-causal: EVERY row sees ``[last - window + 1, last]``, where ``last`` is the - block's final position, so each drafted query sees the whole block. + Non-causal DSpark: every row sees the trailing ``window`` CONTEXT tokens plus all + ``n`` block tokens, matching the paper and vLLM's sparse-index construction. """ end = start_pos + n + if non_causal: + w_lo = max(0, start_pos - window) + width = end - w_lo + return torch.arange(width, device=device).unsqueeze(0).expand(n, width) w_lo = max(0, start_pos - window + 1) ar = torch.arange(window, device=device) - 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) @@ -552,11 +747,11 @@ def window_cols_for_block( __all__ = [ "ConfidenceHead", + "DSparkAdaptiveVerification", "DSparkDrafter", "MarkovHead", - "SpeculativeLoop", - "StepResult", "accepted_prefix", - "draft_width", + "choose_adaptive_draft_width", + "rejection_accept_device", "window_cols_for_block", ] diff --git a/python/freetoken/models/deepseek_v4/model.py b/python/freetoken/models/deepseek_v4/model.py index 3223001..b013a81 100644 --- a/python/freetoken/models/deepseek_v4/model.py +++ b/python/freetoken/models/deepseek_v4/model.py @@ -22,6 +22,7 @@ from __future__ import annotations import os +from dataclasses import dataclass import torch import torch.nn.functional as F @@ -53,6 +54,24 @@ logger = init_logger(__name__) +# FREETOKEN_SPEC_DEBUG=1 also reports where the draft context KV lands. +_DRAFT_CTX_DEBUG = os.environ.get("FREETOKEN_SPEC_DEBUG", "0") == "1" + + +@dataclass(frozen=True) +class DSparkTargetFeatures: + """The target outputs consumed by the next DSpark proposal. + + vLLM passes auxiliary hidden states, target positions, and context slot mappings + together. FreeToken stores request table rows instead of physical slots so the + slots are resolved from the live page map when the feature bundle is consumed. + Every tensor is produced by the same target forward. + """ + + hidden: torch.Tensor + positions: torch.Tensor + table_rows: torch.Tensor + class Block(nn.Module): """Decoder block with manifold-constrained Hyper-Connections (4 residual streams).""" @@ -133,6 +152,28 @@ def decode_step(self, x, pos, rows, cmp_stage_cap, input_ids, wctx=None): x = self.hc_post(x, residual, post, comb) return x + def verify_block( + self, x, pos, rows, cmp_stage_cap, input_ids, num_reqs, span, wctx + ): + """Fixed-width DSpark target block; algebra matches ``decode_step``.""" + residual = x + x, post, comb = self.hc_pre( + x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base + ) + x = self.attn_norm(x) + x = self.attn.verify_block( + x, pos, rows, cmp_stage_cap, num_reqs, span, wctx + ) + x = self.hc_post(x, residual, post, comb) + + residual = x + x, post, comb = self.hc_pre( + x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base + ) + x = self.ffn_norm(x) + x = self.ffn(x, input_ids) + return self.hc_post(x, residual, post, comb) + class Transformer(nn.Module): def __init__(self, args: DeepseekV4Args): @@ -169,9 +210,9 @@ def __init__(self, args: DeepseekV4Args): # 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. + # Layers whose output the drafter reads, as 0-based indices. vLLM converts the + # checkpoint ids [40, 41, 42] to aux ids [41, 42, 43], then captures when + # ``idx + 1`` is in that set: target outputs 40, 41 and 42 exactly. self._aux_layer_ids: frozenset[int] = frozenset() if args.n_draft_layers: from .dspark import DSparkDrafter @@ -181,45 +222,18 @@ 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 - # 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) + ids = tuple(args.dspark_target_layer_ids) self._aux_layer_ids = frozenset(ids) bad = [i for i in ids if not 0 <= i < args.n_layers] if bad: raise ValueError( - f"dspark_target_layer_ids {tuple(args.dspark_target_layer_ids)} read " - f"as {base}-based give {ids}, which fall outside the " + f"0-based dspark_target_layer_ids {ids} 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)" + f"dSpark: tapping target layer outputs {sorted(ids)}" ) - # 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 + self._target_features: DSparkTargetFeatures | None = None def bind(self, pool, device: torch.device) -> None: for layer in self.layers: @@ -227,25 +241,8 @@ def bind(self, pool, device: torch.device) -> None: if self.drafter is not None: self.drafter.bind(pool, device) - def last_aux_hidden(self) -> torch.Tensor | None: - """The tap from the most recent forward, or None if nothing has run yet.""" - return self._last_aux_hidden - - def last_aux_addressing(self) -> tuple[torch.Tensor, torch.Tensor] | None: - """``(positions, table_rows)`` for the rows of :meth:`last_aux_hidden`. - - The tap and its addressing must travel together. A consumer that pairs the tap - with any OTHER batch's positions writes hidden states derived from one set of - tokens into the slots of a different set -- which raises nothing, and shows up - only as a drafter that proposes badly. - """ - if ( - self._last_aux_hidden is None - or self._last_aux_positions is None - or self._last_aux_rows is None - ): - return None - return self._last_aux_positions, self._last_aux_rows + def target_features(self) -> DSparkTargetFeatures | None: + return self._target_features def embed_tokens(self, input_ids: torch.Tensor) -> torch.Tensor: """Vocabulary-parallel lookup. Rows outside this rank's block contribute zero, so @@ -307,13 +304,21 @@ def prefill_batched( # 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 + aux_hidden = ( + torch.cat(aux, dim=-1).view(-1, self.args.dim * len(aux)) + 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 + self._target_features = DSparkTargetFeatures( + aux_hidden, + flat_positions.clone(), + rows.clone(), + ) + else: + self._target_features = None h = self.hc_head(h) h = self.norm(h) # Normally only each request's final token needs logits. A speculative VERIFY @@ -350,14 +355,74 @@ def decode( 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 + aux_hidden = ( + torch.cat(aux, dim=-1).view(-1, self.args.dim * len(aux)) + if aux else None + ) if aux: - self._last_aux_positions = pos.reshape(-1) - self._last_aux_rows = rows.reshape(-1) + table_rows = get_global_ctx().batch.active_table_idx + if table_rows is None: + raise RuntimeError("DSpark decode features require request table rows") + # clone() is captured: replay updates these output buffers without leaving + # them aliased to the graph input buffers that the next batch overwrites. + self._target_features = DSparkTargetFeatures( + aux_hidden, + pos.reshape(-1).clone(), + table_rows[:B].reshape(-1).clone(), + ) + else: + self._target_features = None h = self.hc_head(h) h = self.norm(h) return self.logits(h[:, -1]) + def verify_block( + self, input_ids: torch.Tensor, pos: torch.Tensor, span: int + ) -> torch.Tensor: + """Run one fixed-shape DSpark target verification inside a CUDA graph. + + ``input_ids`` is request-major ``[1, R * span]`` and every request contributes + exactly anchor + gamma rows. The attention backend supplies one persistent + whole-history snapshot per request; ``rows`` maps every token to that request. + """ + total = input_ids.numel() + if span < 1 or total % span: + raise ValueError(f"invalid DSpark verify shape: {total} rows, span={span}") + num_reqs = total // span + rows = torch.arange(num_reqs, device=input_ids.device).repeat_interleave(span) + md = get_global_ctx().batch.attn_metadata + cmp_stage_cap = md.stage_width - 1 + wctx = md.window_ctx(pos, rows) + + h = self.embed_tokens(input_ids) + h = h.unsqueeze(2).repeat(1, 1, self.hc_mult, 1) + aux: list[torch.Tensor] = [] + for i, layer in enumerate(self.layers): + h = layer.verify_block( + h, pos, rows, cmp_stage_cap, input_ids, + num_reqs, span, wctx, + ) + if i in self._aux_layer_ids: + aux.append(h.mean(dim=2)) + + aux_hidden = ( + torch.cat(aux, dim=-1).view(-1, self.args.dim * len(aux)) + if aux else None + ) + if aux: + table_rows = get_global_ctx().batch.active_table_idx + if table_rows is None or table_rows.numel() < num_reqs: + raise RuntimeError("DSpark verify features require one table row per request") + self._target_features = DSparkTargetFeatures( + aux_hidden, + pos.reshape(-1).clone(), + table_rows[:num_reqs].repeat_interleave(span).clone(), + ) + else: + self._target_features = None + h = self.norm(self.hc_head(h)) + return self.logits(h.view(-1, self.args.dim)) + class DeepseekV4ForCausalLM(BaseLLMModel): """Engine adapter: a registered :class:`BaseLLMModel` wrapping the DSV4 transformer. @@ -372,6 +437,9 @@ def __init__(self, config): self._args: DeepseekV4Args = config.dsv4_args self._transformer = Transformer(self._args) self._bound = False + self.speculative_verify_block_size = ( + self._args.dspark_block_size if self._args.n_draft_layers else 0 + ) def _ensure_bound(self) -> None: if self._bound: @@ -418,45 +486,42 @@ def load_state_dict(self, state_dict, *, prefix: str = "", _internal: bool = Fal raise RuntimeError(f"Unexpected keys in state_dict: {list(state_dict.keys())}") self._transformer.load_state_dict(casted, assign=True, strict=False) - def catch_up_draft_context(self, batch) -> None: - """Write the draft layers' KV for the positions the target has just committed. - - The drafter never runs the prompt; its sliding window fills in behind the target, - derived from the target's own hidden state at those positions. + def dspark_target_features(self) -> DSparkTargetFeatures | None: + """Outputs of the most recent target forward, as one addressed bundle.""" + return self._transformer.target_features() - 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 + def catch_up_draft_context(self, features: DSparkTargetFeatures) -> None: + """Precompute draft-layer context KV from one target feature bundle.""" drafter = self._transformer.drafter if drafter is None: return - aux = self._transformer.last_aux_hidden() - addressing = self._transformer.last_aux_addressing() - if aux is None or addressing is None: - return # nothing has run yet; the first block attends over a short window - positions, rows = addressing - flat = aux.view(-1, aux.shape[-1]) - if flat.shape[0] != positions.numel(): - # The tap and its addressing are written together, so this cannot drift -- - # unless a new forward path records one and not the other. + flat = features.hidden + positions = features.positions + rows = features.table_rows + if flat.shape[0] != positions.numel() or rows.numel() != positions.numel(): 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" + "DSpark target features must have one hidden, position, and table row " + f"per token; got {flat.shape[0]}, {positions.numel()}, {rows.numel()}" ) backend = get_global_ctx().attn_backend slots = backend.window_slots_at(rows, positions) + valid = slots >= 0 + flat = flat[valid] + positions = positions[valid] + slots = slots[valid] + if positions.numel() == 0: + return + if _DRAFT_CTX_DEBUG: + logger.debug_rank0( + "draft ctx: %d rows, positions %s..%s, slots %s..%s", + flat.shape[0], int(positions.min()), int(positions.max()), + int(slots.min()), int(slots.max()), + ) drafter.catch_up_context(flat, positions, slots) - def draft(self) -> tuple[torch.Tensor, torch.Tensor] | None: + def draft( + self, sampling_params + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None: """Propose the current batch's block with the dSpark drafter. Runs over the SAME prepared batch as the verify pass that follows -- the same @@ -465,24 +530,19 @@ def draft(self) -> tuple[torch.Tensor, torch.Tensor] | None: 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. + Returns ``(proposed [R*gamma], q [R*gamma,V], confidence [R*gamma])``. """ 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, + sampling_params, ) def forward(self) -> torch.Tensor: @@ -490,6 +550,11 @@ def forward(self) -> torch.Tensor: batch = get_global_ctx().batch input_ids = batch.input_ids.long() md = batch.attn_metadata + if getattr(batch, "spec_verify_decode", False): + span = int(batch.spec_block) + 1 + return self._transformer.verify_block( + input_ids.view(1, -1), batch.positions.long(), span + ) if batch.is_prefill: # Ragged batched prefill (bs >= 1): each request starts from its own cached_len. # A cold segment (start_pos == 0) re-seeds the compressor carry register inside its @@ -525,4 +590,4 @@ def forward(self) -> torch.Tensor: return self._transformer.decode(input_ids.view(B, 1), pos, cmp_stage_cap) -__all__ = ["Transformer", "DeepseekV4ForCausalLM"] +__all__ = ["DSparkTargetFeatures", "Transformer", "DeepseekV4ForCausalLM"] diff --git a/python/freetoken/models/deepseek_v4/moe.py b/python/freetoken/models/deepseek_v4/moe.py index fa29154..bb4617c 100644 --- a/python/freetoken/models/deepseek_v4/moe.py +++ b/python/freetoken/models/deepseek_v4/moe.py @@ -15,6 +15,7 @@ from .args import DeepseekV4Args from .layers import Linear + from .parallel import div_tp, tp_size @@ -138,8 +139,9 @@ def _prefill_routed( # # 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" + 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 diff --git a/python/freetoken/models/deepseek_v4/rollback.py b/python/freetoken/models/deepseek_v4/rollback.py deleted file mode 100644 index d90d762..0000000 --- a/python/freetoken/models/deepseek_v4/rollback.py +++ /dev/null @@ -1,104 +0,0 @@ -"""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/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index a699623..a28b561 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -123,10 +123,6 @@ class OffloadMoeCache: # pcie_bw / cpu_bw ratio so the PCIe fetch and the CPU overflow GEMV take equal # time (perfect overlap): fetched : cpu = pcie : cpu - pcie. hybrid_fetch_fraction: float = 0.0 - # 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} assert self.cache_policy in policy_ids @@ -794,36 +790,18 @@ def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None ) 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. + """Return Equation 4's measured split, unchanged for every exact MoE layer. - 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. + FreeToken defines ``m`` as the number of missing experts in the current layer + and chooses ``q* = m * B_PCIe / B_Host``. ``ensure_experts_hybrid`` already + deduplicates the routes of the whole token batch to obtain that ``m`` before it + applies this fraction. Scaling the ratio again by the token count double-counts + the wider batch and can collapse a speculative verify into pure PCIe fetching. - 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. + ``layer_id`` and ``expert_ids`` remain arguments because the cache-control call + site has this interface; neither changes the hardware bandwidth ratio. """ - # 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) + return self.hybrid_fetch_fraction def materialize_layer(self, layer_id: int) -> None: from freetoken.moe.offload_kernels import materialize_layer diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index e385df2..b7a7eb8 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -227,6 +227,33 @@ def overlap_loop(self, last_data: ForwardData | None) -> ForwardData | None: # table_idx can have its freshly copied prompt clobbered by the prior occupant's # still-pending output write -- corrupting tokens (e.g. dropping an image # placeholder, which the multimodal merge then rejects). + # Speculation cannot be overlapped as written. The loop below schedules and + # LAUNCHES batch N+1 before _process_last_data applies batch N's tokens. Plain + # decode survives that because complete_one() pre-advances device_len by exactly + # one -- the token count is known before the forward runs, which is why + # device_len legitimately leads the buffer by one. + # + # A speculative block contributes 1..k+1 tokens and the count is only known + # AFTER its verify. Overlapped, N+1 is scheduled as though nothing was accepted, + # so it re-drafts positions N already consumed. Measured directly: + # + # draft ctx: 5 rows, positions 18..22 <- block N + # draft ctx: 5 rows, positions 18..22 <- N+1, same positions, same proposals + # draft ctx: 5 rows, positions 20..24 <- then a jump of two + # + # which duplicates tokens in the reply and makes the acceptance figure + # meaningless. + # + # vLLM overlaps and corrects for this by threading num_rejected into the next + # step's input preparation (valid_ctx_end = ctx_end - num_rejected in its + # dflash kernel). Without that correction the only correct option is to drain + # first, which costs the overlap but not the answer. + if self._speculative is not None and last_data is not None: + self.stream.wait_stream(self.engine.stream) + self._process_last_data(last_data) + self._flush_abort_acks() + last_data = None + self.stream.wait_stream(self.engine.stream) forward_input = self._schedule_next_batch() ongoing_data = None @@ -245,8 +272,10 @@ def overlap_loop(self, last_data: ForwardData | None) -> ForwardData | None: # full_to_window INSIDE the captured graph, so an unordered drain can redirect an # in-flight forward. copy_done only covers batch N; order against N+1 explicitly. self.stream.wait_stream(self.engine.stream) - self._process_last_data(last_data) - self._flush_abort_acks() + # None when the speculative drain above already handled it. + if last_data is not None: + self._process_last_data(last_data) + self._flush_abort_acks() return ongoing_data def normal_loop(self) -> None: @@ -359,14 +388,16 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: if reply_tokens is not None else [next_token] ) + block_start = req.input_ids.numel() - len(block) finished = False for pos, tok in enumerate(block): + visible_len = block_start + pos + 1 hit_length = not req.can_decode and pos == len(block) - 1 hit_eos = ( not req.sampling_params.ignore_eos and tok in self.eos_token_ids ) matched_stop = ( - self._match_stop_str(req) + self._match_stop_str(req, end_len=visible_len) if not hit_eos and req.sampling_params.stop_strs else None ) @@ -381,7 +412,7 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: and req.toolcall_anchor_len is None and not finished ): - req.toolcall_anchor_len = req.input_ids.numel() + req.toolcall_anchor_len = visible_len reply.append( DetokenizeMsg( uid=req.uid, @@ -442,17 +473,20 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: ) self.send_result(reply) - def _match_stop_str(self, req: Req) -> str | None: + def _match_stop_str(self, req: Req, end_len: int | None = None) -> str | None: """First stop string present in this request's generated tail, else None. Decodes only a short suffix (bounded by the longest stop string's char length, so a stop of - N chars spans at most N tokens) to keep the per-step cost small.""" + N chars spans at most N tokens) to keep the per-step cost small. ``end_len`` + makes a speculative block visible one token at a time: looking at the already- + appended whole block would let a stop in token 5 terminate the reply at token 1.""" stop_strs = req.sampling_params.stop_strs prompt_len = req.max_device_len - req.output_len - if len(req.input_ids) <= prompt_len: + end_len = len(req.input_ids) if end_len is None else end_len + if end_len <= prompt_len: return None max_chars = max(len(s) for s in stop_strs) - tail_start = max(prompt_len, len(req.input_ids) - (max_chars + 1)) - tail = self.tokenizer.decode(req.input_ids[tail_start:].tolist()) + tail_start = max(prompt_len, end_len - (max_chars + 1)) + tail = self.tokenizer.decode(req.input_ids[tail_start:end_len].tolist()) for s in stop_strs: if s in tail: return s @@ -831,10 +865,19 @@ def _prepare_batch(self, batch: Batch) -> ForwardInput: # built once here instead of rebuilt in each of the 30 GDN layers. For decode # under CUDA graph the persistent cu_seqlens buffer is supplied by set_batch. batch.fla_metadata = build_fla_metadata(batch, self.device) - if batch.is_decode: + if batch.is_decode or getattr(batch, "speculative", False): # This batch's padded per-row page-table rows. Backends that snapshot the table for # a captured replay (DSV4) read them in prepare_metadata / prepare_for_replay. - batch.active_table_idx = input_mapping[0].view(-1) + # A DSpark verify is phase="prefill" but its fixed-shape target graph uses + # the same device-addressed snapshot contract. + if getattr(batch, "speculative", False): + batch.active_table_idx = torch.tensor( + [r.table_idx for r in batch.padded_reqs], + dtype=torch.int64, + device=self.device, + ) + else: + batch.active_table_idx = input_mapping[0].view(-1) self.engine.attn_backend.prepare_metadata(batch) return ForwardInput( batch=batch, @@ -888,21 +931,24 @@ def _maybe_make_speculative(self, batch: Batch) -> None: self._warned_non_greedy = False if spec is None or not batch.is_decode: return - k = spec.block_size - 1 - if k < 1: + gamma = spec.block_size + if gamma < 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: + # Room for gamma proposals AND the bonus token; equality fits exactly. + # A request with less room decodes normally rather than being truncated. + if req.remain_len < gamma + 1: return - noise = torch.full((k,), spec.noise_token_id, dtype=self.token_pool.dtype) + # The target verify has anchor + gamma proposal rows. The drafter takes the + # first gamma of those rows -- anchor + gamma-1 masks -- exactly as DSpark's + # sample-from-anchor layout specifies. + noise = torch.full((gamma,), spec.noise_token_id, dtype=self.token_pool.dtype) for req in batch.reqs: req.append_host(noise) - req.device_len += k + req.device_len += gamma batch.phase = "prefill" batch.speculative = True - batch.spec_block = k + batch.spec_block = gamma # device_len now covers the whole block, so allocate_paged will size for it. # Acceptance keeps only a prefix and must give the rest back through this. batch.release_tail = self.cache_manager.release_speculative_tail @@ -940,8 +986,27 @@ def _forward(self, forward_input: ForwardInput) -> ForwardOutput: # replaces this batch's placeholder token ids in place, so the verify that # follows scores the drafter's actual proposal. batch.draft_confidence = self.engine.draft_into_batch(batch) + self.engine.adapt_speculative_batch(batch) forward_output = self.engine.forward_batch(batch, sample_args) - self.token_pool[output_mapping] = forward_output.next_tokens_gpu + if getattr(batch, "speculative", False): + # output_mapping was prepared before acceptance, when device_len still + # pointed past the whole gamma-wide verify span. A rejection moves each + # request's live bonus position independently; writing through the stale + # mapping leaves noise at the next cycle's anchor. Rebuild the tiny mapping + # from the post-acceptance frontiers. + rows = torch.tensor( + [req.table_idx for req in batch.reqs], + dtype=torch.int64, + device=self.device, + ) + positions = torch.tensor( + [req.device_len - 1 for req in batch.reqs], + dtype=torch.int64, + device=self.device, + ) + self.token_pool[rows, positions] = forward_output.next_tokens_gpu + else: + self.token_pool[output_mapping] = forward_output.next_tokens_gpu self.decode_manager.filter_reqs(forward_input.batch.reqs) return forward_output @@ -949,7 +1014,6 @@ def _forward(self, forward_input: ForwardInput) -> ForwardOutput: class _SpeculativeConfig(NamedTuple): block_size: int noise_token_id: int - confidence_threshold: float def _speculative_config(config) -> "_SpeculativeConfig | None": @@ -967,7 +1031,6 @@ def _speculative_config(config) -> "_SpeculativeConfig | 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), ) diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index 5398d39..42d04e8 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -365,6 +365,30 @@ async def stream_with_cancellation(self, generator, request: Request, uid: int): asyncio.create_task(self.abort_user(uid)) raise + async def wait_for_ack_watching_client(self, uid: int, request): + """``wait_for_ack``, but abort the generation when the client goes away. + + Streaming responses already notice a disconnect, because yielding a chunk into a + dead socket is observable. A non-streaming request awaits the whole completion + and never looks, so a client that hangs up leaves the GPUs generating to + max_tokens for nobody -- and the queue behind it waiting on work whose answer + will be discarded. + + ``request`` may be None for callers with no HTTP request in scope (offline, + internal); then this is exactly ``wait_for_ack``. + """ + try: + async for ack in self.wait_for_ack(uid): + if request is not None and await request.is_disconnected(): + logger.info("Client disconnected for user %s; aborting", uid) + raise asyncio.CancelledError + yield ack + except asyncio.CancelledError: + # create_task, not await: the caller is unwinding, and abort_user sleeps + # before it sends. Mirrors stream_with_cancellation. + asyncio.create_task(self.abort_user(uid)) + raise + async def abort_user(self, uid: int): await asyncio.sleep(0.1) if uid in self.ack_map: diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index be05d90..a55ad22 100644 --- a/python/freetoken/server/generation.py +++ b/python/freetoken/server/generation.py @@ -490,6 +490,19 @@ def _record_generation( ) +def _acks(state, uid, request): + """The ack stream, watching the client when both sides support it. + + getattr, not a bare call: a caller may pass a state that predates the watcher (the + test fakes do), and a missing disconnect check should cost the abort optimisation, + not the request. + """ + watch = getattr(state, "wait_for_ack_watching_client", None) + if request is not None and watch is not None: + return watch(uid, request) + return state.wait_for_ack(uid) + + async def generate_events( uid: int, spec: GenSpec, state: Any, *, source: str | None = None ) -> AsyncIterator[GenEvent]: @@ -521,7 +534,7 @@ async def generate_events( async def generate_full( - uid: int, spec: GenSpec, state: Any, *, source: str | None = None + uid: int, spec: GenSpec, state: Any, *, source: str | None = None, request: Any = None ) -> GenResult: """Wraps `_generate_full_impl` to log the request with its totals; the `finally` also records a `GenerationError` as a failed row.""" @@ -529,7 +542,7 @@ async def generate_full( result: GenResult | None = None error: str | None = None try: - result = await _generate_full_impl(uid, spec, state) + result = await _generate_full_impl(uid, spec, state, request) return result except GenerationError as exc: error = str(exc) @@ -647,6 +660,8 @@ def _route_tool_text(piece: str) -> list[GenEvent]: engine_finish_reason: str | None = None engine_matched_stop: str | None = None + # Streaming: a disconnect is already caught at the transport layer by + # stream_with_cancellation, which sees the write into a dead socket. async for ack in state.wait_for_ack(uid): if getattr(ack, "error", None): raise GenerationError(ack.error, getattr(ack, "error_code", None)) @@ -738,7 +753,9 @@ def _route_tool_text(piece: str) -> list[GenEvent]: ) -async def _generate_full_impl(uid: int, spec: GenSpec, state: Any) -> GenResult: +async def _generate_full_impl( + uid: int, spec: GenSpec, state: Any, request: Any = None +) -> GenResult: """Protocol-neutral non-streaming generation: accumulate, split reasoning, parse tool calls, strip special tokens. The adapters format the GenResult into their wire.""" full_content = "" @@ -747,7 +764,10 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any) -> GenResult: cached_tokens = 0 engine_finish_reason: str | None = None engine_matched_stop: str | None = None - async for ack in state.wait_for_ack(uid): + # Watch the client. A non-streaming request awaits the whole completion, so a + # caller that hangs up otherwise leaves the GPUs generating to max_tokens for an + # answer nobody will read -- with the queue behind it waiting on that work. + async for ack in _acks(state, uid, request): if getattr(ack, "error", None): raise GenerationError(ack.error, getattr(ack, "error_code", None)) prompt_tokens += ack.prompt_tokens_delta diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index b4becd2..a7dbc40 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -23,6 +23,7 @@ from .function_call_parser import ToolCallItem from .request_logger import log_request from .generation import ( + _acks, ContentDelta, GenDone, GenerationError, @@ -199,7 +200,9 @@ async def handle_chat_completion( return StreamingResponse(chunks, media_type="text/event-stream") try: - result = await generate_full(uid, spec, state, source="/v1/chat/completions") + result = await generate_full( + uid, spec, state, source="/v1/chat/completions", request=request + ) except GenerationError as exc: return create_error_response(str(exc), code=exc.code) message: dict[str, Any] = {"role": "assistant", "content": result.content} @@ -413,7 +416,9 @@ async def handle_completion( await state.send_one(TokenizeMsg(uid=uid, text=prompt, sampling_params=_resolve_sampling(req, model_sampling))) text = "" finish_reason = "stop" - async for ack in state.wait_for_ack(uid): + # Watch the client: a non-streaming request otherwise generates to max_tokens + # for a caller that has already hung up. + async for ack in _acks(state, uid, request): if getattr(ack, "error", None): return create_error_response(ack.error) prompt_tokens += ack.prompt_tokens_delta diff --git a/scripts/freetoken-dsv4.sh b/scripts/freetoken-dsv4.sh index 1f773b2..b648180 100644 --- a/scripts/freetoken-dsv4.sh +++ b/scripts/freetoken-dsv4.sh @@ -28,6 +28,19 @@ EXPERT_LOAD="${EXPERT_LOAD:-serial}" # 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}" +# Which base dspark_target_layer_ids is read with (0 or 1). Exported so it reaches the +# rank subprocesses; unset means the model's own default. See the note in model.py -- +# this is settled by measurement, not by the config. +[ -n "${FREETOKEN_DSPARK_LAYER_BASE:-}" ] && export FREETOKEN_DSPARK_LAYER_BASE +# FREETOKEN_DSPARK_MARKOV=0 drops the Markov bias from the draft logits. +[ -n "${FREETOKEN_DSPARK_MARKOV:-}" ] && export FREETOKEN_DSPARK_MARKOV +# FREETOKEN_SPEC_HYBRID_MOE=0 sends speculative verifies down the prefill expert path. +[ -n "${FREETOKEN_SPEC_HYBRID_MOE:-}" ] && export FREETOKEN_SPEC_HYBRID_MOE +# FREETOKEN_SPEC_DEBUG=1 dumps the first few blocks' token flow. +[ -n "${FREETOKEN_SPEC_DEBUG:-}" ] && export FREETOKEN_SPEC_DEBUG +[ -n "${FREETOKEN_SPEC_FORCE_REJECT:-}" ] && export FREETOKEN_SPEC_FORCE_REJECT +[ -n "${FREETOKEN_CMP_FIRST_ONLY:-}" ] && export FREETOKEN_CMP_FIRST_ONLY +[ -n "${FREETOKEN_SPEC_WINDOW_ONLY:-}" ] && export FREETOKEN_SPEC_WINDOW_ONLY LOG="${LOG:-/tmp/freetoken-dsv4.log}" # The TP ranks' torch.distributed rendezvous. Held by every rank, not just the # frontend, so it is the one that lingers after a stop. diff --git a/tests/dsv4/test_dsv4_adaptive_verification.py b/tests/dsv4/test_dsv4_adaptive_verification.py new file mode 100644 index 0000000..cbe56c9 --- /dev/null +++ b/tests/dsv4/test_dsv4_adaptive_verification.py @@ -0,0 +1,124 @@ +"""Paper-faithful adaptive verification tests that do not require CUDA.""" + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.engine.engine import Engine +from freetoken.engine.graph import SharedSpecCarryJournal +from freetoken.models.deepseek_v4.dspark import choose_adaptive_draft_width + + +def test_profile_cliff_stops_before_expensive_graph_bucket(): + width = choose_adaptive_draft_width( + [0.9, 0.9], + draft_cost_ms=0.0, + verify_cost_ms=[1.0, 1.0, 100.0], + ) + assert width == 1 + + +def test_flat_verify_cost_admits_the_whole_prefix(): + width = choose_adaptive_draft_width( + [0.8, 0.7, 0.6], + draft_cost_ms=2.0, + verify_cost_ms=[10.0, 10.0, 10.0, 10.0], + ) + assert width == 3 + + +def test_selector_uses_cumulative_survival_not_raw_confidence(): + # The second token's contribution is 0.5 * 0.9, not 0.9 by itself. + width = choose_adaptive_draft_width( + [0.5, 0.9], + draft_cost_ms=0.0, + verify_cost_ms=[1.0, 1.0, 1.4], + ) + assert width == 1 + + +def test_tie_keeps_the_smaller_prefix_like_argmax(): + assert choose_adaptive_draft_width( + [1.0], draft_cost_ms=0.0, verify_cost_ms=[1.0, 2.0] + ) == 0 + + +def test_cost_curve_must_cover_every_prefix(): + with pytest.raises(ValueError, match="every width"): + choose_adaptive_draft_width( + [0.9, 0.9], draft_cost_ms=1.0, verify_cost_ms=[1.0, 2.0] + ) + + +def test_prefix_graphs_share_the_maximum_carry_journal_storage(): + key = (3, "compress", 128) + storage = {key: [torch.zeros(1) for _ in range(6)]} + journal = SharedSpecCarryJournal(storage) + journal.reset(2) + journal.record(key, torch.tensor([4.0])) + journal.record(key, torch.tensor([5.0])) + + [(seen_key, pieces)] = list(journal.items()) + assert seen_key == key + assert len(pieces) == 2 + assert torch.cat(pieces).tolist() == [4.0, 5.0] + assert storage[key][2].item() == 0.0 + + +def test_shared_carry_journal_rejects_prefix_overflow(): + key = (3, "compress", 128) + journal = SharedSpecCarryJournal({key: [torch.zeros(1)]}) + journal.reset(1) + journal.record(key, torch.ones(1)) + with pytest.raises(RuntimeError, match="overflow"): + journal.record(key, torch.ones(1)) + + +class _ChooseTwo: + block_size = 5 + + def record_and_choose(self, confidence, uid): + assert confidence.tolist() == pytest.approx([0.9] * 5) + assert uid == 7 + return 2 + + +def test_engine_compacts_only_target_views_and_keeps_full_allocation_frontier(): + engine = Engine.__new__(Engine) + engine._adaptive_verification = _ChooseTwo() + ids = torch.arange(20, dtype=torch.int32) + req = SimpleNamespace(input_ids=ids[:15], _ids_buf=ids, uid=7, device_len=15) + metadata = SimpleNamespace(segments=[(0, 6, 3, 9)]) + batch = SimpleNamespace( + reqs=[req], + padded_size=1, + spec_block=5, + draft_confidence=torch.full((5,), 0.9), + draft_tokens=torch.arange(5, dtype=torch.int32), + draft_probs=torch.empty(5, 4), + input_ids=torch.arange(6, dtype=torch.int32), + positions=torch.arange(9, 15, dtype=torch.int32), + out_loc=torch.arange(6, dtype=torch.int32), + attn_metadata=metadata, + ) + + engine.adapt_speculative_batch(batch) + + assert batch.spec_block == 2 + assert batch.input_ids.tolist() == [0, 1, 2] + assert batch.positions.tolist() == [9, 10, 11] + assert batch.out_loc.tolist() == [0, 1, 2] + assert batch.draft_tokens.tolist() == [0, 1] + assert batch.draft_confidence.tolist() == pytest.approx([0.9, 0.9]) + assert metadata.segments == [(0, 3, 3, 9)] + assert req.input_ids.numel() == 12 + assert req.device_len == 15, "tail release still needs the gamma-wide frontier" + + +def test_greedy_verify_reduces_logits_before_crossing_pcie(): + import inspect + + body = inspect.getsource(Engine._finish_speculative) + assert 'logits.argmax(dim=-1).to(torch.int32).to("cpu"' in body + assert 'logits.to("cpu"' not in body diff --git a/tests/dsv4/test_dsv4_aux_addressing.py b/tests/dsv4/test_dsv4_aux_addressing.py index 4c4d4b2..5073829 100644 --- a/tests/dsv4/test_dsv4_aux_addressing.py +++ b/tests/dsv4/test_dsv4_aux_addressing.py @@ -33,31 +33,39 @@ def _fn(name: str) -> str: 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_prefill_and_decode_build_one_feature_bundle(self): + assert SRC.count("DSparkTargetFeatures(") >= 2 def test_addressing_is_exposed_as_one_unit(self): - from freetoken.models.deepseek_v4.model import Transformer + from freetoken.models.deepseek_v4.model import DSparkTargetFeatures - assert hasattr(Transformer, "last_aux_addressing"), ( - "positions and rows must be readable together with the tap" - ) + assert set(DSparkTargetFeatures.__dataclass_fields__) == { + "hidden", "positions", "table_rows" + } def test_addressing_returns_none_until_a_forward_has_run(self): from freetoken.models.deepseek_v4.model import Transformer t = Transformer.__new__(Transformer) - t._last_aux_hidden = None - t._last_aux_positions = None - t._last_aux_rows = None - assert t.last_aux_addressing() is None + t._target_features = None + assert t.target_features() is None + + def test_tap_reshape_reads_the_dimension_from_model_args(self): + # Transformer has no direct ``dim`` field. This typo survives source-only + # wiring tests and crashes only when the first real aux tap runs. + assert "self.dim * len(aux)" not in SRC + assert SRC.count("self.args.dim * len(aux)") == 3 + + def test_cuda_graph_features_are_owned_per_captured_batch_size(self): + graph_src = ( + pathlib.Path(__file__).resolve().parents[2] + / "python" / "freetoken" / "engine" / "graph.py" + ).read_text() + assert "self._dspark_feature_map[bs] = features" in graph_src + assert "self._dspark_feature_map.get(batch.padded_size)" in graph_src + assert "self._spec_feature_map[key] = features" in graph_src + assert "self._spec_carry_map[key]" in graph_src + assert "shared_carry if shared_carry is not None" in graph_src class TestTheCatchUpUsesTheTapsOwnPositions: @@ -69,8 +77,10 @@ def test_it_does_not_read_the_batch_metadata(self): "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_it_takes_one_explicit_feature_bundle(self): + body = _fn("catch_up_draft_context") + for field in ("features.hidden", "features.positions", "features.table_rows"): + assert field in body def test_a_row_count_mismatch_is_an_error_not_a_silent_skip(self): # Truncating to the shorter of the two would write SOME rows at wrong positions, diff --git a/tests/dsv4/test_dsv4_aux_layer_ids.py b/tests/dsv4/test_dsv4_aux_layer_ids.py index 01cbf6a..cd9b469 100644 --- a/tests/dsv4/test_dsv4_aux_layer_ids.py +++ b/tests/dsv4/test_dsv4_aux_layer_ids.py @@ -48,45 +48,19 @@ def test_no_ids_is_not_an_error(self): 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() +class TestTheReferenceMappingIsFixed: + def test_the_resolved_layers_are_logged(self): + assert "dSpark: tapping target layer outputs" in _model_src() - def test_only_zero_or_one_is_accepted(self): - assert "base not in (0, 1)" in _model_src() + def test_no_runtime_base_override_remains(self): + assert "FREETOKEN_DSPARK_LAYER_BASE" not 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 test_vllm_mapping_selects_the_last_three_outputs(self): + # vLLM adds one to the config ids, then captures when idx+1 is selected. + config = (40, 41, 42) + aux_ids = tuple(i + 1 for i in config) + captured = tuple(i for i in range(43) if i + 1 in aux_ids) + assert captured == config def _model_src() -> str: diff --git a/tests/dsv4/test_dsv4_dspark.py b/tests/dsv4/test_dsv4_dspark.py index 50c0329..0767bec 100644 --- a/tests/dsv4/test_dsv4_dspark.py +++ b/tests/dsv4/test_dsv4_dspark.py @@ -1,4 +1,4 @@ -"""dSpark acceptance and adaptive width. +"""DSpark greedy prefix acceptance. These two functions decide what the model EMITS. A wrong acceptance rule does not crash or slow anything down -- it silently changes the output distribution, which is @@ -8,10 +8,9 @@ from __future__ import annotations -import pytest import torch -from freetoken.models.deepseek_v4.dspark import accepted_prefix, draft_width +from freetoken.models.deepseek_v4.dspark import accepted_prefix def _t(*xs: int) -> torch.Tensor: @@ -42,24 +41,3 @@ def test_a_wholly_rejected_block_still_advances_by_one(self): 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 index 0739189..16a5dc5 100644 --- a/tests/dsv4/test_dsv4_dspark_heads.py +++ b/tests/dsv4/test_dsv4_dspark_heads.py @@ -75,8 +75,7 @@ def test_it_scores_one_value_per_position(self): 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. + # Equation 7 defines a probability consumed by the load-aware scheduler. head = ConfidenceHead(_args()) torch.nn.init.normal_(head.proj, std=10.0) # push the pre-sigmoid range wide out = head(torch.randn(16, DIM), torch.randn(16, RANK)) @@ -89,31 +88,64 @@ def test_it_consumes_both_the_hidden_and_the_markov_embedding(self): ) -class TestLogitsUnderTensorParallelism: - """The bug the review caught: the drafter must not use the head weight directly.""" +class TestSequentialMarkovStage: + def test_each_step_reads_the_token_sampled_by_the_previous_step(self): + from freetoken.models.deepseek_v4.dspark import DSparkDrafter + + seen = [] + + class _Markov(torch.nn.Module): + def embed(self, ids): + seen.append(ids.clone()) + return torch.zeros(ids.shape[0], RANK) + + def bias(self, embeds): + return torch.zeros(embeds.shape[0], VOCAB) + + class _Confidence(torch.nn.Module): + def forward(self, hidden, embeds): + return torch.ones(hidden.shape[0]) + + class _Params: + temperature = 0.0 + top_p = 1.0 + top_k = -1 + is_greedy = True + + drafter = DSparkDrafter.__new__(DSparkDrafter) + torch.nn.Module.__init__(drafter) + drafter.markov_head = _Markov() + drafter.confidence_head = _Confidence() + base = torch.full((1, 3, VOCAB), -10.0) + base[0, 0, 2] = 10.0 + base[0, 1, 3] = 10.0 + base[0, 2, 4] = 10.0 + proposed, q, confidence = drafter.sample_block( + base, torch.zeros(1, 3, DIM), torch.tensor([1]), [_Params()] + ) + assert proposed.tolist() == [2, 3, 4] + assert [x.tolist() for x in seen] == [[1], [2], [3]] + assert q.shape == (3, VOCAB) + assert confidence.shape == (3,) + + +class TestDraftHeadNormalization: + def test_backbone_returns_pre_norm_hidden_for_the_confidence_head(self): + import inspect - 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) + body = inspect.getsource(DSparkDrafter.head_hidden) + assert "return self.hc_head(h)" in body + assert "return self.norm(self.hc_head(h))" not in body - # 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) + def test_only_base_logits_apply_the_final_norm(self): + import inspect - 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,) + from freetoken.models.deepseek_v4.dspark import DSparkDrafter + + body = inspect.getsource(DSparkDrafter.propose) + assert "target_logits_fn(self.norm(head_hidden))" in body class TestContextKvPositions: @@ -167,3 +199,37 @@ def test_a_checkpoint_without_a_noise_token_is_refused(self): d, _ = self._drafter(dspark_noise_token_id=-1) with pytest.raises(ValueError, match="noise"): d.block_input_ids(1, torch.device("cpu")) + + def test_propose_does_not_trust_uninitialized_gpu_placeholder_rows(self): + from freetoken.models.deepseek_v4.dspark import DSparkDrafter + + d = DSparkDrafter.__new__(DSparkDrafter) + torch.nn.Module.__init__(d) + d.block_size = 3 + d.noise_token_id = 99 + d.norm = torch.nn.Identity() + seen = {} + d.embed_block = lambda ids: torch.zeros(1, ids.numel(), DIM) + + def head_hidden(_embeds, ids, _segments, _positions): + seen["ids"] = ids.clone() + return torch.zeros(1, ids.numel(), DIM) + + d.head_hidden = head_hidden + d.sample_block = lambda *_args: ( + torch.zeros(6, dtype=torch.long), + torch.zeros(6, VOCAB), + torch.zeros(6), + ) + # Each target span is anchor + gamma proposal slots. The proposal slots here + # deliberately contain stale values, exactly like an unwritten GPU token-pool row. + target_ids = torch.tensor([[11, 7, 8, 9, 22, 70, 80, 90]]) + segments = [(0, 4, 1, 10), (4, 4, 2, 20)] + d.propose( + target_ids, + segments, + torch.arange(8), + lambda h: torch.zeros(h.shape[0], VOCAB), + [object(), object()], + ) + assert seen["ids"].tolist() == [[11, 99, 99, 22, 99, 99]] diff --git a/tests/dsv4/test_dsv4_dspark_masking.py b/tests/dsv4/test_dsv4_dspark_masking.py index b683e51..ab0e424 100644 --- a/tests/dsv4/test_dsv4_dspark_masking.py +++ b/tests/dsv4/test_dsv4_dspark_masking.py @@ -15,6 +15,8 @@ from __future__ import annotations +import inspect + import pytest import torch @@ -51,7 +53,7 @@ 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 + w_lo = start - WIN last_col = (start + n - 1) - w_lo for row in range(n): assert last_col in _visible(cols, row), ( @@ -74,28 +76,36 @@ def test_it_differs_from_the_causal_mask(self): 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 + w_lo = start - WIN highest = max(max(_visible(cols, r)) for r in range(n)) + w_lo assert highest == start + n - 1, "a query must not read past its own block" def test_the_shared_context_stays_visible(self): n, start = 5, 500 cols = window_cols_for_block(start_pos=start, n=n, window=WIN, non_causal=True) - w_lo = start - WIN + 1 + w_lo = start - WIN lowest = min(min(_visible(cols, r)) for r in range(n)) + w_lo assert lowest < start, "the block must still attend over the context before it" @pytest.mark.parametrize("n", [1, 2, 5, 8]) - def test_the_candidate_width_is_the_window_for_any_block_size(self, n): + def test_the_candidate_width_is_context_window_plus_block(self, n): cols = window_cols_for_block(start_pos=900, n=n, window=WIN, non_causal=True) - assert cols.shape == (n, WIN) + assert cols.shape == (n, WIN + n) - 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. + def test_a_single_token_still_keeps_all_context_tokens(self): a = window_cols_for_block(start_pos=700, n=1, window=WIN, non_causal=True) b = window_cols_for_block(start_pos=700, n=1, window=WIN, non_causal=False) - assert torch.equal(a, b) + assert a.shape[-1] == WIN + 1 + assert b.shape[-1] == WIN + + def test_ragged_launcher_labels_the_full_non_causal_width(self): + # The index helper can correctly build 133 columns while the attention call + # still labels only 128 of them as window columns. That only fails at bs > 1. + from freetoken.models.deepseek_v4.attention import Attention + + body = inspect.getsource(Attention.forward_ragged) + assert "max(part.shape[-1] for part in win_parts)" in body + assert "if self.non_causal" in body class TestMaskBounds: @@ -105,7 +115,10 @@ class TestMaskBounds: @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) + w_lo = max( + 0, + start_pos - WIN if non_causal else start_pos - WIN + 1, + ) width = (start_pos + n) - w_lo # slots the caller gathered for this segment assert int(cols.max()) < width, "a column past the gathered slots reads foreign KV" assert int(cols.min()) >= -1, "only -1 encodes 'masked'" diff --git a/tests/dsv4/test_dsv4_dspark_wiring.py b/tests/dsv4/test_dsv4_dspark_wiring.py index d62b92c..825d9d5 100644 --- a/tests/dsv4/test_dsv4_dspark_wiring.py +++ b/tests/dsv4/test_dsv4_dspark_wiring.py @@ -6,9 +6,8 @@ * ``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. +* Per-token compressor states must be selected after acceptance; restoring a single + pre-block snapshot loses any accepted prefix in the same page. * ``catch_up_context`` was written to keep the draft layers' window fresh, and never called, which reads as a weak drafter rather than a missing call. @@ -20,6 +19,7 @@ from __future__ import annotations import ast +import inspect import pathlib import pytest @@ -64,13 +64,12 @@ def _calls_in_package() -> set[str]: ("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"), + ("rejection_accept_device", "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"), + ("_restore_speculative_carry", "acceptance would leave the rejected carry live"), + ("_trim_dspark_target_features", "rejected target rows would become draft context"), ("draft_into_batch", "the block would keep its noise placeholders"), ("_finish_speculative", "nothing would apply acceptance"), ("_maybe_make_speculative", "no batch would ever become speculative"), @@ -93,7 +92,8 @@ class TestSpeculativeStateIsCarried: @pytest.mark.parametrize( "field", ["speculative", "spec_block", "spec_emitted", "draft_probs", - "draft_confidence", "carry_snapshot"], + "draft_confidence", "draft_tokens", "spec_carry_states", + "spec_verify_decode"], ) def test_batch_declares_the_field(self, field): import torch @@ -115,7 +115,30 @@ def test_a_fresh_batch_is_not_speculative(self): assert b.speculative is False assert b.spec_block == 0 assert b.spec_emitted is None - assert b.carry_snapshot is None + assert b.draft_tokens is None + assert b.spec_carry_states is None + assert b.spec_verify_decode is False + + +class TestFixedVerifyGraphWiring: + def test_target_graph_is_selected_before_ordinary_decode_graph(self): + engine = (PKG / "engine" / "engine.py").read_text() + spec_at = engine.index("can_use_spec_cuda_graph") + decode_at = engine.index("can_use_cuda_graph", spec_at) + assert spec_at < decode_at + + def test_graph_replays_graph_owned_partial_states(self): + graph = (PKG / "engine" / "graph.py").read_text() + assert "journal = self._spec_carry_map[key]" in graph + assert "batch.spec_carry_states = journal" in graph + assert "prepare_for_spec_replay" in graph + + def test_verify_uses_device_positions_not_a_host_start_position(self): + from freetoken.models.deepseek_v4.model import Transformer + + body = inspect.getsource(Transformer.verify_block) + assert "pos: torch.Tensor" in body + assert "start_pos" not in body class TestVerifyReturnsEveryPosition: diff --git a/tests/dsv4/test_dsv4_hybrid_balance.py b/tests/dsv4/test_dsv4_hybrid_balance.py index 0a2a743..4b6d130 100644 --- a/tests/dsv4/test_dsv4_hybrid_balance.py +++ b/tests/dsv4/test_dsv4_hybrid_balance.py @@ -1,31 +1,20 @@ -"""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. -""" +"""The hybrid MoE split follows Equation 4 of the FreeToken paper.""" from __future__ import annotations +from types import SimpleNamespace + import pytest import torch class _Cache: - """The two fields _fetch_fraction_for reads, with the real method bound in.""" + """Bind the real split method without constructing a GPU cache.""" - def __init__(self, fraction=0.25, draft_ids=frozenset()): + def __init__(self, fraction=0.25): 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): @@ -33,73 +22,44 @@ def frac(self, layer_id, n_tokens, top_k=6): 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) +class TestTheSplitIsTheProfiledBandwidthRatio: + @pytest.mark.parametrize("n_tokens", [1, 2, 5, 6, 64]) + def test_token_width_does_not_rescale_equation_four(self, n_tokens): + # The cache kernel deduplicates all routes first, so its miss count m already + # reflects the wider batch. Equation 4 applies B_P/B_H to that m exactly once. + assert _Cache(0.28).frac(0, n_tokens) == pytest.approx(0.28) def test_a_disabled_fraction_stays_disabled(self): # 0.0 means "use the fixed hybrid_max_fetch cap instead". Scaling it would turn # the cap path on for speculative steps only. assert _Cache(0.0).frac(0, 6) == 0.0 - def test_a_one_dimensional_id_tensor_counts_as_one_token(self): + def test_tensor_shape_does_not_change_the_ratio(self): c = _Cache(0.3) assert c._fn(c, 0, torch.zeros(6, dtype=torch.int32)) == pytest.approx(0.3) -class 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 +class TestCpuExecutorScratchWidth: + def test_anchor_plus_five_drafts_needs_six_rows(self): + from freetoken.engine.engine import _cpu_moe_max_tokens - 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 + config = SimpleNamespace( + model_config=SimpleNamespace( + dsv4_args=SimpleNamespace(dspark_enabled=True, dspark_block_size=5) + ), + max_running_req=1, + cuda_graph_max_bs=None, + ) + assert _cpu_moe_max_tokens(config) == 6 - from freetoken.models.deepseek_v4.model import DeepseekV4ForCausalLM + def test_bound_scales_with_requests_and_graph_capture(self): + from freetoken.engine.engine import _cpu_moe_max_tokens - 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" + config = SimpleNamespace( + model_config=SimpleNamespace( + dsv4_args=SimpleNamespace(dspark_enabled=True, dspark_block_size=5) + ), + max_running_req=3, + cuda_graph_max_bs=32, ) + assert _cpu_moe_max_tokens(config) == 32 * 6 diff --git a/tests/dsv4/test_dsv4_proposal_writeback.py b/tests/dsv4/test_dsv4_proposal_writeback.py new file mode 100644 index 0000000..d3ea46d --- /dev/null +++ b/tests/dsv4/test_dsv4_proposal_writeback.py @@ -0,0 +1,71 @@ +"""DSpark proposal layout from the paper's gamma queries to target verification.""" + +from __future__ import annotations + +import pathlib + + +def _method(name: str) -> str: + src = ( + pathlib.Path(__file__).resolve().parents[2] + / "python" / "freetoken" / "engine" / "engine.py" + ).read_text() + i = src.index(f" def {name}(") + end = src.find("\n def ", i + 10) + return src[i : len(src) if end < 0 else end] + + +def _scheduler_method(name: str) -> str: + src = ( + pathlib.Path(__file__).resolve().parents[2] + / "python" / "freetoken" / "scheduler" / "scheduler.py" + ).read_text() + i = src.index(f" def {name}(") + end = src.find("\n def ", i + 10) + return src[i : len(src) if end < 0 else end] + + +class TestProposalLayout: + def test_all_gamma_proposals_fill_target_rows_one_through_gamma(self): + body = _method("draft_into_batch") + assert "i * span + 1:(i + 1) * span" in body + assert "proposed[i * k:(i + 1) * k]" in body + + def test_proposals_stay_on_gpu_until_acceptance(self): + body = _method("draft_into_batch") + assert "batch.draft_tokens = proposed" in body + assert "req.input_ids[-k:]" not in body + + def test_acceptance_reads_the_saved_proposals_not_noise_placeholders(self): + body = _method("_finish_speculative") + assert "proposed_gpu = batch.draft_tokens" in body + assert "proposed_cpu[i * k:(i + 1) * k]" in body + + def test_only_the_accepted_prefix_reaches_the_request_buffer(self): + body = _method("_finish_speculative") + assert "req.append_host(proposed[:n_acc]" in body + + def test_target_span_is_anchor_plus_gamma(self): + assert "span = 1 + k" in _method("draft_into_batch") + + def test_draft_masks_are_built_on_gpu_from_the_checkpoint_noise_id(self): + dspark = ( + pathlib.Path(__file__).resolve().parents[2] + / "python" / "freetoken" / "models" / "deepseek_v4" / "dspark.py" + ).read_text() + start = dspark.index(" def propose(") + end = dspark.index("\n def embed_block", start) + body = dspark[start:end] + assert "torch.full(" in body + assert "self.noise_token_id" in body + assert "draft_ids[:, 0] = input_ids" in body + + def test_bonus_writeback_uses_the_post_acceptance_frontier(self): + body = _scheduler_method("_forward") + assert "[req.device_len - 1 for req in batch.reqs]" in body + assert "self.token_pool[rows, positions]" in body + + def test_acceptance_uses_each_requests_sampling_params(self): + body = _method("_finish_speculative") + assert "sp = req.sampling_params" in body + assert "sp = batch.reqs[0].sampling_params" not in body diff --git a/tests/dsv4/test_dsv4_rejection_sampling.py b/tests/dsv4/test_dsv4_rejection_sampling.py index 865a8d9..4b74146 100644 --- a/tests/dsv4/test_dsv4_rejection_sampling.py +++ b/tests/dsv4/test_dsv4_rejection_sampling.py @@ -15,7 +15,10 @@ import pytest import torch -from freetoken.models.deepseek_v4.dspark import rejection_accept +from freetoken.models.deepseek_v4.dspark import ( + rejection_accept, + rejection_accept_device, +) VOCAB = 6 @@ -157,3 +160,29 @@ def test_a_residual_of_zero_falls_back_to_the_target(self): ) assert n_acc == 0 assert tok in (0, 1), "the fallback must still draw from the target's support" + + +class TestDeviceResidentAcceptance: + def test_full_acceptance_and_bonus_stay_exact(self): + q = torch.zeros(2, VOCAB) + q[0, 1] = 1.0 + q[1, 2] = 1.0 + p = torch.zeros(3, VOCAB) + p[0, 1] = 1.0 + p[1, 2] = 1.0 + p[2, 4] = 1.0 + got = rejection_accept_device( + torch.tensor([1, 2]), q, p, generator=torch.Generator().manual_seed(0) + ) + assert got == (2, 4) + + def test_rejection_recovers_from_the_target_residual(self): + q = torch.zeros(1, VOCAB) + q[0, 0] = 1.0 + p = torch.zeros(2, VOCAB) + p[0, 3] = 1.0 + p[1, 4] = 1.0 + got = rejection_accept_device( + torch.tensor([0]), q, p, generator=torch.Generator().manual_seed(0) + ) + assert got == (0, 3) diff --git a/tests/dsv4/test_dsv4_rollback.py b/tests/dsv4/test_dsv4_rollback.py deleted file mode 100644 index d28e19c..0000000 --- a/tests/dsv4/test_dsv4_rollback.py +++ /dev/null @@ -1,134 +0,0 @@ -"""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 deleted file mode 100644 index 5f75c47..0000000 --- a/tests/dsv4/test_dsv4_speculative_loop.py +++ /dev/null @@ -1,195 +0,0 @@ -"""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/kernels/test_dsv4_indexer_decode.py b/tests/kernels/test_dsv4_indexer_decode.py index a3dc63b..d5480e4 100644 --- a/tests/kernels/test_dsv4_indexer_decode.py +++ b/tests/kernels/test_dsv4_indexer_decode.py @@ -41,10 +41,10 @@ def _build(positions, n_stage, seed=0): return q, w, snap, valid -def _reference(q, w, pool, snap, valid, n_stage): +def _reference(q, w, pool, snap, valid, n_stage, row_map=None): """fp32 ground truth: sum_h relu(q_h . k_t) * w_h over the live blocks, -inf past them.""" b = q.shape[0] - rows = torch.arange(b, device=q.device) + rows = torch.arange(b, device=q.device) if row_map is None else row_map blk = torch.arange(n_stage, device=q.device) live = blk[None, :] < valid[:, None] at = snap[rows[:, None], (blk * RATIO)[None, :]] @@ -55,7 +55,8 @@ def _reference(q, w, pool, snap, valid, n_stage): def _check(pool, positions, n_stage, seed=0): q, w, snap, valid = _build(positions, n_stage, seed) - got = indexer_decode_logits(q, w, pool, snap, valid, n_stage, RATIO).float() + rows = torch.arange(q.shape[0], dtype=torch.int64, device=q.device) + got = indexer_decode_logits(q, w, pool, snap, rows, valid, n_stage, RATIO).float() ref = _reference(q, w, pool, snap, valid, n_stage) assert torch.equal(torch.isfinite(got), torch.isfinite(ref)), "live/-inf column set differs" live = torch.isfinite(ref) @@ -98,12 +99,27 @@ def test_per_row_valid(pool): assert torch.isinf(got[i, v:]).all() +def test_multiple_queries_share_one_request_snapshot(pool): + """DSpark verifies anchor+gamma query rows against one snapshot row per request.""" + n_stage = 256 + q, w, _, _ = _build([400] * 6, n_stage, seed=17) + _, _, snap, _ = _build([400], n_stage, seed=19) + rows = torch.zeros(6, dtype=torch.int64, device=q.device) + valid = torch.tensor([1, 7, 19, 41, 73, 100], dtype=torch.int64, device=q.device) + got = indexer_decode_logits(q, w, pool, snap, rows, valid, n_stage, RATIO).float() + ref = _reference(q, w, pool, snap, valid, n_stage, row_map=rows) + assert torch.equal(torch.isfinite(got), torch.isfinite(ref)) + live = torch.isfinite(ref) + torch.testing.assert_close(got[live], ref[live], **TOL) + + def test_cuda_graph_follows_valid(pool): """A captured graph must read the bound from device memory, not from capture time.""" n_stage = 2048 q, w, snap, valid = _build([800], n_stage) out = torch.empty((1, n_stage), dtype=pool.dtype, device="cuda") - call = lambda: indexer_decode_logits(q, w, pool, snap, valid, n_stage, RATIO) + rows = torch.arange(q.shape[0], dtype=torch.int64, device=q.device) + call = lambda: indexer_decode_logits(q, w, pool, snap, rows, valid, n_stage, RATIO) side = torch.cuda.Stream() side.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(side): diff --git a/tests/server/test_client_disconnect_aborts.py b/tests/server/test_client_disconnect_aborts.py new file mode 100644 index 0000000..0769615 --- /dev/null +++ b/tests/server/test_client_disconnect_aborts.py @@ -0,0 +1,136 @@ +"""A client that hangs up must stop the generation. + +Streaming responses already notice: yielding a chunk into a dead socket is observable, +and stream_with_cancellation acts on it. A NON-streaming request awaits the whole +completion and never looks -- so a caller that disconnects leaves the GPUs generating to +max_tokens for an answer nobody will read, with the queue behind it waiting on that work. + +On an offload MoE that is minutes of the whole machine per abandoned request. +""" + +from __future__ import annotations + +import asyncio +import pathlib + +import pytest + + +class _Request: + """Stand-in for a Starlette Request: disconnects after `after` polls.""" + + def __init__(self, after: int | None = None): + self.after = after + self.polls = 0 + + async def is_disconnected(self) -> bool: + self.polls += 1 + return self.after is not None and self.polls > self.after + + +class _Ack: + def __init__(self, finished=False): + self.finished = finished + self.incremental_output = "" + self.prompt_tokens_delta = 0 + self.completion_tokens_delta = 1 + self.cached_tokens = 0 + self.error = None + + +class _State: + """Only the pieces wait_for_ack_watching_client touches.""" + + def __init__(self, n_acks=100): + from freetoken.server.api_server import FrontendManager + + self.n_acks = n_acks + self.yielded = 0 + self.aborted: list[int] = [] + self._fn = FrontendManager.wait_for_ack_watching_client + + async def wait_for_ack(self, uid): + for i in range(self.n_acks): + self.yielded += 1 + yield _Ack(finished=(i == self.n_acks - 1)) + await asyncio.sleep(0) + + async def abort_user(self, uid): + self.aborted.append(uid) + + def watch(self, uid, request): + return self._fn(self, uid, request) + + +async def _drain(state, request): + got = 0 + try: + async for _ack in state.watch(7, request): + got += 1 + except asyncio.CancelledError: + pass + # abort_user is scheduled, not awaited, so let the loop run it. + await asyncio.sleep(0) + await asyncio.sleep(0) + return got + + +class TestADisconnectStopsTheGeneration: + def test_it_stops_early_instead_of_running_to_completion(self): + state = _State(n_acks=100) + got = asyncio.run(_drain(state, _Request(after=3))) + assert got < 100, "the loop ran to completion for a client that had hung up" + assert got <= 5 + + def test_it_aborts_the_request(self): + state = _State(n_acks=100) + asyncio.run(_drain(state, _Request(after=2))) + assert state.aborted == [7], "the engine must be told to drop the request" + + def test_a_connected_client_receives_everything(self): + state = _State(n_acks=10) + got = asyncio.run(_drain(state, _Request(after=None))) + assert got == 10 + assert state.aborted == [] + + def test_no_request_means_no_watching(self): + # Offline / internal callers have no HTTP request; they must still work. + state = _State(n_acks=10) + got = asyncio.run(_drain(state, None)) + assert got == 10 + assert state.aborted == [] + + +class TestTheNonStreamingPathsUseIt: + """The wrapper is worthless if the handlers keep calling the bare loop.""" + + @staticmethod + def _src(name: str) -> str: + return ( + pathlib.Path(__file__).resolve().parents[2] + / "python" / "freetoken" / "server" / name + ).read_text() + + def test_generate_full_watches_the_client(self): + src = self._src("generation.py") + body = src[src.index("async def _generate_full_impl") :] + assert "_acks(state, uid, request)" in body, ( + "a non-streaming ack loop that does not watch the client will generate to " + "max_tokens for a caller that has gone" + ) + + def test_the_helper_tolerates_a_state_without_the_watcher(self): + # Not every state implementation has it (the test fakes do not). A missing + # disconnect check must cost the abort, not the request. + src = self._src("generation.py") + assert 'getattr(state, "wait_for_ack_watching_client", None)' in src + + def test_chat_completions_passes_its_request_down(self): + src = self._src("openai_api.py") + assert "request=request" in src, ( + "generate_full cannot watch a client it was never given" + ) + + def test_completions_watches_too(self): + src = self._src("openai_api.py") + assert "_acks(state, uid, request)" in src From 38d7b2f38f74609edb4fd9e6d69fe85120348aa8 Mon Sep 17 00:00:00 2001 From: Carlos Alvarado Date: Sat, 22 Aug 2026 21:22:56 -0400 Subject: [PATCH 05/13] perf(dsv4): accelerate probabilistic DSpark verification --- python/freetoken/engine/engine.py | 62 ++++++++++++------- python/freetoken/engine/graph.py | 10 +-- python/freetoken/models/deepseek_v4/dspark.py | 38 ++++++++---- tests/dsv4/test_dsv4_adaptive_verification.py | 18 +++++- tests/dsv4/test_dsv4_rejection_sampling.py | 25 ++++++++ 5 files changed, 114 insertions(+), 39 deletions(-) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 7afe8b2..268a92e 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -1147,17 +1147,17 @@ def _finish_speculative( 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. Greedy verification reduces the - # [rows, vocab] logits on the GPU and transfers only one int per row. Copying - # the whole matrix to the CPU was ~3 MiB per six-row cycle and then repeated - # the reduction on the slower side of PCIe. - target_cpu = logits.argmax(dim=-1).to(torch.int32).to("cpu", non_blocking=False) - # The blocking target copy above proves the draft events have completed. - # Record the five startup samples here, after target launch, so timing the - # drafter never introduces a synchronization between draft and verify. - self._record_adaptive_draft_cost(batch) any_sampled = any(not req.sampling_params.is_greedy for req in batch.reqs) + any_greedy = any(req.sampling_params.is_greedy for req in batch.reqs) + # Not self.sampler: a speculative batch has 1+k rows per request while the + # sampling args are sized per request. Greedy verification needs target argmax + # ids. Probabilistic rejection consumes p directly; reducing and copying argmax + # before its p/q test inserted an otherwise unused synchronization every cycle. + target_cpu = ( + logits.argmax(dim=-1).to(torch.int32).to("cpu", non_blocking=False) + if any_greedy + else None + ) if any_sampled and batch.draft_probs is None: raise RuntimeError("sampled DSpark verification is missing draft probabilities") confidence = batch.draft_confidence @@ -1166,7 +1166,18 @@ def _finish_speculative( raise RuntimeError( "DSpark verify is missing the gamma proposals produced by its sequential stage" ) - proposed_cpu = proposed_gpu.to("cpu", non_blocking=False).to(torch.int32) + # Greedy acceptance consumes proposal ids immediately. Sampled acceptance keeps + # them on device until rejection has queued directly behind target verification. + proposed_cpu = ( + proposed_gpu.to("cpu", non_blocking=False).to(torch.int32) + if any_greedy + else None + ) + draft_cost_recorded = False + if any_greedy: + # The blocking id copies prove the draft events have completed. + self._record_adaptive_draft_cost(batch) + draft_cost_recorded = True emitted: list[torch.Tensor] = [] release_tail = getattr(batch, "release_tail", None) @@ -1181,16 +1192,6 @@ def _finish_speculative( # one ahead of the buffer during decode (see the write-back above), so that # form silently yielded k-1 proposals -- a short slice, not an error. start = req.input_ids.numel() - k - proposed = proposed_cpu[i * k:(i + 1) * k] - # A VIEW of _ids_buf, which the accept path overwrites below. Snapshot it - # for the debug line, or the log shows post-mutation values and invents - # agreements that never happened. - proposed_snapshot = proposed.clone() if _SPEC_DEBUG else None - if proposed.numel() != k: - raise RuntimeError( - f"block has {proposed.numel()} proposals, expected {k}; the " - "request's buffer and the batch's block width disagree" - ) # Fixed-width verification is the paper/vLLM fallback when the profiled # hardware-aware scheduler is disabled. A static per-token threshold is # deliberately not used: DSpark section 3.2 schedules a GLOBAL token budget @@ -1207,6 +1208,8 @@ def _finish_speculative( budget = req.max_device_len - start - 1 width = max(0, min(width, budget)) if greedy: + assert proposed_cpu is not None and target_cpu is not None + proposed = proposed_cpu[i * k:(i + 1) * k] n_acc, bonus = accepted_prefix( proposed[:width], target_cpu[off:off + width + 1] ) @@ -1226,6 +1229,22 @@ def _finish_speculative( batch.draft_probs[i * k:i * k + width], p_req, ) + if not draft_cost_recorded: + # The rejection sampler's one small D2H result proves the draft + # timing events have completed; do not add another synchronization. + self._record_adaptive_draft_cost(batch) + draft_cost_recorded = True + proposed = proposed_gpu[i * k:(i + 1) * k].to( + "cpu", non_blocking=False + ).to(torch.int32) + # Snapshot before req._ids_buf is rewritten below. Without it, debug logs + # display post-mutation tokens and invent proposal/target agreements. + proposed_snapshot = proposed.clone() if _SPEC_DEBUG else None + if proposed.numel() != k: + raise RuntimeError( + f"block has {proposed.numel()} proposals, expected {k}; the " + "request's buffer and the batch's block width disagree" + ) keep = start + n_acc req.input_ids = req._ids_buf[:start] if n_acc: @@ -1257,6 +1276,7 @@ def _finish_speculative( [round(float(c), 3) for c in confidence[i * k:(i + 1) * k]], int(req._ids_buf[start - 1]) if start > 0 else None, proposed_snapshot[:width].tolist(), + None if target_cpu is None else target_cpu[off:off + width + 1].tolist(), width, n_acc, bonus, emitted[-1].tolist() if emitted else None, diff --git a/python/freetoken/engine/graph.py b/python/freetoken/engine/graph.py index 29ff569..bfe1dd5 100644 --- a/python/freetoken/engine/graph.py +++ b/python/freetoken/engine/graph.py @@ -250,15 +250,15 @@ def _reset_moe_offload_cache(self) -> None: self.moe_offload_cache.reset() def _profile_graph_ms(self, graph: torch.cuda.CUDAGraph, replays: int = 5) -> float: - """Median replay cost with a cold expert cache, matching vLLM's profiler. + """Median replay cost for DSpark's hardware-capacity curve. - A reset precedes every sample because the real hybrid workload continually - changes routed experts. Pricing a repeated dummy route after its first replay - would measure an all-hot cache and systematically underprice verification. + Begin the series cold, then let ordinary LRU fills and hits participate in + the curve. This prices the cache reuse present during normal engine execution + instead of forcing an artificial cold miss before every timed replay. """ samples: list[float] = [] + self._reset_moe_offload_cache() for _ in range(replays): - self._reset_moe_offload_cache() start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) start.record(self.stream) diff --git a/python/freetoken/models/deepseek_v4/dspark.py b/python/freetoken/models/deepseek_v4/dspark.py index 942efde..20ae23a 100644 --- a/python/freetoken/models/deepseek_v4/dspark.py +++ b/python/freetoken/models/deepseek_v4/dspark.py @@ -613,7 +613,10 @@ def sampling_probs( out = torch.zeros_like(logits) out.scatter_(-1, logits.argmax(dim=-1, keepdim=True), 1.0) return out - scaled = logits.float() / temperature + # DSpark evaluates standard speculative sampling at temperature 1.0. Dividing by + # exactly one is numerically redundant, but it still launches an elementwise kernel + # for every draft and target row in every cycle. + scaled = logits.float() if temperature == 1.0 else logits.float() / temperature if top_k and top_k > 0: kth = scaled.topk(min(top_k, scaled.shape[-1]), dim=-1).values[..., -1:] scaled = scaled.masked_fill(scaled < kth, float("-inf")) @@ -698,25 +701,36 @@ def rejection_accept_device( uniforms = torch.rand( n, device=p.device, dtype=torch.float32, generator=generator ) - rejected = torch.nonzero(p_token <= uniforms * q_token, as_tuple=False) - n_acc = int(rejected[0, 0]) if rejected.numel() else n + # Keep the prefix reduction on device. Converting the first rejection to a + # Python int here would synchronize once for the length and again for the + # recovered token below. vLLM publishes both results together as well. + positions = torch.arange(n, device=p.device, dtype=torch.int64) + n_acc_device = torch.where( + p_token <= uniforms * q_token, + positions, + torch.full_like(positions, n), + ).amin() else: - n_acc = 0 + n_acc_device = torch.zeros((), dtype=torch.int64, device=p.device) - if n_acc < n: - residual = torch.clamp(p[n_acc] - q[n_acc], min=0.0) + if n: + rejected_row = n_acc_device.clamp_max(n - 1) + residual = torch.clamp(p[rejected_row] - q[rejected_row], min=0.0) total = residual.sum() - # torch.where keeps the degenerate fallback on device; clamp_min only - # protects the unselected division branch from producing NaNs. - dist = torch.where( + # clamp_min protects the unselected division branch from producing NaNs. + residual_dist = torch.where( total > 0, residual / total.clamp_min(torch.finfo(residual.dtype).tiny), - p[n_acc], + p[rejected_row], ) + dist = torch.where(n_acc_device < n, residual_dist, p[n]) else: dist = p[n] - bonus = int(torch.multinomial(dist, 1, generator=generator).item()) - return n_acc, bonus + bonus_device = torch.multinomial(dist, 1, generator=generator).squeeze(0) + result = torch.stack((n_acc_device, bonus_device.to(torch.int64))).to( + "cpu", non_blocking=False + ) + return int(result[0]), int(result[1]) def window_cols_for_block( diff --git a/tests/dsv4/test_dsv4_adaptive_verification.py b/tests/dsv4/test_dsv4_adaptive_verification.py index cbe56c9..c168a7f 100644 --- a/tests/dsv4/test_dsv4_adaptive_verification.py +++ b/tests/dsv4/test_dsv4_adaptive_verification.py @@ -6,7 +6,7 @@ import torch from freetoken.engine.engine import Engine -from freetoken.engine.graph import SharedSpecCarryJournal +from freetoken.engine.graph import GraphRunner, SharedSpecCarryJournal from freetoken.models.deepseek_v4.dspark import choose_adaptive_draft_width @@ -122,3 +122,19 @@ def test_greedy_verify_reduces_logits_before_crossing_pcie(): body = inspect.getsource(Engine._finish_speculative) assert 'logits.argmax(dim=-1).to(torch.int32).to("cpu"' in body assert 'logits.to("cpu"' not in body + + +def test_probabilistic_verify_does_not_compute_unused_target_argmax(): + import inspect + + body = inspect.getsource(Engine._finish_speculative) + assert "any_greedy = any(req.sampling_params.is_greedy" in body + assert "if any_greedy" in body + assert "else None" in body + + +def test_capacity_profile_flushes_cache_once_not_before_every_replay(): + import inspect + + body = inspect.getsource(GraphRunner._profile_graph_ms) + assert body.count("self._reset_moe_offload_cache()") == 1 diff --git a/tests/dsv4/test_dsv4_rejection_sampling.py b/tests/dsv4/test_dsv4_rejection_sampling.py index 4b74146..fe86c10 100644 --- a/tests/dsv4/test_dsv4_rejection_sampling.py +++ b/tests/dsv4/test_dsv4_rejection_sampling.py @@ -18,6 +18,7 @@ from freetoken.models.deepseek_v4.dspark import ( rejection_accept, rejection_accept_device, + sampling_probs, ) VOCAB = 6 @@ -186,3 +187,27 @@ def test_rejection_recovers_from_the_target_residual(self): torch.tensor([0]), q, p, generator=torch.Generator().manual_seed(0) ) assert got == (0, 3) + + def test_zero_width_draws_the_target_bonus(self): + p = torch.zeros(1, VOCAB) + p[0, 5] = 1.0 + got = rejection_accept_device( + torch.empty(0, dtype=torch.long), + torch.empty(0, VOCAB), + p, + generator=torch.Generator().manual_seed(0), + ) + assert got == (0, 5) + + def test_accepted_length_and_bonus_share_one_host_transfer(self): + import inspect + + body = inspect.getsource(rejection_accept_device) + assert 'torch.stack((n_acc_device, bonus_device.to(torch.int64))).to(' in body + assert ".item()" not in body + + +def test_temperature_one_probabilities_equal_the_direct_softmax(): + logits = torch.tensor([[1.0, -2.0, 3.0, 0.5]]) + got = sampling_probs(logits, temperature=1.0, top_p=1.0, top_k=-1) + assert torch.equal(got, torch.softmax(logits.float(), dim=-1)) From 986823e991ca12355a523caec08b1113e9d8c234 Mon Sep 17 00:00:00 2001 From: Carlos Alvarado Date: Sun, 23 Aug 2026 12:52:56 -0400 Subject: [PATCH 06/13] fix(dspark): detokenize accepted blocks exactly once --- python/freetoken/llm/llm.py | 6 +- python/freetoken/message/tokenizer.py | 6 ++ python/freetoken/scheduler/scheduler.py | 32 +++++--- python/freetoken/tokenizer/detokenize.py | 62 ++++++++++------ python/freetoken/tokenizer/server.py | 4 +- tests/server/test_message_wire.py | 3 +- .../tokenizer/test_speculative_detokenize.py | 74 +++++++++++++++++++ 7 files changed, 149 insertions(+), 38 deletions(-) create mode 100644 tests/tokenizer/test_speculative_detokenize.py diff --git a/python/freetoken/llm/llm.py b/python/freetoken/llm/llm.py index 1d4aaef..92b215c 100644 --- a/python/freetoken/llm/llm.py +++ b/python/freetoken/llm/llm.py @@ -99,8 +99,10 @@ def offline_send_result(self, reply: List[BaseTokenizerMsg]) -> None: continue assert isinstance(msg, DetokenizeMsg) status = self.status_map[msg.uid] - if not (msg.finished and msg.next_token in self.eos_token_ids): - status.output_ids.append(msg.next_token) + token_ids = msg.token_ids if msg.token_ids is not None else [msg.next_token] + if msg.finished and token_ids[-1] in self.eos_token_ids: + token_ids = token_ids[:-1] + status.output_ids.extend(token_ids) def generate( self, diff --git a/python/freetoken/message/tokenizer.py b/python/freetoken/message/tokenizer.py index 33b75c7..273a229 100644 --- a/python/freetoken/message/tokenizer.py +++ b/python/freetoken/message/tokenizer.py @@ -29,6 +29,12 @@ class DetokenizeMsg(BaseTokenizerMsg): uid: int next_token: int finished: bool + # A speculative verify can accept several tokens in one engine step. Keep + # ``next_token`` as the final token for wire/backward compatibility, while + # carrying the complete ordered block here. This mirrors vLLM's + # EngineCoreOutput.new_token_ids contract and lets the detokenizer advance + # through the block exactly once. + token_ids: list[int] | None = None finish_reason: str | None = None # The stop string that ended generation (if any), so the detokenizer can trim it # and everything after it from the final output. diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index b7a7eb8..ec67706 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -390,6 +390,9 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: ) block_start = req.input_ids.numel() - len(block) finished = False + visible_tokens: list[int] = [] + finish_reason = None + matched_stop = None for pos, tok in enumerate(block): visible_len = block_start + pos + 1 hit_length = not req.can_decode and pos == len(block) - 1 @@ -413,19 +416,26 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: and not finished ): req.toolcall_anchor_len = visible_len - 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, - ) - ) + visible_tokens.append(tok) if finished: break # drop the rest of the block: it is past the reply's end - next_token = block[-1] + # vLLM carries all newly accepted speculative ids in one engine + # output and advances its incremental detokenizer sequentially. Do + # the same here. Sending one DetokenizeMsg per token in a single + # BatchTokenizerMsg made the old batched decoder append cumulative + # prefixes ("pres" + "presidencia" + ...). + next_token = visible_tokens[-1] + reply.append( + DetokenizeMsg( + uid=req.uid, + next_token=next_token, + finished=finished, + token_ids=visible_tokens if len(visible_tokens) > 1 else None, + finish_reason=finish_reason, + matched_stop=matched_stop, + stop_strs=req.sampling_params.stop_strs or None, + ) + ) # NOTE: overlap scheduling may make the request freed twice, skip second free if finished and req not in self.finished_reqs: diff --git a/python/freetoken/tokenizer/detokenize.py b/python/freetoken/tokenizer/detokenize.py index d138ac1..975b00a 100644 --- a/python/freetoken/tokenizer/detokenize.py +++ b/python/freetoken/tokenizer/detokenize.py @@ -90,9 +90,37 @@ def discard(self, uid: int) -> None: accumulated ids and text stay in ``decode_map`` for the life of the worker.""" self.decode_map.pop(uid, None) + def _append_token(self, s: DecodeStatus, token: int, *, skip: bool = False) -> str: + """Advance one request by one token, preserving tokenizer boundary context.""" + if not skip: + s.decoded_ids.append(token) + read_str, surr_str = self.tokenizer.batch_decode( + [ + s.decoded_ids[s.surr_offset :], + s.decoded_ids[s.surr_offset : s.read_offset], + ] + ) + new_text = read_str[len(surr_str) :] + if len(new_text) > 0 and not new_text.endswith("�"): + output_str = s.decoded_str + new_text + s.decoded_str = output_str + s.surr_offset = s.read_offset + s.read_offset = len(s.decoded_ids) + else: + new_text = find_printable_text(new_text) + output_str = s.decoded_str + new_text + return output_str + def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]: - read_ids: List[List[int]] = [] - surr_ids: List[List[int]] = [] + """Detokenize each request's newly emitted token block in order. + + A speculative step can deliver multiple accepted tokens for one uid. The + original two-pass batch implementation appended every token before updating + the uid's offsets, so each decoded prefix was appended cumulatively. vLLM's + speculative output processor advances tokens sequentially; doing so here is + both exact and safe when multiple messages for one uid are coalesced. + """ + incremental_strs: List[str] = [] for msg in msgs: if msg.uid not in self.decode_map: self.decode_map[msg.uid] = DecodeStatus( @@ -103,27 +131,15 @@ def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]: sent_offset=0, ) s = self.decode_map[msg.uid] - if not (msg.finished and msg.next_token in self.eos_token_ids): - s.decoded_ids.append(msg.next_token) - read_ids.append(s.decoded_ids[s.surr_offset :]) - surr_ids.append(s.decoded_ids[s.surr_offset : s.read_offset]) - - read_texts = self.tokenizer.batch_decode(read_ids) - surr_texts = self.tokenizer.batch_decode(surr_ids) - - incremental_strs: List[str] = [] - for msg, read_str, surr_str in zip(msgs, read_texts, surr_texts, strict=True): - s = self.decode_map[msg.uid] - new_text = read_str[len(surr_str) :] - # Streaming chunk: update the decode status - if len(new_text) > 0 and not new_text.endswith("�"): - output_str = s.decoded_str + new_text - s.decoded_str = output_str - s.surr_offset = s.read_offset - s.read_offset = len(s.decoded_ids) - else: - new_text = find_printable_text(new_text) - output_str = s.decoded_str + new_text + tokens = msg.token_ids if msg.token_ids is not None else [msg.next_token] + output_str = s.decoded_str + for pos, token in enumerate(tokens): + final_eos = ( + msg.finished + and pos == len(tokens) - 1 + and token in self.eos_token_ids + ) + output_str = self._append_token(s, token, skip=final_eos) prev_sent = s.sent_offset if msg.finished: diff --git a/python/freetoken/tokenizer/server.py b/python/freetoken/tokenizer/server.py index 530e862..0f69cba 100644 --- a/python/freetoken/tokenizer/server.py +++ b/python/freetoken/tokenizer/server.py @@ -215,7 +215,9 @@ def tokenize_worker( finished=msg.finished, finish_reason=msg.finish_reason, matched_stop=msg.matched_stop, - completion_tokens_delta=1, + completion_tokens_delta=( + len(msg.token_ids) if msg.token_ids is not None else 1 + ), kv_used_pages=msg.kv_used_pages, kv_total_pages=msg.kv_total_pages, mamba_used_slots=msg.mamba_used_slots, diff --git a/tests/server/test_message_wire.py b/tests/server/test_message_wire.py index 3bd7cc6..8c6e652 100644 --- a/tests/server/test_message_wire.py +++ b/tests/server/test_message_wire.py @@ -89,13 +89,14 @@ def test_user_reply_token_deltas_round_trip(): def test_detokenize_msg_carries_kv_usage_round_trip(): msg = DetokenizeMsg( - uid=3, next_token=42, finished=True, + uid=3, next_token=42, finished=True, token_ids=[40, 41, 42], kv_used_pages=10, kv_total_pages=256, gpu_mem_bytes=1 << 30, mamba_used_slots=7, mamba_total_slots=64, swa_used_tokens=8448, swa_total_tokens=76800, ) decoded = BaseTokenizerMsg.decoder(BaseTokenizerMsg.encoder(msg)) assert isinstance(decoded, DetokenizeMsg) + assert decoded.token_ids == [40, 41, 42] assert (decoded.kv_used_pages, decoded.kv_total_pages, decoded.gpu_mem_bytes) == (10, 256, 1 << 30) assert (decoded.mamba_used_slots, decoded.mamba_total_slots) == (7, 64) assert (decoded.swa_used_tokens, decoded.swa_total_tokens) == (8448, 76800) diff --git a/tests/tokenizer/test_speculative_detokenize.py b/tests/tokenizer/test_speculative_detokenize.py new file mode 100644 index 0000000..5a0a41a --- /dev/null +++ b/tests/tokenizer/test_speculative_detokenize.py @@ -0,0 +1,74 @@ +from freetoken.message import DetokenizeMsg +from freetoken.tokenizer.detokenize import DetokenizeManager + + +class _PieceTokenizer: + eos_token_id = 0 + + _pieces = { + 0: "", + 1: "pres", + 2: "idencia", + 3: " de", + 4: " López", + 5: " Obrador", + 6: ".", + } + + def batch_decode(self, rows): + return ["".join(self._pieces[token] for token in row) for row in rows] + + +def _manager() -> DetokenizeManager: + return DetokenizeManager(_PieceTokenizer(), frozenset({0})) + + +def test_speculative_block_emits_each_token_once(): + manager = _manager() + + chunks = manager.detokenize( + [ + DetokenizeMsg( + uid=7, + next_token=5, + token_ids=[1, 2, 3, 4, 5], + finished=False, + ) + ] + ) + + assert chunks == ["presidencia de López Obrador"] + + +def test_coalesced_messages_for_one_request_do_not_repeat_prefixes(): + manager = _manager() + + chunks = manager.detokenize( + [ + DetokenizeMsg(uid=9, next_token=2, token_ids=[1, 2], finished=False), + DetokenizeMsg(uid=9, next_token=5, token_ids=[3, 4, 5], finished=False), + DetokenizeMsg(uid=9, next_token=0, finished=True), + ] + ) + + assert chunks == ["presidencia", " de López Obrador", ""] + assert 9 not in manager.decode_map + + +def test_finished_block_counts_text_before_eos_once(): + manager = _manager() + + chunks = manager.detokenize( + [ + DetokenizeMsg( + uid=11, + next_token=0, + token_ids=[1, 2, 6, 0], + finished=True, + finish_reason="stop", + ) + ] + ) + + assert chunks == ["presidencia."] + assert 11 not in manager.decode_map From 4082dd86306587a35af82892c1dcadda872f8715 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Sun, 23 Aug 2026 10:18:00 -0400 Subject: [PATCH 07/13] perf(dsv4): fuse the HC pre-norm reduction into one kernel ``hc_pre`` and ``hc_head`` compute the inverse RMS as xf = x.flatten(2).float() rsqrt = torch.rsqrt(xf.square().mean(-1, keepdim=True) + eps) ``square()`` materialises a second full fp32 copy of the hidden state purely to reduce it away. At DSV4-Flash shapes (hc_dim 16384, T=8192) that is a 537 MB write plus a 537 MB read on top of the 537 MB read -- to produce [T, 1]. ``inv_rms`` does it in one pass over the bf16 source (the fp32 upcast is exact, so the per-element squares are unchanged, and reading the source halves the bytes again). RTX 6000 Ada, hc_dim 16384: T=8192 1.944 ms -> 0.304 ms (6.4x, ~880 GB/s of a 960 GB/s peak) T=4096 0.972 ms -> 0.153 ms (6.4x) T= 64 0.020 ms -> 0.013 ms (1.5x -- decode benefits too) In an nsys trace of 16-concurrent long-prompt serving, the four kernels this block emits are ~13.7% of GPU kernel time. Not bit-identical, and it cannot be. ATen's reduction order for ``mean(x**2, -1)`` is chosen per shape: a tree/2 fold reproduces it exactly at M=4 and nothing reproduces it at M=64, because the block/grid split changes with M. There is no fixed order to target, so any fused form changes the last ULP -- and on a deterministic server that changes greedy continuations. Two baseline runs produced byte-identical output, so this was verified, not assumed. What the fused form must not do is get *worse*, and this one does not. Against an fp64 reference the mean relative error is 3.65e-08 vs ATen's 3.60e-08 at M=8192, and 3.60e-08 vs 3.63e-08 at M=4096 -- ahead at some shapes, behind at others, always within 8% of ATen's own distance from the truth. A test pins that. Two faster-looking alternatives were rejected on this basis: ``linalg.vector_norm`` is the same speed but consistently ~23% worse (4.47e-08) because sqrt-then-square rounds twice, and ``einsum`` over bf16 accumulates in bf16 (1.9e-03). Bit-identical alternatives exist but none is faster: ``vecdot(xf, xf)`` lowers to the same path at 1.944 ms, and vecdot with fp32 accumulate over bf16 is 4.557 ms. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01GpGe2fQ5pDShGrnSuksfun --- python/freetoken/kernel/triton/dsv4/norm.py | 46 +++++++++++++- python/freetoken/models/deepseek_v4/model.py | 11 ++-- tests/dsv4/test_hc_inv_rms.py | 67 ++++++++++++++++++++ 3 files changed, 119 insertions(+), 5 deletions(-) create mode 100644 tests/dsv4/test_hc_inv_rms.py diff --git a/python/freetoken/kernel/triton/dsv4/norm.py b/python/freetoken/kernel/triton/dsv4/norm.py index b1c4399..523448a 100644 --- a/python/freetoken/kernel/triton/dsv4/norm.py +++ b/python/freetoken/kernel/triton/dsv4/norm.py @@ -52,4 +52,48 @@ def rms_norm(x: torch.Tensor, weight: torch.Tensor | None, eps: float) -> torch. return out.reshape(x.shape) -__all__ = ["rms_norm"] +@triton.jit +def _inv_rms_kernel(x_ptr, o_ptr, D, eps, stride_xm, BLOCK_D: tl.constexpr): + row = tl.program_id(0) + acc = tl.zeros((BLOCK_D,), dtype=tl.float32) + for off in range(0, D, BLOCK_D): + offs = off + tl.arange(0, BLOCK_D) + v = tl.load(x_ptr + row * stride_xm + offs, mask=offs < D, other=0.0).to(tl.float32) + acc += v * v + tl.store(o_ptr + row, tl.rsqrt(tl.sum(acc, axis=0) / D + eps)) + + +def inv_rms(x: torch.Tensor, eps: float) -> torch.Tensor: + """``rsqrt(mean(x^2, -1) + eps)`` -> ``[..., 1]`` fp32, one pass, no fp32 temp. + + :func:`rms_norm` applies the scale and writes a full-size output; callers that need + the *scalar* (DSV4's hyper-connection pre-norm multiplies it into a [.., mix_hc] + tensor, not into x) were spelling it ``x.float().square().mean(-1)``, which + materialises a second full fp32 copy of the hidden state purely to reduce it away. + At hc_dim 16384, T=8192 that is a 537 MB write plus a 537 MB read for a [T, 1] + result. Reading the bf16 source instead of its fp32 upcast halves the bytes again; + the upcast is exact, so the sum of squares is unchanged. + + RTX 6000 Ada, hc_dim 16384: 1.944 ms -> 0.304 ms at T=8192 (6.4x), 0.020 -> 0.013 ms + at T=64. The large end lands at ~880 GB/s against a 960 GB/s peak -- the memory roof. + + NOT bit-identical to ``square().mean()``: the reduction order differs, and ATen's + order is not reproducible anyway (it varies with M -- a tree/2 fold matches at M=4 + and nothing matches at M=64, because the block/grid split is chosen per shape). The + accuracy is equivalent, not merely close: against an fp64 reference the mean relative + error is 3.65e-08 here vs 3.60e-08 for ATen at M=8192, and 3.60e-08 vs 3.63e-08 at + M=4096 -- i.e. it wins at some shapes and loses at others, within 8% of ATen's own + distance from the truth. (``linalg.vector_norm`` is the same speed but consistently + ~23% worse, at 4.47e-08, because sqrt-then-square rounds twice.) + """ + D = x.shape[-1] + x2d = x.reshape(-1, D) + out = torch.empty(x2d.shape[0], dtype=torch.float32, device=x.device) + _inv_rms_kernel[(x2d.shape[0],)]( + x2d, out, D, eps, x2d.stride(0), + BLOCK_D=min(2048, triton.next_power_of_2(D)), num_warps=8, + ) + return out.view(*x.shape[:-1], 1) + + +__all__ = ["rms_norm", "inv_rms"] diff --git a/python/freetoken/models/deepseek_v4/model.py b/python/freetoken/models/deepseek_v4/model.py index b013a81..13bb2af 100644 --- a/python/freetoken/models/deepseek_v4/model.py +++ b/python/freetoken/models/deepseek_v4/model.py @@ -31,6 +31,7 @@ from freetoken.core import get_global_ctx from freetoken.distributed import DistributedCommunicator from freetoken.kernel.triton.dsv4.hc import hc_post_combine, hc_pre_combine +from freetoken.kernel.triton.dsv4.norm import inv_rms from freetoken.kernel.triton.dsv4.sinkhorn import hc_split_sinkhorn from freetoken.models.blocks import BaseLLMModel from freetoken.utils import init_logger @@ -99,8 +100,9 @@ def __init__(self, layer_id: int, args: DeepseekV4Args, compress_ratio: int | No def hc_pre(self, x, hc_fn, hc_scale, hc_base): shape, dtype = x.size(), x.dtype - xf = x.flatten(2).float() - rsqrt = torch.rsqrt(xf.square().mean(-1, keepdim=True) + self.norm_eps) + x2d = x.flatten(2) + xf = x2d.float() + rsqrt = inv_rms(x2d, self.norm_eps) mixes = F.linear(xf, hc_fn) * rsqrt pre, post, comb = hc_split_sinkhorn( mixes.view(-1, mixes.size(-1)), hc_scale, hc_base, self.hc_mult, self.hc_sinkhorn_iters, self.hc_eps @@ -274,8 +276,9 @@ def logits(self, h: torch.Tensor) -> torch.Tensor: def hc_head(self, x): shape, dtype = x.size(), x.dtype dim = self.args.dim - xf = x.flatten(2).float() - rsqrt = torch.rsqrt(xf.square().mean(-1, keepdim=True) + self.norm_eps) + x2d = x.flatten(2) + xf = x2d.float() + rsqrt = inv_rms(x2d, self.norm_eps) mixes = F.linear(xf, self.hc_head_fn) * rsqrt pre = torch.sigmoid(mixes * self.hc_head_scale + self.hc_head_base) + self.hc_eps M = shape[0] * shape[1] diff --git a/tests/dsv4/test_hc_inv_rms.py b/tests/dsv4/test_hc_inv_rms.py new file mode 100644 index 0000000..61eca39 --- /dev/null +++ b/tests/dsv4/test_hc_inv_rms.py @@ -0,0 +1,67 @@ +"""``inv_rms`` must agree with the expression it replaced, and be no less accurate. + +The old form was ``torch.rsqrt(xf.square().mean(-1, keepdim=True) + eps)`` over the +fp32 upcast. ``inv_rms`` reduces the bf16 source directly: the upcast is exact, so the +per-element squares are identical and only the reduction order differs. Bit-equality is +not asserted -- ATen's order varies with M and is not reproducible -- so the contract +is (a) agreement to fp32 rounding and (b) accuracy no worse than ATen's against fp64. +""" +import pytest +import torch + +from freetoken.kernel.triton.dsv4.norm import inv_rms + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="triton kernel") + +HC_DIM = 4 * 4096 # hc_mult * dim at DSV4-Flash +EPS = 1e-6 + + +def _old(xf, eps): + return torch.rsqrt(xf.square().mean(-1, keepdim=True) + eps) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("rows", [1, 7, 512]) +def test_matches_square_mean(dtype, rows): + dev = "cuda" + x = torch.randn(rows, HC_DIM, device=dev, dtype=dtype) + got = inv_rms(x, EPS) + ref = _old(x.float(), EPS) + assert got.shape == ref.shape == (rows, 1) + assert got.dtype == torch.float32 + torch.testing.assert_close(got, ref, rtol=1e-5, atol=0) + + +def test_keeps_leading_dims(): + """hc_pre feeds a [B, T, hc_dim] view and broadcasts the result against [B, T, mix].""" + dev = "cuda" + x = torch.randn(2, 3, HC_DIM, device=dev, dtype=torch.bfloat16) + got = inv_rms(x, EPS) + assert got.shape == (2, 3, 1) + torch.testing.assert_close(got, _old(x.float(), EPS), rtol=1e-5, atol=0) + + +def test_tiny_and_large_magnitudes(): + """Check the scale range a hidden state can reach.""" + dev = "cuda" + for scale in (1e-3, 1.0, 1e3): + x = torch.randn(4, HC_DIM, device=dev, dtype=torch.bfloat16) * scale + torch.testing.assert_close(inv_rms(x, EPS), _old(x.float(), EPS), + rtol=1e-5, atol=0) + + +def test_accuracy_not_worse_than_aten(): + """The point of the Triton reduce over ``linalg.vector_norm``: no sqrt round-trip, + so it stays level with ATen against an fp64 reference instead of ~23% behind.""" + torch.manual_seed(0) + for rows in (64, 1024): + x = torch.randn(rows, HC_DIM, device="cuda", dtype=torch.bfloat16) + xf = x.float() + xd = xf.double() + truth = torch.rsqrt((xd * xd).mean(-1, keepdim=True) + EPS) + err_new = ((inv_rms(x, EPS).double() - truth).abs() / truth.abs()).mean() + err_aten = ((_old(xf, EPS).double() - truth).abs() / truth.abs()).mean() + # same quality of fp32 answer -- allow a small margin either way, but catch a + # real regression such as accumulating in bf16 or re-squaring a norm. + assert err_new < err_aten * 1.25, (rows, err_new.item(), err_aten.item()) From df6c340272ad74d3a0c0ee9f3991b06df55b1934 Mon Sep 17 00:00:00 2001 From: Carlos Alvarado Date: Sun, 23 Aug 2026 11:35:08 -0400 Subject: [PATCH 08/13] fix(dsv4): harden long prefills with fused HC norm Extend Gabriel Devenyi's fused inv_rms path to the DSpark drafter, keep HC combines on the bf16 source, and honor the configured prefill chunk limit so one long OpenWebUI history cannot materialize an unbounded activation. Retain opt-in one-shot verification profiling for deployment cost measurements; it remains disabled during normal serving. Based-on: FlashML-org/FreeToken#101 --- python/freetoken/engine/engine.py | 69 ++++++++++++++----- python/freetoken/models/deepseek_v4/dspark.py | 9 +-- python/freetoken/models/deepseek_v4/model.py | 6 +- scripts/freetoken-dsv4.sh | 4 ++ tests/dsv4/test_dsv4_dspark_wiring.py | 12 ++++ tests/engine/test_cache_budget.py | 9 +++ 6 files changed, 84 insertions(+), 25 deletions(-) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 268a92e..7b0f8ed 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -5,6 +5,7 @@ import glob import os import time +from contextlib import nullcontext from datetime import timedelta from typing import Any, Dict, Iterable, NamedTuple, Tuple @@ -309,6 +310,9 @@ def _materialize_loaded_weight_state_dict( # Diagnostic only: synchronize and report the first speculative cycles by stage. # Disabled by default because the per-cycle synchronization deliberately removes overlap. _SPEC_TIMING = os.environ.get("FREETOKEN_SPEC_TIMING", "0") == "1" +# One-shot operator profile of the target verify. Explicit opt-in: Kineto adds +# substantial CPU overhead and synchronizes CUDA when the table is emitted. +_SPEC_PROFILE = os.environ.get("FREETOKEN_SPEC_PROFILE", "0") == "1" class ForwardOutput(NamedTuple): next_tokens_gpu: torch.Tensor @@ -480,6 +484,7 @@ def __init__(self, config: EngineConfig): self._spec_debug_left = 6 self._spec_drafted = 0 self._spec_timing_left = 12 + self._spec_profile_left = 1 if is_offload_moe_backend(config.moe_backend): self._init_offload_moe_cache(config) if hasattr(self.model, "prepare_for_runtime"): @@ -1447,21 +1452,47 @@ def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput: timing = bool( batch.speculative and _SPEC_TIMING and self._spec_timing_left > 0 ) + profiling = bool( + batch.speculative and _SPEC_PROFILE and self._spec_profile_left > 0 + ) if timing: target_start = torch.cuda.Event(enable_timing=True) target_end = torch.cuda.Event(enable_timing=True) target_start.record(self.stream) - with self.ctx.forward_batch(batch): - if self.graph_runner.can_use_spec_cuda_graph(batch): - logits = self.graph_runner.replay_spec(batch) - target_features = self.graph_runner.dspark_spec_target_features(batch) - elif self.graph_runner.can_use_cuda_graph(batch): - logits = self.graph_runner.replay(batch) - target_features = self.graph_runner.dspark_target_features(batch) - else: - logits = self.model.forward() - get_features = getattr(self.model, "dspark_target_features", None) - target_features = get_features() if get_features is not None else None + profile_ctx = ( + torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + record_shapes=True, + ) + if profiling else nullcontext() + ) + with profile_ctx as prof: + with self.ctx.forward_batch(batch): + if self.graph_runner.can_use_spec_cuda_graph(batch): + logits = self.graph_runner.replay_spec(batch) + target_features = self.graph_runner.dspark_spec_target_features(batch) + elif self.graph_runner.can_use_cuda_graph(batch): + logits = self.graph_runner.replay(batch) + target_features = self.graph_runner.dspark_target_features(batch) + else: + logits = self.model.forward() + get_features = getattr(self.model, "dspark_target_features", None) + target_features = get_features() if get_features is not None else None + if profiling: + torch.cuda.synchronize(self.device) + assert prof is not None + logger.info( + "speculative verify profile (by CUDA time):\n%s", + prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=30), + ) + logger.info( + "speculative verify profile (by CPU time):\n%s", + prof.key_averages().table(sort_by="self_cpu_time_total", row_limit=30), + ) + self._spec_profile_left -= 1 if timing: target_end.record(self.stream) if self.cpu_moe_executor is not None: @@ -1605,8 +1636,8 @@ def _resolve_cache_type(has_linear_attention: bool, requested: str) -> str: def _adjust_dsv4_config(config: EngineConfig, override) -> None: """DSV4 engine-config reconciliation at config-resolution time (before the pool exists). Syncs the resolved runtime config into the opaque ``dsv4_args`` payload, sets - page_size to the window page P, forces single-chunk prefill, and clamps cuda_graph_bs/max_bs to - the DSV4 decode batch size. + page_size to the window page P, and clamps cuda_graph_bs/max_bs to the DSV4 + decode batch size. """ model_config = config.model_config model_config.dsv4_args.max_seq_len = config.max_seq_len @@ -1647,12 +1678,12 @@ def _adjust_dsv4_config(config: EngineConfig, override) -> None: # 'naive' stays naive with the pool's swa currency riding swa_paged. if getattr(config, "cache_type", "radix") != "naive": override("cache_type", "swa_radix") - # 'radix' (SWARadixCache on the full-loc currency, carry-aware re-prefill) is the default and is - # honored, as is an explicit 'naive'. Don't let max_extend_tokens force a second chunk within - # one prompt (the pool's prefill_chunk_budget still chunks prompts larger than the window - # pool); prefill batches ragged (bs>=1), each segment resuming from its own cached_len. - if getattr(config, "max_extend_tokens", 0) < config.max_seq_len: - override("max_extend_tokens", config.max_seq_len) + # Honor --max-prefill-length. DSV4's compressor carry and dSpark context + # catch-up both support continuation chunks; overriding the user's limit to + # max_seq_len made a long prompt materialize the full mHC activation at once. + # On memory-tight deployments that turns one oversized client history into a + # process-fatal prefill OOM. The pool's prefill_chunk_budget remains a second, + # capacity-derived upper bound in Scheduler. # DSV4 decode batches at most max_running_req rows; its full-loc snapshot is sized to that, # so a graph bs above it would exceed the backend's captured snapshot rows. Clamp any diff --git a/python/freetoken/models/deepseek_v4/dspark.py b/python/freetoken/models/deepseek_v4/dspark.py index 20ae23a..fc11daf 100644 --- a/python/freetoken/models/deepseek_v4/dspark.py +++ b/python/freetoken/models/deepseek_v4/dspark.py @@ -37,6 +37,7 @@ from torch import nn from freetoken.kernel.triton.dsv4.hc import hc_pre_combine +from freetoken.kernel.triton.dsv4.norm import inv_rms from freetoken.utils import init_logger from .args import DeepseekV4Args @@ -503,13 +504,13 @@ def block_input_ids(self, last_token: int, device: torch.device) -> torch.Tensor 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 + x2d = x.flatten(2) + xf = x2d.float() + mixes = F.linear(xf, self.hc_head_fn) * inv_rms(x2d, self.norm_eps) pre = torch.sigmoid(mixes * self.hc_head_scale + self.hc_head_base) + self.hc_eps M = shape[0] * shape[1] return hc_pre_combine( - xf.view(M, self.hc_mult, self.dim), pre.view(M, self.hc_mult), dtype + x.view(M, self.hc_mult, self.dim), pre.view(M, self.hc_mult), dtype ).view(*shape[:2], self.dim) def head_hidden( diff --git a/python/freetoken/models/deepseek_v4/model.py b/python/freetoken/models/deepseek_v4/model.py index 13bb2af..7da008a 100644 --- a/python/freetoken/models/deepseek_v4/model.py +++ b/python/freetoken/models/deepseek_v4/model.py @@ -108,7 +108,7 @@ def hc_pre(self, x, hc_fn, hc_scale, hc_base): mixes.view(-1, mixes.size(-1)), hc_scale, hc_base, self.hc_mult, self.hc_sinkhorn_iters, self.hc_eps ) M = shape[0] * shape[1] - y = hc_pre_combine(xf.view(M, self.hc_mult, self.dim), pre, dtype).view(*shape[:2], self.dim) + y = hc_pre_combine(x.view(M, self.hc_mult, self.dim), pre, dtype).view(*shape[:2], self.dim) return y, post.view(M, self.hc_mult), comb.view(M, self.hc_mult, self.hc_mult) def hc_post(self, x, residual, post, comb): @@ -282,7 +282,9 @@ def hc_head(self, x): mixes = F.linear(xf, self.hc_head_fn) * rsqrt pre = torch.sigmoid(mixes * self.hc_head_scale + self.hc_head_base) + self.hc_eps M = shape[0] * shape[1] - return hc_pre_combine(xf.view(M, self.hc_mult, dim), pre.view(M, self.hc_mult), dtype).view(*shape[:2], dim) + return hc_pre_combine( + x.view(M, self.hc_mult, dim), pre.view(M, self.hc_mult), dtype + ).view(*shape[:2], dim) def prefill_batched( self, input_ids: torch.Tensor, segments, flat_positions: torch.Tensor, diff --git a/scripts/freetoken-dsv4.sh b/scripts/freetoken-dsv4.sh index b648180..7475f5e 100644 --- a/scripts/freetoken-dsv4.sh +++ b/scripts/freetoken-dsv4.sh @@ -20,6 +20,7 @@ PORT="${PORT:-8081}" TP_SIZE="${TP_SIZE:-4}" MEMORY_RATIO="${MEMORY_RATIO:-0.90}" MAX_RUNNING_REQUESTS="${MAX_RUNNING_REQUESTS:-1}" +MAX_PREFILL_LENGTH="${MAX_PREFILL_LENGTH:-1024}" EXPERT_LOAD="${EXPERT_LOAD:-serial}" # SPECULATIVE_DSPARK=0 turns OFF the checkpoint's dSpark drafter (the mtp.* stack). # On by default: the draft/verify loop is wired, and dSpark is what the checkpoint was @@ -38,6 +39,8 @@ SPECULATIVE_DSPARK="${SPECULATIVE_DSPARK:-1}" [ -n "${FREETOKEN_SPEC_HYBRID_MOE:-}" ] && export FREETOKEN_SPEC_HYBRID_MOE # FREETOKEN_SPEC_DEBUG=1 dumps the first few blocks' token flow. [ -n "${FREETOKEN_SPEC_DEBUG:-}" ] && export FREETOKEN_SPEC_DEBUG +[ -n "${FREETOKEN_SPEC_TIMING:-}" ] && export FREETOKEN_SPEC_TIMING +[ -n "${FREETOKEN_SPEC_PROFILE:-}" ] && export FREETOKEN_SPEC_PROFILE [ -n "${FREETOKEN_SPEC_FORCE_REJECT:-}" ] && export FREETOKEN_SPEC_FORCE_REJECT [ -n "${FREETOKEN_CMP_FIRST_ONLY:-}" ] && export FREETOKEN_CMP_FIRST_ONLY [ -n "${FREETOKEN_SPEC_WINDOW_ONLY:-}" ] && export FREETOKEN_SPEC_WINDOW_ONLY @@ -125,6 +128,7 @@ cmd_start() { --tensor-parallel-size "$TP_SIZE" \ --memory-ratio "$MEMORY_RATIO" \ --max-running-requests "$MAX_RUNNING_REQUESTS" \ + --max-prefill-length "$MAX_PREFILL_LENGTH" \ --expert-load "$EXPERT_LOAD" \ "${spec[@]}" \ >> "$LOG" 2>&1 < /dev/null & diff --git a/tests/dsv4/test_dsv4_dspark_wiring.py b/tests/dsv4/test_dsv4_dspark_wiring.py index 825d9d5..e48526d 100644 --- a/tests/dsv4/test_dsv4_dspark_wiring.py +++ b/tests/dsv4/test_dsv4_dspark_wiring.py @@ -160,3 +160,15 @@ def test_the_forward_passes_them_for_a_speculative_batch(self): "prefill_batched takes logit_indices but the forward never passes it, so a " "speculative verify scores only each request's last token" ) + + +class TestMHCBoundsLargePrefillResiduals: + def test_target_and_drafter_use_the_fused_inv_rms(self): + model = (PKG / "models" / "deepseek_v4" / "model.py").read_text() + dspark = (PKG / "models" / "deepseek_v4" / "dspark.py").read_text() + assert "inv_rms(x2d, self.norm_eps)" in model + assert "inv_rms(x2d, self.norm_eps)" in dspark + norm = (PKG / "kernel" / "triton" / "dsv4" / "norm.py").read_text() + assert "def _inv_rms_kernel" in norm + assert "x.float().square()" not in model + assert "x.float().square()" not in dspark diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py index a164f0b..7cfbc18 100644 --- a/tests/engine/test_cache_budget.py +++ b/tests/engine/test_cache_budget.py @@ -168,6 +168,15 @@ def test_adjust_config_allows_auto_for_dsv4(): assert cfg.page_size == 128 # DSV4's KV page is the P-token window page +def test_adjust_config_honors_dsv4_prefill_chunk_limit(): + """DSV4 continuation prefill is stateful, so a safety cap must survive resolution.""" + from freetoken.engine.engine import _adjust_config + + cfg = _dsv4_adjust_cfg(max_seq_len=8192, max_extend_tokens=1024) + _adjust_config(cfg) + assert cfg.max_extend_tokens == 1024 + + def test_adjust_config_resolves_num_tokens_for_dsv4(): # --num-tokens resolves AFTER every page_size override, so DSV4's P=128 page divides it. from freetoken.engine.engine import _adjust_config From ad423c66ce6e9a079b465ca0dcb81b0a8cf2c610 Mon Sep 17 00:00:00 2001 From: Carlos Alvarado Date: Sun, 23 Aug 2026 11:33:09 -0400 Subject: [PATCH 09/13] perf(nccl): all-reduce TP tensors in place Avoid the symmetric-window staging buffer and its two device copies for every tensor-parallel reduction. NCCL can reduce the contiguous input directly in place, which also avoids the symmetric collective path that dominated the TP4 verify profile on PCIe RTX A4000 GPUs. --- python/freetoken/kernel/csrc/src/pynccl.cu | 61 ++++------------------ python/freetoken/kernel/pynccl.py | 2 - 2 files changed, 11 insertions(+), 52 deletions(-) diff --git a/python/freetoken/kernel/csrc/src/pynccl.cu b/python/freetoken/kernel/csrc/src/pynccl.cu index 9d41e00..33ffcc0 100644 --- a/python/freetoken/kernel/csrc/src/pynccl.cu +++ b/python/freetoken/kernel/csrc/src/pynccl.cu @@ -72,22 +72,12 @@ inline constexpr auto template_fn = struct NCCLWrapper : public tvm::ffi::Object { public: NCCLWrapper(int rank, int world_size, const size_t max_bytes, NCCLIDList uid) - : m_rank(rank), m_world_size(world_size), m_max_bytes(max_bytes) { + : m_rank(rank), m_world_size(world_size) { + (void)max_bytes; ncclUniqueId id = get_uid(uid); ncclComm_t comm; NCCL_CHECK(::ncclCommInitRank(&comm, m_world_size, id, m_rank)); m_comm = {comm, template_fn<::ncclCommDestroy>}; - - void *buf; - NCCL_CHECK(::ncclMemAlloc(&buf, max_bytes)); - m_sym_mem = {buf, template_fn<::ncclMemFree>}; - - ncclWindow_t win; - NCCL_CHECK(::ncclCommWindowRegister(comm, buf, max_bytes, &win, - NCCL_WIN_COLL_SYMMETRIC)); - m_win = {win, [comm = m_comm](ncclWindow_t w) { - return NCCL_CHECK(::ncclCommWindowDeregister(comm.get(), w)); - }}; } auto all_reduce(tvm::ffi::TensorView t, std::string op) const -> void { @@ -97,40 +87,17 @@ public: RuntimeCheck(t.is_contiguous(), "Tensor must be contiguous"); const auto size_dim = static_cast(t.shape().Product()); const auto dtype = kNCCLDtypeMap.at(t.dtype()); - const auto size_bytes = size_dim * (t.dtype().bits / 8); const auto data_ptr = t.data_ptr(); const auto reduce_op = kNCCLReduceOPMap.at(op); const auto stream = LaunchKernel::resolve_device(t.device()); - - if (size_bytes <= m_max_bytes) { // use internal buffer - const auto buf_ptr = m_sym_mem.get(); - const auto need_memcpy = (buf_ptr != data_ptr); - if (need_memcpy) { - CUDA_CHECK(::cudaMemcpyAsync(buf_ptr, data_ptr, size_bytes, - ::cudaMemcpyDeviceToDevice, stream)); - } - NCCL_CHECK(::ncclAllReduce( - /*sendbuff=*/buf_ptr, - /*recvbuff=*/buf_ptr, - /*count=*/size_dim, - /*datatype=*/dtype, - /*op=*/reduce_op, - /*comm=*/m_comm.get(), - /*stream=*/stream)); - if (need_memcpy) { - CUDA_CHECK(::cudaMemcpyAsync(data_ptr, buf_ptr, size_bytes, - ::cudaMemcpyDeviceToDevice, stream)); - } - } else { - NCCL_CHECK(::ncclAllReduce( - /*sendbuff=*/data_ptr, - /*recvbuff=*/data_ptr, - /*count=*/size_dim, - /*datatype=*/dtype, - /*op=*/reduce_op, - /*comm=*/m_comm.get(), - /*stream=*/stream)); - } + NCCL_CHECK(::ncclAllReduce( + /*sendbuff=*/data_ptr, + /*recvbuff=*/data_ptr, + /*count=*/size_dim, + /*datatype=*/dtype, + /*op=*/reduce_op, + /*comm=*/m_comm.get(), + /*stream=*/stream)); } auto all_gather(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) const @@ -160,18 +127,13 @@ public: /*stream=*/stream)); } - auto get_buffer() const -> void * { return m_sym_mem.get(); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("freetoken.NCCLWrapper", NCCLWrapper, tvm::ffi::Object); private: int m_rank; int m_world_size; - size_t m_max_bytes; shared_obj m_comm; - shared_ptr m_sym_mem; - shared_obj m_win; }; TVM_FFI_STATIC_INIT_BLOCK() { @@ -179,8 +141,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::ObjectDef() .def(refl::init(), "__init__") .def("all_reduce", &NCCLWrapper::all_reduce) - .def("all_gather", &NCCLWrapper::all_gather) - .def("get_buffer", &NCCLWrapper::get_buffer); + .def("all_gather", &NCCLWrapper::all_gather); } TVM_FFI_DLL_EXPORT_TYPED_FUNC(create_nccl_uid, &create_uid); diff --git a/python/freetoken/kernel/pynccl.py b/python/freetoken/kernel/pynccl.py index 23ea573..d2fe424 100644 --- a/python/freetoken/kernel/pynccl.py +++ b/python/freetoken/kernel/pynccl.py @@ -18,8 +18,6 @@ class PyNCCLCommunicator: def all_reduce(self, input: torch.Tensor, op: Literal["sum"]) -> None: ... @abstractmethod def all_gather(self, output: torch.Tensor, input: torch.Tensor) -> None: ... - @abstractmethod - def get_buffer(self) -> int: ... else: PyNCCLCommunicator = Any From 507f172179439af925ad256e953db8d083cd08e2 Mon Sep 17 00:00:00 2001 From: Carlos Alvarado Date: Sun, 23 Aug 2026 12:21:36 -0400 Subject: [PATCH 10/13] feat(engine): report final NUMA-aware model placement Gather the post-capture serving footprint from every TP rank immediately before the scheduler ready acknowledgement. The summary maps ranks to NUMA nodes and GPUs, showing actual VRAM use, host expert-bank bytes, CPU-MoE workers and coordinator cores, per-node RAM pressure, workstation totals, and the counted CPU/GPU placement split. Keep the formatter and bank accounting independently testable without loading a model. --- python/freetoken/engine/distribution.py | 142 ++++++++++++++++++++++ python/freetoken/engine/engine.py | 18 +++ python/freetoken/server/launch.py | 5 + tests/engine/test_distribution_summary.py | 59 +++++++++ 4 files changed, 224 insertions(+) create mode 100644 python/freetoken/engine/distribution.py create mode 100644 tests/engine/test_distribution_summary.py diff --git a/python/freetoken/engine/distribution.py b/python/freetoken/engine/distribution.py new file mode 100644 index 0000000..a54359a --- /dev/null +++ b/python/freetoken/engine/distribution.py @@ -0,0 +1,142 @@ +"""Final CPU/GPU placement summary for a loaded engine. + +The startup banner describes hardware. This module describes where the loaded model +actually landed after host expert banks, GPU caches, and CUDA graphs have reached their +serving sizes. The data is gathered from every TP rank immediately before readiness. +""" + +from __future__ import annotations + +import os +from collections import defaultdict +from typing import Any + +import torch + +from freetoken.utils import numa + + +def _gib(n_bytes: int) -> str: + return f"{n_bytes / (1 << 30):.2f} GiB" + + +def _percent(part: int, whole: int) -> float: + return 100.0 * part / whole if whole > 0 else 0.0 + + +def host_expert_bytes(cache: Any) -> int: + """Pinned host expert bytes owned by one TP rank.""" + if cache is None: + return 0 + return sum( + tensor.numel() * tensor.element_size() + for layers in getattr(cache, "bank_sources", {}).values() + for tensor in layers + ) + + +def _physical_core_count(cpus: list[int]) -> int: + siblings = {numa.thread_siblings(cpu) or f"cpu:{cpu}" for cpu in cpus} + return len(siblings) + + +def _node_memory_total(node: int | None) -> int: + if node is None: + return 0 + try: + with open(f"/sys/devices/system/node/node{node}/meminfo") as f: + for line in f: + if " MemTotal:" in line: + return int(line.split()[-2]) * 1024 + except (OSError, ValueError, IndexError): + pass + return 0 + + +def _system_memory_total() -> int: + try: + return int(os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES")) + except (OSError, ValueError): + return 0 + + +def local_distribution_snapshot(engine: Any) -> dict[str, Any]: + """Collect one rank's final serving placement without synchronizing a hot path.""" + rank = int(engine.config.tp_info.rank) + free, total = torch.cuda.mem_get_info(engine.device) + placed = numa.placement() + node = placed[0] if placed is not None else numa.device_numa_node(rank) + node_cpus = list(placed[1]) if placed is not None else numa._allowed_cpus() + executor = getattr(engine, "cpu_moe_executor", None) + workers = int(getattr(executor, "num_threads", 0) or 0) + coordinator = int(getattr(executor, "_coord_core", -1) >= 0) + return { + "rank": rank, + "gpu": int(engine.device.index if engine.device.index is not None else rank), + "gpu_name": torch.cuda.get_device_name(engine.device), + "gpu_used": int(total - free), + "gpu_total": int(total), + "node": node, + "node_memory_total": _node_memory_total(node), + "node_physical_cores": _physical_core_count(node_cpus), + "cpu_workers": workers, + "cpu_coordinator": coordinator, + "host_experts": host_expert_bytes(getattr(engine, "moe_offload_cache", None)), + "system_memory_total": _system_memory_total(), + } + + +def format_distribution_summary(rows: list[dict[str, Any]]) -> list[str]: + """Format gathered TP snapshots as a compact NUMA-aware workstation map.""" + if not rows: + return [] + rows = sorted(rows, key=lambda row: int(row["rank"])) + by_node: dict[int | None, list[dict[str, Any]]] = defaultdict(list) + for row in rows: + by_node[row.get("node")].append(row) + + lines = ["Model distribution (final, NUMA-aware):"] + for node, node_rows in sorted( + by_node.items(), key=lambda item: (-1 if item[0] is None else int(item[0])) + ): + host = sum(int(row["host_experts"]) for row in node_rows) + memory_total = max(int(row["node_memory_total"]) for row in node_rows) + cores = max(int(row["node_physical_cores"]) for row in node_rows) + assigned = sum( + int(row["cpu_workers"]) + int(row["cpu_coordinator"]) + for row in node_rows + ) + label = "NUMA unknown" if node is None else f"NUMA node {node}" + lines.append( + f" {label}: host experts {_gib(host)}/{_gib(memory_total)} " + f"({_percent(host, memory_total):.1f}% RAM), CPU cores {assigned}/{cores} assigned" + ) + for row in node_rows: + coord = " + 1 coordinator" if row["cpu_coordinator"] else "" + lines.append( + f" rank {row['rank']} -> GPU {row['gpu']} " + f"{_gib(int(row['gpu_used']))}/{_gib(int(row['gpu_total']))} " + f"({_percent(int(row['gpu_used']), int(row['gpu_total'])):.1f}% VRAM), " + f"CPU-MoE {row['cpu_workers']} workers{coord}, " + f"host experts {_gib(int(row['host_experts']))}" + ) + + host = sum(int(row["host_experts"]) for row in rows) + gpu_used = sum(int(row["gpu_used"]) for row in rows) + gpu_total = sum(int(row["gpu_total"]) for row in rows) + system_memory = max(int(row["system_memory_total"]) for row in rows) + counted = host + gpu_used + lines.append( + f" total: CPU host experts {_gib(host)}/{_gib(system_memory)} " + f"({_percent(host, system_memory):.1f}% RAM), GPU {_gib(gpu_used)}/{_gib(gpu_total)} " + f"({_percent(gpu_used, gpu_total):.1f}% VRAM); counted placement " + f"CPU {_percent(host, counted):.1f}% / GPU {_percent(gpu_used, counted):.1f}%" + ) + return lines + + +__all__ = [ + "format_distribution_summary", + "host_expert_bytes", + "local_distribution_snapshot", +] diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 7b0f8ed..95f7fbc 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -577,6 +577,24 @@ def __init__(self, config: EngineConfig): # Prefill runs on the first comma part; warm its autotune cache. self._warmup_prefill() + def log_distribution_summary(self) -> None: + """Gather final TP placement and print it immediately before readiness.""" + from freetoken.engine.distribution import ( + format_distribution_summary, + local_distribution_snapshot, + ) + + tp = get_tp_info() + rows: list[dict[str, Any] | None] = [None] * tp.size + torch.distributed.all_gather_object( + rows, + local_distribution_snapshot(self), + group=self.tp_cpu_group, + ) + if tp.is_primary(): + for line in format_distribution_summary([row for row in rows if row is not None]): + logger.info(line) + def _init_communication(self, config: EngineConfig) -> torch.distributed.ProcessGroup: if config.tp_info.size == 1 or config.use_pynccl: torch.distributed.init_process_group( diff --git a/python/freetoken/server/launch.py b/python/freetoken/server/launch.py index c1f8cca..ef29d95 100644 --- a/python/freetoken/server/launch.py +++ b/python/freetoken/server/launch.py @@ -91,6 +91,11 @@ def _terminate(signum, _frame): try: scheduler = Scheduler(args) scheduler.sync_all_ranks() + # All final serving allocations now exist (host expert banks, GPU caches, + # CUDA graphs). Every rank contributes its real placement before rank 0 + # sends the ready ack, so the log maps the loaded model rather than merely + # describing the workstation's hardware. + scheduler.engine.log_distribution_summary() except KeyboardInterrupt: # A stop during the ~90s load. KeyboardInterrupt is a BaseException, so it # slips past the handler below and prints a traceback from wherever the diff --git a/tests/engine/test_distribution_summary.py b/tests/engine/test_distribution_summary.py new file mode 100644 index 0000000..c864dd6 --- /dev/null +++ b/tests/engine/test_distribution_summary.py @@ -0,0 +1,59 @@ +from types import SimpleNamespace + +import torch + +from freetoken.engine.distribution import format_distribution_summary, host_expert_bytes + + +def test_host_expert_bytes_sums_every_rank_local_bank(): + cache = SimpleNamespace( + bank_sources={ + "gate": [torch.empty(2, 8, dtype=torch.uint8), torch.empty(2, 8, dtype=torch.uint8)], + "down": [torch.empty(2, 4, dtype=torch.bfloat16)], + } + ) + assert host_expert_bytes(cache) == 48 + assert host_expert_bytes(None) == 0 + + +def test_summary_groups_ranks_by_numa_and_reports_cpu_gpu_split(): + gib = 1 << 30 + rows = [ + { + "rank": 0, + "gpu": 0, + "gpu_name": "GPU", + "gpu_used": 12 * gib, + "gpu_total": 16 * gib, + "node": 0, + "node_memory_total": 128 * gib, + "node_physical_cores": 20, + "cpu_workers": 9, + "cpu_coordinator": 1, + "host_experts": 32 * gib, + "system_memory_total": 256 * gib, + }, + { + "rank": 1, + "gpu": 1, + "gpu_name": "GPU", + "gpu_used": 12 * gib, + "gpu_total": 16 * gib, + "node": 0, + "node_memory_total": 128 * gib, + "node_physical_cores": 20, + "cpu_workers": 9, + "cpu_coordinator": 1, + "host_experts": 32 * gib, + "system_memory_total": 256 * gib, + }, + ] + + lines = format_distribution_summary(rows) + + assert lines[0] == "Model distribution (final, NUMA-aware):" + assert "host experts 64.00 GiB/128.00 GiB (50.0% RAM)" in lines[1] + assert "CPU cores 20/20 assigned" in lines[1] + assert "GPU 0 12.00 GiB/16.00 GiB (75.0% VRAM)" in lines[2] + assert "CPU-MoE 9 workers + 1 coordinator" in lines[2] + assert "counted placement CPU 72.7% / GPU 27.3%" in lines[-1] From 9f4595dcc92e1df692994bdb4853c2a60f098df7 Mon Sep 17 00:00:00 2001 From: Carlos Alvarado Date: Sun, 23 Aug 2026 13:37:19 -0400 Subject: [PATCH 11/13] fix(engine): make placement report self-contained --- python/freetoken/engine/distribution.py | 17 +++++++++++++++-- tests/engine/test_distribution_summary.py | 8 ++++++++ 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/python/freetoken/engine/distribution.py b/python/freetoken/engine/distribution.py index a54359a..93ba77b 100644 --- a/python/freetoken/engine/distribution.py +++ b/python/freetoken/engine/distribution.py @@ -35,8 +35,19 @@ def host_expert_bytes(cache: Any) -> int: ) +def _thread_siblings(cpu: int) -> str | None: + """Return the kernel topology key without depending on optional NUMA helpers.""" + try: + with open( + f"/sys/devices/system/cpu/cpu{cpu}/topology/thread_siblings_list" + ) as f: + return f.read().strip() + except OSError: + return None + + def _physical_core_count(cpus: list[int]) -> int: - siblings = {numa.thread_siblings(cpu) or f"cpu:{cpu}" for cpu in cpus} + siblings = {_thread_siblings(cpu) or f"cpu:{cpu}" for cpu in cpus} return len(siblings) @@ -66,7 +77,9 @@ def local_distribution_snapshot(engine: Any) -> dict[str, Any]: free, total = torch.cuda.mem_get_info(engine.device) placed = numa.placement() node = placed[0] if placed is not None else numa.device_numa_node(rank) - node_cpus = list(placed[1]) if placed is not None else numa._allowed_cpus() + node_cpus = ( + list(placed[1]) if placed is not None else sorted(numa.allowed_cpus()) + ) executor = getattr(engine, "cpu_moe_executor", None) workers = int(getattr(executor, "num_threads", 0) or 0) coordinator = int(getattr(executor, "_coord_core", -1) >= 0) diff --git a/tests/engine/test_distribution_summary.py b/tests/engine/test_distribution_summary.py index c864dd6..f472145 100644 --- a/tests/engine/test_distribution_summary.py +++ b/tests/engine/test_distribution_summary.py @@ -2,6 +2,7 @@ import torch +from freetoken.engine import distribution from freetoken.engine.distribution import format_distribution_summary, host_expert_bytes @@ -16,6 +17,13 @@ def test_host_expert_bytes_sums_every_rank_local_bank(): assert host_expert_bytes(None) == 0 +def test_physical_core_count_uses_sysfs_sibling_groups(monkeypatch): + siblings = {0: "0,40", 40: "0,40", 1: "1,41", 41: "1,41"} + monkeypatch.setattr(distribution, "_thread_siblings", siblings.get) + + assert distribution._physical_core_count([0, 1, 40, 41]) == 2 + + def test_summary_groups_ranks_by_numa_and_reports_cpu_gpu_split(): gib = 1 << 30 rows = [ From 2ce00703fa20feba612d23004b3eab53888ba5b2 Mon Sep 17 00:00:00 2001 From: Carlos Alvarado Date: Sun, 23 Aug 2026 11:33:09 -0400 Subject: [PATCH 12/13] perf(moe): prioritize repeated speculative routes Spend the paper-derived hybrid fetch budget on missing experts used by the most rows in the current verify block. Retain expert recency and id as deterministic tie-breakers so the existing steady-state cache policy remains intact. --- python/freetoken/moe/offload_kernels.py | 42 ++++++++++++++++--------- tests/moe/test_hybrid_fetch.py | 18 +++++++++++ 2 files changed, 45 insertions(+), 15 deletions(-) diff --git a/python/freetoken/moe/offload_kernels.py b/python/freetoken/moe/offload_kernels.py index cf513f5..e1f2265 100644 --- a/python/freetoken/moe/offload_kernels.py +++ b/python/freetoken/moe/offload_kernels.py @@ -8,8 +8,11 @@ from flashlib.kernels.slot_cache import lru_ensure # Hybrid backend: which of a step's missing experts to fetch (when capped below the miss -# count). "recency" (default) fetches the experts most-recently active before this step -# (LRU on the expert -> prioritizes recurring misses, lowering the steady miss rate); +# count). "recency" (default) first fetches misses used by the most routes in THIS +# batch, then the experts most-recently active before this step (LRU on the expert). +# The route-count priority spends each paper-derived q* PCIe copy where it removes the +# most repeated CPU GEMVs during a multi-row speculative verify; recency breaks ties and +# preserves the steady-state hit-rate policy. # "lowest_id" fetches the smallest expert ids (the original, routing-blind heuristic). _HYBRID_FETCH_BY_RECENCY = ( os.getenv("FREETOKEN_HYBRID_FETCH", "recency").strip().lower() != "lowest_id" @@ -132,9 +135,11 @@ def _ensure_experts_hybrid_cpu( the GPU path; see tests/test_offload_lru_kernels.py). Fetches at most ``max_fetch`` (or the bandwidth-matched ``~frac_q16/2^16 * misses`` when ``frac_q16`` > 0) of the missing experts; overflow misses are rewritten to -1. With ``BY_RECENCY`` the fetch set is the - most-recently-active misses (ties -> lower id); else the lowest ids.""" + highest-current-route-count, most-recently-active misses (ties -> lower id); + else the lowest ids.""" + flat_experts = expert_ids.view(-1).tolist() seen = [] - for expert in expert_ids.view(-1).tolist(): + for expert in flat_experts: if expert not in seen: seen.append(expert) @@ -150,9 +155,10 @@ def _ensure_experts_hybrid_cpu( cache.usage[slot] = step missing = [e for e in seen if int(cache.slot_for_id[layer_id, e].item()) == -1] + route_counts = {e: flat_experts.count(e) for e in seen} if _HYBRID_FETCH_BY_RECENCY: rec = cache.expert_recency[layer_id].tolist() - missing.sort(key=lambda e: (-rec[e], e)) + missing.sort(key=lambda e: (-route_counts[e], -rec[e], e)) else: missing.sort() if frac_q16 > 0: @@ -321,10 +327,11 @@ def _ensure_experts_hybrid_kernel( = the capped fetch count (copy_missing), ``num_missing_full`` = the pre-cap miss count (stats). - Which misses to fetch is the cap policy. ``BY_RECENCY`` (default) fetches the experts - most-recently active before this step (LRU on the expert, via ``expert_recency``), - breaking ties toward the lower expert id -- this prioritizes *recurring* misses for - caching, lowering the steady miss rate. Otherwise the lowest expert ids are fetched + Which misses to fetch is the cap policy. ``BY_RECENCY`` (default) first prioritizes + the experts used by the most routes in the current batch, so one PCIe copy removes + as many repeated CPU GEMVs as possible during a speculative verify. It then uses + most-recent activity (LRU on the expert, via ``expert_recency``), breaking final ties + toward the lower expert id. Otherwise the lowest expert ids are fetched (``missing_rank``), the original routing-blind heuristic.""" step = tl.load(step_ptr) + 1 tl.store(step_ptr, step) @@ -333,10 +340,11 @@ def _ensure_experts_hybrid_kernel( # ---- Phase 1: active + missing over experts ---- off_e = tl.arange(0, BLOCK_E) e_mask = off_e < num_experts - is_active = tl.zeros((BLOCK_E,), dtype=tl.int1) + route_count = tl.zeros((BLOCK_E,), dtype=tl.int32) for i in tl.range(num_active): e = tl.load(expert_ids_ptr + i) - is_active = is_active | (off_e == e) + route_count += (off_e == e).to(tl.int32) + is_active = route_count > 0 tl.store(active_mask_ptr + off_e, is_active.to(tl.int32), mask=e_mask) slot = tl.load(slot_for_id_ptr + base + off_e, mask=e_mask, other=-1) is_missing = is_active & (slot == -1) & e_mask @@ -358,13 +366,17 @@ def _ensure_experts_hybrid_kernel( is_hit = is_active & (slot >= 0) tl.store(usage_ptr + slot, step, mask=is_hit) - # Fetch-selection priority: encode (recency desc, id asc) into one strictly-ordered - # score so argmax has no ties (rec deltas are multiples of num_experts; the id term - # spans only [0, num_experts), so it can only break exact-recency ties). + # Fetch-selection priority: encode (current route count desc, recency desc, id asc) + # into one strictly-ordered score. Previous recency is <= step-1, so multiplying + # route_count by step+1 makes one additional current route dominate the full recency + # range. The id term spans only [0, num_experts), breaking exact ties only. if BY_RECENCY: rec = tl.load(expert_recency_ptr + base + off_e, mask=e_mask, other=-1).to(tl.int64) score = tl.where( - is_missing, rec * num_experts + (num_experts - 1 - off_e), -1152921504606846976 + is_missing, + (route_count.to(tl.int64) * (step + 1) + rec) * num_experts + + (num_experts - 1 - off_e), + -1152921504606846976, ).to(tl.int64) else: missing_rank = tl.cumsum(is_missing.to(tl.int32)) - 1 diff --git a/tests/moe/test_hybrid_fetch.py b/tests/moe/test_hybrid_fetch.py index e080e16..c072b2f 100644 --- a/tests/moe/test_hybrid_fetch.py +++ b/tests/moe/test_hybrid_fetch.py @@ -106,3 +106,21 @@ def test_hybrid_fixed_cap_unchanged(): cache.ensure_experts_hybrid(0, ids) assert int(cache.num_missing_full.item()) == 8 assert int(cache.num_indices.item()) == 1 + + +def test_hybrid_fetch_spends_cap_on_most_repeated_current_miss(): + """q* still counts unique misses, but its copies should remove the most CPU routes.""" + cache = OffloadMoeCache( + num_layers=1, num_experts=16, cache_size=20, device=torch.device("cpu"), + quant_format="bf16", decode_target="hybrid", hybrid_max_fetch=1, + ) + # Expert 2 was historically recent, but expert 7 occurs three times in this + # multi-row verify. Fetching 7 avoids three CPU GEMVs for the price of one copy. + cache.step.fill_(11) + cache.expert_recency[0, 2] = 10 + ids = torch.tensor([7, 2, 7, 3, 7], dtype=torch.int32) + cache.ensure_experts_hybrid(0, ids) + + assert ids[0] >= 0 and ids[2] == ids[0] and ids[4] == ids[0] + assert ids[1] == -1 and ids[3] == -1 + assert int(cache.src_indices[0]) == 7 From 16ce51a90ec074cdddff385929cc1644f7a5eb77 Mon Sep 17 00:00:00 2001 From: Carlos Alvarado Date: Sun, 23 Aug 2026 16:45:17 -0400 Subject: [PATCH 13/13] perf(dspark): bypass persistently weak drafts --- python/freetoken/engine/config.py | 7 ++ python/freetoken/engine/engine.py | 30 ++++- python/freetoken/models/deepseek_v4/dspark.py | 115 ++++++++++++++++++ python/freetoken/scheduler/scheduler.py | 7 ++ python/freetoken/server/args.py | 25 ++++ scripts/freetoken-dsv4.sh | 14 +++ tests/dsv4/test_dsv4_adaptive_verification.py | 59 ++++++++- 7 files changed, 255 insertions(+), 2 deletions(-) diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 903ec41..7fddd5a 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -82,6 +82,13 @@ class EngineConfig: # DeepSeek-V4 only: build the checkpoint's dSpark drafter for block speculative # decoding. Reaches the model through dsv4_args.dspark_enabled (_adjust_dsv4_config). speculative_dspark: bool = False + # Experimental, request-local DSpark circuit breaker. 0 disables it. Once at least + # dspark_fallback_min_drafted proposals have been measured below this acceptance + # rate, use ordinary target decode for dspark_fallback_steps steps, then probe the + # drafter again. This is based on observed acceptance, not prompt classification. + dspark_fallback_acceptance: float = 0.0 + dspark_fallback_min_drafted: int = 32 + dspark_fallback_steps: int = 64 use_pynccl: bool = True max_seq_len_override: int | None = None num_page_override: int | None = None # if not None, will override the number of pages diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 95f7fbc..78d6edc 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -1067,8 +1067,35 @@ def _init_dspark_adaptive_verification(self) -> None: from freetoken.models.deepseek_v4.dspark import DSparkAdaptiveVerification self._adaptive_verification = DSparkAdaptiveVerification( - block_size, curve, self.device + block_size, + curve, + self.device, + fallback_acceptance=self.config.dspark_fallback_acceptance, + fallback_min_drafted=self.config.dspark_fallback_min_drafted, + fallback_steps=self.config.dspark_fallback_steps, ) + if self.config.dspark_fallback_acceptance > 0: + logger.info_rank0( + "DSpark acceptance fallback enabled: threshold=%.1f%%, " + "min-drafted=%d, target-only-steps=%d", + 100.0 * self.config.dspark_fallback_acceptance, + self.config.dspark_fallback_min_drafted, + self.config.dspark_fallback_steps, + ) + + def should_speculate(self, req: Req) -> bool: + """Return whether this request should pay for a DSpark draft this step.""" + manager = self._adaptive_verification + if manager is None: + return True + return manager.should_speculate(req.uid) + + def _record_dspark_acceptance( + self, req: Req, accepted: int, drafted: int + ) -> None: + manager = self._adaptive_verification + if manager is not None: + manager.record_acceptance(req.uid, accepted, drafted) def adapt_speculative_batch(self, batch: Batch) -> None: """Compact one prepared DSpark block to the paper-selected prefix width. @@ -1306,6 +1333,7 @@ def _finish_speculative( ) self._spec_accepted += n_acc self._spec_drafted += width + self._record_dspark_acceptance(req, n_acc, width) selected_rows.append(off + n_acc) accepted_counts.append(n_acc) off += span diff --git a/python/freetoken/models/deepseek_v4/dspark.py b/python/freetoken/models/deepseek_v4/dspark.py index fc11daf..a2df857 100644 --- a/python/freetoken/models/deepseek_v4/dspark.py +++ b/python/freetoken/models/deepseek_v4/dspark.py @@ -47,6 +47,73 @@ logger = init_logger(__name__) +class DSparkAcceptanceFallback: + """Request-local measured-acceptance circuit breaker. + + DSpark's confidence scheduler chooses how many proposals to verify after the + drafter has run. It cannot recover that draft cost when a high-entropy request + repeatedly rejects the proposals. This controller observes actual target + acceptance over a bounded window, temporarily selects ordinary target decoding, + and then probes DSpark again. It deliberately knows nothing about the prompt. + + ``take_decision`` returns ``(should_speculate, is_resume_probe)``. A tripped + controller skips exactly ``cooldown_steps`` calls before returning one probe. + """ + + def __init__( + self, + threshold: float, + min_drafted: int, + cooldown_steps: int, + ) -> None: + if not math.isfinite(threshold) or not 0 < threshold <= 1: + raise ValueError("DSpark fallback threshold must be in (0, 1]") + if min_drafted < 1: + raise ValueError("DSpark fallback min_drafted must be >= 1") + if cooldown_steps < 1: + raise ValueError("DSpark fallback cooldown_steps must be >= 1") + self.threshold = float(threshold) + self.min_drafted = int(min_drafted) + self.cooldown_steps = int(cooldown_steps) + self.reset() + + def reset(self) -> None: + self.accepted = 0 + self.drafted = 0 + self.cooldown_left = 0 + self._resume_probe = False + + def take_decision(self) -> tuple[bool, bool]: + if self.cooldown_left > 0: + self.cooldown_left -= 1 + if self.cooldown_left == 0: + self._resume_probe = True + return False, False + resumed = self._resume_probe + self._resume_probe = False + return True, resumed + + def record(self, accepted: int, drafted: int) -> tuple[float, int] | None: + if drafted < 0 or accepted < 0 or accepted > drafted: + raise ValueError( + f"invalid DSpark acceptance sample {accepted}/{drafted}" + ) + self.accepted += int(accepted) + self.drafted += int(drafted) + if self.drafted < self.min_drafted: + return None + + measured_drafted = self.drafted + rate = self.accepted / measured_drafted + self.accepted = 0 + self.drafted = 0 + if rate >= self.threshold: + return None + self.cooldown_left = self.cooldown_steps + self._resume_probe = False + return rate, measured_drafted + + def choose_adaptive_draft_width( stale_confidence: Sequence[float] | torch.Tensor, draft_cost_ms: float, @@ -108,6 +175,9 @@ def __init__( block_size: int, verify_curve: Sequence[tuple[int, float]], device: torch.device, + fallback_acceptance: float = 0.0, + fallback_min_drafted: int = 32, + fallback_steps: int = 64, ) -> None: expected_spans = list(range(1, block_size + 2)) curve = sorted((int(span), float(cost)) for span, cost in verify_curve) @@ -136,6 +206,13 @@ def __init__( maxlen=self._DRAFT_PROFILE_SAMPLES ) self._debug_left = 12 + self._acceptance_fallback = ( + DSparkAcceptanceFallback( + fallback_acceptance, fallback_min_drafted, fallback_steps + ) + if fallback_acceptance > 0 + else None + ) @property def needs_draft_profile(self) -> bool: @@ -166,6 +243,43 @@ def _reset_request(self, uid: int) -> None: self._events = [None, None] self._stale_idx = 0 self._active_uid = uid + if self._acceptance_fallback is not None: + self._acceptance_fallback.reset() + + def should_speculate(self, uid: int) -> bool: + if self._active_uid != uid: + self._reset_request(uid) + fallback = self._acceptance_fallback + if fallback is None: + return True + speculate, resumed = fallback.take_decision() + if resumed: + logger.info_rank0( + "DSpark fallback: request %d probing speculation after %d " + "target-only steps", + uid, + fallback.cooldown_steps, + ) + return speculate + + def record_acceptance(self, uid: int, accepted: int, drafted: int) -> None: + if self._active_uid != uid: + self._reset_request(uid) + fallback = self._acceptance_fallback + if fallback is None: + return + trip = fallback.record(accepted, drafted) + if trip is None: + return + rate, measured_drafted = trip + logger.info_rank0( + "DSpark fallback: request %d accepted %.1f%% of %d proposals; " + "target-only for %d steps", + uid, + 100.0 * rate, + measured_drafted, + fallback.cooldown_steps, + ) def record_and_choose(self, confidence: torch.Tensor, uid: int) -> int: """Publish this block's confidence and choose from the stale buffer.""" @@ -762,6 +876,7 @@ def window_cols_for_block( __all__ = [ "ConfidenceHead", + "DSparkAcceptanceFallback", "DSparkAdaptiveVerification", "DSparkDrafter", "MarkovHead", diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index ec67706..8e8d3df 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -941,6 +941,13 @@ def _maybe_make_speculative(self, batch: Batch) -> None: self._warned_non_greedy = False if spec is None or not batch.is_decode: return + should_speculate = getattr( + getattr(self, "engine", None), "should_speculate", None + ) + if should_speculate is not None and any( + not should_speculate(req) for req in batch.reqs + ): + return gamma = spec.block_size if gamma < 1: return diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 8dda7de..c3c67f1 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -234,6 +234,31 @@ def _infer_reasoning_parser(model_path: str) -> str | None: ), ) + parser.add_argument( + "--dspark-fallback-acceptance", + type=_parse_moe_cache_rate, + default=ServerArgs.dspark_fallback_acceptance, + help=( + "Experimental request-local DSpark fallback threshold in [0, 1]. " + "After measured proposal acceptance falls below this value, temporarily " + "use ordinary target decode; 0 disables the fallback." + ), + ) + + parser.add_argument( + "--dspark-fallback-min-drafted", + type=_positive_int, + default=ServerArgs.dspark_fallback_min_drafted, + help="Minimum measured DSpark proposals before evaluating the fallback.", + ) + + parser.add_argument( + "--dspark-fallback-steps", + type=_positive_int, + default=ServerArgs.dspark_fallback_steps, + help="Ordinary target-decode steps before probing DSpark again.", + ) + parser.add_argument( "--distributed-timeout", type=float, diff --git a/scripts/freetoken-dsv4.sh b/scripts/freetoken-dsv4.sh index 7475f5e..cb25348 100644 --- a/scripts/freetoken-dsv4.sh +++ b/scripts/freetoken-dsv4.sh @@ -29,6 +29,11 @@ EXPERT_LOAD="${EXPERT_LOAD:-serial}" # disappear from the log's `layers=` count, which is the quickest way to tell which # mode a run was in. SPECULATIVE_DSPARK="${SPECULATIVE_DSPARK:-1}" +# Experimental request-local circuit breaker. A zero acceptance threshold leaves +# production behavior unchanged. The measured Lenovo A/B starts at 0.60/32/64. +DSPARK_FALLBACK_ACCEPTANCE="${DSPARK_FALLBACK_ACCEPTANCE:-0}" +DSPARK_FALLBACK_MIN_DRAFTED="${DSPARK_FALLBACK_MIN_DRAFTED:-32}" +DSPARK_FALLBACK_STEPS="${DSPARK_FALLBACK_STEPS:-64}" # Which base dspark_target_layer_ids is read with (0 or 1). Exported so it reaches the # rank subprocesses; unset means the model's own default. See the note in model.py -- # this is settled by measurement, not by the config. @@ -118,6 +123,14 @@ cmd_start() { local spec=() [ "$SPECULATIVE_DSPARK" = "1" ] && spec=(--speculative-dspark) + local fallback=() + if [ "$DSPARK_FALLBACK_ACCEPTANCE" != "0" ] && [ "$DSPARK_FALLBACK_ACCEPTANCE" != "0.0" ]; then + fallback=( + --dspark-fallback-acceptance "$DSPARK_FALLBACK_ACCEPTANCE" + --dspark-fallback-min-drafted "$DSPARK_FALLBACK_MIN_DRAFTED" + --dspark-fallback-steps "$DSPARK_FALLBACK_STEPS" + ) + fi echo "starting DeepSeek-V4-Flash on $HOST:$PORT (TP=$TP_SIZE, memory-ratio $MEMORY_RATIO${spec:+, dSpark drafter})" : > "$LOG" @@ -131,6 +144,7 @@ cmd_start() { --max-prefill-length "$MAX_PREFILL_LENGTH" \ --expert-load "$EXPERT_LOAD" \ "${spec[@]}" \ + "${fallback[@]}" \ >> "$LOG" 2>&1 < /dev/null & echo $! > "$PIDFILE" diff --git a/tests/dsv4/test_dsv4_adaptive_verification.py b/tests/dsv4/test_dsv4_adaptive_verification.py index c168a7f..55e61ef 100644 --- a/tests/dsv4/test_dsv4_adaptive_verification.py +++ b/tests/dsv4/test_dsv4_adaptive_verification.py @@ -7,7 +7,64 @@ from freetoken.engine.engine import Engine from freetoken.engine.graph import GraphRunner, SharedSpecCarryJournal -from freetoken.models.deepseek_v4.dspark import choose_adaptive_draft_width +from freetoken.models.deepseek_v4.dspark import ( + DSparkAcceptanceFallback, + choose_adaptive_draft_width, +) +from freetoken.scheduler.scheduler import Scheduler + + +def test_acceptance_fallback_trips_only_after_minimum_sample(): + fallback = DSparkAcceptanceFallback(0.60, min_drafted=8, cooldown_steps=3) + + assert fallback.record(2, 5) is None + trip = fallback.record(1, 3) + + assert trip == pytest.approx((3 / 8, 8)) + + +def test_acceptance_fallback_skips_exact_cooldown_then_probes(): + fallback = DSparkAcceptanceFallback(0.60, min_drafted=4, cooldown_steps=3) + assert fallback.record(1, 4) is not None + + assert [fallback.take_decision() for _ in range(3)] == [ + (False, False), + (False, False), + (False, False), + ] + assert fallback.take_decision() == (True, True) + assert fallback.take_decision() == (True, False) + + +def test_acceptance_fallback_keeps_high_acceptance_speculating(): + fallback = DSparkAcceptanceFallback(0.60, min_drafted=8, cooldown_steps=3) + + assert fallback.record(5, 8) is None + assert fallback.take_decision() == (True, False) + + +def test_acceptance_fallback_reset_clears_request_state(): + fallback = DSparkAcceptanceFallback(0.60, min_drafted=4, cooldown_steps=3) + fallback.record(1, 4) + + fallback.reset() + + assert fallback.take_decision() == (True, False) + assert fallback.record(1, 3) is None + + +def test_scheduler_fallback_bypasses_block_before_appending_noise(): + scheduler = Scheduler.__new__(Scheduler) + scheduler._speculative = SimpleNamespace(block_size=5, noise_token_id=17) + scheduler.engine = SimpleNamespace(should_speculate=lambda _req: False) + scheduler._warned_non_greedy = False + req = SimpleNamespace(remain_len=100, input_ids=torch.tensor([1, 2])) + batch = SimpleNamespace(reqs=[req], is_decode=True, phase="decode") + + scheduler._maybe_make_speculative(batch) + + assert batch.phase == "decode" + assert req.input_ids.tolist() == [1, 2] def test_profile_cliff_stops_before_expensive_graph_bucket():