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-packed, variable-length) format, including context-parallel support. It adds ragged-offset plumbing through the cuDNN frontend graph for Q/K/V/O/Stats tensors, gated on cuDNN ≥ 9.23 and SM ≥ 100, and fixes a pre-existing UB in mha_fill that read device-side cu_seqlens on the host CPU.

fused_attn_fp8.cu: Extends both fwd_impl and bwd_impl with max_b/max_t_q/max_t_kv parameters, ragged-offset tensor construction and workspace accounting, SM120 dense-stats branch, and fixes the cu_seqlens_to_actual_seqlens call to pass actual_b (not b) as the input count.
attention.cpp: Replaces the partial suffix-zero mha_fill (which derived an offset by dereferencing device memory on the host) with a simple stream-ordered full-tensor zero.
context_parallel.py: Wraps gathered FP8 partials as Float8Tensor before dispatch and guards the [2, s//2, …] reshape to skip THD in the backward pass.

Confidence Score: 4/5

Safe to merge for existing configurations; the new THD+FP8 path is hardware-gated and mirrors the validated F16 path, and the FWD seqlen OOB fix is correctly applied.

The FWD/BWD batch-dimension asymmetry relies on an undocumented cuDNN invariant about seqlen=0 elements, and the ragged-offset grid formula lacks a clarifying comment — both are fragile under future maintenance.

Files Needing Attention: transformer_engine/common/fused_attn/fused_attn_fp8.cu — FWD/BWD batch-size branches and ragged-offset kernel grid formula.

Important Files Changed

Filename	Overview
transformer_engine/common/fused_attn/fused_attn_fp8.cu	Core FP8 SDPA extended to support THD ragged format: adds max_b/max_t_q/max_t_kv, ragged-offset tensor setup for Q/K/V/O/Stats, SM120 dense-stats branch, workspace accounting, and fixes the pre-existing seqlen OOB in the FWD path.
transformer_engine/common/fused_attn/fused_attn.cpp	Backend selector updated to allow NVTE_FP8 for THD format when cuDNN >= 9.23 and sm_arch >= 100 with padding mask; removes the now-redundant !requires_64bit_ragged_offset gate from BSHD/SBHD/BHSD paths.
transformer_engine/common/fused_attn/fused_attn_fp8.h	Public header updated to match new fused_attn_fp8_fwd/bwd signatures (num_tokens_q, num_tokens_kv, cu_seqlens_q/kv_padded added).
transformer_engine/pytorch/csrc/extensions/attention.cpp	mha_fill refactored from a partial suffix-zero that read device cu_seqlens on the host (UB) to a full stream-ordered tensor zero.
transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py	FP8+THD forward now wraps q/k/v partials as Float8Tensors; backward skips the [2, s//2] reshape for THD; stale f16-only comment removed.
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Removes hard FP8+THD+CP disable, enabling the newly supported path.
tests/pytorch/attention/test_attention.py	Adds thd_thd_thd layout and thd format to FP8 test matrices; pads seqlen for cuBLAS alignment; sets fast_zero_fill=False.
tests/pytorch/attention/test_attention_with_cp.py	Adds cp_1_6 config for FP8+THD+CP coverage; replaces blanket FP8+THD skip with narrower MXFP8+THD skip.

_{Reviews (11): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | 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

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

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

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

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

3a80b91

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

		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

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!

sudhakarsingh27 commented Jul 31, 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