Skip to content

feat: add native continuous batching and chunked prefill for LLM runtime - #215

Open
ajmalrasi wants to merge 7 commits into
NVIDIA:mainfrom
ajmalrasi:feature/continuous-batching
Open

ajmalrasi wants to merge 7 commits into
NVIDIA:mainfrom
ajmalrasi:feature/continuous-batching

Conversation

@ajmalrasi

Copy link
Copy Markdown

Summary

Adds native continuous batching for the LLM runtime, allowing up to two active
text-generation sequences to share an existing batch-two TensorRT engine.

The implementation is built around persistent sequence-local runtime state,
rather than concurrent calls to the existing single-request API.

What changed

  • Added stable two-slot sequence ownership with generation-tagged handles.
  • Added persistent request-local KV, recurrent, convolution, prompt and output state.
  • Added selected-row native execution so one or both physical slots can run safely.
  • Added bounded chunked prefill with a qualified 128-token policy.
    • Resumed final chunks use 64–128 tokens.
    • Remainders of 129–191 tokens are partitioned safely.
  • Added a decode-first native scheduler with:
    • bounded FIFO admission,
    • round-robin prefill progress,
    • safe-boundary slot reuse,
    • cancellation and failure propagation.
  • Added independent request controls for sampling, output limits, EOS, stop strings,
    thinking, logit bias and logprobs.
  • Added native ticket/channel integration for independent HTTP and SSE responses.
  • Added streaming/non-streaming handling, disconnect cleanup, tool/reasoning parsing,
    usage accounting and overload responses.
  • Added finite startup CUDA graph captures for slot views {0}, {1} and {0,1}.

Validation

  • Host regression suite: 146 passing tests.
  • Native scheduler/state tests passed with ASan/UBSan and ThreadSanitizer coverage.
  • Jetson Orin validation confirmed exact serial/concurrent output equality for the
    qualified execution paths.
  • Chunked-prefill validation covered supported boundaries through 6,144-token prompts.
  • HTTP gates covered staggered requests, mixed streaming, cancellation, EOS/stops,
    logprobs, tool calling, reasoning, overload and cleanup.
  • Performance qualification on a batch-two Orin engine:
    • aggregate output throughput: 22.54 → 33.97 tok/s (1.51×)
    • singleton latency: within 4.3% of the legacy path
    • staggered second-request TTFT: 2.048 → 0.202 seconds

Explicit non-goals

  • No cross-request context/KV reuse or prefix caching.
  • No speculative decoding or MTP.
  • No arbitrary chunk boundaries; the raw short-tail 64+64+1 path remains unsupported.
  • Scope is the validated single-rank, text-only, batch-two runtime.

Signed-off-by: ajmalrasi <kpajmalrasi@gmail.com>
Signed-off-by: ajmalrasi <kpajmalrasi@gmail.com>
Signed-off-by: ajmalrasi <kpajmalrasi@gmail.com>
Signed-off-by: ajmalrasi <kpajmalrasi@gmail.com>
Signed-off-by: ajmalrasi <kpajmalrasi@gmail.com>
Signed-off-by: ajmalrasi <kpajmalrasi@gmail.com>
Signed-off-by: ajmalrasi <kpajmalrasi@gmail.com>
@ajmalrasi
ajmalrasi requested a review from a team September 15, 2026 21:04
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant