diff --git a/python/sglang/srt/layers/radix_attention.py b/python/sglang/srt/layers/radix_attention.py index e440aac24..3ec27dc72 100644 --- a/python/sglang/srt/layers/radix_attention.py +++ b/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: diff --git a/test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py b/test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py index 5034eb4b1..33dce0c45 100644 --- a/test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py +++ b/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.""" diff --git a/test/registered/unit/layers/test_radix_attention.py b/test/registered/unit/layers/test_radix_attention.py index b30bc2f85..aec9f3254 100644 --- a/test/registered/unit/layers/test_radix_attention.py +++ b/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(