9d09488b4a
### What this PR does / why we need it?
This PR implements Context Parallelism (CP) support for the Qwen3-Next
model, including PCP (Parallel Context Parallelism) and DCP
(Dynamic/Data Context Parallelism).
- vLLM version: v0.15.0
- vLLM main:
f176443446
---------
Signed-off-by: SunnyLee219 <3294305115@qq.com>
Signed-off-by: Jingchun Gao <gaojingchun1@huawei.com>
Signed-off-by: 白永斌 <baiyongbin3@h-partners.com>
Signed-off-by: Bai Yongbin <845473182@qq.com>
Co-authored-by: SunnyLee219 <3294305115@qq.com>
Co-authored-by: Jingchun Gao <gaojingchun1@huawei.com>
Co-authored-by: 白永斌 <baiyongbin3@h-partners.com>
Co-authored-by: Mengqing Cao <cmq0113@163.com>
122 lines
3.4 KiB
Python
122 lines
3.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang
|
|
#
|
|
# This file contains code copied from the flash-linear-attention project.
|
|
# The original source code was licensed under the MIT license and included
|
|
# the following copyright notice:
|
|
# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
|
|
|
|
# ruff: noqa: E501
|
|
# mypy: ignore-errors
|
|
|
|
import torch
|
|
from vllm.triton_utils import tl, triton
|
|
|
|
from .utils import prepare_chunk_offsets
|
|
|
|
|
|
@triton.heuristics(
|
|
{
|
|
"IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
|
|
}
|
|
)
|
|
@triton.jit(do_not_specialize=["T"])
|
|
def chunk_fwd_kernel_o_update(
|
|
h,
|
|
h_update,
|
|
updated_h_state,
|
|
cu_seqlens,
|
|
chunk_offsets,
|
|
T,
|
|
H: tl.constexpr,
|
|
Hg: tl.constexpr,
|
|
K: tl.constexpr,
|
|
V: tl.constexpr,
|
|
BT: tl.constexpr,
|
|
BK: tl.constexpr,
|
|
BV: tl.constexpr,
|
|
IS_VARLEN: tl.constexpr,
|
|
):
|
|
i_v, i_nh = tl.program_id(0), tl.program_id(1)
|
|
i_n, i_h = i_nh // H, i_nh % H # splitting by the head of the req
|
|
|
|
if IS_VARLEN:
|
|
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
|
|
T = eos - bos
|
|
NT = tl.cdiv(T, BT)
|
|
boh = tl.load(chunk_offsets + i_n).to(tl.int64)
|
|
else:
|
|
bos, eos = i_n * T, i_n * T + T
|
|
NT = tl.cdiv(T, BT)
|
|
boh = i_n * NT
|
|
|
|
# offset calculation
|
|
updated_h_state += (i_n * H + i_h).to(tl.int64) * K * V
|
|
|
|
for i_t in range(NT):
|
|
i_tg = boh + i_t
|
|
h_base = h + (i_tg * H + i_h).to(tl.int64) * K * V
|
|
hupd_base = h_update + ((i_tg + i_n) * H + i_h).to(tl.int64) * K * K
|
|
|
|
for i_k in range(tl.cdiv(K, BK)):
|
|
p_h = tl.make_block_ptr(h_base, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0))
|
|
p_hupd = tl.make_block_ptr(hupd_base, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BK), (1, 0))
|
|
p_updated_h_state = tl.make_block_ptr(
|
|
updated_h_state, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)
|
|
)
|
|
|
|
# [BK, BV]
|
|
b_h = tl.load(p_h, boundary_check=(0, 1))
|
|
# [BK, BK]
|
|
b_hupd = tl.load(p_hupd, boundary_check=(0, 1))
|
|
# [BK, BV]
|
|
b_updated_h_state = tl.load(p_updated_h_state, boundary_check=(0, 1))
|
|
|
|
b_h += tl.dot(b_hupd.to(tl.bfloat16), b_updated_h_state.to(tl.bfloat16))
|
|
tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1))
|
|
|
|
|
|
def chunk_fwd_o_update(
|
|
q: torch.Tensor,
|
|
v: torch.Tensor,
|
|
h: torch.Tensor,
|
|
h_update: torch.Tensor,
|
|
updated_h_state: torch.Tensor,
|
|
cu_seqlens: torch.LongTensor | None = None,
|
|
chunk_size: int = 64,
|
|
) -> torch.Tensor:
|
|
B, T, Hg, K, V = *q.shape, v.shape[-1]
|
|
H = v.shape[-2]
|
|
BT = chunk_size
|
|
|
|
if cu_seqlens is None:
|
|
N, chunk_offsets = B, None
|
|
else:
|
|
N, chunk_offsets = (
|
|
len(cu_seqlens) - 1,
|
|
prepare_chunk_offsets(cu_seqlens, BT),
|
|
)
|
|
|
|
def grid(meta):
|
|
return (triton.cdiv(V, meta["BV"]), N * H)
|
|
|
|
chunk_fwd_kernel_o_update[grid](
|
|
h=h,
|
|
h_update=h_update,
|
|
updated_h_state=updated_h_state,
|
|
cu_seqlens=cu_seqlens,
|
|
chunk_offsets=chunk_offsets,
|
|
T=T,
|
|
H=H,
|
|
Hg=Hg,
|
|
K=K,
|
|
V=V,
|
|
BT=BT,
|
|
BK=128,
|
|
BV=128,
|
|
num_warps=4,
|
|
num_stages=2,
|
|
)
|
|
return h
|