schedule dspark slim

[zhangxiaolei123456]

Motivation

Modifications

Accuracy Tests

Speed Tests and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ❌ Run #29811258641
Latest PR Test (Extra): ❌ Run #29811258402

[gemini-code-assist]

Code Review

This pull request refactors disaggregation state management by encapsulating various pd_hidden attributes into a new PDHiddenReqState class and moving init_disaggregation from scheduler.py to prefill.py. Feedback highlights a critical missing import of DisaggregationMode in prefill.py that would cause a runtime error. Additionally, the reviewer suggests optimizing performance by storing PDHiddenReqState as a dynamic attribute on Req instead of using a WeakKeyDictionary, and caching the state lookup in local variables to avoid redundant lookups.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Using a global weakref.WeakKeyDictionary introduces significant overhead for weak reference creation, dictionary lookups, and garbage collection tracking in the hot path of request scheduling.

Instead, we can store the PDHiddenReqState as a dynamic attribute directly on the Req object (e.g., req.pd_hidden_state). This keeps the Req class definition slim (no extra attributes initialized by default) while providing much faster, thread-safe, and direct access.

def get_pd_hidden_req_state(req: "Req") -> PDHiddenReqState:
    state = getattr(req, "pd_hidden_state", None)
    if state is None:
        state = PDHiddenReqState()
        req.pd_hidden_state = state
    return state

[gemini-code-assist]

The DisaggregationMode enum is used in init_disaggregation (e.g., self.disaggregation_mode == DisaggregationMode.DECODE), but it is not imported in this file. This will cause a NameError at runtime. Please add DisaggregationMode to the imports from sglang.srt.disaggregation.utils.

Suggested change

	MetadataBuffers,
	ReqToMetadataIdxAllocator,
	TransferBackend,
	DisaggregationMode,
	MetadataBuffers,
	ReqToMetadataIdxAllocator,
	TransferBackend,

[gemini-code-assist]

If we use a dynamic attribute on Req instead of WeakKeyDictionary, the weakref import is no longer needed and can be removed.

Suggested change

	import threading
	import weakref
	import threading

[gemini-code-assist]

To avoid multiple lookups of the request state (whether via WeakKeyDictionary or getattr), retrieve the state once into a local variable and update its attributes.

Suggested change

	pd_hidden_state(req).meta = None
	pd_hidden_state(req).src_indices = None
	pd_hidden_state(req).dst_indices = None
	pd_hidden_state(req).written = None
	pd_hidden_state(req).capture_layer_ids = None
	pd_hidden_state(req).current_src_indices = None
	pd_hidden_state(req).current_start = None
	pd_hidden_state(req).current_row_len = 0
	pd_hidden_state(req).current_is_last = False
	pd_hidden_state(req).owner_direct_sent = False
	state = pd_hidden_state(req)
	state.meta = None
	state.src_indices = None
	state.dst_indices = None
	state.written = None
	state.capture_layer_ids = None
	state.current_src_indices = None
	state.current_start = None
	state.current_row_len = 0
	state.current_is_last = False
	state.owner_direct_sent = False

[gemini-code-assist]

Retrieve the request state once into a local variable to avoid multiple redundant lookups.

Suggested change

	pd_hidden_state(req).meta = dict(dspark_meta)
	pd_hidden_state(req).src_indices = []
	pd_hidden_state(req).dst_indices = []
	pd_hidden_state(req).written = []
	pd_hidden_state(req).owner_direct_sent = False
	state = pd_hidden_state(req)
	state.meta = dict(dspark_meta)
	state.src_indices = []
	state.dst_indices = []
	state.written = []
	state.owner_direct_sent = False

[gemini-code-assist]

Retrieve the request state once into a local variable to avoid multiple redundant lookups.

Suggested change

	pd_hidden_state(req).capture_layer_ids = [int(x) for x in plan.local_layer_ids]
	pd_hidden_state(req).meta = dict(dspark_meta)
	pd_hidden_state(req).src_indices = src_indices
	pd_hidden_state(req).dst_indices = plan.dst_indices
	pd_hidden_state(req).written = (
	None if plan.streaming_hidden else [False] * plan.hidden_len
	)
	req.pd_hidden_owner_direct_sent = False
	pd_hidden_state(req).owner_direct_sent = False
	state = pd_hidden_state(req)
	state.capture_layer_ids = [int(x) for x in plan.local_layer_ids]
	state.meta = dict(dspark_meta)
	state.src_indices = src_indices
	state.dst_indices = plan.dst_indices
	state.written = (
	None if plan.streaming_hidden else [False] * plan.hidden_len
	)
	state.owner_direct_sent = False

[gemini-code-assist]

Retrieve the request state once into a local variable to avoid multiple redundant lookups.

Suggested change

	if streaming_hidden:
	req.pd_hidden_src_indices = None
	req.pd_hidden_current_src_indices = None
	req.pd_hidden_current_start = None
	req.pd_hidden_current_row_len = 0
	req.pd_hidden_current_is_last = False
	req.pd_hidden_owner_direct_sent = True
	pd_hidden_state(req).src_indices = None
	pd_hidden_state(req).current_src_indices = None
	pd_hidden_state(req).current_start = None
	pd_hidden_state(req).current_row_len = 0
	pd_hidden_state(req).current_is_last = False
	pd_hidden_state(req).owner_direct_sent = True
	state = pd_hidden_state(req)
	if streaming_hidden:
	state.src_indices = None
	state.current_src_indices = None
	state.current_start = None
	state.current_row_len = 0
	state.current_is_last = False
	state.owner_direct_sent = True

[gemini-code-assist]

Retrieve the request state once into a local variable to avoid multiple redundant lookups.

            state = pd_hidden_state(req)
            state.current_start = write_start
            state.current_row_len = rows
            state.current_src_indices = write_indices
            state.current_is_last = write_end >= hidden_start + hidden_len
            written = state.written

[gemini-code-assist]

Retrieve the request state once into a local variable to avoid multiple redundant lookups.

Suggested change

	if has_current_pd_hidden and streaming_pd_hidden:
	req.pd_hidden_src_indices = None
	req.pd_hidden_current_src_indices = None
	req.pd_hidden_current_start = None
	req.pd_hidden_current_row_len = 0
	req.pd_hidden_current_is_last = False
	pd_hidden_state(req).src_indices = None
	pd_hidden_state(req).current_src_indices = None
	pd_hidden_state(req).current_start = None
	pd_hidden_state(req).current_row_len = 0
	pd_hidden_state(req).current_is_last = False
	state = pd_hidden_state(req)
	if has_current_pd_hidden and streaming_pd_hidden:
	state.src_indices = None
	state.current_src_indices = None
	state.current_start = None
	state.current_row_len = 0
	state.current_is_last = False