From 6e126c757d12ddd05ec7da616f4ba338a81fa60d Mon Sep 17 00:00:00 2001 From: Jorge Albericio Date: Wed, 12 Aug 2026 01:06:43 -0500 Subject: [PATCH 1/2] Avoid passing -1 sentinels into kernels Signed-off-by: Teodor-Dumitru Ene --- .../inference/contexts/dynamic_context.py | 9 +++++ .../inference/engines/test_dynamic_engine.py | 33 +++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 0b50f50d6de..4a5a2836e5e 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/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: diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index af91409ef7b..02970a8c10c 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/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" From 0dc1b1786592c2f900e71ce1692b2c2d7c20e18a Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Thu, 13 Aug 2026 01:31:10 -0500 Subject: [PATCH 2/2] Fix test Signed-off-by: Teodor-Dumitru Ene --- tests/unit_tests/inference/engines/test_dynamic_engine.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index 02970a8c10c..0336881e22e 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py @@ -1171,6 +1171,7 @@ def test_active_row_block_table_tail_uses_dummy_block(self) -> None: max_prompt_length=8, num_cuda_graphs=1, context_max_requests=4, + max_sequence_length=512, ) env = self._build_test_env(test_config) context = env.engine.context @@ -1180,8 +1181,8 @@ def test_active_row_block_table_tail_uses_dummy_block(self) -> None: 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. + # The 512-token budget spans two 256-token pages. 8-token prompt allocates only the first. + # 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()