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:
@@ -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],
|
||||
|
||||
Reference in New Issue
Block a user