perf(k3): warp-cooperative two-pass MLA absorbed decode kernel - #938
perf(k3): warp-cooperative two-pass MLA absorbed decode kernel#938xiaguan wants to merge 1 commit into
Conversation
Rewrite mla_paged_absorbed_attn_kernel from three serial-scalar sweeps into a two-pass warp-cooperative form: pass 1 computes bf16-landed scores per page with online max/denominator tracking; pass 2 recomputes the scores bit-identically, quantizes probabilities through the same bf16 chain as before (division, not reciprocal-multiply), and attends. Score dot products vectorize over __nv_bfloat162 lanes with a warp reduction; the FFI signature is unchanged and odd layer_offset / page_stride are rejected. The bf16 landings keep the rounding chain inside the calibrated 2.0-ULP noise floor: golden decode gates 13/13, spec_verify 6/6, workspace lib tests green. Serve A/B on the pruned 224-expert EP4 checkpoint (same c1 prompts, per-batch-bucket kernel averages): b=16 verify step: 799us -> 444us per launch (1.80x) b=8: 347us -> 256us (1.36x) A b=16 verify step spends 10.7ms instead of 19.2ms in the 24 MLA layers; round p50 64.65ms -> 60.92ms. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: xiaguan <751080330@qq.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 833344547b
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| // Schedule: two passes, warp-cooperative scores | ||
| // --------------------------------------------------------------------------- | ||
| // 1. q_abs[0..512) = bf16(f32 sum_d q_nope[d] * W_UK[d, j]), d ascending; | ||
| // q_abs[512..576) = q_rope, copied bf16. | ||
| // 2. dot(t) = f32 sum_d q_abs[d] * c_t[d] over 576, d ascending, one thread. | ||
| // 3. scl(t) = f32( bf16(dot) * scale ), the product taken in bf16. | ||
| // 4. m = max_t scl; tot = sum_t exp(scl - m) in f32 (per-thread strided | ||
| // partials in ascending t, then a fixed-order tree reduction). | ||
| // 5. p_t = bf16( exp(scl - m) / tot ). | ||
| // 6. o_lat[j] = bf16( f32 sum_t p_t * c_t[j] ), t ascending (chunk-major). | ||
| // 7. o[dv] = bf16( f32 sum_j W_UV[dv, j] * o_lat[j] ), j ascending. | ||
| // The context is walked twice, page by page. Pass one computes each page's |
There was a problem hiding this comment.
The source now documents and implements a two-pass warp-cooperative schedule, but pegainfer-kernels/KERNELS.md:149 still describes this operation as a three-sweep recompute with serial score semantics. This leaves the repository's kernel inventory inaccurate for anyone auditing performance or numerical behavior; update that entry alongside this implementation change.
AGENTS.md reference: AGENTS.md:L168-L170
Useful? React with 👍 / 👎.
1 participant
What
Rewrites
mla_paged_absorbed_attn_kernel(the absorbed-MLA decode/verify path — plain decode and spec-decode verify both run through it) from three serial-scalar sweeps into a two-pass warp-cooperative kernel:__nv_bfloat162-vectorized q·c dot products with warp reduction, each score landed through bf16 exactly as before; online running max + denominator across pages.bf16(exp(s−m)/l), division not reciprocal-multiply), then attend in ascending token order.FFI signature unchanged; odd
layer_offset/page_strideare rejected. Padding semantics preserved (negative page → score 0 participates in softmax, contributes nothing).Why
EP16 round anatomy (#936 docs) showed the verify step burning ~68ms in local compute, far off the weight-read roofline. The absorbed MLA kernel was the largest single non-GEMM contributor: one thread sweeping scores serially.
Correctness
The bf16 landings keep the new summation order inside the calibrated 2.0-ULP golden noise floor:
golden_decodegates 13/13 (mega + masked-chain transports, incl. step/prefill time snapshots)spec_verify6/6Perf (pruned 224-expert EP4 serve, same c1 prompts, per-batch-bucket kernel averages)
A b=16 verify step spends 10.7ms instead of 19.2ms in its 24 MLA layers; serve round p50 64.65ms → 60.92ms.
🤖 Generated with Claude Code