[main] addrmsnorm + quant fusion optim in Dense Models

[rjg-lyh]

What this PR does / why we need it?

This PR fused addrmsnorm op and w8a8 quant op to get better perf.

Does this PR introduce any user-facing change?

No.

How was this patch tested?

CI passed with new added/existing test.

• vLLM version: v0.10.2
• vLLM main: vllm-project/vllm@0faf3cc

[gemini-code-assist]

Code Review

This pull request introduces a fusion optimization for addrmsnorm and w8a8 quantization on Ascend hardware. The changes add a new environment flag to control the feature and modify the forward context and layer normalization logic to implement the fusion. While the overall direction is good for performance, I've found a few critical issues in the implementation. The core fusion logic in AscendRMSNorm appears to be reversed, and it fails to increment the layer index, which would lead to incorrect behavior for models with more than one layer. Additionally, there are missing arguments in function calls that would cause runtime errors. I've provided suggestions to fix these issues.

[gemini-code-assist]

The initial value for fusion_linear should be "qkv" to correctly fuse with the pre-attention RMSNorm first. The current value "gate_up" causes the fusion logic to be reversed.

Suggested change

	forward_context.fusion_linear = "gate_up"
	forward_context.fusion_linear = "qkv"

[gemini-code-assist]

There appear to be two critical issues in this logic:

1. Incorrect Fusion Logic: The fusion logic seems to be reversed. In a standard transformer block, the first RMSNorm (pre-attention) should be fused with the qkv_proj of the attention layer, and the second RMSNorm (pre-MLP) should be fused with the gate_up_proj of the MLP. The current implementation does the opposite.

2. Missing layer_idx Increment: The layer_idx is fetched from the forward_context but is never incremented. This will cause the fusion to always use the layers from index 0 (prefetch_model.model.layers[0]) for all subsequent transformer layers, which is incorrect.

To fix this, the fusion logic should be corrected, and layer_idx should be incremented after the second RMSNorm of each layer. The initial value for fusion_linear in ascend_forward_context.py should also be adjusted.

Suggested change

	forward_context = get_forward_context()
	prefetch_model = forward_context.prefetch_model
	layer_idx = forward_context.layer_idx
	fusion_linear = forward_context.fusion_linear
	if fusion_linear == "gate_up":
	fusion_linear = prefetch_model.model.layers[layer_idx].mlp.gate_up_proj
	forward_context.fusion_linear = "qkv"
	elif fusion_linear == "qkv":
	fusion_linear = prefetch_model.model.layers[layer_idx].self_attn.qkv_proj
	forward_context.fusion_linear = "gate_up"
	forward_context = get_forward_context()
	prefetch_model = forward_context.prefetch_model
	layer_idx = forward_context.layer_idx
	fusion_target = forward_context.fusion_linear
	if fusion_target == "qkv":
	fusion_linear = prefetch_model.model.layers[layer_idx].self_attn.qkv_proj
	forward_context.fusion_linear = "gate_up"
	elif fusion_target == "gate_up":
	fusion_linear = prefetch_model.model.layers[layer_idx].mlp.gate_up_proj
	forward_context.fusion_linear = "qkv"
	forward_context.layer_idx += 1

[gemini-code-assist]

The calls to _addrmsnorm_w8a8_quant_forward_oot and _addrmsnorm_forward_oot are missing the self argument. These helper functions are defined to take self as their first parameter, but it's not being passed in the calls. This will lead to a TypeError at runtime.

Suggested change

	x, residual = _addrmsnorm_w8a8_quant_forward_oot(x, residual, fusion_linear)
	else:
	x, residual = _addrmsnorm_forward_oot(x, residual)
	else:
	x, _, residual = torch_npu.npu_add_rms_norm(
	x, residual, self.weight, self.variance_epsilon)
	x, residual = _addrmsnorm_forward_oot(x, residual)
	x, residual = _addrmsnorm_w8a8_quant_forward_oot(self, x, residual, fusion_linear)
	else:
	x, residual = _addrmsnorm_forward_oot(self, x, residual)
	else:
	x, residual = _addrmsnorm_forward_oot(self, x, residual)

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

[github-actions]

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

[Angazenn]

Can this PR handle the situation where RMSNorm has an extra bias that is introduced by anti-outlier algorithms?

[github-actions]

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