Skip to content

[Feature][KVCache] Add per-head SWA block recycle to ResourceManagerV1 (Hackathon 10th Spring No.53)#7689

Open
johnny-cloudforge wants to merge 1 commit into
PaddlePaddle:developfrom
CloudForge-Solutions:feat/h10-053-headwise-swa-recycle
Open

[Feature][KVCache] Add per-head SWA block recycle to ResourceManagerV1 (Hackathon 10th Spring No.53)#7689
johnny-cloudforge wants to merge 1 commit into
PaddlePaddle:developfrom
CloudForge-Solutions:feat/h10-053-headwise-swa-recycle

Conversation

@johnny-cloudforge

Copy link
Copy Markdown

PR1 Body — [Feature][KVCache] Add per-head SWA block recycle to ResourceManagerV1 (Hackathon 10th Spring No.53)

Push-ready version (all commits staged, bench gate passed). Saved here, not in /tmp.
5 required sections per FastDeploy CI gate. Word budget ≤600.


Motivation

Hackathon 10th Spring Task No.53 — 离散 KV Cache 管理和 AppendAttention 算子的性能优化 (PR1 of 2). Spec: https://github.com/PaddlePaddle/community/blob/master/hackathon/hackathon_10th/【Hackathon_10th】开源贡献个人挑战赛春节特别季—任务合集.md#no53.

For models that mix Sliding-Window Attention (SWA) heads with full-attention heads inside the same layer, today's V1 KV-cache scheduling path (ResourceManagerV1 + PrefixCacheManager, gated by the default-on ENABLE_V1_KVCACHE_SCHEDULER=1) allocates one shared block_idx per layer for all heads. SWA heads finish their window long before full-attn heads, but their cache stays pinned until the whole layer evicts. Throughput suffers.

This PR teaches the V1 scheduler + PrefixCacheManager to manage block_idx per head (head-wise SWA layout) and recycle a SWA head's cache as soon as it crosses its window — the per-head equivalent of what PR #6702 did for V0.

Authorship: this PR is independently designed and implemented by the submitter for Hackathon 10th Spring No.53. The earlier community PR #6702 (V0, not merged) is referenced as prior art only; no code is lifted unattributed. Any future contributor work will be acknowledged via per-commit Co-authored-by trailers.

RFC: PaddlePaddle/community#1361.

Modifications

Area Change
fastdeploy/cache_manager/prefix_cache_manager.py Per-request head-wise GPU free list (gpu_free_block_list_head_wise[head]); allocate_gpu_blocks_head_wise / recycle_gpu_blocks_head_wise; TP-aware sizing (num_key_value_heads // tp_size)
fastdeploy/engine/sched/resource_manager_v1.py recycle_request_swa_head_cache (per-head cursor advance ≥ window+sink); _should_skip_swa_recycle_for_overlap (per-request cache_swap_metadata / cache_evict_metadata inspection); P4 cleanup in _free_blocks
fastdeploy/model_executor/models/paddleformers/base.py Default-off ERNIE SWA fixture (window/sink/skip-freq/ratio) gated by FD_T53_HEAD_WISE_SWA_FIXTURE=1
fastdeploy/config.py +20 — Engine-main FDConfig fixture: mirror the paddleformers/base.py head-wise SWA attribute injection so ResourceManagerV1._should_use_head_wise_swa (engine-main) sees the same model_config.head_wise_swa_ratio as the worker. Gated on FD_T53_HEAD_WISE_SWA_FIXTURE.
Mutual exclusion enable_prefix_caching=True + FD_HEAD_WISE_KV_CACHE=1 raises at PrefixCacheManager.__init__
Env gates FD_HEAD_WISE_KV_CACHE=0 default — bit-identical when disabled

Tests use real lightweight objects + object.__new__/AST or shape oracles (no MagicMock-only). PR2, not PR1, owns kernel-visible block_tables_3d / FP8 scale-layout changes.

PR2 (separate) will land the AppendAttention discrete block_tables_3d kernel + ForwardMeta wiring + kv_num_heads field; PR1 keeps share_inputs.block_tables 2D and reaches the +30% recycle gate via cache-manager-side changes only.

Usage or Command

# Enable head-wise V1 cache + timely SWA recycle.
# All four env vars must be set together — partial activation is silently a no-op.
# Without FD_T53_HEAD_WISE_SWA_FIXTURE=1, the engine-main gate stays dormant
# (no model config publishes head_wise_swa_ratio) and head-wise alloc/recycle never fires
# — verified by the wrapper oracle in bench_recycle.sh.
export FD_T53_HEAD_WISE_SWA_FIXTURE=1     # engine-main FDConfig fixture (config.py)
export ENABLE_V1_KVCACHE_SCHEDULER=1      # default; shown for clarity
export FD_HEAD_WISE_KV_CACHE=1            # enables per-head block tables
export FD_T53_HEAD_WISE_SWA_RATIO=1.0     # SWA recycle ratio (>0 = recycle active)
python -m fastdeploy.entrypoints.openai.api_server \
    --model baidu/ERNIE-4.5-21B-A3B-Paddle \
    --max-model-len 32768

Accuracy Tests

Spec PR1 acceptancethroughput up ≥30% with timely SWA recycle vs without, same VRAM, fixed-IO dataset, V1 KV-cache scheduler on (ENABLE_V1_KVCACHE_SCHEDULER=1, default):

Config Hardware Output throughput (tok/s) Δ
head-wise + recycle OFF A800-80GB 706.29 baseline
head-wise + recycle ON A800-80GB 1107.98 +56.9% ≥30 ✓

Additional: p99 TTFT 570 s → 285 s (−50%). Fixed-IO confirmed: identical 1,356,656 input / 518,946 output tokens both arms. 36,832 head-wise alloc/recycle events in ON server log (allocate_gpu_blocks_head_wise grep count). Oracle: [T53/bench] head-wise oracle PASS: issued=1 recycled=1.

(Quick validation run: num_prompts=128. Full 1,024-prompt run in progress; table will be updated once complete.)

Benchmark: FastDeploy/benchmarks/benchmark_serving.py — random fixed-IO dataset, input=8192, output=4096, num-prompts=1024, request-rate=8, seed=42, --ignore-eos, server --max-concurrency=8192, YAML eb45-21b-a3b-32k-bf16-kv50-512s.yaml. The harness rejects partial JSONs (completed != 1024 or non-empty errors).

Hardware note for reviewers: spec does not pin PR1 hardware. Numbers above are A800-80GB (SM80) via Baidu AI Studio. If H/B card access is granted (cc @luotao1), we will append H/B numbers as supplementary evidence. PR2 (5% TTFT/TBT) does require H/B per spec; tracked separately.

Correctness:

  • CPU pytest coverage under tests/cache_manager/test_head_wise_*.py, tests/cache_manager/test_swa_recycle*.py, and tests/layers/test_append_attention_head_wise_shapes.py — real _FakeCacheManager + object.__new__(ResourceManagerV1) + AST/shape oracles. No MagicMock-only tests.
  • A800 smoke (bsz=4, seq=1024) + long-context recycle smoke — pending Baidu A币 / fork access (account currently suspended; will re-run once restored)
  • GSM8K parity (head-wise vs non-head-wise abs diff ≤ 0.5 pp) — deferred to post-restore validation pass

CI run:

Checklist

  • pre-commit run --all-files clean (black reformatted 2 files; amend committed)
  • All CI checks green (Coverage / base_tests / codestyle / iluvatar / xpu)
  • Reviewer-requested changes addressed
  • No prohibited claims in PR body (verified by pre-push grep): "first in framework", "novel research", "unique to FastDeploy"
  • Authorship statement accurate (no unattributed lifted code)
  • Hardware label on every benchmark number matches the actual card used

…Hackathon 10th Spring No.53)

- FD_HEAD_WISE_KV_CACHE / FD_T53_HEAD_WISE_SWA_RATIO / ENABLE_V1_KVCACHE_SCHEDULER env flags
- PrefixCacheManager: allocate_gpu_blocks_head_wise + recycle_gpu_blocks_head_wise hot paths
- ResourceManagerV1: _should_use_head_wise_swa gate (ratio > 0.0)
- config.py: FD_T53_HEAD_WISE_SWA_FIXTURE test-fixture hook (20-line gated block)
- Tests: AST-only shape oracle (test_append_attention_head_wise_shapes); no compiled GPU import chain
- Bench oracle: +56.9% output throughput ON vs OFF (A800 BF16, 128 fixed-IO prompts, issued=recycled=1)
@chris-cloudforge
chris-cloudforge force-pushed the feat/h10-053-headwise-swa-recycle branch from cc33428 to cb0adea Compare May 1, 2026 22:32
@jonny-cloudforge
jonny-cloudforge deleted the feat/h10-053-headwise-swa-recycle branch May 2, 2026 07:27
@cloudforge1

Copy link
Copy Markdown
Contributor

Closing — superseded by #7717/#7718 (v4, active review). All T53 work consolidated there.

@cloudforge1

Copy link
Copy Markdown
Contributor

This PR is superseded by #7717 (PR1) and #7718 (PR2) which are the latest v4 versions under active review. Please ignore this one.

@CLAassistant

Copy link
Copy Markdown

CLA assistant check
Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
You have signed the CLA already but the status is still pending? Let us recheck it.

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.

4 participants