Recompiling FlashAttention below the source compiler

This is the how-to companion to the overview’s H100 case study: capture the official flash_attn forward kernel from a running process, configure the six-pass pipeline, apply it to the cubin, and validate the result in an unmodified process. Assumes sass2mlir is built and on your PATH and that you’ve done the capture walkthrough once.

One thing up front, because the overview says it too: two of the six passes apply real byte changes today; the four structural passes run in report mode while they are driven to full application. The config below reflects that honestly — nothing on this page pretends otherwise.

The official forward kernel, as shipped

The target is the CUTLASS FlashAttention-2 forward kernel exactly as the flash_attn package ships it — no source changes, no rebuild. As shipped it asks for 255 registers per thread, the hardware maximum, which caps occupancy at 12.5% and forces 85 register spills to local memory. It moves data with cp.async (48 LDGSTS) plus 144 LDSM, and multiplies with 512 per-warp HMMA tensor-core ops.

The reason those are the numbers is an ISA gate, not the source: the kernel was built for sm_90, but warpgroup HGMMA (wgmma) and TMA are only emittable for the architecture-specific sm_90a target. ptxas cannot emit them for plain sm_90, so an H100 runs Ampere-class code with its newest hardware idle.

The plan: don’t touch the source compiler at all. Lift the SASS ptxas actually emitted, re-optimize it for sm_90a with verified passes, and patch the result back into the original kernel slot.

Capture the kernel from a live process

sass2mlir intercept captures cubins in flight, so the flash_attn process needs no modification — point the output somewhere and filter to the forward kernel:

sass2mlir intercept \
  --out ./captures \
  --filter '*flash_fwd*' \
  -- python bench_flash.py --seqlen 8192 --hdim 64

ls captures/
# flash_fwd_kernel.cubin  flash_fwd_kernel.code  flash_fwd_kernel.mlir  run.csv

run.csv records the launch configuration per call — grid, block, and shared memory — which matters later: the pass config and the validation harness should reflect how the kernel is actually launched.

kernel,grid,block,smem_bytes
flash_fwd_kernel,"(128,8,16)","(128,1,1)",99328

The captured flash_fwd_kernel.cubin is the raw cubin, byte-identical to what the process loaded. If the cubin is already on disk instead — say, extracted with cuobjdump -xelf all flash_attn_2_cuda.so — sass2mlir harvest is the offline equivalent and produces the same .code + .mlir pair.

Six passes, in load-bearing order

The rewrite is an ordered pipeline of six conversion passes, and the order is not cosmetic: free the registers first, lift address and index math onto the uniform datapath, restructure into warpgroups, then convert the tensor-core and load paths that depend on that structure.

scalar-to-uniform is the most validated of the six — applied across the general corpus with byte-exact oracles — so it applies with confidence. ffma-deftz applies but stays opt-in: dropping flush-to-zero is only observable on denormal inputs, and whether that matters is the user’s call, not the toolchain’s. The structural four report where they would fire and mutate nothing — they are still being driven to full application, exactly as the overview’s status note says.

The whole thing is one config file. Passes are targeted per architecture and per function; every kernel the config doesn’t name stays byte-identical:

# flash.toml — the FlashAttention forward pipeline.
# Order is load-bearing: registers first, structure second,
# tensor-core and load-path conversions last.

[[pass]]
name = "spill-removal"
sm   = ["sm_90*"]
fn   = ["*flash_fwd_kernel*"]
mode = "report"          # 85 STL/LDL pairs identified, not yet rewritten

[[pass]]
name = "scalar-to-uniform"
sm   = ["sm_90*"]
fn   = ["*flash_fwd_kernel*"]
mode = "apply"           # corpus-validated, real byte changes

[[pass]]
name = "warp-specialization"
sm   = ["sm_90*"]
fn   = ["*flash_fwd_kernel*"]
mode = "report"          # producer/consumer split computed, not yet applied

[[pass]]
name = "hmma-to-hgmma"
sm   = ["sm_90a"]        # HGMMA is only emittable on the sm_90a re-encode
fn   = ["*flash_fwd_kernel*"]
mode = "report"

[[pass]]
name = "ldsm-to-tma"
sm   = ["sm_90a"]        # same: UTMALDG needs sm_90a
fn   = ["*flash_fwd_kernel*"]
mode = "report"

[[pass]]
name = "ffma-deftz"
sm   = ["sm_90*"]
fn   = ["*flash_fwd_kernel*"]
mode = "apply"           # opt-in: changes denormal handling

What a fully-applied run of the structural passes is working toward — the two conversions the sm_90 build could never emit:

- HMMA.16816.F32.F16  R8, R24, R40, R8   // per-warp MMA, operands in regs
+ HGMMA.64x64x16.F32.F16   // warpgroup async MMA, smem descriptors

- LDGSTS.E.128  [R12], [R4]   // cp.async global→shared
+ UTMALDG.2D  [desc], [R2]    // TMA bulk load + mbarrier

These are not textual swaps — HMMA and HGMMA are different execution models (warp-synchronous vs async plus mbarrier, register operands vs shared-memory descriptors), which is why the rewrite happens on typed IR and not on disassembly text.

Running the apply

One command: lift the captured cubin, load the plugin directory, apply the config, re-encode, patch. --keep-mlir writes the lifted and transformed modules so you can inspect what actually fired.

sass2mlir apply captures/flash_fwd_kernel.cubin \
  --passes build/lib/sass2mlir/passes \
  --config flash.toml \
  --keep-mlir out/mlir \
  -o flash_fwd_kernel.patched.cubin

apply: sm_90 → sm_90a flash_fwd_kernel — 6 passes loaded, 1 function targeted
[apply ] scalar-to-uniform: 41 uniform webs migrated (UR0–UR23, UP0–UP2), 318 ops rewritten
[apply ] ffma-deftz: 212 FFMA.FTZ → FFMA (ftz cleared where inputs proven non-denormal)
[report] spill-removal: 85 STL/LDL spill pairs would be eliminated
[report] warp-specialization: producer/consumer split computed — 2 warpgroups, 4 mbarrier handoff sites
[report] hmma-to-hgmma: 512 HMMA → 128 HGMMA.64x64x16 conversions identified
[report] ldsm-to-tma: 48 LDGSTS + 144 LDSM → 12 UTMALDG + mbarrier sites identified
patch: .text rewritten (530 instructions re-encoded), .nv.info updated
→ flash_fwd_kernel.patched.cubin

Every line marked [report] changed zero bytes — the sites are identified and counted, and the instructions lower verbatim from their $raw. The two [apply] passes invalidated the $raw on the 530 ops they touched, and those were re-encoded structurally by nvisa_encode from their semantic fields alone.

The kept modules show what fired. The before/after pair for the uniform migration — per-thread index math becomes uniform-datapath math:

// out/mlir/flash_fwd_kernel.lifted.mlir — as captured
%27 = sass.IMAD_0x91a %24, %25, %26
        : (!sass.reg, !sass.reg, !sass.reg) -> (!sass.reg)
%31 = sass.FFMA_0x223.FTZ guard(%28) %27, %14, %30
        : (!sass.reg, !sass.reg, !sass.reg) -> (!sass.reg)

// out/mlir/flash_fwd_kernel.transformed.mlir — after the two apply passes
%ur9 = sass.UIMAD_0x191 %ur7, %ur8, %ur6
        : (!sass.ureg, !sass.ureg, !sass.ureg) -> (!sass.ureg)
%31 = sass.FFMA_0x223 guard(%28) %27, %14, %30        // FTZ dropped
        : (!sass.reg, !sass.reg, !sass.reg) -> (!sass.reg)

sass2mlir uniform out/mlir/flash_fwd_kernel.lifted.mlir --fn '*flash_fwd*'
# uniform:   UR6 (param base), UR8 (head idx), UP0 (tile bounds check) ...
# non-uniform: %12 (lane id), %31 (FFMA result), %27 (address lane component) ...

Verifying the patched cubin

The output lands in the original kernel slot — rewritten .text plus updated .nv.info and relocations — with the ABI untouched: the patched kernel reads its parameters from the exact same const-bank offsets the official one did. Three checks, cheap to strict.

First, the byte-level oracles. roundtrip proves the patched cubin still lifts and lowers exactly; crosscheck proves the assembler path on the re-encoded instructions:

sass2mlir roundtrip flash_fwd_kernel.patched.cubin
# roundtrip: sm_90a flash_fwd_kernel: 1042/1042 instructions byte-identical (0 mismatches)

sass2mlir crosscheck out/mlir/flash_fwd_kernel.transformed.mlir
# crosscheck: flash_fwd_kernel — 1042/1042 instructions re-encode clean (0 mismatches)

Second, an independent look with the vendor tools — the diff should be exactly the transform you configured, and nothing else:

cuobjdump -sass captures/flash_fwd_kernel.cubin    > official.sass
cuobjdump -sass flash_fwd_kernel.patched.cubin     > patched.sass
diff official.sass patched.sass | head -12
# <   /*0220*/   IMAD R4, R5, R6, R7 ;
# ---
# >   /*0220*/   UIMAD.MOV.U32 UR4, UR5, UR6, UR7 ;
# <   /*03a0*/   FFMA.FTZ R10, R12, R14, R10 ;
# ---
# >   /*03a0*/   FFMA R10, R12, R14, R10 ;

Third, the one that matters: run an unmodified flash_attn process against the patched cubin. The validation harness interposes on module loads the same way sass2mlir intercept does, but instead of capturing the kernel at load time it substitutes the patched cubin into the kernel slot the process asks for. The test suite then compares against the official kernel’s outputs:

# run under the validation shim, which substitutes the patched cubin at load:
python -m pytest tests/flash_attn/test_fwd.py -q
# tests/flash_attn/test_fwd.py::test_fwd[seqlen=1024-hdim64] PASSED   max|Δ| = 0.0
# tests/flash_attn/test_fwd.py::test_fwd[seqlen=8192-hdim64] PASSED   max|Δ| = 0.0
# tests/flash_attn/test_fwd.py::test_fwd[seqlen=32768-hdim64] PASSED  max|Δ| = 0.0
# 214 passed

The zero diffs are the expected result, not luck: the uniform migration only rewrites webs the analysis proved warp-uniform, so it preserves values bit-for-bit. ffma-deftz is the one exception in principle — FTZ only differs on denormal inputs, which the suite’s distributions don’t generate — and it’s opt-in for exactly that reason.

What it buys

The measured H100 numbers live on the overview’s benchmarks section — same test there as here: official kernel as shipped versus the sass2mlir-patched cubin in the same slot, GPU kernel milliseconds per iteration. The shape of the result: at short sequences (≤2k tokens) the patched kernel runs a few percent behind (0.92–0.98×); from 8k tokens up it pulls ahead to 1.23–1.24×, the long-context regime the pipeline was built for. The tables themselves are not repeated here — the overview is the canonical source and stays in sync with metrics.ts.

Jetson Orin (sm_87) numbers are pending — that sweep is the on-device counterpart of this page and lands as it completes.