diff --git a/python/sglang/srt/layers/dp_attention.py b/python/sglang/srt/layers/dp_attention.py index c513cbdc4..f693a9d48 100644 --- a/python/sglang/srt/layers/dp_attention.py +++ b/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: diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index e66508ea4..026e4e0ce 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/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, ) diff --git a/python/sglang/srt/models/deepseek_v4_nextn.py b/python/sglang/srt/models/deepseek_v4_nextn.py index fe9d9b197..1dd326c6e 100644 --- a/python/sglang/srt/models/deepseek_v4_nextn.py +++ b/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, ) diff --git a/python/sglang/srt/runtime_context.py b/python/sglang/srt/runtime_context.py index bf75292d8..5a0ebed53 100644 --- a/python/sglang/srt/runtime_context.py +++ b/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 diff --git a/test/registered/unit/test_runtime_context.py b/test/registered/unit/test_runtime_context.py index e52cdbced..491c6db94 100644 --- a/test/registered/unit/test_runtime_context.py +++ b/test/registered/unit/test_runtime_context.py @@ -634,6 +634,33 @@ class TestForwardFlags(_IsolatedServerArgs): self.assertIsNone(get_forward().attn_inputs) self.assertFalse(ctx.input_scattered) + def test_dp_buffer_state_split(self): + import torch + + from sglang.srt.layers.dp_attention import _DpGatheredBufferWrapper as wrapper + from sglang.srt.layers.dp_attention import ( + get_dp_dtype, + get_dp_global_num_tokens, + get_global_dp_buffer_len, + is_dp_max_padding, + set_dp_buffer_len, + ) + + reset_context() + # metadata is init-static (flags.dp); sizing is per-forward sticky + wrapper.set_metadata(64, torch.float16, torch.device("cpu")) + self.assertEqual(get_dp_dtype(), torch.float16) + set_dp_buffer_len(128, 32, True, [64, 64]) + self.assertEqual(get_global_dp_buffer_len(), 128) + self.assertTrue(is_dp_max_padding()) + self.assertEqual(get_dp_global_num_tokens(), [64, 64]) + set_dp_buffer_len(256, 64, False) # sticky until the next write + self.assertEqual(get_global_dp_buffer_len(), 256) + self.assertFalse(is_dp_max_padding()) + self.assertIsNone(get_dp_global_num_tokens()) + reset_context() + self.assertIsNone(get_dp_dtype()) + def test_is_extend_in_batch_sticky_within_thread(self): from sglang.srt.layers.dp_attention import ( get_is_extend_in_batch,