dynamo-kv-router 1.3.0

// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

use dynamo_tokens::SequenceHash;
use indexmap::IndexMap;
use parking_lot::RwLock;
#[cfg(test)]
use rustc_hash::FxHashSet;
use rustc_hash::{FxBuildHasher, FxHashMap};
use seqlock::SeqLock;
use std::collections::HashMap;
use std::sync::Arc;
#[cfg(test)]
use std::sync::atomic::{AtomicUsize, Ordering};
use tokio::time::Instant;

use super::PrefillTokenDeltas;
use super::prefill_tracker::{PrefillLoadSnapshot, PrefillTimeLoadError};
use super::prompt_membership_trie::{PromptMembershipTrie, WorkerLookup};
use super::single::PromptMembershipDelta;
use super::topology::WorkerTopologyChange;
use crate::protocols::WorkerWithDpRank;

/// Ephemeral, request-specific view of worker load.
///
/// Values are materialized for one incoming request at a specific instant and
/// should be passed to worker selection immediately. Do not persist, reuse, or
/// publish this projection as durable worker state.
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
pub struct WorkerLoadProjection {
    pub active_prefill_tokens: usize,
    pub active_decode_blocks: usize,
    /// Request blocks not already shared with active sequences on this worker.
    ///
    /// These blocks may still exist in an inactive cache; this field describes
    /// additional active block footprint, not cache misses.
    pub additional_active_blocks: usize,
}

impl WorkerLoadProjection {
    pub fn potential_decode_blocks(self) -> usize {
        self.active_decode_blocks + self.additional_active_blocks
    }
}

pub type PotentialLoadMaps = (
    FxHashMap<WorkerWithDpRank, usize>,
    FxHashMap<WorkerWithDpRank, usize>,
    Option<FxHashMap<WorkerWithDpRank, usize>>,
);

/// Reusable snapshot of a worker's currently tracked execution state.
///
/// `active_blocks` is the worker's unique active decode load in blocks.
/// `prefill` retains the decay state needed to evaluate active prefill load in
/// tokens at a caller-provided instant. Unlike [`WorkerLoadProjection`], this
/// snapshot contains no incoming-request-specific state.
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
pub(super) struct WorkerLoadSnapshot {
    pub(super) active_blocks: usize,
    pub(super) active_requests: usize,
    pub(super) prefill: PrefillLoadSnapshot,
}

impl WorkerLoadSnapshot {
    pub(super) fn active_tokens(&self, decay_now: Instant) -> usize {
        self.prefill.active_tokens_at(decay_now)
    }

    pub(super) fn modeled_remaining_prefill_time_ms(
        &self,
        now: Instant,
    ) -> Result<u64, PrefillTimeLoadError> {
        self.prefill.modeled_remaining_prefill_time_ms_at(now)
    }
}

#[derive(Debug)]
struct WorkerLoadSlot {
    // SeqLock gives the hot projection path a lock-free read of the latest
    // whole-worker snapshot. Writers mark the sequence odd, replace the
    // snapshot, then mark the sequence even. Readers only return after seeing
    // the same even sequence before and after copying the `Copy` payload, so a
    // copy that overlaps a write is retried instead of exposing a mixed
    // snapshot. Writers still serialize with each other, but they do not wait
    // for readers.
    //
    // The upstream crate implements the copy with `read_volatile`, which is
    // technically UB under Rust/LLVM if a writer mutates the same memory at the
    // same time: volatile is not atomic and does not by itself legalize a data
    // race. See https://github.com/Amanieu/seqlock/issues/2#issuecomment-473606523.
    //
    // `WorkerLoadSnapshot` is a small `Copy` value with no references or drop
    // state. The sequence check is what makes the read logically race-free for
    // this slot: a reader either observes a stable whole snapshot or retries.
    // In practice this is the same seqlock/READ_ONCE pattern the crate is
    // designed for, while avoiding the per-worker RwLock read on every request.
    load: SeqLock<WorkerLoadSnapshot>,
}

impl WorkerLoadSlot {
    fn new(load: WorkerLoadSnapshot) -> Self {
        Self {
            load: SeqLock::new(load),
        }
    }

    fn snapshot(&self) -> WorkerLoadSnapshot {
        self.load.read()
    }

    fn replace(&self, load: WorkerLoadSnapshot) {
        *self.load.lock_write() = load;
    }
}

#[derive(Debug)]
struct WorkerLoadTable {
    // IndexMap gives us the dense full-worker scan plus point lookup shape that was previously
    // hand-rolled as Vec<WorkerLoadSlot> + FxHashMap<WorkerWithDpRank, usize>.
    entries: IndexMap<WorkerWithDpRank, WorkerLoadSlot, FxBuildHasher>,
}

impl Default for WorkerLoadTable {
    fn default() -> Self {
        Self {
            entries: IndexMap::with_hasher(FxBuildHasher),
        }
    }
}

impl WorkerLoadTable {
    fn len(&self) -> usize {
        self.entries.len()
    }

    fn iter(&self) -> impl Iterator<Item = (WorkerWithDpRank, WorkerLoadSnapshot)> + '_ {
        self.entries
            .iter()
            .map(|(&worker, slot)| (worker, slot.snapshot()))
    }

    fn ensure_worker(&mut self, worker: WorkerWithDpRank) {
        self.entries
            .entry(worker)
            .or_insert_with(|| WorkerLoadSlot::new(WorkerLoadSnapshot::default()));
    }

    fn update(&self, worker: WorkerWithDpRank, load: WorkerLoadSnapshot) -> bool {
        let Some(slot) = self.entries.get(&worker) else {
            return false;
        };
        slot.replace(load);
        true
    }

    fn upsert(&mut self, worker: WorkerWithDpRank, load: WorkerLoadSnapshot) {
        if let Some(slot) = self.entries.get(&worker) {
            slot.replace(load);
        } else {
            self.entries.insert(worker, WorkerLoadSlot::new(load));
        }
    }

    fn remove(&mut self, worker: WorkerWithDpRank) {
        self.entries.swap_remove(&worker);
    }
}

pub(super) struct PromptRegistry {
    // WARNING: prompt membership and worker load are only eventually consistent.
    // Each mutation still starts from one worker-local source of truth: we mutate the chosen
    // `ActiveSequences`, derive an exact `PromptMembershipDelta` plus `WorkerLoadSnapshot`, then
    // publish them separately here. The trie and load map converge to the correct final state
    // after the write finishes, but reads can still observe a mixed membership/load state that
    // never existed atomically and make a suboptimal routing choice.
    membership: PromptMembershipTrie,
    loads: RwLock<WorkerLoadTable>,
    #[cfg(test)]
    cleanup_attempts: AtomicUsize,
}

impl Default for PromptRegistry {
    fn default() -> Self {
        Self {
            membership: PromptMembershipTrie::new(),
            loads: RwLock::new(WorkerLoadTable::default()),
            #[cfg(test)]
            cleanup_attempts: AtomicUsize::new(0),
        }
    }
}

impl PromptRegistry {
    pub(super) fn new(workers: impl IntoIterator<Item = WorkerWithDpRank>) -> Self {
        let registry = Self::default();
        let mut loads = registry.loads.write();
        for worker in workers {
            loads.ensure_worker(worker);
        }
        drop(loads);
        registry
    }

    pub(super) fn replace_worker_load_state(
        &self,
        worker: WorkerWithDpRank,
        load: WorkerLoadSnapshot,
    ) {
        self.upsert_worker_load(worker, load);
    }

    pub(super) fn apply_membership_delta_and_load(
        &self,
        worker: WorkerWithDpRank,
        lookup: &Arc<parking_lot::RwLock<WorkerLookup>>,
        delta: PromptMembershipDelta,
        load: WorkerLoadSnapshot,
    ) {
        self.apply_membership_delta_and_load_without_cleanup(worker, lookup, delta, load);
        self.maybe_cleanup();
    }

    pub(super) fn apply_membership_delta_and_load_without_cleanup(
        &self,
        worker: WorkerWithDpRank,
        lookup: &Arc<parking_lot::RwLock<WorkerLookup>>,
        delta: PromptMembershipDelta,
        load: WorkerLoadSnapshot,
    ) {
        for remove in delta.removes {
            self.membership.remove_chain(worker, lookup, &remove.hashes);
        }
        for store in delta.stores {
            self.membership
                .store_chain(worker, lookup, store.parent, &store.hashes);
        }
        self.upsert_worker_load(worker, load);
    }

    fn upsert_worker_load(&self, worker: WorkerWithDpRank, load: WorkerLoadSnapshot) {
        if self.loads.read().update(worker, load) {
            return;
        }
        self.loads.write().upsert(worker, load);
    }

    pub(super) fn maybe_cleanup(&self) {
        #[cfg(test)]
        self.cleanup_attempts.fetch_add(1, Ordering::Relaxed);
        self.membership.maybe_cleanup();
    }

    #[cfg(test)]
    pub(super) fn cleanup_attempts(&self) -> usize {
        self.cleanup_attempts.load(Ordering::Relaxed)
    }

    pub(super) fn apply_topology_change(&self, change: WorkerTopologyChange) {
        for removed in change.removed {
            self.membership
                .remove_worker(removed.worker, &removed.trie_lookup);
            self.loads.write().remove(removed.worker);
        }

        for worker in change.added {
            self.loads.write().ensure_worker(worker);
        }
        self.membership.maybe_cleanup();
    }

    fn project_loads_from_membership<const INCLUDE_ACTIVE_REQUESTS: bool>(
        &self,
        query_len: usize,
        matched_depth: &FxHashMap<WorkerWithDpRank, usize>,
        prefill_token_deltas: &PrefillTokenDeltas,
        decay_now: Instant,
    ) -> PotentialLoadMaps {
        let loads = self.loads.read();
        let mut potential_blocks = FxHashMap::with_capacity_and_hasher(loads.len(), FxBuildHasher);
        let mut potential_tokens = FxHashMap::with_capacity_and_hasher(loads.len(), FxBuildHasher);
        let mut active_requests = INCLUDE_ACTIVE_REQUESTS
            .then(|| FxHashMap::with_capacity_and_hasher(loads.len(), FxBuildHasher));

        for (worker, load) in loads.iter() {
            let overlap_depth = matched_depth.get(&worker).copied().unwrap_or(0);
            let new_blocks = query_len.saturating_sub(overlap_depth);
            let active_tokens = load.active_tokens(decay_now);
            let added_tokens = prefill_token_deltas.tokens_for(worker);

            potential_blocks.insert(worker, load.active_blocks + new_blocks);
            potential_tokens.insert(worker, active_tokens + added_tokens);
            if let Some(active_requests) = active_requests.as_mut() {
                active_requests.insert(worker, load.active_requests);
            }
        }

        (potential_blocks, potential_tokens, active_requests)
    }

    pub(super) fn potential_blocks_and_tokens<const INCLUDE_ACTIVE_REQUESTS: bool>(
        &self,
        token_sequence: Option<&[SequenceHash]>,
        prefill_token_deltas: &PrefillTokenDeltas,
        decay_now: Instant,
    ) -> PotentialLoadMaps {
        let query_len = token_sequence.map_or(0, |query| query.len());
        let matched_depth = self.membership.compute_overlap_depths(token_sequence);
        self.project_loads_from_membership::<INCLUDE_ACTIVE_REQUESTS>(
            query_len,
            &matched_depth,
            prefill_token_deltas,
            decay_now,
        )
    }

    pub(super) fn project_worker_loads(
        &self,
        token_sequence: Option<&[SequenceHash]>,
        decay_now: Instant,
    ) -> FxHashMap<WorkerWithDpRank, WorkerLoadProjection> {
        let query_len = token_sequence.map_or(0, |query| query.len());
        let matched_depth = self.membership.compute_overlap_depths(token_sequence);
        let loads = self.loads.read();
        let mut projections = FxHashMap::with_capacity_and_hasher(loads.len(), FxBuildHasher);

        for (worker, load) in loads.iter() {
            let overlap_depth = matched_depth.get(&worker).copied().unwrap_or(0);
            projections.insert(
                worker,
                WorkerLoadProjection {
                    active_prefill_tokens: load.active_tokens(decay_now),
                    active_decode_blocks: load.active_blocks,
                    additional_active_blocks: query_len.saturating_sub(overlap_depth),
                },
            );
        }

        projections
    }

    pub(super) fn active_blocks(&self) -> HashMap<WorkerWithDpRank, usize> {
        self.loads
            .read()
            .iter()
            .map(|(worker, load)| (worker, load.active_blocks))
            .collect()
    }

    pub(super) fn active_request_counts(&self) -> HashMap<WorkerWithDpRank, usize> {
        self.loads
            .read()
            .iter()
            .map(|(worker, load)| (worker, load.active_requests))
            .collect()
    }

    pub(super) fn active_tokens(&self, decay_now: Instant) -> HashMap<WorkerWithDpRank, usize> {
        self.loads
            .read()
            .iter()
            .map(|(worker, load)| (worker, load.active_tokens(decay_now)))
            .collect()
    }

    pub(super) fn modeled_remaining_prefill_times_ms(
        &self,
        now: Instant,
    ) -> HashMap<WorkerWithDpRank, Result<u64, PrefillTimeLoadError>> {
        self.loads
            .read()
            .iter()
            .map(|(worker, load)| (worker, load.modeled_remaining_prefill_time_ms(now)))
            .collect()
    }

    pub(super) fn any_worker_matches_active_tokens(
        &self,
        decay_now: Instant,
        mut predicate: impl FnMut(WorkerWithDpRank, usize) -> bool,
    ) -> bool {
        self.loads
            .read()
            .iter()
            .any(|(worker, load)| predicate(worker, load.active_tokens(decay_now)))
    }

    #[cfg(test)]
    pub(super) fn assert_consistent_with_workers(
        &self,
        expected_loads: &FxHashMap<WorkerWithDpRank, WorkerLoadSnapshot>,
        expected_blocks: &FxHashMap<WorkerWithDpRank, FxHashSet<SequenceHash>>,
    ) {
        let actual_loads: FxHashMap<_, _> = self.loads.read().iter().collect();
        let actual_blocks = self.membership.worker_hashes();
        assert_eq!(
            actual_loads, *expected_loads,
            "prompt registry worker loads drifted from per-worker state",
        );
        assert_eq!(
            actual_blocks, *expected_blocks,
            "prompt registry prompt membership drifted from per-worker state",
        );
    }

    #[cfg(any(test, feature = "bench"))]
    pub(super) fn is_block_index_empty(&self) -> bool {
        self.membership.is_empty()
    }
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use rustc_hash::{FxHashMap, FxHashSet};

    use super::*;
    use crate::protocols::WorkerWithDpRank;
    use crate::sequences::prefill_tracker::AnchoredPrefillSnapshot;
    use crate::sequences::single::{PromptMembershipRemove, PromptMembershipStore};
    use crate::sequences::topology::RemovedWorkerState;

    fn worker(worker_id: u64, dp_rank: u32) -> WorkerWithDpRank {
        WorkerWithDpRank::new(worker_id, dp_rank)
    }

    fn lookup() -> Arc<parking_lot::RwLock<WorkerLookup>> {
        Arc::new(parking_lot::RwLock::new(WorkerLookup::default()))
    }

    fn store(parent: Option<SequenceHash>, hashes: &[SequenceHash]) -> PromptMembershipDelta {
        PromptMembershipDelta {
            stores: vec![PromptMembershipStore {
                parent,
                hashes: hashes.to_vec(),
            }],
            removes: Vec::new(),
        }
    }

    fn worker_load_snapshot(active_blocks: usize) -> WorkerLoadSnapshot {
        WorkerLoadSnapshot {
            active_blocks,
            active_requests: 0,
            prefill: PrefillLoadSnapshot::default(),
        }
    }

    fn anchored_load_snapshot(
        active_blocks: usize,
        prefill_full_tokens_sum: usize,
        anchored_tokens: usize,
        expected_prefill_duration: Option<Duration>,
        anchored_since: Instant,
    ) -> WorkerLoadSnapshot {
        WorkerLoadSnapshot {
            active_blocks,
            active_requests: 0,
            prefill: PrefillLoadSnapshot {
                prefill_full_tokens_sum,
                anchored_prefill: Some(AnchoredPrefillSnapshot {
                    initial_effective_prefill_tokens: anchored_tokens,
                    expected_prefill_duration,
                    anchored_since,
                }),
                total_modeled_prefill_time_ms: expected_prefill_duration
                    .map(|duration| duration.as_millis().min(u64::MAX as u128) as u64),
            },
        }
    }

    fn hash_set(hashes: &[SequenceHash]) -> FxHashSet<SequenceHash> {
        hashes.iter().copied().collect()
    }

    fn naive_potential_loads(
        expected_loads: &FxHashMap<WorkerWithDpRank, WorkerLoadSnapshot>,
        expected_blocks: &FxHashMap<WorkerWithDpRank, FxHashSet<SequenceHash>>,
        token_sequence: Option<&[SequenceHash]>,
        prefill_token_deltas: &PrefillTokenDeltas,
        decay_now: Instant,
    ) -> (
        FxHashMap<WorkerWithDpRank, usize>,
        FxHashMap<WorkerWithDpRank, usize>,
    ) {
        let mut potential_blocks =
            FxHashMap::with_capacity_and_hasher(expected_loads.len(), FxBuildHasher);
        let mut potential_tokens =
            FxHashMap::with_capacity_and_hasher(expected_loads.len(), FxBuildHasher);

        for (&worker, load) in expected_loads {
            let overlap_depth = token_sequence.map_or(0, |query| {
                let worker_blocks = expected_blocks
                    .get(&worker)
                    .expect("worker must have a prompt membership entry");
                query
                    .iter()
                    .position(|hash| !worker_blocks.contains(hash))
                    .unwrap_or(query.len())
            });
            let new_blocks =
                token_sequence.map_or(0, |query| query.len().saturating_sub(overlap_depth));
            let added_tokens = prefill_token_deltas.tokens_for(worker);

            potential_blocks.insert(worker, load.active_blocks + new_blocks);
            potential_tokens.insert(worker, load.active_tokens(decay_now) + added_tokens);
        }

        (potential_blocks, potential_tokens)
    }

    #[test]
    fn removed_hash_can_be_restored_by_later_store() {
        let worker = worker(1, 0);
        let registry = PromptRegistry::new([worker]);
        let lookup = lookup();
        let mut expected_loads = FxHashMap::default();
        let mut expected_blocks = FxHashMap::default();

        registry.apply_membership_delta_and_load(
            worker,
            &lookup,
            store(None, &[42]),
            worker_load_snapshot(1),
        );
        let load = worker_load_snapshot(1);
        registry.apply_membership_delta_and_load(
            worker,
            &lookup,
            PromptMembershipDelta {
                removes: vec![PromptMembershipRemove { hashes: vec![42] }],
                ..Default::default()
            },
            load,
        );
        registry.apply_membership_delta_and_load(worker, &lookup, store(None, &[42]), load);
        expected_loads.insert(worker, load);
        expected_blocks.insert(worker, hash_set(&[42]));

        registry.assert_consistent_with_workers(&expected_loads, &expected_blocks);
    }

    #[test]
    fn staggered_prefix_overlap_matches_naive_projection() {
        let worker_a = worker(1, 0);
        let worker_b = worker(2, 0);
        let worker_c = worker(3, 0);
        let registry = PromptRegistry::new([worker_a, worker_b, worker_c]);
        let lookup_a = lookup();
        let lookup_b = lookup();
        let lookup_c = lookup();
        let decay_now = Instant::now();
        let full_prompt: Vec<SequenceHash> = (1_u64..=96).collect();
        let mut expected_loads = FxHashMap::default();
        let mut expected_blocks = FxHashMap::default();

        for (worker, lookup, prompt_len) in [
            (worker_a, &lookup_a, 96usize),
            (worker_b, &lookup_b, 64),
            (worker_c, &lookup_c, 33),
        ] {
            let blocks = full_prompt[..prompt_len].to_vec();
            let load = worker_load_snapshot(prompt_len);
            registry.apply_membership_delta_and_load(worker, lookup, store(None, &blocks), load);
            expected_loads.insert(worker, load);
            expected_blocks.insert(worker, blocks.into_iter().collect());
        }

        registry.assert_consistent_with_workers(&expected_loads, &expected_blocks);

        let expected = naive_potential_loads(
            &expected_loads,
            &expected_blocks,
            Some(&full_prompt),
            &PrefillTokenDeltas::none(),
            decay_now,
        );
        let actual = registry.potential_blocks_and_tokens::<false>(
            Some(&full_prompt),
            &PrefillTokenDeltas::none(),
            decay_now,
        );

        assert_eq!(actual.0, expected.0);
        assert_eq!(actual.1, expected.1);
    }

    #[test]
    fn load_only_update_preserves_prompt_membership_and_active_token_projection() {
        let worker = worker(1, 0);
        let registry = PromptRegistry::new([worker]);
        let lookup = lookup();
        let now = Instant::now();
        let anchored_since = now.checked_sub(Duration::from_secs(3)).unwrap_or(now);
        let mut expected_loads = FxHashMap::default();
        let mut expected_blocks = FxHashMap::default();

        registry.apply_membership_delta_and_load(
            worker,
            &lookup,
            store(None, &[1, 2, 3]),
            worker_load_snapshot(3),
        );
        expected_blocks.insert(worker, hash_set(&[1, 2, 3]));

        let updated_load =
            anchored_load_snapshot(5, 12, 10, Some(Duration::from_secs(10)), anchored_since);
        registry.replace_worker_load_state(worker, updated_load);
        expected_loads.insert(worker, updated_load);

        registry.assert_consistent_with_workers(&expected_loads, &expected_blocks);
        assert_eq!(registry.active_tokens(now).get(&worker).copied(), Some(9));

        let actual = registry.potential_blocks_and_tokens::<false>(
            Some(&[1, 2, 3]),
            &PrefillTokenDeltas::none(),
            now,
        );
        assert_eq!(actual.0.get(&worker).copied(), Some(5));
        assert_eq!(actual.1.get(&worker).copied(), Some(9));
    }

    #[test]
    fn removing_worker_clears_prompt_membership_and_load_state() {
        let worker_a = worker(1, 0);
        let worker_b = worker(2, 0);
        let registry = PromptRegistry::new([worker_a, worker_b]);
        let lookup_a = lookup();
        let lookup_b = lookup();
        let mut expected_loads = FxHashMap::default();
        let mut expected_blocks = FxHashMap::default();

        let load_a = worker_load_snapshot(3);
        let load_b = worker_load_snapshot(2);
        registry.apply_membership_delta_and_load(
            worker_a,
            &lookup_a,
            store(None, &[1, 2, 3]),
            load_a,
        );
        registry.apply_membership_delta_and_load(worker_b, &lookup_b, store(None, &[1, 2]), load_b);
        expected_loads.insert(worker_a, load_a);
        expected_loads.insert(worker_b, load_b);
        expected_blocks.insert(worker_a, hash_set(&[1, 2, 3]));
        expected_blocks.insert(worker_b, hash_set(&[1, 2]));

        registry.apply_topology_change(WorkerTopologyChange {
            added: Vec::new(),
            removed: vec![RemovedWorkerState {
                worker: worker_a,
                trie_lookup: Arc::clone(&lookup_a),
            }],
        });
        expected_loads.remove(&worker_a);
        expected_blocks.remove(&worker_a);

        registry.assert_consistent_with_workers(&expected_loads, &expected_blocks);
        assert!(!registry.active_blocks().contains_key(&worker_a));

        let actual = registry.potential_blocks_and_tokens::<false>(
            Some(&[1, 2, 3]),
            &PrefillTokenDeltas::none(),
            Instant::now(),
        );
        assert_eq!(actual.0.get(&worker_b).copied(), Some(3));
    }

    #[test]
    fn dp_ranks_with_same_worker_id_remain_isolated() {
        let worker_a = worker(1, 0);
        let worker_b = worker(1, 1);
        let registry = PromptRegistry::new([worker_a, worker_b]);
        let lookup_a = lookup();
        let lookup_b = lookup();
        let decay_now = Instant::now();
        let mut expected_loads = FxHashMap::default();
        let mut expected_blocks = FxHashMap::default();

        let load_a = worker_load_snapshot(3);
        let load_b = worker_load_snapshot(1);
        registry.apply_membership_delta_and_load(
            worker_a,
            &lookup_a,
            store(None, &[1, 2, 3]),
            load_a,
        );
        registry.apply_membership_delta_and_load(worker_b, &lookup_b, store(None, &[1]), load_b);
        expected_loads.insert(worker_a, load_a);
        expected_loads.insert(worker_b, load_b);
        expected_blocks.insert(worker_a, hash_set(&[1, 2, 3]));
        expected_blocks.insert(worker_b, hash_set(&[1]));

        registry.assert_consistent_with_workers(&expected_loads, &expected_blocks);

        let expected = naive_potential_loads(
            &expected_loads,
            &expected_blocks,
            Some(&[1, 2, 3]),
            &PrefillTokenDeltas::none(),
            decay_now,
        );
        let actual = registry.potential_blocks_and_tokens::<false>(
            Some(&[1, 2, 3]),
            &PrefillTokenDeltas::none(),
            decay_now,
        );

        assert_eq!(actual.0, expected.0);
        assert_eq!(actual.1, expected.1);
    }
}