[HIP] [JIT] fp8_mqa_logits: hand-written gfx950 prefill indexer kernel

aiter/jit/optCompilerConfig.json

@@ -1872,5 +1872,16 @@
         "extra_include": [],
         "verbose": "False",
         "blob_gen_cmd": "''"
+    },
+    "module_fp8_mqa_logits": {
+        "srcs": [
+            "f'{AITER_CSRC_DIR}/pybind/fp8_mqa_logits_pybind.cu'",
+            "f'{AITER_CSRC_DIR}/kernels/fp8_mqa_logits.cu'"
+        ],
+        "flags_extra_cc": [],
+        "flags_extra_hip": ["'-fno-honor-nans'"],
+        "extra_ldflags": "None",
+        "extra_include": [],
+        "verbose": "False"
     }
 }

aiter/ops/fp8_mqa_logits.py

@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+"""Hand-written HIP prefill-phase FP8 MQA indexer logits kernel for gfx950.
+
+The dense half of the DeepSeek-V3.2 / GLM-5 sparse-attention lightning indexer:
+
+    logits[m, n] = sum_h relu(Q[m,h,:] . K[n,:]) * w[m,h] * kv_scale[n]
+
+K has already been gathered into a contiguous ``[N, 128]`` buffer, so there is no
+block table. Drop-in for ``aiter.ops.triton.attention.fp8_mqa_logits``, so the two
+can be A/B'd on identical inputs.
+
+Fixed at n_heads=32, head_dim=128 -- the shipped GLM-5-FP8 indexer shape.
+``is_supported()`` gates on that, so a caller that also has to serve other shapes
+can route them to the Triton kernel rather than trip a TORCH_CHECK.
+"""
+
+from torch import Tensor
+
+from ..jit.core import compile_ops
+from ..jit.utils.chip_info import get_gfx
+
+MD_NAME = "module_fp8_mqa_logits"
+
+SUPPORTED_GFX = ("gfx950",)
+NUM_HEADS = 32
+HEAD_DIM = 128
+
+
+def is_supported(num_heads: int, head_dim: int) -> bool:
+    """True when this kernel can run this shape on this device."""
+    return (
+        get_gfx() in SUPPORTED_GFX and num_heads == NUM_HEADS and head_dim == HEAD_DIM
+    )
+
+
+@compile_ops(MD_NAME, fc_name="fp8_mqa_logits")
+def fp8_mqa_logits(
+    q_fp8: Tensor,
+    k_fp8: Tensor,
+    kv_scale: Tensor,
+    weights: Tensor,
+    cu_seqlen_ks: Tensor,
+    cu_seqlen_ke: Tensor,
+    BlockM: int = 0,
+    SplitN: int = 0,
+    num_warps: int = 0,
+    TotalCuCount: int = 256,
+    clean_logits: bool = True,
+    unroll2: int = -1,
+    reverse_rows: int = -1,
+    out: Tensor | None = None,
+) -> Tensor:
+    """Prefill indexer logits over a contiguous K buffer.
+
+    q_fp8       [M, 32, 128] fp8   k_fp8    [N, 128] fp8
+    kv_scale    [N] f32            weights  [M, 32] f32
+    cu_seqlen_ks/ke [M] i32 -- row m is valid on [ks[m], ke[m]). Either bound may
+    legally sit outside [0, N); the row is then empty over the part that does.
+
+    The zero-valued tunables (BlockM, SplitN, num_warps) mean "use the host
+    heuristic"; unroll2 and reverse_rows are tri-state, -1 for the heuristic and
+    0/1 to force off/on. Writes into `out` when given (and returns it), otherwise allocates
+    [M, N] f32; outside each row's window the kernel writes -inf when clean_logits,
+    and leaves the buffer untouched otherwise.
+    """

csrc/include/fp8_mqa_logits.h

@@ -0,0 +1,44 @@
+#pragma once
+// SPDX-License-Identifier: MIT
+// Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+//
+// Hand-written HIP prefill-phase FP8 MQA indexer logits kernel for gfx950.
+//
+// The dense half of the DeepSeek-V3.2 / GLM-5 sparse-attention lightning indexer:
+//
+//   logits[m, n] = sum_h relu(Q[m,h,:] . K[n,:]) * w[m,h] * kv_scale[n]
+//                  for n in [ks[m], ke[m]), -inf elsewhere
+//
+// K has already been gathered out of the paged cache into a contiguous [N, 128]
+// buffer, so there is no block table. Same contract as
+// `aiter.ops.triton.attention.fp8_mqa_logits`, so it drops into the same call site.
+//
+// Fixed at n_heads=32, head_dim=128 -- the shipped GLM-5-FP8 indexer shape.
+
+// aiter builds csrc without hipify, so the ATen CUDA headers (which pull in
+// cuda_runtime_api.h) are unavailable -- use the HIP-flavoured ones, as the rest
+// of csrc does.
+#include <ATen/hip/HIPContext.h>
+#include <ATen/hip/impl/HIPGuardImplMasqueradingAsCUDA.h>
+#include <c10/util/Optional.h>
+#include <torch/all.h>
+#include <torch/extension.h>
+#include <cstdlib>
+
+// Writes into `out` when given (and returns it), otherwise allocates [M, N] f32.
+// Positions outside a row's [ks, ke) window read -inf when `clean_logits`, and are
+// left untouched otherwise.
+torch::Tensor fp8_mqa_logits(torch::Tensor q_fp8,        // [M, 32, 128] fp8
+                             torch::Tensor k_fp8,        // [N, 128]     fp8
+                             torch::Tensor kv_scale,     // [N]          f32
+                             torch::Tensor weights,      // [M, 32]      f32
+                             torch::Tensor cu_seqlen_ks, // [M]          i32
+                             torch::Tensor cu_seqlen_ke, // [M]          i32
+                             int64_t BlockM,
+                             int64_t SplitN,
+                             int64_t num_warps,
+                             int64_t TotalCuCount,
+                             bool clean_logits,
+                             int64_t Unroll2,     // -1 = host heuristic, 0/1 = off/on
+                             int64_t ReverseRows, // -1 = host heuristic, 0/1 = off/on
+                             std::optional<torch::Tensor> out);

csrc/include/rocm_ops.hpp

@@ -2780,3 +2780,21 @@ namespace py = pybind11;
           py::arg("final_lse"),              \
           py::arg("q_scale"),                \
           py::arg("kv_scale"));
+
+#define FP8_MQA_LOGITS_PYBIND                          \
+    m.def("fp8_mqa_logits",                            \
+          &fp8_mqa_logits,                             \
+          py::arg("q_fp8"),                            \
+          py::arg("k_fp8"),                            \
+          py::arg("kv_scale"),                         \
+          py::arg("weights"),                          \
+          py::arg("cu_seqlen_ks"),                     \
+          py::arg("cu_seqlen_ke"),                     \
+          py::arg("BlockM")       = 0,                 \
+          py::arg("SplitN")       = 0,                 \
+          py::arg("num_warps")    = 0,                 \
+          py::arg("TotalCuCount") = 256,               \
+          py::arg("clean_logits") = true,              \
+          py::arg("unroll2")      = -1,                \
+          py::arg("reverse_rows") = -1,                \
+          py::arg("out")          = std::nullopt);