kvcache-ai/sglang

Fork 0

mirror of https://github.com/kvcache-ai/sglang.git synced 2026-11-05 01:07:39 +00:00

[PR #98] fix(kt): reserve VRAM for the generic full-GPU prefill slot, and honour it #98

New issue

Open

opened 2026-10-29 21:07:10 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-29 21:07:10 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/kvcache-ai/sglang/pull/98
Author: @05yuki
Created: 9/25/2026
Status: 🔄 Open

Base: main ← Head: fix/kt-fullgpu-prefill-reservation

📝 Commits (1)

0917f65 fix(kt): reserve VRAM for the generic full-GPU prefill slot, and honour it

📊 Changes

3 files changed (+132 additions, -16 deletions)

View changed files

📝 python/sglang/srt/layers/moe/kt_ep_wrapper.py (+98 -6)
📝 python/sglang/srt/model_executor/model_runner.py (+20 -10)
📝 python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py (+14 -0)

📄 Description

Problem

--kt-gpu-prefill-token-threshold N with any --kt-method other than MXFP4 takes the generic full-GPU prefill path: the first prompt of at least N tokens builds one SharedFullContext, a full copy of one MoE layer's experts on the GPU, shared by every layer. Two things keep that from working on a card that is otherwise full:

The slot is never reserved. ModelRunner only asks for a reservation when kt_method == "MXFP4" (get_mxfp4_layerwise_prefill_reservation_bytes). For FP8, BF16, INT4 and the rest, the KV cache takes the whole mem_fraction_static budget, and the first qualifying prefill dies with CUDA OOM inside SharedFullContext._build_layers. The scheduler goes down with it.
The reservation is only read by the DSv4 pool sizer. Even for MXFP4, mxfp4_layerwise_prefill_reservation_bytes is subtracted only in the DeepSeek-V4 compressed-pool path. profile_max_num_token, which every other model uses, never subtracts it, so the number is computed and then ignored.

Seen on GLM-5-Next (FP8, TP2 on two 16 GB cards): KV 706,432 tokens at startup, then 1.12 GiB requested / 831 MiB free on the first 20k-token prompt.

Change

get_full_gpu_prefill_reservation_bytes() in kt_ep_wrapper.py: repeats what SharedFullContext._build_layers does (copy the layer, set the expert counts to the global total, quant_config.get_quant_method(...).create_weights(...)) on the meta device, and sums the parameter bytes. It gives the exact size for every quant method without allocating anything and without copying each method's shape formulas. On the setup above it returns 3.38 GiB per rank, which matches the hand count for FP8 block-128 (3.375 GiB).
KTEPWrapperMethod.create_weights registers one layer for the sizing when the threshold is set.
ModelRunner: for methods other than MXFP4, it stores that value in the same attribute the MXFP4 reservation already uses.
profile_max_num_token: subtracts the reservation, so every model honours it, not only DSv4.
If the context still cannot be allocated (the fraction is simply too small), the error names the two knobs (--mem-fraction-static, --kt-gpu-prefill-token-threshold 0). A module flag then stops the gate from choosing the full-GPU path again for later requests, instead of failing each one.

Results

Same setup, --kt-method FP8, threshold 512, chunk 2048:

	before	after
KV tokens at startup	706,432	207,168
first 20k-token prompt	CUDA OOM, scheduler exits	113.0 s (CPU-expert prefill: 238.0 s)

With the threshold unset nothing changes: no layer registers and the reservation stays 0.

🤖 Generated with Claude Code

https://claude.ai/code/session_01R57AidCMwH17PTFpX7HfQ5

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/kvcache-ai/sglang/pull/98 **Author:** [@05yuki](https://github.com/05yuki) **Created:** 9/25/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `fix/kt-fullgpu-prefill-reservation` --- ### 📝 Commits (1) - [`0917f65`](https://github.com/kvcache-ai/sglang/commit/0917f65a15a9b66c539727cdddca8935c5a8ccbd) fix(kt): reserve VRAM for the generic full-GPU prefill slot, and honour it ### 📊 Changes **3 files changed** (+132 additions, -16 deletions) <details> <summary>View changed files</summary> 📝 `python/sglang/srt/layers/moe/kt_ep_wrapper.py` (+98 -6) 📝 `python/sglang/srt/model_executor/model_runner.py` (+20 -10) 📝 `python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py` (+14 -0) </details> ### 📄 Description ## Problem `--kt-gpu-prefill-token-threshold N` with any `--kt-method` other than MXFP4 takes the generic full-GPU prefill path: the first prompt of at least N tokens builds one `SharedFullContext`, a full copy of one MoE layer's experts on the GPU, shared by every layer. Two things keep that from working on a card that is otherwise full: 1. **The slot is never reserved.** `ModelRunner` only asks for a reservation when `kt_method == "MXFP4"` (`get_mxfp4_layerwise_prefill_reservation_bytes`). For FP8, BF16, INT4 and the rest, the KV cache takes the whole `mem_fraction_static` budget, and the first qualifying prefill dies with CUDA OOM inside `SharedFullContext._build_layers`. The scheduler goes down with it. 2. **The reservation is only read by the DSv4 pool sizer.** Even for MXFP4, `mxfp4_layerwise_prefill_reservation_bytes` is subtracted only in the DeepSeek-V4 compressed-pool path. `profile_max_num_token`, which every other model uses, never subtracts it, so the number is computed and then ignored. Seen on GLM-5-Next (FP8, TP2 on two 16 GB cards): KV 706,432 tokens at startup, then `1.12 GiB requested / 831 MiB free` on the first 20k-token prompt. ## Change - `get_full_gpu_prefill_reservation_bytes()` in `kt_ep_wrapper.py`: repeats what `SharedFullContext._build_layers` does (copy the layer, set the expert counts to the global total, `quant_config.get_quant_method(...).create_weights(...)`) on the meta device, and sums the parameter bytes. It gives the exact size for every quant method without allocating anything and without copying each method's shape formulas. On the setup above it returns 3.38 GiB per rank, which matches the hand count for FP8 block-128 (3.375 GiB). - `KTEPWrapperMethod.create_weights` registers one layer for the sizing when the threshold is set. - `ModelRunner`: for methods other than MXFP4, it stores that value in the same attribute the MXFP4 reservation already uses. - `profile_max_num_token`: subtracts the reservation, so every model honours it, not only DSv4. - If the context still cannot be allocated (the fraction is simply too small), the error names the two knobs (`--mem-fraction-static`, `--kt-gpu-prefill-token-threshold 0`). A module flag then stops the gate from choosing the full-GPU path again for later requests, instead of failing each one. ## Results Same setup, `--kt-method FP8`, threshold 512, chunk 2048: | | before | after | |---|---|---| | KV tokens at startup | 706,432 | 207,168 | | first 20k-token prompt | CUDA OOM, scheduler exits | 113.0 s (CPU-expert prefill: 238.0 s) | With the threshold unset nothing changes: no layer registers and the reservation stays 0. 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_01R57AidCMwH17PTFpX7HfQ5 --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-29 21:07:10 +00:00

No labels

pull-request

No milestone

No project

No assignees

1 participant

Notifications

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

kvcache-ai/sglang#98

No description provided.

Rows
Columns