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 5, 12:15 UTC
[Bug] Fix Trtllm Fp8 MoE Weight Shuffle Memory Fragamentation (#39054)
Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>
Wei Zhao committed
Apr 7, 2026 at 08:04 UTC
0be9516ea43df0fcb24bf50021e22768a49d61cf
1 file changed
+12
-9
vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+12
-9
@@ -322,20 +322,23 @@ def _shuffle_deepseek_fp8_moe_weights(
322
block_k = 128
323
num_experts = w13.shape[0]
324
325
- w13_shuffled: list[torch.Tensor] = []
326
- w2_shuffled: list[torch.Tensor] = []
325
+ M13, K13 = w13.shape[1], w13.shape[2]
326
+ M2, K2 = w2.shape[1], w2.shape[2]
327
+ w13_out = torch.empty(
328
+ num_experts, K13 // block_k, M13, block_k, dtype=torch.uint8, device=w13.device
329
+ )
330
+ w2_out = torch.empty(
331
+ num_experts, K2 // block_k, M2, block_k, dtype=torch.uint8, device=w2.device
332
+ )
333
+
334
for i in range(num_experts):
335
t13 = shuffle_matrix_a(w13[i].view(torch.uint8), epilogue_tile_m)
329
- t13 = convert_to_block_layout(t13, block_k)
330
- w13_shuffled.append(t13)
336
+ w13_out[i] = convert_to_block_layout(t13, block_k)
337
338
t2 = shuffle_matrix_a(w2[i].view(torch.uint8), epilogue_tile_m)
333
- t2 = convert_to_block_layout(t2, block_k)
334
- w2_shuffled.append(t2)
339
+ w2_out[i] = convert_to_block_layout(t2, block_k)
340
336
- w13_out = torch.stack(w13_shuffled).view(torch.float8_e4m3fn)
337
- w2_out = torch.stack(w2_shuffled).view(torch.float8_e4m3fn)
338
- return w13_out, w2_out
341
+ return w13_out.view(torch.float8_e4m3fn), w2_out.view(torch.float8_e4m3fn)
342
343
344
def _shuffle_mxfp8_moe_weights(