[Feature] Optimize GDN non-spec prefill fallback metadata

[maoxx241]

What this PR does / why we need it?

This PR optimizes the GDN non-spec prefill fallback path for the Ascend Qwen3.5 / Qwen3Next linear-attention backend.

Background:

• causal_conv1d in the current non-spec prefill path still needs host-side metadata, and the old path prepared that metadata lazily in forward.
• chunk_gated_delta_rule also relied on repeated hot-path prepare_* helpers (prepare_chunk_indices, prepare_chunk_offsets, etc.) instead of consuming a step-scoped prebuilt metadata bundle.
• The fallback contract was also too permissive: when required non-spec prefill metadata was missing, the code could silently fall back to legacy recomputation paths instead of failing explicitly.

Changes in this PR:

• Introduce a unified non_spec_prefill_fallback_meta for the GDN non-spec prefill path.
  • causal_conv1d host metadata is prepared during attention metadata build and consumed in forward.
  • chunk_gated_delta_rule prebuilt metadata is also prepared during metadata build and threaded through gdn.py -> chunk.py.
• Add vllm_ascend/ops/triton/gdn_chunk_meta.py to build:
  • chunk_indices_chunk64
  • chunk_offsets_chunk64
  • update_chunk_offsets_chunk64
  • final_chunk_indices_chunk64
  • chunk_indices_large_block
  • block_indices_cumsum
• The builder-side chunk-meta path now has two execution modes:
  • CPU path: generate metadata with the same semantics as the original runtime prepare_* helpers.
  • NPU path: write directly into pooled device tensors, reuse shared seq_lens, and avoid per-call helper fallback logic in the hot path.
• Thread prebuilt chunk_offsets through chunk.py, chunk_o.py, and chunk_o_update.py so wrapper code no longer re-prepares offsets when prebuilt metadata is available.
• Tighten the patched prefill contract:
  • remove legacy silent fallback branches that re-materialize metadata in unexpected paths
  • raise explicitly when required non-spec prefill fallback metadata is missing
• Keep metadata semantics aligned with the original runtime path for mixed spec/non-spec, padding, zero-length, and long-sequence cases.

Why this is needed:

• moves non-spec prefill fallback metadata generation to the metadata-builder stage instead of doing it lazily inside per-layer forward hot paths
• makes the fallback contract explicit and debuggable
• keeps chunk_gated_delta_rule operator inputs consistent with the original prepare_* path while reducing repeated hot-path preparation work
• preserves correctness for mixed/padded edge cases while improving TTFT / TPS on the target workload

Does this PR introduce any user-facing change?

There is no API or interface change.

User-visible impact is limited to backend behavior on Ascend:

• the GDN non-spec prefill path is stricter and will now fail loudly if the expected fallback metadata contract is violated
• performance for the targeted Qwen3.5 / Qwen3Next GDN prefill path is improved on the validated benchmark shape

How was this patch tested?

Unit / correctness coverage:

• tests/ut/patch/worker/patch_common/test_patch_gdn_attn.py
• tests/ut/ops/test_gdn_chunk_meta.py
• tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_gdn_chunk_meta.py

The tests cover:

• causal_conv1d host args matching the original/runtime path
• chunk metadata parity with the original runtime prepare_* helpers
• CPU and NPU chunk-meta generation parity
• mixed spec/non-spec, padding, zero-length, and long-sequence edge cases
• strict failure when required non-spec prefill fallback metadata is missing
• prebuilt chunk_offsets wiring through the FLA wrappers

Additional validation:

• compared the current prebuilt chunk_gated_delta_rule metadata against the original non-prebuilt path on both CPU and NPU; for the tested long-sequence / mixed cases, the operator inputs were identical

Performance validation:

• methodology: for each code state, start the service once, run 5 benchmark rounds against the same warm service, and report the mean of the last 4 rounds
• model: Qwen3.5-35B-A3B
• target: 910B4 32G
• benchmark config:
  • tensor_parallel_size=4
  • --max-model-len 4096
  • --gpu-memory-utilization 0.9
  • --async-scheduling
  • --compilation-config '{"cudagraph_mode":"FULL_DECODE_ONLY","cudagraph_capture_sizes":[4,8,12,16]}'
  • --speculative-config '{"method":"qwen3_5_mtp","num_speculative_tokens":3}'
  • num_prompts=64
  • max_concurrency=16
  • output_len=1500
• comparison states:
  • baseline: pre-PR state dae3c99e
  • original PR implementation: cc3ca6f7
  • current implementation in this PR: 8000a085

Results (mean of last 4 rounds):

State	TTFT	TPS
baseline (dae3c99e)	1048.64 ms	162.26
original (cc3ca6f7)	1008.81 ms	163.60
current (8000a085)	978.58 ms	164.67

Delta:

• current vs baseline:
  • TTFT: -70.06 ms
  • TPS: +2.42
• current vs original:
  • TTFT: -30.22 ms
  • TPS: +1.07

Summary:

• relative to the original PR implementation, the latest scheme keeps improving TTFT and TPS

[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 non-speculative prefill fallback metadata for Gated Delta Networks (GDN). By moving metadata construction to the device and unifying the metadata structures, the changes reduce latency and improve the reliability of prefill operations. The update also includes necessary refactoring to support these changes across existing patch workers.

Highlights

• GDN Metadata Optimization: Introduced a new fallback metadata structure for GDN non-spec prefill, improving efficiency and robustness.
• Device-side Metadata Construction: Added build_chunk_meta_device to construct chunk metadata directly on the NPU, reducing CPU-GPU synchronization overhead.
• Refactored Metadata Handling: Consolidated causal conv1d and chunk metadata into a unified GDNPrefillFallbackMeta structure.
• Enhanced Testing: Added comprehensive unit tests for chunk metadata construction and verified consistency between device-side and runtime helpers.

🧠 New Feature in Public Preview: You can now enable Memory 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 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 counter productive. 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

This pull request optimizes Gated Delta Net (GDN) chunk metadata generation by implementing device-side Triton kernels and refactors the attention metadata builder to use a consolidated fallback metadata structure with buffer pooling. Feedback was provided regarding the need for consistent Triton availability guards and fallback implementations in the new gdn_chunk_meta.py file to prevent potential crashes in environments where Triton is not available.\n\nSuggested PR Title:\n\nmarkdown\n[Ops][Feature] Optimize GDN chunk metadata generation and refactor prefill fallback metadata\n\n\nSuggested PR Summary:\n\nmarkdown\n### What this PR does / why we need it?\nThis PR introduces `build_chunk_meta_device` in `vllm_ascend/ops/triton/gdn_chunk_meta.py` to generate GDN chunk metadata using Triton kernels, improving performance by moving metadata generation to the device. It also refactors `GDNAttentionMetadataBuilder` to use `GDNPrefillFallbackMeta`, which manages host-side metadata for causal conv1d and chunked prefill, and implements buffer pooling to reduce host-side allocation overhead.\n\n### Does this PR introduce _any_ user-facing change?\nNo.\n\n### How was this patch tested?\nThis patch was tested with new unit tests in `tests/ut/ops/test_gdn_chunk_meta.py` and updated tests in `tests/ut/patch/worker/patch_common/test_patch_gdn_attn.py`.\n

[gemini-code-assist]

There's an inconsistency in handling environments where Triton is not available. The call to _build_chunk_counts_kernel has a fallback path, but the subsequent calls to _build_chunk_offsets_kernel and _build_final_chunk_indices_kernel do not. This will cause a crash if the else branch for _build_chunk_counts_kernel is taken.

All Triton kernel calls should be consistently guarded with a check for Triton's availability, and a fallback implementation should be provided. Using torch.cumsum in the fallback path would be consistent with the CPU implementations seen elsewhere in the codebase.

    if chunk_offsets is not None:
        if hasattr(_build_chunk_offsets_kernel, "__getitem__"):
            block_size = 256
            grid = (_cdiv(num_seqs + 1, block_size),)
            _build_chunk_offsets_kernel[grid](
                chunk_counts_ptr=chunk_counts,
                out_offsets_ptr=chunk_offsets,
                num_seqs=num_seqs,
                ADD_ONE=0,
                BLOCK_SIZE=block_size,
            )
        else:
            chunk_offsets[0] = 0
            if num_seqs > 0:
                torch.cumsum(chunk_counts, dim=0, out=chunk_offsets[1:])

    if update_chunk_offsets is not None:
        if hasattr(_build_chunk_offsets_kernel, "__getitem__"):
            block_size = 256
            grid = (_cdiv(num_seqs + 1, block_size),)
            _build_chunk_offsets_kernel[grid](
                chunk_counts_ptr=chunk_counts,
                out_offsets_ptr=update_chunk_offsets,
                num_seqs=num_seqs,
                ADD_ONE=1,
                BLOCK_SIZE=block_size,
            )
        else:
            update_chunk_offsets[0] = 0
            if num_seqs > 0:
                torch.cumsum(chunk_counts + 1, dim=0, out=update_chunk_offsets[1:])

    if out_final_chunk_indices is not None:
        if hasattr(_build_final_chunk_indices_kernel, "__getitem__"):
            block_size = 256
            grid = (_cdiv(num_seqs, block_size),)
            _build_final_chunk_indices_kernel[grid](
                update_chunk_offsets_ptr=update_chunk_offsets,
                out_final_chunk_indices_ptr=out_final_chunk_indices,
                num_seqs=num_seqs,
                BLOCK_SIZE=block_size,
            )
        else:
            if num_seqs > 0:
                torch.cumsum(chunk_counts + 1, dim=0, out=out_final_chunk_indices)
                out_final_chunk_indices.sub_(1)

[weijinqian0]

The operator UT needs to be supplemented.

[weijinqian0]

Supplementing the patch deletion plan.

[github-actions]

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

[github-actions]

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