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

Aiter fp8 kv cache (#13147)

Browse Source
This commit is contained in:
kk
2025-12-09 08:39:53 +08:00
committed by GitHub
parent 119fd956fb
commit c106b54b57
7 changed files with 594 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
python/sglang/srt/model_executor/model_runner.py
View File

@@ -95,6 +95,7 @@ from sglang.srt.layers.dp_attention import (
)
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.layers.pooler import EmbeddingPoolerOutput
from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype
from sglang.srt.layers.sampler import Sampler
from sglang.srt.layers.torchao_utils import apply_torchao_config_to_model
from sglang.srt.lora.lora_manager import LoRAManager
@@ -1629,19 +1630,19 @@ class ModelRunner:
and kv_cache_quant_algo.upper() == "FP8"
):
if _is_hip:
self.kv_cache_dtype = torch.float8_e4m3fnuz
self.kv_cache_dtype = fp8_dtype
else:
self.kv_cache_dtype = torch.float8_e4m3fn
else:
self.kv_cache_dtype = self.dtype
elif self.server_args.kv_cache_dtype == "fp8_e5m2":
if _is_hip: # Using natively supported format
self.kv_cache_dtype = torch.float8_e5m2fnuz
self.kv_cache_dtype = fp8_dtype
else:
self.kv_cache_dtype = torch.float8_e5m2
elif self.server_args.kv_cache_dtype == "fp8_e4m3":
if _is_hip: # Using natively supported format
self.kv_cache_dtype = torch.float8_e4m3fnuz
self.kv_cache_dtype = fp8_dtype
else:
self.kv_cache_dtype = torch.float8_e4m3fn
elif self.server_args.kv_cache_dtype in ("bf16", "bfloat16"):