Skip to content

[AMD] Keep aiter's tuned GEMM for the DeepSeek-V4 compressors #41931

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
HaiShaw merged 3 commits into from
Oct 4, 2026
Merged

[AMD] Keep aiter's tuned GEMM for the DeepSeek-V4 compressors #41931

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
f38697d
Keep aiter's tuned GEMM for the DeepSeek-V4 compressors
RolaoDenthu Sep 29, 2026
2557b25
Merge branch 'main' into dsv4/restore-perf-dsv41
RolaoDenthu Oct 1, 2026
1277a97
Merge branch 'main' into dsv4/restore-perf-dsv41
kkHuang-amd Oct 3, 2026
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
8 changes: 7 additions & 1 deletion python/sglang/srt/layers/attention/dsv4/compressor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,12 +25,15 @@
)
from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
from sglang.srt.model_executor.forward_context import get_attn_backend
from sglang.srt.models.deepseek_v2 import _is_hip
from sglang.srt.models.deepseek_v2 import _is_hip, _use_aiter
from sglang.srt.runtime_context import get_exec
from sglang.srt.utils import add_prefix, is_npu, set_weight_attrs

_is_npu = is_npu()

if _use_aiter:
from aiter.tuned_gemm import tgemm

if TYPE_CHECKING:
from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
from sglang.srt.layers.attention.deepseek_v4_backend import DeepseekV4AttnBackend
Expand Down Expand Up @@ -454,6 +457,9 @@ def compute_kv_score(self, x: torch.Tensor, forward_batch: ForwardBatch):
def _compute_wkv_gate(self, x: torch.Tensor) -> torch.Tensor:
weight = getattr(self.wkv_gate, "weight", None)
if weight is not None:
if _use_aiter and weight.dtype == torch.bfloat16:
# aiter's tuned GEMM for these shapes; kv_score is bf16-rounded
return tgemm.mm(x, weight, otype=x.dtype).float()
return linear_bf16_fp32(x, weight)

from sglang.srt.layers.quantization.gguf import fused_mul_mat_gguf
Expand Down
Loading