[MRV2][Bugfix] Fix two triton ops errors in model_runner_v2 num_nans_kernel and apply_panalties

[AuroraEmiya]

What this PR does / why we need it?

This PR has done two things:

1. This PR fix num_nans wrong import way of libdevice in Ascend env with a right way with CANN
   therefore we add new patch in vllm_ascend/patch/worker/patch_v2/patch_triton.py‎
   metrics_logits.libdevice = triton.language.extra.cann.libdevice
2. This PR also fix over-grid limit with large input in ops apply_palties

Does this PR introduce any user-facing change?

No

How was this patch tested?

uts are added in

1. tests/ut/patch/worker/test_patch_v2_num_nans.py
2. tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_penality.py

• vLLM version: v0.26.0
• vLLM main: vllm-project/vllm@0351e9a

[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 introduces significant improvements to the Triton-based penalty operations within the model runner. By splitting the monolithic penalty kernel into specialized kernels for vocabulary and draft tokens, the implementation achieves better efficiency and robustness. These changes ensure the system can handle larger shapes without crashing and align the penalty logic more closely with reference implementations, while also improving the underlying hardware-specific configuration.

Highlights

• Triton Kernel Refactoring: Refactored the penalty application logic into two distinct kernels: a vocabulary-wide kernel for general penalty application and a draft-specific kernel for handling speculative token penalties.
• Performance and Stability: Improved performance by optimizing memory access patterns and added support for larger shapes to prevent kernel crashes.
• Infrastructure Updates: Added a new test case for verifying CANN libdevice usage and updated the penalty module to include hardware-aware vector core detection.

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.

---

Tip

💡 Consider Linking a Related Issue or RFC

Your PR title contains the [BugFix]/[Feature] tag, indicating a bug fix or new feature.

Linking a related issue or RFC in the PR description is strongly encouraged — it gives reviewers helpful context and speeds up the review. You can use any of these keywords:

• Fixes #<issue_number>
• Closes #<issue_number>
• Resolves #<issue_number>
• Refs #<rfc_or_issue_number> (for RFCs)

🙏 Thanks for helping us keep the project well-organized!

[gemini-code-assist]

Code Review

Suggested PR Title:

[Ops][Feature] Implement optimized Triton penalty kernels

Suggested PR Summary:

### What this PR does / why we need it?
This pull request refactors and optimizes the penalty application logic by introducing two specialized Triton kernels: `_apply_penalties_vocab_kernel` for full-vocabulary penalties and `_apply_penalties_draft_kernel` for draft token corrections. It also updates the bincount kernel and integrates the CANN libdevice for Triton-Ascend.

Feedback:
Two critical issues were identified in the Triton kernels where out-of-bounds token IDs (such as negative padding tokens or values exceeding the vocabulary size) could lead to illegal memory access or out-of-bounds writes on the GPU. Guard conditions should be added to ensure `draft_token` and `prompt_tokens` are within valid bounds before indexing.

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

### How was this patch tested?
The changes are tested using the newly added unit test `test_patch_v2_num_nans_uses_cann_libdevice_only` and the updated e2e test suite in `test_penality.py` including large shape handling.

[gemini-code-assist]

The draft_token loaded from token_ids_ptr is used directly as an index to load from output_counts_row_ptr and logits_row_ptr, and to store back to logits_row_ptr. If draft_token is out of bounds (e.g., negative padding tokens like -1 or values exceeding the vocabulary size), this will cause an illegal memory access on the GPU, leading to a kernel crash. We should guard this by checking that draft_token is within [0, vocab_size). Since logits_stride is equal to vocab_size for the 2D logits tensor, we can use logits_stride as the upper bound.

Suggested change

	if earlier_count == 0:
	if (earlier_count == 0) & (draft_token >= 0) & (draft_token < logits_stride):

[gemini-code-assist]

In _bincount_kernel, prompt_tokens loaded from all_token_ids_ptr can contain out-of-bounds values (such as negative padding tokens like -1 or values exceeding the vocabulary size). Using these directly to compute packed_indices and performing tl.atomic_or will result in out-of-bounds memory writes and GPU crashes. We should guard the mask with a check that prompt_tokens is within [0, vocab_size). Since output_bin_counts_stride is equal to vocab_size, we can use it as the upper bound.

Suggested change

	if block_idx * BLOCK_SIZE < prompt_len:
	mask = block < prompt_len
	prompt_tokens = tl.load(all_token_ids_ptr + req_state_idx * all_token_ids_stride + block, mask=mask)
	idx = prompt_tokens // 32
	bit_idx = prompt_tokens % 32
	bit = tl.full((BLOCK_SIZE,), 1, tl.int32) << bit_idx
	prompt_tokens = tl.load(
	all_token_ids_ptr + req_state_idx * all_token_ids_stride + block,
	mask=mask,
	)
	packed_indices = prompt_tokens // 32
	bit_indices = prompt_tokens % 32
	bits = tl.full((BLOCK_SIZE,), 1, tl.int32) << bit_indices
	tl.atomic_or(
	prompt_bin_mask_ptr + req_state_idx * prompt_bin_mask_stride + idx,
	bit,
	prompt_bin_mask_ptr + req_state_idx * prompt_bin_mask_stride + packed_indices,
	bits,
	mask=mask,
	)
	if block_idx * BLOCK_SIZE < prompt_len:
	mask = block < prompt_len
	prompt_tokens = tl.load(
	all_token_ids_ptr + req_state_idx * all_token_ids_stride + block,
	mask=mask,
	)
	mask = mask & (prompt_tokens >= 0) & (prompt_tokens < output_bin_counts_stride)
	packed_indices = prompt_tokens // 32
	bit_indices = prompt_tokens % 32
	bits = tl.full((BLOCK_SIZE,), 1, tl.int32) << bit_indices
	tl.atomic_or(
	prompt_bin_mask_ptr + req_state_idx * prompt_bin_mask_stride + packed_indices,
	bits,
	mask=mask,
	)

[yiz-liu]

Please see if you need add explanation for the new patch.

[AuroraEmiya]

Please see if you need add explanation for the new patch.

That 's fair. We have added here

What this PR does / why we need it?

This PR has done two things:

1. This PR fix num_nans wrong import way of libdevice in Ascend env with a right way with CANN
   therefore we add new patch in vllm_ascend/patch/worker/patch_v2/patch_triton.py‎
   metrics_logits.libdevice = triton.language.extra.cann.libdevice
2. This PR also fix over-grid limit with large input in ops apply_palties and refactor the kernel gaining a 80x speedup

Does this PR introduce any user-facing change?

No

How was this patch tested?

uts are added in

1. tests/ut/patch/worker/test_patch_v2_num_nans.py
2. tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_penality.py

• vLLM version: v0.26.0
• vLLM main: vllm-project/vllm@d02df74

[github-actions]

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

[yiz-liu]

/rerun
[Bot]: rerun completed.

Failed:

• E2E: still in progress, retry /rerun after completion