chore: import upstream snapshot with attribution

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
+    )