Yongye ZhuandClaude Opus 5.5
a6009ea85b
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
1977 lines
79 KiB
Python
1977 lines
79 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
import vllm.envs as envs
|
|
from vllm.config import CUDAGraphMode, VllmConfig, get_current_vllm_config
|
|
from vllm.distributed import get_dcp_group, get_pcp_group
|
|
from vllm.logger import init_logger
|
|
from vllm.model_executor.warmup.jit_warmup import kernel_launcher, zip_inputs
|
|
from vllm.model_executor.warmup.jit_warmup_triton_helper import (
|
|
LaunchSpec,
|
|
TritonPointerInputVariant,
|
|
TritonWarmupTensor,
|
|
VllmTritonJitKernel,
|
|
triton_scalar_specialization_rep,
|
|
)
|
|
from vllm.platforms import current_platform
|
|
from vllm.triton_utils import tl, triton
|
|
from vllm.utils.deep_gemm import (
|
|
get_paged_mqa_logits_metadata,
|
|
get_paged_mqa_page_sizes,
|
|
has_deep_gemm,
|
|
native_next_n_supported,
|
|
)
|
|
from vllm.utils.math_utils import round_down
|
|
from vllm.utils.platform_utils import num_compute_units
|
|
from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d
|
|
from vllm.v1.attention.backend import (
|
|
AttentionBackend,
|
|
AttentionCGSupport,
|
|
AttentionMetadataBuilder,
|
|
CommonAttentionMetadata,
|
|
MultipleOf,
|
|
)
|
|
from vllm.v1.attention.backends.mla.compressor_utils import get_compressed_slot_mapping
|
|
from vllm.v1.attention.backends.mla.sparse_utils import request_row_bounds
|
|
from vllm.v1.attention.backends.utils import (
|
|
get_dcp_local_seq_lens,
|
|
split_decodes_and_prefills,
|
|
)
|
|
from vllm.v1.kv_cache_interface import (
|
|
AttentionSpec,
|
|
KVCacheLayout,
|
|
KVCacheSpec,
|
|
MLAAttentionSpec,
|
|
)
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
# The DSA indexer K cache is always quantized; "auto" means fp8 (V3.2 layout)
|
|
# and mxfp4 is the opt-in Blackwell and gfx950 path.
|
|
DSA_INDEXER_KV_DTYPES = ("fp8", "mxfp4")
|
|
|
|
|
|
def dsa_indexer_uses_fp4(vllm_config: VllmConfig) -> bool:
|
|
"""Whether the DeepSeek sparse indexer should use the MXFP4 K cache."""
|
|
kv_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype("fp8")
|
|
if kv_dtype not in DSA_INDEXER_KV_DTYPES:
|
|
raise ValueError(
|
|
f"indexer_kv_dtype={kv_dtype!r} is not supported by the DeepSeek "
|
|
f"sparse indexer (expected one of {DSA_INDEXER_KV_DTYPES})."
|
|
)
|
|
use_fp4 = kv_dtype == "mxfp4"
|
|
if use_fp4 and current_platform.is_rocm():
|
|
from vllm._aiter_ops import rocm_aiter_ops
|
|
from vllm.platforms.rocm import get_cdna_version
|
|
|
|
# Only DeepSeek-V4.1 is wired to the ROCm MXFP4 cache; other DSA models
|
|
# would silently keep their FP8 one.
|
|
model_config = vllm_config.model_config
|
|
if model_config is None or model_config.hf_config.model_type != "deepseek_v41":
|
|
raise ValueError(
|
|
"indexer_kv_dtype='mxfp4' on ROCm is only supported for "
|
|
"DeepSeek-V4.1-Flash."
|
|
)
|
|
if get_cdna_version() != 4:
|
|
raise ValueError(
|
|
"indexer_kv_dtype='mxfp4' on ROCm requires CDNA4 (MI350X/MI355X)."
|
|
)
|
|
if not rocm_aiter_ops.is_enabled():
|
|
raise ValueError(
|
|
"indexer_kv_dtype='mxfp4' on ROCm runs on aiter's kernels; enable "
|
|
"aiter with VLLM_ROCM_USE_AITER=1."
|
|
)
|
|
return True
|
|
if use_fp4 and not current_platform.is_device_capability_family(100):
|
|
raise ValueError(
|
|
"indexer_kv_dtype='mxfp4' requires Blackwell datacenter GPUs "
|
|
"(sm_10x, e.g. B200/GB200); sm_120 (consumer Blackwell) and "
|
|
"earlier architectures are not supported."
|
|
)
|
|
return use_fp4
|
|
|
|
|
|
class PrepareUniformDecodeKernel(
|
|
VllmTritonJitKernel["PrepareUniformDecodeKernel.CompileKey"]
|
|
):
|
|
BLOCK_SIZE = 1024
|
|
|
|
@dataclass(frozen=True)
|
|
class CompileKey:
|
|
block_size: int
|
|
block_table_stride: int
|
|
expanded_bt_stride: int
|
|
max_decode_len: int
|
|
|
|
@staticmethod
|
|
@triton.jit
|
|
def kernel(
|
|
seq_lens_ptr,
|
|
decode_seq_lens_ptr,
|
|
block_table_ptr,
|
|
block_table_stride,
|
|
expanded_block_table_ptr,
|
|
expanded_bt_stride,
|
|
decode_lens_ptr,
|
|
max_decode_len,
|
|
BLOCK_SIZE: tl.constexpr,
|
|
):
|
|
idx = tl.program_id(0)
|
|
req_id = idx // max_decode_len
|
|
local_idx = idx % max_decode_len
|
|
|
|
# Compute number of KVs attended to by this token. Padding requests have
|
|
# seq_len == 0, which would otherwise make the first token of each padded
|
|
# request negative (e.g. next_n=2 gives 0-2+0+1 = -1). Downstream kernels
|
|
# read these as uint32, turning -1 into ~4e9.
|
|
seq_len = tl.load(seq_lens_ptr + req_id)
|
|
per_token_seq_len = tl.maximum(seq_len - max_decode_len + local_idx + 1, 0)
|
|
tl.store(decode_seq_lens_ptr + idx, per_token_seq_len)
|
|
|
|
# Copy block table row.
|
|
src = block_table_ptr + req_id * block_table_stride
|
|
dst = expanded_block_table_ptr + idx * expanded_bt_stride
|
|
for i in tl.range(0, expanded_bt_stride, BLOCK_SIZE):
|
|
off = i + tl.arange(0, BLOCK_SIZE)
|
|
mask = off < expanded_bt_stride
|
|
src_block = tl.load(src + off, mask=mask)
|
|
tl.store(dst + off, src_block, mask=mask)
|
|
|
|
# All reqs now have decode_len = 1.
|
|
tl.store(decode_lens_ptr + idx, 1)
|
|
|
|
def dispatch( # type: ignore[override]
|
|
self,
|
|
*,
|
|
block_size: int,
|
|
block_table_stride: int,
|
|
expanded_bt_stride: int,
|
|
max_decode_len: int,
|
|
) -> CompileKey:
|
|
return self.CompileKey(
|
|
block_size=block_size,
|
|
block_table_stride=triton_scalar_specialization_rep(block_table_stride),
|
|
expanded_bt_stride=triton_scalar_specialization_rep(expanded_bt_stride),
|
|
max_decode_len=triton_scalar_specialization_rep(max_decode_len),
|
|
)
|
|
|
|
def get_warmup_keys(self) -> list[CompileKey]:
|
|
return self._trace_dispatch(self.dispatch)(
|
|
block_size=self.BLOCK_SIZE,
|
|
block_table_stride=(1, 2, 16),
|
|
expanded_bt_stride=(1, 2, 16),
|
|
max_decode_len=(1, 2, 16),
|
|
)
|
|
|
|
def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
|
|
int32_ptr = TritonWarmupTensor(torch.int32)
|
|
return dict(
|
|
seq_lens=int32_ptr,
|
|
decode_seq_lens=int32_ptr,
|
|
block_table=TritonWarmupTensor(
|
|
torch.int32,
|
|
shape=(1, compile_key.block_table_stride),
|
|
),
|
|
expanded_block_table=TritonWarmupTensor(
|
|
torch.int32,
|
|
shape=(1, compile_key.expanded_bt_stride),
|
|
),
|
|
decode_lens=int32_ptr,
|
|
num_decode_tokens=1,
|
|
max_decode_len=compile_key.max_decode_len,
|
|
)
|
|
|
|
@kernel_launcher
|
|
def __call__(
|
|
self,
|
|
seq_lens: torch.Tensor,
|
|
decode_seq_lens: torch.Tensor,
|
|
block_table: torch.Tensor,
|
|
expanded_block_table: torch.Tensor,
|
|
decode_lens: torch.Tensor,
|
|
num_decode_tokens: int,
|
|
max_decode_len: int,
|
|
) -> LaunchSpec:
|
|
return (num_decode_tokens,), dict(
|
|
block_table_stride=block_table.stride(0),
|
|
expanded_bt_stride=expanded_block_table.stride(0),
|
|
BLOCK_SIZE=self.BLOCK_SIZE,
|
|
)
|
|
|
|
|
|
# Below this many rows per rank, stay replicated (TP4 engages at 64K rows).
|
|
MIN_TP_SHARD_ROWS_PER_RANK = 16_384
|
|
|
|
|
|
def balanced_prefill_row_shard(
|
|
seq_lens_cpu: torch.Tensor,
|
|
query_lens_cpu: torch.Tensor,
|
|
compress_ratio: int,
|
|
tp_size: int,
|
|
) -> list[int] | None:
|
|
"""Return contiguous TP row counts balanced by indexer MQA work."""
|
|
num_rows = int(query_lens_cpu.sum())
|
|
if tp_size < 2 or num_rows < MIN_TP_SHARD_ROWS_PER_RANK * tp_size:
|
|
return None
|
|
|
|
query_lens = query_lens_cpu.to(torch.int64)
|
|
first_key = torch.repeat_interleave(
|
|
seq_lens_cpu.to(torch.int64) - query_lens + 1, query_lens
|
|
)
|
|
row_in_request = torch.arange(num_rows) - torch.repeat_interleave(
|
|
torch.cumsum(query_lens, 0) - query_lens, query_lens
|
|
)
|
|
cost = torch.cumsum((first_key + row_in_request) // compress_ratio, 0)
|
|
total = int(cost[-1]) if cost.numel() else 0
|
|
if total <= 0:
|
|
return None
|
|
|
|
targets = torch.arange(1, tp_size) * total // tp_size
|
|
bounds = [0, *torch.searchsorted(cost, targets).tolist(), num_rows]
|
|
# Keep >= 1 row per rank (the floor above leaves ample room); this also
|
|
# handles repeated costs from pool compression.
|
|
for i in range(1, tp_size):
|
|
bounds[i] = max(bounds[i], bounds[i - 1] + 1)
|
|
for i in range(tp_size - 1, 0, -1):
|
|
bounds[i] = min(bounds[i], bounds[i + 1] - 1)
|
|
return [bounds[i + 1] - bounds[i] for i in range(tp_size)]
|
|
|
|
|
|
def tp_prefill_row_sharding_supported(
|
|
vllm_config: VllmConfig,
|
|
dcp_world_size: int,
|
|
use_pcp: bool,
|
|
tp_size: int,
|
|
) -> bool:
|
|
"""Whether replicated prefill rows may be sharded across TP ranks."""
|
|
cudagraph_mode = vllm_config.compilation_config.cudagraph_mode or CUDAGraphMode.NONE
|
|
return (
|
|
current_platform.is_cuda()
|
|
and dcp_world_size == 1
|
|
and not use_pcp
|
|
and tp_size > 1
|
|
and not envs.VLLM_DISABLE_PYNCCL
|
|
and not envs.VLLM_USE_NCCL_SYMM_MEM
|
|
and not envs.VLLM_BATCH_INVARIANT
|
|
and cudagraph_mode.mixed_mode() != CUDAGraphMode.FULL
|
|
)
|
|
|
|
|
|
class DeepseekV32IndexerBackend(AttentionBackend):
|
|
@classmethod
|
|
def supports_pcp(cls) -> bool:
|
|
return True
|
|
|
|
@classmethod
|
|
def supports_device_cpu_query_lens_mismatch(cls) -> bool:
|
|
# Only the varlen paged MQA logits kernel takes per-request query
|
|
# lengths from device tensors natively. Hopper can instead flatten each
|
|
# query into a single-token row using device-built metadata.
|
|
return _supports_varlen_paged_mqa_logits() or (
|
|
_supports_flattened_device_query_lens()
|
|
)
|
|
|
|
@staticmethod
|
|
def get_name() -> str:
|
|
return "DEEPSEEK_V32_INDEXER"
|
|
|
|
@staticmethod
|
|
def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
|
|
return [1, MultipleOf(16)] if current_platform.is_rocm() else [64]
|
|
|
|
@classmethod
|
|
def get_supported_head_sizes(cls) -> list[int]:
|
|
return [32, 64, 128]
|
|
|
|
@staticmethod
|
|
def get_builder_cls() -> type["DeepseekV32IndexerMetadataBuilder"]:
|
|
return DeepseekV32IndexerMetadataBuilder
|
|
|
|
|
|
class Glm5NextIndexerBackend(DeepseekV32IndexerBackend):
|
|
@staticmethod
|
|
def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
|
|
if kv_cache_spec is not None:
|
|
# Worker get_kv_cache_spec: runs without the current config.
|
|
index_kpool = int(kv_cache_spec.tokens_per_state)
|
|
else:
|
|
# Platform block-size selection: no spec yet, under the current config.
|
|
model_config = get_current_vllm_config().model_config
|
|
index_kpool = model_config.hf_text_config.index_kpool
|
|
return [index_kpool * n for n in get_paged_mqa_page_sizes()]
|
|
|
|
|
|
class KpoolTailBackend(DeepseekV32IndexerBackend):
|
|
"""Storage-only backend for the GLM-5.3-Flash kpool tail cache."""
|
|
|
|
@classmethod
|
|
def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:
|
|
return (KVCacheLayout.BLHNC,)
|
|
|
|
@staticmethod
|
|
def get_name() -> str:
|
|
return "KPOOL_TAIL"
|
|
|
|
@classmethod
|
|
def get_supported_head_sizes(cls) -> list[int]:
|
|
return []
|
|
|
|
@staticmethod
|
|
def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
|
|
return [MultipleOf(1)]
|
|
|
|
@staticmethod
|
|
def get_builder_cls() -> type["KpoolTailMetadataBuilder"]: # type: ignore[override]
|
|
return KpoolTailMetadataBuilder
|
|
|
|
|
|
class DeepseekV4IndexerBackend(DeepseekV32IndexerBackend):
|
|
@classmethod
|
|
def supports_device_cpu_query_lens_mismatch(cls) -> bool:
|
|
# ROCm runs adaptive verification through the per-token flattened
|
|
# indexer path, which derives row ownership from device decode lengths.
|
|
return (
|
|
current_platform.is_rocm()
|
|
or super().supports_device_cpu_query_lens_mismatch()
|
|
)
|
|
|
|
@staticmethod
|
|
def get_name() -> str:
|
|
return "DEEPSEEK_V4_INDEXER"
|
|
|
|
@classmethod
|
|
def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:
|
|
# DeepSeek-V4 packs the indexer pages beside the MLA latent pages inside
|
|
# each block, so the layer dim must sit inside the block dim.
|
|
return (KVCacheLayout.BLHNC, KVCacheLayout.BLNHC)
|
|
|
|
@staticmethod
|
|
def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
|
|
# Block sizes count uncompressed tokens: C4 indexer pages hold 64 rows.
|
|
return [256]
|
|
|
|
|
|
class DeepseekV41IndexerBackend(DeepseekV4IndexerBackend):
|
|
@classmethod
|
|
def supports_device_cpu_query_lens_mismatch(cls) -> bool:
|
|
# The ROCm flattened-query support above is validated for the
|
|
# DeepSeek-V4 adaptive DSpark path only.
|
|
return DeepseekV32IndexerBackend.supports_device_cpu_query_lens_mismatch()
|
|
|
|
@staticmethod
|
|
def get_name() -> str:
|
|
return "DEEPSEEK_V41_INDEXER"
|
|
|
|
@staticmethod
|
|
def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
|
|
return [64 if current_platform.is_device_capability_family(90) else 128]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class PCPGlobalChunkPlan:
|
|
"""PCP packing for one indexer prefill chunk under PCP + DCP."""
|
|
|
|
# [num_reqs+1] cumsum of the padded per-request context (shard rows x W):
|
|
# the global layout.
|
|
row_start_cu: torch.Tensor
|
|
# [num_reqs+1] like row_start_cu, but only a region's FIRST row carries its
|
|
# extent.
|
|
global_cu: torch.Tensor
|
|
# [num_reqs+1] cumsum of shard rows: this rank's padded layout.
|
|
padded_local_cu: torch.Tensor
|
|
padded_local_total: int
|
|
total: int
|
|
# [total] for each global row, its position in the rank-major gathered buffer.
|
|
deinterleave_idx: torch.Tensor
|
|
|
|
|
|
def build_pcp_global_chunk_plan(
|
|
row_req_idx: np.ndarray,
|
|
row_shard_rows: np.ndarray,
|
|
dcp_world_size: int,
|
|
device: torch.device,
|
|
interleave: int = 1,
|
|
) -> PCPGlobalChunkPlan:
|
|
"""Plan the PCP packing for one chunk from its requests' KV shard rows.
|
|
|
|
The global layout is padded to ``shard rows x W`` per request. Positions
|
|
past a request's true extent lie beyond every token's causal bound, so
|
|
the kernel never reads them.
|
|
"""
|
|
shard_rows = np.ascontiguousarray(row_shard_rows, dtype=np.int64)
|
|
assert row_req_idx.shape == shard_rows.shape
|
|
num_rows = len(shard_rows)
|
|
|
|
row_bounds = request_row_bounds(row_req_idx)
|
|
region_first_row = row_bounds[:-1]
|
|
region_of_row = np.repeat(np.arange(len(region_first_row)), np.diff(row_bounds))
|
|
|
|
region_padded = shard_rows[region_first_row]
|
|
assert np.all(region_padded > 0), (
|
|
f"PCP+DCP prefill got an empty context: {region_padded.tolist()}"
|
|
)
|
|
|
|
# Pinned host buffers, filled in place so the copies below are async.
|
|
cu = torch.zeros((3, num_rows + 1), dtype=torch.int32, pin_memory=PIN_MEMORY)
|
|
row_start_cu, global_cu, padded_cu = cu
|
|
# Only a region's first row carries its extent.
|
|
first_row = torch.from_numpy(region_first_row)
|
|
row_padded = torch.zeros(num_rows, dtype=torch.int32)
|
|
row_padded[first_row] = torch.from_numpy(region_padded).int()
|
|
torch.cumsum(row_padded * dcp_world_size, 0, out=global_cu[1:])
|
|
torch.cumsum(row_padded, 0, out=padded_cu[1:])
|
|
row_start_cu[:num_rows] = global_cu[first_row][torch.from_numpy(region_of_row)]
|
|
row_start_cu[num_rows] = global_cu[num_rows]
|
|
total = int(global_cu[num_rows])
|
|
padded_total = int(padded_cu[num_rows])
|
|
|
|
idx = torch.empty(total, dtype=torch.int64, pin_memory=PIN_MEMORY)
|
|
idx_np = idx.numpy()
|
|
region_start = global_cu[first_row].numpy()
|
|
region_padded_start = padded_cu[first_row].numpy()
|
|
for i in range(len(region_first_row)):
|
|
g = int(region_padded[i]) * dcp_world_size
|
|
t = np.arange(g, dtype=np.int64)
|
|
local = round_down(t // dcp_world_size, interleave) + t % interleave
|
|
# Padded positions are outside causal bounds but must remain in-bounds.
|
|
local = np.minimum(local, region_padded[i] - 1)
|
|
idx_np[region_start[i] : region_start[i] + g] = (
|
|
((t // interleave) % dcp_world_size) * padded_total
|
|
+ region_padded_start[i]
|
|
+ local
|
|
)
|
|
|
|
cu = async_tensor_h2d(cu, device=device)
|
|
return PCPGlobalChunkPlan(
|
|
row_start_cu=cu[0],
|
|
global_cu=cu[1],
|
|
padded_local_cu=cu[2],
|
|
padded_local_total=padded_total,
|
|
total=total,
|
|
deinterleave_idx=async_tensor_h2d(idx, device=device),
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class DeepseekV32IndexerPrefillChunkMetadata:
|
|
block_table: torch.Tensor
|
|
# Under DCP (dcp_world_size > 1) these hold this rank's local row bounds;
|
|
# otherwise they hold the global bounds.
|
|
cu_seqlen_ks: torch.Tensor
|
|
cu_seqlen_ke: torch.Tensor
|
|
cu_seq_lens: torch.Tensor
|
|
token_to_seq: torch.Tensor
|
|
total_seq_lens: int
|
|
# Host upper bound on cu_seqlen_ke - cu_seqlen_ks, sizing the compressed
|
|
# DeepGEMM prefill logits.
|
|
max_seqlen_k: int
|
|
token_start: int
|
|
token_end: int
|
|
num_reqs: int
|
|
skip_kv_gather: bool = False
|
|
local_cu_seq_lens: torch.Tensor | None = None
|
|
local_total_seq_lens: int = 0
|
|
max_local_total_seq_lens: int = 0
|
|
|
|
pcp_deinterleave_idx: torch.Tensor | None = None
|
|
|
|
|
|
class BuildPrefillChunkMetadataKernel(
|
|
VllmTritonJitKernel["BuildPrefillChunkMetadataKernel.CompileKey"]
|
|
):
|
|
BLOCK_SIZE = 1024
|
|
|
|
@dataclass(frozen=True)
|
|
class CompileKey:
|
|
query_slice_start: int
|
|
query_slice_stop: int
|
|
dcp_rank: int
|
|
dcp_world: int
|
|
dcp_interleave: int
|
|
block_size: int
|
|
compress_ratio: int
|
|
input_variant: TritonPointerInputVariant
|
|
|
|
@staticmethod
|
|
@triton.jit
|
|
def kernel(
|
|
# Inputs
|
|
query_start_loc_ptr,
|
|
uncompressed_seq_lens_ptr,
|
|
cu_compressed_seq_lens_ptr,
|
|
# Row-start base for cu_seq_len_ks/ke: local cumulative lens under DCP,
|
|
# aliases cu_compressed_seq_lens_ptr otherwise.
|
|
row_start_cu_compressed_seq_lens_ptr,
|
|
# Outputs
|
|
token_to_seq_ptr,
|
|
cu_compressed_seq_len_ks_ptr,
|
|
cu_compressed_seq_len_ke_ptr,
|
|
query_slice_start,
|
|
query_slice_stop,
|
|
DCP_RANK,
|
|
DCP_WORLD,
|
|
DCP_INTERLEAVE,
|
|
BLOCK_SIZE: tl.constexpr,
|
|
COMPRESS_RATIO: tl.constexpr,
|
|
):
|
|
batch_idx = tl.program_id(0)
|
|
|
|
query_start = tl.load(query_start_loc_ptr + batch_idx)
|
|
query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
|
|
query_len = query_end - query_start
|
|
|
|
seq_start = tl.load(cu_compressed_seq_lens_ptr + batch_idx)
|
|
seq_end = tl.load(cu_compressed_seq_lens_ptr + batch_idx + 1)
|
|
compressed_seq_len = seq_end - seq_start
|
|
|
|
# Row start for the (possibly localized) cu_seq_len_ks/ke. Equals seq_start
|
|
# when DCP is disabled (the pointer aliases cu_compressed_seq_lens_ptr).
|
|
row_start = tl.load(row_start_cu_compressed_seq_lens_ptr + batch_idx)
|
|
|
|
uncompressed_seq_len = tl.load(uncompressed_seq_lens_ptr + batch_idx)
|
|
start_pos = uncompressed_seq_len - query_len
|
|
|
|
for i in range(0, query_len, BLOCK_SIZE):
|
|
offset = i + tl.arange(0, BLOCK_SIZE)
|
|
abs_pos = query_start + offset
|
|
mask = (
|
|
(offset < query_len)
|
|
& (abs_pos >= query_slice_start)
|
|
& (abs_pos < query_slice_stop)
|
|
)
|
|
out_pos = abs_pos - query_slice_start
|
|
|
|
# cu_seq_len_ks: row start in the gathered K buffer.
|
|
tl.store(cu_compressed_seq_len_ks_ptr + out_pos, row_start, mask=mask)
|
|
|
|
# cu_seq_len_ke: row start + per-token context length. Under DCP the
|
|
# global per-token length is sharded across ranks.
|
|
global_ctx = start_pos + 1 + offset
|
|
len_per_token = global_ctx // COMPRESS_RATIO
|
|
if DCP_WORLD > 1:
|
|
# Per-rank local context length under interleave-aware DCP, matching
|
|
# get_dcp_local_seq_lens. K == 1 reduces to (len + world-1-rank)//world.
|
|
base = (len_per_token // DCP_INTERLEAVE // DCP_WORLD) * DCP_INTERLEAVE
|
|
remainder = len_per_token - base * DCP_WORLD
|
|
remainder = tl.minimum(
|
|
tl.maximum(remainder - DCP_RANK * DCP_INTERLEAVE, 0), DCP_INTERLEAVE
|
|
)
|
|
len_per_token = base + remainder
|
|
tl.store(
|
|
cu_compressed_seq_len_ke_ptr + out_pos,
|
|
row_start + len_per_token,
|
|
mask=mask,
|
|
)
|
|
|
|
# Compute token_to_seq
|
|
for i in range(0, compressed_seq_len, BLOCK_SIZE):
|
|
offset = i + tl.arange(0, BLOCK_SIZE)
|
|
mask = offset < compressed_seq_len
|
|
tl.store(token_to_seq_ptr + seq_start + offset, batch_idx, mask=mask)
|
|
|
|
def dispatch( # type: ignore[override]
|
|
self,
|
|
*,
|
|
query_slice_start: int,
|
|
query_slice_stop: int,
|
|
DCP_RANK: int,
|
|
DCP_WORLD: int,
|
|
DCP_INTERLEAVE: int,
|
|
BLOCK_SIZE: int,
|
|
COMPRESS_RATIO: int,
|
|
input_variant: TritonPointerInputVariant,
|
|
) -> CompileKey:
|
|
return self.CompileKey(
|
|
query_slice_start=triton_scalar_specialization_rep(query_slice_start),
|
|
query_slice_stop=triton_scalar_specialization_rep(query_slice_stop),
|
|
dcp_rank=triton_scalar_specialization_rep(DCP_RANK),
|
|
dcp_world=triton_scalar_specialization_rep(DCP_WORLD),
|
|
dcp_interleave=triton_scalar_specialization_rep(DCP_INTERLEAVE),
|
|
block_size=BLOCK_SIZE,
|
|
compress_ratio=COMPRESS_RATIO,
|
|
input_variant=input_variant,
|
|
)
|
|
|
|
def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]:
|
|
max_tokens = max(1, min(vllm_config.scheduler_config.max_num_batched_tokens, 8))
|
|
hf_text_config = vllm_config.model_config.hf_text_config
|
|
parallel_config = vllm_config.parallel_config
|
|
dcp_world = parallel_config.decode_context_parallel_size
|
|
dcp_interleave = parallel_config.cp_kv_cache_interleave_size
|
|
dcp_rank = get_dcp_group().rank_in_group if dcp_world > 1 else 0
|
|
dcp_cases = [dict(DCP_RANK=dcp_rank, DCP_WORLD=dcp_world)]
|
|
if parallel_config.prefill_context_parallel_size > 1 and dcp_world > 1:
|
|
# PCP+DCP prefill constructs a global PCP chunk plan and therefore
|
|
# dispatches this kernel with already-localized row starts. The
|
|
# runtime key is normalized to rank 0/world 1 in that path; warm it
|
|
# explicitly so one PCP rank cannot enter a collective while
|
|
# another rank is still compiling the first real inference batch.
|
|
dcp_cases.append(dict(DCP_RANK=0, DCP_WORLD=1))
|
|
compress_ratios = tuple(
|
|
dict.fromkeys(
|
|
max(1, int(ratio))
|
|
for ratio in (
|
|
*(getattr(hf_text_config, "compress_ratios", None) or (1,)),
|
|
getattr(hf_text_config, "index_kpool", 1) or 1,
|
|
)
|
|
)
|
|
)
|
|
index_kpool = getattr(hf_text_config, "index_kpool", None)
|
|
if index_kpool and index_kpool > 1 and index_kpool not in compress_ratios:
|
|
compress_ratios = compress_ratios + (index_kpool,)
|
|
return self._trace_dispatch(self.dispatch)(
|
|
zip_inputs(*dcp_cases),
|
|
# Cover Triton's divisible, exact-one, and generic i32 classes.
|
|
query_slice_start=(0, 1, 2),
|
|
query_slice_stop=(1, 2 * max_tokens - 1, 2 * max_tokens),
|
|
# Compressed indexers localize in units of compressed states.
|
|
DCP_INTERLEAVE=list(
|
|
dict.fromkeys(
|
|
dcp_interleave // ratio
|
|
for ratio in compress_ratios
|
|
if dcp_interleave % ratio == 0
|
|
)
|
|
)
|
|
or [dcp_interleave],
|
|
BLOCK_SIZE=self.BLOCK_SIZE,
|
|
COMPRESS_RATIO=list(compress_ratios),
|
|
# PCP's global cumulative lengths are the second row of one packed
|
|
# allocation, so their pointer is not always 16-byte aligned.
|
|
# Prefill chunking can independently unalign uncompressed lengths.
|
|
input_variant=(
|
|
TritonPointerInputVariant.from_alignment(
|
|
uncompressed_seq_lens=True,
|
|
cu_compressed_seq_lens=True,
|
|
),
|
|
TritonPointerInputVariant.from_alignment(
|
|
uncompressed_seq_lens=True,
|
|
cu_compressed_seq_lens=False,
|
|
),
|
|
TritonPointerInputVariant.from_alignment(
|
|
uncompressed_seq_lens=False,
|
|
cu_compressed_seq_lens=True,
|
|
),
|
|
TritonPointerInputVariant.from_alignment(
|
|
uncompressed_seq_lens=False,
|
|
cu_compressed_seq_lens=False,
|
|
),
|
|
),
|
|
)
|
|
|
|
def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
|
|
int32_ptr = TritonWarmupTensor(torch.int32)
|
|
return dict(
|
|
query_start_loc=int32_ptr,
|
|
uncompressed_seq_lens=compile_key.input_variant.pointer(
|
|
"uncompressed_seq_lens", torch.int32
|
|
),
|
|
cu_compressed_seq_lens=compile_key.input_variant.pointer(
|
|
"cu_compressed_seq_lens", torch.int32
|
|
),
|
|
row_start_cu_compressed_seq_lens=int32_ptr,
|
|
token_to_seq=int32_ptr,
|
|
cu_compressed_seq_len_ks=int32_ptr,
|
|
cu_compressed_seq_len_ke=int32_ptr,
|
|
query_slice_start=compile_key.query_slice_start,
|
|
query_slice_stop=compile_key.query_slice_stop,
|
|
DCP_RANK=compile_key.dcp_rank,
|
|
DCP_WORLD=compile_key.dcp_world,
|
|
DCP_INTERLEAVE=compile_key.dcp_interleave,
|
|
num_reqs=1,
|
|
COMPRESS_RATIO=compile_key.compress_ratio,
|
|
)
|
|
|
|
@kernel_launcher
|
|
def __call__(
|
|
self,
|
|
query_start_loc: torch.Tensor,
|
|
uncompressed_seq_lens: torch.Tensor,
|
|
cu_compressed_seq_lens: torch.Tensor,
|
|
row_start_cu_compressed_seq_lens: torch.Tensor,
|
|
token_to_seq: torch.Tensor,
|
|
cu_compressed_seq_len_ks: torch.Tensor,
|
|
cu_compressed_seq_len_ke: torch.Tensor,
|
|
query_slice_start: int,
|
|
query_slice_stop: int,
|
|
DCP_RANK: int,
|
|
DCP_WORLD: int,
|
|
DCP_INTERLEAVE: int,
|
|
*,
|
|
num_reqs: int,
|
|
COMPRESS_RATIO: int,
|
|
) -> LaunchSpec:
|
|
return (num_reqs,), dict(
|
|
BLOCK_SIZE=self.BLOCK_SIZE,
|
|
COMPRESS_RATIO=COMPRESS_RATIO,
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class DeepseekV32IndexerPrefillMetadata:
|
|
chunks: list[DeepseekV32IndexerPrefillChunkMetadata]
|
|
max_prefill_seq_len: int = -1
|
|
# Per-TP-rank contiguous row counts, or None for the replicated path.
|
|
row_shard_sizes: list[int] | None = None
|
|
|
|
|
|
@dataclass
|
|
class DeepSeekV32IndexerDecodeMetadata:
|
|
block_table: torch.Tensor
|
|
# seq_lens: per-token effective context lengths.
|
|
# - flatten path / plain decode: 1D (batch_size,)
|
|
# - native MTP path: 2D (B, next_n) where [b,j] = L_b - next_n + j + 1
|
|
# Both fp8_fp4_paged_mqa_logits and the topk kernels accept both shapes.
|
|
seq_lens: torch.Tensor
|
|
decode_lens: torch.Tensor
|
|
requires_padding: bool
|
|
schedule_metadata: torch.Tensor
|
|
global_seq_lens: torch.Tensor | None = None
|
|
per_req_decode_lens: torch.Tensor | None = None
|
|
decode_is_uniform: bool = True
|
|
write_max_decode_len: int = 0
|
|
indices: torch.Tensor | None = None
|
|
|
|
|
|
@dataclass
|
|
class DeepseekV32IndexerMetadata:
|
|
# FIXME (zyongye)
|
|
# hacky way to access the data now, need to be in chunked meta
|
|
seq_lens: torch.Tensor
|
|
max_seq_len: int
|
|
slot_mapping: torch.Tensor
|
|
|
|
# New for MLA (compared to FlashAttention)
|
|
# For handling prefill decode split
|
|
num_decodes: int
|
|
num_decode_tokens: int
|
|
num_prefills: int
|
|
num_prefill_tokens: int
|
|
|
|
decode: DeepSeekV32IndexerDecodeMetadata | None = None
|
|
prefill: DeepseekV32IndexerPrefillMetadata | None = None
|
|
positions: torch.Tensor | None = None
|
|
block_table: torch.Tensor | None = None
|
|
|
|
|
|
@triton.jit(do_not_specialize=["num_reqs", "num_actual_tokens", "num_tokens"])
|
|
def _kpool_tail_slot_mapping_kernel(
|
|
slot_mapping_ptr,
|
|
block_table_ptr,
|
|
block_table_stride,
|
|
query_start_loc_ptr,
|
|
positions_ptr,
|
|
out_ptr,
|
|
num_reqs,
|
|
num_actual_tokens,
|
|
num_tokens,
|
|
kpool,
|
|
BLOCK: tl.constexpr,
|
|
):
|
|
pid = tl.program_id(0)
|
|
if pid < num_reqs:
|
|
start = tl.load(query_start_loc_ptr + pid)
|
|
# Tokens past the last request's boundary (if any) also map to it.
|
|
end = tl.load(query_start_loc_ptr + pid + 1)
|
|
end = tl.where(pid == num_reqs - 1, num_actual_tokens, end)
|
|
own_block = tl.load(block_table_ptr + pid * block_table_stride).to(tl.int64)
|
|
# Every scheduled request owns a tail block; the runners give dummy and
|
|
# padding requests the null block, which must never be written.
|
|
has_block = own_block != 0
|
|
for i in range(start, end, BLOCK):
|
|
offs = i + tl.arange(0, BLOCK)
|
|
mask = offs < end
|
|
pos = tl.load(positions_ptr + offs, mask=mask, other=0).to(tl.int64)
|
|
slot = tl.where(has_block, own_block * kpool + pos % kpool, -1)
|
|
tl.store(out_ptr + offs, slot, mask=mask)
|
|
else:
|
|
offs = num_actual_tokens + (pid - num_reqs) * BLOCK + tl.arange(0, BLOCK)
|
|
mask = offs < num_tokens
|
|
tl.store(out_ptr + offs, tl.load(slot_mapping_ptr + offs, mask=mask), mask=mask)
|
|
|
|
|
|
def compute_kpool_tail_slot_mapping(
|
|
slot_mapping: torch.Tensor,
|
|
block_table: torch.Tensor,
|
|
query_start_loc: torch.Tensor,
|
|
positions: torch.Tensor,
|
|
num_actual_tokens: int,
|
|
num_reqs: int,
|
|
kpool: int,
|
|
out: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
"""Map every token to its request's one circular tail block."""
|
|
if out is None:
|
|
out = torch.empty_like(slot_mapping)
|
|
else:
|
|
assert out.shape == slot_mapping.shape
|
|
if slot_mapping.is_cuda and slot_mapping.dim() == 1 and num_reqs > 0:
|
|
block = 256
|
|
num_tokens = slot_mapping.shape[0]
|
|
num_actual_tokens = min(num_actual_tokens, num_tokens)
|
|
grid = (num_reqs + triton.cdiv(num_tokens - num_actual_tokens, block),)
|
|
_kpool_tail_slot_mapping_kernel[grid](
|
|
slot_mapping,
|
|
block_table,
|
|
block_table.stride(0),
|
|
query_start_loc,
|
|
positions,
|
|
out,
|
|
num_reqs,
|
|
num_actual_tokens,
|
|
num_tokens,
|
|
kpool,
|
|
BLOCK=block,
|
|
num_warps=4,
|
|
)
|
|
return out
|
|
# Torch fallback: CPU tensors (the CPU unit tests), non-1D or empty inputs.
|
|
# Production always takes the Triton path above — spec-decode tokens arrive
|
|
# flattened token-major, so slot_mapping is 1D there too.
|
|
out.copy_(slot_mapping)
|
|
if num_actual_tokens == 0:
|
|
return out
|
|
tokens = torch.arange(num_actual_tokens, device=slot_mapping.device)
|
|
req = torch.searchsorted(query_start_loc, tokens, right=True) - 1
|
|
req = req.clamp_(min=0, max=num_reqs - 1)
|
|
own_block = block_table[:num_reqs, 0].index_select(0, req).to(torch.int64)
|
|
pos = positions[:num_actual_tokens].to(torch.int64)
|
|
slots = own_block * kpool + torch.remainder(pos, kpool)
|
|
out[:num_actual_tokens] = torch.where(own_block != 0, slots, -1)
|
|
return out
|
|
|
|
|
|
class KpoolTailMetadataBuilder(AttentionMetadataBuilder):
|
|
"""Build only the circular slot mapping needed by the storage-only tail."""
|
|
|
|
_cudagraph_support = AttentionCGSupport.ALWAYS
|
|
supports_update_block_table = False
|
|
supports_draft_decode_metadata_update = True
|
|
reorder_batch_threshold = None
|
|
|
|
def __init__(
|
|
self,
|
|
kv_cache_spec: AttentionSpec,
|
|
layer_names: list[str],
|
|
vllm_config: VllmConfig,
|
|
device: torch.device,
|
|
):
|
|
super().__init__(kv_cache_spec, layer_names, vllm_config, device)
|
|
self.slot_mapping_buffer = torch.empty(
|
|
vllm_config.scheduler_config.max_num_batched_tokens,
|
|
dtype=torch.int64,
|
|
device=device,
|
|
)
|
|
|
|
def build(
|
|
self,
|
|
common_prefix_len: int,
|
|
common_attn_metadata: CommonAttentionMetadata,
|
|
fast_build: bool = False,
|
|
) -> DeepseekV32IndexerMetadata:
|
|
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
|
|
split_decodes_and_prefills(common_attn_metadata)
|
|
)
|
|
slot_mapping = common_attn_metadata.slot_mapping
|
|
positions = common_attn_metadata.positions
|
|
if positions is not None:
|
|
slot_mapping_buffer = self.slot_mapping_buffer[
|
|
: slot_mapping.numel()
|
|
].view_as(slot_mapping)
|
|
slot_mapping = compute_kpool_tail_slot_mapping(
|
|
slot_mapping,
|
|
common_attn_metadata.block_table_tensor,
|
|
common_attn_metadata.query_start_loc,
|
|
positions,
|
|
common_attn_metadata.num_actual_tokens,
|
|
common_attn_metadata.num_reqs,
|
|
self.kv_cache_spec.block_size,
|
|
out=slot_mapping_buffer,
|
|
)
|
|
return DeepseekV32IndexerMetadata(
|
|
seq_lens=common_attn_metadata.seq_lens,
|
|
max_seq_len=common_attn_metadata.max_seq_len,
|
|
slot_mapping=slot_mapping,
|
|
num_decodes=num_decodes,
|
|
num_decode_tokens=num_decode_tokens,
|
|
num_prefills=num_prefills,
|
|
num_prefill_tokens=num_prefill_tokens,
|
|
positions=positions,
|
|
)
|
|
|
|
def update_draft_decode_metadata(
|
|
self,
|
|
metadata: DeepseekV32IndexerMetadata,
|
|
) -> None:
|
|
assert metadata.positions is not None
|
|
slot_mapping = metadata.slot_mapping[: metadata.num_decode_tokens]
|
|
slot_mapping.div_(self.kv_cache_spec.block_size, rounding_mode="floor")
|
|
slot_mapping.mul_(self.kv_cache_spec.block_size)
|
|
slot_mapping.add_(
|
|
metadata.positions[: metadata.num_decode_tokens].remainder(
|
|
self.kv_cache_spec.block_size
|
|
)
|
|
)
|
|
|
|
|
|
def get_max_prefill_buffer_size(vllm_config: VllmConfig):
|
|
max_model_len = vllm_config.model_config.max_model_len
|
|
factor = vllm_config.attention_config.sparse_mla_prefill_workspace_factor
|
|
assert factor is not None
|
|
# Each entry is 128 fp8 bytes and 4 scale bytes for a total of 132 bytes.
|
|
# The flashmla_sparse backend uses a workspace size of factor * max_model_len.
|
|
# The memory usage of the workspace there is 576 * 2 bytes; so we size this as
|
|
# (576 * 2 // 132) * factor (40 at factor 5) to maximize this workspace
|
|
# size while still fitting within the flashmla_sparse workspace.
|
|
# For DeepSeek-V3.2, the max_model_len is 163840.
|
|
# 40 * 163840 * 132 = 865075200 bytes = 825 MB
|
|
return max_model_len * (576 * 2 // 132) * factor
|
|
|
|
|
|
def _supports_varlen_paged_mqa_logits() -> bool:
|
|
return (
|
|
current_platform.is_cuda()
|
|
and current_platform.is_device_capability_family(100)
|
|
and has_deep_gemm()
|
|
)
|
|
|
|
|
|
def _supports_flattened_device_query_lens() -> bool:
|
|
return (
|
|
current_platform.is_cuda()
|
|
and current_platform.is_device_capability_family(90)
|
|
and has_deep_gemm()
|
|
)
|
|
|
|
|
|
def _rocm_supports_flattened_device_query_lens(vllm_config: VllmConfig) -> bool:
|
|
model_config = vllm_config.model_config
|
|
return (
|
|
current_platform.is_rocm()
|
|
and model_config is not None
|
|
and "DeepseekV4ForCausalLM" in model_config.architectures
|
|
)
|
|
|
|
|
|
def _supports_native_decode(next_n: int) -> bool:
|
|
"""Whether decode can pass `next_n` Q rows per request to the kernel
|
|
instead of flattening to one single-token row per query, which re-reads
|
|
the KV tile once per row.
|
|
"""
|
|
if not (current_platform.is_cuda() and has_deep_gemm()):
|
|
return next_n in (1, 2)
|
|
if current_platform.is_device_capability_family(100):
|
|
return True
|
|
if current_platform.is_device_capability_family(90):
|
|
return native_next_n_supported(next_n)
|
|
return next_n in (1, 2)
|
|
|
|
|
|
def _use_flattening(vllm_config: VllmConfig) -> bool:
|
|
speculative_config = vllm_config.speculative_config
|
|
next_n = 1 + vllm_config.num_speculative_tokens
|
|
return not _supports_native_decode(next_n) or (
|
|
speculative_config is not None
|
|
and speculative_config.enable_adaptive_verification
|
|
and (
|
|
_supports_flattened_device_query_lens()
|
|
or _rocm_supports_flattened_device_query_lens(vllm_config)
|
|
)
|
|
)
|
|
|
|
|
|
class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
|
|
# The indexer opts out of the shared reorder-threshold vote (see __init__),
|
|
# so this is None; its own split uses self.decode_threshold.
|
|
reorder_batch_threshold: int | None = None
|
|
requires_block_table_width = True
|
|
|
|
@classmethod
|
|
def get_cudagraph_support(
|
|
cls,
|
|
vllm_config: VllmConfig,
|
|
kv_cache_spec: KVCacheSpec,
|
|
) -> AttentionCGSupport:
|
|
if _supports_varlen_paged_mqa_logits() or _use_flattening(vllm_config):
|
|
return AttentionCGSupport.ALWAYS
|
|
return AttentionCGSupport.UNIFORM_BATCH
|
|
|
|
def __init__(self, *args, block_table_width: int, **kwargs) -> None:
|
|
super().__init__(*args, **kwargs)
|
|
scheduler_config = self.vllm_config.scheduler_config
|
|
parallel_config = self.vllm_config.parallel_config
|
|
self.dcp_world_size = parallel_config.decode_context_parallel_size
|
|
self.tp_size = parallel_config.tensor_parallel_size
|
|
self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0
|
|
self.pcp_world_size = parallel_config.prefill_context_parallel_size
|
|
self.use_pcp = self.pcp_world_size > 1
|
|
self.pcp_rank = get_pcp_group().rank_in_group if self.use_pcp else 0
|
|
self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size
|
|
# KV compression (DeepseekV4). Default to 1 for no compression.
|
|
self.compress_ratio = 1
|
|
if isinstance(self.kv_cache_spec, MLAAttentionSpec):
|
|
assert isinstance(self.kv_cache_spec.tokens_per_state, int)
|
|
self.compress_ratio = self.kv_cache_spec.tokens_per_state
|
|
# NOTE(Chen):an estimated max size of flattened_kv. Need to double check.
|
|
# Counted in compressed rows, like the chunker's seq_lens and the
|
|
# workspace.
|
|
self.max_prefill_buffer_size = (
|
|
get_max_prefill_buffer_size(self.vllm_config) // self.compress_ratio
|
|
)
|
|
self.num_speculative_tokens = (
|
|
self.vllm_config.speculative_config.num_speculative_tokens
|
|
if self.vllm_config.speculative_config
|
|
else 0
|
|
)
|
|
self.indexer_uses_fp4 = dsa_indexer_uses_fp4(self.vllm_config)
|
|
|
|
next_n = self.num_speculative_tokens + 1
|
|
self.decode_threshold = next_n
|
|
self.reorder_batch_threshold = None
|
|
self.use_flattening = _use_flattening(self.vllm_config)
|
|
self.supports_varlen = _supports_varlen_paged_mqa_logits()
|
|
logger.info_once(
|
|
"DSA indexer decode path: use_flattening=%s supports_varlen=%s "
|
|
"(next_n=%d, use_fp4_cache=%s)",
|
|
self.use_flattening,
|
|
self.supports_varlen,
|
|
next_n,
|
|
self.indexer_uses_fp4,
|
|
)
|
|
|
|
sm_count = num_compute_units(self.device.index)
|
|
self.num_sms = sm_count
|
|
|
|
self.offsets_buffer = torch.arange(
|
|
next_n, device=self.device, dtype=torch.int32
|
|
)
|
|
self.decode_lens_buffer = torch.zeros(
|
|
(scheduler_config.max_num_batched_tokens,),
|
|
dtype=torch.int32,
|
|
device=self.device,
|
|
)
|
|
self.per_req_decode_lens_buffer = torch.zeros(
|
|
(scheduler_config.max_num_batched_tokens,),
|
|
dtype=torch.int32,
|
|
device=self.device,
|
|
)
|
|
# Shared workspace for decode seq_lens. Native MTP views this as
|
|
# (B, max_decode_len) at runtime, keeping context_lens contiguous even
|
|
# when max_decode_len is smaller than next_n.
|
|
self.decode_seq_lens_buffer = torch.zeros(
|
|
(scheduler_config.max_num_batched_tokens,),
|
|
dtype=torch.int32,
|
|
device=self.device,
|
|
)
|
|
self.global_decode_seq_lens_buffer = torch.zeros(
|
|
(scheduler_config.max_num_batched_tokens,),
|
|
dtype=torch.int32,
|
|
device=self.device,
|
|
)
|
|
self.decode_indices_buffer = torch.zeros(
|
|
(scheduler_config.max_num_batched_tokens,),
|
|
dtype=torch.int32,
|
|
device=self.device,
|
|
)
|
|
self.arange_buffer = torch.arange(
|
|
max(
|
|
scheduler_config.max_num_seqs * next_n,
|
|
scheduler_config.max_num_batched_tokens,
|
|
),
|
|
dtype=torch.int32,
|
|
device=self.device,
|
|
)
|
|
# Materialize the rank on device during builder initialization. Creating
|
|
# this scalar in build() would introduce a GPU<->CPU sync in the decode
|
|
# hot path.
|
|
self.dcp_rank_tensor = torch.tensor(
|
|
self.dcp_rank, dtype=torch.int32, device=self.device
|
|
)
|
|
self.expanded_block_table_buffer = torch.zeros(
|
|
(scheduler_config.max_num_batched_tokens, block_table_width),
|
|
dtype=torch.int32,
|
|
device=self.device,
|
|
)
|
|
|
|
# See: DeepGMM/csrc/apis/attention.hpp. Sized for one slot per SM;
|
|
# build() narrows it to whatever the kernel actually schedules.
|
|
self.scheduler_metadata_buffer = torch.empty(
|
|
(self.num_sms + 1, 2), dtype=torch.int32, device=self.device
|
|
)
|
|
|
|
# Each compressed state must live on a single rank, so the DCP
|
|
# interleave has to cover whole states.
|
|
if (
|
|
self.dcp_world_size > 1
|
|
and self.compress_ratio > 1
|
|
and (
|
|
(
|
|
isinstance(self.kv_cache_spec, MLAAttentionSpec)
|
|
and self.kv_cache_spec.model_version == "deepseek_v4"
|
|
)
|
|
or self.cp_kv_cache_interleave_size % self.compress_ratio != 0
|
|
)
|
|
):
|
|
raise NotImplementedError(
|
|
"DCP with sparse indexer KV compression "
|
|
f"(compress_ratio={self.compress_ratio}) requires "
|
|
"--cp-kv-cache-interleave-size to be a multiple of it; "
|
|
"DeepSeek-V4 is not supported."
|
|
)
|
|
# DCP interleave in units of compressed states.
|
|
self.compressed_cp_interleave_size = (
|
|
self.cp_kv_cache_interleave_size // self.compress_ratio
|
|
if self.dcp_world_size > 1 and self.compress_ratio > 1
|
|
else self.cp_kv_cache_interleave_size
|
|
)
|
|
|
|
# Pre-allocate buffers for CUDA graph compatibility when
|
|
if self.compress_ratio > 1:
|
|
# compress_ratio > 1 (DeepseekV4)
|
|
# Compressed slot mapping output buffer
|
|
self.compressed_slot_mapping_buffer = torch.zeros(
|
|
(scheduler_config.max_num_batched_tokens,),
|
|
dtype=torch.int64,
|
|
device=self.device,
|
|
)
|
|
# Buffer for compressed seq_lens in decode path
|
|
self.expanded_seq_lens_buffer = torch.zeros(
|
|
(scheduler_config.max_num_batched_tokens,),
|
|
dtype=torch.int32,
|
|
device=self.device,
|
|
)
|
|
# DCP not supported yet
|
|
self.supports_draft_decode_metadata_update = self.dcp_world_size == 1
|
|
|
|
def _dcp_localize_decode_seq_lens(
|
|
self,
|
|
seq_lens: torch.Tensor,
|
|
num_decodes: int,
|
|
seq_lens_is_buffer_view: bool,
|
|
) -> torch.Tensor:
|
|
local_seq_lens = get_dcp_local_seq_lens(
|
|
seq_lens,
|
|
self.dcp_world_size,
|
|
self.dcp_rank_tensor,
|
|
self.cp_kv_cache_interleave_size,
|
|
)
|
|
if seq_lens_is_buffer_view:
|
|
seq_lens.copy_(local_seq_lens)
|
|
return seq_lens
|
|
|
|
out = self.decode_seq_lens_buffer[:num_decodes]
|
|
out.copy_(local_seq_lens)
|
|
return out
|
|
|
|
def _prepare_decode_tensors(
|
|
self,
|
|
seq_lens: torch.Tensor,
|
|
block_table: torch.Tensor,
|
|
decode_lens: torch.Tensor,
|
|
decode_lens_cpu: torch.Tensor,
|
|
query_start_loc: torch.Tensor,
|
|
num_decodes: int,
|
|
num_decode_tokens: int,
|
|
use_native: bool,
|
|
next_n: int,
|
|
max_decode_len: int,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, bool]:
|
|
"""Prepare native or per-token flattened decode tensors."""
|
|
spec_config = self.vllm_config.speculative_config
|
|
adaptive = bool(spec_config and spec_config.enable_adaptive_verification)
|
|
min_decode_len = int(decode_lens_cpu.min().item())
|
|
if not use_native:
|
|
assert self.decode_seq_lens_buffer.dim() == 1
|
|
if (
|
|
not self.supports_varlen
|
|
and (num_decodes == 1 or not adaptive)
|
|
and min_decode_len == max_decode_len
|
|
and num_decodes * max_decode_len == num_decode_tokens
|
|
):
|
|
# Uniform decode lengths with no cudagraph token padding.
|
|
_PREPARE_UNIFORM_DECODE_KERNEL(
|
|
seq_lens,
|
|
self.decode_seq_lens_buffer,
|
|
block_table,
|
|
self.expanded_block_table_buffer,
|
|
self.decode_lens_buffer,
|
|
num_decode_tokens,
|
|
max_decode_len,
|
|
)
|
|
self.decode_seq_lens_buffer[num_decode_tokens:] = 0
|
|
seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens]
|
|
block_table = self.expanded_block_table_buffer[:num_decode_tokens]
|
|
decode_lens = self.decode_lens_buffer[:num_decode_tokens]
|
|
return seq_lens, block_table, decode_lens, num_decode_tokens, False
|
|
else:
|
|
# Variable decode lengths.
|
|
# Assume 4 requests with seq_lens [10, 7, 12, 0] (the final req is
|
|
# padding) and decode_lens [3, 1, 4, 0] in the below example comments.
|
|
# The context lengths are therefore
|
|
# [10-3, 7-1, 12-4, 0-0] = [7, 6, 8, 0].
|
|
|
|
# 3 + 1 + 4 + 0 = 8
|
|
actual_expanded = int(decode_lens_cpu.sum().item())
|
|
|
|
# Fuse expanded_base and expanded_starts into a single
|
|
# repeat_interleave:
|
|
# seq_len_i = (context_start[b] - query_start_loc[b]) + arange[i] + 1
|
|
# where context_start[b] = seq_lens[b] - decode_lens[b].
|
|
# Example: offsets = [7-0, 6-3, 8-4, 0-8] = [7, 3, 4, -8]
|
|
# expanded_offsets = [7, 7, 7, 3, 4, 4, 4, 4]
|
|
# result = [8, 9, 10, 7, 9, 10, 11, 12]
|
|
expanded_offsets = torch.repeat_interleave(
|
|
seq_lens - decode_lens - query_start_loc,
|
|
decode_lens,
|
|
output_size=actual_expanded,
|
|
)
|
|
|
|
# [8, 9, 10, 7, 9, 10, 11, 12, ...] where ... is unused buffer space
|
|
self.decode_seq_lens_buffer[:actual_expanded] = (
|
|
expanded_offsets + self.arange_buffer[:actual_expanded] + 1
|
|
)
|
|
self.decode_seq_lens_buffer[actual_expanded:] = 0
|
|
seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens]
|
|
|
|
# Give each of the flattened entries the same block table row as the
|
|
# original request.
|
|
self.expanded_block_table_buffer[:actual_expanded] = (
|
|
torch.repeat_interleave(
|
|
block_table, decode_lens, dim=0, output_size=actual_expanded
|
|
)
|
|
)
|
|
if actual_expanded < num_decode_tokens:
|
|
self.expanded_block_table_buffer[
|
|
actual_expanded:num_decode_tokens, 0
|
|
] = 0
|
|
block_table = self.expanded_block_table_buffer[:num_decode_tokens]
|
|
|
|
# All reqs now have decode_len=1
|
|
self.decode_lens_buffer[:num_decode_tokens] = 1
|
|
decode_lens = self.decode_lens_buffer[:num_decode_tokens]
|
|
return seq_lens, block_table, decode_lens, num_decode_tokens, False
|
|
else:
|
|
# Native path: plain decode (next_n==1) or spec decode
|
|
# with 2D per-token context lengths (next_n > 1).
|
|
#
|
|
# When decode_lens are not truly uniform (e.g. some requests have
|
|
# decode_len < next_n due to padding or short prefills), the simple
|
|
# reshape in sparse_attn_indexer won't work. Use pack_seq_triton
|
|
# (requires_padding) instead.
|
|
requires_padding = min_decode_len != max_decode_len
|
|
if use_native and next_n > 1:
|
|
assert self.decode_seq_lens_buffer.dim() == 1
|
|
# (B, max_decode_len): token j attends to
|
|
# L - max_decode_len + j + 1 KV tokens.
|
|
seq_lens_buffer = self.decode_seq_lens_buffer[
|
|
: num_decodes * max_decode_len
|
|
].view(num_decodes, max_decode_len)
|
|
# Clamp at 0: padding requests have seq_len == 0, which would
|
|
# otherwise make token 0 negative (next_n=2 gives 0-2+1+0 = -1).
|
|
# Downstream kernels read these as uint32, turning -1 into ~4e9.
|
|
seq_lens_buffer[:] = (
|
|
seq_lens.unsqueeze(1)
|
|
- max_decode_len
|
|
+ 1
|
|
+ self.offsets_buffer[:max_decode_len]
|
|
).clamp_(min=0)
|
|
seq_lens = seq_lens_buffer
|
|
return seq_lens, block_table, decode_lens, num_decodes, requires_padding
|
|
|
|
def _prepare_global_decode_seq_lens(
|
|
self,
|
|
global_seq_lens: torch.Tensor | None,
|
|
decode_lens: torch.Tensor,
|
|
decode_lens_cpu: torch.Tensor,
|
|
query_start_loc: torch.Tensor,
|
|
num_decode_tokens: int,
|
|
use_native: bool,
|
|
max_decode_len: int,
|
|
) -> torch.Tensor | None:
|
|
if global_seq_lens is None:
|
|
return None
|
|
if use_native or max_decode_len <= 1:
|
|
return global_seq_lens
|
|
|
|
actual_expanded = int(decode_lens_cpu.sum().item())
|
|
if actual_expanded > 0:
|
|
expanded_offsets = torch.repeat_interleave(
|
|
global_seq_lens - decode_lens - query_start_loc,
|
|
decode_lens,
|
|
output_size=actual_expanded,
|
|
)
|
|
self.global_decode_seq_lens_buffer[:actual_expanded] = (
|
|
expanded_offsets + self.arange_buffer[:actual_expanded] + 1
|
|
)
|
|
self.global_decode_seq_lens_buffer[actual_expanded:num_decode_tokens] = 0
|
|
return self.global_decode_seq_lens_buffer[:num_decode_tokens]
|
|
|
|
def _split_pcp_dcp_prefill_chunks(
|
|
self,
|
|
row_req_idx: np.ndarray,
|
|
row_shard_rows: np.ndarray,
|
|
row_query_lens_cpu: torch.Tensor,
|
|
max_logits_bytes: int,
|
|
request_offset: int,
|
|
) -> list[tuple[slice, slice]]:
|
|
"""Chunk by request rather than by row, so a split prefill's two rows
|
|
charge their shared context once, then widen each chunk back to rows.
|
|
|
|
``row_shard_rows`` holds the whole request's largest DCP shard on
|
|
every row, so the plan is identical on every PCP rank.
|
|
"""
|
|
row_bounds = request_row_bounds(row_req_idx)
|
|
first_rows = row_bounds[:-1]
|
|
# Each request's context, padded to whole DCP shards.
|
|
seq_lens = row_shard_rows[first_rows] * self.dcp_world_size
|
|
# Every rank holds a full-size chunk of a split request (shorter ones
|
|
# are replicated), so rows x longest row is the same on every rank,
|
|
# whichever row holds the short tail.
|
|
row_query_lens = row_query_lens_cpu.numpy()
|
|
query_lens = np.diff(row_bounds) * np.maximum.reduceat(
|
|
row_query_lens, first_rows
|
|
)
|
|
chunk_specs = self._split_indexer_prefill_chunks(
|
|
torch.from_numpy(seq_lens.astype(np.int32)),
|
|
torch.from_numpy(query_lens.astype(np.int32)),
|
|
self.max_prefill_buffer_size,
|
|
max_logits_bytes,
|
|
)
|
|
return [
|
|
(
|
|
slice(
|
|
request_offset + int(row_bounds[request_slice.start]),
|
|
request_offset + int(row_bounds[request_slice.stop]),
|
|
),
|
|
query_slice,
|
|
)
|
|
for request_slice, query_slice in chunk_specs
|
|
]
|
|
|
|
def _prefill_split_seq_lens(self, seq_lens_cpu: torch.Tensor) -> torch.Tensor:
|
|
"""Per-request KV lengths the prefill chunker budgets logits with;
|
|
subclasses whose logits rows are wider than the context override."""
|
|
return seq_lens_cpu
|
|
|
|
@staticmethod
|
|
def _split_indexer_prefill_chunks(
|
|
compressed_seq_lens_cpu: torch.Tensor,
|
|
prefill_query_lens_cpu: torch.Tensor,
|
|
workspace_size: int,
|
|
max_logits_bytes: int,
|
|
request_offset: int = 0,
|
|
) -> list[tuple[slice, slice]]:
|
|
"""Split this step's prefill requests into chunks, respecting:
|
|
- N constraint: total_seq_lens <= workspace_size (existing O(N)
|
|
workspace)
|
|
- Logits constraint: M * N * 4 <= max_logits_bytes
|
|
|
|
When a single request-level chunk still exceeds the logits budget,
|
|
sub-chunks on the query dimension (M) to bound peak memory.
|
|
|
|
Returns list of (req_slice, query_slice) tuples.
|
|
"""
|
|
chunks: list[tuple[slice, slice]] = []
|
|
n = len(compressed_seq_lens_cpu)
|
|
max_logits_elems = max_logits_bytes // 4
|
|
end = 0
|
|
|
|
while end < n:
|
|
start, chunk_m, chunk_n = end, 0, 0
|
|
|
|
while end < n:
|
|
q, s = (
|
|
prefill_query_lens_cpu[end].item(),
|
|
compressed_seq_lens_cpu[end].item(),
|
|
)
|
|
new_m, new_n = chunk_m + q, chunk_n + s
|
|
if new_n <= workspace_size and new_m * new_n <= max_logits_elems:
|
|
chunk_m, chunk_n = new_m, new_n
|
|
end += 1
|
|
else:
|
|
break
|
|
|
|
# A single request can exceed the budget, requiring sub-chunking
|
|
# on the query dimension.
|
|
if end == start:
|
|
chunk_m, chunk_n = (
|
|
prefill_query_lens_cpu[end].item(),
|
|
compressed_seq_lens_cpu[end].item(),
|
|
)
|
|
end += 1
|
|
|
|
req_slice = slice(start + request_offset, end + request_offset)
|
|
max_q = (
|
|
max(1, max_logits_elems // chunk_n) if chunk_n > 0 else max(1, chunk_m)
|
|
)
|
|
for q_off in range(0, chunk_m, max_q):
|
|
sub_m = min(max_q, chunk_m - q_off)
|
|
chunks.append((req_slice, slice(q_off, q_off + sub_m)))
|
|
|
|
return chunks
|
|
|
|
def build(
|
|
self,
|
|
common_prefix_len: int,
|
|
common_attn_metadata: CommonAttentionMetadata,
|
|
fast_build: bool = False,
|
|
) -> DeepseekV32IndexerMetadata:
|
|
num_reqs = common_attn_metadata.num_reqs
|
|
num_tokens = common_attn_metadata.num_actual_tokens
|
|
query_start_loc = common_attn_metadata.query_start_loc
|
|
query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
|
|
seq_lens = common_attn_metadata.seq_lens
|
|
slot_mapping = common_attn_metadata.slot_mapping
|
|
block_table = common_attn_metadata.block_table_tensor
|
|
dcp_local_seq_lens = common_attn_metadata.dcp_local_seq_lens
|
|
|
|
compressed_slot_mapping = slot_mapping
|
|
if self.compress_ratio > 1:
|
|
padded_num_tokens = num_tokens
|
|
local_slot_mapping = slot_mapping
|
|
if self.use_pcp:
|
|
# The gathered layout holds each rank's local tokens, padded, in
|
|
# rank order, so this rank's segment lines up with query_start_loc.
|
|
padded_num_tokens = slot_mapping.shape[0] // self.pcp_world_size
|
|
local_slot_mapping = slot_mapping[
|
|
self.pcp_rank * padded_num_tokens : (self.pcp_rank + 1)
|
|
* padded_num_tokens
|
|
]
|
|
compressed_slot_mapping = get_compressed_slot_mapping(
|
|
num_tokens,
|
|
local_slot_mapping,
|
|
query_start_loc,
|
|
seq_lens,
|
|
block_table,
|
|
self.kv_cache_spec.num_states,
|
|
self.compress_ratio,
|
|
out=self.compressed_slot_mapping_buffer,
|
|
dcp_world_size=self.dcp_world_size,
|
|
cp_interleave=self.cp_kv_cache_interleave_size,
|
|
)
|
|
if self.pcp_world_size > 1:
|
|
compressed_slot_mapping = get_pcp_group().all_gather(
|
|
self.compressed_slot_mapping_buffer[:padded_num_tokens],
|
|
dim=0,
|
|
)
|
|
|
|
# PCP decode sharding keeps a zero-token placeholder on ranks that own
|
|
# no request in a step so collectives retain a uniform rank shape. Do
|
|
# not turn that placeholder into a zero-length indexer decode request.
|
|
if num_tokens == 0:
|
|
return DeepseekV32IndexerMetadata(
|
|
seq_lens=seq_lens,
|
|
max_seq_len=common_attn_metadata.max_seq_len,
|
|
slot_mapping=compressed_slot_mapping,
|
|
num_decodes=0,
|
|
num_decode_tokens=0,
|
|
num_prefills=0,
|
|
num_prefill_tokens=0,
|
|
prefill=None,
|
|
decode=None,
|
|
)
|
|
|
|
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
|
|
split_decodes_and_prefills(
|
|
common_attn_metadata,
|
|
decode_threshold=self.decode_threshold,
|
|
require_uniform=not (self.use_flattening or self.supports_varlen),
|
|
treat_short_extends_as_decodes=not self.use_pcp,
|
|
)
|
|
)
|
|
|
|
assert num_decodes + num_prefills == num_reqs
|
|
assert num_decode_tokens + num_prefill_tokens == num_tokens
|
|
|
|
prefill_metadata = None
|
|
if num_prefills > 0:
|
|
compressed_seq_lens = (
|
|
seq_lens // self.compress_ratio if self.compress_ratio > 1 else seq_lens
|
|
)
|
|
# This CPU value is an upper bound for async-spec extend rows. It
|
|
# is safe for chunking/allocation because CUDA metadata below is
|
|
# built from exact device seq_lens and gather ignores the tail.
|
|
assert common_attn_metadata.seq_lens_cpu_upper_bound is not None
|
|
seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
|
|
compressed_seq_lens_cpu = (
|
|
seq_lens_cpu // self.compress_ratio
|
|
if self.compress_ratio > 1
|
|
else seq_lens_cpu
|
|
)
|
|
prefill_query_lens_cpu = torch.diff(
|
|
query_start_loc_cpu[num_decodes : num_decodes + num_prefills + 1]
|
|
)
|
|
max_logits_bytes = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
|
|
# Upper bound is exact for prefill rows (the `[num_decodes:]`
|
|
# slice below).
|
|
assert common_attn_metadata.seq_lens_cpu_upper_bound is not None
|
|
seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
|
|
req_idx = None
|
|
shard_rows = None
|
|
if self.use_pcp and self.dcp_world_size > 1:
|
|
# The gathered KV must be packed identically on every PCP rank:
|
|
# chunk by request from its DCP shard rows, which every rank
|
|
# holds. A dummy batch bypasses the PCP manager and has one
|
|
# row per request, so its own extent is the request's.
|
|
req_idx = common_attn_metadata.req_idx
|
|
if req_idx is None:
|
|
req_idx = np.arange(num_reqs)
|
|
shard_rows_cpu = common_attn_metadata.dcp_local_seq_lens_cpu_upper_bound
|
|
if shard_rows_cpu is None:
|
|
shard_rows_cpu = get_dcp_local_seq_lens(
|
|
seq_lens_cpu,
|
|
self.dcp_world_size,
|
|
0,
|
|
self.cp_kv_cache_interleave_size,
|
|
)
|
|
shard_rows = shard_rows_cpu.numpy()
|
|
chunk_specs = self._split_pcp_dcp_prefill_chunks(
|
|
req_idx[num_decodes:],
|
|
shard_rows[num_decodes:],
|
|
prefill_query_lens_cpu,
|
|
max_logits_bytes,
|
|
request_offset=num_decodes,
|
|
)
|
|
else:
|
|
chunk_specs = self._split_indexer_prefill_chunks(
|
|
self._prefill_split_seq_lens(compressed_seq_lens_cpu[num_decodes:]),
|
|
prefill_query_lens_cpu,
|
|
self.max_prefill_buffer_size,
|
|
max_logits_bytes,
|
|
request_offset=num_decodes,
|
|
)
|
|
|
|
chunks = []
|
|
for req_slice, query_slice in chunk_specs:
|
|
pcp_plan = None
|
|
if req_idx is not None:
|
|
assert shard_rows is not None
|
|
pcp_plan = build_pcp_global_chunk_plan(
|
|
req_idx[req_slice],
|
|
shard_rows[req_slice],
|
|
self.dcp_world_size,
|
|
self.device,
|
|
self.cp_kv_cache_interleave_size,
|
|
)
|
|
metadata = build_prefill_chunk_metadata(
|
|
req_slice.start,
|
|
req_slice.stop,
|
|
query_start_loc,
|
|
query_start_loc_cpu,
|
|
seq_lens,
|
|
compressed_seq_lens,
|
|
compressed_seq_lens_cpu,
|
|
block_table,
|
|
self.compress_ratio,
|
|
query_slice=query_slice,
|
|
skip_kv_gather=query_slice.start > 0,
|
|
dcp_rank=self.dcp_rank,
|
|
dcp_world_size=self.dcp_world_size,
|
|
cp_kv_cache_interleave_size=self.compressed_cp_interleave_size,
|
|
pcp_plan=pcp_plan,
|
|
)
|
|
# Skip when total_seq_lens is 0 (i.e., no compressed token).
|
|
if metadata is not None:
|
|
chunks.append(metadata)
|
|
row_shard_sizes = None
|
|
if tp_prefill_row_sharding_supported(
|
|
self.vllm_config, self.dcp_world_size, self.use_pcp, self.tp_size
|
|
):
|
|
row_shard_sizes = balanced_prefill_row_shard(
|
|
seq_lens_cpu[num_decodes:],
|
|
prefill_query_lens_cpu,
|
|
self.compress_ratio,
|
|
self.tp_size,
|
|
)
|
|
prefill_metadata = DeepseekV32IndexerPrefillMetadata(
|
|
chunks,
|
|
max_prefill_seq_len=(
|
|
int(seq_lens_cpu[num_decodes:].max().item())
|
|
if num_prefills > 0
|
|
else 0
|
|
),
|
|
row_shard_sizes=row_shard_sizes,
|
|
)
|
|
|
|
decode_metadata = None
|
|
if num_decodes > 0:
|
|
if not self.supports_varlen:
|
|
torch.diff(
|
|
common_attn_metadata.query_start_loc[: num_decodes + 1],
|
|
out=self.decode_lens_buffer[:num_decodes],
|
|
)
|
|
self.per_req_decode_lens_buffer[:num_decodes].copy_(
|
|
self.decode_lens_buffer[:num_decodes]
|
|
)
|
|
decode_lens = self.decode_lens_buffer[:num_decodes]
|
|
decode_lens_cpu = torch.diff(
|
|
common_attn_metadata.query_start_loc_cpu[: num_decodes + 1]
|
|
)
|
|
|
|
# Under DCP the per-token decode bounds must be localized AFTER the
|
|
# per-token expansion below, not before. Expanding from a
|
|
# request-level localized length subtracts decode offsets in local
|
|
# space and yields too-short bounds (e.g. world=2, rank=1, global
|
|
# per-token bounds [8, 9, 10] -> [3, 4, 5] instead of [4, 4, 5]), so
|
|
# the first decode token would run top-k against too short a local KV
|
|
# range and miss valid tokens. Keep the global seq_lens here and
|
|
# localize the expanded bounds further down.
|
|
global_seq_lens_for_decode: torch.Tensor | None = None
|
|
if dcp_local_seq_lens is not None:
|
|
global_seq_lens_for_decode = common_attn_metadata.seq_lens[:num_decodes]
|
|
seq_lens = common_attn_metadata.seq_lens[:num_decodes]
|
|
block_table = common_attn_metadata.block_table_tensor[:num_decodes, ...]
|
|
|
|
max_decode_len = int(decode_lens_cpu.max().item())
|
|
min_decode_len = int(decode_lens_cpu.min().item())
|
|
write_is_uniform = min_decode_len == max_decode_len
|
|
next_n = 1 + self.num_speculative_tokens
|
|
# The kernel sees max_decode_len Q rows, not the configured next_n,
|
|
# so legality is per-step: on SM90 a uniformly 3-deep batch has no
|
|
# native kernel. max_decode_len <= 1 always has one.
|
|
step_next_n_ok = max_decode_len <= 1 or _supports_native_decode(
|
|
max_decode_len
|
|
)
|
|
use_native = (
|
|
not (self.use_flattening or self.supports_varlen)
|
|
and max_decode_len <= next_n
|
|
and step_next_n_ok
|
|
)
|
|
|
|
if not self.supports_varlen:
|
|
global_seq_lens_for_decode = self._prepare_global_decode_seq_lens(
|
|
global_seq_lens=global_seq_lens_for_decode,
|
|
decode_lens=decode_lens,
|
|
decode_lens_cpu=decode_lens_cpu,
|
|
query_start_loc=common_attn_metadata.query_start_loc[:num_decodes],
|
|
num_decode_tokens=num_decode_tokens,
|
|
use_native=use_native,
|
|
max_decode_len=max_decode_len,
|
|
)
|
|
|
|
decode_indices = None
|
|
if self.supports_varlen:
|
|
from vllm.v1.attention.ops.metadata import (
|
|
_indexer_decode_metadata_kernel,
|
|
)
|
|
|
|
capacity = self.decode_seq_lens_buffer.numel()
|
|
grid = max(
|
|
num_decodes,
|
|
num_decode_tokens + triton.cdiv(capacity - num_decode_tokens, 256),
|
|
)
|
|
_indexer_decode_metadata_kernel[(grid,)](
|
|
query_start_loc,
|
|
seq_lens,
|
|
block_table,
|
|
self.decode_seq_lens_buffer,
|
|
self.expanded_block_table_buffer,
|
|
self.decode_lens_buffer,
|
|
self.decode_indices_buffer,
|
|
self.per_req_decode_lens_buffer,
|
|
num_decodes,
|
|
num_decode_tokens,
|
|
capacity,
|
|
block_table.stride(0),
|
|
self.expanded_block_table_buffer.stride(0),
|
|
BLOCK_COLS=block_table.shape[1],
|
|
num_warps=4,
|
|
)
|
|
seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens]
|
|
block_table = self.expanded_block_table_buffer[:num_decode_tokens]
|
|
decode_lens = self.decode_lens_buffer[:num_decode_tokens]
|
|
decode_indices = self.decode_indices_buffer[:num_decode_tokens]
|
|
requires_padding = False
|
|
if global_seq_lens_for_decode is not None and max_decode_len > 1:
|
|
self.global_decode_seq_lens_buffer[:num_decode_tokens].copy_(
|
|
seq_lens
|
|
)
|
|
global_seq_lens_for_decode = self.global_decode_seq_lens_buffer[
|
|
:num_decode_tokens
|
|
]
|
|
else:
|
|
seq_lens, block_table, decode_lens, batch_size, requires_padding = (
|
|
self._prepare_decode_tensors(
|
|
seq_lens=seq_lens,
|
|
block_table=block_table,
|
|
decode_lens=decode_lens,
|
|
decode_lens_cpu=decode_lens_cpu,
|
|
query_start_loc=common_attn_metadata.query_start_loc[
|
|
:num_decodes
|
|
],
|
|
num_decodes=num_decodes,
|
|
num_decode_tokens=num_decode_tokens,
|
|
use_native=use_native,
|
|
next_n=next_n,
|
|
max_decode_len=max_decode_len,
|
|
)
|
|
)
|
|
|
|
# Flattening always returns a buffer view, including single-token
|
|
# batches. Keep its address stable across varlen graph replays.
|
|
seq_lens_is_buffer_view = not use_native or next_n > 1
|
|
|
|
# DCP: localize the now-expanded per-token global bounds to this
|
|
# rank's owned KV. Done here (after expansion) so each token's global
|
|
# causal length is localized individually; see the comment above.
|
|
if dcp_local_seq_lens is not None:
|
|
seq_lens = self._dcp_localize_decode_seq_lens(
|
|
seq_lens, num_decodes, seq_lens_is_buffer_view
|
|
)
|
|
|
|
# For DeepseekV4 (compress_ratio > 1), the indexer KV cache stores
|
|
# compressed tokens. Convert uncompressed seq_lens to compressed.
|
|
if self.compress_ratio > 1:
|
|
if seq_lens_is_buffer_view:
|
|
seq_lens //= self.compress_ratio
|
|
else:
|
|
# Copy to avoid mutating shared state; keeps CG address stable.
|
|
self.expanded_seq_lens_buffer[:num_decodes] = (
|
|
seq_lens // self.compress_ratio
|
|
)
|
|
self.expanded_seq_lens_buffer[num_decodes:num_decode_tokens] = 0
|
|
seq_lens = self.expanded_seq_lens_buffer[:num_decode_tokens]
|
|
|
|
# Non-MTP: deep_gemm paged MQA logits requires 2D context_lens
|
|
# (csrc/apis/attention.hpp). Unsqueeze to (B, 1) so downstream
|
|
# kernels see the same (B, next_n) layout as the MTP path.
|
|
if seq_lens.dim() == 1:
|
|
seq_lens = seq_lens.unsqueeze(-1)
|
|
|
|
# DeepGEMM is required for the paged MQA logits on CUDA devices
|
|
schedule_metadata = self.scheduler_metadata_buffer
|
|
if current_platform.is_cuda() and has_deep_gemm():
|
|
metadata = get_paged_mqa_logits_metadata(
|
|
seq_lens,
|
|
self.kv_cache_spec.num_states,
|
|
self.num_sms,
|
|
indices=decode_indices,
|
|
)
|
|
schedule_metadata = self.scheduler_metadata_buffer[: metadata.shape[0]]
|
|
schedule_metadata[:] = metadata
|
|
|
|
decode_metadata = DeepSeekV32IndexerDecodeMetadata(
|
|
block_table=block_table,
|
|
seq_lens=seq_lens,
|
|
decode_lens=decode_lens,
|
|
requires_padding=requires_padding,
|
|
schedule_metadata=schedule_metadata,
|
|
indices=decode_indices,
|
|
global_seq_lens=global_seq_lens_for_decode,
|
|
per_req_decode_lens=self.per_req_decode_lens_buffer[:num_decodes],
|
|
decode_is_uniform=write_is_uniform,
|
|
write_max_decode_len=max_decode_len,
|
|
)
|
|
|
|
attn_metadata = DeepseekV32IndexerMetadata(
|
|
seq_lens=common_attn_metadata.seq_lens,
|
|
max_seq_len=common_attn_metadata.max_seq_len,
|
|
slot_mapping=compressed_slot_mapping,
|
|
num_decodes=num_decodes,
|
|
num_decode_tokens=num_decode_tokens,
|
|
num_prefills=num_prefills,
|
|
num_prefill_tokens=num_prefill_tokens,
|
|
prefill=prefill_metadata,
|
|
decode=decode_metadata,
|
|
block_table=common_attn_metadata.block_table_tensor,
|
|
)
|
|
|
|
return attn_metadata
|
|
|
|
def update_draft_decode_metadata(
|
|
self,
|
|
metadata: DeepseekV32IndexerMetadata,
|
|
) -> None:
|
|
decode = metadata.decode
|
|
if decode is None or metadata.num_decode_tokens == 0:
|
|
return
|
|
|
|
assert metadata.num_prefills == 0
|
|
assert metadata.num_decodes == metadata.num_decode_tokens
|
|
assert decode.seq_lens.numel() == metadata.num_decode_tokens
|
|
assert metadata.block_table is not None
|
|
assert self.dcp_world_size == 1
|
|
|
|
# One query token per request, so the per-token decode block table is
|
|
# the request block table (at the indexer kernel's page size). build()
|
|
# may copy it into a builder buffer, so re-derive it for this batch.
|
|
decode.block_table.copy_(metadata.block_table[: metadata.num_decode_tokens])
|
|
if decode.indices is not None:
|
|
decode.indices.copy_(self.arange_buffer[: metadata.num_decode_tokens])
|
|
if decode.per_req_decode_lens is not None:
|
|
decode.per_req_decode_lens.fill_(1)
|
|
|
|
if self.compress_ratio > 1:
|
|
get_compressed_slot_mapping(
|
|
metadata.num_decode_tokens,
|
|
self.arange_buffer[: metadata.num_decode_tokens],
|
|
self.arange_buffer[: metadata.num_decode_tokens + 1],
|
|
metadata.seq_lens,
|
|
decode.block_table,
|
|
self.kv_cache_spec.num_states,
|
|
self.compress_ratio,
|
|
out=metadata.slot_mapping,
|
|
)
|
|
torch.div(
|
|
metadata.seq_lens,
|
|
self.compress_ratio,
|
|
rounding_mode="floor",
|
|
out=decode.seq_lens.view(-1),
|
|
)
|
|
else:
|
|
decode.seq_lens.view(-1).copy_(metadata.seq_lens)
|
|
decode.decode_lens.fill_(1)
|
|
|
|
if current_platform.is_cuda() and has_deep_gemm():
|
|
schedule_metadata = get_paged_mqa_logits_metadata(
|
|
decode.seq_lens,
|
|
self.kv_cache_spec.num_states,
|
|
self.num_sms,
|
|
indices=decode.indices,
|
|
)
|
|
assert schedule_metadata.shape == decode.schedule_metadata.shape
|
|
decode.schedule_metadata.copy_(schedule_metadata)
|
|
|
|
|
|
def build_prefill_chunk_metadata(
|
|
start_idx: int,
|
|
end_idx: int,
|
|
query_start_loc: torch.Tensor,
|
|
query_start_loc_cpu: torch.Tensor,
|
|
uncompressed_seq_lens: torch.Tensor,
|
|
compressed_seq_lens: torch.Tensor,
|
|
compressed_seq_lens_cpu: torch.Tensor,
|
|
block_table: torch.Tensor,
|
|
compress_ratio: int,
|
|
query_slice: slice | None = None,
|
|
skip_kv_gather: bool = False,
|
|
dcp_rank: int = 0,
|
|
dcp_world_size: int = 1,
|
|
cp_kv_cache_interleave_size: int = 1,
|
|
pcp_plan: PCPGlobalChunkPlan | None = None,
|
|
) -> DeepseekV32IndexerPrefillChunkMetadata | None:
|
|
if pcp_plan is not None:
|
|
total_seq_lens = pcp_plan.total
|
|
else:
|
|
total_seq_lens = compressed_seq_lens_cpu[start_idx:end_idx].sum().item()
|
|
if total_seq_lens == 0:
|
|
return None
|
|
|
|
num_reqs = end_idx - start_idx
|
|
device = block_table.device
|
|
token_to_seq = torch.empty(total_seq_lens, dtype=torch.int32, device=device)
|
|
|
|
if pcp_plan is not None:
|
|
cu_seq_lens = pcp_plan.global_cu
|
|
else:
|
|
cu_seq_lens = torch.empty(num_reqs + 1, dtype=torch.int32, device=device)
|
|
# Assigning to slice avoids cpu sync.
|
|
cu_seq_lens[:1] = 0
|
|
torch.cumsum(compressed_seq_lens[start_idx:end_idx], dim=0, out=cu_seq_lens[1:])
|
|
|
|
local_cu_seq_lens = cu_seq_lens
|
|
local_total_seq_lens = total_seq_lens
|
|
max_local_total_seq_lens = total_seq_lens
|
|
if pcp_plan is not None:
|
|
local_cu_seq_lens = pcp_plan.padded_local_cu
|
|
local_total_seq_lens = pcp_plan.padded_local_total
|
|
max_local_total_seq_lens = pcp_plan.padded_local_total
|
|
elif dcp_world_size > 1:
|
|
# Per-rank local KV length under interleave-aware DCP sharding, shape
|
|
# [num_reqs, dcp_world_size]. Reuse the canonical CP helper so the
|
|
# sharding matches the rest of the DCP pipeline (decode/prefill).
|
|
local_seq_lens = get_dcp_local_seq_lens(
|
|
compressed_seq_lens[start_idx:end_idx],
|
|
dcp_world_size,
|
|
None,
|
|
cp_kv_cache_interleave_size,
|
|
)
|
|
this_rank_counts = local_seq_lens[:, dcp_rank].to(torch.int32)
|
|
local_cu_seq_lens = torch.zeros(num_reqs + 1, dtype=torch.int32, device=device)
|
|
torch.cumsum(this_rank_counts, dim=0, out=local_cu_seq_lens[1:])
|
|
local_total_seq_lens = int(local_cu_seq_lens[-1].item())
|
|
max_local_total_seq_lens = int(local_seq_lens.sum(dim=0).max().item())
|
|
|
|
query_start_loc = (
|
|
query_start_loc[start_idx : end_idx + 1] - query_start_loc[start_idx]
|
|
)
|
|
|
|
total_query_len = int(
|
|
(query_start_loc_cpu[end_idx] - query_start_loc_cpu[start_idx]).item()
|
|
)
|
|
if query_slice is not None:
|
|
qs_start = min(query_slice.start, total_query_len)
|
|
qs_stop = min(query_slice.stop, total_query_len)
|
|
else:
|
|
qs_start = 0
|
|
qs_stop = total_query_len
|
|
output_query_len = qs_stop - qs_start
|
|
assert 0 <= output_query_len <= total_query_len - qs_start, (
|
|
f"query slice [{qs_start}, {qs_stop}) does not fit this rank's "
|
|
f"{total_query_len} tokens; the metadata tail would stay uninitialized"
|
|
)
|
|
|
|
cu_seq_len_ks = torch.empty(output_query_len, dtype=torch.int32, device=device)
|
|
cu_seq_len_ke = torch.empty(output_query_len, dtype=torch.int32, device=device)
|
|
|
|
if pcp_plan is not None:
|
|
row_start_cu = pcp_plan.row_start_cu
|
|
kernel_dcp_rank, kernel_dcp_world = 0, 1
|
|
else:
|
|
# Under DCP the kernel writes this rank's local row bounds into
|
|
# cu_seq_len_ks/ke; otherwise local_cu_seq_lens aliases cu_seq_lens.
|
|
row_start_cu = local_cu_seq_lens
|
|
kernel_dcp_rank, kernel_dcp_world = dcp_rank, dcp_world_size
|
|
_BUILD_PREFILL_CHUNK_METADATA_KERNEL(
|
|
query_start_loc,
|
|
uncompressed_seq_lens[start_idx:end_idx],
|
|
cu_seq_lens,
|
|
row_start_cu,
|
|
token_to_seq,
|
|
cu_seq_len_ks,
|
|
cu_seq_len_ke,
|
|
qs_start,
|
|
qs_stop,
|
|
kernel_dcp_rank,
|
|
kernel_dcp_world,
|
|
cp_kv_cache_interleave_size,
|
|
num_reqs=num_reqs,
|
|
COMPRESS_RATIO=compress_ratio,
|
|
)
|
|
|
|
pcp_deinterleave_idx = pcp_plan.deinterleave_idx if pcp_plan is not None else None
|
|
|
|
token_start = query_start_loc_cpu[start_idx].item()
|
|
if query_slice is not None:
|
|
token_end = token_start + qs_stop
|
|
token_start = token_start + qs_start
|
|
skip_kv_gather = skip_kv_gather or qs_start > 0
|
|
else:
|
|
token_end = query_start_loc_cpu[end_idx].item()
|
|
|
|
return DeepseekV32IndexerPrefillChunkMetadata(
|
|
pcp_deinterleave_idx=pcp_deinterleave_idx,
|
|
cu_seqlen_ks=cu_seq_len_ks,
|
|
cu_seqlen_ke=cu_seq_len_ke,
|
|
cu_seq_lens=cu_seq_lens,
|
|
token_to_seq=token_to_seq,
|
|
total_seq_lens=total_seq_lens,
|
|
max_seqlen_k=int(compressed_seq_lens_cpu[start_idx:end_idx].max()),
|
|
block_table=block_table[start_idx:end_idx],
|
|
token_start=token_start,
|
|
token_end=token_end,
|
|
num_reqs=num_reqs,
|
|
skip_kv_gather=skip_kv_gather,
|
|
local_cu_seq_lens=local_cu_seq_lens,
|
|
local_total_seq_lens=local_total_seq_lens,
|
|
max_local_total_seq_lens=max_local_total_seq_lens,
|
|
)
|
|
|
|
|
|
_PREPARE_UNIFORM_DECODE_KERNEL = PrepareUniformDecodeKernel()
|
|
_BUILD_PREFILL_CHUNK_METADATA_KERNEL = BuildPrefillChunkMetadataKernel()
|