A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 13, 04:30 UTC
[Performance][DeepGEMM] Estimate expected_m (#28694)
Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com> Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
Varun Sundar Rabindranath committed
Nov 15, 2025 at 00:52 UTC
6965ef436fb398bfbbdce5b6f88dd842c5944771
3 files changed
+73
-17
tests/kernels/moe/test_deepep_deepgemm_moe.py
+34
-12
@@ -7,6 +7,7 @@ fp8 block-quantized case.
7
"""
8
9
import dataclasses
10
+from contextlib import contextmanager
11
12
import pytest
13
import torch.distributed
@@ -14,6 +15,7 @@ from torch.distributed import ProcessGroup
15
from typing_extensions import ParamSpec
16
17
from vllm.config import VllmConfig, set_current_vllm_config
18
+from vllm.forward_context import set_forward_context
19
from vllm.model_executor.layers.fused_moe.config import (
20
FusedMoEQuantConfig,
21
fp8_w8a8_moe_quant_config,
@@ -61,6 +63,23 @@ requires_deep_gemm = pytest.mark.skipif(
63
P = ParamSpec("P")
64
65
66
+@contextmanager
67
+def with_dp_metadata(M: int, world_size: int):
68
+ num_tokens_across_dp = torch.tensor([M] * world_size, device="cpu", dtype=torch.int)
69
+
70
+ vllm_config = VllmConfig()
71
+ vllm_config.parallel_config.data_parallel_size = world_size
72
+ vllm_config.parallel_config.enable_expert_parallel = True
73
+
74
+ with set_forward_context(
75
+ None,
76
+ vllm_config,
77
+ num_tokens=M,
78
+ num_tokens_across_dp=num_tokens_across_dp,
79
+ ):
80
+ yield
81
+
82
+
83
def next_power_of_2(x):
84
import math
85
@@ -285,18 +304,21 @@ def deepep_deepgemm_moe_impl(
304
quant_config=quant_config,
305
)
306
288
- out = mk.forward(
289
- hidden_states=test_tensors.rank_tokens,
290
- w1=w1,
291
- w2=w2,
292
- topk_weights=test_tensors.topk_weights,
293
- topk_ids=test_tensors.topk,
294
- inplace=False,
295
- activation="silu",
296
- global_num_experts=num_experts,
297
- expert_map=build_expert_map(),
298
- apply_router_weight_on_input=False,
299
- )
307
+ with with_dp_metadata(
308
+ M=test_tensors.rank_tokens.size(0), world_size=pgi.world_size
309
+ ):
310
+ out = mk.forward(
311
+ hidden_states=test_tensors.rank_tokens,
312
+ w1=w1,
313
+ w2=w2,
314
+ topk_weights=test_tensors.topk_weights,
315
+ topk_ids=test_tensors.topk,
316
+ inplace=False,
317
+ activation="silu",
318
+ global_num_experts=num_experts,
319
+ expert_map=build_expert_map(),
320
+ apply_router_weight_on_input=False,
321
+ )
322
return out
323
324
vllm/forward_context.py
+4
@@ -221,6 +221,10 @@ def get_forward_context() -> ForwardContext:
221
return _forward_context
222
223
224
+def is_forward_context_available() -> bool:
225
+ return _forward_context is not None
226
+
227
+
228
def create_forward_context(
229
attn_metadata: Any,
230
vllm_config: VllmConfig,
vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py
+35
-5
@@ -5,6 +5,7 @@
5
import torch
6
7
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8
+from vllm.forward_context import get_forward_context, is_forward_context_available
9
from vllm.logger import init_logger
10
from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
11
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
@@ -19,7 +20,7 @@ from vllm.utils.deep_gemm import (
20
get_mk_alignment_for_contiguous_layout,
21
is_deep_gemm_e8m0_used,
22
)
22
-from vllm.utils.math_utils import cdiv
23
+from vllm.utils.math_utils import cdiv, round_up
24
25
logger = init_logger(__name__)
26
@@ -313,6 +314,33 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
314
output = (num_experts, max_num_tokens * num_dispatchers, K)
315
return (workspace13, workspace2, output)
316
317
+ def estimate_expected_m(
318
+ self, global_num_experts: int, max_tokens_per_expert: int, topk: int
319
+ ) -> int:
320
+ dp_meta = (
321
+ get_forward_context().dp_metadata
322
+ if is_forward_context_available()
323
+ else None
324
+ )
325
+ if dp_meta is None:
326
+ logger.warning_once(
327
+ "DPMetadata unavailable. Defaulting expected_m to "
328
+ f"{max_tokens_per_expert}.",
329
+ scope="local",
330
+ )
331
+ return max_tokens_per_expert
332
+
333
+ total_num_tokens = dp_meta.num_tokens_across_dp_cpu.sum().item()
334
+ total_num_tokens_replicated = total_num_tokens * topk
335
+
336
+ # Assume even load balancing
337
+ assert global_num_experts != 0
338
+ estimate = round_up(int(total_num_tokens_replicated // global_num_experts), 16)
339
+ # clamp estimate
340
+ estimate = max(estimate, 16)
341
+ estimate = min(max_tokens_per_expert, estimate)
342
+ return estimate
343
+
344
def apply(
345
self,
346
output: torch.Tensor,
@@ -348,10 +376,12 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
376
377
workspace1 = _resize_cache(workspace13, (E, max_num_tokens, N))
378
351
- # (from deepgemm docs) : A value hint (which is a value on CPU)
352
- # for the M expectation of each batch, correctly setting this value
353
- # may lead to better performance.
354
- expected_m = max_num_tokens
379
+ expected_m = self.estimate_expected_m(
380
+ global_num_experts=global_num_experts,
381
+ max_tokens_per_expert=max_num_tokens,
382
+ topk=topk_ids.size(-1),
383
+ )
384
+
385
fp8_m_grouped_gemm_nt_masked(
386
(a1q, a1q_scale),
387
(w1, self.w1_scale),