Add state-passing context parallelism for Mamba2 - #6808
Open
shcho1118 wants to merge 5 commits into
Open
Conversation
Add the Triton kernels and the fused Conv1d+SSD autograd function that let Mamba2 run context parallelism with the sequence shard kept local. Only the causal boundary is exchanged: the convolution halo and an FP32 state summary of (final state, block decay) per segment, neither of which scales with sequence length. The SSD transform of one causal segment is affine in its initial state, S_out = a_block * S_in + S_ext so each rank computes its own (S_ext, a_block), all-gathers the packed summary, and recovers its true initial state with an exclusive causal scan over the gathered summaries. The gather is issued asynchronously and overlaps the local CB computation. Backward mirrors this with a reverse boundary scan and recomputes everything that scales with sequence length rather than saving it, writing SSD gradients straight into the buffers the convolution backward consumes. Two sequence layouts are supported. A contiguous causal shard is the natural one; the virtual layout instead treats each half of Megatron's balanced front/back shard as an independent causal segment, which lets the caller skip the activation exchange entirely. The permutation helpers that convert between the balanced and contiguous layouts are also here, with p2p and all-to-all backends. Nothing calls this yet; the mixer is wired up separately. Signed-off-by: Sanghun Cho <highcredit1118@gmail.com>
Make the state-passing CP path reachable from training through two new
options:
--use-mamba-state-passing-cp
--mamba-state-passing-cp-load-balancing {none,permute_p2p,permute_a2a,virtual}
The feature is opt-in and the existing all-to-all Mamba CP path stays the
default. When enabled, the mixer skips MambaContextParallel's activation
redistribution and hands the projected zxBCdt straight to the fused
state-passing function.
Standard Megatron CP gives the mixer a front/back balanced shard, so the
mixer rejects the 'none' layout, which expects an already-contiguous shard
and exists for direct calls into the kernels. The permute modes exchange
chunks for a contiguous shard; virtual reinterprets each balanced half as an
independent causal segment and moves no activations.
The path uses the TP-local convolution and SSM parameters rather than
MambaContextParallel's CP slices, since the heads are not sharded across CP
here. RMSNorm and the output projection stay outside the fused function, so
both remain available.
Inference, packed sequences, hybrid and dynamic CP, and
--mamba-training-ssm-states-dtype are not supported yet; each is rejected
with an explicit assertion rather than silently producing wrong results.
Signed-off-by: Sanghun Cho <highcredit1118@gmail.com>
Cover both graph backends for every load-balancing mode: Megatron's local CUDA Graph lifecycle and the Transformer Engine helper. The boundary collectives sit inside the captured region, so capture and replay are worth asserting explicitly rather than inferring from the eager tests. Each test records a graph, replays it, requires bitwise equality with eager execution on the first replay, and then perturbs the input to confirm the static input buffer is actually refreshed rather than a stale result being returned. Both backends need the TE RNG tracker, and the local backward capture accumulates weight gradients into main_grad, which DDP normally provides, so the fixture creates those buffers the way tests/unit_tests/transformer/test_cuda_graphs.py does. Teardown goes through delete_cuda_graphs() to clear the global cudagraph record; leaving it in place makes the next test see cudagraph_created with no matching runners. Signed-off-by: Sanghun Cho <highcredit1118@gmail.com>
Describe the communication trade-off against the existing all-to-all Mamba CP path, the three load-balancing modes, the forward and backward flow including which tensors are saved versus recomputed, and the currently unsupported configurations. Also record that the permute modes overlap with megatron.core.context_parallel_layout, whose CpPartitionMode models the same zigzag/contiguous pair and which Gated DeltaNet already uses, and that consolidating onto that helper is intended follow-up work. Signed-off-by: Sanghun Cho <highcredit1118@gmail.com>
Compare the existing all-to-all Mamba CP path against each state-passing load-balancing mode, so the mode can be chosen from measurements rather than by guessing. Only the post-projection part of the mixer is timed. That is where the paths differ; the projections and RMSNorm are identical across paths and would dilute the comparison. The benchmark reproduces just the mixer attributes MambaContextParallel and MambaStatePassingCPAdapter read, which keeps it independent of the surrounding model, spec, and process-group plumbing. Each iteration is reduced across ranks with MAX rather than averaged, since a CP path is only as fast as its slowest rank, and rank alignment happens on a barrier outside the timed interval so waiting for stragglers is not counted twice. Alongside the human-readable lines, one RESULT line per path is emitted for scripted collection. Signed-off-by: Sanghun Cho <highcredit1118@gmail.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
1 participant
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.
Mamba2 state-passing context parallelism
Adds an opt-in, sequence-sharded context-parallel path for Mamba2 training. The existing all-to-all Mamba CP path stays the default and is untouched.
The existing path redistributes activations so every rank holds the whole sequence for a subset of heads and groups, costing
O(local_sequence * hidden)per rank. The state-passing path keeps the sequence shard local and exchanges only what the causal boundary needs — the convolution halo and an FP32 state summary — neither of which scales with sequence length:The SSD transform of one causal segment is affine in its initial state,
S_out = a_block * S_in + S_ext, so each rank computes its own(S_ext, a_block), all-gathers the packed summary, and recovers its true initial state with an exclusive causal scan. The gather is issued asynchronously and overlaps the localCBcomputation. Backward mirrors this with a reverse boundary scan and recomputes everything that scales with sequence length rather than saving it.Configuration
Standard Megatron CP hands the mixer a front/back balanced (zigzag) shard, so the mixer rejects
none, which expects an already-contiguous shard and exists for direct calls into the kernels.permute_p2ppermute_a2aall_to_all_singlevirtualWhy
virtualis the mode to prefervirtualwas fastest in every valid measured cell, and the reason is structural rather than incidental:O(local_L * hidden)twice in forward (in and out) and again in backward — reintroducing exactly the activation traffic that state passing exists to remove.virtualreinterprets the layout instead, so its only communication is the sequence-length-independent boundary summary.2 * cp_sizeentries instead ofcp_size. That is a handful of KiB and still independent of sequence length.2 * batchprogram instances over the same number of tokens. This is visible at large CP sizes where each rank's local sequence is short and the kernels would otherwise be launch-bound — at CP=8,virtualstill gives 1.2–1.3x while the permute modes are at parity with A2A.The trade-off is a shape constraint: each virtual segment must be a multiple of the SSD
chunk_size, so the sequence length must be a multiple of4 * cp_size * chunk_sizerather than2 * cp_size * chunk_size.Benchmark
Measured on a single B200 node on the development branch this was ported from (2026-07-21, PyTorch
2.11.0a0+a6c236b9fd.nv26.03, CUDA 13.2, NCCL 2.29.7, BF16). Mamba core shapenheads=64, headdim=64, ngroups=8, dstate=128, d_conv=4, chunk_size=128. Each cell is the rank-maximum p50 of 10 forward+backward iterations after 5 warmups, asp50 ms (speedup vs A2A). Only the post-projection mixer core is timed — seeexamples/mamba_state_passing_context_parallel/.CP=2:
CP=4:
CP=8:
TP2-local shape (global heads/groups divided to a
32/4local shape, no TP communication), CP=4, MBS=3, 20 iterations after 5 warmups:virtualtrades a small amount of allocated memory (+0.9% to +2.7% in this workload) for the latency win.A short re-run of the benchmark included in this PR, on the current branch and a different B200 host (CP=4, L=4096, batch 1), reproduces the ordering: A2A 5.89 ms, permute-p2p 5.41 ms (1.088x), permute-a2a 5.16 ms (1.141x), virtual 4.61 ms (1.277x).
Caveat on
permute_a2aThe
FAILcells above are runs that terminated with a CUDA illegal memory access. They appear first and only inpermute_a2aat large exchange payloads. This has not been diagnosed, so the table should not be read as a statement of supported shapes for that mode. It is one more reason to prefervirtual, and reviewers may reasonably ask forpermute_a2ato be dropped until it is understood.Overlap with
context_parallel_layoutpermute_p2pandpermute_a2aconvert the balanced layout to a contiguous causal shard, which is the same conversion Megatron Core already models asCpPartitionMode = Literal["zigzag", "contiguous"]. That abstraction was introduced with DeepSeek-v4 CP support in #5087, refactored intomegatron/core/context_parallel_layout/in #6387 (+ #6515 for main branch), and is exercised by MTP fixes such as #5706 and #6246. Gated DeltaNet already uses it at its module entry point viaconvert_module_input_tensors_cp_partition_mode()when running its chunkwise (sequence-sharded) CP path — the closest existing analogue to what this PR does for Mamba2.The permutation here is implemented separately for two reasons: it is driven from inside the fused Conv+SSD autograd function rather than at the module entry point, and it adds a point-to-point backend the shared helper does not have. This is noted in the module docstring and the feature doc, and consolidating onto the shared helper is intended follow-up work rather than a judgement that the shared path is unsuitable. Happy to do that consolidation in this PR instead if reviewers prefer it; the recommended
virtualmode does not depend on either implementation.Not supported yet
Each is rejected with an explicit assertion rather than silently producing wrong results:
seq_idxand packed sequencesinitial_states--mamba-training-ssm-states-dtypeTesting
All new tests pass on 8xB200, on every rank. Dimensions are Nemotron-3 Nano's mixer shape; the parametrization varies batch size and sequence length, which is what the CP sharding and boundary exchange actually depend on.
tests/unit_tests/ssm/ops/test_ssd_state_passing_cp.pytests/unit_tests/ssm/test_mamba_mixer_state_passing_cp.pytests/unit_tests/ssm/test_mamba_state_passing_cp_cuda_graph.pyGradient comparisons use a relative-RMS criterion rather than elementwise
allclose, because the state-passing path reassociates the SSD scan across CP boundaries, changing BF16 rounding without changing the mathematics. Observed errors are ~1e-6 for the output and most gradients acrossL ∈ {2048, 8192, 32768};ddt_biasis the loosest at ~1.5e-3 against a 3e-3 threshold, stable across sequence lengths.The comparisons are all evaluated before anything is asserted and the verdict is all-reduced, so a mismatch fails every rank identically instead of leaving some ranks inside a collective, which would deadlock the run instead of reporting a failure.
tests/unit_tests/ssm/otherwise unchanged: 122 passed, 150 skipped, and 1 pre-existing failure unrelated to this PR (test_hyper_connection_gdn_gpu_forwardneedsflash-linear-attention, which is absent from the environment used here).pylint10.00/10 andruffclean on the changed files undermegatron/coreandtests/.