[MoE][Kernel] Add optional HPC BF16xFP32 router GEMM

tests/kernels/test_hpc_router_gemm.py

@@ -0,0 +1,155 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm._custom_ops import fp32_router_gemm
+from vllm.utils.hpc import (
+    has_hpc_bf16xfp32_gemm,
+    hpc_gemm_bf16xfp32,
+)
+
+SHAPES = [
+    (3072, 256),
+    (6144, 128),
+    (6144, 256),
+]
+FP32_ROUTER_SHAPES = [
+    (3072, 256),
+    (6144, 128),
+    (6144, 256),
+]
+TOKEN_COUNTS = [1, 2, 8, 16, 32, 128, 512]
+
+ATOL = 1e-1
+RTOL = 8e-2
+MAX_RELATIVE_L2 = 2e-3
+HPC_GEMM_WEIGHT_SCALE = 1.0 / 256.0
+
+
+HPC_GEMM_WEIGHT_SCALE = 1.0 / 256.0
+
+
+def _run_hpc_gemm(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+    weight_high = weight.to(torch.bfloat16)
+    weight_low = ((weight - weight_high.float()) / HPC_GEMM_WEIGHT_SCALE).to(
+        torch.bfloat16
+    )
+    import hpc
+
+    split_flag = hpc.get_gemm_bf16xfp32_workspace(weight.shape[0])
+    return hpc_gemm_bf16xfp32(
+        x=x,
+        weight_high=weight_high,
+        weight_low=weight_low,
+        scale=HPC_GEMM_WEIGHT_SCALE,
+        split_flag=split_flag,
+    )
+
+
+def _requires_hpc_sm90() -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA not available")
+    if torch.cuda.get_device_capability() != (9, 0):
+        pytest.skip("HPC BF16xFP32 router GEMM requires SM90")
+    if not has_hpc_bf16xfp32_gemm():
+        pytest.skip("HPC BF16xFP32 GEMM is unavailable")
+
+
+@pytest.mark.parametrize("hidden_size,num_experts", SHAPES)
+@pytest.mark.parametrize("num_tokens", TOKEN_COUNTS)
+@pytest.mark.parametrize("seed", [1, 17, 42])
+def test_hpc_router_gemm_numerics(
+    hidden_size: int,
+    num_experts: int,
+    num_tokens: int,
+    seed: int,
+) -> None:
+    _requires_hpc_sm90()
+    torch.manual_seed(seed)
+
+    x = torch.randn(
+        num_tokens,
+        hidden_size,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    weight = (
+        torch.randn(
+            num_experts,
+            hidden_size,
+            device="cuda",
+            dtype=torch.float32,
+        )
+        * 0.02
+    )
+
+    output = _run_hpc_gemm(x, weight)
+
+    # FP64 reference avoids TF32/cuBLAS FP32 reference uncertainty.
+    reference = torch.nn.functional.linear(
+        x.double(),
+        weight.double(),
+    ).float()
+
+    assert output.dtype == torch.float32
+    assert output.shape == reference.shape
+    assert torch.isfinite(output).all()
+
+    torch.testing.assert_close(
+        output,
+        reference,
+        atol=ATOL,
+        rtol=RTOL,
+    )
+
+    error_l2 = torch.linalg.vector_norm((output - reference).double())
+    reference_l2 = torch.linalg.vector_norm(reference.double())
+    relative_l2 = error_l2 / reference_l2.clamp_min(1e-12)
+    assert relative_l2.item() < MAX_RELATIVE_L2
+
+
+@pytest.mark.parametrize("hidden_size,num_experts", FP32_ROUTER_SHAPES)
+@pytest.mark.parametrize("num_tokens", [1, 8, 32, 128, 512])
+@pytest.mark.parametrize("seed", range(5))
+def test_hpc_vs_native_router(
+    hidden_size: int,
+    num_experts: int,
+    num_tokens: int,
+    seed: int,
+) -> None:
+    _requires_hpc_sm90()
+    torch.manual_seed(1000 + seed)
+
+    x = torch.randn(
+        num_tokens,
+        hidden_size,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    weight = (
+        torch.randn(
+            num_experts,
+            hidden_size,
+            device="cuda",
+            dtype=torch.float32,
+        )
+        * 0.02
+    )
+
+    hpc_output = _run_hpc_gemm(x, weight)
+    reference = torch.nn.functional.linear(
+        x.double(),
+        weight.double(),
+    )
+
+    hpc_error = torch.linalg.vector_norm((hpc_output - reference).double())
+    reference_norm = torch.linalg.vector_norm(reference.double()).clamp_min(1e-12)
+
+    assert (hpc_error / reference_norm).item() < MAX_RELATIVE_L2
+
+    if num_tokens <= 32:
+        native_output = fp32_router_gemm(x, weight)
+        native_error = torch.linalg.vector_norm((native_output - reference).double())
+        assert (native_error / reference_norm).item() < MAX_RELATIVE_L2

vllm/envs.py

@@ -288,6 +288,7 @@
     VLLM_DISABLE_SHARED_EXPERTS_STREAM: bool = False
     VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD: int = 256
     VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD: int = 1024
+    VLLM_ENABLE_HPC_ROUTER_GEMM: bool = False
     VLLM_COMPILE_CACHE_SAVE_FORMAT: Literal["binary", "unpacked"] = "binary"
     VLLM_USE_V2_MODEL_RUNNER: bool | None = None
     VLLM_LOG_MODEL_INSPECTION: bool = False
@@ -2010,6 +2011,10 @@ def _resolve_rust_cli_path() -> str | None:
     "VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD": lambda: int(
         os.getenv("VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD", "1024")
     ),
+    # Enable the experimental HPC BF16xFP32 router GEMM path on Hopper.
+    "VLLM_ENABLE_HPC_ROUTER_GEMM": lambda: bool(
+        int(os.getenv("VLLM_ENABLE_HPC_ROUTER_GEMM", "0"))
+    ),
     # Format for saving torch.compile cache artifacts
     # - "binary": saves as binary file
     #     Safe for multiple vllm serve processes accessing the same torch compile cache.

vllm/model_executor/layers/fused_moe/router/gate_linear.py

@@ -4,16 +4,27 @@
 from torch.nn.parameter import Parameter
 
 import vllm._custom_ops as ops
+import vllm.envs as envs
 from vllm.config import get_current_vllm_config_or_none
 from vllm.logger import init_logger
 from vllm.model_executor.custom_op import PluggableLayer
-from vllm.model_executor.layers.linear import ReplicatedLinear
+from vllm.model_executor.layers.linear import ReplicatedLinear, UnquantizedLinearMethod
 from vllm.platforms import current_platform
+from vllm.utils.hpc import (
+    has_hpc_bf16xfp32_gemm,
+    hpc_gemm_bf16xfp32,
+)
 from vllm.utils.torch_utils import direct_register_custom_op
 
 logger = init_logger(__name__)
 
 
+class GateLinearWeightProcess(UnquantizedLinearMethod):
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        layer._process_gate_weights_after_loading()
+        super().process_weights_after_loading(layer)
+
+
 @PluggableLayer.register("gate_linear")
 class GateLinear(ReplicatedLinear):
     """MoE gate linear layer with multi-tier GEMM dispatch:
@@ -96,6 +107,11 @@ def __init__(
             vllm_config is not None
             and vllm_config.kernel_config.enable_bf16x3_router_gemm
         )
+        self._hpc_workspace_max_tokens = (
+            vllm_config.scheduler_config.max_num_batched_tokens
+            if vllm_config is not None
+            else _HPC_GEMM_DEFAULT_MAX_TOKENS
+        )
         self.allow_fp32_router_gemm = (
             not bias
             and self.weight.dtype == torch.float32
@@ -144,6 +160,31 @@ def __init__(
                 and is_available()
             )
 
+        self._is_hopper = is_hopper
+        self.allow_hpc_router_gemm = self._compute_allow_hpc_router_gemm()
+
+        if self.allow_hpc_router_gemm:
+            self.quant_method = GateLinearWeightProcess()
+            logger.info_once("Enabled HPC BF16xFP32 router GEMM.")
+
+    def _compute_allow_hpc_router_gemm(self) -> bool:
+        """HPC BF16xFP32 router GEMM eligibility.
+
+        Shared by __init__ and set_out_dtype() since out_dtype may only be
+        known after construction.
+        """
+        return (
+            envs.VLLM_ENABLE_HPC_ROUTER_GEMM
+            and self._router_gemm_no_bias
+            and self.weight.dtype == torch.float32
+            and current_platform.is_cuda()
+            and self._is_hopper
+            and self.input_size % 8 == 0
+            and self.output_size % 64 == 0
+            and has_hpc_bf16xfp32_gemm()
+            and self.out_dtype == torch.float32
+        )
+
     def set_out_dtype(self, out_dtype: torch.dtype) -> None:
         """Set output dtype for the router logits after init.
 
@@ -176,6 +217,35 @@ def set_out_dtype(self, out_dtype: torch.dtype) -> None:
                 and is_available()
             )
 
+        # out_dtype may start as None -> recompute HPC eligibility here too
+        self.allow_hpc_router_gemm = self._compute_allow_hpc_router_gemm()
+        if self.allow_hpc_router_gemm:
+            self.quant_method = GateLinearWeightProcess()
+            logger.info_once("Enabled HPC BF16xFP32 router GEMM.")
+
+    def _process_gate_weights_after_loading(self):
+        if not self.allow_hpc_router_gemm:
+            return
+
+        import hpc
+
+        with torch.no_grad():
+            y = self.weight.data
+            w_high = y.to(torch.bfloat16)
+            w_low = ((y - w_high.float()) / _HPC_GEMM_WEIGHT_SCALE).to(torch.bfloat16)
+
+            if getattr(self, "hpc_w_high", None) is None:
+                # Allocate once: CUDA graph capture and weight-reload flows
+                # depend on these storage addresses staying stable.
+                self.hpc_w_high = w_high
+                self.hpc_w_low = w_low
+                self.hpc_split_flag = hpc.get_gemm_bf16xfp32_workspace(
+                    y.shape[0], max_tokens=self._hpc_workspace_max_tokens
+                )
+            else:
+                self.hpc_w_high.copy_(w_high)
+                self.hpc_w_low.copy_(w_low)
+
     def forward(
         self, x: torch.Tensor
     ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
@@ -197,6 +267,23 @@ def forward(
             )
             return output, None
 
+        # HPC BF16 activation x FP32 router weight.
+        if (
+            self.allow_hpc_router_gemm
+            and x.ndim == 2
+            and x.dtype == torch.bfloat16
+            and x.shape[0] > _FP32_ROUTER_GEMM_MAX_TOKENS
+            and x.shape[0] <= self._hpc_workspace_max_tokens
+        ):
+            output = hpc_gemm_bf16xfp32(
+                x=x,
+                weight_high=self.hpc_w_high,
+                weight_low=self.hpc_w_low,
+                scale=_HPC_GEMM_WEIGHT_SCALE,
+                split_flag=self.hpc_split_flag,
+            )
+            return output, None
+
         # Tier 3: fp32 specialized kernel (H=3072, E=256, M<=32)
         # Dispatch is wrapped in a custom op so that torch.compile/CUDA-graph
         # capture does not freeze the runtime num_tokens branch.
@@ -273,3 +360,13 @@ def fp32_router_gemm_dispatch_fake(
     op_func=fp32_router_gemm_dispatch_impl,
     fake_impl=fp32_router_gemm_dispatch_fake,
 )
+
+
+# The HPC-Ops bf16xfp32 GEMM consumes the fp32 weight decomposed into two
+# bf16 halves: w_high = w.bf16 and w_low = ((w - w_high) / scale).bf16 with
+# scale = 1/256, so that w ~= w_high + scale * w_low.
+_HPC_GEMM_WEIGHT_SCALE = 1.0 / 256.0
+
+# hpc.get_gemm_bf16xfp32_workspace()'s own default when no vLLM config is
+# available to size the split-K workspace from.
+_HPC_GEMM_DEFAULT_MAX_TOKENS = 131072

vllm/utils/hpc.py

@@ -216,8 +216,42 @@ def hpc_fuse_moe_blockwise(
     )
 
 
+@functools.cache
+def has_hpc_bf16xfp32_gemm() -> bool:
+    """Return whether the optional HPC bf16xfp32 GEMM API is available."""
+    if not has_hpc():
+        return False
+
+    import hpc
+
+    required_apis = ("gemm_bf16xfp32", "get_gemm_bf16xfp32_workspace")
+    return all(hasattr(hpc, name) for name in required_apis)
+
+
+def hpc_gemm_bf16xfp32(
+    x: torch.Tensor,
+    weight_high: torch.Tensor,
+    weight_low: torch.Tensor,
+    scale: float,
+    split_flag: torch.Tensor,
+) -> torch.Tensor:
+    import hpc
+
+    return hpc.gemm_bf16xfp32(
+        x,
+        weight_high,
+        weight_low,
+        scale,
+        use_fp32_output=True,
+        use_splitk=True,
+        split_flag=split_flag,
+    )
+
+
 __all__ = [
     "has_hpc",
+    "has_hpc_bf16xfp32_gemm",
     "hpc_fuse_moe",
     "hpc_fuse_moe_blockwise",
+    "hpc_gemm_bf16xfp32",
 ]