server: support inp embd to generate next token

[ngxson]

Overview

Part of #26254

When generating the next token, TTS backbone model outputs both the sampled token (for semantic code) and hidden state - this is similar to MTP where the MTP head takes both the sampled and last embd

The only difference is that MTP gives back draft tokens to tgt model, while TTS model gives back one embd row (= sum of generated codes), so this PR adds support for that case

Note that we cannot reuse the same code path from mtmd-helper (for decode embd), because that doesn't contain the sampling step

Requirements

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

[ngxson]

/bot review

[ggml-gh-bot]

Automated code review

Static review of PR #26313 (server: support inp embd to generate next token).

Blocking

(point 1) ~server_batch() has an invalid-free / UAF once the embd path is actually used. tools/server/server-context.cpp:96-100. The destructor calls llama_batch_free(batch) when batch.pos != nullptr, but after a render() in embd mode batch.token == nullptr and batch.embd points into the std::vector<float> embd member (render() sets batch.embd = embd.data()). llama_batch_free then does if (batch.embd) free(batch.embd) (src/llama-batch.cpp:980), i.e. free() on memory owned by the vector. Worse, members are destroyed in reverse declaration order, and embd is declared after batch, so the vector is already freed before llama_batch_free runs - a hard double/invalid free on shutdown.

This is reachable on the normal path: update_slots() returns after the decode loop without calling batch.clear() (tools/server/server-context.cpp:2871-2882 leaves the batch in the rendered state), so any server shutdown that follows an embd decode corrupts the heap. clear() restores the pointers but the destructor never does. Fix in the destructor, mirroring clear():

~server_batch() {
    if (batch.pos != nullptr) {
        if (batch.token == nullptr) { // we are in embd-rendered state
            batch.token = tokens_ptr;
            batch.embd  = nullptr;
        }
        llama_batch_free(batch);
    }
}

(point 2) slot.inp_embd is never assigned anywhere - grep -rn inp_embd tools/ only finds the declaration, the can_batch_with comparison, and the .empty() read at server-context.cpp:487-488. There is no producer that copies a model's hidden state into inp_embd. As a result the entire embd decode path added by this PR (batch.add(embd), the has_embd flag, batch.embd = embd.data(), the off * n_embd indexing, the spec/embd guard) is never executed by anything in this PR. This is landed, unverified infrastructure, so the blocking bug above and the indexing below cannot be validated. The producer (and ideally a real model + server test) needs to be in this PR, or this should wait for the part of #26254 that wires it up.

Will slow the review

(point 3) add(embd) does not assert embd_in.size() == n_embd. can_batch_with only enforces that all batched slots have equal inp_embd.size(), not that the size equals n_embd. A producer that emits a wrong-sized row for every slot would satisfy can_batch_with, then corrupt every row via batch.embd + off * n_embd in get_view (server-context.cpp:177-179). Add GGML_ASSERT((int32_t) embd_in.size() == n_embd); in the embd add overload.

(point 4) When a generating slot runs in embd mode, the prompt-processing loop silently skips prompt-token slots via can_batch_with (server-context.cpp:3099-3101), so cont_batching is effectively disabled for TTS-style embd generation. That is a real behavioral constraint of this design but is not documented anywhere. Worse, the only thing actually preventing a token add into an embd batch is the GGML_ASSERT(!has_embd) in add(token) (server-context.cpp:114) - i.e. an abort, not a graceful skip, if slot_batched ordering ever fails to enforce the invariant. Make the prompt path also early-out on batch.has_embd instead of relying on the assert, and document the "embd slots don't batch with prompt processing" constraint.

(point 5) The new // TODO @ngxson : dft model may have different n_embd ... comment at server-context.cpp:3626-3627 is fine on length but it and the SRV_ERR("%s", ...)/throw pair describe a not-yet-needed restriction. Per AGENTS.md keep comments to non-obvious why; consider trimming the comment to one line. Not a blocker.

Nits

(point 6) The member std::vector<float> embd; is lexically very close to llama_batch::embd and batch.embd = embd.data() reads confusingly. Renaming the buffer (e.g. embd_buf) would make render()/clear() self-documenting.

(point 7) get_view builds the view from batch.token ? batch.token + off : nullptr and batch.embd ? batch.embd + off * n_embd : nullptr. The int multiplication off * n_embd is fine for realistic shapes, but the ternary on batch.embd (which is nullptr in token mode and a vector pointer in embd mode) is the only thing tying the two modes apart - a one-line comment in get_view stating that exactly one of token/embd is non-null per batch state would help future readers.

Overall: the implementation approach (reuse one batch, swap token/embd pointers) is reasonable and minimal, but it must not land before the destructor free is fixed and the embd path is actually exercised by a model/test, given that every embd-related line is currently dead code.

This review was generated automatically by pi coding agent using zai-org/GLM-5.2. It may contain mistakes. Maintainers make the final call.

[ngxson]

(point 1) ~server_batch() has an invalid-free / UAF once the embd path is actually used.

hmm yeah that need to be fixed

[ngxson]

@ggerganov could you take a quick look? thanks!