[codex] Pad B12X sparse MLA decode heads

[voipmonitor]

Summary

Pads the B12X sparse MLA decode plan/head dimension to an 8-head block and slices decode outputs back to the actual local head count.

Why

GLM-5.1 TP16 shards have only 4 local MLA heads per rank. The existing prefill/extend path already pads small head counts, but decode used the raw 4-head shard directly. That reproduced as coherent short outputs but eventual CJK/garbage corruption under longer/concurrent Estonia-style decode. TP8 has 8 local heads, so it does not hit this path.

The fix keeps the vLLM binding path eager and caller-scratch-owned: plan.bind(scratch=...) -> binding -> kernel. It does not introduce b12x workspaces/arenas or cached workspace bindings.

Validation

• python3 -m py_compile vllm/v1/attention/backends/mla/b12x_mla_sparse.py
• git diff --check -- vllm/v1/attention/backends/mla/b12x_mla_sparse.py
• Runtime validation of the equivalent decode-head padding on GLM-5.1 NVFP4-MTP TP16/DCP1/no-MTP kept B12X_MLA_SPARSE + B12X NvFp4 MoE and logged Padding B12X_MLA_SPARSE decode heads from 4 to 8.
• /mnt/test.py --port 5329 --model GLM-5.1-NVFP4-MTP --max-tokens 220 returned coherent output with chinese_count=0.
• Estonia streaming probe completed 30/30 requests at max_tokens=4096 with total_cjk=0 and no errors.

Summary by CodeRabbit

• Performance
  • Optimized attention kernel decode operations for improved performance and memory efficiency.

[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: 506bf538-ef9b-47fd-91a5-92487c2dd278

📥 Commits

Reviewing files that changed from the base of the PR and between 8f6c0a7 and 01d98f6.

📒 Files selected for processing (1)

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

---

📝 Walkthrough

Walkthrough

This PR introduces decode-head block alignment to optimize kernel execution. A new constant parameterizes block-size padding, the decode head count is computed with alignment, the decode execution plan uses the padded count, and the query tensor is conditionally padded during kernel invocation before outputs are sliced back to the original dimensions.

Changes

Decode-head padding optimization

Layer / File(s)	Summary
Padding configuration and decode-head alignment vllm/v1/attention/backends/mla/b12x_mla_sparse.py	_DECODE_HEADS_PER_BLOCK constant parameterizes block-aligned padding; _decode_num_heads is computed rounded up to that block size with minimum clamping, and logging is extended to reflect decode-head padding.
Decode plan rebuild with aligned padding vllm/v1/attention/backends/mla/b12x_mla_sparse.py	Decode plan is now constructed using the aligned _decode_num_heads rather than workspace-padded head count.
Query padding in decode kernel path vllm/v1/attention/backends/mla/b12x_mla_sparse.py	Query tensor is conditionally padded from num_actual_heads to decode_num_heads before kernel binding and execution.
Output tensor slicing back to original dimensions vllm/v1/attention/backends/mla/b12x_mla_sparse.py	After kernel execution, out (and lse when applicable) are sliced from decode_num_heads down to num_actual_heads and made contiguous in both return paths.

🎯 2 (Simple) | ⏱️ ~12 minutes

🚥 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%.	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: padding B12X sparse MLA decode heads to an 8-head block, which is the primary technical modification in this changeset.
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.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

✅ Unit Test PR creation complete.

• Create PR with unit tests
• Commit unit tests in branch codex/tp16-small-head-decode-padding

---

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 and usage tips.}

[voipmonitor]

@coderabbitai review

[coderabbitai]

✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[coderabbitai]

Note

Unit test generation is a public access feature. Expect some limitations and changes as we gather feedback and continue to improve it.

---

Generating unit tests... This may take up to 20 minutes.

[coderabbitai]

✅ Created PR with unit tests: #6

[lukealonso]

Fixing this another way