config: six more runtime readers ask the bags (#36973)

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
2026-08-29 04:19:14 -07:00
committed by GitHub
co-authored by Claude Opus 5
parent b65e677e48
commit 1a3e152f03
20 changed files with 106 additions and 152 deletions
@@ -14,8 +14,7 @@ from sglang.srt.distributed.parallel_state import (
from sglang.srt.environ import envs
from sglang.srt.eplb.expert_location import get_global_expert_location_metadata
from sglang.srt.managers.io_struct import UpdateExpertBackupReq, sock_recv, sock_send
from sglang.srt.runtime_context import get_exec
from sglang.srt.server_args import ServerArgs
from sglang.srt.runtime_context import get_exec, get_parallel
from sglang.srt.utils.network import get_local_ip_auto
PORT_BASE = envs.SGLANG_BACKUP_PORT_BASE.get()
@@ -34,16 +33,14 @@ class ExpertBackupClient:
def __init__(
self,
*,
server_args: ServerArgs,
model_config,
moe_ep_size: int,
moe_ep_rank: int,
get_model: Callable[[], Any],
):
context = zmq.Context(2)
self.server_args = server_args
self.engine_num = server_args.nnodes
self.engine_rank = server_args.node_rank
self.engine_num = get_parallel().nnodes
self.engine_rank = get_parallel().node_rank
self.recv_list = [None] * self.engine_num
self.ready_sockets = [None] * self.engine_num
self._get_model = get_model
@@ -66,14 +63,14 @@ class ExpertBackupClient:
for i in range(self.engine_num):
self.recv_list[i] = context.socket(zmq.SUB)
self.recv_list[i].connect(
f"tcp://{all_ips[i * get_world_size() // server_args.nnodes]}:{PORT_BASE + i * 2 + 1}"
f"tcp://{all_ips[i * get_world_size() // get_parallel().nnodes]}:{PORT_BASE + i * 2 + 1}"
)
self.recv_list[i].setsockopt(zmq.SUBSCRIBE, b"")
# Synchronization channel to notify the manager when this client is ready.
self.ready_sockets[i] = context.socket(zmq.PUSH)
self.ready_sockets[i].connect(
f"tcp://{all_ips[i * get_world_size() // server_args.nnodes]}:{PORT_BASE + i * 2}"
f"tcp://{all_ips[i * get_world_size() // get_parallel().nnodes]}:{PORT_BASE + i * 2}"
)
sock_send(self.ready_sockets[i], UpdateExpertBackupReq())
+12 -29
View File
@@ -38,7 +38,6 @@ import einops
import torch
import torch.distributed
from sglang.srt.arg_groups.overrides import should_report_expert_balancedness
from sglang.srt.environ import envs
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.observability.metrics_collector import (
@@ -50,8 +49,9 @@ from sglang.srt.runtime_context import get_device as get_device_namespace
from sglang.srt.runtime_context import (
get_exec,
get_schedule,
logs_expert_balancedness_to_server_log,
reports_expert_balancedness,
)
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import Withable, get_device, get_int_env_var
if TYPE_CHECKING:
@@ -85,7 +85,6 @@ class ExpertDistributionRecorder(ABC):
@staticmethod
def init_new(
server_args: ServerArgs,
expert_location_metadata: ExpertLocationMetadata,
rank: int,
):
@@ -95,9 +94,7 @@ class ExpertDistributionRecorder(ABC):
), "ExpertLocationMetadata is required for expert distribution recording. One possible"
"reason is that you are using a model that does not support expert distribution"
"recording. Try setting `get_model_config_for_expert_location` in your model."
return _ExpertDistributionRecorderReal(
server_args, expert_location_metadata, rank
)
return _ExpertDistributionRecorderReal(expert_location_metadata, rank)
else:
return _ExpertDistributionRecorderNoop()
@@ -160,11 +157,9 @@ class _ExpertDistributionRecorderNoop(ExpertDistributionRecorder):
class _ExpertDistributionRecorderReal(ExpertDistributionRecorder):
def __init__(
self,
server_args: ServerArgs,
expert_location_metadata: ExpertLocationMetadata,
rank: int,
):
self._server_args = server_args
self._expert_location_metadata = expert_location_metadata
self._recording = False
@@ -172,15 +167,13 @@ class _ExpertDistributionRecorderReal(ExpertDistributionRecorder):
self._current_forward_pass_id = Withable()
self._current_layer_idx = Withable()
self._current_debug_name = Withable()
self._accumulator = _Accumulator.init_new(
server_args, expert_location_metadata, rank
)
self._accumulator = _Accumulator.init_new(expert_location_metadata, rank)
self._single_pass_gatherers = {
k: _SinglePassGatherer.init_new(server_args, expert_location_metadata, rank)
k: _SinglePassGatherer.init_new(expert_location_metadata, rank)
for k in self._accumulator.get_single_pass_gatherer_keys()
}
if should_report_expert_balancedness(server_args):
if reports_expert_balancedness():
logger.info(
"ExpertDistributionRecorder auto start record since "
f"expert_balancedness_report_mode={get_exec().moe.expert_balancedness_report_mode}"
@@ -329,14 +322,11 @@ def set_global_expert_distribution_recorder(value):
class _SinglePassGatherer(ABC):
@staticmethod
def init_new(
server_args: ServerArgs,
expert_location_metadata: ExpertLocationMetadata,
rank: int,
) -> _SinglePassGatherer:
if get_exec().moe.expert_distribution_recorder_mode == "per_token":
return _DetailSinglePassGatherer(
server_args, expert_location_metadata, rank
)
return _DetailSinglePassGatherer(expert_location_metadata, rank)
if get_exec().moe.moe_a2a_backend == "mori":
return _DeepepLowLatencySinglePassGatherer(expert_location_metadata, rank)
@@ -403,7 +393,6 @@ class _DetailSinglePassGatherer(_SinglePassGatherer):
def __init__(
self,
server_args: ServerArgs,
expert_location_metadata: ExpertLocationMetadata,
rank: int,
):
@@ -668,11 +657,10 @@ _SINGLE_PASS_GATHERER_KEY_PRIMARY = "primary"
class _Accumulator(ABC):
@staticmethod
def init_new(
server_args: ServerArgs,
expert_location_metadata: ExpertLocationMetadata,
rank: int,
) -> _Accumulator:
return _Accumulator.get_class()(server_args, expert_location_metadata, rank)
return _Accumulator.get_class()(expert_location_metadata, rank)
@staticmethod
def get_class() -> Type[_Accumulator]:
@@ -685,11 +673,9 @@ class _Accumulator(ABC):
def __init__(
self,
server_args: ServerArgs,
expert_location_metadata: ExpertLocationMetadata,
rank: int,
):
self._server_args = server_args
self._expert_location_metadata = expert_location_metadata
self._rank = rank
@@ -719,7 +705,7 @@ class _UtilizationRateAccumulatorMixin(_Accumulator):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._enable = should_report_expert_balancedness(self._server_args)
self._enable = reports_expert_balancedness()
if self._enable:
self.window_sizes = EPLB_BALANCEDNESS_WINDOW_SIZES
@@ -727,7 +713,6 @@ class _UtilizationRateAccumulatorMixin(_Accumulator):
self._reset_server_log_history = True
self._rank = torch.distributed.get_rank()
expert_dispatch_cls = resolve_collector_class(
self._server_args,
STAT_LOGGER_ROLE_EXPERT_DISPATCH,
ExpertDispatchCollector,
)
@@ -777,12 +762,10 @@ class _UtilizationRateAccumulatorMixin(_Accumulator):
compute_utilization_rate(gpu_physical_count)
)
should_track_history = not math.isclose(
self._server_args.eplb_min_rebalancing_utilization_threshold, 1.0
get_exec().moe.eplb_min_rebalancing_utilization_threshold, 1.0
)
should_log = (
self._server_args.should_log_expert_balancedness_to_server_log()
)
should_log = logs_expert_balancedness_to_server_log()
outputs["metrics"] = ExpertDistributionMetrics(
forward_pass_id=forward_pass_id,
eplb_balancedness=utilization_rate_gpu,
@@ -954,7 +937,7 @@ class _StatAccumulator(_UtilizationRateAccumulatorMixin):
def _get_global_average_utilization_rate(self):
if not self._enable or math.isclose(
self._server_args.eplb_min_rebalancing_utilization_threshold, 1.0
get_exec().moe.eplb_min_rebalancing_utilization_threshold, 1.0
):
return None
@@ -14,6 +14,7 @@ import torch
from sglang.srt.layers.quantization.base_config import FusedMoEMethodBase
from sglang.srt.runtime_context import (
get_exec,
get_parallel,
get_schedule,
)
@@ -73,7 +74,7 @@ def create_kt_config_from_server_args(
Returns:
KTConfig if KT is configured, None otherwise
"""
if server_args.kt_weight_path is None:
if get_exec().moe.kt_weight_path is None:
return None
from sglang.srt.arg_groups.overrides import model_config_of
@@ -84,13 +85,13 @@ def create_kt_config_from_server_args(
return KTConfig(
layer_idx=layer_idx,
num_gpu_experts=server_args.kt_num_gpu_experts,
cpuinfer_threads=server_args.kt_cpuinfer,
threadpool_count=server_args.kt_threadpool_count,
weight_path=server_args.kt_weight_path,
num_gpu_experts=get_exec().moe.kt_num_gpu_experts,
cpuinfer_threads=get_exec().moe.kt_cpuinfer,
threadpool_count=get_exec().moe.kt_threadpool_count,
weight_path=get_exec().moe.kt_weight_path,
chunked_prefill_size=get_schedule().chunked_prefill_size,
method=server_args.kt_method,
max_deferred_experts_per_token=server_args.kt_max_deferred_experts_per_token,
method=get_exec().moe.kt_method,
max_deferred_experts_per_token=get_exec().moe.kt_max_deferred_experts_per_token,
num_layers=num_layers,
)
@@ -78,7 +78,6 @@ class PrefillDelayer:
dp_size: int,
attn_tp_size: int,
cpu_group,
server_args,
max_delay_passes: int,
token_usage_low_watermark: Optional[float],
metrics_collector: Optional["SchedulerMetricsCollector"] = None,
@@ -91,14 +90,14 @@ class PrefillDelayer:
self._debug_log_enabled = _DEBUG_LOG and debug_log_enabled
# Queue-based trigger is opt-in: activates only when queue_min_ratio
# is explicitly set. Additive with the slot-based trigger.
self._queue_min_ratio = server_args.prefill_delayer_queue_min_ratio
self._queue_min_ratio = get_schedule().prefill_delayer_queue_min_ratio
# Fall back to 5000ms if unset; this is a local safety cap, not a
# semantic default, so we don't surface it via ServerArgs.
self._max_delay_ms = server_args.prefill_delayer_max_delay_ms
self._max_delay_ms = get_schedule().prefill_delayer_max_delay_ms
if self._max_delay_ms is None:
self._max_delay_ms = 5000.0
self._queue_trigger_enabled = self._queue_min_ratio is not None
self._prefill_max_requests = server_args.prefill_max_requests
self._prefill_max_requests = get_schedule().prefill_max_requests
logger.info(
f"PrefillDelayer initialized with "
f"max_delay_passes={self._max_delay_passes} "
+3 -4
View File
@@ -1300,7 +1300,6 @@ class Scheduler(
attn_tp_size=self.ps.attn_tp_size,
cpu_group=self.tp_cpu_group,
device_group=self.tp_group.device_group,
server_args=self.server_args,
metrics_collector=(
self.metrics_collector
if self.metrics_reporter.enable_metrics
@@ -2820,7 +2819,7 @@ class Scheduler(
if self.enable_hicache_storage:
req.init_next_round_input(self.tree_cache, cow_mamba=False)
tree_cache = self.tree_cache
buffer_mode = self.server_args.hicache_host_memory_mode == "buffer_only"
buffer_mode = get_memory().hicache_host_memory_mode == "buffer_only"
last_host_node = req.last_host_node
# Buffer mode host-backups nothing, so match_prefix anchors at
# root; re-anchor on the deepest device node. The anchor is only
@@ -3544,7 +3543,7 @@ class Scheduler(
req.init_next_round_input(self.tree_cache)
if (
self.enable_hicache_storage
and self.server_args.hicache_host_memory_mode == "buffer_only"
and get_memory().hicache_host_memory_mode == "buffer_only"
):
# Buffer mode: surface a staged prefetch as the request's host
# hit (consumed through init_load_back) plus its SWA window,
@@ -4452,7 +4451,7 @@ class Scheduler(
if tc.enable_storage:
idle &= len(tc.ongoing_prefetch) == 0
idle &= len(tc.ongoing_backup) == 0
if self.server_args.hicache_host_memory_mode == "buffer_only":
if get_memory().hicache_host_memory_mode == "buffer_only":
# Queued writes, staged prefetches, and in-flight
# storage writes still hold host staging
# (buffer-mode unified tree only).
@@ -23,11 +23,13 @@ from sglang.srt.observability.metrics_collector import (
compute_routing_key_stats,
)
from sglang.srt.runtime_context import (
exports_expert_balancedness_to_prometheus,
get_context,
get_disagg,
get_observability,
get_parallel,
get_spec,
logs_expert_balancedness_to_server_log,
)
from sglang.srt.utils.device_timer import DeviceTimer
from sglang.srt.utils.scheduler_status_logger import SchedulerStatusLogger
@@ -1009,9 +1011,7 @@ class SchedulerMetricsReporter:
if (m := result.expert_distribution_metrics) is not None:
balancedness = m.eplb_balancedness.item()
if (
self.scheduler.server_args.should_log_expert_balancedness_to_server_log()
):
if logs_expert_balancedness_to_server_log():
if m.reset_server_log_history:
for history in self._eplb_balancedness_history:
history.clear()
@@ -1033,10 +1033,7 @@ class SchedulerMetricsReporter:
f"gpu_physical_count_sum={gpu_physical_count_sum}"
)
if (
self.enable_metrics
and self.scheduler.server_args.should_export_expert_balancedness_to_prometheus()
):
if self.enable_metrics and exports_expert_balancedness_to_prometheus():
assert self.metrics_collector is not None
self.metrics_collector.increment_eplb_balancedness(
forward_mode=batch.forward_mode.name.lower(),
@@ -717,7 +717,6 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
if get_observability().extra_metric_labels:
labels.update(get_observability().extra_metric_labels)
tokenizer_collector_cls = resolve_collector_class(
self.server_args,
STAT_LOGGER_ROLE_TOKENIZER,
TokenizerMetricsCollector,
)
@@ -243,14 +243,10 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
kv_events: Optional[KVCacheEventRecorder] = None
def init_metrics_collector(self):
from sglang.srt.runtime_context import get_server_args
server_args = get_server_args()
labels = {"cache_type": self.__class__.__name__}
if get_observability().extra_metric_labels:
labels.update(get_observability().extra_metric_labels)
radix_cache_cls = resolve_collector_class(
server_args,
STAT_LOGGER_ROLE_RADIX_CACHE,
RadixCacheMetricsCollector,
)
@@ -344,10 +344,8 @@ class HiRadixCache(RadixCache):
labels.update(extra_metric_labels)
existing_collector = getattr(self, "storage_metrics_collector", None)
if existing_collector is None:
from sglang.srt.runtime_context import get_server_args
storage_cls = resolve_collector_class(
get_server_args(),
STAT_LOGGER_ROLE_STORAGE,
StorageMetricsCollector,
)
@@ -327,7 +327,7 @@ def build_kv_only_stack(
storage_backend_extra_config=storage_backend_extra_config,
transfer_layer_num=transfer_layer_num,
enable_storage_metrics=enable_storage_metrics,
host_memory_mode=server_args.hicache_host_memory_mode,
host_memory_mode=get_memory().hicache_host_memory_mode,
)
return host_pool_group, cache_controller
@@ -394,7 +394,7 @@ def build_hybrid_swa_stack(
storage_backend_extra_config=storage_backend_extra_config,
transfer_layer_num=transfer_layer_num,
enable_storage_metrics=enable_storage_metrics,
host_memory_mode=server_args.hicache_host_memory_mode,
host_memory_mode=get_memory().hicache_host_memory_mode,
)
return host_pool_group, cache_controller
@@ -679,7 +679,7 @@ def build_deepseek_v4_hicache_stack(
storage_backend_extra_config=storage_backend_extra_config,
transfer_layer_num=transfer_layer_num,
enable_storage_metrics=enable_storage_metrics,
host_memory_mode=server_args.hicache_host_memory_mode,
host_memory_mode=get_memory().hicache_host_memory_mode,
)
return host_pool_group, cache_controller
@@ -773,7 +773,7 @@ def build_hybrid_mamba_stack(
storage_backend_extra_config=storage_backend_extra_config,
transfer_layer_num=transfer_layer_num,
enable_storage_metrics=enable_storage_metrics,
host_memory_mode=server_args.hicache_host_memory_mode,
host_memory_mode=get_memory().hicache_host_memory_mode,
)
return host_pool_group, cache_controller
@@ -885,7 +885,7 @@ def build_hybrid_mamba_swa_stack(
storage_backend_extra_config=storage_backend_extra_config,
transfer_layer_num=transfer_layer_num,
enable_storage_metrics=enable_storage_metrics,
host_memory_mode=server_args.hicache_host_memory_mode,
host_memory_mode=get_memory().hicache_host_memory_mode,
)
return host_pool_group, cache_controller
@@ -964,7 +964,7 @@ def build_anchor_sidecar_stack(
storage_backend_extra_config=storage_backend_extra_config,
transfer_layer_num=transfer_layer_num,
enable_storage_metrics=enable_storage_metrics,
host_memory_mode=server_args.hicache_host_memory_mode,
host_memory_mode=get_memory().hicache_host_memory_mode,
)
return host_pool_group, cache_controller
@@ -280,10 +280,8 @@ class StorageAttachment:
existing_collector = cache.storage_metrics_collector
if existing_collector is None:
from sglang.srt.runtime_context import get_server_args
storage_cls = resolve_collector_class(
get_server_args(),
STAT_LOGGER_ROLE_STORAGE,
StorageMetricsCollector,
)
@@ -504,7 +504,6 @@ class ModelRunner:
)
set_global_expert_distribution_recorder(
ExpertDistributionRecorder.init_new(
self.server_args,
get_global_expert_location_metadata(),
rank=global_ep_rank,
)
@@ -709,7 +708,6 @@ class ModelRunner:
)
set_global_expert_distribution_recorder(
ExpertDistributionRecorder.init_new(
self.server_args,
get_global_expert_location_metadata(),
rank=expert_rank,
)
@@ -750,7 +748,6 @@ class ModelRunner:
def maybe_init_expert_backup_client(self):
self.expert_backup_client = (
ExpertBackupClient(
server_args=self.server_args,
model_config=self.model_config,
moe_ep_size=self.ps.moe_ep_size,
moe_ep_rank=self.ps.moe_ep_rank,
@@ -971,7 +968,6 @@ class ModelRunner:
)
set_global_expert_distribution_recorder(
ExpertDistributionRecorder.init_new(
self.server_args,
get_global_expert_location_metadata(),
rank=global_ep_rank,
)
@@ -1967,7 +1963,6 @@ class ModelRunner:
return
set_global_expert_distribution_recorder(
ExpertDistributionRecorder.init_new(
self.server_args,
get_global_expert_location_metadata(),
rank=self._elastic_global_rank(),
)
@@ -2032,7 +2027,6 @@ class ModelRunner:
ElasticEPStateManager.on_scale(effective_size, target_size)
set_global_expert_distribution_recorder(
ExpertDistributionRecorder.init_new(
self.server_args,
get_global_expert_location_metadata(),
rank=self._elastic_global_rank(),
)
@@ -27,6 +27,7 @@ from sglang.srt.disaggregation.utils import DisaggregationMode
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.observability.utils import exponential_buckets, generate_buckets
from sglang.srt.runtime_context import (
exports_expert_balancedness_to_prometheus,
get_disagg,
get_observability,
get_schedule,
@@ -203,15 +204,17 @@ STAT_LOGGER_ROLE_RADIX_CACHE = "radix_cache"
STAT_LOGGER_ROLE_EXPERT_DISPATCH = "expert_dispatch"
def resolve_collector_class(
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
`stat_loggers=None`."""
if server_args is None:
def resolve_collector_class(role: str, default_cls: type) -> type:
"""Return the subclass registered for `role` in the published
``observability`` bag, or `default_cls` if none is registered.
An unpublished ``observability`` namespace answers with the default.
"""
from sglang.srt.runtime_context import get_context, get_observability
if not get_context().is_config_namespace_published("observability"):
return default_cls
stat_loggers = getattr(server_args, "stat_loggers", None)
stat_loggers = get_observability().stat_loggers
if not stat_loggers:
return default_cls
return stat_loggers.get(role, default_cls)
@@ -877,10 +880,7 @@ class SchedulerMetricsCollector(_StatLoggerDIMixin):
# =================================================================
# Execution
# =================================================================
if (
labels["moe_ep_rank"] == 0
and server_args.should_export_expert_balancedness_to_prometheus()
):
if labels["moe_ep_rank"] == 0 and exports_expert_balancedness_to_prometheus():
self.eplb_balancedness = Summary(
name="sglang:eplb_balancedness",
documentation="Balancedness of MoE in expert parallelism.",
@@ -1114,7 +1114,7 @@ class SchedulerMetricsCollector(_StatLoggerDIMixin):
if get_observability().extra_metric_labels:
labels.update(get_observability().extra_metric_labels)
scheduler_collector_cls = resolve_collector_class(
server_args, STAT_LOGGER_ROLE_SCHEDULER, cls
STAT_LOGGER_ROLE_SCHEDULER, cls
)
collector = scheduler_collector_cls(
labels=labels,
+18
View File
@@ -1819,6 +1819,24 @@ def process_model_config():
return model_config_of(get_server_args())
def reports_expert_balancedness() -> bool:
"""Whether the expert-balancedness report is on at all.
`overrides.should_report_expert_balancedness` is the pre-publish equivalent.
"""
return get_exec().moe.expert_balancedness_report_mode != "off"
def logs_expert_balancedness_to_server_log() -> bool:
"""Whether the balancedness report goes to the server log."""
return get_exec().moe.expert_balancedness_report_mode in ("server_log", "both")
def exports_expert_balancedness_to_prometheus() -> bool:
"""Whether the balancedness report goes to Prometheus."""
return get_exec().moe.expert_balancedness_report_mode in ("prometheus", "both")
def cutedsl_moe_max_num_tokens() -> int:
"""The CuteDSL A2A per-rank token budget.
-10
View File
@@ -4341,16 +4341,6 @@ class ServerArgs:
descriptor["load_topic"] = LOAD_TOPIC
return descriptor
def should_log_expert_balancedness_to_server_log(self) -> bool:
cfg = resolving_view(self)
return cfg.expert_balancedness_report_mode in ("server_log", "both")
def should_export_expert_balancedness_to_prometheus(self) -> bool:
cfg = resolving_view(self)
return cfg.expert_balancedness_report_mode in ("prometheus", "both")
# --------------------------------------------------------------------------
# Module-level ServerArgs helpers and runtime shims.
@@ -21,7 +21,6 @@ from sglang.srt.managers.overlap_utils import (
)
from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.runtime_context import get_disagg, get_parallel, get_schedule, get_spec
from sglang.srt.server_args import ServerArgs
from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2
from sglang.srt.speculative.dflash_utils import apply_dflash_verify_logits_adjustments
from sglang.srt.speculative.dspark_components.dspark_sps import (
@@ -76,22 +75,20 @@ class DSparkVerifyPlanner:
model_runner,
device,
tp_rank: int,
server_args: ServerArgs,
verify_num_draft_tokens: int,
) -> None:
self.draft_model = draft_model
self.gamma = gamma
self.model_runner = model_runner
self.device = device
self.server_args = server_args
self.verify_num_draft_tokens = verify_num_draft_tokens
self._align_verify_tokens_to_graph_tier = (
server_args.speculative_dspark_align_verify_tokens_to_graph_tier
get_spec().speculative_dspark_align_verify_tokens_to_graph_tier
)
self._confidence_head = getattr(self.draft_model, "confidence_head", None)
sts_path = server_args.speculative_dspark_confidence_sts_path
sts_path = get_spec().speculative_dspark_confidence_sts_path
if sts_path and self._confidence_head is not None:
calibration = load_sts_calibration_from_path(sts_path)
sts_temperatures = torch.tensor(
@@ -148,7 +145,6 @@ class DSparkVerifyPlanner:
f"SGLANG_RAGGED_VERIFY_MODE=static."
)
sps_table = build_sps_cost_table(
server_args=self.server_args,
verify_num_draft_tokens=self.verify_num_draft_tokens,
)
self._is_verify_all = (
@@ -1120,10 +1116,9 @@ class HostConfidenceBudgetPlanner:
def build_sps_cost_table(
*,
server_args: ServerArgs,
verify_num_draft_tokens: int,
) -> Union[SpsCostTable, SpsAdditiveCostTable]:
sps_table_path = server_args.speculative_dspark_sps_table_path
sps_table_path = get_spec().speculative_dspark_sps_table_path
if sps_table_path:
return load_sps_table_from_path(sps_table_path)
max_batch_tokens = max(
@@ -219,7 +219,6 @@ class DSparkWorkerV2(BaseSpecWorker):
model_runner=self.model_runner,
device=self.device,
tp_rank=self.ps.tp_rank,
server_args=self.server_args,
verify_num_draft_tokens=self.verify_num_draft_tokens,
)
if (