diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index e6e6f01cc..2767ebdd2 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -1982,6 +1982,19 @@ class SparseAttnIndexer(CustomOp): self.max_model_len = max_model_len self.max_total_seq_len = max_total_seq_len self.topk_indices_buffer = topk_indices_buffer + if topk_scores_buffer is None and topk_indices_buffer is not None: + try: + from vllm.distributed.parallel_state import get_dcp_group + + dcp_world_size = get_dcp_group().world_size + except Exception: + dcp_world_size = 1 + if dcp_world_size > 1 and use_b12x_sparse_indexer(): + topk_scores_buffer = torch.empty( + topk_indices_buffer.shape, + dtype=torch.float32, + device=topk_indices_buffer.device, + ) self.topk_scores_buffer = topk_scores_buffer self.skip_k_cache_insert = skip_k_cache_insert self.use_fp4_cache = use_fp4_cache diff --git a/vllm/model_executor/models/deepseek_mtp.py b/vllm/model_executor/models/deepseek_mtp.py index 552d6438a..d0fdbaf36 100644 --- a/vllm/model_executor/models/deepseek_mtp.py +++ b/vllm/model_executor/models/deepseek_mtp.py @@ -346,10 +346,7 @@ class DeepSeekMultiTokenPredictorLayer(nn.Module): device=self.device, ) topk_scores_buffer = None - if ( - vllm_config.parallel_config.decode_context_parallel_size > 1 - and use_b12x_sparse_indexer() - ): + if vllm_config.parallel_config.decode_context_parallel_size > 1: topk_scores_buffer = torch.empty( vllm_config.scheduler_config.max_num_batched_tokens, topk_tokens, diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index cbdc9be96..6260ee159 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import os -from dataclasses import dataclass +from dataclasses import dataclass, replace import numpy as np import torch @@ -695,6 +695,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): use_dcp_local_kv = ( self.dcp_world_size > 1 and common_attn_metadata.dcp_local_seq_lens is not None + and not getattr(self.kv_cache_spec, "dcp_replicated", False) ) effective_dcp_world_size = self.dcp_world_size if use_dcp_local_kv else 1 effective_dcp_rank = self.dcp_rank if use_dcp_local_kv else 0 @@ -1021,6 +1022,49 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): return attn_metadata + def build_for_drafting( + self, + common_attn_metadata: CommonAttentionMetadata, + draft_index: int, + ) -> DeepseekV32IndexerMetadata: + if draft_index > 0: + num_tokens = common_attn_metadata.num_actual_tokens + num_reqs = min(common_attn_metadata.num_reqs, num_tokens) + if num_tokens > 0 and num_reqs > 0: + needs_unpadded_decode_view = ( + common_attn_metadata.batch_topology is not None + or common_attn_metadata.max_query_len > 1 + or common_attn_metadata.num_reqs != num_reqs + or int(common_attn_metadata.query_start_loc_cpu[num_reqs]) + != num_tokens + ) + if needs_unpadded_decode_view: + query_start_loc = torch.arange( + num_reqs + 1, + dtype=common_attn_metadata.query_start_loc.dtype, + device=common_attn_metadata.query_start_loc.device, + ) + query_start_loc_cpu = torch.arange( + num_reqs + 1, + dtype=common_attn_metadata.query_start_loc_cpu.dtype, + device=common_attn_metadata.query_start_loc_cpu.device, + ) + common_attn_metadata = replace( + common_attn_metadata.unpadded( + num_actual_tokens=num_tokens, + num_actual_reqs=num_reqs, + ), + query_start_loc=query_start_loc, + query_start_loc_cpu=query_start_loc_cpu, + max_query_len=1, + ) + + return self.build( + common_prefix_len=0, + common_attn_metadata=common_attn_metadata, + fast_build=True, + ) + def build_prefill_chunk_metadata( start_idx: int,