[Bugfix] Fix Kimi K3 NVFP4 MoE weight conversion OOM (#55407)
Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>
Wei Zhao committed
Sep 6, 2026 at 23:45 UTC
6748217fb949952b4db418d09e64952cd7d8e352
1 file changed
+71
-83
vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
+71
-83
@@ -190,25 +190,23 @@ def nvfp4_swizzled_scale_to_cutedsl_mma_view(scale: torch.Tensor) -> torch.Tenso
190
191
192
def prepare_static_weights_for_trtllm_fp4_moe(
193
- # args_dequant,
194
- # args,
195
- gemm1_weights,
196
- gemm2_weights,
197
- gemm1_scales_linear_fp4_bytes,
198
- gemm2_scales_linear_fp4_bytes,
199
- hidden_size,
200
- intermediate_size,
201
- num_experts,
193
+ gemm1_weights: torch.Tensor,
194
+ gemm2_weights: torch.Tensor,
195
+ gemm1_scales_linear_fp4_bytes: torch.Tensor,
196
+ gemm2_scales_linear_fp4_bytes: torch.Tensor,
197
+ hidden_size: int,
198
+ intermediate_size: int,
199
+ num_experts: int,
200
is_gated_activation: bool,
203
-):
201
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
202
+ """Shuffle NVFP4 weights into the FlashInfer TRT-LLM layout."""
203
from flashinfer import nvfp4_block_scale_interleave
204
from flashinfer.fused_moe.core import (
205
_maybe_get_cached_w3_w1_permute_indices,
206
get_w2_permute_indices_with_cache,
207
)
208
210
- _cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
211
- """Prepare quantized weights for kernel (done offline with weights)."""
209
+ permute_indices_cache: dict[tuple[object, ...], torch.Tensor] = {}
210
epilogue_tile_m = 128 # FIXME: this depends on the kernel internals
211
gemm1_intermediate_size = (
212
2 * intermediate_size if is_gated_activation else intermediate_size
@@ -231,85 +229,75 @@ def prepare_static_weights_for_trtllm_fp4_moe(
229
torch.float8_e4m3fn
230
).reshape(num_experts, hidden_size, intermediate_size // 16) # fp8 scaling factors
231
234
- gemm1_weights_fp4_shuffled = []
235
- gemm1_scales_fp4_shuffled = []
236
- gemm2_weights_fp4_shuffled = []
237
- gemm2_scales_fp4_shuffled = []
238
- for i in range(num_experts):
239
- # Calculate the permute indices for the following:
240
- # 1. Reorder rows of W1 and scales for fused gated activation
241
- # 2. Shuffle weights and scaling factors for transposed mma output
242
- # for both w3_w1 and w2 weights and scale factors
243
- permute_indices = _maybe_get_cached_w3_w1_permute_indices(
244
- _cache_permute_indices,
245
- gemm1_weights_fp4[i].view(torch.uint8),
246
- epilogue_tile_m,
247
- is_gated_act_gemm=is_gated_activation,
248
- )
249
- gemm1_weights_fp4_shuffled.append(
250
- gemm1_weights_fp4[i]
251
- .view(torch.uint8)[permute_indices.to(gemm1_weights_fp4.device)]
252
- .contiguous()
253
- )
232
+ gemm1_weight_indices = _maybe_get_cached_w3_w1_permute_indices(
233
+ permute_indices_cache,
234
+ gemm1_weights_fp4[0].view(torch.uint8),
235
+ epilogue_tile_m,
236
+ is_gated_act_gemm=is_gated_activation,
237
+ )
238
+ gemm1_scale_indices = _maybe_get_cached_w3_w1_permute_indices(
239
+ permute_indices_cache,
240
+ gemm1_scales_linear_fp4[0].view(torch.uint8),
241
+ epilogue_tile_m,
242
+ num_elts_per_sf=16,
243
+ is_gated_act_gemm=is_gated_activation,
244
+ )
245
+ gemm2_weight_indices = get_w2_permute_indices_with_cache(
246
+ permute_indices_cache,
247
+ gemm2_weights_fp4[0].view(torch.uint8),
248
+ epilogue_tile_m,
249
+ )
250
+ gemm2_scale_indices = get_w2_permute_indices_with_cache(
251
+ permute_indices_cache,
252
+ gemm2_scales_linear_fp4[0].view(torch.uint8),
253
+ epilogue_tile_m,
254
+ num_elts_per_sf=16,
255
+ )
256
255
- permute_sf_indices = _maybe_get_cached_w3_w1_permute_indices(
256
- _cache_permute_indices,
257
- gemm1_scales_linear_fp4[i].view(torch.uint8),
258
- epilogue_tile_m,
259
- num_elts_per_sf=16,
260
- is_gated_act_gemm=is_gated_activation,
261
- )
262
- gemm1_scales_fp4_shuffled.append(
263
- nvfp4_block_scale_interleave(
264
- gemm1_scales_linear_fp4[i]
265
- .view(torch.uint8)[
266
- permute_sf_indices.to(gemm1_scales_linear_fp4.device)
267
- ]
268
- .contiguous()
269
- )
270
- )
257
+ gemm1_weights_fp4_shuffled = torch.empty_like(gemm1_weights_fp4, dtype=torch.uint8)
258
+ gemm1_scales_fp4_shuffled = torch.empty_like(gemm1_scales_linear_fp4)
259
+ gemm2_weights_fp4_shuffled = torch.empty_like(gemm2_weights_fp4, dtype=torch.uint8)
260
+ gemm2_scales_fp4_shuffled = torch.empty_like(gemm2_scales_linear_fp4)
261
+ gemm1_scale_scratch = torch.empty_like(
262
+ gemm1_scales_linear_fp4[0], dtype=torch.uint8
263
+ )
264
+ gemm2_scale_scratch = torch.empty_like(
265
+ gemm2_scales_linear_fp4[0], dtype=torch.uint8
266
+ )
267
272
- permute_indices = get_w2_permute_indices_with_cache(
273
- _cache_permute_indices,
274
- gemm2_weights_fp4[i].view(torch.uint8),
275
- epilogue_tile_m,
268
+ for expert_id in range(num_experts):
269
+ torch.index_select(
270
+ gemm1_weights_fp4[expert_id].view(torch.uint8),
271
+ 0,
272
+ gemm1_weight_indices,
273
+ out=gemm1_weights_fp4_shuffled[expert_id],
274
)
277
- gemm2_weights_fp4_shuffled.append(
278
- gemm2_weights_fp4[i]
279
- .view(torch.uint8)[permute_indices.to(gemm2_weights_fp4.device)]
280
- .contiguous()
275
+ torch.index_select(
276
+ gemm1_scales_linear_fp4[expert_id].view(torch.uint8),
277
+ 0,
278
+ gemm1_scale_indices,
279
+ out=gemm1_scale_scratch,
280
+ )
281
+ gemm1_scales_fp4_shuffled[expert_id].view(torch.uint8).reshape(-1).copy_(
282
+ nvfp4_block_scale_interleave(gemm1_scale_scratch)
283
)
284
283
- permute_sf_indices = get_w2_permute_indices_with_cache(
284
- _cache_permute_indices,
285
- gemm2_scales_linear_fp4[i].view(torch.uint8),
286
- epilogue_tile_m,
287
- num_elts_per_sf=16,
285
+ torch.index_select(
286
+ gemm2_weights_fp4[expert_id].view(torch.uint8),
287
+ 0,
288
+ gemm2_weight_indices,
289
+ out=gemm2_weights_fp4_shuffled[expert_id],
290
)
289
- gemm2_scales_fp4_shuffled.append(
290
- nvfp4_block_scale_interleave(
291
- gemm2_scales_linear_fp4[i]
292
- .view(torch.uint8)[
293
- permute_sf_indices.to(gemm2_scales_linear_fp4.device)
294
- ]
295
- .contiguous()
296
- )
291
+ torch.index_select(
292
+ gemm2_scales_linear_fp4[expert_id].view(torch.uint8),
293
+ 0,
294
+ gemm2_scale_indices,
295
+ out=gemm2_scale_scratch,
296
+ )
297
+ gemm2_scales_fp4_shuffled[expert_id].view(torch.uint8).reshape(-1).copy_(
298
+ nvfp4_block_scale_interleave(gemm2_scale_scratch)
299
)
300
299
- # Stack weights for all experts
300
- gemm1_weights_fp4_shuffled = torch.stack(gemm1_weights_fp4_shuffled)
301
- gemm1_scales_fp4_shuffled = (
302
- torch.stack(gemm1_scales_fp4_shuffled)
303
- .view(torch.float8_e4m3fn)
304
- .reshape(num_experts, gemm1_intermediate_size, hidden_size // 16)
305
- )
306
-
307
- gemm2_weights_fp4_shuffled = torch.stack(gemm2_weights_fp4_shuffled)
308
- gemm2_scales_fp4_shuffled = (
309
- torch.stack(gemm2_scales_fp4_shuffled)
310
- .view(torch.float8_e4m3fn)
311
- .reshape(num_experts, hidden_size, intermediate_size // 16)
312
- )
301
return (
302
gemm1_weights_fp4_shuffled,
303
gemm1_scales_fp4_shuffled,