A high-throughput and memory-efficient inference and serving engine for LLMs
[Bugfix] Fix MoE model load OOM in FlashInfer_TRTLLM backend with sleep mode (#45589)
Signed-off-by: Dakai An <dakaian108@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Dakai An committed
Jun 17, 2026 at 10:36 UTC
71bc19dbdd0db0154367e27d8072fdcacf02ba4a
2 files changed
+43
-32
vllm/model_executor/layers/fused_moe/oracle/unquantized.py
-2
@@ -269,8 +269,6 @@ def convert_to_unquantized_kernel_format(
269
w13_weight = swap_w13_to_w31(w13_weight)
270
271
elif unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM:
272
- # Swap halves to arrange as [w3; w1] (kernel expectation)
273
- w13_weight = swap_w13_to_w31(w13_weight)
272
_cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
273
w13_weight, w2_weight = convert_moe_weights_to_flashinfer_trtllm_block_layout(
274
_cache_permute_indices,
vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+43
-30
@@ -121,7 +121,6 @@ def convert_moe_weights_to_flashinfer_trtllm_block_layout(
121
122
from flashinfer.fused_moe.core import (
123
_maybe_get_cached_w3_w1_permute_indices,
124
- convert_to_block_layout,
124
get_w2_permute_indices_with_cache,
125
)
126
@@ -131,23 +130,49 @@ def convert_moe_weights_to_flashinfer_trtllm_block_layout(
130
# Reorder rows of W13 and W2 for fused gated activation and convert to the
131
# block layout expected by the FlashInfer kernel.
132
num_experts = w13_weight.shape[0]
134
- device_w13 = w13_weight.device
135
- device_w2 = w2_weight.device
133
137
- w13_weights_shuffled: list[torch.Tensor] = []
138
- w2_weights_shuffled: list[torch.Tensor] = []
134
+ def _copy_permuted_expert_to_block_layout(
135
+ out: torch.Tensor,
136
+ expert_uint8: torch.Tensor,
137
+ source_indices: torch.Tensor,
138
+ ) -> None:
139
+ expert_blocks = expert_uint8.view(
140
+ expert_uint8.shape[0], out.shape[0], block_k
141
+ ).permute(1, 0, 2)
142
+ torch.index_select(
143
+ expert_blocks,
144
+ 1,
145
+ source_indices.to(expert_uint8.device),
146
+ out=out,
147
+ )
148
+
149
+ w13_rows, w13_cols = w13_weight[0].view(torch.uint8).shape
150
+ w2_rows, w2_cols = w2_weight[0].view(torch.uint8).shape
151
+ w13_weights_shuffled_tensor = torch.empty(
152
+ (num_experts, w13_cols // block_k, w13_rows, block_k),
153
+ dtype=torch.uint8,
154
+ device=w13_weight.device,
155
+ )
156
+ w2_weights_shuffled_tensor = torch.empty(
157
+ (num_experts, w2_cols // block_k, w2_rows, block_k),
158
+ dtype=torch.uint8,
159
+ device=w2_weight.device,
160
+ )
161
162
for i in range(num_experts):
163
+ w13_expert_uint8 = w13_weight[i].view(torch.uint8)
164
+
165
permute_indices = _maybe_get_cached_w3_w1_permute_indices(
166
cache_permute_indices,
143
- w13_weight[i].view(torch.uint8),
167
+ w13_expert_uint8,
168
epilogue_tile_m,
169
)
146
- tmp_weights1 = (
147
- w13_weight[i]
148
- .clone()
149
- .view(torch.uint8)[permute_indices.to(device_w13)]
150
- .contiguous()
170
+ rows = w13_expert_uint8.shape[0]
171
+ permute_indices = (permute_indices + rows // 2) % rows
172
+ _copy_permuted_expert_to_block_layout(
173
+ w13_weights_shuffled_tensor[i],
174
+ w13_expert_uint8,
175
+ permute_indices,
176
)
177
178
permute_indices = get_w2_permute_indices_with_cache(
@@ -155,29 +180,17 @@ def convert_moe_weights_to_flashinfer_trtllm_block_layout(
180
w2_weight[i].view(torch.uint8),
181
epilogue_tile_m,
182
)
158
- tmp_weights2 = (
159
- w2_weight[i]
160
- .clone()
161
- .view(torch.uint8)[permute_indices.to(device_w2)]
162
- .contiguous()
183
+ _copy_permuted_expert_to_block_layout(
184
+ w2_weights_shuffled_tensor[i],
185
+ w2_weight[i].view(torch.uint8),
186
+ permute_indices,
187
)
188
165
- tmp_weights1 = convert_to_block_layout(tmp_weights1.view(torch.uint8), block_k)
166
- tmp_weights2 = convert_to_block_layout(tmp_weights2.view(torch.uint8), block_k)
167
-
168
- w13_weights_shuffled.append(tmp_weights1.view(torch.bfloat16))
169
- w2_weights_shuffled.append(tmp_weights2.view(torch.bfloat16))
170
-
171
- # Stack weights for all experts and return as BF16 tensors.
172
- w13_weights_shuffled_tensor = (
173
- torch.stack(w13_weights_shuffled).view(torch.bfloat16).contiguous()
174
- )
175
- w2_weights_shuffled_tensor = (
176
- torch.stack(w2_weights_shuffled).view(torch.bfloat16).contiguous()
189
+ return (
190
+ w13_weights_shuffled_tensor.view(torch.bfloat16),
191
+ w2_weights_shuffled_tensor.view(torch.bfloat16),
192
)
193
179
- return w13_weights_shuffled_tensor, w2_weights_shuffled_tensor
180
-
194
195
def align_fp4_moe_weights_for_fi(
196
w13: torch.Tensor,