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 2, 03:30 UTC
[MoE/b12x] Accept W4A16 (kNvfp4Static, None) in FlashInferB12xExperts supports check (#43332)
Signed-off-by: Junhao Shen <junshen@nvidia.com> Co-authored-by: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com>
Junhao Shen committed
Jun 3, 2026 at 06:20 UTC
a4ac746405f4ddbef553098507210c072b5ba39e
1 file changed
+27
-4
vllm/model_executor/layers/fused_moe/experts/flashinfer_b12x_moe.py
+27
-4
@@ -54,6 +54,11 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular):
54
self.out_dtype = moe_config.in_dtype
55
self.num_local_experts = moe_config.num_local_experts
56
self.ep_rank = moe_config.moe_parallel_config.ep_rank
57
+ # FC2 input scale tensor bound in process_weights_after_loading: the
58
+ # calibrated (now-zeroed) a2_gscale for static-quant checkpoints, or
59
+ # a synthesized uniform-1.0 tensor for W4A16 checkpoints that lack
60
+ # one. Holding it on the instance keeps apply() alloc-free.
61
+ self._fc2_input_scale: torch.Tensor | None = None
62
63
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
64
# Normalise block scales to absorb the per-expert weight global scale
@@ -86,6 +91,18 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular):
91
# its own per-block dynamic scale.
92
if self.a2_gscale is not None:
93
self.a2_gscale.fill_(1.0)
94
+ self._fc2_input_scale = self.a2_gscale
95
+ else:
96
+ # W4A16 NVFP4 checkpoints have no calibrated a2_gscale; b12x
97
+ # performs dynamic per-block FC2-input quantization, so a uniform
98
+ # 1.0 scale per expert is equivalent to the bake-in above for
99
+ # static-quant checkpoints. Allocate once here so apply() stays
100
+ # alloc-free.
101
+ self._fc2_input_scale = torch.ones(
102
+ self.num_local_experts,
103
+ device=layer.w13_weight.device,
104
+ dtype=torch.float32,
105
+ )
106
107
# Precompute MMA-layout views of the weight scale factors once here
108
# rather than recomputing on every forward pass.
@@ -131,7 +148,13 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular):
148
weight_key: QuantKey | None,
149
activation_key: QuantKey | None,
150
) -> bool:
134
- return (weight_key, activation_key) == (kNvfp4Static, kNvfp4Dynamic)
151
+ # b12x performs in-kernel BF16->FP4 activation quant, so W4A16
152
+ # NVFP4 checkpoints (activation_key=None, e.g. mixed-precision
153
+ # compressed-tensors layouts) are runtime-compatible.
154
+ return (weight_key, activation_key) in (
155
+ (kNvfp4Static, kNvfp4Dynamic),
156
+ (kNvfp4Static, None),
157
+ )
158
159
@staticmethod
160
def _supports_activation(activation: MoEActivation) -> bool:
@@ -198,8 +221,8 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular):
221
assert self.g1_alphas is not None and self.g2_alphas is not None, (
222
"g1_alphas and g2_alphas must not be None for FlashInferB12xExperts"
223
)
201
- assert self.a2_gscale is not None, (
202
- "a2_gscale must not be None for FlashInferB12xExperts"
224
+ assert self._fc2_input_scale is not None, (
225
+ "_fc2_input_scale must be set by process_weights_after_loading"
226
)
227
228
top_k = topk_ids.shape[1]
@@ -211,7 +234,7 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular):
234
w1_weight=w1,
235
w1_weight_sf=self.w1_sf_mma,
236
w1_alpha=self.g1_alphas,
214
- fc2_input_scale=self.a2_gscale,
237
+ fc2_input_scale=self._fc2_input_scale,
238
w2_weight=w2,
239
w2_weight_sf=self.w2_sf_mma,
240
w2_alpha=self.g2_alphas,