Repository navigation
[TRTLLM-17359][fix] Preserve per-rank attention-DP KV metrics - #21131
Conversation
32fb605 to
bf54a05
Compare
|
/bot run |
|
PR_Github #80975 [ run ] triggered by Bot. Commit: |
bf54a05 to
e02d1d2
Compare
|
No actionable comments were generated in the recent review.
 |
| Layer / File(s) | Summary |
|---|---|
KV snapshot wire format tensorrt_llm/_torch/pyexecutor/scheduler/rank_state.py, tests/unittest/_torch/executor/test_adp_router.py |
Rank iteration payloads serialize optional KV-cache snapshots with versioned framing. Deserialization validates snapshot headers and payload lengths. Tests cover round trips, wire fields, and malformed payloads. |
Rank-local snapshot fanout tensorrt_llm/_torch/pyexecutor/adp_iter_stats.py, tests/unittest/_torch/executor/test_iter_stats_populate.py |
Payload creation captures self and cross-KV snapshots, and fanout restores available snapshots for each rank. Synthetic idle-rank payloads omit snapshots, and measured payloads replace placeholders for the same iteration. |
Rank-labeled metrics and ownership tensorrt_llm/metrics/collector.py, tests/unittest/metrics/test_collector.py, tests/unittest/grpc/openengine/test_metrics.py, docs/source/commands/trtllm-serve/trtllm-serve.rst |
The collector labels KV-cache and local scheduling metrics by rank, tracks cumulative counters per rank, and logs shared metrics only for rank 0. Tests and documentation cover rank labels, ownership, counter handling, and aggregation. |
Priority:
Estimated code review effort: 4 (Complex) | ~45 minutes
Severity of issue fixed: Medium
Sequence Diagram(s)
sequenceDiagram
participant AttentionDPRankRow
participant MetricsCollector
participant PrometheusMetrics
AttentionDPRankRow->>MetricsCollector: log_iteration_stats row
MetricsCollector->>MetricsCollector: resolve rank and metric ownership
MetricsCollector->>PrometheusMetrics: emit rank-labeled metrics
Suggested reviewers: tanmayv25
Merge Risk: 
bc019
-
No actionable change-specific issue remains from the reviewed evidence. The change is mergeable after normal checks.
Pre-merge checks: 
-
Check name Status Explanation Title check The title clearly identifies the fix and the primary change: preserving per-rank attention-DP KV metrics. It is concise and follows the repository’s ticket and type format. Linked Issues check Issue [ #21124] requires rank-local Attention-DP KV metrics, independent per-rank cumulative KV histories with reset behavior, and ownership rules for shared metrics. The PR captures and transports imm…Out of Scope Changes check The source changes implement the requirements of [ #21124]. The added tests verify snapshot capture, transport, collection, ownership, and serving-facing metrics. The documentation defines rank labels,…Description check The description follows the repository's template in substance. It includes a summary of the Prometheus label-schema change (dp_rank on 28 metric families), scope of the code changes, a dashboard migr… Full details: Linked Issues check
-
Explanation
Issue [
#21124] requires rank-local Attention-DP KV metrics, independent per-rank cumulative KV histories with reset behavior, and ownership rules for shared metrics. The PR captures and transports immutable primary and cross-KV snapshots, emits rank-local series withdp_rank, and maintains per-rank counter baselines. Counter decreases warn, add no increment, and rebaseline the affected rank. Rank-0 ownership prevents duplicate shared speculative-token increments. Tests cover snapshot framing, missing snapshots, rank-local values, counter decreases, shared ownership, request/config label independence, and OpenEngine rank gauges.
Full details: Out of Scope Changes check
-
Explanation
The source changes implement the requirements of [
#21124]. The added tests verify snapshot capture, transport, collection, ownership, and serving-facing metrics. The documentation defines rank labels, freshness, ownership, and aggregation. The summary demonstrates no unrelated change.
Full details: Description check
-
Explanation
The description follows the repository's template in substance. It includes a summary of the Prometheus label-schema change (dp_rank on 28 metric families), scope of the code changes, a dashboard migration note, and a verification section with test counts and the exact pytest command. The PR checklist is marked as reviewed. The template's title and Test Coverage sections are not explicitly labeled. The description also reports a pending full CI run and an inconclusive live check. The release-owner note is clear. Overall the description is mostly complete.
-
Finishing Touches
-
🧪 Generate unit tests (beta)
-
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
|
/bot run --disable-fail-fast |
|
PR_Github #81000 [ run ] triggered by Bot. Commit: |
|
PR_Github #80975 [ run ] completed with state |
|
PR_Github #81000 [ run ] completed with state
|
|
Automatically added "ci: full pre-merge approved" because this PR has satisfied the required GitHub review approvals. Unresolved review conversations and other required checks remain independent merge requirements. |
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Approving — the comments below are optional touch-ups, not blockers.
Two things the description should say up front, since both outlive this PR:
-
This is a Prometheus label-schema change, not only a fix.
dp_rankis added to roughly 25 existing metric families and appears in single-rank serving too. Any existing recording rule or alert that doessum by (model_name)overtrtllm_kv_cache_utilizationorkv_cache_hit_ratenow sums across ranks, which the new rst section correctly calls wrong. The rst covers it, but the[fix]title and the Summary paragraph do not. Please lead with the label change in the description, and consider whetherdocs/source/release-notes.mdshould carry a line for it (if entries are only added at release time, a note to the release owner is enough). -
The fanout also fixed an N× inflation of
numCompletedRequestsand other rank-0-owned counters. Before this change every fanned-out row carried rank 0's scalar fields, socounter_num_requests_completedand the spec-decode token counters were incremented once per rank.test_rank0_owned_fields_are_not_overwritten_or_counted_twicepins the new behavior, but the description only mentions occupancy and reuse/miss. Worth a sentence, because someone diffing dashboards before and after will see completed-request rates drop by the DP size.
On the transport: the -1 header cannot collide with has_iter_stats, the allgather pickles so the float hit rate survives, and mixed compact/versioned frames in one gather are tested. The compact "legacy" frame is only ever produced by the same build (idle ranks), so the length-based legacy inference in RankState.deserialize is complexity without a cross-version consumer. Not blocking, but always emitting the versioned frame (presence flags 0/0 for idle ranks) would delete that branch and the ambiguity it guards against.
Inline comments cover the two behavior choices I would like answered before merge: counting a cumulative decrease as full new work, and leaving idle ranks' KV gauges frozen.
e02d1d2 to
41e8827
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #81333 [ run ] triggered by Bot. Commit: |
|
PR_Github #81333 [ run ] completed with state Resource maintenance in LLM/main/L0_Test-x86_64-Single-GPU #9638: L40S-PyTorch-2 is temporarily paused due to Kubernetes agent provisioning and node issues.
|
Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
41e8827 to
bc01999
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #81641 [ run ] triggered by Bot. Commit: |
|
PR_Github #81641 [ run ] completed with state |
… stats Attention-DP iteration stats fan out one row per rank (NVIDIA#21131): the scheduling and KV fields of each row come from that rank's payload on the ADP rank-state allgather, but specdec_stats was in the list of nested objects copied from rank 0. Every rank's row therefore reported rank 0's draft and accepted token counts and acceptance length, so per-rank acceptance could not be told apart, an idle rank 0 showed zero acceptance for the whole group, and summing the rows multiplied rank 0's numbers by the number of ranks. Any attention-DP deployment with speculative decoding (for example DeepSeek-R1 DEP with MTP) and iteration stats enabled reports these rows, through LLM.get_stats() and the JSON iteration stats of trtllm-serve's /metrics; the Prometheus metrics (/prometheus/metrics) read spec-decoding stats from rank 0's row only and do not change. The acceptance-length gates of the accuracy tests sum every row (acceptance_length_from_iteration_stats), so under attention DP they measured rank 0's acceptance alone and could not see a drafting regression on the other ranks. Found while checking Mistral Large 4 DEP4 acceptance with its DSpark drafter. Each rank now snapshots its draft tokens, accepted tokens and requests with draft tokens (RankSpecDecodingStats) into its payload, and the fanout rebuilds the row's SpecDecodingStats from them, deriving the acceptance length from the rank's own counters. The draft timings stay rank 0's, like the row's other timings. A rank with no snapshot (an explicit zero placeholder) gets no spec-decoding stats rather than rank 0's. The snapshot frame gains a third presence flag, so its version moves from -1 to -2 and its header from four to five fields. Every rank of a job runs the same build, so all ranks write and read the same version; a payload with another version fails the header check instead of being misread. Tests: test_adp_router.py pins the new wire schema and round-trips the snapshot with and without KV snapshots and load fields, keeping zero counters distinct from an absent snapshot. test_iter_stats_populate.py fans out asymmetric work on two ranks (with rank 0 busy or idle, and the padding and CUDA-graph dummies excluded) and checks each row's counts and acceptance length, that the rows add up to the group's acceptance, and that a placeholder rank does not inherit rank 0's drafting. Split out of NVIDIA#21693 (Mistral Large 4 merge-back); originally GitLab 4932efdc24 ("Add native DSpark to StoneWillow Staircase specializations"), rebased onto the versioned snapshot frame of NVIDIA#21131. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
… stats Attention-DP iteration stats fan out one row per rank (NVIDIA#21131): the scheduling and KV fields of each row come from that rank's payload on the ADP rank-state allgather, but specdec_stats was in the list of nested objects copied from rank 0. Every rank's row therefore reported rank 0's draft and accepted token counts and acceptance length, so per-rank acceptance could not be told apart, an idle rank 0 showed zero acceptance for the whole group, and summing the rows multiplied rank 0's numbers by the number of ranks. Any attention-DP deployment with speculative decoding (for example DeepSeek-R1 DEP with MTP) and iteration stats enabled reports these rows, through LLM.get_stats() and the JSON iteration stats of trtllm-serve's /metrics; the Prometheus metrics (/prometheus/metrics) read spec-decoding stats from rank 0's row only and do not change. The acceptance-length gates of the accuracy tests sum every row (acceptance_length_from_iteration_stats), so under attention DP they measured rank 0's acceptance alone and could not see a drafting regression on the other ranks. Found while checking Mistral Large 4 DEP4 acceptance with its DSpark drafter. Each rank now snapshots its draft tokens, accepted tokens and requests with draft tokens (RankSpecDecodingStats) into its payload, and the fanout rebuilds the row's SpecDecodingStats from them, deriving the acceptance length from the rank's own counters. The draft timings stay rank 0's, like the row's other timings. A rank with no snapshot (an explicit zero placeholder) gets no spec-decoding stats rather than rank 0's. The snapshot frame gains a third presence flag, so its version moves from -1 to -2 and its header from four to five fields. Every rank of a job runs the same build, so all ranks write and read the same version; a payload with another version fails the header check instead of being misread. Tests: test_adp_router.py pins the new wire schema and round-trips the snapshot with and without KV snapshots and load fields, keeping zero counters distinct from an absent snapshot. test_iter_stats_populate.py fans out asymmetric work on two ranks (with rank 0 busy or idle, and the padding and CUDA-graph dummies excluded) and checks each row's counts and acceptance length, that the rows add up to the group's acceptance, and that a placeholder rank does not inherit rank 0's drafting. Split out of NVIDIA#21693 (Mistral Large 4 merge-back); originally GitLab 4932efdc24 ("Add native DSpark to StoneWillow Staircase specializations"), rebased onto the versioned snapshot frame of NVIDIA#21131. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
… stats Attention-DP iteration stats fan out one row per rank (NVIDIA#21131): the scheduling and KV fields of each row come from that rank's payload on the ADP rank-state allgather, but specdec_stats was in the list of nested objects copied from rank 0. Every rank's row therefore reported rank 0's draft and accepted token counts and acceptance length, so per-rank acceptance could not be told apart, an idle rank 0 showed zero acceptance for the whole group, and summing the rows multiplied rank 0's numbers by the number of ranks. Any attention-DP deployment with speculative decoding (for example DeepSeek-R1 DEP with MTP) and iteration stats enabled reports these rows, through LLM.get_stats() and the JSON iteration stats of trtllm-serve's /metrics; the Prometheus metrics (/prometheus/metrics) read spec-decoding stats from rank 0's row only and do not change. The acceptance-length gates of the accuracy tests sum every row (acceptance_length_from_iteration_stats), so under attention DP they measured rank 0's acceptance alone and could not see a drafting regression on the other ranks. Found while checking Mistral Large 4 DEP4 acceptance with its DSpark drafter. Each rank now snapshots its draft tokens, accepted tokens and requests with draft tokens (RankSpecDecodingStats) into its payload, and the fanout rebuilds the row's SpecDecodingStats from them, deriving the acceptance length from the rank's own counters. The draft timings stay rank 0's, like the row's other timings. A rank with no snapshot (an explicit zero placeholder) gets no spec-decoding stats rather than rank 0's. The snapshot frame gains a third presence flag, so its version moves from -1 to -2 and its header from four to five fields. Every rank of a job runs the same build, so all ranks write and read the same version; a payload with another version fails the header check instead of being misread. Tests: test_adp_router.py pins the new wire schema and round-trips the snapshot with and without KV snapshots and load fields, keeping zero counters distinct from an absent snapshot. test_iter_stats_populate.py fans out asymmetric work on two ranks (with rank 0 busy or idle, and the padding and CUDA-graph dummies excluded) and checks each row's counts and acceptance length, that the rows add up to the group's acceptance, and that a placeholder rank does not inherit rank 0's drafting. Split out of NVIDIA#21693 (Mistral Large 4 merge-back); originally GitLab 4932efdc24 ("Add native DSpark to StoneWillow Staircase specializations"), rebased onto the versioned snapshot frame of NVIDIA#21131. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
… stats Attention-DP iteration stats fan out one row per rank (NVIDIA#21131): the scheduling and KV fields of each row come from that rank's payload on the ADP rank-state allgather, but specdec_stats was in the list of nested objects copied from rank 0. Every rank's row therefore reported rank 0's draft and accepted token counts and acceptance length, so per-rank acceptance could not be told apart, an idle rank 0 showed zero acceptance for the whole group, and summing the rows multiplied rank 0's numbers by the number of ranks. Any attention-DP deployment with speculative decoding (for example DeepSeek-R1 DEP with MTP) and iteration stats enabled reports these rows, through LLM.get_stats() and the JSON iteration stats of trtllm-serve's /metrics; the Prometheus metrics (/prometheus/metrics) read spec-decoding stats from rank 0's row only and do not change. The acceptance-length gates of the accuracy tests sum every row (acceptance_length_from_iteration_stats), so under attention DP they measured rank 0's acceptance alone and could not see a drafting regression on the other ranks. Found while checking Mistral Large 4 DEP4 acceptance with its DSpark drafter. Each rank now snapshots its draft tokens, accepted tokens and requests with draft tokens (RankSpecDecodingStats) into its payload, and the fanout rebuilds the row's SpecDecodingStats from them, deriving the acceptance length from the rank's own counters. The draft timings stay rank 0's, like the row's other timings. A rank with no snapshot (an explicit zero placeholder) gets no spec-decoding stats rather than rank 0's. The snapshot frame gains a third presence flag, so its version moves from -1 to -2 and its header from four to five fields. Every rank of a job runs the same build, so all ranks write and read the same version; a payload with another version fails the header check instead of being misread. Tests: test_adp_router.py pins the new wire schema and round-trips the snapshot with and without KV snapshots and load fields, keeping zero counters distinct from an absent snapshot. test_iter_stats_populate.py fans out asymmetric work on two ranks (with rank 0 busy or idle, and the padding and CUDA-graph dummies excluded) and checks each row's counts and acceptance length, that the rows add up to the group's acceptance, and that a placeholder rank does not inherit rank 0's drafting. Split out of NVIDIA#21693 (Mistral Large 4 merge-back); originally GitLab 4932efdc24 ("Add native DSpark to StoneWillow Staircase specializations"), rebased onto the versioned snapshot frame of NVIDIA#21131. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
… stats Attention-DP iteration stats fan out one row per rank (NVIDIA#21131): the scheduling and KV fields of each row come from that rank's payload on the ADP rank-state allgather, but specdec_stats was in the list of nested objects copied from rank 0. Every rank's row therefore reported rank 0's draft and accepted token counts and acceptance length, so per-rank acceptance could not be told apart, an idle rank 0 showed zero acceptance for the whole group, and summing the rows multiplied rank 0's numbers by the number of ranks. Any attention-DP deployment with speculative decoding (for example DeepSeek-R1 DEP with MTP) and iteration stats enabled reports these rows, through LLM.get_stats() and the JSON iteration stats of trtllm-serve's /metrics; the Prometheus metrics (/prometheus/metrics) read spec-decoding stats from rank 0's row only and do not change. The acceptance-length gates of the accuracy tests sum every row (acceptance_length_from_iteration_stats), so under attention DP they measured rank 0's acceptance alone and could not see a drafting regression on the other ranks. Found while checking Mistral Large 4 DEP4 acceptance with its DSpark drafter. Each rank now snapshots its draft tokens, accepted tokens and requests with draft tokens (RankSpecDecodingStats) into its payload, and the fanout rebuilds the row's SpecDecodingStats from them, deriving the acceptance length from the rank's own counters. The draft timings stay rank 0's, like the row's other timings. A rank with no snapshot (an explicit zero placeholder) gets no spec-decoding stats rather than rank 0's. The snapshot frame gains a third presence flag, so its version moves from -1 to -2 and its header from four to five fields. Every rank of a job runs the same build, so all ranks write and read the same version; a payload with another version fails the header check instead of being misread. Tests: test_adp_router.py pins the new wire schema and round-trips the snapshot with and without KV snapshots and load fields, keeping zero counters distinct from an absent snapshot. test_iter_stats_populate.py fans out asymmetric work on two ranks (with rank 0 busy or idle, and the padding and CUDA-graph dummies excluded) and checks each row's counts and acceptance length, that the rows add up to the group's acceptance, and that a placeholder rank does not inherit rank 0's drafting. Split out of NVIDIA#21693 (Mistral Large 4 merge-back); originally GitLab 4932efdc24 ("Add native DSpark to StoneWillow Staircase specializations"), rebased onto the versioned snapshot frame of NVIDIA#21131. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
… stats Attention-DP iteration stats fan out one row per rank (NVIDIA#21131): the scheduling and KV fields of each row come from that rank's payload on the ADP rank-state allgather, but specdec_stats was in the list of nested objects copied from rank 0. Every rank's row therefore reported rank 0's draft and accepted token counts and acceptance length, so per-rank acceptance could not be told apart, an idle rank 0 showed zero acceptance for the whole group, and summing the rows multiplied rank 0's numbers by the number of ranks. Any attention-DP deployment with speculative decoding (for example DeepSeek-R1 DEP with MTP) and iteration stats enabled reports these rows, through LLM.get_stats() and the JSON iteration stats of trtllm-serve's /metrics; the Prometheus metrics (/prometheus/metrics) read spec-decoding stats from rank 0's row only and do not change. The acceptance-length gates of the accuracy tests sum every row (acceptance_length_from_iteration_stats), so under attention DP they measured rank 0's acceptance alone and could not see a drafting regression on the other ranks. Found while checking Mistral Large 4 DEP4 acceptance with its DSpark drafter. Each rank now snapshots its draft tokens, accepted tokens and requests with draft tokens (RankSpecDecodingStats) into its payload, and the fanout rebuilds the row's SpecDecodingStats from them, deriving the acceptance length from the rank's own counters. The draft timings stay rank 0's, like the row's other timings. A rank with no snapshot (an explicit zero placeholder) gets no spec-decoding stats rather than rank 0's. The snapshot frame gains a third presence flag, so its version moves from -1 to -2 and its header from four to five fields. Every rank of a job runs the same build, so all ranks write and read the same version; a payload with another version fails the header check instead of being misread. Tests: test_adp_router.py pins the new wire schema and round-trips the snapshot with and without KV snapshots and load fields, keeping zero counters distinct from an absent snapshot. test_iter_stats_populate.py fans out asymmetric work on two ranks (with rank 0 busy or idle, and the padding and CUDA-graph dummies excluded) and checks each row's counts and acceptance length, that the rows add up to the group's acceptance, and that a placeholder rank does not inherit rank 0's drafting. Split out of NVIDIA#21693 (Mistral Large 4 merge-back); originally GitLab 4932efdc24 ("Add native DSpark to StoneWillow Staircase specializations"), rebased onto the versioned snapshot frame of NVIDIA#21131. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Summary
Prometheus label-schema change: 28 existing KV-cache and rank-local scheduling metric families gain
dp_rank, includingdp_rank="0"in single-rank serving. This fixes attention-DP rows copying rank 0's KV measurements, separates cumulative reuse/miss histories, and counts shared speculative-token increments once.Fixes #21124. Tracking: TRTLLM-17359.
Scope
RankKvCacheStats, with framing constants onRankIterStatsPayload.mainateebd9d0fcf, preserving the new collector factory, custom registries, multiprocess gauges, and KV-balanced admission; extend the existing OpenEngine multiprocess scrape test to verify distinct rank gauges and reuse totals.Dashboard migration: retain
dp_rankto inspect individual ranks or aggregate it explicitly. Sum additive counts where appropriate; use maximum utilization for the fullest rank, and do not sum utilization or hit-rate ratios. Documentation explains ownership, freshness, and delivery assumptions.Verification
Published, tested, and reviewed head:
bc01999755c604727016f1d1ac561519e933a261.338 passed, 0 failures, 0 errors, 0 skipped in 112.08 seconds on Linux/Python 3.12.3, NVIDIA PyTorch 2.14.0a0 (26.09 runtime, CUDA 13.4), H200 NVL. All six complete modules ran with repository conftest, real native bindings, and the actual OpenEngine HTTP/multiprocess tests using synthetic engine inputs.
Command:
Both rebase conflicts were import lists. The original sign-off and single-commit history were preserved; range comparison shows the import adaptation and focused OpenEngine regression coverage.
Required changed-file and signed-amendment hooks passed. Two fresh complete eight-file reviews found no further actionable issues after fixing the test readiness race; the test waits for all four expected rank gauges/counters in the same scrape.
Verification loaded the final Python sources via the official precompiled editable install. Artifact #67275's 11,723 native/build inputs matched the final source, including canonical LFS contents; wheel source provenance and all 14 source/test hashes were verified. Tracked source stayed clean. The initial import attempt needed the current source's declared gRPC health-checking dependency in the older runtime; after installation, the entire six-file command passed.
Historical head
41e88279: 294 native tests passed with no failures/errors/skips. Full CI #67015 reported 92,960 passed, zero failed tests, and 14,288 skipped, but the pipeline failed with L40S infrastructure/maintenance and an aborted H100 performance stage. This was a coverage gap, not evidence of a product failure; no source fix was made for those infrastructure failures.Earlier optional live check at
bf54a05: nine Qwen3-0.6B BF16/TP4/ADP/KV-v2 requests completed on four H100 NVL GPUs; 1,667 uniquely matched native/exported rows agreed across nine KV fields, and eight real-collector reuse/miss replay checks passed. The overall check remains inconclusive due to duplicate startup identities and an unobserved retained-cache/null-snapshot idle case. No current-head live-inference pass is claimed.Full CI was requested for this head with
/bot run --disable-fail-fast. Actions #41021946477 passed authorization, vulnerability scan, and dispatch. Bot acknowledgment identifies helper #81641; bounded Jenkins metadata confirms downstream L0_MergeRequest_PR #67304 is running forbc019997as of October 30, 07:57 UTC. Full CI results are pending.Repository PR checklist reviewed; tests and documentation updated. No dependency, ownership, or protected LLM API changes.
Notes
Release-owner note: include the
dp_ranklabel-schema change and dashboard migration in the release notes. Compact and versioned frames are internal same-build formats; this change makes no cross-version or rolling-upgrade compatibility guarantee.