diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index ba054b3e30..de1da4a54e 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -230,8 +230,8 @@ steps: # Catch-all for test files at the tests/kernels root. This job collects # the whole root so new files are wired by default. # Files with dedicated jobs elsewhere in this file are excluded via --ignore -# (test_bf16x3_router_gemm_cutedsl and test_ll_bf16_gemm run in their own jobs / -# Kernels (B200)). +# (test_bf16x3_router_gemm_cutedsl runs in Kernels (B200); test_ll_bf16_gemm and +# test_ll_fp32w_gemm run in Gate Linear CuTeDSL Kernels (B200)). - label: ":nvidia: (B200) Miscellaneous Kernels" key: kernels-root-misc-test-b200 timeout_in_minutes: 30 @@ -260,6 +260,7 @@ steps: --ignore=kernels/test_bf16x3_router_gemm_cutedsl.py --ignore=kernels/test_bf16x3_router_gemm_rocm.py --ignore=kernels/test_ll_bf16_gemm.py + --ignore=kernels/test_ll_fp32w_gemm.py --ignore=kernels/test_shuffle_rows.py # BROKEN on main, pending kernel fixes (B200): # test_shuffle_rows.py (1: test_shuffle_rows_edge_cases) @@ -595,9 +596,6 @@ steps: - vllm/v1/attention/backends/mla/flashinfer_mla.py - vllm/v1/attention/selector.py - vllm/platforms/cuda.py - - vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py - - vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py - - vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_splitk.py - vllm/cute_utils/ - vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/ - vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py @@ -617,7 +615,6 @@ steps: - tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py - tests/kernels/mamba/test_gdn_prefill_cutedsl.py - tests/kernels/test_bf16x3_router_gemm_cutedsl.py - - tests/kernels/test_ll_bf16_gemm.py - tests/kernels/test_top_k_per_row.py - tests/kernels/quantization/test_silu_mul_nvfp4_quant.py - tests/kernels/quantization/test_block_fp8.py @@ -638,7 +635,6 @@ steps: pytest -v -s tests/kernels/moe/test_nvfp4_moe.py pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py pytest -v -s tests/kernels/test_bf16x3_router_gemm_cutedsl.py - pytest -v -s tests/kernels/test_ll_bf16_gemm.py ;; 1) python3 examples/basic/offline_inference/chat.py @@ -695,6 +691,28 @@ steps: - amd-smi - pytest -v -s kernels/quantization/test_nvfp4_emulation.py -k test_nvfp4_dense_emulation --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT +- label: ":nvidia: (B200) Gate Linear CuTeDSL Kernels" + key: kernels-gate-linear-cutedsl-b200 + timeout_in_minutes: 15 + working_dir: "/vllm-workspace/" + device: b200-k8s + source_file_dependencies: + - vllm/cute_utils/ + - vllm/platforms/cuda.py + - vllm/model_executor/layers/fused_moe/router/gate_linear.py + - vllm/model_executor/kernels/linear/cute_dsl/_ll_router_common.py + - vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py + - vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py + - vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_splitk.py + - tests/kernels/test_ll_bf16_gemm.py + - vllm/model_executor/kernels/linear/cute_dsl/ll_fp32w.py + - vllm/model_executor/kernels/linear/cute_dsl/_ll_fp32w_dotprod.py + - tests/kernels/test_ll_fp32w_gemm.py + commands: + - nvidia-smi + - pytest -v -s tests/kernels/test_ll_bf16_gemm.py + - pytest -v -s tests/kernels/test_ll_fp32w_gemm.py + - label: ":nvidia: (DGX) Spark b12x Linear Kernels Nightly" key: b12x-linear-kernels-dgx-spark-nightly timeout_in_minutes: 60 diff --git a/tests/kernels/test_ll_fp32w_gemm.py b/tests/kernels/test_ll_fp32w_gemm.py new file mode 100644 index 0000000000..4d103d185b --- /dev/null +++ b/tests/kernels/test_ll_fp32w_gemm.py @@ -0,0 +1,544 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for cuteDSL low-latency GEMM with fp32 weights.""" + +import pytest +import torch +import torch.nn.functional as F + +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not torch.accelerator.is_available(), reason="CUDA required" +) + + +@pytest.fixture(autouse=True, scope="module") +def _require_sm90_and_cutedsl(): + if not current_platform.has_device_capability(90): + pytest.skip("Requires SM90+ (Hopper/Blackwell)") + from vllm.model_executor.kernels.linear.cute_dsl.ll_fp32w import ( + is_available, + ) + + if not is_available(): + pytest.skip("cuteDSL (CUTLASS Python) not installed") + + +def _ref(a, b): + return torch.mm(a.float(), b.float().T) + + +def _assert_close(out, ref, *, min_cos_sim=0.99, context=""): + assert out.device.type == "cuda", f"{context}: not on CUDA" + assert torch.isfinite(out).all(), f"{context}: NaN/Inf" + cos = F.cosine_similarity( + out.reshape(-1).float(), ref.reshape(-1).float(), dim=0 + ).item() + assert cos > min_cos_sim, ( + f"{context}: cos_sim {cos:.6f} < {min_cos_sim} " + f"(abs_err={(out.float() - ref.float()).abs().max().item():.2e})" + ) + + +def _gemm(a, b, output_dtype=torch.float32): + from vllm.model_executor.kernels.linear.cute_dsl.ll_fp32w import ( + LL_FP32W_GEMM_KERNEL, + ) + + return LL_FP32W_GEMM_KERNEL(a, b, output_dtype) + + +SHAPES = [ + (256, 3072, "MiniMax-M2"), + (256, 7168, "DSV3-sized"), + (128, 5120, "DeepSeek-V2"), + (64, 2880, "non-aligned-K"), + (8, 4096, "Mixtral"), +] + +A_DTYPES = [ + pytest.param(torch.bfloat16, id="bf16"), + pytest.param(torch.float16, id="fp16"), + pytest.param(torch.float32, id="fp32"), +] + + +# ================================================================= +# Correctness: all dtype combos x shapes x M values +# ================================================================= + + +@pytest.mark.parametrize("a_dtype", A_DTYPES) +@pytest.mark.parametrize("M", [1, 2, 4, 8, 16]) +@pytest.mark.parametrize("N,K,desc", SHAPES, ids=[s[2] for s in SHAPES]) +def test_correctness(M, N, K, desc, a_dtype): + torch.manual_seed(42) + a = torch.randn(M, K, device="cuda").to(a_dtype) + b = torch.randn(N, K, dtype=torch.float32, device="cuda") + out = _gemm(a, b) + assert out.dtype == torch.float32 + assert out.shape == (M, N) + _assert_close(out, _ref(a, b), context=f"{a_dtype} M={M} {desc}") + + +# ================================================================= +# Small K (k_main=0, only tail/scalar paths) +# ================================================================= + + +@pytest.mark.parametrize("K", [16, 32, 64, 128, 256, 512, 1024, 1536]) +def test_small_K(K): + torch.manual_seed(42) + a = torch.randn(4, K, dtype=torch.bfloat16, device="cuda") + b = torch.randn(64, K, dtype=torch.float32, device="cuda") + _assert_close(_gemm(a, b), _ref(a, b), context=f"small_K={K}") + + +# ================================================================= +# K sweep (covers main loop + various tail combinations) +# ================================================================= + + +@pytest.mark.parametrize("K", [2048, 2304, 2880, 3072, 4096, 5120, 7168]) +def test_K_sweep(K): + torch.manual_seed(42) + a = torch.randn(8, K, dtype=torch.bfloat16, device="cuda") + b = torch.randn(64, K, dtype=torch.float32, device="cuda") + _assert_close(_gemm(a, b), _ref(a, b), context=f"K={K}") + + +# ================================================================= +# MiniMax-M2 specific shapes +# ================================================================= + + +@pytest.mark.parametrize("M", [1, 4, 8, 16, 32]) +def test_minimax_m2(M): + torch.manual_seed(42) + a = torch.randn(M, 3072, dtype=torch.bfloat16, device="cuda") + b = torch.randn(256, 3072, dtype=torch.float32, device="cuda") + _assert_close(_gemm(a, b), _ref(a, b), context=f"MiniMax-M2 M={M}") + + +# ================================================================= +# Single token (M=1, decode path) +# ================================================================= + + +@pytest.mark.parametrize( + "N,K", + [(256, 3072), (256, 7168), (8, 4096)], + ids=["MiniMax-M2", "DSV3-sized", "Mixtral"], +) +def test_single_token(N, K): + torch.manual_seed(42) + a = torch.randn(1, K, dtype=torch.bfloat16, device="cuda") + b = torch.randn(N, K, dtype=torch.float32, device="cuda") + out = _gemm(a, b) + assert out.shape == (1, N) + _assert_close(out, _ref(a, b), context=f"M=1 {N}x{K}") + + +# ================================================================= +# Arbitrary N +# ================================================================= + + +@pytest.mark.parametrize("N", [1, 3, 16, 64, 256]) +def test_arbitrary_N(N): + torch.manual_seed(42) + a = torch.randn(4, 3072, dtype=torch.bfloat16, device="cuda") + b = torch.randn(N, 3072, dtype=torch.float32, device="cuda") + out = _gemm(a, b) + assert out.shape == (4, N) + _assert_close(out, _ref(a, b), context=f"N={N}") + + +# ================================================================= +# Numerical robustness +# ================================================================= + + +def test_large_values(): + torch.manual_seed(42) + a = torch.randn(4, 5120, dtype=torch.bfloat16, device="cuda") * 100 + b = torch.randn(128, 5120, dtype=torch.float32, device="cuda") * 100 + out = _gemm(a, b) + assert torch.isfinite(out).all() + _assert_close(out, _ref(a, b), context="large_values") + + +def test_near_zero(): + torch.manual_seed(42) + a = torch.randn(4, 4096, dtype=torch.bfloat16, device="cuda") * 1e-4 + b = torch.randn(64, 4096, dtype=torch.float32, device="cuda") * 1e-4 + out = _gemm(a, b) + assert torch.isfinite(out).all() + assert out.abs().max() < 1.0 + + +def test_zeros(): + a = torch.zeros(4, 2048, dtype=torch.bfloat16, device="cuda") + b = torch.randn(64, 2048, dtype=torch.float32, device="cuda") + assert (_gemm(a, b) == 0).all() + + +def test_ones(): + a = torch.ones(1, 4096, dtype=torch.bfloat16, device="cuda") + b = torch.randn(32, 4096, dtype=torch.float32, device="cuda") + _assert_close(_gemm(a, b), _ref(a, b), context="ones") + + +# ================================================================= +# Output dtype +# ================================================================= + + +def test_output_fp32(): + torch.manual_seed(42) + a = torch.randn(4, 5120, dtype=torch.bfloat16, device="cuda") + b = torch.randn(128, 5120, dtype=torch.float32, device="cuda") + assert _gemm(a, b).dtype == torch.float32 + + +# ================================================================= +# Determinism +# ================================================================= + + +@pytest.mark.parametrize("M", [1, 4, 8, 16]) +@pytest.mark.parametrize("K", [1024, 3072, 7168]) +def test_deterministic(M, K): + torch.manual_seed(42) + a = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") + b = torch.randn(64, K, dtype=torch.float32, device="cuda") + torch.testing.assert_close(_gemm(a, b), _gemm(a, b), atol=0, rtol=0) + + +# ================================================================= +# Cross-dtype consistency +# ================================================================= + + +def test_bf16_vs_fp32_activations(): + """bf16 and fp32 activations on same data should roughly agree.""" + torch.manual_seed(42) + K, N = 4096, 128 + a_bf16 = torch.randn(4, K, dtype=torch.bfloat16, device="cuda") + b = torch.randn(N, K, dtype=torch.float32, device="cuda") + out_bf16 = _gemm(a_bf16, b) + a_fp32 = a_bf16.float() + out_fp32 = _gemm(a_fp32, b) + _assert_close( + out_bf16, out_fp32, min_cos_sim=0.999, context="bf16_vs_fp32_activations" + ) + + +# ================================================================= +# CUDA graph +# ================================================================= + + +@pytest.mark.parametrize( + "M,K,N", [(4, 3072, 256), (16, 5120, 128)], ids=["MiniMax", "DeepSeek"] +) +def test_cudagraph(M, K, N): + torch.manual_seed(42) + a = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") + b = torch.randn(N, K, dtype=torch.float32, device="cuda") + _gemm(a, b) + torch.accelerator.synchronize() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out = _gemm(a, b) + for _ in range(5): + g.replay() + torch.accelerator.synchronize() + _assert_close(out, _ref(a, b), context=f"cudagraph M={M} K={K}") + + +def test_cudagraph_20x_replay(): + torch.manual_seed(42) + a = torch.randn(4, 4096, dtype=torch.bfloat16, device="cuda") + b = torch.randn(64, 4096, dtype=torch.float32, device="cuda") + _gemm(a, b) + torch.accelerator.synchronize() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out = _gemm(a, b) + results = [] + for _ in range(20): + g.replay() + torch.accelerator.synchronize() + results.append(out.clone()) + for i in range(1, len(results)): + torch.testing.assert_close( + results[0], results[i], atol=0, rtol=0, msg=f"Replay {i} differs" + ) + + +def test_cudagraph_input_update(): + torch.manual_seed(42) + a = torch.randn(4, 5120, dtype=torch.bfloat16, device="cuda") + b = torch.randn(128, 5120, dtype=torch.float32, device="cuda") + _gemm(a, b) + torch.accelerator.synchronize() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out = _gemm(a, b) + a.copy_(torch.randn_like(a)) + g.replay() + torch.accelerator.synchronize() + _assert_close(out, _ref(a, b), context="cudagraph input update") + + +# ================================================================= +# Negative tests +# ================================================================= + + +def test_invalid_weight_dtype(): + a = torch.randn(4, 4096, dtype=torch.bfloat16, device="cuda") + b = torch.randn(64, 4096, dtype=torch.bfloat16, device="cuda") + with pytest.raises(ValueError): + _gemm(a, b) + + +def test_invalid_device_cpu(): + a = torch.randn(4, 4096, dtype=torch.bfloat16) + b = torch.randn(64, 4096, dtype=torch.float32) + with pytest.raises(ValueError): + _gemm(a, b) + + +def test_invalid_1d_input(): + a = torch.randn(4096, dtype=torch.bfloat16, device="cuda") + b = torch.randn(64, 4096, dtype=torch.float32, device="cuda") + with pytest.raises(ValueError): + _gemm(a, b) + + +def test_invalid_activation_dtype(): + a = torch.ones(4, 4096, dtype=torch.int32, device="cuda") + b = torch.randn(64, 4096, dtype=torch.float32, device="cuda") + with pytest.raises(ValueError): + _gemm(a, b) + + +def test_invalid_output_dtype(): + a = torch.randn(4, 4096, dtype=torch.bfloat16, device="cuda") + b = torch.randn(64, 4096, dtype=torch.float32, device="cuda") + with pytest.raises(ValueError): + _gemm(a, b, torch.bfloat16) + + +def test_invalid_k_mismatch(): + a = torch.randn(4, 4096, dtype=torch.bfloat16, device="cuda") + b = torch.randn(64, 2048, dtype=torch.float32, device="cuda") + with pytest.raises(ValueError): + _gemm(a, b) + + +def test_invalid_non_contiguous(): + base = torch.randn(4, 8192, dtype=torch.bfloat16, device="cuda") + a = base[:, ::2] + b = torch.randn(64, 4096, dtype=torch.float32, device="cuda") + with pytest.raises(ValueError): + _gemm(a, b) + + +# ================================================================= +# GateLinear dispatch integration +# ================================================================= + + +def _make_gate_linear_fp32w( + monkeypatch, + *, + params_dtype, + out_dtype=torch.float32, + input_size=2048, + output_size=64, +): + monkeypatch.setattr( + "vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", + lambda: 0, + ) + monkeypatch.setattr( + "vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", + lambda: 1, + ) + monkeypatch.setattr( + "vllm.model_executor.parameter.get_tensor_model_parallel_rank", + lambda: 0, + ) + monkeypatch.setattr( + "vllm.model_executor.parameter.get_tensor_model_parallel_world_size", + lambda: 1, + ) + monkeypatch.setattr( + "vllm.model_executor.kernels.linear.cute_dsl.ll_bf16.is_available", + lambda: False, + ) + monkeypatch.setattr( + "vllm.model_executor.kernels.linear.cute_dsl.ll_fp32w.is_available", + lambda: True, + ) + + from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear + + return GateLinear( + input_size=input_size, + output_size=output_size, + bias=False, + out_dtype=out_dtype, + params_dtype=params_dtype, + ).cuda() + + +@pytest.mark.parametrize("x_dtype", [torch.bfloat16, torch.float16, torch.float32]) +def test_gate_linear_uses_ll_fp32w_for_fp32_weight_fast_path(monkeypatch, x_dtype): + gate = _make_gate_linear_fp32w(monkeypatch, params_dtype=torch.float32) + x = torch.randn(4, 2048, dtype=x_dtype, device="cuda") + calls = [] + + def fake_ll_fp32w_gemm(hidden_states, router_weight): + calls.append((hidden_states, router_weight)) + return torch.full( + (hidden_states.shape[0], router_weight.shape[0]), + 1.0, + dtype=torch.float32, + device=hidden_states.device, + ) + + monkeypatch.setattr( + "vllm.model_executor.kernels.linear.cute_dsl.ll_fp32w.LL_FP32W_GEMM_KERNEL", + fake_ll_fp32w_gemm, + ) + out, bias = gate(x) + assert bias is None + assert out.shape == (4, 64) + assert out.dtype == torch.float32 + assert len(calls) == 1 + assert calls[0][0] is x + assert calls[0][1] is gate.weight + + +PREFERRED_FP32W_GATE_CASES = [ + *((6144, 128, M) for M in range(1, 33)), + *((6144, 256, M) for M in (1, 8, 12, 14, 16, 20, 22, 25, 27, 32)), +] + + +@pytest.mark.parametrize(("input_size", "output_size", "M"), PREFERRED_FP32W_GATE_CASES) +def test_gate_linear_uses_ll_fp32w_when_cpp_tier_disabled( + monkeypatch, input_size, output_size, M +): + gate = _make_gate_linear_fp32w( + monkeypatch, + params_dtype=torch.float32, + input_size=input_size, + output_size=output_size, + ) + gate.allow_fp32_router_gemm = False + x = torch.randn(M, input_size, dtype=torch.bfloat16, device="cuda") + calls = [] + + def fake_ll_fp32w_gemm(hidden_states, router_weight): + calls.append((hidden_states, router_weight)) + return torch.full( + (hidden_states.shape[0], router_weight.shape[0]), + 2.0, + dtype=torch.float32, + device=hidden_states.device, + ) + + monkeypatch.setattr( + "vllm.model_executor.kernels.linear.cute_dsl.ll_fp32w.LL_FP32W_GEMM_KERNEL", + fake_ll_fp32w_gemm, + ) + out, _ = gate(x) + assert len(calls) == 1 + assert calls[0][0] is x + assert calls[0][1] is gate.weight + assert torch.all(out == 2.0) + + +@pytest.mark.parametrize("M", [8, 16, 32]) +def test_gate_linear_uses_cpp_fp32_router_for_supported_cpp_shape(monkeypatch, M): + gate = _make_gate_linear_fp32w( + monkeypatch, + params_dtype=torch.float32, + input_size=3072, + output_size=256, + ) + assert gate.allow_fp32_router_gemm + x = torch.randn(M, 3072, dtype=torch.bfloat16, device="cuda") + + def fail_ll_fp32w_gemm(hidden_states, router_weight): + raise AssertionError( + "ll_fp32w_gemm should not bypass fp32_router_gemm_dispatch" + ) + + monkeypatch.setattr( + "vllm.model_executor.kernels.linear.cute_dsl.ll_fp32w.LL_FP32W_GEMM_KERNEL", + fail_ll_fp32w_gemm, + ) + out, _ = gate(x) + assert out.shape == (M, 256) + assert out.dtype == torch.float32 + + +def test_gate_linear_bf16_weight_skips_ll_fp32w(monkeypatch): + gate = _make_gate_linear_fp32w(monkeypatch, params_dtype=torch.bfloat16) + assert not gate.allow_ll_fp32w_gemm + x = torch.randn(4, 2048, dtype=torch.bfloat16, device="cuda") + + def fail_ll_fp32w_gemm(hidden_states, router_weight): + raise AssertionError("ll_fp32w_gemm should not run for bf16 weights") + + monkeypatch.setattr( + "vllm.model_executor.kernels.linear.cute_dsl.ll_fp32w.LL_FP32W_GEMM_KERNEL", + fail_ll_fp32w_gemm, + ) + out, _ = gate(x) + assert out.shape == (4, 64) + assert out.dtype == torch.float32 + + +def test_gate_linear_set_out_dtype_enables_ll_fp32w(monkeypatch): + gate = _make_gate_linear_fp32w( + monkeypatch, params_dtype=torch.float32, out_dtype=None + ) + assert not gate.allow_ll_fp32w_gemm + gate.set_out_dtype(torch.float32) + assert gate.allow_ll_fp32w_gemm + + +def test_gate_linear_non_fp32_out_dtype_disables_ll_fp32w(monkeypatch): + gate = _make_gate_linear_fp32w( + monkeypatch, params_dtype=torch.float32, out_dtype=None + ) + gate.set_out_dtype(torch.bfloat16) + assert not gate.allow_ll_fp32w_gemm + + +def test_gate_linear_m_gt_32_skips_ll_fp32w(monkeypatch): + gate = _make_gate_linear_fp32w(monkeypatch, params_dtype=torch.float32) + x = torch.randn(33, 2048, dtype=torch.bfloat16, device="cuda") + + def fail_ll_fp32w_gemm(hidden_states, router_weight): + raise AssertionError("ll_fp32w_gemm should not run for M > 32") + + monkeypatch.setattr( + "vllm.model_executor.kernels.linear.cute_dsl.ll_fp32w.LL_FP32W_GEMM_KERNEL", + fail_ll_fp32w_gemm, + ) + out, _ = gate(x) + assert out.shape == (33, 64) + assert out.dtype == torch.float32 + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/vllm/model_executor/kernels/linear/cute_dsl/_ll_fp32w_dotprod.py b/vllm/model_executor/kernels/linear/cute_dsl/_ll_fp32w_dotprod.py new file mode 100644 index 0000000000..ddac7e868c --- /dev/null +++ b/vllm/model_executor/kernels/linear/cute_dsl/_ll_fp32w_dotprod.py @@ -0,0 +1,457 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import cutlass +import cutlass.cute as cute +from cuda.bindings.driver import CUstream +from cutlass import const_expr + + +class LLFp32WDotprod: + """FP32-weight router GEMM kernel based on CTA-local dot products. + + This kernel computes C[M, N] = A[M, K] @ B[N, K]^T for bf16/fp16/fp32 + activations, fp32 weights, and fp32 output. It launches one CTA per + small output-column tile, distributes K across CTA threads with vectorized + loads, accumulates one fp32 dot product per token/expert, and reduces + through warp shuffles plus shared memory. + + :param k: Compile-time K dimension specialized into the generated kernel. + :type k: int + :param bs: Threads per CTA and K-stripe width used by the reduction. + :type bs: int + :param main_vec_width: Elements loaded per thread in the main loop. + :type main_vec_width: int + :param tail_vec_width: Elements loaded per thread in the vector tail. + :type tail_vec_width: int + :param use_pdl: Whether to launch with Programmatic Dependent Launch. + :type use_pdl: bool + + :note: Supported A/B data types: + - BFloat16/Float32 + - Float16/Float32 + - Float32/Float32 + :note: Supported accumulator data types: + - Float32 + :note: Supported C data types: + - Float32 + :note: Constraints: + - Inputs must be contiguous row-major 2D tensors. + + :note: K is handled as vectorized main/tail loops plus scalar remainder. + + :compile-key: ``(M, K, bs, epb, act_dtype)`` selects the token count, + hidden size, CTA thread/K-stripe width, experts per CTA, and activation + input type. + """ + + def __init__( + self, + m: int, + k: int, + bs: int = 128, + main_vec_width: int = 8, + tail_vec_width: int = 4, + token_groups: int = 1, + epb: int = 1, + use_pdl: bool = False, + ): + """Initialize the dot-product kernel configuration. + + This configuration fixes the CTA thread count, reduction warp count, + vector widths, and K-loop decomposition used by the generated kernel. + + :param m: Token count M selected by the compile key. + :type m: int + :param k: Hidden size K used to specialize the K-loop decomposition. + :type k: int + :param bs: Threads per CTA and K-stripe width for the reduction. + :type bs: int + :param main_vec_width: Elements loaded per thread in the main loop. + :type main_vec_width: int + :param tail_vec_width: Elements loaded per thread in the vector tail. + :type tail_vec_width: int + :param use_pdl: Whether to launch with Programmatic Dependent Launch. + :type use_pdl: bool + """ + self.m = m + self.m_per_group = m // token_groups + self.bs = bs + self.main_vec_width = main_vec_width + self.tail_vec_width = tail_vec_width + self.token_groups = token_groups + self.epb = epb + self.use_pdl = use_pdl + self.num_warps = bs // cute.arch.WARP_SIZE + self._init_k_tiles(k) + + def _vectorized_elems(self, k_extent: int, vec_width: int) -> int: + vector_tile = vec_width * self.bs + return (k_extent // vector_tile) * vector_tile + + def _init_k_tiles(self, k: int) -> None: + """Split K into vector loops, scalar rounds, and ragged tail.""" + self.k_main_elems = self._vectorized_elems(k, self.main_vec_width) + self.k_after_main = k - self.k_main_elems + self.k_tail_elems = self._vectorized_elems( + self.k_after_main, self.tail_vec_width + ) + self.k_done_all = self.k_main_elems + self.k_tail_elems + self.scalar_rem = k - self.k_done_all + self.ks_full = self.scalar_rem // self.bs + self.ks_part = self.scalar_rem % self.bs + self.k_scalar_full = self.ks_full * self.bs + self.k_part_offset = self.k_done_all + self.k_scalar_full + self.main_tiles = self.k_main_elems // (self.main_vec_width * self.bs) + self.tail_tiles = self.k_tail_elems // (self.tail_vec_width * self.bs) + + @cute.jit + def _vector_dotprod( + self, + acc: cute.Tensor, + tA: cute.Tensor, + tB: cute.Tensor, + M: cutlass.Constexpr, + epb: cutlass.Constexpr, + num_tiles: cutlass.Constexpr, + ): + for tile in cutlass.range_constexpr(num_tiles): + bt0 = tB[0, None, tile] + br0 = cute.make_rmem_tensor_like(bt0) + cute.autovec_copy(bt0, br0) + br0_f32 = br0.load().to(cutlass.Float32) + + if const_expr(epb == 2): + bt1 = tB[1, None, tile] + br1 = cute.make_rmem_tensor_like(bt1) + cute.autovec_copy(bt1, br1) + br1_f32 = br1.load().to(cutlass.Float32) + + for m in cutlass.range_constexpr(M): + at = tA[m, None, tile] + ar = cute.make_rmem_tensor_like(at) + cute.autovec_copy(at, ar) + vec_width: cutlass.Constexpr = cute.size(ar) + for v in cutlass.range_constexpr(vec_width): + av = ar[v].to(cutlass.Float32) + acc[m, 0] = acc[m, 0] + av * br0_f32[v] + if const_expr(epb == 2): + acc[m, 1] = acc[m, 1] + av * br1_f32[v] + + def _make_thread_vector_slice( + self, + gA_vec: cute.Tensor, + gB_vec: cute.Tensor, + tidx: cutlass.Int32, + bs: cutlass.Constexpr, + ): + # A: (M, K_TILE, K_LANE, K_VEC); B: (epb, K_TILE, K_LANE, K_VEC). + tA = cute.logical_divide(gA_vec, (None, (None, bs))) + tB = cute.logical_divide(gB_vec, (None, (None, bs))) + return tA[None, (None, (tidx, None))], tB[None, (None, (tidx, None))] + + def _make_m_slice( + self, + gX: cute.Tensor, + m_offset: cutlass.Int32, + m_extent: cutlass.Constexpr, + ): + return cute.local_tile( + cute.domain_offset((m_offset, 0), gX), + (m_extent, cute.size(gX, mode=[1])), + (0, 0), + ) + + def _make_n_tile( + self, + gX: cute.Tensor, + n_tile_idx: cutlass.Int32, + n_extent: cutlass.Constexpr, + ): + return cute.local_tile( + gX, + (n_extent, cute.size(gX, mode=[1])), + (n_tile_idx, 0), + ) + + def _make_sm_m_slice( + self, + sm: cute.Tensor, + m_offset: cutlass.Int32, + m_extent: cutlass.Constexpr, + epb: cutlass.Constexpr, + num_warps: cutlass.Constexpr, + ): + return cute.local_tile( + cute.domain_offset((m_offset, 0, 0), sm), + (m_extent, epb, num_warps), + (0, 0, 0), + ) + + def _make_k_slice( + self, + gX: cute.Tensor, + k_offset: cutlass.Constexpr, + k_extent: cutlass.Constexpr, + ): + k_layout_extent: cutlass.Constexpr = ( + 1 if const_expr(k_extent == 0) else k_extent + ) + + if const_expr(k_offset == 0): + return cute.local_tile( + gX, (cute.size(gX, mode=[0]), k_layout_extent), (0, 0) + ) + return cute.local_tile( + cute.domain_offset((0, k_offset), gX), + (cute.size(gX, mode=[0]), k_layout_extent), + (0, 0), + ) + + @cute.jit + def _compute_token_group( + self, + gA: cute.Tensor, + gB: cute.Tensor, + sm: cute.Tensor, + n_tile_idx: cutlass.Int32, + tidx: cutlass.Int32, + local_wid: cutlass.Int32, + m_per_group: cutlass.Constexpr, + m_offset: cutlass.Int32, + main_vec_width: cutlass.Constexpr, + tail_vec_width: cutlass.Constexpr, + bs: cutlass.Constexpr, + epb: cutlass.Constexpr, + num_warps: cutlass.Constexpr, + k_main_elems: cutlass.Constexpr, + k_tail_elems: cutlass.Constexpr, + k_done_all: cutlass.Constexpr, + ks_full: cutlass.Constexpr, + ks_part: cutlass.Constexpr, + k_scalar_full: cutlass.Constexpr, + k_part_offset: cutlass.Constexpr, + main_tiles: cutlass.Constexpr, + tail_tiles: cutlass.Constexpr, + ): + # Each CTA reuses A for epb adjacent experts. + acc = cute.make_rmem_tensor((m_per_group, epb), cutlass.Float32) + acc.fill(0.0) + + gA_group = self._make_m_slice(gA, m_offset, m_per_group) + gB_epb = self._make_n_tile(gB, n_tile_idx, epb) + sm_group = self._make_sm_m_slice(sm, m_offset, m_per_group, epb, num_warps) + + if const_expr(k_main_elems > 0): + gA_main = self._make_k_slice(gA_group, 0, k_main_elems) + gB_main = self._make_k_slice(gB_epb, 0, k_main_elems) + gA_vec = cute.logical_divide(gA_main, (None, main_vec_width)) + gB_vec = cute.logical_divide(gB_main, (None, main_vec_width)) + tA, tB = self._make_thread_vector_slice(gA_vec, gB_vec, tidx, bs) + self._vector_dotprod(acc, tA, tB, m_per_group, epb, main_tiles) + + if const_expr(k_tail_elems > 0): + gA_tail = self._make_k_slice(gA_group, k_main_elems, k_tail_elems) + gB_tail = self._make_k_slice(gB_epb, k_main_elems, k_tail_elems) + gA_tail_vec = cute.logical_divide(gA_tail, (None, tail_vec_width)) + gB_tail_vec = cute.logical_divide(gB_tail, (None, tail_vec_width)) + tA_t, tB_t = self._make_thread_vector_slice( + gA_tail_vec, gB_tail_vec, tidx, bs + ) + self._vector_dotprod(acc, tA_t, tB_t, m_per_group, epb, tail_tiles) + + if const_expr(ks_full > 0): + gA_scalar = self._make_k_slice(gA_group, k_done_all, k_scalar_full) + gB_scalar = self._make_k_slice(gB_epb, k_done_all, k_scalar_full) + gA_scalar_vec = cute.logical_divide(gA_scalar, (None, 1)) + gB_scalar_vec = cute.logical_divide(gB_scalar, (None, 1)) + tA_s, tB_s = self._make_thread_vector_slice( + gA_scalar_vec, gB_scalar_vec, tidx, bs + ) + self._vector_dotprod(acc, tA_s, tB_s, m_per_group, epb, ks_full) + + if const_expr(ks_part > 0): + gA_part = self._make_k_slice(gA_group, k_part_offset, ks_part) + gB_part = self._make_k_slice(gB_epb, k_part_offset, ks_part) + if tidx < ks_part: + for local_m in cutlass.range_constexpr(m_per_group): + av = gA_part[local_m, tidx].to(cutlass.Float32) + for e in cutlass.range_constexpr(epb): + bv = gB_part[e, tidx].to(cutlass.Float32) + acc[local_m, e] = acc[local_m, e] + av * bv + + for local_m in cutlass.range_constexpr(m_per_group): + for e in cutlass.range_constexpr(epb): + acc[local_m, e] = cute.arch.warp_reduction_sum(acc[local_m, e]) + + with cute.arch.elect_one(): + for local_m in cutlass.range_constexpr(m_per_group): + for e in cutlass.range_constexpr(epb): + sm_group[local_m, e, local_wid] = acc[local_m, e] + + @cute.jit + def _store_outputs( + self, + gC: cute.Tensor, + sm: cute.Tensor, + n_tile_idx: cutlass.Int32, + N_dim: cutlass.Int32, + tidx: cutlass.Int32, + group_idx: cutlass.Int32, + M: cutlass.Constexpr, + epb: cutlass.Constexpr, + ): + n_tile_layout = cute.make_layout( + (cute.ceil_div(N_dim, epb), epb), stride=(epb, 1) + ) + n_coords = cute.make_tensor(0, n_tile_layout) + + if group_idx == 0 and tidx < M: + for e in cutlass.range_constexpr(epb): + n = n_coords[n_tile_idx, e] + if n < N_dim: + partials = sm[tidx, e, None].load() + gC[tidx, n] = partials.reduce( + cute.ReductionOp.ADD, + init_val=cutlass.Float32(0.0), + reduction_profile=0, + ) + + @cute.jit + def __call__( + self, + gA: cute.Tensor, + gB: cute.Tensor, + gC: cute.Tensor, + N_dim: cutlass.Int32, + stream: CUstream, + ): + """Execute the dot-product GEMM operation in steps: + - Launch one CTA per small output-column tile with ``bs`` threads. + - Keep one FP32 accumulator per token/expert in each thread. + - Traverse K with vectorized main/tail, scalar, and ragged-tail loops + from the precomputed K decomposition. + - Reduce each token accumulator first within the warp, then across + warps through shared memory, and store ``C[:, n]``. + + :param gA: Input tensor A with shape ``[M, K]``. + :type gA: cute.Tensor + :param gB: Input tensor B with shape ``[N, K]``. + :type gB: cute.Tensor + :param gC: Output tensor C with shape ``[M, N]``. + :type gC: cute.Tensor + :param N_dim: Output column count used for the launch grid. + :type N_dim: cutlass.Int32 + :param stream: CUDA stream for asynchronous execution. + :type stream: CUstream + """ + self.kernel( + gA, + gB, + gC, + N_dim, + self.m, + self.m_per_group, + self.main_vec_width, + self.tail_vec_width, + self.bs, + self.num_warps, + self.token_groups, + self.epb, + self.k_main_elems, + self.k_tail_elems, + self.k_done_all, + self.ks_full, + self.ks_part, + self.k_scalar_full, + self.k_part_offset, + self.main_tiles, + self.tail_tiles, + ).launch( + grid=[cute.ceil_div(N_dim, self.epb), 1, 1], + block=[self.bs, self.token_groups, 1], + smem=self.m * self.epb * 4 * self.num_warps, + stream=stream, + use_pdl=self.use_pdl, + min_blocks_per_mp=1, + ) + + @cute.kernel + def kernel( + self, + gA: cute.Tensor, + gB: cute.Tensor, + gC: cute.Tensor, + N_dim: cutlass.Int32, + M: cutlass.Constexpr, + m_per_group: cutlass.Constexpr, + main_vec_width: cutlass.Constexpr, + tail_vec_width: cutlass.Constexpr, + bs: cutlass.Constexpr, + num_warps: cutlass.Constexpr, + token_groups: cutlass.Constexpr, + epb: cutlass.Constexpr, + k_main_elems: cutlass.Constexpr, + k_tail_elems: cutlass.Constexpr, + k_done_all: cutlass.Constexpr, + ks_full: cutlass.Constexpr, + ks_part: cutlass.Constexpr, + k_scalar_full: cutlass.Constexpr, + k_part_offset: cutlass.Constexpr, + main_tiles: cutlass.Constexpr, + tail_tiles: cutlass.Constexpr, + ): + tidx, group_idx, _ = cute.arch.thread_idx() + n_tile_idx, _, _ = cute.arch.block_idx() + + local_wid = tidx // cute.arch.WARP_SIZE + + if const_expr(self.use_pdl): + cute.arch.griddepcontrol_wait() + + smem_red_layout = cute.make_layout( + (M, epb, num_warps), stride=(epb * num_warps, num_warps, 1) + ) + smem = cutlass.utils.SmemAllocator() + sm = smem.allocate_tensor(cutlass.Float32, smem_red_layout, byte_alignment=16) + + if const_expr(token_groups == 1): + m_offset: cutlass.Constexpr = 0 + else: + m_offset = group_idx * m_per_group + self._compute_token_group( + gA, + gB, + sm, + n_tile_idx, + tidx, + local_wid, + m_per_group, + m_offset, + main_vec_width, + tail_vec_width, + bs, + epb, + num_warps, + k_main_elems, + k_tail_elems, + k_done_all, + ks_full, + ks_part, + k_scalar_full, + k_part_offset, + main_tiles, + tail_tiles, + ) + + cute.arch.sync_threads() + self._store_outputs(gC, sm, n_tile_idx, N_dim, tidx, group_idx, M, epb) + + if const_expr(self.use_pdl): + cute.arch.griddepcontrol_launch_dependents() + + +def make_host_fp32w( + k_val: int, bs: int = 128, token_groups: int = 1, epb: int = 1, m: int = 1 +): + return LLFp32WDotprod(m=m, k=k_val, bs=bs, token_groups=token_groups, epb=epb) diff --git a/vllm/model_executor/kernels/linear/cute_dsl/_ll_router_common.py b/vllm/model_executor/kernels/linear/cute_dsl/_ll_router_common.py new file mode 100644 index 0000000000..23fcefdac9 --- /dev/null +++ b/vllm/model_executor/kernels/linear/cute_dsl/_ll_router_common.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import logging +from typing import Any + +import torch + +logger = logging.getLogger(__name__) + +_cutedsl_available: bool | None = None +_cute_ctx: tuple[Any, Any] | None = None + + +def is_cutedsl_available(kernel_name: str) -> bool: + global _cutedsl_available + if _cutedsl_available is not None: + return _cutedsl_available + try: + import cutlass # noqa: F401 + import cutlass.cute # noqa: F401 + + _cutedsl_available = True + except ImportError: + _cutedsl_available = False + logger.info( + "cuteDSL (CUTLASS Python) not available, %s disabled", + kernel_name, + ) + return _cutedsl_available + + +def cute_context(): + global _cute_ctx + if _cute_ctx is not None: + return _cute_ctx + import cutlass.cute as cute + from cuda.bindings.driver import CUstream + + _cute_ctx = (cute, CUstream) + return _cute_ctx + + +def current_cuda_stream(): + _, CUstream = cute_context() + from vllm.utils.torch_utils import current_stream + + return CUstream(current_stream().cuda_stream) + + +def use_pdl() -> bool: + from vllm.platforms import current_platform + + return current_platform.is_arch_support_pdl() + + +def cutlass_dtype(dtype: torch.dtype): + from cutlass import BFloat16, Float16, Float32 + + if dtype == torch.bfloat16: + return BFloat16 + if dtype == torch.float16: + return Float16 + if dtype == torch.float32: + return Float32 + raise ValueError(f"unsupported router GEMM dtype: {dtype}") + + +def make_fake_gemm_tensors( + *, + M, + K, + N, + a_dtype: torch.dtype, + b_dtype: torch.dtype, + divisibility: int, +): + from cutlass import Float32 + from quack.compile_utils import make_fake_tensor + + hidden_states = make_fake_tensor( + cutlass_dtype(a_dtype), (M, K), divisibility=divisibility + ) + router_weight = make_fake_tensor( + cutlass_dtype(b_dtype), (N, K), divisibility=divisibility + ) + output = make_fake_tensor(Float32, (M, N), divisibility=1) + return hidden_states, router_weight, output + + +def validate_common_gemm_inputs( + hidden_states: torch.Tensor, + router_weight: torch.Tensor, + output_dtype: torch.dtype, + *, + op_name: str, + k_multiple: int | None = None, +) -> None: + if hidden_states.dim() != 2 or router_weight.dim() != 2: + raise ValueError("hidden_states and router_weight must be 2D tensors") + if hidden_states.device.type != "cuda" or router_weight.device.type != "cuda": + raise ValueError("hidden_states and router_weight must have device_type=cuda") + if hidden_states.device != router_weight.device: + raise ValueError( + "hidden_states and router_weight must be on the same CUDA device" + ) + if output_dtype != torch.float32: + raise ValueError(f"{op_name} only supports output_dtype=torch.float32") + if hidden_states.shape[1] != router_weight.shape[1]: + raise ValueError( + "hidden_states and router_weight must have matching K dimensions" + ) + if k_multiple is not None and hidden_states.shape[1] % k_multiple != 0: + raise ValueError(f"{op_name} requires K to be divisible by {k_multiple}") + if not hidden_states.is_contiguous() or not router_weight.is_contiguous(): + raise ValueError( + "hidden_states and router_weight must be contiguous row-major inputs" + ) diff --git a/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py b/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py index 564abbd84b..95cb78c6a1 100644 --- a/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py +++ b/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py @@ -11,24 +11,20 @@ from typing import Any, Literal import torch -logger = logging.getLogger(__name__) +from ._ll_router_common import ( + current_cuda_stream, + cute_context, + is_cutedsl_available, + make_fake_gemm_tensors, + use_pdl, + validate_common_gemm_inputs, +) -_cutedsl_available: bool | None = None +logger = logging.getLogger(__name__) def is_available() -> bool: - global _cutedsl_available - if _cutedsl_available is not None: - return _cutedsl_available - try: - import cutlass # noqa: F401 - import cutlass.cute # noqa: F401 - - _cutedsl_available = True - except ImportError: - _cutedsl_available = False - logger.info("cuteDSL (CUTLASS Python) not available, ll_bf16_gemm disabled") - return _cutedsl_available + return is_cutedsl_available("ll_bf16_gemm") # Default configs @@ -96,33 +92,6 @@ def _arch_tuned_configs() -> tuple[ return {}, {} -_cute_ctx = None - - -def _cute(): - global _cute_ctx - if _cute_ctx is not None: - return _cute_ctx - import cutlass.cute as cute - from cuda.bindings.driver import CUstream - - _cute_ctx = (cute, CUstream) - return _cute_ctx - - -def _stream(): - _, CUstream = _cute() - from vllm.utils.torch_utils import current_stream - - return CUstream(current_stream().cuda_stream) - - -def _use_pdl() -> bool: - from vllm.platforms import current_platform - - return current_platform.is_arch_support_pdl() - - class LLBf16Gemm: @dataclass(frozen=True, slots=True) class CompileKey: @@ -165,16 +134,17 @@ class LLBf16Gemm: @staticmethod def _fake_gemm_tensors(*, M, K, N, divisibility: int): - from cutlass import BFloat16, Float32 - from quack.compile_utils import make_fake_tensor - - hidden_states = make_fake_tensor(BFloat16, (M, K), divisibility=divisibility) - router_weight = make_fake_tensor(BFloat16, (N, K), divisibility=divisibility) - output = make_fake_tensor(Float32, (M, N), divisibility=1) - return hidden_states, router_weight, output + return make_fake_gemm_tensors( + M=M, + K=K, + N=N, + a_dtype=torch.bfloat16, + b_dtype=torch.bfloat16, + divisibility=divisibility, + ) def _compile_splitk(self, compile_key: CompileKey) -> None: - cute, _ = _cute() + cute, _ = cute_context() from ._ll_bf16_splitk import LLBf16SplitK hidden_states, router_weight, output = self._fake_gemm_tensors( @@ -189,14 +159,14 @@ class LLBf16Gemm: num_stages=compile_key.num_stages, num_dma_warps=4, split_k=compile_key.split_k, - use_pdl=_use_pdl(), + use_pdl=use_pdl(), ) compiled = cute.compile( gemm, hidden_states, router_weight, output, - _stream(), + current_cuda_stream(), options="--enable-tvm-ffi", ) self._splitk_cache[(compile_key.split_k, compile_key.num_stages)] = compiled @@ -207,7 +177,7 @@ class LLBf16Gemm: ) def _compile_dotprod(self, compile_key: CompileKey) -> None: - cute, _ = _cute() + cute, _ = cute_context() from ._ll_bf16_dotprod import LLBf16Dotprod N = cute.sym_int() @@ -221,7 +191,7 @@ class LLBf16Gemm: gemm = LLBf16Dotprod( k=compile_key.K, bs=compile_key.bs, - use_pdl=_use_pdl(), + use_pdl=use_pdl(), prefetch_pdl_weights=self._prefetch_pdl_weights, ) compiled = cute.compile( @@ -232,7 +202,7 @@ class LLBf16Gemm: compile_key.M, compile_key.K, 1, # runtime N placeholder for fake-tensor compile - _stream(), + current_cuda_stream(), options="--enable-tvm-ffi --ptxas-options -maxrregcount=64", ) self._compiled_cache[(compile_key.M, compile_key.K, compile_key.bs)] = compiled @@ -269,34 +239,19 @@ class LLBf16Gemm: router_weight: torch.Tensor, output_dtype: torch.dtype, ) -> None: - if hidden_states.dim() != 2 or router_weight.dim() != 2: - raise ValueError("hidden_states and router_weight must be 2D tensors") if ( hidden_states.dtype != torch.bfloat16 or router_weight.dtype != torch.bfloat16 ): raise ValueError("hidden_states and router_weight must have dtype=bfloat16") - if hidden_states.device.type != "cuda" or router_weight.device.type != "cuda": - raise ValueError( - "hidden_states and router_weight must have device_type=cuda" - ) - if hidden_states.device != router_weight.device: - raise ValueError( - "hidden_states and router_weight must be on the same CUDA device" - ) - if output_dtype != torch.float32: - raise ValueError("ll_bf16_gemm only supports output_dtype=torch.float32") - if hidden_states.shape[1] != router_weight.shape[1]: - raise ValueError( - "hidden_states and router_weight must have matching K dimensions" - ) # Kernels use vectorized bf16 loads and require 16-byte row alignment. - if hidden_states.shape[1] % 8 != 0: - raise ValueError("ll_bf16_gemm requires K to be divisible by 8") - if not hidden_states.is_contiguous() or not router_weight.is_contiguous(): - raise ValueError( - "hidden_states and router_weight must be contiguous row-major inputs" - ) + validate_common_gemm_inputs( + hidden_states, + router_weight, + output_dtype, + op_name="ll_bf16_gemm", + k_multiple=8, + ) def __call__( self, @@ -320,7 +275,7 @@ class LLBf16Gemm: self.compile(compile_key) kernel = self._compiled_cache[dotprod_cache_key] - stream = _stream() + stream = current_cuda_stream() output = torch.empty(M, N, dtype=output_dtype, device=hidden_states.device) if compile_key.backend == "splitk": kernel(hidden_states, router_weight, output, stream, 1.0) diff --git a/vllm/model_executor/kernels/linear/cute_dsl/ll_fp32w.py b/vllm/model_executor/kernels/linear/cute_dsl/ll_fp32w.py new file mode 100644 index 0000000000..edc3073a1f --- /dev/null +++ b/vllm/model_executor/kernels/linear/cute_dsl/ll_fp32w.py @@ -0,0 +1,197 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import logging +import math +from collections.abc import Iterable +from dataclasses import dataclass +from typing import Any + +import torch + +from vllm.model_executor.warmup.jit_warmup import VllmJitKernel, zip_inputs + +from ._ll_router_common import ( + current_cuda_stream, + cute_context, + is_cutedsl_available, + make_fake_gemm_tensors, + use_pdl, + validate_common_gemm_inputs, +) + +logger = logging.getLogger(__name__) + + +def is_available() -> bool: + return is_cutedsl_available("ll_fp32w_gemm") + + +_DEFAULT_DOTPROD_CONFIG = (192, 1, 1) +_DEFAULT_DOTPROD_GROUPED_CONFIG = (192, 2, 1) + + +_TUNED_DOTPROD_CONFIGS: dict[tuple[int, int], dict[int, tuple[int, int, int]]] = { + (6144, 128): { + **{ + M: (384, 1, 1) + for M in ( + *range(1, 6), + *range(9, 14, 2), + 17, + *range(23, 32, 2), + ) + }, + **{M: (256, 1, 1) for M in (7, 15, *range(19, 22, 2))}, + **{M: (384, 2, 1) for M in range(6, 12, 2)}, + **{M: (192, 2, 1) for M in (*range(12, 26, 2), 32)}, + **{M: (128, 2, 1) for M in range(26, 32, 2)}, + }, + (6144, 256): { + **{M: (384, 1, 1) for M in range(1, 4)}, + **{M: (128, 1, 1) for M in (*range(4, 7, 2), 11)}, + **{M: (128, 1, 2) for M in range(5, 11, 2)}, + **{M: (128, 2, 2) for M in range(8, 33, 2)}, + **{M: (256, 1, 2) for M in range(13, 33, 2)}, + }, +} +_SUPPORTED_ACTIVATION_DTYPES = (torch.bfloat16, torch.float16, torch.float32) + + +class LLFp32WGemm(VllmJitKernel["LLFp32WGemm.CompileKey"]): + # Dot-prod is specialized on M/K, tile shape, and activation dtype. + @dataclass(frozen=True, slots=True) + class CompileKey: + m: int + k: int + bs: int + a_dtype: torch.dtype + token_groups: int = 1 + epb: int = 1 + + @staticmethod + def kernel(compile_key: CompileKey) -> Any: + from ._ll_fp32w_dotprod import LLFp32WDotprod + + return LLFp32WDotprod( + m=compile_key.m, + k=compile_key.k, + bs=compile_key.bs, + token_groups=compile_key.token_groups, + epb=compile_key.epb, + use_pdl=use_pdl(), + ) + + def dispatch( # type: ignore[override] + self, + *, + M: int, + K: int, + N: int, + a_dtype: torch.dtype, + ) -> CompileKey: + default_config = ( + _DEFAULT_DOTPROD_GROUPED_CONFIG + if N <= 128 and M >= 12 and M % 2 == 0 + else _DEFAULT_DOTPROD_CONFIG + ) + tuned_configs = _TUNED_DOTPROD_CONFIGS.get((K, N)) + config = ( + tuned_configs.get(M, default_config) if tuned_configs else default_config + ) + return self.CompileKey( + m=M, + k=K, + bs=config[0], + a_dtype=a_dtype, + token_groups=config[1], + epb=config[2], + ) + + def get_warmup_keys( + self, + *, + shapes: Iterable[tuple[int, int]], + m_values: Iterable[int], + a_dtypes: Iterable[torch.dtype] = (torch.bfloat16,), + ) -> list[CompileKey]: + return self._trace_dispatch(self.dispatch)( + zip_inputs(*(dict(K=K, N=N) for K, N in shapes)), + M=tuple(m_values), + a_dtype=tuple(a_dtypes), + ) + + def compile(self, compile_key: CompileKey) -> None: + if compile_key in self._compiled_cache: + return + if compile_key.token_groups not in (1, 2): + raise ValueError("ll_fp32w_gemm supports token_groups 1 or 2") + if compile_key.epb not in (1, 2): + raise ValueError("ll_fp32w_gemm supports epb 1 or 2") + if compile_key.m % compile_key.token_groups != 0: + raise ValueError("M must be divisible by token_groups") + + cute, _ = cute_context() + N = cute.sym_int() + stride_divisibility = math.gcd(8, compile_key.k) + hidden_states, router_weight, output = make_fake_gemm_tensors( + M=compile_key.m, + K=compile_key.k, + N=N, + a_dtype=compile_key.a_dtype, + b_dtype=torch.float32, + divisibility=stride_divisibility, + ) + compiled = cute.compile( + self.kernel(compile_key), + hidden_states, + router_weight, + output, + 1, # runtime N placeholder for fake-tensor compile + current_cuda_stream(), + options="--enable-tvm-ffi", + ) + self._compiled_cache[compile_key] = compiled + + @staticmethod + def _validate_inputs( + hidden_states: torch.Tensor, + router_weight: torch.Tensor, + output_dtype: torch.dtype, + ) -> None: + if hidden_states.dtype not in _SUPPORTED_ACTIVATION_DTYPES: + raise ValueError("hidden_states must have dtype bf16, fp16, or fp32") + if router_weight.dtype != torch.float32: + raise ValueError("router_weight must have dtype=float32") + validate_common_gemm_inputs( + hidden_states, + router_weight, + output_dtype, + op_name="ll_fp32w_gemm", + ) + + def __call__( + self, + hidden_states: torch.Tensor, # [M, K] bf16/fp16/fp32 + router_weight: torch.Tensor, # [N, K] fp32 + output_dtype: torch.dtype = torch.float32, + ) -> torch.Tensor: # [M, N] fp32 + self._validate_inputs(hidden_states, router_weight, output_dtype) + + M, K = hidden_states.shape + N = router_weight.shape[0] + compile_key = self.dispatch(M=M, K=K, N=N, a_dtype=hidden_states.dtype) + compiled = self._get_or_compile( + compile_key, + runtime_context={"M": M, "K": K, "N": N}, + ) + + stream = current_cuda_stream() + output = torch.empty(M, N, dtype=output_dtype, device=hidden_states.device) + compiled(hidden_states, router_weight, output, N, stream) + return output + + +LL_FP32W_GEMM_KERNEL = LLFp32WGemm() diff --git a/vllm/model_executor/layers/fused_moe/router/gate_linear.py b/vllm/model_executor/layers/fused_moe/router/gate_linear.py index 8c4961f269..13eff3889c 100644 --- a/vllm/model_executor/layers/fused_moe/router/gate_linear.py +++ b/vllm/model_executor/layers/fused_moe/router/gate_linear.py @@ -22,14 +22,16 @@ class GateLinear(ReplicatedLinear): """MoE gate linear layer with multi-tier GEMM dispatch: 1. cuteDSL ll_bf16_gemm (SM90+, M<=16, bf16 in, fp32 out, - K divisible by 8) + any shape w/ K divisible by 8) 2. fp32 specialized kernel (SM90+ or gfx950, bf16/fp32 in, fp32 out, M<=32, model-specific shapes) - 3. bf16x3 CuteDSL kernel (SM100, bf16 in, fp32 weight) - 4. ROCm bf16x3 router GEMM (gfx950, bf16 in, fp32 weight, fp32 out, + 3. cuteDSL ll_fp32w_gemm (SM90+, M<=32, bf16/fp32 in, fp32 out, + any shape) + 4. bf16x3 CuteDSL kernel (SM100, bf16 in, fp32 weight) + 5. ROCm bf16x3 router GEMM (gfx950, bf16 in, fp32 weight, fp32 out, M>=2048) - 5. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype) - 6. F.linear via ReplicatedLinear (ultimate fallback) + 6. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype) + 7. F.linear via ReplicatedLinear (ultimate fallback) The ``out_dtype`` attribute is mutable and can be set after init (e.g. when the required dtype depends on the expert quantization @@ -172,20 +174,50 @@ class GateLinear(ReplicatedLinear): self.quant_method = _GateLinearMethod() logger.info_once("Enabled ROCm BF16x3 router GEMM.") - # cuteDSL ll_bf16_gemm eligibility. Any dims supported, but SM90+ required bc: + # cuteDSL ll_bf16_gemm/ll_fp32w_gemm eligibility. Any dims supported, but + # SM90+ required because: # 1. PDL support. Both dot-product and split-K kernels. # 2. Thread Block Clusters. Split-K kernel for cross-CTA reduction. self.allow_ll_bf16_gemm = False + self.allow_ll_fp32w_gemm = False if can_use_specialized_kernels: from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ( - is_available, + is_available as is_ll_bf16_gemm_available, + ) + from vllm.model_executor.kernels.linear.cute_dsl.ll_fp32w import ( + is_available as is_ll_fp32w_gemm_available, ) self.allow_ll_bf16_gemm = ( - self.weight.dtype == torch.bfloat16 + is_ll_bf16_gemm_available() + and self.weight.dtype == torch.bfloat16 and self.out_dtype == torch.float32 - and is_available() ) + self.allow_ll_fp32w_gemm = ( + is_ll_fp32w_gemm_available() + and self.weight.dtype == torch.float32 + and self.out_dtype == torch.float32 + ) + + self._register_ll_fp32w_warmup() + + def _register_ll_fp32w_warmup(self) -> None: + # Based on multi-tier GEMM dispatch. See above. + # This avoids compiling LL_FP32W_GEMM_KERNEL if not used. + if not self.allow_ll_fp32w_gemm or self.allow_fp32_router_gemm: + return + + from vllm.model_executor.kernels.linear.cute_dsl.ll_fp32w import ( + LL_FP32W_GEMM_KERNEL, + ) + + # FP16 remains runtime-supported; JIT monitoring reports unexpected use. + # The kernel supports it, but there is no FP16 activation use case yet. + LL_FP32W_GEMM_KERNEL.register_warmup( + shapes=((self.weight.shape[1], self.weight.shape[0]),), + m_values=range(1, self.FP32_MAX_TOKENS + 1), + a_dtypes=(torch.bfloat16, torch.float32), + ) def set_out_dtype(self, out_dtype: torch.dtype) -> None: """Set output dtype for the router logits after init. @@ -211,14 +243,23 @@ class GateLinear(ReplicatedLinear): if self.allow_specialized_router_gemm: from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ( - is_available, + is_available as is_ll_bf16_gemm_available, + ) + from vllm.model_executor.kernels.linear.cute_dsl.ll_fp32w import ( + is_available as is_ll_fp32w_gemm_available, ) self.allow_ll_bf16_gemm = ( - self.weight.dtype == torch.bfloat16 + is_ll_bf16_gemm_available() + and self.weight.dtype == torch.bfloat16 and out_dtype == torch.float32 - and is_available() ) + self.allow_ll_fp32w_gemm = ( + is_ll_fp32w_gemm_available() + and self.weight.dtype == torch.float32 + and out_dtype == torch.float32 + ) + self._register_ll_fp32w_warmup() def _return( self, output: torch.Tensor @@ -228,7 +269,8 @@ class GateLinear(ReplicatedLinear): def forward( self, x: torch.Tensor ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]: - # Tier 1: cuteDSL ll_bf16_gemm (SM90+, any dims) + # Tier 1: cuteDSL ll_bf16_gemm (SM90+, bf16 in, fp32 out, + # any dims w/ K divisible by 8) if ( self.allow_ll_bf16_gemm and x.shape[0] <= self.LL_BF16_MAX_TOKENS @@ -241,7 +283,7 @@ class GateLinear(ReplicatedLinear): output = ll_bf16_gemm(x, self.weight) return self._return(output) - # Tier 4: ROCm bf16x3 router GEMM. Checked before tier 2 because both + # Tier 5: ROCm bf16x3 router GEMM. Checked before tier 2 because both # tiers claim the same gfx950 shapes and tier 2 returns # unconditionally; the custom op below picks between them on the # runtime num_tokens. @@ -270,7 +312,20 @@ class GateLinear(ReplicatedLinear): ) return self._return(output) - # Tier 3: bf16x3 CuteDSL kernel for fp32 router weights + # Tier 3: cuteDSL ll_fp32w_gemm (SM90+, bf16/fp32 in, fp32 out) + if ( + self.allow_ll_fp32w_gemm + and x.shape[0] <= 32 + and x.dtype in (torch.bfloat16, torch.float16, torch.float32) + ): + from vllm.model_executor.kernels.linear.cute_dsl.ll_fp32w import ( + LL_FP32W_GEMM_KERNEL, + ) + + output = LL_FP32W_GEMM_KERNEL(x, self.weight) + return self._return(output) + + # Tier 4: bf16x3 CuteDSL kernel for fp32 router weights if self.allow_bf16x3_router_gemm and x.dtype == torch.bfloat16: from vllm.model_executor.layers.fused_moe.router.bf16x3_router_gemm_cutedsl import ( # noqa: E501 bf16x3_router_gemm, @@ -279,12 +334,12 @@ class GateLinear(ReplicatedLinear): output = bf16x3_router_gemm(x, self.weight) return self._return(output) - # Tier 5: cuBLAS bf16→fp32 + # Tier 6: cuBLAS bf16→fp32 if self.allow_cublas_router_gemm and x.dtype == torch.bfloat16: output = torch.mm(x, self.weight.T, out_dtype=torch.float32) return self._return(output) - # Tier 6: F.linear (ReplicatedLinear) + # Tier 7: F.linear (ReplicatedLinear) if ( self.out_dtype is not None and self.is_unquantized