[Bug] MLA persistent decode: bf16 query + fp8 KV aborts for gqa_ratio=16, max_seqlen_q > 4

[lizexu123]

Problem Description

Running Kimi-K3 with DSpark MTP speculative decoding (TP8, gqa_ratio=16, bf16 query +
fp8 KV cache) on gfx950, the server aborts during CUDA graph capture with:

get_heuristic_kernel_mla: cannot get heuristic kernel!
  q_type:bf16 kv_type:fp8 gqa:16 ps:1 prefill:0 causal:0 qseqlen:5 lse:0 cprr:0

This happens once num_speculative_tokens=2 pushes the MLA decode verify length
(max_seqlen_q = 1 + 2 * num_speculative_tokens) to 5. num_speculative_tokens=1
(qlen=3) works fine, so the failure is specifically tied to max_seqlen_q > 4 for this
dtype combination — see the attached minimal reproducer, which isolates it from
Kimi-K3/vLLM entirely.

Operating System

Ubuntu 22.04.5 LTS (Jammy Jellyfish), kernel 5.15.0-70-generic

CPU

AMD EPYC 9965 192-Core Processor

GPU

8 x AMD Instinct MI350X (gfx950, device id 0x75a0)

ROCm Version

ROCm 7.2.3

ROCm Component

No response

Steps to Reproduce

This matches the production abort seen running Kimi-K3 TP8 + DSpark MTP with
num_speculative_tokens=2, which pushes the MLA decode verify length
(max_seqlen_q = 1 + 2 * num_speculative_tokens) to 5. The script hardcodes the exact
failing shape: gqa_ratio=16, persistent decode, bf16 query, fp8 KV cache,
causal=False, max_seqlen_q=5.

#!/usr/bin/env python3
"""Minimal repro for the Kimi-K3 crash:

 get_heuristic_kernel_mla: cannot get heuristic kernel!
   q_type:bf16 kv_type:fp8 gqa:16 ps:1 prefill:0 causal:0 qseqlen:5 lse:0 cprr:0

aiter has no persistent MLA decode kernel for bf16 query + fp8 KV cache,
gqa_ratio=16, when max_seqlen_q > 4 (asm_mla.cu only folds the fp8-query
branch to gqa_ratio=32/qseqlen=4 for max_seqlen_q > 4; the bf16-query
branch has no equivalent fold, so it falls through to the kernel lookup
and aborts). This is what DSpark/MTP hits once num_speculative_tokens=2
pushes the verify qlen to 5.

Usage:
 python3 repro_aiter_mla_bf16q_fp8kv.py
"""

import torch

from aiter import dtypes, get_mla_metadata_info_v1, get_mla_metadata_v1
from aiter.jit.utils.chip_info import get_gfx
from aiter.mla import mla_decode_fwd

# Kimi-K3 MLA geometry. At TP8 the model has 12 heads/rank; callers (vLLM,
# SGLang) pad up to 16, so gqa_ratio=16 is what reaches the kernel.
NHEAD = 16
KV_LORA_RANK = 512
QK_ROPE_HEAD_DIM = 64
QK_HEAD_DIM = KV_LORA_RANK + QK_ROPE_HEAD_DIM  # 576
V_HEAD_DIM = KV_LORA_RANK  # 512
PAGE_SIZE = 1
BS = 2
KV_LEN = 1024
MAX_SEQLEN_Q = 5  # DSpark verify qlen = 1 + 2 * num_speculative_tokens(=2)

torch.set_default_device("cuda")

q_dtype = dtypes.bf16
kv_dtype = dtypes.fp8
print(f"gfx={get_gfx()}  q={q_dtype}  kv={kv_dtype}  "
   f"nhead={NHEAD}  max_seqlen_q={MAX_SEQLEN_Q}  persistent=True")

total_s = BS * MAX_SEQLEN_Q
total_kv = BS * KV_LEN

q = torch.randn(total_s, NHEAD, QK_HEAD_DIM, dtype=torch.bfloat16).to(q_dtype)
kv_buffer = torch.randn(
 total_kv, PAGE_SIZE, 1, QK_HEAD_DIM, dtype=torch.bfloat16
).to(kv_dtype)
o = torch.empty(total_s, NHEAD, V_HEAD_DIM, dtype=torch.bfloat16)

qo_indptr = torch.arange(0, (BS + 1) * MAX_SEQLEN_Q, MAX_SEQLEN_Q, dtype=torch.int32)
kv_indptr = torch.arange(0, (BS + 1) * KV_LEN, KV_LEN, dtype=torch.int32)
kv_indices = torch.arange(total_kv, dtype=torch.int32)
kv_last_page_lens = torch.ones(BS, dtype=torch.int32)

# Persistent-mode work metadata, sized by aiter's own helper.
sizes = get_mla_metadata_info_v1(
 BS, MAX_SEQLEN_Q, NHEAD, q_dtype, kv_dtype, is_sparse=False, fast_mode=True,
)
(work_meta_data, work_indptr, work_info_set,
reduce_indptr, reduce_final_map, reduce_partial_map) = [
 torch.empty(size, dtype=dtype) for size, dtype in sizes
]

get_mla_metadata_v1(
 qo_indptr, kv_indptr, kv_last_page_lens, NHEAD, 1, True,
 work_meta_data, work_info_set, work_indptr,
 reduce_indptr, reduce_final_map, reduce_partial_map,
 page_size=PAGE_SIZE, kv_granularity=16,
 max_seqlen_qo=MAX_SEQLEN_Q, uni_seqlen_qo=MAX_SEQLEN_Q,
 fast_mode=True, dtype_q=q_dtype, dtype_kv=kv_dtype,
)

scale = torch.ones(1, dtype=torch.float32)
mla_decode_fwd(
 q, kv_buffer, o, qo_indptr, kv_indptr, kv_indices, kv_last_page_lens,
 MAX_SEQLEN_Q, page_size=PAGE_SIZE,
 work_meta_data=work_meta_data, work_indptr=work_indptr,
 work_info_set=work_info_set, reduce_indptr=reduce_indptr,
 reduce_final_map=reduce_final_map, reduce_partial_map=reduce_partial_map,
 q_scale=None, kv_scale=scale,
 causal=False,
)
torch.cuda.synchronize()
print("OK (unexpected - this should have aborted before reaching here)")

(Optional for Linux users) Output of /opt/rocm/bin/rocminfo --support

rocminfo --support output

Paste output here

Additional Information

No response