From 3508cb78d4c09dff536bd4023016ca486cbde09b Mon Sep 17 00:00:00 2001 From: x41lakazam Date: Thu, 11 Jun 2026 14:17:23 +0300 Subject: [PATCH] [Bugfix] Fix broken profile_modular_kernel.py (#43300) --- .../profile_modular_kernel.py | 70 ++++++++++++++++--- .../moe/test_profile_modular_kernel.py | 38 ++++++++++ 2 files changed, 98 insertions(+), 10 deletions(-) create mode 100644 tests/kernels/moe/test_profile_modular_kernel.py diff --git a/tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py b/tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py index 04e9c2aa459..301aa94e02e 100644 --- a/tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py +++ b/tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py @@ -9,9 +9,19 @@ from typing import Any import torch from vllm.config import VllmConfig +from vllm.forward_context import set_forward_context +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig from vllm.utils.torch_utils import set_random_seed +from vllm.v1.worker.workspace import init_workspace_manager -from .common import Config, RankTensors, WeightTensors, make_modular_kernel +from .common import ( + Config, + RankTensors, + WeightTensors, + _make_gscale, + make_modular_kernel, +) from .parallel_utils import ProcessGroupInfo, parallel_launch_with_config @@ -35,7 +45,7 @@ def do_profile( ) as tprof: fn(**fn_kwargs) device = torch.accelerator.current_device_index() - torch.accelerator.synchronize(device=device) + torch.accelerator.synchronize(device) # TODO (varun): Add a descriptive trace file name tprof.export_chrome_trace( @@ -56,24 +66,60 @@ def profile_modular_kernel( # weights for rank rank_weights = weights.slice_weights(pgi.rank, config.num_local_experts) + if config.quant_dtype == "nvfp4": + gscale = _make_gscale(config.num_local_experts) + else: + gscale = None + + quant_config = FusedMoEQuantConfig.make( + config.quant_dtype, + w1_scale=rank_weights.w1_scale, + w2_scale=rank_weights.w2_scale, + a1_scale=rank_tensors.hidden_states_scale, + g1_alphas=(1 / rank_weights.w1_gs) if rank_weights.w1_gs is not None else None, + g2_alphas=(1 / rank_weights.w2_gs) if rank_weights.w2_gs is not None else None, + a1_gscale=gscale, + a2_gscale=gscale, + block_shape=config.quant_block_shape, + per_act_token_quant=config.is_per_act_token_quant, + per_out_ch_quant=config.is_per_out_ch_quant, + ) + # make modular kernel - mk = make_modular_kernel(config, vllm_config, weights) + mk = make_modular_kernel(config, vllm_config, quant_config) + + topk_ids = rank_tensors.topk_ids.to( + mk.prepare_finalize.topk_indices_dtype() or rank_tensors.topk_ids.dtype + ) + + # impls might update the tensor in place + hidden_states = rank_tensors.hidden_states.clone() mk_kwargs = { - "hidden_states": rank_tensors.hidden_states, + "hidden_states": hidden_states, "w1": rank_weights.w1, "w2": rank_weights.w2, "topk_weights": rank_tensors.topk_weights, - "topk_ids": rank_tensors.topk_ids, + "topk_ids": topk_ids, + "activation": MoEActivation.SILU, "expert_map": rank_tensors.expert_map, - "w1_scale": rank_weights.w1_scale, - "w2_scale": rank_weights.w2_scale, - "a1_scale": rank_tensors.hidden_states_scale, "global_num_experts": config.E, - "apply_router_weight_on_input": config.topk == 1, + "apply_router_weight_on_input": config.topk == 1 + and config.supports_apply_weight_on_input(), } - do_profile(mk.apply, mk_kwargs, pgi, config) + num_tokens = hidden_states.shape[0] + num_tokens_across_dp = torch.tensor( + [num_tokens] * config.world_size, device="cpu", dtype=torch.int + ) + + with set_forward_context( + None, + vllm_config, + num_tokens=num_tokens, + num_tokens_across_dp=num_tokens_across_dp, + ): + do_profile(mk.apply, mk_kwargs, pgi, config) def rank_worker( @@ -85,6 +131,10 @@ def rank_worker( ): set_random_seed(pgi.rank) + # workspace manager is normally initialized by GPUModelRunner; we initialize + # it here for the standalone benchmark process. + init_workspace_manager(torch.device(f"cuda:{pgi.local_rank}")) + # get weights to this device weights.to_current_device() diff --git a/tests/kernels/moe/test_profile_modular_kernel.py b/tests/kernels/moe/test_profile_modular_kernel.py new file mode 100644 index 00000000000..de201057f36 --- /dev/null +++ b/tests/kernels/moe/test_profile_modular_kernel.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm.model_executor.layers.fused_moe import TritonExperts +from vllm.model_executor.layers.fused_moe.prepare_finalize import ( + MoEPrepareAndFinalizeNoDPEPModular, +) + +from .modular_kernel_tools.common import Config +from .modular_kernel_tools.profile_modular_kernel import run + + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="profile_modular_kernel requires a CUDA device", +) +def test_profile_modular_kernel_smoke(tmp_path): + config = Config( + Ms=[16], + K=128, + N=256, + E=4, + topks=[2], + dtype=torch.bfloat16, + quant_config=None, + prepare_finalize_type=MoEPrepareAndFinalizeNoDPEPModular, + fused_experts_type=TritonExperts, + world_size=1, + torch_trace_dir_path=str(tmp_path), + ) + + run(config) + + traces = list(tmp_path.glob("m*_*_trace.json")) + assert traces, "profile_modular_kernel.run did not emit any chrome traces"