Fix PCG MoE Error (#17739)

python/sglang/srt/compilation/compile.py

@@ -88,12 +88,12 @@ class _MaybeIntermediateTensors:
 
 def _mark_dynamic_on_value(val, dims):
     if isinstance(val, torch.Tensor):
-        torch._dynamo.mark_dynamic(val, _normalize_dims(dims, val.ndim))
+        torch._dynamo.maybe_mark_dynamic(val, _normalize_dims(dims, val.ndim))
     else:
         mit = _MaybeIntermediateTensors(val)
         if mit.is_intermediate:
             for t in mit.obj.tensors.values():
-                torch._dynamo.mark_dynamic(t, _normalize_dims(dims, t.ndim))
+                torch._dynamo.maybe_mark_dynamic(t, _normalize_dims(dims, t.ndim))
         # else: ignore (None or non-tensor)
 
 

python/sglang/srt/compilation/piecewise_context_manager.py

@@ -1,10 +1,13 @@
+from __future__ import annotations
+
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import Any, List, Optional
+from typing import TYPE_CHECKING, Any, List, Optional
 
 import torch
 
-from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 
 _in_piecewise_cuda_graph = False
 _in_pcg_torch_compile = False