DeepSpeed
1e6951e4 - feat(kernels): batched decode attention

Commit
25 days ago
feat(kernels): batched decode attention decode_attn was hardcoded to batch 1 (a single shared write_pos, flat [heads, dim] query). One block now covers a (row, query-head) pair with the grid spanning batch * heads, K/V index by row, and write_pos accepts a shared scalar or one entry per row — so batched decode rides the kernel instead of falling back to SDPA against the full-width static cache (the fallback measured 54% of vLLM at b=8, 2K context). The reference implementation and kernel test cover b=1, b=4 shared-pos, and b=4 per-row-pos. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Author
Committer
Parents
Loading