chore: import upstream snapshot with attribution

python/sglang/jit_kernel/minimax_m3/__init__.py

@@ -0,0 +1,25 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Fused Triton kernels for MiniMax-M3 on AMD ROCm (gfx94x / gfx95x).
+
+Model-scoped JIT kernels (mirrors ``jit_kernel/dsv4``), split by op type:
+  * ``rmsnorm`` -- fused fp32 Gemma RMSNorm (plain + fused-add-residual)
+  * ``swiglu``  -- fused fp32 SwiGLU-OAI (split layout)
+"""
+
+from sglang.jit_kernel.minimax_m3.rmsnorm import (
+    _num_warps,
+    gemma_fused_add_rmsnorm,
+    gemma_rmsnorm,
+)
+from sglang.jit_kernel.minimax_m3.swiglu import (
+    swiglu_oai_mxfp8_quant,
+    swiglu_oai_split,
+)
+
+__all__ = [
+    "gemma_rmsnorm",
+    "gemma_fused_add_rmsnorm",
+    "swiglu_oai_split",
+    "swiglu_oai_mxfp8_quant",
+    "_num_warps",
+]

python/sglang/jit_kernel/minimax_m3/qk_norm_rope.py

@@ -0,0 +1,659 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Fused MiniMax-M3 per-head Gemma Q/K RMSNorm + partial RoPE for ROCm."""
+
+from typing import Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _qk_gemma_rmsnorm_rope_kernel(
+    q_ptr,
+    k_ptr,
+    q_out_ptr,
+    k_out_ptr,
+    q_weight_ptr,
+    k_weight_ptr,
+    positions_ptr,
+    cos_sin_cache_ptr,
+    q_stride_m,
+    q_stride_d,
+    k_stride_m,
+    k_stride_d,
+    q_heads: tl.constexpr,
+    k_heads: tl.constexpr,
+    head_dim: tl.constexpr,
+    rotary_dim: tl.constexpr,
+    eps: tl.constexpr,
+    is_neox_style: tl.constexpr,
+    BLOCK_HD: tl.constexpr,
+):
+    token_id = tl.program_id(0)
+    head_program = tl.program_id(1)
+    cols = tl.arange(0, BLOCK_HD)
+    mask = cols < head_dim
+    half_rotary: tl.constexpr = rotary_dim // 2
+
+    is_q = head_program < q_heads
+    head_id = tl.where(is_q, head_program, head_program - q_heads)
+    in_ptr = tl.where(is_q, q_ptr, k_ptr)
+    out_ptr = tl.where(is_q, q_out_ptr, k_out_ptr)
+    weight_ptr = tl.where(is_q, q_weight_ptr, k_weight_ptr)
+    stride_m = tl.where(is_q, q_stride_m, k_stride_m)
+    stride_d = tl.where(is_q, q_stride_d, k_stride_d)
+    n_heads = tl.where(is_q, q_heads, k_heads)
+
+    base_in = in_ptr + token_id * stride_m + head_id * head_dim * stride_d
+    x = tl.load(base_in + cols * stride_d, mask=mask, other=0.0).to(tl.float32)
+    w = tl.load(weight_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+    var = tl.sum(x * x, axis=0) / head_dim
+    rstd = tl.rsqrt(var + eps)
+    normed = x * rstd * (1.0 + w)
+    # Match the unfused path: GemmaRMSNorm writes bf16/fp16, then RoPE reads
+    # that rounded value in the following kernel.
+    normed = normed.to(q_out_ptr.dtype.element_ty).to(tl.float32)
+
+    rotary_mask = cols < rotary_dim
+    if is_neox_style:
+        partner_cols = tl.where(
+            cols < half_rotary, cols + half_rotary, cols - half_rotary
+        )
+        cos_cols = tl.where(cols < half_rotary, cols, cols - half_rotary)
+        sign = tl.where(cols < half_rotary, -1.0, 1.0)
+    else:
+        partner_cols = tl.where((cols % 2) == 0, cols + 1, cols - 1)
+        cos_cols = cols // 2
+        sign = tl.where((cols % 2) == 0, -1.0, 1.0)
+
+    partner_mask = partner_cols < head_dim
+    x_partner = tl.load(
+        base_in + partner_cols * stride_d,
+        mask=partner_mask,
+        other=0.0,
+    ).to(tl.float32)
+    w_partner = tl.load(
+        weight_ptr + partner_cols,
+        mask=partner_mask,
+        other=0.0,
+    ).to(tl.float32)
+    partner_normed = x_partner * rstd * (1.0 + w_partner)
+    partner_normed = partner_normed.to(q_out_ptr.dtype.element_ty).to(tl.float32)
+
+    pos = tl.load(positions_ptr + token_id).to(tl.int64)
+    cos_sin_base = cos_sin_cache_ptr + pos * rotary_dim
+    cos = tl.load(cos_sin_base + cos_cols, mask=rotary_mask, other=1.0).to(tl.float32)
+    sin = tl.load(
+        cos_sin_base + half_rotary + cos_cols,
+        mask=rotary_mask,
+        other=0.0,
+    ).to(tl.float32)
+    rotated = normed * cos + sign * partner_normed * sin
+    out = tl.where(rotary_mask, rotated, normed)
+
+    base_out = out_ptr + token_id * n_heads * head_dim + head_id * head_dim
+    tl.store(base_out + cols, out.to(out_ptr.dtype.element_ty), mask=mask)
+
+
+def qk_gemma_rmsnorm_rope(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    positions: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    eps: float,
+    head_dim: int,
+    rotary_dim: int,
+    is_neox_style: bool,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Return normalized+rotated Q/K tensors with the same shapes as ``q``/``k``."""
+    assert q.dim() == 2 and k.dim() == 2
+    assert positions.dim() == 1
+    assert q.shape[0] == k.shape[0] == positions.shape[0]
+    assert q.shape[1] % head_dim == 0
+    assert k.shape[1] % head_dim == 0
+    assert rotary_dim <= head_dim and rotary_dim % 2 == 0
+
+    q_heads = q.shape[1] // head_dim
+    k_heads = k.shape[1] // head_dim
+    q_out = torch.empty(q.shape, dtype=q.dtype, device=q.device)
+    k_out = torch.empty(k.shape, dtype=k.dtype, device=k.device)
+    block_hd = triton.next_power_of_2(head_dim)
+
+    _qk_gemma_rmsnorm_rope_kernel[(q.shape[0], q_heads + k_heads)](
+        q,
+        k,
+        q_out,
+        k_out,
+        q_weight,
+        k_weight,
+        positions,
+        cos_sin_cache,
+        q.stride(0),
+        q.stride(1),
+        k.stride(0),
+        k.stride(1),
+        q_heads,
+        k_heads,
+        head_dim,
+        rotary_dim,
+        eps,
+        is_neox_style,
+        BLOCK_HD=block_hd,
+        num_warps=4,
+    )
+    return q_out, k_out
+
+
+@triton.jit
+def _sparse_qk_index_gemma_rmsnorm_rope_kernel(
+    q_ptr,
+    k_ptr,
+    idx_q_ptr,
+    idx_k_ptr,
+    q_out_ptr,
+    k_out_ptr,
+    idx_q_out_ptr,
+    idx_k_out_ptr,
+    q_weight_ptr,
+    k_weight_ptr,
+    idx_q_weight_ptr,
+    idx_k_weight_ptr,
+    positions_ptr,
+    cos_sin_cache_ptr,
+    q_stride_m,
+    q_stride_d,
+    k_stride_m,
+    k_stride_d,
+    idx_q_stride_m,
+    idx_q_stride_d,
+    idx_k_stride_m,
+    idx_k_stride_d,
+    q_heads: tl.constexpr,
+    k_heads: tl.constexpr,
+    idx_q_heads: tl.constexpr,
+    head_dim: tl.constexpr,
+    rotary_dim: tl.constexpr,
+    eps: tl.constexpr,
+    is_neox_style: tl.constexpr,
+    BLOCK_HD: tl.constexpr,
+):
+    token_id = tl.program_id(0)
+    head_program = tl.program_id(1)
+    cols = tl.arange(0, BLOCK_HD)
+    mask = cols < head_dim
+    half_rotary: tl.constexpr = rotary_dim // 2
+
+    main_heads: tl.constexpr = q_heads + k_heads
+    idx_k_program: tl.constexpr = q_heads + k_heads + idx_q_heads
+
+    is_q = head_program < q_heads
+    is_k = (head_program >= q_heads) & (head_program < main_heads)
+    is_idx_q = (head_program >= main_heads) & (head_program < idx_k_program)
+
+    head_id = tl.where(
+        is_q,
+        head_program,
+        tl.where(
+            is_k,
+            head_program - q_heads,
+            tl.where(is_idx_q, head_program - main_heads, 0),
+        ),
+    )
+
+    in_ptr = tl.where(
+        is_q,
+        q_ptr,
+        tl.where(is_k, k_ptr, tl.where(is_idx_q, idx_q_ptr, idx_k_ptr)),
+    )
+    out_ptr = tl.where(
+        is_q,
+        q_out_ptr,
+        tl.where(is_k, k_out_ptr, tl.where(is_idx_q, idx_q_out_ptr, idx_k_out_ptr)),
+    )
+    weight_ptr = tl.where(
+        is_q,
+        q_weight_ptr,
+        tl.where(
+            is_k,
+            k_weight_ptr,
+            tl.where(is_idx_q, idx_q_weight_ptr, idx_k_weight_ptr),
+        ),
+    )
+    stride_m = tl.where(
+        is_q,
+        q_stride_m,
+        tl.where(
+            is_k,
+            k_stride_m,
+            tl.where(is_idx_q, idx_q_stride_m, idx_k_stride_m),
+        ),
+    )
+    stride_d = tl.where(
+        is_q,
+        q_stride_d,
+        tl.where(
+            is_k,
+            k_stride_d,
+            tl.where(is_idx_q, idx_q_stride_d, idx_k_stride_d),
+        ),
+    )
+    out_heads = tl.where(
+        is_q,
+        q_heads,
+        tl.where(is_k, k_heads, tl.where(is_idx_q, idx_q_heads, 1)),
+    )
+
+    base_in = in_ptr + token_id * stride_m + head_id * head_dim * stride_d
+    x = tl.load(base_in + cols * stride_d, mask=mask, other=0.0).to(tl.float32)
+    w = tl.load(weight_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+    var = tl.sum(x * x, axis=0) / head_dim
+    rstd = tl.rsqrt(var + eps)
+    normed = x * rstd * (1.0 + w)
+    normed = normed.to(q_out_ptr.dtype.element_ty).to(tl.float32)
+
+    rotary_mask = cols < rotary_dim
+    if is_neox_style:
+        partner_cols = tl.where(
+            cols < half_rotary, cols + half_rotary, cols - half_rotary
+        )
+        cos_cols = tl.where(cols < half_rotary, cols, cols - half_rotary)
+        sign = tl.where(cols < half_rotary, -1.0, 1.0)
+    else:
+        partner_cols = tl.where((cols % 2) == 0, cols + 1, cols - 1)
+        cos_cols = cols // 2
+        sign = tl.where((cols % 2) == 0, -1.0, 1.0)
+
+    partner_mask = partner_cols < head_dim
+    x_partner = tl.load(
+        base_in + partner_cols * stride_d,
+        mask=partner_mask,
+        other=0.0,
+    ).to(tl.float32)
+    w_partner = tl.load(
+        weight_ptr + partner_cols,
+        mask=partner_mask,
+        other=0.0,
+    ).to(tl.float32)
+    partner_normed = x_partner * rstd * (1.0 + w_partner)
+    partner_normed = partner_normed.to(q_out_ptr.dtype.element_ty).to(tl.float32)
+
+    pos = tl.load(positions_ptr + token_id).to(tl.int64)
+    cos_sin_base = cos_sin_cache_ptr + pos * rotary_dim
+    cos = tl.load(cos_sin_base + cos_cols, mask=rotary_mask, other=1.0).to(tl.float32)
+    sin = tl.load(
+        cos_sin_base + half_rotary + cos_cols,
+        mask=rotary_mask,
+        other=0.0,
+    ).to(tl.float32)
+    rotated = normed * cos + sign * partner_normed * sin
+    out = tl.where(rotary_mask, rotated, normed)
+
+    base_out = out_ptr + token_id * out_heads * head_dim + head_id * head_dim
+    tl.store(base_out + cols, out.to(q_out_ptr.dtype.element_ty), mask=mask)
+
+
+def sparse_qk_index_gemma_rmsnorm_rope(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    idx_q: torch.Tensor,
+    idx_k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    idx_q_weight: torch.Tensor,
+    idx_k_weight: torch.Tensor,
+    positions: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    eps: float,
+    head_dim: int,
+    rotary_dim: int,
+    is_neox_style: bool,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Fuse main and sparse-index Gemma Q/K RMSNorm + RoPE into one launch."""
+    assert q.dim() == k.dim() == idx_q.dim() == idx_k.dim() == 2
+    assert positions.dim() == 1
+    assert q.shape[0] == k.shape[0] == idx_q.shape[0] == idx_k.shape[0]
+    assert q.shape[0] == positions.shape[0]
+    assert q.shape[1] % head_dim == 0
+    assert k.shape[1] % head_dim == 0
+    assert idx_q.shape[1] % head_dim == 0
+    assert idx_k.shape[1] == head_dim
+    assert rotary_dim <= head_dim and rotary_dim % 2 == 0
+
+    q_heads = q.shape[1] // head_dim
+    k_heads = k.shape[1] // head_dim
+    idx_q_heads = idx_q.shape[1] // head_dim
+    q_out = torch.empty(q.shape, dtype=q.dtype, device=q.device)
+    k_out = torch.empty(k.shape, dtype=k.dtype, device=k.device)
+    idx_q_out = torch.empty(idx_q.shape, dtype=idx_q.dtype, device=idx_q.device)
+    idx_k_out = torch.empty(idx_k.shape, dtype=idx_k.dtype, device=idx_k.device)
+    block_hd = triton.next_power_of_2(head_dim)
+
+    _sparse_qk_index_gemma_rmsnorm_rope_kernel[
+        (q.shape[0], q_heads + k_heads + idx_q_heads + 1)
+    ](
+        q,
+        k,
+        idx_q,
+        idx_k,
+        q_out,
+        k_out,
+        idx_q_out,
+        idx_k_out,
+        q_weight,
+        k_weight,
+        idx_q_weight,
+        idx_k_weight,
+        positions,
+        cos_sin_cache,
+        q.stride(0),
+        q.stride(1),
+        k.stride(0),
+        k.stride(1),
+        idx_q.stride(0),
+        idx_q.stride(1),
+        idx_k.stride(0),
+        idx_k.stride(1),
+        q_heads,
+        k_heads,
+        idx_q_heads,
+        head_dim,
+        rotary_dim,
+        eps,
+        is_neox_style,
+        BLOCK_HD=block_hd,
+        num_warps=4,
+    )
+    return q_out, k_out, idx_q_out, idx_k_out
+
+
+@triton.jit
+def _sparse_qk_index_gemma_rmsnorm_rope_cache_kernel(
+    q_ptr,
+    k_ptr,
+    v_ptr,
+    idx_q_ptr,
+    idx_k_ptr,
+    q_out_ptr,
+    k_out_ptr,
+    idx_q_out_ptr,
+    idx_k_out_ptr,
+    k_cache_ptr,
+    v_cache_ptr,
+    idx_k_cache_ptr,
+    loc_ptr,
+    q_weight_ptr,
+    k_weight_ptr,
+    idx_q_weight_ptr,
+    idx_k_weight_ptr,
+    positions_ptr,
+    cos_sin_cache_ptr,
+    q_stride_m,
+    q_stride_d,
+    k_stride_m,
+    k_stride_d,
+    v_stride_m,
+    v_stride_d,
+    idx_q_stride_m,
+    idx_q_stride_d,
+    idx_k_stride_m,
+    idx_k_stride_d,
+    k_cache_stride_s,
+    k_cache_stride_h,
+    k_cache_stride_d,
+    v_cache_stride_s,
+    v_cache_stride_h,
+    v_cache_stride_d,
+    idx_k_cache_stride_s,
+    idx_k_cache_stride_h,
+    idx_k_cache_stride_d,
+    q_heads: tl.constexpr,
+    k_heads: tl.constexpr,
+    idx_q_heads: tl.constexpr,
+    head_dim: tl.constexpr,
+    rotary_dim: tl.constexpr,
+    eps: tl.constexpr,
+    is_neox_style: tl.constexpr,
+    BLOCK_HD: tl.constexpr,
+):
+    token_id = tl.program_id(0)
+    head_program = tl.program_id(1)
+    cols = tl.arange(0, BLOCK_HD)
+    mask = cols < head_dim
+    half_rotary: tl.constexpr = rotary_dim // 2
+
+    main_heads: tl.constexpr = q_heads + k_heads
+    idx_k_program: tl.constexpr = q_heads + k_heads + idx_q_heads
+
+    is_q = head_program < q_heads
+    is_k = (head_program >= q_heads) & (head_program < main_heads)
+    is_idx_q = (head_program >= main_heads) & (head_program < idx_k_program)
+
+    head_id = tl.where(
+        is_q,
+        head_program,
+        tl.where(
+            is_k,
+            head_program - q_heads,
+            tl.where(is_idx_q, head_program - main_heads, 0),
+        ),
+    )
+
+    in_ptr = tl.where(
+        is_q,
+        q_ptr,
+        tl.where(is_k, k_ptr, tl.where(is_idx_q, idx_q_ptr, idx_k_ptr)),
+    )
+    out_ptr = tl.where(
+        is_q,
+        q_out_ptr,
+        tl.where(is_k, k_out_ptr, tl.where(is_idx_q, idx_q_out_ptr, idx_k_out_ptr)),
+    )
+    weight_ptr = tl.where(
+        is_q,
+        q_weight_ptr,
+        tl.where(
+            is_k,
+            k_weight_ptr,
+            tl.where(is_idx_q, idx_q_weight_ptr, idx_k_weight_ptr),
+        ),
+    )
+    stride_m = tl.where(
+        is_q,
+        q_stride_m,
+        tl.where(
+            is_k,
+            k_stride_m,
+            tl.where(is_idx_q, idx_q_stride_m, idx_k_stride_m),
+        ),
+    )
+    stride_d = tl.where(
+        is_q,
+        q_stride_d,
+        tl.where(
+            is_k,
+            k_stride_d,
+            tl.where(is_idx_q, idx_q_stride_d, idx_k_stride_d),
+        ),
+    )
+    out_heads = tl.where(
+        is_q,
+        q_heads,
+        tl.where(is_k, k_heads, tl.where(is_idx_q, idx_q_heads, 1)),
+    )
+
+    base_in = in_ptr + token_id * stride_m + head_id * head_dim * stride_d
+    x = tl.load(base_in + cols * stride_d, mask=mask, other=0.0).to(tl.float32)
+    w = tl.load(weight_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+    var = tl.sum(x * x, axis=0) / head_dim
+    rstd = tl.rsqrt(var + eps)
+    normed = x * rstd * (1.0 + w)
+    normed = normed.to(q_out_ptr.dtype.element_ty).to(tl.float32)
+
+    rotary_mask = cols < rotary_dim
+    if is_neox_style:
+        partner_cols = tl.where(
+            cols < half_rotary, cols + half_rotary, cols - half_rotary
+        )
+        cos_cols = tl.where(cols < half_rotary, cols, cols - half_rotary)
+        sign = tl.where(cols < half_rotary, -1.0, 1.0)
+    else:
+        partner_cols = tl.where((cols % 2) == 0, cols + 1, cols - 1)
+        cos_cols = cols // 2
+        sign = tl.where((cols % 2) == 0, -1.0, 1.0)
+
+    partner_mask = partner_cols < head_dim
+    x_partner = tl.load(
+        base_in + partner_cols * stride_d,
+        mask=partner_mask,
+        other=0.0,
+    ).to(tl.float32)
+    w_partner = tl.load(
+        weight_ptr + partner_cols,
+        mask=partner_mask,
+        other=0.0,
+    ).to(tl.float32)
+    partner_normed = x_partner * rstd * (1.0 + w_partner)
+    partner_normed = partner_normed.to(q_out_ptr.dtype.element_ty).to(tl.float32)
+
+    pos = tl.load(positions_ptr + token_id).to(tl.int64)
+    cos_sin_base = cos_sin_cache_ptr + pos * rotary_dim
+    cos = tl.load(cos_sin_base + cos_cols, mask=rotary_mask, other=1.0).to(tl.float32)
+    sin = tl.load(
+        cos_sin_base + half_rotary + cos_cols,
+        mask=rotary_mask,
+        other=0.0,
+    ).to(tl.float32)
+    rotated = normed * cos + sign * partner_normed * sin
+    out = tl.where(rotary_mask, rotated, normed)
+    out_typed = out.to(q_out_ptr.dtype.element_ty)
+
+    base_out = out_ptr + token_id * out_heads * head_dim + head_id * head_dim
+    tl.store(base_out + cols, out_typed, mask=mask)
+
+    loc = tl.load(loc_ptr + token_id)
+    cache_k_base = (
+        k_cache_ptr
+        + loc * k_cache_stride_s
+        + head_id * k_cache_stride_h
+        + cols * k_cache_stride_d
+    )
+    tl.store(cache_k_base, out_typed, mask=mask & is_k)
+
+    v_base = v_ptr + token_id * v_stride_m + head_id * head_dim * v_stride_d
+    v_val = tl.load(v_base + cols * v_stride_d, mask=mask & is_k, other=0.0)
+    cache_v_base = (
+        v_cache_ptr
+        + loc * v_cache_stride_s
+        + head_id * v_cache_stride_h
+        + cols * v_cache_stride_d
+    )
+    tl.store(cache_v_base, v_val, mask=mask & is_k)
+
+    is_idx_k = head_program == idx_k_program
+    idx_cache_base = (
+        idx_k_cache_ptr + loc * idx_k_cache_stride_s + cols * idx_k_cache_stride_d
+    )
+    tl.store(idx_cache_base, out_typed, mask=mask & is_idx_k)
+
+
+def sparse_qk_index_gemma_rmsnorm_rope_cache(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    idx_q: torch.Tensor,
+    idx_k: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    idx_k_cache: torch.Tensor,
+    out_cache_loc: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    idx_q_weight: torch.Tensor,
+    idx_k_weight: torch.Tensor,
+    positions: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    eps: float,
+    head_dim: int,
+    rotary_dim: int,
+    is_neox_style: bool,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Fuse sparse Q/K/index norm+RoPE with main KV and index-K cache stores."""
+    assert q.dim() == k.dim() == v.dim() == idx_q.dim() == idx_k.dim() == 2
+    assert k_cache.dim() == v_cache.dim() == idx_k_cache.dim() == 3
+    assert out_cache_loc.dim() == positions.dim() == 1
+    assert q.shape[0] == k.shape[0] == v.shape[0] == idx_q.shape[0] == idx_k.shape[0]
+    assert q.shape[0] == positions.shape[0] == out_cache_loc.shape[0]
+    assert q.shape[1] % head_dim == 0
+    assert k.shape[1] % head_dim == 0
+    assert v.shape[1] == k.shape[1]
+    assert idx_q.shape[1] % head_dim == 0
+    assert idx_k.shape[1] == head_dim
+    assert rotary_dim <= head_dim and rotary_dim % 2 == 0
+
+    q_heads = q.shape[1] // head_dim
+    k_heads = k.shape[1] // head_dim
+    idx_q_heads = idx_q.shape[1] // head_dim
+    assert k_cache.shape[1] == v_cache.shape[1] == k_heads
+    assert idx_k_cache.shape[1] == 1
+
+    q_out = torch.empty(q.shape, dtype=q.dtype, device=q.device)
+    k_out = torch.empty(k.shape, dtype=k.dtype, device=k.device)
+    idx_q_out = torch.empty(idx_q.shape, dtype=idx_q.dtype, device=idx_q.device)
+    idx_k_out = torch.empty(idx_k.shape, dtype=idx_k.dtype, device=idx_k.device)
+    block_hd = triton.next_power_of_2(head_dim)
+
+    _sparse_qk_index_gemma_rmsnorm_rope_cache_kernel[
+        (q.shape[0], q_heads + k_heads + idx_q_heads + 1)
+    ](
+        q,
+        k,
+        v,
+        idx_q,
+        idx_k,
+        q_out,
+        k_out,
+        idx_q_out,
+        idx_k_out,
+        k_cache,
+        v_cache,
+        idx_k_cache,
+        out_cache_loc,
+        q_weight,
+        k_weight,
+        idx_q_weight,
+        idx_k_weight,
+        positions,
+        cos_sin_cache,
+        q.stride(0),
+        q.stride(1),
+        k.stride(0),
+        k.stride(1),
+        v.stride(0),
+        v.stride(1),
+        idx_q.stride(0),
+        idx_q.stride(1),
+        idx_k.stride(0),
+        idx_k.stride(1),
+        k_cache.stride(0),
+        k_cache.stride(1),
+        k_cache.stride(2),
+        v_cache.stride(0),
+        v_cache.stride(1),
+        v_cache.stride(2),
+        idx_k_cache.stride(0),
+        idx_k_cache.stride(1),
+        idx_k_cache.stride(2),
+        q_heads,
+        k_heads,
+        idx_q_heads,
+        head_dim,
+        rotary_dim,
+        eps,
+        is_neox_style,
+        BLOCK_HD=block_hd,
+        num_warps=4,
+    )
+    return q_out, k_out, idx_q_out, idx_k_out

python/sglang/jit_kernel/minimax_m3/rmsnorm.py

@@ -0,0 +1,148 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Fused Gemma RMSNorm Triton kernels for MiniMax-M3 on AMD ROCm.
+
+Gemma RMSNorm = ``normalize(x) * (1 + weight)``, computed in a single fp32 pass.
+On ROCm with AITER, ``GemmaRMSNorm.forward_hip`` otherwise falls back to a
+~8-op PyTorch sequence: ``sgl_kernel``'s Gemma kernels are CUDA-only, and
+AITER's ``rmsnorm2d_fwd`` requires weight.dtype == activation.dtype (fp32
+weight + bf16 activation silently corrupts on gfx950). These kernels read
+strided inputs, so they serve both the full-hidden norms and the per-head
+q/k/index norms (non-contiguous ``qkv.split`` views).
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _gemma_rmsnorm_kernel(
+    x_ptr,
+    w_ptr,
+    out_ptr,
+    n_cols,
+    stride_row,
+    stride_col,
+    eps,
+    BLOCK_N: tl.constexpr,
+):
+    row = tl.program_id(0)
+    cols = tl.arange(0, BLOCK_N)
+    mask = cols < n_cols
+    x = tl.load(x_ptr + row * stride_row + cols * stride_col, mask=mask, other=0.0).to(
+        tl.float32
+    )
+    var = tl.sum(x * x, axis=0) / n_cols
+    rstd = 1.0 / tl.sqrt(var + eps)
+    w = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+    out = x * rstd * (1.0 + w)
+    tl.store(
+        out_ptr + row * n_cols + cols,
+        out.to(out_ptr.dtype.element_ty),
+        mask=mask,
+    )
+
+
+@triton.jit
+def _gemma_fused_add_rmsnorm_kernel(
+    x_ptr,
+    res_ptr,
+    w_ptr,
+    out_ptr,
+    res_out_ptr,
+    n_cols,
+    stride_xrow,
+    stride_xcol,
+    stride_rrow,
+    stride_rcol,
+    eps,
+    BLOCK_N: tl.constexpr,
+):
+    row = tl.program_id(0)
+    cols = tl.arange(0, BLOCK_N)
+    mask = cols < n_cols
+    x = tl.load(
+        x_ptr + row * stride_xrow + cols * stride_xcol, mask=mask, other=0.0
+    ).to(tl.float32)
+    r = tl.load(
+        res_ptr + row * stride_rrow + cols * stride_rcol, mask=mask, other=0.0
+    ).to(tl.float32)
+    s = x + r
+    # residual_out is the pre-norm sum (consumed by the next layer's add).
+    tl.store(
+        res_out_ptr + row * n_cols + cols,
+        s.to(res_out_ptr.dtype.element_ty),
+        mask=mask,
+    )
+    var = tl.sum(s * s, axis=0) / n_cols
+    rstd = 1.0 / tl.sqrt(var + eps)
+    w = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+    out = s * rstd * (1.0 + w)
+    tl.store(
+        out_ptr + row * n_cols + cols,
+        out.to(out_ptr.dtype.element_ty),
+        mask=mask,
+    )
+
+
+def _num_warps(block_n: int) -> int:
+    if block_n >= 4096:
+        return 16
+    if block_n >= 1024:
+        return 8
+    return 4
+
+
+def gemma_rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
+    """Gemma RMSNorm = normalize(x) * (1 + weight), fp32 math, single pass."""
+    orig_shape = x.shape
+    n = orig_shape[-1]
+    x2 = x.reshape(-1, n)
+    m = x2.shape[0]
+    out = torch.empty((m, n), dtype=x.dtype, device=x.device)
+    block_n = triton.next_power_of_2(n)
+    _gemma_rmsnorm_kernel[(m,)](
+        x2,
+        weight,
+        out,
+        n,
+        x2.stride(0),
+        x2.stride(1),
+        eps,
+        BLOCK_N=block_n,
+        num_warps=_num_warps(block_n),
+    )
+    return out.reshape(orig_shape)
+
+
+def gemma_fused_add_rmsnorm(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+):
+    """Fused (x + residual) then Gemma RMSNorm; returns (normed, pre-norm sum)."""
+    orig_shape = x.shape
+    n = orig_shape[-1]
+    x2 = x.reshape(-1, n)
+    r2 = residual.reshape(-1, n)
+    m = x2.shape[0]
+    out = torch.empty((m, n), dtype=x.dtype, device=x.device)
+    res_out = torch.empty((m, n), dtype=x.dtype, device=x.device)
+    block_n = triton.next_power_of_2(n)
+    _gemma_fused_add_rmsnorm_kernel[(m,)](
+        x2,
+        r2,
+        weight,
+        out,
+        res_out,
+        n,
+        x2.stride(0),
+        x2.stride(1),
+        r2.stride(0),
+        r2.stride(1),
+        eps,
+        BLOCK_N=block_n,
+        num_warps=_num_warps(block_n),
+    )
+    return out.reshape(orig_shape), res_out.reshape(orig_shape)

python/sglang/jit_kernel/minimax_m3/swiglu.py

@@ -0,0 +1,201 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Fused SwiGLU-OAI (split layout) Triton kernel for MiniMax-M3 on AMD ROCm.
+
+SwiGLU-OAI on a ``[*, 2I]`` split-layout tensor (gate = first half, up = second
+half): ``gate * sigmoid(alpha * gate) * (up + beta)`` with optional clamp,
+computed in fp32. Used by the dense MLP / shared experts ``swigluoai``
+activation on ROCm in place of the ``@torch.compile`` bf16 elementwise variant.
+"""
+
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _swiglu_oai_kernel(
+    g_ptr,
+    out_ptr,
+    n_inter,
+    stride_gm,
+    stride_gn,
+    stride_om,
+    stride_on,
+    alpha,
+    beta,
+    limit,
+    HAS_LIMIT: tl.constexpr,
+    BLOCK_I: tl.constexpr,
+):
+    row = tl.program_id(0)
+    pid_i = tl.program_id(1)
+    cols = pid_i * BLOCK_I + tl.arange(0, BLOCK_I)
+    mask = cols < n_inter
+    gate = tl.load(g_ptr + row * stride_gm + cols * stride_gn, mask=mask, other=0.0).to(
+        tl.float32
+    )
+    up = tl.load(
+        g_ptr + row * stride_gm + (n_inter + cols) * stride_gn,
+        mask=mask,
+        other=0.0,
+    ).to(tl.float32)
+    if HAS_LIMIT:
+        gate = tl.minimum(gate, limit)
+        up = tl.minimum(tl.maximum(up, -limit), limit)
+    out = gate * tl.sigmoid(alpha * gate) * (up + beta)
+    tl.store(
+        out_ptr + row * stride_om + cols * stride_on,
+        out.to(out_ptr.dtype.element_ty),
+        mask=mask,
+    )
+
+
+def swiglu_oai_split(
+    gate_up: torch.Tensor,
+    alpha: float,
+    beta: float,
+    limit: Optional[float],
+    out_dtype: Optional[torch.dtype] = None,
+) -> torch.Tensor:
+    """SwiGLU-OAI on a split-layout ``[*, 2I]`` tensor -> ``[*, I]`` (fp32 math)."""
+    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)
+    # Adaptive tile (tuned on gfx950). A 512-wide tile only helps
+    # once the (TP-sharded) per-rank slice is large enough to be bandwidth-bound
+    # (~1.25-1.35x faster than 256 at TP=1 prefill for the dense MLP I=12288).
+    # For small sharded slices (high TP) / decode the kernel is launch-bound, so
+    # fall back to 256. num_warps is pinned to 4 (8 underfills this tile).
+    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)
+
+
+@triton.jit
+def _swiglu_oai_mxfp8_quant_kernel(
+    g_ptr,
+    q_ptr,
+    scale_ptr,
+    n_inter,
+    stride_gm,
+    stride_gn,
+    stride_qm,
+    stride_qn,
+    stride_sm,
+    stride_sn,
+    alpha,
+    beta,
+    limit,
+    HAS_LIMIT: tl.constexpr,
+    BLOCK_I: tl.constexpr,
+):
+    row = tl.program_id(0)
+    pid_i = tl.program_id(1)
+    cols = pid_i * BLOCK_I + tl.arange(0, BLOCK_I)
+    mask = cols < n_inter
+
+    gate = tl.load(g_ptr + row * stride_gm + cols * stride_gn, mask=mask, other=0.0)
+    up = tl.load(
+        g_ptr + row * stride_gm + (n_inter + cols) * stride_gn,
+        mask=mask,
+        other=0.0,
+    )
+    gate = gate.to(tl.float32)
+    up = up.to(tl.float32)
+    if HAS_LIMIT:
+        gate = tl.minimum(gate, limit)
+        up = tl.minimum(tl.maximum(up, -limit), limit)
+
+    # Keep the activation in fp32 all the way to the E8M0 scale selection (no
+    # bf16 round-trip to HBM). Matches the vLLM/ame fused swiglu+quant kernel:
+    # marginally more accurate than the unfused bf16 two-kernel chain.
+    activated = gate * tl.sigmoid(alpha * gate) * (up + beta)
+
+    groups: tl.constexpr = BLOCK_I // 32
+    activated_2d = tl.reshape(activated, (groups, 32))
+    valid_groups = pid_i * groups + tl.arange(0, groups) < (n_inter // 32)
+
+    amax = tl.maximum(tl.max(tl.abs(activated_2d), axis=1), 1e-30)
+    # Round the E8M0 exponent up (ceil(log2(amax / e4m3_max))) so the block amax
+    # stays inside the e4m3 range and the full dynamic range is used.
+    scale_biased = tl.ceil(tl.log2(amax / 448.0)) + 127.0
+    scale_biased = tl.minimum(tl.maximum(scale_biased, 0.0), 254.0)
+    descale = tl.reshape(tl.exp2(scale_biased - 127.0), (groups, 1))
+
+    q_2d = tl.clamp(activated_2d / descale, -448.0, 448.0)
+    q = tl.reshape(q_2d, (BLOCK_I,)).to(q_ptr.dtype.element_ty)
+
+    tl.store(q_ptr + row * stride_qm + cols * stride_qn, q, mask=mask)
+    tl.store(
+        scale_ptr
+        + row * stride_sm
+        + (pid_i * groups + tl.arange(0, groups)) * stride_sn,
+        scale_biased.to(tl.uint8),
+        mask=valid_groups,
+    )
+
+
+def swiglu_oai_mxfp8_quant(
+    gate_up: torch.Tensor,
+    alpha: float,
+    beta: float,
+    limit: Optional[float],
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """SwiGLU-OAI on split layout, then MiniMax MXFP8 quant, in one launch.
+
+    The activation stays in fp32 through the E8M0 scale selection (no bf16
+    round-trip), matching the vLLM/ame fused swiglu+quant kernel.
+    """
+    orig_shape = gate_up.shape
+    two_i = orig_shape[-1]
+    n_inter = two_i // 2
+    assert n_inter % 32 == 0, "MiniMax MXFP8 quant requires I divisible by 32."
+
+    x2 = gate_up.reshape(-1, two_i)
+    m = x2.shape[0]
+    q = torch.empty((m, n_inter), dtype=torch.float8_e4m3fn, device=gate_up.device)
+    scales = torch.empty((m, n_inter // 32), dtype=torch.uint8, device=gate_up.device)
+    block_i = 512 if n_inter >= 2048 else 256
+    grid = (m, triton.cdiv(n_inter, block_i))
+    _swiglu_oai_mxfp8_quant_kernel[grid](
+        x2,
+        q,
+        scales,
+        n_inter,
+        x2.stride(0),
+        x2.stride(1),
+        q.stride(0),
+        q.stride(1),
+        scales.stride(0),
+        scales.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 q.reshape(*orig_shape[:-1], n_inter), scales.reshape(
+        *orig_shape[:-1], n_inter // 32
+    )