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