[MODELRUNNERV2]fix penality ops

[shiyuan680]

What this PR does / why we need it?

fix penality ops for new version, and achieved a 10% performance improvement

Does this PR introduce any user-facing change?

How was this patch tested?

pytest ‎tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_penality.py

• vLLM version: v0.16.0
• vLLM main: vllm-project/vllm@15d76f7

[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 a significant refactoring of the penalty application mechanism within MODELRUNNERV2. The primary goal is to enhance the system's ability to handle penalties correctly in the context of speculative decoding, while also improving the modularity of the sampling process by separating penalty and temperature application. The changes involve updating the Triton kernel, its Python wrapper, and the corresponding test infrastructure to support the new, more granular penalty calculation logic.

Highlights

• Refactored Penalty Application: The core penalty application logic has been refactored, separating penalty calculations from temperature application and removing the SamplingMetadata class for a more direct parameter passing approach.
• Speculative Decoding Support: The penalty calculation now explicitly incorporates support for speculative decoding by considering token_ids and expanded_local_pos to compute cumulative draft counts, ensuring accurate penalty application for draft tokens.
• Triton Kernel Updates: The underlying Triton kernel (_penalties_kernel) has been updated to reflect the new API, handle speculative decoding parameters, and optimize block sizes for NPU performance.
• Test Suite Alignment: The penalty test suite (test_penality.py) has been updated to align with the new penalty application API, including changes to test data generation, function signatures, and the removal of temperature-related test logic.

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

Changelog

• tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_penality.py
  • Updated the test function pytorch_apply_penalties_and_temperature to pytorch_apply_penalties and modified its signature to accept individual penalty tensors and speculative decoding parameters.
  • Removed the SamplingMetadata class and adjusted create_test_data to return individual tensors instead of a metadata object.
  • Renamed NUM_REQS to NUM_TOKENS and added NUM_SPECULATIVE_TOKENS for testing speculative decoding scenarios.
• vllm_ascend/worker/v2/sample/penalties.py
  • Renamed _penalties_and_temperature_kernel to _penalties_kernel and apply_penalties_and_temperature to apply_penalties.
  • Modified the _penalties_kernel to accept new parameters for speculative decoding (token_ids_ptr, expanded_local_pos_ptr, MAX_SPEC_LEN) and to compute total_counts by combining base output counts with draft counts.
  • Separated temperature application logic from the penalty kernel, focusing the kernel solely on penalties.
  • Adjusted BLOCK_SIZE and introduced INNER_BLOCK_SIZE for Triton kernel optimization on NPU.
• vllm_ascend/worker/v2/sample/sampler.py
  • Modified the AscendSampler to call the new apply_penalties function with updated arguments, aligning with the refactored API.
  • Introduced a conditional check to apply penalties only when use_penalty is active, improving efficiency by skipping unnecessary computations.
  • Explicitly separated temperature application into a distinct call to apply_temperature, reflecting the modular design.

Activity

• The pull request was initiated by shiyuan680.
• The changes focus on refactoring penalty operations for a new version of MODELRUNNERV2, as indicated by the PR title and description.
• The modifications span across core penalty logic, Triton kernel implementation, and corresponding unit tests to ensure correctness and performance.

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.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

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 refactors the penalty application logic, separating it from temperature application and adding support for speculative decoding. The changes are implemented in a Triton kernel, its Python wrapper, and the corresponding test file. A critical out-of-bounds read vulnerability has been identified in the apply_penalties function within vllm_ascend/worker/v2/sample/penalties.py. This vulnerability stems from penalty tensors being incorrectly sliced by the batch size (num_tokens) instead of the number of sequence slots, which could lead to memory corruption or information leakage. Additionally, a high-severity performance issue was found in the Triton kernel.

[gemini-code-assist]

A critical out-of-bounds read vulnerability exists here. The penalties tensor is incorrectly constructed by slicing repetition_penalty, frequency_penalty, and presence_penalty up to num_tokens. While num_tokens is the batch size, the Triton kernel indexes these tensors by req_state_idx (sequence slot index), which can exceed num_tokens. This mismatch can lead to an out-of-bounds read, potentially causing information leakage from adjacent GPU memory or a GPU kernel/driver crash. The slicing should be removed to pass the full penalty tensors to the kernel, as suggested.

Suggested change

	penalties = torch.stack([
	repetition_penalty[:num_tokens],
	frequency_penalty[:num_tokens],
	presence_penalty[:num_tokens]
	], dim=1).contiguous()
	penalties = torch.stack([
	repetition_penalty,
	frequency_penalty,
	presence_penalty
	], dim=1).contiguous()

[gemini-code-assist]

The creation of bit_masks is inefficient. It constructs a large tensor of shape (INNER_BLOCK_SIZE // 32, 32) where each row is identical. This can be simplified to create a (1, 32) tensor and rely on Triton's broadcasting capabilities. This change will improve both readability and performance by avoiding the allocation and manipulation of a large, redundant tensor.

Suggested change

	bit_masks = tl.full((INNER_BLOCK_SIZE // 32, 32), 1, dtype=tl.int32) << tl.arange(0, 32)
	bit_masks = (1 << tl.arange(0, 32))[None, :]

[Ronald1995]

The function don't follow vllm's commit 4034c3d32e30d01639459edd3ab486f56993876d.
you don't need to modify sampler, please just follow up #7110

[Ronald1995]

please mention RFC #5208 on your PR.