gahow/xserv
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

speculative: Qwen3 decode graph + gamma sweep (Phase 24 step 2)

Browse Source 
- Split Qwen3::forward_decode_paged into decode_prepare (host-side
  block allocation + table upload) and decode_core (pure-GPU compute
  reading token ids and positions from device buffers via
  embedding_device_ids + rope_inplace_device_pos). This makes the
  entire Qwen3 decode step CUDA-graph-capturable, mirroring the
  gpt_oss.rs architecture.
- Add qwen3_graph.rs: Qwen3DecodeGraph + GraphedQwen3Decoder, a port
  of the gpt_oss_graph.rs whole-step capture pattern. Lazy policy:
  first decode eager (warms pool + cuBLAS), second captures, rest
  replay. Batch>1 always falls back to eager.
- Wire GraphedQwen3Decoder into bench-speculative's draft decode path;
  all 4 draft.forward_decode_paged call sites + replay_draft_tokens
  now route through the graphed decoder. Per-benchmark caches persist
  across prompts for graph reuse.
- Gamma sweep result (10 prompts × 32 tokens, --use-verify-logits):
  γ=1 → 0.57×, γ=2 → 0.57×, γ=4 → 0.49×, γ=6 → 0.41×, γ=8 → 0.36×.
  All matched=true, verify_decode_mismatches=0.
  Acceptance drops sharply with γ (66% → 40% → 25%) because Qwen3-0.6B
  is too inaccurate a draft for Qwen3-8B. Speedup still <1.

Current ceiling analysis: verify costs ~13ms (same as one target decode)
so speculative decoding only wins if acceptance × (tokens/round) >>
(draft_cost + verify_cost) / baseline_decode. With this draft model,
the crossover requires either (a) a much smaller verify cost (batch-GEMM
path, which trades correctness), or (b) a fundamentally better drafter
(EAGLE-style heads, or n-gram lookup).
This commit is contained in:
Gahow Wang
2026-07-01 16:32:17 +08:00
parent e5734b41fa
commit a77239c0c8
4 changed files with 271 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files

87
crates/xserv-model/src/qwen3.rs
View File

@@ -701,45 +701,72 @@ impl Qwen3 {
assert_eq!(seq_slots.len(), batch);
assert!(batch > 0);
// TP: this rank owns a slice of the heads (local_* == full when world==1).
let num_heads = self.local_num_heads;
let num_kv_heads = self.local_num_kv_heads;
let head_dim = self.config.head_dim();
let eps = self.config.rms_norm_eps.unwrap_or(1e-6) as f32;
self.decode_prepare(positions, seq_slots, paged_cache);
// Ensure all slots have enough physical blocks for this token, then
// upload block tables + context_lens once for the whole forward (the
// tables are identical across layers; only the layer's K/V pool changes).
let ids_gpu = upload_u32(tokens);
let positions_u32: Vec<u32> = positions.iter().map(|&p| p as u32).collect();
let pos_gpu = upload_u32(&positions_u32);
let logits = self.decode_core(
ids_gpu.as_ptr() as *const std::ffi::c_void,
pos_gpu.as_ptr() as *const std::ffi::c_void,
batch,
seq_slots,
paged_cache,
);
logits
}
/// Host-side per-step cache bookkeeping: block allocation + uploading block
/// tables / context lens to their (stable-address) GPU buffers. Runs
/// OUTSIDE any CUDA-graph captured region.
pub fn decode_prepare(
&self,
positions: &[usize],
seq_slots: &[usize],
paged_cache: &mut PagedKVCache,
) {
let kv_lens: Vec<i32> = positions.iter().map(|&p| (p + 1) as i32).collect();
for (b, &slot) in seq_slots.iter().enumerate() {
paged_cache.ensure_capacity(slot, positions[b] + 1);
}
paged_cache.sync_active_batch_with_lens(seq_slots, &kv_lens);
}
/// Pure-GPU decode step: embedding → all layers → final norm → logits.
/// Token ids and positions are read from device buffers; every other input
/// (weights, KV pools, block table, context lens) has a stable address —
/// which makes this region CUDA-graph capturable.
pub fn decode_core(
&self,
ids_gpu: *const std::ffi::c_void,
pos_gpu: *const std::ffi::c_void,
batch: usize,
seq_slots: &[usize],
paged_cache: &mut PagedKVCache,
) -> Tensor {
let num_heads = self.local_num_heads;
let num_kv_heads = self.local_num_kv_heads;
let head_dim = self.config.head_dim();
let eps = self.config.rms_norm_eps.unwrap_or(1e-6) as f32;
let bt_ptr = paged_cache.block_table_gpu().as_ptr() as *const i32;
let cl_ptr = paged_cache.context_lens_gpu().as_ptr() as *const i32;
let max_blocks = paged_cache.max_blocks_per_seq();
// RoPE expects `[num_tokens, H, D]` with `num_tokens` positions —
// matches our `[B, H, D]` exactly, so we upload once here.
let positions_u32: Vec<u32> = positions.iter().map(|&p| p as u32).collect();
// Batched embedding: [B, hidden]
let mut x = embedding(&self.embed_tokens, tokens);
let mut x = embedding_device_ids(&self.embed_tokens, ids_gpu, batch);
for (layer_idx, layer) in self.layers.iter().enumerate() {
let residual = x.clone();
let normed = rmsnorm(&x, &layer.input_norm, eps);
// Fused QKV projection: one GEMV instead of three.
let qkv = matmul_2d(&normed, &layer.qkv_proj_wt); // [B, (H+2*KV)*D]
let qkv = matmul_2d(&normed, &layer.qkv_proj_wt);
let q_dim = num_heads * head_dim;
let kv_dim = num_kv_heads * head_dim;
let q_all = qkv.narrow(1, 0, q_dim); // [B, H*D] (view)
let k_all = qkv.narrow(1, q_dim, kv_dim); // [B, KV*D] (view)
let q_all = qkv.narrow(1, 0, q_dim);
let k_all = qkv.narrow(1, q_dim, kv_dim);
let v_all = qkv.narrow(1, q_dim + kv_dim, kv_dim);
// Per-head RMSNorm on contiguous copies (narrow views are strided).
let q_flat = q_all.contiguous().reshape(&[batch * num_heads, head_dim]);
let k_flat = k_all
.contiguous()
@@ -749,16 +776,13 @@ impl Qwen3 {
let q_3d = q_normed.reshape(&[batch, num_heads, head_dim]);
let k_3d = k_normed.reshape(&[batch, num_kv_heads, head_dim]);
rope_inplace(&q_3d, &self.rope_cache, &positions_u32);
rope_inplace(&k_3d, &self.rope_cache, &positions_u32);
rope_inplace_device_pos(&q_3d, &self.rope_cache, pos_gpu);
rope_inplace_device_pos(&k_3d, &self.rope_cache, pos_gpu);
let v_3d = v_all.contiguous().reshape(&[batch, num_kv_heads, head_dim]);
// Single batched scatter for all sequences in the batch.
paged_cache.append_tokens_batched(layer_idx, &k_3d, &v_3d, batch);
// Paged attention reads Q as [B, H, 1, D] — a contiguous view
// of [B, H, D].
let q_4d = q_3d.reshape(&[batch, num_heads, 1, head_dim]);
let k_pool_ptr = paged_cache.k_pool(layer_idx).as_ptr() as *const std::ffi::c_void;
let v_pool_ptr = paged_cache.v_pool(layer_idx).as_ptr() as *const std::ffi::c_void;
@@ -775,27 +799,24 @@ impl Qwen3 {
max_blocks,
);
// attn_out shape [B, H, 1, D] is contiguous-equivalent to [B, H*D].
let attn_merged = attn_out.reshape(&[batch, num_heads * head_dim]);
let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt);
self.all_reduce(&attn_proj); // TP: sum partial attention outputs
self.all_reduce(&attn_proj);
let (normed, x_new) =
xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
let residual = x_new.clone();
// Fused gate+up projection: one GEMV instead of two.
let gate_up = matmul_2d(&normed, &layer.gate_up_proj_wt); // [B, 2*ffn]
let gate_up = matmul_2d(&normed, &layer.gate_up_proj_wt);
let ffn_dim = gate_up.shape()[1] / 2;
let gate = gate_up.narrow(1, 0, ffn_dim).contiguous();
let up = gate_up.narrow(1, ffn_dim, ffn_dim).contiguous();
let hidden_states = xserv_kernels::silu_mul(&gate, &up);
let down = matmul_2d(&hidden_states, &layer.down_proj_wt);
self.all_reduce(&down); // TP: sum partial MLP outputs
self.all_reduce(&down);
x = add_any(&residual, &down);
}
// Advance logical seq_len now that all layers have been written.
for &slot in seq_slots {
paged_cache.advance_seq_len(slot, 1);
}
@@ -1261,6 +1282,14 @@ fn row_view(t: &Tensor, row: usize) -> Tensor {
)
}
/// Upload a u32 slice to a pooled GPU buffer (synchronous H2D).
fn upload_u32(vals: &[u32]) -> xserv_cuda::GpuBuffer {
let bytes = unsafe { std::slice::from_raw_parts(vals.as_ptr() as *const u8, vals.len() * 4) };
let mut buf = xserv_cuda::allocator::cached_alloc(bytes.len()).expect("alloc u32 upload");
buf.copy_from_host(bytes).unwrap();
buf
}
/// Concatenate row tensors [1, cols] into a single [B, cols] tensor via D2D memcpy.
fn concat_rows(rows: &[Tensor]) -> Tensor {
assert!(!rows.is_empty());