[SRT] Clean up no-op compiler pass, dead helpers, and migration tests (#39295)
This commit is contained in:
Xiaoyu Zhang
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
@@ -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))
|
||||
@@ -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)
|
||||
|
||||
@@ -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]})
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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],
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user