[main][Feature] enable prefill with mc2 with additional_config (vllm-project#11022)

### What this PR does / why we need it?

This PR adds support for reserving MC2 token capacity for prefill
batches on Ascend.

A new `additional_config` option, `enable_prefill_mc2`, is introduced.
When it is enabled, `mc2_token_capacity` is calculated based on
`max_num_batched_tokens` instead of the decode-only capacity. This
allows MC2 to be used in prefill-heavy workloads where the required
token capacity can be larger than the decode-only path.

The switch is disabled by default. It is intended as a temporary option
until MC2 operators are complete for all scenarios, after which MC2 can
be enabled by default.

### Does this PR introduce any user-facing change?

Yes.

Users can now enable prefill MC2 capacity reservation through the Ascend
`additional_config` option:

`enable_prefill_mc2: true`

By default, this option is `False`. When enabled, users should ensure
that `max_num_batched_tokens` is configured appropriately. The
documentation recommends that the maximum value of
`max_num_batched_tokens` should be `tp_size * 512`.

### How was this patch tested?

This PR was tested with:

The PR also triggered the project CI checks, including DCO, docs link
checks, E2E workflow, lint-and-select-tests, and selected tests on CPU,
310P, A2, and A3 configurations.

- vLLM version: v0.23.0
- vLLM main:
vllm-project/vllm@a30addc

---------

Signed-off-by: pppeng <liuzepeng6@huawei.com>
Signed-off-by: pppeng <372907983@qq.com>
Co-authored-by: pppeng <liuzepeng6@huawei.com>

Author: ppppeng

.github/workflows/configs/nightly_config.yaml

@@ -134,6 +134,9 @@ a3:
       - name: glm-4.7-w8a8
         os: linux-aarch64-a3-16
         config_file_path: GLM-4.7.yaml
+      - name: glm-5.1-w8a8-prefill-mc2
+        os: linux-aarch64-a3-16
+        config_file_path: GLM-5.1-W8A8-PrefillMC2.yaml
       - name: kimi-k2.5
         os: linux-aarch64-a3-16
         config_file_path: Kimi-K2.5.yaml

docs/source/user_guide/configuration/additional_config.md

@@ -89,6 +89,7 @@ The following table lists additional configuration options available in vLLM Asc
 | `enable_sparse_c8`                  | bool | `False` | Whether to enable KV cache C8 in DSA models (e.g., DeepSeek V3.2 and GLM5). Not supported on Ascend 950 devices now |
 | `c8_enable_reshape_optim`           | bool | `False` | Whether to enable StoreKVBlock operator achieves acceleration under the C8 feature (this means that enable_sparse_c8 needs to be enabled). In the PD separation scenario, only the P node is enabled. |
 | `enable_mc2_hierarchy_comm`         | bool | `False` | Enable dispatch/combine op inter-node communication by ROCE. |
+| `enable_prefill_mc2`                | bool | `False` | Whether to reserve mc2_token_capacity for prefill batches. When enabled, `max_num_batched_tokens` is used to calculate the mc2_token_capacity instead of the decode-only capacity. In this scenario, the recommended maximum value of `max_num_batched_tokens` is `tp_size * 512`. This is a temporary switch; once MC2 operators are complete for all scenarios, this switch will be removed and MC2 will be enabled by default. |
 | `profiling_chunk_config`            | dict | `{}`    | Configuration options for dynamic chunked pipeline parallel. See [Dynamic Chunked Pipeline Parallel](../feature_guide/dynamic_chunk_pipeline_parallel.md) for details. |
 | `enable_balance_scheduling`         | bool | `False` | Whether to enable balance scheduling. Can also be configured via the `VLLM_ASCEND_BALANCE_SCHEDULING` environment variable during the migration period. |
 | `enable_flashcomm1`                 | bool | `False` | Whether to enable FlashComm1 optimization. Can also be configured via the `VLLM_ASCEND_ENABLE_FLASHCOMM1` environment variable during the migration period. |

tests/e2e/nightly/single_node/models/configs/GLM-5.1-W8A8-PrefillMC2.yaml

@@ -0,0 +1,82 @@
+# ==========================================
+# ACTUAL TEST CASES
+# ==========================================
+
+test_cases:
+  - name: "GLM-5.1-W8A8-PrefillMC2"
+    model: "Eco-Tech/GLM-5.1-w8a8"  #need update
+    envs:
+      OMP_PROC_BIND: "false"
+      OMP_NUM_THREADS: "1"
+      HCCL_OP_EXPANSION_MODE: "AIV"
+      PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True"
+      VLLM_USE_V1: "1"
+      HCCL_BUFFSIZE: "1800"
+      ASCEND_AGGREGATE_ENABLE: "1"
+      ASCEND_TRANSPORT_PRINT: "1"
+      ACL_OP_INIT_MODE: "1"
+      ASCEND_A3_ENABLE: "1"
+      VLLM_NIXL_ABORT_REQUEST_TIMEOUT: "300000"
+      TASK_QUEUE_ENABLE: "1"
+      VLLM_ASCEND_ENABLE_MLAPO: "1"
+      VLLM_ASCEND_ENABLE_FLASHCOMM1: "1"
+      SERVER_PORT: "DEFAULT_PORT"
+    server_cmd:
+      - "--tensor-parallel-size"
+      - "16"
+      - "--data-parallel-size"
+      - "1"
+      - "--port"
+      - "$SERVER_PORT"
+      - "--seed"
+      - "1024"
+      - "--max-model-len"
+      - "10240"
+      - "--max-num-batched-tokens"
+      - "4096"
+      - "--max-num-seqs"
+      - "32"
+      - "--trust-remote-code"
+      - "--quantization"
+      - "ascend"
+      - "--async-scheduling"
+      - "--no-enable-prefix-caching"
+      - "--enable-expert-parallel"
+      - "--gpu-memory-utilization"
+      - "0.94"
+      - "--compilation-config"
+      - '{"cudagraph_mode":"FULL_DECODE_ONLY"}'
+      - "--speculative-config"
+      - '{"num_speculative_tokens": 3, "method":"deepseek_mtp", "enforce_eager": true}'
+      - "--additional_config"
+      - '{"enable_prefill_mc2": true}'      
+      - "--reasoning-parser"
+      - "glm45"
+      - "--tool-call-parser"
+      - "glm47"
+      
+    benchmarks:
+      acc_gsm8k:
+        case_type: accuracy
+        dataset_path: vllm-ascend/gsm8k-lite
+        request_conf: vllm_api_general_chat
+        dataset_conf: gsm8k/gsm8k_gen_0_shot_cot_chat_prompt
+        max_out_len: 8192
+        batch_size: 32
+        baseline: 96.88
+        temperature: 1.0
+        top_p: 0.95
+        thinking: true
+        threshold: 5
+
+      perf_2:
+        case_type: performance
+        dataset_path: vllm-ascend/GSM8K-in3500-bs400
+        request_conf: vllm_api_stream_chat
+        dataset_conf: gsm8k/gsm8k_gen_0_shot_cot_str_perf
+        num_prompts: 64
+        max_out_len: 1500
+        batch_size: 32
+        request_rate: 0
+        baseline: 1
+        threshold: 0.97

vllm_ascend/ascend_config.py

@@ -162,6 +162,7 @@ def __init__(self, vllm_config: "VllmConfig"):
         self.enable_cpu_binding = additional_config.get("enable_cpu_binding", True)
         self.enable_sleep_mode_extra_cleanup = additional_config.get("enable_sleep_mode_extra_cleanup", False)
         self.multistream_dsv4_dsa_overlap = additional_config.get("multistream_dsv4_dsa_overlap", True)
+        self.enable_prefill_mc2 = bool(additional_config.get("enable_prefill_mc2", False))
 
         self.enable_matmul_allreduce = self._get_config_value(
             additional_config,

vllm_ascend/ascend_forward_context.py

@@ -203,14 +203,17 @@ def set_mc2_tokens_capacity(vllm_config, max_num_reqs, uniform_decode_query_len)
     global _mc2_tokens_capacity
     if _mc2_tokens_capacity is not None:
         return
-    if vllm_config.compilation_config.cudagraph_capture_sizes:
+    if get_ascend_config().enable_prefill_mc2:
+        max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
+    elif vllm_config.compilation_config.cudagraph_capture_sizes:
         max_num_tokens = vllm_config.compilation_config.max_cudagraph_capture_size
     else:
-        # NOTE: To save memory, we cap the max number of tokens to 512.
-        max_num_tokens = min(max_num_reqs * uniform_decode_query_len, 512)
+        max_num_tokens = max_num_reqs * uniform_decode_query_len
     tp_size = vllm_config.parallel_config.tensor_parallel_size
     # Use integer arithmetic for ceiling division.
     num_tokens_per_tp_rank = (max_num_tokens + tp_size - 1) // tp_size
+    # NOTE: To save memory, we cap the max number of tokens to 512.
+    num_tokens_per_tp_rank = min(num_tokens_per_tp_rank, 512)
     _mc2_tokens_capacity = num_tokens_per_tp_rank * tp_size