[Feature][Ops] Add Triton KeyPool compression and pooled indexing

[lijiahang226]

What this PR does / why we need it?

Refs #15665

Add Triton KeyPool state compression, paged cache writes, and pooled-key selection for GLM-5.3-Flash. Allow IndexerWrapper to select a model-provided backend while preserving checkpoint parameter names. Both operators reuse the shared next_power_of_2 helper.

Include single-operator accuracy tests and documentation covering formulas, parameters, layout constraints, graph replay, and test commands. Model-specific backend integration is handled separately.

Does this PR introduce any user-facing change?

Yes. Provides the Triton pooled-indexing operators and a backend-selection hook. This PR does not enable the model backend by itself or change cache allocation and grouping.

How was this patch tested?

Validated on Atlas A3 with PyTorch 2.10.0, torch-npu 2.10.0.post4, and Triton-Ascend 3.2.0:

• 10 tests passed in tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_glm5next_kpool_triton.py: CPU-reference accuracy, historical windows, rollback, paging, noncontiguous storage, invalid slots, graph padding, empty inputs, and eager/graph replay.
• 9 tests passed in tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_glm5next_pool_key_indexer_triton.py: multiple requests, pool capacities, paged caches, token chunking, causal tails, and eager/graph replay with changing inputs.
• 16 tests passed across tests/ut/models/test_glm5next_indexer_backend.py and tests/ut/ops/test_mla.py.

The CI follow-up adds a type annotation to the test reference data without changing test behavior. Hardware coverage is limited to Atlas A3; these tests establish operator accuracy, not model-level throughput.

• vLLM main: vllm-project/vllm@a97dacb

[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 implements high-performance Triton operators to handle state compression and pooled indexing for the GLM-5.3-Flash model. By moving these operations into Triton, the implementation avoids the overhead of multiple small torch operations. Additionally, the PR enhances the IndexerWrapper to allow for flexible, model-provided backend selection, ensuring that parameter names remain consistent across different model architectures.

Highlights

• Triton Operators: Added two new Triton operators, glm5_next_kpool_state_compress_and_write_cache_triton and glm5_next_lightning_indexer_triton, to optimize state compression and pooled-key indexing for GLM-5.3-Flash.
• Backend Selection: Refactored IndexerWrapper to support dynamic backend selection while ensuring model-specific parameter names are preserved.
• Testing: Introduced comprehensive end-to-end and unit tests to validate the new Triton operators and the backend selection hook.

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:

[Ops][Feature] Add Triton kernels and backend wrapper for GLM5 Next KPool indexer

Suggested PR Summary:

### What this PR does / why we need it?
This pull request introduces Triton fast paths for the GLM5 Next KPool indexer. Specifically, it adds:
1. `glm5_next_kpool_state_compress_and_write_cache_triton` for the pre-compress sequence, which gathers and compresses state windows using softmax and writes them to the paged indexer cache.
2. `glm5_next_lightning_indexer_triton` for scoring compressed pools against head-weighted queries.
3. Updates to `IndexerWrapper` in `vllm_ascend/ops/mla.py` to dynamically select the appropriate backend and preserve checkpoint-visible parameters.
4. Corresponding unit and end-to-end tests.

Feedback:
The review comments identify critical issues in the Triton kernels where out-of-bounds address calculations (e.g., loading from `cum_query_lens_ptr + req_offsets` when `req_offsets >= num_reqs`, or indexing `state_block_table_ptr` and `indexer_block_table_ptr` with out-of-bounds pages) can cause GPU page faults or undefined behavior, even when masked. Clamping these offsets/indices using `tl.where` before address calculation is highly recommended.

### Does this PR introduce _any_ user-facing change?
No, these are internal performance optimizations and backend support for GLM5 Next.

### How was this patch tested?
New unit tests (`test_glm5next_indexer_backend.py`) and end-to-end Triton kernel tests (`test_glm5next_kpool_triton.py`, `test_glm5next_pool_key_indexer_triton.py`) have been added.

[gemini-code-assist]

In Triton, loading from an out-of-bounds address can cause GPU page faults or undefined behavior, even if the load is masked. Since cum_query_lens is not padded to REQ_POW2, cum_query_lens_ptr + req_offsets will access out-of-bounds memory when req_offsets >= num_reqs. We should clamp req_offsets to a safe index (e.g., 0) using tl.where before performing the load.

Suggested change

	req_offsets = tl.arange(0, REQ_POW2)
	query_ends = tl.load(
	cum_query_lens_ptr + req_offsets,
	mask=req_offsets < num_reqs,
	other=2147483647,
	)
	req_offsets = tl.arange(0, REQ_POW2)
	safe_req_offsets = tl.where(req_offsets < num_reqs, req_offsets, 0)
	query_ends = tl.load(
	cum_query_lens_ptr + safe_req_offsets,
	mask=req_offsets < num_reqs,
	other=2147483647,
	)

[gemini-code-assist]

In Triton, even if a load is masked, the address calculation must not point to invalid memory to avoid page faults. When page >= state_max_pages, the address state_block_table_ptr + ... + page * state_block_table_stride_page is out of bounds. We should clamp page to a safe index (e.g., 0) using tl.where before the address calculation.

Suggested change

	physical = tl.load(
	state_block_table_ptr + req_id * state_block_table_stride_req + page * state_block_table_stride_page,
	mask=history_valid,
	other=-1,
	).to(tl.int64)
	safe_page = tl.where(history_valid, page, 0)
	physical = tl.load(
	state_block_table_ptr + req_id * state_block_table_stride_req + safe_page * state_block_table_stride_page,
	mask=history_valid,
	other=-1,
	).to(tl.int64)

[gemini-code-assist]

In Triton, loading from an out-of-bounds address can cause GPU page faults or undefined behavior, even if the load is masked. Since cum_query_lens is not padded to REQ_POW2, cum_query_lens_ptr + req_offsets will access out-of-bounds memory when req_offsets >= num_reqs. We should clamp req_offsets to a safe index (e.g., 0) using tl.where before performing the load.

Suggested change

	req_offsets = tl.arange(0, REQ_POW2)
	query_ends = tl.load(cum_query_lens_ptr + req_offsets, mask=req_offsets < num_reqs, other=2147483647)
	req_offsets = tl.arange(0, REQ_POW2)
	safe_req_offsets = tl.where(req_offsets < num_reqs, req_offsets, 0)
	query_ends = tl.load(cum_query_lens_ptr + safe_req_offsets, mask=req_offsets < num_reqs, other=2147483647)

[gemini-code-assist]

In Triton, even if a load is masked, the address calculation must not point to invalid memory to avoid page faults. When logical_pages is out of bounds (e.g., when in_range is False), the address calculation for indexer_block_table_ptr can access out-of-bounds memory. We should clamp logical_pages to a safe index (e.g., 0) using tl.where before the address calculation.

Suggested change

	physical_blocks = tl.load(
	indexer_block_table_ptr + req_id * block_table_stride_req + logical_pages * block_table_stride_page,
	mask=in_range,
	other=0,
	).to(tl.int64)
	safe_logical_pages = tl.where(in_range, logical_pages, 0)
	physical_blocks = tl.load(
	indexer_block_table_ptr + req_id * block_table_stride_req + safe_logical_pages * block_table_stride_page,
	mask=in_range,
	other=0,
	).to(tl.int64)

[ZT-AIA]

Please supplement the single-operator use cases and the operator documentation.

[lijiahang226]

Please supplement the single-operator use cases and the operator documentation.

Operation ducumentation has been added.