diff --git a/.github/workflows/scripts/test_config.yaml b/.github/workflows/scripts/test_config.yaml index daacd3bf75c1..21af35ea9ac2 100644 --- a/.github/workflows/scripts/test_config.yaml +++ b/.github/workflows/scripts/test_config.yaml @@ -391,6 +391,7 @@ source_file_dependencies: - vllm_ascend/ops/gdn.py - vllm_ascend/ops/gdn_attn_builder.py + - vllm_ascend/ops/kimi_kda.py - vllm_ascend/ops/triton/gdn_chunk_meta.py tests: - tests/ut/ops @@ -913,6 +914,7 @@ estimated_times: tests/ut/ops/a2/test_gdn_layerwise_kv.py: 70 tests/ut/ops/a2/test_token_dispatcher.py: 30 tests/ut/ops/a3_2/test_activation.py: 50 + tests/ut/ops/a3_2/test_kimi_kda_fused_norm_gate.py: 60 tests/ut/ops/a3_2/test_select_experts.py: 20 tests/ut/quantization/methods/a2/test_w4a16.py: 30 tests/ut/quantization/methods/a2/test_w4a4_flatquant.py: 40 diff --git a/tests/ut/ops/a3_2/test_kimi_kda_fused_norm_gate.py b/tests/ut/ops/a3_2/test_kimi_kda_fused_norm_gate.py new file mode 100644 index 000000000000..c9a46fd8f24c --- /dev/null +++ b/tests/ut/ops/a3_2/test_kimi_kda_fused_norm_gate.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm_ascend.ops.triton.kda.kda import rms_norm_gated + + +@torch.inference_mode() +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("tokens, heads, head_dim", [(1, 1, 128), (16, 4, 128), (37, 4, 128), (2, 1, 1024)]) +@pytest.mark.parametrize("strided_gate", [False, True]) +def test_kimi_kda_fused_rms_norm_sigmoid_gate(dtype, tokens, heads, head_dim, strided_gate): + torch.manual_seed(20260801) + eps = 1e-6 + weight = torch.randn(head_dim, dtype=dtype, device="npu") + core_attn_out = torch.randn(1, tokens, heads, head_dim, dtype=dtype, device="npu") + output_gate = torch.randn(tokens, heads, head_dim, dtype=dtype, device="npu") + if strided_gate: + # K3's packed projection leaves gaps between consecutive gate rows. + packed_gate = torch.full((tokens, 2 * heads, head_dim), torch.nan, dtype=dtype, device="npu") + packed_gate[:, :heads].copy_(output_gate) + output_gate = packed_gate[:, :heads] + core_attn_out_before = core_attn_out.clone() + output_gate_before = output_gate.clone() + + actual = rms_norm_gated(core_attn_out, output_gate, weight, None, activation="sigmoid", eps=eps) + + x_float = core_attn_out_before.float() + variance = x_float.square().mean(dim=-1, keepdim=True) + expected = x_float * torch.rsqrt(variance + eps) + expected = expected * weight.float() + expected = expected * output_gate.float().sigmoid().unsqueeze(0) + + torch.testing.assert_close(actual, expected.to(dtype), rtol=2e-3, atol=2e-3) + torch.testing.assert_close(core_attn_out, core_attn_out_before, rtol=0, atol=0) + torch.testing.assert_close(output_gate, output_gate_before, rtol=0, atol=0) + + +@torch.inference_mode() +@pytest.mark.parametrize("residual_dtype", [None, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("elementwise_affine", [False, True]) +def test_fused_rms_norm_silu_gate_preserves_prenorm_contract(residual_dtype, elementwise_affine): + torch.manual_seed(20260801) + x = torch.randn(1, 3, 2, 128, dtype=torch.bfloat16, device="npu") + gate = torch.randn_like(x) + residual = torch.randn_like(x, dtype=residual_dtype) if residual_dtype is not None else None + weight = torch.randn(128, dtype=x.dtype, device="npu") if elementwise_affine else None + before = x.clone() + eps = 1e-6 + + actual, residual_out = rms_norm_gated( + x, gate, weight, None, activation="silu", residual=residual, prenorm=True, residual_in_fp32=True, eps=eps + ) + + summed = before.float() if residual is None else before.float() + residual.float() + expected = summed * torch.rsqrt(summed.square().mean(-1, keepdim=True) + eps) + if weight is not None: + expected *= weight.float() + expected *= gate.float() * gate.float().sigmoid() + expected_residual_dtype = torch.float32 if residual is None else residual.dtype + + torch.testing.assert_close(actual, expected.to(x.dtype), rtol=2e-3, atol=2e-3) + assert residual_out.dtype == expected_residual_dtype + torch.testing.assert_close(residual_out, summed.to(expected_residual_dtype), rtol=0, atol=0) + torch.testing.assert_close(x, before, rtol=0, atol=0) diff --git a/tests/ut/ops/test_gdn_attn_builder.py b/tests/ut/ops/test_gdn_attn_builder.py index 909be456a385..0b0a28191b6c 100644 --- a/tests/ut/ops/test_gdn_attn_builder.py +++ b/tests/ut/ops/test_gdn_attn_builder.py @@ -9,6 +9,7 @@ from vllm.config.compilation import CUDAGraphMode from vllm.third_party.flash_linear_attention.ops import index as _fla_index from vllm.v1.attention.backend import CommonAttentionMetadata +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID from vllm.v1.kv_cache_interface import MambaSpec from vllm_ascend.attention.utils import AscendCommonAttentionMetadata @@ -226,7 +227,10 @@ def _build_attn_metadata( def _assert_chunk_meta_matches_runtime(builder, chunk_meta, cu_seqlens: torch.Tensor) -> None: hf_text_config = getattr(builder.vllm_config.model_config, "hf_text_config", None) - if hf_text_config is not None and hasattr(hf_text_config, "linear_num_value_heads"): + linear_attn_config = getattr(hf_text_config, "linear_attn_config", None) + if isinstance(linear_attn_config, dict) and linear_attn_config.get("num_heads") is not None: + gdn_num_heads = linear_attn_config["num_heads"] // builder.vllm_config.parallel_config.tensor_parallel_size + elif hf_text_config is not None and hasattr(hf_text_config, "linear_num_value_heads"): gdn_num_heads = ( hf_text_config.linear_num_value_heads // builder.vllm_config.parallel_config.tensor_parallel_size ) @@ -287,6 +291,26 @@ def _patch_missing_runtime_cdiv(monkeypatch: pytest.MonkeyPatch) -> None: ) +def test_kimi_chunk_metadata_uses_linear_attention_head_count() -> None: + builder = _make_builder( + device=torch.device("cpu"), + num_heads=128, + num_speculative_tokens=0, + ) + builder.vllm_config.model_config.hf_text_config = SimpleNamespace( + linear_attn_config={"num_heads": 32}, + ) + cu_seqlens = torch.tensor([0, 130], dtype=torch.int32) + + chunk_meta = ascend_gdn_attn_builder._build_non_spec_chunked_prefill_metadata( + builder, + cu_seqlens, + torch.device("cpu"), + ) + + _assert_chunk_meta_matches_runtime(builder, chunk_meta, cu_seqlens) + + def test_ascend_gdn_attention_uses_ascend_backend(): assert AscendGatedDeltaNetAttention.get_attn_backend(object()) is AscendGDNAttentionBackend assert AscendGDNAttentionBackend.get_builder_cls() is AscendGDNAttentionMetadataBuilder @@ -431,7 +455,6 @@ def test_non_spec_prefill_metadata_uses_prefill_tail_for_chunk_metadata( assert torch.equal(conv1d_meta.query_start_loc, torch.tensor([0, 1, 9, 13], dtype=torch.int32)) assert torch.equal(_cache_index_first_column(conv1d_meta.cache_indices), torch.tensor([0, 1, 2], dtype=torch.int32)) assert torch.equal(conv1d_meta.initial_state_mode, torch.tensor([True, True, True])) - assert prefill_metadata.chunk.num_decodes == 0 _assert_chunk_meta_matches_runtime( builder, prefill_metadata.chunk, @@ -570,8 +593,8 @@ def test_full_graph_spec_actual_seq_lengths_use_padded_builder_buffer(): def test_full_graph_non_spec_actual_seq_lengths_use_padded_builder_buffer(): batch_spec = BatchSpec( - seq_lens=[1, 1], - query_lens=[1, 1], + seq_lens=[1, 1, 0, 0], + query_lens=[1, 1, 0, 0], name="full_graph_padded_non_spec_actual_seq_lengths", ) common_attn_metadata = create_common_attn_metadata( @@ -579,7 +602,6 @@ def test_full_graph_non_spec_actual_seq_lengths_use_padded_builder_buffer(): block_size=16, device=torch.device("cpu"), ) - common_attn_metadata.num_actual_tokens = 4 builder = _make_builder( device=torch.device("cpu"), num_heads=32, @@ -759,3 +781,269 @@ def test_builder_skips_prebuilt_meta_without_non_spec_prefill(batch_spec: BatchS spec_decode_metadata.actual_seq_lengths, torch.tensor([0, 4, 4], dtype=torch.int32), ) + + +def test_mixed_spec_prefill_chunk_metadata_preserves_single_token_count( + monkeypatch: pytest.MonkeyPatch, +): + _patch_missing_runtime_cdiv(monkeypatch) + batch_spec = BatchSpec( + seq_lens=[1, 4, 8], + query_lens=[1, 4, 8], + name="mixed_spec_prefill_with_single_token_non_spec", + ) + builder, _, attn_metadata = _build_attn_metadata( + batch_spec, + num_speculative_tokens=3, + num_decode_draft_tokens_cpu=torch.tensor([-1, 3, -1], dtype=torch.int32), + ) + + assert attn_metadata.num_decodes == 0 + assert attn_metadata.num_prefills == 2 + assert torch.equal( + attn_metadata.prefill_query_start_loc, + torch.tensor([0, 1, 9], dtype=torch.int32), + ) + chunk_metadata = attn_metadata.non_spec_prefill_metadata.chunk + _assert_chunk_meta_matches_runtime( + builder, + chunk_metadata, + attn_metadata.prefill_query_start_loc, + ) + + +@pytest.mark.parametrize( + ("seq_len", "expected_decodes", "expected_prefills"), + [ + pytest.param(1, 0, 1, id="first_token_stays_prefill"), + pytest.param(17, 1, 0, id="block-size-plus-one-becomes-decode"), + ], +) +def test_one_token_prefill_selection_respects_recurrent_state( + monkeypatch: pytest.MonkeyPatch, + seq_len: int, + expected_decodes: int, + expected_prefills: int, +): + _patch_missing_runtime_cdiv(monkeypatch) + common_attn_metadata = create_common_attn_metadata( + BatchSpec(seq_lens=[seq_len], query_lens=[1]), + block_size=16, + device=torch.device("cpu"), + ) + # Model a prompt chunk explicitly. The helper normally classifies a + # one-token row as decode when synthesizing test metadata. + common_attn_metadata.is_prefilling = torch.tensor([True]) + builder = _make_builder( + device=torch.device("cpu"), + num_heads=32, + num_speculative_tokens=0, + ) + + attn_metadata = builder.build(0, common_attn_metadata) + + assert common_attn_metadata.is_prefilling.tolist() == [True] + assert attn_metadata.num_decodes == expected_decodes + assert attn_metadata.num_prefills == expected_prefills + + +@pytest.mark.parametrize( + ("seq_len", "expected_spec_decodes", "expected_prefills"), + [ + pytest.param(4, 0, 1, id="first_chunk_stays_prefill"), + pytest.param(8, 1, 0, id="stateful_chunk_folds_into_spec"), + ], +) +def test_spec_sized_prefill_fold_requires_recurrent_state( + monkeypatch: pytest.MonkeyPatch, + seq_len: int, + expected_spec_decodes: int, + expected_prefills: int, +): + _patch_missing_runtime_cdiv(monkeypatch) + common_attn_metadata = create_common_attn_metadata( + BatchSpec(seq_lens=[seq_len], query_lens=[4]), + block_size=16, + device=torch.device("cpu"), + ) + builder = _make_builder( + device=torch.device("cpu"), + num_heads=32, + num_speculative_tokens=3, + ) + + attn_metadata = builder.build( + 0, + common_attn_metadata, + num_accepted_tokens=torch.ones(1, dtype=torch.int32), + num_decode_draft_tokens_cpu=torch.full((1,), -1, dtype=torch.int32), + ) + + assert attn_metadata.num_spec_decodes == expected_spec_decodes + assert attn_metadata.num_prefills == expected_prefills + if expected_spec_decodes: + assert attn_metadata.spec_sequence_masks.tolist() == [True] + assert attn_metadata.num_accepted_tokens.tolist() == [4] + else: + assert attn_metadata.spec_sequence_masks is None + assert attn_metadata.num_accepted_tokens is None + + +def test_full_graph_without_runtime_spec_resets_captured_spec_inputs(): + capture_common_metadata = create_common_attn_metadata( + batch_spec=BatchSpec( + seq_lens=[4, 4], + query_lens=[4, 4], + name="full_graph_spec_capture", + ), + block_size=16, + device=torch.device("cpu"), + ) + capture_common_metadata.num_reqs = 4 + capture_common_metadata.block_table_tensor = torch.tensor( + [[10, 11, 12, 13], [20, 21, 22, 23]], + dtype=torch.int32, + ) + builder = _make_builder( + device=torch.device("cpu"), + num_heads=32, + num_speculative_tokens=3, + cudagraph_mode=CUDAGraphMode.FULL_DECODE_ONLY, + ) + captured_metadata = builder.build( + 0, + capture_common_metadata, + num_accepted_tokens=torch.tensor([2, 4], dtype=torch.int32), + num_decode_draft_tokens_cpu=torch.tensor([3, 3], dtype=torch.int32), + ) + captured_spec_metadata = captured_metadata.spec_decode_metadata + captured_conv1d_metadata = captured_spec_metadata.spec_causal_conv1d + + assert torch.count_nonzero(captured_conv1d_metadata.query_start_loc) > 0 + assert torch.count_nonzero(captured_spec_metadata.actual_seq_lengths) > 0 + + replay_common_metadata = create_common_attn_metadata( + batch_spec=BatchSpec( + seq_lens=[1, 1, 0, 0], + query_lens=[1, 1, 0, 0], + name="full_graph_replay_without_spec", + ), + block_size=16, + device=torch.device("cpu"), + ) + replay_metadata = builder.build( + 0, + replay_common_metadata, + num_accepted_tokens=torch.ones(4, dtype=torch.int32), + num_decode_draft_tokens_cpu=torch.full((4,), -1, dtype=torch.int32), + ) + + assert replay_metadata.spec_sequence_masks is None + assert replay_metadata.spec_decode_metadata is None + assert torch.equal( + captured_conv1d_metadata.cache_indices, + torch.full((4, 4), PAD_SLOT_ID, dtype=torch.int32), + ) + assert torch.count_nonzero(captured_conv1d_metadata.query_start_loc) == 0 + assert torch.count_nonzero(captured_conv1d_metadata.num_accepted_tokens) == 0 + assert torch.count_nonzero(captured_spec_metadata.actual_seq_lengths) == 0 + + +def test_full_graph_idle_dummy_uses_zero_length_recurrent_metadata(): + common_attn_metadata = create_common_attn_metadata( + batch_spec=BatchSpec( + seq_lens=[8, 8, 8, 8], + query_lens=[0, 0, 0, 0], + name="full_graph_idle_dummy", + ), + block_size=16, + device=torch.device("cpu"), + ) + common_attn_metadata.block_table_tensor[:, 0] = torch.tensor( + [10, 11, 98, 99], + dtype=torch.int32, + ) + builder = _make_builder( + device=torch.device("cpu"), + num_heads=32, + num_speculative_tokens=3, + cudagraph_mode=CUDAGraphMode.FULL_DECODE_ONLY, + ) + builder.spec_state_indices_tensor.fill_(77) + builder.spec_query_start_loc.fill_(77) + builder.non_spec_state_indices_tensor.fill_(77) + builder.non_spec_query_start_loc.fill_(77) + + attn_metadata = builder.build(0, common_attn_metadata) + + assert attn_metadata.num_actual_tokens == 0 + assert attn_metadata.num_decode_tokens == 0 + assert torch.count_nonzero(attn_metadata.non_spec_query_start_loc) == 0 + assert torch.all(attn_metadata.non_spec_state_indices_tensor == NULL_BLOCK_ID) + assert torch.count_nonzero(builder.spec_query_start_loc[:5]) == 0 + assert torch.all(builder.spec_state_indices_tensor[:4] == PAD_SLOT_ID) + + +@pytest.mark.parametrize( + ("num_speculative_tokens", "num_decode_draft_tokens_cpu"), + [ + pytest.param(0, None, id="without_spec_decode"), + pytest.param( + 3, + torch.full((4,), -1, dtype=torch.int32), + id="spec_decode_without_runtime_spec_requests", + ), + ], +) +def test_full_graph_non_spec_metadata_nulls_padded_state_indices( + num_speculative_tokens: int, + num_decode_draft_tokens_cpu: torch.Tensor | None, +): + common_attn_metadata = create_common_attn_metadata( + batch_spec=BatchSpec( + seq_lens=[1, 1, 0, 0], + query_lens=[1, 1, 0, 0], + name="full_graph_padded_non_spec_actual_seq_lengths", + ), + block_size=16, + device=torch.device("cpu"), + ) + common_attn_metadata.block_table_tensor[:, 0] = torch.tensor([10, 11, 98, 99]) + builder = _make_builder( + device=torch.device("cpu"), + num_heads=32, + num_speculative_tokens=num_speculative_tokens, + cudagraph_mode=CUDAGraphMode.FULL_DECODE_ONLY, + ) + builder.non_spec_state_indices_tensor.fill_(77) + builder.non_spec_query_start_loc.fill_(77) + builder.non_spec_actual_seq_lengths.fill_(77) + + attn_metadata = builder.build( + 0, + common_attn_metadata, + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, + ) + + assert attn_metadata.num_decodes == 4 + assert attn_metadata.num_decode_tokens == 2 + assert torch.equal( + attn_metadata.non_spec_query_start_loc, + torch.tensor([0, 1, 2, 2, 2], dtype=torch.int32), + ) + assert torch.equal( + attn_metadata.non_spec_state_indices_tensor, + torch.tensor( + [10, 11, NULL_BLOCK_ID, NULL_BLOCK_ID], + dtype=torch.int32, + ), + ) + decode_metadata = attn_metadata.non_spec_decode_metadata + conv1d_metadata = decode_metadata.causal_conv1d + assert conv1d_metadata.query_start_loc.data_ptr() == attn_metadata.non_spec_query_start_loc.data_ptr() + assert conv1d_metadata.cache_indices.data_ptr() == attn_metadata.non_spec_state_indices_tensor.data_ptr() + assert decode_metadata.actual_seq_lengths.data_ptr() == builder.non_spec_actual_seq_lengths.data_ptr() + assert torch.equal( + decode_metadata.actual_seq_lengths, + torch.tensor([0, 1, 1, 0, 0], dtype=torch.int32), + ) diff --git a/tests/ut/ops/test_kimi_kda.py b/tests/ut/ops/test_kimi_kda.py new file mode 100644 index 000000000000..e9171348fcf5 --- /dev/null +++ b/tests/ut/ops/test_kimi_kda.py @@ -0,0 +1,245 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import patch + +import torch +from torch import nn + +from vllm_ascend.ops.kimi_kda import ( + _PACKED_CONV_WEIGHT_NAME, + AscendKimiK3DeltaAttention, + _prepare_beta, + _zero_padded_output, + _zero_padded_recurrent_output, +) + + +def test_zero_padded_recurrent_output_clears_uncovered_tail(): + output = torch.randn(1, 8, 2, 3) + expected = output[:, :5].clone() + output[:, 5:] = torch.nan + + actual = _zero_padded_recurrent_output( + output, + torch.tensor([0, 3, 5, 5], dtype=torch.int32), + ) + + torch.testing.assert_close(actual[:, :5], expected) + assert torch.equal(actual[:, 5:], torch.zeros_like(actual[:, 5:])) + assert torch.isfinite(actual).all() + + +def test_zero_padded_output_uses_combined_live_token_count(): + output = torch.full((1, 8, 1, 1), torch.nan) + output[:, :6] = torch.arange(6).view(1, 6, 1, 1) + + actual = _zero_padded_output(output, torch.tensor(6, dtype=torch.int32)) + + torch.testing.assert_close(actual[:, :6], output[:, :6]) + assert torch.equal(actual[:, 6:], torch.zeros_like(actual[:, 6:])) + + +def test_run_causal_conv1d_returns_declared_output_alias(): + mixed_qkv = torch.randn(3, 8) + conv_weights = torch.randn(4, 8) + conv_state = torch.randn(2, 8, 4) + query_start_loc = torch.tensor([0, 3], dtype=torch.int32) + cache_indices = torch.tensor([1], dtype=torch.int32) + returned_alias = torch.full_like(mixed_qkv, 7) + + with patch.object( + torch.ops._C_ascend, + "npu_causal_conv1d_custom", + return_value=returned_alias, + create=True, + ) as causal_conv: + actual = AscendKimiK3DeltaAttention._run_causal_conv1d( + mixed_qkv, + conv_weights, + conv_state, + query_start_loc, + cache_indices, + None, + run_mode=1, + num_accepted_tokens=torch.tensor([3], dtype=torch.int32), + ) + + assert actual is returned_alias + assert causal_conv.call_args.kwargs["query_start_loc_opt"] is query_start_loc + assert causal_conv.call_args.kwargs["cache_indices_opt"] is cache_indices + assert causal_conv.call_args.kwargs["initial_state_mode_opt"] is None + + +def test_kda_output_norm_uses_checkpoint_epsilon(): + def fake_upstream_init(attention, _config, _vllm_config, _prefix): + nn.Module.__init__(attention) + attention.o_norm = SimpleNamespace(eps=1e-5) + attention.conv_size = 4 + attention.local_projection_size = 2 + attention.model_config = SimpleNamespace(dtype=torch.bfloat16) + attention.conv1d = nn.Module() + attention.conv1d.weight = nn.Parameter(torch.empty(6, 1, 4)) + attention.conv1d.quant_method = SimpleNamespace(process_weights_after_loading=lambda: None) + + config = SimpleNamespace(rms_norm_eps=1e-6) + vllm_config = SimpleNamespace( + model_config=SimpleNamespace( + multimodal_config=None, + enable_prompt_embeds=False, + ) + ) + with patch( + "vllm_ascend.ops.kimi_kda.KimiK3DeltaAttention.__init__", + new=fake_upstream_init, + ): + attention = AscendKimiK3DeltaAttention(config, vllm_config) + + assert attention.o_norm.eps == config.rms_norm_eps + + +def test_prepare_beta_slices_and_applies_sigmoid_in_fp32(): + raw_beta = torch.tensor( + [[[-20.0], [0.0], [20.0], [100.0]]], + dtype=torch.bfloat16, + ) + + beta = _prepare_beta(raw_beta, num_actual_tokens=3) + + assert beta.dtype == torch.float32 + assert beta.shape == (1, 3, 1) + torch.testing.assert_close(beta, raw_beta[:, :3].float().sigmoid()) + assert torch.all((beta >= 0.0) & (beta <= 1.0)) + + +def test_recurrent_gate_uses_unbounded_kda_transform(): + attention = AscendKimiK3DeltaAttention.__new__(AscendKimiK3DeltaAttention) + nn.Module.__init__(attention) + attention.head_dim = 3 + attention.A_log = nn.Parameter(torch.randn(2)) + attention.dt_bias = nn.Parameter(torch.randn(6)) + raw_gate = torch.randn(1, 4, 2, 3) + expected = torch.randn(4, 2, 3) + + with patch( + "vllm_ascend.ops.kimi_kda.fused_kda_gate", + return_value=expected, + ) as fused_gate: + actual = attention._recurrent_gate(raw_gate) + + torch.testing.assert_close(actual, expected.unsqueeze(0)) + fused_gate.assert_called_once() + torch.testing.assert_close(fused_gate.call_args.args[0], raw_gate.reshape(4, 6)) + assert fused_gate.call_args.args[1] is attention.A_log + assert fused_gate.call_args.args[2] == attention.head_dim + assert fused_gate.call_args.kwargs["g_bias"] is attention.dt_bias + + +def test_prefill_accepts_unbounded_gate(): + attention = AscendKimiK3DeltaAttention.__new__(AscendKimiK3DeltaAttention) + nn.Module.__init__(attention) + attention.head_dim = 2 + attention.gate_lower_bound = None + attention.A_log = nn.Parameter(torch.randn(1)) + attention.dt_bias = nn.Parameter(torch.randn(2)) + + q = torch.randn(1, 2, 1, 2) + k = torch.randn_like(q) + v = torch.randn_like(q) + raw_gate = torch.randn_like(q) + beta = torch.randn(1, 2, 1) + recurrent_state = torch.randn(1, 1, 2, 2) + state_indices = torch.tensor([0], dtype=torch.int32) + has_initial_state = torch.tensor([True]) + metadata = SimpleNamespace( + cu_seqlens_host=(0, 2), + cu_seqlens_kern=None, + keep_meta=None, + chunk_indices_chunk64_host=(0, 0), + ) + transformed_gate = torch.randn_like(raw_gate) + gate_cumsum = torch.randn_like(raw_gate, dtype=torch.float32) + output = torch.randn_like(v) + final_state = torch.randn(1, 1, 2, 2) + + with ( + patch("vllm_ascend.ops.kimi_kda.clear_ssm_states"), + patch("vllm_ascend.ops.kimi_kda.l2norm_fwd", side_effect=lambda x: x), + patch.object(attention, "_recurrent_gate", return_value=transformed_gate) as recurrent_gate, + patch.object( + torch.ops._C_ascend, + "kda_gate_cumsum", + return_value=gate_cumsum, + create=True, + ) as kda_gate_cumsum, + patch.object( + torch.ops._C_ascend, + "chunk_kda_fwd", + return_value=(output, final_state), + create=True, + ), + ): + actual = attention._run_prefill( + q, + k, + v, + raw_gate, + beta, + recurrent_state, + state_indices, + has_initial_state, + metadata, + ) + + assert actual is output + recurrent_gate.assert_called_once_with(raw_gate) + assert kda_gate_cumsum.call_args.args[0] is transformed_gate + assert kda_gate_cumsum.call_args.args[1] == 64 + assert "use_gate_in_kernel" not in kda_gate_cumsum.call_args.kwargs + + +def test_kda_empty_forward_context_clears_preallocated_output(): + attention = AscendKimiK3DeltaAttention.__new__(AscendKimiK3DeltaAttention) + core_attn_out = torch.full((1, 4, 2, 3), torch.nan) + + with patch( + "vllm_ascend.ops.kimi_kda.get_forward_context", + return_value=SimpleNamespace(attn_metadata=None), + ): + attention._forward( + mixed_qkv=torch.empty(4, 18), + g1=torch.empty(1, 4, 2, 3), + g2=torch.empty(4, 2, 3), + beta=torch.empty(1, 4, 2), + core_attn_out=core_attn_out, + ) + + assert torch.equal(core_attn_out, torch.zeros_like(core_attn_out)) + + +def test_kda_conv_weight_is_packed_once_in_kernel_layout(): + attention = AscendKimiK3DeltaAttention.__new__(AscendKimiK3DeltaAttention) + nn.Module.__init__(attention) + attention.conv_size = 4 + attention.local_projection_size = 6 + attention.conv1d = nn.Module() + source = torch.arange(18 * 4, dtype=torch.float32).reshape(18, 1, 4) + attention.conv1d.weight = nn.Parameter(source) + attention.register_parameter( + _PACKED_CONV_WEIGHT_NAME, + nn.Parameter(torch.empty(4, 18, dtype=torch.bfloat16), requires_grad=False), + ) + original = attention.get_parameter(_PACKED_CONV_WEIGHT_NAME) + original_ptr = original.data_ptr() + + attention._pack_conv_weights() + + packed = attention.get_parameter(_PACKED_CONV_WEIGHT_NAME) + assert packed.data_ptr() == original_ptr + assert packed.dtype == torch.bfloat16 + assert packed.is_contiguous() + torch.testing.assert_close( + packed, + source[:, 0, :].transpose(0, 1).to(torch.bfloat16), + ) diff --git a/tests/ut/ops/test_layernorm.py b/tests/ut/ops/test_layernorm.py index a296ffa8b5bb..4fa8b72bceaf 100644 --- a/tests/ut/ops/test_layernorm.py +++ b/tests/ut/ops/test_layernorm.py @@ -4,7 +4,9 @@ import torch from vllm.config import set_current_vllm_config from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated +from vllm_ascend.ops.layernorm import AscendFusedRMSNormGated from vllm_ascend.utils import enable_custom_op from vllm_ascend.utils import is_310p as is_310p_hw @@ -83,6 +85,33 @@ def test_RMSNorm_creates_bias_from_quant_description(default_vllm_config): assert not layer.bias.requires_grad +@pytest.mark.parametrize("activation", ["sigmoid", "swish"]) +@pytest.mark.parametrize("prenorm", [False, True]) +def test_FusedRMSNormGated_dispatches_to_ascend_kernel(default_vllm_config, activation, prenorm): + layer = FusedRMSNormGated(hidden_size=8, eps=1e-6, activation=activation) + x = torch.randn(1, 4, 2, 8) + gate = torch.randn(4, 2, 8) + residual = torch.randn_like(x) if prenorm else None + expected = (torch.empty_like(x), torch.empty_like(x)) if prenorm else torch.empty_like(x) + + with patch("vllm_ascend.ops.layernorm.rms_norm_gated", return_value=expected) as fused_norm_gate: + actual = layer(x, gate, residual=residual, prenorm=prenorm, residual_in_fp32=prenorm) + + assert isinstance(layer, AscendFusedRMSNormGated) + assert actual is expected + fused_norm_gate.assert_called_once_with( + x, + gate, + layer.weight, + layer.bias, + activation, + residual=residual, + eps=1e-6, + prenorm=prenorm, + residual_in_fp32=prenorm, + ) + + @pytest.mark.skipif(not is_310p_hw(), reason="310P device unittest case.") @pytest.mark.parametrize("residual", [None, torch.randn(4, 8, dtype=torch.float16)]) @patch("torch_npu.npu_rms_norm", side_effect=mock_rms_norm) diff --git a/vllm_ascend/ops/gdn_attn_builder.py b/vllm_ascend/ops/gdn_attn_builder.py index 0168d1a3e072..3ef3ceb82277 100644 --- a/vllm_ascend/ops/gdn_attn_builder.py +++ b/vllm_ascend/ops/gdn_attn_builder.py @@ -26,6 +26,7 @@ ) from vllm.v1.attention.backends.utils import ( NULL_BLOCK_ID, + PAD_SLOT_ID, compute_causal_conv1d_metadata, mamba_get_block_table_tensor, split_decodes_and_prefills, @@ -51,6 +52,32 @@ def _stable_argsort_for_npu(tensor: torch.Tensor) -> torch.Tensor: return torch.argsort(tensor, stable=True) +def _treat_single_token_prefills_with_state_as_decodes( + common_attn_metadata: CommonAttentionMetadata, +) -> CommonAttentionMetadata: + """Use decode metadata for one-token stateful prompt chunks. + + A final one-token prompt chunk can replay the same fixed graph as an + ordinary decode. Once recurrent state exists, both paths must construct + identical GDN metadata so the graph consumes the current state indices. + First-token prefills remain on the prefill path because they have no state + to update. + """ + is_prefilling = common_attn_metadata.is_prefilling + seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound + if is_prefilling is None or seq_lens_cpu is None: + return common_attn_metadata + + query_lens_cpu = torch.diff(common_attn_metadata.query_start_loc_cpu) + prefill_to_decode = is_prefilling & (query_lens_cpu == 1) & (seq_lens_cpu > 1) + if not torch.any(prefill_to_decode).item(): + return common_attn_metadata + + is_prefilling = is_prefilling.clone() + is_prefilling[prefill_to_decode] = False + return common_attn_metadata.replace(is_prefilling=is_prefilling) + + @dataclass class GDNChunkedPrefillMetadata: cu_seqlens_host: tuple[int, ...] @@ -148,7 +175,10 @@ def _build_non_spec_chunked_prefill_metadata( device: torch.device, ) -> GDNChunkedPrefillMetadata: hf_text_config = getattr(builder.vllm_config.model_config, "hf_text_config", None) - if hf_text_config is not None and hasattr(hf_text_config, "linear_num_value_heads"): + linear_attn_config = getattr(hf_text_config, "linear_attn_config", None) + if isinstance(linear_attn_config, dict) and linear_attn_config.get("num_heads") is not None: + gdn_num_heads = linear_attn_config["num_heads"] // builder.vllm_config.parallel_config.tensor_parallel_size + elif hf_text_config is not None and hasattr(hf_text_config, "linear_num_value_heads"): gdn_num_heads = ( hf_text_config.linear_num_value_heads // builder.vllm_config.parallel_config.tensor_parallel_size ) @@ -297,6 +327,45 @@ def _copy_sequence_indices_to_device( return spec_indices, non_spec_indices + def _pad_non_spec_decode_graph_inputs( + self, + state_indices: torch.Tensor, + query_start_loc: torch.Tensor, + *, + num_decode_tokens: int, + graph_batch_size: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Refresh fixed buffers consumed by a non-spec decode graph.""" + assert num_decode_tokens <= graph_batch_size + + padded_state_indices = self.non_spec_state_indices_tensor[:graph_batch_size] + padded_state_indices[num_decode_tokens:].fill_(NULL_BLOCK_ID) + padded_state_indices[:num_decode_tokens].copy_( + state_indices[:num_decode_tokens], + non_blocking=True, + ) + + padded_query_start_loc = self.non_spec_query_start_loc[: graph_batch_size + 1] + padded_query_start_loc[: num_decode_tokens + 1].copy_( + query_start_loc[: num_decode_tokens + 1], + non_blocking=True, + ) + query_padding = padded_query_start_loc[num_decode_tokens + 1 :] + if query_padding.numel() > 0: + query_padding.copy_( + padded_query_start_loc[num_decode_tokens].expand_as(query_padding), + non_blocking=True, + ) + + return padded_state_indices, padded_query_start_loc + + def _reset_spec_decode_graph_inputs(self, graph_batch_size: int) -> None: + """Make a captured speculative branch a no-op for this replay.""" + self.spec_state_indices_tensor[:graph_batch_size].fill_(PAD_SLOT_ID) + self.spec_query_start_loc[: graph_batch_size + 1].zero_() + self.num_accepted_tokens[:graph_batch_size].zero_() + self.spec_actual_seq_lengths[: graph_batch_size + 1].zero_() + def _attach_non_spec_prefill_metadata( self, attn_metadata: GDNAttentionMetadata, @@ -405,6 +474,42 @@ def _attach_non_spec_decode_metadata( ) return attn_metadata + def _fold_spec_sized_prefill_chunks_into_spec( + self, + common_attn_metadata: CommonAttentionMetadata, + spec_sequence_masks_cpu: torch.Tensor, + num_accepted_tokens: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Advance stateful spec-width prompt chunks through live spec inputs.""" + is_prefilling = common_attn_metadata.is_prefilling + seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound + if is_prefilling is None or seq_lens_cpu is None or num_accepted_tokens is None: + return spec_sequence_masks_cpu, num_accepted_tokens + + num_reqs = min( + spec_sequence_masks_cpu.numel(), + is_prefilling.numel(), + seq_lens_cpu.numel(), + ) + is_prefilling = is_prefilling[:num_reqs] + seq_lens_cpu = seq_lens_cpu[:num_reqs] + query_lens_cpu = torch.diff(common_attn_metadata.query_start_loc_cpu)[:num_reqs] + fold = ( + is_prefilling + & ~spec_sequence_masks_cpu + & (query_lens_cpu == self.num_spec + 1) + & (seq_lens_cpu > query_lens_cpu) + ) + fold_indices = fold.nonzero(as_tuple=True)[0] + if fold_indices.numel() == 0: + return spec_sequence_masks_cpu, num_accepted_tokens + + spec_sequence_masks_cpu = spec_sequence_masks_cpu.clone() + spec_sequence_masks_cpu[fold_indices] = True + num_accepted_tokens = num_accepted_tokens.clone() + num_accepted_tokens[fold_indices.to(num_accepted_tokens.device)] = self.num_spec + 1 + return spec_sequence_masks_cpu, num_accepted_tokens + def build( # type: ignore[override] self, common_prefix_len: int, @@ -413,7 +518,7 @@ def build( # type: ignore[override] num_decode_draft_tokens_cpu: torch.Tensor | None = None, fast_build: bool = False, ) -> GDNAttentionMetadata: - m = common_attn_metadata + m = _treat_single_token_prefills_with_state_as_decodes(common_attn_metadata) query_start_loc = m.query_start_loc query_start_loc_cpu = m.query_start_loc_cpu @@ -436,10 +541,22 @@ def build( # type: ignore[override] else: num_reqs = num_decode_draft_tokens_cpu.numel() spec_sequence_masks_cpu = self.spec_sequence_masks_cpu[:num_reqs] - torch.ge( - num_decode_draft_tokens_cpu, - 0, - out=spec_sequence_masks_cpu, + runtime_draft_tokens = num_decode_draft_tokens_cpu[num_decode_draft_tokens_cpu >= 0] + if runtime_draft_tokens.sum().item() > 0: + torch.ge( + num_decode_draft_tokens_cpu, + 0, + out=spec_sequence_masks_cpu, + ) + else: + # Dynamic speculative decoding can be enabled while this batch + # carries no draft tokens. Treat it as ordinary decode unless a + # stateful spec-width prompt chunk must use the spec branch. + spec_sequence_masks_cpu.zero_() + spec_sequence_masks_cpu, num_accepted_tokens = self._fold_spec_sized_prefill_chunks_into_spec( + m, + spec_sequence_masks_cpu, + num_accepted_tokens, ) num_spec_decodes = spec_sequence_masks_cpu.sum().item() if num_spec_decodes == 0: @@ -584,6 +701,12 @@ def build( # type: ignore[override] spec_sequence_indices, ) + # A FULL graph retains captured speculative conv/recurrent tasks. Clear + # their stable inputs on every no-spec replay so an idle or prefill + # batch cannot mutate state belonging to the preceding request. + if self.use_full_cuda_graph and self.use_spec_decode and num_spec_decodes == 0: + self._reset_spec_decode_graph_inputs(m.num_reqs) + chunk_indices: torch.Tensor | None = None chunk_offsets: torch.Tensor | None = None prefill_query_start_loc: torch.Tensor | None = None @@ -642,8 +765,6 @@ def build( # type: ignore[override] f"num_decodes: {num_decodes}, num_spec_decodes: {num_spec_decodes}" ) - batch_size = m.num_actual_tokens - if ( self.use_full_cuda_graph and num_prefills == 0 @@ -707,22 +828,17 @@ def build( # type: ignore[override] and num_spec_decodes == 0 and num_decodes <= self.decode_cudagraph_max_bs ): - self.non_spec_state_indices_tensor[batch_size:].fill_(NULL_BLOCK_ID) - self.non_spec_state_indices_tensor[:num_decodes].copy_( + graph_batch_size = m.num_reqs + ( non_spec_state_indices_tensor, - non_blocking=True, - ) - non_spec_state_indices_tensor = self.non_spec_state_indices_tensor[:batch_size] - non_spec_state_indices_tensor[num_decodes:].fill_(NULL_BLOCK_ID) - non_spec_conv1d_cache_indices = non_spec_state_indices_tensor - - self.non_spec_query_start_loc[: num_decodes + 1].copy_( non_spec_query_start_loc, - non_blocking=True, + ) = self._pad_non_spec_decode_graph_inputs( + non_spec_state_indices_tensor, + non_spec_query_start_loc, + num_decode_tokens=num_decode_tokens, + graph_batch_size=graph_batch_size, ) - non_spec_num_query_tokens = non_spec_query_start_loc[-1] - non_spec_query_start_loc = self.non_spec_query_start_loc[: batch_size + 1] - non_spec_query_start_loc[num_decodes + 1 :].fill_(non_spec_num_query_tokens) + non_spec_conv1d_cache_indices = non_spec_state_indices_tensor attn_metadata = GDNAttentionMetadata( num_prefills=num_prefills, diff --git a/vllm_ascend/ops/kimi_kda.py b/vllm_ascend/ops/kimi_kda.py new file mode 100644 index 000000000000..10e78e8b9483 --- /dev/null +++ b/vllm_ascend/ops/kimi_kda.py @@ -0,0 +1,569 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Ascend backend for the vLLM 0.27 Kimi K3 delta-attention layer. + +The projections, weight loading, and cache specification stay owned by +upstream vLLM. Only the CUDA-specific convolution and KDA execution is +replaced here with the Ascend metadata builder and AscendC operators. +""" + +from functools import wraps + +import torch +from einops import rearrange +from torch import nn +from vllm.compilation.breakable_cudagraph import eager_break_during_capture +from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, +) +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 +from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata +from vllm.v1.attention.backends.utils import PAD_SLOT_ID + +from vllm_ascend.ops.gdn_attn_builder import AscendGDNAttentionBackend +from vllm_ascend.ops.triton.fla.utils import clear_ssm_states +from vllm_ascend.ops.triton.kda.kda import fused_kda_gate + +_KDA_CHUNK_SIZE = 64 +_PACKED_CONV_WEIGHT_NAME = "ascend_conv1d_weight" + + +def _zero_padded_output( + output: torch.Tensor, + num_live_tokens: torch.Tensor, +) -> torch.Tensor: + """Clear graph-padding rows using a device-side live-token count.""" + token_indices = torch.arange( + output.shape[1], + dtype=num_live_tokens.dtype, + device=output.device, + ) + valid_tokens = token_indices < num_live_tokens + return torch.where(valid_tokens.view(1, -1, 1, 1), output, 0.0) + + +def _zero_padded_recurrent_output( + output: torch.Tensor, + query_start_loc: torch.Tensor, +) -> torch.Tensor: + """Clear graph-padding rows skipped by recurrent KDA.""" + return _zero_padded_output(output, query_start_loc[-1]) + + +def _prepare_beta( + raw_beta: torch.Tensor, + num_actual_tokens: int, +) -> torch.Tensor: + """Convert vLLM 0.27's packed raw beta to the AscendC contract.""" + return raw_beta[:, :num_actual_tokens].float().sigmoid() + + +class AscendKimiK3DeltaAttention(KimiK3DeltaAttention): + """Kimi K3 KDA using AscendC prefill and recurrent kernels.""" + + def __init__(self, config, vllm_config, prefix: str = "") -> None: + quant_config = getattr(vllm_config, "quant_config", None) + uses_mixed_projection = bool( + quant_config is not None + and getattr( + quant_config, + "uses_kimi_k3_mixed_kda_projection", + lambda _prefix: False, + )(f"{prefix}.in_proj_qkvgfab") + ) + super().__init__(config, vllm_config, prefix) + 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. + self.in_proj_qkvgfab = MergedColumnParallelLinear( + self.hidden_size, + [self.projection_size] * 3, + bias=False, + 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, + 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_() + # 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 + # by the validated v0.26 implementation. + self.o_norm.eps = config.rms_norm_eps + # vLLM keeps the checkpoint-compatible FP32 [3C, 1, W] weight, while + # npu_causal_conv1d_custom consumes an activation-dtype [W, 3C] + # tensor. Materialize that kernel layout once after weight loading. + self.register_parameter( + _PACKED_CONV_WEIGHT_NAME, + nn.Parameter( + torch.empty( + self.conv_size, + 3 * self.local_projection_size, + dtype=self.model_config.dtype, + device=self.conv1d.weight.device, + ), + requires_grad=False, + ), + ) + original_process_weights = self.conv1d.quant_method.process_weights_after_loading + + @wraps(original_process_weights) + def process_weights_and_pack(*args, **kwargs): + result = original_process_weights(*args, **kwargs) + self._pack_conv_weights() + return result + + self.conv1d.quant_method.process_weights_after_loading = process_weights_and_pack + + def get_attn_backend(self) -> type[AttentionBackend]: + return AscendGDNAttentionBackend + + def forward( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + ) -> 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) + core_attn_out = torch.empty( + (1, num_tokens, self.local_num_heads, self.head_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + self._forward( + mixed_qkv=mixed_qkv, + g1=g1, + g2=g2, + beta=beta, + core_attn_out=core_attn_out, + ) + 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) + + @staticmethod + def _run_causal_conv1d( + mixed_qkv: torch.Tensor, + conv_weights_t: torch.Tensor, + conv_state: torch.Tensor, + query_start_loc: torch.Tensor, + cache_indices: torch.Tensor, + initial_state_mode: torch.Tensor | None, + *, + run_mode: int, + num_accepted_tokens: torch.Tensor | None = None, + ) -> torch.Tensor: + output = torch.empty_like(mixed_qkv) + # Consume the operator's declared output alias. Returning ``output`` + # independently would let graph functionalization treat the custom-op + # result as dead and expose the uninitialized allocation instead. + return torch.ops._C_ascend.npu_causal_conv1d_custom( + output, + mixed_qkv, + conv_weights_t, + conv_state=conv_state, + bias_opt=None, + query_start_loc_opt=query_start_loc, + cache_indices_opt=cache_indices, + initial_state_mode_opt=initial_state_mode, + num_accepted_tokens_opt=num_accepted_tokens, + activation_mode=1, + pad_slot_id=PAD_SLOT_ID, + run_mode=run_mode, + ) + + @torch.no_grad() + def _pack_conv_weights(self) -> None: + if self.conv1d.weight.is_meta: + return + packed_param = self.get_parameter(_PACKED_CONV_WEIGHT_NAME) + packed_weight = ( + self.conv1d.weight.view(self.conv1d.weight.size(0), self.conv1d.weight.size(2)) + .transpose(0, 1) + .to(device=packed_param.device, dtype=packed_param.dtype) + .contiguous() + ) + replace_parameter( + self, + _PACKED_CONV_WEIGHT_NAME, + packed_weight, + prefer_copy=True, + ) + + def _recurrent_gate(self, raw_gate: torch.Tensor) -> torch.Tensor: + flat_gate = rearrange(raw_gate, "1 n h d -> n (h d)") + gate = fused_kda_gate( + flat_gate, + self.A_log, + self.head_dim, + g_bias=self.dt_bias, + ) + return gate.unsqueeze(0) + + def _run_recurrent( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_gate: torch.Tensor, + beta: torch.Tensor, + recurrent_state: torch.Tensor, + cu_seqlens: torch.Tensor, + state_indices: torch.Tensor, + *, + num_accepted_tokens: torch.Tensor | None = None, + ) -> torch.Tensor: + return torch.ops._C_ascend.recurrent_kda( + q.contiguous(), + k.contiguous(), + v.contiguous(), + raw_gate.contiguous(), + beta.contiguous(), + recurrent_state, + cu_seqlens, + state_indices, + self.A_log.reshape(-1).contiguous(), + self.dt_bias.contiguous(), + num_accepted_tokens=num_accepted_tokens, + scale=self.head_dim**-0.5, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=False, + allow_neg_eigval=False, + safe_gate=self.gate_lower_bound is not None, + lower_bound=(self.gate_lower_bound if self.gate_lower_bound is not None else -5.0), + ) + + def _run_prefill( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_gate: torch.Tensor, + beta: torch.Tensor, + recurrent_state: torch.Tensor, + state_indices: torch.Tensor, + has_initial_state: torch.Tensor, + prebuilt_metadata, + ) -> torch.Tensor: + cu_seqlens = ( + prebuilt_metadata.cu_seqlens_host + if prebuilt_metadata.cu_seqlens_kern is None + else prebuilt_metadata.cu_seqlens_kern + ) + keep = prebuilt_metadata.keep_meta + if keep is not None: + state_indices = state_indices[keep] + has_initial_state = has_initial_state[keep] + + # The recurrent cache is [H, V, K], while chunk_kda_fwd consumes + # [H, K, V]. Keep the conversion at this operator boundary. + initial_state_vk = recurrent_state[state_indices].contiguous() + clear_ssm_states(initial_state_vk, has_initial_state) + initial_state_kv = initial_state_vk.transpose(-1, -2).contiguous() + + q = l2norm_fwd(q.contiguous()) + k = l2norm_fwd(k.contiguous()) + if self.gate_lower_bound is not None: + gate_cumsum = torch.ops._C_ascend.kda_gate_cumsum( + raw_gate.contiguous(), + _KDA_CHUNK_SIZE, + A_log=self.A_log.reshape(-1).contiguous(), + dt_bias=self.dt_bias.contiguous(), + cu_seqlens=cu_seqlens, + use_gate_in_kernel=True, + safe_gate=True, + lower_bound=self.gate_lower_bound, + layout="BSND", + ) + else: + gate = self._recurrent_gate(raw_gate) + gate_cumsum = torch.ops._C_ascend.kda_gate_cumsum( + gate.contiguous(), + _KDA_CHUNK_SIZE, + cu_seqlens=cu_seqlens, + layout="BSND", + ) + result = torch.ops._C_ascend.chunk_kda_fwd( + q, + k, + v.contiguous(), + gate_cumsum, + beta.contiguous(), + self.head_dim**-0.5, + _KDA_CHUNK_SIZE, + layout="BSND", + initial_state=initial_state_kv, + output_final_state=True, + cu_seqlens=cu_seqlens, + chunk_indices=prebuilt_metadata.chunk_indices_chunk64_host, + return_intermediate=False, + ) + recurrent_state[state_indices] = result[1].transpose(-1, -2).contiguous().to(recurrent_state.dtype) + return result[0] + + @eager_break_during_capture + def _forward( + self, + mixed_qkv: torch.Tensor, + g1: torch.Tensor, + g2: torch.Tensor, + beta: torch.Tensor, + core_attn_out: torch.Tensor, + ) -> None: + """Dispatch speculative, prefill, and decode tokens through KDA kernels.""" + forward_context = get_forward_context() + attn_metadata_raw = forward_context.attn_metadata + if attn_metadata_raw is None: + core_attn_out.zero_() + return + + assert isinstance(attn_metadata_raw, dict) + attn_metadata = attn_metadata_raw[self.prefix] + assert isinstance(attn_metadata, GDNAttentionMetadata) + + num_actual_tokens = attn_metadata.num_actual_tokens + mixed_qkv = mixed_qkv[:num_actual_tokens] + g1 = g1[:, :num_actual_tokens] + g2 = g2[:num_actual_tokens] + beta = _prepare_beta(beta, num_actual_tokens) + + conv_state, recurrent_state = self.kv_cache + conv_weights_t = self.get_parameter(_PACKED_CONV_WEIGHT_NAME) + spec_masks = attn_metadata.spec_sequence_masks + spec_token_indices = attn_metadata.spec_token_indx + non_spec_token_indices = attn_metadata.non_spec_token_indx + + if spec_masks is not None: + if attn_metadata.num_prefills == 0 and attn_metadata.num_decodes == 0: + mixed_spec = mixed_qkv + raw_gate_spec = g1 + beta_spec = beta + mixed_non_spec = raw_gate_non_spec = beta_non_spec = None + else: + assert spec_token_indices is not None + assert non_spec_token_indices is not None + mixed_spec = mixed_qkv.index_select(0, spec_token_indices) + raw_gate_spec = g1.index_select(1, spec_token_indices) + beta_spec = beta.index_select(1, spec_token_indices) + mixed_non_spec = mixed_qkv.index_select(0, non_spec_token_indices) + raw_gate_non_spec = g1.index_select(1, non_spec_token_indices) + beta_non_spec = beta.index_select(1, non_spec_token_indices) + else: + mixed_spec = raw_gate_spec = beta_spec = None + mixed_non_spec = mixed_qkv + raw_gate_non_spec = g1 + beta_non_spec = beta + + core_spec = None + if mixed_spec is not None: + spec_meta = attn_metadata.spec_decode_metadata + assert spec_meta is not None + spec_conv_meta = spec_meta.spec_causal_conv1d + mixed_spec = self._run_causal_conv1d( + mixed_spec, + conv_weights_t, + conv_state, + spec_conv_meta.query_start_loc, + spec_conv_meta.cache_indices, + None, + run_mode=1, + num_accepted_tokens=spec_conv_meta.num_accepted_tokens, + ) + q_spec, k_spec, v_spec = ( + rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim) for x in mixed_spec.chunk(3, dim=-1) + ) + assert raw_gate_spec is not None and beta_spec is not None + assert attn_metadata.spec_query_start_loc is not None + assert attn_metadata.spec_state_indices_tensor is not None + core_spec = self._run_recurrent( + q_spec, + k_spec, + v_spec, + raw_gate_spec, + beta_spec, + recurrent_state, + attn_metadata.spec_query_start_loc, + attn_metadata.spec_state_indices_tensor, + num_accepted_tokens=spec_conv_meta.num_accepted_tokens, + ) + core_spec = _zero_padded_recurrent_output( + core_spec, + attn_metadata.spec_query_start_loc, + ) + + core_non_spec = None + if mixed_non_spec is not None and mixed_non_spec.shape[0] > 0: + if attn_metadata.num_prefills > 0: + prefill_meta = attn_metadata.non_spec_prefill_metadata + assert prefill_meta is not None + mixed_non_spec = self._run_causal_conv1d( + mixed_non_spec, + conv_weights_t, + conv_state, + prefill_meta.causal_conv1d.query_start_loc, + prefill_meta.causal_conv1d.cache_indices, + prefill_meta.causal_conv1d.initial_state_mode, + run_mode=0, + ) + elif attn_metadata.num_decodes > 0: + decode_meta = attn_metadata.non_spec_decode_metadata + assert decode_meta is not None + mixed_non_spec = self._run_causal_conv1d( + mixed_non_spec, + conv_weights_t, + conv_state, + decode_meta.causal_conv1d.query_start_loc, + decode_meta.causal_conv1d.cache_indices, + None, + run_mode=1, + ) + + q_non_spec, k_non_spec, v_non_spec = ( + rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim) for x in mixed_non_spec.chunk(3, dim=-1) + ) + assert raw_gate_non_spec is not None + assert beta_non_spec is not None + + split_non_spec = spec_masks is None and attn_metadata.num_prefills > 0 and attn_metadata.num_decodes > 0 + num_decode_tokens = attn_metadata.num_decode_tokens + core_decode = None + if split_non_spec: + assert attn_metadata.non_spec_query_start_loc is not None + assert attn_metadata.non_spec_state_indices_tensor is not None + core_decode = self._run_recurrent( + q_non_spec[:, :num_decode_tokens], + k_non_spec[:, :num_decode_tokens], + v_non_spec[:, :num_decode_tokens], + raw_gate_non_spec[:, :num_decode_tokens], + beta_non_spec[:, :num_decode_tokens], + recurrent_state, + attn_metadata.non_spec_query_start_loc[: attn_metadata.num_decodes + 1], + attn_metadata.non_spec_state_indices_tensor[: attn_metadata.num_decodes], + ) + + if attn_metadata.num_prefills > 0: + if split_non_spec: + q_non_spec = q_non_spec[:, num_decode_tokens:] + k_non_spec = k_non_spec[:, num_decode_tokens:] + v_non_spec = v_non_spec[:, num_decode_tokens:] + raw_gate_non_spec = raw_gate_non_spec[:, num_decode_tokens:] + beta_non_spec = beta_non_spec[:, num_decode_tokens:] + + assert attn_metadata.prefill_state_indices is not None + assert attn_metadata.prefill_has_initial_state is not None + prefill_meta = attn_metadata.non_spec_prefill_metadata + assert prefill_meta is not None + core_prefill = self._run_prefill( + q_non_spec, + k_non_spec, + v_non_spec, + raw_gate_non_spec, + beta_non_spec, + recurrent_state, + attn_metadata.prefill_state_indices, + attn_metadata.prefill_has_initial_state, + prefill_meta.chunk, + ) + core_non_spec = ( + torch.cat((core_decode, core_prefill), dim=1) if core_decode is not None else core_prefill + ) + elif attn_metadata.num_decodes > 0: + assert attn_metadata.non_spec_query_start_loc is not None + assert attn_metadata.non_spec_state_indices_tensor is not None + core_non_spec = self._run_recurrent( + q_non_spec, + k_non_spec, + v_non_spec, + raw_gate_non_spec, + beta_non_spec, + recurrent_state, + attn_metadata.non_spec_query_start_loc[: attn_metadata.num_decodes + 1], + attn_metadata.non_spec_state_indices_tensor, + ) + + if core_non_spec is not None: + assert attn_metadata.non_spec_query_start_loc is not None + core_non_spec = _zero_padded_recurrent_output( + core_non_spec, + attn_metadata.non_spec_query_start_loc, + ) + + if core_spec is None and core_non_spec is None: + # Idle DP dummy runs carry graph-shaped metadata with no live work. + # Do not feed a previous replay's output through the norm gate. + core_attn_out.zero_() + return + + num_live_tokens = None + if core_spec is not None: + assert attn_metadata.spec_query_start_loc is not None + num_live_tokens = attn_metadata.spec_query_start_loc[-1] + if core_non_spec is not None: + assert attn_metadata.non_spec_query_start_loc is not None + num_non_spec_tokens = attn_metadata.non_spec_query_start_loc[-1] + num_live_tokens = num_non_spec_tokens if num_live_tokens is None else num_live_tokens + num_non_spec_tokens + assert num_live_tokens is not None + + # Reuse the caller-owned result buffer. FULL graphs can leave rows + # outside the live spec/non-spec index sets, so define them before the + # two index copies rather than allocating a temporary merged tensor. + core_attn_out[:, :num_actual_tokens].zero_() + if core_spec is not None and core_non_spec is not None: + assert spec_token_indices is not None + assert non_spec_token_indices is not None + assert spec_token_indices.numel() + non_spec_token_indices.numel() <= num_actual_tokens + core_attn_out[:, :num_actual_tokens].index_copy_(1, spec_token_indices, core_spec) + core_attn_out[:, :num_actual_tokens].index_copy_(1, non_spec_token_indices, core_non_spec) + elif core_spec is not None: + core_attn_out[:, :num_actual_tokens] = core_spec + elif core_non_spec is not None: + core_attn_out[:, :num_actual_tokens] = core_non_spec + + # The registered Ascend FusedRMSNormGated uses the fused norm-gate + # kernel while preserving the upstream parameter/loading contract. + normalized = self.o_norm(core_attn_out[:, :num_actual_tokens], g2) + # Mask again after the norm gate: zero * sigmoid(NaN) is still NaN in + # static padding rows whose captured gate values are not live. + core_attn_out[:, :num_actual_tokens].copy_(_zero_padded_output(normalized, num_live_tokens)) + core_attn_out[:, num_actual_tokens:].zero_() diff --git a/vllm_ascend/ops/layernorm.py b/vllm_ascend/ops/layernorm.py index d88246743ed4..cbe73d0b06c0 100644 --- a/vllm_ascend/ops/layernorm.py +++ b/vllm_ascend/ops/layernorm.py @@ -19,8 +19,10 @@ from torch import nn from vllm.config import get_current_vllm_config from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm, RMSNormGated +from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated from vllm_ascend.device.device_op import DeviceOperator +from vllm_ascend.ops.triton.kda.kda import rms_norm_gated from vllm_ascend.ops.triton.layernorm_gated import layer_norm_fwd_npu from vllm_ascend.utils import enable_custom_op @@ -194,3 +196,20 @@ def reset_parameters(self): def forward_oot(self, x, z=None): """If z is not None, we do norm(x) * silu(z) if norm_before_gate, else norm(x * silu(z))""" return LayerNormFn.apply(x, self.weight, self.bias, z, self.eps, self.group_size, self.norm_before_gate, True) + + +class AscendFusedRMSNormGated(FusedRMSNormGated): + """Use Ascend's fused kernel at the upstream FLA CustomOp boundary.""" + + def forward_oot(self, x, g, residual=None, prenorm=False, residual_in_fp32=False): + return rms_norm_gated( + x, + g, + self.weight, + self.bias, + self.activation, + residual=residual, + eps=self.eps, + prenorm=prenorm, + residual_in_fp32=residual_in_fp32, + ) diff --git a/vllm_ascend/ops/triton/kda/kda.py b/vllm_ascend/ops/triton/kda/kda.py index fc09fde39677..313b3b8949b7 100644 --- a/vllm_ascend/ops/triton/kda/kda.py +++ b/vllm_ascend/ops/triton/kda/kda.py @@ -411,6 +411,7 @@ def rms_norm_gated( activation=activation, eps=eps, residual=residual, + out_dtype=x.dtype, # Preserve the input, as in the v0.26 K3 fused norm gate. residual_dtype=residual_dtype, is_rms_norm=True, ) diff --git a/vllm_ascend/utils.py b/vllm_ascend/utils.py index 8841b2458f64..26e1ef72aac6 100644 --- a/vllm_ascend/utils.py +++ b/vllm_ascend/utils.py @@ -674,7 +674,7 @@ def register_ascend_customop(vllm_config: VllmConfig | None = None): from vllm_ascend.ops.fused_moe.fused_moe import AscendMoERunner from vllm_ascend.ops.fused_moe.routed_experts import AscendRoutedExperts from vllm_ascend.ops.gdn import AscendGatedDeltaNetAttention - from vllm_ascend.ops.layernorm import AscendGemmaRMSNorm, AscendRMSNorm, AscendRMSNormGated + from vllm_ascend.ops.layernorm import AscendFusedRMSNormGated, AscendGemmaRMSNorm, AscendRMSNorm, AscendRMSNormGated from vllm_ascend.ops.linear import ( AscendColumnParallelLinear, AscendMergedColumnParallelLinear, @@ -722,6 +722,7 @@ def register_ascend_customop(vllm_config: VllmConfig | None = None): "MMEncoderAttention": AscendMMEncoderAttention, "ApplyRotaryEmb": AscendApplyRotaryEmb, "RMSNormGated": AscendRMSNormGated, + "FusedRMSNormGated": AscendFusedRMSNormGated, "Conv3dLayer": AscendConv3dLayer, "RelPosAttention": AscendRelPosAttention, "CustomQwen2Decoder": AscendCustomQwen2Decoder,