feat(model_runner): remove pool/backend refs from ForwardBatch via ForwardContext (#25983)
Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Sonnet 4.6
parent
44ec2ee18d
commit
c5251a98a9
@@ -360,7 +360,6 @@ class TestDSAIndexer(CustomTestCase):
|
||||
),
|
||||
extend_seq_lens=torch.tensor([q_len] * batch_size, device=self.device),
|
||||
extend_seq_lens_cpu=torch.tensor([q_len] * batch_size, device="cpu"),
|
||||
attn_backend=self.backend,
|
||||
)
|
||||
else: # ForwardMode.DECODE
|
||||
decode_len = 1
|
||||
@@ -379,12 +378,18 @@ class TestDSAIndexer(CustomTestCase):
|
||||
req_pool_indices=torch.arange(batch_size, device=self.device),
|
||||
seq_lens=torch.tensor([total_len] * batch_size, device=self.device),
|
||||
seq_lens_cpu=torch.tensor([total_len] * batch_size, device="cpu"),
|
||||
attn_backend=self.backend,
|
||||
)
|
||||
|
||||
# Add token pools
|
||||
forward_batch.req_to_token_pool = self.model_runner.req_to_token_pool
|
||||
forward_batch.token_to_kv_pool = self.model_runner.token_to_kv_pool
|
||||
# Pool refs + attn_backend are now resolved via the ForwardContext;
|
||||
# publish ``self.backend`` for the duration of this fixture call so
|
||||
# ``get_attn_backend()`` / ``get_token_to_kv_pool()`` /
|
||||
# ``get_req_to_token_pool()`` resolve correctly.
|
||||
from sglang.srt.model_executor.forward_context import (
|
||||
ForwardContext,
|
||||
set_forward_context,
|
||||
)
|
||||
|
||||
set_forward_context(ForwardContext(attn_backend=self.backend))
|
||||
|
||||
# Mock write to req_to_token_pool
|
||||
page_size = self.model_runner.page_size
|
||||
|
||||
Reference in New Issue
Block a user