perf(rocm): enable Liger fused RMSNorm/SwiGLU for AMD DFlash configs - #779
Draft
zhentaocc wants to merge 2 commits into
Draft
perf(rocm): enable Liger fused RMSNorm/SwiGLU for AMD DFlash configs#779zhentaocc wants to merge 2 commits into
zhentaocc wants to merge 2 commits into
Conversation
added 2 commits
August 20, 2026 02:49
… ROCm)
OnlineDFlashModel._forward_draft_blocks (shared by DFlash / Domino / DSpark)
pins the flex_attention kernel option {"BACKEND": "TRITON"} to keep every
DFlash-family batch on the general Triton Flex Attention kernel (AUTO can
otherwise route short-query batches to the flex-decoding kernel, whose config
set can be empty for DFlash's sparse BlockMask).
However, the "BACKEND" kernel option is only understood by the inductor
flex-attention lowering on torch>=2.11, where it is sanitized out of the
generated Triton constexprs. On older builds the string leaks into the kernel
as a bare identifier and compilation fails with NameError: 'TRITON' is not
defined. This affects any pre-2.11 torch, in particular current torch ROCm
wheels (2.9.x-based).
Gate the kernel option on the torch version instead of the vendor:
- torch >= 2.11 -> {"BACKEND": "TRITON"} (unchanged CUDA behaviour)
- older -> {"FORCE_USE_FLEX_ATTENTION": True}
FORCE_USE_FLEX_ATTENTION has been supported since torch 2.5 and selects the
same general Triton Flex Attention kernel (verified by dumping the generated
inductor code on torch 2.11: both options produce a byte-identical
triton_tem_fused_flex_attention kernel, no flex-decoding, no FLASH/CuteDSL).
This fixes every pre-2.11 build, ROCm included, while leaving the CUDA path
identical.
Update test_online_wrapper_forces_standard_triton_flex_backend to assert the
version-appropriate kernel options.
Turn on model.use_liger_kernel for the two AMD Qwen3.5-4B DFlash example
configs (offline + online). This routes the draft model's RMSNorm and
SwiGLU MLP through Liger's fused Triton kernels.
Validated on a single idle MI355X (gfx950), torch 2.9.1+rocm7.2.0, real
`specforge train` (offline, batch_size=2, S=2048, 200 steps, steady-state
mean over 19 post-warmup steps):
use_liger_kernel=false -> 12.68 samples/s, 0.623 s/step, 35.37 GiB peak
use_liger_kernel=true -> 13.97 samples/s, 0.565 s/step, 30.18 GiB peak
=> +10.2% throughput, -5.2 GiB (-15%) peak memory, loss parity to 4
decimals (step200 7.0105 vs 7.0101).
Depends on the flex-backend fix (gate BACKEND=TRITON on torch>=2.11): on
torch<2.11 the Liger path drives flex-attention through the inductor
autotune lowering, which crashes with NameError: 'TRITON' is not defined
without that fix.
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
1 participant
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.
Summary
Turns on Liger fused RMSNorm/SwiGLU (
model.use_liger_kernel: true) for the AMD ROCm DFlash configs. Validated as a clean win on MI355X (gfx950) with realspecforge train, not a profiler estimate.examples/configs/amd/qwen3.5-4b-dflash-offline.yamlexamples/configs/amd/qwen3.5-4b-dflash-online.yamlMeasured on MI355X (Qwen3.5-4B DFlash, offline, batch_size=2, S=2048, num_anchors=512)
Method: container
specforge-514(torch 2.9.1+rocm7.2.0), single card,max_steps=200, discard first 20 warmup steps, steady-state mean ofperf/global_samples_per_second(n=19); peak VRAM polled across all 8 cards (others idle). Same seed, loss checked step-by-step.use_liger_kernel=falseuse_liger_kernel=true+10.2% throughput, −5.2 GiB / −15% peak memory, loss parity.
Note: the CUDA-profiler-estimated FLCE (fused Linear+CE) tier was dropped — it is numerically correct but 2.5× slower on gfx950 (Liger's chunked recompute overhead outweighs the logits memory-bandwidth savings). Only RMSNorm/SwiGLU is enabled here.
Stacked on #773
Requires the flex-backend fix (#773, commit 50c3407): on torch<2.11,
use_liger_kernel=truetriggers the inductor autotune lowering path where the unpatchedkernel_options={"BACKEND":"TRITON"}leaks a bareTRITONidentifier →NameError: 'TRITON' is not defined. That commit appears in this diff until #773 merges; will rebase/retarget afterward.Test plan
use_liger_kernel=false(same seed, per-step)