diff --git a/docs/models/qwen3/dflash-speculative-decoding.md b/docs/models/qwen3/dflash-speculative-decoding.md index fd67f33c7..5b83351b3 100644 --- a/docs/models/qwen3/dflash-speculative-decoding.md +++ b/docs/models/qwen3/dflash-speculative-decoding.md @@ -86,10 +86,12 @@ The gap is **launch exposure**. Under greedy dflash the spec path runs with *no* **Why the whole forward can't go in one graph.** A first attempt captured the entire verify forward; output was correct up to ~token 60, then garbage. Root cause: FlashInfer's paged-prefill attention derives its KV-iteration count (`num_iterations`, from `kv_len`) and that loop bound is **frozen when the graph is recorded**. The verify context grows every step, so once it crosses the captured `CTA_TILE_KV` (~64) boundary the replayed attention under-reads KV. Base-decode's graph is safe only because its *decode* kernel's KV loop is purely device-driven; *prefill*'s is not. FlashInfer ships no graph-safe prefill variant; vLLM hits the same wall and keeps attention out of its piecewise cudagraph. -**The fix: piecewise graph.** Keep attention **eager**; capture only the dense ops, whose dims depend on the fixed `span` row count, never on KV length. `forward_layer_batch_paged` is split into `pre_attn` / `attn` / `post_attn`, and the verify forward becomes `num_layers+1` dense graph segments — `[embed+L0.pre] [L0.attn eager] [L0.post+L1.pre] … [L_last.post+lm_head]` — captured once per batch bucket and replayed (`verify_graph.rs`). The ping-pong residual swap sits inside the captured segments, so its pointer alternation is baked into each graph and reproduces on every replay regardless of layer parity: `run_or_capture` re-runs the CPU swap only on the capture step, and the one eager op (attention) touches just `q/k/v_batch` / `attn_output`, never the swapped `hidden`. +**The fix: piecewise graph.** Keep attention **eager**; capture only the dense ops, whose dims depend on the verify row count (`total_tokens`), never on KV length. `forward_layer_batch_paged` is split into `pre_attn` / `attn` / `post_attn`, and the verify forward becomes `num_layers+1` dense graph segments — `[embed+L0.pre] [L0.attn eager] [L0.post+L1.pre] … [L_last.post+lm_head]` — captured once per batch bucket and replayed (`verify_graph.rs`). The ping-pong residual swap sits inside the captured segments, so its pointer alternation is baked into each graph and reproduces on every replay regardless of layer parity: `run_or_capture` re-runs the CPU swap only on the capture step, and the one eager op (attention) touches just `q/k/v_batch` / `attn_output`, never the swapped `hidden`. Result (5090, greedy, same-session A/B): fixed-buffer eager **250.9 → +graph 274.3 tok/s (+9.3% from the graph alone)**; 237 → 274 (+16%) over the pre-graph batched baseline, **matching vLLM's 278**. Concurrent (no-regression check): c8 1346 → 1525, c16 1868 → 1834 (both still ≥ vLLM). Losslessness gate passes (bf16 tie-flips only) — the dense ops replay bit-identically, and the eager attention is unchanged. +**Capture-shape gotcha (fixed).** `total_tokens` is **not** constant at a given batch bucket: a request near its output budget shortens its verify span (`plan.rs` truncates the span to `max_tokens − generated`), so `total_tokens < batch_size × span`. The captured dense kernels bake their row count into the launch, and `run_or_capture` is capture-once-replay-forever — so a graph first captured at a *short* span and later replayed at a *full* span processes too few rows and leaves the tail-request logits **stale**: a silent losslessness break. It hid from the bs=1 gate (a fresh request's first verify is always a full span ⇒ bucket captured at the max ⇒ only the harmless over-compute direction occurs) **and from the homogeneous c8/c16 benches** (lockstep requests capture every bucket at full span during ramp-up; all truncation comes later as they finish together — still the safe direction). The dangerous direction needs *heterogeneous* progress. Fix: gate the captured-graph path on `total_tokens == batch_size × span` (every request a full span); any truncated step runs eager — making capture-shape ≡ replay-shape an invariant by construction. Regression test `dflash_short_then_long_verify_capture_is_lossless`: a `max_tokens=4` request poisons the bucket-bs=1 graph at a truncated span, then a long request replays it — RED before the gate (diverges to a stale-buffer token), GREEN after. Cost: a batch containing *any* truncated request runs fully eager that step (rare — only a request's final block); a follow-up could pad-to-full + mask to keep the graph. + **Draft-side piecewise graph is the tracked next step.** The draft (5 layers, `dflash.rs`) is the other ~16% of the launch gap; it needs the same pre/attn/post split, with its variable-length contiguous KV (`DFlashLayerCache`) handled at the eager attention boundary. Tracked as its own PR after this one lands. The EAGLE proposer trait is still deferred to when EAGLE actually lands (see "no proposer trait yet" above); both the batched draft and the future CUDA-Graph draft are DFlash-internal changes behind the unchanged `DraftPlan→DraftResult` seam, so neither is thrown away by EAGLE. @@ -136,8 +138,10 @@ The fixed reservation lands exactly in the margin (+2822 MiB) and the per-token ### Review blockers (correctness/usability, independent of the perf work) -Issues surfaced in PR review, largely independent of the draft batching (#2's DFlash wrappers were fixed alongside it): +Issues surfaced in PR review, largely independent of the draft batching. + +1. **Unified path silently skipped DFlash readiness — fixed.** `StepCommand::Unified` (`executor.rs`) captures no DFlash hidden state, and only `execute_prefill` marks a request draft-ready. A greedy request prefilled via a fused Unified step would therefore never become draft-ready and never recover — DFlash silently no-opped for it under mixed load (no wrong tokens, the feature just quietly disabled itself). Fixed by routing capture-eligible pending (greedy, no LoRA, no logprobs) to a **dedicated prefill step** instead of Unified — `build_next_plan`'s `needs_dflash_capture` (`scheduler/plan.rs`), mirroring the existing `needs_prompt_logprobs` precedent — so prefill capture always runs. The Unified decode arm also now drops stale draft context for each decoded request (`execute_unified`, mirroring `execute_decode`), keeping the "readiness comes from prefill capture" invariant closed instead of degrading silently. +2. **Stream-override race — fixed.** The DFlash `qk_norm_rope` / `single_prefill` wrappers (`attention.rs`) already used `active_cu_stream(ctx)`; `copy_hidden_rows_into` and `copy_hidden_token_range_into` (`elementwise.rs`) now launch on `active_cu_stream(ctx)` too, matching the repo convention (`tensor.rs:43`) so a captured copy records on the right stream. Belt-and-braces: DFlash + decode overlap is now **rejected at launch** (`lib.rs`) — the speculative path never takes the unified overlap route, so the combination only burned VRAM the drafter needs. +3. **`gemm_lt_pin_tune` is not a real warmup — still open.** It only pins the heuristic (`linear.cu:497`); it never executes a `cublasLtMatmul` the way the old `gemm_lt_tune_cuda` (`linear.cu:431`) did, so the first real matmul can land inside CUDA-graph capture. `batch_invariance_decode_gemm_graph` backstops it, but the warmup should actually run the matmul. This is the remaining review blocker. -1. **Unified path silently skips DFlash readiness.** `StepCommand::Unified` (`executor.rs`) captures no DFlash hidden state, and only the plain `execute_prefill` post-step marks a request draft-ready (`executor.rs` ~1766). So when active + pending fuse into a Unified step (`scheduler/plan.rs` — the normal mixed-load path), greedy requests routed through Unified never become draft-ready and never recover: DFlash silently no-ops for them. No wrong tokens, but the feature quietly disables itself under mixed load. Crash-early or capture-in-Unified, don't degrade silently. -2. **Stream-override race (partly fixed).** The DFlash `qk_norm_rope` / `single_prefill` wrappers (`attention.rs`) now use `active_cu_stream(ctx)` (fixed in the batched-draft PR). `copy_hidden_rows_into` (`elementwise.rs:209`) still uses `ctx.stream.cu_stream()` instead of the repo convention (`tensor.rs:43`) — under Green-Context / split-stream decode overlap this remains a planted race. -3. **`gemm_lt_pin_tune` is not a real warmup.** It only pins the heuristic (`linear.cu:497`); it never executes a `cublasLtMatmul` the way the old `gemm_lt_tune_cuda` (`linear.cu:431`) did, so the first real matmul can land inside CUDA-graph capture. `batch_invariance_decode_gemm_graph` backstops it, but the warmup should actually run the matmul. +Concurrent, heterogeneous-`max_tokens` losslessness is now covered by `dflash_concurrent_heterogeneous_is_lossless` — several greedy requests at staggered budgets run as one batch, each checked against its own plain-greedy baseline. That exercises the bs>1 draft+verify path the bs=1 gate and the homogeneous c8/c16 benches never reach (a batched-draft indexing or capture-shape regression at bs>1 would surface as a real, non-tie divergence). diff --git a/openinfer-kernels/src/ops/elementwise.rs b/openinfer-kernels/src/ops/elementwise.rs index ab37b8bc2..2f6a0cb06 100644 --- a/openinfer-kernels/src/ops/elementwise.rs +++ b/openinfer-kernels/src/ops/elementwise.rs @@ -206,7 +206,7 @@ pub fn copy_hidden_rows_into( row_offset as i32, src.hidden_dim as i32, src.seq_len as i32, - ctx.stream.cu_stream(), + crate::tensor::active_cu_stream(ctx), ) }; result.result()?; @@ -253,7 +253,7 @@ pub fn copy_hidden_token_range_into( token_count as i32, src.seq_len as i32, dst.seq_len as i32, - ctx.stream.cu_stream(), + crate::tensor::active_cu_stream(ctx), ) }; result.result()?; diff --git a/openinfer-qwen3-4b/src/dflash.rs b/openinfer-qwen3-4b/src/dflash.rs index 96b80d790..3afad0339 100644 --- a/openinfer-qwen3-4b/src/dflash.rs +++ b/openinfer-qwen3-4b/src/dflash.rs @@ -1,16 +1,16 @@ use anyhow::{Context, Result}; use cudarc::driver::CudaSlice; -use log::debug; use crate::config::DFlashConfig; -use crate::weights::{Attention, MLP, Qwen3Model, TransformerBlock}; +use crate::weights::{Qwen3Model, TransformerBlock}; use openinfer_core::ops; use openinfer_core::tensor::HiddenStates; use openinfer_core::tensor::{DeviceContext, DeviceMatrix, DeviceVec}; -use openinfer_core::weight_loader::{ - deserialize_shards, load_shard_info, load_tensor_1d, load_tensor_2d, mmap_shards, - precompute_rope, -}; + +mod loading; +mod reservation; + +pub(crate) use reservation::DFlashMemoryReservation; pub(crate) struct DFlashDraftModel { config: DFlashConfig, @@ -33,87 +33,6 @@ pub(crate) struct DFlashRequestState { max_cache_len: usize, } -/// GPU memory DFlash needs on top of the target KV pool, derived from the draft -/// config so the KV budget can reserve it *before* the draft model loads (the -/// draft buffers live outside the paged `KvCacheManager`). Split by how it scales: -/// -/// - `kv_bytes_per_token` scales with the KV pool (billed by shrinking the target -/// block count): the draft's own KV cache plus the per-request context-projection -/// and pending-context buffers, which currently persist at prompt length per -/// request (see `dflash-speculative-decoding.md` — collapsing that persistence -/// is a tracked follow-up that would shrink this term to the draft KV alone). -/// - `fixed_bytes` does not scale with the pool (billed via the memory margin): -/// the draft weights plus the lane-level batched scratch sized for the whole -/// decode batch. -/// -// TODO: the draft scratch is now a single lane-level `DFlashBatchScratch` -// allocated once (dense buffers sized `max_batch * block_size`, plus one shared -// varlen tail), not a per-request buffer. The per-token `tail_scratch` term and -// the per-request `block_headroom` tail term are therefore over-estimates — kept -// as a conservative upper bound until the accounting is retuned against the -// batched allocation. -pub(crate) struct DFlashMemoryReservation { - pub(crate) kv_bytes_per_token: usize, - pub(crate) fixed_bytes: usize, -} - -impl DFlashMemoryReservation { - pub(crate) fn from_path(draft_path: &str, max_decode_batch_size: usize) -> Result { - let config = DFlashConfig::from_file(draft_path)?; - Ok(Self::from_config(&config, max_decode_batch_size)) - } - - fn from_config(config: &DFlashConfig, max_decode_batch_size: usize) -> Self { - const BF16: usize = 2; - let hidden = config.hidden_size; - let kv_dim = config.num_key_value_heads * config.head_dim; - let q_dim = config.num_attention_heads * config.head_dim; - let inter = config.intermediate_size; - let capture_layers = config.dflash_config.target_layer_ids.len(); - - // Per-sequence-token, pool-scaling buffers. - let draft_kv = config.num_hidden_layers * 2 * kv_dim * BF16; // DFlashLayerCache k+v - // Scratch split by what it tracks: `context_*` grows with the committed - // prefix; `tail_*` (tail_input + k_tail + v_tail) grows with the in-fill - // tail, which is one block past the prefix. - let context_scratch = 2 * hidden * BF16; // context_projected + context_hidden - let tail_scratch = (hidden + 2 * kv_dim) * BF16; // tail_input + k_tail + v_tail - let pending = hidden * capture_layers * BF16; // context_feature_dim - let kv_bytes_per_token = draft_kv + context_scratch + tail_scratch + pending; - - // Lane-level batched dense scratch: every dense buffer is sized for the - // whole decode batch (`max_batch * block_size` rows), allocated once. - // Same total magnitude as the old per-request scratch summed over the - // batch, but now one contiguous allocation. - let dense_scratch_per_block_row = - BF16 * (config.vocab_size + 5 * hidden + 2 * q_dim + 3 * inter); - let scratch_total = dense_scratch_per_block_row * config.block_size * max_decode_batch_size; - - // Draft weights (5 transformer layers + the context projection), +10% slack - // for norms, rope caches, and allocator alignment. - let per_layer = BF16 - * (hidden * (q_dim + 2 * kv_dim) // qkv_proj - + q_dim * hidden // o_proj - + hidden * 2 * inter // gate_up_proj - + inter * hidden); // down_proj - let fc = BF16 * hidden * (hidden * capture_layers); // context projection - let weights = per_layer * config.num_hidden_layers + fc; - let weights = weights + weights / 10; - - // The durable draft KV and the tail scratch are sized to `context + - // block_size` — one in-fill block past the lifetime the KV pool reserves - // for the request. The per-token term bills only the pool's tokens, so - // reserve that one-block headroom per concurrently decoding request to - // keep the reservation an upper bound. - let block_headroom = max_decode_batch_size * config.block_size * (draft_kv + tail_scratch); - - Self { - kv_bytes_per_token, - fixed_bytes: weights + scratch_total + block_headroom, - } - } -} - struct DFlashLayerCache { k: HiddenStates, v: HiddenStates, @@ -375,138 +294,6 @@ impl DFlashBatchScratch { } impl DFlashDraftModel { - pub(crate) fn from_safetensors_for_target( - ctx: &DeviceContext, - model_path: &str, - target: &Qwen3Model, - ) -> Result { - let config = DFlashConfig::from_file(model_path) - .with_context(|| format!("load DFlash config from {model_path}"))?; - config.validate_for_target(target.config())?; - - let (shard_paths, weight_map) = load_shard_info(model_path)?; - debug!( - "Loading DFlash drafter from {model_path}: {} shard(s)", - shard_paths.len() - ); - let mmaps = mmap_shards(&shard_paths)?; - let shards = deserialize_shards(&mmaps)?; - - let mut layers = Vec::with_capacity(config.num_hidden_layers); - for layer_idx in 0..config.num_hidden_layers { - let prefix = format!("layers.{layer_idx}"); - - let q_proj = load_tensor_2d( - ctx, - &shards, - &weight_map, - &format!("{prefix}.self_attn.q_proj.weight"), - )?; - let k_proj = load_tensor_2d( - ctx, - &shards, - &weight_map, - &format!("{prefix}.self_attn.k_proj.weight"), - )?; - let v_proj = load_tensor_2d( - ctx, - &shards, - &weight_map, - &format!("{prefix}.self_attn.v_proj.weight"), - )?; - let q_dim = q_proj.rows; - let kv_dim = k_proj.rows; - let qkv_proj = DeviceMatrix::vstack(ctx, &[&q_proj, &k_proj, &v_proj])?; - drop(q_proj); - drop(k_proj); - drop(v_proj); - - let gate_proj = load_tensor_2d( - ctx, - &shards, - &weight_map, - &format!("{prefix}.mlp.gate_proj.weight"), - )?; - let up_proj = load_tensor_2d( - ctx, - &shards, - &weight_map, - &format!("{prefix}.mlp.up_proj.weight"), - )?; - let gate_up_proj = DeviceMatrix::vstack(ctx, &[&gate_proj, &up_proj])?; - drop(gate_proj); - drop(up_proj); - - layers.push(TransformerBlock { - input_layernorm: load_tensor_1d( - ctx, - &shards, - &weight_map, - &format!("{prefix}.input_layernorm.weight"), - )?, - attention: Attention { - qkv_proj, - o_proj: load_tensor_2d( - ctx, - &shards, - &weight_map, - &format!("{prefix}.self_attn.o_proj.weight"), - )?, - q_norm: load_tensor_1d( - ctx, - &shards, - &weight_map, - &format!("{prefix}.self_attn.q_norm.weight"), - )?, - k_norm: load_tensor_1d( - ctx, - &shards, - &weight_map, - &format!("{prefix}.self_attn.k_norm.weight"), - )?, - q_dim, - kv_dim, - }, - post_attention_layernorm: load_tensor_1d( - ctx, - &shards, - &weight_map, - &format!("{prefix}.post_attention_layernorm.weight"), - )?, - mlp: MLP { - gate_up_proj, - down_proj: load_tensor_2d( - ctx, - &shards, - &weight_map, - &format!("{prefix}.mlp.down_proj.weight"), - )?, - }, - }); - } - - let norm = load_tensor_1d(ctx, &shards, &weight_map, "norm.weight")?; - let hidden_norm = load_tensor_1d(ctx, &shards, &weight_map, "hidden_norm.weight")?; - let fc = load_tensor_2d(ctx, &shards, &weight_map, "fc.weight")?; - let (cos_cache, sin_cache) = precompute_rope( - ctx, - config.head_dim, - config.max_position_embeddings, - config.rope_theta, - )?; - ctx.sync()?; - - Ok(Self { - config, - layers, - norm, - hidden_norm, - fc, - cos_cache, - sin_cache, - }) - } - pub(crate) fn block_size(&self) -> usize { self.config.block_size } diff --git a/openinfer-qwen3-4b/src/dflash/loading.rs b/openinfer-qwen3-4b/src/dflash/loading.rs new file mode 100644 index 000000000..f0db630ac --- /dev/null +++ b/openinfer-qwen3-4b/src/dflash/loading.rs @@ -0,0 +1,146 @@ +use anyhow::{Context, Result}; +use log::debug; + +use crate::config::DFlashConfig; +use crate::weights::{Attention, MLP, Qwen3Model, TransformerBlock}; +use openinfer_core::tensor::{DeviceContext, DeviceMatrix}; +use openinfer_core::weight_loader::{ + deserialize_shards, load_shard_info, load_tensor_1d, load_tensor_2d, mmap_shards, + precompute_rope, +}; + +use super::DFlashDraftModel; + +impl DFlashDraftModel { + pub(crate) fn from_safetensors_for_target( + ctx: &DeviceContext, + model_path: &str, + target: &Qwen3Model, + ) -> Result { + let config = DFlashConfig::from_file(model_path) + .with_context(|| format!("load DFlash config from {model_path}"))?; + config.validate_for_target(target.config())?; + + let (shard_paths, weight_map) = load_shard_info(model_path)?; + debug!( + "Loading DFlash drafter from {model_path}: {} shard(s)", + shard_paths.len() + ); + let mmaps = mmap_shards(&shard_paths)?; + let shards = deserialize_shards(&mmaps)?; + + let mut layers = Vec::with_capacity(config.num_hidden_layers); + for layer_idx in 0..config.num_hidden_layers { + let prefix = format!("layers.{layer_idx}"); + + let q_proj = load_tensor_2d( + ctx, + &shards, + &weight_map, + &format!("{prefix}.self_attn.q_proj.weight"), + )?; + let k_proj = load_tensor_2d( + ctx, + &shards, + &weight_map, + &format!("{prefix}.self_attn.k_proj.weight"), + )?; + let v_proj = load_tensor_2d( + ctx, + &shards, + &weight_map, + &format!("{prefix}.self_attn.v_proj.weight"), + )?; + let q_dim = q_proj.rows; + let kv_dim = k_proj.rows; + let qkv_proj = DeviceMatrix::vstack(ctx, &[&q_proj, &k_proj, &v_proj])?; + drop(q_proj); + drop(k_proj); + drop(v_proj); + + let gate_proj = load_tensor_2d( + ctx, + &shards, + &weight_map, + &format!("{prefix}.mlp.gate_proj.weight"), + )?; + let up_proj = load_tensor_2d( + ctx, + &shards, + &weight_map, + &format!("{prefix}.mlp.up_proj.weight"), + )?; + let gate_up_proj = DeviceMatrix::vstack(ctx, &[&gate_proj, &up_proj])?; + drop(gate_proj); + drop(up_proj); + + layers.push(TransformerBlock { + input_layernorm: load_tensor_1d( + ctx, + &shards, + &weight_map, + &format!("{prefix}.input_layernorm.weight"), + )?, + attention: Attention { + qkv_proj, + o_proj: load_tensor_2d( + ctx, + &shards, + &weight_map, + &format!("{prefix}.self_attn.o_proj.weight"), + )?, + q_norm: load_tensor_1d( + ctx, + &shards, + &weight_map, + &format!("{prefix}.self_attn.q_norm.weight"), + )?, + k_norm: load_tensor_1d( + ctx, + &shards, + &weight_map, + &format!("{prefix}.self_attn.k_norm.weight"), + )?, + q_dim, + kv_dim, + }, + post_attention_layernorm: load_tensor_1d( + ctx, + &shards, + &weight_map, + &format!("{prefix}.post_attention_layernorm.weight"), + )?, + mlp: MLP { + gate_up_proj, + down_proj: load_tensor_2d( + ctx, + &shards, + &weight_map, + &format!("{prefix}.mlp.down_proj.weight"), + )?, + }, + }); + } + + let norm = load_tensor_1d(ctx, &shards, &weight_map, "norm.weight")?; + let hidden_norm = load_tensor_1d(ctx, &shards, &weight_map, "hidden_norm.weight")?; + let fc = load_tensor_2d(ctx, &shards, &weight_map, "fc.weight")?; + let (cos_cache, sin_cache) = precompute_rope( + ctx, + config.head_dim, + config.max_position_embeddings, + config.rope_theta, + )?; + ctx.sync()?; + + Ok(Self { + config, + layers, + norm, + hidden_norm, + fc, + cos_cache, + sin_cache, + }) + } +} diff --git a/openinfer-qwen3-4b/src/dflash/reservation.rs b/openinfer-qwen3-4b/src/dflash/reservation.rs new file mode 100644 index 000000000..3dfe549d1 --- /dev/null +++ b/openinfer-qwen3-4b/src/dflash/reservation.rs @@ -0,0 +1,84 @@ +use anyhow::Result; + +use crate::config::DFlashConfig; + +/// GPU memory DFlash needs on top of the target KV pool, derived from the draft +/// config so the KV budget can reserve it *before* the draft model loads (the +/// draft buffers live outside the paged `KvCacheManager`). Split by how it scales: +/// +/// - `kv_bytes_per_token` scales with the KV pool (billed by shrinking the target +/// block count): the draft's own KV cache plus the per-request context-projection +/// and pending-context buffers, which currently persist at prompt length per +/// request (see `dflash-speculative-decoding.md` — collapsing that persistence +/// is a tracked follow-up that would shrink this term to the draft KV alone). +/// - `fixed_bytes` does not scale with the pool (billed via the memory margin): +/// the draft weights plus the lane-level batched scratch sized for the whole +/// decode batch. +/// +// TODO: the draft scratch is now a single lane-level `DFlashBatchScratch` +// allocated once (dense buffers sized `max_batch * block_size`, plus one shared +// varlen tail), not a per-request buffer. The per-token `tail_scratch` term and +// the per-request `block_headroom` tail term are therefore over-estimates — kept +// as a conservative upper bound until the accounting is retuned against the +// batched allocation. +pub(crate) struct DFlashMemoryReservation { + pub(crate) kv_bytes_per_token: usize, + pub(crate) fixed_bytes: usize, +} + +impl DFlashMemoryReservation { + pub(crate) fn from_path(draft_path: &str, max_decode_batch_size: usize) -> Result { + let config = DFlashConfig::from_file(draft_path)?; + Ok(Self::from_config(&config, max_decode_batch_size)) + } + + pub(crate) fn from_config(config: &DFlashConfig, max_decode_batch_size: usize) -> Self { + const BF16: usize = 2; + let hidden = config.hidden_size; + let kv_dim = config.num_key_value_heads * config.head_dim; + let q_dim = config.num_attention_heads * config.head_dim; + let inter = config.intermediate_size; + let capture_layers = config.dflash_config.target_layer_ids.len(); + + // Per-sequence-token, pool-scaling buffers. + let draft_kv = config.num_hidden_layers * 2 * kv_dim * BF16; // DFlashLayerCache k+v + // Scratch split by what it tracks: `context_*` grows with the committed + // prefix; `tail_*` (tail_input + k_tail + v_tail) grows with the in-fill + // tail, which is one block past the prefix. + let context_scratch = 2 * hidden * BF16; // context_projected + context_hidden + let tail_scratch = (hidden + 2 * kv_dim) * BF16; // tail_input + k_tail + v_tail + let pending = hidden * capture_layers * BF16; // context_feature_dim + let kv_bytes_per_token = draft_kv + context_scratch + tail_scratch + pending; + + // Lane-level batched dense scratch: every dense buffer is sized for the + // whole decode batch (`max_batch * block_size` rows), allocated once. + // Same total magnitude as the old per-request scratch summed over the + // batch, but now one contiguous allocation. + let dense_scratch_per_block_row = + BF16 * (config.vocab_size + 5 * hidden + 2 * q_dim + 3 * inter); + let scratch_total = dense_scratch_per_block_row * config.block_size * max_decode_batch_size; + + // Draft weights (5 transformer layers + the context projection), +10% slack + // for norms, rope caches, and allocator alignment. + let per_layer = BF16 + * (hidden * (q_dim + 2 * kv_dim) // qkv_proj + + q_dim * hidden // o_proj + + hidden * 2 * inter // gate_up_proj + + inter * hidden); // down_proj + let fc = BF16 * hidden * (hidden * capture_layers); // context projection + let weights = per_layer * config.num_hidden_layers + fc; + let weights = weights + weights / 10; + + // The durable draft KV and the tail scratch are sized to `context + + // block_size` — one in-fill block past the lifetime the KV pool reserves + // for the request. The per-token term bills only the pool's tokens, so + // reserve that one-block headroom per concurrently decoding request to + // keep the reservation an upper bound. + let block_headroom = max_decode_batch_size * config.block_size * (draft_kv + tail_scratch); + + Self { + kv_bytes_per_token, + fixed_bytes: weights + scratch_total + block_headroom, + } + } +} diff --git a/openinfer-qwen3-4b/src/executor.rs b/openinfer-qwen3-4b/src/executor.rs index cf8ed1262..c420a457b 100644 --- a/openinfer-qwen3-4b/src/executor.rs +++ b/openinfer-qwen3-4b/src/executor.rs @@ -1912,6 +1912,18 @@ impl ModelExecutor for Qwen3Executor { .expect("request must exist after unified decode"); rkv.apply_decode(req_result.token, self.kv_mgr.pool())?; } + // A plain decode via the fused unified step advances the sequence + // outside the speculative path, so any captured draft context is + // now stale — drop it, mirroring execute_decode. (Eligible pending + // are routed to a dedicated prefill step, so unified prefills never + // need DFlash mark-ready here.) + if self.speculative.is_some() { + for req_result in &result.decode_requests { + if self.dflash_ready_requests.remove(&req_result.request_id) { + self.primary.drop_dflash_request(req_result.request_id)?; + } + } + } for req_result in &result.prefill_requests { self.save_sealed_blocks(req_result.request_id); } diff --git a/openinfer-qwen3-4b/src/lib.rs b/openinfer-qwen3-4b/src/lib.rs index a6bb433f5..3d27ea6b3 100644 --- a/openinfer-qwen3-4b/src/lib.rs +++ b/openinfer-qwen3-4b/src/lib.rs @@ -214,6 +214,13 @@ pub fn launch(model_path: &Path, options: Qwen3LaunchOptions) -> Result { info!( diff --git a/openinfer-qwen3-4b/src/scheduler.rs b/openinfer-qwen3-4b/src/scheduler.rs index 71c918771..0bb500522 100644 --- a/openinfer-qwen3-4b/src/scheduler.rs +++ b/openinfer-qwen3-4b/src/scheduler.rs @@ -1084,7 +1084,7 @@ fn runtime_plan( Some(ExecutionPlan::Prefill { pending }) } } else { - build_next_plan(!active.is_empty(), pending) + build_next_plan(!active.is_empty(), pending, executor.speculative_enabled()) } } diff --git a/openinfer-qwen3-4b/src/scheduler/plan.rs b/openinfer-qwen3-4b/src/scheduler/plan.rs index 5f9859195..19b1bf7ec 100644 --- a/openinfer-qwen3-4b/src/scheduler/plan.rs +++ b/openinfer-qwen3-4b/src/scheduler/plan.rs @@ -48,12 +48,26 @@ pub(super) enum ExecutionArtifacts { pub(super) fn build_next_plan( have_active: bool, pending: Vec, + speculative: bool, ) -> Option { - // echo+logprobs requests need all-position logits, which the unified - // forward does not compute (it passes all_position_logits=None). Route - // them through a dedicated prefill step instead of degrading silently. + // echo+logprobs requests need all-position logits, which the unified forward + // does not compute (it passes all_position_logits=None). And under DFlash + // speculation, an eligible request must capture its target hidden context + // during prefill — the unified forward skips that capture, so a request + // prefilled via Unified would never become draft-ready and DFlash would + // silently no-op for it forever. Either way, route pending through a + // dedicated prefill step instead of degrading silently. let needs_prompt_logprobs = pending.iter().any(|r| r.echo && r.logprobs > 0); - if !pending.is_empty() && have_active && !needs_prompt_logprobs { + // Deliberately a loose superset of the real capture eligibility + // (`dflash_prefill_supported`, which also needs `cached_tokens == 0 && !echo`): + // over-routing an ineligible request to a dedicated prefill only costs one + // fusion, but under-routing a capture-eligible one into Unified would silently + // break its readiness. Never tighten this into the dangerous direction. + let needs_dflash_capture = speculative + && pending + .iter() + .any(|r| r.lora_adapter.is_none() && r.logprobs == 0 && r.params.is_greedy()); + if !pending.is_empty() && have_active && !needs_prompt_logprobs && !needs_dflash_capture { Some(ExecutionPlan::Unified { pending }) } else if !pending.is_empty() { Some(ExecutionPlan::Prefill { pending }) @@ -174,7 +188,11 @@ fn build_speculative_verify_items( // Clamp the verify span to the request's remaining output budget so // a long accepted run can't overshoot max_tokens. let remaining = active.max_tokens.saturating_sub(active.generated_count); - assert!(remaining > 0, "active request must have output budget"); + // A continuing active request always has budget left (resolve emits + // EmitManyAndFinish the moment generated_count hits max_tokens), so + // this is a true invariant, not a runtime condition — don't crash the + // scheduler thread in release on a state we've proven unreachable. + debug_assert!(remaining > 0, "active request must have output budget"); let mut token_ids = draft.token_ids.clone(); token_ids.truncate(remaining); VerifyStepItem::new(draft.request_id, token_ids, active.params) @@ -290,23 +308,26 @@ mod tests { #[test] fn plan_selection_follows_active_and_pending_state() { assert!( - build_next_plan(false, vec![]).is_none(), + build_next_plan(false, vec![], false).is_none(), "idle scheduler (no active, no pending) produces no plan" ); assert!( - matches!(build_next_plan(true, vec![]), Some(ExecutionPlan::Decode)), + matches!( + build_next_plan(true, vec![], false), + Some(ExecutionPlan::Decode) + ), "active-only ticks decode the running batch" ); assert!( matches!( - build_next_plan(false, vec![pending()]), + build_next_plan(false, vec![pending()], false), Some(ExecutionPlan::Prefill { pending }) if pending.len() == 1 ), "pending-only prefills the new arrivals" ); assert!( matches!( - build_next_plan(true, vec![pending()]), + build_next_plan(true, vec![pending()], false), Some(ExecutionPlan::Unified { pending }) if pending.len() == 1 ), "active + pending fuses prefill and decode into one unified step" @@ -318,7 +339,7 @@ mod tests { echo_req.logprobs = 5; assert!( matches!( - build_next_plan(true, vec![echo_req]), + build_next_plan(true, vec![echo_req], false), Some(ExecutionPlan::Prefill { pending }) if pending.len() == 1 ), "active + pending echo+logprobs request routes to prefill not unified" @@ -328,10 +349,35 @@ mod tests { echo_no_lp.echo = true; assert!( matches!( - build_next_plan(true, vec![echo_no_lp]), + build_next_plan(true, vec![echo_no_lp], false), Some(ExecutionPlan::Unified { pending }) if pending.len() == 1 ), "active + pending echo-only request (no logprobs) can use unified" ); + // Under DFlash speculation, an eligible (greedy) pending must capture its + // target context during prefill — the unified forward skips that capture, + // so route it to a dedicated prefill step rather than let DFlash silently + // no-op for it. + assert!( + matches!( + build_next_plan(true, vec![pending()], true), + Some(ExecutionPlan::Prefill { pending }) if pending.len() == 1 + ), + "spec + active + eligible pending routes to prefill so the drafter context is captured" + ); + // A non-greedy pending needs no draft capture, so unified fusion is still + // fine even under speculation. + let mut sampled = pending(); + sampled.params = SamplingParams { + temperature: 1.0, + ..SamplingParams::default() + }; + assert!( + matches!( + build_next_plan(true, vec![sampled], true), + Some(ExecutionPlan::Unified { pending }) if pending.len() == 1 + ), + "spec + active + non-greedy pending (no capture needed) can still use unified" + ); } } diff --git a/openinfer-qwen3-4b/src/verify_graph.rs b/openinfer-qwen3-4b/src/verify_graph.rs index af6d4b391..25c4e626f 100644 --- a/openinfer-qwen3-4b/src/verify_graph.rs +++ b/openinfer-qwen3-4b/src/verify_graph.rs @@ -16,8 +16,10 @@ //! between segments. Attention must stay eager because FlashInfer's paged-prefill //! kernel fixes its KV-iteration count when the graph is recorded; with the verify //! context growing every step, a captured attention would under-read KV and -//! corrupt later tokens. The dense segments are shape-stable in the fixed -//! `span`-row layout, so one capture replays losslessly for the request's life. +//! corrupt later tokens. The dense segments bake their row count, so a captured +//! segment is only ever replayed at the exact full `batch_size * span` shape it +//! was recorded at; a step whose span is truncated near a request's output +//! budget falls back to eager (see [`Qwen3Model::batch_prefill_into`]). use anyhow::Result; use cudarc::driver::CudaSlice; @@ -258,8 +260,19 @@ impl Qwen3Model { // the fixed `span`-row layout. Segments: [embed + L0.pre] [L0.attn] // [L0.post + L1.pre] [L1.attn] ... [L_last.post + lm_head]. --- let num_layers = self.layers.len(); + // Each captured dense segment bakes its row count (`total_tokens`) into + // every kernel launch. A request near its output budget shortens its span + // (the scheduler truncates the verify span to the remaining budget), so + // `total_tokens` varies at a fixed `batch_size`. Replaying a segment that + // was captured at one row count at a *different* count processes the wrong + // number of rows and leaves the tail rows stale — silently corrupting the + // verify logits of the trailing requests. So only the full, maximal + // `batch_size * span` shape (every request contributing a full span) uses + // the graph; any truncated step runs eager. This makes + // capture-shape == replay-shape an invariant by construction. + let full_shape = total_tokens == batch_size * bufs.span; match BATCH_BUCKETS.iter().position(|&b| b == batch_size) { - Some(bidx) => { + Some(bidx) if full_shape => { // Take the bucket's segment graphs out so the capture closures can // borrow `bufs` mutably; restore them after (even on error). let mut segs = std::mem::take(&mut bufs.graphs[bidx]); @@ -280,8 +293,9 @@ impl Qwen3Model { bufs.graphs[bidx] = segs; result?; } - None => { - // Off-bucket batch: run the same segments eager (no capture). + // Off-bucket batch, or a truncated (non-full-span) step: run the same + // segments eager (no capture). + _ => { self.verify_seg_embed_pre(bufs)?; self.verify_attn(0, kv_buffer, layout, bufs)?; for i in 1..num_layers { diff --git a/openinfer-qwen3-4b/tests/dflash_speculative_gate.rs b/openinfer-qwen3-4b/tests/dflash_speculative_gate.rs index 2f2127b50..2cf83760e 100644 --- a/openinfer-qwen3-4b/tests/dflash_speculative_gate.rs +++ b/openinfer-qwen3-4b/tests/dflash_speculative_gate.rs @@ -48,6 +48,7 @@ use openinfer_qwen3_4b::{ DEFAULT_KV_CACHE_MEMORY_MARGIN_BYTES, DEFAULT_MAX_PREFILL_TOKENS, DecodeOverlap, Qwen3LaunchOptions, Qwen3MemoryOptions, Qwen3OffloadOptions, }; +use vllm_text::tokenizer::DynTokenizer; mod common; @@ -126,7 +127,12 @@ struct Step { } /// Submit one greedy request and collect the decoded steps until `Finished`. -fn generate(handle: &EngineHandle, prompt_tokens: Vec, logprobs: usize) -> Vec { +fn generate( + handle: &EngineHandle, + prompt_tokens: Vec, + logprobs: usize, + max_tokens: usize, +) -> Vec { let (token_tx, mut rx) = TokenSink::standalone(); handle .submit(GenerateRequest { @@ -134,7 +140,7 @@ fn generate(handle: &EngineHandle, prompt_tokens: Vec, logprobs: usize) -> queued_at_unix_s: None, prompt_tokens, params: SamplingParams::default(), - max_tokens: GENERATED_TOKENS, + max_tokens, lora_adapter: None, token_tx, logprobs, @@ -158,6 +164,63 @@ fn generate(handle: &EngineHandle, prompt_tokens: Vec, logprobs: usize) -> } } +/// Submit several greedy requests at once, then collect each one's steps. They +/// run concurrently in the one engine — the scheduler batches them — so with +/// heterogeneous `max_tokens` the verify spans differ across a batch, exercising +/// the real bs>1 draft+verify path. Each tuple is `(prompt_tokens, max_tokens)`; +/// logprobs are off so the speculative path stays active. Returns one step list +/// per request, in submission order. +fn generate_concurrent(handle: &EngineHandle, requests: Vec<(Vec, usize)>) -> Vec> { + // Submit all up front so they coexist in the engine and form real batches. + let receivers: Vec<_> = requests + .into_iter() + .map(|(prompt_tokens, max_tokens)| { + let (token_tx, rx) = TokenSink::standalone(); + handle + .submit(GenerateRequest { + request_id: None, + queued_at_unix_s: None, + prompt_tokens, + params: SamplingParams::default(), + max_tokens, + lora_adapter: None, + token_tx, + logprobs: 0, + echo: false, + }) + .expect("submit failed"); + rx + }) + .collect(); + + // Drain each request's channel to completion (events are buffered per-channel, + // so the drain order doesn't matter — they all ran concurrently). + receivers + .into_iter() + .map(|mut rx| { + let mut steps = Vec::new(); + loop { + match rx.blocking_recv().map(|(_, event)| event) { + Some(TokenEvent::Token { id, logprob }) => steps.push(Step { + id, + top_logprobs: logprob.map(|lp| lp.top_logprobs).unwrap_or_default(), + }), + Some(TokenEvent::Scheduled { .. } | TokenEvent::PromptTokens { .. }) => {} + Some(TokenEvent::Finished { .. }) => break, + Some(TokenEvent::Error { message, .. }) => { + panic!("generation failed: {message}") + } + Some(TokenEvent::Rejected { message, .. }) => { + panic!("generation rejected: {message}") + } + None => panic!("scheduler channel closed without Finished"), + } + } + steps + }) + .collect() +} + /// Prefill `context` (echo) and return the next-token distribution the *prefill* /// kernel produces — the kernel the speculative verify path also uses. This is /// the reference the spec pick should match (vs the plain-decode baseline, whose @@ -196,6 +259,133 @@ fn prefill_next(handle: &EngineHandle, context: Vec, logprobs: usize) -> St } } +/// Compare one prompt's speculative `spec` steps against its plain-greedy `base`, +/// tolerating only the benign prefill-vs-decode kernel-gap tie flip (the spec +/// pick sits within `MARGIN_TOL` of the prefill kernel's own argmax, measured in +/// the prefill distribution the verify path actually runs). `Ok(())` ⇒ lossless +/// or a benign tie; `Err(diagnostic)` ⇒ a real spec bug. `handle` must be the +/// live speculative engine — at a divergence it re-prefills the shared context +/// (`prompt_tokens` + the matched prefix) to read that prefill-kernel reference. +fn check_lossless( + handle: &EngineHandle, + tokenizer: &DynTokenizer, + i: usize, + prompt: &str, + prompt_tokens: &[u32], + base: &[Step], + spec: &[Step], +) -> Result<(), String> { + let matched = base + .iter() + .zip(spec) + .take_while(|(b, s)| b.id == s.id) + .count(); + + // Identical sequences (or one a prefix of the other): perfectly lossless. + if matched == base.len().min(spec.len()) { + eprintln!( + "prompt {i} ({prompt:?}): {matched}/{} tokens identical (100% lossless)", + base.len() + ); + return Ok(()); + } + + let spec_id = spec[matched].id; + let decode_argmax = base[matched].top_logprobs[0].0; + + // Diagnostic: show the exact branch point. + { + let lo = matched.saturating_sub(2); + let hi = (matched + 3).min(base.len()).min(spec.len()); + let base_ids: Vec = base[..hi].iter().map(|s| s.id).collect(); + let spec_ids: Vec = spec[..hi].iter().map(|s| s.id).collect(); + eprintln!(" [diag] prompt {i} matched={matched}"); + eprintln!( + " [diag] context+gen base ids {:?} = {:?}", + &base_ids, + tokenizer.decode(&base_ids, false).unwrap_or_default() + ); + eprintln!( + " [diag] base[{lo}..{hi}] = {:?}", + base[lo..hi] + .iter() + .map(|s| (s.id, tokenizer.decode(&[s.id], false).unwrap_or_default())) + .collect::>() + ); + eprintln!( + " [diag] spec[{lo}..{hi}] = {:?}", + spec[lo..hi] + .iter() + .map(|s| (s.id, tokenizer.decode(&[s.id], false).unwrap_or_default())) + .collect::>() + ); + let _ = spec_ids; + } + + // The verify path runs the prefill kernel, so the right reference for the + // spec pick is a plain *prefill* of the same shared context — not the + // plain-decode baseline, whose kernel resolves a bifurcation tie to the + // other side and amplifies the gap. + let mut context = prompt_tokens.to_vec(); + context.extend(base[..matched].iter().map(|s| s.id)); + let prefill_ref = prefill_next(handle, context, LOGPROBS); + + if prefill_ref.id == spec_id { + // Spec faithfully reproduced the prefill-kernel greedy pick; the + // divergence is purely the pre-existing prefill-vs-decode kernel gap. + let decode_lp = base[matched] + .top_logprobs + .iter() + .find(|(t, _)| *t == spec_id) + .map(|(_, lp)| base[matched].top_logprobs[0].1 - lp); + eprintln!( + "prompt {i} ({prompt:?}): kernel-gap flip at token {matched} — verify(prefill)→{spec_id}, \ + decode→{decode_argmax}; spec matches prefill greedy (decode-margin {:?}). Not a spec bug.", + decode_lp + ); + return Ok(()); + } + + // Spec's greedy pick differs from the prefill-kernel argmax too. The verify + // path builds its committed KV incrementally across batched speculative + // spans while this reference prefill builds it in one shot; the two differ by + // a few bf16 ULP. Within MARGIN_TOL of the prefill argmax ⇒ a benign tie + // flip; clearly worse ⇒ the verify/accept/capture logic picked a token the + // forward never favored — a real bug. + let prefill_regret = prefill_ref + .top_logprobs + .iter() + .find(|(t, _)| *t == spec_id) + .map(|(_, lp)| prefill_ref.top_logprobs[0].1 - lp); + + if let Some(regret) = prefill_regret { + if regret <= MARGIN_TOL { + eprintln!( + "prompt {i} ({prompt:?}): tie flip at token {matched} — \ + verify(prefill)→{}, spec→{spec_id}, decode→{decode_argmax}; \ + spec pick is #2 in the prefill distribution (regret {regret:.3} ≤ {MARGIN_TOL}). \ + Not a spec bug.", + prefill_ref.id, + ); + return Ok(()); + } + } + + // Either the spec pick is outside the prefill top-K entirely, or it sits + // clearly below the prefill argmax — neither is a benign tie. + let decode_regret = base[matched] + .top_logprobs + .iter() + .find(|(t, _)| *t == spec_id) + .map(|(_, lp)| base[matched].top_logprobs[0].1 - lp); + Err(format!( + "prompt {i}: at token {matched} spec chose {spec_id} but prefill greedy says {} and \ + decode greedy says {decode_argmax} (spec regret in prefill dist: {prefill_regret:?} > \ + {MARGIN_TOL}; in decode dist: {decode_regret:?}) — real spec bug", + prefill_ref.id, + )) +} + #[test] fn dflash_speculative_greedy_matches_plain_greedy() { let (Some(model_path), Some(draft_path)) = (target_path_or_skip(), draft_path_or_skip()) else { @@ -224,7 +414,7 @@ fn dflash_speculative_greedy_matches_plain_greedy() { .expect("failed to start baseline engine"); let out = encoded .iter() - .map(|t| generate(&handle, t.clone(), LOGPROBS)) + .map(|t| generate(&handle, t.clone(), LOGPROBS, GENERATED_TOKENS)) .collect(); drop(handle); // Let the scheduler thread tear down and free GPU memory before the @@ -244,130 +434,210 @@ fn dflash_speculative_greedy_matches_plain_greedy() { .expect("failed to start speculative engine"); let mut failures = Vec::new(); - for (i, prompt) in prompts.iter().enumerate() { - let base = &baseline[i]; - let spec = generate(&handle, encoded[i].clone(), 0); - let matched = base - .iter() - .zip(&spec) - .take_while(|(b, s)| b.id == s.id) - .count(); - - // Identical sequences (or one a prefix of the other): perfectly lossless. - if matched == base.len().min(spec.len()) { - eprintln!( - "prompt {i} ({prompt:?}): {matched}/{} tokens identical (100% lossless)", - base.len() - ); - continue; + for (i, &prompt) in prompts.iter().enumerate() { + let spec = generate(&handle, encoded[i].clone(), 0, GENERATED_TOKENS); + if let Err(failure) = check_lossless( + &handle, + &tokenizer, + i, + prompt, + &encoded[i], + &baseline[i], + &spec, + ) { + failures.push(failure); } + } - let spec_id = spec[matched].id; - let decode_argmax = base[matched].top_logprobs[0].0; + drop(handle); - // Diagnostic: show the exact branch point. - { - let lo = matched.saturating_sub(2); - let hi = (matched + 3).min(base.len()).min(spec.len()); - let base_ids: Vec = base[..hi].iter().map(|s| s.id).collect(); - let spec_ids: Vec = spec[..hi].iter().map(|s| s.id).collect(); - eprintln!(" [diag] prompt {i} matched={matched}"); - eprintln!( - " [diag] context+gen base ids {:?} = {:?}", - &base_ids, - tokenizer.decode(&base_ids, false).unwrap_or_default() - ); - eprintln!( - " [diag] base[{lo}..{hi}] = {:?}", - base[lo..hi] - .iter() - .map(|s| (s.id, tokenizer.decode(&[s.id], false).unwrap_or_default())) - .collect::>() - ); - eprintln!( - " [diag] spec[{lo}..{hi}] = {:?}", - spec[lo..hi] - .iter() - .map(|s| (s.id, tokenizer.decode(&[s.id], false).unwrap_or_default())) - .collect::>() - ); - let _ = spec_ids; - } + assert!( + failures.is_empty(), + "speculative greedy decode is not lossless:\n{}", + failures.join("\n") + ); +} - // The verify path runs the prefill kernel, so the right reference for the - // spec pick is a plain *prefill* of the same shared context — not the - // plain-decode baseline, whose kernel resolves a bifurcation tie to the - // other side and amplifies the gap. Build that context from the matched - // tokens and ask what the prefill kernel predicts next. - let mut context = encoded[i].clone(); - context.extend(base[..matched].iter().map(|s| s.id)); - let prefill_ref = prefill_next(&handle, context, LOGPROBS); - - if prefill_ref.id == spec_id { - // Spec faithfully reproduced the prefill-kernel greedy pick; the - // divergence is purely the pre-existing prefill-vs-decode kernel gap - // at a near-tie (here decode→{decode_argmax}, prefill→{spec_id}). - let decode_lp = base[matched] - .top_logprobs - .iter() - .find(|(t, _)| *t == spec_id) - .map(|(_, lp)| base[matched].top_logprobs[0].1 - lp); - eprintln!( - "prompt {i} ({prompt:?}): kernel-gap flip at token {matched} — verify(prefill)→{spec_id}, \ - decode→{decode_argmax}; spec matches prefill greedy (decode-margin {:?}). Not a spec bug.", - decode_lp - ); - continue; - } +/// Verify-graph capture-shape regression (heterogeneous `max_tokens`). +/// +/// The piecewise verify CUDA Graph keys its captured dense segments by +/// `batch_size` alone, but a request near its output budget shortens its verify +/// span (`scheduler::plan` truncates the span to the remaining budget), so +/// `total_tokens` — the row count the captured segments bake into their launch +/// grid — varies at a *fixed* batch size. A graph captured at a short span and +/// then replayed at a longer one processes too few rows: the trailing requests +/// read stale logits, silently breaking the lossless contract. +/// +/// Neither existing check can see this. The bs=1 gate above issues each request +/// sequentially, so every fresh request's first verify is a *full* span and +/// bucket bs=1 is always first-captured at the maximal shape (only the harmless +/// over-compute direction occurs). A homogeneous concurrent benchmark is no +/// better: lockstep requests capture every bucket at full span during ramp-up, +/// and all truncation happens later as they finish together (still the safe +/// direction). The dangerous direction needs *heterogeneous* progress. +/// +/// This reproduces it deterministically, single-stream: a `max_tokens=8` request +/// (span < `block_size`) captures the bucket-bs=1 graph at a truncated shape, +/// then a `max_tokens=64` request on the *same* engine replays that poisoned +/// graph at the full span. On the buggy code the long request diverges from +/// plain greedy; with full-shape gating (truncated spans run eager, so the graph +/// is only ever captured/replayed at the maximal shape) it stays lossless. +#[test] +fn dflash_short_then_long_verify_capture_is_lossless() { + let (Some(model_path), Some(draft_path)) = (target_path_or_skip(), draft_path_or_skip()) else { + return; + }; + let _gpu = GPU.lock().unwrap_or_else(|p| p.into_inner()); - // Spec's greedy pick differs from the prefill-kernel argmax too. The - // verify path builds its committed KV incrementally across batched - // speculative spans, while this reference prefill builds it in one - // shot; the two differ by a few bf16 ULP. On a near-tie that flips the - // argmax — benign. So the deciding question is *how far* below the - // prefill argmax the spec pick sits IN THE PREFILL KERNEL'S OWN - // distribution (the kernel the verify path uses). Within MARGIN_TOL ⇒ - // a numerical tie flip, not a bug. Clearly worse ⇒ the verify/accept - // logic picked a token the forward never favored — a real bug. - let prefill_regret = prefill_ref - .top_logprobs + // << block_size (16): the poison request's only verify step is a short + // truncated span, capturing the bucket-bs=1 graph at total_tokens far below + // the full span. The fewer valid rows, the sooner a full-span replay hits the + // stale tail — so the victim diverges early, well clear of its token budget. + const POISON_MAX_TOKENS: usize = 4; + let poison_prompt = "Hello, world! Tell me a story."; + let victim_prompt = "Q: What is 17 multiplied by 23? A: Let's think step by step."; + + let tokenizer = common::load_tokenizer(&model_path); + let poison_tokens = tokenizer + .encode(poison_prompt, false) + .expect("encode failed"); + let victim_tokens = tokenizer + .encode(victim_prompt, false) + .expect("encode failed"); + + // 1. Baseline: the victim's plain-greedy decode (spec off) with logprobs, for + // the regret reference at any divergence. + let baseline = { + let handle = openinfer_qwen3_4b::launch(Path::new(&model_path), launch_options(None)) + .expect("failed to start baseline engine"); + let out = generate(&handle, victim_tokens.clone(), LOGPROBS, GENERATED_TOKENS); + drop(handle); + // Free the target before the speculative engine loads the same 8 GB. + std::thread::sleep(Duration::from_secs(2)); + out + }; + + // 2. Speculative engine, shared across both requests so the bucket-bs=1 + // capture from the poison request persists into the victim's replay. + let handle = openinfer_qwen3_4b::launch( + Path::new(&model_path), + launch_options(Some(PathBuf::from(&draft_path))), + ) + .expect("failed to start speculative engine"); + + // Poison: a short request whose only verify step has total_tokens < span, + // first-capturing the bucket-bs=1 graph at the truncated shape. + let poison = generate(&handle, poison_tokens, 0, POISON_MAX_TOKENS); + assert!( + poison.len() <= POISON_MAX_TOKENS, + "poison request emitted {} tokens, expected <= {POISON_MAX_TOKENS}", + poison.len() + ); + + // Victim: a full-span replay of the poisoned bucket-bs=1 graph. + let spec = generate(&handle, victim_tokens.clone(), 0, GENERATED_TOKENS); + + let result = check_lossless( + &handle, + &tokenizer, + 0, + victim_prompt, + &victim_tokens, + &baseline, + &spec, + ); + drop(handle); + + assert!( + result.is_ok(), + "verify capture-shape bug: the long request diverged from plain greedy after a short \ + request poisoned the bucket-bs=1 graph at a truncated span:\n{}", + result.unwrap_err() + ); +} + +/// Concurrent, heterogeneous-`max_tokens` losslessness coverage for the bs>1 +/// draft+verify path. The bs=1 gate and the homogeneous c8/c16 benches never +/// exercise a real batch with requests at *different* verify-span lengths; this +/// runs several greedy requests concurrently with staggered budgets and asserts +/// each stays lossless vs its own plain-greedy baseline (tolerating only the +/// benign bf16 tie-flip via the shared regret check). A batched-draft indexing +/// regression or a capture-shape mismatch at bs>1 would surface here as a real +/// (non-tie) divergence. +#[test] +fn dflash_concurrent_heterogeneous_is_lossless() { + let (Some(model_path), Some(draft_path)) = (target_path_or_skip(), draft_path_or_skip()) else { + return; + }; + let _gpu = GPU.lock().unwrap_or_else(|p| p.into_inner()); + + // Distinct prompts with staggered budgets: at any tick the in-flight batch + // mixes full and near-budget (truncated) verify spans. + let cases: [(&str, usize); 4] = [ + ("def fibonacci(n):", 64), + ("The three primary colors are", 24), + ( + "Q: What is 17 multiplied by 23? A: Let's think step by step.", + 48, + ), + ("Here is a short story about a dragon. Once upon a time", 40), + ]; + + let tokenizer = common::load_tokenizer(&model_path); + let encoded: Vec> = cases + .iter() + .map(|(p, _)| tokenizer.encode(p, false).expect("encode failed")) + .collect(); + + // 1. Baselines: each prompt's plain-greedy decode (spec off) at ITS budget, + // with logprobs for the regret reference. Sequential, one engine. + let baselines: Vec> = { + let handle = openinfer_qwen3_4b::launch(Path::new(&model_path), launch_options(None)) + .expect("failed to start baseline engine"); + let out = encoded .iter() - .find(|(t, _)| *t == spec_id) - .map(|(_, lp)| prefill_ref.top_logprobs[0].1 - lp); - - if let Some(regret) = prefill_regret { - if regret <= MARGIN_TOL { - eprintln!( - "prompt {i} ({prompt:?}): tie flip at token {matched} — \ - verify(prefill)→{}, spec→{spec_id}, decode→{decode_argmax}; \ - spec pick is #2 in the prefill distribution (regret {regret:.3} ≤ {MARGIN_TOL}). \ - Not a spec bug.", - prefill_ref.id, - ); - continue; - } - } + .zip(&cases) + .map(|(t, (_, max_tokens))| generate(&handle, t.clone(), LOGPROBS, *max_tokens)) + .collect(); + drop(handle); + std::thread::sleep(Duration::from_secs(2)); + out + }; - // Either the spec pick is outside the prefill top-K entirely, or it sits - // clearly below the prefill argmax — neither is a benign tie. - let decode_regret = base[matched] - .top_logprobs + // 2. Speculative engine: submit all four at once so they form real batches. + let handle = openinfer_qwen3_4b::launch( + Path::new(&model_path), + launch_options(Some(PathBuf::from(&draft_path))), + ) + .expect("failed to start speculative engine"); + let specs = generate_concurrent( + &handle, + encoded .iter() - .find(|(t, _)| *t == spec_id) - .map(|(_, lp)| base[matched].top_logprobs[0].1 - lp); - failures.push(format!( - "prompt {i}: at token {matched} spec chose {spec_id} but prefill greedy says {} and \ - decode greedy says {decode_argmax} (spec regret in prefill dist: {prefill_regret:?} > \ - {MARGIN_TOL}; in decode dist: {decode_regret:?}) — real spec bug", - prefill_ref.id, - )); - } + .zip(&cases) + .map(|(t, (_, max_tokens))| (t.clone(), *max_tokens)) + .collect(), + ); + let mut failures = Vec::new(); + for (i, (prompt, _)) in cases.iter().enumerate() { + if let Err(failure) = check_lossless( + &handle, + &tokenizer, + i, + prompt, + &encoded[i], + &baselines[i], + &specs[i], + ) { + failures.push(failure); + } + } drop(handle); assert!( failures.is_empty(), - "speculative greedy decode is not lossless:\n{}", + "concurrent heterogeneous speculative decode is not lossless:\n{}", failures.join("\n") ); }