Add native MLA QH64 fp8 persistent decode kernel for gfx942 by vstakhov-amd · Pull Request #3188 · ROCm/aiter

vstakhov-amd · 2026-05-14T16:04:29Z

Motivation

Technical Details

Add native QH64 (gqa_ratio=64, qSeqLen=1) fp8 persistent decode support for MI300X (gfx942), eliminating the need for the folded QH16 path that reshapes 64 heads into 4x batch expansion.

Kernel binaries:

mla_a8w8_qh64_qseqlen1_gqaratio64_v3_ps.co (decode, no LSE)
mla_a8w8_qh64_qseqlen1_gqaratio64_lse_v3_ps.co (decode, with LSE)

Framework changes:

mla.py: bypass fold path for gfx942 nhead=64 fp8 persistent
attention.py: metadata tile size (ceil(64/128)=1 tile per batch)
v1_2_device.cuh: mark gfx942 nhead=64 as natively_supported to prevent batch expansion in metadata generator
mla_asm.csv: register both kernel variants
test_mla_persistent.py: enable gfx942 in nhead=64 test path

Tested on MI300X across batch={1,4,16,64,128} x ctx={1k,4k,8k,16k,32k}. All 25 configurations pass correctness. Peak: 367 TFLOPS at b128/ctx32k.

Test Plan

Correctness validated across 25 configurations: batch={1,4,16,64,128} x ctx={1k,4k,8k,16k,32k}
All checkAllclose checks pass (fp8 golden reference)
No regressions on existing nhead=16/128 paths

Test Result

┌─────────┬─────────┬─────────┬──────────┬──────────┬───────────┐
│ ctx_len │ batch=1 │ batch=4 │ batch=16 │ batch=64 │ batch=128 │
├─────────┼─────────┼─────────┼──────────┼──────────┼───────────┤
│ 1024    │ 23 us   │ 27 us   │ 35 us    │ 60 us    │ 88 us     │
├─────────┼─────────┼─────────┼──────────┼──────────┼───────────┤
│ 4096    │ 30 us   │ 32 us   │ 62 us    │ 139 us   │ 255 us    │
├─────────┼─────────┼─────────┼──────────┼──────────┼───────────┤
│ 8192    │ 39 us   │ 41 us   │ 86 us    │ 247 us   │ 441 us    │
├─────────┼─────────┼─────────┼──────────┼──────────┼───────────┤
│ 16384   │ 54 us   │ 58 us   │ 140 us   │ 430 us   │ 820 us    │
├─────────┼─────────┼─────────┼──────────┼──────────┼───────────┤
│ 32768   │ 84 us   │ 103 us  │ 243 us   │ 815 us   │ 1591 us   │
└─────────┴─────────┴─────────┴──────────┴──────────┴───────────┘

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

Add native QH64 (gqa_ratio=64, qSeqLen=1) fp8 persistent decode support for MI300X (gfx942), eliminating the need for the folded QH16 path that reshapes 64 heads into 4x batch expansion. Kernel binaries: - mla_a8w8_qh64_qseqlen1_gqaratio64_v3_ps.co (decode, no LSE) - mla_a8w8_qh64_qseqlen1_gqaratio64_lse_v3_ps.co (decode, with LSE) Framework changes: - mla.py: bypass fold path for gfx942 nhead=64 fp8 persistent - attention.py: metadata tile size (ceil(64/128)=1 tile per batch) - v1_2_device.cuh: mark gfx942 nhead=64 as natively_supported to prevent batch expansion in metadata generator - mla_asm.csv: register both kernel variants - test_mla_persistent.py: enable gfx942 in nhead=64 test path Tested on MI300X across batch={1,4,16,64,128} x ctx={1k,4k,8k,16k,32k}. All 25 configurations pass correctness. Peak: 367 TFLOPS at b128/ctx32k.

github-actions · 2026-05-14T16:05:01Z

🏷️ CI Guide

Runs automatically on every PR:

✅ Pre-checks (submodule verification, code formatting)
✅ Aiter op tests (gfx942 + gfx950)
✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
`ci:triton-300x`	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
`ci:sglang`	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
`ci:atom`	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
`ci:atom_full`	ATOM accuracy suite for PR and main models from ATOM `models_accuracy.json`
`ci:vllm`	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
`ci:all`	All standard extended tests (excludes `ci:atom_full`)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 3188 --add-label <label>

JohnNikolay84

LGTM

Byeong-Chan · 2026-05-26T11:36:26Z

Thank you for your hard works!

I have one question.

calling test_mla(ctx_lens                     = 1024,
                 batch_size                   = 1,
                 nhead                        = 64,
                 kv_lora_rank                 = 512,
                 qk_nope_head_dim             = 512,
                 qk_rope_head_dim             = 64,
                 v_head_dim                   = 512,
                 dtype                        = torch.float8_e4m3fnuz,
                 kvtype                       = torch.float8_e4m3fnuz,
                 page_size                    = 64,
                 varlen                       = False,
                 decode_qlen                  = 1,
                 max_split_per_batch          = 32,
                 non_persistent_mode          = False,
                 paged_layout                 = LEGACY,
                 scale_dim                    = 4,
                 return_lse                   = False)
[aiter] import [module_mla_metadata] under /sgl-workspace/aiter/aiter/jit/module_mla_metadata.so
[aiter] LoadKernel: _ZN5aiter39mla_a8w8_qh64_qseqlen1_gqaratio64_v3_psE hsaco: /sgl-workspace/aiter/hsa//gfx942/mla/mla_a8w8_qh64_qseqlen1_gqaratio64_v3_ps.co
[aiter] import [module_mla_reduce] under /sgl-workspace/aiter/aiter/jit/module_mla_reduce.so
[W526 11:20:20.334764147 collection.cpp:1133] Warning: ROCTracer produced duplicate flow start: 3 (function operator())
[aiter] mla_decode-absorb_fp8    [golden vs aiter_asm]:    13.82 us......[checkAllclose atol=0.01 rtol=0.01 failed!]
    a    : torch.Size([1, 64, 512])
           tensor([-0.0079, -0.1338,  0.0330, -0.0292, -0.0192, -0.0128, -0.0569, -0.0344],
       device='cuda:0', dtype=torch.bfloat16)
    b    : torch.Size([1, 64, 512])
           tensor([ 0.2119, -0.1924,  0.3730, -0.8438, -0.6562, -0.0923,  0.2168, -0.1187],
       device='cuda:0', dtype=torch.bfloat16)
    delta:
           tensor([0.2197, 0.0586, 0.3398, 0.8164, 0.6367, 0.0796, 0.2734, 0.0840],
       device='cuda:0', dtype=torch.bfloat16)
[aiter] -->max abs delta:2.046875, delta details: 97.5% (31947 of 32768) elements
[aiter] mla_decode-absorb_fp8    [golden fp8 vs aiter_asm]:    13.82 us......[checkAllclose atol=0.01 rtol=0.01 failed!]
    a    : torch.Size([1, 64, 512])
           tensor([ 0.0164, -0.1040,  0.0610, -0.0315, -0.0016, -0.0152, -0.0278, -0.0415],
       device='cuda:0', dtype=torch.bfloat16)
    b    : torch.Size([1, 64, 512])
           tensor([ 0.2119, -0.1924,  0.3730, -0.8438,  0.0086, -0.6562, -0.0923,  0.2168],
       device='cuda:0', dtype=torch.bfloat16)
    delta:
           tensor([0.1953, 0.0884, 0.3125, 0.8125, 0.0103, 0.6406, 0.0645, 0.2578],
       device='cuda:0', dtype=torch.bfloat16)
[aiter] -->max abs delta:2.046875, delta details: 97.5% (31936 of 32768) elements
[aiter] mla_decode-absorb_fp8    [golden fp8 split_out_ref vs aiter_asm]:    13.82 us......[checkAllclose atol=0.01 rtol=0.01 failed!]
    a    : torch.Size([1, 64, 512])
           tensor([ 0.0164, -0.1040,  0.0610, -0.0315, -0.0016, -0.0152, -0.0278, -0.0415],
       device='cuda:0', dtype=torch.bfloat16)
    b    : torch.Size([1, 64, 512])
           tensor([ 0.2119, -0.1924,  0.3730, -0.8438,  0.0086, -0.6562, -0.0923,  0.2168],
       device='cuda:0', dtype=torch.bfloat16)
    delta:
           tensor([0.1953, 0.0884, 0.3125, 0.8125, 0.0103, 0.6406, 0.0645, 0.2578],
       device='cuda:0', dtype=torch.bfloat16)
[aiter] -->max abs delta:2.046875, delta details: 97.5% (31936 of 32768) elements
Traceback (most recent call last):
  File "/sgl-workspace/aiter/op_tests/test_mla_persistent.py", line 1782, in <module>
    ret = test_mla(
  File "/sgl-workspace/aiter/aiter/test_common.py", line 129, in wrapper
    ret = func(*args, **kwargs)
  File "/sgl-workspace/aiter/op_tests/test_mla_persistent.py", line 1612, in test_mla
    err, us_asm_decode = test_absorb_decode_fp8()
  File "/sgl-workspace/aiter/op_tests/test_mla_persistent.py", line 1486, in test_absorb_decode_fp8
    cal_diff(out_ref, out_asm, "out", True)
  File "/sgl-workspace/aiter/op_tests/test_mla_persistent.py", line 247, in cal_diff
    assert cos_diff < 3e-2

I met the cos_diff error when testing block_size 64.

It's ok about block_size 1.

Could you check precision for block_size 64?

vstakhov-amd · 2026-05-27T09:55:46Z

Could you check precision for block_size 64?

You are right, there is a switch to paged code path that was not correctly triggered. I'm working on the fix so far, thanks for testing!

Updated mla_a8w8_qh64_qseqlen1_gqaratio64{_lse,}_v3_ps.co with corrected output writeback that no longer overruns ptr_RP / ptr_LSEP per launch. Raise kv_granularity floor to 32 for nhead=64 fp8/fp8 in the persistent test so the metadata layer never emits sub-32-token partial works to this kernel (it requires SUB_KV=32 partial granularity).

…_gfx942 # Conflicts: # aiter/mla.py # aiter/ops/attention.py # csrc/kernels/mla/metadata/v1_2_device.cuh # op_tests/test_mla_persistent.py

vstakhov-amd · 2026-06-02T11:35:34Z

@Byeong-Chan I have hopefully fixed the issues with paging and output writing, could you please test that if it's still relevant?

Byeong-Chan · 2026-06-04T10:50:03Z

@vstakhov-amd It has still precision issue in page_size 64.

But if we removed 21 ctx length, then it's ok to precision.

I think short length(<64) causes some issue.

vstakhov-amd · 2026-06-04T10:55:54Z

Is it the same or a different one? I have tested exactly the configuration you have mentioned with ctx_lens=1024, was it different this time?

JohnNikolay84 · 2026-06-04T10:59:29Z

@Byeong-Chan it might be the case indeed that ctx_len < block_size is having an issue. We will have a look at that case.

Byeong-Chan · 2026-06-04T11:00:43Z

Is it the same or a different one? I have tested exactly the configuration you have mentioned with ctx_lens=1024, was it different this time?

it's ok about context length 1024 and non-varlen.

but, if we set short length(<64) or varlen, then it still has precision issue.

varlen case:

[aiter] import [module_aiter_core] under /sgl-workspace/aiter/aiter/jit/module_aiter_core.so
[aiter] 
calling test_mla(ctx_lens                     = 1024,
                 batch_size                   = 3,
                 nhead                        = 64,
                 kv_lora_rank                 = 512,
                 qk_nope_head_dim             = 512,
                 qk_rope_head_dim             = 64,
                 v_head_dim                   = 512,
                 dtype                        = torch.float8_e4m3fnuz,
                 kvtype                       = torch.float8_e4m3fnuz,
                 page_size                    = 64,
                 varlen                       = True,
                 decode_qlen                  = 1,
                 max_split_per_batch          = 32,
                 non_persistent_mode          = False,
                 paged_layout                 = LEGACY,
                 scale_dim                    = 4,
                 return_lse                   = False)
[aiter] import [module_mla_metadata] under /sgl-workspace/aiter/aiter/jit/module_mla_metadata.so
[aiter] LoadKernel: _ZN5aiter39mla_a8w8_qh64_qseqlen1_gqaratio64_v3_psE hsaco: /sgl-workspace/aiter/hsa//gfx942/mla/mla_a8w8_qh64_qseqlen1_gqaratio64_v3_ps.co
[aiter] import [module_mla_reduce] under /sgl-workspace/aiter/aiter/jit/module_mla_reduce.so
[W604 12:00:41.270967635 collection.cpp:1133] Warning: ROCTracer produced duplicate flow start: 3 (function operator())
[aiter] mla_decode-absorb_fp8    [golden vs aiter_asm]:    65.24 us......[checkAllclose atol=0.01 rtol=0.01 failed!]
    a    : torch.Size([3, 64, 512])
           tensor([-0.0051,  0.0146, -0.0486,  0.0505,  0.0391,  0.0449, -0.0187, -0.0569],
       device='cuda:0', dtype=torch.bfloat16)
    b    : torch.Size([3, 64, 512])
           tensor([-0.0176, -0.0023, -0.0374,  0.0615, -0.0182,  0.0723, -0.0425, -0.0459],
       device='cuda:0', dtype=torch.bfloat16)
    delta:
           tensor([0.0125, 0.0170, 0.0112, 0.0110, 0.0571, 0.0273, 0.0238, 0.0110],
       device='cuda:0', dtype=torch.bfloat16)
[aiter] -->max abs delta:0.146484375, delta details: 42.2% (41486 of 98304) elements
[aiter] mla_decode-absorb_fp8    [golden fp8 vs aiter_asm]:    65.24 us......[checkAllclose atol=0.01 rtol=0.01 failed!]
    a    : torch.Size([3, 64, 512])
           tensor([-0.0045,  0.0193, -0.0508,  0.0398,  0.0437, -0.0178, -0.0591,  0.0613],
       device='cuda:0', dtype=torch.bfloat16)
    b    : torch.Size([3, 64, 512])
           tensor([-0.0176, -0.0023, -0.0374, -0.0182,  0.0723, -0.0425, -0.0459,  0.0913],
       device='cuda:0', dtype=torch.bfloat16)
    delta:
           tensor([0.0131, 0.0216, 0.0134, 0.0581, 0.0286, 0.0247, 0.0132, 0.0300],
       device='cuda:0', dtype=torch.bfloat16)
[aiter] -->max abs delta:0.1484375, delta details: 60.5% (59496 of 98304) elements
[aiter] mla_decode-absorb_fp8    [golden fp8 split_out_ref vs aiter_asm]:    65.24 us......[checkAllclose atol=0.01 rtol=0.01 failed!]
    a    : torch.Size([3, 64, 512])
           tensor([-0.0045,  0.0193, -0.0508,  0.0398,  0.0437, -0.0178, -0.0591,  0.0613],
       device='cuda:0', dtype=torch.bfloat16)
    b    : torch.Size([3, 64, 512])
           tensor([-0.0176, -0.0023, -0.0374, -0.0182,  0.0723, -0.0425, -0.0459,  0.0913],
       device='cuda:0', dtype=torch.bfloat16)
    delta:
           tensor([0.0131, 0.0216, 0.0134, 0.0581, 0.0286, 0.0247, 0.0132, 0.0300],
       device='cuda:0', dtype=torch.bfloat16)
[aiter] -->max abs delta:0.1484375, delta details: 60.5% (59496 of 98304) elements
[aiter] 
calling test_mla(ctx_lens                     = 1024,
                 batch_size                   = 3,
                 nhead                        = 64,
                 kv_lora_rank                 = 512,
                 qk_nope_head_dim             = 512,
                 qk_rope_head_dim             = 64,
                 v_head_dim                   = 512,
                 dtype                        = torch.float8_e4m3fnuz,
                 kvtype                       = torch.float8_e4m3fnuz,
                 page_size                    = 64,
                 varlen                       = True,
                 decode_qlen                  = 1,
                 max_split_per_batch          = 32,
                 non_persistent_mode          = False,
                 paged_layout                 = LEGACY,
                 scale_dim                    = 4,
                 return_lse                   = False)
[aiter] mla_decode-absorb_fp8    [golden vs aiter_asm]:    44.54 us......[checkAllclose atol=0.01 rtol=0.01 failed!]
    a    : torch.Size([3, 64, 512])
           tensor([ 0.1660, -0.1934,  0.0583,  0.0038, -0.0581, -0.0156, -0.0277, -0.0693],
       device='cuda:0', dtype=torch.bfloat16)
    b    : torch.Size([3, 64, 512])
           tensor([ 0.2139, -0.1543,  0.1318, -0.0152, -0.1338,  0.0304, -0.0151, -0.1216],
       device='cuda:0', dtype=torch.bfloat16)
    delta:
           tensor([0.0479, 0.0391, 0.0732, 0.0190, 0.0757, 0.0459, 0.0126, 0.0522],
       device='cuda:0', dtype=torch.bfloat16)
[aiter] -->max abs delta:0.21484375, delta details: 60.4% (59409 of 98304) elements
[aiter] mla_decode-absorb_fp8    [golden fp8 vs aiter_asm]:    44.54 us......[checkAllclose atol=0.01 rtol=0.01 failed!]
    a    : torch.Size([3, 64, 512])
           tensor([ 0.1699, -0.2021,  0.0349,  0.0256, -0.0623,  0.0032, -0.0131,  0.0016],
       device='cuda:0', dtype=torch.bfloat16)
    b    : torch.Size([3, 64, 512])
           tensor([ 0.2139, -0.1543,  0.1318, -0.0152, -0.1338,  0.0275,  0.0304, -0.0151],
       device='cuda:0', dtype=torch.bfloat16)
    delta:
           tensor([0.0439, 0.0479, 0.0967, 0.0408, 0.0713, 0.0242, 0.0435, 0.0167],
       device='cuda:0', dtype=torch.bfloat16)
[aiter] -->max abs delta:0.46875, delta details: 75.7% (74446 of 98304) elements
[aiter] mla_decode-absorb_fp8    [golden fp8 split_out_ref vs aiter_asm]:    44.54 us......[checkAllclose atol=0.01 rtol=0.01 failed!]
    a    : torch.Size([3, 64, 512])
           tensor([ 0.1699, -0.2021,  0.0349,  0.0256, -0.0623,  0.0032, -0.0131,  0.0016],
       device='cuda:0', dtype=torch.bfloat16)
    b    : torch.Size([3, 64, 512])
           tensor([ 0.2139, -0.1543,  0.1318, -0.0152, -0.1338,  0.0275,  0.0304, -0.0151],
       device='cuda:0', dtype=torch.bfloat16)
    delta:
           tensor([0.0439, 0.0479, 0.0967, 0.0408, 0.0713, 0.0242, 0.0435, 0.0167],
       device='cuda:0', dtype=torch.bfloat16)
[aiter] -->max abs delta:0.46875, delta details: 75.7% (74446 of 98304) elements
Traceback (most recent call last):
  File "/sgl-workspace/aiter/op_tests/test_mla_persistent.py", line 1800, in <module>
    ret = test_mla(
  File "/sgl-workspace/aiter/aiter/test_common.py", line 129, in wrapper
    ret = func(*args, **kwargs)
  File "/sgl-workspace/aiter/op_tests/test_mla_persistent.py", line 1630, in test_mla
    err, us_asm_decode = test_absorb_decode_fp8()
  File "/sgl-workspace/aiter/op_tests/test_mla_persistent.py", line 1504, in test_absorb_decode_fp8
    cal_diff(out_ref, out_asm, "out", True)
  File "/sgl-workspace/aiter/op_tests/test_mla_persistent.py", line 247, in cal_diff
    assert cos_diff < 3e-2
AssertionError

short length case:

[aiter] import [module_aiter_core] under /sgl-workspace/aiter/aiter/jit/module_aiter_core.so
[aiter] 
calling test_mla(ctx_lens                     = 21,
                 batch_size                   = 1,
                 nhead                        = 64,
                 kv_lora_rank                 = 512,
                 qk_nope_head_dim             = 512,
                 qk_rope_head_dim             = 64,
                 v_head_dim                   = 512,
                 dtype                        = torch.float8_e4m3fnuz,
                 kvtype                       = torch.float8_e4m3fnuz,
                 page_size                    = 64,
                 varlen                       = False,
                 decode_qlen                  = 1,
                 max_split_per_batch          = 32,
                 non_persistent_mode          = False,
                 paged_layout                 = LEGACY,
                 scale_dim                    = 4,
                 return_lse                   = False)
[aiter] import [module_mla_metadata] under /sgl-workspace/aiter/aiter/jit/module_mla_metadata.so
[aiter] LoadKernel: _ZN5aiter39mla_a8w8_qh64_qseqlen1_gqaratio64_v3_psE hsaco: /sgl-workspace/aiter/hsa//gfx942/mla/mla_a8w8_qh64_qseqlen1_gqaratio64_v3_ps.co
[aiter] import [module_mla_reduce] under /sgl-workspace/aiter/aiter/jit/module_mla_reduce.so
[W604 11:00:19.289213803 collection.cpp:1133] Warning: ROCTracer produced duplicate flow start: 3 (function operator())
[aiter] mla_decode-absorb_fp8    [golden vs aiter_asm]:    19.29 us......[checkAllclose atol=0.01 rtol=0.01 failed!]
    a    : torch.Size([1, 64, 512])
           tensor([ 0.5039,  0.0532, -0.3691,  0.1738,  0.0052,  0.1270,  0.4004, -0.2490],
       device='cuda:0', dtype=torch.bfloat16)
    b    : torch.Size([1, 64, 512])
           tensor([ 0.4863,  0.0035, -0.1089,  0.1299, -0.0698,  0.0322,  0.1035, -0.2217],
       device='cuda:0', dtype=torch.bfloat16)
    delta:
           tensor([0.0176, 0.0498, 0.2598, 0.0439, 0.0752, 0.0947, 0.2969, 0.0273],
       device='cuda:0', dtype=torch.bfloat16)
[aiter] -->max abs delta:1.4140625, delta details: 96.6% (31664 of 32768) elements
[aiter] mla_decode-absorb_fp8    [golden fp8 vs aiter_asm]:    19.29 us......[checkAllclose atol=0.01 rtol=0.01 failed!]
    a    : torch.Size([1, 64, 512])
           tensor([-0.7500,  0.0938, -0.4688, -1.5000, -0.2812,  0.6250, -0.0391,  0.3125],
       device='cuda:0', dtype=torch.bfloat16)
    b    : torch.Size([1, 64, 512])
           tensor([ 0.4863,  0.0035, -0.1089,  0.1299, -0.0698,  0.0322,  0.1035, -0.2217],
       device='cuda:0', dtype=torch.bfloat16)
    delta:
           tensor([1.2344, 0.0903, 0.3594, 1.6328, 0.2109, 0.5938, 0.1426, 0.5352],
       device='cuda:0', dtype=torch.bfloat16)
[aiter] -->max abs delta:3.859375, delta details: 99.1% (32480 of 32768) elements
[aiter] mla_decode-absorb_fp8    [golden fp8 split_out_ref vs aiter_asm]:    19.29 us......[checkAllclose atol=0.01 rtol=0.01 failed!]
    a    : torch.Size([1, 64, 512])
           tensor([-0.7500,  0.0938, -0.4688, -1.5000, -0.2812,  0.6250, -0.0391,  0.3125],
       device='cuda:0', dtype=torch.bfloat16)
    b    : torch.Size([1, 64, 512])
           tensor([ 0.4863,  0.0035, -0.1089,  0.1299, -0.0698,  0.0322,  0.1035, -0.2217],
       device='cuda:0', dtype=torch.bfloat16)
    delta:
           tensor([1.2344, 0.0903, 0.3594, 1.6328, 0.2109, 0.5938, 0.1426, 0.5352],
       device='cuda:0', dtype=torch.bfloat16)
[aiter] -->max abs delta:3.859375, delta details: 99.1% (32480 of 32768) elements
Traceback (most recent call last):
  File "/sgl-workspace/aiter/op_tests/test_mla_persistent.py", line 1800, in <module>
    ret = test_mla(
  File "/sgl-workspace/aiter/aiter/test_common.py", line 129, in wrapper
    ret = func(*args, **kwargs)
  File "/sgl-workspace/aiter/op_tests/test_mla_persistent.py", line 1630, in test_mla
    err, us_asm_decode = test_absorb_decode_fp8()
  File "/sgl-workspace/aiter/op_tests/test_mla_persistent.py", line 1504, in test_absorb_decode_fp8
    cal_diff(out_ref, out_asm, "out", True)
  File "/sgl-workspace/aiter/op_tests/test_mla_persistent.py", line 247, in cal_diff
    assert cos_diff < 3e-2
AssertionError

Byeong-Chan · 2026-06-04T11:03:18Z

@JohnNikolay84

Also varlen has same issue.

please, could you check this problem?

JohnNikolay84 · 2026-06-04T11:42:46Z

Also varlen has same issue.

It might be a common problem, but this specific kernel is qseqlen=1 only, there is no varlen.

Byeong-Chan · 2026-06-04T11:54:26Z

@JohnNikolay84

I mean qseqlen=1(decode_qlen), but kvlen(ctx_lens) is varlen in decode.

calling test_mla(ctx_lens                     = 1024,
                 batch_size                   = 3,
                 nhead                        = 64,
                 kv_lora_rank                 = 512,
                 qk_nope_head_dim             = 512,
                 qk_rope_head_dim             = 64,
                 v_head_dim                   = 512,
                 dtype                        = torch.float8_e4m3fnuz,
                 kvtype                       = torch.float8_e4m3fnuz,
                 page_size                    = 64,
                 varlen                       = False,
                 decode_qlen                  = 1,
                 max_split_per_batch          = 32,
                 non_persistent_mode          = False,
                 paged_layout                 = LEGACY,
                 scale_dim                    = 4,
                 return_lse                   = False)
[aiter] mla_decode-absorb_fp8    [golden vs aiter_asm]:    77.23 us......[checkAllclose atol=0.01 rtol=0.01 warning!] a and b results are not all close
[aiter] -->max abs delta:0.013671875, delta details: 0.0% (46 of 98304) elements


calling test_mla(ctx_lens                     = 1024,
                 batch_size                   = 3,
                 nhead                        = 64,
                 kv_lora_rank                 = 512,
                 qk_nope_head_dim             = 512,
                 qk_rope_head_dim             = 64,
                 v_head_dim                   = 512,
                 dtype                        = torch.float8_e4m3fnuz,
                 kvtype                       = torch.float8_e4m3fnuz,
                 page_size                    = 64,
                 varlen                       = True,
                 decode_qlen                  = 1,
                 max_split_per_batch          = 32,
                 non_persistent_mode          = False,
                 paged_layout                 = LEGACY,
                 scale_dim                    = 4,
                 return_lse                   = False)
[aiter] mla_decode-absorb_fp8    [golden vs aiter_asm]:    44.54 us......[checkAllclose atol=0.01 rtol=0.01 failed!]
    a    : torch.Size([3, 64, 512])
           tensor([ 0.1660, -0.1934,  0.0583,  0.0038, -0.0581, -0.0156, -0.0277, -0.0693],
       device='cuda:0', dtype=torch.bfloat16)
    b    : torch.Size([3, 64, 512])
           tensor([ 0.2139, -0.1543,  0.1318, -0.0152, -0.1338,  0.0304, -0.0151, -0.1216],
       device='cuda:0', dtype=torch.bfloat16)
    delta:
           tensor([0.0479, 0.0391, 0.0732, 0.0190, 0.0757, 0.0459, 0.0126, 0.0522],
       device='cuda:0', dtype=torch.bfloat16)
[aiter] -->max abs delta:0.21484375, delta details: 60.4% (59409 of 98304) elements

vstakhov-amd · 2026-06-04T15:03:28Z

I have fixed small context case but I have not tested varlen = True, as I think that this particular kernel does not support this option at all (and I'm also not sure how does it work for 1 token only). I will update .co files shortly. Thanks for testing!

Byeong-Chan · 2026-06-05T07:29:41Z

@vstakhov-amd

But this is very important issue about batching variant context length when decoding phase real workload.

ex)

request1 - prefill done, 128 length kv cache and 1 token decode
request2 - prefill done, 1024 length kv cache and 1 token decode

batching those two case -> precision issue

Rebuild mla_a8w8_qh64_qseqlen1_gqaratio64{_lse,}_v3_ps.co with: - Short paged-tail fix: route partial page64 tails (kv % SUB_KV != 0) through masked tail processing using kv_last_page_lens. Without this the kernel read garbage K/V for positions past the last valid token and emitted NaN/Inf for short ctx_lens (e.g. kv=21 at batch=1). - Phase B reframing for page > SUB_KV: avoid divide-by-zero in int_div_ss by computing the loop count from full SUB_KV passes plus a final partial pass length. Validated on MI300X via op_tests/test_mla_persistent.py: - kv=1024,batch=1,nhead=64,page=64,fp8: PASS (all references) - kv=21,batch=1,nhead=64,page=64,fp8: PASS vs fp8 references, bf16-golden delta within quantization noise floor (err=0).

Rebuild mla_a8w8_qh64_qseqlen1_gqaratio64{_lse,}_v3_ps.co to fix per-batch kv length handling on the varlen decode path. Before: at batch>=2 with varlen=True the kernel produced large systematic errors (~33% of output elements deviated by >1, max delta ~1.6 vs fp8 reference) because per-batch kv_last_page_lens were not applied per work item. After, on MI300X via op_tests/test_mla_persistent.py: kv=1024,batch=3,varlen=True,nhead=64,page=64,fp8: fp8 golden vs aiter_asm: 28/98304 elems off, max delta 0.0146 fp8 split_out_ref: 28/98304 elems off, max delta 0.0146 Regression sweep (kv in {1024,21}, batch in {1,3}, varlen on/off): no regression; all fp8 references pass or warn at quantization noise level.

vstakhov-amd · 2026-06-05T17:40:42Z

@Byeong-Chan I have hopefully fixed the both issues you have mentioned: ctx_lens < page_size and varlen support. Could you please give it another test iteration once you have time? Thanks in advance!

Byeong-Chan · 2026-06-09T06:41:38Z

@vstakhov-amd great! It's ok when page size 64 all test. thanks for your works!

JohnNikolay84 · 2026-06-11T15:50:55Z

@ROCm/team_aiter could someone please review this change. This is a new gfx942 kernel to run GLM-4.5 in context split mode.

# Conflicts: # csrc/kernels/mla/metadata/v1_2_device.cuh

vstakhov-amd requested review from a team and JohnNikolay84 May 14, 2026 16:04

Merge branch 'main' into vstakhov/mla_qh64_fp8_gfx942

9c3c060

vstakhov-amd added the ci:all label May 14, 2026

vstakhov-amd marked this pull request as ready for review May 19, 2026 12:56

JohnNikolay84 requested a review from valarLip May 19, 2026 12:56

JohnNikolay84 reviewed May 19, 2026

View reviewed changes

vstakhov-amd removed the ci:all label May 19, 2026

vstakhov-amd force-pushed the vstakhov/mla_qh64_fp8_gfx942 branch 2 times, most recently from 2d4d526 to ab05216 Compare June 2, 2026 11:24

Merge remote-tracking branch 'origin/main' into vstakhov/mla_qh64_fp8…

0363bea

…_gfx942 # Conflicts: # aiter/mla.py # aiter/ops/attention.py # csrc/kernels/mla/metadata/v1_2_device.cuh # op_tests/test_mla_persistent.py

vstakhov-amd force-pushed the vstakhov/mla_qh64_fp8_gfx942 branch from ab05216 to 0363bea Compare June 2, 2026 11:31

JohnNikolay84 assigned vstakhov-amd Jun 2, 2026

vstakhov-amd added 2 commits June 5, 2026 17:51

Merge branch 'main' into vstakhov/mla_qh64_fp8_gfx942

bd646fb

Merge branch 'main' into vstakhov/mla_qh64_fp8_gfx942

b53f37e

# Conflicts: # csrc/kernels/mla/metadata/v1_2_device.cuh

valarLip approved these changes Jul 10, 2026

View reviewed changes

vstakhov-amd merged commit 04427a5 into main Jul 11, 2026
66 of 67 checks passed

vstakhov-amd deleted the vstakhov/mla_qh64_fp8_gfx942 branch July 11, 2026 10:09

This was referenced Jul 24, 2026

[MLA][gfx942] Native qh64 fp8 persistent decode kernel GPU-faults at page_size=1 (block_size=1) #4363

Open

[Bugfix][MLA] Gate gfx942 native qh64 fp8 decode to page_size=64 (fix page_size=1 GPU fault) #4365

Open

Uh oh!

Conversation

vstakhov-amd commented May 14, 2026

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Uh oh!

github-actions Bot commented May 14, 2026

🏷️ CI Guide

Uh oh!

JohnNikolay84 left a comment

Choose a reason for hiding this comment

Uh oh!

Byeong-Chan commented May 26, 2026

Uh oh!

vstakhov-amd commented May 27, 2026

Uh oh!

vstakhov-amd commented Jun 2, 2026

Uh oh!

Byeong-Chan commented Jun 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

vstakhov-amd commented Jun 4, 2026

Uh oh!

JohnNikolay84 commented Jun 4, 2026

Uh oh!

Byeong-Chan commented Jun 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Byeong-Chan commented Jun 4, 2026

Uh oh!

JohnNikolay84 commented Jun 4, 2026

Uh oh!

Byeong-Chan commented Jun 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

vstakhov-amd commented Jun 4, 2026

Uh oh!

Byeong-Chan commented Jun 5, 2026

Uh oh!

vstakhov-amd commented Jun 5, 2026

Uh oh!

Byeong-Chan commented Jun 9, 2026

Uh oh!

JohnNikolay84 commented Jun 11, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Byeong-Chan commented Jun 4, 2026 •

edited

Loading

Byeong-Chan commented Jun 4, 2026 •

edited

Loading

Byeong-Chan commented Jun 4, 2026 •

edited

Loading