fix(sm120): add labels-contract checker for m-grouped contiguous GEMM

Same fix as the nv_dev-lineage sibling: group boundaries in labels mode
must be multiples of the runtime mk alignment, or a BLOCK_M tile
straddles two groups and silently computes the straddled rows with the
wrong group's B (reported as 'middle empty group corrupts later groups';
the empty group merely made the misalignment likely). The tile
heuristics already guarantee BLOCK_M divides the runtime alignment, so
contract-respecting layouts are always safe; the opt-in checker
(DG_CHECK_CONTIGUOUS_LABELS=1) turns violations into a loud host error.
Regression tests pin middle-empty patterns at al in {64, 128} and the
rejection path.

Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>

Author: lucifer1004

csrc/apis/gemm.hpp

@@ -1,5 +1,7 @@
 #pragma once
 
+#include <format>
+
 #include "../utils/compatibility.hpp"
 
 #include "../jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp"
@@ -182,6 +184,37 @@ static void fp8_fp4_gemm_tt(const std::pair<torch::Tensor, torch::Tensor>& a,
                     d, c, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast, alpha);
 }
 
+// SM120 m-grouped contiguous (labels-mode) layout contract: each group's first row must
+// start at a multiple of `get_mk_alignment_for_contiguous_layout()`. The kernel selects B
+// (and SFB) per BLOCK_M tile from the label of the tile's first row, and the tile
+// heuristics guarantee BLOCK_M divides the runtime alignment, so contract-respecting
+// labels never put a group boundary inside a tile. Labels built at a finer granularity
+// than the runtime alignment violate this and silently compute the straddled rows with
+// the wrong group's B. Set DG_CHECK_CONTIGUOUS_LABELS=1 to turn violations into a loud
+// error (does one GPU->CPU copy of the labels; meant for debugging/integration).
+static void sm120_check_contiguous_labels_contract(const torch::Tensor& grouped_layout,
+                                                   const int& num_groups) {
+    if (not deep_jit::get_env<bool>("DG_CHECK_CONTIGUOUS_LABELS", false))
+        return;
+    const int alignment = heuristics_runtime->get_mk_alignment_for_contiguous_layout();
+    const auto labels = grouped_layout.cpu();
+    const auto* data = labels.data_ptr<int>();
+    const auto m = static_cast<int64_t>(labels.size(0));
+    for (int64_t i = 0, prev_label = -1; i < m; ++ i) {
+        const int label = data[i];
+        DG_HOST_ASSERT(label >= -1 and label < num_groups
+                       and "m-grouped contiguous label out of range");
+        if (label < 0 or label == prev_label)
+            continue;
+        if (i % alignment != 0)
+            DG_HOST_UNREACHABLE(std::format(
+                "m-grouped contiguous group {} starts at row {}, not a multiple of the "
+                "runtime mk alignment {}; rebuild the labels with matching alignment or "
+                "call set_mk_alignment_for_contiguous_layout()", label, i, alignment));
+        prev_label = label;
+    }
+}
+
 static void m_grouped_fp8_fp4_gemm_nt_contiguous(const std::pair<torch::Tensor, torch::Tensor>& a,
                                                  const std::pair<torch::Tensor, torch::Tensor>& b,
                                                  const torch::Tensor& d,
@@ -246,6 +279,8 @@ static void m_grouped_fp8_fp4_gemm_nt_contiguous(const std::pair<torch::Tensor,
                                                      num_groups, m, n, k, gran_k_a, gran_k_b, major_a, major_b,
                                                      compiled_dims, use_psum_layout, ensure_zero_padding, expected_m_for_psum_layout);
     } else if (arch_major == 12 and sfa.scalar_type() == torch::kInt and sfb.scalar_type() == torch::kInt) {
+        if (not use_psum_layout)
+            sm120_check_contiguous_labels_contract(grouped_layout, num_groups);
         const auto b_data = sm120::to_k_major(b.first, major_b, n);
         const bool is_mixed_fp4 = (a.first.scalar_type() != b_data.scalar_type()) and
                                   (a.first.scalar_type() == kPackedFP4 or b_data.scalar_type() == kPackedFP4);
@@ -630,6 +665,8 @@ static void m_grouped_bf16_gemm_nt_contiguous(const torch::Tensor& a, const torc
                                              num_groups, m, n, k, major_a, major_b, compiled_dims,
                                              use_psum_layout, ensure_zero_padding, expected_m_for_psum_layout);
     } else if (arch_major == 12) {
+        if (not use_psum_layout)
+            sm120_check_contiguous_labels_contract(grouped_layout, num_groups);
         sm120_m_grouped_bf16_gemm_contiguous(a, b, d, grouped_layout,
                                              num_groups, m, n, k, major_a, major_b, compiled_dims,
                                              use_psum_layout, expected_m_for_psum_layout);

tests/test_sm120_fp8_fp4.py

@@ -405,3 +405,74 @@ def test_sm120_kgroup_descriptor_reuse_at_default_sms() -> None:
     test_sm120_asymmetric_scale_recipe_swap()
     test_sm120_scale_dtype_validation()
     test_sm120_masked_physical_capacity()
+
+
+@test_filter(lambda: get_arch_major() == 12)
+def test_sm120_contiguous_labels_middle_empty_group() -> None:
+    """Labels-mode contiguous GEMM with an empty group in the middle: groups after the
+    gap must still read their own B (regression for a host/kernel contract where group
+    boundaries must align with the runtime mk alignment)."""
+    old_alignment = deep_gemm.get_mk_alignment_for_contiguous_layout()
+    try:
+        for alignment in (64, 128):
+            deep_gemm.set_mk_alignment_for_contiguous_layout(alignment)
+            for lengths in ([100, 0, 130, 65], [128, 130, 65], [65, 0, 0, 66, 31]):
+                n, k = 128, 256
+                intervals, end = [], 0
+                for length in lengths:
+                    start = (end + alignment - 1) // alignment * alignment
+                    end = start + length
+                    intervals.append((start, end))
+                m = (end + alignment - 1) // alignment * alignment
+                a = (torch.full((m, k), 2.0, device='cuda').to(torch.float8_e4m3fn),
+                     torch.ones(m, k // 128, device='cuda'))
+                b = (torch.cat([torch.full((1, n, k), float(g + 1), device='cuda')
+                                for g in range(len(lengths))]).to(torch.float8_e4m3fn),
+                     torch.ones(len(lengths), n, k // 128, device='cuda'))
+                labels = torch.full((m,), -1, dtype=torch.int32)
+                for g, (start, end) in enumerate(intervals):
+                    labels[start:end] = g
+                d = torch.full((m, n), 3.0, dtype=torch.bfloat16, device='cuda')
+                deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous(
+                    a, b, d, labels.cuda(), recipe=(1, 1, 128))
+                for g, (start, end) in enumerate(intervals):
+                    if start == end:
+                        continue
+                    want = float(k * 2 * (g + 1))
+                    assert torch.all(d[start:end] == want), \
+                        (alignment, lengths, g, d[start:end].unique())
+    finally:
+        deep_gemm.set_mk_alignment_for_contiguous_layout(old_alignment)
+
+
+@test_filter(lambda: get_arch_major() == 12)
+def test_sm120_contiguous_labels_contract_rejection() -> None:
+    """Labels finer than the runtime mk alignment must fail loudly when the opt-in
+    contract checker is enabled (DG_CHECK_CONTIGUOUS_LABELS=1), not corrupt silently."""
+    import os
+    old_env = os.environ.get('DG_CHECK_CONTIGUOUS_LABELS')
+    os.environ['DG_CHECK_CONTIGUOUS_LABELS'] = '1'
+    old_alignment = deep_gemm.get_mk_alignment_for_contiguous_layout()
+    deep_gemm.set_mk_alignment_for_contiguous_layout(128)
+    try:
+        m, n, k = 256, 128, 256
+        a = (torch.zeros(m, k, device='cuda').to(torch.float8_e4m3fn),
+             torch.ones(m, k // 128, device='cuda'))
+        b = (torch.zeros(2, n, k, device='cuda').to(torch.float8_e4m3fn),
+             torch.ones(2, n, k // 128, device='cuda'))
+        labels = torch.full((m,), -1, dtype=torch.int32)
+        labels[0:64] = 0
+        labels[64:128] = 1  # group 1 starts at 64: finer than the runtime alignment 128
+        d = torch.zeros(m, n, dtype=torch.bfloat16, device='cuda')
+        try:
+            deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous(a, b, d, labels.cuda(),
+                                                           recipe=(1, 1, 128))
+            raise AssertionError('expected the contract checker to reject')
+        except RuntimeError as e:
+            assert 'not a multiple of the runtime mk alignment' in str(e), e
+    finally:
+        if old_env is None:
+            os.environ.pop('DG_CHECK_CONTIGUOUS_LABELS', None)
+        else:
+            os.environ['DG_CHECK_CONTIGUOUS_LABELS'] = old_env
+        deep_gemm.set_mk_alignment_for_contiguous_layout(old_alignment)