From 97a00317993e861fe91aa00d9074a8aad7be842c Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Thu, 11 Jun 2026 17:38:05 -0700 Subject: [PATCH] [lint] Enable Ruff UP037 to drop redundant quoted annotations (#27984) --- .pre-commit-config.yaml | 2 +- python/sglang/_mps_stub.py | 4 +-- .../kv_canary/plan/entries_kernel.py | 2 +- python/sglang/jit_kernel/kv_canary/verify.py | 6 ++-- python/sglang/jit_kernel/kv_canary/write.py | 6 ++-- .../tests/kv_canary/_canary_helpers.py | 2 +- .../configs/quantization/nunchaku.py | 4 +-- .../runtime/distributed/cfg_parallel_utils.py | 14 +++++----- .../runtime/distributed/cfg_policy.py | 10 +++---- .../runtime/layers/attention/layer.py | 2 +- .../runtime/layers/attention/turbo_layer.py | 2 +- .../layers/quantization/bitsandbytes.py | 2 +- .../layers/quantization/modelopt_fp8.py | 2 +- .../layers/quantization/modelopt_quant.py | 2 +- .../runtime/layers/quantization/mxfp4_npu.py | 2 +- .../runtime/layers/quantization/mxfp8_npu.py | 2 +- .../managers/dynamic_batch_admission.py | 4 +-- .../runtime/pipelines_core/schedule_batch.py | 8 +++--- .../runtime/pipelines_core/stages/dedup.py | 20 ++++++------- .../model_specific_stages/hunyuan3d/paint.py | 2 +- .../sana_wm/self_forcing.py | 2 +- .../runtime/server_args_auto_tune.py | 2 +- .../test/server/test_disagg_server.py | 2 +- .../test/server/test_server_common.py | 4 +-- python/sglang/srt/batch_overlap/operations.py | 2 +- .../srt/batch_overlap/two_batch_overlap.py | 2 +- .../sglang/srt/compilation/compile_phase.py | 4 +-- .../debug_utils/comparator/meta_overrider.py | 2 +- .../debug_utils/comparator/output_types.py | 2 +- python/sglang/srt/debug_utils/dumper.py | 2 +- .../disaggregation/common/staging_handler.py | 16 +++++------ python/sglang/srt/disaggregation/decode.py | 6 ++-- .../disaggregation/decode_hicache_mixin.py | 18 ++++++------ .../srt/disaggregation/encode_receiver.py | 2 +- .../srt/disaggregation/encode_server.py | 2 +- python/sglang/srt/disaggregation/mori/conn.py | 2 +- python/sglang/srt/disaggregation/nixl/conn.py | 10 +++---- .../sglang/srt/entrypoints/openai/protocol.py | 6 ++-- .../entrypoints/openai/serving_responses.py | 6 ++-- python/sglang/srt/eplb/expert_distribution.py | 6 ++-- python/sglang/srt/eplb/expert_location.py | 2 +- .../cpu/quantization/awq_kernels.py | 6 ++-- .../cpu/quantization/gptq_kernels.py | 6 ++-- .../gpu/quantization/awq_kernels.py | 10 +++---- .../gpu/quantization/gptq_kernels.py | 12 ++++---- python/sglang/srt/hardware_backend/mlx/aot.py | 2 +- .../mlx/kv_cache/attention_wrapper.py | 2 +- .../hardware_backend/mlx/kv_cache/layout.py | 2 +- .../hardware_backend/mlx/scheduler_mixin.py | 16 +++++------ .../musa/attention/flashattention_backend.py | 4 +-- .../npu/attention/ascend_backend.py | 4 +-- .../srt/hardware_backend/npu/moe/fuseep.py | 6 ++-- .../npu/quantization/awq_kernels.py | 6 ++-- .../npu/quantization/gptq_kernels.py | 10 +++---- python/sglang/srt/kv_canary/api.py | 10 +++---- python/sglang/srt/kv_canary/capacities.py | 4 +-- python/sglang/srt/kv_canary/config.py | 2 +- .../sglang/srt/kv_canary/expected_inputs.py | 4 +-- python/sglang/srt/kv_canary/perturb/config.py | 2 +- .../sglang/srt/kv_canary/perturb/manager.py | 18 ++++++------ .../srt/kv_canary/perturb/next_token_swap.py | 2 +- .../kv_canary/perturb/real_kv_post_forward.py | 2 +- .../kv_canary/perturb/real_kv_unused_cache.py | 6 ++-- .../srt/kv_canary/perturb/real_kv_used.py | 8 +++--- .../srt/kv_canary/perturb/req_to_token.py | 4 +-- .../srt/kv_canary/perturb/slot_picker.py | 6 ++-- python/sglang/srt/kv_canary/perturb/utils.py | 2 +- python/sglang/srt/kv_canary/plan_input.py | 6 ++-- .../srt/kv_canary/radix_cache_walker.py | 14 +++++----- .../req_to_expected_token_ids_manager.py | 4 +-- .../srt/kv_canary/runner/canary_manager.py | 18 ++++++------ .../srt/kv_canary/runner/future_tensor.py | 2 +- .../srt/kv_canary/runner/kernel_launcher.py | 2 +- .../srt/kv_canary/runner/swa_divergence.py | 26 ++++++++--------- python/sglang/srt/kv_canary/runner/sweep.py | 4 +-- .../kv_canary/single_forward_manager/data.py | 2 +- .../single_forward_manager/manager.py | 14 +++++----- python/sglang/srt/kv_canary/state.py | 4 +-- .../srt/kv_canary/sweep_plan_builder.py | 2 +- .../srt/kv_canary/token_oracle/install.py | 2 +- .../kv_canary/token_oracle/oracle_manager.py | 4 +-- .../srt/kv_canary/token_oracle/sampler.py | 4 +-- .../srt/layers/attention/aiter_utils.py | 12 ++++---- .../deepseek_v4_backend_hip_radix.py | 6 ++-- .../dsa/dsa_backend_mtp_precompute.py | 6 ++-- .../srt/layers/attention/dsa_backend.py | 4 +-- .../srt/layers/attention/dsv4/metadata.py | 2 +- .../attention/tokenspeed_mla_backend.py | 10 +++---- .../layers/attention/trtllm_mha_backend.py | 2 +- .../layers/attention/trtllm_mla_backend.py | 2 +- python/sglang/srt/layers/moe/mega_moe.py | 6 ++-- .../moe/moe_runner/flashinfer_cutedsl.py | 2 +- .../layers/moe/moe_runner/flashinfer_mxfp4.py | 4 +-- .../layers/moe/moe_runner/triton_kernels.py | 8 +++--- .../triton_utils/fused_moe_triton_kernels.py | 2 +- python/sglang/srt/layers/moe/topk.py | 2 +- python/sglang/srt/layers/pooler.py | 2 +- .../quantization/awq/schemes/awq_cpu.py | 4 +-- .../quantization/awq/schemes/awq_linear.py | 6 ++-- .../quantization/awq/schemes/awq_marlin.py | 4 +-- .../quantization/awq/schemes/awq_moe.py | 8 +++--- .../srt/layers/quantization/base_config.py | 6 ++-- .../srt/layers/quantization/bitsandbytes.py | 2 +- .../compressed_tensors/compressed_tensors.py | 2 +- .../schemes/compressed_tensors_wNa16_moe.py | 4 +-- python/sglang/srt/layers/quantization/fp8.py | 4 +-- python/sglang/srt/layers/quantization/gguf.py | 6 ++-- .../quantization/gptq/schemes/gptq_cpu.py | 8 +++--- .../quantization/gptq/schemes/gptq_linear.py | 6 ++-- .../quantization/gptq/schemes/gptq_marlin.py | 4 +-- .../quantization/gptq/schemes/gptq_moe.py | 8 +++--- .../srt/layers/quantization/marlin_utils.py | 2 +- python/sglang/srt/layers/quantization/mlx.py | 2 +- .../srt/layers/quantization/modelopt_quant.py | 8 +++--- .../schemes/modelslim_w4a4_int4_moe.py | 6 ++-- .../schemes/modelslim_w4a8_int8_moe.py | 6 ++-- .../schemes/modelslim_w8a8_int8_moe.py | 6 ++-- .../sglang/srt/layers/quantization/mxfp4.py | 2 +- .../mxfp4_flashinfer_cutlass_moe.py | 4 +-- .../sglang/srt/layers/quantization/unquant.py | 4 +-- .../srt/lora/deepseek_mla_correction.py | 10 +++---- .../deepseek_mla_correction.py | 10 +++---- .../lora/trtllm_lora_temp/lora_dispatch.py | 16 +++++------ .../srt/lora/trtllm_lora_temp/lora_layer.py | 2 +- .../sglang/srt/managers/cache_controller.py | 2 +- python/sglang/srt/managers/overlap_utils.py | 4 +-- python/sglang/srt/managers/schedule_batch.py | 10 +++---- .../batch_result_processor.py | 28 +++++++++---------- .../managers/scheduler_components/dp_attn.py | 2 +- .../scheduler_components/invariant_checker.py | 2 +- .../kv_events_publisher.py | 6 ++-- .../scheduler_components/load_inquirer.py | 14 +++++----- .../scheduler_components/metrics_reporter.py | 2 +- .../new_token_ratio_tracker.py | 2 +- .../pool_stats_observer.py | 6 ++-- .../scheduler_components/request_receiver.py | 10 +++---- python/sglang/srt/managers/tp_worker.py | 4 +-- .../sglang/srt/mem_cache/base_prefix_cache.py | 4 +-- python/sglang/srt/mem_cache/evict_policy.py | 16 +++++------ .../sglang/srt/mem_cache/hicache_storage.py | 10 +++---- .../sglang/srt/mem_cache/kv_cache_builder.py | 28 +++++++++---------- .../sglang/srt/mem_cache/mamba_radix_cache.py | 8 +++--- python/sglang/srt/mem_cache/memory_pool.py | 20 ++++++------- python/sglang/srt/mem_cache/radix_cache.py | 12 ++++---- python/sglang/srt/mem_cache/registry.py | 2 +- .../storage/file/lru_file_evictor.py | 2 +- .../storage/lmcache/lmc_radix_cache.py | 2 +- .../sglang/srt/mem_cache/swa_radix_cache.py | 4 +-- .../unified_cache_components/swa_component.py | 2 +- .../srt/mem_cache/unified_radix_cache.py | 2 +- .../cuda_graph_buffer_registry.py | 14 +++++----- .../sglang/srt/model_executor/model_runner.py | 12 ++++---- .../runner/prefill_cuda_graph_runner.py | 4 +-- .../model_executor/runner_utils/buffers.py | 6 ++-- python/sglang/srt/model_loader/loader.py | 8 +++--- python/sglang/srt/models/mindspore.py | 6 ++-- python/sglang/srt/models/utils.py | 2 +- .../srt/multimodal/processors/mimo_audio.py | 2 +- .../srt/observability/metrics_collector.py | 12 ++++---- python/sglang/srt/observability/trace.py | 2 +- .../srt/sampling/penaltylib/orchestrator.py | 10 +++---- .../srt/sampling/sampling_batch_info.py | 4 +-- .../sglang/srt/session/session_controller.py | 2 +- python/sglang/srt/speculative/dflash_info.py | 2 +- .../srt/speculative/frozen_kv_mtp_utils.py | 4 +-- .../speculative/frozen_kv_mtp_worker_v2.py | 2 +- .../sglang/srt/speculative/spec_registry.py | 2 +- python/sglang/srt/utils/common.py | 2 +- python/sglang/srt/utils/request_logger.py | 12 ++++---- .../srt/utils/scheduler_status_logger.py | 4 +-- .../test/scripted_runtime/context/api.py | 14 +++++----- .../test/scripted_runtime/context/engine.py | 2 +- .../scripted_runtime/context/http_post.py | 6 ++-- .../context/kv_pool_exhauster.py | 4 +-- .../scripted_runtime/context/lifecycle.py | 12 ++++---- .../context/lock_ref_exhauster.py | 2 +- .../test/scripted_runtime/context/queries.py | 26 ++++++++--------- .../test/scripted_runtime/context/radix.py | 6 ++-- .../scripted_runtime/context/req_starter.py | 2 +- .../test/scripted_runtime/http_server.py | 2 +- .../sglang/test/scripted_runtime/io_struct.py | 2 +- .../test/scripted_runtime/req_handle.py | 4 +-- .../test/scripted_runtime/scheduler_hook.py | 6 ++-- .../kv_canary/test_self_unit_e2e_base.py | 2 +- .../test_customized_info_streaming.py | 4 +-- 185 files changed, 541 insertions(+), 551 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 54f328b0e..66e8c425e 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -35,7 +35,7 @@ repos: hooks: - id: ruff args: - - --select=F401,F821 + - --select=F401,F821,UP037 - --fix files: ^(benchmark/|docs/|examples/|python/sglang/|sgl-model-gateway/py_*|test/) exclude: | diff --git a/python/sglang/_mps_stub.py b/python/sglang/_mps_stub.py index 8a090b398..b463ff6be 100644 --- a/python/sglang/_mps_stub.py +++ b/python/sglang/_mps_stub.py @@ -40,7 +40,7 @@ class Stream: return True # context-manager protocol (``with stream:``) - def __enter__(self) -> "Stream": + def __enter__(self) -> Stream: return self def __exit__(self, *args: Any) -> None: @@ -75,7 +75,7 @@ class StreamContext: def __init__(self, stream: Any = None) -> None: pass - def __enter__(self) -> "StreamContext": + def __enter__(self) -> StreamContext: return self def __exit__(self, *args: Any) -> None: diff --git a/python/sglang/jit_kernel/kv_canary/plan/entries_kernel.py b/python/sglang/jit_kernel/kv_canary/plan/entries_kernel.py index 1a62b8c5b..f63de2484 100644 --- a/python/sglang/jit_kernel/kv_canary/plan/entries_kernel.py +++ b/python/sglang/jit_kernel/kv_canary/plan/entries_kernel.py @@ -13,7 +13,7 @@ if TYPE_CHECKING: @cache_once def _jit_plan_entries_module( has_swa_lut: bool, has_verify_expected_token_pool: bool -) -> "Module": +) -> Module: args = make_cpp_args(has_swa_lut, has_verify_expected_token_pool) return load_jit( "kv_canary_plan_entries", diff --git a/python/sglang/jit_kernel/kv_canary/verify.py b/python/sglang/jit_kernel/kv_canary/verify.py index 2f43f0253..b7fd260bd 100644 --- a/python/sglang/jit_kernel/kv_canary/verify.py +++ b/python/sglang/jit_kernel/kv_canary/verify.py @@ -198,7 +198,7 @@ class VerifyPlan: enable: torch.Tensor @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: raise ValueError( 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), ) - 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 write kernel to avoid hiding bugs.""" self.verify_slot_indices.zero_() @@ -352,7 +352,7 @@ def launch_canary_verify_kernel( @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) return load_jit( "kv_canary_verify", diff --git a/python/sglang/jit_kernel/kv_canary/write.py b/python/sglang/jit_kernel/kv_canary/write.py index d3d0c4382..6926a2dc9 100644 --- a/python/sglang/jit_kernel/kv_canary/write.py +++ b/python/sglang/jit_kernel/kv_canary/write.py @@ -52,7 +52,7 @@ class WritePlan: *, write_req_capacity: int, device: torch.device, - ) -> "WritePlan": + ) -> WritePlan: if write_req_capacity <= 0: raise ValueError( 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), ) - 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 write kernel to avoid hiding bugs.""" self.write_offsets.zero_() @@ -251,7 +251,7 @@ def launch_canary_write_kernel( @cache_once -def _jit_canary_write_module() -> "Module": +def _jit_canary_write_module() -> Module: return load_jit( "kv_canary_write", cuda_files=["kv_canary/canary_write.cuh"], diff --git a/python/sglang/jit_kernel/tests/kv_canary/_canary_helpers.py b/python/sglang/jit_kernel/tests/kv_canary/_canary_helpers.py index dd1ce578c..25b56f8c0 100644 --- a/python/sglang/jit_kernel/tests/kv_canary/_canary_helpers.py +++ b/python/sglang/jit_kernel/tests/kv_canary/_canary_helpers.py @@ -55,7 +55,7 @@ class FakeViolationLog: @classmethod def allocate( cls, *, capacity: int = DEFAULT_RING_CAPACITY, device: torch.device - ) -> "FakeViolationLog": + ) -> FakeViolationLog: return cls( ring=torch.zeros( capacity, consts.VIOLATION_FIELDS, dtype=torch.int64, device=device diff --git a/python/sglang/multimodal_gen/configs/quantization/nunchaku.py b/python/sglang/multimodal_gen/configs/quantization/nunchaku.py index 71fef1b0a..3f89ea55f 100644 --- a/python/sglang/multimodal_gen/configs/quantization/nunchaku.py +++ b/python/sglang/multimodal_gen/configs/quantization/nunchaku.py @@ -69,7 +69,7 @@ class NunchakuSVDQuantArgs: return enable_svdquant, inferred_precision, inferred_rank - def _normalized(self) -> "NunchakuSVDQuantArgs": + def _normalized(self) -> NunchakuSVDQuantArgs: enable_svdquant, inferred_precision, inferred_rank = ( self._infer_from_weights_path() ) @@ -204,7 +204,7 @@ class NunchakuSVDQuantArgs: ) @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). path = ( kwargs.get("transformer_weights_path") diff --git a/python/sglang/multimodal_gen/runtime/distributed/cfg_parallel_utils.py b/python/sglang/multimodal_gen/runtime/distributed/cfg_parallel_utils.py index 39a289516..ddbeeece0 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/cfg_parallel_utils.py +++ b/python/sglang/multimodal_gen/runtime/distributed/cfg_parallel_utils.py @@ -36,7 +36,7 @@ _logged_dispatch_keys: set[tuple[int, int, int]] = set() def _run( - predict_fn: Callable[["CFGBranch"], "torch.Tensor | tuple[torch.Tensor, ...]"], + predict_fn: Callable[[CFGBranch], torch.Tensor | tuple[torch.Tensor, ...]], bid: int, branches, ) -> tuple[torch.Tensor, ...]: @@ -54,9 +54,9 @@ def _run( def run_cfg_parallel( - policy: "CFGPolicy", - predict_fn: Callable[["CFGBranch"], "torch.Tensor | tuple[torch.Tensor, ...]"], -) -> "list[torch.Tensor | tuple[torch.Tensor, ...]]": + policy: CFGPolicy, + predict_fn: Callable[[CFGBranch], torch.Tensor | tuple[torch.Tensor, ...]], +) -> list[torch.Tensor | tuple[torch.Tensor, ...]]: """Dispatch CFG branches across ranks, all-gather results, return in branch order. ``predict_fn`` is a closure capturing all step-varying state @@ -136,12 +136,12 @@ def run_cfg_parallel( def run_two_branch_cfg_parallel( - policy: "CFGPolicy", - predict_fn: Callable[["CFGBranch"], "torch.Tensor | tuple[torch.Tensor, ...]"], + policy: CFGPolicy, + predict_fn: Callable[[CFGBranch], torch.Tensor | tuple[torch.Tensor, ...]], cfg_scale: float, batch, pipeline_config, -) -> "torch.Tensor | tuple[torch.Tensor, ...]": +) -> torch.Tensor | tuple[torch.Tensor, ...]: """Run standard two-pass CFG with the old all-reduce combine. This keeps the existing WAN baselines: it avoids gathering both branch diff --git a/python/sglang/multimodal_gen/runtime/distributed/cfg_policy.py b/python/sglang/multimodal_gen/runtime/distributed/cfg_policy.py index b65e69fde..be2f750cd 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/cfg_policy.py +++ b/python/sglang/multimodal_gen/runtime/distributed/cfg_policy.py @@ -21,7 +21,7 @@ class CFGBranch: is_conditional: bool 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. Override for richer per-branch context (e.g. a branch index instead of @@ -46,11 +46,11 @@ class CFGPolicy: def build( self, - batch: "Req", + batch: Req, image_kwargs: dict[str, Any], pos_cond_kwargs: dict[str, Any], neg_cond_kwargs: dict[str, Any], - ) -> "CFGPolicy": + ) -> CFGPolicy: """Return a new policy with branches populated. Called once before the denoising loop. The returned policy is @@ -66,7 +66,7 @@ class CFGPolicy: def combine( self, predictions: list[torch.Tensor | tuple[torch.Tensor, ...]], - batch: "Req", + batch: Req, cfg_scale: float, pipeline_config: Any, *, @@ -117,7 +117,7 @@ def _unwrap( def _apply_cfg_postprocess( noise_pred: torch.Tensor, noise_pred_cond: torch.Tensor, - batch: "Req", + batch: Req, pipeline_config: Any, ) -> torch.Tensor: if batch.cfg_normalization and float(batch.cfg_normalization) > 0: diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index d3b827c81..8fb2c6b85 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py @@ -435,7 +435,7 @@ class USPAttention(nn.Module): f"but got {backend_enum.name}. " 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.attn_impl = impl_cls( num_heads=num_heads, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/turbo_layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/turbo_layer.py index ce2f3b344..0172c649b 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/turbo_layer.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/turbo_layer.py @@ -252,7 +252,7 @@ class MinimalA2AAttnOp(DistributedAttention): attn_backend = SageSparseLinearAttentionBackend else: 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( num_heads=num_heads, head_size=head_size, diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/bitsandbytes.py b/python/sglang/multimodal_gen/runtime/layers/quantization/bitsandbytes.py index 9e1565895..6895ad000 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/bitsandbytes.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/bitsandbytes.py @@ -105,7 +105,7 @@ class BitsAndBytesConfig(QuantizationConfig): return [] @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): try: value = QuantizationConfig.get_from_keys(config, keys) diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_fp8.py b/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_fp8.py index 7a4bcd6e5..ab016215f 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_fp8.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_fp8.py @@ -74,7 +74,7 @@ class ModelOptFp8Config(QuantizationConfig): return [] @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") if quant_algo is None: raise ValueError( diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py b/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py index 1a3b76e98..a8392c976 100755 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py @@ -187,7 +187,7 @@ class ModelOptFp8Config(ModelOptQuantConfig): return 89 @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") exclude_modules = config.get("ignore") if quant_method is None: diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp4_npu.py b/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp4_npu.py index 3798f36b4..09b4d69c7 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp4_npu.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp4_npu.py @@ -63,7 +63,7 @@ class NPUMXFP4Config(QuantizationConfig): return [] @classmethod - def from_config(cls, config: Dict[str, Any]) -> "NPUMXFP4Config": + def from_config(cls, config: Dict[str, Any]) -> NPUMXFP4Config: return cls() def get_quant_method( diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp8_npu.py b/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp8_npu.py index 17a4370cf..2ba40d1af 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp8_npu.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp8_npu.py @@ -56,7 +56,7 @@ class MXFP8Config(QuantizationConfig): return [] @classmethod - def from_config(cls, config: Dict[str, Any]) -> "MXFP8Config": + def from_config(cls, config: Dict[str, Any]) -> MXFP8Config: return cls() def get_quant_method( diff --git a/python/sglang/multimodal_gen/runtime/managers/dynamic_batch_admission.py b/python/sglang/multimodal_gen/runtime/managers/dynamic_batch_admission.py index 9772fe798..528bf86c4 100644 --- a/python/sglang/multimodal_gen/runtime/managers/dynamic_batch_admission.py +++ b/python/sglang/multimodal_gen/runtime/managers/dynamic_batch_admission.py @@ -78,7 +78,7 @@ class BatchingRule: source: str = "user" @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): raise ValueError( f"batching config rule from {source} must be an object, " @@ -156,7 +156,7 @@ class BatchingRule: class BatchAdmissionController: """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._user_max_batch_size = max(1, int(server_args.batching_max_size)) self._model_path = server_args.model_path diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py index a700d18f8..e52c72497 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py @@ -197,7 +197,7 @@ class Req: VSA_sparsity: float = 0.0 # stage logging - metrics: Optional["RequestMetrics"] = None + metrics: Optional[RequestMetrics] = None # tracing context (TraceReqContext or TraceNullContext) trace_ctx: Union[TraceReqContext, TraceNullContext] = field( @@ -326,7 +326,7 @@ class Req: self.extra["cache_dit_num_inference_steps"] = self.num_inference_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.set_as_warmup(warmup_steps) return req @@ -421,8 +421,8 @@ class OutputBatch: output_file_paths: list[str] | None = None # logged metrics info, directly from Req.timings - metrics: Optional["RequestMetrics"] = None - metrics_list: Optional[list[Optional["RequestMetrics"]]] = None + metrics: Optional[RequestMetrics] = None + metrics_list: Optional[list[Optional[RequestMetrics]]] = None # For ComfyUI integration: noise prediction from denoising stage noise_pred: torch.Tensor | None = None diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/dedup.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/dedup.py index 106aecb48..b3acb2319 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/dedup.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/dedup.py @@ -34,8 +34,8 @@ class StageDedupMixin: def run_grouped_requests( self, - batches: list["Req"], - server_args: "ServerArgs", + batches: list[Req], + server_args: ServerArgs, ) -> list[Any]: """Run this stage for a group of independent requests. @@ -63,7 +63,7 @@ class StageDedupMixin: 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. A fingerprint is the stage-local set of input values that fully @@ -79,10 +79,10 @@ class StageDedupMixin: def run_deduplicated_group( self, - batches: list["Req"], - server_args: "ServerArgs", + batches: list[Req], + server_args: ServerArgs, copy_outputs=None, - ) -> list["Req"]: + ) -> list[Req]: """Run full-stage-equivalent requests once and fan out stage outputs.""" if copy_outputs is None: copy_outputs = self.copy_deduplicated_outputs @@ -102,7 +102,7 @@ class StageDedupMixin: 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. ``deduplicated_output_fields`` uses shallow container copies and shares @@ -175,11 +175,11 @@ class StageDedupMixin: @staticmethod def _group_requests_by_fingerprint( - batches: list["Req"], + batches: list[Req], 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.""" - groups: dict[Any, list[tuple[int, "Req"]]] = {} + groups: dict[Any, list[tuple[int, Req]]] = {} for index, batch in enumerate(batches): fingerprint = fingerprint_fn(batch) groups.setdefault(fingerprint, []).append((index, batch)) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/paint.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/paint.py index 6e1f2e96f..2c6851743 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/paint.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/paint.py @@ -164,7 +164,7 @@ class DDIMSolver: self.ddim_alpha_cumprods = torch.from_numpy(self.ddim_alpha_cumprods) 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_alpha_cumprods = self.ddim_alpha_cumprods.to(device) self.ddim_alpha_cumprods_prev = self.ddim_alpha_cumprods_prev.to(device) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/self_forcing.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/self_forcing.py index f7b7099fa..92fedb50b 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/self_forcing.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/self_forcing.py @@ -53,7 +53,7 @@ class SanaWMSelfForcingSamplerConfig: streaming_cfg_scale: float = 1.0 @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).""" return cls( num_frame_per_block=int(getattr(pcfg, "num_frame_per_block", 3)), diff --git a/python/sglang/multimodal_gen/runtime/server_args_auto_tune.py b/python/sglang/multimodal_gen/runtime/server_args_auto_tune.py index 93e7e2292..dfa4fec50 100644 --- a/python/sglang/multimodal_gen/runtime/server_args_auto_tune.py +++ b/python/sglang/multimodal_gen/runtime/server_args_auto_tune.py @@ -36,7 +36,7 @@ DEFAULT_LAYERWISE_COMPONENT_ARG_NAMES = ( class ServerArgsAutoTuner: """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._explicit_memory_policy = self._has_explicit_memory_policy() self._explicit_layerwise_replacement_policy = ( diff --git a/python/sglang/multimodal_gen/test/server/test_disagg_server.py b/python/sglang/multimodal_gen/test/server/test_disagg_server.py index c3aee6ff3..8ad76c304 100755 --- a/python/sglang/multimodal_gen/test/server/test_disagg_server.py +++ b/python/sglang/multimodal_gen/test/server/test_disagg_server.py @@ -127,7 +127,7 @@ class DisaggCluster: # -- context manager ----------------------------------------------------- - def __enter__(self) -> "DisaggCluster": + def __enter__(self) -> DisaggCluster: for attempt in range(3): try: self._launch_roles() diff --git a/python/sglang/multimodal_gen/test/server/test_server_common.py b/python/sglang/multimodal_gen/test/server/test_server_common.py index 1f8da49e3..894735a2b 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_common.py +++ b/python/sglang/multimodal_gen/test/server/test_server_common.py @@ -63,7 +63,7 @@ logger = init_logger(__name__) # Track test cases missing estimated_full_test_time_s for time measurement output _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( os.environ.get("SGLANG_TEST_OPENAI_REQUEST_TIMEOUT_SECS", "600") ) @@ -518,7 +518,7 @@ class DiffusionServerBase: def _dump_baseline_for_testcase( self, case: DiffusionTestCase, - summary: "PerformanceSummary", + summary: PerformanceSummary, missing_scenario: bool = False, measured_full_time: float | None = None, ) -> None: diff --git a/python/sglang/srt/batch_overlap/operations.py b/python/sglang/srt/batch_overlap/operations.py index 729e00bcb..ea9090e9c 100644 --- a/python/sglang/srt/batch_overlap/operations.py +++ b/python/sglang/srt/batch_overlap/operations.py @@ -117,7 +117,7 @@ class _StageExecutor: debug_name: str, stages: List[Stage], inputs: dict, - child_ctx: Optional["ForwardContext"] = None, + child_ctx: Optional[ForwardContext] = None, ): self._debug_name = debug_name self._stages = stages diff --git a/python/sglang/srt/batch_overlap/two_batch_overlap.py b/python/sglang/srt/batch_overlap/two_batch_overlap.py index dbec58605..af281e6f9 100644 --- a/python/sglang/srt/batch_overlap/two_batch_overlap.py +++ b/python/sglang/srt/batch_overlap/two_batch_overlap.py @@ -266,7 +266,7 @@ def split_spec_info( def compute_split_token_index( split_seq_index: int, - forward_mode: "ForwardMode", + forward_mode: ForwardMode, extend_seq_lens: Optional[Sequence[int]], token_num_per_seq: Optional[int], ) -> int: diff --git a/python/sglang/srt/compilation/compile_phase.py b/python/sglang/srt/compilation/compile_phase.py index 42ef0ce41..152291b19 100644 --- a/python/sglang/srt/compilation/compile_phase.py +++ b/python/sglang/srt/compilation/compile_phase.py @@ -20,7 +20,7 @@ from contextlib import contextmanager import torch _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: @@ -43,7 +43,7 @@ def enable_torch_compile_warmup(): _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 diff --git a/python/sglang/srt/debug_utils/comparator/meta_overrider.py b/python/sglang/srt/debug_utils/comparator/meta_overrider.py index c1ae48eb2..5bbd15cb2 100644 --- a/python/sglang/srt/debug_utils/comparator/meta_overrider.py +++ b/python/sglang/srt/debug_utils/comparator/meta_overrider.py @@ -50,7 +50,7 @@ class MetaOverrider: override_baseline_dims: list[str], override_target_dims: list[str], override_config: Optional[Path], - ) -> "MetaOverrider": + ) -> MetaOverrider: per_side_args: list[tuple[list[str], Literal["both", "baseline", "target"]]] = [ (override_dims, "both"), (override_baseline_dims, "baseline"), diff --git a/python/sglang/srt/debug_utils/comparator/output_types.py b/python/sglang/srt/debug_utils/comparator/output_types.py index d042623ea..e7831810b 100644 --- a/python/sglang/srt/debug_utils/comparator/output_types.py +++ b/python/sglang/srt/debug_utils/comparator/output_types.py @@ -277,7 +277,7 @@ class SummaryRecord(_OutputRecord): errored: int = 0 @model_validator(mode="after") - def _validate_totals(self) -> "SummaryRecord": + def _validate_totals(self) -> SummaryRecord: expected: int = self.passed + self.failed + self.skipped + self.errored if self.total != expected: raise ValueError( diff --git a/python/sglang/srt/debug_utils/dumper.py b/python/sglang/srt/debug_utils/dumper.py index 63af64221..9e11cd3b8 100644 --- a/python/sglang/srt/debug_utils/dumper.py +++ b/python/sglang/srt/debug_utils/dumper.py @@ -250,7 +250,7 @@ class _Dumper: def __init__(self, *, config: DumperConfig): self._config = config self._state = _DumperState() - self._non_intrusives: list["_NonIntrusiveDumper"] = [] + self._non_intrusives: list[_NonIntrusiveDumper] = [] self._grafter = _Grafter(config=config) # ------------------------------- public :: core --------------------------------- diff --git a/python/sglang/srt/disaggregation/common/staging_handler.py b/python/sglang/srt/disaggregation/common/staging_handler.py index 4cc0dc906..a0b6bcbfc 100644 --- a/python/sglang/srt/disaggregation/common/staging_handler.py +++ b/python/sglang/srt/disaggregation/common/staging_handler.py @@ -97,7 +97,7 @@ class DecodeStagingHandler: return 1 @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.""" staging_allocator = kv_manager._staging_ctx.allocator if staging_allocator is None: @@ -132,7 +132,7 @@ class DecodeStagingHandler: # 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 def unregister_decode_req(self, room: int) -> None: @@ -251,13 +251,13 @@ class DecodeStagingHandler: # 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.""" if not getattr(decode_req, "_staging_scatter_done", False): return False 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. Scatter kernels have already been submitted by the decode_thread @@ -292,7 +292,7 @@ class DecodeStagingHandler: staging_offset: int, page_start: int, num_pages: int, - decode_req: "DecodeRequest", + decode_req: DecodeRequest, ) -> bool: """Submit scatter kernels for a staging region to scatter_stream. @@ -347,7 +347,7 @@ class DecodeStagingHandler: 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.""" receiver = decode_req.kv_receiver chunk_infos = getattr(receiver, "chunk_staging_infos", []) @@ -370,7 +370,7 @@ class DecodeStagingHandler: return alloc_id if ok else -1 def _free_and_send_watermark( - self, alloc_id: int, decode_req: "DecodeRequest" + self, alloc_id: int, decode_req: DecodeRequest ) -> None: """Free a staging allocation and broadcast watermark to all prefills.""" self.staging_allocator.free(alloc_id) @@ -474,7 +474,7 @@ class StagingRegisterInfo: @classmethod def from_zmq_fields( cls, msg: list, msg_start_offset: int - ) -> Optional["StagingRegisterInfo"]: + ) -> Optional[StagingRegisterInfo]: i = msg_start_offset base_ptr = ( struct.unpack("Q", msg[i])[0] if len(msg) > i and len(msg[i]) == 8 else 0 diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 1c0825c86..0af05cd01 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -146,7 +146,7 @@ class DecodeReqToTokenPool: def available_size(self): 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 # 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] @@ -170,7 +170,7 @@ class DecodeReqToTokenPool: offset += 1 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" self.free_slots.append(req.req_pool_idx) req.req_pool_idx = None @@ -186,7 +186,7 @@ class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool): max_context_len: int, device: str, enable_memory_saver: bool, - cache_params: "Mamba2CacheParams", + cache_params: Mamba2CacheParams, mamba_layer_ids: List[int], speculative_num_draft_tokens: int, enable_mamba_extra_buffer: bool, diff --git a/python/sglang/srt/disaggregation/decode_hicache_mixin.py b/python/sglang/srt/disaggregation/decode_hicache_mixin.py index bf7c36dda..70c1df8c8 100644 --- a/python/sglang/srt/disaggregation/decode_hicache_mixin.py +++ b/python/sglang/srt/disaggregation/decode_hicache_mixin.py @@ -58,7 +58,7 @@ class HiCacheRestoreResult(Enum): class DecodeHiCachePreallocMixin: """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``. Performs the optional L3 storage hit length query when decode-side @@ -97,7 +97,7 @@ class DecodeHiCachePreallocMixin: ) def _start_hicache_prefetch( - self, req: "Req", prefix_match: Optional["DecodePrefixMatch"] + self, req: Req, prefix_match: Optional[DecodePrefixMatch] ) -> None: """Issue L3 storage prefetch after admission succeeds. @@ -152,7 +152,7 @@ class DecodeHiCachePreallocMixin: class HiCacheRestoreGatedKVReceiver: """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 def poll(self) -> KVPoll: @@ -168,7 +168,7 @@ class HiCacheRestoreGatedKVReceiver: class DecodeHiCacheTransferMixin: """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 ( decode_req.prefix_match is not None and decode_req.prefix_match.prefetch_registered @@ -178,7 +178,7 @@ class DecodeHiCacheTransferMixin: self.tree_cache.dec_lock_ref(decode_req.hicache_restored_node) 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. On success, ``dr.hicache_restored_node`` and ``hicache_restored_kv_indices`` @@ -237,15 +237,13 @@ class DecodeHiCacheTransferMixin: return False return True - def _process_hicache_local_restores( - self, decode_reqs: List["DecodeRequest"] - ) -> None: + def _process_hicache_local_restores(self, decode_reqs: List[DecodeRequest]) -> None: if not hasattr(self.tree_cache, "is_load_back_event_done"): return # Filter once: keep only PENDING reqs that still need restore work; # trivially-done reqs (no prefix_match / nothing to restore) flip to READY. - active: List["DecodeRequest"] = [] + active: List[DecodeRequest] = [] for dr in decode_reqs: if dr.hicache_restore_status != HiCacheRestoreResult.PENDING: continue @@ -291,7 +289,7 @@ class DecodeHiCacheTransferMixin: for dr in queued: 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 if prefix_match is None or not prefix_match.needs_local_restore: return diff --git a/python/sglang/srt/disaggregation/encode_receiver.py b/python/sglang/srt/disaggregation/encode_receiver.py index 2f43b7c50..a4aa338c8 100644 --- a/python/sglang/srt/disaggregation/encode_receiver.py +++ b/python/sglang/srt/disaggregation/encode_receiver.py @@ -83,7 +83,7 @@ class EncoderBootstrapServer: self.port = port self._urls: List[str] = urls if urls is not None else [] 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 = ( health_check_interval if health_check_interval is not None diff --git a/python/sglang/srt/disaggregation/encode_server.py b/python/sglang/srt/disaggregation/encode_server.py index 54fbea330..2ebbbbff8 100644 --- a/python/sglang/srt/disaggregation/encode_server.py +++ b/python/sglang/srt/disaggregation/encode_server.py @@ -2187,7 +2187,7 @@ class EncoderScheduler: self.send_sockets = send_sockets self.max_batch_size = max(1, int(max_batch_size)) 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 def start(self) -> None: diff --git a/python/sglang/srt/disaggregation/mori/conn.py b/python/sglang/srt/disaggregation/mori/conn.py index 81bb7e960..ad33b9b8b 100644 --- a/python/sglang/srt/disaggregation/mori/conn.py +++ b/python/sglang/srt/disaggregation/mori/conn.py @@ -272,7 +272,7 @@ class TransferTarget: @dataclasses.dataclass class _TransferChunk: - sender: "MoriKVSender" + sender: MoriKVSender kv_indices: npt.NDArray[np.int32] index_slice: slice is_last_chunk: bool diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index b69c7cdf2..d16be0065 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/python/sglang/srt/disaggregation/nixl/conn.py @@ -71,7 +71,7 @@ class TransferInfo: decode_prefix_len: Optional[int] = None # for decode radix cache # NOTE: optional staging field; populated via STAGING_RSP. Keep at the # end so positional construction in from_zmq() continues to work. - staging: Optional["StagingTransferInfo"] = None + staging: Optional[StagingTransferInfo] = None def is_dummy(self): # 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) # Keep last: optional, parsed from a variable-length tail of the ZMQ # frame in from_zmq() below, so positional construction stays stable. - staging: Optional["StagingRegisterInfo"] = None + staging: Optional[StagingRegisterInfo] = None @classmethod def from_zmq(cls, msg: List[bytes]): @@ -1293,9 +1293,9 @@ class NixlKVManager(CommonKVManager): def _do_staging_transfer( self, staging_strategy, - kv_chunk: "TransferKVChunk", - req: "TransferInfo", - dst_info: "KVArgsRegisterInfo", + kv_chunk: TransferKVChunk, + req: TransferInfo, + dst_info: KVArgsRegisterInfo, queue: FastQueue, ): """Attempt staging transfer for one chunk. Returns (xfer_handle, deferred). diff --git a/python/sglang/srt/entrypoints/openai/protocol.py b/python/sglang/srt/entrypoints/openai/protocol.py index 5f1efbb11..eae4a312f 100644 --- a/python/sglang/srt/entrypoints/openai/protocol.py +++ b/python/sglang/srt/entrypoints/openai/protocol.py @@ -615,7 +615,7 @@ class Tool(BaseModel): defer_loading: Optional[bool] = None @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: self.function.defer_loading = self.defer_loading return self @@ -1223,7 +1223,7 @@ class TokenizeRequest(BaseModel): ) @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): raise ValueError("Exactly one of 'prompt' or 'messages' must be provided.") return self @@ -1471,7 +1471,7 @@ class ResponsesResponse(BaseModel): ], status: str, usage: Optional[UsageInfo], - ) -> "ResponsesResponse": + ) -> ResponsesResponse: """Create a response from a request.""" # Determine if the output is plain text only to set text.format diff --git a/python/sglang/srt/entrypoints/openai/serving_responses.py b/python/sglang/srt/entrypoints/openai/serving_responses.py index 41aefac68..edda5a471 100644 --- a/python/sglang/srt/entrypoints/openai/serving_responses.py +++ b/python/sglang/srt/entrypoints/openai/serving_responses.py @@ -118,7 +118,7 @@ class OpenAIServingResponses(OpenAIServingChat): # Note: In production, this should use a proper storage backend (Redis, database) # with TTL/expiration to prevent memory leaks self.msg_store: dict[ - str, Union[list[ChatCompletionMessageParam], list["OpenAIMessage"]] + str, Union[list[ChatCompletionMessageParam], list[OpenAIMessage]] ] = {} self.background_tasks: dict[str, asyncio.Task] = {} @@ -631,8 +631,8 @@ class OpenAIServingResponses(OpenAIServingChat): self, request: ResponsesRequest, prev_response: Optional[ResponsesResponse], - ) -> list["OpenAIMessage"]: - messages: list["OpenAIMessage"] = [] + ) -> list[OpenAIMessage]: + messages: list[OpenAIMessage] = [] if prev_response is None: # New conversation. reasoning_effort = request.reasoning.effort if request.reasoning else None diff --git a/python/sglang/srt/eplb/expert_distribution.py b/python/sglang/srt/eplb/expert_distribution.py index 6472dca01..f7d2c5784 100644 --- a/python/sglang/srt/eplb/expert_distribution.py +++ b/python/sglang/srt/eplb/expert_distribution.py @@ -305,7 +305,7 @@ class _SinglePassGatherer(ABC): server_args: ServerArgs, expert_location_metadata: ExpertLocationMetadata, rank: int, - ) -> "_SinglePassGatherer": + ) -> _SinglePassGatherer: if server_args.expert_distribution_recorder_mode == "per_token": return _DetailSinglePassGatherer( server_args, expert_location_metadata, rank @@ -627,13 +627,13 @@ class _Accumulator(ABC): server_args: ServerArgs, expert_location_metadata: ExpertLocationMetadata, rank: int, - ) -> "_Accumulator": + ) -> _Accumulator: return _Accumulator.get_class(server_args)( server_args, expert_location_metadata, rank ) @staticmethod - def get_class(server_args: ServerArgs) -> Type["_Accumulator"]: + def get_class(server_args: ServerArgs) -> Type[_Accumulator]: return { "stat": _StatAccumulator, "stat_approx": _StatAccumulator, diff --git a/python/sglang/srt/eplb/expert_location.py b/python/sglang/srt/eplb/expert_location.py index e5881677c..b5c464cbb 100644 --- a/python/sglang/srt/eplb/expert_location.py +++ b/python/sglang/srt/eplb/expert_location.py @@ -253,7 +253,7 @@ class ExpertLocationMetadata: def update( self, - other: "ExpertLocationMetadata", + other: ExpertLocationMetadata, update_layer_ids: List[int], ): for field in [ diff --git a/python/sglang/srt/hardware_backend/cpu/quantization/awq_kernels.py b/python/sglang/srt/hardware_backend/cpu/quantization/awq_kernels.py index 1ef41d153..4d4c4c483 100644 --- a/python/sglang/srt/hardware_backend/cpu/quantization/awq_kernels.py +++ b/python/sglang/srt/hardware_backend/cpu/quantization/awq_kernels.py @@ -19,7 +19,7 @@ __all__ = ["AWQIntelAMXLinearKernel", "AWQIntelAMXMoEKernel"] class AWQIntelAMXLinearKernel: - def __init__(self, quant_config: "AWQConfig"): + def __init__(self, quant_config: AWQConfig): self.quant_config = quant_config def process_weights_after_loading(self, layer: torch.nn.Module) -> None: @@ -46,7 +46,7 @@ class AWQIntelAMXLinearKernel: class AWQIntelAMXMoEKernel: - def __init__(self, quant_config: "AWQConfig"): + def __init__(self, quant_config: AWQConfig): self.quant_config = quant_config self.moe_runner_config: Optional[MoeRunnerConfig] = None @@ -66,7 +66,7 @@ class AWQIntelAMXMoEKernel: def apply( self, layer: torch.nn.Module, - dispatch_output: "StandardDispatchOutput", + dispatch_output: StandardDispatchOutput, ) -> torch.Tensor: from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput diff --git a/python/sglang/srt/hardware_backend/cpu/quantization/gptq_kernels.py b/python/sglang/srt/hardware_backend/cpu/quantization/gptq_kernels.py index 8923c4b94..53058e905 100644 --- a/python/sglang/srt/hardware_backend/cpu/quantization/gptq_kernels.py +++ b/python/sglang/srt/hardware_backend/cpu/quantization/gptq_kernels.py @@ -19,7 +19,7 @@ __all__ = ["GPTQIntelAMXLinearKernel", "GPTQIntelAMXMoEKernel"] class GPTQIntelAMXLinearKernel: - def __init__(self, quant_config: "GPTQConfig"): + def __init__(self, quant_config: GPTQConfig): self.quant_config = quant_config def process_weights_after_loading(self, layer: torch.nn.Module) -> None: @@ -46,7 +46,7 @@ class GPTQIntelAMXLinearKernel: class GPTQIntelAMXMoEKernel: - def __init__(self, quant_config: "GPTQConfig"): + def __init__(self, quant_config: GPTQConfig): self.quant_config = quant_config self.moe_runner_config: Optional[MoeRunnerConfig] = None @@ -66,7 +66,7 @@ class GPTQIntelAMXMoEKernel: def apply( self, layer: torch.nn.Module, - dispatch_output: "StandardDispatchOutput", + dispatch_output: StandardDispatchOutput, ) -> torch.Tensor: from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput diff --git a/python/sglang/srt/hardware_backend/gpu/quantization/awq_kernels.py b/python/sglang/srt/hardware_backend/gpu/quantization/awq_kernels.py index 59ef742e8..3f61776b8 100644 --- a/python/sglang/srt/hardware_backend/gpu/quantization/awq_kernels.py +++ b/python/sglang/srt/hardware_backend/gpu/quantization/awq_kernels.py @@ -77,7 +77,7 @@ _, scalar_types = get_scalar_types() class AWQLinearKernel: - def __init__(self, quant_config: Optional["QuantizationConfig"] = None): + def __init__(self, quant_config: Optional[QuantizationConfig] = None): self.quant_config = quant_config def process_weights_after_loading(self, layer: torch.nn.Module) -> None: @@ -106,7 +106,7 @@ class AWQLinearKernel: class AWQMarlinLinearKernel: - def __init__(self, quant_config: Optional["QuantizationConfig"] = None): + def __init__(self, quant_config: Optional[QuantizationConfig] = None): self.quant_config = quant_config def process_weights_after_loading(self, layer: torch.nn.Module) -> None: @@ -166,7 +166,7 @@ class AWQMarlinLinearKernel: class AWQMoEKernel: - def __init__(self, quant_config: Optional["QuantizationConfig"] = None): + def __init__(self, quant_config: Optional[QuantizationConfig] = None): self.quant_config = quant_config self.runner: Optional[MoeRunner] = None @@ -236,8 +236,8 @@ class AWQMoEKernel: def apply( self, layer: torch.nn.Module, - dispatch_output: "StandardDispatchOutput", - ) -> "CombineInput": + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: if self.runner is None: raise RuntimeError("moe runner is not initialized") diff --git a/python/sglang/srt/hardware_backend/gpu/quantization/gptq_kernels.py b/python/sglang/srt/hardware_backend/gpu/quantization/gptq_kernels.py index 4902f2221..966e934c2 100644 --- a/python/sglang/srt/hardware_backend/gpu/quantization/gptq_kernels.py +++ b/python/sglang/srt/hardware_backend/gpu/quantization/gptq_kernels.py @@ -83,7 +83,7 @@ def gptq_marlin_moe_repack( class GPTQLinearKernel: - def __init__(self, quant_config: Optional["QuantizationConfig"] = None): + def __init__(self, quant_config: Optional[QuantizationConfig] = None): self.quant_config = quant_config self.use_shuffle = True @@ -129,7 +129,7 @@ class GPTQLinearKernel: class GPTQMarlinLinearKernel: - def __init__(self, quant_config: Optional["QuantizationConfig"] = None): + def __init__(self, quant_config: Optional[QuantizationConfig] = None): self.quant_config = quant_config def process_weights_after_loading(self, layer: torch.nn.Module) -> None: @@ -270,7 +270,7 @@ class GPTQMarlinLinearKernel: class GPTQMarlinMoEKernel: - def __init__(self, quant_config: Optional["QuantizationConfig"] = None): + def __init__(self, quant_config: Optional[QuantizationConfig] = None): self.quant_config = quant_config 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) 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() self.moe_runner_config = moe_runner_config @@ -364,8 +364,8 @@ class GPTQMarlinMoEKernel: def apply( self, layer: torch.nn.Module, - dispatch_output: "StandardDispatchOutput", - ) -> "CombineInput": + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: quant_info = MarlinMoeQuantInfo( w13_qweight=layer.w13_qweight, w2_qweight=layer.w2_qweight, diff --git a/python/sglang/srt/hardware_backend/mlx/aot.py b/python/sglang/srt/hardware_backend/mlx/aot.py index d9be297d2..94158b59e 100644 --- a/python/sglang/srt/hardware_backend/mlx/aot.py +++ b/python/sglang/srt/hardware_backend/mlx/aot.py @@ -204,7 +204,7 @@ class MlxAOTKernelContext: req_pool_idx: dict[str, int], req_to_token_pool: Any | None, layer_caches: list[list[ContiguousAttentionKVCache]], - ) -> "MlxAOTKernelContext": + ) -> MlxAOTKernelContext: """Build optional AOT context for one batched decode step.""" if not aot_kernels.rope.enabled or kv_pool is None: return cls() diff --git a/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_wrapper.py b/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_wrapper.py index 560b5c820..13f731f4f 100644 --- a/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_wrapper.py +++ b/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_wrapper.py @@ -76,7 +76,7 @@ class BatchedDecodeContext: req_to_token_pool: Any | None, attention_layer_indices: list[int] | None = None, attention_pool_index_by_layer: dict[int, int] | None = None, - ) -> "BatchedDecodeContext": + ) -> BatchedDecodeContext: batch_size = len(req_ids) if attention_layer_indices is None: attention_layer_indices = list(range(len(caches[0]))) diff --git a/python/sglang/srt/hardware_backend/mlx/kv_cache/layout.py b/python/sglang/srt/hardware_backend/mlx/kv_cache/layout.py index ac54bea44..5483a4f48 100644 --- a/python/sglang/srt/hardware_backend/mlx/kv_cache/layout.py +++ b/python/sglang/srt/hardware_backend/mlx/kv_cache/layout.py @@ -26,7 +26,7 @@ class MlxModelCacheLayout: cls, layers: Sequence[Any], attention_attrs: Sequence[str | None], - ) -> "MlxModelCacheLayout": + ) -> MlxModelCacheLayout: if len(layers) != len(attention_attrs): raise ValueError( "Layer count and attention attribute count differ: " diff --git a/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py b/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py index ae84db044..af1f7c7c6 100644 --- a/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py +++ b/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py @@ -73,19 +73,19 @@ class MlxPendingJob: """ lazy_tokens: Optional[mx.array] - prefills: list["MlxPendingPrefill"] - extends: list["MlxPendingExtend"] - decode: Optional["MlxPendingDecode"] + prefills: list[MlxPendingPrefill] + extends: list[MlxPendingExtend] + decode: Optional[MlxPendingDecode] mode: str - batch_copy: "ScheduleBatch" - schedule_batch: "ScheduleBatch" + batch_copy: ScheduleBatch + schedule_batch: ScheduleBatch reqs: List[Req] class SchedulerMlxOverlapMixin: """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( pending.prefills, pending.extends, @@ -100,7 +100,7 @@ class SchedulerMlxOverlapMixin: self.process_batch_result(pending.batch_copy, result) @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``. 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_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 # FutureMap relay (decode) before the forward. With deferred input # materialization, get_next_batch_to_run leaves input_ids unset; the diff --git a/python/sglang/srt/hardware_backend/musa/attention/flashattention_backend.py b/python/sglang/srt/hardware_backend/musa/attention/flashattention_backend.py index 01c4fc6bc..b3b364833 100644 --- a/python/sglang/srt/hardware_backend/musa/attention/flashattention_backend.py +++ b/python/sglang/srt/hardware_backend/musa/attention/flashattention_backend.py @@ -39,11 +39,11 @@ _MATE_NO_MLA_SCHEDULER_METADATA_DICT: dict = {} _MATE_NO_MLA_SCHEDULER_METADATA_LOCK = threading.Lock() # Global reference to the current backend instance (set during __init__) -_CURRENT_BACKEND: Optional["MusaFlashAttentionBackend"] = None +_CURRENT_BACKEND: Optional[MusaFlashAttentionBackend] = None def _compute_scheduler_metadata( - backend: "MusaFlashAttentionBackend", + backend: MusaFlashAttentionBackend, cu_seqlens_q: torch.Tensor, cu_seqlens_k_new: Optional[torch.Tensor], cache_seqlens: torch.Tensor, diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py index 07f81e62c..cf0de8f8a 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -530,7 +530,7 @@ class AscendAttnBackend(AttentionBackend): forward_mode: ForwardMode, seq_lens: torch.Tensor, out_cache_loc: Optional[torch.Tensor] = None, - ) -> "ForwardMetadata": + ) -> ForwardMetadata: """Create and store the per-bs ForwardMetadata for CUDA graph capture.""" metadata = ForwardMetadata() metadata.block_tables = self.graph_metadata["block_tables"][:bs, :] @@ -847,7 +847,7 @@ class AscendAttnBackend(AttentionBackend): q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, - layer: "RadixAttention", + layer: RadixAttention, forward_batch: ForwardBatch, ) -> torch.Tensor: """CP-aware attention for standard (non-MLA) models using FIA on Ascend NPU. diff --git a/python/sglang/srt/hardware_backend/npu/moe/fuseep.py b/python/sglang/srt/hardware_backend/npu/moe/fuseep.py index deabbf623..f2ae0f3ca 100644 --- a/python/sglang/srt/hardware_backend/npu/moe/fuseep.py +++ b/python/sglang/srt/hardware_backend/npu/moe/fuseep.py @@ -26,7 +26,7 @@ if TYPE_CHECKING: _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() return DeepEPBuffer.get_deepep_buffer( get_tp_group().device_group, @@ -39,9 +39,9 @@ def _get_fuseep_buffer(layer: "FusedMoE"): def forward_fuseep( - layer: "FusedMoE", + layer: FusedMoE, hidden_states: torch.Tensor, - topk_output: "TopKOutput", + topk_output: TopKOutput, ) -> torch.Tensor: buf = _get_fuseep_buffer(layer) hidden_states, _ = buf.fused_deep_moe( diff --git a/python/sglang/srt/hardware_backend/npu/quantization/awq_kernels.py b/python/sglang/srt/hardware_backend/npu/quantization/awq_kernels.py index 9d66ab722..197c3de31 100644 --- a/python/sglang/srt/hardware_backend/npu/quantization/awq_kernels.py +++ b/python/sglang/srt/hardware_backend/npu/quantization/awq_kernels.py @@ -17,7 +17,7 @@ import torch_npu class AWQAscendLinearKernel: - def __init__(self, quant_config: Optional["QuantizationConfig"] = None): + def __init__(self, quant_config: Optional[QuantizationConfig] = None): self.quant_config = quant_config def process_weights_after_loading(self, layer: torch.nn.Module) -> None: @@ -72,7 +72,7 @@ class AWQAscendLinearKernel: class AWQAscendMoEKernel: - def __init__(self, quant_config: Optional["QuantizationConfig"] = None): + def __init__(self, quant_config: Optional[QuantizationConfig] = None): self.quant_config = quant_config self.kernel = NPUW4A16Int4DynamicMoEMethod() @@ -151,7 +151,7 @@ class AWQAscendMoEKernel: def apply( self, layer: torch.nn.Module, - dispatch_output: "StandardDispatchOutput", + dispatch_output: StandardDispatchOutput, ) -> torch.Tensor: return self.kernel.apply(layer, dispatch_output) diff --git a/python/sglang/srt/hardware_backend/npu/quantization/gptq_kernels.py b/python/sglang/srt/hardware_backend/npu/quantization/gptq_kernels.py index d3484f451..b3ed395d0 100644 --- a/python/sglang/srt/hardware_backend/npu/quantization/gptq_kernels.py +++ b/python/sglang/srt/hardware_backend/npu/quantization/gptq_kernels.py @@ -60,7 +60,7 @@ def unpack_from_int32( class GPTQLinearAscendKernel: - def __init__(self, quant_config: Optional["QuantizationConfig"] = None): + def __init__(self, quant_config: Optional[QuantizationConfig] = None): self.quant_config = quant_config self.use_v2_format = quant_config.checkpoint_format == "gptq_v2" @@ -128,15 +128,15 @@ class GPTQLinearAscendKernel: class GPTQMoEAscendKernel: - def __init__(self, quant_config: Optional["QuantizationConfig"] = None): + def __init__(self, quant_config: Optional[QuantizationConfig] = None): self.quant_config = quant_config 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( self, layer: torch.nn.Module, - moe_runner_config: "MoeRunnerConfig", + moe_runner_config: MoeRunnerConfig, **extra_weight_attrs, ): self.moe_runner_config = moe_runner_config @@ -277,7 +277,7 @@ class GPTQMoEAscendKernel: def apply( self, layer: torch.nn.Module, - dispatch_output: "StandardDispatchOutput", + dispatch_output: StandardDispatchOutput, ) -> torch.Tensor: from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput diff --git a/python/sglang/srt/kv_canary/api.py b/python/sglang/srt/kv_canary/api.py index 143563f01..0a2d7cf57 100644 --- a/python/sglang/srt/kv_canary/api.py +++ b/python/sglang/srt/kv_canary/api.py @@ -29,9 +29,9 @@ logger = logging.getLogger(__name__) def install_canary( *, - server_args: "ServerArgs", - model_runner: "ModelRunner", - token_oracle_manager: Optional["TokenOracleManager"] = None, + server_args: ServerArgs, + model_runner: ModelRunner, + token_oracle_manager: Optional[TokenOracleManager] = None, ) -> Optional[CanaryManager]: config = CanaryConfig.from_env(server_args) if config.mode is CanaryMode.NONE: @@ -101,9 +101,7 @@ def install_canary( return manager -def _patch_model_forward( - *, model_runner: "ModelRunner", manager: CanaryManager -) -> None: +def _patch_model_forward(*, model_runner: ModelRunner, manager: CanaryManager) -> None: def _with_canary_bracketing(original: Callable, *args: Any, **kwargs: Any) -> Any: forward_batch = _extract_forward_batch(args, kwargs) assert ( diff --git a/python/sglang/srt/kv_canary/capacities.py b/python/sglang/srt/kv_canary/capacities.py index 62f17981c..a0d6a4785 100644 --- a/python/sglang/srt/kv_canary/capacities.py +++ b/python/sglang/srt/kv_canary/capacities.py @@ -43,11 +43,11 @@ class CanaryLaunchCapacities: def from_args( cls, *, - server_args: "ServerArgs", + server_args: ServerArgs, req_to_token_pool_size: int, max_seq_len_per_req: int, pool_slot_count: int, - ) -> "CanaryLaunchCapacities": + ) -> CanaryLaunchCapacities: if req_to_token_pool_size <= 0: raise ValueError( "kv-canary: req_to_token_pool_size must be positive, " diff --git a/python/sglang/srt/kv_canary/config.py b/python/sglang/srt/kv_canary/config.py index 0eb0ea194..c970818ce 100644 --- a/python/sglang/srt/kv_canary/config.py +++ b/python/sglang/srt/kv_canary/config.py @@ -60,7 +60,7 @@ class CanaryConfig: stats_print_every_n_steps: int @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() if mode_raw not in ("none", "log", "raise"): raise ValueError( diff --git a/python/sglang/srt/kv_canary/expected_inputs.py b/python/sglang/srt/kv_canary/expected_inputs.py index 9cff403e3..ee78990b7 100644 --- a/python/sglang/srt/kv_canary/expected_inputs.py +++ b/python/sglang/srt/kv_canary/expected_inputs.py @@ -11,13 +11,13 @@ class ExpectedInputs: positions: torch.Tensor @classmethod - def allocate(cls, *, capacity: int, device: torch.device) -> "ExpectedInputs": + def allocate(cls, *, capacity: int, device: torch.device) -> ExpectedInputs: return cls( tokens=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( tokens=self.tokens[:num_tokens], positions=self.positions[:num_tokens], diff --git a/python/sglang/srt/kv_canary/perturb/config.py b/python/sglang/srt/kv_canary/perturb/config.py index dbf3fec8c..400b0e7b0 100644 --- a/python/sglang/srt/kv_canary/perturb/config.py +++ b/python/sglang/srt/kv_canary/perturb/config.py @@ -25,7 +25,7 @@ class PerturbConfig: warmup_steps: int @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_unused_cache_prob = ( envs.SGLANG_KV_CANARY_PERTURB_REAL_KV_UNUSED_CACHE_PROB.get() diff --git a/python/sglang/srt/kv_canary/perturb/manager.py b/python/sglang/srt/kv_canary/perturb/manager.py index 4a21ef26f..cf112b8ca 100644 --- a/python/sglang/srt/kv_canary/perturb/manager.py +++ b/python/sglang/srt/kv_canary/perturb/manager.py @@ -24,7 +24,7 @@ class PerturbManager: self, *, config: PerturbConfig, - req_to_token_pool: "ReqToTokenPool", + req_to_token_pool: ReqToTokenPool, buffer_groups: tuple[CanaryBufferGroup, ...], outer_step_counter_getter: Callable[[], int], swa_window_size: int = 0, @@ -36,18 +36,18 @@ class PerturbManager: self._outer_step_counter_getter = outer_step_counter_getter self._swa_window_size = swa_window_size self._sweep_interval = sweep_interval - self._radix_cache: Optional["BasePrefixCache"] = None + self._radix_cache: Optional[BasePrefixCache] = None self._warmup_gate = WarmupGate( 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 def perturb( self, *, - maybe_inaccurate_forward_batch: Optional["ForwardBatch"], + maybe_inaccurate_forward_batch: Optional[ForwardBatch], ) -> None: self.perturb_req_to_token(maybe_inaccurate_forward_batch) self.perturb_real_kv_used(maybe_inaccurate_forward_batch) @@ -56,12 +56,12 @@ class PerturbManager: def perturb_post_forward( self, *, - maybe_inaccurate_forward_batch: Optional["ForwardBatch"], + maybe_inaccurate_forward_batch: Optional[ForwardBatch], ) -> None: self.perturb_real_kv_post_forward(maybe_inaccurate_forward_batch) def perturb_req_to_token( - self, maybe_inaccurate_forward_batch: Optional["ForwardBatch"] + self, maybe_inaccurate_forward_batch: Optional[ForwardBatch] ) -> None: req_to_token.run( maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch, @@ -71,7 +71,7 @@ class PerturbManager: ) def perturb_real_kv_used( - self, maybe_inaccurate_forward_batch: Optional["ForwardBatch"] + self, maybe_inaccurate_forward_batch: Optional[ForwardBatch] ) -> None: real_kv_used.run( maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch, @@ -83,7 +83,7 @@ class PerturbManager: ) def perturb_real_kv_unused_cache( - self, maybe_inaccurate_forward_batch: Optional["ForwardBatch"] + self, maybe_inaccurate_forward_batch: Optional[ForwardBatch] ) -> None: real_kv_unused_cache.run( maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch, @@ -97,7 +97,7 @@ class PerturbManager: ) def perturb_real_kv_post_forward( - self, maybe_inaccurate_forward_batch: Optional["ForwardBatch"] + self, maybe_inaccurate_forward_batch: Optional[ForwardBatch] ) -> None: real_kv_post_forward.run( maybe_inaccurate_forward_batch=maybe_inaccurate_forward_batch, diff --git a/python/sglang/srt/kv_canary/perturb/next_token_swap.py b/python/sglang/srt/kv_canary/perturb/next_token_swap.py index d17da61d4..274996858 100644 --- a/python/sglang/srt/kv_canary/perturb/next_token_swap.py +++ b/python/sglang/srt/kv_canary/perturb/next_token_swap.py @@ -25,7 +25,7 @@ class NextTokenSwapConfig: warmup_steps: int @classmethod - def from_env(cls) -> "NextTokenSwapConfig": + def from_env(cls) -> NextTokenSwapConfig: return cls( prob=envs.SGLANG_KV_CANARY_PERTURB_NEXT_TOKEN_SWAP_PROB.get(), warmup_steps=envs.SGLANG_KV_CANARY_PERTURB_WARMUP_STEPS.get(), diff --git a/python/sglang/srt/kv_canary/perturb/real_kv_post_forward.py b/python/sglang/srt/kv_canary/perturb/real_kv_post_forward.py index 6db7c3b84..73d9e6a2f 100644 --- a/python/sglang/srt/kv_canary/perturb/real_kv_post_forward.py +++ b/python/sglang/srt/kv_canary/perturb/real_kv_post_forward.py @@ -34,7 +34,7 @@ logger = logging.getLogger(__name__) def run( *, - maybe_inaccurate_forward_batch: Optional["ForwardBatch"], + maybe_inaccurate_forward_batch: Optional[ForwardBatch], config: PerturbConfig, buffer_groups: tuple[CanaryBufferGroup, ...], warmup_gate: WarmupGate, diff --git a/python/sglang/srt/kv_canary/perturb/real_kv_unused_cache.py b/python/sglang/srt/kv_canary/perturb/real_kv_unused_cache.py index b18af9a31..332d0ac73 100644 --- a/python/sglang/srt/kv_canary/perturb/real_kv_unused_cache.py +++ b/python/sglang/srt/kv_canary/perturb/real_kv_unused_cache.py @@ -34,10 +34,10 @@ logger = logging.getLogger(__name__) def run( *, - maybe_inaccurate_forward_batch: Optional["ForwardBatch"], + maybe_inaccurate_forward_batch: Optional[ForwardBatch], config: PerturbConfig, buffer_groups: tuple[CanaryBufferGroup, ...], - radix_cache: Optional["BasePrefixCache"], + radix_cache: Optional[BasePrefixCache], swa_window_size: int, sweep_interval: int, outer_step_counter: int, @@ -114,7 +114,7 @@ def run( def _pick_sweep_slot_for_group( *, - radix_cache: Optional["BasePrefixCache"], + radix_cache: Optional[BasePrefixCache], group: CanaryBufferGroup, swa_window_size: int, ) -> Optional[int]: diff --git a/python/sglang/srt/kv_canary/perturb/real_kv_used.py b/python/sglang/srt/kv_canary/perturb/real_kv_used.py index 2131147a3..1db32ca7e 100644 --- a/python/sglang/srt/kv_canary/perturb/real_kv_used.py +++ b/python/sglang/srt/kv_canary/perturb/real_kv_used.py @@ -37,9 +37,9 @@ logger = logging.getLogger(__name__) def run( *, - maybe_inaccurate_forward_batch: Optional["ForwardBatch"], + maybe_inaccurate_forward_batch: Optional[ForwardBatch], config: PerturbConfig, - req_to_token_pool: "ReqToTokenPool", + req_to_token_pool: ReqToTokenPool, buffer_groups: tuple[CanaryBufferGroup, ...], swa_window_size: int, warmup_gate: WarmupGate, @@ -111,8 +111,8 @@ def run( def _pick_active_slot_for_group( *, - maybe_inaccurate_forward_batch: "ForwardBatch", - req_to_token_pool: "ReqToTokenPool", + maybe_inaccurate_forward_batch: ForwardBatch, + req_to_token_pool: ReqToTokenPool, group: CanaryBufferGroup, swa_window_size: int, ) -> Optional[ReqToTokenEntry]: diff --git a/python/sglang/srt/kv_canary/perturb/req_to_token.py b/python/sglang/srt/kv_canary/perturb/req_to_token.py index 00b207fbc..ee6ed4214 100644 --- a/python/sglang/srt/kv_canary/perturb/req_to_token.py +++ b/python/sglang/srt/kv_canary/perturb/req_to_token.py @@ -26,9 +26,9 @@ logger = logging.getLogger(__name__) def run( *, - maybe_inaccurate_forward_batch: Optional["ForwardBatch"], + maybe_inaccurate_forward_batch: Optional[ForwardBatch], config: PerturbConfig, - req_to_token_pool: "ReqToTokenPool", + req_to_token_pool: ReqToTokenPool, warmup_gate: WarmupGate, ) -> None: if not should_run_perturbation( diff --git a/python/sglang/srt/kv_canary/perturb/slot_picker.py b/python/sglang/srt/kv_canary/perturb/slot_picker.py index e6d9a0ee6..383202dbe 100644 --- a/python/sglang/srt/kv_canary/perturb/slot_picker.py +++ b/python/sglang/srt/kv_canary/perturb/slot_picker.py @@ -21,8 +21,8 @@ class ReqToTokenEntry: def collect_active_slots( *, - maybe_inaccurate_forward_batch: "ForwardBatch", - req_to_token_pool: "ReqToTokenPool", + maybe_inaccurate_forward_batch: ForwardBatch, + req_to_token_pool: ReqToTokenPool, exclude_out_cache_loc: bool = True, ) -> list[ReqToTokenEntry]: """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( - *, maybe_inaccurate_forward_batch: "ForwardBatch" + *, maybe_inaccurate_forward_batch: ForwardBatch ) -> Optional[int]: out_cache_loc = maybe_inaccurate_forward_batch.out_cache_loc if out_cache_loc is None: diff --git a/python/sglang/srt/kv_canary/perturb/utils.py b/python/sglang/srt/kv_canary/perturb/utils.py index 47c4a41c6..320976c73 100644 --- a/python/sglang/srt/kv_canary/perturb/utils.py +++ b/python/sglang/srt/kv_canary/perturb/utils.py @@ -70,7 +70,7 @@ def should_run_perturbation( perturb_name: str, probability: float, warmup_gate: WarmupGate, - maybe_inaccurate_forward_batch: Optional["ForwardBatch"], + maybe_inaccurate_forward_batch: Optional[ForwardBatch], require_forward_batch: bool = True, ) -> bool: if probability <= 0.0: diff --git a/python/sglang/srt/kv_canary/plan_input.py b/python/sglang/srt/kv_canary/plan_input.py index bcbaddeea..9b15a7f63 100644 --- a/python/sglang/srt/kv_canary/plan_input.py +++ b/python/sglang/srt/kv_canary/plan_input.py @@ -53,7 +53,7 @@ class PlanInput: *, bs_capacity: int, device: torch.device, - ) -> "PlanInput": + ) -> PlanInput: return cls( req_pool_indices=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 bs = int(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( *, - forward_batch: "ForwardBatch", + forward_batch: ForwardBatch, out_prefix_lens: torch.Tensor, out_extend_seq_lens: torch.Tensor, bs: int, diff --git a/python/sglang/srt/kv_canary/radix_cache_walker.py b/python/sglang/srt/kv_canary/radix_cache_walker.py index 36f444d94..516220dea 100644 --- a/python/sglang/srt/kv_canary/radix_cache_walker.py +++ b/python/sglang/srt/kv_canary/radix_cache_walker.py @@ -22,7 +22,7 @@ class RadixCacheWalkResult: def walk_radix_cache_for_canary( *, - radix_cache: "BasePrefixCache", + radix_cache: BasePrefixCache, unlocked_only: bool = False, swa_resident_only: bool = False, ) -> RadixCacheWalkResult: @@ -68,8 +68,8 @@ def walk_radix_cache_for_canary( def _walk_radix_subtree( *, - node: "TreeNode", - radix_cache: "BasePrefixCache", + node: TreeNode, + radix_cache: BasePrefixCache, depth: int, parent_last_slot: int, slot_buf: list[int], @@ -123,8 +123,8 @@ def _walk_radix_subtree( def _node_is_unlocked_for_canary( *, - node: "TreeNode", - radix_cache: "BasePrefixCache", + node: TreeNode, + radix_cache: BasePrefixCache, ) -> bool: if type(radix_cache) is RadixCache: return node.lock_ref == 0 @@ -139,8 +139,8 @@ def _node_is_unlocked_for_canary( def _node_is_swa_resident_for_canary( *, - node: "TreeNode", - radix_cache: "BasePrefixCache", + node: TreeNode, + radix_cache: BasePrefixCache, ) -> bool: if type(radix_cache) is not SWARadixCache: return True diff --git a/python/sglang/srt/kv_canary/req_to_expected_token_ids_manager.py b/python/sglang/srt/kv_canary/req_to_expected_token_ids_manager.py index d36ea782e..5321a1d93 100644 --- a/python/sglang/srt/kv_canary/req_to_expected_token_ids_manager.py +++ b/python/sglang/srt/kv_canary/req_to_expected_token_ids_manager.py @@ -15,7 +15,7 @@ if TYPE_CHECKING: def compute_req_all_ids_info( - reqs: "list[Req]", + reqs: list[Req], ) -> tuple[torch.Tensor, torch.Tensor]: """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( *, - forward_batch: "ForwardBatch", + forward_batch: ForwardBatch, req_to_verify_expected_tokens: Optional[torch.Tensor], ) -> None: """Scatter the forward batch's per-req token-id snapshot into the device-side pool.""" diff --git a/python/sglang/srt/kv_canary/runner/canary_manager.py b/python/sglang/srt/kv_canary/runner/canary_manager.py index 04c27eecf..b6d3ca0f2 100644 --- a/python/sglang/srt/kv_canary/runner/canary_manager.py +++ b/python/sglang/srt/kv_canary/runner/canary_manager.py @@ -47,18 +47,18 @@ class CanaryManager: perturb_config: PerturbConfig, buffer_groups: tuple[CanaryBufferGroup, ...], device: torch.device, - req_to_token_pool: "ReqToTokenPool", + req_to_token_pool: ReqToTokenPool, launch_capacities: CanaryLaunchCapacities, swa_window_size: int = 0, token_oracle_manager: Optional[TokenOracleManager] = None, - swa_allocator: Optional["SWATokenToKVPoolAllocator"] = None, + swa_allocator: Optional[SWATokenToKVPoolAllocator] = None, speculative_num_steps: int = 1, is_eagle_draft_decode: bool = False, ) -> None: self.config = config self._req_to_token_pool = req_to_token_pool 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._active_single_forward_manager_index: Optional[int] = None @@ -183,7 +183,7 @@ class CanaryManager: self._active_single_forward_manager_index = None def pre_ops_maybe_inside_graph( - self, forward_batch: "ForwardBatch" + self, forward_batch: ForwardBatch ) -> _PreOpsMaybeInsideGraphOutput: assert self._active_single_forward_manager_index is not None, ( "kv-canary: pre_ops_maybe_inside_graph called without active SingleForwardManager; " @@ -194,7 +194,7 @@ class CanaryManager: def post_ops_maybe_inside_graph( self, - forward_batch: "ForwardBatch", + forward_batch: ForwardBatch, pre_ops_output: _PreOpsMaybeInsideGraphOutput, ) -> None: assert self._active_single_forward_manager_index is not None, ( @@ -209,7 +209,7 @@ class CanaryManager: self, *, single_forward_indices: Sequence[int], - maybe_inaccurate_forward_batch: "ForwardBatch", + maybe_inaccurate_forward_batch: ForwardBatch, ) -> Iterator[None]: self._pre_ops_outside_graph( single_forward_indices=single_forward_indices, @@ -227,7 +227,7 @@ class CanaryManager: self, *, single_forward_indices: Sequence[int], - maybe_inaccurate_forward_batch: "ForwardBatch", + maybe_inaccurate_forward_batch: ForwardBatch, ) -> None: for idx in single_forward_indices: self._single_forward_managers[idx].pre_ops_outside_graph( @@ -241,7 +241,7 @@ class CanaryManager: self, *, single_forward_indices: Sequence[int], - maybe_inaccurate_forward_batch: "ForwardBatch", + maybe_inaccurate_forward_batch: ForwardBatch, ) -> None: for idx in single_forward_indices: self._single_forward_managers[idx].post_ops_outside_graph() @@ -264,7 +264,7 @@ class CanaryManager: single_forward_manager.phase_checker.enable_assert() 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._perturb_manager.attach_radix_cache(radix_cache) diff --git a/python/sglang/srt/kv_canary/runner/future_tensor.py b/python/sglang/srt/kv_canary/runner/future_tensor.py index e13fffdbb..eff94502c 100644 --- a/python/sglang/srt/kv_canary/runner/future_tensor.py +++ b/python/sglang/srt/kv_canary/runner/future_tensor.py @@ -22,7 +22,7 @@ class FutureTensors: @classmethod def device_to_host( cls, xs_device: _TensorOrDict, *, d2h_stream: torch.cuda.Stream - ) -> "FutureTensors": + ) -> FutureTensors: assert not torch.cuda.is_current_stream_capturing(), ( "FutureTensors.device_to_host must not be called during cuda-graph " "capture: the d2h side-stream copy + pinned-host alloc cannot be " diff --git a/python/sglang/srt/kv_canary/runner/kernel_launcher.py b/python/sglang/srt/kv_canary/runner/kernel_launcher.py index 8bf459772..ef5af1087 100644 --- a/python/sglang/srt/kv_canary/runner/kernel_launcher.py +++ b/python/sglang/srt/kv_canary/runner/kernel_launcher.py @@ -64,7 +64,7 @@ def launch_endpoints_per_forward( tag_filter: Callable[[CanaryLaunchTag], bool], verify_plan: VerifyPlan, write_plan: WritePlan, - forward_batch: "ForwardBatch", + forward_batch: ForwardBatch, expected_inputs: ExpectedInputs, violation_log: ViolationLog, real_kv_hash_mode: RealKvHashMode, diff --git a/python/sglang/srt/kv_canary/runner/swa_divergence.py b/python/sglang/srt/kv_canary/runner/swa_divergence.py index d23db2f87..1db37af57 100644 --- a/python/sglang/srt/kv_canary/runner/swa_divergence.py +++ b/python/sglang/srt/kv_canary/runner/swa_divergence.py @@ -32,8 +32,8 @@ class SwaDivergenceReporter: device: torch.device, d2h_stream: torch.cuda.Stream, interval: int, - swa_allocator: Optional["SWATokenToKVPoolAllocator"] = None, - req_to_token_pool: Optional["ReqToTokenPool"] = None, + swa_allocator: Optional[SWATokenToKVPoolAllocator] = None, + req_to_token_pool: Optional[ReqToTokenPool] = None, ) -> None: self._interval = interval self._swa_allocator = swa_allocator @@ -57,7 +57,7 @@ class SwaDivergenceReporter: self, *, outer_step_counter: int, - maybe_inaccurate_forward_batch: Optional["ForwardBatch"], + maybe_inaccurate_forward_batch: Optional[ForwardBatch], ) -> None: self._forward_ct += 1 self._handler.step( @@ -72,7 +72,7 @@ class SwaDivergenceReporter: self, *, outer_step_counter: int, - maybe_inaccurate_forward_batch: Optional["ForwardBatch"], + maybe_inaccurate_forward_batch: Optional[ForwardBatch], ) -> Optional[dict[str, Any]]: if outer_step_counter == 0 or outer_step_counter % self._interval != 0: return None @@ -134,14 +134,14 @@ class SwaDivergenceLog: ) @classmethod - def parse(cls, line: str) -> Optional["SwaDivergenceLog"]: + def parse(cls, line: str) -> Optional[SwaDivergenceLog]: match = _SWA_DIVERGENCE_LINE_RE.search(line) if match is None: return None return cls(**json.loads(match.group(1))) @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 for match in _SWA_DIVERGENCE_LINE_RE.finditer(text): last_match = match @@ -150,7 +150,7 @@ class SwaDivergenceLog: return cls(**json.loads(last_match.group(1))), last_match.group(0) @classmethod - def find_all(cls, text: str) -> list[tuple["SwaDivergenceLog", str]]: + def find_all(cls, text: str) -> list[tuple[SwaDivergenceLog, str]]: return [ (cls(**json.loads(match.group(1))), match.group(0)) for match in _SWA_DIVERGENCE_LINE_RE.finditer(text) @@ -159,9 +159,9 @@ class SwaDivergenceLog: def compute_swa_out_of_window_tokens( *, - swa_allocator: "SWATokenToKVPoolAllocator", - req_to_token_pool: "ReqToTokenPool", - maybe_inaccurate_forward_batch: "ForwardBatch", + swa_allocator: SWATokenToKVPoolAllocator, + req_to_token_pool: ReqToTokenPool, + maybe_inaccurate_forward_batch: ForwardBatch, ) -> torch.Tensor: """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 @@ -180,9 +180,9 @@ def compute_swa_out_of_window_tokens( def compute_swa_full_idx_divergence( *, - swa_allocator: "SWATokenToKVPoolAllocator", - req_to_token_pool: "ReqToTokenPool", - maybe_inaccurate_forward_batch: "ForwardBatch", + swa_allocator: SWATokenToKVPoolAllocator, + req_to_token_pool: ReqToTokenPool, + maybe_inaccurate_forward_batch: ForwardBatch, ) -> torch.Tensor: """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 diff --git a/python/sglang/srt/kv_canary/runner/sweep.py b/python/sglang/srt/kv_canary/runner/sweep.py index 496991e1f..b2fdcf498 100644 --- a/python/sglang/srt/kv_canary/runner/sweep.py +++ b/python/sglang/srt/kv_canary/runner/sweep.py @@ -34,7 +34,7 @@ class SweepOrchestrator: self._endpoints = endpoints self._swa_window_size = swa_window_size 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._sweep_passes: int = 0 @@ -43,7 +43,7 @@ class SweepOrchestrator: def sweep_passes(self) -> int: 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 def maybe_run_sweep(self) -> None: diff --git a/python/sglang/srt/kv_canary/single_forward_manager/data.py b/python/sglang/srt/kv_canary/single_forward_manager/data.py index e28a9c932..0ea5e3286 100644 --- a/python/sglang/srt/kv_canary/single_forward_manager/data.py +++ b/python/sglang/srt/kv_canary/single_forward_manager/data.py @@ -21,7 +21,7 @@ class PostOpsInsideGraphOutputBuffer: num_slot_tags: int, swa_verify_total_count_shape: tuple[int, ...] | None, device: torch.device, - ) -> "PostOpsInsideGraphOutputBuffer": + ) -> PostOpsInsideGraphOutputBuffer: return cls( verify_plan_enable=torch.zeros(1, dtype=torch.int32, device=device), kernel_run_counters=torch.zeros( diff --git a/python/sglang/srt/kv_canary/single_forward_manager/manager.py b/python/sglang/srt/kv_canary/single_forward_manager/manager.py index 6749eb339..83e72cadd 100644 --- a/python/sglang/srt/kv_canary/single_forward_manager/manager.py +++ b/python/sglang/srt/kv_canary/single_forward_manager/manager.py @@ -66,7 +66,7 @@ class SingleForwardManager: device_state: CanaryDeviceState, buffer_groups: tuple[CanaryBufferGroup, ...], endpoints: tuple[CanaryEndpoint, ...], - req_to_token_pool: "ReqToTokenPool", + req_to_token_pool: ReqToTokenPool, swa_window_size: int, per_forward_verify_capacity: int, per_forward_write_req_capacity: int, @@ -119,7 +119,7 @@ class SingleForwardManager: return self._phase_checker def pre_ops_outside_graph( - self, *, maybe_inaccurate_forward_batch: "ForwardBatch" + self, *, maybe_inaccurate_forward_batch: ForwardBatch ) -> None: self._phase_checker.update( expect_phase=_SingleForwardPhase.IDLE, @@ -150,8 +150,8 @@ class SingleForwardManager: ) def pre_ops_maybe_inside_graph( - self, forward_batch: "ForwardBatch" - ) -> "_PreOpsMaybeInsideGraphOutput": + self, forward_batch: ForwardBatch + ) -> _PreOpsMaybeInsideGraphOutput: self._phase_checker.update( expect_phase=_SingleForwardPhase.AFTER_PRE_OUT, next_phase=_SingleForwardPhase.AFTER_PRE_MAYBE_IN, @@ -238,8 +238,8 @@ class SingleForwardManager: def post_ops_maybe_inside_graph( self, - forward_batch: "ForwardBatch", - pre_ops_output: "_PreOpsMaybeInsideGraphOutput", + forward_batch: ForwardBatch, + pre_ops_output: _PreOpsMaybeInsideGraphOutput, ) -> None: self._phase_checker.update( expect_phase=_SingleForwardPhase.AFTER_PRE_MAYBE_IN, @@ -293,7 +293,7 @@ class SingleForwardManager: self._enable_warner.tick(self._output_buffer.verify_plan_enable) def _should_enable_write_input_assert_for_launch( - self, forward_batch: "ForwardBatch" + self, forward_batch: ForwardBatch ) -> bool: if not self._config.enable_write_input_assert: return False diff --git a/python/sglang/srt/kv_canary/state.py b/python/sglang/srt/kv_canary/state.py index 177a27246..304ede4bd 100644 --- a/python/sglang/srt/kv_canary/state.py +++ b/python/sglang/srt/kv_canary/state.py @@ -41,7 +41,7 @@ class ViolationLog: violation_write_index: torch.Tensor @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: raise ValueError( f"kv-canary: ViolationLog ring_capacity must be positive, got {ring_capacity}" @@ -102,7 +102,7 @@ class CanaryDeviceState: num_tags: int, req_to_token_alloc_size: Optional[int] = None, max_context_len: Optional[int] = None, - ) -> "CanaryDeviceState": + ) -> CanaryDeviceState: if num_tags <= 0: raise ValueError( f"kv-canary: CanaryDeviceState num_tags must be positive, got {num_tags}" diff --git a/python/sglang/srt/kv_canary/sweep_plan_builder.py b/python/sglang/srt/kv_canary/sweep_plan_builder.py index 3e9f43168..c4314f05c 100644 --- a/python/sglang/srt/kv_canary/sweep_plan_builder.py +++ b/python/sglang/srt/kv_canary/sweep_plan_builder.py @@ -13,7 +13,7 @@ if TYPE_CHECKING: def build_verify_plan_radix_sweep( *, - radix_cache: "BasePrefixCache", + radix_cache: BasePrefixCache, swa_window_size: int, full_to_swa_index_mapping: Optional[torch.Tensor], unlocked_only: bool = False, diff --git a/python/sglang/srt/kv_canary/token_oracle/install.py b/python/sglang/srt/kv_canary/token_oracle/install.py index 2a57a76ee..d88cc010d 100644 --- a/python/sglang/srt/kv_canary/token_oracle/install.py +++ b/python/sglang/srt/kv_canary/token_oracle/install.py @@ -11,7 +11,7 @@ if TYPE_CHECKING: def install_token_oracle_from_env( - *, server_args: "ServerArgs", vocab_size: int + *, server_args: ServerArgs, vocab_size: int ) -> Optional[TokenOracleManager]: # Must be called before create_sampler() so the factory is present when the # Sampler is first constructed. diff --git a/python/sglang/srt/kv_canary/token_oracle/oracle_manager.py b/python/sglang/srt/kv_canary/token_oracle/oracle_manager.py index 0935607c7..87976352e 100644 --- a/python/sglang/srt/kv_canary/token_oracle/oracle_manager.py +++ b/python/sglang/srt/kv_canary/token_oracle/oracle_manager.py @@ -18,7 +18,7 @@ class TokenOracleManager: def fill_expected_inputs( self, *, - forward_batch: "ForwardBatch", + forward_batch: ForwardBatch, expected_inputs_out: ExpectedInputs, ) -> None: positions = forward_batch.positions @@ -57,7 +57,7 @@ class TokenOracleManager: def _build_generalized_req_id_per_token( *, - forward_batch: "ForwardBatch", + forward_batch: ForwardBatch, num_tokens: int, generalized_req_ids_per_row: torch.Tensor, ) -> torch.Tensor: diff --git a/python/sglang/srt/kv_canary/token_oracle/sampler.py b/python/sglang/srt/kv_canary/token_oracle/sampler.py index 2b96f4882..9ebadad31 100644 --- a/python/sglang/srt/kv_canary/token_oracle/sampler.py +++ b/python/sglang/srt/kv_canary/token_oracle/sampler.py @@ -33,8 +33,8 @@ class _OracleSampler(Sampler): def forward( self, - logits_output: "LogitsProcessorOutput", - sampling_info: "SamplingBatchInfo", + logits_output: LogitsProcessorOutput, + sampling_info: SamplingBatchInfo, return_logprob: bool, top_logprobs_nums: List[int], token_ids_logprobs: List[List[int]], diff --git a/python/sglang/srt/layers/attention/aiter_utils.py b/python/sglang/srt/layers/attention/aiter_utils.py index 844ccff14..b08893648 100644 --- a/python/sglang/srt/layers/attention/aiter_utils.py +++ b/python/sglang/srt/layers/attention/aiter_utils.py @@ -43,12 +43,12 @@ if TYPE_CHECKING: def forward_extend_vectorized_5d( - backend: "AiterAttnBackend", + backend: AiterAttnBackend, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - layer: "RadixAttention", - forward_batch: "ForwardBatch", + layer: RadixAttention, + forward_batch: ForwardBatch, bs0: int, window_size, sinks, @@ -207,10 +207,10 @@ def forward_extend_vectorized_5d( def forward_decode_vectorized_5d( - backend: "AiterAttnBackend", + backend: AiterAttnBackend, q: torch.Tensor, - layer: "RadixAttention", - forward_batch: "ForwardBatch", + layer: RadixAttention, + forward_batch: ForwardBatch, k_cache: torch.Tensor, v_cache: torch.Tensor, o: torch.Tensor, diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py index 6488cb121..426c5749b 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py @@ -996,7 +996,7 @@ class DeepseekV4HipRadixBackend( self.forward_metadata = current_raw def _attach_unified_kv_decode_streams( - self, core: "DSV4AttnMetadata", req_pool_indices: torch.Tensor + self, core: DSV4AttnMetadata, req_pool_indices: torch.Tensor ) -> None: """build the ragged decode index streams once per forward""" from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import ( @@ -1031,7 +1031,7 @@ class DeepseekV4HipRadixBackend( def _attach_unified_kv_prefill_meta( self, - core: "DSV4AttnMetadata", + core: DSV4AttnMetadata, req_pool_indices: torch.Tensor, seq_lens: torch.Tensor, extend_seq_lens: torch.Tensor, @@ -1065,7 +1065,7 @@ class DeepseekV4HipRadixBackend( forward_batch: ForwardBatch, compress_ratio: Literal[0, 4, 128], attn_sink: torch.Tensor, - core_attn_metadata: "DSV4AttnMetadata", + core_attn_metadata: DSV4AttnMetadata, save_kv_cache: bool = True, ) -> torch.Tensor: """unified_kv paged-attention path over the bf16 unified_kv""" diff --git a/python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py b/python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py index 1f25bebed..4df019367 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py @@ -71,8 +71,8 @@ class DeepseekSparseAttnBackendMTPPrecomputeMixin: req_pool_indices: torch.Tensor, seq_lens: torch.Tensor, seq_lens_cpu: torch.Tensor, - forward_mode: "ForwardMode", - spec_info: Optional["SpecInput"], + forward_mode: ForwardMode, + spec_info: Optional[SpecInput], ) -> PrecomputedMetadata: """Precompute all shared metadata for multi-step backends. @@ -252,7 +252,7 @@ class DeepseekSparseAttnBackendMTPPrecomputeMixin: req_pool_indices: torch.Tensor, seq_lens: torch.Tensor, seq_lens_cpu: torch.Tensor, - spec_info: "SpecInput", + spec_info: SpecInput, ) -> PrecomputedMetadata: """Precompute metadata for draft extend mode.""" max_seqlen_k = int(seq_lens_cpu.max().item()) diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index 13eda3172..6f0c5d831 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -116,7 +116,7 @@ class DSAFlashMLAMetadata: num_splits=self.num_splits[sli], ) - def copy_(self, other: "DSAFlashMLAMetadata"): + def copy_(self, other: DSAFlashMLAMetadata): self.flashmla_metadata.copy_(other.flashmla_metadata) self.num_splits.copy_(other.num_splits) @@ -866,7 +866,7 @@ class DeepseekSparseAttnBackend( forward_mode: ForwardMode, spec_info: Optional[SpecInput], 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.""" self.set_dsa_prefill_impl(forward_batch=None) diff --git a/python/sglang/srt/layers/attention/dsv4/metadata.py b/python/sglang/srt/layers/attention/dsv4/metadata.py index 1ffe15f65..e1ce33752 100644 --- a/python/sglang/srt/layers/attention/dsv4/metadata.py +++ b/python/sglang/srt/layers/attention/dsv4/metadata.py @@ -153,7 +153,7 @@ class PagedIndexerMetadata: def max_c4_seq_len(self) -> int: return self.page_table.shape[1] * self.c4_page_size - def copy_(self, other: "PagedIndexerMetadata"): + def copy_(self, other: PagedIndexerMetadata): if is_hip(): copy_fields = ["page_table", "c4_seq_lens"] assign_fields = ["deep_gemm_metadata"] diff --git a/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py b/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py index 006588129..5a7279c3e 100644 --- a/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py +++ b/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py @@ -91,7 +91,7 @@ class TokenspeedMLABackend(TRTLLMMLABackend): def __init__( self, - model_runner: "ModelRunner", + model_runner: ModelRunner, skip_prefill: bool = False, kv_indptr_buf: Optional[torch.Tensor] = None, q_indptr_decode_buf: Optional[torch.Tensor] = None, @@ -221,7 +221,7 @@ class TokenspeedMLABackend(TRTLLMMLABackend): kv_a: torch.Tensor, k_pe: torch.Tensor, positions: torch.Tensor, - layer: "DeepseekV2AttentionMLA", + layer: DeepseekV2AttentionMLA, forward_batch: ForwardBatch, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """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, seq_lens: torch.Tensor, max_seq_len: int, - layer: "RadixAttention", + layer: RadixAttention, ) -> torch.Tensor: k_scale = getattr(layer, "k_scale_float", None) if k_scale is None: @@ -308,7 +308,7 @@ class TokenspeedMLABackend(TRTLLMMLABackend): q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - layer: "RadixAttention", + layer: RadixAttention, batch_size: int, cum_seq_lens_q: torch.Tensor, max_q_len: int, @@ -342,7 +342,7 @@ class TokenspeedMLAMultiStepDraftBackend(TRTLLMMLAMultiStepDraftBackend): """Multi-step draft backend for tokenspeed_mla used by EAGLE.""" 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) # Parent populates self.attn_backends with TRT-LLM instances; replace diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index 4bd4bc0d2..ebfbc42ad 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mha_backend.py @@ -341,7 +341,7 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): forward_mode: ForwardMode, spec_info, device: torch.device, - ) -> "TRTLLMMHAMetadata": + ) -> TRTLLMMHAMetadata: """Create TRTLLMMHAMetadata with pre-allocated buffer slice refs, stored in the dict.""" metadata = TRTLLMMHAMetadata() diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index 99c325e9d..eb0ae35e7 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mla_backend.py @@ -1073,7 +1073,7 @@ class TRTLLMMLAMultiStepDraftBackend(FlashInferMLAMultiStepDraftBackend): def __init__( self, - model_runner: "ModelRunner", + model_runner: ModelRunner, topk: int, speculative_num_steps: int, backend: str = "trtllm-gen", diff --git a/python/sglang/srt/layers/moe/mega_moe.py b/python/sglang/srt/layers/moe/mega_moe.py index 44c5ddf8a..8e8474f16 100644 --- a/python/sglang/srt/layers/moe/mega_moe.py +++ b/python/sglang/srt/layers/moe/mega_moe.py @@ -94,7 +94,7 @@ def _get_mega_moe_symm_buffer( 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(): return 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( - moe: "DeepseekV2MoE", + moe: DeepseekV2MoE, hidden_states: torch.Tensor, forward_batch: Optional[ForwardBatch] = None, input_ids_global: Optional[torch.Tensor] = None, @@ -149,7 +149,7 @@ def forward_mega_moe( def _run_mega_routed( - moe: "DeepseekV2MoE", + moe: DeepseekV2MoE, hidden_states: torch.Tensor, forward_batch: Optional[ForwardBatch], input_ids_global: Optional[torch.Tensor], diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py index 56f90cd6d..02c781fa4 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py @@ -342,7 +342,7 @@ class CuteDslFp4MoeQuantInfo(MoeQuantInfo): use_nvfp4_dispatch: bool = False # 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") diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_mxfp4.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_mxfp4.py index ec00efe17..6d8e141ef 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_mxfp4.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_mxfp4.py @@ -93,10 +93,10 @@ def _flashinfer_cutlass_fused_moe(): @register_fused_func("none", "flashinfer_mxfp4") def fused_experts_none_to_flashinfer_mxfp4( - dispatch_output: "StandardDispatchOutput", + dispatch_output: StandardDispatchOutput, quant_info: MoeQuantInfo, runner_config: MoeRunnerConfig, -) -> "StandardCombineInput": +) -> StandardCombineInput: """SM90 W4A16 MXFP4 fused expert forward pass. Mirrors the legacy ``Mxfp4MoEMethod._apply_sm90_cutlass`` and DSv4's diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_kernels.py b/python/sglang/srt/layers/moe/moe_runner/triton_kernels.py index 258761e50..d5bd82e9d 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton_kernels.py +++ b/python/sglang/srt/layers/moe/moe_runner/triton_kernels.py @@ -42,9 +42,9 @@ class TritonKernelsRunnerInput(RunnerInput): """Input bundle passed to the triton-kernels runner core.""" hidden_states: torch.Tensor - routing_data: "RoutingData" - gather_indx: "GatherIndx" - scatter_indx: "ScatterIndx" + routing_data: RoutingData + gather_indx: GatherIndx + scatter_indx: ScatterIndx @property def runner_backend(self) -> MoeRunnerBackend: @@ -158,7 +158,7 @@ class TritonKernelsRunnerCore(MoeRunnerCore): @register_pre_permute("standard", "triton_kernel") def pre_permute_standard_to_triton_kernels( - dispatch_output: "StandardDispatchOutput", + dispatch_output: StandardDispatchOutput, quant_info: TritonKernelsQuantInfo, runner_config: MoeRunnerConfig, running_state: dict, diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe_triton_kernels.py b/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe_triton_kernels.py index d0ca5abaa..9650ed9c3 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe_triton_kernels.py +++ b/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe_triton_kernels.py @@ -673,7 +673,7 @@ def _set_triton_tma_allocator(): # --- B TensorDescriptor cache (LRU) --- _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): diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index f9d89828d..997823db4 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -314,7 +314,7 @@ class BypassedTopKOutput(NamedTuple): def format(self) -> TopKOutputFormat: 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 topk_ids / topk_weights rather than doing routing internally.""" return select_experts( diff --git a/python/sglang/srt/layers/pooler.py b/python/sglang/srt/layers/pooler.py index a31e60dfd..0245e02b0 100644 --- a/python/sglang/srt/layers/pooler.py +++ b/python/sglang/srt/layers/pooler.py @@ -107,7 +107,7 @@ def pool_at_delimiter_positions( def score_and_pool( score_head: nn.Module, - pooler: "Pooler", + pooler: Pooler, hidden_states: torch.Tensor, forward_batch: ForwardBatch, input_ids: torch.Tensor, diff --git a/python/sglang/srt/layers/quantization/awq/schemes/awq_cpu.py b/python/sglang/srt/layers/quantization/awq/schemes/awq_cpu.py index 3560a88d1..0965a76b4 100644 --- a/python/sglang/srt/layers/quantization/awq/schemes/awq_cpu.py +++ b/python/sglang/srt/layers/quantization/awq/schemes/awq_cpu.py @@ -23,14 +23,14 @@ __all__ = ["AWQIntelAMXLinearScheme", "AWQIntelAMXMoEScheme"] class AWQIntelAMXLinearScheme(AWQLinearScheme): """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) class AWQIntelAMXMoEScheme(AWQMoEScheme): """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) def create_moe_runner( diff --git a/python/sglang/srt/layers/quantization/awq/schemes/awq_linear.py b/python/sglang/srt/layers/quantization/awq/schemes/awq_linear.py index 535776477..30f8bef76 100644 --- a/python/sglang/srt/layers/quantization/awq/schemes/awq_linear.py +++ b/python/sglang/srt/layers/quantization/awq/schemes/awq_linear.py @@ -16,11 +16,11 @@ __all__ = ["AWQLinearScheme", "AWQAscendLinearScheme"] class AWQLinearScheme(AWQLinearSchemeBase): - def __init__(self, quant_config: "AWQConfig"): + def __init__(self, quant_config: AWQConfig): self.quant_config = 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 ( AWQLinearKernel, ) @@ -102,7 +102,7 @@ class AWQLinearScheme(AWQLinearSchemeBase): 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 ( AWQAscendLinearKernel, ) diff --git a/python/sglang/srt/layers/quantization/awq/schemes/awq_marlin.py b/python/sglang/srt/layers/quantization/awq/schemes/awq_marlin.py index 10f115dd5..84de8efca 100644 --- a/python/sglang/srt/layers/quantization/awq/schemes/awq_marlin.py +++ b/python/sglang/srt/layers/quantization/awq/schemes/awq_marlin.py @@ -17,11 +17,11 @@ __all__ = ["AWQMarlinLinearScheme"] class AWQMarlinLinearScheme(AWQLinearSchemeBase): - def __init__(self, quant_config: "AWQMarlinConfig"): + def __init__(self, quant_config: AWQMarlinConfig): self.quant_config = 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 ( AWQMarlinLinearKernel, ) diff --git a/python/sglang/srt/layers/quantization/awq/schemes/awq_moe.py b/python/sglang/srt/layers/quantization/awq/schemes/awq_moe.py index 2233aa99a..9e8467f39 100644 --- a/python/sglang/srt/layers/quantization/awq/schemes/awq_moe.py +++ b/python/sglang/srt/layers/quantization/awq/schemes/awq_moe.py @@ -23,13 +23,13 @@ __all__ = ["AWQMoEScheme", "AWQAscendMoEScheme"] class AWQMoEScheme(AWQMoESchemeBase): - def __init__(self, quant_config: "AWQMarlinConfig"): + def __init__(self, quant_config: AWQMarlinConfig): self.quant_config = quant_config if self.quant_config.weight_bits != 4: raise ValueError("AWQMoEScheme only supports 4bit now.") 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 ( AWQMoEKernel, ) @@ -137,13 +137,13 @@ class AWQMoEScheme(AWQMoESchemeBase): def apply_weights( self, layer: torch.nn.Module, - dispatch_output: "StandardDispatchOutput", + dispatch_output: StandardDispatchOutput, ): return self.kernel.apply(layer, dispatch_output) class AWQAscendMoEScheme(AWQMoEScheme): - def _init_kernel(self, quant_config: "AWQConfig"): + def _init_kernel(self, quant_config: AWQConfig): from sglang.srt.hardware_backend.npu.quantization.awq_kernels import ( AWQAscendMoEKernel, ) diff --git a/python/sglang/srt/layers/quantization/base_config.py b/python/sglang/srt/layers/quantization/base_config.py index cf276bc9f..ccce6169f 100644 --- a/python/sglang/srt/layers/quantization/base_config.py +++ b/python/sglang/srt/layers/quantization/base_config.py @@ -109,7 +109,7 @@ class FusedMoEMethodBase(QuantizeMethodBase): ) -> CombineInput: raise NotImplementedError - def get_triton_quant_info(self, layer: torch.nn.Module) -> "TritonMoeQuantInfo": + def get_triton_quant_info(self, layer: torch.nn.Module) -> TritonMoeQuantInfo: """Return a ``TritonMoeQuantInfo`` describing the quantisation state stored on *layer*. @@ -163,7 +163,7 @@ class QuantizationConfig(ABC): @classmethod @abstractmethod - def from_config(cls, config: Dict[str, Any]) -> "QuantizationConfig": + def from_config(cls, config: Dict[str, Any]) -> QuantizationConfig: """Create a config class from the model's quantization config.""" raise NotImplementedError() @@ -246,7 +246,7 @@ class QuantizationConfig(ABC): raise NotImplementedError() def apply_weight_name_mapper( - self, hf_to_sglang_mapper: "WeightsMapper" + self, hf_to_sglang_mapper: WeightsMapper ): # noqa: B027 """ Interface for models to update module names referenced in diff --git a/python/sglang/srt/layers/quantization/bitsandbytes.py b/python/sglang/srt/layers/quantization/bitsandbytes.py index 87e0d06ff..85bbe8bc3 100644 --- a/python/sglang/srt/layers/quantization/bitsandbytes.py +++ b/python/sglang/srt/layers/quantization/bitsandbytes.py @@ -90,7 +90,7 @@ class BitsAndBytesConfig(QuantizationConfig): return [] @classmethod - def from_config(cls, config: dict[str, Any]) -> "BitsAndBytesConfig": + def from_config(cls, config: dict[str, Any]) -> BitsAndBytesConfig: def get_safe_value(config, keys, default_value=None): try: value = QuantizationConfig.get_from_keys(config, keys) diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py index 28a9d567a..431ffadda 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py @@ -145,7 +145,7 @@ class CompressedTensorsConfig(QuantizationConfig): def get_scaled_act_names(self) -> List[str]: return [] - def apply_weight_name_mapper(self, hf_to_sglang_mapper: "WeightsMapper"): + def apply_weight_name_mapper(self, hf_to_sglang_mapper: WeightsMapper): self.target_scheme_map = hf_to_sglang_mapper.apply_dict(self.target_scheme_map) self.ignore = hf_to_sglang_mapper.apply_list(self.ignore) self.sparsity_scheme_map = hf_to_sglang_mapper.apply_dict( diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py index 58562bb23..581972032 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py @@ -487,8 +487,8 @@ class CompressedTensorsWNA16TritonMoE(CompressedTensorsWNA16MoE): def apply_weights( self, layer: torch.nn.Module, - dispatch_output: "StandardDispatchOutput", - ) -> "CombineInput": + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: assert ( self.moe_runner_config.activation == "silu" ), "Only SiLU activation is supported." diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index ea951303f..437ab2cf0 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -302,7 +302,7 @@ class Fp8Config(QuantizationConfig): def get_scaled_act_names(self) -> List[str]: return [] - def apply_weight_name_mapper(self, hf_to_sglang_mapper: "WeightsMapper"): + def apply_weight_name_mapper(self, hf_to_sglang_mapper: WeightsMapper): if self.ignored_layers: self.ignored_layers = list( dict.fromkeys(hf_to_sglang_mapper.apply_list(self.ignored_layers)) @@ -2094,7 +2094,7 @@ class Fp8MoEMethod(FusedMoEMethodBase): self, layer: torch.nn.Module, no_combine: bool = False, - ) -> Optional["AiterMoeQuantInfo"]: + ) -> Optional[AiterMoeQuantInfo]: if not (_use_aiter or _use_hip_int4): return None assert not no_combine, f"{no_combine=} is not supported." diff --git a/python/sglang/srt/layers/quantization/gguf.py b/python/sglang/srt/layers/quantization/gguf.py index 5c1e7a63e..6b65e1ef7 100644 --- a/python/sglang/srt/layers/quantization/gguf.py +++ b/python/sglang/srt/layers/quantization/gguf.py @@ -81,7 +81,7 @@ class GGUFConfig(QuantizationConfig): def get_scaled_act_names(self) -> List[str]: return [] - def get_name(self) -> "str": + def get_name(self) -> str: return "gguf" def get_supported_act_dtypes(self) -> list[torch.dtype]: @@ -96,7 +96,7 @@ class GGUFConfig(QuantizationConfig): return [] # no extra configs. @classmethod - def from_config(cls, config: dict[str, Any]) -> "GGUFConfig": + def from_config(cls, config: dict[str, Any]) -> GGUFConfig: modules_to_not_convert = cls.get_from_keys_or( config, ["modules_to_not_convert"], None ) @@ -104,7 +104,7 @@ class GGUFConfig(QuantizationConfig): def get_quant_method( self, layer: torch.nn.Module, prefix: str - ) -> Optional["QuantizeMethodBase"]: + ) -> Optional[QuantizeMethodBase]: from sglang.srt.layers.moe.fused_moe_triton import FusedMoE from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding diff --git a/python/sglang/srt/layers/quantization/gptq/schemes/gptq_cpu.py b/python/sglang/srt/layers/quantization/gptq/schemes/gptq_cpu.py index ebf463659..f3ce55600 100644 --- a/python/sglang/srt/layers/quantization/gptq/schemes/gptq_cpu.py +++ b/python/sglang/srt/layers/quantization/gptq/schemes/gptq_cpu.py @@ -29,7 +29,7 @@ if TYPE_CHECKING: __all__ = ["GPTQIntelAMXLinearScheme", "GPTQIntelAMXMoEScheme"] -def _check_cpu_amx_support(quant_config: "GPTQConfig") -> None: +def _check_cpu_amx_support(quant_config: GPTQConfig) -> None: if quant_config.desc_act and not ( quant_config.true_sequential and quant_config.static_groups ): @@ -46,7 +46,7 @@ def _check_cpu_amx_support(quant_config: "GPTQConfig") -> None: class GPTQIntelAMXLinearScheme(GPTQLinearScheme): """Linear scheme for GPTQ on Intel CPU with AMX.""" - def _init_kernel(self, quant_config: "GPTQConfig"): + def _init_kernel(self, quant_config: GPTQConfig): return GPTQIntelAMXLinearKernel(quant_config) def create_weights( @@ -152,7 +152,7 @@ class GPTQIntelAMXLinearScheme(GPTQLinearScheme): class GPTQIntelAMXMoEScheme(GPTQMoESchemeBase): """MoE scheme for GPTQ on Intel CPU with AMX.""" - def __init__(self, quant_config: "GPTQConfig"): + def __init__(self, quant_config: GPTQConfig): self.quant_config = quant_config self.kernel = GPTQIntelAMXMoEKernel(quant_config) @@ -280,6 +280,6 @@ class GPTQIntelAMXMoEScheme(GPTQMoESchemeBase): def apply_weights( self, layer: torch.nn.Module, - dispatch_output: "StandardDispatchOutput", + dispatch_output: StandardDispatchOutput, ): return self.kernel.apply(layer, dispatch_output) diff --git a/python/sglang/srt/layers/quantization/gptq/schemes/gptq_linear.py b/python/sglang/srt/layers/quantization/gptq/schemes/gptq_linear.py index 547d5e5e6..bfcab0e74 100644 --- a/python/sglang/srt/layers/quantization/gptq/schemes/gptq_linear.py +++ b/python/sglang/srt/layers/quantization/gptq/schemes/gptq_linear.py @@ -23,12 +23,12 @@ __all__ = ["GPTQLinearScheme", "GPTQAscendLinearScheme"] class GPTQLinearScheme(GPTQLinearSchemeBase): - def __init__(self, quant_config: "GPTQConfig"): + def __init__(self, quant_config: GPTQConfig): self.quant_config = quant_config self.use_v2_format = quant_config.checkpoint_format == "gptq_v2" self.kernel = self._init_kernel(quant_config) - def _init_kernel(self, quant_config: "GPTQConfig"): + def _init_kernel(self, quant_config: GPTQConfig): from sglang.srt.hardware_backend.gpu.quantization.gptq_kernels import ( GPTQLinearKernel, ) @@ -152,7 +152,7 @@ class GPTQLinearScheme(GPTQLinearSchemeBase): class GPTQAscendLinearScheme(GPTQLinearScheme): - def _init_kernel(self, quant_config: "GPTQConfig"): + def _init_kernel(self, quant_config: GPTQConfig): from sglang.srt.hardware_backend.npu.quantization.gptq_kernels import ( GPTQLinearAscendKernel, ) diff --git a/python/sglang/srt/layers/quantization/gptq/schemes/gptq_marlin.py b/python/sglang/srt/layers/quantization/gptq/schemes/gptq_marlin.py index d54c800df..8fa2b6e99 100644 --- a/python/sglang/srt/layers/quantization/gptq/schemes/gptq_marlin.py +++ b/python/sglang/srt/layers/quantization/gptq/schemes/gptq_marlin.py @@ -27,7 +27,7 @@ __all__ = ["GPTQMarlinLinearScheme"] class GPTQMarlinLinearScheme(GPTQLinearSchemeBase): - def __init__(self, quant_config: "GPTQMarlinConfig"): + def __init__(self, quant_config: GPTQMarlinConfig): self.quant_config = quant_config self.kernel = self._init_kernel(quant_config) @@ -36,7 +36,7 @@ class GPTQMarlinLinearScheme(GPTQLinearSchemeBase): group_size=self.quant_config.group_size, ) - def _init_kernel(self, quant_config: "GPTQMarlinConfig"): + def _init_kernel(self, quant_config: GPTQMarlinConfig): from sglang.srt.hardware_backend.gpu.quantization.gptq_kernels import ( GPTQMarlinLinearKernel, ) diff --git a/python/sglang/srt/layers/quantization/gptq/schemes/gptq_moe.py b/python/sglang/srt/layers/quantization/gptq/schemes/gptq_moe.py index a5c804f55..c83bdba19 100644 --- a/python/sglang/srt/layers/quantization/gptq/schemes/gptq_moe.py +++ b/python/sglang/srt/layers/quantization/gptq/schemes/gptq_moe.py @@ -18,7 +18,7 @@ __all__ = ["GPTQMoEAscendScheme", "GPTQMarlinMoEScheme"] class GPTQMoEAscendScheme(GPTQMoESchemeBase): - def __init__(self, quant_config: "GPTQConfig"): + def __init__(self, quant_config: GPTQConfig): self.quant_config = quant_config from sglang.srt.hardware_backend.npu.quantization.gptq_kernels import ( GPTQMoEAscendKernel, @@ -132,13 +132,13 @@ class GPTQMoEAscendScheme(GPTQMoESchemeBase): def apply_weights( self, layer: torch.nn.Module, - dispatch_output: "StandardDispatchOutput", + dispatch_output: StandardDispatchOutput, ): return self.kernel.apply(layer, dispatch_output) class GPTQMarlinMoEScheme(GPTQMoESchemeBase): - def __init__(self, quant_config: "GPTQMarlinConfig"): + def __init__(self, quant_config: GPTQMarlinConfig): self.quant_config = quant_config from sglang.srt.hardware_backend.gpu.quantization.gptq_kernels import ( GPTQMarlinMoEKernel, @@ -300,6 +300,6 @@ class GPTQMarlinMoEScheme(GPTQMoESchemeBase): def apply_weights( self, layer: torch.nn.Module, - dispatch_output: "StandardDispatchOutput", + dispatch_output: StandardDispatchOutput, ): return self.kernel.apply(layer, dispatch_output) diff --git a/python/sglang/srt/layers/quantization/marlin_utils.py b/python/sglang/srt/layers/quantization/marlin_utils.py index bb2e1159f..2e77d4f93 100644 --- a/python/sglang/srt/layers/quantization/marlin_utils.py +++ b/python/sglang/srt/layers/quantization/marlin_utils.py @@ -679,7 +679,7 @@ class MarlinConfig(QuantizationConfig): return ["quantize_config.json"] @classmethod - def from_config(cls, config: dict[str, Any]) -> "MarlinConfig": + def from_config(cls, config: dict[str, Any]) -> MarlinConfig: group_size = cls.get_from_keys(config, ["group_size"]) lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"], default=False) return cls(group_size, lm_head_quantized) diff --git a/python/sglang/srt/layers/quantization/mlx.py b/python/sglang/srt/layers/quantization/mlx.py index ecb786b6c..b1bf84744 100644 --- a/python/sglang/srt/layers/quantization/mlx.py +++ b/python/sglang/srt/layers/quantization/mlx.py @@ -73,7 +73,7 @@ class MlxQuantizationConfig(QuantizationConfig): return [] @classmethod - def from_config(cls, config: Dict[str, Any]) -> "MlxQuantizationConfig": + def from_config(cls, config: Dict[str, Any]) -> MlxQuantizationConfig: raise NotImplementedError(cls._ERR) @classmethod diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index d362b764e..a73e36b48 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -325,7 +325,7 @@ class ModelOptQuantConfig(QuantizationConfig): return [] def apply_weight_name_mapper( - self, hf_to_sglang_mapper: "WeightsMapper" + self, hf_to_sglang_mapper: WeightsMapper ): # noqa: B027 # Map excluded module patterns from HF layout to sglang layout. # Ref: HF hf_quant_config.json for nvidia/Kimi-K2.5-NVFP4 @@ -613,7 +613,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfig): packed_modules_mapping: Optional[Dict[str, List[str]]], quantized_layers: Dict[str, Dict[str, Any]], fp8_config: ModelOptFp8Config, - nvfp4_config: "ModelOptFp4Config", + nvfp4_config: ModelOptFp4Config, ) -> None: super().__init__(kv_cache_quant_algo, exclude_modules, packed_modules_mapping) self.quantized_layers = quantized_layers @@ -641,7 +641,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfig): return ModelOptFp4Config.get_min_capability() @classmethod - def from_config(cls, config: Dict[str, Any]) -> "ModelOptMixedPrecisionConfig": + def from_config(cls, config: Dict[str, Any]) -> ModelOptMixedPrecisionConfig: kv_cache_quant_algo = None exclude_modules = None quantized_layers = {} @@ -712,7 +712,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfig): nvfp4_config=nvfp4_config, ) - def apply_weight_name_mapper(self, hf_to_sglang_mapper: "WeightsMapper"): + def apply_weight_name_mapper(self, hf_to_sglang_mapper: WeightsMapper): super().apply_weight_name_mapper(hf_to_sglang_mapper) if self.quantized_layers: self.quantized_layers = hf_to_sglang_mapper.apply_dict( diff --git a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_int4_moe.py b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_int4_moe.py index fce46cb1f..95ce416d3 100644 --- a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_int4_moe.py +++ b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_int4_moe.py @@ -109,15 +109,15 @@ class ModelSlimW4A4Int4MoE(ModelSlimMoEScheme): self.kernel.process_weights_after_loading(layer) def create_moe_runner( - self, layer: torch.nn.Module, moe_runner_config: "MoeRunnerConfig" + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig ): self.moe_runner_config = moe_runner_config def apply_weights( self, layer, - dispatch_output: "StandardDispatchOutput", - ) -> "CombineInput": + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: return self.kernel.apply(layer, dispatch_output) def apply_without_routing_weights( diff --git a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a8_int8_moe.py b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a8_int8_moe.py index 4c3cd20f3..2c292d40c 100644 --- a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a8_int8_moe.py +++ b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a8_int8_moe.py @@ -186,15 +186,15 @@ class ModelSlimW4A8Int8MoE(ModelSlimMoEScheme): ) def create_moe_runner( - self, layer: torch.nn.Module, moe_runner_config: "MoeRunnerConfig" + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig ): self.moe_runner_config = moe_runner_config def apply_weights( self, layer, - dispatch_output: "StandardDispatchOutput", - ) -> "CombineInput": + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: # FIXME W4A8 without EP can give 0 accuracy return self.kernel.apply(layer, dispatch_output) diff --git a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w8a8_int8_moe.py b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w8a8_int8_moe.py index b226797f3..01b705fef 100644 --- a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w8a8_int8_moe.py +++ b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w8a8_int8_moe.py @@ -109,15 +109,15 @@ class ModelSlimW8A8Int8MoE(ModelSlimMoEScheme): self.kernel.process_weights_after_loading(layer) def create_moe_runner( - self, layer: torch.nn.Module, moe_runner_config: "MoeRunnerConfig" + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig ): self.moe_runner_config = moe_runner_config def apply_weights( self, layer, - dispatch_output: "StandardDispatchOutput", - ) -> "CombineInput": + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: return self.kernel.apply(layer, dispatch_output) def apply_without_routing_weights( diff --git a/python/sglang/srt/layers/quantization/mxfp4.py b/python/sglang/srt/layers/quantization/mxfp4.py index 1a2577b0d..3a555b87c 100644 --- a/python/sglang/srt/layers/quantization/mxfp4.py +++ b/python/sglang/srt/layers/quantization/mxfp4.py @@ -352,7 +352,7 @@ class Mxfp4Config(QuantizationConfig): def get_quant_method( self, layer: torch.nn.Module, prefix: str - ) -> Optional["QuantizeMethodBase"]: + ) -> Optional[QuantizeMethodBase]: from sglang.srt.layers.linear import LinearBase from sglang.srt.layers.moe.fused_moe_triton import FusedMoE diff --git a/python/sglang/srt/layers/quantization/mxfp4_flashinfer_cutlass_moe.py b/python/sglang/srt/layers/quantization/mxfp4_flashinfer_cutlass_moe.py index 31552cdfe..19d2b1879 100644 --- a/python/sglang/srt/layers/quantization/mxfp4_flashinfer_cutlass_moe.py +++ b/python/sglang/srt/layers/quantization/mxfp4_flashinfer_cutlass_moe.py @@ -215,8 +215,8 @@ class Mxfp4FlashinferCutlassMoEMethod: def apply( self, layer: Module, - dispatch_output: "DispatchOutput", - ) -> "CombineInput": + dispatch_output: DispatchOutput, + ) -> CombineInput: from sglang.srt.layers.moe.moe_runner.flashinfer_mxfp4 import ( FlashInferMxfp4CutlassMoeQuantInfo, ) diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index b21972065..f36e9b6d3 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/python/sglang/srt/layers/quantization/unquant.py @@ -673,7 +673,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): def forward_npu( self, layer: torch.nn.Module, - dispatch_output: "DispatchOutput", + dispatch_output: DispatchOutput, ) -> CombineInput: from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput @@ -767,7 +767,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): def _forward_npu_deepep( self, layer: torch.nn.Module, - dispatch_output: "DispatchOutput", + dispatch_output: DispatchOutput, ) -> CombineInput: from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import ( npu_fused_moe_without_routing_weights_bf16, diff --git a/python/sglang/srt/lora/deepseek_mla_correction.py b/python/sglang/srt/lora/deepseek_mla_correction.py index 8bc564b99..f81d5bcbf 100644 --- a/python/sglang/srt/lora/deepseek_mla_correction.py +++ b/python/sglang/srt/lora/deepseek_mla_correction.py @@ -31,7 +31,7 @@ if TYPE_CHECKING: from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA -def is_kv_b_lora_active(attn_module: "DeepseekV2AttentionMLA") -> bool: +def is_kv_b_lora_active(attn_module: DeepseekV2AttentionMLA) -> bool: """Cheap precondition check used at call sites in the attention forward to skip the entire LoRA-correction path when no ``kv_b_proj`` adapter is wrapped on this module (the common case).""" @@ -39,8 +39,8 @@ def is_kv_b_lora_active(attn_module: "DeepseekV2AttentionMLA") -> bool: def _get_state( - attn_module: "DeepseekV2AttentionMLA", -) -> Optional[Tuple[torch.Tensor, torch.Tensor, "LoRABatchInfo"]]: + attn_module: DeepseekV2AttentionMLA, +) -> Optional[Tuple[torch.Tensor, torch.Tensor, LoRABatchInfo]]: if not is_kv_b_lora_active(attn_module): return None if not hasattr(attn_module.kv_b_proj, "A_buffer"): @@ -61,7 +61,7 @@ def _get_state( def apply_q_correction( - attn_module: "DeepseekV2AttentionMLA", + attn_module: DeepseekV2AttentionMLA, q_nope: torch.Tensor, q_nope_out: torch.Tensor, ) -> torch.Tensor: @@ -86,7 +86,7 @@ def apply_q_correction( def apply_v_correction( - attn_module: "DeepseekV2AttentionMLA", + attn_module: DeepseekV2AttentionMLA, attn_output: torch.Tensor, attn_bmm_flat: torch.Tensor, ) -> torch.Tensor: diff --git a/python/sglang/srt/lora/trtllm_lora_temp/deepseek_mla_correction.py b/python/sglang/srt/lora/trtllm_lora_temp/deepseek_mla_correction.py index 36d2c5b19..b2e066e3a 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/deepseek_mla_correction.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/deepseek_mla_correction.py @@ -42,7 +42,7 @@ if TYPE_CHECKING: from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA -def is_kv_b_lora_active(attn_module: "DeepseekV2AttentionMLA") -> bool: +def is_kv_b_lora_active(attn_module: DeepseekV2AttentionMLA) -> bool: """Cheap precondition check used at call sites in the attention forward to skip the entire LoRA-correction path when no ``kv_b_proj`` adapter is wrapped on this module (the common case).""" @@ -50,8 +50,8 @@ def is_kv_b_lora_active(attn_module: "DeepseekV2AttentionMLA") -> bool: def _get_state( - attn_module: "DeepseekV2AttentionMLA", -) -> Optional[Tuple[torch.Tensor, torch.Tensor, "LoRABatchInfo"]]: + attn_module: DeepseekV2AttentionMLA, +) -> Optional[Tuple[torch.Tensor, torch.Tensor, LoRABatchInfo]]: if not is_kv_b_lora_active(attn_module): return None if not hasattr(attn_module.kv_b_proj, "A_buffer"): @@ -76,7 +76,7 @@ def _get_state( def apply_q_correction( - attn_module: "DeepseekV2AttentionMLA", + attn_module: DeepseekV2AttentionMLA, q_nope: torch.Tensor, q_nope_out: torch.Tensor, ) -> torch.Tensor: @@ -101,7 +101,7 @@ def apply_q_correction( def apply_v_correction( - attn_module: "DeepseekV2AttentionMLA", + attn_module: DeepseekV2AttentionMLA, attn_output: torch.Tensor, attn_bmm_flat: torch.Tensor, ) -> torch.Tensor: diff --git a/python/sglang/srt/lora/trtllm_lora_temp/lora_dispatch.py b/python/sglang/srt/lora/trtllm_lora_temp/lora_dispatch.py index c093b9090..5cd3b0c66 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/lora_dispatch.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/lora_dispatch.py @@ -39,11 +39,11 @@ if TYPE_CHECKING: def fused_experts_none_to_experimental_sgl_trtllm_fp8_lora( - dispatch_output: "StandardDispatchOutput", - quant_info: "FlashInferTrtllmFp8MoeQuantInfo", - runner_config: "MoeRunnerConfig", + dispatch_output: StandardDispatchOutput, + quant_info: FlashInferTrtllmFp8MoeQuantInfo, + runner_config: MoeRunnerConfig, lora_info, -) -> "StandardCombineInput": +) -> StandardCombineInput: from flashinfer.fused_moe import Fp8QuantizationType from sglang.jit_kernel.trtllm_lora_temp import ( @@ -302,11 +302,11 @@ def fused_experts_none_to_experimental_sgl_trtllm_fp8_lora( def fused_experts_none_to_experimental_sgl_trtllm_fp4_lora( - dispatch_output: "StandardDispatchOutput", - quant_info: "FlashInferTrtllmFp4MoeQuantInfo", - runner_config: "MoeRunnerConfig", + dispatch_output: StandardDispatchOutput, + quant_info: FlashInferTrtllmFp4MoeQuantInfo, + runner_config: MoeRunnerConfig, lora_info, -) -> "StandardCombineInput": +) -> StandardCombineInput: """NVFP4 sibling of ``fused_experts_none_to_experimental_sgl_trtllm_fp8_lora``. Decomposed (unfused-activation) MoE-LoRA: routing -> gather -> gate_up grouped diff --git a/python/sglang/srt/lora/trtllm_lora_temp/lora_layer.py b/python/sglang/srt/lora/trtllm_lora_temp/lora_layer.py index 5d568b927..b1c4c0840 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/lora_layer.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/lora_layer.py @@ -125,7 +125,7 @@ def init_experimental_sgl_trtllm_lora(layer, base_layer) -> None: def dispatch_experimental_sgl_trtllm_lora( dispatch_output, quant_info, base_layer, lora_info -) -> "StandardCombineInput": +) -> StandardCombineInput: """Call the trtllm fused-experts LoRA function for a single layer. Looked up at call time so the install-time monkey-patch in diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index ba1e27479..8ab58a755 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -209,7 +209,7 @@ class StorageOperation: self.id = StorageOperation.counter StorageOperation.counter += 1 - def __lt__(self, other: "StorageOperation"): + def __lt__(self, other: StorageOperation): return self.id < other.id diff --git a/python/sglang/srt/managers/overlap_utils.py b/python/sglang/srt/managers/overlap_utils.py index 12eb9fb2d..7f1221ccb 100644 --- a/python/sglang/srt/managers/overlap_utils.py +++ b/python/sglang/srt/managers/overlap_utils.py @@ -20,8 +20,8 @@ if TYPE_CHECKING: def decide_needs_cpu_seq_lens( - server_args: "ServerArgs", - attn_backends: Sequence["AttentionBackend"], + server_args: ServerArgs, + attn_backends: Sequence[AttentionBackend], ) -> bool: """Whether FutureMap must publish seq_lens_cpu / sum. diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 933015b2c..c1bab7d78 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -415,7 +415,7 @@ class MultimodalProcessorOutput: token_type_ids: Optional[torch.Tensor] = None @staticmethod - def from_dict(d: dict) -> "MultimodalProcessorOutput": + def from_dict(d: dict) -> MultimodalProcessorOutput: return MultimodalProcessorOutput( mm_items=d["mm_items"], input_ids=d.get("input_ids"), @@ -498,7 +498,7 @@ class MultimodalInputs: item.feature = None @staticmethod - def from_processor_output(obj: "MultimodalProcessorOutput"): + def from_processor_output(obj: MultimodalProcessorOutput): mm_items = obj.mm_items assert isinstance(mm_items, list) mm_items = [item for item in mm_items if item.is_valid()] @@ -2199,7 +2199,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): def _mamba_radix_cache_v2_req_prepare_for_extend( self, req: Req, - ) -> "_MambaRadixCacheV2TrackEntry": + ) -> _MambaRadixCacheV2TrackEntry: mamba_cache_chunk_size = get_global_server_args().mamba_cache_chunk_size def _force_track_h(i: int) -> int: @@ -2308,7 +2308,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # For split prefill, we need to set the forward mode to SPLIT_PREFILL self.forward_mode = ForwardMode.SPLIT_PREFILL - def mix_with_running(self, running_batch: "ScheduleBatch"): + def mix_with_running(self, running_batch: ScheduleBatch): self.forward_mode = ForwardMode.MIXED running_bs = running_batch.batch_size() @@ -2717,7 +2717,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): has_been_filtered=has_been_filtered, ) - def merge_batch(self, other: "ScheduleBatch"): + def merge_batch(self, other: ScheduleBatch): # Penalizer orchestrator must be merged before Batch.reqs is merged. This is because # orchestrator.merge() depends on Batch.reqs during preparation of each penalizers, so it # needs to be called with pre-merged Batch.reqs. diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py index ed8b586d7..c80030632 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py @@ -61,22 +61,22 @@ logger = logging.getLogger(__name__) @dataclass(kw_only=True, slots=True, frozen=True) class SchedulerBatchResultProcessor: is_generation: bool - disaggregation_mode: "DisaggregationMode" + disaggregation_mode: DisaggregationMode enable_overlap: bool enable_overlap_mlx: bool - server_args: "ServerArgs" - model_config: "ModelConfig" - token_to_kv_pool_allocator: "BaseTokenToKVPoolAllocator" - tree_cache: "BasePrefixCache" - hisparse_coordinator: Optional["HiSparseCoordinator"] - req_to_token_pool: "ReqToTokenPool" - decode_offload_manager: Optional["DecodeKVCacheOffloadManager"] - metrics_collector: "SchedulerMetricsCollector" - metrics_reporter: "SchedulerMetricsReporter" - draft_worker: "BaseTpWorker" - model_worker: "BaseTpWorker" - logprob_result_processor: "SchedulerLogprobResultProcessor" - output_streamer: "SchedulerOutputStreamer" + server_args: ServerArgs + model_config: ModelConfig + token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator + tree_cache: BasePrefixCache + hisparse_coordinator: Optional[HiSparseCoordinator] + req_to_token_pool: ReqToTokenPool + decode_offload_manager: Optional[DecodeKVCacheOffloadManager] + metrics_collector: SchedulerMetricsCollector + metrics_reporter: SchedulerMetricsReporter + draft_worker: BaseTpWorker + model_worker: BaseTpWorker + logprob_result_processor: SchedulerLogprobResultProcessor + output_streamer: SchedulerOutputStreamer abort_request: Callable def process_batch_result_prebuilt(self, batch: ScheduleBatch): diff --git a/python/sglang/srt/managers/scheduler_components/dp_attn.py b/python/sglang/srt/managers/scheduler_components/dp_attn.py index 0ff6dafc4..b473a11ee 100644 --- a/python/sglang/srt/managers/scheduler_components/dp_attn.py +++ b/python/sglang/srt/managers/scheduler_components/dp_attn.py @@ -257,7 +257,7 @@ def prepare_mlp_sync_batch_raw( @dataclass(kw_only=True, slots=True, frozen=True) class SchedulerDPAttnAdapter: - tp_group: "GroupCoordinator" + tp_group: GroupCoordinator req_to_token_pool: ReqToTokenPool token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator tree_cache: BasePrefixCache diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index 4857f5498..33013a424 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/python/sglang/srt/managers/scheduler_components/invariant_checker.py @@ -294,7 +294,7 @@ class SchedulerInvariantChecker: def create_scheduler_watchdog( - scheduler: "Scheduler", watchdog_timeout: float, soft: bool = False + scheduler: Scheduler, watchdog_timeout: float, soft: bool = False ) -> WatchdogRaw: def dump_info() -> str: if scheduler.is_initializing: diff --git a/python/sglang/srt/managers/scheduler_components/kv_events_publisher.py b/python/sglang/srt/managers/scheduler_components/kv_events_publisher.py index 85531e49e..b3301cd1d 100644 --- a/python/sglang/srt/managers/scheduler_components/kv_events_publisher.py +++ b/python/sglang/srt/managers/scheduler_components/kv_events_publisher.py @@ -40,13 +40,13 @@ class KvMetrics: @dataclass(kw_only=True, slots=True) class SchedulerKvEventsPublisher: kv_events_config: Optional[str] - ps: "ParallelState" + ps: ParallelState attn_tp_rank: int attn_cp_rank: int attn_dp_rank: int dp_rank: Optional[int] - tree_cache: "BasePrefixCache" - send_metrics_from_scheduler: Optional["zmq.Socket"] + tree_cache: BasePrefixCache + send_metrics_from_scheduler: Optional[zmq.Socket] max_running_requests: int max_total_num_tokens: int get_stats: Callable diff --git a/python/sglang/srt/managers/scheduler_components/load_inquirer.py b/python/sglang/srt/managers/scheduler_components/load_inquirer.py index 3f10d7eda..e2b83f6bc 100644 --- a/python/sglang/srt/managers/scheduler_components/load_inquirer.py +++ b/python/sglang/srt/managers/scheduler_components/load_inquirer.py @@ -32,15 +32,15 @@ logger = logging.getLogger(__name__) @dataclass(kw_only=True, slots=True, frozen=True) class SchedulerLoadInquirer: - disaggregation_mode: "DisaggregationMode" - ps: "ParallelState" - server_args: "ServerArgs" + disaggregation_mode: DisaggregationMode + ps: ParallelState + server_args: ServerArgs max_total_num_tokens: int max_running_requests: int - pool_stats_observer: "SchedulerPoolStatsObserver" - tp_worker: "BaseTpWorker" - token_to_kv_pool_allocator: "BaseTokenToKVPoolAllocator" - spec_algorithm: "SpeculativeAlgorithm" + pool_stats_observer: SchedulerPoolStatsObserver + tp_worker: BaseTpWorker + token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator + spec_algorithm: SpeculativeAlgorithm get_running_batch: Callable get_waiting_queue: Callable get_stats: Callable diff --git a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py index 566461116..d7bfcbc02 100644 --- a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py +++ b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py @@ -88,7 +88,7 @@ class PrefillStats: @dataclass(kw_only=True) class SchedulerMetricsReporter: - scheduler: "Scheduler" + scheduler: Scheduler tp_rank: int pp_rank: int dp_rank: Optional[int] diff --git a/python/sglang/srt/managers/scheduler_components/new_token_ratio_tracker.py b/python/sglang/srt/managers/scheduler_components/new_token_ratio_tracker.py index 81f7e7801..65e899953 100644 --- a/python/sglang/srt/managers/scheduler_components/new_token_ratio_tracker.py +++ b/python/sglang/srt/managers/scheduler_components/new_token_ratio_tracker.py @@ -18,7 +18,7 @@ class NewTokenRatioTracker: current: float @classmethod - def from_server_args(cls, server_args: ServerArgs) -> "NewTokenRatioTracker": + def from_server_args(cls, server_args: ServerArgs) -> NewTokenRatioTracker: init = min( envs.SGLANG_INIT_NEW_TOKEN_RATIO.get() * server_args.schedule_conservativeness, diff --git a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py index b700bd525..5e811ad5a 100644 --- a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py +++ b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py @@ -140,9 +140,9 @@ class PoolStats: @dataclass(kw_only=True, slots=True, frozen=True) class SchedulerPoolStatsObserver: - tree_cache: "BasePrefixCache" - token_to_kv_pool_allocator: "BaseTokenToKVPoolAllocator" - req_to_token_pool: "ReqToTokenPool" + tree_cache: BasePrefixCache + token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator + req_to_token_pool: ReqToTokenPool session_controller: Any hisparse_coordinator: Any is_hybrid_swa: bool diff --git a/python/sglang/srt/managers/scheduler_components/request_receiver.py b/python/sglang/srt/managers/scheduler_components/request_receiver.py index 608aece5f..3f10a2790 100644 --- a/python/sglang/srt/managers/scheduler_components/request_receiver.py +++ b/python/sglang/srt/managers/scheduler_components/request_receiver.py @@ -42,12 +42,12 @@ if TYPE_CHECKING: @dataclass(kw_only=True, slots=True, frozen=True) class SchedulerRequestReceiver: - recv_from_tokenizer: Union[zmq.Socket, "ScriptedTokenizerRecvProxy"] + recv_from_tokenizer: Union[zmq.Socket, ScriptedTokenizerRecvProxy] recv_from_rpc: Optional[zmq.Socket] recv_skipper: Any input_blocker: Any mm_receiver: Any - ps: "ParallelState" + ps: ParallelState tp_group: Any tp_cpu_group: Any attn_tp_group: Any @@ -55,12 +55,12 @@ class SchedulerRequestReceiver: attn_cp_group: Any attn_cp_cpu_group: Any world_group: Any - server_args: "ServerArgs" - model_config: "ModelConfig" + server_args: ServerArgs + model_config: ModelConfig max_recv_per_poll: int stream_output: Callable[..., None] get_last_forward_mode: Callable[[], Any] - scripted_scheduler_hook: Optional["ScriptedSchedulerHook"] = None + scripted_scheduler_hook: Optional[ScriptedSchedulerHook] = None def recv_limit_reached(self, num_recv_reqs: int) -> bool: if self.max_recv_per_poll < 0: diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index bd9184408..93b669389 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -67,7 +67,7 @@ class BaseTpWorker(ABC): @property @abstractmethod - def model_runner(self) -> "ModelRunner": + def model_runner(self) -> ModelRunner: pass @property @@ -400,7 +400,7 @@ class TpModelWorker(BaseTpWorker): self.dllm_algorithm = None @property - def model_runner(self) -> "ModelRunner": + def model_runner(self) -> ModelRunner: return self._model_runner def register_hicache_layer_transfer_counter(self, counter: LayerDoneCounter): diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index f0bbca6c9..4ef141822 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -111,7 +111,7 @@ class IncLockRefResult: default_factory=dict ) - def to_dec_params(self) -> "DecLockRefParams": + def to_dec_params(self) -> DecLockRefParams: """Convert to the corresponding DecLockRefParams for dec_lock_ref.""" return DecLockRefParams( swa_uuid_for_lock=self.swa_uuid_for_lock, @@ -189,7 +189,7 @@ class MatchResult(NamedTuple): cache_protected_len: Optional[int] = None -def zero_match_result(tree_cache, match_result: "MatchResult") -> "MatchResult": +def zero_match_result(tree_cache, match_result: MatchResult) -> MatchResult: if tree_cache.is_chunk_cache(): # Chunk caches' match_prefix already returns a miss; no root_node to walk back to. return match_result diff --git a/python/sglang/srt/mem_cache/evict_policy.py b/python/sglang/srt/mem_cache/evict_policy.py index 30ab1983f..dad1625fe 100644 --- a/python/sglang/srt/mem_cache/evict_policy.py +++ b/python/sglang/srt/mem_cache/evict_policy.py @@ -9,39 +9,39 @@ if TYPE_CHECKING: class EvictionStrategy(ABC): @abstractmethod - def get_priority(self, node: "TreeNode") -> Union[float, Tuple]: + def get_priority(self, node: TreeNode) -> Union[float, Tuple]: pass class LRUStrategy(EvictionStrategy): - def get_priority(self, node: "TreeNode") -> float: + def get_priority(self, node: TreeNode) -> float: return node.last_access_time class LFUStrategy(EvictionStrategy): - def get_priority(self, node: "TreeNode") -> Tuple[int, float]: + def get_priority(self, node: TreeNode) -> Tuple[int, float]: return (node.hit_count, node.last_access_time) class FIFOStrategy(EvictionStrategy): - def get_priority(self, node: "TreeNode") -> float: + def get_priority(self, node: TreeNode) -> float: return node.creation_time class MRUStrategy(EvictionStrategy): - def get_priority(self, node: "TreeNode") -> float: + def get_priority(self, node: TreeNode) -> float: return -node.last_access_time class FILOStrategy(EvictionStrategy): - def get_priority(self, node: "TreeNode") -> float: + def get_priority(self, node: TreeNode) -> float: return -node.creation_time class PriorityStrategy(EvictionStrategy): """Priority-aware eviction: lower priority values evicted first, then LRU within same priority.""" - def get_priority(self, node: "TreeNode") -> Tuple[int, float]: + def get_priority(self, node: TreeNode) -> Tuple[int, float]: # Return (priority, last_access_time) so lower priority nodes are evicted first return (node.priority, node.last_access_time) @@ -50,7 +50,7 @@ class SLRUStrategy(EvictionStrategy): def __init__(self, protected_threshold: int = 2): self.protected_threshold = protected_threshold - def get_priority(self, node: "TreeNode") -> Tuple[int, float]: + def get_priority(self, node: TreeNode) -> Tuple[int, float]: # Priority Logic: # Smaller value = Evicted earlier. # diff --git a/python/sglang/srt/mem_cache/hicache_storage.py b/python/sglang/srt/mem_cache/hicache_storage.py index 5feb1754f..193be40d2 100644 --- a/python/sglang/srt/mem_cache/hicache_storage.py +++ b/python/sglang/srt/mem_cache/hicache_storage.py @@ -123,7 +123,7 @@ class PoolTransferResult: extra_pool_hit_pages: dict[str, int] @classmethod - def empty(cls) -> "PoolTransferResult": + def empty(cls) -> PoolTransferResult: return cls(0, {}) def update_kv_hit_pages(self, kv_hit_pages: int) -> None: @@ -188,7 +188,7 @@ class HiCacheStorage(ABC): def batch_get_v2( self, transfers: List[PoolTransfer], - extra_info: Optional["HiCacheStorageExtraInfo"] = None, + extra_info: Optional[HiCacheStorageExtraInfo] = None, ) -> dict[str, List[bool]]: """Read data from storage into host memory for each PoolTransfer. @@ -199,7 +199,7 @@ class HiCacheStorage(ABC): def batch_set_v2( self, transfers: List[PoolTransfer], - extra_info: Optional["HiCacheStorageExtraInfo"] = None, + extra_info: Optional[HiCacheStorageExtraInfo] = None, ) -> dict[str, List[bool]]: """Write data from host memory to storage for each PoolTransfer. @@ -589,14 +589,14 @@ class HiCacheFile(HiCacheStorage): def batch_get_v2( self, transfers: List[PoolTransfer], - extra_info: Optional["HiCacheStorageExtraInfo"] = None, + extra_info: Optional[HiCacheStorageExtraInfo] = None, ) -> dict[str, List[bool]]: return self._batch_io_v2(transfers, self._read_page) def batch_set_v2( self, transfers: List[PoolTransfer], - extra_info: Optional["HiCacheStorageExtraInfo"] = None, + extra_info: Optional[HiCacheStorageExtraInfo] = None, ) -> dict[str, List[bool]]: return self._batch_io_v2(transfers, self._write_page) diff --git a/python/sglang/srt/mem_cache/kv_cache_builder.py b/python/sglang/srt/mem_cache/kv_cache_builder.py index 09531eb96..b4d883d4c 100644 --- a/python/sglang/srt/mem_cache/kv_cache_builder.py +++ b/python/sglang/srt/mem_cache/kv_cache_builder.py @@ -45,7 +45,7 @@ if TYPE_CHECKING: def get_draft_kv_pool( *, - draft_worker: "BaseTpWorker", + draft_worker: BaseTpWorker, spec_algorithm: SpeculativeAlgorithm, server_args: ServerArgs, ): @@ -71,8 +71,8 @@ def get_draft_kv_pool( def maybe_register_hicache_draft( *, - tree_cache: "BasePrefixCache", - draft_worker: "BaseTpWorker", + tree_cache: BasePrefixCache, + draft_worker: BaseTpWorker, spec_algorithm: SpeculativeAlgorithm, server_args: ServerArgs, enable_hierarchical_cache: bool, @@ -129,21 +129,21 @@ def maybe_register_hicache_draft( def build_kv_cache( *, - server_args: "ServerArgs", - model_config: "ModelConfig", - tp_worker: "BaseTpWorker", + server_args: ServerArgs, + model_config: ModelConfig, + tp_worker: BaseTpWorker, page_size: int, - spec_algorithm: "SpeculativeAlgorithm", - attn_tp_cpu_group: "ProcessGroup", - tp_cpu_group: "ProcessGroup", - attn_cp_cpu_group: "ProcessGroup", + spec_algorithm: SpeculativeAlgorithm, + attn_tp_cpu_group: ProcessGroup, + tp_cpu_group: ProcessGroup, + attn_cp_cpu_group: ProcessGroup, enable_metrics: bool, enable_kv_cache_events: bool, - ps: "ParallelState", - tp_group: "GroupCoordinator", - pp_group: "GroupCoordinator", + ps: ParallelState, + tp_group: GroupCoordinator, + pp_group: GroupCoordinator, enable_hierarchical_cache: bool, -) -> "KVCacheBuildResult": +) -> KVCacheBuildResult: sliding_window_size: Optional[int] = None full_tokens_per_layer: Optional[int] = None swa_tokens_per_layer: Optional[int] = None diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index b0d80df19..b4c55a6c3 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -146,12 +146,12 @@ class TreeNode: return None return self.hash_value[-1] - def get_prefix_hash_values(self, node: "TreeNode") -> List[str]: + def get_prefix_hash_values(self, node: TreeNode) -> List[str]: if node is None or node.hash_value is None: return [] return node.get_prefix_hash_values(node.parent) + node.hash_value - def __lt__(self, other: "TreeNode"): + def __lt__(self, other: TreeNode): return self.last_access_time < other.last_access_time @@ -319,7 +319,7 @@ class LRUList: return False return node.id in self.cache - def pretty_print(self, tree_cache: Optional["MambaRadixCache"] = None): + def pretty_print(self, tree_cache: Optional[MambaRadixCache] = None): """ Pretty print the lru list """ @@ -358,7 +358,7 @@ class LRUList: return evictable_size # Note: this is expensive, only use for debug or idle check - def sanity_check(self, tree_cache: "MambaRadixCache"): + def sanity_check(self, tree_cache: MambaRadixCache): """ Check if the lru list is valid by rebuilding the lru list from the tree, heapifying it, and checking if the lru list is valid. diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 704892e24..27a0d237f 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -626,14 +626,12 @@ class HybridReqToTokenPool(ReqToTokenPool): ) ) - def register_layer_transfer_counter( - self, layer_transfer_counter: "LayerDoneCounter" - ): + def register_layer_transfer_counter(self, layer_transfer_counter: LayerDoneCounter): self.layer_transfer_counter = layer_transfer_counter # For chunk prefill req, we do not need to allocate mamba cache, # We could use allocated mamba cache instead. - def alloc(self, reqs: List["Req"]) -> Optional[List[int]]: + def alloc(self, reqs: List[Req]) -> Optional[List[int]]: select_index = super().alloc(reqs) if select_index is None: return None @@ -695,7 +693,7 @@ class HybridReqToTokenPool(ReqToTokenPool): else: return mamba_next_track_idx - def get_mamba_ping_pong_keep_idx(self, req: "Req") -> int: + def get_mamba_ping_pong_keep_idx(self, req: Req) -> int: """Return the ping-pong index holding the most recent tracked state. In lazy mode the valid state stays at next_track_idx (no eager swap). @@ -705,7 +703,7 @@ class HybridReqToTokenPool(ReqToTokenPool): return req.mamba_next_track_idx return self.get_mamba_ping_pong_other_idx(req.mamba_next_track_idx) - def _alloc_ping_pong_buffer(self, req: "Req"): + def _alloc_ping_pong_buffer(self, req: Req): """Allocate the ping-pong track buffer for a new request. Lazy mode allocates 1 slot with the second set to -1 (allocated @@ -731,7 +729,7 @@ class HybridReqToTokenPool(ReqToTokenPool): req.mamba_ping_pong_track_buffer = buf req.mamba_next_track_idx = 0 - def set_mamba_ping_pong_slot(self, req: "Req", idx: int, value): + def set_mamba_ping_pong_slot(self, req: Req, idx: int, value): """Update a ping-pong slot value and sync the device-side mapping. The req holds the authoritative buffer; this keeps the @@ -744,7 +742,7 @@ class HybridReqToTokenPool(ReqToTokenPool): ) def donate_mamba_ping_pong_slot( - self, req: "Req", new_slot: torch.Tensor + self, req: Req, new_slot: torch.Tensor ) -> torch.Tensor: """Donate the tracked-state ping-pong slot to the radix cache. @@ -767,7 +765,7 @@ class HybridReqToTokenPool(ReqToTokenPool): return mamba_value_donated def free_mamba_cache( - self, req: "Req", mamba_ping_pong_track_buffer_to_keep: Optional[int] = None + self, req: Req, mamba_ping_pong_track_buffer_to_keep: Optional[int] = None ): mamba_index = req.mamba_pool_idx assert mamba_index is not None, "double free? mamba_index is None" @@ -1860,9 +1858,7 @@ class HybridLinearKVPool(KVCache): ) return self.full_attention_layer_id_mapping[layer_id] - def register_layer_transfer_counter( - self, layer_transfer_counter: "LayerDoneCounter" - ): + def register_layer_transfer_counter(self, layer_transfer_counter: LayerDoneCounter): self.layer_transfer_counter = layer_transfer_counter # The layer-wise wait logic is executed at the Hybrid LinearPool level; # no additional wait is needed in the full_kv_pool diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index bd6adb6e3..3b103582c 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -86,7 +86,7 @@ class RadixKey: else: yield from self.token_ids - def __getitem__(self, idx: Union[int, slice]) -> "RadixKey": + def __getitem__(self, idx: Union[int, slice]) -> RadixKey: # Normalize int -> 1-element slice so the rest handles one shape. if isinstance(idx, int): if idx < 0: @@ -109,7 +109,7 @@ class RadixKey: preview = self.token_ids[:10] return f"RadixKey(extra_key={self.extra_key!r}, token_ids={preview}{'...' if len(self.token_ids) > 10 else ''}, is_bigram={self.is_bigram})" - def page_aligned(self, page_size: int) -> "RadixKey": + def page_aligned(self, page_size: int) -> RadixKey: if page_size == 1: return self aligned_len = len(self) // page_size * page_size @@ -119,7 +119,7 @@ class RadixKey: self, is_eagle: bool, value: Optional[torch.Tensor] = None, - ) -> Tuple["RadixKey", Optional[torch.Tensor]]: + ) -> Tuple[RadixKey, Optional[torch.Tensor]]: # O(1): flip the bigram flag instead of materializing a tuple list. # value is paired with raw tokens and gets truncated to the bigram count. if is_eagle and not self.is_bigram: @@ -128,14 +128,14 @@ class RadixKey: value = value[: len(self)] return self, value - def _check_compatible(self, other: "RadixKey") -> None: + def _check_compatible(self, other: RadixKey) -> None: if self.extra_key != other.extra_key: raise ValueError( f"RadixKey operations require matching extra_key, but got " f"{self.extra_key=} != {other.extra_key=}" ) - def match(self, other: "RadixKey", page_size: int = 1) -> int: + def match(self, other: RadixKey, page_size: int = 1) -> int: """Logical-unit prefix length shared with ``other``. Result is rounded down to ``page_size``.""" self._check_compatible(other) t0, t1 = self.token_ids, other.token_ids @@ -257,7 +257,7 @@ class TreeNode: return node.get_prefix_hash_values(node.parent) + node.hash_value - def __lt__(self, other: "TreeNode"): + def __lt__(self, other: TreeNode): return self.last_access_time < other.last_access_time diff --git a/python/sglang/srt/mem_cache/registry.py b/python/sglang/srt/mem_cache/registry.py index 6901521aa..9e6e7fd7b 100644 --- a/python/sglang/srt/mem_cache/registry.py +++ b/python/sglang/srt/mem_cache/registry.py @@ -141,7 +141,7 @@ def default_radix_cache_factory(ctx: TreeCacheBuildContext) -> BasePrefixCache: def _create_unified_radix_cache( ctx: TreeCacheBuildContext, - server_args: "ServerArgs", + server_args: ServerArgs, params: CacheInitParams, ) -> BasePrefixCache: """Initialize a UnifiedRadixCache with proper components and optional HiCache.""" diff --git a/python/sglang/srt/mem_cache/storage/file/lru_file_evictor.py b/python/sglang/srt/mem_cache/storage/file/lru_file_evictor.py index c782bb4e6..ac5484e1e 100644 --- a/python/sglang/srt/mem_cache/storage/file/lru_file_evictor.py +++ b/python/sglang/srt/mem_cache/storage/file/lru_file_evictor.py @@ -81,7 +81,7 @@ class LRUFileEvictor: self._is_storage_owner = (not is_mla_model) or (tp_rank == 0) # suffixed_key -> file size in bytes; oldest at front. - self._lru: "OrderedDict[str, int]" = OrderedDict() + self._lru: OrderedDict[str, int] = OrderedDict() self._pending_writes: Set[str] = set() self._total_bytes: int = 0 self._lock = threading.Lock() diff --git a/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py b/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py index 88a640041..733ea74d5 100644 --- a/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py +++ b/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py @@ -100,7 +100,7 @@ class LMCRadixCache(RadixCache): def __init__( self, params: CacheInitParams, - model_config: Optional["ModelConfig"] = None, + model_config: Optional[ModelConfig] = None, tp_size: int = 1, rank: int = 0, tp_group: Optional[torch.distributed.ProcessGroup] = None, diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index d4f6afa4c..42a389fbc 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -101,7 +101,7 @@ class TreeNode: def backuped(self): return self.host_value is not None - def __lt__(self, other: "TreeNode"): + def __lt__(self, other: TreeNode): return self.last_access_time < other.last_access_time @@ -288,7 +288,7 @@ class LRUList: return evictable_size # Note: this is expensive, only use for debug or idle check - def sanity_check(self, tree_cache: "SWARadixCache"): + def sanity_check(self, tree_cache: SWARadixCache): """ Check if the lru list is valid by rebuilding the lru list from the tree, heapifying it, and checking if the lru list is valid. diff --git a/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py b/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py index 228f65a6c..99c8451be 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py @@ -681,7 +681,7 @@ class SWAComponent(TreeComponent): ) def _attach_swa_host_value( - self, node: "UnifiedTreeNode", host_indices: torch.Tensor + self, node: UnifiedTreeNode, host_indices: torch.Tensor ) -> None: """Write host_indices into node's SWA host_value and refresh tree state.""" ct = self.component_type diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 0ffa4faa0..baa239257 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -2794,7 +2794,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): def _check_lru_linked_list( self, - lru: "UnifiedLRUList", + lru: UnifiedLRUList, ct: ComponentType, label: str, errors: list[str], diff --git a/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py b/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py index 6cac59ea0..a2703dec8 100644 --- a/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py +++ b/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py @@ -162,12 +162,12 @@ class GraphSlot: pad_value: Optional[Any] = None enabled: bool = True copy_from_fb: bool = True - post_fill: Optional[ - Callable[[torch.Tensor, "ForwardBatch", "FillContext"], None] - ] = None + post_fill: Optional[Callable[[torch.Tensor, ForwardBatch, FillContext], None]] = ( + None + ) slice_fn: Optional[Callable[[torch.Tensor, int], torch.Tensor]] = None source_fn: Optional[ - Callable[["ForwardBatch", "FillContext"], Optional[torch.Tensor]] + Callable[[ForwardBatch, FillContext], Optional[torch.Tensor]] ] = None # runtime @@ -365,7 +365,7 @@ class CudaGraphBufferRegistry: def fill_from( self, - forward_batch: "ForwardBatch", + forward_batch: ForwardBatch, *, raw_bs: int, padded_bs: int, @@ -460,8 +460,8 @@ class CudaGraphBufferRegistry: *, padded_bs: int, padded_num_tokens: int, - forward_batch_template: "ForwardBatch", - ) -> "ForwardBatch": + forward_batch_template: ForwardBatch, + ) -> ForwardBatch: """Return a FB view (``dataclasses.replace`` of ``forward_batch_template``) whose slot fields are buffer views and whose non-slot fields are carried from the template. A plain copy slot whose FB field is ``None`` this iter diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index d761a1b16..4117bc87b 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -353,7 +353,7 @@ class ModelRunnerOutput: @dataclass class _EagerBufferRegistry: # Lazily-built eager input-buffer registry plus the capacity it was sized to. - registry: Optional["CudaGraphBufferRegistry"] = None + registry: Optional[CudaGraphBufferRegistry] = None max_bs: int = 0 max_num_tokens: int = 0 @@ -1976,7 +1976,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): def update_weights_from_tensor( self, - named_tensors: List[Tuple[str, Union[torch.Tensor, "LocalSerializedTensor"]]], + named_tensors: List[Tuple[str, Union[torch.Tensor, LocalSerializedTensor]]], load_format: Optional[str] = None, ): monkey_patch_torch_reductions() @@ -3107,8 +3107,8 @@ class ModelRunner(ModelRunnerKVCacheMixin): cache: _EagerBufferRegistry, raw_bs: int, raw_num_tokens: int, - build: Callable[[int, int], "CudaGraphBufferRegistry"], - ) -> "CudaGraphBufferRegistry": + build: Callable[[int, int], CudaGraphBufferRegistry], + ) -> CudaGraphBufferRegistry: # Built on first use and grown (next power of two) when a batch exceeds # the current capacity. if ( @@ -3126,7 +3126,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): def _ensure_eager_decode_registry( self, raw_bs: int, raw_num_tokens: int - ) -> "CudaGraphBufferRegistry": + ) -> CudaGraphBufferRegistry: is_encoder_decoder = self.model_config.is_encoder_decoder return self._ensure_eager_registry( self._eager_decode_registry, @@ -3162,7 +3162,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): def _ensure_eager_prefill_registry( self, raw_bs: int, raw_num_tokens: int - ) -> "CudaGraphBufferRegistry": + ) -> CudaGraphBufferRegistry: return self._ensure_eager_registry( self._eager_prefill_registry, raw_bs, diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index b6f174fc5..48048d158 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -471,9 +471,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): # ----------------------------------------------------------------- # capture_prepare # ----------------------------------------------------------------- - def capture_prepare( - self, num_tokens: int - ) -> tuple[ForwardBatch, "AttentionBackend"]: + def capture_prepare(self, num_tokens: int) -> tuple[ForwardBatch, AttentionBackend]: """Build a dummy prefill ForwardBatch for capture/warmup at this shape. Default tensor inputs are fresh literals; under a Breakable diff --git a/python/sglang/srt/model_executor/runner_utils/buffers.py b/python/sglang/srt/model_executor/runner_utils/buffers.py index 024bee974..550c6c953 100644 --- a/python/sglang/srt/model_executor/runner_utils/buffers.py +++ b/python/sglang/srt/model_executor/runner_utils/buffers.py @@ -68,7 +68,7 @@ class DecodeInputBuffers(ForwardInputBuffers): global_num_tokens_for_logprob_gpu: torch.Tensor encoder_lens: Optional[torch.Tensor] pp_proxy_tensors: Optional[Dict[str, torch.Tensor]] - ngram_embedding_info: Optional["NgramEmbeddingInfo"] + ngram_embedding_info: Optional[NgramEmbeddingInfo] rids_int: Optional[torch.Tensor] bootstrap_room_ids_int: Optional[torch.Tensor] @@ -94,7 +94,7 @@ class DecodeInputBuffers(ForwardInputBuffers): ne_token_table: Optional[torch.Tensor] = None, is_hybrid_swa: bool = False, hc_hidden_size: Optional[int] = None, - ) -> "DecodeInputBuffers": + ) -> DecodeInputBuffers: with torch.device(device): input_ids = torch.zeros((max_num_token,), dtype=torch.int64) input_embeds = torch.zeros((max_num_token, hidden_size), dtype=dtype) @@ -350,7 +350,7 @@ class PrefillInputBuffers(ForwardInputBuffers): hidden_size: int, dtype: torch.dtype, enable_mamba_track: bool, - ) -> "PrefillInputBuffers": + ) -> PrefillInputBuffers: with torch.device(device): input_ids = torch.zeros((max_num_tokens,), dtype=torch.int64) out_cache_loc = torch.zeros((max_num_tokens,), dtype=cache_loc_dtype) diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 5ce23f861..52bb6aa68 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -351,7 +351,7 @@ class DefaultModelLoader(BaseModelLoader): fall_back_to_pt: bool = True """Whether .pt weights can be used.""" - model_config: Optional["ModelConfig"] = None + model_config: Optional[ModelConfig] = None """The model configuration (for checking architecture, etc).""" @classmethod @@ -517,7 +517,7 @@ class DefaultModelLoader(BaseModelLoader): return hf_folder, hf_weights_files, use_safetensors def _get_weights_iterator( - self, source: "Source" + self, source: Source ) -> Generator[Tuple[str, torch.Tensor], None, None]: """Get an iterator for the model weights based on the load format.""" extra_config = self.load_config.model_loader_extra_config @@ -2863,7 +2863,7 @@ class RunaiModelStreamerLoader(BaseModelLoader): fall_back_to_pt: bool = True """Whether .pt weights can be used.""" - model_config: Optional["ModelConfig"] = None + model_config: Optional[ModelConfig] = None """The model configuration (for checking architecture, etc).""" @classmethod @@ -2963,7 +2963,7 @@ class RunaiModelStreamerLoader(BaseModelLoader): return hf_folder, hf_weights_files def _get_weights_iterator( - self, source: "Source" + self, source: Source ) -> Generator[Tuple[str, torch.Tensor], None, None]: """Get an iterator for the model weights based on the load format.""" from sglang.srt.model_loader.weight_utils import ( diff --git a/python/sglang/srt/models/mindspore.py b/python/sglang/srt/models/mindspore.py index b91197286..3b07a41b5 100644 --- a/python/sglang/srt/models/mindspore.py +++ b/python/sglang/srt/models/mindspore.py @@ -55,7 +55,7 @@ def tensor_torch2ms(x: torch.Tensor): return ms_tensor -def tensor_ms2torch(x: "ms.Tensor"): +def tensor_ms2torch(x: ms.Tensor): if x is None or not isinstance(x, ms.Tensor): return x @@ -152,7 +152,7 @@ class LowerTriangularMask: def gen_attention_mask( self, is_prefill: bool, - position_ids: "ms.Tensor", + position_ids: ms.Tensor, query_lens_np: np.ndarray, seq_lens_np: np.ndarray, ): @@ -316,7 +316,7 @@ class MindSporeForCausalLM(torch.nn.Module): input_ids: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, - ) -> "ms.Tensor": + ) -> ms.Tensor: # prepare base inputs model_inputs = self.prepare_inputs(input_ids, positions, forward_batch) # prepare model inputs diff --git a/python/sglang/srt/models/utils.py b/python/sglang/srt/models/utils.py index 6c8b1ca8c..9465701c6 100644 --- a/python/sglang/srt/models/utils.py +++ b/python/sglang/srt/models/utils.py @@ -55,7 +55,7 @@ class WeightsMapper: orig_to_new_prefix: WeightsMapping = field(default_factory=dict) orig_to_new_suffix: WeightsMapping = field(default_factory=dict) - def __or__(self, other: "WeightsMapper") -> "WeightsMapper": + def __or__(self, other: WeightsMapper) -> WeightsMapper: return WeightsMapper( orig_to_new_substr={**self.orig_to_new_substr, **other.orig_to_new_substr}, orig_to_new_prefix={**self.orig_to_new_prefix, **other.orig_to_new_prefix}, diff --git a/python/sglang/srt/multimodal/processors/mimo_audio.py b/python/sglang/srt/multimodal/processors/mimo_audio.py index 3bae8ea16..ae21abd09 100644 --- a/python/sglang/srt/multimodal/processors/mimo_audio.py +++ b/python/sglang/srt/multimodal/processors/mimo_audio.py @@ -133,7 +133,7 @@ class MiMoAudioPipeline: center=True, ) self._mel_spectrogram = None - self._resamplers: "OrderedDict[int, torchaudio.transforms.Resample]" = ( + self._resamplers: OrderedDict[int, torchaudio.transforms.Resample] = ( OrderedDict() ) self._resamplers_max = max_resamplers diff --git a/python/sglang/srt/observability/metrics_collector.py b/python/sglang/srt/observability/metrics_collector.py index 58d7b1a88..f83020ee8 100644 --- a/python/sglang/srt/observability/metrics_collector.py +++ b/python/sglang/srt/observability/metrics_collector.py @@ -197,7 +197,7 @@ STAT_LOGGER_ROLE_EXPERT_DISPATCH = "expert_dispatch" def resolve_collector_class( - server_args: Optional["ServerArgs"], role: str, default_cls: type + server_args: Optional[ServerArgs], role: str, default_cls: type ) -> type: """Return the subclass registered for `role` on `server_args.stat_loggers`, or `default_cls` if none is registered. Tolerates `server_args=None` and @@ -230,7 +230,7 @@ class SchedulerMetricsCollectorContext: is_stats_logging_rank: bool current_scheduler_metrics_enabled: bool enable_kv_cache_events: bool - collector: Optional["SchedulerMetricsCollector"] + collector: Optional[SchedulerMetricsCollector] class SchedulerMetricsCollector(_StatLoggerDIMixin): @@ -241,7 +241,7 @@ class SchedulerMetricsCollector(_StatLoggerDIMixin): enable_lora: bool = False, enable_hierarchical_cache: bool = False, enable_streaming_session: bool = False, - server_args: Optional["ServerArgs"] = None, + server_args: Optional[ServerArgs] = None, ) -> None: # We need to import prometheus_client after setting the env variable `PROMETHEUS_MULTIPROC_DIR` from prometheus_client import Counter as _PromCounter @@ -1028,7 +1028,7 @@ class SchedulerMetricsCollector(_StatLoggerDIMixin): def init_new( cls, *, - server_args: "ServerArgs", + server_args: ServerArgs, ps: Any, tp_rank: int, pp_rank: int, @@ -1036,7 +1036,7 @@ class SchedulerMetricsCollector(_StatLoggerDIMixin): enable_priority_scheduling: bool, enable_lora: bool, enable_hierarchical_cache: bool, - ) -> "SchedulerMetricsCollectorContext": + ) -> SchedulerMetricsCollectorContext: enable_metrics = server_args.enable_metrics is_stats_logging_rank = ps.attn_tp_rank == 0 current_scheduler_metrics_enabled = enable_metrics and ( @@ -1047,7 +1047,7 @@ class SchedulerMetricsCollector(_StatLoggerDIMixin): and ps.attn_tp_rank == 0 and ps.attn_cp_rank == 0 ) - collector: Optional["SchedulerMetricsCollector"] = None + collector: Optional[SchedulerMetricsCollector] = None if enable_metrics: engine_type = DisaggregationMode.to_engine_type( server_args.disaggregation_mode diff --git a/python/sglang/srt/observability/trace.py b/python/sglang/srt/observability/trace.py index 7c38b6926..fb74fc4bf 100644 --- a/python/sglang/srt/observability/trace.py +++ b/python/sglang/srt/observability/trace.py @@ -407,7 +407,7 @@ class TraceReqContext: ) self.events_cache = [] - def copy_for_thread(self) -> "TraceReqContext": + def copy_for_thread(self) -> TraceReqContext: """ Create a copy of this context for use in another thread. diff --git a/python/sglang/srt/sampling/penaltylib/orchestrator.py b/python/sglang/srt/sampling/penaltylib/orchestrator.py index 650c719f3..5cdbe83fa 100644 --- a/python/sglang/srt/sampling/penaltylib/orchestrator.py +++ b/python/sglang/srt/sampling/penaltylib/orchestrator.py @@ -15,7 +15,7 @@ class BatchedPenalizerOrchestrator: self, vocab_size: int, batch: ScheduleBatch, - penalizers: Set[Type["_BatchedPenalizer"]], + penalizers: Set[Type[_BatchedPenalizer]], ): self.vocab_size = vocab_size self._batch_ref = weakref.ref(batch) @@ -139,13 +139,13 @@ class BatchedPenalizerOrchestrator: self.is_required = False # Context manager support - def __enter__(self) -> "BatchedPenalizerOrchestrator": + def __enter__(self) -> BatchedPenalizerOrchestrator: return self def __exit__(self, exc_type, exc, tb) -> None: self.release() - def merge(self, their: "BatchedPenalizerOrchestrator"): + def merge(self, their: BatchedPenalizerOrchestrator): """ Merge the penalizers of another orchestrator into this one. @@ -227,7 +227,7 @@ class _BatchedPenalizer(abc.ABC): self._filter(keep_indices=keep_indices) - def merge(self, their: "_BatchedPenalizer"): + def merge(self, their: _BatchedPenalizer): if not self._is_prepared and not their._is_prepared: return @@ -281,7 +281,7 @@ class _BatchedPenalizer(abc.ABC): pass @abc.abstractmethod - def _merge(self, their: "_BatchedPenalizer"): + def _merge(self, their: _BatchedPenalizer): """ Merge the penalizer with another penalizer. """ diff --git a/python/sglang/srt/sampling/sampling_batch_info.py b/python/sglang/srt/sampling/sampling_batch_info.py index f0defd283..492bb7dae 100644 --- a/python/sglang/srt/sampling/sampling_batch_info.py +++ b/python/sglang/srt/sampling/sampling_batch_info.py @@ -195,7 +195,7 @@ class SamplingBatchInfo: pass # placeholder for override - def adjusted_merge_batch(self, other: "SamplingBatchInfo"): + def adjusted_merge_batch(self, other: SamplingBatchInfo): pass # placeholder for override @@ -352,7 +352,7 @@ class SamplingBatchInfo: return merged_dict - def merge_batch(self, other: "SamplingBatchInfo"): + def merge_batch(self, other: SamplingBatchInfo): self.penalizer_orchestrator.merge(other.penalizer_orchestrator) # Merge the custom logit processors and custom params lists diff --git a/python/sglang/srt/session/session_controller.py b/python/sglang/srt/session/session_controller.py index 1915c370a..aabb32f94 100644 --- a/python/sglang/srt/session/session_controller.py +++ b/python/sglang/srt/session/session_controller.py @@ -36,7 +36,7 @@ class SessionReqNode: def __init__( self, req: Req, - parent: Optional["SessionReqNode"] = None, + parent: Optional[SessionReqNode] = None, children=None, ): self.req = req diff --git a/python/sglang/srt/speculative/dflash_info.py b/python/sglang/srt/speculative/dflash_info.py index e0a0d7303..51912240c 100644 --- a/python/sglang/srt/speculative/dflash_info.py +++ b/python/sglang/srt/speculative/dflash_info.py @@ -52,7 +52,7 @@ class DFlashVerifyInput(SpecInput): def prepare_for_v2_verify( self, batch: ScheduleBatch, - target_worker: "TpModelWorker", + target_worker: TpModelWorker, ) -> tuple[ForwardBatch, bool]: """Prepare a DFLASH verify forward batch for overlap scheduling. diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_utils.py b/python/sglang/srt/speculative/frozen_kv_mtp_utils.py index 69d121137..c0f48b2b5 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_utils.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_utils.py @@ -30,7 +30,7 @@ if TYPE_CHECKING: def frozen_kv_target_view( forward_batch: ForwardBatch, kv_context: FrozenKVMTPContext, - draft_attn_backend: "AttentionBackend", + draft_attn_backend: AttentionBackend, ): """Build attention metadata against committed target-prefix geometry. @@ -61,7 +61,7 @@ def frozen_kv_target_view( def target_kv_pool_view( forward_batch: ForwardBatch, kv_context: FrozenKVMTPContext, - draft_attn_backend: "AttentionBackend", + draft_attn_backend: AttentionBackend, ): """Run the draft model's forward with the target's frozen KV pool. diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py index fb39f48d6..a15c9e4ad 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py @@ -166,7 +166,7 @@ class FrozenKVMTPDraftWorker(BaseDraftWorker, TpModelWorker): type(self.draft_model_runner.model).__name__, ) - self.kv_context: Optional["FrozenKVMTPContext"] = None + self.kv_context: Optional[FrozenKVMTPContext] = None if hasattr(self.draft_model_runner.model, "bind_frozen_kv_context"): self._bind_kv_context() diff --git a/python/sglang/srt/speculative/spec_registry.py b/python/sglang/srt/speculative/spec_registry.py index 88628a031..82580593d 100644 --- a/python/sglang/srt/speculative/spec_registry.py +++ b/python/sglang/srt/speculative/spec_registry.py @@ -85,7 +85,7 @@ class CustomSpecAlgo: def supports_spec_v2(self) -> bool: return self.supports_overlap - def create_worker(self, server_args: "ServerArgs") -> Type: + def create_worker(self, server_args: ServerArgs) -> Type: if not server_args.disable_overlap_schedule and not self.supports_overlap: raise ValueError( f"Speculative algorithm {self.name} does not support overlap scheduling." diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index bebf16672..9c349b04b 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -455,7 +455,7 @@ class DynamicGradMode(_DecoratorContextManager): else: torch.set_grad_enabled(self.prev) - def clone(self) -> "DynamicGradMode": + def clone(self) -> DynamicGradMode: r""" Create a copy of this class """ diff --git a/python/sglang/srt/utils/request_logger.py b/python/sglang/srt/utils/request_logger.py index 2b3b9bd99..00cbaf529 100644 --- a/python/sglang/srt/utils/request_logger.py +++ b/python/sglang/srt/utils/request_logger.py @@ -34,7 +34,7 @@ WHITELISTED_HEADERS = _DEFAULT_WHITELISTED_HEADERS + [ def _extract_whitelisted_headers( - request: Optional["fastapi.Request"], + request: Optional[fastapi.Request], ) -> Optional[Dict[str, str]]: if request is None: return None @@ -87,9 +87,9 @@ class RequestLogger: def log_received_request( self, - obj: Union["GenerateReqInput", "EmbeddingReqInput"], + obj: Union[GenerateReqInput, EmbeddingReqInput], tokenizer: Any = None, - request: Optional["fastapi.Request"] = None, + request: Optional[fastapi.Request] = None, ) -> None: if not self.log_requests: return @@ -131,7 +131,7 @@ class RequestLogger: def log_openai_received_request( self, obj: Any, - request: Optional["fastapi.Request"] = None, + request: Optional[fastapi.Request] = None, ) -> None: """Log the raw OpenAI request payload before request adaptation/tokenization.""" max_length, _, _ = self.metadata @@ -158,9 +158,9 @@ class RequestLogger: def log_finished_request( self, - obj: Union["GenerateReqInput", "EmbeddingReqInput"], + obj: Union[GenerateReqInput, EmbeddingReqInput], out: Any, - request: Optional["fastapi.Request"] = None, + request: Optional[fastapi.Request] = None, ) -> None: if not self.log_requests: return diff --git a/python/sglang/srt/utils/scheduler_status_logger.py b/python/sglang/srt/utils/scheduler_status_logger.py index 5f1f03caa..dcc65596e 100644 --- a/python/sglang/srt/utils/scheduler_status_logger.py +++ b/python/sglang/srt/utils/scheduler_status_logger.py @@ -20,7 +20,7 @@ class SchedulerStatusLogger: self.rank = dist.get_rank() if dist.is_initialized() else 0 @staticmethod - def maybe_create(enable_metrics: bool) -> Optional["SchedulerStatusLogger"]: + def maybe_create(enable_metrics: bool) -> Optional[SchedulerStatusLogger]: target = envs.SGLANG_LOG_SCHEDULER_STATUS_TARGET.get() if not target: return None @@ -37,7 +37,7 @@ class SchedulerStatusLogger: ) def maybe_dump( - self, running_batch: "ScheduleBatch", waiting_queue: List["Req"] + self, running_batch: ScheduleBatch, waiting_queue: List[Req] ) -> None: now = time.time() if now - self.last_dump_time < self.dump_interval: diff --git a/python/sglang/test/scripted_runtime/context/api.py b/python/sglang/test/scripted_runtime/context/api.py index 8f924748a..bb3a8fbd3 100644 --- a/python/sglang/test/scripted_runtime/context/api.py +++ b/python/sglang/test/scripted_runtime/context/api.py @@ -34,9 +34,9 @@ class ScriptedContext: def __init__( self, *, - scheduler_hook: "ScriptedSchedulerHook", - tokenizer_recv_proxy: Optional["ScriptedTokenizerRecvProxy"], - http_poster: "BackgroundHttpPoster", + scheduler_hook: ScriptedSchedulerHook, + tokenizer_recv_proxy: Optional[ScriptedTokenizerRecvProxy], + http_poster: BackgroundHttpPoster, ) -> None: assert ( scheduler_hook._is_driver @@ -67,7 +67,7 @@ class ScriptedContext: stop_token_ids: Optional[List[int]] = None, temperature: Optional[float] = None, lora_path: Optional[str] = None, - ) -> "ScriptedReqHandle": + ) -> ScriptedReqHandle: return self._req_starter.start_req( prompt_len=prompt_len, max_new_tokens=max_new_tokens, @@ -93,7 +93,7 @@ class ScriptedContext: def abort_all(self) -> None: return lifecycle.abort_all(self) - def abort(self, handle: "ScriptedReqHandle", *, await_arrival: bool = True) -> None: + def abort(self, handle: ScriptedReqHandle, *, await_arrival: bool = True) -> None: return lifecycle.abort(self, rid=handle.rid, await_arrival=await_arrival) def flush_cache(self) -> None: @@ -133,7 +133,7 @@ class ScriptedContext: def last_batch_forward_mode(self) -> Optional[str]: return queries.last_batch_forward_mode(self) - def find_req_by_rid(self, rid: str) -> Optional["Req"]: + def find_req_by_rid(self, rid: str) -> Optional[Req]: return queries.find_req_by_rid(self, rid) def is_finished(self, rid: str) -> bool: @@ -148,7 +148,7 @@ class ScriptedContext: def remaining_prompt_tokens(self, rid: str) -> int: return queries.remaining_prompt_tokens(self, rid) - def list_active_reqs(self) -> List["Req"]: + def list_active_reqs(self) -> List[Req]: return queries.list_active_reqs(self) def chunks_done(self, rid: str) -> int: diff --git a/python/sglang/test/scripted_runtime/context/engine.py b/python/sglang/test/scripted_runtime/context/engine.py index 88a37ea53..73427047f 100644 --- a/python/sglang/test/scripted_runtime/context/engine.py +++ b/python/sglang/test/scripted_runtime/context/engine.py @@ -6,7 +6,7 @@ if TYPE_CHECKING: from sglang.test.scripted_runtime.context.api import ScriptedContext -def engine_stats(ctx: "ScriptedContext") -> Dict[str, int]: +def engine_stats(ctx: ScriptedContext) -> Dict[str, int]: s = ctx.scheduler return { "kv_pool_free": s.token_to_kv_pool_allocator.available_size(), diff --git a/python/sglang/test/scripted_runtime/context/http_post.py b/python/sglang/test/scripted_runtime/context/http_post.py index 54e26cdb4..e5beee3bd 100644 --- a/python/sglang/test/scripted_runtime/context/http_post.py +++ b/python/sglang/test/scripted_runtime/context/http_post.py @@ -12,7 +12,7 @@ RECV_MSG_ARRIVAL_TIMEOUT_S: float = 60.0 def _http_post_and_await_recv_msg( - ctx: "ScriptedContext", + ctx: ScriptedContext, *, path: str, json: Optional[Dict[str, Any]], @@ -29,7 +29,7 @@ def _http_post_and_await_recv_msg( def _http_post_fire_and_forget( - ctx: "ScriptedContext", + ctx: ScriptedContext, *, path: str, json: Optional[Dict[str, Any]], @@ -38,7 +38,7 @@ def _http_post_fire_and_forget( def _submit_post( - ctx: "ScriptedContext", + ctx: ScriptedContext, *, path: str, json: Optional[Dict[str, Any]], diff --git a/python/sglang/test/scripted_runtime/context/kv_pool_exhauster.py b/python/sglang/test/scripted_runtime/context/kv_pool_exhauster.py index 695774bed..3004dffea 100644 --- a/python/sglang/test/scripted_runtime/context/kv_pool_exhauster.py +++ b/python/sglang/test/scripted_runtime/context/kv_pool_exhauster.py @@ -10,9 +10,9 @@ if TYPE_CHECKING: class ScriptedKvPoolExhauster: - def __init__(self, scheduler: "Scheduler") -> None: + def __init__(self, scheduler: Scheduler) -> None: self.scheduler = scheduler - self._held: List["torch.Tensor"] = [] + self._held: List[torch.Tensor] = [] def exhaust(self, *, leave_pages: int) -> None: allocator = self.scheduler.token_to_kv_pool_allocator diff --git a/python/sglang/test/scripted_runtime/context/lifecycle.py b/python/sglang/test/scripted_runtime/context/lifecycle.py index 4d966b470..44078ec66 100644 --- a/python/sglang/test/scripted_runtime/context/lifecycle.py +++ b/python/sglang/test/scripted_runtime/context/lifecycle.py @@ -18,7 +18,7 @@ if TYPE_CHECKING: def _await_control( - ctx: "ScriptedContext", + ctx: ScriptedContext, *, path: str, json, @@ -38,7 +38,7 @@ def _await_control( def pause_generation( - ctx: "ScriptedContext", *, mode: Literal["retract", "in_place"] + ctx: ScriptedContext, *, mode: Literal["retract", "in_place"] ) -> None: _await_control( ctx, @@ -48,7 +48,7 @@ def pause_generation( ) -def continue_generation(ctx: "ScriptedContext", *, torch_empty_cache: bool) -> None: +def continue_generation(ctx: ScriptedContext, *, torch_empty_cache: bool) -> None: _await_control( ctx, path="/continue_generation", @@ -57,7 +57,7 @@ def continue_generation(ctx: "ScriptedContext", *, torch_empty_cache: bool) -> N ) -def abort_all(ctx: "ScriptedContext") -> None: +def abort_all(ctx: ScriptedContext) -> None: _await_control( ctx, path="/abort_request", @@ -66,7 +66,7 @@ def abort_all(ctx: "ScriptedContext") -> None: ) -def abort(ctx: "ScriptedContext", *, rid: str, await_arrival: bool = True) -> None: +def abort(ctx: ScriptedContext, *, rid: str, await_arrival: bool = True) -> None: _await_control( ctx, path="/abort_request", @@ -76,7 +76,7 @@ def abort(ctx: "ScriptedContext", *, rid: str, await_arrival: bool = True) -> No ) -def flush_cache(ctx: "ScriptedContext") -> None: +def flush_cache(ctx: ScriptedContext) -> None: _await_control( ctx, path="/flush_cache", diff --git a/python/sglang/test/scripted_runtime/context/lock_ref_exhauster.py b/python/sglang/test/scripted_runtime/context/lock_ref_exhauster.py index f4257d7dc..67368f55a 100644 --- a/python/sglang/test/scripted_runtime/context/lock_ref_exhauster.py +++ b/python/sglang/test/scripted_runtime/context/lock_ref_exhauster.py @@ -10,7 +10,7 @@ if TYPE_CHECKING: class ScriptedLockRefExhauster: - def __init__(self, scheduler: "Scheduler") -> None: + def __init__(self, scheduler: Scheduler) -> None: self.scheduler = scheduler self._locked: List[Any] = [] diff --git a/python/sglang/test/scripted_runtime/context/queries.py b/python/sglang/test/scripted_runtime/context/queries.py index ce0639532..7cfab1f40 100644 --- a/python/sglang/test/scripted_runtime/context/queries.py +++ b/python/sglang/test/scripted_runtime/context/queries.py @@ -7,7 +7,7 @@ if TYPE_CHECKING: from sglang.test.scripted_runtime.context.api import ScriptedContext -def _get_all_reqs(ctx: "ScriptedContext") -> Iterator["Req"]: +def _get_all_reqs(ctx: ScriptedContext) -> Iterator[Req]: s = ctx.scheduler if s.chunked_req is not None: yield s.chunked_req @@ -23,11 +23,11 @@ def _get_all_reqs(ctx: "ScriptedContext") -> Iterator["Req"]: yield from s.last_batch.reqs -def list_active_reqs(ctx: "ScriptedContext") -> List["Req"]: +def list_active_reqs(ctx: ScriptedContext) -> List[Req]: return list(set(_get_all_reqs(ctx))) -def batch_composition(ctx: "ScriptedContext") -> Dict[str, List[str]]: +def batch_composition(ctx: ScriptedContext) -> Dict[str, List[str]]: s = ctx.scheduler chunked_rid = s.chunked_req.rid if s.chunked_req is not None else None chunked = [chunked_rid] if chunked_rid is not None else [] @@ -50,7 +50,7 @@ def batch_composition(ctx: "ScriptedContext") -> Dict[str, List[str]]: } -def is_idle(ctx: "ScriptedContext") -> bool: +def is_idle(ctx: ScriptedContext) -> bool: s = ctx.scheduler return ( s.chunked_req is None @@ -59,26 +59,26 @@ def is_idle(ctx: "ScriptedContext") -> bool: ) -def is_fully_idle(ctx: "ScriptedContext") -> bool: +def is_fully_idle(ctx: ScriptedContext) -> bool: s = ctx.scheduler return is_idle(ctx) and (s.last_batch is None or s.last_batch.is_empty()) -def last_batch_forward_mode(ctx: "ScriptedContext") -> Optional[str]: +def last_batch_forward_mode(ctx: ScriptedContext) -> Optional[str]: s = ctx.scheduler if s.last_batch is not None and s.last_batch.forward_mode is not None: return s.last_batch.forward_mode.name return None -def find_req_by_rid(ctx: "ScriptedContext", rid: str) -> Optional["Req"]: +def find_req_by_rid(ctx: ScriptedContext, rid: str) -> Optional[Req]: req = next((r for r in _get_all_reqs(ctx) if r.rid == rid), None) if req is not None: ctx._seen_rids.add(rid) return req -def is_finished(ctx: "ScriptedContext", rid: str) -> bool: +def is_finished(ctx: ScriptedContext, rid: str) -> bool: req = find_req_by_rid(ctx, rid) if req is not None: return req.finished() @@ -96,12 +96,12 @@ def is_finished(ctx: "ScriptedContext", rid: str) -> bool: return False -def is_chunking(ctx: "ScriptedContext", rid: str) -> bool: +def is_chunking(ctx: ScriptedContext, rid: str) -> bool: s = ctx.scheduler return s.chunked_req is not None and s.chunked_req.rid == rid -def status(ctx: "ScriptedContext", rid: str) -> str: +def status(ctx: ScriptedContext, rid: str) -> str: s = ctx.scheduler if rid in {r.rid for r in s.waiting_queue}: return "waiting" @@ -113,14 +113,14 @@ def status(ctx: "ScriptedContext", rid: str) -> str: return "unknown" -def remaining_prompt_tokens(ctx: "ScriptedContext", rid: str) -> int: +def remaining_prompt_tokens(ctx: ScriptedContext, rid: str) -> int: req = find_req_by_rid(ctx, rid) if req is None: return 0 return max(0, len(req.origin_input_ids) - req.kv_committed_len) -def chunks_done(ctx: "ScriptedContext", rid: str) -> int: +def chunks_done(ctx: ScriptedContext, rid: str) -> int: log = ctx._scheduler_hook._batch_log held = sum(1 for record in log if record.chunked_rid == rid and rid in record.rids) if held == 0: @@ -131,7 +131,7 @@ def chunks_done(ctx: "ScriptedContext", rid: str) -> int: return held + (1 if completed else 0) -def chunked_parks(ctx: "ScriptedContext", rid: str) -> int: +def chunked_parks(ctx: ScriptedContext, rid: str) -> int: return sum( 1 for record in ctx._scheduler_hook._batch_log diff --git a/python/sglang/test/scripted_runtime/context/radix.py b/python/sglang/test/scripted_runtime/context/radix.py index 6e3294981..45faf2868 100644 --- a/python/sglang/test/scripted_runtime/context/radix.py +++ b/python/sglang/test/scripted_runtime/context/radix.py @@ -8,11 +8,11 @@ if TYPE_CHECKING: from sglang.test.scripted_runtime.context.api import ScriptedContext -def get_all_node_hit_counts(ctx: "ScriptedContext") -> Dict[int, int]: +def get_all_node_hit_counts(ctx: ScriptedContext) -> Dict[int, int]: return _collect_node_attr(ctx, lambda node: node.hit_count) -def get_all_node_lock_refs(ctx: "ScriptedContext") -> Dict[int, int]: +def get_all_node_lock_refs(ctx: ScriptedContext) -> Dict[int, int]: return _collect_node_attr(ctx, _node_lock_ref) @@ -23,7 +23,7 @@ def _node_lock_ref(node: Any) -> int: def _collect_node_attr( - ctx: "ScriptedContext", get_value: Callable[[Any], int] + ctx: ScriptedContext, get_value: Callable[[Any], int] ) -> Dict[int, int]: values: Dict[int, int] = {} stack = list(ctx.scheduler.tree_cache.root_node.children.values()) diff --git a/python/sglang/test/scripted_runtime/context/req_starter.py b/python/sglang/test/scripted_runtime/context/req_starter.py index 99c188ad4..ea21575c4 100644 --- a/python/sglang/test/scripted_runtime/context/req_starter.py +++ b/python/sglang/test/scripted_runtime/context/req_starter.py @@ -13,7 +13,7 @@ if TYPE_CHECKING: class ScriptedContextReqStarter: - def __init__(self, ctx: "ScriptedContext") -> None: + def __init__(self, ctx: ScriptedContext) -> None: self._ctx = ctx self._req_counter = 0 diff --git a/python/sglang/test/scripted_runtime/http_server.py b/python/sglang/test/scripted_runtime/http_server.py index d574c3c7e..03de57289 100644 --- a/python/sglang/test/scripted_runtime/http_server.py +++ b/python/sglang/test/scripted_runtime/http_server.py @@ -52,7 +52,7 @@ class ScriptedHttpServer: self._dirty: Optional[str] = None @classmethod - def start(cls, **engine_kwargs: Any) -> "ScriptedHttpServer": + def start(cls, **engine_kwargs: Any) -> ScriptedHttpServer: out_of_band_error_path = _create_oob_error_file() ctx = zmq.Context() diff --git a/python/sglang/test/scripted_runtime/io_struct.py b/python/sglang/test/scripted_runtime/io_struct.py index d5c6caded..c026329a1 100644 --- a/python/sglang/test/scripted_runtime/io_struct.py +++ b/python/sglang/test/scripted_runtime/io_struct.py @@ -43,7 +43,7 @@ class OutOfBandError: return json.dumps(dataclasses.asdict(self)) @classmethod - def from_json(cls, text: str) -> "OutOfBandError": + def from_json(cls, text: str) -> OutOfBandError: return cls(**json.loads(text)) diff --git a/python/sglang/test/scripted_runtime/req_handle.py b/python/sglang/test/scripted_runtime/req_handle.py index d5d1301a3..ee68a3a16 100644 --- a/python/sglang/test/scripted_runtime/req_handle.py +++ b/python/sglang/test/scripted_runtime/req_handle.py @@ -13,10 +13,10 @@ if TYPE_CHECKING: @dataclass(frozen=True, slots=True) class ScriptedReqHandle: rid: str - context: "ScriptedContext" + context: ScriptedContext @property - def req(self) -> Optional["Req"]: + def req(self) -> Optional[Req]: return self.context.find_req_by_rid(self.rid) @property diff --git a/python/sglang/test/scripted_runtime/scheduler_hook.py b/python/sglang/test/scripted_runtime/scheduler_hook.py index 4411945b3..10b52ab20 100644 --- a/python/sglang/test/scripted_runtime/scheduler_hook.py +++ b/python/sglang/test/scripted_runtime/scheduler_hook.py @@ -119,8 +119,8 @@ class ScriptedSchedulerHook: def __init__( self, *, - scheduler: "Scheduler", - tokenizer_recv_proxy: Optional["ScriptedTokenizerRecvProxy"], + scheduler: Scheduler, + tokenizer_recv_proxy: Optional[ScriptedTokenizerRecvProxy], ) -> None: self.scheduler = scheduler self._is_driver = ( @@ -128,7 +128,7 @@ class ScriptedSchedulerHook: and scheduler.ps.tp_rank == 0 and scheduler.ps.attn_cp_rank == 0 ) - self._batch_log: List["ScriptedBatchRecord"] = [] + self._batch_log: List[ScriptedBatchRecord] = [] if self._is_driver: ensure_script_importable( diff --git a/test/registered/kv_canary/test_self_unit_e2e_base.py b/test/registered/kv_canary/test_self_unit_e2e_base.py index a133389d5..45d1fc936 100644 --- a/test/registered/kv_canary/test_self_unit_e2e_base.py +++ b/test/registered/kv_canary/test_self_unit_e2e_base.py @@ -43,7 +43,7 @@ class _DummyHarness(CanaryE2EBase): class TestAssertSwaDivergenceObserved(CustomTestCase): def _make_harness( self, log_text_or_sequence - ) -> tuple[_DummyHarness, "patch._patch[None]"]: + ) -> tuple[_DummyHarness, patch._patch[None]]: harness = _DummyHarness() harness._stderr_buf = None harness._stdout_buf = None diff --git a/test/registered/unit/managers/test_customized_info_streaming.py b/test/registered/unit/managers/test_customized_info_streaming.py index 52c998069..fc6bf5a9a 100644 --- a/test/registered/unit/managers/test_customized_info_streaming.py +++ b/test/registered/unit/managers/test_customized_info_streaming.py @@ -36,8 +36,8 @@ class CustomizedInfoSampler(Sampler): def forward( self, - logits_output: "LogitsProcessorOutput", - sampling_info: "SamplingBatchInfo", + logits_output: LogitsProcessorOutput, + sampling_info: SamplingBatchInfo, return_logprob: bool, top_logprobs_nums: List[int], token_ids_logprobs: List[List[int]],