fix(mla): declare compressed decode row capacity

[MadeBy561]

Summary

This fixes compressed-MLA DSpark serving above r27's previously qualified C24
envelope for both K5 and K7. At MNS64, target verification can expose 384 rows
with K5 or 512 rows with K7, while the planner previously had a fixed 256-row
capacity. Capturing the larger workload could therefore overwrite graph-owned
split controls and corrupt later FULL graph replays, producing repetitive/BOS
output instead of a valid completion.

• accept an optional integration-declared compressed-MLA decode-row capacity
• preserve the existing planner result through the legacy 256-row envelope
• extend the fixed split/scratch contract only to the declared capacity and reject invalid values

DSpark verifies 1 + draft_tokens rows per request. The integration supplies
the maximum reachable row count; B12X owns the resulting planner and scratch
contract without changing legacy behavior at or below 256 rows.

Testing

• ruff check and ruff format --check pass for all three changed files
• planner/boundary/legacy-identity subset: 77 passed
• CUDA graph capture/replay against the compressed-MLA reference at 384 (MNS64/K5), 512 (MNS64/K7), and 768 (MNS128/K5) rows: 3 passed (max_abs <= 0.10, cosine >= 0.9995)
• exact r27 TP4/MNS64/K5 integration: Estonia C64 64/64, C1 260.4 tok/s, C64 2724.4 tok/s, and 8k-128k prefill 11.2-13.2k tok/s
• same-image TP4/MNS64/K7 integration smoke: Estonia C64 64/64, C1 250.6 tok/s, and observed code emission up to 499.4 tok/s

No draft-depth default, reasoning template, kernel, or sampling behavior changes;
K5 remains the release default unless the separate K7 qualification succeeds.

[coderabbitai]

📝 Walkthrough

Walkthrough

Compressed MLA split planning now accepts an optional decode row capacity. The value is validated, propagated through scratch planning, applied to decode split selection, and covered by parameterized and CUDA-graph tests.

Changes

Compressed MLA decode capacity

Layer / File(s)	Summary
Decode split threshold handling b12x/attention/_shared/mla/compressed_config.py	Split configuration accepts and validates decode_row_capacity. Decode branches use the computed threshold, and chunk planning forwards the value.
Scratch capability integration b12x/attention/compressed_mla/_scratch.py	B12XCompressedMLAScratchCaps stores a normalized positive decode row capacity and passes it to split planning.
Capacity behavior validation tests/attention/test_attention_mla_compressed.py	Tests cover capacity propagation, invalid values, legacy plans, CUDA-graph replay, formatting updates, and deterministic index wrapping.

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

Sequence Diagram(s)

sequenceDiagram
  participant B12XCompressedMLAScratchCaps
  participant compressed_mla_split_config_for_contract
  participant DecodeSplitSelection
  B12XCompressedMLAScratchCaps->>compressed_mla_split_config_for_contract: provide decode_row_capacity
  compressed_mla_split_config_for_contract->>DecodeSplitSelection: apply computed decode row threshold

Loading

Suggested reviewers: lukealonso

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main change: declaring compressed MLA decode row capacity.

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

[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 `@tests/attention/test_attention_mla_compressed.py`:
- Around line 395-485: Add big-page-ID coverage to both compressed-MLA tests: in
tests/attention/test_attention_mla_compressed.py lines 395-485, update
test_declared_decode_capacity_matches_reference_under_graph_replay to use live
page IDs beyond the 2^31 / stride boundary while preserving CUDA Graph replay
and reference comparisons; in tests/attention/test_attention_mla_compressed.py
line 1332, add the equivalent deterministic replay coverage with live pages
beyond that boundary. Use deterministic indices and ensure both paths validate
correct outputs.

🪄 Autofix

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: 866f705e-3392-4cec-9198-c49fd950de9e

📥 Commits

Reviewing files that changed from the base of the PR and between 89f8b9b and 282d61a.

📒 Files selected for processing (3)

• b12x/attention/_shared/mla/compressed_config.py
• b12x/attention/compressed_mla/_scratch.py
• tests/attention/test_attention_mla_compressed.py