Avoid passing -1 sentinels into inference kernels

megatron/core/inference/contexts/dynamic_context.py

@@ -2474,6 +2474,15 @@ def initialize_attention_state(
         self._cpu_mha_block_table[:real_bs] = request_to_kv_block_ids_view[:real_bs]
         if real_bs < padded_bs:
             self._cpu_mha_block_table[real_bs:padded_bs] = self.kv_block_allocator.dummy_block_idx
+        # Real rows must avoid having a -1 sentinel in their trailing columns,
+        # because the kernel treats the -1 sentinel as a real value.
+        # We cannot avoid writing -1 into `request_to_kv_block_ids`; other logic needs it.
+        # The only option is to overwrite the -1 with a dummy block index via `masked_fill`.
+        if real_bs > 0:
+            _real_rows = self._cpu_mha_block_table[:real_bs]
+            # masked_fill_ over the pinned int32 view: no index_put/nonzero
+            # temporaries on the per-step CPU path.
+            _real_rows.masked_fill_(_real_rows < 0, self.kv_block_allocator.dummy_block_idx)
 
         # Max sequence lengths (Python scalars; consumed as kernel launch args).
         if not self.using_cuda_graph_this_step() and real_bs > 0:

tests/unit_tests/inference/engines/test_dynamic_engine.py

@@ -1155,6 +1155,39 @@ def test_cuda_graph_padding_uses_dummy_block(self) -> None:
         assert (padded_rows != -1).all()
         assert (padded_rows == context.kv_block_allocator.dummy_block_idx).all()
 
+    @pytest.mark.internal
+    @pytest.mark.skipif(
+        not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching"
+    )
+    @torch.inference_mode()
+    def test_active_row_block_table_tail_uses_dummy_block(self) -> None:
+        """A real request's unallocated trailing block-table columns must be
+        staged as the dummy block, not the -1 sentinel: graphed decode
+        advertises max_seqlen_k = max_sequence_length, so the kernel's
+        page-table reach can include the tail of a near-limit request."""
+        test_config = DynamicEngineTestConfig(
+            num_requests=1,
+            min_prompt_length=8,
+            max_prompt_length=8,
+            num_cuda_graphs=1,
+            context_max_requests=4,
+        )
+        env = self._build_test_env(test_config)
+        context = env.engine.context
+
+        env.engine._add_request(env.requests[0])
+        self._run_step(env)  # prefill
+        self._run_step(env)  # decode: graphed, one real row
+
+        assert context.using_cuda_graph_this_step()
+        # An 8-token prompt holds a handful of pages of the 512-token budget,
+        # so the real row genuinely has an unallocated tail to pin.
+        block_count = int(context.request_kv_block_counts[0].item())
+        staged_row = context._cpu_mha_block_table[0]
+        assert 0 < block_count < staged_row.numel()
+        assert (staged_row != -1).all()
+        assert (staged_row[block_count:] == context.kv_block_allocator.dummy_block_idx).all()
+
     @pytest.mark.internal
     @pytest.mark.skipif(
         not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching"