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

[lint] Enable Ruff UP037 to drop redundant quoted annotations (#27984)

Browse Source
This commit is contained in:
Cheng Wan
2026-06-11 17:38:05 -07:00
committed by GitHub
parent c0480a88be
commit 97a0031799
185 changed files with 541 additions and 551 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.pre-commit-config.yaml
+1 -1
View File
@@ -35,7 +35,7 @@ repos:
hooks: hooks:
- id: ruff - id: ruff
args: args:
- --select=F401,F821 - --select=F401,F821,UP037
- --fix - --fix
files: ^(benchmark/|docs/|examples/|python/sglang/|sgl-model-gateway/py_*|test/) files: ^(benchmark/|docs/|examples/|python/sglang/|sgl-model-gateway/py_*|test/)
exclude: | exclude: |
python/sglang/_mps_stub.py
+2 -2
View File
@@ -40,7 +40,7 @@ class Stream:
return True return True
# context-manager protocol (``with stream:``) # context-manager protocol (``with stream:``)
def __enter__(self) -> "Stream": def __enter__(self) -> Stream:
return self return self
def __exit__(self, *args: Any) -> None: def __exit__(self, *args: Any) -> None:
@@ -75,7 +75,7 @@ class StreamContext:
def __init__(self, stream: Any = None) -> None: def __init__(self, stream: Any = None) -> None:
pass pass
def __enter__(self) -> "StreamContext": def __enter__(self) -> StreamContext:
return self return self
def __exit__(self, *args: Any) -> None: def __exit__(self, *args: Any) -> None:
python/sglang/jit_kernel/kv_canary/plan/entries_kernel.py
+1 -1
View File
@@ -13,7 +13,7 @@ if TYPE_CHECKING:
@cache_once @cache_once
def _jit_plan_entries_module( def _jit_plan_entries_module(
has_swa_lut: bool, has_verify_expected_token_pool: bool has_swa_lut: bool, has_verify_expected_token_pool: bool
) -> "Module": ) -> Module:
args = make_cpp_args(has_swa_lut, has_verify_expected_token_pool) args = make_cpp_args(has_swa_lut, has_verify_expected_token_pool)
return load_jit( return load_jit(
"kv_canary_plan_entries", "kv_canary_plan_entries",
python/sglang/jit_kernel/kv_canary/verify.py
+3 -3
View File
@@ -198,7 +198,7 @@ class VerifyPlan:
enable: torch.Tensor enable: torch.Tensor
@classmethod @classmethod
def allocate(cls, *, verify_capacity: int, device: torch.device) -> "VerifyPlan": def allocate(cls, *, verify_capacity: int, device: torch.device) -> VerifyPlan:
if verify_capacity <= 0: if verify_capacity <= 0:
raise ValueError( raise ValueError(
f"kv-canary: VerifyPlan verify_capacity must be positive, got {verify_capacity}" f"kv-canary: VerifyPlan verify_capacity must be positive, got {verify_capacity}"
@@ -223,7 +223,7 @@ class VerifyPlan:
enable=torch.ones(1, dtype=torch.int32, device=device), enable=torch.ones(1, dtype=torch.int32, device=device),
) )
def zero_for_testing_(self) -> "VerifyPlan": def zero_for_testing_(self) -> VerifyPlan:
"""WARN: ONLY use it when testing plan kernel. Do not use it when testing verify or """WARN: ONLY use it when testing plan kernel. Do not use it when testing verify or
write kernel to avoid hiding bugs.""" write kernel to avoid hiding bugs."""
self.verify_slot_indices.zero_() self.verify_slot_indices.zero_()
@@ -352,7 +352,7 @@ def launch_canary_verify_kernel(
@cache_once @cache_once
def _jit_canary_verify_module(check_verify_expected_token: bool) -> "Module": def _jit_canary_verify_module(check_verify_expected_token: bool) -> Module:
args = make_cpp_args(check_verify_expected_token) args = make_cpp_args(check_verify_expected_token)
return load_jit( return load_jit(
"kv_canary_verify", "kv_canary_verify",
python/sglang/jit_kernel/kv_canary/write.py
+3 -3
View File
@@ -52,7 +52,7 @@ class WritePlan:
*, *,
write_req_capacity: int, write_req_capacity: int,
device: torch.device, device: torch.device,
) -> "WritePlan": ) -> WritePlan:
if write_req_capacity <= 0: if write_req_capacity <= 0:
raise ValueError( raise ValueError(
f"kv-canary: WritePlan write_req_capacity must be positive, got {write_req_capacity}" f"kv-canary: WritePlan write_req_capacity must be positive, got {write_req_capacity}"
@@ -67,7 +67,7 @@ class WritePlan:
write_num_valid_reqs=torch.empty(1, dtype=torch.int32, device=device), write_num_valid_reqs=torch.empty(1, dtype=torch.int32, device=device),
) )
def zero_for_testing_(self) -> "WritePlan": def zero_for_testing_(self) -> WritePlan:
"""WARN: ONLY use it when testing plan kernel. Do not use it when testing verify or """WARN: ONLY use it when testing plan kernel. Do not use it when testing verify or
write kernel to avoid hiding bugs.""" write kernel to avoid hiding bugs."""
self.write_offsets.zero_() self.write_offsets.zero_()
@@ -251,7 +251,7 @@ def launch_canary_write_kernel(
@cache_once @cache_once
def _jit_canary_write_module() -> "Module": def _jit_canary_write_module() -> Module:
return load_jit( return load_jit(
"kv_canary_write", "kv_canary_write",
cuda_files=["kv_canary/canary_write.cuh"], cuda_files=["kv_canary/canary_write.cuh"],
python/sglang/jit_kernel/tests/kv_canary/_canary_helpers.py
+1 -1
View File
@@ -55,7 +55,7 @@ class FakeViolationLog:
@classmethod @classmethod
def allocate( def allocate(
cls, *, capacity: int = DEFAULT_RING_CAPACITY, device: torch.device cls, *, capacity: int = DEFAULT_RING_CAPACITY, device: torch.device
) -> "FakeViolationLog": ) -> FakeViolationLog:
return cls( return cls(
ring=torch.zeros( ring=torch.zeros(
capacity, consts.VIOLATION_FIELDS, dtype=torch.int64, device=device capacity, consts.VIOLATION_FIELDS, dtype=torch.int64, device=device
python/sglang/multimodal_gen/configs/quantization/nunchaku.py
+2 -2
View File
@@ -69,7 +69,7 @@ class NunchakuSVDQuantArgs:
return enable_svdquant, inferred_precision, inferred_rank return enable_svdquant, inferred_precision, inferred_rank
def _normalized(self) -> "NunchakuSVDQuantArgs": def _normalized(self) -> NunchakuSVDQuantArgs:
enable_svdquant, inferred_precision, inferred_rank = ( enable_svdquant, inferred_precision, inferred_rank = (
self._infer_from_weights_path() self._infer_from_weights_path()
) )
@@ -204,7 +204,7 @@ class NunchakuSVDQuantArgs:
) )
@classmethod @classmethod
def from_dict(cls, kwargs: dict[str, Any]) -> "NunchakuSVDQuantArgs": def from_dict(cls, kwargs: dict[str, Any]) -> NunchakuSVDQuantArgs:
# Map CLI/config keys to dataclass fields (keep backwards compatibility). # Map CLI/config keys to dataclass fields (keep backwards compatibility).
path = ( path = (
kwargs.get("transformer_weights_path") kwargs.get("transformer_weights_path")
python/sglang/multimodal_gen/runtime/distributed/cfg_parallel_utils.py
+7 -7
View File
@@ -36,7 +36,7 @@ _logged_dispatch_keys: set[tuple[int, int, int]] = set()
def _run( def _run(
predict_fn: Callable[["CFGBranch"], "torch.Tensor | tuple[torch.Tensor, ...]"], predict_fn: Callable[[CFGBranch], torch.Tensor | tuple[torch.Tensor, ...]],
bid: int, bid: int,
branches, branches,
) -> tuple[torch.Tensor, ...]: ) -> tuple[torch.Tensor, ...]:
@@ -54,9 +54,9 @@ def _run(
def run_cfg_parallel( def run_cfg_parallel(
policy: "CFGPolicy", policy: CFGPolicy,
predict_fn: Callable[["CFGBranch"], "torch.Tensor | tuple[torch.Tensor, ...]"], predict_fn: Callable[[CFGBranch], torch.Tensor | tuple[torch.Tensor, ...]],
) -> "list[torch.Tensor | tuple[torch.Tensor, ...]]": ) -> list[torch.Tensor | tuple[torch.Tensor, ...]]:
"""Dispatch CFG branches across ranks, all-gather results, return in branch order. """Dispatch CFG branches across ranks, all-gather results, return in branch order.
``predict_fn`` is a closure capturing all step-varying state ``predict_fn`` is a closure capturing all step-varying state
@@ -136,12 +136,12 @@ def run_cfg_parallel(
def run_two_branch_cfg_parallel( def run_two_branch_cfg_parallel(
policy: "CFGPolicy", policy: CFGPolicy,
predict_fn: Callable[["CFGBranch"], "torch.Tensor | tuple[torch.Tensor, ...]"], predict_fn: Callable[[CFGBranch], torch.Tensor | tuple[torch.Tensor, ...]],
cfg_scale: float, cfg_scale: float,
batch, batch,
pipeline_config, pipeline_config,
) -> "torch.Tensor | tuple[torch.Tensor, ...]": ) -> torch.Tensor | tuple[torch.Tensor, ...]:
"""Run standard two-pass CFG with the old all-reduce combine. """Run standard two-pass CFG with the old all-reduce combine.
This keeps the existing WAN baselines: it avoids gathering both branch This keeps the existing WAN baselines: it avoids gathering both branch
python/sglang/multimodal_gen/runtime/distributed/cfg_policy.py
+5 -5
View File
@@ -21,7 +21,7 @@ class CFGBranch:
is_conditional: bool is_conditional: bool
kwargs: dict[str, Any] kwargs: dict[str, Any]
def configure_batch(self, batch: "Req") -> None: def configure_batch(self, batch: Req) -> None:
"""Set batch state before this branch's forward pass. """Set batch state before this branch's forward pass.
Override for richer per-branch context (e.g. a branch index instead of Override for richer per-branch context (e.g. a branch index instead of
@@ -46,11 +46,11 @@ class CFGPolicy:
def build( def build(
self, self,
batch: "Req", batch: Req,
image_kwargs: dict[str, Any], image_kwargs: dict[str, Any],
pos_cond_kwargs: dict[str, Any], pos_cond_kwargs: dict[str, Any],
neg_cond_kwargs: dict[str, Any], neg_cond_kwargs: dict[str, Any],
) -> "CFGPolicy": ) -> CFGPolicy:
"""Return a new policy with branches populated. """Return a new policy with branches populated.
Called once before the denoising loop. The returned policy is Called once before the denoising loop. The returned policy is
@@ -66,7 +66,7 @@ class CFGPolicy:
def combine( def combine(
self, self,
predictions: list[torch.Tensor | tuple[torch.Tensor, ...]], predictions: list[torch.Tensor | tuple[torch.Tensor, ...]],
batch: "Req", batch: Req,
cfg_scale: float, cfg_scale: float,
pipeline_config: Any, pipeline_config: Any,
*, *,
@@ -117,7 +117,7 @@ def _unwrap(
def _apply_cfg_postprocess( def _apply_cfg_postprocess(
noise_pred: torch.Tensor, noise_pred: torch.Tensor,
noise_pred_cond: torch.Tensor, noise_pred_cond: torch.Tensor,
batch: "Req", batch: Req,
pipeline_config: Any, pipeline_config: Any,
) -> torch.Tensor: ) -> torch.Tensor:
if batch.cfg_normalization and float(batch.cfg_normalization) > 0: if batch.cfg_normalization and float(batch.cfg_normalization) > 0:
python/sglang/multimodal_gen/runtime/layers/attention/layer.py
+1 -1
View File
@@ -435,7 +435,7 @@ class USPAttention(nn.Module):
f"but got {backend_enum.name}. " f"but got {backend_enum.name}. "
f"Please ensure your platform supports these backends." f"Please ensure your platform supports these backends."
) )
impl_cls: Type["AttentionImpl"] = attn_backend.get_impl_cls() impl_cls: Type[AttentionImpl] = attn_backend.get_impl_cls()
self.allow_cudnn_sdp = bool(extra_impl_args.get("allow_cudnn_sdp", False)) self.allow_cudnn_sdp = bool(extra_impl_args.get("allow_cudnn_sdp", False))
self.attn_impl = impl_cls( self.attn_impl = impl_cls(
num_heads=num_heads, num_heads=num_heads,
python/sglang/multimodal_gen/runtime/layers/attention/turbo_layer.py
+1 -1
View File
@@ -252,7 +252,7 @@ class MinimalA2AAttnOp(DistributedAttention):
attn_backend = SageSparseLinearAttentionBackend attn_backend = SageSparseLinearAttentionBackend
else: else:
attn_backend = SparseLinearAttentionBackend attn_backend = SparseLinearAttentionBackend
impl_cls: Type["AttentionImpl"] = attn_backend.get_impl_cls() impl_cls: Type[AttentionImpl] = attn_backend.get_impl_cls()
local_attn = impl_cls( local_attn = impl_cls(
num_heads=num_heads, num_heads=num_heads,
head_size=head_size, head_size=head_size,
python/sglang/multimodal_gen/runtime/layers/quantization/bitsandbytes.py
+1 -1
View File
@@ -105,7 +105,7 @@ class BitsAndBytesConfig(QuantizationConfig):
return [] return []
@classmethod @classmethod
def from_config(cls, config: dict[str, Any]) -> "BitsAndBytesConfig": def from_config(cls, config: dict[str, Any]) -> BitsAndBytesConfig:
def get_safe_value(keys, default_value=None): def get_safe_value(keys, default_value=None):
try: try:
value = QuantizationConfig.get_from_keys(config, keys) value = QuantizationConfig.get_from_keys(config, keys)
python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_fp8.py
+1 -1
View File
@@ -74,7 +74,7 @@ class ModelOptFp8Config(QuantizationConfig):
return [] return []
@classmethod @classmethod
def from_config(cls, config: Dict[str, Any]) -> "ModelOptFp8Config": def from_config(cls, config: Dict[str, Any]) -> ModelOptFp8Config:
quant_algo = config.get("quant_algo") quant_algo = config.get("quant_algo")
if quant_algo is None: if quant_algo is None:
raise ValueError( raise ValueError(
python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py
+1 -1
View File
@@ -187,7 +187,7 @@ class ModelOptFp8Config(ModelOptQuantConfig):
return 89 return 89
@classmethod @classmethod
def from_config(cls, config: Dict[str, Any]) -> "ModelOptFp8Config": def from_config(cls, config: Dict[str, Any]) -> ModelOptFp8Config:
quant_method = config.get("quant_algo") quant_method = config.get("quant_algo")
exclude_modules = config.get("ignore") exclude_modules = config.get("ignore")
if quant_method is None: if quant_method is None:
python/sglang/multimodal_gen/runtime/layers/quantization/mxfp4_npu.py
+1 -1
View File
@@ -63,7 +63,7 @@ class NPUMXFP4Config(QuantizationConfig):
return [] return []
@classmethod @classmethod
def from_config(cls, config: Dict[str, Any]) -> "NPUMXFP4Config": def from_config(cls, config: Dict[str, Any]) -> NPUMXFP4Config:
return cls() return cls()
def get_quant_method( def get_quant_method(
python/sglang/multimodal_gen/runtime/layers/quantization/mxfp8_npu.py
+1 -1
View File
@@ -56,7 +56,7 @@ class MXFP8Config(QuantizationConfig):
return [] return []
@classmethod @classmethod
def from_config(cls, config: Dict[str, Any]) -> "MXFP8Config": def from_config(cls, config: Dict[str, Any]) -> MXFP8Config:
return cls() return cls()
def get_quant_method( def get_quant_method(
python/sglang/multimodal_gen/runtime/managers/dynamic_batch_admission.py
+2 -2
View File
@@ -78,7 +78,7 @@ class BatchingRule:
source: str = "user" source: str = "user"
@classmethod @classmethod
def from_dict(cls, data: dict[str, Any], *, source: str) -> "BatchingRule": def from_dict(cls, data: dict[str, Any], *, source: str) -> BatchingRule:
if not isinstance(data, dict): if not isinstance(data, dict):
raise ValueError( raise ValueError(
f"batching config rule from {source} must be an object, " f"batching config rule from {source} must be an object, "
@@ -156,7 +156,7 @@ class BatchingRule:
class BatchAdmissionController: class BatchAdmissionController:
"""Applies configured caps before adding requests to a batch.""" """Applies configured caps before adding requests to a batch."""
def __init__(self, server_args: "ServerArgs", gpu_id: int): def __init__(self, server_args: ServerArgs, gpu_id: int):
self._mode = getattr(server_args, "batching_mode", "dynamic") self._mode = getattr(server_args, "batching_mode", "dynamic")
self._user_max_batch_size = max(1, int(server_args.batching_max_size)) self._user_max_batch_size = max(1, int(server_args.batching_max_size))
self._model_path = server_args.model_path self._model_path = server_args.model_path
python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py
+4 -4
View File
@@ -197,7 +197,7 @@ class Req:
VSA_sparsity: float = 0.0 VSA_sparsity: float = 0.0
# stage logging # stage logging
metrics: Optional["RequestMetrics"] = None metrics: Optional[RequestMetrics] = None
# tracing context (TraceReqContext or TraceNullContext) # tracing context (TraceReqContext or TraceNullContext)
trace_ctx: Union[TraceReqContext, TraceNullContext] = field( trace_ctx: Union[TraceReqContext, TraceNullContext] = field(
@@ -326,7 +326,7 @@ class Req:
self.extra["cache_dit_num_inference_steps"] = self.num_inference_steps self.extra["cache_dit_num_inference_steps"] = self.num_inference_steps
self.num_inference_steps = warmup_steps self.num_inference_steps = warmup_steps
def copy_as_warmup(self, warmup_steps: int = 1) -> "Req": def copy_as_warmup(self, warmup_steps: int = 1) -> Req:
req = deepcopy(self) req = deepcopy(self)
req.set_as_warmup(warmup_steps) req.set_as_warmup(warmup_steps)
return req return req
@@ -421,8 +421,8 @@ class OutputBatch:
output_file_paths: list[str] | None = None output_file_paths: list[str] | None = None
# logged metrics info, directly from Req.timings # logged metrics info, directly from Req.timings
metrics: Optional["RequestMetrics"] = None metrics: Optional[RequestMetrics] = None
metrics_list: Optional[list[Optional["RequestMetrics"]]] = None metrics_list: Optional[list[Optional[RequestMetrics]]] = None
# For ComfyUI integration: noise prediction from denoising stage # For ComfyUI integration: noise prediction from denoising stage
noise_pred: torch.Tensor | None = None noise_pred: torch.Tensor | None = None
python/sglang/multimodal_gen/runtime/pipelines_core/stages/dedup.py
+10 -10
View File
@@ -34,8 +34,8 @@ class StageDedupMixin:
def run_grouped_requests( def run_grouped_requests(
self, self,
batches: list["Req"], batches: list[Req],
server_args: "ServerArgs", server_args: ServerArgs,
) -> list[Any]: ) -> list[Any]:
"""Run this stage for a group of independent requests. """Run this stage for a group of independent requests.
@@ -63,7 +63,7 @@ class StageDedupMixin:
or cls.deduplicated_extra_tensor_tree_output_keys or cls.deduplicated_extra_tensor_tree_output_keys
) )
def build_dedup_fingerprint(self, batch: "Req", server_args: "ServerArgs") -> Any: def build_dedup_fingerprint(self, batch: Req, server_args: ServerArgs) -> Any:
"""Return this stage's semantic input fingerprint for grouped dedup. """Return this stage's semantic input fingerprint for grouped dedup.
A fingerprint is the stage-local set of input values that fully A fingerprint is the stage-local set of input values that fully
@@ -79,10 +79,10 @@ class StageDedupMixin:
def run_deduplicated_group( def run_deduplicated_group(
self, self,
batches: list["Req"], batches: list[Req],
server_args: "ServerArgs", server_args: ServerArgs,
copy_outputs=None, copy_outputs=None,
) -> list["Req"]: ) -> list[Req]:
"""Run full-stage-equivalent requests once and fan out stage outputs.""" """Run full-stage-equivalent requests once and fan out stage outputs."""
if copy_outputs is None: if copy_outputs is None:
copy_outputs = self.copy_deduplicated_outputs copy_outputs = self.copy_deduplicated_outputs
@@ -102,7 +102,7 @@ class StageDedupMixin:
return [result for result in results if result is not None] return [result for result in results if result is not None]
def copy_deduplicated_outputs(self, src: "Req", dst: "Req") -> None: def copy_deduplicated_outputs(self, src: Req, dst: Req) -> None:
"""Copy declared stage outputs from a computed request to a duplicate. """Copy declared stage outputs from a computed request to a duplicate.
``deduplicated_output_fields`` uses shallow container copies and shares ``deduplicated_output_fields`` uses shallow container copies and shares
@@ -175,11 +175,11 @@ class StageDedupMixin:
@staticmethod @staticmethod
def _group_requests_by_fingerprint( def _group_requests_by_fingerprint(
batches: list["Req"], batches: list[Req],
fingerprint_fn, fingerprint_fn,
) -> list[tuple[Any, list[tuple[int, "Req"]]]]: ) -> list[tuple[Any, list[tuple[int, Req]]]]:
"""Group requests by a stage-local fingerprint while preserving order.""" """Group requests by a stage-local fingerprint while preserving order."""
groups: dict[Any, list[tuple[int, "Req"]]] = {} groups: dict[Any, list[tuple[int, Req]]] = {}
for index, batch in enumerate(batches): for index, batch in enumerate(batches):
fingerprint = fingerprint_fn(batch) fingerprint = fingerprint_fn(batch)
groups.setdefault(fingerprint, []).append((index, batch)) groups.setdefault(fingerprint, []).append((index, batch))
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/paint.py
+1 -1
View File
@@ -164,7 +164,7 @@ class DDIMSolver:
self.ddim_alpha_cumprods = torch.from_numpy(self.ddim_alpha_cumprods) self.ddim_alpha_cumprods = torch.from_numpy(self.ddim_alpha_cumprods)
self.ddim_alpha_cumprods_prev = torch.from_numpy(self.ddim_alpha_cumprods_prev) self.ddim_alpha_cumprods_prev = torch.from_numpy(self.ddim_alpha_cumprods_prev)
def to(self, device: torch.device) -> "DDIMSolver": def to(self, device: torch.device) -> DDIMSolver:
self.ddim_timesteps = self.ddim_timesteps.to(device) self.ddim_timesteps = self.ddim_timesteps.to(device)
self.ddim_alpha_cumprods = self.ddim_alpha_cumprods.to(device) self.ddim_alpha_cumprods = self.ddim_alpha_cumprods.to(device)
self.ddim_alpha_cumprods_prev = self.ddim_alpha_cumprods_prev.to(device) self.ddim_alpha_cumprods_prev = self.ddim_alpha_cumprods_prev.to(device)
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/self_forcing.py
+1 -1
View File
@@ -53,7 +53,7 @@ class SanaWMSelfForcingSamplerConfig:
streaming_cfg_scale: float = 1.0 streaming_cfg_scale: float = 1.0
@classmethod @classmethod
def from_pipeline_config(cls, pcfg) -> "SanaWMSelfForcingSamplerConfig": def from_pipeline_config(cls, pcfg) -> SanaWMSelfForcingSamplerConfig:
"""Read the streaming knobs off a pipeline config (note the ``or 1.0`` cfg-scale guard).""" """Read the streaming knobs off a pipeline config (note the ``or 1.0`` cfg-scale guard)."""
return cls( return cls(
num_frame_per_block=int(getattr(pcfg, "num_frame_per_block", 3)), num_frame_per_block=int(getattr(pcfg, "num_frame_per_block", 3)),
python/sglang/multimodal_gen/runtime/server_args_auto_tune.py
+1 -1
View File
@@ -36,7 +36,7 @@ DEFAULT_LAYERWISE_COMPONENT_ARG_NAMES = (
class ServerArgsAutoTuner: class ServerArgsAutoTuner:
"""Auto-tunes the server-arg for the given performance-mode, based on practical deployment experience with different model architectures""" """Auto-tunes the server-arg for the given performance-mode, based on practical deployment experience with different model architectures"""
def __init__(self, server_args: "ServerArgs"): def __init__(self, server_args: ServerArgs):
self.server_args = server_args self.server_args = server_args
self._explicit_memory_policy = self._has_explicit_memory_policy() self._explicit_memory_policy = self._has_explicit_memory_policy()
self._explicit_layerwise_replacement_policy = ( self._explicit_layerwise_replacement_policy = (
python/sglang/multimodal_gen/test/server/test_disagg_server.py
+1 -1
View File
@@ -127,7 +127,7 @@ class DisaggCluster:
# -- context manager ----------------------------------------------------- # -- context manager -----------------------------------------------------
def __enter__(self) -> "DisaggCluster": def __enter__(self) -> DisaggCluster:
for attempt in range(3): for attempt in range(3):
try: try:
self._launch_roles() self._launch_roles()
python/sglang/multimodal_gen/test/server/test_server_common.py
+2 -2
View File
@@ -63,7 +63,7 @@ logger = init_logger(__name__)
# Track test cases missing estimated_full_test_time_s for time measurement output # Track test cases missing estimated_full_test_time_s for time measurement output
_MISSING_ESTIMATED_TIME_CASES: set[str] = set() _MISSING_ESTIMATED_TIME_CASES: set[str] = set()
_PENDING_BASELINE_DUMPS: dict[str, tuple["PerformanceSummary", bool]] = {} _PENDING_BASELINE_DUMPS: dict[str, tuple[PerformanceSummary, bool]] = {}
_OPENAI_REQUEST_TIMEOUT_SECS = float( _OPENAI_REQUEST_TIMEOUT_SECS = float(
os.environ.get("SGLANG_TEST_OPENAI_REQUEST_TIMEOUT_SECS", "600") os.environ.get("SGLANG_TEST_OPENAI_REQUEST_TIMEOUT_SECS", "600")
) )
@@ -518,7 +518,7 @@ class DiffusionServerBase:
def _dump_baseline_for_testcase( def _dump_baseline_for_testcase(
self, self,
case: DiffusionTestCase, case: DiffusionTestCase,
summary: "PerformanceSummary", summary: PerformanceSummary,
missing_scenario: bool = False, missing_scenario: bool = False,
measured_full_time: float | None = None, measured_full_time: float | None = None,
) -> None: ) -> None:
python/sglang/srt/batch_overlap/operations.py
+1 -1
View File
@@ -117,7 +117,7 @@ class _StageExecutor:
debug_name: str, debug_name: str,
stages: List[Stage], stages: List[Stage],
inputs: dict, inputs: dict,
child_ctx: Optional["ForwardContext"] = None, child_ctx: Optional[ForwardContext] = None,
): ):
self._debug_name = debug_name self._debug_name = debug_name
self._stages = stages self._stages = stages
python/sglang/srt/batch_overlap/two_batch_overlap.py
+1 -1
View File
@@ -266,7 +266,7 @@ def split_spec_info(
def compute_split_token_index( def compute_split_token_index(
split_seq_index: int, split_seq_index: int,
forward_mode: "ForwardMode", forward_mode: ForwardMode,
extend_seq_lens: Optional[Sequence[int]], extend_seq_lens: Optional[Sequence[int]],
token_num_per_seq: Optional[int], token_num_per_seq: Optional[int],
) -> int: ) -> int:
python/sglang/srt/compilation/compile_phase.py
+2 -2
View File
@@ -20,7 +20,7 @@ from contextlib import contextmanager
import torch import torch
_in_torch_compile_warmup = False _in_torch_compile_warmup = False
_pcg_capture_stream: "torch.cuda.Stream | None" = None _pcg_capture_stream: torch.cuda.Stream | None = None
def is_in_torch_compile_warmup() -> bool: def is_in_torch_compile_warmup() -> bool:
@@ -43,7 +43,7 @@ def enable_torch_compile_warmup():
_in_torch_compile_warmup = False _in_torch_compile_warmup = False
def get_pcg_capture_stream() -> "torch.cuda.Stream | None": def get_pcg_capture_stream() -> torch.cuda.Stream | None:
return _pcg_capture_stream return _pcg_capture_stream
python/sglang/srt/debug_utils/comparator/meta_overrider.py
+1 -1
View File
@@ -50,7 +50,7 @@ class MetaOverrider:
override_baseline_dims: list[str], override_baseline_dims: list[str],
override_target_dims: list[str], override_target_dims: list[str],
override_config: Optional[Path], override_config: Optional[Path],
) -> "MetaOverrider": ) -> MetaOverrider:
per_side_args: list[tuple[list[str], Literal["both", "baseline", "target"]]] = [ per_side_args: list[tuple[list[str], Literal["both", "baseline", "target"]]] = [
(override_dims, "both"), (override_dims, "both"),
(override_baseline_dims, "baseline"), (override_baseline_dims, "baseline"),
python/sglang/srt/debug_utils/comparator/output_types.py
+1 -1
View File
@@ -277,7 +277,7 @@ class SummaryRecord(_OutputRecord):
errored: int = 0 errored: int = 0
@model_validator(mode="after") @model_validator(mode="after")
def _validate_totals(self) -> "SummaryRecord": def _validate_totals(self) -> SummaryRecord:
expected: int = self.passed + self.failed + self.skipped + self.errored expected: int = self.passed + self.failed + self.skipped + self.errored
if self.total != expected: if self.total != expected:
raise ValueError( raise ValueError(
python/sglang/srt/debug_utils/dumper.py
+1 -1
View File
@@ -250,7 +250,7 @@ class _Dumper:
def __init__(self, *, config: DumperConfig): def __init__(self, *, config: DumperConfig):
self._config = config self._config = config
self._state = _DumperState() self._state = _DumperState()
self._non_intrusives: list["_NonIntrusiveDumper"] = [] self._non_intrusives: list[_NonIntrusiveDumper] = []
self._grafter = _Grafter(config=config) self._grafter = _Grafter(config=config)
# ------------------------------- public :: core --------------------------------- # ------------------------------- public :: core ---------------------------------
python/sglang/srt/disaggregation/common/staging_handler.py
+8 -8
View File
@@ -97,7 +97,7 @@ class DecodeStagingHandler:
return 1 return 1
@classmethod @classmethod
def create(cls, kv_manager, scheduler, tp_rank: int) -> "DecodeStagingHandler": def create(cls, kv_manager, scheduler, tp_rank: int) -> DecodeStagingHandler:
"""Factory: create handler. Raises if staging infra is missing.""" """Factory: create handler. Raises if staging infra is missing."""
staging_allocator = kv_manager._staging_ctx.allocator staging_allocator = kv_manager._staging_ctx.allocator
if staging_allocator is None: if staging_allocator is None:
@@ -132,7 +132,7 @@ class DecodeStagingHandler:
# Registration: called from main thread (DecodeTransferQueue) # Registration: called from main thread (DecodeTransferQueue)
# ------------------------------------------------------------------ # ------------------------------------------------------------------
def register_decode_req(self, room: int, decode_req: "DecodeRequest") -> None: def register_decode_req(self, room: int, decode_req: DecodeRequest) -> None:
self._room_to_decode_req[room] = decode_req self._room_to_decode_req[room] = decode_req
def unregister_decode_req(self, room: int) -> None: def unregister_decode_req(self, room: int) -> None:
@@ -251,13 +251,13 @@ class DecodeStagingHandler:
# Event check + free: called from main thread (pop_transferred) # Event check + free: called from main thread (pop_transferred)
# ------------------------------------------------------------------ # ------------------------------------------------------------------
def is_done(self, decode_req: "DecodeRequest") -> bool: def is_done(self, decode_req: DecodeRequest) -> bool:
"""Return True if staging scatter is complete for this request.""" """Return True if staging scatter is complete for this request."""
if not getattr(decode_req, "_staging_scatter_done", False): if not getattr(decode_req, "_staging_scatter_done", False):
return False return False
return not getattr(decode_req, "_chunk_events", None) return not getattr(decode_req, "_chunk_events", None)
def advance_scatter(self, decode_req: "DecodeRequest") -> None: def advance_scatter(self, decode_req: DecodeRequest) -> None:
"""Check CUDA events and free completed staging allocations. """Check CUDA events and free completed staging allocations.
Scatter kernels have already been submitted by the decode_thread Scatter kernels have already been submitted by the decode_thread
@@ -292,7 +292,7 @@ class DecodeStagingHandler:
staging_offset: int, staging_offset: int,
page_start: int, page_start: int,
num_pages: int, num_pages: int,
decode_req: "DecodeRequest", decode_req: DecodeRequest,
) -> bool: ) -> bool:
"""Submit scatter kernels for a staging region to scatter_stream. """Submit scatter kernels for a staging region to scatter_stream.
@@ -347,7 +347,7 @@ class DecodeStagingHandler:
return True return True
def _submit_last_scatter(self, decode_req: "DecodeRequest") -> int: def _submit_last_scatter(self, decode_req: DecodeRequest) -> int:
"""Submit scatter for the last chunk. Returns alloc_id >= 0, or -1.""" """Submit scatter for the last chunk. Returns alloc_id >= 0, or -1."""
receiver = decode_req.kv_receiver receiver = decode_req.kv_receiver
chunk_infos = getattr(receiver, "chunk_staging_infos", []) chunk_infos = getattr(receiver, "chunk_staging_infos", [])
@@ -370,7 +370,7 @@ class DecodeStagingHandler:
return alloc_id if ok else -1 return alloc_id if ok else -1
def _free_and_send_watermark( def _free_and_send_watermark(
self, alloc_id: int, decode_req: "DecodeRequest" self, alloc_id: int, decode_req: DecodeRequest
) -> None: ) -> None:
"""Free a staging allocation and broadcast watermark to all prefills.""" """Free a staging allocation and broadcast watermark to all prefills."""
self.staging_allocator.free(alloc_id) self.staging_allocator.free(alloc_id)
@@ -474,7 +474,7 @@ class StagingRegisterInfo:
@classmethod @classmethod
def from_zmq_fields( def from_zmq_fields(
cls, msg: list, msg_start_offset: int cls, msg: list, msg_start_offset: int
) -> Optional["StagingRegisterInfo"]: ) -> Optional[StagingRegisterInfo]:
i = msg_start_offset i = msg_start_offset
base_ptr = ( base_ptr = (
struct.unpack("Q", msg[i])[0] if len(msg) > i and len(msg[i]) == 8 else 0 struct.unpack("Q", msg[i])[0] if len(msg) > i and len(msg[i]) == 8 else 0
python/sglang/srt/disaggregation/decode.py
+3 -3
View File
@@ -146,7 +146,7 @@ class DecodeReqToTokenPool:
def available_size(self): def available_size(self):
return len(self.free_slots) return len(self.free_slots)
def alloc(self, reqs: List["Req"]) -> Optional[List[int]]: def alloc(self, reqs: List[Req]) -> Optional[List[int]]:
# Indices of reqs that already have a req_pool_idx and will reuse # Indices of reqs that already have a req_pool_idx and will reuse
# their existing slot (e.g. chunked prefill continuing across chunks). # their existing slot (e.g. chunked prefill continuing across chunks).
reusing = [i for i, r in enumerate(reqs) if r.req_pool_idx is not None] reusing = [i for i, r in enumerate(reqs) if r.req_pool_idx is not None]
@@ -170,7 +170,7 @@ class DecodeReqToTokenPool:
offset += 1 offset += 1
return [r.req_pool_idx for r in reqs] return [r.req_pool_idx for r in reqs]
def free(self, req: "Req"): def free(self, req: Req):
assert req.req_pool_idx is not None, "request must have req_pool_idx" assert req.req_pool_idx is not None, "request must have req_pool_idx"
self.free_slots.append(req.req_pool_idx) self.free_slots.append(req.req_pool_idx)
req.req_pool_idx = None req.req_pool_idx = None
@@ -186,7 +186,7 @@ class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool):
max_context_len: int, max_context_len: int,
device: str, device: str,
enable_memory_saver: bool, enable_memory_saver: bool,
cache_params: "Mamba2CacheParams", cache_params: Mamba2CacheParams,
mamba_layer_ids: List[int], mamba_layer_ids: List[int],
speculative_num_draft_tokens: int, speculative_num_draft_tokens: int,
enable_mamba_extra_buffer: bool, enable_mamba_extra_buffer: bool,
python/sglang/srt/disaggregation/decode_hicache_mixin.py
+8 -10
View File
@@ -58,7 +58,7 @@ class HiCacheRestoreResult(Enum):
class DecodeHiCachePreallocMixin: class DecodeHiCachePreallocMixin:
"""HiCache hooks for ``DecodePreallocQueue``: issue prefetch + reserve tokens.""" """HiCache hooks for ``DecodePreallocQueue``: issue prefetch + reserve tokens."""
def _build_decode_prefix_match(self, req: "Req", result: Any) -> DecodePrefixMatch: def _build_decode_prefix_match(self, req: Req, result: Any) -> DecodePrefixMatch:
"""Convert a ``match_prefix_for_req`` result into ``DecodePrefixMatch``. """Convert a ``match_prefix_for_req`` result into ``DecodePrefixMatch``.
Performs the optional L3 storage hit length query when decode-side Performs the optional L3 storage hit length query when decode-side
@@ -97,7 +97,7 @@ class DecodeHiCachePreallocMixin:
) )
def _start_hicache_prefetch( def _start_hicache_prefetch(
self, req: "Req", prefix_match: Optional["DecodePrefixMatch"] self, req: Req, prefix_match: Optional[DecodePrefixMatch]
) -> None: ) -> None:
"""Issue L3 storage prefetch after admission succeeds. """Issue L3 storage prefetch after admission succeeds.
@@ -152,7 +152,7 @@ class DecodeHiCachePreallocMixin:
class HiCacheRestoreGatedKVReceiver: class HiCacheRestoreGatedKVReceiver:
"""Wraps a kv_receiver so KVPoll.Success is gated on HiCache restore READY.""" """Wraps a kv_receiver so KVPoll.Success is gated on HiCache restore READY."""
def __init__(self, decode_req: "DecodeRequest"): def __init__(self, decode_req: DecodeRequest):
self.decode_req = decode_req self.decode_req = decode_req
def poll(self) -> KVPoll: def poll(self) -> KVPoll:
@@ -168,7 +168,7 @@ class HiCacheRestoreGatedKVReceiver:
class DecodeHiCacheTransferMixin: class DecodeHiCacheTransferMixin:
"""HiCache hooks for ``DecodeTransferQueue``: drive restore state machine.""" """HiCache hooks for ``DecodeTransferQueue``: drive restore state machine."""
def _clean_hicache_prefetch_resources(self, decode_req: "DecodeRequest") -> None: def _clean_hicache_prefetch_resources(self, decode_req: DecodeRequest) -> None:
if ( if (
decode_req.prefix_match is not None decode_req.prefix_match is not None
and decode_req.prefix_match.prefetch_registered and decode_req.prefix_match.prefetch_registered
@@ -178,7 +178,7 @@ class DecodeHiCacheTransferMixin:
self.tree_cache.dec_lock_ref(decode_req.hicache_restored_node) self.tree_cache.dec_lock_ref(decode_req.hicache_restored_node)
decode_req.hicache_restored_node = None decode_req.hicache_restored_node = None
def _try_hicache_queue_load_back(self, dr: "DecodeRequest") -> bool: def _try_hicache_queue_load_back(self, dr: DecodeRequest) -> bool:
"""Queue one L2->L1 load_back op for ``dr``; True iff a DMA was queued. """Queue one L2->L1 load_back op for ``dr``; True iff a DMA was queued.
On success, ``dr.hicache_restored_node`` and ``hicache_restored_kv_indices`` On success, ``dr.hicache_restored_node`` and ``hicache_restored_kv_indices``
@@ -237,15 +237,13 @@ class DecodeHiCacheTransferMixin:
return False return False
return True return True
def _process_hicache_local_restores( def _process_hicache_local_restores(self, decode_reqs: List[DecodeRequest]) -> None:
self, decode_reqs: List["DecodeRequest"]
) -> None:
if not hasattr(self.tree_cache, "is_load_back_event_done"): if not hasattr(self.tree_cache, "is_load_back_event_done"):
return return
# Filter once: keep only PENDING reqs that still need restore work; # Filter once: keep only PENDING reqs that still need restore work;
# trivially-done reqs (no prefix_match / nothing to restore) flip to READY. # trivially-done reqs (no prefix_match / nothing to restore) flip to READY.
active: List["DecodeRequest"] = [] active: List[DecodeRequest] = []
for dr in decode_reqs: for dr in decode_reqs:
if dr.hicache_restore_status != HiCacheRestoreResult.PENDING: if dr.hicache_restore_status != HiCacheRestoreResult.PENDING:
continue continue
@@ -291,7 +289,7 @@ class DecodeHiCacheTransferMixin:
for dr in queued: for dr in queued:
dr.hicache_load_consumer_index = consumer_index dr.hicache_load_consumer_index = consumer_index
def _commit_hicache_local_restore_to_req(self, decode_req: "DecodeRequest") -> None: def _commit_hicache_local_restore_to_req(self, decode_req: DecodeRequest) -> None:
prefix_match = decode_req.prefix_match prefix_match = decode_req.prefix_match
if prefix_match is None or not prefix_match.needs_local_restore: if prefix_match is None or not prefix_match.needs_local_restore:
return return
python/sglang/srt/disaggregation/encode_receiver.py
+1 -1
View File
@@ -83,7 +83,7 @@ class EncoderBootstrapServer:
self.port = port self.port = port
self._urls: List[str] = urls if urls is not None else [] self._urls: List[str] = urls if urls is not None else []
self._lock = threading.Lock() self._lock = threading.Lock()
self._server: Optional["uvicorn.Server"] = None # set in _run_server self._server: Optional[uvicorn.Server] = None # set in _run_server
self._health_check_interval = ( self._health_check_interval = (
health_check_interval health_check_interval
if health_check_interval is not None if health_check_interval is not None
python/sglang/srt/disaggregation/encode_server.py
+1 -1
View File
@@ -2187,7 +2187,7 @@ class EncoderScheduler:
self.send_sockets = send_sockets self.send_sockets = send_sockets
self.max_batch_size = max(1, int(max_batch_size)) self.max_batch_size = max(1, int(max_batch_size))
self.request_timeout = max(1.0, float(request_timeout)) self.request_timeout = max(1.0, float(request_timeout))
self.pending_queue: "asyncio.Queue[PendingRequest]" = asyncio.Queue() self.pending_queue: asyncio.Queue[PendingRequest] = asyncio.Queue()
self._worker_task: Optional[asyncio.Task] = None self._worker_task: Optional[asyncio.Task] = None
def start(self) -> None: def start(self) -> None:
python/sglang/srt/disaggregation/mori/conn.py
+1 -1
View File
@@ -272,7 +272,7 @@ class TransferTarget:
@dataclasses.dataclass @dataclasses.dataclass
class _TransferChunk: class _TransferChunk:
sender: "MoriKVSender" sender: MoriKVSender
kv_indices: npt.NDArray[np.int32] kv_indices: npt.NDArray[np.int32]
index_slice: slice index_slice: slice
is_last_chunk: bool is_last_chunk: bool
python/sglang/srt/disaggregation/nixl/conn.py
+5 -5
View File
@@ -71,7 +71,7 @@ class TransferInfo:
decode_prefix_len: Optional[int] = None # for decode radix cache decode_prefix_len: Optional[int] = None # for decode radix cache
# NOTE: optional staging field; populated via STAGING_RSP. Keep at the # NOTE: optional staging field; populated via STAGING_RSP. Keep at the
# end so positional construction in from_zmq() continues to work. # end so positional construction in from_zmq() continues to work.
staging: Optional["StagingTransferInfo"] = None staging: Optional[StagingTransferInfo] = None
def is_dummy(self): def is_dummy(self):
# A transfer is "dummy" only for CP non-authoritative ranks. # A transfer is "dummy" only for CP non-authoritative ranks.
@@ -124,7 +124,7 @@ class KVArgsRegisterInfo:
dst_state_dim_per_tensor: List[List[int]] = dataclasses.field(default_factory=list) dst_state_dim_per_tensor: List[List[int]] = dataclasses.field(default_factory=list)
# Keep last: optional, parsed from a variable-length tail of the ZMQ # Keep last: optional, parsed from a variable-length tail of the ZMQ
# frame in from_zmq() below, so positional construction stays stable. # frame in from_zmq() below, so positional construction stays stable.
staging: Optional["StagingRegisterInfo"] = None staging: Optional[StagingRegisterInfo] = None
@classmethod @classmethod
def from_zmq(cls, msg: List[bytes]): def from_zmq(cls, msg: List[bytes]):
@@ -1293,9 +1293,9 @@ class NixlKVManager(CommonKVManager):
def _do_staging_transfer( def _do_staging_transfer(
self, self,
staging_strategy, staging_strategy,
kv_chunk: "TransferKVChunk", kv_chunk: TransferKVChunk,
req: "TransferInfo", req: TransferInfo,
dst_info: "KVArgsRegisterInfo", dst_info: KVArgsRegisterInfo,
queue: FastQueue, queue: FastQueue,
): ):
"""Attempt staging transfer for one chunk. Returns (xfer_handle, deferred). """Attempt staging transfer for one chunk. Returns (xfer_handle, deferred).
python/sglang/srt/entrypoints/openai/protocol.py
+3 -3
View File
@@ -615,7 +615,7 @@ class Tool(BaseModel):
defer_loading: Optional[bool] = None defer_loading: Optional[bool] = None
@model_validator(mode="after") @model_validator(mode="after")
def _propagate_defer_loading(self) -> "Tool": def _propagate_defer_loading(self) -> Tool:
if self.defer_loading is not None and self.function.defer_loading is None: if self.defer_loading is not None and self.function.defer_loading is None:
self.function.defer_loading = self.defer_loading self.function.defer_loading = self.defer_loading
return self return self
@@ -1223,7 +1223,7 @@ class TokenizeRequest(BaseModel):
) )
@model_validator(mode="after") @model_validator(mode="after")
def validate_tokenize_input(self) -> "TokenizeRequest": def validate_tokenize_input(self) -> TokenizeRequest:
if (self.prompt is None) == (self.messages is None): if (self.prompt is None) == (self.messages is None):
raise ValueError("Exactly one of 'prompt' or 'messages' must be provided.") raise ValueError("Exactly one of 'prompt' or 'messages' must be provided.")
return self return self
@@ -1471,7 +1471,7 @@ class ResponsesResponse(BaseModel):
], ],
status: str, status: str,
usage: Optional[UsageInfo], usage: Optional[UsageInfo],
) -> "ResponsesResponse": ) -> ResponsesResponse:
"""Create a response from a request.""" """Create a response from a request."""
# Determine if the output is plain text only to set text.format # Determine if the output is plain text only to set text.format
python/sglang/srt/entrypoints/openai/serving_responses.py
+3 -3
View File
@@ -118,7 +118,7 @@ class OpenAIServingResponses(OpenAIServingChat):
# Note: In production, this should use a proper storage backend (Redis, database) # Note: In production, this should use a proper storage backend (Redis, database)
# with TTL/expiration to prevent memory leaks # with TTL/expiration to prevent memory leaks
self.msg_store: dict[ self.msg_store: dict[
str, Union[list[ChatCompletionMessageParam], list["OpenAIMessage"]] str, Union[list[ChatCompletionMessageParam], list[OpenAIMessage]]
] = {} ] = {}
self.background_tasks: dict[str, asyncio.Task] = {} self.background_tasks: dict[str, asyncio.Task] = {}
@@ -631,8 +631,8 @@ class OpenAIServingResponses(OpenAIServingChat):
self, self,
request: ResponsesRequest, request: ResponsesRequest,
prev_response: Optional[ResponsesResponse], prev_response: Optional[ResponsesResponse],
) -> list["OpenAIMessage"]: ) -> list[OpenAIMessage]:
messages: list["OpenAIMessage"] = [] messages: list[OpenAIMessage] = []
if prev_response is None: if prev_response is None:
# New conversation. # New conversation.
reasoning_effort = request.reasoning.effort if request.reasoning else None reasoning_effort = request.reasoning.effort if request.reasoning else None
python/sglang/srt/eplb/expert_distribution.py
+3 -3
View File
@@ -305,7 +305,7 @@ class _SinglePassGatherer(ABC):
server_args: ServerArgs, server_args: ServerArgs,
expert_location_metadata: ExpertLocationMetadata, expert_location_metadata: ExpertLocationMetadata,
rank: int, rank: int,
) -> "_SinglePassGatherer": ) -> _SinglePassGatherer:
if server_args.expert_distribution_recorder_mode == "per_token": if server_args.expert_distribution_recorder_mode == "per_token":
return _DetailSinglePassGatherer( return _DetailSinglePassGatherer(
server_args, expert_location_metadata, rank server_args, expert_location_metadata, rank
@@ -627,13 +627,13 @@ class _Accumulator(ABC):
server_args: ServerArgs, server_args: ServerArgs,
expert_location_metadata: ExpertLocationMetadata, expert_location_metadata: ExpertLocationMetadata,
rank: int, rank: int,
) -> "_Accumulator": ) -> _Accumulator:
return _Accumulator.get_class(server_args)( return _Accumulator.get_class(server_args)(
server_args, expert_location_metadata, rank server_args, expert_location_metadata, rank
) )
@staticmethod @staticmethod
def get_class(server_args: ServerArgs) -> Type["_Accumulator"]: def get_class(server_args: ServerArgs) -> Type[_Accumulator]:
return { return {
"stat": _StatAccumulator, "stat": _StatAccumulator,
"stat_approx": _StatAccumulator, "stat_approx": _StatAccumulator,
python/sglang/srt/eplb/expert_location.py
+1 -1
View File
@@ -253,7 +253,7 @@ class ExpertLocationMetadata:
def update( def update(
self, self,
other: "ExpertLocationMetadata", other: ExpertLocationMetadata,
update_layer_ids: List[int], update_layer_ids: List[int],
): ):
for field in [ for field in [
python/sglang/srt/hardware_backend/cpu/quantization/awq_kernels.py
+3 -3
View File
@@ -19,7 +19,7 @@ __all__ = ["AWQIntelAMXLinearKernel", "AWQIntelAMXMoEKernel"]
class AWQIntelAMXLinearKernel: class AWQIntelAMXLinearKernel:
def __init__(self, quant_config: "AWQConfig"): def __init__(self, quant_config: AWQConfig):
self.quant_config = quant_config self.quant_config = quant_config
def process_weights_after_loading(self, layer: torch.nn.Module) -> None: def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
@@ -46,7 +46,7 @@ class AWQIntelAMXLinearKernel:
class AWQIntelAMXMoEKernel: class AWQIntelAMXMoEKernel:
def __init__(self, quant_config: "AWQConfig"): def __init__(self, quant_config: AWQConfig):
self.quant_config = quant_config self.quant_config = quant_config
self.moe_runner_config: Optional[MoeRunnerConfig] = None self.moe_runner_config: Optional[MoeRunnerConfig] = None
@@ -66,7 +66,7 @@ class AWQIntelAMXMoEKernel:
def apply( def apply(
self, self,
layer: torch.nn.Module, layer: torch.nn.Module,
dispatch_output: "StandardDispatchOutput", dispatch_output: StandardDispatchOutput,
) -> torch.Tensor: ) -> torch.Tensor:
from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
python/sglang/srt/hardware_backend/cpu/quantization/gptq_kernels.py
+3 -3
View File
@@ -19,7 +19,7 @@ __all__ = ["GPTQIntelAMXLinearKernel", "GPTQIntelAMXMoEKernel"]
class GPTQIntelAMXLinearKernel: class GPTQIntelAMXLinearKernel:
def __init__(self, quant_config: "GPTQConfig"): def __init__(self, quant_config: GPTQConfig):
self.quant_config = quant_config self.quant_config = quant_config
def process_weights_after_loading(self, layer: torch.nn.Module) -> None: def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
@@ -46,7 +46,7 @@ class GPTQIntelAMXLinearKernel:
class GPTQIntelAMXMoEKernel: class GPTQIntelAMXMoEKernel:
def __init__(self, quant_config: "GPTQConfig"): def __init__(self, quant_config: GPTQConfig):
self.quant_config = quant_config self.quant_config = quant_config
self.moe_runner_config: Optional[MoeRunnerConfig] = None self.moe_runner_config: Optional[MoeRunnerConfig] = None
@@ -66,7 +66,7 @@ class GPTQIntelAMXMoEKernel:
def apply( def apply(
self, self,
layer: torch.nn.Module, layer: torch.nn.Module,
dispatch_output: "StandardDispatchOutput", dispatch_output: StandardDispatchOutput,
) -> torch.Tensor: ) -> torch.Tensor:
from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
python/sglang/srt/hardware_backend/gpu/quantization/awq_kernels.py
+5 -5
View File
@@ -77,7 +77,7 @@ _, scalar_types = get_scalar_types()
class AWQLinearKernel: class AWQLinearKernel:
def __init__(self, quant_config: Optional["QuantizationConfig"] = None): def __init__(self, quant_config: Optional[QuantizationConfig] = None):
self.quant_config = quant_config self.quant_config = quant_config
def process_weights_after_loading(self, layer: torch.nn.Module) -> None: def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
@@ -106,7 +106,7 @@ class AWQLinearKernel:
class AWQMarlinLinearKernel: class AWQMarlinLinearKernel:
def __init__(self, quant_config: Optional["QuantizationConfig"] = None): def __init__(self, quant_config: Optional[QuantizationConfig] = None):
self.quant_config = quant_config self.quant_config = quant_config
def process_weights_after_loading(self, layer: torch.nn.Module) -> None: def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
@@ -166,7 +166,7 @@ class AWQMarlinLinearKernel:
class AWQMoEKernel: class AWQMoEKernel:
def __init__(self, quant_config: Optional["QuantizationConfig"] = None): def __init__(self, quant_config: Optional[QuantizationConfig] = None):
self.quant_config = quant_config self.quant_config = quant_config
self.runner: Optional[MoeRunner] = None self.runner: Optional[MoeRunner] = None
@@ -236,8 +236,8 @@ class AWQMoEKernel:
def apply( def apply(
self, self,
layer: torch.nn.Module, layer: torch.nn.Module,
dispatch_output: "StandardDispatchOutput", dispatch_output: StandardDispatchOutput,
) -> "CombineInput": ) -> CombineInput:
if self.runner is None: if self.runner is None:
raise RuntimeError("moe runner is not initialized") raise RuntimeError("moe runner is not initialized")
python/sglang/srt/hardware_backend/gpu/quantization/gptq_kernels.py
+6 -6
View File
@@ -83,7 +83,7 @@ def gptq_marlin_moe_repack(
class GPTQLinearKernel: class GPTQLinearKernel:
def __init__(self, quant_config: Optional["QuantizationConfig"] = None): def __init__(self, quant_config: Optional[QuantizationConfig] = None):
self.quant_config = quant_config self.quant_config = quant_config
self.use_shuffle = True self.use_shuffle = True
@@ -129,7 +129,7 @@ class GPTQLinearKernel:
class GPTQMarlinLinearKernel: class GPTQMarlinLinearKernel:
def __init__(self, quant_config: Optional["QuantizationConfig"] = None): def __init__(self, quant_config: Optional[QuantizationConfig] = None):
self.quant_config = quant_config self.quant_config = quant_config
def process_weights_after_loading(self, layer: torch.nn.Module) -> None: def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
@@ -270,7 +270,7 @@ class GPTQMarlinLinearKernel:
class GPTQMarlinMoEKernel: class GPTQMarlinMoEKernel:
def __init__(self, quant_config: Optional["QuantizationConfig"] = None): def __init__(self, quant_config: Optional[QuantizationConfig] = None):
self.quant_config = quant_config self.quant_config = quant_config
def process_weights_after_loading(self, layer: torch.nn.Module) -> None: def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
@@ -355,7 +355,7 @@ class GPTQMarlinMoEKernel:
replace_parameter(layer, "w2_scales", marlin_w2_scales) replace_parameter(layer, "w2_scales", marlin_w2_scales)
def create_moe_runner( def create_moe_runner(
self, layer: torch.nn.Module, moe_runner_config: "MoeRunnerConfig" self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
): ):
assert get_moe_runner_backend().is_auto() assert get_moe_runner_backend().is_auto()
self.moe_runner_config = moe_runner_config self.moe_runner_config = moe_runner_config
@@ -364,8 +364,8 @@ class GPTQMarlinMoEKernel:
def apply( def apply(
self, self,
layer: torch.nn.Module, layer: torch.nn.Module,
dispatch_output: "StandardDispatchOutput", dispatch_output: StandardDispatchOutput,
) -> "CombineInput": ) -> CombineInput:
quant_info = MarlinMoeQuantInfo( quant_info = MarlinMoeQuantInfo(
w13_qweight=layer.w13_qweight, w13_qweight=layer.w13_qweight,
w2_qweight=layer.w2_qweight, w2_qweight=layer.w2_qweight,
python/sglang/srt/hardware_backend/mlx/aot.py
+1 -1
View File
@@ -204,7 +204,7 @@ class MlxAOTKernelContext:
req_pool_idx: dict[str, int], req_pool_idx: dict[str, int],
req_to_token_pool: Any | None, req_to_token_pool: Any | None,
layer_caches: list[list[ContiguousAttentionKVCache]], layer_caches: list[list[ContiguousAttentionKVCache]],
) -> "MlxAOTKernelContext": ) -> MlxAOTKernelContext:
"""Build optional AOT context for one batched decode step.""" """Build optional AOT context for one batched decode step."""
if not aot_kernels.rope.enabled or kv_pool is None: if not aot_kernels.rope.enabled or kv_pool is None:
return cls() return cls()
python/sglang/srt/hardware_backend/mlx/kv_cache/attention_wrapper.py
+1 -1
View File
@@ -76,7 +76,7 @@ class BatchedDecodeContext:
req_to_token_pool: Any | None, req_to_token_pool: Any | None,
attention_layer_indices: list[int] | None = None, attention_layer_indices: list[int] | None = None,
attention_pool_index_by_layer: dict[int, int] | None = None, attention_pool_index_by_layer: dict[int, int] | None = None,
) -> "BatchedDecodeContext": ) -> BatchedDecodeContext:
batch_size = len(req_ids) batch_size = len(req_ids)
if attention_layer_indices is None: if attention_layer_indices is None:
attention_layer_indices = list(range(len(caches[0]))) attention_layer_indices = list(range(len(caches[0])))
python/sglang/srt/hardware_backend/mlx/kv_cache/layout.py
+1 -1
View File
@@ -26,7 +26,7 @@ class MlxModelCacheLayout:
cls, cls,
layers: Sequence[Any], layers: Sequence[Any],
attention_attrs: Sequence[str | None], attention_attrs: Sequence[str | None],
) -> "MlxModelCacheLayout": ) -> MlxModelCacheLayout:
if len(layers) != len(attention_attrs): if len(layers) != len(attention_attrs):
raise ValueError( raise ValueError(
"Layer count and attention attribute count differ: " "Layer count and attention attribute count differ: "
python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py
+8 -8
View File
@@ -73,19 +73,19 @@ class MlxPendingJob:
""" """
lazy_tokens: Optional[mx.array] lazy_tokens: Optional[mx.array]
prefills: list["MlxPendingPrefill"] prefills: list[MlxPendingPrefill]
extends: list["MlxPendingExtend"] extends: list[MlxPendingExtend]
decode: Optional["MlxPendingDecode"] decode: Optional[MlxPendingDecode]
mode: str mode: str
batch_copy: "ScheduleBatch" batch_copy: ScheduleBatch
schedule_batch: "ScheduleBatch" schedule_batch: ScheduleBatch
reqs: List[Req] reqs: List[Req]
class SchedulerMlxOverlapMixin: class SchedulerMlxOverlapMixin:
"""Mixin that adds MLX overlap scheduling to :class:`Scheduler`.""" """Mixin that adds MLX overlap scheduling to :class:`Scheduler`."""
def _finalize_mlx_pending_job(self: "Scheduler", pending: MlxPendingJob): def _finalize_mlx_pending_job(self: Scheduler, pending: MlxPendingJob):
result = self.tp_worker.finalize_mlx_result( result = self.tp_worker.finalize_mlx_result(
pending.prefills, pending.prefills,
pending.extends, pending.extends,
@@ -100,7 +100,7 @@ class SchedulerMlxOverlapMixin:
self.process_batch_result(pending.batch_copy, result) self.process_batch_result(pending.batch_copy, result)
@DynamicGradMode() @DynamicGradMode()
def event_loop_overlap_mlx(self: "Scheduler"): def event_loop_overlap_mlx(self: Scheduler):
"""MLX-specific overlap loop modelled on ``mlx_lm.generate.generate_step``. """MLX-specific overlap loop modelled on ``mlx_lm.generate.generate_step``.
At steady state we keep TWO in-flight MLX graphs queued on the At steady state we keep TWO in-flight MLX graphs queued on the
@@ -142,7 +142,7 @@ class SchedulerMlxOverlapMixin:
pending_curr: Optional[MlxPendingJob] = None pending_curr: Optional[MlxPendingJob] = None
pending_next: Optional[MlxPendingJob] = None pending_next: Optional[MlxPendingJob] = None
def _launch_fresh(batch: "ScheduleBatch") -> MlxPendingJob: def _launch_fresh(batch: ScheduleBatch) -> MlxPendingJob:
# Materialize batch.input_ids from CPU staging (prefill) or the # Materialize batch.input_ids from CPU staging (prefill) or the
# FutureMap relay (decode) before the forward. With deferred input # FutureMap relay (decode) before the forward. With deferred input
# materialization, get_next_batch_to_run leaves input_ids unset; the # materialization, get_next_batch_to_run leaves input_ids unset; the
python/sglang/srt/hardware_backend/musa/attention/flashattention_backend.py
+2 -2
View File
@@ -39,11 +39,11 @@ _MATE_NO_MLA_SCHEDULER_METADATA_DICT: dict = {}
_MATE_NO_MLA_SCHEDULER_METADATA_LOCK = threading.Lock() _MATE_NO_MLA_SCHEDULER_METADATA_LOCK = threading.Lock()
# Global reference to the current backend instance (set during __init__) # Global reference to the current backend instance (set during __init__)
_CURRENT_BACKEND: Optional["MusaFlashAttentionBackend"] = None _CURRENT_BACKEND: Optional[MusaFlashAttentionBackend] = None
def _compute_scheduler_metadata( def _compute_scheduler_metadata(
backend: "MusaFlashAttentionBackend", backend: MusaFlashAttentionBackend,
cu_seqlens_q: torch.Tensor, cu_seqlens_q: torch.Tensor,
cu_seqlens_k_new: Optional[torch.Tensor], cu_seqlens_k_new: Optional[torch.Tensor],
cache_seqlens: torch.Tensor, cache_seqlens: torch.Tensor,
python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py
+2 -2
View File
@@ -530,7 +530,7 @@ class AscendAttnBackend(AttentionBackend):
forward_mode: ForwardMode, forward_mode: ForwardMode,
seq_lens: torch.Tensor, seq_lens: torch.Tensor,
out_cache_loc: Optional[torch.Tensor] = None, out_cache_loc: Optional[torch.Tensor] = None,
) -> "ForwardMetadata": ) -> ForwardMetadata:
"""Create and store the per-bs ForwardMetadata for CUDA graph capture.""" """Create and store the per-bs ForwardMetadata for CUDA graph capture."""
metadata = ForwardMetadata() metadata = ForwardMetadata()
metadata.block_tables = self.graph_metadata["block_tables"][:bs, :] metadata.block_tables = self.graph_metadata["block_tables"][:bs, :]
@@ -847,7 +847,7 @@ class AscendAttnBackend(AttentionBackend):
q: torch.Tensor, q: torch.Tensor,
k_cache: torch.Tensor, k_cache: torch.Tensor,
v_cache: torch.Tensor, v_cache: torch.Tensor,
layer: "RadixAttention", layer: RadixAttention,
forward_batch: ForwardBatch, forward_batch: ForwardBatch,
) -> torch.Tensor: ) -> torch.Tensor:
"""CP-aware attention for standard (non-MLA) models using FIA on Ascend NPU. """CP-aware attention for standard (non-MLA) models using FIA on Ascend NPU.
python/sglang/srt/hardware_backend/npu/moe/fuseep.py
+3 -3
View File
@@ -26,7 +26,7 @@ if TYPE_CHECKING:
_PARAMS_BYTES = 2 # bf16 — Ascend's Dispatch & Combine does not support fp16 _PARAMS_BYTES = 2 # bf16 — Ascend's Dispatch & Combine does not support fp16
def _get_fuseep_buffer(layer: "FusedMoE"): def _get_fuseep_buffer(layer: FusedMoE):
DeepEPBuffer.set_dispatch_mode_as_low_latency() DeepEPBuffer.set_dispatch_mode_as_low_latency()
return DeepEPBuffer.get_deepep_buffer( return DeepEPBuffer.get_deepep_buffer(
get_tp_group().device_group, get_tp_group().device_group,
@@ -39,9 +39,9 @@ def _get_fuseep_buffer(layer: "FusedMoE"):
def forward_fuseep( def forward_fuseep(
layer: "FusedMoE", layer: FusedMoE,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
topk_output: "TopKOutput", topk_output: TopKOutput,
) -> torch.Tensor: ) -> torch.Tensor:
buf = _get_fuseep_buffer(layer) buf = _get_fuseep_buffer(layer)
hidden_states, _ = buf.fused_deep_moe( hidden_states, _ = buf.fused_deep_moe(
python/sglang/srt/hardware_backend/npu/quantization/awq_kernels.py
+3 -3
View File
@@ -17,7 +17,7 @@ import torch_npu
class AWQAscendLinearKernel: class AWQAscendLinearKernel:
def __init__(self, quant_config: Optional["QuantizationConfig"] = None): def __init__(self, quant_config: Optional[QuantizationConfig] = None):
self.quant_config = quant_config self.quant_config = quant_config
def process_weights_after_loading(self, layer: torch.nn.Module) -> None: def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
@@ -72,7 +72,7 @@ class AWQAscendLinearKernel:
class AWQAscendMoEKernel: class AWQAscendMoEKernel:
def __init__(self, quant_config: Optional["QuantizationConfig"] = None): def __init__(self, quant_config: Optional[QuantizationConfig] = None):
self.quant_config = quant_config self.quant_config = quant_config
self.kernel = NPUW4A16Int4DynamicMoEMethod() self.kernel = NPUW4A16Int4DynamicMoEMethod()
@@ -151,7 +151,7 @@ class AWQAscendMoEKernel:
def apply( def apply(
self, self,
layer: torch.nn.Module, layer: torch.nn.Module,
dispatch_output: "StandardDispatchOutput", dispatch_output: StandardDispatchOutput,
) -> torch.Tensor: ) -> torch.Tensor:
return self.kernel.apply(layer, dispatch_output) return self.kernel.apply(layer, dispatch_output)
python/sglang/srt/hardware_backend/npu/quantization/gptq_kernels.py
+5 -5
View File
@@ -60,7 +60,7 @@ def unpack_from_int32(
class GPTQLinearAscendKernel: class GPTQLinearAscendKernel:
def __init__(self, quant_config: Optional["QuantizationConfig"] = None): def __init__(self, quant_config: Optional[QuantizationConfig] = None):
self.quant_config = quant_config self.quant_config = quant_config
self.use_v2_format = quant_config.checkpoint_format == "gptq_v2" self.use_v2_format = quant_config.checkpoint_format == "gptq_v2"
@@ -128,15 +128,15 @@ class GPTQLinearAscendKernel:
class GPTQMoEAscendKernel: class GPTQMoEAscendKernel:
def __init__(self, quant_config: Optional["QuantizationConfig"] = None): def __init__(self, quant_config: Optional[QuantizationConfig] = None):
self.quant_config = quant_config self.quant_config = quant_config
self.use_v2_format = quant_config.checkpoint_format == "gptq_v2" self.use_v2_format = quant_config.checkpoint_format == "gptq_v2"
self.moe_runner_config: Optional["MoeRunnerConfig"] = None self.moe_runner_config: Optional[MoeRunnerConfig] = None
def create_moe_runner( def create_moe_runner(
self, self,
layer: torch.nn.Module, layer: torch.nn.Module,
moe_runner_config: "MoeRunnerConfig", moe_runner_config: MoeRunnerConfig,
**extra_weight_attrs, **extra_weight_attrs,
): ):
self.moe_runner_config = moe_runner_config self.moe_runner_config = moe_runner_config
@@ -277,7 +277,7 @@ class GPTQMoEAscendKernel:
def apply( def apply(
self, self,
layer: torch.nn.Module, layer: torch.nn.Module,
dispatch_output: "StandardDispatchOutput", dispatch_output: StandardDispatchOutput,
) -> torch.Tensor: ) -> torch.Tensor:
from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
python/sglang/srt/kv_canary/api.py
+4 -6
View File
@@ -29,9 +29,9 @@ logger = logging.getLogger(__name__)
def install_canary( def install_canary(
*, *,
server_args: "ServerArgs", server_args: ServerArgs,
model_runner: "ModelRunner", model_runner: ModelRunner,
token_oracle_manager: Optional["TokenOracleManager"] = None, token_oracle_manager: Optional[TokenOracleManager] = None,
) -> Optional[CanaryManager]: ) -> Optional[CanaryManager]:
config = CanaryConfig.from_env(server_args) config = CanaryConfig.from_env(server_args)
if config.mode is CanaryMode.NONE: if config.mode is CanaryMode.NONE:
@@ -101,9 +101,7 @@ def install_canary(
return manager return manager
def _patch_model_forward( def _patch_model_forward(*, model_runner: ModelRunner, manager: CanaryManager) -> None:
*, model_runner: "ModelRunner", manager: CanaryManager
) -> None:
def _with_canary_bracketing(original: Callable, *args: Any, **kwargs: Any) -> Any: def _with_canary_bracketing(original: Callable, *args: Any, **kwargs: Any) -> Any:
forward_batch = _extract_forward_batch(args, kwargs) forward_batch = _extract_forward_batch(args, kwargs)
assert ( assert (
python/sglang/srt/kv_canary/capacities.py
+2 -2
View File
@@ -43,11 +43,11 @@ class CanaryLaunchCapacities:
def from_args( def from_args(
cls, cls,
*, *,
server_args: "ServerArgs", server_args: ServerArgs,
req_to_token_pool_size: int, req_to_token_pool_size: int,
max_seq_len_per_req: int, max_seq_len_per_req: int,
pool_slot_count: int, pool_slot_count: int,
) -> "CanaryLaunchCapacities": ) -> CanaryLaunchCapacities:
if req_to_token_pool_size <= 0: if req_to_token_pool_size <= 0:
raise ValueError( raise ValueError(
"kv-canary: req_to_token_pool_size must be positive, " "kv-canary: req_to_token_pool_size must be positive, "
python/sglang/srt/kv_canary/config.py
+1 -1
View File
@@ -60,7 +60,7 @@ class CanaryConfig:
stats_print_every_n_steps: int stats_print_every_n_steps: int
@classmethod @classmethod
def from_env(cls, server_args: "ServerArgs") -> "CanaryConfig": def from_env(cls, server_args: ServerArgs) -> CanaryConfig:
mode_raw = server_args.kv_canary.strip().lower() mode_raw = server_args.kv_canary.strip().lower()
if mode_raw not in ("none", "log", "raise"): if mode_raw not in ("none", "log", "raise"):
raise ValueError( raise ValueError(
python/sglang/srt/kv_canary/expected_inputs.py
+2 -2
View File
@@ -11,13 +11,13 @@ class ExpectedInputs:
positions: torch.Tensor positions: torch.Tensor
@classmethod @classmethod
def allocate(cls, *, capacity: int, device: torch.device) -> "ExpectedInputs": def allocate(cls, *, capacity: int, device: torch.device) -> ExpectedInputs:
return cls( return cls(
tokens=torch.empty(capacity, dtype=torch.int64, device=device), tokens=torch.empty(capacity, dtype=torch.int64, device=device),
positions=torch.empty(capacity, dtype=torch.int64, device=device), positions=torch.empty(capacity, dtype=torch.int64, device=device),
) )
def slice(self, num_tokens: int) -> "ExpectedInputs": def slice(self, num_tokens: int) -> ExpectedInputs:
return ExpectedInputs( return ExpectedInputs(
tokens=self.tokens[:num_tokens], tokens=self.tokens[:num_tokens],
positions=self.positions[:num_tokens], positions=self.positions[:num_tokens],
python/sglang/srt/kv_canary/perturb/config.py
+1 -1
View File
@@ -25,7 +25,7 @@ class PerturbConfig:
warmup_steps: int warmup_steps: int
@classmethod @classmethod
def from_env(cls) -> "PerturbConfig": def from_env(cls) -> PerturbConfig:
real_kv_used_prob = envs.SGLANG_KV_CANARY_PERTURB_REAL_KV_USED_PROB.get() real_kv_used_prob = envs.SGLANG_KV_CANARY_PERTURB_REAL_KV_USED_PROB.get()
real_kv_unused_cache_prob = ( real_kv_unused_cache_prob = (
envs.SGLANG_KV_CANARY_PERTURB_REAL_KV_UNUSED_CACHE_PROB.get() envs.SGLANG_KV_CANARY_PERTURB_REAL_KV_UNUSED_CACHE_PROB.get()
python/sglang/srt/kv_canary/perturb/manager.py
+9 -9
View File
@@ -24,7 +24,7 @@ class PerturbManager:
self, self,
*, *,
config: PerturbConfig, config: PerturbConfig,
req_to_token_pool: "ReqToTokenPool", req_to_token_pool: ReqToTokenPool,
buffer_groups: tuple[CanaryBufferGroup, ...], buffer_groups: tuple[CanaryBufferGroup, ...],
outer_step_counter_getter: Callable[[], int], outer_step_counter_getter: Callable[[], int],
swa_window_size: int = 0, swa_window_size: int = 0,
@@ -36,18 +36,18 @@ class PerturbManager:
self._outer_step_counter_getter = outer_step_counter_getter self._outer_step_counter_getter = outer_step_counter_getter
self._swa_window_size = swa_window_size self._swa_window_size = swa_window_size
self._sweep_interval = sweep_interval self._sweep_interval = sweep_interval
self._radix_cache: Optional["BasePrefixCache"] = None self._radix_cache: Optional[BasePrefixCache] = None
self._warmup_gate = WarmupGate( self._warmup_gate = WarmupGate(
config=config, outer_step_counter_getter=outer_step_counter_getter config=config, outer_step_counter_getter=outer_step_counter_getter
) )
def attach_radix_cache(self, radix_cache: "BasePrefixCache") -> None: def attach_radix_cache(self, radix_cache: BasePrefixCache) -> None:
self._radix_cache = radix_cache self._radix_cache = radix_cache
def perturb( def perturb(
self, self,
*, *,
maybe_inaccurate_forward_batch: Optional["ForwardBatch"], maybe_inaccurate_forward_batch: Optional[ForwardBatch],
) -> None: ) -> None:
self.perturb_req_to_token(maybe_inaccurate_forward_batch) self.perturb_req_to_token(maybe_inaccurate_forward_batch)
self.perturb_real_kv_used(maybe_inaccurate_forward_batch) self.perturb_real_kv_used(maybe_inaccurate_forward_batch)
@@ -56,12 +56,12 @@ class PerturbManager:
def perturb_post_forward( def perturb_post_forward(
self, self,
*, *,
maybe_inaccurate_forward_batch: Optional["ForwardBatch"], maybe_inaccurate_forward_batch: Optional[ForwardBatch],
) -> None: ) -> None:
self.perturb_real_kv_post_forward(maybe_inaccurate_forward_batch) self.perturb_real_kv_post_forward(maybe_inaccurate_forward_batch)
def perturb_req_to_token( def perturb_req_to_token(
self, maybe_inaccurate_forward_batch: Optional["ForwardBatch"] self, maybe_inaccurate_forward_batch: Optional[ForwardBatch]
) -> None: ) -> None:
req_to_token.run( req_to_token.run(
maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch, maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
@@ -71,7 +71,7 @@ class PerturbManager:
) )
def perturb_real_kv_used( def perturb_real_kv_used(
self, maybe_inaccurate_forward_batch: Optional["ForwardBatch"] self, maybe_inaccurate_forward_batch: Optional[ForwardBatch]
) -> None: ) -> None:
real_kv_used.run( real_kv_used.run(
maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch, maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
@@ -83,7 +83,7 @@ class PerturbManager:
) )
def perturb_real_kv_unused_cache( def perturb_real_kv_unused_cache(
self, maybe_inaccurate_forward_batch: Optional["ForwardBatch"] self, maybe_inaccurate_forward_batch: Optional[ForwardBatch]
) -> None: ) -> None:
real_kv_unused_cache.run( real_kv_unused_cache.run(
maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch, maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
@@ -97,7 +97,7 @@ class PerturbManager:
) )
def perturb_real_kv_post_forward( def perturb_real_kv_post_forward(
self, maybe_inaccurate_forward_batch: Optional["ForwardBatch"] self, maybe_inaccurate_forward_batch: Optional[ForwardBatch]
) -> None: ) -> None:
real_kv_post_forward.run( real_kv_post_forward.run(
maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch, maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch,
python/sglang/srt/kv_canary/perturb/next_token_swap.py
+1 -1
View File
@@ -25,7 +25,7 @@ class NextTokenSwapConfig:
warmup_steps: int warmup_steps: int
@classmethod @classmethod
def from_env(cls) -> "NextTokenSwapConfig": def from_env(cls) -> NextTokenSwapConfig:
return cls( return cls(
prob=envs.SGLANG_KV_CANARY_PERTURB_NEXT_TOKEN_SWAP_PROB.get(), prob=envs.SGLANG_KV_CANARY_PERTURB_NEXT_TOKEN_SWAP_PROB.get(),
warmup_steps=envs.SGLANG_KV_CANARY_PERTURB_WARMUP_STEPS.get(), warmup_steps=envs.SGLANG_KV_CANARY_PERTURB_WARMUP_STEPS.get(),
python/sglang/srt/kv_canary/perturb/real_kv_post_forward.py
+1 -1
View File
@@ -34,7 +34,7 @@ logger = logging.getLogger(__name__)
def run( def run(
*, *,
maybe_inaccurate_forward_batch: Optional["ForwardBatch"], maybe_inaccurate_forward_batch: Optional[ForwardBatch],
config: PerturbConfig, config: PerturbConfig,
buffer_groups: tuple[CanaryBufferGroup, ...], buffer_groups: tuple[CanaryBufferGroup, ...],
warmup_gate: WarmupGate, warmup_gate: WarmupGate,
python/sglang/srt/kv_canary/perturb/real_kv_unused_cache.py
+3 -3
View File
@@ -34,10 +34,10 @@ logger = logging.getLogger(__name__)
def run( def run(
*, *,
maybe_inaccurate_forward_batch: Optional["ForwardBatch"], maybe_inaccurate_forward_batch: Optional[ForwardBatch],
config: PerturbConfig, config: PerturbConfig,
buffer_groups: tuple[CanaryBufferGroup, ...], buffer_groups: tuple[CanaryBufferGroup, ...],
radix_cache: Optional["BasePrefixCache"], radix_cache: Optional[BasePrefixCache],
swa_window_size: int, swa_window_size: int,
sweep_interval: int, sweep_interval: int,
outer_step_counter: int, outer_step_counter: int,
@@ -114,7 +114,7 @@ def run(
def _pick_sweep_slot_for_group( def _pick_sweep_slot_for_group(
*, *,
radix_cache: Optional["BasePrefixCache"], radix_cache: Optional[BasePrefixCache],
group: CanaryBufferGroup, group: CanaryBufferGroup,
swa_window_size: int, swa_window_size: int,
) -> Optional[int]: ) -> Optional[int]:
python/sglang/srt/kv_canary/perturb/real_kv_used.py
+4 -4
View File
@@ -37,9 +37,9 @@ logger = logging.getLogger(__name__)
def run( def run(
*, *,
maybe_inaccurate_forward_batch: Optional["ForwardBatch"], maybe_inaccurate_forward_batch: Optional[ForwardBatch],
config: PerturbConfig, config: PerturbConfig,
req_to_token_pool: "ReqToTokenPool", req_to_token_pool: ReqToTokenPool,
buffer_groups: tuple[CanaryBufferGroup, ...], buffer_groups: tuple[CanaryBufferGroup, ...],
swa_window_size: int, swa_window_size: int,
warmup_gate: WarmupGate, warmup_gate: WarmupGate,
@@ -111,8 +111,8 @@ def run(
def _pick_active_slot_for_group( def _pick_active_slot_for_group(
*, *,
maybe_inaccurate_forward_batch: "ForwardBatch", maybe_inaccurate_forward_batch: ForwardBatch,
req_to_token_pool: "ReqToTokenPool", req_to_token_pool: ReqToTokenPool,
group: CanaryBufferGroup, group: CanaryBufferGroup,
swa_window_size: int, swa_window_size: int,
) -> Optional[ReqToTokenEntry]: ) -> Optional[ReqToTokenEntry]:
python/sglang/srt/kv_canary/perturb/req_to_token.py
+2 -2
View File
@@ -26,9 +26,9 @@ logger = logging.getLogger(__name__)
def run( def run(
*, *,
maybe_inaccurate_forward_batch: Optional["ForwardBatch"], maybe_inaccurate_forward_batch: Optional[ForwardBatch],
config: PerturbConfig, config: PerturbConfig,
req_to_token_pool: "ReqToTokenPool", req_to_token_pool: ReqToTokenPool,
warmup_gate: WarmupGate, warmup_gate: WarmupGate,
) -> None: ) -> None:
if not should_run_perturbation( if not should_run_perturbation(
python/sglang/srt/kv_canary/perturb/slot_picker.py
+3 -3
View File
@@ -21,8 +21,8 @@ class ReqToTokenEntry:
def collect_active_slots( def collect_active_slots(
*, *,
maybe_inaccurate_forward_batch: "ForwardBatch", maybe_inaccurate_forward_batch: ForwardBatch,
req_to_token_pool: "ReqToTokenPool", req_to_token_pool: ReqToTokenPool,
exclude_out_cache_loc: bool = True, exclude_out_cache_loc: bool = True,
) -> list[ReqToTokenEntry]: ) -> list[ReqToTokenEntry]:
"""Collect every (req_pool_idx, position, value) triple for currently-active reqs. """Collect every (req_pool_idx, position, value) triple for currently-active reqs.
@@ -79,7 +79,7 @@ def collect_active_slots(
def pick_out_cache_loc_slot( def pick_out_cache_loc_slot(
*, maybe_inaccurate_forward_batch: "ForwardBatch" *, maybe_inaccurate_forward_batch: ForwardBatch
) -> Optional[int]: ) -> Optional[int]:
out_cache_loc = maybe_inaccurate_forward_batch.out_cache_loc out_cache_loc = maybe_inaccurate_forward_batch.out_cache_loc
if out_cache_loc is None: if out_cache_loc is None:
python/sglang/srt/kv_canary/perturb/utils.py
+1 -1
View File
@@ -70,7 +70,7 @@ def should_run_perturbation(
perturb_name: str, perturb_name: str,
probability: float, probability: float,
warmup_gate: WarmupGate, warmup_gate: WarmupGate,
maybe_inaccurate_forward_batch: Optional["ForwardBatch"], maybe_inaccurate_forward_batch: Optional[ForwardBatch],
require_forward_batch: bool = True, require_forward_batch: bool = True,
) -> bool: ) -> bool:
if probability <= 0.0: if probability <= 0.0:
python/sglang/srt/kv_canary/plan_input.py
+3 -3
View File
@@ -53,7 +53,7 @@ class PlanInput:
*, *,
bs_capacity: int, bs_capacity: int,
device: torch.device, device: torch.device,
) -> "PlanInput": ) -> PlanInput:
return cls( return cls(
req_pool_indices=torch.zeros(bs_capacity, dtype=torch.int64, device=device), req_pool_indices=torch.zeros(bs_capacity, dtype=torch.int64, device=device),
prefix_lens=torch.zeros(bs_capacity, dtype=torch.int64, device=device), prefix_lens=torch.zeros(bs_capacity, dtype=torch.int64, device=device),
@@ -63,7 +63,7 @@ class PlanInput:
), ),
) )
def fill_from_forward_batch(self, *, forward_batch: "ForwardBatch") -> None: def fill_from_forward_batch(self, *, forward_batch: ForwardBatch) -> None:
req_pool_indices = forward_batch.req_pool_indices req_pool_indices = forward_batch.req_pool_indices
bs = int(req_pool_indices.shape[0]) bs = int(req_pool_indices.shape[0])
capacity = int(self.req_pool_indices.shape[0]) capacity = int(self.req_pool_indices.shape[0])
@@ -92,7 +92,7 @@ class PlanInput:
def _extract_prefix_lens_and_extend_seq_lens( def _extract_prefix_lens_and_extend_seq_lens(
*, *,
forward_batch: "ForwardBatch", forward_batch: ForwardBatch,
out_prefix_lens: torch.Tensor, out_prefix_lens: torch.Tensor,
out_extend_seq_lens: torch.Tensor, out_extend_seq_lens: torch.Tensor,
bs: int, bs: int,
python/sglang/srt/kv_canary/radix_cache_walker.py
+7 -7
View File
@@ -22,7 +22,7 @@ class RadixCacheWalkResult:
def walk_radix_cache_for_canary( def walk_radix_cache_for_canary(
*, *,
radix_cache: "BasePrefixCache", radix_cache: BasePrefixCache,
unlocked_only: bool = False, unlocked_only: bool = False,
swa_resident_only: bool = False, swa_resident_only: bool = False,
) -> RadixCacheWalkResult: ) -> RadixCacheWalkResult:
@@ -68,8 +68,8 @@ def walk_radix_cache_for_canary(
def _walk_radix_subtree( def _walk_radix_subtree(
*, *,
node: "TreeNode", node: TreeNode,
radix_cache: "BasePrefixCache", radix_cache: BasePrefixCache,
depth: int, depth: int,
parent_last_slot: int, parent_last_slot: int,
slot_buf: list[int], slot_buf: list[int],
@@ -123,8 +123,8 @@ def _walk_radix_subtree(
def _node_is_unlocked_for_canary( def _node_is_unlocked_for_canary(
*, *,
node: "TreeNode", node: TreeNode,
radix_cache: "BasePrefixCache", radix_cache: BasePrefixCache,
) -> bool: ) -> bool:
if type(radix_cache) is RadixCache: if type(radix_cache) is RadixCache:
return node.lock_ref == 0 return node.lock_ref == 0
@@ -139,8 +139,8 @@ def _node_is_unlocked_for_canary(
def _node_is_swa_resident_for_canary( def _node_is_swa_resident_for_canary(
*, *,
node: "TreeNode", node: TreeNode,
radix_cache: "BasePrefixCache", radix_cache: BasePrefixCache,
) -> bool: ) -> bool:
if type(radix_cache) is not SWARadixCache: if type(radix_cache) is not SWARadixCache:
return True return True
python/sglang/srt/kv_canary/req_to_expected_token_ids_manager.py
+2 -2
View File
@@ -15,7 +15,7 @@ if TYPE_CHECKING:
def compute_req_all_ids_info( def compute_req_all_ids_info(
reqs: "list[Req]", reqs: list[Req],
) -> tuple[torch.Tensor, torch.Tensor]: ) -> tuple[torch.Tensor, torch.Tensor]:
"""Snapshot per-req (origin_input_ids + output_ids) as pinned CPU int64 tensors. """Snapshot per-req (origin_input_ids + output_ids) as pinned CPU int64 tensors.
@@ -38,7 +38,7 @@ def compute_req_all_ids_info(
def populate_req_to_expected_token_ids( def populate_req_to_expected_token_ids(
*, *,
forward_batch: "ForwardBatch", forward_batch: ForwardBatch,
req_to_verify_expected_tokens: Optional[torch.Tensor], req_to_verify_expected_tokens: Optional[torch.Tensor],
) -> None: ) -> None:
"""Scatter the forward batch's per-req token-id snapshot into the device-side pool.""" """Scatter the forward batch's per-req token-id snapshot into the device-side pool."""
python/sglang/srt/kv_canary/runner/canary_manager.py
+9 -9
View File
@@ -47,18 +47,18 @@ class CanaryManager:
perturb_config: PerturbConfig, perturb_config: PerturbConfig,
buffer_groups: tuple[CanaryBufferGroup, ...], buffer_groups: tuple[CanaryBufferGroup, ...],
device: torch.device, device: torch.device,
req_to_token_pool: "ReqToTokenPool", req_to_token_pool: ReqToTokenPool,
launch_capacities: CanaryLaunchCapacities, launch_capacities: CanaryLaunchCapacities,
swa_window_size: int = 0, swa_window_size: int = 0,
token_oracle_manager: Optional[TokenOracleManager] = None, token_oracle_manager: Optional[TokenOracleManager] = None,
swa_allocator: Optional["SWATokenToKVPoolAllocator"] = None, swa_allocator: Optional[SWATokenToKVPoolAllocator] = None,
speculative_num_steps: int = 1, speculative_num_steps: int = 1,
is_eagle_draft_decode: bool = False, is_eagle_draft_decode: bool = False,
) -> None: ) -> None:
self.config = config self.config = config
self._req_to_token_pool = req_to_token_pool self._req_to_token_pool = req_to_token_pool
self._swa_window_size = swa_window_size self._swa_window_size = swa_window_size
self._swa_allocator: Optional["SWATokenToKVPoolAllocator"] = swa_allocator self._swa_allocator: Optional[SWATokenToKVPoolAllocator] = swa_allocator
self._outer_step_counter: int = 0 self._outer_step_counter: int = 0
self._active_single_forward_manager_index: Optional[int] = None self._active_single_forward_manager_index: Optional[int] = None
@@ -183,7 +183,7 @@ class CanaryManager:
self._active_single_forward_manager_index = None self._active_single_forward_manager_index = None
def pre_ops_maybe_inside_graph( def pre_ops_maybe_inside_graph(
self, forward_batch: "ForwardBatch" self, forward_batch: ForwardBatch
) -> _PreOpsMaybeInsideGraphOutput: ) -> _PreOpsMaybeInsideGraphOutput:
assert self._active_single_forward_manager_index is not None, ( assert self._active_single_forward_manager_index is not None, (
"kv-canary: pre_ops_maybe_inside_graph called without active SingleForwardManager; " "kv-canary: pre_ops_maybe_inside_graph called without active SingleForwardManager; "
@@ -194,7 +194,7 @@ class CanaryManager:
def post_ops_maybe_inside_graph( def post_ops_maybe_inside_graph(
self, self,
forward_batch: "ForwardBatch", forward_batch: ForwardBatch,
pre_ops_output: _PreOpsMaybeInsideGraphOutput, pre_ops_output: _PreOpsMaybeInsideGraphOutput,
) -> None: ) -> None:
assert self._active_single_forward_manager_index is not None, ( assert self._active_single_forward_manager_index is not None, (
@@ -209,7 +209,7 @@ class CanaryManager:
self, self,
*, *,
single_forward_indices: Sequence[int], single_forward_indices: Sequence[int],
maybe_inaccurate_forward_batch: "ForwardBatch", maybe_inaccurate_forward_batch: ForwardBatch,
) -> Iterator[None]: ) -> Iterator[None]:
self._pre_ops_outside_graph( self._pre_ops_outside_graph(
single_forward_indices=single_forward_indices, single_forward_indices=single_forward_indices,
@@ -227,7 +227,7 @@ class CanaryManager:
self, self,
*, *,
single_forward_indices: Sequence[int], single_forward_indices: Sequence[int],
maybe_inaccurate_forward_batch: "ForwardBatch", maybe_inaccurate_forward_batch: ForwardBatch,
) -> None: ) -> None:
for idx in single_forward_indices: for idx in single_forward_indices:
self._single_forward_managers[idx].pre_ops_outside_graph( self._single_forward_managers[idx].pre_ops_outside_graph(
@@ -241,7 +241,7 @@ class CanaryManager:
self, self,
*, *,
single_forward_indices: Sequence[int], single_forward_indices: Sequence[int],
maybe_inaccurate_forward_batch: "ForwardBatch", maybe_inaccurate_forward_batch: ForwardBatch,
) -> None: ) -> None:
for idx in single_forward_indices: for idx in single_forward_indices:
self._single_forward_managers[idx].post_ops_outside_graph() self._single_forward_managers[idx].post_ops_outside_graph()
@@ -264,7 +264,7 @@ class CanaryManager:
single_forward_manager.phase_checker.enable_assert() single_forward_manager.phase_checker.enable_assert()
self._device_state.enable_chain_position_assert.fill_(1) self._device_state.enable_chain_position_assert.fill_(1)
def attach_radix_cache(self, radix_cache: "BasePrefixCache") -> None: def attach_radix_cache(self, radix_cache: BasePrefixCache) -> None:
self._sweep_orchestrator.attach_radix_cache(radix_cache) self._sweep_orchestrator.attach_radix_cache(radix_cache)
self._perturb_manager.attach_radix_cache(radix_cache) self._perturb_manager.attach_radix_cache(radix_cache)
python/sglang/srt/kv_canary/runner/future_tensor.py
+1 -1
View File
@@ -22,7 +22,7 @@ class FutureTensors:
@classmethod @classmethod
def device_to_host( def device_to_host(
cls, xs_device: _TensorOrDict, *, d2h_stream: torch.cuda.Stream cls, xs_device: _TensorOrDict, *, d2h_stream: torch.cuda.Stream
) -> "FutureTensors": ) -> FutureTensors:
assert not torch.cuda.is_current_stream_capturing(), ( assert not torch.cuda.is_current_stream_capturing(), (
"FutureTensors.device_to_host must not be called during cuda-graph " "FutureTensors.device_to_host must not be called during cuda-graph "
"capture: the d2h side-stream copy + pinned-host alloc cannot be " "capture: the d2h side-stream copy + pinned-host alloc cannot be "
python/sglang/srt/kv_canary/runner/kernel_launcher.py
+1 -1
View File
@@ -64,7 +64,7 @@ def launch_endpoints_per_forward(
tag_filter: Callable[[CanaryLaunchTag], bool], tag_filter: Callable[[CanaryLaunchTag], bool],
verify_plan: VerifyPlan, verify_plan: VerifyPlan,
write_plan: WritePlan, write_plan: WritePlan,
forward_batch: "ForwardBatch", forward_batch: ForwardBatch,
expected_inputs: ExpectedInputs, expected_inputs: ExpectedInputs,
violation_log: ViolationLog, violation_log: ViolationLog,
real_kv_hash_mode: RealKvHashMode, real_kv_hash_mode: RealKvHashMode,
python/sglang/srt/kv_canary/runner/swa_divergence.py
+13 -13
View File
@@ -32,8 +32,8 @@ class SwaDivergenceReporter:
device: torch.device, device: torch.device,
d2h_stream: torch.cuda.Stream, d2h_stream: torch.cuda.Stream,
interval: int, interval: int,
swa_allocator: Optional["SWATokenToKVPoolAllocator"] = None, swa_allocator: Optional[SWATokenToKVPoolAllocator] = None,
req_to_token_pool: Optional["ReqToTokenPool"] = None, req_to_token_pool: Optional[ReqToTokenPool] = None,
) -> None: ) -> None:
self._interval = interval self._interval = interval
self._swa_allocator = swa_allocator self._swa_allocator = swa_allocator
@@ -57,7 +57,7 @@ class SwaDivergenceReporter:
self, self,
*, *,
outer_step_counter: int, outer_step_counter: int,
maybe_inaccurate_forward_batch: Optional["ForwardBatch"], maybe_inaccurate_forward_batch: Optional[ForwardBatch],
) -> None: ) -> None:
self._forward_ct += 1 self._forward_ct += 1
self._handler.step( self._handler.step(
@@ -72,7 +72,7 @@ class SwaDivergenceReporter:
self, self,
*, *,
outer_step_counter: int, outer_step_counter: int,
maybe_inaccurate_forward_batch: Optional["ForwardBatch"], maybe_inaccurate_forward_batch: Optional[ForwardBatch],
) -> Optional[dict[str, Any]]: ) -> Optional[dict[str, Any]]:
if outer_step_counter == 0 or outer_step_counter % self._interval != 0: if outer_step_counter == 0 or outer_step_counter % self._interval != 0:
return None return None
@@ -134,14 +134,14 @@ class SwaDivergenceLog:
) )
@classmethod @classmethod
def parse(cls, line: str) -> Optional["SwaDivergenceLog"]: def parse(cls, line: str) -> Optional[SwaDivergenceLog]:
match = _SWA_DIVERGENCE_LINE_RE.search(line) match = _SWA_DIVERGENCE_LINE_RE.search(line)
if match is None: if match is None:
return None return None
return cls(**json.loads(match.group(1))) return cls(**json.loads(match.group(1)))
@classmethod @classmethod
def find_last(cls, text: str) -> Optional[tuple["SwaDivergenceLog", str]]: def find_last(cls, text: str) -> Optional[tuple[SwaDivergenceLog, str]]:
last_match: Optional[re.Match] = None last_match: Optional[re.Match] = None
for match in _SWA_DIVERGENCE_LINE_RE.finditer(text): for match in _SWA_DIVERGENCE_LINE_RE.finditer(text):
last_match = match last_match = match
@@ -150,7 +150,7 @@ class SwaDivergenceLog:
return cls(**json.loads(last_match.group(1))), last_match.group(0) return cls(**json.loads(last_match.group(1))), last_match.group(0)
@classmethod @classmethod
def find_all(cls, text: str) -> list[tuple["SwaDivergenceLog", str]]: def find_all(cls, text: str) -> list[tuple[SwaDivergenceLog, str]]:
return [ return [
(cls(**json.loads(match.group(1))), match.group(0)) (cls(**json.loads(match.group(1))), match.group(0))
for match in _SWA_DIVERGENCE_LINE_RE.finditer(text) for match in _SWA_DIVERGENCE_LINE_RE.finditer(text)
@@ -159,9 +159,9 @@ class SwaDivergenceLog:
def compute_swa_out_of_window_tokens( def compute_swa_out_of_window_tokens(
*, *,
swa_allocator: "SWATokenToKVPoolAllocator", swa_allocator: SWATokenToKVPoolAllocator,
req_to_token_pool: "ReqToTokenPool", req_to_token_pool: ReqToTokenPool,
maybe_inaccurate_forward_batch: "ForwardBatch", maybe_inaccurate_forward_batch: ForwardBatch,
) -> torch.Tensor: ) -> torch.Tensor:
"""Count tokens in the live req_to_token range whose SWA mapping is 0 (out-of-window).""" """Count tokens in the live req_to_token range whose SWA mapping is 0 (out-of-window)."""
full_to_swa_index_mapping = swa_allocator.full_to_swa_index_mapping full_to_swa_index_mapping = swa_allocator.full_to_swa_index_mapping
@@ -180,9 +180,9 @@ def compute_swa_out_of_window_tokens(
def compute_swa_full_idx_divergence( def compute_swa_full_idx_divergence(
*, *,
swa_allocator: "SWATokenToKVPoolAllocator", swa_allocator: SWATokenToKVPoolAllocator,
req_to_token_pool: "ReqToTokenPool", req_to_token_pool: ReqToTokenPool,
maybe_inaccurate_forward_batch: "ForwardBatch", maybe_inaccurate_forward_batch: ForwardBatch,
) -> torch.Tensor: ) -> torch.Tensor:
"""Count non-identity (full, swa) index pairs in the live req_to_token range.""" """Count non-identity (full, swa) index pairs in the live req_to_token range."""
full_to_swa_index_mapping = swa_allocator.full_to_swa_index_mapping full_to_swa_index_mapping = swa_allocator.full_to_swa_index_mapping
python/sglang/srt/kv_canary/runner/sweep.py
+2 -2
View File
@@ -34,7 +34,7 @@ class SweepOrchestrator:
self._endpoints = endpoints self._endpoints = endpoints
self._swa_window_size = swa_window_size self._swa_window_size = swa_window_size
self._outer_step_counter_getter = outer_step_counter_getter self._outer_step_counter_getter = outer_step_counter_getter
self._radix_cache: Optional["BasePrefixCache"] = None self._radix_cache: Optional[BasePrefixCache] = None
self._last_sweep_step: int = -1 self._last_sweep_step: int = -1
self._sweep_passes: int = 0 self._sweep_passes: int = 0
@@ -43,7 +43,7 @@ class SweepOrchestrator:
def sweep_passes(self) -> int: def sweep_passes(self) -> int:
return self._sweep_passes return self._sweep_passes
def attach_radix_cache(self, radix_cache: "BasePrefixCache") -> None: def attach_radix_cache(self, radix_cache: BasePrefixCache) -> None:
self._radix_cache = radix_cache self._radix_cache = radix_cache
def maybe_run_sweep(self) -> None: def maybe_run_sweep(self) -> None:
python/sglang/srt/kv_canary/single_forward_manager/data.py
+1 -1
View File
@@ -21,7 +21,7 @@ class PostOpsInsideGraphOutputBuffer:
num_slot_tags: int, num_slot_tags: int,
swa_verify_total_count_shape: tuple[int, ...] | None, swa_verify_total_count_shape: tuple[int, ...] | None,
device: torch.device, device: torch.device,
) -> "PostOpsInsideGraphOutputBuffer": ) -> PostOpsInsideGraphOutputBuffer:
return cls( return cls(
verify_plan_enable=torch.zeros(1, dtype=torch.int32, device=device), verify_plan_enable=torch.zeros(1, dtype=torch.int32, device=device),
kernel_run_counters=torch.zeros( kernel_run_counters=torch.zeros(
python/sglang/srt/kv_canary/single_forward_manager/manager.py
+7 -7
View File
@@ -66,7 +66,7 @@ class SingleForwardManager:
device_state: CanaryDeviceState, device_state: CanaryDeviceState,
buffer_groups: tuple[CanaryBufferGroup, ...], buffer_groups: tuple[CanaryBufferGroup, ...],
endpoints: tuple[CanaryEndpoint, ...], endpoints: tuple[CanaryEndpoint, ...],
req_to_token_pool: "ReqToTokenPool", req_to_token_pool: ReqToTokenPool,
swa_window_size: int, swa_window_size: int,
per_forward_verify_capacity: int, per_forward_verify_capacity: int,
per_forward_write_req_capacity: int, per_forward_write_req_capacity: int,
@@ -119,7 +119,7 @@ class SingleForwardManager:
return self._phase_checker return self._phase_checker
def pre_ops_outside_graph( def pre_ops_outside_graph(
self, *, maybe_inaccurate_forward_batch: "ForwardBatch" self, *, maybe_inaccurate_forward_batch: ForwardBatch
) -> None: ) -> None:
self._phase_checker.update( self._phase_checker.update(
expect_phase=_SingleForwardPhase.IDLE, expect_phase=_SingleForwardPhase.IDLE,
@@ -150,8 +150,8 @@ class SingleForwardManager:
) )
def pre_ops_maybe_inside_graph( def pre_ops_maybe_inside_graph(
self, forward_batch: "ForwardBatch" self, forward_batch: ForwardBatch
) -> "_PreOpsMaybeInsideGraphOutput": ) -> _PreOpsMaybeInsideGraphOutput:
self._phase_checker.update( self._phase_checker.update(
expect_phase=_SingleForwardPhase.AFTER_PRE_OUT, expect_phase=_SingleForwardPhase.AFTER_PRE_OUT,
next_phase=_SingleForwardPhase.AFTER_PRE_MAYBE_IN, next_phase=_SingleForwardPhase.AFTER_PRE_MAYBE_IN,
@@ -238,8 +238,8 @@ class SingleForwardManager:
def post_ops_maybe_inside_graph( def post_ops_maybe_inside_graph(
self, self,
forward_batch: "ForwardBatch", forward_batch: ForwardBatch,
pre_ops_output: "_PreOpsMaybeInsideGraphOutput", pre_ops_output: _PreOpsMaybeInsideGraphOutput,
) -> None: ) -> None:
self._phase_checker.update( self._phase_checker.update(
expect_phase=_SingleForwardPhase.AFTER_PRE_MAYBE_IN, expect_phase=_SingleForwardPhase.AFTER_PRE_MAYBE_IN,
@@ -293,7 +293,7 @@ class SingleForwardManager:
self._enable_warner.tick(self._output_buffer.verify_plan_enable) self._enable_warner.tick(self._output_buffer.verify_plan_enable)
def _should_enable_write_input_assert_for_launch( def _should_enable_write_input_assert_for_launch(
self, forward_batch: "ForwardBatch" self, forward_batch: ForwardBatch
) -> bool: ) -> bool:
if not self._config.enable_write_input_assert: if not self._config.enable_write_input_assert:
return False return False
python/sglang/srt/kv_canary/state.py
+2 -2
View File
@@ -41,7 +41,7 @@ class ViolationLog:
violation_write_index: torch.Tensor violation_write_index: torch.Tensor
@classmethod @classmethod
def allocate(cls, *, ring_capacity: int, device: torch.device) -> "ViolationLog": def allocate(cls, *, ring_capacity: int, device: torch.device) -> ViolationLog:
if ring_capacity <= 0: if ring_capacity <= 0:
raise ValueError( raise ValueError(
f"kv-canary: ViolationLog ring_capacity must be positive, got {ring_capacity}" f"kv-canary: ViolationLog ring_capacity must be positive, got {ring_capacity}"
@@ -102,7 +102,7 @@ class CanaryDeviceState:
num_tags: int, num_tags: int,
req_to_token_alloc_size: Optional[int] = None, req_to_token_alloc_size: Optional[int] = None,
max_context_len: Optional[int] = None, max_context_len: Optional[int] = None,
) -> "CanaryDeviceState": ) -> CanaryDeviceState:
if num_tags <= 0: if num_tags <= 0:
raise ValueError( raise ValueError(
f"kv-canary: CanaryDeviceState num_tags must be positive, got {num_tags}" f"kv-canary: CanaryDeviceState num_tags must be positive, got {num_tags}"
python/sglang/srt/kv_canary/sweep_plan_builder.py
+1 -1
View File
@@ -13,7 +13,7 @@ if TYPE_CHECKING:
def build_verify_plan_radix_sweep( def build_verify_plan_radix_sweep(
*, *,
radix_cache: "BasePrefixCache", radix_cache: BasePrefixCache,
swa_window_size: int, swa_window_size: int,
full_to_swa_index_mapping: Optional[torch.Tensor], full_to_swa_index_mapping: Optional[torch.Tensor],
unlocked_only: bool = False, unlocked_only: bool = False,
python/sglang/srt/kv_canary/token_oracle/install.py
+1 -1
View File
@@ -11,7 +11,7 @@ if TYPE_CHECKING:
def install_token_oracle_from_env( def install_token_oracle_from_env(
*, server_args: "ServerArgs", vocab_size: int *, server_args: ServerArgs, vocab_size: int
) -> Optional[TokenOracleManager]: ) -> Optional[TokenOracleManager]:
# Must be called before create_sampler() so the factory is present when the # Must be called before create_sampler() so the factory is present when the
# Sampler is first constructed. # Sampler is first constructed.
python/sglang/srt/kv_canary/token_oracle/oracle_manager.py
+2 -2
View File
@@ -18,7 +18,7 @@ class TokenOracleManager:
def fill_expected_inputs( def fill_expected_inputs(
self, self,
*, *,
forward_batch: "ForwardBatch", forward_batch: ForwardBatch,
expected_inputs_out: ExpectedInputs, expected_inputs_out: ExpectedInputs,
) -> None: ) -> None:
positions = forward_batch.positions positions = forward_batch.positions
@@ -57,7 +57,7 @@ class TokenOracleManager:
def _build_generalized_req_id_per_token( def _build_generalized_req_id_per_token(
*, *,
forward_batch: "ForwardBatch", forward_batch: ForwardBatch,
num_tokens: int, num_tokens: int,
generalized_req_ids_per_row: torch.Tensor, generalized_req_ids_per_row: torch.Tensor,
) -> torch.Tensor: ) -> torch.Tensor:
python/sglang/srt/kv_canary/token_oracle/sampler.py
+2 -2
View File
@@ -33,8 +33,8 @@ class _OracleSampler(Sampler):
def forward( def forward(
self, self,
logits_output: "LogitsProcessorOutput", logits_output: LogitsProcessorOutput,
sampling_info: "SamplingBatchInfo", sampling_info: SamplingBatchInfo,
return_logprob: bool, return_logprob: bool,
top_logprobs_nums: List[int], top_logprobs_nums: List[int],
token_ids_logprobs: List[List[int]], token_ids_logprobs: List[List[int]],
python/sglang/srt/layers/attention/aiter_utils.py
+6 -6
View File
@@ -43,12 +43,12 @@ if TYPE_CHECKING:
def forward_extend_vectorized_5d( def forward_extend_vectorized_5d(
backend: "AiterAttnBackend", backend: AiterAttnBackend,
q: torch.Tensor, q: torch.Tensor,
k: torch.Tensor, k: torch.Tensor,
v: torch.Tensor, v: torch.Tensor,
layer: "RadixAttention", layer: RadixAttention,
forward_batch: "ForwardBatch", forward_batch: ForwardBatch,
bs0: int, bs0: int,
window_size, window_size,
sinks, sinks,
@@ -207,10 +207,10 @@ def forward_extend_vectorized_5d(
def forward_decode_vectorized_5d( def forward_decode_vectorized_5d(
backend: "AiterAttnBackend", backend: AiterAttnBackend,
q: torch.Tensor, q: torch.Tensor,
layer: "RadixAttention", layer: RadixAttention,
forward_batch: "ForwardBatch", forward_batch: ForwardBatch,
k_cache: torch.Tensor, k_cache: torch.Tensor,
v_cache: torch.Tensor, v_cache: torch.Tensor,
o: torch.Tensor, o: torch.Tensor,
python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py
+3 -3
View File
@@ -996,7 +996,7 @@ class DeepseekV4HipRadixBackend(
self.forward_metadata = current_raw self.forward_metadata = current_raw
def _attach_unified_kv_decode_streams( def _attach_unified_kv_decode_streams(
self, core: "DSV4AttnMetadata", req_pool_indices: torch.Tensor self, core: DSV4AttnMetadata, req_pool_indices: torch.Tensor
) -> None: ) -> None:
"""build the ragged decode index streams once per forward""" """build the ragged decode index streams once per forward"""
from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import ( from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import (
@@ -1031,7 +1031,7 @@ class DeepseekV4HipRadixBackend(
def _attach_unified_kv_prefill_meta( def _attach_unified_kv_prefill_meta(
self, self,
core: "DSV4AttnMetadata", core: DSV4AttnMetadata,
req_pool_indices: torch.Tensor, req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor, seq_lens: torch.Tensor,
extend_seq_lens: torch.Tensor, extend_seq_lens: torch.Tensor,
@@ -1065,7 +1065,7 @@ class DeepseekV4HipRadixBackend(
forward_batch: ForwardBatch, forward_batch: ForwardBatch,
compress_ratio: Literal[0, 4, 128], compress_ratio: Literal[0, 4, 128],
attn_sink: torch.Tensor, attn_sink: torch.Tensor,
core_attn_metadata: "DSV4AttnMetadata", core_attn_metadata: DSV4AttnMetadata,
save_kv_cache: bool = True, save_kv_cache: bool = True,
) -> torch.Tensor: ) -> torch.Tensor:
"""unified_kv paged-attention path over the bf16 unified_kv""" """unified_kv paged-attention path over the bf16 unified_kv"""
python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py
+3 -3
View File
@@ -71,8 +71,8 @@ class DeepseekSparseAttnBackendMTPPrecomputeMixin:
req_pool_indices: torch.Tensor, req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor, seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor, seq_lens_cpu: torch.Tensor,
forward_mode: "ForwardMode", forward_mode: ForwardMode,
spec_info: Optional["SpecInput"], spec_info: Optional[SpecInput],
) -> PrecomputedMetadata: ) -> PrecomputedMetadata:
"""Precompute all shared metadata for multi-step backends. """Precompute all shared metadata for multi-step backends.
@@ -252,7 +252,7 @@ class DeepseekSparseAttnBackendMTPPrecomputeMixin:
req_pool_indices: torch.Tensor, req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor, seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor, seq_lens_cpu: torch.Tensor,
spec_info: "SpecInput", spec_info: SpecInput,
) -> PrecomputedMetadata: ) -> PrecomputedMetadata:
"""Precompute metadata for draft extend mode.""" """Precompute metadata for draft extend mode."""
max_seqlen_k = int(seq_lens_cpu.max().item()) max_seqlen_k = int(seq_lens_cpu.max().item())
python/sglang/srt/layers/attention/dsa_backend.py
+2 -2
View File
@@ -116,7 +116,7 @@ class DSAFlashMLAMetadata:
num_splits=self.num_splits[sli], num_splits=self.num_splits[sli],
) )
def copy_(self, other: "DSAFlashMLAMetadata"): def copy_(self, other: DSAFlashMLAMetadata):
self.flashmla_metadata.copy_(other.flashmla_metadata) self.flashmla_metadata.copy_(other.flashmla_metadata)
self.num_splits.copy_(other.num_splits) self.num_splits.copy_(other.num_splits)
@@ -866,7 +866,7 @@ class DeepseekSparseAttnBackend(
forward_mode: ForwardMode, forward_mode: ForwardMode,
spec_info: Optional[SpecInput], spec_info: Optional[SpecInput],
out_cache_loc: Optional[torch.Tensor] = None, out_cache_loc: Optional[torch.Tensor] = None,
actual_forward_mode: Optional["ForwardMode"] = None, actual_forward_mode: Optional[ForwardMode] = None,
): ):
"""Create and store DSAMetadata for a new batch size during CUDA graph capture.""" """Create and store DSAMetadata for a new batch size during CUDA graph capture."""
self.set_dsa_prefill_impl(forward_batch=None) self.set_dsa_prefill_impl(forward_batch=None)
python/sglang/srt/layers/attention/dsv4/metadata.py
+1 -1
View File
@@ -153,7 +153,7 @@ class PagedIndexerMetadata:
def max_c4_seq_len(self) -> int: def max_c4_seq_len(self) -> int:
return self.page_table.shape[1] * self.c4_page_size return self.page_table.shape[1] * self.c4_page_size
def copy_(self, other: "PagedIndexerMetadata"): def copy_(self, other: PagedIndexerMetadata):
if is_hip(): if is_hip():
copy_fields = ["page_table", "c4_seq_lens"] copy_fields = ["page_table", "c4_seq_lens"]
assign_fields = ["deep_gemm_metadata"] assign_fields = ["deep_gemm_metadata"]
python/sglang/srt/layers/attention/tokenspeed_mla_backend.py
+5 -5
View File
@@ -91,7 +91,7 @@ class TokenspeedMLABackend(TRTLLMMLABackend):
def __init__( def __init__(
self, self,
model_runner: "ModelRunner", model_runner: ModelRunner,
skip_prefill: bool = False, skip_prefill: bool = False,
kv_indptr_buf: Optional[torch.Tensor] = None, kv_indptr_buf: Optional[torch.Tensor] = None,
q_indptr_decode_buf: Optional[torch.Tensor] = None, q_indptr_decode_buf: Optional[torch.Tensor] = None,
@@ -221,7 +221,7 @@ class TokenspeedMLABackend(TRTLLMMLABackend):
kv_a: torch.Tensor, kv_a: torch.Tensor,
k_pe: torch.Tensor, k_pe: torch.Tensor,
positions: torch.Tensor, positions: torch.Tensor,
layer: "DeepseekV2AttentionMLA", layer: DeepseekV2AttentionMLA,
forward_batch: ForwardBatch, forward_batch: ForwardBatch,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Build FP8 (Q, K, V) for the FMHA kernel and write FP8 KV cache.""" """Build FP8 (Q, K, V) for the FMHA kernel and write FP8 KV cache."""
@@ -278,7 +278,7 @@ class TokenspeedMLABackend(TRTLLMMLABackend):
block_tables: torch.Tensor, block_tables: torch.Tensor,
seq_lens: torch.Tensor, seq_lens: torch.Tensor,
max_seq_len: int, max_seq_len: int,
layer: "RadixAttention", layer: RadixAttention,
) -> torch.Tensor: ) -> torch.Tensor:
k_scale = getattr(layer, "k_scale_float", None) k_scale = getattr(layer, "k_scale_float", None)
if k_scale is None: if k_scale is None:
@@ -308,7 +308,7 @@ class TokenspeedMLABackend(TRTLLMMLABackend):
q: torch.Tensor, q: torch.Tensor,
k: torch.Tensor, k: torch.Tensor,
v: torch.Tensor, v: torch.Tensor,
layer: "RadixAttention", layer: RadixAttention,
batch_size: int, batch_size: int,
cum_seq_lens_q: torch.Tensor, cum_seq_lens_q: torch.Tensor,
max_q_len: int, max_q_len: int,
@@ -342,7 +342,7 @@ class TokenspeedMLAMultiStepDraftBackend(TRTLLMMLAMultiStepDraftBackend):
"""Multi-step draft backend for tokenspeed_mla used by EAGLE.""" """Multi-step draft backend for tokenspeed_mla used by EAGLE."""
def __init__( def __init__(
self, model_runner: "ModelRunner", topk: int, speculative_num_steps: int self, model_runner: ModelRunner, topk: int, speculative_num_steps: int
): ):
super().__init__(model_runner, topk, speculative_num_steps) super().__init__(model_runner, topk, speculative_num_steps)
# Parent populates self.attn_backends with TRT-LLM instances; replace # Parent populates self.attn_backends with TRT-LLM instances; replace
python/sglang/srt/layers/attention/trtllm_mha_backend.py
+1 -1
View File
@@ -341,7 +341,7 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
forward_mode: ForwardMode, forward_mode: ForwardMode,
spec_info, spec_info,
device: torch.device, device: torch.device,
) -> "TRTLLMMHAMetadata": ) -> TRTLLMMHAMetadata:
"""Create TRTLLMMHAMetadata with pre-allocated buffer slice refs, stored in the dict.""" """Create TRTLLMMHAMetadata with pre-allocated buffer slice refs, stored in the dict."""
metadata = TRTLLMMHAMetadata() metadata = TRTLLMMHAMetadata()
python/sglang/srt/layers/attention/trtllm_mla_backend.py
+1 -1
View File
@@ -1073,7 +1073,7 @@ class TRTLLMMLAMultiStepDraftBackend(FlashInferMLAMultiStepDraftBackend):
def __init__( def __init__(
self, self,
model_runner: "ModelRunner", model_runner: ModelRunner,
topk: int, topk: int,
speculative_num_steps: int, speculative_num_steps: int,
backend: str = "trtllm-gen", backend: str = "trtllm-gen",
python/sglang/srt/layers/moe/mega_moe.py
+3 -3
View File
@@ -94,7 +94,7 @@ def _get_mega_moe_symm_buffer(
return buf return buf
def should_use_mega_moe(moe: "DeepseekV2MoE", hidden_states: torch.Tensor) -> bool: def should_use_mega_moe(moe: DeepseekV2MoE, hidden_states: torch.Tensor) -> bool:
if not get_moe_a2a_backend().is_megamoe(): if not get_moe_a2a_backend().is_megamoe():
return False return False
if not getattr(moe.experts, "_mega_moe_weights_built", False): if not getattr(moe.experts, "_mega_moe_weights_built", False):
@@ -112,7 +112,7 @@ def should_use_mega_moe(moe: "DeepseekV2MoE", hidden_states: torch.Tensor) -> bo
def forward_mega_moe( def forward_mega_moe(
moe: "DeepseekV2MoE", moe: DeepseekV2MoE,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None, forward_batch: Optional[ForwardBatch] = None,
input_ids_global: Optional[torch.Tensor] = None, input_ids_global: Optional[torch.Tensor] = None,
@@ -149,7 +149,7 @@ def forward_mega_moe(
def _run_mega_routed( def _run_mega_routed(
moe: "DeepseekV2MoE", moe: DeepseekV2MoE,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch], forward_batch: Optional[ForwardBatch],
input_ids_global: Optional[torch.Tensor], input_ids_global: Optional[torch.Tensor],
python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py
+1 -1
View File
@@ -342,7 +342,7 @@ class CuteDslFp4MoeQuantInfo(MoeQuantInfo):
use_nvfp4_dispatch: bool = False use_nvfp4_dispatch: bool = False
# v1 only: SBO down-GEMM overlap args. # v1 only: SBO down-GEMM overlap args.
down_gemm_overlap_args: Optional["DownGemmOverlapArgs"] = None down_gemm_overlap_args: Optional[DownGemmOverlapArgs] = None
@register_fused_func("none", "flashinfer_cutedsl") @register_fused_func("none", "flashinfer_cutedsl")
python/sglang/srt/layers/moe/moe_runner/flashinfer_mxfp4.py
+2 -2
View File
@@ -93,10 +93,10 @@ def _flashinfer_cutlass_fused_moe():
@register_fused_func("none", "flashinfer_mxfp4") @register_fused_func("none", "flashinfer_mxfp4")
def fused_experts_none_to_flashinfer_mxfp4( def fused_experts_none_to_flashinfer_mxfp4(
dispatch_output: "StandardDispatchOutput", dispatch_output: StandardDispatchOutput,
quant_info: MoeQuantInfo, quant_info: MoeQuantInfo,
runner_config: MoeRunnerConfig, runner_config: MoeRunnerConfig,
) -> "StandardCombineInput": ) -> StandardCombineInput:
"""SM90 W4A16 MXFP4 fused expert forward pass. """SM90 W4A16 MXFP4 fused expert forward pass.
Mirrors the legacy ``Mxfp4MoEMethod._apply_sm90_cutlass`` and DSv4's Mirrors the legacy ``Mxfp4MoEMethod._apply_sm90_cutlass`` and DSv4's
python/sglang/srt/layers/moe/moe_runner/triton_kernels.py
+4 -4
View File
@@ -42,9 +42,9 @@ class TritonKernelsRunnerInput(RunnerInput):
"""Input bundle passed to the triton-kernels runner core.""" """Input bundle passed to the triton-kernels runner core."""
hidden_states: torch.Tensor hidden_states: torch.Tensor
routing_data: "RoutingData" routing_data: RoutingData
gather_indx: "GatherIndx" gather_indx: GatherIndx
scatter_indx: "ScatterIndx" scatter_indx: ScatterIndx
@property @property
def runner_backend(self) -> MoeRunnerBackend: def runner_backend(self) -> MoeRunnerBackend:
@@ -158,7 +158,7 @@ class TritonKernelsRunnerCore(MoeRunnerCore):
@register_pre_permute("standard", "triton_kernel") @register_pre_permute("standard", "triton_kernel")
def pre_permute_standard_to_triton_kernels( def pre_permute_standard_to_triton_kernels(
dispatch_output: "StandardDispatchOutput", dispatch_output: StandardDispatchOutput,
quant_info: TritonKernelsQuantInfo, quant_info: TritonKernelsQuantInfo,
runner_config: MoeRunnerConfig, runner_config: MoeRunnerConfig,
running_state: dict, running_state: dict,
python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe_triton_kernels.py
+1 -1
View File
@@ -673,7 +673,7 @@ def _set_triton_tma_allocator():
# --- B TensorDescriptor cache (LRU) --- # --- B TensorDescriptor cache (LRU) ---
_B_DESC_CACHE_MAX = 64 _B_DESC_CACHE_MAX = 64
_B_DESC_CACHE: "OrderedDict[tuple, TensorDescriptor]" = OrderedDict() _B_DESC_CACHE: OrderedDict[tuple, TensorDescriptor] = OrderedDict()
def _get_b_tma_desc_cached(B: torch.Tensor, block_n: int, block_k: int): def _get_b_tma_desc_cached(B: torch.Tensor, block_n: int, block_k: int):
python/sglang/srt/layers/moe/topk.py
+1 -1
View File
@@ -314,7 +314,7 @@ class BypassedTopKOutput(NamedTuple):
def format(self) -> TopKOutputFormat: def format(self) -> TopKOutputFormat:
return TopKOutputFormat.BYPASSED return TopKOutputFormat.BYPASSED
def to_standard(self, layer_id: Optional[int] = None) -> "StandardTopKOutput": def to_standard(self, layer_id: Optional[int] = None) -> StandardTopKOutput:
"""Materialize routing tensors. Used by MoE kernels that need explicit """Materialize routing tensors. Used by MoE kernels that need explicit
topk_ids / topk_weights rather than doing routing internally.""" topk_ids / topk_weights rather than doing routing internally."""
return select_experts( return select_experts(
python/sglang/srt/layers/pooler.py
+1 -1
View File
@@ -107,7 +107,7 @@ def pool_at_delimiter_positions(
def score_and_pool( def score_and_pool(
score_head: nn.Module, score_head: nn.Module,
pooler: "Pooler", pooler: Pooler,
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
forward_batch: ForwardBatch, forward_batch: ForwardBatch,
input_ids: torch.Tensor, input_ids: torch.Tensor,
python/sglang/srt/layers/quantization/awq/schemes/awq_cpu.py
+2 -2
View File
@@ -23,14 +23,14 @@ __all__ = ["AWQIntelAMXLinearScheme", "AWQIntelAMXMoEScheme"]
class AWQIntelAMXLinearScheme(AWQLinearScheme): class AWQIntelAMXLinearScheme(AWQLinearScheme):
"""Linear scheme for AWQ on Intel CPU with AMX.""" """Linear scheme for AWQ on Intel CPU with AMX."""
def _init_kernel(self, quant_config: "AWQConfig"): def _init_kernel(self, quant_config: AWQConfig):
return AWQIntelAMXLinearKernel(quant_config) return AWQIntelAMXLinearKernel(quant_config)
class AWQIntelAMXMoEScheme(AWQMoEScheme): class AWQIntelAMXMoEScheme(AWQMoEScheme):
"""MoE scheme for AWQ on Intel CPU with AMX.""" """MoE scheme for AWQ on Intel CPU with AMX."""
def _init_kernel(self, quant_config: "AWQConfig"): def _init_kernel(self, quant_config: AWQConfig):
return AWQIntelAMXMoEKernel(quant_config) return AWQIntelAMXMoEKernel(quant_config)
def create_moe_runner( def create_moe_runner(
python/sglang/srt/layers/quantization/awq/schemes/awq_linear.py
+3 -3
View File
@@ -16,11 +16,11 @@ __all__ = ["AWQLinearScheme", "AWQAscendLinearScheme"]
class AWQLinearScheme(AWQLinearSchemeBase): class AWQLinearScheme(AWQLinearSchemeBase):
def __init__(self, quant_config: "AWQConfig"): def __init__(self, quant_config: AWQConfig):
self.quant_config = quant_config self.quant_config = quant_config
self.kernel = self._init_kernel(quant_config) self.kernel = self._init_kernel(quant_config)
def _init_kernel(self, quant_config: "AWQConfig"): def _init_kernel(self, quant_config: AWQConfig):
from sglang.srt.hardware_backend.gpu.quantization.awq_kernels import ( from sglang.srt.hardware_backend.gpu.quantization.awq_kernels import (
AWQLinearKernel, AWQLinearKernel,
) )
@@ -102,7 +102,7 @@ class AWQLinearScheme(AWQLinearSchemeBase):
class AWQAscendLinearScheme(AWQLinearScheme): class AWQAscendLinearScheme(AWQLinearScheme):
def _init_kernel(self, quant_config: "AWQConfig"): def _init_kernel(self, quant_config: AWQConfig):
from sglang.srt.hardware_backend.npu.quantization.awq_kernels import ( from sglang.srt.hardware_backend.npu.quantization.awq_kernels import (
AWQAscendLinearKernel, AWQAscendLinearKernel,
) )
python/sglang/srt/layers/quantization/awq/schemes/awq_marlin.py
+2 -2
View File
@@ -17,11 +17,11 @@ __all__ = ["AWQMarlinLinearScheme"]
class AWQMarlinLinearScheme(AWQLinearSchemeBase): class AWQMarlinLinearScheme(AWQLinearSchemeBase):
def __init__(self, quant_config: "AWQMarlinConfig"): def __init__(self, quant_config: AWQMarlinConfig):
self.quant_config = quant_config self.quant_config = quant_config
self.kernel = self._init_kernel(quant_config) self.kernel = self._init_kernel(quant_config)
def _init_kernel(self, quant_config: "AWQMarlinConfig"): def _init_kernel(self, quant_config: AWQMarlinConfig):
from sglang.srt.hardware_backend.gpu.quantization.awq_kernels import ( from sglang.srt.hardware_backend.gpu.quantization.awq_kernels import (
AWQMarlinLinearKernel, AWQMarlinLinearKernel,
) )

Some files were not shown because too many files have changed in this diff Show More