From 0b44e769caa3c38d99d346a3a343a47df68d2e3b Mon Sep 17 00:00:00 2001
From: Martin Vit <martin@voipmonitor.org>
Date: Sun, 9 Aug 2026 06:25:03 +0000
Subject: [PATCH] fix(gemm): preserve dense API contracts with block FP8

---
 b12x/_lib/dense_gemm.py         | 17 +++++++++++++----
 tests/gemm/test_fp6_packed_b.py | 14 ++++++++++++++
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/b12x/_lib/dense_gemm.py b/b12x/_lib/dense_gemm.py
index 1c72c8cbd..818e37b22 100644
--- a/b12x/_lib/dense_gemm.py
+++ b/b12x/_lib/dense_gemm.py
@@ -630,6 +630,7 @@ def __init__(
         mxfp6_fmt_b: Optional[str] = None,
         b_packed: bool = False,
         plain_fp8: bool = False,
+        fused_quant_bf16: Optional[bool] = None,
         block_fp8: bool = False,
     ):
         # When set, A/B operands are MX codes carried in Float8E4M3FN
@@ -1884,10 +1885,12 @@ def kernel(
         accumulators = cute.make_rmem_tensor(acc_shape, self.acc_dtype)
         if cutlass.const_expr(self.block_fp8):
             stage_accumulators = cute.make_rmem_tensor(acc_shape, self.acc_dtype)
-            c_identity = cute.make_identity_tensor(
+            block_c_identity = cute.make_identity_tensor(
                 cute.slice_(self.tile_shape_mnk, (None, None, 0))
             )
-            coord_mn = _reshape_acc_to_mn(thr_mma.partition_C(c_identity))
+            block_coord_mn = _reshape_acc_to_mn(
+                thr_mma.partition_C(block_c_identity)
+            )
 
         # Cluster/thread sync
         if cute.size(self.cluster_shape_mnk) > 1:
@@ -2340,7 +2343,7 @@ def kernel(
                                 self._accumulate_block_fp8_stage(
                                     accumulators,
                                     stage_accumulators,
-                                    coord_mn,
+                                    block_coord_mn,
                                     directSFA_mkl,
                                     directSFB_nkl,
                                     tile_coord_mnl,
@@ -2485,7 +2488,7 @@ def kernel(
                             self._accumulate_block_fp8_stage(
                                 accumulators,
                                 stage_accumulators,
-                                coord_mn,
+                                block_coord_mn,
                                 directSFA_mkl,
                                 directSFB_nkl,
                                 tile_coord_mnl,
@@ -7094,6 +7097,7 @@ def dense_gemm(
     x_bf16: Optional[torch.Tensor] = None,
     w_gscale: Optional[torch.Tensor] = None,
     plain_fp8: bool = False,
+    row_scale: Optional[torch.Tensor] = None,
     block_fp8: bool = False,
 ) -> torch.Tensor:
     """Execute dense block-scaled GEMM for one expert-major batch stack.
@@ -7134,6 +7138,11 @@ def dense_gemm(
     SM12x dense pipeline. The scalar ``alpha`` carries the combined activation
     and weight dequantization scale.
 
+    ``row_scale``: optional contiguous ``(M,)`` tensor in the C dtype, applied
+    per output row in the epilogue. It replaces a separate ``out.mul_(v)``
+    launch and reproduces that multiply bit-for-bit, including its second
+    rounding to the C dtype. MX-FP6 only.
+
     ``block_fp8``: accumulate ordinary E4M3 MMA over each K128 block, then
     apply compact FP32 activation ``[M,K/128]`` and weight
     ``[N/128,K/128]`` scales before adding it to the final accumulator.
diff --git a/tests/gemm/test_fp6_packed_b.py b/tests/gemm/test_fp6_packed_b.py
index 466efd0c6..231bb2279 100644
--- a/tests/gemm/test_fp6_packed_b.py
+++ b/tests/gemm/test_fp6_packed_b.py
@@ -7,6 +7,8 @@
 """
 from __future__ import annotations
 
+import inspect
+
 import pytest
 import torch
 
@@ -15,6 +17,18 @@
 )
 
 
+def test_dense_gemm_preserves_fp6_api_contracts():
+    """Keep the FP6 caller and dense GEMM implementation contracts in sync."""
+    from b12x._lib.dense_gemm import DenseGemmKernel, dense_gemm
+
+    parameter = inspect.signature(dense_gemm).parameters["row_scale"]
+    assert parameter.default is None
+    kernel_parameter = inspect.signature(DenseGemmKernel).parameters[
+        "fused_quant_bf16"
+    ]
+    assert kernel_parameter.default is None
+
+
 def _gemm_operands(m: int, n: int, k: int, source_format: str):
     """Quantized operands for one case, mirroring ``dense_fp6_linear_expanded``."""
     from b12x._lib.fp6 import SF_VEC_SIZE_FP6, as_grouped_mxfp6_scale_view