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 9, 06:15 UTC
Add support for Mistral Large 3 inference with Flashinfer MoE (#33174)
Signed-off-by: Dimitrios Bariamis <12195802+dbari@users.noreply.github.com> Co-authored-by: Dimitrios Bariamis <12195802+dbari@users.noreply.github.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Dimitrios Bariamis committed
Jan 31, 2026 at 07:48 UTC
f0bca83ee4b6aa6d63da66d37dd69929bdcfc1fe
16 files changed
+1104
-31
benchmarks/kernels/benchmark_moe.py
+38
-11
@@ -27,7 +27,6 @@ from vllm.model_executor.layers.fused_moe.fused_moe import *
27
from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
28
TritonOrDeepGemmExperts,
29
)
30
-from vllm.platforms import current_platform
30
from vllm.transformers_utils.config import get_config
31
from vllm.triton_utils import triton
32
from vllm.utils.argparse_utils import FlexibleArgumentParser
@@ -482,6 +481,8 @@ class BenchmarkWorker:
481
block_quant_shape: list[int] = None,
482
use_deep_gemm: bool = False,
483
) -> tuple[dict[str, int], float]:
484
+ # local import to allow serialization by ray
485
+
486
set_random_seed(self.seed)
487
dtype_str = _get_config_dtype_str(
488
dtype, use_int8_w8a16=use_int8_w8a16, use_fp8_w8a8=use_fp8_w8a8
@@ -535,6 +536,9 @@ class BenchmarkWorker:
536
block_quant_shape: list[int],
537
use_deep_gemm: bool,
538
) -> dict[str, int]:
539
+ # local import to allow serialization by ray
540
+ from vllm.platforms import current_platform
541
+
542
best_config = None
543
best_time = float("inf")
544
if current_platform.is_rocm():
@@ -646,20 +650,28 @@ def save_configs(
650
f.write("\n")
651
652
653
+def get_compressed_tensors_block_structure(config, default_value=None):
654
+ config_groups = config.get("config_groups", {})
655
+ if len(config_groups) != 1:
656
+ return default_value
657
+ group = next(iter(config_groups.values()))
658
+ weights = group.get("weights", {})
659
+ block_structure = weights.get("block_structure", default_value)
660
+ return block_structure
661
+
662
+
663
def get_weight_block_size_safety(config, default_value=None):
664
quantization_config = getattr(config, "quantization_config", {})
665
if isinstance(quantization_config, dict):
652
- return quantization_config.get("weight_block_size", default_value)
666
+ if "weight_block_size" in quantization_config:
667
+ return quantization_config["weight_block_size"]
668
+ return get_compressed_tensors_block_structure(
669
+ quantization_config, default_value
670
+ )
671
return default_value
672
673
656
-def main(args: argparse.Namespace):
657
- print(args)
658
-
659
- config = get_config(model=args.model, trust_remote_code=args.trust_remote_code)
660
- if args.model_prefix:
661
- config = getattr(config, args.model_prefix)
662
-
674
+def get_model_params(config):
675
if config.architectures[0] == "DbrxForCausalLM":
676
E = config.ffn_config.moe_num_experts
677
topk = config.ffn_config.moe_top_k
@@ -677,6 +689,7 @@ def main(args: argparse.Namespace):
689
"Glm4MoeForCausalLM",
690
"Glm4MoeLiteForCausalLM",
691
"NemotronHForCausalLM",
692
+ "MistralLarge3ForCausalLM",
693
):
694
E = config.n_routed_experts
695
topk = config.num_experts_per_tok
@@ -697,16 +710,20 @@ def main(args: argparse.Namespace):
710
topk = text_config.num_experts_per_tok
711
intermediate_size = text_config.moe_intermediate_size
712
hidden_size = text_config.hidden_size
700
- elif config.architectures[0] in ("HunYuanMoEV1ForCausalLM"):
713
+ elif config.architectures[0] == "HunYuanMoEV1ForCausalLM":
714
E = config.num_experts
715
topk = config.moe_topk[0]
716
intermediate_size = config.moe_intermediate_size[0]
717
hidden_size = config.hidden_size
705
- elif config.architectures[0] in ["Qwen3OmniMoeForConditionalGeneration"]:
718
+ elif config.architectures[0] == "Qwen3OmniMoeForConditionalGeneration":
719
E = config.thinker_config.text_config.num_experts
720
topk = config.thinker_config.text_config.num_experts_per_tok
721
intermediate_size = config.thinker_config.text_config.moe_intermediate_size
722
hidden_size = config.thinker_config.text_config.hidden_size
723
+ elif config.architectures[0] == "PixtralForConditionalGeneration":
724
+ # Pixtral can contain different LLM architectures,
725
+ # recurse to get their parameters
726
+ return get_model_params(config.get_text_config())
727
else:
728
# Support for llama4
729
config = config.get_text_config()
@@ -715,6 +732,16 @@ def main(args: argparse.Namespace):
732
topk = config.num_experts_per_tok
733
intermediate_size = config.intermediate_size
734
hidden_size = config.hidden_size
735
+ return E, topk, intermediate_size, hidden_size
736
+
737
+
738
+def main(args: argparse.Namespace):
739
+ print(args)
740
+
741
+ config = get_config(model=args.model, trust_remote_code=args.trust_remote_code)
742
+ if args.model_prefix:
743
+ config = getattr(config, args.model_prefix)
744
+ E, topk, intermediate_size, hidden_size = get_model_params(config)
745
enable_ep = bool(args.enable_expert_parallel)
746
if enable_ep:
747
ensure_divisibility(E, args.tp_size, "Number of experts")
docker/Dockerfile
+1
-1
@@ -586,7 +586,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
586
# This is ~1.1GB and only changes when FlashInfer version bumps
587
# https://docs.flashinfer.ai/installation.html
588
# From versions.json: .flashinfer.version
589
-ARG FLASHINFER_VERSION=0.6.1
589
+ARG FLASHINFER_VERSION=0.6.2
590
RUN --mount=type=cache,target=/root/.cache/uv \
591
uv pip install --system flashinfer-cubin==${FLASHINFER_VERSION} \
592
&& uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
docker/Dockerfile.nightly_torch
+2
-2
@@ -221,13 +221,13 @@ RUN pip install setuptools==75.6.0 packaging==23.2 ninja==1.11.1.3 build==1.2.2.
221
222
223
# build flashinfer for torch nightly from source around 10 mins
224
-# release version: v0.6.1
224
+# release version: v0.6.2
225
# todo(elainewy): cache flashinfer build result for faster build
226
ENV CCACHE_DIR=/root/.cache/ccache
227
RUN --mount=type=cache,target=/root/.cache/ccache \
228
--mount=type=cache,target=/root/.cache/uv \
229
echo "git clone flashinfer..." \
230
- && git clone --depth 1 --branch v0.6.1 --recursive https://github.com/flashinfer-ai/flashinfer.git \
230
+ && git clone --depth 1 --branch v0.6.2 --recursive https://github.com/flashinfer-ai/flashinfer.git \
231
&& cd flashinfer \
232
&& git submodule update --init --recursive \
233
&& echo "finish git clone flashinfer..." \
docker/versions.json
+1
-1
@@ -68,7 +68,7 @@
68
"default": "true"
69
},
70
"FLASHINFER_VERSION": {
71
- "default": "0.6.1"
71
+ "default": "0.6.2"
72
},
73
"GDRCOPY_CUDA_VERSION": {
74
"default": "12.8"
requirements/cuda.txt
+1
-1
@@ -10,4 +10,4 @@ torchaudio==2.9.1
10
# These must be updated alongside torch
11
torchvision==0.24.1 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
12
# FlashInfer should be updated together with the Dockerfile
13
-flashinfer-python==0.6.1
13
+flashinfer-python==0.6.2
tests/kernels/moe/test_flashinfer.py
+2
@@ -135,6 +135,8 @@ class TestData:
135
layer.w2_input_scale,
136
)
137
layer.custom_routing_function = Llama4MoE.custom_routing_function
138
+ layer.routing_method_type = RoutingMethodType.Llama4
139
+ layer.renormalize = False
140
layer.intermediate_size_per_partition = n
141
layer.ep_rank = 0
142
layer.local_num_experts = e
vllm/model_executor/layers/fused_moe/configs/E=128,N=512,device_name=NVIDIA_B200,dtype=fp8_w8a8.json
new
+147
@@ -0,0 +1,147 @@
1
+{
2
+ "triton_version": "3.4.0",
3
+ "1": {
4
+ "BLOCK_SIZE_M": 16,
5
+ "BLOCK_SIZE_N": 32,
6
+ "BLOCK_SIZE_K": 256,
7
+ "GROUP_SIZE_M": 64,
8
+ "num_warps": 4,
9
+ "num_stages": 3
10
+ },
11
+ "2": {
12
+ "BLOCK_SIZE_M": 16,
13
+ "BLOCK_SIZE_N": 64,
14
+ "BLOCK_SIZE_K": 128,
15
+ "GROUP_SIZE_M": 1,
16
+ "num_warps": 4,
17
+ "num_stages": 5
18
+ },
19
+ "4": {
20
+ "BLOCK_SIZE_M": 16,
21
+ "BLOCK_SIZE_N": 128,
22
+ "BLOCK_SIZE_K": 128,
23
+ "GROUP_SIZE_M": 1,
24
+ "num_warps": 4,
25
+ "num_stages": 5
26
+ },
27
+ "8": {
28
+ "BLOCK_SIZE_M": 16,
29
+ "BLOCK_SIZE_N": 128,
30
+ "BLOCK_SIZE_K": 128,
31
+ "GROUP_SIZE_M": 1,
32
+ "num_warps": 4,
33
+ "num_stages": 5
34
+ },
35
+ "16": {
36
+ "BLOCK_SIZE_M": 16,
37
+ "BLOCK_SIZE_N": 128,
38
+ "BLOCK_SIZE_K": 128,
39
+ "GROUP_SIZE_M": 1,
40
+ "num_warps": 4,
41
+ "num_stages": 3
42
+ },
43
+ "24": {
44
+ "BLOCK_SIZE_M": 16,
45
+ "BLOCK_SIZE_N": 256,
46
+ "BLOCK_SIZE_K": 256,
47
+ "GROUP_SIZE_M": 1,
48
+ "num_warps": 8,
49
+ "num_stages": 2
50
+ },
51
+ "32": {
52
+ "BLOCK_SIZE_M": 16,
53
+ "BLOCK_SIZE_N": 128,
54
+ "BLOCK_SIZE_K": 128,
55
+ "GROUP_SIZE_M": 1,
56
+ "num_warps": 4,
57
+ "num_stages": 3
58
+ },
59
+ "48": {
60
+ "BLOCK_SIZE_M": 16,
61
+ "BLOCK_SIZE_N": 128,
62
+ "BLOCK_SIZE_K": 256,
63
+ "GROUP_SIZE_M": 1,
64
+ "num_warps": 8,
65
+ "num_stages": 3
66
+ },
67
+ "64": {
68
+ "BLOCK_SIZE_M": 16,
69
+ "BLOCK_SIZE_N": 128,
70
+ "BLOCK_SIZE_K": 128,
71
+ "GROUP_SIZE_M": 1,
72
+ "num_warps": 4,
73
+ "num_stages": 4
74
+ },
75
+ "96": {
76
+ "BLOCK_SIZE_M": 16,
77
+ "BLOCK_SIZE_N": 128,
78
+ "BLOCK_SIZE_K": 128,
79
+ "GROUP_SIZE_M": 1,
80
+ "num_warps": 4,
81
+ "num_stages": 4
82
+ },
83
+ "128": {
84
+ "BLOCK_SIZE_M": 16,
85
+ "BLOCK_SIZE_N": 256,
86
+ "BLOCK_SIZE_K": 256,
87
+ "GROUP_SIZE_M": 1,
88
+ "num_warps": 8,
89
+ "num_stages": 2
90
+ },
91
+ "256": {
92
+ "BLOCK_SIZE_M": 16,
93
+ "BLOCK_SIZE_N": 256,
94
+ "BLOCK_SIZE_K": 256,
95
+ "GROUP_SIZE_M": 1,
96
+ "num_warps": 8,
97
+ "num_stages": 2
98
+ },
99
+ "512": {
100
+ "BLOCK_SIZE_M": 16,
101
+ "BLOCK_SIZE_N": 128,
102
+ "BLOCK_SIZE_K": 128,
103
+ "GROUP_SIZE_M": 1,
104
+ "num_warps": 4,
105
+ "num_stages": 3
106
+ },
107
+ "1024": {
108
+ "BLOCK_SIZE_M": 64,
109
+ "BLOCK_SIZE_N": 256,
110
+ "BLOCK_SIZE_K": 64,
111
+ "GROUP_SIZE_M": 1,
112
+ "num_warps": 8,
113
+ "num_stages": 5
114
+ },
115
+ "1536": {
116
+ "BLOCK_SIZE_M": 64,
117
+ "BLOCK_SIZE_N": 256,
118
+ "BLOCK_SIZE_K": 64,
119
+ "GROUP_SIZE_M": 1,
120
+ "num_warps": 4,
121
+ "num_stages": 5
122
+ },
123
+ "2048": {
124
+ "BLOCK_SIZE_M": 64,
125
+ "BLOCK_SIZE_N": 128,
126
+ "BLOCK_SIZE_K": 128,
127
+ "GROUP_SIZE_M": 1,
128
+ "num_warps": 4,
129
+ "num_stages": 3
130
+ },
131
+ "3072": {
132
+ "BLOCK_SIZE_M": 128,
133
+ "BLOCK_SIZE_N": 256,
134
+ "BLOCK_SIZE_K": 128,
135
+ "GROUP_SIZE_M": 1,
136
+ "num_warps": 8,
137
+ "num_stages": 4
138
+ },
139
+ "4096": {
140
+ "BLOCK_SIZE_M": 128,
141
+ "BLOCK_SIZE_N": 256,
142
+ "BLOCK_SIZE_K": 128,
143
+ "GROUP_SIZE_M": 1,
144
+ "num_warps": 8,
145
+ "num_stages": 4
146
+ }
147
+}
vllm/model_executor/layers/fused_moe/configs/E=128,N=512,device_name=NVIDIA_B200.json
new
+147
@@ -0,0 +1,147 @@
1
+{
2
+ "triton_version": "3.4.0",
3
+ "1": {
4
+ "BLOCK_SIZE_M": 16,
5
+ "BLOCK_SIZE_N": 32,
6
+ "BLOCK_SIZE_K": 64,
7
+ "GROUP_SIZE_M": 32,
8
+ "num_warps": 4,
9
+ "num_stages": 5
10
+ },
11
+ "2": {
12
+ "BLOCK_SIZE_M": 16,
13
+ "BLOCK_SIZE_N": 32,
14
+ "BLOCK_SIZE_K": 64,
15
+ "GROUP_SIZE_M": 1,
16
+ "num_warps": 4,
17
+ "num_stages": 4
18
+ },
19
+ "4": {
20
+ "BLOCK_SIZE_M": 16,
21
+ "BLOCK_SIZE_N": 128,
22
+ "BLOCK_SIZE_K": 64,
23
+ "GROUP_SIZE_M": 64,
24
+ "num_warps": 4,
25
+ "num_stages": 4
26
+ },
27
+ "8": {
28
+ "BLOCK_SIZE_M": 16,
29
+ "BLOCK_SIZE_N": 128,
30
+ "BLOCK_SIZE_K": 64,
31
+ "GROUP_SIZE_M": 1,
32
+ "num_warps": 8,
33
+ "num_stages": 4
34
+ },
35
+ "16": {
36
+ "BLOCK_SIZE_M": 16,
37
+ "BLOCK_SIZE_N": 128,
38
+ "BLOCK_SIZE_K": 64,
39
+ "GROUP_SIZE_M": 32,
40
+ "num_warps": 4,
41
+ "num_stages": 3
42
+ },
43
+ "24": {
44
+ "BLOCK_SIZE_M": 16,
45
+ "BLOCK_SIZE_N": 128,
46
+ "BLOCK_SIZE_K": 64,
47
+ "GROUP_SIZE_M": 16,
48
+ "num_warps": 4,
49
+ "num_stages": 4
50
+ },
51
+ "32": {
52
+ "BLOCK_SIZE_M": 16,
53
+ "BLOCK_SIZE_N": 64,
54
+ "BLOCK_SIZE_K": 64,
55
+ "GROUP_SIZE_M": 16,
56
+ "num_warps": 4,
57
+ "num_stages": 3
58
+ },
59
+ "48": {
60
+ "BLOCK_SIZE_M": 16,
61
+ "BLOCK_SIZE_N": 128,
62
+ "BLOCK_SIZE_K": 64,
63
+ "GROUP_SIZE_M": 64,
64
+ "num_warps": 4,
65
+ "num_stages": 4
66
+ },
67
+ "64": {
68
+ "BLOCK_SIZE_M": 16,
69
+ "BLOCK_SIZE_N": 128,
70
+ "BLOCK_SIZE_K": 64,
71
+ "GROUP_SIZE_M": 64,
72
+ "num_warps": 4,
73
+ "num_stages": 4
74
+ },
75
+ "96": {
76
+ "BLOCK_SIZE_M": 16,
77
+ "BLOCK_SIZE_N": 64,
78
+ "BLOCK_SIZE_K": 64,
79
+ "GROUP_SIZE_M": 16,
80
+ "num_warps": 4,
81
+ "num_stages": 3
82
+ },
83
+ "128": {
84
+ "BLOCK_SIZE_M": 16,
85
+ "BLOCK_SIZE_N": 64,
86
+ "BLOCK_SIZE_K": 64,
87
+ "GROUP_SIZE_M": 16,
88
+ "num_warps": 4,
89
+ "num_stages": 3
90
+ },
91
+ "256": {
92
+ "BLOCK_SIZE_M": 16,
93
+ "BLOCK_SIZE_N": 64,
94
+ "BLOCK_SIZE_K": 64,
95
+ "GROUP_SIZE_M": 16,
96
+ "num_warps": 4,
97
+ "num_stages": 3
98
+ },
99
+ "512": {
100
+ "BLOCK_SIZE_M": 32,
101
+ "BLOCK_SIZE_N": 256,
102
+ "BLOCK_SIZE_K": 64,
103
+ "GROUP_SIZE_M": 1,
104
+ "num_warps": 8,
105
+ "num_stages": 4
106
+ },
107
+ "1024": {
108
+ "BLOCK_SIZE_M": 64,
109
+ "BLOCK_SIZE_N": 128,
110
+ "BLOCK_SIZE_K": 64,
111
+ "GROUP_SIZE_M": 16,
112
+ "num_warps": 8,
113
+ "num_stages": 4
114
+ },
115
+ "1536": {
116
+ "BLOCK_SIZE_M": 64,
117
+ "BLOCK_SIZE_N": 128,
118
+ "BLOCK_SIZE_K": 64,
119
+ "GROUP_SIZE_M": 1,
120
+ "num_warps": 8,
121
+ "num_stages": 2
122
+ },
123
+ "2048": {
124
+ "BLOCK_SIZE_M": 128,
125
+ "BLOCK_SIZE_N": 128,
126
+ "BLOCK_SIZE_K": 64,
127
+ "GROUP_SIZE_M": 1,
128
+ "num_warps": 8,
129
+ "num_stages": 3
130
+ },
131
+ "3072": {
132
+ "BLOCK_SIZE_M": 128,
133
+ "BLOCK_SIZE_N": 128,
134
+ "BLOCK_SIZE_K": 64,
135
+ "GROUP_SIZE_M": 1,
136
+ "num_warps": 8,
137
+ "num_stages": 3
138
+ },
139
+ "4096": {
140
+ "BLOCK_SIZE_M": 128,
141
+ "BLOCK_SIZE_N": 128,
142
+ "BLOCK_SIZE_K": 64,
143
+ "GROUP_SIZE_M": 16,
144
+ "num_warps": 8,
145
+ "num_stages": 3
146
+ }
147
+}
vllm/model_executor/layers/fused_moe/configs/E=128,N=512,device_name=NVIDIA_GB200,dtype=fp8_w8a8.json
new
+147
@@ -0,0 +1,147 @@
1
+{
2
+ "triton_version": "3.4.0",
3
+ "1": {
4
+ "BLOCK_SIZE_M": 16,
5
+ "BLOCK_SIZE_N": 32,
6
+ "BLOCK_SIZE_K": 256,
7
+ "GROUP_SIZE_M": 64,
8
+ "num_warps": 4,
9
+ "num_stages": 3
10
+ },
11
+ "2": {
12
+ "BLOCK_SIZE_M": 16,
13
+ "BLOCK_SIZE_N": 64,
14
+ "BLOCK_SIZE_K": 128,
15
+ "GROUP_SIZE_M": 1,
16
+ "num_warps": 4,
17
+ "num_stages": 5
18
+ },
19
+ "4": {
20
+ "BLOCK_SIZE_M": 16,
21
+ "BLOCK_SIZE_N": 128,
22
+ "BLOCK_SIZE_K": 128,
23
+ "GROUP_SIZE_M": 1,
24
+ "num_warps": 4,
25
+ "num_stages": 5
26
+ },
27
+ "8": {
28
+ "BLOCK_SIZE_M": 16,
29
+ "BLOCK_SIZE_N": 128,
30
+ "BLOCK_SIZE_K": 128,
31
+ "GROUP_SIZE_M": 1,
32
+ "num_warps": 4,
33
+ "num_stages": 5
34
+ },
35
+ "16": {
36
+ "BLOCK_SIZE_M": 16,
37
+ "BLOCK_SIZE_N": 128,
38
+ "BLOCK_SIZE_K": 128,
39
+ "GROUP_SIZE_M": 1,
40
+ "num_warps": 4,
41
+ "num_stages": 3
42
+ },
43
+ "24": {
44
+ "BLOCK_SIZE_M": 16,
45
+ "BLOCK_SIZE_N": 256,
46
+ "BLOCK_SIZE_K": 256,
47
+ "GROUP_SIZE_M": 1,
48
+ "num_warps": 8,
49
+ "num_stages": 2
50
+ },
51
+ "32": {
52
+ "BLOCK_SIZE_M": 16,
53
+ "BLOCK_SIZE_N": 128,
54
+ "BLOCK_SIZE_K": 128,
55
+ "GROUP_SIZE_M": 1,
56
+ "num_warps": 4,
57
+ "num_stages": 3
58
+ },
59
+ "48": {
60
+ "BLOCK_SIZE_M": 16,
61
+ "BLOCK_SIZE_N": 128,
62
+ "BLOCK_SIZE_K": 256,
63
+ "GROUP_SIZE_M": 1,
64
+ "num_warps": 8,
65
+ "num_stages": 3
66
+ },
67
+ "64": {
68
+ "BLOCK_SIZE_M": 16,
69
+ "BLOCK_SIZE_N": 128,
70
+ "BLOCK_SIZE_K": 128,
71
+ "GROUP_SIZE_M": 1,
72
+ "num_warps": 4,
73
+ "num_stages": 4
74
+ },
75
+ "96": {
76
+ "BLOCK_SIZE_M": 16,
77
+ "BLOCK_SIZE_N": 128,
78
+ "BLOCK_SIZE_K": 128,
79
+ "GROUP_SIZE_M": 1,
80
+ "num_warps": 4,
81
+ "num_stages": 4
82
+ },
83
+ "128": {
84
+ "BLOCK_SIZE_M": 16,
85
+ "BLOCK_SIZE_N": 256,
86
+ "BLOCK_SIZE_K": 256,
87
+ "GROUP_SIZE_M": 1,
88
+ "num_warps": 8,
89
+ "num_stages": 2
90
+ },
91
+ "256": {
92
+ "BLOCK_SIZE_M": 16,
93
+ "BLOCK_SIZE_N": 256,
94
+ "BLOCK_SIZE_K": 256,
95
+ "GROUP_SIZE_M": 1,
96
+ "num_warps": 8,
97
+ "num_stages": 2
98
+ },
99
+ "512": {
100
+ "BLOCK_SIZE_M": 16,
101
+ "BLOCK_SIZE_N": 128,
102
+ "BLOCK_SIZE_K": 128,
103
+ "GROUP_SIZE_M": 1,
104
+ "num_warps": 4,
105
+ "num_stages": 3
106
+ },
107
+ "1024": {
108
+ "BLOCK_SIZE_M": 64,
109
+ "BLOCK_SIZE_N": 256,
110
+ "BLOCK_SIZE_K": 64,
111
+ "GROUP_SIZE_M": 1,
112
+ "num_warps": 8,
113
+ "num_stages": 5
114
+ },
115
+ "1536": {
116
+ "BLOCK_SIZE_M": 64,
117
+ "BLOCK_SIZE_N": 256,
118
+ "BLOCK_SIZE_K": 64,
119
+ "GROUP_SIZE_M": 1,
120
+ "num_warps": 4,
121
+ "num_stages": 5
122
+ },
123
+ "2048": {
124
+ "BLOCK_SIZE_M": 64,
125
+ "BLOCK_SIZE_N": 128,
126
+ "BLOCK_SIZE_K": 128,
127
+ "GROUP_SIZE_M": 1,
128
+ "num_warps": 4,
129
+ "num_stages": 3
130
+ },
131
+ "3072": {
132
+ "BLOCK_SIZE_M": 128,
133
+ "BLOCK_SIZE_N": 256,
134
+ "BLOCK_SIZE_K": 128,
135
+ "GROUP_SIZE_M": 1,
136
+ "num_warps": 8,
137
+ "num_stages": 4
138
+ },
139
+ "4096": {
140
+ "BLOCK_SIZE_M": 128,
141
+ "BLOCK_SIZE_N": 256,
142
+ "BLOCK_SIZE_K": 128,
143
+ "GROUP_SIZE_M": 1,
144
+ "num_warps": 8,
145
+ "num_stages": 4
146
+ }
147
+}
vllm/model_executor/layers/fused_moe/configs/E=128,N=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json
new
+147
@@ -0,0 +1,147 @@
1
+{
2
+ "triton_version": "3.5.0",
3
+ "1": {
4
+ "BLOCK_SIZE_M": 16,
5
+ "BLOCK_SIZE_N": 128,
6
+ "BLOCK_SIZE_K": 128,
7
+ "GROUP_SIZE_M": 1,
8
+ "num_warps": 4,
9
+ "num_stages": 4
10
+ },
11
+ "2": {
12
+ "BLOCK_SIZE_M": 16,
13
+ "BLOCK_SIZE_N": 128,
14
+ "BLOCK_SIZE_K": 128,
15
+ "GROUP_SIZE_M": 1,
16
+ "num_warps": 4,
17
+ "num_stages": 4
18
+ },
19
+ "4": {
20
+ "BLOCK_SIZE_M": 16,
21
+ "BLOCK_SIZE_N": 128,
22
+ "BLOCK_SIZE_K": 128,
23
+ "GROUP_SIZE_M": 64,
24
+ "num_warps": 4,
25
+ "num_stages": 4
26
+ },
27
+ "8": {
28
+ "BLOCK_SIZE_M": 16,
29
+ "BLOCK_SIZE_N": 128,
30
+ "BLOCK_SIZE_K": 128,
31
+ "GROUP_SIZE_M": 64,
32
+ "num_warps": 4,
33
+ "num_stages": 4
34
+ },
35
+ "16": {
36
+ "BLOCK_SIZE_M": 16,
37
+ "BLOCK_SIZE_N": 128,
38
+ "BLOCK_SIZE_K": 256,
39
+ "GROUP_SIZE_M": 16,
40
+ "num_warps": 4,
41
+ "num_stages": 3
42
+ },
43
+ "24": {
44
+ "BLOCK_SIZE_M": 16,
45
+ "BLOCK_SIZE_N": 128,
46
+ "BLOCK_SIZE_K": 128,
47
+ "GROUP_SIZE_M": 64,
48
+ "num_warps": 4,
49
+ "num_stages": 3
50
+ },
51
+ "32": {
52
+ "BLOCK_SIZE_M": 16,
53
+ "BLOCK_SIZE_N": 128,
54
+ "BLOCK_SIZE_K": 128,
55
+ "GROUP_SIZE_M": 1,
56
+ "num_warps": 4,
57
+ "num_stages": 5
58
+ },
59
+ "48": {
60
+ "BLOCK_SIZE_M": 16,
61
+ "BLOCK_SIZE_N": 128,
62
+ "BLOCK_SIZE_K": 128,
63
+ "GROUP_SIZE_M": 1,
64
+ "num_warps": 4,
65
+ "num_stages": 3
66
+ },
67
+ "64": {
68
+ "BLOCK_SIZE_M": 16,
69
+ "BLOCK_SIZE_N": 128,
70
+ "BLOCK_SIZE_K": 256,
71
+ "GROUP_SIZE_M": 16,
72
+ "num_warps": 4,
73
+ "num_stages": 3
74
+ },
75
+ "96": {
76
+ "BLOCK_SIZE_M": 16,
77
+ "BLOCK_SIZE_N": 128,
78
+ "BLOCK_SIZE_K": 128,
79
+ "GROUP_SIZE_M": 1,
80
+ "num_warps": 4,
81
+ "num_stages": 3
82
+ },
83
+ "128": {
84
+ "BLOCK_SIZE_M": 16,
85
+ "BLOCK_SIZE_N": 128,
86
+ "BLOCK_SIZE_K": 128,
87
+ "GROUP_SIZE_M": 1,
88
+ "num_warps": 4,
89
+ "num_stages": 3
90
+ },
91
+ "256": {
92
+ "BLOCK_SIZE_M": 16,
93
+ "BLOCK_SIZE_N": 128,
94
+ "BLOCK_SIZE_K": 256,
95
+ "GROUP_SIZE_M": 1,
96
+ "num_warps": 4,
97
+ "num_stages": 3
98
+ },
99
+ "512": {
100
+ "BLOCK_SIZE_M": 16,
101
+ "BLOCK_SIZE_N": 128,
102
+ "BLOCK_SIZE_K": 128,
103
+ "GROUP_SIZE_M": 1,
104
+ "num_warps": 4,
105
+ "num_stages": 3
106
+ },
107
+ "1024": {
108
+ "BLOCK_SIZE_M": 64,
109
+ "BLOCK_SIZE_N": 128,
110
+ "BLOCK_SIZE_K": 256,
111
+ "GROUP_SIZE_M": 16,
112
+ "num_warps": 4,
113
+ "num_stages": 4
114
+ },
115
+ "1536": {
116
+ "BLOCK_SIZE_M": 64,
117
+ "BLOCK_SIZE_N": 128,
118
+ "BLOCK_SIZE_K": 128,
119
+ "GROUP_SIZE_M": 16,
120
+ "num_warps": 4,
121
+ "num_stages": 4
122
+ },
123
+ "2048": {
124
+ "BLOCK_SIZE_M": 64,
125
+ "BLOCK_SIZE_N": 128,
126
+ "BLOCK_SIZE_K": 128,
127
+ "GROUP_SIZE_M": 32,
128
+ "num_warps": 4,
129
+ "num_stages": 4
130
+ },
131
+ "3072": {
132
+ "BLOCK_SIZE_M": 64,
133
+ "BLOCK_SIZE_N": 128,
134
+ "BLOCK_SIZE_K": 256,
135
+ "GROUP_SIZE_M": 32,
136
+ "num_warps": 4,
137
+ "num_stages": 3
138
+ },
139
+ "4096": {
140
+ "BLOCK_SIZE_M": 64,
141
+ "BLOCK_SIZE_N": 128,
142
+ "BLOCK_SIZE_K": 256,
143
+ "GROUP_SIZE_M": 1,
144
+ "num_warps": 4,
145
+ "num_stages": 3
146
+ }
147
+}
vllm/model_executor/layers/fused_moe/configs/E=128,N=512,device_name=NVIDIA_H200.json
new
+147
@@ -0,0 +1,147 @@
1
+{
2
+ "triton_version": "3.4.0",
3
+ "1": {
4
+ "BLOCK_SIZE_M": 16,
5
+ "BLOCK_SIZE_N": 32,
6
+ "BLOCK_SIZE_K": 64,
7
+ "GROUP_SIZE_M": 1,
8
+ "num_warps": 4,
9
+ "num_stages": 5
10
+ },
11
+ "2": {
12
+ "BLOCK_SIZE_M": 16,
13
+ "BLOCK_SIZE_N": 32,
14
+ "BLOCK_SIZE_K": 64,
15
+ "GROUP_SIZE_M": 1,
16
+ "num_warps": 4,
17
+ "num_stages": 5
18
+ },
19
+ "4": {
20
+ "BLOCK_SIZE_M": 16,
21
+ "BLOCK_SIZE_N": 128,
22
+ "BLOCK_SIZE_K": 64,
23
+ "GROUP_SIZE_M": 64,
24
+ "num_warps": 8,
25
+ "num_stages": 4
26
+ },
27
+ "8": {
28
+ "BLOCK_SIZE_M": 16,
29
+ "BLOCK_SIZE_N": 32,
30
+ "BLOCK_SIZE_K": 64,
31
+ "GROUP_SIZE_M": 1,
32
+ "num_warps": 4,
33
+ "num_stages": 2
34
+ },
35
+ "16": {
36
+ "BLOCK_SIZE_M": 16,
37
+ "BLOCK_SIZE_N": 64,
38
+ "BLOCK_SIZE_K": 64,
39
+ "GROUP_SIZE_M": 64,
40
+ "num_warps": 8,
41
+ "num_stages": 4
42
+ },
43
+ "24": {
44
+ "BLOCK_SIZE_M": 16,
45
+ "BLOCK_SIZE_N": 128,
46
+ "BLOCK_SIZE_K": 64,
47
+ "GROUP_SIZE_M": 1,
48
+ "num_warps": 4,
49
+ "num_stages": 4
50
+ },
51
+ "32": {
52
+ "BLOCK_SIZE_M": 16,
53
+ "BLOCK_SIZE_N": 64,
54
+ "BLOCK_SIZE_K": 64,
55
+ "GROUP_SIZE_M": 64,
56
+ "num_warps": 4,
57
+ "num_stages": 3
58
+ },
59
+ "48": {
60
+ "BLOCK_SIZE_M": 16,
61
+ "BLOCK_SIZE_N": 128,
62
+ "BLOCK_SIZE_K": 64,
63
+ "GROUP_SIZE_M": 16,
64
+ "num_warps": 4,
65
+ "num_stages": 4
66
+ },
67
+ "64": {
68
+ "BLOCK_SIZE_M": 16,
69
+ "BLOCK_SIZE_N": 64,
70
+ "BLOCK_SIZE_K": 64,
71
+ "GROUP_SIZE_M": 1,
72
+ "num_warps": 4,
73
+ "num_stages": 3
74
+ },
75
+ "96": {
76
+ "BLOCK_SIZE_M": 16,
77
+ "BLOCK_SIZE_N": 64,
78
+ "BLOCK_SIZE_K": 64,
79
+ "GROUP_SIZE_M": 16,
80
+ "num_warps": 4,
81
+ "num_stages": 3
82
+ },
83
+ "128": {
84
+ "BLOCK_SIZE_M": 16,
85
+ "BLOCK_SIZE_N": 64,
86
+ "BLOCK_SIZE_K": 64,
87
+ "GROUP_SIZE_M": 1,
88
+ "num_warps": 4,
89
+ "num_stages": 4
90
+ },
91
+ "256": {
92
+ "BLOCK_SIZE_M": 16,
93
+ "BLOCK_SIZE_N": 64,
94
+ "BLOCK_SIZE_K": 64,
95
+ "GROUP_SIZE_M": 1,
96
+ "num_warps": 4,
97
+ "num_stages": 4
98
+ },
99
+ "512": {
100
+ "BLOCK_SIZE_M": 32,
101
+ "BLOCK_SIZE_N": 256,
102
+ "BLOCK_SIZE_K": 64,
103
+ "GROUP_SIZE_M": 1,
104
+ "num_warps": 8,
105
+ "num_stages": 3
106
+ },
107
+ "1024": {
108
+ "BLOCK_SIZE_M": 64,
109
+ "BLOCK_SIZE_N": 128,
110
+ "BLOCK_SIZE_K": 64,
111
+ "GROUP_SIZE_M": 1,
112
+ "num_warps": 8,
113
+ "num_stages": 4
114
+ },
115
+ "1536": {
116
+ "BLOCK_SIZE_M": 64,
117
+ "BLOCK_SIZE_N": 128,
118
+ "BLOCK_SIZE_K": 64,
119
+ "GROUP_SIZE_M": 1,
120
+ "num_warps": 8,
121
+ "num_stages": 4
122
+ },
123
+ "2048": {
124
+ "BLOCK_SIZE_M": 128,
125
+ "BLOCK_SIZE_N": 128,
126
+ "BLOCK_SIZE_K": 64,
127
+ "GROUP_SIZE_M": 1,
128
+ "num_warps": 8,
129
+ "num_stages": 3
130
+ },
131
+ "3072": {
132
+ "BLOCK_SIZE_M": 128,
133
+ "BLOCK_SIZE_N": 128,
134
+ "BLOCK_SIZE_K": 64,
135
+ "GROUP_SIZE_M": 1,
136
+ "num_warps": 8,
137
+ "num_stages": 3
138
+ },
139
+ "4096": {
140
+ "BLOCK_SIZE_M": 128,
141
+ "BLOCK_SIZE_N": 128,
142
+ "BLOCK_SIZE_K": 64,
143
+ "GROUP_SIZE_M": 1,
144
+ "num_warps": 8,
145
+ "num_stages": 3
146
+ }
147
+}
vllm/model_executor/layers/fused_moe/configs/E=16,N=4096,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128,128].json
new
+147
@@ -0,0 +1,147 @@
1
+{
2
+ "triton_version": "3.5.1",
3
+ "1": {
4
+ "BLOCK_SIZE_M": 16,
5
+ "BLOCK_SIZE_N": 128,
6
+ "BLOCK_SIZE_K": 128,
7
+ "GROUP_SIZE_M": 64,
8
+ "num_warps": 4,
9
+ "num_stages": 5
10
+ },
11
+ "2": {
12
+ "BLOCK_SIZE_M": 16,
13
+ "BLOCK_SIZE_N": 128,
14
+ "BLOCK_SIZE_K": 128,
15
+ "GROUP_SIZE_M": 1,
16
+ "num_warps": 4,
17
+ "num_stages": 3
18
+ },
19
+ "4": {
20
+ "BLOCK_SIZE_M": 16,
21
+ "BLOCK_SIZE_N": 128,
22
+ "BLOCK_SIZE_K": 256,
23
+ "GROUP_SIZE_M": 1,
24
+ "num_warps": 8,
25
+ "num_stages": 4
26
+ },
27
+ "8": {
28
+ "BLOCK_SIZE_M": 16,
29
+ "BLOCK_SIZE_N": 128,
30
+ "BLOCK_SIZE_K": 128,
31
+ "GROUP_SIZE_M": 64,
32
+ "num_warps": 4,
33
+ "num_stages": 5
34
+ },
35
+ "16": {
36
+ "BLOCK_SIZE_M": 16,
37
+ "BLOCK_SIZE_N": 128,
38
+ "BLOCK_SIZE_K": 128,
39
+ "GROUP_SIZE_M": 1,
40
+ "num_warps": 4,
41
+ "num_stages": 3
42
+ },
43
+ "24": {
44
+ "BLOCK_SIZE_M": 16,
45
+ "BLOCK_SIZE_N": 128,
46
+ "BLOCK_SIZE_K": 128,
47
+ "GROUP_SIZE_M": 1,
48
+ "num_warps": 4,
49
+ "num_stages": 3
50
+ },
51
+ "32": {
52
+ "BLOCK_SIZE_M": 16,
53
+ "BLOCK_SIZE_N": 128,
54
+ "BLOCK_SIZE_K": 128,
55
+ "GROUP_SIZE_M": 1,
56
+ "num_warps": 4,
57
+ "num_stages": 3
58
+ },
59
+ "48": {
60
+ "BLOCK_SIZE_M": 16,
61
+ "BLOCK_SIZE_N": 128,
62
+ "BLOCK_SIZE_K": 128,
63
+ "GROUP_SIZE_M": 1,
64
+ "num_warps": 4,
65
+ "num_stages": 3
66
+ },
67
+ "64": {
68
+ "BLOCK_SIZE_M": 16,
69
+ "BLOCK_SIZE_N": 128,
70
+ "BLOCK_SIZE_K": 128,
71
+ "GROUP_SIZE_M": 64,
72
+ "num_warps": 4,
73
+ "num_stages": 4
74
+ },
75
+ "96": {
76
+ "BLOCK_SIZE_M": 64,
77
+ "BLOCK_SIZE_N": 128,
78
+ "BLOCK_SIZE_K": 128,
79
+ "GROUP_SIZE_M": 32,
80
+ "num_warps": 4,
81
+ "num_stages": 4
82
+ },
83
+ "128": {
84
+ "BLOCK_SIZE_M": 64,
85
+ "BLOCK_SIZE_N": 128,
86
+ "BLOCK_SIZE_K": 128,
87
+ "GROUP_SIZE_M": 32,
88
+ "num_warps": 4,
89
+ "num_stages": 4
90
+ },
91
+ "256": {
92
+ "BLOCK_SIZE_M": 64,
93
+ "BLOCK_SIZE_N": 256,
94
+ "BLOCK_SIZE_K": 256,
95
+ "GROUP_SIZE_M": 64,
96
+ "num_warps": 8,
97
+ "num_stages": 4
98
+ },
99
+ "512": {
100
+ "BLOCK_SIZE_M": 64,
101
+ "BLOCK_SIZE_N": 256,
102
+ "BLOCK_SIZE_K": 128,
103
+ "GROUP_SIZE_M": 64,
104
+ "num_warps": 8,
105
+ "num_stages": 4
106
+ },
107
+ "1024": {
108
+ "BLOCK_SIZE_M": 64,
109
+ "BLOCK_SIZE_N": 128,
110
+ "BLOCK_SIZE_K": 128,
111
+ "GROUP_SIZE_M": 1,
112
+ "num_warps": 4,
113
+ "num_stages": 3
114
+ },
115
+ "1536": {
116
+ "BLOCK_SIZE_M": 64,
117
+ "BLOCK_SIZE_N": 128,
118
+ "BLOCK_SIZE_K": 128,
119
+ "GROUP_SIZE_M": 64,
120
+ "num_warps": 4,
121
+ "num_stages": 4
122
+ },
123
+ "2048": {
124
+ "BLOCK_SIZE_M": 64,
125
+ "BLOCK_SIZE_N": 256,
126
+ "BLOCK_SIZE_K": 256,
127
+ "GROUP_SIZE_M": 64,
128
+ "num_warps": 8,
129
+ "num_stages": 4
130
+ },
131
+ "3072": {
132
+ "BLOCK_SIZE_M": 64,
133
+ "BLOCK_SIZE_N": 128,
134
+ "BLOCK_SIZE_K": 256,
135
+ "GROUP_SIZE_M": 64,
136
+ "num_warps": 4,
137
+ "num_stages": 4
138
+ },
139
+ "4096": {
140
+ "BLOCK_SIZE_M": 64,
141
+ "BLOCK_SIZE_N": 256,
142
+ "BLOCK_SIZE_K": 256,
143
+ "GROUP_SIZE_M": 64,
144
+ "num_warps": 8,
145
+ "num_stages": 4
146
+ }
147
+}
vllm/model_executor/layers/fused_moe/configs/E=16,N=4096,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json
new
+147
@@ -0,0 +1,147 @@
1
+{
2
+ "triton_version": "3.5.0",
3
+ "1": {
4
+ "BLOCK_SIZE_M": 16,
5
+ "BLOCK_SIZE_N": 128,
6
+ "BLOCK_SIZE_K": 128,
7
+ "GROUP_SIZE_M": 1,
8
+ "num_warps": 4,
9
+ "num_stages": 4
10
+ },
11
+ "2": {
12
+ "BLOCK_SIZE_M": 16,
13
+ "BLOCK_SIZE_N": 128,
14
+ "BLOCK_SIZE_K": 128,
15
+ "GROUP_SIZE_M": 1,
16
+ "num_warps": 4,
17
+ "num_stages": 4
18
+ },
19
+ "4": {
20
+ "BLOCK_SIZE_M": 16,
21
+ "BLOCK_SIZE_N": 128,
22
+ "BLOCK_SIZE_K": 128,
23
+ "GROUP_SIZE_M": 16,
24
+ "num_warps": 4,
25
+ "num_stages": 5
26
+ },
27
+ "8": {
28
+ "BLOCK_SIZE_M": 16,
29
+ "BLOCK_SIZE_N": 128,
30
+ "BLOCK_SIZE_K": 128,
31
+ "GROUP_SIZE_M": 1,
32
+ "num_warps": 4,
33
+ "num_stages": 3
34
+ },
35
+ "16": {
36
+ "BLOCK_SIZE_M": 16,
37
+ "BLOCK_SIZE_N": 128,
38
+ "BLOCK_SIZE_K": 128,
39
+ "GROUP_SIZE_M": 1,
40
+ "num_warps": 4,
41
+ "num_stages": 3
42
+ },
43
+ "24": {
44
+ "BLOCK_SIZE_M": 16,
45
+ "BLOCK_SIZE_N": 128,
46
+ "BLOCK_SIZE_K": 128,
47
+ "GROUP_SIZE_M": 1,
48
+ "num_warps": 4,
49
+ "num_stages": 3
50
+ },
51
+ "32": {
52
+ "BLOCK_SIZE_M": 16,
53
+ "BLOCK_SIZE_N": 128,
54
+ "BLOCK_SIZE_K": 128,
55
+ "GROUP_SIZE_M": 1,
56
+ "num_warps": 4,
57
+ "num_stages": 3
58
+ },
59
+ "48": {
60
+ "BLOCK_SIZE_M": 16,
61
+ "BLOCK_SIZE_N": 128,
62
+ "BLOCK_SIZE_K": 128,
63
+ "GROUP_SIZE_M": 1,
64
+ "num_warps": 4,
65
+ "num_stages": 3
66
+ },
67
+ "64": {
68
+ "BLOCK_SIZE_M": 16,
69
+ "BLOCK_SIZE_N": 128,
70
+ "BLOCK_SIZE_K": 128,
71
+ "GROUP_SIZE_M": 1,
72
+ "num_warps": 4,
73
+ "num_stages": 4
74
+ },
75
+ "96": {
76
+ "BLOCK_SIZE_M": 16,
77
+ "BLOCK_SIZE_N": 128,
78
+ "BLOCK_SIZE_K": 128,
79
+ "GROUP_SIZE_M": 1,
80
+ "num_warps": 4,
81
+ "num_stages": 3
82
+ },
83
+ "128": {
84
+ "BLOCK_SIZE_M": 16,
85
+ "BLOCK_SIZE_N": 128,
86
+ "BLOCK_SIZE_K": 128,
87
+ "GROUP_SIZE_M": 1,
88
+ "num_warps": 4,
89
+ "num_stages": 3
90
+ },
91
+ "256": {
92
+ "BLOCK_SIZE_M": 64,
93
+ "BLOCK_SIZE_N": 128,
94
+ "BLOCK_SIZE_K": 128,
95
+ "GROUP_SIZE_M": 64,
96
+ "num_warps": 4,
97
+ "num_stages": 3
98
+ },
99
+ "512": {
100
+ "BLOCK_SIZE_M": 64,
101
+ "BLOCK_SIZE_N": 128,
102
+ "BLOCK_SIZE_K": 128,
103
+ "GROUP_SIZE_M": 32,
104
+ "num_warps": 4,
105
+ "num_stages": 2
106
+ },
107
+ "1024": {
108
+ "BLOCK_SIZE_M": 16,
109
+ "BLOCK_SIZE_N": 128,
110
+ "BLOCK_SIZE_K": 128,
111
+ "GROUP_SIZE_M": 1,
112
+ "num_warps": 4,
113
+ "num_stages": 2
114
+ },
115
+ "1536": {
116
+ "BLOCK_SIZE_M": 16,
117
+ "BLOCK_SIZE_N": 128,
118
+ "BLOCK_SIZE_K": 128,
119
+ "GROUP_SIZE_M": 1,
120
+ "num_warps": 4,
121
+ "num_stages": 2
122
+ },
123
+ "2048": {
124
+ "BLOCK_SIZE_M": 64,
125
+ "BLOCK_SIZE_N": 256,
126
+ "BLOCK_SIZE_K": 128,
127
+ "GROUP_SIZE_M": 32,
128
+ "num_warps": 8,
129
+ "num_stages": 3
130
+ },
131
+ "3072": {
132
+ "BLOCK_SIZE_M": 16,
133
+ "BLOCK_SIZE_N": 128,
134
+ "BLOCK_SIZE_K": 128,
135
+ "GROUP_SIZE_M": 1,
136
+ "num_warps": 4,
137
+ "num_stages": 2
138
+ },
139
+ "4096": {
140
+ "BLOCK_SIZE_M": 64,
141
+ "BLOCK_SIZE_N": 256,
142
+ "BLOCK_SIZE_K": 256,
143
+ "GROUP_SIZE_M": 32,
144
+ "num_warps": 8,
145
+ "num_stages": 3
146
+ }
147
+}
vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py
+6
-3
@@ -70,9 +70,12 @@ def _supports_routing_method(
70
RoutingMethodType.RenormalizeNaive,
71
]
72
elif (weight_key, activation_key) == (kFp8StaticTensorSym, kFp8StaticTensorSym):
73
- # NOTE(rob): kernel requires Llama4.
74
- return routing_method == RoutingMethodType.Llama4
75
-
73
+ # NOTE(dbari): as above, potentially allow others here.
74
+ return routing_method in [
75
+ RoutingMethodType.Llama4,
76
+ RoutingMethodType.Renormalize,
77
+ RoutingMethodType.RenormalizeNaive,
78
+ ]
79
else:
80
raise ValueError("Unsupported quantization scheme.")
81
vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+13
-9
@@ -113,15 +113,19 @@ def apply_fi_trtllm_fp8_per_tensor_moe(
113
and hasattr(layer, "output2_scales_scalar")
114
)
115
116
- # Added to the layer by: register_scales_for_trtllm_fp8_per_tensor_moe
117
- assert (
118
- hasattr(layer, "output1_scales_scalar")
119
- and hasattr(layer, "output1_scales_gate_scalar")
120
- and hasattr(layer, "output2_scales_scalar")
121
- )
116
+ if layer.routing_method_type == RoutingMethodType.Llama4:
117
+ assert (
118
+ not layer.renormalize
119
+ and layer.custom_routing_function == Llama4MoE.custom_routing_function
120
+ ), (
121
+ "FusedMoE flashinfer kernels with Llama4 routing method are only "
122
+ "supported for Llama4"
123
+ )
124
+ else:
125
+ assert layer.custom_routing_function is None, (
126
+ "Custom routing function is only supported for Llama4"
127
+ )
128
123
- is_llama4 = layer.custom_routing_function == Llama4MoE.custom_routing_function
124
- assert is_llama4, "FusedMoE flashinfer kernels are only supported for Llama4"
129
return torch.ops.vllm.fi_trtllm_fp8_per_tensor_moe(
130
routing_logits=router_logits,
131
routing_bias=routing_bias,
@@ -140,7 +144,7 @@ def apply_fi_trtllm_fp8_per_tensor_moe(
144
local_expert_offset=layer.ep_rank * layer.local_num_experts,
145
local_num_experts=layer.local_num_experts,
146
use_routing_scales_on_input=apply_router_weight_on_input,
143
- routing_method_type=RoutingMethodType.Llama4,
147
+ routing_method_type=layer.routing_method_type,
148
)
149
150
vllm/model_executor/models/deepseek_v2.py
+11
-3
@@ -295,6 +295,14 @@ class DeepseekV2MoE(nn.Module):
295
prefix=f"{prefix}.shared_experts",
296
)
297
298
+ n_group = getattr(config, "n_group", 1)
299
+ topk_group = getattr(config, "topk_group", 1)
300
+ use_grouped_topk = True
301
+ if (n_group, topk_group) == (1, 1):
302
+ n_group = None
303
+ topk_group = None
304
+ use_grouped_topk = False
305
+
306
self.experts = SharedFusedMoE(
307
shared_experts=self.shared_experts,
308
gate=self.gate,
@@ -305,9 +313,9 @@ class DeepseekV2MoE(nn.Module):
313
reduce_results=False,
314
renormalize=config.norm_topk_prob,
315
quant_config=quant_config,
308
- use_grouped_topk=True,
309
- num_expert_group=getattr(config, "n_group", 1),
310
- topk_group=getattr(config, "topk_group", 1),
316
+ use_grouped_topk=use_grouped_topk,
317
+ num_expert_group=n_group,
318
+ topk_group=topk_group,
319
prefix=f"{prefix}.experts",
320
scoring_func=getattr(config, "scoring_func", "softmax"),
321
# we do scaling outside, set factor to 1.0 to avoid double mul