Support FP8 E4M3 KV Cache (#2786)

Co-authored-by: root <bjmsong@126.com>
2025-01-13 13:17:11 +08:00
parent 85b2e05770
commit 0bb0f76311
9 changed files with 205 additions and 10 deletions
--- a/python/sglang/srt/layers/radix_attention.py
+++ b/python/sglang/srt/layers/radix_attention.py
@@ -47,6 +47,8 @@ class RadixAttention(nn.Module):
        self.logit_cap = logit_cap
        self.sliding_window_size = sliding_window_size or -1
        self.is_cross_attention = is_cross_attention
+        self.k_scale = 1.0
+        self.v_scale = 1.0

    def forward(
        self,