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 plumbs FP8 attention support for the THD (token × head × dim, ragged/packed variable-length) format, gated on cuDNN ≥ 9.23 and sm100+. It also fixes a pre-existing out-of-bounds read in the FWD path where b instead of actual_b was passed as the read count to cu_seqlens_to_actual_seqlens.

Backend selection (fused_attn.cpp): THD is now admitted for the NVTE_FP8 backend with an explicit mask-type guard (padding or padding_causal only) and a supported_ragged_offset_size check for large-tensor int64 offsets.
FWD/BWD impl (fused_attn_fp8.cu): Adds ragged-offset cuDNN graph tensors (Q, K, V, O, Stats), a new cu_seqlens_padded_to_offsets workspace region, and quantised batch/token dimensions. SM120 follows a separate branch that keeps dense stats and skips batch quantisation.
Tests: thd_thd_thd layout and thd format added to the FP8-vs-F16 comparison sweep; fp8_output and fast_zero_fill parameters forwarded to the attention call.

Confidence Score: 4/5

The change is substantial but well-structured; the pre-existing OOB read is fixed and the ragged-offset workspace arithmetic checks out. The main area needing a second look is the asymmetry between FWD and BWD in how b is quantised when use_cu_seqlens_directly is active.

The OOB-read fix (actual_b vs b in cu_seqlens_to_actual_seqlens) is correctly applied in FWD. Workspace size calculations for both the actual-seqlen and ragged-offset regions are consistent with kernel bounds checks in cu_seqlens_padded_to_offsets. SM120 dense-stats and non-120 ragged-stats branches are cleanly separated. The FWD/BWD asymmetry in b promotion is by design rather than a bug.

transformer_engine/common/fused_attn/fused_attn_fp8.cu — the FWD/BWD impl interaction around use_cu_seqlens_directly, actual_b vs max_b, and ragged-offset workspace layout; the SM120 branch path should be validated end-to-end.

Important Files Changed

Filename	Overview
transformer_engine/common/fused_attn/fused_attn_fp8.cu	Core THD+FP8 plumbing: adds ragged-offset graph tensors (Q/K/V/O/Stats), workspace for offset buffers, quantised batch/token sizes, and fixes the pre-existing cu_seqlens OOB read. SM120 follows a separate dense-stats branch.
transformer_engine/common/fused_attn/fused_attn.cpp	Backend selection updated to admit THD for FP8 on cuDNN >= 9.23 / sm100+, with correct mask-type guard. FWD and BWD dispatch updated to pass new t_q/t_kv and padded-seqlens arguments.
transformer_engine/common/fused_attn/fused_attn_fp8.h	Header updated to match new fused_attn_fp8_fwd/bwd signatures with num_tokens_q/kv and cu_seqlens_q/kv_padded parameters.
tests/pytorch/attention/test_attention.py	Extends FP8-vs-F16 tests to cover thd_thd_thd layout and thd format; adds fp8_output and fast_zero_fill (THD-only) arguments.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[nvte_fused_attn_fwd] --> B{backend == NVTE_FP8?}
    B -- yes --> C[fused_attn_fp8_fwd]
    B -- no --> D[F16 path]
    C --> E{qkv_format == THD?}
    E -- yes --> F[compute max_batch_size and max_tokens]
    E -- no --> G[skip quantisation]
    F --> H[fused_attn_fp8_fwd_impl]
    G --> H
    H --> I{THD and cuDNN >= 9.23?}
    I -- no --> J[cu_seqlens to actual_seqlens]
    I -- yes --> K{sm_arch != 120?}
    K -- sm100 non-120 --> L[b=max_b s_q=max_t_q ragged-offset graph use_ragged_stats=true]
    K -- SM120 --> M[b=actual_b ragged-offset graph use_ragged_stats=false]
    L --> N[workspace: plan + actual_seqlen + offset buffers]
    M --> N
    N --> O[cu_seqlens_padded_to_offsets kernel]
    J --> P[cu_seqlens_to_actual_seqlens kernel]
    O --> Q[cuDNN SDPA FP8 kernel with ragged tensors]
    P --> Q

_{Reviews (6): Last reviewed commit: "Fix FP8 direct-seqlen merge integration" | 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

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

		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: 4/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!

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