From 6b74bcbe1ef034d0571828796c63af826d2a29dd Mon Sep 17 00:00:00 2001 From: "Lin, Soga" Date: Thu, 1 Oct 2026 03:39:25 +0000 Subject: [PATCH] [AMD][Diffusion] Route unquantized linears through aiter tgemm on ROCm Add SGLANG_USE_ROCM_AITER_TGEMM (default off). When set on ROCm, apply_unquantized_linear sends bf16/fp16 linears to aiter's tgemm.mm, which picks the backend (torch/hipBLASLt or Triton) per (M, N, K) from aiter's bf16_tuned_gemm CSVs. Mixed-dtype calls keep using F.linear. This covers every DiT linear in Wan (QKV/out projections, FFN, text and time embedders), which all reach ColumnParallelLinear/RowParallelLinear. Wan2.2-T2V-A14B, 720x1280, 33 frames, 8 steps, 1x gfx1250, median of 3 interleaved runs, with the companion aiter Wan2.2 gfx1250 configs: | Stage | hipBLASLt | aiter tgemm | Speedup | | Encoding | 1.01 s | 0.96 s | 1.05x | | Denoise | 228.74 s | 64.41 s | 3.55x | | Decoding | 7.37 s | 7.00 s | 1.05x | | Total | 237.39 s | 72.64 s | 3.27x | Output video is bit-identical (same mp4 md5) to the F.linear path. --- python/sglang/multimodal_gen/envs.py | 4 ++++ .../multimodal_gen/runtime/layers/linear.py | 16 ++++++++++++++++ 2 files changed, 20 insertions(+) diff --git a/python/sglang/multimodal_gen/envs.py b/python/sglang/multimodal_gen/envs.py index 0a0ee5b0e92f..c16f9e7cc32c 100644 --- a/python/sglang/multimodal_gen/envs.py +++ b/python/sglang/multimodal_gen/envs.py @@ -108,6 +108,7 @@ SGLANG_USE_ROCM_CUDNN_BENCHMARK: bool = False SGLANG_USE_ROCM_VAE_CONV2D: bool = False SGLANG_USE_ROCM_VAE_CONV2D_BF16: bool = False + SGLANG_USE_ROCM_AITER_TGEMM: bool = False def get_default_cache_root() -> str: @@ -485,6 +486,9 @@ def _getter(): "SGLANG_USE_ROCM_VAE_CONV2D": _lazy_bool("SGLANG_USE_ROCM_VAE_CONV2D"), # ROCm: use BF16 compute for the Conv2D replacement (implies CONV2D=true) "SGLANG_USE_ROCM_VAE_CONV2D_BF16": _lazy_bool("SGLANG_USE_ROCM_VAE_CONV2D_BF16"), + # ROCm: route unquantized bf16/fp16 linears through aiter tgemm, which picks + # the backend per (M, N, K) from aiter's bf16_tuned_gemm CSVs + "SGLANG_USE_ROCM_AITER_TGEMM": _lazy_bool("SGLANG_USE_ROCM_AITER_TGEMM"), } # Add cache-dit Secondary Transformer Env Vars via programmatic generation to reduce duplication diff --git a/python/sglang/multimodal_gen/runtime/layers/linear.py b/python/sglang/multimodal_gen/runtime/layers/linear.py index ec8ae6b68a98..ff7884ed020b 100644 --- a/python/sglang/multimodal_gen/runtime/layers/linear.py +++ b/python/sglang/multimodal_gen/runtime/layers/linear.py @@ -11,6 +11,7 @@ from torch.nn.parameter import Parameter from sglang.kernels.kernel_api_logging import wrap_method_with_debug_kernel_once +from sglang.multimodal_gen import envs from sglang.multimodal_gen.runtime.distributed import ( divide, get_tp_group, @@ -42,6 +43,11 @@ logger = init_logger(__name__) IS_AMP_SUPPORTED = current_platform.is_amp_supported() +_USE_AITER_TGEMM = current_platform.is_rocm() and envs.SGLANG_USE_ROCM_AITER_TGEMM +if _USE_AITER_TGEMM: + from aiter.tuned_gemm import tgemm +# aiter gemm_a16w16 kernels index operands with int32 offsets and fault beyond this. +_AITER_TGEMM_MAX_NUMEL = 2**31 WEIGHT_LOADER_V2_SUPPORTED = [ "CompressedTensorsLinearMethod", "AWQMarlinLinearMethod", @@ -144,6 +150,16 @@ def apply_unquantized_linear( None if bias is None else bias.to(torch.float32), ).to(x.dtype) + if ( + _USE_AITER_TGEMM + and x.dtype in (torch.bfloat16, torch.float16) + and weight.dtype == x.dtype + and (bias is None or bias.dtype == x.dtype) + and x.numel() < _AITER_TGEMM_MAX_NUMEL + and x.numel() // x.shape[-1] * weight.shape[0] < _AITER_TGEMM_MAX_NUMEL + ): + return tgemm.mm(x, weight, bias) + return ( F.linear(x, weight, bias) if IS_AMP_SUPPORTED or bias is None