kvarn_k4v2_g128 throughput does not scale with batch size (Sinkhorn kernel JIT-recompiles per step under dynamic batching)

[sztlink]

Summary

Under concurrent serving, kvarn_k4v2_g128 aggregate throughput saturates at ~80 tok/s regardless of batch size, while fp16 KV and another KV-quant backend scale to ~2000 tok/s on the same machine. Single-stream decode is fine (~44 tok/s), so this is specific to batched/concurrent serving.

Scaling curve (Qwen3-4B, warm runs, aggregate output tok/s)

concurrency N	1	8	16	32	64	128
kvarn_k4v2_g128	27	80	78	83	83	81
fp16 KV	66	362	663	1087	(noisy)	2172

KVarN flatlines from N=8 onward; fp16 scales ~linearly. A second KV-quant backend on the same box/flags also scaled to ~1950 tok/s at N=128, so this is not a WSL / pin_memory artifact.

Root cause (your own jit_monitor flags it)

During the concurrent burst the engine logs:

WARNING [jit_monitor] Triton kernel JIT compilation during inference:
_sinkhorn_log_kernel. This causes a latency spike;
consider extending warmup to cover this shape/config.

The running-set shape changes every decode step under dynamic batching, so _sinkhorn_log_kernel recompiles continuously. The scheduler never runs more than ~22 concurrent requests with 60+ waiting and GPU KV-cache usage at ~3% (so KV capacity is not the bottleneck):

Avg generation throughput: 1.5 tokens/s, Running: 15 reqs, Waiting: 53 reqs, GPU KV cache usage: 2.6%

generation throughput oscillates between ~1.5 and ~277 tok/s as the kernel recompiles.

Warmup does not fix it

Running each concurrency level twice (a discarded warmup burst, then a measured burst) gives essentially the same throughput cold vs warm at every N (e.g. N=128: 76 cold vs 81 warm). So this is per-step recompilation driven by the changing running-set shape, not a one-time first-compile cost.

Env

RTX 4090 (WSL2, Ubuntu 24.04), Qwen3-4B fp16, vLLM dev build (v0.1.dev58+gd784952a9), flags: --enforce-eager --no-enable-prefix-caching --no-enable-chunked-prefill --gpu-memory-utilization 0.9 --max-model-len 4096 --max-num-seqs 256.

Suggested direction

Shape-specialize / bucket the Sinkhorn kernel over a fixed set of batch sizes (pad the running-set dimension to the next bucket) so the compiled kernel is reused across decode steps instead of recompiling each step. Happy to share the full probe script and per-N engine logs.

[Dmtrii-tesla]

Hi @sztlink, could you please clarify which second KV-quant backend you used for comparison (the one that scaled to ~1950 tok/s)?

[philippebich]

Thanks. Tomorrow I will fix it. It should be an easy fix.

Yes, it would be nice if you could share the script so I can better replicate numbers.

[sztlink]

@Dmtrii-tesla the second backend is TurboQuant (TheTom's vllm-turboquant fork), preset turboquant_k4v2_nc. fp16 KV is the third line.

@philippebich thanks for the quick turnaround. Full repro (probe + serve commands + per-N table): https://github.com/sztlink/turboquant-cuda-bench/blob/main/bench/kvarn-vs-turboquant-2026-06-07/THROUGHPUT.md

Short version: Qwen3-4B on a 4090, same flags for all three (--no-enable-prefix-caching --no-enable-chunked-prefill --max-model-len 4096 --gpu-memory-utilization 0.90 --max-num-seqs 256, CUDA graphs on). Fire N concurrent /v1/completions, measure aggregate tok/s with a monotonic clock. KVarN held ~96 tok/s from N=8 to N=128; TurboQuant and fp16 scaled to ~1920 and ~2150. Happy to re-run any config or model once your fix lands.

[philippebich]

It should be fixed. From my tests (using different hardware from yours):

N	KVarN (tps, main, default)	TurboQuant (tps)
1	153	143
8	1000	952
32	2567	2419
64	3423	3347
128	4443	3919

[sztlink]

Confirmed your #15 fix scales on native-class hardware. One platform gap from our side: on WSL2 (RTX 4090) throughput is still capped post-fix, flat ~41 tok/s under load while fp16 and TurboQuant scale to ~2150 / ~1921 on the same box. py-spy (60s, steady decode) traced it to two per-step host-device syncs that are cheap on native CUDA but ~100x on WSL's paravirtualized GPU path:

• _mixed_batch_path: int(prefill_seq_lens.max().item()) (60.5%), an .item() that drains the queue, serializing decode before prefill setup.
• _batched_flush: per-chunk torch.as_tensor(..., device=dev) (20.5%).

Opened a small PR for both: #16. On WSL it gives 4-5x and restores scaling (N=128: 41 -> 176; N=32: 36 -> 204), action-router fidelity unchanged (106/120, 0 errors). Numerically identical; it also cuts syncs on native.

Re-profiling the patched build, the remaining host-side cost is the allocator: block_table_tensor.tolist() in build (~55%) and the slots as_tensor H2D (~20%). That looks like the device-side decode path your comments call "Stage 4". Is that already underway? If not, I would be glad to help with a device-side allocator, with your guidance so we do not duplicate.

[sztlink]

Follow-up: I posted a correction on #16. The WSL throughput numbers I reported earlier in this thread were a measurement artifact on my side (the probe undercounted by ~10x); re-measured with the engine's own counter, the patch shows no material throughput change and KVarN already scaled under load after your fix here. Details in the PR.

[sztlink]

Found the actual WSL bottleneck. It is memory residency, not host-device syncs. Config-level fix included (4-5x), with your table reproduced on WSL.

Following up on my correction in #16: after retracting the sync claim I kept digging into why KVarN specifically was slow on WSL2 while TurboQuant on the same box was not. Root cause found, with a reproducible fix.

Symptom (engine-reported generation throughput, Qwen3-4B, RTX 4090, WSL2, 128 concurrent decodes):

• KVarN, --gpu-memory-utilization 0.90: ~692 tok/s
• TurboQuant fork, same flags, same box: ~3382 tok/s

Evidence chain:

1. During KVarN load: GPU 100% util but only ~104 W power draw and ~0% memory-bandwidth utilization (TQ: ~98% util, ~425 W, 33-68% mem). Kernels resident but stalled.
2. torch.profiler in-process, with TQ as control under the identical harness: the same cuBLAS GEMM kernels, same names, same call counts, run 50-130x slower in the KVarN process (e.g. ampere_fp16_s1688gemm_128x128: 49.8 ms avg vs 0.94 ms in TQ). So it is not KVarN's kernels (_sinkhorn_log_kernel is ~5% of GPU time); the plain model GEMMs are crawling.
3. nvidia-smi during the KVarN run: 24102 / 24564 MiB allocated (98%).

Mechanism: on WSL2 (WDDM), allocations beyond the GPU's resident limit do not fail; Windows silently spills pages to "shared GPU memory" (system RAM over PCIe). GEMMs touching spilled pages run at PCIe speed, which matches the 50-130x slowdown, the 104 W draw, and the near-zero DRAM traffic. KVarN crosses the line because weights + KV cache + the fp16 tail pool at gpu_memory_utilization=0.90 total ~98% of VRAM; TurboQuant has no tail pool and stays below it. On native Linux this cannot happen (over-allocation OOMs instead of spilling), which is why it is invisible on your hardware.

Fix: lower gpu_memory_utilization so weights + KV + tail pool stay resident. Knee mapped on the 24 GB card:

gpu_memory_utilization	KVarN, N=128 (tok/s)	note
0.90	692	spills (98% VRAM)
0.85	~800, unstable (60-1574)	still spills
0.80	3523	resident, full concurrency
0.70	2864	resident, smaller KV caps concurrency ~100

Output quality is unaffected, as expected for a memory-budget change: my 120-case action-trace probe scores 106/120 with 0 errors, identical to the 0.90 baseline.

Your table reproduced on WSL with the fix (KVarN at 0.80, TQ at 0.90, same box):

N	KVarN, your table (native)	TurboQuant, your table (native)	KVarN, ours (WSL, gmu 0.80)	TurboQuant, ours (WSL, gmu 0.90)
1	153	143	51	49
8	1000	952	378	372
32	2567	2419	1181	1043
64	3423	3347	2133	2197
128	4443	3919	3523	3382

Each WSL cell is the mean of steady-state 10-second engine-logger windows (ramp and drain samples excluded).

With the fix, KVarN matches TurboQuant at every concurrency on WSL, mirroring the parity you see on native (including KVarN slightly ahead at N=128). The remaining gap to your absolute numbers is uniform across both methods and shrinks from ~3x at N=1 to ~1.3x at N=128, consistent with fixed per-step launch overhead on WSL GPU-PV (we measured cudaLaunchKernel at ~247us vs ~5-10us native) on top of faster hardware on your side. Nothing KVarN-specific remains.

This also retroactively explains the earlier confusion in this thread: the slowness I was chasing was real, but it lived in memory residency, not in the host-side syncs I patched.

Possible upstream actions, if you think they are worth it: a README note for WSL users, or accounting the tail pool more conservatively when sizing the KV cache so the default config stays resident. Happy to test any approach on this setup.

[sztlink]

Measured both of today's additions (42899a5) on the same rig and model as above (RTX 4090, WSL2, Qwen2.5-7B-Instruct, gpu_memory_utilization 0.80, engine-counter method):

	k4v2_g128	k4v2_g64	k4v4_g128
KV capacity	201,856 tok	184,704 tok	120,320 tok
behavior (action-trace, 120 cases)	96/120	93/120	102/120
decode tok/s @ N=128, steady	~3005	~2650	~2877

(behavior runs re-run identical for g64 and k4v4; fp16 reference on this rig is 100/120)

Two observations:

1. k4v4 recovers full behavior parity (102/120, same as fp16-level) at 1.5x capacity. The behavior cost of k4v2 lives in the 2-bit V, consistent with the decode-horizon results earlier in this thread. For agent-style workloads where behavior is the hard constraint, k4v4 looks like the preset to recommend.

2. On this setup g128 dominates g64 on all three axes, which makes mechanical sense (more per-tile metadata and more flush cycles at 64). g64 reads like a compatibility option (page size 64 for hybrids / vLLM defaults) rather than a quality knob. Is that the intent? If g64 is expected to win anywhere (MLA, hybrid models, different head_dim), happy to test that config on this rig.

[philippebich]

Thanks for your help, it is very much appreciated!

Regarding your questions:

1. K4V4 is definitely an option I will consider.

2. The reason for G64 is that the standard page size in vLLM is quite small (e.g., 16). Smaller page sizes generally result in higher overhead, but they also provide better granularity and cache efficiency. More importantly, smaller page sizes typically improve prefix cache hit rates, which is why many users prefer them.

Given your experience, if you want to add a note in the README for WSL users, I will be happy to read it. Thanks!

[sztlink]

Done: #19. Thanks for the g64 context, the prefix-cache angle makes sense; added the page-size trade-off to my notes.