[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>

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)

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
                 )

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)