diff --git a/.claude/skills/sglang-runtime-context/SKILL.md b/.claude/skills/sglang-runtime-context/SKILL.md index 3c8d651fe..6733281bc 100644 --- a/.claude/skills/sglang-runtime-context/SKILL.md +++ b/.claude/skills/sglang-runtime-context/SKILL.md @@ -73,11 +73,9 @@ with what the operator typed, not with what resolution decided.** `TokenizerManager.record_config_updates(source, **fields)`, a named wrapper over `get_context().override`. One process keeps one log: the request dumps ship `get_context().overrides_log()`, and `config_value(name)` / - `resolved_config_dict(base)` answer from the bags. The exposure ratchet - resolves the wrapper, so a field recorded through it joins the post-publish - override surface exactly like a direct `override` and needs the same ordering - judgment against any supplied-instance read of it - (`test_supplied_instance_exposure_ratchet.py`). + `resolved_config_dict(base)` answer from the bags. Fields recorded through + the wrapper follow the same ordering rules as a direct `override`: a later + read must observe the updated bag, not the startup record. - **`model_path` and `served_model_name` are answered off the manager.** Both are `NS` leaves and `override` accepts them, but the tokenizer-side weight reload records only `load_format` and writes the two path fields as `TokenizerManager` @@ -154,8 +152,7 @@ bag to override at all. **retracted** — owner ruling (2026-08-15): a process holds at most one live config at a time (concurrent multi-Engine is unsupported; sequential rebuild stays legal, unit tests rely on it). Nothing in those files reads the instance - any more -- the exposure ratchet's pin set is empty, so the next such read is a - new entry that has to argue for itself. What + any more; review new instance reads against the raw-input contract. What genuinely stays per-instance is what differs per *worker* within one engine: `base_gpu_id` travels as a constructor argument (`MMEncoder(gpu_id=...)`; `BaseMultimodalProcessor._fast_image_processor_device` is the shape to copy). @@ -172,12 +169,9 @@ bag to override at all. attention pair and the encode-server `gpu_id` above are both this). The per-instance boundaries above are **not** exempt from this unless-clause (the multi-Engine exemption is retracted); each one gets its own disposition. - `test_supplied_instance_exposure_ratchet.py` - pins that set (empty today) — three spellings of the read: `server_args.field`, - literal-name `getattr(server_args, "field", default)`, and the parked form - (`self.x = server_args` in a method that takes the parameter, read as - `self.x.field` anywhere in the class) — and fails on a new one, so the - disposition gets picked when the read is written. Two shapes stay parameter-form on purpose: a helper the + Check direct attributes, `getattr`, and records stored on `self`; validate + the resolved value and any later overrides in behavior tests. + Two shapes stay parameter-form on purpose: a helper the *resolution pipeline* calls with a `resolved_view` (its parameter happens to be named `server_args`), and a factory whose contract is "build X from the record you are handed" (`create_kt_config_from_server_args`, `DllmConfig.from_server_args`). @@ -402,12 +396,10 @@ through a view instead: - `resolved_view(server_args)` — snapshots the overlay when built, which is what a post-process pass wants: it reads the state at *its* slot. -`test_resolution_reads_the_declarations` pins direct field reads at zero over the -two scopes it can derive exactly (every `arg_groups` function taking a config, -every `ServerArgs` handler the dispatcher reaches). Readers the pipeline calls -from elsewhere (`ModelConfig`, the platform defaults, the spec-algo hook) have -moved to the view as well — a field read there is the same bug, just one the -derivation cannot enumerate. +Resolution hooks and the helpers they call (`ModelConfig`, platform defaults, +the spec-algo hook) must read through these views too. Keep coverage in +`test_resolution_declarations.py`, `test_resolution_is_reproducible.py`, and +`test_record_holds_the_raw_input.py` focused on the values callers observe. One consequence worth knowing: because the fields are the raw input, resolving a bare `dataclasses.replace` copy lands in the same place as the parent — the diff --git a/python/sglang/srt/arg_groups/attention_hook.py b/python/sglang/srt/arg_groups/attention_hook.py index e4f968285..b5ef4ed97 100644 --- a/python/sglang/srt/arg_groups/attention_hook.py +++ b/python/sglang/srt/arg_groups/attention_hook.py @@ -174,11 +174,8 @@ def handle_attention_backend_compatibility(server_args: Any): # AMD platforms backends if resolved_view(server_args).attention_backend == "aiter": if model_config.context_len > 8192: - # The record, via the input snapshot rather than the field: a - # hook may not read a field off the record (the guard in - # `test_resolution_reads_the_declarations.py`), and what this - # needs is the input anyway -- whether the operator asked for a - # memory fraction, not the value in effect. + # Check whether the operator supplied a memory fraction using + # the raw input snapshot; resolution may have filled the value in. explicit_mem_fraction = ( getattr(server_args, "_raw_input", None) or {} ).get("mem_fraction_static") is not None diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 1531ee1f3..5813f8fc5 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -128,17 +128,8 @@ def run_post_process_pass(server_args: Any, fn: Callable[..., dict]) -> None: f"got {type(declared).__name__}" ) if declared: - # Refused only once there is something to record. A pass that declares - # nothing is a validation, and `check_server_args` runs those again on - # a rebuild: `Engine(server_args=sa)` after `Engine.shutdown()` hands - # back the same instance while the context still holds it, and - # refusing on identity alone would fail that launch. - # Only a non-empty return is a declaration. An empty one is a - # validation and may run on the published instance -- see above -- so it - # must not reach the guard in `declare_resolution`. - if declared: - declare_resolution(server_args, fn.__qualname__, **declared) - validate_declarations(server_args, [(fn.__qualname__, dict(declared))]) + declare_resolution(server_args, fn.__qualname__, **declared) + validate_declarations(server_args, [(fn.__qualname__, dict(declared))]) def declare_resolution(server_args: Any, source: str, **fields: Any) -> None: @@ -147,8 +138,7 @@ def declare_resolution(server_args: Any, source: str, **fields: Any) -> None: The stash *is* the resolution result: the bags are projected from it, `resolution_result` answers from it, and no field is written. A resolver reading a field another resolver may have decided must read `resolving_view` - (or `resolved_view(server_args)`), which - `test_resolution_reads_the_declarations` pins. + (or `resolved_view(server_args)`). Every declaration goes through here, whenever it is made: inside ``__post_init__``, at launcher stage (LoRA normalization, the auto-detected diff --git a/python/sglang/srt/compilation/backend.py b/python/sglang/srt/compilation/backend.py index 5f40fd296..8c934add5 100644 --- a/python/sglang/srt/compilation/backend.py +++ b/python/sglang/srt/compilation/backend.py @@ -396,7 +396,6 @@ class SGLangBackend: self.compile_config = config def configure_post_pass(self): - self.post_grad_pass_manager.configure() self.inductor_config["post_grad_custom_post_pass"] = self.post_grad_pass_manager def __call__(self, graph: fx.GraphModule, example_inputs) -> Callable: diff --git a/python/sglang/srt/compilation/fix_functionalization.py b/python/sglang/srt/compilation/fix_functionalization.py deleted file mode 100644 index 5fdbea269..000000000 --- a/python/sglang/srt/compilation/fix_functionalization.py +++ /dev/null @@ -1,136 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# Adapted from https://github.com/vllm-project/vllm/blob/v0.10.0/vllm/compilation/fix_functionalization.py - -import logging -import operator -from collections.abc import Iterable -from typing import Optional, Union - -import torch -from torch._higher_order_ops.auto_functionalize import auto_functionalized - -from sglang.srt.compilation.fx_utils import is_func -from sglang.srt.compilation.inductor_pass import SGLangInductorPass - -logger = logging.getLogger(__name__) - - -class FixFunctionalizationPass(SGLangInductorPass): - """ - This pass defunctionalizes certain nodes to avoid redundant tensor copies. - After this pass, DCE (dead-code elimination) should never be run, - as de-functionalized nodes may appear as dead code. - - To add new nodes to defunctionalize, add to the if-elif chain in __call__. - """ - - def __call__(self, graph: torch.fx.Graph): - self.begin() - self.dump_graph(graph, "before_fix_functionalization") - - self.nodes_to_remove: list[torch.fx.Node] = [] - count = 0 - for node in graph.nodes: - if not is_func(node, auto_functionalized): - continue # Avoid deep if-elif nesting - count += 1 - - self.dump_graph(graph, "before_fix_functionalization_cleanup") - - # Remove the nodes all at once - count_removed = len(self.nodes_to_remove) - for node in self.nodes_to_remove: - graph.erase_node(node) - - logger.debug( - "De-functionalized %s nodes, removed %s nodes", count, count_removed - ) - self.dump_graph(graph, "after_fix_functionalization") - self.end_and_log() - - def _remove(self, node_or_nodes: Union[torch.fx.Node, Iterable[torch.fx.Node]]): - """ - Stage a node (or nodes) for removal at the end of the pass. - """ - if isinstance(node_or_nodes, torch.fx.Node): - self.nodes_to_remove.append(node_or_nodes) - else: - self.nodes_to_remove.extend(node_or_nodes) - - def defunctionalize( - self, - graph: torch.fx.Graph, - node: torch.fx.Node, - mutated_args: dict[int, Union[torch.fx.Node, str]], - args: Optional[tuple[Union[torch.fx.Node, str], ...]] = None, - ): - """ - De-functionalize a node by replacing it with a call to the original. - It also replaces the getitem users with the mutated arguments. - See replace_users_with_mutated_args and insert_defunctionalized. - """ - self.replace_users_with_mutated_args(node, mutated_args) - self.insert_defunctionalized(graph, node, args=args) - self._remove(node) - - def replace_users_with_mutated_args( - self, node: torch.fx.Node, mutated_args: dict[int, Union[torch.fx.Node, str]] - ): - """ - Replace all getitem users of the auto-functionalized node with the - mutated arguments. - :param node: The auto-functionalized node - :param mutated_args: The mutated arguments, indexed by getitem index. - If the value of an arg is a string, `node.kwargs[arg]` is used. - """ - for idx, user in self.getitem_users(node).items(): - arg = mutated_args[idx] - arg = node.kwargs[arg] if isinstance(arg, str) else arg - user.replace_all_uses_with(arg) - self._remove(user) - - def getitem_users(self, node: torch.fx.Node) -> dict[int, torch.fx.Node]: - """ - Returns the operator.getitem users of the auto-functionalized node, - indexed by the index they are getting. - """ - users = {} - for user in node.users: - if is_func(user, operator.getitem): - idx = user.args[1] - users[idx] = user - return users - - def insert_defunctionalized( - self, - graph: torch.fx.Graph, - node: torch.fx.Node, - args: Optional[tuple[Union[torch.fx.Node, str], ...]] = None, - ): - """ - Insert a new defunctionalized node into the graph before node. - If one of the kwargs is 'out', provide args directly, - as node.kwargs cannot be used. - See https://github.com/pytorch/pytorch/blob/a00faf440888ffb724bad413f329a49e2b6388e7/torch/_inductor/lowering.py#L351 - - :param graph: Graph to insert the defunctionalized node into - :param node: The auto-functionalized node to defunctionalize - :param args: If we cannot use kwargs, specify args directly. - If an arg is a string, `node.kwargs[arg]` is used. - """ # noqa: E501 - assert is_func(node, auto_functionalized), ( - f"node must be auto-functionalized, is {node} instead" - ) - - # Create a new call to the original function - with graph.inserting_before(node): - function = node.args[0] - if args is None: - graph.call_function(function, kwargs=node.kwargs) - else: - # Args passed as strings refer to items in node.kwargs - args = tuple( - node.kwargs[arg] if isinstance(arg, str) else arg for arg in args - ) - graph.call_function(function, args=args) diff --git a/python/sglang/srt/compilation/fx_utils.py b/python/sglang/srt/compilation/fx_utils.py deleted file mode 100644 index cd078f840..000000000 --- a/python/sglang/srt/compilation/fx_utils.py +++ /dev/null @@ -1,85 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# Adapted from https://github.com/vllm-project/vllm/blob/v0.10.0/vllm/compilation/fx_utils.py - -import operator -from collections.abc import Iterable, Iterator -from typing import Optional - -from torch import fx -from torch._higher_order_ops.auto_functionalize import auto_functionalized -from torch._ops import OpOverload - - -def is_func(node: fx.Node, target) -> bool: - return node.op == "call_function" and node.target == target - - -def is_auto_func(node: fx.Node, op: OpOverload) -> bool: - return is_func(node, auto_functionalized) and node.args[0] == op - - -# Returns the first specified node with the given op (if it exists) -def find_specified_fn_maybe( - nodes: Iterable[fx.Node], op: OpOverload -) -> Optional[fx.Node]: - for node in nodes: - if node.target == op: - return node - return None - - -# Returns the first specified node with the given op -def find_specified_fn(nodes: Iterable[fx.Node], op: OpOverload) -> fx.Node: - node = find_specified_fn_maybe(nodes, op) - assert node is not None, f"Could not find {op} in nodes {nodes}" - return node - - -# Returns the first auto_functionalized node with the given op (if it exists) -def find_auto_fn_maybe(nodes: Iterable[fx.Node], op: OpOverload) -> Optional[fx.Node]: - for node in nodes: - if is_func(node, auto_functionalized) and node.args[0] == op: # noqa - return node - return None - - -# Returns the first auto_functionalized node with the given op -def find_auto_fn(nodes: Iterable[fx.Node], op: OpOverload) -> fx.Node: - node = find_auto_fn_maybe(nodes, op) - assert node is not None, f"Could not find {op} in nodes {nodes}" - return node - - -# Returns the getitem node that extracts the idx-th element from node -# (if it exists) -def find_getitem_maybe(node: fx.Node, idx: int) -> Optional[fx.Node]: - for user in node.users: - if is_func(user, operator.getitem) and user.args[1] == idx: - return user - return None - - -# Returns the getitem node that extracts the idx-th element from node -def find_getitem(node: fx.Node, idx: int) -> fx.Node: - ret = find_getitem_maybe(node, idx) - assert ret is not None, f"Could not find getitem {idx} in node {node}" - return ret - - -# An auto-functionalization-aware utility for finding nodes with a specific op -def find_op_nodes(op: OpOverload, graph: fx.Graph) -> Iterator[fx.Node]: - if not op._schema.is_mutable: - yield from graph.find_nodes(op="call_function", target=op) - - for n in graph.find_nodes(op="call_function", target=auto_functionalized): - if n.args[0] == op: - yield n - - -# Asserts that the node only has one user and returns it -# Even if a node has only 1 user, it might share storage with another node, -# which might need to be taken into account. -def get_only_user(node: fx.Node) -> fx.Node: - assert len(node.users) == 1 - return next(iter(node.users)) diff --git a/python/sglang/srt/compilation/inductor_pass.py b/python/sglang/srt/compilation/inductor_pass.py index 48c1f4b95..a43cb38af 100644 --- a/python/sglang/srt/compilation/inductor_pass.py +++ b/python/sglang/srt/compilation/inductor_pass.py @@ -9,10 +9,9 @@ import logging import time import types from contextlib import contextmanager -from typing import Any, Callable, Optional, Union +from typing import Any, Optional, Union import torch -from torch import fx from torch._dynamo.utils import lazy_format_graph_code from torch._inductor.custom_graph_pass import CustomGraphPass @@ -93,25 +92,6 @@ class InductorPass(CustomGraphPass): return True -class CallableInductorPass(InductorPass): - """ - This class is a wrapper for a callable that automatically provides an - implementation of the UUID. - """ - - def __init__( - self, callable: Callable[[fx.Graph], None], uuid: Optional[Any] = None - ): - self.callable = callable - self._uuid = self.hash_source(callable) if uuid is None else uuid - - def __call__(self, graph: torch.fx.Graph): - self.callable(graph) - - def uuid(self) -> Any: - return self._uuid - - class SGLangInductorPass(InductorPass): def __init__( self, @@ -128,12 +108,3 @@ class SGLangInductorPass(InductorPass): self._end_time = time.perf_counter_ns() duration_ms = float(self._end_time - self._start_time) / 1.0e6 logger.debug("%s completed in %.1f ms", self.pass_name, duration_ms) - - -class PrinterInductorPass(SGLangInductorPass): - def __init__(self, name: str): - super().__init__() - self.name = name - - def __call__(self, graph: torch.fx.Graph): - self.dump_graph(graph, self.name) diff --git a/python/sglang/srt/compilation/pass_manager.py b/python/sglang/srt/compilation/pass_manager.py index 34e355169..184b2ba73 100644 --- a/python/sglang/srt/compilation/pass_manager.py +++ b/python/sglang/srt/compilation/pass_manager.py @@ -2,11 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # Adapted from https://github.com/vllm-project/vllm/blob/v0.10.0/vllm/compilation/pass_manager.py -import logging - from torch import fx as fx -from sglang.srt.compilation.fix_functionalization import FixFunctionalizationPass from sglang.srt.compilation.inductor_pass import ( CustomGraphPass, InductorPass, @@ -14,22 +11,11 @@ from sglang.srt.compilation.inductor_pass import ( get_pass_context, ) -logger = logging.getLogger(__name__) - class PostGradPassManager(CustomGraphPass): - """ - The pass manager for post-grad passes. - It handles configuration, adding custom passes, and running passes. - It supports uuid for the Inductor code cache. That includes torch<2.6 - support using pickling (in .inductor_pass.CustomGraphPass). + """Run post-grad passes in insertion order for the current runtime shape. - The order of the post-grad post-passes is: - 1. passes (constructor parameter) - 2. default passes (NoopEliminationPass, FusionPass) - 3. config["post_grad_custom_post_pass"] (if it exists) - 4. fix_functionalization - This way, all passes operate on a functionalized graph. + The ordered pass UUIDs identify the manager in Inductor's code cache. """ def __init__(self): @@ -41,15 +27,6 @@ class PostGradPassManager(CustomGraphPass): if pass_.is_applicable_for_shape(shape): pass_(graph) - # always run fix_functionalization last - self.fix_functionalization(graph) - - def configure( - self, - ): - self.pass_config = dict() - self.fix_functionalization = FixFunctionalizationPass() - def add(self, pass_: InductorPass): assert isinstance(pass_, InductorPass) self.passes.append(pass_) @@ -58,11 +35,6 @@ class PostGradPassManager(CustomGraphPass): """ The PostGradPassManager is set as a custom pass in the Inductor and affects compilation caching. Its uuid depends on the UUIDs of all - dependent passes and the pass config. See InductorPass for more info. + dependent passes. See InductorPass for more info. """ - pass_manager_uuid = "fshdakhsa" - state = {"pass_config": pass_manager_uuid, "passes": []} - for pass_ in self.passes: - state["passes"].append(pass_.uuid()) - state["passes"].append(self.fix_functionalization.uuid()) - return InductorPass.hash_dict(state) + return InductorPass.hash_dict({"passes": [p.uuid() for p in self.passes]}) diff --git a/python/sglang/srt/managers/mm_utils.py b/python/sglang/srt/managers/mm_utils.py index 4a9d7a29e..0e0373fd2 100644 --- a/python/sglang/srt/managers/mm_utils.py +++ b/python/sglang/srt/managers/mm_utils.py @@ -756,66 +756,6 @@ def general_mm_embed_routine( return hidden_states -def get_multimodal_data_bounds( - input_ids: torch.Tensor, pad_values: List[int], token_pairs: List[Tuple[int, int]] -) -> torch.Tensor: - """ - Returns a tensor indicating the bounds of multimodal data (images, video, audio, etc.) - - Returns: - [bounds_count, 2] - """ - # All the multimodal data in the batch should share the same special bound token ids. - start_tokens = {s for s, _e in token_pairs} - end_tokens = {e for _s, e in token_pairs} - - assert all(isinstance(t, int) for t in start_tokens) - assert all(isinstance(t, int) for t in end_tokens) - - start_cond = torch.isin( - input_ids, torch.as_tensor(start_tokens, device=input_ids.device) - ) - end_cond = torch.isin( - input_ids, torch.as_tensor(end_tokens, device=input_ids.device) - ) - - (data_start_tokens,) = torch.where(start_cond) - (data_end_tokens,) = torch.where(end_cond) - - data_start_tokens_cpu = data_start_tokens.cpu().tolist() - data_end_tokens_cpu = data_end_tokens.cpu().tolist() - - # the im_start_id sometimes can be cached as prefix, but it is needed for the embedding of the multimodal data - if len(data_start_tokens_cpu) != len(data_end_tokens_cpu): - if ( - len(data_start_tokens_cpu) + 1 == len(data_end_tokens_cpu) - and input_ids[0].item() in pad_values - and data_end_tokens_cpu - and data_start_tokens_cpu - and data_end_tokens_cpu[0] < data_start_tokens_cpu[0] - ): - data_start_tokens_cpu.insert(0, 0) - valid_mm_data_nums = min(len(data_start_tokens_cpu), len(data_end_tokens_cpu)) - - if valid_mm_data_nums == 0: - return torch.zeros((0, 2), device=input_ids.device) - - # Filter out pairs where start_token >= end_token - valid_pairs = [] - for i in range(valid_mm_data_nums): - start_token = data_start_tokens_cpu[i] - end_token = data_end_tokens_cpu[i] - if start_token < end_token: - valid_pairs.append((start_token + 1, end_token - 1)) - - if not valid_pairs: - return torch.zeros((0, 2), device=input_ids.device) - - # Convert valid pairs to tensor - valid_pairs_tensor = torch.as_tensor(valid_pairs, device=input_ids.device) - return valid_pairs_tensor - - def data_hash(data) -> int: hash_bytes = hashlib.sha256(data).digest()[:8] return int.from_bytes(hash_bytes, byteorder="big", signed=False) diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 18fa575ad..844d5f1b6 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -1368,24 +1368,6 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): f"{token_id}; valid range is [0, {vocab_size})." ) - def _validate_input_ids_in_vocab( - self, input_ids: Union[List[int], List[List[int]]], vocab_size: int - ) -> None: - # Handle both single sequence and batch of sequences - if isinstance(input_ids[0], list): - # Batch of sequences - for seq in input_ids: - if any(id >= vocab_size for id in seq): - raise ValueError( - f"The input_ids {seq} contains values greater than the vocab size ({vocab_size})." - ) - else: - # Single sequence - if any(id >= vocab_size for id in input_ids): - raise ValueError( - f"The input_ids {input_ids} contains values greater than the vocab size ({vocab_size})." - ) - def _create_tokenized_object( self, obj: Union[GenerateReqInput, EmbeddingReqInput], diff --git a/python/sglang/srt/mem_cache/buffer_mode/buffer_page_cache.py b/python/sglang/srt/mem_cache/buffer_mode/buffer_page_cache.py deleted file mode 100644 index 883e6c71b..000000000 --- a/python/sglang/srt/mem_cache/buffer_mode/buffer_page_cache.py +++ /dev/null @@ -1,396 +0,0 @@ -"""Refcounted content-addressed page cache over the buffer-mode host pool. - -NOT WIRED YET: staged spans register as ``(pool, page_hash) -> (slots, -refcount)`` so prefetches can be served zero-copy from local staging -(write-around / promote-on-read retention, zero-ref LRU reclaim); keys are -the content-chained page hashes, so entries survive node deletion, splits, -and recompute. Counterpart of ``StorageExistenceCache`` (beliefs about -STORAGE, dedupes writes); this tracks LOCAL HOST RAM and dedupes loads. - -TP determinism: replicas must stay identical across attention ranks — the -cache feeds scheduler-visible structure, so divergence is a collective -hang, not a soft miss. Preconditions when wiring: (1) mutate only on the -scheduler thread at lockstep points; (2) rank-reduce any fold anchored by a -per-rank storage outcome (hit count, revoke) before it picks a mutation; -(3) controller queues stay FIFO and single-threaded so MIN-count drains -process the same prefix on every rank. -""" - -from __future__ import annotations - -from collections import OrderedDict -from typing import Callable, Optional, Sequence - -import torch - -from sglang.srt.mem_cache.hicache_storage import ( - PoolHitPolicy, - PoolName, - PoolTransfer, -) - - -class _PageRef: - """One cached page: host slot span, reader refcount, and whether to - retain at refs==0 (write-around: write staging frees at its storage ack - unless a read promoted it). ``first_slot`` mirrors ``slots[0]`` as a - plain int because per-page tensor-scalar reads are too slow in - ``release``.""" - - __slots__ = ("slots", "first_slot", "refs", "retain") - - def __init__( - self, slots: torch.Tensor, first_slot: int, refs: int = 1, retain: bool = True - ): - self.slots = slots - self.first_slot = first_slot - self.refs = refs - self.retain = retain - - -class BufferPageCache: - def __init__(self) -> None: - # (pool, page_hash) -> _PageRef; slots stay allocated in the host - # pool for as long as the entry exists. - self._entries: dict[tuple[str, str], _PageRef] = {} - # Per-pool zero-ref LRU (head = coldest): reclaim victims. - self._zero_ref: dict[str, OrderedDict[str, None]] = {} - # Per-pool slot tokens held by the cache (refed + zero-ref). - self._held_tokens: dict[str, int] = {} - # Per-pool slot tokens at refs=0 (reclaimable under pressure). - self._zero_ref_tokens: dict[str, int] = {} - - def __len__(self) -> int: - return len(self._entries) - - def num_zero_ref_pages(self) -> int: - return sum(len(lru) for lru in self._zero_ref.values()) - - def held_tokens(self, pool: str) -> int: - return self._held_tokens.get(pool, 0) - - def zero_ref_tokens(self, pool: str) -> int: - """Slot tokens reclaimable right now (zero-ref cached pages). - Occupancy/rate-limit gates must treat these as free-able, not used: - a pool full of zero-ref cache is one reclaim away from empty.""" - return self._zero_ref_tokens.get(pool, 0) - - def register( - self, - pool: str, - hashes: Sequence[str], - host_indices: torch.Tensor, - page_size: int, - retain: bool = True, - ) -> int: - """Cache a staged span, one entry per page, refs=1 (the staging op); - returns the number of pages newly cached. ``retain=False`` = - write-around (freed at last ref unless a read hit promotes it); a - duplicate hash keeps the existing entry and the newcomer's slots - stay op-owned for a raw free at release.""" - assert len(host_indices) == len(hashes) * page_size - registered = 0 - entries = self._entries - # One batched read of the page-boundary slot ids (see _PageRef). - first_slots = host_indices[::page_size].tolist() - for i, page_hash in enumerate(hashes): - key = (pool, page_hash) - existing = entries.get(key) - if existing is not None: - existing.retain = existing.retain or retain - continue - entries[key] = _PageRef( - host_indices[i * page_size : (i + 1) * page_size], - first_slots[i], - retain=retain, - ) - registered += 1 - if registered: - self._held_tokens[pool] = ( - self._held_tokens.get(pool, 0) + registered * page_size - ) - return registered - - def contains(self, pool: str, page_hash: str) -> bool: - """Non-mutating presence probe (no LRU touch).""" - return (pool, page_hash) in self._entries - - def peek_run_len(self, pool: str, hashes: Sequence[str]) -> int: - """Length of the leading run of cached pages. Non-mutating.""" - entries = self._entries - run = 0 - for page_hash in hashes: - if (pool, page_hash) not in entries: - break - run += 1 - return run - - def acquire(self, pool: str, hashes: Sequence[str]) -> Optional[torch.Tensor]: - """refs++ on every page and return the gathered slot tensor (pages - expanded to token slots, in page order). All-or-nothing: returns - None without mutating if any page is missing.""" - entries = self._entries - refs = [] - for page_hash in hashes: - entry = entries.get((pool, page_hash)) - if entry is None: - return None - refs.append(entry) - zero_ref = self._zero_ref.get(pool) - for page_hash, entry in zip(hashes, refs): - if entry.refs == 0 and zero_ref is not None: - if zero_ref.pop(page_hash, None) is not None: - self._zero_ref_tokens[pool] -= len(entry.slots) - entry.refs += 1 - # Read demand proven: promote write-around pages to retained. - entry.retain = True - return torch.cat([entry.slots for entry in refs]) - - def release( - self, - pool: str, - hashes: Sequence[str], - host_indices: torch.Tensor, - page_size: int, - ) -> Optional[torch.Tensor]: - """Drop one ref per page of a span: at refs==0 retained pages move - to the zero-ref LRU tail while write-around pages return their - slots for an immediate free. Duplicate-staging slots (canonical - entry lives elsewhere) are returned for a raw free without touching - the canonical refcount.""" - assert len(host_indices) == len(hashes) * page_size - leftover: list[torch.Tensor] = [] - entries = self._entries - # One batched read of the page-boundary slot ids (see _PageRef). - first_slots = host_indices[::page_size].tolist() - for i, page_hash in enumerate(hashes): - entry = entries.get((pool, page_hash)) - if entry is None or entry.first_slot != first_slots[i]: - leftover.append(host_indices[i * page_size : (i + 1) * page_size]) - continue - assert entry.refs > 0, "release without a matching acquire/register" - entry.refs -= 1 - if entry.refs == 0: - if entry.retain: - self._zero_ref.setdefault(pool, OrderedDict())[page_hash] = None - self._zero_ref_tokens[pool] = self._zero_ref_tokens.get( - pool, 0 - ) + len(entry.slots) - else: - del entries[(pool, page_hash)] - self._held_tokens[pool] -= len(entry.slots) - leftover.append(entry.slots) - if not leftover: - return None - return torch.cat(leftover) - - def reclaim( - self, - pool: str, - need_tokens: int, - free: Callable[[torch.Tensor], int], - ) -> int: - """Pop zero-ref LRU heads, free their slots back to the host pool, - and drop the entries. Called under allocation pressure only. Returns - the number of slot tokens freed (may undershoot when everything - left is refed).""" - zero_ref = self._zero_ref.get(pool) - if not zero_ref or need_tokens <= 0: - return 0 - freed = 0 - batch: list[torch.Tensor] = [] - while zero_ref and freed < need_tokens: - page_hash, _ = zero_ref.popitem(last=False) - entry = self._entries.pop((pool, page_hash)) - batch.append(entry.slots) - freed += len(entry.slots) - if batch: - free(torch.cat(batch)) - self._held_tokens[pool] -= freed - self._zero_ref_tokens[pool] -= freed - return freed - - -class BufferPageCacheOps: - """Pool-facing operations over a :class:`BufferPageCache`: span/hold - registration and release keyed the way the storage write keys them, - pressure reclaim, and the SWA-folded continuation fold. The caller owns - the collectives — rank-reduce any fold anchored by a per-rank storage - outcome before acting on it (see the module docstring).""" - - def __init__( - self, - page_cache: BufferPageCache, - mem_pool_host, - sw_window_pages_fn: Callable[[], int], - ): - # Rebound by the owner when the structure is recreated (reset). - self.page_cache = page_cache - self._mem_pool_host = mem_pool_host - # SWA window in KV pages when SWA stages through a host pool - # (0 = KV-only: no trailing window in the fold). - self._sw_window_pages_fn = sw_window_pages_fn - - def aux_window_keys( - self, hash_values: list[str], transfer: PoolTransfer - ) -> Optional[list[str]]: - """Trailing KV page hashes keying an aux transfer's staged window - (one key per aux-pool page), recomputed from the rank-synced span - hashes so registration and release always agree across ranks.""" - if transfer.host_indices is None or transfer.host_indices.numel() == 0: - return None - if transfer.indices_from_pool is not None: - return None # sidecar rides another pool's slots; nothing to key - entry = self._mem_pool_host.entry_map.get(transfer.name) - if entry is None: - return None - pool_page_size = entry.host_pool.page_size - num_keys = len(transfer.host_indices) // pool_page_size - if num_keys == 0 or num_keys > len(hash_values): - return None - return hash_values[-num_keys:] - - def register_span( - self, - pool: PoolName, - hashes: list[str], - host_indices: torch.Tensor, - retain: bool = True, - ) -> None: - """Cache a page-aligned staged span (refs=1 for the staging op).""" - if not hashes: - return - entry = self._mem_pool_host.entry_map.get(pool) - if entry is None: - return - self.page_cache.register( - pool, hashes, host_indices, entry.host_pool.page_size, retain=retain - ) - - def release_span( - self, - pool: PoolName, - hashes: list[str], - host_indices: torch.Tensor, - ) -> None: - """Drop the staging op's ref on a span; zero-ref pages stay cached - (servable) until pressure reclaims them. Op-owned duplicate slots - (their hash was cached elsewhere) are freed raw, as before.""" - if host_indices is None or host_indices.numel() == 0: - return - entry = self._mem_pool_host.entry_map.get(pool) - if entry is None: - return - if not hashes: - entry.host_pool.free(host_indices) - return - leftover = self.page_cache.release( - pool, hashes, host_indices, entry.host_pool.page_size - ) - if leftover is not None and leftover.numel() > 0: - entry.host_pool.free(leftover) - - def register_hold( - self, - hash_values: list[str], - host_indices: torch.Tensor, - aux_xfers: list[PoolTransfer], - retain: bool = True, - ) -> None: - """Register a staged KV span plus its aux windows (SWA/Mamba states - keyed by their trailing KV page hashes, same keying the storage - write uses). ``retain=False`` = write-around: servable only while - the staging op pins the slots, freed at the last release unless a - read hit promotes it.""" - self.register_span(PoolName.KV, hash_values, host_indices, retain=retain) - for transfer in aux_xfers: - keys = self.aux_window_keys(hash_values, transfer) - if keys is not None: - self.register_span( - transfer.name, keys, transfer.host_indices, retain=retain - ) - - def release_hold( - self, - hash_values: list[str], - host_indices: torch.Tensor, - aux_xfers: list[PoolTransfer], - ) -> None: - """Mirror of register_hold for every hold retirement path - (storage-ack, fill H2D-ack, staged drop, abort).""" - self.release_span(PoolName.KV, hash_values, host_indices) - for transfer in aux_xfers: - if transfer.indices_from_pool is not None: - continue - keys = self.aux_window_keys(hash_values, transfer) - self.release_span(transfer.name, keys or [], transfer.host_indices) - - def reclaim(self, pool: PoolName, num_tokens: int) -> int: - """Free just enough zero-ref cached pages for an allocation of - num_tokens to succeed. Scheduler-thread only (lockstep pressure - points: staging-hit alloc, prepare_prefetch, cc.write).""" - entry = self._mem_pool_host.entry_map.get(pool) - if entry is None: - return 0 - shortfall = num_tokens - entry.host_pool.available_size() - if shortfall <= 0: - return 0 - return self.page_cache.reclaim(pool, shortfall, entry.host_pool.free) - - def continuation_run(self, chain: list[str], start_pages: int) -> int: - """Longest cached run continuing the span at page ``start_pages`` - (0 = leading run), folded for SWA: the joint span's trailing window - must be fully cache-servable, mirroring batch_exists_v2's - trailing_pages fold. Non-mutating and rank-deterministic.""" - page_cache = self.page_cache - kv_run = page_cache.peek_run_len(PoolName.KV, chain[start_pages:]) - if kv_run == 0: - return 0 - sw_pages = self._sw_window_pages_fn() - if sw_pages == 0: - return kv_run - for cont in range(kv_run, 0, -1): - joint = start_pages + cont - window = min(sw_pages, joint) - if cont < window: - # Window straddles into the head; only possible for - # anchored runs, and shrinking cont cannot fix it. - break - if all( - page_cache.contains(PoolName.SWA, chain[i]) - for i in range(joint - window, joint) - ): - return cont - return 0 - - def acquire_span( - self, chain: list[str], start_pages: int, cont_pages: int - ) -> Optional[tuple[torch.Tensor, list[PoolTransfer]]]: - """Acquire a folded continuation run: its KV pages plus the JOINT - span's trailing SWA window (refs++ on every page). Returns - (kv_slots, aux_xfers) or None (with no refs held) if a page - vanished since the fold — defensive; fold and acquire run in the - same lockstep step.""" - page_cache = self.page_cache - cont_hashes = list(chain[start_pages : start_pages + cont_pages]) - kv_slots = page_cache.acquire(PoolName.KV, cont_hashes) - if kv_slots is None: - return None - aux_xfers: list[PoolTransfer] = [] - sw_pages = self._sw_window_pages_fn() - if sw_pages > 0: - joint = start_pages + cont_pages - window_hashes = list(chain[joint - min(sw_pages, joint) : joint]) - swa_slots = page_cache.acquire(PoolName.SWA, window_hashes) - if swa_slots is None: - self.release_span(PoolName.KV, cont_hashes, kv_slots) - return None - aux_xfers.append( - PoolTransfer( - name=PoolName.SWA, - host_indices=swa_slots, - keys=window_hashes, - hit_policy=PoolHitPolicy.TRAILING_PAGES, - ) - ) - return kv_slots, aux_xfers diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 5d832fc97..6c0009840 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -1092,29 +1092,6 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): lens_to_track = self.mamba_track_seqlens - self.extend_prefix_lens return (lens_to_track // chunk_size) * chunk_size - def merge_mm_inputs(self) -> Optional[MultimodalInputs]: - """ - Merge all multimodal inputs in the batch into a single MultiModalInputs object. - - Returns: - if none, current batch contains no multimodal input - - """ - if not self.mm_inputs or all(x is None for x in self.mm_inputs): - return None - # Filter out None values - valid_inputs = [x for x in self.mm_inputs if x is not None] - - # TODO: is it expensive? - # a workaround to avoid importing `MultimodalInputs` - merged = valid_inputs[0].__class__(mm_items=[]) - - # Merge remaining inputs - for mm_input in valid_inputs: - merged.merge(mm_input) - - return merged - def contains_image_inputs(self) -> bool: if self.mm_inputs is None: return False diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py index 6bc867c35..bb6cbba5c 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/python/sglang/srt/model_loader/weight_utils.py @@ -17,10 +17,8 @@ import re import struct import tempfile import threading -from collections import defaultdict from pathlib import Path from typing import ( - Any, Callable, Dict, Generator, @@ -172,57 +170,6 @@ def get_lock( return lock -def _shared_pointers(tensors): - ptrs = defaultdict(list) - for k, v in tensors.items(): - ptrs[v.data_ptr()].append(k) - failing = [] - for _, names in ptrs.items(): - if len(names) > 1: - failing.append(names) - return failing - - -def convert_bin_to_safetensor_file( - pt_filename: str, - sf_filename: str, -) -> None: - loaded = torch.load(pt_filename, map_location="cpu", weights_only=True) - if "state_dict" in loaded: - loaded = loaded["state_dict"] - shared = _shared_pointers(loaded) - for shared_weights in shared: - for name in shared_weights[1:]: - loaded.pop(name) - - # For tensors to be contiguous - loaded = {k: v.contiguous() for k, v in loaded.items()} - - dirname = os.path.dirname(sf_filename) - os.makedirs(dirname, exist_ok=True) - - from safetensors.torch import save_file - - save_file(loaded, sf_filename, metadata={"format": "pt"}) - - # check file size - sf_size = os.stat(sf_filename).st_size - pt_size = os.stat(pt_filename).st_size - if (sf_size - pt_size) / pt_size > 0.01: - raise RuntimeError(f"""The file size different is more than 1%: - - {sf_filename}: {sf_size} - - {pt_filename}: {pt_size} - """) - - # check if the tensors are the same - reloaded = safetensors.torch.load_file(sf_filename) - for k in loaded: - pt_tensor = loaded[k] - sf_tensor = reloaded[k] - if not torch.equal(pt_tensor, sf_tensor): - raise RuntimeError(f"The output tensors do not match for key {k}") - - def replace_prefix(key: str, prefix_mapping: dict[str, str]) -> str: for prefix, new_prefix in prefix_mapping.items(): if key.startswith(prefix): @@ -1263,13 +1210,10 @@ def fastsafetensors_weights_iterator( loader.add_filenames(rank_file_map) try: fb = loader.copy_files_to_device() - try: - keys = list(fb.key_to_rank_lidx.keys()) - for k in keys: - t = fb.get_tensor(k) - yield k, t - finally: - pass + keys = list(fb.key_to_rank_lidx.keys()) + for k in keys: + t = fb.get_tensor(k) + yield k, t finally: loader.close() if drop_cache_after_load: @@ -1277,50 +1221,6 @@ def fastsafetensors_weights_iterator( _drop_file_cache_after_load(loaded_file) -def multi_thread_safetensors_weights_iterator( - hf_weights_files: List[str], - max_workers: int, - disable_mmap: bool = False, - drop_cache_after_load: bool = False, -) -> Generator[Tuple[str, torch.Tensor], None, None]: - """Multi-Thread iterate over the weights in the model safetensor files.""" - enable_tqdm = ( - not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0 - ) - - def _load_file(st_file: str): - if disable_mmap: - with open(st_file, "rb") as f: - result = safetensors.torch.load(f.read()) - else: - with safetensors.safe_open(st_file, framework="pt", device="cpu") as f: - result = {k: f.get_tensor(k) for k in f.keys()} - - return st_file, result - - with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: - futures = [executor.submit(_load_file, st_file) for st_file in hf_weights_files] - - if enable_tqdm: - futures_iter = tqdm( - concurrent.futures.as_completed(futures), - total=len(hf_weights_files), - desc="Multi-thread loading shards", - disable=not enable_tqdm, - bar_format=BAR_FORMAT, - ) - else: - futures_iter = concurrent.futures.as_completed(futures) - - for future in futures_iter: - st_file, state_dict = future.result() - for name, param in state_dict.items(): - yield name, param - del state_dict - if drop_cache_after_load: - _drop_file_cache_after_load(st_file) - - def buffered_multi_thread_safetensors_weights_iterator( hf_weights_files: List[str], max_workers: int, @@ -1569,55 +1469,19 @@ def gguf_quant_weights_iterator( yield name, param -def convert_pyslice_to_tensor(x: Any) -> torch.Tensor: - """convert PySafeSlice object from safetensors to torch.Tensor - - PySafeSlice object supports indexing, which is done before loading the - actual tensor and can reduce the amount of memory being read into the - memory. However, it does not support more advanced functionalities - like `.view()` or `.t()`. Therefore, if we need to modify the loaded - tensor with these more complicated operators, we need to convert to - tensor first. - """ - if not isinstance(x, torch.Tensor): - x = x[:] - return x - - def default_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: """Default weight loader.""" - try: - if param.numel() == 1 and loaded_weight.numel() == 1: - # Sometimes scalar values aren't considered tensors with shapes - # so if both param and loaded_weight are a scalar, - # "broadcast" instead of copy - param.data.fill_(loaded_weight.item()) - else: - assert param.size() == loaded_weight.size(), ( - f"Attempted to load weight ({loaded_weight.size()}) " - f"into parameter ({param.size()})" - ) - - param.data.copy_(loaded_weight) - except Exception: - # NOTE: This exception is added for the purpose of setting breakpoint to - # debug weight loading issues. - raise - - -def row_parallel_weight_loader( - param: torch.Tensor, loaded_weight: torch.Tensor -) -> None: - """Load weights that are row-parallelized.""" - tp_rank = get_parallel().tp_rank - shard_dim = 0 if param.dim() != 1 else None - - if shard_dim is not None: - shard_size = param.data.shape[shard_dim] - start_idx = tp_rank * shard_size - loaded_weight = loaded_weight.narrow(shard_dim, start_idx, shard_size) - - return default_weight_loader(param, loaded_weight) + if param.numel() == 1 and loaded_weight.numel() == 1: + # Sometimes scalar values aren't considered tensors with shapes + # so if both param and loaded_weight are a scalar, + # "broadcast" instead of copy + param.data.fill_(loaded_weight.item()) + else: + assert param.size() == loaded_weight.size(), ( + f"Attempted to load weight ({loaded_weight.size()}) " + f"into parameter ({param.size()})" + ) + param.data.copy_(loaded_weight) LoaderFunction = Callable[[torch.Tensor, torch.Tensor], torch.Tensor] diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 62ad5159f..5509adfde 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -1145,29 +1145,6 @@ def get_cuda_driver_bindings(): return cuda_driver -def get_physical_device_id(pytorch_device_id: int) -> int: - """ - Convert PyTorch logical device ID to physical device ID. - - When CUDA_VISIBLE_DEVICES is set, maps the logical device ID (as seen by PyTorch) - to the actual physical device ID. If CUDA_VISIBLE_DEVICES is not set, returns - the device ID unchanged. - - Args: - pytorch_device_id: The logical device ID from PyTorch (e.g., torch.cuda.current_device()) - - Returns: - The physical device ID - """ - device_idx = int(pytorch_device_id) - cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", None) - if cuda_visible_devices: - device_list = cuda_visible_devices.split(",") - return int(device_list[device_idx]) - else: - return device_idx - - def get_device_sm_nvidia_smi(): try: # Run nvidia-smi command and capture output @@ -1431,25 +1408,6 @@ def mark_end(name): time_infos[name].pretty_print() -def calculate_time(show=False, min_cost_ms=0.0): - def wrapper(func): - def inner_func(*args, **kwargs): - torch.cuda.synchronize() - if show: - start_time = time.perf_counter() - result = func(*args, **kwargs) - torch.cuda.synchronize() - if show: - cost_time = (time.perf_counter() - start_time) * 1000 - if cost_time > min_cost_ms: - print(f"Function {func.__name__} took {cost_time} ms to run.") - return result - - return inner_func - - return wrapper - - class LayerFn(Protocol): def __call__(self, idx: int, prefix: str) -> torch.nn.Module: ... @@ -1498,24 +1456,6 @@ def make_layers( return modules, start_layer, end_layer -def make_layers_non_pp( - num_hidden_layers: int, - layer_fn: LayerFn, - prefix: str = "", -) -> torch.nn.ModuleList: - from sglang.srt.utils.offloader import get_offloader - - layers = torch.nn.ModuleList( - get_offloader().wrap_modules( - ( - layer_fn(idx=idx, prefix=add_prefix(idx, prefix)) - for idx in range(num_hidden_layers) - ) - ) - ) - return layers - - def set_random_seed(seed: int) -> None: """Set the random seed for all libraries.""" random.seed(seed) @@ -2818,11 +2758,6 @@ def init_custom_process_group( return pg -def crash_on_warnings(): - # Crash on warning if we are running CI tests - return get_bool_env_var("SGLANG_IS_IN_CI") - - @functools.lru_cache(None) def print_warning_once(msg: str) -> None: # Set the stacklevel to 2 to print the caller's line info @@ -2956,26 +2891,6 @@ def set_gpu_proc_affinity( logger.info(f"Process {pid} gpu_id {gpu_id} is running on CPUs: {p.cpu_affinity()}") -def permute_weight(x: torch.Tensor) -> torch.Tensor: - b_ = x.shape[0] - n_ = x.shape[1] - k_ = x.shape[2] - - x_ = x - if x.dtype == torch.bfloat16 or x.dtype == torch.float16: - x_ = x_.view(int(b_), int(n_ / 16), 16, int(k_ / 32), 4, 8) - elif x.dtype == torch.float8_e4m3fnuz or x.dtype == torch.int8: - x_ = x_.view(int(b_), int(n_ / 16), 16, int(k_ / 64), 4, 16) - else: - # return x_ - x_ = x_.view(int(b_), int(n_ / 16), 16, int(k_ / 8), 2, 4) - - x_ = x_.permute(0, 1, 3, 4, 2, 5) - x_ = x_.contiguous() - x_ = x_.view(*x.shape) - return x_ - - class MultiprocessingSerializer: @staticmethod def serialize(obj, output_str: bool = False): @@ -3141,30 +3056,6 @@ def safe_pickle_loads(data): return SafeUnpickler(io.BytesIO(buf)).load() -def debug_timing(func): - # todo: replace with a more organized instrumentation - def wrapper(*args, **kwargs): - if logger.isEnabledFor(logging.DEBUG): - tic = torch.cuda.Event(enable_timing=True) - toc = torch.cuda.Event(enable_timing=True) - tic.record() - result = func(*args, **kwargs) - toc.record() - toc.synchronize() # Wait for the function to complete without synchronizing all ops on the GPU - elapsed = tic.elapsed_time(toc) - indices = kwargs.get("indices", args[1] if len(args) > 1 else None) - num_tokens = len(indices) if indices is not None else 0 - throughput = num_tokens / elapsed * 1000 if elapsed > 0 else 0 - logger.debug( - f"Transfer time: {elapsed} ms, throughput: {throughput} tokens/s" - ) - return result - else: - return func(*args, **kwargs) - - return wrapper - - def nullable_str(val: str): if not val or val == "None": return None @@ -3705,35 +3596,6 @@ def is_no_spec_infer_or_topk_one(cfg): ) -def is_fa3_default_architecture(hf_config): - architectures = getattr(hf_config, "architectures", None) - if not isinstance(architectures, list) or not architectures: - return False - default_archs = { - "Llama4ForConditionalGeneration", - "LlamaForCausalLM", - "Olmo2ForCausalLM", - "Gemma2ForCausalLM", - "Gemma3ForConditionalGeneration", - "MixtralForCausalLM", - "Qwen2ForCausalLM", - "Qwen3ForCausalLM", - "Qwen3MoeForCausalLM", - "Qwen3VLForConditionalGeneration", - "Qwen3VLMoeForConditionalGeneration", - "Glm4MoeForCausalLM", - "Glm4vForConditionalGeneration", - "Glm4vMoeForConditionalGeneration", - "GlmOcrForConditionalGeneration", - "Step3VLForConditionalGeneration", - "StepVLForConditionalGeneration", - "Step3p7ForConditionalGeneration", - "MiMoV2ForCausalLM", - "MiMoV2FlashForCausalLM", - } - return architectures[0] in default_archs - - # Can be more general if it is used in multiple places (keep it simple and thus not general now) class BumpAllocator: def __init__(self, buffer_size: int, dtype, device): diff --git a/test/registered/unit/server_args/test_resolution_reads_the_declarations.py b/test/registered/unit/server_args/test_resolution_reads_the_declarations.py deleted file mode 100644 index c608c43ef..000000000 --- a/test/registered/unit/server_args/test_resolution_reads_the_declarations.py +++ /dev/null @@ -1,595 +0,0 @@ -"""Resolution reads its own decisions, not the record's fields. - -`declare_resolution` records a decision in the declaration stash and writes -nothing. The fields keep what the caller passed, so a resolver that reads a -field another resolver may have decided reads the raw input -- silently, and -only on the configurations where that other resolver fires. The whole pipeline -therefore reads through `resolving_view` (or `resolved_view`, which is -the same view after resolution has finished), and this pins that there is -nothing left reading a field directly. - -Subjects: every function in `arg_groups/` that takes a config, every -`ServerArgs` handler the dispatcher reaches, and every member of `ServerArgs` / -`PortArgs` -- the members are reached from the hooks and from business code, -which the handler walk cannot see, and a member that recomputes from a raw field -decides from what was typed. All three -are derived -- a new hook file, a new handler or a new member is covered the -moment it is written. Readers *outside* those -two -- the platform defaults, `ModelConfig`, the spec-algo hook -- are reached by -resolution too and have moved to the view as well, but enumerating them needs -the call-graph derivation `test_resolution_reads_no_bag` owns; this file pins -the two scopes it can derive exactly. -""" - -import ast -import pathlib -import re - -import msgspec -import msgspec.structs - -import sglang -from sglang.srt.server_args import ServerArgs -from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import CustomTestCase - -register_cpu_ci(est_time=45, suite="base-a-test-cpu") - -_SRT = pathlib.Path(sglang.__file__).resolve().parent / "srt" -_FIELDS = frozenset(field.name for field in msgspec.structs.fields(ServerArgs)) - -# Names a config travels under. `args` is included because the platform hooks -# use it; a false positive would be a function taking an argparse Namespace and -# reading an attribute that happens to be a ServerArgs field name, which the -# allowlist below would then have to carry. -_HOLDER_NAMES = frozenset({"server_args", "sa", "args"}) - - -def _holders(fn): - names = { - arg.arg - for arg in list(fn.args.posonlyargs) - + list(fn.args.args) - + list(fn.args.kwonlyargs) - if arg.arg in _HOLDER_NAMES - } - for arg in ( - list(fn.args.posonlyargs) + list(fn.args.args) + list(fn.args.kwonlyargs) - ): - annotation = arg.annotation - text = ( - annotation.value - if isinstance(annotation, ast.Constant) - else ( - annotation.id - if isinstance(annotation, ast.Name) - else annotation.attr - if isinstance(annotation, ast.Attribute) - else None - ) - ) - if text == "ServerArgs": - names.add(arg.arg) - return names - - -def _field_reads(fn, holders): - for node in ast.walk(fn): - if ( - isinstance(node, ast.Attribute) - and node.attr in _FIELDS - and isinstance(node.value, ast.Name) - and node.value.id in holders - and isinstance(node.ctx, ast.Load) - ): - yield node.lineno, node.attr - - -_DECLARERS = frozenset( - { - "declare_resolution", - "record_foreign_defaults", - } -) - - -def _declared_fields(): - """The fields resolution decides, read off every shape that reaches the stash. - - A keyword on a `declare_*` call is only one shape: the model-override and - post-process passes build a mapping instead (`MODEL_OVERRIDES` literals, - `overrides["dtype"] = ...`, a returned dict), and late resolution splats a - variable-keyed one. Deriving from keywords alone leaves nineteen fields - outside the subject set, `dtype` and `reasoning_parser` among them. - """ - fields = set() - # The declaration calls live wherever a resolver does; the mapping channels - # only exist where the override providers and post-process passes are. - keyword_sources = [_SRT / "server_args.py"] - for sub in ("arg_groups", "hardware_backend", "parser"): - keyword_sources += sorted((_SRT / sub).rglob("*.py")) - mapping_sources = {_SRT / "server_args.py", *(_SRT / "arg_groups").rglob("*.py")} - for path in keyword_sources: - tree = ast.parse(path.read_text(encoding="utf-8-sig")) - for node in ast.walk(tree): - # 1. `declare_resolution(sa, src, page_size=64)` and its siblings - if isinstance(node, ast.Call): - name = ( - node.func.id - if isinstance(node.func, ast.Name) - else getattr(node.func, "attr", None) - ) - if name in _DECLARERS: - for keyword in node.keywords: - if keyword.arg: - fields.add(keyword.arg) - elif isinstance(keyword.value, ast.Dict): - fields.update(_string_keys(keyword.value)) - # 2. every mapping literal in the files that declare through one: - # the MODEL_OVERRIDES tables, the dicts the override providers - # return, the ones the post-process passes build. Scanning - # unrelated files here would collect a plain kwarg dict - # (`tokenizer_config={"trust_remote_code": ...}`) and turn a - # passthrough read into a violation. - if isinstance(node, ast.Dict) and path in mapping_sources: - fields.update(_string_keys(node)) - # 3. `overrides["field"] = ...` - if ( - path in mapping_sources - and isinstance(node, ast.Assign) - and isinstance(node.targets[0], ast.Subscript) - and isinstance(node.targets[0].slice, ast.Constant) - and isinstance(node.targets[0].slice.value, str) - ): - fields.add(node.targets[0].slice.value) - return frozenset(fields & _FIELDS) - - -def _string_keys(node: ast.Dict) -> set: - return { - key.value - for key in node.keys - if isinstance(key, ast.Constant) and isinstance(key.value, str) - } - - -def _record_members(): - """Every member of `ServerArgs` / `PortArgs`, by class and name.""" - tree = ast.parse((_SRT / "server_args.py").read_text(encoding="utf-8-sig")) - members = {} - for node in ast.walk(tree): - if isinstance(node, ast.ClassDef) and node.name in ("ServerArgs", "PortArgs"): - for member in node.body: - if isinstance(member, (ast.FunctionDef, ast.AsyncFunctionDef)): - members[f"{node.name}.{member.name}"] = member - return members - - -def _config_reading_helpers(): - """Module functions that load a decided field off the config they are handed. - - A member that hands them `self`, or a call site that hands them a record, - reads the raw input through the callee -- the shape neither an attribute - scan nor a `getattr` scan can see, because the field name is spelled in the - helper and the record is spelled at the call site. - """ - decided = _declared_fields() - helpers = {} - sources = [_SRT / "server_args.py"] + sorted((_SRT / "arg_groups").rglob("*.py")) - for path in sources: - tree = ast.parse(path.read_text(encoding="utf-8-sig")) - for fn in ast.walk(tree): - if not isinstance(fn, (ast.FunctionDef, ast.AsyncFunctionDef)): - continue - params = { - arg.arg - for arg in list(fn.args.posonlyargs) - + list(fn.args.args) - + list(fn.args.kwonlyargs) - } - {"self", "cls"} - if not params: - continue - reads = { - node.attr - for node in ast.walk(fn) - if isinstance(node, ast.Attribute) - and isinstance(node.value, ast.Name) - and node.value.id in params - and isinstance(node.ctx, ast.Load) - and node.attr in decided - } - if reads: - helpers[fn.name] = sorted(reads) - return helpers - - -# The accessors that hand back the process-global record itself. A helper that -# is handed one of these reads the raw input exactly as a bare `self` would. -_RECORD_ACCESSORS = frozenset({"get_server_args", "global_server_args"}) - -# `self._server_args` is the same record under a private name; the scan has to -# see it or a reader inside the context object escapes every shape above. -_RECORD_ATTR = re.compile(r"^_*(server_args|sa)$") - - -def _record_arguments(node, aliases=frozenset()): - """The bare-record arguments of a call. - - Four spellings reach a helper with a record: the bare name (`self`, `sa`), - an attribute (`runner.server_args`), the process-global accessor called - inline (`get_server_args()`), and a local bound to either of the last two - earlier in the same function. - """ - out = [] - for arg in node.args: - if isinstance(arg, ast.Name) and arg.id in ("self", "server_args", "sa"): - out.append(arg.id) - elif isinstance(arg, ast.Attribute) and _RECORD_ATTR.match(arg.attr or ""): - out.append(ast.unparse(arg)) - elif ( - isinstance(arg, ast.Call) - and isinstance(arg.func, ast.Name) - and arg.func.id in _RECORD_ACCESSORS - ): - out.append(ast.unparse(arg)) - elif isinstance(arg, ast.Name) and arg.id in aliases: - out.append(arg.id) - return out - - -def _record_aliases(function): - """Locals bound to the record under a name of their own. - - `_sa = getattr(runner, "server_args", None)`, `cfg = get_server_args()` and - `engine_args = ServerArgs.from_cli_args(args)` all put the record behind a - name the argument scan does not recognise, so a later - `getattr(_sa, "")` reads what the operator typed. - """ - aliases = set() - for node in ast.walk(function): - if not isinstance(node, ast.Assign) or len(node.targets) != 1: - continue - target = node.targets[0] - if not isinstance(target, ast.Name): - continue - value = node.value - if isinstance(value, ast.Attribute): - if _RECORD_ATTR.match(value.attr or ""): - aliases.add(target.id) - continue - if not isinstance(value, ast.Call): - continue - func = value.func - if isinstance(func, ast.Name): - if func.id in _RECORD_ACCESSORS or func.id == "ServerArgs": - aliases.add(target.id) - elif ( - func.id == "getattr" - and len(value.args) >= 2 - and isinstance(value.args[1], ast.Constant) - and _RECORD_ATTR.match(str(value.args[1].value)) - ): - aliases.add(target.id) - elif ( - isinstance(func, ast.Attribute) - and func.attr == "from_cli_args" - and isinstance(func.value, ast.Name) - and func.value.id == "ServerArgs" - ): - aliases.add(target.id) - return aliases - - -# The one reader for which the raw field is the right answer. The gateway sizes -# its worker pool from the operator's requested replica count; `--dwdp-size` -# makes resolution declare a `dp_size` describing one multi-rank server's -# internal topology, so reading the decision there would spawn dp_size -# single-rank children and ask for dp_size^2 GPUs. A new entry here needs that -# kind of reason next to it. -_NO_RESOLVED_SURFACE = frozenset( - { - ( - "sgl-model-gateway/bindings/python/src/sglang_router/launch_server.py", - "server_args.dp_size", - ), - } -) - - -def _is_record_base(node, aliases): - """Is this expression the record itself? - - A local bound to one, a parameter that carries one (`server_args`, `sa`, - `engine_args`), or an attribute holding one (`self._server_args`). - """ - if isinstance(node, ast.Name): - return node.id in aliases - if isinstance(node, ast.Attribute): - return bool(_RECORD_ATTR.match(node.attr or "")) - return False - - -def _record_handoff_offenders(rel, tree, helpers, decided, is_record=False): - """Every way a decided leaf is reached through a record in one module. - - Two shapes, both scanned under the record aliases the function binds: - handing the record to a helper that loads a decided field, and loading one - off the alias directly (`alias.` or `getattr(alias, "")`). The - second is what the MiniMax backend spelled, and an argument scan cannot see - it -- the leaf never appears at a call site. - """ - offenders, seen = [], set() - - def record(lineno, text): - if (lineno, text) in seen: - return - seen.add((lineno, text)) - offenders.append(f"{rel}:{lineno} {text}") - - scopes = [(tree, frozenset())] + [ - (fn, _record_aliases(fn)) - for fn in ast.walk(tree) - if isinstance(fn, (ast.FunctionDef, ast.AsyncFunctionDef)) - ] - for scope, aliases in scopes: - for node in ast.walk(scope): - if isinstance(node, ast.Call): - name = ( - node.func.id - if isinstance(node.func, ast.Name) - else getattr(node.func, "attr", None) - ) - if name in helpers: - for arg in _record_arguments(node, aliases): - if arg == "self" and not is_record: - continue - record( - node.lineno, - f"{name}({arg}) reads {', '.join(helpers[name])}", - ) - if ( - isinstance(node.func, ast.Name) - and node.func.id == "getattr" - and len(node.args) >= 2 - and isinstance(node.args[0], ast.Name) - and node.args[0].id in aliases - and isinstance(node.args[1], ast.Constant) - and node.args[1].value in decided - ): - record( - node.lineno, - f'getattr({node.args[0].id}, "{node.args[1].value}")', - ) - elif ( - isinstance(node, ast.Attribute) - and isinstance(node.ctx, ast.Load) - and node.attr in decided - and _is_record_base(node.value, aliases) - ): - record(node.lineno, f"{ast.unparse(node.value)}.{node.attr}") - return offenders - - -# Source the scanner must read the same way whether or not the tree happens to -# contain these shapes today. The first four are the spellings that reached -# production and were converted; the last two are the legal forms next to them, -# which have to stay quiet or the guard is unusable. -_SPELLINGS = """ -def hands_the_alias_to_a_helper(runner): - _sa = getattr(runner, "server_args", None) - return m3_fp8_attn_gemm_enabled(_sa) - - -def loads_a_leaf_off_the_alias(runner): - _sa = getattr(runner, "server_args", None) - return getattr(_sa, "speculative_num_draft_tokens", None) - - -def reads_a_leaf_through_the_alias(runner): - sa_local = runner.server_args - return sa_local.attention_backend - - -def hands_the_accessor_to_a_helper(): - return attention_backends_of(get_server_args()) - - -def reads_the_view(runner): - cfg = resolving_view(runner.server_args) - return cfg.attention_backend - - -def reads_an_undecided_leaf(runner): - _sa = runner.server_args - return _sa.tp_size - - -def reads_a_leaf_off_a_private_attribute(self): - return self._server_args.attention_backend - - -def reads_a_leaf_off_a_constructed_record(cli): - engine_args = ServerArgs.from_cli_args(cli) - engine_args.resolve_once() - return engine_args.attention_backend -""" - - -class TestResolutionReadsTheDeclarations(CustomTestCase): - def test_no_hook_reads_a_field_off_the_record(self): - offenders = [] - files = sorted((_SRT / "arg_groups").glob("*.py")) - self.assertGreater(len(files), 5, "the hook scan found almost nothing") - for path in files: - rel = f"arg_groups/{path.name}" - tree = ast.parse(path.read_text(encoding="utf-8-sig")) - for fn in ast.walk(tree): - if not isinstance(fn, (ast.FunctionDef, ast.AsyncFunctionDef)): - continue - holders = _holders(fn) - if not holders: - continue - for lineno, field in _field_reads(fn, holders): - offenders.append(f"{rel}:{lineno} {fn.name} reads .{field}") - self.assertEqual( - offenders, - [], - "a resolution hook reads a field off the record; the field holds the " - "raw input, so this decides from what was typed rather than from " - "what resolution decided. Read `resolving_view(server_args)`:\n " - + "\n ".join(offenders), - ) - - def test_the_record_hosts_no_resolution_handler(self): - """The pipeline and every step it runs live under `arg_groups/`. - - While a step was a method, it could read a raw field off `self` and - `test_no_handler_reads_a_field_off_self` had to say it could not. There - is no such method left, so the invariant is now the stronger one: the - record hosts none of them. What the steps read is checked on the - package side, by `test_no_hook_reads_a_field_off_the_record`. - """ - handlers = sorted( - name - for name, node in _record_members().items() - if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) - and ( - name.split(".")[-1].startswith(("_handle_", "_validate_")) - or "resolution_pipeline" in name - ) - ) - self.assertEqual( - handlers, - [], - "a resolution handler is back on the record; it belongs in an " - "`arg_groups` family, where the package-side guards can see it:\n " - + "\n ".join(handlers), - ) - - def test_no_member_recomputes_from_a_raw_field(self): - decided = _declared_fields() - self.assertGreater( - len(decided), 100, f"the declaration set derived only {len(decided)} fields" - ) - members = _record_members() - # The floor is here to catch the scan collapsing, not to pin the - # class's size -- it drops as derived members move to their namespaces - # and become declarations rather than methods on the record. - self.assertGreater(len(members), 10, f"only {len(members)} members were found") - offenders = [] - for name, fn in sorted(members.items()): - holders = _holders(fn) | {"self"} - for lineno, field in _field_reads(fn, holders): - if field in decided: - offenders.append(f"server_args.py:{lineno} {name} reads .{field}") - self.assertEqual( - offenders, - [], - "a record member recomputes from a field resolution decides; the " - "field holds the raw input, so the member answers for what was " - "typed. Bind `cfg = resolving_view(self)` and read that:\n " - + "\n ".join(offenders), - ) - - def test_no_reader_hands_the_record_to_a_config_helper(self): - helpers = _config_reading_helpers() - self.assertGreater( - len(helpers), 5, f"the helper derivation found only {len(helpers)}" - ) - decided = _declared_fields() - offenders = [] - # `scripts/`, `examples/` and the gateway binding are outside the - # package but hold records they resolve themselves, and every reader - # this scan found in them was reading a field resolution fills in. - _REPO = _SRT.parent.parent.parent - roots = ( - [_SRT] - + [_SRT.parent / d for d in ("benchmark", "lang")] - + [ - _REPO / d - for d in ( - "scripts", - "examples", - "sgl-model-gateway/bindings/python/src", - ) - ] - ) - for root in roots: - if not root.exists(): - continue - for path in sorted(root.rglob("*.py")): - # Package files keep their `srt/...` spelling (the skips below - # key on it); the repo-level roots are named from the repo. - try: - rel = path.relative_to(_SRT.parent).as_posix() - except ValueError: - rel = path.relative_to(_REPO).as_posix() - if rel.startswith(("srt/arg_groups/", "multimodal_gen/")): - continue - try: - tree = ast.parse(path.read_text(encoding="utf-8-sig")) - except SyntaxError: - continue - offenders += _record_handoff_offenders( - rel, tree, helpers, decided, is_record=rel == "srt/server_args.py" - ) - offenders = [ - line - for line in offenders - if (line.split(":", 1)[0], line.split(" ", 1)[1]) - not in _NO_RESOLVED_SURFACE - ] - self.assertEqual( - offenders, - [], - "a caller hands the record to a helper that loads a field " - "resolution decides; the helper then reads the raw input. Hand it " - "`resolving_view(record)` (or the published bag):\n " - + "\n ".join(offenders), - ) - - def test_the_scan_sees_every_spelling_that_reached_production(self): - """Every spelling that reached production, pinned next to the scanner. - - A shape the scan stops seeing is a silent hole, so each one is listed - here with the legal forms beside it and the flagged set compared - exactly. - - What it does not reach: a record that arrives as a *parameter* and was - resolved by the caller (`scripts/playground/bench_speculative.py` hands - `main(args, server_args)` one). Binding that would need the call graph, - and naming a parameter `server_args` is also how the resolution-time - readers spell a view. - """ - helpers = _config_reading_helpers() - decided = _declared_fields() - for name in ("m3_fp8_attn_gemm_enabled", "attention_backends_of"): - self.assertIn(name, helpers, f"the helper derivation lost {name}") - for field in ("speculative_num_draft_tokens", "attention_backend"): - self.assertIn(field, decided, f"the declared set lost {field}") - - offenders = _record_handoff_offenders( - "sample.py", ast.parse(_SPELLINGS), helpers, decided - ) - flagged = {line.split(" ", 1)[1] for line in offenders} - self.assertEqual( - flagged, - { - "m3_fp8_attn_gemm_enabled(_sa)" - " reads " + ", ".join(helpers["m3_fp8_attn_gemm_enabled"]), - 'getattr(_sa, "speculative_num_draft_tokens")', - "sa_local.attention_backend", - "self._server_args.attention_backend", - "engine_args.attention_backend", - "attention_backends_of(get_server_args())" - " reads " + ", ".join(helpers["attention_backends_of"]), - }, - "the scan lost a spelling, or started flagging a legal one:\n " - + "\n ".join(sorted(flagged)), - ) - - -if __name__ == "__main__": - import unittest - - unittest.main() diff --git a/test/registered/unit/test_chain_read_ratchet.py b/test/registered/unit/test_chain_read_ratchet.py deleted file mode 100644 index 362eb008b..000000000 --- a/test/registered/unit/test_chain_read_ratchet.py +++ /dev/null @@ -1,583 +0,0 @@ -"""Nobody reaches the startup record through another object for a resolved value. - -The supplied-instance census counts three spellings, all of which start from a -`server_args` parameter -- the caller chose the object, which is the contract -that makes those reads defensible. This pins the fourth: `model_runner. -server_args.field`, `self.scheduler.server_args.field`, `tokenizer_manager. -server_args.field`. A reference lifted off whatever object happened to hold the -record carries no contract at all, and it was invisible to every census, which -is how it grew to 105 reads across 36 files unnoticed. - -They are gone, and this is what keeps them gone. Only fields resolution writes -are pinned: reading `model_runner.server_args.host` off the record answers with -what the caller asked for, which is what the record is for. The written set is -derived from the declaration sites rather than listed, so a field that stops -being resolution-written drops out on its own -- and a field that stops being -*declared* cannot slip out that way, because bare assignment during resolution -is refused by `server_args/test_resolution_declarations.py`. -""" - -import ast -import pathlib -import unittest - -import sglang -from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import CustomTestCase - -register_cpu_ci(est_time=28, suite="base-a-test-cpu") - -_PACKAGE = pathlib.Path(sglang.__file__).resolve().parent -_SRT = _PACKAGE / "srt" - -# The pipeline and its extension points: reading the in-flight record is their -# job, and they run before anything is published. -_OWNERS = ("server_args.py", "runtime_context.py", "arg_groups/") - -# Where the writers are. Resolution lives in `srt` -- nothing outside it -# declares -- so the written-field derivations scan `srt` while the *reads* are -# counted across the whole shipped package: a borrowed read answers with the -# startup default wherever it is written, and `benchmark/` ships too. -_READS_SCANNED = _PACKAGE - -_DECLARERS = ("declare_resolution",) - - -def _declared_by_keyword(): - """Fields named as a keyword at a declaration site.""" - written = set() - for path in sorted(_SRT.rglob("*.py")): - source = path.read_text(encoding="utf-8-sig") - if not any(declarer in source for declarer in _DECLARERS): - continue - try: - tree = ast.parse(source) - except SyntaxError: - raise AssertionError(f"unparsable module in the census: {path}") - for node in ast.walk(tree): - if not isinstance(node, ast.Call): - continue - if isinstance(node.func, ast.Attribute): - name = node.func.attr - elif isinstance(node.func, ast.Name): - name = node.func.id - else: - continue - if name in _DECLARERS: - written |= {kw.arg for kw in node.keywords if kw.arg} - return written - - -def _returned_field_names(function): - """Field names a provider/pass writes: the keys of the mapping it returns. - - Only the returned mapping counts -- walking every `ast.Dict` in the body - also collects a dict-valued field's *nested* keys and any unrelated local - mapping, and those stray names would reject valid borrowed reads of fields - resolution never writes. The mapping is traced through four spellings: a - returned literal, assignments (annotated or not) to a returned name, a - literal-key subscript write on it, and `.update(field=...)` on it. A - spelling this cannot see raises instead of skipping. - - ``overrides[name]`` is also accepted when ``name`` comes from - ``for name in ("a", "b", ...)`` -- the keys stay statically enumerable. - """ - names = set() - returned = set() - # for x in ("a", "b"): ... -> {"x": {"a", "b"}} - loop_keys = { - node.target.id: {elt.value for elt in node.iter.elts} - for node in ast.walk(function) - if isinstance(node, ast.For) - and isinstance(node.target, ast.Name) - and isinstance(node.iter, (ast.Tuple, ast.List)) - and node.iter.elts - and all( - isinstance(elt, ast.Constant) and isinstance(elt.value, str) - for elt in node.iter.elts - ) - } - - def top_level_keys(mapping): - for key in mapping.keys: - if not (isinstance(key, ast.Constant) and isinstance(key.value, str)): - raise AssertionError(f"non-literal key in {function.name}") - names.add(key.value) - - def add_subscript_key(key): - if isinstance(key, ast.Constant): - names.add(key.value) - elif isinstance(key, ast.Name) and key.id in loop_keys: - names.update(loop_keys[key.id]) - else: - raise AssertionError(f"non-literal key in {function.name}") - - for node in ast.walk(function): - if isinstance(node, ast.Return) and node.value is not None: - value = node.value - if isinstance(value, ast.Dict): - top_level_keys(value) - elif isinstance(value, ast.Name): - returned.add(value.id) - elif isinstance(value, ast.Constant) and value.value is None: - pass - else: - raise AssertionError( - f"opaque return in {function.name}: {ast.unparse(value)}" - ) - for node in ast.walk(function): - if isinstance(node, (ast.Assign, ast.AnnAssign)): - targets = node.targets if isinstance(node, ast.Assign) else [node.target] - for target in targets: - if ( - isinstance(target, ast.Name) - and target.id in returned - and isinstance(node.value, ast.Dict) - ): - top_level_keys(node.value) - if isinstance(target, ast.Subscript) and ( - isinstance(target.value, ast.Name) and target.value.id in returned - ): - add_subscript_key(target.slice) - if ( - isinstance(node, ast.Call) - and isinstance(node.func, ast.Attribute) - and node.func.attr == "update" - and isinstance(node.func.value, ast.Name) - and node.func.value.id in returned - ): - names |= {kw.arg for kw in node.keywords if kw.arg} - # A positional dict literal has to be read here. The Dict walk above - # only reaches literals that are *returned* or assigned to a returned - # name, so `d.update({"field": value})` was being type-checked and - # then dropped -- silently, under a comment claiming otherwise. - for arg in node.args: - if isinstance(arg, ast.Dict): - top_level_keys(arg) - else: - raise AssertionError(f"opaque update() argument in {function.name}") - if any(kw.arg is None for kw in node.keywords): - raise AssertionError(f"**kwargs update() in {function.name}") - return names - - -def _declared_by_registry_and_passes(): - """Fields the model-override registry and the post-process passes write. - - These field names are *data* -- dict keys, not keywords -- so a keyword - scan misses every one of them. The callables are collected from the live - registries rather than by matching decorator names: 26 of the 27 providers - register through a `_register_for(...)` helper, so a scan for - `@register_model_override*` sees exactly one of them and reports a healthy - census over a channel it cannot see. - """ - import sys - - from sglang.srt.arg_groups import overrides - - # Resolve each callable's body in the file it actually lives in. The - # declarations are spread over `arg_groups/model_overrides/`, one module per - # model family, and a scan hard-coded to `overrides.py` would find none of - # them -- and, worse, would keep reporting a healthy census while doing it. - bodies_by_module = {} - - def _bodies(module_name): - if module_name not in bodies_by_module: - path = getattr(sys.modules[module_name], "__file__", None) - assert path, f"{module_name} has no source file" - module_tree = ast.parse(pathlib.Path(path).read_text(encoding="utf-8-sig")) - bodies_by_module[module_name] = { - node.name: node - for node in ast.walk(module_tree) - if isinstance(node, ast.FunctionDef) - } - return bodies_by_module[module_name] - - callables = {fn for fns in overrides._MODEL_OVERRIDE_FNS.values() for fn in fns} - callables |= { - fn for _predicate, fn in getattr(overrides, "_PREDICATE_OVERRIDE_FNS", ()) - } - callables |= set(overrides.POST_PROCESS_PASSES) - - fields = set() - for fn in callables: - name = getattr(fn, "__name__", "") - body = _bodies(fn.__module__).get(name) - # Loud, not silent: a body this scan cannot find is a field census it - # is not taking, and a narrower census makes every check downstream of - # it quietly vacuous. - assert body is not None, f"{fn.__module__}.{name} has no body to scan" - fields |= _returned_field_names(body) - - # The literal arch -> {field: value} table, which has no callable at all. - # It lives with the rest of the registry, in `model_override_base`. - from sglang.srt.arg_groups import model_override_base - - table_tree = ast.parse( - pathlib.Path(model_override_base.__file__).read_text(encoding="utf-8-sig") - ) - seen_table = False - for node in table_tree.body: - target = None - if isinstance(node, ast.Assign) and isinstance(node.targets[0], ast.Name): - target = node.targets[0].id - elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name): - target = node.target.id - if target != "MODEL_OVERRIDES" or node.value is None: - continue - for inner in ast.walk(node.value): - if not isinstance(inner, ast.Dict): - continue - for key, value in zip(inner.keys, inner.values): - if isinstance(value, ast.Dict): - continue - if not isinstance(key, ast.Constant): - raise AssertionError("non-literal override key") - fields.add(key.value) - seen_table = True - assert seen_table, "MODEL_OVERRIDES is not where this scan looks for it" - return fields - - -def _written_after_publish(): - """Fields the runtime overrides once the bags exist. - - Imported from the supplied-instance ratchet rather than re-derived: it - already enumerates `get_context().override(...)` and its named wrapper, and - a second derivation of the same channel is what drifts narrower. A borrowed - read of one of these answers with the startup value the same way a - resolution-written one does -- the write just lands later. - """ - import importlib.util - - companion = ( - pathlib.Path(__file__).resolve().parent - / "test_supplied_instance_exposure_ratchet.py" - ) - spec = importlib.util.spec_from_file_location("_exposure_for_ratchet", companion) - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - return set(module.TestSuppliedInstanceExposure._override_written_fields()) - - -def _resolution_written(): - """Every field the startup record answers wrong once it stays raw.""" - return ( - _declared_by_keyword() - | _declared_by_registry_and_passes() - | _written_after_publish() - ) - - -def _borrowed_parking_spans(tree): - """[(span, parked attribute names)] for classes that park a borrowed record. - - The supplied-instance census counts `self.server_args.` only where - the class was handed the record as a parameter -- `self.server_args = - server_args` inside a method that takes one. A class that borrows it off - another object instead (`self.server_args = scheduler.server_args`) is - covered by neither census, and a read through that attribute is exactly the - borrowed-record chain read this file is about. The parked name is whatever - the class chose -- `self.args = scheduler.server_args` hides the same read, - so the assignment target is recorded, not assumed. - """ - spans = [] - for cls in (n for n in ast.walk(tree) if isinstance(n, ast.ClassDef)): - parked_names = set() - for node in ast.walk(cls): - # Both assignment spellings -- `self.args = x.server_args` and the - # annotated `self.args: ServerArgs = x.server_args`. - if isinstance(node, ast.Assign): - targets = node.targets - elif isinstance(node, ast.AnnAssign) and node.value is not None: - targets = [node.target] - else: - continue - # A bare name on the right is the parameter the companion census - # follows; an attribute chain ending in `.server_args` is a record - # taken off another object. - if not ( - isinstance(node.value, ast.Attribute) - and node.value.attr == "server_args" - ): - continue - for target in targets: - if ( - isinstance(target, ast.Attribute) - and isinstance(target.value, ast.Name) - and target.value.id == "self" - ): - parked_names.add(target.attr) - if parked_names: - spans.append(((cls.lineno, cls.end_lineno), parked_names)) - return spans - - -def _subtrees_with_their_own_record(): - """Top-level package directories that define a second `ServerArgs`. - - `multimodal_gen` ships one, so `x.server_args.` inside it names a - field of *that* class -- and `model_path` is a field of both. Keying on the - spelling alone once put a resolution call into a diffusion entry point. - Derived from the class definitions rather than named here, so a third - record would be excluded the same way instead of silently counting. - """ - roots = set() - for path in _PACKAGE.rglob("*.py"): - rel = path.relative_to(_PACKAGE).as_posix() - if rel.startswith("srt/") or "/" not in rel: - continue - source = path.read_text(encoding="utf-8-sig") - if "ServerArgs" not in source: - continue - try: - tree = ast.parse(source) - except SyntaxError: - continue - if any( - isinstance(node, ast.ClassDef) and node.name == "ServerArgs" - for node in ast.walk(tree) - ): - roots.add(rel.split("/")[0]) - return roots - - -def _reads_the_startup_record(tree): - """True when this module's `server_args` is the one `srt` resolves. - - Inside a subtree that owns another record, only a module that imports the - startup record is talking about it. - """ - for node in ast.walk(tree): - if isinstance(node, ast.ImportFrom) and node.module: - if node.module.startswith("sglang.srt"): - # The *original* names: `ServerArgs as SrtServerArgs` is still - # the startup record, whatever this module calls it. - names = {alias.name for alias in node.names} - if names & {"ServerArgs", "server_args", "prepare_server_args"}: - return True - return False - - -def _chain_reads(written): - """`.server_args.` where the field is resolution-written.""" - found = [] - other_records = _subtrees_with_their_own_record() - for path in sorted(_READS_SCANNED.rglob("*.py")): - rel = path.relative_to(_READS_SCANNED).as_posix() - in_srt = rel.startswith("srt/") - if in_srt: - under_srt = rel[len("srt/") :] - if path.name in _OWNERS or under_srt.startswith(_OWNERS[-1]): - continue - source = path.read_text(encoding="utf-8-sig") - if "server_args" not in source: - continue - try: - tree = ast.parse(source) - except SyntaxError: - raise AssertionError(f"unparsable module in the census: {rel}") - if ( - not in_srt - and rel.split("/")[0] in other_records - and not _reads_the_startup_record(tree) - ): - continue - parked = _borrowed_parking_spans(tree) - - def parked_alias(lineno, name): - return any( - start <= lineno <= end and name in names - for (start, end), names in parked - ) - - for node in ast.walk(tree): - if not ( - isinstance(node, ast.Attribute) - and isinstance(node.ctx, ast.Load) - and node.attr in written - ): - continue - base = node.value - if not isinstance(base, ast.Attribute): - continue - through_self = isinstance(base.value, ast.Name) and base.value.id == "self" - if base.attr == "server_args": - # `self.server_args.field` is the parked spelling the - # supplied-instance census counts -- but only where the record - # arrived as a parameter. - if through_self and not parked_alias(node.lineno, base.attr): - continue - elif not (through_self and parked_alias(node.lineno, base.attr)): - # Any other attribute counts only as a recorded parked alias - # (`self.args = scheduler.server_args` and later `self.args.x`). - continue - suffix = " (parked borrowed record)" if through_self else "" - found.append(f"{rel}:{node.lineno} {ast.unparse(base)}.{node.attr}{suffix}") - # `getattr(model_runner.server_args, "field", default)` reads the same - # borrowed record through a `Call`, with the same stale default. - for node in ast.walk(tree): - if not ( - isinstance(node, ast.Call) - and isinstance(node.func, ast.Name) - and node.func.id == "getattr" - and len(node.args) >= 2 - and isinstance(node.args[1], ast.Constant) - and node.args[1].value in written - ): - continue - base = node.args[0] - if not isinstance(base, ast.Attribute): - continue - through_self = isinstance(base.value, ast.Name) and base.value.id == "self" - if base.attr == "server_args": - if through_self and not parked_alias(node.lineno, base.attr): - continue - elif not (through_self and parked_alias(node.lineno, base.attr)): - continue - found.append( - f"{rel}:{node.lineno} getattr({ast.unparse(base)}, " - f"{node.args[1].value!r})" - ) - return sorted(found) - - -def _passes_named_at_call_sites() -> set: - """Names passed to ``run_post_process_pass(sa, fn)`` anywhere in the tree. - - A call whose pass is not a bare name is a hard failure, not a skip: this - scan is the ground truth every registry-driven check below is derived from, - so `run_post_process_pass(self, overrides._new_pass)` (an `ast.Attribute`) - or `run_post_process_pass(self, fn=_new_pass)` (a keyword) would otherwise - walk past all of them silently. Keeping the call shape uniform is the - price of the scan being complete. - """ - names = set() - for path in sorted(pathlib.Path(next(iter(sglang.__path__))).rglob("*.py")): - source = path.read_text(encoding="utf-8-sig") - if "run_post_process_pass" not in source: - continue - try: - tree = ast.parse(source) - except SyntaxError: - continue - for node in ast.walk(tree): - if not isinstance(node, ast.Call): - continue - func = node.func - if isinstance(func, ast.Name): - called = func.id - elif isinstance(func, ast.Attribute): - called = func.attr - else: - called = None - if called != "run_post_process_pass": - continue - if ( - len(node.args) != 2 - or node.keywords - or not isinstance(node.args[1], ast.Name) - ): - raise AssertionError( - f"{path}:{node.lineno}: run_post_process_pass takes the pass " - "as a bare name in its second positional argument; " - f"{ast.unparse(node)!r} is invisible to this scan and to " - "every registry-driven check derived from it" - ) - names.add(node.args[1].id) - return names - - -class TestEveryInvokedPassIsRegistered(CustomTestCase): - """The registry is what the scans above enumerate, so a pass missing from it - is a pass nothing checks. - - Being invoked and being registered are two edits, and `_a2a_fusion_adjustments` - shipped with only the first: it ran in production while the registry-driven - scans walked past it. The call sites are the ground truth here -- the registry - is derived from a decorator someone has to remember. - """ - - def test_the_registry_covers_every_call_site(self): - from sglang.srt.arg_groups import overrides - - invoked = _passes_named_at_call_sites() - self.assertGreater( - len(invoked), - 20, - f"only {len(invoked)} call sites found; the scan is broken, not the tree", - ) - registered = {fn.__name__ for fn in overrides.POST_PROCESS_PASSES} - self.assertEqual( - set(), - invoked - registered, - "these passes are invoked but carry no @register_post_process, so " - "every check that walks POST_PROCESS_PASSES skips them", - ) - self.assertEqual( - set(), - registered - invoked, - "these passes carry @register_post_process but no slot invokes " - "them; deleting a call site and leaving the decorator behind " - "leaves a pass that only the scans can see", - ) - - -class TestNoChainReadsOfResolvedConfig(CustomTestCase): - def test_the_census_has_something_to_count(self): - """A written set that collapsed would make the pin vacuous. - - Each mechanism is checked on its own, because they fail - independently: the keyword scan cannot see a field name that is data, - and a scan for `@register_model_override*` sees one provider out of - twenty-seven because the rest register through a helper. A hand-written - expectation of the resulting field names is what hid that -- it stayed - green while a whole channel went unscanned -- so each mechanism is - pinned by a floor derived from the live registry instead. - """ - from sglang.srt.arg_groups import overrides - - by_keyword = _declared_by_keyword() - by_data = _declared_by_registry_and_passes() - - self.assertGreater( - len(by_keyword), - 100, - f"only {len(by_keyword)} fields are declared by keyword; the scan broke", - ) - providers = {fn for fns in overrides._MODEL_OVERRIDE_FNS.values() for fn in fns} - providers |= { - fn for _predicate, fn in getattr(overrides, "_PREDICATE_OVERRIDE_FNS", ()) - } - self.assertGreater( - len(providers) + len(overrides.POST_PROCESS_PASSES), - 50, - "the registry and pass tables collapsed; the data-channel scan is " - "reading an empty registry", - ) - self.assertGreater( - len(by_data), - 25, - f"only {len(by_data)} fields come from the registry and the passes, " - f"across {len(providers)} providers and " - f"{len(overrides.POST_PROCESS_PASSES)} passes; the scan of the " - "dict-key channel broke", - ) - # The data channel is not the keyword scan's subset: if it became one, - # that scan would be doing all the work and a regression here would be - # invisible. - self.assertTrue(by_data - by_keyword, "the data channel adds nothing") - - def test_nothing_reads_a_resolved_field_off_a_borrowed_record(self): - found = _chain_reads(_resolution_written()) - self.assertEqual( - found, - [], - "these reach the startup record through another object for a value " - "resolution decides, so they answer with the CLI default once the " - "record stays raw; read the config bag instead:\n " + "\n ".join(found), - ) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/unit/test_supplied_instance_exposure_ratchet.py b/test/registered/unit/test_supplied_instance_exposure_ratchet.py deleted file mode 100644 index 5ad79c344..000000000 --- a/test/registered/unit/test_supplied_instance_exposure_ratchet.py +++ /dev/null @@ -1,874 +0,0 @@ -"""The supplied-instance surface is measured on two axes, and may only shrink. - -A callee that takes ``server_args`` keeps the supplied-instance contract: the -caller chose the object, so no global-read ratchet counts it. Step 12 changes -what that object *carries* — the instance stays at the user's raw input — so a -callee reading a field **resolution fills in** would start seeing the CLI default -instead of the effective value. - -This pins that intersection. Each entry is one (file, field) pair where a -parameter named ``server_args`` is read for a field resolution writes; the plan -doc carries the proposed disposition per field -(``global_context/12-raw-input-config.md``, "the supplied-instance conversion -list"). New pairs fail: a new one is new step-12 work, and the moment to decide -where the value should come from is when the read is written, not during the -flip. Pairs that disappear also fail, with the entry to delete — the list is the -measurement, not a memory of one. - -The written-field set is derived here rather than hardcoded: the -representative configs in ``_MATRIX`` (one per resolution family it exercises) -are resolved and compared against the dataclass defaults, the same matrix the -context repo's audit tool uses. Ambient environment is normalized per entry -- -resolution branches on CI detection and leaves sticky process state, so each -entry resolves from the pristine snapshot, and the CI shape is an explicit -entry rather than an accident of the runner. The read scan mirrors that -tool's three shapes — a parameter attribute, ``getattr(server_args, "literal")``, -and the parameter parked on ``self`` — because two implementations of one census -that disagree are worse than either alone. - -The second axis is **already wrong today**, not after a flip. Some config is -decided *after* publish and recorded with ``get_context().override(...)`` — -elastic-EP resizing `ep_size`, a weight update rewriting `model_path` / -`load_format`, HiCache attach naming a storage backend, adaptive speculative -decoding moving `speculative_num_steps`. That write reaches the bags and never -the record, so a supplied-instance read of one of those fields answers with the -startup value from the moment the override lands. Whether that is a defect -depends on ordering — a value copied at construction, before any override, is -fine — so this axis is pinned as a measurement with the same growth guard rather -than as a list of bugs. One of them *was* a defect and is fixed at the base of -this stack: the linear-attn dispatch table rebuilt itself from the record after -the SM100 GDN prefill decision had been recorded in the bag, so a second runner's -rebuild dropped it. That choice is a per-runner stamp now and is not recorded -process-wide at all, so neither the read nor the field is on this axis. -""" - -import ast -import json -import os -import shutil -import tempfile -import unittest -from pathlib import Path - -import msgspec -import msgspec.structs - -import sglang -from sglang.srt.server_args import ServerArgs -from sglang.test.ci.ci_register import register_cpu_ci, register_cuda_ci -from sglang.test.test_utils import CustomTestCase - -register_cpu_ci(est_time=23, suite="base-a-test-cpu") -# Also on a CUDA runner: the written set is derived by resolving on the running -# host, and `is_cuda()` / capability gates only open on real hardware. The pin -# is split by host so both registrations stay exact: `_EXPOSED` is asserted -# everywhere, and a pair whose write only happens on CUDA belongs in -# `_EXPOSED_CUDA_ONLY` -- pinned on the CUDA runner, invisible to the CPU -# assertion. Without the split, one shared exact list could not hold such a -# pair at all: pinning it fails the CPU run as "gone", omitting it fails the -# CUDA run as "new". (No AMD registration: an `is_hip()`-gated write would -# shift the exact sets in ways none of the pinning hosts can verify; the ROCm -# resolution surface is covered by `test_resolution_is_reproducible.py` -# instead, whose assertion is device-agnostic.) -register_cuda_ci(est_time=16, stage="base-b", runner_config="1-gpu-small") - -_PACKAGE_ROOT = Path(next(iter(sglang.__path__))) / "srt" - -# The config the resolution pipeline owns; reading the in-flight record is their -# job, not a supplied-instance read. -_OWNERS = ("server_args.py", "runtime_context.py", "arg_groups/") - -_MINI_CONFIG = { - "architectures": ["LlamaForCausalLM"], - "model_type": "llama", - "hidden_size": 16, - "intermediate_size": 32, - "num_attention_heads": 2, - "num_key_value_heads": 2, - "num_hidden_layers": 2, - "vocab_size": 128, - "max_position_embeddings": 2048, -} - -# One config resolves only its own decisions, so the written set is a union. -_MATRIX = ( - {}, - { - "speculative_algorithm": "EAGLE", - "speculative_num_steps": 3, - "speculative_eagle_topk": 1, - "speculative_num_draft_tokens": 4, - }, - {"dp_size": 2, "tp_size": 2, "enable_dp_attention": True}, - # DWDP resolves dp_size and enable_dp_attention *itself* -- the plain DP - # entry above passes them in, and passed-in fields are excluded from the - # written set, so without this entry the dp_size readers would never pin. - {"tp_size": 2, "dwdp_size": 2}, - {"enable_hierarchical_cache": True, "hicache_ratio": 2.0}, - {"disaggregation_mode": "prefill"}, - {"tp_size": 2, "attn_cp_size": 2}, - {"enable_lora": True, "max_lora_rank": 16}, - {"kv_cache_dtype": "fp8_e4m3", "page_size": 64}, - # MIS resolves disable_radix_cache (and friends) itself; the backend is - # passed in because the handler asserts flashinfer rather than switching. - {"enable_mis": True, "attention_backend": "flashinfer"}, -) - -# `declare_resolution` call sites whose keyword expansion is built -# dynamically; the written fields are spelled out here and drift-guarded. -_LATE_RESOLUTION_DYNAMIC_SITES = { - "parser/template_detection.py": frozenset({"reasoning_parser", "tool_call_parser"}), -} - -# `get_context().override(...)` declares through the same seam, but the fields -# are the caller's -- a test names them one call at a time. There is no static -# set to collect, and nothing resolution decides: whatever a caller overrides -# there is exposure only through that caller's own reads. -_CALLER_SUPPLIED_LATE_SITES = frozenset({"runtime_context.py"}) - -# Resolution also branches on ambient environment; those shapes are explicit -# entries so the written set is the same on every host. `SGLANG_IS_IN_CI` -# makes resolution fill `soft_watchdog_timeout`. -_ENV_MATRIX = (({}, {"SGLANG_IS_IN_CI": "true"}),) - -# Only true constructor inputs: `tokenizer_path` / `served_model_name` are -# resolution-written (filled from `model_path` when unset), so their readers -# are step-12 exposure like any other pair. -_PASSED = frozenset({"model_path", "device", "random_seed"}) - -# Empty. A pair belongs here when a reader has no bag to read -- it runs before -# its process publishes -- and cannot use `resolving_view` either. The launcher's -# pre-publish reads (`_set_envs_and_config`, the auto-parser gate) and the -# late-resolution detection it calls all read the declarations now, so nothing -# qualifies. A new entry needs that kind of reason next to it. -_EXPOSED: frozenset = frozenset() - -# Pairs whose resolution write only happens on a CUDA host (capability or -# `is_cuda()` gated): asserted on the CUDA registration, invisible to the CPU -# one. Empty today -- the current written sets coincide across the two hosts -- -# but this is where a GPU-only write's readers get pinned without breaking the -# CPU-exact assertion. -_EXPOSED_CUDA_ONLY: frozenset = frozenset() - - -# Axis two: (file, field) pairs where a supplied-instance read names a field that -# some code overrides post-publish. Each needs an ordering judgment, not a blanket -# conversion; the list exists so a new one is a decision made when it is written. -_OVERRIDDEN_AND_READ: frozenset = frozenset() - - -def _expanded_override_keys(rel, tree, call, kw) -> set: - """The statically visible keys behind an ``override(..., **expr)``. - - Handles a dict literal, a conditional between dict literals, and a name - bound to a dict literal in the enclosing function (plus constant-subscript - stores onto it -- the HiCache attach shape). One expansion is unresolvable - by design and exempted by name: ``update_server_args`` forwards - operator-chosen fields, so its key set is the API's, not this file's. - Anything else unresolvable fails -- a silently skipped expansion would - shrink the written set. - """ - for a in call.args: - if isinstance(a, ast.Constant) and a.value == "update_server_args": - return set() - for k in call.keywords: - if ( - k.arg == "source" - and isinstance(k.value, ast.Constant) - and k.value.value == "update_server_args" - ): - return set() - - def loop_variable_values(name: str) -> set: - """The values a `for name, ... in ()` loop binds. - - A handler that records one field per loop iteration spells the field - names in the loop's own literal, so they are still static. - """ - values = set() - for node in ast.walk(tree): - if not isinstance(node, ast.For): - continue - target = node.target - names = ( - [target] - if isinstance(target, ast.Name) - else list(getattr(target, "elts", [])) - ) - if not names or not isinstance(names[0], ast.Name) or names[0].id != name: - continue - if not (node.lineno <= call.lineno <= (node.end_lineno or node.lineno)): - continue - for item in getattr(node.iter, "elts", []): - first = ( - item.elts[0] if isinstance(item, ast.Tuple) and item.elts else item - ) - if isinstance(first, ast.Constant) and isinstance(first.value, str): - values.add(first.value) - return values - - def dict_keys(node) -> set: - assert isinstance(node, ast.Dict), ( - f"non-literal dict in override expansion at {rel}:{call.lineno}" - ) - keys = set() - for key in node.keys: - if isinstance(key, ast.Constant): - keys.add(key.value) - continue - assert isinstance(key, ast.Name), ( - f"non-literal dict key in override expansion at {rel}:{call.lineno}" - ) - bound = loop_variable_values(key.id) - assert bound, ( - f"dict key {key.id!r} at {rel}:{call.lineno} is not bound by a " - "literal loop; extend the resolver" - ) - keys |= bound - return keys - - if isinstance(kw.value, ast.Dict): - return dict_keys(kw.value) - if isinstance(kw.value, ast.IfExp): - keys = set() - for branch in (kw.value.body, kw.value.orelse): - if isinstance(branch, ast.Dict) and branch.keys: - keys |= dict_keys(branch) - elif isinstance(branch, ast.Dict): - pass - else: - raise AssertionError( - f"unresolvable override expansion at {rel}:{call.lineno}" - ) - return keys - assert isinstance(kw.value, ast.Name), ( - f"unresolvable override expansion at {rel}:{call.lineno}" - ) - name = kw.value.id - enclosing = None - for fn in ast.walk(tree): - if isinstance(fn, (ast.FunctionDef, ast.AsyncFunctionDef)): - if ( - fn.lineno - <= call.lineno - <= max(getattr(fn, "end_lineno", fn.lineno), fn.lineno) - ): - if enclosing is None or fn.lineno > enclosing.lineno: - enclosing = fn - assert enclosing is not None, ( - f"override expansion outside any function at {rel}:{call.lineno}" - ) - keys = set() - found = False - for node in ast.walk(enclosing): - if ( - isinstance(node, ast.Assign) - and len(node.targets) == 1 - and isinstance(node.targets[0], ast.Name) - and node.targets[0].id == name - and isinstance(node.value, ast.Dict) - ): - found = True - keys |= dict_keys(node.value) - elif ( - isinstance(node, ast.Assign) - and len(node.targets) == 1 - and isinstance(node.targets[0], ast.Subscript) - and isinstance(node.targets[0].value, ast.Name) - and node.targets[0].value.id == name - and isinstance(node.targets[0].slice, ast.Constant) - ): - keys.add(node.targets[0].slice.value) - assert found, ( - f"override expansion '{name}' at {rel}:{call.lineno} has no " - "dict-literal assignment in its function; extend the resolver" - ) - return keys - - -class TestSuppliedInstanceExposure(CustomTestCase): - def _callTestMethod(self, method): - # No CI retry: a failed first attempt has already resolved the matrix - # and mutated process state; a retry against that contamination could - # pass on a drifted written set or mask a real drift. - return unittest.TestCase._callTestMethod(self, method) - - def setUp(self): - # Resolving the matrix writes process state on the way through (the - # multimodal transport handler sets SGLANG_USE_CUDA_IPC_TRANSPORT, and - # `EnvField.set()` flips a descriptor flag `os.environ` does not carry). - # Leaking it makes *later* files in the same worker fail, which is how - # this was found -- so the case restores what it touched. - super().setUp() - state = (dict(os.environ), self._env_field_flags()) - self.addCleanup(self._restore_process_state, state) - - @staticmethod - def _env_field_flags() -> dict: - from sglang.srt.environ import EnvField, envs - - flags = {} - for klass in reversed(type(envs).__mro__): - for name, field in vars(klass).items(): - if isinstance(field, EnvField): - flags[name] = field._set_to_none - return flags - - @staticmethod - def _restore_process_state(state) -> None: - from sglang.srt.environ import envs - - saved_environ, saved_flags = state - os.environ.clear() - os.environ.update(saved_environ) - for name, was_none in saved_flags.items(): - getattr(type(envs), name)._set_to_none = was_none - - def _config_dir(self) -> str: - config_dir = tempfile.mkdtemp(prefix="supplied_instance_") - self.addCleanup(shutil.rmtree, config_dir, ignore_errors=True) - with open(os.path.join(config_dir, "config.json"), "w") as handle: - json.dump(_MINI_CONFIG, handle) - return config_dir - - def _resolution_written_fields(self) -> set: - """The union of what resolution fills in across the matrix. - - Every entry must resolve. A silently skipped one would shrink this set, - which makes pinned pairs look like they disappeared -- the list would - then drift by environment rather than by code, and the failure would - point at the wrong thing. Each entry resolves from the pristine - process snapshot (resolution writes env and EnvField flags on the way - through, and DWDP flips `SGLANG_SCHEDULER_SKIP_ALL_GATHER`), so the - union does not depend on matrix order; and the ambient CI marker is - cleared, so a runner's identity cannot leak into the measurement -- - the CI-conditioned writes come from `_ENV_MATRIX`'s explicit entry. - Declarers outside `arg_groups/` count too: the parser auto-detection - runs at launcher stage and the NPU helper is called by the pipeline, so - their target fields are collected statically from the call sites -- - resolution writes by definition, just not reached by the matrix. - """ - pristine = (dict(os.environ), self._env_field_flags()) - written = set() - - def resolve_one(extra, env): - self._restore_process_state(pristine) - os.environ.pop("SGLANG_IS_IN_CI", None) - os.environ.update(env) - model_path = self._config_dir() - try: - resolved = ServerArgs( - model_path=model_path, device="cuda", random_seed=42, **extra - ) - resolved.resolve_once() - except Exception as exc: - self.fail( - f"the matrix entry {extra} (env={env}) did not resolve in " - f"this environment ({type(exc).__name__}: {exc}); the " - "written-field union would be short and the pinned list " - "would drift" - ) - defaults = {} - for field in msgspec.structs.fields(resolved): - if field.default is not msgspec.NODEFAULT: - defaults[field.name] = field.default - elif field.default_factory is not msgspec.NODEFAULT: - defaults[field.name] = field.default_factory() - for field_name, default in defaults.items(): - if field_name in _PASSED or field_name in extra: - continue - if getattr(resolved, field_name) != default: - written.add(field_name) - - for extra in _MATRIX: - resolve_one(extra, {}) - for extra, env in _ENV_MATRIX: - resolve_one(extra, env) - self._restore_process_state(pristine) - written |= self._declared_outside_the_pipeline() - written |= self._hook_assignment_targets() - written |= self._record_method_assignment_targets() - written |= self._declarative_override_fields() - return written - - def _hook_assignment_targets(self) -> set: - """Fields any resolution hook can write, collected statically. - - The matrix can only enumerate families someone thought to add -- the - DFLASH hole (its hook is the sole writer of - `speculative_draft_attention_backend`, and no entry ran it) showed - that a family nobody listed leaves its readers unpinned. The hook - modules under `arg_groups/` are the resolution pipeline's extension - points -- along with the NPU default helper, which the pipeline calls - the same way -- and their write surface is the may-write set, - family-blind by - construction. A hook writes two ways: `server_args.field = ...`, and - `declare_resolution(server_args, source, field=...)`, which records - the write in the declaration stash on its way to the field. Counting - only the assignment would read a hook's conversion to a declaration as - the field having stopped being written. Collected like the - late-resolution keywords: statically, failing loudly on an - unparsable module. Underscore-prefixed targets are pipeline - bookkeeping, not config leaves. - """ - targets = set() - modules = sorted((_PACKAGE_ROOT / "arg_groups").glob("*.py")) - modules.append(_PACKAGE_ROOT / "hardware_backend/npu/utils.py") - for path in modules: - try: - tree = ast.parse(path.read_text(encoding="utf-8-sig")) - except SyntaxError: - self.fail(f"unparsable hook module in the census: {path.name}") - for node in ast.walk(tree): - if isinstance(node, ast.Assign): - tgts = node.targets - elif isinstance(node, (ast.AnnAssign, ast.AugAssign)): - tgts = [node.target] - elif ( - isinstance(node, ast.Call) - and isinstance(node.func, ast.Name) - and node.func.id == "declare_resolution" - ): - targets |= { - kw.arg - for kw in node.keywords - if kw.arg and not kw.arg.startswith("_") - } - continue - else: - continue - for tgt in tgts: - if ( - isinstance(tgt, ast.Attribute) - and isinstance(tgt.value, ast.Name) - and tgt.value.id == "server_args" - and not tgt.attr.startswith("_") - ): - targets.add(tgt.attr) - return targets - - def _record_method_assignment_targets(self) -> set: - """Fields ``ServerArgs``'s own methods can write, collected statically. - - The record's handlers are as family-conditional as the hooks -- the - mooncake/layer_first layout rewrite, the deepseek-EP mode defaults, - the seed fill that only runs when the caller did *not* supply one (so - construct-and-diff can never see it: measuring requires supplying). - A write site that can never fire is a dead branch to delete upstream, - not a census exemption. Only names that are declared dataclass fields - count; underscore bookkeeping does not. - - Two spellings write: an assignment, and ``self._declare(source, - field=value)``, which records the write in the declaration stash on - its way to the field. Counting only assignments would read a handler's - conversion to a declaration as the field having stopped being written, - which would quietly retire every pinned pair that reads it. - """ - tree = ast.parse( - (_PACKAGE_ROOT / "server_args.py").read_text(encoding="utf-8-sig") - ) - sa_class = next( - node - for node in tree.body - if isinstance(node, ast.ClassDef) and node.name == "ServerArgs" - ) - declared = { - node.target.id - for node in sa_class.body - if isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name) - } - targets = set() - for node in ast.walk(sa_class): - if isinstance(node, ast.Assign): - tgts = node.targets - elif isinstance(node, (ast.AnnAssign, ast.AugAssign)): - tgts = [node.target] - elif ( - isinstance(node, ast.Call) - and isinstance(node.func, ast.Name) - and node.func.id == "declare_resolution" - ): - targets |= { - kw.arg - for kw in node.keywords - if kw.arg in declared and not kw.arg.startswith("_") - } - continue - else: - continue - for tgt in tgts: - if ( - isinstance(tgt, ast.Attribute) - and isinstance(tgt.value, ast.Name) - and tgt.value.id == "self" - and tgt.attr in declared - and not tgt.attr.startswith("_") - ): - targets.add(tgt.attr) - # The deprecated-alias normalization declares through `**renamed`, so - # the keyword scan sees no names; its field set is pinned here. - alias_fields = { - "attention_backend", - "decode_attention_backend", - "prefill_attention_backend", - "speculative_draft_attention_backend", - } - - # The handler lives in `arg_groups/serving_hook.py`, reached either as a - # record method or as a bare-name call, so look the loop up by both. - def _deprecated_alias_handler(): - for node in ast.walk(sa_class): - if ( - isinstance(node, ast.FunctionDef) - and node.name == "_handle_deprecated_args" - and any(isinstance(n, ast.For) for n in ast.walk(node)) - ): - return node - for path in sorted((_PACKAGE_ROOT / "arg_groups").glob("*.py")): - tree = ast.parse(path.read_text(encoding="utf-8-sig")) - for node in tree.body: - if ( - isinstance(node, ast.FunctionDef) - and node.name == "handle_deprecated_args" - ): - return node - raise AssertionError("the deprecated-alias handler was not found") - - deprecated = _deprecated_alias_handler() - found_tuples = [ - {elt.value for elt in node.iter.elts if isinstance(elt, ast.Constant)} - for node in ast.walk(deprecated) - if isinstance(node, ast.For) and isinstance(node.iter, ast.Tuple) - ] - self.assertIn( - alias_fields, - found_tuples, - "the deprecated-alias normalization loop moved or changed its " - "field tuple; update alias_fields to match", - ) - return targets | alias_fields - - def _declarative_override_fields(self) -> set: - """Fields the declarative override registry can write. - - ``MODEL_OVERRIDES`` maps arch -> {field: value}, and the - ``@register_model_override``(-``_predicate``) providers return (or - build by subscript) {field: value} dicts, which go straight into the - declaration stash, so no assignment scan sees these writes and a - llama-only matrix never triggers them. Keys must be - string literals; anything else fails loudly. - """ - tree = ast.parse( - (_PACKAGE_ROOT / "arg_groups" / "overrides.py").read_text( - encoding="utf-8-sig" - ) - ) - fields = set() - for node in tree.body: - target = None - if isinstance(node, ast.Assign) and isinstance(node.targets[0], ast.Name): - target = node.targets[0].id - elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name): - target = node.target.id - if target != "MODEL_OVERRIDES" or node.value is None: - continue - for inner in ast.walk(node.value): - if not isinstance(inner, ast.Dict): - continue - for key, value in zip(inner.keys, inner.values): - if isinstance(value, ast.Dict): - continue # arch -> {…} outer layer - self.assertIsInstance(key, ast.Constant, "non-literal override key") - fields.add(key.value) - for node in ast.walk(tree): - if not isinstance(node, ast.FunctionDef): - continue - if not any( - isinstance(dec, ast.Call) - and isinstance(dec.func, ast.Name) - and dec.func.id.startswith("register_model_override") - for dec in node.decorator_list - ): - continue - for inner in ast.walk(node): - if isinstance(inner, ast.Assign) and isinstance( - inner.targets[0], ast.Subscript - ): - key = inner.targets[0].slice - self.assertIsInstance( - key, ast.Constant, f"non-literal override key in {node.name}" - ) - fields.add(key.value) - if isinstance(inner, ast.Dict): - for key in inner.keys: - self.assertIsInstance( - key, - ast.Constant, - f"non-literal override key in {node.name}", - ) - fields.add(key.value) - return fields - - def _declared_outside_the_pipeline(self) -> set: - """Fields declared by a `declare_resolution` caller outside - `arg_groups/`, collected statically. - - Resolution's launcher-stage writes live here -- the auto-detected - parsers need a tokenizer or chat-template load, so they cannot run in - `__post_init__` -- alongside the NPU default helper and the expert-pack - loader, which the pipeline calls the same way. The construct-and-diff - pass above never sees any of them. - - `arg_groups/` is deliberately excluded: `_hook_assignment_targets` - covers it exactly, and it resolves the pipeline's own computed - expansions (`record_foreign_defaults` declares a `**` dict this - collector's resolver cannot read). The keywords at the call sites are - the written fields; an expansion this cannot resolve fails loudly like - the override collector's, except the named dynamic sites below, whose - field sets are spelled out and drift-guarded (each name must still - appear as a constant in the file).""" - written = set() - root = _PACKAGE_ROOT - for path in sorted(root.rglob("*.py")): - rel = path.relative_to(root).as_posix() - if rel.startswith("arg_groups/"): - continue - source = path.read_text(encoding="utf-8-sig") - if "declare_resolution" not in source: - continue - try: - tree = ast.parse(source) - except SyntaxError: - self.fail(f"unparsable module in the census: {rel}") - for node in ast.walk(tree): - if not ( - isinstance(node, ast.Call) - and ( - ( - isinstance(node.func, ast.Name) - and node.func.id == "declare_resolution" - ) - or ( - isinstance(node.func, ast.Attribute) - and node.func.attr - in ("declare_resolution", "_declare_resolution") - ) - ) - ): - continue - if all(kw.arg is None for kw in node.keywords) and any( - isinstance(kw.value, ast.Name) and kw.value.id == "fields" - for kw in node.keywords - ): - # The forwarding shim (`ServerArgs._late_resolution` / - # the helper's own body) re-expands its caller's kwargs; - # the write sites are the callers. - continue - for kw in node.keywords: - if kw.arg and kw.arg != "source": - written.add(kw.arg) - elif kw.arg is None: - if rel in _CALLER_SUPPLIED_LATE_SITES: - continue - dynamic = _LATE_RESOLUTION_DYNAMIC_SITES.get(rel) - if dynamic is not None: - constants = { - c.value - for c in ast.walk(tree) - if isinstance(c, ast.Constant) - } - missing = dynamic - constants - self.assertFalse( - missing, - f"{rel}: the declared dynamic field set drifted " - f"from the file ({sorted(missing)} not found)", - ) - written |= dynamic - else: - written |= _expanded_override_keys(rel, tree, node, kw) - return written - - _READS_CACHE = None - - def _supplied_instance_reads(self) -> set: - """Three spellings of the same read: ``server_args.field`` off the - parameter, ``getattr(server_args, "field", default)`` with a literal - name, and the *parked* form -- ``self.x = server_args`` in a method - that takes the parameter, read as ``self.x.field`` anywhere in the - class. Parking under a different object, a container, or a computed - name stays invisible, like in every census of this family. The - loudest boundary *was* the *chain* spelling, - ``model_runner.server_args.field`` off some other object, which this - census still does not count -- but those reads are gone for every - resolution-written field and ``test_chain_read_ratchet.py`` holds them - at zero, so the gap is no longer where the risk is.""" - if TestSuppliedInstanceExposure._READS_CACHE is not None: - return TestSuppliedInstanceExposure._READS_CACHE - pairs = set() - for path in sorted(_PACKAGE_ROOT.rglob("*.py")): - rel = path.relative_to(_PACKAGE_ROOT).as_posix() - if rel.startswith(_OWNERS): - continue - source = path.read_text(encoding="utf-8-sig") - if "server_args" not in source: - continue - try: - tree = ast.parse(source) - except SyntaxError: - # A silently dropped module shrinks `found` and reads as - # intentional surface shrinkage under the bidirectional pin. - self.fail(f"unparsable module in the census: {rel}") - for fn in ast.walk(tree): - if not isinstance(fn, (ast.FunctionDef, ast.AsyncFunctionDef)): - continue - params = {a.arg for a in list(fn.args.args) + list(fn.args.kwonlyargs)} - if "server_args" not in params: - continue - for node in ast.walk(fn): - if ( - isinstance(node, ast.Attribute) - and isinstance(node.value, ast.Name) - and node.value.id == "server_args" - and isinstance(node.ctx, ast.Load) - ): - pairs.add((rel, node.attr)) - elif ( - isinstance(node, ast.Call) - and isinstance(node.func, ast.Name) - and node.func.id == "getattr" - and len(node.args) >= 2 - and isinstance(node.args[0], ast.Name) - and node.args[0].id == "server_args" - and isinstance(node.args[1], ast.Constant) - and isinstance(node.args[1].value, str) - ): - # The same read in optional clothing. Only a literal - # name is censusable; a computed one is not. - pairs.add((rel, node.args[1].value)) - for cls in ast.walk(tree): - if not isinstance(cls, ast.ClassDef): - continue - # Parked: `self.x = server_args` in a method that takes the - # parameter, read as `self.x.field` anywhere in the class. - parked = set() - for fn in cls.body: - if not isinstance(fn, (ast.FunctionDef, ast.AsyncFunctionDef)): - continue - if "server_args" not in { - a.arg for a in list(fn.args.args) + list(fn.args.kwonlyargs) - }: - continue - for node in ast.walk(fn): - if ( - isinstance(node, ast.Assign) - and len(node.targets) == 1 - and isinstance(node.targets[0], ast.Attribute) - and isinstance(node.targets[0].value, ast.Name) - and node.targets[0].value.id == "self" - and isinstance(node.value, ast.Name) - and node.value.id == "server_args" - ): - parked.add(node.targets[0].attr) - for node in ast.walk(cls): - if ( - isinstance(node, ast.Attribute) - and isinstance(node.ctx, ast.Load) - and isinstance(node.value, ast.Attribute) - and node.value.attr in parked - and isinstance(node.value.value, ast.Name) - and node.value.value.id == "self" - ): - pairs.add((rel, node.attr)) - TestSuppliedInstanceExposure._READS_CACHE = pairs - return pairs - - @staticmethod - def _override_written_fields() -> set: - """Fields written post-publish through ``get_context().override(...)``.""" - written = set() - for path in sorted(_PACKAGE_ROOT.rglob("*.py")): - rel = path.relative_to(_PACKAGE_ROOT).as_posix() - source = path.read_text(encoding="utf-8-sig") - if "record_config_updates" not in source and not ( - "get_context" in source and "override" in source - ): - continue - try: - tree = ast.parse(source) - except SyntaxError: - raise AssertionError(f"unparsable module in the census: {rel}") - for node in ast.walk(tree): - if not ( - isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) - ): - continue - base = node.func.value - is_override = node.func.attr == "override" and ( - isinstance(base, ast.Call) - and isinstance(base.func, ast.Name) - and base.func.id == "get_context" - ) - # `record_config_updates` is a named wrapper over override, so - # its call sites are override sites. Its body forwards **kwargs - # and names no field, so skip the forwarding call itself. - is_wrapper = node.func.attr == "record_config_updates" - if not (is_override or is_wrapper): - continue - inside_wrapper = any( - isinstance(fn, (ast.FunctionDef, ast.AsyncFunctionDef)) - and fn.name == "record_config_updates" - and fn.lineno <= node.lineno <= (fn.end_lineno or fn.lineno) - for fn in ast.walk(tree) - ) - if inside_wrapper: - continue - for kw in node.keywords: - if kw.arg == "source": - # Override metadata, not a config field. - continue - if kw.arg: - written.add(kw.arg) - else: - written |= _expanded_override_keys(rel, tree, node, kw) - return written - - def test_the_post_publish_override_surface_matches_the_pinned_list(self): - written = self._override_written_fields() - self.assertGreater( - len(written), 5, "found almost no override targets; the scan broke" - ) - found = {pair for pair in self._supplied_instance_reads() if pair[1] in written} - new = sorted(found - _OVERRIDDEN_AND_READ) - gone = sorted(_OVERRIDDEN_AND_READ - found) - self.assertEqual( - ([], []), - (new, gone), - "the post-publish override surface drifted. A read here answers with " - "the startup value once the override lands, so a new pair needs an " - "ordering judgment: copied before any override (fine), or read after " - "one (then it must come from the bags).\n" - f" new: {new}\n" - f" gone (delete from _OVERRIDDEN_AND_READ): {gone}", - ) - - def test_the_exposed_set_matches_the_pinned_list(self): - import torch - - written = self._resolution_written_fields() - found = {pair for pair in self._supplied_instance_reads() if pair[1] in written} - expected = set(_EXPOSED) - if torch.cuda.is_available(): - expected |= _EXPOSED_CUDA_ONLY - new = sorted(found - expected) - gone = sorted(expected - found) - self.assertEqual( - ([], []), - (new, gone), - "the supplied-instance step-12 surface drifted.\n" - f" new (decide where the resolved value comes from): {new}\n" - f" gone (delete from _EXPOSED / _EXPOSED_CUDA_ONLY): {gone}", - ) - - -if __name__ == "__main__": - unittest.main()