| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081 |
- --- baked-ours/v1/worker/gpu/block_table.py 2026-09-10 19:56:14.000000000 +0200
- +++ baked/v1/worker/gpu/block_table.py 2026-09-11 10:13:47.000000000 +0200
- @@ -27,6 +27,7 @@
- cp_rank: int = 0,
- cp_interleave: int = 1,
- slot_mapping_enabled: list[bool] | None = None,
- + cp_sharded: list[bool] | None = None,
- ):
- self.block_sizes = block_sizes
- self.kernel_block_sizes = kernel_block_sizes
- @@ -44,6 +45,13 @@
- slot_mapping_enabled = [True] * self.num_kv_cache_groups
- assert len(slot_mapping_enabled) == self.num_kv_cache_groups
- self._slot_mapping_enabled = slot_mapping_enabled
- + # Per group: True = positions are round-robined across the cp_size
- + # ranks (this rank stores 1/cp_size of them); False = DCP-replicated
- + # group, every rank stores every position at its plain block geometry.
- + if cp_sharded is None:
- + cp_sharded = [True] * self.num_kv_cache_groups
- + assert len(cp_sharded) == self.num_kv_cache_groups
- + self._cp_sharded = cp_sharded
-
- self.blocks_per_kv_block = [
- bs // kbs for bs, kbs in zip(block_sizes, kernel_block_sizes)
- @@ -108,6 +116,9 @@
- self.slot_mapping_enabled = torch.tensor(
- self._slot_mapping_enabled, dtype=torch.bool, device=self.device
- )
- + self.cp_sharded = torch.tensor(
- + self._cp_sharded, dtype=torch.bool, device=self.device
- + )
- self.input_block_table_ptrs = self._make_ptr_tensor(self.input_block_tables)
-
- def append_block_ids(
- @@ -211,6 +222,7 @@
- self.block_sizes_tensor,
- self.kernel_block_sizes_tensor,
- self.slot_mapping_enabled,
- + self.cp_sharded,
- slot_mappings,
- slot_mappings.stride(0),
- self.cp_rank,
- @@ -284,6 +296,7 @@
- block_sizes, # [num_kv_cache_groups]
- kernel_block_sizes, # [num_kv_cache_groups]
- slot_mapping_enabled, # [num_kv_cache_groups]
- + cp_sharded, # [num_kv_cache_groups]
- slot_mappings_ptr, # [num_kv_cache_groups, max_num_tokens]
- slot_mappings_stride,
- cp_rank,
- @@ -313,6 +326,7 @@
- kv_block_size = tl.load(block_sizes + group_id)
- kernel_block_size = tl.load(kernel_block_sizes + group_id)
- mapping_enabled = tl.load(slot_mapping_enabled + group_id)
- + group_sharded = tl.load(cp_sharded + group_id)
-
- req_state_idx = tl.load(idx_mapping + batch_idx)
- start_idx = tl.load(query_start_loc + batch_idx)
- @@ -326,15 +340,19 @@
- local_positions = positions
- is_local = True
- else:
- - # Context parallelism is used.
- + # Context parallelism is used. Sharded groups round-robin positions
- + # across the CP ranks; DCP-replicated groups (cp_sharded False)
- + # keep every position on every rank at the plain block geometry.
- virtual_block_size = kv_block_size * CP_SIZE
- virtual_block_indices = positions // virtual_block_size
- virtual_block_offsets = positions % virtual_block_size
- - is_local = virtual_block_offsets // CP_INTERLEAVE % CP_SIZE == cp_rank
- + owner_is_me = virtual_block_offsets // CP_INTERLEAVE % CP_SIZE == cp_rank
- rounds = virtual_block_offsets // (CP_INTERLEAVE * CP_SIZE)
- remainder = virtual_block_offsets % CP_INTERLEAVE
- local_offsets = rounds * CP_INTERLEAVE + remainder
- - local_positions = virtual_block_indices * kv_block_size + local_offsets
- + sharded_positions = virtual_block_indices * kv_block_size + local_offsets
- + local_positions = tl.where(group_sharded, sharded_positions, positions)
- + is_local = tl.where(group_sharded, owner_is_me, True)
-
- block_indices = tl.where(
- mapping_enabled, local_positions // kernel_block_size, 0
|
