From 3c116d5e5a3b77c0b79bb91a211a270e438230e0 Mon Sep 17 00:00:00 2001
From: kun-llfl <i@imux.top>
Date: Sat, 20 Dec 2025 08:25:46 +0800
Subject: [PATCH] [router] bugfix: cache_aware in grpc inbalance forward
 (#15473)

Signed-off-by: Kun(llfl) <i@imux.top>
---
 .../grpc/common/stages/request_execution.rs   | 26 ++++++++++++++++++-
 sgl-model-gateway/src/routers/grpc/context.rs | 15 ++++++++++-
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/sgl-model-gateway/src/routers/grpc/common/stages/request_execution.rs b/sgl-model-gateway/src/routers/grpc/common/stages/request_execution.rs
index ddb8f4efa..e976bd145 100644
--- a/sgl-model-gateway/src/routers/grpc/common/stages/request_execution.rs
+++ b/sgl-model-gateway/src/routers/grpc/common/stages/request_execution.rs
@@ -8,7 +8,7 @@ use super::PipelineStage;
 use crate::routers::{
     error,
     grpc::{
-        context::{ClientSelection, ExecutionResult, RequestContext},
+        context::{ClientSelection, ExecutionResult, LoadGuards, RequestContext, WorkerSelection},
         proto_wrapper::{ProtoGenerateRequest, ProtoStream},
     },
 };
@@ -56,6 +56,30 @@ impl PipelineStage for RequestExecutionStage {
             )
         })?;
 
+        // Create load guards for worker load tracking (increment load when created)
+        // They will be automatically dropped (and decrement load) when RequestContext is dropped
+        let workers = ctx.state.workers.as_ref().ok_or_else(|| {
+            error!(
+                function = "RequestExecutionStage::execute",
+                "Worker selection not completed"
+            );
+            error::internal_error(
+                "worker_selection_not_completed",
+                "Worker selection not completed",
+            )
+        })?;
+
+        let load_guards = match workers {
+            WorkerSelection::Single { worker } => {
+                LoadGuards::Single(crate::core::WorkerLoadGuardV2::new(worker.clone()))
+            }
+            WorkerSelection::Dual { prefill, decode } => LoadGuards::Dual {
+                prefill: crate::core::WorkerLoadGuardV2::new(prefill.clone()),
+                decode: crate::core::WorkerLoadGuardV2::new(decode.clone()),
+            },
+        };
+        ctx.state.load_guards = Some(load_guards);
+
         // Extract dispatch metadata for tracing span
         let request_id = ctx
             .state
diff --git a/sgl-model-gateway/src/routers/grpc/context.rs b/sgl-model-gateway/src/routers/grpc/context.rs
index f01eaa949..36f25673f 100644
--- a/sgl-model-gateway/src/routers/grpc/context.rs
+++ b/sgl-model-gateway/src/routers/grpc/context.rs
@@ -14,7 +14,7 @@ use super::{
     proto_wrapper::{ProtoGenerateComplete, ProtoGenerateRequest, ProtoStream},
 };
 use crate::{
-    core::Worker,
+    core::{Worker, WorkerLoadGuardV2},
     protocols::{
         chat::{ChatCompletionRequest, ChatCompletionResponse},
         generate::{GenerateRequest, GenerateResponse},
@@ -76,6 +76,9 @@ pub struct ProcessingState {
     // Stage 5: Dispatch metadata
     pub dispatch: Option<DispatchMetadata>,
 
+    // Load guard for worker load tracking (created at execution stage)
+    pub load_guards: Option<LoadGuards>,
+
     // Stage 6: Response processing state
     pub response: ResponseState,
 }
@@ -143,6 +146,16 @@ pub struct DispatchMetadata {
     pub is_streaming: bool,
 }
 
+/// Load guards for worker load tracking
+/// Automatically decrements load when dropped
+pub enum LoadGuards {
+    Single(WorkerLoadGuardV2),
+    Dual {
+        prefill: WorkerLoadGuardV2,
+        decode: WorkerLoadGuardV2,
+    },
+}
+
 /// Response processing state (Step 6)
 #[derive(Default)]
 pub struct ResponseState {