feat(model_runner): remove pool/backend refs from ForwardBatch via ForwardContext (#25983)

Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>

test/manual/attention/test_flashattn_backend.py

@@ -11,6 +11,10 @@ from sglang.srt.layers.attention.torch_native_backend import TorchNativeAttnBack
 from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.model_executor.forward_context import (
+    ForwardContext,
+    set_forward_context,
+)
 from sglang.test.test_utils import CustomTestCase
 
 
@@ -109,6 +113,9 @@ class TestFlashAttentionBackend(CustomTestCase):
         self.backend = FlashAttentionBackend(self.model_runner)
         self.ref_backend = TorchNativeAttnBackend(self.model_runner)
         self.model_runner.model_config.num_attention_heads = self.num_heads
+        # Publish the backend for any RadixAttention.forward path the tests
+        # exercise; tearDown is unnecessary here since each test re-inits.
+        set_forward_context(ForwardContext(attn_backend=self.backend))
 
     def _mock_write_to_req_to_token_pool(self, batch_size, seq_len, page_size):
         # if page_size > 1, the token pool stores the index to the page.
@@ -223,7 +230,6 @@ class TestFlashAttentionBackend(CustomTestCase):
                 extend_seq_lens_cpu=torch.tensor(
                     [q_len] * self.batch_size, device="cpu"
                 ),
-                attn_backend=self.backend,
             )
             if attn_cp_size > 1:
                 forward_batch.attn_cp_metadata = type(
@@ -273,16 +279,11 @@ class TestFlashAttentionBackend(CustomTestCase):
                     [total_len] * self.batch_size, device=self.device
                 ),
                 seq_lens_cpu=torch.tensor([total_len] * self.batch_size, device="cpu"),
-                attn_backend=self.backend,
             )
 
-        # Add token pool
-        forward_batch.req_to_token_pool = self.model_runner.req_to_token_pool
-
-        # Write current batch's req_to_token to req_to_token_pool
+        # Pool refs are resolved via the active ForwardContext (published in
+        # setUp). Write the test fixture's req_to_token mapping.
         self._mock_write_to_req_to_token_pool(self.batch_size, total_len, page_size)
-        # Add kv pool for this forward batch
-        forward_batch.token_to_kv_pool = self.model_runner.token_to_kv_pool
 
         return forward_batch
 
@@ -307,7 +308,7 @@ class TestFlashAttentionBackend(CustomTestCase):
         )
 
         # Set the prefix KV cache
-        forward_batch.token_to_kv_pool.set_kv_buffer(
+        self.model_runner.token_to_kv_pool.set_kv_buffer(
             layer,
             torch.arange(self.batch_size * cache_len, device=self.device),
             cache_k,

test/manual/attention/test_flashattn_mla_backend.py

@@ -8,6 +8,10 @@ from sglang.srt.layers.attention.torch_native_backend import TorchNativeAttnBack
 from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.model_executor.forward_context import (
+    ForwardContext,
+    set_forward_context,
+)
 from sglang.test.test_utils import CustomTestCase
 
 
@@ -112,6 +116,8 @@ class TestFlashAttentionMLABackend(CustomTestCase):
         self.backend = FlashAttentionBackend(self.model_runner)
         self.ref_backend = TorchNativeAttnBackend(self.model_runner)
         self.num_local_heads = 2
+        # Publish the backend so RadixAttention.forward resolves correctly.
+        set_forward_context(ForwardContext(attn_backend=self.backend))
 
     def _init_model_runner(self):
         self.model_runner = MockModelRunner(
@@ -192,7 +198,6 @@ class TestFlashAttentionMLABackend(CustomTestCase):
                 extend_seq_lens_cpu=torch.tensor(
                     [q_len] * self.batch_size, device="cpu"
                 ),
-                attn_backend=self.backend,
             )
 
         else:  # ForwardMode.DECODE
@@ -216,15 +221,10 @@ class TestFlashAttentionMLABackend(CustomTestCase):
                     [total_len] * self.batch_size, device=self.device
                 ),
                 seq_lens_cpu=torch.tensor([total_len] * self.batch_size, device="cpu"),
-                attn_backend=self.backend,
             )
 
-        # Add token pool from model runner to forward batch
-        forward_batch.req_to_token_pool = self.model_runner.req_to_token_pool
-
-        # Add KV cache from model runner to forward batch
-        forward_batch.token_to_kv_pool = self.model_runner.token_to_kv_pool
-
+        # Pool refs are resolved via the active ForwardContext (published in
+        # setUp); the fixture no longer needs to attach them to forward_batch.
         return forward_batch
 
     def _setup_kv_cache(self, forward_batch, layer, cache_len):
@@ -250,7 +250,7 @@ class TestFlashAttentionMLABackend(CustomTestCase):
         )
 
         # Set the prefix KV cache using MLA-specific method
-        forward_batch.token_to_kv_pool.set_mla_kv_buffer(
+        self.model_runner.token_to_kv_pool.set_mla_kv_buffer(
             layer,
             torch.arange(self.batch_size * cache_len, device=self.device),
             cache_k_nope,

test/manual/attention/test_prefix_chunk_info.py

@@ -110,12 +110,10 @@ class MockReqToTokenPool:
 
 
 # Test correctness of triton kernel for computing kv indices
-def check_kv_indices(forward_batch):
+def check_kv_indices(forward_batch, req_to_token_pool):
     for i in range(forward_batch.num_prefix_chunks):
         computed_kv_indices = forward_batch.prefix_chunk_kv_indices[i]
-        req_to_token = forward_batch.req_to_token_pool.req_to_token[
-            : forward_batch.batch_size, :
-        ]
+        req_to_token = req_to_token_pool.req_to_token[: forward_batch.batch_size, :]
         ref_kv_indices = torch.empty(
             forward_batch.prefix_chunk_num_tokens[i],
             dtype=torch.int32,
@@ -205,8 +203,20 @@ class TestPrefixChunkInfo(CustomTestCase):
                 extend_prefix_lens=prefix_lens,
                 extend_prefix_lens_cpu=prefix_lens_cpu,
             )
-            forward_batch.req_to_token_pool = self.req_to_token_pool
-            forward_batch.token_to_kv_pool = self.token_to_kv_pool
+            # Pool refs are resolved via the active ForwardContext; mock an
+            # attn_backend that carries the pools (Pattern A invariant).
+            from types import SimpleNamespace
+
+            from sglang.srt.model_executor.forward_context import (
+                ForwardContext,
+                set_forward_context,
+            )
+
+            mock_backend = SimpleNamespace(
+                req_to_token_pool=self.req_to_token_pool,
+                token_to_kv_pool=self.token_to_kv_pool,
+            )
+            set_forward_context(ForwardContext(attn_backend=mock_backend))
 
             forward_batch.prepare_chunked_prefix_cache_info(self.device)
             assert forward_batch.get_max_chunk_capacity() == max_chunk_capacity
@@ -221,7 +231,7 @@ class TestPrefixChunkInfo(CustomTestCase):
                 test_case["prefix_chunk_seq_lens"].to(self.device),
             )
 
-            check_kv_indices(forward_batch)
+            check_kv_indices(forward_batch, self.req_to_token_pool)
 
 
 if __name__ == "__main__":

test/manual/attention/test_trtllm_mla_backend.py

@@ -20,6 +20,10 @@ from sglang.srt.layers.attention.utils import get_num_page_per_block_flashmla
 from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.model_executor.forward_context import (
+    ForwardContext,
+    set_forward_context,
+)
 from sglang.srt.server_args import (
     ServerArgs,
     get_global_server_args,
@@ -434,10 +438,9 @@ class TestTRTLLMMLA(CustomTestCase):
             req_pool_indices=torch.arange(batch_size, device=config["device"]),
             seq_lens=seq_lens,
             seq_lens_cpu=seq_lens.cpu(),
-            attn_backend=backend,
         )
-        fb.req_to_token_pool = model_runner.req_to_token_pool
-        fb.token_to_kv_pool = model_runner.token_to_kv_pool
+        # Publish backend for RadixAttention dispatch.
+        set_forward_context(ForwardContext(attn_backend=backend))
 
         # Add position information for RoPE
         fb.positions = torch.arange(batch_size, device=config["device"])
@@ -1167,10 +1170,9 @@ class TestTRTLLMMLA(CustomTestCase):
                         seq_lens_cpu=seq_lens.cpu(),
                         attn_attend_prefix_cache=False,
                         mha_return_lse=False,
-                        attn_backend=backend,
                     )
-                    fb.req_to_token_pool = model_runner.req_to_token_pool
-                    fb.token_to_kv_pool = model_runner.token_to_kv_pool
+                    # Publish backend for RadixAttention dispatch.
+                    set_forward_context(ForwardContext(attn_backend=backend))
 
                     # Add position information for RoPE
                     fb.positions = torch.arange(batch_size, device=config["device"])