Fix: stabilize MTP logits validation

[Little-oil]

Problem

MTP compared device logits (projected from device hidden_out) with CPU logits (projected from the independent CPU hidden_out). The permitted hidden-state error can be amplified by the LM-head projection, producing intermittent logits-only failures even when hidden_out, pre_hc_hidden_out, and kv_cache pass. Input generation was also unseeded.

Fix

• Validate LM-head output against a reference projected from the actual device hidden_out, preserving the strict 1% outlier budget.
• Retain an end-to-end max-absolute-difference guard (0.5) against the original CPU golden to catch catastrophic upstream errors.
• Reject non-finite conditional reference logits explicitly.
• Seed MTP input generation by default and expose --seed.
• Make the LM-head CPU reference use the runtime owner count and only one copy of each TP weight shard.
• Add focused tests for projection order, EP2/4/8 owner counts, NaN/Inf and corruption handling, outlier budget, and the end-to-end guard.

Validation

• pytest -q tests/golden: 180 passed
• Ruff, py_compile, and git diff --check: passed
• Current upstream/toolchain (pypto-lib c0751b7, PyPTO 6730e96, runtime dccb837), even cards 0,2:
  • task_20260803_030315_106269014303: kv_cache, hidden_out, pre_hc_hidden_out, and logits all passed
• Historical CI-era toolchain, even cards only:
  • seed 0: task_20260803_021409_220873612205 (all four outputs passed)
  • seed 1: task_20260803_022257_28353452942 (all four outputs passed)
  • seed 3: task_20260803_022301_284110020704 (all four outputs passed)
  • swimlane seed 0: task_20260803_023142_3476347660 (all four outputs passed)
• standalone lm_head: task_20260803_022857_321955919156 (logits and sampled_ids passed)

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 18b332d7-9570-4299-a3d4-2d7d152745bb

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

DeepSeek-V4 Flash now uses shared LM-head validation helpers. The validation checks tensor-parallel projections, finite values, ratio limits, and maximum end-to-end differences. Prefill execution also supports deterministic seeds.

Changes

DeepSeek-V4 LM-head validation

Layer / File(s)	Summary
Shared LM-head validation helpers models/deepseek/v4-flash/lm_head_validation.py	Adds reference-logit projection across tensor-parallel shards and device-aware validation with ratio and end-to-end absolute-difference checks.
Golden and prefill integration models/deepseek/v4-flash/lm_head.py, models/deepseek/v4-flash/prefill_mtp.py	Routes golden logits through the shared helper. Replaces ratio-only comparison and adds deterministic seed handling.
Projection and comparator tests tests/golden/test_lm_head_validation.py	Tests hidden-state usage, corruption, NaN values, catastrophic errors, padding, outlier limits, and runtime owner counts.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Sequence Diagram(s)

sequenceDiagram
  participant PrefillMTP
  participant GoldenLMHead
  participant LMHeadValidation
  participant TPWeights
  PrefillMTP->>GoldenLMHead: generate reference logits
  GoldenLMHead->>LMHeadValidation: compute_lm_head_logits(hidden states, row indices, TP size)
  LMHeadValidation->>TPWeights: read vocabulary shards
  TPWeights-->>LMHeadValidation: sharded projection weights
  LMHeadValidation-->>GoldenLMHead: reference logits
  PrefillMTP->>LMHeadValidation: compare device logits and hidden outputs
  LMHeadValidation-->>PrefillMTP: validation result

Loading

Possibly related PRs

• hw-native-sys/pypto-lib#835: Updates lm_head.py to use validation helpers with related row and shard semantics.
• hw-native-sys/pypto-lib#864: Concerns DeepSeek-V4 Flash LM-head handling in shared callers.

Poem

A rabbit checks logits row by row,
Across the shards, the values flow.
NaNs are caught and errors shown,
Seeds make each test’s path known.
“Hop!” says validation, crisp and bright—
The LM head now checks its flight.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 14.29% 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 clearly summarizes the primary change: stabilizing MTP logits validation.
Description check	✅ Passed	The description directly explains the validation fix, deterministic seeding, reference updates, tests, and validation results.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@models/deepseek/v4-flash/lm_head_validation.py`:
- Around line 64-79: Validate the LM-head reference produced by
compute_lm_head_logits before calling base_cmp, rejecting it when it contains
any NaN or infinite values; preserve the existing comparison flow for finite
references. Add a regression test using a selected NaN hidden_out value with
matching finite actual and _expected logits to ensure the comparison fails.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 3acde632-8d7e-490c-9112-e75b234a33bb

📥 Commits

Reviewing files that changed from the base of the PR and between c0751b7 and 6266f31.

📒 Files selected for processing (4)

• models/deepseek/v4-flash/lm_head.py
• models/deepseek/v4-flash/lm_head_validation.py
• models/deepseek/v4-flash/prefill_mtp.py
• tests/golden/test_lm_head_validation.py

[Little-oil]

CI retry triage for empty commit 437b3ed:

• Run 30806297190 did not reproduce the logits precision mismatch.
• a2a3 job 91665122786: prefill_mtp.py reached runtime, then failed with S1:running-stalled / AICPU 507018 / runtime -100; tensor copy-back and all output comparisons were skipped.
• a2a3sim job 91665122784 and a5sim job 91665122715: the preceding examples passed, then both stalled in unchanged decode_mtp.py runtime and hit the 30-minute job limit. Neither reached precision validation.
• unit-tests, pre-commit, detect-changes, and serving-deepseek passed.

The first attempt also passed the hardware a2a3 and serving-deepseek jobs, and the current upstream base reproduces the simulator runtime stall locally. The remaining red checks are runtime/simulator stalls rather than an output-validation regression.