[Performance][KDA] Compose and overlap gate projections on main

[Dawn952]

What this PR does / why we need it?

This recreates #15168 without the unrelated mla_v1.py head-padding change and its corresponding test changes.

For the existing mixed-precision Kimi K3 KDA layout, this change:

• composes f_proj.weight = f_b_proj.weight @ f_a_proj.weight after checkpoint loading and later source-weight reloads;
• packs beta, the composed F projection, and the output gate into one floating-point BFG projection;
• runs a two-stage auxiliary-stream schedule:
  1. main-stream MXFP DynamicQuant overlaps the auxiliary BFG GEMM;
  2. main-stream QKV GEMM overlaps auxiliary B/F/G split, beta FP32 sigmoid, and gate reshaping;
• marks mixed-path beta as preprocessed so the v0.27 dispatch path only slices it, while ordinary upstream raw beta retains its existing FP32 sigmoid path;
• preserves the configured W8A8 MXFP8 scale_alg when DynamicQuant is split from the linear method;
• routes global or local F shards through the v0.27 packed loader;
• leaves the ordinary upstream in_proj_qkvgfab path unchanged.

Compared with #15168, this PR does not modify:

• vllm_ascend/attention/mla_v1.py;
• tests/ut/attention/a2/test_mla_v1.py.

ACL graph and stream-order assessment

The auxiliary stream records bfg_ready as its final node, and the main stream waits for that exact tail event. No auxiliary operation is queued after the join point.

The four-event schedule is intentional:

1. hidden_states_ready forks BFG from main;
2. bfg_projection_ready serializes the BFG and QKV Cube matmuls;
3. quant_ready starts auxiliary beta vector work only after main DynamicQuant;
4. bfg_ready joins the complete auxiliary tail before KDA consumes its outputs.

The live-token slices for mixed_qkv, g1, and g2 remain at the eager dispatch boundary. They are host-side basic views: after QKV GEMM and the auxiliary tail wait are enqueued, Python creates those views while the asynchronous device work is still running. Moving num_actual_tokens into the projection path would not add an NPU kernel overlap and would weaken the runtime metadata boundary.

Static assessment: the stream is fully joined for ACLGraph capture and mixed beta is transformed exactly once. A real NPU ACLGraph capture/replay has not been run for this main port.

How was this patch tested?

No additional local or NPU validation was run for this replacement PR. The retained KDA implementation and focused tests are unchanged from #15168; this replacement only excludes the MLA implementation and test diffs.

Does this PR introduce any user-facing change?

No API or configuration change. Runtime behavior changes only for Kimi K3 KDA layers using the existing mixed-precision full-rank gate layout.

• 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 optimizes the Kimi K3 Delta Attention (KDA) performance by introducing a fused BFG projection and a two-stage auxiliary stream schedule. By packing beta, F projection, and output gate operations and overlapping them with main-stream computations like MXFP DynamicQuant and QKV GEMM, the implementation reduces latency and improves hardware utilization on NPU devices.

Highlights

• Fused BFG Projection: Packed beta, the composed F projection, and the output gate into a single floating-point BFG projection to streamline operations.
• Two-Stage Auxiliary Stream Schedule: Implemented a new scheduling mechanism that overlaps main-stream MXFP DynamicQuant and QKV GEMM with auxiliary BFG operations to improve NPU utilization.
• Weight Composition: Added logic to compose F projection weights after checkpoint loading and source-weight reloads.
• Beta Preprocessing: Optimized beta handling by marking it as preprocessed, allowing the v0.27 dispatch path to skip redundant sigmoid calculations.
• Test Suite Updates: Updated Kimi K3 adapter and KDA operation tests to validate the new fused BFG structure and stream scheduling logic.

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.

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.

[gemini-code-assist]

Code Review

Suggested PR Title:\n\nmarkdown\n[Ops][Feature] Optimize Kimi K3 Delta Attention with Fused BFG Projection and Multi-Stream Overlap\n\n\nSuggested PR Summary:\n\nmarkdown\n### What this PR does / why we need it?\nThis pull request optimizes the Kimi K3 Delta Attention (KDA) implementation on Ascend NPU. It splits the precision groups for QKV and BFG projections, introducing a fused BFG linear layer (`_KDAFusedBFGLinear`) that offline-composes the F projection and packs it with beta and the output gate. It also implements multi-stream execution (`_run_overlapped_qkv_bfg`) to overlap the dynamic quantization of QKV with the BFG projection, improving hardware utilization.\n\nFeedback on the implementation highlights critical issues when initializing models on the `meta` device. Specifically, loading weights when parameters are on the `meta` device will raise `RuntimeError` during copy operations, and the fusion logic (`_maybe_fuse_f_proj`) can be bypassed or fail if weights are loaded out of order.\n\n### Does this PR introduce _any_ user-facing change?\nNo, this is an internal performance optimization and refactoring of the Kimi K3 attention layer.\n\n### How was this patch tested?\nThe changes are covered by new unit tests in `tests/ut/models/test_kimi_k3_adapter.py` and `tests/ut/ops/test_kimi_kda.py` verifying weight loading, fused projection correctness, stream synchronization, and quantization integration.\n

[gemini-code-assist]

If the checkpoint weights are loaded in an order where f_a_proj and f_b_proj are loaded before self.weight, _maybe_fuse_f_proj will be called while self.weight is still on the meta device, causing a copy failure or skipping fusion without a way to re-trigger it.

To prevent this, wrap self.weight.weight_loader to automatically trigger _maybe_fuse_f_proj once self.weight is loaded and moved off the meta device.

        self.f_a_weight.weight_loader = self._load_f_a_weight
        self.f_b_weight.weight_loader = self._load_f_b_weight
        original_weight_loader = self.weight.weight_loader
        def wrapped_weight_loader(*args, **kwargs):
            original_weight_loader(*args, **kwargs)
            self._maybe_fuse_f_proj()
        self.weight.weight_loader = wrapped_weight_loader
        self._f_a_loaded = False
        self._f_b_loaded = False

[gemini-code-assist]

When the model is initialized on the meta device, self.f_a_weight is also on the meta device. Calling param.data.copy_(loaded_weight) directly on a meta tensor will raise a RuntimeError.

We should check if the parameter is on the meta device and allocate it on the correct device (either self.weight.device if it is already loaded, or CPU/NPU) before copying.

    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.device == torch.device("meta"):
            device = self.weight.device if self.weight.device != torch.device("meta") else "cpu"
            param.data = torch.empty(param.shape, device=device, dtype=param.dtype)
        if param.shape != loaded_weight.shape:

[gemini-code-assist]

When the model is initialized on the meta device, self.f_b_weight is also on the meta device. Calling param.data.copy_(local_weight) directly on a meta tensor will raise a RuntimeError.

We should check if the parameter is on the meta device and allocate it on the correct device before copying.

Suggested change

	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:
	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 param.device == torch.device("meta"):
	device = self.weight.device if self.weight.device != torch.device("meta") else "cpu"
	param.data = torch.empty(param.shape, device=device, dtype=param.dtype)
	if loaded_weight.shape == param.shape:

[gemini-code-assist]

If self.weight is still on the meta device when _maybe_fuse_f_proj is called, we should skip the fusion and wait for self.weight to be loaded (which will trigger the fusion via the wrapped weight loader).

Suggested change

	@torch.no_grad()
	def _maybe_fuse_f_proj(self) -> None:
	if not self._f_a_loaded or not self._f_b_loaded:
	return
	@torch.no_grad()
	def _maybe_fuse_f_proj(self) -> None:
	if not self._f_a_loaded or not self._f_b_loaded:
	return
	if self.weight.device == torch.device("meta"):
	return

[Dawn952]

/cherry-pick releases/v0.27.1rc
[Bot]: cherry-pick completed successfully. New PR created: #15869