Support piecewise cuda graph for MLA (#11812)

python/sglang/srt/compilation/compile.py

@@ -112,6 +112,9 @@ def _infer_dynamic_arg_dims_from_annotations(forward_fn):
             for a in getattr(ann, "__args__", [])
         ):
             dyn[name] = 0
+        elif ann == "torch.Tensor" or ann == "Optional[torch.Tensor]":
+            # For future import annotations (e.g. from __future__ import annotations), the annotation is a string
+            dyn[name] = 0
     if not dyn:
         raise ValueError("No dynamic dims inferred; pass dynamic_arg_dims explicitly.")
     return dyn

python/sglang/srt/compilation/piecewise_context_manager.py

@@ -30,11 +30,10 @@ def get_forward_context() -> Optional[ForwardContext]:
 @contextmanager
 def set_forward_context(forward_batch: ForwardBatch, attention_layers: List[Any]):
     global _forward_context
-    prev_forward_context = _forward_context
     _forward_context = ForwardContext()
     _forward_context.set_forward_batch(forward_batch)
     _forward_context.set_attention_layers(attention_layers)
     try:
         yield
     finally:
-        _forward_context = prev_forward_context
+        _forward_context = None