From 3d91a569ce480f5548ca836799e84b7ee7fe97c4 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Fri, 24 Jul 2026 19:46:11 +0800 Subject: [PATCH] [MoE Backend] Add HPC-Ops FP8 MoE runner backend (#30541) Co-authored-by: Claude Fable 5 Co-authored-by: Halcyon <56064364+VAthree@users.noreply.github.com> --- .../srt/layers/moe/fused_moe_triton/layer.py | 19 ++ .../srt/layers/moe/moe_runner/hpc_ops.py | 208 ++++++++++++++++ .../srt/layers/moe/moe_runner/runner.py | 22 +- .../layers/moe/token_dispatcher/standard.py | 2 + python/sglang/srt/layers/moe/utils.py | 4 + python/sglang/srt/layers/quantization/fp8.py | 77 ++++++ python/sglang/srt/lora/mem_pool.py | 1 + python/sglang/srt/server_args.py | 1 + test/registered/moe/test_hpc_ops_moe.py | 234 ++++++++++++++++++ .../layers/moe/test_hpc_ops_runner_guard.py | 63 +++++ .../unit/lora/test_mem_pool_ep_unit.py | 1 + 11 files changed, 631 insertions(+), 1 deletion(-) create mode 100644 python/sglang/srt/layers/moe/moe_runner/hpc_ops.py create mode 100644 test/registered/moe/test_hpc_ops_moe.py create mode 100644 test/registered/unit/layers/moe/test_hpc_ops_runner_guard.py diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index 1c2aa6f81..e8007a4b7 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -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() diff --git a/python/sglang/srt/layers/moe/moe_runner/hpc_ops.py b/python/sglang/srt/layers/moe/moe_runner/hpc_ops.py new file mode 100644 index 000000000..cc7259116 --- /dev/null +++ b/python/sglang/srt/layers/moe/moe_runner/hpc_ops.py @@ -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) diff --git a/python/sglang/srt/layers/moe/moe_runner/runner.py b/python/sglang/srt/layers/moe/moe_runner/runner.py index a02a46d79..a0a3a92c5 100644 --- a/python/sglang/srt/layers/moe/moe_runner/runner.py +++ b/python/sglang/srt/layers/moe/moe_runner/runner.py @@ -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}") diff --git a/python/sglang/srt/layers/moe/token_dispatcher/standard.py b/python/sglang/srt/layers/moe/token_dispatcher/standard.py index 009221228..94db5a49e 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/standard.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/standard.py @@ -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 diff --git a/python/sglang/srt/layers/moe/utils.py b/python/sglang/srt/layers/moe/utils.py index 1fd5f7ee3..0f3f48f1e 100644 --- a/python/sglang/srt/layers/moe/utils.py +++ b/python/sglang/srt/layers/moe/utils.py @@ -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 diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 27c3ef7db..0a4c7e6df 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -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: diff --git a/python/sglang/srt/lora/mem_pool.py b/python/sglang/srt/lora/mem_pool.py index 1a09e0385..587d3a620 100644 --- a/python/sglang/srt/lora/mem_pool.py +++ b/python/sglang/srt/lora/mem_pool.py @@ -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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 710011a6f..0e408cf18 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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 = [ diff --git a/test/registered/moe/test_hpc_ops_moe.py b/test/registered/moe/test_hpc_ops_moe.py new file mode 100644 index 000000000..a37b25508 --- /dev/null +++ b/test/registered/moe/test_hpc_ops_moe.py @@ -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() diff --git a/test/registered/unit/layers/moe/test_hpc_ops_runner_guard.py b/test/registered/unit/layers/moe/test_hpc_ops_runner_guard.py new file mode 100644 index 000000000..95175cdc9 --- /dev/null +++ b/test/registered/unit/layers/moe/test_hpc_ops_runner_guard.py @@ -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"])) diff --git a/test/registered/unit/lora/test_mem_pool_ep_unit.py b/test/registered/unit/lora/test_mem_pool_ep_unit.py index 2e88ff8ec..280061d34 100644 --- a/test/registered/unit/lora/test_mem_pool_ep_unit.py +++ b/test/registered/unit/lora/test_mem_pool_ep_unit.py @@ -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,