@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 182 lines 5.38 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import copy
5 from collections.abc import Callable
6 from itertools import product
7 from typing import Any
8
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
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
28 def do_profile(
29 fn: Callable,
30 fn_kwargs: dict[Any, Any],
31 pgi: ProcessGroupInfo,
32 config: Config,
33 num_warmups: int = 5,
34 ):
35 for _ in range(num_warmups):
36 fn(**fn_kwargs)
37
38 with torch.profiler.profile(
39 activities=[
40 torch.profiler.ProfilerActivity.CPU,
41 torch.profiler.ProfilerActivity.CUDA,
42 ],
43 with_stack=True,
44 record_shapes=True,
45 ) as tprof:
46 fn(**fn_kwargs)
47 device = torch.accelerator.current_device_index()
48 torch.accelerator.synchronize(device)
49
50 # TODO (varun): Add a descriptive trace file name
51 tprof.export_chrome_trace(
52 f"{config.torch_trace_dir_path}/m{config.M}_{pgi.rank}_trace.json"
53 )
54
55
56 def profile_modular_kernel(
57 pgi: ProcessGroupInfo,
58 vllm_config: VllmConfig,
59 config: Config,
60 weights: WeightTensors,
61 rank_tensors: RankTensors,
62 ) -> None:
63 assert isinstance(config.Ms, int)
64 assert isinstance(config.topks, int)
65
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
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 = {
99 "hidden_states": hidden_states,
100 "w1": rank_weights.w1,
101 "w2": rank_weights.w2,
102 "topk_weights": rank_tensors.topk_weights,
103 "topk_ids": topk_ids,
104 "activation": MoEActivation.SILU,
105 "expert_map": rank_tensors.expert_map,
106 "global_num_experts": config.E,
107 "apply_router_weight_on_input": config.topk == 1
108 and config.supports_apply_weight_on_input(),
109 }
110
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(
126 pgi: ProcessGroupInfo,
127 vllm_config: VllmConfig,
128 cpu_group,
129 config: Config,
130 weights: WeightTensors,
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
141 Ms = config.Ms
142 assert isinstance(Ms, list)
143 TOPKs = config.topks
144 assert isinstance(TOPKs, list)
145
146 for m, topk in product(Ms, TOPKs):
147 print(f"Running m={m}, topk={topk} ...")
148 # override m and topk
149 cfgx = copy.deepcopy(config)
150 cfgx.Ms = m
151 cfgx.topks = topk
152
153 # inputs for rank
154 rank_tensors = RankTensors.make(cfgx, pgi)
155 profile_modular_kernel(pgi, vllm_config, cfgx, weights, rank_tensors)
156
157
158 def run(config: Config):
159 weights: WeightTensors = WeightTensors.make(config)
160 vllm_config, env_dict = config.make_env_data()
161 parallel_launch_with_config(
162 config.world_size, rank_worker, vllm_config, env_dict, config, weights
163 )
164
165
166 if __name__ == "__main__":
167 from .cli_args import make_config, make_config_arg_parser
168
169 parser = make_config_arg_parser(
170 description=(
171 "Run single prepare-finalize & fused-experts combination test"
172 "Example : python3 -m tests.kernels.moe.modular_kernel_tools.profile_modular_kernel " # noqa: E501
173 "--pf-type DeepEPLLPrepareAndFinalize --experts-type BatchedTritonExperts"
174 )
175 )
176 args = parser.parse_args()
177 assert args.torch_trace_dir_path is not None, (
178 "Please pass in a directory to store torch traces"
179 )
180 config = make_config(args)
181
182 run(config)