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