[Feature][KDA] Add Kimi K3 KDA execution

[maoxx241]

What this PR does / why we need it?

• Adds the Ascend KDA attention builder and runtime.
• Implements chunk-prefill and recurrent-decode dispatch through the KDA AscendC operators.
• Handles causal-convolution state packing and mixed W8A8/BF16 projection loading.
• Registers upstream FusedRMSNormGated with the Ascend CustomOp backend and reuses the existing fused Triton kernel for output normalization and gating. Weight loading, checkpoint epsilon, and kernel arithmetic/tiling are unchanged; the result is allocated separately to preserve the input.
• Adds focused builder, KDA execution, CustomOp dispatch, and fused norm-gate numerical tests.

Dependencies

• [Feature][Ops] Add Kimi K3 AscendC operators #14426 provides the KDA AscendC operators and bindings.
• [Feature][Model] Register Kimi K3 model adapters #14600 consumes this implementation through the Kimi K3 model adapter.

How was this patch tested?

• git diff --check and bash format.sh ci
• On each of vLLM v0.27.1 and the pinned upstream revision: 40 targeted CPU tests pass (4 existing skips), and 22 real A3 NPU numerical cases pass. Norm-gate coverage includes FP16/BF16, sigmoid/SiLU, packed gate strides, affine and residual/prenorm paths, and input preservation.
• A5 performance and full-model serving were not rerun for the fused norm-gate change.
• Python bytecode compilation for changed implementation and test files
• Unit coverage in tests/ut/ops/test_gdn_attn_builder.py and tests/ut/ops/test_kimi_kda.py
• Integrated KDA prefill/decode coverage with the Kimi K3 checkpoint

Does this PR introduce any user-facing change?

No model is registered by this PR alone. It provides KDA execution for Kimi K3.

• vLLM version: v0.27.1
• vLLM main: vllm-project/vllm@ba07e4a

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces the Ascend attention backend integration for the Kimi K3 model. It provides essential support for KDA and MLA mechanisms, including specialized prefill and recurrent execution wrappers, output gating, and optimized A5 dispatch. These changes facilitate the Kimi K3 integration stack by enabling necessary attention-layer functionality while maintaining compatibility with existing vLLM contracts.

Highlights

• KDA Integration: Added KDA prefill and recurrent execution wrappers for Kimi K3 model support.
• MLA Enhancements: Implemented no-RoPE MLA metadata handling, output gating, and attention residual support.
• A5 Dispatch: Added A5 native MLA Prolog V3 dispatch for BF16 and optional RoPE configurations.
• Parallel Drafting: Implemented sequence-length finalization for parallel-drafting at the MLA boundary.

New Features

🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Ignored Files

• Ignored by pattern: .github/workflows/** (1)
  • .github/workflows/scripts/test_config.yaml

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[github-actions]

👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:‌‌

• A PR should do only one thing, smaller PRs enable faster reviews.
• Every PR should include unit tests and end-to-end tests ‌to ensure it works and is not broken by other future PRs.
• Write the commit message by fulfilling the PR description to help reviewer and future developers understand.

If CI fails, you can run linting and testing checks locally according Contributing and Testing.

---

Tip

💡 Consider Linking a Related Issue or RFC

Your PR title contains the [Feature] tag, indicating a bug fix or new feature.

Linking a related issue or RFC in the PR description is strongly encouraged — it gives reviewers helpful context and speeds up the review. You can use any of these keywords:

• Fixes #<issue_number>
• Closes #<issue_number>
• Resolves #<issue_number>
• Refs #<rfc_or_issue_number> (for RFCs)

🙏 Thanks for helping us keep the project well-organized!

[gemini-code-assist]

Code Review

Suggested PR Title:

[Attention][Feature] Add Ascend support for Kimi K3 Delta-Attention and enhance MLA features

Suggested PR Summary:

### What this PR does / why we need it?
This PR introduces Ascend backend support for the Kimi K3 delta-attention (KDA) layer by implementing `AscendKimiK3DeltaAttention` using AscendC prefill and recurrent kernels. It also enhances the MLA (Multi-Head Latent Attention) implementation to support optional RoPE inputs, output gating, and deferred parallel draft reject finalization. Additionally, it updates the GDN attention builder to correctly extract head counts from `linear_attn_config`.

### Does this PR introduce _any_ user-facing change?
Yes, it adds support for Kimi K3 models and enhances MLA capabilities on Ascend devices.

### How was this patch tested?
Tested using newly added unit tests in `tests/ut/attention/a2/test_mla_v1.py` and `tests/ut/ops/test_gdn_attn_builder.py`.

### Review Feedback
1. **MLA RoPE Mode Lookup**: In `mla_v1.py`, the `layer_name` lookup in `static_forward_context` might fail if the keys do not match exactly. It is recommended to fall back to checking `layer_name.rsplit(".", 1)[0]`.
2. **Tensor Padding**: In `mla_v1.py`, using `torch.cat` with a zero tensor is preferred over `F.pad` with 6 elements on a 3D tensor for better compatibility and performance on `torch_npu`.
3. **Safe Namespace Access**: In `kimi_kda.py`, use `getattr(torch.ops, "_C_ascend", None)` to safely check for the custom namespace registration to avoid unhandled `AttributeError`s.
4. **Distributed Check**: In `kimi_kda.py`, avoid calling `get_pcp_group().world_size` unconditionally as it can crash in non-distributed environments; check `self.vllm_config.parallel_config.prefill_context_parallel_size > 1` instead.
5. **GDN Head Count Detection**: In `gdn_attn_builder.py`, support both dictionary and attribute access for `linear_attn_config` to ensure robust head count detection.

[maoxx241]

Implementation walkthrough

Metadata construction

vllm_ascend/ops/gdn_attn_builder.py builds separate metadata for:

• chunk prefill;
• recurrent decode;
• speculative decode;
• causal-convolution state.

The builder compacts empty segments before launching the prefill operator, pads/reset graph inputs for stable ACLGraph replay, and keeps the sequence reorder indices on device.

The one-token boundary is state-aware: a one-token request with existing recurrent state is decode, while a genuine one-token prompt without state remains prefill. This prevents the same query_len == 1 shape from selecting the wrong KDA path.

KDA execution

AscendKimiK3DeltaAttention in vllm_ascend/ops/kimi_kda.py provides the Ascend implementation used by the upstream Kimi KDA module:

1. project and pack the Q/K/V, gate, and convolution inputs;
2. run the custom causal-convolution operator and update its cache;
3. for prefill, build the cumulative gate and call chunk_kda_fwd;
4. for decode, normalize the recurrent gate and call recurrent_kda;
5. zero graph-padding rows before returning the output.

_pack_conv_weights() prepares the convolution weights once in the layout required by the NPU operator. The projection path supports the K3 mixed layout where Q/K/V are quantized while gate-related projections remain floating point.

Backend registration

AscendGDNAttentionBackend returns the new metadata builder, allowing the existing vLLM attention-layer registration to select this implementation without a K3-specific model-runner branch.

Coverage

The tests cover prefill/decode classification, one-token boundaries, empty segments, graph padding/reset, causal-convolution metadata, mixed projection loading, chunk prefill, recurrent decode, and padded-output zeroing.

[github-actions]

This pull request has conflicts, please resolve those before we can evaluate the pull request.

[maoxx241]

Closing this child PR following the merge of parent #14454.