diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 1b8fa13a519a..6a131f8aaa70 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -404,6 +404,24 @@ def _get_kv_b_proj_input_dtype( return weight_dtype +def split_kv_b_proj( + kv_b_proj: nn.Module, + out_dtype: torch.dtype, + kv_lora_rank: int, + num_heads: int, + qk_nope_head_dim: int, + v_head_dim: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Dequantize ``kv_b_proj`` and return ``W_UK [L,N,P]``, ``W_UV [L,N,V]``.""" + weight = get_and_maybe_dequant_weights(kv_b_proj, out_dtype=out_dtype).T + assert weight.shape == ( + kv_lora_rank, + num_heads * (qk_nope_head_dim + v_head_dim), + ), f"kv_b_proj weight {tuple(weight.shape)} vs {kv_lora_rank=} {num_heads=}" + weight = weight.view(kv_lora_rank, num_heads, qk_nope_head_dim + v_head_dim) + return weight.split([qk_nope_head_dim, v_head_dim], dim=-1) + + class MLAAttention(nn.Module, AttentionLayerBase): """Multi-Head Latent Attention layer. @@ -1221,13 +1239,6 @@ def process_weights_after_loading(self, act_dtype: torch.dtype): ) return - # we currently do not have quantized bmm's which are needed for - # `W_UV` and `W_UK_T`, we just store fp16/bf16 copies and perform - # the bmm's in 16-bit, the extra memory overhead of this is fairly low - kv_b_proj_weight = get_and_maybe_dequant_weights( - self.kv_b_proj, out_dtype=act_dtype - ).T - if self.dcp_q_replicate: # qrep wired here: validate unsupported decode backends once. assert self.q_pad_num_heads in (None, self.num_heads), ( @@ -1243,24 +1254,13 @@ def process_weights_after_loading(self, act_dtype: torch.dtype): "FP4/FP8 MLA BMM paths." ) - assert kv_b_proj_weight.shape == ( - self.kv_lora_rank, - self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), - ), ( - f"{kv_b_proj_weight.shape=}, " - f"{self.kv_lora_rank=}, " - f"{self.num_heads=}, " - f"{self.qk_nope_head_dim=}, " - f"{self.v_head_dim=}" - ) - kv_b_proj_weight = kv_b_proj_weight.view( + W_UK, W_UV = split_kv_b_proj( + self.kv_b_proj, + act_dtype, self.kv_lora_rank, self.num_heads, - self.qk_nope_head_dim + self.v_head_dim, - ) - - W_UK, W_UV = kv_b_proj_weight.split( - [self.qk_nope_head_dim, self.v_head_dim], dim=-1 + self.qk_nope_head_dim, + self.v_head_dim, ) # If kv_b_proj_weight is unquantized, quantize it to mxfp4 if supported diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py index cb9cb34158c2..ce4520cbbe39 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py @@ -153,12 +153,10 @@ def process_weights_after_loading(self, layer) -> None: layer.logical_widths, getattr(layer, "input_scale", None), ) - weight = weight.t() elif self.strategy == QuantizationStrategy.CHANNEL: weight, weight_scale, input_scale = process_fp8_weight_channel_strategy( layer.weight, layer.weight_scale, getattr(layer, "input_scale", None) ) - weight = weight.t() elif self.strategy == QuantizationStrategy.BLOCK: assert self.is_static_input_scheme is False @@ -177,8 +175,7 @@ def process_weights_after_loading(self, layer) -> None: # required by torch.compile to be torch.nn.Parameter layer.weight = Parameter(weight.data, requires_grad=False) - # Preserve the dim tags dropped by the transpose so layout-aware - # kernels (humming) see (K, N) instead of assuming (N, K). + # The helpers return (K, N); retag so layout-aware kernels see it. layer.weight.input_dim = 0 layer.weight.output_dim = 1 layer.weight_scale = Parameter(weight_scale.data, requires_grad=False) diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index c85a2c288e76..47df735e89df 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -416,7 +416,6 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: if self.act_q_static: assert input_scale is not None input_scale = input_scale.max() - weight = weight.t() # Update layer with new values. replace_parameter(layer, "weight", weight.data) diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 364c52b4c879..a9ac30715928 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -2129,7 +2129,7 @@ def process(self, layer, role) -> None: weight, weight_scale, _ = process_fp8_weight_channel_strategy( layer.weight, layer.weight_scale.data ) - layer.weight = Parameter(weight.t(), requires_grad=False) + layer.weight = Parameter(weight, requires_grad=False) layer.weight_scale = Parameter(weight_scale, requires_grad=False) diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py index 79859a69d440..c9888dc3cb6f 100644 --- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py @@ -1373,7 +1373,7 @@ def process_fp8_weight_tensor_strategy( logical_widths: list[int], input_scale: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: - """Process weights for tensor-wise quantization strategy.""" + """Requantize fused shards to one scale and return ``(K, N)`` weight.""" from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( normalize_e4m3fn_to_e4m3fnuz, requantize_with_max_scale, @@ -1391,7 +1391,7 @@ def process_fp8_weight_tensor_strategy( logical_widths=logical_widths, ) - weight = _maybe_pad_fp8_weight(weight) + weight = _maybe_pad_fp8_weight(weight).t() return weight, weight_scale, input_scale @@ -1400,7 +1400,7 @@ def process_fp8_weight_channel_strategy( weight_scale: torch.Tensor, input_scale: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: - """Process weights for channel-wise quantization strategy.""" + """Normalize FNUZ if needed and return ``(K, N)`` weight.""" from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( normalize_e4m3fn_to_e4m3fnuz, ) @@ -1410,14 +1410,14 @@ def process_fp8_weight_channel_strategy( weight=weight, weight_scale=weight_scale, input_scale=input_scale ) - return weight, weight_scale, input_scale + return weight.t(), weight_scale, input_scale def process_fp8_weight_block_strategy( weight: torch.Tensor, weight_scale: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: - """Process weights for block-wise quantization strategy.""" + """Normalize FNUZ if needed and return ``(N, K)`` weight (no transpose).""" from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( normalize_e4m3fn_to_e4m3fnuz, ) diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py index c13537513f45..188f72e66016 100644 --- a/vllm/models/kimi_k3/nvidia/mla.py +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -58,6 +58,7 @@ accumulate_mla_context_chunk, init_mla_context_partial, neutralize_empty_context_partials, + split_kv_b_proj, ) from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.layernorm import RMSNorm @@ -73,9 +74,6 @@ QuantizationConfig, resolve_quant_method, ) -from vllm.model_executor.layers.quantization.utils.quant_utils import ( - get_and_maybe_dequant_weights, -) from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding, get_rope from vllm.model_executor.utils import replace_parameter from vllm.models.common.ops import fused_q_kv_rmsnorm @@ -441,20 +439,13 @@ def process_weights_after_loading(self, act_dtype: torch.dtype) -> None: projected into latent space by ``W_UK_T`` and the attention output is projected back to ``v`` by ``W_UV`` -- avoiding materializing full K/V. """ - kv_b_proj_weight = get_and_maybe_dequant_weights( - self.kv_b_proj, out_dtype=act_dtype - ).T - assert kv_b_proj_weight.shape == ( - self.kv_lora_rank, - self.num_local_heads * (self.qk_nope_head_dim + self.v_head_dim), - ), f"{kv_b_proj_weight.shape=}" - kv_b_proj_weight = kv_b_proj_weight.view( + W_UK, W_UV = split_kv_b_proj( + self.kv_b_proj, + act_dtype, self.kv_lora_rank, self.num_local_heads, - self.qk_nope_head_dim + self.v_head_dim, - ) - W_UK, W_UV = kv_b_proj_weight.split( - [self.qk_nope_head_dim, self.v_head_dim], dim=-1 + self.qk_nope_head_dim, + self.v_head_dim, ) # (L, N, V) -> (N, L, V) replace_parameter(self, "W_UV", W_UV.transpose(0, 1), prefer_copy=True)