[MRV2][Performance][Kernel] Use windowed block-table gather for V2 slot mappings

[Liamup777]

What this PR does / why we need it?

Replace full block-table row staging in the Ascend V2 slot-mapping Triton kernel with a bounded contiguous window load followed by tl.gather.

This bounds UB use independently of block-table row length. The window is sized from the smallest kernel block size, because the upstream V0.28.0 block_sizes_tensor used by slot mapping stores kernel_block_sizes.

The kernel preserves the upstream CP mapping behavior, narrows positions to INT32, and computes block offsets with multiply/subtract instead of remainder. The V2 operator documentation is moved under ops/triton/docs/v2.

Does this PR introduce any user-facing change?

No.

How was this patch tested?

• bash format.sh — passed.

• bash format.sh ci — passed.

• SHELLCHECK_OPTS="--exclude=SC2046,SC2006,SC2086" pre-commit run --all-files --hook-stage manual --show-diff-on-failure — passed.

• python -m compileall -q vllm_ascend/ops/triton/v2/block_table/compute_slot_mappings.py vllm_ascend/worker/v2/block_table.py tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_compute_slot_mapping.py — passed.

• NPU UT on lab-worker-a3-01 passed on the preceding revision:

  python -m pytest -sv tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_compute_slot_mapping.py

  Result: 4 passed, 14 warnings in 17.74s. This covers CP=1/2/4, empty requests, cross-block/cross-tile ranges, mixed kernel block sizes, padding, and the out interface. Re-run is pending after the reviewer-requested tile-size source deduplication; the change preserves the configured value and kernel algorithm.

• Per-case NPU msprof op comparison against the V0.28.0 upstream kernel. Task Duration(us):

  Case	Grid	Upstream	Ascend	Speedup
  tile-1	1x2	114.598	3.680	31.14x
  tile-64	1x65	229.255	6.680	34.32x
  tile-1024	1x1025	2934.781	39.399	74.49x
  long-1x8192	1x2	786.244	10.920	72.00x
  long-64x8192	1x65	1573.109	20.960	75.05x
  long-64x32768	1x65	6181.396	70.719	87.41x
  cp2-tile-64-rank0	1x65	498.110	43.579	11.43x
  cp4-tile-64-rank2	1x65	493.110	44.019	11.20x

• vLLM main: vllm-project/vllm@b2f6858

[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]

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 Ascend V2 slot-mapping Triton kernel by transitioning from full block-table row staging to a bounded contiguous window load approach. This change significantly improves memory efficiency by decoupling Unified Buffer (UB) usage from the length of block-table rows. The implementation now dynamically determines window sizes based on kernel block sizes, ensuring consistent performance. Additionally, documentation has been reorganized to align with the new structure, and the test suite has been updated to verify the correctness of the new windowing logic.

Highlights

• Kernel Optimization: Implemented a windowed block-table gather in the Ascend V2 slot-mapping Triton kernel to replace full-row staging.
• UB Usage Management: Bounded Unified Buffer (UB) usage independently of block-table row length by utilizing a contiguous window load.
• Dynamic Window Sizing: Refactored AscendBlockTables to calculate window sizes dynamically based on kernel block sizes rather than static staging limits.
• Documentation: Migrated and updated documentation for the slot-mapping kernel to the new ops/triton/docs/v2 directory.
• Test Suite Updates: Updated test cases to validate the new windowing logic and removed obsolete staging parameters.

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. ↩

[gemini-code-assist]

Code Review

Suggested PR Title:

[Ops][Feature] Implement bounded window gather for Ascend block table slot mapping

Suggested PR Summary:

### What this PR does / why we need it?
This PR replaces the full-row block-table staging mechanism with a bounded window gather approach in `_compute_slot_mappings_kernel` for Ascend. Instead of staging the entire block-table row (which could exceed UB for large rows), it calculates a power-of-two `BLOCK_TABLE_WINDOW_SIZE` based on the smallest kernel block size and stages only the contiguous portion of the request row used by the current token tile. This removes the need for `_MAX_STAGED_BLOCK_TABLE_PAD_SIZE` and `USE_BLOCK_TABLE_STAGING`.

Feedback:
The reviewer suggests avoiding hardcoding the token tile size `1024` in both `__init__` and the kernel launch by defining an instance variable `self._triton_block_size = 1024` to ensure consistency and prevent potential out-of-bounds indexing.

### Does this PR introduce _any_ user-facing change?
No.

### How was this patch tested?
The patch was tested using the existing single-operator test suite: `pytest -sv tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_compute_slot_mapping.py`.

[gemini-code-assist]

The token tile size 1024 is hardcoded both in the window size calculation here and in the kernel launch parameters (TRITON_BLOCK_SIZE=1024 on line 101). If the block size is changed in one place but not the other, it can lead to silent out-of-bounds indexing or correctness bugs.

To prevent this, define an instance variable self._triton_block_size = 1024 in __init__ and use it in both places.

Suggested change

	# block_sizes_tensor stores kernel_block_sizes, which determine the
	# number of block-table entries touched by one token tile. Use the
	# smallest kernel block size to form one safe constexpr window for all
	# groups, without staging a whole row.
	min_kernel_block_size = min(kernel_block_sizes)
	window_size = (1024 + min_kernel_block_size - 1) // min_kernel_block_size + 1
	self._block_table_window_size = triton.next_power_of_2(window_size)
	self._triton_block_size = 1024
	# block_sizes_tensor stores kernel_block_sizes, which determine the
	# number of block-table entries touched by one token tile. Use the
	# smallest kernel block size to form one safe constexpr window for all
	# groups, without staging a whole row.
	min_kernel_block_size = min(kernel_block_sizes)
	window_size = (self._triton_block_size + min_kernel_block_size - 1) // min_kernel_block_size + 1
	self._block_table_window_size = triton.next_power_of_2(window_size)

[Liamup777]

Thanks, agreed. I introduced self._triton_block_size = 1024 in init and use it for both window-size calculation and the kernel TRITON_BLOCK_SIZE launch argument. The manually constructed out-interface test fixture now initializes the same attribute.

[gemini-code-assist]

Use the newly defined self._triton_block_size instead of the hardcoded 1024 constant to ensure consistency with the window size calculation.

Suggested change

	CP_INTERLEAVE=self.cp_interleave,
	PAD_ID=PAD_SLOT_ID,
	TRITON_BLOCK_SIZE=1024,
	BLOCK_TABLE_PAD_SIZE=self._block_table_pad_size,
	USE_BLOCK_TABLE_STAGING=(self._block_table_pad_size <= _MAX_STAGED_BLOCK_TABLE_PAD_SIZE),
	BLOCK_TABLE_WINDOW_SIZE=self._block_table_window_size,
	)
	CP_INTERLEAVE=self.cp_interleave,
	PAD_ID=PAD_SLOT_ID,
	TRITON_BLOCK_SIZE=self._triton_block_size,
	BLOCK_TABLE_WINDOW_SIZE=self._block_table_window_size,
	)

[Liamup777]

Updated in the same change: the launch now passes TRITON_BLOCK_SIZE=self._triton_block_size, so it cannot diverge from the window-size calculation.

[AuroraEmiya]

Check if Gemini say it right : )

[Liamup777]

Confirmed. Gemini’s consistency concern was valid, and the requested tile-size deduplication has now been applied.

[AuroraEmiya]

Check if there is any performance fallback compared to which you already submitted

[AuroraEmiya]

LGTM