Repository navigation
[mirror #321] common : --fit grows the ubatch of MoE models into spare VRAM, for up to 2.5x faster prompt processing - #8
Conversation
A MoE ubatch reads the weights of nearly every expert, and uploads them when the experts are in host memory, so prompt processing pays one pass over all experts per ubatch. Without -ub, MoE models now use the logical batch (up to 4096) instead of 512. -ub still overrides it. With a draft or MTP model the larger default applies only when the tensor overrides keep the experts in host memory. The draft context follows the main context's ubatch, so both compute buffers grow, and with the whole model on the GPU that can overrun VRAM.
Replaces the fixed MoE default from the previous commit. Without -ub, --fit now places the weights and the context at the default ubatch as before, then doubles n_ubatch towards n_batch (up to 4096) while the larger compute buffers still fit in the memory left over on every device. The fit never gives up context or device weights for a bigger ubatch, and a draft or MTP context that follows the main ubatch is measured with it. -ub still overrides it, and -fit off keeps the default.
The growth stays on by default. -fitub off (or LLAMA_ARG_FIT_UBATCH=off) keeps the ubatch at its default while --fit still places the weights and the context, so the behaviour can be switched off without pinning -ub 512.
-fitub becomes -fitu (as -fitp, -fitt, -fitc), the option moves next to the other --fit options, and the log line names it as the way to turn the growth off.
|
@codex review |
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ae1a7faafc
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| constexpr int64_t MiB = 1024*1024; | ||
|
|
||
| // a draft context follows the ubatch of the main context unless it was given its own | ||
| const bool extra_follows = extra != nullptr && extra->cparams->n_ubatch == cparams->n_ubatch; |
There was a problem hiding this comment.
When --spec-draft-n-ubatch is explicitly set to the target's initial ubatch (for example, 512), this equality check treats the draft setting as implicit and raises it along with the target. That discards the user's draft VRAM cap; track whether the draft ubatch was omitted instead of inferring it from equality.
Useful? React with 👍 / 👎.
…ubatch The growth inferred that the draft follows the target from equal ubatch values, so --spec-draft-n-ubatch set to the target's ubatch was measured at the grown size while the draft context is created at the user's value. Pass whether the draft follows from the caller instead.
|
Pushed Before/after on an M5 (Metal), granite-3.1-1b-a400m Q4_K_M as target and draft,
The user's draft cap was never lost at runtime; the fit over-counted the draft's compute buffer by about 132 MiB here, which could reject a growth step that fits. @codex review |
|
Codex Review: Didn't find any major issues. What shall we delve into next? Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
Mirror of unslothai#321 at
ae1a7faafc, base5c19a5a756(base/upstream-5c19a5a75). A sandbox for the Codex review and fork CI; fixes land on the source branch and are synced here.