Initialize newly private value-cache pages before use

megatron/core/inference/contexts/dynamic_context.py

@@ -1118,6 +1118,14 @@ def _allocate_memory_buffer(self):
                 self.memory_buffer, device="cpu"
             ).pin_memory()
 
+    def _initialize_mha_value_pages(self, block_ids: Sequence[int] | Tensor) -> None:
+        """Zero newly acquired MHA value-cache pages before their first model write."""
+        # Paged attention can read masked V tails; initialize only newly private pages.
+        if self.cache_mla_latent or len(block_ids) == 0:
+            return
+        block_ids = torch.as_tensor(block_ids, dtype=torch.int64, device=self.memory_buffer.device)
+        self.memory_buffer[1].index_fill_(1, block_ids, 0)
+
     def _allocate_mamba_states(self):
         """Allocate Mamba states for hybrid models."""
         if self.is_hybrid_model:
@@ -3524,6 +3532,7 @@ def add_request(
                 if matched_tensor is not None:
                     self.kv_block_allocator.block_ref_counts[matched_tensor] -= 1
                 raise BlockOverflowError(req.request_id)
+            self._initialize_mha_value_pages(new_block_ids)
 
         # Track prefix cache hits only after allocation succeeds. Matched blocks
         # measure KV reuse, while num_cached_tokens accumulates the prefill tokens
@@ -4075,6 +4084,7 @@ def resume_paused_requests(
                     assert (
                         block_ids is not None and block_ids.numel() == num_new_blocks
                     ), f"failed to allocate {num_new_blocks} blocks for resumed requests"
+                    self._initialize_mha_value_pages(block_ids)
 
                     # Apply updates only to the requests that required a new block
                     relative_row_idx = torch.nonzero(needs_new_block).squeeze(1)
@@ -4307,6 +4317,7 @@ def prepare_requests(self) -> None:
                     raise RuntimeError(
                         "Async scheduling cannot evict requests to allocate new blocks."
                     )
+                self._initialize_mha_value_pages(block_ids)
 
             self.active_token_count = active_request_count * tokens_per_request
             active_token_slice = slice(0, self.active_token_count)

megatron/core/inference/contexts/mtp_context_mixin.py

@@ -229,6 +229,7 @@ def _allocate_mtp_decode_blocks(self, request_slice: slice) -> None:
             return
         block_ids = self.kv_block_allocator.allocate_memory_blocks(total)
         assert block_ids is not None and block_ids.numel() == total
+        self._initialize_mha_value_pages(block_ids)
         rows = torch.repeat_interleave(
             torch.arange(request_slice.start, request_slice.stop), counts
         )

megatron/core/inference/disaggregation/inference_state_handoff.py

@@ -986,6 +986,7 @@ def _try_start_kv_handoff_import(self, handoff: DeferredKvHandoff) -> bool:
         allocated_blocks = [int(block) for block in allocated_blocks_tensor.tolist()]
         imported_blocks = allocated_blocks[:num_blocks_to_import]
         continuation_blocks = allocated_blocks[num_blocks_to_import:]
+        self.context._initialize_mha_value_pages(continuation_blocks)
         local_blocks = cached_blocks + imported_blocks
         owned_blocks_tensor = torch.tensor(
             local_blocks + continuation_blocks, dtype=torch.int32, device="cpu"

tests/test_utils/recipes/h100/unit-tests.yaml

@@ -51,6 +51,9 @@ spec:
     # test_gtp_muon.py/test_gtp_muon_qkv.py, and will skip here instead).
     if [[ "$BUCKET" == "tests/unit_tests/generalized_tensor_parallel/**/*.py" ]]; then
         export GPUS_PER_NODE=4
+    elif [[ "$BUCKET" == "tests/unit_tests/inference/engines/test_batch_invariant_handoff.py" ]]; then
+        # One prefill worker and one decode worker exchange a real NIXL handoff.
+        export GPUS_PER_NODE=2
     fi
 
     if [[ "$TAG" == "latest" ]]; then
@@ -83,6 +86,15 @@ spec:
     fi
 
 products:
+  # Run the two-worker handoff separately; the broad inference bucket excludes it.
+  - test_case: [tests/unit_tests/inference/engines/test_batch_invariant_handoff.py]
+    products:
+      - environment: [lts, dev]
+        tag: [latest]
+        scope: [unit-tests]
+        gpus: [2]
+        n_repeat: [1]
+        time_limit: [1800]
   # This file runs in isolated one-rank processes; other buckets exclude it.
   - test_case: [tests/unit_tests/inference/test_batch_invariant_fp8.py]
     products:

tests/unit_tests/inference/contexts/test_dynamic_context.py

@@ -1271,6 +1271,7 @@ def test_async_sched_prepare_requests_success(
             last_block_offsets=last_offsets,
         )
         original_tokens = ctx.token_to_input_ids[:active_request_count].clone()
+        ctx.memory_buffer[1].fill_(5)
 
         ctx.prepare_requests()
 
@@ -1291,6 +1292,9 @@ def test_async_sched_prepare_requests_success(
         if last_offsets and last_offsets[0] == ctx.block_size_tokens - 1:
             assert ctx.request_kv_block_counts[0] == 2
             assert ctx.token_to_block_idx[0] == ctx.request_last_kv_block_id[0]
+            new_value_page = ctx.memory_buffer[1, :, ctx.request_last_kv_block_id[0].item()]
+            assert torch.count_nonzero(new_value_page) == 0
+            assert torch.count_nonzero(ctx.memory_buffer[1] != 5) == new_value_page.numel()
 
     @pytest.mark.internal
     @rounder_override(8)

tests/unit_tests/inference/contexts/test_dynamic_context_mtp_kv_cache.py

@@ -1280,10 +1280,13 @@ def test_async_prepare_checks_draft_headroom_before_mutating_state(self):
         with pytest.raises(RuntimeError, match="cannot pause requests"):
             context.prepare_requests()
         torch.testing.assert_close(context.request_kv_length_offsets, before)
+        context.memory_buffer[1, :, held[0]].fill_(float("nan"))
         context.kv_block_allocator.release_memory_blocks(held[:1])
         assert context.can_prepare_requests()
         context.prepare_requests()
         assert context.request_kv_block_counts[0] == 2
+        assert context.request_to_kv_block_ids[0, 1] == held[0]
+        assert torch.count_nonzero(context.memory_buffer[1, :, held[0]]) == 0
         self._assert_main_token_mapping(context)
         context.kv_block_allocator.release_memory_blocks(held[1:])
 

tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py

@@ -535,6 +535,11 @@ def test_add_request_full_cache_partial_hit_pins_matched_blocks(self):
         assert alloc.block_ref_counts[s0].item() == 0
         assert alloc.block_ref_counts[s1].item() == 0
         assert alloc.block_ref_counts[sx].item() == 0
+        untouched_block_id = next(iter(set(range(alloc.pool_size - 1)) - {s0, s1, sx}))
+        ctx.memory_buffer[1, :, s0] = 1
+        ctx.memory_buffer[1, :, s1] = 2
+        ctx.memory_buffer[1, :, sx] = float("nan")
+        ctx.memory_buffer[1, :, untouched_block_id] = 4
 
         # Force a full pool: the new block for H2 can only come from eviction.
         alloc.pool_avail = 0
@@ -560,6 +565,10 @@ def test_add_request_full_cache_partial_hit_pins_matched_blocks(self):
         assert alloc.block_hashes[s0].item() == h0
         assert alloc.block_hashes[s1].item() == h1
         assert alloc.block_hashes[sx].item() == h2
+        assert torch.all(ctx.memory_buffer[1, :, s0] == 1)
+        assert torch.all(ctx.memory_buffer[1, :, s1] == 2)
+        assert torch.count_nonzero(ctx.memory_buffer[1, :, sx]) == 0
+        assert torch.all(ctx.memory_buffer[1, :, untouched_block_id] == 4)
         # Parent bookkeeping is stored as resolved block ids: S1's parent is S0
         # and SX's parent is S1 along the H0 -> H1 -> H2 chain.
         assert alloc.block_parent_id[s1].item() == s0
@@ -666,6 +675,12 @@ def test_resume_boundary_crossing_evicts_lru_block_when_free_pool_empty(self):
         assert alloc.block_ref_counts[original_block_id].item() == 1
 
         cached_block_id, cached_hash = self._fill_pool_with_one_evictable_block(ctx)
+        untouched_block_id = next(
+            iter(set(range(alloc.pool_size - 1)) - {original_block_id, cached_block_id})
+        )
+        ctx.memory_buffer[1, :, original_block_id] = 2
+        ctx.memory_buffer[1, :, cached_block_id] = float("nan")
+        ctx.memory_buffer[1, :, untouched_block_id] = 4
 
         result = ctx.update_requests(
             torch.ones(1, device=torch.cuda.current_device(), dtype=torch.int32),
@@ -682,6 +697,9 @@ def test_resume_boundary_crossing_evicts_lru_block_when_free_pool_empty(self):
         assert new_block_id != original_block_id
         assert cached_hash not in alloc.kv_hash_to_block_id
         assert alloc.block_hashes[new_block_id].item() == -1
+        assert torch.all(ctx.memory_buffer[1, :, original_block_id] == 2)
+        assert torch.count_nonzero(ctx.memory_buffer[1, :, new_block_id]) == 0
+        assert torch.all(ctx.memory_buffer[1, :, untouched_block_id] == 4)
         assert alloc.block_ref_counts[original_block_id].item() == 1
         assert alloc.block_ref_counts[new_block_id].item() == 1
         assert alloc.pool_avail == 0
@@ -1716,6 +1734,10 @@ def test_unaligned_continuation_redirects_inherited_partial_block(self):
         prompt = self._prompt(bs * 4)
 
         ctx.add_request(self._req(ctx, prompt.clone()))
+        shared_blocks = self._block_ids(ctx, 0, 4)
+        for value, block_id in enumerate(shared_blocks, start=1):
+            ctx.memory_buffer[1, :, block_id] = value
+        shared_values = ctx.memory_buffer[1, :, shared_blocks].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
@@ -1741,6 +1763,7 @@ def test_unaligned_continuation_redirects_inherited_partial_block(self):
         # 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
+        assert torch.equal(ctx.memory_buffer[1, :, shared_blocks], shared_values)
 
 
 @pytest.mark.internal