[yuan-luo]

Motivation

This PR is to fill gaps in FA4 path on Blackwell:

1. The softmax row max ignores an sm_103 instruction. Blackwell Ultra adds tcgen05.ld.red, a TMEM load that also returns a per-thread reduction of the loaded values. We can use it to get each row's max for free while reading S out of TMEM. Currently SGLang FA4 forward does not: on B300 every softmax step still reduces 128 values per row with FMNMX3.
2. ViT attention on B300 falls back to Triton. Use triton_attn as default vision attention on B300 (SM103) #25570 switched the default multimodal attention backend on SM103 from fa4 to triton_attn because "FA4 is not yet validated on B300". Measured on B300, FA4's kernel is 2–20x faster than the Triton kernel on ViT shapes (head dims 64/72/80/128) and matches it numerically, so B300 ViTs pay a large, avoidable cost.
3. The ViT CUDA graph runners do not support fa4. SGLANG_VIT_ENABLE_CUDA_GRAPH=1 only handles triton_attn and fa3; with fa4 (already the default on SM100) the server dies at startup with AssertionError: cu_seqlens must be a torch.Tensor. Changing the SM103 default without fixing this would break every B300 deployment that turned the ViT graph on. The graph is also what removes FA4's per-call host overhead for small images (see Speed Tests).

Modifications

• flash_fwd_sm100.py: on sm_103, build a second TMEM copy with tcgen05.LdRed32x32bOp(redOp=MAX). In softmax_step, tiles that need no masking (mask_fn is None) in kernels without score_mod/bias load S with it — one cute.copy_atom_call per 32-column chunk, each also returning that chunk's max — and pass the reduced max to the softmax. Masked tiles keep the plain load and the software max, because the hardware max is taken over the raw S. On the tiles where it is used the hardware max equals the software reduction exactly, so outputs are bit-identical to the current kernel. The first-iteration tile is always masked (mask_seqlen=True), so only steady-state tiles change.
• softmax.py: SoftmaxSm100.update_row_max accepts an optional precomputed tile_row_max; everything after the max (fast exp2, rescale threshold) is unchanged.
• interface.py: SGLANG_FA4_TMEM_LOAD_RED_MAX (default 1; 0 restores the FMNMX reduction). It is part of the compile key, so both variants can coexist in one process. The flag is ignored on non-sm_103 GPUs.
• vision.py: the platform default multimodal attention backend is fa4 for all compute capability 10.x (SM100 and SM103). Explicit --mm-attention-backend and per-model overrides (e.g. Kimi-K3-VL's own SM103 heuristic) are unchanged.
• vision.py (VisionFlash4Attention): accepts the graph runners' [cu_seqlens, max_seqlen] list. Unlike the fa3/triton wrappers it keys off the argument type, not the env var, so models without a graph runner keep passing tensors when SGLANG_VIT_ENABLE_CUDA_GRAPH=1 is set globally.
• vit_cuda_graph_runner.py, internvl_vit_cuda_graph_runner.py: pass fa4 the same layout as fa3. FA4 JIT-compiles a kernel variant the first time it sees a configuration, which must not happen inside stream capture; for windowed ViTs (Qwen2.5-VL) the generic runner therefore warms up both the windowed and the full-attention layout before capturing (it previously warmed up only block 0's layout). The InternVL runner already warms up every layer eagerly.
• test_vision_backend_selection.py: covers the CUDA default per compute capability (9.0→fa3, 10.0→fa4, 10.3→fa4, 12.0→triton_attn) and VisionFlash4Attention with list-form and tensor seqlens, with the graph env var on and off.

Accuracy Tests

All on 8x B300 (sm_103, 148 SMs), CUDA 13.0, torch 2.13.0, nvidia-cutlass-dsl 4.6.2.

FA4 kernel, bit-exactness. 16 configurations run on the unpatched tree and on this branch with the flag on; outputs and LSE are identical in every case (torch.equal): causal / non-causal, varlen, chunked prefill (q < k), GQA and pack_gqa, head dims 64 / 96 / 128 / (192,128) (incl. the 2-CTA path), sliding window, learnable sink, softcap (reducing load disabled there), split-KV, paged KV (page 128 TMA and page 64 cp.async), and fp8 e4m3 with descales. The compiled PTX (.target sm_103a) contains 4 tcgen05.ld.red instructions with the flag on and none with it off.

FA4 unit tests. test/registered/kernels/ops/attention/test_flash_attention_4.py: 330 passed, 24 skipped, 0 failed.

ViT attention, FA4 vs Triton. sglang's own VisionFlash4Attention and VisionTritonAttention on the same inputs: max abs difference 5e-4 – 1.6e-2 in bf16 across GLM-5.3 / DeepSeek-V4.1 shapes (d=64) and Qwen2.5-VL shapes (d=80, windowed and full attention), consistent with bf16 rounding.

End to end, MMMU validation (900 questions) Two servers from the same base commit, one per GPU:

model (serving config)	ViT	triton_attn (current B300 default)	fa4 (this PR)	same final answer	same-config baseline
Qwen2.5-VL-7B-Instruct (1 GPU)	eager	51.7%	51.8%	871 / 900 (96.8%)	eager vs graph, same backend: 887–891 / 900
Qwen2.5-VL-7B-Instruct (1 GPU)	CUDA graph	52.0%	51.6%	875 / 900 (97.2%)	(same as above)
DeepSeek-V4.1-Flash (TP4 + EP4, cookbook flags, chat mode; 32-layer d=64 ViT)	eager	67.8% / 66.9% / 67.2% (3 runs)	68.0%	830 / 900 (92.2%)	two identical triton_attn servers: 834 / 900 (92.7%)

Speed Tests and Profiling

All on B300 (sm_103, 148 SMs, 1 GPU per measurement unless noted). Kernel numbers are CUDA-graph timed unless marked eager; "eager" includes the Python wrapper, as an un-graphed ViT sees it.

tcgen05.ld.red.max in the FA4 forward, ViT shapes

Command: python benchmark/kernels/attention/bench_flash_attention_tmem_red_max.py. It calls FA4 the way VisionFlash4Attention does (varlen, non-causal, no LSE, 16 heads, bf16).

Each of 10 rounds times the flag off and on back to back, alternating which goes first. "off" and "on" are the best round of each. "paired" is the median of the per-round off/on ratios, and "wins" counts the rounds where on was faster. Outputs are bit-identical for every shape.

input	head dim	off	on	paired speedup	wins
8 images x 4096 tokens	80 (Qwen2.5-VL)	648.8 us	588.0 us	1.075x	10/10
8 images x 4096 tokens	72 (Qwen3-VL)	703.9 us	681.6 us	1.039x	10/10
8 images x 4096 tokens	64 (InternViT, GLM / DeepSeek-V4 ViT)	527.4 us	518.2 us	1.023x	8/10
64 images x 1024 tokens	80	438.2 us	415.7 us	1.048x	10/10
64 images x 1024 tokens	64	360.5 us	356.8 us	1.010x	10/10
1 image x 32768 tokens	80	4861.1 us	4572.3 us	1.058x	10/10
1 image x 32768 tokens	72	5217.1 us	5070.5 us	1.030x	10/10
1 image x 32768 tokens	64	3974.7 us	3887.8 us	1.022x	10/10

ViT attention kernel, fa4 vs triton_attn (the kernels behind the default switch):

Head-dim sweep, 16 heads, non-causal varlen, speedup of fa4 over triton_attn (FA4 from the base commit, i.e. without item 1; the two wrapper tables below are on this branch):

layout	d=64	d=72	d=80	d=128
8 images x 4096 tokens	5.4x	20.5x	7.9x	6.6x
1 image x 16384 tokens	5.4x	19.3x	7.8x	6.4x
64 images x 1024 tokens	4.2x	16.2x	6.1x	5.4x

GLM-5.3 / DeepSeek-V4.1 ViT shapes (16 heads, d=64), sglang's VisionFlash4Attention / VisionTritonAttention, per call:

input	fa4 eager	triton eager	eager ratio	fa4 graph	triton graph	graph ratio
1 image 224² (256 tok)	62.4 us	38.4 us	0.62x	7.0 us	14.4 us	2.05x
1 image 448² (1024 tok)	61.5 us	49.3 us	0.80x	13.1 us	44.4 us	3.38x
4 images 448²	62.0 us	98.3 us	1.59x	25.5 us	93.5 us	3.67x
16 images 448²	94.2 us	370.2 us	3.93x	90.7 us	364.8 us	4.02x
64 images 448²	361.2 us	1522.1 us	4.21x	357.0 us	1516.7 us	4.25x
1 image 896² (4096 tok)	75.9 us	357.7 us	4.71x	72.5 us	352.8 us	4.87x
4 images 896²	258.4 us	1351.9 us	5.23x	254.6 us	1347.1 us	5.29x
1 image 1792² (16384 tok)	904.4 us	5200.8 us	5.75x	944.3 us	5187.9 us	5.49x
8 mixed-size images (11136 tok)	124.9 us	610.0 us	4.88x	120.6 us	606.1 us	5.03x
video, 32 frames x 256 tok	61.2 us	62.1 us	1.01x	28.7 us	57.8 us	2.01x

Qwen2.5-VL-7B ViT (16 heads, d=80; windows of 64 patches, 28 windowed + 4 full-attention blocks), per layer and per image:

image	windowed layer fa4 / triton (eager)	full layer fa4 / triton (eager)	per image, eager: fa4 / triton	per image, ViT graph: fa4 / triton
448x448 (16 windows)	62.6 / 37.2 us	61.2 / 42.3 us	2.00 / 1.21 ms	0.38 / 0.55 ms
720p (84 windows)	62.3 / 64.9 us	138.5 / 783.8 us	2.30 / 4.95 ms	2.11 / 4.81 ms
1080p (180 windows)	147.5 / 150.6 us	478.0 / 3725.6 us	6.04 / 19.12 ms	5.98 / 18.79 ms
4K (700 windows)	977.7 / 561.5 us	7399.8 / 60559.9 us	57.0 / 258.0 ms	57.6 / 257.7 ms

Where fa4 loses: (a) eager calls on tiny inputs (a single image of ≤1024 tokens), where its ~60 us per-call host overhead exceeds the kernel time; the ViT CUDA graph removes it (448x448: 0.38 ms vs 0.55 ms per image), and SM100 already carries this trade-off with fa4 as its default. (b) Very many 64-token windows (Qwen2.5-VL at 4K, 700 windows): the windowed layer is 1.7x slower than Triton because FA4's 128-row tiles are half empty; the full-attention layers more than make up for it (4.5x per image). A smaller Q tile for short segments would be a follow-up, not part of this PR.

End to end, eager ViT (Qwen2.5-VL-7B-Instruct, 1 GPU per server, c=1, bench_serving --dataset-name image, fresh images per run):

request	triton_attn median TTFT	fa4 median TTFT
1 image 4K (2 runs)	943 / 1021 ms	752 / 868 ms (−15–20%)
1 image 1080p (3 interleaved runs)	235 / 169 / 211 ms	194 / 171 / 159 ms

At 1080p the per-image attention saving (~13 ms) is below the run-to-run noise of host-side image preprocessing (medians of the same server swing by ~60 ms), so it is not claimed from this table. GLM-5.3-Flash (TP4, 24-layer d=64 ViT) is neutral end to end within noise. Concurrent image benchmarks (c=16/32) were dominated by host-side image preprocessing (the same configuration varied 1.8x between runs) and are not used.

End to end with the ViT CUDA graph (SGLANG_VIT_ENABLE_CUDA_GRAPH=1, Qwen2.5-VL-7B-Instruct, c=1, fresh images per run). Preprocessing noise only ever adds latency, so the per-request p25 (2 interleaved runs of 40 / 20 requests) tracks the floor better than the median:

request	triton_attn + graph, p25 TTFT	fa4 + graph, p25 TTFT	difference	attention-only prediction (item 2)
1 image 448x448	38.3 / 38.6 ms	38.9 / 38.3 ms	none	−0.2 ms
1 image 1080p	181.9 / 182.3 ms	168.7 / 166.7 ms	−13 to −16 ms	−13 ms
1 image 4K	866.7 / 901.7 ms	673.9 / 747.5 ms	−150 to −190 ms	−200 ms

---

CI States

Latest PR Test (Base): ✅ Run #36254948407
Latest PR Test (Extra): ❌ Run #36254948158
Latest PR Test (AMD ROCm 10): ❌ Run #36254948317

[BBuf]

ruff-format is failing on this call; it wants it on one line.

[yuan-luo]

Revised.

[BBuf]

I think we can remove these two test functions since they don’t add much value.

[BBuf]

For kernel change, should we add some micro benchmark results or correcless results?

[yuan-luo]

Updated benchmark case and correctness test case with results updated in PR description.

[DarkSharpness]

some very tiny comments

[DarkSharpness]

Why use new environ flag? either expose to environ.py or remove it (if it always brings perf gain, we should remove it and make it default behavior)

[yuan-luo]

Moved environ flag to environ.py as envs.SGLANG_FA4_TMEM_LOAD_RED_MAX.

[DarkSharpness]

maybe remove this and directly use the envs.SGLANG_FA4_TMEM_LOAD_RED_MAX