Skip to content

perf(kernel): avoid layout conversions in nvfp4 marlin decode - #631

Merged
jason-fxz merged 1 commit into
mainfrom
perf/nvfp4-marlin-decode
Oct 8, 2026
Merged

perf(kernel): avoid layout conversions in nvfp4 marlin decode#631
jason-fxz merged 1 commit into
mainfrom
perf/nvfp4-marlin-decode

Conversation

@jason-fxz

@jason-fxz jason-fxz commented Oct 8, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Part of #629: fixes the Triton 3.8 slowdown of the NVFP4 MoE decode GEMV.

Call path: TritonNvfp4MoEKernel.apply (decode) -> fused_experts_decode_nvfp4_marlin -> _decode_gemm_marlin -> _decode_nvfp4_marlin_kernel, once for gate_up and once for down. Prefill, the activation kernel and the sum-reduce are unchanged. "Marlin" here is the name of FreeToken's Triton GEMV (int32 wide loads, deferred K reduction), not vLLM's Marlin backend.

On Triton 3.8 the K loop of _decode_nvfp4_marlin_kernel compiles with 14 convert_layout ops and 2 KB of shared memory, against 1 and 32 B on 3.6 (default tile; the same change shows up on sm_80, sm_86, sm_89, sm_90, sm_100 and sm_120 in a compile-only check).

Changes:

  • Every load in the K loop now gets the register layout of the int32 word tile: the fp8 block scales are loaded as a half tile and expanded with tl.interleave; the e2m1 codes are decoded with the existing _nvfp4_pair_f32 from nvfp4_linear instead of a LUT gather; with contiguous, 16-byte-aligned bf16/fp16 activations, the 8 activations of a word are read as two int64 values and unpacked. Other activations keep the per-element loads.
  • Fixed configs, no autotune: K >= 2048 uses BLOCK_N 8, BLOCK_KW 32, UNROLL 2; smaller K uses BLOCK_N 16, BLOCK_KW 16, UNROLL 2; 4 warps. K == 2048 (Qwen3.6 gate_up) moves to the first config (the threshold was K > 2048).

With the change the loop compiles to 1 convert_layout and 16 to 32 B of shared memory on all six targets under both Triton versions, with no spills.

Kernel time

µs, triton.testing.do_bench_cudagraph, best of two medians, cold weights over 8 rotating routings, GPU used by the benchmark only. old = main 377a3bc. 3.6 = torch 2.11.0 / Triton 3.6.0, 3.8 = torch 2.14.1 / Triton 3.8.0 (#630). H100 80GB HBM3 (driver 580.95.05) and RTX PRO 6000 Blackwell Server Edition (sm_120, driver 595.91.07).

shape GPU old@3.6 old@3.8 new@3.6 new@3.8
Qwen3.6-35B-A3B gate_up M1 H100 16.81 24.18 8.49 7.73
RTX PRO 6000 9.51 16.76 5.08 5.22
Qwen3.6-35B-A3B down M1 H100 7.01 11.29 5.24 4.67
RTX PRO 6000 5.00 8.55 3.13 3.24
GLM-4.7 gate_up M8 H100 395.68 949.87 281.39 260.34
RTX PRO 6000 336.89 648.57 324.70 308.62
whole op, Qwen3.6-35B-A3B M1 H100 28.52 40.43 18.17 17.61
RTX PRO 6000 20.13 29.17 12.41 12.03
whole op, Gemma-4-26B-A4B M1 H100 35.18 63.81 28.69 27.92
RTX PRO 6000 28.02 47.66 21.53 21.23

Over all 30 GEMM shapes (Qwen3.6-35B-A3B, Qwen3.8-Flash-Next, MiniMax-M2.5, Gemma-4-26B-A4B and GLM-4.7; gate_up and down; M = 1, 4, 8), relative to old@3.6:

GPU old@3.8 new@3.6 new@3.8
H100 1.44 to 2.40 0.51 to 0.83 0.46 to 0.80
RTX PRO 6000 1.56 to 2.32 0.53 to 1.04 0.55 to 0.99

On the RTX PRO 6000 under Triton 3.6, six MiniMax-M2.5 and GLM-4.7 shapes at M = 4 and 8 are 0.4% to 3.8% slower than old; three separate runs gave the same ratios. Under Triton 3.8 all 30 shapes are faster than old@3.6 on both GPUs.

All shapes, H100
case old36 new36 old38 new38 new36/old36 old38/old36 new38/old36
Qwen3.6-35B gate_up M1 N1024 K2048 16.81 8.49 24.18 7.73 0.50 1.44 0.46
Qwen3.6-35B gate_up M4 N1024 K2048 39.04 24.00 83.19 22.23 0.61 2.13 0.57
Qwen3.6-35B gate_up M8 N1024 K2048 69.95 43.56 159.27 40.13 0.62 2.28 0.57
Qwen3.6-35B down M1 N2048 K512 7.01 5.24 11.29 4.67 0.75 1.61 0.67
Qwen3.6-35B down M4 N2048 K512 20.34 14.53 36.37 13.62 0.71 1.79 0.67
Qwen3.6-35B down M8 N2048 K512 36.73 25.73 67.68 24.20 0.70 1.84 0.66
Qwen3.6-35B op M1 28.52 18.17 40.43 17.61 0.64 1.42 0.62
Qwen3.6-35B op M4 63.74 42.89 123.08 40.11 0.67 1.93 0.63
Flash-Next gate_up M1 N1280 K2560 20.43 15.33 40.97 14.53 0.75 2.01 0.71
Flash-Next gate_up M4 N1280 K2560 66.59 42.56 145.29 39.35 0.64 2.18 0.59
Flash-Next gate_up M8 N1280 K2560 128.06 80.05 284.08 73.89 0.63 2.22 0.58
Flash-Next down M1 N2560 K640 12.32 10.24 20.53 9.38 0.83 1.67 0.76
Flash-Next down M4 N2560 K640 35.71 28.31 64.28 26.47 0.79 1.80 0.74
Flash-Next down M8 N2560 K640 66.79 52.53 124.44 49.03 0.79 1.86 0.73
MiniMax-M2.5 gate_up M1 N3072 K3072 35.28 26.70 76.83 24.71 0.76 2.18 0.70
MiniMax-M2.5 gate_up M4 N3072 K3072 124.87 90.82 289.51 83.12 0.73 2.32 0.67
MiniMax-M2.5 gate_up M8 N3072 K3072 244.16 175.81 572.33 161.86 0.72 2.34 0.66
MiniMax-M2.5 down M1 N3072 K1536 24.01 15.47 42.87 14.61 0.64 1.79 0.61
MiniMax-M2.5 down M4 N3072 K1536 74.16 50.99 148.69 44.87 0.69 2.01 0.61
MiniMax-M2.5 down M8 N3072 K1536 140.51 94.56 292.42 86.15 0.67 2.08 0.61
Gemma-4-26B gate_up M1 N1408 K2816 18.30 14.18 38.53 13.35 0.77 2.10 0.73
Gemma-4-26B gate_up M4 N1408 K2816 59.11 44.07 128.59 40.52 0.75 2.18 0.69
Gemma-4-26B gate_up M8 N1408 K2816 113.36 82.84 254.40 76.09 0.73 2.24 0.67
Gemma-4-26B down M1 N2816 K704 12.19 9.78 20.30 9.74 0.80 1.67 0.80
Gemma-4-26B down M4 N2816 K704 35.41 26.15 67.65 25.96 0.74 1.91 0.73
Gemma-4-26B down M8 N2816 K704 66.39 48.43 130.98 48.16 0.73 1.97 0.73
Gemma-4-26B op M1 35.18 28.69 63.81 27.92 0.82 1.81 0.79
Gemma-4-26B op M4 99.28 75.57 201.86 71.97 0.76 2.03 0.72
GLM-4.7 gate_up M1 N3072 K5120 56.00 41.82 125.89 38.68 0.75 2.25 0.69
GLM-4.7 gate_up M4 N3072 K5120 201.83 144.94 479.71 134.55 0.72 2.38 0.67
GLM-4.7 gate_up M8 N3072 K5120 395.68 281.39 949.87 260.34 0.71 2.40 0.66
GLM-4.7 down M1 N5120 K1536 38.05 23.70 66.76 21.58 0.62 1.75 0.57
GLM-4.7 down M4 N5120 K1536 118.39 80.46 244.21 71.77 0.68 2.06 0.61
GLM-4.7 down M8 N5120 K1536 228.81 151.94 481.58 138.75 0.66 2.10 0.61
All shapes, RTX PRO 6000
case old36 new36 old38 new38 new36/old36 old38/old36 new38/old36
Qwen3.6-35B gate_up M1 N1024 K2048 9.51 5.08 16.76 5.22 0.53 1.76 0.55
Qwen3.6-35B gate_up M4 N1024 K2048 28.31 18.89 63.62 18.89 0.67 2.25 0.67
Qwen3.6-35B gate_up M8 N1024 K2048 51.98 33.55 110.30 33.40 0.65 2.12 0.64
Qwen3.6-35B down M1 N2048 K512 5.00 3.13 8.55 3.24 0.63 1.71 0.65
Qwen3.6-35B down M4 N2048 K512 12.83 7.95 25.76 7.81 0.62 2.01 0.61
Qwen3.6-35B down M8 N2048 K512 24.85 15.68 50.48 15.20 0.63 2.03 0.61
Qwen3.6-35B op M1 20.13 12.41 29.17 12.03 0.62 1.45 0.60
Qwen3.6-35B op M4 46.60 35.18 95.86 35.73 0.75 2.06 0.77
Flash-Next gate_up M1 N1280 K2560 16.97 9.98 29.27 10.03 0.59 1.72 0.59
Flash-Next gate_up M4 N1280 K2560 54.75 48.53 105.30 47.34 0.89 1.92 0.86
Flash-Next gate_up M8 N1280 K2560 105.85 89.75 202.36 87.51 0.85 1.91 0.83
Flash-Next down M1 N2560 K640 7.59 5.36 14.22 5.31 0.71 1.87 0.70
Flash-Next down M4 N2560 K640 26.00 22.38 50.04 21.95 0.86 1.92 0.84
Flash-Next down M8 N2560 K640 49.33 38.94 95.62 38.35 0.79 1.94 0.78
MiniMax-M2.5 gate_up M1 N3072 K3072 31.30 29.26 55.78 28.58 0.93 1.78 0.91
MiniMax-M2.5 gate_up M4 N3072 K3072 109.24 109.44 201.15 107.46 1.00 1.84 0.98
MiniMax-M2.5 gate_up M8 N3072 K3072 208.35 206.35 393.83 201.83 0.99 1.89 0.97
MiniMax-M2.5 down M1 N3072 K1536 17.64 12.40 33.56 11.55 0.70 1.90 0.65
MiniMax-M2.5 down M4 N3072 K1536 61.99 61.83 111.08 48.94 1.00 1.79 0.79
MiniMax-M2.5 down M8 N3072 K1536 138.21 139.91 215.19 97.05 1.01 1.56 0.70
Gemma-4-26B gate_up M1 N1408 K2816 15.06 9.02 25.62 8.95 0.60 1.70 0.59
Gemma-4-26B gate_up M4 N1408 K2816 47.22 32.14 90.30 31.40 0.68 1.91 0.66
Gemma-4-26B gate_up M8 N1408 K2816 91.69 57.15 175.47 56.69 0.62 1.91 0.62
Gemma-4-26B down M1 N2816 K704 7.05 4.96 15.02 5.08 0.70 2.13 0.72
Gemma-4-26B down M4 N2816 K704 22.19 15.26 51.52 15.69 0.69 2.32 0.71
Gemma-4-26B down M8 N2816 K704 43.69 28.79 100.36 30.29 0.66 2.30 0.69
Gemma-4-26B op M1 28.02 21.53 47.66 21.23 0.77 1.70 0.76
Gemma-4-26B op M4 75.87 59.44 147.99 58.67 0.78 1.95 0.77
GLM-4.7 gate_up M1 N3072 K5120 49.80 46.90 90.74 44.66 0.94 1.82 0.90
GLM-4.7 gate_up M4 N3072 K5120 178.16 184.97 330.68 175.86 1.04 1.86 0.99
GLM-4.7 gate_up M8 N3072 K5120 336.89 324.70 648.57 308.62 0.96 1.93 0.92
GLM-4.7 down M1 N5120 K1536 28.92 25.71 51.40 22.04 0.89 1.78 0.76
GLM-4.7 down M4 N5120 K1536 102.85 104.28 180.07 79.23 1.01 1.75 0.77
GLM-4.7 down M8 N5120 K1536 218.63 223.51 351.76 156.04 1.02 1.61 0.71

Correctness and tests

  • Against an fp32 dequant reference, old and new have the same max abs and rel error under both Triton versions on both GPUs, for every shape.
  • New vs old: at most 21 of 196608 output elements per shape differ, by at most 0.0039 (the final reduction order changed). New is bit-identical between Triton 3.6 and 3.8 on H100.
  • fp16, strided bf16 and fp32 activations (the per-element load path) stay within the bf16/fp16 output rounding of the fp32 reference on H100 under both versions.
  • tests/moe/test_nvfp4_backends.py and tests/kernels/test_e4m3_compat.py pass under both versions: 13 passed and 7 skipped on H100, 15 passed and 5 skipped on the RTX PRO 6000.

End to end

AIME25 problem 1, then a long-prompt probe: TTFT and greedy 256-token decode after an 8k-token and a 32k-token prompt, median of 3. old@3.8 is main with the dependency change of #630.

model, MoE strategy GPU decode tok/s old@3.6 old@3.8 new@3.8
Qwen3.6-35B-A3B-NVFP4, fused H100 8k / 32k 221.9 / 212.2 199.5 / 192.1 244.2 / 232.9
RTX PRO 6000 8k / 32k 224.6 / 211.5 209.7 / 195.7 238.9 / 224.1
Qwen3.6-35B-A3B-NVFP4, hybrid H100 8k 144.0 140.4 144.4
RTX PRO 6000 8k 159.6 156.5 163.7
Gemma-4-26B-A4B-NVFP4, hybrid H100 8k 120.0 116.3 119.6
RTX PRO 6000 8k 117.5 113.9 117.4

Every run of this branch answered AIME25 problem 1 correctly (70). Gemma-4 also runs Triton attention, whose Triton 3.8 change is handled in #632; with both changes Gemma-4 decodes at 121.7 tok/s at 8k on H100.

Not covered: timing on other GPUs (sm_80, sm_86, sm_89 and sm_100 were only compiled), ROCm.

This was referenced Oct 8, 2026
Closed
Merged
@jason-fxz
jason-fxz merged commit 7644984 into main Oct 8, 2026
3 checks passed
@github-actions github-actions Bot mentioned this pull request Oct 8, 2026
Draft
pretor pushed a commit to pretor/FreeToken that referenced this pull request Oct 8, 2026
pretor
Merge upstream/main: PR FlashML-org#631, FlashML-org#632, FlashML-org…
283a737 
…#633, FlashML-org#634, DeepSeek V4.1
This was referenced Oct 8, 2026
Open
Open
KarrAcaRn pushed a commit to KarrAcaRn/FreeToken-ByAI that referenced this pull request Oct 9, 2026
KarrAcaRn
Merge branch 'main' into next
6b9a04a 
Upstream's kernel perf PRs overlap our own FP4/KV work:
- mxfp4 split-K GEMV (FlashML-org#653): keep upstream's column-pair int16 reads, decode
  the codes with our bit-placement e2m1_code_f32 instead of the LUT.
- nvfp4 Marlin decode (FlashML-org#631): take upstream's kernel body (same fp16 bit
  trick plus register-layout and wide activation loads); import both
  _nvfp4_pair_f16 and _nvfp4_pair_f32.
- Triton extend attention (FlashML-org#632): keep upstream's mask-free tile ranges and
  exp2 math; quantized KV caches (fp8 codes with per-token scales, nvfp4) go
  through the new _extend_attend_kv_quant_tile inside the same loops, bf16
  stays on upstream's _extend_attend_tile. Decode keeps our per-kv-head grid
  (group % VALID_BLOCK_H != 0) plus upstream's 64-key sliding-window tiles.
- CUDA JIT (FlashML-org#638, FlashML-org#134): upstream's sm_90/ROCm PDL guards and kDLROCM device
  matchers, with our MSVC-safe __forceinline__ and .template spellings.
- docs/install_amd.md: keep both ROCm sections.

Not GPU-tested (no GPU on the box); CPU test runs were unreliable on the
current unstable host.

Assisted-by: Claude Opus 5.5
@jason-fxz jason-fxz mentioned this pull request Oct 10, 2026
Merged
trcwebdesign pushed a commit to trcwebdesign/FreeToken that referenced this pull request Oct 14, 2026
@jason-fxz
perf(kernel): avoid layout conversions in nvfp4 marlin decode (FlashM…
d67d76d 
…L-org#631)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@jason-fxz