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 adds FP8 fused-attention support for THD ragged inputs.

Extends FP8 backend selection to THD layouts on supported Blackwell and cuDNN configurations, restricted to padding masks.
Plumbs token counts, padded sequence metadata, and 64-bit ragged offsets through FP8 forward and backward execution.
Corrects forward sequence-length conversion to distinguish the actual and quantized batch sizes.
Expands PyTorch attention coverage to include THD and FP8 output paths.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains; the prior mask-selection mismatch, debug output, and forward sequence-length out-of-bounds issue are fixed at the current head.

Important Files Changed

Filename	Overview
tests/pytorch/attention/test_attention.py	Extends FP8 attention test coverage to THD layouts and configures FP8 output with the safe zero-fill path.
transformer_engine/common/fused_attn/fused_attn.cpp	Selects the FP8 backend for supported THD padding-mask configurations and forwards token and padded-sequence metadata.
transformer_engine/common/fused_attn/fused_attn_fp8.cu	Implements THD-aware FP8 forward and backward graphs, ragged offsets, workspace allocation, and safe actual-batch sequence conversion.
transformer_engine/common/fused_attn/fused_attn_fp8.h	Updates internal FP8 fused-attention declarations with token counts and padded sequence-length tensors.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[PyTorch FP8 attention with THD tensors] --> B[FP8 backend selection]
  B -->|cuDNN 9.23+, supported SM, padding mask| C[FP8 fused-attention dispatcher]
  C --> D[Convert sequence lengths]
  C --> E[Generate 64-bit ragged offsets]
  D --> F[cuDNN FP8 SDPA graph]
  E --> F
  F --> G[THD output and training context]
  G --> H[FP8 backward graph]

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

greptile-apps · 2026-05-14T19:20:15Z

+    print(f"qkv_format: {qkv_format}")
+    print(f"inp shape: {inp[0].shape}, {inp[1].shape}, {inp[2].shape}")


Debug prints committed to test code

Two print statements were left in _run_dpa_fp8_vs_f16, which will produce noise on every test run and in CI logs. These look like leftover debug instrumentation from development.

Suggested change

print(f"qkv_format: {qkv_format}")

print(f"inp shape: {inp[0].shape}, {inp[1].shape}, {inp[2].shape}")

with autocast(enabled=fp8_dpa, recipe=fp8_recipe):

We can probably remove these debug prints once everything is verified to be fine?

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:38:39Z

+  bool is_ragged_kv = (kv_format == NVTE_QKV_Format::NVTE_THD);
+  const int device_id = cuda::current_device();
+  const int sm_arch_ = cuda::sm_arch(device_id);
+  bool use_ragged_stats = is_ragged_q && cudnn_runtime_version >= 90600 && sm_arch_ != 120;


I wonder if all the 90600 occurrences should be 92300, given that we only start to support FP8 THD from 9.23?

cyanguwa · 2026-06-08T21:39:34Z

+  bool use_ragged_stats = is_ragged_q && cudnn_runtime_version >= 90600 && sm_arch_ != 120;
+
+  NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout);
+  const DType ragged_offset_type = cudnn_runtime_version >= 90500 ? DType::kInt64 : DType::kInt32;


Could you check with cuDNN see if FP8 THD is supported with both types, or just one, i.e. kInt64? Also, should the runtime version here really be 92300 (if both types were supported)?

cyanguwa · 2026-06-08T21:40:11Z

+  const DType ragged_offset_type = cudnn_runtime_version >= 90500 ? DType::kInt64 : DType::kInt32;
+
+  int64_t actual_b = b;
+  if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600) {


Same as above.

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

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

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

		print(f"qkv_format: {qkv_format}")
		print(f"inp shape: {inp[0].shape}, {inp[1].shape}, {inp[2].shape}")

	print(f"qkv_format: {qkv_format}")
	print(f"inp shape: {inp[0].shape}, {inp[1].shape}, {inp[2].shape}")
	with autocast(enabled=fp8_dpa, recipe=fp8_recipe):

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

Important Files Changed

Flowchart

Uh oh!

greptile-apps Bot May 14, 2026

Choose a reason for hiding this comment

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!

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!

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!

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!

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!

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

cyanguwa Jun 8, 2026 •

edited

Loading