Repository navigation
feat(cake_moe_grouped_gemm): fuse the fp32 weight-gradient accumulate into main_grad on sm_100a, sm_103a and sm_107a - #8602
Merged
yyihuang merged 8 commits intoNov 4, 2026
Conversation
grouped_gemm_wgrad / prepare_grouped_gemm_wgrad take accumulate=True: the launch adds G_e^T X_e into the caller's out[e] (out <- out + dW). The fused route binds the accumulating build of the selected main stage (K tile 512, 32-byte aligned out); every other plan computes into a scratch of out's dtype and adds it (composed route). plan["accumulate"] names the route. Both routes check out before any launch: dtype, row and expert strides (an expanded out is rejected) and no overlap with the operands. The fused route bumps out's version counter. cake_grouped_mm(..., fuse_wgrad_accumulation=True) accumulates the weight gradient into w.main_grad in the backward (Megatron-Core gradient_accumulation_fusion convention) and returns the placeholder weight gradient. The facade routes accumulate=True on the deterministic route through the deterministic weight gradient plus an add. Adds the accumulation tests, the benchmark arms and the README / API docs sections. Defaults are unchanged.
… builds Regenerate the package's generated files (the PROGRAMS and ROUTES tables of cake_jit.py, build.json and csrc/cake_moe_grouped_gemm/) with the fp32 accumulate builds of the three half-release forms of the 512-wide weight gradient, which add the weight gradient to out (out += dW): - 3 programs, each served as a stage of the fp32 512-wide route only: main_halves_acc (sm_100a, sm_103a, sm_107a), main_wave_halves_acc (sm_100a, sm_103a) and main_halves_k256t_acc (sm_107a). The existing stages keep their order and programs; the bf16 route is unchanged. - 6 build.json entries, min_cuda 13.4 (the fp32 .noftz reduction needs PTX ISA 9.4). The accumulate programs share their own binding templates. - The existing generated kernel units drop unused macros and reword comments; their SASS and resource usage are unchanged on every architecture they build for. No existing program is renamed or removed.
…ey load Follow-up of the fp32 accumulate builds for the host, tests, docs and benchmark: - Route rule: an accumulate call is fused only when the 512-wide k tile adds its split tail in place (never beside an overwriting tail-reduce stage: reason "tail"), the route serves the stage's accumulate build, its library can load by the loader's static rule (a prebuilt library under aot/, or a FlashInfer JIT toolkit at the build's min_cuda; nothing is compiled or loaded to decide it: reason "min_cuda", one warning) and out is 32-byte aligned. An fp32 accumulate on a toolkit below CUDA 13.4 now composes instead of failing at the preparation. generated_program_available(accumulate=True) applies the same rule and is False for a dtype the weight gradient does not write. - Docs: the fused builds add into an fp32 out only; a bf16 out composes (the bf16 overwrite and a same-dtype add_: AccumulateGrad's bits and cost). The numerical contract states that subnormals are kept in the add into C, not in the split tail's partial sum; the composed scratch is owned by the prepared launch (keep it alive while a captured graph may replay); the peak memory saving is stated against an fp32 dW plus add_ (2 x E x N x K bytes with the placeholder gradient) and against the bf16 w.grad path (equal). - Tests: the registry's accumulate stages must be exactly the shipped set (also stated by the README and rst feature rows), the build.json entries are required, a shipped build is required to fuse wherever it can load, the loader rule and the ordered-tail rule have CPU tests, every benchmark row is checked on the fused fp32 route as planned (fp64 budget, no lost update, empty groups bitwise, bitwise out0.add_(dW32) outside the split tail) and on the composed bf16 route, and the bf16 budget of an accumulate is operand-scaled. - Benchmark: the deterministic package's baselines (det_acc: its fp32 dW plus main_grad.add_, also in the autograd step; det_mgB: its bf16 w.grad plus the fp32 add) and a backward-only accumulation row (bwd_acc).
…y device The sanitize entry searched its split case over four small groups that never split a tail, so it printed "no shape" and sanitized unsplit tails only. The split case now takes the groups the split-tail tests use, and the entry fails when a shipped build has no split or unsplit case.
…backward only With fuse_wgrad_accumulation=True the forward checked w.main_grad on every call. In nested non-reentrant checkpoints, torch recomputes the inner region inside the backward with an alias of w that has no Python attributes, so the recompute raised "w has no main_grad tensor". The forward now checks main_grad only outside the autograd engine (graph task id -1). The backward of the original forward reads its own ctx.w and checks main_grad there, so a reentrant checkpoint still raises at the backward.
… count The split-tail test and the sanitize entry searched one set of group sizes, which splits no k256t tail at 160 or 208 SMs (R200 208-SM board: "no candidate shape splits the main_halves_k256t tail"). Both now try SPLIT_SIZES in order (the old sizes first, so other devices keep their case) and take the first whose plan splits the tail.
…autograd records With fuse_wgrad_accumulation=True, a forward under torch.no_grad() or inference mode raised "w has no main_grad tensor" for a w that requires grad but has no main_grad, although no backward reads it then. The check now runs in cake_grouped_mm before apply, when w requires grad, grad mode is on and the call is not a checkpoint recompute inside a backward. It leaves CakeGroupedMm.forward, which runs with grad mode off and cannot see the caller's mode. The backward still checks w.main_grad, so a direct apply is checked there. The host refusal test covers no_grad and inference mode.
Contributor
|
Warning Review limit reachedYou've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository. Next included review available in 15 minutes. View limit details
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. |
Collaborator
Author
|
@flashinfer-bot run tests/experimental/test_cake_moe_grouped_gemm.py tests/experimental/test_cake_moe_grouped_gemm_acc.py tests/experimental/test_cake_moe_grouped_gemm_api.py tests/experimental/test_cake_moe_grouped_gemm_reproducibility.py tests/experimental/test_cake_moe_grouped_gemm_standalone.py tests/experimental/test_cake_moe_grouped_det_gemm.py tests/experimental/test_cake_moe_grouped_det_gemm_alias.py tests/experimental/test_cake_moe_grouped_det_gemm_api.py tests/experimental/test_cake_moe_grouped_det_gemm_standalone.py |
1 task done
Collaborator
Author
|
@flashinfer-bot run tests/experimental/test_cake_moe_grouped_gemm.py tests/experimental/test_cake_moe_grouped_gemm_acc.py tests/experimental/test_cake_moe_grouped_gemm_api.py tests/experimental/test_cake_moe_grouped_gemm_reproducibility.py tests/experimental/test_cake_moe_grouped_gemm_standalone.py tests/experimental/test_cake_moe_grouped_det_gemm.py tests/experimental/test_cake_moe_grouped_det_gemm_alias.py tests/experimental/test_cake_moe_grouped_det_gemm_api.py tests/experimental/test_cake_moe_grouped_det_gemm_standalone.py |
yyihuang
marked this pull request as ready for review
November 4, 2026 07:25
yyihuang
requested review from
Anerudhan,
Jingfan-Sun,
YangXu1990uiuc,
aleozlx,
bkryu,
flashinfer-bot and
yzh119
as code owners
November 4, 2026 07:25
yyihuang
enabled auto-merge (squash)
November 4, 2026 07:25
Collaborator
Author
|
@flashinfer-bot run tests/experimental/test_cake_moe_grouped_gemm.py tests/experimental/test_cake_moe_grouped_gemm_acc.py tests/experimental/test_cake_moe_grouped_gemm_api.py tests/experimental/test_cake_moe_grouped_gemm_reproducibility.py tests/experimental/test_cake_moe_grouped_gemm_standalone.py tests/experimental/test_cake_moe_grouped_det_gemm.py tests/experimental/test_cake_moe_grouped_det_gemm_alias.py tests/experimental/test_cake_moe_grouped_det_gemm_api.py tests/experimental/test_cake_moe_grouped_det_gemm_standalone.py |
yzh119
approved these changes
Nov 4, 2026
yzh119
left a comment
yzh119
left a comment
Collaborator
There was a problem hiding this comment.
LGTM. One nit: cake_backend.py imports _jit_toolchain, _local_cuda_version, _version, libraries_by_name and prebuilt_path from cake_build. Make sure they all exist there, since that file isn't in the diff.
yyihuang
added a commit
that referenced
this pull request
Nov 6, 2026
…ic=False, and deterministic weight-gradient accumulation (#8924) ## Summary This PR changes the Cake MoE grouped GEMM in FlashInfer host code only. Kernels and programs are unchanged. 1. **Det-only names.** The deterministic grouped GEMM can now be reached only through names that contain `det`. All old det-free names and compatibility aliases are removed with no shim. Old imports raise `ModuleNotFoundError` and old attributes raise `AttributeError`. 2. **Explicit `deterministic=False`.** The non-deterministic package and the `flashinfer.experimental.cake.grouped_mm` facade run only when `deterministic=False` is passed explicitly. PyTorch's non-determinism alert fires only right before an order-varying launch. 3. **Deterministic weight-gradient accumulation.** `grouped_det_gemm_wgrad(..., out=main_grad, accumulate=True)` and `cake_grouped_det_mm(..., fuse_wgrad_accumulation=True)` accumulate into `main_grad`. The result is bitwise equal to AccumulateGrad of the deterministic dW and bitwise identical from run to run. In this description, the *det package* is `flashinfer.experimental.cake_moe_grouped_det_gemm` and the *non-det package* is `flashinfer.experimental.cake.cake_moe_grouped_gemm`. ## Breaking change: det-only names No alias, shim, module `__getattr__` or compatibility entry is left anywhere. | Old (removed) | Error now | Det name to use | | --- | --- | --- | | package `flashinfer.experimental.cake_moe_grouped_gemm` and submodules `cake_autograd`, `cake_backend`, `cake_build`, `cake_jit` (directory deleted) | `ModuleNotFoundError` | `flashinfer.experimental.cake_moe_grouped_det_gemm` (same submodules) | | module `flashinfer.grouped_mm.cake` (renamed with `git mv`) | `ModuleNotFoundError` | `flashinfer.grouped_mm.cake_det` | | `flashinfer.grouped_mm.{grouped_gemm_fwd, grouped_gemm_dgrad, grouped_gemm_wgrad, cake_grouped_mm}` | `AttributeError` | `flashinfer.grouped_mm.{grouped_det_gemm_fwd, grouped_det_gemm_dgrad, grouped_det_gemm_wgrad, cake_grouped_det_mm}` | | det package root `grouped_gemm_{fwd,dgrad,wgrad}`, `prepare_grouped_gemm_{fwd,dgrad,wgrad}`, `GroupedGemmLaunch`, `CakeGroupedMm`, `cake_grouped_mm` | `AttributeError` | `grouped_det_gemm_{fwd,dgrad,wgrad}`, `prepare_grouped_det_gemm_{fwd,dgrad,wgrad}`, `GroupedDetGemmLaunch`, `CakeGroupedDetMm`, `cake_grouped_det_mm` | | det `cake_backend.{GroupedGemmLaunch, grouped_mm_bf16_cake, prepare_grouped_gemm_*}` | `AttributeError` | `cake_backend.{GroupedDetGemmLaunch, grouped_det_mm_bf16_cake, prepare_grouped_det_gemm_*}` | | det `cake_autograd.{CakeGroupedMm, cake_grouped_mm}` | `AttributeError` | `cake_autograd.{CakeGroupedDetMm, cake_grouped_det_mm}` | | `_RENAMED` maps and alias `__getattr__` in det `__init__`, `cake_autograd`, `cake_backend`, `grouped_mm/cake(_det).py`, `grouped_mm/__init__.py` (and its `_Any` import) | deleted | - | | `deterministic` parameter of every det entry (`cake_grouped_det_mm` in the det package and in `grouped_mm.cake_det`) | `TypeError` | omit it | | `CakeGroupedDetMm.apply(x, w, offs, <mode>)` (stale 4th positional mode) | `TypeError` "takes no fourth argument" | `apply(x, w, offs)`; fuse: `apply(x, w, offs, None, True)` | | test-helper key `PACKAGES["cake_moe_grouped_gemm"]` (MoE entry of the shared `tests/test_helpers/cake_package.py` only) | `KeyError` | `"cake_moe_grouped_det_gemm"` (and `"cake.cake_moe_grouped_gemm"` for the non-det package) | | `tests/experimental/test_cake_moe_grouped_det_gemm_alias.py` | deleted | absence tests in `test_cake_moe_grouped_det_gemm_api.py` / `test_cake_moe_grouped_gemm_api.py` | | det benchmark `benchmarks/bench_cake_moe_grouped_det_gemm.py`: arms `cake`, `cake_inforward`; functions `cake_arm`, `cake_inforward_arm`, `_cake_vs_inforward`; JSON key `cake_vs_inforward` | `parser.error` / `AttributeError` | arms `cake_det`, `cake_det_inforward`; `cake_det_arm`, `cake_det_inforward_arm`, `_cake_det_vs_inforward`; key `cake_det_vs_inforward` | | non-det `cake_backend._check_accumulate_call` (used only by the old facade det route) | deleted | - | **Unchanged names:** - `flashinfer.grouped_mm.grouped_mm_bf16(..., backend="cake")` keeps the value `"cake"`. It runs the forward pass only, and the det and non-det forward programs are bitwise identical. - The non-det package, the facade `flashinfer.experimental.cake.grouped_mm` and the non-det program names stay as they are. **Editable installs:** a stale `flashinfer/experimental/cake_moe_grouped_gemm/__pycache__` makes the removed path import as an empty namespace package. To fix this, run `git clean -fdX flashinfer/experimental/cake_moe_grouped_gemm`. Fresh clones, wheels and CI are not affected. ## Mode rules | Entry | `deterministic` | `torch.use_deterministic_algorithms(True)` | `warn_only=True` | | --- | --- | --- | --- | | det package, `flashinfer.grouped_mm.cake_det`, `flashinfer.grouped_mm.{grouped_det_gemm_*, cake_grouped_det_mm}` | No parameter. Passing it (`True`/`False`/`None`) raises `TypeError`. | No alert, no warning | same | | `CakeGroupedDetMm.apply(x, w, offs[, None[, fuse]])` | A 4th argument other than `None` (e.g. a stale `True`/`False`/`0`/`1`) raises `TypeError` before any launch. A non-bool `fuse` raises `ValueError`. | No change | same | | `grouped_mm_bf16(backend="cake")` | No parameter. Passing it raises `TypeError`. | Det forward, no alert | same | | det accumulate: `grouped_det_gemm_wgrad(..., out=out, accumulate=True)` / `cake_grouped_det_mm(..., fuse_wgrad_accumulation=True)` | n/a | No alert. Result has the AccumulateGrad bits of the det dW and is bitwise identical run to run. | same | | non-det package entries (`grouped_gemm_*`, `prepare_grouped_gemm_*`, `cake_grouped_mm`) and the four facade entries | Keyword-only `deterministic: bool = True`. Only the bool `False` runs. Omitted/`True`/`None`/`0`/`1`/`"False"`/`np.False_` raise `ValueError`, which names `flashinfer.experimental.cake_moe_grouped_det_gemm.<det entry>` and says "pass deterministic=False for the non-deterministic path". | `RuntimeError` naming the entry and its det entry, raised right before an order-varying launch only | `UserWarning` with the same names, then the launch runs | | `CakeGroupedMm.apply(x, w, offs[, mode[, None[, fuse]]])` | The 4th argument is the mode, with default `True`, so omitting it raises `ValueError` like every value other than `False`. Anything but `None` in the 5th (internal) slot raises `TypeError`, so `apply(x, w, offs, False, True)` no longer drops the fuse flag. The 6th argument is the fuse flag. | as above | as above | - **No routing to det.** Nothing routes to the det package and there is no fallback. The facade's old det route is gone; call the det entry directly. - **Order-varying launches.** Only the weight gradient's split tail varies in order: the 512-wide k tile with `unordered_tail` and `tail_splits > 1`, in overwrite or accumulate mode. It is reached through: - `grouped_gemm_wgrad` - `GroupedGemmLaunch.launch()`, checked at every launch - `CakeGroupedMm.backward` when it computes dW - **What the alert names.** The alert names the det replacement: `grouped_det_gemm_wgrad`, `GroupedDetGemmLaunch.launch` or `CakeGroupedDetMm`. - **What never alerts.** fwd, dgrad, k256 or unsplit wgrad, and preparation never alert. - **Where the warning points.** Under `warn_only`, the warning points at the caller's line for direct calls and at a PyTorch autograd frame for the backward. - **No shared state between packages.** No tape, saved tensor or workspace crosses between the two packages. - **Other Cake families.** This is the rule the other Cake families adopt. ## Deterministic weight-gradient accumulation ### API - **`grouped_det_gemm_wgrad`.** New signature `grouped_det_gemm_wgrad(g, x, offs, out=None, out_dtype=None, num_groups=None, *, accumulate=False)`, available in the det package, `flashinfer.grouped_mm.cake_det` and `flashinfer.grouped_mm`. - **`cake_grouped_det_mm`.** New signature `cake_grouped_det_mm(x, w, offs, *, fuse_wgrad_accumulation=False)`, in the same three places, plus `CakeGroupedDetMm.forward(ctx, x, w, offs, _launch=None, fuse_wgrad_accumulation=False)`. With `fuse_wgrad_accumulation=True` the backward accumulates into `w.main_grad`, following the Megatron `main_grad` / `grad_added_to_main_grad` convention used by the non-det `cake_grouped_mm`. - **Prepared launches.** `cake_backend.prepare_grouped_det_gemm_wgrad(..., *, accumulate=False)`, `cake_backend.ACCUMULATE_MODES = ("off", "composed")`, `plan["accumulate"]` and `GroupedDetGemmLaunch.add_from`. - **Benchmark.** The non-det benchmark `benchmarks/bench_cake_moe_grouped_gemm.py` gains the arm `det_mgX` and the ratios `det_mgX/det_acc`, `det_mgX/det_mgB` and `mgX/det_mgX`. ### Design and numerics - **Composed accumulation.** The unchanged deterministic dW programs write to a launch-owned scratch of `out`'s dtype, and then `out.add_` runs on the same stream. Accumulation uses exactly the same plan and stages as the overwrite path. - **Bitwise equality.** For fp32 `out` the result is `RN32(out + dW)`. For bf16 `out`, the bf16 dW is added in bf16. Both are bitwise equal to AccumulateGrad of the deterministic dW and bitwise identical from run to run. - **Overlap checks.** A `w.main_grad` with overlapping rows or experts (expanded, row stride < K) raises `ValueError` at the recorded forward. Memory shared with the operands is still refused at the backward. - **CUDA graphs.** Both autograd backwards (det and non-det) query the current stream only outside a capture, guarded by `torch.cuda.is_current_stream_capturing()`. As a result, capturing the fused backward in a CUDA graph no longer fails at a stream query. - **Scratch lifetime.** The scratch has no `record_stream`. If you launch a prepared launch on a stream other than the one it was prepared on, keep the launch alive until that launch completes. ### Cost The transient scratch takes `4*E*N*K` bytes for an fp32 `out`. On the #8602 MoE shapes: | shape | E=4 | E=8 | E=16 | E=32 | | --- | --- | --- | --- | --- | | gate_up (N=4096, K=6144) | 384 MiB | 768 MiB | 1.5 GiB | 3 GiB | | down (N=6144, K=2048) | 192 MiB | 384 MiB | 768 MiB | 1.5 GiB | Under `torch.use_deterministic_algorithms(True)` with `fill_uninitialized_memory` on, PyTorch NaN-fills the scratch when it is allocated, which adds the cost of that fill. The extra write and read of the scratch cost time; see Performance. A fused deterministic accumulate program, which would not need the scratch, is future work. ## Performance Measured per architecture (B200 sm_100a, GB300 sm_103a, R200 sm_107a); table follows. ## Tests CPU suites pass; GPU suites on B200 / GB300 / R200: results follow. --------- Co-authored-by: Yingyi Huang <averyh@nvidia.com>
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
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.
The Cake ragged BF16 grouped GEMM for MoE expert training has two variants: deterministic
flashinfer.experimental.cake_moe_grouped_det_gemmand non-deterministicflashinfer.experimental.cake.cake_moe_grouped_gemm(selected withdeterministic=False). Both use FP32 accumulation. Each has forward (Y = X W^T), data-gradient (dX = G W) and weight-gradient (dW = G^T X, BF16 or FP32 output) programs. The facadeflashinfer.experimental.cake.grouped_mmreaches the kernel, and its default stays deterministic.This PR adds weight-gradient accumulation into an existing buffer,
out[e] += G_e^T X_e. It is opt-in and applies to the non-deterministic variant only. These are the semantics of Megatron'sgradient_accumulation_fusionand Transformer Engine'sfuse_wgrad_accumulation. The public surface isgrouped_gemm_wgrad(..., out=out, accumulate=True),prepare_grouped_gemm_wgrad(..., accumulate=True)andcake_grouped_mm(x, w, offs, deterministic=False, fuse_wgrad_accumulation=True). The last one adds intow.main_gradin the backward and supports Megatron'sgrad_added_to_main_grad.accumulateandfuse_wgrad_accumulationboth default to off. With the defaults, the routes, programs, plans and bits are unchanged. For an fp32outormain_grad, the add is fused into the epilogue of the k512 half-release weight-gradient programs. On sm_100a and sm_103a these are formsmain_halvesandmain_wave_halves; on sm_107a they aremain_halvesandmain_halves_k256t. This adds six fp32 accumulate builds that need CUDA toolkit 13.4 or newer. Every other case runs composed (the dW, then one add): a bf16main_grad, the 256-wide k tile, other forms, a misalignedout, or an accumulate library that cannot load.plan["accumulate"]reports"fused"or"composed". As before for this variant, the summation order of a split tail tile can differ between runs. Against an fp32 dW followed bymain_grad.add_, the fused route withgrad_added_to_main_gradsaves2 * Sbytes of peak memory (S = E * N * K elements).Results (point estimates, new non-deterministic Cake against each baseline):
m_splits, so it needs the group sizes on the host. This kernel reads device offsets with no host sync. Closing the standalone wgrad accumulate gap to TE is follow-up work.Baselines and measurement. Speedup = baseline latency / new latency, so a value above 1 means the new kernel is faster. Brackets are bootstrap 95% CIs of the paired per-group median ratio.
main_grad.add_. In full backward and checkpointed 2×fwd + bwd, it is the deterministic autograd plus the add.GroupedLinear(fuse_wgrad_accumulation=True)on the legacy per-expert cuBLAS path with hostm_splits.m_splits(cuBLASLt grouped GEMM).cublasGemmExwith beta=1 intomain_grad.torch._grouped_mmwith fp32 output is unsupported in torch 2.14.1.main_grad += dW" is the accumulate call alone (CUPTI cold-L2 span, 12 groups).main_gradis fp32.GB300 sm_103a
main_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWB200 sm_100a
main_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWR200 sm_107a
main_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWmain_grad += dWThe family also ships the deterministic variant, BF16 and FP32 dW outputs, and the forward and data-gradient programs on sm_100a (B200), sm_103a (GB300) and sm_107a (R200), and this PR's fused accumulate builds cover all three architectures.