MLA kernel flydsl bf16 - #4616
Open
ahmed-bsod wants to merge 3 commits into
Open
Conversation
ahmed-bsod
force-pushed
the
ahmed/mla-latest-rebase
branch
from
August 6, 2026 22:15
7f1e500 to
7bdad45
Compare
Contributor
There was a problem hiding this comment.
Pull request overview
Warning
Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.
Adds an A16W16 FlyDSL MLA (Multi-head Latent Attention) decode kernel targeting gfx1250, plus correctness/scale tests and a Python API wrapper to validate and launch the kernel.
Changes:
- Introduce
flydsl_mla_decodePython wrapper with shape/device/dtype checks and kernel dispatch. - Add gfx1250 shuffled-layout FlyDSL main + reduce kernels for MLA decode.
- Add pytest correctness tests and large-shape stress tests for the new kernel.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 7 comments.
| File | Description |
|---|---|
| op_tests/flydsl_tests/test_flydsl_mla_decode.py | Adds correctness + “large” test coverage for the new MLA decode kernel (including KV shuffle helper + torch reference). |
| aiter/ops/flydsl/mla_decode.py | Adds the public flydsl_mla_decode API with validation, partitioning logic, and kernel launches. |
| aiter/ops/flydsl/kernels/mla_decode_shuffled_gfx1250.py | Implements the gfx1250 FlyDSL MLA decode main and reduce kernels (shuffled KV layout). |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+285
to
+303
| _LARGE_CASES = [ | ||
| (1024, 8192), | ||
| (1024, 16384), | ||
| (1024, 32768), | ||
| ] | ||
|
|
||
|
|
||
| @pytest.mark.parametrize("num_seqs,ctx_len", _LARGE_CASES) | ||
| @pytest.mark.parametrize("num_q_heads", _NUM_Q_HEADS) | ||
| @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) | ||
| def test_flydsl_mla_decode_large(num_seqs, ctx_len, num_q_heads, dtype): | ||
| _run_decode_case( | ||
| num_seqs=num_seqs, | ||
| ctx_len=ctx_len, | ||
| dtype=dtype, | ||
| varlen=True, | ||
| block_size=64, | ||
| num_query_heads=num_q_heads, | ||
| ) |
Comment on lines
+269
to
+282
| @pytest.mark.parametrize("num_seqs,ctx_len", _CASES) | ||
| @pytest.mark.parametrize("num_q_heads", _NUM_Q_HEADS) | ||
| @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) | ||
| @pytest.mark.parametrize("varlen", [True, False]) | ||
| @pytest.mark.parametrize("block_size", _BLOCK_SIZES) | ||
| def test_flydsl_mla_decode(num_seqs, ctx_len, num_q_heads, dtype, varlen, block_size): | ||
| _run_decode_case( | ||
| num_seqs=num_seqs, | ||
| ctx_len=ctx_len, | ||
| dtype=dtype, | ||
| varlen=varlen, | ||
| block_size=block_size, | ||
| num_query_heads=num_q_heads, | ||
| ) |
Comment on lines
+163
to
+170
| if varlen: | ||
| lens = [ | ||
| int(max(random.normalvariate(ctx_len, ctx_len / 2), ctx_len)) | ||
| for _ in range(num_seqs) | ||
| ] | ||
| seq_lens = torch.tensor(lens, dtype=torch.int32, device=device) | ||
| else: | ||
| seq_lens = torch.full((num_seqs,), ctx_len, dtype=torch.int32, device=device) |
Comment on lines
+62
to
+71
| def shuffle(kvb, h): | ||
| kvb = kvb.view( | ||
| -1, | ||
| num_kv_heads, | ||
| block_size // num_lanes, | ||
| num_lanes, | ||
| h // (2 * num_elements_per_thread), | ||
| 2, # 2 thread groups: t0..t15 and t16..t31 | ||
| num_elements_per_thread, | ||
| ) |
Comment on lines
+129
to
+131
| num_seqs, num_q_heads, qk_head_dim = query.shape | ||
| # Pre-shuffled cache axis order: [nb, kv_heads, block, head]. | ||
| num_blocks, num_kv_heads, block_size, qk_head_dim_kv = kv_cache.shape |
Comment on lines
+141
to
+142
| if num_kv_heads != 1: | ||
| raise ValueError(f"MLA expects num_kv_heads == 1, got {num_kv_heads}") |
Comment on lines
+134
to
+137
| # using fmath.exp2 lowers to llvm.exp2.f32 which adds extra instructions | ||
| # to guard x < -126 underflow. We do not need to waste instrucitons on the | ||
| # guard here because ULP of sum is big enough that we donot care. | ||
| # instead we use bare hardware exp2 |
ahmed-bsod
force-pushed
the
ahmed/mla-latest-rebase
branch
from
August 7, 2026 01:25
7bdad45 to
f3aead4
Compare
ahmed-bsod
force-pushed
the
ahmed/mla-latest-rebase
branch
from
August 7, 2026 01:27
f3aead4 to
4e4be05
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Add A16W16 flydsl MLA kernel