A high-throughput and memory-efficient inference and serving engine for LLMs
[Bugfix] Fix broken profile_modular_kernel.py (#43300)
x41lakazam committed
Jun 11, 2026 at 14:17 UTC
3508cb78d4c09dff536bd4023016ca486cbde09b
2 files changed
+99
-11
tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py
+61
-11
@@ -9,9 +9,19 @@ from typing import Any
9
import torch
10
11
from vllm.config import VllmConfig
12
+from vllm.forward_context import set_forward_context
13
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
14
+from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
15
from vllm.utils.torch_utils import set_random_seed
13
-
14
-from .common import Config, RankTensors, WeightTensors, make_modular_kernel
16
+from vllm.v1.worker.workspace import init_workspace_manager
17
+
18
+from .common import (
19
+ Config,
20
+ RankTensors,
21
+ WeightTensors,
22
+ _make_gscale,
23
+ make_modular_kernel,
24
+)
25
from .parallel_utils import ProcessGroupInfo, parallel_launch_with_config
26
27
@@ -35,7 +45,7 @@ def do_profile(
45
) as tprof:
46
fn(**fn_kwargs)
47
device = torch.accelerator.current_device_index()
38
- torch.accelerator.synchronize(device=device)
48
+ torch.accelerator.synchronize(device)
49
50
# TODO (varun): Add a descriptive trace file name
51
tprof.export_chrome_trace(
@@ -56,24 +66,60 @@ def profile_modular_kernel(
66
# weights for rank
67
rank_weights = weights.slice_weights(pgi.rank, config.num_local_experts)
68
69
+ if config.quant_dtype == "nvfp4":
70
+ gscale = _make_gscale(config.num_local_experts)
71
+ else:
72
+ gscale = None
73
+
74
+ quant_config = FusedMoEQuantConfig.make(
75
+ config.quant_dtype,
76
+ w1_scale=rank_weights.w1_scale,
77
+ w2_scale=rank_weights.w2_scale,
78
+ a1_scale=rank_tensors.hidden_states_scale,
79
+ g1_alphas=(1 / rank_weights.w1_gs) if rank_weights.w1_gs is not None else None,
80
+ g2_alphas=(1 / rank_weights.w2_gs) if rank_weights.w2_gs is not None else None,
81
+ a1_gscale=gscale,
82
+ a2_gscale=gscale,
83
+ block_shape=config.quant_block_shape,
84
+ per_act_token_quant=config.is_per_act_token_quant,
85
+ per_out_ch_quant=config.is_per_out_ch_quant,
86
+ )
87
+
88
# make modular kernel
60
- mk = make_modular_kernel(config, vllm_config, weights)
89
+ mk = make_modular_kernel(config, vllm_config, quant_config)
90
+
91
+ topk_ids = rank_tensors.topk_ids.to(
92
+ mk.prepare_finalize.topk_indices_dtype() or rank_tensors.topk_ids.dtype
93
+ )
94
+
95
+ # impls might update the tensor in place
96
+ hidden_states = rank_tensors.hidden_states.clone()
97
98
mk_kwargs = {
63
- "hidden_states": rank_tensors.hidden_states,
99
+ "hidden_states": hidden_states,
100
"w1": rank_weights.w1,
101
"w2": rank_weights.w2,
102
"topk_weights": rank_tensors.topk_weights,
67
- "topk_ids": rank_tensors.topk_ids,
103
+ "topk_ids": topk_ids,
104
+ "activation": MoEActivation.SILU,
105
"expert_map": rank_tensors.expert_map,
69
- "w1_scale": rank_weights.w1_scale,
70
- "w2_scale": rank_weights.w2_scale,
71
- "a1_scale": rank_tensors.hidden_states_scale,
106
"global_num_experts": config.E,
73
- "apply_router_weight_on_input": config.topk == 1,
107
+ "apply_router_weight_on_input": config.topk == 1
108
+ and config.supports_apply_weight_on_input(),
109
}
110
76
- do_profile(mk.apply, mk_kwargs, pgi, config)
111
+ num_tokens = hidden_states.shape[0]
112
+ num_tokens_across_dp = torch.tensor(
113
+ [num_tokens] * config.world_size, device="cpu", dtype=torch.int
114
+ )
115
+
116
+ with set_forward_context(
117
+ None,
118
+ vllm_config,
119
+ num_tokens=num_tokens,
120
+ num_tokens_across_dp=num_tokens_across_dp,
121
+ ):
122
+ do_profile(mk.apply, mk_kwargs, pgi, config)
123
124
125
def rank_worker(
@@ -85,6 +131,10 @@ def rank_worker(
131
):
132
set_random_seed(pgi.rank)
133
134
+ # workspace manager is normally initialized by GPUModelRunner; we initialize
135
+ # it here for the standalone benchmark process.
136
+ init_workspace_manager(torch.device(f"cuda:{pgi.local_rank}"))
137
+
138
# get weights to this device
139
weights.to_current_device()
140
tests/kernels/moe/test_profile_modular_kernel.py
new
+38
@@ -0,0 +1,38 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+from vllm.model_executor.layers.fused_moe import TritonExperts
8
+from vllm.model_executor.layers.fused_moe.prepare_finalize import (
9
+ MoEPrepareAndFinalizeNoDPEPModular,
10
+)
11
+
12
+from .modular_kernel_tools.common import Config
13
+from .modular_kernel_tools.profile_modular_kernel import run
14
+
15
+
16
+@pytest.mark.skipif(
17
+ not torch.cuda.is_available(),
18
+ reason="profile_modular_kernel requires a CUDA device",
19
+)
20
+def test_profile_modular_kernel_smoke(tmp_path):
21
+ config = Config(
22
+ Ms=[16],
23
+ K=128,
24
+ N=256,
25
+ E=4,
26
+ topks=[2],
27
+ dtype=torch.bfloat16,
28
+ quant_config=None,
29
+ prepare_finalize_type=MoEPrepareAndFinalizeNoDPEPModular,
30
+ fused_experts_type=TritonExperts,
31
+ world_size=1,
32
+ torch_trace_dir_path=str(tmp_path),
33
+ )
34
+
35
+ run(config)
36
+
37
+ traces = list(tmp_path.glob("m*_*_trace.json"))
38
+ assert traces, "profile_modular_kernel.run did not emit any chrome traces"