chore: import upstream snapshot with attribution

python/sglang/srt/layers/moe/rocm_moe_utils.py

@@ -0,0 +1,329 @@
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.9.1rc2/vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from enum import IntEnum
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.utils import get_bool_env_var, is_hip
+from sglang.srt.utils.custom_op import register_custom_op
+
+_is_hip = is_hip()
+_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
+
+
+class ActivationMethod(IntEnum):
+    # This allows interfacing with AITER ActivationType enum
+    # without importing the ActivationType enum from AITER globally.
+    SILU = 0
+    GELU = 1
+
+
+# NOTE: for non _use_aiter case, use lazy registration to avoid overhead
+# (registration may not be trigger actually, since it will not be called)
+@register_custom_op(out_shape="hidden_states", eager=_use_aiter)
+def rocm_aiter_asm_moe_tkw1(
+    hidden_states: torch.Tensor,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    fc1_scale: Optional[torch.Tensor] = None,
+    fc2_scale: Optional[torch.Tensor] = None,
+    fc1_smooth_scale: Optional[torch.Tensor] = None,
+    fc2_smooth_scale: Optional[torch.Tensor] = None,
+    a16: bool = False,
+    per_tensor_quant_scale: Optional[torch.Tensor] = None,
+    expert_mask: Optional[torch.Tensor] = None,
+    activation_method: int = ActivationMethod.SILU.value,
+) -> torch.Tensor:
+
+    from aiter import ActivationType
+    from aiter.fused_moe_bf16_asm import asm_moe_tkw1
+
+    activation = ActivationType(activation_method)
+
+    return asm_moe_tkw1(
+        hidden_states,
+        w1,
+        w2,
+        topk_weights,
+        topk_ids,
+        fc1_scale=fc1_scale,
+        fc2_scale=fc2_scale,
+        fc1_smooth_scale=fc1_smooth_scale,
+        fc2_smooth_scale=fc2_smooth_scale,
+        a16=a16,
+        per_tensor_quant_scale=per_tensor_quant_scale,
+        expert_mask=expert_mask,
+        activation=activation,
+    )
+
+
+def rocm_fused_experts_tkw1(
+    hidden_states: torch.Tensor,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    activation: str = "silu",
+    apply_router_weight_on_input: bool = False,
+    use_fp8_w8a8: bool = False,
+    per_channel_quant: bool = False,
+    w1_scale: Optional[torch.Tensor] = None,
+    w2_scale: Optional[torch.Tensor] = None,
+    a1_scale: Optional[torch.Tensor] = None,
+    a2_scale: Optional[torch.Tensor] = None,
+    block_shape: Optional[list[int]] = None,
+) -> torch.Tensor:
+
+    activation_method = (
+        ActivationMethod.SILU if activation == "silu" else ActivationMethod.GELU
+    )
+    # All AITER Fused MoE kernels are expecting the following datatypes
+    topk_weights = topk_weights.to(torch.float32)
+    topk_ids = topk_ids.to(torch.int32)
+
+    # w8a8 per-channel quantization
+    if per_channel_quant and apply_router_weight_on_input and use_fp8_w8a8:
+        # AITER tkw1 kernel for FP8 models with `apply_router_weight_on_input`
+        # This applies topk_weights on the GEMM output of the first FC layer
+        #  rather than the second FC.
+        assert (
+            topk_weights.dim() == 2
+        ), "`topk_weights` should be in shape (num_tokens, topk)"
+        assert topk_weights.shape[-1] == 1, (
+            "Only support topk=1 when" " `apply_router_weight_on_input` is True"
+        )
+
+        return rocm_aiter_asm_moe_tkw1(
+            hidden_states,
+            w1,
+            w2,
+            topk_weights,
+            topk_ids,
+            fc1_scale=w1_scale,
+            fc2_scale=w2_scale,
+            fc1_smooth_scale=None,
+            fc2_smooth_scale=None,
+            a16=False,
+            per_tensor_quant_scale=None,
+            expert_mask=None,
+            activation_method=activation_method,
+        )
+    else:
+        assert False, "This should not be called."
+
+
+@triton.jit
+def upscale_kernel(
+    A_ptr,  # *fp16 / *fp32
+    scale_ptr,  # *fp16 / *fp32
+    Out_ptr,  # *fp16 / *fp32
+    M,
+    N,
+    recv_token_num,
+    stride_am,
+    stride_an,
+    stride_sm,
+    stride_sn,
+    stride_om,
+    stride_on,
+    BLOCK_N: tl.constexpr,
+):
+    pid_m = tl.program_id(0)  # row id
+    pid_n = tl.program_id(1)  # block id along N
+
+    recv_token_num_val = tl.load(recv_token_num)
+
+    if pid_m >= recv_token_num_val:
+        return
+
+    # column offsets
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    mask = offs_n < N
+
+    # A[m, n]
+    a_ptrs = A_ptr + pid_m * stride_am + offs_n * stride_an
+    a = tl.load(a_ptrs, mask=mask, other=0.0)
+
+    # scale index: n // 128
+    scale_idx = offs_n // 128
+    s_ptrs = scale_ptr + pid_m * stride_sm + scale_idx * stride_sn
+    s = tl.load(s_ptrs, mask=mask, other=1.0)
+
+    out = a * s
+
+    out_ptrs = Out_ptr + pid_m * stride_om + offs_n * stride_on
+    tl.store(out_ptrs, out, mask=mask)
+
+
+def upscale(hidden_state, hidden_state_scale, recv_token_num, output_dtype):
+    M, N = hidden_state.shape
+
+    Out = torch.empty_like(hidden_state, dtype=output_dtype)
+
+    BLOCK_N = 256
+
+    grid = (M, triton.cdiv(N, BLOCK_N))
+
+    upscale_kernel[grid](
+        hidden_state,
+        hidden_state_scale,
+        Out,
+        M,
+        N,
+        recv_token_num,
+        hidden_state.stride(0),
+        hidden_state.stride(1),
+        hidden_state_scale.stride(0),
+        hidden_state_scale.stride(1),
+        Out.stride(0),
+        Out.stride(1),
+        BLOCK_N=BLOCK_N,
+    )
+
+    return Out
+
+
+@triton.jit
+def upscale_fp4x2_block32_kernel(
+    A_u8_ptr,  # *uint8  (view from float4_e2m1fn_x2)
+    S_u8_ptr,  # *uint8  (view from float8_e8m0fnu), shape (M, N_fp4/32)
+    Out_ptr,  # *fp16/fp32/bf16, shape (M, N_fp4)
+    N_FP4: tl.constexpr,
+    recv_token_num,
+    stride_am,
+    stride_an,  # A strides (in uint8 elements) for (M, packed_N)
+    stride_sm,
+    stride_sn,  # S strides (in uint8 elements) for (M, N_FP4/32)
+    stride_om,
+    stride_on,  # Out strides (in output elements) for (M, N_FP4)
+    BLOCK_N: tl.constexpr,
+    OUT_DTYPE: tl.constexpr,  # tl.float16 / tl.float32 / tl.bfloat16
+):
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+
+    recv_token_num_val = tl.load(recv_token_num)
+    if pid_m >= recv_token_num_val:
+        return
+
+    offs = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    mask = offs < N_FP4
+
+    # --------------------------
+    # Load packed fp4x2 byte
+    # --------------------------
+    byte_idx = offs >> 1  # offs // 2
+    is_hi = (offs & 1) != 0  # select high nibble?
+
+    a_ptrs = A_u8_ptr + pid_m * stride_am + byte_idx * stride_an
+    a_byte = tl.load(a_ptrs, mask=mask, other=0).to(tl.int32)
+
+    lo = a_byte & 0xF
+    hi = (a_byte >> 4) & 0xF
+    code = tl.where(is_hi, hi, lo).to(tl.int32)  # 0..15
+
+    # --------------------------
+    # Decode float4_e2m1fn
+    # layout: [sign|exp(2)|mant(1)]
+    # bias=1, finite-only
+    # --------------------------
+    sign = (code >> 3) & 0x1
+    exp = (code >> 1) & 0x3
+    mant = code & 0x1
+
+    mant_f = mant.to(tl.float32) * 0.5
+    is_sub = exp == 0
+
+    # normal: 2^(exp-bias) * (1 + mant/2), bias=1
+    e_norm = (exp - 1).to(tl.float32)
+    val_norm = tl.exp2(e_norm) * (1.0 + mant_f)
+
+    # subnorm/zero: mant/2 * 2^(1-bias) = mant/2
+    val_sub = mant_f
+
+    val = tl.where(is_sub, val_sub, val_norm)
+    val = tl.where(sign != 0, -val, val)  # apply sign
+
+    # --------------------------
+    # Per-token block32 scale: scale_idx = offs // 32
+    # scale dtype: float8_e8m0fnu stored in uint8
+    # decode: e==0 -> 0
+    #         e in [1..254] -> 2^(e-127)
+    #         e==255 -> clamp to 254
+    # --------------------------
+    scale_idx = offs >> 5  # offs // 32
+
+    s_ptrs = S_u8_ptr + pid_m * stride_sm + scale_idx * stride_sn
+    e = tl.load(s_ptrs, mask=mask, other=0).to(tl.int32)
+
+    e = tl.minimum(e, 254)  # clamp 255->254
+    is_zero = e == 0
+    exp_s = (e - 127).to(tl.float32)
+    s = tl.exp2(exp_s)
+    s = tl.where(is_zero, 0.0, s)
+
+    out = (val * s).to(OUT_DTYPE)
+
+    out_ptrs = Out_ptr + pid_m * stride_om + offs * stride_on
+    tl.store(out_ptrs, out, mask=mask)
+
+
+def upscale_mxfp4(hidden_state, hidden_state_scale, recv_token_num, output_dtype):
+    """
+    hidden_state: (M, packed_N) torch.float4_e2m1fn_x2
+    hidden_state_scale: (M, packed_N*2/32) = (M, N_fp4/32) torch.float8_e8m0fnu
+    output: (M, N_fp4) output_dtype
+    """
+    assert hidden_state.dtype == torch.float4_e2m1fn_x2, hidden_state.dtype
+    assert hidden_state_scale.dtype == torch.float8_e8m0fnu, hidden_state_scale.dtype
+    assert hidden_state.is_contiguous() or True  # stride-based load OK
+
+    M, packed_N = hidden_state.shape
+    N_fp4 = packed_N * 2
+
+    # scale second dim must be N_fp4/32
+    assert hidden_state_scale.shape[0] == M
+    assert hidden_state_scale.shape[1] == (N_fp4 // 32), (
+        hidden_state_scale.shape,
+        N_fp4,
+    )
+
+    # Triton doesn't (reliably) accept torch.float4/float8 pointers directly.
+    # Use raw uint8 views.
+    A_u8 = hidden_state.view(torch.uint8)
+    S_u8 = hidden_state_scale.view(torch.uint8)
+
+    Out = torch.empty((M, N_fp4), dtype=output_dtype, device=hidden_state.device)
+
+    BLOCK_N = 256
+    grid = (M, triton.cdiv(N_fp4, BLOCK_N))
+
+    OUT_TL = (
+        tl.float16
+        if output_dtype == torch.float16
+        else tl.bfloat16 if output_dtype == torch.bfloat16 else tl.float32
+    )
+
+    upscale_fp4x2_block32_kernel[grid](
+        A_u8,
+        S_u8,
+        Out,
+        N_FP4=N_fp4,
+        recv_token_num=recv_token_num,
+        stride_am=A_u8.stride(0),
+        stride_an=A_u8.stride(1),
+        stride_sm=S_u8.stride(0),
+        stride_sn=S_u8.stride(1),
+        stride_om=Out.stride(0),
+        stride_on=Out.stride(1),
+        BLOCK_N=BLOCK_N,
+        OUT_DTYPE=OUT_TL,
+        num_warps=4,
+    )
+    return Out