| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111 |
- 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,
|
