[Bugfix] Fix hidden state size in EAGLE PD disaggregation buffers (#13590)
Co-authored-by: ZeldaHuang <hzm414167@alibaba-inc.com> Co-authored-by: Shangming Cai <csmthu@gmail.com>
This commit is contained in:
co-authored by
ZeldaHuang
Shangming Cai
parent
90a0133515
commit
b30f63c40f
@@ -338,7 +338,7 @@ class Scheduler(
|
|||||||
|
|
||||||
# Draft workers are looked up via `SpeculativeAlgorithm` registry; new
|
# Draft workers are looked up via `SpeculativeAlgorithm` registry; new
|
||||||
# algorithms should register their factory instead of patching this code.
|
# 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
|
draft_worker_kwargs["enable_overlap"] = self.enable_overlap
|
||||||
self.draft_worker = self.spec_algorithm.create_draft_worker(
|
self.draft_worker = self.spec_algorithm.create_draft_worker(
|
||||||
**draft_worker_kwargs
|
**draft_worker_kwargs
|
||||||
@@ -864,8 +864,16 @@ class Scheduler(
|
|||||||
)
|
)
|
||||||
self.disagg_metadata_buffers = MetadataBuffers(
|
self.disagg_metadata_buffers = MetadataBuffers(
|
||||||
buffer_size,
|
buffer_size,
|
||||||
hidden_size=self.model_config.hf_text_config.hidden_size,
|
hidden_size=(
|
||||||
hidden_states_dtype=self.model_config.dtype,
|
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(),
|
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(
|
self.disagg_metadata_buffers = MetadataBuffers(
|
||||||
buffer_size,
|
buffer_size,
|
||||||
hidden_size=self.model_config.hf_text_config.hidden_size,
|
hidden_size=(
|
||||||
hidden_states_dtype=self.model_config.dtype,
|
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(),
|
custom_mem_pool=self.token_to_kv_pool_allocator.get_kvcache().maybe_get_custom_mem_pool(),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user