From 1b9b13e0cc2df9aab5840b51e1c4b59832ee3938 Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Tue, 9 Dec 2025 15:37:17 -0800 Subject: [PATCH 01/20] update_filter_fp8_thd_attention Signed-off-by: Sudhakar Singh --- tests/pytorch/attention/test_attention.py | 28 ++++++++++++----------- 1 file changed, 15 insertions(+), 13 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 4aedcff1b8..dfd797cd66 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -1757,23 +1757,23 @@ def get_model(dtype, config): model_configs_fp8_vs_f16 = { # test: ModelConfig(b, sq, hq, dqk) - "fp8_9": ModelConfig(2, 2048, 16, 128), - "fp8_10": ModelConfig(2, 2048, 24, 128, num_gqa_groups=12), - "fp8_11": ModelConfig(1, 8192, 32, 128, num_gqa_groups=4), - "fp8_12": ModelConfig(2, 2048, 16, 128, attn_mask_type="causal"), - "fp8_13": ModelConfig(2, 2048, 24, 128, num_gqa_groups=12, attn_mask_type="causal"), - "fp8_14": ModelConfig(1, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="causal"), - "fp8_15": ModelConfig(2, 2048, 16, 128, attn_mask_type="padding"), - "fp8_16": ModelConfig(2, 2048, 24, 128, num_gqa_groups=12, attn_mask_type="padding"), - "fp8_17": ModelConfig(1, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="padding"), - "fp8_18": ModelConfig(2, 2048, 16, 128, attn_mask_type="padding_causal"), + # "fp8_9": ModelConfig(2, 2048, 16, 128), + # "fp8_10": ModelConfig(2, 2048, 24, 128, num_gqa_groups=12), + # "fp8_11": ModelConfig(1, 8192, 32, 128, num_gqa_groups=4), + # "fp8_12": ModelConfig(2, 2048, 16, 128, attn_mask_type="causal"), + # "fp8_13": ModelConfig(2, 2048, 24, 128, num_gqa_groups=12, attn_mask_type="causal"), + # "fp8_14": ModelConfig(1, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="causal"), + # "fp8_15": ModelConfig(2, 2048, 16, 128, attn_mask_type="padding"), + # "fp8_16": ModelConfig(2, 2048, 24, 128, num_gqa_groups=12, attn_mask_type="padding"), + # "fp8_17": ModelConfig(1, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="padding"), + # "fp8_18": ModelConfig(2, 2048, 16, 128, attn_mask_type="padding_causal"), "fp8_19": ModelConfig(2, 2048, 24, 128, num_gqa_groups=12, attn_mask_type="padding_causal"), - "fp8_20": ModelConfig(1, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="padding_causal"), + # "fp8_20": ModelConfig(1, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="padding_causal"), } param_types_fp8_vs_f16 = [torch.float16, torch.bfloat16] -qkv_layout_fp8_vs_f16 = ["sbh3d", "bshd_bshd_bshd", "sbhd_sbhd_sbhd"] -qkv_format_fp8_vs_f16 = ["bshd", "sbhd"] +qkv_layout_fp8_vs_f16 = ["sbh3d", "bshd_bshd_bshd", "sbhd_sbhd_sbhd", "thd_thd_thd"] +qkv_format_fp8_vs_f16 = ["bshd", "sbhd", "thd"] @pytest.mark.skipif(get_cudnn_version() < (9, 2, 1), reason="cuDNN 9.2.1+ is required.") @@ -2300,6 +2300,8 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: out_grad_shape_new = [*out_grad_shape[:-2], out_grad_shape[-2] * out_grad_shape[-1]] out_grad = torch.randn(out_grad_shape_new, dtype=dtype, device="cuda") + 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): out = dpa( inp[0], From dc4a172fb41a2b61993d185359ae150fdbdfa8db Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Fri, 2 Jan 2026 15:22:51 -0800 Subject: [PATCH 02/20] add for fp8+thd debug Signed-off-by: Sudhakar Singh --- tests/pytorch/attention/test_attention.py | 1 + .../common/fused_attn/fused_attn.cpp | 3 +-- .../common/fused_attn/fused_attn_fp8.cu | 27 +++++-------------- 3 files changed, 8 insertions(+), 23 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 553d7b30db..539eaa9ea4 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -2316,6 +2316,7 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: checkpoint_core_attention=False, core_attention_bias_type=config.attn_bias_type, fp8_output=fp8_dpa, + fast_zero_fill=False, ) if is_training: out.backward(out_grad) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index fde0d38921..1a9ccc26ee 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -252,7 +252,6 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK || attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK))) && - (qkv_format == NVTE_QKV_Format::NVTE_BSHD || qkv_format == NVTE_QKV_Format::NVTE_SBHD) && !requires_64bit_ragged_offset && (softmax_type == NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX) && // 9.10.0: known bugs with SDPA FP8 (cudnn_runtime_version != 91000) && !return_max_logit) { @@ -539,7 +538,7 @@ void nvte_fused_attn_fwd_qkvpacked(const NVTETensor QKV, const NVTETensor Bias, h = input_QKV->data.shape[ndim - 3]; } else { NVTE_ERROR("nvte_fused_attn_fwd_qkvpacked only supports H3D and 3HD layouts!"); - } + } size_t d = input_QKV->data.shape[ndim - 1]; size_t t = 0; NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 3630041ccf..9cdbb49492 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -2474,7 +2474,8 @@ void fused_attn_fp8_fwd(size_t batch, size_t num_attn_heads, size_t num_gqa_grou size_t workspace_size = 0; NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); - if ((qkv_format == NVTE_QKV_Format::NVTE_BSHD) || (qkv_format == NVTE_QKV_Format::NVTE_SBHD)) { + if ((qkv_format == NVTE_QKV_Format::NVTE_BSHD) || (qkv_format == NVTE_QKV_Format::NVTE_SBHD) + || (qkv_format == NVTE_QKV_Format::NVTE_THD)) { fused_attn::fused_attn_fp8_fwd_impl_v1( batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim, is_training, attn_scale, p_dropout, qkv_layout, bias_type, mask_type, devPtrQ, devPtrK, devPtrV, devPtrM, @@ -2482,16 +2483,8 @@ void fused_attn_fp8_fwd(size_t batch, size_t num_attn_heads, size_t num_gqa_grou devPtrScaleS, devPtrScaleO, devPtrAmaxO, devPtrAmaxS, devPtrcuSeqlensQ, devPtrcuSeqlensKV, devPtrDropoutSeed, devPtrDropoutOffset, get_cudnn_fe_dtype(QKV_type), get_cudnn_fe_dtype(O_type), workspace->data.dptr, &workspace_size, stream, handle); - } else if (qkv_layout == NVTE_QKV_Layout::NVTE_T3HD) { - fused_attn::fused_attn_fp8_fwd_impl( - batch, num_attn_heads, max_seqlen_q, max_seqlen_kv, head_dim, is_training, attn_scale, - p_dropout, qkv_layout, devPtrQ, devPtrK, devPtrV, devPtrM, devPtrZInv, devPtrO, - devPtrDescaleQ, devPtrDescaleK, devPtrDescaleV, devPtrDescaleS, devPtrScaleS, devPtrScaleO, - devPtrAmaxO, devPtrAmaxS, devPtrcuSeqlensQ, devPtrcuSeqlensKV, devPtrDropoutSeed, - devPtrDropoutOffset, get_cudnn_dtype(QKV_type), workspace->data.dptr, &workspace_size, - stream, handle); } else { - NVTE_ERROR("FP8 fused attention only supports qkv_layout=t3hd or qkv_format=bshd/sbhd. \n"); + NVTE_ERROR("FP8 fused attention only supports qkv_format=bshd/sbhd/thd. \n"); } if (workspace_size > 0) { @@ -2569,7 +2562,8 @@ void fused_attn_fp8_bwd(size_t batch, size_t num_attn_heads, size_t num_gqa_grou size_t workspace_size = 0; NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); - if ((qkv_format == NVTE_QKV_Format::NVTE_BSHD) || (qkv_format == NVTE_QKV_Format::NVTE_SBHD)) { + if ((qkv_format == NVTE_QKV_Format::NVTE_BSHD) || (qkv_format == NVTE_QKV_Format::NVTE_SBHD) + || (qkv_format == NVTE_QKV_Format::NVTE_THD)) { fused_attn::fused_attn_fp8_bwd_impl_v1( batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim, attn_scale, p_dropout, qkv_layout, bias_type, mask_type, devPtrQ, devPtrK, devPtrV, devPtrM, devPtrZInv, @@ -2580,17 +2574,8 @@ void fused_attn_fp8_bwd(size_t batch, size_t num_attn_heads, size_t num_gqa_grou devPtrDropoutSeed, devPtrDropoutOffset, get_cudnn_fe_dtype(QKV_type), get_cudnn_fe_dtype(O_type), get_cudnn_fe_dtype(dO_type), get_cudnn_fe_dtype(dQKV_type), workspace->data.dptr, &workspace_size, stream, handle); - } else if (qkv_layout == NVTE_QKV_Layout::NVTE_T3HD) { - fused_attn::fused_attn_fp8_bwd_impl( - batch, num_attn_heads, max_seqlen_q, max_seqlen_kv, head_dim, attn_scale, p_dropout, - qkv_layout, devPtrQ, devPtrK, devPtrV, devPtrM, devPtrZInv, devPtrO, devPtrdO, devPtrdQ, - devPtrdK, devPtrdV, devPtrDescaleQ, devPtrDescaleK, devPtrDescaleV, devPtrDescaleO, - devPtrDescaledO, devPtrDescaleS, devPtrDescaledP, devPtrScaleS, devPtrScaledP, - devPtrScaledQ, devPtrScaledK, devPtrScaledV, devPtrAmaxdP, devPtrAmaxdQ, devPtrAmaxdK, - devPtrAmaxdV, devPtrcuSeqlensQ, devPtrcuSeqlensKV, devPtrDropoutSeed, devPtrDropoutOffset, - get_cudnn_dtype(QKV_type), workspace->data.dptr, &workspace_size, stream, handle); } else { - NVTE_ERROR("FP8 fused attention only supports qkv_layout=t3hd or qkv_format=bshd/sbhd. \n"); + NVTE_ERROR("FP8 fused attention only supports qkv_format=bshd/sbhd/thd. \n"); } if (workspace_size > 0) { From 10aa1940a53bb3b16b0ffa043903596f9b1e3eb7 Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Fri, 2 Jan 2026 15:49:16 -0800 Subject: [PATCH 03/20] uncomment the configs Signed-off-by: Sudhakar Singh --- tests/pytorch/attention/test_attention.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 539eaa9ea4..9b940caae5 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -1757,18 +1757,18 @@ def get_model(dtype, config): model_configs_fp8_vs_f16 = { # test: ModelConfig(b, sq, hq, dqk) - # "fp8_9": ModelConfig(2, 2048, 16, 128), - # "fp8_10": ModelConfig(2, 2048, 24, 128, num_gqa_groups=12), - # "fp8_11": ModelConfig(1, 8192, 32, 128, num_gqa_groups=4), - # "fp8_12": ModelConfig(2, 2048, 16, 128, attn_mask_type="causal"), - # "fp8_13": ModelConfig(2, 2048, 24, 128, num_gqa_groups=12, attn_mask_type="causal"), - # "fp8_14": ModelConfig(1, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="causal"), - # "fp8_15": ModelConfig(2, 2048, 16, 128, attn_mask_type="padding"), - # "fp8_16": ModelConfig(2, 2048, 24, 128, num_gqa_groups=12, attn_mask_type="padding"), - # "fp8_17": ModelConfig(1, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="padding"), - # "fp8_18": ModelConfig(2, 2048, 16, 128, attn_mask_type="padding_causal"), + "fp8_9": ModelConfig(2, 2048, 16, 128), + "fp8_10": ModelConfig(2, 2048, 24, 128, num_gqa_groups=12), + "fp8_11": ModelConfig(1, 8192, 32, 128, num_gqa_groups=4), + "fp8_12": ModelConfig(2, 2048, 16, 128, attn_mask_type="causal"), + "fp8_13": ModelConfig(2, 2048, 24, 128, num_gqa_groups=12, attn_mask_type="causal"), + "fp8_14": ModelConfig(1, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="causal"), + "fp8_15": ModelConfig(2, 2048, 16, 128, attn_mask_type="padding"), + "fp8_16": ModelConfig(2, 2048, 24, 128, num_gqa_groups=12, attn_mask_type="padding"), + "fp8_17": ModelConfig(1, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="padding"), + "fp8_18": ModelConfig(2, 2048, 16, 128, attn_mask_type="padding_causal"), "fp8_19": ModelConfig(2, 2048, 24, 128, num_gqa_groups=12, attn_mask_type="padding_causal"), - # "fp8_20": ModelConfig(1, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="padding_causal"), + "fp8_20": ModelConfig(1, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="padding_causal"), } param_types_fp8_vs_f16 = [torch.float16, torch.bfloat16] From ad50a2e7b8daf3ff7268ea8b62bd4b827360450d Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Tue, 12 May 2026 16:48:34 -0700 Subject: [PATCH 04/20] Re-plumb FP8+THD ragged-offset support on top of merged main 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 --- .../common/fused_attn/fused_attn.cpp | 34 +- .../common/fused_attn/fused_attn_fp8.cu | 429 ++++++++++++++++-- .../common/fused_attn/fused_attn_fp8.h | 40 +- 3 files changed, 422 insertions(+), 81 deletions(-) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index d2eb1a831c..df70bcce66 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -281,12 +281,15 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK))) && // pre-9.21: {bshd, sbhd}, {vanilla} // 9.21+: {bshd, sbhd, bhsd}, {vanilla, off-by-one, learnable} + // 9.23+ sm100+: {thd} (ragged/packed variable-length) ((cudnn_runtime_version < 92100 && (qkv_format == NVTE_QKV_Format::NVTE_BSHD || qkv_format == NVTE_QKV_Format::NVTE_SBHD) && softmax_type == NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX) || (cudnn_runtime_version >= 92100 && (qkv_format == NVTE_QKV_Format::NVTE_BSHD || qkv_format == NVTE_QKV_Format::NVTE_SBHD || - qkv_format == NVTE_QKV_Format::NVTE_BHSD))) && + qkv_format == NVTE_QKV_Format::NVTE_BHSD)) || + (cudnn_runtime_version >= 92300 && sm_arch_ >= 100 && + qkv_format == NVTE_QKV_Format::NVTE_THD && supported_ragged_offset_size)) && !requires_64bit_ragged_offset && // 9.10.0: known bugs with SDPA FP8 (cudnn_runtime_version != 91000) && !return_max_logit) { @@ -623,12 +626,14 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso input_cu_seqlens_kv, input_cu_seqlens_q_padded, input_cu_seqlens_kv_padded, input_page_table_k, input_page_table_v, input_rng_state, wkspace, stream, handle); } else if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_FP8) { - fused_attn_fp8_fwd(b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, is_training, - attn_scale, dropout, qkv_layout, o_format, qkv_scale_inv_format, bias_type, - attn_mask_type, softmax_type, window_size_left, window_size_right, - bottom_right_diagonal, input_Q, input_K, input_V, input_SoftmaxOffset, - input_output_S, output_O, Aux_CTX_Tensors, input_cu_seqlens_q, - input_cu_seqlens_kv, input_rng_state, wkspace, stream, handle); + fused_attn_fp8_fwd(b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, t_q, t_kv, + is_training, attn_scale, dropout, qkv_layout, o_format, + qkv_scale_inv_format, bias_type, attn_mask_type, softmax_type, + window_size_left, window_size_right, bottom_right_diagonal, input_Q, + input_K, input_V, input_SoftmaxOffset, input_output_S, output_O, + Aux_CTX_Tensors, input_cu_seqlens_q, input_cu_seqlens_kv, + input_cu_seqlens_q_padded, input_cu_seqlens_kv_padded, input_rng_state, + wkspace, stream, handle); } else { NVTE_ERROR("Invalid combination of data type and sequence length for fused attention. \n"); } @@ -725,13 +730,14 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso if (input_dO->scaling_mode == NVTE_MXFP8_1D_SCALING) { input_dO_f16 = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); } - fused_attn_fp8_bwd(b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, attn_scale, dropout, - qkv_layout, o_format, do_format, dqkv_layout, qkv_scale_inv_format, - do_scale_inv_format, bias_type, attn_mask_type, softmax_type, - window_size_left, window_size_right, bottom_right_diagonal, deterministic, - input_Q, input_K, input_V, input_O, input_dO, input_dO_f16, input_M, input_S, - input_SoftmaxOffset, input_output_dP, output_dQ, output_dK, output_dV, - output_dSoftmaxOffset, input_cu_seqlens_q, input_cu_seqlens_kv, + fused_attn_fp8_bwd(b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, t_q, t_kv, + attn_scale, dropout, qkv_layout, o_format, do_format, dqkv_layout, + qkv_scale_inv_format, do_scale_inv_format, bias_type, attn_mask_type, + softmax_type, window_size_left, window_size_right, bottom_right_diagonal, + deterministic, input_Q, input_K, input_V, input_O, input_dO, input_dO_f16, + input_M, input_S, input_SoftmaxOffset, input_output_dP, output_dQ, + output_dK, output_dV, output_dSoftmaxOffset, input_cu_seqlens_q, + input_cu_seqlens_kv, input_cu_seqlens_q_padded, input_cu_seqlens_kv_padded, input_rng_state, wkspace, stream, handle); } else { NVTE_ERROR("Invalid combination of data type and sequence length for fused attention. \n"); diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index eab1ae02e6..978a6c0035 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -6,6 +6,7 @@ #include "../common.h" #include "../cudnn_utils.h" +#include "../util/cuda_runtime.h" #include "../util/system.h" #include "fused_attn_fp8.h" #include "utils.h" @@ -18,6 +19,7 @@ using namespace transformer_engine; // fused attention FWD FP8 with FE 1.0+ void fused_attn_fp8_fwd_impl( int64_t b, int64_t h, int64_t hg, int64_t s_q, int64_t s_kv, int64_t d_qk, int64_t d_v, + int64_t max_b, int64_t max_t_q, int64_t max_t_kv, bool is_training, float scaling_factor, float dropout_probability, NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, @@ -25,7 +27,8 @@ void fused_attn_fp8_fwd_impl( void* devPtrSoftmaxOffset, void* devPtrM, void* devPtrO, void* devPtrDescaleQ, void* devPtrDescaleK, void* devPtrDescaleV, void* devPtrDescaleS, void* devPtrScaleS, void* devPtrScaleO, void* devPtrAmaxO, void* devPtrAmaxS, void* devPtrcuSeqlensQ, - void* devPtrcuSeqlensKV, void* devPtrDropoutSeed, void* devPtrDropoutOffset, + void* devPtrcuSeqlensKV, void* devPtrSeqOffsetsQ, void* devPtrSeqOffsetsKV, + void* devPtrDropoutSeed, void* devPtrDropoutOffset, cudnn_frontend::DataType_t qkv_tensor_type, cudnn_frontend::DataType_t o_tensor_type, NVTEScalingMode scaling_mode, NVTE_QKV_Format qkv_scale_inv_format, void* workspace, size_t* workspace_size, cudaStream_t stream, cudnnHandle_t handle) { @@ -60,6 +63,27 @@ void fused_attn_fp8_fwd_impl( NVTE_CHECK(!is_mxfp8 || cudnn_runtime_version >= 92100, "MXFP8 fused attention requires cuDNN 9.21.0 or later!"); + NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); + NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); + bool is_ragged_q = (q_format == NVTE_QKV_Format::NVTE_THD); + 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; + + 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; + + int64_t actual_b = b; + if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600) { + NVTE_CHECK(is_padding, "Ragged QKV input requires padding or padding_causal mask!"); + if (sm_arch_ != 120) { + b = max_b; + s_q = is_ragged_q ? max_t_q : s_q; + s_kv = is_ragged_kv ? max_t_kv : s_kv; + } + } + try { FADescriptor_v1 descriptor{b, h, @@ -121,6 +145,11 @@ void fused_attn_fp8_fwd_impl( std::shared_ptr, // softmax_offset std::shared_ptr, // seq_q std::shared_ptr, // seq_kv + std::shared_ptr, // offset_q + std::shared_ptr, // offset_o + std::shared_ptr, // offset_k + std::shared_ptr, // offset_v + std::shared_ptr, // offset_stats std::shared_ptr, // dropout_seed std::shared_ptr>; // dropout_offset @@ -146,6 +175,8 @@ void fused_attn_fp8_fwd_impl( std::shared_ptr descale_q, descale_k, descale_v; std::shared_ptr descale_s, scale_s, scale_o; std::shared_ptr bias, softmax_offset, seq_q, seq_kv; + std::shared_ptr offset_q, offset_k, offset_v, offset_o, + offset_stats; std::shared_ptr dropout_seed, dropout_offset; // Q, K, V, attn_scale @@ -157,6 +188,14 @@ void fused_attn_fp8_fwd_impl( .set_dim({b, h, s_q, d_qk}) .set_stride(q_strides) .set_data_type(qkv_tensor_type)); + if (is_ragged_q) { + offset_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_q") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + Q->set_ragged_offset(offset_q); + } K = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("K") .set_dim({b, hg, s_kv, d_qk}) @@ -167,6 +206,20 @@ void fused_attn_fp8_fwd_impl( .set_dim({b, hg, s_kv, d_v}) .set_stride(v_strides) .set_data_type(qkv_tensor_type)); + if (is_ragged_kv) { + offset_k = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_k") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + offset_v = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_v") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + K->set_ragged_offset(offset_k); + V->set_ragged_offset(offset_v); + } attn_scale = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("attn_scale") .set_dim({1, 1, 1, 1}) @@ -323,15 +376,33 @@ void fused_attn_fp8_fwd_impl( .set_dim({b, h, s_q, d_v}) .set_stride(o_strides) .set_data_type(o_tensor_type); + if (is_ragged_q) { + offset_o = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_o") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + O->set_ragged_offset(offset_o); + } amax_o->set_output(!is_mxfp8) .set_dim({1, 1, 1, 1}) .set_stride({1, 1, 1, 1}) .set_data_type(fe::DataType_t::FLOAT); - Stats->set_output(true) - .set_data_type(fe::DataType_t::FLOAT) - .set_dim({b, h, s_q, 1}) - .set_stride({h * s_q, s_q, 1, 1}); + if (use_ragged_stats) { + offset_stats = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_stats") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + } + Stats->set_output(true).set_data_type(fe::DataType_t::FLOAT).set_dim({b, h, s_q, 1}); + if (use_ragged_stats) { + Stats->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); + } else { + Stats->set_stride({h * s_q, s_q, 1, 1}); + } std::tuple, // Q std::shared_ptr, // K @@ -357,6 +428,12 @@ void fused_attn_fp8_fwd_impl( is_softmax_offset ? std::make_tuple(softmax_offset) : std::make_tuple(nullptr); auto padding_tuple = is_padding ? std::make_tuple(seq_q, seq_kv) : std::make_tuple(nullptr, nullptr); + auto offset_qo_tuple = + is_ragged_q ? std::make_tuple(offset_q, offset_o) : std::make_tuple(nullptr, nullptr); + auto offset_kv_tuple = + is_ragged_kv ? std::make_tuple(offset_k, offset_v) : std::make_tuple(nullptr, nullptr); + auto offset_s_tuple = + use_ragged_stats ? std::make_tuple(offset_stats) : std::make_tuple(nullptr); auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) : std::make_tuple(nullptr, nullptr); @@ -367,22 +444,37 @@ void fused_attn_fp8_fwd_impl( NVTE_CHECK_CUDNN_FE(mha_graph->build_plans(handle)); auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, Stats_tuple, bias_tuple, - softmax_offset_tuple, padding_tuple, dropout_tuple); + softmax_offset_tuple, padding_tuple, offset_qo_tuple, offset_kv_tuple, + offset_s_tuple, dropout_tuple); cache.insert({descriptor, return_tuple}); return return_tuple; }; auto [mha_graph, Q, K, V, descale_q, descale_k, descale_v, descale_s, scale_s, scale_o, - attn_scale, O, amax_s, amax_o, Stats, bias, softmax_offset, seq_q, seq_kv, dropout_seed, - dropout_offset] = get_graph(sdpa_fp8_fprop_cache, descriptor); - - 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] = + get_graph(sdpa_fp8_fprop_cache, descriptor); + + auto plan_workspace_size = alignTo<16>(mha_graph->get_workspace_size()); + const size_t num_bytes_per_seqlen = alignTo<16>(b * sizeof(int32_t)); + const size_t actual_seqlen_workspace_size = is_padding ? 2 * num_bytes_per_seqlen : 0; + const size_t num_bytes_per_ragged_offset = + alignTo<16>(((b + 1) * typeToNumBits(ragged_offset_type)) / 8); + size_t seqlen_offsets_workspace_size = 0; + if (is_ragged_q || is_ragged_kv) { + size_t count = 2 * (static_cast(is_ragged_q) + static_cast(is_ragged_kv)); + if (use_ragged_stats) { + seqlen_offsets_workspace_size = (count + 1) * num_bytes_per_ragged_offset; + } else { + seqlen_offsets_workspace_size = count * num_bytes_per_ragged_offset; + } + } // Exit to request upper level API to allocate memory if needed - size_t actual_seqlen_workspace_size = 2 * b * sizeof(int32_t); if (workspace == nullptr) { - *workspace_size = plan_workspace_size + actual_seqlen_workspace_size; + *workspace_size = + plan_workspace_size + actual_seqlen_workspace_size + seqlen_offsets_workspace_size; return; } @@ -420,9 +512,9 @@ void fused_attn_fp8_fwd_impl( constexpr size_t nthreads_per_block = 128; const size_t grid = (b + nthreads_per_block - 1) / nthreads_per_block; void* devActualSeqlenQ = static_cast(workspace) + plan_workspace_size; - void* devActualSeqlenKV = static_cast(devActualSeqlenQ) + b * sizeof(int32_t); + void* devActualSeqlenKV = static_cast(devActualSeqlenQ) + num_bytes_per_seqlen; cu_seqlens_to_actual_seqlens<<>>( - b, b, static_cast(devPtrcuSeqlensQ), // TODO(pass max_b) + actual_b, b, static_cast(devPtrcuSeqlensQ), static_cast(devPtrcuSeqlensKV), static_cast(devActualSeqlenQ), static_cast(devActualSeqlenKV)); NVTE_CHECK_CUDA(cudaGetLastError()); @@ -430,6 +522,49 @@ void fused_attn_fp8_fwd_impl( variant_pack[seq_kv] = devActualSeqlenKV; } + if (is_ragged_q || is_ragged_kv) { + constexpr size_t nthreads_per_block = 128; + const size_t grid = (b + nthreads_per_block) / nthreads_per_block; + void* devOffsets = + static_cast(workspace) + plan_workspace_size + actual_seqlen_workspace_size; + void* devOffsetsQ = nullptr; + void* devOffsetsO = nullptr; + if (is_ragged_q) { + devOffsetsQ = devOffsets; + devOffsetsO = static_cast(devOffsetsQ) + num_bytes_per_ragged_offset; + } + void* devOffsetsK = nullptr; + void* devOffsetsV = nullptr; + if (is_ragged_kv) { + devOffsetsK = static_cast(devOffsets) + + static_cast(is_ragged_q) * 2 * num_bytes_per_ragged_offset; + devOffsetsV = static_cast(devOffsetsK) + num_bytes_per_ragged_offset; + } + void* devOffsetsS = nullptr; + if (use_ragged_stats) { + devOffsetsS = static_cast(devOffsets) + + (static_cast(is_ragged_q) + static_cast(is_ragged_kv)) * 2 * + num_bytes_per_ragged_offset; + } + cu_seqlens_padded_to_offsets<<>>( + layout_group, actual_b, b, h, hg, d_qk, d_v, + static_cast(devPtrSeqOffsetsQ), + static_cast(devPtrSeqOffsetsKV), ragged_offset_type, devOffsetsQ, devOffsetsK, + devOffsetsV, devOffsetsO, devOffsetsS); + NVTE_CHECK_CUDA(cudaGetLastError()); + if (is_ragged_q) { + variant_pack[offset_q] = devOffsetsQ; + variant_pack[offset_o] = devOffsetsO; + } + if (is_ragged_kv) { + variant_pack[offset_k] = devOffsetsK; + variant_pack[offset_v] = devOffsetsV; + } + if (use_ragged_stats) { + variant_pack[offset_stats] = devOffsetsS; + } + } + if (is_dropout) { variant_pack[dropout_seed] = devPtrDropoutSeed; variant_pack[dropout_offset] = devPtrDropoutOffset; @@ -448,6 +583,7 @@ void fused_attn_fp8_fwd_impl( // fused attention BWD FP8 with FE 1.0+ void fused_attn_fp8_bwd_impl( int64_t b, int64_t h, int64_t hg, int64_t s_q, int64_t s_kv, int64_t d_qk, int64_t d_v, + int64_t max_b, int64_t max_t_q, int64_t max_t_kv, float scaling_factor, float dropout_probability, NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, @@ -460,7 +596,8 @@ void fused_attn_fp8_bwd_impl( void* devPtrScaledV, void* devPtrAmaxdP, void* devPtrAmaxdQ, void* devPtrAmaxdK, void* devPtrAmaxdV, void* devPtrQ_t, void* devPtrK_t, void* devPtrdO_f16, void* devPtrdO_t, void* devPtrDescaleQ_t, void* devPtrDescaleK_t, void* devPtrDescaledO_t, void* devPtrcuSeqlensQ, - void* devPtrcuSeqlensKV, void* devPtrDropoutSeed, void* devPtrDropoutOffset, + void* devPtrcuSeqlensKV, void* devPtrSeqOffsetsQ, void* devPtrSeqOffsetsKV, + void* devPtrDropoutSeed, void* devPtrDropoutOffset, cudnn_frontend::DataType_t qkv_tensor_type, cudnn_frontend::DataType_t o_tensor_type, cudnn_frontend::DataType_t do_tensor_type, cudnn_frontend::DataType_t dqkv_tensor_type, NVTEScalingMode scaling_mode, NVTE_QKV_Format qkv_scale_inv_format, @@ -497,6 +634,27 @@ void fused_attn_fp8_bwd_impl( NVTE_CHECK(!is_mxfp8 || cudnn_runtime_version >= 92100, "MXFP8 fused attention requires cuDNN 9.21.0 or later!"); + NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); + NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); + bool is_ragged_q = (q_format == NVTE_QKV_Format::NVTE_THD); + 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; + + 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; + + int64_t actual_b = b; + if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600) { + NVTE_CHECK(is_padding, "Ragged QKV input requires padding or padding_causal mask!"); + if (sm_arch_ != 120) { + b = max_b; + s_q = is_ragged_q ? max_t_q : s_q; + s_kv = is_ragged_kv ? max_t_kv : s_kv; + } + } + bool is_O_in_F16 = (o_tensor_type == cudnn_frontend::DataType_t::HALF || o_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); @@ -582,6 +740,11 @@ void fused_attn_fp8_bwd_impl( std::shared_ptr, // d_softmax_offset std::shared_ptr, // seq_q std::shared_ptr, // seq_kv + std::shared_ptr, // offset_q + std::shared_ptr, // offset_o + std::shared_ptr, // offset_k + std::shared_ptr, // offset_v + std::shared_ptr, // offset_stats std::shared_ptr, // dropout_seed std::shared_ptr>; // dropout_offset @@ -614,6 +777,8 @@ void fused_attn_fp8_bwd_impl( std::shared_ptr scale_dQ, scale_dK, scale_dV; std::shared_ptr bias, dBias, softmax_offset, d_softmax_offset; std::shared_ptr seq_q, seq_kv; + std::shared_ptr offset_q, offset_k, offset_v, offset_o, + offset_stats; std::shared_ptr dropout_seed, dropout_offset; // Q, K, V, O, dO, stats, attn_scale @@ -627,6 +792,19 @@ void fused_attn_fp8_bwd_impl( .set_dim({b, h, s_q, d_qk}) .set_stride(q_strides) .set_data_type(qkv_tensor_type)); + if (is_ragged_q) { + offset_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_q") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + offset_o = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_o") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + Q->set_ragged_offset(offset_q); + } K = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("K") .set_dim({b, hg, s_kv, d_qk}) @@ -637,21 +815,53 @@ void fused_attn_fp8_bwd_impl( .set_dim({b, hg, s_kv, d_v}) .set_stride(v_strides) .set_data_type(qkv_tensor_type)); + if (is_ragged_kv) { + offset_k = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_k") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + offset_v = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_v") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + K->set_ragged_offset(offset_k); + V->set_ragged_offset(offset_v); + } O = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("O") .set_dim({b, h, s_q, d_v}) .set_stride(o_strides) .set_data_type(o_tensor_type)); + if (is_ragged_q) { + O->set_ragged_offset(offset_o); + } dO = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("dO") .set_dim({b, h, s_q, d_v}) .set_stride(dO_strides) .set_data_type(do_tensor_type)); + if (is_ragged_q) { + dO->set_ragged_offset(offset_o); + } + if (use_ragged_stats) { + offset_stats = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_stats") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + } Stats = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("Stats") .set_dim({b, h, s_q, 1}) - .set_stride({h * s_q, s_q, 1, 1}) .set_data_type(fe::DataType_t::FLOAT)); + if (use_ragged_stats) { + Stats->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); + } else { + Stats->set_stride({h * s_q, s_q, 1, 1}); + } attn_scale = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("attn_scale") .set_dim({1, 1, 1, 1}) @@ -901,6 +1111,9 @@ void fused_attn_fp8_bwd_impl( .set_dim({b, h, s_q, d_qk}) .set_stride(dq_strides) .set_data_type(dqkv_tensor_type); + if (is_ragged_q) { + dQ->set_ragged_offset(offset_q); + } dK->set_output(true) .set_dim({b, hg, s_kv, d_qk}) .set_stride(dk_strides) @@ -909,6 +1122,10 @@ void fused_attn_fp8_bwd_impl( .set_dim({b, hg, s_kv, d_v}) .set_stride(dv_strides) .set_data_type(dqkv_tensor_type); + if (is_ragged_kv) { + dK->set_ragged_offset(offset_k); + dV->set_ragged_offset(offset_v); + } amax_dQ->set_output(!is_mxfp8) .set_dim({1, 1, 1, 1}) .set_stride({1, 1, 1, 1}) @@ -967,6 +1184,12 @@ void fused_attn_fp8_bwd_impl( : std::make_tuple(nullptr, nullptr); auto padding_tuple = is_padding ? std::make_tuple(seq_q, seq_kv) : std::make_tuple(nullptr, nullptr); + auto offset_qo_tuple = + is_ragged_q ? std::make_tuple(offset_q, offset_o) : std::make_tuple(nullptr, nullptr); + auto offset_kv_tuple = + is_ragged_kv ? std::make_tuple(offset_k, offset_v) : std::make_tuple(nullptr, nullptr); + auto offset_s_tuple = + use_ragged_stats ? std::make_tuple(offset_stats) : std::make_tuple(nullptr); auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) : std::make_tuple(nullptr, nullptr); @@ -978,7 +1201,8 @@ void fused_attn_fp8_bwd_impl( auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, mxfp8_tensors_tuple, - bias_tuple, softmax_offset_tuple, padding_tuple, dropout_tuple); + bias_tuple, softmax_offset_tuple, padding_tuple, offset_qo_tuple, + offset_kv_tuple, offset_s_tuple, dropout_tuple); cache.insert({descriptor, return_tuple}); return return_tuple; @@ -987,14 +1211,27 @@ void fused_attn_fp8_bwd_impl( descale_dO, descale_s, descale_dP, scale_s, scale_dQ, scale_dK, scale_dV, scale_dP, dQ, dK, dV, amax_dQ, amax_dK, amax_dV, amax_dP, Q_t, K_t, dO_f16, dO_t, descale_q_t, descale_k_t, descale_dO_t, bias, dBias, softmax_offset, d_softmax_offset, seq_q, seq_kv, - dropout_seed, dropout_offset] = get_graph(sdpa_fp8_bprop_cache, descriptor); - - auto plan_workspace_size = mha_graph->get_workspace_size(); + offset_q, offset_o, offset_k, offset_v, offset_stats, dropout_seed, dropout_offset] = + get_graph(sdpa_fp8_bprop_cache, descriptor); + + auto plan_workspace_size = alignTo<16>(mha_graph->get_workspace_size()); + const size_t num_bytes_per_seqlen = alignTo<16>(b * sizeof(int32_t)); + const size_t actual_seqlen_workspace_size = is_padding ? 2 * num_bytes_per_seqlen : 0; + const size_t num_bytes_per_ragged_offset = + alignTo<16>(((b + 1) * typeToNumBits(ragged_offset_type)) / 8); + size_t seqlen_offsets_workspace_size = 0; + if (is_ragged_q || is_ragged_kv) { + size_t count = 2 * (static_cast(is_ragged_q) + static_cast(is_ragged_kv)); + if (use_ragged_stats) { + seqlen_offsets_workspace_size = (count + 1) * num_bytes_per_ragged_offset; + } else { + seqlen_offsets_workspace_size = count * num_bytes_per_ragged_offset; + } + } - // Exit to request upper level API to allocate memory if needed - size_t actual_seqlen_workspace_size = 2 * b * sizeof(int32_t); if (workspace == nullptr) { - *workspace_size = plan_workspace_size + actual_seqlen_workspace_size; + *workspace_size = + plan_workspace_size + actual_seqlen_workspace_size + seqlen_offsets_workspace_size; return; } @@ -1060,9 +1297,9 @@ void fused_attn_fp8_bwd_impl( constexpr size_t nthreads_per_block = 128; const size_t grid = (b + nthreads_per_block - 1) / nthreads_per_block; void* devActualSeqlenQ = static_cast(workspace) + plan_workspace_size; - void* devActualSeqlenKV = static_cast(devActualSeqlenQ) + b * sizeof(int32_t); + void* devActualSeqlenKV = static_cast(devActualSeqlenQ) + num_bytes_per_seqlen; cu_seqlens_to_actual_seqlens<<>>( - b, b, static_cast(devPtrcuSeqlensQ), // TODO(pass max_b) + actual_b, b, static_cast(devPtrcuSeqlensQ), static_cast(devPtrcuSeqlensKV), static_cast(devActualSeqlenQ), static_cast(devActualSeqlenKV)); NVTE_CHECK_CUDA(cudaGetLastError()); @@ -1070,6 +1307,49 @@ void fused_attn_fp8_bwd_impl( variant_pack[seq_kv] = devActualSeqlenKV; } + if (is_ragged_q || is_ragged_kv) { + constexpr size_t nthreads_per_block = 128; + const size_t grid = (b + nthreads_per_block) / nthreads_per_block; + void* devOffsets = + static_cast(workspace) + plan_workspace_size + actual_seqlen_workspace_size; + void* devOffsetsQ = nullptr; + void* devOffsetsO = nullptr; + if (is_ragged_q) { + devOffsetsQ = devOffsets; + devOffsetsO = static_cast(devOffsetsQ) + num_bytes_per_ragged_offset; + } + void* devOffsetsK = nullptr; + void* devOffsetsV = nullptr; + if (is_ragged_kv) { + devOffsetsK = static_cast(devOffsets) + + static_cast(is_ragged_q) * 2 * num_bytes_per_ragged_offset; + devOffsetsV = static_cast(devOffsetsK) + num_bytes_per_ragged_offset; + } + void* devOffsetsS = nullptr; + if (use_ragged_stats) { + devOffsetsS = static_cast(devOffsets) + + (static_cast(is_ragged_q) + static_cast(is_ragged_kv)) * 2 * + num_bytes_per_ragged_offset; + } + cu_seqlens_padded_to_offsets<<>>( + layout_group, actual_b, b, h, hg, d_qk, d_v, + static_cast(devPtrSeqOffsetsQ), + static_cast(devPtrSeqOffsetsKV), ragged_offset_type, devOffsetsQ, devOffsetsK, + devOffsetsV, devOffsetsO, devOffsetsS); + NVTE_CHECK_CUDA(cudaGetLastError()); + if (is_ragged_q) { + variant_pack[offset_q] = devOffsetsQ; + variant_pack[offset_o] = devOffsetsO; + } + if (is_ragged_kv) { + variant_pack[offset_k] = devOffsetsK; + variant_pack[offset_v] = devOffsetsV; + } + if (use_ragged_stats) { + variant_pack[offset_stats] = devOffsetsS; + } + } + if (is_dropout) { variant_pack[dropout_seed] = devPtrDropoutSeed; variant_pack[dropout_offset] = devPtrDropoutOffset; @@ -1091,14 +1371,16 @@ void fused_attn_fp8_bwd_impl( // fused attention FWD FP8 with separate Q, K, V void fused_attn_fp8_fwd( size_t batch, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, - size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, bool is_training, float attn_scale, - float p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, - NVTE_QKV_Format qkv_scale_inv_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, - NVTE_Softmax_Type softmax_type, size_t window_size_left, size_t window_size_right, - bool bottom_right_diagonal, const Tensor* input_Q, const Tensor* input_K, const Tensor* input_V, + size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, size_t num_tokens_q, + size_t num_tokens_kv, bool is_training, float attn_scale, float p_dropout, + NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_QKV_Format qkv_scale_inv_format, + NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, + size_t window_size_left, size_t window_size_right, bool bottom_right_diagonal, + const Tensor* input_Q, const Tensor* input_K, const Tensor* input_V, const Tensor* input_SoftmaxOffset, Tensor* input_output_S, Tensor* output_O, NVTETensorPack* Aux_CTX_Tensors, const Tensor* cu_seqlens_q, const Tensor* cu_seqlens_kv, - const Tensor* rng_state, Tensor* workspace, cudaStream_t stream, cudnnHandle_t handle) { + const Tensor* cu_seqlens_q_padded, const Tensor* cu_seqlens_kv_padded, const Tensor* rng_state, + Tensor* workspace, cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; void *devPtrQ = nullptr, *devPtrK = nullptr, *devPtrV = nullptr; void *devPtrDescaleQ = nullptr, *devPtrDescaleK = nullptr, *devPtrDescaleV = nullptr; @@ -1125,12 +1407,36 @@ void fused_attn_fp8_fwd( if (softmax_type != NVTE_VANILLA_SOFTMAX) { devPtrSoftmaxOffset = input_SoftmaxOffset->data.dptr; } + NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); + NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); + const auto cudnn_runtime_version = cudnnGetVersion(); + + void* devPtrSeqOffsetsQ = cu_seqlens_q_padded->data.dptr; + void* devPtrSeqOffsetsKV = cu_seqlens_kv_padded->data.dptr; + + size_t max_batch_size = 0; + size_t max_tokens_q = 0; + size_t max_tokens_kv = 0; + if (q_format == NVTE_QKV_Format::NVTE_THD || kv_format == NVTE_QKV_Format::NVTE_THD) { + max_batch_size = fused_attn::get_max_batch_size(batch); + } + if (q_format == NVTE_QKV_Format::NVTE_THD) { + max_tokens_q = fused_attn::get_max_tokens(num_tokens_q); + } + if (kv_format == NVTE_QKV_Format::NVTE_THD) { + max_tokens_kv = fused_attn::get_max_tokens(num_tokens_kv); + } + void* devPtrM = nullptr; if (Aux_CTX_Tensors->size == 0) { int i = 0; Tensor* output_M = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_M->data.dptr = nullptr; - output_M->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; + if (q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) { + output_M->data.shape = {num_tokens_q, num_attn_heads, 1}; + } else { + output_M->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; + } output_M->data.dtype = DType::kFloat32; Tensor* output_rng_state = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_rng_state->data.dptr = nullptr; @@ -1172,18 +1478,20 @@ void fused_attn_fp8_fwd( NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); if ((qkv_format == NVTE_QKV_Format::NVTE_BSHD) || (qkv_format == NVTE_QKV_Format::NVTE_SBHD) || - (qkv_format == NVTE_QKV_Format::NVTE_BHSD)) { + (qkv_format == NVTE_QKV_Format::NVTE_BHSD) || (qkv_format == NVTE_QKV_Format::NVTE_THD)) { fused_attn::fused_attn_fp8_fwd_impl( batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, + max_batch_size, max_tokens_q, max_tokens_kv, is_training, attn_scale, p_dropout, qkv_layout, o_format, bias_type, mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, devPtrQ, devPtrK, devPtrV, devPtrSoftmaxOffset, devPtrM, devPtrO, devPtrDescaleQ, devPtrDescaleK, devPtrDescaleV, devPtrDescaleS, devPtrScaleS, devPtrScaleO, devPtrAmaxO, devPtrAmaxS, - devPtrcuSeqlensQ, devPtrcuSeqlensKV, devPtrDropoutSeed, devPtrDropoutOffset, + devPtrcuSeqlensQ, devPtrcuSeqlensKV, devPtrSeqOffsetsQ, devPtrSeqOffsetsKV, + devPtrDropoutSeed, devPtrDropoutOffset, get_cudnn_fe_dtype(QKV_type), get_cudnn_fe_dtype(O_type), input_Q->scaling_mode, qkv_scale_inv_format, workspace->data.dptr, &workspace_size, stream, handle); } else { - NVTE_ERROR("FP8 fused attention only supports qkv_format=BSHD, SBHD, or BHSD.\n"); + NVTE_ERROR("FP8 fused attention only supports qkv_format=BSHD, SBHD, BHSD, or THD.\n"); } if (workspace_size > 0) { @@ -1201,18 +1509,20 @@ void fused_attn_fp8_fwd( // fused attention BWD FP8 with separate Q, K, V void fused_attn_fp8_bwd( size_t batch, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, - size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, float attn_scale, float p_dropout, - NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, - NVTE_QKV_Layout dqkv_layout, NVTE_QKV_Format qkv_scale_inv_format, - NVTE_QKV_Format do_scale_inv_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, - NVTE_Softmax_Type softmax_type, size_t window_size_left, size_t window_size_right, - bool bottom_right_diagonal, bool deterministic, const Tensor* input_Q, const Tensor* input_K, - const Tensor* input_V, const Tensor* input_O, const Tensor* input_dO, - const Tensor* input_dO_f16, const Tensor* input_M, const Tensor* input_S, - const Tensor* input_SoftmaxOffset, Tensor* input_output_dP, const Tensor* output_dQ, - const Tensor* output_dK, const Tensor* output_dV, Tensor* output_dSoftmaxOffset, - const Tensor* cu_seqlens_q, const Tensor* cu_seqlens_kv, const Tensor* rng_state, - Tensor* workspace, cudaStream_t stream, cudnnHandle_t handle) { + size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, size_t num_tokens_q, + size_t num_tokens_kv, float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout, + NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, + NVTE_QKV_Format qkv_scale_inv_format, NVTE_QKV_Format do_scale_inv_format, + NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, + size_t window_size_left, size_t window_size_right, bool bottom_right_diagonal, + bool deterministic, const Tensor* input_Q, const Tensor* input_K, const Tensor* input_V, + const Tensor* input_O, const Tensor* input_dO, const Tensor* input_dO_f16, + const Tensor* input_M, const Tensor* input_S, const Tensor* input_SoftmaxOffset, + Tensor* input_output_dP, const Tensor* output_dQ, const Tensor* output_dK, + const Tensor* output_dV, Tensor* output_dSoftmaxOffset, const Tensor* cu_seqlens_q, + const Tensor* cu_seqlens_kv, const Tensor* cu_seqlens_q_padded, + const Tensor* cu_seqlens_kv_padded, const Tensor* rng_state, Tensor* workspace, + cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; void* devPtrQ = input_Q->data.dptr; void* devPtrK = input_K->data.dptr; @@ -1277,6 +1587,25 @@ void fused_attn_fp8_bwd( devPtrScaledV = output_dV->scale.dptr; } + NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); + NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); + + void* devPtrSeqOffsetsQ = cu_seqlens_q_padded->data.dptr; + void* devPtrSeqOffsetsKV = cu_seqlens_kv_padded->data.dptr; + + size_t max_batch_size = 0; + size_t max_tokens_q = 0; + size_t max_tokens_kv = 0; + if (q_format == NVTE_QKV_Format::NVTE_THD || kv_format == NVTE_QKV_Format::NVTE_THD) { + max_batch_size = fused_attn::get_max_batch_size(batch); + } + if (q_format == NVTE_QKV_Format::NVTE_THD) { + max_tokens_q = fused_attn::get_max_tokens(num_tokens_q); + } + if (kv_format == NVTE_QKV_Format::NVTE_THD) { + max_tokens_kv = fused_attn::get_max_tokens(num_tokens_kv); + } + void* devPtrcuSeqlensQ = reinterpret_cast(reinterpret_cast(cu_seqlens_q->data.dptr)); void* devPtrcuSeqlensKV = @@ -1293,9 +1622,10 @@ void fused_attn_fp8_bwd( NVTE_QKV_Format dqkv_format = nvte_get_qkv_format(dqkv_layout); if ((dqkv_format == NVTE_QKV_Format::NVTE_BSHD) || (dqkv_format == NVTE_QKV_Format::NVTE_SBHD) || - (dqkv_format == NVTE_QKV_Format::NVTE_BHSD)) { + (dqkv_format == NVTE_QKV_Format::NVTE_BHSD) || (dqkv_format == NVTE_QKV_Format::NVTE_THD)) { fused_attn::fused_attn_fp8_bwd_impl( batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, + max_batch_size, max_tokens_q, max_tokens_kv, attn_scale, p_dropout, qkv_layout, o_format, do_format, dqkv_layout, bias_type, mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, deterministic, devPtrQ, devPtrK, devPtrV, devPtrM, devPtrO, devPtrdO, devPtrSoftmaxOffset, devPtrdQ, @@ -1304,12 +1634,13 @@ void fused_attn_fp8_bwd( devPtrScaledP, devPtrScaledQ, devPtrScaledK, devPtrScaledV, devPtrAmaxdP, devPtrAmaxdQ, devPtrAmaxdK, devPtrAmaxdV, devPtrQ_t, devPtrK_t, devPtrdO_f16, devPtrdO_t, devPtrDescaleQ_t, devPtrDescaleK_t, devPtrDescaledO_t, devPtrcuSeqlensQ, devPtrcuSeqlensKV, + devPtrSeqOffsetsQ, devPtrSeqOffsetsKV, devPtrDropoutSeed, devPtrDropoutOffset, get_cudnn_fe_dtype(QKV_type), get_cudnn_fe_dtype(O_type), get_cudnn_fe_dtype(dO_type), get_cudnn_fe_dtype(dQKV_type), input_dO->scaling_mode, qkv_scale_inv_format, do_scale_inv_format, workspace->data.dptr, &workspace_size, stream, handle); } else { - NVTE_ERROR("FP8 fused attention only supports dqkv_format=BSHD, SBHD, or BHSD.\n"); + NVTE_ERROR("FP8 fused attention only supports dqkv_format=BSHD, SBHD, BHSD, or THD.\n"); } if (workspace_size > 0) { diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.h b/transformer_engine/common/fused_attn/fused_attn_fp8.h index b9660128ca..5b3b2fff14 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.h +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.h @@ -15,28 +15,32 @@ namespace transformer_engine { // fused attention FWD FP8 with separate Q, K, V void fused_attn_fp8_fwd( size_t batch, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, - size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, bool is_training, float attn_scale, - float p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, - NVTE_QKV_Format qkv_scale_inv_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, - NVTE_Softmax_Type softmax_type, size_t window_size_left, size_t window_size_right, - bool bottom_right_diagonal, const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, + size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, size_t num_tokens_q, + size_t num_tokens_kv, bool is_training, float attn_scale, float p_dropout, + NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_QKV_Format qkv_scale_inv_format, + NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, + size_t window_size_left, size_t window_size_right, bool bottom_right_diagonal, + const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, const Tensor *input_SoftmaxOffset, Tensor *input_output_S, Tensor *output_O, NVTETensorPack *Aux_CTX_Tensors, const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, - const Tensor *rng_state, Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); + const Tensor *cu_seqlens_q_padded, const Tensor *cu_seqlens_kv_padded, const Tensor *rng_state, + Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); // fused attention BWD FP8 with separate Q, K, V void fused_attn_fp8_bwd( size_t batch, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, - size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, float attn_scale, float p_dropout, - NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, - NVTE_QKV_Layout dqkv_layout, NVTE_QKV_Format qkv_scale_inv_format, - NVTE_QKV_Format do_scale_inv_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, - NVTE_Softmax_Type softmax_type, size_t window_size_left, size_t window_size_right, - bool bottom_right_diagonal, bool deterministic, const Tensor *input_Q, const Tensor *input_K, - const Tensor *input_V, const Tensor *input_O, const Tensor *input_dO, - const Tensor *input_dO_f16, const Tensor *input_M, const Tensor *input_S, - const Tensor *input_SoftmaxOffset, Tensor *input_output_dP, const Tensor *output_dQ, - const Tensor *output_dK, const Tensor *output_dV, Tensor *output_dSoftmaxOffset, - const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, const Tensor *rng_state, - Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); + size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, size_t num_tokens_q, + size_t num_tokens_kv, float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout, + NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, + NVTE_QKV_Format qkv_scale_inv_format, NVTE_QKV_Format do_scale_inv_format, + NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, + size_t window_size_left, size_t window_size_right, bool bottom_right_diagonal, + bool deterministic, const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, + const Tensor *input_O, const Tensor *input_dO, const Tensor *input_dO_f16, + const Tensor *input_M, const Tensor *input_S, const Tensor *input_SoftmaxOffset, + Tensor *input_output_dP, const Tensor *output_dQ, const Tensor *output_dK, + const Tensor *output_dV, Tensor *output_dSoftmaxOffset, const Tensor *cu_seqlens_q, + const Tensor *cu_seqlens_kv, const Tensor *cu_seqlens_q_padded, + const Tensor *cu_seqlens_kv_padded, const Tensor *rng_state, Tensor *workspace, + cudaStream_t stream, cudnnHandle_t handle); } // namespace transformer_engine From bfbafe9dbcb6ace0e1a378dafe93b6ccfdd1a7a9 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 14 May 2026 19:10:21 +0000 Subject: [PATCH 05/20] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../common/fused_attn/fused_attn.cpp | 33 ++++--- .../common/fused_attn/fused_attn_fp8.cu | 86 +++++++++---------- 2 files changed, 56 insertions(+), 63 deletions(-) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index df70bcce66..4742d1bdd7 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -626,14 +626,13 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso input_cu_seqlens_kv, input_cu_seqlens_q_padded, input_cu_seqlens_kv_padded, input_page_table_k, input_page_table_v, input_rng_state, wkspace, stream, handle); } else if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_FP8) { - fused_attn_fp8_fwd(b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, t_q, t_kv, - is_training, attn_scale, dropout, qkv_layout, o_format, - qkv_scale_inv_format, bias_type, attn_mask_type, softmax_type, - window_size_left, window_size_right, bottom_right_diagonal, input_Q, - input_K, input_V, input_SoftmaxOffset, input_output_S, output_O, - Aux_CTX_Tensors, input_cu_seqlens_q, input_cu_seqlens_kv, - input_cu_seqlens_q_padded, input_cu_seqlens_kv_padded, input_rng_state, - wkspace, stream, handle); + fused_attn_fp8_fwd(b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, t_q, t_kv, is_training, + attn_scale, dropout, qkv_layout, o_format, qkv_scale_inv_format, bias_type, + attn_mask_type, softmax_type, window_size_left, window_size_right, + bottom_right_diagonal, input_Q, input_K, input_V, input_SoftmaxOffset, + input_output_S, output_O, Aux_CTX_Tensors, input_cu_seqlens_q, + input_cu_seqlens_kv, input_cu_seqlens_q_padded, input_cu_seqlens_kv_padded, + input_rng_state, wkspace, stream, handle); } else { NVTE_ERROR("Invalid combination of data type and sequence length for fused attention. \n"); } @@ -730,15 +729,15 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso if (input_dO->scaling_mode == NVTE_MXFP8_1D_SCALING) { input_dO_f16 = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); } - fused_attn_fp8_bwd(b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, t_q, t_kv, - attn_scale, dropout, qkv_layout, o_format, do_format, dqkv_layout, - qkv_scale_inv_format, do_scale_inv_format, bias_type, attn_mask_type, - softmax_type, window_size_left, window_size_right, bottom_right_diagonal, - deterministic, input_Q, input_K, input_V, input_O, input_dO, input_dO_f16, - input_M, input_S, input_SoftmaxOffset, input_output_dP, output_dQ, - output_dK, output_dV, output_dSoftmaxOffset, input_cu_seqlens_q, - input_cu_seqlens_kv, input_cu_seqlens_q_padded, input_cu_seqlens_kv_padded, - input_rng_state, wkspace, stream, handle); + fused_attn_fp8_bwd(b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, t_q, t_kv, attn_scale, + dropout, qkv_layout, o_format, do_format, dqkv_layout, qkv_scale_inv_format, + do_scale_inv_format, bias_type, attn_mask_type, softmax_type, + window_size_left, window_size_right, bottom_right_diagonal, deterministic, + input_Q, input_K, input_V, input_O, input_dO, input_dO_f16, input_M, input_S, + input_SoftmaxOffset, input_output_dP, output_dQ, output_dK, output_dV, + output_dSoftmaxOffset, input_cu_seqlens_q, input_cu_seqlens_kv, + input_cu_seqlens_q_padded, input_cu_seqlens_kv_padded, input_rng_state, + wkspace, stream, handle); } else { NVTE_ERROR("Invalid combination of data type and sequence length for fused attention. \n"); } diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 978a6c0035..f29f0d92be 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -19,16 +19,15 @@ using namespace transformer_engine; // fused attention FWD FP8 with FE 1.0+ void fused_attn_fp8_fwd_impl( int64_t b, int64_t h, int64_t hg, int64_t s_q, int64_t s_kv, int64_t d_qk, int64_t d_v, - int64_t max_b, int64_t max_t_q, int64_t max_t_kv, - bool is_training, float scaling_factor, float dropout_probability, NVTE_QKV_Layout qkv_layout, - NVTE_QKV_Format o_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, - NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, void* devPtrQ, void* devPtrK, void* devPtrV, - void* devPtrSoftmaxOffset, void* devPtrM, void* devPtrO, void* devPtrDescaleQ, - void* devPtrDescaleK, void* devPtrDescaleV, void* devPtrDescaleS, void* devPtrScaleS, - void* devPtrScaleO, void* devPtrAmaxO, void* devPtrAmaxS, void* devPtrcuSeqlensQ, - void* devPtrcuSeqlensKV, void* devPtrSeqOffsetsQ, void* devPtrSeqOffsetsKV, - void* devPtrDropoutSeed, void* devPtrDropoutOffset, + int64_t max_b, int64_t max_t_q, int64_t max_t_kv, bool is_training, float scaling_factor, + float dropout_probability, NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, + NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, + int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, void* devPtrQ, + void* devPtrK, void* devPtrV, void* devPtrSoftmaxOffset, void* devPtrM, void* devPtrO, + void* devPtrDescaleQ, void* devPtrDescaleK, void* devPtrDescaleV, void* devPtrDescaleS, + void* devPtrScaleS, void* devPtrScaleO, void* devPtrAmaxO, void* devPtrAmaxS, + void* devPtrcuSeqlensQ, void* devPtrcuSeqlensKV, void* devPtrSeqOffsetsQ, + void* devPtrSeqOffsetsKV, void* devPtrDropoutSeed, void* devPtrDropoutOffset, cudnn_frontend::DataType_t qkv_tensor_type, cudnn_frontend::DataType_t o_tensor_type, NVTEScalingMode scaling_mode, NVTE_QKV_Format qkv_scale_inv_format, void* workspace, size_t* workspace_size, cudaStream_t stream, cudnnHandle_t handle) { @@ -547,8 +546,7 @@ void fused_attn_fp8_fwd_impl( num_bytes_per_ragged_offset; } cu_seqlens_padded_to_offsets<<>>( - layout_group, actual_b, b, h, hg, d_qk, d_v, - static_cast(devPtrSeqOffsetsQ), + layout_group, actual_b, b, h, hg, d_qk, d_v, static_cast(devPtrSeqOffsetsQ), static_cast(devPtrSeqOffsetsKV), ragged_offset_type, devOffsetsQ, devOffsetsK, devOffsetsV, devOffsetsO, devOffsetsS); NVTE_CHECK_CUDA(cudaGetLastError()); @@ -583,13 +581,13 @@ void fused_attn_fp8_fwd_impl( // fused attention BWD FP8 with FE 1.0+ void fused_attn_fp8_bwd_impl( int64_t b, int64_t h, int64_t hg, int64_t s_q, int64_t s_kv, int64_t d_qk, int64_t d_v, - int64_t max_b, int64_t max_t_q, int64_t max_t_kv, - float scaling_factor, float dropout_probability, NVTE_QKV_Layout qkv_layout, - NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, - NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, - int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, - bool deterministic, void* devPtrQ, void* devPtrK, void* devPtrV, void* devPtrM, void* devPtrO, - void* devPtrdO, void* devPtrSoftmaxOffset, void* devPtrdQ, void* devPtrdK, void* devPtrdV, + int64_t max_b, int64_t max_t_q, int64_t max_t_kv, float scaling_factor, + float dropout_probability, NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, + NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, NVTE_Bias_Type bias_type, + NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, + int64_t window_size_right, bool bottom_right_diagonal, bool deterministic, void* devPtrQ, + void* devPtrK, void* devPtrV, void* devPtrM, void* devPtrO, void* devPtrdO, + void* devPtrSoftmaxOffset, void* devPtrdQ, void* devPtrdK, void* devPtrdV, void* devPtrdSoftmaxOffset, void* devPtrDescaleQ, void* devPtrDescaleK, void* devPtrDescaleV, void* devPtrDescaleO, void* devPtrDescaledO, void* devPtrDescaleS, void* devPtrDescaledP, void* devPtrScaleS, void* devPtrScaledP, void* devPtrScaledQ, void* devPtrScaledK, @@ -597,12 +595,11 @@ void fused_attn_fp8_bwd_impl( void* devPtrAmaxdV, void* devPtrQ_t, void* devPtrK_t, void* devPtrdO_f16, void* devPtrdO_t, void* devPtrDescaleQ_t, void* devPtrDescaleK_t, void* devPtrDescaledO_t, void* devPtrcuSeqlensQ, void* devPtrcuSeqlensKV, void* devPtrSeqOffsetsQ, void* devPtrSeqOffsetsKV, - void* devPtrDropoutSeed, void* devPtrDropoutOffset, - cudnn_frontend::DataType_t qkv_tensor_type, cudnn_frontend::DataType_t o_tensor_type, - cudnn_frontend::DataType_t do_tensor_type, cudnn_frontend::DataType_t dqkv_tensor_type, - NVTEScalingMode scaling_mode, NVTE_QKV_Format qkv_scale_inv_format, - NVTE_QKV_Format do_scale_inv_format, void* workspace, size_t* workspace_size, - cudaStream_t stream, cudnnHandle_t handle) { + void* devPtrDropoutSeed, void* devPtrDropoutOffset, cudnn_frontend::DataType_t qkv_tensor_type, + cudnn_frontend::DataType_t o_tensor_type, cudnn_frontend::DataType_t do_tensor_type, + cudnn_frontend::DataType_t dqkv_tensor_type, NVTEScalingMode scaling_mode, + NVTE_QKV_Format qkv_scale_inv_format, NVTE_QKV_Format do_scale_inv_format, void* workspace, + size_t* workspace_size, cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; const auto cudnn_runtime_version = cudnnGetVersion(); bool is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); @@ -1332,8 +1329,7 @@ void fused_attn_fp8_bwd_impl( num_bytes_per_ragged_offset; } cu_seqlens_padded_to_offsets<<>>( - layout_group, actual_b, b, h, hg, d_qk, d_v, - static_cast(devPtrSeqOffsetsQ), + layout_group, actual_b, b, h, hg, d_qk, d_v, static_cast(devPtrSeqOffsetsQ), static_cast(devPtrSeqOffsetsKV), ragged_offset_type, devOffsetsQ, devOffsetsK, devOffsetsV, devOffsetsO, devOffsetsS); NVTE_CHECK_CUDA(cudaGetLastError()); @@ -1481,15 +1477,14 @@ void fused_attn_fp8_fwd( (qkv_format == NVTE_QKV_Format::NVTE_BHSD) || (qkv_format == NVTE_QKV_Format::NVTE_THD)) { fused_attn::fused_attn_fp8_fwd_impl( batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, - max_batch_size, max_tokens_q, max_tokens_kv, - is_training, attn_scale, p_dropout, qkv_layout, o_format, bias_type, mask_type, - softmax_type, window_size_left, window_size_right, bottom_right_diagonal, devPtrQ, devPtrK, - devPtrV, devPtrSoftmaxOffset, devPtrM, devPtrO, devPtrDescaleQ, devPtrDescaleK, - devPtrDescaleV, devPtrDescaleS, devPtrScaleS, devPtrScaleO, devPtrAmaxO, devPtrAmaxS, - devPtrcuSeqlensQ, devPtrcuSeqlensKV, devPtrSeqOffsetsQ, devPtrSeqOffsetsKV, - devPtrDropoutSeed, devPtrDropoutOffset, - get_cudnn_fe_dtype(QKV_type), get_cudnn_fe_dtype(O_type), input_Q->scaling_mode, - qkv_scale_inv_format, workspace->data.dptr, &workspace_size, stream, handle); + max_batch_size, max_tokens_q, max_tokens_kv, is_training, attn_scale, p_dropout, qkv_layout, + o_format, bias_type, mask_type, softmax_type, window_size_left, window_size_right, + bottom_right_diagonal, devPtrQ, devPtrK, devPtrV, devPtrSoftmaxOffset, devPtrM, devPtrO, + devPtrDescaleQ, devPtrDescaleK, devPtrDescaleV, devPtrDescaleS, devPtrScaleS, devPtrScaleO, + devPtrAmaxO, devPtrAmaxS, devPtrcuSeqlensQ, devPtrcuSeqlensKV, devPtrSeqOffsetsQ, + devPtrSeqOffsetsKV, devPtrDropoutSeed, devPtrDropoutOffset, get_cudnn_fe_dtype(QKV_type), + get_cudnn_fe_dtype(O_type), input_Q->scaling_mode, qkv_scale_inv_format, + workspace->data.dptr, &workspace_size, stream, handle); } else { NVTE_ERROR("FP8 fused attention only supports qkv_format=BSHD, SBHD, BHSD, or THD.\n"); } @@ -1625,16 +1620,15 @@ void fused_attn_fp8_bwd( (dqkv_format == NVTE_QKV_Format::NVTE_BHSD) || (dqkv_format == NVTE_QKV_Format::NVTE_THD)) { fused_attn::fused_attn_fp8_bwd_impl( batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, - max_batch_size, max_tokens_q, max_tokens_kv, - attn_scale, p_dropout, qkv_layout, o_format, do_format, dqkv_layout, bias_type, mask_type, - softmax_type, window_size_left, window_size_right, bottom_right_diagonal, deterministic, - devPtrQ, devPtrK, devPtrV, devPtrM, devPtrO, devPtrdO, devPtrSoftmaxOffset, devPtrdQ, - devPtrdK, devPtrdV, devPtrdSoftmaxOffset, devPtrDescaleQ, devPtrDescaleK, devPtrDescaleV, - devPtrDescaleO, devPtrDescaledO, devPtrDescaleS, devPtrDescaledP, devPtrScaleS, - devPtrScaledP, devPtrScaledQ, devPtrScaledK, devPtrScaledV, devPtrAmaxdP, devPtrAmaxdQ, - devPtrAmaxdK, devPtrAmaxdV, devPtrQ_t, devPtrK_t, devPtrdO_f16, devPtrdO_t, - devPtrDescaleQ_t, devPtrDescaleK_t, devPtrDescaledO_t, devPtrcuSeqlensQ, devPtrcuSeqlensKV, - devPtrSeqOffsetsQ, devPtrSeqOffsetsKV, + max_batch_size, max_tokens_q, max_tokens_kv, attn_scale, p_dropout, qkv_layout, o_format, + do_format, dqkv_layout, bias_type, mask_type, softmax_type, window_size_left, + window_size_right, bottom_right_diagonal, deterministic, devPtrQ, devPtrK, devPtrV, devPtrM, + devPtrO, devPtrdO, devPtrSoftmaxOffset, devPtrdQ, devPtrdK, devPtrdV, devPtrdSoftmaxOffset, + devPtrDescaleQ, devPtrDescaleK, devPtrDescaleV, devPtrDescaleO, devPtrDescaledO, + devPtrDescaleS, devPtrDescaledP, devPtrScaleS, devPtrScaledP, devPtrScaledQ, devPtrScaledK, + devPtrScaledV, devPtrAmaxdP, devPtrAmaxdQ, devPtrAmaxdK, devPtrAmaxdV, devPtrQ_t, devPtrK_t, + devPtrdO_f16, devPtrdO_t, devPtrDescaleQ_t, devPtrDescaleK_t, devPtrDescaledO_t, + devPtrcuSeqlensQ, devPtrcuSeqlensKV, devPtrSeqOffsetsQ, devPtrSeqOffsetsKV, devPtrDropoutSeed, devPtrDropoutOffset, get_cudnn_fe_dtype(QKV_type), get_cudnn_fe_dtype(O_type), get_cudnn_fe_dtype(dO_type), get_cudnn_fe_dtype(dQKV_type), input_dO->scaling_mode, qkv_scale_inv_format, do_scale_inv_format, workspace->data.dptr, From 46153a41f0e79ef3129515a9601df306cd221ba1 Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Tue, 9 Jun 2026 22:43:45 -0700 Subject: [PATCH 06/20] Address FP8 THD review feedback 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 --- tests/pytorch/attention/test_attention.py | 4 +- .../common/fused_attn/fused_attn.cpp | 5 ++- .../common/fused_attn/fused_attn_fp8.cu | 45 ++++++++++--------- 3 files changed, 29 insertions(+), 25 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 4dec26af75..dfd963f281 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -2512,8 +2512,6 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: out_grad_shape_new = [*out_grad_shape[:-2], out_grad_shape[-2] * out_grad_shape[-1]] out_grad = torch.randn(out_grad_shape_new, dtype=dtype, device="cuda") - 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): out = dpa( inp[0], @@ -2529,7 +2527,7 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: checkpoint_core_attention=False, core_attention_bias_type=config.attn_bias_type, fp8_output=fp8_dpa, - fast_zero_fill=False, + fast_zero_fill=fp8_dpa and qkv_format == "thd", ) if is_training: out.backward(out_grad) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 4742d1bdd7..8b01ae0223 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -289,8 +289,9 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( (qkv_format == NVTE_QKV_Format::NVTE_BSHD || qkv_format == NVTE_QKV_Format::NVTE_SBHD || qkv_format == NVTE_QKV_Format::NVTE_BHSD)) || (cudnn_runtime_version >= 92300 && sm_arch_ >= 100 && - qkv_format == NVTE_QKV_Format::NVTE_THD && supported_ragged_offset_size)) && - !requires_64bit_ragged_offset && + qkv_format == NVTE_QKV_Format::NVTE_THD && supported_ragged_offset_size && + (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK || + attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK))) && // 9.10.0: known bugs with SDPA FP8 (cudnn_runtime_version != 91000) && !return_max_logit) { backend = NVTE_Fused_Attn_Backend::NVTE_FP8; diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index f29f0d92be..73da2df80c 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -16,6 +16,8 @@ namespace fused_attn { using namespace transformer_engine; +constexpr size_t kFP8THDRaggedCudnnVersion = 92300; + // fused attention FWD FP8 with FE 1.0+ void fused_attn_fp8_fwd_impl( int64_t b, int64_t h, int64_t hg, int64_t s_q, int64_t s_kv, int64_t d_qk, int64_t d_v, @@ -68,13 +70,14 @@ void fused_attn_fp8_fwd_impl( 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; + bool use_ragged_stats = + is_ragged_q && cudnn_runtime_version >= kFP8THDRaggedCudnnVersion && 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; + const DType ragged_offset_type = DType::kInt64; int64_t actual_b = b; - if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600) { + if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= kFP8THDRaggedCudnnVersion) { NVTE_CHECK(is_padding, "Ragged QKV input requires padding or padding_causal mask!"); if (sm_arch_ != 120) { b = max_b; @@ -145,9 +148,9 @@ void fused_attn_fp8_fwd_impl( std::shared_ptr, // seq_q std::shared_ptr, // seq_kv std::shared_ptr, // offset_q - std::shared_ptr, // offset_o std::shared_ptr, // offset_k std::shared_ptr, // offset_v + std::shared_ptr, // offset_o std::shared_ptr, // offset_stats std::shared_ptr, // dropout_seed std::shared_ptr>; // dropout_offset @@ -427,10 +430,10 @@ void fused_attn_fp8_fwd_impl( is_softmax_offset ? std::make_tuple(softmax_offset) : std::make_tuple(nullptr); auto padding_tuple = is_padding ? std::make_tuple(seq_q, seq_kv) : std::make_tuple(nullptr, nullptr); - auto offset_qo_tuple = - is_ragged_q ? std::make_tuple(offset_q, offset_o) : std::make_tuple(nullptr, nullptr); + 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); auto offset_s_tuple = use_ragged_stats ? std::make_tuple(offset_stats) : std::make_tuple(nullptr); auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) @@ -443,8 +446,8 @@ void fused_attn_fp8_fwd_impl( NVTE_CHECK_CUDNN_FE(mha_graph->build_plans(handle)); auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, Stats_tuple, bias_tuple, - softmax_offset_tuple, padding_tuple, offset_qo_tuple, offset_kv_tuple, - offset_s_tuple, dropout_tuple); + softmax_offset_tuple, padding_tuple, offset_q_tuple, offset_kv_tuple, + offset_o_tuple, offset_s_tuple, dropout_tuple); cache.insert({descriptor, return_tuple}); return return_tuple; @@ -452,7 +455,7 @@ void fused_attn_fp8_fwd_impl( auto [mha_graph, Q, K, V, descale_q, descale_k, descale_v, descale_s, scale_s, scale_o, 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] = + offset_k, offset_v, offset_o, offset_stats, dropout_seed, dropout_offset] = get_graph(sdpa_fp8_fprop_cache, descriptor); auto plan_workspace_size = alignTo<16>(mha_graph->get_workspace_size()); @@ -576,7 +579,7 @@ void fused_attn_fp8_fwd_impl( } catch (cudnn_frontend::cudnnException& e) { NVTE_ERROR(e.what()); } -} +} // NOLINT(readability/fn_size) // fused attention BWD FP8 with FE 1.0+ void fused_attn_fp8_bwd_impl( @@ -637,13 +640,14 @@ void fused_attn_fp8_bwd_impl( 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; + bool use_ragged_stats = + is_ragged_q && cudnn_runtime_version >= kFP8THDRaggedCudnnVersion && 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; + const DType ragged_offset_type = DType::kInt64; int64_t actual_b = b; - if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600) { + if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= kFP8THDRaggedCudnnVersion) { NVTE_CHECK(is_padding, "Ragged QKV input requires padding or padding_causal mask!"); if (sm_arch_ != 120) { b = max_b; @@ -738,9 +742,9 @@ void fused_attn_fp8_bwd_impl( std::shared_ptr, // seq_q std::shared_ptr, // seq_kv std::shared_ptr, // offset_q - std::shared_ptr, // offset_o std::shared_ptr, // offset_k std::shared_ptr, // offset_v + std::shared_ptr, // offset_o std::shared_ptr, // offset_stats std::shared_ptr, // dropout_seed std::shared_ptr>; // dropout_offset @@ -1181,10 +1185,10 @@ void fused_attn_fp8_bwd_impl( : std::make_tuple(nullptr, nullptr); auto padding_tuple = is_padding ? std::make_tuple(seq_q, seq_kv) : std::make_tuple(nullptr, nullptr); - auto offset_qo_tuple = - is_ragged_q ? std::make_tuple(offset_q, offset_o) : std::make_tuple(nullptr, nullptr); + 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); auto offset_s_tuple = use_ragged_stats ? std::make_tuple(offset_stats) : std::make_tuple(nullptr); auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) @@ -1198,8 +1202,8 @@ void fused_attn_fp8_bwd_impl( auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, mxfp8_tensors_tuple, - bias_tuple, softmax_offset_tuple, padding_tuple, offset_qo_tuple, - offset_kv_tuple, offset_s_tuple, dropout_tuple); + bias_tuple, softmax_offset_tuple, padding_tuple, offset_q_tuple, + offset_kv_tuple, offset_o_tuple, offset_s_tuple, dropout_tuple); cache.insert({descriptor, return_tuple}); return return_tuple; @@ -1208,7 +1212,7 @@ void fused_attn_fp8_bwd_impl( descale_dO, descale_s, descale_dP, scale_s, scale_dQ, scale_dK, scale_dV, scale_dP, dQ, dK, dV, amax_dQ, amax_dK, amax_dV, amax_dP, Q_t, K_t, dO_f16, dO_t, descale_q_t, descale_k_t, descale_dO_t, bias, dBias, softmax_offset, d_softmax_offset, seq_q, seq_kv, - offset_q, offset_o, offset_k, offset_v, offset_stats, dropout_seed, dropout_offset] = + offset_q, offset_k, offset_v, offset_o, offset_stats, dropout_seed, dropout_offset] = get_graph(sdpa_fp8_bprop_cache, descriptor); auto plan_workspace_size = alignTo<16>(mha_graph->get_workspace_size()); @@ -1428,7 +1432,8 @@ void fused_attn_fp8_fwd( int i = 0; Tensor* output_M = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_M->data.dptr = nullptr; - if (q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) { + if (q_format == NVTE_QKV_Format::NVTE_THD && + cudnn_runtime_version >= kFP8THDRaggedCudnnVersion) { output_M->data.shape = {num_tokens_q, num_attn_heads, 1}; } else { output_M->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; From 14a951827ce410cfda55e240a37992d3af97f5b3 Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Tue, 9 Jun 2026 23:49:19 -0700 Subject: [PATCH 07/20] Fix FP8 THD ragged version scope Signed-off-by: Sudhakar Singh --- transformer_engine/common/fused_attn/fused_attn_fp8.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 73da2df80c..58d5d56389 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -1433,7 +1433,7 @@ void fused_attn_fp8_fwd( Tensor* output_M = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_M->data.dptr = nullptr; if (q_format == NVTE_QKV_Format::NVTE_THD && - cudnn_runtime_version >= kFP8THDRaggedCudnnVersion) { + cudnn_runtime_version >= fused_attn::kFP8THDRaggedCudnnVersion) { output_M->data.shape = {num_tokens_q, num_attn_heads, 1}; } else { output_M->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; From 01144a02b2e850a7a12c2656359c9ba5db3d3621 Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Wed, 22 Jul 2026 16:23:02 -0700 Subject: [PATCH 08/20] Fix FP8 ragged offset kernel calls after main merge 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 --- transformer_engine/common/fused_attn/fused_attn_fp8.cu | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 55d12f6254..d6262fc60f 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -594,8 +594,9 @@ void fused_attn_fp8_fwd_impl( (static_cast(is_ragged_q) + static_cast(is_ragged_kv)) * 2 * num_bytes_per_ragged_offset; } + const RaggedOffsetMultipliers offset_mults(layout_group, h, hg, d_qk, d_v); cu_seqlens_padded_to_offsets<<>>( - layout_group, actual_b, b, h, hg, d_qk, d_v, static_cast(devPtrSeqOffsetsQ), + offset_mults, actual_b, b, static_cast(devPtrSeqOffsetsQ), static_cast(devPtrSeqOffsetsKV), ragged_offset_type, devOffsetsQ, devOffsetsK, devOffsetsV, devOffsetsO, devOffsetsS); NVTE_CHECK_CUDA(cudaGetLastError()); @@ -1378,8 +1379,9 @@ void fused_attn_fp8_bwd_impl( (static_cast(is_ragged_q) + static_cast(is_ragged_kv)) * 2 * num_bytes_per_ragged_offset; } + const RaggedOffsetMultipliers offset_mults(layout_group, h, hg, d_qk, d_v); cu_seqlens_padded_to_offsets<<>>( - layout_group, actual_b, b, h, hg, d_qk, d_v, static_cast(devPtrSeqOffsetsQ), + offset_mults, actual_b, b, static_cast(devPtrSeqOffsetsQ), static_cast(devPtrSeqOffsetsKV), ragged_offset_type, devOffsetsQ, devOffsetsK, devOffsetsV, devOffsetsO, devOffsetsS); NVTE_CHECK_CUDA(cudaGetLastError()); From a4612c79e38e471024310d7a66d4cf85adcfcae2 Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Wed, 22 Jul 2026 18:39:57 -0700 Subject: [PATCH 09/20] Fix FP8 direct-seqlen merge integration 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 --- .../common/fused_attn/fused_attn_fp8.cu | 44 ++++++++++--------- 1 file changed, 23 insertions(+), 21 deletions(-) diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index d6262fc60f..18f963e69e 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -76,20 +76,9 @@ void fused_attn_fp8_fwd_impl( NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); const DType ragged_offset_type = DType::kInt64; - int64_t actual_b = b; - if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= kFP8THDRaggedCudnnVersion) { - NVTE_CHECK(is_padding, "Ragged QKV input requires padding or padding_causal mask!"); - if (sm_arch_ != 120) { - b = max_b; - s_q = is_ragged_q ? max_t_q : s_q; - s_kv = is_ragged_kv ? max_t_kv : s_kv; - } - } // Newer versions of cuDNN SDPA can accept sequence lengths directly as a cumulative - // tensor. Take advantage of this if possible to avoid 1 extra kernel call. (Unlike - // the F16 path, the FP8 path has no THD/ragged-offset support, so only the - // cu_seqlens_to_actual_seqlens conversion applies here. Also note that the - // needed versions of cuDNN backend and frontend are higher than for F16.) + // tensor. Take advantage of this if possible to avoid the actual-seqlen conversion; + // THD inputs still use their separate ragged-offset tensors. const bool use_cu_seqlens_directly = // Frontend 1.26 supports fp8+cu_seqlens (for the C++ API). // Note: For the Python API, 1.27 is required. @@ -103,6 +92,20 @@ void fused_attn_fp8_fwd_impl( // (which doesn't support cu_seqlens). Remove this restriction when possible. !is_dropout; + int64_t actual_b = b; + if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= kFP8THDRaggedCudnnVersion) { + NVTE_CHECK(is_padding, "Ragged QKV input requires padding or padding_causal mask!"); + if (sm_arch_ != 120) { + // cuDNN reads the user's [actual_b+1] cu_seqlens buffers directly, so a quantized + // batch dimension would read out of bounds on the direct path. + if (!use_cu_seqlens_directly) { + b = max_b; + } + s_q = is_ragged_q ? max_t_q : s_q; + s_kv = is_ragged_kv ? max_t_kv : s_kv; + } + } + try { FADescriptor_v1 descriptor{b, h, @@ -498,7 +501,8 @@ void fused_attn_fp8_fwd_impl( auto plan_workspace_size = alignTo<16>(mha_graph->get_workspace_size()); const size_t num_bytes_per_seqlen = alignTo<16>(b * sizeof(int32_t)); - const size_t actual_seqlen_workspace_size = is_padding ? 2 * num_bytes_per_seqlen : 0; + const size_t actual_seqlen_workspace_size = + (is_padding && !use_cu_seqlens_directly) ? 2 * num_bytes_per_seqlen : 0; const size_t num_bytes_per_ragged_offset = alignTo<16>(((b + 1) * typeToNumBits(ragged_offset_type)) / 8); size_t seqlen_offsets_workspace_size = 0; @@ -511,10 +515,6 @@ void fused_attn_fp8_fwd_impl( } } - // Exit to request upper level API to allocate memory if needed. - // When passing cu_seqlens* directly to cuDNN SDPA, no conversion workspace is - // needed: cuDNN consumes the user's cu_seqlens buffers as-is. - size_t actual_seqlen_workspace_size = use_cu_seqlens_directly ? 0 : 2 * b * sizeof(int32_t); if (workspace == nullptr) { *workspace_size = plan_workspace_size + actual_seqlen_workspace_size + seqlen_offsets_workspace_size; @@ -559,9 +559,9 @@ void fused_attn_fp8_fwd_impl( constexpr size_t nthreads_per_block = 128; const size_t grid = (b + nthreads_per_block - 1) / nthreads_per_block; void* devActualSeqlenQ = static_cast(workspace) + plan_workspace_size; - void* devActualSeqlenKV = static_cast(devActualSeqlenQ) + b * sizeof(int32_t); + void* devActualSeqlenKV = static_cast(devActualSeqlenQ) + num_bytes_per_seqlen; cu_seqlens_to_actual_seqlens<<>>( - b, b, static_cast(devPtrcuSeqlensQ), // TODO(pass max_b) + actual_b, b, static_cast(devPtrcuSeqlensQ), static_cast(devPtrcuSeqlensKV), static_cast(devActualSeqlenQ), static_cast(devActualSeqlenKV)); NVTE_CHECK_CUDA(cudaGetLastError()); @@ -1458,6 +1458,7 @@ void fused_attn_fp8_fwd( NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); const auto cudnn_runtime_version = cudnnGetVersion(); + const int sm_arch_ = cuda::sm_arch(cuda::current_device()); void* devPtrSeqOffsetsQ = cu_seqlens_q_padded->data.dptr; void* devPtrSeqOffsetsKV = cu_seqlens_kv_padded->data.dptr; @@ -1480,8 +1481,9 @@ void fused_attn_fp8_fwd( int i = 0; Tensor* output_M = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_M->data.dptr = nullptr; + // SM120 uses dense stats in the graph, so its allocation must remain [b, h, s_q, 1]. if (q_format == NVTE_QKV_Format::NVTE_THD && - cudnn_runtime_version >= fused_attn::kFP8THDRaggedCudnnVersion) { + cudnn_runtime_version >= fused_attn::kFP8THDRaggedCudnnVersion && sm_arch_ != 120) { output_M->data.shape = {num_tokens_q, num_attn_heads, 1}; } else { output_M->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; From 7d785e811ad5aab671eeb804ee32ae90186e0088 Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Fri, 24 Jul 2026 11:24:38 -0700 Subject: [PATCH 10/20] Disable fast zero fill in FP8 THD comparison 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 --- tests/pytorch/attention/test_attention.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 341c99346d..f82cd73ba9 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -2526,7 +2526,8 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: checkpoint_core_attention=False, core_attention_bias_type=config.attn_bias_type, fp8_output=fp8_dpa, - fast_zero_fill=fp8_dpa and qkv_format == "thd", + # The optimized zero-fill path dereferences device memory on the host. + fast_zero_fill=False, ) if is_training: out.backward(out_grad) From 6df2aaa84a74d2d6c9ec5810ab6177a5466bd2b1 Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Thu, 30 Jul 2026 14:05:24 -0700 Subject: [PATCH 11/20] Align FP8 THD MHA test token counts 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 --- tests/pytorch/attention/test_attention.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index d68fae8987..1febc546b7 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -2249,6 +2249,10 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: seqlens_kv = torch.full( [config.batch_size], config.max_seqlen_kv, dtype=torch.int32, device="cuda" ) + if qkv_format == "thd": + # FP8 Linear flattens THD input to [t, h*d], so align total tokens for cuBLAS. + seqlens_q[-1] += -seqlens_q.sum() % 8 + seqlens_kv[-1] += -seqlens_kv.sum() % 8 cu_seqlens_q = torch.zeros(config.batch_size + 1, dtype=torch.int32, device="cuda") cu_seqlens_kv = torch.zeros(config.batch_size + 1, dtype=torch.int32, device="cuda") cu_seqlens_q[1:] = torch.cumsum(seqlens_q, dim=0) @@ -2287,6 +2291,8 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: rotary_pos_emb=rotary_pos_emb, cu_seqlens_q=cu_seqlens_q, cu_seqlens_kv=cu_seqlens_kv, + # The optimized zero-fill path dereferences device memory on the host. + fast_zero_fill=False, ) if is_training: out.backward(out_grad) From 579cd2e636a1ff836c0a651fb5095e0e08087899 Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Fri, 31 Jul 2026 16:49:42 -0700 Subject: [PATCH 12/20] Enable FP8 THD context parallel attention 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 --- .../attention/test_attention_with_cp.py | 7 ++-- .../dot_product_attention/context_parallel.py | 12 +++++-- .../attention/dot_product_attention/utils.py | 6 ---- .../pytorch/csrc/extensions/attention.cpp | 32 ++++--------------- 4 files changed, 22 insertions(+), 35 deletions(-) diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index 681ee5e6e0..c3b88c74a6 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -407,6 +407,8 @@ def test_cp_with_flash_attention(cp_pool, dtype, model, qkv_format, cp_comm_type 2, 4096, 12, 128, attn_bias_type="post_scale_bias", bias_shape="bhss" ), # MHA "cp_1_5": ModelConfig(2, 4096, 12, 128, attn_mask_type="causal", window_size=(512, 512)), # MHA + # Noncausal MHA without bias/max-logit provides an FP8+THD+CP backend-compatible row. + "cp_1_6": ModelConfig(2, 4096, 12, 128), "cp_2_0": ModelConfig( 2, 4096, @@ -489,6 +491,7 @@ def test_cp_with_flash_attention(cp_pool, dtype, model, qkv_format, cp_comm_type if test_essential: configs = [ "cp_1_0", + "cp_1_6", "cp_2_0", "cp_2_1", "cp_2_2", @@ -548,8 +551,6 @@ def test_cp_with_fused_attention( if dtype != "fp8" and (fp8_mha or fp8_dpa): pytest.skip("dtype!=fp8 requires fp8_dpa=False and fp8_mha=False!") - if dtype == "fp8" and qkv_format == "thd": - pytest.skip("No support for FP8 attention with THD format!") if dtype == "fp8" and config.attn_bias_type != "no_bias": pytest.skip("No support for FP8 attention with bias!") @@ -597,6 +598,8 @@ def test_cp_with_fused_attention( pytest.skip("scaling_mode=delayed requires f16_O=False!") if scaling_mode == "mxfp8" and not f16_O: pytest.skip("scaling_mode=mxfp8 requires f16_O=True!") + if scaling_mode == "mxfp8" and qkv_format == "thd": + pytest.skip("MXFP8 quantization does not support THD format!") if scaling_mode == "mxfp8" and fp8_mha: pytest.skip("No support for scaling_mode=mxfp8 with fp8_mha=True!") diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 995dd2e90a..91cdeb7d29 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -3235,7 +3235,7 @@ def forward( # is large enough to outlast cp_stream's launch (e.g. bucket128k @ cp=8). cp_stream.wait_stream(torch.cuda.current_stream()) - # THD all_gather only reaches this path for f16/bf16 attention today. + # Shapes before per-step slicing and FP8 metadata wrapping. # q: [b, 2, s//2, h, d] or [2, s//2, b, h, d] # k: [s, b, h, d] # v: [s, b, h, d] @@ -3403,6 +3403,13 @@ def forward( ) 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(): + q_part, k_part, v_part = [ + Float8Tensor.make_like(x, data=y, dtype=fwd_nominal_dtype) + for x, y in zip( + [q_fp8, k_fp8, v_fp8], [q_part, k_part, v_part] + ) + ] if use_fused_attention: # Set per-step parameters for THD vs bshd/sbhd if qkv_format == "thd": @@ -3735,7 +3742,8 @@ def backward(ctx, dout, *_args): # v: [s, b, h, d] if ctx.fp8 and not ctx.fp8_recipe.mxfp8(): q, k, v = [x._data for x in [q_fp8, k_fp8, v_fp8]] - if not ctx.qkv_reshaped: + # BSHD/SBHD split the sequence into two chunks; THD stays token-major [t, h, d]. + if not ctx.qkv_reshaped and ctx.qkv_format != "thd": q = q.view( *q.shape[:seq_dim_qkv], 2, q.shape[seq_dim_qkv] // 2, *q.shape[(seq_dim_qkv + 1) :] ) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 8a4c6ce286..2cbb136d56 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -1212,12 +1212,6 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt " bias for THD format" ) use_fused_attention = False - elif fp8 and fp8_meta["recipe"].fp8_dpa and qkv_format == "thd": - logger.debug( - "Disabling FusedAttention as it does not support context parallelism with FP8" - " attention and THD format" - ) - use_fused_attention = False elif fp8 and fp8_meta["recipe"].fp8_dpa and core_attention_bias_type != "no_bias": logger.debug( "Disabling FusedAttention as it does not support context parallelism with FP8" diff --git a/transformer_engine/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index eb8813d4a0..3b95202eaa 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -12,27 +12,9 @@ namespace { constexpr int block_size = 512; -// fast zero-fills of tensors -void mha_fill(const transformer_engine::TensorWrapper &self, const at::Tensor &start_index) { - std::vector shape = transformer_engine::pytorch::convertShape(self.shape()); - - auto max_tokens = shape[0]; - auto fcd_size = 1; - for (size_t i = 1; i <= shape.size(); i++) { - fcd_size *= shape[i]; - } - - NVTE_CHECK(fcd_size % block_size == 0, "input size not aligned to block size"); - - size_t element_size_bits = transformer_engine::pytorch::typeToNumBits(self.dtype()); - int32_t start_row = start_index.data_ptr()[0]; - void *base_ptr = static_cast(self.get_rowwise_data().data_ptr) + - static_cast(start_row) * fcd_size * element_size_bits / 8; - size_t num_rows_to_zero = max_tokens - start_row; - size_t total_bytes = num_rows_to_zero * fcd_size * element_size_bits / 8; - - 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) { + self.zero_(at::cuda::getCurrentCUDAStream()); } } // namespace @@ -165,7 +147,7 @@ std::vector fused_attn_fwd( // FP8 if (set_zero && (o_format == NVTE_QKV_Format::NVTE_THD)) { if ((h * d) % block_size == 0) { - mha_fill(te_O, cu_seqlens_q.index({torch::indexing::Slice(-1, torch::indexing::None)})); + mha_fill(te_O); } else { te_O.zero_(at::cuda::getCurrentCUDAStream()); } @@ -470,7 +452,7 @@ std::vector fused_attn_bwd( if (set_zero) { if (dq_format == NVTE_QKV_Format::NVTE_THD) { if (((h_q * d_qk) % block_size == 0) && dQ.is_contiguous()) { - mha_fill(te_dQ, cu_seqlens_q.index({torch::indexing::Slice(-1, torch::indexing::None)})); + mha_fill(te_dQ); } else { dQ.fill_(0); } @@ -478,8 +460,8 @@ std::vector fused_attn_bwd( if (dkv_format == NVTE_QKV_Format::NVTE_THD) { if (((h_kv * d_qk) % block_size == 0) && ((h_kv * d_v) % block_size == 0) && dK.is_contiguous() && dV.is_contiguous()) { - mha_fill(te_dK, cu_seqlens_kv.index({torch::indexing::Slice(-1, torch::indexing::None)})); - mha_fill(te_dV, cu_seqlens_kv.index({torch::indexing::Slice(-1, torch::indexing::None)})); + mha_fill(te_dK); + mha_fill(te_dV); } else { dK.fill_(0); dV.fill_(0); From 3a80b910c1847dfd400c0552a0c0c2eed138f70f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 31 Jul 2026 23:51:03 +0000 Subject: [PATCH 13/20] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../attention/dot_product_attention/context_parallel.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 91cdeb7d29..981100b927 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -3406,9 +3406,7 @@ def forward( if fp8 and not fp8_recipe.mxfp8(): q_part, k_part, v_part = [ Float8Tensor.make_like(x, data=y, dtype=fwd_nominal_dtype) - for x, y in zip( - [q_fp8, k_fp8, v_fp8], [q_part, k_part, v_part] - ) + for x, y in zip([q_fp8, k_fp8, v_fp8], [q_part, k_part, v_part]) ] if use_fused_attention: # Set per-step parameters for THD vs bshd/sbhd From 3bc700bf5736f35e04f0385d8200385552ca3460 Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Mon, 10 Aug 2026 16:23:27 -0700 Subject: [PATCH 14/20] Fix FP8 THD CP gradient padding 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 --- .../dot_product_attention/context_parallel.py | 37 +++++++++++-------- 1 file changed, 21 insertions(+), 16 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 981100b927..ace6ded43a 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -93,6 +93,22 @@ def get_bsh_dims(tensor_format): return batch_dim, seq_dim, head_dim +def _zero_thd_padding(tensor, cu_seqlens, cu_seqlens_padded): + """Zero inter-sequence padding without copying CUDA cu-seqlens to the host.""" + if tensor is None or cu_seqlens is None or cu_seqlens_padded is None: + return + rows = torch.arange(tensor.shape[0], device=tensor.device) + padding_mask = torch.zeros(tensor.shape[0], dtype=torch.bool, device=tensor.device) + for batch_idx in range(cu_seqlens.numel() - 1): + valid_end = ( + cu_seqlens_padded[batch_idx] + + cu_seqlens[batch_idx + 1] + - cu_seqlens[batch_idx] + ) + padding_mask |= (rows >= valid_end) & (rows < cu_seqlens_padded[batch_idx + 1]) + tensor[padding_mask] = 0 + + def flash_attn_p2p_communicate( rank, send_tensor, send_dst, recv_tensor, recv_src, cp_group, batch_p2p_comm ): @@ -2923,22 +2939,11 @@ def backward(ctx, dout, *_args): nvtx_range_pop(f"{nvtx_label}") - # Zero-fill dQ/dK/dV at positions beyond the actual sequence end (THD CUDA Graph). - # cu_seqlens_*_padded are already local to this CP rank in the THD path. - # Use Q's padded boundary for dQ and KV's padded boundary for dK/dV. - # Skip the corresponding zero-fill when its padded cu_seqlens is absent. - if ctx.qkv_format == "thd": - if cu_seqlens_q_padded is not None and isinstance(dq, torch.Tensor) and dq.shape[0] > 0: - q_pad_mask = torch.arange(dq.shape[0], device=dq.device) >= cu_seqlens_q_padded[-1] - dq[q_pad_mask] = 0 - if cu_seqlens_kv_padded is not None: - kv_actual_t = cu_seqlens_kv_padded[-1] - for d_tensor in [dk, dv]: - if isinstance(d_tensor, torch.Tensor) and d_tensor.shape[0] > 0: - kv_pad_mask = ( - torch.arange(d_tensor.shape[0], device=d_tensor.device) >= kv_actual_t - ) - d_tensor[kv_pad_mask] = 0 + # FP8 partial-gradient reduction can write THD inter-sequence padding. + if ctx.qkv_format == "thd" and ctx.fp8: + _zero_thd_padding(dq, cu_seqlens_q_per_step[0], cu_seqlens_q_padded) + _zero_thd_padding(dk, cu_seqlens_kv_per_step[0], cu_seqlens_kv_padded) + _zero_thd_padding(dv, cu_seqlens_kv_per_step[0], cu_seqlens_kv_padded) return ( None, From 0a41f7668c08179b9da03a82241dbb0f984a38bf Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 10 Aug 2026 23:24:43 +0000 Subject: [PATCH 15/20] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../attention/dot_product_attention/context_parallel.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index ace6ded43a..ed0881b53c 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -100,11 +100,7 @@ def _zero_thd_padding(tensor, cu_seqlens, cu_seqlens_padded): rows = torch.arange(tensor.shape[0], device=tensor.device) padding_mask = torch.zeros(tensor.shape[0], dtype=torch.bool, device=tensor.device) for batch_idx in range(cu_seqlens.numel() - 1): - valid_end = ( - cu_seqlens_padded[batch_idx] - + cu_seqlens[batch_idx + 1] - - cu_seqlens[batch_idx] - ) + valid_end = cu_seqlens_padded[batch_idx] + cu_seqlens[batch_idx + 1] - cu_seqlens[batch_idx] padding_mask |= (rows >= valid_end) & (rows < cu_seqlens_padded[batch_idx + 1]) tensor[padding_mask] = 0 From 038607931b0ee09dd7b9be85442101f505dc27f5 Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Mon, 10 Aug 2026 21:24:13 -0700 Subject: [PATCH 16/20] Enable FP8 THD inference on Hopper with cuDNN 9.25 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 --- transformer_engine/common/fused_attn/fused_attn.cpp | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index d983758e5f..93747ea14d 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -283,14 +283,15 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK))) && // pre-9.21: {bshd, sbhd}, {vanilla} // 9.21+: {bshd, sbhd, bhsd}, {vanilla, off-by-one, learnable} - // 9.23+ sm100+: {thd} (ragged/packed variable-length) + // 9.23+ sm100+: {thd}; 9.25+ sm90+: {thd} inference-only ((cudnn_runtime_version < 92100 && (qkv_format == NVTE_QKV_Format::NVTE_BSHD || qkv_format == NVTE_QKV_Format::NVTE_SBHD) && softmax_type == NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX) || (cudnn_runtime_version >= 92100 && (qkv_format == NVTE_QKV_Format::NVTE_BSHD || qkv_format == NVTE_QKV_Format::NVTE_SBHD || qkv_format == NVTE_QKV_Format::NVTE_BHSD)) || - (cudnn_runtime_version >= 92300 && sm_arch_ >= 100 && + (((cudnn_runtime_version >= 92300 && sm_arch_ >= 100) || + (cudnn_runtime_version >= 92500 && !is_training)) && qkv_format == NVTE_QKV_Format::NVTE_THD && supported_ragged_offset_size && (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK || attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK))) && From 02dbd4a83532f59a0b10ea492085ff5df1adfcea Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Mon, 10 Aug 2026 21:44:26 -0700 Subject: [PATCH 17/20] Cover FP8 THD CP inference on Hopper 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 --- .../attention/test_attention_with_cp.py | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index 0734ae8b73..70419b226a 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -533,6 +533,20 @@ def test_cp_with_fused_attention( config.context_parallel = True config.cp_comm_type = cp_comm_type + hopper_fp8_thd_forward = ( + get_device_compute_capability() == (9, 0) + and dtype == "fp8" + and model == "cp_1_6" + and qkv_format == "thd" + and not fp8_bwd + and fp8_dpa + and not fp8_mha + and scaling_mode == "delayed" + and not f16_O + ) + if hopper_fp8_thd_forward and get_cudnn_version() < (9, 25, 0): + pytest.skip("FP8+THD inference on Hopper requires cuDNN 9.25+.") + if config.head_dim_qk == 256 and config.head_dim_v == 256: # D=256 uses this generic CP runner, but only a subset of its axes is supported. if get_device_compute_capability() not in ((10, 0), (10, 3)): @@ -641,8 +655,9 @@ def test_cp_with_fused_attention( MXFP8BlockScaling(fp8_format=Format.E4M3, fp8_dpa=True), ] - # For 111s, dbias calculation is not supported as of cuDNN 9.18, hence, test fwd only for 111s. - is_training = False if config.bias_shape == "111s" else True + # 111s runs forward-only because its dbias is unsupported. Reuse otherwise-skipped Hopper + # FP8+THD nodes for forward-only coverage of the supported cuDNN path. + is_training = False if config.bias_shape == "111s" or hopper_fp8_thd_forward else True available_backends, _, fused_attn_backends = get_available_attention_backends( config, qkv_dtype=dtypes[dtype] if dtype != "fp8" else torch.float8_e4m3fn, From b5da8e50ca54bf2f78c3394a6c67fa5e188f4cc2 Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Mon, 10 Aug 2026 23:21:33 -0700 Subject: [PATCH 18/20] Support bottom-right causal FP8 THD attention 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 --- .../common/fused_attn/fused_attn.cpp | 8 ++++-- .../common/fused_attn/fused_attn_fp8.cu | 12 ++++++--- .../dot_product_attention/context_parallel.py | 27 +++++++++++++++---- 3 files changed, 36 insertions(+), 11 deletions(-) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 93747ea14d..472264d75a 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -274,7 +274,9 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( (attn_mask_type == NVTE_Mask_Type::NVTE_NO_MASK || attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)) || + attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK || + (sm_arch_ >= 100 && + attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK))) || // 9.21: d_qk=192, d_v=128 (cudnn_runtime_version >= 92100 && sm_arch_ >= 100 && head_dim_qk <= 192 && head_dim_v <= 128 && head_dim_qk % 16 == 0 && head_dim_v % 16 == 0 && @@ -294,7 +296,9 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( (cudnn_runtime_version >= 92500 && !is_training)) && qkv_format == NVTE_QKV_Format::NVTE_THD && supported_ragged_offset_size && (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK))) && + attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK || + (sm_arch_ >= 100 && + attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK)))) && // 9.10.0: known bugs with SDPA FP8 (cudnn_runtime_version != 91000) && !return_max_logit) { backend = NVTE_Fused_Attn_Backend::NVTE_FP8; diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 18f963e69e..2ef0ac3393 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -38,9 +38,11 @@ void fused_attn_fp8_fwd_impl( bool is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); bool is_alibi = (bias_type == NVTE_Bias_Type::NVTE_ALIBI); bool is_causal = ((mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)); + (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK) || + (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK)); bool is_padding = ((mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)); + (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK) || + (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK)); bool is_dropout = (is_training && dropout_probability != 0.0f); bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); auto bias_b = b; @@ -655,9 +657,11 @@ void fused_attn_fp8_bwd_impl( bool is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); bool is_alibi = (bias_type == NVTE_Bias_Type::NVTE_ALIBI); bool is_causal = ((mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)); + (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK) || + (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK)); bool is_padding = ((mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)); + (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK) || + (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK)); bool is_dropout = (dropout_probability != 0.0f); bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); auto bias_b = b; diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 6bee2c66ae..2fa711189d 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -105,6 +105,23 @@ def _zero_thd_padding(tensor, cu_seqlens, cu_seqlens_padded): tensor[padding_mask] = 0 +def _copy_fp8_thd_half(dst, src, cu_seqlens, half_idx): + """Place a raw FP8 THD half-gradient and zero the inactive half.""" + seq_dim = 0 if dst.dim() == 3 else 1 + half_cu_seqlens = cu_seqlens // 2 + src_rows = torch.arange(src.shape[seq_dim], device=src.device) + seq_ids = torch.searchsorted(half_cu_seqlens[1:], src_rows, right=True) + first_half_rows = src_rows + half_cu_seqlens[seq_ids] + second_half_rows = src_rows + half_cu_seqlens[seq_ids + 1] + copy_rows, zero_rows = ( + (first_half_rows, second_half_rows) + if half_idx == 0 + else (second_half_rows, first_half_rows) + ) + dst.index_copy_(seq_dim, copy_rows, src) + dst.index_fill_(seq_dim, zero_rows, 0) + + def flash_attn_p2p_communicate( rank, send_tensor, send_dst, recv_tensor, recv_src, cp_group, batch_p2p_comm ): @@ -2672,8 +2689,8 @@ def backward(ctx, dout, *_args): elif ctx.qkv_format == "sbhd": dq[0].fill_(0) dq[1].copy_(dq_) - else: - dq.copy_(dq_) + elif ctx.qkv_format == "thd": + _copy_fp8_thd_half(dq, dq_, cu_seqlens_q_padded, 1) elif causal: if i > (cp_size - rank - 1): dq.add_(dq_) @@ -2759,9 +2776,9 @@ def backward(ctx, dout, *_args): dk[1].fill_(0) dv[0].copy_(dv_) dv[1].fill_(0) - else: - dk.copy_(dk_) - dv.copy_(dv_) + elif ctx.qkv_format == "thd": + _copy_fp8_thd_half(dk, dk_, cu_seqlens_kv_padded, 0) + _copy_fp8_thd_half(dv, dv_, cu_seqlens_kv_padded, 0) else: dk.copy_(dk_) dv.copy_(dv_) From c90b09b8042a51b06ec33901674f3240de95fee2 Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Mon, 10 Aug 2026 23:33:52 -0700 Subject: [PATCH 19/20] Fuse delayed FP8 THD half-gradient placement 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 --- tests/pytorch/attention/test_cp_utils.py | 20 +++++++ .../common/fused_attn/context_parallel.cu | 57 ++++++++++++++----- .../include/transformer_engine/fused_attn.h | 8 ++- .../dot_product_attention/context_parallel.py | 29 +++------- 4 files changed, 76 insertions(+), 38 deletions(-) diff --git a/tests/pytorch/attention/test_cp_utils.py b/tests/pytorch/attention/test_cp_utils.py index c3a423cef5..8ae43441fe 100644 --- a/tests/pytorch/attention/test_cp_utils.py +++ b/tests/pytorch/attention/test_cp_utils.py @@ -881,6 +881,26 @@ def test_thd_read_half_tensor_reads_each_sequence_half(self): torch.equal(kv_second, torch.stack([expected_second, expected_second + 128])) ) + def test_thd_grad_correction_copies_byte_half_and_zeros_inactive_half(self): + cu_seqlens = torch.tensor([0, 8, 20], dtype=torch.int32, device="cuda") + grad_per_step = torch.arange(10 * 2 * 8, dtype=torch.uint8, device="cuda").view(10, 2, 8) + first_half_rows = torch.tensor( + [0, 1, 2, 3, 8, 9, 10, 11, 12, 13], device="cuda" + ) + second_half_rows = torch.tensor( + [4, 5, 6, 7, 14, 15, 16, 17, 18, 19], device="cuda" + ) + + grad = torch.full((20, 2, 8), 255, dtype=torch.uint8, device="cuda") + tex.thd_grad_correction(grad, grad_per_step, cu_seqlens, "copy", "zero") + self.assertTrue(torch.equal(grad[first_half_rows], grad_per_step)) + self.assertEqual(torch.count_nonzero(grad[second_half_rows]).item(), 0) + + grad.fill_(255) + tex.thd_grad_correction(grad, grad_per_step, cu_seqlens, "zero", "copy") + self.assertEqual(torch.count_nonzero(grad[first_half_rows]).item(), 0) + self.assertTrue(torch.equal(grad[second_half_rows], grad_per_step)) + def test_thd_read_second_half_lse_handles_packed_and_batch_major_lse(self): cu_seqlens = torch.tensor([0, 8, 16], dtype=torch.int32, device="cuda") lse = torch.arange(2 * 2 * 8, dtype=torch.float32, device="cuda").view(2, 2, 8) diff --git a/transformer_engine/common/fused_attn/context_parallel.cu b/transformer_engine/common/fused_attn/context_parallel.cu index 0f7b820bbd..b4fd7e1256 100644 --- a/transformer_engine/common/fused_attn/context_parallel.cu +++ b/transformer_engine/common/fused_attn/context_parallel.cu @@ -42,6 +42,12 @@ struct CopyFunctor { } }; +struct ZeroFunctor { + __forceinline__ __device__ static void run(void *token, void *token_per_step, int idx) { + reinterpret_cast(token)[idx] = make_float4(0.f, 0.f, 0.f, 0.f); + } +}; + template struct AddFunctor { __forceinline__ __device__ static void run(dtype *token, dtype *token_per_step, int idx) { @@ -357,24 +363,27 @@ __global__ void thd_grad_correction_kernel(dtype *grad, dtype *grad_per_step, in for (int token_id = group_id; token_id < num_total_tokens; token_id += num_groups) { int seq_id = binary_search(token_id, cu_seqlens_s, batch + 1); - int token_offset; - bool is_first_half; if constexpr (functor_idx < 2) { - token_offset = cu_seqlens_s[seq_id + functor_idx]; - is_first_half = (functor_idx == 0); + dtype *first_half_token = + &grad[(token_id + cu_seqlens_s[seq_id]) * static_cast(hidden_size)]; + dtype *second_half_token = + &grad[(token_id + cu_seqlens_s[seq_id + 1]) * static_cast(hidden_size)]; + dtype *token_per_step = &grad_per_step[token_id * static_cast(hidden_size)]; + for (int idx = lane_id; idx < num_inner_loops; idx += group_size) { + Functor_0::run(first_half_token, token_per_step, idx); + Functor_1::run(second_half_token, token_per_step, idx); + } } else { - token_offset = 0; int len = cu_seqlens_s[seq_id + 1] - cu_seqlens_s[seq_id]; - is_first_half = (token_id - cu_seqlens_s[seq_id]) < (len / 2); - } - - dtype *token = &grad[(token_id + token_offset) * static_cast(hidden_size)]; - dtype *token_per_step = &grad_per_step[token_id * static_cast(hidden_size)]; - for (int idx = lane_id; idx < num_inner_loops; idx += group_size) { - if (is_first_half) { - Functor_0::run(token, token_per_step, idx); - } else { - Functor_1::run(token, token_per_step, idx); + bool is_first_half = (token_id - cu_seqlens_s[seq_id]) < (len / 2); + dtype *token = &grad[token_id * static_cast(hidden_size)]; + dtype *token_per_step = &grad_per_step[token_id * static_cast(hidden_size)]; + for (int idx = lane_id; idx < num_inner_loops; idx += group_size) { + if (is_first_half) { + Functor_0::run(token, token_per_step, idx); + } else { + Functor_1::run(token, token_per_step, idx); + } } } } @@ -707,6 +716,12 @@ static void thd_grad_dispatcher(Tensor grad, const Tensor &grad_per_step, const } else if (first_half == "none" && second_half == "copy") { thd_grad_correction_helper(grad, grad_per_step, cu_seqlens, stream); + } else if (first_half == "copy" && second_half == "zero") { + thd_grad_correction_helper(grad, grad_per_step, cu_seqlens, + stream); + } else if (first_half == "zero" && second_half == "copy") { + thd_grad_correction_helper(grad, grad_per_step, cu_seqlens, + stream); } else if (first_half == "add" && second_half == "copy") { thd_grad_correction_helper, CopyFunctor, 2>(grad, grad_per_step, cu_seqlens, stream); @@ -722,6 +737,18 @@ void thd_grad_correction(Tensor grad, const Tensor &grad_per_step, const Tensor const std::string &first_half, const std::string &second_half, cudaStream_t stream) { using namespace transformer_engine; + if (grad.dtype() == DType::kByte) { + if (first_half == "copy" && second_half == "zero") { + thd_grad_correction_helper(grad, grad_per_step, cu_seqlens, + stream); + } else if (first_half == "zero" && second_half == "copy") { + thd_grad_correction_helper(grad, grad_per_step, cu_seqlens, + stream); + } else { + NVTE_ERROR("Byte gradients require copy/zero or zero/copy correction\n"); + } + return; + } TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( grad.dtype(), dtype, thd_grad_dispatcher(grad, grad_per_step, cu_seqlens, first_half, second_half, diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index 41e4b136bd..4f8b2d284e 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -509,9 +509,11 @@ void nvte_cp_thd_out_correction(NVTETensor out, const NVTETensor &out_per_step, * \param[out] grad Output tensor. * \param[in] grad_per_step THD format gradient of context parallelism. * \param[in] cu_seqlens Cumulative sequence lengths, [batch_size + 1]. - * \param[in] first_half One of ("add", "copy", "none") correction op for first half. - * \param[in] second_half One of ("add", "copy", "none") correction op for second half. - Must be different from first_half. + * \param[in] first_half One of ("add", "copy", "none", "zero") correction op for + * 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. * \param[in] stream CUDA stream used for this operation. */ void nvte_cp_thd_grad_correction(NVTETensor grad, const NVTETensor &grad_per_step, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 2fa711189d..de8251af04 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -105,23 +105,6 @@ def _zero_thd_padding(tensor, cu_seqlens, cu_seqlens_padded): tensor[padding_mask] = 0 -def _copy_fp8_thd_half(dst, src, cu_seqlens, half_idx): - """Place a raw FP8 THD half-gradient and zero the inactive half.""" - seq_dim = 0 if dst.dim() == 3 else 1 - half_cu_seqlens = cu_seqlens // 2 - src_rows = torch.arange(src.shape[seq_dim], device=src.device) - seq_ids = torch.searchsorted(half_cu_seqlens[1:], src_rows, right=True) - first_half_rows = src_rows + half_cu_seqlens[seq_ids] - second_half_rows = src_rows + half_cu_seqlens[seq_ids + 1] - copy_rows, zero_rows = ( - (first_half_rows, second_half_rows) - if half_idx == 0 - else (second_half_rows, first_half_rows) - ) - dst.index_copy_(seq_dim, copy_rows, src) - dst.index_fill_(seq_dim, zero_rows, 0) - - def flash_attn_p2p_communicate( rank, send_tensor, send_dst, recv_tensor, recv_src, cp_group, batch_p2p_comm ): @@ -2690,7 +2673,9 @@ def backward(ctx, dout, *_args): dq[0].fill_(0) dq[1].copy_(dq_) elif ctx.qkv_format == "thd": - _copy_fp8_thd_half(dq, dq_, cu_seqlens_q_padded, 1) + tex.thd_grad_correction( + dq, dq_, cu_seqlens_q_padded, "zero", "copy" + ) elif causal: if i > (cp_size - rank - 1): dq.add_(dq_) @@ -2777,8 +2762,12 @@ def backward(ctx, dout, *_args): dv[0].copy_(dv_) dv[1].fill_(0) elif ctx.qkv_format == "thd": - _copy_fp8_thd_half(dk, dk_, cu_seqlens_kv_padded, 0) - _copy_fp8_thd_half(dv, dv_, cu_seqlens_kv_padded, 0) + tex.thd_grad_correction( + dk, dk_, cu_seqlens_kv_padded, "copy", "zero" + ) + tex.thd_grad_correction( + dv, dv_, cu_seqlens_kv_padded, "copy", "zero" + ) else: dk.copy_(dk_) dv.copy_(dv_) From 31539f8819c090a055eab046f261ca5746024d69 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 11 Aug 2026 06:42:01 +0000 Subject: [PATCH 20/20] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/attention/test_cp_utils.py | 8 ++------ .../dot_product_attention/context_parallel.py | 12 +++--------- 2 files changed, 5 insertions(+), 15 deletions(-) diff --git a/tests/pytorch/attention/test_cp_utils.py b/tests/pytorch/attention/test_cp_utils.py index 8ae43441fe..1630b2d8c7 100644 --- a/tests/pytorch/attention/test_cp_utils.py +++ b/tests/pytorch/attention/test_cp_utils.py @@ -884,12 +884,8 @@ def test_thd_read_half_tensor_reads_each_sequence_half(self): def test_thd_grad_correction_copies_byte_half_and_zeros_inactive_half(self): cu_seqlens = torch.tensor([0, 8, 20], dtype=torch.int32, device="cuda") grad_per_step = torch.arange(10 * 2 * 8, dtype=torch.uint8, device="cuda").view(10, 2, 8) - first_half_rows = torch.tensor( - [0, 1, 2, 3, 8, 9, 10, 11, 12, 13], device="cuda" - ) - second_half_rows = torch.tensor( - [4, 5, 6, 7, 14, 15, 16, 17, 18, 19], device="cuda" - ) + first_half_rows = torch.tensor([0, 1, 2, 3, 8, 9, 10, 11, 12, 13], device="cuda") + second_half_rows = torch.tensor([4, 5, 6, 7, 14, 15, 16, 17, 18, 19], device="cuda") grad = torch.full((20, 2, 8), 255, dtype=torch.uint8, device="cuda") tex.thd_grad_correction(grad, grad_per_step, cu_seqlens, "copy", "zero") diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index de8251af04..dd4a91de0a 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -2673,9 +2673,7 @@ def backward(ctx, dout, *_args): dq[0].fill_(0) dq[1].copy_(dq_) elif ctx.qkv_format == "thd": - tex.thd_grad_correction( - dq, dq_, cu_seqlens_q_padded, "zero", "copy" - ) + tex.thd_grad_correction(dq, dq_, cu_seqlens_q_padded, "zero", "copy") elif causal: if i > (cp_size - rank - 1): dq.add_(dq_) @@ -2762,12 +2760,8 @@ def backward(ctx, dout, *_args): dv[0].copy_(dv_) dv[1].fill_(0) elif ctx.qkv_format == "thd": - tex.thd_grad_correction( - dk, dk_, cu_seqlens_kv_padded, "copy", "zero" - ) - tex.thd_grad_correction( - dv, dv_, cu_seqlens_kv_padded, "copy", "zero" - ) + tex.thd_grad_correction(dk, dk_, cu_seqlens_kv_padded, "copy", "zero") + tex.thd_grad_correction(dv, dv_, cu_seqlens_kv_padded, "copy", "zero") else: dk.copy_(dk_) dv.copy_(dv_)