[bugfix][torchair] fix kv_nz accuracy problem and remove redundant reshape_and_cache operation

[linfeng-yuan]

What this PR does / why we need it?

This PR removes redundant calling of reshape_and_cache operation at prefilling stage with torchair graph mode. This reduces prefilling latency as well as fixes accuracy problem while enable_kv_nz is enabled. Although #2988 fixes enable_kv_nz accuracy problem, the output tokens with deepseek is inaccurate, leading to a decline in benchmark scoring.

Does this PR introduce any user-facing change?

No.

How was this patch tested?

We run e2e online serving and accuracy test containing eager mode with enable_shared_expert_dp and torchair graph mode with enable_kv_nz.

• vLLM version: v0.10.2
• vLLM main: vllm-project/vllm@6d8246a

[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 correctly removes a redundant _npu_reshape_and_cache operation that was being called in torchair graph mode. This is a good simplification, as caching is already handled by npu_kv_rmsnorm_rope_cache in that scenario. However, an important assertion checking the kv_cache size was also removed. I've recommended re-adding it to ensure code robustness and prevent potential runtime errors.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 71.96%. Comparing base (1bbb20e) to head (d745a8e).
⚠️ Report is 81 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3066      +/-   ##
==========================================
- Coverage   74.76%   71.96%   -2.81%     
==========================================
  Files         150      168      +18     
  Lines       20891    23544    +2653     
==========================================
+ Hits        15620    16943    +1323     
- Misses       5271     6601    +1330     

Flag	Coverage Δ
unittests	71.96% <100.00%> (-2.81%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[jianzs]

@linfeng-yuan Thanks a lot!!!

[jianzs]

@linfeng-yuan This pull request only fixes one accuracy problem. Tests show accuracy is fine without KV NZ enabled, but still problematic when it's on, even after applying this change. The GSM-8K benchmark scores are still too low with KV NZ active....

[jianzs]

@linfeng-yuan The torch_npu.atb.npu_ring_mla and torch_npu._npu_flash_attention functions were used in the prefill stage, but the code doesn't seem to have any adaptations for KV NZ. This might be the cause of the problem?

[jianzs]

This PR cannot be merged until kv_nz is supported during the prefill phase.