[main] flashcomm_v1 optim in Qwen Dense Models (#2802)

### What this PR does / why we need it?
Flashcomm_v1 optim in Qwen Dense Models.

### 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.1.1
- vLLM main:
5e537f45b4

Co-authored-by: 1024daniel <xxltju324@gmail.com>

tests/e2e/multicard/test_offline_inference_distributed.py

@@ -23,6 +23,7 @@ Run `pytest tests/test_offline_inference.py`.
 import os
 from unittest.mock import patch
 
+import pytest
 from modelscope import snapshot_download  # type: ignore
 from vllm import SamplingParams
 
@@ -30,6 +31,8 @@ from tests.e2e.conftest import VllmRunner
 
 os.environ["PYTORCH_NPU_ALLOC_CONF"] = "max_split_size_mb:256"
 
+QWEN_DENSE_MODELS = ["Qwen/QwQ-32B", "Qwen/Qwen-32B"]
+
 
 def test_models_distributed_QwQ():
     example_prompts = [
@@ -150,3 +153,23 @@ def test_sp_for_qwen3_moe() -> None:
                     enable_expert_parallel=True,
                     enforce_eager=True) as vllm_model:
         vllm_model.generate(example_prompts, sampling_params)
+
+
+@pytest.mark.parametrize("enforce_eager", [True, False])
+@pytest.mark.parametrize("model", QWEN_DENSE_MODELS)
+@patch.dict(os.environ, {"VLLM_ASCEND_ENABLE_DENSE_OPTIMIZE": "1"})
+@patch.dict(os.environ, {"VLLM_ASCEND_ENABLE_FLASHCOMM": "1"})
+def test_models_distributed_Qwen_Dense_with_flashcomm_v1(model, enforce_eager):
+    example_prompts = [
+        "Hello, my name is",
+    ]
+    max_tokens = 5
+
+    with VllmRunner(
+            snapshot_download(model),
+            max_model_len=8192,
+            enforce_eager=enforce_eager,
+            dtype="auto",
+            tensor_parallel_size=4,
+    ) as vllm_model:
+        vllm_model.generate_greedy(example_prompts, max_tokens)

tests/ut/ops/test_layernorm.py

@@ -10,6 +10,13 @@ def dummy_tensor():
     return torch.randn(4, 8, dtype=torch.float16)
 
 
+def mock_maybe_chunk_residual(x, residual):
+    if x.size(0) != residual.size(0):
+        return residual[:4]
+
+    return residual
+
+
 def mock_rms_norm(x, weight, eps):
     return x + 1, None
 
@@ -23,11 +30,13 @@ def mock_add_rms_norm(x, residual, weight, eps):
                          [None, torch.randn(4, 8, dtype=torch.float32)])
 @patch("torch_npu.npu_rms_norm", side_effect=mock_rms_norm)
 @patch("torch_npu.npu_add_rms_norm", side_effect=mock_add_rms_norm)
-def test_RMSNorm_forward(mock_add_rmsnorm, mock_rmsnorm, is_310p_return,
-                         residual, dummy_tensor):
+@patch("torch.ops.vllm.maybe_chunk_residual",
+       side_effect=mock_maybe_chunk_residual)
+def test_RMSNorm_forward(mock_maybe_chunk_residual, mock_add_rmsnorm,
+                         mock_rmsnorm, is_310p_return, residual, dummy_tensor):
 
     with patch("vllm_ascend.utils.is_310p", return_value=is_310p_return):
-        layer = RMSNorm(hidden_size=32, eps=1e-05)
+        layer = RMSNorm(hidden_size=8, eps=1e-05)
         if residual is not None:
             out_x, out_residual = layer.forward_oot(dummy_tensor, residual)
 
@@ -51,3 +60,25 @@ def test_RMSNorm_forward(mock_add_rmsnorm, mock_rmsnorm, is_310p_return,
 
             mock_rmsnorm.assert_called_once()
             assert torch.allclose(out_x, expected_out_x)
+
+
+@patch("vllm_ascend.utils.is_310p", return_value=False)
+@patch("torch_npu.npu_add_rms_norm", side_effect=mock_add_rms_norm)
+@patch("torch.ops.vllm.maybe_chunk_residual",
+       side_effect=mock_maybe_chunk_residual)
+def test_RMSNorm_forward_with_flashcomm_v1(mock_maybe_chunk_residual,
+                                           mock_add_rms_norm, mock_is310p):
+    x = torch.randn(4, 512, dtype=torch.bfloat16)
+    residual = torch.randn(16, 512, dtype=torch.bfloat16)
+    layer = RMSNorm(hidden_size=512, eps=1e-05)
+
+    out_x, out_residual = layer.forward_oot(x, residual)
+
+    expected_out_x = 2 * x
+    expected_out_residual = 2 * residual[:4]
+
+    mock_maybe_chunk_residual.assert_called_once()
+    mock_add_rms_norm.assert_called_once()
+    assert out_residual.size(0) == 4
+    assert torch.allclose(out_x, expected_out_x)
+    assert torch.allclose(out_residual, expected_out_residual)

tests/ut/test_utils.py

@@ -303,13 +303,13 @@ class TestUtils(TestBase):
         # ascend custom op is not registered
         utils.register_ascend_customop()
         # should call register_oot three
-        self.assertEqual(mock_customop.register_oot.call_count, 12)
+        self.assertEqual(mock_customop.register_oot.call_count, 13)
         self.assertTrue(utils._ASCEND_CUSTOMOP_IS_REIGISTERED)
 
         # ascend custom op is already registered
         utils.register_ascend_customop()
         # should not register_oot again, thus only called three in this ut
-        self.assertEqual(mock_customop.register_oot.call_count, 12)
+        self.assertEqual(mock_customop.register_oot.call_count, 13)
 
 
 class TestProfileExecuteDuration(TestBase):