[BugFix] weights processing peak memory reduction for nvfp4 MoE layers (#46276)
Signed-off-by: Jimmy Lee <hirejimmylee@gmail.com>
Jimmy Lee committed
Jul 10, 2026 at 19:05 UTC
092387963c09ff3dde648e4ff620174d5d32c5d3
2 files changed
+36
-29
tests/kernels/moe/test_flashinfer_b12x_moe.py
+12
-21
@@ -41,6 +41,9 @@ from vllm.model_executor.layers.fused_moe.config import nvfp4_moe_quant_config
41
from vllm.model_executor.layers.fused_moe.experts.flashinfer_b12x_moe import (
42
FlashInferB12xExperts,
43
)
44
+from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
45
+ reorder_w1w3_to_w3w1,
46
+)
47
from vllm.utils.torch_utils import set_random_seed
48
49
# Dimensions chosen to satisfy FP4 alignment requirements (k multiple of 256,
@@ -53,23 +56,6 @@ MNK_FACTORS = [
56
]
57
58
56
-def _reorder_gate_up_to_up_gate(
57
- w: torch.Tensor,
58
- w_s: torch.Tensor,
59
-) -> tuple[torch.Tensor, torch.Tensor]:
60
- """Swap gate and up-projection halves along dim=1 to [up, gate] order.
61
-
62
- The B12x kernel expects weights in [up (w3), gate (w1)] order while the
63
- BF16 reference uses [gate (w1), up (w3)]. This replicates the reordering
64
- done at model-load time by ``prepare_nvfp4_moe_layer_for_fi_or_cutlass``.
65
- """
66
- n = w.shape[1] // 2
67
- return (
68
- torch.cat([w[:, n:, :], w[:, :n, :]], dim=1),
69
- torch.cat([w_s[:, n:, :], w_s[:, :n, :]], dim=1),
70
- )
71
-
72
-
59
def _process_b12x_weights(
60
experts: FlashInferB12xExperts,
61
w1_scale: torch.Tensor,
@@ -142,9 +128,14 @@ def test_flashinfer_b12x_moe(
128
sf_vec_size = 16
129
130
# W1: reorder BF16 from [gate, up] → [up, gate], then quantise.
145
- w1_reordered = torch.cat(
146
- [w1_bf16[:, n:, :], w1_bf16[:, :n, :]], dim=1
147
- ) # shape (e, 2n, k), [up, gate]
131
+ # Note: in reorder_w1w3_to_w3w1, "w1" refers to the gate projection
132
+ # and "w3" refers to the up projection.
133
+ # A dummy scale is passed and discarded; real scales come from
134
+ # fp4_quantize after reordering.
135
+ w1_reordered, _ = reorder_w1w3_to_w3w1(
136
+ w1_bf16.clone(),
137
+ torch.ones((e, 2 * n, 1), device="cuda", dtype=torch.float32),
138
+ )
139
w1_flat = w1_reordered.reshape(e * 2 * n, k)
140
w1_q_flat, w1_sf_flat = fp4_quantize(
141
w1_flat,
@@ -190,6 +181,7 @@ def test_flashinfer_b12x_moe(
181
moe_config=moe_config,
182
quant_config=quant_config,
183
)
184
+
185
_process_b12x_weights(
186
experts,
187
w1_blockscale,
@@ -324,7 +316,6 @@ def test_flashinfer_b12x_moe_relu2(
316
use_monolithic=False,
317
),
318
experts,
327
- inplace=False,
319
)
320
321
score = torch.randn((m, e), device="cuda", dtype=dtype)
vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
+24
-8
@@ -23,7 +23,6 @@ if TYPE_CHECKING:
23
24
logger = init_logger(__name__)
25
26
-
26
__all__ = [
27
"reorder_w1w3_to_w3w1",
28
]
@@ -32,18 +31,35 @@ __all__ = [
31
def reorder_w1w3_to_w3w1(
32
weight: torch.Tensor, scale: torch.Tensor, dim: int = -2
33
) -> tuple[torch.Tensor, torch.Tensor]:
35
- """Re-order the concatenated `[w1, w3]` tensors to `[w3, w1]`"""
34
+ """Re-order concatenated `[w1, w3]` tensors to `[w3, w1]` in-place.
35
+
36
+ `weight` and `scale` must be contiguous; they remain contiguous on return.
37
+ """
38
+ assert weight.is_contiguous(), "weight must be contiguous"
39
+ assert scale.is_contiguous(), "scale must be contiguous"
40
size = weight.size(dim)
41
assert size % 2 == 0, f"Expected even size in dim {dim}, got {size}"
42
half = size // 2
43
+ d = dim % weight.dim()
44
40
- w1, w3 = weight.split(half, dim=dim)
41
- s1, s3 = scale.split(half, dim=dim)
42
-
43
- return (
44
- torch.cat([w3, w1], dim=dim).contiguous(),
45
- torch.cat([s3, s1], dim=dim).contiguous(),
45
+ # 64 MB transient cap
46
+ bytes_per_row = max(
47
+ weight.numel() // size * weight.element_size(),
48
+ scale.numel() // size * scale.element_size(),
49
)
50
+ chunk = max(1, min(half, (64 << 20) // max(bytes_per_row, 1)))
51
+
52
+ fa, fb = [slice(None)] * weight.dim(), [slice(None)] * weight.dim()
53
+ for off in range(0, half, chunk):
54
+ end = min(off + chunk, half)
55
+ fa[d], fb[d] = slice(off, end), slice(half + off, half + end)
56
+ a, b = tuple(fa), tuple(fb)
57
+ for t in (weight, scale):
58
+ tmp = t[b].clone()
59
+ t[b] = t[a]
60
+ t[a] = tmp
61
+
62
+ return weight, scale
63
64
65
def interleave_linear_and_gate(