[ngxson]

Overview

Fix efe3626#commitcomment-204575113

Not 100% sure but seems like there is a problem with the case: non-causal + hybrid + multi seq

This seems to fix the problem

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: yes

[ggerganov]

Yes, when computing non-causal data, it has to fit in a single ubatch. The proposed fix should be good.

It can be improved by allowing split_seq to fit multiple sequences in the ubatch, as long as they can fit fully. This way for example, with an ubatch of 4096 tokens, one can process in a single pass 4 sequences of length 1024, instead of splitting them in 4 ubatches.

Add TODO to support that:

diff --git a/src/llama-batch.h b/src/llama-batch.h
index 8882c9941e..e1e43979b5 100644
--- a/src/llama-batch.h
+++ b/src/llama-batch.h
@@ -186,6 +186,7 @@ public:
     llama_ubatch split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail);
 
     // sequence-set-wise split - each ubatch contains a single sequence-set
+    // TODO: extend to support multiple _full_ sequences in the ubatch
     llama_ubatch split_seq(uint32_t n_ubatch);
 
     // a helper method for creating a well-defined ubatch of tokens