| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748 |
- --- baked-ours/models/deepseek_v4_1/attention.py 2026-09-10 20:13:20.000000000 +0200
- +++ baked/models/deepseek_v4_1/attention.py 2026-09-11 10:13:47.000000000 +0200
- @@ -42,7 +42,7 @@
- VllmConfig,
- get_current_vllm_config,
- )
- -from vllm.distributed import get_tensor_model_parallel_world_size
- +from vllm.distributed import get_dcp_group, get_tensor_model_parallel_world_size
- from vllm.forward_context import get_forward_context
- from vllm.logger import init_logger
- from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
- @@ -172,6 +172,13 @@
- # workspace allocated in _forward_prefill and is also read by the dummy-run
- # path to pre-reserve that workspace.
- PREFILL_CHUNK_SIZE: ClassVar[int] = 4
- + # DCP contract, mirroring AttentionImpl (v1/attention/backend.py): the
- + # V4.1 layer runs its own forward (no AttentionImpl), so
- + # check_attention_cp_compatibility reads these from the layer. A platform
- + # subclass that implements the per-rank partial attention + LSE combine
- + # sets can_return_lse_for_decode = True (SM120 FlashInfer).
- + can_return_lse_for_decode: ClassVar[bool] = False
- + supports_mtp_with_cp_non_trivial_interleave_size: ClassVar[bool] = False
-
- @classmethod
- @abstractmethod
- @@ -226,6 +233,22 @@
- self.n_heads = config.num_attention_heads
- assert self.n_heads % tp_size == 0
- self.n_local_heads = self.n_heads // tp_size
- + # Decode context parallelism: compressed/indexer caches are sharded over
- + # compressed states across the DCP group (a subset of the TP ranks),
- + # SWA caches are replicated; the platform subclass gathers queries and
- + # combines per-rank partial attention (DCP2-PLAN.md).
- + parallel_config = vllm_config.parallel_config
- + self.dcp_world_size = parallel_config.decode_context_parallel_size
- + self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0
- + self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size
- + self.need_to_return_lse_for_decode = (
- + self.dcp_world_size > 1 and self.can_return_lse_for_decode
- + )
- + if self.dcp_world_size > 1 and self.cp_kv_cache_interleave_size != 1:
- + raise NotImplementedError(
- + "DeepSeek-V4.1 DCP supports cp_kv_cache_interleave_size=1 only "
- + f"(got {self.cp_kv_cache_interleave_size})."
- + )
- self.q_lora_rank = config.q_lora_rank
- self.o_lora_rank = config.o_lora_rank
- self.head_dim = config.head_dim
|
