[BugFxi] Fix DeepGEMM FP8 workspace over allocation (#53914)
Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
Lucas Wilkinson committed
Sep 17, 2026 at 09:35 UTC
fe284ea17d606ceba4a5d3d55d3e3d9a34f478d9
3 files changed
+73
-3
tests/kernels/moe/test_deepgemm.py
+40
@@ -26,6 +26,10 @@ from vllm.model_executor.layers.fused_moe.config import (
26
)
27
from vllm.model_executor.layers.fused_moe.deep_gemm_utils import (
28
deepgemm_moe_permute,
29
+ ep_gather,
30
+)
31
+from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import (
32
+ _fp8_workspace_shape,
33
)
34
from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import (
35
TritonOrDeepGemmExperts,
@@ -43,6 +47,42 @@ from vllm.utils.deep_gemm import (
47
BLOCK_SIZE = [128, 128]
48
49
50
+def test_ep_gather_uses_64_bit_row_offsets():
51
+ hidden_size = 4096
52
+ source_row = 1 << 18
53
+ input_tensor = torch.empty(
54
+ (source_row + 1, hidden_size),
55
+ device="cuda",
56
+ dtype=torch.bfloat16,
57
+ )
58
+ input_tensor[source_row].fill_(1)
59
+ output = torch.empty((1, hidden_size), device="cuda", dtype=torch.bfloat16)
60
+
61
+ ep_gather(
62
+ input_tensor=input_tensor,
63
+ recv_topk_ids=torch.zeros((1, 1), device="cuda", dtype=torch.int64),
64
+ recv_topk_weight=torch.ones((1, 1), device="cuda", dtype=torch.float32),
65
+ input_index=torch.full((1, 1), source_row, device="cuda", dtype=torch.int32),
66
+ expert_map=None,
67
+ output_tensor=output,
68
+ )
69
+
70
+ torch.testing.assert_close(output, torch.ones_like(output), rtol=0, atol=0)
71
+
72
+
73
+@pytest.mark.parametrize("workspace_dtype", [torch.float16, torch.bfloat16])
74
+@pytest.mark.parametrize("num_columns", [2048, 6144, 6145])
75
+def test_fp8_workspace_shape(workspace_dtype, num_columns):
76
+ num_rows = 17
77
+ shape = _fp8_workspace_shape(num_rows, num_columns, workspace_dtype)
78
+
79
+ allocated_bytes = math.prod(shape) * workspace_dtype.itemsize
80
+ required_bytes = num_rows * num_columns * torch.float8_e4m3fn.itemsize
81
+
82
+ assert allocated_bytes >= required_bytes
83
+ assert allocated_bytes - required_bytes < num_rows * workspace_dtype.itemsize
84
+
85
+
86
@pytest.mark.skipif(not is_deep_gemm_supported(), reason="Requires deep_gemm kernels")
87
def test_deepgemm_moe_permute_initializes_padding_scales(workspace_init):
88
hidden_states = torch.randn(2, 128, device="cuda", dtype=torch.bfloat16)
vllm/model_executor/layers/fused_moe/deep_gemm_utils.py
+2
-1
@@ -393,12 +393,13 @@ def _fwd_kernel_ep_gather(
393
source_token_index = tl.load(
394
input_index + cur_token * input_index_stride0 + topk_index
395
)
396
+ source_token_index_i64 = source_token_index.to(tl.int64)
397
acc_weight = tl.load(
398
recv_topk_weight + cur_token * recv_topk_weight_stride0 + topk_index
399
)
400
tmp = tl.load(
401
input_tensor
401
- + source_token_index * input_tensor_stride0
402
+ + source_token_index_i64 * input_tensor_stride0
403
+ cur_block * BLOCK_D
404
+ off_d
405
)
vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py
+31
-2
@@ -50,6 +50,21 @@ from vllm.utils.import_utils import has_deep_gemm
50
logger = init_logger(__name__)
51
52
53
+def _fp8_workspace_shape(
54
+ num_rows: int, num_columns: int, workspace_dtype: torch.dtype
55
+) -> tuple[int, int]:
56
+ """Size an FP8 byte buffer stored in a workspace of another dtype."""
57
+ bytes_per_workspace_element = workspace_dtype.itemsize
58
+ fp8_columns_per_workspace_element = (
59
+ bytes_per_workspace_element // torch.float8_e4m3fn.itemsize
60
+ )
61
+ assert bytes_per_workspace_element % torch.float8_e4m3fn.itemsize == 0
62
+ return (
63
+ num_rows,
64
+ -(-num_columns // fp8_columns_per_workspace_element),
65
+ )
66
+
67
+
68
def _valid_deep_gemm_shape(M: int, N: int, K: int) -> bool:
69
align = get_mk_alignment_for_contiguous_layout()[0]
70
return align <= M and N % align == 0 and K % align == 0
@@ -216,7 +231,14 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
231
assert M_sum % align_used == 0
232
233
activation_out_dim = self.adjust_N_for_activation(N, activation)
219
- workspace1 = (M_sum, max(activation_out_dim, K))
234
+ # workspace1 is allocated in the activation dtype by the workspace
235
+ # manager, but is only ever viewed and used as FP8 in apply(). Size it
236
+ # by bytes instead of reserving one BF16/FP16 element per FP8 element.
237
+ workspace1 = _fp8_workspace_shape(
238
+ M_sum,
239
+ max(activation_out_dim, K),
240
+ self.workspace_dtype(self.moe_config.in_dtype),
241
+ )
242
workspace2 = (M_sum, max(N, K))
243
output = (M, K)
244
return (workspace1, workspace2, output)
@@ -480,7 +502,14 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular):
502
assert M_sum % align_used == 0
503
504
activation_out_dim = self.adjust_N_for_activation(N, activation)
483
- workspace1 = (M_sum, max(activation_out_dim, K))
505
+ # workspace1 holds the permuted and requantized FP8 activations. The
506
+ # workspace manager allocates it in the model activation dtype, so
507
+ # account for the dtype sizes rather than overallocating BF16/FP16.
508
+ workspace1 = _fp8_workspace_shape(
509
+ M_sum,
510
+ max(activation_out_dim, K),
511
+ self.workspace_dtype(self.moe_config.in_dtype),
512
+ )
513
workspace2 = (M_sum, max(N, K))
514
output = (M, K)
515
return (workspace1, workspace2, output)