133
docs/source/developer_guide/Design_Documents/ACL_Graph.md
Normal file
133
docs/source/developer_guide/Design_Documents/ACL_Graph.md
Normal file
@@ -0,0 +1,133 @@
|
||||
# ACL Graph
|
||||
|
||||
## Overview
|
||||
|
||||
ACL Graph is the Ascend realization of vLLM static graph execution. Upstream vLLM and PyTorch documents already describe the generic graph model, including `CUDAGraphMode`, runtime dispatch, batch descriptors, bucketing and padding, and the definitions of full graph and piecewise graph. This document focuses on what is specific to Ascend in `vllm-ascend`: the platform integration points, the extra constraints introduced by ACL graph capture, and the mechanisms used to keep attention parameters correct during replay.
|
||||
|
||||
On Ascend, the design goal is the same as upstream static graph execution: reduce host launch overhead for small and medium runtime shapes. The implementation boundary is different. vLLM provides the generic dispatch path, while `vllm-ascend` supplies the platform wrapper, capture-size trimming, and attention-specific update logic needed by ACL graph replay.
|
||||
|
||||
## Prerequisites and References
|
||||
|
||||
- Upstream vLLM design doc for generic graph concepts: [CUDA Graphs](https://docs.vllm.ai/en/latest/design/cuda_graphs/).
|
||||
- PyTorch graph documentation for generic capture and replay semantics: [Accelerating PyTorch with CUDA Graphs](https://pytorch.org/blog/accelerating-pytorch-with-cuda-graphs/).
|
||||
- Ascend user guide for operational enablement: [Graph Mode Guide](https://docs.vllm.ai/projects/ascend/en/latest/user_guide/feature_guide/graph_mode.html).
|
||||
- Existing repo design note: [ACL Graph](https://docs.vllm.ai/projects/ascend/zh-cn/latest/developer_guide/Design_Documents/ACL_Graph.html)
|
||||
|
||||
This document intentionally does not re-explain upstream topics such as graph mode selection, dispatcher behavior, batch descriptor construction, capture bucketing, padding policy, or the generic meaning of full versus piecewise execution.
|
||||
|
||||
## How ACL Graph Fits into vLLM
|
||||
|
||||
vLLM owns the generic static graph flow. On Ascend, `NPUPlatform.get_static_graph_wrapper_cls()` returns `vllm_ascend.compilation.acl_graph.ACLGraphWrapper`, which is the platform-specific wrapper used when vLLM enables static graph mode.
|
||||
|
||||
`ACLGraphWrapper` is responsible for:
|
||||
|
||||
- reading the runtime mode and `batch_descriptor` from the forward context,
|
||||
- deciding whether to run eagerly, capture a new ACL graph, or replay a cached ACL graph,
|
||||
- caching graph entries per batch descriptor,
|
||||
- preserving the graph pool and replay bookkeeping needed by the Ascend backend.
|
||||
|
||||
The wrapper does not define the upstream dispatch policy. It assumes the runtime mode and batch descriptor have already been chosen correctly by vLLM, then applies Ascend capture or replay to that concrete runtime shape.
|
||||
|
||||
## Capture Sizes and Bucketing
|
||||
|
||||
vLLM graph replay requires stable runtime shapes, so vLLM does not try to capture every possible batch shape. Instead, it prepares a finite set of capture sizes and dispatches a runtime batch to the nearest supported size. If the runtime batch is larger than the largest configured capture size, graph mode is skipped and execution falls back to eager mode.
|
||||
|
||||
By default, vLLM builds capture sizes as:
|
||||
|
||||
- `1`, `2`, `4`
|
||||
- multiples of `8` from `8` up to `255`
|
||||
- multiples of `16` from `256` up to `max_cudagraph_capture_size`
|
||||
|
||||
Conceptually, the default list looks like:
|
||||
|
||||
```text
|
||||
[1, 2, 4, 8, 16, 24, 32, ..., 248, 256, 272, 288, ...]
|
||||
```
|
||||
|
||||
The smaller step at small batch sizes reduces padding overhead where latency is most sensitive, while the larger step at bigger sizes keeps the number of captured graphs under control.
|
||||
|
||||
On Ascend, this generic upstream bucketing strategy is still the starting point, but the final capture sizes may be reduced further by platform-specific constraints:
|
||||
|
||||
- sequence-parallel filtering may remove unsupported sizes,
|
||||
- runtime resource limits may still prevent some configured sizes from being captured,
|
||||
- some runtime modes may be normalized before capture begins.
|
||||
|
||||
## Ascend-Specific Design Constraints
|
||||
|
||||
### Capture breadth is still constrained by runtime resources
|
||||
|
||||
Unlike CUDA Graph on CUDA devices, ACL graph capture on Ascend can still fail when the selected graph sizes consume more runtime resources than the current backend can supply. Piecewise mode is the most sensitive case because it captures many subgraphs and the total capture cost scales with model depth and configured size coverage.
|
||||
|
||||
Older versions of vLLM Ascend applied a local `update_aclgraph_sizes()` heuristic to shrink the PIECEWISE capture-size set before final capture. That heuristic has been removed. The current implementation keeps upstream sizing and dispatch behavior intact, then intercepts the confirmed capture-time stream-resource signature in `vllm_ascend/compilation/acl_graph.py` and re-raises it with clearer mitigation guidance.
|
||||
|
||||
In practice, this means users should treat `cudagraph_capture_sizes` and `max_cudagraph_capture_size` as the primary tuning levers when capture fails. Newer HDK/CANN combinations can materially improve ACL graph capacity, while communication-heavy configurations may still require a smaller configured size set.
|
||||
|
||||
### Platform mode normalization is stricter than generic upstream behavior
|
||||
|
||||
Ascend currently narrows some generic upstream modes in `vllm_ascend.platform.NPUPlatform.check_and_update_config()`.
|
||||
|
||||
- Encoder-decoder models are forced to `PIECEWISE`.
|
||||
- `use_inductor` is disabled for ACL graph paths.
|
||||
- `ASCEND_LAUNCH_BLOCKING=1` is rejected when ACL graph is enabled.
|
||||
- Xlite graph mode can disable ACL graph full mode or fall back to `FULL_DECODE_ONLY`, depending on configuration.
|
||||
|
||||
These checks document the subset of upstream graph behavior that the current Ascend backend can execute safely. Some of them are long-term platform constraints, while others are clearly transitional in the current implementation.
|
||||
|
||||
## Key Ascend-Specific Mechanisms
|
||||
|
||||
### Host-side attention parameter update for full graph replay
|
||||
|
||||
Full graph replay on Ascend has an extra problem that upstream generic documentation does not cover in detail: some attention operators need runtime metadata updates even when the overall graph is static. The Ascend implementation handles this by separating graph capture from host-side task parameter updates.
|
||||
|
||||
The flow is:
|
||||
|
||||
1. During capture, attention backends record per-graph task handles, events, workspaces, and weak references to the tensors or metadata that must be refreshed.
|
||||
2. Before replay, `update_full_graph_params()` calls the backend specific `update_graph_params()` implementation.
|
||||
3. That backend runs parameter refresh on an update stream with `torch.npu.graph_task_update_begin(...)` and `torch.npu.graph_task_update_end(...)` around the underlying attention operator launch.
|
||||
4. `torch.npu.ExternalEvent` objects are used to enforce ordering between the host-side update stream and the replay stream.
|
||||
|
||||
This mechanism is implemented in attention backends such as:
|
||||
|
||||
- `vllm_ascend/attention/attention_v1.py`
|
||||
- `vllm_ascend/attention/mla_v1.py`
|
||||
- `vllm_ascend/attention/context_parallel/attention_cp.py`
|
||||
- `vllm_ascend/attention/context_parallel/mla_cp.py`
|
||||
|
||||
The important design point is that Ascend full graph support depends on backend-provided `update_graph_params()` hooks. Without that hook, capture alone is not enough to replay the correct attention state.
|
||||
|
||||
### Replay ordering and synchronization
|
||||
|
||||
`ACLGraphWrapper` synchronizes the current stream before replay in the common path to ensure that host-side parameter updates stay aligned with the graph execution that will consume them. This is especially relevant in asynchronous scheduling or multi-threaded execution.
|
||||
|
||||
If ordering is not preserved, the parameter update for iteration *i* can be observed by the replay of iteration *i-1*, or the replay of iteration *i* can start before its own parameter update has completed. In practice, this means the attention operator may run with mismatched runtime metadata, which can cause incorrect results, precision issues, or even hangs. The code keeps a narrower path for the main full-graph eagle case, but the general design assumption is the same: replay must not overtake pending parameter update work.
|
||||
|
||||
## Full vs Piecewise on Ascend
|
||||
|
||||
Upstream docs already define full graph and piecewise graph semantically. On Ascend, the practical difference is driven by backend support and resource cost.
|
||||
|
||||
### Piecewise mode
|
||||
|
||||
Piecewise mode is the conservative path. It relies on the generic vLLM split execution strategy, then applies ACL graph capture to the non-attention segments selected by the compilation path. On Ascend, this mode is currently the more widely supported option, but it is also the most sensitive to stream pressure because the number of captured graphs scales with model depth.
|
||||
|
||||
### Full graph mode
|
||||
|
||||
Full graph mode is the more performance-oriented path when the attention backend can support runtime parameter patching through `update_graph_params()`. On Ascend, full graph support is tied to those attention-specific update hooks, workspace caching, and replay ordering guarantees.
|
||||
|
||||
## Diagnostics and Operational Notes
|
||||
|
||||
- The simplest way to confirm that graph mode is active is to enable cudagraph metrics and keep log stats enabled. In CLI usage, use `--cudagraph-metrics` and do not pass `--disable-log-stats`. In Python usage, set `cudagraph_metrics=True` and `disable_log_stats=False`. Then inspect the emitted metrics and logs.
|
||||
- Profiling can also confirm whether replay is happening, and developers can add temporary prints before replay when debugging locally, but those are secondary methods and are not expanded here.
|
||||
- Capture-size selection primarily follows upstream configuration and dispatch behavior; only the confirmed stream-resource capture failure is rewritten with user-facing guidance at runtime.
|
||||
- In debug mode, `ACLGraphWrapper` asserts that replay uses the same tensor addresses recorded during capture.
|
||||
- `ASCEND_LAUNCH_BLOCKING=1` is incompatible with ACL graph enablement in the current implementation.
|
||||
- For debugging inside graph execution, the repo also provides graph-aware print helpers in `vllm_ascend.utils`, but those are developer diagnostics rather than part of the execution design.
|
||||
|
||||
## Related Files
|
||||
|
||||
- `vllm_ascend/platform.py`, mode normalization, platform hooks, and static graph wrapper selection.
|
||||
- `vllm_ascend/compilation/acl_graph.py`, ACL graph wrapper, capture and replay cache, graph parameter containers, and full graph update dispatch.
|
||||
- `vllm_ascend/compilation/acl_graph.py`, runtime ACL graph capture, replay, and capture-failure guidance.
|
||||
- `vllm_ascend/attention/attention_v1.py`, full graph attention parameter capture and update logic.
|
||||
- `vllm_ascend/attention/mla_v1.py`, MLA (Multi-Head Latent Attention) specific full graph parameter capture and update logic.
|
||||
- `vllm_ascend/attention/context_parallel/attention_cp.py`, context parallel attention update path.
|
||||
- `vllm_ascend/attention/context_parallel/mla_cp.py`, context parallel MLA update path.
|
||||
@@ -0,0 +1,91 @@
|
||||
# KV Cache Pool
|
||||
|
||||
## Why KV Cache Pool?
|
||||
|
||||
Prefix caching is an important feature in LLM inference that can reduce prefill computation time drastically.
|
||||
|
||||
However, the performance gain from prefix caching is highly dependent on the cache hit rate, while the cache hit rate can be limited if one only uses on-chip memory for KV cache storage.
|
||||
|
||||
Hence, KV Cache Pool is proposed to utilize various types of storage including on-chip memory, DRAM, and SSD, making a pool for KV Cache storage while making the prefix of requests visible across all nodes, increasing the cache hit rate for all requests.
|
||||
|
||||
vLLM Ascend currently supports [MooncakeStore](https://github.com/kvcache-ai/Mooncake), one of the most recognized KV Cache storage engines.
|
||||
|
||||
While one can utilize MooncakeStore in vLLM V1 engine by setting it as a remote backend of LMCache with GPU (see [Tutorial](https://github.com/LMCache/LMCache/blob/dev/examples/kv_cache_reuse/remote_backends/mooncakestore/README.md)), we find it would be better to integrate a connector that directly supports MooncakeStore and can utilize the data transfer strategy that best fits Huawei NPU hardware.
|
||||
|
||||
Hence, we propose to integrate MooncakeStore with a brand new **MooncakeStoreConnectorV1**, which is indeed largely inspired by **LMCacheConnectorV1** (see the [How is MooncakeStoreConnectorV1 Implemented?](#how-is-mooncakestoreconnectorv1-implemented) section).
|
||||
|
||||
## Usage
|
||||
|
||||
vLLM Ascend currently supports MooncakeStore for KV Cache Pool. To enable MooncakeStore, one needs to configure `kv-transfer-config` and choose `MooncakeStoreConnector` as the KV Connector.
|
||||
|
||||
For step-by-step deployment and configuration, please refer to the [KV Pool User Guide](https://docs.vllm.ai/projects/ascend/en/latest/user_guide/feature_guide/kv_pool.html).
|
||||
|
||||
## How it works?
|
||||
|
||||
The KV Cache Pool integrates multiple memory tiers (on-chip memory, DRAM, SSD, etc.) through a connector-based architecture.
|
||||
|
||||
Each connector implements a unified interface for storing, retrieving, and transferring KV blocks between tiers, depending on access frequency and hardware bandwidth.
|
||||
|
||||
When combined with vLLM's Prefix Caching mechanism, the pool enables efficient caching both locally (in on-chip memory) and globally (via Mooncake), ensuring that frequently used prefixes remain hot while less frequently accessed KV data can spill over to lower-cost memory.
|
||||
|
||||
### 1. Combining KV Cache Pool with on-chip memory Prefix Caching
|
||||
|
||||
Prefix Caching with on-chip memory is already supported by the vLLM V1 Engine.
|
||||
By introducing KV Connector V1, users can seamlessly combine on-chip memory-based Prefix Caching with Mooncake-backed KV Pool.
|
||||
|
||||
The user can enable both features simply by enabling Prefix Caching, which is enabled by default in vLLM V1 unless the `--no-enable-prefix-caching` flag is set, and setting up the KV Connector for KV Pool (e.g., the MooncakeStoreConnector).
|
||||
|
||||
**Workflow**:
|
||||
|
||||
1. The engine first checks for prefix hits in the on-chip memory cache.
|
||||
|
||||
2. After getting the number of hit tokens on on-chip memory, it queries the KV Pool via the connector. If there are additional hits in the KV Pool, we get the **additional blocks only** from the KV Pool, and get the rest of the blocks directly from on-chip memory to minimize the data transfer latency.
|
||||
|
||||
3. After the KV Caches in the KV Pool are loaded into on-chip memory, the remaining process is the same as Prefix Caching in on-chip memory.
|
||||
|
||||
### 2. Combining KV Cache Pool with Mooncake PD Disaggregation
|
||||
|
||||
When used together with Mooncake PD (Prefill-Decode) Disaggregation, the KV Cache Pool can further decouple prefill and decode stages across devices or nodes.
|
||||
|
||||
Currently, we only perform put and get operations of KV Pool for **Prefill Nodes**, and Decode Nodes get their KV Cache from Mooncake P2P KV Connector, i.e., MooncakeConnector.
|
||||
|
||||
The key benefit of doing this is that we can keep the gain in performance by computing less with Prefix Caching from on-chip memory and KV Pool for Prefill Nodes, while not sacrificing the data transfer efficiency between Prefill and Decode nodes with P2P KV Connector that transfers KV Caches between NPU devices directly.
|
||||
|
||||
To enable this feature, we need to set up both Mooncake Connector and MooncakeStore Connector with a Multi Connector, which is a KV Connector class provided by vLLM that can call multiple KV Connectors in a specific order.
|
||||
|
||||
For details, please also refer to the [Mooncake connector deployment guide](https://github.com/vllm-project/vllm-ascend/blob/main/examples/disaggregated_prefill_v1/mooncake_connector_deployment_guide.md).
|
||||
|
||||
## How is MooncakeStoreConnectorV1 Implemented?
|
||||
|
||||
**MooncakeStoreConnectorV1** inherits the KV Connector V1 class in vLLM V1: through implementing the required methods defined in the KV connector V1 base class, one can integrate a third-party KV cache transfer/storage backend into the vLLM framework.
|
||||
|
||||
MooncakeStoreConnectorV1 is also largely inspired by LMCacheConnectorV1 in terms of the `Lookup Engine`/`Lookup Client` design for looking up KV cache keys, and the `ChunkedTokenDatabase` class for processing tokens into prefix-aware hashes as well as other hashing related designs. On top of this, we have also added our own design including `KVTransferThread` that allows async `get` and `put` of KV caches with multi-threading, and NPU-related data transfer optimization such as removing the `LocalBuffer` in LMCache to remove redundant data transfer.
|
||||
|
||||
The KV Connector methods that need to be implemented can be categorized into scheduler-side methods that are called in V1 scheduler and worker-side methods that are called in V1 worker, namely:
|
||||
|
||||
### KV Connector Scheduler-Side Methods
|
||||
|
||||
`get_num_new_matched_tokens`: Get prefix cache hit in number of tokens through looking up into the KV pool.
|
||||
`update_states_after_alloc`: Update KVConnector state after temporary buffer alloc.
|
||||
`build_connector_meta`: Attach the connector metadata to the request object.
|
||||
`request_finished`: Once a request is finished, determine whether request blocks should be freed now or will be sent asynchronously and freed later.
|
||||
|
||||
### Connector Worker-Side Methods
|
||||
|
||||
`register_kv_caches`: Register KV cache buffers needed for KV cache transfer.
|
||||
`start_load_kv`: Perform KV cache load operation that transfers KV cache from storage to device.
|
||||
`wait_for_layer_load`: Optional; Wait for layer load in layerwise + async KV load scenario.
|
||||
`save_kv_layer`: Optional; Do layerwise KV cache put into KV Pool.
|
||||
`wait_for_save`: Wait for KV Save to finish if async KV cache save/put.
|
||||
`get_finished`: Get request that finished KV transfer, `done_sending` if `put` finished, `done_receiving` if `get` finished.
|
||||
|
||||
## DFX
|
||||
|
||||
1. When looking up a key in KV Pool, if we cannot find the key, there is no Cache Hit for this specific block; we return no hit for this block and do not look up further blocks for the current request.
|
||||
2. Similarly, when we are trying to put a block into KV Pool and it fails, we do not put further blocks (subject to change).
|
||||
|
||||
## Limitations
|
||||
|
||||
1. Currently, MooncakeStore for vLLM Ascend only supports DRAM as the storage for KV Cache Pool.
|
||||
|
||||
2. For now, if we successfully looked up a key and found it exists, but failed to get it when calling KV Pool's get function, we just output a log indicating the get operation failed and keep going; hence, the accuracy of that specific request may be affected. We will handle this situation by falling back the request and re-compute everything assuming there's no prefix cache hit (or even better, revert only one block and keep using the Prefix Caches before that).
|
||||
@@ -0,0 +1,286 @@
|
||||
# Prepare inputs for model forwarding
|
||||
|
||||
## Purpose
|
||||
|
||||
Information required to perform model forward pass:
|
||||
|
||||
- the inputs
|
||||
- the corresponding attention metadata of the inputs
|
||||
|
||||
The following diagram shows what we should prepare for model inference.
|
||||
|
||||
```text
|
||||
+---------------+
|
||||
inputs --> | |
|
||||
| model | --> output
|
||||
attn_meta --> | |
|
||||
+---------------+
|
||||
```
|
||||
|
||||
Therefore, as long as we have these two pieces of information mentioned above, we can perform the model's forward propagation.
|
||||
|
||||
This document will explain **how we obtain the inputs and their corresponding attention metadata**.
|
||||
|
||||
## Overview
|
||||
|
||||
### 1. Obtain inputs
|
||||
|
||||
The workflow of obtaining inputs:
|
||||
|
||||
1. Get `token positions`: relative position of each token within its request sequence.
|
||||
|
||||
2. Get `token indices`: index of each scheduled token in the token table.
|
||||
|
||||
3. Get `Token IDs`: using token indices to retrieve the Token IDs from **token id table**.
|
||||
|
||||
At last, these `Token IDs` are required to be fed into a model, and `positions` should also be sent into the model to create `RoPE` (Rotary positional embedding). Both of them are the inputs of the model.
|
||||
|
||||
**Note**: The `Token IDs` are the inputs of a model, so we also call them `Input IDs`.
|
||||
|
||||
### 2. Build inputs attention metadata
|
||||
|
||||
A model requires these attention metadata during the forward pass:
|
||||
|
||||
- `query start location`: start and end location of each request corresponding to the scheduled tokens.
|
||||
- `sequence length`: length of each request including both computed tokens and newly scheduled tokens.
|
||||
- `number of computed tokens`: number of computed tokens for each request.
|
||||
- `number of requests`: number of requests in this batch.
|
||||
- `number of tokens`: total number of scheduled tokens in this batch.
|
||||
- **`block table`**: translates the logical address (within its sequence) of each block to its global physical address in the device's memory.
|
||||
- `max query len`: the longest scheduled tokens length in this request batch.
|
||||
- `slot mapping`: indices of each token that input token will be stored into.
|
||||
- `attention mask`: mask matrix applied to attention scores before softmax to control which tokens can attend to each other (usually a causal attention).
|
||||
|
||||
## Before start
|
||||
|
||||
There are mainly three types of variables.
|
||||
|
||||
- token level: represents one attribute corresponding to each scheduled token, so the length of this variable is the number of scheduled tokens.
|
||||
- request level: represents one attribute of each scheduled request, whose length usually is the number of scheduled requests. (`query start location` is a special case, which has one more element.)
|
||||
- system level:
|
||||
1. **Token IDs table**: stores the token IDs (i.e. the inputs of a model) of each request. The shape of this table is `(max num request, max model len)`. Here, `max num request` is the maximum count of concurrent requests allowed in a forward batch and `max model len` is the maximum token count that can be handled at one request sequence in this model.
|
||||
2. **Block table**: translates the logical address (within its sequence) of each block to its global physical address in the device's memory. The shape of this table is `(max num request, max model len / block size)`
|
||||
|
||||
**Note**: Both of these two tables come from the `_update_states` method before **preparing inputs**. You can take a look if you need more inspiration.
|
||||
|
||||
### Tips
|
||||
|
||||
Simply put, a `token ID` is an **integer** (usually `int32`), which represents a token.
|
||||
Example of `Token ID`:
|
||||
|
||||
```shell
|
||||
| Token ID | Token |
|
||||
|--------------|---------------|
|
||||
| 0 | [PAD] |
|
||||
| 1 | <|endoftext|> |
|
||||
| 2 | <|start|> |
|
||||
| 3 | [SEP] |
|
||||
| 4 | I |
|
||||
| 5 | the |
|
||||
| 6 | be |
|
||||
| 7 | of |
|
||||
| 8 | and |
|
||||
| ... | ... |
|
||||
| ... | ... |
|
||||
| vocab_size-1 | <|im_end|> |
|
||||
```
|
||||
|
||||
## Go through details
|
||||
|
||||
Assumptions:
|
||||
|
||||
- maximum number of tokens that can be scheduled at once: 10
|
||||
- `block size`: 2
|
||||
- Totally schedule 3 requests. Their prompt lengths are 3, 2, and 8 respectively.
|
||||
- `max model length`: 12 (the maximum token count that can be handled at one request sequence in a model).
|
||||
|
||||
These assumptions are configured at the beginning when starting vLLM. They are not fixed, so you can manually set them.
|
||||
|
||||
### Step 1: All requests in the prefill phase
|
||||
|
||||
#### Obtain inputs
|
||||
|
||||
As the maximum number of tokens that can be scheduled is 10, the scheduled tokens of each request can be represented as `{'0': 3, '1': 2, '2': 5}`. Note that `request_2` uses chunked prefill, leaving 3 prompt tokens unscheduled.
|
||||
|
||||
##### 1. Get token positions
|
||||
|
||||
First, determine which request each token belongs to: tokens 0–2 are assigned to **request_0**, tokens 3–4 to **request_1**, and tokens 5–9 to **request_2**. To represent this mapping, we use `request indices`, for example, `request indices`: `[0, 0, 0, 1, 1, 2, 2, 2, 2, 2]`.
|
||||
|
||||
For each request, use **the number of computed tokens** + **the relative position of current scheduled tokens** (`request_0: [0 + 0, 0 + 1, 0 + 2]`, `request_1: [0 + 0, 0 + 1]`, `request_2: [0 + 0, 0 + 1,..., 0 + 4]`) and then concatenate them together (`[0, 1, 2, 0, 1, 0, 1, 2, 3, 4]`).
|
||||
|
||||
Note: there is a more efficient way (using `request indices`) to create positions in actual code.
|
||||
|
||||
Finally, `token positions` can be obtained as `[0, 1, 2, 0, 1, 0, 1, 2, 3, 4]`. This variable is **token level**.
|
||||
|
||||
##### 2. Get token indices
|
||||
|
||||
The shape of the current **Token IDs table** is `(max num request, max model len)`.
|
||||
|
||||
Why are these `T_3_5`, `T_3_6`, `T_3_7` in this table without being scheduled?
|
||||
|
||||
- We fill all Token IDs in one request sequence to this table at once, but we only retrieve the tokens we scheduled this time. Then we retrieve the remaining Token IDs next time.
|
||||
|
||||
```shell
|
||||
| T_0_0 | T_0_1 | T_0_2 | ? | ? | ? | ? | ? | ? | ? | ? | ? |
|
||||
| T_1_0 | T_1_1 | ? | ? | ? | ? | ? | ? | ? | ? | ? | ? |
|
||||
| T_2_0 | T_2_1 | T_3_2 | T_3_3 | T_3_4 | T_3_5 | T_3_6 | T_3_7 | ? | ? | ? | ? |
|
||||
| ? | ? | ? | ? | ? | ? | ? | ? | ? | ? | ? | ? |
|
||||
......
|
||||
......
|
||||
......
|
||||
```
|
||||
|
||||
Note that `T_x_x` is an `int32`.
|
||||
|
||||
Let's say `M = max model len`. Then we can use `token positions` together with `request indices` of each token to construct `token indices`.
|
||||
|
||||
So `token indices` = `[0 + 0 * M, 1 + 0 * M, 2 + 0 * M, 0 + 1 * M, 1 + 1 * M, 0 + 2 * M, 1 + 2 * M, 2 + 2 * M, 3 + 2 * M, 4 + 2 * M]` = `[0, 1, 2, 12, 13, 24, 25, 26, 27, 28]`
|
||||
|
||||
##### 3. Retrieve the Token IDs
|
||||
|
||||
We use `token indices` to select out the corresponding `Input IDs` from the token table. The pseudocode is as follows:
|
||||
|
||||
```shell
|
||||
input_ids = token_table[token_indices]
|
||||
```
|
||||
|
||||
As mentioned before, we refer to these `Token IDs` as `Input IDs`.
|
||||
|
||||
- `Input IDs` = `[T_0_0, T_0_1, T_0_2, T_1_0, T_1_1, T_2_0, T_2_1, T_3_2, T_3_3, T_3_4]`
|
||||
|
||||
#### Build inputs attention metadata
|
||||
|
||||
In the current **Block Table**, we use the first block (i.e. block_0) to mark the unused block. The shape of the block is `(max num request, max model len / block size)`, where `max model len / block size = 12 / 2 = 6`.
|
||||
|
||||
```shell
|
||||
| 1 | 2 | 0 | 0 | 0 | 0 |
|
||||
| 3 | 0 | 0 | 0 | 0 | 0 |
|
||||
| 4 | 5 | 6 | 0 | 0 | 0 |
|
||||
| 0 | 0 | 0 | 0 | 0 | 0 |
|
||||
......
|
||||
......
|
||||
......
|
||||
```
|
||||
|
||||
The KV cache block in the device memory is like:
|
||||
|
||||
```shell
|
||||
| 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | ......
|
||||
```
|
||||
|
||||
Let's say `K = max model len / block size = 6`, and we can get token `device block number`.
|
||||
|
||||
The workflow of achieving slot mapping:
|
||||
|
||||
1. Get `block table indices` using `K`, `positions` and `request indices`.
|
||||
|
||||
Purpose: For each token, it could be used to select `device block number` from `block table`.
|
||||
|
||||
2. Get `device block number` using `block table indices`.
|
||||
|
||||
Purpose: `device block number` indicates which device block each token belongs to.
|
||||
|
||||
3. Get `block offsets` using `positions` and `block size`.
|
||||
|
||||
Purpose: `block offsets` indicates the offsets of each token within a block.
|
||||
|
||||
4. construct `slot mapping` using `device block number` and `block offsets`.
|
||||
|
||||
Purpose: we can use `slot mapping` to store Token IDs into token slots.
|
||||
|
||||
Details:
|
||||
|
||||
1. (**Token level**) Use a simple formula to calculate `block table indices`: `request indices * K + positions / block size`. So it equals `[0 * 6 + 0 / 2, 0 * 6 + 1 / 2, 0 * 6 + 2 / 2, 1 * 6 + 0 / 2, 1 * 6 + 1 / 2, 2 * 6 + 0 / 2, 2 * 6 + 1 / 2, 2 * 6 + 2 / 2, 2 * 6 + 3 / 2, 2 * 6 + 4 / 2] = [0, 0, 1, 6, 6, 12, 12, 13, 13, 14]`. This could be used to select `device block number` from `block table`.
|
||||
2. (**Token level**) Use `block table indices` to select out `device block number` for each scheduled token. The pseudocode is `block_numbers = block_table[block_table_indices]`. So `device block number=[1, 1, 2, 3, 3, 4, 4, 5, 5, 6]`
|
||||
3. (**Token level**) `block offsets` could be computed by `block offsets = positions % block size = [0, 1, 0, 0, 1, 0, 1, 0, 1, 0]`.
|
||||
4. Finally, use `block offsets` and `device block number` to create `slot mapping`: `device block number * block size + block_offsets = [2, 3, 4, 6, 7, 8, 9, 10, 11, 12]`
|
||||
|
||||
(**Request level**) As we know the scheduled token count is `[3, 2, 5]`:
|
||||
|
||||
- (**Request level**) Use prefix sum to calculate `query start location`: `[0, 3, 5, 10]`.
|
||||
- (**Request level**) All tokens in step 1 are in the prefill stage, and the computed tokens count is 0; then `sequence length` = `[3, 2, 5]`.
|
||||
- (**Request level**) As mentioned above, `number of computed tokens` are all 0s: `[0, 0, 0]`.
|
||||
- `number of requests`: `3`
|
||||
- (**Request level**) `number of tokens`: `[3, 2, 5]`
|
||||
- `max query len`: `5`
|
||||
- (**Token level**) `slot mapping`: `[2, 3, 4, 6, 7, 8, 9, 10, 11, 12]`
|
||||
- `attention mask`: For all requests that initiate a prefill process, we simply create only one mask matrix for reuse across different requests. The shape of this mask matrix is `5 * 5`:
|
||||
|
||||
### Step 2: Chunked prefill
|
||||
|
||||
In Step 2, we no longer provide explanations or perform calculations; instead, we directly present the final result.
|
||||
|
||||
#### Obtain inputs
|
||||
|
||||
Scheduled token of each request: `{'0': 1, '1': 1, '2': 3}`
|
||||
|
||||
1. `request indices`: `[0, 1, 2, 2, 2]`
|
||||
2. `token positions`: `[3, 2, 5, 6, 7]`
|
||||
|
||||
Current **Token IDs table**:
|
||||
|
||||
```shell
|
||||
| T_0_0 | T_0_1 | T_0_2 | T_0_3 | ? | ? | ? | ? | ? | ? | ? | ? |
|
||||
| T_1_0 | T_1_1 | T_1_2 | ? | ? | ? | ? | ? | ? | ? | ? | ? |
|
||||
| T_2_0 | T_2_1 | T_3_2 | T_3_3 | T_3_4 | T_3_5 | T_3_6 | T_3_7 | ? | ? | ? | ? |
|
||||
| ? | ? | ? | ? | ? | ? | ? | ? | ? | ? | ? | ? |
|
||||
......
|
||||
......
|
||||
......
|
||||
```
|
||||
|
||||
**Note**: **T_0_3**, **T_1_2** are new Token IDs of **request_0** and **request_1** respectively. They are sampled from the output of the model.
|
||||
|
||||
3. `token indices`: `[3, 14, 29, 30, 31]`
|
||||
4. `Input IDs`: `[T_0_3, T_1_2, T_3_5, T_3_6, T_3_7]`
|
||||
|
||||
#### Build inputs attention metadata
|
||||
|
||||
We allocate the blocks `7` and `8` to `request_1` and `request_2` respectively, as they need more space in device to store KV cache following token generation or chunked prefill.
|
||||
|
||||
Current **Block Table**:
|
||||
|
||||
```shell
|
||||
| 1 | 2 | 0 | 0 | 0 | 0 |
|
||||
| 3 | 7 | 0 | 0 | 0 | 0 |
|
||||
| 4 | 5 | 6 | 8 | 0 | 0 |
|
||||
| 0 | 0 | 0 | 0 | 0 | 0 |
|
||||
......
|
||||
......
|
||||
......
|
||||
```
|
||||
|
||||
KV cache block in the device memory:
|
||||
|
||||
```shell
|
||||
| 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | ......
|
||||
```
|
||||
|
||||
1. (**Token level**) `block table indices`: `[1, 7, 14, 15, 15]`
|
||||
2. (**Token level**) `device block number`: `[2, 7, 6, 8, 8]`
|
||||
3. (**Token level**) `block offsets`: `[1, 0, 1, 0, 1]`
|
||||
4. (**Token level**) `slot mapping`: `[5, 14, 13, 16, 17]`
|
||||
|
||||
Scheduled token count: `[1, 1, 3]`
|
||||
|
||||
- `query start location`: `[0, 1, 2, 5]`
|
||||
|
||||
- `sequence length`: `[4, 3, 8]`
|
||||
|
||||
- `number of computed tokens`: `[3, 2, 5]`
|
||||
|
||||
- `number of requests`: `3`
|
||||
|
||||
- `max query len`: `3`
|
||||
|
||||
- `slot mapping`: `[5, 14, 13, 16, 17]`
|
||||
|
||||
- `attention mask`: `5 * 8`
|
||||
|
||||
Each token has a `1 * 8` vector, and there are 5 scheduled tokens.
|
||||
|
||||
## At last
|
||||
|
||||
If you understand step 1 and step 2, you will know all the following steps.
|
||||
|
||||
Hope this document helps you better understand how vLLM prepares inputs for model forwarding. If you have any good ideas, you are welcome to contribute to us.
|
||||
@@ -0,0 +1,25 @@
|
||||
# Adding a custom aclnn operation
|
||||
|
||||
This document describes how to add a custom aclnn operation to vllm-ascend.
|
||||
|
||||
## How custom aclnn operation works in vllm-ascend?
|
||||
|
||||
Custom aclnn operations are built and installed into `vllm_ascend/cann_ops_custom` directory during the build process of vllm-ascend. Then the aclnn operators are bound to `torch.ops._C_ascend` module, enabling users to invoke them in vllm-ascend python code.
|
||||
|
||||
To enable custom operations, use the following code:
|
||||
|
||||
```python
|
||||
from vllm_ascend.utils import enable_custom_op
|
||||
|
||||
enable_custom_op()
|
||||
```
|
||||
|
||||
## How to add a custom aclnn operation?
|
||||
|
||||
1. Create a new operation folder under `csrc` directory.
|
||||
2. Create `op_host` and `op_kernel` directories for host and kernel source code.
|
||||
3. Add build options in `csrc/build_aclnn.sh` for supported SOC. Note that multiple ops should be separated with `;`, i.e. `CUSTOM_OPS="op1;op2;op3"`.
|
||||
4. Bind aclnn operators to torch.ops._C_ascend module in `csrc/torch_binding.cpp`.
|
||||
5. Write a meta implementation in `csrc/torch_binding_meta.cpp` for the op to be captured into the aclgraph.
|
||||
|
||||
After a successful build of vllm-ascend, the custom aclnn operation can be invoked in python code.
|
||||
197
docs/source/developer_guide/Design_Documents/context_parallel.md
Normal file
197
docs/source/developer_guide/Design_Documents/context_parallel.md
Normal file
@@ -0,0 +1,197 @@
|
||||
# Context Parallel (CP)
|
||||
|
||||
TL;DR: PCP accelerates prefill via sequence splitting. DCP eliminates KV cache redundancy.
|
||||
|
||||
|
||||
|
||||
For the main discussions during the development process, please refer to the [RFC](https://github.com/vllm-project/vllm/issues/25749) and the relevant links referenced by or referencing this RFC.
|
||||
|
||||
## What is CP?
|
||||
|
||||
**Context Parallel (CP)** is a strategy for parallelizing computation along the sequence dimension across multiple devices.
|
||||
|
||||
**Prefill Context Parallel (PCP)** expands the world size of devices and uses dedicated communication domains.
|
||||
Its primary goal is to partition the sequence dimension during the prefill phase, enabling different devices to compute distinct chunks of the sequence simultaneously.
|
||||
The KV cache is sharded along the sequence dimension across devices.
|
||||
This approach impacts the computational logic of both the Prefill and Decode stages to varying degrees.
|
||||
|
||||
**Decode Context Parallel (DCP)** reuses the communication domain of Tensor Parallelism (TP) and does not require additional devices.
|
||||
Its main objective is to eliminate duplicated storage of the KV cache by sharding it along the sequence dimension across devices within the TP domain that would otherwise hold redundant copies.
|
||||
DCP primarily influences the Decode logic, as well as the logic for chunked prefill and cached prefill.
|
||||
|
||||
## How to Use CP?
|
||||
|
||||
Please refer to the [context parallel user guide](../../user_guide/feature_guide/context_parallel.md) for detailed information.
|
||||
|
||||
## How It Works?
|
||||
|
||||
### Device Distribution
|
||||
|
||||
We introduce new communication domains for PCP and reuse TP for DCP, and this is the new layout of devices for PCP2, DCP2, and TP4.
|
||||
|
||||
|
||||
### Block Table
|
||||
|
||||
CP performs sequence sharding on the KV cache storage. To facilitate efficient storage and access, tokens are stored in an interleaved manner across devices, with the interleaving granularity determined by `cp_kv_cache_interleave_size`, whose default value is `cp_kv_cache_interleave_size=1`, a.k.a. 'token interleave'.
|
||||
|
||||
Given that PCP and DCP behave similarly for KV cache sharding, we refer to them collectively as CP. Specifically, `cp_size = pcp_size * dcp_size`, and `cp_rank = pcp_rank * dcp_size + dcp_rank`.
|
||||
|
||||
As illustrated, a virtual block is defined in the block table, where blocks within the same CP device group form a virtual block. The virtual block size is `virtual_block_size = block_size * cp_size`.
|
||||
|
||||
For any token `x`, referencing the following figure, its (virtual) block index is `x // virtual_block_size`, and the offset within the virtual block is `offset_within_virtual_block = x % virtual_block_size`.
|
||||
The local block index is `local_block_index = offset_within_virtual_block // cp_kv_cache_interleave_size`, and the device number is `target_rank = local_block_index % cp_size`.
|
||||
The offset within the local block is `(local_block_index // cp_size) * cp_kv_cache_interleave_size + offset_within_virtual_block % cp_kv_cache_interleave_size`.
|
||||
|
||||
|
||||
|
||||
Based on the logic above, the `slot_mapping` calculation process is adjusted, and the `slot_mapping` values on each device are modified to ensure the KV cache is sharded along the sequence dimension and stored across different devices as expected.
|
||||
|
||||
The current implementation requires that `block_size % cp_kv_cache_interleave_size == 0`.
|
||||
|
||||
### Decode Context Parallel (DCP)
|
||||
|
||||
As mentioned above, the primary function of DCP is to shard the KV cache along the sequence dimension for storage. Its impact lies in the logic of the decode and chunked prefill phases.
|
||||
|
||||
**Prefill Phase:**
|
||||
As illustrated, during the Chunked Prefill computation, two distinct logic implementations are employed for MLA and GQA backends.
|
||||
|
||||
- In the **MLA backend**, a Context KV Cache `all_gather` operation is performed to aggregate the full KV values.
|
||||
These are then used for attention computation with the Q values of the current chunk.
|
||||
Note that in multi-request scenarios, the directly gathered KV results are interleaved across requests.
|
||||
The `reorg_kvcache` function is used to reorganize the KV cache, ensuring that the KV cache of the same request is stored contiguously.
|
||||
|
||||
- In the **GQA backend**, an `all_gather` is performed along the head dimension for Q.
|
||||
This is because DCP overlaps with the TP communication domain, and the Q heads within a DCP group differ.
|
||||
However, they need to exchange results with the locally computed KV cache for online Softmax updates.
|
||||
To ensure correctness during result updates, the Q values are synchronized across the DCP group via head-dimension `all_gather`.
|
||||
During the result update process, `cp_lse_ag_out_rs` is invoked to aggregate `attn_output` and `attn_lse`, update the results, and perform a reduce-scatter operation on the outputs.
|
||||
Alternatively, we can use an all-to-all communication to exchange the output and LSE results, followed by direct local updates. This approach aligns with the logic adapted for PCP compatibility.
|
||||
|
||||
|
||||
|
||||
**Decode Phase:**
|
||||
The logic during the decode phase is consistent with that of GQA's chunked prefill: an all-gather operation is first performed along the Q head dimension to ensure consistency within the DCP group.
|
||||
After computing the results with the local KV cache, the results are updated via the `cp_lse_ag_out_rs` function.
|
||||
|
||||
|
||||
|
||||
### GLM-5.2 SFA DCP Replicated Indexer
|
||||
|
||||
GLM-5.2 uses Sparse Flash Attention (SFA) with a LightningIndexer. For DCP,
|
||||
the indexer needs a full-sequence view to select the same sparse top-k blocks
|
||||
as non-DCP SFA, while the much larger SFA KV cache should remain sharded to
|
||||
retain DCP's memory benefit. The replicated-indexer path provides this split
|
||||
layout:
|
||||
|
||||
- The LightningIndexer cache is replicated on every DCP rank. Index selection
|
||||
therefore uses the complete sequence and produces globally consistent sparse
|
||||
top-k indices.
|
||||
- The SFA KV cache remains DCP-local. The global indices from the replicated
|
||||
indexer view are remapped to local KV indices before SFA runs.
|
||||
- During prefill or a mixed batch, only the KV blocks referenced by the sparse
|
||||
block table are compacted and all-gathered after the current layer has
|
||||
written its KV cache. The gathered KV uses a remapped block table for SFA,
|
||||
so this path does not all-gather Q and does not need LSE or output
|
||||
post-processing.
|
||||
- Decode-only batches retain the DCP SFA Q-gather and result-merge path.
|
||||
|
||||
This mode is selected automatically for SFA sparse models when
|
||||
`prefill_context_parallel_size=1` and `decode_context_parallel_size>1`. It
|
||||
requires `decode_context_parallel_size == tensor_parallel_size`; PCP combined
|
||||
with this replicated-indexer path is not supported.
|
||||
|
||||
For a GLM-5.2 DSA-CP deployment, enable FlashComm1 and DSA CP and keep the CP
|
||||
interleave size equal to the KV-cache block size:
|
||||
|
||||
```bash
|
||||
export VLLM_ASCEND_ENABLE_FLASHCOMM1=1
|
||||
|
||||
vllm serve <glm-5.2-model> \
|
||||
--tensor-parallel-size <N> \
|
||||
--prefill-context-parallel-size 1 \
|
||||
--decode-context-parallel-size <N> \
|
||||
--block-size <B> \
|
||||
--cp-kv-cache-interleave-size <B> \
|
||||
--additional-config '{"enable_dsa_cp": true}'
|
||||
```
|
||||
|
||||
The replicated indexer increases indexer-cache memory in proportion to the
|
||||
DCP world size; the SFA KV cache itself remains sharded. For this SFA CP path,
|
||||
`cp_kv_cache_interleave_size` must equal `block_size`. A mismatched setting is
|
||||
overridden during configuration validation, but deployments should set both
|
||||
values explicitly to avoid relying on that fallback.
|
||||
|
||||
### Prefill Context Parallel (PCP)
|
||||
|
||||
**Tokens Partition in Head-Tail Style**
|
||||
|
||||
PCP requires splitting the input sequence and ensuring balanced computational load across devices during the prefill phase.
|
||||
We employ a head-tail style for splitting and concatenation: specifically, the sequence is first padded to a length of `2*pcp_size`, then divided into `2*pcp_size` equal parts.
|
||||
The first part is merged with the last part, the second part with the second last part, and so on, thereby assigning computationally balanced chunks to each device.
|
||||
Additionally, since allgather aggregation of KV or Q results in interleaved chunks from different requests, we compute `pcp_allgather_restore_idx` to quickly restore the original order.
|
||||
|
||||
These logics are implemented in the function `_update_tokens_for_pcp`.
|
||||
|
||||
|
||||
|
||||
**Prefill Phase:**
|
||||
|
||||
During the Prefill phase (excluding chunked prefill), we employ an all-gather KV approach to address the issue of incomplete sequences on individual GPUs.
|
||||
It is important to note that we only aggregate the KV values for the current layer at a time, and these are discarded immediately after use, avoiding excessive peak memory usage.
|
||||
This method can also be directly applied to KV cache storage (since the KV cache partitioning method differs from PCP sequence partitioning, it is inevitable that each GPU requires a complete copy of the KV values).
|
||||
All attention backends maintain consistency in this logic.
|
||||
|
||||
Note: While a Ring Attention approach could also facilitate information exchange with lower peak memory and enable computation-communication overlap, we prioritized the all-gather KV implementation after evaluating that the development complexity was high and the benefits of overlap were limited.
|
||||
|
||||
|
||||
|
||||
**Decode Phase:**
|
||||
|
||||
During the decode phase, we only need to add an allgather within the PCP group after the DCP all-to-all communication exchanges the output and LSE, before proceeding with the output update.
|
||||
|
||||
|
||||
|
||||
**Chunked Prefill:**
|
||||
|
||||
Currently, there are three viable approaches for Chunked Prefill compatibility: **AllGatherQ**, **AllGatherKV**, and **Ring-Attn**.
|
||||
Since PCP performs sequence sharding on both the query sequence and the KV cache, we need to ensure that one side has complete information or employ a method like Ring-Attn to perform computations sequentially.
|
||||
The advantages and disadvantages of Ring-Attn will not be elaborated here.
|
||||
|
||||
We have implemented the **AllGatherQ** approach in the GQA attention backend and the **AllGatherKV** approach in the MLA attention backend.
|
||||
The workflow after **AllGatherQ** is identical to the decode phase, while the workflow after **AllGatherKV** is the same as the standard prefill phase.
|
||||
For details, please refer to the diagram below; specific steps will not be repeated.
|
||||
|
||||
One important note: **AllGatherKV** may lead to significant peak memory usage when the context length becomes excessively long.
|
||||
To mitigate this, we adopt a segmented processing strategy.
|
||||
By predefining the maximum amount of KV cache processed per round, we sequentially complete the attention computation and online softmax updates for each segment.
|
||||
|
||||
|
||||
|
||||
### SFA DSA-CP Mixed `o_proj` Path
|
||||
|
||||
SFA DSA-CP mixed execution intentionally reuses the normal TP-sharded `o_proj`.
|
||||
This is part of the DSA-CP mixed data path, not a standalone user-facing `o_proj` TP switch.
|
||||
The mixed path is used when one instance may handle both decode-only and prefill/mixed batches, so `o_proj` must support two layouts at runtime:
|
||||
|
||||
- **Decode-only batches** keep the decode TP path.
|
||||
SFA outputs are exchanged with an all-to-all in the TP group, then the original TP-sharded `o_proj` runs normally.
|
||||
- **Prefill or mixed batches** produce SFA outputs that are not directly compatible with the TP-sharded `o_proj` input layout.
|
||||
Before `o_proj` forward, each rank all-gathers the TP-sharded `o_proj` weight and all input-sharded quantization parameters into temporary full-weight buffers.
|
||||
The full-weight `o_proj` forward runs once for that batch, and the module is then restored to the TP parameter aliases.
|
||||
|
||||
The storage invariant is that the original TP-sharded `o_proj` parameter remains the only persistent source of truth.
|
||||
`o_proj_tp_*` tensors are aliases of the original parameter storage.
|
||||
`o_proj_full_*` tensors are reusable communication buffers for prefill/mixed full-gather execution only.
|
||||
They must not become a second persistent copy of the TP weight.
|
||||
|
||||
This coupling preserves the existing decode TP behavior, supports prefill/mixed DSA-CP batches, and avoids adding an extra configuration path whose state can drift from DSA-CP mixed execution.
|
||||
|
||||
### Related Files
|
||||
|
||||
- slot_mapping computation: `vllm_ascend/worker/block_table.py`
|
||||
- sequences splitting and metadata prepare: `vllm_ascend/worker/model_runner_v1.py`
|
||||
- PCP token splitting and metadata generation: `vllm_ascend/worker/pcp_utils.py`
|
||||
- GQA backend: `vllm_ascend/attention/context_parallel/attention_cp.py`
|
||||
- MLA backend: `vllm_ascend/attention/context_parallel/mla_cp.py`
|
||||
- DSA backend: `vllm_ascend/attention/context_parallel/dsa_cp.py`
|
||||
- SFA backend: `vllm_ascend/attention/context_parallel/sfa_cp.py`
|
||||
281
docs/source/developer_guide/Design_Documents/cpu_binding.md
Normal file
281
docs/source/developer_guide/Design_Documents/cpu_binding.md
Normal file
@@ -0,0 +1,281 @@
|
||||
# CPU Binding
|
||||
|
||||
## Overview
|
||||
|
||||
CPU binding is an **Ascend-native host-side optimization** for vLLM workers on
|
||||
ARM servers. **Starting from vllm-ascend v0.18.0rc1, it is enabled by default
|
||||
through `enable_cpu_binding=True`.**
|
||||
|
||||
The feature does not change model execution logic or numerical results. It only
|
||||
controls CPU placement for the worker process, key runtime threads, memory
|
||||
pages, and NPU IRQs when the host environment allows it. By keeping the main
|
||||
worker, ACL, and release threads on dedicated CPU ranges, it **helps reduce
|
||||
context-switch overhead from scheduler preemption on busy hosts.**
|
||||
|
||||
## Why CPU Binding?
|
||||
|
||||
On multi-socket ARM systems, the Linux scheduler may place worker threads on
|
||||
CPUs far from the NPU that the worker drives. This can increase cross-NUMA
|
||||
traffic, increase thread preemption, and introduce latency jitter. The Ascend
|
||||
backend therefore owns a CPU allocation policy to **reduce cross-NUMA traffic,
|
||||
reduce thread preemption, and improve latency stability** instead of relying on
|
||||
upstream GPU NUMA binding flags.
|
||||
|
||||
This is also why upstream NUMA flags are adapted on Ascend:
|
||||
|
||||
- `--numa-bind` is converted to `additional_config={"enable_cpu_binding": true}`.
|
||||
- `--numa-bind-nodes` and `--numa-bind-cpus` are ignored because Ascend computes CPU pools from NPU topology or global logical NPU IDs.
|
||||
|
||||
## How It Works?
|
||||
|
||||
The allocator derives its plan from runtime host state:
|
||||
|
||||
| Input | Source | Purpose |
|
||||
| --- | --- | --- |
|
||||
| Allowed CPUs | `/proc/self/status` `Cpus_allowed_list` | The only CPUs eligible for binding. Container cpusets are respected. |
|
||||
| Logical NPU map | `npu-smi info -m` | Maps card/chip IDs to global logical NPU IDs and gives `total_logic_npus`. On Ascend 950, `Chip Logic ID` is not reported, so `NPU ID` is used as the logical ID. |
|
||||
| Running NPUs | `npu-smi info` process table, filtered by `ASCEND_RT_VISIBLE_DEVICES` | Identifies the logical NPUs used by this worker process. A2/A3 process rows use `NPU Chip`; Ascend 950 process rows use `NPU ID`. |
|
||||
| Topology affinity | `npu-smi info -t topo` | Provides NPU-to-CPU affinity for `topo_affinity` mode. |
|
||||
| CPU NUMA map | `lscpu -e=CPU,NODE` | Used to extend single-NUMA affinity pools to the next NUMA node. |
|
||||
| Thread topology | `lscpu` `Thread(s) per core` | Determines Ascend 950 cluster size: 8 CPUs for 1 thread per core, 16 CPUs for 2 threads per core. |
|
||||
| UVB polling threads | `ps -Te` | Finds host `uvb_poll_window_thread` threads for Ascend 950 UVB CPU binding. Docker containers must use `--pid=host` to see these host threads. |
|
||||
|
||||
### Strategy Selection
|
||||
|
||||
The binding strategy is selected by Ascend device type:
|
||||
|
||||
| Device type | Strategy | Reason |
|
||||
| --- | --- | --- |
|
||||
| A3 | `global_slice` | A3 uses HCCS card-to-card interconnect. Each NPU is nearly equidistant from all NUMA nodes, so there is no strong NPU-to-NUMA affinity signal. Global logical NPU ID based slicing gives deterministic, non-overlapping CPU pools and CPU/NUMA isolation between workers. |
|
||||
| Ascend 950 | `topo_affinity` | Ascend 950 uses NPU-to-CPU affinity from `npu-smi info -t topo` to choose an affinity NUMA node, then assigns one CPU cluster from that NUMA node to each worker. It also reports process rows by `NPU ID` instead of `NPU Chip`, skips IRQ binding, and binds host UVB polling threads. |
|
||||
| A2 and Atlas 300 inference products | `topo_affinity` | A2 and Atlas 300 inference products provide NPU-to-CPU affinity information through `npu-smi info -t topo`, so they use this topology signal when available. |
|
||||
|
||||
If `topo_affinity` is selected but topo affinity is unavailable, the allocator falls back to `global_slice`.
|
||||
|
||||
### CPU Pool Construction
|
||||
|
||||
#### global_slice
|
||||
|
||||
`global_slice` is designed for devices without a useful NPU-to-CPU affinity
|
||||
signal, including A3. Because A3's **HCCS interconnect makes the distance
|
||||
from each NPU to each NUMA node nearly the same**, topology affinity is not a
|
||||
useful placement signal. The allocator therefore partitions the sorted
|
||||
`allowed_cpus` list by global logical NPU ID.
|
||||
|
||||
1. Determine `total_npus` in this order:
|
||||
- `total_logic_npus` from `npu-smi info -m`
|
||||
- number of topo affinity entries
|
||||
- number of running NPUs
|
||||
2. Compute:
|
||||
- `base = len(allowed_cpus) // total_npus`
|
||||
- `extra = len(allowed_cpus) % total_npus`
|
||||
3. Each logical NPU gets a deterministic slice:
|
||||
- NPU IDs `< extra` receive `base + 1` CPUs.
|
||||
- Remaining NPU IDs receive `base` CPUs.
|
||||
4. Only running NPUs are materialized into `npu_cpu_pool`.
|
||||
|
||||
This is the key property: two independent worker processes with the same cpuset
|
||||
but different visible NPU IDs still get **non-overlapping CPU pools** because
|
||||
both processes slice against the same global NPU ID space. With a NUMA-aligned
|
||||
cpuset, this also provides **CPU/NUMA isolation between workers**, so one worker
|
||||
does not share the same CPU or NUMA slice with another worker.
|
||||
|
||||
`global_slice` requires enough CPUs for the selected device's role split:
|
||||
|
||||
- Devices with IRQ binding require `base >= 5`:
|
||||
2 CPUs for SQ/CQ IRQ binding, at least 1 CPU for the main worker, 1 CPU for
|
||||
ACL thread, and 1 CPU for release thread.
|
||||
|
||||
#### topo_affinity
|
||||
|
||||
`topo_affinity` is designed for A2, Atlas 300 inference products, Ascend 950,
|
||||
and other non-A3 device types. A2 and Atlas 300 inference products expose
|
||||
**meaningful NPU-to-CPU affinity information**, so the allocator starts from NPU
|
||||
topology affinity when it is available and then avoids overlap for shared
|
||||
affinity groups.
|
||||
|
||||
1. Build candidate NPUs from all logical NPUs:
|
||||
- always include running NPUs
|
||||
- include non-running NPUs only when their affinity overlaps this process's allowed cpuset
|
||||
2. For each candidate NPU, intersect topo affinity with `allowed_cpus`.
|
||||
3. If the intersection is empty for a candidate, binding fails for this rank.
|
||||
4. If the affinity CPUs are all on one NUMA node, extend the pool with CPUs from the next NUMA node, constrained by `allowed_cpus`.
|
||||
5. Group NPUs with identical extended pools and split each shared pool evenly across that group.
|
||||
6. Keep only running NPUs in the final `npu_cpu_pool`.
|
||||
|
||||
The non-running candidate step is intentional. It prevents two independent
|
||||
single-card workers from selecting the same CPU range when their visible NPUs
|
||||
share the same topology affinity.
|
||||
|
||||
For Ascend 950, topology affinity is used differently:
|
||||
|
||||
1. Bind all visible host `uvb_poll_window_thread` threads to NUMA0 CPUs except CPU0, constrained by `allowed_cpus`. Docker containers must use `--pid=host` to make these host threads visible.
|
||||
2. Use topo affinity to identify each NPU's single affinity NUMA node.
|
||||
3. Parse `Thread(s) per core` from `lscpu` and set cluster size to 8 CPUs when it is 1, or 16 CPUs when it is 2.
|
||||
4. Split each affinity NUMA's sorted allowed CPU list into contiguous clusters.
|
||||
5. Assign clusters by sorted logical NPU ID, including hidden NPUs that share the same affinity NUMA.
|
||||
6. Keep only running NPUs in the final `npu_cpu_pool`.
|
||||
|
||||
If Ascend 950 topo affinity is missing, spans multiple NUMA nodes, has too few
|
||||
clusters, or reports an unsupported `Thread(s) per core`, worker CPU binding is
|
||||
skipped without raising to the worker process.
|
||||
|
||||
### Role Split
|
||||
|
||||
After a CPU pool is built, the allocator splits it by role:
|
||||
|
||||
For devices with IRQ binding:
|
||||
|
||||
| Role | CPUs |
|
||||
| --- | --- |
|
||||
| SQ/CQ IRQ | `pool[0]`, `pool[1]` |
|
||||
| Main worker process and subthreads | `pool[2:-2]` |
|
||||
| ACL thread | `pool[-2]` |
|
||||
| Release thread | `pool[-1]` |
|
||||
|
||||
For Ascend 950:
|
||||
|
||||
| Role | CPUs |
|
||||
| --- | --- |
|
||||
| Main worker process and subthreads | the whole assigned cluster |
|
||||
| ACL thread | not separately pinned |
|
||||
| Release thread | not separately pinned |
|
||||
|
||||
If a final pool has fewer CPUs than the selected role split requires, binding
|
||||
fails for this rank and the worker logs a warning from the caller. The minimum
|
||||
is 5 CPUs per NPU for devices with IRQ binding. Ascend 950 requires one full
|
||||
cluster per worker.
|
||||
|
||||
## Conditional Host Tuning
|
||||
|
||||
After CPU affinity is applied, CPU binding can also apply two host-side tuning
|
||||
steps when the environment supports them:
|
||||
|
||||
- Memory migration uses `migratepages` to move the worker process's existing
|
||||
pages to the selected NUMA node. This keeps the worker closer to the memory it
|
||||
reads and reduces remote-NUMA memory read latency.
|
||||
- IRQ binding places NPU IRQ handling on the CPUs reserved for the corresponding
|
||||
NPU when `/proc/irq` is writable and IRQ files can be resolved.
|
||||
Ascend 950 skips this step.
|
||||
|
||||
These are conditional parts of CPU binding, not separate feature switches. If a
|
||||
host prerequisite is missing, that step is skipped while CPU thread binding
|
||||
still proceeds. Missing `migratepages` can still leave pages on remote NUMA
|
||||
nodes, so **latency or throughput may regress compared with a full CPU binding
|
||||
setup.**
|
||||
|
||||
## Examples
|
||||
|
||||
### A3 inference server with 640 CPUs and 16 NPUs
|
||||
|
||||
Inputs:
|
||||
|
||||
- `allowed_cpus = [0..639]`
|
||||
- `total_logic_npus = 16`
|
||||
- `running_npu_list = [0..15]`
|
||||
|
||||
Computation:
|
||||
|
||||
- `base = 640 // 16 = 40`
|
||||
- `extra = 0`
|
||||
- Worker `i` driving logical NPU `i` receives CPU slice
|
||||
`[i * 40 .. i * 40 + 39]`.
|
||||
|
||||
Global slice view:
|
||||
|
||||
```text
|
||||
CPU range: 0 639
|
||||
|-- worker0/NPU0 --|-- worker1/NPU1 --| ... |-- worker15/NPU15 --|
|
||||
| 0-39 | 40-79 | ... | 600-639 |
|
||||
```
|
||||
|
||||
Role split inside each worker slice:
|
||||
|
||||
```text
|
||||
40-CPU worker slice
|
||||
| IRQ CPUs | main worker process and subthreads | ACL thread | release thread |
|
||||
| c0-c1 | c2-c37 | c38 | c39 |
|
||||
```
|
||||
|
||||
Concrete examples:
|
||||
|
||||
| Worker | Logical NPU | CPU pool | IRQ CPUs | Main CPUs | ACL CPU | Release CPU |
|
||||
| --- | --- | --- | --- | --- | --- | --- |
|
||||
| 0 | 0 | 0-39 | 0-1 | 2-37 | 38 | 39 |
|
||||
| 1 | 1 | 40-79 | 40-41 | 42-77 | 78 | 79 |
|
||||
| ... | ... | ... | ... | ... | ... | ... |
|
||||
| 15 | 15 | 600-639 | 600-601 | 602-637 | 638 | 639 |
|
||||
|
||||
This layout remains deterministic even when different worker processes share
|
||||
the same cpuset, because slicing is based on the global logical NPU ID.
|
||||
|
||||
### A2 topo_affinity with hidden same-affinity NPUs
|
||||
|
||||
Inputs from an A2 topology:
|
||||
|
||||
- NPU0 affinity: 144-167
|
||||
- NPU2 affinity: 144-167
|
||||
- Process A sees only NPU0
|
||||
- Process B sees only NPU2
|
||||
- Both processes have `allowed_cpus = [144..191]`
|
||||
|
||||
The allocator includes the hidden same-affinity NPU as a candidate in each
|
||||
process, splits the shared extended pool, and then keeps only the visible NPU in
|
||||
the final pool.
|
||||
|
||||
Final pools:
|
||||
|
||||
| Process | Visible NPU | Final CPU pool |
|
||||
| --- | --- | --- |
|
||||
| A | 0 | 144-167 |
|
||||
| B | 2 | 168-191 |
|
||||
|
||||
This avoids overlapping CPU pools even when the two workers are launched as independent single-card services.
|
||||
|
||||
## Logs
|
||||
|
||||
The allocator logs the selected mode and allocation plan:
|
||||
|
||||
```text
|
||||
[cpu_bind_mode] mode=topo_affinity rank=0 visible_npus=[0]
|
||||
The CPU allocation plan is as follows:
|
||||
NPU0: main=[...] acl=[...] release=[...]
|
||||
```
|
||||
|
||||
Ascend 950 uses a different role split, so its plan log does not include ACL or
|
||||
release fields. UVB polling thread binding is reported separately when matching
|
||||
threads are found:
|
||||
|
||||
```text
|
||||
[cpu_bind_mode] mode=topo_affinity rank=0 visible_npus=[0]
|
||||
The CPU allocation plan is as follows:
|
||||
Ascend 950 NPU0: worker=[...]
|
||||
[cpu_bind_ascend_950] uvb_poll_window_thread tids=[...] cpus=[...]
|
||||
```
|
||||
|
||||
## Limitations
|
||||
|
||||
- CPU binding runs only on ARM. It is skipped on x86_64.
|
||||
- Each final NPU pool must have enough CPUs for its role split: at least 5 CPUs
|
||||
for devices with IRQ binding. Ascend 950 requires one complete CPU cluster per worker.
|
||||
- `global_slice` is deterministic and provides CPU/NUMA isolation when the
|
||||
cpuset is NUMA-aligned, but it cannot guarantee NUMA-local pools when CPU
|
||||
numbering or cpuset layout crosses NUMA boundaries.
|
||||
- `topo_affinity` depends on usable output from `npu-smi info -t topo`.
|
||||
- IRQ binding requires writable `/proc/irq` and resolvable PCI/IRQ information.
|
||||
Ascend 950 skips IRQ binding even when `/proc/irq` is writable.
|
||||
- Ascend 950 UVB polling thread binding requires visibility into the host PID
|
||||
namespace. Docker containers must be created with `--pid=host`; otherwise
|
||||
`uvb_poll_window_thread` may not be found.
|
||||
- Memory migration requires `migratepages`; otherwise only memory migration is
|
||||
skipped. CPU affinity still applies, but performance may degrade because
|
||||
existing pages are not moved to the target NUMA node and may be read through
|
||||
higher-latency remote NUMA access.
|
||||
- If an exception escapes the binding flow, `NPUWorker` logs a warning and skips CPU binding for that rank.
|
||||
|
||||
## References
|
||||
|
||||
- Implementation: `vllm_ascend/cpu_binding.py`
|
||||
- Worker integration: `vllm_ascend/worker/worker.py`
|
||||
- Config: `vllm_ascend/ascend_config.py` and `docs/source/user_guide/configuration/additional_config.md`
|
||||
- Tests: `tests/ut/device_allocator/test_cpu_binding.py`
|
||||
@@ -0,0 +1,105 @@
|
||||
# Disaggregated-prefill
|
||||
|
||||
## Why disaggregated-prefill?
|
||||
|
||||
This feature addresses the need to optimize the **Time Per Output Token (TPOT)** and **Time To First Token (TTFT)** in large-scale inference tasks. The motivation is two-fold:
|
||||
|
||||
1. **Adjusting Parallel Strategy and Instance Count for P and D Nodes**
|
||||
Using the disaggregated-prefill strategy, this feature allows the system to flexibly adjust the parallelization strategy (e.g., data parallelism (dp), tensor parallelism (tp), and expert parallelism (ep)) and the instance count for both P (Prefiller) and D (Decoder) nodes. This leads to better system performance tuning, particularly for **TTFT** and **TPOT**.
|
||||
|
||||
2. **Optimizing TPOT**
|
||||
Without the disaggregated-prefill strategy, prefill tasks are inserted during decoding, which results in inefficiencies and delays. Disaggregated-prefill solves this by allowing for better control over the system's **TPOT**. By managing chunked prefill tasks effectively, the system avoids the challenge of determining the optimal chunk size and provides more reliable control over the time taken for generating output tokens.
|
||||
|
||||
---
|
||||
|
||||
## Usage
|
||||
|
||||
vLLM Ascend currently supports two types of connectors for handling KV cache management:
|
||||
|
||||
- **MooncakeConnector**: D nodes pull KV cache from P nodes.
|
||||
- **MooncakeLayerwiseConnector**: P nodes push KV cache to D nodes in a layered manner.
|
||||
|
||||
For step-by-step deployment and configuration, refer to the following guide:
|
||||
[https://docs.vllm.ai/projects/ascend/en/latest/tutorials/features/pd_disaggregation_mooncake_multi_node.html](https://docs.vllm.ai/projects/ascend/en/latest/tutorials/features/pd_disaggregation_mooncake_multi_node.html)
|
||||
|
||||
---
|
||||
|
||||
## How It Works
|
||||
|
||||
### 1. Design Approach
|
||||
|
||||
Under the disaggregated-prefill, a global proxy receives external requests, forwarding prefill to P nodes and decode to D nodes; the KV cache (key-value cache) is exchanged between P and D nodes via peer-to-peer (P2P) communication.
|
||||
|
||||
### 2. Implementation Design
|
||||
|
||||
Our design diagram is shown below, illustrating the pull and push schemes respectively.
|
||||
|
||||
|
||||
|
||||
#### Mooncake Connector
|
||||
|
||||
1. The request is sent to the Proxy's `_handle_completions` endpoint.
|
||||
2. The Proxy calls `select_prefiller` to choose a P node and forwards the request, configuring `kv_transfer_params` with `do_remote_decode=True`, `max_completion_tokens=1`, and `min_tokens=1`.
|
||||
3. After the P node's scheduler finishes prefill, `update_from_output` invokes the schedule connector's `request_finished` to defer KV cache release, constructs `kv_transfer_params` with `do_remote_prefill=True`, and returns to the Proxy.
|
||||
4. The Proxy calls `select_decoder` to choose a D node and forwards the request.
|
||||
5. On the D node, the scheduler marks the request as `RequestStatus.WAITING_FOR_REMOTE_KVS`, pre-allocates KV cache, calls `kv_connector_no_forward` to pull the remote KV cache, then notifies the P node to release KV cache and proceeds with decoding to return the result.
|
||||
|
||||
#### Mooncake Layerwise Connector
|
||||
|
||||
1. The request is sent to the Proxy's `_handle_completions` endpoint.
|
||||
2. The Proxy calls `select_decoder` to choose a D node and forwards the request, configuring `kv_transfer_params` with `do_remote_prefill=True` and setting the `metaserver` endpoint.
|
||||
3. On the D node, the scheduler uses `kv_transfer_params` to mark the request as `RequestStatus.WAITING_FOR_REMOTE_KVS`, pre-allocates KV cache, then calls `kv_connector_no_forward` to send a request to the metaserver and waits for the KV cache transfer to complete.
|
||||
4. The Proxy's `metaserver` endpoint receives the request, calls `select_prefiller` to choose a P node, and forwards it with `kv_transfer_params` set to `do_remote_decode=True`, `max_completion_tokens=1`, and `min_tokens=1`.
|
||||
5. During processing, the P node's scheduler pushes KV cache layer-wise; once all layers pushing is complete, it releases the request and notifies the D node to begin decoding.
|
||||
6. The D node performs decoding and returns the result.
|
||||
|
||||
### 3. Interface Design
|
||||
|
||||
Taking MooncakeConnector as an example, the system is organized into three primary classes:
|
||||
|
||||
- **MooncakeConnector**: Base class that provides core interfaces.
|
||||
- **MooncakeConnectorScheduler**: Interface for scheduling the connectors within the engine core, responsible for managing KV cache transfer requirements and completion.
|
||||
- **MooncakeConnectorWorker**: Interface for managing KV cache registration and transfer in worker processes.
|
||||
|
||||
### 4. Specifications Design
|
||||
|
||||
This feature is flexible and supports various configurations, including setups with MLA and GQA models. It is compatible with A2 and A3 hardware configurations and facilitates scenarios involving equal TP setups and certain unequal TP setups across multiple P and D nodes.
|
||||
|
||||
| Feature | Status |
|
||||
|-------------------------------|----------------|
|
||||
| A2 | 🟢 Functional |
|
||||
| A3 | 🟢 Functional |
|
||||
| equal TP configuration | 🟢 Functional |
|
||||
| unequal TP configuration | 🟢 Functional |
|
||||
| MLA | 🟢 Functional |
|
||||
| GQA | 🟢 Functional |
|
||||
|
||||
- 🟢 Functional: Fully operational, with ongoing optimizations.
|
||||
- 🔵 Experimental: Experimental support, interfaces and functions may change.
|
||||
- 🚧 WIP: Under active development, will be supported soon.
|
||||
- 🟡 Planned: Scheduled for future implementation (some may have open PRs/RFCs).
|
||||
- 🔴 NO plan/Deprecated: No plan or deprecated by vLLM.
|
||||
|
||||
---
|
||||
|
||||
## DFX Analysis
|
||||
|
||||
### 1. Config Parameter Validation
|
||||
|
||||
Validate KV transfer config by checking whether the kv_connector type is supported. On transfer failures, emit clear error logs for diagnostics.
|
||||
|
||||
### 2. Port Conflict Detection
|
||||
|
||||
Before startup, perform a port-usage check on configured ports (e.g., rpc_port, metrics_port, http_port/metaserver) by attempting to bind. If a port is already in use, fail fast and log an error.
|
||||
|
||||
### 3. PD Ratio Validation
|
||||
|
||||
Under non-symmetric PD scenarios, validate the P-to-D tp ratio against expected and scheduling constraints to ensure correct and reliable operation.
|
||||
|
||||
---
|
||||
|
||||
## Limitations
|
||||
|
||||
- Heterogeneous P and D nodes are not supported, for example, running P nodes on A2 and D nodes on A3.
|
||||
|
||||
- In non-symmetric TP configurations, only cases where the P nodes have a higher TP degree than the D nodes and the P TP count is an integer multiple of the D TP count are supported (i.e., P_tp > D_tp and P_tp % D_tp = 0).
|
||||
@@ -0,0 +1,158 @@
|
||||
# Dynamic Chunked Pipeline Parallel (CPP)
|
||||
|
||||
TL;DR CPP uses profiling-based dynamic chunking to equalize per-chunk latency and eliminate pipeline bubbles in PP scenarios.
|
||||
|
||||
## Background
|
||||
|
||||
### Problem Statement
|
||||
|
||||
In Pipeline Parallelism (PP) + Chunked Prefill scenarios, long sequences are split into fixed-size chunks that pass through the pipeline sequentially. Due to the O(n²) computational complexity of Self-Attention, **chunks of the same size take increasingly longer to process as the prefix sequence grows**:
|
||||
|
||||
```text
|
||||
Chunk 1 (history=0): ██████ → Time T1
|
||||
Chunk 2 (history=4K): ████████ → Time T2 > T1
|
||||
Chunk 3 (history=8K): ██████████ → Time T3 > T2
|
||||
Chunk 4 (history=12K): ████████████ → Time T4 > T3
|
||||
```
|
||||
|
||||
This time variance propagates across pipeline stages, causing increased idle waiting (Pipeline Bubble) and significantly reducing GPU utilization.
|
||||
|
||||
### Solution Overview
|
||||
|
||||
Dynamic Chunked Pipeline Parallel uses a **profile-first, then predict** strategy:
|
||||
|
||||
```text
|
||||
Fixed Chunking (equal chunk size, unequal time):
|
||||
|
||||
Stage 0 |■■■■|■■■■■■|■■■■■■■■|■■■■■■■■■■|
|
||||
Stage 1 | |■■■■ |■■■■■■ |■■■■■■■■ |■■■■■■■■■■|
|
||||
↑ bubble ↑ bubble ↑ bubble
|
||||
|
||||
Dynamic Chunking (unequal chunk size, equal time):
|
||||
|
||||
Stage 0 |■■■■■■|■■■■■■|■■■■■■|■■■■■■|
|
||||
Stage 1 | |■■■■■■|■■■■■■|■■■■■■|■■■■■■|
|
||||
↑ no bubble — stages stay in sync
|
||||
```
|
||||
|
||||
The core idea is borrowed from [SGLang's dynamic chunking mechanism](https://lmsys.org/blog/2026-01-15-chunked-pipeline/), with additional enhancements such as online calibration.
|
||||
|
||||
## Design
|
||||
|
||||
### Quadratic Latency Model
|
||||
|
||||
Transformer prefill latency grows quadratically with sequence length due to the O(n²) Self-Attention mechanism:
|
||||
|
||||
$$f(l) = a \cdot l^2 + b \cdot l + c$$
|
||||
|
||||
Where:
|
||||
|
||||
- $a \cdot l^2$: Attention overhead (quadratic)
|
||||
- $b \cdot l$: Linear operations (FFN, projection)
|
||||
- $c$: Fixed overhead (kernel launch)
|
||||
|
||||
### Startup Phase: Profiling
|
||||
|
||||
During engine initialization, the system profiles actual model performance:
|
||||
|
||||
1. **Sampling**: Uniformly sample 64 different chunk sizes from `base_chunk_size` down to near 0
|
||||
2. **Execution**: Perform real model forward passes for each chunk size and precisely measure latency (milliseconds)
|
||||
3. **Fitting**: Fit the quadratic model using least squares
|
||||
4. **Target Setting**: Calculate target per-chunk latency based on `base_chunk_size`
|
||||
|
||||
In PP mode, all workers execute forward passes to stay synchronized, but only the first PP rank's timing results are used for scheduling decisions.
|
||||
|
||||
### Runtime Phase: Dynamic Prediction
|
||||
|
||||
Given current prefix length $L$ and target latency $T = f(\text{base\_chunk\_size}) - f(0)$, the system solves for the next chunk size $x$:
|
||||
|
||||
$$f(L + x) - f(L) = T$$
|
||||
|
||||
Expanding to:
|
||||
|
||||
$$a \cdot x^2 + (2aL + b) \cdot x - T = 0$$
|
||||
|
||||
Solved using the quadratic formula:
|
||||
|
||||
$$x = \frac{-(2aL + b) + \sqrt{(2aL + b)^2 + 4aT}}{2a}$$
|
||||
|
||||
The result goes through post-processing:
|
||||
|
||||
1. **Smoothing**: Blend predicted chunk size with `base_chunk_size` using `smooth_factor`
|
||||
2. **Alignment**: Round down to multiple of `page_size` (minimum 64)
|
||||
3. **Constraints**: Not exceeding `max_model_len - history_len` and `max_num_scheduled_tokens`
|
||||
|
||||
### Online Calibration
|
||||
|
||||
Since profiling only covers sequences up to `max_num_batched_tokens` (typically shorter than real workloads), the system continuously refines the model at runtime.
|
||||
|
||||
**Extended Model (two variables):**
|
||||
|
||||
$$f(C, H) = a \cdot C(C+H) + b \cdot (C+H) + c$$
|
||||
|
||||
Where $C$ is chunk size and $H$ is prefix history length.
|
||||
|
||||
After each batch, feature vectors `[Σ(C+H)·C, Σ(C+H), N]` and actual execution time are recorded. Once enough data points accumulate (5-30), model parameters are updated using least squares.
|
||||
|
||||
## Architecture
|
||||
|
||||
### Key Components
|
||||
|
||||
| Component | Location | Responsibility |
|
||||
|-----------|----------|---------------|
|
||||
| **ChunkSizePredictor** | `vllm_ascend/core/profiling_chunk_predictor.py` | Quadratic model fitting and prediction |
|
||||
| **ProfilingChunkManager** | `vllm_ascend/core/profiling_chunk_predictor.py` | Manage profiling workflow and predictor |
|
||||
| **Scheduler** | `vllm_ascend/core/scheduler_profiling_chunk.py` | Integrate CPP scheduling |
|
||||
| **EngineCore** | `vllm_ascend/patch/platform/patch_profiling_chunk.py` | Startup profiling, record execution time |
|
||||
| **NPUWorker** | `vllm_ascend/worker/worker.py` | Execute real forward pass profiling |
|
||||
| **NPUModelRunner** | `vllm_ascend/worker/model_runner_v1.py` | `profile_cpp=True` mode |
|
||||
|
||||
### Workflow
|
||||
|
||||
```text
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ Startup Phase │
|
||||
├─────────────────────────────────────────────────────────────┤
|
||||
│ 1. EngineCore.init() triggers profiling │
|
||||
│ 2. ProfilingChunkManager samples 64 chunk sizes │
|
||||
│ 3. NPUWorker executes forward passes │
|
||||
│ 4. ChunkSizePredictor fits quadratic model │
|
||||
│ 5. Target latency = f(base_chunk_size) - f(0) │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
↓
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ Runtime Phase │
|
||||
├─────────────────────────────────────────────────────────────┤
|
||||
│ For each prefill chunk: │
|
||||
│ 1. Scheduler queries ChunkSizePredictor │
|
||||
│ 2. Given history length L, solve for optimal chunk size │
|
||||
│ 3. Apply smoothing and alignment │
|
||||
│ 4. Execute chunk │
|
||||
│ 5. Record actual timing for online calibration │
|
||||
│ 6. Update model if enough samples collected │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
## Comparison with SGLang
|
||||
|
||||
| Feature | SGLang Dynamic Chunking | Dynamic Chunked Pipeline Parallel |
|
||||
|---------|------------------------|-----------------------------------|
|
||||
| Profiling method | Preset quadratic function | Real forward pass profiling at startup |
|
||||
| Model fitting | $f(l) = a \cdot l^2 + b \cdot l + c$ | Same + online calibration $f(C,H)$ |
|
||||
| Online updates | None | History-based fitting |
|
||||
| Accuracy | May deviate on different hardware | Adapts to actual hardware performance |
|
||||
| Startup cost | None | ~64 forward passes (tens of seconds) |
|
||||
|
||||
## Constraints
|
||||
|
||||
- **Pipeline Parallelism Required**: Must set `--pipeline-parallel-size > 1`
|
||||
- **Chunked Prefill Required**: Must enable `--enable-chunked-prefill`
|
||||
- **Incompatible with Balance Scheduling**: Cannot enable `VLLM_ASCEND_BALANCE_SCHEDULING`
|
||||
- **Startup Overhead**: Profiling phase adds tens of seconds to initialization
|
||||
- **Memory**: No additional runtime memory overhead; profiling reuses existing dummy_run mechanism
|
||||
|
||||
## References
|
||||
|
||||
- [SGLang Dynamic Chunking Blog](https://lmsys.org/blog/2026-01-15-chunked-pipeline/)
|
||||
- [User Guide](../../user_guide/feature_guide/dynamic_chunk_pipeline_parallel.md)
|
||||
- [Tutorial](../../tutorials/features/dynamic_chunked_pipeline_parallel.md)
|
||||
@@ -0,0 +1,245 @@
|
||||
# Expert Parallelism Load Balancer (EPLB)
|
||||
|
||||
## Why We Need EPLB?
|
||||
|
||||
When using Expert Parallelism (EP), different experts are assigned to different NPUs. Given that the load of various experts may vary depending on the current workload, it is crucial to maintain balanced loads across different NPUs. We adopt a redundant experts strategy by duplicating heavily-loaded experts. Then, we heuristically pack these duplicated experts onto NPUs to ensure load balancing across them. Moreover, thanks to the group-limited expert routing used in MoE models, we also attempt to place experts of the same group on the same node to reduce inter-node data traffic, whenever possible.
|
||||
|
||||
To facilitate reproduction and deployment, vLLM Ascend supports the deployed EP load balancing algorithm in `vllm_ascend/eplb/core/policy`. The algorithm computes a balanced expert replication and placement plan based on the estimated expert loads. Note that the exact method for predicting expert loads is outside the scope of this repository. A common method is to use a moving average of historical statistics.
|
||||
|
||||
|
||||
|
||||
## How to Use EPLB?
|
||||
|
||||
Please refer to the EPLB section of the user guide for detailed information: [How to Use EPLB](../../user_guide/feature_guide/expert_parallelism_load_balancer.md)
|
||||
|
||||
## How It Works?
|
||||
|
||||
**EPLB Module Architecture**
|
||||
|
||||
```shell
|
||||
vllm_ascend
|
||||
├── eplb
|
||||
│ ├── adaptor
|
||||
│ │ └── vllm_adaptor.py
|
||||
│ ├── core
|
||||
│ │ ├── policy
|
||||
│ │ │ ├── policy_abstract.py
|
||||
│ │ │ ├── policy_default_eplb.py
|
||||
│ │ │ ├── policy_factory.py
|
||||
│ │ │ ├── policy_flashlb.py
|
||||
│ │ │ ├── policy_random.py
|
||||
│ │ │ └── policy_swift_balancer.py
|
||||
│ │ ├── eplb_device_transfer_loader.py
|
||||
│ │ ├── eplb_utils.py
|
||||
│ │ └── eplb_worker.py
|
||||
│ ├── eplb_updator.py
|
||||
│ └── utils.py
|
||||
└───────────
|
||||
```
|
||||
|
||||
**1. Adaptor Module**
|
||||
*Handles registration and adaptation for different MoE model types*
|
||||
|
||||
- `vllm_adaptor.py`
|
||||
Implementation supporting Qwen3-MoE and DeepSeek models, standardizing parameter handling for policy algorithms
|
||||
|
||||
**2. Core Module**
|
||||
*Implements core algorithms, updates, and asynchronous processing*
|
||||
|
||||
- **Policy Submodule**
|
||||
*Load balancing algorithms with factory pattern instantiation*
|
||||
- `policy_abstract.py`
|
||||
Abstract class for load balancing strategy interfaces
|
||||
- `policy_default_eplb.py`
|
||||
Default implementation of open-source EPLB paper algorithm
|
||||
- `policy_swift_balancer.py`
|
||||
Enhanced version optimizing expert swaps for low-bandwidth devices (e.g., A2)
|
||||
- `policy_flashlb.py`
|
||||
Threshold-based adjustment reducing operational costs through layer-wise fluctuation detection
|
||||
- `policy_random.py`
|
||||
Random policy for basic testing
|
||||
- `policy_factory.py`
|
||||
Strategy factory for automatic algorithm instantiation
|
||||
|
||||
- `eplb_device_transfer_loader.py`
|
||||
Manages expert table/weight transmission and updates
|
||||
- `eplb_utils.py`
|
||||
Utilities for expert table initialization and mapping
|
||||
- `eplb_worker.py`
|
||||
Asynchronous algorithm orchestration and result processing
|
||||
|
||||
**3. System Components**
|
||||
|
||||
- `eplb_updator.py`
|
||||
Central coordinator for load balancing during inference workflows
|
||||
- `utils.py`
|
||||
General utilities for EPLB interface registration
|
||||
|
||||
*Key Optimizations:*
|
||||
|
||||
1. Maintained original structure while improving technical clarity
|
||||
2. Standardized terminology
|
||||
3. Enhanced algorithm differentiation through concise descriptors
|
||||
4. Improved scoping through hierarchical presentation
|
||||
5. Preserved file/class relationships while optimizing readability
|
||||
|
||||
### Default Algorithm
|
||||
|
||||
#### Hierarchical Load Balancing
|
||||
|
||||
When the number of server nodes evenly divides the number of expert groups, we use the hierarchical load balancing policy to leverage group-limited expert routing. We first pack the expert groups onto nodes evenly, ensuring balanced loads across different nodes. Then, we replicate the experts within each node. Finally, we pack the replicated experts onto individual NPUs to ensure load balancing across them. The hierarchical load balancing policy can be used in the prefilling stage with a smaller expert-parallel size.
|
||||
|
||||
#### Global Load Balancing
|
||||
|
||||
In other cases, we use the global load balancing policy, which replicates experts globally regardless of expert groups, and packs the replicated experts onto individual NPUs. This policy can be adopted in the decoding stage with a larger expert-parallel size.
|
||||
|
||||
### Add a New EPLB Policy
|
||||
|
||||
If you want to add a new eplb policy to vllm_ascend, you must follow these steps:
|
||||
|
||||
1. Inherit the `EplbPolicy` abstract class of `policy_abstract.py` and override the `rebalance_experts` interface, ensuring consistent input parameters `current_expert_table`, `expert_workload` and return types `newplacement`.
|
||||
For example:
|
||||
|
||||
```python
|
||||
class RandomLoadBalance(EplbPolicy):
|
||||
def rebalance_experts(self, current_expert_table, expert_workload):
|
||||
new_table = copy.deepcopy(current_expert_table)
|
||||
num_layers = len(current_expert_table)
|
||||
|
||||
for i in range(num_layers):
|
||||
# randomly choose two card
|
||||
# indices = random.sample(range(num_card), 2)
|
||||
indices = [3, 1]
|
||||
|
||||
# swap redundant experts
|
||||
expert_id_to_exchange = new_table[i][indices[0]][-1].clone()
|
||||
new_table[i][indices[0]][-1] = new_table[i][indices[1]][-1]
|
||||
new_table[i][indices[1]][-1] = expert_id_to_exchange
|
||||
|
||||
return 1, [-i for i in range(num_layers)], new_table
|
||||
```
|
||||
|
||||
2. To add a new EPLB algorithm, include the policy type and its corresponding implementation class in the `PolicyFactory` of `policy_factory.py`.
|
||||
|
||||
### Add a New MoE Model
|
||||
|
||||
**Implementation Guide for Model Integration**
|
||||
|
||||
1. **Adapter File Modification**
|
||||
- Inherit or modify `vllm_ascend/eplb/adaptor/vllm_adaptor.py`
|
||||
- Add processing logic for key parameters:
|
||||
- `num_dense_layers`
|
||||
- `global_expert_num`
|
||||
- `num_roe_layers`
|
||||
- Ensure parameter synchronization in the `model_register` function.
|
||||
|
||||
For example:
|
||||
|
||||
Modify `__init__` of `vllm_adaptor.py` to add a new moe model eplb params:
|
||||
|
||||
```python
|
||||
if self.model.config.model_type == "qwen3_moe":
|
||||
self.num_dense_layers = 0
|
||||
self.global_expert_num = self.model.config.num_experts
|
||||
```
|
||||
|
||||
Modify `model_register` of `vllm_adaptor.py` to register eplb params for new moe model:
|
||||
|
||||
```python
|
||||
if config.model_type == "qwen3_moe":
|
||||
model.num_moe_layers = config.num_hidden_layers
|
||||
```
|
||||
|
||||
2. **MoE Feature Integration**
|
||||
- Extend `vllm_ascend/eplb/utils.py` with MoE-specific methods
|
||||
- Implement required functionality for expert routing or weight management
|
||||
|
||||
3. **Registration Logic Update**
|
||||
- Add patch logic within the `model_register` function
|
||||
- Maintain backward compatibility with existing model types
|
||||
|
||||
4. **Validation & Testing**
|
||||
- Verify parameter consistency across layers
|
||||
- Test cross-device communication for expert tables
|
||||
- Benchmark against baseline implementations (e.g., Qwen3-MoE)
|
||||
|
||||
*Key Implementation Notes:*
|
||||
|
||||
- Preserve existing interface contracts in abstract classes
|
||||
- Use decorators for non-intrusive patch integration
|
||||
- Leverage `eplb_utils.py` for shared expert mapping operations
|
||||
|
||||
## DFX
|
||||
|
||||
### Parameter Validation
|
||||
|
||||
#### Integer Parameters
|
||||
|
||||
All integer input parameters must explicitly specify their maximum and minimum values and be subject to valid value validation. For example, `expert_heat_collection_interval` must be greater than 0:
|
||||
|
||||
```python
|
||||
@staticmethod
|
||||
def check_iterations(iterations):
|
||||
if not isinstance(iterations, int):
|
||||
raise TypeError(f"The {iterations} is not int.")
|
||||
if iterations <= 0:
|
||||
raise ValueError(
|
||||
f"The {iterations} can not be less than or equal to 0.")
|
||||
if iterations > sys.maxsize:
|
||||
raise ValueError(
|
||||
f"The {iterations} can not be larger than {sys.maxsize}")
|
||||
```
|
||||
|
||||
#### File Path
|
||||
|
||||
The file path for EPLB must be checked for legality, such as whether the file path is valid and whether it has appropriate read and write permissions. For example:
|
||||
|
||||
```python
|
||||
@staticmethod
|
||||
def check_expert_map_path(expert_map):
|
||||
if expert_map is None:
|
||||
return
|
||||
if not isinstance(expert_map, str):
|
||||
raise TypeError("The expert_map is not str.")
|
||||
if not expert_map.strip():
|
||||
raise ValueError("The expert_map is not empty.")
|
||||
_, ext = os.path.splitext(expert_map)
|
||||
if ext.lower() != ".json":
|
||||
raise TypeError("The expert_map is not json.")
|
||||
if not os.path.exists(expert_map):
|
||||
raise ValueError("The expert_map does not exist.")
|
||||
try:
|
||||
with open(expert_map, "w", encoding='utf-8') as f:
|
||||
f.read()
|
||||
except Exception as e:
|
||||
raise IOError(
|
||||
f"Fail read expert info from {expert_map}, please check the reading permission of {expert_map} : {e}"
|
||||
)
|
||||
|
||||
```
|
||||
|
||||
### Function Specifications
|
||||
|
||||
#### Initialization Function
|
||||
|
||||
All EPLB parameters must be initialized by default during initialization, with specified parameter types and default values for proper handling.
|
||||
|
||||
#### General Functions
|
||||
|
||||
All method arguments must specify parameter types and default values, and functions must include default return value handling for default arguments. It is recommended to use `try-except` blocks to handle the function body, specifying the type of exception captured and the failure handling (e.g., logging exceptions or returning a failure status).
|
||||
|
||||
### Consistency
|
||||
|
||||
#### Expert Map
|
||||
|
||||
The expert map must be globally unique during initialization and update. In a multi-node scenario during initialization, distributed communication should be used to verify the consistency of expert maps across each rank. If they are inconsistent, the user should be notified of which ranks have inconsistent maps.
|
||||
During the update process, if only a few layers or the expert table of a certain rank has been changed, the updated expert table must be synchronized with the EPLB's context to ensure global consistency.
|
||||
|
||||
#### Expert Weight
|
||||
|
||||
When updating expert weights, ensure that the memory allocated for the expert weights has been released, or that the expert (referring to the old version) is no longer in use.
|
||||
|
||||
## Limitations
|
||||
|
||||
Before using EPLB, start the script and add `export DYNAMIC_EPLB="true"`.
|
||||
Before performing load data collection (or performance data collection), start the script and add `export EXPERT_MAP_RECORD="true"`.
|
||||
20
docs/source/developer_guide/Design_Documents/index.md
Normal file
20
docs/source/developer_guide/Design_Documents/index.md
Normal file
@@ -0,0 +1,20 @@
|
||||
# Design Documents
|
||||
|
||||
This section provides an overview of the features implemented in vLLM Ascend. Developers can refer to this guide to understand how vLLM Ascend works.
|
||||
|
||||
:::{toctree}
|
||||
:caption: Design Documents
|
||||
:maxdepth: 1
|
||||
patch
|
||||
cpu_binding
|
||||
ModelRunner_prepare_inputs
|
||||
disaggregated_prefill
|
||||
eplb_swift_balancer
|
||||
ACL_Graph
|
||||
KV_Cache_Pool_Guide
|
||||
add_custom_aclnn_op
|
||||
context_parallel
|
||||
dynamic_chunked_pipeline_parallel
|
||||
quantization
|
||||
npugraph_ex
|
||||
:::
|
||||
105
docs/source/developer_guide/Design_Documents/npugraph_ex.md
Normal file
105
docs/source/developer_guide/Design_Documents/npugraph_ex.md
Normal file
@@ -0,0 +1,105 @@
|
||||
# Npugraph_ex
|
||||
|
||||
## How Does It Work?
|
||||
|
||||
This is an optimization based on FX graphs, which can be considered an acceleration solution for the aclgraph mode.
|
||||
|
||||
You can get its code [code](https://gitcode.com/Ascend/torchair)
|
||||
|
||||
```{note}
|
||||
Atlas 300I DUO and Atlas 200I Pro do not support `enable_npugraph_ex`. Set --additional-config '{"ascend_compilation_config": {"enable_npugraph_ex":false}}'.
|
||||
```
|
||||
|
||||
## Default FX Graph Optimization
|
||||
|
||||
### FX Graph pass
|
||||
|
||||
- For the intermediate nodes of the model, replace the non-in-place operators contained in the nodes with in-place operators to reduce memory movement during computation and improve performance.
|
||||
- For the original input parameters of the model, if they include in-place operators, Dynamo's Functionalize process will replace the in-place operators with a form of non-in-place operators + copy operators. npugraph_ex will reverse this process, restoring the in-place operators and reducing memory movement.
|
||||
|
||||
### FX fusion pass
|
||||
|
||||
npugraph_ex now provides some operator fusion passes, and more will be added in the future.
|
||||
|
||||
Operator combinations that meet the replacement rules can be replaced with the corresponding fused operators.
|
||||
|
||||
You can get the default [fusion pass list](https://www.hiascend.com/document/detail/zh/Pytorch/2600/modthirdparty/torchairuseguide/docs/zh/npugraph_ex/basic/pattern_fusion_pass.md#功能简介)
|
||||
|
||||
## Custom fusion pass
|
||||
|
||||
Users can register a custom graph fusion pass in npugraph_ex to modify PyTorch FX graphs. The registration relies on the register_replacement API.
|
||||
|
||||
Below is the declaration of this API and a demo of its usage.
|
||||
|
||||
```python
|
||||
register_replacement(search_fn, replace_fn, example_inputs, trace_fn=fwd_only, extra_check=_return_true, search_fn_pattern=None)
|
||||
```
|
||||
|
||||
|Parameter Name| Input/Output |Explanation|Is necessary|
|
||||
|--|--------------|---|-------|
|
||||
|search_fn|Input|This function is the operator combination or calculation logic that you want to recognize in the FX graph, such as the operator combination that needs to be fused|Yes|
|
||||
|replace_fn|Input|When the combination corresponding to search_fn is found in the target graph, this function's computation logic will replace the original subgraph to achieve operator fusion or optimization.|Yes|
|
||||
|example_inputs|Input|Example input tensors used to track search_fn and replace_fn. The shape and dtype of the input should match the actual scenario.|Yes|
|
||||
|trace_fn|Input|By default, only the forward computation graph is tracked, which is suitable for optimization during the inference phase; if training scenarios need to be supported, a function that supports backward tracking can be provided.|No|
|
||||
|extra_check|Input|Find the extra verification function after operator fusion. The function's input parameter must be a Match object from torch._inductor.pattern_matcher, and it is used for further custom checks on the matching result, such as checking whether the fused operators are on the same stream, checking the device type, checking the input shapes, and so on.|No|
|
||||
|search_fn_pattern|Input|A custom pattern object is generally unnecessary to provide. Its definition follows the rules of the native PyTorch MultiOutputPattern object. After passing this parameter, search_fn will no longer be used to match operator combinations; instead, this parameter will be used directly as the matching rule.|No|
|
||||
|
||||
### Usage Example
|
||||
|
||||
```python
|
||||
import functools
|
||||
import torch, torch_npu, npugraph_ex
|
||||
|
||||
from torch._inductor.pattern_matcher import Match
|
||||
from torch._subclasses.fake_tensor import FakeTensorMode
|
||||
from npugraph_ex.core.utils import logger
|
||||
|
||||
# Assume fusing the add operator and the npu_rms_norm operator into the npu_add_rms_norm operator
|
||||
# Define a search_fn to find the operator combinations in the original FX graph before fusion.
|
||||
def search_fn(x1, x2, gamma):
|
||||
xOut = torch.add(x1, x2)
|
||||
y, _ = torch_npu.npu_rms_norm(xOut, gamma)
|
||||
return y, xOut
|
||||
|
||||
# Define a replace_fn, that is, a fusion operator, used to replace operator combinations in the FX graph
|
||||
def replace_fn(x1, x2, gamma):
|
||||
y, _, xOut = torch_npu.npu_add_rms_norm(
|
||||
x1, x2, gamma
|
||||
)
|
||||
return y, xOut
|
||||
|
||||
# extra_check can pass in additional validation logic. Here, it is used to check whether the last dimension of the first input parameter x1 is a specific value; if it is not the specific value, fusion is not allowed.
|
||||
def extra_check(match: Match):
|
||||
x1 = match.kwargs.get("x1")
|
||||
|
||||
if x1 is None:
|
||||
return False
|
||||
if not hasattr(x1, "meta") or "val" not in x1.meta:
|
||||
return False
|
||||
|
||||
a_shape = x1.meta["val"].shape
|
||||
return a_shape[-1] == 7168
|
||||
|
||||
|
||||
# Define some sample inputs to trace search_fn and replace_fn into an FX graph
|
||||
fake_mode = FakeTensorMode()
|
||||
with fake_mode:
|
||||
# sizes/values don't actually matter for initial trace
|
||||
# once we get a possible match we re-trace with the actual values and verify the match still holds
|
||||
input_tensor = functools.partial(torch.empty, (1, 1, 2), device="npu", dtype=torch.float16)
|
||||
kwargs_tensor = functools.partial(torch.empty, 2, device="npu", dtype=torch.float16)
|
||||
|
||||
# Call the npugraph_ex.register_replacement API with search_fn, replace_fn, and example_inputs. If there are additional validations, you can pass them in as extra_check.
|
||||
npugraph_ex.register_replacement(
|
||||
search_fn=search_fn,
|
||||
replace_fn=replace_fn,
|
||||
example_inputs=(input_tensor(), input_tensor(), kwargs_tensor()),
|
||||
extra_check=extra_check
|
||||
)
|
||||
```
|
||||
|
||||
The default fusion pass in npugraph_ex is also implemented based on this API. You can see more examples of using this API in the vllm-ascend and npugraph_ex code repositories.
|
||||
|
||||
### DFX
|
||||
|
||||
By reusing the TORCH_COMPILE_DEBUG environment variable from the PyTorch community, when TORCH_COMPILE_DEBUG=1 is set, it will output the FX graphs throughout the entire process.
|
||||
75
docs/source/developer_guide/Design_Documents/patch.md
Normal file
75
docs/source/developer_guide/Design_Documents/patch.md
Normal file
@@ -0,0 +1,75 @@
|
||||
# Patch in vLLM Ascend
|
||||
|
||||
vLLM Ascend is a platform plugin for vLLM. Due to the different release cycle of vLLM and vLLM Ascend and their hardware limitations, we need to patch some code in vLLM to make it compatible with vLLM Ascend.
|
||||
|
||||
In vLLM Ascend code, we provide a patch module `vllm_ascend/patch` to adapt to changes in vLLM.
|
||||
|
||||
## Principle
|
||||
|
||||
We should keep in mind that Patch is not the best way to make vLLM Ascend compatible. It's just a temporary solution. The best way is to contribute the change to vLLM to make it compatible with vLLM Ascend initially. In vLLM Ascend, we have the basic principle for Patch strategy:
|
||||
|
||||
1. Less is more. Please do not patch unless it's the only way currently.
|
||||
2. Once a patch is added, it's required to describe the future plan for removing the patch.
|
||||
3. Anytime, cleaning the patch code is welcome.
|
||||
|
||||
## How it works
|
||||
|
||||
In `vllm_ascend/patch`, you can see the code structure as follows:
|
||||
|
||||
```shell
|
||||
vllm_ascend/
|
||||
└── patch/
|
||||
├── platform/
|
||||
│ └── patch_xxx.py
|
||||
└── worker/
|
||||
└── patch_yyy.py
|
||||
```
|
||||
|
||||
- **platform**: The patch code in this directory is for patching the code in vLLM Main process. It's called by `vllm_ascend/platform::NPUPlatform::pre_register_and_update` very early when vLLM is initialized.
|
||||
- For online mode, vLLM process calls the platform patch in `vllm/vllm/engine/arg_utils.py::AsyncEngineArgs.add_cli_args` when parsing the CLI args.
|
||||
- For offline mode, vLLM process calls the platform patch in `vllm/vllm/engine/arg_utils.py::EngineArgs.create_engine_config` when parsing the input parameters.
|
||||
- **worker**: The patch code in this directory is for patching the code in vLLM worker process. It's called by `vllm_ascend/worker/worker::NPUWorker::__init__` when the vLLM Worker process is initialized.
|
||||
- For both online and offline mode, vLLM EngineCore process calls the worker patch in `vllm/vllm/worker/worker_base.py::WorkerWrapperBase.init_worker` when initializing the worker process.
|
||||
|
||||
## How to write a patch
|
||||
|
||||
Before writing a patch, following the principle above, we should patch the least code. If it's necessary, we can patch the code in either **platform** or **worker** folder. Here is an example to patch `distributed` module in vLLM.
|
||||
|
||||
1. Decide which version of vLLM we should patch. For example, after analysis, here we want to patch both `0.10.0` and `main` of vLLM.
|
||||
2. Decide which process we should patch. For example, here `distributed` belongs to the vLLM main process, so we should patch `platform`.
|
||||
3. Create the patch file in the right folder. The file should be named as `patch_{module_name}.py`. The example here is `vllm_ascend/patch/platform/patch_distributed.py`.
|
||||
4. Write your patch code in the new file. Here is an example:
|
||||
|
||||
```python
|
||||
import vllm
|
||||
|
||||
def patch_destroy_model_parallel():
|
||||
# your patch code
|
||||
...
|
||||
|
||||
vllm.distributed.parallel_state.destroy_model_parallel = patch_destroy_model_parallel
|
||||
```
|
||||
|
||||
5. Import the patch file in `__init__.py`. In this example, add `import vllm_ascend.patch.platform.patch_distributed` into `vllm_ascend/patch/platform/__init__.py`.
|
||||
6. Add the description of the patch in `vllm_ascend/patch/__init__.py`. The description format is as follows:
|
||||
|
||||
```python
|
||||
# ** File: <The patch file name> **
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# 1. `<The target patch module in vLLM>`
|
||||
# Why:
|
||||
# <Describe the reason why we need to patch>
|
||||
# How:
|
||||
# <Describe the way to patch>
|
||||
# Related PR (if no, explain why):
|
||||
# <Add a link to the related PR in vLLM. If there is no related PR, explain why>
|
||||
# Future Plan:
|
||||
# <Describe the future plan to remove the patch>
|
||||
```
|
||||
|
||||
7. Add the Unit Test and E2E Test. Any newly added code in vLLM Ascend should contain the Unit Test and E2E Test as well. You can find more details in [test guide](../contribution/testing.md)
|
||||
|
||||
## Limitations
|
||||
|
||||
1. In V1 Engine, vLLM starts three kinds of processes: Main process, EngineCore process and Worker process. Now vLLM Ascend can only patch the code in Main process and Worker process by default. If you want to patch the code running in EngineCore process, you should patch EngineCore process entirely during setup. Find the entire code in `vllm.v1.engine.core`. Please override `EngineCoreProc` and `DPEngineCoreProc` entirely.
|
||||
2. If you are running edited vLLM code, the version of vLLM may be changed automatically. For example, if you run the edited vLLM based on v0.9.n, the version of vLLM may be changed to v0.9.nxxx. In this case, the patch for v0.9.n in vLLM Ascend would not work as expected, because vLLM Ascend can't distinguish the version of the vLLM you're using. In this case, you can set the environment variable `VLLM_VERSION` to specify the version of the vLLM you're using, and then the patch for that version (e.g., v0.9.n) should work.
|
||||
114
docs/source/developer_guide/Design_Documents/quantization.md
Normal file
114
docs/source/developer_guide/Design_Documents/quantization.md
Normal file
@@ -0,0 +1,114 @@
|
||||
# Quantization Adaptation Guide
|
||||
|
||||
This document provides guidance for adapting quantization algorithms and models related to **ModelSlim**.
|
||||
|
||||
## Quantization Feature Introduction
|
||||
|
||||
### Quantization Inference Process
|
||||
|
||||
The current process for registering and obtaining quantization methods in vLLM Ascend is as follows:
|
||||
|
||||
|
||||
|
||||
vLLM Ascend registers a custom Ascend quantization method. By configuring the `--quantization ascend` parameter (or `quantization="ascend"` for offline), the quantization feature is enabled. When constructing the `quant_config`, the registered `AscendModelSlimConfig` is initialized and `get_quant_method` is called to obtain the quantization method corresponding to each weight part, stored in the `quant_method` attribute.
|
||||
|
||||
Currently supported quantization methods include `AscendLinearMethod`, `AscendFusedMoEMethod`, `AscendEmbeddingMethod`, and their corresponding non-quantized methods:
|
||||
|
||||
|
||||
|
||||
The quantization method base class defined by vLLM and the overall call flow of quantization methods are as follows:
|
||||
|
||||
|
||||
|
||||
The `embedding` method is generally not implemented for quantization, focusing only on the other three methods.
|
||||
|
||||
The `create_weights` method is used for weight initialization; the `process_weights_after_loading` method is used for weight post-processing, such as transposition, format conversion, data type conversion, etc.; the `apply` method is used to perform activation quantization and quantized matrix multiplication calculations during the forward process.
|
||||
|
||||
We need to implement the `create_weights`, `process_weights_after_loading`, and `apply` methods for different **layers** (**attention**, **mlp**, **MoE (Mixture of Experts)**).
|
||||
|
||||
**Supplement**: When loading the model, the quantized model's description file **quant_model_description.json** needs to be read. This file describes the quantization configuration and parameters for each part of the model weights, for example:
|
||||
|
||||
```json
|
||||
{
|
||||
"model.layers.0.linear_attn.dt_bias": "FLOAT",
|
||||
"model.layers.0.linear_attn.A_log": "FLOAT",
|
||||
"model.layers.0.linear_attn.conv1d.weight": "FLOAT",
|
||||
"model.layers.0.linear_attn.in_proj_qkvz.weight": "W8A8_DYNAMIC",
|
||||
"model.layers.0.linear_attn.in_proj_qkvz.weight_scale": "W8A8_DYNAMIC",
|
||||
"model.layers.0.linear_attn.in_proj_qkvz.weight_offset": "W8A8_DYNAMIC",
|
||||
"model.layers.0.linear_attn.in_proj_ba.weight": "FLOAT",
|
||||
"model.layers.0.linear_attn.norm.weight": "FLOAT",
|
||||
"model.layers.0.linear_attn.out_proj.weight": "FLOAT",
|
||||
"model.layers.0.mlp.gate.weight": "FLOAT",
|
||||
"model.layers.0.mlp.experts.0.gate_proj.weight": "W8A8_DYNAMIC",
|
||||
"model.layers.0.mlp.experts.0.gate_proj.weight_scale": "W8A8_DYNAMIC",
|
||||
"model.layers.0.mlp.experts.0.gate_proj.weight_offset": "W8A8_DYNAMIC"
|
||||
}
|
||||
```
|
||||
|
||||
Based on the above content, we present a brief description of the adaptation process for quantization algorithms and quantized models.
|
||||
|
||||
### Quantization Algorithm Adaptation
|
||||
|
||||
- **Step 1: Algorithm Design**. Define the algorithm ID (e.g., `W4A8_DYNAMIC`), determine supported layers (linear, moe, attention), and design the quantization scheme (static/dynamic, pertensor/perchannel/pergroup).
|
||||
- **Step 2: Registration**. Use the `@register_scheme` decorator in `vllm_ascend/quantization/methods/registry.py` to register your quantization scheme class.
|
||||
|
||||
```python
|
||||
from vllm_ascend.quantization.methods import register_scheme, AscendLinearScheme, AscendMoEScheme
|
||||
|
||||
@register_scheme("W4A8_DYNAMIC", "linear")
|
||||
class AscendW4A8DynamicLinearMethod(AscendLinearScheme):
|
||||
...
|
||||
|
||||
@register_scheme("W4A8_DYNAMIC", "moe")
|
||||
class AscendW4A8DynamicFusedMoEMethod(AscendMoEScheme):
|
||||
...
|
||||
```
|
||||
|
||||
- **Step 3: Implementation**. Create an algorithm implementation file, such as `vllm_ascend/quantization/methods/w4a8.py`, and implement the method class and logic.
|
||||
- **Step 4: Testing**. Use your algorithm to generate quantization configurations and verify correctness and performance on target models and hardware.
|
||||
|
||||
### Quantized Model Adaptation
|
||||
|
||||
Adapting a new quantized model requires ensuring the following three points:
|
||||
|
||||
- The original model has been successfully adapted in `vLLM Ascend`.
|
||||
- **Fused Module Mapping**: Add the model's `model_type` to `packed_modules_model_mapping` in `vllm_ascend/quantization/modelslim_config.py` (e.g., `qkv_proj`, `gate_up_proj`, `experts`) to ensure sharding consistency and correct loading.
|
||||
|
||||
```python
|
||||
packed_modules_model_mapping = {
|
||||
"qwen3_moe": {
|
||||
"qkv_proj": [
|
||||
"q_proj",
|
||||
"k_proj",
|
||||
"v_proj",
|
||||
],
|
||||
"gate_up_proj": [
|
||||
"gate_proj",
|
||||
"up_proj",
|
||||
],
|
||||
"experts":
|
||||
["experts.0.gate_proj", "experts.0.up_proj", "experts.0.down_proj"],
|
||||
},
|
||||
}
|
||||
```
|
||||
|
||||
- All quantization algorithms used by the quantized model have been integrated into the `quantization` module.
|
||||
|
||||
## Currently Supported Quantization Algorithms
|
||||
|
||||
vLLM Ascend supports multiple quantization algorithms. The following table provides an overview of each quantization algorithm based on the implementation in the `vllm_ascend.quantization` module:
|
||||
|
||||
| Algorithm | Weight | Activation | Weight Granularity | Activation Granularity | Type | Description |
|
||||
| ------------------------ | ------ | ---------- | ------------------ | ---------------------- | ------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
|
||||
| `W4A16` | INT4 | FP16/BF16 | Per-Group | Per-Tensor | Static | 4-bit weight quantization with 16-bit activation precision, specifically designed for MoE model expert layers, supporting int32 format weight packing |
|
||||
| `W8A16` | INT8 | FP16/BF16 | Per-Channel | Per-Tensor | Static | 8-bit weight quantization with 16-bit activation precision, balancing accuracy and performance, suitable for linear layers |
|
||||
| `W8A8` | INT8 | INT8 | Per-Channel | Per-Tensor | Static | Static activation quantization, suitable for scenarios requiring high precision |
|
||||
| `W8A8_DYNAMIC` | INT8 | INT8 | Per-Channel | Per-Token | Dynamic | Dynamic activation quantization with per-token scaling factor calculation |
|
||||
| `W4A8_DYNAMIC` | INT4 | INT8 | Per-Group | Per-Token | Dynamic | Supports both direct per-channel quantization to 4-bit and two-step quantization (per-channel to 8-bit then per-group to 4-bit) |
|
||||
| `W4A4_FLATQUANT_DYNAMIC` | INT4 | INT4 | Per-Channel | Per-Token | Dynamic | Uses FlatQuant for activation distribution smoothing before 4-bit dynamic quantization, with additional matrix multiplications for precision preservation |
|
||||
| `W8A8_MIX` | INT8 | INT8 | Per-Channel | Per-Tensor/Token | Mixed | We support two deployment modes: PD Colocation (dynamic quantization for both P and D) and PD Disaggregation (dynamic-quant P and static-quant D) |
|
||||
|
||||
**Static vs Dynamic:** Static quantization uses pre-computed scaling factors with better performance, while dynamic quantization computes scaling factors on-the-fly for each token/activation tensor with higher precision.
|
||||
|
||||
**Granularity:** Refers to the scope of scaling factor computation (e.g., per-tensor, per-channel, per-group).
|
||||
Reference in New Issue
Block a user