[wayrise]

Summary

presets/fast.env, which scripts/train.sh loads by default since #58, sets OPENWAM_SDPA_CUDNN=1, so masked SDPA runs on cuDNN first.

The cuDNN SDPA backward bundled with PyTorch 2.7.1 (cudnn-frontend 1.10, cuDNN 9.7.1) has a bug. When s_kv % 128 != 0, it computes the last 128-key tile without a padding mask. For a query row whose visible logits are all very negative (logsumexp below about −88.7), the backward then evaluates exp(0 − lse) = inf on the out-of-range key columns and writes NaN into that row of dQ. The forward output is still correct. See NVIDIA/cudnn-frontend#160 and pytorch/pytorch#166211; cudnn-frontend ≥ 1.15 applies a padding mask itself.

OpenWAM hits this. At 384×320 with 33-frame windows, the joint MoT attention has 392 keys and the text cross-attention has 64. Training dual-system joint_self_attn from the raw Wan2.2-TI2V-5B weights with attention_mask_mode=action_sees_video on LIBERO went NaN after 5–21 steps:

• At step 21 the loss is still finite (1.611), but DeepSpeed reports a grad norm of exactly 1.0, which is its −1 sentinel for a non-finite norm.
• bf16 has no overflow skip, so the update is applied.
• Every loss from step 22 onward is NaN.

Fix. When cuDNN is preferred and the call will be differentiated, sdpa() pads K/V with zero keys up to the next multiple of 128. The padded keys are hidden in the mask: False for a bool mask, -inf for an additive mask, or a new key-padding mask when none is given. They receive exactly zero attention, so outputs and dQ/dK/dV are unchanged. K/V keep their physical layout, for the same reason given in _zero_masked_rows. Calls without autograd are not padded, because the forward kernel is unaffected.

Validation (H200, torch 2.7.1+cu128, cuDNN 9.7.1)

Minimal reproduction. cuDNN forced. One query row is anti-aligned with keys that share a large component (row logsumexp ≈ −111 to −120). B=2, H=4, D=128.

s_q × s_kv	mask	non-finite dQ elements, main	this PR
392 × 392	bool / additive / none	1024 / 1024 / 1024	0 / 0 / 0
360 × 64	bool / additive / none	1024 / 1024 / 1024	0 / 0 / 0
100 × 77	bool / additive / none	1024 / 1024 / 1024	0 / 0 / 0
392 × 512	bool / additive / none	0	0

• 1024 = B·H·D, i.e. exactly the affected row. dK, dV and the forward output are finite in every case.
• With the fix, outputs match an fp32 MATH reference within bf16 rounding (max abs ≤ 0.016) and gradients within ≤ 1 % relative error.
• With a mask, the forward output is bit-identical to main.

Training.

• The run that diverged (dual JSA, action_sees_video, Wan2.2 init, LIBERO, 7×H200, 32/GPU, seed 42, presets/fast.env):
  • main: NaN from step 22.
  • This PR: same losses through step 21 (1.6111 on both), then finite past the warm-up peak to step 340 (loss 0.205).
  • Speed: 1.343 s/step over steps 100–330, vs 1.347 s/step with OPENWAM_SDPA_CUDNN=0.
• Official LIBERO recipe (fine-tune the OpenWAM-α foundation model, mutual, cuDNN on):
  • Losses match main within 0.0018 over 40 steps.
  • Speed: 1.36 vs 1.32 s/step (+3 %).
  • Peak memory: +2.8 GB from the padded K/V copies.

Tests.

• New tests/test_sdpa_kv_padding.py:
  • CPU float64 exactness of the padding with bool, additive, no and key-broadcast masks; 3-D inputs keep their rank.
  • Layout preservation, including BSHD buffers viewed as BHSD and broadcast axes.
  • GPU regression with forced cuDNN: 5 shapes (including s_q on the tile with s_kv off it) × 3 mask kinds × 2 layouts = 30 cases.
  • 4 more GPU cases put the low row in the last, partial query tile, with head_dim 64 (the Cosmos joint_cross_attn bridge heads) and 128.
  • All 34 GPU regression cases fail on main and pass here.
  • Only differentiated calls are padded: each half of the guard is tested separately, and calls that are not padded stay bit-identical to stock cuDNN.
• tests/test_sdpa_empty_rows.py and tests/test_fast_kernels.py, including their GPU tests: 77 passed together with the new file.
• CPU suite (-m "not gpu"): 2152 passed, 1 failed; the failure, test_cosmos_predict25_build_with_bad_path_raises_clear_error, fails identically on main.

Notes.

• The padding only runs on the cuDNN-preferred path (OPENWAM_SDPA_CUDNN=1).
• K/V are copied into a zero-filled buffer with copy_, which is vectorized for bf16, unlike torch.cat in PyTorch 2.7.
• The padding can be dropped once PyTorch bundles cudnn-frontend ≥ 1.16. numeric bug in CUDNN_ATTENTION, NaN pytorch/pytorch#166211 reports the bug in 2.8.0 as well.

🤖 Generated with Claude Code