diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py index 5bb1aa0279..151fbbf840 100644 --- a/tests/quantization/test_modelopt.py +++ b/tests/quantization/test_modelopt.py @@ -650,7 +650,9 @@ def test_modelopt_linear_exposes_humming_layer_attrs(dist_init, monkeypatch): from vllm.config.quantization import QuantSpec from vllm.model_executor.layers.quantization import modelopt as mo - monkeypatch.setattr(mo, "select_linear_kernel", lambda spec, layer, rt: Mock()) + monkeypatch.setattr( + mo, "select_linear_kernel", lambda spec, layer, rt, **kwargs: Mock() + ) monkeypatch.setattr(mo, "expose_input_quant_key", lambda layer, kernel: None) def build(layer): @@ -829,7 +831,8 @@ def test_modelopt_fp8_pb_wo_hides_output_padding(monkeypatch): ) kernel = Mock() - monkeypatch.setattr(mo, "select_linear_kernel", lambda spec, layer, rt: kernel) + init_fp8_linear_kernel = Mock(return_value=kernel) + monkeypatch.setattr(mo, "init_fp8_linear_kernel", init_fp8_linear_kernel) monkeypatch.setattr(mo, "expose_input_quant_key", lambda layer, k: None) method = ModelOptLinearMethod.__new__(ModelOptLinearMethod) @@ -860,6 +863,7 @@ def test_modelopt_fp8_pb_wo_hides_output_padding(monkeypatch): # loaded at logical size; scale is cdiv(2624, 128) = 21 block rows assert layer.weight.shape == (2624, 128) assert layer.weight_scale.shape == (21, 1, 1, 1) + assert init_fp8_linear_kernel.call_args.kwargs["weight_shape"] == (2688, 128) layer.weight.data.fill_(1) method.process_weights_after_loading(layer) diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 465e8bd573..e2ab74c369 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -2215,7 +2215,12 @@ def maybe_fuse_global_scales(layer) -> None: ) -def select_linear_kernel(spec: QuantSpec, layer, rt: RuntimeDtypes): +def select_linear_kernel( + spec: QuantSpec, + layer, + rt: RuntimeDtypes, + weight_shape: tuple[int, int] | None = None, +): """Thin family dispatcher on the weight key: nvfp4 / mxfp8 / fp8.""" w = spec.weight assert isinstance(w, QuantKey), f"resolve() must supply a weight key, got {w!r}" @@ -2235,7 +2240,7 @@ def select_linear_kernel(spec: QuantSpec, layer, rt: RuntimeDtypes): weight_quant_key=w, input_dtype=rt.input_dtype, out_dtype=rt.out_dtype, - weight_shape=layer.weight.shape, + weight_shape=weight_shape or layer.weight.shape, module_name=type(layer).__name__, ) @@ -2255,6 +2260,10 @@ class FormatScheme: def extra_weights(self, layer, shapes: "Shapes", ctx: CkptCtx, wl) -> None: """Register format-level params (after the key schemes' weights).""" + def kernel_weight_shape(self, layer) -> tuple[int, int]: + """Return the weight shape the selected kernel will receive.""" + return tuple(layer.weight.shape) + def pre_process(self, layer) -> None: """Run before the key schemes' ``process``.""" @@ -2284,6 +2293,10 @@ class _Fp8PbWoPartialBlock(FormatScheme): the weight rows and the output need adjusting. """ + def kernel_weight_shape(self, layer) -> tuple[int, int]: + out, in_ = layer.weight.shape + return (cdiv(out, 128) * 128, in_) + def post_process(self, layer) -> None: w = layer.weight out = w.shape[0] @@ -2405,7 +2418,12 @@ class ModelOptLinearMethod(LinearMethodBase): self.fmt.extra_weights(layer, shapes, self.ctx, weight_loader) rt = RuntimeDtypes(self.input_dtype, self.out_dtype, self.marlin_input_dtype) - self.kernel = select_linear_kernel(self.spec, layer, rt) + self.kernel = select_linear_kernel( + self.spec, + layer, + rt, + weight_shape=self.fmt.kernel_weight_shape(layer), + ) expose_input_quant_key(layer, self.kernel) def process_weights_after_loading(self, layer) -> None: diff --git a/vllm/models/kimi_k3/nvidia/kda.py b/vllm/models/kimi_k3/nvidia/kda.py index 3b17445a39..bfc3572a18 100644 --- a/vllm/models/kimi_k3/nvidia/kda.py +++ b/vllm/models/kimi_k3/nvidia/kda.py @@ -33,6 +33,9 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( from vllm.model_executor.layers.mamba.ops.gather_initial_states import ( gather_initial_states, ) +from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptMixedPrecisionConfig, +) from vllm.model_executor.model_loader.weight_utils import ( default_weight_loader, sharded_weight_loader, @@ -428,8 +431,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): ) self._projection_overlap_max_tokens = 0 - # Keep f_a before the narrow beta shard, then pad each TP-local row - # to select the aligned BF16 GEMM path. + # Keep f_a before the narrow beta shard, then align each TP-local row. qkvg_output_sizes = [self.projection_size] * 4 in_proj_output_sizes = qkvg_output_sizes + [ self.head_dim, @@ -438,7 +440,14 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): local_output_size = ( 4 * self.local_projection_size + self.head_dim + self.local_num_heads ) - self.in_proj_padding = -local_output_size % 16 + in_proj_prefix = f"{prefix}.in_proj_qkvgfab" + alignment = ( + 128 + if isinstance(self.quant_config, ModelOptMixedPrecisionConfig) + and self.quant_config._resolve_quant_algo(in_proj_prefix) == "FP8_PB_WO" + else 16 + ) + self.in_proj_padding = -local_output_size % alignment if self.in_proj_padding: in_proj_output_sizes.append(self.in_proj_padding * self.tp_size) self.in_proj_qkvgfab = _KimiGDNMergedColumnParallelLinear( @@ -448,7 +457,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): tp_size=self.tp_size, bias=False, quant_config=self.quant_config, - prefix=f"{prefix}.in_proj_qkvgfab", + prefix=in_proj_prefix, ) if self.in_proj_padding: self.in_proj_qkvgfab.weight.data[-self.in_proj_padding :].zero_() @@ -588,10 +597,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): if gemm_rs_ar.can_run(self.o_proj): self.gemm_rs_ar = gemm_rs_ar else: - logger.warning_once( - "GEMM-RS/AR is disabled for %s due to an incompatible projection.", - prefix, - ) + gemm_rs_ar.warn_incompatible_projection() compilation_config = vllm_config.compilation_config if prefix in compilation_config.static_forward_context: raise ValueError(f"Duplicate layer name: {prefix}") diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py index 9cd261c73a..87bf59749c 100644 --- a/vllm/models/kimi_k3/nvidia/mla.py +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -309,10 +309,7 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase): if gemm_rs_ar.can_run(self.o_proj): self.gemm_rs_ar = gemm_rs_ar else: - logger.warning_once( - "GEMM-RS/AR is disabled for %s due to an incompatible projection.", - prefix, - ) + gemm_rs_ar.warn_incompatible_projection() # ---- Attention backend / impl / KV cache ---- self.quant_config = quant_config diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 2fc20ecf21..0eb473623d 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -283,10 +283,7 @@ class KimiMLP(nn.Module): if gemm_rs_ar.can_run(self.down_proj): self.gemm_rs_ar = gemm_rs_ar else: - logger.warning_once( - "GEMM-RS/AR is disabled for %s due to an incompatible projection.", - prefix, - ) + gemm_rs_ar.warn_incompatible_projection() if hidden_act == "silu": self.act_fn = SiluAndMul() elif hidden_act == "situ": diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py index 923f76b6c9..f8a90cc1ab 100644 --- a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py @@ -22,8 +22,11 @@ from cutlass.utils import get_smem_capacity_in_bytes from vllm.cute_utils import _tcgen05, mbarrier, simple_tma_copy, to_cta0_smem from vllm.distributed import get_tp_group +from vllm.logger import init_logger from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod +logger = init_logger(__name__) + @dsl_user_op def nanosleep(ns: int, *, loc=None, ip=None) -> None: @@ -779,6 +782,13 @@ class GemmRsAr: and w.is_contiguous() ) + def warn_incompatible_projection(self) -> None: + logger.warning_once( + "Some projections are incompatible with GEMM-RS/AR; using the " + "unfused path instead.", + scope="global", + ) + def should_run(self, x: torch.Tensor) -> bool: # Use the same threshold for RS and AR for now. Small-M shapes are # supported but faster on the existing LL path.