diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index ec492bb4d..e8a63e8bb 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -43,6 +43,7 @@ if TYPE_CHECKING: try: from aiter import ( + flash_attn_varlen_fp8_pertensor_func, flash_attn_varlen_func, get_mla_metadata_info_v1, get_mla_metadata_v1, @@ -2337,6 +2338,43 @@ class AiterAttnBackend(AttentionBackend): if layer.sliding_window_size is not None and layer.sliding_window_size > -1: window_size = (layer.sliding_window_size, -1) + if ( + get_bool_env_var("SGLANG_AITER_FMHA_FP8_ASM", "False") + and is_gfx95_supported() + and forward_batch.forward_mode.is_extend() + and forward_batch.extend_prefix_lens_cpu is not None + and not any(forward_batch.extend_prefix_lens_cpu) + and window_size == (-1, -1) + and sinks is None + and self.logits_soft_cap == 0.0 + and layer.qk_head_dim == 256 + and layer.v_head_dim == 256 + and self.kv_cache_dtype == fp8_dtype + ): + q_c = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim) + k_c = k.contiguous().view(-1, layer.tp_k_head_num, layer.head_dim) + v_c = v.contiguous().view(-1, layer.tp_v_head_num, layer.v_head_dim) + fp8_q_descale = ( + layer.k_scale if layer.k_scale is not None else self.k_scale + ) + o = flash_attn_varlen_fp8_pertensor_func( + q_c.to(fp8_dtype), + k_c.to(fp8_dtype), + v_c.to(fp8_dtype), + fp8_q_descale, + k_descale, + v_descale, + self.qo_indptr[:bs0], + self.qo_indptr[:bs0], + self.forward_metadata.max_q_len, + self.forward_metadata.max_q_len, + softmax_scale=layer.scaling, + causal=True, + ) + if o.dtype != self.input_dtype: + o = o.to(self.input_dtype) + return o.view(-1, layer.tp_q_head_num * layer.head_dim) + if self.kv_cache_is_vectorized_5d: return forward_extend_vectorized_5d( self,