Repository navigation
[None][perf] Build FlashMLA without -rdc and speed up FP8 block-scale MoE expansion/activation - #21178
Draft
ruocheng-nv wants to merge 2 commits into
Draft
[None][perf] Build FlashMLA without -rdc and speed up FP8 block-scale MoE expansion/activation#21178ruocheng-nv wants to merge 2 commits into
ruocheng-nv wants to merge 2 commits into
Conversation
The flash_mla_cpp_tllm target was built with CUDA_SEPARABLE_COMPILATION, i.e. nvcc -rdc=true. Nothing in FlashMLA needs cross-TU device linking (its own setup.py builds without -rdc), and -rdc makes ptxas spill in the attention kernels: sm90 sparse prefill (d_qk 512) 0 -> 2044 B spill stores, 296 B stack sm90 sparse fp8 decode (h64) 0 -> 220 B sm100 sparse prefill (h128) 0 -> 44 B sm100 sparse decode (h64) 0 -> 12 B H200, flash_mla_sparse_fwd with 13312 queries x 64 heads x 2176 top-k rows (d_qk 512): 17.07 ms -> 7.85 ms, outputs bit-identical. Same FlashMLA commit and nvcc; only the -rdc flag differs. Signed-off-by: Ruocheng Jia <ruochengj@nvidia.com>
Two of the CUTLASS MoE kernels around the FP8 block-scale grouped GEMMs were latency-bound at prefill sizes: - expandInputRowsFp8BlockScaleKernel binary-searched the expert offsets for every permuted row it moved. Each block now takes a contiguous run of permuted rows, searches once for the first row and advances the expert linearly (rows are sorted by expert). - doActivation always launched 256-thread blocks, but each thread covers eight channels, so a 512-wide intermediate slice (tensor parallelism) left three quarters of every block idle. The block is now sized to the row width in whole warps. GLM-5.3-Flash FP8, H200 TP4, 13K-token prefill (45 MoE layers, 106K expanded rows): row expansion 29.6 -> 18.9 ms, activation 13.7 -> 6.6 ms; TTFT 522 -> 507 ms at concurrency 1. GSM8K unchanged. Signed-off-by: Ruocheng Jia <ruochengj@nvidia.com>
This was referenced Oct 29, 2026
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
C++/CMake part of #20878 (GLM-5.3-Flash FP8 on Hopper). The Python changes are in #21179 and the
Triton FP8-activation MoE is in #21250. These changes are model-agnostic and do not depend on either.
Build FlashMLA without relocatable device code
flash_mla_cpp_tllmwas built withCUDA_SEPARABLE_COMPILATION ON(nvcc -rdc=true). Nothing inFlashMLA needs cross-TU device linking (its own
setup.pybuilds without-rdc), and-rdcmakesptxas spill in the attention kernels:
flash_mla_sparse_fwdwith 13312 queries x 64 heads x 2176 top-k rows (d_qk 512):17.07 -> 7.85 ms per call, outputs bit-identical. Same FlashMLA commit and nvcc; only
-rdcdiffers.FP8 block-scale CUTLASS MoE: row expansion and activation
expandInputRowsFp8BlockScaleKernelbinary-searched the expert offsets for every permuted row.Each block now takes a contiguous run of permuted rows, searches once for the first row and
advances the expert linearly (rows are sorted by expert).
doActivationalways launched 256-thread blocks with eight channels per thread, so a 512-wideintermediate slice (tensor parallelism) left three quarters of each block idle. The block is now
sized to the row width in whole warps.
expansion 29.6 -> 18.9 ms, activation 13.7 -> 6.6 ms; C1 TTFT 522 -> 507 ms. GSM8K unchanged.
Measurements were taken on the #20878 branch (older main); this branch is the same two commits
cherry-picked onto current main without conflicts.
Test Coverage
for the FlashMLA change; the MoE kernels change scheduling only).
PR Checklist