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

docs/platforms/ascend_npu_pass_development.md

@@ -0,0 +1,30 @@
+## How to transform model instances with PyTorch FX Toolkit in SGLang for NPU
+
+### PassManager
+`PassManager` is implemented here: [PassManager](https://github.com/eshoguli/sglang/blob/eshogulin/pass_manager/python/sglang/srt/hardware_backend/npu/graph_runner/compilation/pass_manager.py)
+
+
+You can explore `PassManager` usage in [`NpuGraphCompilerBackend`](https://github.com/eshoguli/sglang/blob/eshogulin/pass_manager/python/sglang/srt/hardware_backend/npu/graph_runner/compilation/npu_graph_compiler_backend.py) compiler backend.
+
+### Pass development
+There are two approaches to develop passes for SGLang NPU PassManager:
+
+1. Matches all possible non-overlapping sets of operators and their data dependencies with `torch.fx.replace_pattern` api.
+Pass example: [NpuAddRmsNormQuantFuse](https://github.com/eshoguli/sglang/blob/3365d711fd5aa0d6191c32769163320fe41e27f2/python/sglang/srt/hardware_backend/npu/graph_runner/compilation/passes/w8a8_int8.py#L82).
+You can find details on official FX toolkit web site: https://docs.pytorch.org/docs/stable/fx.html#subgraph-rewriting-with-replace-pattern
+
+2. Direct Graph Manipulation.
+Pass example: [EraseCopy](https://github.com/eshoguli/sglang/blob/3365d711fd5aa0d6191c32769163320fe41e27f2/python/sglang/srt/hardware_backend/npu/graph_runner/compilation/passes/w8a8_int8.py#L28).
+You can find details on official FX toolkit web site: https://docs.pytorch.org/docs/stable/fx.html#direct-graph-manipulation
+
+### Compiler backend update
+After pass development you should create `PassManager` instance, add the pass and call `apply` method:
+```
+def apply_passes(self, graph_module: torch.fx.GraphModule):
+    passManager = PassManager(graph_module)
+    passManager.add(NpuAddRmsNormQuantFuse)
+    passManager.apply()
+    graph_module.recompile()
+```
+
+You can explore [`NpuGraphCompilerBackend`](https://github.com/eshoguli/sglang/blob/eshogulin/pass_manager/python/sglang/srt/hardware_backend/npu/graph_runner/compilation/npu_graph_compiler_backend.py) as example.

docs/platforms/ascend_npu_support.rst

@@ -8,6 +8,7 @@ Ascend NPUs
    ascend_npu_support_models.md
    ascend_npu_support_features.md
    ascend_npu_deepseek_example.md
+   ascend_npu_pass_development.md
    ascend_npu_qwen3_examples.md
    ascend_contribution_guide.md
    ascend_npu_best_practice.md

python/sglang/srt/compilation/compilation_config.py

@@ -18,7 +18,7 @@ def decorator(op_func: Callable):
 class CompilationConfig:
     def __init__(
         self,
-        capture_sizes: List[int],
+        capture_sizes: List[int] = None,
         compiler: str = "eager",
         enable_debug_mode: bool = False,
     ):

python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py

@@ -222,7 +222,6 @@ def __init__(self, model_runner: ModelRunner):
         self.req_to_token = model_runner.req_to_token_pool.req_to_token
         self.graph_mode = False
         self.use_fia = get_bool_env_var("ASCEND_USE_FIA", "False")
-        self.enable_torch_compile = model_runner.server_args.enable_torch_compile
         self.speculative_num_draft_tokens = (
             model_runner.server_args.speculative_num_draft_tokens
         )
@@ -238,6 +237,13 @@ def __init__(self, model_runner: ModelRunner):
         if self.use_mla:
             self.ringmla_mask = self.ascend_attn_mask_builder.ringmla_mask
 
+        self.enable_npu_torchair_compile = (
+            model_runner.server_args.enable_npu_torchair_compile
+        )
+        if self.enable_npu_torchair_compile:
+            max_total_tokens = model_runner.max_total_num_tokens
+            self.max_seqlen_pad = max_total_tokens // model_runner.server_args.page_size
+
     def get_verify_buffers_to_fill_after_draft(self):
         """
         Return buffers for verify attention kernels that needs to be filled after draft.
@@ -257,12 +263,29 @@ def init_forward_metadata(self, forward_batch: ForwardBatch):
         seq_lens_max = forward_batch.seq_lens.max()
         if forward_batch.forward_mode.is_target_verify():
             seq_lens_max += self.speculative_num_draft_tokens
-        self.forward_metadata.block_tables = (
+
+        block_tables = (
             forward_batch.req_to_token_pool.req_to_token[
                 forward_batch.req_pool_indices, :seq_lens_max
             ][:, :: self.page_size]
             // self.page_size
         )
+
+        if (
+            self.enable_npu_torchair_compile
+            and forward_batch.forward_mode.is_decode_or_idle()
+        ):
+            bs = forward_batch.input_ids.size(0)
+            device = forward_batch.input_ids.device
+            self.forward_metadata.block_tables = torch.full(
+                (bs, self.max_seqlen_pad), -1, dtype=torch.int32, device=device
+            )
+            self.forward_metadata.block_tables[:, : block_tables.size(1)].copy_(
+                block_tables
+            )
+        else:
+            self.forward_metadata.block_tables = block_tables
+
         if forward_batch.extend_seq_lens is not None:
             self.forward_metadata.extend_seq_lens = forward_batch.extend_seq_lens
             self.forward_metadata.extend_seq_lens_cpu_int = (
@@ -1283,7 +1306,7 @@ def forward_decode(
                 topk_indices,
             )
 
-        if self.graph_mode and (not self.enable_torch_compile):
+        if self.graph_mode and (not self.enable_npu_torchair_compile):
             return self.forward_decode_graph(
                 q,
                 k,

python/sglang/srt/hardware_backend/npu/cmo.py

@@ -1,5 +1,7 @@
 import torch
 
+from sglang.srt.layers.parameter import ModelWeightParameter
+
 cmo_stream = None
 
 
@@ -18,6 +20,12 @@ def set_cmo_stream(stream):
     cmo_stream = stream
 
 
+def get_weight_cache(layer):
+    if isinstance(layer.weight, ModelWeightParameter):
+        return layer.weight_data
+    return layer.weight
+
+
 def prepare_weight_cache(handle, cache, PREFETCH_MAX_SIZE=1000000000):
     """
     PREFETCH_MAX_SIZE: maximum size (bytes) for each prefetch operation.

python/sglang/srt/hardware_backend/npu/cmo_custom_ops.py

@@ -0,0 +1,30 @@
+from typing import List
+
+import torch
+
+import sglang.srt.hardware_backend.npu.cmo
+from sglang.srt.utils.custom_op import register_custom_op
+
+
+@torch.library.custom_op("sglang::wait_cmo_stream", mutates_args=())
+def wait_cmo_stream() -> None:
+    if sglang.srt.hardware_backend.npu.cmo.get_cmo_stream():
+        sglang.srt.hardware_backend.npu.cmo.wait_cmo_stream()
+
+
+@wait_cmo_stream.register_fake
+def wait_cmo_stream_fake() -> None:
+    pass
+
+
+def get_cmo_stream() -> bool:
+    return True
+
+
+def prepare_weight_cache_fake(handle: torch.Tensor, cache: List[torch.Tensor]) -> None:
+    pass
+
+
+@register_custom_op(fake_impl=prepare_weight_cache_fake)
+def prepare_weight_cache(handle: torch.Tensor, cache: List[torch.Tensor]) -> None:
+    sglang.srt.hardware_backend.npu.cmo.prepare_weight_cache(handle, cache)

python/sglang/srt/hardware_backend/npu/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/hardware_backend/npu/graph_runner/compilation/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/hardware_backend/npu/graph_runner/compilation/custom_ops.py

@@ -0,0 +1,62 @@
+from typing import List
+
+import sgl_kernel_npu.norm.split_qkv_rmsnorm_rope as sgl_kernel_npu
+import torch
+
+from sglang.srt.utils.custom_op import register_custom_op
+
+
+def split_qkv_rmsnorm_rope_fake(
+    input: torch.Tensor,
+    sin: torch.Tensor,
+    cos: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    q_hidden_size: int,
+    kv_hiddem_size: int,
+    head_dim: int,
+    eps: float,
+    q_bias: torch.Tensor,
+    k_bias: torch.Tensor,
+) -> List[torch.Tensor]:
+    # TODO: generalize shape
+    q = torch.empty(
+        (input.shape[0], q_hidden_size), dtype=input.dtype, device=input.device
+    )
+    k = torch.empty(
+        (input.shape[0], kv_hiddem_size), dtype=input.dtype, device=input.device
+    )
+    v = torch.empty(
+        (input.shape[0], kv_hiddem_size), dtype=input.dtype, device=input.device
+    )
+    return [q, k, v]
+
+
+@register_custom_op(fake_impl=split_qkv_rmsnorm_rope_fake)
+def split_qkv_rmsnorm_rope(
+    input: torch.Tensor,
+    sin: torch.Tensor,
+    cos: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    q_hidden_size: int,
+    kv_hiddem_size: int,
+    head_dim: int,
+    eps: float,
+    q_bias: torch.Tensor,
+    k_bias: torch.Tensor,
+) -> List[torch.Tensor]:
+    q, k, v = sgl_kernel_npu.split_qkv_rmsnorm_rope(
+        input,
+        sin,
+        cos,
+        q_weight,
+        k_weight,
+        q_hidden_size,
+        kv_hiddem_size,
+        head_dim,
+        eps,
+        q_bias,
+        k_bias,
+    )
+    return [q, k, v]

python/sglang/srt/hardware_backend/npu/graph_runner/compilation/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,
+        batch_size: int,
+    ):
+        torch._dynamo.reset()
+
+        if compilation_config is None:
+            compilation_config = CompilationConfig(compiler="npugraph")
+
+        backend = get_compiler_backend(
+            compilation_config=compilation_config, model_runner=model_runner
+        )
+        backend.init(model_runner.model_config, batch_size)
+
+        self.compiled_callable = torch.compile(
+            model, fullgraph=True, dynamic=False, backend=backend
+        )