minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[AMD]Qwen3.5 integration gfx950 fmha fp8 hd256 (#32046)

Browse Source
This commit is contained in:
amd-oshkarav
2026-08-03 00:25:31 -07:00
committed by GitHub
parent 45c00daa1b
commit 92999d84f4
1 changed files with 38 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/aiter_backend.py
+38
View File
@@ -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,