NVIDIA / NVIDIA/recsys-examples

[FEA] Add KV-cache-aware Triton orchestration and move offload/reap off the inference critical path`

Open
#474 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
322
Forks
83
Avg merge
3d 3h
Merged PRs (30d)
8

Description

Is your feature request related to a problem? Please describe.
The current Triton AOTI deployment runs a single stateful model instance and disables Triton scheduler-level batching:

The exported forward embeds lookup, allocation, onboarding launch/wait, and HSTU compute in one AOTI call: ExportKVCachedInferenceRankingGR.forward. Offload/reap operators exist in the C++ runtime, but they are not part of this serving forward path. If they are added directly to the same AOTI call, FlexKV RPC, GPU-to-host transfer, and task reaping can extend the request critical path and block the only model instance.

The broad scheduling problem is already tracked by #461. This issue should be a concrete implementation subtask focused on separating the KV-cache lifecycle from the compute graph.

Describe the solution you'd like
A clear and concise description of what you want to happen.

Describe alternatives you've considered
A clear and concise description of any alternative solutions or features you've considered.

Additional context
Add any other context, code examples, or references to existing implementations about the feature request here.


By submitting this issue, you agree to follow our code of conduct and our contributing guidelines.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with examples/hstu/inference_aoti/triton_aoti/hstu_gr_ranking_kvcache/config.pbtxt and examples/hstu/model/export_kvcached_inference_ranking_gr.py, then inspect the existing C++ offload/reap runtime and related scheduling work in #461. Define a lifecycle that separates KV-cache offload and reap from the exported compute forward. Done means the serving path no longer blocks the model instance on those operations while preserving KV-cache behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
ai, backend
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.