vllm.models.minimax_m3.amd.ops

AMD/ROCm fused Triton ops for MiniMax-M3.

These replace per-element PyTorch fallbacks (FlashInfer / fused HIP kernels are unavailable on ROCm) with single-pass Triton kernels to cut launch overhead and intermediate-tensor traffic during decode.

Modules:

• gemma_rmsnorm –

  Fused Gemma-style RMSNorm for AMD ROCm via Triton.

• index_topk –

  Triton kernels for MiniMax M3 lightning-indexer block scoring + top-k.

• indexer_context_parallel –

  Opt-in M3 compute-only index context partitioning, independent of global DCP.

• indexer_cp_exchange –

  Candidate exchange for MiniMax-M3 Triton indexer CP (graph-safe).

• sparse_attn –

  ROCm gfx942/gfx950 block-sparse GQA prefill kernel for MiniMax-M3.

• sparse_pa –

  AITER page-16 sparse paged-attention helpers for MiniMax-M3 on ROCm.

• swiglu_oai –

  Fused SwiGLU-OAI activation (split layout) for AMD ROCm via Triton.

Functions:

• gemma_fused_add_rmsnorm_mxfp8 –

  gemma_fused_add_rmsnorm with the normed output quantized to MXFP8.

• swiglu_oai_quantize_mxfp8 –

  SwiGLU-OAI on split-layout [M, 2I] fused with MXFP8 activation-quant.

• swiglu_oai_split –

  SwiGLU-OAI on a split-layout [*, 2I] tensor -> [*, I].

gemma_fused_add_rmsnorm_mxfp8(x, residual, weight, eps)

gemma_fused_add_rmsnorm with the normed output quantized to MXFP8.

Bit-identical to mxfp8_e4m3_quantize(gemma_fused_add_rmsnorm(...)[0]).

Returns:

• Tensor –

  (x_q, x_scale, residual_out): x_q float8_e4m3fn [M, N] and

• Tensor –

  x_scale uint8 E8M0 [M, N // 32] (row-major, not swizzled), with

• Tensor –

  the leading dims of x flattened into M; residual_out has

• tuple[Tensor, Tensor, Tensor] –

  the shape of x.

Source code in vllm/models/minimax_m3/amd/ops/gemma_rmsnorm.py

def gemma_fused_add_rmsnorm_mxfp8(
    x: torch.Tensor,
    residual: torch.Tensor,
    weight: torch.Tensor,
    eps: float,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """``gemma_fused_add_rmsnorm`` with the normed output quantized to MXFP8.

    Bit-identical to ``mxfp8_e4m3_quantize(gemma_fused_add_rmsnorm(...)[0])``.

    Returns:
        ``(x_q, x_scale, residual_out)``: ``x_q`` float8_e4m3fn ``[M, N]`` and
        ``x_scale`` uint8 E8M0 ``[M, N // 32]`` (row-major, not swizzled), with
        the leading dims of ``x`` flattened into ``M``; ``residual_out`` has
        the shape of ``x``.

    """
    n = x.shape[-1]
    assert n % MXFP8_BLOCK_SIZE == 0, (
        f"MXFP8 needs the hidden size to be a multiple of {MXFP8_BLOCK_SIZE}, got {n}"
    )
    x_q, x_scale, res_out = _gemma_fused_add_rmsnorm(x, residual, weight, eps, True)
    assert x_scale is not None
    return x_q, x_scale, res_out

swiglu_oai_quantize_mxfp8(gate_up, alpha, beta, limit, block_m=64)

SwiGLU-OAI on split-layout [M, 2I] fused with MXFP8 activation-quant.

Returns (act_q [M, I] float8_e4m3fn, act_scale [M, I//32] uint8 E8M0), identical to mxfp8_e4m3_quantize(swiglu_oai_split(gate_up)) but in a single Triton pass (no bf16 intermediate). Used between the two GEMMs of the native MXFP8 MoE. Numerically equivalent to the unfused chain (bit-exact on measured MoE shapes); marginally more accurate (fp32 act, no bf16 round-trip).

Source code in vllm/models/minimax_m3/amd/ops/swiglu_oai.py

def swiglu_oai_quantize_mxfp8(
    gate_up: torch.Tensor,
    alpha: float,
    beta: float,
    limit: float | None,
    block_m: int = 64,
) -> tuple[torch.Tensor, torch.Tensor]:
    """SwiGLU-OAI on split-layout ``[M, 2I]`` fused with MXFP8 activation-quant.

    Returns ``(act_q [M, I] float8_e4m3fn, act_scale [M, I//32] uint8 E8M0)``,
    identical to ``mxfp8_e4m3_quantize(swiglu_oai_split(gate_up))`` but in a
    single Triton pass (no bf16 intermediate). Used between the two GEMMs of the
    native MXFP8 MoE. Numerically equivalent to the unfused chain (bit-exact on
    measured MoE shapes); marginally more accurate (fp32 act, no bf16 round-trip).
    """
    from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
        MXFP8_BLOCK_SIZE,
        MXFP8_SCALE_DTYPE,
        MXFP8_VALUE_DTYPE,
    )

    two_i = gate_up.shape[-1]
    n_inter = two_i // 2
    assert n_inter % MXFP8_BLOCK_SIZE == 0, (
        f"fused swiglu+quant needs I % {MXFP8_BLOCK_SIZE} == 0, got I={n_inter}"
    )
    g1 = gate_up.reshape(-1, two_i).contiguous()
    M = g1.shape[0]
    aq = torch.empty((M, n_inter), dtype=MXFP8_VALUE_DTYPE, device=g1.device)
    asc = torch.empty(
        (M, n_inter // MXFP8_BLOCK_SIZE), dtype=MXFP8_SCALE_DTYPE, device=g1.device
    )
    grid = (triton.cdiv(M, block_m), n_inter // MXFP8_BLOCK_SIZE)
    _swiglu_oai_quant_kernel[grid](
        g1,
        aq,
        asc,
        M,
        n_inter,
        g1.stride(0),
        g1.stride(1),
        aq.stride(0),
        aq.stride(1),
        asc.stride(0),
        asc.stride(1),
        float(alpha),
        float(beta),
        0.0 if limit is None else float(limit),
        HAS_LIMIT=limit is not None,
        BLOCK_M=block_m,
        FP8_MAX=float(torch.finfo(MXFP8_VALUE_DTYPE).max),
        TINY=float(torch.finfo(torch.float32).tiny),
        num_warps=4,
    )
    return aq, asc

swiglu_oai_split(gate_up, alpha, beta, limit, out_dtype=None)

SwiGLU-OAI on a split-layout [*, 2I] tensor -> [*, I].

Source code in vllm/models/minimax_m3/amd/ops/swiglu_oai.py

def swiglu_oai_split(
    gate_up: torch.Tensor,
    alpha: float,
    beta: float,
    limit: float | None,
    out_dtype: torch.dtype | None = None,
) -> torch.Tensor:
    """SwiGLU-OAI on a split-layout ``[*, 2I]`` tensor -> ``[*, I]``."""
    orig_shape = gate_up.shape
    two_i = orig_shape[-1]
    n_inter = two_i // 2
    x2 = gate_up.reshape(-1, two_i)
    m = x2.shape[0]
    dt = out_dtype if out_dtype is not None else gate_up.dtype
    out = torch.empty((m, n_inter), dtype=dt, device=gate_up.device)
    # Tile tuned on gfx950. The SwiGLU intermediate is sharded across tensor
    # parallel ranks (per-rank n_inter = I / tp: dense I=12288, MoE I=3072), and
    # a 512-wide tile (4 warps, ~2 elems/lane) only helps once the per-rank slice
    # is large enough to be bandwidth-bound — at TP=1 prefill that is ~1.25-1.35x
    # faster than 256. For small sharded slices (high TP) the kernel is launch-
    # bound (~12us) and a wide tile can slightly regress, so fall back to 256.
    # Decode is launch-bound at every TP. num_warps=8 underfills this tile, so it
    # is pinned to 4.
    block_i = 512 if n_inter >= 2048 else 256
    grid = (m, triton.cdiv(n_inter, block_i))
    _swiglu_oai_kernel[grid](
        x2,
        out,
        n_inter,
        x2.stride(0),
        x2.stride(1),
        out.stride(0),
        out.stride(1),
        float(alpha),
        float(beta),
        0.0 if limit is None else float(limit),
        HAS_LIMIT=limit is not None,
        BLOCK_I=block_i,
        num_warps=4,
    )
    return out.reshape(*orig_shape[:-1], n_inter)