[HIP] [JIT] Fused MiniMaxM3 QKNorm+RoPE+CacheInsert

aiter/__init__.py

@@ -113,6 +113,7 @@ def getLogger():
     from .ops.moe_sorting_opus import *
     from .ops.moe_mxfp4_aux import *
     from .ops.pa_sparse_prefill_opus import *
+    from .ops.minimax_m3_fused_qknorm_rope import *
     from .ops.pos_encoding import *
     from .ops.cache import *
     from .ops.rmsnorm import *

aiter/jit/optCompilerConfig.json

@@ -1200,6 +1200,23 @@
         "verbose": "False",
         "blob_gen_cmd": "''"
     },
+    "module_minimax_m3_fused_qknorm_rope_cache_shuffle": {
+        "srcs": [
+            "f'{AITER_CSRC_DIR}/pybind/minimax_m3_fused_qknorm_rope_cache_shuffle_pybind.cu'",
+            "f'{AITER_CSRC_DIR}/kernels/minimax_m3_fused_qknorm_rope_cache_shuffle.cu'"
+        ],
+        "flags_extra_cc": [],
+        "flags_extra_hip": [
+            "'-DENABLE_FP8'"
+        ],
+        "extra_ldflags": "None",
+        "extra_include": [
+            "f'{AITER_CSRC_DIR}/include/ck_tile'",
+            "f'{AITER_CSRC_DIR}/include/opus'"
+        ],
+        "verbose": "False",
+        "blob_gen_cmd": "''"
+    },
     "module_fused_qk_norm_rope_cache_quant_shuffle": {
         "srcs": [
             "f'{AITER_CSRC_DIR}/pybind/fused_qk_norm_rope_cache_quant_pybind.cu'",

aiter/ops/minimax_m3_fused_qknorm_rope.py

@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved.
+
+
+from torch import Tensor
+
+from ..jit.core import compile_ops
+
+
+@compile_ops("module_minimax_m3_fused_qknorm_rope_cache_shuffle", develop=True)
+def minimax_m3_qknorm_rope_cache_shuffle_insert(
+    qkv: Tensor,
+    q_norm_weight: Tensor,
+    k_norm_weight: Tensor,
+    cos_sin_cache: Tensor,
+    positions: Tensor,
+    num_heads: int,
+    num_kv_heads: int,
+    num_index_heads: int,
+    rotary_dim: int,
+    eps: float,
+    slot_mapping: Tensor,
+    k_cache: Tensor,
+    v_cache: Tensor,
+    q_out: Tensor,
+    index_q_norm_weight: Tensor | None = None,
+    index_k_norm_weight: Tensor | None = None,
+    index_slot_mapping: Tensor | None = None,
+    index_cache: Tensor | None = None,
+    index_q_out: Tensor | None = None,
+    kv_cache_dtype: str = "auto",
+    k_scale: Tensor | None = None,
+    v_scale: Tensor | None = None,
+    skip_index_branch: bool = False,
+) -> None:
+    """Fused MiniMax-M3 QK-norm + partial NeoX RoPE + page-16 SHUFFLE KV insert.
+
+    Consumes the sparse layer's packed projection row
+    ``[q | k | v | index_q | index_k]`` (all head_dim=128) and, in one pass:
+
+    * writes Gemma-normed + roped ``q`` into ``q_out`` and ``index_q`` into
+      ``index_q_out``,
+    * scatters normed + roped ``k`` and verbatim ``v`` into the paged caches in
+      the layout ``pa_decode_gluon`` reads::
+
+          k_cache: [num_pages, num_kv_heads, head_dim // x,  page_size, x]
+          v_cache: [num_pages, num_kv_heads, page_size // x, head_dim,  x]
+
+      where ``x = 16 // k_cache.element_size()`` and ``page_size`` is
+      ``k_cache.shape[3]`` (16 for the gluon decode path),
+    * scatters normed + roped ``index_k`` into ``index_cache``
+      ``[num_pages, page_size, 128]``.
+
+    ``k``/``v``/``index_k`` are *not* written back into ``qkv``; they are
+    consumed only by the cache scatter. ``skip_index_branch=True`` skips all
+    index work (the row still carries the index sub-blocks).
+
+    ``k_scale``/``v_scale`` are optional single-element fp32 **device** tensors
+    applied as ``value / scale`` when ``kv_cache_dtype`` is quantized, matching
+    ``reshape_and_cache``.
+    """

csrc/include/minimax_m3_fused_qknorm_rope_cache_shuffle.h

@@ -0,0 +1,46 @@
+#pragma once
+// SPDX-License-Identifier: MIT
+// Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved.
+
+#include <optional>
+#include <string>
+
+#include "aiter_tensor.h"
+
+namespace aiter {
+
+// Fused MiniMax-M3 attention pre-processing that writes the paged K/V caches
+// directly in the page-16 SHUFFLE layout consumed by pa_decode_gluon.
+//
+//   k_cache: [num_pages, num_kv_heads, head_dim/x,  page_size, x]
+//   v_cache: [num_pages, num_kv_heads, page_size/x, head_dim,  x]   x = 16/elem_size
+//
+// Replaces the three-kernel sequence (fused QK-norm/RoPE -> reshape_and_cache
+// with asm_layout=True -> index-cache scatter) with a single pass over the
+// fused ``qkv`` row [q | k | v | index_q | index_k].
+void minimax_m3_qknorm_rope_cache_shuffle_insert(
+    aiter_tensor_t& qkv,           // [num_tokens, qkv_row] fp16/bf16
+    aiter_tensor_t& q_norm_weight, // [head_dim]
+    aiter_tensor_t& k_norm_weight, // [head_dim]
+    aiter_tensor_t& cos_sin_cache, // [max_pos, rotary_dim]
+    aiter_tensor_t& positions,     // [num_tokens] i64
+    int64_t num_heads,
+    int64_t num_kv_heads,
+    int64_t num_index_heads,
+    int64_t rotary_dim,
+    double eps,
+    aiter_tensor_t& slot_mapping, // [num_tokens] i64
+    aiter_tensor_t& k_cache,
+    aiter_tensor_t& v_cache,
+    aiter_tensor_t& q_out, // [num_tokens, num_heads * head_dim]
+    std::optional<aiter_tensor_t> index_q_norm_weight,
+    std::optional<aiter_tensor_t> index_k_norm_weight,
+    std::optional<aiter_tensor_t> index_slot_mapping, // [num_tokens] i64
+    std::optional<aiter_tensor_t> index_cache,        // [pages, page_size, head_dim]
+    std::optional<aiter_tensor_t> index_q_out,        // [num_tokens, niq*head_dim]
+    const std::string& kv_cache_dtype,
+    std::optional<aiter_tensor_t> k_scale, // fp32 device scalar
+    std::optional<aiter_tensor_t> v_scale, // fp32 device scalar
+    bool skip_index_branch);
+
+} // namespace aiter

csrc/include/rocm_ops.hpp

@@ -2154,6 +2154,33 @@ namespace py = pybind11;
           py::arg("quant_group_size") = 128,           \
           py::arg("scale_layout")     = 0);
 
+#define MINIMAX_M3_FUSED_QKNORM_ROPE_CACHE_SHUFFLE_PYBIND      \
+    m.def("minimax_m3_qknorm_rope_cache_shuffle_insert",       \
+          &aiter::minimax_m3_qknorm_rope_cache_shuffle_insert, \
+          py::arg("qkv"),                                      \
+          py::arg("q_norm_weight"),                            \
+          py::arg("k_norm_weight"),                            \
+          py::arg("cos_sin_cache"),                            \
+          py::arg("positions"),                                \
+          py::arg("num_heads"),                                \
+          py::arg("num_kv_heads"),                             \
+          py::arg("num_index_heads"),                          \
+          py::arg("rotary_dim"),                               \
+          py::arg("eps"),                                      \
+          py::arg("slot_mapping"),                             \
+          py::arg("k_cache"),                                  \
+          py::arg("v_cache"),                                  \
+          py::arg("q_out"),                                    \
+          py::arg("index_q_norm_weight") = std::nullopt,       \
+          py::arg("index_k_norm_weight") = std::nullopt,       \
+          py::arg("index_slot_mapping")  = std::nullopt,       \
+          py::arg("index_cache")         = std::nullopt,       \
+          py::arg("index_q_out")         = std::nullopt,       \
+          py::arg("kv_cache_dtype")      = "auto",             \
+          py::arg("k_scale")             = std::nullopt,       \
+          py::arg("v_scale")             = std::nullopt,       \
+          py::arg("skip_index_branch")   = false);
+
 #define SMOOTHQUANT_PYBIND                      \
     m.def("smoothquant_fwd", &smoothquant_fwd); \
     m.def("moe_smoothquant_fwd", &moe_smoothquant_fwd);