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

speculative: batched-GEMV kernel for verify path (Phase 24 step 1)

Browse Source 
Add launch_gemv_bf16_batched: runs M m=1 GEMVs in a single 3D grid
launch (z = batch row) with numerically identical output to M sequential
launch_gemv_bf16 calls — same K-block partial accumulation, same
fixed-order reduction. Verified on dash5 with 10 prompts × 32 tokens:
matched=true, verify_decode_mismatches=0.

Expose as matmul_batched_gemv(a: [M,K], b: [K,N]) → [M,N] in
xserv-kernels. Replace the old matmul_rows_gemv helper in qwen3
forward_verify_paged_decode_attention; the per-row loop over matmul_2d +
concat_rows is replaced by a single matmul_batched_gemv call that
allocates the partials buffer in one shot and launches 2 kernels instead
of 2*M.

Current speedup_e2e is 0.47× (same ballpark as Phase 23 0.44×);
the batched launch saves ~3 ms overhead but this is small relative to
the total 28 ms spec cost. The path forward (per docs/24 §4) is
higher acceptance rate or cheaper draft, not further kernel optimization.
This commit is contained in:
Gahow Wang
2026-07-01 16:13:37 +08:00
parent 42e13f33dd
commit e5734b41fa
4 changed files with 155 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -923,7 +923,7 @@ impl Qwen3 {
let residual = x.clone();
let normed = rmsnorm(&x, &layer.input_norm, eps);
let qkv = matmul_rows_gemv(&normed, &layer.qkv_proj_wt);
let qkv = matmul_batched_gemv(&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);
@@ -966,25 +966,25 @@ impl Qwen3 {
);
let attn_merged = attn_out.reshape(&[new_tokens, num_heads * head_dim]);
let attn_proj = matmul_rows_gemv(&attn_merged, &layer.o_proj_wt);
let attn_proj = matmul_batched_gemv(&attn_merged, &layer.o_proj_wt);
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();
let gate_up = matmul_rows_gemv(&normed, &layer.gate_up_proj_wt);
let gate_up = matmul_batched_gemv(&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_rows_gemv(&hidden_states, &layer.down_proj_wt);
let down = matmul_batched_gemv(&hidden_states, &layer.down_proj_wt);
self.all_reduce(&down);
x = add_any(&residual, &down);
}
let x = rmsnorm(&x, &self.norm, eps);
matmul_rows_gemv(&x, &self.lm_head_t)
matmul_batched_gemv(&x, &self.lm_head_t)
}
/// Forward with GPU-resident KV cache and GPU transpose/reshape kernels.
@@ -1261,20 +1261,6 @@ fn row_view(t: &Tensor, row: usize) -> Tensor {
)
}
/// Run a 2D matmul row by row so each row uses the same GEMV kernel as
/// single-token decode. Used by speculative verify parity, where near-tie
/// logits must follow decode's BF16 rounding path.
fn matmul_rows_gemv(a: &Tensor, b: &Tensor) -> Tensor {
assert_eq!(a.ndim(), 2);
assert!(a.is_contiguous());
let rows = a.shape()[0];
if rows == 1 {
return matmul_2d(a, b);
}
let out_rows: Vec<Tensor> = (0..rows).map(|i| matmul_2d(&row_view(a, i), b)).collect();
concat_rows(&out_rows)
}
/// Concatenate row tensors [1, cols] into a single [B, cols] tensor via D2D memcpy.
fn concat_rows(rows: &[Tensor]) -> Tensor {
assert!(!rows.is_empty());