Port _fused_fp8_set_kv_buffer_kernel

[apinge]

Motivation

This PR ports the _fused_fp8_set_kv_buffer_kernel from zejunchen-zejun#276

Modifications

Accuracy Tests

Origin Acc

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


model=/models/Qwen/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 \
        --disable-radix-cache 2>&1 | tee launch_qwen3.5-397B-fp8_tp8_prefix_cache_origin_ali2.sh.log

 python3 benchmark/gsm8k/bench_sglang.py    --port 7080       --tokenizer-path /models/Qwen/Qwen3.5-397B-A17B-PTPC-FP8  --max-new-tokens 4096 2>&1 | tee gsm8k.log
100%|██████████| 200/200 [01:48<00:00,  1.85it/s]
Accuracy: 0.945
Invalid: 0.030
Latency: 108.039 s
Output throughput: 357.444 token/s

Dflash Acc

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 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

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 2>&1 | tee gsm8k.log
100%|██████████| 200/200 [01:20<00:00,  2.49it/s]
Accuracy: 0.950
Invalid: 0.020
Latency: 80.414 s
Output throughput: 443.769 token/s

Speed Tests and Profiling

C1 8k in prefill before

### C1 8k in decode before

C1 8k in prefill after

C1 8k in decode after

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 ports an ROCm/AITER-specific fused Triton kernel (_fused_fp8_set_kv_buffer_kernel) and wires it into the KV cache write path to accelerate FP8 KV cache population (notably for head_dim=256, NHD layout).

Changes:

• Add additional eligibility checks for DFlash fused QKV weight slicing to avoid invalid/unsupported layouts.
• Add an ROCm FP8 KV-cache “fast path” in MHATokenToKVPool.set_kv_buffer, including a warmup initializer.
• Introduce a new Triton kernel module (aiter_fp8_kv_kernel.py) that fuses scale division + FP8 KV-cache writes.

Reviewed changes

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

File	Description
python/sglang/srt/speculative/dflash_utils.py	Tightens validation for when DFlash is allowed to slice KV weights from fused QKV weights.
python/sglang/srt/mem_cache/memory_pool.py	Adds ROCm FP8 fused KV write initialization and a fast-path in set_kv_buffer.
python/sglang/srt/layers/attention/triton_ops/aiter_fp8_kv_kernel.py	New Triton fused kernel + Python wrapper for FP8 KV buffer writes.

---

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

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated no new comments.

Suppressed comments (2)

python/sglang/srt/mem_cache/memory_pool.py:1744

• The fused FP8 KV write path accepts any float32 tensor scale, but the Triton kernel loads tl.load(k_scale_ptr) / tl.load(v_scale_ptr) without an offset, so it only supports a single-element (scalar) scale tensor. As written, a multi-element tensor scale would silently use only the first element and produce incorrect KV values. Tighten the gate to require numel()==1 for tensor scales (matching the per-tensor FP8 KV scale semantics elsewhere).

            or (
                isinstance(scale, torch.Tensor)
                and scale.device == cache_k.device
                and scale.dtype == torch.float32
            )

python/sglang/srt/layers/attention/triton_ops/aiter_fp8_kv_kernel.py:114

• fused_fp8_set_kv_buffer() doesn’t validate that the destination cache tensors are actually FP8 (or that cache_loc is an integer 1D index tensor). If a caller accidentally passes the uint8 storage view (used elsewhere for FP8 caches) or a non-integer cache_loc, the kernel will still run and will store incorrectly-typed data or read invalid indices. Add explicit dtype/shape validation to fail fast with a clear error.

    if k_cache.ndim != 3 or v_cache.ndim != 3:
        raise ValueError("AITER fused FP8 KV write expects 3D K/V cache tensors")
    if k.ndim != 3 or v.ndim != 3:
        raise ValueError("AITER fused FP8 KV write expects 3D K/V input tensors")
    if cache_loc.numel() == 0: