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

Add get_parallel(): a structured accessor for parallel-topology state (#28567)

Browse Source
This commit is contained in:
Cheng Wan
2026-06-17 20:23:43 -07:00
committed by GitHub
parent d27d8b24de
commit 53318911ca
184 changed files with 1871 additions and 1733 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py
+3 -2
View File
@@ -7,7 +7,6 @@ import torch.nn.functional as F
from torch import nn
from sglang.srt.configs.model_config import AttentionArch
from sglang.srt.layers import dp_attention as _dp_attention
from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, ReqToTokenPool
@@ -19,6 +18,7 @@ from sglang.srt.model_executor.cuda_graph_config import (
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.runtime_context import get_parallel
from sglang.srt.server_args import set_global_server_args_for_scheduler
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
@@ -26,7 +26,8 @@ from ..mock_server_args import make_mock_server_args
# Unit tests run without distributed initialization. Backends that size buffers by
# attention tensor-parallel degree should see the single-rank default.
_dp_attention.get_attention_tp_size = lambda: 1
_parallel_override = get_parallel().override(attn_tp_size=1)
_parallel_override.__enter__()
DEFAULT_HEAD_DIM = 16
DEFAULT_HIDDEN_SIZE = 64
python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py
+3 -3
View File
@@ -6,7 +6,6 @@ import torch
from torch import nn
from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS
from sglang.srt.layers.attention.dsa import utils as _dsa_utils
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool, ReqToTokenPool
from sglang.srt.model_executor.cuda_graph_config import (
@@ -17,6 +16,7 @@ from sglang.srt.model_executor.cuda_graph_config import (
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.runtime_context import get_parallel
from sglang.srt.server_args import set_global_server_args_for_scheduler
from ..mock_server_args import make_mock_server_args
@@ -37,8 +37,8 @@ from .dense_attention import (
# Unit tests run without distributed initialization. DSA context-parallel probes
# should see the single-rank default.
_dsa_utils.get_attention_cp_size = lambda: 1
_dsa_utils.get_attention_cp_rank = lambda: 0
_parallel_override = get_parallel().override(attn_cp_size=1, attn_cp_rank=0)
_parallel_override.__enter__()
DSA_PAGE_SIZE = 64
DSA_INDEX_HEAD_DIM = 128
python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py
+5 -5
View File
@@ -19,7 +19,6 @@ from typing import Any
import torch
from torch import nn
from sglang.srt.layers import dp_attention as _dp_attention
from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS
from sglang.srt.layers.attention.dsv4.quant_k_cache import (
quant_to_nope_fp8_rope_bf16_pack_triton,
@@ -34,15 +33,16 @@ from sglang.srt.model_executor.cuda_graph_config import (
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
from sglang.srt.runtime_context import get_parallel
from sglang.srt.server_args import set_global_server_args_for_scheduler
from ..mock_server_args import make_mock_server_args
# DSV4 backend pre-resolves attention TP at construction; pin to single-rank.
_dp_attention.get_attention_tp_size = lambda: 1
_dp_attention.get_attention_tp_rank = lambda: 0
_dp_attention.get_attention_cp_size = lambda: 1
_dp_attention.get_attention_cp_rank = lambda: 0
_parallel_override = get_parallel().override(
attn_tp_size=1, attn_tp_rank=0, attn_cp_size=1, attn_cp_rank=0
)
_parallel_override.__enter__()
# DSV4 hard-coded geometry. Do not change.
DSV4_PAGE_SIZE = 256
python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py
+3 -1
View File
@@ -18,6 +18,7 @@ from sglang.srt.model_executor.cuda_graph_config import (
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.runtime_context import get_parallel
from sglang.srt.server_args import set_global_server_args_for_scheduler
from ..mock_server_args import make_mock_server_args
@@ -78,7 +79,8 @@ DUAL_CHUNK_SPARSE_SUB_WINDOW_CONFIG = {
# Unit tests run without distributed initialization. Sparse dual-chunk config
# lookup should see the single-rank default.
_dual_chunk_backend.get_tensor_model_parallel_rank = lambda: 0
_parallel_override = get_parallel().override(tp_rank=0)
_parallel_override.__enter__()
@dataclass(frozen=True)
python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py
+3 -2
View File
@@ -10,7 +10,6 @@ from sglang.srt.configs.mamba_utils import (
Mamba2StateShape,
)
from sglang.srt.configs.model_config import AttentionArch
from sglang.srt.layers import dp_attention as _dp_attention
from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS
from sglang.srt.layers.attention.hybrid_linear_attn_backend import (
HybridLinearAttnBackend,
@@ -30,10 +29,12 @@ from sglang.srt.model_executor.cuda_graph_config import (
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.runtime_context import get_parallel
from ..mock_server_args import make_mock_server_args
_dp_attention.get_attention_tp_size = lambda: 1
_parallel_override = get_parallel().override(attn_tp_size=1)
_parallel_override.__enter__()
DEFAULT_HEAD_K_DIM = 32
DEFAULT_HEAD_V_DIM = 32
python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py
+3 -2
View File
@@ -10,7 +10,6 @@ from sglang.srt.configs.mamba_utils import (
Mamba2StateDType,
)
from sglang.srt.configs.model_config import AttentionArch
from sglang.srt.layers import dp_attention as _dp_attention
from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS
from sglang.srt.layers.attention.hybrid_linear_attn_backend import (
HybridLinearAttnBackend,
@@ -30,10 +29,12 @@ from sglang.srt.model_executor.cuda_graph_config import (
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.runtime_context import get_parallel
from ..mock_server_args import make_mock_server_args
_dp_attention.get_attention_tp_size = lambda: 1
_parallel_override = get_parallel().override(attn_tp_size=1)
_parallel_override.__enter__()
DEFAULT_HEAD_K_DIM = 32
DEFAULT_HEAD_V_DIM = 32
python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py
+3 -3
View File
@@ -10,7 +10,6 @@ from sglang.srt.configs.mamba_utils import (
Mamba2StateShape,
)
from sglang.srt.configs.model_config import AttentionArch
from sglang.srt.layers import dp_attention as _dp_attention
from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS
from sglang.srt.layers.attention.linear.lightning_backend import (
LightningAttentionBackend,
@@ -29,11 +28,12 @@ from sglang.srt.model_executor.cuda_graph_config import (
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.runtime_context import get_parallel
from ..mock_server_args import make_mock_server_args
_dp_attention.get_attention_tp_size = lambda: 1
_dp_attention.get_attention_tp_rank = lambda: 0
_parallel_override = get_parallel().override(attn_tp_size=1, attn_tp_rank=0)
_parallel_override.__enter__()
# seg_la kernel constraints (see seg_la.py:683-694):
# - decode (`seg_la_d_kernel`): K_SPLIT_DIM = 128, so head_dim must be >= 128
python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py
+5 -14
View File
@@ -6,22 +6,13 @@ import torch.nn.functional as F
from torch import nn
# Patch TP world size / rank before importing modules that read them at __init__.
import sglang.srt.distributed as _distributed
import sglang.srt.layers.attention.mamba.mamba as _mamba_mod
import sglang.srt.layers.attention.mamba.mixer2_rms_norm_gated as _norm_mod
import sglang.srt.layers.linear as _linear_mod
from sglang.srt.layers import dp_attention as _dp_attention
from sglang.srt.runtime_context import get_parallel
_distributed.get_tensor_model_parallel_world_size = lambda: 1
_distributed.get_tensor_model_parallel_rank = lambda: 0
_mamba_mod.get_tensor_model_parallel_world_size = lambda: 1
_mamba_mod.get_tensor_model_parallel_rank = lambda: 0
_norm_mod.get_tensor_model_parallel_world_size = lambda: 1
_norm_mod.get_tensor_model_parallel_rank = lambda: 0
_linear_mod.get_tensor_model_parallel_world_size = lambda: 1
_linear_mod.get_tensor_model_parallel_rank = lambda: 0
_dp_attention.get_attention_tp_size = lambda: 1
_dp_attention.get_attention_tp_rank = lambda: 0
_parallel_override = get_parallel().override(
tp_size=1, tp_rank=0, attn_tp_size=1, attn_tp_rank=0
)
_parallel_override.__enter__()
# RowParallelLinear.forward calls get_tp_group() to manage symmetric memory.
# Provide a stub group with world_size=1 so use_symmetric_memory short-circuits.
python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py
+3 -2
View File
@@ -7,7 +7,6 @@ import torch.nn.functional as F
from torch import nn
from sglang.srt.configs.model_config import AttentionArch
from sglang.srt.layers import dp_attention as _dp_attention
from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool, ReqToTokenPool
@@ -23,11 +22,13 @@ from sglang.srt.model_executor.forward_context import (
get_token_to_kv_pool,
)
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.runtime_context import get_parallel
from sglang.srt.server_args import set_global_server_args_for_scheduler
from ..mock_server_args import make_mock_server_args
_dp_attention.get_attention_tp_size = lambda: 1
_parallel_override = get_parallel().override(attn_tp_size=1)
_parallel_override.__enter__()
DEFAULT_HIDDEN_SIZE = 64
DEFAULT_KV_LORA_RANK = 32
python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py
+2 -4
View File
@@ -14,6 +14,7 @@ from sglang.srt.model_executor.forward_batch_info import (
ForwardBatch,
)
from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
from sglang.srt.runtime_context import get_parallel
from sglang.srt.speculative.draft_utils import DraftBackendFactory
from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import (
EAGLEDraftExtendCudaGraphRunner,
@@ -560,10 +561,7 @@ def _capture_eagle_draft_extend_graph_runner(
"sglang.srt.model_executor.runner.decode_cuda_graph_runner.get_available_gpu_memory",
lambda *args, **kwargs: 0.0,
),
patch(
"sglang.srt.model_executor.runner.base_cuda_graph_runner.get_attention_cp_size",
lambda: 1,
),
get_parallel().override(attn_cp_size=1),
):
_reset_cuda_graph_test_buffers()
return EAGLEDraftExtendCudaGraphRunner(
python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py
+3 -8
View File
@@ -16,6 +16,7 @@ from sglang.srt.model_executor.forward_batch_info import (
)
from sglang.srt.model_executor.input_buffers import _forward_input_buffer_pool
from sglang.srt.model_executor.runner import set_global_graph_memory_pool
from sglang.srt.runtime_context import get_parallel
from sglang.srt.server_args import set_global_server_args_for_scheduler
from sglang.srt.speculative.draft_utils import DraftBackendFactory
from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
@@ -450,10 +451,7 @@ def _capture_eagle_draft_graph_runner(
"sglang.srt.model_executor.runner.decode_cuda_graph_runner.get_available_gpu_memory",
lambda *args, **kwargs: 0.0,
),
patch(
"sglang.srt.model_executor.runner.base_cuda_graph_runner.get_attention_cp_size",
lambda: 1,
),
get_parallel().override(attn_cp_size=1),
):
_reset_cuda_graph_test_buffers()
return EAGLEDraftCudaGraphRunner(
@@ -479,10 +477,7 @@ def _capture_frozen_kv_mtp_graph_runner(
"sglang.srt.model_executor.runner.decode_cuda_graph_runner.get_available_gpu_memory",
lambda *args, **kwargs: 0.0,
),
patch(
"sglang.srt.model_executor.runner.base_cuda_graph_runner.get_attention_cp_size",
lambda: 1,
),
get_parallel().override(attn_cp_size=1),
):
_reset_cuda_graph_test_buffers()
return FrozenKVMTPCudaGraphRunner(worker)