[II] Preserve the final Kimi-K3 AttnRes block

[voipmonitor]

Behavior

Kimi-K3 AttnRes output reuse now stops at the boundary that commits the final
residual block. The old prefix remains in the residual bank, while the
post-attention mixture receives separate output storage. Earlier block
boundaries and the remaining KDA, MLA, dense-MLP, and routed-MoE caller-output
paths retain their existing storage reuse.

Status: implemented and qualified.

Technical reason

The official Kimi-K3 text transformer has 93 layers, an AttnRes block size of
12, and eight residual-bank slots. The eighth slot can hold temporary output
before layer 84 because it is not yet a residual source. Layer 84 commits that
slot. Reusing the same storage for the post-attention output then overwrites a
residual source consumed by layers 85 through 92 and the final AttnRes mixture.

The fix identifies the layer that commits the final block and allocates one
separate output tensor at that boundary. A 4,096-token BF16 tensor with hidden
width 7,168 occupies 58,720,256 bytes. No persistent workspace is added.

Compatibility

Models without AttnRes and platforms without AttnRes output reuse are
unchanged. Tensor values, kernels, operation order, cache layout, KDA, MLA,
MoE, DCP, and decode collectives are unchanged. The only additional allocation
occurs when the final AttnRes block becomes live.

Validation

• Two focused unit tests pass and verify ordinary block reuse plus preservation
  of the final committed block.
• A 93-layer lifetime harness at shape [4096, 7168] reproduces the defect:
  unrestricted reuse changes 7,266,461 of 29,360,128 final-output elements.
• Preserving the final block is bitwise identical to separately allocated
  output for all 29,360,128 output elements and all 234,881,024 residual-bank
  elements in three complete 93-layer repetitions.
• Official Kimi-K3 MXFP4, TP16/DCP16, FP8 KV, and 4,096-token scheduler chunks
  complete a captured 208,026-token prompt plus 8,192 greedy output tokens
  without the deterministic repetition failure.
• The same runtime completes a 524,288-token prompt and retains 1,058,823
  physical target-KV tokens.
• No-speculation decode for 256 input and 512 output tokens measures 55.799
  tok/s median across three runs; the source-locked comparison measured 55.769
  tok/s.
• ruff check, ruff format --check, git diff --check, and Python bytecode
  validation pass for the changed files.

Qualification artifacts are stored under
/mnt/luke/kimi-k3-runs/merge-lse-alias-fix-20260822/full-r31-attnres-final-block-fix-1m/
on the 16-GPU qualification host.

Summary by CodeRabbit

• Bug Fixes
  • Fixed residual state handling at the final attention block boundary.
  • Ensured committed residual blocks are preserved while post-attention normalization completes.
  • Added coverage to verify correct output allocation and state propagation.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 89e49043-5e59-4cc5-b5aa-5dc4ec889c70

📥 Commits

Reviewing files that changed from the base of the PR and between b5f995e and d19ef45.

📒 Files selected for processing (2)

• tests/models/kimi_k3/test_eagle3.py
• vllm/models/kimi_k3/nvidia/model.py

Included review availability: Your plan provides up to 1 included review per hour; 0 remain after this review.

---

📝 Walkthrough

Walkthrough

The Kimi decoder now identifies the final AttnRes block-write layer and avoids reusing prefix_sum as output storage at that boundary. Tests cover the new layer attribute and verify preservation of the committed block.

Changes

Kimi AttnRes preservation

Layer / File(s)	Summary
Track final block boundary vllm/models/kimi_k3/nvidia/model.py, tests/models/kimi_k3/test_eagle3.py	KimiDecoderLayer records the final block-write status. The test fixture initializes the new attribute.
Preserve final committed block vllm/models/kimi_k3/nvidia/model.py, tests/models/kimi_k3/test_eagle3.py	The final boundary passes output=None to attn_res. The test verifies allocated output usage, attention-prefix propagation, and block preservation.

Estimated code review effort: 2 (Simple) | ~10 minutes

Merge Risk: ⚪ Minimal · up to d19ef

The change preserves the final AttnRes residual block with a localized additional allocation while retaining existing behavior for other paths. No actionable merge-blocking risk remains after normal checks and review.

Suggested reviewers: zjy0516, lukealonso

🚥 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%. Docstring coverage is scoped to functions touched by this diff. Analyzed 6 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main change: preserving the final Kimi-K3 AttnRes block.
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.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.}