Fix stateful KV tail at the NGEN limit

[winklemad]

Summary

When generation stopped exactly at the NGEN limit, spec_decode emitted the final token but returned before forwarding it into the stateful KV cache. As a result, :more resumed from the preceding token and .coli_kv omitted the visible tail of the answer.

Let stateful callers take the existing zero-draft forward path for the capped token before kv_out advances. One-shot generation keeps its existing saved terminal forward.

Add a dependency-free zero-layer regression covering NGEN 0/1, continued generation and :more, persisted history, accepted speculative tokens, EOS, and normal non-limit termination.

Validation

No CUDA paths or performance claims are involved.

• make -C c check
• CUDA changes were tested with make -C c cuda-test (if applicable)
• Performance claims include hardware, commands, and repeatable measurements

Compatibility

• The default CPU build remains dependency-free
• No model files, generated binaries, or benchmark artifacts are included

[JustVugg]

Thanks for this — the stateful-KV-tail bug is real (the last token dropped from persistence and re-forwarded on MORE), and shipping a dependency-free regression test is exactly right. But the test itself is failing in CI on both Linux and Windows, so I can't merge yet:

tests/test_spec_decode_state.c:99: check failed: out.n == 1 && emitted[0] == 2 && all[2] == 2
tests/test_spec_decode_state.c:216: check failed: test_ngen_zero_and_one_shot() == 0

It's the one-shot case (kv_out == NULL) that trips, which is the path the fix is supposed to leave byte-for-byte unchanged. So either the fix changes one-shot behavior after all, or the test's expectation for that case is off. Worth running make -C c test-c locally (or just tests/test_spec_decode_state) to see which — the CI is running the exact same thing.

Once the test is green I'll merge it — the fix direction is correct and the coverage is good; it's just that the gate has to pass before it lands on dev (a red test in the suite blocks everyone). Ping me when it's updated.

[ZacharyZcR]

This is a real bug and !kv_out is the correct predicate — and I owe you a heads-up, because #581 (single-slot serve MTP, merged to dev a few hours ago) both collides with this PR and adds a second caller your fix covers.

Why it's now CONFLICTING: #581 gave spec_decode a 10th parameter — float **logit_out — so the mux scheduler can own the continuation logits. The signature and every call site changed:

static int spec_decode(Model *m, int *all, int kv, int n_new, int eos, float *logit,
                       void (*emit)(int,void*), void *ud, int *kv_out, float **logit_out);

Your one-line change to line ~4721 is untouched by that (I didn't move the break), so it rebases cleanly. The conflict is your new test_spec_decode_state.c: its spec_decode(...) calls are 9-argument and won't compile against the new signature. Add a trailing NULL (the non-logit_out path, exactly like the four production callers) and they're correct again.

And here's why your fix matters more after the rebase, not less: #581 added a fifth spec_decode caller — the single-slot mux serve path — and it is a stateful caller with the exact exposure you're fixing:

int prod = spec_decode(m, sc->hist, sc->len, r->maximum - r->emitted, eos, lg,
                       mux_spec_emit, &ud, &sc->len, NULL);   // kv_out = &sc->len (non-NULL)

At an NGEN cap this hits the same if(emitted>=n_new) break before the final token is forwarded, so mux_done's kv_disk_append persists sc->len without the last visible token — the mux .coli_kv / prefix-match version of the :more truncation you documented. Your && !kv_out guard fixes this caller for free, since it passes kv_out non-NULL.

So the rebase is worth it: (1) trailing NULL on the test's calls, and (2) if you want the regression to cover the caller #581 introduced, a case driving the mux single-slot spec path (or at minimum a comment noting spec_decode's stateful callers now include serve, not just chat/run_serve). Happy to help wire that test case — it's my exposure your fix closes.

+1. Correct root-cause fix; the merge order just means it needs a touch-up against the new signature.

[winklemad]

Thanks, Zachary — this is really helpful, especially the context around the new single-slot mux caller.

I also found the cause of the earlier Linux/Windows CI failure. It is in my test helper, not the stateful fix: choose() allocates the logits with falloc() (which uses malloc) but initializes only the selected slot. The remaining entries contain uninitialized values, so another token can accidentally win on some platforms. I will initialize the full logits buffer before setting the chosen token.

I will rebase onto the current dev branch, add the trailing NULL logit_out argument to every spec_decode call in the test, and rerun make -C c test-c and make -C c check before pushing the update.

The mux path is useful confirmation that this fix now protects more than the original :more case. I can also add a focused mux regression, or adjust the existing coverage, if you think that would be useful — just let me know what you would prefer.