From 89d0b3a34781d36b1eaeda761c68404a92e84053 Mon Sep 17 00:00:00 2001 From: Sanghun Cho Date: Mon, 24 Aug 2026 09:14:35 +0000 Subject: [PATCH 1/5] Add state-passing context-parallel kernels for Mamba2 SSD 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 --- megatron/core/ssm/ops/ssd_state_passing_cp.py | 1572 +++++++++++++++++ .../ssm/mamba_state_passing_cp_utils.py | 136 ++ .../ssm/ops/test_ssd_state_passing_cp.py | 305 ++++ 3 files changed, 2013 insertions(+) create mode 100644 megatron/core/ssm/ops/ssd_state_passing_cp.py create mode 100644 tests/unit_tests/ssm/mamba_state_passing_cp_utils.py create mode 100644 tests/unit_tests/ssm/ops/test_ssd_state_passing_cp.py diff --git a/megatron/core/ssm/ops/ssd_state_passing_cp.py b/megatron/core/ssm/ops/ssd_state_passing_cp.py new file mode 100644 index 00000000000..af7a716df6f --- /dev/null +++ b/megatron/core/ssm/ops/ssd_state_passing_cp.py @@ -0,0 +1,1572 @@ +"""Context-parallel state passing for Mamba2 SSD. + +The production autograd path combines causal Conv1d and the local SSD scan, then +exchanges only FP32 state summaries and scalar block decays. Backward recomputes +the conv output, applies the reverse state-passing CP boundary scan, and writes SSD gradients +directly into the buffers consumed by causal Conv1d backward. + +Verified by tests/unit_tests/ssm/ops/test_ssd_state_passing_cp.py. + +Note on the ``permute_p2p`` / ``permute_a2a`` load-balancing modes: converting +Megatron's balanced (zigzag) CP layout to a contiguous causal shard overlaps with +``megatron.core.context_parallel_layout``, whose ``CpPartitionMode`` models the +same zigzag/contiguous pair and which GDN already uses via +``convert_module_input_tensors_cp_partition_mode``. The permutation here is kept +self-contained because it is driven from inside the fused Conv+SSD autograd +function rather than at the module entry point, and because it adds a P2P backend +that the shared helper does not have. Consolidating onto the shared helper is +intended follow-up work, not a decision that the shared path is unsuitable. +""" + +from typing import List, Optional, Tuple + +import torch +import torch.distributed as dist +import triton +import triton.language as tl +from einops import rearrange +from mamba_ssm.ops.triton.ssd_combined import ( + _bmm_chunk_bwd, + _bmm_chunk_fwd, + _chunk_cumsum_bwd, + _chunk_cumsum_fwd, + _chunk_scan_bwd_dC, + _chunk_scan_bwd_dcb, + _chunk_scan_bwd_ddAcs_stable, + _chunk_scan_bwd_dstates, + _chunk_scan_bwd_dz, + _chunk_scan_chunk_state_bwd_dx, + _chunk_scan_fwd, + _chunk_state_bwd_db, + _chunk_state_fwd, + _state_passing_bwd, + _state_passing_fwd, + causal_conv1d_bwd_function, + causal_conv1d_fwd_function, + ensure_stride, +) +from mamba_ssm.utils.torch import custom_bwd, custom_fwd + + +def _all_gather_stack(x, group, async_op=False): + """Gather directly into a stacked contiguous tensor.""" + world = dist.get_world_size(group) + x = x.contiguous() + gathered = torch.empty((world, *x.shape), device=x.device, dtype=x.dtype) + work = dist.all_gather_into_tensor(gathered, x, group=group, async_op=async_op) + return gathered, work + + +@triton.jit +def _state_passing_summary_fwd_kernel( + # Pointers to matrices + states_ptr, + dA_cs_ptr, + summary_ptr, + # Matrix dimensions + dim: tl.constexpr, + nchunks: tl.constexpr, + nheads: tl.constexpr, + state_numel: tl.constexpr, + # Strides + stride_states_batch: tl.constexpr, + stride_states_chunk: tl.constexpr, + stride_states_head: tl.constexpr, + stride_states_dim: tl.constexpr, + stride_dA_cs_batch: tl.constexpr, + stride_dA_cs_head: tl.constexpr, + stride_dA_cs_chunk: tl.constexpr, + # Meta-parameters + BATCH_STRIDE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid_b = tl.program_id(axis=1) + pid_h = tl.program_id(axis=2) + pid_m = tl.program_id(axis=0) + batch_idx = pid_b * BATCH_STRIDE + states_ptr += batch_idx * stride_states_batch + pid_h * stride_states_head + dA_cs_ptr += batch_idx * stride_dA_cs_batch + pid_h * stride_dA_cs_head + + offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + states_ptrs = states_ptr + offs_m * stride_states_dim + state = tl.zeros((BLOCK_SIZE,), dtype=tl.float32) + total_decay = 0.0 + for _ in range(nchunks): + new_state = tl.load(states_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32) + dA_cs = tl.load(dA_cs_ptr).to(tl.float32) + state = tl.exp(dA_cs) * state + new_state + total_decay += dA_cs + states_ptrs += stride_states_chunk + dA_cs_ptr += stride_dA_cs_chunk + + state_offset = (batch_idx * nheads + pid_h) * dim + tl.store(summary_ptr + state_offset + offs_m, state, mask=offs_m < dim) + decay_offset = state_numel + batch_idx * nheads + pid_h + tl.store(summary_ptr + decay_offset, tl.exp(total_decay), mask=pid_m == 0) + + +def _state_passing_summary_fwd(states, dA_chunk_cumsum, *, active_batch=None, batch_stride=1): + """Pack final states and block decays into one contiguous FP32 payload. + + active_batch and batch_stride select the interleaved front slots used by the + virtual state-passing CP rank-0 summary optimization. + """ + assert states.dtype == torch.float32 + assert dA_chunk_cumsum.dtype == torch.float32 + batch, nchunks, nheads, dim = states.shape + active_batch = batch if active_batch is None else active_batch + assert active_batch > 0 + assert (active_batch - 1) * batch_stride < batch + + state_numel = batch * nheads * dim + summary_numel = state_numel + batch * nheads + if active_batch < batch: + summary = torch.zeros(summary_numel, device=states.device, dtype=torch.float32) + else: + summary = torch.empty(summary_numel, device=states.device, dtype=torch.float32) + + block_size = 256 + grid = (triton.cdiv(dim, block_size), active_batch, nheads) + _state_passing_summary_fwd_kernel[grid]( + states, + dA_chunk_cumsum, + summary, + dim=dim, + nchunks=nchunks, + nheads=nheads, + state_numel=state_numel, + stride_states_batch=states.stride(0), + stride_states_chunk=states.stride(1), + stride_states_head=states.stride(2), + stride_states_dim=states.stride(3), + stride_dA_cs_batch=dA_chunk_cumsum.stride(0), + stride_dA_cs_head=dA_chunk_cumsum.stride(1), + stride_dA_cs_chunk=dA_chunk_cumsum.stride(2), + BATCH_STRIDE=batch_stride, + BLOCK_SIZE=block_size, + ) + return summary, state_numel + + +@triton.jit +def _state_passing_summary_bwd_kernel( + # Pointers to matrices + dout_ptr, + dA_cs_ptr, + dinitstates_ptr, + # Matrix dimensions + dim: tl.constexpr, + nchunks: tl.constexpr, + # Strides + stride_dout_batch: tl.constexpr, + stride_dout_chunk: tl.constexpr, + stride_dout_head: tl.constexpr, + stride_dout_dim: tl.constexpr, + stride_dA_cs_batch: tl.constexpr, + stride_dA_cs_head: tl.constexpr, + stride_dA_cs_chunk: tl.constexpr, + stride_dinitstates_batch: tl.constexpr, + stride_dinitstates_head: tl.constexpr, + stride_dinitstates_dim: tl.constexpr, + # Meta-parameters + BLOCK_SIZE: tl.constexpr, +): + pid_b = tl.program_id(axis=1) + pid_h = tl.program_id(axis=2) + pid_m = tl.program_id(axis=0) + dout_ptr += ( + pid_b * stride_dout_batch + pid_h * stride_dout_head + (nchunks - 1) * stride_dout_chunk + ) + dA_cs_ptr += ( + pid_b * stride_dA_cs_batch + pid_h * stride_dA_cs_head + (nchunks - 1) * stride_dA_cs_chunk + ) + + offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + dout_ptrs = dout_ptr + offs_m * stride_dout_dim + dstates = tl.zeros((BLOCK_SIZE,), dtype=tl.float32) + for _ in range(nchunks): + dout = tl.load(dout_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32) + dA_cs = tl.load(dA_cs_ptr).to(tl.float32) + dstates = tl.exp(dA_cs) * dstates + dout + dout_ptrs -= stride_dout_chunk + dA_cs_ptr -= stride_dA_cs_chunk + + dinitstates_ptr += pid_b * stride_dinitstates_batch + pid_h * stride_dinitstates_head + tl.store(dinitstates_ptr + offs_m * stride_dinitstates_dim, dstates, mask=offs_m < dim) + + +def _state_passing_summary_bwd(dout, dA_chunk_cumsum): + """Compute only the gradient entering the first local chunk.""" + assert dout.dtype == torch.float32 + assert dA_chunk_cumsum.dtype == torch.float32 + batch, nchunks, nheads, dim = dout.shape + assert dA_chunk_cumsum.shape == (batch, nheads, nchunks) + + dinitstates = torch.empty((batch, nheads, dim), device=dout.device, dtype=torch.float32) + block_size = 256 + grid = (triton.cdiv(dim, block_size), batch, nheads) + _state_passing_summary_bwd_kernel[grid]( + dout, + dA_chunk_cumsum, + dinitstates, + dim=dim, + nchunks=nchunks, + stride_dout_batch=dout.stride(0), + stride_dout_chunk=dout.stride(1), + stride_dout_head=dout.stride(2), + stride_dout_dim=dout.stride(3), + stride_dA_cs_batch=dA_chunk_cumsum.stride(0), + stride_dA_cs_head=dA_chunk_cumsum.stride(1), + stride_dA_cs_chunk=dA_chunk_cumsum.stride(2), + stride_dinitstates_batch=dinitstates.stride(0), + stride_dinitstates_head=dinitstates.stride(1), + stride_dinitstates_dim=dinitstates.stride(2), + BLOCK_SIZE=block_size, + ) + return dinitstates + + +@triton.jit +def _state_passing_boundary_scan_kernel( + # Pointers to matrices + boundary_values_ptr, + block_decays_ptr, + boundary_output_ptr, + # Matrix dimensions + dim: tl.constexpr, + world: tl.constexpr, + rank: tl.constexpr, + # Strides + stride_values_rank: tl.constexpr, + stride_values_batch: tl.constexpr, + stride_values_head: tl.constexpr, + stride_decays_rank: tl.constexpr, + stride_decays_batch: tl.constexpr, + stride_decays_head: tl.constexpr, + stride_output_batch: tl.constexpr, + stride_output_head: tl.constexpr, + # Meta-parameters + VIRTUAL_STATE_PASSING_CP: tl.constexpr, + REVERSE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid_b = tl.program_id(axis=1) + pid_h = tl.program_id(axis=2) + pid_m = tl.program_id(axis=0) + offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + state = tl.zeros((BLOCK_SIZE,), dtype=tl.float32) + + nsegments = 2 * world if VIRTUAL_STATE_PASSING_CP else world + for step in range(nsegments): + segment = nsegments - 1 - step if REVERSE else step + if VIRTUAL_STATE_PASSING_CP: + if segment == rank: + tl.store( + boundary_output_ptr + + (2 * pid_b) * stride_output_batch + + pid_h * stride_output_head + + offs_m, + state, + mask=offs_m < dim, + ) + if segment == nsegments - 1 - rank: + tl.store( + boundary_output_ptr + + (2 * pid_b + 1) * stride_output_batch + + pid_h * stride_output_head + + offs_m, + state, + mask=offs_m < dim, + ) + source_rank = segment if segment < world else nsegments - 1 - segment + source_batch = 2 * pid_b if segment < world else 2 * pid_b + 1 + include = True + else: + source_rank = segment + source_batch = pid_b + include = source_rank > rank if REVERSE else source_rank < rank + + if include: + new_state = tl.load( + boundary_values_ptr + + source_rank * stride_values_rank + + source_batch * stride_values_batch + + pid_h * stride_values_head + + offs_m, + mask=offs_m < dim, + other=0.0, + ).to(tl.float32) + scale = tl.load( + block_decays_ptr + + source_rank * stride_decays_rank + + source_batch * stride_decays_batch + + pid_h * stride_decays_head + ).to(tl.float32) + state = scale * state + new_state + + if not VIRTUAL_STATE_PASSING_CP: + tl.store( + boundary_output_ptr + pid_b * stride_output_batch + pid_h * stride_output_head + offs_m, + state, + mask=offs_m < dim, + ) + + +def _state_passing_boundary_scan( + boundary_values, block_decays, rank, state_passing_cp_virtual, reverse +): + """Scan gathered summaries and return this rank's exclusive causal boundary.""" + assert boundary_values.dtype == torch.float32 + assert block_decays.dtype == torch.float32 + world, batch, nheads = boundary_values.shape[:3] + state_shape = boundary_values.shape[3:] + boundary_values_flat = boundary_values.view(world, batch, nheads, -1) + dim = boundary_values_flat.shape[-1] + boundary_output = torch.empty( + batch, nheads, dim, device=boundary_values.device, dtype=torch.float32 + ) + + block_size = 256 + if state_passing_cp_virtual: + assert batch % 2 == 0 + scan_batch = batch // 2 + else: + scan_batch = batch + grid = (triton.cdiv(dim, block_size), scan_batch, nheads) + _state_passing_boundary_scan_kernel[grid]( + boundary_values_flat, + block_decays, + boundary_output, + dim=dim, + world=world, + rank=rank, + stride_values_rank=boundary_values_flat.stride(0), + stride_values_batch=boundary_values_flat.stride(1), + stride_values_head=boundary_values_flat.stride(2), + stride_decays_rank=block_decays.stride(0), + stride_decays_batch=block_decays.stride(1), + stride_decays_head=block_decays.stride(2), + stride_output_batch=boundary_output.stride(0), + stride_output_head=boundary_output.stride(1), + VIRTUAL_STATE_PASSING_CP=state_passing_cp_virtual, + REVERSE=reverse, + BLOCK_SIZE=block_size, + ) + return boundary_output.view(batch, nheads, *state_shape) + + +def _mamba_chunk_scan_combined_state_passing_cp_fwd( + x, + dt, + A, + B, + C, + chunk_size, + D=None, + z=None, + dt_bias=None, + initial_states=None, + seq_idx=None, + cu_seqlens=None, + dt_softplus=False, + dt_limit=(0.0, float("inf")), + state_passing_cp_group=None, + state_passing_cp_virtual=False, +): + """Mamba combined forward with state-passing CP summary and boundary exchange.""" + assert state_passing_cp_group is not None + batch, seqlen, nheads, headdim = x.shape + _, _, ngroups, dstate = B.shape + assert nheads % ngroups == 0 + assert B.shape == (batch, seqlen, ngroups, dstate) + assert x.shape == (batch, seqlen, nheads, headdim) + assert dt.shape == (batch, seqlen, nheads) + assert A.shape == (nheads,) + assert C.shape == B.shape + if z is not None: + assert z.shape == x.shape + if D is not None: + assert D.shape == (nheads, headdim) or D.shape == (nheads,) + if seq_idx is not None: + assert seq_idx.shape == (batch, seqlen) + if B.stride(-1) != 1: + B = B.contiguous() + if C.stride(-1) != 1: + C = C.contiguous() + if x.stride(-1) != 1 and x.stride(1) != 1: + x = x.contiguous() + if z is not None and z.stride(-1) != 1 and z.stride(1) != 1: + z = z.contiguous() + if D is not None and D.stride(-1) != 1: + D = D.contiguous() + if initial_states is not None: + assert initial_states.shape == (batch, nheads, headdim, dstate) + + assert ( + initial_states is None + ), "external initial_states are not implemented for state-passing CP" + assert seq_idx is None, "seq_idx is not implemented for state-passing CP" + assert cu_seqlens is None, "cu_seqlens is not implemented for state-passing CP" + state_passing_cp_size = dist.get_world_size(state_passing_cp_group) + rank = dist.get_rank(state_passing_cp_group) + if state_passing_cp_virtual: + assert ( + batch % 2 == 0 + ), "virtual state-passing CP expects front/back packed on the batch axis" + virtual_batch = batch // 2 if state_passing_cp_virtual else None + + dA_cumsum, dt = _chunk_cumsum_fwd( + dt, A, chunk_size, dt_bias=dt_bias, dt_softplus=dt_softplus, dt_limit=dt_limit + ) + states = _chunk_state_fwd(B, x, dt, dA_cumsum, seq_idx=seq_idx, states_in_fp32=True) + state_passing_initial_states = None + state_passing_gathered_decays = None + if state_passing_cp_size > 1: + state_dim = headdim * dstate + if not state_passing_cp_virtual and rank == state_passing_cp_size - 1: + # Collective shapes are fixed even though no later rank consumes this summary. + state_numel = batch * nheads * state_dim + packed_summary = torch.zeros( + state_numel + batch * nheads, device=x.device, dtype=torch.float32 + ) + else: + if state_passing_cp_virtual and rank == 0: + # Rank 0's causal-last segment has no successor. + active_batch, batch_stride = virtual_batch, 2 + else: + active_batch, batch_stride = None, 1 + packed_summary, state_numel = _state_passing_summary_fwd( + rearrange(states, "... p n -> ... (p n)"), + dA_cumsum[:, :, :, -1], + active_batch=active_batch, + batch_stride=batch_stride, + ) + + gathered_summary, gather_work = _all_gather_stack( + packed_summary, state_passing_cp_group, async_op=True + ) + CB = _bmm_chunk_fwd(C, B, chunk_size, seq_idx=seq_idx, output_dtype=torch.float32) + gather_work.wait() + gathered_states = gathered_summary[:, :state_numel].view( + state_passing_cp_size, batch, nheads, headdim, dstate + ) + state_passing_gathered_decays = gathered_summary[:, state_numel:].view( + state_passing_cp_size, batch, nheads + ) + state_passing_initial_states = _state_passing_boundary_scan( + gathered_states, + state_passing_gathered_decays, + rank, + state_passing_cp_virtual=state_passing_cp_virtual, + reverse=False, + ) + state_passing_gathered_decays = state_passing_gathered_decays.contiguous() + + states, final_states = _state_passing_fwd( + rearrange(states, "... p n -> ... (p n)"), + dA_cumsum[:, :, :, -1], + initial_states=( + rearrange(state_passing_initial_states, "... p n -> ... (p n)") + if state_passing_initial_states is not None + else None + ), + seq_idx=seq_idx, + chunk_size=chunk_size, + out_dtype=C.dtype, + ) + states, final_states = [ + rearrange(t, "... (p n) -> ... p n", n=dstate) for t in [states, final_states] + ] + if state_passing_cp_size == 1: + CB = _bmm_chunk_fwd(C, B, chunk_size, seq_idx=seq_idx, output_dtype=torch.float32) + # CB comes from the state-passing branch above when cp_size > 1, and from the + # branch just above otherwise; the two cases are exhaustive. + # pylint: disable-next=possibly-used-before-assignment + out, out_x = _chunk_scan_fwd(CB, x, dt, dA_cumsum, C, states, D=D, z=z, seq_idx=seq_idx) + return ( + out, + out_x, + dt, + dA_cumsum, + states, + final_states, + state_passing_initial_states, + state_passing_gathered_decays, + ) + + +def _mamba_chunk_scan_combined_state_passing_cp_bwd( + dout, + x, + dt, + A, + B, + C, + out, + chunk_size, + D=None, + z=None, + dt_bias=None, + initial_states=None, + dfinal_states=None, + seq_idx=None, + dt_softplus=False, + dt_limit=(0.0, float("inf")), + dx=None, + ddt=None, + dB=None, + dC=None, + dz=None, + recompute_output=False, + state_passing_cp_group=None, + state_passing_initial_states=None, + state_passing_gathered_decays=None, + state_passing_cp_virtual=False, +): + """Mamba combined backward with a state-passing CP reverse-boundary handoff.""" + assert state_passing_cp_group is not None + if dout.stride(-1) != 1: + dout = dout.contiguous() + batch, seqlen, nheads, headdim = x.shape + _, _, ngroups, dstate = B.shape + assert dout.shape == (batch, seqlen, nheads, headdim) + assert dt.shape == (batch, seqlen, nheads) + assert A.shape == (nheads,) + assert nheads % ngroups == 0 + assert B.shape == (batch, seqlen, ngroups, dstate) + assert C.shape == B.shape + assert out.shape == x.shape + if initial_states is not None: + assert initial_states.shape == (batch, nheads, headdim, dstate) + if seq_idx is not None: + assert seq_idx.shape == (batch, seqlen) + if dx is not None: + assert dx.shape == x.shape + if dB is not None: + assert dB.shape == B.shape + dB_given = dB + else: + dB_given = torch.empty_like(B) + if dC is not None: + assert dC.shape == C.shape + dC_given = dC + else: + dC_given = torch.empty_like(C) + if dz is not None: + assert z is not None + assert dz.shape == z.shape + if ddt is not None: + assert ddt.shape == dt.shape + ddt_given = ddt + else: + ddt_given = torch.empty_like(dt) + + assert ( + initial_states is None + ), "external initial_states are not implemented for state-passing CP" + assert dfinal_states is None, "dfinal_states are not implemented for state-passing CP" + assert seq_idx is None, "seq_idx is not implemented for state-passing CP" + state_passing_cp_size = dist.get_world_size(state_passing_cp_group) + rank = dist.get_rank(state_passing_cp_group) + + dt_in = dt.clone() + dA_cumsum, dt = _chunk_cumsum_fwd( + dt_in, A, chunk_size, dt_bias=dt_bias, dt_softplus=dt_softplus, dt_limit=dt_limit + ) + CB = _bmm_chunk_fwd(C, B, chunk_size, seq_idx=seq_idx, output_dtype=torch.float32) + states = _chunk_state_fwd(B, x, dt, dA_cumsum, seq_idx=seq_idx, states_in_fp32=True) + states, _ = _state_passing_fwd( + rearrange(states, "... p n -> ... (p n)"), + dA_cumsum[:, :, :, -1], + initial_states=( + rearrange(state_passing_initial_states, "... p n -> ... (p n)") + if state_passing_initial_states is not None + else None + ), + seq_idx=seq_idx, + chunk_size=chunk_size, + ) + states = rearrange(states, "... (p n) -> ... p n", n=dstate) + + if z is not None: + dz, dout, dD, *rest = _chunk_scan_bwd_dz( + x, + z, + out, + dout, + chunk_size=chunk_size, + has_ddAcs=False, + D=D, + dz=dz, + recompute_output=recompute_output, + ) + outz = rest[0] if recompute_output else out + else: + dz = None + outz = out + + dstates = _chunk_scan_bwd_dstates(C, dA_cumsum, dout, seq_idx=seq_idx, dtype=states.dtype) + + # Everything in this block is state-passing CP-specific. The local kernels run while the + # reverse-boundary all-gather is in flight. + dC_local = dCB = ddA = None + final_state_grad = None + if state_passing_cp_size > 1: + assert ( + state_passing_initial_states is not None and state_passing_gathered_decays is not None + ) + initial_state_grad = _state_passing_summary_bwd( + rearrange(dstates, "... p n -> ... (p n)"), dA_cumsum[:, :, :, -1] + ) + initial_state_grad = rearrange(initial_state_grad, "... (p n) -> ... p n", n=dstate) + gathered_initial_state_grads, gather_work = _all_gather_stack( + initial_state_grad, state_passing_cp_group, async_op=True + ) + + states_for_scan = states.to(x.dtype) + dC_local, ddA_cumsum_prev = _chunk_scan_bwd_dC( + states_for_scan, dA_cumsum, dout, seq_idx=seq_idx, C=C, ngroups=ngroups + ) + del states_for_scan + dCB = _chunk_scan_bwd_dcb(x, dt, dA_cumsum, dout, seq_idx=seq_idx, ngroups=ngroups).to( + CB.dtype + ) + ddA = _chunk_scan_bwd_ddAcs_stable(x, dt, dA_cumsum, dout, CB) + gather_work.wait() + final_state_grad = _state_passing_boundary_scan( + gathered_initial_state_grads, + state_passing_gathered_decays, + rank, + state_passing_cp_virtual=state_passing_cp_virtual, + reverse=True, + ) + + state_bwd_outputs = _state_passing_bwd( + rearrange(states, "... p n -> ... (p n)"), + dA_cumsum[:, :, :, -1], + rearrange(dstates, "... p n -> ... (p n)"), + dfinal_states=( + rearrange(final_state_grad, "... p n -> ... (p n)") + if final_state_grad is not None + else None + ), + seq_idx=seq_idx, + has_initial_states=state_passing_initial_states is not None, + dstates_dtype=x.dtype, + # dC already consumed the converted states in the overlapped state-passing CP path. + states_dtype=x.dtype if dC_local is None else None, + chunk_size=chunk_size, + ) + dstates, ddA_chunk_cumsum, _ = state_bwd_outputs[:3] + dstates = rearrange(dstates, "... (p n) -> ... p n", n=dstate) + + if dC_local is None: + states = rearrange(state_bwd_outputs[3], "... (p n) -> ... p n", n=dstate) + dC_local, ddA_cumsum_prev = _chunk_scan_bwd_dC( + states.to(x.dtype), dA_cumsum, dout, seq_idx=seq_idx, C=C, ngroups=ngroups + ) + dCB = _chunk_scan_bwd_dcb(x, dt, dA_cumsum, dout, seq_idx=seq_idx, ngroups=ngroups).to( + CB.dtype + ) + ddA = _chunk_scan_bwd_ddAcs_stable(x, dt, dA_cumsum, dout, CB) + + dx, ddt, dD_from_x = _chunk_scan_chunk_state_bwd_dx( + x, dt, dA_cumsum, B, CB, dout, dstates, D=D, seq_idx=seq_idx, dx=dx + ) + dB, ddA_next = _chunk_state_bwd_db( + x, dt, dA_cumsum, dstates, seq_idx=seq_idx, B=B, ngroups=ngroups + ) + _bmm_chunk_bwd(C, dCB, residual=dB, out=dB_given) + _bmm_chunk_bwd(B, rearrange(dCB, "... l s -> ... s l"), residual=dC_local, out=dC_given) + if z is None: + dD = dD_from_x + + # ddA_cumsum_prev is set by whichever of the two exhaustive dC_local branches + # ran, and dD by whichever of the two exhaustive `z` branches ran. + # pylint: disable-next=possibly-used-before-assignment + ddA_cumsum_prev[..., -1] += ddA_chunk_cumsum + ddA_prev = ddA_cumsum_prev.flip([-1]).cumsum(dim=-1).flip([-1]) + ddA += ddA_next + ddA_prev + ddt_given, dA, ddt_bias = _chunk_cumsum_bwd( + ddA, + ddt, + dt_in, + A, + dt_bias=dt_bias, + dt_softplus=dt_softplus, + dt_limit=dt_limit, + ddt=ddt_given, + ) + + dinitial_states = None + # pylint: disable-next=possibly-used-before-assignment + return_vals = (dx, ddt_given, dA, dB_given, dC_given, dD, dz, ddt_bias, dinitial_states) + return return_vals if not recompute_output else (*return_vals, outz) + + +# Mixer integration, conv boundary exchange, and state-passing CP layout handling. +def _pack_state_passing_cp_virtual_segments(x: torch.Tensor) -> torch.Tensor: + """View balanced chunks as interleaved front/back segments on the batch axis.""" + assert x.shape[1] % 2 == 0, "virtual state-passing CP requires two equal local sequence chunks" + batch = x.shape[0] + half = x.shape[1] // 2 + return x.reshape(2 * batch, half, *x.shape[2:]) + + +def _unpack_state_passing_cp_virtual_segments(x: torch.Tensor) -> torch.Tensor: + """View interleaved front/back segments as the balanced local sequence layout.""" + assert x.shape[0] % 2 == 0 + batch = x.shape[0] // 2 + return x.reshape(batch, 2 * x.shape[1], *x.shape[2:]) + + +def _state_passing_cp_chunk_owner_slot( + chunk_id: int, state_passing_cp_size: int, layout: str +) -> Tuple[int, int]: + if layout == "contiguous": + return chunk_id // 2, chunk_id % 2 + if chunk_id < state_passing_cp_size: + return chunk_id, 0 + return 2 * state_passing_cp_size - 1 - chunk_id, 1 + + +def _state_passing_cp_local_chunk_ids( + rank: int, state_passing_cp_size: int, layout: str +) -> Tuple[int, int]: + if layout == "contiguous": + return 2 * rank, 2 * rank + 1 + return rank, 2 * state_passing_cp_size - 1 - rank + + +def permute_state_passing_cp_sequence_chunks( + x: torch.Tensor, + state_passing_cp_group: torch.distributed.ProcessGroup, + undo_load_balancing: bool, + backend: str = "p2p", +) -> torch.Tensor: + """Exchange sequence chunks between the balanced and contiguous CP layouts. + + ``undo_load_balancing=True`` converts Megatron's balanced front/back layout + into a contiguous causal shard; ``False`` restores the balanced layout. The + backward pass applies the inverse permutation. + + This mirrors the zigzag/contiguous conversion in + ``megatron.core.context_parallel_layout``; see this module's docstring for why + it is currently implemented separately. + """ + if dist.get_world_size(state_passing_cp_group) == 1: + return x + if backend not in ("p2p", "a2a"): + raise ValueError(f"Unsupported state-passing CP sequence permutation backend {backend!r}") + return _StatePassingCPSequenceChunkPermutationFn.apply( + x, state_passing_cp_group, undo_load_balancing, backend + ) + + +def undo_state_passing_cp_load_balancing( + x: torch.Tensor, state_passing_cp_group: torch.distributed.ProcessGroup, backend: str = "p2p" +) -> torch.Tensor: + """Convert a balanced CP shard into a contiguous causal shard.""" + return permute_state_passing_cp_sequence_chunks(x, state_passing_cp_group, True, backend) + + +def redo_state_passing_cp_load_balancing( + x: torch.Tensor, state_passing_cp_group: torch.distributed.ProcessGroup, backend: str = "p2p" +) -> torch.Tensor: + """Convert a contiguous causal shard back into the balanced CP layout.""" + return permute_state_passing_cp_sequence_chunks(x, state_passing_cp_group, False, backend) + + +class _StatePassingCPSequenceChunkPermutationFn(torch.autograd.Function): + @staticmethod + def forward( + ctx, + x, + state_passing_cp_group: torch.distributed.ProcessGroup, + undo_load_balancing: bool, + backend: str, + ): + ctx.state_passing_cp_group = state_passing_cp_group + ctx.undo_load_balancing = undo_load_balancing + ctx.backend = backend + return _permute_state_passing_cp_sequence_chunks( + x, state_passing_cp_group, undo_load_balancing, backend + ) + + @staticmethod + def backward(ctx, grad_output): + grad_x = _permute_state_passing_cp_sequence_chunks( + grad_output, ctx.state_passing_cp_group, not ctx.undo_load_balancing, ctx.backend + ) + return grad_x, None, None, None + + +def _permute_state_passing_cp_sequence_chunks( + x: torch.Tensor, + state_passing_cp_group: torch.distributed.ProcessGroup, + undo_load_balancing: bool, + backend: str = "p2p", +) -> torch.Tensor: + if backend == "a2a": + return _permute_state_passing_cp_sequence_chunks_a2a( + x, state_passing_cp_group, undo_load_balancing + ) + if backend != "p2p": + raise ValueError(f"Unsupported state-passing CP sequence permutation backend {backend!r}") + return _permute_state_passing_cp_sequence_chunks_p2p( + x, state_passing_cp_group, undo_load_balancing + ) + + +def _permute_state_passing_cp_sequence_chunks_p2p( + x: torch.Tensor, + state_passing_cp_group: torch.distributed.ProcessGroup, + undo_load_balancing: bool, +) -> torch.Tensor: + state_passing_cp_size = dist.get_world_size(state_passing_cp_group) + rank = dist.get_rank(state_passing_cp_group) + assert ( + x.size(0) % 2 == 0 + ), "State-passing CP load-balanced sequence shard must contain two chunks" + chunk_len = x.size(0) // 2 + input_layout = "balanced" if undo_load_balancing else "contiguous" + output_layout = "contiguous" if undo_load_balancing else "balanced" + input_chunks = _state_passing_cp_local_chunk_ids(rank, state_passing_cp_size, input_layout) + output_chunks = _state_passing_cp_local_chunk_ids(rank, state_passing_cp_size, output_layout) + + # P2P receive buffers must be non-overlapping and dense. In backward, + # grad_output can inherit a non-standard stride from the downstream layout. + out = torch.empty_like(x, memory_format=torch.contiguous_format) + ops = [] + sends = [] + + for out_slot, chunk_id in enumerate(output_chunks): + src_rank, src_slot = _state_passing_cp_chunk_owner_slot( + chunk_id, state_passing_cp_size, input_layout + ) + out_slice = out[out_slot * chunk_len : (out_slot + 1) * chunk_len] + if src_rank == rank: + src = x[src_slot * chunk_len : (src_slot + 1) * chunk_len] + out_slice.copy_(src) + else: + ops.append( + dist.P2POp(dist.irecv, out_slice, group=state_passing_cp_group, group_peer=src_rank) + ) + + for in_slot, chunk_id in enumerate(input_chunks): + dst_rank, _ = _state_passing_cp_chunk_owner_slot( + chunk_id, state_passing_cp_size, output_layout + ) + if dst_rank != rank: + send = x[in_slot * chunk_len : (in_slot + 1) * chunk_len].contiguous() + sends.append(send) + ops.append( + dist.P2POp(dist.isend, send, group=state_passing_cp_group, group_peer=dst_rank) + ) + + for work in dist.batch_isend_irecv(ops): + work.wait() + return out + + +def _permute_state_passing_cp_sequence_chunks_a2a( + x: torch.Tensor, + state_passing_cp_group: torch.distributed.ProcessGroup, + undo_load_balancing: bool, +) -> torch.Tensor: + state_passing_cp_size = dist.get_world_size(state_passing_cp_group) + rank = dist.get_rank(state_passing_cp_group) + assert ( + x.size(0) % 2 == 0 + ), "State-passing CP load-balanced sequence shard must contain two chunks" + chunk_len = x.size(0) // 2 + input_layout = "balanced" if undo_load_balancing else "contiguous" + output_layout = "contiguous" if undo_load_balancing else "balanced" + input_chunks = _state_passing_cp_local_chunk_ids(rank, state_passing_cp_size, input_layout) + + local_destinations = [ + _state_passing_cp_chunk_owner_slot(chunk_id, state_passing_cp_size, output_layout) + for chunk_id in input_chunks + ] + send_slot_order = sorted(range(2), key=lambda slot: local_destinations[slot]) + local_chunks = (x[:chunk_len], x[chunk_len:]) + send_buf = torch.cat([local_chunks[slot] for slot in send_slot_order], dim=0).contiguous() + + input_split_chunks = [0] * state_passing_cp_size + for dst_rank, _ in local_destinations: + input_split_chunks[dst_rank] += 1 + + output_split_chunks = [0] * state_passing_cp_size + recv_slots_by_source = [[] for _ in range(state_passing_cp_size)] + for src_rank in range(state_passing_cp_size): + src_chunks = _state_passing_cp_local_chunk_ids( + src_rank, state_passing_cp_size, input_layout + ) + src_destinations = [ + _state_passing_cp_chunk_owner_slot(chunk_id, state_passing_cp_size, output_layout) + for chunk_id in src_chunks + ] + for src_slot in sorted(range(2), key=lambda slot: src_destinations[slot]): + dst_rank, dst_slot = src_destinations[src_slot] + if dst_rank == rank: + output_split_chunks[src_rank] += 1 + recv_slots_by_source[src_rank].append(dst_slot) + + input_split_sizes = [count * chunk_len for count in input_split_chunks] + output_split_sizes = [count * chunk_len for count in output_split_chunks] + recv_buf = torch.empty_like(x, memory_format=torch.contiguous_format) + dist.all_to_all_single( + recv_buf, + send_buf, + output_split_sizes=output_split_sizes, + input_split_sizes=input_split_sizes, + group=state_passing_cp_group, + ) + + target_slots: List[Optional[torch.Tensor]] = [None, None] + offset = 0 + for src_rank in range(state_passing_cp_size): + for dst_slot in recv_slots_by_source[src_rank]: + target_slots[dst_slot] = recv_buf[offset : offset + chunk_len] + offset += chunk_len + assert all( + slot is not None for slot in target_slots + ), "Incomplete state-passing CP sequence A2A reassembly" + return torch.cat(target_slots, dim=0) + + +@triton.jit +def _route_virtual_state_passing_cp_conv_boundary_kernel( + gathered, + output, + batch: tl.constexpr, + channels: tl.constexpr, + halo: tl.constexpr, + rank: tl.constexpr, + world: tl.constexpr, + stride_g_rank: tl.constexpr, + stride_g_batch: tl.constexpr, + stride_g_channel: tl.constexpr, + stride_g_halo: tl.constexpr, + stride_o_batch: tl.constexpr, + stride_o_channel: tl.constexpr, + stride_o_halo: tl.constexpr, + REVERSE: tl.constexpr, + BLOCK_C: tl.constexpr, + BLOCK_H: tl.constexpr, +): + pid_b = tl.program_id(0) + pid_c = tl.program_id(1) + slot = pid_b % 2 + local_batch = pid_b // 2 + + if not REVERSE: + if slot == 0: + valid = rank > 0 + source_rank = rank - 1 + source_batch = 2 * local_batch + else: + valid = True + source_rank = rank + 1 + source_batch = 2 * local_batch + 1 + if rank == world - 1: + source_rank = rank + source_batch = 2 * local_batch + else: + if slot == 0: + valid = True + source_rank = rank + 1 + source_batch = 2 * local_batch + if rank == world - 1: + source_rank = rank + source_batch = 2 * local_batch + 1 + else: + valid = rank > 0 + source_rank = rank - 1 + source_batch = 2 * local_batch + 1 + + offsets_c = pid_c * BLOCK_C + tl.arange(0, BLOCK_C) + offsets_h = tl.arange(0, BLOCK_H) + mask = (offsets_c[:, None] < channels) & (offsets_h[None, :] < halo) + values = tl.load( + gathered + + source_rank * stride_g_rank + + source_batch * stride_g_batch + + offsets_c[:, None] * stride_g_channel + + offsets_h[None, :] * stride_g_halo, + mask=mask & valid, + other=0.0, + ) + tl.store( + output + + pid_b * stride_o_batch + + offsets_c[:, None] * stride_o_channel + + offsets_h[None, :] * stride_o_halo, + values, + mask=mask, + ) + + +def _route_virtual_state_passing_cp_conv_boundary( + gathered: torch.Tensor, + batch: int, + rank: int, + world: int, + reverse: bool, + channel_last_output: bool, +) -> torch.Tensor: + channels, halo = gathered.shape[-2:] + if channel_last_output: + output = torch.empty( + 2 * batch, halo, channels, device=gathered.device, dtype=gathered.dtype + ).transpose(1, 2) + else: + output = torch.empty( + 2 * batch, channels, halo, device=gathered.device, dtype=gathered.dtype + ) + block_c = 128 + block_h = triton.next_power_of_2(halo) + _route_virtual_state_passing_cp_conv_boundary_kernel[ + (2 * batch, triton.cdiv(channels, block_c)) + ]( + gathered, + output, + batch=batch, + channels=channels, + halo=halo, + rank=rank, + world=world, + stride_g_rank=gathered.stride(0), + stride_g_batch=gathered.stride(1), + stride_g_channel=gathered.stride(2), + stride_g_halo=gathered.stride(3), + stride_o_batch=output.stride(0), + stride_o_channel=output.stride(1), + stride_o_halo=output.stride(2), + REVERSE=reverse, + BLOCK_C=block_c, + BLOCK_H=block_h, + ) + return output + + +def _gather_causal_conv1d_state_passing_cp_boundary( + boundary: torch.Tensor, + state_passing_cp_group: torch.distributed.ProcessGroup, + state_passing_cp_virtual: bool, + reverse: bool, + channel_last_output: bool, +) -> torch.Tensor: + """Gather a predecessor state or successor gradient for causal conv.""" + world = dist.get_world_size(state_passing_cp_group) + rank = dist.get_rank(state_passing_cp_group) + batch, channels, halo = boundary.shape + gathered = boundary.new_empty(world * batch, channels, halo) + dist.all_gather_into_tensor(gathered, boundary.contiguous(), group=state_passing_cp_group) + gathered = gathered.view(world, batch, channels, halo) + + if state_passing_cp_virtual: + assert batch % 2 == 0 + return _route_virtual_state_passing_cp_conv_boundary( + gathered, + batch // 2, + rank, + world, + reverse=reverse, + channel_last_output=channel_last_output, + ) + + if channel_last_output: + output = torch.empty( + batch, halo, channels, device=boundary.device, dtype=boundary.dtype + ).transpose(1, 2) + else: + output = torch.empty_like(boundary, memory_format=torch.contiguous_format) + source_rank = rank + 1 if reverse else rank - 1 + if 0 <= source_rank < world: + output.copy_(gathered[source_rank]) + else: + output.zero_() + return output + + +def _causal_conv1d_state_passing_cp_fwd( + x, + weight, + bias, + seq_idx=None, + initial_states=None, + final_states_out=None, + activation=None, + state_passing_cp_group=None, + state_passing_cp_virtual=False, +): + """Run causal Conv1d using a gathered state-passing CP predecessor state.""" + assert causal_conv1d_fwd_function is not None + assert state_passing_cp_group is not None + assert seq_idx is None, "seq_idx is not implemented for state-passing CP" + assert final_states_out is None + halo = weight.shape[-1] - 1 + assert halo > 0 and x.size(-1) >= halo + state_passing_conv_initial_states = initial_states + if state_passing_conv_initial_states is None: + state_passing_conv_initial_states = _gather_causal_conv1d_state_passing_cp_boundary( + x[..., -halo:], + state_passing_cp_group, + state_passing_cp_virtual=state_passing_cp_virtual, + reverse=False, + channel_last_output=True, + ) + out = causal_conv1d_fwd_function( + x, + weight, + bias, + seq_idx, + state_passing_conv_initial_states, + final_states_out, + activation in ["silu", "swish"], + ) + return out, state_passing_conv_initial_states + + +def _causal_conv1d_state_passing_cp_bwd( + x, + weight, + bias, + dout, + seq_idx=None, + initial_states=None, + dfinal_states=None, + dx=None, + return_dinitial_states=False, + activation=None, + state_passing_cp_group=None, + state_passing_cp_virtual=False, +): + """Run causal Conv1d backward and route its state-passing CP halo gradient.""" + assert causal_conv1d_bwd_function is not None + assert state_passing_cp_group is not None + assert seq_idx is None, "seq_idx is not implemented for state-passing CP" + assert dfinal_states is None + assert not return_dinitial_states, "the state-passing boundary gradient is consumed internally" + if dout.stride(2) != 1 and dout.stride(1) != 1: + dout = dout.contiguous() + dx, dweight, dbias, dinitial_states = causal_conv1d_bwd_function( + x, + weight, + bias, + dout, + seq_idx, + initial_states, + dfinal_states, + dx, + True, + activation in ["silu", "swish"], + ) + halo = weight.shape[-1] - 1 + grad_tail = _gather_causal_conv1d_state_passing_cp_boundary( + dinitial_states, + state_passing_cp_group, + state_passing_cp_virtual=state_passing_cp_virtual, + reverse=True, + channel_last_output=False, + ) + dx[..., -halo:].add_(grad_tail) + return dx, dweight, dbias + + +class MambaSplitConv1dScanCombinedStatePassingCPFn(torch.autograd.Function): + """State-passing CP counterpart of ``MambaSplitConv1dScanCombinedFn``.""" + + @staticmethod + @custom_fwd + def forward( + ctx, + zxbcdt, + conv1d_weight, + conv1d_bias, + dt_bias, + A, + D, + chunk_size, + initial_states=None, + seq_idx=None, + dt_limit=(0.0, float("inf")), + return_final_states=False, + activation="silu", + rmsnorm_weight=None, + rmsnorm_eps=1e-6, + outproj_weight=None, + outproj_bias=None, + headdim=None, + ngroups=1, + norm_before_gate=True, + state_passing_cp_group=None, + state_passing_cp_virtual=False, + ): + assert activation in [None, "silu", "swish"] + assert state_passing_cp_group is not None + assert ( + initial_states is None + ), "external initial_states are not implemented for state-passing CP" + assert seq_idx is None, "seq_idx packed input is not implemented for state-passing CP" + assert ( + not return_final_states + ), "return_final_states is not implemented for state-passing CP" + assert rmsnorm_weight is None, "fused RMSNorm is not implemented for state-passing CP" + assert outproj_weight is None and outproj_bias is None + if D.dim() == 1: + assert headdim is not None + (nheads,) = D.shape + else: + nheads, headdim = D.shape + batch, seqlen, _ = zxbcdt.shape + dim = nheads * headdim + assert nheads % ngroups == 0 + dstate = (conv1d_weight.shape[0] - dim) // ngroups // 2 + d_nonssm = (zxbcdt.shape[-1] - 2 * dim - 2 * ngroups * dstate - nheads) // 2 + assert d_nonssm >= 0 + assert d_nonssm == 0, "non-SSM channels are not implemented for state-passing CP" + assert zxbcdt.shape == ( + batch, + seqlen, + 2 * d_nonssm + 2 * dim + 2 * ngroups * dstate + nheads, + ) + assert dt_bias.shape == (nheads,) + assert A.shape == (nheads,) + zx0, z, xBC, dt = torch.split( + zxbcdt, [2 * d_nonssm, dim, dim + 2 * ngroups * dstate, nheads], dim=-1 + ) + seq_idx = seq_idx.contiguous() if seq_idx is not None else None + xBC_conv, state_passing_conv_initial_states = _causal_conv1d_state_passing_cp_fwd( + rearrange(ensure_stride(xBC), "b s d -> b d s"), + conv1d_weight, + conv1d_bias, + seq_idx, + None, + None, + activation, + state_passing_cp_group, + state_passing_cp_virtual, + ) + xBC_conv = rearrange(xBC_conv, "b d s -> b s d") + x, B, C = torch.split(xBC_conv, [dim, ngroups * dstate, ngroups * dstate], dim=-1) + x = rearrange(x, "b l (h p) -> b l h p", h=nheads) + B = rearrange(B, "b l (g n) -> b l g n", g=ngroups) + C = rearrange(C, "b l (g n) -> b l g n", g=ngroups) + z = rearrange(z, "b l (h p) -> b l h p", h=nheads) if z is not None else None + + ( + out, + out_x, + dt_out, + dA_cumsum, + states, + final_states, + state_passing_initial_states, + state_passing_gathered_decays, + ) = _mamba_chunk_scan_combined_state_passing_cp_fwd( + x, + dt, + A, + B, + C, + chunk_size=chunk_size, + D=D, + z=z, + dt_bias=dt_bias, + initial_states=initial_states, + seq_idx=seq_idx, + dt_softplus=True, + dt_limit=dt_limit, + state_passing_cp_group=state_passing_cp_group, + state_passing_cp_virtual=state_passing_cp_virtual, + ) + out = rearrange(out, "b s h p -> b s (h p)") + rstd = None + + ctx.save_for_backward( + zxbcdt, + conv1d_weight, + conv1d_bias, + out_x, + A, + D, + dt_bias, + initial_states, + seq_idx, + rmsnorm_weight, + rstd, + outproj_weight, + outproj_bias, + state_passing_conv_initial_states, + state_passing_initial_states, + state_passing_gathered_decays, + ) + ctx.dt_limit = dt_limit + ctx.return_final_states = return_final_states + ctx.activation = activation + ctx.rmsnorm_eps = rmsnorm_eps + ctx.norm_before_gate = norm_before_gate + ctx.chunk_size = chunk_size + ctx.headdim = headdim + ctx.ngroups = ngroups + ctx.state_passing_cp_group = state_passing_cp_group + ctx.state_passing_cp_virtual = state_passing_cp_virtual + return out if not return_final_states else (out, final_states) + + @staticmethod + @custom_bwd + def backward(ctx, dout, *args): + ( + zxbcdt, + conv1d_weight, + conv1d_bias, + out, + A, + D, + dt_bias, + initial_states, + seq_idx, + rmsnorm_weight, + rstd, + outproj_weight, + outproj_bias, + state_passing_conv_initial_states, + state_passing_initial_states, + state_passing_gathered_decays, + ) = ctx.saved_tensors + dfinal_states = args[0] if ctx.return_final_states else None + headdim = ctx.headdim + nheads = D.shape[0] + dim = nheads * headdim + assert nheads % ctx.ngroups == 0 + dstate = (conv1d_weight.shape[0] - dim) // ctx.ngroups // 2 + d_nonssm = (zxbcdt.shape[-1] - 2 * dim - 2 * ctx.ngroups * dstate - nheads) // 2 + assert d_nonssm >= 0 + assert d_nonssm == 0 + recompute_output = outproj_weight is not None + + zx0, z, xBC, dt = torch.split( + zxbcdt, [2 * d_nonssm, dim, dim + 2 * ctx.ngroups * dstate, nheads], dim=-1 + ) + # Recompute x, B, C. + xBC_conv, _ = _causal_conv1d_state_passing_cp_fwd( + rearrange(ensure_stride(xBC), "b s d -> b d s"), + conv1d_weight, + conv1d_bias, + seq_idx, + state_passing_conv_initial_states, + None, + ctx.activation, + ctx.state_passing_cp_group, + ctx.state_passing_cp_virtual, + ) + xBC_conv = rearrange(xBC_conv, "b d s -> b s d") + x, B, C = torch.split(xBC_conv, [dim, ctx.ngroups * dstate, ctx.ngroups * dstate], dim=-1) + x = rearrange(x, "b l (h p) -> b l h p", h=nheads) + B = rearrange(B, "b l (g n) -> b l g n", g=ctx.ngroups) + C = rearrange(C, "b l (g n) -> b l g n", g=ctx.ngroups) + + dzxbcdt = torch.empty_like(zxbcdt) + dzx0, dz, dxBC_given, ddt_given = torch.split( + dzxbcdt, [2 * d_nonssm, dim, dim + 2 * ctx.ngroups * dstate, nheads], dim=-1 + ) + dxBC = torch.empty_like(xBC) + dx, dB, dC = torch.split(dxBC, [dim, ctx.ngroups * dstate, ctx.ngroups * dstate], dim=-1) + z = rearrange(z, "b l (h p) -> b l h p", h=nheads) + dx = rearrange(dx, "b l (h p) -> b l h p", h=nheads) + dB = rearrange(dB, "b l (g n) -> b l g n", g=ctx.ngroups) + dC = rearrange(dC, "b l (g n) -> b l g n", g=ctx.ngroups) + dout = rearrange(dout, "b s (h p) -> b s h p", p=headdim) + dz = rearrange(dz, "b l (h p) -> b l h p", h=nheads) + + dx, ddt, dA, dB, dC, dD, dz, ddt_bias, dinitial_states = ( + _mamba_chunk_scan_combined_state_passing_cp_bwd( + dout, + x, + dt, + A, + B, + C, + out, + ctx.chunk_size, + D=D, + z=z, + dt_bias=dt_bias, + initial_states=initial_states, + dfinal_states=dfinal_states, + seq_idx=seq_idx, + dt_softplus=True, + dt_limit=ctx.dt_limit, + dx=dx, + ddt=ddt_given, + dB=dB, + dC=dC, + dz=dz, + recompute_output=recompute_output, + state_passing_cp_group=ctx.state_passing_cp_group, + state_passing_initial_states=state_passing_initial_states, + state_passing_gathered_decays=state_passing_gathered_decays, + state_passing_cp_virtual=ctx.state_passing_cp_virtual, + ) + ) + + dxBC_given_update, dweight, dbias = _causal_conv1d_state_passing_cp_bwd( + rearrange(ensure_stride(xBC), "b s d -> b d s"), + conv1d_weight, + conv1d_bias, + rearrange(ensure_stride(dxBC), "b s d -> b d s"), + seq_idx, + state_passing_conv_initial_states, + None, + rearrange(ensure_stride(dxBC_given), "b s d -> b d s"), + False, + ctx.activation, + ctx.state_passing_cp_group, + ctx.state_passing_cp_virtual, + ) + dxBC_given_update = rearrange(dxBC_given_update, "b d s -> b s d") + if dxBC_given.stride() != dxBC_given_update.stride(): + dxBC_given.copy_(dxBC_given_update) + else: + dxBC_given = dxBC_given_update + + drmsnorm_weight = None + doutproj_weight = None + doutproj_bias = None + return ( + dzxbcdt, + dweight, + dbias, + ddt_bias, + dA, + dD, + None, + dinitial_states, + None, + None, + None, + None, + drmsnorm_weight, + None, + doutproj_weight, + doutproj_bias, + None, + None, + None, + None, + None, + ) + + +def mamba_split_conv1d_scan_combined_state_passing_cp( + zxbcdt, + conv1d_weight, + conv1d_bias, + dt_bias, + A, + D, + chunk_size, + initial_states=None, + seq_idx=None, + dt_limit=(0.0, float("inf")), + return_final_states=False, + activation="silu", + rmsnorm_weight=None, + rmsnorm_eps=1e-6, + outproj_weight=None, + outproj_bias=None, + headdim=None, + ngroups=1, + norm_before_gate=True, + state_passing_cp_group=None, + state_passing_cp_virtual=False, +): + """State-passing CP counterpart of ``mamba_split_conv1d_scan_combined``.""" + return MambaSplitConv1dScanCombinedStatePassingCPFn.apply( + zxbcdt, + conv1d_weight, + conv1d_bias, + dt_bias, + A, + D, + chunk_size, + initial_states, + seq_idx, + dt_limit, + return_final_states, + activation, + rmsnorm_weight, + rmsnorm_eps, + outproj_weight, + outproj_bias, + headdim, + ngroups, + norm_before_gate, + state_passing_cp_group, + state_passing_cp_virtual, + ) + + +class MambaStatePassingCPAdapter: + """Adapter between ``MambaMixer`` and the state-passing CP production path.""" + + def __init__(self, mixer): + self.mixer = mixer + + def forward(self, zxBCdt: torch.Tensor) -> torch.Tensor: + """Run the fused state-passing Conv+SSD path on a projected activation.""" + mixer = self.mixer + state_passing_cp_load_balancing = mixer.config.mamba_state_passing_cp_load_balancing + assert state_passing_cp_load_balancing in ("none", "permute_p2p", "permute_a2a", "virtual") + if state_passing_cp_load_balancing.startswith("permute_"): + state_passing_cp_permute_backend = state_passing_cp_load_balancing.removeprefix( + "permute_" + ) + zxBCdt = undo_state_passing_cp_load_balancing( + zxBCdt, mixer.cp.cp_group, backend=state_passing_cp_permute_backend + ) + zxBCdt = rearrange(zxBCdt, "l b d -> b l d").contiguous() + state_passing_cp_virtual = state_passing_cp_load_balancing == "virtual" + if state_passing_cp_virtual: + zxBCdt = _pack_state_passing_cp_virtual_segments(zxBCdt) + A = -torch.exp(mixer.A_log.float()) + D = ( + rearrange(mixer.D.float(), "(h p) -> h p", p=mixer.headdim) + if mixer.D_has_hdim + else mixer.D + ) + + assert ( + causal_conv1d_fwd_function is not None and causal_conv1d_bwd_function is not None + ), "Mamba state-passing CP requires causal-conv1d" + assert mixer.activation in ["silu", "swish"] + y = mamba_split_conv1d_scan_combined_state_passing_cp( + zxBCdt, + rearrange(mixer.conv1d_weight, "d 1 w -> d w"), + mixer.conv1d_bias, + mixer.dt_bias.float(), + A, + D, + mixer.chunk_size, + activation=mixer.activation, + headdim=None if mixer.D_has_hdim else mixer.headdim, + ngroups=mixer.ngroups_local_tp, + norm_before_gate=mixer.norm_before_gate, + state_passing_cp_group=mixer.cp.cp_group, + state_passing_cp_virtual=state_passing_cp_virtual, + ) + + if mixer.rmsnorm: + # Match MambaMixer's fused path: materialize the BF16 gated scan + # output before applying RMSNorm. + y = mixer.norm(y) + if state_passing_cp_virtual: + y = _unpack_state_passing_cp_virtual_segments(y) + y = rearrange(y, "b l d -> l b d").contiguous() + if state_passing_cp_load_balancing.startswith("permute_"): + y = redo_state_passing_cp_load_balancing( + y, mixer.cp.cp_group, backend=state_passing_cp_permute_backend + ) + return y diff --git a/tests/unit_tests/ssm/mamba_state_passing_cp_utils.py b/tests/unit_tests/ssm/mamba_state_passing_cp_utils.py new file mode 100644 index 00000000000..f9fec03bb8e --- /dev/null +++ b/tests/unit_tests/ssm/mamba_state_passing_cp_utils.py @@ -0,0 +1,136 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Shared helpers for Mamba state-passing context-parallel unit tests.""" + +from dataclasses import dataclass + +import torch +import torch.distributed as dist + +# Relative RMS error tolerance used for BF16 comparisons against a +# full-sequence reference. +DEFAULT_ERROR_RATIO = 3e-3 + +def balanced_cp_chunk_ids(rank: int, cp_size: int) -> tuple[int, int]: + """Return the front/back chunk indices Megatron's balanced CP assigns to ``rank``.""" + assert cp_size > 0 and 0 <= rank < cp_size + return rank, 2 * cp_size - 1 - rank + + +def select_balanced_cp_shard( + tensor: torch.Tensor, rank: int, cp_size: int, *, sequence_dim: int = 0 +) -> torch.Tensor: + """Select Megatron's front/back balanced CP shard without packing the batch axis.""" + assert tensor.shape[sequence_dim] % (2 * cp_size) == 0 + chunks = torch.chunk(tensor, 2 * cp_size, dim=sequence_dim) + front, back = balanced_cp_chunk_ids(rank, cp_size) + return torch.cat((chunks[front], chunks[back]), dim=sequence_dim).contiguous() + + +def select_contiguous_cp_shard( + tensor: torch.Tensor, rank: int, cp_size: int, *, sequence_dim: int = 0 +) -> torch.Tensor: + """Select the two causally adjacent chunks owned by a contiguous CP rank.""" + assert tensor.shape[sequence_dim] % (2 * cp_size) == 0 + chunks = torch.chunk(tensor, 2 * cp_size, dim=sequence_dim) + return torch.cat((chunks[2 * rank], chunks[2 * rank + 1]), dim=sequence_dim).contiguous() + + +def select_state_passing_cp_shard( + tensor: torch.Tensor, + rank: int, + cp_size: int, + *, + virtual: bool, + batch_dim: int = 0, + sequence_dim: int = 1, +) -> torch.Tensor: + """Select a contiguous shard, or pack balanced front/back chunks on the batch axis. + + The state-passing kernels consume either a contiguous causal shard or, in + ``virtual`` mode, an interleaved view where each balanced front/back chunk + becomes an independent virtual batch entry. + """ + assert cp_size > 0 and 0 <= rank < cp_size + assert tensor.shape[sequence_dim] % cp_size == 0 + if not virtual: + local_length = tensor.shape[sequence_dim] // cp_size + return tensor.narrow(sequence_dim, rank * local_length, local_length).contiguous() + + assert batch_dim != sequence_dim + moved = tensor.movedim((batch_dim, sequence_dim), (0, 1)) + assert moved.shape[1] % (2 * cp_size) == 0 + chunks = torch.chunk(moved, 2 * cp_size, dim=1) + front, back = balanced_cp_chunk_ids(rank, cp_size) + packed = torch.stack((chunks[front], chunks[back]), dim=1).flatten(0, 1).contiguous() + return packed.movedim((0, 1), (batch_dim, sequence_dim)) + + +def relative_rms_error(actual: torch.Tensor, expected: torch.Tensor) -> float: + """Return the RMS error of ``actual`` relative to the RMS magnitude of ``expected``. + + A relative criterion is used instead of elementwise ``allclose`` because the + state-passing path reassociates the SSD scan across CP boundaries, which + changes BF16 rounding without changing the mathematical result. + """ + actual_float = actual.float() + expected_float = expected.float() + assert torch.isfinite(actual_float).all(), "state-passing result contains non-finite values" + assert torch.isfinite(expected_float).all(), "reference contains non-finite values" + difference = actual_float - expected_float + if difference.abs().max().item() == 0.0: + return 0.0 + expected_rms = expected_float.square().mean().sqrt().item() + return difference.square().mean().sqrt().item() / (expected_rms + 1e-8) + + +def assert_all_close_rms( + checks: dict[str, tuple[torch.Tensor, torch.Tensor]], + group: dist.ProcessGroup, + ratio: float = DEFAULT_ERROR_RATIO, +) -> None: + """Compare every (actual, expected) pair and fail identically on all CP ranks. + + Every comparison is evaluated before anything is asserted, and the verdict is + all-reduced. Asserting per comparison would let one rank leave the test while + the others are still inside a collective, which deadlocks the whole run + instead of reporting a failure. + """ + errors = {name: relative_rms_error(*pair) for name, pair in checks.items()} + names = list(errors) + local = torch.tensor([errors[name] for name in names], device=torch.cuda.current_device()) + dist.all_reduce(local, op=dist.ReduceOp.MAX, group=group) + failures = [ + f"{name}: relative RMS error {error:.3e} exceeds {ratio:.3e}" + for name, error in zip(names, local.tolist()) + if not error < ratio + ] + assert not failures, "state-passing CP mismatch on at least one rank:\n" + "\n".join(failures) + + +@dataclass(frozen=True) +class MambaModelShape: + """TP-local Mamba mixer shape used to build test models.""" + + hidden_size: int = 2688 + nheads: int = 64 + head_dim: int = 64 + state_dim: int = 128 + ngroups: int = 8 + chunk_size: int = 128 + d_conv: int = 4 + + @property + def d_inner(self) -> int: + """Inner (gated) width of the mixer.""" + return self.nheads * self.head_dim + + @property + def conv_dim(self) -> int: + """Number of channels the depthwise causal convolution operates on.""" + return self.d_inner + 2 * self.ngroups * self.state_dim + + @property + def projected_width(self) -> int: + """Width of the ``in_proj`` output (z, x, B, C, dt).""" + return 2 * self.d_inner + 2 * self.ngroups * self.state_dim + self.nheads diff --git a/tests/unit_tests/ssm/ops/test_ssd_state_passing_cp.py b/tests/unit_tests/ssm/ops/test_ssd_state_passing_cp.py new file mode 100644 index 00000000000..2ca8c4bd132 --- /dev/null +++ b/tests/unit_tests/ssm/ops/test_ssd_state_passing_cp.py @@ -0,0 +1,305 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Kernel-level tests for the Mamba2 state-passing context-parallel path. + +Each test builds a full-sequence reference on every rank, shards it the way +Megatron's context parallelism would, and checks that the state-passing kernels +reproduce the reference shard for both the forward output and every gradient. + +The mixer dimensions are Nemotron-3 Nano's (64 heads of 64, 8 groups, state 128, +chunk 128); the batch size and sequence length are what the parametrization +varies, because those are what the CP sharding and the boundary exchange +actually depend on. +""" + +import pytest +import torch + +from tests.unit_tests.ssm.mamba_state_passing_cp_utils import ( + MambaModelShape, + assert_all_close_rms, + select_balanced_cp_shard, + select_contiguous_cp_shard, + select_state_passing_cp_shard, +) +from tests.unit_tests.test_utilities import Utils + +try: + from causal_conv1d import causal_conv1d_fn + from mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined + + from megatron.core.ssm.ops.ssd_state_passing_cp import ( + _causal_conv1d_state_passing_cp_bwd, + _causal_conv1d_state_passing_cp_fwd, + _mamba_chunk_scan_combined_state_passing_cp_bwd, + _mamba_chunk_scan_combined_state_passing_cp_fwd, + redo_state_passing_cp_load_balancing, + undo_state_passing_cp_load_balancing, + ) + + HAVE_STATE_PASSING_CP = True +except ImportError: + HAVE_STATE_PASSING_CP = False + +pytestmark = [ + pytest.mark.internal, + pytest.mark.skipif( + not HAVE_STATE_PASSING_CP, reason="mamba_ssm and causal_conv1d are required" + ), + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required"), +] + +SHAPE = MambaModelShape() + +# (batch, sequence_length). Every causal segment must align with the SSD chunk +# size, and ``virtual`` mode halves the segment length, so the sequence length +# must be a multiple of ``4 * cp_size * chunk_size``. +BATCH_AND_SEQUENCE_LENGTHS = ((1, 2048), (3, 4096)) + + +@pytest.fixture +def cp_context(request): + """Initialize a context-parallel group of the requested size and tear it down.""" + cp_size = request.param + if Utils.world_size % cp_size != 0: + pytest.skip(f"world size {Utils.world_size} is not a multiple of cp_size {cp_size}") + Utils.initialize_model_parallel(context_parallel_size=cp_size) + from megatron.core import parallel_state + + group = parallel_state.get_context_parallel_group() + yield group, group.rank(), cp_size, torch.device(torch.cuda.current_device()) + Utils.destroy_model_parallel() + + +@pytest.mark.parametrize("cp_context", (2, 4), indirect=True) +@pytest.mark.parametrize("backend", ("p2p", "a2a")) +def test_load_balancing_permutation_round_trip(cp_context, backend): + """The balanced-to-contiguous permutation and its backward must be exact.""" + group, rank, cp_size, device = cp_context + sequence_length, hidden_size = 2048, SHAPE.hidden_size + + total_chunks = 2 * cp_size + chunk_length = sequence_length // total_chunks + full = torch.arange(sequence_length * hidden_size, device=device, dtype=torch.float32).view( + sequence_length, 1, hidden_size + ) + grad_full = torch.empty_like(full) + for chunk_id in range(total_chunks): + grad_full[chunk_id * chunk_length : (chunk_id + 1) * chunk_length].fill_(chunk_id + 1) + + balanced = select_balanced_cp_shard(full, rank, cp_size).clone().requires_grad_(True) + contiguous = undo_state_passing_cp_load_balancing(balanced, group, backend=backend) + restored = redo_state_passing_cp_load_balancing(contiguous, group, backend=backend) + (contiguous * select_contiguous_cp_shard(grad_full, rank, cp_size)).sum().backward() + + torch.testing.assert_close( + contiguous, select_contiguous_cp_shard(full, rank, cp_size), rtol=0, atol=0 + ) + torch.testing.assert_close(restored, balanced.detach(), rtol=0, atol=0) + torch.testing.assert_close( + balanced.grad, select_balanced_cp_shard(grad_full, rank, cp_size), rtol=0, atol=0 + ) + + +@pytest.mark.parametrize("cp_context", (2, 4), indirect=True) +@pytest.mark.parametrize("virtual", (False, True)) +@pytest.mark.parametrize("batch,sequence_length", BATCH_AND_SEQUENCE_LENGTHS) +def test_causal_conv1d_matches_full_sequence(cp_context, virtual, batch, sequence_length): + """The halo-exchanging causal conv must match a full-sequence causal_conv1d.""" + group, rank, cp_size, device = cp_context + channels = SHAPE.conv_dim + dtype = torch.bfloat16 + + generator = torch.Generator(device=device).manual_seed(1234) + x = torch.randn( + batch, channels, sequence_length, device=device, dtype=dtype, generator=generator + ) + weight = torch.randn(channels, SHAPE.d_conv, device=device, dtype=dtype, generator=generator) + bias = torch.randn(channels, device=device, dtype=dtype, generator=generator) + # Every input must come from the explicit seeded generator: the per-rank RNG + # that ``randn_like`` would use differs across ranks, which would give each + # rank a different "full-sequence" reference. + grad_output = torch.randn( + batch, channels, sequence_length, device=device, dtype=dtype, generator=generator + ) + + x_reference = x.clone().requires_grad_(True) + weight_reference = weight.clone().requires_grad_(True) + bias_reference = bias.clone().requires_grad_(True) + y_reference = causal_conv1d_fn(x_reference, weight_reference, bias_reference, activation="silu") + (y_reference * grad_output).sum().backward() + + def local_view(tensor): + return select_state_passing_cp_shard( + tensor, rank, cp_size, virtual=virtual, batch_dim=0, sequence_dim=2 + ) + + # causal_conv1d requires a channel-last-in-memory layout. + x_local = local_view(x).transpose(1, 2).contiguous().transpose(1, 2) + y_local, conv_initial_states = _causal_conv1d_state_passing_cp_fwd( + x_local, + weight.clone(), + bias.clone(), + activation="silu", + state_passing_cp_group=group, + state_passing_cp_virtual=virtual, + ) + dx_local, dweight_local, dbias_local = _causal_conv1d_state_passing_cp_bwd( + x_local, + weight.clone(), + bias.clone(), + local_view(grad_output), + initial_states=conv_initial_states, + activation="silu", + state_passing_cp_group=group, + state_passing_cp_virtual=virtual, + ) + + # Weight and bias are replicated across CP, so their gradients are summed. + dweight = dweight_local.detach().float() + dbias = dbias_local.detach().float() + torch.distributed.all_reduce(dweight, group=group) + torch.distributed.all_reduce(dbias, group=group) + + assert_all_close_rms( + { + "conv output": (y_local, local_view(y_reference)), + "conv dx": (dx_local, local_view(x_reference.grad)), + "conv dweight": (dweight, weight_reference.grad), + "conv dbias": (dbias, bias_reference.grad), + }, + group, + ) + + +@pytest.mark.parametrize("cp_context", (2, 4), indirect=True) +@pytest.mark.parametrize("virtual", (False, True)) +@pytest.mark.parametrize("batch,sequence_length", BATCH_AND_SEQUENCE_LENGTHS) +def test_chunk_scan_combined_matches_full_sequence(cp_context, virtual, batch, sequence_length): + """The boundary-scanning SSD kernel must match a full-sequence chunk scan.""" + group, rank, cp_size, device = cp_context + nheads, headdim, ngroups = SHAPE.nheads, SHAPE.head_dim, SHAPE.ngroups + dstate, chunk_size = SHAPE.state_dim, SHAPE.chunk_size + segment_length = sequence_length // (cp_size * (2 if virtual else 1)) + assert ( + segment_length % chunk_size == 0 + ), f"each causal segment ({segment_length}) must align with the SSD chunk size" + dtype = torch.bfloat16 + kernel_kwargs = dict(dt_softplus=True, dt_limit=(0.0, float("inf"))) + + generator = torch.Generator(device=device).manual_seed(1234) + x = torch.randn( + batch, sequence_length, nheads, headdim, device=device, dtype=dtype, generator=generator + ) + dt = ( + torch.rand(batch, sequence_length, nheads, device=device, dtype=dtype, generator=generator) + * 0.5 + + 0.01 + ) + A_log = torch.log( + torch.rand(nheads, device=device, dtype=torch.float32, generator=generator) + 0.1 + ) + B = torch.randn( + batch, sequence_length, ngroups, dstate, device=device, dtype=dtype, generator=generator + ) + C = torch.randn( + batch, sequence_length, ngroups, dstate, device=device, dtype=dtype, generator=generator + ) + # Every input must come from the explicit seeded generator: the per-rank RNG + # that ``randn_like`` would use differs across ranks, which would give each + # rank a different "full-sequence" reference. + z = torch.randn( + batch, sequence_length, nheads, headdim, device=device, dtype=dtype, generator=generator + ) + D = torch.randn(nheads, device=device, dtype=torch.float32, generator=generator) + dt_bias = torch.randn(nheads, device=device, dtype=dtype, generator=generator) * 0.1 + grad_output = torch.randn( + batch, sequence_length, nheads, headdim, device=device, dtype=dtype, generator=generator + ) + + def local_view(tensor): + return select_state_passing_cp_shard( + tensor, rank, cp_size, virtual=virtual, batch_dim=0, sequence_dim=1 + ) + + reference_inputs = [tensor.clone().requires_grad_(True) for tensor in (x, dt, B, C, z)] + A_log_reference = A_log.clone().requires_grad_(True) + A_reference = -torch.exp(A_log_reference) + A_reference.retain_grad() + D_reference = D.clone().requires_grad_(True) + dt_bias_reference = dt_bias.clone().requires_grad_(True) + y_reference = mamba_chunk_scan_combined( + reference_inputs[0], + reference_inputs[1], + A_reference, + reference_inputs[2], + reference_inputs[3], + chunk_size, + D=D_reference, + z=reference_inputs[4], + dt_bias=dt_bias_reference, + **kernel_kwargs, + ) + (y_reference * grad_output).sum().backward() + + A = -torch.exp(A_log) + cp_inputs = [local_view(tensor).clone().requires_grad_(True) for tensor in (x, dt, B, C, z)] + y_cp, out_x, _, _, _, _, initial_states, gathered_decays = ( + _mamba_chunk_scan_combined_state_passing_cp_fwd( + cp_inputs[0], + cp_inputs[1], + A, + cp_inputs[2], + cp_inputs[3], + chunk_size, + D=D, + z=cp_inputs[4], + dt_bias=dt_bias, + state_passing_cp_group=group, + state_passing_cp_virtual=virtual, + **kernel_kwargs, + ) + ) + cp_grads = _mamba_chunk_scan_combined_state_passing_cp_bwd( + local_view(grad_output), + cp_inputs[0], + cp_inputs[1], + A, + cp_inputs[2], + cp_inputs[3], + out_x, + chunk_size, + D=D, + z=cp_inputs[4], + dt_bias=dt_bias, + state_passing_initial_states=initial_states, + state_passing_gathered_decays=gathered_decays, + state_passing_cp_group=group, + state_passing_cp_virtual=virtual, + **kernel_kwargs, + ) + dx, ddt, dA, dB, dC, dD, dz, ddt_bias = cp_grads[:8] + + # A, D, and dt_bias are replicated across CP, so their gradients are summed. + reduced = [] + for gradient in (dA, dD, ddt_bias): + value = gradient.detach().float() + torch.distributed.all_reduce(value, group=group) + reduced.append(value) + dA_reduced, dD_reduced, ddt_bias_reduced = reduced + + assert_all_close_rms( + { + "ssd output": (y_cp, local_view(y_reference)), + "ssd dx": (dx, local_view(reference_inputs[0].grad)), + "ssd ddt": (ddt, local_view(reference_inputs[1].grad)), + "ssd dB": (dB, local_view(reference_inputs[2].grad)), + "ssd dC": (dC, local_view(reference_inputs[3].grad)), + "ssd dz": (dz, local_view(reference_inputs[4].grad)), + "ssd dA": (dA_reduced, A_reference.grad), + "ssd dA_log": (dA_reduced * A, A_log_reference.grad), + "ssd dD": (dD_reduced, D_reference.grad), + "ssd ddt_bias": (ddt_bias_reduced, dt_bias_reference.grad), + }, + group, + ) From ab5032a30921769bf534fa4265b473b9771a68ae Mon Sep 17 00:00:00 2001 From: Sanghun Cho Date: Mon, 24 Aug 2026 09:14:53 +0000 Subject: [PATCH 2/5] Wire state-passing context parallelism into MambaMixer 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 --- megatron/core/ssm/mamba_mixer.py | 50 +++++- .../core/transformer/transformer_config.py | 23 +++ .../ssm/mamba_state_passing_cp_utils.py | 142 +++++++++++++++++- .../ssm/test_mamba_mixer_state_passing_cp.py | 138 +++++++++++++++++ 4 files changed, 350 insertions(+), 3 deletions(-) create mode 100644 tests/unit_tests/ssm/test_mamba_mixer_state_passing_cp.py diff --git a/megatron/core/ssm/mamba_mixer.py b/megatron/core/ssm/mamba_mixer.py index c8b3ef583fe..1d3587a9e29 100644 --- a/megatron/core/ssm/mamba_mixer.py +++ b/megatron/core/ssm/mamba_mixer.py @@ -47,6 +47,14 @@ from .mamba_context_parallel import MambaContextParallel +try: + from .ops.ssd_state_passing_cp import MambaStatePassingCPAdapter + + HAVE_MAMBA_STATE_PASSING_CP = True +except ImportError: + MambaStatePassingCPAdapter = None + HAVE_MAMBA_STATE_PASSING_CP = False + try: from causal_conv1d import causal_conv1d_fn from causal_conv1d.causal_conv1d_varlen import causal_conv1d_varlen_states @@ -221,6 +229,10 @@ def __init__( "mamba_training_ssm_states_dtype is set, but the installed mamba_ssm does " "not accept the `state_dtype` argument. Upgrade mamba_ssm or unset the option." ) + self.use_mamba_state_passing_cp = self.config.use_mamba_state_passing_cp + self.mamba_state_passing_cp_load_balancing = ( + self.config.mamba_state_passing_cp_load_balancing + ) self.d_state = self.config.mamba_state_dim self.headdim = self.config.mamba_head_dim self.ngroups = self.config.mamba_num_groups @@ -459,6 +471,9 @@ def __init__( D_has_hdim=self.D_has_hdim, ) self.tp_group = pg_collection.tp + self.state_passing_cp_adapter = ( + MambaStatePassingCPAdapter(self) if HAVE_MAMBA_STATE_PASSING_CP else None + ) def forward( self, @@ -499,9 +514,16 @@ def forward( zxBCdt, _ = self.in_proj(hidden_states) - zxBCdt = self.cp.pre_conv_ssm(zxBCdt, packed_seq_params) + use_state_passing_cp = self._use_mamba_state_passing_cp( + in_inference_mode, packed_seq_params + ) + if not use_state_passing_cp: + zxBCdt = self.cp.pre_conv_ssm(zxBCdt, packed_seq_params) - if in_inference_mode or not self.use_mem_eff_path: + if use_state_passing_cp: + assert ssm_state is None + y = self.state_passing_cp_adapter.forward(zxBCdt) + elif in_inference_mode or not self.use_mem_eff_path: # TODO(ksanthanam): Consider deprecating this path for training assert packed_seq_params is None, ( "Training with packed sequences is not supported " @@ -727,6 +749,30 @@ def _decode( return out, out_bias + def _use_mamba_state_passing_cp( + self, in_inference_mode: bool, packed_seq_params: Optional[PackedSeqParams] + ) -> bool: + if not self.use_mamba_state_passing_cp or self.cp.cp_size == 1: + return False + assert HAVE_MAMBA_STATE_PASSING_CP, "state-passing Mamba CP helper is unavailable" + assert self.use_mem_eff_path, "state-passing Mamba CP is only wired for training fast path" + assert not in_inference_mode, "state-passing Mamba CP is not supported for inference yet" + assert ( + packed_seq_params is None + ), "state-passing Mamba CP does not support packed sequences yet" + assert ( + not self.config.hybrid_context_parallel and not self.config.dynamic_context_parallel + ), "state-passing Mamba CP does not support hybrid/dynamic (variable-length) CP yet" + assert ( + self.config.mamba_training_ssm_states_dtype is None + ), "state-passing Mamba CP does not support --mamba-training-ssm-states-dtype yet" + assert self.mamba_state_passing_cp_load_balancing != "none", ( + "standard Megatron CP always produces front/back load-balanced sequence shards; " + "use 'permute_p2p', 'permute_a2a', or 'virtual'. Mode 'none' is only valid for " + "direct state-passing calls whose input is already contiguous." + ) + return True + def _ssm_training( self, zxBCdt: torch.Tensor, packed_seq_params: Optional[PackedSeqParams] = None ) -> torch.Tensor: diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index d6afac69c48..6c09fd57181 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1362,6 +1362,29 @@ class TransformerConfig(ModelParallelConfig): ) """Controls usage of the memory efficient path for Mamba layers.""" + use_mamba_state_passing_cp: bool = field( + default=False, metadata={"argparse_meta": {"arg_names": ["--use-mamba-state-passing-cp"]}} + ) + """Use experimental sequence-sharded state-passing context parallel for Mamba training.""" + + mamba_state_passing_cp_load_balancing: Literal[ + "none", "permute_p2p", "permute_a2a", "virtual" + ] = field( + default="permute_p2p", + metadata={ + "argparse_meta": { + "arg_names": ["--mamba-state-passing-cp-load-balancing"], + "choices": ["none", "permute_p2p", "permute_a2a", "virtual"], + } + }, + ) + """How state-passing Mamba CP handles the standard balanced CP sequence layout. + + ``none`` expects contiguous causal shards, ``permute_p2p`` and ``permute_a2a`` + exchange activation chunks with the selected communication method, and ``virtual`` + treats each balanced front/back chunk as an independent virtual rank. + """ + mlp_chunks_for_prefill: int = 1 """The number of chunks along the sequence dimension to use for MLP computation during prefill.""" diff --git a/tests/unit_tests/ssm/mamba_state_passing_cp_utils.py b/tests/unit_tests/ssm/mamba_state_passing_cp_utils.py index f9fec03bb8e..0e8357b25a9 100644 --- a/tests/unit_tests/ssm/mamba_state_passing_cp_utils.py +++ b/tests/unit_tests/ssm/mamba_state_passing_cp_utils.py @@ -2,15 +2,26 @@ """Shared helpers for Mamba state-passing context-parallel unit tests.""" -from dataclasses import dataclass +from dataclasses import dataclass, replace import torch import torch.distributed as dist +from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.ssm.mamba_block import MambaStackSubmodules +from megatron.core.ssm.mamba_layer import MambaLayer, MambaLayerSubmodules +from megatron.core.ssm.mamba_mixer import MambaMixer, MambaMixerSubmodules +from megatron.core.transformer import TransformerConfig + # Relative RMS error tolerance used for BF16 comparisons against a # full-sequence reference. DEFAULT_ERROR_RATIO = 3e-3 +# Load-balancing modes exercised by the state-passing CP tests. +STATE_PASSING_CP_MODES = ("permute_p2p", "permute_a2a", "virtual") + + def balanced_cp_chunk_ids(rank: int, cp_size: int) -> tuple[int, int]: """Return the front/back chunk indices Megatron's balanced CP assigns to ``rank``.""" assert cp_size > 0 and 0 <= rank < cp_size @@ -108,6 +119,33 @@ def assert_all_close_rms( assert not failures, "state-passing CP mismatch on at least one rank:\n" + "\n".join(failures) +def broadcast_module_parameters(module: torch.nn.Module, group: dist.ProcessGroup) -> None: + """Make every CP rank hold the same weights. + + The CP path combines activations produced from every rank's parameters, so a + per-rank difference in initialization would invalidate any comparison against + a rank-local full-sequence reference. + """ + source = dist.get_global_rank(group, 0) + for tensor in list(module.parameters()) + list(module.buffers()): + dist.broadcast(tensor.data, src=source, group=group) + + +def collect_parameter_grads( + module: torch.nn.Module, reduce_group: dist.ProcessGroup | None = None +) -> dict[str, torch.Tensor]: + """Collect parameter grads as FP32, optionally summing replicated grads across CP.""" + gradients = {} + for name, parameter in module.named_parameters(): + if parameter.grad is None: + continue + gradient = parameter.grad.detach().float().clone() + if reduce_group is not None: + dist.all_reduce(gradient, group=reduce_group) + gradients[name] = gradient + return gradients + + @dataclass(frozen=True) class MambaModelShape: """TP-local Mamba mixer shape used to build test models.""" @@ -134,3 +172,105 @@ def conv_dim(self) -> int: def projected_width(self) -> int: """Width of the ``in_proj`` output (z, x, B, C, dt).""" return 2 * self.d_inner + 2 * self.ngroups * self.state_dim + self.nheads + + +def build_mamba_config( + shape: MambaModelShape, + cp_size: int, + *, + load_balancing: str, + use_state_passing_cp: bool = True, + cuda_graph_impl: str = "none", + cuda_graph_warmup_steps: int = 2, +) -> TransformerConfig: + """Build a single-layer BF16 Mamba config for state-passing CP tests.""" + return TransformerConfig( + hidden_size=shape.hidden_size, + num_layers=1, + num_attention_heads=1, + bf16=True, + params_dtype=torch.bfloat16, + pipeline_dtype=torch.bfloat16, + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + context_parallel_size=cp_size, + mamba_num_heads=shape.nheads, + mamba_head_dim=shape.head_dim, + mamba_state_dim=shape.state_dim, + mamba_num_groups=shape.ngroups, + use_cpu_initialization=True, + use_mamba_mem_eff_path=True, + use_mamba_state_passing_cp=use_state_passing_cp, + mamba_state_passing_cp_load_balancing=load_balancing, + cuda_graph_impl=cuda_graph_impl, + cuda_graph_scope=[], + cuda_graph_warmup_steps=cuda_graph_warmup_steps, + # Both CUDA Graph backends require the TE RNG tracker: the default + # tracker refuses to be captured. + use_te_rng_tracker=cuda_graph_impl != "none", + hidden_dropout=0.0, + ) + + +def _mamba_layer_spec(): + assert isinstance(mamba_stack_spec.submodules, MambaStackSubmodules) + spec = mamba_stack_spec.submodules.mamba_layer + assert isinstance(spec.submodules, MambaLayerSubmodules) + return spec + + +def build_mamba_mixer( + shape: MambaModelShape, cp_group: dist.ProcessGroup, tp_group: dist.ProcessGroup, **kwargs +) -> MambaMixer: + """Build a ``MambaMixer`` bound to ``cp_group`` and move it to the current device.""" + config = build_mamba_config(shape, cp_group.size(), **kwargs) + mixer_spec = _mamba_layer_spec().submodules.mixer + assert isinstance(mixer_spec.submodules, MambaMixerSubmodules) + mixer = MambaMixer( + config, + mixer_spec.submodules, + shape.hidden_size, + d_conv=shape.d_conv, + chunk_size=shape.chunk_size, + layer_number=1, + pg_collection=ProcessGroupCollection(tp=tp_group, cp=cp_group), + ).to(torch.cuda.current_device()) + mixer.train() + assert mixer.conv1d_weight.shape[0] == shape.conv_dim + assert mixer.in_proj.weight.shape[0] == shape.projected_width + return mixer + + +def build_mamba_layer( + shape: MambaModelShape, cp_group: dist.ProcessGroup, tp_group: dist.ProcessGroup, **kwargs +) -> tuple[MambaLayer, TransformerConfig]: + """Build a ``MambaLayer`` bound to ``cp_group`` and move it to the current device.""" + config = build_mamba_config(shape, cp_group.size(), **kwargs) + spec = _mamba_layer_spec() + mixer_spec = replace( + spec.submodules.mixer, + params={ + **spec.submodules.mixer.params, + "d_conv": shape.d_conv, + "chunk_size": shape.chunk_size, + }, + ) + layer = MambaLayer( + config, + replace(spec.submodules, mixer=mixer_spec), + layer_number=1, + pg_collection=ProcessGroupCollection(tp=tp_group, cp=cp_group), + ).to(torch.cuda.current_device()) + layer.train() + return layer, config + + +def set_state_passing_cp_mode(mixer: MambaMixer, mode: str | None) -> None: + """Switch a built mixer between the A2A CP path (``None``) and a state-passing mode.""" + if mode is None: + mixer.use_mamba_state_passing_cp = False + return + assert mode in STATE_PASSING_CP_MODES + mixer.use_mamba_state_passing_cp = True + mixer.mamba_state_passing_cp_load_balancing = mode + mixer.config.mamba_state_passing_cp_load_balancing = mode diff --git a/tests/unit_tests/ssm/test_mamba_mixer_state_passing_cp.py b/tests/unit_tests/ssm/test_mamba_mixer_state_passing_cp.py new file mode 100644 index 00000000000..0f8d44cd47b --- /dev/null +++ b/tests/unit_tests/ssm/test_mamba_mixer_state_passing_cp.py @@ -0,0 +1,138 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""End-to-end tests for ``MambaMixer`` with state-passing context parallelism. + +Each state-passing load-balancing mode is compared against a single-rank +full-sequence ``MambaMixer`` and against the existing all-to-all CP path, for +the output, the input gradient, and every parameter gradient. +""" + +import pytest +import torch + +from megatron.core import parallel_state +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from tests.unit_tests.ssm.mamba_state_passing_cp_utils import ( + STATE_PASSING_CP_MODES, + MambaModelShape, + assert_all_close_rms, + broadcast_module_parameters, + build_mamba_mixer, + collect_parameter_grads, + select_balanced_cp_shard, + set_state_passing_cp_mode, +) +from tests.unit_tests.test_utilities import Utils + +try: + import causal_conv1d # noqa: F401 + import mamba_ssm # noqa: F401 + + from megatron.core.ssm.mamba_mixer import HAVE_MAMBA_STATE_PASSING_CP +except ImportError: + HAVE_MAMBA_STATE_PASSING_CP = False + +pytestmark = [ + pytest.mark.internal, + pytest.mark.skipif( + not HAVE_MAMBA_STATE_PASSING_CP, reason="mamba_ssm and causal_conv1d are required" + ), + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required"), +] + + +def _run_mixer(mixer, mode, hidden_states, grad_output, parameter_reduce_group=None): + """Run one forward/backward of ``mixer`` in ``mode`` and collect all gradients.""" + mixer.zero_grad(set_to_none=True) + local_input = hidden_states.detach().clone().requires_grad_(True) + set_state_passing_cp_mode(mixer, mode) + output, output_bias = mixer(local_input) + assert output_bias is None, "this test expects a bias-free output projection" + torch.autograd.backward(output, grad_output) + return ( + output.detach(), + local_input.grad.detach(), + collect_parameter_grads(mixer, parameter_reduce_group), + ) + + +@pytest.mark.parametrize("cp_size", (2, 4)) +@pytest.mark.parametrize("batch_size,sequence_length", ((1, 2048), (3, 4096))) +def test_mixer_state_passing_cp_matches_full_sequence_and_a2a(cp_size, batch_size, sequence_length): + """All state-passing modes must match the full-sequence and A2A CP references.""" + if Utils.world_size % cp_size != 0: + pytest.skip(f"world size {Utils.world_size} is not a multiple of cp_size {cp_size}") + Utils.initialize_model_parallel(context_parallel_size=cp_size) + try: + torch.manual_seed(1234) + model_parallel_cuda_manual_seed(1234, force_reset_rng=True) + cp_group = parallel_state.get_context_parallel_group() + tp_group = parallel_state.get_tensor_model_parallel_group() + assert tp_group.size() == 1, "the full-sequence reference needs a size-1 CP group" + rank = cp_group.rank() + device = torch.device(torch.cuda.current_device()) + + shape = MambaModelShape() + # ``virtual`` mode halves the segment length, and every causal segment + # must align with the SSD chunk size. + assert sequence_length % (4 * cp_size * shape.chunk_size) == 0 + + mixer = build_mamba_mixer(shape, cp_group, tp_group, load_balancing="permute_p2p") + broadcast_module_parameters(mixer, cp_group) + # A size-1 CP group makes this mixer process the whole sequence. + reference_mixer = build_mamba_mixer( + shape, tp_group, tp_group, load_balancing="permute_p2p", use_state_passing_cp=False + ) + reference_mixer.load_state_dict(mixer.state_dict()) + + generator = torch.Generator(device=device).manual_seed(5678) + global_input = torch.randn( + sequence_length, + batch_size, + shape.hidden_size, + device=device, + dtype=torch.bfloat16, + generator=generator, + ) + global_grad_output = torch.randn( + global_input.shape, device=device, dtype=torch.bfloat16, generator=generator + ) + local_input = select_balanced_cp_shard(global_input, rank, cp_size) + local_grad_output = select_balanced_cp_shard(global_grad_output, rank, cp_size) + + reference_output, reference_input_grad, reference_parameter_grads = _run_mixer( + reference_mixer, None, global_input, global_grad_output + ) + expected_output = select_balanced_cp_shard(reference_output, rank, cp_size) + expected_input_grad = select_balanced_cp_shard(reference_input_grad, rank, cp_size) + + # The existing all-to-all CP path is the second reference: it is the + # numerical behaviour production already ships. + a2a_output, a2a_input_grad, a2a_parameter_grads = _run_mixer( + mixer, None, local_input, local_grad_output, cp_group + ) + checks = { + "a2a output vs full sequence": (a2a_output, expected_output), + "a2a input grad vs full sequence": (a2a_input_grad, expected_input_grad), + } + for mode in STATE_PASSING_CP_MODES: + output, input_grad, parameter_grads = _run_mixer( + mixer, mode, local_input, local_grad_output, cp_group + ) + checks[f"{mode} output vs full sequence"] = (output, expected_output) + checks[f"{mode} input grad vs full sequence"] = (input_grad, expected_input_grad) + checks[f"{mode} output vs a2a"] = (output, a2a_output) + checks[f"{mode} input grad vs a2a"] = (input_grad, a2a_input_grad) + assert parameter_grads.keys() == reference_parameter_grads.keys() + for name, gradient in parameter_grads.items(): + checks[f"{mode} {name} vs full sequence"] = ( + gradient, + reference_parameter_grads[name], + ) + checks[f"{mode} {name} vs a2a"] = (gradient, a2a_parameter_grads[name]) + + # Asserted in one place so that a mismatch fails every rank identically + # instead of leaving some ranks inside a collective. + assert_all_close_rms(checks, cp_group) + finally: + Utils.destroy_model_parallel() From fb1006ec87e50be5b5b9b8d69a886ec798fbdb50 Mon Sep 17 00:00:00 2001 From: Sanghun Cho Date: Mon, 24 Aug 2026 09:15:10 +0000 Subject: [PATCH 3/5] Test state-passing Mamba CP under CUDA Graphs 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 --- .../test_mamba_state_passing_cp_cuda_graph.py | 267 ++++++++++++++++++ 1 file changed, 267 insertions(+) create mode 100644 tests/unit_tests/ssm/test_mamba_state_passing_cp_cuda_graph.py diff --git a/tests/unit_tests/ssm/test_mamba_state_passing_cp_cuda_graph.py b/tests/unit_tests/ssm/test_mamba_state_passing_cp_cuda_graph.py new file mode 100644 index 00000000000..f8d89de28bb --- /dev/null +++ b/tests/unit_tests/ssm/test_mamba_state_passing_cp_cuda_graph.py @@ -0,0 +1,267 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""CUDA Graph tests for Mamba state-passing context parallelism. + +Both graph backends are covered: Megatron's local CUDA Graph lifecycle +(``cuda_graph_impl="local"``) and the Transformer Engine helper +(``cuda_graph_impl="transformer_engine"``). Each test records a graph, replays +it, checks bitwise parity with eager execution on the first replay, and then +changes the input to confirm that the static input buffer is actually updated. +""" + +import gc + +import pytest +import torch +import torch.distributed as dist +import torch.nn as nn + +from megatron.core import parallel_state +from megatron.core.num_microbatches_calculator import ( + init_num_microbatches_calculator, + unset_num_microbatches_calculator, +) +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.cuda_graphs import ( + TECudaGraphHelper, + create_cudagraphs, + delete_cuda_graphs, +) +from tests.unit_tests.ssm.mamba_state_passing_cp_utils import ( + STATE_PASSING_CP_MODES, + MambaModelShape, + build_mamba_layer, + select_balanced_cp_shard, +) +from tests.unit_tests.test_utilities import Utils + +try: + import causal_conv1d # noqa: F401 + import mamba_ssm # noqa: F401 + + from megatron.core.ssm.mamba_mixer import HAVE_MAMBA_STATE_PASSING_CP +except ImportError: + HAVE_MAMBA_STATE_PASSING_CP = False + +pytestmark = [ + pytest.mark.internal, + pytest.mark.skipif( + not HAVE_MAMBA_STATE_PASSING_CP, reason="mamba_ssm and causal_conv1d are required" + ), + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required"), +] + +CP_SIZE = 2 +INPUT_GRAD_RTOL = 1e-3 +INPUT_GRAD_ATOL = 3e-6 + + +class _DecoderShell(nn.Module): + """Minimal decoder interface used by ``TECudaGraphHelper`` discovery.""" + + def __init__(self, layer): + super().__init__() + self.layers = nn.ModuleList([layer]) + + +class _ModelChunkShell(nn.Module): + """Minimal training model chunk required by ``TECudaGraphHelper``.""" + + def __init__(self, layer, config): + super().__init__() + self.config = config + self.decoder = _DecoderShell(layer) + + def zero_grad_buffer(self): + """Match the training loop's interface for clearing gradient buffers.""" + self.zero_grad(set_to_none=True) + + +def _unwrap_output(output): + """Normalize graph replay's optional one-tensor tuple to a tensor.""" + if isinstance(output, tuple): + assert len(output) == 1 + return output[0] + return output + + +def _assert_module_grads_finite(module): + for name, parameter in module.named_parameters(): + assert ( + parameter.grad is None or torch.isfinite(parameter.grad).all() + ), f"non-finite gradient for {name}" + + +def _setup(mode, cuda_graph_impl, warmup_steps): + """Initialize CP, build a graphed MambaLayer, and return it with its local input.""" + Utils.initialize_model_parallel(context_parallel_size=CP_SIZE) + torch.manual_seed(1234) + model_parallel_cuda_manual_seed(1234, te_rng_tracker=True, force_reset_rng=True) + cp_group = parallel_state.get_context_parallel_group() + tp_group = parallel_state.get_tensor_model_parallel_group() + device = torch.device(torch.cuda.current_device()) + + shape = MambaModelShape() + sequence_length = shape.chunk_size * 2 * CP_SIZE + batch_size = 1 + layer, config = build_mamba_layer( + shape, + cp_group, + tp_group, + load_balancing=mode, + cuda_graph_impl=cuda_graph_impl, + cuda_graph_warmup_steps=warmup_steps, + ) + generator = torch.Generator(device=device).manual_seed(5678) + global_input = torch.randn( + sequence_length, + batch_size, + shape.hidden_size, + device=device, + dtype=torch.bfloat16, + generator=generator, + ) + local_input = select_balanced_cp_shard(global_input, cp_group.rank(), CP_SIZE) + + # Cudagraph backward capture accumulates wgrads into ``main_grad``, which + # DDP normally provides; this test drives a bare module, so create them here + # (the same setup tests/unit_tests/transformer/test_cuda_graphs.py uses). + for parameter in layer.parameters(): + parameter.main_grad = torch.zeros_like(parameter) + + return layer, config, cp_group, local_input, sequence_length, batch_size + + +def _replay_and_check(layer, local_input, eager_output, eager_input_grad): + """Replay the graph twice, checking eager parity and static-buffer refresh.""" + replay_outputs = [] + for replay_index in range(2): + layer.zero_grad(set_to_none=True) + replay_input = (local_input * (1.0 + 0.01 * replay_index)).detach().requires_grad_(True) + replay_output = _unwrap_output(layer(hidden_states=replay_input, attention_mask=None)) + replay_output.float().square().mean().backward() + torch.cuda.synchronize() + + assert torch.isfinite(replay_output).all() + assert torch.isfinite(replay_input.grad).all() + _assert_module_grads_finite(layer) + if replay_index == 0: + torch.testing.assert_close(replay_output, eager_output, rtol=0, atol=0) + torch.testing.assert_close( + replay_input.grad, eager_input_grad, rtol=INPUT_GRAD_RTOL, atol=INPUT_GRAD_ATOL + ) + replay_outputs.append(replay_output.detach().clone()) + + assert not torch.equal( + replay_outputs[0], replay_outputs[1] + ), "CUDA Graph replay did not consume the changed input" + + +@pytest.mark.parametrize("mode", STATE_PASSING_CP_MODES) +def test_local_cuda_graph(mode): + """Megatron's local CUDA Graph lifecycle must capture and replay the CP path.""" + if Utils.world_size % CP_SIZE != 0: + pytest.skip(f"world size {Utils.world_size} is not a multiple of {CP_SIZE}") + layer, _, cp_group, local_input, _, _ = _setup(mode, "local", warmup_steps=2) + runner = None + try: + # A normal Megatron call records the runners and establishes eager references. + recorded_input = local_input.detach().clone().requires_grad_(True) + recorded_output = layer(hidden_states=recorded_input, attention_mask=None) + recorded_output.float().square().mean().backward() + eager_output = recorded_output.detach().clone() + eager_input_grad = recorded_input.grad.detach().clone() + torch.cuda.synchronize() + dist.barrier(group=cp_group) + + assert len(layer.cudagraph_manager.cudagraph_runners) == 1 + runner = layer.cudagraph_manager.cudagraph_runners[0] + assert runner.fwd_graph_recorded and runner.bwd_graph_recorded + + create_cudagraphs() + torch.cuda.synchronize() + dist.barrier(group=cp_group) + assert runner.cudagraph_created + assert runner.fwd_graph is not None and runner.bwd_graph is not None + + _replay_and_check(layer, local_input, eager_output, eager_input_grad) + dist.barrier(group=cp_group) + finally: + # Release the graph execs and the global cudagraph record before + # destroying the NCCL process groups. Leaving the global record in place + # makes the next test see `cudagraph_created` with no matching runners. + torch.cuda.synchronize() + delete_cuda_graphs() + del layer, runner + gc.collect() + torch.cuda.empty_cache() + Utils.destroy_model_parallel() + + +@pytest.mark.parametrize("mode", STATE_PASSING_CP_MODES) +def test_te_cuda_graph(mode): + """The Transformer Engine graph helper must capture and replay the CP path.""" + if Utils.world_size % CP_SIZE != 0: + pytest.skip(f"world size {Utils.world_size} is not a multiple of {CP_SIZE}") + warmup_steps = 2 + layer, config, cp_group, local_input, sequence_length, batch_size = _setup( + mode, "transformer_engine", warmup_steps + ) + helper = None + try: + init_num_microbatches_calculator( + rank=torch.distributed.get_rank(), + rampup_batch_size=None, + global_batch_size=batch_size, + micro_batch_size=batch_size, + data_parallel_size=1, + decrease_batch_size_if_needed=False, + ) + model_chunk = _ModelChunkShell(layer, config).to(torch.cuda.current_device()) + + eager_output = None + eager_input_grad = None + for _ in range(warmup_steps): + layer.zero_grad(set_to_none=True) + eager_input = local_input.detach().clone().requires_grad_(True) + eager_result = layer(hidden_states=eager_input, attention_mask=None) + eager_result.float().square().mean().backward() + torch.cuda.synchronize() + eager_output = eager_result.detach().clone() + eager_input_grad = eager_input.grad.detach().clone() + + # Training releases the previous iteration's autograd graph before TE + # capture. Keeping these tensors alive makes an old AccumulateGrad node + # retain its original stream, which can invalidate capture. + del eager_result, eager_input + layer.zero_grad(set_to_none=True) + gc.collect() + torch.cuda.synchronize() + dist.barrier(group=cp_group) + + helper = TECudaGraphHelper( + model=[model_chunk], + config=config, + seq_length=sequence_length, + micro_batch_size=batch_size, + optimizers=[], + ) + assert helper.flattened_callables == [layer] + helper.create_cudagraphs() + torch.cuda.synchronize() + dist.barrier(group=cp_group) + assert helper.capture_finished() and helper.graphs_created() + assert len(layer.cuda_graphs) == 1 + + _replay_and_check(layer, local_input, eager_output, eager_input_grad) + dist.barrier(group=cp_group) + finally: + if helper is not None: + helper.delete_cuda_graphs() + delete_cuda_graphs() + del layer + torch.cuda.synchronize() + gc.collect() + torch.cuda.empty_cache() + unset_num_microbatches_calculator() + Utils.destroy_model_parallel() From ef6d12c8c1da3e22dfe6a2af450bb46173539844 Mon Sep 17 00:00:00 2001 From: Sanghun Cho Date: Mon, 24 Aug 2026 09:15:10 +0000 Subject: [PATCH 4/5] Document Mamba2 state-passing context parallelism 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 --- docs/index.md | 1 + docs/user-guide/features/index.md | 1 + .../mamba_state_passing_context_parallel.md | 279 ++++++++++++++++++ 3 files changed, 281 insertions(+) create mode 100644 docs/user-guide/features/mamba_state_passing_context_parallel.md diff --git a/docs/index.md b/docs/index.md index 11337315588..133317dc449 100644 --- a/docs/index.md +++ b/docs/index.md @@ -67,6 +67,7 @@ models/index user-guide/features/moe user-guide/features/context_parallel +user-guide/features/mamba_state_passing_context_parallel user-guide/features/megatron_fsdp user-guide/features/dist_optimizer user-guide/features/optimizer_cpu_offload diff --git a/docs/user-guide/features/index.md b/docs/user-guide/features/index.md index cb2e895afdc..ce0fb95938b 100644 --- a/docs/user-guide/features/index.md +++ b/docs/user-guide/features/index.md @@ -18,6 +18,7 @@ cuda_graph fine_grained_activation_offloading moe context_parallel +mamba_state_passing_context_parallel megatron_fsdp dist_optimizer optimizer_cpu_offload diff --git a/docs/user-guide/features/mamba_state_passing_context_parallel.md b/docs/user-guide/features/mamba_state_passing_context_parallel.md new file mode 100644 index 00000000000..f57622ea049 --- /dev/null +++ b/docs/user-guide/features/mamba_state_passing_context_parallel.md @@ -0,0 +1,279 @@ + + +# Mamba2 State-Passing Context Parallelism + +Sequence-sharded, state-passing context parallelism (CP) for Mamba2 training. +The production implementation lives in +[`megatron/core/ssm/ops/ssd_state_passing_cp.py`](../../../megatron/core/ssm/ops/ssd_state_passing_cp.py) +and is opt-in; the existing all-to-all Mamba CP path remains the default. + +## Overview + +The existing Mamba CP path redistributes activations with an all-to-all (A2A) so +that every rank holds the full sequence for a subset of heads and groups. +State-passing CP instead keeps the sequence shard local and communicates only +what the causal boundary requires: the convolution halo and the SSD state +summary. + +| Path | Sequence handling | Dominant communication | +|---|---|---| +| No CP | one rank processes the whole sequence | none | +| Mamba A2A CP | full sequence, head/group shard | activation A2A proportional to sequence length | +| State-passing CP | local sequence shard is kept | conv halo and SSD boundary all-gather | + +Approximate rank-local payloads: + +```text +A2A: + O(local_sequence * hidden) + +State-passing boundary: + Conv: O(batch * channels * (d_conv - 1)) + SSD : O(batch * heads * headdim * dstate) + O(batch * heads) for decay +``` + +The boundary collective is an all-gather. Its size does not scale with sequence +length, but the received volume grows with the CP size. + +## Configuration + +- `--use-mamba-state-passing-cp` enables the path. +- `--mamba-state-passing-cp-load-balancing {none,permute_p2p,permute_a2a,virtual}` + selects how the standard balanced CP layout is handled. + +`none` expects an already-contiguous causal shard and is only valid for direct +calls into the state-passing kernels. Standard Megatron CP hands the mixer a +front/back balanced layout, so `MambaMixer` requires `permute_p2p`, +`permute_a2a`, or `virtual`. + +`MambaMixer` dispatch: + +```text +hidden_states + -> in_proj -> zxBCdt + -> MambaMixer._use_mamba_state_passing_cp() + | + +-- cp_size == 1 + | -> standard fused Mamba + | + +-- state-passing OFF + | -> MambaContextParallel.pre_conv_ssm() + | -> fused Mamba + | -> MambaContextParallel.post_conv_ssm() + | + `-- state-passing ON + -> MambaStatePassingCPAdapter.forward() + -> out_proj +``` + +## Load-balancing modes + +Megatron's balanced CP assigns one front chunk and one back chunk to each rank: + +```text +CP=3 balanced ownership: + rank 0: [1, 6] + rank 1: [2, 5] + rank 2: [3, 4] + +causal order: + [1, 2, 3, 4, 5, 6] +``` + +### Permute + +`permute_p2p` and `permute_a2a` physically exchange the balanced activations for +a contiguous causal shard: + +```text +balanced: + rank 0 [1,6], rank 1 [2,5], rank 2 [3,4] + +contiguous: + rank 0 [1,2], rank 1 [3,4], rank 2 [5,6] +``` + +`undo_state_passing_cp_load_balancing()` runs on the forward input and +`redo_state_passing_cp_load_balancing()` on the output. A custom autograd +function restores gradient ownership with the inverse permutation in the +backward pass. + +- `permute_p2p` sends and receives the required remote chunks point-to-point. +- `permute_a2a` exchanges them in one unequal-split `all_to_all_single`. + +Both backends produce the same layout and the same mathematics; only the +communication implementation differs. + +### Relation to `context_parallel_layout` + +Megatron Core already models the balanced/contiguous distinction as +`CpPartitionMode = Literal["zigzag", "contiguous"]` in +`megatron.core.context_parallel_layout`, and Gated DeltaNet uses it at its module +entry point through `convert_module_input_tensors_cp_partition_mode()` when it +runs its chunkwise (sequence-sharded) CP path. `permute_p2p` and `permute_a2a` +perform the same conversion, so the two overlap. + +They are implemented separately here for two reasons: the permutation is driven +from inside the fused Conv+SSD autograd function rather than at the module entry +point, and `permute_p2p` adds a point-to-point backend that the shared helper +does not provide. Consolidating onto the shared helper is intended follow-up +work. + +### Virtual + +`virtual` moves no activations. Each physical rank's front and back chunks are +interpreted as two independent virtual causal segments, and the input is viewed +as an interleaved virtual batch: + +```text +[local_L, batch, ...] -> [2 * batch, local_L / 2, ...] +``` + +Boundary routing then selects the predecessor and successor virtual segment out +of the rank-major gather result. + +## Production flow + +The path fuses the convolution and the SSD scan into a single custom autograd +function, `MambaSplitConv1dScanCombinedStatePassingCPFn`. Its public API keeps +the upstream `MambaSplitConv1dScanCombinedFn` argument order and appends +`state_passing_cp_group` and `state_passing_cp_virtual`. + +### Forward + +```text +zxBCdt [local_L, batch, packed] + -> load-balancing handling + -> virtual: view as [2 * batch, local_L / 2, packed] + -> split z / xBC / dt + +Conv + -> all-gather local tail + -> select predecessor tail + -> causal_conv1d(initial_states=predecessor_tail) + -> split x / B / C + +SSD + -> dt cumsum and local chunk contribution + -> packed summary [S_ext, a_block] + -> async all-gather(summary) || local CB computation + -> exclusive causal boundary scan + -> S_in + -> local state passing and chunk scan + +Output + -> virtual unpack or permutation redo + -> [local_L, batch, d_inner] +``` + +The SSD transform of one causal segment is affine: + +```text +S_out = a_block * S_in + S_ext +``` + +- `S_ext` is the final state the local segment produces from a zero initial state. +- `a_block` is the per-head decay across the whole segment. +- `S_in` is the exclusive prefix composition of the preceding segment summaries. + +`_state_passing_summary_fwd_kernel()` writes `S_ext` and the decay straight into +an FP32 packed buffer. `_state_passing_boundary_scan_kernel()` handles both the +forward and the backward causal order through the compile-time `VIRTUAL_CP` and +`REVERSE` flags. + +### Convolution boundary + +The causal convolution needs the last `d_conv - 1` tokens of the preceding +segment: + +```text +forward: + local tail -> all-gather -> predecessor tail + -> causal_conv1d_fwd_function(initial_states=...) + +backward: + causal_conv1d_bwd_function(return_dinitial_states=True) + -> all-gather d(initial_state) + -> select successor gradient + -> in-place add into the local dx tail +``` + +### Backward + +```text +re-split xBC from zxbcdt + -> recompute the causal conv forward + -> recompute dt cumsum, CB, chunk contribution, recurrence states + -> local chunk dstates + -> local dS_in summary + -> async all-gather(dS_in) || local dC / dCB / ddA computation + -> exclusive reverse boundary scan + -> dS_ext from the successor + -> one _state_passing_bwd(dfinal_states=dS_ext) + -> dx / dB / dC / ddt / dA / dD / dz + -> causal conv backward + -> apply the reverse conv boundary gradient to the dx tail +``` + +The post-conv gradient the SSD backward produces is written into a preallocated +`dzxBCdt` view and consumed directly by the convolution backward; there is no +separate autograd accumulation to re-join the split gradients. + +### Saved tensors and dtypes + +The forward context keeps the upstream tensor order and appends +`state_passing_conv_initial_states`, `state_passing_initial_states`, and +`state_passing_gathered_decays`. Quantities that scale with sequence length +(`xBC_conv`, the split `x`/`B`/`C`, `dt_proc`, the chunk states, and `CB`) are +recomputed in the backward pass instead of being saved. + +| Value | dtype | +|---|---| +| `S_ext`, `a_block` | FP32 | +| `S_in`, recurrence states | FP32 | +| `dS_in`, `dS_ext` | FP32 | +| scan/BMM boundary tensors | upstream input dtype | + +## Supported scope + +Supported: + +- the training memory-efficient path +- fixed-length, non-packed sequences +- the standard balanced CP input layout +- `permute_p2p`, `permute_a2a`, and `virtual` +- BF16 activations with FP32 boundary state + +Not supported yet: + +- inference +- `seq_idx` and packed sequences +- hybrid, dynamic, or otherwise variable-length CP +- an externally supplied SSD `initial_states` +- `--mamba-training-ssm-states-dtype` +- fused RMSNorm and output projection inside the custom function + +The mixer adapter runs RMSNorm and the output projection outside the custom +function, so both remain available on the full `MambaMixer` training path. + +Shape constraints: each virtual segment length must be a multiple of the SSD +`chunk_size`, and each convolution segment must hold at least `d_conv - 1` +tokens. + +## Tests + +- [`tests/unit_tests/ssm/ops/test_ssd_state_passing_cp.py`](../../../tests/unit_tests/ssm/ops/test_ssd_state_passing_cp.py) + checks the load-balancing permutation for exactness and the convolution and + SSD kernels against a full-sequence reference. +- [`tests/unit_tests/ssm/test_mamba_mixer_state_passing_cp.py`](../../../tests/unit_tests/ssm/test_mamba_mixer_state_passing_cp.py) + checks every mode end to end against a full-sequence mixer and against the + A2A CP path. +- [`tests/unit_tests/ssm/test_mamba_state_passing_cp_cuda_graph.py`](../../../tests/unit_tests/ssm/test_mamba_state_passing_cp_cuda_graph.py) + covers the Megatron local and Transformer Engine CUDA Graph backends. From 622a768f954668839dbd5471ad3a634cfca6b70e Mon Sep 17 00:00:00 2001 From: Sanghun Cho Date: Mon, 24 Aug 2026 09:26:53 +0000 Subject: [PATCH 5/5] Add a benchmark for the Mamba2 context-parallel paths 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 --- .../README.md | 88 ++++ .../benchmark_mamba_state_passing_cp.py | 415 ++++++++++++++++++ 2 files changed, 503 insertions(+) create mode 100644 examples/mamba_state_passing_context_parallel/README.md create mode 100644 examples/mamba_state_passing_context_parallel/benchmark_mamba_state_passing_cp.py diff --git a/examples/mamba_state_passing_context_parallel/README.md b/examples/mamba_state_passing_context_parallel/README.md new file mode 100644 index 00000000000..a1986370bec --- /dev/null +++ b/examples/mamba_state_passing_context_parallel/README.md @@ -0,0 +1,88 @@ + + +# Mamba2 State-Passing Context Parallelism Benchmark + +Compares the Mamba2 context-parallel paths against each other so the +`--mamba-state-passing-cp-load-balancing` mode can be chosen from measurements +rather than by guessing. See +[the feature guide](../../docs/user-guide/features/mamba_state_passing_context_parallel.md) +for what the paths do. + +## What is measured + +Only the post-projection part of `MambaMixer` — the causal convolution and the +SSD scan — is timed. That is where the CP paths differ; `in_proj`, the output +projection, and RMSNorm are identical across paths and would dilute the +comparison, so they are excluded. + +Four paths are measured: + +| Path | Description | +|---|---| +| `a2a` | The existing Mamba CP path: activations are redistributed so every rank holds the whole sequence for a subset of heads and groups | +| `permute-p2p` | State passing; the balanced shard is exchanged for a contiguous causal shard point-to-point | +| `permute-a2a` | State passing; the same exchange in one unequal-split `all_to_all_single` | +| `virtual` | State passing with no activation exchange; each balanced half is treated as an independent causal segment | + +## Running it + +One process per GPU; the world size is the CP size. No tensor-parallel group is +created — `--tp-size` only divides the head and group counts so that a TP-local +shard shape can be measured with one rank per CP position. + +```bash +# CP=4, 32K sequence +PYTHONPATH=. torchrun --standalone --nproc_per_node=4 \ + examples/mamba_state_passing_context_parallel/benchmark_mamba_state_passing_cp.py \ + --L 32768 --iters 20 + +# Sweep sequence lengths and batch sizes in one run +PYTHONPATH=. torchrun --standalone --nproc_per_node=8 \ + examples/mamba_state_passing_context_parallel/benchmark_mamba_state_passing_cp.py \ + --sequence-lengths 32768 131072 --batch-sizes 1 3 + +# Forward only, and just the two paths worth comparing for a decision +PYTHONPATH=. torchrun --standalone --nproc_per_node=4 \ + examples/mamba_state_passing_context_parallel/benchmark_mamba_state_passing_cp.py \ + --forward-only --paths a2a virtual + +# TP2-local shard shape under CP4 +PYTHONPATH=. torchrun --standalone --nproc_per_node=4 \ + examples/mamba_state_passing_context_parallel/benchmark_mamba_state_passing_cp.py \ + --tp-size 2 --L 32768 --batch 3 +``` + +Defaults are Nemotron-3 Nano's mixer shape: 64 heads of 64, 8 groups, state 128, +SSD chunk 128, `d_conv` 4, BF16 activations. + +## Timing method + +Each iteration is timed with CUDA events and reduced across ranks with `MAX`, +because a CP path is only as fast as its slowest rank. Rank alignment happens on +a barrier outside the timed interval, so waiting for stragglers is not counted +twice. Warmup iterations are discarded, and peak allocated memory is reported +alongside latency. + +Output is one human-readable line per path plus a machine-readable `RESULT,...` +line per path for scripted collection: + +```text +Mamba CP paths (fwd+bwd) L=32768 batch=1 local_L=8192 cp=4 ... + a2a: 12.34 ms +/- 0.12 p50= 12.30 peak= 1.23 GiB + permute-p2p: 9.87 ms +/- 0.09 p50= 9.85 peak= 0.98 GiB +p50 speedup vs a2a: permute-p2p=1.249x, ... +RESULT,cp=4,L=32768,batch=1,tp_size=1,path=a2a,mean_ms=...,p50_ms=...,peak_gib=... +``` + +## Constraints + +The sequence length must divide the CP size, each rank's local length must be +even (balanced CP gives each rank two chunks), and each local half must be a +multiple of the SSD chunk size. The script asserts all three before measuring. diff --git a/examples/mamba_state_passing_context_parallel/benchmark_mamba_state_passing_cp.py b/examples/mamba_state_passing_context_parallel/benchmark_mamba_state_passing_cp.py new file mode 100644 index 00000000000..e842de89af8 --- /dev/null +++ b/examples/mamba_state_passing_context_parallel/benchmark_mamba_state_passing_cp.py @@ -0,0 +1,415 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Benchmark the Mamba2 context-parallel paths against each other. + +Measures the post-projection part of ``MambaMixer`` -- the causal convolution +and the SSD scan, which is where the CP paths differ -- for the existing +all-to-all path and for each state-passing load-balancing mode. The projection +layers are excluded on purpose: they are identical across paths and would dilute +the comparison. + +Run from the repository root, one process per GPU: + + PYTHONPATH=. torchrun --standalone --nproc_per_node=4 \\ + examples/mamba_state_passing_context_parallel/benchmark_mamba_state_passing_cp.py \\ + --L 32768 --iters 20 + +Sweep several shapes in one run: + + PYTHONPATH=. torchrun --standalone --nproc_per_node=8 \\ + examples/mamba_state_passing_context_parallel/benchmark_mamba_state_passing_cp.py \\ + --sequence-lengths 32768 131072 --batch-sizes 1 3 + +The world size is the CP size; no tensor-parallel group is created. ``--tp-size`` +only divides the head and group counts so a TP-local shard shape can be measured +on a single rank per CP position. +""" + +import argparse +import statistics +from dataclasses import dataclass +from types import SimpleNamespace +from typing import Callable, Dict + +import torch +import torch.distributed as dist +import torch.nn as nn +from einops import rearrange +from mamba_ssm.ops.triton.ssd_combined import mamba_split_conv1d_scan_combined + +from megatron.core.ssm.mamba_context_parallel import MambaContextParallel +from megatron.core.ssm.ops.ssd_state_passing_cp import MambaStatePassingCPAdapter + +CP_PATH_NAMES = ("a2a", "permute-p2p", "permute-a2a", "virtual") + + +@dataclass(frozen=True) +class MambaCoreShape: + """TP-local Mamba mixer shape driving the measured kernels.""" + + nheads: int = 64 + headdim: int = 64 + ngroups: int = 8 + d_state: int = 128 + d_conv: int = 4 + chunk_size: int = 128 + + @property + def d_inner(self) -> int: + """Inner (gated) width of the mixer.""" + return self.nheads * self.headdim + + @property + def projected_width(self) -> int: + """Width of the ``in_proj`` output (z, x, B, C, dt).""" + return 2 * self.d_inner + 2 * self.ngroups * self.d_state + self.nheads + + @property + def conv_dim(self) -> int: + """Number of channels the depthwise causal convolution operates on.""" + return self.d_inner + 2 * self.ngroups * self.d_state + + +@dataclass(frozen=True) +class BenchmarkMetrics: + """Per-iteration latency statistics and peak memory for one CP path.""" + + mean_ms: float + stdev_ms: float + p50_ms: float + peak_gib: float + + +@dataclass +class MambaCPInputs: + """Rank-local activation and output gradient for one measured configuration.""" + + projected_input: torch.Tensor + grad_output: torch.Tensor + + +class MambaCPWorkload(nn.Module): + """Minimal post-projection mixer satisfying the contract both CP paths expect. + + ``MambaContextParallel`` and ``MambaStatePassingCPAdapter`` read a specific + set of attributes off the mixer. Reproducing just those keeps the benchmark + independent of the surrounding model, spec, and process-group plumbing. + """ + + def __init__(self, shape: MambaCoreShape, device: torch.device, cp_group: dist.ProcessGroup): + super().__init__() + self.shape = shape + dtype = torch.bfloat16 + generator = torch.Generator(device=device).manual_seed(4321) + + self.nheads_local_tp = shape.nheads + self.d_inner_local_tp = shape.d_inner + self.ngroups_local_tp = shape.ngroups + self.d_state = shape.d_state + self.d_conv = shape.d_conv + self.headdim = shape.headdim + self.chunk_size = shape.chunk_size + self.activation = "silu" + self.act = nn.SiLU() + # RMSNorm and the output projection are outside the measured region. + self.rmsnorm = False + self.norm_before_gate = False + self.D_has_hdim = False + self.config = SimpleNamespace(mamba_state_passing_cp_load_balancing="permute_p2p") + + def randn(*sizes, dtype=dtype, scale=1.0): + return torch.randn(*sizes, device=device, dtype=dtype, generator=generator) * scale + + self.conv1d_weight = nn.Parameter(randn(shape.conv_dim, 1, self.d_conv, scale=0.02)) + self.conv1d_bias = nn.Parameter(randn(shape.conv_dim, scale=0.02)) + self.dt_bias = nn.Parameter(randn(self.nheads_local_tp)) + self.A_log = nn.Parameter(randn(self.nheads_local_tp, dtype=torch.float32)) + self.D = nn.Parameter(randn(self.nheads_local_tp, dtype=torch.float32)) + + self.cp = MambaContextParallel( + cp_group=cp_group, + d_inner_local_tp=self.d_inner_local_tp, + nheads_local_tp=self.nheads_local_tp, + ngroups_local_tp=self.ngroups_local_tp, + d_state=self.d_state, + conv1d_weight_cp1=self.conv1d_weight, + conv1d_bias_cp1=self.conv1d_bias, + conv1d_padding=self.d_conv - 1, + dt_bias_cp1=self.dt_bias, + A_log_cp1=self.A_log, + D_cp1=self.D, + D_has_hdim=self.D_has_hdim, + ) + self.state_passing_cp_adapter = MambaStatePassingCPAdapter(self) + + +def parse_args(): + """Parse benchmark shape, sweep, and iteration-count arguments.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--L", dest="sequence_length", type=int, default=32768) + parser.add_argument("--batch", dest="batch_size", type=int, default=1) + parser.add_argument("--sequence-lengths", type=int, nargs="+") + parser.add_argument("--batch-sizes", type=int, nargs="+") + parser.add_argument("--nheads", type=int, default=64) + parser.add_argument("--headdim", type=int, default=64) + parser.add_argument("--ngroups", type=int, default=8) + parser.add_argument("--d-state", type=int, default=128) + parser.add_argument("--d-conv", type=int, default=4) + parser.add_argument("--chunk", type=int, default=128) + parser.add_argument( + "--tp-size", + type=int, + default=1, + help=( + "emulate a tensor-parallel local Mamba shard by dividing --nheads and " + "--ngroups; no TP process group or TP communication is created" + ), + ) + parser.add_argument("--iters", type=int, default=20) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--forward-only", action="store_true") + parser.add_argument("--paths", nargs="+", choices=CP_PATH_NAMES, default=CP_PATH_NAMES) + return parser.parse_args() + + +def shape_from_args(args) -> MambaCoreShape: + """Build the TP-local shape the benchmark measures.""" + assert args.tp_size > 0 + assert args.nheads % args.tp_size == 0, "global nheads must divide TP size" + assert args.ngroups % args.tp_size == 0, "global ngroups must divide TP size" + return MambaCoreShape( + nheads=args.nheads // args.tp_size, + headdim=args.headdim, + ngroups=args.ngroups // args.tp_size, + d_state=args.d_state, + d_conv=args.d_conv, + chunk_size=args.chunk, + ) + + +def validate_shape(shape: MambaCoreShape, sequence_length: int, cp_size: int) -> int: + """Check the CP and SSD alignment constraints and return the local length.""" + assert shape.nheads % shape.ngroups == 0 + assert sequence_length % cp_size == 0, "global sequence length must divide CP size" + local_length = sequence_length // cp_size + assert local_length % 2 == 0, "load-balanced CP requires two local chunks" + assert ( + local_length // 2 + ) % shape.chunk_size == 0, "each local load-balanced chunk must align with SSD chunk size" + return local_length + + +def create_inputs( + workload: MambaCPWorkload, + sequence_length: int, + batch_size: int, + rank: int, + device: torch.device, + cp_size: int, +) -> MambaCPInputs: + """Allocate the rank-local projected activation and output gradient.""" + local_length = validate_shape(workload.shape, sequence_length, cp_size) + generator = torch.Generator(device=device).manual_seed(1234 + rank) + projected_input = torch.randn( + local_length, + batch_size, + workload.shape.projected_width, + device=device, + dtype=torch.bfloat16, + generator=generator, + requires_grad=True, + ) + grad_output = torch.randn( + local_length, + batch_size, + workload.d_inner_local_tp, + device=device, + dtype=torch.bfloat16, + generator=generator, + ) + return MambaCPInputs(projected_input=projected_input, grad_output=grad_output) + + +def build_cp_path_runners( + workload: MambaCPWorkload, inputs: MambaCPInputs +) -> Dict[str, Callable[[], torch.Tensor]]: + """Return one output-producing closure per measured CP path.""" + + def a2a(): + projected = workload.cp.pre_conv_ssm(inputs.projected_input) + projected = rearrange(projected, "l b d -> b l d").contiguous() + output = mamba_split_conv1d_scan_combined( + projected, + rearrange(workload.cp.get_conv1d_weight(), "d 1 w -> d w"), + workload.cp.get_conv1d_bias(), + workload.cp.get_dt_bias().float(), + -torch.exp(workload.cp.get_A_log().float()), + D=workload.cp.get_D(), + chunk_size=workload.chunk_size, + activation=workload.activation, + headdim=workload.headdim, + ngroups=workload.cp.ngroups_local_tpcp, + norm_before_gate=workload.norm_before_gate, + ) + output = rearrange(output, "b l d -> l b d").contiguous() + return workload.cp.post_conv_ssm(output) + + def state_passing(load_balancing): + workload.config.mamba_state_passing_cp_load_balancing = load_balancing + return workload.state_passing_cp_adapter.forward(inputs.projected_input) + + return { + "a2a": a2a, + "permute-p2p": lambda: state_passing("permute_p2p"), + "permute-a2a": lambda: state_passing("permute_a2a"), + "virtual": lambda: state_passing("virtual"), + } + + +def benchmark_cp_path( + name: str, + runner: Callable[[], torch.Tensor], + *, + warmup: int, + iterations: int, + forward_only: bool, + workload: MambaCPWorkload, + inputs: MambaCPInputs, + group: dist.ProcessGroup, + device: torch.device, + rank: int, +) -> BenchmarkMetrics: + """Time one CP path, reporting the per-iteration maximum across ranks. + + A CP path is only as fast as its slowest rank, so each iteration is reduced + with MAX rather than averaged. Rank alignment happens on a barrier outside + the timed interval so that waiting for stragglers is not counted twice. + """ + times = [] + peaks = [] + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + + torch.cuda.synchronize(device) + dist.barrier(group=group) + for iteration in range(warmup + iterations): + inputs.projected_input.grad = None + workload.zero_grad(set_to_none=True) + dist.barrier(group=group) + torch.cuda.reset_peak_memory_stats(device) + start.record() + output = runner() + if not forward_only: + torch.autograd.backward(output, inputs.grad_output) + end.record() + end.synchronize() + + metrics = torch.tensor( + [start.elapsed_time(end), torch.cuda.max_memory_allocated(device) / 1024**3], + device=device, + dtype=torch.float64, + ) + dist.all_reduce(metrics, op=dist.ReduceOp.MAX, group=group) + if iteration >= warmup: + times.append(metrics[0].item()) + peaks.append(metrics[1].item()) + del output + torch.cuda.synchronize(device) + dist.barrier(group=group) + + result = BenchmarkMetrics( + mean_ms=statistics.mean(times), + stdev_ms=statistics.stdev(times) if len(times) > 1 else 0.0, + p50_ms=statistics.median(times), + peak_gib=max(peaks), + ) + if rank == 0: + print( + f"{name:>18}: {result.mean_ms:8.2f} ms +/- {result.stdev_ms:6.2f} " + f"p50={result.p50_ms:8.2f} peak={result.peak_gib:6.2f} GiB", + flush=True, + ) + return result + + +def run_configuration(args, workload, context, sequence_length, batch_size): + """Benchmark every requested path for one (sequence length, batch size).""" + rank, cp_size, device, group = context + local_length = validate_shape(workload.shape, sequence_length, cp_size) + inputs = create_inputs(workload, sequence_length, batch_size, rank, device, cp_size) + runners = build_cp_path_runners(workload, inputs) + + if rank == 0: + mode = "fwd" if args.forward_only else "fwd+bwd" + print( + f"Mamba CP paths ({mode}) L={sequence_length} batch={batch_size} " + f"local_L={local_length} cp={cp_size} tp={args.tp_size}(local-shape-only) " + f"d_inner={workload.d_inner_local_tp} nheads={workload.nheads_local_tp} " + f"ngroups={workload.ngroups_local_tp}", + flush=True, + ) + + results = { + name: benchmark_cp_path( + name, + runners[name], + warmup=args.warmup, + iterations=args.iters, + forward_only=args.forward_only, + workload=workload, + inputs=inputs, + group=group, + device=device, + rank=rank, + ) + for name in args.paths + } + if rank != 0: + return + + if "a2a" in results: + baseline = results["a2a"].p50_ms + speedups = ", ".join( + f"{name}={baseline / metrics.p50_ms:.3f}x" + for name, metrics in results.items() + if name != "a2a" + ) + print(f"p50 speedup vs a2a: {speedups}", flush=True) + + for name, metrics in results.items(): + print( + f"RESULT,cp={cp_size},L={sequence_length},batch={batch_size}," + f"tp_size={args.tp_size},path={name},mean_ms={metrics.mean_ms:.6f}," + f"stdev_ms={metrics.stdev_ms:.6f},p50_ms={metrics.p50_ms:.6f}," + f"peak_gib={metrics.peak_gib:.6f}", + flush=True, + ) + print("Times are per-iteration max latency across ranks; lower is better.", flush=True) + + +def main(): + """Initialize one process per GPU and sweep the requested configurations.""" + args = parse_args() + assert args.iters > 0 and args.warmup >= 0 + + local_rank = dist.get_node_local_rank() + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + dist.init_process_group("nccl", device_id=device) + group = dist.group.WORLD + context = (dist.get_rank(), dist.get_world_size(), device, group) + + workload = MambaCPWorkload(shape_from_args(args), device, group) + for sequence_length in args.sequence_lengths or [args.sequence_length]: + for batch_size in args.batch_sizes or [args.batch_size]: + run_configuration(args, workload, context, sequence_length, batch_size) + torch.cuda.empty_cache() + torch.cuda.synchronize(device) + dist.barrier(group=group) + + dist.destroy_process_group() + + +if __name__ == "__main__": + main()