Fix padded positions in breakable CUDA Graph attention (#33253)

TestBreakableCUDAGraphBasic and all NVIDIA CI tests pass.

python/sglang/srt/layers/radix_attention.py

@@ -351,9 +351,12 @@ def _unified_attention_with_output_impl(
         kwargs["topk_indices"] = topk_indices[:real_query_num_tokens]
 
     original_out_cache_loc = forward_batch.out_cache_loc
+    original_positions = forward_batch.positions
     # Keep the original ForwardBatch object and only narrow cache locations for
     # this backend call so model/backend state is still written to the same batch.
     forward_batch.out_cache_loc = original_out_cache_loc[:real_query_num_tokens]
+    if original_positions is not None:
+        forward_batch.positions = original_positions[:real_query_num_tokens]
 
     # Store pre-allocated output for FA backend to write directly into.
     # Must slice to real_query_num_tokens to match the narrowed query shape —
@@ -370,6 +373,7 @@ def _unified_attention_with_output_impl(
         **kwargs,
     )
     forward_batch.out_cache_loc = original_out_cache_loc
+    forward_batch.positions = original_positions
 
     lse = None
     if return_lse:

test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py

@@ -8,6 +8,7 @@ Two test classes:
 """
 
 import unittest
+from unittest.mock import patch
 
 import torch
 
@@ -215,6 +216,56 @@ class TestBreakableCUDAGraphBasic(CustomTestCase):
             "eager output bridge buffer must be strongly captured",
         )
 
+    def test_attention_narrows_padded_positions(self):
+        from sglang.srt.layers.radix_attention import unified_attention_with_output
+
+        num_tokens = 3
+        padded_num_tokens = 5
+        forward_batch = SimpleNamespace(
+            num_token_non_padded_cpu=num_tokens,
+            out_cache_loc=torch.arange(padded_num_tokens, device=self.device),
+            positions=torch.arange(padded_num_tokens, device=self.device),
+        )
+        context = SimpleNamespace(
+            forward_batch=forward_batch,
+            attention_layers=[object()],
+            mha_companion_layers=None,
+            num_tokens=padded_num_tokens,
+            raw_num_tokens=num_tokens,
+        )
+        observed = {}
+
+        def attention_forward(query, key, value, layer, batch, save_kv_cache):
+            observed["positions"] = batch.positions.clone()
+            observed["out_cache_loc"] = batch.out_cache_loc.clone()
+            return torch.ones_like(query)
+
+        output = torch.full((padded_num_tokens, 2), float("nan"), device=self.device)
+        with (
+            patch(
+                "sglang.srt.layers.radix_attention.get_tc_piecewise_forward_context",
+                return_value=context,
+            ),
+            patch(
+                "sglang.srt.layers.radix_attention.get_attn_backend",
+                return_value=SimpleNamespace(forward=attention_forward),
+            ),
+        ):
+            unified_attention_with_output(
+                torch.zeros((padded_num_tokens, 2), device=self.device),
+                torch.zeros((padded_num_tokens, 1, 2), device=self.device),
+                torch.zeros((padded_num_tokens, 1, 2), device=self.device),
+                output,
+                True,
+                0,
+            )
+
+        expected = torch.arange(num_tokens, device=self.device)
+        torch.testing.assert_close(observed["positions"], expected)
+        torch.testing.assert_close(observed["out_cache_loc"], expected)
+        self.assertEqual(forward_batch.positions.shape[0], padded_num_tokens)
+        self.assertEqual(forward_batch.out_cache_loc.shape[0], padded_num_tokens)
+
 
 class TestCopyOutput(CustomTestCase):
     """Test the _copy_output helper for structured output writeback."""

test/registered/unit/layers/test_radix_attention.py

@@ -71,6 +71,7 @@ class TestRadixAttentionGraphInterface(CustomTestCase):
         forward_batch = SimpleNamespace(
             num_token_non_padded_cpu=real_num_tokens,
             out_cache_loc=torch.arange(num_tokens, dtype=torch.int64),
+            positions=torch.arange(num_tokens, dtype=torch.int64),
             _attn_output=None,
         )
         return SimpleNamespace(