[egilliam-nv]

Summary

#3186 passes cu_seqlens directly to cuDNN SDPA fprop. Because cuDNN reads [actual_b+1] entries from the user's buffers, that PR builds the forward graph at the exact batch size instead of the bucketed one. For packed THD training the number of documents per batch varies, so every new count misses the graph cache and rebuilds a graph. This shows up as a large throughput regression in Llama-3 training (about -48% on H100 across #3186 in our bisect).

Change

• Bucket the forward graph batch size again on the direct path (config_and_params.cpp).
• Add pad_cu_seqlens_tails, a small kernel that extends the cu_seqlens and ragged-offset buffers to the bucketed batch by repeating the last value (zero-length sequences), so cuDNN never reads out of bounds. It is launched only when the bucketed batch differs from the actual one, uses one launch (vs. two conversion kernels before [Common] Pass cu_seqlens and token-unit ragged offsets directly to cuDNN SDPA fprop #3186), and needs 4 * (b_bucket+1) * 4 bytes of extra workspace.
• Wired into both the f16/bf16 and FP8 forward launchers.

Status (draft)

Tested on H100 (sm90), cuDNN 9.24.0, bf16, built with NVTE_CUDA_ARCHS=90:

• Fused THD forward with 17 distinct document counts (3 to 129, fixed total tokens) builds only 4 forward graphs (one per batch bucket) per NVTE_FUSED_ATTN_CACHE_DEBUG=1, and matches a per-document SDPA reference within bf16 noise (max abs error <= 0.0078).
• pytest tests/pytorch/attention/test_attention.py -k thd: 79 passed, 1532 skipped (hardware), 0 failed.

TODO:

• Throughput comparison vs. pre-[Common] Pass cu_seqlens and token-unit ragged offsets directly to cuDNN SDPA fprop #3186 and [Common] Pass cu_seqlens and token-unit ragged offsets directly to cuDNN SDPA fprop #3186 (Llama-3 style packed THD).
• Other archs (sm100) and FP8 THD forward.
• A graph-cache regression test for varying document counts.

🤖 Generated with Claude Code