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

[AMD][Quantization] Add int4fp8_moe online quantization on ROCm (#7392)

Browse Source 
Co-authored-by: Dehua Tang <dehtang@amd.com>
Co-authored-by: HAI <hixiao@gmail.com>
Co-authored-by: YC Tseng <yctseng@amd.com>
This commit is contained in:
fxmarty-amd
2026-01-14 10:44:40 +01:00
committed by GitHub
parent feae615b11
commit 5af84c8af5
12 changed files with 615 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
python/sglang/srt/model_loader/utils.py
View File

@@ -87,7 +87,13 @@ def get_model_architecture(model_config: ModelConfig) -> Tuple[Type[nn.Module],
architectures = getattr(model_config.hf_config, "architectures", [])
# Special handling for quantized Mixtral.
# FIXME(woosuk): This is a temporary hack.
mixtral_supported = ["fp8", "compressed-tensors", "gptq_marlin", "awq_marlin"]
mixtral_supported = [
"fp8",
"compressed-tensors",
"gptq_marlin",
"awq_marlin",
"quark_int4fp8_moe",
]
if (
model_config.quantization is not None

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

@@ -44,7 +44,12 @@ from sglang.srt.model_loader.ci_weight_validation import (
ci_download_with_validation_and_retry,
ci_validate_and_cleanup_local_snapshot,
)
from sglang.srt.utils import find_local_repo_dir, log_info_on_rank0, print_warning_once
from sglang.srt.utils import (
BAR_FORMAT,
find_local_repo_dir,
log_info_on_rank0,
print_warning_once,
)
from sglang.utils import is_in_ci
try:
@@ -608,13 +613,6 @@ def filter_files_not_needed_for_inference(hf_weights_files: List[str]) -> List[s
return hf_weights_files
# explicitly use pure text format, with a newline at the end
# this makes it impossible to see the animation in the progress bar
# but will avoid messing up with ray or multiprocessing, which wraps
# each line of output with some prefix.
_BAR_FORMAT = "{desc}: {percentage:3.0f}% Completed | {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]\n" # noqa: E501
def np_cache_weights_iterator(
model_name_or_path: str,
cache_dir: Optional[str],
@@ -642,7 +640,8 @@ def np_cache_weights_iterator(
hf_weights_files,
desc="Loading np_cache checkpoint shards",
disable=not enable_tqdm,
bar_format=_BAR_FORMAT,
bar_format=BAR_FORMAT,
position=tqdm._get_free_pos(),
):
state = torch.load(bin_file, map_location="cpu", weights_only=True)
for name, param in state.items():
@@ -699,7 +698,8 @@ def safetensors_weights_iterator(
hf_weights_files,
desc="Loading safetensors checkpoint shards",
disable=not enable_tqdm,
bar_format=_BAR_FORMAT,
bar_format=BAR_FORMAT,
position=tqdm._get_free_pos(),
):
if disable_mmap:
with open(st_file, "rb") as f:
@@ -811,7 +811,7 @@ def multi_thread_safetensors_weights_iterator(
total=len(hf_weights_files),
desc="Multi-thread loading shards",
disable=not enable_tqdm,
bar_format=_BAR_FORMAT,
bar_format=BAR_FORMAT,
)
else:
futures_iter = concurrent.futures.as_completed(futures)
@@ -853,7 +853,8 @@ def pt_weights_iterator(
hf_weights_files,
desc="Loading pt checkpoint shards",
disable=not enable_tqdm,
bar_format=_BAR_FORMAT,
bar_format=BAR_FORMAT,
position=tqdm._get_free_pos(),
):
state = _load_pt_file(bin_file)
yield from state.items()
@@ -880,7 +881,7 @@ def multi_thread_pt_weights_iterator(
total=len(hf_weights_files),
desc="Multi-thread loading pt checkpoint shards",
disable=not enable_tqdm,
bar_format=_BAR_FORMAT,
bar_format=BAR_FORMAT,
)
else:
futures_iter = concurrent.futures.as_completed(futures)
@@ -1033,7 +1034,8 @@ def runai_safetensors_weights_iterator(
hf_weights_files,
desc="Loading safetensors using Runai Model Streamer",
disable=not enable_tqdm,
bar_format=_BAR_FORMAT,
bar_format=BAR_FORMAT,
position=tqdm._get_free_pos(),
):
streamer.stream_file(st_file)
yield from streamer.get_tensors()