Skip to content

fix(pd): account for KV transfer pages in mem_fraction - #1495

Open
sufubao wants to merge 3 commits into
ModelTC:mainfrom
sufubao:fix/pd-kv-page-memory-budget
Open

fix(pd): account for KV transfer pages in mem_fraction#1495
sufubao wants to merge 3 commits into
ModelTC:mainfrom
sufubao:fix/pd-kv-page-memory-budget

Conversation

@sufubao

@sufubao sufubao commented Aug 25, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • reserve the future paged KV transfer buffer when auto-profiling token KV capacity on Prefill and Decode workers
  • share one buffer-shape calculation between memory profiling and the actual transfer allocation, including the DeepSeek MLA layout
  • fail early with actionable configuration guidance when the reservation leaves no token KV capacity

Motivation

The PD transfer process allocates its GPU page buffer after model initialization and KV-cache profiling, so mem_fraction could consume that space before the transfer process starts. With Qwen3.5-27B, 16 pages, page size 4096, 17 full-attention KV layers, 4 global KV heads, head dimension 256, and BF16, the unprofiled allocation is 4.25 GiB per GPU.

This change subtracts the exact allocation size from the automatic KV-cache budget. Normal and PD-master modes are unchanged, and an explicit max_total_token_num continues to bypass automatic profiling.

Validation

  • Black 21.12b0 check passed on modified files
  • flake8 6.1.0 passed on modified files
  • compileall and git diff --check passed
  • no new unit tests are included

Validation note

ACP GPU end-to-end validation was not run from the local macOS development environment.

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