diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 82791084d8..d79906a684 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -2096,8 +2096,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.") @@ -2302,6 +2302,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) @@ -2340,6 +2344,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) @@ -2673,6 +2679,9 @@ 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, + # The optimized zero-fill path dereferences device memory on the host. + fast_zero_fill=False, ) if is_training: out.backward(out_grad) diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index 2c0a5d9217..70419b226a 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -402,6 +402,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, @@ -486,6 +488,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", @@ -530,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)): @@ -564,8 +581,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!") @@ -607,6 +622,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!") @@ -638,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, diff --git a/tests/pytorch/attention/test_cp_utils.py b/tests/pytorch/attention/test_cp_utils.py index c3a423cef5..1630b2d8c7 100644 --- a/tests/pytorch/attention/test_cp_utils.py +++ b/tests/pytorch/attention/test_cp_utils.py @@ -881,6 +881,22 @@ 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/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 1ac7a36383..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 && @@ -283,13 +285,20 @@ 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}; 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))) && - !requires_64bit_ragged_offset && + qkv_format == NVTE_QKV_Format::NVTE_BHSD)) || + (((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 || + (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; @@ -645,12 +654,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"); } @@ -747,14 +757,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..2ef0ac3393 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) { @@ -34,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; @@ -60,11 +66,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 +94,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 +169,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 +199,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 +212,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 +230,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 +421,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 +473,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 +489,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 +561,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 +572,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,37 +628,40 @@ 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); 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; @@ -543,6 +685,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 +792,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 +829,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 +844,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 +867,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 +1163,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 +1174,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 +1236,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 +1253,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 +1263,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 +1349,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 +1359,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 +1423,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 +1459,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 +1533,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 +1563,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 +1641,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 +1676,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 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 b444f034c8..dd4a91de0a 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,18 @@ 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 ): @@ -2660,8 +2672,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": + tex.thd_grad_correction(dq, dq_, cu_seqlens_q_padded, "zero", "copy") elif causal: if i > (cp_size - rank - 1): dq.add_(dq_) @@ -2747,9 +2759,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": + 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_) @@ -2923,22 +2935,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, @@ -3235,7 +3236,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 +3404,11 @@ 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 +3741,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 6eb3ce54f1..979dde45ec 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);