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

Adds the prefill half of the DeepSeek-V3.2 / GLM-5 sparse-attention lightning
indexer as a HIP kernel, alongside the existing Triton/Gluon one:

    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 is already gathered into a contiguous [N, 128] buffer, so there is no block
table. Same call contract as aiter.ops.triton.attention.fp8_mqa_logits.

  aiter/ops/fp8_mqa_logits.py      the op, plus is_supported()
  csrc/kernels/fp8_mqa_logits.cu   module_fp8_mqa_logits
  op_tests/test_fp8_mqa_logits.py  triton vs hip, one fp32 reference

The kernel streams K HBM->register with no LDS staging and contracts 32 heads x
32 tokens per mfma_scale_f32_32x32x64_f8f6f4 tile. Heads reduce across lanes
with v_permlane32_swap_b32, two query rows at a time, and the per-token scale is
applied once after that reduction (kv_scale >= 0, so it commutes with the ReLU).
BLOCK_M query rows share one K stream; the warps of a block split the tile range
rather than the rows, which keeps several K streams in flight per block. With
clean_logits the kernel writes the -inf outside each window itself, so the caller
needs no separate fill pass over the whole output.

Row groups are dispatched high-m first for N > 2048. Under causal masking a row
group's work grows with m, so in natural order the longest-running blocks are
dispatched LAST and become the tail; reversing starts them first and lets the
short ones fill in behind. Worth +6% at (4096, 4096) and +3% at (8192, 8192),
neutral elsewhere, and free -- it only reorders block dispatch. `unroll2` and
`reverse_rows` are tri-state (-1 heuristic, 0/1 off/on) so the host can tell
"off" from "caller did not choose".

It is gfx950-only and fixed at n_heads=32/head_dim=128 (the shipped GLM-5-FP8
indexer shape). is_supported() gates on that so a caller serving other shapes can
route them to the Triton kernel rather than trip a TORCH_CHECK.

The test sweeps both kernels over the same shapes and six window modes. Beyond
causal and cp it covers misaligned, empty, past-end and multi-request windows --
all legal indexer input at a chunk boundary, and where the masking and the -inf
fill are easiest to get wrong. It poisons the output buffer with NaN first, so a
position no kernel writes fails the -inf mask check instead of passing on
allocator leftovers -- that poisoning is a separate call from the timed one, so
neither candidate is charged a full-output memset per iteration. Shapes include
the real chunk sizes a request is split into, and the reference runs in
query-row chunks so those fit.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Author: anhcvt

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,65 @@
+# 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);