diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index f2ab4a8495..341c99346d 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -1948,8 +1948,8 @@ def get_model(dtype, config): } 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.") @@ -2525,6 +2525,8 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: attn_mask_type=config.attn_mask_type, 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", ) 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 fc21771297..1e8ed96e2c 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -281,13 +281,17 @@ 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))) && - !requires_64bit_ragged_offset && + 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 && + (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; @@ -627,12 +631,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, is_training, + 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_rng_state, wkspace, stream, handle); + 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"); } @@ -729,14 +734,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, attn_scale, dropout, - qkv_layout, o_format, do_format, dqkv_layout, qkv_scale_inv_format, + 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_rng_state, wkspace, stream, handle); + 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 000af41aee..18f963e69e 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" @@ -15,17 +16,20 @@ 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, - 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* 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) { @@ -60,11 +64,21 @@ 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 >= kFP8THDRaggedCudnnVersion && sm_arch_ != 120; + + NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); + const DType ragged_offset_type = DType::kInt64; + // 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. @@ -78,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, @@ -139,6 +167,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_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 @@ -164,6 +197,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 @@ -175,6 +210,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}) @@ -185,6 +228,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}) @@ -362,15 +419,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 @@ -396,6 +471,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_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) : std::make_tuple(nullptr, nullptr); @@ -406,24 +487,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_q_tuple, offset_kv_tuple, + offset_o_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_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()); + const size_t num_bytes_per_seqlen = alignTo<16>(b * sizeof(int32_t)); + 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; + 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. - // 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; + *workspace_size = + plan_workspace_size + actual_seqlen_workspace_size + seqlen_offsets_workspace_size; return; } @@ -465,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()); @@ -476,6 +570,49 @@ void fused_attn_fp8_fwd_impl( } } + 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; + } + const RaggedOffsetMultipliers offset_mults(layout_group, h, hg, d_qk, d_v); + cu_seqlens_padded_to_offsets<<>>( + offset_mults, actual_b, b, 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; @@ -489,29 +626,30 @@ 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( 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, - 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, 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, - 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* 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) { using namespace transformer_engine; const auto cudnn_runtime_version = cudnnGetVersion(); bool is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); @@ -543,6 +681,28 @@ 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 >= kFP8THDRaggedCudnnVersion && sm_arch_ != 120; + + 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; + } + } + bool is_O_in_F16 = (o_tensor_type == cudnn_frontend::DataType_t::HALF || o_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); @@ -628,6 +788,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_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 @@ -660,6 +825,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 @@ -673,6 +840,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}) @@ -683,21 +863,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}) @@ -947,6 +1159,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) @@ -955,6 +1170,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}) @@ -1013,6 +1232,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_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) : std::make_tuple(nullptr, nullptr); @@ -1024,7 +1249,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_q_tuple, + offset_kv_tuple, offset_o_tuple, offset_s_tuple, dropout_tuple); cache.insert({descriptor, return_tuple}); return return_tuple; @@ -1033,14 +1259,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_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()); + 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; } @@ -1106,9 +1345,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()); @@ -1116,6 +1355,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; + } + const RaggedOffsetMultipliers offset_mults(layout_group, h, hg, d_qk, d_v); + cu_seqlens_padded_to_offsets<<>>( + offset_mults, actual_b, b, 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; @@ -1137,14 +1419,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; @@ -1171,12 +1455,39 @@ 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(); + 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; + + 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}; + // 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 && 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}; + } output_M->data.dtype = DType::kFloat32; Tensor* output_rng_state = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_rng_state->data.dptr = nullptr; @@ -1218,18 +1529,19 @@ 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, - 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, - 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, or BHSD.\n"); + NVTE_ERROR("FP8 fused attention only supports qkv_format=BSHD, SBHD, BHSD, or THD.\n"); } if (workspace_size > 0) { @@ -1247,18 +1559,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; @@ -1323,6 +1637,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 = @@ -1339,23 +1672,24 @@ 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, - 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, + 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, &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