[Feat] Flashcomm2 use o_shared linear

[zzhx1]

What this PR does / why we need it?

It is mentioned in the flashcomm2 technical report that FC2 will introduce full redundant storage of the o_proj matrix, which will put pressure on the memory. Therefore, the technical report proposed a compromise solution using otp2, but it will introduce additional reduce-scatter communication.

We propose a shared linear feature (#2931 ) that supports distributing weights layer by layer to each card, avoiding the need for TP splitting, and can solve the memory issue.

This PR depends on #3232 and #2931

Flashcomm2 flowchart

Does this PR introduce any user-facing change?

Use environment variables

export VLLM_ASCEND_FLASHCOMM2_PARALLEL_SIZE=1
export VLLM_ASCEND_ENABLE_FLASHCOMM2_OSHARED=1

• vLLM version: v0.12.0
• vLLM main: vllm-project/vllm@ad32e3e

[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 introduces support for shared o_proj linear layers for Flashcomm2, which involves changes across configuration, distributed state management, and the attention mechanism. The core logic for shared weights is implemented in vllm_ascend/torchair/ops/shared_weight_layer.py, which has been refactored for better usability.

My review focuses on ensuring the correctness and robustness of the new feature. I've identified a few critical issues:

• Incorrect validation logic for the new flashcomm2_oproj_shared configuration that could lead to silent failures.
• A potential crash in the shared weight layer logic when handling a series with a single layer.

I have provided suggestions to fix these issues. The rest of the changes look good and the refactoring of the shared weight layer API is a nice improvement.

[gemini-code-assist]

The validation if self.flashcomm2_oproj_tensor_parallel_size is None: is incorrect. The value of self.flashcomm2_oproj_tensor_parallel_size is an integer returned from get_flashcomm2_config_and_validate (which gets it from an environment variable with a default of 0), so it will never be None. The check should be against 0, as flashcomm2_oproj_shared requires flashcomm2_oproj_tensor_parallel_size to be greater than 0.

Suggested change

	if self.flashcomm2_oproj_tensor_parallel_size is None:
	raise AssertionError(
	"flashcomm2_oproj_shared must be enabled simultaneously with flashcomm2_oproj_tensor_parallel_size"
	)
	if self.flashcomm2_oproj_tensor_parallel_size == 0:
	raise AssertionError(
	"flashcomm2_oproj_shared must be enabled with flashcomm2_oproj_tensor_parallel_size > 0"
	)

[gemini-code-assist]

The assertion self.prefetch_step <= self.num_layers - 2 will cause a crash if a shared weight series contains only one layer (self.num_layers == 1), because self.num_layers - 2 would be -1. For a single-layer series, prefetching is not applicable, and prefetch_step should be 0. To prevent this crash, the assertion should be adjusted to handle this edge case.

Suggested change

	assert self.prefetch_step >= 0 and self.prefetch_step <= self.num_layers - 2, "prefetch_step must be in [0, num_layers - 2]"
	assert self.prefetch_step >= 0 and self.prefetch_step <= max(0, self.num_layers - 2), "prefetch_step must be in [0, num_layers - 2]"

[gemini-code-assist]

This validation logic for flashcomm2_oproj_shared is redundant with the logic in vllm_ascend/ascend_config.py. It's better to have validation in one place to avoid inconsistencies. Since ascend_config.py is the configuration entry point, it's a better place for this check. Additionally, the check if flashcomm2_oproj_tp_size is None: is incorrect, as flashcomm2_oproj_tp_size is an integer. I've suggested a fix in ascend_config.py and recommend removing this redundant block.

[github-actions]

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

[zzhx1]

@wangxiyuan this PR is ready, please help merge it in.

[wangxiyuan]

please update the doc as well https://docs.vllm.ai/projects/ascend/en/latest/tutorials/DeepSeek-V3.1.html

[wangxiyuan]

Add the note to describe how to use this env

[dsxsteven]

How can adapt this to PCP?