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

[AMD]Reuse fused FP8 KV cache write on standard aiter prefill/decode (#26852)

Browse Source
This commit is contained in:
zhengyao
2026-07-16 01:05:04 -07:00
committed by GitHub
parent b55228cfdb
commit 01b003255a
2 changed files with 479 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/aiter_backend.py
+42 -6
View File
@@ -1873,6 +1873,17 @@ class AiterAttnBackend(AttentionBackend):
# base class NotImplementedError.
pass
def _use_fused_fp8_kv_write(self, layer: RadixAttention) -> bool:
# Fused write reuses K's num_heads/head_dim for V, so it needs FP8,
# non-MLA, non-SWA, and matching K/V head count + head_dim.
return (
self.kv_cache_dtype == fp8_dtype
and not self.use_mla
and not self.use_sliding_window_kv_pool
and layer.tp_k_head_num == layer.tp_v_head_num
and layer.qk_head_dim == layer.v_head_dim
)
def forward_extend(
self,
q: torch.Tensor,
@@ -1895,7 +1906,7 @@ class AiterAttnBackend(AttentionBackend):
v_descale = None
if self.kv_cache_dtype == fp8_dtype:
k_descale = layer.k_scale if layer.k_scale is not None else self.k_scale
v_descale = layer.v_scale if layer.v_scale is not None else self.k_scale
v_descale = layer.v_scale if layer.v_scale is not None else self.v_scale
if k is not None:
assert v is not None
@@ -1947,6 +1958,24 @@ class AiterAttnBackend(AttentionBackend):
)
elif self.use_mla:
self.token_to_kv_pool.set_kv_buffer(layer, cache_loc, k, v)
elif self._use_fused_fp8_kv_write(layer):
# FP8: fuse bf16->fp8 cast + paged write in one kernel.
k_cache, v_cache = self.token_to_kv_pool.get_kv_buffer(
layer.layer_id
)
launch_reshape_and_cache_flash(
k.view(-1, layer.tp_k_head_num, layer.qk_head_dim),
v.view(-1, layer.tp_v_head_num, layer.v_head_dim),
k_cache.view(
-1, self.page_size, layer.tp_k_head_num, layer.qk_head_dim
),
v_cache.view(
-1, self.page_size, layer.tp_v_head_num, layer.v_head_dim
),
cache_loc,
k_scale=k_descale,
v_scale=v_descale,
)
else:
self.token_to_kv_pool.set_kv_buffer(
layer,
@@ -2374,7 +2403,7 @@ class AiterAttnBackend(AttentionBackend):
v_descale = None
if self.kv_cache_dtype == fp8_dtype:
k_descale = layer.k_scale if layer.k_scale is not None else self.k_scale
v_descale = layer.v_scale if layer.v_scale is not None else self.k_scale
v_descale = layer.v_scale if layer.v_scale is not None else self.v_scale
if save_kv_cache:
# SHUFFLE 5D pool path — see forward_extend for rationale.
@@ -2414,10 +2443,13 @@ class AiterAttnBackend(AttentionBackend):
k_scale=k_descale,
v_scale=v_descale,
)
elif self.use_triton_unified_attention and self.kv_cache_dtype == fp8_dtype:
# [PATCH] FP8 non-SWA: use launch_reshape_and_cache_flash to
# fuse bf16→fp8 cast + paged write in one Triton kernel,
# eliminating separate float8_copy + store_kvcache overhead.
elif self.use_mla:
# MLA pool has its own set_kv_buffer (no scale args).
self.token_to_kv_pool.set_kv_buffer(
layer, forward_batch.out_cache_loc, k, v
)
elif self._use_fused_fp8_kv_write(layer):
# FP8: fuse bf16->fp8 cast + paged write in one kernel.
token_to_kv_pool = self.token_to_kv_pool
k_cache, v_cache = token_to_kv_pool.get_kv_buffer(layer.layer_id)
launch_reshape_and_cache_flash(
@@ -2430,6 +2462,8 @@ class AiterAttnBackend(AttentionBackend):
-1, self.page_size, layer.tp_v_head_num, layer.v_head_dim
),
forward_batch.out_cache_loc,
k_scale=k_descale,
v_scale=v_descale,
)
else:
self.token_to_kv_pool.set_kv_buffer(
@@ -2440,6 +2474,8 @@ class AiterAttnBackend(AttentionBackend):
),
k,
v,
k_descale,
v_descale,
)
if self.use_mla: