chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

fix(disagg): tensor-safe prefix check + graceful per-layer register (lever A)

Browse Source 
E2E caught it: registered=8 (per-layer activated correctly) then the prefill
crashed with "Boolean value of Tensor with no values is ambiguous" — req.prefix_indices
is a tensor and `prefix or []` evaluated its truthiness. Use a None+len check
(safe for None/list/tensor), and wrap each request's registration in try/except so
a setup error degrades to the monolithic transfer instead of crashing the scheduler.
Also guard start_send_idx with getattr.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
孙远航
2026-06-07 01:15:23 +00:00
parent ae18e3adc8
commit 501faa8f5b
1 changed files with 23 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
python/sglang/srt/disaggregation/prefill.py
View File

@@ -659,18 +659,29 @@ class SchedulerDisaggregationPrefillMixin:
return
page_size = self.token_to_kv_pool_allocator.page_size
for req in batch.reqs:
if getattr(req, "disagg_kv_sender", None) is None:
continue
if getattr(req, "is_chunked", 0) != 0 or req.start_send_idx != 0:
continue # chunked / partially-sent: this forward isn't the full range
if len(getattr(req, "prefix_indices", []) or []) != 0:
continue # cached prefix is HiCache-loaded, not forward-written (lever B)
end_idx = min(len(req.fill_ids), len(req.origin_input_ids))
page_indices = _kv_locs_to_page_indices_cpu(
self.req_to_token_pool.req_to_token[req.req_pool_idx, 0:end_idx],
page_size,
)
kv_manager.register_per_layer_transfer(req.bootstrap_room, page_indices)
try:
if getattr(req, "disagg_kv_sender", None) is None:
continue
if getattr(req, "is_chunked", 0) != 0 or getattr(req, "start_send_idx", 0) != 0:
continue # chunked / partially-sent: this forward isn't the full range
prefix = getattr(req, "prefix_indices", None)
# NB: prefix_indices is a tensor — never use `or`/truthiness on it.
if prefix is not None and len(prefix) != 0:
continue # cached prefix is HiCache-loaded, not forward-written -> lever B
end_idx = min(len(req.fill_ids), len(req.origin_input_ids))
page_indices = _kv_locs_to_page_indices_cpu(
self.req_to_token_pool.req_to_token[req.req_pool_idx, 0:end_idx],
page_size,
)
kv_manager.register_per_layer_transfer(req.bootstrap_room, page_indices)
except Exception as e:
# Never let lever-A setup crash the scheduler; fall back to the
# monolithic post-forward transfer for this request.
logger.warning(
"[CP_PER_LAYER_TRANSFER] register skipped for room %s: %r",
getattr(req, "bootstrap_room", None),
e,
)
@torch.no_grad()
def event_loop_normal_disagg_prefill(self: Scheduler) -> None: