fix: Fix CUDA graph execution for non-batching models

[pskiran1]

This PR fixes CUDA graph lookup for non-batching models (max_batch_size: 0).
For these models, the captured graph key used 0 as its batch-size prefix, while the runtime lookup key used 1 because one non-batched request is treated as one execution unit. The mismatch prevented Triton from finding the captured graph, causing inference to fall back to regular TensorRT execution.

• Uses 1 as the capture-key prefix for non-batching models, matching runtime lookup.
• Keeps CUDA graph lower-bound keys aligned with input dimensions for non-batching models.

TensorRT input shapes and batching behavior are unchanged.

Closes triton-inference-server/server#7150

[greptile-apps]

Greptile Summary

This PR fixes a long-standing CUDA graph lookup miss for non-batching models (max_batch_size: 0). At graph capture time, the key's batch-size prefix was 0, but at inference time it was 1 (derived from total_batch_size_), so the map lookup always failed and the engine fell back to ordinary TRT execution.

• Key prefix fix: SetCudaGraphShape now uses 1 as the batch-size prefix when batch_size == 0, matching how the runtime key is constructed at line 969 (std::vector<int64_t> input_dims{(int64_t)payload_->total_batch_size_}).
• lower_bound_key alignment: The lower_bound_key.push_back(lower_bound_key[0]) call that inserts a per-input batch lower-bound element is moved inside the if (batch_size != 0) guard, removing spurious elements that caused the key and its lower-bound vector to differ in length for non-batching models.

Confidence Score: 5/5

Safe to merge; the changes are surgical, well-commented, and only affect the non-batching CUDA graph path.

Both hunks make narrow, correctly reasoned changes. The key-prefix fix directly matches how input_dims is initialised at inference time. The lower_bound_key guard prevents a structural length mismatch that would break the inexact-match bounds check for non-batching models. Batching models are unaffected because both new statements remain inside if (batch_size != 0).

No files require special attention.

Important Files Changed

Filename	Overview
src/instance_state.cc	Two targeted changes to SetCudaGraphShape: aligns the capture-time key prefix to 1 (matching runtime), and gates the per-input batch lower-bound push behind the batching guard — both correct and well-commented.

Sequence Diagram

sequenceDiagram
    participant GS as Graph Capture (SetCudaGraphShape)
    participant Map as cuda_graph_execs_ (sorted map)
    participant RT as Inference (RunInfer)
    participant LK as FindClosestCudaGraph

    Note over GS: max_batch_size == 0 (non-batching model)
    GS->>GS: "batch_size = graph_spec.batch_size_ → 0"
    GS->>GS: "(BEFORE FIX) key prefix = {0}"
    GS->>GS: "(AFTER FIX)  key prefix = {1}"
    GS->>GS: Per-input: append raw shape dims (no batch dim prepended)
    GS->>GS: lower_bound_key: skip push_back(lower_bound_key[0]) for non-batching
    GS->>Map: "insert key = {1, shape_dims...}"

    Note over RT: Single non-batched request arrives
    RT->>RT: "input_dims = {total_batch_size_} = {1}"
    RT->>RT: "For each input: append batchn_shape (no batch dim, support_batching_=false)"
    RT->>RT: "input_dims = {1, shape_dims...}"

    RT->>LK: "FindClosestCudaGraph(input_dims={1, shape_dims...})"
    LK->>Map: "lower_bound({1, shape_dims...})"
    Map-->>LK: (BEFORE FIX) miss — stored key starts with 0
    Map-->>LK: (AFTER FIX) hit — stored key starts with 1
    LK-->>RT: cuda_graph (found)

Loading

_{Reviews (1): Last reviewed commit: "Update" | Re-trigger Greptile}