ce5544bfc1
### What this PR does / why we need it?
To support prefix cache for Qwen3.5/Next in vLLM-Ascend, this PR mainly
follows the design in
[#30877](https://github.com/vllm-project/vllm/pull/30877) and inherits
changes to functions which are overridden in vLLM-Ascend.
Note:
1. `--mamba-cache-mode align` && PD disaggregation is still not
supported yet in vLLM v0.17.0(see
https://github.com/vllm-project/vllm/blob/main/vllm/v1/core/sched/scheduler.py#L295).
2. The current implementation of hybrid kv cache might result in a very
large block_size when scheduling. For example, if we run Qwen3.5-35B-A3B
with `-tp 2`, the block_size is adjusted to 2048, which means that any
prefix shorter than 2048 will never be cached. Although this behavior is
consistent with vLLM, it still needs improvements in the future.
3. `--mamba-cache-mode align` requires to copy mamba states during
forward steps. vLLM uses a triton kernel to implement it. However, the
original version run into some bugs on Ascend hardwares. Thus we patch a
new triton kernel to avoid this bug.
### Does this PR introduce _any_ user-facing change?
To use mamba prefix cache, set `--enable-prefix-caching` and
`--mamba-cache-mode align`. Note that the mamba state copy function(see
[do_mamba_copy_block](https://github.com/vllm-project/vllm/blob/main/vllm/v1/worker/mamba_utils.py#L132))
does not provide a torch native version, thus it might have trouble if
users can't use triton.
- vLLM version: v0.16.0
- vLLM main:
4034c3d32e
---------
Signed-off-by: Angazenn <supperccell@163.com>
32 lines
1.1 KiB
Python
32 lines
1.1 KiB
Python
# Adapt from https://github.com/vllm-project/vllm/blob/main/vllm/v1/worker/mamba_utils.py
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
|
|
from vllm.triton_utils import tl, triton
|
|
|
|
|
|
@triton.jit
|
|
def batch_memcpy_kernel(src_ptrs, dst_ptrs, sizes, BLOCK_SIZE: tl.constexpr):
|
|
pid = tl.program_id(0)
|
|
|
|
src_ptr = tl.load(src_ptrs + pid)
|
|
dst_ptr = tl.load(dst_ptrs + pid)
|
|
size = tl.load(sizes + pid)
|
|
|
|
# We need to mv pointer_type cast outside the loop.
|
|
# Otherwise it causes potential bugs.
|
|
src_ptr = src_ptr.to(tl.pointer_type(tl.uint8))
|
|
dst_ptr = dst_ptr.to(tl.pointer_type(tl.uint8))
|
|
|
|
offsets = tl.arange(0, BLOCK_SIZE)
|
|
for i in range(0, size, BLOCK_SIZE):
|
|
mask = (i + offsets) < size
|
|
|
|
curr_src_ptr = src_ptr + i + offsets
|
|
curr_dst_ptr = dst_ptr + i + offsets
|
|
|
|
# cache_modifier=".cg" bypasses L1 cache for streaming data.
|
|
data = tl.load(curr_src_ptr, mask=mask, cache_modifier=".cg")
|
|
tl.store(curr_dst_ptr, data, mask=mask, cache_modifier=".cg")
|