server: improve user message detection and create checkpoints at every user message

[aldehir]

Overview

Scan for user message boundaries directly on the tokens and avoid doing a second pass that translates from byte offsets.

Create checkpoints at the start of every user message, as opposed to only the last message.

cc. @ggerganov @jacekpoplawski

Additional information

The checkpoint logic makes sense to me, but I haven't thoroughly tested it for edge cases.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: Yes - used to update the chat params for every model. Also used to review the checkpoint changes.

[jacekpoplawski]

Could you explain the use case where I could observe the behavior this PR fixes?

I think my problem is that I implemented #22929 for my own use case, and I don't see any issues with the current implementation. So I assume this is either for a different kind of model (not Qwen 3.6 27B) or for different behavior (not agentic coding in pi).

I remember ggerganov mentioned /tree in pi, but I use /tree in pi often and there is no reprocessing, because each step is stored in the checkpoint.

[aldehir]

I remember ggerganov mentioned /tree in pi, but I use /tree in pi often and there is no reprocessing, because each step is stored in the checkpoint.

If you resume an older chat session, or restart the server, then using /tree after sending a message will trigger it.

[ggerganov]

I remember ggerganov mentioned /tree in pi, but I use /tree in pi often and there is no reprocessing, because each step is stored in the checkpoint.

If you resume an older chat session, or restart the server, then using /tree after sending a message will trigger it.

Yes, that's the main use case.

[ggerganov]

I am testing and something I didn't realize before is that each tool response is actually a "user" message. So this results in a checkpoint created for each tool call. Instead, I was thinking about making the checkpoints only at the points where the user manually provides input. But I realize that maybe we cannot distinguish between the two?

[aldehir]

But I realize that maybe we cannot distinguish between the two?

A tool response message in this case contains the user message preamble as a prefix, so we simply need to check for a tool response first. I can easily do this with the specialized parsers, but @pwilkin will need to assist in providing it for the autoparser variants.

[ggerganov]

But I realize that maybe we cannot distinguish between the two?

A tool response message in this case contains the user message preamble as a prefix, so we simply need to check for a tool response first. I can easily do this with the specialized parsers, but @pwilkin will need to assist in providing it for the autoparser variants.

@aldehir An simpler alternative approach would be:

• Keep the logic as it is - i.e. checkpoint on all user and tool response messages
• Bump the default checkpoint_min_step from 256 -> 8192
• Add extra condition to ignore the min-step condition for the last user message

[pwilkin]

I'd tune it to 4096, I think 8192 is a bit too much with agentic sessions with large turnover, but yeah, definitely 256 is too little, the way it works currently is that if the agent calls a lot of small tools at once, then that immediately clobbers the checkpoints pool even at the large default size of 32. The way I see it, a rule of "always keep last user checkpoint, after that, respect the hard limit" would solve most of these problems.

[aldehir]

• Keep the logic as it is - i.e. checkpoint on all user and tool response messages

• Bump the default checkpoint_min_step from 256 -> 8192

• Add extra condition to ignore the min-step condition for the last user message

Sounds good, the latest commit should reflect that.

I'd tune it to 4096, I think 8192 is a bit too much with agentic sessions with large turnover

I set it 8192, can change if @ggerganov agrees.

[rankaiyx]

I rebased this PR onto the main branch—which includes #24746 —and put it through a full day of intensive testing. Everything seems to be working fine.

I'll set it to 256; the harness I use trims old tool messages, which can improve the KV cache reuse rate.

[ggerganov]

I'd tune it to 4096, I think 8192 is a bit too much with agentic sessions with large turnover

I set it 8192, can change if @ggerganov agrees.

With 8192, on average we will be reprocessing half of the step (i.e. 4096) tokens when restoring a checkpoint. I think reprocessing 4096 tokens should be OK in majority of cases so default of 8192 is better IMO - we are more interested in sparing memory usage.

[ggerganov]

Restoring an old session now works correctly.

Note that when running a current session and we do many small tool calls one after the other, the logic creates a checkpoint for each one because it is always the last user message on each turn. To improve that, we should add a logic to remove the penultimate checkpoint if it does not satisfy the step-size condition. This should make the checkpoints to be both distributed sparsely and also respect the user boundaries. With this improvement I think we should be gucci. Can also do it in a follow-up PR.

[aldehir]

@ggerganov let's do it in a follow up PR.