Remove the torchao integration (--torchao-config) (#34304)

Co-authored-by: Brayden Zhong <brayden@radixark.ai>
2026-08-14 21:49:11 +08:00
parent fe0c18effd
commit 5e65dd01a7
23 changed files with 6 additions and 309 deletions
@@ -58,7 +58,6 @@ runtime_common = [
  "soundfile==0.13.1",
  "tiktoken",
  "timm==1.0.16",
-  "torchao==0.9.0",
  "tqdm",
  "transformers==5.12.1",
  "uvicorn",
@@ -282,7 +282,6 @@ ARG SETUPTOOLS_SCM_PRETEND_VERSION
 RUN pip install IPython \
    && pip install orjson \
    && pip install python-multipart \
-    && pip install torchao==0.9.0 \
    && pip install pybind11

 # Rust toolchain — needed by setuptools-rust to build the sglang-mm extension
@@ -69,7 +69,7 @@ ENV PATH="$VIRTUAL_ENV/bin:$PATH"
 WORKDIR /sgl-workspace

 RUN  pip install --no-cache-dir msgspec blake3 py-cpuinfo compressed_tensors gguf partial_json_parser einops tabulate --root-user-action=ignore && \
-     pip install --no-cache-dir torch==2.12.0+xpu torchao==0.17.0+xpu torchvision==0.27.0+xpu torchaudio==2.11.0+xpu --index-url https://download.pytorch.org/whl/xpu
+     pip install --no-cache-dir torch==2.12.0+xpu torchvision==0.27.0+xpu torchaudio==2.11.0+xpu --index-url https://download.pytorch.org/whl/xpu

 RUN echo "Cloning ${SG_LANG_BRANCH} from ${SG_LANG_REPO}" && \
    git clone --branch ${SG_LANG_BRANCH} --single-branch ${SG_LANG_REPO} sglang && \
@@ -1,7 +1,7 @@
 ---
 title: "Quantization"
 metatags:
-    description: "SGLang quantization: FP8, FP4, AWQ, GPTQ, ModelOpt, torchao. Offline and online quantization methods for efficient LLM inference."
+    description: "SGLang quantization: FP8, FP4, AWQ, GPTQ, ModelOpt. Offline and online quantization methods for efficient LLM inference."
 ---
 SGLang supports various quantization methods, including offline quantization and online dynamic quantization.

@@ -182,13 +182,6 @@ The following table summarizes quantization method support across NVIDIA and AMD
      <td>No</td>
      <td>Depends on bitsandbytes ROCm support</td>
    </tr>
-    <tr>
-      <td><code>torchao</code> (<code>int4wo</code>, etc.)</td>
-      <td>Yes</td>
-      <td>Partial</td>
-      <td>No</td>
-      <td><code>int4wo</code> not supported on AMD; other methods may work</td>
-    </tr>
    <tr>
      <td><code>modelslim</code></td>
      <td>No</td>
@@ -840,29 +833,6 @@ python3 -m sglang.launch_server \

 Our team is working on supporting more online quantization methods. SGLang will soon support methods including but not limited to `["awq", "gptq", "marlin", "gptq_marlin", "awq_marlin", "bitsandbytes", "gguf"]`.

-### torchao online quantization method
-
-SGLang also supports quantization methods based on [torchao](https://github.com/pytorch/ao). You can simply specify `--torchao-config` in the command line to support this feature. For example, if you want to enable `int4wo-128` for model `meta-llama/Meta-Llama-3.1-8B-Instruct`, you can launch the server with the following command:
-
-```bash Command
-python3 -m sglang.launch_server \
-    --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \
-    --torchao-config int4wo-128 \
-    --port 30000 --host 0.0.0.0
-```
-
-SGLang supports the following quantization methods based on torchao `["int8dq", "int8wo", "fp8wo", "fp8dq-per_tensor", "fp8dq-per_row", "int4wo-32", "int4wo-64", "int4wo-128", "int4wo-256"]`.
-
-Note: According to [this issue](https://github.com/sgl-project/sglang/issues/2219#issuecomment-2561890230), `"int8dq"` method currently has some bugs when using together with cuda graph capture. So we suggest to disable cuda graph capture when using `"int8dq"` method. Namely, please use the following command:
-
-```bash Command
-python3 -m sglang.launch_server \
-    --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \
-    --torchao-config int8dq \
-    --disable-cuda-graph \
-    --port 30000 --host 0.0.0.0
-```
-
 ### `nvfp4_online` online quantization method

 Use `--quantization nvfp4_online` to convert eligible BF16, FP16, or FP8 MoE expert weights to NVFP4 at load time with per-token FP32 activation scales. Use `modelopt_fp4` for serialized NVFP4 checkpoints or the same load-time conversion with per-tensor FP32 activation scales.
@@ -969,7 +939,6 @@ For the full quantization + format conversion workflow and a complete list of su
 - [NVIDIA Model Optimizer (ModelOpt)](https://github.com/NVIDIA/Model-Optimizer)
 - [NVIDIA Model Optimizer LLM PTQ](https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/llm_ptq)
 - [Petit: NVFP4 on ROCm](https://github.com/causalflow-ai/petit-kernel) — [LMSYS blog](https://lmsys.org/blog/2025-09-21-petit-amdgpu/), [AMD ROCm blog](https://rocm.blogs.amd.com/artificial-intelligence/fp4-mixed-precision/README.html)
- [Torchao: PyTorch Architecture Optimization](https://github.com/pytorch/ao)
 - [vLLM Quantization](https://docs.vllm.ai/en/latest/quantization/)
 - [auto-round](https://github.com/intel/auto-round)
 - [ModelSlim](https://gitcode.com/Ascend/msmodelslim)
@@ -2621,12 +2621,6 @@ Please consult the documentation below and [server_args.py](https://github.com/s
      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>4096</code></td>
      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Type: int</td>
    </tr>
-    <tr>
-      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--torchao-config`</td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Optimize the model with torchao. Experimental feature. Current choices are: int8dq, int8wo, int4wo-&lt;group_size&gt;, fp8wo, fp8dq-per_tensor, fp8dq-per_row</td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`""`</td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Type: str</td>
-    </tr>

    <tr>
      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--enable-p2p-check`</td>
@@ -2246,12 +2246,6 @@ If the value is int8, you must also set the environment variable:DEEP_NORMAL_MOD
      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Type: int</td>
      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>A2, A3</td>
    </tr>
-    <tr>
-      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--torchao-config`</td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>``</td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Type: str</td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Special for GPU</td>
-    </tr>
    <tr>
      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--enable-p2p-check`</td>
      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`False`</td>
@@ -52,23 +52,6 @@ python -m sglang.launch_server \
 The quantization and limited context length (`--dtype half --context-length 8192`) are due to the limited computational resources in [Nvidia jetson kit](https://www.nvidia.com/en-us/autonomous-machines/embedded-systems/jetson-orin/). A detailed explanation can be found in [Server Arguments](../advanced_features/server_arguments).

 After launching the engine, refer to [Chat completions](../basic_usage/openai_api_completions#usage) to test the usability.
-* * * * *
-Running quantization with TorchAO
-------------------------------------
-TorchAO is suggested to NVIDIA Jetson Orin.
-```bash Command
-python -m sglang.launch_server \
-    --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \
-    --device cuda \
-    --dtype bfloat16 \
-    --attention-backend flashinfer \
-    --mem-fraction-static 0.8 \
-    --context-length 8192 \
-    --torchao-config int4wo-128
-```
-This enables TorchAO's int4 weight-only quantization with a 128-group size. The usage of `--torchao-config int4wo-128` is also for memory efficiency.
-
-
 * * * * *
 Structured output with XGrammar
 -------------------------------
@@ -53,7 +53,7 @@ conda create -n sgl-xpu python=3.12 -y
 conda activate sgl-xpu

 # Set PyTorch XPU as primary pip install channel to avoid installing the larger CUDA-enabled version and prevent potential runtime issues.
-pip3 install torch==2.12.0+xpu torchao==0.17.0+xpu torchvision==0.27.0+xpu torchaudio==2.11.0+xpu --index-url https://download.pytorch.org/whl/xpu
+pip3 install torch==2.12.0+xpu torchvision==0.27.0+xpu torchaudio==2.11.0+xpu --index-url https://download.pytorch.org/whl/xpu
 pip3 install xgrammar --no-deps # xgrammar will introduce CUDA-enabled triton which might conflict with XPU
 pip3 install apache-tvm-ffi # xgrammar requires apache-tvm-ffi

@@ -79,7 +79,6 @@ dependencies = [
  "tokenspeed_mla==0.1.8",
  "torch==2.13.0",
  "torch_memory_saver>=0.0.9.post1",
-  "torchao==0.17.0",
  "torchaudio==2.11.0",
  "torchcodec==0.15.0 ; sys_platform != 'linux' or (sys_platform == 'linux' and platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')", # Not available on Linux ARM.
  "torchvision",
@@ -61,7 +61,6 @@ dependencies = [
  "tiktoken",
  "timm==1.0.16",
  "torch==2.12.0",
-  "torchao==0.17.0",
  "torchaudio==2.11.0",
  "torchvision==0.27.0",
  "tqdm",
@@ -60,7 +60,6 @@ dependencies = [
  "soundfile==0.13.1",
  "tiktoken",
  "timm==1.0.16",
-  "torchao==0.9.0",
  "tqdm",
  "transformers==5.12.1",
  "uvicorn",
@@ -73,7 +73,6 @@ runtime_common = [
  "compressed-tensors",
  "outlines==0.1.11",
  "timm==1.0.16",
-  "torchao==0.9.0",
  "xgrammar==0.2.1",
 ]

@@ -159,7 +158,6 @@ srt_mps = [
    "mlx-lm",
    "sglang[runtime_common]",
    "torch==2.11.0",
-    "torchao==0.9.0",
    "torchaudio==2.11.0",
    "torchvision",
 ]
@@ -60,7 +60,6 @@ dependencies = [
  "tiktoken",
  "timm==1.0.16",
  "torch==2.12.0+xpu",
-  "torchao==0.17.0+xpu",
  "torchaudio==2.11.0+xpu",
  "torchcodec==0.12.0 ; sys_platform != 'linux' or (sys_platform == 'linux' and platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')", # torch==2.12.0 on XPU uses torchcodec 0.12.0
  "torchvision==0.27.0+xpu",
@@ -40,7 +40,6 @@ PACKAGE_LIST = [
    "pydantic",
    "python-multipart",
    "pyzmq",
-    "torchao",
    "uvicorn",
    "uvloop",
    "vllm",
@@ -1,95 +0,0 @@
-"""
-Common utilities for torchao.
-"""
-
-import logging
-from typing import Callable, Optional
-
-import torch
-
-logger = logging.getLogger(__name__)
-
-
-def proj_filter(
-    module: torch.nn.Module,
-    fqn: str,
-):
-    """Filter function for quantizing projection layers."""
-    return "proj" in fqn
-
-
-# TODO: implement a more general filter function
-def proj_filter_conv3d(
-    module: torch.nn.Module,
-    fqn: str,
-):
-    if isinstance(module, torch.nn.Conv3d):
-        logger.warning(f"Quantize: skipping {fqn} because it's a Conv3d")
-        return False
-    return "proj" in fqn
-
-
-def apply_torchao_config_to_model(
-    model: torch.nn.Module,
-    torchao_config: str,
-    filter_fn: Optional[Callable] = proj_filter,
-):
-    """Quantize a modelwith torchao quantization specified by torchao_config
-
-    Args:
-       `model`: a model to be quantized based on torchao_config
-       `torchao_config` (str): type of quantization and their arguments we want to use to
-        quantize the model, e.g. int4wo-128 means int4 weight only quantization with group_size
-        128
-    """
-    if torchao_config == "" or torchao_config is None:
-        return model
-
-    # Lazy import to suppress some warnings
-    from torchao.quantization import (
-        float8_dynamic_activation_float8_weight,
-        float8_weight_only,
-        int4_weight_only,
-        int8_dynamic_activation_int8_weight,
-        int8_weight_only,
-        quantize_,
-    )
-    from torchao.quantization.observer import PerRow, PerTensor
-
-    if "int8wo" in torchao_config:
-        quantize_(model, int8_weight_only(), filter_fn=proj_filter_conv3d)
-    elif "int8dq" in torchao_config:
-        quantize_(model, int8_dynamic_activation_int8_weight(), filter_fn=filter_fn)
-    elif "int4wo" in torchao_config:
-        group_size = int(torchao_config.split("-")[-1])
-        assert group_size in [
-            32,
-            64,
-            128,
-            256,
-        ], f"int4wo groupsize needs to be one of [32, 64, 128, 256] but got {group_size}"
-        quantize_(model, int4_weight_only(group_size=group_size), filter_fn=filter_fn)
-    elif "fp8wo" in torchao_config:
-        # this requires newer hardware
-        # [rank0]: AssertionError: fp8e4nv data type is not supported on CUDA arch < 89
-        quantize_(model, float8_weight_only(), filter_fn=proj_filter_conv3d)
-    elif "fp8dq" in torchao_config:
-        granularity = torchao_config.split("-")[-1]
-        GRANULARITY_MAP = {
-            "per_row": PerRow(),
-            "per_tensor": PerTensor(),
-        }
-        assert (
-            granularity in GRANULARITY_MAP
-        ), f"Supported granularity are: {GRANULARITY_MAP.keys()}, got {granularity}"
-        quantize_(
-            model,
-            float8_dynamic_activation_float8_weight(
-                granularity=GRANULARITY_MAP[granularity]
-            ),
-            filter_fn=proj_filter_conv3d,
-        )
-    else:
-        raise ValueError(f"Unexpected config: {torchao_config}")
-
-    return model
@@ -80,7 +80,6 @@ from sglang.srt.layers.cp.utils import (
 )
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 from sglang.srt.layers.sampler import create_sampler
-from sglang.srt.layers.torchao_utils import apply_torchao_config_to_model
 from sglang.srt.layers.utils.cp_utils import is_mla_prefill_cp_enabled
 from sglang.srt.lora.lora_manager import LoRAManager, init_lora_cuda_graph_moe_buffers
 from sglang.srt.lora.lora_registry import LoRARef
@@ -754,10 +753,6 @@ class ModelRunner:
        )

    def maybe_apply_post_load_model_transforms(self):
-        # In layered loading, torchao may have been applied
-        torchao_applied = getattr(self.model, "torchao_applied", False)
-        if not torchao_applied:
-            apply_torchao_config_to_model(self.model, get_exec().graph.torchao_config)
        supports_torch_tp = getattr(self.model, "supports_torch_tp", False)
        if self.ps.tp_size > 1 and supports_torch_tp:
            self.apply_torch_tp()
@@ -1079,9 +1079,6 @@ class LayeredModelLoader(DefaultModelLoader):
        model_config: ModelConfig,
        device_config: DeviceConfig,
    ) -> nn.Module:
-        from sglang.srt.layers.torchao_utils import apply_torchao_config_to_model
-
-        torchao_config = get_exec().graph.torchao_config
        target_device = torch.device(device_config.device)
        quant_config = _get_quantization_config(model_config, self.load_config)

@@ -1122,18 +1119,10 @@ class LayeredModelLoader(DefaultModelLoader):
                    fqn_path,
                    weights,
                )
-                # Quantize weights if applicable
-                if torchao_config and "proj" in fqn_path:
-                    # Note: `None` here is needed to indicate no filter, see
-                    # `apply_torchao_config_to_model` for details.
-                    apply_torchao_config_to_model(module, torchao_config, None)

            # Start calling on root module
            fill_module(model, [], weights)

-        if torchao_config:
-            model.torchao_applied = True
-
        return model.eval()


@@ -290,7 +290,7 @@ class SiglipMLP(nn.Module):

        self.config = config
        self.activation_fn = get_act_fn(config.hidden_act)
-        if quant_config and quant_config.get_name() in ["bitsandbytes", "torchao"]:
+        if quant_config and quant_config.get_name() == "bitsandbytes":
            quantizable = True
        else:
            quantizable = (
@@ -2025,7 +2025,7 @@ class ServerArgs:
    ] = False

    # -------------------------------------------------------------------------
-    # Torch compile and torchao
+    # Torch compile
    # -------------------------------------------------------------------------
    enable_torch_compile: A[
        bool,
@@ -2038,12 +2038,6 @@ class ServerArgs:
    torch_compile_max_bs: A[
        int, "Set the maximum batch size when using torch compile.", NS("exec.graph")
    ] = 32
-    torchao_config: A[
-        str,
-        "Optimize the model with torchao. Experimental feature. Current choices are: int8dq, int8wo, int4wo-<group_size>, fp8wo, fp8dq-per_tensor, fp8dq-per_row",
-        NS("exec.graph"),
-    ] = ""
-
    # -------------------------------------------------------------------------
    # Speculative decoding
    # -------------------------------------------------------------------------
@@ -577,7 +577,6 @@ class SRTRunner:
        speculative_num_draft_tokens: Optional[int] = None,
        disable_overlap_schedule: bool = False,
        disable_custom_all_reduce: bool = False,
-        torchao_config: Optional[str] = None,
        cuda_graph_max_bs_decode: int = 4,
        sleep_on_idle=False,
        max_lora_rank: Optional[int] = None,
@@ -615,7 +614,6 @@ class SRTRunner:
            dtype=get_dtype_str(torch_dtype),
            port=port,
            model_impl=model_impl,
-            torchao_config=torchao_config,
            mem_fraction_static=mem_fraction_static,
            trust_remote_code=trust_remote_code,
            is_embedding=not self.is_generation,
@@ -413,7 +413,7 @@ uninstall_stale_flashinfer() {

 install_pytorch_stack() {
    PYTORCH_SPECS=()
-    for package in torch torchaudio torchvision torchao torchcodec; do
+    for package in torch torchaudio torchvision torchcodec; do
        spec=$(grep -Po -m1 "\"${package}([<>=!~ ;][^\"]*)?\"" python/pyproject.toml | tr -d '"' || true)
        if [ -n "$spec" ]; then
            PYTORCH_SPECS+=("$spec")
@@ -1,92 +0,0 @@
-import unittest
-
-import requests
-from transformers import AutoProcessor
-
-from sglang import Engine
-from sglang.srt.utils import kill_process_tree
-from sglang.test.kits.eval_accuracy_kit import MMLUMixin
-from sglang.test.test_utils import (
-    DEFAULT_IMAGE_URL,
-    DEFAULT_MODEL_NAME_FOR_TEST,
-    DEFAULT_SMALL_VLM_MODEL_NAME_FOR_TEST,
-    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-    DEFAULT_URL_FOR_TEST,
-    CustomTestCase,
-    build_vlm_image_prompt,
-    is_in_amd_ci,
-    popen_launch_server,
-)
-
-
-class TestTorchAO(CustomTestCase, MMLUMixin):
-    mmlu_score_threshold = 0.60
-    mmlu_num_examples = 64
-    mmlu_num_threads = 32
-
-    @classmethod
-    def setUpClass(cls):
-        cls.model = DEFAULT_MODEL_NAME_FOR_TEST
-        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=["--torchao-config", "int4wo-128"],
-        )
-
-    @classmethod
-    def tearDownClass(cls):
-        kill_process_tree(cls.process.pid)
-
-    def run_decode(self, max_new_tokens):
-        response = requests.post(
-            self.base_url + "/generate",
-            json={
-                "text": "The capital of France is",
-                "sampling_params": {
-                    "temperature": 0,
-                    "max_new_tokens": max_new_tokens,
-                },
-                "ignore_eos": True,
-            },
-        )
-        return response.json()
-
-    def test_throughput(self):
-        import time
-
-        max_tokens = 256
-
-        tic = time.perf_counter()
-        res = self.run_decode(max_tokens)
-        tok = time.perf_counter()
-        print(res["text"])
-        throughput = max_tokens / (tok - tic)
-        print(f"Throughput: {throughput} tokens/s")
-        if is_in_amd_ci():
-            assert throughput >= 150
-        else:
-            assert throughput >= 210
-
-
-class TestTorchAOForVLM(CustomTestCase):
-    def test_vlm_generate(self):
-        model_path = DEFAULT_SMALL_VLM_MODEL_NAME_FOR_TEST
-        text = build_vlm_image_prompt(
-            AutoProcessor.from_pretrained(model_path), "What is in this picture?"
-        )
-
-        engine = Engine(
-            model_path=model_path,
-            max_total_tokens=512,
-            enable_multimodal=True,
-            torchao_config="fp8wo",
-        )
-        out = engine.generate([text], image_data=[DEFAULT_IMAGE_URL])
-        engine.shutdown()
-        self.assertGreater(len(out), 0)
-
-
-if __name__ == "__main__":
-    unittest.main()
@@ -32,29 +32,6 @@ class TestSRTEngineWithQuantArgs(CustomTestCase):
            engine.generate(prompt, sampling_params)
            engine.shutdown()

-    def test_2_torchao_args(self):
-
-        # we don't test int8dq because currently there is conflict between int8dq and capture cuda graph
-        torchao_args_list = [
-            # "int8dq",
-            "int8wo",
-            "fp8wo",
-            "fp8dq-per_tensor",
-            "fp8dq-per_row",
-        ] + [f"int4wo-{group_size}" for group_size in [32, 64, 128, 256]]
-
-        prompt = "Today is a sunny day and I like"
-        model_path = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
-
-        sampling_params = {"temperature": 0, "max_new_tokens": 8}
-
-        for torchao_config in torchao_args_list:
-            engine = sgl.Engine(
-                model_path=model_path, random_seed=42, torchao_config=torchao_config
-            )
-            engine.generate(prompt, sampling_params)
-            engine.shutdown()
-

 if __name__ == "__main__":
    unittest.main()