From 1260f73dbb4accaff2c8e7e1259eff54a02aa424 Mon Sep 17 00:00:00 2001 From: Dawn952 Date: Fri, 28 Aug 2026 10:36:20 +0800 Subject: [PATCH 1/6] perf(kda): overlap Kimi K3 projection preprocessing Port #14497 onto the Kimi K3 implementation now available on main. Compose the full-rank F projection at load time, overlap the float BFG path with MXFP QKV preprocessing, and keep the auxiliary stream fully joined for ACL graph capture. Preserve the configured W8A8 MXFP8 dynamic quantization scale algorithm. Signed-off-by: Dawn952 (cherry picked from commit 37df0a276703dca143446bf43480ebaad732aa5a) (cherry picked from commit 858b61937e81de059ac6ed593fb7aec2941f6f47) Signed-off-by: Dawn952 --- tests/ut/models/test_kimi_k3_adapter.py | 36 +++- tests/ut/ops/test_kimi_kda.py | 258 +++++++++++++++++++++++- vllm_ascend/models/kimi_k3.py | 16 +- vllm_ascend/ops/kimi_kda.py | 257 ++++++++++++++++++++--- 4 files changed, 519 insertions(+), 48 deletions(-) diff --git a/tests/ut/models/test_kimi_k3_adapter.py b/tests/ut/models/test_kimi_k3_adapter.py index 0627ac012a3d..9439d15050d0 100644 --- a/tests/ut/models/test_kimi_k3_adapter.py +++ b/tests/ut/models/test_kimi_k3_adapter.py @@ -69,9 +69,11 @@ def test_kimi_mixed_kda_gate_weights_use_upstream_packed_loader(monkeypatch): nn.Module.__init__(model) layer = nn.Module() layer.self_attn = nn.Module() - layer.self_attn.in_proj_gfab = nn.Module() + layer.self_attn.fused_bfg_proj = nn.Module() packed_weight = nn.Parameter(torch.empty(6, 4)) - layer.self_attn.in_proj_gfab.register_parameter("weight", packed_weight) + layer.self_attn.fused_bfg_proj.register_parameter("weight", packed_weight) + layer.self_attn.fused_bfg_proj.register_parameter("f_a_weight", nn.Parameter(torch.empty(1))) + layer.self_attn.fused_bfg_proj.register_parameter("f_b_weight", nn.Parameter(torch.empty(1))) layer.router = nn.Linear(4, 1, bias=False) model.layers = nn.ModuleList([layer]) @@ -90,26 +92,40 @@ def fake_upstream_load_weights(_self, weights): ("layers.0.router.weight", torch.full((1, 4), 0.5)), ("layers.0.self_attn.g_proj.weight", torch.full((1,), 1.0)), ("layers.0.self_attn.f_a_proj.weight", torch.full((1,), 2.0)), - ("layers.0.self_attn.b_proj.weight", torch.full((1,), 3.0)), - ("layers.0.self_attn.o_proj.weight", torch.full((1,), 4.0)), + ("layers.0.self_attn.f_b_proj.weight", torch.full((1,), 3.0)), + ("layers.0.self_attn.b_proj.weight", torch.full((1,), 4.0)), + ("layers.0.self_attn.o_proj.weight", torch.full((1,), 5.0)), ] loaded = model.load_weights(iter(source_weights)) assert remaining[0] == source_weights[0] assert remaining[-1] == source_weights[-1] - assert [name for name, _, _ in remaining[1:4]] == [ - "layers.0.self_attn.in_proj_gfab.weight", - ] * 3 - assert [loaded_weight.item() for _, loaded_weight, _ in remaining[1:4]] == [1.0, 2.0, 3.0] - assert [kwargs["loaded_shard_id"] for _, _, kwargs in remaining[1:4]] == [0, 1, 2] + assert [name for name, _, _ in remaining[1:5]] == [ + "layers.0.self_attn.fused_bfg_proj.weight", + "layers.0.self_attn.fused_bfg_proj.f_a_weight", + "layers.0.self_attn.fused_bfg_proj.f_b_weight", + "layers.0.self_attn.fused_bfg_proj.weight", + ] + assert [loaded_weight.item() for _, loaded_weight, _ in remaining[1:5]] == [1.0, 2.0, 3.0, 4.0] + assert [kwargs["loaded_shard_id"] for _, _, kwargs in remaining[1:5]] == [2, None, None, 0] assert loaded == { - "layers.0.self_attn.in_proj_gfab.weight", + "layers.0.self_attn.fused_bfg_proj.weight", + "layers.0.self_attn.fused_bfg_proj.f_a_weight", + "layers.0.self_attn.fused_bfg_proj.f_b_weight", "layers.0.router.weight", "layers.0.self_attn.o_proj.weight", } +def test_kimi_model_declares_fused_bfg_checkpoint_mapping(): + assert AscendKimiLinearModel.packed_modules_mapping["fused_bfg_proj"] == [ + "b_proj", + "f_a_proj", + "g_proj", + ] + + def test_kimi_attention_residual_stays_sequence_sharded(monkeypatch): class IdentityAttention(nn.Module): def forward(self, *, hidden_states, positions): diff --git a/tests/ut/ops/test_kimi_kda.py b/tests/ut/ops/test_kimi_kda.py index fc1ae8a096d3..845e5bebeff6 100644 --- a/tests/ut/ops/test_kimi_kda.py +++ b/tests/ut/ops/test_kimi_kda.py @@ -2,18 +2,71 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import MagicMock, patch +import pytest import torch from torch import nn from vllm_ascend.ops.kimi_kda import ( _PACKED_CONV_WEIGHT_NAME, AscendKimiK3DeltaAttention, + _KDAFusedBFGLinear, _prepare_beta, _zero_padded_output, _zero_padded_recurrent_output, ) +from vllm_ascend.quantization.methods.w4a8_mxfp4 import ( + AscendW4A8MXFPDynamicLinearMethod, +) +from vllm_ascend.quantization.methods.w8a8_mxfp8 import ( + AscendW8A8MXFP8DynamicLinearMethod, +) + + +class _RecordingLinear(nn.Module): + def __init__(self, output: torch.Tensor) -> None: + super().__init__() + self.output = output + + def forward(self, _input: torch.Tensor): + return self.output, None + + +class _RecordingStream: + def __init__(self, name: str, event_names: list[str], trace: list[str]) -> None: + self.name = name + self.event_names = iter(event_names) + self.trace = trace + + def record_event(self) -> str: + event = next(self.event_names) + self.trace.append(f"{self.name}.record:{event}") + return event + + def wait_event(self, event: str) -> None: + self.trace.append(f"{self.name}.wait:{event}") + + +class _RecordingTensor: + def __init__(self, name: str, trace: list[str]) -> None: + self.name = name + self.trace = trace + + def record_stream(self, stream: _RecordingStream) -> None: + self.trace.append(f"{self.name}.record_stream:{stream.name}") + + +class _RecordingStreamSwitch: + def __init__(self, stream: _RecordingStream, trace: list[str]) -> None: + self.stream = stream + self.trace = trace + + def __enter__(self) -> None: + self.trace.append(f"enter:{self.stream.name}") + + def __exit__(self, *args) -> None: + self.trace.append(f"exit:{self.stream.name}") def test_zero_padded_recurrent_output_clears_uncovered_tail(): @@ -82,6 +135,209 @@ def test_prepare_beta_slices_and_applies_sigmoid_in_fp32(): assert torch.all((beta >= 0.0) & (beta <= 1.0)) +@pytest.mark.parametrize("f_b_is_local", [False, True]) +def test_fused_bfg_linear_composes_f_and_packs_bfg(f_b_is_local: bool): + with ( + patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=4), + patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", return_value=2), + patch("vllm.model_executor.parameter.get_tensor_model_parallel_rank", return_value=2), + patch("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", return_value=4), + ): + linear = _KDAFusedBFGLinear( + hidden_size=6, + num_heads=8, + head_dim=3, + tp_size=4, + quant_config=None, + prefix="model.layers.0.self_attn.in_proj_gfab", + ) + + linear.weight.data.zero_() + b_weight = torch.arange(8 * 6, dtype=linear.weight.dtype).reshape(8, 6) + f_a_weight = torch.arange(3 * 6, dtype=linear.weight.dtype).reshape(3, 6) + 100 + global_f_b_weight = torch.arange(24 * 3, dtype=linear.weight.dtype).reshape(24, 3) + 200 + local_f_b_weight = global_f_b_weight[12:18] + g_weight = torch.arange(24 * 6, dtype=linear.weight.dtype).reshape(24, 6) + 200 + + linear.weight.weight_loader(linear.weight, b_weight, 0) + linear.f_a_weight.weight_loader(linear.f_a_weight, f_a_weight) + linear.f_b_weight.weight_loader( + linear.f_b_weight, + local_f_b_weight if f_b_is_local else global_f_b_weight, + ) + linear.weight.weight_loader(linear.weight, g_weight, 2) + + expected_f = (local_f_b_weight.float() @ f_a_weight.float()).to(linear.weight.dtype) + assert tuple(linear.weight.shape) == (14, 6) + torch.testing.assert_close(linear.weight[:2], b_weight[4:6]) + torch.testing.assert_close(linear.weight[2:8], expected_f) + torch.testing.assert_close(linear.weight[8:], g_weight[12:18]) + + +def test_fused_bfg_linear_recomposes_f_after_source_reload(): + with ( + patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=1), + patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", return_value=0), + patch("vllm.model_executor.parameter.get_tensor_model_parallel_rank", return_value=0), + patch("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", return_value=1), + ): + linear = _KDAFusedBFGLinear( + hidden_size=4, + num_heads=2, + head_dim=2, + tp_size=1, + quant_config=None, + prefix="model.layers.0.self_attn.in_proj_gfab", + ) + + linear.weight.data.zero_() + first_f_a = torch.arange(8, dtype=linear.weight.dtype).reshape(2, 4) + first_f_b = torch.arange(8, dtype=linear.weight.dtype).reshape(4, 2) + linear.f_a_weight.weight_loader(linear.f_a_weight, first_f_a) + torch.testing.assert_close(linear.weight[2:6], torch.zeros_like(linear.weight[2:6])) + linear.f_b_weight.weight_loader(linear.f_b_weight, first_f_b) + torch.testing.assert_close(linear.weight[2:6], first_f_b.float() @ first_f_a.float()) + + reloaded_f_a = first_f_a + 10 + reloaded_f_b = first_f_b + 20 + linear.f_a_weight.weight_loader(linear.f_a_weight, reloaded_f_a) + linear.f_b_weight.weight_loader(linear.f_b_weight, reloaded_f_b) + torch.testing.assert_close(linear.weight[2:6], reloaded_f_b.float() @ reloaded_f_a.float()) + + +def test_fused_bfg_projection_preserves_staged_outputs(): + attention = AscendKimiK3DeltaAttention.__new__(AscendKimiK3DeltaAttention) + nn.Module.__init__(attention) + attention.head_dim = 3 + attention._fused_bfg_output_sizes = (2, 6, 6) + hidden_states = torch.randn(4, 5) + fused_output = torch.arange(56, dtype=torch.float32).reshape(4, 14) + attention.fused_bfg_proj = _RecordingLinear(fused_output) + + beta, raw_gate, output_gate = attention._project_bfg(hidden_states) + torch.testing.assert_close(beta, fused_output[:, :2]) + torch.testing.assert_close(raw_gate, fused_output[:, 2:8]) + torch.testing.assert_close(output_gate, fused_output[:, 8:]) + + beta, raw_gate, output_gate = attention._postprocess_bfg(beta, raw_gate, output_gate) + torch.testing.assert_close(beta, fused_output[:, :2].unsqueeze(0)) + torch.testing.assert_close(raw_gate, fused_output[:, 2:8].reshape(4, 2, 3).unsqueeze(0)) + torch.testing.assert_close(output_gate, fused_output[:, 8:].reshape(4, 2, 3)) + + +def test_overlapped_qkv_bfg_joins_auxiliary_tail_for_graph_capture(): + attention = AscendKimiK3DeltaAttention.__new__(AscendKimiK3DeltaAttention) + nn.Module.__init__(attention) + trace: list[str] = [] + main_stream = _RecordingStream("main", ["hidden_ready", "quant_ready"], trace) + bfg_stream = _RecordingStream("bfg", ["bfg_projection_ready", "bfg_ready"], trace) + hidden_states = _RecordingTensor("hidden", trace) + raw_bfg = tuple(_RecordingTensor(name, trace) for name in ("beta_raw", "raw_gate_raw", "gate_raw")) + processed_bfg = tuple(_RecordingTensor(name, trace) for name in ("beta", "raw_gate", "output_gate")) + quantized_qkv = object() + qkv = object() + + attention._project_bfg = MagicMock(side_effect=lambda _: (trace.append("project_bfg"), raw_bfg)[1]) + attention._quantize_fused_qkv = MagicMock(side_effect=lambda _: (trace.append("dynamic_quant"), quantized_qkv)[1]) + attention._matmul_fused_qkv = MagicMock(side_effect=lambda _: (trace.append("qkv_matmul"), qkv)[1]) + attention._postprocess_bfg = MagicMock(side_effect=lambda *_: (trace.append("postprocess_bfg"), processed_bfg)[1]) + + with ( + patch("vllm_ascend.ops.kimi_kda.torch.npu.current_stream", return_value=main_stream), + patch("vllm_ascend.ops.kimi_kda._kda_bfg_stream", return_value=bfg_stream), + patch( + "vllm_ascend.ops.kimi_kda.npu_stream_switch", + side_effect=lambda stream: _RecordingStreamSwitch(stream, trace), + ), + ): + actual = attention._run_overlapped_qkv_bfg(hidden_states) + + assert actual == (qkv, *processed_bfg) + assert trace == [ + "main.record:hidden_ready", + "hidden.record_stream:bfg", + "enter:bfg", + "bfg.wait:hidden_ready", + "project_bfg", + "bfg.record:bfg_projection_ready", + "exit:bfg", + "dynamic_quant", + "main.record:quant_ready", + "main.wait:bfg_projection_ready", + "qkv_matmul", + "enter:bfg", + "bfg.wait:quant_ready", + "postprocess_bfg", + "bfg.record:bfg_ready", + "exit:bfg", + "beta.record_stream:main", + "raw_gate.record_stream:main", + "output_gate.record_stream:main", + "main.wait:bfg_ready", + ] + + +@pytest.mark.parametrize( + "quant_method_type", + [AscendW4A8MXFPDynamicLinearMethod, AscendW8A8MXFP8DynamicLinearMethod], +) +def test_fused_qkv_splits_mxfp_dynamic_quant_from_matmul(quant_method_type): + attention = AscendKimiK3DeltaAttention.__new__(AscendKimiK3DeltaAttention) + nn.Module.__init__(attention) + inner_quant_method = quant_method_type.__new__(quant_method_type) + if isinstance(inner_quant_method, AscendW8A8MXFP8DynamicLinearMethod): + inner_quant_method.dynamic_mx_quant_scale_alg = "floor" + adapter = SimpleNamespace( + quant_method=inner_quant_method, + apply=MagicMock(return_value=torch.randn(4, 18)), + ) + attention.in_proj_qkvgfab = SimpleNamespace(quant_method=adapter) + hidden_states = torch.randn(4, 6, dtype=torch.bfloat16) + quantized = torch.empty(4, 6, dtype=torch.float8_e4m3fn) + dynamic_scale = torch.empty(4, 1, dtype=torch.uint8) + + with patch( + "vllm_ascend.ops.kimi_kda.torch_npu.npu_dynamic_mx_quant", + return_value=(quantized, dynamic_scale), + ) as dynamic_quant: + qkv_input = attention._quantize_fused_qkv(hidden_states) + + assert isinstance(qkv_input, tuple) + assert qkv_input[0] is quantized + assert qkv_input[1] is dynamic_scale + if isinstance(inner_quant_method, AscendW8A8MXFP8DynamicLinearMethod): + dynamic_quant.assert_called_once_with( + hidden_states, + dst_type=torch.float8_e4m3fn, + scale_alg="floor", + ) + else: + dynamic_quant.assert_called_once_with(hidden_states, dst_type=torch.float8_e4m3fn) + output = attention._matmul_fused_qkv(qkv_input) + assert output is adapter.apply.return_value + adapter.apply.assert_called_once_with(attention.in_proj_qkvgfab, qkv_input, bias=None) + + +def test_fused_qkv_keeps_non_mxfp_quantization_in_linear_apply(): + attention = AscendKimiK3DeltaAttention.__new__(AscendKimiK3DeltaAttention) + nn.Module.__init__(attention) + adapter = SimpleNamespace( + quant_method=object(), + apply=MagicMock(return_value=torch.randn(4, 18)), + ) + attention.in_proj_qkvgfab = SimpleNamespace(quant_method=adapter) + hidden_states = torch.randn(4, 6) + + with patch("vllm_ascend.ops.kimi_kda.torch_npu.npu_dynamic_mx_quant") as dynamic_quant: + qkv_input = attention._quantize_fused_qkv(hidden_states) + + assert qkv_input is hidden_states + dynamic_quant.assert_not_called() + output = attention._matmul_fused_qkv(qkv_input) + assert output is adapter.apply.return_value + adapter.apply.assert_called_once_with(attention.in_proj_qkvgfab, hidden_states, bias=None) + + def test_prefill_fuses_raw_gate_and_updates_v_first_state(): attention = AscendKimiK3DeltaAttention.__new__(AscendKimiK3DeltaAttention) nn.Module.__init__(attention) diff --git a/vllm_ascend/models/kimi_k3.py b/vllm_ascend/models/kimi_k3.py index 622b82892cb8..c5e84a7db5be 100644 --- a/vllm_ascend/models/kimi_k3.py +++ b/vllm_ascend/models/kimi_k3.py @@ -524,7 +524,14 @@ def forward_attn_residual( class AscendKimiLinearModel(UpstreamKimiLinearModel): """Kimi text model assembled from the Ascend decoder layer.""" - packed_modules_mapping = UpstreamPackedKimiLinearModel.packed_modules_mapping + packed_modules_mapping = { + name: list(shards) for name, shards in UpstreamPackedKimiLinearModel.packed_modules_mapping.items() + } + packed_modules_mapping["fused_bfg_proj"] = [ + "b_proj", + "f_a_proj", + "g_proj", + ] # Legacy Qwen3 GQA DSpark checkpoints consume the materialized input # to each selected Kimi layer. MLA DSpark checkpoints consume the raw # prefix-sum stream used by upstream vLLM, so keep that as the default. @@ -595,9 +602,10 @@ def load_weights(self, weights): """Route mixed-precision KDA gates through vLLM's packed loader.""" params_dict = dict(self.named_parameters()) gate_mapping = ( - (".g_proj", ".in_proj_gfab", 0), - (".f_a_proj", ".in_proj_gfab", 1), - (".b_proj", ".in_proj_gfab", 2), + (".b_proj.weight", ".fused_bfg_proj.weight", 0), + (".f_a_proj.weight", ".fused_bfg_proj.f_a_weight", None), + (".f_b_proj.weight", ".fused_bfg_proj.f_b_weight", None), + (".g_proj.weight", ".fused_bfg_proj.weight", 2), ) def remap_mixed_gate_weights(): diff --git a/vllm_ascend/ops/kimi_kda.py b/vllm_ascend/ops/kimi_kda.py index 344189cc8655..602245801a38 100644 --- a/vllm_ascend/ops/kimi_kda.py +++ b/vllm_ascend/ops/kimi_kda.py @@ -10,6 +10,7 @@ from functools import wraps import torch +import torch_npu from einops import rearrange from torch import nn from vllm.compilation.breakable_cudagraph import eager_break_during_capture @@ -20,7 +21,6 @@ from vllm.model_executor.utils import replace_parameter from vllm.models.kimi_k3.nvidia.kda import ( KimiK3DeltaAttention, - _KimiGDNMergedColumnParallelLinear, ) from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd from vllm.v1.attention.backend import AttentionBackend @@ -29,9 +29,125 @@ from vllm_ascend.ops.gdn_attn_builder import AscendGDNAttentionBackend from vllm_ascend.ops.triton.fla.utils import clear_ssm_states +from vllm_ascend.quantization.methods.w4a8_mxfp4 import ( + AscendW4A8MXFPDynamicLinearMethod, +) +from vllm_ascend.quantization.methods.w8a8_mxfp8 import ( + AscendW8A8MXFP8DynamicLinearMethod, +) +from vllm_ascend.utils import npu_stream_switch _KDA_CHUNK_SIZE = 64 _PACKED_CONV_WEIGHT_NAME = "ascend_conv1d_weight" +_F_PROJ_SHARD_ID = 1 +_KDA_BFG_STREAM: torch.npu.Stream | None = None + + +def _kda_bfg_stream() -> torch.npu.Stream: + global _KDA_BFG_STREAM + if _KDA_BFG_STREAM is None: + _KDA_BFG_STREAM = torch_npu.npu.Stream() + return _KDA_BFG_STREAM + + +class _KDAFusedBFGLinear(MergedColumnParallelLinear): + """Pack beta, an offline-composed F projection, and the output gate.""" + + def __init__( + self, + hidden_size: int, + num_heads: int, + head_dim: int, + tp_size: int, + quant_config, + prefix: str, + ) -> None: + projection_size = num_heads * head_dim + super().__init__( + input_size=hidden_size, + output_sizes=[num_heads, projection_size, projection_size], + bias=False, + quant_config=quant_config, + prefix=prefix, + ) + if self.tp_size != tp_size: + raise ValueError(f"KDA fused BFG TP mismatch: layer={self.tp_size}, attention={tp_size}") + local_projection_size = projection_size // tp_size + self.f_a_weight = nn.Parameter( + self.weight.new_empty((head_dim, hidden_size)), + requires_grad=False, + ) + self.f_b_weight = nn.Parameter( + self.weight.new_empty((local_projection_size, head_dim)), + requires_grad=False, + ) + self.f_a_weight.weight_loader = self._load_f_a_weight + self.f_b_weight.weight_loader = self._load_f_b_weight + self._f_a_loaded = False + self._f_b_loaded = False + + def _load_f_a_weight( + self, + param: nn.Parameter, + loaded_weight: torch.Tensor, + loaded_shard_id: tuple[int, ...] | int | None = None, + ) -> None: + del loaded_shard_id + if param.shape != loaded_weight.shape: + raise ValueError( + "KDA f_a_proj checkpoint shape mismatch: " + f"expected {tuple(param.shape)}, got {tuple(loaded_weight.shape)}" + ) + param.data.copy_(loaded_weight) + self._f_a_loaded = True + self._maybe_fuse_f_proj() + + def _load_f_b_weight( + self, + param: nn.Parameter, + loaded_weight: torch.Tensor, + loaded_shard_id: tuple[int, ...] | int | None = None, + ) -> None: + del loaded_shard_id + if loaded_weight.shape == param.shape: + local_weight = loaded_weight + else: + expected_shape = (param.shape[0] * self.tp_size, param.shape[1]) + if loaded_weight.shape != expected_shape: + raise ValueError( + "KDA f_b_proj checkpoint shape mismatch: " + f"expected {expected_shape} or {tuple(param.shape)}, " + f"got {tuple(loaded_weight.shape)}" + ) + local_weight = loaded_weight.narrow( + 0, + self.tp_rank * param.shape[0], + param.shape[0], + ) + param.data.copy_(local_weight) + self._f_b_loaded = True + self._maybe_fuse_f_proj() + + @torch.no_grad() + def _maybe_fuse_f_proj(self) -> None: + if not self._f_a_loaded or not self._f_b_loaded: + return + output_dim = getattr(self.weight, "output_dim", None) + if output_dim is None: + raise ValueError("KDA fused f_proj requires an output-sharded parameter") + shard_offset = sum(self.output_sizes[:_F_PROJ_SHARD_ID]) // self.tp_size + shard_size = self.output_sizes[_F_PROJ_SHARD_ID] // self.tp_size + param_shard = self.weight.data.narrow(output_dim, shard_offset, shard_size) + fused_weight = torch.matmul( + self.f_b_weight.float(), + self.f_a_weight.float(), + ).to(dtype=param_shard.dtype) + if fused_weight.shape != param_shard.shape: + raise ValueError( + "KDA composed f_proj shape mismatch: " + f"expected {tuple(param_shard.shape)}, got {tuple(fused_weight.shape)}" + ) + param_shard.copy_(fused_weight) def _zero_padded_output( @@ -81,9 +197,9 @@ def __init__(self, config, vllm_config, prefix: str = "") -> None: self.uses_mixed_projection = uses_mixed_projection if uses_mixed_projection: # vLLM 0.27 packs all KDA input projections into one linear. A - # QuaRot checkpoint instead stores q/k/v as W8A8 and keeps the - # three gates in floating point, so form one fused GEMM per - # precision group instead of falling back to four projections. + # QuaRot checkpoint instead stores q/k/v as W8A8 and keeps B/F/G + # in floating point. Split those precision groups so DynamicQuant + # can overlap the composed BFG projection. self.in_proj_qkvgfab = MergedColumnParallelLinear( self.hidden_size, [self.projection_size] * 3, @@ -91,24 +207,20 @@ def __init__(self, config, vllm_config, prefix: str = "") -> None: quant_config=quant_config, prefix=f"{prefix}.in_proj_qkv", ) - gate_output_sizes = [ - self.projection_size, - self.head_dim, - self.num_heads, - ] - if self.in_proj_padding: - gate_output_sizes.append(self.in_proj_padding * self.tp_size) - self.in_proj_gfab = _KimiGDNMergedColumnParallelLinear( - self.hidden_size, - gate_output_sizes, - replicated_shard_id=1, + del self.f_b_proj + self.fused_bfg_proj = _KDAFusedBFGLinear( + hidden_size=self.hidden_size, + num_heads=self.num_heads, + head_dim=self.head_dim, tp_size=self.tp_size, - bias=False, quant_config=quant_config, prefix=f"{prefix}.in_proj_gfab", ) - if self.in_proj_padding: - self.in_proj_gfab.weight.data[-self.in_proj_padding :].zero_() + self._fused_bfg_output_sizes = ( + self.local_num_heads, + self.local_projection_size, + self.local_projection_size, + ) # Upstream's FusedRMSNormGated constructor defaults to 1e-5, while # Kimi K3 checkpoints use the model-configured RMS epsilon (1e-6 for # the production checkpoint). Preserve the checkpoint contract used @@ -149,21 +261,7 @@ def forward( ) -> torch.Tensor: if self.uses_mixed_projection: num_tokens = hidden_states.size(0) - mixed_qkv = self.in_proj_qkvgfab(hidden_states)[0] - projected_gfab = self.in_proj_gfab(hidden_states)[0] - split_sizes = [ - self.local_projection_size, - self.head_dim, - self.local_num_heads, - ] - if self.in_proj_padding: - split_sizes.append(self.in_proj_padding) - g_proj_states, f_a, beta = projected_gfab.split(split_sizes, dim=-1)[:3] - beta = beta.unsqueeze(0) - - g1 = self.f_b_proj(f_a)[0] - g1 = rearrange(g1, "n (h d) -> 1 n h d", d=self.head_dim) - g2 = rearrange(g_proj_states, "... (h d) -> ... h d", d=self.head_dim) + mixed_qkv, beta, g1, g2 = self._run_overlapped_qkv_bfg(hidden_states) core_attn_out = torch.empty( (1, num_tokens, self.local_num_heads, self.head_dim), dtype=hidden_states.dtype, @@ -180,6 +278,99 @@ def forward( return self.o_proj(core_attn_out)[0] return super().forward(hidden_states, positions) + def _run_overlapped_qkv_bfg( + self, + hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Fork BFG work to an auxiliary stream and join it back to main.""" + main_stream = torch.npu.current_stream() + bfg_stream = _kda_bfg_stream() + + hidden_states_ready = main_stream.record_event() + hidden_states.record_stream(bfg_stream) + with npu_stream_switch(bfg_stream): + bfg_stream.wait_event(hidden_states_ready) + raw_bfg = self._project_bfg(hidden_states) + bfg_projection_ready = bfg_stream.record_event() + + quantized_qkv = self._quantize_fused_qkv(hidden_states) + quant_ready = main_stream.record_event() + + main_stream.wait_event(bfg_projection_ready) + mixed_qkv = self._matmul_fused_qkv(quantized_qkv) + + with npu_stream_switch(bfg_stream): + bfg_stream.wait_event(quant_ready) + beta, g1, g2 = self._postprocess_bfg(*raw_bfg) + bfg_ready = bfg_stream.record_event() + + for tensor in (beta, g1, g2): + tensor.record_stream(main_stream) + # bfg_ready is the auxiliary stream tail. Joining that exact event is + # required for multi-stream ACL graph capture as well as eager reuse. + main_stream.wait_event(bfg_ready) + return mixed_qkv, beta, g1, g2 + + def _project_bfg( + self, + hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + fused_bfg = self.fused_bfg_proj(hidden_states)[0] + return fused_bfg.split(self._fused_bfg_output_sizes, dim=-1) + + def _postprocess_bfg( + self, + beta: torch.Tensor, + raw_gate: torch.Tensor, + output_gate: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + beta = beta.unsqueeze(0) + raw_gate = rearrange(raw_gate, "n (h d) -> 1 n h d", d=self.head_dim) + output_gate = rearrange(output_gate, "n (h d) -> n h d", d=self.head_dim) + return beta, raw_gate, output_gate + + def _quantize_fused_qkv( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + quant_method = self.in_proj_qkvgfab.quant_method + inner_quant_method = getattr(quant_method, "quant_method", quant_method) + if ( + isinstance( + inner_quant_method, + ( + AscendW4A8MXFPDynamicLinearMethod, + AscendW8A8MXFP8DynamicLinearMethod, + ), + ) + and hidden_states.dtype == torch.bfloat16 + and hidden_states.ndim == 2 + ): + if isinstance(inner_quant_method, AscendW8A8MXFP8DynamicLinearMethod): + return torch_npu.npu_dynamic_mx_quant( + hidden_states, + dst_type=torch.float8_e4m3fn, + scale_alg=inner_quant_method.dynamic_mx_quant_scale_alg, + ) + return torch_npu.npu_dynamic_mx_quant( + hidden_states, + dst_type=torch.float8_e4m3fn, + ) + return hidden_states + + def _matmul_fused_qkv( + self, + qkv_input: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + quant_method = self.in_proj_qkvgfab.quant_method + if quant_method is None: + raise RuntimeError("KDA fused QKV quantization method is not initialized") + return quant_method.apply( + self.in_proj_qkvgfab, + qkv_input, + bias=None, + ) + @staticmethod def _run_causal_conv1d( mixed_qkv: torch.Tensor, From eea77c73ee3d416f9c239624a7de7d1a92264603 Mon Sep 17 00:00:00 2001 From: Dawn952 Date: Fri, 28 Aug 2026 11:54:05 +0800 Subject: [PATCH 2/6] perf(kda): restore two-stage BFG vector overlap Move fused BFG splitting and beta FP32 sigmoid onto the auxiliary stream after QKV matmul is enqueued. Mark mixed-path beta as preprocessed so the v0.27 dispatch path only slices it, while ordinary raw beta retains its existing sigmoid contract. Keep the graph-safe auxiliary tail join and make the stream-order test mypy-safe. Signed-off-by: Dawn952 (cherry picked from commit fd8125e531052237a99db346a88d9bced1b06510) (cherry picked from commit ca19043f2ef7ce798a18982f29970a2aa6e18d1e) Signed-off-by: Dawn952 --- tests/ut/ops/test_kimi_kda.py | 83 +++++++++++++++++++++++++++++------ vllm_ascend/ops/kimi_kda.py | 45 +++++++++++++------ 2 files changed, 101 insertions(+), 27 deletions(-) diff --git a/tests/ut/ops/test_kimi_kda.py b/tests/ut/ops/test_kimi_kda.py index 845e5bebeff6..26751c461c77 100644 --- a/tests/ut/ops/test_kimi_kda.py +++ b/tests/ut/ops/test_kimi_kda.py @@ -135,6 +135,24 @@ def test_prepare_beta_slices_and_applies_sigmoid_in_fp32(): assert torch.all((beta >= 0.0) & (beta <= 1.0)) +def test_prepare_beta_does_not_repeat_auxiliary_sigmoid(): + raw_beta = torch.tensor( + [[[-20.0], [0.0], [20.0], [100.0]]], + dtype=torch.bfloat16, + ) + preprocessed_beta = raw_beta.float().sigmoid() + + beta = _prepare_beta( + preprocessed_beta, + num_actual_tokens=3, + is_preprocessed=True, + ) + + assert beta.dtype == torch.float32 + assert beta.shape == (1, 3, 1) + torch.testing.assert_close(beta, preprocessed_beta[:, :3]) + + @pytest.mark.parametrize("f_b_is_local", [False, True]) def test_fused_bfg_linear_composes_f_and_packs_bfg(f_b_is_local: bool): with ( @@ -211,36 +229,75 @@ def test_fused_bfg_projection_preserves_staged_outputs(): attention.head_dim = 3 attention._fused_bfg_output_sizes = (2, 6, 6) hidden_states = torch.randn(4, 5) - fused_output = torch.arange(56, dtype=torch.float32).reshape(4, 14) + fused_output = torch.arange(56, dtype=torch.float32).reshape(4, 14).to(torch.bfloat16) attention.fused_bfg_proj = _RecordingLinear(fused_output) - beta, raw_gate, output_gate = attention._project_bfg(hidden_states) - torch.testing.assert_close(beta, fused_output[:, :2]) - torch.testing.assert_close(raw_gate, fused_output[:, 2:8]) - torch.testing.assert_close(output_gate, fused_output[:, 8:]) + projected_bfg = attention._project_bfg(hidden_states) + assert projected_bfg is fused_output - beta, raw_gate, output_gate = attention._postprocess_bfg(beta, raw_gate, output_gate) - torch.testing.assert_close(beta, fused_output[:, :2].unsqueeze(0)) + beta, raw_gate, output_gate = attention._postprocess_bfg(projected_bfg) + assert beta.dtype == torch.float32 + torch.testing.assert_close(beta, fused_output[:, :2].float().sigmoid().unsqueeze(0)) torch.testing.assert_close(raw_gate, fused_output[:, 2:8].reshape(4, 2, 3).unsqueeze(0)) torch.testing.assert_close(output_gate, fused_output[:, 8:].reshape(4, 2, 3)) -def test_overlapped_qkv_bfg_joins_auxiliary_tail_for_graph_capture(): +def test_mixed_forward_marks_auxiliary_beta_as_preprocessed(): + attention = AscendKimiK3DeltaAttention.__new__(AscendKimiK3DeltaAttention) + nn.Module.__init__(attention) + attention.uses_mixed_projection = True + attention.local_num_heads = 2 + attention.head_dim = 3 + hidden_states = torch.randn(4, 6) + positions = torch.arange(4) + mixed_qkv = torch.randn(4, 18) + beta = torch.rand(1, 4, 2, dtype=torch.float32) + raw_gate = torch.randn(1, 4, 2, 3) + output_gate = torch.randn(4, 2, 3) + projected = torch.randn(4, 6) + attention._run_overlapped_qkv_bfg = MagicMock(return_value=(mixed_qkv, beta, raw_gate, output_gate)) + attention._forward = MagicMock() + attention.o_proj = _RecordingLinear(projected) + + actual = attention.forward(hidden_states, positions) + + assert actual is projected + assert attention._forward.call_args.kwargs["beta"] is beta + assert attention._forward.call_args.kwargs["beta_is_preprocessed"] is True + + +def test_overlapped_qkv_bfg_keeps_two_stage_vector_cube_overlap(): attention = AscendKimiK3DeltaAttention.__new__(AscendKimiK3DeltaAttention) nn.Module.__init__(attention) trace: list[str] = [] main_stream = _RecordingStream("main", ["hidden_ready", "quant_ready"], trace) bfg_stream = _RecordingStream("bfg", ["bfg_projection_ready", "bfg_ready"], trace) hidden_states = _RecordingTensor("hidden", trace) - raw_bfg = tuple(_RecordingTensor(name, trace) for name in ("beta_raw", "raw_gate_raw", "gate_raw")) + fused_bfg = _RecordingTensor("fused_bfg", trace) processed_bfg = tuple(_RecordingTensor(name, trace) for name in ("beta", "raw_gate", "output_gate")) quantized_qkv = object() qkv = object() - attention._project_bfg = MagicMock(side_effect=lambda _: (trace.append("project_bfg"), raw_bfg)[1]) - attention._quantize_fused_qkv = MagicMock(side_effect=lambda _: (trace.append("dynamic_quant"), quantized_qkv)[1]) - attention._matmul_fused_qkv = MagicMock(side_effect=lambda _: (trace.append("qkv_matmul"), qkv)[1]) - attention._postprocess_bfg = MagicMock(side_effect=lambda *_: (trace.append("postprocess_bfg"), processed_bfg)[1]) + def record_project_bfg(_hidden_states: object) -> _RecordingTensor: + trace.append("project_bfg") + return fused_bfg + + def record_dynamic_quant(_hidden_states: object) -> object: + trace.append("dynamic_quant") + return quantized_qkv + + def record_qkv_matmul(_qkv_input: object) -> object: + trace.append("qkv_matmul") + return qkv + + def record_postprocess_bfg(*_args: object) -> tuple[_RecordingTensor, ...]: + trace.append("postprocess_bfg") + return processed_bfg + + attention._project_bfg = MagicMock(side_effect=record_project_bfg) + attention._quantize_fused_qkv = MagicMock(side_effect=record_dynamic_quant) + attention._matmul_fused_qkv = MagicMock(side_effect=record_qkv_matmul) + attention._postprocess_bfg = MagicMock(side_effect=record_postprocess_bfg) with ( patch("vllm_ascend.ops.kimi_kda.torch.npu.current_stream", return_value=main_stream), diff --git a/vllm_ascend/ops/kimi_kda.py b/vllm_ascend/ops/kimi_kda.py index 602245801a38..4d10518d9ba0 100644 --- a/vllm_ascend/ops/kimi_kda.py +++ b/vllm_ascend/ops/kimi_kda.py @@ -137,7 +137,7 @@ def _maybe_fuse_f_proj(self) -> None: raise ValueError("KDA fused f_proj requires an output-sharded parameter") shard_offset = sum(self.output_sizes[:_F_PROJ_SHARD_ID]) // self.tp_size shard_size = self.output_sizes[_F_PROJ_SHARD_ID] // self.tp_size - param_shard = self.weight.data.narrow(output_dim, shard_offset, shard_size) + param_shard = self.weight.narrow(output_dim, shard_offset, shard_size) fused_weight = torch.matmul( self.f_b_weight.float(), self.f_a_weight.float(), @@ -173,11 +173,14 @@ def _zero_padded_recurrent_output( def _prepare_beta( - raw_beta: torch.Tensor, + beta: torch.Tensor, num_actual_tokens: int, + *, + is_preprocessed: bool = False, ) -> torch.Tensor: - """Convert vLLM 0.27's packed raw beta to the AscendC contract.""" - return raw_beta[:, :num_actual_tokens].float().sigmoid() + """Slice beta and apply sigmoid unless the auxiliary stream already did.""" + beta = beta[:, :num_actual_tokens] + return beta if is_preprocessed else beta.float().sigmoid() class AscendKimiK3DeltaAttention(KimiK3DeltaAttention): @@ -273,6 +276,7 @@ def forward( g2=g2, beta=beta, core_attn_out=core_attn_out, + beta_is_preprocessed=True, ) core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)") return self.o_proj(core_attn_out)[0] @@ -290,18 +294,24 @@ def _run_overlapped_qkv_bfg( hidden_states.record_stream(bfg_stream) with npu_stream_switch(bfg_stream): bfg_stream.wait_event(hidden_states_ready) - raw_bfg = self._project_bfg(hidden_states) + fused_bfg = self._project_bfg(hidden_states) bfg_projection_ready = bfg_stream.record_event() quantized_qkv = self._quantize_fused_qkv(hidden_states) quant_ready = main_stream.record_event() + # Stage 1 join: DynamicQuant on main overlaps the BFG GEMM on the + # auxiliary stream, but the two Cube matmuls remain serialized. main_stream.wait_event(bfg_projection_ready) mixed_qkv = self._matmul_fused_qkv(quantized_qkv) with npu_stream_switch(bfg_stream): + # Stage 2: after both first-stage branches complete, overlap the + # QKV Cube matmul with beta's FP32 conversion and sigmoid vector + # work. Split and reshape the F/output gates here as well so all + # BFG output handling occurs after the QKV matmul is enqueued. bfg_stream.wait_event(quant_ready) - beta, g1, g2 = self._postprocess_bfg(*raw_bfg) + beta, g1, g2 = self._postprocess_bfg(fused_bfg) bfg_ready = bfg_stream.record_event() for tensor in (beta, g1, g2): @@ -314,17 +324,18 @@ def _run_overlapped_qkv_bfg( def _project_bfg( self, hidden_states: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - fused_bfg = self.fused_bfg_proj(hidden_states)[0] - return fused_bfg.split(self._fused_bfg_output_sizes, dim=-1) + ) -> torch.Tensor: + return self.fused_bfg_proj(hidden_states)[0] def _postprocess_bfg( self, - beta: torch.Tensor, - raw_gate: torch.Tensor, - output_gate: torch.Tensor, + fused_bfg: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - beta = beta.unsqueeze(0) + beta, raw_gate, output_gate = fused_bfg.split( + self._fused_bfg_output_sizes, + dim=-1, + ) + beta = beta.float().sigmoid().unsqueeze(0) raw_gate = rearrange(raw_gate, "n (h d) -> 1 n h d", d=self.head_dim) output_gate = rearrange(output_gate, "n (h d) -> n h d", d=self.head_dim) return beta, raw_gate, output_gate @@ -516,6 +527,8 @@ def _forward( g2: torch.Tensor, beta: torch.Tensor, core_attn_out: torch.Tensor, + *, + beta_is_preprocessed: bool = False, ) -> None: """Dispatch speculative, prefill, and decode tokens through KDA kernels.""" forward_context = get_forward_context() @@ -532,7 +545,11 @@ def _forward( mixed_qkv = mixed_qkv[:num_actual_tokens] g1 = g1[:, :num_actual_tokens] g2 = g2[:num_actual_tokens] - beta = _prepare_beta(beta, num_actual_tokens) + beta = _prepare_beta( + beta, + num_actual_tokens, + is_preprocessed=beta_is_preprocessed, + ) conv_state, recurrent_state = self.kv_cache conv_weights_t = self.get_parameter(_PACKED_CONV_WEIGHT_NAME) From 34776f35c15ec68cb9599bdfc71a482ce9a02c9d Mon Sep 17 00:00:00 2001 From: Dawn952 Date: Mon, 31 Aug 2026 21:53:18 +0800 Subject: [PATCH 3/6] fix(ci): sort KDA test imports Signed-off-by: Dawn952 (cherry picked from commit 7513b78278e4aa1f43da6a9082f18f50884235da) Signed-off-by: Dawn952 --- tests/ut/ops/test_kimi_kda.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/tests/ut/ops/test_kimi_kda.py b/tests/ut/ops/test_kimi_kda.py index 26751c461c77..f1b84ec81fe0 100644 --- a/tests/ut/ops/test_kimi_kda.py +++ b/tests/ut/ops/test_kimi_kda.py @@ -8,6 +8,13 @@ import torch from torch import nn +from vllm_ascend.quantization.methods.w4a8_mxfp4 import ( + AscendW4A8MXFPDynamicLinearMethod, +) +from vllm_ascend.quantization.methods.w8a8_mxfp8 import ( + AscendW8A8MXFP8DynamicLinearMethod, +) + from vllm_ascend.ops.kimi_kda import ( _PACKED_CONV_WEIGHT_NAME, AscendKimiK3DeltaAttention, @@ -16,12 +23,6 @@ _zero_padded_output, _zero_padded_recurrent_output, ) -from vllm_ascend.quantization.methods.w4a8_mxfp4 import ( - AscendW4A8MXFPDynamicLinearMethod, -) -from vllm_ascend.quantization.methods.w8a8_mxfp8 import ( - AscendW8A8MXFP8DynamicLinearMethod, -) class _RecordingLinear(nn.Module): From 483a97dfff2e4106c89c14ba4bea6eb5f67eb576 Mon Sep 17 00:00:00 2001 From: Dawn952 Date: Tue, 1 Sep 2026 10:59:50 +0800 Subject: [PATCH 4/6] fix(ci): follow quantization package layout Signed-off-by: Dawn952 (cherry picked from commit 0059617c5d9e85c4ee2185110ee2cec950698533) Signed-off-by: Dawn952 --- tests/ut/ops/test_kimi_kda.py | 13 ++++++------- vllm_ascend/ops/kimi_kda.py | 4 ++-- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/tests/ut/ops/test_kimi_kda.py b/tests/ut/ops/test_kimi_kda.py index f1b84ec81fe0..7859f91603a8 100644 --- a/tests/ut/ops/test_kimi_kda.py +++ b/tests/ut/ops/test_kimi_kda.py @@ -8,13 +8,6 @@ import torch from torch import nn -from vllm_ascend.quantization.methods.w4a8_mxfp4 import ( - AscendW4A8MXFPDynamicLinearMethod, -) -from vllm_ascend.quantization.methods.w8a8_mxfp8 import ( - AscendW8A8MXFP8DynamicLinearMethod, -) - from vllm_ascend.ops.kimi_kda import ( _PACKED_CONV_WEIGHT_NAME, AscendKimiK3DeltaAttention, @@ -23,6 +16,12 @@ _zero_padded_output, _zero_padded_recurrent_output, ) +from vllm_ascend.quantization.methods.w4a8.w4a8_mxfp4 import ( + AscendW4A8MXFPDynamicLinearMethod, +) +from vllm_ascend.quantization.methods.w8a8.w8a8_mxfp8 import ( + AscendW8A8MXFP8DynamicLinearMethod, +) class _RecordingLinear(nn.Module): diff --git a/vllm_ascend/ops/kimi_kda.py b/vllm_ascend/ops/kimi_kda.py index 4d10518d9ba0..3ce4ded20931 100644 --- a/vllm_ascend/ops/kimi_kda.py +++ b/vllm_ascend/ops/kimi_kda.py @@ -29,10 +29,10 @@ from vllm_ascend.ops.gdn_attn_builder import AscendGDNAttentionBackend from vllm_ascend.ops.triton.fla.utils import clear_ssm_states -from vllm_ascend.quantization.methods.w4a8_mxfp4 import ( +from vllm_ascend.quantization.methods.w4a8.w4a8_mxfp4 import ( AscendW4A8MXFPDynamicLinearMethod, ) -from vllm_ascend.quantization.methods.w8a8_mxfp8 import ( +from vllm_ascend.quantization.methods.w8a8.w8a8_mxfp8 import ( AscendW8A8MXFP8DynamicLinearMethod, ) from vllm_ascend.utils import npu_stream_switch From 5b1753ed6774ae52d7dbfb2a95db758ba73e6e4e Mon Sep 17 00:00:00 2001 From: linfeng-yuan <1102311262@qq.com> Date: Mon, 7 Sep 2026 21:47:45 +0800 Subject: [PATCH 5/6] cut v0.27.1 CI test Signed-off-by: linfeng-yuan <1102311262@qq.com> (cherry picked from commit da4846002da04346772c12774418bc2cc8d40ed4) Signed-off-by: Dawn952 --- .github/workflows/pr_test.yaml | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/.github/workflows/pr_test.yaml b/.github/workflows/pr_test.yaml index 1e07621481b4..44604b9917df 100644 --- a/.github/workflows/pr_test.yaml +++ b/.github/workflows/pr_test.yaml @@ -46,7 +46,7 @@ jobs: container: image: quay.io/ascend-ci/vllm-ascend:lint outputs: - main_commit: ${{ steps.vllm.outputs.main_commit }} + release_tag: ${{ steps.vllm.outputs.release_tag }} steps: - name: Validate PR title prefix run: | @@ -84,11 +84,13 @@ jobs: id: vllm run: | main_commit="$(tr -d '[:space:]' < .github/vllm-main-verified.commit)" + release_tag="$(tr -d '[:space:]' < .github/vllm-release-tag.commit)" [[ "${main_commit}" =~ ^[0-9a-f]{7,40}$ ]] || { echo "::error file=.github/vllm-main-verified.commit::invalid vLLM main commit: ${main_commit}" exit 1 } echo "main_commit=${main_commit}" >> "$GITHUB_OUTPUT" + echo "release_tag=${release_tag}" >> "$GITHUB_OUTPUT" - name: cp problem matchers run: | @@ -516,7 +518,7 @@ jobs: fail-fast: false matrix: vllm_version: - - ${{ needs.pre-commit.outputs.main_commit }} + - ${{ needs.pre-commit.outputs.release_tag }} uses: ./.github/workflows/_selected_tests.yaml with: vllm: ${{ matrix.vllm_version }} @@ -543,7 +545,7 @@ jobs: strategy: matrix: vllm_version: - - ${{ needs.pre-commit.outputs.main_commit }} + - ${{ needs.pre-commit.outputs.release_tag }} uses: ./.github/workflows/_selected_tests_upstream.yaml with: vllm: ${{ matrix.vllm_version }} @@ -567,7 +569,7 @@ jobs: strategy: matrix: vllm_version: - - ${{ needs.pre-commit.outputs.main_commit }} + - ${{ needs.pre-commit.outputs.release_tag }} uses: ./.github/workflows/_selected_tests.yaml with: vllm: ${{ matrix.vllm_version }} From 64c51dde55c5a5dd4203810ea529560497e45eee Mon Sep 17 00:00:00 2001 From: Dawn952 Date: Tue, 8 Sep 2026 06:12:08 +0800 Subject: [PATCH 6/6] fix(platform): support v0.27.1 attention selector config Signed-off-by: Dawn952 --- tests/ut/test_platform.py | 31 ++----------------------------- vllm_ascend/platform.py | 2 +- 2 files changed, 3 insertions(+), 30 deletions(-) diff --git a/tests/ut/test_platform.py b/tests/ut/test_platform.py index cceee46c86a1..3e5ff0060763 100644 --- a/tests/ut/test_platform.py +++ b/tests/ut/test_platform.py @@ -1792,30 +1792,16 @@ def test_get_attn_backend_cls_selects_cp_backend(self, mock_get_ascend_config): ( True, True, - False, "vllm_ascend.attention.mla_v1.AscendMLABackend", ), ( - False, - True, - False, - "vllm_ascend.attention.attention_v1.AscendAttentionBackend", - ), - ( - True, - False, - True, - "vllm_ascend.attention.mla_v1.AscendMLABackend", - ), - ( - False, False, True, "vllm_ascend.attention.attention_v1.AscendAttentionBackend", ), ) - for use_mla, use_pcp, use_dcp, expected_backend in cases: - with self.subTest(use_mla=use_mla, use_pcp=use_pcp, use_dcp=use_dcp): + for use_mla, use_pcp, expected_backend in cases: + with self.subTest(use_mla=use_mla, use_pcp=use_pcp): attn_selector_config = AttentionSelectorConfig( dtype=torch.float16, head_size=0, @@ -1824,23 +1810,10 @@ def test_get_attn_backend_cls_selects_cp_backend(self, mock_get_ascend_config): use_mla=use_mla, use_sparse=False, use_pcp=use_pcp, - use_dcp=use_dcp, ) result = self.platform.get_attn_backend_cls("ascend", attn_selector_config) self.assertEqual(result, expected_backend) - def test_get_attn_backend_cls_rejects_pcp_and_dcp(self): - attn_selector_config = AttentionSelectorConfig( - dtype=torch.float16, - head_size=0, - kv_cache_dtype=None, - block_size=128, - use_pcp=True, - use_dcp=True, - ) - with self.assertRaisesRegex(NotImplementedError, "does not support PCP and DCP simultaneously"): - self.platform.get_attn_backend_cls("ascend", attn_selector_config) - def test_get_attn_backend_cls_selects_sfa_pcp_backend(self): attn_selector_config = AttentionSelectorConfig( dtype=torch.float16, diff --git a/vllm_ascend/platform.py b/vllm_ascend/platform.py index 786bd2a458cb..e5b3a1bc11f5 100644 --- a/vllm_ascend/platform.py +++ b/vllm_ascend/platform.py @@ -235,7 +235,7 @@ def get_attn_backend_cls(cls, selected_backend, attn_selector_config, num_heads: key = (attn_selector_config.use_mla, attn_selector_config.use_sparse) backend_key = (*key, use_compress) - if attn_selector_config.use_pcp and attn_selector_config.use_dcp: + if attn_selector_config.use_pcp and getattr(attn_selector_config, "use_dcp", False): raise NotImplementedError("Ascend MRV2 does not support PCP and DCP simultaneously yet.") if not attn_selector_config.use_pcp and _validate_fa3_backend(key, attn_selector_config):