Skip to content

[AMD][Diffusion] Route unquantized linears through aiter tgemm on gfx1250 #41988

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yichiche merged 1 commit into from
Oct 15, 2026
Merged

[AMD][Diffusion] Route unquantized linears through aiter tgemm on gfx1250 #41988

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions python/sglang/multimodal_gen/envs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand Down
16 changes: 16 additions & 0 deletions python/sglang/multimodal_gen/runtime/layers/linear.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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)

@sushildubey171 sushildubey171 Oct 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we should add x.ndim >= 2 and x.is_contiguous() to avoid unsuported and copy call ?

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
Expand Down
Loading