Support vectorized_5d kv cache for HybridLinearKVPool in Qwen3.5

[apinge]

Motivation

This PR functionally supports vectorized_5d for HybridLinearKVPool in Qwen3.5 by using export SGLANG_AITER_KV_CACHE_LAYOUT=vectorized_5d.

Currently Dflash speculative decoding cannot work with this feature.

Modifications

Accuracy Tests

export SGLANG_DISABLE_CUDNN_CHECK=1
export SGLANG_USE_CUDA_IPC_TRANSPORT=1
export SGLANG_VLM_CACHE_SIZE_MB=8192 #阿里用0
export HSA_NO_SCRATCH_RECLAIM=1

export SGLANG_USE_AITER=1
export SGLANG_ROCM_USE_AITER_LINEAR_SHUFFLE=1
export SGLANG_ROCM_USE_AITER_LINEAR_FP8HIPB=1
export AITER_QUICK_REDUCE_QUANTIZATION=INT6
export USE_AITER_COMM=1

export SGLANG_USE_AITER_NEW_CA=false
export SGLANG_USE_IPC_POOL_HANDLE_CACHE=1
export HIP_GDN_SORT_IDX_BS=32768


export SGLANG_AITER_KV_CACHE_LAYOUT=vectorized_5d
# export SGLANG_USE_AITER_UNIFIED_ATTN=1
export TVM_FFI_DISABLE_TORCH_C_DLPACK=1 # pip uninstall torch-c-dlpack-ext

model=/model/Qwen3.5-397B-A17B-PTPC-FP8
python3 -m sglang.launch_server \
        --port 7080 \
        --model-path ${model} \
        --tp-size 8 \
        --reasoning-parser qwen3 \
        --tool-call-parser qwen3_coder \
        --enable-multimodal \
        --trust-remote-code \
        --chunked-prefill-size 32768 \
        --mem-fraction-static 0.9 \
        --max-prefill-tokens 32768 \
        --max-running-requests 128 \
        --attention-backend aiter \
        --mm-attention-backend aiter_attn \
        --kv-cache-dtype fp8_e4m3 \
        --cuda-graph-max-bs 128 \
        --linear-attn-backend aiter \
        --linear-attn-decode-backend aiter \
        --linear-attn-prefill-backend aiter \
        --watchdog-timeout 1200 \
        --page-size 64 \
        --disable-radix-cache  2>&1 | tee launch_qwen3.5-397B-fp8_tp8_disable_prefix_cache.sh.log

gsm8k

python3 benchmark/gsm8k/bench_sglang.py    --port 7080  --tokenizer-path /model/Qwen3.5-397B-A17B-PTPC-FP8  --max-new-tokens 4096 
100%|██████████| 200/200 [01:19<00:00,  2.53it/s]
Accuracy: 0.960
Invalid: 0.020
Latency: 79.184 s
Output throughput: 451.909 token/s

Speed Tests and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

[Copilot]

🟡 Changes recommended

Composite-pool detection misclassifies unified-memory buffers, and hybrid prefix reads bypass layer-transfer synchronization.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Adds vectorized 5D KV-cache support for Qwen3.5’s HybridLinearKVPool.

Changes:

• Resolves wrapped MHA sub-pools and layer mappings.
• Recursively detects vectorized 5D layouts.
• Supplies query_scale to AITER decode.

File summaries

File	Description
python/sglang/srt/layers/attention/aiter_utils.py	Adds hybrid sub-pool resolution and updates decode arguments.
python/sglang/srt/layers/attention/aiter_backend.py	Detects vectorized layouts through composite pools.

Review details

• Files reviewed: 2/2 changed files
• Comments generated: 2
• Review effort level: Balanced

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[apinge]

--enable-unified-memory conflicts with the 5d KV cache layout.
SGLANG_AITER_KV_CACHE_LAYOUT=vectorized_5d requires the MHA KV cache to use a true 5D physical tensor layout:

K: (num_blocks, H, D // X, page_size, X)
V: (num_blocks, H, page_size // X, D, X)

However, under --enable-unified-memory, UnifiedMHATokenToKVPool exposes only a 4d page-major view rather than a true 5d layout. In other words, it may inherit the kv_cache_layout="vectorized_5d" string, but the actual k_buffer /v_buffer are not 5D.

Therefore, --enable-unified-memory and vectorized_5d are currently incompatible.

I added an extra check to reject the 5d layout in this case.

I also add warning in server_args.py

[apinge]

Makes sense, I've applied the fix. As Copilot noted, this issue occurs during async kv cache transfers when --enable-hierarchical-cache and --enable-lmcache are enabled. Although we aren't using these options right now, it's better to fix it proactively.