wehub-resource-sync
296 lines
11 KiB
Python
296 lines
11 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
from types import SimpleNamespace
|
|
from typing import TYPE_CHECKING, Any, cast
|
|
|
|
import torch
|
|
|
|
from vllm.config import VllmConfig, get_layers_from_vllm_config
|
|
from vllm.distributed import get_dcp_group, get_pcp_group
|
|
from vllm.logger import init_logger
|
|
from vllm.v1.attention.backend import CommonAttentionMetadata
|
|
from vllm.v1.attention.backends.utils import split_decodes_prefills_and_extends
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
|
|
else:
|
|
AttentionLayerBase = object
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
def check_attention_cp_compatibility(vllm_config: VllmConfig) -> None:
|
|
pcp_size = vllm_config.parallel_config.prefill_context_parallel_size
|
|
dcp_size = vllm_config.parallel_config.decode_context_parallel_size
|
|
interleave_size = vllm_config.parallel_config.cp_kv_cache_interleave_size
|
|
if pcp_size * dcp_size > 1:
|
|
layer_type = cast(type[Any], AttentionLayerBase)
|
|
layers = get_layers_from_vllm_config(vllm_config, layer_type)
|
|
for layer in layers.values():
|
|
layer_impl = getattr(layer, "impl", None)
|
|
if layer_impl is None:
|
|
continue
|
|
if vllm_config.speculative_config is not None and interleave_size > 1:
|
|
assert layer_impl.supports_mtp_with_cp_non_trivial_interleave_size, (
|
|
"MTP with cp_kv_cache_interleave_size > 1 is not "
|
|
f"supported in {layer_impl.__class__.__name__}."
|
|
)
|
|
if dcp_size > 1:
|
|
assert layer_impl.need_to_return_lse_for_decode, (
|
|
"Decode Context Parallelism (DCP) requires attention "
|
|
"implementations to return the softmax LSE during decode, "
|
|
f"but {layer_impl.__class__.__name__} does not. "
|
|
"Try a different backend by setting "
|
|
"--attention-backend or disable DCP."
|
|
)
|
|
|
|
if pcp_size > 1:
|
|
assert layer_impl.supports_pcp, (
|
|
"PCP requires attention impls' support, "
|
|
f"but the impl {layer_impl.__class__.__name__} "
|
|
"does not support PCP."
|
|
)
|
|
|
|
|
|
def get_total_cp_world_size():
|
|
try:
|
|
pcp_world_size = get_pcp_group().world_size
|
|
except AssertionError:
|
|
# PCP might not be initialized in testing
|
|
pcp_world_size = 1
|
|
try:
|
|
dcp_world_size = get_dcp_group().world_size
|
|
except AssertionError:
|
|
# DCP might not be initialized in testing
|
|
dcp_world_size = 1
|
|
return dcp_world_size * pcp_world_size
|
|
|
|
|
|
def get_dcp_dummy_context_len(
|
|
dcp_world_size: int,
|
|
cp_kv_cache_interleave_size: int,
|
|
has_kv_cache_config: bool,
|
|
create_mixed_batch: bool,
|
|
is_graph_capturing: bool,
|
|
uniform_decode: bool,
|
|
) -> int:
|
|
if (
|
|
dcp_world_size <= 1
|
|
or not has_kv_cache_config
|
|
or not (create_mixed_batch or (is_graph_capturing and uniform_decode))
|
|
):
|
|
return 0
|
|
return dcp_world_size * cp_kv_cache_interleave_size
|
|
|
|
|
|
def prepare_dcp_dummy_context_metadata(
|
|
*,
|
|
input_batch: Any,
|
|
kv_cache_config: Any,
|
|
query_pos: Any,
|
|
positions: torch.Tensor,
|
|
query_start_loc: Any,
|
|
num_reqs: int,
|
|
num_tokens_unpadded: int,
|
|
dcp_dummy_context_len: int,
|
|
) -> None:
|
|
"""Populate valid fake KV metadata for DCP CUDA graph warmup/capture."""
|
|
if dcp_dummy_context_len == 0:
|
|
return
|
|
|
|
# DCP graph warmup may exercise context attention, so block-table entries
|
|
# must point at allocated KV blocks.
|
|
assert kv_cache_config is not None
|
|
max_valid_block_id = kv_cache_config.num_blocks - 1
|
|
assert max_valid_block_id > 0
|
|
for blk_table in input_batch.block_table.block_tables:
|
|
max_row_blocks = (
|
|
blk_table.max_num_blocks_per_req // blk_table.blocks_per_kv_block
|
|
)
|
|
block_ids = [
|
|
(block_idx % max_valid_block_id) + 1 for block_idx in range(max_row_blocks)
|
|
]
|
|
for req_idx in range(num_reqs):
|
|
blk_table.add_row(block_ids, req_idx)
|
|
blk_table.commit_block_table(num_reqs)
|
|
|
|
query_pos.copy_to_gpu(num_tokens_unpadded)
|
|
positions[:num_tokens_unpadded] = (
|
|
query_pos.gpu[:num_tokens_unpadded] + dcp_dummy_context_len
|
|
)
|
|
input_batch.block_table.compute_slot_mapping(
|
|
num_reqs,
|
|
query_start_loc.gpu[: num_reqs + 1],
|
|
positions[:num_tokens_unpadded],
|
|
)
|
|
|
|
|
|
def should_skip_dcp_context_attention(context_kv_lens_cpu: torch.Tensor) -> bool:
|
|
"""Whether DCP context attention can be skipped for this batch.
|
|
|
|
Must be computed from rank-invariant inputs only (the global context
|
|
lengths, NOT this rank's local share from get_dcp_local_seq_lens): the
|
|
non-skip path in _forward_with_dcp issues DCP collectives (query
|
|
all-gather + LSE combine), so every DCP rank must take the same branch.
|
|
A rank can hold zero local context tokens while other ranks still hold
|
|
context for the same batch.
|
|
"""
|
|
return int(context_kv_lens_cpu.max().item()) == 0
|
|
|
|
|
|
def split_dcp_context_queries(
|
|
query_start_loc: torch.Tensor,
|
|
seq_lens_cpu_upper_bound: torch.Tensor | None,
|
|
max_query_len: int,
|
|
num_actual_tokens: int,
|
|
) -> tuple[int, int, int, int]:
|
|
"""Split reordered DCP context queries into decode and extend regions."""
|
|
num_reqs = query_start_loc.shape[0] - 1
|
|
if max_query_len <= 1:
|
|
return num_reqs, 0, num_actual_tokens, 0
|
|
if seq_lens_cpu_upper_bound is None:
|
|
return 0, num_reqs, 0, num_actual_tokens
|
|
|
|
common_attn_metadata = cast(
|
|
CommonAttentionMetadata,
|
|
SimpleNamespace(
|
|
max_query_len=max_query_len,
|
|
num_reqs=num_reqs,
|
|
num_actual_tokens=num_actual_tokens,
|
|
query_start_loc_cpu=query_start_loc,
|
|
seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
|
|
is_prefilling=None,
|
|
),
|
|
)
|
|
(
|
|
num_decodes,
|
|
num_extends,
|
|
_num_prefills,
|
|
num_decode_tokens,
|
|
num_extend_tokens,
|
|
_num_prefill_tokens,
|
|
) = split_decodes_prefills_and_extends(common_attn_metadata)
|
|
return num_decodes, num_extends, num_decode_tokens, num_extend_tokens
|
|
|
|
|
|
def should_split_fa2_dcp_context_attention(
|
|
fa_version: int | None,
|
|
max_query_len: int,
|
|
num_reqs: int,
|
|
num_decode_reqs: int,
|
|
num_context_prefill_reqs: int,
|
|
) -> bool:
|
|
num_prefills = num_reqs - num_decode_reqs
|
|
# TODO: Remove this FA2-only DCP compatibility path once FA4 supports
|
|
# the Qwen3.5 head_size=256 shape on Blackwell and can be used here.
|
|
# FA2 paged-varlen context attention can fail for DCP mixed batches when
|
|
# decode rows, context-bearing extend rows, and zero-context pure prefill
|
|
# rows are submitted together.
|
|
return (
|
|
fa_version == 2
|
|
and max_query_len > 1
|
|
and num_prefills > 0
|
|
and (num_decode_reqs > 0 or num_context_prefill_reqs < num_prefills)
|
|
)
|
|
|
|
|
|
def run_split_fa2_dcp_context_attention(
|
|
flash_attn_varlen_func: Any,
|
|
query_across_dcp: torch.Tensor,
|
|
key_cache: torch.Tensor,
|
|
value_cache: torch.Tensor,
|
|
dcp_context_out: torch.Tensor,
|
|
cu_seqlens_q: torch.Tensor,
|
|
max_seqlen_q: int,
|
|
dcp_context_kv_lens: torch.Tensor,
|
|
max_dcp_context_kv_len: int,
|
|
softmax_scale: float,
|
|
alibi_slopes: torch.Tensor | None,
|
|
sliding_window_size: list[int] | None,
|
|
block_table: torch.Tensor,
|
|
softcap: float,
|
|
fa_version: int,
|
|
q_descale: torch.Tensor | None,
|
|
k_descale: torch.Tensor | None,
|
|
v_descale: torch.Tensor | None,
|
|
max_num_splits: int,
|
|
num_heads: int,
|
|
dcp_world_size: int,
|
|
num_decode_reqs: int,
|
|
num_context_prefill_reqs: int,
|
|
num_decode_tokens: int,
|
|
num_context_prefill_tokens: int,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
dcp_context_out.zero_()
|
|
context_lse = torch.full(
|
|
(num_heads * dcp_world_size, query_across_dcp.shape[0]),
|
|
-torch.inf,
|
|
dtype=torch.float32,
|
|
device=query_across_dcp.device,
|
|
)
|
|
|
|
if num_decode_tokens > 0:
|
|
_, decode_context_lse = flash_attn_varlen_func(
|
|
q=query_across_dcp[:num_decode_tokens],
|
|
k=key_cache,
|
|
v=value_cache,
|
|
out=dcp_context_out[:num_decode_tokens],
|
|
cu_seqlens_q=cu_seqlens_q[: num_decode_reqs + 1],
|
|
max_seqlen_q=1,
|
|
seqused_k=dcp_context_kv_lens[:num_decode_reqs],
|
|
max_seqlen_k=max_dcp_context_kv_len,
|
|
softmax_scale=softmax_scale,
|
|
causal=False,
|
|
alibi_slopes=alibi_slopes,
|
|
window_size=sliding_window_size,
|
|
block_table=block_table[:num_decode_reqs],
|
|
softcap=softcap,
|
|
return_softmax_lse=True,
|
|
scheduler_metadata=None,
|
|
fa_version=fa_version,
|
|
q_descale=q_descale[:num_decode_reqs] if q_descale is not None else None,
|
|
k_descale=k_descale[:num_decode_reqs] if k_descale is not None else None,
|
|
v_descale=v_descale[:num_decode_reqs] if v_descale is not None else None,
|
|
num_splits=max_num_splits,
|
|
)
|
|
context_lse[:, :num_decode_tokens] = decode_context_lse
|
|
|
|
if num_context_prefill_tokens > 0:
|
|
prefill_start = num_decode_tokens
|
|
prefill_end = prefill_start + num_context_prefill_tokens
|
|
prefill_query_start_loc = (
|
|
cu_seqlens_q[
|
|
num_decode_reqs : num_decode_reqs + num_context_prefill_reqs + 1
|
|
]
|
|
- num_decode_tokens
|
|
)
|
|
prefill_req_slice = slice(
|
|
num_decode_reqs, num_decode_reqs + num_context_prefill_reqs
|
|
)
|
|
_, prefill_context_lse = flash_attn_varlen_func(
|
|
q=query_across_dcp[prefill_start:prefill_end],
|
|
k=key_cache,
|
|
v=value_cache,
|
|
out=dcp_context_out[prefill_start:prefill_end],
|
|
cu_seqlens_q=prefill_query_start_loc,
|
|
max_seqlen_q=max_seqlen_q,
|
|
seqused_k=dcp_context_kv_lens[prefill_req_slice],
|
|
max_seqlen_k=max_dcp_context_kv_len,
|
|
softmax_scale=softmax_scale,
|
|
causal=False,
|
|
alibi_slopes=alibi_slopes,
|
|
window_size=sliding_window_size,
|
|
block_table=block_table[prefill_req_slice],
|
|
softcap=softcap,
|
|
return_softmax_lse=True,
|
|
scheduler_metadata=None,
|
|
fa_version=fa_version,
|
|
q_descale=q_descale[prefill_req_slice] if q_descale is not None else None,
|
|
k_descale=k_descale[prefill_req_slice] if k_descale is not None else None,
|
|
v_descale=v_descale[prefill_req_slice] if v_descale is not None else None,
|
|
num_splits=max_num_splits,
|
|
)
|
|
context_lse[:, prefill_start:prefill_end] = prefill_context_lse
|
|
|
|
return dcp_context_out, context_lse
|