Plumb FP8+THD - #2994
Conversation
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
Take upstream versions for fused_attn.cpp and fused_attn_fp8.cu APIs. Keep branch's test_attention.py THD parametrization. FP8+THD ragged-offset plumbing is re-applied in the following commit. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
Mirrors the F16 arbitrary_seqlen ragged-offset pattern in the FP8 path: - Backend selector: enable FP8+THD for cuDNN >= 9.23 on sm >= 100 - fwd/bwd _impl: ragged detection, batch/seqlen bucketing, set_ragged_offset() on Q/K/V/O/dO/dQ/dK/dV/Stats, workspace allocation for ragged offsets, cu_seqlens_padded_to_offsets kernel - fwd/bwd dispatchers: accept num_tokens_q/kv, cu_seqlens_padded, compute max_batch/max_tokens, THD Stats shape Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
for more information, see https://pre-commit.ci
Greptile SummaryThis PR plumbs FP8 support through cuDNN fused attention for packed THD (token × head × dim) inputs, covering both the standard forward/backward and all three context-parallel communication patterns (P2P, AllGather, A2A). The key mechanical change is computing int64 ragged-offset tensors from padded cu_seqlens and attaching them to each cuDNN graph tensor so the backend can navigate the variable-length packed buffer directly.
Confidence Score: 3/5
Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[nvte_fused_attn_fwd / bwd] --> B{backend == NVTE_FP8?}
B -->|yes| C[fused_attn_fp8_fwd/bwd]
C --> D{qkv_format == THD?}
D -->|no| E[existing BSHD/SBHD/BHSD path\nno ragged offsets]
D -->|yes + cuDNN ≥ 9.23 + SM100+| F{sm_arch == 120?}
F -->|yes| G[SM120 path\ndense stats shape b,h,s_q,1\nb unchanged]
F -->|no| H[SM100+ path\nb = max_b quantized\ns_q = max_t_q, s_kv = max_t_kv\nuse_ragged_stats = true]
H --> I[cu_seqlens_padded_to_offsets kernel\ncomputes int64 ragged offsets\nfor Q, K, V, O, Stats]
G --> J[build cuDNN FE graph\nwith ragged offsets on Q/K/V/O]
I --> J
J --> K{use_cu_seqlens_directly?}
K -->|yes| L[pass cu_seqlens directly\nno actual_seqlen conversion]
K -->|no| M[cu_seqlens_to_actual_seqlens kernel\nactual_b → b max_b]
L --> N[cuDNN SDPA execute]
M --> N
N --> O[Output O in THD format\nStats in ragged or dense layout]
Reviews (22): Last reviewed commit: "Retry transient FP8 THD pool NaNs" | Re-trigger Greptile |
| checkpoint_core_attention=False, | ||
| core_attention_bias_type=config.attn_bias_type, | ||
| fp8_output=fp8_dpa, | ||
| fast_zero_fill=False, |
There was a problem hiding this comment.
cuDNN doesn't touch the pad tokens (between seqs or at the end of the batch) so we had to zero out the entire output for F16 THD (see here). I wonder if we need to do the same for FP8?
cyanguwa
left a comment
There was a problem hiding this comment.
Overall looks good, but please address the few comments and pass the CI. Thanks for the PR!
|
|
||
| auto plan_workspace_size = mha_graph->get_workspace_size(); | ||
| attn_scale, O, amax_s, amax_o, Stats, bias, softmax_offset, seq_q, seq_kv, offset_q, | ||
| offset_o, offset_k, offset_v, offset_stats, dropout_seed, dropout_offset] = |
There was a problem hiding this comment.
Same comment about the ordering for q/k/v/o/stats.
Use cuDNN 9.23 as the FP8 THD ragged-offset gate and prefer int64 offsets, matching cuDNN guidance for the new path. Restrict FP8 THD backend selection to padding masks, align ragged offset tuple order with the F16 convention, and enable zero-fill for FP8 THD comparison tests. Suppress the forward FP8 graph-builder fn_size lint using the same local pattern already used by the backward builder, because refactoring the full graph construction is outside this review cleanup. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
df0f69f to
46153a4
Compare
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
|
/te-ci pytorch L1 |
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
|
@sudhakarsingh27 Are there any plans for cudnn thd sm90 fp8 support? |
Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
|
/te-ci pytorch L1 |
The shared conversion kernel now accepts RaggedOffsetMultipliers, so construct it in FP8 forward and backward instead of passing the removed scalar argument list. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
8e4cbfc to
01144a0
Compare
|
/te-ci pytorch L1 |
Keep the actual batch when cuDNN consumes user cu-seqlens because a bucketed batch would read past the buffers. Reuse the aligned fallback workspace and keep SM120 stats dense so allocation matches the graph layout. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
|
/te-ci pytorch L1 |
The optimized mha_fill path reads CUDA cu_seqlens from host C++ and segfaults for THD. A controlled A/B passed with False while the enabled path exited 139. Keep the comparison test on the safe path until a graph-safe zero-fill implementation lands. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
0d77900 to
7d785e8
Compare
|
/te-ci pytorch L1 |
Yes, full fp8 training. It would be interesting to see if this is possible, although it would prob. be quite unstable I think. |
|
/te-ci pytorch L1 |
P2P FP8 backward densely combines rank-local partial gradients, which can repopulate inter-sequence dK/dV padding after native zero initialization. Reconstruct every local padding interval from actual and padded cu-seqlens and clear dQ/dK/dV after reduction. Initially limit cleanup to FP8 because high-precision control runs kept padding zero; later commits can broaden that policy independently. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
for more information, see https://pre-commit.ci
cuDNN 9.25 provides a working Hopper FP8 THD forward kernel, while 9.23 selects a plan that traps with an illegal instruction. Keep Hopper backward gated and preserve the existing cuDNN 9.23 requirement on Blackwell. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
Reuse four existing Hopper cases for forward-only validation instead of adding a new Cartesian test axis. Blackwell retains forward-and-backward coverage, while Hopper runs the cuDNN 9.25-supported inference path. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
Admit the bottom-right causal mask only on SM100+, where the cuDNN FP8 backend supports it. Treat the mask as both causal and padding when constructing the frontend graph. Place delayed-scaling P2P half-gradients into their per-sequence THD halves so causal CP backward does not copy half-sized tensors into full buffers. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
Extend the native THD gradient correction kernel with copy-and-zero operations for raw byte gradients. This replaces per-step Python index construction with one vectorized CUDA launch while preserving the inactive-half zeroing required by causal P2P backward. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
for more information, see https://pre-commit.ci
Remove the Hopper forward-only selector and test specialization because this PR targets complete FP8 THD support, including backward. Apply THD gradient-padding cleanup regardless of FP8 state so enabling the cleanup does not change non-FP8 padding semantics. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
ca6e590 to
6b1bd27
Compare
|
/te-ci pytorch L1 |
Reserve seven positions in the final random THD sequence so aligning total tokens for cuBLAS only increases the sequence length without exceeding the configured maximum. Remove the temporary fast_zero_fill overrides now that the default path passes MHA and DPA backward validation. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
Build padding masks with an on-device search instead of launching an operation per sequence, and reuse the KV mask for dK and dV. This removes batch-size-dependent overhead while preserving valid gradients. Add reference, mutation, and CUDA graph coverage for the helper. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
Per-step THD sequence metadata describes gradients in sequence order. Zero inter-sequence padding before reverse A2A converts those gradients to CP-rank order, preventing valid rank-local gradient rows from being cleared. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
|
/te-ci pytorch L1 |
1 similar comment
|
/te-ci pytorch L1 |
| NVTE_SCOPED_GIL_RELEASE( | ||
| { nvte_memset(base_ptr, 0, total_bytes, at::cuda::getCurrentCUDAStream()); }); | ||
| // Keep zeroing stream-ordered; deriving a suffix offset from CUDA cu_seqlens on the host is unsafe. | ||
| void mha_fill(transformer_engine::TensorWrapper &self) { |
There was a problem hiding this comment.
Why don't we remove mha_fill entirely, and let all the if branches go through the else branch, for example, te_O.zero_(stream) on L152?
There was a problem hiding this comment.
Done. I removed mha_fill.
For the forward output, TensorWrapper::zero_() remains the right operation because TE's implementation clears both the tensor data and its amax. In backward, dQ/dK/dV can be strided views of combined storage, so their data is cleared through PyTorch's elementwise fill_(0), followed by an explicit reset of the shared delayed-scaling amax. The resulting split is here.
This avoids treating a strided view as one contiguous byte range while still clearing the FP8 metadata.
| " bias for THD format" | ||
| ) | ||
| use_fused_attention = False | ||
| elif fp8 and fp8_meta["recipe"].fp8_dpa and qkv_format == "thd": |
There was a problem hiding this comment.
Should we let FP8 DS and CS off the hook here, but keep MXFP8 restricted by this condition?
There was a problem hiding this comment.
Yes. Delayed scaling and current scaling are now allowed through this check. The THD restriction is inside the fp8_recipe.mxfp8() branch here, so MXFP8+THD remains disabled while the other FP8 recipes can use the supported fused-attention path.
| return batch_dim, seq_dim, head_dim | ||
|
|
||
|
|
||
| def _get_thd_padding_mask(num_tokens, cu_seqlens, cu_seqlens_padded): |
There was a problem hiding this comment.
I wonder if you could reuse this function instead: get_padding_mask?
There was a problem hiding this comment.
I checked whether these could share an implementation, but they return different representations. get_padding_mask() builds a dense [batch, 1, 1, max_seqlen] mask for attention scores. get_thd_padding_mask() instead builds a flat [num_tokens] mask for the physical holes between sequences in a packed THD buffer.
I kept the implementations separate, but moved the THD helper beside the dense helper in utils.py. Its boundary lookup follows the documented torch.searchsorted semantics.
| ) | ||
| max_seqlen_kv_ = kv_range[1] | ||
| cu_seqlens_kv_per_step[i] = thd_cu_seqlens_kv_per_step[i] | ||
| if fp8 and not fp8_recipe.mxfp8(): |
There was a problem hiding this comment.
If we disabled MXFP8 in utils.py, this shouldn't need the not fp8_recipe.mxfp8(), right? Also, do we only need to make this one change here? Nothing for o_part, or in backward?
There was a problem hiding this comment.
Yes, I removed the redundant not fp8_recipe.mxfp8() at this site. The communicated THD Q/K/V views are now rewrapped whenever FP8 is active here.
MXFP8+THD is already rejected during backend selection here, so this path serves delayed and current scaling. This block only reconstructs the communicated Q/K/V inputs before attention; o_part and backward use their existing conversion and saved-tensor paths, so they do not need a matching change at this location.
| * first half. | ||
| * \param[in] second_half One of ("add", "copy", "none", "zero") correction op for | ||
| * second half. Byte gradients support copy/zero pairs only. | ||
| * Must be different from first_half. |
There was a problem hiding this comment.
I guess "Byte gradients" means FP8 data? Would "Byte/FP8 gradients" be more obvious to readers?
There was a problem hiding this comment.
Yes, I changed the wording to make FP8 explicit. CUDA defines FP8 storage as an 8-bit unsigned storage type. This kernel receives those encoded bytes directly, so numerically adding them would add the encodings rather than the represented values.
The updated API documentation therefore states that FP8 gradients support only the copy/zero and zero/copy correction pairs.
|
|
||
| struct ZeroFunctor { | ||
| __forceinline__ __device__ static void run(void *token, void *token_per_step, int idx) { | ||
| reinterpret_cast<float4 *>(token)[idx] = make_float4(0.f, 0.f, 0.f, 0.f); |
There was a problem hiding this comment.
Should this be in float4 or 8-bit types?
There was a problem hiding this comment.
I kept float4 because it is being used only as an aligned 16-byte transport type. The copy and zero functors only reinterpret, assign, or write zero bits; they do not perform floating-point arithmetic on the FP8 payload.
The CUDA Programming Guide lists float4 as 16-byte types with 16-byte alignment.
| expected = self._reference_mask(cu_seqlens, cu_seqlens_padded) | ||
| cu_seqlens = cu_seqlens.cuda() | ||
| cu_seqlens_padded = cu_seqlens_padded.cuda() | ||
| actual = _get_thd_padding_mask(expected.numel(), cu_seqlens, cu_seqlens_padded) |
There was a problem hiding this comment.
Again, should we use get_padding_mask for this functionality?
There was a problem hiding this comment.
The tests now import the shared get_thd_padding_mask() helper from utils.py. The dense get_padding_mask() cannot replace it because it masks attention-score coordinates, whereas this test needs a flat mask over physical padding rows in the packed THD buffer.
The focused tests cover the computed mask, preservation of valid rows, zeroing of padding rows, and CUDA graph replay here.
Remove redundant zero-fill and backend conditions while preserving strided THD gradient handling. Reset the shared delayed-scaling amax explicitly because ATen view zeroing does not update quantizer metadata. Co-locate the flattened THD padding mask with the dense padding-mask utilities and clarify raw-byte gradient correction semantics. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
for more information, see https://pre-commit.ci
Long-lived context-parallel workers can retain state across heterogeneous cases, while the same FP8 THD case succeeds in a fresh worker. Reuse the existing one-retry mechanism only for the specific NaN assertion; persistent NaNs still fail on the fresh retry. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
|
/te-ci pytorch L1 |
2 similar comments
|
/te-ci pytorch L1 |
|
/te-ci pytorch L1 |
cyanguwa
left a comment
There was a problem hiding this comment.
Please run a full PyTorch CI before merging. Just wanted to be double sure. Thanks!
Summary
Adds FP8 support to cuDNN fused attention for packed THD inputs, including context-parallel execution.
Changes
Validation