diff --git a/aphrodite/config/speculative.py b/aphrodite/config/speculative.py index ef09c40a17..852a5b1033 100644 --- a/aphrodite/config/speculative.py +++ b/aphrodite/config/speculative.py @@ -32,6 +32,22 @@ logger = init_logger(__name__) + +def resolve_draft_kv_cache_dtype( + speculative_config: "SpeculativeConfig", + target_cache_dtype: CacheDType, +) -> CacheDType: + """Resolve the draft cache dtype without leaking target-only layouts.""" + if speculative_config.kv_cache_dtype is not None: + return speculative_config.kv_cache_dtype + + draft_model_config = speculative_config.draft_model_config + if target_cache_dtype == "fp8_ds_mla" and not draft_model_config.use_mla: + return "fp8_e4m3" + + return target_cache_dtype + + MTPModelTypes = Literal[ "deepseek_mtp", "mimo_mtp", diff --git a/aphrodite/model_executor/kernels/attention/dsa/sm120_indexer.py b/aphrodite/model_executor/kernels/attention/dsa/sm120_indexer.py new file mode 100644 index 0000000000..8cc7414d8f --- /dev/null +++ b/aphrodite/model_executor/kernels/attention/dsa/sm120_indexer.py @@ -0,0 +1,245 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""SM120 FP8 sparse-attention indexer score kernels.""" + +from functools import cache + +import torch + +from aphrodite.platforms import current_platform +from aphrodite.triton_utils import tl, triton + + +@cache +def use_sm120_dsa_indexer() -> bool: + """Whether the native SM120 sparse-indexer kernels should be used.""" + return current_platform.is_cuda() and current_platform.is_device_capability_family(120) + + +@triton.jit +def _sm120_fp8_mqa_logits_kernel( + q_ptr, + k_ptr, + k_scale_ptr, + weights_ptr, + starts_ptr, + ends_ptr, + logits_ptr, + num_kv_tokens, + stride_q_m: tl.int64, + stride_q_h: tl.int64, + stride_q_d: tl.int64, + stride_k_n: tl.int64, + stride_k_d: tl.int64, + stride_w_m: tl.int64, + stride_w_h: tl.int64, + stride_o_m: tl.int64, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_KV: tl.constexpr, +): + row = tl.program_id(0) + tile = tl.program_id(1) + kv = tile * BLOCK_KV + tl.arange(0, BLOCK_KV) + heads = tl.arange(0, NUM_HEADS)[:, None] + dims = tl.arange(0, HEAD_DIM) + + start = tl.load(starts_ptr + row) + end = tl.load(ends_ptr + row) + valid = (kv >= start) & (kv < end) & (kv < num_kv_tokens) + + q = tl.load( + q_ptr + row * stride_q_m + heads * stride_q_h + dims[None, :] * stride_q_d, + ) + k = tl.load( + k_ptr + dims[:, None] * stride_k_d + kv[None, :] * stride_k_n, + mask=valid[None, :], + other=0.0, + ) + scale = tl.load(k_scale_ptr + kv, mask=valid, other=0.0).to(tl.float32) + weights = tl.load(weights_ptr + row * stride_w_m + heads * stride_w_h).to(tl.float32) + + scores = tl.dot(q, k, input_precision="ieee").to(tl.float32) + scores = tl.maximum(scores * scale[None, :], 0.0) + scores = tl.sum(scores * weights, axis=0) + scores = tl.where(valid, scores, -float("inf")) + tl.store(logits_ptr + row * stride_o_m + kv, scores, mask=kv < num_kv_tokens) + + +@triton.jit +def _sm120_fp8_paged_mqa_logits_kernel( + q_ptr, + k_ptr, + k_scale_ptr, + weights_ptr, + context_lens_ptr, + block_tables_ptr, + logits_ptr, + max_model_len, + stride_q_b: tl.int64, + stride_q_n: tl.int64, + stride_q_h: tl.int64, + stride_q_d: tl.int64, + stride_k_block: tl.int64, + stride_k_token: tl.int64, + stride_k_d: tl.int64, + stride_s_block: tl.int64, + stride_s_token: tl.int64, + stride_ctx_b: tl.int64, + stride_ctx_n: tl.int64, + stride_bt_b: tl.int64, + stride_bt_block: tl.int64, + stride_w_m: tl.int64, + stride_w_h: tl.int64, + stride_o_m: tl.int64, + NEXT_N: tl.constexpr, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + PAGE_SIZE: tl.constexpr, + BLOCK_KV: tl.constexpr, +): + row = tl.program_id(0) + tile = tl.program_id(1) + batch = row // NEXT_N + q_index = row % NEXT_N + logical = tile * BLOCK_KV + tl.arange(0, BLOCK_KV) + heads = tl.arange(0, NUM_HEADS)[:, None] + dims = tl.arange(0, HEAD_DIM) + + context_len = tl.load(context_lens_ptr + batch * stride_ctx_b + q_index * stride_ctx_n) + valid = (logical < context_len) & (logical < max_model_len) + logical_block = logical // PAGE_SIZE + block_offset = logical % PAGE_SIZE + physical_block = tl.load( + block_tables_ptr + batch * stride_bt_b + logical_block * stride_bt_block, + mask=valid, + other=0, + ) + + q = tl.load( + q_ptr + batch * stride_q_b + q_index * stride_q_n + heads * stride_q_h + dims[None, :] * stride_q_d, + ) + k = tl.load( + k_ptr + + physical_block[None, :] * stride_k_block + + block_offset[None, :] * stride_k_token + + dims[:, None] * stride_k_d, + mask=valid[None, :], + other=0.0, + ) + scale = tl.load( + k_scale_ptr + physical_block * stride_s_block + block_offset * stride_s_token, + mask=valid, + other=0.0, + ).to(tl.float32) + weights = tl.load(weights_ptr + row * stride_w_m + heads * stride_w_h).to(tl.float32) + + scores = tl.dot(q, k, input_precision="ieee").to(tl.float32) + scores = tl.maximum(scores * scale[None, :], 0.0) + scores = tl.sum(scores * weights, axis=0) + scores = tl.where(valid, scores, -float("inf")) + tl.store(logits_ptr + row * stride_o_m + logical, scores, mask=logical < max_model_len) + + +def sm120_fp8_mqa_logits( + q: torch.Tensor, + k: torch.Tensor, + k_scales: torch.Tensor, + weights: torch.Tensor, + starts: torch.Tensor, + ends: torch.Tensor, +) -> torch.Tensor: + """Compute ragged FP8 indexer logits on SM120.""" + num_rows, num_heads, head_dim = q.shape + num_kv_tokens = k.shape[0] + logits = torch.empty((num_rows, num_kv_tokens), dtype=torch.float32, device=q.device) + block_kv = 64 + _sm120_fp8_mqa_logits_kernel[(num_rows, triton.cdiv(num_kv_tokens, block_kv))]( + q, + k, + k_scales.reshape(-1), + weights, + starts, + ends, + logits, + num_kv_tokens, + *q.stride(), + *k.stride(), + *weights.stride(), + logits.stride(0), + NUM_HEADS=num_heads, + HEAD_DIM=head_dim, + BLOCK_KV=block_kv, + num_warps=4, + num_stages=2, + ) + return logits + + +def sm120_fp8_paged_mqa_logits( + q: torch.Tensor, + kv_cache: torch.Tensor, + weights: torch.Tensor, + context_lens: torch.Tensor, + block_tables: torch.Tensor, + max_model_len: int, +) -> torch.Tensor: + """Compute paged FP8 indexer logits on SM120.""" + batch_size, next_n, num_heads, head_dim = q.shape + page_size = kv_cache.shape[1] + num_pages = kv_cache.shape[0] + page_stride = kv_cache.stride(0) + cache_bytes = kv_cache.view(torch.uint8) + + # indexer_k_quant_and_cache stores a page in planar form: all + # ``page_size * head_dim`` FP8 key bytes first, followed by one FP32 scale + # per token. The logical cache tensor has a per-token trailing width, but + # slicing that dimension would incorrectly interpret the scale bytes as + # interleaved with each key row. + k_bytes = torch.as_strided( + cache_bytes, + size=(num_pages, page_size, head_dim), + stride=(page_stride, head_dim, 1), + ) + scale_bytes = torch.as_strided( + cache_bytes, + size=(num_pages, page_size, 4), + stride=(page_stride, 4, 1), + storage_offset=page_size * head_dim, + ) + k = k_bytes.view(torch.float8_e4m3fn) + k_scales = scale_bytes.view(torch.float32).squeeze(-1) + if context_lens.ndim == 1: + context_lens = context_lens[:, None].expand(-1, next_n) + + logits = torch.empty( + (batch_size * next_n, max_model_len), + dtype=torch.float32, + device=q.device, + ) + block_kv = 64 + _sm120_fp8_paged_mqa_logits_kernel[(batch_size * next_n, triton.cdiv(max_model_len, block_kv))]( + q, + k, + k_scales, + weights, + context_lens, + block_tables, + logits, + max_model_len, + *q.stride(), + *k.stride(), + *k_scales.stride(), + *context_lens.stride(), + *block_tables.stride(), + *weights.stride(), + logits.stride(0), + NEXT_N=next_n, + NUM_HEADS=num_heads, + HEAD_DIM=head_dim, + PAGE_SIZE=page_size, + BLOCK_KV=block_kv, + num_warps=4, + num_stages=2, + ) + return logits diff --git a/aphrodite/model_executor/layers/attention/attention.py b/aphrodite/model_executor/layers/attention/attention.py index 2679ebf303..2b0d60917c 100644 --- a/aphrodite/model_executor/layers/attention/attention.py +++ b/aphrodite/model_executor/layers/attention/attention.py @@ -231,11 +231,17 @@ def __init__( mm_prefix_clamp_sliding_window: bool = False, attn_backend: type[AttentionBackend] | None = None, head_size_v: int | None = None, + cache_num_kv_heads: int | None = None, **extra_impl_args, ) -> None: """ The KV cache is stored inside this class and is accessed via `self.kv_cache`. + + ``cache_num_kv_heads`` can increase the number of KV heads stored in + each cache block without changing the number projected by the + attention implementation. DFlash and DSpark use this under DCP to + replicate the draft KV heads needed by gathered query heads. """ super().__init__() sliding_window: int | None @@ -310,6 +316,10 @@ def __init__( self.head_size = head_size self.head_size_v = self.head_size if head_size_v is None else head_size_v self.num_kv_heads = num_kv_heads + self.cache_num_kv_heads = num_kv_heads if cache_num_kv_heads is None else cache_num_kv_heads + assert self.cache_num_kv_heads % num_kv_heads == 0, ( + f"cache_num_kv_heads ({self.cache_num_kv_heads}) must be a multiple of num_kv_heads ({num_kv_heads})" + ) self.sliding_window = sliding_window self.has_sink = extra_impl_args.get("sinks") is not None @@ -585,7 +595,7 @@ def get_kv_cache_spec(self, aphrodite_config: AphroditeConfig) -> KVCacheSpec | shared_page = aphrodite_config.cache_config.skip_page_size_padded sw_per_token = SlidingWindowSpec( block_size=1, - num_kv_heads=self.num_kv_heads, + num_kv_heads=self.cache_num_kv_heads, head_size=self.head_size, head_size_v=self.head_size_v, dtype=self.kv_cache_torch_dtype, @@ -595,7 +605,7 @@ def get_kv_cache_spec(self, aphrodite_config: AphroditeConfig) -> KVCacheSpec | sw_block_size = _largest_kernel_block_within(self.attn_backend, sw_per_token, shared_page, block_size) return SlidingWindowSpec( block_size=sw_block_size, - num_kv_heads=self.num_kv_heads, + num_kv_heads=self.cache_num_kv_heads, head_size=self.head_size, head_size_v=self.head_size_v, dtype=self.kv_cache_torch_dtype, @@ -612,7 +622,7 @@ def get_kv_cache_spec(self, aphrodite_config: AphroditeConfig) -> KVCacheSpec | tq_config = TurboQuantConfig.from_cache_dtype(self.kv_cache_dtype, self.head_size) return TQFullAttentionSpec( block_size=block_size, - num_kv_heads=self.num_kv_heads, + num_kv_heads=self.cache_num_kv_heads, head_size=self.head_size, head_size_v=self.head_size, dtype=self.kv_cache_torch_dtype, @@ -621,7 +631,7 @@ def get_kv_cache_spec(self, aphrodite_config: AphroditeConfig) -> KVCacheSpec | else: return FullAttentionSpec( block_size=block_size, - num_kv_heads=self.num_kv_heads, + num_kv_heads=self.cache_num_kv_heads, head_size=self.head_size, head_size_v=self.head_size_v, dtype=self.kv_cache_torch_dtype, diff --git a/aphrodite/model_executor/layers/fused_moe/runner/moe_runner.py b/aphrodite/model_executor/layers/fused_moe/runner/moe_runner.py index eb145b13c9..3c791066cd 100644 --- a/aphrodite/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/aphrodite/model_executor/layers/fused_moe/runner/moe_runner.py @@ -552,6 +552,7 @@ def _apply_quant_method( router_logits: torch.Tensor, shared_experts_input: torch.Tensor | None, input_ids: torch.Tensor | None = None, + shared_experts_overlapping: bool = False, ) -> tuple[torch.Tensor | None, torch.Tensor]: """Run expert routing and the fused MoE kernel via the quant method. @@ -585,10 +586,9 @@ def _apply_quant_method( shared_experts_input=shared_experts_input, ) - self._maybe_apply_shared_experts( - shared_experts_input, - SharedExpertsOrder.MULTI_STREAM_OVERLAPPED, - ) + if shared_experts_overlapping: + assert self._shared_experts is not None + self._shared_experts.wait() return ( self._shared_experts.output if self._shared_experts is not None else None, @@ -606,18 +606,6 @@ def _sequence_parallel_context(self): ctx = get_forward_context() return ctx.dp_metadata.sp_local_sizes(self.moe_config.sp_size) if ctx.dp_metadata else nullcontext() - def _maybe_sync_shared_experts_stream( - self, - shared_experts_input: torch.Tensor | None, - ): - # If router/gate provided, then apply it here. - # (Note: This code runs only when "overlapped mode" is on to allow - # parallel execution of shared experts with the FusedMoEFactory via - # separate cuda stream) - if self._shared_experts is not None: - assert shared_experts_input is not None - self._shared_experts.maybe_sync_shared_experts_stream(shared_experts_input) - def _maybe_add_zero_expert_output( self, result: torch.Tensor, @@ -796,8 +784,11 @@ def _forward_impl( # TODO(bnell): this can be removed after MK migration is complete. self.routed_experts._ensure_moe_quant_config_init() - # Sync aux and main stream for shared expert multi-stream overlap. - self._maybe_sync_shared_experts_stream(shared_experts_input) + # Launch shared experts before routed dispatch so both paths overlap. + shared_experts_overlapping = False + if self._shared_experts is not None: + assert shared_experts_input is not None + shared_experts_overlapping = self._shared_experts.maybe_forward_async(shared_experts_input) # If the Runner holds the gate, apply it after the stream sync, # so it can run overlapped with the @@ -823,6 +814,7 @@ def _forward_impl( router_logits=router_logits, shared_experts_input=shared_experts_input, input_ids=input_ids, + shared_experts_overlapping=shared_experts_overlapping, ) return self._maybe_combine( diff --git a/aphrodite/model_executor/layers/fused_moe/runner/shared_experts.py b/aphrodite/model_executor/layers/fused_moe/runner/shared_experts.py index b781a40943..8f519c9757 100644 --- a/aphrodite/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/aphrodite/model_executor/layers/fused_moe/runner/shared_experts.py @@ -71,6 +71,11 @@ def __init__( if self._stream is not None: logger.debug_once("Enabled separate cuda stream for MoE shared_experts") + if self._stream is not None: + # One pair per DBO ubatch id. + self._input_ready_event = [torch.cuda.Event(), torch.cuda.Event()] + self._output_ready_event = [torch.cuda.Event(), torch.cuda.Event()] + # TODO(bnell): Hack for elastic_ep. Get rid of this def _set_moe_config(self, new_moe_config: FusedMoEConfig): self.moe_config = new_moe_config @@ -106,38 +111,28 @@ def _determine_shared_experts_order( else: return SharedExpertsOrder.NO_OVERLAP - def maybe_sync_shared_experts_stream( - self, - shared_experts_input: torch.Tensor, - ): - experts_order = self._determine_shared_experts_order(shared_experts_input) - - if experts_order == SharedExpertsOrder.MULTI_STREAM_OVERLAPPED: - assert self._stream is not None - - # Record that the clone will be used by shared_experts_stream - # to avoid gc issue from deallocation of hidden_states_clone - # For more details: https://docs.pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html # noqa: E501 - # NOTE: We don't need shared_output.record_stream(current_stream()) - # because we synch the streams before using shared_output. - shared_experts_input.record_stream(self._stream) - - # Mark sync start point for the aux stream since we will - # run in parallel with router/gate. - self._stream.wait_stream(current_stream()) - - def _run_in_aux_stream( - self, - shared_experts_input: torch.Tensor, - ) -> torch.Tensor: - # TODO: assert that maybe_sync_shared_experts_stream has been called. - - # Run shared experts in parallel on a separate stream. + def maybe_forward_async(self, shared_experts_input: torch.Tensor) -> bool: + """Enqueue shared experts on the auxiliary stream. + + Return whether work was enqueued. Call :meth:`wait` before reading + ``output`` when this returns true. + """ + if self._determine_shared_experts_order(shared_experts_input) != SharedExpertsOrder.MULTI_STREAM_OVERLAPPED: + return False + assert self._stream is not None + idx = self._output_idx + assert self._output[idx] is None + self._input_ready_event[idx].record(current_stream()) with torch.cuda.stream(self._stream): - output = self._layer(shared_experts_input) - current_stream().wait_stream(self._stream) + self._input_ready_event[idx].wait(self._stream) + self._output[idx] = self._layer(shared_experts_input) + self._output_ready_event[idx].record(self._stream) + return True - return output + def wait(self) -> None: + """Make the current stream wait for the asynchronous output.""" + assert self._stream is not None + self._output_ready_event[self._output_idx].wait(current_stream()) @property def _output_idx(self) -> int: @@ -162,9 +157,6 @@ def forward( assert self._output[self._output_idx] is None - if order == SharedExpertsOrder.MULTI_STREAM_OVERLAPPED: - self._output[self._output_idx] = self._run_in_aux_stream(shared_experts_input) - else: - self._output[self._output_idx] = self._layer(shared_experts_input) + self._output[self._output_idx] = self._layer(shared_experts_input) assert self._output[self._output_idx] is not None diff --git a/aphrodite/model_executor/layers/sparse_attn_indexer.py b/aphrodite/model_executor/layers/sparse_attn_indexer.py index 83bb701048..012fed96ed 100644 --- a/aphrodite/model_executor/layers/sparse_attn_indexer.py +++ b/aphrodite/model_executor/layers/sparse_attn_indexer.py @@ -13,6 +13,11 @@ from aphrodite.forward_context import get_forward_context from aphrodite.logger import init_logger from aphrodite.model_executor.custom_op import CustomOp +from aphrodite.model_executor.kernels.attention.dsa.sm120_indexer import ( + sm120_fp8_mqa_logits, + sm120_fp8_paged_mqa_logits, + use_sm120_dsa_indexer, +) from aphrodite.model_executor.layers.attention.pcp import maybe_gather_indexer_k from aphrodite.model_executor.layers.quantization.utils.quant_utils import ( get_fp8_min_max, @@ -44,6 +49,7 @@ logger = init_logger(__name__) RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024 +DCP_TOPK_GATHER_MAX_BYTES = 256 * 1024 * 1024 # MXFP4 layout: 2 values packed per byte, ue8m0 (1-byte) scale per block of 32. MXFP4_BLOCK_SIZE = 32 @@ -101,22 +107,29 @@ def _merge_dcp_topk_global( stable_topk_from_gathered_candidates_cutedsl, ) - packed = torch.empty( - (*topk_indices.shape, 2), - dtype=torch.float32, - device=topk_indices.device, - ) - pack_dcp_topk_candidates_cutedsl( - logits, - topk_indices, - packed, - dcp_rank, - dcp_world_size, - cp_interleave, - row_starts, - ) - gathered = get_dcp_group().all_gather(packed, dim=1) - stable_topk_from_gathered_candidates_cutedsl(gathered, topk_tokens, out=topk_indices) + gathered_bytes_per_row = dcp_world_size * topk_tokens * 2 * torch.float32.itemsize + rows_per_collective = max(1, DCP_TOPK_GATHER_MAX_BYTES // gathered_bytes_per_row) + for row_start in range(0, topk_indices.shape[0], rows_per_collective): + row_end = min(row_start + rows_per_collective, topk_indices.shape[0]) + logits_chunk = logits[row_start:row_end] + indices_chunk = topk_indices[row_start:row_end] + starts_chunk = row_starts[row_start:row_end] if row_starts is not None else None + packed = torch.empty( + (*indices_chunk.shape, 2), + dtype=torch.float32, + device=indices_chunk.device, + ) + pack_dcp_topk_candidates_cutedsl( + logits_chunk, + indices_chunk, + packed, + dcp_rank, + dcp_world_size, + cp_interleave, + starts_chunk, + ) + gathered = get_dcp_group().all_gather(packed, dim=1) + stable_topk_from_gathered_candidates_cutedsl(gathered, topk_tokens, out=indices_chunk) @triton.jit @@ -549,6 +562,16 @@ def sparse_attn_indexer( cu_seqlen_ke, logits, ) + elif use_sm120_dsa_indexer(): + assert q_scale_slice is None + logits = sm120_fp8_mqa_logits( + q_slice_cast, + k_quant_cast, + k_scale_cast, + weights[chunk.token_start : chunk.token_end], + cu_seqlen_ks, + cu_seqlen_ke, + ) else: logits = fp8_fp4_mqa_logits( (q_slice_cast, q_scale_slice), @@ -592,7 +615,7 @@ def sparse_attn_indexer( # each page holds block_size fp8 key rows then block_size fp32 # scales) rather than the per-token quant view deep_gemm uses. kv_cache = kv_cache.view(torch.uint8).view(kv_cache.shape[0], -1) - else: + elif not use_sm120_dsa_indexer(): kv_cache = kv_cache_as_quant_view(kv_cache, head_dim, use_fp4_cache) decode_lens = decode_metadata.decode_lens if num_decode_tokens == 0: @@ -670,6 +693,16 @@ def sparse_attn_indexer( logits, False, ) + elif use_sm120_dsa_indexer(): + assert padded_q_scale is None + logits = sm120_fp8_paged_mqa_logits( + padded_q_quant_cast, + kv_cache, + weights[:num_padded_tokens], + seq_lens, + decode_metadata.block_table, + max_model_len, + ) else: logits = fp8_fp4_paged_mqa_logits( (padded_q_quant_cast, padded_q_scale), @@ -845,12 +878,12 @@ def __init__( # 64-token page of alignment padding per request. Used only to # reserve workspace during the profiling run; 0 disables it. self.dcp_local_prefill_shadow_rows = 0 - if self.dcp_world_size > 1 and _use_sm89_dsa(): + if self.dcp_world_size > 1 and (_use_sm89_dsa() or use_sm120_dsa_indexer()): scheduler_config = get_current_aphrodite_config().scheduler_config self.dcp_local_prefill_shadow_rows = ( round_up(scheduler_config.max_num_batched_tokens, 64) + 64 * scheduler_config.max_num_seqs ) - if current_platform.is_cuda() and not has_deep_gemm() and not _use_sm89_dsa(): + if current_platform.is_cuda() and not has_deep_gemm() and not _use_sm89_dsa() and not use_sm120_dsa_indexer(): raise RuntimeError( "Sparse Attention Indexer CUDA op requires DeepGEMM support in the current Aphrodite environment." ) diff --git a/aphrodite/model_executor/models/qwen3_dflash.py b/aphrodite/model_executor/models/qwen3_dflash.py index d998669d7e..5df21a968a 100644 --- a/aphrodite/model_executor/models/qwen3_dflash.py +++ b/aphrodite/model_executor/models/qwen3_dflash.py @@ -13,6 +13,7 @@ from aphrodite.compilation.decorators import support_torch_compile from aphrodite.config import AphroditeConfig, CacheConfig, get_current_aphrodite_config from aphrodite.distributed import ( + get_dcp_group, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, ) @@ -135,6 +136,16 @@ def _resolve_layer_attention(config: Qwen3Config, layer_idx: int) -> tuple[int | return sliding_window, _dflash_layer_causal(config, layer_idx) +def dcp_kv_head_replicas(total_num_kv_heads: int) -> int: + """Return the draft-cache KV-head replication factor required by DCP.""" + parallel_config = get_current_aphrodite_config().parallel_config + dcp_size = parallel_config.decode_context_parallel_size + tp_size = parallel_config.tensor_parallel_size + if dcp_size > 1 and total_num_kv_heads >= tp_size: + return dcp_size + return 1 + + class DFlashQwen3Attention(nn.Module): """Attention for DFlash speculative decoding. @@ -208,6 +219,7 @@ def __init__( ) self.sliding_window = sliding_window + self.dcp_kv_replicas = dcp_kv_head_replicas(self.total_num_kv_heads) self.attn = Attention( self.num_heads, self.head_dim, @@ -219,6 +231,7 @@ def __init__( prefix=f"{prefix}.attn", attn_type=attn_type, sinks=self.attention_sink_bias, + cache_num_kv_heads=self.num_kv_heads * self.dcp_kv_replicas, ) self.causal = causal self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) @@ -451,6 +464,7 @@ def _build_fused_kv_buffers(self) -> None: self._head_dim = attn0.head_dim self._num_kv_heads = attn0.num_kv_heads self._rms_norm_eps = attn0.q_norm.variance_epsilon + self._dcp_kv_replicas = getattr(attn0, "dcp_kv_replicas", 1) # Validation that all layers have the same attention config for attn in layers_attn[1:]: assert ( @@ -574,8 +588,13 @@ def precompute_and_store_context_kv( if context_slot_mapping is None: return - # --- Per-layer cache insert --- all_k_final = all_k_flat.view(L, num_ctx, nkv, hd) + if self._dcp_kv_replicas > 1: + dcp_group = get_dcp_group() + all_k_final = dcp_group.all_gather(all_k_final.contiguous(), dim=2) + all_v = dcp_group.all_gather(all_v.contiguous(), dim=2) + + # --- Per-layer cache insert --- per_layer = isinstance(context_slot_mapping, (list, tuple)) for i in range(L): slot_mapping = context_slot_mapping[i] if per_layer else context_slot_mapping diff --git a/aphrodite/models/deepseek_v32/attention.py b/aphrodite/models/deepseek_v32/attention.py index cb2e804214..b1a48736be 100644 --- a/aphrodite/models/deepseek_v32/attention.py +++ b/aphrodite/models/deepseek_v32/attention.py @@ -6,7 +6,7 @@ from aphrodite.compilation.breakable_cudagraph import eager_break_during_capture from aphrodite.config import AphroditeConfig, CacheConfig -from aphrodite.distributed import get_tensor_model_parallel_world_size +from aphrodite.distributed import get_dcp_group, get_tensor_model_parallel_world_size from aphrodite.forward_context import get_forward_context from aphrodite.model_executor.layers.attention import MLAAttention from aphrodite.model_executor.layers.layernorm import LayerNorm, RMSNorm @@ -33,6 +33,8 @@ from aphrodite.model_executor.models.utils import extract_layer_index from aphrodite.models.deepseek_v32.common.kernels import fused_norm_rope, fused_q from aphrodite.utils.torch_utils import is_quantized_kv_cache +from aphrodite.v1.attention.ops.common import cp_lse_ag_out_rs +from aphrodite.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce class DeepseekV32Indexer(nn.Module): @@ -459,6 +461,11 @@ def _fused_attention( use_fp4_cache=False, # fused_norm_rope already cleared the topk buffer this forward. skip_topk_buffer_clear=True, + # This fused path bypasses SparseAttnIndexer.forward_cuda, so + # forward its run-constant DCP geometry explicitly. + dcp_rank=self.indexer.indexer_op.dcp_rank, + dcp_world_size=self.indexer.indexer_op.dcp_world_size, + cp_kv_cache_interleave_size=(self.indexer.indexer_op.cp_kv_cache_interleave_size), ) if attn_metadata is None: @@ -474,9 +481,39 @@ def _fused_attention( mqa_q_arg: torch.Tensor | tuple[torch.Tensor, torch.Tensor] = mqa_q[:num_actual] else: mqa_q_arg = (ql_nope[:num_actual], mqa_q[:num_actual]) - attn_out, _ = self.impl.forward_mqa( # type: ignore[attr-defined] + + if self.impl.dcp_world_size > 1: + if self.use_pcp: + raise NotImplementedError( + "The NVIDIA DeepSeek-v3.2/GLM-5.2 override does not yet support combined PCP and DCP attention." + ) + if isinstance(mqa_q_arg, tuple): + mqa_q_arg = torch.cat(mqa_q_arg, dim=-1) + # Every DCP rank evaluates all query heads against its local KV + # shard. The per-rank outputs are LSE-combined below. + mqa_q_arg = get_dcp_group().all_gather(mqa_q_arg, dim=1) + + attn_out, lse = self.impl.forward_mqa( # type: ignore[attr-defined] mqa_q_arg, kv_cache, attn_metadata, self ) + if self.impl.dcp_world_size > 1: + if lse is None: + raise RuntimeError("The NVIDIA DCP attention path requires per-head LSE from the sparse MLA backend.") + dcp_group = get_dcp_group() + if self.dcp_a2a: + attn_out = dcp_a2a_lse_reduce( + attn_out, + lse, + dcp_group, + is_lse_base_on_e=self.impl.lse_base_on_e, + ) + else: + attn_out = cp_lse_ag_out_rs( + attn_out, + lse, + dcp_group, + is_lse_base_on_e=self.impl.lse_base_on_e, + ) x = attn_out.view(num_actual, self.num_local_heads, self.kv_lora_rank).transpose(0, 1) out = output[:num_actual].view(num_actual, self.num_local_heads, self.v_head_dim).transpose(0, 1) torch.bmm(x, self.W_UV, out=out) diff --git a/aphrodite/models/deepseek_v32/common/kernels.py b/aphrodite/models/deepseek_v32/common/kernels.py index 9c54411c55..0bb0750a35 100644 --- a/aphrodite/models/deepseek_v32/common/kernels.py +++ b/aphrodite/models/deepseek_v32/common/kernels.py @@ -169,20 +169,25 @@ def _fused_norm_rope_kernel( if slot_mapping_ptr is None: # Memory profiling run. return - slot_idx = tl.load(slot_mapping_ptr + tok_idx) - if slot_idx < 0: - # Padding - return if pid == 2: - # Q RMS norm + # Query normalization is independent of KV-cache ownership. Under DCP, + # non-owner ranks have a negative slot mapping but still contribute a + # query-head shard to the subsequent all-gather. q_block = tl.arange(0, Q_BLOCK_SIZE) q_mask = q_block < Q_DIM q_c = tl.load(q_c_ptr + tok_idx * q_c_stride + q_block, mask=q_mask, other=0.0) q_c_rms_w = tl.load(q_rms_norm_w_ptr + q_block, mask=q_mask) q_c = _rms_norm(q_c, q_c_rms_w, q_rms_eps, Q_DIM) tl.store(q_c_out_ptr + tok_idx * q_c_out_stride + q_block, q_c, mask=q_mask) - elif pid == 1: + return + + slot_idx = tl.load(slot_mapping_ptr + tok_idx) + if slot_idx < 0: + # Padding or a token owned by another DCP rank. + return + + if pid == 1: # KV RMS Norm + KV RoPE + MLA concat_and_cache. # Merged so the normed kv_c and RoPE'd k_pe can be written # to the MLA KV cache directly without a separate kernel. diff --git a/aphrodite/v1/attention/backends/flash_attn.py b/aphrodite/v1/attention/backends/flash_attn.py index 9bc0daf013..bc60805f48 100755 --- a/aphrodite/v1/attention/backends/flash_attn.py +++ b/aphrodite/v1/attention/backends/flash_attn.py @@ -1008,7 +1008,41 @@ def _forward_with_dcp( max_seqlen_q = attn_metadata.max_query_len block_table = attn_metadata.block_table + # A non-quantized DCP draft cache can store more KV heads than the + # implementation projects. Its descales are no-op scalars, and leaving + # them populated makes FlashAttention validate them against the larger + # replicated cache-head dimension. + context_q_descale = q_descale + context_k_descale = k_descale + context_v_descale = v_descale + if not is_quantized_kv_cache(self.kv_cache_dtype): + context_q_descale = context_k_descale = context_v_descale = None + q_descale = k_descale = v_descale = None query = query.contiguous() + if attn_metadata.max_dcp_context_kv_len == 0: + flash_attn_varlen_func( + q=query, + k=key, + v=value, + out=output, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=max_seqlen_q, + cu_seqlens_k=cu_seqlens_q, + max_seqlen_k=max_seqlen_q, + softmax_scale=self.scale, + causal=attn_metadata.causal, + alibi_slopes=self.alibi_slopes, + window_size=list(self.sliding_window) if self.sliding_window is not None else None, + softcap=self.logits_soft_cap, + return_softmax_lse=True, + fa_version=self.vllm_flash_attn_version, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + num_splits=attn_metadata.max_num_splits, + ) + return output + query_across_dcp = get_dcp_group().all_gather(query, dim=1) sliding_window_size = list(self.sliding_window) if self.sliding_window is not None else None n = query_across_dcp.shape[0] @@ -1036,9 +1070,9 @@ def _forward_with_dcp( return_softmax_lse=True, scheduler_metadata=attn_metadata.scheduler_metadata, fa_version=self.vllm_flash_attn_version, - q_descale=q_descale, - k_descale=k_descale, - v_descale=v_descale, + q_descale=context_q_descale, + k_descale=context_k_descale, + v_descale=context_v_descale, num_splits=attn_metadata.max_num_splits, ) # FA returns LSE in shape [ H, B ] but DCP combine wants [ B, H ] diff --git a/aphrodite/v1/attention/backends/mla/flashinfer_mla_sparse.py b/aphrodite/v1/attention/backends/mla/flashinfer_mla_sparse.py index 07ac27082a..6bf9bc3340 100644 --- a/aphrodite/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/aphrodite/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -273,18 +273,85 @@ def __init__( supports_dcp_with_varlen=True, ) + dcp_world_size = aphrodite_config.parallel_config.decode_context_parallel_size + if dcp_world_size > 1: + _get_workspace_buffer( + device, + _required_workspace_bytes( + dcp_world_size, + num_q_heads, + aphrodite_config.scheduler_config.max_num_batched_tokens, + ), + ) + # Global workspace buffer (lazily initialized) _fi_sparse_workspace: torch.Tensor | None = None +_TRTLLM_GEN_SOFTMAX_STAT_BYTES = 8 +_TRTLLM_GEN_SOFTMAX_SLOTS_PER_TOKEN = 256 +_TRTLLM_GEN_SOFTMAX_GUARD_BYTES = 1024 * 1024 +_DEFAULT_WORKSPACE_BUFFER_SIZE = 394 * 1024 * 1024 + + +def compute_trtllm_sparse_mla_workspace_bytes( + base_workspace_bytes: int, + dcp_world_size: int, + num_heads_per_rank: int, + max_num_batched_tokens: int, +) -> int: + """Return workspace bytes required by sparse MLA decode under DCP.""" + if dcp_world_size <= 1: + return base_workspace_bytes + softmax_bytes = ( + _TRTLLM_GEN_SOFTMAX_STAT_BYTES + * (num_heads_per_rank * dcp_world_size) + * max_num_batched_tokens + * _TRTLLM_GEN_SOFTMAX_SLOTS_PER_TOKEN + + _TRTLLM_GEN_SOFTMAX_GUARD_BYTES + ) + return base_workspace_bytes + softmax_bytes + + +def _required_workspace_bytes( + dcp_world_size: int, + num_heads_per_rank: int, + max_num_batched_tokens: int, +) -> int: + computed = compute_trtllm_sparse_mla_workspace_bytes( + _DEFAULT_WORKSPACE_BUFFER_SIZE, + dcp_world_size, + num_heads_per_rank, + max_num_batched_tokens, + ) + if not envs.is_set("APHRODITE_FLASHINFER_WORKSPACE_BUFFER_SIZE"): + return computed + + env_bytes = envs.APHRODITE_FLASHINFER_WORKSPACE_BUFFER_SIZE + if env_bytes < computed: + logger.warning_once( + "APHRODITE_FLASHINFER_WORKSPACE_BUFFER_SIZE=%d is below the " + "%d bytes required for sparse MLA with DCP=%d, %d heads per " + "rank, and max_num_batched_tokens=%d. The kernel may overflow " + "its workspace; set the value to at least %d or unset it.", + env_bytes, + computed, + dcp_world_size, + num_heads_per_rank, + max_num_batched_tokens, + computed, + ) + return env_bytes + -def _get_workspace_buffer(device: torch.device) -> torch.Tensor: +def _get_workspace_buffer(device: torch.device, min_bytes: int | None = None) -> torch.Tensor: global _fi_sparse_workspace - if _fi_sparse_workspace is None: + required = min_bytes if min_bytes is not None else envs.APHRODITE_FLASHINFER_WORKSPACE_BUFFER_SIZE + if _fi_sparse_workspace is None or _fi_sparse_workspace.numel() < required: # FlashInfer's CuteDSL MLA-decode tactic requires an int8 workspace; # the trtllm-gen path views it as uint8, so int8 is safe for all backends. _fi_sparse_workspace = torch.zeros( - envs.APHRODITE_FLASHINFER_WORKSPACE_BUFFER_SIZE, + required, dtype=torch.int8, device=device, ) diff --git a/aphrodite/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py b/aphrodite/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py index 7c60b696fa..71776feab7 100644 --- a/aphrodite/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py +++ b/aphrodite/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """SM120 implementation variant for ``FLASHINFER_MLA_SPARSE_SM120``.""" -from typing import TYPE_CHECKING, cast +from typing import TYPE_CHECKING import torch @@ -17,6 +17,7 @@ ) from aphrodite.v1.attention.backends.mla.sparse_utils import ( triton_convert_req_index_to_global_index, + triton_filter_and_convert_dcp_index, ) if TYPE_CHECKING: @@ -33,6 +34,11 @@ class FlashInferMLASparseSM120Impl(MLAAttentionImpl[FlashInferMLASparseMetadata] """SM120 FlashInfer sparse-MLA implementation.""" is_sparse = True + # The SM120 launcher provides only sparse MQA. Dense/masked prefill is not + # implemented for GLM's head geometry and packed FP8 cache. + supports_dense_mha_prefill = False + can_return_lse_for_decode = True + lse_base_on_e = False def __init__( self, @@ -70,6 +76,10 @@ def __init__( self.kv_lora_rank: int = mla_args["kv_lora_rank"] self.qk_nope_head_dim: int = mla_args["qk_nope_head_dim"] self.qk_rope_head_dim: int = mla_args["qk_rope_head_dim"] + self.topk_indices_buffer: torch.Tensor | None = ( + indexer.topk_indices_buffer if indexer is not None else mla_args.get("topk_indices_buffer") + ) + from aphrodite.config import get_current_aphrodite_config aphrodite_config = get_current_aphrodite_config() @@ -78,11 +88,6 @@ def __init__( model_type = getattr(aphrodite_config.model_config.hf_text_config, "model_type", None) self.kv_scale_format = _kv_scale_format_for_model(model_type) - # Skip-topk layers are built with indexer=None and get the shared - # buffer via mla_args instead (cf. FLASHMLA_SPARSE). - self.topk_indices_buffer: torch.Tensor | None = ( - indexer.topk_indices_buffer if indexer is not None else mla_args.get("topk_indices_buffer") - ) from aphrodite.utils.flashinfer import has_flashinfer_sparse_mla_sm120 if not has_flashinfer_sparse_mla_sm120(): @@ -107,19 +112,30 @@ def forward_mqa( assert self.topk_indices_buffer is not None topk_indices = self.topk_indices_buffer[:num_actual_toks] - topk_indices_physical = cast( - torch.Tensor, - triton_convert_req_index_to_global_index( + if self.dcp_world_size > 1: + topk_indices_physical, seq_lens = triton_filter_and_convert_dcp_index( attn_metadata.req_id_per_token[:num_actual_toks], attn_metadata.block_table, topk_indices, + dcp_size=self.dcp_world_size, + dcp_rank=self.dcp_rank, + cp_kv_cache_interleave_size=attn_metadata.cp_kv_cache_interleave_size, BLOCK_SIZE=attn_metadata.block_size, NUM_TOPK_TOKENS=topk_indices.shape[1], - ), - ) + return_valid_counts=True, + ) + else: + topk_indices_physical, seq_lens = triton_convert_req_index_to_global_index( + attn_metadata.req_id_per_token[:num_actual_toks], + attn_metadata.block_table, + topk_indices, + BLOCK_SIZE=attn_metadata.block_size, + NUM_TOPK_TOKENS=topk_indices.shape[1], + return_valid_counts=True, + ) output = q.new_empty( - (num_actual_toks, self.num_heads, self.kv_lora_rank), + (num_actual_toks, q.shape[-2], self.kv_lora_rank), dtype=q.dtype, ) @@ -138,12 +154,44 @@ def forward_mqa( kv_lora_rank=self.kv_lora_rank, qk_rope_head_dim=self.qk_rope_head_dim, block_tables=topk_indices_physical.unsqueeze(1), - seq_lens=None, + # Preserve FlashInfer's original SM120 contract outside DCP. The + # explicit valid-count vector is required only after DCP filters + # the global top-k down to the current rank's compact local set. + seq_lens=seq_lens if self.dcp_world_size > 1 else None, max_seq_len=attn_metadata.topk_tokens, out=output.unsqueeze(1), bmm1_scale=self.scale, bmm2_scale=1.0, sparse_mla_top_k=attn_metadata.topk_tokens, kv_scale_format=self.kv_scale_format, + return_lse=self.need_to_return_lse_for_decode, ) - return out.squeeze(1), None + if self.need_to_return_lse_for_decode: + assert isinstance(out, tuple) + output, lse = out + else: + assert isinstance(out, torch.Tensor) + output, lse = out, None + + output = output.squeeze(1) + if lse is not None: + lse = self._normalize_lse(lse, output.shape[0], output.shape[1]) + empty_rows = (topk_indices_physical == -1).all(dim=-1) + output.masked_fill_(empty_rows.view(-1, 1, 1), 0.0) + lse.masked_fill_(empty_rows.view(-1, 1), float("-inf")) + return output, lse + + @staticmethod + def _normalize_lse(lse: torch.Tensor, num_tokens: int, num_heads: int) -> torch.Tensor: + if lse.dim() == 3: + if lse.shape[-1] == 1: + lse = lse.squeeze(-1) + elif lse.shape[1] == 1: + lse = lse.squeeze(1) + elif lse.shape[0] * lse.shape[1] == num_tokens: + lse = lse.reshape(num_tokens, lse.shape[-1]) + if lse.shape != (num_tokens, num_heads): + raise RuntimeError( + f"Unexpected FlashInfer sparse MLA LSE shape: {tuple(lse.shape)}, expected ({num_tokens}, {num_heads})." + ) + return lse diff --git a/aphrodite/v1/attention/backends/mla/indexer.py b/aphrodite/v1/attention/backends/mla/indexer.py index a771374e39..c9eb0c3852 100644 --- a/aphrodite/v1/attention/backends/mla/indexer.py +++ b/aphrodite/v1/attention/backends/mla/indexer.py @@ -10,6 +10,9 @@ from aphrodite.config import AphroditeConfig from aphrodite.distributed import get_dcp_group, get_pcp_group from aphrodite.logger import init_logger +from aphrodite.model_executor.kernels.attention.dsa.sm120_indexer import ( + use_sm120_dsa_indexer, +) from aphrodite.model_executor.warmup.jit_warmup import ( AphroditeJitKernel, WarmupIntRange, @@ -41,6 +44,7 @@ split_decodes_and_prefills, ) from aphrodite.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec +from aphrodite.v1.worker.block_table import get_block_table_width from aphrodite.v1.worker.cp_utils import get_kv_cache_shard_count logger = init_logger(__name__) @@ -443,7 +447,10 @@ def __init__(self, *args, **kwargs): # caches). Outside the SM100 family the FP8 # paged MQA logits kernel only supports next_n in (1, 2) # (deepgemm smxx_fp8_fp4_paged_mqa_logits.hpp:233), so flatten there. - self.use_flattening = not current_platform.is_device_capability_family(100) and next_n not in (1, 2) + self.use_sm120_dsa = use_sm120_dsa_indexer() + self.use_flattening = ( + not current_platform.is_device_capability_family(100) and not self.use_sm120_dsa and next_n not in (1, 2) + ) logger.info_once( "DSA indexer decode path: use_flattening=%s (next_n=%d, use_fp4_indexer_cache=%s)", self.use_flattening, @@ -488,9 +495,12 @@ def __init__(self, *args, **kwargs): dtype=torch.int32, device=self.device, ) - max_num_blocks_per_req = cdiv( - self.aphrodite_config.model_config.max_model_len, - self.kv_cache_spec.block_size * get_kv_cache_shard_count(), + max_num_blocks_per_req = get_block_table_width( + cdiv( + self.aphrodite_config.model_config.max_model_len, + self.kv_cache_spec.block_size * get_kv_cache_shard_count(), + ), + self.kv_cache_spec.block_size, ) self.expanded_block_table_buffer = torch.zeros( ( @@ -514,7 +524,7 @@ def __init__(self, *args, **kwargs): f"DCP is not supported with sparse indexer KV compression (compress_ratio={self.compress_ratio})." ) - if self.dcp_world_size > 1 and self.use_sm89_dsa: + if self.dcp_world_size > 1 and (self.use_sm89_dsa or self.use_sm120_dsa): self._warmup_dcp_kernels() # Pre-allocate buffers for CUDA graph compatibility when @@ -612,7 +622,7 @@ def _use_dcp_local_prefill( seq_lens_cpu: torch.Tensor, prefill_query_lens_cpu: torch.Tensor, ) -> bool: - """Whether this batch qualifies for the sm89 DCP local-prefill path. + """Whether this batch qualifies for the native DCP local-prefill path. Fresh prompts only (context length 0 for every prefill request). The full prompt K is then computable on-rank from the replicated @@ -625,9 +635,14 @@ def _use_dcp_local_prefill( """ parallel_config = self.aphrodite_config.parallel_config if ( - not self.use_sm89_dsa + not (self.use_sm89_dsa or self.use_sm120_dsa) or self.dcp_world_size <= 1 or parallel_config.prefill_context_parallel_size > 1 + # Speculative decoding can turn the first scheduler step into a + # mixed prefill/decode sequence before the next metadata build. + # Keep it on the regular DCP gather-and-merge path; the ephemeral + # full-prompt shadow cache is only safe for non-speculative + # prefills. or self.num_speculative_tokens > 0 or num_prefills == 0 ): @@ -1030,7 +1045,7 @@ def build( seq_lens.shape[1], ) # DeepGEMM is required for the paged MQA logits on CUDA devices - elif current_platform.is_cuda() and has_deep_gemm(): + elif current_platform.is_cuda() and has_deep_gemm() and not self.use_sm120_dsa: self.scheduler_metadata_buffer[:] = get_paged_mqa_logits_metadata( seq_lens, self.kv_cache_spec.storage_block_size, diff --git a/aphrodite/v1/attention/backends/mla/sparse_utils.py b/aphrodite/v1/attention/backends/mla/sparse_utils.py index 9083aaca98..1d6c81f9fd 100644 --- a/aphrodite/v1/attention/backends/mla/sparse_utils.py +++ b/aphrodite/v1/attention/backends/mla/sparse_utils.py @@ -23,6 +23,9 @@ def _convert_req_index_to_global_index_kernel( BLOCK_N: tl.constexpr, # tile width along columns HAS_PREFILL: tl.constexpr, COUNT_VALID: tl.constexpr, # whether to count valid indices + # BLOCK_N == NUM_TOPK_TOKENS: one program owns the row, so the valid count + # is an in-register reduction and needs no atomic. + SINGLE_TILE: tl.constexpr, # When set, scatter valid slots to a contiguous prefix [0, valid_count) using # valid_count_ptr as an atomic slot allocator (DCP filtering leaves interior # -1 gaps; the trtllm-gen sparse kernel reads the first valid_count entries). @@ -105,7 +108,11 @@ def _convert_req_index_to_global_index_kernel( is_valid = (~is_invalid_tok).to(tl.int32) local_offset = tl.cumsum(is_valid) - is_valid tile_valid_count = tl.sum(is_valid) - base = tl.atomic_add(valid_count_ptr + token_id, tile_valid_count) + if SINGLE_TILE: + base = 0 + tl.store(valid_count_ptr + token_id, tile_valid_count) + else: + base = tl.atomic_add(valid_count_ptr + token_id, tile_valid_count) dest = base + local_offset out_ptr_dest = out_ptr + token_id * out_stride0 + dest * out_stride1 tl.store(out_ptr_dest, out_val, mask=is_valid == 1) @@ -117,7 +124,18 @@ def _convert_req_index_to_global_index_kernel( # Count valid indices in this tile and atomically add to row total if COUNT_VALID: tile_valid_count = tl.sum((~is_invalid_tok).to(tl.int32)) - tl.atomic_add(valid_count_ptr + token_id, tile_valid_count) + if SINGLE_TILE: + tl.store(valid_count_ptr + token_id, tile_valid_count) + else: + tl.atomic_add(valid_count_ptr + token_id, tile_valid_count) + + +def _remap_tiling(NUM_TOPK_TOKENS: int, BLOCK_N: int, count_valid: bool) -> tuple[bool, int, int, int]: + """Pick the column tiling for the sparse-index remap kernel.""" + single_tile = count_valid and triton.next_power_of_2(NUM_TOPK_TOKENS) == NUM_TOPK_TOKENS + if single_tile: + return True, NUM_TOPK_TOKENS, 1, 8 + return False, BLOCK_N, NUM_TOPK_TOKENS // BLOCK_N, 4 def triton_convert_req_index_to_global_index( @@ -175,7 +193,7 @@ def triton_convert_req_index_to_global_index( num_tokens = req_id.shape[0] max_num_blocks_per_req = block_table.shape[1] - tiles_per_row = NUM_TOPK_TOKENS // BLOCK_N + single_tile, block_n, tiles_per_row, num_warps = _remap_tiling(NUM_TOPK_TOKENS, BLOCK_N, return_valid_counts) # Ensure contiguous tensors on the same device req_id_c = req_id.contiguous() @@ -186,7 +204,8 @@ def triton_convert_req_index_to_global_index( # Allocate valid count buffer if needed (must be zero-initialized for atomics) valid_counts: torch.Tensor | None = None if return_valid_counts: - valid_counts = torch.zeros(num_tokens, dtype=torch.int32, device=token_indices.device) + alloc = torch.empty if single_tile else torch.zeros + valid_counts = alloc(num_tokens, dtype=torch.int32, device=token_indices.device) # Strides in elements bt_stride0, bt_stride1 = block_table_c.stride() @@ -214,9 +233,10 @@ def triton_convert_req_index_to_global_index( # shapes / constexprs max_num_blocks_per_req, BLOCK_SIZE, - BLOCK_N, + block_n, HAS_PREFILL_WORKSPACE, return_valid_counts, + single_tile, False, # COMPACT_TO_FRONT: keep input column == output column # DCP disabled (no-op de-interleave) 1, @@ -229,6 +249,7 @@ def triton_convert_req_index_to_global_index( ti_stride1, out_stride0, out_stride1, + num_warps=num_warps, ) if return_valid_counts: @@ -287,8 +308,6 @@ def triton_filter_and_convert_dcp_index( num_tokens = req_id.shape[0] max_num_blocks_per_req = block_table.shape[1] - tiles_per_row = NUM_TOPK_TOKENS // BLOCK_N - req_id_c = req_id.contiguous() block_table_c = block_table.contiguous() token_indices_c = token_indices.contiguous() @@ -296,6 +315,7 @@ def triton_filter_and_convert_dcp_index( # The compaction uses the valid-count buffer as an atomic slot allocator, so # it requires counting. Pre-fill out with -1 so the unwritten tail stays -1. count_valid = return_valid_counts or compact_valid_to_front + single_tile, block_n, tiles_per_row, num_warps = _remap_tiling(NUM_TOPK_TOKENS, BLOCK_N, count_valid) if compact_valid_to_front: out = torch.full_like(token_indices_c, -1) else: @@ -303,7 +323,8 @@ def triton_filter_and_convert_dcp_index( valid_counts: torch.Tensor | None = None if count_valid: - valid_counts = torch.zeros(num_tokens, dtype=torch.int32, device=token_indices.device) + alloc = torch.empty if single_tile else torch.zeros + valid_counts = alloc(num_tokens, dtype=torch.int32, device=token_indices.device) bt_stride0, bt_stride1 = block_table_c.stride() ti_stride0, ti_stride1 = token_indices_c.stride() @@ -320,9 +341,10 @@ def triton_filter_and_convert_dcp_index( None, max_num_blocks_per_req, BLOCK_SIZE, - BLOCK_N, + block_n, False, # HAS_PREFILL count_valid, + single_tile, compact_valid_to_front, dcp_size, dcp_rank, @@ -333,6 +355,7 @@ def triton_filter_and_convert_dcp_index( ti_stride1, out_stride0, out_stride1, + num_warps=num_warps, ) if return_valid_counts: diff --git a/aphrodite/v1/executor/multiproc_executor.py b/aphrodite/v1/executor/multiproc_executor.py index eb767ee9a0..2bab6c2aae 100644 --- a/aphrodite/v1/executor/multiproc_executor.py +++ b/aphrodite/v1/executor/multiproc_executor.py @@ -773,6 +773,14 @@ def worker_main(*args, **kwargs): """Worker initialization and execution loops. This runs a background process""" + # Triton's cache writes are not atomic across processes. Give each + # local worker a stable subdirectory so simultaneous first-use JITs do + # not observe another worker's partially written metadata. Keeping the + # rank in the path preserves reuse across server restarts. + rank = kwargs.get("rank", 0) + triton_cache_root = os.environ.get("TRITON_CACHE_DIR", os.path.expanduser("~/.triton/cache")) + os.environ["TRITON_CACHE_DIR"] = os.path.join(triton_cache_root, f"worker-{rank}") + # Signal handler used for graceful termination. # SystemExit exception is only raised once to allow this and worker # processes to terminate without error @@ -816,7 +824,6 @@ def signal_handler(signum, frame): try: # Initialize tracer - rank = kwargs.get("rank", 0) maybe_init_worker_tracer( instrumenting_module_name="aphrodite.worker", process_kind="worker", diff --git a/aphrodite/v1/spec_decode/llm_base_proposer.py b/aphrodite/v1/spec_decode/llm_base_proposer.py index 495df1f1b7..d6a0e201f5 100644 --- a/aphrodite/v1/spec_decode/llm_base_proposer.py +++ b/aphrodite/v1/spec_decode/llm_base_proposer.py @@ -1209,14 +1209,18 @@ def _create_draft_aphrodite_config(self) -> AphroditeConfig: ), ) - if spec_cfg.kv_cache_dtype is not None: - base = replace( - base, - cache_config=replace( - base.cache_config, - cache_dtype=spec_cfg.kv_cache_dtype, + from aphrodite.config.speculative import resolve_draft_kv_cache_dtype + + base = replace( + base, + cache_config=replace( + base.cache_config, + cache_dtype=resolve_draft_kv_cache_dtype( + spec_cfg, + base.cache_config.cache_dtype, ), - ) + ), + ) return base diff --git a/aphrodite/v1/worker/block_table.py b/aphrodite/v1/worker/block_table.py index 770b033e7d..711e4366e7 100644 --- a/aphrodite/v1/worker/block_table.py +++ b/aphrodite/v1/worker/block_table.py @@ -15,6 +15,31 @@ logger = init_logger(__name__) +def get_block_table_width( + max_num_blocks: int, + block_size: int, + kernel_block_size: int | None = None, + *, + token_alignment: int | None = 128, +) -> int: + """Return the kernel-visible block-table width. + + The allocation block count is aligned before allocation blocks are split + into smaller kernel blocks. Metadata scratch buffers must use the same + width or speculative decode can read beyond their final row. + """ + if kernel_block_size is None: + kernel_block_size = block_size + if block_size % kernel_block_size != 0: + raise ValueError(f"kernel_block_size {kernel_block_size} must divide block_size {block_size}") + if token_alignment is not None: + if token_alignment <= 0: + raise ValueError("token_alignment must be positive") + block_alignment = token_alignment // np.gcd(token_alignment, block_size) + max_num_blocks = cdiv(max_num_blocks, block_alignment) * block_alignment + return max_num_blocks * block_size // kernel_block_size + + class BlockTable: def __init__( self, @@ -242,11 +267,7 @@ def __init__( f"max_num_blocks length ({len(max_num_blocks)}) must match block_sizes length ({len(block_sizes)})" ) - # Align to a multiple of (128 / block_size) as required - # by some attention backends such as TRTLLM (#39324) - max_num_blocks = [ - cdiv(n, 128 // bs) * (128 // bs) if bs <= 128 else n for n, bs in zip(max_num_blocks, block_sizes) - ] + max_num_blocks = [get_block_table_width(n, bs, bs) for n, bs in zip(max_num_blocks, block_sizes)] self.block_tables = [ BlockTable( diff --git a/aphrodite/v1/worker/gpu/attn_utils.py b/aphrodite/v1/worker/gpu/attn_utils.py index 7dba3edcf3..9324d76478 100644 --- a/aphrodite/v1/worker/gpu/attn_utils.py +++ b/aphrodite/v1/worker/gpu/attn_utils.py @@ -109,8 +109,8 @@ def init_attn_backend( layer_type = cast(type[Any], AttentionLayerBase) attn_layers = get_layers_from_aphrodite_config(aphrodite_config, layer_type, layer_names) - group_map: dict[tuple[tuple[str, str], KVCacheSpec, int], AttentionGroup] = {} - group_order: list[tuple[tuple[str, str], KVCacheSpec, int]] = [] + group_map: dict[tuple[tuple[str, str], KVCacheSpec, int, str | None], AttentionGroup] = {} + group_order: list[tuple[tuple[str, str], KVCacheSpec, int, str | None]] = [] for layer_name in layer_names: attn_backend = attn_layers[layer_name].get_attn_backend() @@ -123,9 +123,16 @@ def init_attn_backend( # counts (e.g. a spec-decode draft head and its target) get separate # metadata builders. num_heads_q = getattr(attn_layers[layer_name], "num_heads", 0) - key = (attn_backend.full_cls_name(), layer_kv_cache_spec, num_heads_q) + layer_cache_dtype = getattr(attn_layers[layer_name], "kv_cache_dtype", None) + key = (attn_backend.full_cls_name(), layer_kv_cache_spec, num_heads_q, layer_cache_dtype) if key not in group_map: - group_map[key] = AttentionGroup(attn_backend, [layer_name], layer_kv_cache_spec, kv_cache_group_id) + group_map[key] = AttentionGroup( + attn_backend, + [layer_name], + layer_kv_cache_spec, + kv_cache_group_id, + cache_dtype=layer_cache_dtype, + ) group_order.append(key) else: group_map[key].layer_names.append(layer_name) diff --git a/aphrodite/v1/worker/gpu/model_runner.py b/aphrodite/v1/worker/gpu/model_runner.py index 60e1914d19..765fadfad3 100644 --- a/aphrodite/v1/worker/gpu/model_runner.py +++ b/aphrodite/v1/worker/gpu/model_runner.py @@ -496,6 +496,8 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: self.input_buffers, self.attn_groups, ) + if hasattr(self.speculator, "set_num_cached_tokens"): + self.speculator.set_num_cached_tokens(self.req_states.num_cached_tokens.gpu) if self.speculator is not None: # After set_attn, so the speculator can size its cudagraph mode # to its own attention support. diff --git a/aphrodite/v1/worker/gpu/spec_decode/dflash/speculator.py b/aphrodite/v1/worker/gpu/spec_decode/dflash/speculator.py index 9c016e401c..d38aef8f39 100644 --- a/aphrodite/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/aphrodite/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import copy from collections.abc import Mapping -from typing import Any +from typing import Any, cast import numpy as np import torch @@ -65,6 +66,7 @@ def __init__(self, aphrodite_config: AphroditeConfig, device: torch.device): # prepare_dflash_inputs, and processed by the model's # precompute_and_store_context_kv method. NOT captured by CUDA graphs. self.context_positions = torch.zeros(self.max_num_tokens, dtype=torch.int64, device=device) + self.num_cached_tokens = torch.zeros(self.max_num_reqs, dtype=torch.int32, device=device) # Per-mask-token sampling buffers. Flattened from (num_reqs, num_spec_tokens). max_num_sampled_tokens = self.max_num_reqs * self.num_speculative_steps @@ -83,15 +85,31 @@ def __init__(self, aphrodite_config: AphroditeConfig, device: torch.device): @property def attn_aphrodite_config(self) -> AphroditeConfig: # The draft's attention differs from the target's in causality. - return replace( - self.aphrodite_config, - attention_config=replace( - self.aphrodite_config.attention_config, - use_non_causal=self.requires_non_causal, - ), + from aphrodite.v1.worker.gpu.spec_decode.dflash.utils import ( + resolve_dflash_attention_backend, + ) + + config = copy.copy(self.aphrodite_config) + config.attention_config = replace( + self.aphrodite_config.attention_config, + use_non_causal=self.requires_non_causal, + backend=resolve_dflash_attention_backend(self.aphrodite_config), ) + return config def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: + from aphrodite.v1.worker.gpu.spec_decode.dflash.utils import ( + dflash_needs_dcp_kv_head_replication, + ) + + if dflash_needs_dcp_kv_head_replication(self.aphrodite_config): + logger.warning( + "%s with DCP-replicated draft KV heads runs the draft eagerly; " + "full CUDA graph replay does not support the decoupled cache-head layout.", + self._speculator_name, + ) + cudagraph_mode = CUDAGraphMode.NONE + wants_full = cudagraph_mode.decode_mode() == CUDAGraphMode.FULL supports_full = self.attn_cg_support.min_cg_support.value >= AttentionCGSupport.UNIFORM_BATCH.value if wants_full and not supports_full: @@ -139,6 +157,10 @@ def load_draft_model( ) -> nn.Module: return load_dflash_model(target_model, self.aphrodite_config) + def set_num_cached_tokens(self, num_cached_tokens: torch.Tensor) -> None: + """Register cache-restored token counts indexed by request-state slot.""" + self.num_cached_tokens = num_cached_tokens + def set_attn( self, model_state: ModelState, @@ -155,6 +177,23 @@ def set_attn( target_attn_groups, ) + # FlashAttention's AOT scheduler reads these values from the config, + # which otherwise still describes the target model. Match them to the + # actual draft tensors so target/draft GQA differences cannot select + # incompatible scheduler metadata. + parallel_config = self.aphrodite_config.parallel_config + draft_heads_q = self.draft_model_config.get_num_attention_heads(parallel_config) + draft_heads_kv = self.draft_model_config.get_num_kv_heads(parallel_config) + draft_head_dim = self.draft_model_config.get_head_size() + for groups in self.attn_groups: + for group in groups: + builder = group.get_metadata_builder(0) + if hasattr(builder, "num_heads_q"): + draft_builder = cast(Any, builder) + draft_builder.num_heads_q = draft_heads_q + draft_builder.num_heads_kv = draft_heads_kv + draft_builder.headdim = draft_head_dim + self.draft_kv_cache_group_ids = [gid for gid, g in enumerate(self.attn_groups) if g] assert self.draft_kv_cache_group_ids, "No draft attention groups found." self.draft_kv_cache_group_id = self.draft_kv_cache_group_ids[0] @@ -358,15 +397,33 @@ def propose( seeds, self.block_tables.input_block_tables[gid], self.block_tables.kernel_block_sizes[gid], + self.num_cached_tokens, self.parallel_drafting_token_id, self.num_query_per_req, self.num_speculative_steps, self.max_num_reqs, self.max_num_tokens, self.max_model_len, + self.block_tables.cp_size, + self.block_tables.cp_rank, + self.block_tables.cp_interleave, self.sample_from_anchor, ) + # Cache-restored target tokens have no corresponding draft KV. Shift + # their whole DCP-virtual blocks out of the draft's view after slot + # mappings have been computed from the original table. + if not dummy_run: + for gid in self.draft_kv_cache_group_ids: + shift_draft_block_tables( + self.block_tables.input_block_tables[gid], + input_batch.idx_mapping, + self.num_cached_tokens, + self.input_buffers.seq_lens, + self.block_tables.kernel_block_sizes[gid], + self.block_tables.cp_size, + ) + # Pre-insert context K/V into the cache. Runs eagerly outside the captured graph # because the context shape varies per step. During dummy runs the block tables # are placeholders, so we skip the cache write to avoid clobbering real entries. @@ -432,6 +489,34 @@ def propose( return self.draft_tokens[:num_reqs] +@triton.jit +def _pos_to_slot( + pos, + block_table_row_ptr, + block_table_stride, + mask, + block_size, + cp_rank, + CP_SIZE: tl.constexpr, + CP_INTERLEAVE: tl.constexpr, + PAD_SLOT_ID: tl.constexpr, + FORCE_PAD_UNDER_CP: tl.constexpr = False, +): + """Map a global position to this DCP rank's interleaved cache slot.""" + block_num = pos // (block_size * CP_SIZE) + block_off = pos % (block_size * CP_SIZE) + block_num = tl.minimum(block_num, block_table_stride - 1) + block_id = tl.load(block_table_row_ptr + block_num, mask=mask, other=0).to(tl.int64) + if CP_SIZE == 1: + return block_id * block_size + block_off + if FORCE_PAD_UNDER_CP: + return block_id * 0 + PAD_SLOT_ID + is_local = block_off // CP_INTERLEAVE % CP_SIZE == cp_rank + local_off = block_off // (CP_INTERLEAVE * CP_SIZE) * CP_INTERLEAVE + block_off % CP_INTERLEAVE + slot = block_id * block_size + local_off + return tl.where(is_local & mask, slot, PAD_SLOT_ID) + + @triton.jit def _prepare_dflash_inputs_kernel( # Outputs @@ -461,6 +546,7 @@ def _prepare_dflash_inputs_kernel( # Block table for slot mapping lookup. block_table_ptr, block_table_stride, + num_cached_tokens_ptr, # Scalars parallel_drafting_token_id, block_size, @@ -469,6 +555,9 @@ def _prepare_dflash_inputs_kernel( max_num_reqs, max_num_tokens, max_model_len, + cp_rank, + CP_SIZE: tl.constexpr, + CP_INTERLEAVE: tl.constexpr, SAMPLE_FROM_ANCHOR: tl.constexpr, PAD_SLOT_ID: tl.constexpr, BLOCK_SIZE: tl.constexpr, @@ -503,14 +592,17 @@ def _prepare_dflash_inputs_kernel( # --- Context positions / slots --- ctx_pos_idx = ctx_start + tl.where(is_ctx, j, 0) ctx_pos = tl.load(target_positions_ptr + ctx_pos_idx, mask=is_ctx, other=0) - ctx_block_num = ctx_pos // block_size - ctx_block_num = tl.minimum(ctx_block_num, block_table_stride - 1) - ctx_block_id = tl.load( - block_table_ptr + req_idx * block_table_stride + ctx_block_num, - mask=is_ctx, - other=0, - ).to(tl.int64) - ctx_slot = ctx_block_id * block_size + (ctx_pos % block_size) + ctx_slot = _pos_to_slot( + ctx_pos, + block_table_ptr + req_idx * block_table_stride, + block_table_stride, + is_ctx, + block_size, + cp_rank, + CP_SIZE, + CP_INTERLEAVE, + PAD_SLOT_ID, + ) tl.store(out_context_positions_ptr + ctx_start + j, ctx_pos, mask=is_ctx) tl.store(out_context_slot_mapping_ptr + ctx_start + j, ctx_slot, mask=is_ctx) @@ -520,14 +612,18 @@ def _prepare_dflash_inputs_kernel( is_bonus = is_query & (query_off == 0) input_id = tl.where(is_bonus, bonus_token, parallel_drafting_token_id) - q_block_num = query_pos // block_size - q_block_num = tl.minimum(q_block_num, block_table_stride - 1) - q_block_id = tl.load( - block_table_ptr + req_idx * block_table_stride + q_block_num, - mask=is_query, - other=0, - ).to(tl.int64) - q_slot = q_block_id * block_size + (query_pos % block_size) + q_slot = _pos_to_slot( + query_pos, + block_table_ptr + req_idx * block_table_stride, + block_table_stride, + is_query, + block_size, + cp_rank, + CP_SIZE, + CP_INTERLEAVE, + PAD_SLOT_ID, + FORCE_PAD_UNDER_CP=True, + ) tl.store(out_input_ids_ptr + query_idx, input_id, mask=is_query) clamped_query_pos = tl.minimum(query_pos, max_model_len - 1) @@ -549,10 +645,19 @@ def _prepare_dflash_inputs_kernel( if block_idx == 0: tl.store(out_query_start_loc_ptr + req_idx, query_base) - # seq_lens is the absolute sequence length the draft attention - # reads up to (context + query), not just the count of accepted - # tokens this step. - tl.store(out_seq_lens_ptr + req_idx, last_valid_pos + 1 + num_query_per_req) + # Hide whole DCP-virtual blocks restored without a target forward; + # those slots never received draft KV. Keep global sequence lengths, + # because the dense attention metadata builder localizes them for DCP. + num_cached = tl.load(num_cached_tokens_ptr + req_state_idx) + virtual_block_size = block_size * CP_SIZE + num_shifted_tokens = (num_cached // virtual_block_size) * virtual_block_size + tl.store( + out_seq_lens_ptr + req_idx, + tl.maximum( + last_valid_pos + 1 + num_query_per_req - num_shifted_tokens, + num_query_per_req, + ), + ) # Copy sampling state. tl.store( out_temperature_ptr + req_state_idx, @@ -618,12 +723,16 @@ def prepare_dflash_inputs( # [max_num_reqs, max_num_blocks] block_table: torch.Tensor, block_size: int, + num_cached_tokens: torch.Tensor, parallel_drafting_token_id: int, num_query_per_req: int, num_speculative_steps: int, max_num_reqs: int, max_num_tokens: int, max_model_len: int, + cp_size: int = 1, + cp_rank: int = 0, + cp_interleave: int = 1, sample_from_anchor: bool = False, ) -> None: num_reqs = input_batch.num_reqs @@ -658,6 +767,7 @@ def prepare_dflash_inputs( input_seeds, block_table, block_table.stride(0), + num_cached_tokens, parallel_drafting_token_id, block_size, num_query_per_req, @@ -665,7 +775,59 @@ def prepare_dflash_inputs( max_num_reqs, max_num_tokens, max_model_len, + cp_rank, + CP_SIZE=cp_size, + CP_INTERLEAVE=cp_interleave, SAMPLE_FROM_ANCHOR=sample_from_anchor, PAD_SLOT_ID=PAD_SLOT_ID, BLOCK_SIZE=BLOCK_SIZE, ) + + +@triton.jit +def _shift_draft_block_tables_kernel( + block_table_ptr, + block_table_stride, + idx_mapping_ptr, + num_cached_tokens_ptr, + seq_lens_ptr, + block_size, + CP_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + req_idx = tl.program_id(0) + req_state_idx = tl.load(idx_mapping_ptr + req_idx) + virtual_block_size = block_size * CP_SIZE + shift = tl.load(num_cached_tokens_ptr + req_state_idx) // virtual_block_size + if shift == 0: + return + row_ptr = block_table_ptr + req_idx.to(tl.int64) * block_table_stride + seq_len = tl.load(seq_lens_ptr + req_idx) + num_needed = (seq_len + virtual_block_size - 1) // virtual_block_size + num_remaining = tl.minimum(block_table_stride - shift, num_needed) + for i in tl.range(0, num_remaining, BLOCK_SIZE): + offset = i + tl.arange(0, BLOCK_SIZE) + mask = offset < num_remaining + block_ids = tl.load(row_ptr + offset + shift, mask=mask, other=0) + tl.store(row_ptr + offset, block_ids, mask=mask) + + +def shift_draft_block_tables( + block_table: torch.Tensor, + idx_mapping: torch.Tensor, + num_cached_tokens: torch.Tensor, + seq_lens: torch.Tensor, + block_size: int, + cp_size: int = 1, +) -> None: + """Hide cache-restored virtual blocks that contain no draft KV.""" + _shift_draft_block_tables_kernel[(idx_mapping.shape[0],)]( + block_table, + block_table.stride(0), + idx_mapping, + num_cached_tokens, + seq_lens, + block_size, + CP_SIZE=cp_size, + BLOCK_SIZE=1024, + ) diff --git a/aphrodite/v1/worker/gpu/spec_decode/dflash/utils.py b/aphrodite/v1/worker/gpu/spec_decode/dflash/utils.py index 2c9ea8df76..7a012eb0b7 100644 --- a/aphrodite/v1/worker/gpu/spec_decode/dflash/utils.py +++ b/aphrodite/v1/worker/gpu/spec_decode/dflash/utils.py @@ -3,13 +3,93 @@ import torch.nn as nn from aphrodite.config import AphroditeConfig, replace +from aphrodite.config.speculative import resolve_draft_kv_cache_dtype from aphrodite.distributed.parallel_state import get_pp_group +from aphrodite.logger import init_logger from aphrodite.model_executor.model_loader import get_model from aphrodite.v1.worker.gpu.spec_decode.eagle.utils import ( _should_share, get_target_lm_head, ) +logger = init_logger(__name__) + +_DCP_KV_REPLICATION_ARCHS = {"DFlashDraftModel", "Qwen3DSparkModel"} + + +def resolve_dflash_attention_backend(aphrodite_config: AphroditeConfig): + """Select the DCP-capable backend for a KV-head-replicated draft.""" + speculative_config = aphrodite_config.speculative_config + assert speculative_config is not None + backend = speculative_config.attention_backend + needs_head_replication = dflash_needs_dcp_kv_head_replication(aphrodite_config) + if not needs_head_replication: + return backend + + from aphrodite.v1.attention.backends.registry import AttentionBackendEnum + + if backend is None: + return AttentionBackendEnum.FLASH_ATTN + if backend != AttentionBackendEnum.FLASH_ATTN: + raise ValueError( + "DFlash/DSpark with DCP-replicated draft KV heads requires " + f"attention_backend=FLASH_ATTN, got {backend.name}." + ) + return backend + + +def dflash_needs_dcp_kv_head_replication(aphrodite_config: AphroditeConfig) -> bool: + speculative_config = aphrodite_config.speculative_config + assert speculative_config is not None + parallel_config = aphrodite_config.parallel_config + dcp_size = parallel_config.decode_context_parallel_size + if dcp_size == 1: + return False + + draft_model_config = speculative_config.draft_model_config + total_kv_heads = draft_model_config.get_total_num_kv_heads() + tp_size = parallel_config.tensor_parallel_size + needs_replication = total_kv_heads >= tp_size + if needs_replication and not (set(draft_model_config.architectures or []) & _DCP_KV_REPLICATION_ARCHS): + raise NotImplementedError( + "DFlash/DSpark with sharded draft KV heads under DCP requires a " + "draft architecture that replicates its KV cache across the DCP " + "group, but got " + f"{draft_model_config.architectures}." + ) + if not needs_replication: + total_q_heads = draft_model_config.model_arch_config.total_num_attention_heads + kv_replicated_across_dcp = ( + tp_size > total_kv_heads + and dcp_size <= tp_size // total_kv_heads + and (total_q_heads // total_kv_heads) % dcp_size == 0 + ) + if not kv_replicated_across_dcp: + raise NotImplementedError( + "DFlash/DSpark under DCP requires draft KV heads to be fully " + "sharded or replicated across the DCP group; got " + f"q_heads={total_q_heads}, kv_heads={total_kv_heads}, " + f"tp={tp_size}, dcp={dcp_size}." + ) + return needs_replication + + +def resolve_dflash_cache_dtype(aphrodite_config: AphroditeConfig): + """Use a non-quantized cache when DCP replicates draft KV heads.""" + if dflash_needs_dcp_kv_head_replication(aphrodite_config): + logger.warning_once( + "DFlash/DSpark with DCP requires replicated draft KV heads; " + "using the draft model dtype for its KV cache. The target KV " + "cache dtype is unchanged." + ) + return "auto" + speculative_config = aphrodite_config.speculative_config + assert speculative_config is not None + return resolve_draft_kv_cache_dtype( + speculative_config, + aphrodite_config.cache_config.cache_dtype, + ) + def load_dflash_model(target_model: nn.Module, aphrodite_config: AphroditeConfig) -> nn.Module: from aphrodite.compilation.backends import set_model_tag @@ -25,15 +105,11 @@ def load_dflash_model(target_model: nn.Module, aphrodite_config: AphroditeConfig attention_config=replace( aphrodite_config.attention_config, use_non_causal=dflash_has_any_non_causal(draft_model_config.hf_config), - backend=speculative_config.attention_backend, + backend=resolve_dflash_attention_backend(aphrodite_config), ), - cache_config=( - replace( - aphrodite_config.cache_config, - cache_dtype=speculative_config.kv_cache_dtype, - ) - if speculative_config.kv_cache_dtype is not None - else aphrodite_config.cache_config + cache_config=replace( + aphrodite_config.cache_config, + cache_dtype=resolve_dflash_cache_dtype(aphrodite_config), ), ) with set_model_tag("dflash_head"): diff --git a/aphrodite/v1/worker/gpu/spec_decode/dspark/utils.py b/aphrodite/v1/worker/gpu/spec_decode/dspark/utils.py index c04405936e..26af72657b 100644 --- a/aphrodite/v1/worker/gpu/spec_decode/dspark/utils.py +++ b/aphrodite/v1/worker/gpu/spec_decode/dspark/utils.py @@ -6,6 +6,10 @@ from aphrodite.config import AphroditeConfig, replace from aphrodite.distributed.parallel_state import get_pp_group from aphrodite.model_executor.model_loader import get_model +from aphrodite.v1.worker.gpu.spec_decode.dflash.utils import ( + resolve_dflash_attention_backend, + resolve_dflash_cache_dtype, +) from aphrodite.v1.worker.gpu.spec_decode.eagle.utils import ( _should_share, get_target_lm_head, @@ -25,15 +29,11 @@ def load_dspark_model(target_model: nn.Module, aphrodite_config: AphroditeConfig attention_config=replace( aphrodite_config.attention_config, use_non_causal=dflash_has_any_non_causal(draft_model_config.hf_config), - backend=speculative_config.attention_backend, + backend=resolve_dflash_attention_backend(aphrodite_config), ), - cache_config=( - replace( - aphrodite_config.cache_config, - cache_dtype=speculative_config.kv_cache_dtype, - ) - if speculative_config.kv_cache_dtype is not None - else aphrodite_config.cache_config + cache_config=replace( + aphrodite_config.cache_config, + cache_dtype=resolve_dflash_cache_dtype(aphrodite_config), ), ) diff --git a/aphrodite/v1/worker/gpu/spec_decode/eagle/utils.py b/aphrodite/v1/worker/gpu/spec_decode/eagle/utils.py index 7262a07e30..6af3e8b68e 100644 --- a/aphrodite/v1/worker/gpu/spec_decode/eagle/utils.py +++ b/aphrodite/v1/worker/gpu/spec_decode/eagle/utils.py @@ -4,6 +4,7 @@ import torch.nn as nn from aphrodite.config import AphroditeConfig, replace +from aphrodite.config.speculative import resolve_draft_kv_cache_dtype from aphrodite.distributed.parallel_state import get_pp_group from aphrodite.lora.layers.base import BaseLayerWithLoRA from aphrodite.model_executor.model_loader import get_model @@ -37,14 +38,16 @@ def load_eagle_model(target_model: nn.Module, aphrodite_config: AphroditeConfig) speculative_config = aphrodite_config.speculative_config assert speculative_config is not None draft_model_config = speculative_config.draft_model_config - if speculative_config.kv_cache_dtype is not None: - aphrodite_config = replace( - aphrodite_config, - cache_config=replace( - aphrodite_config.cache_config, - cache_dtype=speculative_config.kv_cache_dtype, + aphrodite_config = replace( + aphrodite_config, + cache_config=replace( + aphrodite_config.cache_config, + cache_dtype=resolve_draft_kv_cache_dtype( + speculative_config, + aphrodite_config.cache_config.cache_dtype, ), - ) + ), + ) with set_model_tag("eagle_head"): eagle_model = get_model(aphrodite_config=aphrodite_config, model_config=draft_model_config) diff --git a/aphrodite/v1/worker/gpu/states.py b/aphrodite/v1/worker/gpu/states.py index 14add07e63..ddfb1332de 100644 --- a/aphrodite/v1/worker/gpu/states.py +++ b/aphrodite/v1/worker/gpu/states.py @@ -56,6 +56,10 @@ def __init__( self.num_computed_tokens = StagedWriteTensor(self.max_num_reqs, dtype=torch.int32, device=device) # Optimistic CPU mirror of num_computed_tokens (upper bound on GPU value). self.num_computed_tokens_np = np.zeros(self.max_num_reqs, dtype=np.int32) + # Tokens restored at the request's latest admission. DFlash and + # DSpark cannot derive draft KV for these tokens from target hidden + # states because the target forward did not process them. + self.num_cached_tokens = StagedWriteTensor(self.max_num_reqs, dtype=torch.int32, device=device) # Last sampled tokens. self.last_sampled_tokens = torch.zeros(self.max_num_reqs, 1, dtype=torch.int64, device=device) @@ -105,6 +109,7 @@ def add_request( self.num_computed_prefill_tokens[req_idx] = num_computed_tokens self.num_computed_tokens_np[req_idx] = num_computed_tokens self.num_computed_tokens.stage_write_elem(req_idx, num_computed_tokens) + self.num_cached_tokens.stage_write_elem(req_idx, num_computed_tokens) self.draft_tokens[req_idx].zero_() @@ -114,6 +119,7 @@ def apply_staged_writes(self) -> None: self.total_len.apply_write() self.all_token_ids.apply_write() self.num_computed_tokens.apply_write() + self.num_cached_tokens.apply_write() def remove_request(self, req_id: str) -> int | None: """Return the freed slot index, or None if the request was not found.""" diff --git a/aphrodite/v1/worker/utils.py b/aphrodite/v1/worker/utils.py index 60c6d1baf7..0d11c95746 100644 --- a/aphrodite/v1/worker/utils.py +++ b/aphrodite/v1/worker/utils.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import copy import math from collections import defaultdict from collections.abc import Iterable, Sequence @@ -210,6 +211,9 @@ class AttentionGroup: layer_names: list[str] kv_cache_spec: KVCacheSpec kv_cache_group_id: int + # A speculative draft can override a target-only cache layout while both + # models share the top-level AphroditeConfig. + cache_dtype: str | None = None # When ubatching is enabled we will have a metadata builder for each ubatch # so that if they use internal persistent buffers for cudagraphs, and they # won't have to worry about conflicting with the other ubatches. @@ -222,6 +226,11 @@ def create_metadata_builders( kernel_block_size: int | None = None, num_metadata_builders: int = 1, ): + builder_config = aphrodite_config + if self.cache_dtype is not None and self.cache_dtype != aphrodite_config.cache_config.cache_dtype: + builder_config = copy.copy(aphrodite_config) + builder_config.cache_config = copy.copy(aphrodite_config.cache_config) + builder_config.cache_config.cache_dtype = self.cache_dtype kv_cache_spec_builder = ( self.kv_cache_spec.copy_with_new_block_size(kernel_block_size) if kernel_block_size is not None @@ -231,7 +240,7 @@ def create_metadata_builders( self.backend.get_builder_cls()( kv_cache_spec_builder, self.layer_names, - aphrodite_config, + builder_config, device, ) for _ in range(num_metadata_builders) diff --git a/tests/config/test_speculative_draft_kv_cache.py b/tests/config/test_speculative_draft_kv_cache.py new file mode 100644 index 0000000000..f3299e5a1a --- /dev/null +++ b/tests/config/test_speculative_draft_kv_cache.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +from aphrodite.config.speculative import resolve_draft_kv_cache_dtype + + +def _spec_config(*, draft_uses_mla: bool, kv_cache_dtype: str | None = None): + return SimpleNamespace( + kv_cache_dtype=kv_cache_dtype, + draft_model_config=SimpleNamespace(use_mla=draft_uses_mla), + ) + + +def test_non_mla_draft_does_not_inherit_packed_mla_cache(): + spec_config = _spec_config(draft_uses_mla=False) + + assert resolve_draft_kv_cache_dtype(spec_config, "fp8_ds_mla") == "fp8_e4m3" + + +def test_mla_draft_inherits_packed_mla_cache(): + spec_config = _spec_config(draft_uses_mla=True) + + assert resolve_draft_kv_cache_dtype(spec_config, "fp8_ds_mla") == "fp8_ds_mla" + + +def test_explicit_draft_cache_dtype_takes_precedence(): + spec_config = _spec_config(draft_uses_mla=False, kv_cache_dtype="bfloat16") + + assert resolve_draft_kv_cache_dtype(spec_config, "fp8_ds_mla") == "bfloat16" diff --git a/tests/kernels/attention/test_sm120_dsa_indexer.py b/tests/kernels/attention/test_sm120_dsa_indexer.py new file mode 100644 index 0000000000..65bcac398f --- /dev/null +++ b/tests/kernels/attention/test_sm120_dsa_indexer.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Correctness tests for the native SM120 sparse-indexer score kernels.""" + +import pytest +import torch + +from aphrodite.model_executor.kernels.attention.dsa.sm120_indexer import ( + sm120_fp8_mqa_logits, + sm120_fp8_paged_mqa_logits, +) +from aphrodite.platforms import current_platform + +PAGE_SIZE = 64 +HEAD_DIM = 128 +NUM_HEADS = 64 + + +def _sm120_available() -> bool: + if not current_platform.is_cuda(): + return False + try: + return current_platform.is_device_capability_family(120) + except RuntimeError: + return False + + +pytestmark = pytest.mark.skipif(not _sm120_available(), reason="requires an SM120 GPU") + + +def _reference( + q: torch.Tensor, + k: torch.Tensor, + scales: torch.Tensor, + weights: torch.Tensor, +) -> torch.Tensor: + scores = torch.einsum("mhd,nd->mhn", q.float(), k.float()) + return (scores.mul(scales[None, None, :]).relu() * weights[:, :, None]).sum(1) + + +def test_sm120_fp8_mqa_logits() -> None: + torch.manual_seed(0) + num_rows, num_tokens = 5, 193 + q = torch.randn((num_rows, NUM_HEADS, HEAD_DIM), device="cuda", dtype=torch.bfloat16).to(torch.float8_e4m3fn) + k = torch.randn((num_tokens, HEAD_DIM), device="cuda", dtype=torch.bfloat16).to(torch.float8_e4m3fn) + scales = torch.rand(num_tokens, device="cuda") * 0.02 + weights = torch.randn((num_rows, NUM_HEADS), device="cuda") + starts = torch.tensor([0, 3, 61, 100, 192], dtype=torch.int32, device="cuda") + ends = torch.tensor([193, 129, 162, 193, 193], dtype=torch.int32, device="cuda") + + actual = sm120_fp8_mqa_logits(q, k, scales, weights, starts, ends) + expected = _reference(q, k, scales, weights) + positions = torch.arange(num_tokens, device="cuda")[None, :] + valid = (positions >= starts[:, None]) & (positions < ends[:, None]) + expected.masked_fill_(~valid, -torch.inf) + + torch.testing.assert_close(actual[valid], expected[valid], rtol=2e-2, atol=2e-2) + assert torch.equal(torch.isneginf(actual), torch.isneginf(expected)) + + +@pytest.mark.parametrize("next_n", [1, 2, 7]) +def test_sm120_fp8_paged_mqa_logits(next_n: int) -> None: + torch.manual_seed(1) + batch_size, num_pages = 2, 5 + max_model_len = 192 + q = torch.randn( + (batch_size, next_n, NUM_HEADS, HEAD_DIM), + device="cuda", + dtype=torch.bfloat16, + ).to(torch.float8_e4m3fn) + keys = torch.randn((num_pages, PAGE_SIZE, HEAD_DIM), device="cuda", dtype=torch.bfloat16).to(torch.float8_e4m3fn) + scales = torch.rand((num_pages, PAGE_SIZE), device="cuda") * 0.02 + cache = torch.empty( + (num_pages, PAGE_SIZE, 1, HEAD_DIM + 4), + dtype=torch.uint8, + device="cuda", + ) + # Match indexer_k_quant_and_cache's planar page layout: all key rows, + # followed by all per-token scales. + page_bytes = cache.view(num_pages, -1) + page_bytes[:, : PAGE_SIZE * HEAD_DIM].copy_(keys.view(torch.uint8).reshape(num_pages, -1)) + page_bytes[:, PAGE_SIZE * HEAD_DIM :].copy_(scales.view(torch.uint8).reshape(num_pages, -1)) + weights = torch.randn((batch_size * next_n, NUM_HEADS), device="cuda") + context_lens = torch.tensor( + [[129 + token for token in range(next_n)], [70 + token for token in range(next_n)]], + dtype=torch.int32, + device="cuda", + ) + block_tables = torch.tensor([[2, 0, 4], [3, 1, 0]], dtype=torch.int32, device="cuda") + + actual = sm120_fp8_paged_mqa_logits(q, cache, weights, context_lens, block_tables, max_model_len) + expected = torch.full_like(actual, -torch.inf) + for batch in range(batch_size): + pages = block_tables[batch].long() + logical_k = keys[pages].reshape(-1, HEAD_DIM) + logical_scales = scales[pages].reshape(-1) + for token in range(next_n): + row = batch * next_n + token + length = int(context_lens[batch, token]) + expected[row, :length] = _reference( + q[batch, token].unsqueeze(0), + logical_k[:length], + logical_scales[:length], + weights[row].unsqueeze(0), + )[0] + + finite = torch.isfinite(expected) + torch.testing.assert_close(actual[finite], expected[finite], rtol=2e-2, atol=2e-2) + assert torch.equal(torch.isneginf(actual), torch.isneginf(expected)) diff --git a/tests/kernels/test_fused_deepseek_v32_norm_rope.py b/tests/kernels/test_fused_deepseek_v32_norm_rope.py index 2a29d35bab..5da843bdc2 100644 --- a/tests/kernels/test_fused_deepseek_v32_norm_rope.py +++ b/tests/kernels/test_fused_deepseek_v32_norm_rope.py @@ -229,6 +229,72 @@ def test_fused_norm_rope(num_tokens: int, index_interleave: bool, mla_fp8: bool) assert (topk == -1).all(), "topk buffer not cleared on indexer layer" +def test_fused_norm_rope_normalizes_query_without_local_cache_slots(): + """DCP non-owner ranks still produce their query shard.""" + torch.manual_seed(7) + dev = "cuda" + num_tokens = 4 + max_pos = 16 + pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) + + q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16) + kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16) + k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16) + qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16) + kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16) + ik = torch.randn(num_tokens, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16) + ikw = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32) + ikb = torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32) + cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev) + + mla_cache = torch.zeros( + 1, + max_pos, + KV_LORA + ROPE_DIM, + device=dev, + dtype=torch.bfloat16, + ) + idx_row = INDEX_HEAD_DIM + INDEX_HEAD_DIM // 128 * 4 + idx_cache = torch.zeros(1, max_pos, idx_row, device=dev, dtype=torch.uint8) + no_local_slots = torch.full( + (num_tokens,), + -1, + device=dev, + dtype=torch.int64, + ) + topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32) + + q_out = K.fused_norm_rope( + pos, + q_c, + qw, + EPS, + kv_c, + kvw, + EPS, + k_pe, + cos_sin, + ik, + ikw, + ikb, + EPS, + cos_sin, + topk, + slot_mapping=no_local_slots, + indexer_k_cache=idx_cache, + mla_kv_cache=mla_cache, + mla_kv_cache_dtype="auto", + mla_k_scale=None, + has_indexer=True, + index_rope_interleave=True, + ) + + assert_bf16(q_out, rms_norm(q_c, qw), "q_c rmsnorm without local cache slots") + assert not mla_cache.any(), "non-owner rank wrote the MLA KV cache" + assert not idx_cache.any(), "non-owner rank wrote the indexer KV cache" + assert (topk == -1).all(), "topk buffer not cleared on non-owner rank" + + @pytest.mark.parametrize("num_tokens", [1, 17, 512]) def test_fused_norm_rope_no_indexer(num_tokens: int): """Shared (no-indexer) layer: q + kv/MLA only; top-k buffer untouched.""" diff --git a/tests/v1/attention/test_flashinfer_sparse_mla_workspace.py b/tests/v1/attention/test_flashinfer_sparse_mla_workspace.py new file mode 100644 index 0000000000..2d26440c37 --- /dev/null +++ b/tests/v1/attention/test_flashinfer_sparse_mla_workspace.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU-only tests for FlashInfer sparse MLA workspace sizing.""" + +from aphrodite.v1.attention.backends.mla.flashinfer_mla_sparse import ( + _DEFAULT_WORKSPACE_BUFFER_SIZE, + compute_trtllm_sparse_mla_workspace_bytes, +) + + +def test_dcp_workspace_covers_softmax_stats() -> None: + size = compute_trtllm_sparse_mla_workspace_bytes( + base_workspace_bytes=_DEFAULT_WORKSPACE_BUFFER_SIZE, + dcp_world_size=8, + num_heads_per_rank=8, + max_num_batched_tokens=16384, + ) + + expected_softmax_bytes = 8 * (8 * 8) * 16384 * 256 + 1024 * 1024 + assert size == _DEFAULT_WORKSPACE_BUFFER_SIZE + expected_softmax_bytes + + +def test_non_dcp_workspace_keeps_default_size() -> None: + size = compute_trtllm_sparse_mla_workspace_bytes( + base_workspace_bytes=_DEFAULT_WORKSPACE_BUFFER_SIZE, + dcp_world_size=1, + num_heads_per_rank=128, + max_num_batched_tokens=65536, + ) + + assert size == _DEFAULT_WORKSPACE_BUFFER_SIZE diff --git a/tests/v1/spec_decode/test_dflash_dcp_slot_mapping.py b/tests/v1/spec_decode/test_dflash_dcp_slot_mapping.py new file mode 100644 index 0000000000..7c0eb51207 --- /dev/null +++ b/tests/v1/spec_decode/test_dflash_dcp_slot_mapping.py @@ -0,0 +1,195 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DCP-aware slot mapping for DFlash and DSpark draft inputs.""" + +from types import SimpleNamespace + +import pytest +import torch + +from aphrodite.platforms import current_platform +from aphrodite.v1.attention.backends.utils import PAD_SLOT_ID +from aphrodite.v1.worker.gpu.input_batch import InputBuffers +from aphrodite.v1.worker.gpu.spec_decode.dflash.speculator import ( + prepare_dflash_inputs, +) + +pytest.importorskip("triton") +if not current_platform.is_cuda_alike(): + pytest.skip("CUDA required for DFlash kernel tests", allow_module_level=True) + +DEVICE = "cuda" +BLOCK_SIZE = 16 +NUM_SPECULATIVE_STEPS = 3 +NUM_QUERY_PER_REQ = 1 + NUM_SPECULATIVE_STEPS +MAX_NUM_REQS = 8 +MAX_NUM_TOKENS = 512 +MAX_NUM_BLOCKS = 64 +MAX_MODEL_LEN = 4096 + + +def _ref_slots( + positions: torch.Tensor, + block_row: torch.Tensor, + cp_size: int, + cp_rank: int, + cp_interleave: int, +) -> torch.Tensor: + virtual_block = BLOCK_SIZE * cp_size + block_indices = (positions // virtual_block).clamp(max=block_row.shape[0] - 1) + block_offsets = positions % virtual_block + block_ids = block_row[block_indices].long() + if cp_size == 1: + return block_ids * BLOCK_SIZE + block_offsets + is_local = (block_offsets // cp_interleave) % cp_size == cp_rank + local_offsets = block_offsets // (cp_interleave * cp_size) * cp_interleave + block_offsets % cp_interleave + slots = block_ids * BLOCK_SIZE + local_offsets + return torch.where(is_local, slots, torch.full_like(slots, PAD_SLOT_ID)) + + +def _run_prepare(cp_size: int, cp_rank: int, cp_interleave: int): + context_positions = [ + torch.arange(100, 120, device=DEVICE), + torch.arange(0, 1, device=DEVICE), + torch.arange(37, 100, device=DEVICE), + ] + num_rejected = torch.tensor([2, 0, 0], dtype=torch.int32, device=DEVICE) + num_sampled = torch.tensor([1, 1, 0], dtype=torch.int32, device=DEVICE) + num_reqs = len(context_positions) + + positions = torch.cat(context_positions) + num_context = torch.tensor([len(p) for p in context_positions], device=DEVICE) + query_start_loc = torch.zeros(num_reqs + 1, dtype=torch.int32, device=DEVICE) + query_start_loc[1:] = num_context.cumsum(0) + num_tokens = int(query_start_loc[-1]) + + input_batch = SimpleNamespace( + num_reqs=num_reqs, + num_scheduled_tokens=num_context.cpu().numpy(), + positions=positions, + query_start_loc=query_start_loc, + idx_mapping=torch.arange(MAX_NUM_REQS, dtype=torch.int32, device=DEVICE), + ) + input_buffers = InputBuffers(MAX_NUM_REQS, MAX_NUM_TOKENS, torch.device(DEVICE)) + generator = torch.Generator(device=DEVICE).manual_seed(42) + block_table = ( + torch.randperm( + MAX_NUM_REQS * MAX_NUM_BLOCKS, + generator=generator, + device=DEVICE, + ) + .view(MAX_NUM_REQS, MAX_NUM_BLOCKS) + .to(torch.int32) + ) + + query_slot_mapping = torch.zeros(MAX_NUM_TOKENS, dtype=torch.int64, device=DEVICE) + output_context_positions = torch.zeros(MAX_NUM_TOKENS, dtype=torch.int64, device=DEVICE) + context_slot_mapping = torch.zeros(MAX_NUM_TOKENS, dtype=torch.int64, device=DEVICE) + max_num_sampled = MAX_NUM_REQS * NUM_SPECULATIVE_STEPS + sample_indices = torch.zeros(max_num_sampled, dtype=torch.int64, device=DEVICE) + sample_pos = torch.zeros(max_num_sampled, dtype=torch.int64, device=DEVICE) + sample_idx_mapping = torch.zeros(max_num_sampled, dtype=torch.int32, device=DEVICE) + temperature = torch.zeros(MAX_NUM_REQS, dtype=torch.float32, device=DEVICE) + seeds = torch.zeros(MAX_NUM_REQS, dtype=torch.int64, device=DEVICE) + last_sampled = torch.full((MAX_NUM_REQS,), 7, dtype=torch.int64, device=DEVICE) + next_prefill_tokens = torch.full((MAX_NUM_REQS,), 11, dtype=torch.int64, device=DEVICE) + + prepare_dflash_inputs( + input_buffers, + query_slot_mapping, + output_context_positions, + context_slot_mapping, + sample_indices, + sample_pos, + sample_idx_mapping, + temperature, + seeds, + input_batch, + num_sampled, + num_rejected, + last_sampled, + next_prefill_tokens, + temperature, + seeds, + block_table, + BLOCK_SIZE, + torch.zeros(MAX_NUM_REQS, dtype=torch.int32, device=DEVICE), + parallel_drafting_token_id=1, + num_query_per_req=NUM_QUERY_PER_REQ, + num_speculative_steps=NUM_SPECULATIVE_STEPS, + max_num_reqs=MAX_NUM_REQS, + max_num_tokens=MAX_NUM_TOKENS, + max_model_len=MAX_MODEL_LEN, + cp_size=cp_size, + cp_rank=cp_rank, + cp_interleave=cp_interleave, + ) + return SimpleNamespace( + context_positions=context_positions, + num_rejected=num_rejected, + num_tokens=num_tokens, + num_reqs=num_reqs, + query_start_loc=query_start_loc, + block_table=block_table, + input_buffers=input_buffers, + query_slot_mapping=query_slot_mapping, + output_context_positions=output_context_positions, + context_slot_mapping=context_slot_mapping, + ) + + +@pytest.mark.parametrize("cp_size,cp_interleave", [(1, 1), (2, 1), (4, 1), (2, 16), (4, 8)]) +def test_dflash_slot_mapping_matches_dcp_layout(cp_size: int, cp_interleave: int): + for cp_rank in range(cp_size): + result = _run_prepare(cp_size, cp_rank, cp_interleave) + for req in range(result.num_reqs): + start = int(result.query_start_loc[req]) + end = int(result.query_start_loc[req + 1]) + context_positions = result.context_positions[req] + expected_context_slots = _ref_slots( + context_positions, + result.block_table[req], + cp_size, + cp_rank, + cp_interleave, + ) + torch.testing.assert_close( + result.context_slot_mapping[start:end], + expected_context_slots, + rtol=0, + atol=0, + ) + torch.testing.assert_close( + result.output_context_positions[start:end], + context_positions, + rtol=0, + atol=0, + ) + + query_start = req * NUM_QUERY_PER_REQ + query_slots = result.query_slot_mapping[query_start : query_start + NUM_QUERY_PER_REQ] + if cp_size == 1: + last_valid_pos = int(context_positions[-1 - int(result.num_rejected[req])]) + query_positions = torch.arange( + last_valid_pos + 1, + last_valid_pos + 1 + NUM_QUERY_PER_REQ, + device=DEVICE, + ) + expected_query_slots = _ref_slots( + query_positions, + result.block_table[req], + cp_size, + cp_rank, + cp_interleave, + ) + torch.testing.assert_close(query_slots, expected_query_slots, rtol=0, atol=0) + else: + assert (query_slots == PAD_SLOT_ID).all() + + +@pytest.mark.parametrize("cp_size,cp_interleave", [(2, 1), (4, 1), (4, 8)]) +def test_dflash_context_slots_partition_across_dcp(cp_size: int, cp_interleave: int): + per_rank = [_run_prepare(cp_size, rank, cp_interleave) for rank in range(cp_size)] + for token_idx in range(per_rank[0].num_tokens): + owners = sum(int(result.context_slot_mapping[token_idx] != PAD_SLOT_ID) for result in per_rank) + assert owners == 1 diff --git a/tests/v1/spec_decode/test_dflash_prefix_cache_masking.py b/tests/v1/spec_decode/test_dflash_prefix_cache_masking.py new file mode 100644 index 0000000000..307d388879 --- /dev/null +++ b/tests/v1/spec_decode/test_dflash_prefix_cache_masking.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DFlash and DSpark draft masking for cache-restored prefixes.""" + +import pytest +import torch + +from aphrodite.platforms import current_platform +from aphrodite.v1.worker.gpu.spec_decode.dflash.speculator import ( + shift_draft_block_tables, +) + +pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") + +DEVICE = "cuda" +BLOCK_SIZE = 16 +MAX_BLOCKS = 64 +MAX_NUM_REQS = 8 + + +def _make_block_table(num_reqs: int) -> torch.Tensor: + return ( + torch.arange(MAX_NUM_REQS * MAX_BLOCKS, dtype=torch.int32, device=DEVICE) + .view(MAX_NUM_REQS, MAX_BLOCKS)[:num_reqs] + .contiguous() + ) + + +@pytest.mark.parametrize("cp_size", [1, 2, 4]) +@pytest.mark.parametrize("cached_virtual_blocks", [0, 1, 3]) +def test_shift_single_request(cp_size: int, cached_virtual_blocks: int): + block_table = _make_block_table(1) + original = block_table.clone() + idx_mapping = torch.zeros(1, dtype=torch.int32, device=DEVICE) + num_cached_tokens = torch.full( + (MAX_NUM_REQS,), + cached_virtual_blocks * BLOCK_SIZE * cp_size, + dtype=torch.int32, + device=DEVICE, + ) + seq_lens = torch.full( + (1,), + (MAX_BLOCKS - cached_virtual_blocks) * BLOCK_SIZE * cp_size, + dtype=torch.int32, + device=DEVICE, + ) + + shift_draft_block_tables( + block_table, + idx_mapping, + num_cached_tokens, + seq_lens, + BLOCK_SIZE, + cp_size, + ) + + kept = MAX_BLOCKS - cached_virtual_blocks + torch.testing.assert_close( + block_table[0, :kept], + original[0, cached_virtual_blocks:], + ) + + +def test_shift_ignores_partial_dcp_virtual_block(): + cp_size = 4 + block_table = _make_block_table(1) + original = block_table.clone() + idx_mapping = torch.zeros(1, dtype=torch.int32, device=DEVICE) + num_cached_tokens = torch.full( + (MAX_NUM_REQS,), + BLOCK_SIZE * cp_size - 1, + dtype=torch.int32, + device=DEVICE, + ) + seq_lens = torch.full((1,), MAX_BLOCKS * BLOCK_SIZE * cp_size, dtype=torch.int32, device=DEVICE) + + shift_draft_block_tables( + block_table, + idx_mapping, + num_cached_tokens, + seq_lens, + BLOCK_SIZE, + cp_size, + ) + + torch.testing.assert_close(block_table, original) diff --git a/tests/v1/worker/test_block_table_width.py b/tests/v1/worker/test_block_table_width.py new file mode 100644 index 0000000000..3389b5b78a --- /dev/null +++ b/tests/v1/worker/test_block_table_width.py @@ -0,0 +1,21 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from aphrodite.v1.worker.block_table import get_block_table_width + + +def test_block_table_width_matches_dcp_alignment() -> None: + # 80K tokens / (64-token blocks * DCP4) needs 313 local blocks. The + # runtime table aligns this to a 128-token boundary, hence 314 entries. + assert get_block_table_width(313, 64) == 314 + + +def test_block_table_width_accounts_for_kernel_block_splitting() -> None: + assert get_block_table_width(7, 32, 16) == 16 + + +def test_block_table_width_rejects_incompatible_kernel_size() -> None: + with pytest.raises(ValueError, match="must divide"): + get_block_table_width(8, 24, 16)