fix(mla): use packaged compact KV cache writer

[davidsyoung]

Dependency

Requires local-inference-lab/b12x#37, which adds the public compact-cache writer API.

Summary

• replace the KV_FP8_ROPE=1 external shared-library loader with b12x.attention.mla.kv_cache.concat_and_cache_nvfp4_mla_fp8_rope
• bind the package callable once during B12xMLASparseImpl construction only when the compact 368-byte cache format is enabled
• preserve the KV_FP8_ROPE=0 stock ops.concat_and_cache_mla path without importing the new b12x submodule
• preserve the empty-cache no-op and the non-NVFP4 cache dtype guard
• pass the existing latent tensor, squeezed 64-d RoPE tensor, flattened slot map, cache, and identical k_scale object to the package API
• remove the downstream duplicate shape/record-width checks; the b12x public wrapper now owns the record ABI and input validation
• fail closed with a targeted error when the requested package API is unavailable; non-import initialization failures propagate unchanged

This removes the deployment-only KV_FP8_ROPE_WRITER_LIB and torch.ops.load_library dependency. The ordinary 432-byte writer and all KV_FP8_ROPE=0 behavior remain unchanged.

Compatibility and rollout

This PR must land with a b12x revision containing b12x#37. With an older b12x package, KV_FP8_ROPE=1 raises a targeted startup error rather than silently falling back to the stock 432-byte writer. Disabled deployments do not import the new submodule and continue through the existing stock writer.

This is a package-boundary correction, not a new performance candidate; no serving-throughput claim is made. The writer bytes and both production reader families are covered in b12x#37.

Test plan

• pytest tests/v1/attention/test_b12x_mla_fp8_rope_writer.py -vv — 6 passed in the gg-v18 Python 3.12 / torch environment
  • disabled construction never imports/calls the compact writer and reaches the stock writer
  • enabled construction binds/calls the public package API with exact squeeze/flatten/scale semantics
  • wrong cache dtype fails before mutation
  • empty cache remains a no-op
  • missing public API fails closed
  • non-ImportError initialization failures propagate
• ruff check vllm/v1/attention/backends/mla/b12x_mla_sparse.py tests/v1/attention/test_b12x_mla_fp8_rope_writer.py — passed with the repository-pinned ruff 0.14.0
• ruff format --check vllm/v1/attention/backends/mla/b12x_mla_sparse.py tests/v1/attention/test_b12x_mla_fp8_rope_writer.py — passed

Summary by CodeRabbit

• New Features

  • Added integration with the packaged KV-cache writer for supported FP8 RoPE configurations.
  • Added validation to prevent unsupported cache formats from being processed.

• Bug Fixes

  • Improved handling of empty cache tensors.
  • Added fail-fast behavior when the required writer integration is unavailable or cannot initialize.

• Tests

  • Added comprehensive coverage for enabled and disabled writer modes, argument forwarding, validation, fallback behavior, and initialization failures.

[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: d23c42cf-478b-4c85-ac12-20c22965566b

📥 Commits

Reviewing files that changed from the base of the PR and between 5b116d1 and 77123d3.

📒 Files selected for processing (2)

• tests/v1/attention/test_b12x_mla_fp8_rope_writer.py
• vllm/v1/attention/backends/mla/b12x_mla_sparse.py

---

📝 Walkthrough

Walkthrough

The MLA backend now uses the packaged b12x FP8 RoPE KV-cache writer when enabled, while preserving the stock path otherwise. New tests cover argument forwarding, dtype and empty-cache behavior, missing APIs, and initialization failures.

Changes

FP8 RoPE writer integration

Layer / File(s)	Summary
Bind and invoke the packaged writer vllm/v1/attention/backends/mla/b12x_mla_sparse.py	The backend replaces standalone writer loading and direct torch-op invocation with the packaged b12x writer API.
Validate writer routing and failure behavior tests/v1/attention/test_b12x_mla_fp8_rope_writer.py	Tests cover disabled and enabled routing, transformed arguments, NVFP4 validation, empty caches, missing APIs, and initialization failures.

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

Sequence Diagram(s)

sequenceDiagram
  participant B12xMLASparseImpl
  participant b12x_writer as b12x FP8 RoPE writer
  participant KVCache as KV cache tensors
  B12xMLASparseImpl->>b12x_writer: invoke concat_and_cache_nvfp4_mla_fp8_rope
  b12x_writer->>KVCache: update NVFP4 cache

Loading

Possibly related PRs

• local-inference-lab/vllm#95: Earlier FP8 RoPE KV-cache writer routing changes in the same MLA backend.

🚥 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: switching MLA KV cache writing to the packaged writer.
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.}

[davidsyoung]

CI note: pre-run-check requires one of verified, ready, or ready-run-all-tests for external authors, but none of those labels currently exists in this repository's label catalog. The focused candidate tests are already green (6/6 vLLM seam tests, 18/18 dependent b12x writer/reader tests, ruff check/format clean). Maintainer action is required to create/apply one of the workflow-recognized labels or adjust/re-run the gate.