feat(determinism): grouped gemm dsrelu deterministic dprob - #521
Merged
Anerudhan merged 15 commits intoAug 17, 2026
Merged
Conversation
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThe PR adds deterministic ChangesGrouped GEMM dsReLU determinism
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant grouped_gemm_dsrelu_wrapper_sm100
participant BlockScaledMoEGroupedGemmQuantBwdKernel
participant dprob_slots
participant dprob_output
Caller->>grouped_gemm_dsrelu_wrapper_sm100: enable deterministic dprob
grouped_gemm_dsrelu_wrapper_sm100->>BlockScaledMoEGroupedGemmQuantBwdKernel: launch deterministic backward kernel
BlockScaledMoEGroupedGemmQuantBwdKernel->>dprob_slots: write per-N-tile partial results
grouped_gemm_dsrelu_wrapper_sm100->>dprob_output: reduce slots in canonical order
dprob_output-->>Caller: return public dprob tensor
Suggested labels: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 7
🧹 Nitpick comments (1)
test/python/fe_api/grouped_gemm/test_grouped_gemm_dsrelu.py (1)
393-426: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider covering
use_dsrelu_reuse=Truein the deterministic tests.The wrapper documents
use_dsrelu_reuseas "Reuse relu(C)^2 between d_srelu and dprob". That option shares intermediate results with the dprob path, so it can interact with the new per-subtile parking and fixed-order reduction. The helper accepts**wrapper_kwargs, so adding one parametrization is cheap.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@test/python/fe_api/grouped_gemm/test_grouped_gemm_dsrelu.py` around lines 393 - 426, Extend the deterministic grouped GEMM test cases that call _run_dsrelu_case to include a use_dsrelu_reuse=True parametrization. Reuse the existing **wrapper_kwargs pass-through and assertions so both reuse-enabled repeated runs and the default path verify deterministic results.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/fe-oss-apis/gemm_fusions/grouped_gemm_dsrelu.md`:
- Around line 319-320: The documentation around the kernel output determinism
statement incorrectly excludes dbias. Update the text to state that the
reproducibility guarantee applies only to dprob when generate_dbias is enabled,
or document deterministic dbias handling if the implementation supports it.
- Around line 347-349: The tensor parameter documentation must distinguish
wrapper and class APIs: update the class API dprob/dprob_tensor shape entry to
explicitly use one slot per N-tile with the deterministic shape, while retaining
`(valid_m, 1, 1)` for the wrapper API. Ensure the class entry also states that
reduction over dimension 1 is performed by the caller.
In `@python/cudnn/gemm/cutedsl/grouped/dsrelu/api.py`:
- Line 1745: Update execute() so _reduce_dprob_slots() runs under the supplied
current_stream context, ensuring its torch.sum and add_ operations wait for the
compiled kernel writes; alternatively establish an equivalent event dependency
before the reduction. Add a CUDA test using a non-current stream to verify the
reduction observes completed kernel output.
- Line 1555: Update the full-dynamic branch in the cache-key construction around
dynamic_m_tensor_signature so it also uses dprob_tensor.shape[1:] metadata,
matching the non-full-dynamic path instead of dynamic_tensor_signature. Add a
cache-reuse test covering n_out=512 and n_out=513, verifying distinct dprob slot
layouts do not reuse an incompatible compiled kernel.
In `@test/python/fe_api/grouped_gemm/test_grouped_gemm_dsrelu.py`:
- Around line 539-552: The deterministic DSReLU tests need architecture and
backend-version gates before invoking the SM100 wrapper. In
test/python/fe_api/grouped_gemm/test_grouped_gemm_dsrelu.py lines 539-552, add
torch.cuda.get_device_capability() and cudnn.backend_version() checks after the
import guard, and limit the ValueError/NotImplementedError handling to only the
baseline _run_dsrelu_case call. Apply the same capability and backend-version
gates after the import guard at lines 571-575.
- Around line 471-479: Update _torch_deterministic_algorithms to capture the
prior warn-only setting with
torch.is_deterministic_algorithms_warn_only_enabled() alongside the determinism
flag, then pass that captured value when restoring
torch.use_deterministic_algorithms in the finally block.
In `@test/python/test_utils.py`:
- Around line 22-33: Update assert_bitwise_runs to validate that repeats is at
least 2 before constructing or indexing runs, raising a clear assertion or
validation error when fewer runs are requested; preserve the existing bitwise
comparison behavior for valid repeat counts.
---
Nitpick comments:
In `@test/python/fe_api/grouped_gemm/test_grouped_gemm_dsrelu.py`:
- Around line 393-426: Extend the deterministic grouped GEMM test cases that
call _run_dsrelu_case to include a use_dsrelu_reuse=True parametrization. Reuse
the existing **wrapper_kwargs pass-through and assertions so both reuse-enabled
repeated runs and the default path verify deterministic results.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 1c7c7511-c670-4811-957c-d9296c07101b
📒 Files selected for processing (7)
docs/fe-oss-apis/gemm_fusions/grouped_gemm_dsrelu.mddocs/fe-oss-apis/overview.mdpython/cudnn/gemm/cutedsl/grouped/dsrelu/api.pypython/cudnn/gemm/cutedsl/grouped/dsrelu/moe_blockscaled_grouped_gemm_dsrelu_quant.pytest/python/fe_api/grouped_gemm/test_grouped_gemm_dsrelu.pytest/python/linear_attention/common.pytest/python/test_utils.py
Anerudhan
self-requested a review
August 11, 2026 02:36
Anerudhan
added
orig-nv-eng
Collaborator
|
@cudnn-ci-bot oss |
|
cuDNN CI bot commands
Only allowlisted maintainers can use |
Collaborator
|
@cudnn-ci-bot run oss |
ZhiyuLi-Nvidia
force-pushed
the
zhiyul/grouped-gemm-dsrelu-deterministic-dprob
branch
from
August 14, 2026 09:41
a8fd219 to
b39a5b4
Compare
dprob is the only output of this kernel that is not reproducible run to run.
Every other output writes each element once; dprob is a float reduction, and it
is non-deterministic at two levels:
1. Within a CTA the N-subtile loop runs forward or reversed depending on the
accumulator pipeline phase, so a running fp32 sum is over an order that
varies between runs.
2. Across CTAs every N-tile atomically accumulates into the same dprob[token],
so the summation order follows tile scheduling.
deterministic=True fixes both, and neither fix is sufficient alone. Each
subtile's partial is parked in a slot indexed by the actual subtile and summed
canonically after the loop; dprob gets one slot per N-tile so every
(token, tile_n) pair has exactly one writer, and the wrapper reduces those slots
in fixed order.
Default is off, so the shipped path is untouched: _dprob_n_slots returns 1 when
deterministic=False, leaving shapes, cache-key signatures and the accumulation
code identical to before. The extra grid_n dprob workspace, the reduction and
the subtile_cnt registers are only paid for when asked, and the two modes cache
separately.
dprob_tensor keeps its (valid_m, 1, 1) shape either way -- the workspace is
internal to the wrapper -- and the reduction accumulates into it rather than
overwriting, matching what the kernel's atomic does on the other path. The class
API is the low-level form and takes the wide tensor directly.
The reduction runs inside the kernel, so no process-level determinism setting
reaches it; CUDNN_FE_GROUPED_GEMM_DSRELU_DETERMINISTIC exists so the behaviour
can still be selected where the call site is not editable.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…rete mode The determinism tests so far only compared the deterministic path against itself and against the default path. Self-consistency is not correctness: a partial dropped by a gap in the per-subtile slots would be stably wrong, and comparing two runs of the same kernel cannot see it. Run the deterministic outputs through check_ref_grouped_gemm_dsrelu, the same reference the default wrapper tests use, which covers the whole backward output set -- dprob, dA row/col, d_srelu, amax and the scale factors. Also add a discrete-weight case. Discrete mode builds its cache key on a separate branch, so the dprob slot-count signature there was previously uncovered; the dense tests could not have caught a mistake in it. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Both were reported green while covering nothing. The fp8 cases never reached the kernel: the helper omitted norm_const_tensor, and the fp8 path allocates sfd_row/sfd_col, so the API rejected the descriptor set with "must be all None or all not None". That surfaced as an Unsupported testcase skip rather than a failure, because the helper turns ValueError into pytest.skip. The env-var test used the uint8 raw-FP4 config, which the shared allocators skip wherever CUTLASS is unavailable. Moved it to the fp8 config so it actually exercises the environment opt-in. On a GB300 this takes the deterministic coverage from 3 passed / 3 skipped to 5 passed / 1 skipped; the remaining skip is the uint8 raw-FP4 path, which fails identically with and without this branch's changes. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
… default The slot count was ceil_div(n, TILE_N), but MoEPersistentTileScheduler counts whole clusters and then expands to CTAs, so tile_n_idx reaches ceil_div(n, TILE_N * cluster_n) * cluster_n - 1. Those agree only at cluster_n == 1, which every current test config happens to use; check_support accepts cluster_shape_mn[1] up to 4, and any of those would have indexed past the end of the workspace. Derived from the scheduler's own formula now, with _resolve_cluster_shape_mn shared by the class and the wrapper so the default cannot drift between them. Drop CUDNN_FE_GROUPED_GEMM_DSRELU_DETERMINISTIC in favour of torch.are_deterministic_algorithms_enabled(). csa/compressor/api.py already established that FE OSS kernels react to torch's global flag; it can only reject, because it has no deterministic path, while this kernel can switch. A per-kernel env var would also have forced anyone wanting a deterministic MoE backward to set one name per kernel. Explicit deterministic=True still works independently of the torch setting. Also from the cleanup pass: - the wrapper allocated a (M,1,1) buffer, zeroed it, and added the reduction onto it -- three launches to add zero. One reduction now, in _reduce_dprob_slots, shared with the zero-token early return so that path returns the documented shape. - the per-subtile slot array is only built when overlapping_accum is on, which is the only thing that reverses the loop; other tile shapes keep the plain running sum, which is already in canonical order. - one store site for both dprob producers instead of two copies, and the cache key reads the slot count off the tensor rather than recomputing it. - tests: assert_bitwise_runs (8 repeats, compares bit patterns, so a +0.0/-0.0 flip cannot pass) replaces a hand-rolled 2-run torch.equal; promoted it from linear_attention/common.py to test_utils.py and re-exported. The env-var test became a torch-flag check inside the main test, and the dense and discrete tests share one assertion helper. - d_col is excluded from the unchanged-outputs check unless the fp8 scale-factor path is active. The kernel does not write it otherwise, so the old assertion compared uninitialized memory and failed depending on what ran before it. - the two-level rationale lives in the docs page now, with pointers from the docstrings instead of three near-verbatim copies; overview.md gained a determinism convention entry. GB300, cuDNN 9.25: determinism tests 3 passed / 1 skipped (the skip is the uint8 raw-FP4 config, which fails identically without this branch). Whole dsrelu file at L0 gives a failure set byte-identical to the same file built from 96548b4 in the same container. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Two numbers that were sitting in the code without saying where they came from. _resolve_cluster_shape_mn tested mma_tiler_mn[0] == 256 directly. 256 is not arbitrary -- it is the 2-CTA MMA tile, which is why the default cluster is (2, 1): both CTAs of the pair have to sit in one cluster along M. The code it replaced said that in words, via use_2cta_instrs, and lifting it to module scope dropped the name. _uses_2cta_instrs restores it, and __init__ now shares it instead of repeating the literal. The dprob closeness check used rtol=1e-4 with no derivation. The two modes sum the same n terms in a different order, so the gap between them is reordering error alone, bounded by (n-1) * u for u = 2**-24; the tolerance is 4x that. At the default n=512 this is 1.2e-4, which is what the hardcoded value was -- but it now tracks n, which is settable with --grouped-gemm-nkl and would have false-failed a wider problem. atol stays separate and scaled by the largest dprob, because rtol says nothing useful about an entry that cancelled to near zero, whose absolute error is bounded by the summands rather than the result. GB300: determinism tests 3 passed / 1 skipped, whole dsrelu file byte-identical to the 96548b4 control. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
tcgen05 offers exactly two MMA forms -- 1-SM (CtaGroup.ONE, M=128) and 2-SM (CtaGroup.TWO, M=256, two CTAs of a cluster cooperating on one MMA) -- so the M tile and the CTA group carry the same information, which is why check_support asserts the pairing both ways and requires an even cluster_shape_mn[0]. Also note what the docstring cannot show from one line: the neighbouring 128s and 256s are unrelated constants that merely collide in value -- the N tile, the scale-factor layout atom, FIX_PAD_SIZE, and byte alignments all move independently. Tying them to a single shared constant would couple things that are free to diverge. Docstring only; no executable change. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…ot count Two real defects from PR review. The reduction ran on torch's current stream while execute() ran the kernel on the caller's `current_stream`. A caller passing any other stream could have the sum read dprob_tensor before the kernel finished writing it. It now runs inside a stream context, the same shape csa/compressor and deepseek_sparse_attention already use (kept local rather than imported so a GEMM kernel does not depend on an attention module). Every existing test passed torch's own stream, so none of them could see this; the new side-stream test does not enter torch.cuda.stream(), because inside that context the two streams are the same one and the bug cannot appear. The dprob slot count was missing from the cache key on the use_full_dynamic path, which drops tensor shapes -- yet the slot extent is static in the compiled descriptor. n=512 and n=768 give 2 and 3 slots with identical stride order, so the second call reused a kernel built for the wrong extent. Added as a scalar rather than by widening the tensor signature: dprob's strides include valid_m * slots, so a shape-derived signature would change on every new M and thrash the cache. Docs said dprob was the only non-reproducible output. Wrong when generate_dbias=True -- dbias is atomically accumulated across CTA tiles too (moe_blockscaled_grouped_gemm_dsrelu_quant.py:1044). Scoped the claim and said so plainly. Also split the dprob shape entry, which quoted (valid_m, 1, 1) for both APIs though the class API takes one slot per N-tile under deterministic. Smaller: restore warn_only when unwinding torch's determinism flag instead of forcing it on; require repeats >= 2 in assert_bitwise_runs, below which the comparison loop is empty and the assertion passes vacuously; and narrow the skip-on-unsupported guard to the baseline probe so a failure of the deterministic run cannot come back as a green skip. GB300: 5 passed / 1 skipped (the uint8 raw-FP4 config, which fails identically without this branch), whole dsrelu file byte-identical to the 96548b4 control. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
It was a one-line comparison behind a function call carrying a ten-line docstring, most of which was background rather than anything needed to read the code. _resolve_cluster_shape_mn keeps its place -- the wrapper and the class constructor both need the default and the wrapper cannot reach self, so one copy is what stops them drifting -- but the predicate is now a named local inside it. Also reverts the __init__ line to what it was: routing a pre-existing assignment through the new helper widened the diff without changing behaviour. Ten lines shorter, no behaviour change. GB300: 5 passed / 1 skipped, whole dsrelu file byte-identical to the 96548b4 control. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The derived bound and its twelve lines of justification were more machinery than the check needs. The tolerance is not delicate: both modes sum the same terms, so the honest difference is fp32 reordering, around 1e-5 at this problem size, while the failure being guarded against -- a partial dropped by a gap in the per-subtile slots, or double-counted by a second writer -- moves dprob by tens of percent. Anything between those catches it without false alarms. Also drops the max(scale, 1.0) floor on atol, which only ever made the bound looser for small-magnitude tensors and so weakened the assertion. GB300: 5 passed / 1 skipped, whole dsrelu file byte-identical to the 96548b4 control. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…t the kernel Three things the dprob work left behind, all reachable only from configurations the 120B recipe happens not to use -- so none of this changes that run. Initialise the accumulators on the caller's stream. dprob, dbias and amax are atomic read-modify-write targets, so their initial value has to be visible to the kernel. Allocating them on torch's current stream puts that memset on a stream unordered against a caller-supplied current_stream: the kernel could then accumulate onto uninitialised memory and deterministic=True would silently stop being deterministic. They are now allocated inside the caller's stream, with torch.zeros/torch.full and no separate fill step -- use_deterministic_algorithms (True) makes torch fill every user-visible torch.empty with NaN, so anything built on torch.empty pays a second pass over the buffer on exactly the path this serves, and on dprob's grid_n-times-larger workspace. That trades some allocator-pool locality for one fewer pass; the pool cost is a guess, the extra pass is not. The write-only outputs still come from torch's stream and so get record_stream, the pattern csa/compressor and the sibling grouped-GEMM APIs already use. The stream context itself is now imported from grouped/backend_utils rather than redefined -- the shared one already fast-paths the common case where the caller passes torch's own stream. Cover dbias. It contends across M-tiles rather than dprob's N-tiles, so it gets its own fp32 workspace keyed by absolute M-block -- the index the SFD-col path already builds as global_sfd_m -- plus a per-expert segment sum. That sum is a one-hot matmul: index_add_/scatter_add_ are non-deterministic on CUDA, and slicing per expert would need padded_offsets on the host, a sync in the training loop. Keeping fp32 partials and narrowing once also makes this path more accurate than the default, which rounds to bf16 on every tile. check_support now requires m_aligned % (cta_tile_m * cluster_m) == 0, without which the scheduler emits tiles past an expert's range whose slots alias the next expert's. This reverses the conclusion of a resolved review thread on the dbias paragraph of grouped_gemm_dsrelu.md, which settled on "dbias remains scheduling-dependent when generate_dbias=True". That was accurate for the code as it stood; the doc is updated rather than left contradicting the implementation. Refuse deterministic=True under jax. The slot reductions are torch ops on torch streams, and jax has no use_deterministic_algorithms global to inherit from, so the default resolves to False there and an explicit request raises instead of going half-enforced. Tests: a forced side-stream regression that poisons a same-sized allocator block and occupies torch's stream, so a misplaced memset is caught rather than merely maybe-caught; deterministic dbias against the default path and across repeats; and a segment check that compares each expert's dbias to a column sum of the kernel's own d_row, which a wrong one-hot would fail while still looking finite. Not executed: this machine has no GPU or torch. Formatting and syntax only.
Job 459169 failed all three dbias tests. Two of the three were the tests, not the kernel. The segment test compared dbias to a column sum of d_row. In that config d_dtype is float8_e4m3fn, so d_row is quantized and carries its scale factors separately, while dbias sums the pre-quantization values -- 94% of elements mismatched, which is what summing scaled against unscaled looks like. run_grouped_gemm_dsrelu_ref already builds the right thing, ref_dbias, per expert from the unquantized dA; check_ref_grouped_gemm_dsrelu asked for it and then never asserted on it. Assert it there, so every dbias test gets the segment mapping checked rather than just this one, and reduce the segment test to running two n values through it. The det-vs-non-det comparison used a flat atol of 2e-2 on values whose magnitude is group_m times d_row's. It failed on a 16.0 difference, about 2 ULP of bf16 at that scale, and on a relative inf where the bf16 baseline had rounded to zero. Scale the bound to the values involved. The two modes genuinely differ: the default rounds to bf16 on every M-tile, the deterministic path keeps fp32 partials and narrows once, so it is the more accurate of the pair and pinning it to the less accurate one was the wrong assertion. Correctness is the reference check's job. Whether the deterministic dbias values are right is still open -- these tests could not answer it before. 24 compile_execute_fp4 failures in the same run are unrelated (dbias_tensor: None, "expected dtype=float4_e2m1fnx2") and are being classified by a control run at HEAD~1.
Second /simplify pass. One of these is a real regression, the rest are cleanups. The cache key. tensor_signature(dbias_kernel_tensor) put the workspace's dim 0 -- ceil_div(valid_m, cta_tile_m) -- into the compile cache key, so a deterministic dbias run would recompile on every distinct token count and grow the object cache without bound. Every other entry in that key is deliberately M-independent, and _make_dbias_fake already keeps the extent symbolic precisely so one kernel serves all valid_m; the key contradicted the fake. Keyed like dprob now: M dynamic, the rest static. The tests could not have caught it -- both dbias tests take the full-dynamic branch, which drops the shape entirely. _record_streams moves to grouped/backend_utils next to _torch_stream_context, with the handle-to-torch-stream resolution factored into one _resolve_torch_stream both use. It was a sixth copy of that resolution inside grouped/, and the two copies living in one file had already drifted: backend_utils mapped a default-stream handle to torch's default stream, the local one wrapped it in an ExternalStream. The class API's deterministic dbias contract existed only in the .md and in validation error strings. sample_dbias's docstring still promised (expert_cnt, n, 1) bf16, which is wrong under the flag, and deterministic's still described dprob alone. Both now say what the caller has to pass and that a plain sum will not reduce it. Smaller: _make_dbias_fake uses APIBase's dynamic_mode parameter instead of splicing sym_int into the shape by hand; the absolute-M-block index is computed once instead of twice with a comment pointing at the other copy; the same fact stated in three places in the kernel is stated once; the dbias test reuses _assert_dprob_deterministic, which restores the "other outputs unchanged" check the hand-rolled copy had dropped; the segments test drops its n=512 case, already covered, and its rationale, which was wrong -- block count follows group_m, not n. Skipped: replacing the two scalar fp32 atomics with a vector store (real, needs hardware measurement), rounding the workspace to a fixed size class to stabilise the allocator, and promoting the segment sum to a shared module (no second consumer yet).
Challenged the correctness claim rather than the code, and the tests turned out to be the weak part. The determinism tests were vacuous at their own config. Measured directly (job 466159, 16 launches per config): at l=4 / [256] * 4 / n=512 -- what every determinism test in this file used -- the NON-deterministic dprob and dbias are both already bit-stable. assert_bitwise_runs could not fail there, so "passes" meant nothing. It varies 15/15 for dbias at a ragged distribution, and for both outputs at l=8 / [1024] * 8 / n=2048. The dbias test moves to the ragged config, and a new L1 test covers the one case where both outputs are genuinely unstable by default. This applies to the pre-existing dprob tests too, not only the new ones. The segment sum was only ever tested against equal-sized experts. group_m_list defaults to [256] * l, so every expert owns exactly two M-blocks and a mapping that assumed a fixed stride between experts would have passed everything. Added ragged distributions (2/4/2/8 and 4/2/8/2 blocks) plus a per-expert non-empty check, so a collapsed or off-by-one mapping cannot hide. _assert_dprob_deterministic's flat 1e-4 does not transfer across n. dprob sums n terms, so at n=2048 one element in 8192 lands at 1.8e-4 relative on reordering alone -- reordering, not a dropped partial, which moves dprob by tens of percent. The tolerance is now a parameter, default unchanged, loosened only by the at-scale caller. Also added: a zero-token deterministic dbias case, where the reduction runs on an empty workspace and an empty contraction; and a guard that deterministic dbias adds no token-count dependence to the compile cache key. That last one is asserted as "no more entries than the default path" rather than "exactly one" -- job 466362 diffed the keys and found this config already recompiles per valid_m without dbias at all, via the d_col batch stride and the sfd_col shape. That is pre-existing and worth its own look; what the test pins is that the flag adds nothing on top of it. 13 determinism tests pass at L0+L1 (job 466365). The 25 remaining failures in the file are the fp4 ones the HEAD~1 control already showed are inherited.
ZhiyuLi-Nvidia
force-pushed
the
zhiyul/grouped-gemm-dsrelu-deterministic-dprob
branch
from
August 14, 2026 16:49
b39a5b4 to
88c7fab
Compare
added 2 commits
August 14, 2026 12:08
Two changes, both prompted by asking why the deterministic API looked nothing like the default one. The slots did not need to be fp32. Reproducibility comes from each (M-block, n) having a single writer and from the fixed-order reduction, not from a wider accumulator, so fp32 bought accuracy and nothing else while costing double the workspace (4 MiB -> 2 MiB at valid_m=64k, n=2048) and splitting one packed bf16x2 store into two scalar atomics. bf16 slots still beat the default on accuracy -- there each M-tile's atomic rounds the running sum, here each slot rounds once and the segment matmul accumulates in fp32 -- and they make the two modes emit the same store instruction, so dbias_store and its branch are gone. Only the destination row differs now, which is what the docstring claimed all along and was not true until this commit. The workspaces are their own arguments. Before, `deterministic` changed sample_dprob's extent and sample_dbias's extent, dim-0 meaning and dtype, so the flag rewrote the type of arguments the caller supplies -- and for dbias the bridge back was a segment sum that is private, and wrong quietly when hand-rolled. sample_dprob is now (valid_m, 1, 1) float32 and sample_dbias (expert_cnt, n, 1) bf16 in both modes; the flag adds sample_dprob_workspace and sample_dbias_workspace (plus the execute() counterparts) and nothing else. That follows sdpa/bwd/api_dsl.py, where deterministic changes a scratch size the API reports rather than the caller's output contract. With it: dprob_workspace_shape, dbias_workspace_shape, reduce_dprob_workspace and reduce_dbias_workspace, public so nobody re-derives _cta_tile_m or reimplements the segment sum from prose. Tests: a class-API case, which was the gap flagged twice and never covered -- every other determinism test drives the wrapper, which hides all of this. It builds the op directly, reduces through the public helpers, and requires the result bit-identical to the wrapper's. 14 determinism tests pass at L0+L1 (job 468385); the 25 remaining failures in the file are the fp4 ones the HEAD~1 control showed are inherited.
…tics All three findings hold. One needed correcting on severity, and fixing the second made a chunk of this branch unnecessary. 1. The dbias cache key drops n on the default path. Confirmed: under full dynamic every tensor drops its shape because the kernel compiles with symbolic extents, but _make_dbias_fake frees only dim 0, so dbias's n stays baked, and nothing else in the dense key carries n -- b_tensor.shape[2] is l, and dprob_n_slots is too coarse (n=384 and n=512 both give 2). dbias_signature is now used on both arms. The review's severity was wrong, though: it does not write dbias wrong. CuTeDSL validates the compiled descriptor's static extents at call time, so the second n raises "Mismatched dbias_tensor.shape[1] ... expected to be 384" (job 469460). Fail-loud, not silent corruption. Pre-existing on the default path -- reproduced with deterministic=False -- so the new test is parametrized over both modes. 2. The write-only outputs were still allocated on torch's stream while the kernel writes them on the caller's. record_stream does not cover that direction: it defers reuse after *this* tensor is freed, and says nothing about whether the block was safe to hand out. torch's allocator recycles a block freed on the allocating stream on the strength of that stream's ordering, which does not constrain a kernel on current_stream. d_row/d_col/d_srelu and the sfd tensors now allocate inside the stream context, like the accumulators. That makes _record_streams dead: nothing is allocated on torch's stream any more, so recording the caller's stream is a no-op. Removed, and with its only caller gone it comes back out of grouped/backend_utils -- that shared module is now untouched by this branch, 0 lines changed against upstream. 3. The two reductions disagreed. The kernel atomically *adds* into dbias when the flag is off, so a caller accumulating across micro-batches got accumulate semantics one way and overwrite the other. _reduce_dbias_slots now accumulates like _reduce_dprob_slots, reduce_dbias_workspace takes the output tensor for symmetry with reduce_dprob_workspace, and the class-API test drops the copy_ that was papering over it. 15 determinism tests pass, both arms of the new n-collision test pass, and the 25 remaining failures are the fp4 ones the HEAD~1 control showed are inherited (job 469466).
Contributor
Author
|
from @yanqinz2 PR #521 review — top 3 findings
Line numbers are from the PR head and will drift. 1. dbias cache key drops
|
Contributor
Author
Confirmed and fixed — dbias_signature is now used on both arms of the key. One correction on severity: it doesn't write dbias wrong. CuTeDSL validates the compiled descriptor's static extents at call time, so the second n raises Mismatched dbias_tensor.shape[1] ... expected to be 384 rather than corrupting silently. Fail-loud, not silent. Also confirmed pre-existing on the default path — it reproduces with deterministic=False — so the new regression test is parametrized over both modes: two n values sharing a dprob_n_slots must both come back correct.
Confirmed and fixed — d_row/d_col/d_srelu and the sfd_* tensors now allocate inside _torch_stream_context, like the accumulators. Agreed on the mechanism: record_stream can't cover this direction, since it defers reuse after this tensor is freed and says nothing about whether the block was safe to hand out. Follow-on: with nothing left allocated on torch's stream, _record_streams became dead and is removed. With its only caller gone it also comes back out of grouped/backend_utils.py — that shared module is now untouched by this PR, 0 lines against upstream.
Confirmed and fixed — _reduce_dbias_slots now accumulates like reduce_dprob_slots, matching the atomic the kernel uses when the flag is off. reduce_dbias_workspace takes the output tensor for symmetry, and the class-API test drops the copy that was masking it. All three verified on hardware: 15 determinism tests pass. |
vedaanta
approved these changes
Aug 14, 2026
Anerudhan
approved these changes
Aug 17, 2026
Collaborator
|
@cudnn-ci-bot run oss |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-521-f1c3623 |
hxbai
pushed a commit
to hxbai/cudnn-frontend
that referenced
this pull request
Aug 18, 2026
* Add opt-in deterministic dprob to the grouped GEMM dsReLU backward
dprob is the only output of this kernel that is not reproducible run to run.
Every other output writes each element once; dprob is a float reduction, and it
is non-deterministic at two levels:
1. Within a CTA the N-subtile loop runs forward or reversed depending on the
accumulator pipeline phase, so a running fp32 sum is over an order that
varies between runs.
2. Across CTAs every N-tile atomically accumulates into the same dprob[token],
so the summation order follows tile scheduling.
deterministic=True fixes both, and neither fix is sufficient alone. Each
subtile's partial is parked in a slot indexed by the actual subtile and summed
canonically after the loop; dprob gets one slot per N-tile so every
(token, tile_n) pair has exactly one writer, and the wrapper reduces those slots
in fixed order.
Default is off, so the shipped path is untouched: _dprob_n_slots returns 1 when
deterministic=False, leaving shapes, cache-key signatures and the accumulation
code identical to before. The extra grid_n dprob workspace, the reduction and
the subtile_cnt registers are only paid for when asked, and the two modes cache
separately.
dprob_tensor keeps its (valid_m, 1, 1) shape either way -- the workspace is
internal to the wrapper -- and the reduction accumulates into it rather than
overwriting, matching what the kernel's atomic does on the other path. The class
API is the low-level form and takes the wide tensor directly.
The reduction runs inside the kernel, so no process-level determinism setting
reaches it; CUDNN_FE_GROUPED_GEMM_DSRELU_DETERMINISTIC exists so the behaviour
can still be selected where the call site is not editable.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* Check the deterministic dprob path against the reference, and in discrete mode
The determinism tests so far only compared the deterministic path against itself
and against the default path. Self-consistency is not correctness: a partial
dropped by a gap in the per-subtile slots would be stably wrong, and comparing
two runs of the same kernel cannot see it.
Run the deterministic outputs through check_ref_grouped_gemm_dsrelu, the same
reference the default wrapper tests use, which covers the whole backward output
set -- dprob, dA row/col, d_srelu, amax and the scale factors.
Also add a discrete-weight case. Discrete mode builds its cache key on a
separate branch, so the dprob slot-count signature there was previously
uncovered; the dense tests could not have caught a mistake in it.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* Fix two determinism tests that were skipping instead of running
Both were reported green while covering nothing.
The fp8 cases never reached the kernel: the helper omitted norm_const_tensor,
and the fp8 path allocates sfd_row/sfd_col, so the API rejected the descriptor
set with "must be all None or all not None". That surfaced as an Unsupported
testcase skip rather than a failure, because the helper turns ValueError into
pytest.skip.
The env-var test used the uint8 raw-FP4 config, which the shared allocators skip
wherever CUTLASS is unavailable. Moved it to the fp8 config so it actually
exercises the environment opt-in.
On a GB300 this takes the deterministic coverage from 3 passed / 3 skipped to
5 passed / 1 skipped; the remaining skip is the uint8 raw-FP4 path, which fails
identically with and without this branch's changes.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* Size the dprob workspace from the scheduler, and follow torch for the default
The slot count was ceil_div(n, TILE_N), but MoEPersistentTileScheduler counts
whole clusters and then expands to CTAs, so tile_n_idx reaches
ceil_div(n, TILE_N * cluster_n) * cluster_n - 1. Those agree only at
cluster_n == 1, which every current test config happens to use; check_support
accepts cluster_shape_mn[1] up to 4, and any of those would have indexed past
the end of the workspace. Derived from the scheduler's own formula now, with
_resolve_cluster_shape_mn shared by the class and the wrapper so the default
cannot drift between them.
Drop CUDNN_FE_GROUPED_GEMM_DSRELU_DETERMINISTIC in favour of
torch.are_deterministic_algorithms_enabled(). csa/compressor/api.py already
established that FE OSS kernels react to torch's global flag; it can only
reject, because it has no deterministic path, while this kernel can switch. A
per-kernel env var would also have forced anyone wanting a deterministic MoE
backward to set one name per kernel. Explicit deterministic=True still works
independently of the torch setting.
Also from the cleanup pass:
- the wrapper allocated a (M,1,1) buffer, zeroed it, and added the reduction
onto it -- three launches to add zero. One reduction now, in
_reduce_dprob_slots, shared with the zero-token early return so that path
returns the documented shape.
- the per-subtile slot array is only built when overlapping_accum is on, which
is the only thing that reverses the loop; other tile shapes keep the plain
running sum, which is already in canonical order.
- one store site for both dprob producers instead of two copies, and the cache
key reads the slot count off the tensor rather than recomputing it.
- tests: assert_bitwise_runs (8 repeats, compares bit patterns, so a +0.0/-0.0
flip cannot pass) replaces a hand-rolled 2-run torch.equal; promoted it from
linear_attention/common.py to test_utils.py and re-exported. The env-var test
became a torch-flag check inside the main test, and the dense and discrete
tests share one assertion helper.
- d_col is excluded from the unchanged-outputs check unless the fp8
scale-factor path is active. The kernel does not write it otherwise, so the
old assertion compared uninitialized memory and failed depending on what ran
before it.
- the two-level rationale lives in the docs page now, with pointers from the
docstrings instead of three near-verbatim copies; overview.md gained a
determinism convention entry.
GB300, cuDNN 9.25: determinism tests 3 passed / 1 skipped (the skip is the
uint8 raw-FP4 config, which fails identically without this branch). Whole
dsrelu file at L0 gives a failure set byte-identical to the same file built
from 96548b4 in the same container.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* Name the 2-CTA tile check, and derive the reordering tolerance from n
Two numbers that were sitting in the code without saying where they came from.
_resolve_cluster_shape_mn tested mma_tiler_mn[0] == 256 directly. 256 is not
arbitrary -- it is the 2-CTA MMA tile, which is why the default cluster is
(2, 1): both CTAs of the pair have to sit in one cluster along M. The code it
replaced said that in words, via use_2cta_instrs, and lifting it to module
scope dropped the name. _uses_2cta_instrs restores it, and __init__ now shares
it instead of repeating the literal.
The dprob closeness check used rtol=1e-4 with no derivation. The two modes sum
the same n terms in a different order, so the gap between them is reordering
error alone, bounded by (n-1) * u for u = 2**-24; the tolerance is 4x that.
At the default n=512 this is 1.2e-4, which is what the hardcoded value was --
but it now tracks n, which is settable with --grouped-gemm-nkl and would have
false-failed a wider problem. atol stays separate and scaled by the largest
dprob, because rtol says nothing useful about an entry that cancelled to near
zero, whose absolute error is bounded by the summands rather than the result.
GB300: determinism tests 3 passed / 1 skipped, whole dsrelu file byte-identical
to the 96548b4 control.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* Record that the 2-CTA M tile is an ISA constant, not a tuning choice
tcgen05 offers exactly two MMA forms -- 1-SM (CtaGroup.ONE, M=128) and 2-SM
(CtaGroup.TWO, M=256, two CTAs of a cluster cooperating on one MMA) -- so the M
tile and the CTA group carry the same information, which is why check_support
asserts the pairing both ways and requires an even cluster_shape_mn[0].
Also note what the docstring cannot show from one line: the neighbouring 128s
and 256s are unrelated constants that merely collide in value -- the N tile,
the scale-factor layout atom, FIX_PAD_SIZE, and byte alignments all move
independently. Tying them to a single shared constant would couple things that
are free to diverge.
Docstring only; no executable change.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* Order the dprob reduction against the kernel's stream, and key the slot count
Two real defects from PR review.
The reduction ran on torch's current stream while execute() ran the kernel on
the caller's `current_stream`. A caller passing any other stream could have the
sum read dprob_tensor before the kernel finished writing it. It now runs inside
a stream context, the same shape csa/compressor and deepseek_sparse_attention
already use (kept local rather than imported so a GEMM kernel does not depend on
an attention module). Every existing test passed torch's own stream, so none of
them could see this; the new side-stream test does not enter
torch.cuda.stream(), because inside that context the two streams are the same
one and the bug cannot appear.
The dprob slot count was missing from the cache key on the use_full_dynamic
path, which drops tensor shapes -- yet the slot extent is static in the compiled
descriptor. n=512 and n=768 give 2 and 3 slots with identical stride order, so
the second call reused a kernel built for the wrong extent. Added as a scalar
rather than by widening the tensor signature: dprob's strides include
valid_m * slots, so a shape-derived signature would change on every new M and
thrash the cache.
Docs said dprob was the only non-reproducible output. Wrong when
generate_dbias=True -- dbias is atomically accumulated across CTA tiles too
(moe_blockscaled_grouped_gemm_dsrelu_quant.py:1044). Scoped the claim and said
so plainly. Also split the dprob shape entry, which quoted (valid_m, 1, 1) for
both APIs though the class API takes one slot per N-tile under deterministic.
Smaller: restore warn_only when unwinding torch's determinism flag instead of
forcing it on; require repeats >= 2 in assert_bitwise_runs, below which the
comparison loop is empty and the assertion passes vacuously; and narrow the
skip-on-unsupported guard to the baseline probe so a failure of the
deterministic run cannot come back as a green skip.
GB300: 5 passed / 1 skipped (the uint8 raw-FP4 config, which fails identically
without this branch), whole dsrelu file byte-identical to the 96548b4 control.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* Fold _uses_2cta_instrs into its only real caller
It was a one-line comparison behind a function call carrying a ten-line
docstring, most of which was background rather than anything needed to read the
code. _resolve_cluster_shape_mn keeps its place -- the wrapper and the class
constructor both need the default and the wrapper cannot reach self, so one copy
is what stops them drifting -- but the predicate is now a named local inside it.
Also reverts the __init__ line to what it was: routing a pre-existing assignment
through the new helper widened the diff without changing behaviour.
Ten lines shorter, no behaviour change. GB300: 5 passed / 1 skipped, whole
dsrelu file byte-identical to the 96548b4 control.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* Use a plain tolerance for the det-vs-non-det comparison
The derived bound and its twelve lines of justification were more machinery than
the check needs. The tolerance is not delicate: both modes sum the same terms,
so the honest difference is fp32 reordering, around 1e-5 at this problem size,
while the failure being guarded against -- a partial dropped by a gap in the
per-subtile slots, or double-counted by a second writer -- moves dprob by tens
of percent. Anything between those catches it without false alarms.
Also drops the max(scale, 1.0) floor on atol, which only ever made the bound
looser for small-magnitude tensors and so weakened the assertion.
GB300: 5 passed / 1 skipped, whole dsrelu file byte-identical to the 96548b4
control.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* Extend the determinism flag to dbias, and order the workspaces against the kernel
Three things the dprob work left behind, all reachable only from configurations
the 120B recipe happens not to use -- so none of this changes that run.
Initialise the accumulators on the caller's stream. dprob, dbias and amax are
atomic read-modify-write targets, so their initial value has to be visible to the
kernel. Allocating them on torch's current stream puts that memset on a stream
unordered against a caller-supplied current_stream: the kernel could then
accumulate onto uninitialised memory and deterministic=True would silently stop
being deterministic. They are now allocated inside the caller's stream, with
torch.zeros/torch.full and no separate fill step -- use_deterministic_algorithms
(True) makes torch fill every user-visible torch.empty with NaN, so anything built
on torch.empty pays a second pass over the buffer on exactly the path this serves,
and on dprob's grid_n-times-larger workspace. That trades some allocator-pool
locality for one fewer pass; the pool cost is a guess, the extra pass is not. The
write-only outputs still come from torch's stream and so get record_stream, the
pattern csa/compressor and the sibling grouped-GEMM APIs already use. The stream
context itself is now imported from grouped/backend_utils rather than redefined --
the shared one already fast-paths the common case where the caller passes torch's
own stream.
Cover dbias. It contends across M-tiles rather than dprob's N-tiles, so it gets
its own fp32 workspace keyed by absolute M-block -- the index the SFD-col path
already builds as global_sfd_m -- plus a per-expert segment sum. That sum is a
one-hot matmul: index_add_/scatter_add_ are non-deterministic on CUDA, and
slicing per expert would need padded_offsets on the host, a sync in the training
loop. Keeping fp32 partials and narrowing once also makes this path more accurate
than the default, which rounds to bf16 on every tile. check_support now requires
m_aligned % (cta_tile_m * cluster_m) == 0, without which the scheduler emits
tiles past an expert's range whose slots alias the next expert's.
This reverses the conclusion of a resolved review thread on the dbias paragraph
of grouped_gemm_dsrelu.md, which settled on "dbias remains scheduling-dependent
when generate_dbias=True". That was accurate for the code as it stood; the doc is
updated rather than left contradicting the implementation.
Refuse deterministic=True under jax. The slot reductions are torch ops on torch
streams, and jax has no use_deterministic_algorithms global to inherit from, so
the default resolves to False there and an explicit request raises instead of
going half-enforced.
Tests: a forced side-stream regression that poisons a same-sized allocator block
and occupies torch's stream, so a misplaced memset is caught rather than merely
maybe-caught; deterministic dbias against the default path and across repeats;
and a segment check that compares each expert's dbias to a column sum of the
kernel's own d_row, which a wrong one-hot would fail while still looking finite.
Not executed: this machine has no GPU or torch. Formatting and syntax only.
* Check dbias against the reference instead of against the other mode
Job 459169 failed all three dbias tests. Two of the three were the tests, not the
kernel.
The segment test compared dbias to a column sum of d_row. In that config d_dtype
is float8_e4m3fn, so d_row is quantized and carries its scale factors separately,
while dbias sums the pre-quantization values -- 94% of elements mismatched, which
is what summing scaled against unscaled looks like. run_grouped_gemm_dsrelu_ref
already builds the right thing, ref_dbias, per expert from the unquantized dA;
check_ref_grouped_gemm_dsrelu asked for it and then never asserted on it. Assert
it there, so every dbias test gets the segment mapping checked rather than just
this one, and reduce the segment test to running two n values through it.
The det-vs-non-det comparison used a flat atol of 2e-2 on values whose magnitude
is group_m times d_row's. It failed on a 16.0 difference, about 2 ULP of bf16 at
that scale, and on a relative inf where the bf16 baseline had rounded to zero.
Scale the bound to the values involved. The two modes genuinely differ: the
default rounds to bf16 on every M-tile, the deterministic path keeps fp32
partials and narrows once, so it is the more accurate of the pair and pinning it
to the less accurate one was the wrong assertion. Correctness is the reference
check's job.
Whether the deterministic dbias values are right is still open -- these tests
could not answer it before. 24 compile_execute_fp4 failures in the same run are
unrelated (dbias_tensor: None, "expected dtype=float4_e2m1fnx2") and are being
classified by a control run at HEAD~1.
* Key the dbias workspace dynamically, and share the stream helpers
Second /simplify pass. One of these is a real regression, the rest are cleanups.
The cache key. tensor_signature(dbias_kernel_tensor) put the workspace's dim 0 --
ceil_div(valid_m, cta_tile_m) -- into the compile cache key, so a deterministic
dbias run would recompile on every distinct token count and grow the object cache
without bound. Every other entry in that key is deliberately M-independent, and
_make_dbias_fake already keeps the extent symbolic precisely so one kernel serves
all valid_m; the key contradicted the fake. Keyed like dprob now: M dynamic, the
rest static. The tests could not have caught it -- both dbias tests take the
full-dynamic branch, which drops the shape entirely.
_record_streams moves to grouped/backend_utils next to _torch_stream_context,
with the handle-to-torch-stream resolution factored into one _resolve_torch_stream
both use. It was a sixth copy of that resolution inside grouped/, and the two
copies living in one file had already drifted: backend_utils mapped a
default-stream handle to torch's default stream, the local one wrapped it in an
ExternalStream.
The class API's deterministic dbias contract existed only in the .md and in
validation error strings. sample_dbias's docstring still promised
(expert_cnt, n, 1) bf16, which is wrong under the flag, and deterministic's still
described dprob alone. Both now say what the caller has to pass and that a plain
sum will not reduce it.
Smaller: _make_dbias_fake uses APIBase's dynamic_mode parameter instead of
splicing sym_int into the shape by hand; the absolute-M-block index is computed
once instead of twice with a comment pointing at the other copy; the same fact
stated in three places in the kernel is stated once; the dbias test reuses
_assert_dprob_deterministic, which restores the "other outputs unchanged" check
the hand-rolled copy had dropped; the segments test drops its n=512 case, already
covered, and its rationale, which was wrong -- block count follows group_m, not n.
Skipped: replacing the two scalar fp32 atomics with a vector store (real, needs
hardware measurement), rounding the workspace to a fixed size class to stabilise
the allocator, and promoting the segment sum to a shared module (no second
consumer yet).
* Make the determinism tests capable of failing
Challenged the correctness claim rather than the code, and the tests turned out
to be the weak part.
The determinism tests were vacuous at their own config. Measured directly (job
466159, 16 launches per config): at l=4 / [256] * 4 / n=512 -- what every
determinism test in this file used -- the NON-deterministic dprob and dbias are
both already bit-stable. assert_bitwise_runs could not fail there, so "passes"
meant nothing. It varies 15/15 for dbias at a ragged distribution, and for both
outputs at l=8 / [1024] * 8 / n=2048. The dbias test moves to the ragged config,
and a new L1 test covers the one case where both outputs are genuinely unstable
by default. This applies to the pre-existing dprob tests too, not only the new
ones.
The segment sum was only ever tested against equal-sized experts. group_m_list
defaults to [256] * l, so every expert owns exactly two M-blocks and a mapping
that assumed a fixed stride between experts would have passed everything. Added
ragged distributions (2/4/2/8 and 4/2/8/2 blocks) plus a per-expert non-empty
check, so a collapsed or off-by-one mapping cannot hide.
_assert_dprob_deterministic's flat 1e-4 does not transfer across n. dprob sums n
terms, so at n=2048 one element in 8192 lands at 1.8e-4 relative on reordering
alone -- reordering, not a dropped partial, which moves dprob by tens of percent.
The tolerance is now a parameter, default unchanged, loosened only by the
at-scale caller.
Also added: a zero-token deterministic dbias case, where the reduction runs on an
empty workspace and an empty contraction; and a guard that deterministic dbias
adds no token-count dependence to the compile cache key. That last one is
asserted as "no more entries than the default path" rather than "exactly one" --
job 466362 diffed the keys and found this config already recompiles per valid_m
without dbias at all, via the d_col batch stride and the sfd_col shape. That is
pre-existing and worth its own look; what the test pins is that the flag adds
nothing on top of it.
13 determinism tests pass at L0+L1 (job 466365). The 25 remaining failures in the
file are the fp4 ones the HEAD~1 control already showed are inherited.
* Keep the output args mode-invariant, and make the dbias slots bf16
Two changes, both prompted by asking why the deterministic API looked nothing
like the default one.
The slots did not need to be fp32. Reproducibility comes from each (M-block, n)
having a single writer and from the fixed-order reduction, not from a wider
accumulator, so fp32 bought accuracy and nothing else while costing double the
workspace (4 MiB -> 2 MiB at valid_m=64k, n=2048) and splitting one packed
bf16x2 store into two scalar atomics. bf16 slots still beat the default on
accuracy -- there each M-tile's atomic rounds the running sum, here each slot
rounds once and the segment matmul accumulates in fp32 -- and they make the two
modes emit the same store instruction, so dbias_store and its branch are gone.
Only the destination row differs now, which is what the docstring claimed all
along and was not true until this commit.
The workspaces are their own arguments. Before, `deterministic` changed
sample_dprob's extent and sample_dbias's extent, dim-0 meaning and dtype, so the
flag rewrote the type of arguments the caller supplies -- and for dbias the
bridge back was a segment sum that is private, and wrong quietly when hand-rolled.
sample_dprob is now (valid_m, 1, 1) float32 and sample_dbias (expert_cnt, n, 1)
bf16 in both modes; the flag adds sample_dprob_workspace and
sample_dbias_workspace (plus the execute() counterparts) and nothing else. That
follows sdpa/bwd/api_dsl.py, where deterministic changes a scratch size the API
reports rather than the caller's output contract.
With it: dprob_workspace_shape, dbias_workspace_shape, reduce_dprob_workspace and
reduce_dbias_workspace, public so nobody re-derives _cta_tile_m or reimplements
the segment sum from prose.
Tests: a class-API case, which was the gap flagged twice and never covered --
every other determinism test drives the wrapper, which hides all of this. It
builds the op directly, reduces through the public helpers, and requires the
result bit-identical to the wrapper's.
14 determinism tests pass at L0+L1 (job 468385); the 25 remaining failures in the
file are the fp4 ones the HEAD~1 control showed are inherited.
* Address PR review: dbias cache key, output alloc stream, reduce semantics
All three findings hold. One needed correcting on severity, and fixing the second
made a chunk of this branch unnecessary.
1. The dbias cache key drops n on the default path. Confirmed: under full dynamic
every tensor drops its shape because the kernel compiles with symbolic extents,
but _make_dbias_fake frees only dim 0, so dbias's n stays baked, and nothing else
in the dense key carries n -- b_tensor.shape[2] is l, and dprob_n_slots is too
coarse (n=384 and n=512 both give 2). dbias_signature is now used on both arms.
The review's severity was wrong, though: it does not write dbias wrong. CuTeDSL
validates the compiled descriptor's static extents at call time, so the second n
raises "Mismatched dbias_tensor.shape[1] ... expected to be 384" (job 469460).
Fail-loud, not silent corruption. Pre-existing on the default path -- reproduced
with deterministic=False -- so the new test is parametrized over both modes.
2. The write-only outputs were still allocated on torch's stream while the kernel
writes them on the caller's. record_stream does not cover that direction: it
defers reuse after *this* tensor is freed, and says nothing about whether the
block was safe to hand out. torch's allocator recycles a block freed on the
allocating stream on the strength of that stream's ordering, which does not
constrain a kernel on current_stream. d_row/d_col/d_srelu and the sfd tensors now
allocate inside the stream context, like the accumulators.
That makes _record_streams dead: nothing is allocated on torch's stream any more,
so recording the caller's stream is a no-op. Removed, and with its only caller
gone it comes back out of grouped/backend_utils -- that shared module is now
untouched by this branch, 0 lines changed against upstream.
3. The two reductions disagreed. The kernel atomically *adds* into dbias when the
flag is off, so a caller accumulating across micro-batches got accumulate
semantics one way and overwrite the other. _reduce_dbias_slots now accumulates
like _reduce_dprob_slots, reduce_dbias_workspace takes the output tensor for
symmetry with reduce_dprob_workspace, and the class-API test drops the copy_ that
was papering over it.
15 determinism tests pass, both arms of the new n-collision test pass, and the 25
remaining failures are the fp4 ones the HEAD~1 control showed are inherited
(job 469466).
---------
Co-authored-by: Zhiyu Li <zhiyul@oci-aga-slurm-1-dm-02.cm.cluster>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
13 tasks
Merged
3 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Before submitting
pre-commit runand committed any formatting changes.cat-*, one or moremod-*, and oneorig-*(see label list).Affected area
Summary
Why
Related issues
API and compatibility impact
Testing
Summary by CodeRabbit
New Features
dprobresults across repeated runs.dbiasresults remain nondeterministic.Documentation
Tests