[feat] Add ROCm ATOM model impl backend

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/).

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:

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)

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, 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')
+
+        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]

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

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()