feat(cake_moe_grouped_gemm): fuse the fp32 weight-gradient accumulate into main_grad on sm_100a, sm_103a and sm_107a

docs/api/cake_moe_grouped_gemm.rst

@@ -58,6 +58,55 @@ The operator is local: on a context-parallel rank it runs on the rank's local ro
 no collective, and leaves the reduction of a parameter gradient across ranks to the
 caller.
 
+Weight-gradient accumulation
+----------------------------
+
+``grouped_gemm_wgrad(..., out=out, accumulate=True)`` adds the weight gradient into the
+caller's bf16 or fp32 ``out`` (``out[e] += G[e].T @ X[e]``, Megatron's gradient-accumulation
+fusion) and ``cake_grouped_mm(..., fuse_wgrad_accumulation=True)`` adds it into
+``w.main_grad`` in its backward.  Both default to off; with the defaults nothing changes.
+The fused accumulate builds of this release add into an fp32 ``out`` only; a bf16 ``out``
+always takes the composed route (the bf16 overwrite, then the same-dtype ``out.add_``:
+autograd's AccumulateGrad bits and cost).  The host plans the launch exactly as the
+overwrite of ``out``'s dtype and records the route in ``plan["accumulate"]``:
+
+============== ======================================================================== =========================================================
+route          when                                                                     what runs
+============== ======================================================================== =========================================================
+``"fused"``    512-wide k tile with its unordered split tail, the selected stage has an  one launch whose epilogue adds the tile into ``out``
+               accumulate build on the architecture whose library can load, ``out``
+               32-byte aligned
+``"composed"`` otherwise (``plan["accumulate_reason"]``, the first that fails:          the overwrite into a scratch of ``out``'s dtype, then
+               ``"tile_k"``, ``"tail"``, ``"stage"`` -- every bf16 ``out`` --,          ``out.add_``
+               ``"min_cuda"``, ``"alignment"``; the last two warn once per process)
+============== ======================================================================== =========================================================
+
+An accumulate build is bound by a static rule (nothing is compiled, loaded or test-built to
+decide it): its prebuilt library exists, or the version of FlashInfer's JIT toolkit reaches
+its ``build.json`` ``min_cuda``.  The fp32 accumulate builds need CUDA 13.4; below it the
+call composes instead of failing.  ``out`` must not overlap itself, ``g``, ``x`` or ``offs``
+(``ValueError`` on both routes, before anything runs).  Every launch and every CUDA graph
+replay adds again.  Numerical contract: the fused route adds the tile's fp32 sum into the
+fp32 ``out`` with one round-to-nearest add that keeps subnormals (bitwise
+``out0.add_(dW32)`` on every whole tile); a split tail tile adds its unordered sum ``S``
+once (``S`` is the overwrite's partial sum, whose unit reductions flush subnormals; the
+nondeterminism scope is unchanged); the composed route is ``out.add_`` of the overwrite in
+``out``'s dtype; an empty group is not touched by the fused route, while the composed
+route has ``out.add_``'s semantics (values kept, ``-0.0`` may become ``+0.0``).  The
+composed scratch (the size of the weight gradient) is owned by the prepared launch: keep
+the launch alive while a CUDA graph that captured it may replay.  The autograd flag returns
+a placeholder gradient for ``w`` when ``w`` has ``grad_added_to_main_grad``
+(``torch.zeros`` under ``zero_out_wgrad``), else none; it adds at every backward, also when
+``torch.autograd.grad``'s ``inputs`` exclude ``w`` (split backward-weight schedules are not
+supported), and raises ``NotImplementedError`` under ``torch.jit.trace``, ``torch.compile``
+and a ``create_graph`` backward.
+
+Through the facade, ``accumulate=True`` on the deterministic route (``deterministic=True``,
+or torch's deterministic mode with its one-time warning) computes the deterministic weight
+gradient in ``out``'s dtype and adds it with ``out.add_`` after the same checks;
+``fuse_wgrad_accumulation=True`` on the deterministic route raises ``NotImplementedError``
+(the deterministic package has not adopted it yet).
+
 Functions
 ---------
 
@@ -69,9 +118,11 @@ autodoc before graduation):
 - ``grouped_gemm_dgrad(g, w, offs, out=None, *, deterministic=True)``: ``dX = G @ W[e]``
   per group (``W`` read in place), bf16 ``[sum_m, K]``;
 - ``grouped_gemm_wgrad(g, x, offs, out=None, out_dtype=None, num_groups=None, *,
-  deterministic=True)``: ``dW[e] = G[e].T @ X[e]``, bf16 or fp32 ``[E, N, K]``;
-- ``cake_grouped_mm(x, w, offs, deterministic=True)``: the autograd wrapper (``Y``, with
-  ``dX`` and ``dW`` in the backward, ``dW`` in the dtype of ``w``).
+  deterministic=True, accumulate=False)``: ``dW[e] = G[e].T @ X[e]``, bf16 or fp32
+  ``[E, N, K]`` (with ``accumulate=True``, ``out[e] += dW[e]``);
+- ``cake_grouped_mm(x, w, offs, deterministic=True, *, fuse_wgrad_accumulation=False)``: the
+  autograd wrapper (``Y``, with ``dX`` and ``dW`` in the backward, ``dW`` in the dtype of
+  ``w``, or added into ``w.main_grad``).
 
 The shapes, the group offsets (``offs[E]`` end offsets or ``m_indptr[E + 1]``) and the
 alignment rules are those of :ref:`apicake_moe_grouped_det_gemm`.  The prepare-once /
@@ -123,5 +174,7 @@ feature                                         sm_100a         sm_103a
 ``wgrad`` 512 half-release builds               planned, slack  planned, slack  planned, slack
 ``wgrad`` 512 k256-tail twin                    --              --              slack tail
 ``wgrad`` output dtypes                         bf16, fp32      bf16, fp32      bf16, fp32
+``wgrad`` accumulate, fused, fp32 ``out``       planned, wave   planned, wave   planned, k256t
 CUDA toolkit floor                              12.8            12.9            13.4
+CUDA toolkit floor, fp32 accumulate builds      13.4            13.4            13.4
 =============================================== =============== =============== ===============

flashinfer/experimental/cake/cake_moe_grouped_gemm/README.md

@@ -19,7 +19,7 @@ an int32 **device** tensor `offs[E]` of cumulative end offsets, or
 | --- | --- | --- |
 | `grouped_gemm_fwd(x, w, offs, out=None, *, deterministic=True)` | `Y[offs[e-1]:offs[e]] = X[offs[e-1]:offs[e]] @ W[e].T` | `Y [sum_m, N]` bf16 |
 | `grouped_gemm_dgrad(g, w, offs, out=None, *, deterministic=True)` | `dX[offs[e-1]:offs[e]] = G[offs[e-1]:offs[e]] @ W[e]` (W read in place) | `dX [sum_m, K]` bf16 |
-| `grouped_gemm_wgrad(g, x, offs, out=None, out_dtype=None, num_groups=None, *, deterministic=True)` | `dW[e] = G[offs[e-1]:offs[e]].T @ X[offs[e-1]:offs[e]]` | `dW [E, N, K]` bf16 or fp32 |
+| `grouped_gemm_wgrad(g, x, offs, out=None, out_dtype=None, num_groups=None, *, deterministic=True, accumulate=False)` | `dW[e] = G[offs[e-1]:offs[e]].T @ X[offs[e-1]:offs[e]]`; with `accumulate=True`, `out[e] += dW[e]` | `dW [E, N, K]` bf16 or fp32 |
 
 ## Numerical mode and dispatch
 
@@ -83,6 +83,134 @@ This is a local operator: on a context-parallel rank it runs on that rank's
 local rows; it adds no collective and no KCP form, and a parameter gradient's
 reduction across ranks stays with the caller.
 
+## Weight-gradient accumulation
+
+`grouped_gemm_wgrad(..., out=out, accumulate=True)` (and
+`prepare_grouped_gemm_wgrad(..., accumulate=True)`) adds the weight gradient
+into the caller's `out` instead of overwriting it: `out[e] += G[e].T @ X[e]`,
+Megatron's `gradient_accumulation_fusion` into an fp32 or bf16 `main_grad`.
+`cake_grouped_mm(..., fuse_wgrad_accumulation=True)` does the same in its
+backward into `w.main_grad` (below). Both default to off; with the defaults
+the routes, programs, plans and bits are unchanged, and `plan["accumulate"]`
+is `"off"`.
+
+The fused accumulate builds of this release add into an **fp32** `out` only:
+the 512-wide k tile's half-release forms have fp32 accumulate builds (feature
+matrix below). A **bf16** `out` always takes the composed route: the bf16
+overwrite into a scratch, then the same-dtype `out.add_`, which is
+autograd's AccumulateGrad (the same bits and the same cost as the default
+bf16 `w.grad` accumulation).
+
+The host plans an accumulating launch exactly as the overwrite of `out.dtype`
+(route, k tile, stage, grid, tail split, raster, wave mode) and records how it
+adds in `plan["accumulate"]`:
+
+| `plan["accumulate"]` | when | what runs |
+| --- | --- | --- |
+| `"fused"` | the 512-wide k tile with its unordered split tail, the route serves the selected stage's accumulate build (`stages == (stage + "_acc",)`) on the architecture, that build's library can load (below), and `out` is 32-byte aligned (base, row pitch, expert pitch) | one launch whose epilogue reads C and adds the tile into it; no new workspace, nothing allocated after the first launch on a stream |
+| `"composed"` | otherwise; `plan["accumulate_reason"]` is the first condition that fails: `"tile_k"` (the 256-wide k tile), `"tail"` (an ordered tail), `"stage"` (no accumulate build of the stage for the dtype and architecture: every bf16 `out`, and the `main` / `main_wave` forms), `"min_cuda"` (the build cannot load) or `"alignment"`; the last two warn once per process | the overwrite into a contiguous scratch of `out.dtype`, then `out.add_(scratch)` on the same stream |
+
+The route is chosen by a static rule (nothing is compiled, loaded or
+test-built to decide it): an accumulate build is bound when its prebuilt
+library exists under `aot/`, or else when the version of FlashInfer's JIT
+toolkit reaches the build's `build.json` `min_cuda` (the loader's own rule). The fp32 accumulate builds use the PTX ISA 9.4
+`red.add.noftz.f32` form, so their `min_cuda` is 13.4 on every architecture;
+an ahead-of-time build with an older toolkit ships no library for them, and
+with a JIT toolkit below 13.4 the call composes (`"min_cuda"`) instead of
+failing.
+
+`out` is required (`out_dtype` None or `out.dtype`) and must not overlap
+itself (row stride at least `K`, expert stride at least `N` row strides: an
+expanded `out` is refused) nor share memory with `g`, `x` or `offs`; both
+routes check this before running anything (`ValueError`). Row and expert
+padding is allowed and never read or written. `launch()` and every CUDA graph
+replay add again: a prepared accumulate launch is not idempotent. A fused
+launch bumps `out._version` when it is launched or captured (the composed
+`add_` bumps it itself). `sum_m == 0` launches nothing and leaves `out`
+unchanged. The composed route's scratch (`E * N * K` elements of `out.dtype`,
+the size of the weight gradient) is owned by the prepared launch, as its
+descriptor slots are: keep the launch alive while any CUDA graph that captured
+it may replay (the graph does not keep the scratch alive), and do not have one
+prepared composed launch in flight on two streams at once. Inside a captured
+backward, the scratch is allocated from, and returned to, the graph's pool
+within that backward. As for any in-place operation, two launches into the
+same `out` on different streams must be ordered by the caller.
+
+Numerical contract:
+
+- Fused (fp32 `out`): `out <- out + dW` with one round-to-nearest fp32 add of
+  the tile's fp32 sum into C. Every whole tile is bitwise `out0.add_(dW32)`,
+  where `dW32` is the fp32 weight gradient of the same plan; that add keeps
+  subnormals (no flush to zero).
+- A split tail tile (512-wide k tile, `tail_splits > 1`) adds `C + S` once,
+  where `S` is the tile's unordered fp32 sum of its split units, so it varies
+  run to run as the overwrite does; the nondeterminism scope is unchanged.
+  `S` is the overwrite's partial sum, whose unit reductions flush subnormal
+  values to zero (above); subnormals are kept only in the `C + S` add.
+- Composed: the overwrite's `dW` in `out.dtype` added by `out.add_`. For a
+  bf16 `out` that is `RN_bf16(out + dW_bf16)`, AccumulateGrad's bits.
+- An empty group's `out[e]` is not read or written on the fused route (its
+  bits, `-0.0` and NaN payloads included, stay). The composed route has the
+  semantics of `out.add_`: the values are kept, but `-0.0` may become `+0.0`
+  and a signaling NaN is quieted. The published contract is the composed one;
+  bitwise-untouched empty groups are a property of the fused route.
+
+`cake_grouped_mm(x, w, offs, deterministic=False, fuse_wgrad_accumulation=True)`
+follows Megatron and Transformer Engine: `w.main_grad` must be an fp32 or bf16
+tensor of `w`'s shape on its device with unit stride along the last dimension
+(checked at the forward call when autograd records it, that is when `w`
+requires grad and grad mode is on, so not under `torch.no_grad()`; and again
+at the backward, which reads `w.main_grad` then, so a re-pointed buffer is
+honoured).
+The backward adds into it (fused for an fp32 `main_grad` where the route
+serves it, composed for a bf16 one) and returns for `w` a fresh `torch.empty`
+of `w`'s shape and dtype when `w` has the attribute `grad_added_to_main_grad`
+(it is set to `True`; `torch.zeros` when `w.zero_out_wgrad` is true), else
+`None` (`w.grad` stays untouched and no gradient hook of `w` fires). `x.grad`
+is unchanged, bitwise. Every backward adds: `retain_graph` twice adds twice,
+and a `torch.autograd.grad(..., inputs=...)` that excludes `w` still adds into
+`w.main_grad`, so zero-bubble (split backward-input / backward-weight)
+schedules are not supported. A `w` that does not require grad is never added
+to. The flag raises `NotImplementedError` under `torch.jit.trace`,
+`torch.compile` and in a `create_graph` backward; under a reentrant checkpoint
+`w` arrives detached without `main_grad` (`ValueError`; use
+`use_reentrant=False`).
+
+Peak memory of one weight gradient into an fp32 `main_grad` (`S = E * N * K`
+elements): against an fp32 `dW` followed by `main_grad.add_` (a `4 * S`-byte
+temporary), the fused route with `grad_added_to_main_grad` saves `2 * S` bytes
+(its bf16 placeholder of `2 * S` bytes remains; without the attribute it saves
+`4 * S`); against Megatron's unfused path (a bf16 `w.grad` of `2 * S` bytes,
+then the fp32 add) the peak is equal. The benchmark's accumulation arms report
+the measured peak.
+
+The facade `flashinfer.experimental.cake.grouped_mm` takes the same keywords:
+
+| entry | `deterministic` | torch det. mode | runs |
+| --- | --- | --- | --- |
+| `grouped_gemm_wgrad(..., accumulate=True)` | `False` | off | this package's accumulate mode |
+| `grouped_gemm_wgrad(..., accumulate=True)` | `True`, or `False` with det. mode on (warns) | any | the deterministic `dW` in `out.dtype` into a scratch, then `out.add_` (the same checks first) |
+| `cake_grouped_mm(..., fuse_wgrad_accumulation=True)` | `False` | off | this package's wrapper |
+| `cake_grouped_mm(..., fuse_wgrad_accumulation=True)` | `True`, or `False` with det. mode on (warns) | any | `NotImplementedError`: the deterministic package has not adopted the accumulation yet |
+
+```python
+from flashinfer.experimental.cake.cake_moe_grouped_gemm import prepare_grouped_gemm_wgrad
+from flashinfer.experimental.cake.grouped_mm import cake_grouped_mm
+
+# Megatron-style fused accumulation into an fp32 main_grad
+w = torch.nn.Parameter(torch.randn(E, N, K, dtype=torch.bfloat16, device="cuda"))
+w.main_grad = torch.zeros(E, N, K, dtype=torch.float32, device="cuda")
+# (Megatron's DDP also sets w.grad_added_to_main_grad; the backward then returns a
+# placeholder gradient for its hooks.  Without the attribute w.grad stays None.)
+for x, g in micro_batches:
+    y = cake_grouped_mm(x, w, offs, deterministic=False, fuse_wgrad_accumulation=True)
+    y.backward(g)  # w.main_grad[e] += G[e].T @ X[e]; x.grad as usual
+
+# Low level: add into any fp32 / bf16 buffer
+launch = prepare_grouped_gemm_wgrad(g, x, offs, out=main_grad, deterministic=False, accumulate=True)
+launch.launch()  # main_grad += dW; launch.plan["accumulate"] is "fused" or "composed"
+```
+
 ## Usage
 
 ```python
@@ -198,7 +326,10 @@ only); the 512-wide k tile's routes bind the unordered-tail builds.
 | `wgrad` k512 unordered tail: `main_wave`, `main_wave_halves` (wp8) | yes | yes | - |
 | `wgrad` k512 unordered tail: `main_halves_k256t` | - | - | yes (slack tail) |
 | `wgrad` output dtypes | bf16, fp32 | bf16, fp32 | bf16, fp32 |
+| `wgrad` accumulate, fused (`_acc` builds, fp32 `out` only) | `main_halves`, `main_wave_halves` | `main_halves`, `main_wave_halves` | `main_halves`, `main_halves_k256t` |
+| `wgrad` accumulate, composed | bf16 `out`; k256 tile; `main` / `main_wave`; misaligned `out`; `_acc` library not loadable | bf16 `out`; k256 tile; `main` / `main_wave`; misaligned `out`; `_acc` library not loadable | bf16 `out`; k256 tile; `main`; misaligned `out`; `_acc` library not loadable |
 | CUDA toolkit floor (`min_cuda`) | 12.8 | 12.9 | 13.4 |
+| CUDA toolkit floor of the fp32 `_acc` builds | 13.4 | 13.4 | 13.4 |
 
 The slack tail placement runs a launch's split tail whole on the clusters the
 full waves leave least loaded when every tail tile fits there; the kernel
@@ -218,7 +349,7 @@ loads without a FlashInfer installation.
 ## Tests and benchmark
 
 ```bash
-pytest tests/experimental/test_cake_moe_grouped_gemm.py tests/experimental/test_cake_moe_grouped_gemm_api.py tests/experimental/test_cake_moe_grouped_gemm_standalone.py tests/experimental/test_cake_moe_grouped_gemm_reproducibility.py
+pytest tests/experimental/test_cake_moe_grouped_gemm.py tests/experimental/test_cake_moe_grouped_gemm_api.py tests/experimental/test_cake_moe_grouped_gemm_standalone.py tests/experimental/test_cake_moe_grouped_gemm_reproducibility.py tests/experimental/test_cake_moe_grouped_gemm_acc.py
 python benchmarks/bench_cake_moe_grouped_gemm.py --help
 ```
 
@@ -235,7 +366,16 @@ checks that `deterministic=True` gives the deterministic package's bits, that
 `fwd` / `dgrad` give them with `deterministic=False`, and reports the
 run-to-run variation of the unordered weight gradient.
 `test_cake_moe_grouped_gemm_standalone.py` imports a copy with `flashinfer`
-blocked. The benchmark times the deterministic and non-deterministic packages
+blocked. `test_cake_moe_grouped_gemm_acc.py` covers the accumulate mode: the
+`_acc` stages against the shipped set and their twins, the plans, the routes
+and the checks on the CPU; the fused and composed results against the fp64
+reference and the bitwise oracles (also on every benchmark row as planned),
+empty groups, padding, short groups, streams, graph replay, allocations and
+the version counter on a GPU; and the autograd flag. The benchmark's `acc`
+(checkpointed micro-step), `bwd_acc` (its backward) and `wgrad_acc` (one
+weight gradient) rows time the accumulation arms (`ovw`, `acc`, `mgB`, `mgF`,
+`mgO`, `mgX`, `mgXn`, `mgX_zero`, `accX`, and the deterministic package's
+`det_acc` and `det_mgB`) with peak memory. The benchmark times the deterministic and non-deterministic packages
 (fwd, dgrad, wgrad and an autograd `2 x fwd + bwd` step) with peak memory,
 against `torch._grouped_mm` and a per-group cuBLAS loop.
 

flashinfer/experimental/cake/cake_moe_grouped_gemm/__init__.py

@@ -157,6 +157,7 @@ def grouped_gemm_wgrad(
     num_groups: Optional[int] = None,
     *,
     deterministic: bool = True,
+    accumulate: bool = False,
 ) -> torch.Tensor:
     """``dW[e] = G[offs[e-1]:offs[e]].T @ X[offs[e-1]:offs[e]]`` for every group ``e``.
 
@@ -172,6 +173,14 @@ def grouped_gemm_wgrad(
     non-determinism alert.  ``N % 256 == 0`` and ``K`` a
     multiple of the selected k tile (256, or 512 when chosen) are required.
     See ``grouped_gemm_fwd`` for the prepared form.
+
+    ``accumulate=True`` adds into ``out`` instead, ``out[e] += G[e].T @ X[e]``
+    (Megatron's gradient-accumulation fusion; ``out`` required, bf16 or fp32,
+    ``out_dtype`` None or its dtype), and returns ``out``: fused into the
+    512-wide k tile's epilogue where the route serves that stage's accumulate
+    build, else the overwrite into a scratch followed by ``out.add_``
+    (``prepare_grouped_gemm_wgrad`` and the package README give the routes and
+    the numerical contract).
     """
     from .cake_backend import (
         _alert_not_deterministic,
@@ -181,7 +190,9 @@ def grouped_gemm_wgrad(
 
     _require_deterministic_false(deterministic, "grouped_gemm_wgrad")
     _alert_not_deterministic("grouped_gemm_wgrad")
-    return _prepare_wgrad(g, x, offs, out, out_dtype, num_groups)._run()
+    return _prepare_wgrad(
+        g, x, offs, out, out_dtype, num_groups, accumulate=accumulate
+    )._run()
 
 
 def __getattr__(name: str) -> Any: