Repository navigation
[Kernel] Feed merged qkv into Kimi K3 KDA attention via a vllm-ascend-local kda_attention_qkv op - #15143
[Kernel] Feed merged qkv into Kimi K3 KDA attention via a vllm-ascend-local kda_attention_qkv op#15143Levi-JQ wants to merge 1 commit into
Conversation
|
👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:
If CI fails, you can run linting and testing checks locally according Contributing and Testing. |
Summary of ChangesHello, I'm Gemini Code Assist1! 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 KDA attention layer by removing unnecessary I/O reshaping operations (split/concat) surrounding the fused CausalConv1d. By updating the underlying NPU operator to natively handle the split q/k/v output layout, the PR eliminates significant latency in the attention sub-graph while maintaining compatibility with existing workflows. Highlights
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 AssistThe 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
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 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
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
Suggested PR Title:
[Ops][Feature] Support split-output variant for CausalConv1d in KDASuggested PR Summary:
### What this PR does / why we need it?
This PR introduces a split-output variant of the `CausalConv1d` operator (`npu_causal_conv1d_custom_3io`) to support KDA. By writing `q`, `k`, and `v` directly into three separate pre-allocated tensors, it eliminates the need for external chunking, rearranging, and contiguous operations, improving performance. The PR updates the operator definition, shape inference, PyTorch bindings, and meta functions, and refactors the KDA attention forward pass to pass the merged `qkv` tensor directly.
However, there are critical issues in the implementation:
1. In `npu_causal_conv1d_custom_3io`, the `EXEC_NPU_CMD` call for `aclnnCausalConv1d` omits the first output `y`, which shifts the subsequent outputs (`y_q`, `y_k`, `y_v`) into incorrect parameter slots, leading to compilation or runtime errors. An undefined `at::Tensor()` placeholder must be passed for `y`. Additionally, the existing `npu_causal_conv1d_custom` function must be updated to match the new signature.
2. The shape inference logic in `InferShapeCausalConv1d` assumes the input tensor is exactly 2D. It should be made dimension-agnostic to handle other configurations robustly.
### Does this PR introduce _any_ user-facing change?
No, this is an internal optimization for the Ascend backend's KDA implementation.
### How was this patch tested?
No tests were added in this PR. It is recommended to add unit tests verifying the correctness of the new `npu_causal_conv1d_custom_3io` operator and its shape inference logic.| EXEC_NPU_CMD(aclnnCausalConv1d, | ||
| x, | ||
| weight, | ||
| bias_opt, | ||
| conv_state, | ||
| query_start_loc_opt, | ||
| cache_indices_opt, | ||
| initial_state_mode_opt, | ||
| num_accepted_tokens_opt, | ||
| activation_mode, | ||
| pad_slot_id, | ||
| run_mode, | ||
| /*split_qkv=*/true, | ||
| y_q, | ||
| y_k, | ||
| y_v | ||
| ); |
There was a problem hiding this comment.
Adding the split_qkv attribute and optional outputs y_q, y_k, y_v to the CausalConv1d OpDef changes the generated aclnnCausalConv1d API signature. The generated signature expects all attributes and outputs in the exact order defined in the OpDef.
In npu_causal_conv1d_custom_3io, the first output y is omitted from the EXEC_NPU_CMD call, which shifts the remaining outputs (y_q, y_k, y_v) into the wrong parameter slots and leaves the last parameter missing. This will cause a compilation error or runtime crash.
To fix this, pass an undefined at::Tensor() for the unused y output.
Additionally, please note that the existing npu_causal_conv1d_custom function (not modified in this diff hunk) must also be updated to pass /*split_qkv=*/false and three undefined at::Tensor() placeholders for y_q, y_k, and y_v to match the new signature and avoid compilation failures.
EXEC_NPU_CMD(aclnnCausalConv1d,
x,
weight,
bias_opt,
conv_state,
query_start_loc_opt,
cache_indices_opt,
initial_state_mode_opt,
num_accepted_tokens_opt,
activation_mode,
pad_slot_id,
run_mode,
/*split_qkv=*/true,
at::Tensor(), // y is unused when split_qkv=true
y_q,
y_k,
y_v
);| gert::Shape singleShape = *xShape; | ||
| singleShape.SetDimNum(2); | ||
| int64_t last = xShape->GetDim(1); | ||
| singleShape.SetDim(1, last / 3); |
There was a problem hiding this comment.
The shape inference logic assumes that the input tensor x is always exactly 2D. If x is 3D or 4D (which is common in other configurations or future extensions), singleShape.SetDimNum(2) will incorrectly truncate the dimensions, and xShape->GetDim(1) will retrieve the sequence/token dimension instead of the channel dimension, leading to incorrect shape inference. Additionally, if x has fewer than 2 dimensions, GetDim(1) will cause an out-of-bounds access.
We should make the shape inference robust and dimension-agnostic by querying the actual number of dimensions and modifying the last dimension (the channel dimension).
| gert::Shape singleShape = *xShape; | |
| singleShape.SetDimNum(2); | |
| int64_t last = xShape->GetDim(1); | |
| singleShape.SetDim(1, last / 3); | |
| gert::Shape singleShape = *xShape; | |
| size_t dimNum = xShape->GetDimNum(); | |
| if (dimNum > 0) { | |
| int64_t last = xShape->GetDim(dimNum - 1); | |
| singleShape.SetDim(dimNum - 1, last / 3); | |
| } |
Keep fused_qkv's merged [N, 3*C] output all the way to the conv instead of splitting it just to re-concatenate inside the op. Register a vllm-ascend-local vllm::kda_attention_qkv custom op with the merged-qkv signature (same opaque-boundary mechanics as the upstream vllm::kda_attention: resolve the layer from no_compile_layers and call its _forward eagerly) instead of changing the shared upstream op signature — the upstream op, the reference path on other platforms, and stock vLLM compatibility stay untouched. The front 3 Slice + 1 ConcatD collapse to a single trim Slice; the result is bit-identical (qkv == cat(split(qkv))). The op is appended to compilation_config.splitting_ops at layer init so piecewise cudagraphs still split the compiled graph at the attention boundary (model construction runs before compilation; FULL cudagraph and eager modes don't consult the list).
Levi-JQ
changed the title
Levi-JQ
changed the title
What this PR does / why we need it?
NPU profiling of the Kimi K3 KDA (gated delta-net) layer shows the merged
fused_qkvprojection output is split into q/k/v just so the attention op can take three separate tensors, after which the pieces are re-concatenated (torch.cat) before the fusedCausalConv1d. The front 3 Slices + 1 ConcatD add no semantic value — the conv weights are already packed as[conv_size, 3*C]and the kernel natively processes q/k/v together — and are part of asplit → cat → conv → chunklayout round-trip that dominates the conv sub-graph time.This PR keeps the merged qkv all the way to the conv instead of splitting it just to re-concatenate. It registers a vllm-ascend-local
vllm::kda_attention_qkvcustom op with the merged-qkv signature, using the same opaque-boundary mechanics as the upstreamvllm::kda_attentionop (the impl resolves the layer fromno_compile_layersand calls its_forwardeagerly). The shared upstream op signature is deliberately not changed, so the upstream op, its(q, k, v)schema, and the reference path on other platforms stay untouched and vllm-ascend keeps working against stock vLLM. The result is bit-identical (qkv == cat(split(qkv))).The op is appended to
compilation_config.splitting_opsat layer init so piecewise cudagraphs still split the compiled graph at the attention boundary (model construction runs before compilation; FULL cudagraph and eager modes don't consult the list).The change is pure Python (
vllm_ascend/ops/kimi_kda.pyonly) — no kernel or.sorebuild required.Does this PR introduce any user-facing change?
Additive only:
vllm::kda_attention_qkv(registered on the vllm-ascend side); the upstreamvllm::kda_attentionop and its(q, k, v)schema are untouched.AscendKimiGatedDeltaNetAttention._forward(private method) signature changes to take the mergedqkv.vllm::kda_attention_qkvis appended tocompilation_config.splitting_opsat KDA layer init (no user action needed).How was this patch tested?
FULL_DECODE_ONLY, dspark speculative decoding): the same greedy prompt run on pre-PR code and on this PR produces token-identical outputs./start_profile→ identical decode request →/stop_profile,torch_npu.profiler.analyse, rank0op_statistic.csv): the ConcatD /CausalConv1d/RecurrentKdanumbers above;CausalConv1dandRecurrentKdainvocation counts and durations flat, confirming conv and SSM run exactly as before (pure I/O optimization).