minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Give the attention-DP width and rank one home (#40067)

Browse Source
This commit is contained in:
Cheng Wan
2026-09-18 17:34:33 -07:00
committed by GitHub
parent 8ea0ee300d
commit d0730a0e8b
11 changed files with 174 additions and 115 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/debug_utils/dumper.py
+5 -4
View File
@@ -1739,7 +1739,7 @@ class _SGLangPlugin(_FrameworkPlugin):
info["moe_tp_size"] = parallel.moe_tp_size
info["moe_dp_rank"] = parallel.moe_dp_rank
info["moe_dp_size"] = self._dp_attn.get_moe_cp_size()
except (AttributeError, AssertionError, ValueError):
except (AttributeError, AssertionError, ValueError, RuntimeError):
info["distributed_error"] = True
try:
@@ -1747,11 +1747,12 @@ class _SGLangPlugin(_FrameworkPlugin):
info["enable_dp_attention"] = self._dp_attn.is_dp_attention_enabled()
info["attn_tp_rank"] = parallel.attn_tp_rank
info["attn_tp_size"] = parallel.attn_tp_size
info["attn_dp_rank"] = self._dp_attn.get_attention_dp_rank()
info["attn_dp_size"] = self._dp_attn.get_attention_dp_size()
info["attn_dp_rank"] = parallel.attn_dp_rank
info["attn_dp_size"] = parallel.attn_dp_size
info["attn_cp_rank"] = parallel.attn_cp_rank
info["attn_cp_size"] = parallel.attn_cp_size
except (AttributeError, AssertionError, ValueError):
# An unstamped topology name raises RuntimeError.
except (AttributeError, AssertionError, ValueError, RuntimeError):
info["dp_attention_error"] = True
return info
python/sglang/srt/disaggregation/common/conn.py
+2 -6
View File
@@ -33,10 +33,6 @@ from sglang.srt.disaggregation.utils import (
)
from sglang.srt.distributed import get_pp_group, get_world_group
from sglang.srt.environ import envs
from sglang.srt.layers.dp_attention import (
get_attention_dp_rank,
get_attention_dp_size,
)
from sglang.srt.runtime_context import (
get_disagg,
get_parallel,
@@ -197,8 +193,8 @@ class CommonKVManager(BaseKVManager):
self.attn_cp_rank = parallel.attn_cp_rank
self.dcp_size = parallel.attn_dcp_size
self.dcp_rank = parallel.attn_dcp_rank
self.attn_dp_size = get_attention_dp_size()
self.attn_dp_rank = get_attention_dp_rank()
self.attn_dp_size = parallel.attn_dp_size
self.attn_dp_rank = parallel.attn_dp_rank
self.system_dp_size = (
1 if get_parallel().enable_dp_attention else get_parallel().dp_size
)
python/sglang/srt/layers/dp_attention.py
+18 -51
View File
@@ -2,7 +2,6 @@ from __future__ import annotations
import functools
import logging
from contextlib import contextmanager
from enum import IntEnum, auto
from typing import TYPE_CHECKING, List, Optional, Tuple
@@ -53,9 +52,6 @@ logger = logging.getLogger(__name__)
if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
_ATTN_DP_RANK: Optional[int] = None
_ATTN_DP_SIZE: Optional[int] = None
def world_dp_gather_enabled() -> bool:
"""Whether DP gathers should use expanded WORLD after joiner admission."""
@@ -68,10 +64,9 @@ def enable_joiner_all_gather():
def update_dp_attention_post_scale(new_dp_size: int, new_dp_rank: int):
global _ATTN_DP_SIZE, _ATTN_DP_RANK
_ATTN_DP_SIZE = new_dp_size
_ATTN_DP_RANK = new_dp_rank
get_parallel().override_permanently(attn_dp_size=new_dp_size)
get_parallel().override_permanently(
attn_dp_size=new_dp_size, attn_dp_rank=new_dp_rank
)
get_flags().dp.use_world_group_for_gather = True
logger.debug(
"[Elastic EP] dp_attention switched to WORLD: dp_size=%d dp_rank=%d",
@@ -101,7 +96,7 @@ class DpPaddingMode(IntEnum):
def get_dp_padding_mode(
cls, is_extend_in_batch, global_num_tokens: List[int]
) -> DpPaddingMode:
dp_size = get_attention_dp_size()
dp_size = get_parallel().attn_dp_size
# (trangdough) pplx-kernels a2a is a symmetric collective: every EP rank
# must dispatch the same number of tokens or the device-side handshake
@@ -375,7 +370,6 @@ def initialize_dp_attention(
server_args: ServerArgs,
model_config: ModelConfig,
):
global _ATTN_DP_RANK, _ATTN_DP_SIZE
dp = get_flags().dp
dp.max_len_with_idle = (
getattr(model_config.hf_config, "hybrid_override_pattern", None) is not None
@@ -389,13 +383,12 @@ def initialize_dp_attention(
tp_rank = get_tensor_model_parallel_rank()
tp_size = get_tensor_model_parallel_world_size()
_, _, _ATTN_DP_RANK, _ATTN_DP_SIZE = compute_dp_attention_world_info(
_, _, attn_dp_rank, attn_dp_size = compute_dp_attention_world_info(
enable_dp_attention, tp_rank, tp_size, dp_size, attn_cp_size
)
get_parallel().override_permanently(attn_dp_size=_ATTN_DP_SIZE)
if get_exec().moe.elastic_ep_backend is not None and get_parallel().max_ep_size:
_ATTN_DP_RANK = tp_rank + get_parallel().ep_join_rank_offset
attn_dp_rank = tp_rank + get_parallel().ep_join_rank_offset
# Reads the resolution, not a bag: this runs under
# `initialize_dp_attention`, which the weight-cache daemon calls from
# `_init_distributed` -- and other callers reach it from processes
@@ -404,6 +397,13 @@ def initialize_dp_attention(
if ep_scale_joiner_of(resolving_view(server_args)):
dp.joiner_skip_all_gather = True
# Stamped together, after the elastic adjustment: the width and the rank
# describe one topology, and a reader that caught them mid-update would
# see this process placed in a group it is not in.
get_parallel().override_permanently(
attn_dp_size=attn_dp_size, attn_dp_rank=attn_dp_rank
)
_DpGatheredBufferWrapper.set_metadata(
hidden_size=model_config.hidden_size,
dtype=model_config.dtype,
@@ -419,42 +419,9 @@ def is_allocation_symmetric() -> bool:
return not is_dp_attention_enabled() or is_dp_max_padding()
def get_attention_dp_rank() -> int:
assert _ATTN_DP_RANK is not None, "dp attention not initialized!"
return _ATTN_DP_RANK
def get_attention_dp_size() -> int:
assert _ATTN_DP_SIZE is not None, "dp attention not initialized!"
return _ATTN_DP_SIZE
@contextmanager
def disable_dp_size():
"""Run without DP attention until this scope ends.
This is for draft workers of speculative decoding, which run the draft model
at a different width from the target model's workers.
The scope replaces both the module global that ``get_attention_dp_size()``
reads and the derived width the runtime context answers with, so the two
spellings of the name cannot disagree inside it.
"""
global _ATTN_DP_SIZE
assert _ATTN_DP_SIZE is not None, "dp attention not initialized!"
old_dp_size = _ATTN_DP_SIZE
_ATTN_DP_SIZE = 1
try:
with get_parallel().override(attn_dp_size=1):
yield
finally:
_ATTN_DP_SIZE = old_dp_size
def get_dp_local_info(forward_batch: ForwardBatch) -> Tuple[torch.Tensor, torch.Tensor]:
# `get_dp_local_info` is only called in global DP gather and scatter. We use global DP rank here.
dp_rank = get_attention_dp_rank()
dp_rank = get_parallel().attn_dp_rank
if forward_batch.dp_local_start_pos is None:
cumtokens = torch.cumsum(forward_batch.global_num_tokens_gpu, dim=0)
@@ -478,7 +445,7 @@ def get_dp_local_slice_cpu(
# CPU (start, length) slice for DP-local data in a rank-padded buffer.
# Returns Python ints (no D2H sync) and handles the cuda-graph-padded layout.
global_num_tokens = forward_batch.global_num_tokens_cpu
dp_rank = get_attention_dp_rank()
dp_rank = get_parallel().attn_dp_rank
local_num_tokens = global_num_tokens[dp_rank]
if can_run_graph:
local_start_pos = dp_rank * cuda_graph_batch
@@ -758,7 +725,7 @@ def is_dp_gatherv_active() -> bool:
_USE_DP_GATHERV
and not world_dp_gather_enabled()
and get_attn_tensor_model_parallel_world_size() == 1
and get_tensor_model_parallel_world_size() == get_attention_dp_size()
and get_tensor_model_parallel_world_size() == get_parallel().attn_dp_size
and not _DpGatheredBufferWrapper.is_dp_max_padding()
)
@@ -792,7 +759,7 @@ def _dp_gather_via_all_gatherv(
# each rank's local tensor up to sizes[rank] with zeros (matching the
# buffer's reserved per-rank slot) so sum(sizes) == buffer rows and there
# is no uninitialized tail for the MoE to read.
rank = get_attention_dp_rank()
rank = get_parallel().attn_dp_rank
local_rows = sizes[rank]
if local_tokens.shape[0] == local_rows:
local_real = local_tokens
@@ -917,7 +884,7 @@ def dp_reduce_scatter_tensor(output: torch.Tensor, input: torch.Tensor):
if sizes is not None:
get_tp_group().reduce_scatterv(input, output=output, sizes=sizes)
return
if get_tensor_model_parallel_world_size() == get_attention_dp_size():
if get_tensor_model_parallel_world_size() == get_parallel().attn_dp_size:
get_tp_group().reduce_scatter_tensor(output, input)
else:
scattered_local_tokens = input.tensor_split(
python/sglang/srt/layers/engram.py
+2 -3
View File
@@ -40,7 +40,6 @@ from sglang.srt.layers.dp_attention import (
dp_gather_replicate,
dp_reduce_scatter_tensor,
dp_scatter,
get_attention_dp_size,
get_global_dp_buffer_len,
is_dp_gatherv_active,
)
@@ -766,7 +765,7 @@ class EngramEmbedding(nn.Module):
attn_cp_all_gather_into_tensor(all_indices, indices.contiguous())
start = parallel.attn_cp_rank * local_rows
return self._lookup(all_indices)[start : start + local_rows]
if self.tp_size > 1 and get_attention_dp_size() > 1:
if self.tp_size > 1 and get_parallel().attn_dp_size > 1:
return self._dp_sharded_lookup(indices, forward_batch)
return self._lookup(indices)
@@ -832,7 +831,7 @@ class EngramEmbedding(nn.Module):
if (
padding is not None
and padding.is_max_len()
and self.tp_size == get_attention_dp_size()
and self.tp_size == get_parallel().attn_dp_size
and rows == self.tp_size * local.shape[0]
) or is_dp_gatherv_active():
dp_reduce_scatter_tensor(local, values)
python/sglang/srt/managers/cache_controller.py
+1 -2
View File
@@ -37,7 +37,6 @@ if TYPE_CHECKING:
from sglang.srt.mem_cache.pool_host import HostKVCache
from sglang.srt.layers.dp_attention import (
get_attention_dp_rank,
is_dp_attention_enabled,
)
from sglang.srt.mem_cache.l2_transfer import L2Transfer, L2TransferEngine
@@ -696,7 +695,7 @@ class HiCacheController:
if is_dp_attention_enabled():
self.tp_rank = get_parallel().attn_tp_rank
self.tp_size = get_parallel().attn_tp_size
self.dp_rank = get_attention_dp_rank()
self.dp_rank = get_parallel().attn_dp_rank
else:
self.tp_rank = get_parallel().tp_rank
self.tp_size = get_parallel().tp_size
python/sglang/srt/managers/scheduler_components/dp_attn.py
+1 -2
View File
@@ -58,9 +58,8 @@ def _resolve_elastic_world_dp_size(
return dp_size
from sglang.srt.elastic_ep.elastic_ep import ElasticEPStateManager
from sglang.srt.layers.dp_attention import get_attention_dp_size
live_dp_size = get_attention_dp_size()
live_dp_size = get_parallel().attn_dp_size
effective_ep_size = ElasticEPStateManager.get_effective_ep_size()
world_size = torch.distributed.get_world_size(group)
python/sglang/srt/managers/scheduler_components/dynamic_chunk_sizer.py
+3 -4
View File
@@ -12,14 +12,13 @@ from tqdm import tqdm
from sglang.srt.environ import envs
from sglang.srt.layers.dp_attention import (
get_attention_dp_rank,
get_attention_dp_size,
is_dp_attention_enabled,
set_is_extend_in_batch,
)
from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
from sglang.srt.mem_cache.common import release_kv_cache
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.runtime_context import get_parallel
from sglang.srt.sampling.sampling_params import SamplingParams
from sglang.srt.utils import broadcast_pyobj
from sglang.srt.utils.common import get_device_module
@@ -196,9 +195,9 @@ class DynamicChunkSizer:
if is_dp_attention_enabled():
# Profiling runs one request on this rank; other DP ranks report 0.
dp_size = get_attention_dp_size()
dp_size = get_parallel().attn_dp_size
global_num_tokens = [0] * dp_size
dp_rank = get_attention_dp_rank()
dp_rank = get_parallel().attn_dp_rank
global_num_tokens[dp_rank] = current_seq_len
batch.global_num_tokens = global_num_tokens
batch.global_num_tokens_for_logprob = global_num_tokens
python/sglang/srt/mem_cache/storage/umbp/umbp_store.py
+8 -10
View File
@@ -310,15 +310,14 @@ class UMBPStore(HiCacheStorage):
if dp_rank_hint is None:
try:
from sglang.srt.layers.dp_attention import (
get_attention_dp_rank,
get_attention_dp_size,
is_dp_attention_enabled,
)
from sglang.srt.runtime_context import get_parallel
if is_dp_attention_enabled():
dp_rank_hint = get_attention_dp_rank()
dp_size_hint = get_attention_dp_size()
except (ImportError, AssertionError):
dp_rank_hint = get_parallel().attn_dp_rank
dp_size_hint = get_parallel().attn_dp_size
except (ImportError, AssertionError, RuntimeError):
pass
if local_rank_hint is not None:
@@ -807,14 +806,13 @@ class UMBPStore(HiCacheStorage):
try:
from sglang.srt.layers.dp_attention import (
get_attention_dp_rank,
get_attention_dp_size,
is_dp_attention_enabled,
)
from sglang.srt.runtime_context import get_parallel
if is_dp_attention_enabled():
dp_rank = get_attention_dp_rank()
dp_size = get_attention_dp_size()
dp_rank = get_parallel().attn_dp_rank
dp_size = get_parallel().attn_dp_size
dp_rank_hint = dp_rank
dp_size_hint = dp_size
if cfg.ssd.enabled:
@@ -883,7 +881,7 @@ class UMBPStore(HiCacheStorage):
dp_size,
cfg.ssd.storage_dir,
)
except (ImportError, AssertionError):
except (ImportError, AssertionError, RuntimeError):
pass
if (
python/sglang/srt/models/qwen4_exp.py
+2 -3
View File
@@ -27,7 +27,6 @@ from sglang.srt.layers.dp_attention import (
attn_tp_all_reduce,
dp_gather_replicate,
dp_scatter,
get_attention_dp_size,
get_dp_global_num_tokens,
get_global_dp_buffer,
get_local_dp_buffer,
@@ -505,7 +504,7 @@ class Qwen4ExpNGramEmbedding(nn.Module):
self.use_attn_tp_ngram = _use_attn_tp_ngram()
self.gather_dp_tokens = (
is_dp_attention_enabled()
and get_attention_dp_size() > 1
and get_parallel().attn_dp_size > 1
and not self.use_attn_tp_ngram
)
ngram_prefix = f"{prefix}.ngram_embedding" if prefix else "ngram_embedding"
@@ -1360,7 +1359,7 @@ class Qwen4ExpLayerExtensionMixin:
return hidden_states, residual
def _qwen4_exp_use_dp_moe_gather(self) -> bool:
return get_attention_dp_size() > 1 and get_moe_a2a_backend().is_none()
return get_parallel().attn_dp_size > 1 and get_moe_a2a_backend().is_none()
def _qwen4_exp_use_attn_tp_a2a_scatter(self) -> bool:
return get_parallel().attn_tp_size > 1 and not get_moe_a2a_backend().is_none()
python/sglang/srt/runtime_context.py
+48 -12
View File
@@ -290,22 +290,53 @@ class ParallelContext:
raise AttributeError(f"ParallelContext has no {name!r}")
def _v(self, name, getter):
overrides = self._overrides
return overrides[name] if name in overrides else getter()
"""Scoped override, else the permanent stamp, else the live group.
def override_permanently(self, **widths) -> None:
"""Permanently correct a derived width the published bag can't answer
One priority order for ranks and widths alike (`_derived_width`), so a
stamped value wins over the coordinator for both.
"""
overrides = self._overrides
if name in overrides:
return overrides[name]
derived = self._derived
if name in derived:
return derived[name]
return getter()
def _stamped(self, name, why):
"""A per-process fact no configuration implies: scoped override, else
the permanent stamp, else fail.
Unlike a width, this has nothing to fall back on -- the configuration
does not carry this process's rank, and there is no group to ask --
so an unstamped read is a missing initialization rather than a
missing override, and says so.
"""
overrides = self._overrides
if name in overrides:
return overrides[name]
derived = self._derived
if name in derived:
return derived[name]
raise RuntimeError(f"parallel rank {name!r} is not available: {why}")
def override_permanently(self, **values) -> None:
"""Permanently record a width or rank the published bag can't answer
or no longer answers correctly -- not `RuntimeContext.override`,
because a derived width is not a resolved config leaf and this must
work with no config published at all (`multimodal_gen` lends a TP
group to `srt` layers with no `srt` config to publish against).
because neither is a resolved config leaf and this must work with no
config published at all (`multimodal_gen` lends a TP group to `srt`
layers with no `srt` config to publish against).
Widths are quotients of the configured leaves, so the bag can usually
answer and this only corrects it; a rank is a per-process fact the
configuration never carries, so for those this is the only source.
Lives beside, not inside, the `@contextmanager` `override` above -- a
name it cannot also have on this class -- because these are permanent
for the process, not scoped to a `with` block: none of the real
callers ever restore the value they set here.
"""
self._derived.update(widths)
self._derived.update(values)
def clear_derived_widths(self) -> None:
self._derived.clear()
@@ -401,7 +432,12 @@ class ParallelContext:
@property
def attn_dp_rank(self) -> int:
return self._v("attn_dp_rank", _dp().get_attention_dp_rank)
return self._stamped(
"attn_dp_rank",
"it is computed from this process's `tp_rank` when the attention "
"topology is initialized, so a process that never ran "
"`initialize_dp_attention` has no answer to give",
)
@property
def world_group(self) -> Any:
@@ -560,9 +596,9 @@ class MoeFlags(_FlagGroupBase):
class DpFlags(_FlagGroupBase):
"""DP-attention runtime flags, materialized by ``initialize_dp_attention``
(after distributed setup; reads the model config). Topology values
(sizes/ranks) stay on ``layers.dp_attention`` until the parallel vertical
migrates them."""
(after distributed setup; reads the model config). The topology values it
also computes -- the attention-DP width and rank -- are stamped on
``get_parallel()``, not kept here."""
enabled: bool = False
use_world_group_for_gather: bool = False