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