[refactor] Split the DP gathered-buffer state between flags.dp and ctx.forward (#30491)

python/sglang/srt/layers/dp_attention.py

@@ -100,10 +100,14 @@ class DpPaddingMode(IntEnum):
 
 
 class _DpGatheredBufferWrapper:
+    """Facade for the DP gathered-buffer state: allocation metadata lives on
+    ``flags.dp`` (set once at initialize_dp_attention). The per-forward
+    sizing quartet stays as class attributes: the values are read inside
+    torch.compile-traced model code, and attribute-source ints get dynamo's
+    automatic-dynamic treatment, while contextvars are untraceable and dict
+    slots value-guard into the recompile limit (one recompile per distinct
+    size)."""
 
-    _hidden_size: int
-    _dtype: torch.dtype
-    _device: torch.device
     _global_dp_buffer_len: int
     _local_dp_buffer_len: int
     _dp_max_padding: bool
@@ -111,9 +115,12 @@ class _DpGatheredBufferWrapper:
 
     @classmethod
     def set_metadata(cls, hidden_size: int, dtype: torch.dtype, device: torch.device):
-        cls._hidden_size = hidden_size
-        cls._dtype = dtype
-        cls._device = device
+        from sglang.srt.runtime_context import get_flags
+
+        dp = get_flags().dp
+        dp.buffer_hidden_size = hidden_size
+        dp.buffer_dtype = dtype
+        dp.buffer_device = device
 
     @classmethod
     def set_dp_buffer_len(
@@ -130,21 +137,27 @@ class _DpGatheredBufferWrapper:
 
     @classmethod
     def get_global_dp_buffer(cls, group: GroupCoordinator) -> torch.Tensor:
+        from sglang.srt.runtime_context import get_flags
+
+        dp = get_flags().dp
         with use_symmetric_memory(group, disabled=not cls._dp_max_padding):
             buffer = torch.empty(
-                (cls._global_dp_buffer_len, cls._hidden_size),
-                dtype=cls._dtype,
-                device=cls._device,
+                (cls._global_dp_buffer_len, dp.buffer_hidden_size),
+                dtype=dp.buffer_dtype,
+                device=dp.buffer_device,
             )
         return buffer
 
     @classmethod
     def get_local_dp_buffer(cls, group: GroupCoordinator) -> torch.Tensor:
+        from sglang.srt.runtime_context import get_flags
+
+        dp = get_flags().dp
         with use_symmetric_memory(group, disabled=not cls._dp_max_padding):
             buffer = torch.empty(
-                (cls._local_dp_buffer_len, cls._hidden_size),
-                dtype=cls._dtype,
-                device=cls._device,
+                (cls._local_dp_buffer_len, dp.buffer_hidden_size),
+                dtype=dp.buffer_dtype,
+                device=dp.buffer_device,
             )
         return buffer
 
@@ -162,15 +175,21 @@ class _DpGatheredBufferWrapper:
 
     @classmethod
     def get_dp_hidden_size(cls) -> int:
-        return cls._hidden_size
+        from sglang.srt.runtime_context import get_flags
+
+        return get_flags().dp.buffer_hidden_size
 
     @classmethod
     def get_dp_dtype(cls) -> torch.dtype:
-        return cls._dtype
+        from sglang.srt.runtime_context import get_flags
+
+        return get_flags().dp.buffer_dtype
 
     @classmethod
     def get_dp_device(cls) -> torch.device:
-        return cls._device
+        from sglang.srt.runtime_context import get_flags
+
+        return get_flags().dp.buffer_device
 
     @classmethod
     def is_dp_max_padding(cls) -> bool:

python/sglang/srt/models/deepseek_v4.py

@@ -53,7 +53,6 @@ from sglang.srt.layers.deepseek_v4_rope import (
     v4_rope_inplace_npu,
 )
 from sglang.srt.layers.dp_attention import (
-    _DpGatheredBufferWrapper,
     _tbo_event,
     attn_tp_all_gather,
     attn_tp_all_reduce,
@@ -2055,7 +2054,7 @@ class DeepseekV4Model(nn.Module):
 
         if get_parallel().attn_dp_size > 1 and get_moe_a2a_backend().is_none():
             input_ids_global = torch.empty(
-                (_DpGatheredBufferWrapper._global_dp_buffer_len, 1),
+                (get_global_dp_buffer_len(), 1),
                 dtype=input_ids.dtype,
                 device=input_ids.device,
             )

python/sglang/srt/models/deepseek_v4_nextn.py

@@ -14,8 +14,8 @@ from sglang.srt.layers.attention.dsa.utils import (
     is_dsa_prefill_cp_round_robin_split,
 )
 from sglang.srt.layers.dp_attention import (
-    _DpGatheredBufferWrapper,
     dp_gather_partial,
+    get_global_dp_buffer_len,
     is_dp_attention_enabled,
 )
 from sglang.srt.layers.layernorm import RMSNorm
@@ -158,7 +158,7 @@ class DeepseekV4ModelNextN(nn.Module):
 
         if get_parallel().attn_dp_size > 1 and get_moe_a2a_backend().is_none():
             input_ids_global = torch.empty(
-                (_DpGatheredBufferWrapper._global_dp_buffer_len, 1),
+                (get_global_dp_buffer_len(), 1),
                 dtype=input_ids.dtype,
                 device=input_ids.device,
             )

python/sglang/srt/runtime_context.py

@@ -304,6 +304,11 @@ class DpFlags(_FlagGroupBase):
     # Hybrid-SSM models materialize idle ranks via the MAX_LEN fabricated-row
     # conversion (set when hf_config has hybrid_override_pattern).
     max_len_with_idle: bool = False
+    # DP gathered-buffer allocation metadata (model hidden size / dtype /
+    # device), set by initialize_dp_attention alongside the flags above.
+    buffer_hidden_size: Any = None
+    buffer_dtype: Any = None
+    buffer_device: Any = None
 
 
 @dataclasses.dataclass