From eb9e23be0d63d7419d328aa4243f2520f738e05e Mon Sep 17 00:00:00 2001 From: "Claude Fable 5.1" Date: Wed, 9 Sep 2026 10:21:04 +0000 Subject: [PATCH 1/3] MiniMax-H3: use cuDNN SDPA for attention on sm_90+ GPUs The repo-wide dispatch picks FlashAttention-2 whenever flash_attn is installed. FA2's kernels are an sm80 design; on B200 at this model's shape (S~16.5k, 56 heads x 128, bf16) they reach ~440 TFLOP/s while torch's cuDNN SDPA backend reaches 1.3 PFLOP/s fwd+bwd with the same error against an fp32 reference. Route the H3 attention segments through SDPA with cuDNN first (flash/efficient/math as fallbacks) when the device is Hopper or newer; older GPUs keep the existing dispatch. Co-Authored-By: Claude Fable 5.1 --- diffsynth/models/minimax_h3_dit.py | 26 +++++++++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/diffsynth/models/minimax_h3_dit.py b/diffsynth/models/minimax_h3_dit.py index 3ed30de6..b707eac1 100644 --- a/diffsynth/models/minimax_h3_dit.py +++ b/diffsynth/models/minimax_h3_dit.py @@ -67,6 +67,30 @@ def _modulate_gate(x, gate, other, indices): return (x + gate.index_select(0, indices) * other).to(x.dtype) +def _prefer_cudnn_sdpa(device) -> bool: + """On Hopper/Blackwell the cuDNN fused attention in torch SDPA is several times faster than the + FlashAttention-2 kernels the repo-wide dispatch would pick (FA2 is an sm80 design; measured 3.4x on + B200 at this model's shape, same error vs an fp32 reference). Older GPUs keep the default dispatch.""" + if device.type != "cuda": + return False + cached = _CUDNN_PREF.get(device.index) + if cached is None: + major, _ = torch.cuda.get_device_capability(device) + cached = _CUDNN_PREF[device.index] = major >= 9 + return cached + + +_CUDNN_PREF = {} + + +def _segment_attention(seg_q, seg_k, seg_v, softmax_scale): + if _prefer_cudnn_sdpa(seg_q.device): + from torch.nn.attention import sdpa_kernel, SDPBackend + with sdpa_kernel([SDPBackend.CUDNN_ATTENTION, SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH], set_priority=True): + return torch.nn.functional.scaled_dot_product_attention(seg_q, seg_k, seg_v, scale=softmax_scale) + return attention_forward(seg_q, seg_k, seg_v, scale=softmax_scale) + + def _sdpa_varlen_attention(q, k, v, cu_seqlens, softmax_scale): out = torch.empty_like(q) bounds = cu_seqlens.tolist() @@ -76,7 +100,7 @@ def _sdpa_varlen_attention(q, k, v, cu_seqlens, softmax_scale): seg_q = q[start:stop].transpose(0, 1).unsqueeze(0) seg_k = k[start:stop].transpose(0, 1).unsqueeze(0) seg_v = v[start:stop].transpose(0, 1).unsqueeze(0) - seg_out = attention_forward(seg_q, seg_k, seg_v, scale=softmax_scale) + seg_out = _segment_attention(seg_q, seg_k, seg_v, softmax_scale) out[start:stop] = seg_out.squeeze(0).transpose(0, 1) return out From f355123d0dc21987560bb4ded2454d9d38e84b15 Mon Sep 17 00:00:00 2001 From: "Claude Fable 5.1" Date: Wed, 9 Sep 2026 10:52:53 +0000 Subject: [PATCH 2/3] MiniMax-H3: compile the DiT block body; one cu_seqlens sync per forward Profile of the LoRA training step: 20% of GPU time was un-fused element-wise work inside the 50 DiT blocks (RoPE cat/neg/slice, six [S,5376] AdaLN gathers, RMSNorm, SwiGLU, residual gates), ~600 launches per block execution, each a memory-bound pass over a 160-850 MB activation. - MiniMaxH3DiTBlock.forward runs its body through one torch.compile'd function shared by all blocks (dynamic=False) when grad is enabled, i.e. training; inference keeps the eager path. DIFFSYNTH_COMPILE_DIT=0 turns it off. Attention stays a library op (cuDNN SDPA); the GEMMs stay cuBLAS. - The per-block cu_seqlens.tolist() device sync (102 per step with checkpoint recompute) moves to once per DiT forward; blocks receive the bounds as python ints. Single-segment sequences skip the scratch buffer. Co-Authored-By: Claude Fable 5.1 --- diffsynth/models/minimax_h3_dit.py | 46 +++++++++++++++++++++++++----- 1 file changed, 39 insertions(+), 7 deletions(-) diff --git a/diffsynth/models/minimax_h3_dit.py b/diffsynth/models/minimax_h3_dit.py index b707eac1..c87c5403 100644 --- a/diffsynth/models/minimax_h3_dit.py +++ b/diffsynth/models/minimax_h3_dit.py @@ -4,6 +4,7 @@ import torch import torch.nn as nn +from torch.nn.attention import sdpa_kernel, SDPBackend from ..core.attention import attention_forward from ..core.gradient import gradient_checkpoint_forward @@ -85,15 +86,18 @@ def _prefer_cudnn_sdpa(device) -> bool: def _segment_attention(seg_q, seg_k, seg_v, softmax_scale): if _prefer_cudnn_sdpa(seg_q.device): - from torch.nn.attention import sdpa_kernel, SDPBackend with sdpa_kernel([SDPBackend.CUDNN_ATTENTION, SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH], set_priority=True): return torch.nn.functional.scaled_dot_product_attention(seg_q, seg_k, seg_v, scale=softmax_scale) return attention_forward(seg_q, seg_k, seg_v, scale=softmax_scale) -def _sdpa_varlen_attention(q, k, v, cu_seqlens, softmax_scale): +def _sdpa_varlen_attention(q, k, v, cu_seqlens, softmax_scale, seq_bounds=None): + # seq_bounds: cu_seqlens as python ints, computed once per model forward so that the 50 blocks (and + # their checkpoint recomputes) do not each pay a device sync here. + bounds = cu_seqlens.tolist() if seq_bounds is None else seq_bounds + if len(bounds) == 2: + return _segment_attention(q.transpose(0, 1).unsqueeze(0), k.transpose(0, 1).unsqueeze(0), v.transpose(0, 1).unsqueeze(0), softmax_scale).squeeze(0).transpose(0, 1) out = torch.empty_like(q) - bounds = cu_seqlens.tolist() for start, stop in zip(bounds[:-1], bounds[1:]): if stop == start: continue @@ -157,7 +161,7 @@ def __init__(self, hidden_size, num_attention_heads, attention_head_dim, qk_norm self.k_norm = _norm(attention_head_dim, eps=qk_norm_eps) self.out_proj = nn.Linear(inner_dim, hidden_size, bias=False) - def forward(self, x, *, rope_freqs, cu_seqlens, max_seqlen=None): + def forward(self, x, *, rope_freqs, cu_seqlens, max_seqlen=None, seq_bounds=None): total = x.shape[0] qkv = self.qkv_proj(x) qkv = qkv.view(total, self.num_heads, 3, self.head_dim) @@ -169,7 +173,7 @@ def forward(self, x, *, rope_freqs, cu_seqlens, max_seqlen=None): if rope_freqs is not None: q = _apply_rope(q, rope_freqs) k = _apply_rope(k, rope_freqs) - out = _sdpa_varlen_attention(q, k, v, cu_seqlens=cu_seqlens, softmax_scale=self.softmax_scale) + out = _sdpa_varlen_attention(q, k, v, cu_seqlens=cu_seqlens, softmax_scale=self.softmax_scale, seq_bounds=seq_bounds) out = out.reshape(total, self.num_heads * self.head_dim) return self.out_proj(out) @@ -244,12 +248,17 @@ def __init__(self, hidden_size, num_attention_heads, attention_head_dim, ffn_hid self.mlp = MiniMaxH3MLP(hidden_size, ffn_hidden_size) self.adaln_proj = MiniMaxH3AdalnProj(hidden_size, time_embed_dim, adaln_out_features, expand_ratio=6, modality_num=MINIMAX_H3_ADALN_MODALITY_NUM) - def forward(self, x, *, t_emb, combined_indices, rope_freqs, cu_seqlens, max_seqlen): + def forward(self, x, *, t_emb, combined_indices, rope_freqs, cu_seqlens, max_seqlen, seq_bounds=None): + if _compile_blocks() and seq_bounds is not None: + return _compiled_block_forward(self, x, t_emb, combined_indices, rope_freqs, cu_seqlens, max_seqlen, seq_bounds) + return self._forward(x, t_emb=t_emb, combined_indices=combined_indices, rope_freqs=rope_freqs, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, seq_bounds=seq_bounds) + + def _forward(self, x, *, t_emb, combined_indices, rope_freqs, cu_seqlens, max_seqlen, seq_bounds=None): shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln_proj(t_emb) residual = x h = self.norm1(x) h = _modulate_scale_shift(h, shift_msa, scale_msa, combined_indices) - h = self.attn(h, rope_freqs=rope_freqs, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen) + h = self.attn(h, rope_freqs=rope_freqs, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, seq_bounds=seq_bounds) x = _modulate_gate(residual, gate_msa, h, combined_indices) residual = x h = self.norm2(x) @@ -258,6 +267,27 @@ def forward(self, x, *, t_emb, combined_indices, rope_freqs, cu_seqlens, max_seq return _modulate_gate(residual, gate_mlp, h, combined_indices) +_COMPILE_BLOCKS = None + + +def _compile_blocks() -> bool: + """torch.compile the DiT block body (norm + AdaLN modulation + RoPE + SwiGLU fuse into a few kernels; + the GEMMs and attention stay library calls). Off with DIFFSYNTH_COMPILE_DIT=0. Only used for CUDA + with grad enabled, i.e. training; inference keeps the eager path.""" + global _COMPILE_BLOCKS + if _COMPILE_BLOCKS is None: + import os + _COMPILE_BLOCKS = os.environ.get("DIFFSYNTH_COMPILE_DIT", "1") != "0" and torch.cuda.is_available() + return _COMPILE_BLOCKS and torch.is_grad_enabled() + + +def _block_body(block, x, t_emb, combined_indices, rope_freqs, cu_seqlens, max_seqlen, seq_bounds): + return block._forward(x, t_emb=t_emb, combined_indices=combined_indices, rope_freqs=rope_freqs, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, seq_bounds=seq_bounds) + + +_compiled_block_forward = torch.compile(_block_body, dynamic=False) + + class MiniMaxH3FinalLayer(nn.Module): def __init__(self, hidden_size, time_embed_dim, final_adaln_out_features, latents_dim, audio_latents_dim, patch_size, final_norm_eps): super().__init__() @@ -396,6 +426,7 @@ def forward( hidden = decoder_input cu_seqlens = cu_seqlens.to(device) + seq_bounds = tuple(cu_seqlens.tolist()) # one sync per forward instead of one per block execution for block_id, block in enumerate(self.blocks): hidden = gradient_checkpoint_forward( block, @@ -407,6 +438,7 @@ def forward( rope_freqs=rope_freqs, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, + seq_bounds=seq_bounds, ) if control_hints is not None and block_id in control_hints: hidden = hidden + control_hints[block_id].to(hidden.device, hidden.dtype) From 1e8661d07035b3c08cfdafe4621ba45f072af8a5 Mon Sep 17 00:00:00 2001 From: "Claude Fable 5.1" Date: Wed, 9 Sep 2026 11:01:11 +0000 Subject: [PATCH 3/3] MiniMax-H3: selective activation checkpointing, attention output kept With attention on cuDNN and the block body compiled, the checkpoint recompute of each block (a full second forward) is a large share of the step, and the attention forward is its most expensive op relative to the memory its output takes (~240 MB per block in bf16). Move the checkpoint inside the compiled function and use a selective policy: the fused SDPA output is saved, everything else in the block is recomputed in backward. Peak GPU memory 75.4 -> 83.6 GB per rank on the 124f 480x832 job (+11%). Falls back to the repo's gradient_checkpoint_forward for the offload and DeepSpeed variants and when compilation is off. Co-Authored-By: Claude Fable 5.1 --- diffsynth/models/minimax_h3_dit.py | 45 ++++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/diffsynth/models/minimax_h3_dit.py b/diffsynth/models/minimax_h3_dit.py index c87c5403..c1814f6a 100644 --- a/diffsynth/models/minimax_h3_dit.py +++ b/diffsynth/models/minimax_h3_dit.py @@ -287,6 +287,45 @@ def _block_body(block, x, t_emb, combined_indices, rope_freqs, cu_seqlens, max_s _compiled_block_forward = torch.compile(_block_body, dynamic=False) +# Selective activation checkpointing for the training path: the block is recomputed in backward except +# the fused attention output (the one op whose recompute is expensive relative to the memory it saves: +# ~240 MB per block, ~12 GB for 50 blocks, vs ~3 GB per block for a full set of activations). With the +# checkpoint inside the compiled function the partitioner applies the policy to the compiled graph. +_SAC_SAVE_OPS = { + torch.ops.aten._scaled_dot_product_cudnn_attention.default, + torch.ops.aten._scaled_dot_product_flash_attention.default, + torch.ops.aten._scaled_dot_product_efficient_attention.default, +} + + +def _sac_policy(ctx, op, *args, **kwargs): + from torch.utils.checkpoint import CheckpointPolicy + return CheckpointPolicy.MUST_SAVE if op in _SAC_SAVE_OPS else CheckpointPolicy.PREFER_RECOMPUTE + + +def _block_body_checkpointed(block, x, t_emb, combined_indices, rope_freqs, cu_seqlens, max_seqlen, seq_bounds): + import functools + from torch.utils.checkpoint import checkpoint, create_selective_checkpoint_contexts + return checkpoint( + _block_body, block, x, t_emb, combined_indices, rope_freqs, cu_seqlens, max_seqlen, seq_bounds, + use_reentrant=False, context_fn=functools.partial(create_selective_checkpoint_contexts, _sac_policy), + ) + + +_compiled_block_checkpointed = torch.compile(_block_body_checkpointed, dynamic=False) + + +def _use_compiled_checkpoint(use_gradient_checkpointing, use_gradient_checkpointing_offload): + if not (use_gradient_checkpointing and not use_gradient_checkpointing_offload and _compile_blocks()): + return False + try: + from ..core.gradient.gradient_checkpoint import _HAS_DEEPSPEED, deepspeed + if _HAS_DEEPSPEED and deepspeed.checkpointing.is_configured(): + return False # the repo's DeepSpeed checkpoint path keeps its own behaviour + except ImportError: + pass + return True + class MiniMaxH3FinalLayer(nn.Module): def __init__(self, hidden_size, time_embed_dim, final_adaln_out_features, latents_dim, audio_latents_dim, patch_size, final_norm_eps): @@ -427,7 +466,13 @@ def forward( hidden = decoder_input cu_seqlens = cu_seqlens.to(device) seq_bounds = tuple(cu_seqlens.tolist()) # one sync per forward instead of one per block execution + compiled_ckpt = _use_compiled_checkpoint(use_gradient_checkpointing, use_gradient_checkpointing_offload) for block_id, block in enumerate(self.blocks): + if compiled_ckpt: + hidden = _compiled_block_checkpointed(block, hidden, t_emb, combined_indices, rope_freqs, cu_seqlens, max_seqlen, seq_bounds) + if control_hints is not None and block_id in control_hints: + hidden = hidden + control_hints[block_id].to(hidden.device, hidden.dtype) + continue hidden = gradient_checkpoint_forward( block, use_gradient_checkpointing,