[GH-ISSUE #2227] Borrowed Prefill: lend the GPU's idle dense-weight VRAM to the prefill chunk — 3-4x faster hybrid prefill on 16 GB cards, same KV pool #1308

New issue

Open

opened 2026-10-29 14:51:06 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-29 14:51:06 +00:00

Owner

Originally created by @05yuki on GitHub (Oct 7, 2026).
Original GitHub issue: https://github.com/kvcache-ai/ktransformers/issues/2227

Sharing a patch set for the KT hybrid (CPU experts in kt-kernel, SGLang on the GPU) in case it is useful to others here, and to ask whether any of it is worth upstreaming.

The problem. With the GPU-streamed prefill, every chunk pays a fixed cost to move a layer's experts over PCIe, so a bigger chunk is faster. On a 16 GB card the chunk is capped by what decode leaves free: attention weights, hot experts, KV pool, CUDA graphs. Meanwhile, during a prefill, the decoder-layer weights that sit in VRAM do nothing until their layer runs.

What the patch does. Before a long prefill it pauses the decoder-layer weights with torch_memory_saver (physical memory released, virtual addresses kept, so the decode CUDA graphs stay valid), streams each layer's weights back from a pinned host copy just before that layer runs, and puts everything back before the next decode step. Prefill-only buffers live in a second region that is resident only during prefill. The KV pool is untouched by default: a calibration launch measures the chunk that fits, and a guard refuses any chunk that would shrink the pool.

Numbers (2x RTX 5070 Ti 16 GB, dual EPYC 7452, single stream, base and lend measured the same day with the same prompts):

Model	chunk	38K prompt	~110K prompt	decode	output
Qwen3.8-Flash-Next NVFP4	2048 -> 16384	1,337 -> 4,154 tok/s	1,367 -> 4,596	58.5 -> 59.3	bit-identical
Qwen3.5-35B-A3B bf16	2048 -> 32768	1,568 -> 6,500	1,533 -> 6,406	91-97 -> 98-100	bit-identical
MiMo-V2.6-Flash MXFP4	2048 -> 22528	525 -> 1,970	496 -> 1,550	30-32 -> 32	top-20 unchanged
DeepSeek V4.1-Flash, 1M ctx	2048 -> 6144	255 -> 576	259 -> 596	25.5 -> 27.8	NLL z -1.9
DeepSeek V4-Flash-Vision	2048 -> 16384	601 -> 910	598 -> 849	38 -> 44	NLL z -0.6
GLM-5.3-Flash NVFP4 experts	2048 -> 14336	163 -> 326	155 -> 292	18.3 -> 20.0	NLL z -0.4 / -1.0

V4.1 prefills a 989K-token prompt to the end at the full 1,048,576-token KV pool on these two cards.

Repo, patches per tree, and the write-up: https://github.com/05yuki/kt-prefill-lend

Upstream question. The lend itself sits on our GPU-streamed prefill, which is not in upstream SGLang, so the whole thing does not transplant as a PR. Two pieces are independent of it and I would be glad to send either if there is interest:

kt-kernel's expert loader reads the checkpoint through the page faults of its memcpy into the memfd arenas, one readahead window at a time (about 1 GB/s on a drive that does 6). Resolving each layer's byte ranges from the safetensors headers and preading them in parallel, with the next layer read while this one copies, took a V4-Flash-Vision launch from 266 to 170 s and V4.1 from 740 to 243 s (with two more load fixes). Env-gated, ~150 lines in kt_kernel/utils/loader.py.
In deepseek_v4.py's load_weights, the ~100K routed-expert tensors that live on the CPU each go through the name remap, the expert-mapping scan and a weight_loader call that returns at once; skipping them right after the remap (using the same logical -> physical -> local mapping the loader applies) removes most of that Python time.

Happy to open either as a PR against the branch you prefer.

Originally created by @05yuki on GitHub (Oct 7, 2026). Original GitHub issue: https://github.com/kvcache-ai/ktransformers/issues/2227 Sharing a patch set for the KT hybrid (CPU experts in kt-kernel, SGLang on the GPU) in case it is useful to others here, and to ask whether any of it is worth upstreaming. **The problem.** With the GPU-streamed prefill, every chunk pays a fixed cost to move a layer's experts over PCIe, so a bigger chunk is faster. On a 16 GB card the chunk is capped by what decode leaves free: attention weights, hot experts, KV pool, CUDA graphs. Meanwhile, during a prefill, the decoder-layer weights that sit in VRAM do nothing until their layer runs. **What the patch does.** Before a long prefill it pauses the decoder-layer weights with torch_memory_saver (physical memory released, virtual addresses kept, so the decode CUDA graphs stay valid), streams each layer's weights back from a pinned host copy just before that layer runs, and puts everything back before the next decode step. Prefill-only buffers live in a second region that is resident only during prefill. The KV pool is untouched by default: a calibration launch measures the chunk that fits, and a guard refuses any chunk that would shrink the pool. **Numbers** (2x RTX 5070 Ti 16 GB, dual EPYC 7452, single stream, base and lend measured the same day with the same prompts): | Model | chunk | 38K prompt | ~110K prompt | decode | output | |---|---|---|---|---|---| | Qwen3.8-Flash-Next NVFP4 | 2048 -> 16384 | 1,337 -> **4,154** tok/s | 1,367 -> **4,596** | 58.5 -> 59.3 | bit-identical | | Qwen3.5-35B-A3B bf16 | 2048 -> 32768 | 1,568 -> **6,500** | 1,533 -> **6,406** | 91-97 -> 98-100 | bit-identical | | MiMo-V2.6-Flash MXFP4 | 2048 -> 22528 | 525 -> **1,970** | 496 -> **1,550** | 30-32 -> 32 | top-20 unchanged | | DeepSeek V4.1-Flash, 1M ctx | 2048 -> 6144 | 255 -> **576** | 259 -> **596** | 25.5 -> 27.8 | NLL z -1.9 | | DeepSeek V4-Flash-Vision | 2048 -> 16384 | 601 -> **910** | 598 -> **849** | 38 -> 44 | NLL z -0.6 | | GLM-5.3-Flash NVFP4 experts | 2048 -> 14336 | 163 -> **326** | 155 -> **292** | 18.3 -> 20.0 | NLL z -0.4 / -1.0 | V4.1 prefills a 989K-token prompt to the end at the full 1,048,576-token KV pool on these two cards. Repo, patches per tree, and the write-up: https://github.com/05yuki/kt-prefill-lend **Upstream question.** The lend itself sits on our GPU-streamed prefill, which is not in upstream SGLang, so the whole thing does not transplant as a PR. Two pieces are independent of it and I would be glad to send either if there is interest: - kt-kernel's expert loader reads the checkpoint through the page faults of its memcpy into the memfd arenas, one readahead window at a time (about 1 GB/s on a drive that does 6). Resolving each layer's byte ranges from the safetensors headers and preading them in parallel, with the next layer read while this one copies, took a V4-Flash-Vision launch from 266 to 170 s and V4.1 from 740 to 243 s (with two more load fixes). Env-gated, ~150 lines in `kt_kernel/utils/loader.py`. - In `deepseek_v4.py`'s `load_weights`, the ~100K routed-expert tensors that live on the CPU each go through the name remap, the expert-mapping scan and a weight_loader call that returns at once; skipping them right after the remap (using the same logical -> physical -> local mapping the loader applies) removes most of that Python time. Happy to open either as a PR against the branch you prefer.