diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py old mode 100644 new mode 100755 index 97d6ee573..fd8661234 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -203,7 +203,10 @@ class AiterAttnBackend(AttentionBackend): # need to fall back to non-persist # only use mla_ps_kernel when fp8 kv_cache # for non-fp8 kv_cache on tp8, use non-persist kernel to avoid performance degradation - if self.num_head == 16 and self.kv_cache_dtype is not fp8_dtype: + # head_num=16 (tp8 perf issue), head_num=128 (unsupported, like tp1 or --enable-dp-attention with tp8-dp8) + if ( + self.num_head == 16 or self.num_head == 128 + ) and self.kv_cache_dtype is not fp8_dtype: _use_mla_ps_kernel = False fast_mode = False intra_batch_mode = False