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

Support kv8 (FP8) with torch_native attention backend (#12596)

Browse Source 
Signed-off-by: Ho-Ren (Jack) Chuang <horenchuang@bytedance.com>
This commit is contained in:
Ho-Ren (Jack) Chuang
2025-12-27 18:48:30 -08:00
committed by GitHub
parent 183b65190a
commit 349ce2dd19
2 changed files with 11 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
python/sglang/srt/layers/attention/torch_native_backend.py
View File

@@ -93,6 +93,11 @@ class TorchNativeAttnBackend(AttentionBackend):
per_req_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
per_req_value = v_cache[per_req_tokens].movedim(0, query.dim() - 2)
if not (per_req_query.dtype == per_req_key.dtype == per_req_value.dtype):
# scaled_dot_product_attention() expects query, key, and value to have the same dtype
per_req_key = per_req_key.to(per_req_query.dtype)
per_req_value = per_req_value.to(per_req_query.dtype)
per_req_out_redudant = (
scaled_dot_product_attention(
per_req_query_redudant.unsqueeze(0),
@@ -162,6 +167,11 @@ class TorchNativeAttnBackend(AttentionBackend):
per_req_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
per_req_value = v_cache[per_req_tokens].movedim(0, query.dim() - 2)
if not (per_req_query.dtype == per_req_key.dtype == per_req_value.dtype):
# scaled_dot_product_attention() expects query, key, and value to have the same dtype
per_req_key = per_req_key.to(per_req_query.dtype)
per_req_value = per_req_value.to(per_req_query.dtype)
per_req_out = (
scaled_dot_product_attention(
per_req_query.unsqueeze(0),