Skip to content

[Native] Add chunked prefill and interactive chat to the LLM runner - #24639

Draft
metascroy wants to merge 1 commit into
offgraph-kv-runnerfrom
offgraph-kv-interactive
Draft

metascroy wants to merge 1 commit into
offgraph-kv-runnerfrom
offgraph-kv-interactive

Conversation

@metascroy

@metascroy metascroy commented Nov 5, 2026 •

Copy link
Copy Markdown
Contributor

Stack

Off-graph KV cache for the native ggml engine, through serving Muse Glimmer on it. Review bottom-up; each PR builds and passes its tests on its own.

Summary

Chunked prefill: prompts longer than the method's token input used to be rejected. They're now prefilled in chunks of at most that many tokens. This applies to in-graph and off-graph packages alike. generate now takes the starting position and returns the last token it chose but never fed, and whether generation stopped on an EOS.

--interactive chat: needs an off-graph package.

  • Each stdin line is wrapped in the Llama 3 chat template; BOS goes only on the first turn.
  • The conversation stays in the KV cache, so each turn feeds only its new tokens. The previous reply's last token is fed at the start of the next turn.
  • The next turn always closes the previous reply with <|eot_id|>. If the reply hit --max_new_tokens, its last token is fed first. If it ended on an EOS, the EOS itself is not fed, so <|end_of_text|> can't drift from the template.
  • --warmup_prefill is rejected with --interactive; it would prefill every turn twice. So is --dump_first_logits, since every new conversation prefills from position 0 again.
  • /reset clears the cache, a full cache starts a new conversation, and Ctrl-D quits.
run_hf_llm_ggml --device gpu --package llama1b_q4km_offgraph.ptn \
  --tokenizer tokenizer.model --interactive --kv_capacity 2048 --kv_dtype fp16

Test plan

On an Apple M4 Max, Llama 3.2 1B Q4_K_M:

  • An 821-token prompt (method input 512): the off-graph and in-graph packages give the same answer.
  • Two-turn chat: the second turn recalls facts from the first; /reset forgets them.
  • A three-turn chat with --max_new_tokens 24 and a 16-row initial cache: turn 1 hits the token limit, turn 2 ends on EOS, and the conversation stays coherent across both.
  • --interactive on an in-graph package, or combined with --warmup_prefill or --dump_first_logits, fails with a clear error.

@pytorch-bot

pytorch-bot Bot commented Nov 5, 2026 •

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24639

Note: Links to docs will display an error until the docs builds have been completed.

❌ 133 New Failures, 11 Unrelated Failures, 10 Unclassified Failures

As of commit 923c00e with merge base c7f7931 (image):

NEW FAILURES - The following jobs have failed:

UNCLASSIFIED FAILURES - DrCI could not classify the following jobs because the workflow did not run on the merge base. The failures may be pre-existing on trunk or introduced by this PR:

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Nov 5, 2026
@metascroy
metascroy force-pushed the offgraph-kv-interactive branch from ce847ce to baa5ab8 Compare November 5, 2026 22:29
@metascroy
metascroy force-pushed the offgraph-kv-interactive branch from baa5ab8 to cb061b2 Compare November 11, 2026 18:05
continue;
}
}
const size_t limit = std::min(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This loop computes its limit from capacity alone, while the existing one-shot path also caps at --max_seq_length through generation_limit. Should chat apply the same cap?

Prompts longer than the method's token input are now prefilled in chunks
instead of rejected, for both in-graph and off-graph packages.

--interactive chats on stdin with the Llama 3 chat template, keeping the
conversation in the off-graph KV cache so each turn feeds only its new
tokens. /reset clears the cache, a full cache starts a new conversation,
and Ctrl-D quits.

This branch had an error being deployed

1 failed deployment
cadence — 923c00ed Deployed Nov 14, 2026 by metascroy via vision-build / vision #38332
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants