From fe284ea17d606ceba4a5d3d55d3e3d9a34f478d9 Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Thu, 17 Sep 2026 09:35:34 -0400 Subject: [PATCH] [BugFxi] Fix DeepGEMM FP8 workspace over allocation (#53914) Signed-off-by: Lucas Wilkinson Signed-off-by: Lucas Wilkinson Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> --- tests/kernels/moe/test_deepgemm.py | 40 +++++++++++++++++++ .../layers/fused_moe/deep_gemm_utils.py | 3 +- .../layers/fused_moe/experts/deep_gemm_moe.py | 33 ++++++++++++++- 3 files changed, 73 insertions(+), 3 deletions(-) diff --git a/tests/kernels/moe/test_deepgemm.py b/tests/kernels/moe/test_deepgemm.py index f6df64c51b..cee0faa592 100644 --- a/tests/kernels/moe/test_deepgemm.py +++ b/tests/kernels/moe/test_deepgemm.py @@ -26,6 +26,10 @@ from vllm.model_executor.layers.fused_moe.config import ( ) from vllm.model_executor.layers.fused_moe.deep_gemm_utils import ( deepgemm_moe_permute, + ep_gather, +) +from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import ( + _fp8_workspace_shape, ) from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import ( TritonOrDeepGemmExperts, @@ -43,6 +47,42 @@ from vllm.utils.deep_gemm import ( BLOCK_SIZE = [128, 128] +def test_ep_gather_uses_64_bit_row_offsets(): + hidden_size = 4096 + source_row = 1 << 18 + input_tensor = torch.empty( + (source_row + 1, hidden_size), + device="cuda", + dtype=torch.bfloat16, + ) + input_tensor[source_row].fill_(1) + output = torch.empty((1, hidden_size), device="cuda", dtype=torch.bfloat16) + + ep_gather( + input_tensor=input_tensor, + recv_topk_ids=torch.zeros((1, 1), device="cuda", dtype=torch.int64), + recv_topk_weight=torch.ones((1, 1), device="cuda", dtype=torch.float32), + input_index=torch.full((1, 1), source_row, device="cuda", dtype=torch.int32), + expert_map=None, + output_tensor=output, + ) + + torch.testing.assert_close(output, torch.ones_like(output), rtol=0, atol=0) + + +@pytest.mark.parametrize("workspace_dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("num_columns", [2048, 6144, 6145]) +def test_fp8_workspace_shape(workspace_dtype, num_columns): + num_rows = 17 + shape = _fp8_workspace_shape(num_rows, num_columns, workspace_dtype) + + allocated_bytes = math.prod(shape) * workspace_dtype.itemsize + required_bytes = num_rows * num_columns * torch.float8_e4m3fn.itemsize + + assert allocated_bytes >= required_bytes + assert allocated_bytes - required_bytes < num_rows * workspace_dtype.itemsize + + @pytest.mark.skipif(not is_deep_gemm_supported(), reason="Requires deep_gemm kernels") def test_deepgemm_moe_permute_initializes_padding_scales(workspace_init): hidden_states = torch.randn(2, 128, device="cuda", dtype=torch.bfloat16) diff --git a/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py b/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py index 4ea864eaac..33a2296955 100644 --- a/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py +++ b/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py @@ -393,12 +393,13 @@ def _fwd_kernel_ep_gather( source_token_index = tl.load( input_index + cur_token * input_index_stride0 + topk_index ) + source_token_index_i64 = source_token_index.to(tl.int64) acc_weight = tl.load( recv_topk_weight + cur_token * recv_topk_weight_stride0 + topk_index ) tmp = tl.load( input_tensor - + source_token_index * input_tensor_stride0 + + source_token_index_i64 * input_tensor_stride0 + cur_block * BLOCK_D + off_d ) diff --git a/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py index fe95b36858..699b029bad 100644 --- a/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py @@ -50,6 +50,21 @@ from vllm.utils.import_utils import has_deep_gemm logger = init_logger(__name__) +def _fp8_workspace_shape( + num_rows: int, num_columns: int, workspace_dtype: torch.dtype +) -> tuple[int, int]: + """Size an FP8 byte buffer stored in a workspace of another dtype.""" + bytes_per_workspace_element = workspace_dtype.itemsize + fp8_columns_per_workspace_element = ( + bytes_per_workspace_element // torch.float8_e4m3fn.itemsize + ) + assert bytes_per_workspace_element % torch.float8_e4m3fn.itemsize == 0 + return ( + num_rows, + -(-num_columns // fp8_columns_per_workspace_element), + ) + + def _valid_deep_gemm_shape(M: int, N: int, K: int) -> bool: align = get_mk_alignment_for_contiguous_layout()[0] return align <= M and N % align == 0 and K % align == 0 @@ -216,7 +231,14 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular): assert M_sum % align_used == 0 activation_out_dim = self.adjust_N_for_activation(N, activation) - workspace1 = (M_sum, max(activation_out_dim, K)) + # workspace1 is allocated in the activation dtype by the workspace + # manager, but is only ever viewed and used as FP8 in apply(). Size it + # by bytes instead of reserving one BF16/FP16 element per FP8 element. + workspace1 = _fp8_workspace_shape( + M_sum, + max(activation_out_dim, K), + self.workspace_dtype(self.moe_config.in_dtype), + ) workspace2 = (M_sum, max(N, K)) output = (M, K) return (workspace1, workspace2, output) @@ -480,7 +502,14 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular): assert M_sum % align_used == 0 activation_out_dim = self.adjust_N_for_activation(N, activation) - workspace1 = (M_sum, max(activation_out_dim, K)) + # workspace1 holds the permuted and requantized FP8 activations. The + # workspace manager allocates it in the model activation dtype, so + # account for the dtype sizes rather than overallocating BF16/FP16. + workspace1 = _fp8_workspace_shape( + M_sum, + max(activation_out_dim, K), + self.workspace_dtype(self.moe_config.in_dtype), + ) workspace2 = (M_sum, max(N, K)) output = (M, K) return (workspace1, workspace2, output)