prefillnocache branch use fia op. #4396
Conversation
Signed-off-by: wangxiaoxin-sherie <[email protected]>
|
👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:
If CI fails, you can run linting and testing checks locally according Contributing and Testing. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request updates the prefill-no-cache attention mechanism to use the npu_fused_infer_attention_score operator. This is a positive change, likely for performance or feature alignment. However, the refactoring has left behind some dead code and unused function parameters in _forward_prefill_no_cache, which I've commented on. The related change in model_runner_v1.py to update the attention mask generation is consistent with this new operator.
| output, _ = torch_npu.npu_fused_infer_attention_score.out( | ||
| query=query, | ||
| key=key, | ||
| value=value, | ||
| atten_mask=attn_metadata.attn_mask, | ||
| block_table=None, | ||
| input_layout="TND", | ||
| block_size=128, | ||
| actual_seq_lengths=attn_metadata.actual_seq_lengths_q, | ||
| actual_seq_lengths_kv=attn_metadata.seq_lens_list, | ||
| num_key_value_heads=self.num_kv_heads, | ||
| num_heads=self.num_heads, | ||
| scale=self.scale, | ||
| sparse_mode=3, | ||
| ) | ||
| assert output is not None | ||
| return output[:num_tokens, :, :] | ||
| return output |
There was a problem hiding this comment.
The switch to npu_fused_infer_attention_score has left some surrounding code in an inconsistent state, which impacts maintainability.
Specifically:
- The
_forward_prefill_no_cachefunction's parametersoutputandnum_tokensare now unused. - The preceding
if is_310p():block (lines 347-356) contains dead code:- The
outputtensor is aligned, but then theoutputvariable is immediately reassigned, making the alignment useless. - The
maskvariable is computed, but the new operator call usesattn_metadata.attn_maskinstead.
- The
This unused code and parameters should be removed to avoid confusion and improve code clarity.
1 participant
What this PR does / why we need it?
prefillnocache branch use fia op.
Does this PR introduce any user-facing change?
How was this patch tested?