refactor(attn): init hisparse_coordinator before attn_backend; replace lazy property with init-time capture (#26012)

Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
2026-05-21 16:03:42 -07:00
committed by GitHub
co-authored by Claude Sonnet 4.6
parent c5251a98a9
commit d765dfd043
8 changed files with 9 additions and 42 deletions
@@ -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] = []
@@ -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] = []
@@ -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 = []
@@ -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":