Skip to content

[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
NVIDIA:mainfrom
ruocheng-nv:perf/glm53-flash-hopper-cpp
Draft

ruocheng-nv wants to merge 2 commits into
NVIDIA:mainfrom
ruocheng-nv:perf/glm53-flash-hopper-cpp

Conversation

@ruocheng-nv

@ruocheng-nv ruocheng-nv commented Oct 28, 2026 •

Copy link
Copy Markdown
Collaborator

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_tllm was built with CUDA_SEPARABLE_COMPILATION ON (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 -> 7.85 ms per call, outputs bit-identical. Same FlashMLA commit and nvcc; only -rdc differs.

FP8 block-scale CUTLASS MoE: row expansion and activation

  • expandInputRowsFp8BlockScaleKernel binary-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).
  • doActivation always launched 256-thread blocks with eight channels per thread, so a 512-wide
    intermediate slice (tensor parallelism) left three quarters of each 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; 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

  • Existing FlashMLA sparse attention and CUTLASS FP8 block-scale MoE tests (outputs bit-identical
    for the FlashMLA change; the MoE kernels change scheduling only).
  • End-to-end GSM8K on GLM-5.3-Flash FP8 (H200 TP4) unchanged.

PR Checklist

  • PR description clearly explains what and why.
  • No API changes.

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 branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant