diff --git a/python/sglang/srt/arg_groups/model_hook.py b/python/sglang/srt/arg_groups/model_hook.py index 72b462a73..e49392131 100644 --- a/python/sglang/srt/arg_groups/model_hook.py +++ b/python/sglang/srt/arg_groups/model_hook.py @@ -3,7 +3,9 @@ from __future__ import annotations +import json import logging +import os from typing import Any from sglang.srt.arg_groups.overrides import ( @@ -81,6 +83,59 @@ def _rocm_fp8_wo_a_supported() -> bool: return False +def _probe_wo_a_weight_dtype(model_config: Any, download_dir: str | None) -> str | None: + """Read one indexed wo_a dtype without downloading a weight shard.""" + try: + from huggingface_hub import ( + parse_local_safetensors_file_metadata, + parse_safetensors_file_metadata, + ) + from transformers.utils.hub import cached_file + + model_path = model_config.model_path + revision = ( + getattr(model_config.hf_config, "_commit_hash", None) + or model_config.revision + ) + index_path = cached_file( + model_path, + "model.safetensors.index.json", + revision=revision, + cache_dir=download_dir, + ) + with open(index_path) as f: + weight_map = json.load(f).get("weight_map", {}) + name = next((key for key in weight_map if key.endswith(".wo_a.weight")), None) + if name is None: + return None + + shard = weight_map[name] + local_shard = os.path.join(os.path.dirname(index_path), shard) + metadata = ( + parse_local_safetensors_file_metadata(local_shard) + if os.path.isfile(local_shard) + else parse_safetensors_file_metadata(model_path, shard, revision=revision) + ) + return getattr(metadata.tensors.get(name), "dtype", None) + except Exception: + logger.debug("Unable to inspect the checkpoint wo_a dtype", exc_info=True) + return None + + +def _configure_rocm_fp8_wo_a_gemm(model_config: Any, download_dir: str | None) -> None: + flag = envs.SGLANG_OPT_FP8_WO_A_GEMM + if not _rocm_fp8_wo_a_supported(): + flag.set(False) + return + if flag.is_set(): + return + + dtype = _probe_wo_a_weight_dtype(model_config, download_dir) + if dtype is not None and dtype != "F8_E4M3": + flag.set(False) + logger.info("Disabled ROCm fp8 wo_a GEMM for checkpoint dtype %s", dtype) + + def handle_model_specific_adjustments(server_args: Any): cfg = resolving_view(server_args) @@ -387,11 +442,7 @@ def handle_model_specific_adjustments(server_args: Any): envs.SGLANG_OPT_USE_TILELANG_INDEXER.set(True) elif get_platform().is_hip: envs.SGLANG_OPT_DEEPGEMM_HC_PRENORM.set(False) - # The fp8 wo_a GEMM is DeepGEMM-based on CUDA. ROCm has an aiter - # e8m0 block-scale equivalent, but only on gfx950 -- everywhere else - # keeps the bf16 absorb GEMM. - if not _rocm_fp8_wo_a_supported(): - envs.SGLANG_OPT_FP8_WO_A_GEMM.set(False) + _configure_rocm_fp8_wo_a_gemm(model_config, cfg.download_dir) envs.SGLANG_OPT_USE_JIT_INDEXER_METADATA.set(False) envs.SGLANG_OPT_USE_TOPK_V2.set(True) envs.SGLANG_OPT_USE_AITER_INDEXER.set(True)