Add graph runner support with torch compile on CPU

.github/workflows/pr-test-xeon.yml

@@ -70,7 +70,7 @@ jobs:
 
       - name: Run unit tests
         if: steps.check_amx.outcome == 'success'
-        timeout-minutes: 30
+        timeout-minutes: 36
         run: |
           docker exec -w /sglang-checkout/ ci_sglang_xeon \
             bash -c "cd ./test/srt && python3 run_suite.py --suite per-commit-cpu"

docs/platforms/cpu_server.md

@@ -134,7 +134,12 @@ Notes:
     export SGLANG_CPU_OMP_THREADS_BIND="0-39|43-82|86-125|128-167|171-210|214-253"
     ```
 
-3. A warmup step is automatically triggered when the service is started.
+3. For optimizing decoding with torch.compile, please add the flag `--enable-torch-compile`.
+    To specify the maximum batch size when using torch compile, set the flag `--torch-compile-max-bs`.
+    For example, `--enable-torch-compile --torch-compile-max-bs 4` means using torch compile and setting the
+    maximum batch size to 4.
+
+4. A warmup step is automatically triggered when the service is started.
 The server is ready when you see the log `The server is fired up and ready to roll!`.
 
 ## Benchmarking with Requests

python/sglang/srt/disaggregation/prefill.py

@@ -778,7 +778,7 @@ def event_loop_pp_disagg_prefill(self: Scheduler):
                         extend_input_len_per_req=None,
                         extend_logprob_start_len_per_req=None,
                         bid=bids[next_mb_id],
-                        can_run_cuda_graph=result.can_run_cuda_graph,
+                        can_run_graph=result.can_run_graph,
                     )
                     self.process_batch_result_disagg_prefill(
                         mbs[next_mb_id], output_result

python/sglang/srt/distributed/parallel_state.py

@@ -60,6 +60,9 @@ class GraphCaptureContext:
 
 TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"])
 
+# use int value instead of ReduceOp.SUM to support torch compile
+REDUCE_OP_SUM = int(torch.distributed.ReduceOp.SUM)
+
 
 def _split_tensor_dict(
     tensor_dict: Dict[str, Union[torch.Tensor, Any]]
@@ -482,9 +485,7 @@ def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
 
         if input_.is_cpu:
             if is_shm_available(input_.dtype, self.world_size, self.local_size):
-                torch.ops.sgl_kernel.shm_allreduce(
-                    input_, torch.distributed.ReduceOp.SUM
-                )
+                torch.ops.sgl_kernel.shm_allreduce(input_, REDUCE_OP_SUM)
             else:
                 torch.distributed.all_reduce(input_, group=self.device_group)
             return input_

python/sglang/srt/layers/attention/intel_amx_backend.py

@@ -49,6 +49,9 @@ def init_forward_metadata(self, forward_batch: ForwardBatch):
             max_extend_len = torch.max(forward_batch.extend_seq_lens).item()
         self.forward_metadata = (attn_logits, max_extend_len)
 
+    def get_graph_seq_len_fill_value(self):
+        return 1
+
     def forward_extend(
         self,
         q,

python/sglang/srt/layers/quantization/fp8.py

@@ -344,6 +344,9 @@ def process_weights_after_loading(self, layer: Module) -> None:
                     _is_cpu_amx_available
                 ), "Fp8LinearMethod on CPU requires that CPU has AMX support"
                 _amx_process_weight_after_loading(layer, ["weight"])
+                layer.weight_scale_inv = torch.nn.Parameter(
+                    layer.weight_scale_inv.data, requires_grad=False
+                )
                 return
             else:
                 weight, weight_scale = layer.weight.data, layer.weight_scale_inv.data

python/sglang/srt/layers/quantization/w8a8_int8.py

@@ -339,9 +339,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
                 _is_cpu_amx_available
             ), "W8A8Int8LinearMethod on CPU requires that CPU has AMX support"
             _amx_process_weight_after_loading(layer, ["weight"])
-            return
-
-        layer.weight = Parameter(layer.weight.t(), requires_grad=False)
+        else:
+            layer.weight = Parameter(layer.weight.t(), requires_grad=False)
         layer.weight_scale = Parameter(layer.weight_scale.data, requires_grad=False)
 
     def create_weights(
@@ -472,10 +471,9 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
                 _is_cpu_amx_available
             ), "W8A8Int8MoEMethod on CPU requires that CPU has AMX support"
             _amx_process_weight_after_loading(layer, ["w13_weight", "w2_weight"])
-            return
-
-        layer.w13_weight = Parameter(layer.w13_weight, requires_grad=False)
-        layer.w2_weight = Parameter(layer.w2_weight, requires_grad=False)
+        else:
+            layer.w13_weight = Parameter(layer.w13_weight, requires_grad=False)
+            layer.w2_weight = Parameter(layer.w2_weight, requires_grad=False)
         layer.w13_weight_scale = Parameter(
             layer.w13_weight_scale.data, requires_grad=False
         )

python/sglang/srt/managers/scheduler.py

@@ -181,7 +181,7 @@ class GenerationBatchResult:
     extend_input_len_per_req: List[int]
     extend_logprob_start_len_per_req: List[int]
     bid: int
-    can_run_cuda_graph: bool
+    can_run_graph: bool
 
 
 @dataclass
@@ -398,7 +398,7 @@ def __init__(
                 f"max_prefill_tokens={self.max_prefill_tokens}, "
                 f"max_running_requests={self.max_running_requests}, "
                 f"context_len={self.model_config.context_len}, "
-                f"available_gpu_mem={avail_mem:.2f} GB"
+                f"{'available_cpu_mem' if self.device == 'cpu' else 'available_gpu_mem'}={avail_mem:.2f} GB"
             )
 
         # Init memory pool and cache
@@ -929,7 +929,7 @@ def event_loop_pp(self):
                             "extend_logprob_start_len_per_req", None
                         ),
                         bid=bids[next_mb_id],
-                        can_run_cuda_graph=result.can_run_cuda_graph,
+                        can_run_graph=result.can_run_graph,
                     )
                     self.process_batch_result(mbs[next_mb_id], output_result)
                     last_mbs[next_mb_id] = mbs[next_mb_id]
@@ -1778,11 +1778,11 @@ def run_batch(
                     model_worker_batch.hicache_consumer_index
                 )
                 if self.pp_group.is_last_rank:
-                    logits_output, next_token_ids, can_run_cuda_graph = (
+                    logits_output, next_token_ids, can_run_graph = (
                         self.tp_worker.forward_batch_generation(model_worker_batch)
                     )
                 else:
-                    pp_hidden_states_proxy_tensors, _, can_run_cuda_graph = (
+                    pp_hidden_states_proxy_tensors, _, can_run_graph = (
                         self.tp_worker.forward_batch_generation(model_worker_batch)
                     )
                 bid = model_worker_batch.bid
@@ -1792,7 +1792,7 @@ def run_batch(
                     next_token_ids,
                     bid,
                     num_accepted_tokens,
-                    can_run_cuda_graph,
+                    can_run_graph,
                 ) = self.draft_worker.forward_batch_speculative_generation(batch)
                 bs = batch.batch_size()
                 self.spec_num_total_accepted_tokens += num_accepted_tokens + bs
@@ -1827,7 +1827,7 @@ def run_batch(
                 extend_input_len_per_req=extend_input_len_per_req,
                 extend_logprob_start_len_per_req=extend_logprob_start_len_per_req,
                 bid=bid,
-                can_run_cuda_graph=can_run_cuda_graph,
+                can_run_graph=can_run_graph,
             )
         else:  # embedding or reward model
             model_worker_batch = batch.get_model_worker_batch()
@@ -2288,10 +2288,9 @@ def get_internal_state(self, recv_req: GetInternalStateReq):
             "token_capacity": int(self.max_total_num_tokens),
         }
 
-        if not _is_cpu:
-            ret["memory_usage"]["cuda_graph"] = round(
-                self.tp_worker.worker.model_runner.cuda_graph_mem_usage, 2
-            )
+        ret["memory_usage"]["graph"] = round(
+            self.tp_worker.worker.model_runner.graph_mem_usage, 2
+        )
 
         if not self.spec_algorithm.is_none() and self.cum_spec_accept_count > 0:
             ret["avg_spec_accept_length"] = (

python/sglang/srt/managers/scheduler_metrics_mixin.py

@@ -138,7 +138,7 @@ def log_prefill_stats(
         self._publish_kv_events()
 
     def log_decode_stats(
-        self, can_run_cuda_graph: bool, running_batch: ScheduleBatch = None
+        self, can_run_graph: bool, running_batch: ScheduleBatch = None
     ):
         batch = running_batch or self.running_batch
 
@@ -193,7 +193,7 @@ def log_decode_stats(
             msg += f"#retracted-req: {len(self.disagg_decode_prealloc_queue.retracted_queue)}, "
 
         msg += (
-            f"cuda graph: {can_run_cuda_graph}, "
+            f"{'cpu graph' if self.device == 'cpu' else 'cuda graph'}: {can_run_graph}, "
             f"gen throughput (token/s): {self.last_gen_throughput:.2f}, "
             f"#queue-req: {len(self.waiting_queue)}, "
         )

python/sglang/srt/managers/scheduler_output_processor_mixin.py

@@ -197,15 +197,15 @@ def process_batch_result_decode(
         result: GenerationBatchResult,
         launch_done: Optional[threading.Event] = None,
     ):
-        logits_output, next_token_ids, can_run_cuda_graph = (
+        logits_output, next_token_ids, can_run_graph = (
             result.logits_output,
             result.next_token_ids,
-            result.can_run_cuda_graph,
+            result.can_run_graph,
         )
         self.num_generated_tokens += len(batch.reqs)
 
         if self.enable_overlap:
-            logits_output, next_token_ids, can_run_cuda_graph = (
+            logits_output, next_token_ids, can_run_graph = (
                 self.tp_worker.resolve_last_batch_result(launch_done)
             )
             next_token_logprobs = logits_output.next_token_logprobs
@@ -293,7 +293,7 @@ def process_batch_result_decode(
             self.current_scheduler_metrics_enabled()
             and self.forward_ct_decode % self.server_args.decode_log_interval == 0
         ):
-            self.log_decode_stats(can_run_cuda_graph, running_batch=batch)
+            self.log_decode_stats(can_run_graph, running_batch=batch)
 
     def add_input_logprob_return_values(
         self: Scheduler,

python/sglang/srt/managers/tp_worker.py

@@ -236,7 +236,7 @@ def forward_batch_generation(
             )
 
         if self.pp_group.is_last_rank:
-            logits_output, can_run_cuda_graph = self.model_runner.forward(
+            logits_output, can_run_graph = self.model_runner.forward(
                 forward_batch, pp_proxy_tensors=pp_proxy_tensors
             )
             if launch_done is not None:
@@ -249,13 +249,13 @@ def forward_batch_generation(
                     logits_output, model_worker_batch
                 )
 
-            return logits_output, next_token_ids, can_run_cuda_graph
+            return logits_output, next_token_ids, can_run_graph
         else:
-            pp_proxy_tensors, can_run_cuda_graph = self.model_runner.forward(
+            pp_proxy_tensors, can_run_graph = self.model_runner.forward(
                 forward_batch,
                 pp_proxy_tensors=pp_proxy_tensors,
             )
-            return pp_proxy_tensors.tensors, None, can_run_cuda_graph
+            return pp_proxy_tensors.tensors, None, can_run_graph
 
     def forward_batch_embedding(self, model_worker_batch: ModelWorkerBatch):
         forward_batch = ForwardBatch.init_new(model_worker_batch, self.model_runner)

python/sglang/srt/managers/tp_worker_overlap_thread.py

@@ -172,7 +172,7 @@ def forward_thread_func_(self):
             # update the consumer index of hicache to the running batch
             self.set_hicache_consumer(model_worker_batch.hicache_consumer_index)
             # Run forward
-            logits_output, next_token_ids, can_run_cuda_graph = (
+            logits_output, next_token_ids, can_run_graph = (
                 self.worker.forward_batch_generation(
                     model_worker_batch, model_worker_batch.launch_done
                 )
@@ -201,15 +201,15 @@ def forward_thread_func_(self):
             copy_done.record()
 
             self.output_queue.put(
-                (copy_done, logits_output, next_token_ids, can_run_cuda_graph)
+                (copy_done, logits_output, next_token_ids, can_run_graph)
             )
 
     def resolve_last_batch_result(self, launch_done: Optional[threading.Event] = None):
         """
         This function is called to resolve the last batch result and
         wait for the current batch to be launched. Used in overlap mode.
         """
-        copy_done, logits_output, next_token_ids, can_run_cuda_graph = (
+        copy_done, logits_output, next_token_ids, can_run_graph = (
             self.output_queue.get()
         )
 
@@ -226,7 +226,7 @@ def resolve_last_batch_result(self, launch_done: Optional[threading.Event] = Non
                     logits_output.input_token_logprobs.tolist()
                 )
         next_token_ids = next_token_ids.tolist()
-        return logits_output, next_token_ids, can_run_cuda_graph
+        return logits_output, next_token_ids, can_run_graph
 
     def forward_batch_generation(
         self, model_worker_batch: ModelWorkerBatch