feat(kernel): add gfx950 Gluon bf16 MLA decode and prefill kernels#664
feat(kernel): add gfx950 Gluon bf16 MLA decode and prefill kernels#664Yu-Zhewen wants to merge 1 commit into
Conversation
Add absorbed MLA decode (kv_lora_rank=512, rope=64) with two regimes --
bh16bn64 for <=16 heads and bh64 for {64,128} heads on a 3-D XCD-aware
grid -- plus a non-absorbed MLA prefill kernel (qk=192, v=128, causal
and non-causal). Register both under the gfx950 gluon solution, extend
the mla_decode_with_kvcache traits with num_q_heads for regime dispatch,
and add parity tests against the reference kernels.
Signed-off-by: Yu-Zhewen <zhewenyu@amd.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: d11449bcb8
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| traits={ | ||
| "q_len": frozenset({1}), | ||
| # bh16bn64 (<=16 heads) + bh64 ({64, 128} heads, large batch). | ||
| "num_q_heads": frozenset(range(1, 17)) | {64, 128}, |
There was a problem hiding this comment.
When num_q_heads is 64 or 128 but the batch is not divisible by 64 (for example normal single-request decode), this registration still wins because the traits only check head count. The selected gluon_mla_decode_bf16_gfx950 implementation then raises NotImplementedError for that batch shape, so default dispatch fails instead of falling back to the existing Triton MLA decode. Add a dispatch trait/gate for the batch constraint or avoid advertising the 64/128-head regime unless the dispatcher can prove the large-batch requirement.
Useful? React with 👍 / 👎.
| if WITHIN_2GB: | ||
| gl.amd.cdna4.async_copy.buffer_load_to_shared(bufs_kv0, Kv_c_cache, offs_k_c0, mask=offs_n_nope0[None, :] < split_kv_end) | ||
| else: | ||
| gl.amd.cdna4.async_copy.global_load_to_shared(bufs_kv0, Kv_c_cache + offs_k_c0) |
There was a problem hiding this comment.
When WITHIN_2GB is false, over-provisioned max_seqlen_k can launch empty trailing splits (split_kv_start == split_kv_end), and the page-table load above is fully masked. This fallback still performs an unmasked global_load_to_shared using the resulting undefined page numbers, unlike the masked buffer_load path, so large KV caches can dereference arbitrary addresses for short/empty splits before the reduce masks them out. The global-load fallback needs an equivalent guard or to avoid issuing loads for empty splits.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
if the gluon is backed by the shared mla kernel api. maybe we can re-use the existing tests in tokenspeed-kernel/test/ops/test_attention_mla.py?
| # buffer_load uses a scalar base + 32-bit offsets; KV pools > 2 GB fall back | ||
| # to global_load (64-bit pointers). | ||
| max_kv_bytes = kv_c.shape[0] * kv_c.stride(0) * kv_c.element_size() | ||
| within_2gb = max_kv_bytes <= 0x80000000 |
There was a problem hiding this comment.
is this within_2gb estimation safe? maybe we should have a test to cover it
| stride_mid_lse_s, | ||
| Final_lse, # RETURN_LSE only: merged fp32 lse [B, H] (else None) | ||
| stride_final_lse_b, | ||
| stride_final_lse_h, |
There was a problem hiding this comment.
are all these strides variables not constexpr?
There was a problem hiding this comment.
maybe we can restructure decode like prefill (tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/attention/gluon/mla_prefill_bf16_gfx950.py) to split config, program and pipeline structure. so it is easier to read.
2 participants
Kernels included
gluon_mla_decode_bf16_gfx950— absorbed MLA decode over a paged compressed KV cache (kv_lora_rank=512,qk_rope=64,page_size=64,q_len=1). Two regimes selected by head count:bh16bn64—num_q_heads ≤ 16, 2-D(batch, split)grid.bh64—num_q_heads ∈ {64, 128}, 3-D XCD-aware grid (requiresbatch % 64 == 0).return_lse;logit_capis not supported (guarded).gluon_mla_prefill_bf16_gfx950— non-absorbed varlen MLA prefill (qk_head_dim=192,v_head_dim=128), causal and non-causal, with optionalreturn_lse.logit_capis not supported (guarded).Performance comparison
Decode — bh16bn64 regime (num_q_heads ≤ 16, 2-D grid)
h= number of query heads,b= batch size (number of sequences),s= sequence length in tokens (KV context length for decode; per-sequence query/KV length for prefill).speedup= triton ref latency ÷ gluon latency (higher = gluon faster).h=8 b=1 s=512h=8 b=1 s=1024h=8 b=1 s=2048h=8 b=1 s=8192h=8 b=1 s=16384h=8 b=1 s=32768h=8 b=8 s=512h=8 b=8 s=1024h=8 b=8 s=2048h=8 b=8 s=8192h=8 b=8 s=16384h=8 b=8 s=32768h=8 b=32 s=512h=8 b=32 s=1024h=8 b=32 s=2048h=8 b=32 s=8192h=8 b=32 s=16384h=8 b=32 s=32768h=16 b=1 s=512h=16 b=1 s=1024h=16 b=1 s=2048h=16 b=1 s=8192h=16 b=1 s=16384h=16 b=1 s=32768h=16 b=8 s=512h=16 b=8 s=1024h=16 b=8 s=2048h=16 b=8 s=8192h=16 b=8 s=16384h=16 b=8 s=32768h=16 b=32 s=512h=16 b=32 s=1024h=16 b=32 s=2048h=16 b=32 s=8192h=16 b=32 s=16384h=16 b=32 s=32768Decode — bh64 regime (num_q_heads ∈ {64,128}, 3-D XCD grid, batch % 64 == 0)
h=64 b=64 s=512h=64 b=64 s=1024h=64 b=64 s=2048h=64 b=64 s=8192h=64 b=64 s=16384h=64 b=128 s=512h=64 b=128 s=1024h=64 b=128 s=2048h=64 b=128 s=8192h=64 b=128 s=16384h=128 b=64 s=512h=128 b=64 s=1024h=128 b=64 s=2048h=128 b=64 s=8192h=128 b=64 s=16384h=128 b=128 s=512h=128 b=128 s=1024h=128 b=128 s=2048h=128 b=128 s=8192h=128 b=128 s=16384Prefill — causal
h=16 b=1 s=1024h=16 b=1 s=2048h=16 b=1 s=4096h=16 b=1 s=8192h=16 b=1 s=16384h=16 b=2 s=1024h=16 b=2 s=2048h=16 b=2 s=4096h=16 b=2 s=8192h=16 b=2 s=16384h=16 b=4 s=1024h=16 b=4 s=2048h=16 b=4 s=4096h=16 b=4 s=8192h=16 b=4 s=16384h=16 b=8 s=1024h=16 b=8 s=2048h=16 b=8 s=4096h=16 b=8 s=8192h=16 b=8 s=16384h=16 b=16 s=1024h=16 b=16 s=2048h=16 b=16 s=4096h=16 b=16 s=8192h=16 b=16 s=16384h=64 b=1 s=1024h=64 b=1 s=2048h=64 b=1 s=4096h=64 b=1 s=8192h=64 b=1 s=16384h=64 b=2 s=1024h=64 b=2 s=2048h=64 b=2 s=4096h=64 b=2 s=8192h=64 b=2 s=16384h=64 b=4 s=1024h=64 b=4 s=2048h=64 b=4 s=4096h=64 b=4 s=8192h=64 b=4 s=16384h=64 b=8 s=1024h=64 b=8 s=2048h=64 b=8 s=4096h=64 b=8 s=8192h=64 b=8 s=16384h=64 b=16 s=1024h=64 b=16 s=2048h=64 b=16 s=4096h=64 b=16 s=8192h=64 b=16 s=16384Signed-off-by: Yu-Zhewen zhewenyu@amd.com