Skip to content

[Engine] Fit a 512K context: give the prompt loan back to the K/V, keep big chunks staged - #4433

Closed
xvcezrqpbxtpzblrt wants to merge 1 commit into
Niko1221:mainfrom
xvcezrqpbxtpzblrt:engine/loan-release-027
Closed

xvcezrqpbxtpzblrt wants to merge 1 commit into
Niko1221:mainfrom
xvcezrqpbxtpzblrt:engine/loan-release-027

Conversation

@xvcezrqpbxtpzblrt

Copy link
Copy Markdown
Contributor

Problem

A 512K context on a 16 GB card could not fit in one run. The prompt path borrows part of the expert cache as a staging buffer (the loan), and the K/V's ceiling stopped at the loan's first slot: the borrowed region was excluded from the cache for the whole run. When the K/V needed those slots, the engine asked the driver for new VRAM, and on a card that has none left the request ended with and new VRAM, the cache being at its floor.

Result (measured)

k8v4, 16 GB RTX 3080 Ti Laptop, --max-context 524288, 510K prompt, 1.5 GiB reserve, --prefill 6144, staged chunks:

run prompt prefill K/V new VRAM
1 522,000 846.8 tok/s 5.30 GiB none (log checked)
2 522,000 831.7 tok/s 5.30 GiB none (K/V grown line has no and new VRAM)

The loan shrinks from 6,144 to 4,608 tokens and the freed slots become K/V cells; nothing is requested from the driver. 128K on a cool card: 1463.5 tok/s (previous build 1514.5, within run-to-run heat).

What changed

  • kvg_ensure, at its floor and before any new VRAM: asks the prompt path to re-lay out in the largest chunk whose freed granules cover the K/V's need, then unmaps the tail into the cache (kvg_fit_loan).
  • New option --prefill-stream-min N: chunks of N tokens or more stream every expert instead of staging the routed ones; an explicit value wins over the CPU share's own limit. Default behaviour is unchanged (the built-in 1,024 threshold).
  • Short contexts never take the new path (remain is 0), so their behaviour is the same as before.

Validation

  • 510K runs above, full logs kept (K/V grown without and new VRAM).
  • 32K: 1,672.7 tok/s; 128K: 1,463.5 tok/s (cool card).
  • Not run: 8 GB and 24 GB cards, multi-GPU (the path is single-participant), HIP, SYCL (this source has no such path), Metal (elastic K/V is off there).

Limits

  • CUDA and HIP only; the SYCL tree is not touched by this change.
  • Requires --kv-grow (the elastic K/V path).
  • On a 1.5 GiB reserve, int8 at 510K still asks for about 1.04 GiB of new VRAM; k8v4 does not. int8 needs a 1 GiB reserve to stay inside the cache.

@xvcezrqpbxtpzblrt

Copy link
Copy Markdown
Contributor Author

Superseded by #4606, which carries this change with the rest of the 0.2.7 branch. Closing so the review has one place.

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.

1 participant