@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 160 lines 6.4 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from types import SimpleNamespace
5
6 import torch
7
8 from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import NvFp4MoeBackend
9 from vllm.model_executor.layers.quantization.utils import flashinfer_fp4_moe
10 from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
11 prepare_nvfp4_moe_layer_for_fi_or_cutlass,
12 )
13 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
14 align_fp4_moe_hidden_dim_for_fi,
15 align_fp4_moe_weights_for_fi,
16 )
17
18
19 def test_shared_nvfp4_input_scales_have_writable_storage(monkeypatch):
20 monkeypatch.setattr(flashinfer_fp4_moe, "swizzle_blockscale", lambda x: x)
21
22 num_experts = 3
23 layer = SimpleNamespace(
24 activation=SimpleNamespace(is_gated=False),
25 moe_config=SimpleNamespace(
26 moe_parallel_config=SimpleNamespace(enable_eplb=False)
27 ),
28 )
29 w13 = torch.zeros((num_experts, 2, 1), dtype=torch.uint8)
30 w2 = torch.zeros((num_experts, 2, 1), dtype=torch.uint8)
31 w13_scale = torch.zeros((num_experts, 2, 1), dtype=torch.float8_e4m3fn)
32 w2_scale = torch.zeros((num_experts, 2, 1), dtype=torch.float8_e4m3fn)
33 weight_scale = torch.ones(num_experts)
34
35 outputs = prepare_nvfp4_moe_layer_for_fi_or_cutlass(
36 backend=NvFp4MoeBackend.FLASHINFER_CUTLASS,
37 layer=layer,
38 w13=w13,
39 w13_scale=w13_scale,
40 w13_scale_2=weight_scale,
41 a13_scale=torch.tensor([1.0, 2.0, 3.0]),
42 w2=w2,
43 w2_scale=w2_scale,
44 w2_scale_2=weight_scale,
45 a2_scale=torch.tensor([4.0, 5.0, 6.0]),
46 is_act_and_mul=False,
47 )
48 a13_scale, a2_scale = outputs[3], outputs[7]
49
50 torch.testing.assert_close(a13_scale, torch.full((num_experts,), 3.0))
51 torch.testing.assert_close(a2_scale, torch.full((num_experts,), 6.0))
52 distinct_values = torch.arange(num_experts, dtype=torch.float32)
53 a13_scale.copy_(distinct_values)
54 a2_scale.copy_(distinct_values)
55 torch.testing.assert_close(a13_scale, distinct_values)
56 torch.testing.assert_close(a2_scale, distinct_values)
57
58
59 def test_align_trtllm_fp4_moe_hidden_dim_noop():
60 w13 = torch.arange(2 * 8 * 256, dtype=torch.uint8).reshape(2, 8, 256)
61 w13_scale = torch.arange(2 * 8 * 32, dtype=torch.uint8).reshape(2, 8, 32)
62 w2 = torch.arange(2 * 512 * 4, dtype=torch.uint8).reshape(2, 512, 4)
63 w2_scale = torch.arange(2 * 512 * 1, dtype=torch.uint8).reshape(2, 512, 1)
64
65 out_w13, out_w13_scale, out_w2, out_w2_scale, padded_hidden = (
66 align_fp4_moe_hidden_dim_for_fi(w13, w13_scale, w2, w2_scale)
67 )
68
69 assert padded_hidden == 512
70 assert out_w13 is w13
71 assert out_w13_scale is w13_scale
72 assert out_w2 is w2
73 assert out_w2_scale is w2_scale
74
75
76 def test_align_trtllm_fp4_moe_hidden_dim_pads_to_256_multiple():
77 hidden_dim = 2688
78 padded_hidden_dim = 2816
79
80 w13 = torch.arange(2 * 12 * (hidden_dim // 2), dtype=torch.uint8).reshape(
81 2, 12, hidden_dim // 2
82 )
83 w13_scale = torch.arange(2 * 12 * (hidden_dim // 16), dtype=torch.uint8).reshape(
84 2, 12, hidden_dim // 16
85 )
86
87 w2 = torch.arange(2 * hidden_dim * 6, dtype=torch.uint8).reshape(2, hidden_dim, 6)
88 w2_scale = torch.arange(2 * hidden_dim * 2, dtype=torch.uint8).reshape(
89 2, hidden_dim, 2
90 )
91
92 out_w13, out_w13_scale, out_w2, out_w2_scale, out_hidden_dim = (
93 align_fp4_moe_hidden_dim_for_fi(w13, w13_scale, w2, w2_scale)
94 )
95
96 assert out_hidden_dim == padded_hidden_dim
97 assert out_w13.shape == (2, 12, padded_hidden_dim // 2)
98 assert out_w13_scale.shape == (2, 12, padded_hidden_dim // 16)
99 assert out_w2.shape == (2, padded_hidden_dim, 6)
100 assert out_w2_scale.shape == (2, padded_hidden_dim, 2)
101
102 torch.testing.assert_close(out_w13[:, :, : hidden_dim // 2], w13)
103 torch.testing.assert_close(out_w13_scale[:, :, : hidden_dim // 16], w13_scale)
104 torch.testing.assert_close(out_w2[:, :hidden_dim, :], w2)
105 torch.testing.assert_close(out_w2_scale[:, :hidden_dim, :], w2_scale)
106
107 assert torch.count_nonzero(out_w13[:, :, hidden_dim // 2 :]) == 0
108 assert torch.count_nonzero(out_w13_scale[:, :, hidden_dim // 16 :]) == 0
109 assert torch.count_nonzero(out_w2[:, hidden_dim:, :]) == 0
110 assert torch.count_nonzero(out_w2_scale[:, hidden_dim:, :]) == 0
111
112
113 def test_align_trtllm_fp4_moe_intermediate_pads_gate_and_up_separately():
114 # Qwen3.8-Flash-Next has intermediate_size=640. At TP4 each rank owns
115 # 160 values, which TRTLLM pads to the next 64-value boundary (192).
116 intermediate = 160
117 padded_intermediate = 192
118 hidden_dim = 32
119
120 gate = torch.ones((1, intermediate, hidden_dim // 2), dtype=torch.uint8)
121 up = torch.full_like(gate, 2)
122 w13 = torch.cat((gate, up), dim=1)
123 gate_scale = torch.full((1, intermediate, hidden_dim // 16), 3, dtype=torch.uint8)
124 up_scale = torch.full_like(gate_scale, 4)
125 w13_scale = torch.cat((gate_scale, up_scale), dim=1)
126 w2 = torch.full((1, hidden_dim, intermediate // 2), 5, dtype=torch.uint8)
127 w2_scale = torch.full((1, hidden_dim, intermediate // 16), 6, dtype=torch.uint8)
128
129 out_w13, out_w13_scale, out_w2, out_w2_scale, out_intermediate = (
130 align_fp4_moe_weights_for_fi(
131 w13,
132 w13_scale,
133 w2,
134 w2_scale,
135 is_act_and_mul=True,
136 min_alignment=64,
137 )
138 )
139
140 assert out_intermediate == padded_intermediate
141 torch.testing.assert_close(out_w13[:, :intermediate], gate)
142 torch.testing.assert_close(
143 out_w13[:, padded_intermediate : padded_intermediate + intermediate], up
144 )
145 torch.testing.assert_close(out_w13_scale[:, :intermediate], gate_scale)
146 torch.testing.assert_close(
147 out_w13_scale[:, padded_intermediate : padded_intermediate + intermediate],
148 up_scale,
149 )
150
151 assert torch.count_nonzero(out_w13[:, intermediate:padded_intermediate]) == 0
152 assert torch.count_nonzero(out_w13[:, padded_intermediate + intermediate :]) == 0
153 assert torch.count_nonzero(out_w13_scale[:, intermediate:padded_intermediate]) == 0
154 assert (
155 torch.count_nonzero(out_w13_scale[:, padded_intermediate + intermediate :]) == 0
156 )
157 torch.testing.assert_close(out_w2[:, :, : intermediate // 2], w2)
158 torch.testing.assert_close(out_w2_scale[:, :, : intermediate // 16], w2_scale)
159 assert torch.count_nonzero(out_w2[:, :, intermediate // 2 :]) == 0
160 assert torch.count_nonzero(out_w2_scale[:, :, intermediate // 16 :]) == 0