[NPU]ACLGraph Compilation support and PassManager with AddRmsNorm & Quantize fuse. TorchAir compiler backend support.

python/sglang/srt/compilation/compilation_config.py

@@ -1,20 +1,23 @@
 # Adapted from https://github.com/vllm-project/vllm/blob/v0.10.0/vllm/compilation/compilation_config.py
 
+import json
 from typing import List
 
 
 # TODO(Yuwei): support better compile config support
 class CompilationConfig:
     def __init__(
         self,
-        capture_sizes: List[int],
+        capture_sizes: List[int] = [],
         compiler: str = "eager",
         enable_debug_mode: bool = False,
+        splitting_ops: List[str] = [],
     ):
         self.traced_files = set()
         self.capture_sizes = capture_sizes
         self.compiler = compiler
         self.enable_debug_mode = enable_debug_mode
+        self.splitting_ops = splitting_ops
 
     def add_traced_file(self, file_path: str):
         self.traced_files.add(file_path)
@@ -25,5 +28,13 @@ def get_traced_files(self):
     def get_capture_sizes(self):
         return self.capture_sizes
 
+    @classmethod
+    def from_cli(cls, args) -> "CompilationConfig":
+        args_dict = json.loads(args)
+        return CompilationConfig(**args_dict)
+
     def get_enable_debug_mode(self):
         return self.enable_debug_mode
+
+    def get_splitting_ops(self):
+        return self.splitting_ops

python/sglang/srt/compilation/custom_ops.py

@@ -0,0 +1,52 @@
+# Copyright 2025 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.
+# ==============================================================================
+
+from typing import List, Optional
+
+import torch
+
+import sglang.srt.layers.dp_attention
+
+
+@torch.library.custom_op("sglang::_set_dp_buffer_len", mutates_args=())
+def _set_dp_buffer_len(
+    global_dp_buffer_len: Optional[int],
+    num_tokens: Optional[int],
+    is_max_len: bool,
+    global_num_tokens: Optional[List[int]] = None,
+) -> None:
+    global set_dp_buffer_len_original
+    sglang.srt.layers.dp_attention.set_dp_buffer_len(
+        global_dp_buffer_len, num_tokens, is_max_len, global_num_tokens
+    )
+
+
+@_set_dp_buffer_len.register_fake
+def _set_dp_buffer_len_fake(
+    global_dp_buffer_len: Optional[int],
+    num_tokens: Optional[int],
+    is_max_len: bool,
+    global_num_tokens: Optional[List[int]] = None,
+) -> None:
+    pass
+
+
+@torch.library.custom_op("sglang::_set_is_extend_in_batch", mutates_args=())
+def _set_is_extend_in_batch(is_extend_in_batch: bool) -> None:
+    sglang.srt.layers.dp_attention.set_is_extend_in_batch(is_extend_in_batch)
+
+
+@_set_is_extend_in_batch.register_fake
+def _set_is_extend_in_batch_fake(is_extend_in_batch: bool) -> None:
+    pass

python/sglang/srt/compilation/npu/compilation_context.py

@@ -0,0 +1,20 @@
+# Copyright 2023-2025 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.
+# ==============================================================================
+
+import torch_npu
+
+
+class CompilationContext:
+    graph_memory_pool = None
+    stream: torch_npu.npu.Stream = None

python/sglang/srt/compilation/npu/npu_graph_compiler.py

@@ -0,0 +1,41 @@
+# Copyright 2025 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.
+# ==============================================================================
+
+import torch
+
+from sglang.srt.compilation.compilation_config import CompilationConfig
+from sglang.srt.utils.common import get_compiler_backend
+
+
+class NpuGraphCompiler:
+    def __init__(
+        self,
+        model_runner,
+        model: torch.nn.Module,
+        compilation_config: CompilationConfig,
+    ):
+        torch._dynamo.reset()
+
+        backend = get_compiler_backend(
+            (
+                "npugraph_fused"
+                if compilation_config is None or compilation_config.compiler is None
+                else compilation_config.compiler
+            ),
+            model_runner.model_config.dtype,
+        )
+        backend.init(model_runner.model_config)
+        self.compiled_callable = torch.compile(
+            model, fullgraph=True, dynamic=False, backend=backend
+        )

python/sglang/srt/compilation/npu/npu_graph_compiler_backend.py

@@ -0,0 +1,81 @@
+# Copyright 2025 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.
+# ==============================================================================
+
+from typing import Callable
+
+import torch
+
+from sglang.srt.compilation.npu.pass_manager import PassManager
+from sglang.srt.compilation.npu.passes.fp16 import SplitQkvRmsnormRopeFuse
+from sglang.srt.compilation.npu.passes.w8a8_int8 import (
+    DivFuse,
+    EraseCopy,
+    NpuAddRmsNormDynamicQuantFuse,
+    NpuAddRmsNormQuantFuse,
+)
+from sglang.srt.layers.dp_attention import get_attention_tp_size
+
+
+class NpuGraphCompilerBackend:
+    def __init__(self, model_type: torch.dtype):
+        self.model_type = model_type
+
+    def __call__(self, graph: torch.fx.GraphModule, example_inputs) -> Callable:
+        if self.model_type == torch.bfloat16:
+            self.apply_passes(graph)
+        return graph
+
+    def init(self, config):
+        config = config.hf_config
+
+        hidden_size = config.hidden_size
+
+        num_heads = config.num_attention_heads
+        num_kv_heads = config.num_key_value_heads
+
+        head_dim = getattr(config, "head_dim", None)
+        self.rms_norm_eps = config.rms_norm_eps
+
+        total_num_heads = num_heads
+        attn_tp_size = get_attention_tp_size()
+
+        assert total_num_heads % attn_tp_size == 0
+        num_heads = total_num_heads // attn_tp_size
+        total_num_kv_heads = num_kv_heads
+        num_kv_heads = max(1, total_num_kv_heads // attn_tp_size)
+
+        self.head_dim = head_dim or hidden_size // total_num_heads
+        self.q_size = num_heads * self.head_dim
+        self.kv_size = num_kv_heads * self.head_dim
+
+        self.q_shape = (self.head_dim, self.q_size)
+        self.k_shape = (self.head_dim, self.kv_size)
+
+    def apply_passes(self, graph_module: torch.fx.GraphModule):
+        passManager = PassManager(graph_module)
+        passManager.add(
+            SplitQkvRmsnormRopeFuse,
+            q_size=self.q_size,
+            kv_size=self.kv_size,
+            head_dim=self.head_dim,
+            q_shape=self.q_shape,
+            k_shape=self.k_shape,
+            variance_epsilon=self.rms_norm_eps,
+        )
+        passManager.add(NpuAddRmsNormQuantFuse)
+        passManager.add(NpuAddRmsNormDynamicQuantFuse)
+        passManager.add(DivFuse)
+        passManager.add(EraseCopy)
+        passManager.apply()
+        graph_module.recompile()

python/sglang/srt/compilation/npu/pass_manager.py

@@ -0,0 +1,46 @@
+# Copyright 2025 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.
+# ==============================================================================
+
+import torch
+
+
+class PassManager:
+    def __init__(self, graph_module: torch.fx.GraphModule):
+        self.graph_module = graph_module
+        self.passes = []
+
+    def add(self, pass_, **kwargs):
+        self.passes.append((pass_, kwargs))
+
+    def apply(self):
+        updated = False
+        for pass_, kwargs in self.passes:
+            pass_instance = pass_(**kwargs)
+            results = []
+            try:
+                if callable(pass_instance):
+                    results = pass_instance(self.graph_module)
+                else:
+                    results = torch.fx.replace_pattern(
+                        self.graph_module, pass_.pattern, pass_.replacement
+                    )
+            except:
+                # pass was not applied
+                pass
+
+            if not updated:
+                updated = len(results) != 0
+
+        if updated:
+            self.graph_module.recompile()