speculative: Qwen3 draft-model v0 with paged verify parity

Phase 22 lands a correctness-only speculative decoding loop for Qwen3
target + Qwen3 small draft (batch=1, greedy, gamma=4). Phase 23 turns
verify logits into the authoritative acceptance signal so mirror-decode
per accepted token is no longer needed.

- paged_kv_cache: truncate_sequence(slot, new_len) shrinks a registered
  sequence, freeing whole physical blocks no longer reachable and
  leaving the slot registered. Covered by a CUDA-gated unit test.
- qwen3: forward_verify_paged_decode_attention writes the draft window
  into the target cache, runs the same paged decode attention kernel per
  draft token, and uses matmul_rows_gemv so linear layers follow the
  single-token decode BF16 rounding path.
- bench-speculative: new bench binary drives the state machine with
  --gamma / --gen-tokens / --prompts / --use-verify-logits /
  --verify-path flash|paged-decode / --dump-verify-mismatches, and
  compares baseline vs spec token sequences plus TPOT / tok/s / speedup.
- docs/22 records the decode-authoritative v0 result and dash5 numbers
  (matched=true, speedup_e2e ~0.29x, verify_decode_mismatches>0 under
  --use-verify-logits).
- docs/23 records the paged-decode verify path (matched=true,
  verify_decode_mismatches=0, 50x64 speedup_e2e ~0.44x) and the
  next-step performance TODO.
This commit is contained in:
2026-07-01 14:15:39 +08:00
parent 5b350ee5f0
commit ce7229f4fe
5 changed files with 1444 additions and 0 deletions

View File

@@ -486,6 +486,35 @@ impl PagedKVCache {
state.seq_len += num_tokens;
}
/// Roll a registered sequence back to `new_len` tokens.
///
/// This only changes cache metadata and frees whole physical blocks that are
/// no longer reachable. Bytes inside retained blocks are left untouched; the
/// logical `seq_len` prevents attention from reading them, and later writes
/// to the same positions overwrite them.
pub fn truncate_sequence(&mut self, slot: usize, new_len: usize) -> Result<(), &'static str> {
if slot >= self.max_seqs {
return Err("truncate_sequence: slot out of range");
}
let state = self.seq_states[slot]
.as_mut()
.ok_or("truncate_sequence: empty slot")?;
if new_len > state.seq_len {
return Err("truncate_sequence: cannot extend");
}
let needed_blocks = ((new_len + BLOCK_SIZE - 1) / BLOCK_SIZE).max(1);
while state.block_ids.len() > needed_blocks {
let block = state.block_ids.pop().expect("checked len");
match state.location {
Location::Gpu => self.allocator.free(block),
Location::Cpu => self.cpu_allocator.free(block),
}
}
state.seq_len = new_len;
Ok(())
}
/// Refresh the host-side block table + context lens from `seq_states`,
/// then upload to GPU. Call once per decode step before the paged kernel.
pub fn sync_to_gpu(&mut self) {
@@ -748,6 +777,71 @@ impl PagedKVCache {
}
}
#[cfg(test)]
mod tests {
use super::*;
fn tiny_config() -> ModelConfig {
serde_json::from_value(serde_json::json!({
"model_type": "qwen3",
"hidden_size": 8,
"intermediate_size": 16,
"num_attention_heads": 1,
"num_key_value_heads": 1,
"num_hidden_layers": 1,
"vocab_size": 32,
"max_position_embeddings": 64
}))
.unwrap()
}
#[test]
fn truncate_sequence_frees_whole_blocks_and_keeps_slot_registered() {
if xserv_cuda::device::set_device(0).is_err() {
eprintln!("skipping CUDA-backed PagedKVCache test: device 0 unavailable");
return;
}
let config = tiny_config();
let mut cache = PagedKVCache::new(&config, 5, 0, 1, 4, DType::BF16, 0);
assert_eq!(
cache.truncate_sequence(1, 0),
Err("truncate_sequence: slot out of range")
);
assert_eq!(
cache.truncate_sequence(0, 0),
Err("truncate_sequence: empty slot")
);
cache.register_sequence(0).unwrap();
cache.ensure_capacity(0, BLOCK_SIZE * 3 + 1);
cache.advance_seq_len(0, BLOCK_SIZE * 3 + 1);
assert_eq!(cache.seq_len(0), BLOCK_SIZE * 3 + 1);
assert_eq!(cache.block_count(0), 4);
assert_eq!(cache.free_blocks(), 0);
cache.truncate_sequence(0, BLOCK_SIZE + 1).unwrap();
assert_eq!(cache.seq_len(0), BLOCK_SIZE + 1);
assert_eq!(cache.block_count(0), 2);
assert_eq!(cache.free_blocks(), 2);
cache.truncate_sequence(0, BLOCK_SIZE).unwrap();
assert_eq!(cache.seq_len(0), BLOCK_SIZE);
assert_eq!(cache.block_count(0), 1);
assert_eq!(cache.free_blocks(), 3);
cache.truncate_sequence(0, 0).unwrap();
assert_eq!(cache.seq_len(0), 0);
assert_eq!(cache.block_count(0), 1);
assert_eq!(cache.free_blocks(), 3);
assert_eq!(
cache.truncate_sequence(0, 1),
Err("truncate_sequence: cannot extend")
);
}
}
unsafe fn tensor_from_owned_buf(
buf: GpuBuffer,
shape: &[usize],