feat(grpc): add FlushCache and profile RPCs with worker-abstracted admin ops

crates/grpc_client/proto/common.proto

@@ -56,3 +56,39 @@ message KvBlocksRemoved {
 }
 
 message KvCacheCleared {}
+
+// =====================
+// Admin Operations
+// =====================
+
+message FlushCacheRequest {
+  // Seconds to wait for the scheduler to go idle before flushing.
+  // 0 = flush immediately (fails if requests are in flight).
+  float timeout_s = 1;
+}
+
+message FlushCacheResponse {
+  bool success = 1;
+  string message = 2;
+}
+
+message StartProfileRequest {
+  // Directory where profiler traces are written (backend default if unset).
+  optional string output_dir = 1;
+  // Forward step at which profiling starts (immediately if unset).
+  optional int32 start_step = 2;
+  // Number of forward steps to profile (until StopProfile if unset).
+  optional int32 num_steps = 3;
+  // Profiler activities, e.g. "CPU", "GPU", "MEM". Backend default if empty.
+  repeated string activities = 4;
+  optional bool with_stack = 5;
+  optional bool record_shapes = 6;
+  bool profile_by_stage = 7;
+}
+
+message StopProfileRequest {}
+
+message ProfileResponse {
+  bool success = 1;
+  string message = 2;
+}

crates/grpc_client/proto/sglang_scheduler.proto

@@ -30,6 +30,15 @@ service SglangScheduler {
   // Get comprehensive load metrics
   rpc GetLoads(GetLoadsRequest) returns (GetLoadsResponse);
 
+  // Flush the KV cache on all scheduler processes
+  rpc FlushCache(smg.grpc.common.FlushCacheRequest) returns (smg.grpc.common.FlushCacheResponse);
+
+  // Start the profiler on all scheduler processes
+  rpc StartProfile(smg.grpc.common.StartProfileRequest) returns (smg.grpc.common.ProfileResponse);
+
+  // Stop the profiler and export traces
+  rpc StopProfile(smg.grpc.common.StopProfileRequest) returns (smg.grpc.common.ProfileResponse);
+
   // Get tokenizer artifacts for remote construction
   rpc GetTokenizer(smg.grpc.common.GetTokenizerRequest) returns (stream smg.grpc.common.GetTokenizerChunk);
 

crates/grpc_client/src/lib.rs

@@ -57,6 +57,100 @@ macro_rules! impl_get_tokenizer {
 }
 pub(crate) use impl_get_tokenizer;
 
+/// Extra local-deadline margin for `flush_cache` on top of the timeout
+/// forwarded to the backend. The servicer bounds its own scheduler
+/// round-trip at `max(30, timeout_s + 10)` seconds, so the margin must
+/// cover that budget plus transport overhead.
+pub const FLUSH_RPC_DEADLINE_MARGIN: std::time::Duration = std::time::Duration::from_secs(45);
+
+/// Local deadline for profile start/stop RPCs. Stopping a profile can take
+/// a long time while the backend serializes large traces.
+pub const PROFILE_RPC_DEADLINE: std::time::Duration = std::time::Duration::from_secs(630);
+
+/// Shared admin-op implementations (`flush_cache`, `start_profile`,
+/// `stop_profile`) for engine clients whose protos expose the common
+/// admin RPCs (request/response messages live in `common.proto`).
+///
+/// Every call enforces a local deadline so an unresponsive backend cannot
+/// hang the gateway, and injects trace context for distributed tracing.
+macro_rules! impl_admin_ops {
+    () => {
+        /// Flush the KV cache on the backend scheduler.
+        ///
+        /// `timeout_s` is forwarded to the backend: 0 = flush immediately
+        /// (fails if requests are in flight), >0 = wait up to that many
+        /// seconds for the scheduler to go idle first.
+        pub async fn flush_cache(
+            &self,
+            timeout_s: f32,
+        ) -> Result<$crate::common_proto::FlushCacheResponse, tonic::Status> {
+            tracing::debug!("Requesting cache flush (timeout_s={timeout_s})");
+            let mut request =
+                tonic::Request::new($crate::common_proto::FlushCacheRequest { timeout_s });
+            if let Err(e) = self.trace_injector.inject(request.metadata_mut()) {
+                tracing::warn!("Failed to inject trace context: {}", e);
+            }
+            let deadline = std::time::Duration::from_secs_f32(timeout_s.max(0.0))
+                + $crate::FLUSH_RPC_DEADLINE_MARGIN;
+            let mut client = self.client.clone();
+            let response = tokio::time::timeout(deadline, client.flush_cache(request))
+                .await
+                .map_err(|_| {
+                    tonic::Status::deadline_exceeded(format!(
+                        "FlushCache did not complete within {deadline:?}"
+                    ))
+                })??;
+            Ok(response.into_inner())
+        }
+
+        /// Start the profiler on the backend scheduler.
+        pub async fn start_profile(
+            &self,
+            req: $crate::common_proto::StartProfileRequest,
+        ) -> Result<$crate::common_proto::ProfileResponse, tonic::Status> {
+            tracing::debug!("Requesting profile start");
+            let mut request = tonic::Request::new(req);
+            if let Err(e) = self.trace_injector.inject(request.metadata_mut()) {
+                tracing::warn!("Failed to inject trace context: {}", e);
+            }
+            let mut client = self.client.clone();
+            let response =
+                tokio::time::timeout($crate::PROFILE_RPC_DEADLINE, client.start_profile(request))
+                    .await
+                    .map_err(|_| {
+                        tonic::Status::deadline_exceeded(format!(
+                            "StartProfile did not complete within {:?}",
+                            $crate::PROFILE_RPC_DEADLINE
+                        ))
+                    })??;
+            Ok(response.into_inner())
+        }
+
+        /// Stop the profiler on the backend scheduler and export traces.
+        pub async fn stop_profile(
+            &self,
+        ) -> Result<$crate::common_proto::ProfileResponse, tonic::Status> {
+            tracing::debug!("Requesting profile stop");
+            let mut request = tonic::Request::new($crate::common_proto::StopProfileRequest {});
+            if let Err(e) = self.trace_injector.inject(request.metadata_mut()) {
+                tracing::warn!("Failed to inject trace context: {}", e);
+            }
+            let mut client = self.client.clone();
+            let response =
+                tokio::time::timeout($crate::PROFILE_RPC_DEADLINE, client.stop_profile(request))
+                    .await
+                    .map_err(|_| {
+                        tonic::Status::deadline_exceeded(format!(
+                            "StopProfile did not complete within {:?}",
+                            $crate::PROFILE_RPC_DEADLINE
+                        ))
+                    })??;
+            Ok(response.into_inner())
+        }
+    };
+}
+pub(crate) use impl_admin_ops;
+
 /// Shared `subscribe_kv_events()` implementation for all engine clients.
 ///
 /// Each engine's generated proto client has a `subscribe_kv_events` RPC method

crates/grpc_client/src/sglang_scheduler.rs

@@ -141,6 +141,7 @@ impl SglangSchedulerClient {
 
     crate::impl_get_tokenizer!();
     crate::impl_subscribe_kv_events!();
+    crate::impl_admin_ops!();
 
     /// Build a single SGLang EmbedRequest
     #[expect(

crates/protocols/src/worker.rs

@@ -995,9 +995,64 @@ pub struct FlushCacheResult {
     pub failed: Vec<(String, String)>,
     pub total_workers: usize,
     pub http_workers: usize,
+    #[serde(default)]
+    pub grpc_workers: usize,
+    pub message: String,
+}
+
+/// Options for starting a profiling run on workers.
+///
+/// Mirrors the engines' native profile parameters: serialized verbatim as
+/// the JSON body for HTTP workers and mapped to the `StartProfile` RPC for
+/// gRPC workers. Unset fields fall back to backend defaults.
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+#[serde(default)]
+pub struct ProfileOptions {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub output_dir: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub start_step: Option<i32>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub num_steps: Option<i32>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub activities: Option<Vec<String>>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub with_stack: Option<bool>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub record_shapes: Option<bool>,
+    pub profile_by_stage: bool,
+}
+
+/// Result from profile start/stop operations across workers
+#[derive(Debug, Clone, Deserialize, Serialize)]
+pub struct ProfileResult {
+    pub successful: Vec<String>,
+    pub failed: Vec<(String, String)>,
+    pub total_workers: usize,
     pub message: String,
 }
 
+/// Request body for the gateway `/start_profile` route: profile options
+/// plus an optional worker URL to target a single worker (e.g. one
+/// PD-disaggregation role). All workers are profiled when `url` is unset.
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+#[serde(default)]
+pub struct StartProfileRequest {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub url: Option<String>,
+    #[serde(flatten)]
+    pub options: ProfileOptions,
+}
+
+/// Request body for the gateway `/stop_profile` route: optional worker URL
+/// to target a single worker.
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+#[serde(default)]
+pub struct StopProfileRequest {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub url: Option<String>,
+}
+
 /// Result from getting worker loads
 #[derive(Debug, Clone, Deserialize, Serialize)]
 pub struct WorkerLoadsResult {
@@ -1101,6 +1156,43 @@ impl IntoResponse for FlushCacheResult {
             "message": self.message,
             "workers_flushed": self.successful.len(),
             "total_http_workers": self.http_workers,
+            "total_grpc_workers": self.grpc_workers,
+            "total_workers": self.total_workers
+        });
+
+        if !self.failed.is_empty() {
+            body["successful"] = json!(self.successful);
+            body["failed"] = json!(self
+                .failed
+                .into_iter()
+                .map(|(url, err)| json!({"worker": url, "error": err}))
+                .collect::<Vec<_>>());
+        }
+
+        (status, Json(body)).into_response()
+    }
+}
+
+#[cfg(feature = "axum")]
+impl IntoResponse for ProfileResult {
+    fn into_response(self) -> Response {
+        let status = if self.total_workers == 0 {
+            StatusCode::NOT_FOUND
+        } else if self.failed.is_empty() {
+            StatusCode::OK
+        } else {
+            StatusCode::PARTIAL_CONTENT
+        };
+
+        let status_str = match status {
+            StatusCode::OK => "success",
+            StatusCode::PARTIAL_CONTENT => "partial_success",
+            _ => "error",
+        };
+        let mut body = json!({
+            "status": status_str,
+            "message": self.message,
+            "workers_profiled": self.successful.len(),
             "total_workers": self.total_workers
         });