[AMD] Enable FP4 indexer for Deepseek V4

[RolaoDenthu]

Co-authored-by: @AMD-yanfeiwang @1am9trash

This PR is originally based on the work of @AMD-yanfeiwang.

• Original FP4 indexer enablement on ROCm gfx950 and the early integration work are from #36581.
• We also rely on the AITER fix from @AMD-yanfeiwang.
  • aiter#5034: Merged. Cherry-pick at #37083.
  • aiter#5126: Not merged. Apply change in this PR.

Motivation

Enable the DeepSeek V4 FP4 indexer on gfx95x GPUs through AITER.

AITER’s FP4 paged-MQA kernels use a different K-cache layout from the existing fused SGLang path. The 64-byte FP4 payload and 4-byte UE8M0 scale for each compressed token are stored in separate preshuffled tensors:

• payload: [num_pages, 1, 4, 64, 16]
• scale: [num_pages, 1, 4, 64]

Supporting this layout requires more than switching the kernel dispatch. This PR adds the corresponding device-pool representation, routes the Q/K RoPE, Hadamard rotation, FP4 quantization, cache writes, and paged-MQA logits through AITER, and introduces reusable scheduling metadata for prefill, decode, and EAGLE MTP.

The cache sizing is updated to reflect the 68-byte FP4 indexer entry. HiCache is also made aware of the split representation by mirroring the payload and scale as separate page-level regions, ensuring that host backup and restore preserve AITER’s physical layout.

Modifications

ROCm dependency

• docker/rocm.Dockerfile
  • Patch fix aiter#5126.

FP4 kernel adapters

• python/sglang/kernels/ops/attention/dsv4/fp4_indexer_hip.py
  • Add the gfx95 AITER adapter for Q RoPE/Hadamard/FP4 quantization.
  • Add guarded page-table construction and decode/prefill schedule workspaces.
  • Dispatch decode and prefill FP4 paged-MQA logits.
  • Build per-forward K write positions and cache slots.
  • Dispatch fused RMSNorm, RoPE, Hadamard, FP4 quantization, and paged cache writes.
• python/sglang/kernels/ops/attention/dsv4/compress.py
  • Route HIP FP4 indexer cache writes through the AITER fused writer.
  • Pass the separate payload/scale buffers, precomputed RoPE tables, and prebuilt write metadata.

Attention runtime

• python/sglang/srt/layers/attention/dsv4/indexer.py
  • Select the AITER FP4 path on HIP while preserving the existing CUDA/DeepGEMM and FP8 paths.
  • Quantize indexer Q through AITER.
  • Read the split K payload/scale cache with AITER paged-MQA kernels.
  • Reuse per-forward decode and prefill workspaces.
  • Keep padded FP4 query rows at zero context length.
• python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py
  • Attach FP4 decode/prefill workspaces, K write metadata, and widened Q positions to DSV4 metadata.
  • Size target-verification metadata for draft tokens crossing a page boundary.
  • Derive sequence lengths and output cache locations independently for every EAGLE draft step.
  • Build decode schedules during graph capture while keeping prefill schedule construction outside capture.
  • Reuse page tables and schedules across all C4 layers in a forward pass.
• python/sglang/srt/layers/attention/dsv4/compressor.py
  • Maintain a BF16 mirror of the indexer RMSNorm weight required by the AITER K writer.
  • Refresh the mirror through the normal weight-loader path.
• python/sglang/srt/layers/attention/dsv4/compressor_v2.py
  • Emit BF16 compressed indexer K directly on HIP.
  • Select the split FP4 payload/scale cache.
  • Forward cached RoPE tables and per-forward write metadata to the fused writer.
• python/sglang/srt/layers/attention/dsv4/metadata.py
  • Skip unnecessary DeepGEMM metadata construction for HIP indexer backends.
• python/sglang/srt/models/deepseek_v4.py
  • Reuse the model-owned BF16 RoPE cos/sin cache for FP4 Q and K transforms.
  • Refresh the indexer’s RoPE references after module _apply operations.
• python/sglang/srt/arg_groups/serving_hook.py
  • Allow --enable-deepseek-v4-fp4-indexer on supported gfx95 GPUs.

Device memory and sizing

• python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py
  • Share the FP8/FP4 indexer bytes-per-token calculation.
  • Allocate AITER FP4 K payload and UE8M0 scale buffers separately.
  • Add payload/scale accessors for the attention and compressor paths.
  • Expose page-row views for cache transfer paths.
• python/sglang/srt/model_executor/pool_configurator.py
  • Account for the FP4 indexer’s 68-byte compressed-token footprint instead of the FP8 132-byte footprint.

HiCache integration

• python/sglang/srt/mem_cache/hicache_storage.py
  • Add a dedicated pool identifier for the FP4 indexer scale region.
• python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py
  • Resolve fused FP8 and split FP4 indexer layouts into page-contiguous regions.
  • Build separate host pools for FP4 payload and scale data.
  • Register the scale region as a KV-derived HiCache sidecar.
• python/sglang/srt/mem_cache/hybrid_cache/linker_pool_assembler.py
  • Export each FP4 indexer region through the direct external-linker device-pool group.
• python/sglang/srt/mem_cache/memory_pool_host.py
  • Mark grouped FP4 indexer rows as whole-page-only transfers.
  • Reject unsupported token-granular payload/scale transfers instead of copying with the fused-layout kernel.
• python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py
  • Register Mooncake component keys for the FP4 indexer scale pool.

Tests

• test/registered/kernels/ops/attention/test_fp4_indexer_hip.py
• test/registered/unit/mem_cache/test_mem_pool_host.py

Accuracy Tests

• sglang gsm8k 1319: 0.948
• inferenceX lm-eval:

  |Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
  |-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
  |gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9674|±  |0.0049|
  |     |       |strict-match    |     5|exact_match|↑  |0.9682|±  |0.0048|
  

Speed Tests and Profiling

Fixed length benchmark:

• input length: 68k
• output length: 200

InferenceX Agentic:

• conc=48, using hicache

Metric	Baseline (SA curve)	FP4 indexer	Delta
Throughput / GPU	20,251 tok/s	21,430 tok/s	+5.8%
TTFT (median)	887 ms	838 ms	-5.5%
GPU KV cache capacity	10,221,824 tokens	11,006,208 tokens	+7.7%
GPU cache hit rate	91.3%	93.0%	+1.7%
CPU cache hit rate	5.3%	3.8%	-1.5%

• conc=16, no hicache

Metric	Official baseline	FP4 indexer	Delta
Throughput / GPU	9,667 tok/s	9,918 tok/s	+2.6%
TTFT (median)	542 ms	501 ms	-7.6%
GPU KV cache capacity	10,755,584 tokens	11,170,816 tokens	+3.9%
GPU cache hit rate	96.7%	96.7%	0.0 pt
CPU cache hit rate	n/a	n/a	n/a

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.

---

CI States

Latest PR Test (Base): ❌ Run #33563869960
Latest PR Test (Extra): ❌ Run #33563869631
Latest PR Test (AMD ROCm 7.2): 🚫 Run #33563869924

[HaiShaw]

LGTM

[1am9trash]

#33237 broke amd v4 code path.
Fixed in fe6c7db.

[1am9trash]

Update branch for #37406 fix (15 mins ago).
All CI is broken by #35118 changes (5 hours ago).

[HaiShaw]

NV CI signal is clean.
https://github.com/sgl-project/sglang/actions/runs/33563869960/job/100169694068?pr=37353 is unrelated, same as https://github.com/sgl-project/sglang/actions/runs/33573528101/job/100072782014