[https://nvbugs/6707518][fix] Skip logprobs for the SA spec run in the Kimi-K3 SA harness - #18618
[https://nvbugs/6707518][fix] Skip logprobs for the SA spec run in the Kimi-K3 SA harness#18618trtllm-agent wants to merge 1 commit into
Conversation
Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
|
[Repair Bot][Two-Leg Repro Comparison:6707518-30316da103a0-1788386971169698926] Reproduction comparison: |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review. WalkthroughThe integration harness disables logprob collection for SA speculative runs. Baseline and other runs continue collecting logprobs. Logits parity uses the existing one-sided comparison when SA output lacks aligned logprobs. ChangesSA logprob handling
Estimated code review effort: 2 (Simple) | ~10 minutes Merge Risk: ⚪ Minimal · up to This localized test-harness change avoids an unsupported SA logprobs request while preserving baseline coverage and other speculative modes; no actionable merge-blocking risk remains beyond normal checks and review. 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Full details: Description checkExplanation The description clearly explains the root cause, implementation, affected test, validation, regression checks, and bug link. It does not use the template headings exactly and omits the full PR checklist, but it contains the critical information and is mostly complete.
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
/bot run |
|
PR_Github #71109 [ run ] triggered by Bot. Commit: |
fredricz-20070104
left a comment
fredricz-20070104
left a comment
There was a problem hiding this comment.
Review summary - CONCERNS
Verdict: The one-line guard is the right shape for the described root cause and the change is correctly scoped to the harness, but it is applied inconsistently — line 509 uses the new spec_want_logprobs while line 512 still forwards the old want_logprobs — and it silently turns the sa case of a *_logits_parity test into a check that can no longer fail. I would fix the line 512 skew and add a signal for the degraded parity before merging; mergeable_state is reported as blocked in any case, so CI/approval gates still have to clear.
Issues
- [MAJOR]
tests/integration/defs/kimi_k3_sa_harness.py:512-_dump_completionsstill told logprobs exist after they were suppressed - [MAJOR]
tests/integration/defs/kimi_k3_sa_harness.py:508-salogits-parity check silently degrades to one-sided with no signal - [MINOR]
tests/integration/defs/kimi_k3_sa_harness.py:508- guard keyed on the bare literal"sa"rather than the capability - [NIT]
tests/integration/defs/kimi_k3_sa_harness.py:502- six-line comment for a one-line condition
QA view
- Test coverage: partial - the only changed file is the harness itself. No new test case, no test-list entry. Still uncovered: nothing asserts the
saspec request is now admitted (a re-regression would look like the same opaquesanity harness reported FAIL); theoutput_json+want_logprobs=True+spec_mode="sa"combination at line 512 is untested; the one-sided parity branch that thesacase now always takes is not verified to be able to fail;test_kimi_k3_specdec.pyis not in the diff, so its markers and assertions are unknown to me. - SM coverage: architecture-independent - no
get_sm_version(), nosm_versioncomparison, no__CUDA_ARCH__, no arch-guarded kernel, no fp8/nvfp4 path. Pure Python boolean in a test harness. The PR body's claim of verification on the original GPU type is not something the diff can show. - Test code: argument skew between lines 509 and 512; the
saparity assertion can no longer fail on the property the test is named for; no warning/marker/xfail recording the degradation, so a degraded run is indistinguishable from a verified one in CI logs; this is effectively a capability waive inside the harness with no follow-up tracker referenced; magic string"sa"as a mode discriminator; the comment asserts behaviour of_compare_logits_parity, which is not in this diff and can rot unnoticed. - Test time: none - one boolean plus six comment lines; no new cases, parametrisations, timeouts, or larger models. Skipping logprob collection may make the
sarun marginally cheaper, but the diff cannot quantify that. - Needs
/qa-verify: yes - a bug fix whose test does not reproduce or assert the original failure, a change to test infrastructure itself, an effective waive that degrades the parity check, plus an untestedoutput_jsonpath. Please run the harness with logprobs enabled and--output-jsonunderspec_mode=sa, confirm admission succeeds and the JSON is well formed, and confirm one non-saspec mode still collects logprobs.
Does this actually fix nvbugs/6707518?
Partial. NVBugs resolved the id, but the synopsis is truncated ([TensorRT-LLM][L0][Post-Merge][main] Test failed: test_kimi_k3_specdec.py::test_kimi_k3_sa_specde...) and says only that the test failed — it does not state that the failure was the logprobs admission RequestError. The PR's own reproduction block reports the generic signature AssertionError: sanity harness reported FAIL with Signature relation: needs_agent, so the signature was not matched mechanically either. What the diff does is consistent with the stated root cause: spec_want_logprobs = want_logprobs and spec_mode != "sa" stops the spec-side _generate from requesting logprobs for exactly the mode the failing test exercises, which would remove the admission rejection. It is partial because (a) line 512 still forwards want_logprobs to _dump_completions, so the post-generation path can still fail or emit an inconsistent artifact on the same configuration, and (b) the guard matches only the exact literal "sa", so any sibling one-engine mode with a tokens-only sample state still hits the identical error.
Possible new issues
- With
want_logprobs=True,spec_mode="sa"andoutput_jsonset,_dump_completionsis told to serialise logprobs that_generatenever produced: either an exception (surfacing as the same harness FAIL this PR targets) or a JSON file whoselogprobsentries are empty while claiming to contain them, so downstream comparison silently agrees with nothing. - An SA-sampler logits regression (e.g. wrong draft acceptance changing the distribution) now lands in a one-sided comparison and can report PASS, with no log line indicating the two-sided check was skipped.
- A future spec mode with the same sampler limitation, or a differently spelled/suffixed
savariant, reproduces the original admissionRequestErrorand will look like a new bug. - The harness now routes around the hard admission check introduced by commit 22e4cbc rather than exercising it, so later relaxation or extension of that rule will not be surfaced by this test.
What I could not verify
_generate, _dump_completions, _compare_logits_parity, the SpecSampler admission check, and test_kimi_k3_specdec.py itself are all outside this diff, so I am reasoning about the line 512 skew and the one-sided-parity claim from the call sites and the comment only — if _dump_completions tolerates a True flag over logprob-less completions, that finding drops to a NIT. I also cannot confirm the PR body's claims that the fix was verified on the original GPU type or that related tests were checked for regressions, nor the runtime status of any CI job referenced outside the diff.
Findings not attached to a line
- [MINOR]
tests/integration/defs/kimi_k3_sa_harness.py:508- Guard keyed on the bare literal "sa" rather than the capability- The condition hardcodes the single string
"sa". The underlying constraint is not 'the mode is named sa' but 'this sampler's sample state carries tokens only, so logprobs are rejected at admission'. Any future or existing sibling mode with the same one-engine SpecSampler (or a differently-cased / suffixed spelling such as"SA"or"sa_..."reaching this code path) will re-hit the exact RequestError this PR is fixing, and the failure will look like a brand-new bug rather than a known limitation. If the harness has a set of spec modes it recognises, prefer a named constant or a small membership set so the next mode with the same limitation is a one-line addition:
- The condition hardcodes the single string
# spec modes whose sample state carries tokens only
_MODES_WITHOUT_LOGPROBS = {"sa"}
spec_want_logprobs = want_logprobs and spec_mode not in _MODES_WITHOUT_LOGPROBSNon-blocking; the current form is correct for the mode in the failing test.
Automated review by NVCortex Lite, run by @fredricz-20070104.
| completions = _generate(llm, prompt_texts, max_tokens, spec_want_logprobs) | ||
| llm.shutdown() | ||
| if output_json: | ||
| _dump_completions(output_json, completions, want_logprobs) |
There was a problem hiding this comment.
[MAJOR] _dump_completions still told logprobs exist after they were suppressed
completions on line 509 is now produced with spec_want_logprobs, but line 512 still forwards the original want_logprobs to _dump_completions. Concrete trigger: invoke the harness with logprobs enabled (want_logprobs=True), spec_mode="sa", and --output-json set (output_json truthy). _generate was called with False, so every completion carries no logprob field, while _dump_completions is told to serialise logprobs. Depending on how it reads them, the outcome is either an AttributeError/KeyError/TypeError on the absent field — which the caller surfaces as the same 'sanity harness reported FAIL' assertion this PR is trying to fix — or a JSON artifact whose logprobs entries are silently null/empty while claiming to contain them, so any downstream parity/diff tooling reading that file compares against nothing and reports spurious agreement. This is the same inconsistency the fix removed one line earlier, left in place one line later.
| _dump_completions(output_json, completions, want_logprobs) | |
| _dump_completions(output_json, completions, spec_want_logprobs) |
If _dump_completions deliberately needs the requested flag rather than the available one, that asymmetry deserves an explicit comment, because the two arguments now diverge for exactly the sa path this PR adds.
| # than silently dropping them, so this run must not ask for them at all. | ||
| # _compare_logits_parity already handles a spec run without logprobs via | ||
| # one-sided parity. | ||
| spec_want_logprobs = want_logprobs and spec_mode != "sa" |
There was a problem hiding this comment.
[MAJOR] sa logits-parity check silently degrades to one-sided with no signal
For spec_mode == "sa" the spec run no longer produces logprobs, so per the comment _compare_logits_parity falls back to 'one-sided parity'. The concrete consequence: test_kimi_k3_sa_specdec_logits_parity — a test whose entire purpose is detecting logits divergence between the baseline and the SA speculative path — now runs its primary comparison with only one side's logprobs available. A real SA-sampler logits regression (wrong draft acceptance producing different distributions) would land in a test that still reports PASS, i.e. the assertion for the property under test can no longer fail on the sa configuration. Nothing in this diff records that degradation: no logger.warning, no pytest.xfail/skip, no marker, so from CI output the SA case is indistinguishable from a fully-verified one.
Suggested minimum: emit an explicit warning (or have the harness print a parity_mode: one-sided field into the summary/JSON) when want_logprobs and spec_mode == "sa", so a reader of the log knows the two-sided check was not performed:
spec_want_logprobs = want_logprobs and spec_mode != "sa"
if want_logprobs and not spec_want_logprobs:
print("[harness] spec_mode=sa cannot return logprobs; "
"logits parity degrades to one-sided", flush=True)If one-sided parity genuinely cannot detect regressions, then skipping/xfailing the sa parity case is more honest than a green result, and the follow-up should be tracked.
|
|
||
| llm = _build_llm(ckpt, tp, spec_mode, adp) | ||
| completions = _generate(llm, prompt_texts, max_tokens, want_logprobs) | ||
| # The baseline above keeps its logprobs (logits parity hard-requires an |
There was a problem hiding this comment.
[NIT] Six-line comment for a one-line condition
The rationale block (lines 502-507) is six lines of prose in front of a single boolean expression, and it embeds a claim about another function's behaviour ('_compare_logits_parity already handles a spec run without logprobs via one-sided parity') that will silently rot if that function changes. Two lines are enough to carry the non-obvious part — that sa rejects logprobs at admission instead of dropping them — and the parity-degradation note is better placed as the runtime signal requested in the finding on line 508, where it cannot go stale unnoticed.
fredricz-20070104
left a comment
fredricz-20070104
left a comment
There was a problem hiding this comment.
Review summary - Approve (non-blocking)
Approving so this is not blocked on me. The points raised in my review comment above are non-blocking — please read them and address what you agree with before merging.
Worth doing before this is relied on: Multiple triggers at once: (1) it is a bug fix whose test does not reproduce or assert the original failure mode — nothing checks that the sa spec request is now admitted, so a green run does not prove the RequestError is gone; (2) it is a change to test infrastructure itself (the harness), so the safety net is what changed; (3) it is effectively a capability waive that silently degrades the sa logits-parity check to one-sided, and a human should confirm the test still has value in that mode rather than passing vacuously; (4) the output_json + want_logprobs + spec_mode="sa" combination at line 512 is an untested path that may still fail; (5) the PR's own reproduction block reports the generic signature 'AssertionError: sanity harness reported FAIL' with 'Signature relation: needs_agent', so the mechanical signature match did not confirm the root cause. QA should run the harness with logprobs enabled and --output-json under spec_mode=sa, confirm admission succeeds, confirm the JSON is well-formed, and confirm at least one non-sa spec mode still collects logprobs.
Automated review by NVCortex Lite, run by @fredricz-20070104.
|
PR_Github #71109 [ run ] completed with state
|
|
Duplicate of #18682, which fixes a few minor issues of this one, so closing |
6 participants
Summary
kimi_k3_sa_harness.pylogits-parity harness asks forlogprobs=5on both the baseline run and the speculative-decoding run. Under the one-enginesaspec mode the sampler's sample state carries tokens only, and commit 22e4cbc changed that combination from a warning that silently droppedlogprobsinto a hardRequestErrorat request admission. As a result the spec-side generation now aborts during admission and the test fails before any logits comparison happens.spec_want_logprobs = want_logprobs and spec_mode != "sa", so only thesapath stops asking for logprobs while every other spec mode and the baseline keep theirs. The baseline must retain logprobs because logits parity requires an aligned reference, and_compare_logits_parityalready degrades to one-sided parity when the spec run has no logprobs — so dropping the unsupported request is sufficient and keeps the test meaningful without relaxing the new admission check.pytest tests/integration/defs/test_kimi_k3_specdec.py::test_kimi_k3_sa_specdec_logits_parity -vTest plan
Links
Reproduction comparison
Signature: AssertionError: sanity harness reported FAIL
Signature: AssertionError: sanity harness reported FAIL
Dev Engineer Review
QA Engineer Review
test_kimi_k3_sa_specdec_logits_parity.