diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index bf743a2c2..6f48fb817 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -354,14 +354,9 @@ class DeepseekV4AttnBackend( self.page_size = model_runner.page_size assert self.page_size == 256, "the system hardcodes page_size=256" - # Pool refs — captured at construction so they survive deletion of the - # corresponding ForwardBatch fields. self.req_to_token_pool = model_runner.req_to_token_pool self.token_to_kv_pool: DeepSeekV4TokenToKVPool = model_runner.token_to_kv_pool - # Keep a runner ref to read live state set after backend construction - # (e.g. hisparse_coordinator is built in model_runner *after* - # init_attention_backend()). - self.model_runner = model_runner + self.hisparse_coordinator = model_runner.hisparse_coordinator self.req_to_token = model_runner.req_to_token_pool.req_to_token self.MAX_SEQ_LEN_FOR_CAPTURE = self.req_to_token.shape[1] @@ -385,12 +380,6 @@ class DeepseekV4AttnBackend( ] = None self._replay_forward_batch: Optional[ForwardBatch] = None # FIXME: out-of-band - @property - def hisparse_coordinator(self): - # Live read: model_runner builds the coordinator *after* - # init_attention_backend(), so we cannot capture at __init__ time. - return self.model_runner.hisparse_coordinator - def _move_to_device(self, x: List[int]) -> torch.Tensor: pin_tensor = torch.tensor(x, dtype=torch.int32, pin_memory=True) return pin_tensor.to(self.device, non_blocking=True) @@ -1196,7 +1185,6 @@ class DeepseekV4MultiStepBackend(DeepseekV4AttnBackend): self, model_runner: ModelRunner, topk: int, speculative_num_steps: int ): super().__init__(model_runner) - self.model_runner = model_runner self.topk = topk self.speculative_num_steps = speculative_num_steps self.attn_backends: List[DeepseekV4AttnBackend] = [] diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py index 7b8c9dc59..033b877ba 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py @@ -348,14 +348,9 @@ class DeepseekV4HipRadixBackend( self.page_size = model_runner.page_size assert self.page_size == 256, "the system hardcodes page_size=256" - # Pool refs — captured at construction so they survive deletion of the - # corresponding ForwardBatch fields. self.req_to_token_pool = model_runner.req_to_token_pool self.token_to_kv_pool: DeepSeekV4TokenToKVPool = model_runner.token_to_kv_pool - # Keep a runner ref to read live state set after backend construction - # (e.g. hisparse_coordinator is built in model_runner *after* - # init_attention_backend()). - self.model_runner = model_runner + self.hisparse_coordinator = model_runner.hisparse_coordinator self.req_to_token = model_runner.req_to_token_pool.req_to_token self.MAX_SEQ_LEN_FOR_CAPTURE = self.req_to_token.shape[1] @@ -379,12 +374,6 @@ class DeepseekV4HipRadixBackend( ] = None self._replay_forward_batch: Optional[ForwardBatch] = None # FIXME: out-of-band - @property - def hisparse_coordinator(self): - # Live read: model_runner builds the coordinator *after* - # init_attention_backend(), so we cannot capture at __init__ time. - return self.model_runner.hisparse_coordinator - def _move_to_device(self, x: List[int]) -> torch.Tensor: pin_tensor = torch.tensor(x, dtype=torch.int32, pin_memory=True) return pin_tensor.to(self.device, non_blocking=True) @@ -1196,7 +1185,6 @@ class DeepseekV4MultiStepBackend(DeepseekV4HipRadixBackend): self, model_runner: ModelRunner, topk: int, speculative_num_steps: int ): super().__init__(model_runner) - self.model_runner = model_runner self.topk = topk self.speculative_num_steps = speculative_num_steps self.attn_backends: List[DeepseekV4HipRadixBackend] = [] diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index 9e5c8f134..7bdd94007 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -330,14 +330,9 @@ class DeepseekSparseAttnBackend( self.qk_rope_head_dim = model_runner.model_config.qk_rope_head_dim assert model_runner.req_to_token_pool is not None - # Pool refs — captured at construction so they survive deletion of the - # corresponding ForwardBatch fields. self.req_to_token_pool = model_runner.req_to_token_pool self.token_to_kv_pool = model_runner.token_to_kv_pool - # Keep a runner ref to read live state set after backend construction - # (e.g. hisparse_coordinator is built in model_runner *after* - # init_attention_backend()). - self.model_runner = model_runner + self.hisparse_coordinator = model_runner.hisparse_coordinator self.req_to_token = model_runner.req_to_token_pool.req_to_token self.use_mha: bool = False @@ -400,12 +395,6 @@ class DeepseekSparseAttnBackend( else: self.workspace_buffer = None - @property - def hisparse_coordinator(self): - # Live read: model_runner builds the coordinator *after* - # init_attention_backend(), so we cannot capture at __init__ time. - return self.model_runner.hisparse_coordinator - def get_device_int32_arange(self, l: int) -> torch.Tensor: if l > len(self._arange_buf): next_pow_of_2 = 1 << (l - 1).bit_length() @@ -2318,7 +2307,6 @@ class DeepseekSparseAttnMultiStepBackend: def __init__( self, model_runner: ModelRunner, topk: int, speculative_num_steps: int ): - self.model_runner = model_runner self.topk = topk self.speculative_num_steps = speculative_num_steps self.attn_backends = [] diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 52a483802..278469a4b 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -761,9 +761,6 @@ class ModelRunner(ModelRunnerKVCacheMixin): if self.device == "cuda" or self.device == "musa": self.init_cublas() - self.init_attention_backend() - self.kernel_warmup() - # Init hisparse coordinator (must happen before CUDA graph capture) if self.enable_hisparse: from sglang.srt.managers.hisparse_coordinator import HiSparseCoordinator from sglang.srt.mem_cache.sparsity import parse_hisparse_config @@ -785,6 +782,8 @@ class ModelRunner(ModelRunnerKVCacheMixin): ), host_to_device_ratio=hisparse_cfg.host_to_device_ratio, ) + self.init_attention_backend() + self.kernel_warmup() self._pre_initialize_flashinfer_allreduce_workspace() self.init_device_graphs() elif self.device == "cpu": diff --git a/test/manual/attention/test_flashattn_backend.py b/test/manual/attention/test_flashattn_backend.py index 863d5e36e..466871df6 100644 --- a/test/manual/attention/test_flashattn_backend.py +++ b/test/manual/attention/test_flashattn_backend.py @@ -91,6 +91,7 @@ class MockModelRunner: device=self.device, enable_memory_saver=False, ) + self.hisparse_coordinator = None @unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA") diff --git a/test/manual/attention/test_flashattn_mla_backend.py b/test/manual/attention/test_flashattn_mla_backend.py index f1d53fcf9..fb44bf0f3 100644 --- a/test/manual/attention/test_flashattn_mla_backend.py +++ b/test/manual/attention/test_flashattn_mla_backend.py @@ -76,6 +76,7 @@ class MockModelRunner: device=self.device, enable_memory_saver=False, ) + self.hisparse_coordinator = None class MockReqToTokenPool: diff --git a/test/manual/attention/test_trtllm_mla_backend.py b/test/manual/attention/test_trtllm_mla_backend.py index e800ed582..25cfe2a6d 100755 --- a/test/manual/attention/test_trtllm_mla_backend.py +++ b/test/manual/attention/test_trtllm_mla_backend.py @@ -268,6 +268,7 @@ class MockModelRunner: device=self.device, enable_memory_saver=False, ) + self.hisparse_coordinator = None def compare_outputs(trtllm_out, reference_out, tolerance=1e-2): diff --git a/test/registered/kernels/test_dsa_indexer.py b/test/registered/kernels/test_dsa_indexer.py index 979d9e60d..83aff7d74 100644 --- a/test/registered/kernels/test_dsa_indexer.py +++ b/test/registered/kernels/test_dsa_indexer.py @@ -252,6 +252,7 @@ class MockModelRunner: "dsa_decode_backend": "fa3", }, )() + self.hisparse_coordinator = None @unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")