[peihu-nv]

@coderabbitai summary

Description

K3 checkpoint-planning foundation only. Full support requires the relevant Rubin-advance merge-back and porting the integrated implementation from internal MR !10793. Retention additionally requires #20001.

K3 cache reuse needs the complete attention prefix and the recurrent/convolution state at the same checkpoint endpoint. This draft adds a transfer planner using the existing connector layout and block hash chain. It describes every payload and scale region, separates shared attention objects from endpoint-private state, and validates a completion manifest before restore.

This is the preparatory checkpoint contract, not working K3 Mooncake serving support. Native snapshot pinning, transfer retirement, request adoption, and executor integration remain to be implemented and validated. The producer must supply an independently pinned exact checkpoint and a representation identity covering weights and layout. Sliding-window layouts are rejected; existing serving restrictions remain intact.

The intended integration uses the generic connector work in #19171, the lending lifecycle in #19722, and private-checkpoint retention in #20001. This draft does not copy their transport, allocator, or retention implementations and can be reviewed independently.

Test Coverage

• 31 CPU cases passed in tests/unittest/_torch/executor/test_mooncake_k3_checkpoint.py: payload/scales coverage, endpoint and representation isolation, address-independent manifests, partial/failed completion, invalid slots, and unsupported layouts.
• Local execution used real PyTorch and connector layout metadata with the native package bootstrap bypassed; unused native conversion helpers were fail-fast stubs. This is not a full installed-runtime test.
• Changed-file pre-commit checks passed. GPU/RDMA, native ownership, full-model correctness, and E2E performance are not validated by this draft.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[peihu-nv]

/bot run

[tensorrt-cicd]

PR_Github #76788 [ run ] triggered by Bot. Commit: 2a2eda8 Link to invocation

[peihu-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #77064 [ run ] triggered by Bot. Commit: 2a2eda8 Link to invocation

[tensorrt-cicd]

PR_Github #77064 [ run ] completed with state SUCCESS. Commit: 2a2eda8
/LLM/main/L0_MergeRequest_PR pipeline #63532 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[2ez4bz]

@peihu-nv how much of this is specific to K3 vs generic for other mamba / hybrid models?