[https://nvbugs/6599150][fix] Initialize dt_bias in KDA verify-parity test by brnguyen2 · Pull Request #17584 · NVIDIA/TensorRT-LLM

brnguyen2 · 2026-08-12T21:59:45Z

Dev Engineer Review

Initialized runtime.mixer.dt_bias with finite random values before verification.
Prevented uninitialized torch.empty data from causing intermittent NaN failures.
Removed the obsolete waiver for test_kda_verify_matches_sequential_decode[2-1].
Preserved the existing mixer construction and public API.
No performance or error-handling regressions identified.

QA Engineer Review

Modified test_kda_verify_matches_sequential_decode.
The test is covered by the existing test configuration.
Removed its waiver entry from tests/integration/test_lists/waives.txt.
All four parametrizations passed in a single-process run.
The test passed 200/200 fresh-process runs after the fix.
Verdict: sufficient.

Description

test_kda_verify_matches_sequential_decode failed intermittently in the DGX_B200 PyTorch CI stage with NaN outputs (nvbug 6599150), while SBSA CI on the same commits passed.

Root cause: the test builds KimiKDARuntime with random weights and no checkpoint, but KimiKDAMixer.dt_bias is constructed with torch.empty and only filled by load_weights(), so in this test it held uninitialized host memory (torch.manual_seed does not affect torch.empty). When that memory happened to contain NaN/Inf bit patterns, both the decode-reference and verify paths produced identical NaN outputs, and torch.testing.assert_close (default equal_nan=False) reported a 100% mismatch. Finite garbage preserved parity because both paths share the parameter — hence the nondeterministic, platform-dependent flake rather than a real kernel divergence.

The fix initializes dt_bias explicitly after constructing the runtime, as the sibling KDA parity tests (test_kimi_kda_fused_verify_parity, test_kda_mtp_decode_cute_parity) already do. Also removes the corresponding waives.txt entry for test_kda_verify_matches_sequential_decode[2-1], since the flake is fixed.

The torch.empty construction in the mixer itself is left unchanged; it matches the convention for checkpoint-loaded parameters (e.g. mamba2_mixer).

Evidence for the root cause: across 240 instrumented fresh-process test runs on an x86_64 B200 node, NaN/Inf content in dt_bias at construction time correlated perfectly with test failure (78/78 failures, 162/162 passes, zero off-diagonal cases).

Test Coverage

tests/unittest/_torch/modeling/test_kimi_kda_verify_parity.py (the fixed test): 200/200 fresh-process runs of the [2-1]/[2-3] parametrizations on an x86_64 B200 node after the fix; before the fix ~64% of fresh processes failed on the same node. All four parametrizations pass in a single-process run.

PR Checklist

PR description clearly explains what and why.
PR follows TRT-LLM coding guidelines.
Test cases are provided (test-only change).
No API changes.
No new dependencies.

… test The test builds KimiKDARuntime with random weights and no checkpoint, but dt_bias is constructed with torch.empty and only filled by load_weights(), so it held uninitialized host memory. When that memory contained NaN/Inf bit patterns, both the decode and verify paths produced identical NaN outputs and assert_close (equal_nan=False by default) reported a 100% mismatch — the intermittent DGX_B200 CI failure. Finite garbage preserved parity because both paths share the parameter, which is why the failure was nondeterministic and platform-dependent. Initialize dt_bias explicitly after construction, as the sibling KDA parity tests already do. Verified: 200/200 fresh-process runs of the [2-1]/[2-3] parametrizations on an x86_64 B200 node (previously ~64% of processes failed), plus a full-file run of all four parametrizations. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>

coderabbitai · 2026-08-12T22:03:16Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 8199d321-d0c0-4477-b832-4f02bc6fdb8c

📥 Commits

Reviewing files that changed from the base of the PR and between ec3e1a1 and 944c4c0.

📒 Files selected for processing (2)

tests/integration/test_lists/waives.txt
tests/unittest/_torch/modeling/test_kimi_kda_verify_parity.py

💤 Files with no reviewable changes (1)

tests/integration/test_lists/waives.txt

🚧 Files skipped from review as they are similar to previous changes (1)

tests/unittest/_torch/modeling/test_kimi_kda_verify_parity.py

Walkthrough

The parity test initializes runtime.mixer.dt_bias with finite random values before verification. The waiver for the affected parity case is removed.

Changes

Kimi KDA parity test

Layer / File(s)	Summary
Initialize runtime mixer bias and enable the test case `tests/unittest/_torch/modeling/test_kimi_kda_verify_parity.py`, `tests/integration/test_lists/waives.txt`	The test initializes `runtime.mixer.dt_bias` with a normal distribution. The waiver for `test_kda_verify_matches_sequential_decode[2-1]` is removed.

Estimated code review effort: 1 (Trivial) | ~2 minutes

Mergeability Score: ⚪ Minimal · up to 944c4

This localized test-only change initializes a previously uninitialized test parameter and removes the waived test entry; no actionable merge-blocking risk remains beyond normal checks and review.

Possibly related PRs

NVIDIA/TensorRT-LLM#17444: Removes a waiver for the same KDA parity test and updates its initialization.
NVIDIA/TensorRT-LLM#17561: Addresses initialization of dt_bias for the same parity test.

Suggested reviewers: bowenfu

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title follows the required ticket and type format and clearly describes the dt_bias initialization fix.
Description check	✅ Passed	The description explains the root cause, solution, test coverage, and relevant checklist items in sufficient detail.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

pranav-nvidia

LGTM!

brnguyen2 · 2026-08-12T22:51:29Z

/bot run

tensorrt-cicd · 2026-08-12T22:57:59Z

PR_Github #65705 [ run ] triggered by Bot. Commit: 5a31fb7 Link to invocation

…ty-dt-bias-init Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>

The flake is fixed by initializing dt_bias in the test; remove the DGX_B200 waive for test_kda_verify_matches_sequential_decode[2-1]. Verified after rebasing onto current main: 40/40 fresh-process runs of the [2-1]/[2-3] parametrizations on an x86_64 B200 node, plus all four parametrizations passing in a single-process run. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>

coderabbitai · 2026-08-12T23:05:18Z

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

tensorrt-cicd · 2026-08-13T00:42:48Z

PR_Github #65705 [ run ] completed with state SUCCESS. Commit: 5a31fb7
/LLM/main/L0_MergeRequest_PR pipeline #53422 completed with status: 'SUCCESS'

CI Report

Link to invocation

brnguyen2 · 2026-08-13T01:54:35Z

/bot reuse-pipeline

tensorrt-cicd · 2026-08-13T02:00:01Z

PR_Github #65744 [ reuse-pipeline ] triggered by Bot. Commit: 21dde19 Link to invocation

tensorrt-cicd · 2026-08-13T02:12:25Z

PR_Github #65744 [ reuse-pipeline ] completed with state SUCCESS. Commit: 21dde19
Reusing PR_Github #65705 for commit 21dde19

Link to invocation

…DIA#17584) test_kda_verify_matches_sequential_decode[2-1] was waived under nvbugs/6599150, but that bug is closed-fixed: PR NVIDIA#17584 (merged) fixed the test by initializing dt_bias in the parity test and removed the need for the waive. Main carries no such waive. Drop the re-added line so the now-passing test runs. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>

brnguyen2 requested a review from a team as a code owner August 12, 2026 21:59

brnguyen2 requested a review from pranav-nvidia August 12, 2026 21:59

github-actions Bot assigned brnguyen2 Aug 12, 2026

pranav-nvidia approved these changes Aug 12, 2026

View reviewed changes

brnguyen2 enabled auto-merge (squash) August 12, 2026 22:51

brnguyen2 added 2 commits August 12, 2026 16:04

Merge remote-tracking branch 'upstream/main' into fix/kda-verify-pari…

37e6857

…ty-dt-bias-init Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>

BowenFu approved these changes Aug 13, 2026

View reviewed changes

Merge branch 'main' into fix/kda-verify-parity-dt-bias-init

21dde19

brnguyen2 merged commit 564dab1 into NVIDIA:main Aug 13, 2026
7 checks passed

Navigation Menu

Uh oh!

[https://nvbugs/6599150][fix] Initialize dt_bias in KDA verify-parity test - #17584

[https://nvbugs/6599150][fix] Initialize dt_bias in KDA verify-parity test#17584
brnguyen2 merged 4 commits into
NVIDIA:mainfrom
brnguyen2:fix/kda-verify-parity-dt-bias-init

brnguyen2 commented Aug 12, 2026 •

edited by coderabbitai Bot

Loading

Uh oh!

coderabbitai Bot commented Aug 12, 2026 •

edited

Loading

❌ Failed checks (1 warning)

Uh oh!

pranav-nvidia left a comment

Uh oh!

brnguyen2 commented Aug 12, 2026

Uh oh!

tensorrt-cicd commented Aug 12, 2026

Uh oh!

coderabbitai Bot commented Aug 12, 2026

Uh oh!

tensorrt-cicd commented Aug 13, 2026

Uh oh!

brnguyen2 commented Aug 13, 2026

Uh oh!

tensorrt-cicd commented Aug 13, 2026

Uh oh!

tensorrt-cicd commented Aug 13, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

brnguyen2 commented Aug 12, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Dev Engineer Review

QA Engineer Review

Description

Test Coverage

PR Checklist

Uh oh!

coderabbitai Bot commented Aug 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

❌ Failed checks (1 warning)

Uh oh!

pranav-nvidia left a comment

Choose a reason for hiding this comment

Uh oh!

brnguyen2 commented Aug 12, 2026

Uh oh!

tensorrt-cicd commented Aug 12, 2026

Uh oh!

coderabbitai Bot commented Aug 12, 2026

Uh oh!

tensorrt-cicd commented Aug 13, 2026

Uh oh!

brnguyen2 commented Aug 13, 2026

Uh oh!

tensorrt-cicd commented Aug 13, 2026

Uh oh!

tensorrt-cicd commented Aug 13, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

brnguyen2 commented Aug 12, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Aug 12, 2026 •

edited

Loading