[DeepSeek-V4.1] Replace Hopper prefill Torch indexer fallback with triton

[vladnosiv]

Motivation

Based on #38798

V4.1 prefill path on Hopper currently reaches _low_ratio_index_topk_torch fallback.

This change replaces that Hopper prefill fallback with a Triton FP8 path.

Known limitation: the path is noticeably more optimal than the current torch-folback, but it still materializes fp32 scores for the entire context and then calls the top-k kernel.

Modifications

• Add an SM90 Triton path for FP4 index-K unpacking and FP8 prefill scoring
• Send its scores through the existing ragged top-k kernel (prefill-only change, decode path is unchanged)

Accuracy Tests

gsm8k

baseline (torch fallback)

Accuracy: 0.880
Invalid: 0.000
Latency: 38.305 s
Output throughput: 492.101 token/s

current changes

Accuracy: 0.900
Invalid: 0.000
Latency: 36.707 s
Output throughput: 510.689 token/s

Speed Tests and Profiling

Cold Prefill Benchmark (4xH200)

Input	Torch fallback	SM90 Triton	Throughput
2K	10,649 tok/s	14,341 tok/s	+34.7%
4K	13,779 tok/s	16,681 tok/s	+21.1%
8K	15,607 tok/s	17,871 tok/s	+14.5%
16K	15,255 tok/s	17,728 tok/s	+16.2%
32K	14,481 tok/s	17,535 tok/s	+21.1%
64K	13,100 tok/s	16,937 tok/s	+29.3%
128K	10,904 tok/s	15,855 tok/s	+45.4%
256K	8,186 tok/s	14,004 tok/s	+71.1%

Serve command (4xH200):

SGLANG_ENABLE_DSV41_ENGRAM_HOST_TABLE=1 \
SGLANG_DSV41_ENGRAM_HOST_TABLE_LAYOUT=private \
SGLANG_DSV41_ENGRAM_HOST_TABLE_PIN=1 \
sglang serve \
  --model-path deepseek-ai/DeepSeek-V4.1-Flash \
  --served-model-name deepseek-ai/DeepSeek-V4.1-Flash \
  --host 0.0.0.0 \
  --trust-remote-code \
  --tp 4 \
  --ep-size 4 \
  --enable-decoder-swa-bounded-replay \
  --chunked-prefill-size 8192 \
  --mem-fraction-static 0.80 \
  --reasoning-parser auto \
  --tool-call-parser auto \
  --flashinfer-mxfp4-moe-precision fp8

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 #34830001786
Latest PR Test (Extra): ❌ Run #34830001482
Latest PR Test (AMD ROCm 10): ❌ Run #34830001629