feat(gemm): opt-in labels-contract checker for m-grouped contiguous GEMM - #450
lucifer1004 wants to merge 1 commit into
Conversation
| if (label < 0 or label == prev_label) | ||
| continue; |
There was a problem hiding this comment.
🟡 warning: 检查 padding 后同一标签重新出现的边界: 开启检查且 alignment=128 时，[0]*128 + [1]*128 + [-1] + [1]*127 会通过检查，因为跳过 padding 后 prev_label 仍为 1。然而第 256 行所在 tile 的首标签是 -1，Scheduler::get_global_idx 会选择 B[0]，使第 257–383 行错误地使用 B[0] 而非 B[1]，仍然静默产生错误结果。请不要跨 padding 沿用前一个标签而跳过边界检查。
🤖 v6
| else: | ||
| # Unsupported architectures fail after the contract check, at kernel dispatch | ||
| assert 'not a multiple of the runtime mk alignment' not in raised, raised |
There was a problem hiding this comment.
🟡 warning: 让受支持架构上的正向测试拒绝非预期异常: 在 sm90/sm100 等受支持架构上，匹配 alignment 的调用若因标签范围检查、SF 转换或 CUDA/JIT 错误抛出 RuntimeError，这里仍会判定测试通过，只要异常不包含指定的 alignment 文案。这会掩盖合法输入路径的回归；已通过注入标签范围异常确认该测试仍然通过。应仅在明确不支持的架构上允许预期的 dispatch 异常，受支持架构必须执行成功并验证数值。
🤖 v6
| 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) { |
There was a problem hiding this comment.
🔵 suggestion: prev_label is never reset when a -1 is encountered, so a non-negative label run that resumes after padding at an unaligned row is not flagged. Example with alignment 128: [0]*128 + [-1]*3 + [0]*125 — tile 1's first row is -1, is_computation_valid skips the tile, and rows 131..255 are silently never computed. Note that simply resetting prev_label on -1 would introduce false positives ([0]*50 + [-1]*10 + [0]*68 is computed correctly since tile 0's first row is 0). The exact rule is per-tile: for each alignment-sized tile with first-row label L, all rows must be -1 if L == -1, otherwise all rows must be L or -1. That is about the same amount of code and precise rather than merely sufficient. Realistic callers (sorted expert assignments) never build such layouts, so this is optional, not blocking.
🤖 v5
| actual_ms = [int(expected_m_per_group * random.uniform(0.7, 1.3)) for _ in range(num_groups)] | ||
| quant_config: Optional[QuantConfig] = None, | ||
| actual_ms_override: Optional[List[int]] = None): | ||
| if actual_ms_override is not None: |
There was a problem hiding this comment.
🔵 suggestion: Nit: actual_ms_override silently overwrites the caller-provided num_groups. Fine for the tests as written (they pass len(actual_ms)), but an assert num_groups == len(actual_ms_override) or a short comment would make the override contract explicit and prevent a future caller from passing a mismatched num_groups without noticing.
🤖 v5
|
|
||
|
|
||
| def test_m_grouped_gemm_contiguous_labels_contract_rejection() -> None: | ||
| print('Testing m-grouped contiguous labels contract rejection:') |
There was a problem hiding this comment.
🔵 suggestion: The rejection test only exercises the fp8/fp4 entry point. Since check_contiguous_labels_contract is a shared helper this is acceptable, but a single-call bf16 variant (or a parametrised loop over both entry points) would guard against a future refactor dropping the call site in m_grouped_bf16_gemm_nt_contiguous.
🤖 v5
| const int& num_groups) { | ||
| if (not deep_jit::get_env<int>("DG_CHECK_CONTIGUOUS_LABELS")) | ||
| return; | ||
| const int alignment = heuristics_runtime->get_mk_alignment_for_contiguous_layout(); |
There was a problem hiding this comment.
🟡 warning: i % alignment at line 200 divides by zero if the runtime MK alignment is ever set to 0 (the setter set_mk_alignment_for_contiguous_layout performs no validation). Other paths would also misbehave, but the checker runs first and would crash the process; consider DG_HOST_ASSERT(alignment > 0) here, or validate the value in the setter.
🤖 v4
| 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 |
There was a problem hiding this comment.
🔵 suggestion: The checker only accepts -1 as the padding sentinel, while the device scheduler treats any negative label as padding (grouped_layout[...] >= 0 for validity and cute::max(0, ...) for the B offset in deep_gemm/include/deep_gemm/scheduler/gemm.cuh). A caller that pads with a different negative value would be rejected even though the kernel handles it. Consider relaxing to label < num_groups (any negative = padding), or explicitly document that only -1 is supported. This is the one place where the checker is stricter than the actual contract.
🤖 v4
| print() | ||
|
|
||
|
|
||
| def test_m_grouped_gemm_contiguous_labels_contract_rejection() -> None: |
There was a problem hiding this comment.
🔵 suggestion: test_m_grouped_gemm_contiguous_labels_contract_rejection only exercises the alignment branch. The range-validation branch added by the checker (label < -1 or label >= num_groups) is untested; please add a case that enables the env var and passes an out-of-range label, asserting the assertion fires with a useful message.
🤖 v4
| print() | ||
|
|
||
|
|
||
| def test_m_grouped_gemm_contiguous_middle_empty_groups() -> None: |
There was a problem hiding this comment.
🔵 suggestion: The contract checker is also wired into m_grouped_bf16_gemm_nt_contiguous, but the rejection test only exists for fp8/fp4. Add a bf16 rejection case (or parametrize the fp8 test over both entry points) so the bf16 call site cannot silently regress.
🤖 v4
🤖 ds-review-bot Code Reviewv6检查器存在可导致静默错误结果的漏检，正向测试也会吞掉非预期异常。已用独立 host 用例和 mock 验证；当前环境未安装 PyTorch，未运行 GPU 测试。 v5Adds an opt-in (DG_CHECK_CONTIGUOUS_LABELS=1) host-side contract checker for labels-mode m-grouped contiguous GEMM, plus tests for middle-empty-group patterns and loud rejection of misaligned labels. Verified against the kernel and heuristics code: the encoded contract is correct. Findings: (1) README.md documents every DG_* env var including debug-only ones, but Overall: correct, minimal, zero-cost when disabled, no device-code change. Approve after adding the README entry. Tests were not executed here (no GPU available); conclusions are from code inspection. v4This change adds an opt-in host-side contract checker for labels-mode m-grouped contiguous GEMM. The implementation is well-scoped and correct: the checker lives in the shared labels branch of both Files reviewed: 4 📍 未定位到 diff 的评论🟡 warning 🔵 suggestion |
Labels-mode m-grouped contiguous GEMM requires every run of a non-padding label to start at a multiple of the runtime mk alignment: the kernels select B (and SFB) per BLOCK_M tile from the tile's first row label, and heuristics guarantee BLOCK_M divides the runtime alignment. Labels built at a finer granularity put a group boundary inside a tile and silently compute the straddled rows with the wrong group's B, or skip the tile entirely when padding lands on its first row (reported downstream as 'an empty middle group corrupts later groups'; an empty group merely makes a misaligned boundary likely -- [128, 130, 65] built at 64 with runtime 128 corrupts identically with no empty group). With DG_CHECK_CONTIGUOUS_LABELS=1, check_contiguous_labels_contract() validates labels on the host (one GPU->CPU copy; opt-in for integration debugging) and raises a DGException naming the group, row, and alignment with the remedy: any negative label is padding (matching the device scheduler), any label >= num_groups is rejected, and every non-padding label run must start at an aligned row, including a run that resumes after padding. Arch-generic: the call sites are in the shared labels layout-check branch of both m_grouped_fp8_fp4_gemm_nt_contiguous and m_grouped_bf16_gemm_nt_contiguous, ahead of arch dispatch. Tests: generate_m_grouped_contiguous gains actual_ms_override (default path byte-identical); new tests pin middle-empty-group patterns (the previously uncovered class) for fp8/fp4 and bf16, and the loud rejection of misaligned, resumed-after-padding, and out-of-range labels plus the matching-alignment pass. Validated on sm_103: new tests pass and the pre-existing m-grouped contiguous/masked suites are unchanged (checker inert unless enabled). Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
c624490 to
ba1bcf5
Compare
|
Thanks @ds-review-bot for the careful rounds — all findings addressed in 🟡 warnings — fixed
🔵 suggestions
Validation (fresh JIT cache,
|
| print() | ||
|
|
||
|
|
||
| def test_m_grouped_gemm_contiguous_labels_contract_rejection() -> None: |
There was a problem hiding this comment.
🔵 suggestion: Nit (non-blocking): the bf16 rejection test only exercises the misaligned-run case, whereas the fp8/fp4 variant also covers resume-after-padding, out-of-range, and the matching-alignment positive half. Since check_contiguous_labels_contract is shared, this is sufficient as a call-site guard; just noting the asymmetry in case you want the two tests to mirror each other.
🤖 v5
| // M-grouped contiguous (labels-mode) layout contract: every run of a non-padding label | ||
| // must start at a multiple of `get_mk_alignment_for_contiguous_layout()`. The kernels | ||
| // select 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 |
There was a problem hiding this comment.
🔵 suggestion: Nit (non-blocking): the comment says the heuristics guarantee BLOCK_M divides the runtime alignment; in practice sm90.hpp:33 and sm100.hpp:40 set BLOCK_M == get_mk_alignment_for_contiguous_layout(). 'Divides' is a correct (weaker) statement so no change is required, but 'equals' would be more precise if you touch this again.
🤖 v5
| // ROW's label (not the previous non-padding one), so a label run resuming after | ||
| // padding at an unaligned row is still caught. | ||
| int64_t prev_label = -1; | ||
| for (int64_t i = 0; i < m; ++ i) { |
There was a problem hiding this comment.
🔵 suggestion: The run-based rule is strictly stronger than the device's per-tile rule and can produce false positives. At alignment 128, [0]*50 + [-1]*10 + [0]*68 is computed correctly: tile 0's first row is 0 and every non-padding row in the tile is group 0, so the scheduler selects B[0] for the whole tile. But prev_label is -1 at row 60 and 60 % 128 != 0, so the checker raises. The v6 review explicitly called this out as a layout the kernel handles. Consider the exact per-tile rule so the checker mirrors the scheduler precisely: for each alignment-sized tile, if the first-row label is padding, every row must be padding; otherwise every row must equal the first-row label or be padding. It is the same amount of code and removes the false positives without ever accepting a silently-wrong layout.
🤖 v4
| def test_m_grouped_gemm_contiguous_labels_contract_rejection() -> None: | ||
| print('Testing m-grouped contiguous labels contract rejection:') | ||
|
|
||
| saved_env = os.environ.get('DG_CHECK_CONTIGUOUS_LABELS') |
There was a problem hiding this comment.
🔵 suggestion: The bf16 rejection test only exercises the alignment branch; the resume-after-padding and out-of-range branches are only covered through the fp8/fp4 entry point. Because check_contiguous_labels_contract is shared this is acceptable, but parametrizing the fp8 rejection test over both entry points (or adding the resume/range cases here) would guard the bf16 call site and the shared helper together.
🤖 v4
| else: | ||
| # Unsupported architectures fail after the contract check, at kernel dispatch; | ||
| # a supported architecture must run to completion and is checked above | ||
| assert get_arch_major() == 12, raised |
There was a problem hiding this comment.
🔵 suggestion: The dispatch-failure branch hardcodes get_arch_major() == 12. That is the only major without an m-grouped contiguous dispatch on this tree, but the assertion will be wrong for a future unsupported major (or if a kernel is removed). Consider allowing any major that has no dispatch branch, or add a short comment tying 12 to 'main has no sm12x dispatch by design'.
🤖 v4
| - `DG_JIT_DUMP_SASS`: `0` or `1`, dump SASS output, `0` by default | ||
| - `DG_COMM_KERNEL_DEBUG`: `0` or `1`, zero symmetric buffer before each Mega MoE call for debugging, `0` by default | ||
| - `DG_USE_NVIDIA_TOOLS`: `0` or `1`, skip internal profiling when running under external NVIDIA tools, `0` by default | ||
| - `DG_CHECK_CONTIGUOUS_LABELS`: `0` or `1`, validate on the host that m-grouped contiguous labels respect the mk-alignment layout contract, turning silent miscomputation into a loud error (one GPU-to-CPU copy per call, meant for debugging/integration), `0` by default |
There was a problem hiding this comment.
🔵 suggestion: The variable is listed under the section introduced by 'Each DG_JIT_* variable falls back to the corresponding global DJ_JIT_* variable when unset', but DG_CHECK_CONTIGUOUS_LABELS is not a DG_JIT_* variable and has no DJ_ fallback. Other non-JIT debug vars are already listed there, so this is only a wording nit.
🤖 v4
🤖 ds-review-bot Code Reviewv6未发现新的可操作缺陷，历史报告的漏检和受支持架构测试吞异常问题均已修复。Python 语法检查及主机侧规则模型验证通过；当前环境缺少 PyTorch/CUDA，未运行 GPU 测试。 v5v7 ( All findings from v4–v6 are resolved in this revision and were verified against the kernel/heuristics code:
Independent check of the fixed rule (reject iff Other details confirmed: call sites sit in the labels-only branch after the Verdict: Approve. Only two non-blocking nits remain (see comments). Tests were not executed here (no GPU/PyTorch in this environment); conclusions are from code inspection, and the author reports passing runs on sm_103 and sm_120. v4ba1bcf5 adds an opt-in (DG_CHECK_CONTIGUOUS_LABELS=1) host-side contract checker for labels-mode m-grouped contiguous GEMM plus tests, with no device-code change. I verified the two v6 warnings are genuinely fixed: (1) prev_label now tracks the previous ROW (padding included), so a non-padding run that resumes after padding at an unaligned row is rejected -- both regression patterns ([0]*128+[-1]*3+[0]*125 and [0]*128+[1]*128+[-1]+[1]*127 at alignment 128) are caught; (2) the positive half no longer swallows unexpected errors: it only tolerates a dispatch failure on arch 12 and requires a numerical match on supported architectures. The alignment>0 guard, README entry, len(actual_ms_override) assert, relaxed label < num_groups range check, and the bf16 rejection test are all present. The checker has no false negatives (it rejects every layout the per-tile scheduler would miscompute) and is inert unless the env var is set. Remaining items are non-blocking. Files reviewed: 5 |
2 participants
Summary
Add an opt-in host-side contract checker for labels-mode m-grouped contiguous GEMM (
m_grouped_fp8_fp4_gemm_nt_contiguous,m_grouped_bf16_gemm_nt_contiguous, and thennforwarders), turning a silent wrong-results contract violation into a loud, actionable error.The contract and the silent failure
The contiguous kernels select B (and SFB) per
BLOCK_Mtile from the label of the tile's first row (grouped_layout[m_block_idx * BLOCK_M]). The tile heuristics guaranteeBLOCK_Mdivides the runtimemk_alignment_for_contiguous_layout, so labels whose group boundaries are aligned to the runtime alignment can never put two groups inside one tile.A caller that builds labels at a finer granularity than the current runtime alignment (e.g. labels packed at 64 while the runtime default is 128) puts group boundaries inside a tile. The tile then computes all its rows with the first row's group B — silently corrupting every group after the first misaligned boundary. This was reported downstream as "an empty middle group corrupts later groups" (an empty group merely makes a misaligned boundary much more likely; no empty group is required —
[128, 130, 65]built at 64 with runtime 128 corrupts the same way).This is a caller-contract violation, and today it fails silently on every arch (sm90/sm100/sm12x all select B per tile from the first row's label).
What this PR does
csrc/apis/gemm.hpp:check_contiguous_labels_contract()— whenDG_CHECK_CONTIGUOUS_LABELSis set, validates on the host (one GPU→CPU copy of the labels) that every group's first row starts at a multiple of the runtime mk alignment, and that labels are in range; violations raise aDGExceptionnaming the group, row, and alignment, with the remedy. Opt-in: zero cost in production.tests/generators.py:generate_m_grouped_contiguousgains an optionalactual_ms_override(default path byte-identical).tests/test_fp8_fp4.py,tests/test_bf16.py:test_m_grouped_gemm_contiguous_middle_empty_groups— empty groups in the middle followed by non-empty groups (e.g.[100, 0, 130, 65]), the previously uncovered pattern class, verified contract-respecting;test_m_grouped_gemm_contiguous_labels_contract_rejection— misaligned labels raise loudly with the checker on, and pass with the checker off plus matching alignment.Validation
test_m_grouped_gemm_contiguous(fp8_fp4 + bf16) andtest_m_grouped_gemm_maskedsuites pass unchanged (the checker is inert unless enabled).