[spec-overlap] bugfix for pd disaggregation and npu (#14088)

Co-authored-by: Even Zhou <even.y.zhou@outlook.com>
This commit is contained in:
liupeng374
2025-12-01 22:58:20 +08:00
committed by GitHub
parent e55731b6e8
commit 2e8f54e61e
9 changed files with 32 additions and 38 deletions

View File

@@ -824,8 +824,10 @@ class Scheduler(
draft_token_to_kv_pool = (
self.draft_worker.draft_worker.draft_runner.token_to_kv_pool
)
model_config = self.draft_worker.draft_worker.draft_runner.model_config
else:
draft_token_to_kv_pool = self.draft_worker.model_runner.token_to_kv_pool
model_config = self.draft_worker.model_config
if (
self.disaggregation_mode == DisaggregationMode.DECODE
@@ -837,12 +839,12 @@ class Scheduler(
self.disagg_metadata_buffers = MetadataBuffers(
buffer_size,
hidden_size=(
self.draft_worker.model_config.hidden_size
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
model_config.dtype
if self.spec_algorithm.is_eagle()
else torch.float32
),
@@ -890,12 +892,12 @@ class Scheduler(
self.disagg_metadata_buffers = MetadataBuffers(
buffer_size,
hidden_size=(
self.draft_worker.model_config.hidden_size
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
model_config.dtype
if self.spec_algorithm.is_eagle()
else torch.float32
),