Skip to content

fix(qwen3-dflash): post-merge correctness fixes — capture-shape gate, unified readiness, stream #442

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
xiaguan merged 7 commits into from
Jun 23, 2026
Merged

fix(qwen3-dflash): post-merge correctness fixes — capture-shape gate, unified readiness, stream #442

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
648bcba
fix(qwen3-dflash): gate verify CUDA Graph on the full span shape
xiaguan Jun 23, 2026
ad95ea4
fix(kernels): launch DFlash hidden-copy kernels on the active stream
xiaguan Jun 23, 2026
633ee9b
fix(qwen3-dflash): capture draft context on prefill, never the unifie…
xiaguan Jun 23, 2026
6832481
fix(qwen3-dflash): reject decode overlap combined with speculation at…
xiaguan Jun 23, 2026
cf7384c
refactor(qwen3-dflash): split weight loading and memory accounting in…
xiaguan Jun 23, 2026
596c8ce
test(qwen3-dflash): cover concurrent heterogeneous-budget losslessness
xiaguan Jun 23, 2026
e5283a4
docs(qwen3): mark DFlash unified-readiness and stream-race blockers r…
xiaguan Jun 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 9 additions & 5 deletions docs/models/qwen3/dflash-speculative-decoding.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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).
4 changes: 2 additions & 2 deletions openinfer-kernels/src/ops/elementwise.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()?;
Expand Down Expand Up @@ -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()?;
Expand Down
225 changes: 6 additions & 219 deletions openinfer-qwen3-4b/src/dflash.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -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<Self> {
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,
Expand Down Expand Up @@ -375,138 +294,6 @@ impl DFlashBatchScratch {
}

impl DFlashDraftModel {
pub(crate) fn from_safetensors_for_target(
ctx: &DeviceContext,
model_path: &str,
target: &Qwen3Model,
) -> Result<Self> {
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
}
Expand Down
Loading