[NIXL][TURBOQUANT] Support turboquant in NIXL KV connector

[skavulya]

Add support for turboquant KV cache to NIXL workers

Purpose

Test Plan

Test Result

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[skavulya]

@xinyu-intel @jikunshang @xuechendi Please review

[gemini-code-assist]

Code Review

This pull request adds support for TurboQuant (TQ) in the KV transfer connector, specifically for handling packed K+V tensors and per-layer block length validation in models with mixed layer types. The review feedback highlights a logic error in register_kv_caches where tensor_size_bytes validation could fail for mixed-layer models. Additionally, the reviewer suggests refactoring the tp_ratio scaling validation in _validate_remote_agent_handshake to remove redundant code between the TQ and standard execution paths.

[NickLucche]

Hey thanks for the work.

Afaiu, the state of this feature is mostly useful for consumer/local deployments with very limited hbm, trading-off e2e perf @mgoin .
Therefore I am very hesitant to consider this feature here until we can prove benefits that target medium/large-scale high-end interconnected GPUs, which is the whole point of this connector.

Happy to re-evaluate when that's the case!

[skavulya]

Thanks @NickLucche for your insight. We are working on this feature as part of our efforts to enable heterogeneous disaggregated serving with Intel XPUs and Nvidia GPUs. We would be happy to collect benchmarking data for your consideration. There is a request for this feature on Dynamo repo on ai-dynamo/dynamo#8271

[xuechendi]

Thanks, @skavulya , overall, I think this PR is clear and LGTM.

Please provide background / tested_platform / test_script / accuracy check / coverage (full_attn, MLA, hybrid_attn) / perf impact

@NickLucche , PTAL

[skavulya]

@NickLucche @mgoin @xuechendi @vibhavagarwal5 I ran a benchmark on H200 NVL with Qwen/Qwen3-30B-A3B-Thinking-2507 (TP=1, 8K ISL/1K OSL). At max-concurrency 256, turboquant_4bit_nc TTFT is lower because of smaller queuing delays. However, the issue is turboquant's decode latency where TPOT is 7x slower than FP8. _tq_decode_stage1 is the bottleneck based on the pytorch profiles below.

export MODEL="Qwen/Qwen3-30B-A3B-Thinking-2507"

# prefill
CUDA_VISIBLE_DEVICES=0 \
UCX_NET_DEVICES=all \
VLLM_NIXL_SIDE_CHANNEL_PORT=5559 \
vllm serve "$MODEL" \
  --host 127.0.0.1 --port 8100 \
  --tensor-parallel-size 1 \
  --max-model-len 10240 \
  --gpu-memory-utilization 0.9 \
  --no-enable-prefix-caching \
  --kv-cache-dtype turboquant_4bit_nc \
  --kv-transfer-config \
      '{"kv_connector":"NixlConnector","kv_role":"kv_producer","kv_load_failure_policy":"fail",   "enable_permute_local_kv":"True","kv_connector_extra_config":{"enforce_handshake_compat":false}}'

# decode
CUDA_VISIBLE_DEVICES=1 \
UCX_NET_DEVICES=all \
VLLM_NIXL_SIDE_CHANNEL_PORT=5659 \
vllm serve "$MODEL" \
  --host 127.0.0.1 --port 8200 \
  --tensor-parallel-size 1 \
  --max-model-len 10240 \
  --gpu-memory-utilization 0.9 \
  --no-enable-prefix-caching \
  --kv-cache-dtype turboquant_4bit_nc \
  --kv-transfer-config \
      '{"kv_connector":"NixlConnector","kv_role":"kv_consumer","kv_load_failure_policy":"fail",   "enable_permute_local_kv":"True","kv_connector_extra_config":{"enforce_handshake_compat":false}}'

# proxy
  python tests/v1/kv_connector/nixl_integration/toy_proxy_server.py \
  --host 127.0.0.1 --port 9256 \
  --prefiller-hosts 127.0.0.1 --prefiller-ports 8100 \
  --decoder-hosts 127.0.0.1 --decoder-ports 8200

# benchmark
vllm bench serve --model=Qwen/Qwen3-30B-A3B-Thinking-2507  --host=127.0.0.1 --port=9256 --backend=vllm --dataset-name=random --random-input-len=8192 --random-output-len=1024 --num-prompts="$NUM_PROMPTS" --max-concurrency=${CONC} --request-rate=inf --ignore-eos

<style> </style>

kv_cache_dtype	ISL	OSL	concurrency	num_prompts	req_s	output_tok_s	total_tok_s	median_ttft_ms	median_tpot_ms	median_e2el_ms
bf16	8192	1024	4	40	0.45	457.76	4119.81	578	7.96	8768.73
auto	8192	1024	8	80	0.7	719.82	6478.42	562.09	10.57	11425.43
auto	8192	1024	16	160	1.05	1076.56	9689.02	554.16	14.1	14981.61
auto	8192	1024	32	320	1.51	1544.84	13903.59	593.34	19.75	20845.42
auto	8192	1024	64	640	2.12	2168.46	19516.15	651.49	28.03	29351.55
auto	8192	1024	128	1280	2.21	2259.76	20337.81	25476.4	35.25	59303.67
auto	8192	1024	256	2560	2.21	2263.95	20375.56	75391.26	35.32	111532.94
fp8	8192	1024	4	40	0.44	453.05	4077.43	512.81	8	8686.66
fp8	8192	1024	8	80	0.71	724.74	6522.7	514.84	10.33	11107.11
fp8	8192	1024	16	160	1.07	1092.3	9830.67	510.83	14.06	15015.27
fp8	8192	1024	32	320	1.58	1618.33	14564.93	548.06	18.8	19774.64
fp8	8192	1024	64	640	2.3	2359.6	21236.36	607.29	25.6	26810.75
fp8	8192	1024	128	1280	3.12	3199.82	28798.39	664.83	37.76	39299.73
fp8	8192	1024	256	2560	3.13	3207.28	28865.49	27873.88	47.56	77166.94
turboquant_4bit_nc	8192	1024	4	40	0.17	177.97	1601.77	520.26	21.89	23000.84
turboquant_4bit_nc	8192	1024	8	80	0.27	272.26	2450.3	558.59	28.79	30060.4
turboquant_4bit_nc	8192	1024	16	160	0.32	329.69	2967.2	634.42	47.78	49569.91
turboquant_4bit_nc	8192	1024	32	320	0.37	378.03	3402.27	815.16	83.45	86222.85
turboquant_4bit_nc	8192	1024	64	640	0.4	413.93	3725.41	1208.34	153.18	157892.79
turboquant_4bit_nc	8192	1024	128	1280	0.43	436.73	3930.61	2033.85	290.86	299606.63
turboquant_4bit_nc	8192	1024	256	2560	0.44	446.51	4018.63	3845.06	559.41	576198.13

FP8 KV cache Decoder Profile — 8.29ms/step

<style> </style>

Kernel	Time	%
Fused MoE	2.355s	28.80%
Cutlass GEMM	1.479s	18.10%
Flash Attention 3 (decode)	1.226s	15.00%
nvjet GEMMs	1.116s	13.60%
reshape_and_cache_flash (FP8)	190ms	2.30%

TurboQuant 4-bit KV Cache Decoder Profile — 21.84ms/step

<style> </style>

Kernel	Time	%
_tq_decode_stage1	11.191s	51.90%
sm80_xmma_gemm (FP32 dequant)	2.673s	12.40%
Fused MoE	1.951s	9.10%
Cutlass GEMM	1.281s	5.90%
_fwd_kernel_stage2	464ms	2.20%
nvjet GEMMs	1.120s	5.20%
sm80_xmma_gemm (small tile)	243ms	1.10%
_tq_fused_store_mse	128ms	0.60%
reduce_kernel	101ms	0.50%

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @skavulya.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mgoin]

@skavulya can we not specialize tq within the nixl connector? We also have nvfp4 kv cache in flashinfer and per-token int8 in triton attention, so it would be good to use a standard approach for all of these based on the block structure

[skavulya]

Hi @mgoin. Thanks for the review. I had implemented those changes to support scenarios with mixed full-attention and TurboQuant layers, since TurboQuant skips certain layers to preserve accuracy.

I rebased my PR and updated it to align with the recent NIXL connector changes in main, which have significantly simplified the implementation for your review

[mergify]

Documentation preview: https://vllm--40858.org.readthedocs.build/en/40858/

[mergify]

Hi @skavulya, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[skavulya]

@mgoin @NickLucche The nixl kv connector in main no longer needs special handling for turboquant because it deleted the assertion that enforced that tensor bytes were identical across all layers.

The only pending change from this PR the turboquant tests. I excluded AMD GPUs from the accuracy test the attention backend used by AMD CI for Hybrid SSM did not support turboquant kv_cache_dtype