[Triton/Gluon] [ASM] [HIP] add mla v4 prefill asm kernel

.github/scripts/split_tests.sh

@@ -103,7 +103,7 @@ if [[ "$TEST_TYPE" == "aiter" ]]; then
     FILE_TIMES[op_tests/test_batched_gemm_a8w8.py]=51
     FILE_TIMES[op_tests/test_mha_varlen_large_kv.py]=46
     FILE_TIMES[op_tests/test_mla_reduce.py]=44
-    FILE_TIMES[op_tests/test_pa_sparse_prefill_opus.py]=44
+    FILE_TIMES[op_tests/test_pa_sparse_prefill.py]=44
     FILE_TIMES[op_tests/test_pa_ragged.py]=40
     FILE_TIMES[op_tests/test_moeTopkSoftmax.py]=39
     FILE_TIMES[op_tests/test_moe_sorting_mxfp4.py]=39

aiter/__init__.py

@@ -112,6 +112,7 @@ def getLogger():
     from .ops.moe_sorting import *
     from .ops.moe_sorting_opus import *
     from .ops.moe_mxfp4_aux import *
+    from .ops.mla_sparse_prefill import *
     from .ops.pa_sparse_prefill_opus import *
     from .ops.pos_encoding import *
     from .ops.cache import *

aiter/jit/optCompilerConfig.json

@@ -103,6 +103,17 @@
         "verbose": "False",
         "blob_gen_cmd": "f'{AITER_META_DIR}/hsa/codegen.py -m mla_v4 --output_dir {{}}'"
     },
+    "module_mla_sparse_prefill_asm": {
+        "srcs": [
+            "f'{AITER_CSRC_DIR}/py_itfs_cu/asm_mla_sparse_prefill.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/ops/mla_sparse_prefill.py

@@ -0,0 +1,143 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+
+import torch
+
+from ..jit.core import compile_ops
+from ..jit.utils.chip_info import get_gfx_runtime
+from ..jit.utils.torch_guard import torch_compile_guard
+
+MD_NAME = "module_mla_sparse_prefill_asm"
+
+
+# NOTE: ctypes binds positionally off this signature -- the argument order here
+# must match `mla_sparse_prefill_fp8_asm_fwd` in
+# csrc/py_itfs_cu/asm_mla_sparse_prefill.cu one for one. A silent reorder here
+# is a silent wrong-pointer launch, not a compile error.
+@compile_ops(MD_NAME, fc_name="mla_sparse_prefill_fp8_asm_fwd", ffi_type="ctypes")
+def mla_sparse_prefill_fp8_asm_fwd(
+    q_nope: torch.Tensor,
+    q_rope: torch.Tensor,
+    unified_kv_nope: torch.Tensor,
+    unified_kv_rope: torch.Tensor,
+    kv_indices_prefix: torch.Tensor,
+    kv_indptr_prefix: torch.Tensor,
+    kv_nope: torch.Tensor,
+    kv_rope: torch.Tensor,
+    kv_indices_extend: torch.Tensor,
+    kv_indptr_extend: torch.Tensor,
+    attn_sink: torch.Tensor,
+    out: torch.Tensor,
+    softmax_scale: float,
+) -> None: ...
+
+
+def _mla_sparse_prefill_fp8_asm_fake(
+    q_nope: torch.Tensor,
+    q_rope: torch.Tensor,
+    unified_kv_nope: torch.Tensor,
+    unified_kv_rope: torch.Tensor,
+    kv_indices_prefix: torch.Tensor,
+    kv_indptr_prefix: torch.Tensor,
+    kv_nope: torch.Tensor,
+    kv_rope: torch.Tensor,
+    kv_indices_extend: torch.Tensor,
+    kv_indptr_extend: torch.Tensor,
+    attn_sink: torch.Tensor,
+    softmax_scale: float,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is not None:
+        return out
+    t, h, _ = q_nope.shape
+    return torch.empty((t, h, 512), dtype=torch.bfloat16, device=q_nope.device)
+
+
+@torch_compile_guard(mutates_args=["out"], gen_fake=_mla_sparse_prefill_fp8_asm_fake)
+def mla_sparse_prefill_fp8_asm(
+    q_nope: torch.Tensor,
+    q_rope: torch.Tensor,
+    unified_kv_nope: torch.Tensor,
+    unified_kv_rope: torch.Tensor,
+    kv_indices_prefix: torch.Tensor,
+    kv_indptr_prefix: torch.Tensor,
+    kv_nope: torch.Tensor,
+    kv_rope: torch.Tensor,
+    kv_indices_extend: torch.Tensor,
+    kv_indptr_extend: torch.Tensor,
+    attn_sink: torch.Tensor,
+    softmax_scale: float,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Sparse prefill attention with split fp8 NoPE and bf16 RoPE inputs (asm).
+
+    Signature-compatible with
+    :func:`aiter.ops.pa_sparse_prefill_opus.pa_sparse_prefill_fp8_opus`.
+
+    Args:
+      q_nope:            ``[T, H, 512]`` fp8 query without positional encoding.
+      q_rope:            ``[T, H, 64]`` bf16 query RoPE encoding part.
+      unified_kv_nope:   ``[total_pages, 512]`` fp8 prefix KV NoPE source.
+      unified_kv_rope:   ``[total_pages, 64]`` bf16 prefix KV RoPE source.
+      kv_indices_prefix: ``[total_prefix]`` int32 row indices into the prefix
+        sources, concatenated per token.
+      kv_indptr_prefix:  ``[T+1]`` int32 CSR row pointers.
+      kv_nope:           ``[total_tokens, 512]`` fp8 extend KV NoPE source.
+      kv_rope:           ``[total_tokens, 64]`` bf16 extend KV RoPE source.
+      kv_indices_extend: ``[total_extend]`` int32 row indices into the extend
+        sources, concatenated per token.
+      kv_indptr_extend:  ``[T+1]`` int32 CSR row pointers.
+      attn_sink:         ``[H]`` fp32 per-head softmax-denom bias.
+      softmax_scale:     float scalar applied to the combined QK^T scores.
+      out:               Optional ``[T, H, 512]`` bf16 output buffer; allocated
+        if ``None``.
+
+    Returns:
+      ``out`` (``[T, H, 512]`` bf16).
+    """
+    gfx = get_gfx_runtime()
+    if gfx != "gfx1250":
+        raise RuntimeError(f"mla_sparse_prefill_fp8_asm requires gfx1250, got {gfx}")
+
+    if q_nope.dtype != unified_kv_nope.dtype or q_nope.dtype != kv_nope.dtype:
+        raise RuntimeError(
+            f"NoPE dtype mismatch: q_nope={q_nope.dtype}, "
+            f"unified_kv_nope={unified_kv_nope.dtype}, kv_nope={kv_nope.dtype}"
+        )
+    if q_rope.dtype != torch.bfloat16:
+        raise RuntimeError(f"q_rope must be bf16, got {q_rope.dtype}")
+
+    t, h = q_nope.shape[0], q_nope.shape[1]
+    if h != 128:
+        # Hard constraint, not a dispatch miss.
+        raise RuntimeError(f"mla_sparse_prefill_fp8_asm requires H == 128, got {h}")
+    if out is None:
+        out = torch.empty((t, h, 512), dtype=torch.bfloat16, device=q_nope.device)
+    elif out.shape != (t, h, 512) or out.dtype != torch.bfloat16:
+        raise RuntimeError(
+            f"out shape/dtype mismatch: got shape={tuple(out.shape)} dtype={out.dtype}, "
+            f"expected shape={(t, h, 512)} dtype={torch.bfloat16}"
+        )
+
+    mla_sparse_prefill_fp8_asm_fwd(
+        q_nope,
+        q_rope,
+        unified_kv_nope,
+        unified_kv_rope,
+        kv_indices_prefix,
+        kv_indptr_prefix,
+        kv_nope,
+        kv_rope,
+        kv_indices_extend,
+        kv_indptr_extend,
+        attn_sink,
+        out,
+        float(softmax_scale),
+    )
+    return out
+
+
+__all__ = [
+    "mla_sparse_prefill_fp8_asm",
+    "mla_sparse_prefill_fp8_asm_fwd",
+]

aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4.py

@@ -291,7 +291,12 @@ def _sparse_attn_prefill_kernel(
     BLOCK_D: tl.constexpr,
     BLOCK_K: tl.constexpr,
 ):
-    query_idx = tl.program_id(0)
+    # 64-bit before the multiply, same reasoning as `slot_off` below: the
+    # program id fits 32 bits, but `query_idx * q_stride_t` does not once
+    # num_queries passes 32K, because q_stride_t is num_heads * head_dim
+    # (128 * 512 = 65536) in the V4 layout. Unlike the pool read, the wrapped
+    # offset lands outside the q/out allocations, so it page-faults.
+    query_idx = tl.program_id(0).to(tl.int64)
     pid_h = tl.program_id(1)
 
     head_offsets = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)