[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_a8w8_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

@@ -1028,21 +1028,14 @@ def mla_decode_fwd(
             and is_experimental_enabled()
         )
 
-        # 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.
-        opus_is_fp8 = (
-            q.dtype == dtypes.fp8
-            and kv_buffer.dtype == dtypes.fp8
-            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
         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)
+            and (
+                (q.dtype == dtypes.bf16 and kv_buffer.dtype == dtypes.bf16)
+                or (q.dtype == dtypes.fp8 and kv_buffer.dtype == dtypes.fp8)
+            )
         )
 
         # Head counts with code objects exported from the FlyDSL PS1 kernel
@@ -1099,28 +1092,6 @@ 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(
-                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,
-                o,
-                final_lse,
-                q_scale,
-                kv_scale,
-                causal,
-            )
         elif use_opus:
             aiter.opus_mla_decode_fwd(
                 q,
@@ -1136,9 +1107,16 @@ def mla_decode_fwd(
                 attn_lse,
                 o,
                 final_lse,
+                q_scale=q_scale,
+                kv_scale=kv_scale,
+                max_seqlen_q=max_seqlen_q,
+                causal=causal,
+                g_kv_indptr=g_kv_indptr,
+                cp_world_size=cp_world_size,
+                cp_rank=cp_rank,
             )
         elif use_hk:
-            aiter.hk_mla_decode_fwd(
+            aiter.hk_mla_v32_decode_fwd(
                 q,
                 kv_buffer,
                 qo_indptr,

aiter/ops/attention.py

@@ -1970,8 +1970,8 @@ 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,
@@ -1983,28 +1983,11 @@ def opus_mla_decode_fwd(
     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
+    max_seqlen_q: int = -1,  # fp8: required, picks the kernel; bf16: unused
+    causal: bool = True,  # fp8: mask across the max_seqlen_q query tokens
+    g_kv_indptr: torch.Tensor | None = None,  # fp8 round-robin CP: GLOBAL kv_indptr
+    cp_world_size: int = 1,  # fp8 round-robin CP: number of ranks, 1 = no CP
+    cp_rank: int = 0,  # fp8 round-robin CP: this rank
 ) -> None: ...

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

@@ -0,0 +1,33 @@
+#!/bin/bash
+# Builds hsa/gfx950/mla_opus/opus_mla_decode_a8w8_32mx2_32nx2.co with the named-register pin
+# toolchain (https://github.com/yuyzhang512/llvm-project, branch amdgpu-named-reg-pin-std), driven
+# through hipcc so the ROCm device libs and --genco still apply. Besides clang and lld, --genco
+# needs that build's clang-offload-bundler and llvm-objcopy:
+#   ninja clang lld clang-offload-bundler llvm-objcopy
+#
+#   PIN_LLVM=/path/to/llvm-project/build/bin ./build/build_mla_decode_a8w8_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-std/build/bin}"
+OUT="${OUT:-$AITER/hsa/gfx950/mla_opus/opus_mla_decode_a8w8_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_a8w8_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_a8w8_32mx4_64nx1_co.sh

@@ -0,0 +1,33 @@
+#!/bin/bash
+# Builds hsa/gfx950/mla_opus/opus_mla_decode_a8w8_32mx4_64nx1.co with the named-register pin
+# toolchain (https://github.com/yuyzhang512/llvm-project, branch amdgpu-named-reg-pin-std), driven
+# through hipcc so the ROCm device libs and --genco still apply. Besides clang and lld, --genco
+# needs that build's clang-offload-bundler and llvm-objcopy:
+#   ninja clang lld clang-offload-bundler llvm-objcopy
+#
+#   PIN_LLVM=/path/to/llvm-project/build/bin ./build/build_mla_decode_a8w8_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-std/build/bin}"
+OUT="${OUT:-$AITER/hsa/gfx950/mla_opus/opus_mla_decode_a8w8_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_a8w8_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"