[ASM] [HIP] [JIT] MI355 MLA PS mode kernel add 4 wave fp8 kernel

aiter/jit/optCompilerConfig.json

@@ -2065,7 +2065,8 @@
         "srcs": [
             "f'{AITER_CSRC_DIR}/kernels/mla/opus_mla_decode_fwd.cu'",
             "f'{AITER_CSRC_DIR}/kernels/mla/opus/mla_decode_mxfp8_16mx8_32nx1.hpp'",
-            "f'{AITER_CSRC_DIR}/kernels/mla/opus/mla_decode_fp8_16mx8_32nx1.hpp'",
+            "f'{AITER_CSRC_DIR}/kernels/mla/opus/mla_decode_fp8_16mx1_32nx4.hpp'",
+            "f'{AITER_CSRC_DIR}/kernels/mla/opus/mla_global_load.hpp'",
             "f'{AITER_CSRC_DIR}/kernels/mla/opus/mla_decode_kargs.h'",
             "f'{AITER_CSRC_DIR}/kernels/mla/opus/mla_decode_traits.h'"
         ],

aiter/mla.py

@@ -1031,17 +1031,33 @@ def mla_decode_fwd(
         # Opt-in opus merged-buffer fp8 path (gfx950). Requires a single
         # merged d=576 fp8 q/kv buffer and per-tensor scalar float q/kv scales,
         # or a bf16 q/kv buffer of the same shape.
+        # Shapes with an opus fp8 build: what asm runs on its qh32 / qseqlen4 kernel goes to
+        # 32mx4, what it runs on qh16 / qseqlen4 to 32mx2, and nhead 16 x 1 to 32nx4 -- the
+        # table in opus_fp8_route (csrc/kernels/mla/opus_mla_decode_fwd.cu). Anything else
+        # stays on the asm path. nhead is the folded one here.
+        opus_fp8_shape = (
+            (nhead == 32 and max_seqlen_q >= 2)
+            or nhead == 64
+            or nhead == 128
+            or (nhead == 96 and max_seqlen_q <= 6)
+            or (nhead == 16 and max_seqlen_q in (1, 3, 4))
+        )
         opus_is_fp8 = (
             q.dtype == dtypes.fp8
             and kv_buffer.dtype == dtypes.fp8
+            and opus_fp8_shape
+            # No opus fp8 build addresses a real block table: none has been ported past one
+            # token per page.
+            and page_size == 1
             and q_scale is not None
             and kv_scale is not None
         )
-        opus_is_bf16 = q.dtype == dtypes.bf16 and kv_buffer.dtype == dtypes.bf16
+        opus_is_bf16 = (
+            q.dtype == dtypes.bf16 and kv_buffer.dtype == dtypes.bf16 and page_size == 1
+        )
         use_opus = (
             os.environ.get("AITER_MLA_USE_OPUS", "0") == "1"
             and get_gfx() == "gfx950"
-            and page_size == 1
             and (opus_is_fp8 or opus_is_bf16)
         )
 
@@ -1099,19 +1115,17 @@ def mla_decode_fwd(
                 cp_world_size=cp_world_size,
                 cp_rank=cp_rank,
             )
-        elif use_opus and opus_is_fp8:
-            aiter.opus_mla_decode_fp8_fwd(
+        elif use_opus:
+            aiter.opus_mla_decode_fwd(
                 q,
                 kv_buffer,
                 qo_indptr,
                 kv_indptr,
                 kv_indices,
-                kv_last_page_lens,
                 work_indptr,
                 work_info_set,
                 max_seqlen_q,
                 page_size,
-                nhead_kv,
                 sm_scale,
                 logits,
                 attn_lse,
@@ -1121,24 +1135,8 @@ def mla_decode_fwd(
                 kv_scale,
                 causal,
             )
-        elif use_opus:
-            aiter.opus_mla_decode_fwd(
-                q,
-                kv_buffer,
-                qo_indptr,
-                kv_indptr,
-                kv_indices,
-                work_indptr,
-                work_info_set,
-                page_size,
-                sm_scale,
-                logits,
-                attn_lse,
-                o,
-                final_lse,
-            )
         elif use_hk:
-            aiter.hk_mla_decode_fwd(
+            aiter.hk_mla_v32_decode_fwd(
                 q,
                 kv_buffer,
                 qo_indptr,

aiter/ops/attention.py

@@ -1970,41 +1970,21 @@ def opus_mla_decode_mxfp8_fwd(
 
 @compile_ops("module_mla_decode_opus", ffi_type="ctypes")
 def opus_mla_decode_fwd(
-    q: torch.Tensor,  # [total_q, H, 576] bf16
-    kv: torch.Tensor,  # [num_page, 1, 1, 576] bf16, page_size == 1
+    q: torch.Tensor,  # [total_q, H, 576] bf16 or fp8 (merged nope+rope)
+    kv: torch.Tensor,  # [total_tokens, ..., 576] same dtype as q, page_size == 1
     qo_indptr: torch.Tensor,
     kv_indptr: torch.Tensor,
     kv_indices: torch.Tensor,
     work_indptr: torch.Tensor,
     work_info_set: torch.Tensor,
+    max_seqlen_q: int,
     page_size: int,
     softmax_scale: float,
     logits: torch.Tensor,  # aiter split_output [num_partials,1,H,512] fp32
     attn_lse: torch.Tensor,  # aiter split_lse    [num_partials,1,H,1]   fp32
     out: torch.Tensor,  # final [total_q, H, 512] bf16
     final_lse: torch.Tensor | None = None,  # [total_q, H] fp32
-) -> None: ...
-
-
-@compile_ops("module_mla_decode_opus", ffi_type="ctypes")
-def opus_mla_decode_fp8_fwd(
-    q: torch.Tensor,  # [B, H, 576]           fp8 (merged nope+rope)
-    kv: torch.Tensor,  # [total_tokens, 576]   fp8 (merged nope+rope)
-    qo_indptr: torch.Tensor,
-    kv_indptr: torch.Tensor,
-    kv_indices: torch.Tensor,
-    kv_last_page_lens: torch.Tensor,
-    work_indptr: torch.Tensor,
-    work_info_set: torch.Tensor,
-    max_seqlen_q: int,
-    page_size: int,
-    nhead_kv: int,
-    softmax_scale: float,
-    logits: torch.Tensor,  # aiter split_output [num_partials,1,H,512] fp32
-    attn_lse: torch.Tensor,  # aiter split_lse    [num_partials,1,H,1]   fp32
-    out: torch.Tensor,  # final [B, H, 512] bf16
-    final_lse: torch.Tensor | None = None,  # [B, H] fp32
-    q_scale: torch.Tensor | None = None,  # float[1] per-tensor descale
-    kv_scale: torch.Tensor | None = None,  # float[1] per-tensor descale
-    causal: bool = True,  # mask across the max_seqlen_q query tokens
+    q_scale: torch.Tensor | None = None,  # fp8: float[1] per-tensor descale
+    kv_scale: torch.Tensor | None = None,  # fp8: float[1] per-tensor descale
+    causal: bool = True,  # fp8: mask across the max_seqlen_q query tokens
 ) -> None: ...

csrc/kernels/mla/opus/build/build_mla_decode_fp8_32mx2_32nx2_co.sh

@@ -0,0 +1,31 @@
+#!/bin/bash
+# Builds hsa/gfx950/mla_opus/opus_mla_decode_fp8_32mx2_32nx2.co with the amdgpu-pin-op-dst
+# toolchain (git@github.com:yuyzhang512/llvm-project.git, branch amdgpu-pin-op-dst), driven
+# through hipcc so the ROCm device libs and --genco still apply.
+#
+#   PIN_LLVM=/path/to/llvm-project/build/bin ./build/build_mla_decode_fp8_32mx2_32nx2_co.sh
+#
+# The flags are the code generation aiter's JIT gives its in-tree kernels; -Rpass-analysis
+# prints each entry's register budget.
+set -euo pipefail
+
+HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+OPUS="$(cd "$HERE/.." && pwd)"
+AITER="$(cd "$OPUS/../../../.." && pwd)"
+PIN_LLVM="${PIN_LLVM:-$HOME/llvm-project/build/bin}"
+OUT="${OUT:-$AITER/hsa/gfx950/mla_opus/opus_mla_decode_fp8_32mx2_32nx2.co}"
+
+[ -x "$PIN_LLVM/clang++" ] || { echo "no pin toolchain at $PIN_LLVM" >&2; exit 1; }
+
+HIP_CLANG_PATH="$PIN_LLVM" /opt/rocm/bin/hipcc -x hip "$OPUS/mla_decode_fp8_32mx2_32nx2_co.hip" \
+    -I"$AITER/csrc/include" -I"$AITER/csrc/kernels/mla" \
+    -std=c++20 -O3 --offload-arch=gfx950 -D__HIPCC_RTC__ \
+    -ffast-math -fgpu-flush-denormals-to-zero -fno-offload-uniform-block \
+    -mllvm --amdgpu-kernarg-preload-count=32 -mllvm --lsr-drop-solution=1 \
+    -mllvm -amdgpu-early-inline-all=true -mllvm -amdgpu-function-calls=false \
+    -mllvm -enable-post-misched=0 \
+    -Rpass-analysis=kernel-resource-usage ${EXTRA_FLAGS:-} \
+    --genco -o "$OUT" 2>&1 | grep -E "error|Function Name|VGPRs:|AGPRs:|ScratchSize|Occupancy|SGPRs:" || true
+
+[ -s "$OUT" ] || { echo "build failed" >&2; exit 1; }
+echo "wrote $OUT"

csrc/kernels/mla/opus/build/build_mla_decode_fp8_32mx4_64nx1_co.sh

@@ -0,0 +1,31 @@
+#!/bin/bash
+# Builds hsa/gfx950/mla_opus/opus_mla_decode_fp8_32mx4_64nx1.co with the amdgpu-pin-op-dst
+# toolchain (git@github.com:yuyzhang512/llvm-project.git, branch amdgpu-pin-op-dst), driven
+# through hipcc so the ROCm device libs and --genco still apply.
+#
+#   PIN_LLVM=/path/to/llvm-project/build/bin ./build/build_mla_decode_fp8_32mx4_64nx1_co.sh
+#
+# The flags are the code generation aiter's JIT gives its in-tree kernels; -Rpass-analysis
+# prints each entry's register budget.
+set -euo pipefail
+
+HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+OPUS="$(cd "$HERE/.." && pwd)"
+AITER="$(cd "$OPUS/../../../.." && pwd)"
+PIN_LLVM="${PIN_LLVM:-$HOME/llvm-project/build/bin}"
+OUT="${OUT:-$AITER/hsa/gfx950/mla_opus/opus_mla_decode_fp8_32mx4_64nx1.co}"
+
+[ -x "$PIN_LLVM/clang++" ] || { echo "no pin toolchain at $PIN_LLVM" >&2; exit 1; }
+
+HIP_CLANG_PATH="$PIN_LLVM" /opt/rocm/bin/hipcc -x hip "$OPUS/mla_decode_fp8_32mx4_64nx1_co.hip" \
+    -I"$AITER/csrc/include" -I"$AITER/csrc/kernels/mla" \
+    -std=c++20 -O3 --offload-arch=gfx950 -D__HIPCC_RTC__ \
+    -ffast-math -fgpu-flush-denormals-to-zero -fno-offload-uniform-block \
+    -mllvm --amdgpu-kernarg-preload-count=32 -mllvm --lsr-drop-solution=1 \
+    -mllvm -amdgpu-early-inline-all=true -mllvm -amdgpu-function-calls=false \
+    -mllvm -enable-post-misched=0 \
+    -Rpass-analysis=kernel-resource-usage ${EXTRA_FLAGS:-} \
+    --genco -o "$OUT" 2>&1 | grep -E "error|Function Name|VGPRs:|AGPRs:|ScratchSize|Occupancy|SGPRs:" || true
+
+[ -s "$OUT" ] || { echo "build failed" >&2; exit 1; }
+echo "wrote $OUT"