|
@@ -0,0 +1,111 @@
|
|
|
|
|
+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,
|