From 349ce2dd196e9d6f0dca37f919c4323807e2f28e Mon Sep 17 00:00:00 2001 From: "Ho-Ren (Jack) Chuang" Date: Sat, 27 Dec 2025 18:48:30 -0800 Subject: [PATCH] Support kv8 (FP8) with torch_native attention backend (#12596) Signed-off-by: Ho-Ren (Jack) Chuang --- docs/advanced_features/attention_backend.md | 2 +- .../srt/layers/attention/torch_native_backend.py | 10 ++++++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/docs/advanced_features/attention_backend.md b/docs/advanced_features/attention_backend.md index 7052b4bbd..6fc385f8c 100644 --- a/docs/advanced_features/attention_backend.md +++ b/docs/advanced_features/attention_backend.md @@ -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** | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | diff --git a/python/sglang/srt/layers/attention/torch_native_backend.py b/python/sglang/srt/layers/attention/torch_native_backend.py index 6a67ea947..85f0cb42f 100644 --- a/python/sglang/srt/layers/attention/torch_native_backend.py +++ b/python/sglang/srt/layers/attention/torch_native_backend.py @@ -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),