Fix early exit condition in attention kernel (#788)

816022e4 · Li Zhang · GitHub · d3e2cee4 · 816022e4
Unverified Commit 816022e4 authored Dec 02, 2023 by Li Zhang Committed by GitHub Dec 02, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

src/turbomind/kernels/decoder_multihead_attention/decoder_multihead_attention_template.h ...ultihead_attention/decoder_multihead_attention_template.h +1 -1

No files found.
--- a/src/turbomind/kernels/decoder_multihead_attention/decoder_multihead_attention_template.h
+++ b/src/turbomind/kernels/decoder_multihead_attention/decoder_multihead_attention_template.h
@@ -736,7 +736,7 @@ struct DecoderMultiHeadAttentionKernel {
        }

        // early exit if split if out of bound
-        if (kSplitK && step_begin_ >= step_end_) {
+        if (kSplitK && step_begin_ && step_begin_ >= step_end_) {
            return;
        }