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

add max_t for KV Signed-off-by: Charlene Yang <[email protected]>
NVIDIA · Dec 17, 2024 · f4f35c2 · f4f35c2
1 parent 0196ed4
commit f4f35c2
Showing 1 changed file with 1 addition and 0 deletions.
diff --git 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) {