help to understand data structure behind add_seq_to_ubatch

[CharlesJu1]

I am having a difficult time trying to understand the code of add_seq_to_ubatch. From the comments, it sounds like ubatch can only accept equal length batch. However, when I step through the code, the ubatch.equal_seqs is false and it goes to the "simple split" case, which handles "unknown number of sequences of unequal lengths" according to the comments. And later, there is a concept of "virtual sequences for simple splits", which increments the ubatch.n_seqs by "length".

Could someone help outline the logic behind add_seq_to_ubatch? The relationship between seq，ubatch, "virtual sequence"?

    void add_seq_to_ubatch(llama_ubatch & ubatch, llama_sbatch_seq & seq, size_t length) {
        GGML_ASSERT(batch != nullptr);
        GGML_ASSERT(length <= seq.length);
        // Can only add sequences of equal lengths to a batch,
        // otherwise it isn't clear to which sequence a token belongs
        GGML_ASSERT(seq.n_seq_id == 0 || ubatch.n_seqs == 0 || length == (size_t) ubatch.n_tokens / ubatch.n_seqs);
        GGML_ASSERT((seq.n_seq_id != 0) == ubatch.equal_seqs);
        // NOTE: loops are separated for cache-friendliness
        if (batch->token) {
            if (ubatch.equal_seqs) {
                for (size_t i = 0; i < length; ++i) {
                    ubatch.token[ubatch.n_tokens + i] = batch->token[ids[seq.offset + i]];
                }
            } else {
                // simple split
                ubatch.token = batch->token + seq.offset;
            }
        } else {
            ubatch.token = nullptr;
        }

       ... 

        ubatch.n_tokens += length;
        ubatch.n_seqs += ubatch.equal_seqs ? 1 : length; // virtual sequences for simple splits
        seq.offset += length;
        seq.length -= length;
        n_tokens -= length;
        GGML_ASSERT(ubatch.n_tokens == ubatch.n_seq_tokens * ubatch.n_seqs);
    }

    // simple split, unknown number of sequences of unequal lengths
    llama_ubatch split_simple(size_t n_ubatch) {
        n_ubatch = n_tokens < n_ubatch ? n_tokens : n_ubatch;
        llama_ubatch ubatch = reserve_ubatch(n_ubatch, /* has_embd */ batch->embd != nullptr);
        ubatch.equal_seqs = false;
        if (!seq.empty()) {
            llama_sbatch_seq & s = seq[0];
            size_t length = s.length < n_ubatch ? s.length : n_ubatch;
            GGML_ASSERT(seq.size() == 1 && s.n_seq_id == 0); // don't mix with other splits
            add_seq_to_ubatch(ubatch, s, length);
        }
        return ubatch;
    }