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, 04:45 UTC
[Kernel] Support TRTLLM GEN NVFP4 MoE for non-512-aligned hidden dims via weight padding (#39510)
Signed-off-by: root <root@lyris0017.lyris.clusters.nvidia.com> Signed-off-by: Daniel Afrimi <dafrimi@nvidia.com> Co-authored-by: root <root@lyris0017.lyris.clusters.nvidia.com>
danielafrimi committed
Apr 14, 2026 at 21:49 UTC
240f2636ca4954941dd314ba961023100f804307
4 files changed
+124
-3
tests/quantization/test_trtllm_nvfp4_hidden_dim_padding.py
new
+62
@@ -0,0 +1,62 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+
6
+from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
7
+ align_trtllm_fp4_moe_hidden_dim_for_fi,
8
+)
9
+
10
+
11
+def test_align_trtllm_fp4_moe_hidden_dim_noop():
12
+ w13 = torch.arange(2 * 8 * 256, dtype=torch.uint8).reshape(2, 8, 256)
13
+ w13_scale = torch.arange(2 * 8 * 32, dtype=torch.uint8).reshape(2, 8, 32)
14
+ w2 = torch.arange(2 * 512 * 4, dtype=torch.uint8).reshape(2, 512, 4)
15
+ w2_scale = torch.arange(2 * 512 * 1, dtype=torch.uint8).reshape(2, 512, 1)
16
+
17
+ out_w13, out_w13_scale, out_w2, out_w2_scale, padded_hidden = (
18
+ align_trtllm_fp4_moe_hidden_dim_for_fi(w13, w13_scale, w2, w2_scale)
19
+ )
20
+
21
+ assert padded_hidden == 512
22
+ assert out_w13 is w13
23
+ assert out_w13_scale is w13_scale
24
+ assert out_w2 is w2
25
+ assert out_w2_scale is w2_scale
26
+
27
+
28
+def test_align_trtllm_fp4_moe_hidden_dim_pads_to_256_multiple():
29
+ hidden_dim = 2688
30
+ padded_hidden_dim = 2816
31
+
32
+ w13 = torch.arange(2 * 12 * (hidden_dim // 2), dtype=torch.uint8).reshape(
33
+ 2, 12, hidden_dim // 2
34
+ )
35
+ w13_scale = torch.arange(2 * 12 * (hidden_dim // 16), dtype=torch.uint8).reshape(
36
+ 2, 12, hidden_dim // 16
37
+ )
38
+
39
+ w2 = torch.arange(2 * hidden_dim * 6, dtype=torch.uint8).reshape(2, hidden_dim, 6)
40
+ w2_scale = torch.arange(2 * hidden_dim * 2, dtype=torch.uint8).reshape(
41
+ 2, hidden_dim, 2
42
+ )
43
+
44
+ out_w13, out_w13_scale, out_w2, out_w2_scale, out_hidden_dim = (
45
+ align_trtllm_fp4_moe_hidden_dim_for_fi(w13, w13_scale, w2, w2_scale)
46
+ )
47
+
48
+ assert out_hidden_dim == padded_hidden_dim
49
+ assert out_w13.shape == (2, 12, padded_hidden_dim // 2)
50
+ assert out_w13_scale.shape == (2, 12, padded_hidden_dim // 16)
51
+ assert out_w2.shape == (2, padded_hidden_dim, 6)
52
+ assert out_w2_scale.shape == (2, padded_hidden_dim, 2)
53
+
54
+ torch.testing.assert_close(out_w13[:, :, : hidden_dim // 2], w13)
55
+ torch.testing.assert_close(out_w13_scale[:, :, : hidden_dim // 16], w13_scale)
56
+ torch.testing.assert_close(out_w2[:, :hidden_dim, :], w2)
57
+ torch.testing.assert_close(out_w2_scale[:, :hidden_dim, :], w2_scale)
58
+
59
+ assert torch.count_nonzero(out_w13[:, :, hidden_dim // 2 :]) == 0
60
+ assert torch.count_nonzero(out_w13_scale[:, :, hidden_dim // 16 :]) == 0
61
+ assert torch.count_nonzero(out_w2[:, hidden_dim:, :]) == 0
62
+ assert torch.count_nonzero(out_w2_scale[:, hidden_dim:, :]) == 0
vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
+12
-3
@@ -50,6 +50,9 @@ class TrtLlmNvFp4ExpertsBase:
50
moe_config.intermediate_size_per_partition
51
)
52
self.hidden_dim = moe_config.hidden_dim
53
+ self.hidden_dim_unpadded = (
54
+ moe_config.hidden_dim_unpadded or moe_config.hidden_dim
55
+ )
56
self.local_num_experts = moe_config.num_local_experts
57
self.ep_rank = moe_config.moe_parallel_config.ep_rank
58
@@ -114,8 +117,12 @@ class TrtLlmNvFp4ExpertsBase:
117
118
@staticmethod
119
def _supports_shape(hidden_dim: int) -> bool:
117
- """Requires hidden dim to be multiple of 512."""
118
- return hidden_dim % 512 == 0
120
+ # Weights are zero-padded to 256-alignment at load time and the MoE
121
+ # runner pads activations via _maybe_pad_hidden_states, so any
122
+ # hidden_dim is accepted.
123
+ # NOTE: non-256-aligned dims will trigger a warning log and may
124
+ # cause performance degradation due to activation slicing.
125
+ return True
126
127
@staticmethod
128
def activation_format() -> mk.FusedMoEActivationFormat:
@@ -194,7 +201,7 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
201
import vllm.utils.flashinfer as fi_utils
202
203
if fi_utils._is_fi_autotuning:
197
- return hidden_states
204
+ return
205
206
# Invoke kernel.
207
flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe(
@@ -324,6 +331,8 @@ class TrtLlmNvFp4ExpertsMonolithic(
331
e_score_correction_bias = e_score_correction_bias.to(torch.bfloat16)
332
333
# Invoke kernel.
334
+ # NOTE: Activation padding and output
335
+ # truncation are handled by the MoE runner's
336
return flashinfer.fused_moe.trtllm_fp4_block_scale_moe(
337
routing_logits=router_logits,
338
routing_bias=e_score_correction_bias,
vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
+8
@@ -10,6 +10,7 @@ import vllm.envs as envs
10
from vllm.logger import init_logger
11
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
12
align_fp4_moe_weights_for_fi,
13
+ align_trtllm_fp4_moe_hidden_dim_for_fi,
14
)
15
from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
16
swizzle_blockscale,
@@ -341,6 +342,13 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass(
342
343
# Shuffle weights and scales for FI TRTLLM NVFP4 MoE kernels.
344
if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM:
345
+ w13, w13_scale, w2, w2_scale, padded_hidden = (
346
+ align_trtllm_fp4_moe_hidden_dim_for_fi(w13, w13_scale, w2, w2_scale)
347
+ )
348
+ if layer.moe_config.hidden_dim_unpadded is None:
349
+ layer.moe_config.hidden_dim_unpadded = layer.moe_config.hidden_dim
350
+ layer.moe_config.hidden_dim = padded_hidden
351
+
352
# Align weights for FI NVFP4 MoE kernels.
353
min_alignment = 16 if is_gated else 128
354
w13, w13_scale, w2, w2_scale, padded_intermediate = (
vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+42
@@ -265,6 +265,48 @@ def align_fp4_moe_weights_for_fi(
265
return padded_w13, padded_w13_scale, padded_w2, padded_w2_scale, padded_intermediate
266
267
268
+def align_trtllm_fp4_moe_hidden_dim_for_fi(
269
+ w13: torch.Tensor,
270
+ w13_scale: torch.Tensor,
271
+ w2: torch.Tensor,
272
+ w2_scale: torch.Tensor,
273
+ min_alignment: int = 256,
274
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]:
275
+ num_experts, gate_up_dim, packed_hidden_size = w13.shape
276
+ hidden_size = packed_hidden_size * 2
277
+ padded_hidden_size = round_up(hidden_size, min_alignment)
278
+
279
+ if padded_hidden_size == hidden_size:
280
+ return w13, w13_scale, w2, w2_scale, hidden_size
281
+
282
+ logger.warning_once(
283
+ "Padding hidden size from %d to %d for TRTLLM NVFP4 MoE weights. "
284
+ "This requires activation slicing at runtime and may cause "
285
+ "performance degradation.",
286
+ hidden_size,
287
+ padded_hidden_size,
288
+ scope="local",
289
+ )
290
+
291
+ padded_w13 = w13.new_zeros((num_experts, gate_up_dim, padded_hidden_size // 2))
292
+ padded_w13[:, :, :packed_hidden_size] = w13
293
+
294
+ padded_w13_scale = w13_scale.new_zeros(
295
+ (num_experts, gate_up_dim, padded_hidden_size // 16)
296
+ )
297
+ padded_w13_scale[:, :, : w13_scale.shape[2]] = w13_scale
298
+
299
+ padded_w2 = w2.new_zeros((num_experts, padded_hidden_size, w2.shape[2]))
300
+ padded_w2[:, : w2.shape[1], :] = w2
301
+
302
+ padded_w2_scale = w2_scale.new_zeros(
303
+ (num_experts, padded_hidden_size, w2_scale.shape[2])
304
+ )
305
+ padded_w2_scale[:, : w2_scale.shape[1], :] = w2_scale
306
+
307
+ return padded_w13, padded_w13_scale, padded_w2, padded_w2_scale, padded_hidden_size
308
+
309
+
310
def align_fp8_moe_weights_for_fi(
311
w13: torch.Tensor, w2: torch.Tensor, is_act_and_mul: bool, min_alignment: int = 16
312
) -> tuple[torch.Tensor, torch.Tensor, int]: