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

[AMD] Fp8 prefill integration with radix cache path for dpsk models (#20187)

Browse Source
This commit is contained in:
Thomas Wang
2026-03-10 02:49:47 -07:00
committed by GitHub
parent 751c454099
commit 6407891b4f
1 changed files with 129 additions and 102 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/aiter_backend.py
+108 -81
View File
@@ -432,6 +432,81 @@ class AiterAttnBackend(AttentionBackend):
is_causal=is_causal, is_causal=is_causal,
) )
def mla_fp8_prefill_attn(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
layer: RadixAttention,
):
total_q = q.shape[0]
nhead = layer.tp_q_head_num
v_head_dim = layer.v_head_dim
if q.dtype != fp8_dtype:
q = q.to(fp8_dtype)
if k.dtype != fp8_dtype:
k = k.to(fp8_dtype)
if v.dtype != fp8_dtype:
v = v.to(fp8_dtype)
one_scale = torch.ones((), dtype=torch.float32, device=q.device)
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_q, nhead),
dtype=torch.float32,
device=q.device,
)
output = q.new_empty(
(total_q, nhead, v_head_dim),
dtype=self.input_dtype,
)
mla_prefill_ps_asm_fwd(
q,
k,
v,
self.forward_metadata.qo_indptr,
self.forward_metadata.kv_indptr,
self.forward_metadata.fp8_prefill_kv_indices,
self.forward_metadata.work_indptr,
self.forward_metadata.work_info_set,
self.forward_metadata.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,
)
return output
def init_forward_metadata(self, forward_batch: ForwardBatch): def init_forward_metadata(self, forward_batch: ForwardBatch):
"""Init auxiliary variables for triton attention backend.""" """Init auxiliary variables for triton attention backend."""
@@ -749,6 +824,7 @@ class AiterAttnBackend(AttentionBackend):
max_q_len = self.mla_indices_updater_prefill.max_q_len max_q_len = self.mla_indices_updater_prefill.max_q_len
qo_indptr = self.mla_indices_updater_prefill.qo_indptr qo_indptr = self.mla_indices_updater_prefill.qo_indptr
kv_indptr = self.mla_indices_updater_prefill.kv_indptr
work_metadata = None work_metadata = None
work_indptr = None work_indptr = None
@@ -774,7 +850,7 @@ class AiterAttnBackend(AttentionBackend):
self.make_mla_prefill_ps_meta_data( self.make_mla_prefill_ps_meta_data(
qo_indptr, qo_indptr,
qo_indptr, kv_indptr,
forward_batch.seq_lens, forward_batch.seq_lens,
work_metadata, work_metadata,
work_indptr, work_indptr,
@@ -785,7 +861,7 @@ class AiterAttnBackend(AttentionBackend):
is_causal=True, is_causal=True,
) )
total_s = int(forward_batch.extend_seq_lens.sum()) total_s = forward_batch.seq_lens_sum
fp8_prefill_kv_indices = torch.arange( fp8_prefill_kv_indices = torch.arange(
total_s, device=self.device, dtype=torch.int32 total_s, device=self.device, dtype=torch.int32
) )
@@ -1456,77 +1532,11 @@ class AiterAttnBackend(AttentionBackend):
extend_no_prefix = not any(forward_batch.extend_prefix_lens_cpu) extend_no_prefix = not any(forward_batch.extend_prefix_lens_cpu)
if kv_indices.shape[0] == 0 or extend_no_prefix: if kv_indices.shape[0] == 0 or extend_no_prefix:
if _use_fp8_prefill_attn: if _use_fp8_prefill_attn:
total_s = q.shape[0] output = self.mla_fp8_prefill_attn(
nhead = layer.tp_q_head_num
v_head_dim = layer.v_head_dim
# q is cast here (after RoPE).
# k/v are already FP8 for MXFP4 main model (fused kv_b_proj),
# but need casting for FP8/BF16 weights (e.g. MTP draft model).
if q.dtype != fp8_dtype:
q = q.to(fp8_dtype)
if k.dtype != fp8_dtype:
k = k.to(fp8_dtype)
if v.dtype != fp8_dtype:
v = v.to(fp8_dtype)
one_scale = torch.ones((), dtype=torch.float32, device=q.device)
kv_indptr_asm = qo_indptr
kv_indices_asm = self.forward_metadata.fp8_prefill_kv_indices
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, q,
k, k,
v, v,
qo_indptr, layer,
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: else:
output = flash_attn_varlen_func( output = flash_attn_varlen_func(
@@ -1553,17 +1563,34 @@ class AiterAttnBackend(AttentionBackend):
kvc = kvc.to(dtype) kvc = kvc.to(dtype)
k_pe = k_pe.to(dtype) k_pe = k_pe.to(dtype)
kvprefix = layer.kv_b_proj(kvc.contiguous())[0] if (
_use_fp8_prefill_attn
and layer.kv_b_proj.weight.dtype == torch.uint8
):
# MXFP4 weights + FP8 prefill: fuse GEMM, nope/v split, and k_pe cat
# into a single kernel (fused_gemm_afp4wfp4_split_cat) that writes k and v
# directly in FP8, avoiding a separate elementwise cast
k, v = layer.kv_b_proj(
(
kvc.squeeze(1),
k_pe.expand(-1, layer.tp_k_head_num, -1),
qk_nope_head_dim,
layer.v_head_dim,
fp8_dtype,
)
)[0]
else:
kv = layer.kv_b_proj(kvc.contiguous())[0]
kvprefix = kvprefix.view( kv = kv.view(
-1, layer.tp_k_head_num, qk_nope_head_dim + layer.v_head_dim -1, layer.tp_k_head_num, qk_nope_head_dim + layer.v_head_dim
) )
k_prefix, v_prefix = torch.split( k, v = torch.split(
kvprefix, [qk_nope_head_dim, layer.v_head_dim], dim=-1 kv, [qk_nope_head_dim, layer.v_head_dim], dim=-1
) )
k_prefix = torch.cat( k = torch.cat(
[ [
k_prefix, k,
torch.broadcast_to( torch.broadcast_to(
k_pe, k_pe,
(k_pe.shape[0], layer.tp_k_head_num, k_pe.shape[2]), (k_pe.shape[0], layer.tp_k_head_num, k_pe.shape[2]),
@@ -1571,15 +1598,16 @@ class AiterAttnBackend(AttentionBackend):
], ],
dim=-1, dim=-1,
) )
assert ( assert (
forward_batch.extend_prefix_lens.shape forward_batch.extend_prefix_lens.shape
== forward_batch.extend_seq_lens.shape == forward_batch.extend_seq_lens.shape
) )
k = k_prefix if _use_fp8_prefill_attn:
v = v_prefix return self.mla_fp8_prefill_attn(q, k, v, layer)
else:
o = flash_attn_varlen_func( return flash_attn_varlen_func(
q, q,
k, k,
v, v,
@@ -1590,7 +1618,6 @@ class AiterAttnBackend(AttentionBackend):
softmax_scale=layer.scaling, softmax_scale=layer.scaling,
causal=True, causal=True,
) )
return o
else: else:
if layer.qk_head_dim != layer.v_head_dim: if layer.qk_head_dim != layer.v_head_dim: