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

[Bugfix] Fix hidden state size in EAGLE PD disaggregation buffers (#13590)

Browse Source 
Co-authored-by: ZeldaHuang <hzm414167@alibaba-inc.com>
Co-authored-by: Shangming Cai <csmthu@gmail.com>
This commit is contained in:
Michele Marzollo
2025-11-21 08:32:40 +01:00
committed by GitHub
parent 90a0133515
commit b30f63c40f
1 changed files with 21 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
python/sglang/srt/managers/scheduler.py
View File

@@ -338,7 +338,7 @@ class Scheduler(
# Draft workers are looked up via `SpeculativeAlgorithm` registry; new
# algorithms should register their factory instead of patching this code.
if self.spec_algorithm.name in {"EAGLE", "EAGLE3"}:
if self.spec_algorithm.is_eagle():
draft_worker_kwargs["enable_overlap"] = self.enable_overlap
self.draft_worker = self.spec_algorithm.create_draft_worker(
**draft_worker_kwargs
@@ -864,8 +864,16 @@ class Scheduler(
)
self.disagg_metadata_buffers = MetadataBuffers(
buffer_size,
hidden_size=self.model_config.hf_text_config.hidden_size,
hidden_states_dtype=self.model_config.dtype,
hidden_size=(
self.draft_worker.model_config.hidden_size
if self.spec_algorithm.is_eagle()
else 16 # minimal padding size for RDMA
),
hidden_states_dtype=(
self.draft_worker.model_config.dtype
if self.spec_algorithm.is_eagle()
else torch.float32
),
custom_mem_pool=self.token_to_kv_pool_allocator.get_kvcache().maybe_get_custom_mem_pool(),
)
@@ -909,8 +917,16 @@ class Scheduler(
)
self.disagg_metadata_buffers = MetadataBuffers(
buffer_size,
hidden_size=self.model_config.hf_text_config.hidden_size,
hidden_states_dtype=self.model_config.dtype,
hidden_size=(
self.draft_worker.model_config.hidden_size
if self.spec_algorithm.is_eagle()
else 16 # minimal padding size for RDMA
),
hidden_states_dtype=(
self.draft_worker.model_config.dtype
if self.spec_algorithm.is_eagle()
else torch.float32
),
custom_mem_pool=self.token_to_kv_pool_allocator.get_kvcache().maybe_get_custom_mem_pool(),
)