Unverified Commit db8e52f2 authored by Yan Ru Pei's avatar Yan Ru Pei Committed by GitHub
Browse files

fix: correct prefill/decode block defaults when no overlaps (#3811)


Signed-off-by: default avatarPeaBrane <yanrpei@gmail.com>
parent 87ea346d
...@@ -826,17 +826,22 @@ impl ActiveSequencesMultiWorker { ...@@ -826,17 +826,22 @@ impl ActiveSequencesMultiWorker {
let token_sequence_shared = token_sequence.map(Arc::new); let token_sequence_shared = token_sequence.map(Arc::new);
let mut receivers = Vec::new(); let mut receivers = Vec::new();
// Iterate through overlaps to process each WorkerWithDpRank // Iterate through all workers, not just those with overlap
for (worker, overlap) in overlaps.scores.iter() { // This ensures we properly account for active tokens/blocks on all workers
// Check if the worker has a sender for sender_entry in self.senders.iter() {
if let Some(sender) = self.senders.get(worker) { let worker = *sender_entry.key();
let sender = sender_entry.value();
// Get overlap for this worker (defaults to 0 if not in overlaps)
let overlap = *overlaps.scores.get(&worker).unwrap_or(&0);
let (resp_tx, resp_rx) = tokio::sync::oneshot::channel(); let (resp_tx, resp_rx) = tokio::sync::oneshot::channel();
receivers.push((*worker, resp_rx)); receivers.push((worker, resp_rx));
if let Err(e) = sender.send(UpdateSequences::PotentialBlocksAndTokens { if let Err(e) = sender.send(UpdateSequences::PotentialBlocksAndTokens {
token_sequence: token_sequence_shared.clone(), token_sequence: token_sequence_shared.clone(),
isl, isl,
overlap: *overlap, overlap,
resp_tx, resp_tx,
}) { }) {
tracing::error!( tracing::error!(
...@@ -846,7 +851,6 @@ impl ActiveSequencesMultiWorker { ...@@ -846,7 +851,6 @@ impl ActiveSequencesMultiWorker {
); );
} }
} }
}
// Collect results from all workers // Collect results from all workers
for (worker, receiver) in receivers { for (worker, receiver) in receivers {
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment