EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Add support for nvidia modelopt fp8 kv cache (#3223)

Browse Source
This commit is contained in:
Zhiyu
2025-02-21 15:04:58 -08:00
committed by GitHub
parent 20b765a26e
commit c66b2c9cf1
4 changed files with 65 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
python/sglang/srt/model_loader/weight_utils.py
View File

@@ -644,9 +644,20 @@ def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> Optional[str]:
return remapped_name
possible_scale_names = [".k_scale", ".v_scale"]
modelopt_scale_names = [".self_attn.k_proj.k_scale", ".self_attn.v_proj.v_scale"]
for scale_name in possible_scale_names:
if name.endswith(scale_name):
remapped_name = name.replace(scale_name, f".attn{scale_name}")
# Check and remap the name based on modelopt scale names
if any(
modelopt_scale_name in name
for modelopt_scale_name in modelopt_scale_names
):
remapped_name = name.replace(
f".self_attn.{scale_name[1]}_proj{scale_name}",
f".self_attn.attn{scale_name}",
)
else:
remapped_name = name.replace(scale_name, f".attn{scale_name}")
if remapped_name not in params_dict:
print_warning_once(
f"Found {scale_name} in the checkpoint (e.g. {name}), "