[MLA ASM kernel] add v4 MLA decode for mi350 (recompile mi300)

The v4 new-model kernel uses an 18-slot kernarg ABI that's incompatible
with v3's 14-slot layout:

- slot 8  raw `gqa_ratio` (not `s_MQA = gqa_ratio * max_seqlen_q`)
- slot 10 `total_kv = kv_seq_lens * num_seqs`
- 4 new tail slots: ptr_STP, out_16_nosplit, ptr_QROPE, ptr_KVROPE
- scalar = 1/sqrt(kV4DimNope + kV4DimRope) = 1/sqrt(512), independent of
  head_size
- kernel hardcodes wave64 (gfx950) block dim = wv_tg * 64 = 256

Integration follows the canonical aiter JIT module pattern (peer of
asm_mla.cu / module_mla_asm, not the compile_template_op jinja path):

- csrc/py_itfs_cu/asm_mla_v4.cu: C-ABI dispatcher using aiter_tensor_t*,
  wrapped in AITER_CTYPES_DEFINE_ENTRYPOINT_VOID so AITER_CHECK failures
  surface as clean Python RuntimeError via the aiter_get_last_error TLS
  bridge (no std::abort()).
- hsa/gfx950/mla_v4/mla_v4_asm.csv: CSV registry consumed by
  hsa/codegen.py -m mla_v4 -> asm_mla_v4_configs.hpp / cfg_mla_v4_asm.
  Schema (qType, kvType, gqa, sub_Q, page_size, num_kv_splits, prefill,
  causal, knl_name, co_name) is independent of mla_asm.csv so v3 / v4
  dispatchers don't collide.
- hsa/gfx950/mla_v4/mla_a8w8_qh64_qseqlen4_gqaratio16_nm_recmp.co:
  compiled shader. Built from poc_kl/mi350/mla_asm/shaders/mla_a8w8_qh64
  _1tg_16mx4_16nx1_nm_recompile.s with the literal `mla_kernel_func`
  symbol sed-replaced to `_ZN5aiter42mla_a8w8_qh64_qseqlen4_gqaratio16
  _nm_recmpE` then reassembled with `clang -x assembler
  --offload-arch=gfx950`.
- aiter/jit/optCompilerConfig.json: module_mla_v4_asm entry with
  blob_gen_cmd = codegen.py -m mla_v4.
- aiter/ops/attention.py: mla_decode_v4_asm Python signature
  (@compile_ops ffi_type="ctypes").
- aiter/mla.py: thin wrapper `mla_decode_fwd_v4_nm` routes through
  aiter.mla_decode_v4_asm and allocates the canonical 5D
  [num_seqs, num_kv_splits, num_kv_heads, gqa*max_seqlen_q, v_head_dim]
  split-logits / attn_lse internally.
- aiter/utility/dtypes.py: _torch_to_aiter_dtype now accepts both
  torch.float8_e4m3fn (OCP) and torch.float8_e4m3fnuz (NUZ); both map
  to AITER_DTYPE_fp8, so callers on gfx942 / gfx950 can use either
  variant transparently.

Tests (op_tests/test_mla_v4_nm.py):

- test_v4_nm_smoke_default_shape    : dispatcher -> kernel launch ->
                                       SENTINEL coverage check.
- test_v4_nm_no_half_zero_pattern   : wave64-launch regression guard
                                       (the historic 256-NaN + 256-zero
                                       output pattern on gfx950).
- test_v4_nm_determinism            : two back-to-back launches are
                                       bit-identical.
- test_v4_nm_out_16_nosplit_arg_accepted
                                    : out_16_nosplit=1 path doesn't
                                       crash the dispatcher.
- test_v4_nm_unknown_variant_raises : sub_Q=128 (not in cfg_mla_v4_asm)
                                       triggers a clean RuntimeError via
                                       the C-ABI error bridge, no
                                       silently-loaded wrong .co.
- test_v4_nm_kernarg_scalar_slots   : env-gated kernarg hexdump
                                       captured via capfd; scalar slots
                                       7-12 + 15 byte-exact match
                                       poc_kl's MlaV4HipKernelArgs.

Numerical correctness (op_tests/test_mla_v4_nm_golden.py):

- Byte-exact compare against poc_kl gpu_SPLIT_DATA.hex. Inputs loaded
  via raw hipMemcpy H2D from device-side `.bin` dumps produced by an
  env-gated patch on poc_kl side
  (POC_KL_V4_DUMP_DEVBINS=1 in mla_execute_v4_hip.inl). Bypasses the
  brittle path of reconstructing inputs from hex dumps; sidesteps the
  v4 nm kernel's degenerate code path on uniform e8m0 scale bytes
  (which would produce 100% NaN for synthetic all-1.0 inputs).
- CI-safe: libamdhip64.so is lazy-loaded inside the test body, and the
  test SKIPs when POC_KL_DUMP_DIR / devbins / gpu_SPLIT_DATA.hex are
  missing or when the GPU is not gfx950.

Variant currently shipped:
  gqa_ratio=16, sub_Q=64, page_size=1, q_dtype=fp8, kv_dtype=fp8

Validated on gfx950 (MI355X):
- 6 baseline pytest cases pass.
- Byte-exact golden test passes (0/65536 differ vs poc_kl).

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: liyjiang

aiter/jit/optCompilerConfig.json

@@ -138,6 +138,17 @@
         "verbose": "False",
         "blob_gen_cmd": "f'{AITER_META_DIR}/hsa/codegen.py -m mla --output_dir {{}}'"
     },
+    "module_mla_v4_asm": {
+        "srcs": [
+            "f'{AITER_CSRC_DIR}/py_itfs_cu/asm_mla_v4.cu'"
+        ],
+        "flags_extra_cc": [],
+        "flags_extra_hip": [],
+        "extra_ldflags": "None",
+        "extra_include": [],
+        "verbose": "False",
+        "blob_gen_cmd": "f'{AITER_META_DIR}/hsa/codegen.py -m mla_v4 --output_dir {{}}'"
+    },
     "module_cache": {
         "srcs": [
             "f'{AITER_CSRC_DIR}/pybind/cache_pybind.cu'",

aiter/mla.py

@@ -973,3 +973,74 @@ def mla_prefill_reduce(
             )  # [tile_q, v_head_dim]
 
             output[qo_start:qo_end, head_idx, :] = final_output[:q_len, :]
+
+
+# ---------------------------------------------------------------------------
+# DSv4 MLA — additive entry point. Does NOT touch any existing
+#   gqa_ratio=16, sub_Q=64, page_size=1, q_dtype=fp8, kv_dtype=fp8
+# ---------------------------------------------------------------------------
+def mla_decode_fwd_v4_nm(
+    q,                    # [total_query_len, num_heads, head_size]   FP8 packed Q+e8m0
+    qrope,                # [total_query_len, num_heads, kv_rotary]   BF16
+    kv_buffer,            # [num_page, page_size, num_kv_heads, dim_qk_packed]
+    kvrope,               # [num_page, page_size, num_kv_heads, kv_rotary]
+    output,               # [total_query_len, num_heads, v_head_dim]  BF16 (used for out_16_nosplit=1)
+    qo_indptr,            # [num_seqs+1]
+    kv_indptr,            # [num_seqs+1]
+    kv_page_indices,      # [num_page_used]
+    kv_last_page_lens,    # [num_seqs]
+    split_indptr,         # [num_seqs+1]
+    max_seqlen_q,
+    sm_scale=None,        # ignored on v4 nm; kernel hardcodes 1/sqrt(512)
+    out_16_nosplit=0,
+    num_kv_splits=1,
+    sub_Q=64,
+    logits=None,
+    attn_lse=None,
+):
+    """v4 nm-recompile MLA decode forward (mi350 / gfx950 wave64).
+
+    Routes through the canonical aiter JIT C-ABI module
+    `module_mla_v4_asm` (csrc/py_itfs_cu/asm_mla_v4.cu). Returns
+    (logits, attn_lse) — caller is responsible for any post-reduce /
+    final-O work.
+
+    logits/attn_lse may be pre-allocated (e.g. for SENTINEL pre-fill in
+    correctness tests); if not, we allocate the canonical 5D layout
+    `[num_seqs, num_kv_splits, num_kv_heads, gqa*max_seqlen_q, v_head_dim]`
+    (and `[..., 1]` for attn_lse) inferred from the input shapes.
+    """
+    num_seqs = qo_indptr.shape[0] - 1
+    num_heads = q.size(1)
+    v_head_dim = output.size(2)
+    num_kv_heads = kv_buffer.size(2)
+    gqa_ratio = num_heads // num_kv_heads
+    q_seq_lens_internal = gqa_ratio * max_seqlen_q
+
+    if logits is None:
+        logits = torch.empty(
+            (num_seqs, num_kv_splits, num_kv_heads, q_seq_lens_internal, v_head_dim),
+            dtype=dtypes.fp32, device=q.device,
+        )
+    if attn_lse is None:
+        attn_lse = torch.empty(
+            (num_seqs, num_kv_splits, num_kv_heads, q_seq_lens_internal, 1),
+            dtype=dtypes.fp32, device=q.device,
+        )
+
+    # softmax_scale is ignored by the v4 nm kernel (hardcodes 1/sqrt(512));
+    # we still pass *something* through to satisfy the C ABI.
+    sm_scale_arg = 0.0 if sm_scale is None else float(sm_scale)
+
+    aiter.mla_decode_v4_asm(
+        q, qrope, kv_buffer, kvrope,
+        qo_indptr, kv_indptr, kv_page_indices, kv_last_page_lens,
+        split_indptr,
+        max_seqlen_q,
+        sm_scale_arg,
+        int(out_16_nosplit),
+        int(sub_Q),
+        int(num_kv_splits),
+        logits, attn_lse, output,
+    )
+    return logits, attn_lse

aiter/ops/attention.py

@@ -688,6 +688,48 @@ def mla_decode_stage1_asm_fwd(
 ) -> None: ...
 
 
+MD_NAME_V4 = "module_mla_v4_asm"
+
+
+@compile_ops(MD_NAME_V4, ffi_type="ctypes")
+def mla_decode_v4_asm(
+    # [total_query_len, num_heads, head_size]   FP8 packed Q + e8m0 scale region
+    Q: torch.Tensor,
+    # [total_query_len, num_heads, kv_rotary]   BF16
+    qrope: torch.Tensor,
+    # [num_page, page_size, num_kv_heads, head_size]  FP8
+    KV: torch.Tensor,
+    # [num_page, page_size, num_kv_heads, kv_rotary]  BF16
+    kvrope: torch.Tensor,
+    # [num_seqs+1]
+    qo_indptr: torch.Tensor,
+    # [num_seqs+1]
+    kv_indptr: torch.Tensor,
+    # [num_page_used]
+    kv_page_indices: torch.Tensor,
+    # [num_seqs]
+    kv_last_page_lens: torch.Tensor,
+    # [num_seqs+1]
+    split_indptr: torch.Tensor,
+    max_seqlen_q: int,
+    # ignored on v4 nm; kernel hardcodes 1/sqrt(kV4DimNope+kV4DimRope)=1/sqrt(512)
+    softmax_scale: float,
+    # 0 = fp32 split-out path; 1 = bf16 nosplit reduce path
+    out_16_nosplit: int,
+    # poc_kl `sub_Q` (= per-WG Q tile); only 64 currently shipped
+    sub_Q: int,
+    # poc_kl `passes`
+    num_kv_splits: int,
+    # outputs
+    # [num_seqs, num_kv_splits, num_kv_heads, gqa*max_seqlen_q, v_head_dim] FP32
+    splitData: torch.Tensor,
+    # [num_seqs, num_kv_splits, num_kv_heads, gqa*max_seqlen_q, 1]          FP32
+    splitLse: torch.Tensor,
+    # [total_query_len, num_heads, v_head_dim] BF16 (used when out_16_nosplit==1)
+    output: torch.Tensor,
+) -> None: ...
+
+
 @compile_ops(MD_NAME, ffi_type="ctypes")
 def mla_prefill_asm_fwd(
     # [num_seqs, num_heads, head_size]

aiter/utility/dtypes.py

@@ -39,6 +39,15 @@ def get_dtype_fp8():
 globals().update({f"AITER_DTYPE_{name}": idx for name, idx in aiter_dtypes.items()})
 _torch_to_aiter_dtype = {globals()[name]: idx for name, idx in aiter_dtypes.items()}
 
+# Both e4m3fn (OCP) and e4m3fnuz (ROCm NUZ) are valid FP8 variants at the
+# byte level for kernels that just read raw FP8 bytes. Map both torch dtypes
+# to the same AITER_DTYPE_fp8 enum so the strict dtype check in
+# torch_to_aiter() / torch_to_aiter_pybind() accepts whichever variant the
+# caller has — letting v3/v4 MLA tests use either dtype interchangeably.
+if "fp8" in aiter_dtypes:
+    for _alt_fp8 in (torch.float8_e4m3fn, torch.float8_e4m3fnuz):
+        _torch_to_aiter_dtype.setdefault(_alt_fp8, aiter_dtypes["fp8"])
+
 
 def torch_to_aiter_pybind(tensor: torch.Tensor):
     """Convert torch.Tensor to pybind aiter_tensor_t for passing to C++ ops.