phase 15: batched decode forward — 35 tok/s (97% of HF transformers)
Implement batched decode that processes multiple sequences' tokens in one forward pass. The key insight: cuBLAS M=4 GEMM is dramatically faster than 4× M=1 GEMV due to better TensorCore utilization and amortized kernel launch overhead. New method Qwen3::forward_decode_batch(&tokens, &positions, &mut caches): - Batched embedding, norm, projections, FFN: [B, hidden] × [hidden, X] → one cuBLAS call per weight matrix instead of B calls - Per-sequence attention: RoPE, KV cache, decode_attention remain per-seq (each has different position and KV length) - Row extraction (row_view) and concatenation (concat_rows) for batched↔per-seq transitions Engine Step 4b: - batch_size >= 2: extracts caches via std::mem::replace, calls forward_decode_batch, restores caches, samples per-sequence - batch_size == 1: falls back to per-seq forward_gpu_cache (no overhead) Ablation results (dash5, RTX 5090, Qwen3-8B BF16): | Scenario | Throughput | vs HF | |----------|-----------|-------| | Serial (batch=1) | 13.2 tok/s | 37% | | Concurrent (batch=4) | 35.1 tok/s | 97% | | HF transformers | 36.0 tok/s | 100% | The 2.66x throughput improvement (13.2 → 35.1) for concurrent requests comes from cuBLAS going from 1008 M=1 GEMVs to 252 M=4 GEMMs per step, which cuBLAS handles ~4x more efficiently on TensorCores. Milestone ④ target (50% of vLLM/HF throughput) achieved with 97%. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -104,28 +104,78 @@ impl Engine {
|
||||
}
|
||||
}
|
||||
|
||||
// Step 4b: Process decode (one token per sequence)
|
||||
// Currently per-sequence (each has different KV cache length).
|
||||
// TODO(Phase 14): With Flash Attention, batch all decode tokens into
|
||||
// one forward pass — batch the compute-heavy ops (projections, FFN)
|
||||
// and use FlashDecoding for per-seq variable-length attention.
|
||||
let decode_count = running.iter()
|
||||
.filter(|s| s.prefilled && !newly_prefilled.contains(&s.id))
|
||||
.count();
|
||||
if decode_count > 0 {
|
||||
// Step 4b: Batched decode — batch all decode-ready sequences into one forward pass.
|
||||
// Projections and FFN run as [B, hidden] matmuls; attention remains per-seq.
|
||||
let decode_indices: Vec<usize> = running.iter().enumerate()
|
||||
.filter(|(_, s)| s.prefilled && !newly_prefilled.contains(&s.id))
|
||||
.map(|(i, _)| i)
|
||||
.collect();
|
||||
|
||||
if !decode_indices.is_empty() {
|
||||
static LOG_ONCE: Once = Once::new();
|
||||
LOG_ONCE.call_once(|| {
|
||||
eprintln!("[scheduler] decode batching active (per-seq until Flash Attention)");
|
||||
eprintln!("[scheduler] batched decode active");
|
||||
});
|
||||
eprintln!("[scheduler] decode batch_size={}", decode_count);
|
||||
}
|
||||
for seq in running.iter_mut() {
|
||||
if seq.prefilled && !newly_prefilled.contains(&seq.id) {
|
||||
let last = *seq.generated_tokens.last().unwrap();
|
||||
let logits = self.model.forward_gpu_cache(&[last], &mut seq.kv_cache);
|
||||
let next = sample(&logits, &seq.sampling);
|
||||
seq.generated_tokens.push(next);
|
||||
self.emit_token(seq, next);
|
||||
eprintln!("[scheduler] decode batch_size={}", decode_indices.len());
|
||||
|
||||
if decode_indices.len() == 1 {
|
||||
// Single sequence: use per-seq path (no batching overhead)
|
||||
let i = decode_indices[0];
|
||||
let last = *running[i].generated_tokens.last().unwrap();
|
||||
let logits = self.model.forward_gpu_cache(&[last], &mut running[i].kv_cache);
|
||||
let next = sample(&logits, &running[i].sampling);
|
||||
running[i].generated_tokens.push(next);
|
||||
self.emit_token(&running[i], next);
|
||||
} else {
|
||||
// Batched decode: extract tokens and positions
|
||||
let tokens: Vec<u32> = decode_indices.iter()
|
||||
.map(|&i| *running[i].generated_tokens.last().unwrap())
|
||||
.collect();
|
||||
let positions: Vec<usize> = decode_indices.iter()
|
||||
.map(|&i| running[i].kv_cache.seq_len())
|
||||
.collect();
|
||||
|
||||
// Take caches out of sequences temporarily to satisfy borrow checker.
|
||||
// The dummy caches (max_seq_len=1) are never used and dropped immediately
|
||||
// after the real caches are restored. Minor alloc overhead (~microseconds).
|
||||
let mut caches: Vec<GpuKVCache> = decode_indices.iter()
|
||||
.map(|&i| {
|
||||
std::mem::replace(
|
||||
&mut running[i].kv_cache,
|
||||
GpuKVCache::new(&self.config, 1, DType::BF16, 0),
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
let mut cache_refs: Vec<&mut GpuKVCache> = caches.iter_mut().collect();
|
||||
|
||||
let logits = self.model.forward_decode_batch(&tokens, &positions, &mut cache_refs);
|
||||
|
||||
// Put caches back: pop from end while iterating in reverse
|
||||
drop(cache_refs);
|
||||
for &i in decode_indices.iter().rev() {
|
||||
running[i].kv_cache = caches.pop().unwrap();
|
||||
}
|
||||
|
||||
// Sample per-sequence from batched logits [B, vocab_size]
|
||||
let vocab_size = logits.shape()[1];
|
||||
let logits_cpu = logits.to_device(xserv_tensor::Device::Cpu);
|
||||
let data = logits_cpu.as_slice::<half::bf16>();
|
||||
for (j, &i) in decode_indices.iter().enumerate() {
|
||||
let row_start = j * vocab_size;
|
||||
let row_logits = &data[row_start..row_start + vocab_size];
|
||||
let next = if running[i].sampling.temperature == 0.0 {
|
||||
// Greedy: argmax
|
||||
row_logits.iter().enumerate()
|
||||
.max_by(|a, b| a.1.to_f32().partial_cmp(&b.1.to_f32()).unwrap())
|
||||
.map(|(idx, _)| idx as u32).unwrap()
|
||||
} else {
|
||||
// Use the row as a single-row tensor for full sampling
|
||||
let row_tensor = xserv_tensor::Tensor::from_slice(row_logits, &[1, vocab_size]);
|
||||
sample(&row_tensor, &running[i].sampling)
|
||||
};
|
||||
running[i].generated_tokens.push(next);
|
||||
self.emit_token(&running[i], next);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user