[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:
Xiaoyu Zhang
2026-07-24 19:46:11 +08:00
committed by GitHub
co-authored by Claude Fable 5 Halcyon
parent 4d5917e744
commit 3d91a569ce
11 changed files with 631 additions and 1 deletions
@@ -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
+4
View File
@@ -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:
+1
View File
@@ -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
+1
View File
@@ -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 = [
+234
View File
@@ -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,