chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Fp8 prefill attn kernel integration (#18528)

Browse Source 
Co-authored-by: kkHuang-amd <wunhuang@amd.com>
This commit is contained in:
Thomas Wang
2026-02-11 15:23:48 +08:00
committed by GitHub
parent 2cc235e795
commit a8eef53dc4
1 changed files with 230 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

246
python/sglang/srt/layers/attention/aiter_backend.py
View File

@@ -19,6 +19,7 @@ from sglang.srt.layers.dp_attention import (
is_dp_attention_enabled,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.utils import is_gfx95_supported
if TYPE_CHECKING:
from sglang.srt.layers.radix_attention import RadixAttention
@@ -30,7 +31,11 @@ try:
flash_attn_varlen_func,
get_mla_metadata_info_v1,
get_mla_metadata_v1,
get_ps_metadata_info_v1,
get_ps_metadata_v1,
mha_batch_prefill_func,
mla_prefill_ps_asm_fwd,
mla_reduce_v1,
paged_attention_ragged,
)
from aiter.mla import mla_decode_fwd, mla_prefill_fwd
@@ -49,6 +54,11 @@ logger = logging.getLogger(__name__)
# Use aiter mla persist design for fp8-kv cache
_use_mla_ps_kernel = get_bool_env_var("SGLANG_AITER_MLA_PERSIST", "True")
# Use fp8 prefill only on gfx95
_use_fp8_prefill_attn = (
get_bool_env_var("SGLANG_AITER_FP8_PREFILL_ATTN", "True") and is_gfx95_supported()
)
# Persist
# fast_mode=True if _use_mla_ps_kernel else False
# intra_batch_mode=False if _use_mla_ps_kernel else True
@@ -308,6 +318,94 @@ class AiterAttnBackend(AttentionBackend):
dtype_kv=dtype,
)
def make_mla_prefill_ps_meta_data_buffer(
self, batch_size: int, max_qlen: int, qlen_granularity: int
):
(
(work_meta_data_size, work_meta_data_type),
(work_indptr_size, work_indptr_type),
(work_info_size, work_info_type),
(reduce_indptr_size, reduce_indptr_type),
(reduce_final_map_size, reduce_final_map_type),
(reduce_partial_map_size, reduce_partial_map_type),
) = get_ps_metadata_info_v1(
batch_size=batch_size,
num_head_k=self.num_kv_head,
max_qlen=max_qlen,
qlen_granularity=qlen_granularity,
)
device = self.device
work_metadata_ptrs = torch.empty(
work_meta_data_size, dtype=work_meta_data_type, device=device
)
work_indptr = torch.empty(
work_indptr_size, dtype=work_indptr_type, device=device
)
work_info = torch.empty(work_info_size, dtype=work_info_type, device=device)
reduce_indptr = torch.empty(
reduce_indptr_size, dtype=reduce_indptr_type, device=device
)
reduce_final_map = torch.empty(
reduce_final_map_size, dtype=reduce_final_map_type, device=device
)
reduce_partial_map = torch.empty(
reduce_partial_map_size, dtype=reduce_partial_map_type, device=device
)
return (
work_metadata_ptrs,
work_indptr,
work_info,
reduce_indptr,
reduce_final_map,
reduce_partial_map,
)
def make_mla_prefill_ps_meta_data(
self,
qo_indptr: torch.Tensor,
kv_indptr: torch.Tensor,
seq_lens: torch.Tensor,
work_metadata: torch.Tensor,
work_indptr: torch.Tensor,
work_info: torch.Tensor,
reduce_indptr: torch.Tensor,
reduce_final_map: torch.Tensor,
reduce_partial_map: torch.Tensor,
is_causal: bool = True,
):
gqa_ratio = self.num_head // self.num_kv_head
num_heads_k = self.num_kv_head
tile_q = 256
qhead_granularity = gqa_ratio
qlen_granularity = tile_q // qhead_granularity
kvlen_granularity = max(128, self.page_size)
block_size = self.page_size
qo_indptr_cpu = qo_indptr.to("cpu", dtype=torch.int32)
kv_indptr_cpu = kv_indptr.to("cpu", dtype=torch.int32)
seq_lens_cpu = seq_lens.to("cpu", dtype=torch.int32)
get_ps_metadata_v1(
qo_indptr_cpu,
kv_indptr_cpu,
seq_lens_cpu,
gqa_ratio,
num_heads_k,
work_metadata,
work_indptr,
work_info,
reduce_indptr,
reduce_final_map,
reduce_partial_map,
qhead_granularity=qhead_granularity,
qlen_granularity=qlen_granularity,
kvlen_granularity=kvlen_granularity,
block_size=block_size,
is_causal=is_causal,
)
def init_forward_metadata(self, forward_batch: ForwardBatch):
"""Init auxiliary variables for triton attention backend."""
@@ -587,15 +685,56 @@ class AiterAttnBackend(AttentionBackend):
spec_info=None,
)
kv_indices = self.mla_indices_updater_prefill.kv_indices
max_q_len = self.mla_indices_updater_prefill.max_q_len
qo_indptr = self.mla_indices_updater_prefill.qo_indptr
work_metadata = None
work_indptr = None
work_info_set = None
reduce_indptr = None
reduce_final_map = None
reduce_partial_map = None
if _use_fp8_prefill_attn:
tile_q = 256
qlen_granularity = tile_q // (self.num_head // self.num_kv_head)
(
work_metadata,
work_indptr,
work_info_set,
reduce_indptr,
reduce_final_map,
reduce_partial_map,
) = self.make_mla_prefill_ps_meta_data_buffer(
bs, max_q_len, qlen_granularity
)
self.make_mla_prefill_ps_meta_data(
qo_indptr,
qo_indptr,
forward_batch.seq_lens,
work_metadata,
work_indptr,
work_info_set,
reduce_indptr,
reduce_final_map,
reduce_partial_map,
is_causal=True,
)
self.forward_metadata = ForwardMetadata(
self.mla_indices_updater_prefill.kv_indptr,
kv_indices,
self.mla_indices_updater_prefill.qo_indptr,
self.mla_indices_updater_prefill.kv_indices,
qo_indptr,
self.kv_last_page_len[:bs],
self.mla_indices_updater_prefill.max_q_len,
max_q_len,
self.mla_indices_updater_prefill.max_kv_len,
work_metadata=work_metadata,
work_info_set=work_info_set,
work_indptr=work_indptr,
reduce_indptr=reduce_indptr,
reduce_final_map=reduce_final_map,
reduce_partial_map=reduce_partial_map,
)
else:
self.indices_updater_prefill.update(
@@ -1047,18 +1186,93 @@ class AiterAttnBackend(AttentionBackend):
):
extend_no_prefix = not any(forward_batch.extend_prefix_lens_cpu)
if kv_indices.shape[0] == 0 or extend_no_prefix:
o = flash_attn_varlen_func(
q,
k,
v,
qo_indptr,
qo_indptr,
max_q_len,
max_q_len,
softmax_scale=layer.scaling,
causal=True,
)
return o
if _use_fp8_prefill_attn:
total_s = q.shape[0]
nhead = layer.tp_q_head_num
v_head_dim = layer.v_head_dim
if q.dtype != fp8_dtype:
q = q.float().to(fp8_dtype)
if k.dtype != fp8_dtype:
k = k.float().to(fp8_dtype)
if v.dtype != fp8_dtype:
v = v.float().to(fp8_dtype)
one_scale = torch.tensor(
1.0, dtype=torch.float32, device=q.device
)
kv_indptr_asm = qo_indptr
kv_indices_asm = torch.arange(
total_s, device=q.device, dtype=torch.int32
)
tile_q = 256
reduce_indptr = self.forward_metadata.reduce_indptr
reduce_final_map = self.forward_metadata.reduce_final_map
reduce_partial_map = self.forward_metadata.reduce_partial_map
logits = torch.empty(
(reduce_partial_map.size(0) * tile_q, nhead, v_head_dim),
dtype=torch.float32,
device=q.device,
)
attn_lse = torch.empty(
(reduce_partial_map.size(0) * tile_q, nhead),
dtype=torch.float32,
device=q.device,
)
final_lse = torch.empty(
(total_s, nhead),
dtype=torch.float32,
device=q.device,
)
output = q.new_empty(
(total_s, nhead, v_head_dim),
dtype=self.input_dtype,
)
mla_prefill_ps_asm_fwd(
q,
k,
v,
qo_indptr,
kv_indptr_asm,
kv_indices_asm,
self.forward_metadata.work_indptr,
self.forward_metadata.work_info_set,
max_q_len,
layer.scaling,
True,
logits,
attn_lse,
output,
one_scale,
one_scale,
one_scale,
)
mla_reduce_v1(
logits,
attn_lse,
reduce_indptr,
reduce_final_map,
reduce_partial_map,
tile_q,
output,
final_lse,
)
else:
output = flash_attn_varlen_func(
q,
k,
v,
qo_indptr,
qo_indptr,
max_q_len,
max_q_len,
softmax_scale=layer.scaling,
causal=True,
)
return output
elif layer.qk_head_dim != (kv_lora_rank + qk_rope_head_dim):
K_Buffer = torch.index_select(K_Buffer, 0, kv_indices)
kvc, k_pe = torch.split(