@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 9352cdb56d
Fork 0

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 15, 12:15 UTC

[Hardware][AMD] Improve OAM device ID + llama4 Maverick MOE tuning (#16263)

Signed-off-by: Lu Fang <lufang@fb.com> Co-authored-by: Lu Fang <lufang@fb.com>

Xiaodong Wang committed May 2, 2025 at 12:44 UTC 9352cdb56d70bd52d4e6ea88d991bf5f4cc93393
3 files changed +231 -3
benchmarks/kernels/benchmark_moe.py
+17 -2
@@ -442,8 +442,14 @@ class BenchmarkWorker:
442 hidden_size, search_space,
443 is_fp16, topk)
444
445 - with torch.cuda.device(self.device_id) if current_platform.is_rocm(
446 - ) else nullcontext():
445 + need_device_guard = False
446 + if current_platform.is_rocm():
447 + visible_device = os.environ.get("ROCR_VISIBLE_DEVICES", None)
448 + if visible_device != f"{self.device_id}":
449 + need_device_guard = True
450 +
451 + with torch.cuda.device(
452 + self.device_id) if need_device_guard else nullcontext():
453 for config in tqdm(search_space):
454 try:
455 kernel_time = benchmark_config(
@@ -578,6 +584,15 @@ def main(args: argparse.Namespace):
584
585 use_deep_gemm = bool(args.use_deep_gemm)
586
587 + if current_platform.is_rocm() and "HIP_VISIBLE_DEVICES" in os.environ:
588 + # Ray will set ROCR_VISIBLE_DEVICES for device visibility
589 + logger.warning(
590 + "Ray uses ROCR_VISIBLE_DEVICES to control device accessibility."
591 + "Replacing HIP_VISIBLE_DEVICES with ROCR_VISIBLE_DEVICES.")
592 + val = os.environ["HIP_VISIBLE_DEVICES"]
593 + os.environ["ROCR_VISIBLE_DEVICES"] = val
594 + del os.environ["HIP_VISIBLE_DEVICES"]
595 +
596 ray.init()
597 num_gpus = int(ray.available_resources()["GPU"])
598 workers = [BenchmarkWorker.remote(args.seed) for _ in range(num_gpus)]
vllm/model_executor/layers/fused_moe/configs/E=128,N=1024,device_name=AMD_Instinct_MI300X.json new
+200
@@ -0,0 +1,200 @@
1 +{
2 + "1": {
3 + "BLOCK_SIZE_M": 16,
4 + "BLOCK_SIZE_N": 16,
5 + "BLOCK_SIZE_K": 256,
6 + "GROUP_SIZE_M": 1,
7 + "num_warps": 4,
8 + "num_stages": 2,
9 + "waves_per_eu": 0,
10 + "matrix_instr_nonkdim": 16,
11 + "kpack": 2
12 + },
13 + "2": {
14 + "BLOCK_SIZE_M": 16,
15 + "BLOCK_SIZE_N": 16,
16 + "BLOCK_SIZE_K": 256,
17 + "GROUP_SIZE_M": 1,
18 + "num_warps": 4,
19 + "num_stages": 2,
20 + "waves_per_eu": 0,
21 + "matrix_instr_nonkdim": 16,
22 + "kpack": 1
23 + },
24 + "4": {
25 + "BLOCK_SIZE_M": 16,
26 + "BLOCK_SIZE_N": 16,
27 + "BLOCK_SIZE_K": 256,
28 + "GROUP_SIZE_M": 1,
29 + "num_warps": 2,
30 + "num_stages": 2,
31 + "waves_per_eu": 0,
32 + "matrix_instr_nonkdim": 16,
33 + "kpack": 2
34 + },
35 + "8": {
36 + "BLOCK_SIZE_M": 16,
37 + "BLOCK_SIZE_N": 64,
38 + "BLOCK_SIZE_K": 128,
39 + "GROUP_SIZE_M": 1,
40 + "num_warps": 4,
41 + "num_stages": 2,
42 + "waves_per_eu": 0,
43 + "matrix_instr_nonkdim": 16,
44 + "kpack": 2
45 + },
46 + "16": {
47 + "BLOCK_SIZE_M": 16,
48 + "BLOCK_SIZE_N": 128,
49 + "BLOCK_SIZE_K": 128,
50 + "GROUP_SIZE_M": 1,
51 + "num_warps": 4,
52 + "num_stages": 2,
53 + "waves_per_eu": 0,
54 + "matrix_instr_nonkdim": 16,
55 + "kpack": 2
56 + },
57 + "24": {
58 + "BLOCK_SIZE_M": 16,
59 + "BLOCK_SIZE_N": 64,
60 + "BLOCK_SIZE_K": 128,
61 + "GROUP_SIZE_M": 1,
62 + "num_warps": 4,
63 + "num_stages": 2,
64 + "waves_per_eu": 0,
65 + "matrix_instr_nonkdim": 16,
66 + "kpack": 1
67 + },
68 + "32": {
69 + "BLOCK_SIZE_M": 16,
70 + "BLOCK_SIZE_N": 128,
71 + "BLOCK_SIZE_K": 128,
72 + "GROUP_SIZE_M": 1,
73 + "num_warps": 8,
74 + "num_stages": 2,
75 + "waves_per_eu": 0,
76 + "matrix_instr_nonkdim": 16,
77 + "kpack": 2
78 + },
79 + "48": {
80 + "BLOCK_SIZE_M": 16,
81 + "BLOCK_SIZE_N": 64,
82 + "BLOCK_SIZE_K": 128,
83 + "GROUP_SIZE_M": 1,
84 + "num_warps": 2,
85 + "num_stages": 2,
86 + "waves_per_eu": 0,
87 + "matrix_instr_nonkdim": 16,
88 + "kpack": 2
89 + },
90 + "64": {
91 + "BLOCK_SIZE_M": 16,
92 + "BLOCK_SIZE_N": 64,
93 + "BLOCK_SIZE_K": 128,
94 + "GROUP_SIZE_M": 4,
95 + "num_warps": 2,
96 + "num_stages": 2,
97 + "waves_per_eu": 0,
98 + "matrix_instr_nonkdim": 16,
99 + "kpack": 2
100 + },
101 + "96": {
102 + "BLOCK_SIZE_M": 16,
103 + "BLOCK_SIZE_N": 32,
104 + "BLOCK_SIZE_K": 128,
105 + "GROUP_SIZE_M": 1,
106 + "num_warps": 1,
107 + "num_stages": 2,
108 + "waves_per_eu": 0,
109 + "matrix_instr_nonkdim": 16,
110 + "kpack": 1
111 + },
112 + "128": {
113 + "BLOCK_SIZE_M": 16,
114 + "BLOCK_SIZE_N": 64,
115 + "BLOCK_SIZE_K": 128,
116 + "GROUP_SIZE_M": 1,
117 + "num_warps": 4,
118 + "num_stages": 2,
119 + "waves_per_eu": 0,
120 + "matrix_instr_nonkdim": 16,
121 + "kpack": 1
122 + },
123 + "256": {
124 + "BLOCK_SIZE_M": 16,
125 + "BLOCK_SIZE_N": 32,
126 + "BLOCK_SIZE_K": 128,
127 + "GROUP_SIZE_M": 1,
128 + "num_warps": 2,
129 + "num_stages": 2,
130 + "waves_per_eu": 0,
131 + "matrix_instr_nonkdim": 16,
132 + "kpack": 2
133 + },
134 + "512": {
135 + "BLOCK_SIZE_M": 16,
136 + "BLOCK_SIZE_N": 128,
137 + "BLOCK_SIZE_K": 128,
138 + "GROUP_SIZE_M": 8,
139 + "num_warps": 8,
140 + "num_stages": 2,
141 + "waves_per_eu": 0,
142 + "matrix_instr_nonkdim": 16,
143 + "kpack": 2
144 + },
145 + "1024": {
146 + "BLOCK_SIZE_M": 16,
147 + "BLOCK_SIZE_N": 128,
148 + "BLOCK_SIZE_K": 128,
149 + "GROUP_SIZE_M": 1,
150 + "num_warps": 8,
151 + "num_stages": 2,
152 + "waves_per_eu": 0,
153 + "matrix_instr_nonkdim": 16,
154 + "kpack": 2
155 + },
156 + "1536": {
157 + "BLOCK_SIZE_M": 32,
158 + "BLOCK_SIZE_N": 128,
159 + "BLOCK_SIZE_K": 128,
160 + "GROUP_SIZE_M": 8,
161 + "num_warps": 8,
162 + "num_stages": 2,
163 + "waves_per_eu": 0,
164 + "matrix_instr_nonkdim": 16,
165 + "kpack": 2
166 + },
167 + "2048": {
168 + "BLOCK_SIZE_M": 64,
169 + "BLOCK_SIZE_N": 64,
170 + "BLOCK_SIZE_K": 128,
171 + "GROUP_SIZE_M": 16,
172 + "num_warps": 8,
173 + "num_stages": 2,
174 + "waves_per_eu": 0,
175 + "matrix_instr_nonkdim": 16,
176 + "kpack": 2
177 + },
178 + "3072": {
179 + "BLOCK_SIZE_M": 64,
180 + "BLOCK_SIZE_N": 128,
181 + "BLOCK_SIZE_K": 128,
182 + "GROUP_SIZE_M": 8,
183 + "num_warps": 8,
184 + "num_stages": 2,
185 + "waves_per_eu": 0,
186 + "matrix_instr_nonkdim": 16,
187 + "kpack": 1
188 + },
189 + "4096": {
190 + "BLOCK_SIZE_M": 64,
191 + "BLOCK_SIZE_N": 128,
192 + "BLOCK_SIZE_K": 128,
193 + "GROUP_SIZE_M": 32,
194 + "num_warps": 8,
195 + "num_stages": 2,
196 + "waves_per_eu": 0,
197 + "matrix_instr_nonkdim": 16,
198 + "kpack": 2
199 + }
200 +}
vllm/platforms/rocm.py
+14 -1
@@ -58,6 +58,15 @@ _ROCM_PARTIALLY_SUPPORTED_MODELS: Dict[str, str] = {
58 "excessive use of shared memory. If this happens, disable Triton FA "
59 "by setting `VLLM_USE_TRITON_FLASH_ATTN=0`")
60 }
61 +_ROCM_DEVICE_ID_NAME_MAP: Dict[str, str] = {
62 + "0x74a0": "AMD_Instinct_MI300A",
63 + "0x74a1": "AMD_Instinct_MI300X",
64 + "0x74b5": "AMD_Instinct_MI300X", # MI300X VF
65 + "0x74a5": "AMD_Instinct_MI325X",
66 + "0x74b9": "AMD_Instinct_MI325X", # MI325X VF
67 + "0x74a9": "AMD_Instinct_MI300X_HF",
68 + "0x74bd": "AMD_Instinct_MI300X_HF",
69 +}
70
71 # Prevent use of clashing `{CUDA/HIP}_VISIBLE_DEVICES``
72 if "HIP_VISIBLE_DEVICES" in os.environ:
@@ -225,7 +234,11 @@ class RocmPlatform(Platform):
234 def get_device_name(cls, device_id: int = 0) -> str:
235 physical_device_id = device_id_to_physical_device_id(device_id)
236 handle = amdsmi_get_processor_handles()[physical_device_id]
228 - return amdsmi_get_gpu_asic_info(handle)["market_name"]
237 + asic_info = amdsmi_get_gpu_asic_info(handle)
238 + device_name: str = asic_info["device_id"]
239 + if device_name in _ROCM_DEVICE_ID_NAME_MAP:
240 + return _ROCM_DEVICE_ID_NAME_MAP[device_name]
241 + return asic_info["market_name"]
242
243 @classmethod
244 def get_device_total_memory(cls, device_id: int = 0) -> int: