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..7859f91603a8 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.w4a8_mxfp4 import ( + AscendW4A8MXFPDynamicLinearMethod, +) +from vllm_ascend.quantization.methods.w8a8.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,266 @@ 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 ( + 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).to(torch.bfloat16) + attention.fused_bfg_proj = _RecordingLinear(fused_output) + + projected_bfg = attention._project_bfg(hidden_states) + assert projected_bfg is fused_output + + 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_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) + 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() + + 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), + 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..3ce4ded20931 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.w4a8_mxfp4 import ( + AscendW4A8MXFPDynamicLinearMethod, +) +from vllm_ascend.quantization.methods.w8a8.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.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( @@ -57,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): @@ -81,9 +200,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 +210,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 +264,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, @@ -175,11 +276,112 @@ 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] 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) + 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(fused_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, + ) -> torch.Tensor: + return self.fused_bfg_proj(hidden_states)[0] + + def _postprocess_bfg( + self, + fused_bfg: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + 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 + + 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, @@ -325,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() @@ -341,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)