[SMG] Fix cache-aware policy pool isolation in PD mode (#25184)

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

sgl-model-gateway/src/core/steps/worker/local/remove_from_policy_registry.rs

@@ -38,11 +38,17 @@ impl StepExecutor<WorkerRemovalWorkflowData> for RemoveFromPolicyRegistryStep {
             let model_id = worker.model_id().to_string();
             let worker_url = worker.url();
 
-            // Remove from cache-aware policy
+            // Remove from cache-aware policy (model_policies; covers the regular Router)
             app_context
                 .policy_registry
                 .remove_worker_from_cache_aware(&model_id, worker_url);
 
+            // PD mode keeps prefill/decode cache-aware policies separate from
+            // model_policies, so also drop the worker from the matching pool's policy.
+            app_context
+                .policy_registry
+                .remove_pd_worker_from_cache_aware(worker.as_ref());
+
             // Notify policy registry
             app_context.policy_registry.on_worker_removed(&model_id);
         }

sgl-model-gateway/src/core/steps/worker/shared/update_policies.rs

@@ -141,6 +141,23 @@ impl<D: WorkerRegistrationData + WorkflowData> StepExecutor<D> for UpdatePolicie
             }
         }
 
+        // Initialize cache-aware policies for PD mode (prefill_policy / decode_policy
+        // are separate instances from model_policies and are not touched by the
+        // per-model loop above). `init_workers` is idempotent so re-running on each
+        // registration is safe.
+        let decode_workers = app_context.worker_registry.get_decode_workers();
+        let prefill_is_cache_aware =
+            app_context.policy_registry.get_prefill_policy().name() == "cache_aware";
+        let decode_is_cache_aware =
+            app_context.policy_registry.get_decode_policy().name() == "cache_aware";
+        if (prefill_is_cache_aware && !prefill_workers.is_empty())
+            || (decode_is_cache_aware && !decode_workers.is_empty())
+        {
+            app_context
+                .policy_registry
+                .init_pd_cache_aware_policies(&prefill_workers, &decode_workers);
+        }
+
         debug!(
             "Updated policies for {} workers across {} models",
             workers.len(),

sgl-model-gateway/src/policies/cache_aware.rs

@@ -71,13 +71,40 @@ use super::{
     get_healthy_worker_indices, normalize_model_key, tree::Tree, utils::PeriodicTask,
     CacheAwareConfig, LoadBalancingPolicy, SelectWorkerInfo,
 };
-use crate::core::{Worker, UNKNOWN_MODEL_ID};
+use crate::core::{Worker, WorkerType, UNKNOWN_MODEL_ID};
+
+/// Tag used to isolate prefill/decode/regular worker pools in the cache_aware tree key.
+///
+/// Trees are keyed by `pool::model` so that an alternating prefill→decode call sequence
+/// for the same model cannot evict each other's tenants. Without this isolation, the
+/// `tree.insert(text, url)` at the end of every `select_worker` call would overwrite
+/// the previous pool's tenant for the same prompt and collapse cache_aware into a
+/// flip-flop between pools.
+fn pool_tag(worker_type: &WorkerType) -> &'static str {
+    match worker_type {
+        WorkerType::Regular => "regular",
+        WorkerType::Prefill { .. } => "prefill",
+        WorkerType::Decode => "decode",
+    }
+}
+
+fn make_tree_key(pool: &str, model: &str) -> String {
+    format!("{}::{}", pool, model)
+}
+
+fn tree_key_for_worker(worker: &dyn Worker) -> String {
+    make_tree_key(
+        pool_tag(worker.worker_type()),
+        normalize_model_key(worker.model_id()),
+    )
+}
 
 /// Cache-aware routing policy
 ///
 /// Routes requests based on cache affinity when load is balanced,
 /// switches to shortest-queue routing when load is imbalanced.
-/// Maintains separate trees per model for multi-model support.
+/// Maintains separate trees per `(pool, model)` so that prefill, decode, and
+/// regular worker pools cannot evict each other's tenants.
 /// Supports mesh synchronization of tree operations across cluster nodes.
 /// When mesh is not enabled, the policy works independently without synchronization.
 #[derive(Debug)]
@@ -106,13 +133,13 @@ impl CacheAwarePolicy {
                 "Eviction",
                 move || {
                     for tree_ref in trees_clone.iter() {
-                        let model_id = tree_ref.key();
+                        let tree_key = tree_ref.key();
                         let tree = tree_ref.value();
                         tree.evict_tenant_by_size(max_tree_size);
 
                         debug!(
-                            "Cache eviction completed for model {}, max_size: {}",
-                            model_id, max_tree_size
+                            "Cache eviction completed for {}, max_size: {}",
+                            tree_key, max_tree_size
                         );
                     }
                 },
@@ -139,24 +166,22 @@ impl CacheAwarePolicy {
 
     /// Initialize the tree with worker URLs (used only during initial setup)
     pub fn init_workers(&self, workers: &[Arc<dyn Worker>]) {
-        // Group workers by model
-        let mut model_workers: std::collections::HashMap<String, Vec<&Arc<dyn Worker>>> =
+        // Group workers by (pool, model) so each pool gets its own isolated tree.
+        let mut grouped: std::collections::HashMap<String, Vec<&Arc<dyn Worker>>> =
             std::collections::HashMap::new();
         for worker in workers {
-            let tree_key = normalize_model_key(worker.model_id());
-            model_workers
-                .entry(tree_key.to_string())
+            grouped
+                .entry(tree_key_for_worker(worker.as_ref()))
                 .or_default()
                 .push(worker);
         }
 
-        // Initialize tree for each model
-        for (tree_key, model_workers) in model_workers {
+        for (tree_key, pool_workers) in grouped {
             let tree = self
                 .trees
                 .entry(tree_key)
                 .or_insert_with(|| Arc::new(Tree::new()));
-            for worker in model_workers {
+            for worker in pool_workers {
                 tree.insert("", worker.url());
             }
         }
@@ -164,27 +189,18 @@ impl CacheAwarePolicy {
 
     /// Add a single worker to the tree (incremental update)
     pub fn add_worker(&self, worker: &dyn Worker) {
-        let tree_key = normalize_model_key(worker.model_id());
+        let tree_key = tree_key_for_worker(worker);
         let tree = self
             .trees
-            .entry(tree_key.to_string())
+            .entry(tree_key)
             .or_insert_with(|| Arc::new(Tree::new()));
         tree.insert("", worker.url());
     }
 
-    /// Add a worker by URL and model (for backward compatibility)
-    pub fn add_worker_by_url(&self, url: &str, model_id: &str) {
-        let tree = self
-            .trees
-            .entry(model_id.to_string())
-            .or_insert_with(|| Arc::new(Tree::new()));
-        tree.insert("", url);
-    }
-
     /// Remove a worker from the tree
     pub fn remove_worker(&self, worker: &dyn Worker) {
-        let tree_key = normalize_model_key(worker.model_id());
-        if let Some(tree) = self.trees.get(tree_key) {
+        let tree_key = tree_key_for_worker(worker);
+        if let Some(tree) = self.trees.get(&tree_key) {
             tree.remove_tenant(worker.url());
         }
     }
@@ -207,11 +223,11 @@ impl CacheAwarePolicy {
             // In a full implementation, we might want to query mesh for all tree states
 
             for tree_ref in self.trees.iter() {
-                let model_id = tree_ref.key();
-                if let Some(tree_state) = mesh_sync.get_tree_state(model_id) {
+                let tree_key = tree_ref.key();
+                if let Some(tree_state) = mesh_sync.get_tree_state(tree_key) {
                     debug!(
-                        "Restoring tree state for model {} with {} operations",
-                        model_id,
+                        "Restoring tree state for {} with {} operations",
+                        tree_key,
                         tree_state.operations.len()
                     );
 
@@ -232,20 +248,30 @@ impl CacheAwarePolicy {
         }
     }
 
-    /// Normalize model_id for mesh synchronization
-    /// Converts empty model_id to UNKNOWN_MODEL_ID for consistency
-    fn normalize_mesh_model_id(model_id: &str) -> &str {
-        if model_id.is_empty() {
+    /// Normalize a tree key for mesh synchronization, converting an accidentally
+    /// empty key to `UNKNOWN_MODEL_ID` for consistency. In current code the
+    /// composite `pool::model` key is never empty, so this is defensive.
+    fn normalize_mesh_model_id(tree_key: &str) -> &str {
+        if tree_key.is_empty() {
             UNKNOWN_MODEL_ID
         } else {
-            model_id
+            tree_key
         }
     }
 
-    /// Apply remote tree operation from mesh
-    /// This is called when receiving tree state updates from other nodes
-    pub fn apply_remote_tree_operation(&self, model_id: &str, operation: &TreeOperation) {
-        let tree_key = Self::normalize_mesh_model_id(model_id);
+    /// Apply remote tree operation from mesh.
+    ///
+    /// `mesh_key` is the opaque key the operation was originally synced under;
+    /// `select_worker` / `select_worker_min_load` send tree operations to mesh
+    /// keyed by the composite `pool::model`, and any future receive path is
+    /// expected to forward that same string back here unchanged. The argument
+    /// is kept as `&str` so the mesh layer can stay key-agnostic.
+    ///
+    /// Note: `PolicyRegistry::apply_remote_tree_operation` (the only forwarder)
+    /// currently has no in-process callers; the receive path is not yet wired,
+    /// so this method is reachable only via tests today.
+    pub fn apply_remote_tree_operation(&self, mesh_key: &str, operation: &TreeOperation) {
+        let tree_key = Self::normalize_mesh_model_id(mesh_key);
 
         let tree = self
             .trees
@@ -256,15 +282,15 @@ impl CacheAwarePolicy {
             TreeOperation::Insert(insert_op) => {
                 tree.insert(&insert_op.text, &insert_op.tenant);
                 debug!(
-                    "Applied remote tree insert: model={}, text={}, tenant={}",
-                    model_id, insert_op.text, insert_op.tenant
+                    "Applied remote tree insert: key={}, text={}, tenant={}",
+                    mesh_key, insert_op.text, insert_op.tenant
                 );
             }
             TreeOperation::Remove(remove_op) => {
                 tree.remove_tenant(&remove_op.tenant);
                 debug!(
-                    "Applied remote tree remove: model={}, tenant={}",
-                    model_id, remove_op.tenant
+                    "Applied remote tree remove: key={}, tenant={}",
+                    mesh_key, remove_op.tenant
                 );
             }
         }
@@ -273,13 +299,10 @@ impl CacheAwarePolicy {
     /// Run cache eviction to prevent unbounded growth
     pub fn evict_cache(&self, max_size: usize) {
         for tree_ref in self.trees.iter() {
-            let model_id = tree_ref.key();
+            let tree_key = tree_ref.key();
             let tree = tree_ref.value();
             tree.evict_tenant_by_size(max_size);
-            debug!(
-                "Cache eviction for model {}, max_size: {}",
-                model_id, max_size
-            );
+            debug!("Cache eviction for {}, max_size: {}", tree_key, max_size);
         }
     }
 
@@ -288,7 +311,7 @@ impl CacheAwarePolicy {
         workers: &[Arc<dyn Worker>],
         request_text: &Option<&str>,
         healthy_indices: &[usize],
-        model_id: &str,
+        tree_key: &str,
         max_load: usize,
         min_load: usize,
     ) -> Option<usize> {
@@ -312,7 +335,7 @@ impl CacheAwarePolicy {
         if let Some(text) = request_text {
             // Get the tree reference without locking the entire HashMap
             // DashMap only locks the specific shard containing this key
-            let tree = self.trees.get(model_id).map(|entry| entry.value().clone());
+            let tree = self.trees.get(tree_key).map(|entry| entry.value().clone());
 
             if let Some(tree) = tree {
                 let worker_url = workers[min_load_idx].url();
@@ -326,15 +349,17 @@ impl CacheAwarePolicy {
                         text: text.to_string(),
                         tenant: worker_url.to_string(),
                     });
-                    let mesh_model_id = Self::normalize_mesh_model_id(model_id);
-                    if let Err(e) = mesh_sync.sync_tree_operation(mesh_model_id.to_string(), op) {
+                    let mesh_key = Self::normalize_mesh_model_id(tree_key);
+                    if let Err(e) = mesh_sync.sync_tree_operation(mesh_key.to_string(), op) {
                         warn!("Failed to sync tree insert operation to mesh: {}", e);
                     }
                 }
             } else {
-                debug!(
-                    "Warning: No tree found for model '{}', skipping cache update",
-                    model_id
+                warn!(
+                    "cache_aware: no tree found for key '{}', skipping cache update — \
+                     pool tree was not seeded (init_pd_cache_aware_policies missed or \
+                     a race during worker registration)",
+                    tree_key
                 );
             }
         }
@@ -360,9 +385,10 @@ impl LoadBalancingPolicy for CacheAwarePolicy {
             return None;
         }
 
-        // Determine the model for this set of workers (router pre-filters by model)
-        // All workers should be from the same model
-        let model_id = normalize_model_key(workers[healthy_indices[0]].model_id());
+        // Determine the (pool, model) key for this set of workers — the router pre-filters
+        // so every healthy worker here belongs to the same pool and same model.
+        let pivot = workers[healthy_indices[0]].as_ref();
+        let tree_key = tree_key_for_worker(pivot);
 
         // Get current load statistics - compute min/max in single pass without allocation
         let (min_load, max_load) = workers.iter().fold((usize::MAX, 0usize), |(min, max), w| {
@@ -380,7 +406,7 @@ impl LoadBalancingPolicy for CacheAwarePolicy {
                 workers,
                 &request_text,
                 &healthy_indices,
-                model_id,
+                &tree_key,
                 max_load,
                 min_load,
             );
@@ -391,7 +417,7 @@ impl LoadBalancingPolicy for CacheAwarePolicy {
 
         // Get the tree reference without locking the entire HashMap
         // DashMap only locks the specific shard containing this key
-        let tree = self.trees.get(model_id).map(|entry| entry.value().clone());
+        let tree = self.trees.get(&tree_key).map(|entry| entry.value().clone());
 
         if let Some(tree) = tree {
             // Now we work with the tree without holding the HashMap lock
@@ -430,8 +456,8 @@ impl LoadBalancingPolicy for CacheAwarePolicy {
                         text: text.to_string(),
                         tenant: workers[idx].url().to_string(),
                     });
-                    let mesh_model_id = Self::normalize_mesh_model_id(model_id);
-                    if let Err(e) = mesh_sync.sync_tree_operation(mesh_model_id.to_string(), op) {
+                    let mesh_key = Self::normalize_mesh_model_id(&tree_key);
+                    if let Err(e) = mesh_sync.sync_tree_operation(mesh_key.to_string(), op) {
                         warn!("Failed to sync tree insert operation to mesh: {}", e);
                     }
                 }
@@ -454,8 +480,8 @@ impl LoadBalancingPolicy for CacheAwarePolicy {
                     let op = TreeOperation::Remove(TreeRemoveOp {
                         tenant: tenant_url.to_string(),
                     });
-                    let mesh_model_id = Self::normalize_mesh_model_id(model_id);
-                    if let Err(e) = mesh_sync.sync_tree_operation(mesh_model_id.to_string(), op) {
+                    let mesh_key = Self::normalize_mesh_model_id(&tree_key);
+                    if let Err(e) = mesh_sync.sync_tree_operation(mesh_key.to_string(), op) {
                         warn!("Failed to sync tree remove operation to mesh: {}", e);
                     }
                 }
@@ -464,12 +490,14 @@ impl LoadBalancingPolicy for CacheAwarePolicy {
             // Fallback to first healthy worker
             healthy_indices.first().copied()
         } else {
-            // No tree for this model, log warning and use random selection
-            debug!(
-                "Warning: No tree found for model '{}', using random worker selection",
-                model_id
+            warn!(
+                "cache_aware: no tree found for key '{}', falling back to random \
+                 worker selection — pool tree was not seeded \
+                 (init_pd_cache_aware_policies missed or a race during worker \
+                 registration); cache affinity is effectively disabled until this \
+                 clears",
+                tree_key
             );
-            // Return a random healthy worker
             let mut rng = rand::rng();
             let random_idx = rng.random_range(0..healthy_indices.len());
             Some(healthy_indices[random_idx])
@@ -711,8 +739,11 @@ mod tests {
             .await
             .unwrap();
 
-        // Verify tree operation was synced to mesh (under UNKNOWN_MODEL_ID since no model was specified)
-        let tree_state = mesh_sync.get_tree_state(UNKNOWN_MODEL_ID);
+        // Verify tree operation was synced to mesh under the composite `pool::model`
+        // key — workers here are Regular and no model was specified, so the key is
+        // `regular::UNKNOWN_MODEL_ID`.
+        let expected_key = format!("regular::{}", UNKNOWN_MODEL_ID);
+        let tree_state = mesh_sync.get_tree_state(&expected_key);
         assert!(tree_state.is_some());
         let tree = tree_state.unwrap();
         assert!(!tree.operations.is_empty());
@@ -884,4 +915,597 @@ mod tests {
             .unwrap();
         assert_eq!(idx, 0);
     }
+
+    fn make_prefill(url: &str) -> Arc<dyn Worker> {
+        Arc::new(
+            BasicWorkerBuilder::new(url)
+                .worker_type(WorkerType::Prefill {
+                    bootstrap_port: Some(9000),
+                })
+                .build(),
+        )
+    }
+
+    fn make_decode(url: &str) -> Arc<dyn Worker> {
+        Arc::new(
+            BasicWorkerBuilder::new(url)
+                .worker_type(WorkerType::Decode)
+                .build(),
+        )
+    }
+
+    /// PD setup with two separate `CacheAwarePolicy` instances — the production
+    /// wiring. Each pool's tree is seeded only with its own workers. Across a
+    /// 4-turn growing prompt, each pool must stick to one worker.
+    #[tokio::test]
+    async fn test_pd_pool_isolation_two_policies() {
+        let config = CacheAwareConfig {
+            eviction_interval_secs: 0,
+            ..Default::default()
+        };
+        let prefill_policy = CacheAwarePolicy::with_config(config.clone());
+        let decode_policy = CacheAwarePolicy::with_config(config);
+
+        let prefill_workers: Vec<Arc<dyn Worker>> = vec![
+            make_prefill("http://prefill0:8000"),
+            make_prefill("http://prefill1:8000"),
+        ];
+        let decode_workers: Vec<Arc<dyn Worker>> = vec![
+            make_decode("http://decode0:8000"),
+            make_decode("http://decode1:8000"),
+        ];
+        prefill_policy.init_workers(&prefill_workers);
+        decode_policy.init_workers(&decode_workers);
+
+        let turns = [
+            "turn1",
+            "turn1 turn2",
+            "turn1 turn2 turn3",
+            "turn1 turn2 turn3 turn4",
+        ];
+
+        let mut prefill_idx: Option<usize> = None;
+        let mut decode_idx: Option<usize> = None;
+        for prompt in turns {
+            let info = SelectWorkerInfo {
+                request_text: Some(prompt),
+                ..Default::default()
+            };
+            let p = prefill_policy
+                .select_worker(&prefill_workers, &info)
+                .await
+                .expect("prefill pool returns a worker");
+            let d = decode_policy
+                .select_worker(&decode_workers, &info)
+                .await
+                .expect("decode pool returns a worker");
+            match prefill_idx {
+                None => prefill_idx = Some(p),
+                Some(pinned) => assert_eq!(
+                    p, pinned,
+                    "prefill should stay pinned across turns (prompt={prompt:?})"
+                ),
+            }
+            match decode_idx {
+                None => decode_idx = Some(d),
+                Some(pinned) => assert_eq!(
+                    d, pinned,
+                    "decode should stay pinned across turns (prompt={prompt:?})"
+                ),
+            }
+        }
+    }
+
+    /// Regression: even if a single `CacheAwarePolicy` instance is incorrectly
+    /// wired to both pools, pool-aware tree keys must keep their state disjoint.
+    /// The pre-fix code shared one trie keyed by `model_id`, so alternating
+    /// prefill/decode `tree.insert` calls overwrote each other and the policy
+    /// degenerated into worker-flipping random selection.
+    #[tokio::test]
+    async fn test_pd_pool_isolation_shared_policy_regression() {
+        let config = CacheAwareConfig {
+            eviction_interval_secs: 0,
+            ..Default::default()
+        };
+        let policy = CacheAwarePolicy::with_config(config);
+
+        let prefill_workers: Vec<Arc<dyn Worker>> = vec![
+            make_prefill("http://prefill0:8000"),
+            make_prefill("http://prefill1:8000"),
+        ];
+        let decode_workers: Vec<Arc<dyn Worker>> = vec![
+            make_decode("http://decode0:8000"),
+            make_decode("http://decode1:8000"),
+        ];
+
+        // One instance, mixed init — pool-aware keys split the trees internally.
+        let mut combined: Vec<Arc<dyn Worker>> = Vec::new();
+        combined.extend(prefill_workers.iter().cloned());
+        combined.extend(decode_workers.iter().cloned());
+        policy.init_workers(&combined);
+
+        let turns = [
+            "turn1",
+            "turn1 turn2",
+            "turn1 turn2 turn3",
+            "turn1 turn2 turn3 turn4",
+        ];
+
+        let mut prefill_idx: Option<usize> = None;
+        let mut decode_idx: Option<usize> = None;
+        for prompt in turns {
+            let info = SelectWorkerInfo {
+                request_text: Some(prompt),
+                ..Default::default()
+            };
+            let p = policy
+                .select_worker(&prefill_workers, &info)
+                .await
+                .expect("prefill pool returns a worker");
+            let d = policy
+                .select_worker(&decode_workers, &info)
+                .await
+                .expect("decode pool returns a worker");
+
+            assert!(
+                prefill_workers[p].url().starts_with("http://prefill"),
+                "prefill call must return a prefill index, got {} (prompt={prompt:?})",
+                prefill_workers[p].url()
+            );
+            assert!(
+                decode_workers[d].url().starts_with("http://decode"),
+                "decode call must return a decode index, got {} (prompt={prompt:?})",
+                decode_workers[d].url()
+            );
+
+            match prefill_idx {
+                None => prefill_idx = Some(p),
+                Some(pinned) => assert_eq!(
+                    p, pinned,
+                    "prefill should stay pinned across turns (prompt={prompt:?})"
+                ),
+            }
+            match decode_idx {
+                None => decode_idx = Some(d),
+                Some(pinned) => assert_eq!(
+                    d, pinned,
+                    "decode should stay pinned across turns (prompt={prompt:?})"
+                ),
+            }
+        }
+    }
+
+    /// Removing a PD worker via the composite-key `remove_worker(&dyn Worker)` path
+    /// must drop it from its own pool's tree without touching the other pool. This
+    /// covers `PolicyRegistry::remove_pd_worker_from_cache_aware`, which routes the
+    /// removal here based on `worker.worker_type()`.
+    #[tokio::test]
+    async fn test_pd_pool_isolation_remove_worker() {
+        let config = CacheAwareConfig {
+            eviction_interval_secs: 0,
+            ..Default::default()
+        };
+        let policy = CacheAwarePolicy::with_config(config);
+
+        let prefill0 = make_prefill("http://prefill0:8000");
+        let prefill1 = make_prefill("http://prefill1:8000");
+        let decode0 = make_decode("http://decode0:8000");
+        let decode1 = make_decode("http://decode1:8000");
+
+        let prefill_workers: Vec<Arc<dyn Worker>> = vec![prefill0.clone(), prefill1.clone()];
+        let decode_workers: Vec<Arc<dyn Worker>> = vec![decode0.clone(), decode1.clone()];
+        let combined: Vec<Arc<dyn Worker>> = vec![
+            prefill0.clone(),
+            prefill1.clone(),
+            decode0.clone(),
+            decode1.clone(),
+        ];
+        policy.init_workers(&combined);
+
+        // Seed both trees with affinity for one prompt.
+        let prompt = "shared prefix to seed the cache_aware trees";
+        let info = SelectWorkerInfo {
+            request_text: Some(prompt),
+            ..Default::default()
+        };
+        policy
+            .select_worker(&prefill_workers, &info)
+            .await
+            .expect("seed prefill");
+        policy
+            .select_worker(&decode_workers, &info)
+            .await
+            .expect("seed decode");
+
+        let prefill_key = format!("prefill::{}", UNKNOWN_MODEL_ID);
+        let decode_key = format!("decode::{}", UNKNOWN_MODEL_ID);
+
+        let prefill_before = policy
+            .trees
+            .get(&prefill_key)
+            .map(|t| t.value().prefix_match_with_counts(prompt).tenant)
+            .expect("prefill tree seeded");
+        let decode_before = policy
+            .trees
+            .get(&decode_key)
+            .map(|t| t.value().prefix_match_with_counts(prompt).tenant)
+            .expect("decode tree seeded");
+        assert!(
+            prefill_before.starts_with("http://prefill"),
+            "prefill tree should hold a prefill tenant before removal, got {prefill_before}"
+        );
+        assert!(
+            decode_before.starts_with("http://decode"),
+            "decode tree should hold a decode tenant before removal, got {decode_before}"
+        );
+
+        // Drop prefill0 via the composite-key removal path.
+        policy.remove_worker(prefill0.as_ref());
+
+        // The prefill tree must no longer point at prefill0 for the seeded prompt.
+        let prefill_after = policy
+            .trees
+            .get(&prefill_key)
+            .map(|t| t.value().prefix_match_with_counts(prompt).tenant)
+            .expect("prefill tree still exists");
+        assert_ne!(
+            &*prefill_after,
+            prefill0.url(),
+            "prefill0 should be gone from the prefill tree"
+        );
+
+        // The decode tree must be byte-for-byte unchanged.
+        let decode_after = policy
+            .trees
+            .get(&decode_key)
+            .map(|t| t.value().prefix_match_with_counts(prompt).tenant)
+            .expect("decode tree still exists");
+        assert_eq!(
+            decode_after, decode_before,
+            "removing a prefill worker must not touch the decode tree"
+        );
+    }
+
+    /// Shared setup for `PolicyRegistry::remove_pd_worker_from_cache_aware` tests:
+    /// build a registry whose prefill and decode policies are separate
+    /// `CacheAwarePolicy` instances seeded with the matching pool's workers, then
+    /// return the registry, the per-pool policy handles (for tree inspection), and
+    /// representative workers from each pool.
+    #[allow(clippy::type_complexity)]
+    fn pd_registry_with_cache_aware_pools() -> (
+        Arc<crate::policies::PolicyRegistry>,
+        Arc<CacheAwarePolicy>,
+        Arc<CacheAwarePolicy>,
+        Arc<dyn Worker>,
+        Arc<dyn Worker>,
+        Arc<dyn Worker>,
+        Arc<dyn Worker>,
+    ) {
+        let registry = Arc::new(crate::policies::PolicyRegistry::new(
+            crate::config::types::PolicyConfig::RoundRobin,
+        ));
+        let no_eviction = CacheAwareConfig {
+            eviction_interval_secs: 0,
+            ..Default::default()
+        };
+        let prefill_ca = Arc::new(CacheAwarePolicy::with_config(no_eviction.clone()));
+        let decode_ca = Arc::new(CacheAwarePolicy::with_config(no_eviction));
+        registry.set_prefill_policy(prefill_ca.clone() as Arc<dyn LoadBalancingPolicy>);
+        registry.set_decode_policy(decode_ca.clone() as Arc<dyn LoadBalancingPolicy>);
+
+        let prefill0 = make_prefill("http://prefill0:8000");
+        let prefill1 = make_prefill("http://prefill1:8000");
+        let decode0 = make_decode("http://decode0:8000");
+        let decode1 = make_decode("http://decode1:8000");
+
+        let prefill_workers: Vec<Arc<dyn Worker>> = vec![prefill0.clone(), prefill1.clone()];
+        let decode_workers: Vec<Arc<dyn Worker>> = vec![decode0.clone(), decode1.clone()];
+        registry.init_pd_cache_aware_policies(&prefill_workers, &decode_workers);
+
+        (
+            registry, prefill_ca, decode_ca, prefill0, prefill1, decode0, decode1,
+        )
+    }
+
+    /// Seed both pool trees so each has a known tenant for `prompt`, then return
+    /// the (prefill_tenant, decode_tenant) snapshot to compare against after a
+    /// dispatched removal.
+    async fn seed_pd_pools(
+        prefill_ca: &CacheAwarePolicy,
+        decode_ca: &CacheAwarePolicy,
+        prefill_workers: &[Arc<dyn Worker>],
+        decode_workers: &[Arc<dyn Worker>],
+        prompt: &str,
+    ) -> (Arc<str>, Arc<str>) {
+        let info = SelectWorkerInfo {
+            request_text: Some(prompt),
+            ..Default::default()
+        };
+        prefill_ca
+            .select_worker(prefill_workers, &info)
+            .await
+            .expect("prefill seed");
+        decode_ca
+            .select_worker(decode_workers, &info)
+            .await
+            .expect("decode seed");
+
+        let prefill_key = format!("prefill::{}", UNKNOWN_MODEL_ID);
+        let decode_key = format!("decode::{}", UNKNOWN_MODEL_ID);
+        let prefill_tenant = prefill_ca
+            .trees
+            .get(&prefill_key)
+            .map(|t| t.value().prefix_match_with_counts(prompt).tenant)
+            .expect("prefill tree seeded");
+        let decode_tenant = decode_ca
+            .trees
+            .get(&decode_key)
+            .map(|t| t.value().prefix_match_with_counts(prompt).tenant)
+            .expect("decode tree seeded");
+        (prefill_tenant, decode_tenant)
+    }
+
+    /// A `Prefill` worker passed to `remove_pd_worker_from_cache_aware` must hit
+    /// the registry's `prefill_policy` and leave `decode_policy` untouched.
+    /// Catches a dispatch swap like `Prefill => self.decode_policy.get()`.
+    #[tokio::test]
+    async fn test_registry_remove_pd_worker_prefill_dispatches_to_prefill_policy() {
+        let (registry, prefill_ca, decode_ca, prefill0, prefill1, decode0, decode1) =
+            pd_registry_with_cache_aware_pools();
+        let prefill_workers: Vec<Arc<dyn Worker>> = vec![prefill0.clone(), prefill1.clone()];
+        let decode_workers: Vec<Arc<dyn Worker>> = vec![decode0.clone(), decode1.clone()];
+
+        let prompt = "prefix used to seed both pool trees";
+        let (_prefill_before, decode_before) = seed_pd_pools(
+            &prefill_ca,
+            &decode_ca,
+            &prefill_workers,
+            &decode_workers,
+            prompt,
+        )
+        .await;
+
+        registry.remove_pd_worker_from_cache_aware(prefill0.as_ref());
+
+        let prefill_key = format!("prefill::{}", UNKNOWN_MODEL_ID);
+        let decode_key = format!("decode::{}", UNKNOWN_MODEL_ID);
+        let prefill_after = prefill_ca
+            .trees
+            .get(&prefill_key)
+            .map(|t| t.value().prefix_match_with_counts(prompt).tenant)
+            .expect("prefill tree still exists");
+        assert_ne!(
+            &*prefill_after,
+            prefill0.url(),
+            "registry dispatch must drop prefill0 from the prefill pool's tree"
+        );
+        let decode_after = decode_ca
+            .trees
+            .get(&decode_key)
+            .map(|t| t.value().prefix_match_with_counts(prompt).tenant)
+            .expect("decode tree still exists");
+        assert_eq!(
+            decode_after, decode_before,
+            "removing a prefill worker must not touch the decode pool's tree"
+        );
+    }
+
+    /// Mirror of the prefill dispatch test for `Decode`. Catches a dispatch swap
+    /// in the other direction (`Decode => self.prefill_policy.get()`).
+    #[tokio::test]
+    async fn test_registry_remove_pd_worker_decode_dispatches_to_decode_policy() {
+        let (registry, prefill_ca, decode_ca, prefill0, prefill1, decode0, decode1) =
+            pd_registry_with_cache_aware_pools();
+        let prefill_workers: Vec<Arc<dyn Worker>> = vec![prefill0.clone(), prefill1.clone()];
+        let decode_workers: Vec<Arc<dyn Worker>> = vec![decode0.clone(), decode1.clone()];
+
+        let prompt = "prefix used to seed both pool trees";
+        let (prefill_before, _decode_before) = seed_pd_pools(
+            &prefill_ca,
+            &decode_ca,
+            &prefill_workers,
+            &decode_workers,
+            prompt,
+        )
+        .await;
+
+        registry.remove_pd_worker_from_cache_aware(decode0.as_ref());
+
+        let prefill_key = format!("prefill::{}", UNKNOWN_MODEL_ID);
+        let decode_key = format!("decode::{}", UNKNOWN_MODEL_ID);
+        let decode_after = decode_ca
+            .trees
+            .get(&decode_key)
+            .map(|t| t.value().prefix_match_with_counts(prompt).tenant)
+            .expect("decode tree still exists");
+        assert_ne!(
+            &*decode_after,
+            decode0.url(),
+            "registry dispatch must drop decode0 from the decode pool's tree"
+        );
+        let prefill_after = prefill_ca
+            .trees
+            .get(&prefill_key)
+            .map(|t| t.value().prefix_match_with_counts(prompt).tenant)
+            .expect("prefill tree still exists");
+        assert_eq!(
+            prefill_after, prefill_before,
+            "removing a decode worker must not touch the prefill pool's tree"
+        );
+    }
+
+    /// `remove_pd_worker_from_cache_aware` must short-circuit on `Regular`
+    /// workers and silently ignore non-cache_aware policies (`name() != "cache_aware"`).
+    /// Both branches are no-ops: neither pool tree changes, and no downcast panic.
+    #[tokio::test]
+    async fn test_registry_remove_pd_worker_regular_and_non_cache_aware_noop() {
+        // (a) Regular worker: should early-return regardless of policy state.
+        let (registry, prefill_ca, decode_ca, prefill0, prefill1, decode0, decode1) =
+            pd_registry_with_cache_aware_pools();
+        let prefill_workers: Vec<Arc<dyn Worker>> = vec![prefill0.clone(), prefill1.clone()];
+        let decode_workers: Vec<Arc<dyn Worker>> = vec![decode0.clone(), decode1.clone()];
+
+        let prompt = "regular-noop seed prompt";
+        let (prefill_before, decode_before) = seed_pd_pools(
+            &prefill_ca,
+            &decode_ca,
+            &prefill_workers,
+            &decode_workers,
+            prompt,
+        )
+        .await;
+
+        let regular: Arc<dyn Worker> = Arc::new(
+            BasicWorkerBuilder::new("http://regular0:8000")
+                .worker_type(WorkerType::Regular)
+                .build(),
+        );
+        registry.remove_pd_worker_from_cache_aware(regular.as_ref());
+
+        let prefill_key = format!("prefill::{}", UNKNOWN_MODEL_ID);
+        let decode_key = format!("decode::{}", UNKNOWN_MODEL_ID);
+        let prefill_after = prefill_ca
+            .trees
+            .get(&prefill_key)
+            .map(|t| t.value().prefix_match_with_counts(prompt).tenant)
+            .expect("prefill tree still exists");
+        let decode_after = decode_ca
+            .trees
+            .get(&decode_key)
+            .map(|t| t.value().prefix_match_with_counts(prompt).tenant)
+            .expect("decode tree still exists");
+        assert_eq!(
+            prefill_after, prefill_before,
+            "Regular worker dispatch must not touch the prefill tree"
+        );
+        assert_eq!(
+            decode_after, decode_before,
+            "Regular worker dispatch must not touch the decode tree"
+        );
+
+        // (b) Non-cache_aware policy: PD pool is round_robin. The downcast must
+        // be skipped (no panic) and the call must be a no-op.
+        let registry =
+            crate::policies::PolicyRegistry::new(crate::config::types::PolicyConfig::RoundRobin);
+        let rr_prefill: Arc<dyn LoadBalancingPolicy> =
+            Arc::new(crate::policies::RoundRobinPolicy::new());
+        let rr_decode: Arc<dyn LoadBalancingPolicy> =
+            Arc::new(crate::policies::RoundRobinPolicy::new());
+        registry.set_prefill_policy(rr_prefill);
+        registry.set_decode_policy(rr_decode);
+        // No panic, no downcast — this would fault if the guard
+        // `policy.name() == "cache_aware"` were dropped.
+        registry.remove_pd_worker_from_cache_aware(prefill0.as_ref());
+        registry.remove_pd_worker_from_cache_aware(decode0.as_ref());
+    }
+
+    /// `init_pd_cache_aware_policies` must seed only the pool whose policy is
+    /// cache_aware AND whose worker list is non-empty. Covers all four corners:
+    /// both seeded, only-prefill-cache_aware, empty-worker short-circuit, and the
+    /// non-cache_aware side staying a no-op.
+    #[tokio::test]
+    async fn test_registry_init_pd_cache_aware_policies_gating() {
+        let no_eviction = CacheAwareConfig {
+            eviction_interval_secs: 0,
+            ..Default::default()
+        };
+        let prefill_key = format!("prefill::{}", UNKNOWN_MODEL_ID);
+        let decode_key = format!("decode::{}", UNKNOWN_MODEL_ID);
+
+        let prefill0 = make_prefill("http://prefill0:8000");
+        let decode0 = make_decode("http://decode0:8000");
+        let prefill_workers: Vec<Arc<dyn Worker>> = vec![prefill0.clone()];
+        let decode_workers: Vec<Arc<dyn Worker>> = vec![decode0.clone()];
+
+        // (a) Both pools are cache_aware with workers → both trees seeded under
+        // the correct composite key.
+        {
+            let registry = crate::policies::PolicyRegistry::new(
+                crate::config::types::PolicyConfig::RoundRobin,
+            );
+            let prefill_ca = Arc::new(CacheAwarePolicy::with_config(no_eviction.clone()));
+            let decode_ca = Arc::new(CacheAwarePolicy::with_config(no_eviction.clone()));
+            registry.set_prefill_policy(prefill_ca.clone() as Arc<dyn LoadBalancingPolicy>);
+            registry.set_decode_policy(decode_ca.clone() as Arc<dyn LoadBalancingPolicy>);
+
+            registry.init_pd_cache_aware_policies(&prefill_workers, &decode_workers);
+
+            assert!(
+                prefill_ca.trees.contains_key(&prefill_key),
+                "prefill cache_aware policy must be seeded under '{prefill_key}'"
+            );
+            assert!(
+                decode_ca.trees.contains_key(&decode_key),
+                "decode cache_aware policy must be seeded under '{decode_key}'"
+            );
+            assert!(
+                !prefill_ca.trees.contains_key(&decode_key),
+                "prefill_workers must not seed the decode tree key"
+            );
+            assert!(
+                !decode_ca.trees.contains_key(&prefill_key),
+                "decode_workers must not seed the prefill tree key"
+            );
+        }
+
+        // (b) Only prefill is cache_aware (decode is round_robin) → prefill seeded,
+        // decode side skipped silently (no downcast, no panic).
+        {
+            let registry = crate::policies::PolicyRegistry::new(
+                crate::config::types::PolicyConfig::RoundRobin,
+            );
+            let prefill_ca = Arc::new(CacheAwarePolicy::with_config(no_eviction.clone()));
+            let decode_rr: Arc<dyn LoadBalancingPolicy> =
+                Arc::new(crate::policies::RoundRobinPolicy::new());
+            registry.set_prefill_policy(prefill_ca.clone() as Arc<dyn LoadBalancingPolicy>);
+            registry.set_decode_policy(decode_rr);
+
+            registry.init_pd_cache_aware_policies(&prefill_workers, &decode_workers);
+
+            assert!(
+                prefill_ca.trees.contains_key(&prefill_key),
+                "prefill cache_aware side must seed even when decode side is non-cache_aware"
+            );
+        }
+
+        // (c) Both cache_aware but prefill worker list is empty → prefill tree
+        // NOT seeded (the inner `!is_empty()` guard short-circuits); decode side
+        // is still seeded.
+        {
+            let registry = crate::policies::PolicyRegistry::new(
+                crate::config::types::PolicyConfig::RoundRobin,
+            );
+            let prefill_ca = Arc::new(CacheAwarePolicy::with_config(no_eviction.clone()));
+            let decode_ca = Arc::new(CacheAwarePolicy::with_config(no_eviction.clone()));
+            registry.set_prefill_policy(prefill_ca.clone() as Arc<dyn LoadBalancingPolicy>);
+            registry.set_decode_policy(decode_ca.clone() as Arc<dyn LoadBalancingPolicy>);
+
+            registry.init_pd_cache_aware_policies(&[], &decode_workers);
+
+            assert!(
+                prefill_ca.trees.is_empty(),
+                "empty prefill worker list must not seed the prefill tree"
+            );
+            assert!(
+                decode_ca.trees.contains_key(&decode_key),
+                "decode side must still seed when only the prefill list is empty"
+            );
+        }
+
+        // (d) Both worker lists empty → neither pool seeded (init is a full no-op).
+        {
+            let registry = crate::policies::PolicyRegistry::new(
+                crate::config::types::PolicyConfig::RoundRobin,
+            );
+            let prefill_ca = Arc::new(CacheAwarePolicy::with_config(no_eviction.clone()));
+            let decode_ca = Arc::new(CacheAwarePolicy::with_config(no_eviction));
+            registry.set_prefill_policy(prefill_ca.clone() as Arc<dyn LoadBalancingPolicy>);
+            registry.set_decode_policy(decode_ca.clone() as Arc<dyn LoadBalancingPolicy>);
+
+            registry.init_pd_cache_aware_policies(&[], &[]);
+
+            assert!(prefill_ca.trees.is_empty());
+            assert!(decode_ca.trees.is_empty());
+        }
+    }
 }

sgl-model-gateway/src/policies/registry.rs

@@ -369,6 +369,31 @@ impl PolicyRegistry {
         }
     }
 
+    /// Remove a PD worker from the matching pool's cache-aware policy (lock-free).
+    ///
+    /// `prefill_policy` and `decode_policy` are separate `CacheAwarePolicy` instances
+    /// from `model_policies`, so the regular `remove_worker_from_cache_aware` does not
+    /// touch them. This dispatches to the right pool based on `worker.worker_type()`.
+    pub fn remove_pd_worker_from_cache_aware(&self, worker: &dyn Worker) {
+        let policy = match worker.worker_type() {
+            crate::core::WorkerType::Prefill { .. } => self.prefill_policy.get(),
+            crate::core::WorkerType::Decode => self.decode_policy.get(),
+            crate::core::WorkerType::Regular => return,
+        };
+        if let Some(policy) = policy {
+            if policy.name() == "cache_aware" {
+                if let Some(cache_aware) = policy.as_any().downcast_ref::<CacheAwarePolicy>() {
+                    cache_aware.remove_worker(worker);
+                    debug!(
+                        "Removed PD worker {} ({}) from cache-aware policy",
+                        worker.url(),
+                        worker.worker_type()
+                    );
+                }
+            }
+        }
+    }
+
     /// Initialize bucket policies for PD mode - lock-free
     pub fn init_pd_bucket_policies(&self, prefill_workers: &[Arc<dyn Worker>]) {
         // Initialize prefill policy if it's bucket (lock-free via OnceLock::get)