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 7, 22:45 UTC
[Model] Add NVFP4 quantization support for Step3.5-Flash (#34478)
Signed-off-by: tacos8me <ian@cloudhabit.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
tacos8me committed
Feb 22, 2026 at 14:30 UTC
b7892a3beff05971f7e1ed3519aec96c3d89bfb0
5 files changed
+204
-4
tests/kernels/moe/test_nvfp4_moe.py
+126
@@ -14,6 +14,7 @@ from tests.kernels.utils import torch_moe
14
from vllm import _custom_ops as ops
15
from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
16
from vllm.model_executor.layers.fused_moe import fused_topk
17
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
18
from vllm.model_executor.layers.fused_moe.config import nvfp4_moe_quant_config
19
from vllm.model_executor.layers.fused_moe.cutlass_moe import (
20
CutlassExpertsFp4,
@@ -147,5 +148,130 @@ def test_cutlass_fp4_moe_no_graph(
148
torch.testing.assert_close(torch_output, cutlass_output, atol=1e-1, rtol=1e-1)
149
150
151
+# step3.5-flash uses swiglustep activation (clipped SwiGLU with limit=7.0)
152
+# for MoE layers 43-44. This tests the non-fused activation fallback path
153
+# in run_cutlass_moe_fp4 (apply_moe_activation + separate fp4 quantization).
154
+# Model dims: e=288, topk=8, n=1280 (moe_intermediate_size), k=4096 (hidden)
155
+SWIGLUSTEP_MNK_FACTORS = [
156
+ (2, 1280, 4096),
157
+ (64, 1280, 4096),
158
+ (224, 1280, 4096),
159
+]
160
+
161
+
162
+@pytest.mark.parametrize("m,n,k", SWIGLUSTEP_MNK_FACTORS)
163
+@pytest.mark.parametrize("e", [64, 288])
164
+@pytest.mark.parametrize("topk", [1, 8])
165
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
166
+@torch.inference_mode()
167
+def test_cutlass_fp4_moe_swiglustep(
168
+ m: int, n: int, k: int, e: int, topk: int, dtype: torch.dtype, workspace_init
169
+):
170
+ set_random_seed(7)
171
+ with set_current_vllm_config(
172
+ VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
173
+ ):
174
+ quant_blocksize = 16
175
+
176
+ a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
177
+
178
+ (_, w1_q, w1_blockscale, w1_gs), (_, w2_q, w2_blockscale, w2_gs) = (
179
+ make_test_weights(
180
+ e,
181
+ n,
182
+ k,
183
+ in_dtype=dtype,
184
+ quant_dtype="nvfp4",
185
+ block_shape=None,
186
+ per_out_ch_quant=False,
187
+ )
188
+ )
189
+
190
+ score = torch.randn((m, e), device="cuda", dtype=dtype)
191
+ topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False)
192
+
193
+ a1_gs = torch.ones((e,), device="cuda", dtype=torch.float32)
194
+ a2_gs = torch.ones((e,), device="cuda", dtype=torch.float32)
195
+
196
+ assert w1_gs is not None
197
+ assert w2_gs is not None
198
+ assert w1_blockscale is not None
199
+ assert w2_blockscale is not None
200
+
201
+ quant_config = nvfp4_moe_quant_config(
202
+ g1_alphas=(1 / w1_gs),
203
+ g2_alphas=(1 / w2_gs),
204
+ a1_gscale=a1_gs,
205
+ a2_gscale=a2_gs,
206
+ w1_scale=w1_blockscale,
207
+ w2_scale=w2_blockscale,
208
+ )
209
+
210
+ kernel = mk.FusedMoEModularKernel(
211
+ MoEPrepareAndFinalizeNoEP(),
212
+ CutlassExpertsFp4(
213
+ moe_config=make_dummy_moe_config(),
214
+ quant_config=quant_config,
215
+ ),
216
+ inplace=False,
217
+ )
218
+
219
+ cutlass_output = kernel(
220
+ hidden_states=a,
221
+ w1=w1_q,
222
+ w2=w2_q,
223
+ topk_weights=topk_weights,
224
+ topk_ids=topk_ids,
225
+ activation=MoEActivation.SWIGLUSTEP,
226
+ )
227
+
228
+ # Reference: dequantize everything and run torch_moe with swiglustep
229
+ a_global_scale = (
230
+ (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(a.flatten(), dim=-1)
231
+ ).to(torch.float32)
232
+ a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(a, a_global_scale)
233
+
234
+ a_in_dtype = dequantize_nvfp4_to_dtype(
235
+ a_fp4,
236
+ a_scale_interleaved,
237
+ a_global_scale,
238
+ dtype=a.dtype,
239
+ device=a.device,
240
+ block_size=quant_blocksize,
241
+ )
242
+
243
+ w1_d = torch.empty((e, 2 * n, k), device="cuda", dtype=dtype)
244
+ w2_d = torch.empty((e, k, n), device="cuda", dtype=dtype)
245
+
246
+ for idx in range(0, e):
247
+ w1_d[idx] = dequantize_nvfp4_to_dtype(
248
+ w1_q[idx],
249
+ w1_blockscale[idx],
250
+ w1_gs[idx],
251
+ dtype=dtype,
252
+ device=w1_q.device,
253
+ block_size=quant_blocksize,
254
+ )
255
+ w2_d[idx] = dequantize_nvfp4_to_dtype(
256
+ w2_q[idx],
257
+ w2_blockscale[idx],
258
+ w2_gs[idx],
259
+ dtype=dtype,
260
+ device=w2_q.device,
261
+ block_size=quant_blocksize,
262
+ )
263
+
264
+ torch_output = torch_moe(
265
+ a_in_dtype,
266
+ w1_d,
267
+ w2_d,
268
+ score,
269
+ topk,
270
+ activation=MoEActivation.SWIGLUSTEP,
271
+ )
272
+
273
+ torch.testing.assert_close(torch_output, cutlass_output, atol=1e-1, rtol=1e-1)
274
+
275
+
276
if __name__ == "__main__":
277
test_cutlass_fp4_moe_no_graph((2, 1024, 1024), 40, 1, torch.half)
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+4
@@ -690,10 +690,14 @@ class CutlassExpertsFp4(mk.FusedMoEPermuteExpertsUnpermute):
690
691
@staticmethod
692
def _supports_activation(activation: MoEActivation) -> bool:
693
+ # SILU uses a fused silu+mul+fp4_quant kernel path.
694
+ # Other gated activations use the generic apply_moe_activation()
695
+ # fallback + separate fp4 quantization in run_cutlass_moe_fp4().
696
return activation in [
697
MoEActivation.SILU,
698
MoEActivation.GELU,
699
MoEActivation.SWIGLUOAI,
700
+ MoEActivation.SWIGLUSTEP,
701
]
702
703
@staticmethod
vllm/model_executor/layers/fused_moe/fused_marlin_moe.py
+3
@@ -586,10 +586,13 @@ class MarlinExpertsBase(mk.FusedMoEPermuteExpertsUnpermute):
586
587
@staticmethod
588
def _supports_activation(activation: MoEActivation) -> bool:
589
+ # Marlin uses apply_moe_activation() callback for activation,
590
+ # so any activation supported there can be used here.
591
return activation in [
592
MoEActivation.SILU,
593
MoEActivation.GELU,
594
MoEActivation.SWIGLUOAI,
595
+ MoEActivation.SWIGLUSTEP,
596
MoEActivation.SILU_NO_MUL,
597
MoEActivation.GELU_NO_MUL,
598
MoEActivation.RELU2_NO_MUL,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
-3
@@ -652,9 +652,6 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
652
shared_experts_input: torch.Tensor | None,
653
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
654
assert not self.is_monolithic
655
- assert layer.activation == MoEActivation.SILU, (
656
- f"Only SiLU activation is supported, not {layer.activation}."
657
- )
655
656
# EPLB path
657
if self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM:
vllm/model_executor/models/step3p5.py
+71
-1
@@ -2,7 +2,8 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Inference-only Jurassic model."""
4
5
-from collections.abc import Iterable
5
+import typing
6
+from collections.abc import Callable, Iterable
7
from typing import Any
8
9
import torch
@@ -231,6 +232,7 @@ class Step3p5Attention(nn.Module):
232
hidden_size,
233
self.total_num_heads,
234
bias=False,
235
+ quant_config=quant_config,
236
prefix=f"{prefix}.g_proj",
237
)
238
@@ -640,12 +642,22 @@ class Step3p5Model(nn.Module):
642
params_dict = dict(self.named_parameters())
643
loaded_params: set[str] = set()
644
645
+ # Old packed 3D format: .moe.gate_proj.weight [num_experts, out, in]
646
expert_params_mapping = [
647
(".moe.experts.w13_weight", ".moe.gate_proj.weight", "w1"),
648
(".moe.experts.w13_weight", ".moe.up_proj.weight", "w3"),
649
(".moe.experts.w2_weight", ".moe.down_proj.weight", "w2"),
650
]
651
652
+ # New per-expert format: .moe.experts.E.gate_proj.weight_packed [out, in]
653
+ per_expert_mapping = FusedMoE.make_expert_params_mapping(
654
+ self,
655
+ ckpt_gate_proj_name="gate_proj",
656
+ ckpt_down_proj_name="down_proj",
657
+ ckpt_up_proj_name="up_proj",
658
+ num_experts=self.moe_num_experts,
659
+ )
660
+
661
disable_moe_stacked_params = [data[1] for data in expert_params_mapping]
662
663
for name, loaded_weight in weights:
@@ -668,6 +680,54 @@ class Step3p5Model(nn.Module):
680
if layer_idx >= config.num_hidden_layers:
681
continue
682
683
+ # Per-expert MoE weights (new format from LLM Compressor):
684
+ # .moe.experts.{E}.{gate,up,down}_proj.{weight_packed,scale,...}
685
+ # Each weight is individual per-expert, not stacked 3D.
686
+ if ".moe.experts." in local_name:
687
+ is_expert_weight = False
688
+ for mapping in per_expert_mapping:
689
+ param_name, weight_name, expert_id, shard_id = mapping
690
+ if weight_name not in local_name:
691
+ continue
692
+ is_expert_weight = True
693
+ name_mapped = local_name.replace(weight_name, param_name)
694
+ if is_pp_missing_parameter(name_mapped, self):
695
+ continue
696
+ if name_mapped not in params_dict:
697
+ continue
698
+ param = params_dict[name_mapped]
699
+ weight_loader = typing.cast(
700
+ Callable[..., bool], param.weight_loader
701
+ )
702
+ success = weight_loader(
703
+ param,
704
+ loaded_weight,
705
+ name_mapped,
706
+ shard_id=shard_id,
707
+ expert_id=expert_id,
708
+ return_success=True,
709
+ )
710
+ if success:
711
+ loaded_params.add(name_mapped)
712
+ break
713
+ else:
714
+ if (
715
+ not is_expert_weight
716
+ and not is_pp_missing_parameter(local_name, self)
717
+ and local_name in params_dict
718
+ ):
719
+ # Not an expert proj — use default loader
720
+ # (e.g. share_expert weights if they matched)
721
+ param = params_dict[local_name]
722
+ weight_loader = getattr(
723
+ param,
724
+ "weight_loader",
725
+ default_weight_loader,
726
+ )
727
+ weight_loader(param, loaded_weight)
728
+ loaded_params.add(local_name)
729
+ continue
730
+
731
for param_name, weight_name, shard_id in stacked_params_mapping:
732
if weight_name not in local_name:
733
continue
@@ -703,6 +763,16 @@ class Step3p5Model(nn.Module):
763
param = params_dict[replaced_name]
764
weight_loader = param.weight_loader
765
moe_expert_num = self.moe_num_experts
766
+ # Per-tensor global scales (e.g. weight_global_scale)
767
+ # have shape [1] in compressed-tensors NVFP4 checkpoints.
768
+ # Expand to per-expert before the iteration loop.
769
+ if (
770
+ loaded_weight.shape[0] == 1
771
+ and loaded_weight.shape[0] != moe_expert_num
772
+ ):
773
+ loaded_weight = loaded_weight.expand(
774
+ moe_expert_num, *loaded_weight.shape[1:]
775
+ )
776
assert loaded_weight.shape[0] == moe_expert_num
777
for expert_id in range(moe_expert_num):
778
loaded_weight_expert = loaded_weight[expert_id]