From f38697d9ecbbaeeb92220ad0e52f294d42514ba9 Mon Sep 17 00:00:00 2001 From: RolaoDenthu Date: Tue, 29 Sep 2026 04:41:01 +0000 Subject: [PATCH] Keep aiter's tuned GEMM for the DeepSeek-V4 compressors #41019 moved the aiter branch of linear_bf16_fp32 from tgemm.mm(...).float() to torch.mm(out_dtype=float32) for every caller. The same helper serves the DeepSeek-V4 c4/c128 compressors (Compressor._compute_wkv_gate, about 90 GEMMs per V4-Pro forward), which moved from aiter's tuned FlyDSL kernels onto hipBLASLt. At decode row counts that is 1.3-2x slower on MI355X (K=7168, graph replay): N M=8 tgemm / torch.mm M=28 tgemm / torch.mm 2048 8.6 / 11.0 us 10.2 / 11.2 us 1024 7.2 / 14.1 us 8.3 / 10.7 us 512 6.5 / 10.1 us 7.4 / 7.7 us or 0.13-0.55 ms per DSpark verify step at TP8 conc 4. With SGLANG_USE_AITER and a bf16 weight, route the compressor projection back through tgemm. kv_score is bf16-rounded as before #41019. linear_bf16_fp32 keeps its fp32 contract for its other callers, including the V4.1 compressors in dsv41_sparse.py. --- python/sglang/srt/layers/attention/dsv4/compressor.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/attention/dsv4/compressor.py b/python/sglang/srt/layers/attention/dsv4/compressor.py index 81cfd1af2cdd..d005a1b355bf 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor.py +++ b/python/sglang/srt/layers/attention/dsv4/compressor.py @@ -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 @@ -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