[common] Add max_t support for KV in THD (#1370)

add max_t for KV Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

[common] Add max_t support for KV in THD (#1370)
add max_t for KV Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
f4f35c2f · Charlene Yang · GitHub · 0196ed44 · f4f35c2f
Unverified Commit f4f35c2f authored Dec 16, 2024 by Charlene Yang Committed by GitHub Dec 16, 2024
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 0 deletions

transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu ...gine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu +1 -0

No files found.
--- a/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu
+++ b/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu
@@ -661,6 +661,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl(

      if (is_ragged && cudnn_runtime_version >= 90600) {
        sdpa_backward_options.set_max_total_seq_len_q(s_q);
+        sdpa_backward_options.set_max_total_seq_len_kv(s_kv);
      }

      if (cudnn_runtime_version >= 90200 && window_size_left != -1) {