feat(qwen3): make KV page size configurable

[sparkzky]

What

Replace the hardcoded page_size = 16 in Qwen3's kv_budget_geometry() with a configurable knob threaded through Qwen3MemoryOptions. The new --kv-page-size CLI flag (default 16, only 16/64 valid per FlashInfer constraint) flows end-to-end: CLI → config → Qwen3MemoryOptions → kv_budget_geometry() → KvBudget.block_size → KvCacheManager → PagedKvLayout → FlashInfer kernel.

Closes #545.

Why

In vLLM-compat P/D mode (#540) the block hash granularity, seal granularity, and the 1:1 remote-block-to-GPU-page load mapping all derive from budget.block_size. Lifting the hardcoded 16 into a config knob allows the vLLM prefill peer to align to either granularity — and GLM natively uses 64-token pages.

Changes

• weights.rs: DEFAULT_KV_PAGE_SIZE=16 + VALID_KV_PAGE_SIZES=[16,64] constants; page_size field on Qwen3MemoryOptions with validate() enforcing the 16/64 constraint; kv_budget_geometry(&self, page_size) now parametric; 3 unit tests (valid/invalid/default).
• config.rs + main.rs: --kv-page-size CLI flag (default 16), threaded into Qwen3MemoryOptions.
• dynamo-backend: mirrored --kv-page-size arg + DEFAULT_KV_PAGE_SIZE import.
• dflash tests: updated Qwen3MemoryOptions::new call sites for the new 3rd param.

Verification

• cargo build --release — clean (0 warnings)
• 3 unit tests (page_size suite) — all pass
• page_size=16 (default): 362 blocks (814 MB), greedy output correct — " Paris. The capital of Germany is Berlin. The"
• page_size=64: accepted by CLI, passes validation; OOM on 12GB RTX 3060 during profiling (each block is 4x larger → higher temp alloc floor). Expected to work on ≥16GB cards.

Trade-offs (from #545)

• Larger pages coarsen prefix-cache reuse (up to page_size - 1 wasted tokens per boundary).
• 64-token pages quadruple per-(layer, block) transfer segment (64 KiB → 256 KiB).
• FlashInfer supports both 16 and 64; accuracy gates and perf baselines need re-run per page size.

[xiaguan]

@FeathBow could u help review this? thks

[FeathBow]

Thanks for picking up #545 — the direction is sound. What needs work is the state of the branch.

Blocking — the pushed tree doesn't build (it's already on post-#495 main, so the hunks need fixing, not a rebase):

• openinfer_qwen3_4b:: references in engine.rs, both dflash tests, and config.rs — that crate is openinfer-qwen3 since refactor(qwen3): rename crate openinfer-qwen3-4b -> openinfer-qwen3 #495.
• kv_page_size is gated on #[cfg(feature = "qwen3-4b")] while the server feature is qwen3, so the field is compiled out and main.rs:220 reads it → the default cargo build --release fails with "no field kv_page_size".
• Three malformed import blocks — a dangling }; at engine.rs:39, dflash_speculative_gate.rs:51, dflash_speculative_perf.rs:27 (cargo fmt --all --check catches these).
• Qwen3MemoryOptions::new at tp_concurrent_decode.rs:59 still passes 2 args — the last call site; your Default update already covers the ::default() ones.
• The flashinfer submodule pointer moved 57ba7eeb → d768c14e — unrelated to this PR and a downgrade; please restore it: git checkout origin/main -- openinfer-kernels/third_party/flashinfer.

The feature itself works: I applied just the kv_budget_geometry 16→64 change on top of main and ran it on a 48 GB card — output is token-identical to the 16-page baseline and the engine reports KvCapacity { total_blocks: 3823, block_size: 64 }, so block_size flows through correctly and your OOM was a VRAM limit, not a functional problem. (I also checked split-KV chunking — it derives from token counts, not pages, so no batch-invariance interaction.) Still owed once it compiles: hf_golden_gate at --kv-page-size 64 (my run is a smoke test, not a gate). Perf isn't free either — larger pages coarsen prefix-cache reuse and round each request up to whole blocks (up to 63 wasted tokens), and the decode kernel's page-size sensitivity is unmeasured — so keeping the default at 16 is right; treat 64 as opt-in for #540 until a 16-vs-64 bench exists.

On the OOM specifically: the memory profile gives every profiled decode row a full block, so both the profiling transient and the 64-block minimum floor scale ~4× with page size — those block-count constants were tuned for 16-token pages and are worth revisiting while you're here.

Minor:

• The load line KV cache (profiled): N blocks (M MB, ...) (weights.rs:1069) doesn't print the page size — please add it so --kv-page-size 64 is visible (a P/D page-size mismatch otherwise only surfaces as silent prefix-cache misses).
• The feat(qwen3): vLLM-prefill P/D — cross-engine KV compat, verified end-to-end #540 rationale (hash/seal/1:1 remote-block mapping deriving from block_size) isn't exercised by any test at 64 — a good follow-up once this lands.

[FeathBow]

LGTM. It seems I'm unable to trigger CI due to a settings issue, so this hasn't been validated by CI. the formal grant is yours @xiaguan :)

Also, I reviewed on an sm89 GPU (CUDA 12.9).

Environment: sm89 · x86_64 · CUDA 12.9 · Qwen3-4B

Startup log (blocks scale with page size, capacity unchanged):

• page 16 → 15296 blocks (34416 MB, page size 16)
• page 64 → 3824 blocks (34416 MB, page size 64) — 4x fewer blocks, same MB

The KvBudget -> KvCacheManager -> PagedKvLayout -> FlashInfer chain derives correctly from block_size.

Performance — decode (--no-prefix-cache, two-point TPOT)

Metric	page 16	page 64	Δ
single-stream TPOT	9.84 ms	9.81 ms	-0.3%
conc1 throughput	100.8 tok/s	101.0 tok/s	+0.2%
conc16 throughput	1378.6 tok/s	1360.4 tok/s	-1.3%
conc16 latency p50	1.490 s	1.494 s	+0.3%
conc16 latency p99	1.506 s	1.509 s	+0.2%
conc32 throughput	2378.3 tok/s	2384.4 tok/s	+0.3%
conc32 latency p99	1.710 s	1.705 s	-0.3%

All within noise — page size does not move decode.

Performance — prefix cache (cache on)

Re-send same prompt (warm hit):

prompt_tokens	cached@16	cached@64	warm TTFT@16	warm TTFT@64
69	64	64	13.1 ms	12.9 ms
137	128	128	13.0 ms	12.7 ms
273	272	256	12.6 ms	12.5 ms

Two-turn (A -> R, then A+R+suffix, turn2 = 253 tokens):

Page size	cached	warm TTFT
16	240	13.6 ms
64	192	13.6 ms

The granularity cost is real (up to 48 tokens, the "up to page-1 wasted" effect), but the extra prefill is sub-millisecond, so warm TTFT is unchanged.

Performance — fragmentation (KV util=0.30 -> 5.3 GB pool, short requests, high concurrency)

Scenario	page 16	page 64
KV pool	2364 blocks	591 blocks (both 37824 token slots)
64 concurrent / 128 reqs	0 rejections · 142.4 req/s	0 rejections · 142.3 req/s
128 concurrent / 256 reqs	0 rejections · 207.4 req/s	0 rejections · 201.5 req/s (-3%)

Same total token capacity; normal concurrency never reaches KV pressure, no rejections.

Verdict

64 is correct, accuracy-lossless, and carries no measurable serving cost on any of the three axes. Keeping 16 as the default with 64 opt-in is the right call.

One line worth adding to the flag help text (not a blocker): the memory profile gives each profiled decode row a full block, so the startup profiling transient scales ~4x with page size — that is why 64 ran out of memory on the smaller card. --kv-page-size 64 wants a >=16 GB card to start; worth documenting so it does not read as a bug.

[xiaguan]

LGTM — CPU checks are green (FeathBow's earlier CI-trigger issue resolved itself once the run went through), and @FeathBow already did the detailed pass. Thanks @sparkzky for closing out #545!