[Perf] Kimi K3 nvfp4 Align in_proj weights by 128 to avoid elementwise copy (#55242)

Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>
Co-authored-by: OpenAI Codex <codex@openai.com>

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)

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:

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}")

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

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":

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.