From 84e3772caf5d828416dc140f9d927c4ac02dfee2 Mon Sep 17 00:00:00 2001 From: zejunchen-zejun Date: Fri, 26 Dec 2025 16:17:45 +0800 Subject: [PATCH 1/3] [feat] Add ATOM model impl backend Signed-off-by: zejunchen-zejun --- docs/supported_models/atom_models.md | 56 ++++++++++++++++ python/sglang/srt/configs/model_config.py | 1 + python/sglang/srt/model_loader/utils.py | 2 + python/sglang/srt/models/atom.py | 80 +++++++++++++++++++++++ python/sglang/srt/server_args.py | 4 ++ 5 files changed, 143 insertions(+) create mode 100644 docs/supported_models/atom_models.md create mode 100644 python/sglang/srt/models/atom.py diff --git a/docs/supported_models/atom_models.md b/docs/supported_models/atom_models.md new file mode 100644 index 000000000000..49f1d16acfc0 --- /dev/null +++ b/docs/supported_models/atom_models.md @@ -0,0 +1,56 @@ +# ATOM Models + +## Introduction + +ATOM is a high-performance model implementation backend specifically optimized for AMD Instinct GPUs. This doc guides users to run ATOM-accelerated models in SGLang, delivering optimal performance on AMD Instinct GPU devices (e.g., MI300 series). + +## Requirements + +Users need to install ATOM + +## Supported Models + +ATOM currently provides optimized implementations for the following model architectures on AMD Instinct GPUs: + +- **Llama Models**: Llama 3.x +- **Qwen Models**: Qwen serial models +- **DeepSeek**: deepseek-r1, deepseek-v3.2 + +## Install ATOM + +```bash +# Clone the ATOM repository +git clone https://github.com/ROCm/atom.git +cd atom +pip install -e . +``` + +## Run Model + +Launch server with multiple AMD GPUs: + +```bash +python3 -m sglang.launch_server \ + --model-path /path/to/your/model \ + --host 0.0.0.0 \ + --port 30000 \ + --tensor-parallel-size 8 \ + --expert-parallel-size 8 \ + --model-impl atom +``` + +## Performance Benefits + +ATOM provides significant performance improvements on AMD Instinct GPUs: + +- **Optimized Kernels**: Custom-tuned kernels for AMD GPUs +- **Customized Fusions**: ATOM model provides customized fusion patterns + +## Support + +For ATOM-specific issues: +- Refer to the [ATOM GitHub Repository](https://github.com/ROCm/atom) +- ROCm documentation: [AMD ROCm Documentation](https://rocm.docs.amd.com/) +- SGLang issues: [SGLang GitHub Issues](https://github.com/sgl-project/sglang/issues) + +For general SGLang support, please refer to the main [SGLang documentation](https://sgl-project.github.io/docs/). diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 41ad51ce004f..9879cce10751 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -49,6 +49,7 @@ class ModelImpl(str, Enum): SGLANG = "sglang" TRANSFORMERS = "transformers" MINDSPORE = "mindspore" + ATOM = "atom" def is_deepseek_nsa(config: PretrainedConfig) -> bool: diff --git a/python/sglang/srt/model_loader/utils.py b/python/sglang/srt/model_loader/utils.py index 18739ed6954e..fbaf6a96732c 100644 --- a/python/sglang/srt/model_loader/utils.py +++ b/python/sglang/srt/model_loader/utils.py @@ -112,6 +112,8 @@ def get_model_architecture(model_config: ModelConfig) -> Tuple[Type[nn.Module], if model_config.model_impl == ModelImpl.MINDSPORE: architectures = ["MindSporeForCausalLM"] + elif model_config.model_impl == ModelImpl.ATOM: + architectures = ["ATOMForCausalLM"] elif not is_native_supported or model_config.model_impl == ModelImpl.TRANSFORMERS: architectures = resolve_transformers_arch(model_config, architectures) return ModelRegistry.resolve_model_cls(architectures) diff --git a/python/sglang/srt/models/atom.py b/python/sglang/srt/models/atom.py new file mode 100644 index 000000000000..bf18676749f6 --- /dev/null +++ b/python/sglang/srt/models/atom.py @@ -0,0 +1,80 @@ +# Copyright 2026 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Wrapper around `atom` models""" +import logging +from typing import Iterable, Optional, Tuple + +import torch +from torch import nn + +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.layers.logits_processor import LogitsProcessor, LogitsProcessorOutput + +logger = logging.getLogger(__name__) + + +class ATOMForCausalLM(nn.Module): + + def __init__( + self, + config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + logger.info("Using Atom backend.") + + self.quant_config = quant_config + self.config = config + self.vocab_size = config.vocab_size + self.unpadded_vocab_size = config.vocab_size + + import atom + self.model = atom.prepare_model(config=config, framework="sglang") + if self.model is None: + model_arch = config.model_config.architectures[0] + raise ValueError(f'This model{model_arch} is not supported by atom') + + self.logits_processor = LogitsProcessor(config) + + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + get_embedding: bool = False, + ) -> LogitsProcessorOutput: + aux_hidden_states = None + hidden_states = self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=None, + inputs_embeds=input_embeds, + forward_batch=forward_batch, + get_embedding=get_embedding, + pp_proxy_tensors=None, + ) + + return self.logits_processor( + input_ids, hidden_states, self.model.lm_head, forward_batch, aux_hidden_states + ) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + self.model.load_weights(weights) + +EntryClass = [ATOMForCausalLM] diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index bcaa015f2a8a..d32c5db9ea72 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -3324,6 +3324,7 @@ def add_cli_args(parser: argparse.ArgumentParser): '* "sglang" will use the SGLang model implementation.\n' '* "transformers" will use the Transformers model ' '* "mindspore" will use the MindSpore model ' + '* "atom" will use the Atom model ' "implementation.\n", ) @@ -5729,6 +5730,9 @@ def check_server_args(self): if self.model_impl == "mindspore": assert is_npu(), "MindSpore model impl is only supported on Ascend npu." + if self.model_impl == "atom": + assert is_hip(), "Atom model impl is only supported on ROCm." + # Check metrics labels if ( not self.tokenizer_metrics_custom_labels_header From e3ce42e6dd87340f023780264d7bec74b73592e1 Mon Sep 17 00:00:00 2001 From: zejunchen-zejun Date: Mon, 2 Feb 2026 15:10:07 +0800 Subject: [PATCH 2/3] add Signed-off-by: zejunchen-zejun --- python/sglang/srt/models/atom.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/sglang/srt/models/atom.py b/python/sglang/srt/models/atom.py index bf18676749f6..bb974022c9d6 100644 --- a/python/sglang/srt/models/atom.py +++ b/python/sglang/srt/models/atom.py @@ -42,7 +42,7 @@ def __init__( self.unpadded_vocab_size = config.vocab_size import atom - self.model = atom.prepare_model(config=config, framework="sglang") + self.model = atom.prepare_model(config=config, engine="sglang") if self.model is None: model_arch = config.model_config.architectures[0] raise ValueError(f'This model{model_arch} is not supported by atom') From 5d8ee20cd3fd7ecee70b9312af68a36b9c4c1d29 Mon Sep 17 00:00:00 2001 From: Zhiwei Date: Tue, 10 Mar 2026 16:32:05 +0800 Subject: [PATCH 3/3] ATOM Plugin UT Test(#208) ATOM Plugin UT Test --- test/registered/models/test_atom_models.py | 57 ++++++++++++++++++++++ 1 file changed, 57 insertions(+) create mode 100644 test/registered/models/test_atom_models.py diff --git a/test/registered/models/test_atom_models.py b/test/registered/models/test_atom_models.py new file mode 100644 index 000000000000..ff4f9b5901d3 --- /dev/null +++ b/test/registered/models/test_atom_models.py @@ -0,0 +1,57 @@ +# ATOM Plugin tests + +import unittest +from types import SimpleNamespace + + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_amd_ci +from sglang.test.few_shot_gsm8k import run_eval +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +register_amd_ci(est_time=600, suite="stage-b-test-small-1-gpu-amd") + +class TestQwen3MoeFp8(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = "Qwen/Qwen3-30B-A3B-FP8" + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--model-impl", + "atom", + "--kv-cache-dtype", + "fp8_e4m3", + "--page-size", + "1024" + ], + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + args = SimpleNamespace( + num_shots=5, + data_path=None, + num_questions=200, + max_new_tokens=512, + parallel=128, + host="http://127.0.0.1", + port=int(self.base_url.split(":")[-1]), + ) + metrics = run_eval(args) + print(f"{metrics=}") + self.assertGreater(metrics["accuracy"], 0.78) + +if __name__ == "__main__": + unittest.main()