mla: split latent_scale identity/dynamic compile-cache entries

[MadeBy561]

CuTeDSL folds the NVFP4 outer-scale multiply out of the MLA kernels when
latent_scale traces at exactly 1.0. The explicit compile keys
(attention.mla.sm120.decode v17, attention.mla.sm120.prefill_mg v3) hash
neither kernel source nor scalar values, so a persistent cache populated by
scale-1.0 runs replays identity cubins for launches with real per-layer
scales — the restore is silently dropped.

Repro (GLM-5.2, SM120, TP4+DCP4, nvfp4_ds_mla + VLLM_NVFP4_MLA_SCALES_FILE):
teacher-forced prefill KLD 0.13 → 6.7–7.4 on a stale cache; 0.13 on a fresh
cache, identical config. Short prompts through serving degenerate to
near-immediate EOS.

Fix: add latent_scale_identity as a compile-spec fact and bump both spec
versions (decode 17→18, prefill_mg 3→4), following the v17 FP8-RoPE ABI-bump
pattern. Two entries max per config; both directions correct.

🤖 Generated with Claude Code

Summary by CodeRabbit

• Bug Fixes
  • Fixed kernel compilation caching to correctly distinguish identity and non-identity latent scaling configurations.
  • Prevented incompatible compiled kernels from being reused across different scaling settings.
  • Improved reliability for attention decoding and prefill operations.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 95b99c3e-702b-4885-9aab-2df9416f305a

📥 Commits

Reviewing files that changed from the base of the PR and between 957b9fd and 133bbb1.

📒 Files selected for processing (2)

• b12x/attention/mla/kernel.py
• b12x/attention/mla/prefill_mg.py

---

Walkthrough

The MLA decode and prefill kernel compile-spec cache keys now distinguish latent_scale == 1.0 from other latent-scale values, with compile-spec versions bumped accordingly.

Changes

MLA compile cache updates

Layer / File(s)	Summary
Separate latent-scale compilation variants b12x/attention/mla/kernel.py, b12x/attention/mla/prefill_mg.py	Decode and prefill compile specifications add the latent_scale_identity cache key and bump their compile-spec versions to prevent incompatible cubin reuse.

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

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the main change: separating compile-cache entries for identity versus dynamic latent_scale cases.
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.

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

[MadeBy561]

Closing: authored against the pre-restructure b12x layout (b12x/attention/mla/) and now conflicting. Superseded by the sparkinfer-layout port with the identical fix and spec-version bumps.