[Feature] Optimize forward metadata collection across dp ranks

[jianzs]

What this PR does / why we need it?

This PR introduces two optimizations for cases where data parallel size > 1:

1. Eliminates DP communication in set_forward_context
2. Implements HCCL for DP metadata communication, resulting in significant performance improvements for large DP configurations
   • Achieves ~20ms latency reduction with DP size of 64

Does this PR introduce any user-facing change?

no

How was this patch tested?

CI passed.

[jianzs]

@NeverRaR PTAL

[NeverRaR]

lgtm

[Copilot]

Pull Request Overview

This PR optimizes how forward-pass metadata is collected and communicated across data-parallel ranks by removing the previous all-reduce and introducing an HCCL-based all-gather approach.

• Enforce that dummy batch execution only runs under data parallelism and refactor execute_dummy_batch to use per-rank metadata.
• Replace dist.all_reduce with HCCL all_gather in _get_forward_metadata_across_dp and update callers to handle the Tensor of per-rank token counts.
• Propagate num_tokens_across_dp through dummy runs and forward contexts, masking sentinel values before the pass.

Reviewed Changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
vllm_ascend/worker/worker_v1.py	Added assertion for dp_size > 1, refactored dummy-run logic to use HCCL per-rank metadata.
vllm_ascend/worker/model_runner_v1.py	Swapped all_reduce for all_gather under get_dp_group(), changed method signature and updated callers to handle a Tensor of metadata.

Comments suppressed due to low confidence (1)

vllm_ascend/worker/model_runner_v1.py:622

• Add unit or integration tests covering the dp_size > 1 aggregation path to verify that all_gather produces the correct combined metadata and that the masked_fill_ logic correctly replaces sentinel values.

            local_forward_metadata)

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 52.36%. Comparing base (c30ddb8) to head (6100e0d).
Report is 75 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #1593       +/-   ##
===========================================
+ Coverage   27.39%   52.36%   +24.96%     
===========================================
  Files          56       78       +22     
  Lines        6191     9631     +3440     
===========================================
+ Hits         1696     5043     +3347     
- Misses       4495     4588       +93     

Flag	Coverage Δ
unittests	52.36% <ø> (+24.96%)	⬆️

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]

@Yikun @wangxiyuan @ApsarasX @ganyi1996ppo ready to merge.

[Angazenn]

Hello, I tried your PR and run offline inference with DP > 1, and I found that line 630 would throw a "dim out of range" error. It seems that you are using all_gather to get forward metadata across all dp ranks, but local_forward_metadata is actually an 1D tensor. Therefore, global_forward_metadata is also an 1D tensor. Perhaps codes should be changed like this:

Suggested change

	local_forward_metadata = torch.tensor([num_tokens, with_prefill],
	device="npu",
	dtype=torch.int32)
	global_forward_metadata = get_dp_group().all_gather(
	local_forward_metadata)
	local_forward_metadata = torch.tensor([num_tokens, with_prefill],
	device="npu",
	dtype=torch.int32).unsqueeze(0)
	global_forward_metadata = get_dp_group().all_gather(
	local_forward_metadata, dim=0)

[Angazenn]

why is it 1?

[jianzs]

If graph mode is off, a dummy run only needs to be executed; computational requirements are not a factor.