From ffa318ea6f1d0a0487bbef9703383965e7735769 Mon Sep 17 00:00:00 2001 From: Sumit Mishra Date: Thu, 17 Sep 2026 15:38:58 +0530 Subject: [PATCH 1/2] fix(frontend): prepend KServe batch dim for Triton metadata Triton ModelConfig dims omit the batch axis. When max_batch_size > 0, ModelMetadata must prepend -1 so KServe v2 clients accept batched inputs. Fixes #14969 Signed-off-by: Sumit Mishra --- lib/llm/src/grpc/service/kserve.rs | 14 ++++++++--- lib/llm/src/grpc/service/tensor.rs | 37 ++++++++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 3 deletions(-) diff --git a/lib/llm/src/grpc/service/kserve.rs b/lib/llm/src/grpc/service/kserve.rs index c39cd6ce6a09..3580afca0458 100644 --- a/lib/llm/src/grpc/service/kserve.rs +++ b/lib/llm/src/grpc/service/kserve.rs @@ -67,7 +67,9 @@ impl GrpcTuningConfig { use crate::grpc::service::dispatch_error_status; use crate::grpc::service::openai::completion_response_stream; -use crate::grpc::service::tensor::{ExtendedNvCreateTensorResponse, tensor_response_stream}; +use crate::grpc::service::tensor::{ + ExtendedNvCreateTensorResponse, kserve_metadata_shape, tensor_response_stream, +}; use std::convert::{TryFrom, TryInto}; use tonic::{Request, Response, Status, transport::Server}; @@ -686,7 +688,10 @@ impl GrpcInferenceService for KserveService { .and_then(|dt| dt.oip_name()) .unwrap_or("TYPE_INVALID") .to_string(), - shape: input.dims.clone(), + shape: kserve_metadata_shape( + &input.dims, + model_config.max_batch_size, + ), }) .collect(), outputs: model_config @@ -700,7 +705,10 @@ impl GrpcInferenceService for KserveService { .and_then(|dt| dt.oip_name()) .unwrap_or("TYPE_INVALID") .to_string(), - shape: output.dims.clone(), + shape: kserve_metadata_shape( + &output.dims, + model_config.max_batch_size, + ), }, ) .collect(), diff --git a/lib/llm/src/grpc/service/tensor.rs b/lib/llm/src/grpc/service/tensor.rs index a2951c2ffadf..f7cbf5398c8d 100644 --- a/lib/llm/src/grpc/service/tensor.rs +++ b/lib/llm/src/grpc/service/tensor.rs @@ -845,6 +845,20 @@ impl DataType { } } +/// KServe v2 `ModelMetadata` wire shape for a Triton tensor. +/// +/// Triton's `ModelConfig` stores per-tensor `dims` without the batch dimension. +/// When `max_batch_size > 0`, native Triton's KServe adapter prepends `-1` +/// (variable batch) so clients can send batched inputs. Copying `dims` +/// verbatim drops that axis and KServe v2 clients reject the tensor. +pub fn kserve_metadata_shape(dims: &[i64], max_batch_size: i32) -> Vec { + if max_batch_size > 0 { + std::iter::once(-1).chain(dims.iter().copied()).collect() + } else { + dims.to_vec() + } +} + impl std::fmt::Display for tensor::DataType { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match *self { @@ -1213,3 +1227,26 @@ mod oip_datatype_tests { } } } + +#[cfg(test)] +mod kserve_metadata_shape_tests { + use super::kserve_metadata_shape; + + /// Native Triton's KServe adapter prepends a variable batch dim (`-1`) + /// when `max_batch_size > 0`. Empty `dims` still get that leading axis. + #[test] + fn prepends_variable_batch_dim_when_batching_enabled() { + assert_eq!(kserve_metadata_shape(&[-1], 4), vec![-1, -1]); + assert_eq!(kserve_metadata_shape(&[768], 8), vec![-1, 768]); + assert_eq!(kserve_metadata_shape(&[], 1), vec![-1]); + } + + /// `max_batch_size == 0` means the model is not batched; wire shape is + /// `dims` unchanged, including the empty-dims case. + #[test] + fn leaves_dims_unchanged_when_batching_disabled() { + assert_eq!(kserve_metadata_shape(&[-1], 0), vec![-1]); + assert_eq!(kserve_metadata_shape(&[3, 224, 224], 0), vec![3, 224, 224]); + assert_eq!(kserve_metadata_shape(&[], 0), Vec::::new()); + } +} From 8bad4780ab472ef0f2e9a2b66d8cad0f591a6434 Mon Sep 17 00:00:00 2001 From: Sumit Mishra Date: Thu, 17 Sep 2026 17:01:22 +0530 Subject: [PATCH 2/2] style: drop restating kserve_metadata_shape test rustdocs [skip ci] Signed-off-by: Sumit Mishra --- lib/llm/src/grpc/service/tensor.rs | 4 ---- 1 file changed, 4 deletions(-) diff --git a/lib/llm/src/grpc/service/tensor.rs b/lib/llm/src/grpc/service/tensor.rs index f7cbf5398c8d..619b4be6e4c4 100644 --- a/lib/llm/src/grpc/service/tensor.rs +++ b/lib/llm/src/grpc/service/tensor.rs @@ -1232,8 +1232,6 @@ mod oip_datatype_tests { mod kserve_metadata_shape_tests { use super::kserve_metadata_shape; - /// Native Triton's KServe adapter prepends a variable batch dim (`-1`) - /// when `max_batch_size > 0`. Empty `dims` still get that leading axis. #[test] fn prepends_variable_batch_dim_when_batching_enabled() { assert_eq!(kserve_metadata_shape(&[-1], 4), vec![-1, -1]); @@ -1241,8 +1239,6 @@ mod kserve_metadata_shape_tests { assert_eq!(kserve_metadata_shape(&[], 1), vec![-1]); } - /// `max_batch_size == 0` means the model is not batched; wire shape is - /// `dims` unchanged, including the empty-dims case. #[test] fn leaves_dims_unchanged_when_batching_disabled() { assert_eq!(kserve_metadata_shape(&[-1], 0), vec![-1]);