[MoE Backend] Add HPC-Ops FP8 MoE runner backend (#30541)
Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Halcyon <56064364+VAthree@users.noreply.github.com>
This commit is contained in:
co-authored by
Claude Fable 5
Halcyon
parent
4d5917e744
commit
3d91a569ce
@@ -155,6 +155,24 @@ class FusedMoeWeightScaleSupported(Enum):
|
||||
BLOCK = "block"
|
||||
|
||||
|
||||
def _validate_hpc_ops_quant_method(quant_method) -> None:
|
||||
"""--moe-runner-backend hpc_ops makes the standard dispatcher keep global
|
||||
expert ids for every MoE layer, so the resolved quant method must be the
|
||||
FP8 one the hpc_ops runner supports. Quant methods that never construct a
|
||||
MoeRunner (e.g. W4AFp8 calls its kernel directly from apply()) bypass the
|
||||
MoeRunner-level guard, so validate here at layer init.
|
||||
"""
|
||||
if get_moe_runner_backend().is_hpc_ops() and not isinstance(
|
||||
quant_method, Fp8MoEMethod
|
||||
):
|
||||
raise ValueError(
|
||||
"--moe-runner-backend hpc_ops only supports Fp8MoEMethod "
|
||||
"(FP8 blockwise or per-tensor MoE), but this layer selected "
|
||||
f"{type(quant_method).__name__}. Remove --moe-runner-backend "
|
||||
"hpc_ops for this model."
|
||||
)
|
||||
|
||||
|
||||
class FusedMoE(torch.nn.Module):
|
||||
"""FusedMoE layer for MoE models.
|
||||
|
||||
@@ -331,6 +349,7 @@ class FusedMoE(torch.nn.Module):
|
||||
self.use_flashinfer_trtllm_moe,
|
||||
self.use_deep_gemm,
|
||||
)
|
||||
_validate_hpc_ops_quant_method(self.quant_method)
|
||||
self.supports_deferred_finalize = (
|
||||
envs.SGLANG_ENABLE_MOE_DEFERRED_FINALIZE.get()
|
||||
and get_moe_runner_backend().is_flashinfer_trtllm()
|
||||
|
||||
@@ -0,0 +1,208 @@
|
||||
from __future__ import annotations
|
||||
|
||||
"""
|
||||
MoE runner backend powered by HPC-Ops (https://github.com/Tencent/hpc-ops),
|
||||
a production-grade operator library for LLM inference developed by the
|
||||
Tencent Hunyuan AI Infra team.
|
||||
|
||||
The backend wraps the monolithic FP8 fused-MoE kernels ``fuse_moe_blockwise``
|
||||
(128x128 block-quantized weights + per-token-group-128 activations, e.g.
|
||||
Qwen3-FP8 style checkpoints) and ``fuse_moe`` (per-tensor weight and static
|
||||
per-tensor activation quantization, e.g. Hy3-FP8 style checkpoints). Both
|
||||
kernels fuse
|
||||
gather -> grouped gate_up GEMM -> SiLU-and-mul -> grouped down GEMM -> weighted
|
||||
reduce into one call and consume *global* top-k expert ids together with
|
||||
``rank_ep`` / ``num_expert_total``, so expert parallelism with contiguous
|
||||
expert partitioning works without a local-expert remap.
|
||||
|
||||
Only supported on NVIDIA Hopper / Blackwell (sm90+). Note that the HPC-Ops
|
||||
kernels are currently tuned primarily for H20: on other GPUs (H100/H200/B200,
|
||||
...) the speedup over the default MoE runner may be limited or absent. Enable
|
||||
it explicitly with ``--moe-runner-backend hpc_ops``.
|
||||
"""
|
||||
|
||||
import functools
|
||||
import importlib.util
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, List, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.moe.moe_runner.base import MoeQuantInfo, register_fused_func
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
|
||||
from sglang.srt.layers.moe.token_dispatcher.standard import (
|
||||
StandardCombineInput,
|
||||
StandardDispatchOutput,
|
||||
)
|
||||
|
||||
# The HPC-Ops group GEMM tiles N by 128 and the blockwise path quantizes
|
||||
# activations in groups of 128, so every participating dim must be 128-aligned.
|
||||
HPC_OPS_BLOCK_SIZE = 128
|
||||
# The K dim of the blockwise weight scales must be padded to a multiple of 4
|
||||
# (see hpc-ops tests: (k // 128 + 3) // 4 * 4).
|
||||
_SCALE_K_ALIGN = 4
|
||||
|
||||
|
||||
@functools.cache
|
||||
def has_hpc_ops() -> bool:
|
||||
"""Return True if the ``hpc`` package (HPC-Ops) is installed."""
|
||||
return importlib.util.find_spec("hpc") is not None
|
||||
|
||||
|
||||
def pad_hpc_ops_block_scale(scale: torch.Tensor) -> torch.Tensor:
|
||||
"""Pad the K dim (last dim) of a [E, N/128, K/128] block scale to %4."""
|
||||
k = scale.shape[-1]
|
||||
k_pad = (k + _SCALE_K_ALIGN - 1) // _SCALE_K_ALIGN * _SCALE_K_ALIGN
|
||||
if k == k_pad:
|
||||
return scale.contiguous()
|
||||
padded = scale.new_zeros((*scale.shape[:-1], k_pad))
|
||||
padded[..., :k].copy_(scale)
|
||||
return padded
|
||||
|
||||
|
||||
@dataclass
|
||||
class HpcOpsMoeQuantInfo(MoeQuantInfo):
|
||||
"""Quant payload for the HPC-Ops fused MoE kernels.
|
||||
|
||||
``block_quant`` selects between the two kernels:
|
||||
- True: ``fuse_moe_blockwise`` with ``w13_weight_scale_inv`` /
|
||||
``w2_weight_scale_inv`` ([E, N/128, K/128], K dim padded to %4) and
|
||||
dynamic per-token-group-128 activation quantization.
|
||||
- False: ``fuse_moe`` with per-expert dequant alphas
|
||||
``gate_up_alphas = w13_weight_scale * w13_input_scale`` ([E]),
|
||||
``down_alphas = w2_weight_scale * w2_input_scale`` ([E]) and the static
|
||||
activation scales ``w13_input_scale`` / ``w2_input_scale`` (scalars).
|
||||
"""
|
||||
|
||||
w13_weight: torch.Tensor
|
||||
w2_weight: torch.Tensor
|
||||
block_quant: bool
|
||||
global_num_experts: int
|
||||
moe_ep_rank: int
|
||||
# Blockwise path
|
||||
w13_weight_scale_inv: Optional[torch.Tensor] = None
|
||||
w2_weight_scale_inv: Optional[torch.Tensor] = None
|
||||
block_shape: Optional[List[int]] = None
|
||||
# Per-tensor path
|
||||
gate_up_alphas: Optional[torch.Tensor] = None
|
||||
down_alphas: Optional[torch.Tensor] = None
|
||||
w13_input_scale: Optional[torch.Tensor] = None
|
||||
w2_input_scale: Optional[torch.Tensor] = None
|
||||
|
||||
|
||||
def _check_runner_config_supported(runner_config: MoeRunnerConfig) -> None:
|
||||
if runner_config.activation != "silu" or not runner_config.is_gated:
|
||||
raise ValueError(
|
||||
"The hpc_ops MoE runner backend only supports the gated silu "
|
||||
f"activation, got activation={runner_config.activation}, "
|
||||
f"is_gated={runner_config.is_gated}."
|
||||
)
|
||||
if runner_config.num_fused_shared_experts != 0:
|
||||
raise ValueError(
|
||||
"The hpc_ops MoE runner backend does not support fused shared experts."
|
||||
)
|
||||
if runner_config.apply_router_weight_on_input:
|
||||
raise ValueError(
|
||||
"The hpc_ops MoE runner backend does not support "
|
||||
"apply_router_weight_on_input."
|
||||
)
|
||||
if runner_config.no_combine:
|
||||
raise ValueError(
|
||||
"The hpc_ops MoE runner backend does not support no_combine "
|
||||
"(the fused kernel always reduces over top-k experts)."
|
||||
)
|
||||
if (
|
||||
runner_config.gemm1_alpha is not None
|
||||
or runner_config.gemm1_clamp_limit is not None
|
||||
or runner_config.swiglu_limit is not None
|
||||
):
|
||||
raise ValueError(
|
||||
"The hpc_ops MoE runner backend runs a plain SiLU-and-mul; it does "
|
||||
"not support gemm1_alpha / gemm1_clamp_limit / swiglu_limit."
|
||||
)
|
||||
|
||||
|
||||
@register_fused_func("none", "hpc_ops")
|
||||
def fused_experts_none_to_hpc_ops(
|
||||
dispatch_output: StandardDispatchOutput,
|
||||
quant_info: HpcOpsMoeQuantInfo,
|
||||
runner_config: MoeRunnerConfig,
|
||||
) -> StandardCombineInput:
|
||||
import hpc
|
||||
|
||||
from sglang.kernels.ops.quantization.fp8_kernel import (
|
||||
scaled_fp8_quant,
|
||||
sglang_per_token_group_quant_fp8,
|
||||
)
|
||||
from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
|
||||
|
||||
if not isinstance(quant_info, HpcOpsMoeQuantInfo):
|
||||
raise ValueError(
|
||||
"The hpc_ops MoE runner backend only supports FP8-quantized MoE "
|
||||
"models (Fp8MoEMethod); got quant info "
|
||||
f"{type(quant_info).__name__}. Note that with expert parallelism "
|
||||
"this backend also expects global top-k ids, so other quant "
|
||||
"methods must not run with --moe-runner-backend hpc_ops."
|
||||
)
|
||||
assert (
|
||||
quant_info.w13_weight.dtype == torch.float8_e4m3fn
|
||||
), f"expected fp8 w13_weight, got {quant_info.w13_weight.dtype}"
|
||||
assert (
|
||||
quant_info.w2_weight.dtype == torch.float8_e4m3fn
|
||||
), f"expected fp8 w2_weight, got {quant_info.w2_weight.dtype}"
|
||||
_check_runner_config_supported(runner_config)
|
||||
|
||||
x = dispatch_output.hidden_states
|
||||
topk_weights, topk_ids, _ = dispatch_output.topk_output
|
||||
|
||||
assert x.dtype == torch.bfloat16, (
|
||||
"The hpc_ops MoE runner backend only supports bf16 hidden states, "
|
||||
f"got {x.dtype}."
|
||||
)
|
||||
|
||||
topk_ids = topk_ids.to(torch.int32)
|
||||
topk_weights = topk_weights.to(torch.float32)
|
||||
|
||||
if quant_info.block_quant:
|
||||
assert quant_info.block_shape == [
|
||||
HPC_OPS_BLOCK_SIZE,
|
||||
HPC_OPS_BLOCK_SIZE,
|
||||
], (
|
||||
"The hpc_ops MoE runner backend only supports 128x128 block "
|
||||
f"quantization, got {quant_info.block_shape}."
|
||||
)
|
||||
x_q, x_scale = sglang_per_token_group_quant_fp8(x, HPC_OPS_BLOCK_SIZE)
|
||||
output = hpc.fuse_moe_blockwise(
|
||||
x_q,
|
||||
x_scale,
|
||||
quant_info.w13_weight,
|
||||
quant_info.w13_weight_scale_inv,
|
||||
quant_info.w2_weight,
|
||||
quant_info.w2_weight_scale_inv,
|
||||
topk_ids,
|
||||
topk_weights,
|
||||
quant_info.moe_ep_rank,
|
||||
quant_info.global_num_experts,
|
||||
)
|
||||
else:
|
||||
x_q, _ = scaled_fp8_quant(x, quant_info.w13_input_scale)
|
||||
act_and_mul_scale = 1.0 / quant_info.w2_input_scale.reshape(1)
|
||||
output = hpc.fuse_moe(
|
||||
x_q,
|
||||
quant_info.w13_weight,
|
||||
quant_info.w2_weight,
|
||||
quant_info.gate_up_alphas,
|
||||
quant_info.down_alphas,
|
||||
act_and_mul_scale,
|
||||
topk_ids,
|
||||
topk_weights,
|
||||
quant_info.moe_ep_rank,
|
||||
quant_info.global_num_experts,
|
||||
)
|
||||
|
||||
if runner_config.routed_scaling_factor is not None:
|
||||
output *= runner_config.routed_scaling_factor
|
||||
|
||||
return StandardCombineInput(hidden_states=output)
|
||||
@@ -12,7 +12,7 @@ from sglang.srt.layers.moe.moe_runner.base import (
|
||||
from sglang.srt.layers.moe.moe_runner.deep_gemm import DeepGemmRunnerCore
|
||||
from sglang.srt.layers.moe.moe_runner.triton import TritonRunnerCore
|
||||
from sglang.srt.layers.moe.moe_runner.triton_kernels import TritonKernelsRunnerCore
|
||||
from sglang.srt.layers.moe.utils import get_moe_a2a_backend
|
||||
from sglang.srt.layers.moe.utils import get_moe_a2a_backend, get_moe_runner_backend
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.batch_overlap.single_batch_overlap import DownGemmOverlapArgs
|
||||
@@ -35,6 +35,21 @@ class MoeRunner:
|
||||
self.config = config
|
||||
self.lora_enabled = lora_enabled
|
||||
|
||||
# --moe-runner-backend hpc_ops makes the standard dispatcher keep
|
||||
# global expert ids (skip_local_expert_mapping), so every MoE layer
|
||||
# must actually run the hpc_ops runner. A quant method that falls
|
||||
# back to another runner here (e.g. an unquantized MoE never enters
|
||||
# the FP8 path) would consume global ids as local ones and misroute
|
||||
# tokens under EP>1, so fail loudly at startup instead.
|
||||
if get_moe_runner_backend().is_hpc_ops() and not runner_backend.is_hpc_ops():
|
||||
raise ValueError(
|
||||
"--moe-runner-backend hpc_ops was requested, but this MoE "
|
||||
f"layer's quantization method selected the "
|
||||
f"'{runner_backend.value}' runner (hpc_ops only supports FP8 "
|
||||
"blockwise / per-tensor quantized MoE). Remove "
|
||||
"--moe-runner-backend hpc_ops for this model."
|
||||
)
|
||||
|
||||
self.fused_func = None
|
||||
|
||||
if runner_backend.is_triton():
|
||||
@@ -80,6 +95,11 @@ class MoeRunner:
|
||||
)
|
||||
elif runner_backend.is_cutlass():
|
||||
self.runner_core = None # CUTLASS uses the direct cutlass_moe_fp4 path
|
||||
elif runner_backend.is_hpc_ops():
|
||||
self.runner_core = None # HPC-Ops only supports the fused path
|
||||
# Import here (not at module top, to avoid a circular import) to
|
||||
# register the hpc_ops fused func before the pool lookup.
|
||||
from sglang.srt.layers.moe.moe_runner import hpc_ops # noqa: F401
|
||||
else:
|
||||
raise NotImplementedError(f"Unsupported runner backend: {runner_backend}")
|
||||
|
||||
|
||||
@@ -100,12 +100,14 @@ class StandardDispatcher(BaseDispatcher):
|
||||
# Skip local expert mapping when the backend handles EP with global expert IDs:
|
||||
# - cutlass / cutedsl / trtllm_routed handle EP internally
|
||||
# - mxfp4 dispatcher mapping is already global
|
||||
# - hpc_ops consumes global ids together with rank_ep / num_expert_total
|
||||
self.skip_local_expert_mapping = (
|
||||
backend.is_flashinfer_cutlass()
|
||||
or backend.is_flashinfer_cutedsl()
|
||||
or backend.is_flashinfer_trtllm()
|
||||
or backend.is_experimental_sgl_trtllm()
|
||||
or backend.is_flashinfer_trtllm_routed()
|
||||
or backend.is_hpc_ops()
|
||||
or self.enable_flashinfer_mxfp4_moe
|
||||
)
|
||||
self.num_experts = moe_runner_config.num_experts
|
||||
|
||||
@@ -105,10 +105,14 @@ class MoeRunnerBackend(Enum):
|
||||
HUMMING = "humming"
|
||||
EXPERIMENTAL_SGL_MARLIN = "experimental_sgl_marlin"
|
||||
AITER = "aiter"
|
||||
HPC_OPS = "hpc_ops"
|
||||
|
||||
def is_auto(self):
|
||||
return self == MoeRunnerBackend.AUTO
|
||||
|
||||
def is_hpc_ops(self):
|
||||
return self == MoeRunnerBackend.HPC_OPS
|
||||
|
||||
def is_deep_gemm(self):
|
||||
return self == MoeRunnerBackend.DEEP_GEMM
|
||||
|
||||
|
||||
@@ -2095,9 +2095,83 @@ class Fp8MoEMethod(FusedMoEMethodBase):
|
||||
|
||||
align_fp8_moe_weights_for_flashinfer_trtllm(layer)
|
||||
|
||||
if get_moe_runner_backend().is_hpc_ops():
|
||||
self._prepare_hpc_ops_weights(layer)
|
||||
|
||||
if hasattr(layer, "dispatcher"):
|
||||
layer.dispatcher.set_quant_config({"weight_dtype": layer.w13_weight.dtype})
|
||||
|
||||
def _prepare_hpc_ops_weights(self, layer: Module) -> None:
|
||||
"""Precompute the scale layouts consumed by the HPC-Ops fused MoE kernels.
|
||||
|
||||
- Blockwise FP8: the kernel wants [E, N/128, K/128] float32 dequant
|
||||
scales with the K dim padded to a multiple of 4.
|
||||
- Per-tensor FP8: the kernel wants per-expert dequant alphas
|
||||
(weight_scale * input_scale) and a static w2 input scale; this
|
||||
requires the static activation scheme.
|
||||
"""
|
||||
from sglang.srt.layers.moe.moe_runner.hpc_ops import pad_hpc_ops_block_scale
|
||||
|
||||
if self.block_quant:
|
||||
layer.hpc_ops_w13_weight_scale = pad_hpc_ops_block_scale(
|
||||
layer.w13_weight_scale_inv.data.float()
|
||||
)
|
||||
layer.hpc_ops_w2_weight_scale = pad_hpc_ops_block_scale(
|
||||
layer.w2_weight_scale_inv.data.float()
|
||||
)
|
||||
else:
|
||||
if layer.w13_input_scale is None or layer.w2_input_scale is None:
|
||||
raise ValueError(
|
||||
"The hpc_ops MoE runner backend requires static activation "
|
||||
"scales for per-tensor FP8 models (activation_scheme="
|
||||
"'static' in the checkpoint quantization config)."
|
||||
)
|
||||
layer.hpc_ops_gate_up_alphas = (
|
||||
layer.w13_weight_scale.data.float() * layer.w13_input_scale.data.float()
|
||||
)
|
||||
layer.hpc_ops_down_alphas = (
|
||||
layer.w2_weight_scale.data.float() * layer.w2_input_scale.data.float()
|
||||
)
|
||||
|
||||
def _get_hpc_ops_quant_info(self, layer: torch.nn.Module):
|
||||
from sglang.srt.layers.moe.moe_runner.hpc_ops import HpcOpsMoeQuantInfo
|
||||
|
||||
# The HPC-Ops fused kernels take no per-expert GEMM bias; refuse
|
||||
# instead of silently dropping it.
|
||||
if (
|
||||
getattr(layer, "w13_weight_bias", None) is not None
|
||||
or getattr(layer, "w2_weight_bias", None) is not None
|
||||
):
|
||||
raise ValueError(
|
||||
"The hpc_ops MoE runner backend does not support MoE GEMM "
|
||||
"biases (w13_weight_bias / w2_weight_bias); use another "
|
||||
"--moe-runner-backend for this model."
|
||||
)
|
||||
|
||||
if self.block_quant:
|
||||
return HpcOpsMoeQuantInfo(
|
||||
w13_weight=layer.w13_weight,
|
||||
w2_weight=layer.w2_weight,
|
||||
block_quant=True,
|
||||
global_num_experts=int(layer.num_experts),
|
||||
moe_ep_rank=int(layer.moe_ep_rank),
|
||||
w13_weight_scale_inv=layer.hpc_ops_w13_weight_scale,
|
||||
w2_weight_scale_inv=layer.hpc_ops_w2_weight_scale,
|
||||
block_shape=self.quant_config.weight_block_size,
|
||||
)
|
||||
else:
|
||||
return HpcOpsMoeQuantInfo(
|
||||
w13_weight=layer.w13_weight,
|
||||
w2_weight=layer.w2_weight,
|
||||
block_quant=False,
|
||||
global_num_experts=int(layer.num_experts),
|
||||
moe_ep_rank=int(layer.moe_ep_rank),
|
||||
gate_up_alphas=layer.hpc_ops_gate_up_alphas,
|
||||
down_alphas=layer.hpc_ops_down_alphas,
|
||||
w13_input_scale=layer.w13_input_scale,
|
||||
w2_input_scale=layer.w2_input_scale,
|
||||
)
|
||||
|
||||
def process_weights_hip_int4(self, layer: Module):
|
||||
# TODO: _use_aiter: add after triton kernel added
|
||||
# INT4-FP8 (INT4 MoE Weight, FP8 Compute)
|
||||
@@ -2195,6 +2269,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
|
||||
or moe_runner_backend.is_aiter()
|
||||
or moe_runner_backend.is_flashinfer_trtllm()
|
||||
or moe_runner_backend.is_flashinfer_trtllm_routed()
|
||||
or moe_runner_backend.is_hpc_ops()
|
||||
):
|
||||
self.runner = MoeRunner(moe_runner_backend, moe_runner_config)
|
||||
else:
|
||||
@@ -2452,6 +2527,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
|
||||
),
|
||||
activation_type=activation_type,
|
||||
)
|
||||
elif self.runner.runner_backend.is_hpc_ops():
|
||||
quant_info = self._get_hpc_ops_quant_info(layer)
|
||||
elif self.runner.runner_backend.is_triton():
|
||||
quant_info = self.get_triton_quant_info(layer)
|
||||
else:
|
||||
|
||||
@@ -121,6 +121,7 @@ def _moe_runner_keeps_global_expert_ids() -> bool:
|
||||
or b.is_flashinfer_trtllm()
|
||||
or b.is_flashinfer_trtllm_routed()
|
||||
or b.is_flashinfer_mxfp4()
|
||||
or b.is_hpc_ops()
|
||||
)
|
||||
except Exception: # pragma: no cover - backend not initialized
|
||||
return False
|
||||
|
||||
@@ -266,6 +266,7 @@ MOE_RUNNER_BACKEND_CHOICES = [
|
||||
"marlin",
|
||||
"humming",
|
||||
"experimental_sgl_marlin",
|
||||
"hpc_ops", # HPC-Ops (https://github.com/Tencent/hpc-ops), FP8 MoE on Hopper+
|
||||
]
|
||||
|
||||
MOE_A2A_BACKEND_CHOICES = [
|
||||
|
||||
@@ -0,0 +1,234 @@
|
||||
"""Numerical tests for the HPC-Ops FP8 MoE runner backend.
|
||||
|
||||
Compares the hpc_ops fused func (hpc.fuse_moe_blockwise) against the triton
|
||||
fused_experts reference and an fp32 exact reference on realistic blockwise
|
||||
FP8 quantized weights. Skipped when HPC-Ops (https://github.com/Tencent/hpc-ops)
|
||||
is not installed or the GPU is older than sm90.
|
||||
"""
|
||||
|
||||
import os
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.distributed.parallel_state import (
|
||||
init_distributed_environment,
|
||||
initialize_model_parallel,
|
||||
model_parallel_is_initialized,
|
||||
)
|
||||
from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig
|
||||
from sglang.srt.layers.moe.moe_runner.hpc_ops import (
|
||||
HpcOpsMoeQuantInfo,
|
||||
fused_experts_none_to_hpc_ops,
|
||||
has_hpc_ops,
|
||||
pad_hpc_ops_block_scale,
|
||||
)
|
||||
from sglang.srt.layers.moe.token_dispatcher.standard import StandardDispatchOutput
|
||||
from sglang.srt.layers.moe.topk import StandardTopKOutput
|
||||
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=60, stage="base-b", runner_config="1-gpu-large")
|
||||
|
||||
# Qwen3-30B-A3B-FP8 MoE shapes.
|
||||
E, TOPK, H, I = 128, 8, 2048, 768
|
||||
|
||||
|
||||
def _sm90_or_newer() -> bool:
|
||||
if not torch.cuda.is_available():
|
||||
return False
|
||||
major, _ = torch.cuda.get_device_capability()
|
||||
return major >= 9
|
||||
|
||||
|
||||
def _ensure_dist_initialized() -> None:
|
||||
"""Single-rank gloo distributed + model-parallel groups (TP=1, EP=1).
|
||||
|
||||
The triton fused_experts reference allocates its output under
|
||||
``use_symmetric_memory(get_tp_group(), ...)``, which requires the TP
|
||||
group even when symmetric allocation is disabled.
|
||||
"""
|
||||
os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
|
||||
os.environ.setdefault("MASTER_PORT", "29633")
|
||||
os.environ.setdefault("RANK", "0")
|
||||
os.environ.setdefault("WORLD_SIZE", "1")
|
||||
os.environ.setdefault("LOCAL_RANK", "0")
|
||||
if not torch.distributed.is_initialized():
|
||||
init_distributed_environment(world_size=1, rank=0, local_rank=0, backend="gloo")
|
||||
if not model_parallel_is_initialized():
|
||||
initialize_model_parallel(
|
||||
tensor_model_parallel_size=1,
|
||||
expert_model_parallel_size=1,
|
||||
pipeline_model_parallel_size=1,
|
||||
backend="gloo",
|
||||
)
|
||||
|
||||
|
||||
def _quant_blockwise(w: torch.Tensor, block: int = 128):
|
||||
"""Proper 128x128 blockwise fp8 quantization of a fp32 weight [E, N, K]."""
|
||||
num_experts, n, k = w.shape
|
||||
wb = w.view(num_experts, n // block, block, k // block, block)
|
||||
amax = wb.abs().amax(dim=(2, 4), keepdim=True).clamp(min=1e-6)
|
||||
scale = amax / 448.0
|
||||
wq = (wb / scale).clamp(-448, 448).view(num_experts, n, k).to(torch.float8_e4m3fn)
|
||||
return wq, scale.squeeze(-1).squeeze(2)
|
||||
|
||||
|
||||
@unittest.skipUnless(
|
||||
has_hpc_ops() and _sm90_or_newer(),
|
||||
"requires HPC-Ops (install from source: https://github.com/Tencent/hpc-ops) and sm90+",
|
||||
)
|
||||
class TestHpcOpsMoeBlockwise(CustomTestCase):
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
|
||||
_ensure_dist_initialized()
|
||||
torch.manual_seed(0)
|
||||
|
||||
def _make_case(self, num_tokens: int):
|
||||
device = "cuda"
|
||||
x = torch.randn(num_tokens, H, dtype=torch.bfloat16, device=device)
|
||||
w13_f = torch.randn(E, 2 * I, H, dtype=torch.float32, device=device) * 0.02
|
||||
w2_f = torch.randn(E, H, I, dtype=torch.float32, device=device) * 0.02
|
||||
w13, w13_scale = _quant_blockwise(w13_f)
|
||||
w2, w2_scale = _quant_blockwise(w2_f)
|
||||
logits = torch.randn(num_tokens, E, dtype=torch.float32, device=device)
|
||||
topk_weights, topk_ids = torch.topk(logits.softmax(dim=-1), TOPK, dim=-1)
|
||||
topk_weights = (topk_weights / topk_weights.sum(dim=-1, keepdim=True)).float()
|
||||
return x, w13, w2, w13_scale, w2_scale, topk_weights, topk_ids.to(torch.int32)
|
||||
|
||||
def _runner_config(self):
|
||||
return MoeRunnerConfig(
|
||||
num_experts=E,
|
||||
num_local_experts=E,
|
||||
hidden_size=H,
|
||||
intermediate_size_per_partition=I,
|
||||
layer_id=0,
|
||||
top_k=TOPK,
|
||||
num_fused_shared_experts=0,
|
||||
params_dtype=torch.bfloat16,
|
||||
activation="silu",
|
||||
is_gated=True,
|
||||
inplace=False,
|
||||
)
|
||||
|
||||
def _run_hpc_ops(self, x, w13, w2, w13_scale, w2_scale, topk_weights, topk_ids):
|
||||
dispatch_output = StandardDispatchOutput(
|
||||
hidden_states=x,
|
||||
hidden_states_scale=None,
|
||||
topk_output=StandardTopKOutput(topk_weights, topk_ids, None),
|
||||
)
|
||||
quant_info = HpcOpsMoeQuantInfo(
|
||||
w13_weight=w13,
|
||||
w2_weight=w2,
|
||||
block_quant=True,
|
||||
global_num_experts=E,
|
||||
moe_ep_rank=0,
|
||||
w13_weight_scale_inv=pad_hpc_ops_block_scale(w13_scale),
|
||||
w2_weight_scale_inv=pad_hpc_ops_block_scale(w2_scale),
|
||||
block_shape=[128, 128],
|
||||
)
|
||||
return fused_experts_none_to_hpc_ops(
|
||||
dispatch_output, quant_info, self._runner_config()
|
||||
).hidden_states
|
||||
|
||||
def _run_triton(self, x, w13, w2, w13_scale, w2_scale, topk_weights, topk_ids):
|
||||
from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import (
|
||||
fused_experts,
|
||||
)
|
||||
|
||||
return fused_experts(
|
||||
hidden_states=x.clone(),
|
||||
w1=w13,
|
||||
w2=w2,
|
||||
topk_output=StandardTopKOutput(topk_weights, topk_ids, None),
|
||||
moe_runner_config=self._runner_config(),
|
||||
use_fp8_w8a8=True,
|
||||
w1_scale=w13_scale,
|
||||
w2_scale=w2_scale,
|
||||
block_shape=[128, 128],
|
||||
)
|
||||
|
||||
def test_blockwise_fp8_matches_triton(self):
|
||||
for num_tokens in (7, 64, 512):
|
||||
with self.subTest(num_tokens=num_tokens):
|
||||
case = self._make_case(num_tokens)
|
||||
out_hpc = self._run_hpc_ops(*case).float()
|
||||
out_triton = self._run_triton(*case).float()
|
||||
cos = torch.nn.functional.cosine_similarity(
|
||||
out_hpc.flatten(), out_triton.flatten(), dim=0
|
||||
)
|
||||
self.assertGreater(cos.item(), 0.999)
|
||||
torch.testing.assert_close(out_hpc, out_triton, rtol=0.05, atol=0.05)
|
||||
|
||||
def test_per_tensor_fp8_matches_triton(self):
|
||||
|
||||
device = "cuda"
|
||||
for num_tokens in (7, 64, 512):
|
||||
with self.subTest(num_tokens=num_tokens):
|
||||
x = torch.randn(num_tokens, H, dtype=torch.bfloat16, device=device)
|
||||
w13_f = torch.randn(E, 2 * I, H, device=device) * 0.02
|
||||
w2_f = torch.randn(E, H, I, device=device) * 0.02
|
||||
# Per-tensor per-expert weight quant.
|
||||
w13_scale = w13_f.abs().amax(dim=(1, 2)) / 448.0
|
||||
w2_scale = w2_f.abs().amax(dim=(1, 2)) / 448.0
|
||||
w13 = (w13_f / w13_scale[:, None, None]).to(torch.float8_e4m3fn)
|
||||
w2 = (w2_f / w2_scale[:, None, None]).to(torch.float8_e4m3fn)
|
||||
# Static activation scales.
|
||||
a1_scale = torch.tensor(x.float().abs().max() / 448.0, device=device)
|
||||
a2_scale = torch.tensor(0.005, device=device)
|
||||
logits = torch.randn(num_tokens, E, device=device)
|
||||
topk_weights, topk_ids = torch.topk(logits.softmax(dim=-1), TOPK, -1)
|
||||
topk_weights = (
|
||||
topk_weights / topk_weights.sum(dim=-1, keepdim=True)
|
||||
).float()
|
||||
topk_ids = topk_ids.to(torch.int32)
|
||||
|
||||
dispatch_output = StandardDispatchOutput(
|
||||
hidden_states=x,
|
||||
hidden_states_scale=None,
|
||||
topk_output=StandardTopKOutput(topk_weights, topk_ids, None),
|
||||
)
|
||||
quant_info = HpcOpsMoeQuantInfo(
|
||||
w13_weight=w13,
|
||||
w2_weight=w2,
|
||||
block_quant=False,
|
||||
global_num_experts=E,
|
||||
moe_ep_rank=0,
|
||||
gate_up_alphas=w13_scale * a1_scale,
|
||||
down_alphas=w2_scale * a2_scale,
|
||||
w13_input_scale=a1_scale,
|
||||
w2_input_scale=a2_scale,
|
||||
)
|
||||
out_hpc = fused_experts_none_to_hpc_ops(
|
||||
dispatch_output, quant_info, self._runner_config()
|
||||
).hidden_states.float()
|
||||
|
||||
from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import (
|
||||
fused_experts,
|
||||
)
|
||||
|
||||
out_triton = fused_experts(
|
||||
hidden_states=x.clone(),
|
||||
w1=w13,
|
||||
w2=w2,
|
||||
topk_output=StandardTopKOutput(topk_weights, topk_ids, None),
|
||||
moe_runner_config=self._runner_config(),
|
||||
use_fp8_w8a8=True,
|
||||
w1_scale=w13_scale,
|
||||
w2_scale=w2_scale,
|
||||
a1_scale=a1_scale,
|
||||
a2_scale=a2_scale,
|
||||
).float()
|
||||
|
||||
cos = torch.nn.functional.cosine_similarity(
|
||||
out_hpc.flatten(), out_triton.flatten(), dim=0
|
||||
)
|
||||
self.assertGreater(cos.item(), 0.99)
|
||||
torch.testing.assert_close(out_hpc, out_triton, rtol=0.10, atol=0.10)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,63 @@
|
||||
"""The hpc_ops MoE runner backend makes the standard dispatcher keep global
|
||||
expert ids, so a quant method that silently falls back to another runner
|
||||
(e.g. an unquantized MoE) would misroute tokens under EP>1. MoeRunner must
|
||||
reject that combination loudly at startup.
|
||||
"""
|
||||
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
|
||||
from sglang.srt.layers.moe.moe_runner.runner import MoeRunner
|
||||
from sglang.srt.layers.moe.utils import MoeRunnerBackend
|
||||
from sglang.srt.runtime_context import get_flags
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-c-test-cpu")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _runner_backend_flag():
|
||||
moe = get_flags().moe
|
||||
saved = moe.runner_backend
|
||||
yield moe
|
||||
moe.runner_backend = saved
|
||||
|
||||
|
||||
def test_non_hpc_runner_rejected_when_hpc_ops_requested(_runner_backend_flag):
|
||||
_runner_backend_flag.runner_backend = MoeRunnerBackend.HPC_OPS
|
||||
with pytest.raises(ValueError, match="hpc_ops"):
|
||||
MoeRunner(MoeRunnerBackend.TRITON, MoeRunnerConfig())
|
||||
|
||||
|
||||
def test_triton_runner_allowed_without_hpc_ops(_runner_backend_flag):
|
||||
_runner_backend_flag.runner_backend = MoeRunnerBackend.TRITON
|
||||
runner = MoeRunner(MoeRunnerBackend.TRITON, MoeRunnerConfig())
|
||||
assert runner.runner_core is not None
|
||||
|
||||
|
||||
def test_direct_kernel_quant_method_rejected_when_hpc_ops_requested(
|
||||
_runner_backend_flag,
|
||||
):
|
||||
# W4AFp8MoEMethod never constructs a MoeRunner (apply() calls its kernel
|
||||
# directly), so it bypasses the MoeRunner-level guard; the layer-level
|
||||
# check must reject it.
|
||||
from sglang.srt.layers.moe.fused_moe_triton.layer import (
|
||||
_validate_hpc_ops_quant_method,
|
||||
)
|
||||
from sglang.srt.layers.quantization.fp8 import Fp8MoEMethod
|
||||
from sglang.srt.layers.quantization.w4afp8 import W4AFp8MoEMethod
|
||||
|
||||
_runner_backend_flag.runner_backend = MoeRunnerBackend.HPC_OPS
|
||||
with pytest.raises(ValueError, match="hpc_ops"):
|
||||
_validate_hpc_ops_quant_method(object.__new__(W4AFp8MoEMethod))
|
||||
# The FP8 method (the one the hpc_ops runner supports) passes.
|
||||
_validate_hpc_ops_quant_method(object.__new__(Fp8MoEMethod))
|
||||
# Without hpc_ops requested, any quant method passes.
|
||||
_runner_backend_flag.runner_backend = MoeRunnerBackend.TRITON
|
||||
_validate_hpc_ops_quant_method(object.__new__(W4AFp8MoEMethod))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
@@ -913,6 +913,7 @@ class TestModuleLevelHelpers(unittest.TestCase):
|
||||
backends.FLASHINFER_CUTLASS,
|
||||
backends.FLASHINFER_MXFP4,
|
||||
backends.FLASHINFER_CUTEDSL,
|
||||
backends.HPC_OPS,
|
||||
}
|
||||
config = types.SimpleNamespace(
|
||||
num_experts=8,
|
||||
|
||||
Reference in New Issue
Block a user