Module combined_backward 

Recompute-based gradient math (the memory-efficient backward).

Recompute-based gradient math for the Mamba-2 SSD

The analytic backward of the five-kernel serial scan, mirroring _mamba_chunk_scan_combined_bwd in the reference ssd_combined.py. The forward intermediates (K1–K4) are recomputed from the saved leaf inputs rather than stashed, then a reverse per-chunk loop fuses the K5 and K4 backwards; K1/K2/K3 backwards run as batched ops once the loop has gathered the per-chunk slices. Comment colours (BLUE / ORANGE / …) tag the corresponding terms of the chunk-scan gradient, matching the reference.

Everything here operates on backend primitives through the rank-tagged [F] wrapper: a custom Backward node runs with a generic backend B, so the high-level Tensor (pinned to the global Dispatch backend) is unavailable and the math must use B’s float_* ops. The recomputed K1/K2/K4 kernels are local primitive ports of the high-level crate::mamba2::ssd::serial kernels.

Structs

CombinedGrads

Per-input gradients produced by combined_backward (one field per differentiable forward input).

Functions

combined_backward

Memory-efficient backward for the Mamba-2 chunkwise SSD.

k3_ssd_chunk_state_extended

Same as k3_ssd_chunk_state but also returns intermediates needed by the custom backward: