Plumb FP8+THD by sudhakarsingh27 · Pull Request #2994 · NVIDIA/TransformerEngine

sudhakarsingh27 · 2026-05-14T19:09:33Z

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

Please list the changes introduced in this PR:

Change A
Change B

Checklist:

I have read and followed the contributing guidelines
The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

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-apps · 2026-05-14T19:20:12Z

Greptile Summary

This PR enables FP8 attention with THD (token × head × dim, ragged/variable-length) format for both standalone and context-parallel (CP) attention paths by wiring ragged-offset tensors into the cuDNN FP8 SDPA graph, adding quantized batch/token sizing helpers, and lifting the blanket fp8+thd skip in the backend selector and CP utility.

C++ kernel changes (fused_attn_fp8.cu, fused_attn.cpp): FWD and BWD fused_attn_fp8_*_impl now accept padded cu_seqlens, compute ragged offsets via a new cu_seqlens_padded_to_offsets kernel, and pass them to the cuDNN graph as ragged-offset tensors; also fixes the pre-existing out-of-bounds read in cu_seqlens_to_actual_seqlens (now correctly passes actual_b, max_b).
Python/CP changes (context_parallel.py, utils.py): removes the FP8+THD+CP block, wraps per-step slices in Float8Tensor for the all-gather path, and adds _zero_thd_padding to clean FP8 inter-sequence padding in the P2P backward.
attention.cpp: replaces the unsafe host-side device-pointer dereference in mha_fill with a safe full zero_() call, and threads padded cu_seqlens through to the FP8 dispatch.

Confidence Score: 4/5

The C++ kernel plumbing is solid and the pre-existing out-of-bounds seqlen read is fixed, but the new _zero_thd_padding helper introduces O(batch_size) GPU synchronizations in the FP8 CP P2P backward despite its docstring claiming the opposite.
The backward pass helper _zero_thd_padding scalar-indexes CUDA tensors inside a Python for-loop, triggering ~3×batch_size implicit device-to-host copies per backward step. Its own docstring explicitly states it avoids host copies, confirming this was unintentional. For large batches this degrades training throughput meaningfully. Everything else — kernel logic, workspace sizing, SM-arch gating, test coverage, the mha_fill simplification — looks correct.
transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py — the _zero_thd_padding function needs a vectorized rewrite to match its docstring and avoid per-element GPU syncs.

Important Files Changed

Filename	Overview
transformer_engine/common/fused_attn/fused_attn_fp8.cu	Major change: adds full THD/ragged-offset support to both FWD and BWD FP8 SDPA paths; fixes the pre-existing `cu_seqlens_to_actual_seqlens` out-of-bounds call (now passes `actual_b, max_b` instead of `b, b`); introduces SM-arch gating (sm_arch_==120 takes dense stats path); overall logic is consistent between FWD and BWD though the BWD never short-circuits the actual-seqlen workspace (missing the `use_cu_seqlens_directly` optimisation the FWD has, minor inefficiency).
transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py	Adds `_zero_thd_padding` helper and calls it in the FP8+THD+CP P2P backward path; the helper's docstring claims it avoids host copies but its Python for-loop scalar-indexes CUDA tensors, causing O(batch_size) GPU synchronizations per call — a real performance defect in the backward pass.
transformer_engine/common/fused_attn/fused_attn.cpp	Backend-selection adds THD sub-condition with correct cuDNN-version and SM-arch gating and enforces padding mask types; also threads `t_q`/`t_kv` and padded cu_seqlens into both FWD and BWD FP8 dispatch.
transformer_engine/pytorch/csrc/extensions/attention.cpp	Simplifies `mha_fill` to a full zero_ of the tensor (removing the unsafe host-side dereference of a device pointer) and threads `cu_seqlens_*_padded` through to the FP8 kernel; change is correct and safer.

Sequence Diagram

sequenceDiagram
    participant Caller
    participant fused_attn_fp8_fwd as fused_attn_fp8_fwd (C++)
    participant impl as fused_attn_fp8_fwd_impl
    participant q_kernel as cu_seqlens_to_actual_seqlens
    participant off_kernel as cu_seqlens_padded_to_offsets
    participant cuDNN as cuDNN SDPA Graph

    Caller->>fused_attn_fp8_fwd: batch, max_seqlen, num_tokens, cu_seqlens, cu_seqlens_padded
    fused_attn_fp8_fwd->>fused_attn_fp8_fwd: detect THD format → get_max_batch_size / get_max_tokens
    fused_attn_fp8_fwd->>impl: max_b, max_t_q, max_t_kv, devPtrSeqOffsetsQ/KV
    impl->>impl: "b = max_b (if !use_cu_seqlens_directly)"
    impl->>impl: build cuDNN graph with ragged_offset tensors (offset_q/k/v/o/stats)
    impl->>q_kernel: actual_b, max_b → fill actual-seqlen array
    impl->>off_kernel: actual_b, max_b, cu_seqlens_padded → fill int64 ragged offsets
    impl->>cuDNN: variant_pack with ragged offsets + seq_q/kv + FP8 scales
    cuDNN-->>impl: FP8 output O, amax, stats
    impl-->>Caller: output tensors

_{Reviews (14): Last reviewed commit: "Merge branch 'main' into fp8_thd_attenti..." | Re-trigger Greptile}

cyanguwa · 2026-06-08T21:33:12Z

            checkpoint_core_attention=False,
            core_attention_bias_type=config.attn_bias_type,
+            fp8_output=fp8_dpa,
+            fast_zero_fill=False,


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

Overall looks good, but please address the few comments and pass the CI. Thanks for the PR!

cyanguwa · 2026-06-08T21:41:36Z

+                   std::shared_ptr<fe::graph::Tensor_attributes>,   // offset_o
+                   std::shared_ptr<fe::graph::Tensor_attributes>,   // offset_k
+                   std::shared_ptr<fe::graph::Tensor_attributes>,   // offset_v
+                   std::shared_ptr<fe::graph::Tensor_attributes>,   // offset_stats


Nit: do we want to follow the same order of listing as in F16, i.e. offset_q/k/v/o/stats?

cyanguwa · 2026-06-08T21:44:41Z

-
-    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] =


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>

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

sudhakarsingh27 · 2026-06-10T06:45:21Z

/te-ci pytorch L1

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

cyanguwa · 2026-06-10T22:17:07Z

+      auto offset_q_tuple = is_ragged_q ? std::make_tuple(offset_q) : std::make_tuple(nullptr);
      auto offset_kv_tuple =
          is_ragged_kv ? std::make_tuple(offset_k, offset_v) : std::make_tuple(nullptr, nullptr);
+      auto offset_o_tuple = is_ragged_q ? std::make_tuple(offset_o) : std::make_tuple(nullptr);


Ok, maybe I was making a bigger deal out of this than it is. It looks like we are following the q/k/v/o/stats order in the shared_ptr declaration but not in the tuple creation, in the F16 file. So feel free to revert this change, or just change the shared_ptr order if you like. Thanks.

bbuschkaemper · 2026-06-12T12:17:09Z

@sudhakarsingh27 Are there any plans for cudnn thd sm90 fp8 support?

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

sudhakarsingh27 · 2026-07-22T22:56:41Z

/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>

sudhakarsingh27 · 2026-07-22T23:24:19Z

/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>

sudhakarsingh27 · 2026-07-23T19:22:08Z

/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>

sudhakarsingh27 · 2026-07-24T21:12:17Z

/te-ci pytorch L1

sudhakarsingh27 · 2026-07-29T23:28:33Z

/te-ci pytorch L1

FP8 Linear flattens THD inputs to [t, h*d], so adjust generated sequence lengths before building cu_seqlens to make the total token count divisible by eight for both forward and backward. Keep fast zero fill disabled in the MHA helper to avoid the known host dereference path. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

sudhakarsingh27 · 2026-07-31T08:03:44Z

/te-ci pytorch L1

Remove obsolete FP8+THD+CP gates now that the fused-attention path supports this combination. Replace the unsafe host-side suffix calculation with stream-ordered zeroing, and preserve FP8 metadata and token-major THD layout in the all-gather path. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

for more information, see https://pre-commit.ci

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. Gate the cleanup on FP8 backward because EOS H100 and Prenyx B200 controls showed BF16 and high-precision backward padding already remains zero. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

for more information, see https://pre-commit.ci

sudhakarsingh27 · 2026-08-10T23:32:49Z

/te-ci pytorch L1

sudhakarsingh27 · 2026-08-10T23:33:24Z

@bbuschkaemper, are you looking for both fwd and bwd in fp8+thd? SM90 in cudnn only supports fp8+thd in fwd right now

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>

sudhakarsingh27 added 6 commits December 9, 2025 15:37

update_filter_fp8_thd_attention

1b9b13e

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

Merge branch 'main' of github.com:NVIDIA/TransformerEngine

977bfb6

add for fp8+thd debug

dc4a172

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

uncomment the configs

10aa194

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

Merge main into fp8_thd_attention_fixes

148fe40

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>

sudhakarsingh27 requested a review from cyanguwa as a code owner May 14, 2026 19:09

[pre-commit.ci] auto fixes from pre-commit.com hooks

bfbafe9

for more information, see https://pre-commit.ci

greptile-apps Bot reviewed May 14, 2026

View reviewed changes

Comment thread tests/pytorch/attention/test_attention.py Outdated

Comment thread transformer_engine/common/fused_attn/fused_attn.cpp Outdated

cyanguwa reviewed Jun 8, 2026

View reviewed changes

cyanguwa requested changes Jun 8, 2026

View reviewed changes

sudhakarsingh27 force-pushed the fp8_thd_attention_try2 branch from df0f69f to 46153a4 Compare June 10, 2026 05:44

Merge remote-tracking branch 'origin/main' into fp8_thd_attention_try2

70f7b08

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

Fix FP8 THD ragged version scope

14a9518

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

cyanguwa reviewed Jun 10, 2026

View reviewed changes

cyanguwa added the 2.18 label Jul 13, 2026

resolve conflicts with main

f0e4fcf

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

Fix FP8 ragged offset kernel calls after main merge

01144a0

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>

sudhakarsingh27 force-pushed the fp8_thd_attention_try2 branch from 8e4cbfc to 01144a0 Compare July 22, 2026 23:23

sudhakarsingh27 force-pushed the fp8_thd_attention_try2 branch from 0d77900 to 7d785e8 Compare July 24, 2026 18:47

sudhakarsingh27 self-assigned this Jul 29, 2026

sudhakarsingh27 added 2.19 and removed 2.18 labels Jul 29, 2026

Merge branch 'main' into fp8_thd_attention_try2

fbb9b1b

sudhakarsingh27 requested a review from ksivaman as a code owner July 31, 2026 23:49

pre-commit-ci Bot and others added 4 commits July 31, 2026 23:51

[pre-commit.ci] auto fixes from pre-commit.com hooks

3a80b91

for more information, see https://pre-commit.ci

[pre-commit.ci] auto fixes from pre-commit.com hooks

0a41f76

for more information, see https://pre-commit.ci

Merge branch 'main' into fp8_thd_attention_try2

044ad35

sudhakarsingh27 added 3 commits August 10, 2026 21:24

Merge branch 'main' into fp8_thd_attention_try2

8c08530

Uh oh!

Conversation

sudhakarsingh27 commented May 14, 2026

Description

Type of change

Changes

Checklist:

Uh oh!

greptile-apps Bot commented May 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 4/5

Important Files Changed

Sequence Diagram

Uh oh!

Uh oh!

Uh oh!

cyanguwa Jun 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cyanguwa left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cyanguwa Jun 8, 2026

Choose a reason for hiding this comment

Uh oh!

cyanguwa Jun 8, 2026

Choose a reason for hiding this comment

Uh oh!

sudhakarsingh27 commented Jun 10, 2026

Uh oh!

cyanguwa Jun 10, 2026

Choose a reason for hiding this comment

Uh oh!

bbuschkaemper commented Jun 12, 2026

Uh oh!

sudhakarsingh27 commented Jul 22, 2026

Uh oh!

sudhakarsingh27 commented Jul 22, 2026

Uh oh!

sudhakarsingh27 commented Jul 23, 2026

Uh oh!

sudhakarsingh27 commented Jul 24, 2026

Uh oh!

sudhakarsingh27 commented Jul 29, 2026

Uh oh!

sudhakarsingh27 commented Jul 31, 2026

Uh oh!

sudhakarsingh27 commented Aug 10, 2026

Uh oh!

sudhakarsingh27 commented Aug 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

greptile-apps Bot commented May 14, 2026 •

edited

Loading

cyanguwa Jun 8, 2026 •

edited

Loading

sudhakarsingh27 commented Aug 10, 2026 •

edited

Loading