@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 725c3bc808
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

[ROCm][Perf] Enable W4A16 FlyDSL MoE (#44400)

Signed-off-by: amd-asalykov <asalykov@amd.com> Signed-off-by: Amanzhol Salykov <asalykov@amd.com>

Amanzhol Salykov committed Jun 14, 2026 at 09:14 UTC 725c3bc808c6eb5a572bdb37ed8a84bd11aad24a
15 files changed +2173
benchmarks/kernels/benchmark_flydsl_moe_w4a16.py new
+277
@@ -0,0 +1,277 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +# Copyright (c) 2025 FlyDSL Project Contributors
4 +
5 +import json
6 +import os
7 +
8 +import torch
9 +from aiter.test_common import run_perftest
10 +
11 +from vllm.model_executor.layers.fused_moe import fused_experts
12 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation
13 +from vllm.model_executor.layers.fused_moe.config import (
14 + int4_w4a16_moe_quant_config,
15 +)
16 +from vllm.model_executor.layers.fused_moe.fused_flydsl_moe import fused_flydsl_moe
17 +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501
18 + compressed_tensors_moe_w4a16_flydsl,
19 +)
20 +from vllm.platforms import current_platform
21 +
22 +RoutingBuffers = tuple[
23 + torch.Tensor, # sorted_token_ids
24 + torch.Tensor, # sorted_weights
25 + torch.Tensor, # sorted_expert_ids
26 + torch.Tensor, # num_valid_ids (shape [1], i32)
27 + int, # sorted_size
28 + int, # blocks
29 +]
30 +
31 +MODEL_PARAMS_TO_TUNE = [
32 + # (num_experts, inter_dim, hidden_size, topk)
33 + (384, 256, 7168, 8), # Kimi K2.5 TP=8
34 + (384, 512, 7168, 8), # Kimi K2.5 TP=4
35 +]
36 +
37 +NUM_TOKENS_TO_TUNE = [
38 + 1,
39 + 2,
40 + 4,
41 + 8,
42 + 16,
43 + 24,
44 + 32,
45 + 48,
46 + 64,
47 + 128,
48 + 256,
49 + 512,
50 + 1024,
51 + 2048,
52 + 4096,
53 + 8192,
54 +]
55 +
56 +TILE_M_SEARCH_SPACE = [16, 32, 64, 128, 256]
57 +TILE_N_SEARCH_SPACE = [16, 32, 64, 128, 256]
58 +TILE_K_SEARCH_SPACE = [16, 32, 64, 128, 256, 512]
59 +TILE_N2_SEARCH_SPACE = [16, 32, 64, 128, 256]
60 +TILE_K2_SEARCH_SPACE = [16, 32, 64, 128, 256, 512]
61 +
62 +TILE_CONFIGS = []
63 +for tile_m in TILE_M_SEARCH_SPACE:
64 + for tile_n in TILE_N_SEARCH_SPACE:
65 + for tile_k in TILE_K_SEARCH_SPACE:
66 + for tile_n2 in TILE_N2_SEARCH_SPACE:
67 + for tile_k2 in TILE_K2_SEARCH_SPACE:
68 + TILE_CONFIGS.append(
69 + {
70 + "tile_m": tile_m,
71 + "tile_n": tile_n,
72 + "tile_k": tile_k,
73 + "tile_n2": tile_n2,
74 + "tile_k2": tile_k2,
75 + }
76 + )
77 +
78 +
79 +def tune_flydsl_moe_w4a16(
80 + device: str = "cuda", num_iters: int = 100, num_warmup: int = 10
81 +):
82 + packed_factor = 8
83 + w13_num_shards = 2
84 + params_dtype = torch.bfloat16
85 + group_size = 32
86 + scale_factor = 0.01
87 +
88 + for model_params in MODEL_PARAMS_TO_TUNE:
89 + num_experts = model_params[0]
90 + inter_dim = model_params[1]
91 + hidden_size = model_params[2]
92 + topk = model_params[3]
93 + print(
94 + f"\nTuning: num_experts={num_experts}, inter_dim={inter_dim}, "
95 + f"hidden_size={hidden_size}, topk={topk}...\n"
96 + )
97 +
98 + w2_scales_size = inter_dim
99 + num_groups_w2 = w2_scales_size // group_size
100 + num_groups_w13 = hidden_size // group_size
101 +
102 + w13_weight = torch.randint(
103 + 0,
104 + 255,
105 + (num_experts, hidden_size // packed_factor, w13_num_shards * inter_dim),
106 + dtype=torch.int32,
107 + device=device,
108 + )
109 +
110 + w2_weight = torch.randint(
111 + 0,
112 + 255,
113 + (num_experts, inter_dim // packed_factor, hidden_size),
114 + dtype=torch.int32,
115 + device=device,
116 + )
117 + w13_scale = scale_factor * torch.randn(
118 + num_experts,
119 + num_groups_w13,
120 + w13_num_shards * inter_dim,
121 + dtype=params_dtype,
122 + device=device,
123 + )
124 + w2_scale = scale_factor * torch.randn(
125 + num_experts, num_groups_w2, hidden_size, dtype=params_dtype, device=device
126 + )
127 +
128 + w13 = w13_weight
129 + w13 = compressed_tensors_moe_w4a16_flydsl._gptq_int32_to_flydsl_packed(w13)
130 + w13 = w13.view(-1).contiguous()
131 +
132 + w2 = w2_weight
133 + w2 = compressed_tensors_moe_w4a16_flydsl._gptq_int32_to_flydsl_packed(w2)
134 + w2 = w2.view(-1).contiguous()
135 +
136 + w13_scale_flydsl = w13_scale
137 + w2_scale_flydsl = w2_scale
138 +
139 + if group_size > 0 and w13_scale.dim() == 3 and w13_scale.shape[1] > 1:
140 + E, G, N = w13_scale.shape
141 + w13_scale_flydsl = (
142 + w13_scale_flydsl.view(E, G // 2, 2, N)
143 + .permute(0, 1, 3, 2)
144 + .contiguous()
145 + .view(-1)
146 + .contiguous()
147 + )
148 + elif w13_scale.dim() == 3 and w13_scale.shape[1] == 1:
149 + w13_scale_flydsl = w13_scale_flydsl.squeeze(1)
150 +
151 + if group_size > 0 and w2_scale.dim() == 3 and w2_scale.shape[1] > 1:
152 + E, G, N = w2_scale.shape
153 + w2_scale_flydsl = (
154 + w2_scale_flydsl.view(E, G // 2, 2, N)
155 + .permute(0, 1, 3, 2)
156 + .contiguous()
157 + .view(-1)
158 + .contiguous()
159 + )
160 + elif w2_scale.dim() == 3 and w2_scale.shape[1] == 1:
161 + w2_scale_flydsl = w2_scale_flydsl.squeeze(1)
162 +
163 + w13_scale_flydsl = w13_scale_flydsl.contiguous()
164 + w2_scale_flydsl = w2_scale_flydsl.contiguous()
165 +
166 + w13.is_shuffled = True
167 + w2.is_shuffled = True
168 +
169 + w13_weight_scale = w13_scale.transpose(1, 2).contiguous()
170 + w2_weight_scale = w2_scale.transpose(1, 2).contiguous()
171 + w13_weight_packed = w13_weight.transpose(1, 2).contiguous().view(torch.uint8)
172 + w2_weight_packed = w2_weight.transpose(1, 2).contiguous().view(torch.uint8)
173 +
174 + moe_quant_config = int4_w4a16_moe_quant_config(
175 + w1_scale=w13_weight_scale,
176 + w2_scale=w2_weight_scale,
177 + w1_zp=None,
178 + w2_zp=None,
179 + block_shape=[0, group_size],
180 + )
181 +
182 + tuned_config = {}
183 +
184 + for num_tokens in NUM_TOKENS_TO_TUNE:
185 + score = torch.rand(
186 + (num_tokens, num_experts), device=device, dtype=torch.float32
187 + )
188 + topk_vals, topk_ids = torch.topk(score, k=topk, dim=1)
189 + topk_weights = torch.softmax(topk_vals, dim=1).to(torch.float32)
190 + x = torch.randn(
191 + (num_tokens, hidden_size), dtype=torch.bfloat16, device=device
192 + )
193 + us_best = float("inf")
194 + for tile_config in TILE_CONFIGS:
195 + try:
196 + tile_m = tile_config["tile_m"]
197 + tile_n = tile_config["tile_n"]
198 + tile_k = tile_config["tile_k"]
199 + tile_n2 = tile_config["tile_n2"]
200 + tile_k2 = tile_config["tile_k2"]
201 +
202 + model_dim = x.shape[1]
203 + assert model_dim % 64 == 0
204 + assert model_dim % tile_k == 0
205 + assert inter_dim % tile_n == 0
206 + assert model_dim % tile_n2 == 0
207 + assert inter_dim % tile_k2 == 0
208 + assert ((tile_m * tile_k2) % 256) == 0
209 + bytes_per_thread_x = (tile_m * tile_k2) // 256
210 + assert (bytes_per_thread_x % 4) == 0
211 +
212 + out, _us = run_perftest(
213 + fused_flydsl_moe,
214 + x,
215 + w13,
216 + w2,
217 + num_experts,
218 + inter_dim,
219 + topk_weights,
220 + topk_ids,
221 + num_iters=num_iters,
222 + num_warmup=num_warmup,
223 + w1_scale=w13_scale_flydsl,
224 + w2_scale=w2_scale_flydsl,
225 + topk=topk_weights.shape[-1],
226 + group_size=group_size,
227 + doweight_stage1=False,
228 + scale_is_bf16=True,
229 + config=tile_config,
230 + )
231 + torch.accelerator.synchronize()
232 + except Exception:
233 + torch.accelerator.synchronize()
234 + continue
235 + else:
236 + us = _us.item()
237 + if us < us_best:
238 + out_ref = fused_experts(
239 + x,
240 + w13_weight_packed,
241 + w2_weight_packed,
242 + topk_weights=topk_weights,
243 + topk_ids=topk_ids,
244 + activation=MoEActivation.SILU,
245 + apply_router_weight_on_input=False,
246 + global_num_experts=num_experts,
247 + expert_map=None,
248 + quant_config=moe_quant_config,
249 + )
250 + try:
251 + assert torch.allclose(out, out_ref, atol=0.5, rtol=0.1)
252 + except Exception:
253 + continue
254 + else:
255 + print(
256 + f"For [num_tokens={num_tokens}, num_experts={num_experts}, " # noqa: E501
257 + f"inter_dim={inter_dim}] found new best " # noqa: E501
258 + f"config={tile_config}, us={us:0.3f}"
259 + )
260 + us_best = us
261 + tuned_config[str(num_tokens)] = tile_config
262 + device_name = current_platform.get_device_name().replace(" ", "_")
263 + tuned_config_file_name = (
264 + f"E={num_experts},N={inter_dim},device_name={device_name},"
265 + f"dtype=int4_w4a16,backend=flydsl.json"
266 + )
267 + tuner_dir_path = os.path.dirname(os.path.realpath(__file__))
268 + store_path = os.path.join(tuner_dir_path, tuned_config_file_name)
269 + with open(store_path, "w") as f:
270 + json.dump(tuned_config, f, indent=4)
271 + print(
272 + f"\nTuned config for num_tokens={num_tokens} was stored at {store_path}\n" # noqa: E501
273 + )
274 +
275 +
276 +if __name__ == "__main__":
277 + tune_flydsl_moe_w4a16(device="cuda")
tests/kernels/moe/test_flydsl_moe.py new
+179
@@ -0,0 +1,179 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +# Copyright (c) 2025 FlyDSL Project Contributors
4 +
5 +
6 +import importlib.util
7 +
8 +import pytest
9 +import torch
10 +
11 +from vllm.model_executor.layers.fused_moe import fused_experts
12 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation
13 +from vllm.model_executor.layers.fused_moe.config import (
14 + int4_w4a16_moe_quant_config,
15 +)
16 +from vllm.platforms import current_platform
17 +from vllm.platforms.rocm import on_gfx950
18 +
19 +if not (current_platform.is_rocm() and on_gfx950()):
20 + pytest.skip("This test can only run on ROCm and gfx950.", allow_module_level=True)
21 +
22 +aiter_available = importlib.util.find_spec("aiter") is not None
23 +
24 +if not aiter_available:
25 + pytest.skip("These tests require AITER to run.", allow_module_level=True)
26 +
27 +from vllm.model_executor.layers.fused_moe.fused_flydsl_moe import ( # noqa: E402
28 + fused_flydsl_moe,
29 +)
30 +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E402, E501
31 + compressed_tensors_moe_w4a16_flydsl,
32 +)
33 +
34 +RoutingBuffers = tuple[
35 + torch.Tensor, # sorted_token_ids
36 + torch.Tensor, # sorted_weights
37 + torch.Tensor, # sorted_expert_ids
38 + torch.Tensor, # num_valid_ids (shape [1], i32)
39 + int, # sorted_size
40 + int, # blocks
41 +]
42 +
43 +
44 +@pytest.mark.parametrize(
45 + "num_tokens", [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384]
46 +)
47 +@pytest.mark.parametrize("inter_dim", [256, 512])
48 +def test_flydsl_moe(num_tokens: int, inter_dim: int):
49 + device = "cuda"
50 + topk = 8
51 + num_experts = 384
52 + hidden_size = 7168
53 + packed_factor = 8
54 + w13_num_shards = 2
55 + params_dtype = torch.bfloat16
56 + group_size = 32
57 + w2_scales_size = inter_dim
58 + scale_factor = 0.01
59 +
60 + num_groups_w2 = w2_scales_size // group_size
61 + num_groups_w13 = hidden_size // group_size
62 +
63 + w13_weight = torch.randint(
64 + 0,
65 + 255,
66 + (num_experts, hidden_size // packed_factor, w13_num_shards * inter_dim),
67 + dtype=torch.int32,
68 + device=device,
69 + )
70 +
71 + w2_weight = torch.randint(
72 + 0,
73 + 255,
74 + (num_experts, inter_dim // packed_factor, hidden_size),
75 + dtype=torch.int32,
76 + device=device,
77 + )
78 + w13_scale = scale_factor * torch.randn(
79 + num_experts,
80 + num_groups_w13,
81 + w13_num_shards * inter_dim,
82 + dtype=params_dtype,
83 + device=device,
84 + )
85 + w2_scale = scale_factor * torch.randn(
86 + num_experts, num_groups_w2, hidden_size, dtype=params_dtype, device=device
87 + )
88 +
89 + w13_weight_packed = w13_weight.transpose(1, 2).contiguous().view(torch.uint8)
90 + w2_weight_packed = w2_weight.transpose(1, 2).contiguous().view(torch.uint8)
91 + w13_weight_scale = w13_scale.transpose(1, 2).contiguous()
92 + w2_weight_scale = w2_scale.transpose(1, 2).contiguous()
93 +
94 + moe_quant_config = int4_w4a16_moe_quant_config(
95 + w1_scale=w13_weight_scale,
96 + w2_scale=w2_weight_scale,
97 + w1_zp=None,
98 + w2_zp=None,
99 + block_shape=[0, group_size],
100 + )
101 + score = torch.rand((num_tokens, num_experts), device=device, dtype=torch.float32)
102 + topk_vals, topk_ids = torch.topk(score, k=topk, dim=1)
103 + topk_weights = torch.softmax(topk_vals, dim=1).to(torch.float32)
104 + x = torch.randn((num_tokens, hidden_size), dtype=torch.bfloat16, device=device)
105 + out_ref = fused_experts(
106 + x,
107 + w13_weight_packed,
108 + w2_weight_packed,
109 + topk_weights=topk_weights,
110 + topk_ids=topk_ids,
111 + activation=MoEActivation.SILU,
112 + apply_router_weight_on_input=False,
113 + global_num_experts=num_experts,
114 + expert_map=None,
115 + quant_config=moe_quant_config,
116 + )
117 +
118 + w13 = w13_weight
119 + w13 = compressed_tensors_moe_w4a16_flydsl._gptq_int32_to_flydsl_packed(w13)
120 + w13 = w13.view(-1).contiguous()
121 +
122 + w2 = w2_weight
123 + w2 = compressed_tensors_moe_w4a16_flydsl._gptq_int32_to_flydsl_packed(w2)
124 + w2 = w2.view(-1).contiguous()
125 +
126 + w13_scale_flydsl = w13_scale
127 + w2_scale_flydsl = w2_scale
128 +
129 + if group_size > 0 and w13_scale.dim() == 3 and w13_scale.shape[1] > 1:
130 + E, G, N = w13_scale.shape
131 + w13_scale_flydsl = (
132 + w13_scale_flydsl.view(E, G // 2, 2, N)
133 + .permute(0, 1, 3, 2)
134 + .contiguous()
135 + .view(-1)
136 + .contiguous()
137 + )
138 + elif w13_scale.dim() == 3 and w13_scale.shape[1] == 1:
139 + w13_scale_flydsl = w13_scale_flydsl.squeeze(1)
140 +
141 + if group_size > 0 and w2_scale.dim() == 3 and w2_scale.shape[1] > 1:
142 + E, G, N = w2_scale.shape
143 + w2_scale_flydsl = (
144 + w2_scale_flydsl.view(E, G // 2, 2, N)
145 + .permute(0, 1, 3, 2)
146 + .contiguous()
147 + .view(-1)
148 + .contiguous()
149 + )
150 + elif w2_scale.dim() == 3 and w2_scale.shape[1] == 1:
151 + w2_scale_flydsl = w2_scale_flydsl.squeeze(1)
152 +
153 + w13_scale_flydsl = w13_scale_flydsl.contiguous()
154 + w2_scale_flydsl = w2_scale_flydsl.contiguous()
155 +
156 + w13.is_shuffled = True
157 + w2.is_shuffled = True
158 +
159 + out = fused_flydsl_moe(
160 + x,
161 + w13,
162 + w2,
163 + num_experts,
164 + inter_dim,
165 + topk_weights,
166 + topk_ids,
167 + w1_scale=w13_scale_flydsl,
168 + w2_scale=w2_scale_flydsl,
169 + topk=topk_weights.shape[-1],
170 + group_size=group_size,
171 + doweight_stage1=False,
172 + scale_is_bf16=True,
173 + )
174 +
175 + assert torch.allclose(out, out_ref, atol=0.5, rtol=0.1)
176 +
177 +
178 +if __name__ == "__main__":
179 + test_flydsl_moe(512, 256)
vllm/config/kernel.py
+2
@@ -133,6 +133,7 @@ MoEBackend = Literal[
133 "humming",
134 "triton_unfused",
135 "aiter",
136 + "flydsl",
137 "emulation",
138 ]
139
@@ -186,6 +187,7 @@ class KernelConfig:
187 - "humming": Use Humming Mixed Precision kernels
188 - "triton_unfused": Use Triton unfused MoE kernels
189 - "aiter": Use AMD AITer kernels (ROCm only)
190 + - "flydsl": Use AMD FlyDSL kernels (ROCm only)
191 - "emulation": use BF16/FP16 GEMM, dequantizing weights and
192 running QDQ on activations.
193 """
vllm/model_executor/layers/fused_moe/configs/E=384,N=256,device_name=AMD_Instinct_MI350X,dtype=int4_w4a16,backend=flydsl.json new
+114
@@ -0,0 +1,114 @@
1 +{
2 + "1": {
3 + "tile_m": 16,
4 + "tile_n": 64,
5 + "tile_k": 512,
6 + "tile_n2": 256,
7 + "tile_k2": 256
8 + },
9 + "2": {
10 + "tile_m": 16,
11 + "tile_n": 64,
12 + "tile_k": 512,
13 + "tile_n2": 256,
14 + "tile_k2": 128
15 + },
16 + "4": {
17 + "tile_m": 16,
18 + "tile_n": 64,
19 + "tile_k": 512,
20 + "tile_n2": 256,
21 + "tile_k2": 128
22 + },
23 + "8": {
24 + "tile_m": 16,
25 + "tile_n": 64,
26 + "tile_k": 512,
27 + "tile_n2": 256,
28 + "tile_k2": 256
29 + },
30 + "16": {
31 + "tile_m": 16,
32 + "tile_n": 64,
33 + "tile_k": 128,
34 + "tile_n2": 128,
35 + "tile_k2": 256
36 + },
37 + "24": {
38 + "tile_m": 16,
39 + "tile_n": 64,
40 + "tile_k": 128,
41 + "tile_n2": 256,
42 + "tile_k2": 256
43 + },
44 + "32": {
45 + "tile_m": 16,
46 + "tile_n": 64,
47 + "tile_k": 128,
48 + "tile_n2": 256,
49 + "tile_k2": 256
50 + },
51 + "48": {
52 + "tile_m": 16,
53 + "tile_n": 64,
54 + "tile_k": 128,
55 + "tile_n2": 256,
56 + "tile_k2": 256
57 + },
58 + "64": {
59 + "tile_m": 16,
60 + "tile_n": 64,
61 + "tile_k": 128,
62 + "tile_n2": 128,
63 + "tile_k2": 128
64 + },
65 + "128": {
66 + "tile_m": 16,
67 + "tile_n": 64,
68 + "tile_k": 128,
69 + "tile_n2": 256,
70 + "tile_k2": 256
71 + },
72 + "256": {
73 + "tile_m": 16,
74 + "tile_n": 128,
75 + "tile_k": 128,
76 + "tile_n2": 256,
77 + "tile_k2": 256
78 + },
79 + "512": {
80 + "tile_m": 16,
81 + "tile_n": 64,
82 + "tile_k": 128,
83 + "tile_n2": 256,
84 + "tile_k2": 256
85 + },
86 + "1024": {
87 + "tile_m": 32,
88 + "tile_n": 64,
89 + "tile_k": 128,
90 + "tile_n2": 256,
91 + "tile_k2": 256
92 + },
93 + "2048": {
94 + "tile_m": 64,
95 + "tile_n": 64,
96 + "tile_k": 64,
97 + "tile_n2": 256,
98 + "tile_k2": 64
99 + },
100 + "4096": {
101 + "tile_m": 32,
102 + "tile_n": 64,
103 + "tile_k": 128,
104 + "tile_n2": 256,
105 + "tile_k2": 256
106 + },
107 + "8192": {
108 + "tile_m": 64,
109 + "tile_n": 64,
110 + "tile_k": 64,
111 + "tile_n2": 256,
112 + "tile_k2": 64
113 + }
114 +}
\ No newline at end of file
vllm/model_executor/layers/fused_moe/configs/E=384,N=256,device_name=AMD_Instinct_MI350_OAM,dtype=int4_w4a16,backend=flydsl.json new
+114
@@ -0,0 +1,114 @@
1 +{
2 + "1": {
3 + "tile_m": 16,
4 + "tile_n": 64,
5 + "tile_k": 512,
6 + "tile_n2": 256,
7 + "tile_k2": 256
8 + },
9 + "2": {
10 + "tile_m": 16,
11 + "tile_n": 64,
12 + "tile_k": 512,
13 + "tile_n2": 256,
14 + "tile_k2": 128
15 + },
16 + "4": {
17 + "tile_m": 16,
18 + "tile_n": 64,
19 + "tile_k": 512,
20 + "tile_n2": 256,
21 + "tile_k2": 128
22 + },
23 + "8": {
24 + "tile_m": 16,
25 + "tile_n": 64,
26 + "tile_k": 512,
27 + "tile_n2": 256,
28 + "tile_k2": 256
29 + },
30 + "16": {
31 + "tile_m": 16,
32 + "tile_n": 64,
33 + "tile_k": 128,
34 + "tile_n2": 128,
35 + "tile_k2": 256
36 + },
37 + "24": {
38 + "tile_m": 16,
39 + "tile_n": 64,
40 + "tile_k": 128,
41 + "tile_n2": 256,
42 + "tile_k2": 256
43 + },
44 + "32": {
45 + "tile_m": 16,
46 + "tile_n": 64,
47 + "tile_k": 128,
48 + "tile_n2": 256,
49 + "tile_k2": 256
50 + },
51 + "48": {
52 + "tile_m": 16,
53 + "tile_n": 64,
54 + "tile_k": 128,
55 + "tile_n2": 256,
56 + "tile_k2": 256
57 + },
58 + "64": {
59 + "tile_m": 16,
60 + "tile_n": 64,
61 + "tile_k": 128,
62 + "tile_n2": 128,
63 + "tile_k2": 128
64 + },
65 + "128": {
66 + "tile_m": 16,
67 + "tile_n": 64,
68 + "tile_k": 128,
69 + "tile_n2": 256,
70 + "tile_k2": 256
71 + },
72 + "256": {
73 + "tile_m": 16,
74 + "tile_n": 128,
75 + "tile_k": 128,
76 + "tile_n2": 256,
77 + "tile_k2": 256
78 + },
79 + "512": {
80 + "tile_m": 16,
81 + "tile_n": 64,
82 + "tile_k": 128,
83 + "tile_n2": 256,
84 + "tile_k2": 256
85 + },
86 + "1024": {
87 + "tile_m": 32,
88 + "tile_n": 64,
89 + "tile_k": 128,
90 + "tile_n2": 256,
91 + "tile_k2": 256
92 + },
93 + "2048": {
94 + "tile_m": 64,
95 + "tile_n": 64,
96 + "tile_k": 64,
97 + "tile_n2": 256,
98 + "tile_k2": 64
99 + },
100 + "4096": {
101 + "tile_m": 32,
102 + "tile_n": 64,
103 + "tile_k": 128,
104 + "tile_n2": 256,
105 + "tile_k2": 256
106 + },
107 + "8192": {
108 + "tile_m": 64,
109 + "tile_n": 64,
110 + "tile_k": 64,
111 + "tile_n2": 256,
112 + "tile_k2": 64
113 + }
114 +}
\ No newline at end of file
vllm/model_executor/layers/fused_moe/configs/E=384,N=256,device_name=AMD_Instinct_MI355X,dtype=int4_w4a16,backend=flydsl.json new
+114
@@ -0,0 +1,114 @@
1 +{
2 + "1": {
3 + "tile_m": 16,
4 + "tile_n": 64,
5 + "tile_k": 512,
6 + "tile_n2": 256,
7 + "tile_k2": 256
8 + },
9 + "2": {
10 + "tile_m": 16,
11 + "tile_n": 64,
12 + "tile_k": 512,
13 + "tile_n2": 256,
14 + "tile_k2": 128
15 + },
16 + "4": {
17 + "tile_m": 16,
18 + "tile_n": 64,
19 + "tile_k": 512,
20 + "tile_n2": 256,
21 + "tile_k2": 128
22 + },
23 + "8": {
24 + "tile_m": 16,
25 + "tile_n": 64,
26 + "tile_k": 512,
27 + "tile_n2": 256,
28 + "tile_k2": 256
29 + },
30 + "16": {
31 + "tile_m": 16,
32 + "tile_n": 64,
33 + "tile_k": 128,
34 + "tile_n2": 128,
35 + "tile_k2": 256
36 + },
37 + "24": {
38 + "tile_m": 16,
39 + "tile_n": 64,
40 + "tile_k": 128,
41 + "tile_n2": 256,
42 + "tile_k2": 256
43 + },
44 + "32": {
45 + "tile_m": 16,
46 + "tile_n": 64,
47 + "tile_k": 128,
48 + "tile_n2": 256,
49 + "tile_k2": 256
50 + },
51 + "48": {
52 + "tile_m": 16,
53 + "tile_n": 64,
54 + "tile_k": 128,
55 + "tile_n2": 256,
56 + "tile_k2": 256
57 + },
58 + "64": {
59 + "tile_m": 16,
60 + "tile_n": 64,
61 + "tile_k": 128,
62 + "tile_n2": 128,
63 + "tile_k2": 128
64 + },
65 + "128": {
66 + "tile_m": 16,
67 + "tile_n": 64,
68 + "tile_k": 128,
69 + "tile_n2": 256,
70 + "tile_k2": 256
71 + },
72 + "256": {
73 + "tile_m": 16,
74 + "tile_n": 128,
75 + "tile_k": 128,
76 + "tile_n2": 256,
77 + "tile_k2": 256
78 + },
79 + "512": {
80 + "tile_m": 16,
81 + "tile_n": 64,
82 + "tile_k": 128,
83 + "tile_n2": 256,
84 + "tile_k2": 256
85 + },
86 + "1024": {
87 + "tile_m": 32,
88 + "tile_n": 64,
89 + "tile_k": 128,
90 + "tile_n2": 256,
91 + "tile_k2": 256
92 + },
93 + "2048": {
94 + "tile_m": 64,
95 + "tile_n": 64,
96 + "tile_k": 64,
97 + "tile_n2": 256,
98 + "tile_k2": 64
99 + },
100 + "4096": {
101 + "tile_m": 32,
102 + "tile_n": 64,
103 + "tile_k": 128,
104 + "tile_n2": 256,
105 + "tile_k2": 256
106 + },
107 + "8192": {
108 + "tile_m": 64,
109 + "tile_n": 64,
110 + "tile_k": 64,
111 + "tile_n2": 256,
112 + "tile_k2": 64
113 + }
114 +}
\ No newline at end of file
vllm/model_executor/layers/fused_moe/configs/E=384,N=256,device_name=AMD_Instinct_MI355_OAM,dtype=int4_w4a16,backend=flydsl.json new
+114
@@ -0,0 +1,114 @@
1 +{
2 + "1": {
3 + "tile_m": 16,
4 + "tile_n": 64,
5 + "tile_k": 512,
6 + "tile_n2": 256,
7 + "tile_k2": 256
8 + },
9 + "2": {
10 + "tile_m": 16,
11 + "tile_n": 64,
12 + "tile_k": 512,
13 + "tile_n2": 256,
14 + "tile_k2": 128
15 + },
16 + "4": {
17 + "tile_m": 16,
18 + "tile_n": 64,
19 + "tile_k": 512,
20 + "tile_n2": 256,
21 + "tile_k2": 128
22 + },
23 + "8": {
24 + "tile_m": 16,
25 + "tile_n": 64,
26 + "tile_k": 512,
27 + "tile_n2": 256,
28 + "tile_k2": 256
29 + },
30 + "16": {
31 + "tile_m": 16,
32 + "tile_n": 64,
33 + "tile_k": 128,
34 + "tile_n2": 128,
35 + "tile_k2": 256
36 + },
37 + "24": {
38 + "tile_m": 16,
39 + "tile_n": 64,
40 + "tile_k": 128,
41 + "tile_n2": 256,
42 + "tile_k2": 256
43 + },
44 + "32": {
45 + "tile_m": 16,
46 + "tile_n": 64,
47 + "tile_k": 128,
48 + "tile_n2": 256,
49 + "tile_k2": 256
50 + },
51 + "48": {
52 + "tile_m": 16,
53 + "tile_n": 64,
54 + "tile_k": 128,
55 + "tile_n2": 256,
56 + "tile_k2": 256
57 + },
58 + "64": {
59 + "tile_m": 16,
60 + "tile_n": 64,
61 + "tile_k": 128,
62 + "tile_n2": 128,
63 + "tile_k2": 128
64 + },
65 + "128": {
66 + "tile_m": 16,
67 + "tile_n": 64,
68 + "tile_k": 128,
69 + "tile_n2": 256,
70 + "tile_k2": 256
71 + },
72 + "256": {
73 + "tile_m": 16,
74 + "tile_n": 128,
75 + "tile_k": 128,
76 + "tile_n2": 256,
77 + "tile_k2": 256
78 + },
79 + "512": {
80 + "tile_m": 16,
81 + "tile_n": 64,
82 + "tile_k": 128,
83 + "tile_n2": 256,
84 + "tile_k2": 256
85 + },
86 + "1024": {
87 + "tile_m": 32,
88 + "tile_n": 64,
89 + "tile_k": 128,
90 + "tile_n2": 256,
91 + "tile_k2": 256
92 + },
93 + "2048": {
94 + "tile_m": 64,
95 + "tile_n": 64,
96 + "tile_k": 64,
97 + "tile_n2": 256,
98 + "tile_k2": 64
99 + },
100 + "4096": {
101 + "tile_m": 32,
102 + "tile_n": 64,
103 + "tile_k": 128,
104 + "tile_n2": 256,
105 + "tile_k2": 256
106 + },
107 + "8192": {
108 + "tile_m": 64,
109 + "tile_n": 64,
110 + "tile_k": 64,
111 + "tile_n2": 256,
112 + "tile_k2": 64
113 + }
114 +}
\ No newline at end of file
vllm/model_executor/layers/fused_moe/configs/E=384,N=512,device_name=AMD_Instinct_MI350X,dtype=int4_w4a16,backend=flydsl.json new
+114
@@ -0,0 +1,114 @@
1 +{
2 + "1": {
3 + "tile_m": 16,
4 + "tile_n": 64,
5 + "tile_k": 512,
6 + "tile_n2": 256,
7 + "tile_k2": 128
8 + },
9 + "2": {
10 + "tile_m": 16,
11 + "tile_n": 64,
12 + "tile_k": 512,
13 + "tile_n2": 128,
14 + "tile_k2": 128
15 + },
16 + "4": {
17 + "tile_m": 16,
18 + "tile_n": 64,
19 + "tile_k": 512,
20 + "tile_n2": 128,
21 + "tile_k2": 128
22 + },
23 + "8": {
24 + "tile_m": 16,
25 + "tile_n": 64,
26 + "tile_k": 128,
27 + "tile_n2": 128,
28 + "tile_k2": 128
29 + },
30 + "16": {
31 + "tile_m": 16,
32 + "tile_n": 64,
33 + "tile_k": 128,
34 + "tile_n2": 128,
35 + "tile_k2": 128
36 + },
37 + "24": {
38 + "tile_m": 16,
39 + "tile_n": 64,
40 + "tile_k": 128,
41 + "tile_n2": 256,
42 + "tile_k2": 128
43 + },
44 + "32": {
45 + "tile_m": 16,
46 + "tile_n": 64,
47 + "tile_k": 128,
48 + "tile_n2": 256,
49 + "tile_k2": 128
50 + },
51 + "48": {
52 + "tile_m": 16,
53 + "tile_n": 64,
54 + "tile_k": 128,
55 + "tile_n2": 256,
56 + "tile_k2": 128
57 + },
58 + "64": {
59 + "tile_m": 16,
60 + "tile_n": 64,
61 + "tile_k": 128,
62 + "tile_n2": 128,
63 + "tile_k2": 128
64 + },
65 + "128": {
66 + "tile_m": 16,
67 + "tile_n": 64,
68 + "tile_k": 128,
69 + "tile_n2": 256,
70 + "tile_k2": 128
71 + },
72 + "256": {
73 + "tile_m": 16,
74 + "tile_n": 64,
75 + "tile_k": 128,
76 + "tile_n2": 256,
77 + "tile_k2": 128
78 + },
79 + "512": {
80 + "tile_m": 16,
81 + "tile_n": 64,
82 + "tile_k": 128,
83 + "tile_n2": 128,
84 + "tile_k2": 128
85 + },
86 + "1024": {
87 + "tile_m": 32,
88 + "tile_n": 64,
89 + "tile_k": 128,
90 + "tile_n2": 128,
91 + "tile_k2": 64
92 + },
93 + "2048": {
94 + "tile_m": 64,
95 + "tile_n": 64,
96 + "tile_k": 64,
97 + "tile_n2": 128,
98 + "tile_k2": 64
99 + },
100 + "4096": {
101 + "tile_m": 32,
102 + "tile_n": 64,
103 + "tile_k": 64,
104 + "tile_n2": 128,
105 + "tile_k2": 64
106 + },
107 + "8192": {
108 + "tile_m": 64,
109 + "tile_n": 64,
110 + "tile_k": 64,
111 + "tile_n2": 128,
112 + "tile_k2": 128
113 + }
114 +}
\ No newline at end of file
vllm/model_executor/layers/fused_moe/configs/E=384,N=512,device_name=AMD_Instinct_MI350_OAM,dtype=int4_w4a16,backend=flydsl.json new
+114
@@ -0,0 +1,114 @@
1 +{
2 + "1": {
3 + "tile_m": 16,
4 + "tile_n": 64,
5 + "tile_k": 512,
6 + "tile_n2": 256,
7 + "tile_k2": 128
8 + },
9 + "2": {
10 + "tile_m": 16,
11 + "tile_n": 64,
12 + "tile_k": 512,
13 + "tile_n2": 128,
14 + "tile_k2": 128
15 + },
16 + "4": {
17 + "tile_m": 16,
18 + "tile_n": 64,
19 + "tile_k": 512,
20 + "tile_n2": 128,
21 + "tile_k2": 128
22 + },
23 + "8": {
24 + "tile_m": 16,
25 + "tile_n": 64,
26 + "tile_k": 128,
27 + "tile_n2": 128,
28 + "tile_k2": 128
29 + },
30 + "16": {
31 + "tile_m": 16,
32 + "tile_n": 64,
33 + "tile_k": 128,
34 + "tile_n2": 128,
35 + "tile_k2": 128
36 + },
37 + "24": {
38 + "tile_m": 16,
39 + "tile_n": 64,
40 + "tile_k": 128,
41 + "tile_n2": 256,
42 + "tile_k2": 128
43 + },
44 + "32": {
45 + "tile_m": 16,
46 + "tile_n": 64,
47 + "tile_k": 128,
48 + "tile_n2": 256,
49 + "tile_k2": 128
50 + },
51 + "48": {
52 + "tile_m": 16,
53 + "tile_n": 64,
54 + "tile_k": 128,
55 + "tile_n2": 256,
56 + "tile_k2": 128
57 + },
58 + "64": {
59 + "tile_m": 16,
60 + "tile_n": 64,
61 + "tile_k": 128,
62 + "tile_n2": 128,
63 + "tile_k2": 128
64 + },
65 + "128": {
66 + "tile_m": 16,
67 + "tile_n": 64,
68 + "tile_k": 128,
69 + "tile_n2": 256,
70 + "tile_k2": 128
71 + },
72 + "256": {
73 + "tile_m": 16,
74 + "tile_n": 64,
75 + "tile_k": 128,
76 + "tile_n2": 256,
77 + "tile_k2": 128
78 + },
79 + "512": {
80 + "tile_m": 16,
81 + "tile_n": 64,
82 + "tile_k": 128,
83 + "tile_n2": 128,
84 + "tile_k2": 128
85 + },
86 + "1024": {
87 + "tile_m": 32,
88 + "tile_n": 64,
89 + "tile_k": 128,
90 + "tile_n2": 128,
91 + "tile_k2": 64
92 + },
93 + "2048": {
94 + "tile_m": 64,
95 + "tile_n": 64,
96 + "tile_k": 64,
97 + "tile_n2": 128,
98 + "tile_k2": 64
99 + },
100 + "4096": {
101 + "tile_m": 32,
102 + "tile_n": 64,
103 + "tile_k": 64,
104 + "tile_n2": 128,
105 + "tile_k2": 64
106 + },
107 + "8192": {
108 + "tile_m": 64,
109 + "tile_n": 64,
110 + "tile_k": 64,
111 + "tile_n2": 128,
112 + "tile_k2": 128
113 + }
114 +}
\ No newline at end of file
vllm/model_executor/layers/fused_moe/configs/E=384,N=512,device_name=AMD_Instinct_MI355X,dtype=int4_w4a16,backend=flydsl.json new
+114
@@ -0,0 +1,114 @@
1 +{
2 + "1": {
3 + "tile_m": 16,
4 + "tile_n": 64,
5 + "tile_k": 512,
6 + "tile_n2": 256,
7 + "tile_k2": 128
8 + },
9 + "2": {
10 + "tile_m": 16,
11 + "tile_n": 64,
12 + "tile_k": 512,
13 + "tile_n2": 128,
14 + "tile_k2": 128
15 + },
16 + "4": {
17 + "tile_m": 16,
18 + "tile_n": 64,
19 + "tile_k": 512,
20 + "tile_n2": 128,
21 + "tile_k2": 128
22 + },
23 + "8": {
24 + "tile_m": 16,
25 + "tile_n": 64,
26 + "tile_k": 128,
27 + "tile_n2": 128,
28 + "tile_k2": 128
29 + },
30 + "16": {
31 + "tile_m": 16,
32 + "tile_n": 64,
33 + "tile_k": 128,
34 + "tile_n2": 128,
35 + "tile_k2": 128
36 + },
37 + "24": {
38 + "tile_m": 16,
39 + "tile_n": 64,
40 + "tile_k": 128,
41 + "tile_n2": 256,
42 + "tile_k2": 128
43 + },
44 + "32": {
45 + "tile_m": 16,
46 + "tile_n": 64,
47 + "tile_k": 128,
48 + "tile_n2": 256,
49 + "tile_k2": 128
50 + },
51 + "48": {
52 + "tile_m": 16,
53 + "tile_n": 64,
54 + "tile_k": 128,
55 + "tile_n2": 256,
56 + "tile_k2": 128
57 + },
58 + "64": {
59 + "tile_m": 16,
60 + "tile_n": 64,
61 + "tile_k": 128,
62 + "tile_n2": 128,
63 + "tile_k2": 128
64 + },
65 + "128": {
66 + "tile_m": 16,
67 + "tile_n": 64,
68 + "tile_k": 128,
69 + "tile_n2": 256,
70 + "tile_k2": 128
71 + },
72 + "256": {
73 + "tile_m": 16,
74 + "tile_n": 64,
75 + "tile_k": 128,
76 + "tile_n2": 256,
77 + "tile_k2": 128
78 + },
79 + "512": {
80 + "tile_m": 16,
81 + "tile_n": 64,
82 + "tile_k": 128,
83 + "tile_n2": 128,
84 + "tile_k2": 128
85 + },
86 + "1024": {
87 + "tile_m": 32,
88 + "tile_n": 64,
89 + "tile_k": 128,
90 + "tile_n2": 128,
91 + "tile_k2": 64
92 + },
93 + "2048": {
94 + "tile_m": 64,
95 + "tile_n": 64,
96 + "tile_k": 64,
97 + "tile_n2": 128,
98 + "tile_k2": 64
99 + },
100 + "4096": {
101 + "tile_m": 32,
102 + "tile_n": 64,
103 + "tile_k": 64,
104 + "tile_n2": 128,
105 + "tile_k2": 64
106 + },
107 + "8192": {
108 + "tile_m": 64,
109 + "tile_n": 64,
110 + "tile_k": 64,
111 + "tile_n2": 128,
112 + "tile_k2": 128
113 + }
114 +}
\ No newline at end of file
vllm/model_executor/layers/fused_moe/configs/E=384,N=512,device_name=AMD_Instinct_MI355_OAM,dtype=int4_w4a16,backend=flydsl.json new
+114
@@ -0,0 +1,114 @@
1 +{
2 + "1": {
3 + "tile_m": 16,
4 + "tile_n": 64,
5 + "tile_k": 512,
6 + "tile_n2": 256,
7 + "tile_k2": 128
8 + },
9 + "2": {
10 + "tile_m": 16,
11 + "tile_n": 64,
12 + "tile_k": 512,
13 + "tile_n2": 128,
14 + "tile_k2": 128
15 + },
16 + "4": {
17 + "tile_m": 16,
18 + "tile_n": 64,
19 + "tile_k": 512,
20 + "tile_n2": 128,
21 + "tile_k2": 128
22 + },
23 + "8": {
24 + "tile_m": 16,
25 + "tile_n": 64,
26 + "tile_k": 128,
27 + "tile_n2": 128,
28 + "tile_k2": 128
29 + },
30 + "16": {
31 + "tile_m": 16,
32 + "tile_n": 64,
33 + "tile_k": 128,
34 + "tile_n2": 128,
35 + "tile_k2": 128
36 + },
37 + "24": {
38 + "tile_m": 16,
39 + "tile_n": 64,
40 + "tile_k": 128,
41 + "tile_n2": 256,
42 + "tile_k2": 128
43 + },
44 + "32": {
45 + "tile_m": 16,
46 + "tile_n": 64,
47 + "tile_k": 128,
48 + "tile_n2": 256,
49 + "tile_k2": 128
50 + },
51 + "48": {
52 + "tile_m": 16,
53 + "tile_n": 64,
54 + "tile_k": 128,
55 + "tile_n2": 256,
56 + "tile_k2": 128
57 + },
58 + "64": {
59 + "tile_m": 16,
60 + "tile_n": 64,
61 + "tile_k": 128,
62 + "tile_n2": 128,
63 + "tile_k2": 128
64 + },
65 + "128": {
66 + "tile_m": 16,
67 + "tile_n": 64,
68 + "tile_k": 128,
69 + "tile_n2": 256,
70 + "tile_k2": 128
71 + },
72 + "256": {
73 + "tile_m": 16,
74 + "tile_n": 64,
75 + "tile_k": 128,
76 + "tile_n2": 256,
77 + "tile_k2": 128
78 + },
79 + "512": {
80 + "tile_m": 16,
81 + "tile_n": 64,
82 + "tile_k": 128,
83 + "tile_n2": 128,
84 + "tile_k2": 128
85 + },
86 + "1024": {
87 + "tile_m": 32,
88 + "tile_n": 64,
89 + "tile_k": 128,
90 + "tile_n2": 128,
91 + "tile_k2": 64
92 + },
93 + "2048": {
94 + "tile_m": 64,
95 + "tile_n": 64,
96 + "tile_k": 64,
97 + "tile_n2": 128,
98 + "tile_k2": 64
99 + },
100 + "4096": {
101 + "tile_m": 32,
102 + "tile_n": 64,
103 + "tile_k": 64,
104 + "tile_n2": 128,
105 + "tile_k2": 64
106 + },
107 + "8192": {
108 + "tile_m": 64,
109 + "tile_n": 64,
110 + "tile_k": 64,
111 + "tile_n2": 128,
112 + "tile_k2": 128
113 + }
114 +}
\ No newline at end of file
vllm/model_executor/layers/fused_moe/fused_flydsl_moe.py new
+430
@@ -0,0 +1,430 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Fused MoE Triton kernels."""
4 +
5 +import functools
6 +import json
7 +import os
8 +
9 +import flydsl.compiler as flyc
10 +import torch
11 +from aiter.fused_moe import moe_sorting as aiter_moe_sorting
12 +from aiter.ops.flydsl.kernels.moe_gemm_2stage import (
13 + compile_moe_gemm1,
14 + compile_moe_gemm2,
15 +)
16 +
17 +from vllm.logger import init_logger
18 +from vllm.platforms import current_platform
19 +from vllm.utils.torch_utils import direct_register_custom_op
20 +
21 +logger = init_logger(__name__)
22 +
23 +_FLYDSL_MOE_GEMM1_CACHE: dict = {}
24 +_FLYDSL_MOE_GEMM2_CACHE: dict = {}
25 +
26 +_FLYDSL_MOE_DEFAULT_CONFIG = {
27 + 1: {"tile_m": 16, "tile_n": 64, "tile_k": 512, "tile_n2": 256, "tile_k2": 256},
28 + 2: {"tile_m": 16, "tile_n": 64, "tile_k": 512, "tile_n2": 256, "tile_k2": 128},
29 + 4: {"tile_m": 16, "tile_n": 64, "tile_k": 512, "tile_n2": 256, "tile_k2": 128},
30 + 8: {"tile_m": 16, "tile_n": 64, "tile_k": 512, "tile_n2": 256, "tile_k2": 256},
31 + 16: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 256},
32 + 24: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 256, "tile_k2": 256},
33 + 32: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 256, "tile_k2": 256},
34 + 48: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 256, "tile_k2": 256},
35 + 64: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128},
36 + 128: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 256, "tile_k2": 256},
37 + 256: {"tile_m": 16, "tile_n": 128, "tile_k": 128, "tile_n2": 256, "tile_k2": 256},
38 + 512: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 256, "tile_k2": 256},
39 + 1024: {"tile_m": 32, "tile_n": 64, "tile_k": 128, "tile_n2": 256, "tile_k2": 256},
40 + 2048: {"tile_m": 64, "tile_n": 64, "tile_k": 64, "tile_n2": 256, "tile_k2": 64},
41 + 4096: {"tile_m": 32, "tile_n": 64, "tile_k": 128, "tile_n2": 256, "tile_k2": 256},
42 + 8192: {"tile_m": 64, "tile_n": 64, "tile_k": 64, "tile_n2": 256, "tile_k2": 64},
43 +}
44 +
45 +
46 +def moe_sorting(
47 + topk_ids: torch.Tensor,
48 + topk_weights: torch.Tensor,
49 + *,
50 + num_experts: int,
51 + model_dim: int,
52 + block_m: int,
53 +):
54 + topk_ids_i32 = topk_ids.to(torch.int32)
55 + topk_w_f32 = topk_weights.to(torch.float32)
56 + sorted_ids, sorted_w, sorted_expert_ids, num_valid_ids, _moe_buf = (
57 + aiter_moe_sorting(
58 + topk_ids_i32,
59 + topk_w_f32,
60 + num_experts,
61 + model_dim,
62 + torch.float16,
63 + block_m,
64 + )
65 + )
66 + if num_valid_ids.numel() > 1:
67 + num_valid_ids = num_valid_ids[:1].contiguous()
68 + return sorted_ids, sorted_w, sorted_expert_ids, num_valid_ids
69 +
70 +
71 +def build_routing_buffers(
72 + *,
73 + topk_ids: torch.Tensor,
74 + topk_weights: torch.Tensor,
75 + num_experts: int,
76 + model_dim: int,
77 + tile_m: int,
78 +):
79 + res = moe_sorting(
80 + topk_ids,
81 + topk_weights,
82 + num_experts=num_experts,
83 + model_dim=model_dim,
84 + block_m=tile_m,
85 + )
86 + if res is None:
87 + raise RuntimeError(
88 + "aiter moe_sorting failed/unavailable; cannot build routing buffers."
89 + )
90 + sorted_token_ids, sorted_weights, sorted_expert_ids, num_valid_ids = res
91 +
92 + sorted_token_ids = sorted_token_ids.contiguous()
93 + sorted_weights = sorted_weights.contiguous()
94 + sorted_expert_ids = sorted_expert_ids.contiguous()
95 + sorted_size = int(sorted_token_ids.numel())
96 + blocks = int(sorted_expert_ids.numel())
97 + return (
98 + sorted_token_ids,
99 + sorted_weights,
100 + sorted_expert_ids,
101 + num_valid_ids,
102 + sorted_size,
103 + blocks,
104 + )
105 +
106 +
107 +@functools.lru_cache
108 +def try_get_optimal_config(num_experts, inter_dim):
109 + device_name = current_platform.get_device_name().replace(" ", "_")
110 + json_file_name = (
111 + f"E={num_experts},N={inter_dim},device_name={device_name},"
112 + "dtype=int4_w4a16,backend=flydsl.json"
113 + )
114 + config_file_path = os.path.join(
115 + os.path.dirname(os.path.realpath(__file__)), "configs", json_file_name
116 + )
117 + if os.path.exists(config_file_path):
118 + with open(config_file_path) as f:
119 + logger.info_once(
120 + "Using tuned FlyDSL MoE config from %s",
121 + config_file_path,
122 + scope="global",
123 + )
124 + tuned_config = json.load(f)
125 + return {int(key): val for key, val in tuned_config.items()}
126 +
127 + logger.warning_once(
128 + "Using default FlyDSL MoE config. Performance might be sub-optimal! "
129 + "Config file not found at %s",
130 + config_file_path,
131 + scope="local",
132 + )
133 + return _FLYDSL_MOE_DEFAULT_CONFIG
134 +
135 +
136 +def fused_flydsl_moe_impl(
137 + hidden_states: torch.Tensor,
138 + w1: torch.Tensor,
139 + w2: torch.Tensor,
140 + num_experts: int,
141 + inter_dim: int,
142 + topk_weights: torch.Tensor,
143 + topk_ids: torch.Tensor,
144 + w1_scale: torch.Tensor | None = None,
145 + w2_scale: torch.Tensor | None = None,
146 + topk: int = 8,
147 + group_size: int = 32,
148 + doweight_stage1: bool = False,
149 + in_dtype: str = "int4_bf16",
150 + out_dtype: str = "bf16",
151 + scale_is_bf16: bool = True,
152 + tile_m: int | None = None,
153 + tile_n: int | None = None,
154 + tile_k: int | None = None,
155 + tile_n2: int | None = None,
156 + tile_k2: int | None = None,
157 +) -> torch.Tensor:
158 + device = hidden_states.device
159 + tokens = hidden_states.shape[0]
160 + model_dim = hidden_states.shape[1]
161 +
162 + tuned_config = {}
163 + if tile_m and tile_n and tile_k and tile_n2 and tile_k2:
164 + tuned_config["tile_m"] = tile_m
165 + tuned_config["tile_n"] = tile_n
166 + tuned_config["tile_k"] = tile_k
167 + tuned_config["tile_n2"] = tile_n2
168 + tuned_config["tile_k2"] = tile_k2
169 + else:
170 + tuned_config = try_get_optimal_config(num_experts, inter_dim)
171 + tuned_config = tuned_config[
172 + min(tuned_config.keys(), key=lambda x: abs(x - tokens))
173 + ]
174 + out_torch_dtype = torch.bfloat16 if out_dtype == "bf16" else torch.float16
175 +
176 + tile_m = tuned_config["tile_m"]
177 + tile_n = tuned_config["tile_n"]
178 + tile_k = tuned_config["tile_k"]
179 + tile_n2 = tuned_config["tile_n2"]
180 + tile_k2 = tuned_config["tile_k2"]
181 +
182 + routing = build_routing_buffers(
183 + topk_ids=topk_ids,
184 + topk_weights=topk_weights,
185 + num_experts=num_experts,
186 + model_dim=model_dim,
187 + tile_m=tile_m,
188 + )
189 + (
190 + sorted_token_ids,
191 + sorted_weights,
192 + sorted_expert_ids,
193 + num_valid_ids,
194 + sorted_size,
195 + blocks,
196 + ) = routing
197 +
198 + scale_x_1d = torch.empty((0,), device=device, dtype=torch.float32)
199 + sorted_weights_1d = sorted_weights.view(-1).contiguous()
200 + out_stage1 = torch.empty(
201 + (tokens, topk, inter_dim), device=device, dtype=out_torch_dtype
202 + )
203 +
204 + stream = torch.cuda.current_stream()
205 +
206 + key1 = (
207 + model_dim,
208 + inter_dim,
209 + num_experts,
210 + topk,
211 + in_dtype,
212 + out_dtype,
213 + group_size,
214 + tile_m,
215 + tile_n,
216 + tile_k,
217 + bool(doweight_stage1),
218 + False,
219 + )
220 +
221 + compiled_exe1 = _FLYDSL_MOE_GEMM1_CACHE.get(key1)
222 + if compiled_exe1 is None:
223 + exe1 = compile_moe_gemm1(
224 + model_dim=model_dim,
225 + inter_dim=inter_dim,
226 + experts=num_experts,
227 + topk=topk,
228 + in_dtype=in_dtype,
229 + out_dtype=out_dtype,
230 + group_size=group_size,
231 + tile_m=tile_m,
232 + tile_n=tile_n,
233 + tile_k=tile_k,
234 + doweight_stage1=bool(doweight_stage1),
235 + use_cshuffle_epilog=False,
236 + scale_is_bf16=scale_is_bf16,
237 + )
238 + compiled_exe1 = flyc.compile(
239 + exe1,
240 + out_stage1,
241 + hidden_states,
242 + w1,
243 + scale_x_1d,
244 + w1_scale,
245 + sorted_token_ids,
246 + sorted_expert_ids,
247 + sorted_weights_1d,
248 + num_valid_ids,
249 + tokens,
250 + inter_dim,
251 + model_dim,
252 + int(blocks),
253 + stream,
254 + )
255 + _FLYDSL_MOE_GEMM1_CACHE[key1] = compiled_exe1
256 +
257 + compiled_exe1(
258 + out_stage1,
259 + hidden_states,
260 + w1,
261 + scale_x_1d,
262 + w1_scale,
263 + sorted_token_ids,
264 + sorted_expert_ids,
265 + sorted_weights_1d,
266 + num_valid_ids,
267 + tokens,
268 + inter_dim,
269 + model_dim,
270 + int(blocks),
271 + stream,
272 + )
273 +
274 + a2_1d = out_stage1.view(-1).contiguous()
275 + a2_scale_1d = torch.empty((0,), device=device, dtype=torch.float32)
276 + out_stage2 = torch.empty((tokens, model_dim), device=device, dtype=out_torch_dtype)
277 + doweight_stage2 = not bool(doweight_stage1)
278 +
279 + key2 = (
280 + model_dim,
281 + inter_dim,
282 + num_experts,
283 + topk,
284 + in_dtype,
285 + out_dtype,
286 + group_size,
287 + tile_m,
288 + tile_n2,
289 + tile_k2,
290 + bool(doweight_stage2),
291 + )
292 +
293 + compiled_exe2 = _FLYDSL_MOE_GEMM2_CACHE.get(key2)
294 + if compiled_exe2 is None:
295 + exe2 = compile_moe_gemm2(
296 + model_dim=model_dim,
297 + inter_dim=inter_dim,
298 + experts=num_experts,
299 + topk=topk,
300 + in_dtype=in_dtype,
301 + out_dtype=out_dtype,
302 + group_size=group_size,
303 + tile_m=tile_m,
304 + tile_n=tile_n2,
305 + tile_k=tile_k2,
306 + doweight_stage2=bool(doweight_stage2),
307 + scale_is_bf16=scale_is_bf16,
308 + )
309 + compiled_exe2 = flyc.compile(
310 + exe2,
311 + out_stage2,
312 + a2_1d,
313 + w2,
314 + a2_scale_1d,
315 + w2_scale,
316 + sorted_token_ids,
317 + sorted_expert_ids,
318 + sorted_weights_1d,
319 + num_valid_ids,
320 + tokens,
321 + model_dim,
322 + inter_dim,
323 + int(blocks),
324 + stream,
325 + )
326 + _FLYDSL_MOE_GEMM2_CACHE[key2] = compiled_exe2
327 +
328 + out_stage2.zero_()
329 + compiled_exe2(
330 + out_stage2,
331 + a2_1d,
332 + w2,
333 + a2_scale_1d,
334 + w2_scale,
335 + sorted_token_ids,
336 + sorted_expert_ids,
337 + sorted_weights_1d,
338 + num_valid_ids,
339 + tokens,
340 + model_dim,
341 + inter_dim,
342 + int(blocks),
343 + stream,
344 + )
345 + return out_stage2
346 +
347 +
348 +def fused_flydsl_moe_impl_fake(
349 + hidden_states: torch.Tensor,
350 + w1: torch.Tensor,
351 + w2: torch.Tensor,
352 + num_experts: int,
353 + inter_dim: int,
354 + topk_weights: torch.Tensor,
355 + topk_ids: torch.Tensor,
356 + w1_scale: torch.Tensor | None = None,
357 + w2_scale: torch.Tensor | None = None,
358 + topk: int = 8,
359 + group_size: int = 32,
360 + doweight_stage1: bool = False,
361 + in_dtype: str = "int4_bf16",
362 + out_dtype: str = "bf16",
363 + scale_is_bf16: bool = True,
364 + tile_m: int | None = None,
365 + tile_n: int | None = None,
366 + tile_k: int | None = None,
367 + tile_n2: int | None = None,
368 + tile_k2: int | None = None,
369 +) -> torch.Tensor:
370 + return torch.empty_like(hidden_states)
371 +
372 +
373 +direct_register_custom_op(
374 + op_name="fused_flydsl_moe_impl",
375 + op_func=fused_flydsl_moe_impl,
376 + fake_impl=fused_flydsl_moe_impl_fake,
377 +)
378 +
379 +
380 +def fused_flydsl_moe(
381 + hidden_states: torch.Tensor,
382 + w1: torch.Tensor,
383 + w2: torch.Tensor,
384 + num_experts: int,
385 + inter_dim: int,
386 + topk_weights: torch.Tensor,
387 + topk_ids: torch.Tensor,
388 + w1_scale: torch.Tensor | None = None,
389 + w2_scale: torch.Tensor | None = None,
390 + topk: int = 8,
391 + group_size: int = 32,
392 + doweight_stage1: bool = False,
393 + in_dtype: str = "int4_bf16",
394 + out_dtype: str = "bf16",
395 + scale_is_bf16: bool = True,
396 + config: dict | None = None,
397 +) -> torch.Tensor:
398 + tile_m = None
399 + tile_n = None
400 + tile_k = None
401 + tile_n2 = None
402 + tile_k2 = None
403 + if config is not None:
404 + tile_m = config.get("tile_m")
405 + tile_n = config.get("tile_n")
406 + tile_k = config.get("tile_k")
407 + tile_n2 = config.get("tile_n2")
408 + tile_k2 = config.get("tile_k2")
409 + return torch.ops.vllm.fused_flydsl_moe_impl(
410 + hidden_states=hidden_states,
411 + w1=w1,
412 + w2=w2,
413 + num_experts=num_experts,
414 + inter_dim=inter_dim,
415 + topk_weights=topk_weights,
416 + topk_ids=topk_ids,
417 + w1_scale=w1_scale,
418 + w2_scale=w2_scale,
419 + topk=topk,
420 + group_size=group_size,
421 + doweight_stage1=doweight_stage1,
422 + in_dtype=in_dtype,
423 + out_dtype=out_dtype,
424 + scale_is_bf16=scale_is_bf16,
425 + tile_m=tile_m,
426 + tile_n=tile_n,
427 + tile_k=tile_k,
428 + tile_n2=tile_n2,
429 + tile_k2=tile_k2,
430 + )
vllm/model_executor/layers/fused_moe/routed_experts.py
+2
@@ -197,6 +197,7 @@ class RoutedExperts(PluggableLayer):
197 "AutoGPTQMoEMethod",
198 "CompressedTensorsWNA16MarlinMoEMethod",
199 "CompressedTensorsWNA16MoEMethod",
200 + "CompressedTensorsW4A16FlydslMoEMethod",
201 )
202
203 def _ensure_moe_quant_config_init(self):
@@ -610,6 +611,7 @@ class RoutedExperts(PluggableLayer):
611 "CompressedTensorsWNA16MarlinMoEMethod",
612 "CompressedTensorsWNA16MoEMethod",
613 "CompressedTensorsWNA16RDNA3MoEMethod",
614 + "CompressedTensorsW4A16FlydslMoEMethod",
615 ):
616 if is_transposed:
617 loaded_weight = loaded_weight.t().contiguous()
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py
+23
@@ -7,8 +7,10 @@ from compressed_tensors import CompressionFormat
7 from compressed_tensors.quantization import (
8 ActivationOrdering,
9 QuantizationStrategy,
10 + QuantizationType,
11 )
12
13 +from vllm.config import get_current_vllm_config
14 from vllm.logger import init_logger
15 from vllm.model_executor.layers.fused_moe import (
16 FusedMoEMethodBase,
@@ -115,7 +117,28 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
117 return rocm_moe_rdna.make_method(
118 weight_quant, input_quant, layer.moe_config
119 )
120 + from vllm.platforms.rocm import on_gfx950
121 +
122 + vllm_config = get_current_vllm_config()
123 + is_lora_disabled = vllm_config.lora_config is None
124 + moe_backend = vllm_config.kernel_config.moe_backend
125 + if (
126 + weight_quant.strategy == QuantizationStrategy.GROUP
127 + and weight_quant.type == QuantizationType.INT
128 + and group_size == 32
129 + and weight_quant.num_bits == 4
130 + and is_lora_disabled
131 + and on_gfx950()
132 + and moe_backend == "flydsl"
133 + ):
134 + from .compressed_tensors_moe_w4a16_flydsl import (
135 + CompressedTensorsW4A16FlydslMoEMethod,
136 + )
137
138 + logger.info_once("Using CompressedTensorsW4A16FlydslMoEMethod")
139 + return CompressedTensorsW4A16FlydslMoEMethod(
140 + weight_quant, input_quant, layer.moe_config
141 + )
142 from .compressed_tensors_moe_wna16 import (
143 CompressedTensorsWNA16MoEMethod,
144 )
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a16_flydsl.py new
+348
@@ -0,0 +1,348 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +
5 +import torch
6 +from aiter.ops.shuffle import shuffle_weight
7 +from compressed_tensors.quantization import (
8 + QuantizationArgs,
9 +)
10 +
11 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
12 +from vllm.logger import init_logger
13 +from vllm.model_executor.layers.fused_moe import (
14 + RoutedExperts,
15 + SharedExperts,
16 +)
17 +from vllm.model_executor.layers.fused_moe.config import (
18 + FusedMoEConfig,
19 + FusedMoEQuantConfig,
20 + int4_w4a16_moe_quant_config,
21 +)
22 +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501
23 + CompressedTensorsMoEMethod,
24 +)
25 +from vllm.model_executor.utils import set_weight_attrs
26 +
27 +logger = init_logger(__name__)
28 +
29 +
30 +def _pack_shuffled_int8_to_packed_int4_no_perm(x_shuf_i8: torch.Tensor) -> torch.Tensor:
31 + """Pack a preshuffled int8 tensor (values in [-8, 7]) into packed int4 bytes.
32 + Each contiguous 8-value block [v0..v7] -> 4 bytes:
33 + b0=(v4<<4)|v0, b1=(v5<<4)|v1, b2=(v6<<4)|v2, b3=(v7<<4)|v3.
34 + This matches the 7-op in-kernel unpack sequence and avoids any v_perm.
35 + """
36 + flat = x_shuf_i8.contiguous().view(-1).to(torch.int16)
37 + assert flat.numel() % 8 == 0
38 + u = (flat & 0xF).to(torch.uint8).view(-1, 8)
39 + out = torch.empty((u.shape[0], 4), device=u.device, dtype=torch.uint8)
40 + out[:, 0] = u[:, 0] | (u[:, 4] << 4)
41 + out[:, 1] = u[:, 1] | (u[:, 5] << 4)
42 + out[:, 2] = u[:, 2] | (u[:, 6] << 4)
43 + out[:, 3] = u[:, 3] | (u[:, 7] << 4)
44 + return out.view(-1).to(torch.int8)
45 +
46 +
47 +def _unpack_gptq_int32_to_signed_int4(w_int32):
48 + """Unpack GPTQ int32 [E, K//8, N] to signed int4 values [E, N, K] (as int8).
49 + Shared by both the packed-int4 and bf16-dequant paths.
50 + """
51 + E = w_int32.shape[0]
52 + # [E, K//8, N] -> transpose -> [E, N, K//8]
53 + w = w_int32.transpose(1, 2).contiguous()
54 + N = w.shape[1]
55 + K_div8 = w.shape[2]
56 + K = K_div8 * 8
57 +
58 + # Unpack int32 -> 8 x uint4 values along K
59 + w_expanded = w.unsqueeze(-1).expand(E, N, K_div8, 8) # [E, N, K//8, 8]
60 + shifts = torch.arange(8, device=w.device) * 4 # [0, 4, 8, ..., 28]
61 + nibbles = ((w_expanded >> shifts) & 0xF).to(torch.int8) # [E, N, K//8, 8]
62 + nibbles = nibbles.reshape(E, N, K) # [E, N, K] unsigned int4 as int8
63 +
64 + # Convert unsigned [0,15] to signed [-8,7]
65 + signed = nibbles.to(torch.int16) - 8
66 + signed = signed.to(torch.int8) # [E, N, K] signed int4 as int8
67 + return signed
68 +
69 +
70 +def _gptq_int32_to_flydsl_packed(w_int32):
71 + """Convert GPTQ int32 [E, K//8, N] to FlyDSL shuffled packed int4 [E, N, K//2].
72 + Steps:
73 + 1. Unpack int32 to individual signed int4 values (as int8)
74 + 2. Apply FlyDSL preshuffle (on individual int8 values)
75 + 3. Pack with FlyDSL's interleaved int4 packing
76 + """
77 + signed = _unpack_gptq_int32_to_signed_int4(w_int32)
78 + E, N, K = signed.shape
79 +
80 + # FlyDSL preshuffle (operates on individual values)
81 + shuffled = shuffle_weight(signed, layout=(16, 16))
82 +
83 + # FlyDSL interleaved int4 packing
84 + packed = _pack_shuffled_int8_to_packed_int4_no_perm(shuffled).contiguous()
85 + return packed.view(E, N, K // 2)
86 +
87 +
88 +class CompressedTensorsW4A16FlydslMoEMethod(CompressedTensorsMoEMethod):
89 + def __init__(
90 + self,
91 + weight_quant: QuantizationArgs,
92 + input_quant: QuantizationArgs | None,
93 + moe: FusedMoEConfig,
94 + layer_name: str | None = None,
95 + ):
96 + super().__init__(moe)
97 + self.weight_quant = weight_quant
98 + self.input_quant = input_quant
99 + # Extract properties from weight_quant
100 + assert weight_quant.num_bits == 4
101 + self.num_bits = weight_quant.num_bits
102 + self.packed_factor = 32 // weight_quant.num_bits
103 + self.strategy = weight_quant.strategy
104 + # channelwise is not supported by this kernel
105 + assert weight_quant.strategy == "group"
106 + assert weight_quant.group_size == 32
107 + self.group_size = weight_quant.group_size
108 + # grouped actorder isn't supported by this kernel
109 + assert weight_quant.actorder != "group"
110 + assert weight_quant.symmetric, (
111 + "Only symmetric quantization is supported for MoE"
112 + )
113 +
114 + def create_weights(
115 + self,
116 + layer: torch.nn.Module,
117 + num_experts: int,
118 + hidden_size: int,
119 + intermediate_size_per_partition: int,
120 + params_dtype: torch.dtype,
121 + **extra_weight_attrs,
122 + ):
123 + self.num_experts = num_experts
124 + self.inter_dim = intermediate_size_per_partition
125 + # Will transpose the loaded weight along the
126 + # intermediate and hidden dim sizes. Will
127 + # shard for TP along the transposed dims
128 + extra_weight_attrs.update(
129 + {"is_transposed": True, "quant_method": self.strategy}
130 + )
131 + w13_num_shards = 2 if self.moe.is_act_and_mul else 1
132 + w13_weight = torch.nn.Parameter(
133 + torch.empty(
134 + num_experts,
135 + hidden_size // self.packed_factor,
136 + w13_num_shards * intermediate_size_per_partition,
137 + dtype=torch.int32,
138 + ),
139 + requires_grad=False,
140 + )
141 + layer.register_parameter("w13_weight_packed", w13_weight)
142 + set_weight_attrs(w13_weight, extra_weight_attrs)
143 +
144 + w2_weight = torch.nn.Parameter(
145 + torch.empty(
146 + num_experts,
147 + intermediate_size_per_partition // self.packed_factor,
148 + hidden_size,
149 + dtype=torch.int32,
150 + ),
151 + requires_grad=False,
152 + )
153 + layer.register_parameter("w2_weight_packed", w2_weight)
154 + set_weight_attrs(w2_weight, extra_weight_attrs)
155 +
156 + w2_scales_size = intermediate_size_per_partition
157 +
158 + if self.strategy == "channel":
159 + num_groups_w2 = num_groups_w13 = 1
160 + self.group_size = -1
161 + else:
162 + num_groups_w2 = w2_scales_size // self.group_size
163 + num_groups_w13 = hidden_size // self.group_size
164 +
165 + w13_scale = torch.nn.Parameter(
166 + torch.ones(
167 + num_experts,
168 + num_groups_w13,
169 + w13_num_shards * intermediate_size_per_partition,
170 + dtype=params_dtype,
171 + ),
172 + requires_grad=False,
173 + )
174 + layer.register_parameter("w13_weight_scale", w13_scale)
175 + set_weight_attrs(w13_scale, extra_weight_attrs)
176 +
177 + w2_scale = torch.nn.Parameter(
178 + torch.ones(num_experts, num_groups_w2, hidden_size, dtype=params_dtype),
179 + requires_grad=False,
180 + )
181 + layer.register_parameter("w2_weight_scale", w2_scale)
182 + set_weight_attrs(w2_scale, extra_weight_attrs)
183 + set_weight_attrs(w2_scale, {"load_full_w2": False})
184 +
185 + w2_weight_shape = torch.nn.Parameter(
186 + torch.empty(num_experts, 2), requires_grad=False
187 + )
188 + layer.register_parameter("w2_weight_shape", w2_weight_shape)
189 + set_weight_attrs(w2_weight_shape, extra_weight_attrs)
190 + w13_weight_shape = torch.nn.Parameter(
191 + torch.empty(num_experts, 2), requires_grad=False
192 + )
193 +
194 + layer.register_parameter("w13_weight_shape", w13_weight_shape)
195 + set_weight_attrs(w13_weight_shape, extra_weight_attrs)
196 +
197 + w13_g_idx = torch.nn.Parameter(
198 + torch.empty(
199 + num_experts,
200 + hidden_size,
201 + dtype=torch.int32,
202 + ),
203 + requires_grad=False,
204 + )
205 + layer.register_parameter("w13_weight_g_idx", w13_g_idx)
206 + set_weight_attrs(w13_g_idx, extra_weight_attrs)
207 +
208 + w2_g_idx = torch.nn.Parameter(
209 + torch.empty(
210 + num_experts,
211 + intermediate_size_per_partition,
212 + dtype=torch.int32,
213 + ),
214 + requires_grad=False,
215 + )
216 + layer.register_parameter("w2_weight_g_idx", w2_g_idx)
217 + set_weight_attrs(w2_g_idx, extra_weight_attrs)
218 +
219 + w13_g_idx_sort_indices = torch.nn.Parameter(
220 + torch.empty(
221 + num_experts,
222 + hidden_size,
223 + dtype=torch.int32,
224 + ),
225 + requires_grad=False,
226 + )
227 + layer.register_parameter("w13_g_idx_sort_indices", w13_g_idx_sort_indices)
228 + set_weight_attrs(w13_g_idx_sort_indices, extra_weight_attrs)
229 +
230 + w2_g_idx_sort_indices = torch.nn.Parameter(
231 + torch.empty(
232 + num_experts,
233 + intermediate_size_per_partition,
234 + dtype=torch.int32,
235 + ),
236 + requires_grad=False,
237 + )
238 + layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices)
239 + set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs)
240 +
241 + layer.a13_scale = None
242 + layer.a2_scale = None
243 +
244 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
245 + # Reconfigure packed weights and scales to match flydsl_w4a16 format
246 +
247 + # Convert w13 weights
248 + w13 = layer.w13_weight_packed.data
249 + w13 = _gptq_int32_to_flydsl_packed(w13)
250 + w13 = w13.view(-1).contiguous()
251 + layer.w13_weight_packed = torch.nn.Parameter(w13, requires_grad=False)
252 +
253 + # Convert w2 weights
254 + w2 = layer.w2_weight_packed.data
255 + w2 = _gptq_int32_to_flydsl_packed(w2)
256 + w2 = w2.view(-1).contiguous()
257 + layer.w2_weight_packed = torch.nn.Parameter(w2, requires_grad=False)
258 +
259 + # Convert scales for FlyDSL:
260 + # per-row: [E, 1, N] -> squeeze -> [E, N]
261 + # groupwise: [E, K//gs, N] -> keep as-is (Opt 0: cache-friendly layout)
262 + w13_scale = layer.w13_weight_scale.data
263 + if self.group_size > 0 and w13_scale.dim() == 3 and w13_scale.shape[1] > 1:
264 + E, G, N = w13_scale.shape
265 + w13_scale = (
266 + w13_scale.view(E, G // 2, 2, N)
267 + .permute(0, 1, 3, 2)
268 + .contiguous()
269 + .view(-1)
270 + .contiguous()
271 + )
272 + elif w13_scale.dim() == 3 and w13_scale.shape[1] == 1:
273 + # Per-row: squeeze [E, 1, N] -> [E, N]
274 + w13_scale = w13_scale.squeeze(1)
275 + layer.w13_weight_scale = torch.nn.Parameter(
276 + w13_scale.contiguous(), requires_grad=False
277 + )
278 +
279 + w2_scale = layer.w2_weight_scale.data
280 + if self.group_size > 0 and w2_scale.dim() == 3 and w2_scale.shape[1] > 1:
281 + E, G, N = w2_scale.shape
282 + w2_scale = (
283 + w2_scale.view(E, G // 2, 2, N)
284 + .permute(0, 1, 3, 2)
285 + .contiguous()
286 + .view(-1)
287 + .contiguous()
288 + )
289 + elif w2_scale.dim() == 3 and w2_scale.shape[1] == 1:
290 + # Per-row: squeeze [E, 1, N] -> [E, N]
291 + w2_scale = w2_scale.squeeze(1)
292 + layer.w2_weight_scale = torch.nn.Parameter(
293 + w2_scale.contiguous(), requires_grad=False
294 + )
295 +
296 + layer.w13_weight_packed.is_shuffled = True
297 + layer.w2_weight_packed.is_shuffled = True
298 + layer.is_aiter_converted = True
299 +
300 + def get_fused_moe_quant_config(
301 + self, layer: torch.nn.Module
302 + ) -> FusedMoEQuantConfig | None:
303 + assert self.num_bits == 4
304 + return int4_w4a16_moe_quant_config(
305 + w1_scale=layer.w13_weight_scale,
306 + w2_scale=layer.w2_weight_scale,
307 + w1_zp=None,
308 + w2_zp=None,
309 + block_shape=[0, self.group_size],
310 + )
311 +
312 + def select_gemm_impl(
313 + self,
314 + prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular,
315 + layer: torch.nn.Module,
316 + ) -> mk.FusedMoEExpertsModular:
317 + raise NotImplementedError
318 +
319 + def apply(
320 + self,
321 + layer: RoutedExperts,
322 + x: torch.Tensor,
323 + topk_weights: torch.Tensor,
324 + topk_ids: torch.Tensor,
325 + shared_experts: SharedExperts | None,
326 + shared_experts_input: torch.Tensor | None,
327 + ) -> torch.Tensor:
328 + from vllm.model_executor.layers.fused_moe.fused_flydsl_moe import (
329 + fused_flydsl_moe,
330 + )
331 +
332 + assert self.moe_quant_config is not None
333 +
334 + return fused_flydsl_moe(
335 + x,
336 + layer.w13_weight_packed,
337 + layer.w2_weight_packed,
338 + self.num_experts,
339 + self.inter_dim,
340 + topk_weights,
341 + topk_ids,
342 + w1_scale=self.moe_quant_config.w1_scale,
343 + w2_scale=self.moe_quant_config.w2_scale,
344 + topk=topk_weights.shape[-1],
345 + group_size=self.group_size,
346 + doweight_stage1=layer.apply_router_weight_on_input,
347 + scale_is_bf16=True,
348 + )