Skip to content

Add fused QK-Norm + RoPE kernel for Qwen3#1286

Open
zhangxin81 wants to merge 1 commit into
linkedin:mainfrom
zhangxin81:qk-norm-rope-v2
Open

Add fused QK-Norm + RoPE kernel for Qwen3#1286
zhangxin81 wants to merge 1 commit into
linkedin:mainfrom
zhangxin81:qk-norm-rope-v2

Conversation

@zhangxin81

Copy link
Copy Markdown

Summary

Add a fused Triton QK-Norm + RoPE path for Qwen3 and wire it into the optional Qwen3 monkey patch path.

This PR:

  • adds a new fused qk_norm_rope Triton op with forward/backward
  • exposes a functional wrapper and a Qwen3 attention forward replacement
  • adds an opt-in qk_norm_rope flag to apply_liger_kernel_to_qwen3
  • adds correctness tests and a profiling script
  • updates monkey-patch coverage and the README patching table

Testing Done

  • Hardware Type: NVIDIA H100 80GB HBM3 (from the provided benchmark report)
  • run make test to ensure correctness
  • run make checkstyle to ensure code style (targeted local ruff/import/format checks on changed Python files)
  • run make test-convergence to ensure convergence

Additional validation completed locally:

  • imported liger_kernel successfully in a scratch Python 3.12 venv
  • ruff check passed for all changed Python files
  • ruff format --check passed for all changed Python files
  • py_compile passed for the new Python files

Note: full make test / convergence runs were not executed in this shell environment because the default local Python environment did not have the CUDA/PyTorch stack installed.

Benchmark Summary

From the supplied performance report (report file intentionally not included in this PR):

Correctness

  • 55 tests passed
  • The fused path keeps the RMSNorm -> RoPE chain in fp32 and is measured against an fp32 gold reference
  • For bf16, the fused path was reported as not worse than the unfused baseline and often more accurate because it avoids extra rounding

Performance (bsz=4, seq_len=2048, H100 80GB)

BF16

  • qwen3_0.6b: forward 20.3x, full 7.2x, memory 577M -> 242M (58% reduction)
  • qwen3_4b/8b: forward 20.2x, full 10.8x, memory 1089M -> 402M (63% reduction)
  • qwen3_14b: forward 19.3x, full 13.2x, memory 1345M -> 483M (64% reduction)
  • qwen3_32b: forward 20.4x, full 16.7x, memory 2113M -> 723M (66% reduction)

FP32

  • qwen3_0.6b: forward 10.5x, full 7.2x, memory 770M -> 483M (37% reduction)
  • qwen3_4b/8b: forward 10.6x, full 10.5x, memory 1287M -> 803M (38% reduction)
  • qwen3_14b: forward 10.5x, full 9.9x, memory 1570M -> 964M (39% reduction)
  • qwen3_32b: forward 10.5x, full 12.0x, memory 2434M -> 1444M (41% reduction)

Profiler summary

For the qwen3 4B/8B bf16 case, the report states the unfused path spent ~84.8 ms self CUDA time across many elementwise/cast kernels, while the fused path reduced that to ~10.2 ms dominated by the fused Triton forward/backward kernels.

Notes for reviewers

  • The benchmark script is included in the PR.
  • The benchmark/report markdown file is intentionally excluded, per request.
  • The Qwen3 fused path is opt-in via qk_norm_rope=False by default.

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