| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for the SM100 BF16x3 router GEMM.""" |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | |
| 8 | from vllm.utils.import_utils import has_cutedsl |
| 9 | |
| 10 | |
| 11 | def _requires_sm100_cutedsl(): |
| 12 | if not torch.cuda.is_available(): |
| 13 | pytest.skip("CUDA not available") |
| 14 | major, _ = torch.cuda.get_device_capability() |
| 15 | if major != 10: |
| 16 | pytest.skip("bf16x3 router GEMM requires SM100-class GPU") |
| 17 | if not has_cutedsl(): |
| 18 | pytest.skip("cutedsl (cutlass) not installed") |
| 19 | |
| 20 | |
| 21 | @pytest.mark.parametrize( |
| 22 | ("num_tokens", "hidden_dim", "num_experts"), |
| 23 | [ |
| 24 | (48, 6144, 128), |
| 25 | (96, 3072, 256), |
| 26 | (129, 3072, 17), |
| 27 | # long-K cases exercise the multi-accumulation path (the split-K |
| 28 | # heuristic leaves chains of 15 and 32 K-tiles here, above the |
| 29 | # kernel's num_tmem_acc bound) |
| 30 | (1024, 8192, 256), |
| 31 | (2048, 8192, 256), |
| 32 | ], |
| 33 | ) |
| 34 | def test_bf16x3_router_gemm_matches_reference( |
| 35 | num_tokens: int, hidden_dim: int, num_experts: int |
| 36 | ): |
| 37 | _requires_sm100_cutedsl() |
| 38 | from vllm.model_executor.layers.fused_moe.router.bf16x3_router_gemm_cutedsl import ( # noqa: E501 |
| 39 | bf16x3_router_gemm, |
| 40 | ) |
| 41 | |
| 42 | torch.manual_seed(42) |
| 43 | x = torch.randn(num_tokens, hidden_dim, dtype=torch.bfloat16, device="cuda") |
| 44 | w = torch.randn(num_experts, hidden_dim, dtype=torch.float32, device="cuda") |
| 45 | # Match the observed router weight scale |
| 46 | w *= 0.053 |
| 47 | out = bf16x3_router_gemm(x, w) |
| 48 | # FP64 reference: the FP32 reference itself drifts by ~5e-6 at N=2048 |
| 49 | ref = torch.nn.functional.linear(x.double(), w.double()) |
| 50 | |
| 51 | assert out.shape == (num_tokens, num_experts) |
| 52 | assert out.dtype == torch.float32 |
| 53 | assert torch.mean(torch.abs(out.double() - ref)).item() < 5e-6 |