[spec-overlap] bugfix for pd disaggregation and npu (#14088)
Co-authored-by: Even Zhou <even.y.zhou@outlook.com>
This commit is contained in:
@@ -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
|
||||
),
|
||||
|
||||
Reference in New Issue
Block a user