From cbcc5aab52f0963dcb1450aeed508b31e668669d Mon Sep 17 00:00:00 2001
From: Yingge He <yinggeh@nvidia.com>
Date: Tue, 15 Sep 2026 02:18:04 -0700
Subject: [PATCH 1/4] feat(llm): carry KServe requested outputs on tensor
 requests

Forward ModelInferRequest.outputs (name and parameters) onto
NvCreateTensorRequest so workers can honor classification and
output selection.

Signed-off-by: Yingge He <yinggeh@nvidia.com>
---
 lib/llm/src/grpc/service/tensor.rs | 99 +++++++++++++++++++++++++++++-
 lib/llm/src/protocols/tensor.rs    | 13 ++++
 2 files changed, 111 insertions(+), 1 deletion(-)

diff --git a/lib/llm/src/grpc/service/tensor.rs b/lib/llm/src/grpc/service/tensor.rs
index f1d628fb0042..d488c1ea89c8 100644
--- a/lib/llm/src/grpc/service/tensor.rs
+++ b/lib/llm/src/grpc/service/tensor.rs
@@ -28,7 +28,7 @@ use crate::http::service::{
 
 use crate::protocols::tensor;
 use crate::protocols::tensor::{
-    NvCreateTensorRequest, NvCreateTensorResponse, Tensor, TensorMetadata,
+    NvCreateTensorRequest, NvCreateTensorResponse, RequestedOutput, Tensor, TensorMetadata,
 };
 
 use crate::grpc::service::kserve::inference;
@@ -303,10 +303,20 @@ impl TryFrom<inference::ModelInferRequest> for NvCreateTensorRequest {
             },
             model: request.model_name.clone(),
             tensors: Vec::new(),
+            outputs: Vec::with_capacity(request.outputs.len()),
             parameters,
             nvext: None,
         };
 
+        // Requested outputs select which model outputs come back and how; their
+        // parameters (e.g. `classification`) are interpreted by the worker.
+        for output in &request.outputs {
+            tensor_request.outputs.push(RequestedOutput {
+                name: output.name.clone(),
+                parameters: convert_kserve_to_dynamo_params(&output.parameters)?,
+            });
+        }
+
         // iterate through inputs
         for (idx, input) in request.inputs.into_iter().enumerate() {
             // Extract per-tensor parameters
@@ -853,6 +863,93 @@ impl FromStr for tensor::DataType {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use inference::model_infer_request::InferRequestedOutputTensor;
+    use std::collections::HashMap;
+
+    fn param(choice: Option<ParameterChoice>) -> inference::InferParameter {
+        inference::InferParameter {
+            parameter_choice: choice,
+        }
+    }
+
+    fn classification(k: i64) -> HashMap<String, inference::InferParameter> {
+        HashMap::from([(
+            "classification".to_string(),
+            param(Some(ParameterChoice::Int64Param(k))),
+        )])
+    }
+
+    fn requested_output(
+        name: &str,
+        parameters: HashMap<String, inference::InferParameter>,
+    ) -> InferRequestedOutputTensor {
+        InferRequestedOutputTensor {
+            name: name.to_string(),
+            parameters,
+        }
+    }
+
+    fn infer_request(outputs: Vec<InferRequestedOutputTensor>) -> inference::ModelInferRequest {
+        inference::ModelInferRequest {
+            model_name: "classifier".to_string(),
+            outputs,
+            ..Default::default()
+        }
+    }
+
+    #[test]
+    fn requested_outputs_keep_their_parameters() {
+        let converted = NvCreateTensorRequest::try_from(infer_request(vec![
+            requested_output("OUTPUT0", classification(3)),
+            requested_output("OUTPUT1", HashMap::new()),
+        ]))
+        .unwrap();
+
+        assert_eq!(
+            converted.outputs,
+            vec![
+                RequestedOutput {
+                    name: "OUTPUT0".to_string(),
+                    parameters: tensor::Parameters::from([(
+                        "classification".to_string(),
+                        tensor::ParameterValue::Int64(3),
+                    )]),
+                },
+                RequestedOutput {
+                    name: "OUTPUT1".to_string(),
+                    parameters: tensor::Parameters::new(),
+                },
+            ]
+        );
+        assert_eq!(
+            serde_json::to_value(&converted.outputs[0]).unwrap(),
+            serde_json::json!({"name": "OUTPUT0", "parameters": {"classification": {"int64": 3}}})
+        );
+    }
+
+    #[test]
+    fn requested_outputs_empty_are_omitted_from_json() {
+        let converted = NvCreateTensorRequest::try_from(infer_request(vec![])).unwrap();
+        assert!(converted.outputs.is_empty());
+        assert!(
+            serde_json::to_value(&converted)
+                .unwrap()
+                .get("outputs")
+                .is_none()
+        );
+    }
+
+    #[test]
+    fn requested_outputs_valueless_parameter_is_rejected() {
+        let error = NvCreateTensorRequest::try_from(infer_request(vec![requested_output(
+            "OUTPUT0",
+            HashMap::from([("classification".to_string(), param(None))]),
+        )]))
+        .unwrap_err();
+
+        assert_eq!(error.code(), tonic::Code::InvalidArgument);
+        assert!(error.message().contains("classification"));
+    }
 
     // Golden bytes hardcoded (not derived from to_le_bytes) so a symmetric
     // break in both encode/decode still fires this test.
diff --git a/lib/llm/src/protocols/tensor.rs b/lib/llm/src/protocols/tensor.rs
index a3241973a623..38dd61c3fabd 100644
--- a/lib/llm/src/protocols/tensor.rs
+++ b/lib/llm/src/protocols/tensor.rs
@@ -154,6 +154,15 @@ pub struct Tensor {
     pub data: FlattenTensor,
 }
 
+#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
+#[serde(deny_unknown_fields)]
+pub struct RequestedOutput {
+    pub name: String,
+
+    #[serde(skip_serializing_if = "HashMap::is_empty", default)]
+    pub parameters: Parameters,
+}
+
 impl validator::Validate for Tensor {
     fn validate(&self) -> Result<(), validator::ValidationErrors> {
         use validator::{ValidationError, ValidationErrors};
@@ -207,6 +216,10 @@ pub struct NvCreateTensorRequest {
     #[validate(nested)]
     pub tensors: Vec<Tensor>,
 
+    /// Outputs the client asked for. An empty list asks for every model output.
+    #[serde(skip_serializing_if = "Vec::is_empty", default)]
+    pub outputs: Vec<RequestedOutput>,
+
     /// Optional request-level parameters
     #[serde(skip_serializing_if = "HashMap::is_empty", default)]
     pub parameters: Parameters,

From 507fb6c12a6425e1b4eef21a28bb2f27db08a838 Mon Sep 17 00:00:00 2001
From: Yingge He <yinggeh@nvidia.com>
Date: Thu, 10 Sep 2026 08:20:45 -0700
Subject: [PATCH 2/4] feat(triton): honor KServe classification on tensor
 outputs

Interpret requested-output classification as Triton's top-K class
strings, and return only the outputs the client asked for.

Signed-off-by: Yingge He <yinggeh@nvidia.com>
---
 .../src/dynamo/triton/classification.py       | 128 ++++++
 components/src/dynamo/triton/handlers.py      |  53 ++-
 .../tests/test_triton_classification.py       | 425 ++++++++++++++++++
 .../triton/tests/test_triton_handlers.py      |  15 +-
 .../triton/tests/test_triton_health_check.py  |   1 +
 .../backends/triton/overview.md               |   2 +-
 6 files changed, 615 insertions(+), 9 deletions(-)
 create mode 100644 components/src/dynamo/triton/classification.py
 create mode 100644 components/src/dynamo/triton/tests/test_triton_classification.py

diff --git a/components/src/dynamo/triton/classification.py b/components/src/dynamo/triton/classification.py
new file mode 100644
index 000000000000..c46e3fdfd57d
--- /dev/null
+++ b/components/src/dynamo/triton/classification.py
@@ -0,0 +1,128 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""KServe classification (top-K) for Triton model outputs.
+
+A requested output carrying the KServe ``classification`` parameter asks for the
+output's top-K classes as ``"<score>:<index>[:<label>]"`` strings instead of raw
+values. Triton implements this in its own HTTP / gRPC frontends rather than in
+the core, so the worker reproduces it here: the same descending-score ordering,
+score formatting, and BYTES output shape.
+"""
+
+from collections.abc import Mapping
+from typing import Any, Optional
+
+import numpy as np
+
+CLASSIFICATION_PARAMETER = "classification"
+
+# Bound on the elements sorted per batch entry, matching Triton's own limit. A
+# larger output would make the full sort below pathologically expensive.
+MAX_CLASSIFICATION_ELEMENTS = 1_000_000
+
+
+def class_count(parameters: Mapping[str, Any]) -> Optional[int]:
+    """Number of classes a requested output asked for.
+
+    Returns None when the output carries no ``classification`` parameter, which
+    means it asked for raw values.
+    """
+    value = parameters.get(CLASSIFICATION_PARAMETER)
+    if value is None:
+        return None
+
+    # Dynamo serializes a tensor parameter as a single-entry {tag: value} map.
+    # For example, {"classification": {"int64": 10}} becomes 10.
+    if isinstance(value, Mapping):
+        value = value.get("int64")
+    if isinstance(value, bool) or not isinstance(value, int):
+        raise ValueError(
+            f"invalid value type for '{CLASSIFICATION_PARAMETER}' parameter, "
+            "expected int64_param"
+        )
+    if value < 1:
+        raise ValueError(
+            f"invalid value for '{CLASSIFICATION_PARAMETER}' parameter, expected >= 1"
+        )
+    return value
+
+
+def classification_label(
+    response: Any, output_name: str, class_index: int
+) -> Optional[str]:
+    """Label for ``output_name`` at ``class_index``, or None if unlabelled.
+
+    Labels sit on the C ``TRITONSERVER_InferenceResponse``, exposed as
+    each tensor's ``memory_buffer.owner``. That binding takes an output
+    index, in the order ``outputs`` was filled — the same
+    ``(response, output_idx, class_index)`` lookup ``TopkClassifications``
+    makes. The high-level Python ``InferenceResponse`` does not wrap this
+    yet; if the owner has no lookup, class strings stay unlabeled.
+    """
+    try:
+        output_idx = list(response.outputs).index(output_name)
+        owner = response.outputs[output_name].memory_buffer.owner
+    except (AttributeError, KeyError, ValueError):
+        return None
+    lookup = getattr(owner, "output_classification_label", None)
+    if lookup is None:
+        return None
+    label = lookup(output_idx, class_index)
+    return label or None
+
+
+def top_k_classifications(
+    array: np.ndarray,
+    req_class_count: int,
+    triton_dtype: str,
+    batched: bool,
+    response: Any = None,
+    output_name: str = "",
+) -> np.ndarray:
+    """Top-K classes of one output as an object array of BYTES class strings.
+
+    ``batched`` marks the leading dimension of ``array`` as a batch dimension,
+    in which case the top-K is taken per batch entry and the result keeps that
+    dimension. Class indices are relative to a batch entry, as Triton reports
+    them. Labels are read from ``response`` when it is provided.
+    ``triton_dtype`` is the model's output datatype (``FP32``, ``INT32``, …),
+    the same string Triton's frontends pass into ``TopkClassifications``.
+    """
+    # FP32/FP64 and INT/UINT, matching TopkClassifications. FP16 is rejected.
+    is_float = triton_dtype in ("FP32", "FP64")
+    if not is_float and not triton_dtype.startswith(("UINT", "INT")):
+        raise ValueError(
+            "class result not available for output due to unsupported type "
+            f"'{triton_dtype}'"
+        )
+
+    flat = array.reshape(-1)
+    batch_size = int(array.shape[0]) if batched and array.ndim > 0 else 0
+    element_cnt = flat.size // batch_size if batch_size else flat.size
+    if element_cnt > MAX_CLASSIFICATION_ELEMENTS:
+        raise ValueError("classification output tensor too large")
+
+    class_cnt = min(req_class_count, element_cnt)
+    class_strs = []
+    for bs in range(max(1, batch_size)):
+        probs = flat[bs * element_cnt : (bs + 1) * element_cnt]
+        idx = np.argsort(probs)[::-1][:class_cnt]
+        for k in range(class_cnt):
+            kth_class_index = idx[k]
+            # Format the score as a string, matching Triton's std::to_string.
+            score = (
+                f"{probs[kth_class_index]:.6f}"
+                if is_float
+                else str(int(probs[kth_class_index]))
+            )
+            class_str = f"{score}:{kth_class_index}"
+            if response is not None:
+                label = classification_label(
+                    response, output_name, int(kth_class_index)
+                )
+                if label:
+                    class_str += f":{label}"
+            class_strs.append(class_str.encode())
+
+    shape = (batch_size, class_cnt) if batch_size else (class_cnt,)
+    return np.array(class_strs, dtype=object).reshape(shape)
diff --git a/components/src/dynamo/triton/handlers.py b/components/src/dynamo/triton/handlers.py
index 2f1605bf76bf..8fce7667b68d 100644
--- a/components/src/dynamo/triton/handlers.py
+++ b/components/src/dynamo/triton/handlers.py
@@ -12,6 +12,7 @@
 from tritonserver import TritonError
 
 from dynamo.common.backend.health_check import is_probe
+from dynamo.triton.classification import class_count, top_k_classifications
 from dynamo.triton.util import (
     TRITON_TO_DYNAMO_DTYPE,
     dynamo_tensor_to_numpy,
@@ -29,10 +30,12 @@ class RequestHandler:
     def __init__(self, server: TritonServer, model: TritonModel):
         self._server = server
         self._model = model
-        # Output schema is fixed at load time; cache to avoid a per-response lookup.
+        # Output schema and batching are fixed at load time; cache to avoid a
+        # per-response lookup.
         self._output_dtypes: dict[str, str] = {
             out["name"]: out["datatype"] for out in model.metadata()["outputs"]
         }
+        self._batched = model.config().get("max_batch_size", 0) > 0
 
     async def generate(self, request: dict) -> AsyncGenerator[dict, None]:
         logger.debug(f"Received request: {request}")
@@ -43,6 +46,9 @@ async def generate(self, request: dict) -> AsyncGenerator[dict, None]:
             return
 
         self._validate_request(request)
+        requested_outputs = {
+            output["name"]: output for output in request.get("outputs") or []
+        }
 
         inference_request = self._model.create_request()
         for tensor in request["tensors"]:
@@ -54,11 +60,16 @@ async def generate(self, request: dict) -> AsyncGenerator[dict, None]:
         async for inference_response in inference_responses:
             response_tensors = []
             for output_name, output_tensor in inference_response.outputs.items():
+                requested_output = requested_outputs.get(output_name)
+                if requested_outputs and requested_output is None:
+                    continue
+
                 triton_dtype = self._output_dtypes[output_name]
-                if triton_dtype not in TRITON_TO_DYNAMO_DTYPE.keys():
-                    raise TypeError(
-                        f"Requested dtype '{triton_dtype}', for output '{output_name}', is not supported."
-                    )
+                req_class_count = (
+                    class_count(requested_output.get("parameters") or {})
+                    if requested_output is not None
+                    else None
+                )
 
                 if triton_dtype == "BYTES":
                     # String/BYTES tensors are not DLPack-compatible; pull them as
@@ -75,6 +86,21 @@ async def generate(self, request: dict) -> AsyncGenerator[dict, None]:
 
                     response_arr = np.from_dlpack(output_tensor)
 
+                if req_class_count is not None:
+                    response_arr = top_k_classifications(
+                        response_arr,
+                        req_class_count,
+                        triton_dtype,
+                        self._batched,
+                        inference_response,
+                        output_name,
+                    )
+                    triton_dtype = "BYTES"
+                elif triton_dtype not in TRITON_TO_DYNAMO_DTYPE.keys():
+                    raise TypeError(
+                        f"Requested dtype '{triton_dtype}', for output '{output_name}', is not supported."
+                    )
+
                 dtype_str = TRITON_TO_DYNAMO_DTYPE.get(triton_dtype, triton_dtype)
                 response_tensors.append(
                     {
@@ -99,7 +125,9 @@ async def generate(self, request: dict) -> AsyncGenerator[dict, None]:
             yield response
 
     def _validate_request(self, request: dict) -> None:
-        """Reject non-tensor requests and unsupported input dtypes early."""
+        """Reject non-tensor requests, unsupported input dtypes, requested
+        outputs the model does not produce, and an ill-formed classification
+        parameter."""
         if "tensors" not in request:
             raise ValueError(
                 "dynamo.triton only accepts tensor requests; missing 'tensors' "
@@ -116,6 +144,19 @@ def _validate_request(self, request: dict) -> None:
             raise ValueError(
                 f"dynamo does not support {' / '.join(bad_dtypes)} input tensors."
             )
+        unknown = sorted(
+            output["name"]
+            for output in request.get("outputs") or []
+            if output["name"] not in self._output_dtypes
+        )
+        if unknown:
+            names = " / ".join(f"'{name}'" for name in unknown)
+            raise ValueError(
+                f"unexpected inference output {names} for model "
+                f"'{self._model.name}'"
+            )
+        for output in request.get("outputs") or []:
+            class_count(output.get("parameters") or {})
 
     def _probe(self) -> dict:
         """Report readiness without invoking Triton inference.
diff --git a/components/src/dynamo/triton/tests/test_triton_classification.py b/components/src/dynamo/triton/tests/test_triton_classification.py
new file mode 100644
index 000000000000..67932d73c92b
--- /dev/null
+++ b/components/src/dynamo/triton/tests/test_triton_classification.py
@@ -0,0 +1,425 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""KServe classification (class_count / top-K) through the Triton worker.
+
+Covers the two halves separately:
+
+* ``dynamo.triton.classification`` -- ``top_k_classifications``: score formatting,
+  class ordering, output shape, and label lookup.
+* ``RequestHandler`` -- honoring a request's requested outputs, turning a
+  ``classification`` parameter into a BYTES tensor of class strings, and
+  rejecting outputs the model does not produce.
+"""
+
+import asyncio
+import types
+from typing import Any
+from unittest.mock import MagicMock
+
+import numpy as np
+import pytest
+from test_triton_handlers import (
+    _MockModel,
+    assert_dynamo_response,
+    build_dynamo_request,
+    build_triton_response,
+    run_handler_generate,
+)
+
+from dynamo.triton import classification, handlers
+
+pytestmark = [
+    pytest.mark.unit,
+    pytest.mark.triton,
+    pytest.mark.gpu_0,
+    pytest.mark.pre_merge,
+]
+
+
+# --- class_count ------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "parameters, expected",
+    [
+        ({}, None),
+        ({"other": {"bool": True}}, None),
+        ({"classification": {"int64": 3}}, 3),
+        # A plain integer, as an in-process caller would build it.
+        ({"classification": 3}, 3),
+    ],
+    ids=["absent", "unrelated-parameter", "int64-parameter", "plain-int"],
+)
+def test_class_count_reads_the_classification_parameter(parameters, expected):
+    assert classification.class_count(parameters) == expected
+
+
+@pytest.mark.parametrize(
+    "parameters, message",
+    [
+        ({"classification": {"string": "3"}}, "expected int64_param"),
+        ({"classification": {"double": 3.0}}, "expected int64_param"),
+        ({"classification": True}, "expected int64_param"),
+        ({"classification": {"int64": 0}}, "expected >= 1"),
+        ({"classification": {"int64": -1}}, "expected >= 1"),
+    ],
+    ids=["string", "double", "bool", "zero", "negative"],
+)
+def test_class_count_rejects_invalid_parameters(parameters, message):
+    with pytest.raises(ValueError, match=message):
+        classification.class_count(parameters)
+
+
+# --- top_k_classifications -------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "array, triton_dtype, expected",
+    [
+        (
+            np.array([0.1, 0.5, 0.3, 0.2], np.float32),
+            "FP32",
+            [b"0.500000:1", b"0.300000:2"],
+        ),
+        (np.array([1, 3, 2], np.int32), "INT32", [b"3:1", b"2:2"]),
+    ],
+    ids=["fp32", "int32"],
+)
+def test_top_k_classifications_returns_top_scores_in_descending_order(
+    array, triton_dtype, expected
+):
+    classes = classification.top_k_classifications(
+        array, 2, triton_dtype, batched=False
+    )
+
+    assert classes.shape == (2,)
+    assert classes.tolist() == expected
+
+
+def test_top_k_classifications_appends_labels_when_the_output_has_them():
+    owner = types.SimpleNamespace(
+        output_classification_label=lambda _out, class_index: f"label{class_index}"
+    )
+    output = types.SimpleNamespace(memory_buffer=types.SimpleNamespace(owner=owner))
+    response = types.SimpleNamespace(outputs={"OUTPUT0": output})
+
+    classes = classification.top_k_classifications(
+        np.array([0.1, 0.5], np.float32),
+        2,
+        "FP32",
+        batched=False,
+        response=response,
+        output_name="OUTPUT0",
+    )
+
+    assert classes.tolist() == [b"0.500000:1:label1", b"0.100000:0:label0"]
+
+
+def test_top_k_classifications_takes_the_top_k_per_batch_entry():
+    """A batched output is split on its leading dimension, and class indices are
+    relative to the entry, as Triton reports them."""
+    classes = classification.top_k_classifications(
+        np.array([[1, 3, 2], [9, 7, 8]], np.int32), 2, "INT32", batched=True
+    )
+
+    assert classes.shape == (2, 2)
+    assert classes.tolist() == [[b"3:1", b"2:2"], [b"9:0", b"8:2"]]
+
+
+def test_top_k_classifications_indexes_an_unbatched_output_flat():
+    classes = classification.top_k_classifications(
+        np.array([[1, 3, 2], [9, 7, 8]], np.int32), 2, "INT32", batched=False
+    )
+
+    assert classes.shape == (2,)
+    assert classes.tolist() == [b"9:3", b"8:5"]
+
+
+def test_top_k_classifications_clamps_the_count_to_the_elements_available():
+    classes = classification.top_k_classifications(
+        np.array([0.1, 0.2], np.float32), 5, "FP32", batched=False
+    )
+
+    assert classes.shape == (2,)
+
+
+def test_top_k_classifications_handles_an_empty_batch():
+    classes = classification.top_k_classifications(
+        np.zeros((0, 3), np.float32), 2, "FP32", batched=True
+    )
+
+    assert classes.shape == (0,)
+    assert classes.tolist() == []
+
+
+# Triton's TopkClassifications switch: UINT/INT/FP32/FP64 only.
+_UNSUPPORTED_CLASSIFICATION_DTYPES = ("BOOL", "BYTES", "FP16", "BF16")
+
+
+@pytest.mark.parametrize("triton_dtype", _UNSUPPORTED_CLASSIFICATION_DTYPES)
+def test_top_k_classifications_rejects_types_without_an_ordering(triton_dtype):
+    with pytest.raises(ValueError, match=f"unsupported type '{triton_dtype}'"):
+        classification.top_k_classifications(
+            np.array([0.0], np.float32), 1, triton_dtype, batched=False
+        )
+
+
+def test_top_k_classifications_rejects_an_output_too_large_to_sort():
+    oversized = np.zeros(classification.MAX_CLASSIFICATION_ELEMENTS + 1, np.float32)
+
+    with pytest.raises(ValueError, match="too large"):
+        classification.top_k_classifications(oversized, 1, "FP32", batched=False)
+
+
+# --- classification_label ---------------------------------------------------
+
+
+def test_classification_label_looks_up_by_the_outputs_position_in_the_response():
+    """The low-level lookup keys on an output's index, so the helper has to
+    map the name a caller uses back to that position."""
+    owner = types.SimpleNamespace(
+        output_classification_label=lambda output_index, class_index: (
+            f"raw-{output_index}-{class_index}"
+        )
+    )
+    output = types.SimpleNamespace(memory_buffer=types.SimpleNamespace(owner=owner))
+    response = types.SimpleNamespace(
+        outputs={"OUTPUT0": output, "OUTPUT1": output},
+    )
+
+    assert classification.classification_label(response, "OUTPUT1", 3) == "raw-1-3"
+
+
+def test_classification_label_returns_none_when_the_owner_has_no_lookup():
+    """The C response exposes output_classification_label; a host copy
+    whose owner is just the array must not raise."""
+    response = types.SimpleNamespace(
+        outputs={
+            "OUTPUT0": types.SimpleNamespace(
+                memory_buffer=types.SimpleNamespace(owner=object())
+            )
+        },
+    )
+
+    assert classification.classification_label(response, "OUTPUT0", 0) is None
+
+
+def test_classification_label_maps_an_unlabelled_class_to_none():
+    owner = types.SimpleNamespace(output_classification_label=lambda _out, _cls: "")
+    response = types.SimpleNamespace(
+        outputs={
+            "OUTPUT0": types.SimpleNamespace(
+                memory_buffer=types.SimpleNamespace(owner=owner)
+            )
+        },
+    )
+
+    assert classification.classification_label(response, "OUTPUT0", 0) is None
+
+
+# --- RequestHandler ---------------------------------------------------------
+
+
+def _classification_request(*outputs: tuple[str, int | None]) -> dict[str, Any]:
+    """A one-input request asking for the given (output name, class count) pairs."""
+    request = build_dynamo_request(("IN", "Float32", [1], [0.0]))
+    request["outputs"] = [
+        {"name": name}
+        if count is None
+        else {"name": name, "parameters": {"classification": {"int64": count}}}
+        for name, count in outputs
+    ]
+    return request
+
+
+class _LabelledOutput:
+    """A response output that reads as DLPack and carries the low-level response
+    the label lookup goes through, the way a tritonserver output tensor does."""
+
+    def __init__(self, array: np.ndarray) -> None:
+        self._array = array
+        self.memory_buffer = types.SimpleNamespace(
+            owner=types.SimpleNamespace(
+                output_classification_label=(
+                    lambda _output_index, class_index: f"label{class_index}"
+                )
+            )
+        )
+
+    def __dlpack__(self, *args: Any, **kwargs: Any) -> Any:
+        return self._array.__dlpack__(*args, **kwargs)
+
+    def __dlpack_device__(self) -> Any:
+        return self._array.__dlpack_device__()
+
+    def to_bytes_array(self) -> np.ndarray:
+        return self._array
+
+
+def _labelled_response(outputs: dict[str, Any]) -> types.SimpleNamespace:
+    return build_triton_response(
+        "req-id",
+        "classifier",
+        {name: _LabelledOutput(array) for name, array in outputs.items()},
+    )
+
+
+def test_generate_returns_class_strings_for_a_classification_output():
+    _, dynamo_responses = run_handler_generate(
+        [{"name": "OUTPUT0", "datatype": "FP32"}],
+        [_labelled_response({"OUTPUT0": np.array([0.1, 0.5, 0.3], np.float32)})],
+        _classification_request(("OUTPUT0", 2)),
+    )
+
+    assert_dynamo_response(
+        dynamo_responses[0],
+        "req-id",
+        "classifier",
+        {
+            "OUTPUT0": (
+                "Bytes",
+                [2],
+                [list(b"0.500000:1:label1"), list(b"0.300000:2:label2")],
+            )
+        },
+    )
+
+
+def test_generate_keeps_the_batch_dimension_for_a_batchable_model():
+    _, dynamo_responses = run_handler_generate(
+        [{"name": "OUTPUT0", "datatype": "INT32"}],
+        [_labelled_response({"OUTPUT0": np.array([[1, 3], [9, 7]], np.int32)})],
+        _classification_request(("OUTPUT0", 1)),
+        max_batch_size=8,
+    )
+
+    assert_dynamo_response(
+        dynamo_responses[0],
+        "req-id",
+        "classifier",
+        {"OUTPUT0": ("Bytes", [2, 1], [list(b"3:1:label1"), list(b"9:0:label0")])},
+    )
+
+
+def test_generate_mixes_raw_and_classification_outputs():
+    """A request may ask for one output raw and another as classes."""
+    _, dynamo_responses = run_handler_generate(
+        [
+            {"name": "OUTPUT0", "datatype": "FP32"},
+            {"name": "OUTPUT1", "datatype": "FP32"},
+        ],
+        [
+            _labelled_response(
+                {
+                    "OUTPUT0": np.array([0.25, 0.75], np.float32),
+                    "OUTPUT1": np.array([0.5], np.float32),
+                }
+            )
+        ],
+        _classification_request(("OUTPUT0", 1), ("OUTPUT1", None)),
+    )
+
+    assert_dynamo_response(
+        dynamo_responses[0],
+        "req-id",
+        "classifier",
+        {
+            "OUTPUT0": ("Bytes", [1], [list(b"0.750000:1:label1")]),
+            "OUTPUT1": ("Float32", [1], [0.5]),
+        },
+    )
+
+
+def test_generate_returns_only_the_requested_outputs():
+    _, dynamo_responses = run_handler_generate(
+        [
+            {"name": "OUTPUT0", "datatype": "FP32"},
+            {"name": "OUTPUT1", "datatype": "INT32"},
+        ],
+        [
+            build_triton_response(
+                "req-id",
+                "classifier",
+                {
+                    "OUTPUT0": np.array([1.5], np.float32),
+                    "OUTPUT1": np.array([7], np.int32),
+                },
+            )
+        ],
+        _classification_request(("OUTPUT1", None)),
+    )
+
+    assert_dynamo_response(
+        dynamo_responses[0],
+        "req-id",
+        "classifier",
+        {"OUTPUT1": ("Int32", [1], [7])},
+    )
+
+
+def test_generate_returns_every_output_when_outputs_is_empty():
+    """An empty requested-output list is the KServe 'all outputs' default."""
+    request = build_dynamo_request(("IN", "Float32", [1], [0.0]))
+    request["outputs"] = []
+    _, dynamo_responses = run_handler_generate(
+        [
+            {"name": "OUTPUT0", "datatype": "FP32"},
+            {"name": "OUTPUT1", "datatype": "INT32"},
+        ],
+        [
+            build_triton_response(
+                "req-id",
+                "classifier",
+                {
+                    "OUTPUT0": np.array([1.5], np.float32),
+                    "OUTPUT1": np.array([7], np.int32),
+                },
+            )
+        ],
+        request,
+    )
+
+    assert set(
+        tensor["metadata"]["name"] for tensor in dynamo_responses[0]["tensors"]
+    ) == {"OUTPUT0", "OUTPUT1"}
+
+
+def test_generate_rejects_an_output_the_model_does_not_produce():
+    with pytest.raises(ValueError, match="unexpected inference output 'OUTPUT1'"):
+        run_handler_generate(
+            [{"name": "OUTPUT0", "datatype": "FP32"}],
+            [build_triton_response("req-id", "classifier", {})],
+            _classification_request(("OUTPUT1", 1)),
+        )
+
+
+def test_generate_rejects_an_invalid_classification_parameter_before_infer():
+    """classification: 0 is a request error; the model must not be invoked."""
+    model = _MockModel(
+        [{"name": "OUTPUT0", "datatype": "FP32"}],
+        [build_triton_response("req-id", "classifier", {})],
+    )
+    handler = handlers.RequestHandler(MagicMock(), model)
+
+    with pytest.raises(ValueError, match="expected >= 1"):
+        asyncio.run(
+            handler.generate(_classification_request(("OUTPUT0", 0))).__anext__()
+        )
+
+    assert model.last_request is None
+
+
+@pytest.mark.parametrize("triton_dtype", _UNSUPPORTED_CLASSIFICATION_DTYPES)
+def test_generate_rejects_classification_on_an_unsupported_output_dtype(triton_dtype):
+    array = (
+        np.array([b"a"], dtype=object)
+        if triton_dtype == "BYTES"
+        else np.array([0.0], np.float32)
+    )
+    with pytest.raises(ValueError, match=f"unsupported type '{triton_dtype}'"):
+        run_handler_generate(
+            [{"name": "OUTPUT0", "datatype": triton_dtype}],
+            [_labelled_response({"OUTPUT0": array})],
+            _classification_request(("OUTPUT0", 1)),
+        )
diff --git a/components/src/dynamo/triton/tests/test_triton_handlers.py b/components/src/dynamo/triton/tests/test_triton_handlers.py
index db2cb5ef10ce..65720fe26a52 100644
--- a/components/src/dynamo/triton/tests/test_triton_handlers.py
+++ b/components/src/dynamo/triton/tests/test_triton_handlers.py
@@ -27,10 +27,16 @@ class _MockModel:
     """Records the request it builds and replays a fixed response stream."""
 
     def __init__(
-        self, output_metadata: list[dict[str, Any]], responses: list[Any]
+        self,
+        output_metadata: list[dict[str, Any]],
+        responses: list[Any],
+        max_batch_size: int = 0,
+        name: str = "mock-model",
     ) -> None:
         self._output_metadata = output_metadata
         self._responses = responses
+        self._max_batch_size = max_batch_size
+        self.name = name
         self.last_request: types.SimpleNamespace | None = None
 
     def create_request(self) -> types.SimpleNamespace:
@@ -40,6 +46,9 @@ def create_request(self) -> types.SimpleNamespace:
     def metadata(self) -> dict[str, Any]:
         return {"outputs": self._output_metadata}
 
+    def config(self) -> dict[str, Any]:
+        return {"max_batch_size": self._max_batch_size}
+
     def async_infer(self, _inference_request: Any) -> AsyncIterator[Any]:
         async def _stream() -> AsyncIterator[Any]:
             for response in self._responses:
@@ -78,9 +87,10 @@ def run_handler_generate(
     triton_output_metadata: list[dict[str, Any]],
     triton_responses: list[Any],
     dynamo_request: dict[str, Any],
+    max_batch_size: int = 0,
 ) -> tuple[_MockModel, list[dict[str, Any]]]:
     """Build a RequestHandler over a _MockModel and drive generate to completion."""
-    model = _MockModel(triton_output_metadata, triton_responses)
+    model = _MockModel(triton_output_metadata, triton_responses, max_batch_size)
     handler = handlers.RequestHandler(MagicMock(), model)
 
     async def _collect() -> list[dict[str, Any]]:
@@ -292,6 +302,7 @@ def _readiness(value: bool | Exception) -> MagicMock:
     model.ready = _readiness(model_ready)
     model.name = model_name
     model.metadata = MagicMock(return_value={"outputs": output_metadata})
+    model.config = MagicMock(return_value={})
     model.async_infer = MagicMock(
         side_effect=AssertionError("async_infer must not be called")
     )
diff --git a/components/src/dynamo/triton/tests/test_triton_health_check.py b/components/src/dynamo/triton/tests/test_triton_health_check.py
index 8ed33f24495a..3f9d4c06b67a 100644
--- a/components/src/dynamo/triton/tests/test_triton_health_check.py
+++ b/components/src/dynamo/triton/tests/test_triton_health_check.py
@@ -65,6 +65,7 @@ def _readiness(value: bool | Exception) -> MagicMock:
     model.ready = _readiness(model_ready)
     model.name = model_name
     model.metadata = MagicMock(return_value={"outputs": output_metadata})
+    model.config = MagicMock(return_value={})
     model.async_infer = MagicMock(
         side_effect=AssertionError("async_infer must not be called")
     )
diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/backends/triton/overview.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/backends/triton/overview.md
index d11cca3fa220..4d0888d1ec5a 100644
--- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/backends/triton/overview.md
+++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/backends/triton/overview.md
@@ -17,6 +17,7 @@ That image is the build's `RUNTIME_IMAGE`, so the result is a single "Dynamo + T
 | Feature                                        | Status  | Notes                                                           |
 | :--------------------------------------------- | :-----: | :-------------------------------------------------------------- |
 | Tensor (KServe gRPC) Serving                   |  Ready  | Multiple models per worker                                      |
+| Classification (`class_count` / top-K)         |  Ready  | Top-K `"<score>:<index>[:<label>]"` class strings               |
 | Service Discovery / Routing                    |  Ready  | Via the Dynamo Frontend                                         |
 | Triton backends (TensorRT, ONNX, PyTorch, ...) |  Ready  | Whatever the Triton release image ships                         |
 | TensorRT Plugins                               |  Ready  | Via `--backend-config='tensorrt,plugins=...'`                   |
@@ -31,7 +32,6 @@ The Triton and KServe features below are not yet supported end-to-end through th
 | Limitation                             | Effect                                                                                                                                                                                                                                              |
 | :------------------------------------- | :-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
 | FP16 / BFloat16 tensors                | Requests with half-precision input/output tensors are rejected (`unknown variant 'Float16'` / `Invalid data type`). The Dynamo tensor `DataType` enum has no `Float16`/`BFloat16` variant.                                                          |
-| Classification (`class_count` / top-K) | The KServe requested-output `classification` parameter is dropped, so models return raw output logits instead of the top-K `"<score>:<index>:<label>"` class strings. `NvCreateTensorRequest` does not carry requested outputs or their parameters. |
 | Model version selection                | Dynamo routes by model name only (`triton.tritonserver.<model_name>`) with no version concept, so version-specific requests are served by the model's default version.                                                                              |
 | Shared-memory tensor I/O               | System / CUDA shared-memory inputs and outputs are unsupported — the frontend does not expose the KServe shared-memory region control RPCs. Send tensors inline.                                                                                    |
 

From a519b27513023da00e68da7636628240a4ef2e8a Mon Sep 17 00:00:00 2001
From: Yingge He <yinggeh@nvidia.com>
Date: Thu, 10 Sep 2026 09:35:06 -0700
Subject: [PATCH 3/4] test(triton): slim handler dtype coverage and drop
 restating docstrings

Keep the full unsupported-dtype matrix on the helper; the handler only
needs BYTES plus one numeric reject path.

Signed-off-by: Yingge He <yinggeh@nvidia.com>
---
 .../src/dynamo/triton/tests/test_triton_classification.py     | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/components/src/dynamo/triton/tests/test_triton_classification.py b/components/src/dynamo/triton/tests/test_triton_classification.py
index 67932d73c92b..091d14f37ebe 100644
--- a/components/src/dynamo/triton/tests/test_triton_classification.py
+++ b/components/src/dynamo/triton/tests/test_triton_classification.py
@@ -303,7 +303,6 @@ def test_generate_keeps_the_batch_dimension_for_a_batchable_model():
 
 
 def test_generate_mixes_raw_and_classification_outputs():
-    """A request may ask for one output raw and another as classes."""
     _, dynamo_responses = run_handler_generate(
         [
             {"name": "OUTPUT0", "datatype": "FP32"},
@@ -395,7 +394,6 @@ def test_generate_rejects_an_output_the_model_does_not_produce():
 
 
 def test_generate_rejects_an_invalid_classification_parameter_before_infer():
-    """classification: 0 is a request error; the model must not be invoked."""
     model = _MockModel(
         [{"name": "OUTPUT0", "datatype": "FP32"}],
         [build_triton_response("req-id", "classifier", {})],
@@ -410,7 +408,7 @@ def test_generate_rejects_an_invalid_classification_parameter_before_infer():
     assert model.last_request is None
 
 
-@pytest.mark.parametrize("triton_dtype", _UNSUPPORTED_CLASSIFICATION_DTYPES)
+@pytest.mark.parametrize("triton_dtype", ("BYTES", "FP16"))
 def test_generate_rejects_classification_on_an_unsupported_output_dtype(triton_dtype):
     array = (
         np.array([b"a"], dtype=object)

From 6a109febf308823e24f55cfb8a01c71686618571 Mon Sep 17 00:00:00 2001
From: Yingge He <yinggeh@nvidia.com>
Date: Tue, 15 Sep 2026 12:37:50 -0700
Subject: [PATCH 4/4] fix test

---
 components/src/dynamo/triton/tests/test_triton_main.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/components/src/dynamo/triton/tests/test_triton_main.py b/components/src/dynamo/triton/tests/test_triton_main.py
index 20a4eea8c9c8..961438334ece 100644
--- a/components/src/dynamo/triton/tests/test_triton_main.py
+++ b/components/src/dynamo/triton/tests/test_triton_main.py
@@ -47,6 +47,7 @@ def test_register_and_serve_registers_and_serves(patched_worker, tmp_path):
     config.server_id = "triton"
 
     loaded_model = MagicMock(name="model")
+    loaded_model.config.return_value = {}
     server = MagicMock(name="server")
     server.model.return_value = loaded_model