[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>
Wei Zhao committed
Sep 4, 2026 at 22:46 UTC
16328c7a775ff20d17d9750096a094e8cf2b45f2
6 files changed
+53
-21
tests/quantization/test_modelopt.py
+6
-2
@@ -650,7 +650,9 @@ def test_modelopt_linear_exposes_humming_layer_attrs(dist_init, monkeypatch):
650
from vllm.config.quantization import QuantSpec
651
from vllm.model_executor.layers.quantization import modelopt as mo
652
653
- monkeypatch.setattr(mo, "select_linear_kernel", lambda spec, layer, rt: Mock())
653
+ monkeypatch.setattr(
654
+ mo, "select_linear_kernel", lambda spec, layer, rt, **kwargs: Mock()
655
+ )
656
monkeypatch.setattr(mo, "expose_input_quant_key", lambda layer, kernel: None)
657
658
def build(layer):
@@ -829,7 +831,8 @@ def test_modelopt_fp8_pb_wo_hides_output_padding(monkeypatch):
831
)
832
833
kernel = Mock()
832
- monkeypatch.setattr(mo, "select_linear_kernel", lambda spec, layer, rt: kernel)
834
+ init_fp8_linear_kernel = Mock(return_value=kernel)
835
+ monkeypatch.setattr(mo, "init_fp8_linear_kernel", init_fp8_linear_kernel)
836
monkeypatch.setattr(mo, "expose_input_quant_key", lambda layer, k: None)
837
838
method = ModelOptLinearMethod.__new__(ModelOptLinearMethod)
@@ -860,6 +863,7 @@ def test_modelopt_fp8_pb_wo_hides_output_padding(monkeypatch):
863
# loaded at logical size; scale is cdiv(2624, 128) = 21 block rows
864
assert layer.weight.shape == (2624, 128)
865
assert layer.weight_scale.shape == (21, 1, 1, 1)
866
+ assert init_fp8_linear_kernel.call_args.kwargs["weight_shape"] == (2688, 128)
867
868
layer.weight.data.fill_(1)
869
method.process_weights_after_loading(layer)
vllm/model_executor/layers/quantization/modelopt.py
+21
-3
@@ -2215,7 +2215,12 @@ def maybe_fuse_global_scales(layer) -> None:
2215
)
2216
2217
2218
-def select_linear_kernel(spec: QuantSpec, layer, rt: RuntimeDtypes):
2218
+def select_linear_kernel(
2219
+ spec: QuantSpec,
2220
+ layer,
2221
+ rt: RuntimeDtypes,
2222
+ weight_shape: tuple[int, int] | None = None,
2223
+):
2224
"""Thin family dispatcher on the weight key: nvfp4 / mxfp8 / fp8."""
2225
w = spec.weight
2226
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):
2240
weight_quant_key=w,
2241
input_dtype=rt.input_dtype,
2242
out_dtype=rt.out_dtype,
2238
- weight_shape=layer.weight.shape,
2243
+ weight_shape=weight_shape or layer.weight.shape,
2244
module_name=type(layer).__name__,
2245
)
2246
@@ -2255,6 +2260,10 @@ class FormatScheme:
2260
def extra_weights(self, layer, shapes: "Shapes", ctx: CkptCtx, wl) -> None:
2261
"""Register format-level params (after the key schemes' weights)."""
2262
2263
+ def kernel_weight_shape(self, layer) -> tuple[int, int]:
2264
+ """Return the weight shape the selected kernel will receive."""
2265
+ return tuple(layer.weight.shape)
2266
+
2267
def pre_process(self, layer) -> None:
2268
"""Run before the key schemes' ``process``."""
2269
@@ -2284,6 +2293,10 @@ class _Fp8PbWoPartialBlock(FormatScheme):
2293
the weight rows and the output need adjusting.
2294
"""
2295
2296
+ def kernel_weight_shape(self, layer) -> tuple[int, int]:
2297
+ out, in_ = layer.weight.shape
2298
+ return (cdiv(out, 128) * 128, in_)
2299
+
2300
def post_process(self, layer) -> None:
2301
w = layer.weight
2302
out = w.shape[0]
@@ -2405,7 +2418,12 @@ class ModelOptLinearMethod(LinearMethodBase):
2418
self.fmt.extra_weights(layer, shapes, self.ctx, weight_loader)
2419
2420
rt = RuntimeDtypes(self.input_dtype, self.out_dtype, self.marlin_input_dtype)
2408
- self.kernel = select_linear_kernel(self.spec, layer, rt)
2421
+ self.kernel = select_linear_kernel(
2422
+ self.spec,
2423
+ layer,
2424
+ rt,
2425
+ weight_shape=self.fmt.kernel_weight_shape(layer),
2426
+ )
2427
expose_input_quant_key(layer, self.kernel)
2428
2429
def process_weights_after_loading(self, layer) -> None:
vllm/models/kimi_k3/nvidia/kda.py
+14
-8
@@ -33,6 +33,9 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
33
from vllm.model_executor.layers.mamba.ops.gather_initial_states import (
34
gather_initial_states,
35
)
36
+from vllm.model_executor.layers.quantization.modelopt import (
37
+ ModelOptMixedPrecisionConfig,
38
+)
39
from vllm.model_executor.model_loader.weight_utils import (
40
default_weight_loader,
41
sharded_weight_loader,
@@ -428,8 +431,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
431
)
432
self._projection_overlap_max_tokens = 0
433
431
- # Keep f_a before the narrow beta shard, then pad each TP-local row
432
- # to select the aligned BF16 GEMM path.
434
+ # Keep f_a before the narrow beta shard, then align each TP-local row.
435
qkvg_output_sizes = [self.projection_size] * 4
436
in_proj_output_sizes = qkvg_output_sizes + [
437
self.head_dim,
@@ -438,7 +440,14 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
440
local_output_size = (
441
4 * self.local_projection_size + self.head_dim + self.local_num_heads
442
)
441
- self.in_proj_padding = -local_output_size % 16
443
+ in_proj_prefix = f"{prefix}.in_proj_qkvgfab"
444
+ alignment = (
445
+ 128
446
+ if isinstance(self.quant_config, ModelOptMixedPrecisionConfig)
447
+ and self.quant_config._resolve_quant_algo(in_proj_prefix) == "FP8_PB_WO"
448
+ else 16
449
+ )
450
+ self.in_proj_padding = -local_output_size % alignment
451
if self.in_proj_padding:
452
in_proj_output_sizes.append(self.in_proj_padding * self.tp_size)
453
self.in_proj_qkvgfab = _KimiGDNMergedColumnParallelLinear(
@@ -448,7 +457,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
457
tp_size=self.tp_size,
458
bias=False,
459
quant_config=self.quant_config,
451
- prefix=f"{prefix}.in_proj_qkvgfab",
460
+ prefix=in_proj_prefix,
461
)
462
if self.in_proj_padding:
463
self.in_proj_qkvgfab.weight.data[-self.in_proj_padding :].zero_()
@@ -588,10 +597,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
597
if gemm_rs_ar.can_run(self.o_proj):
598
self.gemm_rs_ar = gemm_rs_ar
599
else:
591
- logger.warning_once(
592
- "GEMM-RS/AR is disabled for %s due to an incompatible projection.",
593
- prefix,
594
- )
600
+ gemm_rs_ar.warn_incompatible_projection()
601
compilation_config = vllm_config.compilation_config
602
if prefix in compilation_config.static_forward_context:
603
raise ValueError(f"Duplicate layer name: {prefix}")
vllm/models/kimi_k3/nvidia/mla.py
+1
-4
@@ -309,10 +309,7 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):
309
if gemm_rs_ar.can_run(self.o_proj):
310
self.gemm_rs_ar = gemm_rs_ar
311
else:
312
- logger.warning_once(
313
- "GEMM-RS/AR is disabled for %s due to an incompatible projection.",
314
- prefix,
315
- )
312
+ gemm_rs_ar.warn_incompatible_projection()
313
314
# ---- Attention backend / impl / KV cache ----
315
self.quant_config = quant_config
vllm/models/kimi_k3/nvidia/model.py
+1
-4
@@ -283,10 +283,7 @@ class KimiMLP(nn.Module):
283
if gemm_rs_ar.can_run(self.down_proj):
284
self.gemm_rs_ar = gemm_rs_ar
285
else:
286
- logger.warning_once(
287
- "GEMM-RS/AR is disabled for %s due to an incompatible projection.",
288
- prefix,
289
- )
286
+ gemm_rs_ar.warn_incompatible_projection()
287
if hidden_act == "silu":
288
self.act_fn = SiluAndMul()
289
elif hidden_act == "situ":
vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py
+10
@@ -22,8 +22,11 @@ from cutlass.utils import get_smem_capacity_in_bytes
22
23
from vllm.cute_utils import _tcgen05, mbarrier, simple_tma_copy, to_cta0_smem
24
from vllm.distributed import get_tp_group
25
+from vllm.logger import init_logger
26
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
27
28
+logger = init_logger(__name__)
29
+
30
31
@dsl_user_op
32
def nanosleep(ns: int, *, loc=None, ip=None) -> None:
@@ -779,6 +782,13 @@ class GemmRsAr:
782
and w.is_contiguous()
783
)
784
785
+ def warn_incompatible_projection(self) -> None:
786
+ logger.warning_once(
787
+ "Some projections are incompatible with GEMM-RS/AR; using the "
788
+ "unfused path instead.",
789
+ scope="global",
790
+ )
791
+
792
def should_run(self, x: torch.Tensor) -> bool:
793
# Use the same threshold for RS and AR for now. Small-M shapes are
794
# supported but faster on the existing LL path.