@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 973a3be780

[Refactor][Quantization] Make FP8 and MLA weight transforms reusable pure functions (#57732)

Signed-off-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>

aoshen02 committed Sep 23, 2026 at 16:16 UTC 973a3be780bd8b0d2f69c3727fb022e90a06ed74

6 files changed +36 -49

vllm/model_executor/layers/attention/mla_attention.py

+23 -23

@@ -404,6 +404,24 @@ def _get_kv_b_proj_input_dtype(

404 return weight_dtype

405

406

407 +def split_kv_b_proj(

408 + kv_b_proj: nn.Module,

409 + out_dtype: torch.dtype,

410 + kv_lora_rank: int,

411 + num_heads: int,

412 + qk_nope_head_dim: int,

413 + v_head_dim: int,

414 +) -> tuple[torch.Tensor, torch.Tensor]:

415 + """Dequantize ``kv_b_proj`` and return ``W_UK [L,N,P]``, ``W_UV [L,N,V]``."""

416 + weight = get_and_maybe_dequant_weights(kv_b_proj, out_dtype=out_dtype).T

417 + assert weight.shape == (

418 + kv_lora_rank,

419 + num_heads * (qk_nope_head_dim + v_head_dim),

420 + ), f"kv_b_proj weight {tuple(weight.shape)} vs {kv_lora_rank=} {num_heads=}"

421 + weight = weight.view(kv_lora_rank, num_heads, qk_nope_head_dim + v_head_dim)

422 + return weight.split([qk_nope_head_dim, v_head_dim], dim=-1)

423 +

424 +

425 class MLAAttention(nn.Module, AttentionLayerBase):

426 """Multi-Head Latent Attention layer.

427

@@ -1221,13 +1239,6 @@ class MLAAttention(nn.Module, AttentionLayerBase):

1239 )

1240 return

1241

1224 - # we currently do not have quantized bmm's which are needed for

1225 - # `W_UV` and `W_UK_T`, we just store fp16/bf16 copies and perform

1226 - # the bmm's in 16-bit, the extra memory overhead of this is fairly low

1227 - kv_b_proj_weight = get_and_maybe_dequant_weights(

1228 - self.kv_b_proj, out_dtype=act_dtype

1229 - ).T

1230 -

1242 if self.dcp_q_replicate:

1243 # qrep wired here: validate unsupported decode backends once.

1244 assert self.q_pad_num_heads in (None, self.num_heads), (

@@ -1243,24 +1254,13 @@ class MLAAttention(nn.Module, AttentionLayerBase):

1254 "FP4/FP8 MLA BMM paths."

1255 )

1256

1246 - assert kv_b_proj_weight.shape == (

1247 - self.kv_lora_rank,

1248 - self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),

1249 - ), (

1250 - f"{kv_b_proj_weight.shape=}, "

1251 - f"{self.kv_lora_rank=}, "

1252 - f"{self.num_heads=}, "

1253 - f"{self.qk_nope_head_dim=}, "

1254 - f"{self.v_head_dim=}"

1255 - )

1256 - kv_b_proj_weight = kv_b_proj_weight.view(

1257 + W_UK, W_UV = split_kv_b_proj(

1258 + self.kv_b_proj,

1259 + act_dtype,

1260 self.kv_lora_rank,

1261 self.num_heads,

1259 - self.qk_nope_head_dim + self.v_head_dim,

1260 - )

1261 -

1262 - W_UK, W_UV = kv_b_proj_weight.split(

1263 - [self.qk_nope_head_dim, self.v_head_dim], dim=-1

1262 + self.qk_nope_head_dim,

1263 + self.v_head_dim,

1264 )

1265

1266 # If kv_b_proj_weight is unquantized, quantize it to mxfp4 if supported

vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py

+1 -4

@@ -153,12 +153,10 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):

153 layer.logical_widths,

154 getattr(layer, "input_scale", None),

155 )

156 - weight = weight.t()

156 elif self.strategy == QuantizationStrategy.CHANNEL:

157 weight, weight_scale, input_scale = process_fp8_weight_channel_strategy(

158 layer.weight, layer.weight_scale, getattr(layer, "input_scale", None)

159 )

161 - weight = weight.t()

160

161 elif self.strategy == QuantizationStrategy.BLOCK:

162 assert self.is_static_input_scheme is False

@@ -177,8 +175,7 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):

175

176 # required by torch.compile to be torch.nn.Parameter

177 layer.weight = Parameter(weight.data, requires_grad=False)

180 - # Preserve the dim tags dropped by the transpose so layout-aware

181 - # kernels (humming) see (K, N) instead of assuming (N, K).

178 + # The helpers return (K, N); retag so layout-aware kernels see it.

179 layer.weight.input_dim = 0

180 layer.weight.output_dim = 1

181 layer.weight_scale = Parameter(weight_scale.data, requires_grad=False)

vllm/model_executor/layers/quantization/fp8.py

-1

@@ -416,7 +416,6 @@ class Fp8LinearMethod(LinearMethodBase):

416 if self.act_q_static:

417 assert input_scale is not None

418 input_scale = input_scale.max()

419 - weight = weight.t()

419

420 # Update layer with new values.

421 replace_parameter(layer, "weight", weight.data)

vllm/model_executor/layers/quantization/modelopt.py

+1 -1

@@ -2129,7 +2129,7 @@ class KFp8StaticChannel(QuantKeyScheme):

2129 weight, weight_scale, _ = process_fp8_weight_channel_strategy(

2130 layer.weight, layer.weight_scale.data

2131 )

2132 - layer.weight = Parameter(weight.t(), requires_grad=False)

2132 + layer.weight = Parameter(weight, requires_grad=False)

2133 layer.weight_scale = Parameter(weight_scale, requires_grad=False)

2134

2135

vllm/model_executor/layers/quantization/utils/fp8_utils.py

+5 -5

@@ -1373,7 +1373,7 @@ def process_fp8_weight_tensor_strategy(

1373 logical_widths: list[int],

1374 input_scale: torch.Tensor | None = None,

1375 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:

1376 - """Process weights for tensor-wise quantization strategy."""

1376 + """Requantize fused shards to one scale and return ``(K, N)`` weight."""

1377 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (

1378 normalize_e4m3fn_to_e4m3fnuz,

1379 requantize_with_max_scale,

@@ -1391,7 +1391,7 @@ def process_fp8_weight_tensor_strategy(

1391 logical_widths=logical_widths,

1392 )

1393

1394 - weight = _maybe_pad_fp8_weight(weight)

1394 + weight = _maybe_pad_fp8_weight(weight).t()

1395 return weight, weight_scale, input_scale

1396

1397

@@ -1400,7 +1400,7 @@ def process_fp8_weight_channel_strategy(

1400 weight_scale: torch.Tensor,

1401 input_scale: torch.Tensor | None = None,

1402 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:

1403 - """Process weights for channel-wise quantization strategy."""

1403 + """Normalize FNUZ if needed and return ``(K, N)`` weight."""

1404 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (

1405 normalize_e4m3fn_to_e4m3fnuz,

1406 )

@@ -1410,14 +1410,14 @@ def process_fp8_weight_channel_strategy(

1410 weight=weight, weight_scale=weight_scale, input_scale=input_scale

1411 )

1412

1413 - return weight, weight_scale, input_scale

1413 + return weight.t(), weight_scale, input_scale

1414

1415

1416 def process_fp8_weight_block_strategy(

1417 weight: torch.Tensor,

1418 weight_scale: torch.Tensor,

1419 ) -> tuple[torch.Tensor, torch.Tensor]:

1420 - """Process weights for block-wise quantization strategy."""

1420 + """Normalize FNUZ if needed and return ``(N, K)`` weight (no transpose)."""

1421 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (

1422 normalize_e4m3fn_to_e4m3fnuz,

1423 )

vllm/models/kimi_k3/nvidia/mla.py

+6 -15

@@ -58,6 +58,7 @@ from vllm.model_executor.layers.attention.mla_attention import (

58 accumulate_mla_context_chunk,

59 init_mla_context_partial,

60 neutralize_empty_context_partials,

61 + split_kv_b_proj,

62 )

63 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

64 from vllm.model_executor.layers.layernorm import RMSNorm

@@ -73,9 +74,6 @@ from vllm.model_executor.layers.quantization import (

74 QuantizationConfig,

75 resolve_quant_method,

76 )

76 -from vllm.model_executor.layers.quantization.utils.quant_utils import (

77 - get_and_maybe_dequant_weights,

78 -)

77 from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding, get_rope

78 from vllm.model_executor.utils import replace_parameter

79 from vllm.models.common.ops import fused_q_kv_rmsnorm

@@ -441,20 +439,13 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):

439 projected into latent space by ``W_UK_T`` and the attention output is

440 projected back to ``v`` by ``W_UV`` -- avoiding materializing full K/V.

441 """

444 - kv_b_proj_weight = get_and_maybe_dequant_weights(

445 - self.kv_b_proj, out_dtype=act_dtype

446 - ).T

447 - assert kv_b_proj_weight.shape == (

448 - self.kv_lora_rank,

449 - self.num_local_heads * (self.qk_nope_head_dim + self.v_head_dim),

450 - ), f"{kv_b_proj_weight.shape=}"

451 - kv_b_proj_weight = kv_b_proj_weight.view(

442 + W_UK, W_UV = split_kv_b_proj(

443 + self.kv_b_proj,

444 + act_dtype,

445 self.kv_lora_rank,

446 self.num_local_heads,

454 - self.qk_nope_head_dim + self.v_head_dim,

455 - )

456 - W_UK, W_UV = kv_b_proj_weight.split(

457 - [self.qk_nope_head_dim, self.v_head_dim], dim=-1

447 + self.qk_nope_head_dim,

448 + self.v_head_dim,

449 )

450 # (L, N, V) -> (N, L, V)

451 replace_parameter(self, "W_UV", W_UV.transpose(0, 1), prefer_copy=True)