Skip to content

Flash Next on CUDA: warm turns, so the next turn resumes past the reply - #3219

Open
ashhart wants to merge 3 commits into
mainfrom
lead/flashnext-cuda-warm-turns
Open

ashhart wants to merge 3 commits into
mainfrom
lead/flashnext-cuda-warm-turns

Conversation

@ashhart

@ashhart ashhart commented Nov 13, 2026

Copy link
Copy Markdown
Owner

Flash Next on CUDA now takes the background pass the server sends after a chat reply, as Nemotron on CUDA and Flash Next on Metal already do. The next turn resumes past the reply instead of prefilling it again.

  • The family declares warm_quiet, 150 ms as on Nemotron. Its CUDA cache rules then wait for quiet and keep the pass's state one token short of its end, since the MTP head reads one token ahead. The family keeps and resumes states at any row, so the pass needs no message-start cuts and none are added.
  • It runs on one rank only. With --tp the store's warm rule is off and the server sends no pass. With --prompt-cache-gib 0 nothing changes.
  • The pass gives way to real work. Before each chunk it checks whether a request is waiting, a stream is decoding or another prompt is still going in. If so it stops, and the cache keeps what it filled. Its chunks use the busy chunk's rows, 1,536 by default, so on a one-stream server a request that arrives waits for at most that many rows, not 8,192.
  • In the lane host, a request whose prompt is still filling a chunk at a time now counts as work for a warm pass. Before, a pass taken beside one waited out its quiet while that prompt stood still, then ran whole. Warm passes on every host with a fill queue now give way sooner, and replies are unchanged.
  • Requests that are not warm passes take the same path and rows as before.

Host tests: a warm pass taken while a request's prompt is filling runs no chunk, or yields at its first, and the request finishes (this test fails without the lane host change); a Flash Next CUDA warm pass steps aside before its first chunk through the real lane opener; Flash Next's CUDA cache rules keep warm marks. test-lane-api-host, test-cuda-host and the Linux CUDA test compile pass.

Device check on a GB10 (running): cold prompt and decode A B B A with replies equal to main, and a five-turn chat at 1 and 4 streams, cache on against cache off, greedy and seeded, drafted and not, with cached tokens past each reply from turn 2 and a request sent during a warm pass. Results go here before this lands.

Fixes #3176

  • Synthena (Ash's Agent)

@devopsfortunato

Copy link
Copy Markdown
Contributor

Device check of this head on a DGX Spark (GB10, 128 GB unified memory), Qwen3.8 Flash Next modelopt-mixed.

9ab8aed1a against its parent on main a3aecd8d0, parent first, a fresh server per build:
--context 65536 --parallel 2 --max-tokens 32768 --prompt-cache-gib 12 (4 for the shared-cache run),
TENSORFOLD_MEMORY_RESERVE_GIB=3, same prompts on both builds (14k-19k tokens; the stream's conversation 3.4k-5.5k), temperature 0, one run per case.

Replies were byte-identical between the two builds in all five runs (68 requests per build, 0 errors).

Run Tokens read again (parent / PR) TTFT on follow-ups (parent -> PR) Mean difference
Follow-up 1 s after the reply reply + ~40 (80-1,095) / 28-43 short replies 0.25-0.27 -> 0.18-0.20 s; 880-1,100-token replies 0.65-0.75 -> 0.18-0.20 s -0.218 s
Follow-up 0.3 s after - - -0.009 s
Follow-up 0.15 s after same on both after short replies 0.24-0.27 -> 0.28-0.33 s +0.021 s
Five conversations in turn, 4 GiB cache (full) - every follow-up found its state -0.074 s
Short follow-ups beside a 2,000-token stream same on both (78-108) 0.25-0.33 s on both 0.00 s
  • The stream's speed: alone 49.2 (parent) and 49.1 (PR) tokens/s; beside the follow-ups 54.4 and 54.5.
  • The stream's own follow-up, 1 s after its 2,000-token reply, met the pass part way: 511 tokens read instead of
    2,047, TTFT 0.49 s instead of 1.29 s. The prefix read before the pass gave way is kept and used.
  • Costs seen: seven of nine follow-ups 0.15 s after a reply were 0.03-0.07 s slower (they meet the pass as it
    starts), and the cache
    evicts more often, since each pass keeps one more state: evicted 44 against 27 over the 12 GiB runs and 38
    against 18 on the 4 GiB run. No miss was added: misses 1 on both builds in both servers.

Not run: --prompt-cache-disk, --tp, more than one GPU.

…the reply as on Nemotron

After a chat reply the server sends a background prompt-only pass over the reply as the next turn's prompt renders
it. Flash Next on CUDA did not take it, so every next turn resumed at its history's end and prefilled the reply
again. The family now declares warm_quiet, 150 ms as Nemotron's. Its CUDA cache rules then wait for quiet and keep
the pass's state one token short of its end, since the MTP head reads one token on. The native CUDA host turns warm
turns on when the prompt cache is on and one rank serves. The family keeps and resumes states at any row, so the
pass needs no message-start cuts. With --tp the store's warm rule is off and the server sends no pass.

The pass gives way to real work. Its whole-prompt prefill asks before each chunk whether to step aside, and when it
does it keeps the state it reached. Its chunks take the busy chunk's rows, 1,536 by default, where a one-stream
server's solo pass takes 8,192, so a request that arrives waits for at most that many rows.

The lane host now also counts another request's prompt that is still filling a chunk a round as work. Before, a
pass taken beside one saw no queued request and no decoding stream. It waited out its quiet on the lane thread
while that prompt stood still, then ran whole. Warm passes on every host with a fill queue now give way sooner, and
replies are unchanged.

Requests that are not warm passes take the same path and the same rows as before. Warm states go through the
store's budget and eviction like every other mark.

Host tests: a warm pass taken while a request's prompt is filling runs no chunk, or yields at its first, and the
request finishes. On main's lane host the pass ran whole and the test fails with expected .cancelled, found
.length. A Flash Next CUDA warm pass steps aside before its first chunk through the real lane opener, and Flash
Next's CUDA cache rules keep warm marks.
@ashhart
ashhart force-pushed the lead/flashnext-cuda-warm-turns branch from 9ab8aed to 668492e Compare November 14, 2026 03:31

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Flash Next on CUDA: warm turns, so the next turn resumes past the reply

2 participants