From 9413f2c54e662d72c04469c7104666cb8c993fac Mon Sep 17 00:00:00 2001 From: key4ng Date: Thu, 13 Aug 2026 12:03:21 -0700 Subject: [PATCH] feat(grpc): engines declare how decode constraints interact with reasoning MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The router cannot observe whether an engine enforces a decode constraint (json_schema/regex/grammar) from the first output token or only after the reasoning block — yet response parsing depends on it: a grammar bound from token 0 makes reasoning in the completion impossible, while a reasoning- aware grammar guarantees a think-end token precedes the payload. Today the router infers this from chat-template heuristics. Make it a declared capability instead: servicers advertise constrained_decoding_mode ("from_first_token" | "after_reasoning") via GetServerInfo, discovery carries it through the label pipeline onto worker metadata, and dispatch resolves it per request onto DispatchMetadata from the generating (decode) worker. - sglang/tokenspeed servicers: declare via the existing server_args struct (no proto change); both enforce grammars from the first token. - vllm: new GetServerInfoResponse.constrained_decoding_mode field, derived from structured_outputs_config (reasoning parser configured => after_reasoning); guarded for older smg-grpc-proto packages. - gateway: extract the label for all backends, parse into ConstrainedDecodingMode, resolve per dispatch. The response pipeline's reasoning/tool decision sites consume this in a follow-up (after #2122): declared from_first_token skips reasoning parsing under a tool constraint outright; undeclared workers keep the heuristic-plus-recovery fallback. Signed-off-by: key4ng --- crates/grpc_client/proto/vllm_engine.proto | 9 ++++ .../smg_grpc_servicer/sglang/servicer.py | 5 ++ .../smg_grpc_servicer/tokenspeed/servicer.py | 7 ++- .../smg_grpc_servicer/vllm/servicer.py | 16 +++++- .../tests/test_constrained_decoding_mode.py | 52 +++++++++++++++++++ model_gateway/src/routers/grpc/client.rs | 41 ++++++++++++++- .../grpc/common/stages/dispatch_metadata.rs | 25 +++++---- model_gateway/src/routers/grpc/context.rs | 48 +++++++++++++++++ 8 files changed, 191 insertions(+), 12 deletions(-) create mode 100644 grpc_servicer/tests/test_constrained_decoding_mode.py diff --git a/crates/grpc_client/proto/vllm_engine.proto b/crates/grpc_client/proto/vllm_engine.proto index fd637f6355..e55a3b12c0 100644 --- a/crates/grpc_client/proto/vllm_engine.proto +++ b/crates/grpc_client/proto/vllm_engine.proto @@ -360,6 +360,15 @@ message GetServerInfoResponse { // the router can verify a shared /dev/shm before using the SHM tensor // transport under `auto`. Empty when it can't be determined. string shm_namespace_id = 10; + + // How the engine applies decode constraints (json_schema/regex/grammar): + // "from_first_token" - grammar enforced from the first output token; the + // completion cannot contain reasoning/think tokens. + // "after_reasoning" - grammar activates after the reasoning block; the + // payload is preceded by reasoning + think-end token. + // Empty when unknown. The router uses this to decide whether reasoning + // parsing applies to grammar-constrained completions. + string constrained_decoding_mode = 11; } // ===================== diff --git a/grpc_servicer/smg_grpc_servicer/sglang/servicer.py b/grpc_servicer/smg_grpc_servicer/sglang/servicer.py index 6e9320269f..1a184be059 100644 --- a/grpc_servicer/smg_grpc_servicer/sglang/servicer.py +++ b/grpc_servicer/smg_grpc_servicer/sglang/servicer.py @@ -568,6 +568,11 @@ def make_serializable(obj): return str(obj) serializable_args = make_serializable(server_args_dict) + # Decode constraints (json_schema/regex/ebnf) are enforced by the + # grammar backend from the first output token; `require_reasoning` + # does not delay grammar activation in gRPC mode. Declared so the + # router knows constrained completions cannot contain reasoning. + serializable_args["constrained_decoding_mode"] = "from_first_token" server_args_struct.update(serializable_args) # Convert scheduler_info to Struct diff --git a/grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py b/grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py index 18c3405b63..0119939459 100644 --- a/grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py +++ b/grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py @@ -535,7 +535,12 @@ async def GetServerInfo( else: server_args_dict = dict(getattr(self.server_args, "__dict__", {})) server_args_struct = Struct() - server_args_struct.update(_make_json_serializable(server_args_dict)) + server_args_dict = _make_json_serializable(server_args_dict) + # Decode constraints are enforced by the grammar backend from the first + # output token. Declared so the router knows constrained completions + # cannot contain reasoning. + server_args_dict["constrained_decoding_mode"] = "from_first_token" + server_args_struct.update(server_args_dict) scheduler_info_struct = Struct() scheduler_info_struct.update(_make_json_serializable(dict(self.scheduler_info))) diff --git a/grpc_servicer/smg_grpc_servicer/vllm/servicer.py b/grpc_servicer/smg_grpc_servicer/vllm/servicer.py index 1b14495181..6b3ba9cbe7 100755 --- a/grpc_servicer/smg_grpc_servicer/vllm/servicer.py +++ b/grpc_servicer/smg_grpc_servicer/vllm/servicer.py @@ -484,13 +484,27 @@ async def GetServerInfo( # the router derives the suffix from the rank it pins per request kv_engine_id = getattr(kv_transfer_config, "engine_id", "") or "" - return vllm_engine_pb2.GetServerInfoResponse( + response = vllm_engine_pb2.GetServerInfoResponse( kv_connector=kv_connector, kv_role=kv_role, kv_engine_id=kv_engine_id, data_parallel_size=parallel.data_parallel_size, shm_namespace_id=mm_shm.shm_namespace_id(), ) + # Declare how decode constraints interact with reasoning: with a + # reasoning parser configured, vLLM activates grammars only after the + # reasoning block; otherwise the grammar binds from the first output + # token and constrained completions cannot contain reasoning. Guarded + # so older smg-grpc-proto packages (without the field) keep working. + if "constrained_decoding_mode" in response.DESCRIPTOR.fields_by_name: + response.constrained_decoding_mode = self._constrained_decoding_mode() + return response + + def _constrained_decoding_mode(self) -> str: + """Derive the engine's constrained-decoding behavior from its config.""" + structured_outputs = getattr(self.engine.vllm_config, "structured_outputs_config", None) + reasoning_parser = getattr(structured_outputs, "reasoning_parser", "") or "" + return "after_reasoning" if reasoning_parser else "from_first_token" async def GetLoads( self, diff --git a/grpc_servicer/tests/test_constrained_decoding_mode.py b/grpc_servicer/tests/test_constrained_decoding_mode.py new file mode 100644 index 0000000000..a61c194980 --- /dev/null +++ b/grpc_servicer/tests/test_constrained_decoding_mode.py @@ -0,0 +1,52 @@ +"""Unit tests for the engine-declared constrained-decoding capability. + +Run with: pytest grpc_servicer/tests/test_constrained_decoding_mode.py +""" + +from types import SimpleNamespace + +import pytest + +pytest.importorskip("smg_grpc_proto") +from smg_grpc_proto import vllm_engine_pb2 # noqa: E402 + +_HAS_FIELD = ( + "constrained_decoding_mode" in vllm_engine_pb2.GetServerInfoResponse.DESCRIPTOR.fields_by_name +) + + +@pytest.mark.skipif(not _HAS_FIELD, reason="installed smg-grpc-proto predates the field") +class TestGetServerInfoResponseConstrainedDecodingMode: + def test_defaults_to_empty_when_undeclared(self): + info = vllm_engine_pb2.GetServerInfoResponse() + assert info.constrained_decoding_mode == "" + + def test_mode_roundtrips(self): + info = vllm_engine_pb2.GetServerInfoResponse(constrained_decoding_mode="after_reasoning") + parsed = vllm_engine_pb2.GetServerInfoResponse.FromString(info.SerializeToString()) + assert parsed.constrained_decoding_mode == "after_reasoning" + + +class TestVllmDerivation: + """`_constrained_decoding_mode` derives honestly from the engine config.""" + + @staticmethod + def _servicer_with(structured_outputs_config): + vllm_servicer = pytest.importorskip("smg_grpc_servicer.vllm.servicer") + servicer = vllm_servicer.VllmEngineServicer.__new__(vllm_servicer.VllmEngineServicer) + servicer.engine = SimpleNamespace( + vllm_config=SimpleNamespace(structured_outputs_config=structured_outputs_config) + ) + return servicer + + def test_from_first_token_without_reasoning_parser(self): + servicer = self._servicer_with(SimpleNamespace(reasoning_parser="")) + assert servicer._constrained_decoding_mode() == "from_first_token" + + def test_from_first_token_without_structured_outputs_config(self): + servicer = self._servicer_with(None) + assert servicer._constrained_decoding_mode() == "from_first_token" + + def test_after_reasoning_with_reasoning_parser(self): + servicer = self._servicer_with(SimpleNamespace(reasoning_parser="qwen3")) + assert servicer._constrained_decoding_mode() == "after_reasoning" diff --git a/model_gateway/src/routers/grpc/client.rs b/model_gateway/src/routers/grpc/client.rs index 2bd1221b57..a7b4614395 100644 --- a/model_gateway/src/routers/grpc/client.rs +++ b/model_gateway/src/routers/grpc/client.rs @@ -852,6 +852,7 @@ const SGLANG_GRPC_KEYS: &[&str] = &[ "is_embedding", "vocab_size", "weight_version", + "constrained_decoding_mode", ]; /// Keys worth extracting from TokenSpeed gRPC `server_args` (post-rename: bare @@ -870,6 +871,7 @@ const TOKENSPEED_GRPC_KEYS: &[&str] = &[ "is_embedding", "vocab_size", "weight_version", + "constrained_decoding_mode", ]; // --------------------------------------------------------------------------- @@ -942,7 +944,7 @@ fn pick_prost_fields(labels: &mut HashMap, s: &prost_types::Stru mod tests { use std::collections::BTreeMap; - use smg_grpc_client::{sglang_proto, tokenspeed_proto}; + use smg_grpc_client::{sglang_proto, tokenspeed_proto, vllm_proto}; use super::{trtllm_status_healthy, ModelInfo, ServerInfo}; @@ -1057,6 +1059,43 @@ mod tests { assert!(!labels.contains_key("api_key")); } + /// The engine-declared constrained-decoding behavior travels as a label + /// for every backend: via `server_args` for SGLang/TokenSpeed, via the + /// flat proto field for vLLM (absent when the engine doesn't declare it). + #[test] + fn server_info_to_labels_carries_constrained_decoding_mode() { + let info = ServerInfo::Sglang(Box::new(sglang_proto::GetServerInfoResponse { + server_args: Some(prost_types::Struct { + fields: BTreeMap::from([( + "constrained_decoding_mode".to_string(), + string_value("from_first_token"), + )]), + }), + ..Default::default() + })); + assert_eq!( + info.to_labels() + .get("constrained_decoding_mode") + .map(String::as_str), + Some("from_first_token") + ); + + let info = ServerInfo::Vllm(vllm_proto::GetServerInfoResponse { + constrained_decoding_mode: "after_reasoning".to_string(), + ..Default::default() + }); + assert_eq!( + info.to_labels() + .get("constrained_decoding_mode") + .map(String::as_str), + Some("after_reasoning") + ); + + // Engines that don't declare (older servicers) produce no label. + let info = ServerInfo::Vllm(vllm_proto::GetServerInfoResponse::default()); + assert!(!info.to_labels().contains_key("constrained_decoding_mode")); + } + /// `GetModelInfoResponse` is flat for every backend, so it serializes via /// `flat_labels`: empty strings and zero numbers are skipped, booleans are /// kept, arrays are JSON-encoded. diff --git a/model_gateway/src/routers/grpc/common/stages/dispatch_metadata.rs b/model_gateway/src/routers/grpc/common/stages/dispatch_metadata.rs index 73215c019d..800ba05319 100644 --- a/model_gateway/src/routers/grpc/common/stages/dispatch_metadata.rs +++ b/model_gateway/src/routers/grpc/common/stages/dispatch_metadata.rs @@ -9,7 +9,9 @@ use tracing::error; use super::PipelineStage; use crate::routers::{ error, - grpc::context::{DispatchMetadata, RequestContext, RequestType, WorkerSelection}, + grpc::context::{ + ConstrainedDecodingMode, DispatchMetadata, RequestContext, RequestType, WorkerSelection, + }, }; /// Dispatch metadata stage: Prepare metadata for dispatch @@ -43,17 +45,21 @@ impl PipelineStage for DispatchMetadataStage { RequestType::Messages(req) => req.model.clone(), }; - let weight_version = ctx - .state - .workers - .as_ref() - .map(|w| match w { - WorkerSelection::Single { worker } => worker, - WorkerSelection::Disaggregated { decode, .. } => decode, - }) + // For PD disaggregation the decode leg generates the tokens, so its + // labels describe the completion (weight version, decode behavior). + let generating_worker = ctx.state.workers.as_ref().map(|w| match w { + WorkerSelection::Single { worker } => worker, + WorkerSelection::Disaggregated { decode, .. } => decode, + }); + + let weight_version = generating_worker .and_then(|w| w.metadata().spec.labels.get("weight_version").cloned()) .unwrap_or_else(|| "default".to_string()); + let constrained_decoding_mode = generating_worker + .and_then(|w| w.metadata().spec.labels.get("constrained_decoding_mode")) + .and_then(|value| ConstrainedDecodingMode::from_label(value)); + let created = SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap_or_default() @@ -64,6 +70,7 @@ impl PipelineStage for DispatchMetadataStage { model, created, weight_version: Some(weight_version), + constrained_decoding_mode, }); Ok(None) diff --git a/model_gateway/src/routers/grpc/context.rs b/model_gateway/src/routers/grpc/context.rs index 614aaad846..7c23d21263 100644 --- a/model_gateway/src/routers/grpc/context.rs +++ b/model_gateway/src/routers/grpc/context.rs @@ -412,6 +412,30 @@ pub(crate) enum ClientSelection { }, } +/// Engine-declared behavior for grammar-constrained decoding, advertised via +/// `GetServerInfo` and carried as the `constrained_decoding_mode` worker label. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum ConstrainedDecodingMode { + /// The grammar binds from the first output token: a constrained + /// completion cannot contain reasoning. + FromFirstToken, + /// The grammar activates after the reasoning block: reasoning and a + /// think-end token precede the constrained payload. + AfterReasoning, +} + +impl ConstrainedDecodingMode { + /// Parse the worker label. Unknown values resolve to `None` so the + /// response pipeline falls back to its capability-agnostic behavior. + pub(crate) fn from_label(value: &str) -> Option { + match value { + "from_first_token" => Some(Self::FromFirstToken), + "after_reasoning" => Some(Self::AfterReasoning), + _ => None, + } + } +} + /// Dispatch metadata (Step 5) #[derive(Clone)] pub(crate) struct DispatchMetadata { @@ -419,6 +443,15 @@ pub(crate) struct DispatchMetadata { pub model: String, pub created: u64, pub weight_version: Option, + /// How the dispatched worker applies decode constraints, when declared. + /// `None` for workers that don't advertise the label (older servicers). + #[expect( + dead_code, + reason = "read by the reasoning/tool decision sites in the follow-up \ + that lands after the required-tool-choice fix merges; the \ + unfulfilled expectation then forces removal of this attribute" + )] + pub constrained_decoding_mode: Option, } /// Load guards for worker load tracking @@ -1045,4 +1078,19 @@ mod tests { assert_eq!(plan.request_type(), "generate"); assert_eq!(plan.mode_label(), "prefill_decode"); } + + #[test] + fn constrained_decoding_mode_parses_known_labels_only() { + assert_eq!( + ConstrainedDecodingMode::from_label("from_first_token"), + Some(ConstrainedDecodingMode::FromFirstToken) + ); + assert_eq!( + ConstrainedDecodingMode::from_label("after_reasoning"), + Some(ConstrainedDecodingMode::AfterReasoning) + ); + // Unknown/future values fall back to capability-agnostic behavior. + assert_eq!(ConstrainedDecodingMode::from_label("adaptive"), None); + assert_eq!(ConstrainedDecodingMode::from_label(""), None); + } }