[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
co-authored by
Claude Fable 5
Halcyon
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 = [
|
||||
|
||||
Reference in New Issue
Block a user