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