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

TestBreakableCUDAGraphBasic and all NVIDIA CI tests pass.

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."""