Prevent overwriting matched prefix cache blocks with recomputed KV

megatron/core/inference/contexts/dynamic_context.py

@@ -1791,6 +1791,7 @@ def append_key_value_cache(self, layer_number: int, key: Tensor, value: Tensor)
                 padded_active_token_count=self.padded_active_token_count,
                 token_to_block_idx=self.gpu_view.token_to_block_idx,
                 token_to_local_position_within_kv_block=self.gpu_view.token_to_local_position_within_kv_block,
+                dummy_block_idx=self.kv_block_allocator.dummy_block_idx,
             )
 
         block_idx = self.gpu_view.token_to_block_idx[: self.padded_active_token_count]
@@ -3414,6 +3415,63 @@ def add_request(
         self.token_to_block_idx[
             self.active_token_count : self.active_token_count + effective_prefill_chunk_length
         ] = self.request_to_kv_block_ids[current_id][token_offset_range // self.block_size_tokens]
+        if num_matched_blocks > 0:
+            # Some tokens we are about to compute may land inside a block we matched
+            # by hash. That block already holds the correct KV for exactly these
+            # tokens and is shared with whoever cached it, so send those writes to the
+            # dummy block rather than perturb a concurrent reader's values.
+            #
+            # Only the write mapping moves. `request_to_kv_block_ids` still points at
+            # the real block, so attention reads the cached KV through the block table,
+            # and `token_to_block_idx` is rebuilt every step so nothing needs restoring.
+            #
+            # All positions below are absolute token positions within the request.
+            # A chunk resuming mid-block, with one fresh block past the match:
+            #
+            #   token      0     B     2B    3B    4B
+            #   blocks     |--0--|--1--|--2--|--3--|
+            #   matched    [=================]        blocks 0-2 were hash-matched
+            #   chunk               [==============]  tokens computed this step
+            #   redirect            [========]        the overlap -> dummy block
+            #   write                        [=====]  the rest -> real block 3
+
+            # 1. The blocks matched here sit right after the blocks this request
+            #    already owns. The partial block carried over from the previous
+            #    chunk joins them when it was itself matched earlier -- the tokens
+            #    completing it are cached too, and it is still shared with whoever
+            #    cached it. Everything before that is a prefix of matched blocks, so
+            #    starting the span at 0 is safe; the intersection below clips it.
+            matched_prefix_start_token = (
+                0
+                if req.num_matched_prefix_blocks >= already_allocated_blocks
+                else already_allocated_blocks * self.block_size_tokens
+            )
+            matched_prefix_end_token = (
+                already_allocated_blocks + num_matched_blocks
+            ) * self.block_size_tokens
+
+            # 2. This chunk computes the tokens starting at `effective_kv_offset`
+            #    (the prompt offset left over after skipping the cached prefix).
+            chunk_start_token = effective_kv_offset
+            chunk_end_token = effective_kv_offset + effective_prefill_chunk_length
+
+            # 3. The redundant tokens are where those two spans overlap. Non-empty
+            #    whenever we matched more blocks than we skipped tokens for: the
+            #    ">= 2 computed tokens" clamp, the Mamba back-off, or memory-only
+            #    hybrid mode where nothing is skipped but blocks are still shared.
+            overlap_start_token = max(matched_prefix_start_token, chunk_start_token)
+            overlap_end_token = min(matched_prefix_end_token, chunk_end_token)
+
+            # 4. Rebase the overlap onto this chunk's slots and redirect it. The
+            #    overlap starts past the chunk start only when this chunk's own
+            #    matched blocks begin later than the tokens it computes, i.e. the
+            #    chunk opens inside a freshly allocated block that it must write.
+            if overlap_end_token > overlap_start_token:
+                redirect_start = self.active_token_count + overlap_start_token - chunk_start_token
+                redirect_end = self.active_token_count + overlap_end_token - chunk_start_token
+                self.token_to_block_idx[redirect_start:redirect_end] = (
+                    self.kv_block_allocator.dummy_block_idx
+                )
         self.token_to_local_position_within_kv_block[
             self.active_token_count : self.active_token_count + effective_prefill_chunk_length
         ] = (token_offset_range % self.block_size_tokens)
@@ -3497,6 +3555,13 @@ def _register_range(start: int, end: int):
                 overall_required_blocks,
             )
 
+        # Extend the request's run of matched blocks, so a later chunk knows not to
+        # rewrite the partial block it inherits from this one. Only extend when this
+        # chunk's matches continue the run: a gap means the blocks in between were
+        # computed by this request and the tokens completing them must be written.
+        if num_matched_blocks > 0 and req.num_matched_prefix_blocks >= already_allocated_blocks:
+            req.num_matched_prefix_blocks = already_allocated_blocks + num_matched_blocks
+
         self.active_token_count += effective_prefill_chunk_length
         self.lifetime_prefill_token_count += effective_prefill_chunk_length
         if self.enable_prefix_caching:

megatron/core/inference/contexts/fused_kv_append_kernel.py

@@ -42,6 +42,7 @@ def _append_kv_cache_kernel(
     n_tokens: tl.int32,
     num_heads: tl.int32,
     H_DIM: tl.int32,
+    dummy_block_idx: tl.int32,
     # --- Compile-Time Constants ---
     BLOCK_SIZE_H: tl.constexpr,
 ):
@@ -86,8 +87,11 @@ def _append_kv_cache_kernel(
     value_dest_ptr = value_cache_ptr + dest_offset
 
     # --- Store the head data into the cache ---
-    tl.store(key_dest_ptr + offs_h * stride_cache_hdim, key_to_write, mask=mask_h)
-    tl.store(value_dest_ptr + offs_h * stride_cache_hdim, value_to_write, mask=mask_h)
+    # Tokens routed to the dummy block are padding or recomputed prefix; drop their
+    # stores rather than spend bandwidth on a scratch block nothing reads.
+    write_mask = mask_h & (block_idx != dummy_block_idx)
+    tl.store(key_dest_ptr + offs_h * stride_cache_hdim, key_to_write, mask=write_mask)
+    tl.store(value_dest_ptr + offs_h * stride_cache_hdim, value_to_write, mask=write_mask)
 
 
 def triton_append_key_value_cache(
@@ -98,6 +102,7 @@ def triton_append_key_value_cache(
     padded_active_token_count: int,
     token_to_block_idx: Tensor,
     token_to_local_position_within_kv_block: Tensor,
+    dummy_block_idx: int,
 ) -> None:
     """
     Append to KV cache using a high-performance, standalone Triton kernel.
@@ -112,6 +117,7 @@ def triton_append_key_value_cache(
         the cache.
         token_to_local_position_within_kv_block (Tensor): Tensor mapping token index
         to its position within a block.
+        dummy_block_idx (int): Sentinel block index whose writes are skipped.
     """
     # --- Input Validation and Preparation ---
     assert (
@@ -182,6 +188,7 @@ def triton_append_key_value_cache(
         n_tokens=n_tokens,
         num_heads=num_heads,
         H_DIM=h_dim,
+        dummy_block_idx=dummy_block_idx,
         # Compile-time constant
         BLOCK_SIZE_H=BLOCK_SIZE_H,
     )

megatron/core/inference/engines/dynamic_engine.py

@@ -996,6 +996,7 @@ def suspend(self):
                 if req.finished_chunk_token_count > 0:
                     req.remaining_prompt_tokens = req.prompt_tokens
                     req.finished_chunk_token_count = 0
+                    req.num_matched_prefix_blocks = 0
 
             # Reset the chunked prefill request id
             self.chunked_prefill_request_id = -1

megatron/core/inference/inference_request.py

@@ -526,6 +526,10 @@ class DynamicInferenceRequest(InferenceRequest):
     block_size_tokens: Optional[int] = None  # Block size for hash computation
     enable_prefix_caching: bool = False  # Whether prefix caching is enabled
     num_cached_tokens: int = 0  # Tokens served from prefix cache (set by context on first match)
+    # Length of the leading run of this request's blocks that was obtained by hash
+    # match rather than computed. Accumulated across prefill chunks by the context,
+    # which uses it to avoid rewriting KV into blocks that already hold it.
+    num_matched_prefix_blocks: int = 0
 
     # Computed field - not passed by caller
     precomputed_block_hashes: List[int] = field(default_factory=list)

tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py

@@ -1421,6 +1421,196 @@ def test_mixed_batch(self, model_type):
         assert len(log_probs_list[4]) == fresh_ql
 
 
+class TestMatchedBlockWriteRedirect(PrefixCachingTestBase):
+    """Tokens recomputed inside a hash-matched block must write to the dummy block.
+
+    The matched block is shared with the request that cached it and already holds
+    the correct KV for exactly those tokens, so `add_request` redirects their
+    `token_to_block_idx` entries away from the real block.
+    """
+
+    @staticmethod
+    def _write_targets(ctx, start, end):
+        """Block ids this step will write KV to, for token slots [start, end)."""
+        return ctx.token_to_block_idx[start:end].tolist()
+
+    def _mctx(self, **kwargs):
+        """Hybrid context with a Mamba state cache, so prefill skipping is possible."""
+        defaults = dict(
+            mamba_config=self._mamba_config(),
+            prefix_caching_mamba_gb=0.01,
+            block_size_tokens=256,
+            max_sequence_length=4096,
+        )
+        defaults.update(kwargs)
+        return self._ctx(**defaults)
+
+    def _add_chunk(self, ctx, req, chunk_length=None):
+        """Admit one prefill chunk; return the token slots it claimed.
+
+        Mirrors the continuation bookkeeping the engine does between chunks:
+        advance the prompt window and hand the same request row back to the
+        next `add_request` call.
+        """
+        start = ctx.active_token_count
+        ctx.add_request(req, prefill_chunk_length=chunk_length)
+        end = ctx.active_token_count
+        if chunk_length is not None:
+            req.remaining_prompt_tokens = req.remaining_prompt_tokens[chunk_length:]
+            req.finished_chunk_token_count += chunk_length
+            # `update_requests` decrements this for a continuing chunked prefill so
+            # the next chunk reuses the same request row.
+            ctx.total_request_count -= 1
+        return start, end
+
+    @pytest.mark.internal
+    def test_fully_cached_repeat_redirects_recomputed_block(self):
+        """The `>= 2 computed tokens` clamp forces a matched block to be recomputed."""
+        ctx = self._ctx()
+        bs = ctx.block_size_tokens
+        dummy = ctx.kv_block_allocator.dummy_block_idx
+        prompt = self._prompt(bs * 3)
+
+        ctx.add_request(self._req(ctx, prompt.clone()))
+        cached_blocks = self._block_ids(ctx, 0, 3)
+
+        # Re-sending the prompt verbatim matches all 3 blocks, but skipping all 96
+        # tokens would leave a 0-token chunk, so the skip is clamped down to 2 blocks
+        # and the third block's 32 tokens are recomputed.
+        req2 = self._req(ctx, prompt.clone(), request_id=2)
+        start, end = self._add_chunk(ctx, req2)
+        assert end - start == bs
+        assert ctx.request_kv_length_offsets[1].item() == bs * 2
+
+        # Every recomputed token lands in the matched third block, so all of the
+        # chunk's writes are redirected.
+        assert self._write_targets(ctx, start, end) == [dummy] * bs
+        # The block table is untouched: attention still reads the cached KV.
+        assert self._block_ids(ctx, 1, 3) == cached_blocks
+
+    @pytest.mark.internal
+    def test_partial_match_redirects_only_matched_region(self):
+        """A prompt that extends a cached prefix writes fresh blocks normally."""
+        ctx = self._ctx()
+        bs = ctx.block_size_tokens
+        dummy = ctx.kv_block_allocator.dummy_block_idx
+
+        cached_prompt = self._prompt(bs * 2)
+        ctx.add_request(self._req(ctx, cached_prompt.clone()))
+
+        # Two cached blocks plus two fresh ones. The skip covers both matched
+        # blocks exactly, so nothing is recomputed and nothing is redirected.
+        extended = torch.cat([cached_prompt, self._prompt(bs * 2, offset=5000)])
+        req2 = self._req(ctx, extended, request_id=2)
+        start, end = self._add_chunk(ctx, req2)
+        assert end - start == bs * 2
+        assert ctx.request_kv_length_offsets[1].item() == bs * 2
+
+        fresh_blocks = self._block_ids(ctx, 1, 4)[2:]
+        assert dummy not in fresh_blocks
+        targets = self._write_targets(ctx, start, end)
+        assert targets == [fresh_blocks[0]] * bs + [fresh_blocks[1]] * bs
+
+    @pytest.mark.internal
+    def test_no_match_writes_every_token_to_a_real_block(self):
+        """Control: with no prefix hit there is nothing to redirect."""
+        ctx = self._ctx()
+        bs = ctx.block_size_tokens
+        dummy = ctx.kv_block_allocator.dummy_block_idx
+
+        ctx.add_request(self._req(ctx, self._prompt(bs * 2)))
+        req2 = self._req(ctx, self._prompt(bs * 2, offset=9000), request_id=2)
+        start, end = self._add_chunk(ctx, req2)
+
+        assert end - start == bs * 2
+        assert dummy not in self._write_targets(ctx, start, end)
+
+    @pytest.mark.internal
+    def test_hybrid_memory_only_redirects_all_matched_tokens(self):
+        """Hybrid models share blocks without skipping prefill, so all writes redirect."""
+        ctx = self._ctx(mamba_config=self._mamba_config())
+        bs = ctx.block_size_tokens
+        dummy = ctx.kv_block_allocator.dummy_block_idx
+        prompt = self._prompt(bs * 3)
+        assert ctx.is_hybrid_model
+
+        ctx.add_request(self._req(ctx, prompt.clone()))
+        cached_blocks = self._block_ids(ctx, 0, 3)
+
+        # No Mamba budget here, so prefix_skip_tokens is 0: the whole prompt is
+        # recomputed even though all 3 blocks were matched and shared.
+        req2 = self._req(ctx, prompt.clone(), request_id=2)
+        start, end = self._add_chunk(ctx, req2)
+        assert end - start == bs * 3
+        assert ctx.request_kv_length_offsets[1].item() == 0
+
+        assert self._write_targets(ctx, start, end) == [dummy] * (bs * 3)
+        assert self._block_ids(ctx, 1, 3) == cached_blocks
+
+    @pytest.mark.parametrize("mamba_cached_blocks", [1, 3])
+    @pytest.mark.internal
+    def test_mamba_skip_shorter_than_kv_match_redirects_remainder(self, mamba_cached_blocks):
+        """Mamba state, not the KV match, bounds the skip -- the rest is recomputed.
+
+        With 1 of 3 blocks holding Mamba state the skip stops after that block; with
+        all 3 it would leave a zero-token chunk, so it backs off to block 2. Either
+        way the recomputed tokens land in KV blocks that were matched.
+        """
+        ctx = self._mctx()
+        bs = ctx.block_size_tokens
+        dummy = ctx.kv_block_allocator.dummy_block_idx
+        prompt = self._prompt(bs * 3)
+
+        ctx.add_request(self._req(ctx, prompt.clone()))
+        cached_blocks = self._block_ids(ctx, 0, 3)
+        self._mamba_allocate_and_register(ctx, cached_blocks[:mamba_cached_blocks])
+
+        req2 = self._req(ctx, prompt.clone(), request_id=2)
+        start, end = self._add_chunk(ctx, req2)
+
+        skipped_blocks = 1 if mamba_cached_blocks == 1 else 2
+        assert ctx.request_kv_length_offsets[1].item() == skipped_blocks * bs
+        assert end - start == (3 - skipped_blocks) * bs
+
+        assert self._write_targets(ctx, start, end) == [dummy] * (end - start)
+        assert self._block_ids(ctx, 1, 3) == cached_blocks
+
+    @pytest.mark.internal
+    def test_unaligned_continuation_redirects_inherited_partial_block(self):
+        """A chunk resuming mid-block must not write into it if it was matched too."""
+        ctx = self._ctx()
+        bs = ctx.block_size_tokens
+        dummy = ctx.kv_block_allocator.dummy_block_idx
+        prompt = self._prompt(bs * 4)
+
+        ctx.add_request(self._req(ctx, prompt.clone()))
+
+        # Chunk 1 stops at token 40, i.e. 8 tokens into block 1, leaving
+        # `finished_chunk_token_count` unaligned for chunk 2. Both blocks it spans
+        # are matched, so its 8 computed tokens redirect.
+        req2 = self._req(ctx, prompt.clone(), request_id=2)
+        chunk1_length = bs + 8
+        start1, end1 = self._add_chunk(ctx, req2, chunk_length=chunk1_length)
+        assert end1 - start1 == 8
+        assert req2.finished_chunk_token_count == chunk1_length
+        assert req2.num_matched_prefix_blocks == 2
+        assert self._write_targets(ctx, start1, end1) == [dummy] * 8
+
+        # Chunk 2 covers tokens [40, 128). The first 24 complete block 1 -- which
+        # chunk 1 obtained by match and still shares with the first request -- so
+        # they redirect alongside the 64 tokens in matched blocks 2-3.
+        chunk2_length = bs * 4 - chunk1_length
+        start2, end2 = self._add_chunk(ctx, req2, chunk_length=chunk2_length)
+        assert end2 - start2 == chunk2_length
+        assert ctx.request_kv_length_offsets[1].item() == chunk1_length
+        assert req2.num_matched_prefix_blocks == 4
+
+        assert self._write_targets(ctx, start2, end2) == [dummy] * chunk2_length
+        # Every block this request holds is shared with the first request.
+        for block_id in self._block_ids(ctx, 1, 4):
+            assert ctx.kv_block_allocator.block_ref_counts[block_id].item() == 2
+
+
 def _make_cpu_mamba_slot_allocator(
     monkeypatch, *, total_blocks: int, max_slots: int
 ) -> MambaSlotAllocator: