| 1234567891011121314151617181920212223242526272829303132333435363738394041 |
- --- baked-ours/v1/attention/backend.py 2026-09-10 19:56:14.000000000 +0200
- +++ baked/v1/attention/backend.py 2026-09-11 10:13:47.000000000 +0200
- @@ -611,6 +611,26 @@
- self.vllm_config = vllm_config
- self.device = device
- self.kernel_block_size: int | None = None
- + # The DCP geometry this builder's group actually runs at: sharded
- + # groups use the process DCP size; DCP-replicated groups (every rank
- + # holds the full KV) behave as dcp=1 for metadata purposes.
- + from vllm.v1.kv_cache_interface import is_dcp_replicated_spec
- +
- + process_dcp_size = vllm_config.parallel_config.decode_context_parallel_size
- + self.dcp_replicated = is_dcp_replicated_spec(kv_cache_spec)
- + if process_dcp_size > 1 and not self.dcp_replicated:
- + from vllm.distributed.parallel_state import get_dcp_group
- +
- + try:
- + self.dcp_world_size = get_dcp_group().world_size
- + self.dcp_rank = get_dcp_group().rank_in_group
- + except AssertionError:
- + # DCP group not initialized (unit tests).
- + self.dcp_world_size = process_dcp_size
- + self.dcp_rank = 0
- + else:
- + self.dcp_world_size = 1
- + self.dcp_rank = 0
-
- def set_kernel_block_size(self, kernel_block_size: int) -> None:
- self.kernel_block_size = kernel_block_size
- @@ -650,10 +670,7 @@
- max_num_queries_for_spec,
- )
-
- - if (
- - self.vllm_config.parallel_config.decode_context_parallel_size > 1
- - and not supports_dcp_with_varlen
- - ):
- + if getattr(self, "dcp_world_size", 1) > 1 and not supports_dcp_with_varlen:
- self.reorder_batch_threshold = 1
-
- @abstractmethod
|
