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 14, 05:30 UTC
[Quantization/NVFP4] Speed up TRTLLM NVFP4 MOE weight loading and fix K/V scale loading for MLA Attn (#25968)
Signed-off-by: Pavani Majety <pmajety@nvidia.com>
Pavani Majety committed
Oct 3, 2025 at 12:35 UTC
a26917332fabf5fee6544f2215e211f59d27a774
3 files changed
+77
-55
vllm/model_executor/layers/quantization/kv_cache.py
+4
-4
@@ -86,7 +86,7 @@ class BaseKVCacheMethod(QuantizeMethodBase):
86
logger.warning_once(
87
"Checkpoint does not provide a q scaling factor. "
88
"Setting it to k_scale. This only matters for "
89
- "the flash-attn backend.")
89
+ "FP8 Attention backends (flash-attn or flashinfer).")
90
layer._q_scale.copy_(k_scale)
91
layer._q_scale_float = k_scale
92
@@ -98,9 +98,9 @@ class BaseKVCacheMethod(QuantizeMethodBase):
98
if (k_scale == 1.0 and v_scale == 1.0
99
and "e5m2" not in layer.kv_cache_dtype):
100
logger.warning_once(
101
- "Using KV cache scaling factor 1.0 for fp8_e4m3. This "
102
- "may cause accuracy issues. Please make sure k/v_scale "
103
- "scaling factors are available in the fp8 checkpoint.")
101
+ "Using KV cache scaling factor 1.0 for fp8_e4m3. "
102
+ "If this is unintended, verify that k/v_scale "
103
+ "scaling factors are properly set in the checkpoint.")
104
105
if layer.q_scale > 0.0:
106
q_scale = layer.q_scale
vllm/model_executor/layers/quantization/modelopt.py
+66
-50
@@ -1064,7 +1064,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1064
self.allow_flashinfer = _nvfp4.allow_flashinfer
1065
self.use_marlin = _nvfp4.use_marlin
1066
self.flashinfer_moe_backend = None
1067
-
1067
+ self._cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
1068
if self.allow_flashinfer:
1069
self.flashinfer_moe_backend = get_flashinfer_moe_backend()
1070
logger.info_once(
@@ -1197,19 +1197,23 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1197
weight_loader=weight_loader)
1198
layer.register_parameter("w2_input_scale", w2_input_scale)
1199
1200
- def prepare_static_weight_layouts_for_trtllm_moe(
1200
+ def prepare_static_weights_for_trtllm_fp4_moe(
1201
self,
1202
- gemm1_weights: torch.Tensor,
1203
- gemm2_weights: torch.Tensor,
1204
- gemm1_scales_linear_fp4_bytes: torch.Tensor,
1205
- gemm2_scales_linear_fp4_bytes: torch.Tensor,
1206
- hidden_size: int,
1207
- intermediate_size: int,
1208
- num_experts: int,
1209
- ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
1202
+ # args_dequant,
1203
+ # args,
1204
+ gemm1_weights,
1205
+ gemm2_weights,
1206
+ gemm1_scales_linear_fp4_bytes,
1207
+ gemm2_scales_linear_fp4_bytes,
1208
+ hidden_size,
1209
+ intermediate_size,
1210
+ num_experts,
1211
+ ):
1212
+ from flashinfer import nvfp4_block_scale_interleave
1213
+ from flashinfer.fused_moe.core import (
1214
+ _maybe_get_cached_w2_permute_indices,
1215
+ _maybe_get_cached_w3_w1_permute_indices)
1216
"""Prepare quantized weights for kernel (done offline with weights)."""
1211
- from flashinfer import (reorder_rows_for_gated_act_gemm,
1212
- shuffle_matrix_a, shuffle_matrix_sf_a)
1217
epilogue_tile_m = 128 # FIXME: this depends on the kernel internals
1218
1219
# Convert quantized weights to proper formats
@@ -1227,48 +1231,54 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1231
intermediate_size //
1232
16) # fp8 scaling factors
1233
1230
- # Reorder rows of W1 and scales for fused gated activation
1231
- gemm1_weights_fp4_interleaved = []
1232
- gemm1_scales_fp4_interleaved = []
1233
- for i in range(num_experts):
1234
- gemm1_weights_fp4_interleaved.append(
1235
- reorder_rows_for_gated_act_gemm(gemm1_weights_fp4[i].clone()))
1236
- gemm1_scales_fp4_interleaved.append(
1237
- reorder_rows_for_gated_act_gemm(
1238
- gemm1_scales_linear_fp4[i].clone()))
1239
-
1240
- # Stack weights and scales for all experts
1241
- gemm1_weights_fp4_interleaved = torch.stack(
1242
- gemm1_weights_fp4_interleaved).reshape(num_experts,
1243
- 2 * intermediate_size,
1244
- hidden_size // 2)
1245
- gemm1_scales_fp4_interleaved = torch.stack(
1246
- gemm1_scales_fp4_interleaved).reshape(num_experts,
1247
- 2 * intermediate_size,
1248
- hidden_size // 16)
1249
-
1250
- # Shuffle weights and scaling factors for transposed mma output
1234
gemm1_weights_fp4_shuffled = []
1235
gemm1_scales_fp4_shuffled = []
1236
gemm2_weights_fp4_shuffled = []
1237
gemm2_scales_fp4_shuffled = []
1238
for i in range(num_experts):
1256
- gemm1_weights_fp4_shuffled.append(
1257
- shuffle_matrix_a(
1258
- gemm1_weights_fp4_interleaved[i].view(torch.uint8),
1259
- epilogue_tile_m))
1239
+ # Calculate the permute indices for the following:
1240
+ # 1. Reorder rows of W1 and scales for fused gated activation
1241
+ # 2. Shuffle weights and scaling factors for transposed mma output
1242
+ # for both w3_w1 and w2 weights and scale factors
1243
+ permute_indices = _maybe_get_cached_w3_w1_permute_indices(
1244
+ self._cache_permute_indices,
1245
+ gemm1_weights_fp4[i].view(torch.uint8),
1246
+ epilogue_tile_m,
1247
+ )
1248
+ gemm1_weights_fp4_shuffled.append(gemm1_weights_fp4[i].view(
1249
+ torch.uint8)[permute_indices.to(
1250
+ gemm1_weights_fp4.device)].contiguous())
1251
+
1252
+ permute_sf_indices = _maybe_get_cached_w3_w1_permute_indices(
1253
+ self._cache_permute_indices,
1254
+ gemm1_scales_linear_fp4[i].view(torch.uint8),
1255
+ epilogue_tile_m,
1256
+ num_elts_per_sf=16,
1257
+ )
1258
gemm1_scales_fp4_shuffled.append(
1261
- shuffle_matrix_sf_a(
1262
- gemm1_scales_fp4_interleaved[i].view(torch.uint8),
1263
- epilogue_tile_m))
1264
-
1265
- gemm2_weights_fp4_shuffled.append(
1266
- shuffle_matrix_a(gemm2_weights_fp4[i].view(torch.uint8),
1267
- epilogue_tile_m))
1259
+ nvfp4_block_scale_interleave(gemm1_scales_linear_fp4[i].view(
1260
+ torch.uint8)[permute_sf_indices.to(
1261
+ gemm1_scales_linear_fp4.device)].contiguous()))
1262
+
1263
+ permute_indices = _maybe_get_cached_w2_permute_indices(
1264
+ self._cache_permute_indices,
1265
+ gemm2_weights_fp4[i].view(torch.uint8),
1266
+ epilogue_tile_m,
1267
+ )
1268
+ gemm2_weights_fp4_shuffled.append(gemm2_weights_fp4[i].view(
1269
+ torch.uint8)[permute_indices.to(
1270
+ gemm2_weights_fp4.device)].contiguous())
1271
+
1272
+ permute_sf_indices = _maybe_get_cached_w2_permute_indices(
1273
+ self._cache_permute_indices,
1274
+ gemm2_scales_linear_fp4[i].view(torch.uint8),
1275
+ epilogue_tile_m,
1276
+ num_elts_per_sf=16,
1277
+ )
1278
gemm2_scales_fp4_shuffled.append(
1269
- shuffle_matrix_sf_a(
1270
- gemm2_scales_linear_fp4[i].view(torch.uint8),
1271
- epilogue_tile_m))
1279
+ nvfp4_block_scale_interleave(gemm2_scales_linear_fp4[i].view(
1280
+ torch.uint8)[permute_sf_indices.to(
1281
+ gemm2_scales_linear_fp4.device)].contiguous()))
1282
1283
# Stack weights for all experts
1284
gemm1_weights_fp4_shuffled = torch.stack(gemm1_weights_fp4_shuffled)
@@ -1283,8 +1293,12 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1293
torch.stack(gemm2_scales_fp4_shuffled).view(
1294
torch.float8_e4m3fn).reshape(num_experts, hidden_size,
1295
intermediate_size // 16))
1286
- return (gemm1_weights_fp4_shuffled, gemm1_scales_fp4_shuffled,
1287
- gemm2_weights_fp4_shuffled, gemm2_scales_fp4_shuffled)
1296
+ return (
1297
+ gemm1_weights_fp4_shuffled,
1298
+ gemm1_scales_fp4_shuffled,
1299
+ gemm2_weights_fp4_shuffled,
1300
+ gemm2_scales_fp4_shuffled,
1301
+ )
1302
1303
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
1304
# GEMM 1 processing
@@ -1334,9 +1348,10 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1348
if self.allow_flashinfer and \
1349
self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM:
1350
# Prepare static weights for TRT-LLM kernel
1351
+ # alternate: prepare_static_weight_layouts_for_trtllm_moe
1352
(gemm1_weights_fp4_shuffled, gemm1_scales_fp4_shuffled,
1353
gemm2_weights_fp4_shuffled, gemm2_scales_fp4_shuffled
1339
- ) = self.prepare_static_weight_layouts_for_trtllm_moe(
1354
+ ) = self.prepare_static_weights_for_trtllm_fp4_moe(
1355
layer.w13_weight,
1356
layer.w2_weight,
1357
layer.w13_weight_scale,
@@ -1345,6 +1360,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1360
layer.w13_weight.size(-2) // 2, # intermediate_size
1361
layer.w13_weight.size(0), # num_experts
1362
)
1363
+ logger.debug_once("Finished shuffling weights for TRT-LLM MOE")
1364
1365
layer.gemm1_weights_fp4_shuffled = Parameter(
1366
gemm1_weights_fp4_shuffled, requires_grad=False)
vllm/model_executor/model_loader/weight_utils.py
+7
-1
@@ -1003,12 +1003,18 @@ def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> Optional[str]:
1003
return None
1004
return remapped_name
1005
1006
+ if any("mla_attn" in key for key in params_dict):
1007
+ attn_str = "mla_attn.mla_attn"
1008
+ logger.debug_once(f"Found mla_attn with k_scale and v_scale in "
1009
+ f"the checkpoint, using {attn_str} as attn_str")
1010
+ else:
1011
+ attn_str = "attn"
1012
# Define scale name mapping patterns in order of precedence
1013
scale_mapping_patterns = [
1014
# ModelOpt format: .self_attn.{k,v}_proj.{k,v}_scale ->
1015
# .self_attn.attn.{k,v}_scale
1016
(r"\.self_attn\.([kv])_proj\.([kv])_scale$",
1011
- r".self_attn.attn.\2_scale"),
1017
+ rf".self_attn.{attn_str}.\2_scale"),
1018
# QKV proj format: .self_attn.qkv_proj.{k,v}_scale ->
1019
# .self_attn.attn.{k,v}_scale
1020
(r"\.self_attn\.qkv_proj\.([kv])_scale$", r".self_attn.attn.\1_scale"),