chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 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

2
docs/advanced_features/attention_backend.md
View File

@@ -21,7 +21,7 @@ The support matrix is split into two parts: MHA (standard attention) and MLA (mu
| **FA3 (FlashAttention 3)** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ |
| **FA4 (FlashAttention 4)** | 128 | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **Triton** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ |
| **Torch Native (SDPA)** | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ |
| **Torch Native (SDPA)** | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ |
| **FlexAttention (PyTorch)** | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **TRTLLM MHA** | 16, 32 or 64 | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ |
| **Dual Chunk FlashAttention** | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |

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),