Fix qkv slice in dflash

[apinge]

Motivation

Modifications

Accuracy Tests

397B

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


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 USE_HIP_LINEAR_ATTN=1
export SGLANG_USE_AITER_NEW_CA=false
export SGLANG_USE_IPC_POOL_HANDLE_CACHE=1

#export AITER_MOE_PADDING_SIZE=192
export HIP_GDN_SORT_IDX_BS=32768
export TVM_FFI_DISABLE_TORCH_C_DLPACK=1 # pip uninstall torch-c-dlpack-ext

model=/models/Qwen/Qwen3.5-397B-A17B-PTPC-FP8
model1=/models/Qwen/Qwen3.5-397B-A17B-Dflash
#  --kv-cache-dtype fp8_e4m3 \
# --disable-radix-cache
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 \
 --speculative-algorithm DFLASH \
 --speculative-draft-model-path ${model1} \
 --speculative-num-draft-tokens 16 \
 --speculative-draft-attention-backend triton \
 --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 \
 --mamba-scheduler-strategy extra_buffer \
 --kv-cache-dtype fp8_e4m3 \
 --page-size 64 2>&1 | tee qwen3.5-397B-fp8_tp8_dflash.log

 python3 benchmark/gsm8k/bench_sglang.py --port 7080 --tokenizer-path /models/Qwen/Qwen3.5-397B-A17B-PTPC-FP8  --max-new-tokens 4096                     
100%|████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:32<00:00,  6.22it/s]
Accuracy: 0.965
Invalid: 0.015
Latency: 32.179 s

35B

export SGLANG_DISABLE_CUDNN_CHECK=1
export SGLANG_USE_CUDA_IPC_TRANSPORT=1
export SGLANG_VLM_CACHE_SIZE_MB=8192 #阿里用0
export SGLANG_USE_AITER=1
export SGLANG_ROCM_USE_AITER_LINEAR_SHUFFLE=1 # internal branch
export SGLANG_ROCM_USE_AITER_LINEAR_FP8HIPB=1
export USE_AITER_COMM=1
export AITER_MOE_SMALL_BATCH=1
#export USE_HIP_LINEAR_ATTN=1
export SGLANG_USE_AITER_NEW_CA=false
export SGLANG_USE_IPC_POOL_HANDLE_CACHE=1
#export AITER_MOE_PADDING_SIZE=192
export HIP_GDN_SORT_IDX_BS=32768
export TVM_FFI_DISABLE_TORCH_C_DLPACK=1 # pip uninstall torch-c-dlpack-ext



model=/models/Qwen/Qwen3.5-35B-A3B-PTPC-FP8
model1=/models/Qwen/Qwen3.5-35B-A3B-DFlash

python3 -m sglang.launch_server \
 --port 7080 \
 --model-path ${model} \
 --tp-size 4 \
 --reasoning-parser qwen3 \
 --tool-call-parser qwen3_coder \
 --enable-multimodal \
 --trust-remote-code \
  --speculative-algorithm DFLASH \
 --speculative-draft-model-path ${model1} \
 --speculative-num-draft-tokens 16 \
 --speculative-draft-attention-backend triton \
 --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 \
 --disable-custom-all-reduce \
 --disable-radix-cache  2>&1 | tee qwen3.5-35B-bf16_tp4_Dflash.log

# python3 benchmark/gsm8k/bench_sglang.py --port 7080 --tokenizer-path /models/Qwen/Qwen3.5-35B-A3B-PTPC-FP8  --max-new-tokens 4096 
100%|████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [01:23<00:00,  2.41it/s]
Accuracy: 0.845
Invalid: 0.005
Latency: 83.108 s
Output throughput: 1973.023 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]

Pull request overview

This PR improves DFlash’s “KV-only projection” fast path by correctly handling (and explicitly validating) different fused-QKV weight layouts, including the ROCm/AITER-transposed weight case.

Changes:

• Tightens eligibility checks for slicing KV weights from a fused QKV linear layer (adds shape/layout validation and rejects AITER-transposed weights for the row-slice path).
• Adds a dedicated eligibility check for AITER-transposed QKV weights (column-slice path).
• Adds a new KV-only projection fast path in DFlashAttention.kv_proj_only() that slices AITER-transposed weights and runs the ROCm AITER GEMM.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
python/sglang/srt/speculative/dflash_utils.py	Adds stricter sliceability validation and introduces a new validator for AITER-transposed QKV weights.
python/sglang/srt/models/dflash.py	Adds an AITER-transposed fast path for KV-only projection using ROCm AITER GEMM.

Suppressed comments (1)

python/sglang/srt/speculative/dflash_utils.py:576

• Same issue as above: getattr(qkv_proj, "input_size") / getattr(qkv_proj, "output_size_per_partition") are used without defaults, which can raise and break DFlash initialization for models that don't expose these exact attributes. This should return (False, reason) rather than throwing so the code can safely fall back.

    expected_shape = (
        int(getattr(qkv_proj, "input_size")),
        int(getattr(qkv_proj, "output_size_per_partition")),
    )

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.