diff --git a/tests/kernels/moe/test_block_fp8.py b/tests/kernels/moe/test_block_fp8.py index 77b638a7642..0a4fd90f093 100644 --- a/tests/kernels/moe/test_block_fp8.py +++ b/tests/kernels/moe/test_block_fp8.py @@ -17,6 +17,7 @@ from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.fused_moe import ( fused_experts, fused_topk, + override_config, ) from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.all2all_utils import ( @@ -32,6 +33,9 @@ from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import ( from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import ( TritonOrDeepGemmExperts, ) +from vllm.model_executor.layers.fused_moe.fused_moe import ( + try_get_optimal_moe_config, +) from vllm.model_executor.layers.quantization.utils.fp8_utils import ( per_token_group_quant_fp8, ) @@ -260,6 +264,57 @@ def test_w8a8_block_fp8_fused_moe( torch.testing.assert_close(m_out, ref_out_m, atol=tol, rtol=tol) +@pytest.mark.parametrize("M", [24, 1024]) +@pytest.mark.parametrize("modular", [False, True]) +@torch.inference_mode() +def test_h200_block_fp8_up_and_down_tma_matches_pointer(M, modular, workspace_init): + if "H200" not in torch.cuda.get_device_name(): + pytest.skip("This block FP8 TMA test targets H200") + + E, N, K, topk = 256, 256, 128, 8 + torch.manual_seed(123) + a = torch.randn((M, K), dtype=torch.bfloat16) / 10 + score = torch.randn((M, E), dtype=torch.bfloat16) + w1, w2, quant_config = make_test_quant_config( + E, N, K, torch.bfloat16, torch.float8_e4m3fn, block_shape=[128, 128] + ) + topk_weights, topk_ids, _ = fused_topk(a, score.float(), topk, False) + + config_args = (w1.shape, w2.shape, topk, "fp8_w8a8", M) + config = try_get_optimal_moe_config(*config_args, block_shape=[128, 128]) + assert w2.is_contiguous() + + if modular: + moe = modular_triton_fused_moe(make_dummy_moe_config(), quant_config) + + def run(): + return moe.apply( + a, + w1, + w2, + topk_weights, + topk_ids, + activation=MoEActivation.SILU, + apply_router_weight_on_input=False, + expert_map=None, + global_num_experts=E, + ) + + else: + + def run(): + return fused_experts( + a, w1, w2, topk_weights, topk_ids, quant_config=quant_config + ) + + with set_current_vllm_config(vllm_config): + with override_config({**config, "USE_TMA": False}): + pointer_output = run() + tma_output = run() + + torch.testing.assert_close(tma_output, pointer_output, atol=0, rtol=0) + + @pytest.mark.parametrize("clamp_limit", [None, 10.0]) @torch.inference_mode() def test_w8a8_block_fp8_clamp_skips_fused_silu_quant( diff --git a/tests/kernels/moe/test_moe.py b/tests/kernels/moe/test_moe.py index d279c15c2c7..26e9f87f101 100644 --- a/tests/kernels/moe/test_moe.py +++ b/tests/kernels/moe/test_moe.py @@ -75,7 +75,7 @@ def test_triton_moe_launcher_passes_scalar_scale_as_pointer(monkeypatch) -> None class FakeKernel: def __getitem__(self, grid): def launch(*args, **kwargs) -> None: - captured["a_scale"] = args[4] + captured["a_scale"] = args[5] return launch diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py index e06288a3e92..7e9f0e28f6a 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe.py @@ -299,6 +299,7 @@ def fused_moe_kernel( # Pointers to matrices a_ptr, b_ptr, + b_desc, c_ptr, b_bias_ptr, a_scale_ptr, @@ -351,6 +352,7 @@ def fused_moe_kernel( SWAP_AB: tl.constexpr, # Tensor-descriptor path for the A gather and B load in the K-loop. USE_TD: tl.constexpr = False, + USE_TMA_B: tl.constexpr = False, ): """Implements the fused computation for a Mixture of Experts (MOE) using token and expert matrices. @@ -441,6 +443,7 @@ def fused_moe_kernel( offs_k = tl.arange(0, BLOCK_SIZE_K) # TD gather and the SWAP_AB accumulator layout are mutually exclusive. tl.static_assert(not (USE_TD and SWAP_AB)) + tl.static_assert(not (USE_TD and USE_TMA_B)) if USE_TD: # ``tt.descriptor_gather`` requires block_shape[0] == 1 and i32 idx. m_td = num_valid_tokens // top_k @@ -523,16 +526,26 @@ def fused_moe_kernel( b = b_desc.load([pid_n * BLOCK_SIZE_N, k * BLOCK_SIZE_K]).T elif SWAP_AB: a_mask = (offs_k[:, None] < K - k * BLOCK_SIZE_K) & token_mask[None, :] - b_mask = offs_k[None, :] < K - k * BLOCK_SIZE_K a = tl.load(a_ptrs, mask=a_mask, other=0.0) - b = tl.load(b_ptrs, mask=b_mask, other=0.0) + if not USE_TMA_B: + b_mask = offs_k[None, :] < K - k * BLOCK_SIZE_K + b = tl.load(b_ptrs, mask=b_mask, other=0.0) else: a = tl.load( a_ptrs, mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K), other=0.0, ) - b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) + if not USE_TMA_B: + b = tl.load( + b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0 + ) + if USE_TMA_B: + b = b_desc.load( + [off_experts.to(tl.int32), pid_n * BLOCK_SIZE_N, k * BLOCK_SIZE_K] + ).reshape((BLOCK_SIZE_N, BLOCK_SIZE_K)) + if not SWAP_AB: + b = b.T # We accumulate along the K dimension. if use_int8_w8a16: accumulator = tl.dot(a, b.to(compute_type), acc=accumulator) @@ -842,6 +855,26 @@ def invoke_fused_moe_triton_kernel( BLOCK_SIZE_K = config.pop("BLOCK_SIZE_K") if block_shape is not None: BLOCK_SIZE_K = min(BLOCK_SIZE_K, min(block_shape[0], block_shape[1])) + use_tma_b = ( + bool(config.pop("USE_TMA", True)) + and use_fp8_w8a8 + and block_shape is not None + and all(block_shape) + and current_platform.is_cuda() + and current_platform.is_device_capability(90) + and B.device.type == "cuda" + and B.is_contiguous() + and B.size(1) % config["BLOCK_SIZE_N"] == 0 + and B.size(2) % BLOCK_SIZE_K == 0 + ) + b_desc = None + if use_tma_b: + from triton.tools.tensor_descriptor import TensorDescriptor + + set_triton_allocator(B.device) + b_desc = TensorDescriptor.from_tensor( + B, [1, config["BLOCK_SIZE_N"], BLOCK_SIZE_K] + ) if use_td and A.size(1) % BLOCK_SIZE_K != 0: # TD gather/load feeding tl.dot with a non-block-aligned K # miscompiles (~74% of output elements wrong) on real HW; @@ -864,6 +897,7 @@ def invoke_fused_moe_triton_kernel( fused_moe_kernel[grid]( A, B, + b_desc, C, B_bias, A_scale, @@ -906,6 +940,7 @@ def invoke_fused_moe_triton_kernel( BLOCK_SIZE_K=BLOCK_SIZE_K, SWAP_AB=SWAP_AB, USE_TD=use_td, + USE_TMA_B=use_tma_b, **config, )