fix(gemm): preserve dense API contracts with block FP8

[voipmonitor]

Summary

• preserve the existing fused_quant_bf16 and FP6 row_scale API contracts when adding the regular K128 block-FP8 path
• keep block-FP8 accumulator coordinates in distinct CuTe DSL locals
• add a focused signature regression test for both legacy parameters

Root cause

Commit fe92b95 added block_fp8, but two signatures replaced existing parameters instead of appending the new one:

• DenseGemmKernel.fused_quant_bf16
• dense_gemm.row_scale

The implementation and callers still used those parameters. This caused model warmup to fail on FP6 row-scale paths and removed the explicit fused-quant compile-key input.

The same change also introduced a block-FP8 local named c_identity in the shared CuTe DSL kernel. Legacy fused-quant compilation uses that name in other compile-time branches. CuTe therefore observed None on one control-flow path and a tensor on another and rejected the first real DSpark request with TYPE_UNSTABLE_JOIN. Distinct block_c_identity / block_coord_mn locals remove that cross-family type join without changing either kernel's arithmetic.

Validation

• Ruff and git diff --check: pass
• targeted API, FP6 row-scale, block-FP8 correctness and graph tests: 9 passed
• legacy fused-quant GPU oracle plus CUDA graph replay: passed
• DS4-Flash-0731 TP4/K5 full-model boot: target, DSpark and context-KV FULL CUDA graph capture passed
• 64 concurrent deterministic requests: 64/64 correct, no runtime errors
• matched remote-host performance versus r31:
  • C1: 238.7 vs 229.7 tok/s
  • C64: 2551.3 vs 2540.5 tok/s
  • prefill 8k/64k: 15187/14904 tok/s

The full-model test used GPUs 4-7 on the same direct-attached host as the r31 reference.

Summary by CodeRabbit

• New Features

  • Added optional per-row scaling support for MX-FP6 dense matrix multiplication.
  • Added per-launch control over fused BF16 quantization.
  • Preserved the existing output rounding behavior for scaled operations.

• Documentation

  • Documented the new row-scaling behavior and optional controls.

• Tests

  • Added regression coverage to verify the public API defaults remain available.

[coderabbitai]

📝 Walkthrough

Walkthrough

Dense GEMM adds optional fused_quant_bf16 and row_scale parameters, documents MX-FP6 row scaling, renames block-FP8 coordinate identities, and adds runtime API contract tests.

Changes

Dense GEMM API and kernel updates

Layer / File(s)	Summary
API contracts and validation b12x/_lib/dense_gemm.py, tests/gemm/test_fp6_packed_b.py	DenseGemmKernel accepts fused_quant_bf16=None. dense_gemm accepts row_scale=None. The API documents MX-FP6 per-row epilogue scaling and its preserved second C-dtype rounding. A regression test validates both defaults.
Block-FP8 coordinate naming b12x/_lib/dense_gemm.py	Block-FP8 accumulator coordinate identities and their use sites now use block_coord_mn.

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

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 preserving dense GEMM API contracts while adding block FP8 support.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix/dense-gemm-row-scale-api-20260809

---

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

[lukealonso]

Closing this PR because it was opened before August 15, 2026. Please reopen it if the change is still valid.