diff --git a/proto/sglang/runtime/v1/sglang.proto b/proto/sglang/runtime/v1/sglang.proto index fb8e4f7a1750..15251ef8010b 100644 --- a/proto/sglang/runtime/v1/sglang.proto +++ b/proto/sglang/runtime/v1/sglang.proto @@ -11,6 +11,7 @@ service SglangService { rpc Tokenize(TokenizeRequest) returns (TokenizeResponse); rpc Detokenize(DetokenizeRequest) returns (DetokenizeResponse); rpc HealthCheck(HealthCheckRequest) returns (HealthCheckResponse); + rpc WatchEngineState(WatchEngineStateRequest) returns (stream EngineStateSnapshot); rpc GetModelInfo(GetModelInfoRequest) returns (GetModelInfoResponse); rpc GetServerInfo(GetServerInfoRequest) returns (GetServerInfoResponse); rpc ListModels(ListModelsRequest) returns (ListModelsResponse); @@ -174,6 +175,23 @@ message HealthCheckResponse { bool healthy = 1; } +// ---- Engine state ---- + +message WatchEngineStateRequest {} + +// A complete discovery and lifecycle snapshot. The JSON discovery payloads +// intentionally retain the same wire shape as GetModelInfo/GetServerInfo. +message EngineStateSnapshot { + // Unix time in nanoseconds, captured once for this engine process. + uint64 instance_id = 1; + // Starts at one and increases for each snapshot from this instance. + uint64 revision = 2; + bool healthy = 3; + bool is_pause = 4; + GetModelInfoResponse model_info = 5; + GetServerInfoResponse server_info = 6; +} + // ---- Model info ---- message GetModelInfoRequest {} diff --git a/python/pyproject.toml b/python/pyproject.toml index 171994664b49..739bf3907ffe 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -97,8 +97,8 @@ dependencies = [ "uvloop", "watchfiles", "xgrammar==0.2.1", - "xxhash", - "zstandard", + "xxhash", # /pull_weights delta checksum + "zstandard", # /pull_weights delta codec ] [[tool.uv.index]] @@ -141,7 +141,7 @@ diffusion-qvg = [ ] ray = [ - "ray[default]>=2.55.1", + "ray[default]>=2.56.0", ] tracing = [ diff --git a/python/sglang/kernels/ops/moe/fused_moe_lora_kernel.py b/python/sglang/kernels/ops/moe/fused_moe_lora_kernel.py index b0c85481e9c8..073b1927d2ac 100644 --- a/python/sglang/kernels/ops/moe/fused_moe_lora_kernel.py +++ b/python/sglang/kernels/ops/moe/fused_moe_lora_kernel.py @@ -403,8 +403,11 @@ def _fused_moe_lora_expand( IS_PRIMARY=False, **expand_config, ) + # the base output may be padded (Marlin), so its slices sit at output width / num_slices + slice_stride = output.shape[-1] // num_slices for i in range(num_slices): - output[:, :, i * N + offset : (i + 1) * N + offset] += b_intermediate_cache1[i] + start = i * slice_stride + offset + output[:, :, start : start + N] += b_intermediate_cache1[i] @torch.inference_mode() diff --git a/python/sglang/kernels/ops/moe/virtual_experts.py b/python/sglang/kernels/ops/moe/virtual_experts.py index e8b35e4465d5..74e22ef02549 100644 --- a/python/sglang/kernels/ops/moe/virtual_experts.py +++ b/python/sglang/kernels/ops/moe/virtual_experts.py @@ -55,7 +55,9 @@ def _fused_virtual_topk_ids_kernel( # a real virtual-expert slot belonging to another adapter and trigger OOB # loads in downstream LoRA kernels. shifted = base + safe_lora * num_experts_for_weight - result = tl.where(base < 0, base, shifted) + # Reject an out-of-range base ID before it can alias the next adapter. + result = tl.where(base >= num_experts_for_weight, -1, shifted) + result = tl.where(base < 0, base, result) tl.store(virtual_topk_ids_ptr + offs, result, mask=valid) # Write mask once per row (at first k position) @@ -517,7 +519,7 @@ def _merged_experts_fused_moe_lora_add_impl( output: torch.Tensor, hidden_states: torch.Tensor, lora_a: torch.Tensor, - lora_b: torch.Tensor, + lora_b: torch.Tensor | list[torch.Tensor] | tuple[torch.Tensor, ...], topk_ids: torch.Tensor, topk_weights: torch.Tensor, token_lora_mapping: torch.Tensor, @@ -526,13 +528,30 @@ def _merged_experts_fused_moe_lora_add_impl( experts_shared_outer_loras_b: bool, routing_cache: dict | None = None, ) -> None: + """Fused virtual-experts LoRA delta add. + + ``lora_b`` accepts either a single tensor or a sequence of tensors stacked + along the output dim. Length-2 is the gate_up case where A has rank ``2*r`` + (gate's A and up's A concatenated along rank) and each B has rank ``r``. + The shrink runs once over the full ``2*r`` rank; the expand runs once per + B, each reading its half of the intermediate and writing to its slice of + ``output``. """ - 1. Prepare virtual expert routing metadata from topk_ids + token_lora_mapping * num_experts. - 2. Flatten LoRA weights from [max_loras, num_experts, ...] to [max_loras * num_experts, ...]. - 3. Run regular SGLang fused-MoE kernels for LoRA A and LoRA B. - 4. Mask out tokens with token_lora_mapping == -1 on the add path. - """ + lora_b_list: list[torch.Tensor] = ( + list(lora_b) if isinstance(lora_b, (list, tuple)) else [lora_b] + ) + n_b = len(lora_b_list) + assert n_b in (1, 2), f"lora_b must be length 1 or 2, got {n_b}" + b_rank = lora_b_list[0].shape[3] + for b in lora_b_list[1:]: + assert b.shape == lora_b_list[0].shape, ( + f"all lora_b tensors must share shape; got {[tuple(t.shape) for t in lora_b_list]}" + ) + max_loras, _, max_lora_rank, _ = lora_a.shape + assert max_lora_rank == n_b * b_rank, ( + f"lora_a rank {max_lora_rank} != n_b ({n_b}) * lora_b rank {b_rank}" + ) input_top_k = 1 if hidden_states.shape[0] == topk_ids.numel() else topk_ids.shape[1] def _merge_lora_expert_weight(t: torch.Tensor) -> torch.Tensor: @@ -616,16 +635,18 @@ def _get_routing( block_size=block_size, num_experts=virtual_num_experts, ) - # _align_block_size uses a worst-case padded allocation. Trim the routing buffers - # to a tighter upper bound so we keep the real routed work but drop unused padding - num_tokens = topk_ids.numel() - max_nonempty = min(num_tokens, virtual_num_experts) - tight_padded = ( - triton.cdiv(num_tokens + max_nonempty * (block_size - 1), block_size) - * block_size - ) - sorted_token_ids = sorted_token_ids[:tight_padded] - expert_ids = expert_ids[: tight_padded // block_size] + # NOTE: do NOT trim sorted_token_ids / expert_ids to a tighter upper bound here. + # The downstream kernels (_moe_lora_shrink_splitk_kernel, fused_moe_kernel) read + # sorted_token_ids[pid_m*BLOCK : +BLOCK] and expert_ids[pid_m] WITHOUT a bounds mask + # for every block up to num_tokens_post_padded (a GPU-side count loaded at run time). + # num_tokens_post_padded comes from _align_block_size with `virtual_num_experts` buckets + # and can exceed a tighter `numel + min(numel,virtual_num_experts)*(block-1)` bound + # (most so for shared-outer, where virtual_num_experts = max_loras is small), so trimming + # made those unmasked reads land PAST the view. In eager mode the slack still lives inside + # the same _align_block_size allocation (garbage, masked out downstream) so it worked; under + # CUDA-graph capture/replay the graph mempool packs tensors tightly and that slack may belong + # to another pooled tensor / lie past a page -> cudaErrorIllegalInstruction during capture. + # Keep the full worst-case-allocated buffers so every unmasked read stays in-allocation. expert_ids = fused_sanitize_expert_ids(expert_ids, virtual_num_experts) result = ( sorted_token_ids, @@ -644,13 +665,24 @@ def _get_routing( ) lora_a_virtual = _merge_lora_expert_weight(lora_a) - lora_b_virtual = _merge_lora_expert_weight(lora_b) + lora_b_virtuals = [_merge_lora_expert_weight(b) for b in lora_b_list] num_experts_a = lora_a.shape[1] - num_experts_b = lora_b.shape[1] - + num_experts_b = lora_b_list[0].shape[1] + half_out = lora_b_list[0].shape[2] + + # The kernels index token_lora_mapping / intermediate by token ids up to + # topk_ids.shape[0] (the DP-gathered token count under --enable-dp-attention). An + # under-sized mapping means unmasked OOB reads/writes that surface as a sticky, + # hard-to-attribute CUDA IMA — fail loudly on the host instead. + assert token_lora_mapping.shape[0] >= topk_ids.shape[0], ( + f"token_lora_mapping covers {token_lora_mapping.shape[0]} tokens but the MoE runs on " + f"{topk_ids.shape[0]} (DP-gathered?) tokens; mapping was sized before the dp gather " + f"length was known (see get_gathered_moe_num_tokens)" + ) + # Split-K atomic additions must accumulate before rounding to the activation dtype. intermediate = torch.zeros( - [token_lora_mapping.shape[0], topk_ids.shape[1], max_lora_rank], - dtype=hidden_states.dtype, + [topk_ids.shape[0], topk_ids.shape[1], max_lora_rank], + dtype=torch.float32, device=hidden_states.device, ) @@ -680,7 +712,8 @@ def _get_routing( a_stage_config, ) - b_stage_config = _get_stage_config(lora_b_virtual, 1) + intermediate = intermediate.to(hidden_states.dtype) + b_stage_config = _get_stage_config(lora_b_virtuals[0], 1) ( sorted_token_ids, expert_ids, @@ -694,33 +727,53 @@ def _get_routing( b_stage_config["BLOCK_SIZE_M"], ) - invoke_fused_moe_kernel( - intermediate.view(-1, max_lora_rank), - lora_b_virtual, - None, - output, - None, - None, - None, - topk_weights, - topk_ids, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - mul_routed_weight, - 1, - b_stage_config, - tl.bfloat16 if hidden_states.dtype == torch.bfloat16 else tl.float16, - False, - False, - False, - False, - False, - None, - fuse_add_to_output=True, - add_output_mask=token_lora_mask, - router_topk=topk_ids.shape[1], - ) + # n_b expands. For len 1: K=b_rank covers full intermediate, write full output. + # For len 2 (gate_up): split intermediate along rank into [gate, up] halves + # (each contiguous, K=b_rank=r) and output along last dim into [gate, up] + # halves (each of width half_out). Each B in lora_b_virtuals is its own + # half's weight tensor, naturally K=b_rank. + for b_idx, b_virtual in enumerate(lora_b_virtuals): + if n_b == 1: + inter_arg = intermediate.view(-1, b_rank) + out_arg = output + else: + inter_arg = ( + intermediate[..., b_idx * b_rank : (b_idx + 1) * b_rank] + .contiguous() + .view(-1, b_rank) + ) + out_arg = output[ + ..., b_idx * half_out : (b_idx + 1) * half_out + ].contiguous() + invoke_fused_moe_kernel( + inter_arg, + b_virtual, + None, + out_arg, + None, + None, + None, + topk_weights, + topk_ids, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + mul_routed_weight, + 1, + b_stage_config, + tl.bfloat16 if hidden_states.dtype == torch.bfloat16 else tl.float16, + False, + False, + False, + False, + False, + None, + fuse_add_to_output=True, + add_output_mask=token_lora_mask, + router_topk=topk_ids.shape[1], + ) + if n_b != 1: + output[..., b_idx * half_out : (b_idx + 1) * half_out].copy_(out_arg) def _merged_experts_fused_moe_lora_add_op( @@ -763,7 +816,7 @@ def merged_experts_fused_moe_lora_add( output: torch.Tensor, hidden_states: torch.Tensor, lora_a: torch.Tensor, - lora_b: torch.Tensor, + lora_b: torch.Tensor | list[torch.Tensor] | tuple[torch.Tensor, ...], topk_ids: torch.Tensor, topk_weights: torch.Tensor, token_lora_mapping: torch.Tensor, @@ -772,7 +825,12 @@ def merged_experts_fused_moe_lora_add( experts_shared_outer_loras_b: bool, routing_cache: dict | None = None, ) -> None: - """Public API: wraps the registered op with routing_cache support.""" + """Public API: wraps the registered op with routing_cache support. + + ``lora_b`` accepts a sequence of length 2 for the gate_up case (each B + holds one half of the stacked output, rank ``r``, with A's rank ``2*r``); + a single tensor is used for the down case. + """ _merged_experts_fused_moe_lora_add_impl( output, hidden_states, diff --git a/python/sglang/multimodal_gen/configs/sample/minimax_h3.py b/python/sglang/multimodal_gen/configs/sample/minimax_h3.py index c78216d5b11b..d5d79aba56a0 100644 --- a/python/sglang/multimodal_gen/configs/sample/minimax_h3.py +++ b/python/sglang/multimodal_gen/configs/sample/minimax_h3.py @@ -241,14 +241,15 @@ def _validate(self) -> None: "video/audio denoise loop has no lossless TeaCache contract" ) if self.rollout: - raise ValueError( - "MiniMax H3 does not support rollout: its coupled video/audio " - "scheduler has no SchedulerRLMixin contract" - ) + task = str(self.task or "t2va").lower() + if task not in ("t2va",): + raise ValueError( + f"MiniMax H3 rollout currently supports task=t2va only, got {task!r}" + ) if self.return_trajectory_latents or self.return_trajectory_decoded: raise ValueError( - "MiniMax H3 does not support trajectory output for its coupled " - "video/audio denoise state" + "MiniMax H3 does not support return_trajectory_latents/decoded; " + "use rollout=True with rollout_return_dit_trajectory instead" ) seeds = self.seed if isinstance(self.seed, list) else [self.seed] for seed in seeds: diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/io_struct.py b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/io_struct.py index b558be3e4a21..de01bc49283b 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/io_struct.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/io_struct.py @@ -79,6 +79,8 @@ class RolloutRequest(BaseModel): fps: Optional[int] = None rollout: bool = True + # "uint8": quantise the video engine-side; None: ship unchanged + rollout_video_dtype: Optional[str] = None rollout_sde_type: str = "sde" rollout_noise_level: float = 0.7 rollout_log_prob_no_const: bool = False diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/request_timing.py b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/request_timing.py new file mode 100644 index 000000000000..6404a1c75076 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/request_timing.py @@ -0,0 +1,83 @@ +"""Wall-clock marks for the rollout HTTP path. + +The client reconstructs one waterfall per request across three processes, so it +needs absolute marks, not durations. They ride a response header: every mark is +taken before the response headers are sent, and a header leaves the msgpack +body contract untouched for clients that ignore it. +""" + +from __future__ import annotations + +import json +import time +from contextlib import contextmanager + +# Compact JSON rides under each header, e.g. +# x-sgld-timing: {"srv_recv":1787472609.104512,...,"msgpack_end":1787472691.63172,"request_id":"abc123"} +# mark name -> epoch seconds (6 dp), plus this request's id +# x-sgld-stages: {"TextEncodingStage":91.2,"DenoisingStage":448512.301,"DecodingStage":8123.457} +# pipeline stage class -> milliseconds (3 dp) +TIMING_HEADER = "x-sgld-timing" +STAGES_HEADER = "x-sgld-stages" + +MARKS = ( + "srv_recv", + "forward_start", + "forward_end", + "build_start", + "build_end", + "dump_end", + "msgpack_end", +) + + +class RequestStamps: + """Absolute wall-clock marks for one rollout request.""" + + __slots__ = ("request_id", "_marks") + + def __init__(self, request_id: str = "") -> None: + self.request_id = request_id + self._marks: dict[str, float] = {} + + def mark(self, name: str) -> None: + assert name in MARKS, f"unknown timing mark {name!r}" + self._marks[name] = time.time() + + @contextmanager + def span(self, name: str): + """Mark ``{name}_start`` on entry and ``{name}_end`` on exit. + + dump/msgpack have no start mark: each begins where the previous span + ended, so only the end boundary is recorded. + """ + start = f"{name}_start" + if start in MARKS: + self.mark(start) + try: + yield + finally: + self.mark(f"{name}_end") + + def to_header(self) -> str: + payload: dict[str, object] = { + name: round(t, 6) for name, t in self._marks.items() + } + if self.request_id: + payload["request_id"] = self.request_id + return json.dumps(payload, separators=(",", ":")) + + +def stages_header(metrics) -> str: + """The per-stage milliseconds the engine already recorded, verbatim. + + The client sees the whole forward as one number, so without these it cannot + tell a slow denoise from a slow VAE decode. Only stage totals travel, never + the per-step list, which would grow the header with the step count. + """ + if metrics is None or not metrics.stages: + return "" + return json.dumps( + {name: round(ms, 3) for name, ms in metrics.stages.items()}, + separators=(",", ":"), + ) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/rollout_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/rollout_api.py index 71492523dcd7..9c3cba5c0d29 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/rollout_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/rollout_api.py @@ -2,11 +2,12 @@ from __future__ import annotations +import asyncio from typing import Any -import msgspec import torch from fastapi import APIRouter, HTTPException, Response +from fastapi.responses import StreamingResponse from sglang.multimodal_gen.configs.sample.sampling_params import generate_request_id from sglang.multimodal_gen.runtime.entrypoints.openai.utils import build_sampling_params @@ -14,8 +15,16 @@ RolloutRequest, RolloutResponse, ) +from sglang.multimodal_gen.runtime.entrypoints.post_training.request_timing import ( + STAGES_HEADER, + TIMING_HEADER, + RequestStamps, + stages_header, +) from sglang.multimodal_gen.runtime.entrypoints.post_training.utils import ( _maybe_serialize, + _quantize_video_uint8, + msgpack_encode_spliced, ) from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch @@ -131,6 +140,7 @@ def _slice_rollout_trajectory_for_sample( latents=_extract_single_sample_tensor(dit.latents, sample_idx, batch_size), timesteps=dit.timesteps, sigmas=dit.sigmas, + latent_step_indices=dit.latent_step_indices, ) return RolloutTrajectoryData( rollout_log_probs=log_probs, @@ -185,6 +195,11 @@ def _serialize_rollout_trajectory( ), "timesteps": serialized_dit_timesteps, "sigmas": serialized_dit_sigmas, + "latent_step_indices": ( + _maybe_serialize(dit.latent_step_indices) + if dit.latent_step_indices is not None + else None + ), } return ( serialized_log_probs, @@ -195,7 +210,12 @@ def _serialize_rollout_trajectory( def _build_response( - request_id: str, prompt: str, seed: int, rollout: bool, result: OutputBatch + request_id: str, + prompt: str, + seed: int, + rollout: bool, + result: OutputBatch, + video_dtype: str | None = None, ) -> list[RolloutResponse]: """ rollout: bool - set to False when evaluating the model @@ -227,6 +247,8 @@ def _build_response( for sample_idx in range(batch_size): out_i = result.output[sample_idx] if isinstance(out_i, torch.Tensor): + if video_dtype == "uint8": + out_i = _quantize_video_uint8(out_i) out_i = out_i.contiguous() serialized_generated_output = _maybe_serialize(out_i) if not rollout: @@ -317,7 +339,10 @@ def _build_sampling_kwargs(request: RolloutRequest) -> dict: }, ) async def rollout_generate(request: RolloutRequest): + stamps = RequestStamps() + stamps.mark("srv_recv") request_id = generate_request_id() + stamps.request_id = request_id server_args = get_global_server_args() sampling_kwargs = _build_sampling_kwargs(request) try: @@ -330,9 +355,10 @@ async def rollout_generate(request: RolloutRequest): server_args=server_args, sampling_params=sampling_params ) try: - output_batch: OutputBatch = await async_scheduler_client.forward( - pipeline_request - ) + with stamps.span("forward"): + output_batch: OutputBatch = await async_scheduler_client.forward( + pipeline_request + ) except Exception as exc: logger.error("Rollout generation failed: %s", exc, exc_info=True) raise HTTPException( @@ -340,11 +366,36 @@ async def rollout_generate(request: RolloutRequest): ) from exc if output_batch.error: raise HTTPException(status_code=500, detail=output_batch.error) - rollout_responses = _build_response( - request_id, request.prompt, request.seed, request.rollout, output_batch - ) - payload = [r.model_dump() for r in rollout_responses] - return Response( - content=msgspec.msgpack.encode(payload), + + def _serialize_response() -> list[bytes]: + with stamps.span("build"): + rollout_responses = _build_response( + request_id, + request.prompt, + request.seed, + request.rollout, + output_batch, + video_dtype=request.rollout_video_dtype, + ) + with stamps.span("dump"): + payload = [r.model_dump() for r in rollout_responses] + with stamps.span("msgpack"): + return msgpack_encode_spliced(payload) + + # Off the event loop: building the response copies hundreds of MB + parts = await asyncio.to_thread(_serialize_response) + + # Timing rides a header: the last marks postdate the encoded body. The + # explicit content-length keeps uvicorn on identity framing, not chunked. + headers = { + TIMING_HEADER: stamps.to_header(), + "content-length": str(sum(len(p) for p in parts)), + } + stages = stages_header(output_batch.metrics) + if stages: + headers[STAGES_HEADER] = stages + return StreamingResponse( + iter(parts), media_type="application/msgpack", + headers=headers, ) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/utils.py index 495fadf88d0b..8fa5c5f3e6c4 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/utils.py @@ -4,10 +4,13 @@ from typing import Any +import msgspec import numpy as np import torch from safetensors.torch import load, save +_SPLICE_MIN_BYTES = 1 << 20 + def tensor_to_bytes(t: torch.Tensor) -> bytes: return save({"t": t.detach().contiguous().cpu()}) @@ -34,6 +37,48 @@ def _maybe_serialize(obj: Any) -> Any: return obj +def _container_header(tag_fix: int, tag16: bytes, tag32: bytes, n: int) -> bytes: + if n <= 15: + return bytes([tag_fix | n]) + if n <= 0xFFFF: + return tag16 + n.to_bytes(2, "big") + return tag32 + n.to_bytes(4, "big") + + +def msgpack_encode_spliced(obj: Any, threshold: int = _SPLICE_MIN_BYTES) -> list[bytes]: + """Encode to msgpack as a list of buffers, splicing large ``bytes`` values. + + Container and bin32 headers are hand-written so a large ``bytes`` payload + lands in the output by reference, never copied through the encoder; the + concatenated parts are byte-identical to ``msgspec.msgpack.encode(obj)``. + """ + parts: list[bytes] = [] + small = bytearray() + + def _emit(value: Any) -> None: + if isinstance(value, bytes) and len(value) >= threshold: + small.extend(b"\xc6" + len(value).to_bytes(4, "big")) + parts.append(bytes(small)) + small.clear() + parts.append(value) + elif isinstance(value, dict): + small.extend(_container_header(0x80, b"\xde", b"\xdf", len(value))) + for key, item in value.items(): + small.extend(msgspec.msgpack.encode(key)) + _emit(item) + elif isinstance(value, (list, tuple)): + small.extend(_container_header(0x90, b"\xdc", b"\xdd", len(value))) + for item in value: + _emit(item) + else: + small.extend(msgspec.msgpack.encode(value)) + + _emit(obj) + if small: + parts.append(bytes(small)) + return parts + + def _maybe_deserialize(obj: Any) -> Any: if isinstance(obj, dict): if obj.get("__tensor__"): @@ -42,3 +87,11 @@ def _maybe_deserialize(obj: Any) -> Any: if isinstance(obj, (list, tuple)): return [_maybe_deserialize(v) for v in obj] return obj + + +def _quantize_video_uint8(video: torch.Tensor) -> torch.Tensor: + """Map the decoded [0,1] float video to 0..255 uint8; consumers divide by 255.""" + out = video.float() + if float(out.max()) <= 1.0 + 1e-3: + out = out * 255.0 + return out.clamp(0, 255).to(torch.uint8) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py index db1d29a2d5d2..31b3f73bc40a 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py @@ -59,6 +59,7 @@ ) from sglang.multimodal_gen.runtime.layers.linear import ( ColumnParallelLinear, + LinearBase, MergedColumnParallelLinear, RowParallelLinear, ) @@ -370,6 +371,8 @@ def _rotate_half(x: torch.Tensor) -> torch.Tensor: def _accepts_mxfp8_input(linear: nn.Module) -> bool: + if not isinstance(linear, LinearBase): + return False return linear.quant_method is not None and linear.quant_method.accepts_mxfp8_input( linear ) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py index cea45be71c0c..9d7357dd017a 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py @@ -919,6 +919,7 @@ def _get_added_qkv_projections( if ( self._unquantized_added_qkv_is_packed and not qwen_image_added_qkv_active(self) + and isinstance(self.to_added_qkv, MergedColumnParallelLinear) ): return _split_unquantized_merged_linear( self.to_added_qkv, encoder_hidden_states diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/lora/pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines_core/lora/pipeline.py index 5f6620b8b3f0..f009e10d1ba0 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/lora/pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/lora/pipeline.py @@ -53,12 +53,10 @@ logger = init_logger(__name__) -def _swap_peft_swiglu_fc1_lora_b( +def swap_peft_swiglu_fc1_lora_b( source_name: str, target_name: str, weight: torch.Tensor ) -> torch.Tensor: - # Only the PEFT -> native H3 FFN rewrite: ff.net.0.proj [value; gate] - # onto mlp.fc1 [gate; value]. Native fused mlp.fc1 and other models' - # ff.net.0.proj (e.g. Flux) must not match. + """Rewrite PEFT H3 FFN lora_B from [value; gate] to native [gate; value].""" if ( weight.dim() != 2 or ".ff.net.0.proj.lora_B" not in source_name @@ -934,7 +932,7 @@ def load_lora_adapter( else: continue - weight = _swap_peft_swiglu_fc1_lora_b(name, target_name, weight) + weight = swap_peft_swiglu_fc1_lora_b(name, target_name, weight) if target_name in self.lora_adapters[lora_nickname]: raise ValueError( f"Dit target weight name {target_name} already exists in lora_adapters[{lora_nickname}]" diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/denoise_loop.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/denoise_loop.py index 77f389b08ffe..62599da199bc 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/denoise_loop.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/denoise_loop.py @@ -462,6 +462,7 @@ def minimax_h3_denoise_loop( attn_metadata: AttentionMetadata | None = None, on_step: Callable[[int, torch.Tensor, torch.Tensor], None] | None = None, step_profiler: Callable[[int], AbstractContextManager] | None = None, + rollout_ctx=None, ) -> tuple[torch.Tensor, torch.Tensor]: """Run the full denoise loop; returns final (video_rows, audio_rows). @@ -557,11 +558,22 @@ def minimax_h3_denoise_loop( audio_one_minus_sigma_ratios = 1.0 - audio_sigma_ratios video_denoised_scratch = torch.empty_like(video_rows[video_target_slice]) audio_denoised_scratch = torch.empty_like(audio_rows[audio_target_slice]) + if rollout_ctx is not None: + from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.minimax_h3_rollout import ( + minimax_h3_rollout_update_video_target, + ) + + rollout_ctx.batch._h3_rollout_sigma_max = float(max(sigmas_video)) + video_target = video_rows[video_target_slice] + rollout_ctx.collector.record_initial(video_target.detach().clone()) for step in range(num_steps): step_cm = step_profiler(step) if step_profiler is not None else nullcontext() with step_cm: s_v = sigmas_video[step] s_a = sigmas_audio[step] + s_n_v = sigmas_video[step + 1] + if rollout_ctx is not None: + rollout_ctx.batch._rollout_loop_step_index = step fk = positive.forward_kwargs( video_rows=video_rows, @@ -592,15 +604,38 @@ def minimax_h3_denoise_loop( mv_audio_t = v_audio[audio_target_slice].float() video_target = video_rows[video_target_slice] - _minimax_h3_update_target_rows_( - video_target, - mv_video_t, - sigma_t=video_sigma_t[step], - sigma_curr=s_v, - sigma_ratio=video_sigma_ratios[step], - one_minus_sigma_ratio=video_one_minus_sigma_ratios[step], - denoised_scratch=video_denoised_scratch, - ) + if rollout_ctx is not None: + ( + updated, + log_sum, + log_count, + rollout_ctx.noise_buffer, + ) = minimax_h3_rollout_update_video_target( + video_target, + mv_video_t, + sigma_curr=s_v, + sigma_next=s_n_v, + batch=rollout_ctx.batch, + generator=rollout_ctx.generator, + loop_step_index=step, + noise_buffer=rollout_ctx.noise_buffer, + ) + video_target.copy_(updated) + rollout_ctx.collector.record_step( + video_target.detach().clone(), + log_sum, + log_count, + ) + else: + _minimax_h3_update_target_rows_( + video_target, + mv_video_t, + sigma_t=video_sigma_t[step], + sigma_curr=s_v, + sigma_ratio=video_sigma_ratios[step], + one_minus_sigma_ratio=video_one_minus_sigma_ratios[step], + denoised_scratch=video_denoised_scratch, + ) audio_target = audio_rows[audio_target_slice] _minimax_h3_update_target_rows_( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/minimax_h3_rollout.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/minimax_h3_rollout.py new file mode 100644 index 000000000000..2bccab0e0ca4 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/minimax_h3_rollout.py @@ -0,0 +1,242 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Rollout (RL) helpers for MiniMax H3 video-only denoising.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, field +from typing import Any + +import torch + +from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import ( + RolloutDenoisingEnv, + RolloutDitTrajectory, + RolloutTrajectoryData, +) + +_LOG_SQRT_2PI = math.log(math.sqrt(2 * math.pi)) + + +def _effective_sde_type(batch, loop_step_index: int) -> str: + sde_type = str(getattr(batch, "rollout_sde_type", "sde")) + if sde_type == "ode": + return "ode" + sde_step_indices = getattr(batch, "rollout_sde_step_indices", None) + if sde_step_indices is not None and loop_step_index not in sde_step_indices: + return "ode" + return sde_type + + +def minimax_h3_rollout_update_video_target( + video_target: torch.Tensor, + velocity: torch.Tensor, + *, + sigma_curr: float, + sigma_next: float, + batch, + generator: torch.Generator, + loop_step_index: int, + noise_buffer: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Stochastic or deterministic update for video target rows during rollout. + + H3 velocity ``v`` relates to flow-matching model output as ``model_output = -v``. + Returns ``(updated_target, log_prob_sum[B=1], log_prob_count[B=1], noise_buffer)``. + """ + sde_type = _effective_sde_type(batch, loop_step_index) + noise_level = float(getattr(batch, "rollout_noise_level", 0.0)) + log_prob_no_const = bool(getattr(batch, "rollout_log_prob_no_const", False)) + + if sde_type == "ode" or noise_level == 0.0: + from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.denoise_loop import ( + _minimax_h3_update_target_rows_, + ) + + out = video_target.clone() + sigma_t = out.new_tensor(sigma_curr) + ratio = out.new_tensor(sigma_next / sigma_curr if sigma_curr != 0.0 else 1.0) + one_minus_ratio = 1.0 - ratio + scratch = torch.empty_like(out) + _minimax_h3_update_target_rows_( + out, + velocity.float(), + sigma_t=sigma_t, + sigma_curr=sigma_curr, + sigma_ratio=ratio, + one_minus_sigma_ratio=one_minus_ratio, + denoised_scratch=scratch, + ) + elem_count = float(out.numel()) + log_prob_sum = torch.zeros(1, device=out.device, dtype=torch.float32) + log_prob_count = out.new_tensor([elem_count]) + return out, log_prob_sum, log_prob_count, noise_buffer + + model_output = (-velocity).float() + sample = video_target.float() + current_sigma = sample.new_tensor(float(sigma_curr)) + next_sigma = sample.new_tensor(float(sigma_next)) + + if ( + noise_buffer is None + or noise_buffer.shape != sample.shape + or noise_buffer.device != sample.device + ): + noise_buffer = torch.empty_like(sample) + variance_noise = torch.randn( + sample.shape, + generator=generator, + device=sample.device, + dtype=torch.float32, + ) + noise_buffer.copy_(variance_noise) + + if sde_type == "cps": + std_dev_t = next_sigma * math.sin(noise_level * math.pi / 2) + noise_std_dev = std_dev_t + pred_original = sample - current_sigma * model_output + noise_estimate = sample + model_output * (1.0 - current_sigma) + prev_mean = pred_original * (1.0 - next_sigma) + noise_estimate * torch.sqrt( + torch.clamp(next_sigma**2 - std_dev_t**2, min=1e-12) + ) + prev_sample = prev_mean + variance_noise * noise_std_dev + log_prob_no_const_val = -((variance_noise * noise_std_dev) ** 2) + elif sde_type == "sde": + dt = next_sigma - current_sigma + sigma_max = float(getattr(batch, "_h3_rollout_sigma_max", 1.0)) + std_dev_t = ( + torch.sqrt( + current_sigma + / ( + 1.0 + - torch.where( + torch.isclose(current_sigma, current_sigma.new_tensor(1.0)), + current_sigma.new_tensor(sigma_max), + current_sigma, + ) + ) + ) + * noise_level + ) + noise_std_dev = std_dev_t * torch.sqrt(-1.0 * dt) + prev_mean = ( + sample * (1.0 + std_dev_t**2 / (2.0 * current_sigma) * dt) + + model_output + * (1.0 + std_dev_t**2 * (1.0 - current_sigma) / (2.0 * current_sigma)) + * dt + ) + prev_sample = prev_mean + variance_noise * noise_std_dev + log_prob_no_const_val = -((variance_noise * noise_std_dev) ** 2) + else: + raise ValueError(f"Unsupported rollout_sde_type for H3: {sde_type!r}") + + if log_prob_no_const or sde_type == "ode": + log_prob_sum = log_prob_no_const_val.sum().unsqueeze(0) + else: + log_prob_sum = ( + ( + log_prob_no_const_val / (2.0 * (noise_std_dev**2)) + - torch.log(noise_std_dev) + - _LOG_SQRT_2PI + ) + .sum() + .unsqueeze(0) + ) + log_prob_count = sample.new_tensor([float(sample.numel())]) + return ( + prev_sample.to(dtype=video_target.dtype), + log_prob_sum, + log_prob_count, + noise_buffer, + ) + + +@dataclass +class MiniMaxH3RolloutCtx: + """Per-denoise-loop rollout state for H3.""" + + batch: Any + generator: torch.Generator + sigmas_video: list[float] + collector: MiniMaxH3RolloutCollector + noise_buffer: torch.Tensor | None = None + denoising_env_kwargs: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class MiniMaxH3RolloutCollector: + """Accumulates video-target trajectory and per-step log probs.""" + + sigmas_video: list[float] + return_step_indices: set[int] | None = None + latent_steps: list[torch.Tensor] = field(default_factory=list) + latent_step_indices: list[int] = field(default_factory=list) + log_prob_sums: list[torch.Tensor] = field(default_factory=list) + log_prob_counts: list[torch.Tensor] = field(default_factory=list) + pos_cond_kwargs: dict[str, Any] = field(default_factory=dict) + _next_latent_index: int = 0 + + def _record_latent(self, video_target: torch.Tensor) -> None: + index = self._next_latent_index + self._next_latent_index += 1 + if ( + self.return_step_indices is not None + and index not in self.return_step_indices + ): + return + self.latent_steps.append(video_target.detach().cpu().clone()) + self.latent_step_indices.append(index) + + def record_initial(self, video_target: torch.Tensor) -> None: + self._record_latent(video_target) + + def record_step( + self, + video_target: torch.Tensor, + log_prob_sum: torch.Tensor, + log_prob_count: torch.Tensor, + ) -> None: + self._record_latent(video_target) + self.log_prob_sums.append(log_prob_sum.detach().cpu()) + self.log_prob_counts.append(log_prob_count.detach().cpu()) + + def build_trajectory_data(self) -> RolloutTrajectoryData: + # latents: [B=1, T+1, num_video_target_rows, width] + stacked = torch.stack(self.latent_steps, dim=0).unsqueeze(0) + divisor = 1000.0 + step_sigmas = torch.tensor( + [float(s) for s in self.sigmas_video], + dtype=torch.float32, + ) + timesteps = step_sigmas * divisor + sigmas = torch.tensor( + [float(s) for s in self.sigmas_video], + dtype=torch.float32, + ) + log_probs = None + if self.log_prob_sums: + sums = torch.stack(self.log_prob_sums, dim=0) + counts = torch.stack(self.log_prob_counts, dim=0) + per_step = sums.squeeze(-1) / counts.squeeze(-1).clamp(min=1.0) + log_probs = per_step.unsqueeze(0) + return RolloutTrajectoryData( + rollout_log_probs=log_probs, + denoising_env=RolloutDenoisingEnv( + pos_cond_kwargs=self.pos_cond_kwargs, + ), + dit_trajectory=RolloutDitTrajectory( + latents=stacked, + timesteps=timesteps, + sigmas=sigmas, + latent_step_indices=torch.tensor( + self.latent_step_indices, dtype=torch.long + ), + ), + ) + + +__all__ = [ + "MiniMaxH3RolloutCollector", + "MiniMaxH3RolloutCtx", + "minimax_h3_rollout_update_video_target", +] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py index ab65f65c54fa..71998ec3b0c7 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py @@ -780,6 +780,54 @@ def _run_full_loop(self, batch: Req, server_args: ServerArgs) -> None: device=device, ) initial_video, initial_audio = _expand_initial_rows(ctx, positive) + rollout_ctx = None + if getattr(batch, "rollout", False): + from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.minimax_h3_rollout import ( + MiniMaxH3RolloutCollector, + MiniMaxH3RolloutCtx, + ) + from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import ( + RolloutTrajectoryData, + ) + + task = str(getattr(batch.sampling_params, "task", "") or "t2va").lower() + if task not in ("t2va",): + raise ValueError( + f"MiniMax H3 rollout currently supports task=t2va only, got {task!r}" + ) + generator = torch.Generator(device=device) + seed = getattr(batch.sampling_params, "seed", 0) + if isinstance(seed, list): + seed = seed[0] + generator.manual_seed(int(seed)) + return_step_indices = getattr( + batch, "rollout_return_step_indices", None + ) + collector = MiniMaxH3RolloutCollector( + sigmas_video=sigmas_video, + return_step_indices=( + set(return_step_indices) + if return_step_indices is not None + else None + ), + ) + packed_cpu = { + k: (v.detach().cpu() if isinstance(v, torch.Tensor) else v) + for k, v in packed.items() + } + collector.pos_cond_kwargs = { + "encoder_hidden_states": emb["hidden_states"].detach().cpu(), + "h3_packed_layout": packed_cpu, + "h3_token_tags": tags.detach().cpu(), + "h3_video_target_start": positive.video_target_start, + } + rollout_ctx = MiniMaxH3RolloutCtx( + batch=batch, + generator=generator, + sigmas_video=sigmas_video, + collector=collector, + ) + batch.rollout_trajectory_data = RolloutTrajectoryData() with ( maybe_nvtx_range("denoising_loop", self.current_use_nvtx), self.progress_bar( @@ -818,7 +866,12 @@ def on_step(_step, _video_rows, _audio_rows): self._profile_denoising_step, batch=batch, ), + rollout_ctx=rollout_ctx, ) + if rollout_ctx is not None: + batch.rollout_trajectory_data = ( + rollout_ctx.collector.build_trajectory_data() + ) finally: self._finish_active_component_use() _publish_full_loop_outputs( diff --git a/python/sglang/multimodal_gen/runtime/post_training/rl_dataclasses.py b/python/sglang/multimodal_gen/runtime/post_training/rl_dataclasses.py index 1882776d573d..c48a4fcb3bdd 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/rl_dataclasses.py +++ b/python/sglang/multimodal_gen/runtime/post_training/rl_dataclasses.py @@ -51,9 +51,12 @@ class RolloutDitTrajectory: # [B, T+1, ...]: per-step noisy latents x_{t_0..t_{T-1}} followed by the # final denoised latent x_{t_T} (last scheduler.step output). latents: torch.Tensor | None = None - timesteps: torch.Tensor | None = None # [T] + timesteps: torch.Tensor | None = None # [T+1], includes the terminal timestep # [T+1] scheduler.sigmas snapshot (post-shift, includes terminal 0). sigmas: torch.Tensor | None = None + # [K] original step index of each kept latent (0..T); None means the full + # 0..T trajectory. Set whenever rollout_return_step_indices filtered it. + latent_step_indices: torch.Tensor | None = None @dataclass diff --git a/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py b/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py index 674fb1a87933..fde2d199a5e8 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py +++ b/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py @@ -138,6 +138,7 @@ def _maybe_init_denoising_env_collection( batch._rollout_denoising_env_state = { "env": env, "step_latents": [], + "step_latent_indices": [], "step_timesteps": [], "pos_cond_kwargs_src": pos_src, "neg_cond_kwargs_src": neg_src, @@ -156,12 +157,13 @@ def _maybe_append_dit_trajectory_step( if state is None: return + state["step_timesteps"].append(timestep_value.detach().cpu()) return_step_indices = getattr(batch, "rollout_return_step_indices", None) if return_step_indices is not None and step_index not in return_step_indices: return state["step_latents"].append(latents.detach()) - state["step_timesteps"].append(timestep_value.detach().cpu()) + state["step_latent_indices"].append(step_index) def _maybe_finalize_denoising_env_collection(self, batch, pipeline_config) -> None: state = getattr(batch, "_rollout_denoising_env_state", None) @@ -186,6 +188,9 @@ def _maybe_finalize_denoising_env_collection(self, batch, pipeline_config) -> No latents=step_latents_tensor.cpu(), timesteps=torch.stack(step_timesteps, dim=0).cpu(), sigmas=batch.scheduler.sigmas.detach().cpu().clone(), + latent_step_indices=torch.tensor( + state["step_latent_indices"], dtype=torch.long + ), ) if env is not None and batch.rollout_return_denoising_env: diff --git a/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py b/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py index 75359a82cefe..bf92ab67d73a 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py +++ b/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py @@ -65,6 +65,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.lora.pipeline import ( LoRAPipeline, stack_or_compose_fused_lora, + swap_peft_swiglu_fc1_lora_b, ) from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import maybe_download_model from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger @@ -263,12 +264,18 @@ def _resolve_lora_ipc_layer_dict_key( if map_name is None: return None, layer_prefix, None - mapped_name, merge_index = map_name(f"{layer_prefix}.weight") - mapped = _strip_param_weight_suffix(mapped_name) - if mapped != layer_prefix: - layer = layer_dict.get(mapped) - if layer is not None: - return layer, mapped, merge_index + # miles-h3 disk mapping is LoRA-suffix-only; main also maps ".weight". + for suffix in (".weight", ".lora_A"): + mapped_name, merge_index = map_name(f"{layer_prefix}{suffix}") + mapped = ( + mapped_name[: -len(suffix)] + if mapped_name.endswith(suffix) + else _strip_param_weight_suffix(mapped_name) + ) + if mapped != layer_prefix: + layer = layer_dict.get(mapped) + if layer is not None: + return layer, mapped, merge_index return None, layer_prefix, None @@ -707,7 +714,7 @@ def _update_lora_from_tensor( plain_pairs: list[tuple[torch.Tensor, torch.Tensor, Any]] = [] fused_sections: dict[Any, dict[int, tuple[torch.Tensor, torch.Tensor]]] = {} for layer_name, (lora_a, lora_b) in pairs.items(): - layer, _resolved_key, merge_index = _resolve_lora_ipc_layer_dict_key( + layer, resolved_key, merge_index = _resolve_lora_ipc_layer_dict_key( layer_name, layer_dict, dit_module ) if layer is None: @@ -719,6 +726,15 @@ def _update_lora_from_tensor( unknown_layers.append(layer_name) skipped += 1 continue + # Same PEFT → native rewrite as disk load_lora_adapter: fused QKV + # uses merge_index; H3 gated FFN swaps [up, gate] → [gate, up]. + # Native names already in lora_layers skip both transforms. + lora_a = swap_peft_swiglu_fc1_lora_b( + f"{layer_name}.lora_A", f"{resolved_key}.lora_A", lora_a + ) + lora_b = swap_peft_swiglu_fc1_lora_b( + f"{layer_name}.lora_B", f"{resolved_key}.lora_B", lora_b + ) inferred_rank = int(lora_a.shape[-2]) if lora_rank is not None and lora_rank != inferred_rank: logger.warning( diff --git a/python/sglang/multimodal_gen/test/unit/test_fused_lora_compose.py b/python/sglang/multimodal_gen/test/unit/test_fused_lora_compose.py index cdb221e31693..7dc45e100701 100644 --- a/python/sglang/multimodal_gen/test/unit/test_fused_lora_compose.py +++ b/python/sglang/multimodal_gen/test/unit/test_fused_lora_compose.py @@ -8,6 +8,7 @@ import torch import torch.nn.functional as F +from sglang.multimodal_gen.configs.models.dits.minimax_h3 import MiniMaxH3DiTArchConfig from sglang.multimodal_gen.runtime.layers.lora.linear import ( MergedColumnParallelLinearWithLoRA, ) @@ -15,6 +16,7 @@ LoRAPipeline, _store_fused_lora_groups, stack_or_compose_fused_lora, + swap_peft_swiglu_fc1_lora_b, ) from sglang.multimodal_gen.runtime.post_training.weights_updater import ( _resolve_lora_ipc_layer_dict_key, @@ -397,3 +399,44 @@ def test_apply_composed_adapter_end_to_end(): + _reference_delta(x, a_list, b_list, adapter_alpha) * strength ) torch.testing.assert_close(out, expected, rtol=1e-5, atol=1e-5) + + +def test_h3_ipc_reuses_disk_mapping_and_ffn_swap(): + module = torch.nn.Module() + module.param_names_mapping = MiniMaxH3DiTArchConfig().param_names_mapping + qkv, fc1 = object(), object() + layer_dict = { + "blocks.0.attn.qkv_proj": qkv, + "blocks.0.mlp.fc1": fc1, + "token_refiner.blocks.1.attn.qkv_proj": qkv, + } + cases = [ + ("transformer_blocks.0.attn.to_k", qkv, "blocks.0.attn.qkv_proj", 1), + ("transformer_blocks.0.ff.net.0.proj", fc1, "blocks.0.mlp.fc1", None), + ( + "token_refiner.refiner_blocks.1.attn.to_q", + qkv, + "token_refiner.blocks.1.attn.qkv_proj", + 0, + ), + ] + for src, layer, key, merge_index in cases: + assert _resolve_lora_ipc_layer_dict_key(src, layer_dict, module) == ( + layer, + key, + merge_index, + ) + + peft_b = torch.arange(8, dtype=torch.float32).reshape(4, 2) + swapped = swap_peft_swiglu_fc1_lora_b( + "transformer_blocks.0.ff.net.0.proj.lora_B", + "blocks.0.mlp.fc1.lora_B", + peft_b, + ) + torch.testing.assert_close(swapped, torch.cat([peft_b[2:], peft_b[:2]])) + assert ( + swap_peft_swiglu_fc1_lora_b( + "blocks.0.mlp.fc1.lora_B", "blocks.0.mlp.fc1.lora_B", swapped + ) + is swapped + ) diff --git a/python/sglang/multimodal_gen/test/unit/test_msgpack_splice.py b/python/sglang/multimodal_gen/test/unit/test_msgpack_splice.py new file mode 100644 index 000000000000..faa6dd0487f4 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_msgpack_splice.py @@ -0,0 +1,67 @@ +"""Unit tests for msgpack_encode_spliced. + +Mental model (what we exercise): + + payload = {meta..., "data": , nested: [...]} + | + msgpack_encode_spliced(payload) + | + [small-parts bytes][big bytes BY REFERENCE][small-parts bytes] + | + b"".join(parts) == msgspec.msgpack.encode(payload) (byte-identical) + +Covered: +1. Joined parts are byte-identical to a plain msgspec encode, across nested + dicts/lists, >15-element containers (map16/array16 headers), and all + scalar leaves the rollout payload uses. +2. Large bytes land in the parts list by reference (zero copy), small bytes + stay inline. +""" + +import unittest + +import msgspec + +from sglang.multimodal_gen.runtime.entrypoints.post_training.utils import ( + msgpack_encode_spliced, +) + + +def _sample_payload(big: bytes) -> list: + wide_map = {f"k{i}": i for i in range(20)} + return [ + { + "request_id": "req-0", + "seed": 7, + "scale": 1.5, + "flag": True, + "missing": None, + "dit_trajectory": { + "timesteps": {"__tensor__": True, "data": b"tiny", "shape": [4]}, + "latents": {"__tensor__": True, "data": big, "shape": [2, 3]}, + }, + "wide": wide_map, + "long_list": list(range(30)), + } + ] + + +class TestMsgpackEncodeSpliced(unittest.TestCase): + def test_byte_identical_to_msgspec(self): + big = bytes(range(256)) * 8192 + payload = _sample_payload(big) + parts = msgpack_encode_spliced(payload, threshold=1 << 10) + self.assertEqual(b"".join(parts), msgspec.msgpack.encode(payload)) + + def test_large_bytes_spliced_by_reference(self): + big = b"\x00" * (1 << 20) + payload = _sample_payload(big) + parts = msgpack_encode_spliced(payload, threshold=1 << 10) + self.assertTrue(any(part is big for part in parts)) + decoded = msgspec.msgpack.decode(b"".join(parts)) + self.assertEqual(decoded[0]["dit_trajectory"]["latents"]["data"], big) + self.assertEqual(decoded[0]["dit_trajectory"]["timesteps"]["data"], b"tiny") + + +if __name__ == "__main__": + unittest.main() diff --git a/python/sglang/multimodal_gen/test/unit/test_rollout_api.py b/python/sglang/multimodal_gen/test/unit/test_rollout_api.py index 64b0298c0f12..e4f84d27f03a 100644 --- a/python/sglang/multimodal_gen/test/unit/test_rollout_api.py +++ b/python/sglang/multimodal_gen/test/unit/test_rollout_api.py @@ -1,10 +1,24 @@ -"""Unit tests for the rollout generate API (serialization, io_struct, rollout_api).""" +"""Unit tests for the rollout generate API (serialization, io_struct, rollout_api). +Request timing model (TestRequestStamps / TestStagesHeader): + + srv_recv .. forward .. build .. dump .. msgpack (absolute wall-clock marks) + -> all marks ride the x-sgld-timing response header + stages_header: engine per-stage milliseconds -> x-sgld-stages header +""" + +import json +import time import types import unittest import torch +from sglang.multimodal_gen.runtime.entrypoints.post_training.request_timing import ( + MARKS, + RequestStamps, + stages_header, +) from sglang.multimodal_gen.runtime.entrypoints.post_training.utils import ( _maybe_deserialize, _maybe_serialize, @@ -416,5 +430,66 @@ def test_sampling_params_exposes_filters_via_req_getattr(self): self.assertEqual(req.rollout_return_step_indices, [1, 3]) +class TestRequestStamps(unittest.TestCase): + """The client joins these marks with its own, so what matters is that the + header decodes, keeps wall-clock order, and stays readable when a request + returned early with only some marks taken.""" + + def _decode(self, stamps: RequestStamps) -> dict: + header = stamps.to_header() + self.assertEqual(header, header.encode("ascii").decode()) + return json.loads(header) + + def test_marks_decode_in_wall_clock_order(self): + before = time.time() + stamps = RequestStamps("req-1") + for name in MARKS: + stamps.mark(name) + decoded = self._decode(stamps) + self.assertEqual(decoded.pop("request_id"), "req-1") + self.assertEqual(set(decoded), set(MARKS)) + values = [decoded[name] for name in MARKS] + self.assertEqual(values, sorted(values)) + self.assertGreaterEqual(values[0], round(before, 6) - 1e-6) + + def test_span_pairs_start_and_end_marks(self): + stamps = RequestStamps() + with stamps.span("forward"): + pass + with stamps.span("dump"): + pass + self.assertEqual( + set(json.loads(stamps.to_header())), + {"forward_start", "forward_end", "dump_end"}, + ) + + def test_partial_marks_only_report_what_was_taken(self): + stamps = RequestStamps("req-3") + stamps.mark("srv_recv") + self.assertEqual(set(self._decode(stamps)), {"srv_recv", "request_id"}) + + def test_unknown_mark_is_rejected(self): + with self.assertRaises(AssertionError): + RequestStamps().mark("no_such_mark") + + +class TestStagesHeader(unittest.TestCase): + """The engine's own stage breakdown is the only view inside the forward the + client waits on, so it has to survive the trip and stay absent when unset.""" + + def test_stages_are_reported_in_milliseconds(self): + metrics = types.SimpleNamespace( + stages={"decoding": 8123.4567, "text_encoding": 91.2} + ) + self.assertEqual( + json.loads(stages_header(metrics)), + {"decoding": 8123.457, "text_encoding": 91.2}, + ) + + def test_no_metrics_or_no_stages_yields_no_header(self): + self.assertEqual(stages_header(None), "") + self.assertEqual(stages_header(types.SimpleNamespace(stages={})), "") + + if __name__ == "__main__": unittest.main() diff --git a/python/sglang/srt/arg_groups/attention_hook.py b/python/sglang/srt/arg_groups/attention_hook.py index e4f96828561b..db600d889c5f 100644 --- a/python/sglang/srt/arg_groups/attention_hook.py +++ b/python/sglang/srt/arg_groups/attention_hook.py @@ -31,6 +31,10 @@ from sglang.srt.environ import envs from sglang.srt.model_executor.cuda_graph_config import Backend, Phase, with_phase from sglang.srt.runtime_context import get_platform +from sglang.srt.true_on_policy.contracts import ( + resolve_true_on_policy_runtime_policy, + validate_true_on_policy_contract, +) from sglang.srt.utils.common import ( parse_connector_type, ) @@ -497,20 +501,55 @@ def handle_deterministic_inference(server_args: Any): ) cfg = resolving_view(server_args) + validate_true_on_policy_contract(cfg) + + if cfg.enable_prefill_only_deterministic_inference: + declare_resolution( + server_args, + "_handle_deterministic_inference", + enable_deterministic_inference=True, + ) + if cfg.rl_on_policy_target is not None: - logger.warning("Enable deterministic inference because of rl_on_policy_target.") + logger.warning( + "Enable deterministic inference because of legacy rl_on_policy_target." + ) declare_resolution( server_args, "_handle_deterministic_inference", enable_deterministic_inference=True, ) + if cfg.true_on_policy_contract is not None: + logger.warning( + "Enable deterministic inference because of true_on_policy_contract." + ) + declare_resolution( + server_args, + "_handle_deterministic_inference", + enable_deterministic_inference=True, + ) # For VLM envs.SGLANG_VLM_CACHE_SIZE_MB.set(0) - # TODO remove this environment variable as a whole - envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.set(True) + + if ( + resolve_true_on_policy_runtime_policy( + cfg + ).disable_flashinfer_allreduce_fusion + and cfg.enable_flashinfer_allreduce_fusion + ): + declare_resolution( + server_args, + "_handle_deterministic_inference", + enable_flashinfer_allreduce_fusion=False, + ) + logger.warning( + "Disable flashinfer allreduce fusion because of " + "true_on_policy_contract with TP rollout." + ) if cfg.enable_deterministic_inference: + envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.set("1") if cfg.enable_aiter_allreduce_fusion: logger.warning( "Disable --enable-aiter-allreduce-fusion because deterministic inference is enabled." diff --git a/python/sglang/srt/arg_groups/choices.py b/python/sglang/srt/arg_groups/choices.py index 8ce2c00ecbef..91ffe5183c5a 100644 --- a/python/sglang/srt/arg_groups/choices.py +++ b/python/sglang/srt/arg_groups/choices.py @@ -189,7 +189,7 @@ RADIX_EVICTION_POLICY_CHOICES = ["lru", "lfu", "slru", "priority"] -RL_ON_POLICY_TARGET_CHOICES = ["fsdp"] +RL_ON_POLICY_TARGET_CHOICES = ["fsdp", "fsdp_tp"] LINEAR_ATTN_KERNEL_BACKEND_CHOICES = [ "triton", diff --git a/python/sglang/srt/arg_groups/field_order.py b/python/sglang/srt/arg_groups/field_order.py index 992a91b1728a..2d2c51f41e08 100644 --- a/python/sglang/srt/arg_groups/field_order.py +++ b/python/sglang/srt/arg_groups/field_order.py @@ -404,6 +404,7 @@ "experts_shared_outer_loras", "lora_use_virtual_experts", "lora_strict_loading", + "lora_no_cpu_backup", "lora_drain_wait_threshold", "enable_two_batch_overlap", "enable_single_batch_overlap", diff --git a/python/sglang/srt/arg_groups/fields/exec_.py b/python/sglang/srt/arg_groups/fields/exec_.py index e70086e61119..35d7d61cd017 100644 --- a/python/sglang/srt/arg_groups/fields/exec_.py +++ b/python/sglang/srt/arg_groups/fields/exec_.py @@ -506,6 +506,9 @@ class ExecGraph(msgspec.Struct): "Disable the decode-phase CUDA graph. Convenience for --cuda-graph-backend-decode=disabled.", ] = False disable_cuda_graph: A[bool, Arg(no_cli=True)] = False + disable_draft_cuda_graph: A[ + bool, "Disable cuda graph for draft model in speculative decoding." + ] = False disable_cuda_graph_padding: A[ bool, "Disable cuda graph when padding is needed. Still uses cuda graph when padding is not needed.", @@ -939,6 +942,15 @@ class ExecDeterministic(msgspec.Struct): enable_deterministic_inference: A[ bool, "Enable deterministic inference mode with batch invariant ops." ] = False + enable_prefill_only_deterministic_inference: A[ + bool, + "Enable prefill-only deterministic inference mode with batch invariant ops.", + ] = False + true_on_policy_contract: A[ + Optional[str], + "Internal true-on-policy parity contract selected by the launcher. " + "Normal users should prefer the Miles true_on_policy switch.", + ] = None rl_on_policy_target: A[ Optional[str], Arg( diff --git a/python/sglang/srt/arg_groups/fields/lora.py b/python/sglang/srt/arg_groups/fields/lora.py index 674aa504f51a..c9b1b858dd8f 100644 --- a/python/sglang/srt/arg_groups/fields/lora.py +++ b/python/sglang/srt/arg_groups/fields/lora.py @@ -24,10 +24,6 @@ ) from sglang.srt.arg_groups.argparse_actions import LoRAPathAction from sglang.srt.lora.lora_registry import LoRARef -from sglang.srt.utils.common import ( - LORA_TARGET_ALL_MODULES, - SUPPORTED_LORA_TARGET_MODULES, -) class Lora(msgspec.Struct): @@ -53,9 +49,12 @@ class Lora(msgspec.Struct): lora_target_modules: A[ Optional[Union[set[str], List[str]]], Arg( - help="The union set of all target modules where LoRA should be applied. If not specified, it will be automatically inferred from the adapters provided in --lora-paths. If 'all' is specified, all supported modules will be targeted.", + help=( + "HF module names or paths where LoRA should be applied; fused names are " + "resolved by the loaded model. If not specified, infer targets from " + "--lora-paths. Use 'all' to target all supported modules." + ), nargs="*", - choices=SUPPORTED_LORA_TARGET_MODULES + [LORA_TARGET_ALL_MODULES], ), ] = None lora_paths: A[ @@ -113,6 +112,10 @@ class Lora(msgspec.Struct): action=argparse.BooleanOptionalAction, ), ] = False + lora_no_cpu_backup: A[ + bool, + "Keep streamed LoRA weights on the GPU and drop the staged copy once the adapter is installed in the memory pool, instead of holding a host copy per TP rank. Evicting an adapter from the pool is then refused, so every loaded adapter needs its own slot.", + ] = False lora_drain_wait_threshold: A[ float, "When any LoRA adapter request waits longer than this threshold (in seconds), the scheduler will selectively drain one running adapter to make room. This mitigates extreme tail latency under high or skewed workloads by preventing a small set of adapters from monopolizing batch slots. Set to 0 to disable draining (default).", diff --git a/python/sglang/srt/arg_groups/fields/model.py b/python/sglang/srt/arg_groups/fields/model.py index 606279c01a8a..5aab8b92060b 100644 --- a/python/sglang/srt/arg_groups/fields/model.py +++ b/python/sglang/srt/arg_groups/fields/model.py @@ -285,6 +285,10 @@ class Model(msgspec.Struct): nargs="*", ), ] = None + custom_pull_weights_pre_read_hook: A[ + Optional[str], + "Import path of a hook(source_dir, target_version) that /pull_weights calls before reading the published weights. POSIX shared filesystems need no hook; object-store-backed mounts often lack cross-host read-after-write consistency, so another host's writes only become visible after an explicit refresh.", + ] = None weight_loader_disable_mmap: A[ bool, "Disable mmap while loading weight using safetensors.", @@ -322,6 +326,14 @@ class Model(msgspec.Struct): bool, "Start seed server via transfer engine backend for remote instance weight loader.", ] = False + enable_engine_info_bootstrap: A[ + bool, + "Start the EngineInfoBootstrapServer and register per-rank parallelism config, without the mooncake/verbs P2P transfer-engine seeding.", + ] = False + enable_rdt_weight_sync: A[ + bool, + "Expose SchedulerActor.pull_weights for RDT (Ray Direct Transport / NIXL) weight sync. Requires --use-ray; implies --enable-engine-info-bootstrap.", + ] = False engine_info_bootstrap_port: A[ int, "Port for the engine info bootstrap server. Default is 6789. Must be set explicitly when running multiple instances on the same node.", diff --git a/python/sglang/srt/arg_groups/model_overrides/qwen3_moe.py b/python/sglang/srt/arg_groups/model_overrides/qwen3_moe.py index 1747e5192ddb..b95e16cb67dc 100644 --- a/python/sglang/srt/arg_groups/model_overrides/qwen3_moe.py +++ b/python/sglang/srt/arg_groups/model_overrides/qwen3_moe.py @@ -86,9 +86,15 @@ def _qwen3_moe_family_overrides(server_args: Any, hf_config: Any) -> dict: and cfg.moe_a2a_backend == "none" and cfg.moe_runner_backend == "auto" ): - overrides["moe_runner_backend"] = "flashinfer_trtllm" + # The fused TRT-LLM BF16 path cannot inject expert LoRA deltas. + use_lora = quantization is None and ( + cfg.enable_lora or (cfg.enable_lora is None and bool(cfg.lora_paths)) + ) + overrides["moe_runner_backend"] = ( + "triton" if use_lora else "flashinfer_trtllm" + ) logger.info( - "Use flashinfer_trtllm as MoE runner backend on sm100 for " + f"Use {overrides['moe_runner_backend']} as MoE runner backend on sm100 for " f"{hf_config.architectures[0]}" ) return overrides diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 6d26836debea..8abe38beae10 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -989,7 +989,8 @@ def _sparse_head_overlap_disable(view: Any) -> dict: { "DeepseekV3ForCausalLM", "DeepseekV32ForCausalLM", - "DeepseekV4ForCausalLM", + # DeepseekV4 is absent on purpose: under EP + NextN its workspace + # creation barriers against the cuda-graph capture loop's own barrier. "GptOssForCausalLM", "GlmMoeDsaForCausalLM", "Glm4MoeForCausalLM", @@ -1422,7 +1423,12 @@ def _page_size_default(view: Any) -> dict: @register_post_process def _data_parallelism_defaults(view: Any) -> dict: if view.dp_size == 1 and view.ep_join_mode != "scale": - return {"enable_dp_attention": False, "enable_dp_lm_head": False} + overrides = {"enable_dp_attention": False} + # Keep the dp LM head when attention context parallelism is enabled, + # even without data parallelism (context-parallel LM head). + if not (view.enable_dp_lm_head and view.attn_cp_size > 1): + overrides["enable_dp_lm_head"] = False + return overrides return {} @@ -1481,8 +1487,9 @@ def _dp_lm_head_validation(view: Any) -> dict: dp LM head and the TP LM-head all-to-all path. Reads the mid-resolution values through the view.""" if view.enable_dp_lm_head: - assert view.enable_dp_attention, ( - "Please enable dp attention when setting enable_dp_lm_head. " + assert view.enable_dp_attention or view.attn_cp_size > 1, ( + "Please enable dp attention when setting enable_dp_lm_head, " + "unless attention context parallelism is enabled." ) if view.enable_tp_lm_head_all_to_all: assert view.enable_dp_attention, ( diff --git a/python/sglang/srt/arg_groups/pipeline.py b/python/sglang/srt/arg_groups/pipeline.py index 9b982c232d85..d5bc0c3c34a2 100644 --- a/python/sglang/srt/arg_groups/pipeline.py +++ b/python/sglang/srt/arg_groups/pipeline.py @@ -81,6 +81,7 @@ def run_resolution_pipeline(server_args: Any) -> None: handle_multimodal, handle_other_validations, handle_prefill_delayer_env_compat, + handle_rdt_weight_sync, handle_return_hidden_states_mode, handle_ssl_validation, handle_tokenizer_batching, @@ -97,6 +98,7 @@ def run_resolution_pipeline(server_args: Any) -> None: from sglang.srt.arg_groups.memory_hook import handle_offload_compatibility run_hook(handle_offload_compatibility, server_args) + run_hook(handle_rdt_weight_sync, server_args) from sglang.srt.arg_groups.validation_hook import ( default_unset_prefill_decode_interval, validate_experimental_sgl_marlin, diff --git a/python/sglang/srt/arg_groups/serving_hook.py b/python/sglang/srt/arg_groups/serving_hook.py index 2015ebf0e14b..bfcb6019f8cb 100644 --- a/python/sglang/srt/arg_groups/serving_hook.py +++ b/python/sglang/srt/arg_groups/serving_hook.py @@ -389,6 +389,20 @@ def handle_deprecated_args(server_args: Any): ) +def handle_rdt_weight_sync(server_args: Any): + cfg = resolving_view(server_args) + if not cfg.enable_rdt_weight_sync: + return + assert cfg.use_ray, ( + "--enable-rdt-weight-sync requires --use-ray: the trainer pulls weights through SchedulerActors." + ) + declare_resolution( + server_args, + "_handle_rdt_weight_sync", + enable_engine_info_bootstrap=True, + ) + + def handle_environment_variables(server_args: Any): cfg = resolving_view(server_args) handle_multimodal_feature_transport(server_args) diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index e7ce9277aba2..9a363044cd21 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -346,7 +346,10 @@ def dsa_layer_skips_topk(config: PretrainedConfig, layer_id: int) -> bool: def get_dsa_index_n_heads(config: PretrainedConfig) -> int: - assert is_deepseek_dsa(config) + # Permit both DSA (V3.2-family) and V4: both carry the indexer (index_n_heads) and this must + # match get_dsa_index_head_dim's contract, else LoRA buffer init for indexer.wq_b / + # indexer.weights_proj on a V4 model asserts here while indexer.wk (which uses head_dim) succeeds. + assert is_deepseek_dsa(config) or is_deepseek_v4(config) return config.index_n_heads diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 6d806fda3388..6bf54a9c6423 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -2202,11 +2202,11 @@ def _commit_transfer_to_req(self, decode_req: DecodeRequest): # under the current weights and sampled a fresh handoff token, but when # there is a remembered boundary token we are *replaying* an # already-emitted token. Override the handoff with it, and skip - # re-committing a logprob for it -- it keeps its original behavior - # logprob from before the retract (we never re-score generated tokens - # under the new policy). A rebootstrap with no boundary token (retracted - # before emitting any output) falls through to the normal path so its - # first token and logprob are committed as usual. + # re-committing a logprob or sampling mask for it -- it keeps its original + # behavior logprob and sampling mask from before the retract (we never + # re-score generated tokens under the new policy). A rebootstrap with no + # boundary token (retracted before emitting any output) falls through to + # the normal path so its first token and logprob are committed as usual. replayed_boundary = ( decode_req.is_rebootstrap and decode_req.req.pd_rebootstrap_forced_output_id is not None @@ -2275,21 +2275,20 @@ def _commit_transfer_to_req(self, decode_req: DecodeRequest): : decode_req.req.logprob.top_logprobs_num ].tolist() ) - if decode_req.req.return_sampling_mask: + if decode_req.req.return_sampling_mask and not replayed_boundary: assert output_token_sampling_mask_idx is not None, ( "sampling mask buffer disabled on decode side" ) sampling_mask_len = int(output_token_sampling_mask_len[0].item()) - if sampling_mask_len < 0: - decode_req.req.output_token_sampling_mask.append(None) - decode_req.req.output_token_sampling_logprobs.append(None) - else: - decode_req.req.output_token_sampling_mask.append( - output_token_sampling_mask_idx[:sampling_mask_len].cpu().tolist() - ) - decode_req.req.output_token_sampling_logprobs.append( - float(output_token_sampling_logprobs[0].item()) - ) + num_logprobs = ( + sampling_mask_len + if decode_req.req.sampling_logprobs_mode == "support" + else 1 + ) + decode_req.req.sampling_mask_rows.append( + output_token_sampling_mask_idx[:sampling_mask_len].cpu().numpy(), + output_token_sampling_logprobs[:num_logprobs].cpu().numpy(), + ) decode_req.kv_receiver.clear() decode_req.kv_receiver = None diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 56604a281a55..fb4ee5021196 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -998,9 +998,11 @@ def process_disagg_prefill_inflight_queue( # todo: set Transferring correctly in backend undone_reqs.append(req) elif poll == KVPoll.Success: # transfer done + if req.return_routed_experts: + self.batch_result_processor._maybe_collect_routed_experts(req) + release_kv_cache(req, self.tree_cache) # unlock the tree if not isinstance(req.finished_reason, FINISH_ABORT): req.finished_reason = FINISH_LENGTH(length=0) - release_kv_cache(req, self.tree_cache) # unlock the tree self.tree_cache.finish( req.cache_request_handle, CacheRequestOutcome.SUCCESS ) diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index c6bd1080c502..5e519d6c45db 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -388,7 +388,9 @@ def __init__( (size, max_sampling_mask_tokens), dtype=torch.int32, device=device ) self.output_token_sampling_logprobs = torch.zeros( - (size, 16), dtype=torch.float32, device=device + (size, max_sampling_mask_tokens), + dtype=torch.float32, + device=device, ) # For PD + spec decode self.output_topk_p = torch.zeros( @@ -528,38 +530,17 @@ def set_buf(self, req: Req): device="cpu", ) if req.return_sampling_mask: - # Sentinel -1: the decode side records None for this handoff token. - self.output_token_sampling_mask_len[req.metadata_buffer_index][0] = -1 - sampling_masks = req.output_token_sampling_mask - sampling_logprobs = req.output_token_sampling_logprobs - if sampling_masks: - sampling_mask = sampling_masks[0] - sampling_logprob = sampling_logprobs[0] if sampling_logprobs else None - if sampling_mask is not None and sampling_logprob is not None: - mask_len = len(sampling_mask) - max_mask_len = self.output_token_sampling_mask_idx.shape[1] - if mask_len > max_mask_len: - raise RuntimeError( - f"Sampling mask length {mask_len} exceeds disaggregation " - f"metadata capacity {max_mask_len}. Increase " - "--sampling-mask-max-tokens." - ) - self.output_token_sampling_mask_len[req.metadata_buffer_index][ - 0 - ] = mask_len - if mask_len: - self.output_token_sampling_mask_idx[ - req.metadata_buffer_index, :mask_len - ].copy_( - torch.tensor( - sampling_mask, - dtype=torch.int32, - device=self.output_token_sampling_mask_idx.device, - ) - ) - self.output_token_sampling_logprobs[req.metadata_buffer_index][ - 0 - ] = float(sampling_logprob) + # Prefill streams a request only once its KV transfer ends or it aborts, + # so the first token's row is the only one queued here. + chunk = req.sampling_mask_rows.view() + mask_len = len(chunk.token_ids) + self.output_token_sampling_mask_len[req.metadata_buffer_index][0] = mask_len + self.output_token_sampling_mask_idx[ + req.metadata_buffer_index, :mask_len + ].copy_(torch.from_numpy(chunk.token_ids)) + self.output_token_sampling_logprobs[ + req.metadata_buffer_index, : len(chunk.logprobs) + ].copy_(torch.from_numpy(chunk.logprobs)) # For PD + spec decode if req.hidden_states_tensor is not None: # speculative_eagle_topk should not be greater than 16 currently diff --git a/python/sglang/srt/distributed/communication_op.py b/python/sglang/srt/distributed/communication_op.py index 31752f89b06f..6a1e9ebb854b 100644 --- a/python/sglang/srt/distributed/communication_op.py +++ b/python/sglang/srt/distributed/communication_op.py @@ -7,6 +7,8 @@ import torch import torch.distributed +from sglang.srt.tp_invariant_ops import tree_all_reduce_sum +from sglang.srt.true_on_policy import should_use_tp_invariant_tree_all_reduce from sglang.srt.utils import broadcast_pyobj from .parallel_state import ( @@ -20,6 +22,8 @@ def tensor_model_parallel_all_reduce(input_: torch.Tensor) -> torch.Tensor: """All-reduce the input tensor across model parallel group.""" + if should_use_tp_invariant_tree_all_reduce(): + return tree_all_reduce_sum(input_, device_group=get_tp_group().device_group) return get_tp_group().all_reduce(input_) @@ -118,6 +122,10 @@ def attn_cp_tp_broadcast_pyobj(data: List[Any]) -> List[Any]: def attention_tensor_model_parallel_all_reduce(input_: torch.Tensor) -> torch.Tensor: """All-reduce the input tensor across attention parallel group.""" + if should_use_tp_invariant_tree_all_reduce(): + return tree_all_reduce_sum( + input_, device_group=get_attn_tp_group().device_group + ) return get_attn_tp_group().all_reduce(input_) diff --git a/python/sglang/srt/distributed/device_communicators/triton_symm_mem_ag.py b/python/sglang/srt/distributed/device_communicators/triton_symm_mem_ag.py index a7cf40bd7e02..a403ae6fd677 100644 --- a/python/sglang/srt/distributed/device_communicators/triton_symm_mem_ag.py +++ b/python/sglang/srt/distributed/device_communicators/triton_symm_mem_ag.py @@ -16,6 +16,7 @@ import triton import triton.language as tl +from sglang.srt.environ import envs from sglang.srt.runtime_context import ( get_parallel, get_schedule, @@ -465,6 +466,8 @@ def __init__( ): self._max_tokens = int(max_tokens) self._skip_entry_sync = skip_entry_sync + if envs.SGLANG_DISABLE_MULTIMEM_AG.get(): + enabled = False # None => always NCCL; _UNINIT => build on first eager call. self._state = self._UNINIT if enabled else None if self._state is self._UNINIT: diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 8383132c0457..4489a10bdaf1 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -31,11 +31,11 @@ import weakref from collections import namedtuple from contextlib import contextmanager, nullcontext -from dataclasses import dataclass +from dataclasses import asdict, dataclass from datetime import timedelta from multiprocessing import shared_memory from typing import Any, Callable, Dict, List, Optional, Tuple, Union -from unittest.mock import patch +from unittest.mock import MagicMock, patch import torch import torch.distributed @@ -2990,7 +2990,10 @@ def get_dcp_rank(): def get_tensor_model_parallel_rank(): """Return my rank for the tensor model parallel group.""" - return get_tp_group().rank_in_group + try: + return get_tp_group().rank_in_group + except Exception: + return 0 # ATTN_TP @@ -3264,3 +3267,184 @@ def monkey_patch_vllm_parallel_state(reverse: bool = False): setattr(vllm_parallel_state, "get_pp_group", get_pp_group) setattr(vllm_parallel_state, "get_tp_group", get_tp_group) setattr(vllm_parallel_state, "get_world_group", get_world_group) + + +@dataclass +class RankParallelismConfig: + """ + Complete parallelism configuration for a single inference rank. + + This configuration captures all the parallelism settings needed to recreate + a model shard outside of sglang. It supports: + - TP/PP/EP for model parallelism + - MoE-TP/Attn-TP/Attn-DP for MoE and DP attention. + """ + + tp_size: int = 1 + tp_rank: int = 0 + pp_size: int = 1 + pp_rank: int = 0 + ep_size: int = 1 + ep_rank: int = 0 + moe_tp_size: int = 1 + moe_tp_rank: int = 0 + attn_tp_size: int = 1 + attn_tp_rank: int = 0 + attn_dp_size: int = 1 + attn_dp_rank: int = 0 + attn_cp_size: int = 1 + attn_cp_rank: int = 0 + moe_dp_size: int = 1 + moe_dp_rank: int = 0 + + world_size: int = 1 + global_rank: int = 0 + local_rank: int = 0 + + def to_dict(self) -> Dict[str, Any]: + """Convert to dictionary for serialization.""" + return asdict(self) + + @classmethod + def from_dict(cls, data: Dict[str, Any]) -> "RankParallelismConfig": + """Create from dictionary, filtering unknown fields.""" + import dataclasses + + valid_fields = {f.name for f in dataclasses.fields(cls)} + filtered_data = {k: v for k, v in data.items() if k in valid_fields} + return cls(**filtered_data) + + @classmethod + def from_parallel_state(cls, local_rank: int = 0) -> "RankParallelismConfig": + """Extract current parallelism settings from the global parallel state.""" + tp_size = get_tensor_model_parallel_world_size() + tp_rank = get_tensor_model_parallel_rank() + + # v0.5.16 retired the get_attention_{tp,dp,cp}_* helpers from dp_attention; + # the attention-side ranks/sizes now live on the runtime ParallelState. + # Imported lazily to avoid a circular import. + from sglang.srt.runtime_context import get_parallel + + ps = get_parallel() + + return cls( + tp_size=tp_size, + tp_rank=tp_rank, + pp_size=get_pipeline_model_parallel_world_size(), + pp_rank=get_pipeline_model_parallel_rank(), + ep_size=get_moe_expert_parallel_world_size(), + ep_rank=get_moe_expert_parallel_rank(), + moe_tp_size=get_moe_tensor_parallel_world_size(), + moe_tp_rank=get_moe_tensor_parallel_rank(), + attn_tp_size=ps.attn_tp_size, + attn_tp_rank=ps.attn_tp_rank, + attn_dp_size=ps.attn_dp_size, + attn_dp_rank=ps.attn_dp_rank, + attn_cp_size=ps.attn_cp_size, + attn_cp_rank=ps.attn_cp_rank, + moe_dp_size=get_moe_data_parallel_world_size(), + moe_dp_rank=get_moe_data_parallel_rank(), + world_size=( + torch.distributed.get_world_size() + if torch.distributed.is_initialized() + else 1 + ), + global_rank=( + torch.distributed.get_rank() + if torch.distributed.is_initialized() + else 0 + ), + local_rank=local_rank, + ) + + +# Globals on parallel_state module to save/restore +_PS_GLOBALS = ("_TP", "_PP", "_MOE_EP", "_MOE_TP", "_ATTN_TP", "_ATTN_CP", "_MOE_DP") +# Globals on dp_attention module to save/restore. v0.5.16 moved the dp-attention +# enable flag out of this module onto the runtime flags (get_flags().dp.enabled), +# so it is saved/restored separately below rather than as a module attribute. +_DA_GLOBALS = ("_ATTN_DP_RANK", "_ATTN_DP_SIZE") + + +class ParallelismContext: + """ + Context manager for creating model replicas with specific parallelism settings. + + Temporarily sets global variables to allow creating model shards outside of a + real distributed environment. + Usage: + with ParallelismContext(RankParallelismConfig.from_dict(parallelism_info)): + model = get_model(...) + """ + + def __init__(self, parallelism_config: RankParallelismConfig): + self.config = parallelism_config + self._original_globals: Dict[str, Any] = {} + + def _create_mock_group(self, world_size: int, rank_in_group: int): + """Create a mock group coordinator with all necessary properties.""" + mock_group = MagicMock() + mock_group.world_size = world_size + mock_group.rank_in_group = rank_in_group + mock_group.rank = rank_in_group + mock_group.local_rank = rank_in_group + mock_group.ranks = list(range(world_size)) + mock_group.first_rank = 0 + mock_group.last_rank = world_size - 1 + mock_group.is_first_rank = rank_in_group == 0 + mock_group.is_last_rank = rank_in_group == world_size - 1 + mock_group.next_rank = mock_group.ranks[(rank_in_group + 1) % world_size] + mock_group.prev_rank = mock_group.ranks[(rank_in_group - 1) % world_size] + return mock_group + + def __enter__(self): + conf = self.config + + from sglang.srt.distributed import parallel_state + from sglang.srt.layers import dp_attention + from sglang.srt.runtime_context import get_flags + + # Save original globals + for name in _PS_GLOBALS: + self._original_globals[name] = getattr(parallel_state, name, None) + for name in _DA_GLOBALS: + self._original_globals[name] = getattr(dp_attention, name, None) + + # Build and set mock group objects on parallel_state + _ps_new_values = { + "_TP": self._create_mock_group(conf.tp_size, conf.tp_rank), + "_PP": self._create_mock_group(conf.pp_size, conf.pp_rank), + "_MOE_EP": self._create_mock_group(conf.ep_size, conf.ep_rank), + "_MOE_TP": self._create_mock_group(conf.moe_tp_size, conf.moe_tp_rank), + "_ATTN_TP": self._create_mock_group(conf.attn_tp_size, conf.attn_tp_rank), + "_ATTN_CP": self._create_mock_group(conf.attn_cp_size, conf.attn_cp_rank), + "_MOE_DP": self._create_mock_group(conf.moe_dp_size, conf.moe_dp_rank), + } + for name, value in _ps_new_values.items(): + setattr(parallel_state, name, value) + + # Set dp_attention scalar globals + dp_attention._ATTN_DP_RANK = conf.attn_dp_rank + dp_attention._ATTN_DP_SIZE = conf.attn_dp_size + # is_dp_attention_enabled() reads this, not a dp_attention module global. + dp_flags = get_flags().dp + self._original_dp_enabled = dp_flags.enabled + dp_flags.enabled = conf.attn_dp_size > 1 + + logger.info(f"[ParallelismContext] Activated: {conf}") + return self + + def __exit__(self, *args): + from sglang.srt.distributed import parallel_state + from sglang.srt.layers import dp_attention + from sglang.srt.runtime_context import get_flags + + # Restore original globals + for name in _PS_GLOBALS: + setattr(parallel_state, name, self._original_globals.get(name)) + for name in _DA_GLOBALS: + setattr(dp_attention, name, self._original_globals.get(name)) + get_flags().dp.enabled = self._original_dp_enabled + + logger.info("[ParallelismContext] Deactivated") + return False diff --git a/python/sglang/srt/entrypoints/anthropic/serving.py b/python/sglang/srt/entrypoints/anthropic/serving.py index b20bbdbab5c4..c8c09eb4d67a 100644 --- a/python/sglang/srt/entrypoints/anthropic/serving.py +++ b/python/sglang/srt/entrypoints/anthropic/serving.py @@ -10,7 +10,7 @@ import json import logging import uuid -from typing import TYPE_CHECKING, Any, AsyncGenerator, Optional, Union +from typing import TYPE_CHECKING, Any, AsyncGenerator, Callable, Optional, Union from fastapi import Request from fastapi.responses import JSONResponse, StreamingResponse @@ -181,6 +181,597 @@ def _scrub_error_message(message: str, status_code: int) -> str: return cleaned or "Request failed" +def convert_to_chat_completion_request( + anthropic_request: AnthropicMessagesRequest, + *, + merge_inline_system: bool, + wrap_reasoning_history: Optional[Callable[[str], str]] = None, + supports_native_reasoning_history: Optional[Callable[[], bool]] = None, + apply_reasoning_enabled: Optional[ + Callable[[ChatCompletionRequest, bool], None] + ] = None, +) -> ChatCompletionRequest: + """Convert an Anthropic Messages request to an OpenAI ChatCompletion request.""" + openai_messages = [] + + def _convert_anthropic_image_source_to_openai_part( + source: Any, + ) -> Optional[dict]: + # Source may arrive as a Pydantic model (typed ImageBlock.source) + # or as a raw dict when parsed from a nested tool_result payload. + if isinstance(source, BaseModel): + source = source.model_dump(exclude_none=True) + if not isinstance(source, dict): + return None + + source_type = source.get("type") + if source_type == "base64": + media_type = source.get("media_type", "image/png") + data = source.get("data", "") + if not data: + return None + return { + "type": "image_url", + "image_url": { + "url": f"data:{media_type};base64,{data}", + }, + } + + url = source.get("url") + if url: + return { + "type": "image_url", + "image_url": { + "url": url, + }, + } + + return None + + def _text_from_search_result(item: dict[str, Any]) -> str: + search_parts = [] + title = item.get("title") + if title: + search_parts.append(f"Title: {title}") + + source = item.get("source") + if isinstance(source, dict): + source_text = source.get("url") or source.get("text") + if source_text: + search_parts.append(f"Source: {source_text}") + elif source: + search_parts.append(f"Source: {source}") + + content = item.get("content") + content_parts = [] + if isinstance(content, str): + content_parts.append(content) + elif isinstance(content, list): + for part in content: + if not isinstance(part, dict): + continue + if part.get("type") == "text" and part.get("text"): + content_parts.append(part["text"]) + if content_parts: + search_parts.append("Content: " + "\n".join(content_parts)) + + return "\n".join(search_parts) + + def _convert_tool_result_content( + content: Any, + ) -> tuple[list[Union[str, list[dict]]], str]: + if isinstance(content, list): + tool_content_parts = [] + tool_text_parts = [] + + for raw_item in content: + # Items may be typed Pydantic blocks (after request + # validation) or raw dicts (from legacy callers). Coerce + # to dict so the existing key-based logic still works. + if isinstance(raw_item, BaseModel): + item = raw_item.model_dump(exclude_none=True) + elif isinstance(raw_item, dict): + item = raw_item + else: + continue + + item_type = item.get("type") + if item_type == "text": + text = item.get("text", "") + if text: + tool_text_parts.append(text) + tool_content_parts.append({"type": "text", "text": text}) + elif item_type == "image": + image_part = _convert_anthropic_image_source_to_openai_part( + item.get("source") + ) + if image_part is not None: + tool_content_parts.append(image_part) + elif item_type == "tool_reference": + # Anthropic uses `tool_name`; the SGLang chat template + # matches on `name`. Translate at the boundary. + ref_name = item.get("tool_name") or item.get("name") + if ref_name: + tool_content_parts.append( + {"type": "tool_reference", "name": ref_name} + ) + elif item_type == "search_result": + search_text = _text_from_search_result(item) + if search_text: + tool_text_parts.append(search_text) + tool_content_parts.append({"type": "text", "text": search_text}) + + tool_text = "\n".join(tool_text_parts) + # GLM templates expand references only at the start of a tool + # message, so isolate reference runs without changing part order. + tool_content_groups: list[list[dict]] = [] + for part in tool_content_parts: + is_reference = part["type"] == "tool_reference" + if ( + not tool_content_groups + or (tool_content_groups[-1][0]["type"] == "tool_reference") + != is_reference + ): + tool_content_groups.append([]) + tool_content_groups[-1].append(part) + + tool_contents: list[Union[str, list[dict]]] = [] + for group in tool_content_groups: + if len(group) == 1 and group[0]["type"] == "text": + tool_contents.append(group[0]["text"]) + else: + tool_contents.append(group) + return tool_contents or [""], tool_text + + tool_text = str(content) if content else "" + return [tool_text], tool_text + + def _convert_assistant_thinking_blocks( + blocks: list[AnthropicContentBlock], + ) -> tuple[Optional[str], Optional[str]]: + """Reconstruct prior-turn thinking as ``(reasoning_content, text)``. + + At most one is set: encoders that frame the reasoning channel take + it as ``reasoning_content``, everything else gets it re-wrapped and + spliced into content. + + ``redacted_thinking`` carries encrypted bytes that no local + parser can interpret, so we raise rather than silently drop it. + On non-reasoning models (no detector configured) the rewrap is + best-effort: we log a warning and drop the thinking text so a + history echo doesn't 400 the whole request — the prior thinking + is opaque context the model didn't need anyway. + """ + if any(block.type == "redacted_thinking" for block in blocks): + raise ValueError("Anthropic redacted_thinking history is not supported") + + thinking_parts = [ + block.thinking + for block in blocks + if block.type == "thinking" and block.thinking + ] + if not thinking_parts: + return None, None + + reasoning_text = "\n".join(thinking_parts) + if ( + supports_native_reasoning_history is not None + and supports_native_reasoning_history() + ): + return reasoning_text, None + if wrap_reasoning_history is None: + raise ValueError( + "Cannot convert thinking history without a reasoning adapter" + ) + try: + return None, wrap_reasoning_history(reasoning_text) + except ValueError as e: + logger.warning( + "Dropping prior-turn thinking history (%d blocks): %s", + len(thinking_parts), + e, + ) + return None, None + + system_parts: list[str] = [] + if anthropic_request.system: + if isinstance(anthropic_request.system, str): + if anthropic_request.system.strip(): + system_parts.append(anthropic_request.system) + else: + for block in anthropic_request.system: + if block.type == "text" and block.text: + system_parts.append(block.text) + + if merge_inline_system: + for msg in anthropic_request.messages: + if msg.role != "system": + continue + text = _extract_system_text(msg.content) + if text: + system_parts.append(text) + + if system_parts: + openai_messages.append({"role": "system", "content": "\n".join(system_parts)}) + + def _emit_user_message(parts: list[dict]) -> None: + """Append accumulated parts as a user message, then clear them. + + Used to flush content collected BEFORE a tool_result so the + wire order stays user(pre) → tool → user(post). Without this + flush, text/image parts that appeared before a tool_result + block would be moved AFTER the tool message at end of loop. + """ + if not parts: + return + if len(parts) == 1 and parts[0]["type"] == "text": + openai_messages.append({"role": "user", "content": parts[0]["text"]}) + else: + openai_messages.append({"role": "user", "content": list(parts)}) + parts.clear() + + # Convert messages + for msg in anthropic_request.messages: + if msg.role == "system" and merge_inline_system: + continue + if isinstance(msg.content, str): + openai_messages.append({"role": msg.role, "content": msg.content}) + continue + + # Complex content with blocks + openai_msg = {"role": msg.role} + content_parts: list[dict] = [] + tool_calls: list[dict] = [] + + if msg.role == "assistant": + reasoning_content, reasoning_history = _convert_assistant_thinking_blocks( + msg.content + ) + if reasoning_content is not None: + openai_msg["reasoning_content"] = reasoning_content + if reasoning_history is not None: + content_parts.append({"type": "text", "text": reasoning_history}) + + for block in msg.content: + # ``thinking``/``redacted_thinking`` blocks are surfaced via + # the reasoning-history reconstruction above; skip them here + # to avoid double-injecting their text into the prompt. + if block.type in ("thinking", "redacted_thinking"): + continue + + # ``is not None`` (not truthy) so an empty-string text block + # still produces a placeholder text part — without it, an + # assistant turn whose only content is "" vanishes and + # subsequent user→user pairs trip strict chat templates. + if block.type == "text" and block.text is not None: + content_parts.append({"type": "text", "text": block.text}) + + elif block.type == "image" and block.source: + image_part = _convert_anthropic_image_source_to_openai_part( + block.source + ) + if image_part is not None: + content_parts.append(image_part) + + elif block.type == "search_result": + search_text = _text_from_search_result(block.model_dump()) + if search_text: + content_parts.append({"type": "text", "text": search_text}) + + elif block.type == "tool_use": + tool_call = { + "id": block.id or f"call_{uuid.uuid4().hex}", + "type": "function", + "function": { + "name": block.name or "", + "arguments": json.dumps(block.input or {}), + }, + } + tool_calls.append(tool_call) + + elif block.type == "tool_result": + tool_contents, tool_text = _convert_tool_result_content(block.content) + + # Use tool_use_id (per spec) with fallback to id + tool_call_id = block.tool_use_id or block.id or "" + + # Tool results from user become separate tool messages. + # Flush any pending text/image first so the wire order + # is preserved (a tool_result that arrived AFTER a text + # block must come AFTER that text in OpenAI form too). + if msg.role == "user": + _emit_user_message(content_parts) + for tool_content in tool_contents: + openai_messages.append( + { + "role": "tool", + "tool_call_id": tool_call_id, + "content": tool_content, + } + ) + else: + content_parts.append( + { + "type": "text", + "text": f"Tool result: {tool_text}", + } + ) + + # Attach tool calls to assistant messages + if tool_calls: + openai_msg["tool_calls"] = tool_calls + + # Attach content + if content_parts: + if len(content_parts) == 1 and content_parts[0]["type"] == "text": + openai_msg["content"] = content_parts[0]["text"] + else: + openai_msg["content"] = content_parts + openai_messages.append(openai_msg) + elif tool_calls: + openai_messages.append(openai_msg) + elif msg.role == "user": + # User turn that was entirely tool_results — the tool + # messages were already emitted above, nothing left. + continue + else: + # Assistant turn with no content and no tool_calls: emit + # an empty-string placeholder so strict templates still + # see a valid role-alternation sequence. + openai_msg["content"] = "" + openai_messages.append(openai_msg) + + # Build ChatCompletionRequest + request_data = { + "messages": openai_messages, + "model": anthropic_request.model, + "max_tokens": anthropic_request.max_tokens, + "stream": anthropic_request.stream or False, + } + + if anthropic_request.temperature is not None: + request_data["temperature"] = anthropic_request.temperature + if anthropic_request.top_p is not None: + request_data["top_p"] = anthropic_request.top_p + if anthropic_request.top_k is not None: + request_data["top_k"] = anthropic_request.top_k + if anthropic_request.stop_sequences is not None: + request_data["stop"] = anthropic_request.stop_sequences + + # Enable usage in stream so we can report it + if anthropic_request.stream: + request_data["stream_options"] = StreamOptions( + include_usage=True, + continuous_usage_stats=True, + ) + + chat_request = ChatCompletionRequest(**request_data) + + if anthropic_request.thinking is not None: + # The protocol layer already enforces SDK shape: + # enabled -> budget_tokens required (>=1024), display optional + # disabled -> neither budget_tokens nor display allowed + # adaptive -> budget_tokens forbidden, display optional + # So by the time we get here ``budget_tokens`` can only be + # set on ``enabled``. The local backend has no equivalent + # hard-cap knob, so we log a WARNING instead of rejecting — + # the Anthropic SDK would have accepted the request and we + # mirror that. Operators see the unenforced budget in logs. + if anthropic_request.thinking.budget_tokens is not None: + logger.warning( + "Anthropic thinking.budget_tokens=%d is accepted for " + "SDK compatibility but the local backend has no " + "equivalent hard-cap knob — the budget is not enforced", + anthropic_request.thinking.budget_tokens, + ) + # Claude 4.7's ``adaptive`` is treated identically to ``enabled`` + # because the local backend has no auto-throttle equivalent. + # Anything other than ``disabled`` enables reasoning. + enabled = anthropic_request.thinking.type != "disabled" + if anthropic_request.thinking.display == "omitted": + # Anthropic 4.7 spec: keep reasoning ON but hide reasoning + # text from the client. The OpenAI streaming pipeline has + # no equivalent suppress knob — log so operators can see + # the request, then proceed with normal reasoning emission. + logger.warning( + "Anthropic thinking.display='omitted' is accepted for " + "SDK compatibility but reasoning text will still be " + "emitted to the client" + ) + if apply_reasoning_enabled is None: + raise ValueError( + "Cannot convert Anthropic thinking without a reasoning adapter" + ) + apply_reasoning_enabled(chat_request, enabled) + + # Claude 4.7 ``output_config``: map ``effort`` onto the OpenAI + # ``reasoning_effort`` knob. ``xhigh`` collapses to ``max`` because + # the OpenAI Literal does not include the Anthropic-only ``xhigh``. + # ``task_budget`` is a soft hint forwarded as a custom param so the + # model can see it without it becoming a hard cap (``max_tokens`` + # is still the hard cap). + if anthropic_request.output_config is not None: + oc = anthropic_request.output_config + if oc.effort is not None: + chat_request.reasoning_effort = "max" if oc.effort == "xhigh" else oc.effort + if oc.task_budget is not None: + # Custom params are silently ignored by backends that + # don't recognise them; logging it makes the propagation + # visible. + logger.info( + "Anthropic output_config.task_budget hint: %d %s", + oc.task_budget.total, + oc.task_budget.type, + ) + + # ``betas`` is the Anthropic SDK's opt-in feature list (e.g. + # ``["thinking-2025-08-04"]``). The local backend has no + # equivalent beta system; accept-and-log so requests don't 400. + if anthropic_request.betas: + logger.info( + "Anthropic request opted into betas %s — no-op locally", + anthropic_request.betas, + ) + + # Convert tools. Deferred tools stay in the list with defer_loading=True; + # the chat template hides them from the initial block and renders + # them on demand when a tool_reference block names them. + if anthropic_request.tools: + converted_tools = [] + for tool in anthropic_request.tools: + if is_server_tool(tool): + # Anthropic server-side tools (web_search_*, computer_*, + # bash_*, text_editor_*) have no client-side input_schema + # because Anthropic provides the implementation. We can't + # forward them to the OpenAI tools array (which requires a + # schema), so skip with a visible log. + logger.info( + "Skipping built-in Anthropic server tool %r (type=%r): " + "no native support in the OpenAI-compatible backend", + tool.name, + tool.type, + ) + continue + + # Custom tools always have a validated input_schema + # (enforced at Pydantic parse time). + converted_tools.append( + Tool( + type="function", + defer_loading=tool.defer_loading, + function={ + "name": tool.name, + "description": tool.description or "", + "parameters": tool.input_schema, + }, + ) + ) + + if converted_tools: + chat_request.tools = converted_tools + + # Convert tool choice. ``any``/``tool`` express a hard requirement + # ("the model MUST call a tool"); if every requested tool was a + # server-side Anthropic built-in that we just skipped, there is + # no tool the model could call. Silently downgrading to "no tool" + # would deceive the caller, so raise an explicit 400. + if anthropic_request.tool_choice is not None: + tc_type = anthropic_request.tool_choice.type + if tc_type == "none": + chat_request.tool_choice = "none" + elif chat_request.tools: + if tc_type == "auto": + chat_request.tool_choice = "auto" + elif tc_type == "any": + chat_request.tool_choice = "required" + elif tc_type == "tool": + tool_name = anthropic_request.tool_choice.name + # ``Tool.function`` is a ``Function`` Pydantic model, not + # a dict — access by attribute. A dict ``.get`` would + # AttributeError and surface as a 500 instead of the + # intended 400 / happy path. + if not any(t.function.name == tool_name for t in chat_request.tools): + raise ValueError( + f"tool_choice references tool {tool_name!r} but it " + f"is not in the forwarded tools list " + f"(server-side Anthropic tools cannot be selected)" + ) + chat_request.tool_choice = ToolChoice( + type="function", + function=ToolChoiceFuncName(name=tool_name), + ) + elif tc_type in ("any", "tool"): + raise ValueError( + f"tool_choice={tc_type!r} requires at least one custom " + f"tool; all supplied tools were server-side Anthropic " + f"built-ins which the OpenAI-compatible backend cannot " + f"invoke" + ) + elif chat_request.tools: + chat_request.tool_choice = "auto" + + return chat_request + + +def convert_response( + response: ChatCompletionResponse, +) -> AnthropicMessagesResponse: + """Convert an OpenAI ChatCompletionResponse to an Anthropic Messages response.""" + if not response.choices: + return AnthropicMessagesResponse( + content=[TextBlock(text="")], + model=response.model, + stop_reason="end_turn", + usage=AnthropicUsage(input_tokens=0, output_tokens=0), + ) + + choice = response.choices[0] + content: list[AnthropicContentBlock] = [] + + # Add reasoning content as a thinking block. signature is omitted + # entirely when the backend doesn't provide one — empty strings + # would fail downstream Anthropic signature verifiers. + if choice.message.reasoning_content: + content.append(ThinkingBlock(thinking=choice.message.reasoning_content)) + + # Add text content + if choice.message.content: + content.append(TextBlock(text=choice.message.content)) + + # Add tool calls + if choice.message.tool_calls: + for tool_call in choice.message.tool_calls: + raw_args = tool_call.function.arguments + try: + tool_input = json.loads(raw_args) + except (json.JSONDecodeError, TypeError): + # Surface invalid tool arguments so an empty-dict + # tool call is never indistinguishable from a real + # one when something downstream goes wrong. + logger.warning( + "Tool %r emitted invalid JSON arguments: %r — " + "defaulting to empty input", + tool_call.function.name, + (raw_args or "")[:200], + ) + tool_input = {} + + content.append( + ToolUseBlock( + id=tool_call.id, + name=tool_call.function.name, + input=tool_input, + ) + ) + + # Map stop reason + finish_reason = choice.finish_reason or "stop" + if finish_reason not in STOP_REASON_MAP: + logger.warning( + "Unmapped OpenAI finish_reason %r; defaulting to end_turn", + finish_reason, + ) + stop_reason = STOP_REASON_MAP.get(finish_reason, "end_turn") + + # Anthropic requires ``content`` to contain at least one block. + # Empty string completions (max_tokens=1 stop, content filter, etc.) + # would otherwise ship ``content=[]`` and break strict SDK parsers. + if not content: + content.append(TextBlock(text="")) + + return AnthropicMessagesResponse( + id=f"msg_{uuid.uuid4().hex}", + content=content, + model=response.model, + stop_reason=stop_reason, + usage=_anthropic_usage_from_openai( + response.usage, + include_input=True, + include_output=True, + ), + ) + + class AnthropicServing: """Handler for Anthropic Messages API requests. @@ -229,508 +820,17 @@ async def handle_messages( def _convert_to_chat_completion_request( self, anthropic_request: AnthropicMessagesRequest ) -> ChatCompletionRequest: - """Convert an Anthropic Messages request to an OpenAI ChatCompletion request.""" - openai_messages = [] - - def _convert_anthropic_image_source_to_openai_part( - source: Any, - ) -> Optional[dict]: - # Source may arrive as a Pydantic model (typed ImageBlock.source) - # or as a raw dict when parsed from a nested tool_result payload. - if isinstance(source, BaseModel): - source = source.model_dump(exclude_none=True) - if not isinstance(source, dict): - return None - - source_type = source.get("type") - if source_type == "base64": - media_type = source.get("media_type", "image/png") - data = source.get("data", "") - if not data: - return None - return { - "type": "image_url", - "image_url": { - "url": f"data:{media_type};base64,{data}", - }, - } - - url = source.get("url") - if url: - return { - "type": "image_url", - "image_url": { - "url": url, - }, - } - - return None - - def _text_from_search_result(item: dict[str, Any]) -> str: - search_parts = [] - title = item.get("title") - if title: - search_parts.append(f"Title: {title}") - - source = item.get("source") - if isinstance(source, dict): - source_text = source.get("url") or source.get("text") - if source_text: - search_parts.append(f"Source: {source_text}") - elif source: - search_parts.append(f"Source: {source}") - - content = item.get("content") - content_parts = [] - if isinstance(content, str): - content_parts.append(content) - elif isinstance(content, list): - for part in content: - if not isinstance(part, dict): - continue - if part.get("type") == "text" and part.get("text"): - content_parts.append(part["text"]) - if content_parts: - search_parts.append("Content: " + "\n".join(content_parts)) - - return "\n".join(search_parts) - - def _convert_tool_result_content( - content: Any, - ) -> tuple[list[Union[str, list[dict]]], str]: - if isinstance(content, list): - tool_content_parts = [] - tool_text_parts = [] - - for raw_item in content: - # Items may be typed Pydantic blocks (after request - # validation) or raw dicts (from legacy callers). Coerce - # to dict so the existing key-based logic still works. - if isinstance(raw_item, BaseModel): - item = raw_item.model_dump(exclude_none=True) - elif isinstance(raw_item, dict): - item = raw_item - else: - continue - - item_type = item.get("type") - if item_type == "text": - text = item.get("text", "") - if text: - tool_text_parts.append(text) - tool_content_parts.append({"type": "text", "text": text}) - elif item_type == "image": - image_part = _convert_anthropic_image_source_to_openai_part( - item.get("source") - ) - if image_part is not None: - tool_content_parts.append(image_part) - elif item_type == "tool_reference": - # Anthropic uses `tool_name`; the SGLang chat template - # matches on `name`. Translate at the boundary. - ref_name = item.get("tool_name") or item.get("name") - if ref_name: - tool_content_parts.append( - {"type": "tool_reference", "name": ref_name} - ) - elif item_type == "search_result": - search_text = _text_from_search_result(item) - if search_text: - tool_text_parts.append(search_text) - tool_content_parts.append( - {"type": "text", "text": search_text} - ) - - tool_text = "\n".join(tool_text_parts) - # GLM templates expand references only at the start of a tool - # message, so isolate reference runs without changing part order. - tool_content_groups: list[list[dict]] = [] - for part in tool_content_parts: - is_reference = part["type"] == "tool_reference" - if ( - not tool_content_groups - or (tool_content_groups[-1][0]["type"] == "tool_reference") - != is_reference - ): - tool_content_groups.append([]) - tool_content_groups[-1].append(part) - - tool_contents: list[Union[str, list[dict]]] = [] - for group in tool_content_groups: - if len(group) == 1 and group[0]["type"] == "text": - tool_contents.append(group[0]["text"]) - else: - tool_contents.append(group) - return tool_contents or [""], tool_text - - tool_text = str(content) if content else "" - return [tool_text], tool_text - - def _convert_assistant_thinking_blocks( - blocks: list[AnthropicContentBlock], - ) -> tuple[Optional[str], Optional[str]]: - """Reconstruct prior-turn thinking as ``(reasoning_content, text)``. - - At most one is set: encoders that frame the reasoning channel take - it as ``reasoning_content``, everything else gets it re-wrapped and - spliced into content. - - ``redacted_thinking`` carries encrypted bytes that no local - parser can interpret, so we raise rather than silently drop it. - On non-reasoning models (no detector configured) the rewrap is - best-effort: we log a warning and drop the thinking text so a - history echo doesn't 400 the whole request — the prior thinking - is opaque context the model didn't need anyway. - """ - if any(block.type == "redacted_thinking" for block in blocks): - raise ValueError("Anthropic redacted_thinking history is not supported") - - thinking_parts = [ - block.thinking - for block in blocks - if block.type == "thinking" and block.thinking - ] - if not thinking_parts: - return None, None - - reasoning_text = "\n".join(thinking_parts) - if self.openai_serving_chat.supports_native_reasoning_history(): - return reasoning_text, None - - try: - return None, self.openai_serving_chat.wrap_reasoning_history( - reasoning_text - ) - except ValueError as e: - logger.warning( - "Dropping prior-turn thinking history (%d blocks): %s", - len(thinking_parts), - e, - ) - return None, None - - system_parts: list[str] = [] - if anthropic_request.system: - if isinstance(anthropic_request.system, str): - if anthropic_request.system.strip(): - system_parts.append(anthropic_request.system) - else: - for block in anthropic_request.system: - if block.type == "text" and block.text: - system_parts.append(block.text) - - if self._merge_inline_system: - for msg in anthropic_request.messages: - if msg.role != "system": - continue - text = _extract_system_text(msg.content) - if text: - system_parts.append(text) - - if system_parts: - openai_messages.append( - {"role": "system", "content": "\n".join(system_parts)} - ) - - def _emit_user_message(parts: list[dict]) -> None: - """Append accumulated parts as a user message, then clear them. - - Used to flush content collected BEFORE a tool_result so the - wire order stays user(pre) → tool → user(post). Without this - flush, text/image parts that appeared before a tool_result - block would be moved AFTER the tool message at end of loop. - """ - if not parts: - return - if len(parts) == 1 and parts[0]["type"] == "text": - openai_messages.append({"role": "user", "content": parts[0]["text"]}) - else: - openai_messages.append({"role": "user", "content": list(parts)}) - parts.clear() - - # Convert messages - for msg in anthropic_request.messages: - if msg.role == "system" and self._merge_inline_system: - continue - if isinstance(msg.content, str): - openai_messages.append({"role": msg.role, "content": msg.content}) - continue - - # Complex content with blocks - openai_msg = {"role": msg.role} - content_parts: list[dict] = [] - tool_calls: list[dict] = [] - - if msg.role == "assistant": - reasoning_content, reasoning_history = ( - _convert_assistant_thinking_blocks(msg.content) - ) - if reasoning_content is not None: - openai_msg["reasoning_content"] = reasoning_content - if reasoning_history is not None: - content_parts.append({"type": "text", "text": reasoning_history}) - - for block in msg.content: - # ``thinking``/``redacted_thinking`` blocks are surfaced via - # the reasoning-history reconstruction above; skip them here - # to avoid double-injecting their text into the prompt. - if block.type in ("thinking", "redacted_thinking"): - continue - - # ``is not None`` (not truthy) so an empty-string text block - # still produces a placeholder text part — without it, an - # assistant turn whose only content is "" vanishes and - # subsequent user→user pairs trip strict chat templates. - if block.type == "text" and block.text is not None: - content_parts.append({"type": "text", "text": block.text}) - - elif block.type == "image" and block.source: - image_part = _convert_anthropic_image_source_to_openai_part( - block.source - ) - if image_part is not None: - content_parts.append(image_part) - - elif block.type == "search_result": - search_text = _text_from_search_result(block.model_dump()) - if search_text: - content_parts.append({"type": "text", "text": search_text}) - - elif block.type == "tool_use": - tool_call = { - "id": block.id or f"call_{uuid.uuid4().hex}", - "type": "function", - "function": { - "name": block.name or "", - "arguments": json.dumps(block.input or {}), - }, - } - tool_calls.append(tool_call) - - elif block.type == "tool_result": - tool_contents, tool_text = _convert_tool_result_content( - block.content - ) - - # Use tool_use_id (per spec) with fallback to id - tool_call_id = block.tool_use_id or block.id or "" - - # Tool results from user become separate tool messages. - # Flush any pending text/image first so the wire order - # is preserved (a tool_result that arrived AFTER a text - # block must come AFTER that text in OpenAI form too). - if msg.role == "user": - _emit_user_message(content_parts) - for tool_content in tool_contents: - openai_messages.append( - { - "role": "tool", - "tool_call_id": tool_call_id, - "content": tool_content, - } - ) - else: - content_parts.append( - { - "type": "text", - "text": f"Tool result: {tool_text}", - } - ) - - # Attach tool calls to assistant messages - if tool_calls: - openai_msg["tool_calls"] = tool_calls - - # Attach content - if content_parts: - if len(content_parts) == 1 and content_parts[0]["type"] == "text": - openai_msg["content"] = content_parts[0]["text"] - else: - openai_msg["content"] = content_parts - openai_messages.append(openai_msg) - elif tool_calls: - openai_messages.append(openai_msg) - elif msg.role == "user": - # User turn that was entirely tool_results — the tool - # messages were already emitted above, nothing left. - continue - else: - # Assistant turn with no content and no tool_calls: emit - # an empty-string placeholder so strict templates still - # see a valid role-alternation sequence. - openai_msg["content"] = "" - openai_messages.append(openai_msg) - - # Build ChatCompletionRequest - request_data = { - "messages": openai_messages, - "model": anthropic_request.model, - "max_tokens": anthropic_request.max_tokens, - "stream": anthropic_request.stream or False, - } - - if anthropic_request.temperature is not None: - request_data["temperature"] = anthropic_request.temperature - if anthropic_request.top_p is not None: - request_data["top_p"] = anthropic_request.top_p - if anthropic_request.top_k is not None: - request_data["top_k"] = anthropic_request.top_k - if anthropic_request.stop_sequences is not None: - request_data["stop"] = anthropic_request.stop_sequences - - # Enable usage in stream so we can report it - if anthropic_request.stream: - request_data["stream_options"] = StreamOptions( - include_usage=True, - continuous_usage_stats=True, - ) - - chat_request = ChatCompletionRequest(**request_data) - - if anthropic_request.thinking is not None: - # The protocol layer already enforces SDK shape: - # enabled -> budget_tokens required (>=1024), display optional - # disabled -> neither budget_tokens nor display allowed - # adaptive -> budget_tokens forbidden, display optional - # So by the time we get here ``budget_tokens`` can only be - # set on ``enabled``. The local backend has no equivalent - # hard-cap knob, so we log a WARNING instead of rejecting — - # the Anthropic SDK would have accepted the request and we - # mirror that. Operators see the unenforced budget in logs. - if anthropic_request.thinking.budget_tokens is not None: - logger.warning( - "Anthropic thinking.budget_tokens=%d is accepted for " - "SDK compatibility but the local backend has no " - "equivalent hard-cap knob — the budget is not enforced", - anthropic_request.thinking.budget_tokens, - ) - # Claude 4.7's ``adaptive`` is treated identically to ``enabled`` - # because the local backend has no auto-throttle equivalent. - # Anything other than ``disabled`` enables reasoning. - enabled = anthropic_request.thinking.type != "disabled" - if anthropic_request.thinking.display == "omitted": - # Anthropic 4.7 spec: keep reasoning ON but hide reasoning - # text from the client. The OpenAI streaming pipeline has - # no equivalent suppress knob — log so operators can see - # the request, then proceed with normal reasoning emission. - logger.warning( - "Anthropic thinking.display='omitted' is accepted for " - "SDK compatibility but reasoning text will still be " - "emitted to the client" - ) - self.openai_serving_chat.apply_reasoning_enabled(chat_request, enabled) - - # Claude 4.7 ``output_config``: map ``effort`` onto the OpenAI - # ``reasoning_effort`` knob. ``xhigh`` collapses to ``max`` because - # the OpenAI Literal does not include the Anthropic-only ``xhigh``. - # ``task_budget`` is a soft hint forwarded as a custom param so the - # model can see it without it becoming a hard cap (``max_tokens`` - # is still the hard cap). - if anthropic_request.output_config is not None: - oc = anthropic_request.output_config - if oc.effort is not None: - chat_request.reasoning_effort = ( - "max" if oc.effort == "xhigh" else oc.effort - ) - if oc.task_budget is not None: - # Custom params are silently ignored by backends that - # don't recognise them; logging it makes the propagation - # visible. - logger.info( - "Anthropic output_config.task_budget hint: %d %s", - oc.task_budget.total, - oc.task_budget.type, - ) - - # ``betas`` is the Anthropic SDK's opt-in feature list (e.g. - # ``["thinking-2025-08-04"]``). The local backend has no - # equivalent beta system; accept-and-log so requests don't 400. - if anthropic_request.betas: - logger.info( - "Anthropic request opted into betas %s — no-op locally", - anthropic_request.betas, - ) - - # Convert tools. Deferred tools stay in the list with defer_loading=True; - # the chat template hides them from the initial block and renders - # them on demand when a tool_reference block names them. - if anthropic_request.tools: - converted_tools = [] - for tool in anthropic_request.tools: - if is_server_tool(tool): - # Anthropic server-side tools (web_search_*, computer_*, - # bash_*, text_editor_*) have no client-side input_schema - # because Anthropic provides the implementation. We can't - # forward them to the OpenAI tools array (which requires a - # schema), so skip with a visible log. - logger.info( - "Skipping built-in Anthropic server tool %r (type=%r): " - "no native support in the OpenAI-compatible backend", - tool.name, - tool.type, - ) - continue - - # Custom tools always have a validated input_schema - # (enforced at Pydantic parse time). - converted_tools.append( - Tool( - type="function", - defer_loading=tool.defer_loading, - function={ - "name": tool.name, - "description": tool.description or "", - "parameters": tool.input_schema, - }, - ) - ) - - if converted_tools: - chat_request.tools = converted_tools - - # Convert tool choice. ``any``/``tool`` express a hard requirement - # ("the model MUST call a tool"); if every requested tool was a - # server-side Anthropic built-in that we just skipped, there is - # no tool the model could call. Silently downgrading to "no tool" - # would deceive the caller, so raise an explicit 400. - if anthropic_request.tool_choice is not None: - tc_type = anthropic_request.tool_choice.type - if tc_type == "none": - chat_request.tool_choice = "none" - elif chat_request.tools: - if tc_type == "auto": - chat_request.tool_choice = "auto" - elif tc_type == "any": - chat_request.tool_choice = "required" - elif tc_type == "tool": - tool_name = anthropic_request.tool_choice.name - # ``Tool.function`` is a ``Function`` Pydantic model, not - # a dict — access by attribute. A dict ``.get`` would - # AttributeError and surface as a 500 instead of the - # intended 400 / happy path. - if not any( - t.function.name == tool_name for t in chat_request.tools - ): - raise ValueError( - f"tool_choice references tool {tool_name!r} but it " - f"is not in the forwarded tools list " - f"(server-side Anthropic tools cannot be selected)" - ) - chat_request.tool_choice = ToolChoice( - type="function", - function=ToolChoiceFuncName(name=tool_name), - ) - elif tc_type in ("any", "tool"): - raise ValueError( - f"tool_choice={tc_type!r} requires at least one custom " - f"tool; all supplied tools were server-side Anthropic " - f"built-ins which the OpenAI-compatible backend cannot " - f"invoke" - ) - elif chat_request.tools: - chat_request.tool_choice = "auto" - - return chat_request + return convert_to_chat_completion_request( + anthropic_request, + merge_inline_system=self._merge_inline_system, + wrap_reasoning_history=lambda reasoning_text: ( + self.openai_serving_chat.wrap_reasoning_history(reasoning_text) + ), + supports_native_reasoning_history=self.openai_serving_chat.supports_native_reasoning_history, + apply_reasoning_enabled=lambda request, enabled: ( + self.openai_serving_chat.apply_reasoning_enabled(request, enabled) + ), + ) async def _handle_non_streaming( self, @@ -1258,80 +1358,7 @@ def _parse_upstream_error(data_str: str) -> Optional[tuple[str, str]]: def _convert_response( self, response: ChatCompletionResponse ) -> AnthropicMessagesResponse: - """Convert an OpenAI ChatCompletionResponse to an Anthropic Messages response.""" - if not response.choices: - return AnthropicMessagesResponse( - content=[TextBlock(text="")], - model=response.model, - stop_reason="end_turn", - usage=AnthropicUsage(input_tokens=0, output_tokens=0), - ) - - choice = response.choices[0] - content: list[AnthropicContentBlock] = [] - - # Add reasoning content as a thinking block. signature is omitted - # entirely when the backend doesn't provide one — empty strings - # would fail downstream Anthropic signature verifiers. - if choice.message.reasoning_content: - content.append(ThinkingBlock(thinking=choice.message.reasoning_content)) - - # Add text content - if choice.message.content: - content.append(TextBlock(text=choice.message.content)) - - # Add tool calls - if choice.message.tool_calls: - for tool_call in choice.message.tool_calls: - raw_args = tool_call.function.arguments - try: - tool_input = json.loads(raw_args) - except (json.JSONDecodeError, TypeError): - # Surface invalid tool arguments so an empty-dict - # tool call is never indistinguishable from a real - # one when something downstream goes wrong. - logger.warning( - "Tool %r emitted invalid JSON arguments: %r — " - "defaulting to empty input", - tool_call.function.name, - (raw_args or "")[:200], - ) - tool_input = {} - - content.append( - ToolUseBlock( - id=tool_call.id, - name=tool_call.function.name, - input=tool_input, - ) - ) - - # Map stop reason - finish_reason = choice.finish_reason or "stop" - if finish_reason not in STOP_REASON_MAP: - logger.warning( - "Unmapped OpenAI finish_reason %r; defaulting to end_turn", - finish_reason, - ) - stop_reason = STOP_REASON_MAP.get(finish_reason, "end_turn") - - # Anthropic requires ``content`` to contain at least one block. - # Empty string completions (max_tokens=1 stop, content filter, etc.) - # would otherwise ship ``content=[]`` and break strict SDK parsers. - if not content: - content.append(TextBlock(text="")) - - return AnthropicMessagesResponse( - id=f"msg_{uuid.uuid4().hex}", - content=content, - model=response.model, - stop_reason=stop_reason, - usage=_anthropic_usage_from_openai( - response.usage, - include_input=True, - include_output=True, - ), - ) + return convert_response(response) def _convert_openai_error_response(self, response) -> JSONResponse: """Forward an upstream OpenAI-handler error as an Anthropic error. diff --git a/python/sglang/srt/entrypoints/anthropic/utils.py b/python/sglang/srt/entrypoints/anthropic/utils.py new file mode 100644 index 000000000000..10cd9726f79a --- /dev/null +++ b/python/sglang/srt/entrypoints/anthropic/utils.py @@ -0,0 +1,219 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Anthropic error conversion and fake-SSE response helpers. + +Request and response format conversion lives in serving.py as public, +runtime-independent functions. This module contains the two adaptations used +by external frontends: an Anthropic error DTO and an eager event sequence for +an already-complete OpenAI response. +""" + +from __future__ import annotations + +import json +import logging +import uuid +from typing import Callable, Optional + +from sglang.srt.entrypoints.anthropic.protocol import ( + AnthropicError, + AnthropicErrorResponse, + AnthropicMessageEndDelta, + AnthropicMessagesResponse, + AnthropicStreamEvent, + ContentBlockDeltaEvent, + ContentBlockStartEvent, + ContentBlockStopEvent, + InputJsonDelta, + MessageDeltaEvent, + MessageStartEvent, + MessageStopEvent, + TextBlock, + TextDelta, + ThinkingBlock, + ThinkingDelta, + ToolUseBlock, +) +from sglang.srt.entrypoints.anthropic.serving import ( + ERROR_TYPE_MAP as _SERVING_ERROR_TYPE_MAP, +) +from sglang.srt.entrypoints.anthropic.serving import ( + STOP_REASON_MAP, + _anthropic_usage_from_openai, + _scrub_error_message, +) +from sglang.srt.entrypoints.openai.protocol import ChatCompletionResponse + +logger = logging.getLogger(__name__) + +# ``serving.py``'s map plus the two statuses the HTTP layer maps in its +# ``/v1/messages`` exception handler (``http_server.py``): FastAPI raises +# 413/422 before ``AnthropicServing`` runs, so ``serving.py`` never needed +# them. Unlisted statuses fall through to ``api_error`` in both sources. +ERROR_TYPE_MAP = { + **_SERVING_ERROR_TYPE_MAP, + 413: "request_too_large", + 422: "invalid_request_error", +} + + +def _anthropic_message_id() -> str: + """Message-ID factory with the same format the server generates.""" + return f"msg_{uuid.uuid4().hex}" + + +def to_anthropic_error(status_code: int, body: bytes) -> AnthropicErrorResponse: + """Map an upstream error status and body to an Anthropic error envelope. + + Same payload parsing and scrub policy as ``serving.py``'s + ``_convert_openai_error_response``, but against the composite + ``ERROR_TYPE_MAP`` above and returning the envelope DTO — HTTP response + construction stays with the caller. 4xx keeps a sanitized upstream + message and honors an upstream ``error.type``; 5xx is always the + generic scrubbed message. + """ + body = body or b"" + error_type = ERROR_TYPE_MAP.get(status_code, "api_error") + + upstream_message: Optional[str] = None + try: + payload = json.loads(body.decode("utf-8")) if body else None + except (json.JSONDecodeError, UnicodeDecodeError): + # Non-JSON body (HTML gateway error, plain text, ...). Use a bounded + # slice of the raw body so the client still has a useful hint. + upstream_message = body.decode("utf-8", errors="replace")[:500] + else: + if isinstance(payload, dict): + error_payload = payload.get("error", payload) + if isinstance(error_payload, dict): + upstream_message = error_payload.get("message") or payload.get( + "message" + ) + # Honor the upstream error.type only for 4xx; 5xx is + # normalized by the scrub below. + if status_code < 500: + upstream_type = error_payload.get("type") + if isinstance(upstream_type, str) and upstream_type: + error_type = upstream_type + elif isinstance(error_payload, str): + upstream_message = error_payload + elif isinstance(payload.get("message"), str): + upstream_message = payload["message"] + + message = _scrub_error_message(upstream_message or "", status_code) + return AnthropicErrorResponse( + error=AnthropicError(type=error_type, message=message) + ) + + +def to_anthropic_fake_sse_events( + response: ChatCompletionResponse, + *, + model: str, + id_factory: Callable[[], str] = _anthropic_message_id, +) -> tuple[AnthropicStreamEvent, ...]: + """Eagerly synthesize the Anthropic SSE event sequence from one complete + OpenAI response. + + Event schema, block ordering (thinking → text → tool_use), index + accounting, usage split (input on ``message_start``, output on + ``message_delta``) and stop-reason mapping follow the server's live + stream; collapsing each block into a single delta is the documented + difference. ``model`` must be the original Anthropic request model, not + the backend's possibly-aliased response model. + """ + events: list[AnthropicStreamEvent] = [ + MessageStartEvent( + message=AnthropicMessagesResponse( + id=id_factory(), + content=[], + model=model, + usage=_anthropic_usage_from_openai( + response.usage, + include_input=True, + include_output=True, + force_zero_output=True, + ), + ) + ) + ] + + message = response.choices[0].message if response.choices else None + index = 0 + + if message is not None and message.reasoning_content: + events.append( + ContentBlockStartEvent( + index=index, content_block=ThinkingBlock(thinking="") + ) + ) + events.append( + ContentBlockDeltaEvent( + index=index, delta=ThinkingDelta(thinking=message.reasoning_content) + ) + ) + events.append(ContentBlockStopEvent(index=index)) + index += 1 + + if message is not None and message.content: + events.append( + ContentBlockStartEvent(index=index, content_block=TextBlock(text="")) + ) + events.append( + ContentBlockDeltaEvent(index=index, delta=TextDelta(text=message.content)) + ) + events.append(ContentBlockStopEvent(index=index)) + index += 1 + + if message is not None and message.tool_calls: + for tool_call in message.tool_calls: + events.append( + ContentBlockStartEvent( + index=index, + content_block=ToolUseBlock( + id=tool_call.id or f"toolu_{uuid.uuid4().hex}", + name=tool_call.function.name, + input={}, + ), + ) + ) + if tool_call.function.arguments: + events.append( + ContentBlockDeltaEvent( + index=index, + delta=InputJsonDelta(partial_json=tool_call.function.arguments), + ) + ) + events.append(ContentBlockStopEvent(index=index)) + index += 1 + + finish_reason = ( + (response.choices[0].finish_reason or "stop") if response.choices else "stop" + ) + if finish_reason not in STOP_REASON_MAP: + logger.warning( + "Unmapped OpenAI finish_reason %r; defaulting to end_turn", finish_reason + ) + events.append( + MessageDeltaEvent( + delta=AnthropicMessageEndDelta( + stop_reason=STOP_REASON_MAP.get(finish_reason, "end_turn") + ), + usage=_anthropic_usage_from_openai( + response.usage, include_input=False, include_output=True + ), + ) + ) + events.append(MessageStopEvent()) + return tuple(events) diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index 9afd8c480e74..3ce2ce1873fd 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -67,18 +67,20 @@ ) from sglang.srt.managers.detokenizer_manager import run_detokenizer_process from sglang.srt.managers.io_struct import ( + BeginWeightUpdateReqInput, CloseSessionReqInput, DestroyWeightsUpdateGroupReqInput, EmbeddingReqInput, + EndWeightUpdateReqInput, GenerateReqInput, GetWeightsByNameReqInput, InitWeightsUpdateGroupReqInput, - LoadLoRAAdapterFromTensorsReqInput, LoadLoRAAdapterReqInput, MultimodalDataInputFormat, OpenSessionReqInput, ProfileReq, ProfileReqType, + RegisterLoRAAdapterReqInput, ReleaseMemoryOccupationReqInput, ResumeMemoryOccupationReqInput, ReturnHiddenStatesMode, @@ -240,10 +242,6 @@ class Engine(EngineScoreMixin, EngineBase): run_scheduler_process_func: Callable = staticmethod(run_scheduler_process) run_detokenizer_process_func: Callable = staticmethod(run_detokenizer_process) - # Backend-specific launch handle: the Ray engine schedules its actors onto a - # placement group. Not config — a live cluster object. - _placement_group = None - def __init__(self, **kwargs): """ The arguments of this function is the same as `sglang/srt/server_args.py::ServerArgs`. @@ -303,7 +301,6 @@ def __init__(self, **kwargs): init_tokenizer_manager_func=self.init_tokenizer_manager_func, run_scheduler_process_func=self.run_scheduler_process_func, run_detokenizer_process_func=self.run_detokenizer_process_func, - placement_group=self._placement_group, ) self.tokenizer_manager = tokenizer_manager self.template_manager = template_manager @@ -849,8 +846,6 @@ def _launch_scheduler_processes( server_args: ServerArgs, port_args: PortArgs, run_scheduler_process_func: Callable, - *, - placement_group=None, ) -> Tuple[SchedulerInitResult, Optional[List]]: """Launch scheduler processes using multiprocessing. Override in subclasses for different backends (e.g. Ray). @@ -1055,7 +1050,6 @@ def _launch_subprocesses( run_scheduler_process_func: Callable, run_detokenizer_process_func: Callable, port_args: Optional[PortArgs] = None, - placement_group=None, ) -> Tuple[ TokenizerManager, TemplateManager, @@ -1111,7 +1105,7 @@ def _launch_subprocesses( # Start the engine info bootstrap server if per-rank info is needed. engine_info_bootstrap_server = None if ( - get_model().remote_instance_weight_loader_start_seed_via_transfer_engine + server_args.needs_engine_info_bootstrap() and get_parallel().node_rank == 0 ): bootstrap_port = get_model().engine_info_bootstrap_port @@ -1137,13 +1131,8 @@ def _launch_subprocesses( raise # Launch scheduler processes - # Passed only when there is one: this hook is an override point, and a - # subclass written against the three-argument signature must keep working. - launch_kwargs = ( - {"placement_group": placement_group} if placement_group is not None else {} - ) scheduler_init_result, scheduler_procs = cls._launch_scheduler_processes( - server_args, port_args, run_scheduler_process_func, **launch_kwargs + server_args, port_args, run_scheduler_process_func ) scheduler_init_result.engine_info_bootstrap_server = ( engine_info_bootstrap_server @@ -1443,6 +1432,39 @@ def destroy_weights_update_group( self.tokenizer_manager.destroy_weights_update_group(obj, None) ) + def begin_weight_update(self, selector: str = "all", sync_base: bool = True): + """Open a weight-update session: unpack in-place-quantized weights on the + selected runners so update_weights_from_{distributed,tensor} can load into + them. sync_base=False declares an adapter-only session (no unpack; base + tensors rejected). Must be closed with end_weight_update().""" + obj = BeginWeightUpdateReqInput(selector=selector, sync_base=sync_base) + return self.loop.run_until_complete( + self.tokenizer_manager.begin_weight_update(obj, None) + ) + + def end_weight_update(self, expected_lora_checksums: Optional[Dict] = None): + """Close the session opened by begin_weight_update(): finalize quantized + weights into kernel layout (sync_base sessions) and apply the streamed + LoRA stash (optionally checksum-verified).""" + obj = EndWeightUpdateReqInput(expected_lora_checksums=expected_lora_checksums) + return self.loop.run_until_complete( + self.tokenizer_manager.end_weight_update(obj, None) + ) + + def register_lora_adapter( + self, lora_name: str, config_dict: Dict, pinned: bool = False + ): + """Create-or-refresh a LoRA adapter's identity and config. Weights are + untouched by the caller (new adapters start zeroed); the bytes arrive as + '{lora_name}:{hf_key}'-prefixed tensors in the update_weights_from_* + stream and are applied at end_weight_update.""" + obj = RegisterLoRAAdapterReqInput( + lora_name=lora_name, config_dict=config_dict, pinned=pinned + ) + return self.loop.run_until_complete( + self.tokenizer_manager.register_lora_adapter(obj, None) + ) + def update_weights_from_distributed( self, names: list[str], @@ -1547,26 +1569,6 @@ def _serialize_tensors_per_rank( for _ in range(get_parallel().tp_size) ] - def load_lora_adapter_from_tensors( - self, - lora_name: str, - tensors: Union[Dict[str, torch.Tensor], List[SerializedTensorPayload]], - config_dict: Dict, - load_format: Optional[str] = None, - ): - serialized_named_tensors = self._serialize_tensors_per_rank( - tensors, load_format - ) - lora_req = LoadLoRAAdapterFromTensorsReqInput( - lora_name=lora_name, - config_dict=config_dict, - serialized_named_tensors=serialized_named_tensors, - load_format=load_format, - ) - return self.loop.run_until_complete( - self.tokenizer_manager.load_lora_adapter_from_tensors(lora_req, None) - ) - def load_lora_adapter(self, lora_name: str, lora_path: str, pinned: bool = False): """Load a new LoRA adapter without re-launching the engine.""" diff --git a/python/sglang/srt/entrypoints/engine_info_bootstrap_server.py b/python/sglang/srt/entrypoints/engine_info_bootstrap_server.py index 77de7fc7d030..88e48075a644 100644 --- a/python/sglang/srt/entrypoints/engine_info_bootstrap_server.py +++ b/python/sglang/srt/entrypoints/engine_info_bootstrap_server.py @@ -31,7 +31,8 @@ class EngineInfoBootstrapServer: accesses the collected info directly in-process; external consumers can query via HTTP GET. - Currently supports transfer engine memory registration info. + Currently supports transfer engine memory registration info and + per-rank parallelism configuration. """ def __init__(self, host: str, port: int): @@ -40,6 +41,8 @@ def __init__(self, host: str, port: int): # Storage: {tp_rank: (session_id, weights_info_dict)} self.transfer_engine_info: Dict[int, Tuple] = {} + # Storage: {tp_rank: parallelism_config_dict} + self.parallelism_config: Dict[int, dict] = {} self.lock = threading.Lock() app = FastAPI() @@ -89,6 +92,38 @@ def get_transfer_engine_info(rank: int): config = uvicorn.Config(app, host=host, port=port, log_level="warning") self._server = uvicorn.Server(config) + + @app.put("/register_parallelism_config") + def register_parallelism_config(data: dict): + try: + tp_rank = data["tp_rank"] + config = data["parallelism_config"] + + with self.lock: + self.parallelism_config[tp_rank] = config + + logger.info(f"Registered parallelism config for tp_rank={tp_rank}") + return PlainTextResponse("OK") + except Exception as e: + logger.error(f"Failed to register parallelism config: {e}") + raise HTTPException(status_code=400, detail=str(e)) + + @app.get("/get_parallelism_config") + def get_parallelism_config(rank: int): + if rank < 0: + raise HTTPException(status_code=400, detail="Invalid rank parameter") + + with self.lock: + config = self.parallelism_config.get(rank) + + if config is None: + raise HTTPException( + status_code=404, + detail=f"No parallelism config for rank {rank}", + ) + + return config + self._thread = threading.Thread( target=self._server.run, daemon=True, @@ -103,3 +138,7 @@ def close(self): def get_transfer_engine_info(self, rank: int) -> Optional[Tuple]: """Direct in-process access for co-located HTTP server (no HTTP round-trip).""" return self.transfer_engine_info.get(rank) + + def get_parallelism_config_info(self, rank: int) -> Optional[dict]: + """Direct in-process access for parallelism config (no HTTP round-trip).""" + return self.parallelism_config.get(rank) diff --git a/python/sglang/srt/entrypoints/grpc_bridge.py b/python/sglang/srt/entrypoints/grpc_bridge.py index 48d5b0b8d8ba..234ab8bb7c58 100644 --- a/python/sglang/srt/entrypoints/grpc_bridge.py +++ b/python/sglang/srt/entrypoints/grpc_bridge.py @@ -84,6 +84,9 @@ def __init__( self.tokenizer_manager.auto_create_handle_loop() self._event_loop = self.tokenizer_manager.event_loop + def set_engine_state_changed_callback(self, callback) -> None: + self.tokenizer_manager.set_engine_state_changed_callback(callback) + @property def _tm_loop(self): """Return the TokenizerManager loop used by communicator RPCs.""" @@ -446,6 +449,10 @@ def health_check(self) -> bool: ServerStatus.UnHealthy, ) + def is_pause(self) -> bool: + """Return the tokenizer manager's authoritative generation pause state.""" + return self.tokenizer_manager.is_pause + def tokenize(self, text: str, add_special_tokens: bool = True) -> str: tokenizer = self.tokenizer_manager.tokenizer tokens = tokenizer.encode(text, add_special_tokens=add_special_tokens) diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index dd93d65459ca..17ec94067460 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -117,6 +117,7 @@ from sglang.srt.managers.io_struct import ( AbortReq, AttachHiCacheStorageReqInput, + BeginWeightUpdateReqInput, CheckWeightsReqInput, CloseSessionReqInput, ConfigureLoggingReq, @@ -124,16 +125,18 @@ DestroyWeightsUpdateGroupReqInput, DumperControlReqInput, EmbeddingReqInput, + EndWeightUpdateReqInput, GenerateReqInput, GetWeightsByNameReqInput, InitWeightsSendGroupForRemoteInstanceReqInput, InitWeightsUpdateGroupReqInput, - LoadLoRAAdapterFromTensorsReqInput, LoadLoRAAdapterReqInput, OpenSessionReqInput, ParseFunctionCallReq, PauseGenerationReqInput, ProfileReq, + PullWeightsReqInput, + RegisterLoRAAdapterReqInput, ReleaseMemoryOccupationReqInput, ResumeMemoryOccupationReqInput, SendWeightsToRemoteInstanceReqInput, @@ -659,6 +662,13 @@ async def validate_json_request(raw_request: Request): ##### Native API endpoints ##### +@app.get("/ready") +async def ready() -> Response: + """Report whether the server is ready to receive new requests.""" + status_code = 200 if _global_state.tokenizer_manager.is_ready() else 503 + return Response(status_code=status_code) + + @app.get("/health") @app.get("/health_generate") async def health_generate(request: Request) -> Response: @@ -1278,6 +1288,19 @@ async def update_weights_from_disk( ) +@app.post("/pull_weights") +@auth_level(AuthLevel.ADMIN_OPTIONAL) +async def pull_weights(obj: Annotated[PullWeightsReqInput, Body()], request: Request): + """Have every host of this deployment pull published weight deltas into its + local checkpoint (materialized from the model path on first use).""" + success, message = await _global_state.tokenizer_manager.pull_weights(obj, request) + + content = {"success": success, "message": message} + return ORJSONResponse( + content, status_code=HTTPStatus.OK if success else HTTPStatus.BAD_REQUEST + ) + + @app.post("/init_weights_send_group_for_remote_instance") @auth_level(AuthLevel.ADMIN_OPTIONAL) async def init_weights_send_group_for_remote_instance( @@ -1353,6 +1376,31 @@ async def remote_instance_transfer_engine_info(rank: int = None): ) +@app.get("/parallelism_config") +@auth_level(AuthLevel.ADMIN_OPTIONAL) +async def parallelism_config(rank: int = None): + """Get per-rank parallelism config from the bootstrap server.""" + if rank is None or rank < 0: + return ORJSONResponse( + {"error": {"message": "Missing or invalid rank parameter"}}, + status_code=HTTPStatus.BAD_REQUEST, + ) + + server_args = _global_state.tokenizer_manager.server_args + try: + resp = requests.get( + f"{server_args.engine_info_bootstrap_url}/get_parallelism_config", + params={"rank": rank}, + timeout=5, + ) + if resp.status_code == 200: + return resp.json() + except Exception: + pass + + return Response(status_code=HTTPStatus.BAD_REQUEST) + + @app.post("/init_weights_update_group") @auth_level(AuthLevel.ADMIN_OPTIONAL) async def init_weights_update_group( @@ -1407,6 +1455,36 @@ async def update_weights_from_tensor( ) +@app.post("/begin_weight_update") +@auth_level(AuthLevel.ADMIN_OPTIONAL) +async def begin_weight_update( + obj: Annotated[BeginWeightUpdateReqInput, Body()], request: Request +): + """Open a weight-update session so in-place-quantized weights become loadable.""" + success, message = await _global_state.tokenizer_manager.begin_weight_update( + obj, request + ) + return ORJSONResponse( + {"success": success, "message": message}, + status_code=HTTPStatus.OK if success else HTTPStatus.BAD_REQUEST, + ) + + +@app.post("/end_weight_update") +@auth_level(AuthLevel.ADMIN_OPTIONAL) +async def end_weight_update( + obj: Annotated[EndWeightUpdateReqInput, Body()], request: Request +): + """Close the weight-update session and finalize quantized weights.""" + success, message = await _global_state.tokenizer_manager.end_weight_update( + obj, request + ) + return ORJSONResponse( + {"success": success, "message": message}, + status_code=HTTPStatus.OK if success else HTTPStatus.BAD_REQUEST, + ) + + @app.post("/update_weights_from_distributed") @auth_level(AuthLevel.ADMIN_OPTIONAL) async def update_weights_from_distributed( @@ -1565,14 +1643,13 @@ async def load_lora_adapter( return ORJSONResponse(msgspec_to_builtins(result), status_code=status_code) -@app.api_route("/load_lora_adapter_from_tensors", methods=["POST"]) -async def load_lora_adapter_from_tensors( - obj: Annotated[LoadLoRAAdapterFromTensorsReqInput, Body()], request: Request +@app.api_route("/register_lora_adapter", methods=["POST"]) +@auth_level(AuthLevel.ADMIN_OPTIONAL) +async def register_lora_adapter( + obj: Annotated[RegisterLoRAAdapterReqInput, Body()], request: Request ): - """Load a new LoRA adapter from tensors without re-launching the server.""" - result = await _global_state.tokenizer_manager.load_lora_adapter_from_tensors( - obj, request - ) + """Create-or-refresh a LoRA adapter's identity and config (weights zeroed).""" + result = await _global_state.tokenizer_manager.register_lora_adapter(obj, request) status_code = HTTPStatus.OK if result.success else HTTPStatus.BAD_REQUEST return ORJSONResponse(msgspec_to_builtins(result), status_code=status_code) @@ -1628,7 +1705,7 @@ async def abort_request(obj: Annotated[AbortReq, Body()], request: Request): """Abort a request.""" try: _global_state.tokenizer_manager.abort_request( - rid=obj.rid, abort_all=obj.abort_all + rid=obj.rid, abort_all=obj.abort_all, prefix=obj.prefix ) return Response(status_code=200) except Exception as e: diff --git a/python/sglang/srt/entrypoints/openai/protocol.py b/python/sglang/srt/entrypoints/openai/protocol.py index d1509a7ccc4b..924ec1610fde 100644 --- a/python/sglang/srt/entrypoints/openai/protocol.py +++ b/python/sglang/srt/entrypoints/openai/protocol.py @@ -894,6 +894,7 @@ class ChatCompletionRequest(BaseModel): return_input_ids_in_sglext: bool = False return_output_ids_in_sglext: bool = False return_sampling_mask: bool = False + sampling_logprobs_mode: Optional[Literal["selected", "support"]] = None reasoning_effort: ReasoningEffortType = Field( default=None, description="Constrains effort on reasoning for reasoning models. " diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index 89cac804aff9..13e40412e12c 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -1162,6 +1162,7 @@ def _convert_to_internal_request( logprob_start_len=-1, top_logprobs_num=request.top_logprobs or 0, return_sampling_mask=request.return_sampling_mask, + sampling_logprobs_mode=request.sampling_logprobs_mode, stream=request.stream, return_text_in_logprobs=True, modalities=processed_messages.modalities, @@ -2367,6 +2368,15 @@ def _process_tool_call_id( """Process for generating a new and unique `tool_call_id`""" if self.tool_call_parser == "kimi_k3": return f"{call_item.name}:{history_tool_calls_cnt + call_item.tool_index}" + if self.tool_call_parser == "kimi_k2_raw_id": + # RL training needs the model-emitted tool_call_id round-tripped verbatim, + # so we skip the history-based renumbering below and return whatever the + # detector captured. Fall back to the canonical Kimi-K2 reconstruction + # (without history offset) if for any reason the detector did not record + # a raw id — the raw id field is best-effort but the format is stable. + if call_item.tool_call_id: + return call_item.tool_call_id + return f"functions.{call_item.name}:{call_item.tool_index}" if self.tool_call_parser != "kimi_k2": # A simple uuid is sufficient for all models except for Kimi-K2. tool_call_id = f"call_{uuid.uuid4().hex[:24]}" diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 46e9c70dc3a4..301a24e6893f 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1696,6 +1696,7 @@ class Envs: # =================================================================== SGLANG_SYMM_MEM_PREALLOC_GB_SIZE = EnvInt(-1) SGLANG_DEBUG_SYMM_MEM = EnvBool(False) + SGLANG_DISABLE_MULTIMEM_AG = EnvBool(False) # Qwen3.5 and GDN SGLANG_ENABLE_GDN_DECODE_FUSED_PROJ_CONV = EnvBool(True) diff --git a/python/sglang/srt/function_call/core_types.py b/python/sglang/srt/function_call/core_types.py index 1ea87df798c8..297dd2712ef3 100644 --- a/python/sglang/srt/function_call/core_types.py +++ b/python/sglang/srt/function_call/core_types.py @@ -10,6 +10,12 @@ class ToolCallItem(BaseModel): tool_index: int name: Optional[str] = None parameters: str # JSON string + # The tool_call_id string emitted by the model, captured verbatim. + # Only populated by detectors whose downstream consumers need the exact + # model-emitted id (e.g. RL training trajectories). Existing detectors + # leave this as None and the serving layer falls back to its usual id + # generation strategy. + tool_call_id: Optional[str] = None class StreamingParseResult(BaseModel): diff --git a/python/sglang/srt/function_call/function_call_parser.py b/python/sglang/srt/function_call/function_call_parser.py index 2bf5dad8411d..a1b06e100d7f 100644 --- a/python/sglang/srt/function_call/function_call_parser.py +++ b/python/sglang/srt/function_call/function_call_parser.py @@ -34,7 +34,10 @@ from sglang.srt.function_call.inkling_detector import InklingDetector from sglang.srt.function_call.internlm_detector import InternlmDetector from sglang.srt.function_call.k2_v3_detector import K2V3Detector -from sglang.srt.function_call.kimik2_detector import KimiK2Detector +from sglang.srt.function_call.kimik2_detector import ( + KimiK2Detector, + KimiK2RawIdDetector, +) from sglang.srt.function_call.kimik3_detector import KimiK3Detector from sglang.srt.function_call.lfm2_detector import Lfm2Detector from sglang.srt.function_call.ling3_detector import Ling3Detector @@ -84,6 +87,7 @@ class FunctionCallParser: "k2_horizon": K2V3Detector, "kimi_k2": KimiK2Detector, "kimi_k3": KimiK3Detector, + "kimi_k2_raw_id": KimiK2RawIdDetector, "lfm2": Lfm2Detector, "ling3": Ling3Detector, "llama3": Llama32Detector, diff --git a/python/sglang/srt/function_call/kimik2_detector.py b/python/sglang/srt/function_call/kimik2_detector.py index a8054b587337..eccc986bea90 100644 --- a/python/sglang/srt/function_call/kimik2_detector.py +++ b/python/sglang/srt/function_call/kimik2_detector.py @@ -198,6 +198,7 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult tool_index=local_tool_index, name=function_name, parameters=function_args, + tool_call_id=function_id, ) ) local_tool_index += 1 @@ -329,6 +330,9 @@ def parse_streaming_increment( else None ), parameters=argument_diff, + # Capture the model-emitted id on the name-carrying + # delta so kimi_k2_raw_id can round-trip it (RL). + tool_call_id=function_id if name_just_resolved else None, ) ) if argument_diff: @@ -472,3 +476,23 @@ def get_structural_tag( def get_structural_tag_name(self) -> str: return "kimi" + + +class KimiK2RawIdDetector(KimiK2Detector): + """ + Variant of KimiK2Detector that preserves the model-emitted tool_call_id verbatim. + + The default kimi_k2 path renumbers ids via `history_tool_calls_cnt + tool_index` + in the serving layer so that multi-turn conversations get globally unique, + monotonically increasing ids (see PR #10600). That is the right behavior for + chat use cases. + + RL training has the opposite requirement: the trajectory must round-trip the + exact tool_call_id the model produced (e.g. `functions.foo:5`), so that the + follow-up tool result turn references the same id the policy emitted. This + subclass exists purely as a marker so the serving layer can branch on the + parser name and use `ToolCallItem.tool_call_id` directly. Parsing logic + is identical to KimiK2Detector. + """ + + pass diff --git a/python/sglang/srt/layers/activation.py b/python/sglang/srt/layers/activation.py index b620fd701eff..2da5ea27fc01 100644 --- a/python/sglang/srt/layers/activation.py +++ b/python/sglang/srt/layers/activation.py @@ -33,7 +33,8 @@ Phase, check_cuda_graph_backend, ) -from sglang.srt.runtime_context import get_exec, get_parallel, publish_role +from sglang.srt.runtime_context import get_parallel, publish_role +from sglang.srt.true_on_policy import is_true_on_policy_enabled from sglang.srt.utils import ( cpu_has_amx_support, get_bool_env_var, @@ -130,10 +131,7 @@ def _(x): class SiluAndMul(BaseFusedOp): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - if ( - publish_role() is not None - and get_exec().deterministic.rl_on_policy_target is not None - ): + if publish_role() is not None and is_true_on_policy_enabled(): self._forward_method = self.forward_native elif _use_aiter and envs.SGLANG_OPT_USE_AITER_SILU_MUL.get(): self._forward_method = self.forward_aiter diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py index 26a457f62df9..8f1bf7be1934 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py @@ -46,6 +46,7 @@ get_schedule, ) from sglang.srt.state_capturer.indexer_topk import ( + get_global_indexer_capturer, maybe_capture_indexer_topk, ) from sglang.srt.utils import ( @@ -976,7 +977,14 @@ def _chunked_fp8_paged_mqa_logits( # NOTE(dark): logits should be cleaned in topk_transform self._mask_init_and_local_tokens(logits, seqlens_32) - topk_result = metadata.topk_transform(logits, self.index_topk) + capture = get_global_indexer_capturer() is not None + if capture: + topk_result, raw_result = metadata.topk_transform( + logits, self.index_topk, return_raw_indices=True + ) + else: + topk_result = metadata.topk_transform(logits, self.index_topk) + raw_result = None # Restore possible padding exist in the hidden states. if not _is_hip and q_offset < q_fp8.shape[0]: pad_len = q_fp8.shape[0] - q_offset @@ -987,7 +995,9 @@ def _chunked_fp8_paged_mqa_logits( device=topk_result.device, ) topk_result = torch.cat([topk_result, padding], dim=0) - return topk_result + if capture: + raw_result = torch.cat([raw_result, padding], dim=0) + return topk_result, raw_result def _get_mqa_logits_budget_bytes(self, device_index: int) -> int: free_mem_fraction = self._mqa_logits_free_mem_fraction() @@ -1113,8 +1123,10 @@ def _get_topk_ragged( topk_result = torch.full( (token_nums, self.index_topk), -1, device=device, dtype=torch.int32 ) + capture = get_global_indexer_capturer() is not None + raw_result = torch.full_like(topk_result, -1) if capture else None if batch_size == 0: - return topk_result + return topk_result, raw_result ks, ke = metadata.get_indexer_kvcache_range() @@ -1190,9 +1202,17 @@ def _get_topk_ragged( assert logits.shape[1] == k_offset self._mask_init_and_local_tokens(logits, seq_lens_expanded, ks) - raw_topk_result = metadata.topk_transform(logits, self.index_topk, ks=ks) - topk_result[:q_offset] = raw_topk_result - return topk_result + if capture: + transformed, raw = metadata.topk_transform( + logits, self.index_topk, ks=ks, return_raw_indices=True + ) + topk_result[:q_offset] = transformed + raw_result[:q_offset] = raw + else: + topk_result[:q_offset] = metadata.topk_transform( + logits, self.index_topk, ks=ks + ) + return topk_result, raw_result bytes_per_row = k_offset * self._MQA_LOGITS_BYTES_PER_ELEM max_rows = max(1, int(logits_budget_bytes // max(bytes_per_row, 1))) @@ -1267,19 +1287,32 @@ def _get_topk_ragged( cu_seqlens_q_chunk = cu_seqlens_q_full[start:end] batch_idx_chunk = token_to_batch_idx[start:end] - raw_topk_chunk = metadata.topk_transform( - logits_chunk, - self.index_topk, - ks=ks[start:end], - cu_seqlens_q=cu_seqlens_q_chunk, - ke_offset=lengths_chunk, - batch_idx_list=batch_idx_chunk, - topk_indices_offset_override=topk_offset_chunk, - ) - topk_result[start:end] = raw_topk_chunk + if capture: + transformed, raw = metadata.topk_transform( + logits_chunk, + self.index_topk, + ks=ks[start:end], + cu_seqlens_q=cu_seqlens_q_chunk, + ke_offset=lengths_chunk, + batch_idx_list=batch_idx_chunk, + topk_indices_offset_override=topk_offset_chunk, + return_raw_indices=True, + ) + topk_result[start:end] = transformed + raw_result[start:end] = raw + else: + topk_result[start:end] = metadata.topk_transform( + logits_chunk, + self.index_topk, + ks=ks[start:end], + cu_seqlens_q=cu_seqlens_q_chunk, + ke_offset=lengths_chunk, + batch_idx_list=batch_idx_chunk, + topk_indices_offset_override=topk_offset_chunk, + ) start = end - return topk_result + return topk_result, raw_result def _forward_cuda_k_only( self, @@ -1347,7 +1380,7 @@ def _forward_cuda_k_only( # MHA doesn't need topk_indices if not return_indices: - return None + return None, None # MLA: use dummy logits with topk kernel's fast path to generate indices # When length <= 2048, naive_topk_cuda directly generates [0,1,...,length-1,-1,...] @@ -1358,13 +1391,19 @@ def _forward_cuda_k_only( dtype=torch.float32, device=x_meta.device, ) - raw_topk_result = metadata.topk_transform(dummy_logits, self.index_topk) + if get_global_indexer_capturer() is not None: + raw_topk_result, raw_result = metadata.topk_transform( + dummy_logits, self.index_topk, return_raw_indices=True + ) + else: + raw_topk_result = metadata.topk_transform(dummy_logits, self.index_topk) + raw_result = None if topk_result is not None: # PCG/BCG: fill the valid prefix of the padded static buffer and # leave padded rows at the -1 sentinel. topk_result[: raw_topk_result.shape[0]] = raw_topk_result - return None - return raw_topk_result + return None, raw_result + return raw_topk_result, raw_result def _store_index_k_cache( self, @@ -1466,6 +1505,15 @@ def forward_xpu( x, q_lora, positions, forward_batch, layer_id, return_indices ) + def _capture_and_return(self, layer_id, topk_result, raw_result): + # Capture the model's natural (sequence-relative) topk for rollout R3 + # replay; the attention kernel still receives the transformed + # (paged/ragged kv-cache) indices in topk_result. + maybe_capture_indexer_topk( + layer_id, raw_result if raw_result is not None else topk_result + ) + return topk_result + def forward_cuda( self, x: torch.Tensor, @@ -1519,7 +1567,7 @@ def forward_cuda( # Optimization: fast path when skipping topk computation if skip_logits_computation and (not self.dsa_enable_prefill_cp): - topk_result = self._forward_cuda_k_only( + topk_result, raw_result = self._forward_cuda_k_only( x, positions, forward_batch, @@ -1529,7 +1577,7 @@ def forward_cuda( return_indices, ) topk_result = _broadcast_indexer_topk_from_rank0(topk_result) - return maybe_capture_indexer_topk(layer_id, topk_result) + return self._capture_and_return(layer_id, topk_result, raw_result) # When weights_proj is LoRA-wrapped, use an eager module call so the # wrapper owns base+delta and no LoRA kernel runs under torch.compile. @@ -1713,6 +1761,7 @@ def forward_cuda( else: weights = self._get_logits_head_gate(x_for_gate, q_scale) + raw_result = None if _is_cuda or _is_hip or _is_xpu: # In piecewise/breakable CUDA graph, any access to seq_lens_cpu # creates a Dynamo shape guard. These graph modes never have empty @@ -1731,14 +1780,14 @@ def forward_cuda( device=x_meta.device, ) topk_result = _broadcast_indexer_topk_from_rank0(topk_result) - return maybe_capture_indexer_topk(layer_id, topk_result) + return self._capture_and_return(layer_id, topk_result, None) if ( forward_batch.forward_mode.is_decode_or_idle() or forward_batch.forward_mode.is_target_verify() or forward_batch.forward_mode.is_draft_extend_v2() ): - topk_result = self._get_topk_paged( + topk_result, raw_result = self._get_topk_paged( forward_batch, layer_id, q_fp8, weights, metadata ) else: @@ -1749,7 +1798,7 @@ def forward_cuda( "Internal error: in-graph DSA prefill must go through the " "graph DSA split-op dispatch" ) - topk_result = self._get_topk_ragged( + topk_result, raw_result = self._get_topk_ragged( enable_dual_stream, forward_batch, layer_id, @@ -1760,4 +1809,4 @@ def forward_cuda( else: raise NotImplementedError("DSA indexer only supports CUDA, HIP, and NPU") topk_result = _broadcast_indexer_topk_from_rank0(topk_result) - return maybe_capture_indexer_topk(layer_id, topk_result) + return self._capture_and_return(layer_id, topk_result, raw_result) diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer_metadata.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer_metadata.py index 134b65aaf519..77de6879c829 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer_metadata.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer_metadata.py @@ -135,6 +135,7 @@ def topk_transform( ke_offset: Optional[torch.Tensor] = None, batch_idx_list: Optional[List[int]] = None, topk_indices_offset_override: Optional[torch.Tensor] = None, + return_raw_indices: bool = False, ) -> torch.Tensor: if topk_indices_offset_override is not None: cu_topk_indices_offset = topk_indices_offset_override @@ -166,4 +167,5 @@ def topk_transform( row_starts=ks, batch_idx_list=batch_idx_list, force_unfused_topk=self.force_unfused_topk, + return_raw_indices=return_raw_indices, ) diff --git a/python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py b/python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py index 2ff9256f9936..4a4b115eb976 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py @@ -102,9 +102,21 @@ def topk_transform( row_starts: Optional[torch.Tensor] = None, batch_idx_list: Optional[List[int]] = None, force_unfused_topk: bool = False, + return_raw_indices: bool = False, ) -> torch.Tensor: if not envs.SGLANG_DSA_FUSE_TOPK.get() or force_unfused_topk: - return self.topk_func(logits, lengths, topk, row_starts=row_starts) + result = self.topk_func(logits, lengths, topk, row_starts=row_starts) + return (result, result) if return_raw_indices else result + + # Sequence-relative selection, before the kv-cache coordinate remap done + # by the fused transform below. Rollout R3 replay compares against the + # model's natural (sequence-relative) topk, so capture must use these raw + # indices rather than the transformed (paged/ragged) result. + raw_indices = ( + self.topk_func(logits, lengths, topk, row_starts=row_starts) + if return_raw_indices + else None + ) # Decode-shaped PAGED top-k for the SGL backend (plain decode AND spec # verify / draft-extend, whose expanded rows match the same shape) routes @@ -127,7 +139,8 @@ def topk_transform( == logits.shape[0] == attn_metadata.real_page_table.shape[0] ): - return _topk_transform_v2_paged(logits, lengths, topk, attn_metadata) + result = _topk_transform_v2_paged(logits, lengths, topk, attn_metadata) + return (result, raw_indices) if return_raw_indices else result # Extend-shaped RAGGED top-k for the SGL backend routes to the same v2 # kernel through its ragged entry point: no page table (the columns are @@ -165,7 +178,7 @@ def topk_transform( if batch_idx_list is not None else attn_metadata.page_table_1 ) - return fast_topk_transform_fused( + result = fast_topk_transform_fused( score=logits, lengths=lengths, page_table_size_1=page_table_size_1, @@ -173,22 +186,22 @@ def topk_transform( topk=topk, row_starts=row_starts, ) - if topk_transform_method == TopkTransformMethod.RAGGED: + elif topk_transform_method == TopkTransformMethod.RAGGED: if topk_indices_offset is None: raise RuntimeError( "RAGGED topk_transform requires topk_indices_offset; " "expected extend-without-speculative metadata." ) - return fast_topk_transform_ragged_fused( + result = fast_topk_transform_ragged_fused( score=logits, lengths=lengths, topk_indices_offset=topk_indices_offset, topk=topk, row_starts=row_starts, ) - raise RuntimeError(f"Unsupported {topk_transform_method = }.") - - if self.is_flashinfer(): + else: + raise RuntimeError(f"Unsupported {topk_transform_method = }.") + elif self.is_flashinfer(): import flashinfer if topk_transform_method == TopkTransformMethod.PAGED: @@ -200,7 +213,7 @@ def topk_transform( device=logits.device, num_rows=logits.shape[0], ) - return flashinfer.top_k_page_table_transform( + result = flashinfer.top_k_page_table_transform( logits.contiguous(), attn_metadata.page_table_1.contiguous(), lengths.contiguous(), @@ -212,13 +225,13 @@ def topk_transform( row_starts=row_starts, page_table_row_starts=page_table_row_starts, ) - if topk_transform_method == TopkTransformMethod.RAGGED: + elif topk_transform_method == TopkTransformMethod.RAGGED: if topk_indices_offset is None: raise RuntimeError( "RAGGED topk_transform requires topk_indices_offset; " "expected extend-without-speculative metadata." ) - return flashinfer.top_k_ragged_transform( + result = flashinfer.top_k_ragged_transform( logits.contiguous(), topk_indices_offset.contiguous(), lengths.contiguous(), @@ -228,9 +241,12 @@ def topk_transform( dsa_graph_safe=True, row_starts=row_starts, ) - raise RuntimeError(f"Unsupported {topk_transform_method = }.") + else: + raise RuntimeError(f"Unsupported {topk_transform_method = }.") + else: + raise RuntimeError(f"Unsupported {self = } for SGLANG_DSA_FUSE_TOPK.") - raise RuntimeError(f"Unsupported {self = } for SGLANG_DSA_FUSE_TOPK.") + return (result, raw_indices) if return_raw_indices else result def _topk_unfused( diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index 8fdd7ba12dda..0a382cf92119 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -1,6 +1,7 @@ from __future__ import annotations import logging +from contextlib import nullcontext from dataclasses import dataclass from typing import ( TYPE_CHECKING, @@ -15,6 +16,7 @@ import torch from sglang.srt.configs.model_config import get_dsa_index_topk, is_deepseek_dsa +from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH from sglang.srt.runtime_context import ( get_buffer, get_exec, @@ -91,6 +93,7 @@ is_xpu, print_warning_once, ) +from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter _IS_GFX95 = is_gfx95_supported() @@ -337,6 +340,7 @@ def __init__( self.needs_cpu_seq_lens = self.dsa_index_kpool > 1 self._init_kpool_metadata_fusion() self.max_context_len = model_runner.model_config.context_len + self.enable_memory_saver = model_runner.server_args.enable_memory_saver self.num_q_heads = ( model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size ) @@ -1219,6 +1223,16 @@ def _cal_indexer_k_start_end( token_to_batch_idx = split_per_token(token_to_batch_idx) return (ks, ke), token_to_batch_idx + def _cuda_graph_memory_region(self): + """Region whose tensors are released while the engine is paused.""" + adapter = TorchMemorySaverAdapter.create( + enable=self.enable_memory_saver + and envs.SGLANG_MEMORY_SAVER_CUDA_GRAPH.get() + ) + if not adapter.enabled: + return nullcontext() + return adapter.region(tag=GPU_MEMORY_TYPE_CUDA_GRAPH) + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): """Initialize CUDA graph state for the attention backend. @@ -1256,6 +1270,31 @@ def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): ) max_ctx_len = self.req_to_token.shape[1] + # The wide page_table (max_num_tokens x max_ctx_len int32) and the + # flashmla metadata are rewritten in full before every replay, so they + # can live in the pausable cuda-graph region. The others are written + # once at init and would come back zeroed after a resume. + with self._cuda_graph_memory_region(): + page_table = ( + None + if self.dsa_drop_wide_page_table + else torch.zeros( + max_num_tokens, + max_ctx_len, + dtype=torch.int32, + device=self.device, + ) + ) + flashmla_metadata = ( + self._compute_flashmla_metadata( + cache_seqlens=torch.ones( + max_num_tokens, dtype=torch.int32, device=self.device + ), + seq_len_q=1, + ) + if self.dsa_decode_impl == "flashmla_kv" + else None + ) self.decode_cuda_graph_metadata: Dict = { "cache_seqlens": torch.ones( max_num_tokens, dtype=torch.int32, device=self.device @@ -1282,26 +1321,8 @@ def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): if self.dsa_drop_wide_page_table else None ), - "page_table": ( - None - if self.dsa_drop_wide_page_table - else torch.zeros( - max_num_tokens, - max_ctx_len, - dtype=torch.int32, - device=self.device, - ) - ), - "flashmla_metadata": ( - self._compute_flashmla_metadata( - cache_seqlens=torch.ones( - max_num_tokens, dtype=torch.int32, device=self.device - ), - seq_len_q=1, - ) - if self.dsa_decode_impl == "flashmla_kv" - else None - ), + "page_table": page_table, + "flashmla_metadata": flashmla_metadata, } def _build_forward_metadata_cuda_graph( diff --git a/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py b/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py index 369654de6803..3cfd762e8f35 100644 --- a/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py +++ b/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py @@ -254,10 +254,16 @@ def prepare_mixed( num_prefill_tokens = int(sum(extend_seq_lens_cpu)) else: num_prefill_tokens = int(forward_batch.extend_num_tokens) - batch_size = getattr(forward_batch, "_original_batch_size", None) - if batch_size is None: - batch_size = len(forward_batch.seq_lens) - num_decodes = max(0, batch_size - num_prefills) + if forward_batch._original_forward_mode is not None: + # mlp-sync converted the whole batch to EXTEND, so there are no decode + # rows; on an idle rank _original_batch_size is 0 and subtracting here + # would go negative. + num_decodes = 0 + else: + batch_size = forward_batch._original_batch_size + if batch_size is None: + batch_size = len(forward_batch.seq_lens) + num_decodes = batch_size - num_prefills context_lens_tensor = forward_batch.extend_prefix_lens assert context_lens_tensor is not None has_initial_states = context_lens_tensor > 0 diff --git a/python/sglang/srt/layers/attention/qsa/mqa.py b/python/sglang/srt/layers/attention/qsa/mqa.py index 6446ed145cbe..5d4580529751 100644 --- a/python/sglang/srt/layers/attention/qsa/mqa.py +++ b/python/sglang/srt/layers/attention/qsa/mqa.py @@ -120,7 +120,6 @@ def _tilelang_qsa_mqa_prefill_kernel( head_dim: int, block_n: int = 64, block_q: int = 32, - num_stages: int = 3, threads: int = 512, ): rows = T.dynamic("rows") @@ -150,9 +149,8 @@ def kernel( end_max = T.max(end_max, T.min(Ends[row_base + qi], keys)) T.copy(Q[row_base * heads, 0], q_shared) - for ni in T.Pipelined( - T.ceildiv(end_max - start_min, block_n), num_stages=num_stages - ): + # a serial loop: tilelang >= 0.1.14 only pipelines constant trip counts + for ni in T.serial(T.ceildiv(end_max - start_min, block_n)): T.copy(K[start_min + ni * block_n, 0], k_shared) T.gemm( k_shared, diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py index bff80a1f1a09..6a1d96405903 100644 --- a/python/sglang/srt/layers/attention/vision.py +++ b/python/sglang/srt/layers/attention/vision.py @@ -19,11 +19,11 @@ from sglang.srt.models.utils import apply_qk_norm from sglang.srt.runtime_context import ( get_context, - get_exec, get_mm, get_parallel, get_platform, ) +from sglang.srt.true_on_policy import is_true_on_policy_enabled from sglang.srt.utils import ( cpu_has_amx_support, get_bool_env_var, @@ -1248,7 +1248,7 @@ def _init_qk_norm( weight_dtype=torch.float32, cast_x_before_out_mul=True, ) - if get_exec().deterministic.rl_on_policy_target is not None + if is_true_on_policy_enabled() else {} ) q_norm = RMSNorm( @@ -1390,10 +1390,7 @@ def forward( if x.dim() == 2: x = x.unsqueeze(0) assert x.dim() == 3, x.shape - if ( - get_exec().deterministic.rl_on_policy_target is not None - and position_embeddings is not None - ): + if is_true_on_policy_enabled() and position_embeddings is not None: assert isinstance(position_embeddings, tuple), ( "expected position_embeddings to be a tuple of two tensors,\n" f"but got {type(position_embeddings)}, change if needed" diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index 323d4c309707..6d2a01765ba6 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/python/sglang/srt/layers/communicator.py @@ -82,6 +82,10 @@ get_spec, ) from sglang.srt.speculative.spec_info import SpeculativeAlgorithm +from sglang.srt.true_on_policy import ( + should_disable_mlp_allreduce_fusion_for_on_policy, + should_disable_reduce_scatter_for_on_policy, +) from sglang.srt.utils import ( get_bool_env_var, is_cuda, @@ -905,6 +909,9 @@ def postprocess_layer( ) def should_use_reduce_scatter(self, forward_batch: ForwardBatch): + if should_disable_reduce_scatter_for_on_policy(): + return False + if not self.allow_reduce_scatter: return False if ( @@ -925,6 +932,9 @@ def should_use_reduce_scatter(self, forward_batch: ForwardBatch): def should_fuse_mlp_allreduce_with_next_layer( self, forward_batch: ForwardBatch ) -> bool: + if should_disable_mlp_allreduce_fusion_for_on_policy(): + return False + # When MOE_FULL is active (moe_cp allgather), fusion must be disabled because # the fusion path skips postprocess_layer which contains the moe_cp scatter. # Without scatter, hidden_states remain at MOE_FULL size while residual is at diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py index eb3e5988d4ec..965ef1140abf 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py @@ -43,7 +43,13 @@ # Force redirect deep_gemm cache_dir. Defaults under SGLANG_CACHE_DIR so it # sits with the other compiled-kernel caches; SGLANG_DG_CACHE_DIR still wins. -os.environ["DG_JIT_CACHE_DIR"] = envs.SGLANG_DG_CACHE_DIR.get() +_dg_cache_dir = envs.SGLANG_DG_CACHE_DIR.get() +if os.getenv("SGLANG_DG_CACHE_DIR_PER_PROCESS", "0").lower() in ("1", "true", "yes"): + # Colocated RL engines share a host; a per-process dir keeps their JIT caches + # from racing on the same files. + _dg_cache_dir = os.path.join(_dg_cache_dir, f"pid_{os.getpid()}") +os.makedirs(_dg_cache_dir, exist_ok=True) +os.environ["DG_JIT_CACHE_DIR"] = _dg_cache_dir # Refer to https://github.com/deepseek-ai/DeepGEMM/commit/d75b218b7b8f4a5dd5406ac87905039ead3ae42f # NVRTC may have performance loss with some cases. diff --git a/python/sglang/srt/layers/dp_attention.py b/python/sglang/srt/layers/dp_attention.py index 71a937bd0d08..bcf42bf178fc 100644 --- a/python/sglang/srt/layers/dp_attention.py +++ b/python/sglang/srt/layers/dp_attention.py @@ -475,16 +475,21 @@ def get_dp_local_slice_cpu( can_run_graph: bool, cuda_graph_batch: Optional[int], ) -> Tuple[int, int]: - # CPU (start, length) slice for DP-local data in a rank-padded buffer. - # Returns Python ints (no D2H sync) and handles the cuda-graph-padded layout. + # CPU (start, length) slice for this rank's real rows, no D2H sync. A graph + # forward strides by the replayed graph's padded size, not by + # max(global_num_tokens): a prefill graph pads 300 tokens to a 320 capture. + # The length excludes that padding -- every pad row holds the routing of a + # zero hidden state, and its out_cache_loc entry points at a live token. global_num_tokens = forward_batch.global_num_tokens_cpu dp_rank = get_attention_dp_rank() - local_num_tokens = global_num_tokens[dp_rank] if can_run_graph: + assert cuda_graph_batch is not None, "graph forward without its size" local_start_pos = dp_rank * cuda_graph_batch else: local_start_pos = sum(global_num_tokens[:dp_rank]) - return local_start_pos, local_num_tokens + original = forward_batch.original_global_num_tokens_cpu + counts = original if original is not None else global_num_tokens + return local_start_pos, counts[dp_rank] from sglang.kernels.ops.memory.memcpy_triton import memcpy_triton diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index 5bb45333a398..6a6b1ef45c8f 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/python/sglang/srt/layers/layernorm.py @@ -33,6 +33,10 @@ check_cuda_graph_backend, ) from sglang.srt.runtime_context import get_exec, get_parallel +from sglang.srt.true_on_policy import ( + get_on_policy_rms_norm_kwargs, + is_true_on_policy_enabled, +) from sglang.srt.utils import ( cpu_has_amx_support, get_bool_env_var, @@ -226,7 +230,7 @@ def _forward_with_allreduce_fusion( if world_size > 1: if post_residual_addition is not None: - residual = residual + post_residual_addition + x = x + post_residual_addition # Prefer AITER fused AR+RMSNorm when enabled on AMD. if _use_aiter: @@ -443,14 +447,33 @@ def __init__( eps: float = 1e-6, var_hidden_size: Optional[int] = None, cast_x_before_out_mul: bool = False, - fp32_residual: bool = False, + fp32_residual: bool = True, has_weight: bool = True, - weight_dtype: Optional = None, - override_orig_dtype: Optional = None, + weight_dtype: Optional[torch.dtype] = None, + override_orig_dtype: Optional[torch.dtype] = None, x_pad_to_multiple: int = 0, force_native: bool = False, + true_on_policy_weight_dtype: Optional[torch.dtype] = None, + true_on_policy_override_orig_dtype: Optional[torch.dtype] = None, + true_on_policy_fp32_residual: bool = False, ) -> None: super().__init__() + true_on_policy_kwargs = get_on_policy_rms_norm_kwargs( + weight_dtype=true_on_policy_weight_dtype, + override_orig_dtype=true_on_policy_override_orig_dtype, + fp32_residual=true_on_policy_fp32_residual, + ) + if not cast_x_before_out_mul: + cast_x_before_out_mul = true_on_policy_kwargs.get( + "cast_x_before_out_mul", cast_x_before_out_mul + ) + fp32_residual = true_on_policy_kwargs.get("fp32_residual", fp32_residual) + if weight_dtype is None: + weight_dtype = true_on_policy_kwargs.get("weight_dtype", weight_dtype) + if override_orig_dtype is None: + override_orig_dtype = true_on_policy_kwargs.get( + "override_orig_dtype", override_orig_dtype + ) self.has_weight = has_weight self.cast_x_before_out_mul = cast_x_before_out_mul self.fp32_residual = fp32_residual @@ -502,11 +525,17 @@ def forward_cuda( return x if self.variance_size_override is not None: return self.forward_native(x, residual, post_residual_addition) + if ( + self.weight.dtype != x.dtype + or self.cast_x_before_out_mul + or self.override_orig_dtype is not None + ): + return self.forward_native(x, residual, post_residual_addition) if is_batch_invariant_mode_enabled(): if ( residual is not None or self.cast_x_before_out_mul - or get_exec().deterministic.rl_on_policy_target == "fsdp" + or is_true_on_policy_enabled() ): return self.forward_native(x, residual, post_residual_addition) original_shape = x.shape @@ -704,10 +733,10 @@ def forward_aiter( ) return result if residual is not None: - residual_out = torch.empty_like(x) - output = torch.empty_like(x) if post_residual_addition is not None: residual = residual + post_residual_addition + residual_out = torch.empty_like(x) + output = torch.empty_like(x) fused_add_rms_norm( output, x, @@ -754,10 +783,10 @@ def forward_hip( # NOTE: Remove this if aiter kernel supports discontinuous input x = x.contiguous() if residual is not None: - out = torch.empty_like(x) - residual_out = torch.empty_like(x) if post_residual_addition is not None: residual = residual + post_residual_addition + out = torch.empty_like(x) + residual_out = torch.empty_like(x) fused_add_rms_norm( out, x, residual_out, residual, self.weight.data, self.variance_epsilon ) @@ -800,15 +829,18 @@ def forward_native( if not x.is_contiguous(): x = x.contiguous() orig_dtype = self.override_orig_dtype or x.dtype + + if residual is not None and not self.fp32_residual: + x = x + residual + if post_residual_addition is not None: + x = x + post_residual_addition + residual = x.clone() x = x.to(torch.float32) - if residual is not None: + if residual is not None and self.fp32_residual: x = x + residual.to(torch.float32) if post_residual_addition is not None: x = x + post_residual_addition.to(torch.float32) - if self.fp32_residual: - residual = x.clone() - else: - residual = x.to(orig_dtype) + residual = x.to(orig_dtype) hidden_size = x.shape[-1] if hidden_size != self.hidden_size: @@ -1119,7 +1151,7 @@ def forward_native( orig_dtype = x.dtype if residual is not None: if post_residual_addition is not None: - residual = residual + post_residual_addition + x = x + post_residual_addition x = x + residual residual = x diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py index 1f3ee1f23b21..af000e1085a2 100644 --- a/python/sglang/srt/layers/linear.py +++ b/python/sglang/srt/layers/linear.py @@ -41,6 +41,7 @@ ) from sglang.srt.layers.utils import pad_or_narrow_weight from sglang.srt.runtime_context import get_exec, get_forward, get_parallel +from sglang.srt.true_on_policy import should_use_tp_invariant_row_linear from sglang.srt.utils import get_bool_env_var, is_cpu, is_hip, is_npu, set_weight_attrs if TYPE_CHECKING: @@ -1587,6 +1588,15 @@ def weight_loader_v2(self, param: BasevLLMParameter, loaded_weight: torch.Tensor assert loaded_weight.numel() == 1 loaded_weight = loaded_weight.reshape(1) + if ( + isinstance(param, BlockQuantScaleParameter) + and getattr(param, "format_ue8m0", False) + and loaded_weight.dtype == torch.int32 + and not self.use_presharded_weights + ): + self._load_ue8m0_packed_scale(param, loaded_weight) + return + if isinstance(param, RowvLLMParameter): # This `BasevLLMParameter` is defined in sglang/srt/layers/parameter.py, # It supports additional parameters like tp_rank and use_presharded_weights. @@ -1609,6 +1619,29 @@ def weight_loader_v2(self, param: BasevLLMParameter, loaded_weight: torch.Tensor # Fallback for parameters that don't accept additional args param.load_row_parallel_weight(loaded_weight) + def _load_ue8m0_packed_scale( + self, param: BlockQuantScaleParameter, loaded_weight: torch.Tensor + ): + # TODO: This is hacky, better to have a more stable solution + # The loaded scale is a full-size DeepGEMM UE8M0 packed tensor (int32, + # mn-major, 4 scale columns per word; see transform_scale_ue8m0). The + # input-dim shard boundary sits at sub-word granularity whenever + # (k_local / block_k) % 4 != 0, so the packed form cannot be narrowed + # directly: unpack to fp32, narrow in scale-column units, then repack. + from sglang.srt.layers.quantization.fp8_utils import ( + inverse_transform_scale_ue8m0, + transform_scale_ue8m0, + ) + + block_k = self.quant_method.quant_config.weight_block_size[1] + mn = param.data.shape[0] + sf_fp32 = inverse_transform_scale_ue8m0(loaded_weight, mn=mn) + shard_size = (self.input_size_per_partition + block_k - 1) // block_k + sf_local = sf_fp32.narrow(1, self.tp_rank * shard_size, shard_size).contiguous() + packed = transform_scale_ue8m0(sf_local, mn=mn) + assert param.data.shape == packed.shape, f"{param.data.shape=} {packed.shape=}" + param.data.copy_(packed) + def forward( self, input_, @@ -1651,7 +1684,27 @@ def forward( get_tp_group(), disabled=not is_allocation_symmetric() ) with symm_ctx: - if output_tensor is None: + from sglang.srt.layers.quantization.fp8 import Fp8LinearMethod + + is_tuple = isinstance(input_parallel, tuple) + is_fp8 = isinstance(self.quant_method, Fp8LinearMethod) + k_size = ( + input_parallel[0].shape[-1] if is_tuple else input_parallel.shape[-1] + ) + if should_use_tp_invariant_row_linear(k_size): + if output_tensor is not None: + raise RuntimeError( + "tp-invariant row linear cannot write into a caller-owned " + "output tensor" + ) + if is_fp8: + raise NotImplementedError( + "FP8 tp-invariant row-linear not yet supported" + ) + output_parallel = torch.ops.tp_inv_ops.matmul_tp_inv( + input_parallel, self.weight.t(), bias_ + ) + elif output_tensor is None: output_parallel = self.quant_method.apply( self, input_parallel, bias=bias_ ) diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index 6aac6eb2157a..b107900b3901 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -19,6 +19,7 @@ from enum import IntEnum from typing import Any, Dict, List, Optional, Tuple, Union +import numpy as np import torch from torch import nn @@ -58,6 +59,7 @@ ForwardMode, ) from sglang.srt.runtime_context import get_exec, get_parallel +from sglang.srt.true_on_policy import should_force_bfloat16_lm_head from sglang.srt.utils.common import ( is_cpu, is_npu, @@ -94,17 +96,20 @@ class SamplingMaskStatus(IntEnum): @dataclasses.dataclass class SamplingMaskOutput: - """Tensor result for opted-in rows in batch order.""" + """Sampling-support IDs and optional full-support behavior logprobs.""" token_ids: torch.Tensor lengths: torch.Tensor selected_logprobs: torch.Tensor + support_logprobs: Optional[torch.Tensor] statuses: torch.Tensor def map_device_tensors(self, fn) -> None: self.token_ids = fn(self.token_ids) self.lengths = fn(self.lengths) self.selected_logprobs = fn(self.selected_logprobs) + if self.support_logprobs is not None: + self.support_logprobs = fn(self.support_logprobs) self.statuses = fn(self.statuses) @@ -221,11 +226,11 @@ class LogitsProcessorOutput: List[Union[List[float], torch.Tensor]] ] = None next_token_token_ids_logprobs_idx: Optional[List] = None - # Post-filter support IDs, bounded by server capacity, and selected-token - # logprob over the full realized support. + # Post-filter support IDs and requested behavior logprobs, bounded by server + # capacity. Logprobs are normalized over the full realized support. sampling_mask_output: Optional[SamplingMaskOutput] = None - next_token_sampling_mask_idx: Optional[List[Optional[List[int]]]] = None - next_token_sampling_logprobs: Optional[List[Optional[float]]] = None + next_token_sampling_mask_idx: Optional[List[Optional[np.ndarray]]] = None + next_token_sampling_logprobs: Optional[List[Optional[np.ndarray]]] = None next_token_sampling_mask_status: Optional[List[Optional[int]]] = None ## Part 3: Prefill-only. This part will be assigned in python/sglang/srt/layers/logits_processor.py::LogitsProcessor @@ -309,10 +314,15 @@ class LogitsMetadata: @classmethod def from_forward_batch(cls, forward_batch: ForwardBatch): + # MLP-sync may turn an idle rank into a dummy EXTEND for DP prefill + # graphs. It still has no real request whose last token needs logits. + forward_mode = forward_batch.forward_mode + if forward_batch._original_forward_mode == ForwardMode.IDLE: + forward_mode = ForwardMode.IDLE if ( - forward_batch.forward_mode.is_extend() + forward_mode.is_extend() and forward_batch.return_logprob - and not forward_batch.forward_mode.is_target_verify() + and not forward_mode.is_target_verify() ): extend_return_top_logprob = any( x > 0 for x in forward_batch.top_logprobs_nums @@ -340,7 +350,7 @@ def from_forward_batch(cls, forward_batch: ForwardBatch): draft_extend_select_index = None return cls( - forward_mode=forward_batch.forward_mode, + forward_mode=forward_mode, capture_hidden_mode=forward_batch.capture_hidden_mode, next_token_logits_buffer=forward_batch.next_token_logits_buffer, extend_return_logprob=extend_return_logprob, @@ -915,6 +925,11 @@ def _compute_lm_head( hidden_states.to(torch.float32), lm_head.weight.to(torch.float32).T, ) + elif should_force_bfloat16_lm_head(use_fp32_lm_head=self.use_fp32_lm_head): + logits = torch.matmul( + hidden_states.to(torch.bfloat16), + lm_head.weight.to(torch.bfloat16).T, + ).to(torch.bfloat16) elif use_intel_amx_backend(lm_head): logits = torch.ops.sgl_kernel.weight_packed_linear( hidden_states.to(lm_head.weight.dtype), @@ -922,11 +937,6 @@ def _compute_lm_head( None, # bias True, # is_vnni ) - elif self.rl_on_policy_target is not None: - # Due to tie-weight, we may not be able to change lm_head's weight dtype - logits = torch.matmul( - hidden_states.bfloat16(), lm_head.weight.T.bfloat16() - ) else: logits = torch.matmul( hidden_states.to(lm_head.weight.dtype), lm_head.weight.T @@ -1083,6 +1093,8 @@ def _copy_logits_to_buffer( assert logits_buffer.dtype == torch.float logits_buffer.copy_(logits) logits = logits_buffer + elif should_force_bfloat16_lm_head(use_fp32_lm_head=self.use_fp32_lm_head): + logits = logits[:, : self.vocab_size].to(torch.bfloat16) else: logits = logits.float() return logits diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py b/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py index 7f76b2616efc..21be981278d1 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py @@ -131,6 +131,61 @@ def swiglu_gpt_oss_sigmoid_alpha_contiguous( output.copy_(gate * torch.sigmoid(gate * gemm1_alpha) * (up + 1)) +def is_mxfp4_marlin_weights( + num_bits: int, + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + w1_zeros: Optional[torch.Tensor], + w2_zeros: Optional[torch.Tensor], +) -> bool: + return ( + num_bits == 4 + and w1_zeros is None + and w2_zeros is None + and w1_scale.dtype == torch.float8_e8m0fnu + and w2_scale.dtype == torch.float8_e8m0fnu + ) + + +def marlin_moe_use_atomic_add( + hidden_states: torch.Tensor, is_mxfp4_marlin: bool +) -> bool: + # the E8M0 Marlin kernels have no atomic-add path + return ( + hidden_states.dtype == torch.half + or torch.cuda.get_device_capability(hidden_states.device)[0] >= 9 + ) and not is_mxfp4_marlin + + +def gated_marlin_activation( + activation: str, + gemm1_alpha: Optional[float], + clamp_limit: Optional[float], + gate_up: torch.Tensor, + output: torch.Tensor, +) -> None: + """Apply the gated activation of a Marlin MoE layer to `gate_up` (`[..., 2N]`) into `output`.""" + if activation == "silu" and gemm1_alpha is not None: + if clamp_limit is None: + raise ValueError("GPT-OSS Marlin activation requires clamp_limit.") + swiglu_gpt_oss_sigmoid_alpha_contiguous( + output, gate_up, gemm1_alpha, clamp_limit + ) + elif activation == "silu" and clamp_limit is not None: + swiglu_limit_func(output, gate_up, clamp_limit) + elif activation == "silu": + silu_and_mul(gate_up, output) + elif activation == "situ": + situ_and_mul( + output, + gate_up, + situ_beta=gemm1_alpha if gemm1_alpha is not None else 4.0, + linear_beta=clamp_limit, + ) + else: + raise ValueError(f"Unsupported gated activation: {activation=}") + + @register_custom_op(out_shape="hidden_states") def fused_marlin_moe( hidden_states: torch.Tensor, @@ -201,12 +256,8 @@ def fused_marlin_moe( assert w1.is_contiguous(), "Expert weights1 must be contiguous" assert w2.is_contiguous(), "Expert weights2 must be contiguous" assert hidden_states.dtype in [torch.float16, torch.bfloat16] - is_mxfp4_marlin = ( - num_bits == 4 - and w1_zeros is None - and w2_zeros is None - and w1_scale.dtype == torch.float8_e8m0fnu - and w2_scale.dtype == torch.float8_e8m0fnu + is_mxfp4_marlin = is_mxfp4_marlin_weights( + num_bits, w1_scale, w2_scale, w1_zeros, w2_zeros ) is_nvfp4_marlin = ( num_bits == 4 @@ -297,10 +348,7 @@ def fused_marlin_moe( intermediate_cache3 = intermediate_cache13[: M * topk_ids.shape[1] * K] intermediate_cache3 = intermediate_cache3.view(-1, K) - use_atomic_add = ( - hidden_states.dtype == torch.half - or torch.cuda.get_device_capability(hidden_states.device)[0] >= 9 - ) and (not is_mxfp4_marlin) + use_atomic_add = marlin_moe_use_atomic_add(hidden_states, is_mxfp4_marlin) intermediate_cache1 = moe_wna16_marlin_gemm( hidden_states, @@ -331,29 +379,13 @@ def fused_marlin_moe( is_zp_float=False, ) - if activation == "silu" and is_gated and gemm1_alpha is not None: - if clamp_limit is None: - raise ValueError("GPT-OSS Marlin activation requires clamp_limit.") - swiglu_gpt_oss_sigmoid_alpha_contiguous( - intermediate_cache2, - intermediate_cache1.view(-1, gemm1_n), + if is_gated: + gated_marlin_activation( + activation, gemm1_alpha, clamp_limit, - ) - elif activation == "silu" and is_gated and clamp_limit is not None: - swiglu_limit_func( - intermediate_cache2, intermediate_cache1.view(-1, gemm1_n), - clamp_limit, - ) - elif activation == "silu" and is_gated: - silu_and_mul(intermediate_cache1.view(-1, gemm1_n), intermediate_cache2) - elif activation == "situ" and is_gated: - situ_and_mul( intermediate_cache2, - intermediate_cache1.view(-1, gemm1_n), - situ_beta=gemm1_alpha if gemm1_alpha is not None else 4.0, - linear_beta=clamp_limit, ) elif activation == "silu" and not is_gated: intermediate_cache2 = F.silu(intermediate_cache1.view(-1, N)) diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py b/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py index ce665e936444..f283db073e33 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py +++ b/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py @@ -122,7 +122,11 @@ def _validate_fused_swiglu_interleaved( def _use_moe_sum_reduce_torch_compile(num_tokens: int) -> bool: - return num_tokens <= 32 and not is_batch_invariant_mode_enabled() + return ( + num_tokens <= 32 + and not is_batch_invariant_mode_enabled() + and not get_exec().deterministic.enable_deterministic_inference + ) @register_custom_op(mutates_args=["hidden_states"]) diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 36b8aa92d8e9..4294bdafd2ee 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -2107,19 +2107,29 @@ def capture_routed_experts_if_allowed( topk_config: TopKConfig, layer_id: Optional[int], topk_ids: torch.Tensor, + num_token_non_padded: Optional[torch.Tensor] = None, ) -> None: """Single capture site for every backend, gated by the per-config opt-out. Routing all backends through here keeps the draft-side opt-out from being bypassed by an inlined capturer call. + + Capture runs ahead of every padded-region mask in ``_post_process_topk_ids`` + (and ROCm's masks to 0, a valid expert id), so mask a copy here: replay skips + -1 rows but reads a padded one as a real selection. """ if not topk_config.allow_routed_experts_capture: return - if (cap := get_global_experts_capturer()) is not None: - cap.capture( - layer_id=layer_id, - topk_indices=topk_ids, - ) + cap = get_global_experts_capturer() + if cap is None: + return + if num_token_non_padded is not None: + topk_ids = topk_ids.clone() + _mask_topk_ids_padded_region(topk_ids, num_token_non_padded) + cap.capture( + layer_id=layer_id, + topk_indices=topk_ids, + ) def _post_process_topk_ids( @@ -2138,7 +2148,9 @@ def _post_process_topk_ids( fused_shared_experts_scaling_factor = ( topk_config.fused_shared_experts_scaling_factor ) - capture_routed_experts_if_allowed(topk_config, layer_id, topk_ids) + capture_routed_experts_if_allowed( + topk_config, layer_id, topk_ids, num_token_non_padded + ) recorder_topk_ids = None _fold_pad_into_append = False if _is_cuda: diff --git a/python/sglang/srt/layers/on_policy_utils.py b/python/sglang/srt/layers/on_policy_utils.py new file mode 100644 index 000000000000..234397b9b33c --- /dev/null +++ b/python/sglang/srt/layers/on_policy_utils.py @@ -0,0 +1,6 @@ +"""Compatibility imports for SGLang true-on-policy helpers. + +New code should import from :mod:`sglang.srt.true_on_policy`. +""" + +from sglang.srt.true_on_policy import * # noqa: F403 diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py index b6a23177e172..627b842d1294 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py @@ -1243,6 +1243,10 @@ def __init__(self, quantization_config: CompressedTensorsConfig): def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.scheme.process_weights_after_loading(layer) + def restore_weights_before_loading(self, layer: torch.nn.Module) -> None: + if hasattr(layer.scheme, "restore_weights_before_loading"): + layer.scheme.restore_weights_before_loading(layer) + def create_weights( self, layer: torch.nn.Module, diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index b9ca03d8c25f..a1b905787d8b 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -558,11 +558,19 @@ def validate_block_quant_shapes( f"{input_size_per_partition} is not divisible by " f"weight quantization block_k = {block_k}." ) - # Required by column parallel or enabling merged weights - if ( + # Required by column parallel or enabling merged weights. + is_tp_split = ( tp_size > 1 and output_size // output_size_per_partition == tp_size - ) or len(output_partition_sizes) > 1: - for output_partition_size in output_partition_sizes: + ) + is_merged_gemm = len(output_partition_sizes) > 1 + if is_tp_split or is_merged_gemm: + sizes_to_check = output_partition_sizes + if not is_tp_split and is_merged_gemm: + # Match validate_fp8_block_shape: merged weights may have a + # ragged final logical matrix, and scale tensors are already + # allocated with ceil-divided block counts. + sizes_to_check = output_partition_sizes[:-1] + for output_partition_size in sizes_to_check: if output_partition_size % block_n != 0: raise ValueError( f"Weight output_partition_size = " @@ -2311,10 +2319,15 @@ def _prepare_flashinfer_trtllm_activation_params(self, layer: Module) -> None: """Materialize optional TRT-LLM SwiGLU parameters once per expert.""" num_experts = int(layer.num_local_experts) device = layer.w13_weight.device + clamp_limit = ( + self.moe_runner_config.gemm1_clamp_limit + if self.moe_runner_config.gemm1_clamp_limit is not None + else self.moe_runner_config.swiglu_limit + ) for name, value in ( ("gemm1_alpha", self.moe_runner_config.gemm1_alpha), ("gemm1_beta", self.moe_runner_config.gemm1_beta), - ("gemm1_clamp_limit", self.moe_runner_config.gemm1_clamp_limit), + ("gemm1_clamp_limit", clamp_limit), ): tensor = ( None diff --git a/python/sglang/srt/layers/quantization/marlin_utils.py b/python/sglang/srt/layers/quantization/marlin_utils.py index 9b432028d10b..1e3dd5eb7a4a 100644 --- a/python/sglang/srt/layers/quantization/marlin_utils.py +++ b/python/sglang/srt/layers/quantization/marlin_utils.py @@ -300,6 +300,51 @@ def marlin_sort_g_idx(g_idx: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: return g_idx[g_idx_sort_indices], g_idx_sort_indices +def marlin_permute_weights(q_w, size_k, size_n, perm, tile=GPTQ_MARLIN_TILE): + assert q_w.shape == (size_k, size_n) + assert size_k % tile == 0, f"size_k = {size_k}, tile = {tile}" + assert size_n % tile == 0, f"size_k = {size_n}, tile = {tile}" + + # Permute weights to 16x64 marlin tiles + q_w = q_w.reshape((size_k // tile, tile, size_n // tile, tile)) + q_w = q_w.permute((0, 2, 1, 3)) + q_w = q_w.reshape((size_k // tile, size_n * tile)) + + q_w = q_w.reshape((-1, perm.numel()))[:, perm].reshape(q_w.shape) + + return q_w + + +def get_weight_perm(num_bits: int): + perm_list: list[int] = [] + for i in range(32): + perm1: list[int] = [] + col = i // 4 + for block in [0, 1]: + for row in [ + 2 * (i % 4), + 2 * (i % 4) + 1, + 2 * (i % 4 + 4), + 2 * (i % 4 + 4) + 1, + ]: + perm1.append(16 * row + col + 8 * block) + for j in range(4): + perm_list.extend([p + 256 * j for p in perm1]) + + perm = numpy.array(perm_list) + + if num_bits == 4: + interleave = numpy.array([0, 2, 4, 6, 1, 3, 5, 7]) + elif num_bits == 8: + interleave = numpy.array([0, 2, 1, 3]) + else: + raise Exception("num_bits must be 4 or 8, got {}".format(num_bits)) + + perm = perm.reshape((-1, len(interleave)))[:, interleave].ravel() + perm = torch.from_numpy(perm) + return perm + + def get_scale_perms(): scale_perm: list[int] = [] for i in range(8): diff --git a/python/sglang/srt/layers/quantization/marlin_utils_fp4.py b/python/sglang/srt/layers/quantization/marlin_utils_fp4.py index 003282452427..1e99d6710f40 100644 --- a/python/sglang/srt/layers/quantization/marlin_utils_fp4.py +++ b/python/sglang/srt/layers/quantization/marlin_utils_fp4.py @@ -13,6 +13,7 @@ should_use_atomic_add_reduce, ) from sglang.srt.layers.quantization.utils import get_scalar_types +from sglang.srt.layers.utils import copy_or_rebind_param from sglang.srt.utils import is_cuda from sglang.srt.utils.custom_op import register_custom_op @@ -258,7 +259,25 @@ def _get_optional_param(layer: torch.nn.Module, *names: str) -> torch.Tensor | N return None -def deinterleave_moe_mxfp4_w13_for_marlin(layer: torch.nn.Module) -> None: +def _view_tensor_storage( + tensor: torch.Tensor, shape: tuple[int, ...], dtype: torch.dtype +) -> torch.Tensor: + if not tensor.is_contiguous(): + raise ValueError("MXFP4 Marlin in-place reload requires contiguous parameters.") + expected_nbytes = torch.empty((), dtype=dtype).element_size() + for dim in shape: + expected_nbytes *= dim + if tensor.nbytes != expected_nbytes: + raise ValueError( + f"Cannot reuse {tensor.nbytes} bytes as shape={shape}, dtype={dtype} " + f"({expected_nbytes} bytes)." + ) + return tensor.view(torch.uint8).reshape(-1).view(dtype).reshape(shape) + + +def deinterleave_moe_mxfp4_w13_for_marlin( + layer: torch.nn.Module, reuse_parameter_storage: bool = False +) -> None: """Convert GPT-OSS interleaved w13 rows to Marlin's contiguous halves. GPT-OSS stores gate/up rows as [gate0, up0, gate1, up1, ...]. The Marlin @@ -273,9 +292,21 @@ def deinterleave_moe_mxfp4_w13_for_marlin(layer: torch.nn.Module) -> None: raise ValueError(f"Expected even w13 row dimension, got {w13.shape}.") e, n, k = w13.shape - layer.w13_weight.data = ( - w13.view(e, n // 2, 2, k).permute(0, 2, 1, 3).contiguous().view(e, n, k) - ) + # in-place reload works one expert at a time so the scratch stays one expert wide + if reuse_parameter_storage: + for expert_idx in range(e): + deinterleaved = ( + w13[expert_idx] + .view(n // 2, 2, k) + .permute(1, 0, 2) + .contiguous() + .view(n, k) + ) + w13[expert_idx].copy_(deinterleaved) + else: + layer.w13_weight.data = ( + w13.view(e, n // 2, 2, k).permute(0, 2, 1, 3).contiguous().view(e, n, k) + ) if w13_scale is not None: scale = w13_scale.data @@ -283,18 +314,38 @@ def deinterleave_moe_mxfp4_w13_for_marlin(layer: torch.nn.Module) -> None: raise ValueError( f"Expected w13 scale row dimension {n}, got {scale.shape}." ) - w13_scale.data = ( - scale.view(e, n // 2, 2, scale.shape[-1]) - .permute(0, 2, 1, 3) - .contiguous() - .view(e, n, scale.shape[-1]) - ) + if reuse_parameter_storage: + for expert_idx in range(e): + deinterleaved = ( + scale[expert_idx] + .view(n // 2, 2, scale.shape[-1]) + .permute(1, 0, 2) + .contiguous() + .view(n, scale.shape[-1]) + ) + scale[expert_idx].copy_(deinterleaved) + else: + w13_scale.data = ( + scale.view(e, n // 2, 2, scale.shape[-1]) + .permute(0, 2, 1, 3) + .contiguous() + .view(e, n, scale.shape[-1]) + ) if w13_bias is not None: bias = w13_bias.data if bias.shape[1] != n: raise ValueError(f"Expected w13 bias row dimension {n}, got {bias.shape}.") - w13_bias.data = bias.view(e, n // 2, 2).permute(0, 2, 1).contiguous().view(e, n) + if reuse_parameter_storage: + for expert_idx in range(e): + deinterleaved = ( + bias[expert_idx].view(n // 2, 2).permute(1, 0).contiguous().view(n) + ) + bias[expert_idx].copy_(deinterleaved) + else: + w13_bias.data = ( + bias.view(e, n // 2, 2).permute(0, 2, 1).contiguous().view(e, n) + ) def _repack_moe_fp4_weight_for_marlin( @@ -343,7 +394,9 @@ def _permute_moe_fp4_scales_for_marlin( return torch.stack(tensor_list) -def prepare_moe_mxfp4_layer_for_marlin(layer: torch.nn.Module) -> None: +def prepare_moe_mxfp4_layer_for_marlin( + layer: torch.nn.Module, reuse_parameter_storage: bool = False +) -> None: group_size = 32 w13 = layer.w13_weight.data w2 = layer.w2_weight.data @@ -379,7 +432,13 @@ def prepare_moe_mxfp4_layer_for_marlin(layer: torch.nn.Module) -> None: ) device = w13.device - layer.workspace = marlin_make_workspace(device, 4) + if reuse_parameter_storage: + if not hasattr(layer, "workspace"): + raise ValueError( + "MXFP4 Marlin workspace is missing during in-place reload." + ) + else: + layer.workspace = marlin_make_workspace(device, 4) perm = torch.empty(0, dtype=torch.int, device=device) def _pad_w13(x: torch.Tensor) -> torch.Tensor: @@ -398,15 +457,23 @@ def _pad_w2(x: torch.Tensor, packing: int) -> torch.Tensor: x, (0, (padded_intermediate_size - intermediate_size) // packing) ) + # a padded reload would need storage the checkpoint-shaped parameters do not have + if reuse_parameter_storage and padded_intermediate_size != intermediate_size: + raise ValueError( + "MXFP4 Marlin in-place reload requires checkpoint parameters " + "to include Marlin padding." + ) + w13 = _pad_w13(w13) w2 = _pad_w2(w2, packing=2) - w13_scale_data = _pad_w13(_normalize_scale_tensor(w13_scale_data, param_dtype)) - w2_scale_data = _pad_w2( - _normalize_scale_tensor(w2_scale_data, param_dtype), - packing=group_size, - ) - if w13_bias_data is not None: - w13_bias_data = _pad_w13(w13_bias_data.unsqueeze(-1)).squeeze(-1) + if not reuse_parameter_storage: + w13_scale_data = _pad_w13(_normalize_scale_tensor(w13_scale_data, param_dtype)) + w2_scale_data = _pad_w2( + _normalize_scale_tensor(w2_scale_data, param_dtype), + packing=group_size, + ) + if w13_bias_data is not None: + w13_bias_data = _pad_w13(w13_bias_data.unsqueeze(-1)).squeeze(-1) w13_size_n, w13_size_k = padded_intermediate_size * 2, hidden_size w2_size_n, w2_size_k = hidden_size, padded_intermediate_size @@ -422,6 +489,61 @@ def _permute_bias(bias: torch.Tensor | None) -> torch.Tensor | None: tensor_list.append(marlin_permute_bias(bias[i].to(param_dtype))) return torch.stack(tensor_list) + def _transform_parameter_in_place(name: str, transform) -> None: + param = getattr(layer, name) + source = param.data + if source.shape[0] != num_experts: + raise ValueError( + f"Expected {num_experts} experts in {name}, got {source.shape}." + ) + + # expert i keeps its byte range, so writing it only overwrites bytes its transform read + first = transform(source[0]) + output_shape = (num_experts, *first.shape) + output = _view_tensor_storage(source, output_shape, first.dtype) + output[0].copy_(first) + for expert_idx in range(1, num_experts): + output[expert_idx].copy_(transform(source[expert_idx])) + param.data = output + + if reuse_parameter_storage: + + def _repack_expert(size_n: int, size_k: int): + return lambda weight: _repack_moe_fp4_weight_for_marlin( + weight[None], num_experts=1, size_n=size_n, size_k=size_k, perm=perm + )[0] + + def _permute_expert_scale(size_n: int, size_k: int): + return lambda scale: _permute_moe_fp4_scales_for_marlin( + _normalize_scale_tensor(scale, param_dtype)[None], + num_experts=1, + size_n=size_n, + size_k=size_k, + group_size=group_size, + process_scales=_process_scales, + )[0] + + _transform_parameter_in_place( + "w13_weight", _repack_expert(w13_size_n, w13_size_k) + ) + _transform_parameter_in_place("w2_weight", _repack_expert(w2_size_n, w2_size_k)) + _transform_parameter_in_place( + "w13_weight_scale", _permute_expert_scale(w13_size_n, w13_size_k) + ) + _transform_parameter_in_place( + "w2_weight_scale", _permute_expert_scale(w2_size_n, w2_size_k) + ) + if w13_bias_data is not None: + _transform_parameter_in_place( + "w13_weight_bias", + lambda bias: marlin_permute_bias(bias.to(param_dtype)), + ) + if w2_bias_data is not None: + _transform_parameter_in_place( + "w2_weight_bias", lambda bias: marlin_permute_bias(bias.to(param_dtype)) + ) + return + w13_marlin = _repack_moe_fp4_weight_for_marlin( w13, num_experts=num_experts, size_n=w13_size_n, size_k=w13_size_k, perm=perm ) @@ -445,19 +567,16 @@ def _permute_bias(bias: torch.Tensor | None) -> torch.Tensor | None: process_scales=_process_scales, ) - layer.w13_weight = torch.nn.Parameter(w13_marlin, requires_grad=False) - layer.w2_weight = torch.nn.Parameter(w2_marlin, requires_grad=False) - layer.w13_weight_scale = torch.nn.Parameter(w13_scale_marlin, requires_grad=False) - layer.w2_weight_scale = torch.nn.Parameter(w2_scale_marlin, requires_grad=False) + # rebinding keeps the Parameter objects, and with them the weight_loader a reload needs + copy_or_rebind_param(layer, "w13_weight", w13_marlin) + copy_or_rebind_param(layer, "w2_weight", w2_marlin) + copy_or_rebind_param(layer, "w13_weight_scale", w13_scale_marlin) + copy_or_rebind_param(layer, "w2_weight_scale", w2_scale_marlin) if w13_bias_data is not None: - layer.w13_weight_bias = torch.nn.Parameter( - _permute_bias(w13_bias_data), requires_grad=False - ) + copy_or_rebind_param(layer, "w13_weight_bias", _permute_bias(w13_bias_data)) if w2_bias_data is not None: - layer.w2_weight_bias = torch.nn.Parameter( - _permute_bias(w2_bias_data), requires_grad=False - ) + copy_or_rebind_param(layer, "w2_weight_bias", _permute_bias(w2_bias_data)) # Marlin uses the repacked scales; release the loader-format parameters. for stale in ("w13_weight_scale_inv", "w2_weight_scale_inv"): diff --git a/python/sglang/srt/layers/quantization/mxfp4.py b/python/sglang/srt/layers/quantization/mxfp4.py index b94894ea7b8b..697cb2f84b26 100644 --- a/python/sglang/srt/layers/quantization/mxfp4.py +++ b/python/sglang/srt/layers/quantization/mxfp4.py @@ -367,6 +367,15 @@ def get_scaled_act_names(self) -> List[str]: class Mxfp4MoEMethod(FusedMoEMethodBase): + _MARLIN_CHECKPOINT_PARAMETER_NAMES = ( + "w13_weight", + "w13_weight_scale", + "w13_weight_bias", + "w2_weight", + "w2_weight_scale", + "w2_weight_bias", + ) + def __init__( self, prefix: str, @@ -614,6 +623,45 @@ def create_weights( layer.register_parameter("w2_weight_bias", w2_weight_bias) set_weight_attrs(w2_weight_bias, extra_weight_attrs) + if self.use_marlin and not self.use_mega_moe: + layer._mxfp4_marlin_checkpoint_parameter_specs = { + name: (tuple(getattr(layer, name).shape), getattr(layer, name).dtype) + for name in self._marlin_parameter_names(layer) + } + layer._mxfp4_marlin_reload_in_place = False + + def _marlin_parameter_names(self, layer): + # the biases only exist when the checkpoint carries them + return [ + name + for name in self._MARLIN_CHECKPOINT_PARAMETER_NAMES + if hasattr(layer, name) + ] + + def restore_weights_before_loading(self, layer): + if not (self.use_marlin and not self.use_mega_moe): + return + + from sglang.srt.layers.quantization.marlin_utils_fp4 import ( + _view_tensor_storage, + ) + + specs = layer._mxfp4_marlin_checkpoint_parameter_specs + runtime_specs = layer._mxfp4_marlin_runtime_parameter_specs + for name in self._marlin_parameter_names(layer): + shape, dtype = specs[name] + param = getattr(layer, name) + current_spec = (tuple(param.shape), param.dtype) + if current_spec not in ((shape, dtype), runtime_specs[name]): + raise ValueError( + f"Unexpected MXFP4 Marlin parameter state for {name}: " + f"shape={param.shape}, dtype={param.dtype}." + ) + # the loader fills only the unpadded slice, so the padding must start at zero + param.data = _view_tensor_storage(param.data, shape, dtype) + param.data.zero_() + layer._mxfp4_marlin_reload_in_place = True + def process_weights_after_loading(self, layer): if self.use_marlin and not self.use_mega_moe: from sglang.srt.layers.quantization.marlin_utils import ( @@ -635,11 +683,29 @@ def process_weights_after_loading(self, layer): "Current MXFP4 MoE layer is not supported by Marlin." ) + reuse_parameter_storage = layer._mxfp4_marlin_reload_in_place if self.moe_runner_config.gemm1_alpha is not None and getattr( self.moe_runner_config, "gate_up_interleaved", True ): - deinterleave_moe_mxfp4_w13_for_marlin(layer) - prepare_moe_mxfp4_layer_for_marlin(layer) + deinterleave_moe_mxfp4_w13_for_marlin( + layer, reuse_parameter_storage=reuse_parameter_storage + ) + prepare_moe_mxfp4_layer_for_marlin( + layer, reuse_parameter_storage=reuse_parameter_storage + ) + runtime_specs = { + name: (tuple(getattr(layer, name).shape), getattr(layer, name).dtype) + for name in self._marlin_parameter_names(layer) + } + if reuse_parameter_storage: + if runtime_specs != layer._mxfp4_marlin_runtime_parameter_specs: + raise ValueError( + "MXFP4 Marlin in-place reload changed the runtime layout: " + f"{runtime_specs} vs {layer._mxfp4_marlin_runtime_parameter_specs}." + ) + else: + layer._mxfp4_marlin_runtime_parameter_specs = runtime_specs + layer._mxfp4_marlin_reload_in_place = False layer._mxfp4_backend = "marlin" return @@ -1481,6 +1547,15 @@ def _apply_sm90_cutlass(self, layer, dispatch_output): ) return self.runner.run(dispatch_output, quant_info) + def get_marlin_quant_info(self, layer): + # consumed by FusedMoEWithLoRA's marlin branch + from sglang.srt.layers.quantization.mxfp4_marlin_moe import ( + build_marlin_moe_quant_info, + ) + + assert self.use_marlin, "Marlin quant info requires the Marlin backend" + return build_marlin_moe_quant_info(layer) + def _apply_marlin(self, layer, dispatch_output): """MXFP4 x BF16 MoE via the Marlin runner. The quant_info (incl. the dispatcher's EP mapping) is shared with diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index c0dc373a4bc9..00caa048b9b4 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/python/sglang/srt/layers/quantization/unquant.py @@ -906,7 +906,7 @@ def maybe_restore_flashinfer_trtllm_bf16_weight_shape_for_load( block layout. During weight update, checkpoint tensors are in canonical layout and need a temporary shape restore for copy. """ - if not get_moe_runner_backend().is_flashinfer_trtllm_routed(): + if not self.use_flashinfer_trtllm_moe: return expected_shape = None diff --git a/python/sglang/srt/layers/rotary_embedding/base.py b/python/sglang/srt/layers/rotary_embedding/base.py index a8bda5656bdd..d0e7b0f3b4ca 100644 --- a/python/sglang/srt/layers/rotary_embedding/base.py +++ b/python/sglang/srt/layers/rotary_embedding/base.py @@ -11,7 +11,8 @@ from sglang.srt.environ import envs from sglang.srt.layers.rotary_embedding.utils import apply_rotary_emb from sglang.srt.platforms import current_platform -from sglang.srt.runtime_context import get_exec, publish_role +from sglang.srt.runtime_context import publish_role +from sglang.srt.true_on_policy import is_true_on_policy_enabled from sglang.srt.utils import ( cpu_has_amx_support, get_bool_env_var, @@ -102,10 +103,7 @@ def __init__( self.base = base self.is_neox_style = is_neox_style self.dtype = dtype - self._force_native = ( - publish_role() is not None - and get_exec().deterministic.rl_on_policy_target is not None - ) + self._force_native = publish_role() is not None and is_true_on_policy_enabled() cache = self._compute_cos_sin_cache() # NOTE(ByronHsu): cache needs to be in FP32 for numerical stability. @@ -143,10 +141,6 @@ def __init__( # XXX (MUSA): Implement sgl_kernel.rotary_embedding support for MUSA backend if self._force_native or _is_musa: self._forward_method = self.forward_native - self._apply_rotary_emb_wrapped = torch.compile( - dynamic=True, - disable=_is_npu, - )(apply_rotary_emb) self.position_cos, self.position_sin = None, None def _match_cos_sin_cache_dtype(self, query: torch.Tensor) -> None: diff --git a/python/sglang/srt/layers/rotary_embedding/mrope.py b/python/sglang/srt/layers/rotary_embedding/mrope.py index 777edc7c5e09..646c56145abb 100644 --- a/python/sglang/srt/layers/rotary_embedding/mrope.py +++ b/python/sglang/srt/layers/rotary_embedding/mrope.py @@ -104,6 +104,10 @@ def __init__( ) self.register_buffer("axis_map", self._build_axis_map(), persistent=False) + if self.axis_map is not None: + # A lane-to-axis lookup derived from mrope_section, not a weight: the RL + # weight checker must neither poison nor compare it. + self.axis_map._skip_weight_check = True if self._force_native: self._forward_method = self.forward_native diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py index d5ee8d8267e6..32c002b58727 100644 --- a/python/sglang/srt/layers/sampler.py +++ b/python/sglang/srt/layers/sampler.py @@ -23,6 +23,7 @@ from sglang.srt.runtime_context import get_exec, get_parallel, get_server_args from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo from sglang.srt.sampling.sampling_params import TOP_K_ALL +from sglang.srt.true_on_policy import resolve_true_on_policy_runtime_policy from sglang.srt.utils.async_probe import sanitize_nan_logits from sglang.srt.utils.common import ( get_bool_env_var, @@ -115,13 +116,13 @@ def __init__(self): self.tp_sync_group = get_parallel().attn_tp_group.device_group self.cp_sync_group = get_parallel().attn_cp_group.device_group - self.rl_on_policy_target = get_exec().deterministic.rl_on_policy_target + true_on_policy = resolve_true_on_policy_runtime_policy(get_server_args()) # In RL on-policy mode, deterministic inference is automatically enabled. self.enable_deterministic = ( get_exec().deterministic.enable_deterministic_inference ) # In RL on-policy mode, we use log_softmax to compute logprobs to match the trainer. - self.use_log_softmax_logprob = self.rl_on_policy_target is not None + self.use_log_softmax_logprob = true_on_policy.enabled self.use_ascend_backend = get_exec().kernel.sampling_backend == "ascend" self.sampling_mask_max_tokens = get_exec().features.sampling_mask_max_tokens @@ -175,6 +176,9 @@ def forward( _trace_e2e_sampler("preprocess_returned") sampling_mask_batch_indices = sampling_info.sampling_mask_batch_indices return_sampling_mask = sampling_mask_batch_indices is not None + sampling_support_logprobs_capture_indices = ( + sampling_info.sampling_support_logprobs_capture_indices + ) sampling_mask_capture = None if sampling_info.is_all_greedy: @@ -204,17 +208,13 @@ def forward( if return_logprob and SGLANG_RETURN_ORIGINAL_LOGPROB: original_logprobs = torch.log_softmax(logits, dim=-1) + # Post process logits + logits.div_(sampling_info.temperatures) + # In RL on-policy mode, we use log_softmax to compute logprobs to match the trainer. logprobs_via_logsoftmax_kernel = None - if self.rl_on_policy_target is not None: - # TODO: use more inplace ops to save memory - logits_div_temperature = ( - logits.bfloat16().div(sampling_info.temperatures).bfloat16() - ) - logprobs_via_logsoftmax_kernel = torch.log_softmax( - logits_div_temperature, dim=-1 - ) - del logits_div_temperature + if self.use_log_softmax_logprob: + logprobs_via_logsoftmax_kernel = torch.log_softmax(logits, dim=-1) if self.use_ascend_backend: # Ascend backend: sample from logits directly. @@ -240,8 +240,6 @@ def forward( logprobs = logprobs_via_logsoftmax_kernel else: # Standard path: do softmax and sample from probs. - logits.div_(sampling_info.temperatures) - # Deterministic inference must derive the returned logprobs # from F.log_softmax — the same kernel prefill rescoring uses — # not log(softmax(x)) below: the two disagree at ~1e-6 despite @@ -295,7 +293,9 @@ def forward( if sampling_info.is_all_greedy: logits_output.sampling_mask_output = ( self._build_greedy_sampling_mask_output( - sampling_mask_batch_indices, batch_next_token_ids + sampling_mask_batch_indices, + batch_next_token_ids, + sampling_support_logprobs_capture_indices, ) ) else: @@ -308,6 +308,7 @@ def forward( logits_output.sampling_mask_output = self._build_sampling_mask_output( batch_next_token_ids, sampling_mask_capture, + sampling_support_logprobs_capture_indices, ) _trace_e2e_sampler("forward_returned") @@ -460,6 +461,7 @@ def _build_greedy_sampling_mask_output( self, batch_indices: torch.Tensor, batch_next_token_ids: torch.Tensor, + support_capture_indices: Optional[torch.Tensor], ) -> SamplingMaskOutput: token_ids = batch_next_token_ids.index_select(0, batch_indices).to(torch.int32) num_requests = batch_indices.numel() @@ -471,6 +473,15 @@ def _build_greedy_sampling_mask_output( selected_logprobs=torch.zeros( num_requests, dtype=torch.float32, device=token_ids.device ), + support_logprobs=( + torch.zeros( + (support_capture_indices.numel(), 1), + dtype=torch.float32, + device=token_ids.device, + ) + if support_capture_indices is not None + else None + ), statuses=torch.full( (num_requests,), SamplingMaskStatus.OK, @@ -483,6 +494,7 @@ def _build_sampling_mask_output( self, batch_next_token_ids: torch.Tensor, sampling_mask_capture: _SamplingMaskCapture, + support_capture_indices: Optional[torch.Tensor], ) -> SamplingMaskOutput: """Pack captured positive support into the fixed-cap device result.""" batch_indices = sampling_mask_capture.batch_rows @@ -548,18 +560,27 @@ def _build_sampling_mask_output( packed_size = min(self.sampling_mask_max_tokens, weights.shape[-1]) if token_ids is None: - _, packed_positions = torch.topk( + packed_weights, packed_positions = torch.topk( weights, k=packed_size, dim=-1, largest=True, sorted=True ) packed_token_ids = packed_positions.to(torch.int32) else: # The PyTorch producer already sorts weights and IDs together. packed_token_ids = token_ids[:, :packed_size].contiguous() + packed_weights = weights[:, :packed_size] + + support_logprobs = None + if support_capture_indices is not None: + support_logprobs = torch.log( + packed_weights.index_select(0, support_capture_indices).float() + / support_mass.index_select(0, support_capture_indices).unsqueeze(-1) + ) return SamplingMaskOutput( token_ids=packed_token_ids, lengths=realized_lengths.clamp(max=packed_size), selected_logprobs=selected_logprobs, + support_logprobs=support_logprobs, statuses=statuses, ) @@ -629,13 +650,13 @@ def _forward_ascend_backend( """Handle the full Ascend backend sampling path. Ascend backend has fused kernels that handle softmax internally, - so we sample directly from temperature-scaled logits. + so we sample directly from temperature-scaled logits. Temperature + scaling is already applied by the caller before branch dispatch. Returns: A tuple of (batch_next_token_ids, logprobs). logprobs is None when return_logprob is False or SGLANG_RETURN_ORIGINAL_LOGPROB is set. """ - logits.div_(sampling_info.temperatures) batch_next_token_ids = self._sample_from_logits( logits, sampling_info, simple_sampling_case, positions ) diff --git a/python/sglang/srt/lora/backend/base_backend.py b/python/sglang/srt/lora/backend/base_backend.py index 99c2bc69079c..e36e4fa68096 100644 --- a/python/sglang/srt/lora/backend/base_backend.py +++ b/python/sglang/srt/lora/backend/base_backend.py @@ -11,6 +11,39 @@ get_batch_token_counts, ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.utils.common import ceil_align + + +def get_gathered_moe_num_tokens(forward_batch: ForwardBatch, num_tokens: int) -> int: + """Token count the MoE-LoRA mapping must cover: gathered under --enable-dp-attention, else per-rank. + + In the eager path prepare_lora_batch runs from ForwardBatch.init_new, BEFORE + prepare_mlp_sync_batch assigns forward_batch.global_dp_buffer_len (only the cuda-graph + capture path pre-sets it), so that field alone under-sizes the MoE-LoRA token mapping to the + per-rank length and the MoE-LoRA kernels index past it (sticky CUDA IMA). When it is unset, + derive an upper bound of the gathered length from global_num_tokens_cpu (assigned in + init_new before prepare_lora_batch runs), mirroring prepare_mlp_sync_batch's attn-tp/cp + alignment; max*n covers both SUM_LEN and MAX_LEN padding modes. Over-allocation is harmless: + the kernels index at most the actual gathered length. + """ + if forward_batch.global_dp_buffer_len is not None: + return max(forward_batch.global_dp_buffer_len, num_tokens) + global_num_tokens = getattr(forward_batch, "global_num_tokens_cpu", None) + if not global_num_tokens: + return num_tokens + # Local import: a module-level cp import here is circular (see forward_batch_info). + # v0.5.16 moved this helper layers.utils.cp_utils -> layers.cp.padding. + from sglang.srt.layers.cp.padding import get_cp_padding_align_size + from sglang.srt.runtime_context import get_parallel + + # v0.5.16 retired dp_attention.get_attention_tp_size(); it lives on ParallelState now. + attn_tp_size = get_parallel().attn_tp_size + cp_align_size = get_cp_padding_align_size() + upper = max( + ceil_align(ceil_align(t, attn_tp_size), cp_align_size) + for t in global_num_tokens + ) * len(global_num_tokens) + return max(upper, num_tokens) class BaseLoRABackend(LoRABackendLmHeadMixing): @@ -233,11 +266,26 @@ def init_cuda_graph_moe_buffers( """ base = moe_layer.base_layer top_k = base.top_k - device = moe_layer._quant_info.w13_weight.device + # Take the device from the base FusedMoE rather than quant-specific tensors, + # so this works for any scheme (FP, WNA16, Marlin-packed, etc.). + device = next(base.parameters()).device num_experts = base.num_experts + # Under --enable-dp-attention the MoE runs on DP-GATHERED tokens (the global DP buffer of + # length up to max_bs * attn_dp_size), not the per-rank batch. The per-token LoRA routing + # buffers below are indexed by that gathered token count, so size them for the gathered + # maximum; otherwise the MoE-LoRA kernels read token_lora_mapping / sorted_token_ids past a + # per-rank-sized buffer -> cudaErrorIllegalInstruction during cuda-graph capture. Expert- and + # adapter-indexed buffers (cumsum_buffer, adapter_enabled, lora_ids) are unaffected. + from sglang.srt.layers.dp_attention import get_attention_dp_size + + dp_size = max(1, get_attention_dp_size()) + max_moe_tokens = max_bs * dp_size + block_size_m = 64 - max_num_tokens_padded = max_bs * top_k + num_experts * (block_size_m - 1) + max_num_tokens_padded = max_moe_tokens * top_k + num_experts * ( + block_size_m - 1 + ) max_num_tokens_padded = ( (max_num_tokens_padded + block_size_m - 1) // block_size_m ) * block_size_m @@ -265,12 +313,12 @@ def init_cuda_graph_moe_buffers( device=device, ), "token_mask": torch.empty( - (max_loras * max_bs * top_k,), + (max_loras * max_moe_tokens * top_k,), dtype=torch.int32, device=device, ), "token_lora_mapping": torch.full( - (max_bs,), -1, dtype=torch.int32, device=device + (max_moe_tokens,), -1, dtype=torch.int32, device=device ), } @@ -322,6 +370,19 @@ def _add_moe_lora_info( seg_indptr = batch_info.seg_indptr[: num_moe_segments + 1] req_to_lora = batch_info.weight_indices[:num_moe_segments] + # --enable-dp-attention all-gathers tokens into the MoE, so the MoE-LoRA kernels index + # token_lora_mapping by the GATHERED token count, not the per-rank num_tokens. Size the + # mapping to (an upper bound of) the gathered count so those reads stay in-bounds — see + # get_gathered_moe_num_tokens for why global_dp_buffer_len alone is NOT enough in the + # eager path (the per-rank segments still fill only [0, num_tokens); the tail stays -1). + moe_num_tokens = get_gathered_moe_num_tokens(forward_batch, num_tokens) + if batch_info.use_cuda_graph: + # Static capture buffers hold max_bs*dp tokens; a REAL replay's gathered length never + # exceeds that (the captured graph could not address it), so cap the upper bound at + # the buffer size. Batches whose gathered bound exceeds it are demoted to the eager + # prep path in LoRAManager.prepare_lora_batch before we get here. + moe_num_tokens = min(moe_num_tokens, token_lora_mapping.shape[0]) + adapter_enabled, token_lora_mapping = _compute_moe_lora_info( num_tokens, seg_indptr, @@ -330,8 +391,34 @@ def _add_moe_lora_info( adapter_enabled, token_lora_mapping, max_len=max_len, + mapping_len=moe_num_tokens, ) + # Tier-1 (colocate RL, exactly one adapter loaded): the DP-gathered tail + # [num_tokens, moe_num_tokens) covers OTHER dp ranks' tokens, which under colocate RL all + # use that single adapter. The per-rank fill above only wrote [0, num_tokens), leaving the + # tail -1 (adapter-disabled) -> cross-rank gathered tokens would miss the LoRA delta on + # this rank's local experts. The stamps below are ARMED BY LoRAManager.prepare_lora_batch + # (policy lives there; both are host ints, no GPU sync -> cuda-graph safe, written each + # batch in the eager prep path): + # * _single_loaded_buffer_id: armed only when exactly one adapter is loaded AND this + # rank's local requests actively use it -> stamp the tail with that adapter. + # * _idle_rank_active_buffer_id: armed only on a true idle rank (num_tokens == 0) -> + # _compute_moe_lora_info left the whole mapping -1 / adapter_enabled all-0, so stamp + # the whole gathered buffer and enable the adapter. + # Multi-adapter batches and base-only local batches arm neither stamp and keep the -1 + # tail: foreign tokens get base rather than a delta this rank cannot attribute. + if moe_num_tokens > num_tokens: + if num_tokens > 0: + single_bid = getattr(self, "_single_loaded_buffer_id", None) + if single_bid is not None: + token_lora_mapping[num_tokens:moe_num_tokens].fill_(single_bid) + else: + idle_bid = getattr(self, "_idle_rank_active_buffer_id", None) + if idle_bid is not None: + token_lora_mapping.fill_(idle_bid) + adapter_enabled[idle_bid] = 1 + batch_info.moe_lora_info = MoELoRABatchInfo( seg_indptr=seg_indptr, req_to_lora=req_to_lora, @@ -412,19 +499,31 @@ def _compute_moe_lora_info( adapter_enabled: torch.Tensor | None, token_lora_mapping: torch.Tensor | None, max_len: int, + mapping_len: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: + # ``num_tokens`` is the PER-RANK fill count (segments cover this rank's tokens). ``mapping_len`` + # is the length of the token_lora_mapping the MoE-LoRA kernels actually index -- under + # --enable-dp-attention the MoE runs on DP-GATHERED tokens (mapping_len = global_dp_buffer_len + # >= num_tokens), so the returned mapping must span the gathered count to keep those kernels + # in-bounds. The DP-gathered tail [num_tokens, mapping_len) defaults to -1 (adapter-disabled). + if mapping_len is None: + mapping_len = num_tokens + assert mapping_len >= num_tokens if token_lora_mapping is not None: - assert num_tokens <= token_lora_mapping.shape[0], ( - "num_tokens must be less than or equal to the shape of token_lora_mapping" + assert mapping_len <= token_lora_mapping.shape[0], ( + "mapping_len must be less than or equal to the shape of token_lora_mapping" ) # Clear padded replay rows left by a larger batch. - if num_tokens < token_lora_mapping.shape[0]: - token_lora_mapping[num_tokens:].fill_(-1) - token_lora_mapping = token_lora_mapping[:num_tokens] + if mapping_len < token_lora_mapping.shape[0]: + token_lora_mapping[mapping_len:].fill_(-1) + token_lora_mapping = token_lora_mapping[:mapping_len] else: token_lora_mapping = torch.empty( - (num_tokens,), dtype=torch.int32, device=seg_indptr.device + (mapping_len,), dtype=torch.int32, device=seg_indptr.device ) + if mapping_len > num_tokens: + # clean the gathered tail before the per-rank fill writes [0, num_tokens) + token_lora_mapping.fill_(-1) if adapter_enabled is not None: assert len(lora_ranks) <= adapter_enabled.shape[0], ( @@ -484,8 +583,12 @@ def _compute_moe_lora_info( torch.searchsorted(seg_indptr.to(torch.int32), token_positions, right=True) - 1 ) - token_lora_mapping = torch.index_select( - weight_indices.to(torch.int32), 0, req_indices, out=token_lora_mapping + # Fill only the per-rank prefix [0, num_tokens); the gathered tail keeps the -1 set above. + torch.index_select( + weight_indices.to(torch.int32), + 0, + req_indices, + out=token_lora_mapping[:num_tokens], ) return adapter_enabled, token_lora_mapping diff --git a/python/sglang/srt/lora/layers.py b/python/sglang/srt/lora/layers.py index 1e0e604bb996..07a049375e4a 100644 --- a/python/sglang/srt/lora/layers.py +++ b/python/sglang/srt/lora/layers.py @@ -56,6 +56,13 @@ def __init__( # LoRA-wrapped. if hasattr(self.base_layer, "reduce_results"): self.reduce_results = self.base_layer.reduce_results + # Alias remaining base-layer parameters onto the wrapper so + # `named_parameters(remove_duplicate=True)` yields them at the outer + # path — weight loaders (e.g. FusedMoE's `w13_weight_packed`) lookup + # names without the `.base_layer.` segment. + for _name, _param in base_layer.named_parameters(recurse=False): + if not hasattr(self, _name): + setattr(self, _name, _param) def forward(self, x: torch.Tensor): return self.base_layer.forward(x) @@ -813,6 +820,10 @@ def forward(self, input_: torch.Tensor, skip_all_reduce=False, forward_batch=Non ) input_parallel = splitted_input[tp_rank].contiguous() + # K3 gates o_proj's input in a forward patch this path bypasses; apply it here + if hasattr(self.base_layer, "lora_input_transform"): + input_parallel = self.base_layer.lora_input_transform(input_parallel) + lora_active = self.lora_active if lora_active: self.start_lora_a_overlap(input_parallel) @@ -990,6 +1001,13 @@ def __init__( self.experts_shared_outer_loras: bool = False self.lora_use_virtual_experts: bool = False + # Forward for the model's own forward-path dispatch — the outer model + # reads several FusedMoE attributes (e.g. `self.experts.moe_runner_config`, + # `self.experts.dispatcher`, `self.experts.num_local_experts`, + # `self.experts.quant_method`) directly on the wrapper. Quant + # post-processing iterators skip LoRA wrappers via + # `isinstance(module, BaseLayerWithLoRA)` so the packed params on the + # inner FusedMoE get processed there, not here. self.quant_method = base_layer.quant_method self.moe_runner_config = base_layer.moe_runner_config # Don't let the MoE runner overwrite hidden_states with its output: @@ -1002,15 +1020,18 @@ def __init__( self.should_fuse_routed_scaling_factor_in_topk = ( base_layer.should_fuse_routed_scaling_factor_in_topk ) + if hasattr(base_layer, "scheme"): + self.scheme = base_layer.scheme self.tp_size = base_layer.moe_tp_size self.tp_rank = base_layer.moe_tp_rank self.intermediate_size_per_partition = ( base_layer.intermediate_size_per_partition ) - # Stock MoE LoRA buffers are split gate/up except for GPT-OSS-style weights. + # K3 sets gemm1_alpha (SiTU beta) but loads w1/w3 contiguously (gate_up_interleaved=False) self._uses_interleaved_gate_up = ( base_layer.moe_runner_config.gemm1_alpha is not None + and base_layer.moe_runner_config.gate_up_interleaved ) # Initialize triton_lora moe runner for batches with lora enabled @@ -1064,17 +1085,7 @@ def __init__( ) if runner_backend.is_marlin(): - from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors import ( - CompressedTensorsFusedMoEMethod, - ) - from sglang.srt.layers.quantization.modelopt_quant import ( - ModelOptNvFp4FusedMoEMethod, - ) - - assert isinstance( - base_layer.quant_method, - (CompressedTensorsFusedMoEMethod, ModelOptNvFp4FusedMoEMethod), - ), ( + assert hasattr(base_layer.quant_method, "get_marlin_quant_info"), ( f"Marlin MoE backend requires a quant method exposing " f"get_marlin_quant_info, got {type(base_layer.quant_method).__name__}" ) @@ -1192,6 +1203,9 @@ def _forward_with_lora( # Use pre-computed quant info (doesn't change so not sure why we need to pass it in every time) quant_info = self._quant_info + quant_info.expert_map = getattr( + base_layer.dispatcher, "local_expert_mapping", None + ) # ===== TO BE REFACTORED ==== if self._lora_runner_backend.is_experimental_sgl_trtllm(): diff --git a/python/sglang/srt/lora/lora.py b/python/sglang/srt/lora/lora.py index c154e1c5581b..e0230377067e 100644 --- a/python/sglang/srt/lora/lora.py +++ b/python/sglang/srt/lora/lora.py @@ -29,6 +29,7 @@ from sglang.srt.layers.utils import get_layer_id from sglang.srt.lora.backend.base_backend import BaseLoRABackend from sglang.srt.lora.lora_config import LoRAConfig +from sglang.srt.lora.utils import get_normalized_target_modules from sglang.srt.model_loader.loader import DefaultModelLoader from sglang.srt.utils.hf_transformers_utils import AutoConfig @@ -60,10 +61,15 @@ def __init__( load_config: LoadConfig, lora_backend: BaseLoRABackend, base_model: Optional[torch.nn.Module] = None, + keep_weights_on_device: bool = False, ): super().__init__() self.uid: str = uid + self.keep_weights_on_device: bool = keep_weights_on_device self.config: LoRAConfig = config + self.normalized_target_modules = get_normalized_target_modules( + config.target_modules, base_model=base_model + ) assert self.config.hf_config["peft_type"].lower() == "lora" self.base_hf_config: AutoConfig = base_hf_config self.load_config: LoadConfig = load_config @@ -161,13 +167,23 @@ def initialize_weights_from_tensors(self, tensors: Dict[str, torch.Tensor]): self._normalize_weights() - def _process_weight(self, name: str, loaded_weight: torch.Tensor): - from sglang.srt.lora.utils import get_normalized_target_modules + def _stage_weight(self, loaded_weight: torch.Tensor) -> torch.Tensor: + # streamed weights are IPC views into the trainer's buckets; a full-adapter clone would not fit + if self.keep_weights_on_device and loaded_weight.is_cuda: + return loaded_weight + return loaded_weight.cpu() - normalized_target_modules = get_normalized_target_modules( - self.config.target_modules - ) + def release_staged_weights(self): + """Drop the staged copies once installed in the memory pool; the adapter cannot be reinstalled.""" + for layer in self.layers: + layer.weights.clear() + layer.pinned_weights.clear() + self.embedding_layers.clear() + self.pinned_embedding_layers.clear() + self.added_tokens_embeddings.clear() + self.pinned_added_tokens_embeddings.clear() + def _process_weight(self, name: str, loaded_weight: torch.Tensor): # Remap PEFT "unembed_tokens" key to "lm_head" so the weight is # recognized and loaded into the correct buffer. if "unembed_tokens" in name: @@ -175,7 +191,7 @@ def _process_weight(self, name: str, loaded_weight: torch.Tensor): layer_id = get_layer_id(name) if layer_id is not None: - self.layers[layer_id].weights[name] = loaded_weight.cpu() + self.layers[layer_id].weights[name] = self._stage_weight(loaded_weight) elif "embed_tokens" in name or "lm_head" in name: # Check if this module is declared in target_modules before loading. # When normalized_target_modules is {"all"} (e.g. target_modules was @@ -183,17 +199,17 @@ def _process_weight(self, name: str, loaded_weight: torch.Tensor): # --lora-target-modules will govern which modules are active. module_name = "embed_tokens" if "embed_tokens" in name else "lm_head" if ( - "all" in normalized_target_modules - or module_name in normalized_target_modules + "all" in self.normalized_target_modules + or module_name in self.normalized_target_modules ): - self.embedding_layers[name] = loaded_weight.cpu() + self.embedding_layers[name] = self._stage_weight(loaded_weight) else: logger.debug( f"Skipping {name} as '{module_name}' is not in adapter's target_modules: {self.config.target_modules}" ) elif "input_embeddings" in name or "output_embeddings" in name: # added/extra token emb - self.added_tokens_embeddings[name] = loaded_weight.cpu() + self.added_tokens_embeddings[name] = self._stage_weight(loaded_weight) assert loaded_weight.shape[0] == self.config.lora_added_tokens_size, ( f"LoRA adapter {self.uid} has lora_added_tokens_size {self.config.lora_added_tokens_size} specified in the config, " f"but the loaded weight '{name}' has shape {loaded_weight.shape[0]} in first dimension" @@ -210,6 +226,8 @@ def _normalize_weights(self): self._normalize_in_proj(layer.weights) # Stack in_proj_q + in_proj_k + in_proj_v + in_proj_z → in_proj_qkvz for GDN layers self._normalize_in_proj_qkvz(layer.weights) + # Stack in_proj_b + in_proj_a → in_proj_ba for GDN layers + self._normalize_in_proj_ba(layer.weights) weight_names = list(layer.weights.keys()) self.normalize_gate_up_proj(weight_names, layer.weights) weight_names = list(layer.weights.keys()) @@ -450,6 +468,25 @@ def _normalize_in_proj_qkvz(self, weights: Dict[str, torch.Tensor]): weights.pop(k_name) weights.pop(v_name) weights.pop(z_name) + elif "in_proj_qkv." in weight_name: + # 2-way split (Megatron-Bridge adapter export): in_proj_qkv + # covers the q|k|v rows and in_proj_z the z rows, with one + # shared lora_A. The stacked buffer expects one A block per + # slice, so repeat the qkv A block 3x (q, k, v) before + # appending the z block; B rows concatenate directly. + z_name = weight_name.replace("in_proj_qkv.", "in_proj_z.") + if z_name not in weights: + continue + qkvz_name = weight_name.replace("in_proj_qkv.", "in_proj_qkvz.") + cat_dim = weights[weight_name].dim() - 2 + qkv_w = weights[weight_name] + if "lora_A" in weight_name: + repeat_dims = [1] * qkv_w.dim() + repeat_dims[cat_dim] = 3 + qkv_w = qkv_w.repeat(*repeat_dims) + weights[qkvz_name] = torch.cat((qkv_w, weights[z_name]), cat_dim) + weights.pop(weight_name) + weights.pop(z_name) elif "in_proj_qkvz" in weight_name and "lora_A" in weight_name: # Already-merged adapter: replicate the shared A across the 4 # stacked slots the buffer expects (q, k, v, z). @@ -459,6 +496,45 @@ def _normalize_in_proj_qkvz(self, weights: Dict[str, torch.Tensor]): weights[weight_name] = weights[weight_name].repeat(*repeat_dims) # else (in_proj_qkvz lora_B, or unrelated): no-op. + def _normalize_in_proj_ba(self, weights: Dict[str, torch.Tensor]): + """Normalize in_proj_ba weights for GDN (GatedDeltaNet) layers like + Qwen3.5. + + Two adapter formats are handled: + + 1. Split: ``in_proj_b + in_proj_a`` (HF checkpoint naming, also the + Megatron-Bridge adapter export) are present as separate weights → + concatenate them into ``in_proj_ba`` (B rows b|a; A blocks b, a). + + 2. Already-merged: the adapter has a single ``in_proj_ba`` weight + (PEFT trained against SGLang's fused Linear). The stacked buffer + expects two per-slice ``A`` blocks, so repeat ``lora_A`` 2x along + the rank dim. ``lora_B`` is already full-output-dim and matches + the buffer directly. + """ + for weight_name in list(weights.keys()): + # NB: match with the trailing dot so the merged "in_proj_ba." + # names don't take the split branch. + if "in_proj_b." in weight_name: + a_name = weight_name.replace("in_proj_b.", "in_proj_a.") + if a_name not in weights: + continue + ba_name = weight_name.replace("in_proj_b.", "in_proj_ba.") + cat_dim = weights[weight_name].dim() - 2 + weights[ba_name] = torch.cat( + (weights[weight_name], weights[a_name]), cat_dim + ) + weights.pop(weight_name) + weights.pop(a_name) + elif "in_proj_ba" in weight_name and "lora_A" in weight_name: + # Already-merged adapter: replicate the shared A across the 2 + # stacked slots the buffer expects (b, a). + ndim = weights[weight_name].dim() + repeat_dims = [1] * ndim + repeat_dims[ndim - 2] = 2 + weights[weight_name] = weights[weight_name].repeat(*repeat_dims) + # else (in_proj_ba lora_B, or unrelated): no-op. + def normalize_gate_up_proj( self, weight_names: List[str], weights: Dict[str, torch.Tensor] ): diff --git a/python/sglang/srt/lora/lora_manager.py b/python/sglang/srt/lora/lora_manager.py index 682116f4283a..a31f785b2220 100644 --- a/python/sglang/srt/lora/lora_manager.py +++ b/python/sglang/srt/lora/lora_manager.py @@ -37,6 +37,7 @@ from sglang.srt.lora.lora_registry import LoRARef from sglang.srt.lora.mem_pool import LoRAMemoryPool from sglang.srt.lora.utils import ( + _KNOWN_LORA_TARGET_MODULES, DSA_INDEXER_LORA_NAMES, EMBEDDING_NAMES, LoRAType, @@ -105,6 +106,7 @@ def __init__( ) self.lora_use_virtual_experts: bool = get_lora().lora_use_virtual_experts self.lora_strict_loading: bool = get_lora().lora_strict_loading + self.lora_no_cpu_backup: bool = get_lora().lora_no_cpu_backup self.speculative_algorithm: Optional[str] = get_spec().speculative_algorithm # LoRA backend for running sgemm kernels @@ -293,9 +295,19 @@ def _load_lora_adapter(self, lora_ref: LoRARef) -> LoRAUpdateOutput: return self.create_lora_update_result(success=True) - def validate_new_adapter(self, lora_config: LoRAConfig, lora_ref: LoRARef): + def validate_new_adapter( + self, + lora_config: LoRAConfig, + lora_ref: LoRARef, + is_update: bool = False, + old_ref: Optional[LoRARef] = None, + ): """ Validate if an adapter can be loaded into the current LoRA memory pool and generate error if it is incompatible. + + For an in-place refresh (``is_update``), pass the currently-loaded ref as + ``old_ref`` so checks that count loaded adapters exclude the adapter's + own current state. """ if lora_config.lora_added_tokens_size > 0: raise ValueError( @@ -307,18 +319,19 @@ def validate_new_adapter(self, lora_config: LoRAConfig, lora_ref: LoRARef): f"Failed to load {lora_ref.lora_name} because LoRA serving currently doesn't support DoRA adapters" ) - # Check if this LoRA adapter is already loaded - for existing_lora_ref in self.lora_refs.values(): - if lora_ref.lora_name == existing_lora_ref.lora_name: - raise ValueError( - f"Failed to load LoRA adapter {lora_ref.lora_name} because it is already loaded" - ) + # Reject duplicates unless refreshing an existing adapter in place. + if not is_update: + for existing_lora_ref in self.lora_refs.values(): + if lora_ref.lora_name == existing_lora_ref.lora_name: + raise ValueError( + f"Failed to load LoRA adapter {lora_ref.lora_name} because it is already loaded" + ) - if lora_ref.lora_path == existing_lora_ref.lora_path: - logger.warning( - f"{lora_ref.lora_path} is already loaded with name: {existing_lora_ref.lora_name}, " - f"but another copy is being loaded with name: {lora_ref.lora_name}" - ) + if lora_ref.lora_path == existing_lora_ref.lora_path: + logger.warning( + f"{lora_ref.lora_path} is already loaded with name: {existing_lora_ref.lora_name}, " + f"but another copy is being loaded with name: {lora_ref.lora_name}" + ) # Check if the LoRA adapter shape is compatible with the current LoRA memory pool configuration. memory_pool = getattr(self, "memory_pool", None) @@ -331,7 +344,13 @@ def validate_new_adapter(self, lora_config: LoRAConfig, lora_ref: LoRARef): ) # Ensure pinned LoRA adapters does not exceed maximal limit or cause starvation. - if lora_ref.pinned and self.num_pinned_loras >= self.max_loras_per_batch - 1: + # On an in-place refresh the adapter's own pin is already counted in + # num_pinned_loras; refreshing occupies no additional pinned slot, so + # exclude it or a pinned adapter at the limit could never be refreshed. + num_other_pinned_loras = self.num_pinned_loras + if is_update and old_ref is not None: + num_other_pinned_loras -= int(bool(old_ref.pinned)) + if lora_ref.pinned and num_other_pinned_loras >= self.max_loras_per_batch - 1: raise ValueError( f"Failed to load LoRA adapter {lora_ref.lora_name} as a pinned adapter. It is not allowed to pin all slots " "in the LoRA memory pool to avoid starvation for unpinned adapters and base models. Please increase your " @@ -452,6 +471,7 @@ def reset_lora_batch(self): self.lora_backend.reset_batch_state() def prepare_lora_batch(self, forward_batch: ForwardBatch): + self.memory_pool.check_valid() # Some internal-only backends (currently UNO) use explicit token-row # routing for their adapted forwards and want all-base batches to run # through the plain model path. Clear any routing retained by the @@ -470,6 +490,23 @@ def prepare_lora_batch(self, forward_batch: ForwardBatch): and bs <= self.max_bs_in_cuda_graph and forward_batch.forward_mode.is_cuda_graph() ) + if use_cuda_graph and self.lora_backend.is_moe_lora: + # This flag is a HEURISTIC computed before the runner's real replay decision. Under + # --enable-dp-attention a batch that is graph-eligible on THIS rank (small local bs) + # can still have a DP-gathered token count exceeding the static MoE capture buffers + # (max_bs*dp) when other ranks carry more tokens — such a batch cannot replay the + # captured graph and runs eager, so prep the eager (freshly sized) buffers for it + # instead of an under-sized static mapping. + from sglang.srt.lora.backend.base_backend import get_gathered_moe_num_tokens + + moe_cg_buffers = getattr(self.lora_backend, "moe_cg_buffers", None) + if ( + moe_cg_buffers is not None + and get_gathered_moe_num_tokens(forward_batch, bs) + > moe_cg_buffers["token_lora_mapping"].shape[0] + ): + use_cuda_graph = False + # Eligible extend batches refresh the static prefill batch info in # place so captured kernels read current values at replay. use_prefill_cuda_graph = not use_cuda_graph and self.can_use_prefill_cuda_graph( @@ -487,6 +524,55 @@ def prepare_lora_batch(self, forward_batch: ForwardBatch): lora = self.loras[uid] lora_ranks[weight_indices[i]] = lora.config.r scalings[weight_indices[i]] = lora.scaling + + local_active = any(lora_ranks[wi] > 0 for wi in weight_indices) + + # MoE-expert LoRA under --enable-dp-attention: the MoE runs on DP-GATHERED tokens, so this + # rank's local experts also process OTHER ranks' tokens, whose adapter identity is not known + # host-side. Under Tier-1 (exactly one adapter loaded — the colocate RL case) the gathered + # tail can be attributed to that single adapter. Two mutually-exclusive stamps, both armed + # here (policy) and mechanically applied by the backend in _add_moe_lora_info: + # * idle rank (forward_mode.is_idle(), zero local tokens): inject the adapter into + # lora_ranks/scalings (nothing else carries them into the batch tensors) and record its + # buffer id so the backend stamps the WHOLE gathered mapping and enables the adapter — + # otherwise foreign tokens routed to this rank's experts silently lose the LoRA delta. + # * active rank (local requests DO use the adapter): record the buffer id so the backend + # stamps the gathered tail [num_tokens, moe_num_tokens) with it instead of copying an + # arbitrary local token's slot. Base-only local batches (local_active False) and + # multi-adapter batches arm NOTHING: the tail stays -1 (adapter-disabled), foreign + # tokens get base — never a delta this rank cannot attribute. + # Buffer ids are host ints (no GPU sync -> cuda-graph safe). NOTE: gate on + # global_num_tokens_cpu too, not just global_dp_buffer_len — the eager path runs + # prepare_lora_batch from ForwardBatch.init_new BEFORE prepare_mlp_sync_batch assigns + # global_dp_buffer_len (only cuda-graph capture pre-sets it). + idle_rank_active_buffer_id = None + single_active_buffer_id = None + if self.lora_backend.is_moe_lora and ( + getattr(forward_batch, "global_dp_buffer_len", None) is not None + or len(getattr(forward_batch, "global_num_tokens_cpu", None) or []) > 1 + ): + loaded = [ + (uid, bid) + for uid, bid in self.memory_pool.uid_to_buffer_id.items() + if uid is not None + and uid in self.loras + and self.loras[uid].config.r > 0 + ] + if len(loaded) == 1: + uid, bid = loaded[0] + if forward_batch.forward_mode.is_idle(): + lora = self.loras[uid] + lora_ranks[bid] = lora.config.r + scalings[bid] = lora.scaling + idle_rank_active_buffer_id = bid + elif local_active: + single_active_buffer_id = bid + + # Pass the stamps to the backend (read in _add_moe_lora_info); reset each batch so a + # previous batch's stamp never leaks into a later one. + self.lora_backend._idle_rank_active_buffer_id = idle_rank_active_buffer_id + self.lora_backend._single_loaded_buffer_id = single_active_buffer_id + # Do in-place updates when CUDA graph is enabled and the batch forward mode # could use CUDA graph. self.lora_backend.prepare_lora_batch( @@ -497,8 +583,8 @@ def prepare_lora_batch(self, forward_batch: ForwardBatch): use_cuda_graph=use_cuda_graph, use_prefill_cuda_graph=use_prefill_cuda_graph, ) - self.lora_backend.batch_info.has_active_lora = any( - lora_ranks[wi] > 0 for wi in weight_indices + self.lora_backend.batch_info.has_active_lora = local_active or ( + idle_rank_active_buffer_id is not None ) def prepare_lora_token_segments( @@ -736,7 +822,9 @@ def init_lora_shapes( ) target_modules = self.target_modules elif target_modules: - self.target_modules = get_normalized_target_modules(target_modules) + self.target_modules = get_normalized_target_modules( + target_modules, base_model=self.base_model + ) else: self.target_modules = set() @@ -780,7 +868,7 @@ def init_lora_shapes( ) adapter_target_modules = get_normalized_target_modules( - config.target_modules + config.target_modules, base_model=self.base_model ) if target_modules is not None: @@ -798,6 +886,12 @@ def init_lora_shapes( # Otherwise, infer target_modules from adapter configs. self.target_modules.update(adapter_target_modules) + unsupported_targets = self.target_modules - _KNOWN_LORA_TARGET_MODULES + if unsupported_targets: + raise ValueError( + f"Unsupported normalized LoRA targets: {sorted(unsupported_targets)}" + ) + # Fusion folds wk + weights_proj into wk_weights_proj, so the modules # LoRA wraps are absent and an indexer-targeted adapter is silently dropped. indexer_targets = self.target_modules & DSA_INDEXER_LORA_NAMES @@ -867,13 +961,23 @@ def load_lora_weights_from_tensors( """ Load the weights of a LoRA adapter from tensors to CPU memory. """ + self.loras[lora_ref.lora_id] = self._create_lora_adapter_from_tensors( + lora_ref, self.configs[lora_ref.lora_id], tensors + ) + + def _create_lora_adapter_from_tensors( + self, lora_ref: LoRARef, config: LoRAConfig, tensors: Dict[str, torch.Tensor] + ) -> LoRAAdapter: + """Build and initialize a LoRAAdapter without touching served state, + so callers can stage it and commit only after every fallible step.""" lora_adapter = LoRAAdapter( lora_ref.lora_id, - self.configs[lora_ref.lora_id], + config, self.base_hf_config, self.load_config, self.lora_backend, base_model=self.base_model, + keep_weights_on_device=self.lora_no_cpu_backup, ) lora_adapter.initialize_weights_from_tensors(tensors) @@ -882,18 +986,50 @@ def load_lora_weights_from_tensors( embedding_layer_names=lora_adapter.embedding_layers.keys(), speculative_algorithm=self.speculative_algorithm, ) - self.loras[lora_ref.lora_id] = lora_adapter + return lora_adapter + + def register_lora_adapter( + self, lora_ref: LoRARef, config_dict: Dict + ) -> LoRAUpdateOutput: + """Create-or-refresh an adapter's identity and config; re-registration + re-zeroes — a slot's new tenant must not serve its predecessor's weights.""" + return self.load_lora_adapter_from_tensors( + lora_ref, tensors={}, config_dict=config_dict, upsert=True + ) + + def require_registered(self, lora_name: str) -> LoRARef: + for ref in self.lora_refs.values(): + if ref.lora_name == lora_name: + return ref + raise ValueError( + f"streamed weights for unregistered LoRA adapter {lora_name!r}; " + "call /register_lora_adapter before streaming its tensors" + ) + + def apply_streamed_adapter( + self, lora_name: str, tensors: Dict[str, torch.Tensor] + ) -> LoRAUpdateOutput: + """Apply a complete streamed tensor set to a registered adapter, reusing + the config stored at registration.""" + try: + lora_ref = self.require_registered(lora_name) + except ValueError as e: + return self.create_lora_update_result(success=False, error_message=str(e)) + return self.load_lora_adapter_from_tensors( + lora_ref, tensors, config_dict=None, upsert=True + ) def load_lora_adapter_from_tensors( self, lora_ref: LoRARef, tensors: Dict[str, torch.Tensor], - config_dict: Dict, + config_dict: Optional[Dict], added_tokens_config: Optional[Dict] = None, + upsert: bool = False, ) -> LoRAUpdateOutput: logger.info(f"LoRA adapter loading from tensors starts: {lora_ref}.") result = self._load_lora_adapter_from_tensors( - lora_ref, tensors, config_dict, added_tokens_config + lora_ref, tensors, config_dict, added_tokens_config, upsert=upsert ) logger.info(f"LoRA adapter loading from tensors completes: {lora_ref}.") return result @@ -904,36 +1040,129 @@ def _load_lora_adapter_from_tensors( tensors: Dict[str, torch.Tensor], config_dict: Dict, added_tokens_config: Optional[Dict] = None, + upsert: bool = False, ) -> LoRAUpdateOutput: """ Load a single LoRA adapter from tensors and config dict. + + With ``upsert``, an already-loaded adapter has its weights refreshed in + place (reusing its lora_id); otherwise the adapter must not be loaded yet. """ assert lora_ref.lora_name is not None and lora_ref.lora_path is not None, ( "LoRARef must have both lora_name and lora_path set for loading." ) - assert lora_ref.lora_id not in self.loras, ( - f"LoRA adapter with ID {lora_ref.lora_id} is already loaded. This should have been verified before request is sent to the backend." - ) - - try: - new_adapter = LoRAConfig.from_dict( - config_dict, - added_tokens_config, - base_vocab_size=self.base_hf_config.vocab_size, + is_update = upsert and (lora_ref.lora_id in self.loras) + if not is_update: + assert lora_ref.lora_id not in self.loras, ( + f"LoRA adapter with ID {lora_ref.lora_id} is already loaded. This should have been verified before request is sent to the backend." ) - self.validate_new_adapter(new_adapter, lora_ref) - self.configs[lora_ref.lora_id] = new_adapter - self.load_lora_weights_from_tensors(lora_ref, tensors) + uid = lora_ref.lora_id + old_config = self.configs.get(uid) + old_lora = self.loras.get(uid) + old_ref = self.lora_refs.get(uid) - self.lora_refs[lora_ref.lora_id] = lora_ref - self.num_pinned_loras += int(lora_ref.pinned) + # Stage every fallible step before mutating served state: a failed + # request must not leave a live adapter with new metadata over old + # weights (or leak unreachable entries on a fresh insert). + try: + if config_dict is None: + # Streamed apply: reuse the config stored at registration. + assert old_config is not None, ( + f"config_dict=None requires a registered adapter (got {lora_ref})" + ) + new_config = old_config + else: + new_config = LoRAConfig.from_dict( + config_dict, + added_tokens_config, + base_vocab_size=self.base_hf_config.vocab_size, + ) + self.validate_new_adapter( + new_config, lora_ref, is_update=is_update, old_ref=old_ref + ) + new_lora = self._create_lora_adapter_from_tensors( + lora_ref, new_config, tensors + ) except Exception as e: return self.create_lora_update_result( success=False, error_message=str(e), ) + if ( + self.lora_no_cpu_backup + and getattr(self, "memory_pool", None) is not None + and (tensors or (is_update and uid in self.memory_pool.uid_to_buffer_id)) + ): + try: + if self.device.type == "cuda": + torch.cuda.synchronize(self.device) + self.memory_pool.install_streamed_adapter( + uid, + new_lora, + self.lora_modules, + self.embed_tokens_module, + self.lm_head_module, + ) + except Exception as e: + return self.create_lora_update_result( + success=False, error_message=str(e) + ) + elif ( + is_update + and getattr(self, "memory_pool", None) is not None + and uid in self.memory_pool.uid_to_buffer_id + ): + buffer_id = self.memory_pool.uid_to_buffer_id[uid] + try: + # The served buffer slot is rewritten in place. Wait for + # already-launched kernels first: under overlap scheduling / + # CUDA-graph replay a forward pass can still be executing on + # another stream, and it must not read a half-rewritten slot. + if self.device.type == "cuda": + torch.cuda.synchronize(self.device) + self.memory_pool.load_lora_weight_to_buffer( + uid, + buffer_id, + new_lora, + self.lora_modules, + self.embed_tokens_module, + self.lm_head_module, + ) + except Exception as e: + # Roll back so the adapter keeps serving its previous weights + # instead of a torn half-old/half-new buffer. + self.configs[uid] = old_config + self.loras[uid] = old_lora + try: + self.memory_pool.load_lora_weight_to_buffer( + uid, + buffer_id, + old_lora, + self.lora_modules, + self.embed_tokens_module, + self.lm_head_module, + ) + except Exception: + logger.exception( + f"Failed to restore previous weights for LoRA adapter " + f"{lora_ref.lora_name} (buffer slot {buffer_id}) after a " + "failed upsert; the served buffer may be corrupted." + ) + return self.create_lora_update_result( + success=False, + error_message=str(e), + ) + + self.configs[uid] = new_config + self.loras[uid] = new_lora + self.lora_refs[uid] = lora_ref + # An upsert may change ``pinned``; track the delta against the replaced + # ref so the counter stays consistent with lora_refs. + old_pinned = int(bool(old_ref.pinned)) if is_update else 0 + self.num_pinned_loras += int(bool(lora_ref.pinned)) - old_pinned + return self.create_lora_update_result(success=True) def init_memory_pool(self): @@ -953,6 +1182,7 @@ def init_memory_pool(self): experts_shared_outer_loras=self.experts_shared_outer_loras, strict_loading=self.lora_strict_loading, enable_lora_overlap_loading=self.enable_lora_overlap_loading, + lora_no_cpu_backup=self.lora_no_cpu_backup, ) # Initializing memory pool with base model @@ -972,6 +1202,8 @@ def init_lora_modules(self): self.embed_tokens_module: Optional[BaseLayerWithLoRA] = None self.lm_head_module: Optional[BaseLayerWithLoRA] = None + # One MoE layer's worth of warning is enough; the condition is model-wide. + warned_one_sided_moe_targets = False # When tie_word_embeddings=True, lm_head is the same Python object as # embed_tokens. PyTorch's named_modules() deduplicates by object identity, @@ -1063,9 +1295,36 @@ def init_lora_modules(self): ) continue - if isinstance(module, (FusedMoE, InklingBatchDenseMLP)) and all( - x in self.target_modules for x in ["gate_up_proj", "down_proj"] - ): + if isinstance(module, (FusedMoE, InklingBatchDenseMLP)): + expert_projections = ("gate_up_proj", "down_proj") + missing = [ + x for x in expert_projections if x not in self.target_modules + ] + if len(missing) == 1 and not warned_one_sided_moe_targets: + warned_one_sided_moe_targets = True + # The MoE-LoRA hooks inject a delta after gate_up and after + # down together, so wrapping for only one of them is not + # implemented and the layer stays unwrapped. Adapters that only + # train the dense/shared MLP still work, but one that carries + # expert weights for the targeted projection has them loaded + # into the pool and never applied. Warn once -- silently + # dropping part of an adapter reads as a serving/training + # mismatch with no symptom. + present = next(x for x in expert_projections if x != missing[0]) + logger.warning( + "LoRA target modules include %r but not %r, so MoE expert layers " + "(e.g. %s) get no adapter at all: expert LoRA needs both projections. " + "Any %r expert weights in a loaded adapter will be ignored. Add %r to " + "the target modules to enable expert LoRA -- an adapter that does not " + "train it keeps zero-filled buffers and contributes no delta.", + present, + missing[0], + module_name, + present, + missing[0], + ) + if missing: + continue layer_id = get_layer_id(module_name) if layer_id is None: if module_name.startswith("model.meta_mlp."): diff --git a/python/sglang/srt/lora/lora_moe_runner_marlin.py b/python/sglang/srt/lora/lora_moe_runner_marlin.py index 405643a0ec36..0d973811035a 100644 --- a/python/sglang/srt/lora/lora_moe_runner_marlin.py +++ b/python/sglang/srt/lora/lora_moe_runner_marlin.py @@ -24,13 +24,15 @@ _is_cuda = is_cuda() if _is_cuda: - from sglang.kernels.ops.activation import silu_and_mul from sglang.kernels.ops.moe.fused_moe_triton_kernels import ( moe_sum_reduce_triton, ) from sglang.kernels.ops.moe.moe_wna16_marlin import moe_wna16_marlin_gemm from sglang.srt.layers.moe.fused_moe_triton.fused_marlin_moe import ( + gated_marlin_activation, get_scalar_type, + is_mxfp4_marlin_weights, + marlin_moe_use_atomic_add, ) from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import ( moe_align_block_size, @@ -74,7 +76,9 @@ def run_from_dispatch( topk_weights = topk_output.topk_weights topk_ids = topk_output.topk_ids - assert runner_config.activation == "silu", "Only SiLU activation is supported." + assert runner_config.is_gated, ( + "MarlinLoraRunnerCore only supports gated experts" + ) assert torch.cuda.get_device_capability(hidden_states.device)[0] >= 9, ( "MarlinLoraRunnerCore requires CUDA compute capability >= 9" ) @@ -85,6 +89,16 @@ def run_from_dispatch( N = quant_info.w2_qweight.shape[1] * 16 topk = topk_ids.shape[1] num_bits = quant_info.weight_bits + use_atomic_add = marlin_moe_use_atomic_add( + hidden_states, + is_mxfp4_marlin_weights( + num_bits, + quant_info.w13_scales, + quant_info.w2_scales, + quant_info.w13_qzeros, + quant_info.w2_qzeros, + ), + ) for block_size_m in [8, 16, 32, 48, 64]: if M * topk / E / block_size_m < 0.9: @@ -146,7 +160,7 @@ def run_from_dispatch( size_n=2 * N, size_k=K, is_k_full=quant_info.is_k_full, - use_atomic_add=True, + use_atomic_add=use_atomic_add, use_fp32_reduce=True, is_zp_float=False, ) @@ -162,7 +176,13 @@ def run_from_dispatch( intermediate_cache2 = torch.empty( (M * topk, N), device=hidden_states.device, dtype=hidden_states.dtype ) - silu_and_mul(intermediate_cache1.view(-1, 2 * N), intermediate_cache2) + gated_marlin_activation( + runner_config.activation, + runner_config.gemm1_alpha, + runner_config.gemm1_clamp_limit, + intermediate_cache1.view(-1, 2 * N), + intermediate_cache2, + ) # Stage 3: Down (Marlin) intermediate_cache3 = torch.empty( @@ -195,7 +215,7 @@ def run_from_dispatch( size_n=K, size_k=N, is_k_full=quant_info.is_k_full, - use_atomic_add=True, + use_atomic_add=use_atomic_add, use_fp32_reduce=True, is_zp_float=False, ) diff --git a/python/sglang/srt/lora/lora_moe_runners.py b/python/sglang/srt/lora/lora_moe_runners.py index 8b1d3e5e4533..a7950694caa0 100644 --- a/python/sglang/srt/lora/lora_moe_runners.py +++ b/python/sglang/srt/lora/lora_moe_runners.py @@ -328,7 +328,6 @@ def _add_lora_gate_up_delta( r = lora_info.max_lora_rank gate_up_a = lora_info.gate_up_lora_a_weights gate_up_b = lora_info.gate_up_lora_b_weights - if lora_info.experts_shared_outer_loras and not lora_info.lora_use_virtual_experts: gate_up_a = gate_up_a.expand(-1, lora_info.num_experts, -1, -1) @@ -337,7 +336,18 @@ def _add_lora_gate_up_delta( is_gated = gate_up_a.shape[2] > r if is_gated: inter_size = gate_up_b.shape[2] // 2 + # Marlin may pad the intermediate: the base output is [gate | up] at the padded width + padded_inter_size = gate_up_dim // 2 + assert inter_size <= padded_inter_size, ( + f"gate/up LoRA output {inter_size} is wider than the base intermediate {padded_inter_size}" + ) + if lora_info.lora_use_virtual_experts and inter_size != padded_inter_size: + raise NotImplementedError( + "virtual-expert MoE LoRA does not support padded Marlin outputs" + ) lora_a_stacked = [gate_up_a[:, :, :r, :], gate_up_a[:, :, r : 2 * r, :]] + # B halves are also the tuple form the virtual-experts kernel wants + # (one shrink at K=2*r, two expands at K=r each). lora_b_stacked = [ gate_up_b[:, :, :inter_size, :], gate_up_b[:, :, inter_size:, :], @@ -351,7 +361,7 @@ def _add_lora_gate_up_delta( output=intermediate_cache, hidden_states=hidden_states, lora_a=gate_up_a, - lora_b=gate_up_b, + lora_b=tuple(lora_b_stacked) if is_gated else gate_up_b, topk_ids=topk_ids, topk_weights=topk_weights, token_lora_mapping=token_lora_mapping, @@ -417,7 +427,16 @@ def _add_lora_down_delta( down_lora_a = lora_info.down_lora_a_weights down_lora_b = lora_info.down_lora_b_weights + # the Marlin activation carries the padded intermediate; LoRA A is only the logical width + inter_size = down_lora_a.shape[-1] + assert inter_size <= intermediate_input.shape[-1], ( + f"down LoRA input {inter_size} is wider than the base activation {intermediate_input.shape[-1]}" + ) + if inter_size != intermediate_input.shape[-1]: + intermediate_input = intermediate_input[:, :inter_size] if lora_info.experts_shared_outer_loras and not lora_info.lora_use_virtual_experts: + # fused_moe_lora requires B's expert_dim to match A's; expand the + # shared B view. down_lora_b = down_lora_b.expand(-1, lora_info.num_experts, -1, -1) if lora_info.fully_sharded and lora_info.tp_size > 1: diff --git a/python/sglang/srt/lora/lora_registry.py b/python/sglang/srt/lora/lora_registry.py index 09f0c771a8f0..f50dc0aec2ce 100644 --- a/python/sglang/srt/lora/lora_registry.py +++ b/python/sglang/srt/lora/lora_registry.py @@ -15,11 +15,11 @@ import asyncio from collections import OrderedDict -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional, Tuple, Union from uuid import NAMESPACE_URL, uuid4, uuid5 import msgspec -from msgspec.structs import fields +from msgspec.structs import fields, replace from sglang.srt.utils import ConcurrentCounter from sglang.srt.utils.aio_rwlock import RWLock @@ -38,6 +38,12 @@ class LoRARef(msgspec.Struct, frozen=True, array_like=True): lora_name: Optional[str] = None lora_path: Optional[str] = None pinned: Optional[bool] = None + # False for adapters whose weights arrived over the wire (lora_path + # "__stream__"): there is no disk artifact to reload from, so they must + # never be LRU-evicted nor implicitly reloaded. + # Trailing field with a default keeps the array_like wire format + # compatible with refs encoded before this field existed. + reloadable: bool = True def __post_init__(self): if self.lora_id is None: @@ -123,6 +129,48 @@ async def unregister(self, lora_name: str) -> str: return lora_ref.lora_id + async def get_lora_id(self, lora_name: str) -> Optional[str]: + """Return the ``lora_id`` of a registered adapter, or ``None``.""" + async with self._registry_lock.reader_lock: + lora_ref = self._registry.get(lora_name, None) + return lora_ref.lora_id if lora_ref is not None else None + + async def register_or_reuse( + self, lora_ref: LoRARef, upsert: bool = False + ) -> Tuple[LoRARef, bool]: + """Resolve which identity a load request should use. + + Returns ``(ref, reused)``. With ``upsert`` and a same-name adapter + already registered, the returned ref adopts the existing ``lora_id`` + (``reused=True``) so the backend refreshes that adapter in place; + otherwise ``lora_ref`` is returned unchanged (``reused=False``). + Nothing is registered here: the caller commits the resolved ref with + ``register`` / ``refresh`` once the backend load succeeded, keeping + failed loads invisible to the registry. + """ + if not upsert: + return lora_ref, False + async with self._registry_lock.reader_lock: + existing = self._registry.get(lora_ref.lora_name, None) + if existing is None: + return lora_ref, False + return replace(lora_ref, lora_id=existing.lora_id), True + + async def refresh(self, lora_ref: LoRARef): + """Replace a registered adapter's ref after a successful upsert. + + Keeps the id (asserted) while adopting the new path/pinned metadata, + and counts as a use for LRU ordering. + """ + async with self._registry_lock.writer_lock: + existing = self._registry.get(lora_ref.lora_name, None) + assert existing is not None and existing.lora_id == lora_ref.lora_id, ( + f"refresh() must target a registered adapter with the same lora_id; " + f"got {lora_ref}, registered: {existing}" + ) + self._registry[lora_ref.lora_name] = lora_ref + self._registry.move_to_end(lora_ref.lora_name) + async def acquire(self, lora_name: Union[str, List[str]]) -> Union[str, List[str]]: """ Queries registry for LoRA IDs based on LoRA names and start tracking the usage of the corresponding LoRA adapters @@ -142,24 +190,28 @@ def _lookup(name: str) -> str: self._registry.move_to_end(name) return lora_ref.lora_id + # Lookup and increment must be one atomic admission step under the + # writer lock: with the increment outside, an unload can slip between + # them — unregister the adapter, observe the counter at zero, delete + # the counter — leaving this request to KeyError or to run on an + # adapter that is already gone. wait_for_unload never awaits inside + # this lock, so holding it across the increment cannot deadlock. if isinstance(lora_name, str): async with self._registry_lock.writer_lock: lora_id = _lookup(lora_name) - - await self._counters[lora_id].increment(notify_all=False) + await self._counters[lora_id].increment(notify_all=False) return lora_id elif isinstance(lora_name, list): async with self._registry_lock.writer_lock: lora_ids = [_lookup(name) for name in lora_name] - - # Increment the counters only after all IDs are looked up. - await asyncio.gather( - *[ - self._counters[id].increment(notify_all=False) - for id in lora_ids - if id is not None - ] - ) + # Increment the counters only after all IDs are looked up. + await asyncio.gather( + *[ + self._counters[id].increment(notify_all=False) + for id in lora_ids + if id is not None + ] + ) return lora_ids else: raise TypeError("lora_name must be either a string or a list of strings.") @@ -224,14 +276,15 @@ async def lru_lora_name(self, exclude_pinned=False): If exclude_pinned is True, then return the LRU LoRA adapter that isn't pinned. """ async with self._registry_lock.reader_lock: - if not exclude_pinned: - return next(iter(self._registry), None) - for lora_name, lora_ref in self._registry.items(): - if not lora_ref.pinned: - return lora_name - else: - return None + # Evicting a non-reloadable (wire-loaded) adapter is always + # destructive: there is no artifact to reload it from. + if not lora_ref.reloadable: + continue + if exclude_pinned and lora_ref.pinned: + continue + return lora_name + return None def _register_adapter(self, lora_ref: LoRARef): """ diff --git a/python/sglang/srt/lora/mem_pool.py b/python/sglang/srt/lora/mem_pool.py index 0270ef3e6a00..f70232025698 100644 --- a/python/sglang/srt/lora/mem_pool.py +++ b/python/sglang/srt/lora/mem_pool.py @@ -37,6 +37,7 @@ get_normalized_target_modules, get_stacked_multiply, get_target_module_name, + shared_experts_count, ) from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import is_pin_memory_available @@ -147,7 +148,13 @@ def __init__( experts_shared_outer_loras: bool = False, strict_loading: bool = False, enable_lora_overlap_loading: bool = False, + lora_no_cpu_backup: bool = False, ): + assert not (lora_no_cpu_backup and enable_lora_overlap_loading), ( + "--lora-no-cpu-backup drops the staged weights right after install, which overlap loading still reads" + ) + self.lora_no_cpu_backup: bool = lora_no_cpu_backup + self.invalid_uids: Set[str] = set() self.base_hf_config: AutoConfig = base_hf_config self.num_layer: int = base_hf_config.num_hidden_layers self.max_loras_per_batch: int = max_loras_per_batch @@ -248,7 +255,9 @@ def _can_support(config: LoRAConfig) -> bool: return False if config.lora_added_tokens_size > self.lora_added_tokens_size: return False - target_module_names = get_normalized_target_modules(config.target_modules) + target_module_names = get_normalized_target_modules( + config.target_modules, base_model=self.base_model + ) if "all" in target_module_names: return True return target_module_names.issubset(self.target_modules) @@ -317,6 +326,38 @@ def _has_moe_module(base_model: torch.nn.Module) -> bool: return any(isinstance(m, FusedMoE) for m in base_model.modules()) + @staticmethod + def _get_fused_shared_moe(base_model: torch.nn.Module, layer_idx: int): + from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE + + return next( + ( + module + for module in base_model.modules() + if isinstance(module, FusedMoE) + and module.layer_id == layer_idx + and module.num_fused_shared_experts > 0 + ), + None, + ) + + def _get_moe_pool_expert_dim(self, base_model, layer_idx): + fused_moe = self._get_fused_shared_moe(base_model, layer_idx) + if fused_moe is None: + return self._get_num_local_experts(base_model) + assert not self.experts_shared_outer_loras, ( + "fused shared experts require per-expert LoRA factors" + ) + assert self.moe_use_local_expert_ids or ( + fused_moe.num_experts + == fused_moe._num_global_routed + fused_moe.num_fused_shared_experts + ), "LoRA does not support global IDs with per-rank fused shared slots" + return ( + fused_moe.num_local_experts + if self.moe_use_local_expert_ids + else fused_moe.num_experts + ) + def _get_num_local_experts(self, base_model: torch.nn.Module) -> int: """Experts owned by this rank. Equals the global count when EP is off, the runner keeps global IDs, or the split isn't even (all @@ -341,6 +382,7 @@ def _iter_local_expert_weights( cache_keys: Union[str, Dict[int, str]], *, localize: bool = True, + fused_moe=None, ) -> Iterator[Tuple[int, torch.Tensor, str]]: """Yield `(expert_id, weight, cache_key)` triples for MoE LoRA A/B weights. @@ -350,7 +392,12 @@ def _iter_local_expert_weights( if isinstance(weights, dict): assert isinstance(cache_keys, dict) for gid, w in weights.items(): - lid = self._global_to_local_expert_id(gid) if localize else gid + if fused_moe is not None and localize and self.moe_use_local_expert_ids: + lid = fused_moe._map_global_expert_id_to_local_expert_id(gid) + if lid < 0: + continue + else: + lid = self._global_to_local_expert_id(gid) if localize else gid if lid is not None: yield lid, w, cache_keys[gid] return @@ -375,6 +422,157 @@ def _iter_local_expert_weights( f"Expected dict or 3D torch.Tensor, got {type(weights).__name__}." ) + def _row_parallel_shard_tp( + self, module_name: str, base_model: torch.nn.Module, layer_idx: int + ) -> int: + """Shard count for a non-MoE row-parallel module's activation axis. + + Probes the base module's ``input_size // input_size_per_partition`` so + the LoRA buffer matches the actual shard regardless of which TP group + owns it — covers DP-attention (``o_proj`` uses ``attn_tp_size``) and + shared-expert dense-vs-MoE per-layer-TP differences. Falls back to + ``self.tp_size``. Cached per ``(module_name, layer_idx)``. + + MoE-internal names go through ``self.moe_tp_size`` upstream. + """ + cache = getattr(self, "_row_parallel_tp_cache", None) + if cache is None: + cache = {} + setattr(self, "_row_parallel_tp_cache", cache) + key = (module_name, layer_idx) + if key in cache: + return cache[key] + + layer_markers = (f".layers.{layer_idx}.", f"layers.{layer_idx}.") + + def _probe(m): + in_size = getattr(m, "input_size", None) + per_part = getattr(m, "input_size_per_partition", None) + if in_size is not None and per_part is not None and per_part > 0: + return max(1, in_size // per_part) + inner = getattr(m, "base_layer", None) + if inner is not None and inner is not m: + return _probe(inner) + return None + + suffix = f".{module_name}" + found = None + for _name, module in base_model.named_modules(): + if not _name.endswith(suffix): + continue + if not any(marker in _name for marker in layer_markers): + continue + r = _probe(module) + if r is not None: + found = r + break + + out = found if found is not None else self.tp_size + cache[key] = out + return out + + def _column_parallel_shard_tp( + self, module_name: str, base_model: torch.nn.Module, layer_idx: int + ) -> int: + """Shard count for a non-MoE column-parallel module's output axis. + + Probes the base module's ``output_size // output_size_per_partition``. + The input-axis probe used for row-parallel modules is poisoned here: + quantized linear methods (e.g. ``Fp8LinearMethod.create_weights``) + stamp ``input_size_per_partition == input_size`` on column-parallel + layers (whose input is never sharded), which reports a shard count of + 1 and leaves the LoRA-B buffer at the full output dim while the base + stays TP-sharded -- ``set_lora_info`` then fails with "LoRA B output + dim != base partition prefix dim" (e.g. blockwise-fp8 GLM-5.2 + ``shared_experts.gate_up_proj``). The output-axis ratio is + quantization-independent: ``output_size_per_partition`` is set in + ``ColumnParallelLinear.__init__`` before the quant method runs. Falls + back to ``self.tp_size``. Cached per ``(module_name, layer_idx)``. + """ + cache = getattr(self, "_col_parallel_tp_cache", None) + if cache is None: + cache = {} + setattr(self, "_col_parallel_tp_cache", cache) + key = (module_name, layer_idx) + if key in cache: + return cache[key] + + layer_markers = (f".layers.{layer_idx}.", f"layers.{layer_idx}.") + + def _probe(m): + out_size = getattr(m, "output_size", None) + per_part = getattr(m, "output_size_per_partition", None) + if out_size is not None and per_part is not None and per_part > 0: + return max(1, out_size // per_part) + inner = getattr(m, "base_layer", None) + if inner is not None and inner is not m: + return _probe(inner) + return None + + suffix = f".{module_name}" + found = None + for _name, module in base_model.named_modules(): + if not _name.endswith(suffix): + continue + if not any(marker in _name for marker in layer_markers): + continue + r = _probe(module) + if r is not None: + found = r + break + + out = found if found is not None else self.tp_size + cache[key] = out + return out + + def _column_parallel_out_partition( + self, module_name: str, base_model: torch.nn.Module, layer_idx: int + ): + """Actual per-rank output dim of a non-MoE column-parallel base module. + + Reads ``output_size_per_partition`` from the matching base linear -- the + ground truth that ``set_lora_info`` validates against. Critically handles + the dense MLP ``gate_up_proj`` that is fully REPLICATED under + ``--moe-dense-tp-size 1`` (``output_size_per_partition == output_size``), + where dividing ``get_hidden_dim``'s output by the global ``tp_size`` + undersizes LoRA-B and raises "LoRA B output dim != base partition prefix + dim". Returns ``None`` if no base module is found. Cached per + ``(module_name, layer_idx)``. + """ + cache = getattr(self, "_col_parallel_out_cache", None) + if cache is None: + cache = {} + setattr(self, "_col_parallel_out_cache", cache) + key = (module_name, layer_idx) + if key in cache: + return cache[key] + + layer_markers = (f".layers.{layer_idx}.", f"layers.{layer_idx}.") + + def _probe(m): + ops = getattr(m, "output_size_per_partition", None) + if ops is not None and ops > 0: + return ops + inner = getattr(m, "base_layer", None) + if inner is not None and inner is not m: + return _probe(inner) + return None + + suffix = f".{module_name}" + found = None + for _name, module in base_model.named_modules(): + if not _name.endswith(suffix): + continue + if not any(marker in _name for marker in layer_markers): + continue + r = _probe(module) + if r is not None: + found = r + break + + cache[key] = found + return found + def _get_standard_shape( self, module_name: str, @@ -387,8 +585,12 @@ def _get_standard_shape( module_name, self.base_hf_config, base_model, layer_idx ) c = get_stacked_multiply(module_name, base_model) - if self.tp_size > 1 and module_name in ROW_PARALLELISM_LINEAR_LORA_NAMES: - input_dim = divide(input_dim, self.tp_size) + # Non-MoE row-parallel modules: probe the actual shard size so o_proj / + # down_proj match attn_tp under DP-attention and the shared-experts + # dense-vs-MoE per-layer-TP differences. + row_tp = self._row_parallel_shard_tp(module_name, base_model, layer_idx) + if row_tp > 1 and module_name in ROW_PARALLELISM_LINEAR_LORA_NAMES: + input_dim = divide(input_dim, row_tp) return (self.max_loras_per_batch, max_lora_dim * c, input_dim) def get_lora_A_shape( @@ -421,7 +623,7 @@ def get_lora_A_shape( if self.is_shared_moe_module(module_name): expert_dim = self._get_num_shared_experts(base_model) else: - expert_dim = self._get_num_local_experts(base_model) + expert_dim = self._get_moe_pool_expert_dim(base_model, layer_idx) if self.experts_shared_outer_loras and module_name in ( "gate_up_proj_moe", "gate_up_proj_shared_moe", @@ -513,16 +715,32 @@ def get_lora_B_shape( and module_name not in ROW_PARALLELISM_LINEAR_LORA_NAMES and module_name not in REPLICATED_LINEAR_LORA_NAMES ): - output_dim = self._column_parallel_lora_b_per_rank_dim( - module_name, output_dim, effective_tp_size + # If the base column-parallel module is fully REPLICATED (its actual + # output_size_per_partition still equals the full output_dim -- e.g. the + # dense MLP gate_up under --moe-dense-tp-size 1), its output is NOT + # sharded, so keep LoRA-B at the full output dim. Dividing by the global + # tp_size here undersizes B and crashes set_lora_info ("LoRA B output dim + # != base partition prefix dim"). Non-MoE only; MoE shards by moe_tp_size. + probed_out = ( + None + if self.is_moe_module(module_name) + else self._column_parallel_out_partition( + module_name, base_model, layer_idx + ) ) + if probed_out is not None and probed_out == output_dim: + pass # replicated base: keep full B output dim + else: + output_dim = self._column_parallel_lora_b_per_rank_dim( + module_name, output_dim, effective_tp_size + ) # Check if MoE module and return appropriate shape if self.is_moe_module(module_name): if self.is_shared_moe_module(module_name): expert_dim = self._get_num_shared_experts(base_model) else: - expert_dim = self._get_num_local_experts(base_model) + expert_dim = self._get_moe_pool_expert_dim(base_model, layer_idx) if self.experts_shared_outer_loras and module_name in ( "down_proj_moe", "down_proj_shared_moe", @@ -572,9 +790,8 @@ def init_buffer( if hasattr(cfg, "get_text_config"): cfg = cfg.get_text_config() has_shared_experts = ( - hasattr(cfg, "shared_expert_intermediate_size") - and cfg.shared_expert_intermediate_size > 0 - ) or (getattr(cfg, "n_shared_experts", 0) or 0) > 0 + getattr(cfg, "shared_expert_intermediate_size", 0) or 0 + ) > 0 or (shared_experts_count(cfg) or 0) > 0 has_moe = self._has_moe_module(base_model) # Shape functions automatically handle both 3D (standard) and 4D (MoE) @@ -598,16 +815,24 @@ def init_buffer( # MoE expert version (4D) moe_key = f"{module_name}_moe" - buffer[moe_key] = [ - torch.zeros( - get_lora_shape_fn( - moe_key, base_model, self.max_lora_rank, idx - ), - dtype=self.dtype, - device=device, + buffer[moe_key] = [] + for idx in range(self.num_layer): + shape = get_lora_shape_fn( + moe_key, base_model, self.max_lora_rank, idx + ) + fused_moe = self._get_fused_shared_moe(base_model, idx) + if fused_moe is not None: + num_experts = ( + fused_moe.num_local_experts + if self.moe_use_local_expert_ids + else fused_moe.num_experts + ) + assert shape[1] == num_experts, ( + "LoRA/base expert layout mismatch" + ) + buffer[moe_key].append( + torch.zeros(shape, dtype=self.dtype, device=device) ) - for idx in range(self.num_layer) - ] # Shared-expert MoE version (4D, separate sink namespace). if self._has_shared_fused_moe(base_model): @@ -746,6 +971,7 @@ def prepare_lora_batch( lora_embed_tokens_module: Optional[BaseLayerWithLoRA], lora_lm_head_module: Optional[BaseLayerWithLoRA], ): + self.check_valid() # Python hash seeds differ by TP process; slot and LRU updates must not. ordered_uids = sorted(cur_uids, key=lambda uid: (uid is not None, uid or "")) @@ -791,6 +1017,12 @@ def get_available_buffer_slot(): # Select victim using eviction policy victim_uid = self.eviction_policy.select_victim(candidates_to_use) + if victim_uid is not None and self.lora_no_cpu_backup: + raise RuntimeError( + f"Refusing to evict LoRA adapter '{victim_uid}': --lora-no-cpu-backup dropped its " + "staged weights after installation, so it could never be reloaded. Increase " + "--max-loras-per-batch." + ) # Evict the selected victim victim_buffer_id = self.uid_to_buffer_id[victim_uid] @@ -820,6 +1052,60 @@ def get_available_buffer_slot(): ) self.uid_to_buffer_id[uid] = buffer_id self.buffer_id_to_uid[buffer_id] = uid + if self.lora_no_cpu_backup and lora_adapter is not None: + lora_adapter.release_staged_weights() + + def check_valid(self): + if self.invalid_uids: + raise RuntimeError( + f"LoRA adapters {sorted(self.invalid_uids)} have incomplete GPU installs; " + "reload or unload them before resuming inference." + ) + + def install_streamed_adapter( + self, + uid: str, + lora_adapter: LoRAAdapter, + lora_modules: List[Dict[str, torch.nn.Module]], + lora_embed_tokens_module: Optional[BaseLayerWithLoRA], + lora_lm_head_module: Optional[BaseLayerWithLoRA], + ) -> None: + """Validate a streamed adapter before overwriting its only resident copy.""" + if uid in self.uid_to_buffer_id: + buffer_id = self.uid_to_buffer_id[uid] + elif EMPTY_SLOT in self.buffer_id_to_uid: + buffer_id = self.buffer_id_to_uid.index(EMPTY_SLOT) + elif None in self.uid_to_buffer_id: + buffer_id = self.uid_to_buffer_id[None] + else: + raise RuntimeError( + f"No free LoRA memory pool slot for streamed adapter '{uid}'; " + "--lora-no-cpu-backup cannot evict, increase --max-loras-per-batch." + ) + load_args = ( + uid, + buffer_id, + lora_adapter, + lora_modules, + lora_embed_tokens_module, + lora_lm_head_module, + ) + self.load_lora_weight_to_buffer(*load_args, validate_only=True) + + # No backup exists: a failed write must block inference until a complete reload. + self.invalid_uids.add(uid) + if self.buffer_id_to_uid[buffer_id] is None: + del self.uid_to_buffer_id[None] + self.eviction_policy.remove(None) + self.uid_to_buffer_id[uid] = buffer_id + self.buffer_id_to_uid[buffer_id] = uid + self.load_lora_weight_to_buffer(*load_args) + device = next(self.base_model.parameters()).device + if device.type == "cuda": + torch.cuda.current_stream(device).synchronize() + self.invalid_uids.remove(uid) + self.eviction_policy.mark_used(uid) + lora_adapter.release_staged_weights() def _clear_buffer_slot_for_base(self, buffer_id: int) -> None: """Make an evicted slot safe for graph-captured base-model replay.""" @@ -846,6 +1132,7 @@ def remove_lora(self, uid: str) -> Optional[int]: del self.uid_to_buffer_id[uid] self.buffer_id_to_uid[buffer_id] = EMPTY_SLOT self.eviction_policy.remove(uid) + self.invalid_uids.discard(uid) return buffer_id def load_lora_weight_to_buffer( @@ -856,21 +1143,25 @@ def load_lora_weight_to_buffer( lora_modules: List[Dict[str, torch.nn.Module]], lora_embed_tokens_module: Optional[BaseLayerWithLoRA], lora_lm_head_module: Optional[BaseLayerWithLoRA], + *, + validate_only: bool = False, ): def load_lora_weight_tensor( buffer_view: torch.Tensor, weight: Optional[torch.Tensor] ): - if weight is None: - # If the particular weight is not present in the adapter, we initialize the buffer to zero - # to avoid contamination from the residual weight of the evicted adapters. - buffer_view.zero_() - else: + if weight is not None: assert buffer_view.shape == weight.shape, ( f"LoRA buffer shape {buffer_view.shape} does not match weight shape {weight.shape}." ) + if validate_only: + return + if weight is None: + buffer_view.zero_() + else: copy_weight_into_buffer(buffer_view, weight) if uid is None: + assert not validate_only self._clear_buffer_slot_for_base(buffer_id) return @@ -939,12 +1230,22 @@ def load_lora_weight_tensor( target_module: None for target_module in self.B_buffer } + fused_moe = self._get_fused_shared_moe(self.base_model, layer_id) for name, weights in layer_weights.items(): + cache_name = name + if fused_moe is not None and re.search(r"\.shared_experts?\.", name): + assert fused_moe.num_fused_shared_experts == 1 + assert weights.dim() == 2 + name = re.sub( + r"\.shared_experts?\.", + f".experts.{fused_moe._num_global_routed}.", + name, + ) target_module = get_target_module_name(name, self.target_modules) # Check if this is an MoE weight (has expert index in name) expert_match = re.search(r"experts\.(\d+)\.", name) - is_shared_expert = "shared_experts." in name + is_shared_expert = re.search(r"\.shared_experts?\.", name) is not None shared_moe_target = f"{target_module}_shared_moe" if is_shared_expert and ( @@ -963,49 +1264,92 @@ def load_lora_weight_tensor( expert_id = int(expert_match.group(1)) if "lora_A" in name: temp_A_buffer[target_module][expert_id] = weights - temp_A_cache_keys[target_module][expert_id] = name + temp_A_cache_keys[target_module][expert_id] = cache_name else: temp_B_buffer[target_module][expert_id] = weights - temp_B_cache_keys[target_module][expert_id] = name + temp_B_cache_keys[target_module][expert_id] = cache_name elif "lora_A" in name: temp_A_buffer[target_module] = weights - temp_A_cache_keys[target_module] = name + temp_A_cache_keys[target_module] = cache_name else: temp_B_buffer[target_module] = weights - temp_B_cache_keys[target_module] = name + temp_B_cache_keys[target_module] = cache_name elif expert_match: - # Per-expert MoE weight — 2D tensors, one per expert + # Per-expert MoE weight — 2D tensors, one per expert. + # Init A and B INDEPENDENTLY (both buffer and cache_keys). Under + # ``experts_shared_outer_loras`` one side of a projection is a + # shared 3D Tensor (set by the dim()==3 branch below) and the + # other is this per-expert dict: fc1 = shared A + per-expert B, + # fc2 = the opposite. The old coupled init keyed every dict off + # ``temp_A_buffer is None``, which either left the per-expert + # side's cache_keys as None (-> TypeError at + # ``[expert_id] = name``) or clobbered the shared side the 3D + # branch already populated. So each side now guards its own + # buffer + cache_keys and never touches the other side. target_module = target_module + "_moe" - if temp_A_buffer[target_module] is None: - temp_A_buffer[target_module] = {} - temp_B_buffer[target_module] = {} - temp_A_cache_keys[target_module] = {} - temp_B_cache_keys[target_module] = {} - expert_id = int(expert_match.group(1)) if "lora_A" in name: + assert not isinstance( + temp_A_buffer[target_module], torch.Tensor + ), ( + f"{target_module} lora_A already holds a shared-outer 3D " + f"tensor but also got per-expert weight '{name}'; a " + f"projection side must use one layout, not both." + ) + if temp_A_buffer[target_module] is None: + temp_A_buffer[target_module] = {} + if temp_A_cache_keys[target_module] is None: + temp_A_cache_keys[target_module] = {} temp_A_buffer[target_module][expert_id] = weights - temp_A_cache_keys[target_module][expert_id] = name + temp_A_cache_keys[target_module][expert_id] = cache_name else: + assert not isinstance( + temp_B_buffer[target_module], torch.Tensor + ), ( + f"{target_module} lora_B already holds a shared-outer 3D " + f"tensor but also got per-expert weight '{name}'; a " + f"projection side must use one layout, not both." + ) + if temp_B_buffer[target_module] is None: + temp_B_buffer[target_module] = {} + if temp_B_cache_keys[target_module] is None: + temp_B_cache_keys[target_module] = {} temp_B_buffer[target_module][expert_id] = weights - temp_B_cache_keys[target_module][expert_id] = name + temp_B_cache_keys[target_module][expert_id] = cache_name elif "experts" in name and weights.dim() == 3: - # Shared outer MoE weight — 3D tensor [expert_dim, rank, hidden] + # Packed routed or shared-outer MoE factors target_module = target_module + "_moe" - if "lora_A" in name: + if fused_moe is not None: + # Merge packed routed factors with separately named shared factors. + buffer, cache_keys = ( + (temp_A_buffer, temp_A_cache_keys) + if "lora_A" in name + else (temp_B_buffer, temp_B_cache_keys) + ) + if buffer[target_module] is None: + buffer[target_module] = {} + cache_keys[target_module] = {} + for expert_id, weight in enumerate(weights): + buffer[target_module][expert_id] = weight + cache_keys[target_module][expert_id] = ( + append_cache_key_suffix( + cache_name, f"expert{expert_id}" + ) + ) + elif "lora_A" in name: temp_A_buffer[target_module] = weights - temp_A_cache_keys[target_module] = name + temp_A_cache_keys[target_module] = cache_name else: temp_B_buffer[target_module] = weights - temp_B_cache_keys[target_module] = name + temp_B_cache_keys[target_module] = cache_name else: # Standard weight — single tensor per module if "lora_A" in name: temp_A_buffer[target_module] = weights - temp_A_cache_keys[target_module] = name + temp_A_cache_keys[target_module] = cache_name else: temp_B_buffer[target_module] = weights - temp_B_cache_keys[target_module] = name + temp_B_cache_keys[target_module] = cache_name # Track which buffer keys correspond to a real wrapped module on # this layer. `temp_A/B_buffer` is seeded with every key in the @@ -1163,7 +1507,7 @@ def load_lora_weight_tensor( # Place each stacked component at max_rank-spaced # positions so the kernel's [:max_r] / [max_r:2*max_r] # slicing is correct. - target_buffer[buffer_id, 0].zero_() + load_lora_weight_tensor(target_buffer[buffer_id, 0], None) if representative_weight is not None: for ci in range(c): buffer_view = target_buffer[ @@ -1176,14 +1520,14 @@ def load_lora_weight_tensor( ], ) elif weights is None: - target_buffer[buffer_id].zero_() + load_lora_weight_tensor(target_buffer[buffer_id], None) elif isinstance(weights, (torch.Tensor, dict)): # Zero first so any local-expert slot the adapter # doesn't fill (e.g. out-of-rank under EP) is clean; # then load owned slots at max_rank-spaced offsets so # the MoE kernel's [:max_r] / [max_r:2*max_r] slicing # is correct. - target_buffer[buffer_id].zero_() + load_lora_weight_tensor(target_buffer[buffer_id], None) assert isinstance(weights_cache_key, (str, dict)) for ( local_eid, @@ -1193,6 +1537,7 @@ def load_lora_weight_tensor( weights, weights_cache_key, localize=not self.is_shared_moe_module(name), + fused_moe=fused_moe, ): if expert_weight is None: continue @@ -1292,14 +1637,16 @@ def load_lora_weight_tensor( f"shape={weights.shape if isinstance(weights, torch.Tensor) else 'N/A'}" ) # Zero beyond loaded rank — MoE kernel reads full max_rank. - target_buffer[buffer_id, 0, :, lora_rank:].zero_() + load_lora_weight_tensor( + target_buffer[buffer_id, 0, :, lora_rank:], None + ) elif weights is None: - target_buffer[buffer_id].zero_() + load_lora_weight_tensor(target_buffer[buffer_id], None) elif isinstance(weights, (torch.Tensor, dict)): # Zero out slots this rank owns but the adapter # doesn't fill (padded-out / out-of-rank experts); # then scale+load the ones it does. - target_buffer[buffer_id].zero_() + load_lora_weight_tensor(target_buffer[buffer_id], None) assert isinstance(weights_cache_key, (str, dict)) for ( local_eid, @@ -1309,6 +1656,7 @@ def load_lora_weight_tensor( weights, weights_cache_key, localize=not self.is_shared_moe_module(name), + fused_moe=fused_moe, ): if w is not None: w = w * lora_adapter.scaling @@ -1334,7 +1682,9 @@ def load_lora_weight_tensor( if _SGLANG_EXPERIMENTAL_LORA_OPTI: # Zero beyond loaded rank: the experimental dense LoRA-B kernel # contracts over the full padded max_rank, so the tail must be clean. - target_buffer[buffer_id, :, lora_rank:].zero_() + load_lora_weight_tensor( + target_buffer[buffer_id, :, lora_rank:], None + ) if lora_adapter.embedding_layers: org_vocab_size = self.base_hf_config.vocab_size @@ -1465,18 +1815,20 @@ def load_lora_weight_tensor( # Zero out embedding/lm_head buffers for adapters without embedding LoRA # to avoid using garbage values from uninitialized memory for k in self.embedding_A_buffer.keys(): - self.embedding_A_buffer[k][buffer_id].zero_() + load_lora_weight_tensor(self.embedding_A_buffer[k][buffer_id], None) for k in self.embedding_B_buffer.keys(): - self.embedding_B_buffer[k][buffer_id].zero_() + load_lora_weight_tensor(self.embedding_B_buffer[k][buffer_id], None) for k in self.lm_head_A_buffer.keys(): - self.lm_head_A_buffer[k][buffer_id].zero_() + load_lora_weight_tensor(self.lm_head_A_buffer[k][buffer_id], None) for k in self.lm_head_B_buffer.keys(): - self.lm_head_B_buffer[k][buffer_id].zero_() + load_lora_weight_tensor(self.lm_head_B_buffer[k][buffer_id], None) if ( self.lora_added_tokens_size > 0 and "input_embeddings" in self.new_embeddings_buffer ): - self.new_embeddings_buffer["input_embeddings"][buffer_id].zero_() + load_lora_weight_tensor( + self.new_embeddings_buffer["input_embeddings"][buffer_id], None + ) def get_embedding_tensor( self, target_module: str, lora_type: LoRAType diff --git a/python/sglang/srt/lora/utils.py b/python/sglang/srt/lora/utils.py index ce764275a32a..e982d1d01123 100644 --- a/python/sglang/srt/lora/utils.py +++ b/python/sglang/srt/lora/utils.py @@ -153,6 +153,13 @@ def get_hidden_dim( ) +def shared_experts_count(config: AutoConfig) -> Optional[int]: + # DeepSeek-style configs say n_shared_experts, Kimi K3 says num_shared_experts + return getattr(config, "n_shared_experts", None) or getattr( + config, "num_shared_experts", None + ) + + def get_default_hidden_dim( module_name: str, config: AutoConfig, @@ -188,7 +195,7 @@ def get_default_hidden_dim( moe_freq = getattr(config, "moe_layer_freq", 1) if first_k is not None and layer_idx >= first_k and layer_idx % moe_freq == 0: moe_inter = getattr(config, "moe_intermediate_size", None) - n_shared = getattr(config, "n_shared_experts", None) + n_shared = shared_experts_count(config) if moe_inter is not None and n_shared is not None: inter = moe_inter * n_shared return config.hidden_size, inter * 2 @@ -198,7 +205,7 @@ def get_default_hidden_dim( moe_freq = getattr(config, "moe_layer_freq", 1) if first_k is not None and layer_idx >= first_k and layer_idx % moe_freq == 0: moe_inter = getattr(config, "moe_intermediate_size", None) - n_shared = getattr(config, "n_shared_experts", None) + n_shared = shared_experts_count(config) if moe_inter is not None and n_shared is not None: inter = moe_inter * n_shared return inter, config.hidden_size @@ -240,12 +247,19 @@ def get_default_hidden_dim( moe_inter = ( getattr(config, "moe_intermediate_size", None) or config.intermediate_size ) - return config.hidden_size, moe_inter * 2 + # Latent MoE (Kimi K3) routes a compressed activation into the experts. + moe_hidden = ( + getattr(config, "routed_expert_hidden_size", None) or config.hidden_size + ) + return moe_hidden, moe_inter * 2 elif module_name == "down_proj_moe": moe_inter = ( getattr(config, "moe_intermediate_size", None) or config.intermediate_size ) - return moe_inter, config.hidden_size + moe_hidden = ( + getattr(config, "routed_expert_hidden_size", None) or config.hidden_size + ) + return moe_inter, moe_hidden elif module_name == "embed_tokens": # For embedding: input is vocab_size (as embedding lookup), output is hidden_size # if contain extra tokens will be added; otherwise is 0. @@ -260,6 +274,7 @@ def get_default_hidden_dim( def get_normalized_target_modules( target_modules: Union[str, Iterable[str]], + base_model: Optional[torch.nn.Module] = None, ) -> set[str]: """ Mapping a list of target module name to names of the normalized LoRA weights. @@ -285,6 +300,13 @@ def get_normalized_target_modules( "v_proj": "qkv_proj", "gate_proj": "gate_up_proj", "up_proj": "gate_up_proj", + "in_proj_q": "in_proj_qkvz", + "in_proj_k": "in_proj_qkvz", + "in_proj_v": "in_proj_qkvz", + "in_proj_qkv": "in_proj_qkvz", + "in_proj_z": "in_proj_qkvz", + "in_proj_b": "in_proj_ba", + "in_proj_a": "in_proj_ba", "out_proj": "out_proj", "embed_tokens": "embed_tokens", "vocab_emb": "embed_tokens", @@ -305,6 +327,11 @@ def get_normalized_target_modules( result = set() for name in target_modules: + # Models may fuse adapter projections differently from the common HF layout. + if base_model is not None and hasattr( + base_model, "get_lora_target_module_name" + ): + name = base_model.get_lora_target_module_name(name) base_name = name.split(".")[-1] normalized_name = params_mapping.get(base_name, base_name) result.add(normalized_name) @@ -323,6 +350,7 @@ def get_stacked_multiply( stacked_rank = { "qkv_proj": 3, "in_proj_qkvz": 4, # GDN packed input projection + "in_proj_ba": 2, # GDN packed b/a input projection "gate_up_proj": 2, "gate_up_proj_moe": 2, "gate_up_proj_shared_moe": 2, @@ -402,6 +430,7 @@ def get_target_module_name(full_module_name: str, target_modules: Set[str]) -> s "out_proj", "in_proj", "in_proj_qkvz", + "in_proj_ba", "up_proj", "gate_up_proj", "down_proj", @@ -458,7 +487,7 @@ def auto_detect_lora_target_modules(model: "torch.nn.Module") -> set: else: raw_names.add(leaf_name) - normalized = get_normalized_target_modules(raw_names) + normalized = get_normalized_target_modules(raw_names, base_model=model) result = normalized & _KNOWN_LORA_TARGET_MODULES # Allow models to declare additional LoRA-compatible modules that @@ -492,6 +521,8 @@ def get_lm_head_lora_b_shard_size(output_dim: int, shard_indices=None) -> int: def get_batch_token_counts(forward_batch: ForwardBatch) -> Tuple[int, int]: """(total tokens, max tokens per request) for LoRA segment math.""" mode = forward_batch.forward_mode + if mode.is_idle(): + return 0, 0 if mode.is_decode(): return forward_batch.batch_size, 1 if mode.is_target_verify(): @@ -508,7 +539,9 @@ def generate_sequence_lengths( device = torch.get_default_device() if device is None else device with torch.device(device): - if forward_batch.forward_mode.is_decode(): + if forward_batch.forward_mode.is_idle(): + seg_lens = torch.empty(0, dtype=torch.int32) + elif forward_batch.forward_mode.is_decode(): seg_lens = torch.ones(forward_batch.batch_size, dtype=torch.int32) elif forward_batch.forward_mode.is_target_verify(): seg_lens = torch.full( diff --git a/python/sglang/srt/managers/detokenizer_manager.py b/python/sglang/srt/managers/detokenizer_manager.py index a15bb1d7b136..7d6731d5ed0e 100644 --- a/python/sglang/srt/managers/detokenizer_manager.py +++ b/python/sglang/srt/managers/detokenizer_manager.py @@ -495,10 +495,10 @@ def handle_batch_token_id_out(self, recv_obj: BatchTokenIDOutput): output_token_ids_logprobs_idx=recv_obj.output_token_ids_logprobs_idx, output_token_entropy_val=recv_obj.output_token_entropy_val, output_token_sampling_mask=recv_obj.output_token_sampling_mask, - output_token_sampling_logprobs=recv_obj.output_token_sampling_logprobs, output_hidden_states=recv_obj.output_hidden_states, routed_experts=routed_experts, indexer_topk=indexer_topk, + indexer_topk_num_layers=recv_obj.indexer_topk_num_layers, customized_info=recv_obj.customized_info, placeholder_tokens_idx=None, placeholder_tokens_val=None, diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index d3f13eba8917..0c7ae04b2328 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -58,9 +58,11 @@ Modality, MultimodalProcessorOutput, ReturnHiddenStatesMode, + SamplingLogprobsMode, get_return_hidden_states_mode, ) from sglang.srt.multimodal.mm_utils import has_valid_data +from sglang.srt.sampling.sampling_mask import SamplingMaskChunk from sglang.srt.sampling.sampling_params import SamplingParams from sglang.srt.utils import ImageData, VideoData from sglang.srt.utils.field_validators import validate_optional_list_i64_1d_2d @@ -231,8 +233,14 @@ class GenerateReqInput: top_logprobs_num: Optional[Union[List[int], int]] = None # If return logprobs, the token ids to return logprob for. token_ids_logprob: Optional[Union[List[List[int]], List[int]]] = None - # Whether to return output-token sampling support and renormalized logprobs. + # Whether to return each output token's sampling support and behavior logprob. return_sampling_mask: Optional[Union[List[bool], bool]] = None + # Return either the selected token's behavior logprob or behavior logprobs for + # the full sampling support. When omitted, selected mode is used if + # return_sampling_mask is enabled. + sampling_logprobs_mode: Optional[ + Union[List[Optional[SamplingLogprobsMode]], SamplingLogprobsMode] + ] = None # Whether to detokenize tokens in text in the returned logprobs. return_text_in_logprobs: bool = False # Return prompt top logprobs as flat arrays plus shape metadata instead of @@ -265,6 +273,8 @@ class GenerateReqInput: lora_path: Optional[Union[List[Optional[str]], str]] = None # The uid of LoRA adaptors, should be initialized by tokenizer manager lora_id: Optional[Union[List[Optional[str]], str]] = None + # {lora_name: disk path}: backfill source when this engine does not hold the adapter + lora_backfill_paths: Optional[Dict[str, str]] = None # Custom logit processor for advanced sampling control. Must be a serialized instance # of `CustomLogitProcessor` in python/sglang/srt/sampling/custom_logit_processor.py @@ -751,6 +761,9 @@ def normalize_param(param, default_value, param_name): self.return_sampling_mask = normalize_param( self.return_sampling_mask, False, "return_sampling_mask" ) + self.sampling_logprobs_mode = normalize_param( + self.sampling_logprobs_mode, None, "sampling_logprobs_mode" + ) # Handle token_ids_logprob specially due to its nested structure if not self.token_ids_logprob: # covers both None and [] @@ -912,6 +925,7 @@ def __getitem__(self, i): top_logprobs_num=self.top_logprobs_num[i], token_ids_logprob=self.token_ids_logprob[i], return_sampling_mask=self.return_sampling_mask[i], + sampling_logprobs_mode=self.sampling_logprobs_mode[i], return_text_in_logprobs=self.return_text_in_logprobs, return_flat_raw_top_logprobs=self.return_flat_raw_top_logprobs, return_flat_raw_top_logprobs_b64=self.return_flat_raw_top_logprobs_b64, @@ -1078,6 +1092,10 @@ class TokenizedGenerateReqInput(BaseReq, kw_only=True): # Cache namespace used to isolate otherwise-identical prefixes. cache_salt: Optional[str] = None + # Shape of output_token_sampling_logprobs for each output token. This is a + # defaulted tail field so older IPC senders decode as selected mode. + sampling_logprobs_mode: SamplingLogprobsMode = "selected" + def wrap_pickle_fields(self): self.time_stats = wrap_as_pickle(self.time_stats) @@ -1143,6 +1161,8 @@ class EmbeddingReqInput: lora_path: Optional[Union[List[Optional[str]], str]] = None # The uid of LoRA adaptors, should be initialized by tokenizer manager lora_id: Optional[Union[List[Optional[str]], str]] = None + # {lora_name: disk path}: backfill source when this engine does not hold the adapter + lora_backfill_paths: Optional[Dict[str, str]] = None # Resolved embedding overrides with positions (set by tokenizer manager or score mixin). # Runtime type: Optional[Union[PositionalEmbeds, List[Optional[PositionalEmbeds]]]] positional_embed_overrides: Any = None @@ -1466,12 +1486,7 @@ class BatchTokenIDOutput(BaseBatchReq, kw_only=True): output_token_ids_logprobs_val: TokenIdsLogprobValues output_token_ids_logprobs_idx: TokenIdsLogprobIndices output_token_entropy_val: Optional[List[Optional[float]]] - # Per-request chunks of output-token sampling supports. None when no request - # in the batch asks for return_sampling_mask. - output_token_sampling_mask: Optional[List[List]] - # Per-request chunks of selected-token logprobs renormalized over the - # corresponding sampling supports. None when sampling masks are not returned. - output_token_sampling_logprobs: Optional[List[List]] + output_token_sampling_mask: Optional[List[Optional[SamplingMaskChunk]]] # Hidden states output_hidden_states: OutputHiddenStates @@ -1513,6 +1528,10 @@ class BatchTokenIDOutput(BaseBatchReq, kw_only=True): # Pickled Optional[List[SchedulerReqTimeStats]] time_stats: Optional[PickleWrapper] = None + # Number of indexer layers, set when indexer_topk is non-empty. The + # rollout-indexer-replay consumer needs the layer count to reshape the + # flattened per-token topk back into (token, layer, topk). + indexer_topk_num_layers: Optional[int] = None # Multimodal prompt token counts (image/audio/video). None when not applicable. image_tokens: Optional[List[int]] = None audio_tokens: Optional[List[int]] = None @@ -1566,10 +1585,7 @@ class BatchStrOutput(BaseBatchReq, kw_only=True): output_token_ids_logprobs_val: TokenIdsLogprobValues output_token_ids_logprobs_idx: TokenIdsLogprobIndices output_token_entropy_val: Optional[List[Optional[float]]] - # Detokenizer pass-through for BatchTokenIDOutput.output_token_sampling_*. - # None when sampling masks are not returned. - output_token_sampling_mask: Optional[List[List]] - output_token_sampling_logprobs: Optional[List[List]] + output_token_sampling_mask: Optional[List[Optional[SamplingMaskChunk]]] # Hidden states output_hidden_states: OutputHiddenStates @@ -1610,6 +1626,9 @@ class BatchStrOutput(BaseBatchReq, kw_only=True): # Pickled Optional[List[SchedulerReqTimeStats]] time_stats: Optional[PickleWrapper] = None + # Number of indexer layers, set when indexer_topk is non-empty (see + # BatchTokenIDOutput.indexer_topk_num_layers). + indexer_topk_num_layers: Optional[int] = None # Multimodal prompt token counts (image/audio/video). None when not applicable. image_tokens: Optional[List[int]] = None audio_tokens: Optional[List[int]] = None @@ -1815,6 +1834,22 @@ class UpdateWeightFromDiskReqOutput(BaseReq, kw_only=True): num_paused_requests: int = 0 +class PullWeightsReqInput(BaseReq, kw_only=True): + # Host-local checkpoint dir the pulled weights land in; seeded from the + # server's model path when the published stream has no full version. + local_checkpoint_dir: str + # Shared dir the publisher writes weight_v{N:06d}/ version dirs under; each + # version is a full HF checkpoint or a delta against the previous version. + source_dir: str + # The version to bring the local checkpoint up to. + target_version: int + + +class PullWeightsReqOutput(BaseReq, kw_only=True): + success: bool + message: str + + class UpdateWeightsFromDistributedReqInput(BaseReq, kw_only=True): names: List[str] dtypes: List[str] @@ -1829,6 +1864,9 @@ class UpdateWeightsFromDistributedReqInput(BaseReq, kw_only=True): weight_version: Optional[str] = None # Optional format specification for loading load_format: Optional[str] = None + # Which model runners to update: "target" (target model only), "draft" (draft + # worker(s) only), or "all" (default). + selector: Literal["target", "draft", "all"] = "all" # Whether to call torch.cuda.empty_cache() during flush torch_empty_cache: bool = False @@ -1855,8 +1893,9 @@ class UpdateWeightsFromTensorReqInput(BaseReq, kw_only=True): abort_all_requests: bool = False # Optional: Update weight version along with weights weight_version: Optional[str] = None - # Optional: Determine whether to disable updating the draft model - disable_draft_model: Optional[bool] = None + # Which model runners to update: "target" (target model only), "draft" (draft + # worker(s) only), or "all" (default). + selector: Literal["target", "draft", "all"] = "all" # Whether to call torch.cuda.empty_cache() during flush torch_empty_cache: bool = False @@ -2008,15 +2047,56 @@ class ResumeMemoryOccupationReqOutput(BaseReq, kw_only=True): pass +class BeginWeightUpdateReqInput(BaseReq, kw_only=True): + """Open a weight-update session: unpack in-place-quantized weights on the + selected runners so fresh weights can be loaded into them.""" + + selector: Literal["target", "draft", "all"] = "all" + # Session scope: False = no base bytes will land this session (adapter-only + # sync), so the quant unpack/repack round-trip is skipped and a base-named + # tensor arriving in-session is an error. + sync_base: bool = True + + +class BeginWeightUpdateReqOutput(BaseReq, kw_only=True): + success: bool + message: str + + +class EndWeightUpdateReqInput(BaseReq, kw_only=True): + """Close the weight-update session opened by BeginWeightUpdateReqInput: + re-finalize base weights (sync_base sessions only) and apply the streamed + LoRA stash accumulated by update_weights_from_* during the session.""" + + # {lora_name: {hf_key: sha256}}; when set, each stashed adapter is verified + # (set equality + per-tensor checksum) before it is applied. + expected_lora_checksums: Optional[Dict[str, Dict[str, str]]] = None + # discard the streamed LoRA stash and deferred publications; in-place base writes are not rolled back + abort: bool = False + + +class EndWeightUpdateReqOutput(BaseReq, kw_only=True): + success: bool + message: str + + class CheckWeightsReqInput(BaseReq, kw_only=True): action: str = "checksum" allow_quant_error: bool = False + # Substrings of tensor names to exclude from reset/compare/checksum. + skip_tensor_list: Optional[List[str]] = None + # Which model runners to update: "target" (target model only), "draft" (draft + # worker(s) only), or "all" (default). + selector: Literal["target", "draft", "all"] = "all" # Wire versions of the pydantic ParallelismInfo/ChecksumInfo in # sglang.srt.utils.weight_checker. Not array_like: the payload is read by field # name and re-serialized to JSON, so it must stay a {field: value} map. class ParallelismInfo(msgspec.Struct, kw_only=True): + # Which runner this describes: "target", or a draft role such as "draft" / + # "draft_step_0". One entry per runner the checksum covers. + role: str tp_rank: int tp_size: int dp_rank: int @@ -2030,7 +2110,10 @@ class ParallelismInfo(msgspec.Struct, kw_only=True): class ChecksumInfo(msgspec.Struct, kw_only=True): checksums: Dict[str, str] per_gpu_checksum: str - parallelism_info: ParallelismInfo + # One entry per role the checksum covers: the target model plus, under + # speculative decoding, each draft runner. All roles on a rank share the GPU + # rank; consumers key off it to merge the shards. + parallelism_info: List[ParallelismInfo] class CheckWeightsReqOutput(BaseReq, kw_only=True): @@ -2056,6 +2139,13 @@ class AbortReq(BaseReq, kw_only=True): finished_reason: Optional[FinishReasonDict] = None abort_message: Optional[str] = None weight_versions: Optional[WeightVersionSpans] = None + # Treat ``rid`` as a namespace prefix. Note this only relaxes the + # tokenizer-side gate (which otherwise requires an exact live rid, with a + # single tokenizer worker): once the request reaches the scheduler, + # matching is prefix-based (``rid.startswith``) regardless of this flag, + # because batch requests derive child rids as ``f"{rid}_{i}"`` and an + # abort for the parent rid must cover them. + prefix: bool = False def __post_init__(self): # FIXME: This is a hack to keep the same with the old code @@ -2298,27 +2388,35 @@ def to_ref(self) -> LoRARef: ) -class LoadLoRAAdapterFromTensorsReqInput(BaseReq, kw_only=True): +class RegisterLoRAAdapterReqInput(BaseReq, kw_only=True): + """Create-or-refresh an adapter's identity and config (control plane). + + Weights are untouched by the caller: a new adapter starts zeroed, and + re-registering an existing name re-zeroes it — a slot's new tenant must + not serve its predecessor's weights. The weight bytes arrive later as + ``{lora_name}:{hf_key}``-prefixed tensors in the ordinary + update_weights_from_* stream and are applied at end_weight_update.""" + lora_name: str - # The PEFT adapter_config.json, already JSON — a tighter type would only add - # decode strictness with no benefit. + # The PEFT adapter_config.json fields (r, lora_alpha, target_modules, ...). config_dict: Dict[str, Any] - # One serialized copy of the adapter tensors per TP rank; each rank - # deserializes only its own copy. Same normalization conventions as - # UpdateWeightsFromTensorReqInput.serialized_named_tensors. - serialized_named_tensors: Annotated[List[bytes], Base64Bytes()] + # Unpinned: the pool refills a registered adapter lazily from its CPU copy, + # and pinning every slot would trip the anti-starvation check. pinned: bool = False - added_tokens_config: Optional[Dict[str, int]] = None lora_id: Optional[str] = None - load_format: Optional[str] = None - expected_checksums: Optional[Dict[str, str]] = None + # Disk artifact holding the same adapter (PEFT dir). With a path the + # adapter is reloadable: it may be LRU-evicted and refilled from disk. + lora_path: Optional[str] = None + # keep the name unservable until end_weight_update commits its session; fresh names only + defer_publish: bool = False def to_ref(self) -> LoRARef: return LoRARef( lora_id=self.lora_id, lora_name=self.lora_name, - lora_path="__tensor__", + lora_path=self.lora_path or "__stream__", pinned=self.pinned, + reloadable=self.lora_path is not None, ) @@ -2326,11 +2424,13 @@ class LoRAUpdateOutput(BaseReq, kw_only=True): success: bool error_message: Optional[str] = None loaded_adapters: Optional[Dict[str, Union[str, LoRARef]]] = None + # acks a defer_publish registration; the trainer fails closed without it + pending: bool = False -LoadLoRAAdapterReqOutput = UnloadLoRAAdapterReqOutput = ( - LoadLoRAAdapterFromTensorsReqOutput -) = LoRAUpdateOutput +LoadLoRAAdapterReqOutput = UnloadLoRAAdapterReqOutput = RegisterLoRAAdapterReqOutput = ( + LoRAUpdateOutput +) class BlockReqType(Enum): diff --git a/python/sglang/srt/managers/multi_tokenizer_mixin.py b/python/sglang/srt/managers/multi_tokenizer_mixin.py index ff88f27e8d67..150ee49055b5 100644 --- a/python/sglang/srt/managers/multi_tokenizer_mixin.py +++ b/python/sglang/srt/managers/multi_tokenizer_mixin.py @@ -247,9 +247,6 @@ def _handle_output_by_index(output, i): output_token_sampling_mask=_extract_field_by_index( output, "output_token_sampling_mask", i, check_length=False ), - output_token_sampling_logprobs=_extract_field_by_index( - output, "output_token_sampling_logprobs", i, check_length=False - ), output_hidden_states=_extract_field_by_index( output, "output_hidden_states", i, check_length=False ), @@ -259,6 +256,7 @@ def _handle_output_by_index(output, i): indexer_topk=_extract_field_by_index( output, "indexer_topk", i, check_length=False ), + indexer_topk_num_layers=getattr(output, "indexer_topk_num_layers", None), retraction_counts=_extract_field_by_index(output, "retraction_counts", i), weight_versions=_extract_field_by_index(output, "weight_versions", i), placeholder_tokens_idx=None, @@ -368,9 +366,6 @@ def _handle_output_by_index(output, i): output_token_sampling_mask=_extract_field_by_index( output, "output_token_sampling_mask", i, check_length=False ), - output_token_sampling_logprobs=_extract_field_by_index( - output, "output_token_sampling_logprobs", i, check_length=False - ), output_hidden_states=_extract_field_by_index( output, "output_hidden_states", i, check_length=False ), @@ -380,6 +375,7 @@ def _handle_output_by_index(output, i): indexer_topk=_extract_field_by_index( output, "indexer_topk", i, check_length=False ), + indexer_topk_num_layers=getattr(output, "indexer_topk_num_layers", None), customized_info=_extract_field_by_index( output, "customized_info", i, check_length=False ), diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 7811e4477a8f..038f6b1de94e 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -139,6 +139,7 @@ SchedulerReqTimeStats, ) from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo +from sglang.srt.sampling.sampling_mask import SamplingMaskRows from sglang.srt.sampling.sampling_params import SamplingParams from sglang.srt.utils import flatten_nested_list from sglang.srt.utils.token_sequence_matcher import TokenSequenceMatcher @@ -164,6 +165,7 @@ ReturnHiddenStatesMode = Union[bool, Literal["last"]] +SamplingLogprobsMode = Literal["selected", "support"] def get_return_hidden_states_mode( @@ -943,6 +945,7 @@ def __init__( dllm_config: Optional[DllmConfig] = None, token_ids_logprob: List[int] = None, return_sampling_mask: bool = False, + sampling_logprobs_mode: SamplingLogprobsMode = "selected", return_flat_raw_top_logprobs: bool = False, stream: bool = False, origin_input_ids_unpadded: Optional[array[int]] = None, @@ -1166,7 +1169,6 @@ def __init__( # TODO (Byron): send_output_token_logprobs_offset and send_decode_id_offset can be different in disaggregation mode # because the decode server does not have the first output token logprobs self.send_output_token_logprobs_offset: int = 0 - self.send_output_sampling_mask_offset: int = 0 # Logprobs (arguments) self.return_logprob = return_logprob @@ -1179,6 +1181,7 @@ def __init__( self.temp_scaled_logprobs = False self.top_p_normalized_logprobs = False self.return_sampling_mask = return_sampling_mask + self.sampling_logprobs_mode = sampling_logprobs_mode self.return_flat_raw_top_logprobs = return_flat_raw_top_logprobs # Logprobs (return values) @@ -1201,12 +1204,9 @@ def __init__( # Can contain either lists or GPU tensors (delayed copy optimization for prefill-only scoring) self.logprob.output_token_ids_logprobs_val = [] self.logprob.output_token_ids_logprobs_idx = [] - if return_sampling_mask: - self.output_token_sampling_mask = [] - self.output_token_sampling_logprobs = [] - else: - self.output_token_sampling_mask = None - self.output_token_sampling_logprobs = None + self.sampling_mask_rows: Optional[SamplingMaskRows] = ( + SamplingMaskRows() if return_sampling_mask else None + ) self.hidden_states: List[List[float]] = [] self.hidden_states_tensor = None # Note: use tensor instead of list to transfer hidden_states when PD + MTP self.output_topk_p = None @@ -3176,8 +3176,10 @@ def retract_decode(self) -> Tuple[List[Req], float, List[Req]]: while first_iter or ( not self.check_decode_mem(selected_indices=sorted_indices) ): - if len(sorted_indices) == 1: - # Always keep at least one request + # We should allow all requests to be retracted in decode disaggregation mode + # because there can be prealloc prefill requests. + num_minimum_reqs = 0 if get_disagg().disaggregation_mode == "decode" else 1 + if len(sorted_indices) == num_minimum_reqs: break first_iter = False diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 743c8afe067f..5b5ef7a38adc 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -72,6 +72,7 @@ is_deepseek_v4, is_minimax_sparse, ) +from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE from sglang.srt.constrained.grammar_manager import GrammarManager from sglang.srt.debug_utils.pr_fix_toggle import maybe_revert_pr_fix from sglang.srt.disaggregation.decode import ( @@ -125,6 +126,7 @@ AttachHiCacheStorageReqOutput, BatchTokenizedEmbeddingReqInput, BatchTokenizedGenerateReqInput, + BeginWeightUpdateReqInput, CheckWeightsReqInput, ClearHiCacheReqInput, ClearHiCacheReqOutput, @@ -136,6 +138,7 @@ DetachHiCacheStorageReqOutput, DumperControlReqInput, DumperControlReqOutput, + EndWeightUpdateReqInput, ExpertDistributionReq, ExpertDistributionReqOutput, ExpertDistributionReqType, @@ -151,14 +154,15 @@ InitWeightsUpdateGroupReqInput, ListExternalCorporaReqInput, ListExternalCorporaReqOutput, - LoadLoRAAdapterFromTensorsReqInput, - LoadLoRAAdapterFromTensorsReqOutput, LoadLoRAAdapterReqInput, LoadLoRAAdapterReqOutput, MMInputsProcessError, OpenSessionReqInput, PauseGenerationReqInput, ProfileReq, + PullWeightsReqInput, + RegisterLoRAAdapterReqInput, + RegisterLoRAAdapterReqOutput, ReleaseMemoryOccupationReqInput, RemoveExternalCorpusReqInput, RemoveExternalCorpusReqOutput, @@ -856,6 +860,7 @@ def init_idle_sleeper(self) -> None: self.ipc_channels.recv_from_tokenizer, self.ipc_channels.recv_from_rpc, ], + can_empty_cache=lambda: not self._engine_paused, ) else: self.idle_sleeper = None @@ -1760,6 +1765,14 @@ def init_request_dispatcher(self): SendWeightsToRemoteInstanceReqInput, self.send_weights_to_remote_instance, ), + ( + BeginWeightUpdateReqInput, + self.weight_updater.begin_weight_update, + ), + ( + EndWeightUpdateReqInput, + self.weight_updater.end_weight_update, + ), ( UpdateWeightsFromDistributedReqInput, self.weight_updater.update_weights_from_distributed, @@ -1792,6 +1805,10 @@ def init_request_dispatcher(self): CheckWeightsReqInput, self.weight_updater.check_weights, ), + ( + PullWeightsReqInput, + self.weight_updater.pull_weights, + ), (SlowDownReqInput, self.slow_down), ( ProfileReq, @@ -1804,10 +1821,7 @@ def init_request_dispatcher(self): (RpcReqInput, self.handle_rpc_request), (ExpertDistributionReq, self.expert_distribution_handle), (LoadLoRAAdapterReqInput, self.load_lora_adapter), - ( - LoadLoRAAdapterFromTensorsReqInput, - self.load_lora_adapter_from_tensors, - ), + (RegisterLoRAAdapterReqInput, self.register_lora_adapter), (UnloadLoRAAdapterReqInput, self.unload_lora_adapter), (PauseGenerationReqInput, self.pause_generation), (ContinueGenerationReqInput, self.continue_generation), @@ -2282,7 +2296,9 @@ def maybe_init_rust_server(self) -> None: # drains its request ring (rust_server_mode) instead of a zmq socket. self.recv_from_tokenizer = rust_server # Park the idle loop on the request ring within the rank-0 rust-server - self.idle_sleeper = RustServerIdleSleeper(rust_server) + self.idle_sleeper = RustServerIdleSleeper( + rust_server, can_empty_cache=lambda: not self._engine_paused + ) def rust_server_tokenizer_path(self) -> str: return get_serving().tokenizer_path @@ -2461,6 +2477,7 @@ def init_output_streamer(self) -> None: ps=self.ps, server_args=self.server_args, is_generation=self.is_generation, + is_multimodal_gen=getattr(self.model_config, "is_multimodal_gen", False), spec_algorithm=self.spec_algorithm, disaggregation_mode=self.disaggregation_mode, enable_hicache_storage=lambda: self.enable_hicache_storage, @@ -2755,6 +2772,7 @@ def handle_generate_request( top_logprobs_num=recv_req.top_logprobs_num, token_ids_logprob=recv_req.token_ids_logprob, return_sampling_mask=recv_req.return_sampling_mask, + sampling_logprobs_mode=recv_req.sampling_logprobs_mode, return_flat_raw_top_logprobs=recv_req.return_flat_raw_top_logprobs, stream=recv_req.stream, lora_id=recv_req.lora_id, @@ -4805,7 +4823,7 @@ def _record_scheduler_state_for_paused_engine(self) -> None: else: self.metrics_reporter.record_scheduler_active() - def is_fully_idle(self, for_health_check=False) -> bool: + def is_fully_idle(self, for_health_check=False, ignore_waiting=False) -> bool: # Health check piggybacks on running requests in process_output. # Only running_batch + waiting_queue guarantee active GPU processing; # disagg queues (bootstrap/prealloc/transfer) may have items without @@ -4822,7 +4840,8 @@ def is_fully_idle(self, for_health_check=False) -> bool: ) # Waiting queues: waiting + bootstrapping + preallocation + kv transfer (decode) - idle &= len(self.waiting_queue) == 0 + if not ignore_waiting: + idle &= len(self.waiting_queue) == 0 if ( for_health_check @@ -4980,7 +4999,10 @@ def detach_hicache_storage_wrapped( def flush_cache(self, empty_cache: bool = True): """Flush memory pools (e.g., KV cache, Mamba cache) and optionally empty device allocator cache.""" - if self.is_fully_idle(): + if GPU_MEMORY_TYPE_KV_CACHE in self.weight_updater.offload_tags: + logger.info("Cache flush skipped: KV cache memory is released.") + return True + if self.is_fully_idle(ignore_waiting=self._engine_paused): self.cur_batch_for_debug = None self.last_batch = None self.tree_cache.reset() @@ -5374,6 +5396,7 @@ def _pause_engine(self) -> Tuple[List[Req], int]: def pause_generation(self, recv_req: PauseGenerationReqInput): assert recv_req.mode in ("in_place", "retract") + self._engine_paused = True if recv_req.mode == "in_place": @@ -5570,20 +5593,22 @@ def load_lora_adapter( result = self.tp_worker.load_lora_adapter(recv_req) return result - def load_lora_adapter_from_tensors( - self, recv_req: LoadLoRAAdapterFromTensorsReqInput - ) -> LoadLoRAAdapterFromTensorsReqOutput: - """In-place loading a new lora adapter from serialized tensors.""" - - result = self.tp_worker.load_lora_adapter_from_tensors(recv_req) - return result - def unload_lora_adapter( self, recv_req: UnloadLoRAAdapterReqInput ) -> UnloadLoRAAdapterReqOutput: """Unload the lora adapter.""" result = self.tp_worker.unload_lora_adapter(recv_req) + self.weight_updater.forget_lora_adapter(recv_req.lora_name) + return result + + def register_lora_adapter( + self, recv_req: RegisterLoRAAdapterReqInput + ) -> RegisterLoRAAdapterReqOutput: + """Create-or-refresh a LoRA adapter's identity and config (weights zeroed).""" + + result = self.tp_worker.register_lora_adapter(recv_req) + self.weight_updater.forget_lora_adapter(recv_req.lora_name) return result def init_weights_send_group_for_remote_instance( diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py index 386e8d10f6cd..91e71d7c55b6 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py @@ -1162,11 +1162,9 @@ def add_sampling_mask_return_values( output: LogitsProcessorOutput, ) -> None: """Attach sparse sampling support metadata to the return values.""" - mask = output.next_token_sampling_mask_idx - logprobs = output.next_token_sampling_logprobs - req.output_token_sampling_mask.append(None if mask is None else mask[i]) - req.output_token_sampling_logprobs.append( - None if logprobs is None else logprobs[i] + req.sampling_mask_rows.append( + output.next_token_sampling_mask_idx[i], + output.next_token_sampling_logprobs[i], ) @staticmethod @@ -1174,14 +1172,13 @@ def materialize_sampling_mask_output( reqs: List[Req], output: Optional[LogitsProcessorOutput], ) -> None: - """Convert opted-in tensor rows to batch-aligned Python results.""" + """Convert opted-in tensor rows to batch-aligned host rows.""" if output is None or output.sampling_mask_output is None: return sampling_output = output.sampling_mask_output batch_indices = [i for i, req in enumerate(reqs) if req.return_sampling_mask] lengths = sampling_output.lengths.tolist() - selected_logprobs = sampling_output.selected_logprobs.tolist() statuses = sampling_output.statuses.tolist() assert len(batch_indices) == len(lengths) @@ -1189,17 +1186,29 @@ def materialize_sampling_mask_output( masks = [None] * batch_size logprobs = [None] * batch_size status_by_batch = [None] * batch_size - token_ids = sampling_output.token_ids.cpu() - packed_width = token_ids.shape[1] + token_ids = sampling_output.token_ids.cpu().numpy() + selected_logprobs = sampling_output.selected_logprobs.cpu().numpy() + support_logprobs = ( + None + if sampling_output.support_logprobs is None + else sampling_output.support_logprobs.cpu().numpy() + ) + support_row = 0 for row, batch_index in enumerate(batch_indices): + returns_support_logprobs = ( + reqs[batch_index].sampling_logprobs_mode == "support" + ) status = int(statuses[row]) length = int(lengths[row]) - if status == SamplingMaskStatus.OK and not (0 <= length <= packed_width): - status = SamplingMaskStatus.INVALID status_by_batch[batch_index] = status if status == SamplingMaskStatus.OK: - masks[batch_index] = token_ids[row, :length].tolist() - logprobs[batch_index] = float(selected_logprobs[row]) + masks[batch_index] = token_ids[row, :length] + if returns_support_logprobs: + logprobs[batch_index] = support_logprobs[support_row, :length] + else: + logprobs[batch_index] = selected_logprobs[row : row + 1] + if returns_support_logprobs: + support_row += 1 output.next_token_sampling_mask_idx = masks output.next_token_sampling_logprobs = logprobs diff --git a/python/sglang/srt/managers/scheduler_components/dp_attn.py b/python/sglang/srt/managers/scheduler_components/dp_attn.py index fcaf5ef2da5e..2487f708bd17 100644 --- a/python/sglang/srt/managers/scheduler_components/dp_attn.py +++ b/python/sglang/srt/managers/scheduler_components/dp_attn.py @@ -340,6 +340,8 @@ def _local_prefill_cuda_graph_vote( if prefill_graph_runner is None: return True + if not isinstance(prefill_graph_runner, PrefillCudaGraphRunner): + return False return prefill_graph_runner.can_replay_locally( batch_size=local_batch.batch_size(), num_tokens=num_tokens, diff --git a/python/sglang/srt/managers/scheduler_components/idle_sleeper.py b/python/sglang/srt/managers/scheduler_components/idle_sleeper.py index 54099ae3ab06..83127e824711 100644 --- a/python/sglang/srt/managers/scheduler_components/idle_sleeper.py +++ b/python/sglang/srt/managers/scheduler_components/idle_sleeper.py @@ -24,9 +24,10 @@ class IdleSleeper: data that needs handling immediately. """ - def __init__(self, sockets): + def __init__(self, sockets, can_empty_cache=None): self.poller = zmq.Poller() self.last_empty_time = real_time() + self.can_empty_cache = can_empty_cache for s in sockets: self.poller.register(s, zmq.POLLIN) @@ -39,7 +40,8 @@ def maybe_sleep(self): and real_time() - self.last_empty_time > self.empty_cache_interval ): self.last_empty_time = real_time() - current_platform.empty_cache() + if self.can_empty_cache is None or self.can_empty_cache(): + current_platform.empty_cache() class RustServerIdleSleeper: @@ -52,10 +54,13 @@ class RustServerIdleSleeper: requests — or the timeout elapses. """ - def __init__(self, rust_server: RustServer, timeout_ms: int = 1000): + def __init__( + self, rust_server: RustServer, timeout_ms: int = 1000, can_empty_cache=None + ): self.rust_server = rust_server self.timeout_ms = timeout_ms self.last_empty_time = real_time() + self.can_empty_cache = can_empty_cache self.empty_cache_interval = envs.SGLANG_EMPTY_CACHE_INTERVAL.get() def maybe_sleep(self): @@ -65,4 +70,5 @@ def maybe_sleep(self): and real_time() - self.last_empty_time > self.empty_cache_interval ): self.last_empty_time = real_time() - current_platform.empty_cache() + if self.can_empty_cache is None or self.can_empty_cache(): + current_platform.empty_cache() diff --git a/python/sglang/srt/managers/scheduler_components/output_streamer.py b/python/sglang/srt/managers/scheduler_components/output_streamer.py index 15659cf2dcf3..05e2f21a5ee2 100644 --- a/python/sglang/srt/managers/scheduler_components/output_streamer.py +++ b/python/sglang/srt/managers/scheduler_components/output_streamer.py @@ -33,6 +33,7 @@ ) from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache from sglang.srt.runtime_context import get_observability, get_serving +from sglang.srt.sampling.sampling_mask import SamplingMaskChunk from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.utils.weight_versions import compute_weight_version_spans @@ -56,6 +57,7 @@ class SchedulerOutputStreamer: ps: ParallelState server_args: ServerArgs is_generation: bool + is_multimodal_gen: bool spec_algorithm: SpeculativeAlgorithm disaggregation_mode: DisaggregationMode enable_hicache_storage: Callable[[], bool] @@ -204,6 +206,8 @@ def _stream_output_generation( acc.customized_info[key] = values # Send to detokenizer + if self.is_multimodal_gen: + return payload = acc.to_payload( dp_rank=self.ps.dp_rank, is_idle_batch=is_idle_batch, @@ -376,8 +380,7 @@ class _GenerationStreamAccumulator: input_token_ids_logprobs_idx: Optional[list] = None output_token_ids_logprobs_val: Optional[list] = None output_token_ids_logprobs_idx: Optional[list] = None - output_token_sampling_mask: Optional[list] = None - output_token_sampling_logprobs: Optional[list] = None + output_token_sampling_mask: Optional[list[Optional[SamplingMaskChunk]]] = None # Rust server mode: the Rust detokenizer reconstructs text/ids from the raw # output tokens itself and never consumes the scheduler's incremental-detok # offsets (decode_ids / read_offset), so that per-step bookkeeping is skipped. @@ -409,7 +412,6 @@ def __post_init__(self) -> None: self.output_token_ids_logprobs_idx = [] if self.return_sampling_mask: self.output_token_sampling_mask = [] - self.output_token_sampling_logprobs = [] def _beam_admits(self, *, req: Req) -> bool: # Only the leader is ever streamed, and only at group finish. @@ -626,22 +628,9 @@ def accept(self, *, req: Req) -> None: if self.return_sampling_mask: if req.return_sampling_mask: - send_output_sampling_mask_offset = req.send_output_sampling_mask_offset - sampling_mask_end = len(req.output_token_sampling_mask) - self.output_token_sampling_mask.append( - req.output_token_sampling_mask[ - send_output_sampling_mask_offset:sampling_mask_end - ] - ) - self.output_token_sampling_logprobs.append( - req.output_token_sampling_logprobs[ - send_output_sampling_mask_offset:sampling_mask_end - ] - ) - req.send_output_sampling_mask_offset = sampling_mask_end + self.output_token_sampling_mask.append(req.sampling_mask_rows.take()) else: - self.output_token_sampling_mask.append([]) - self.output_token_sampling_logprobs.append([]) + self.output_token_sampling_mask.append(None) if self.return_hidden_states: if req.return_hidden_states: @@ -693,6 +682,15 @@ def to_payload( if not (self.rids or is_idle_batch): return None dp_ranks = [dp_rank] * len(self.rids) if self.rids else None + indexer_topk_num_layers = None + if self.return_indexer_topk: + from sglang.srt.state_capturer.indexer_topk import ( + get_global_indexer_capturer, + ) + + _idx_cap = get_global_indexer_capturer() + if _idx_cap is not None: + indexer_topk_num_layers = _idx_cap.num_indexer_layers return BatchTokenIDOutput( rids=self.rids, http_worker_ipcs=self.http_worker_ipcs, @@ -750,10 +748,10 @@ def to_payload( output_token_ids_logprobs_idx=self.output_token_ids_logprobs_idx, output_token_entropy_val=None, output_token_sampling_mask=self.output_token_sampling_mask, - output_token_sampling_logprobs=self.output_token_sampling_logprobs, output_hidden_states=self.output_hidden_states, routed_experts=self.routed_experts, indexer_topk=self.indexer_topk, + indexer_topk_num_layers=indexer_topk_num_layers, customized_info=( wrap_as_pickle(self.customized_info) if self.customized_info else None ), diff --git a/python/sglang/srt/managers/scheduler_components/weight_updater.py b/python/sglang/srt/managers/scheduler_components/weight_updater.py index 9c838e59ff8a..7349033a8d8b 100644 --- a/python/sglang/srt/managers/scheduler_components/weight_updater.py +++ b/python/sglang/srt/managers/scheduler_components/weight_updater.py @@ -6,7 +6,7 @@ import traceback from contextlib import contextmanager from dataclasses import dataclass, field -from typing import Any, Callable, Dict, Iterator, Optional, Tuple +from typing import Any, Callable, Dict, Iterator, List, Optional, Set, Tuple import msgspec import torch @@ -19,15 +19,21 @@ ) from sglang.srt.disaggregation.utils import DisaggregationMode from sglang.srt.managers.io_struct import ( + BeginWeightUpdateReqInput, + BeginWeightUpdateReqOutput, ChecksumInfo, CheckWeightsReqInput, CheckWeightsReqOutput, DestroyWeightsUpdateGroupReqInput, DestroyWeightsUpdateGroupReqOutput, + EndWeightUpdateReqInput, + EndWeightUpdateReqOutput, GetWeightsByNameReqInput, GetWeightsByNameReqOutput, InitWeightsUpdateGroupReqInput, InitWeightsUpdateGroupReqOutput, + PullWeightsReqInput, + PullWeightsReqOutput, ReleaseMemoryOccupationReqInput, ReleaseMemoryOccupationReqOutput, ResumeMemoryOccupationReqInput, @@ -41,36 +47,69 @@ UpdateWeightsFromTensorReqInput, UpdateWeightsFromTensorReqOutput, ) +from sglang.srt.model_executor.model_runner_components.weight_updater import ( + LocalSerializedTensor, +) from sglang.srt.runtime_context import get_model +from sglang.srt.utils import MultiprocessingSerializer +from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions +from sglang.srt.utils.weight_checker import overall_checksum +from sglang.srt.weight_sync.tensor_bucket import FlattenedTensorBucket logger = logging.getLogger(__name__) -def _get_draft_model_runner(draft_worker): - # DFlash / FrozenKVMTP workers expose draft_model_runner directly - runner = getattr(draft_worker, "draft_model_runner", None) - if runner is not None: - return runner - # EAGLEWorkerV2: _draft_worker.draft_runner - inner = getattr(draft_worker, "_draft_worker", None) - if inner is not None: - runner = getattr(inner, "draft_runner", None) - if runner is not None: - return runner - return None - - -def _merge_checksum_payloads(target: Dict, draft: Dict) -> Dict: - merged_checksums = dict(target["checksums"]) - for name, chk in draft["checksums"].items(): - merged_checksums[f"draft.{name}"] = chk - h = hashlib.sha256() - for name in sorted(merged_checksums): - h.update(name.encode()) - h.update(merged_checksums[name].encode()) - target["checksums"] = merged_checksums - target["per_gpu_checksum"] = h.hexdigest() - return target +def _split_lora_named_tensors(named_tensors): + """Partition a weight-update payload on the ``{lora_name}:`` prefix. + + The prefix is unconditional on streamed adapter tensors (it is both the + routing discriminator and the stash grouping key), so a lora_A/lora_B + tensor without one is a sender bug, not base data.""" + base_tensors, lora_tensors = [], [] + for name, tensor in named_tensors: + if ":" in name: + lora_tensors.append((name, tensor)) + else: + assert ".lora_A." not in name and ".lora_B." not in name, ( + f"LoRA tensor {name!r} arrived without a '{{lora_name}}:' prefix" + ) + base_tensors.append((name, tensor)) + return base_tensors, lora_tensors + + +def _sha256_tensor(tensor: torch.Tensor) -> str: + return hashlib.sha256( + tensor.detach().cpu().contiguous().flatten().view(torch.uint8).numpy().tobytes() + ).hexdigest() + + +def _merge_checksum_payloads(role_payloads: List[Tuple[str, Dict]]) -> Dict: + merged: Dict[str, str] = {} + parallelism_infos = [] + for role, p in role_payloads: + for name, chk in p["checksums"].items(): + # Only non-target roles are prefixed, so target keys stay stable. + key = name if role == "" else f"{role}.{name}" + if key in merged: + raise ValueError(f"checksum key collision: {key}") + merged[key] = chk + parallelism_infos.append({"role": role or "target", **p["parallelism_info"]}) + return { + "checksums": merged, + "per_gpu_checksum": overall_checksum(merged), + "parallelism_info": parallelism_infos, + } + + +def _parse_runner_selector(selector: str) -> Set[str]: + """Map a {target, draft, all} weight-op selector to the set of roles it covers.""" + if selector == "all": + return {"target", "draft"} + if selector in ("target", "draft"): + return {selector} + raise ValueError( + f"invalid selector {selector!r}; expected 'target', 'draft', or 'all'" + ) @dataclass(kw_only=True, slots=True) @@ -85,6 +124,24 @@ class SchedulerWeightUpdaterManager: metrics_collector: Optional[Any] = None offload_tags: set = field(default_factory=set) stashed_model_static_state: Any = None + _weight_update_in_progress: bool = False + _weight_update_loaded: bool = False + # Runner selector for the open session, recorded at begin_weight_update and + # reused by end_weight_update so the same set is restored and finalized. + _weight_update_selector: str = "all" + # Session scope, recorded at begin_weight_update: only sync_base sessions + # unpack/re-finalize base weights, and only they accept base tensors. + _weight_update_sync_base: bool = True + # Streamed adapter tensors accumulated this session: {lora_name: {hf_key: tensor}}. + # Applied and cleared by end_weight_update; discarded by the next begin. + _lora_stash: Dict[str, Dict[str, torch.Tensor]] = field(default_factory=dict) + # Name set of each adapter's last applied stream. The set is static for an + # adapter's lifetime; a change means a partial stream, which the manager's + # whole-adapter replace would silently zero. + _lora_applied_names: Dict[str, frozenset] = field(default_factory=dict) + # Version reported by the buckets of the open session; recorded only when + # end_weight_update commits, so a version never names a half-applied update. + _weight_update_pending_version: Optional[str] = None @contextmanager def _observe_weight_load(self, source: str) -> Iterator[None]: @@ -109,6 +166,10 @@ def flush_cache_after_weight_update(self, recv_req) -> None: assert flush_cache_success, "Cache flush failed after updating weights" def record_weight_version_after_update(self, weight_version: Optional[str]) -> None: + if self._weight_update_in_progress: + if weight_version is not None: + self._weight_update_pending_version = weight_version + return self.scheduler.record_weight_version_change(new_version=weight_version) def update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput): @@ -128,6 +189,44 @@ def update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput): success=success, message=message, num_paused_requests=0 ) + def pull_weights(self, recv_req: PullWeightsReqInput): + """Sync this host's local checkpoint up to recv_req.target_version. + + Every rank runs the pull; a per-host file lock collapses co-located + ranks to one pull. Success is gathered across the TP group (all nodes), + so the reply only reports success once every host holds a verified + checkpoint. + """ + from sglang.srt.weight_sync import local_checkpoint + + server_args = self.tp_worker.model_runner.server_args + try: + local_checkpoint.pull( + local_checkpoint_dir=recv_req.local_checkpoint_dir, + base_dir=server_args.model_path, + source_dir=recv_req.source_dir, + target_version=recv_req.target_version, + pre_read_hook=server_args.custom_pull_weights_pre_read_hook, + ) + success, message = True, "Success." + except Exception: + success, message = False, traceback.format_exc() + logger.error(message) + + tp_size = ( + torch.distributed.get_world_size(group=self.tp_cpu_group) + if torch.distributed.is_initialized() + else 1 + ) + if tp_size > 1: + results = [None] * tp_size + torch.distributed.all_gather_object( + results, (success, message), group=self.tp_cpu_group + ) + success = all(ok for ok, _ in results) + message = "; ".join(msg for ok, msg in results if not ok) or message + return PullWeightsReqOutput(success=success, message=message) + def init_weights_update_group(self, recv_req: InitWeightsUpdateGroupReqInput): """Initialize the online model parameter update group.""" success, message = self.tp_worker.init_weights_update_group(recv_req) @@ -141,30 +240,115 @@ def destroy_weights_update_group( success, message = self.tp_worker.destroy_weights_update_group(recv_req) return DestroyWeightsUpdateGroupReqOutput(success=success, message=message) + def iter_weight_update_workers( + self, selector: str = "all" + ) -> List[Tuple[str, Any]]: + """Resolve a {target, draft, all} selector to (role, worker) pairs, target + first. This is the worker-level inclusion decision; each worker then + contributes its own runners via iter_runners().""" + parsed = _parse_runner_selector(selector) + workers: List[Tuple[str, Any]] = [] + if "target" in parsed: + workers.append(("target", self.tp_worker)) + if "draft" in parsed and self.draft_worker is not None: + workers.append(("draft", self.draft_worker)) + return workers + + def get_model_runners(self, selector: str = "all") -> List[Tuple[str, Any]]: + """Resolve a {target, draft, all} selector to (role, ModelRunner) pairs, + target first. Role is "" for the target runner; draft roles come from the + draft worker's iter_runners().""" + runners: List[Tuple[str, Any]] = [] + for _, worker in self.iter_weight_update_workers(selector): + runners += worker.iter_runners() + return runners + def update_weights_from_distributed( self, recv_req: UpdateWeightsFromDistributedReqInput, ) -> Tuple[bool, str]: - """Update the online model parameter.""" + """Update the online model parameter, fanning out to the selected runners.""" + assert self._weight_update_in_progress, ( + "update_weights_from_distributed requires an open begin_weight_update session" + ) with self._observe_weight_load("distributed"): - success, message = self.tp_worker.update_weights_from_distributed(recv_req) + # Only the target (main) model joined this process's update group, so it + # receives the broadcast once; the received weights are then loaded into + # each selected runner locally. Draft runners never join the group. + try: + weights = self.tp_worker.model_runner.weight_updater.receive_weights_from_distributed( + recv_req.names, + recv_req.dtypes, + recv_req.shapes, + recv_req.group_name, + recv_req.load_format, + ) + base_tensors, lora_tensors = _split_lora_named_tensors(weights) + self._stash_lora_tensors(lora_tensors) + if base_tensors: + assert self._weight_update_sync_base, ( + "base tensors arrived in a sync_base=False weight-update session" + ) + for _, runner in self.get_model_runners(recv_req.selector): + runner.weight_updater.load_weights(base_tensors) + success, message = True, "Succeeded to update parameter online." + except Exception as e: + success = False + message = ( + f"Failed to update parameter online: {e}. The full weights of the " + "ModelRunner are partially updated. Please discard the whole weights." + ) + logger.error(message) if success: + if base_tensors: + self._weight_update_loaded = True self.flush_cache_after_weight_update(recv_req) self.record_weight_version_after_update(recv_req.weight_version) - else: - logger.error(message) return UpdateWeightsFromDistributedReqOutput( success=success, message=message ) def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput): - """Update the online model parameter from tensors.""" + """Update the online model parameter from tensors, fanning out to the + selected runners.""" + assert self._weight_update_in_progress, ( + "update_weights_from_tensor requires an open begin_weight_update session" + ) with self._observe_weight_load("tensor"): - if recv_req.disable_draft_model: - worker = self.tp_worker + monkey_patch_torch_reductions() + named_tensors = MultiprocessingSerializer.deserialize( + recv_req.serialized_named_tensors[self.tp_worker.ps.tp_rank] + ) + load_format = recv_req.load_format + if load_format == "flattened_bucket": + # Names live inside the bucket metadata, so the LoRA split must + # run on the reconstructed tensors. A pure-base bucket keeps the + # original payload and format (no behavior change). + reconstructed = FlattenedTensorBucket( + flattened_tensor=named_tensors["flattened_tensor"], + metadata=named_tensors["metadata"], + ).reconstruct_tensors() + base_tensors, lora_tensors = _split_lora_named_tensors(reconstructed) + if lora_tensors: + named_tensors, load_format = base_tensors, None else: - worker = self.draft_worker or self.tp_worker - success, message = worker.update_weights_from_tensor(recv_req) + base_tensors, lora_tensors = _split_lora_named_tensors(named_tensors) + named_tensors = base_tensors + self._stash_lora_tensors(lora_tensors, copy_tensors=True) + success, message = True, "Success" + if base_tensors: + assert self._weight_update_sync_base, ( + "base tensors arrived in a sync_base=False weight-update session" + ) + for _, runner in self.get_model_runners(recv_req.selector): + success, message = runner.weight_updater.update_weights_from_tensor( + named_tensors=named_tensors, + load_format=load_format, + ) + if not success: + break + if success: + self._weight_update_loaded = True if success: self.flush_cache_after_weight_update(recv_req) self.record_weight_version_after_update(recv_req.weight_version) @@ -209,21 +393,130 @@ def _assert_weight_cache_inactive(self, op: str) -> None: f"Restart with --weight-cache-mode off to use this operation." ) - def release_memory_occupation(self, recv_req: ReleaseMemoryOccupationReqInput): - assert self.is_fully_idle(), ( - "release_memory_occupation should be called only when server is idle." + def begin_weight_update(self, recv_req: BeginWeightUpdateReqInput): + """Begin a weight-update session: restore in-place-packed weights to a + loadable state on the selected runners (target and/or draft), so the draft + model is prepared identically to the target. The selector is recorded and + reused by end_weight_update so the same set is finalized.""" + assert not self._weight_update_in_progress, ( + "begin_weight_update called while a weight-update session is already open" + ) + self._weight_update_selector = recv_req.selector + self._weight_update_sync_base = recv_req.sync_base + self._lora_stash = {} + self._weight_update_pending_version = None + if recv_req.sync_base: + for _, runner in self.get_model_runners(recv_req.selector): + runner.begin_weight_update() + self._weight_update_in_progress = True + self._weight_update_loaded = False + torch.distributed.barrier(group=self.tp_cpu_group) + return BeginWeightUpdateReqOutput(success=True, message="Success") + + def end_weight_update(self, recv_req: EndWeightUpdateReqInput): + """End the weight-update session on the runners begin_weight_update opened + (its recorded selector): quant finalize on each, plus model.post_load_weights + only when load_weights was bypassed this session (e.g. P2P/RDMA).""" + assert self._weight_update_in_progress, ( + "end_weight_update called without begin_weight_update" ) + if self._weight_update_sync_base: + run_post_load = not self._weight_update_loaded + for _, runner in self.get_model_runners(self._weight_update_selector): + runner.end_weight_update(run_post_load=run_post_load) + if recv_req.abort: + self._lora_stash = {} + success, message = True, "Aborted: streamed adapters discarded" + else: + success, message = self._apply_lora_stash(recv_req.expected_lora_checksums) + self._weight_update_in_progress = False + if success and not recv_req.abort: + self.record_weight_version_after_update(self._weight_update_pending_version) + self._weight_update_pending_version = None + torch.distributed.barrier(group=self.tp_cpu_group) + return EndWeightUpdateReqOutput(success=success, message=message) + + def forget_lora_adapter(self, lora_name: str) -> None: + """Drop the partial-stream guard entry: a re-registered or unloaded name + is a new adapter identity and may stream a different tensor set.""" + self._lora_applied_names.pop(lora_name, None) + + def _stash_lora_tensors(self, lora_tensors, *, copy_tensors: bool = False) -> None: + copied_devices = set() + for prefixed_name, tensor in lora_tensors: + lora_name, hf_key = prefixed_name.split(":", 1) + if isinstance(tensor, LocalSerializedTensor): + tensor = tensor.get(self.tp_worker.ps.tp_rank) + assert isinstance(tensor, torch.Tensor), ( + f"streamed LoRA tensor {prefixed_name!r} must arrive as a plain tensor" + ) + if copy_tensors: + # the stash outlives this RPC, but an IPC sender may reuse the bucket once it replies + tensor = tensor.clone() + if tensor.is_cuda: + copied_devices.add(tensor.device) + self._lora_stash.setdefault(lora_name, {})[hf_key] = tensor + for device in copied_devices: + torch.cuda.current_stream(device).synchronize() + + def _apply_lora_stash( + self, expected_checksums: Optional[Dict[str, Dict[str, str]]] + ) -> Tuple[bool, str]: + """Hand each streamed adapter to the LoRA manager's whole-adapter entry + point (config from registration, upsert in place), then clear the stash.""" + if expected_checksums is not None and set(expected_checksums) != set( + self._lora_stash + ): + return False, ( + f"[LORA-CHECK] streamed adapters {sorted(self._lora_stash)} do not " + f"match the expected manifest {sorted(expected_checksums)}" + ) + if not self._lora_stash: + return True, "Success" + # The attribute only exists when LoRA is enabled at server start. + lora_manager = getattr(self.tp_worker.model_runner, "lora_manager", None) + if lora_manager is None: + return False, "streamed LoRA tensors require --enable-lora" + for lora_name in sorted(self._lora_stash): + tensors = self._lora_stash[lora_name] + names = frozenset(tensors) + if expected_checksums is not None: + expected = expected_checksums[lora_name] + if set(expected) != set(tensors): + return False, ( + f"[LORA-CHECK] adapter {lora_name!r}: streamed tensor names " + f"do not match the expected manifest" + ) + for name in sorted(tensors): + if _sha256_tensor(tensors[name]) != expected[name]: + return False, ( + f"[LORA-CHECK] adapter {lora_name!r}: checksum mismatch for {name!r}" + ) + prev = self._lora_applied_names.get(lora_name) + if prev is not None and prev != names: + return False, ( + f"streamed adapter {lora_name!r} arrived with a different tensor " + f"set than its previous sync (partial stream?)" + ) + result = lora_manager.apply_streamed_adapter(lora_name, tensors) + if not result.success: + return False, result.error_message + self._lora_applied_names[lora_name] = names + self._lora_stash = {} + return True, "Success" + + def release_memory_occupation(self, recv_req: ReleaseMemoryOccupationReqInput): + scheduler = self.scheduler + assert self.is_fully_idle( + ignore_waiting=scheduler is not None and scheduler._engine_paused + ), "release_memory_occupation should be called only when server is idle." tags = recv_req.tags if tags is None or len(tags) == 0: tags = GPU_MEMORY_ALL_TYPES - for tag in tags: - self.offload_tags.add(tag) - if GPU_MEMORY_TYPE_KV_CACHE in tags: - scheduler = self.scheduler if scheduler is not None: if scheduler.disaggregation_mode == DisaggregationMode.DECODE: for queue_name in ( @@ -237,9 +530,17 @@ def release_memory_occupation(self, recv_req: ReleaseMemoryOccupationReqInput): queue = getattr(scheduler, "disagg_prefill_bootstrap_queue", None) if queue is not None: queue.release_memory_occupation() - self.memory_saver_adapter.pause(GPU_MEMORY_TYPE_KV_CACHE) + # flush_cache is a no-op once KV is tagged as offloaded, so flush before tagging self.flush_cache() + for tag in tags: + self.offload_tags.add(tag) + + if GPU_MEMORY_TYPE_KV_CACHE in tags: + # Cache-clearing kernels must finish before KV memory is unmapped. + torch.get_device_module().synchronize() + self.memory_saver_adapter.pause(GPU_MEMORY_TYPE_KV_CACHE) + if GPU_MEMORY_TYPE_WEIGHTS in tags: self._assert_weight_cache_inactive("release_memory_occupation") self.stashed_model_static_state = _export_static_state( @@ -298,19 +599,16 @@ def resume_memory_occupation(self, recv_req: ResumeMemoryOccupationReqInput): def check_weights(self, recv_req: CheckWeightsReqInput): try: - payload = self.tp_worker.model_runner.check_weights( - action=recv_req.action, allow_quant_error=recv_req.allow_quant_error - ) - - if self.draft_worker is not None: - draft_runner = _get_draft_model_runner(self.draft_worker) - if draft_runner is not None: - draft_payload = draft_runner.check_weights( - action=recv_req.action, - allow_quant_error=recv_req.allow_quant_error, - ) - if payload is not None and draft_payload is not None: - payload = _merge_checksum_payloads(payload, draft_payload) + role_payloads = [] + for role, runner in self.get_model_runners(recv_req.selector): + p = runner.check_weights( + action=recv_req.action, + allow_quant_error=recv_req.allow_quant_error, + skip_tensor_list=recv_req.skip_tensor_list, + ) + if p is not None: + role_payloads.append((role, p)) + payload = _merge_checksum_payloads(role_payloads) if role_payloads else None tp_size = torch.distributed.get_world_size(group=self.tp_cpu_group) if tp_size > 1 and payload is not None: diff --git a/python/sglang/srt/managers/tokenizer_control_mixin.py b/python/sglang/srt/managers/tokenizer_control_mixin.py index 69ef86803b8d..67dd4424797b 100644 --- a/python/sglang/srt/managers/tokenizer_control_mixin.py +++ b/python/sglang/srt/managers/tokenizer_control_mixin.py @@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple import fastapi +import msgspec from sglang.srt.managers.communicator import FanOutCommunicator from sglang.srt.managers.io_struct import ( @@ -15,6 +16,8 @@ AddExternalCorpusReqOutput, AttachHiCacheStorageReqInput, AttachHiCacheStorageReqOutput, + BeginWeightUpdateReqInput, + BeginWeightUpdateReqOutput, ChecksumInfo, CheckWeightsReqInput, CheckWeightsReqOutput, @@ -27,6 +30,8 @@ DetachHiCacheStorageReqOutput, DumperControlReqInput, DumperControlReqOutput, + EndWeightUpdateReqInput, + EndWeightUpdateReqOutput, ExpertDistributionReq, ExpertDistributionReqOutput, ExpertDistributionReqType, @@ -42,8 +47,6 @@ InitWeightsUpdateGroupReqOutput, ListExternalCorporaReqInput, ListExternalCorporaReqOutput, - LoadLoRAAdapterFromTensorsReqInput, - LoadLoRAAdapterFromTensorsReqOutput, LoadLoRAAdapterReqInput, LoadLoRAAdapterReqOutput, LoRAUpdateOutput, @@ -51,6 +54,10 @@ ProfileReq, ProfileReqOutput, ProfileReqType, + PullWeightsReqInput, + PullWeightsReqOutput, + RegisterLoRAAdapterReqInput, + RegisterLoRAAdapterReqOutput, ReleaseMemoryOccupationReqInput, ReleaseMemoryOccupationReqOutput, RemoveExternalCorpusReqInput, @@ -114,6 +121,7 @@ ("release_memory_occupation", ReleaseMemoryOccupationReqOutput), ("resume_memory_occupation", ResumeMemoryOccupationReqOutput), ("check_weights", CheckWeightsReqOutput), + ("pull_weights", PullWeightsReqOutput), ("slow_down", SlowDownReqOutput), ("flush_cache", FlushCacheReqOutput), ("add_external_corpus", AddExternalCorpusReqOutput), @@ -126,6 +134,8 @@ ("get_internal_state", GetInternalStateReqOutput), ("set_internal_state", SetInternalStateReqOutput), ("expert_distribution", ExpertDistributionReqOutput), + ("begin_weight_update", BeginWeightUpdateReqOutput), + ("end_weight_update", EndWeightUpdateReqOutput), ("update_lora_adapter", LoRAUpdateOutput), ("dumper_control", DumperControlReqOutput), ("scale_elastic_ep", ScaleElasticEPReqOutput), @@ -444,6 +454,119 @@ async def destroy_weights_update_group( results = await self.destroy_weights_update_group_communicator(obj) return FanOutCommunicator.merge_results(results) + async def _weight_update_session_call( + self: TokenizerManager, communicator, obj, staged: bool = False + ) -> Tuple[bool, str]: + """Run one weight-update RPC with pause-aware locking (a paused engine + already holds the writer lock). A staged session takes the reader lock + and runs alongside generation: deferred names are not servable and base + tensors are rejected; other sessions keep the writer lock.""" + self.auto_create_handle_loop() + async with self.is_pause_cond: + is_paused = self.is_pause + if is_paused: + results = await communicator(obj) + if not is_paused: + lock = ( + self.model_update_lock.reader_lock + if staged + else self.model_update_lock.writer_lock + ) + async with lock: + results = await communicator(obj) + return FanOutCommunicator.merge_results(results) + + def _staged_serving_state_error(self: TokenizerManager, obj) -> Optional[str]: + # the reader lock gave up exclusivity; serving-state mutations need the writer lock + if self._weight_update_staged_session and ( + obj.flush_cache or obj.abort_all_requests or obj.weight_version is not None + ): + return "a staged adapter session cannot change serving state" + return None + + async def begin_weight_update( + self: TokenizerManager, + obj: BeginWeightUpdateReqInput, + request: Optional[fastapi.Request] = None, + ) -> Tuple[bool, str]: + staged = not obj.sync_base and bool(self._pending_lora_publications) + success, message = await self._weight_update_session_call( + self.begin_weight_update_communicator, obj, staged=staged + ) + if success: + self._weight_update_session_open = True + self._weight_update_staged_session = staged + self._weight_update_pending_version = None + return success, message + + async def end_weight_update( + self: TokenizerManager, + obj: EndWeightUpdateReqInput, + request: Optional[fastapi.Request] = None, + ) -> Tuple[bool, str]: + pending = dict(self._pending_lora_publications) + # fail closed: without a full manifest a lost bucket would publish an incomplete adapter + missing = ( + set() + if obj.abort + else set(pending) - set(obj.expected_lora_checksums or {}) + ) + obj.abort = obj.abort or bool(missing) + success, message = await self._weight_update_session_call( + self.end_weight_update_communicator, + obj, + staged=self._weight_update_staged_session, + ) + self._weight_update_session_open = False + self._weight_update_staged_session = False + if success and not obj.abort: + self._update_weight_version_if_provided(self._weight_update_pending_version) + self._weight_update_pending_version = None + if obj.abort or not success: + await self._discard_pending_publications() + elif pending: + await self._publish_pending_adapters() + if missing: + return ( + False, + f"deferred adapters {sorted(missing)} have no checksum manifest; session aborted", + ) + return success, message + + async def _discard_pending_publications(self: TokenizerManager) -> None: + """Drop deferred publications after an aborted or failed session: the + backends must forget the zeroed identities the names never served.""" + pending, self._pending_lora_publications = self._pending_lora_publications, {} + async with self.lora_update_lock: + for ref in pending.values(): + result = _merge_lora_update_results( + await self.update_lora_adapter_communicator( + UnloadLoRAAdapterReqInput( + lora_name=ref.lora_name, lora_id=ref.lora_id + ) + ) + ) + if not result.success: + logger.warning( + "Failed to discard unpublished LoRA adapter %r: %s", + ref.lora_name, + result.error_message, + ) + + async def _publish_pending_adapters(self: TokenizerManager) -> None: + """Commit deferred publications: the names become servable atomically.""" + pending, self._pending_lora_publications = self._pending_lora_publications, {} + async with self.lora_update_lock: + for name, ref in pending.items(): + await self.lora_registry.register(ref) + self.lora_ref_cache[name] = ref + try: + await self._evict_lru_over_cap() + except ValueError: + # The publish is already committed; an unevictable over-cap + # pool is a capacity problem, not a publication failure. + logger.exception("LRU eviction after LoRA publication failed") + async def update_weights_from_distributed( self: TokenizerManager, obj: UpdateWeightsFromDistributedReqInput, @@ -454,20 +577,16 @@ async def update_weights_from_distributed( "dp_size must be 1 or dp attention must be enabled for update weights from distributed" ) + if error := self._staged_serving_state_error(obj): + return False, error if obj.abort_all_requests: self.abort_request(abort_all=True) - # Hold is_pause_cond while updating to prevent unpause from racing. - async with self.is_pause_cond: - is_paused = self.is_pause - if is_paused: - results = await self.update_weights_from_distributed_communicator(obj) - - if not is_paused: - async with self.model_update_lock.writer_lock: - results = await self.update_weights_from_distributed_communicator(obj) - - success, message = FanOutCommunicator.merge_results(results) + success, message = await self._weight_update_session_call( + self.update_weights_from_distributed_communicator, + obj, + staged=self._weight_update_staged_session, + ) if success and obj.flush_cache and self.mm_processor is not None: self.mm_processor.clear_preprocess_cache() if success and obj.weight_version is not None: @@ -514,6 +633,8 @@ async def update_weights_from_tensor( "dp_size must be 1 or dp attention must be enabled for update weights from tensor" ) + if error := self._staged_serving_state_error(obj): + return False, error if obj.abort_all_requests: self.abort_request(abort_all=True) @@ -521,16 +642,11 @@ async def update_weights_from_tensor( obj.serialized_named_tensors ) - async with self.is_pause_cond: - is_paused = self.is_pause - if is_paused: - results = await self.update_weights_from_tensor_communicator(obj) - - if not is_paused: - async with self.model_update_lock.writer_lock: - results = await self.update_weights_from_tensor_communicator(obj) - - success, message = FanOutCommunicator.merge_results(results) + success, message = await self._weight_update_session_call( + self.update_weights_from_tensor_communicator, + obj, + staged=self._weight_update_staged_session, + ) if success and obj.flush_cache and self.mm_processor is not None: self.mm_processor.clear_preprocess_cache() if success and obj.weight_version is not None: @@ -638,35 +754,8 @@ async def load_lora_adapter( await self.lora_registry.register(new_adapter) self.lora_ref_cache[obj.lora_name] = new_adapter - if get_lora().max_loaded_loras is not None: - while ( - self.lora_registry.num_registered_loras - > get_lora().max_loaded_loras - ): - lru_lora_name = await self.lora_registry.lru_lora_name( - exclude_pinned=True - ) - if lru_lora_name is None: - raise ValueError( - "Didn't find any LoRA adapters when trying to evict LRU LoRA adapter. " - f"LoRA registry is: {self.lora_registry._registry}" - ) - - logger.info( - f"Unloading least recently used LoRA adapter '{lru_lora_name}' " - f"(current number of adapters: {self.lora_registry.num_registered_loras}, " - f"max allowed: {get_lora().max_loaded_loras})" - ) - - unload_result = await self._unload_lora_adapter_locked( - UnloadLoRAAdapterReqInput(lora_name=lru_lora_name) - ) - if not unload_result.success: - raise ValueError( - f"Error while unloading LRU LoRA adapter '{lru_lora_name}': " - f"{unload_result.error_message}" - ) - del result.loaded_adapters[lru_lora_name] + for evicted in await self._evict_lru_over_cap(): + del result.loaded_adapters[evicted] return result except ValueError as e: @@ -675,11 +764,59 @@ async def load_lora_adapter( error_message=str(e), ) - async def load_lora_adapter_from_tensors( + async def _evict_lru_over_cap(self: TokenizerManager) -> list: + """Unload LRU adapters until the registry is back under + --max-loaded-loras; returns the evicted names. Caller must hold + lora_update_lock.""" + evicted = [] + if get_lora().max_loaded_loras is None: + return evicted + while self.lora_registry.num_registered_loras > get_lora().max_loaded_loras: + lru_lora_name = await self.lora_registry.lru_lora_name(exclude_pinned=True) + if lru_lora_name is None: + raise ValueError( + "Didn't find any LoRA adapters when trying to evict LRU LoRA adapter. " + f"LoRA registry is: {self.lora_registry._registry}" + ) + + logger.info( + f"Unloading least recently used LoRA adapter '{lru_lora_name}' " + f"(current number of adapters: {self.lora_registry.num_registered_loras}, " + f"max allowed: {get_lora().max_loaded_loras})" + ) + + unload_result = await self._unload_lora_adapter_locked( + UnloadLoRAAdapterReqInput(lora_name=lru_lora_name) + ) + if not unload_result.success: + raise ValueError( + f"Error while unloading LRU LoRA adapter '{lru_lora_name}': " + f"{unload_result.error_message}" + ) + evicted.append(lru_lora_name) + return evicted + + def _validate_lora_upsert_supported(self: TokenizerManager) -> None: + """Upsert resolves lora_name -> lora_id through this process's registry. + + With multiple tokenizer workers each HTTP worker process holds its own + registry, so the resolution depends on which worker the router picks: + a worker that never served the original load would mint a fresh id and + die on the backend duplicate check. Fail loudly instead. + """ + if get_serving().tokenizer_worker_num > 1: + raise ValueError( + "LoRA upsert is not supported with tokenizer_worker_num > 1: " + "each HTTP worker resolves lora_name against its own registry, " + "making upsert nondeterministic across workers." + ) + + async def register_lora_adapter( self: TokenizerManager, - obj: LoadLoRAAdapterFromTensorsReqInput, + obj: RegisterLoRAAdapterReqInput, _: Optional[fastapi.Request] = None, - ) -> LoadLoRAAdapterFromTensorsReqOutput: + ) -> RegisterLoRAAdapterReqOutput: + """Create-or-refresh an adapter's identity and config (control plane).""" self.auto_create_handle_loop() try: @@ -687,66 +824,76 @@ async def load_lora_adapter_from_tensors( raise ValueError( "LoRA is not enabled. Please set `--enable-lora` to enable LoRA." ) - assert get_parallel().dp_size == 1 or get_parallel().enable_dp_attention, ( "dp_size must be 1 or dp attention must be enabled for dynamic lora loading" ) - logger.info( - "Start load Lora adapter from tensors. Lora name=%s", - obj.lora_name, - ) - - obj.serialized_named_tensors = normalize_serialized_named_tensor_payloads( - obj.serialized_named_tensors - ) - + if ":" in obj.lora_name: + raise ValueError( + f"LoRA adapter name '{obj.lora_name}' must not contain ':': the " + "character separates the adapter name from the tensor key in " + "streamed weight updates." + ) async with self.lora_update_lock: - new_adapter = LoRARef( + self._validate_lora_upsert_supported() + ref = LoRARef( lora_name=obj.lora_name, - lora_path="__tensor__", + lora_path=obj.lora_path or "__stream__", pinned=obj.pinned, + reloadable=obj.lora_path is not None, ) + if obj.lora_name in self._pending_lora_publications: + raise ValueError( + f"LoRA adapter '{obj.lora_name}' is awaiting publication; " + "commit or abort its session before registering it again." + ) + if obj.defer_publish: + if await self.lora_registry.get_lora_id(obj.lora_name) is not None: + raise ValueError( + f"defer_publish requires a fresh adapter name, but " + f"'{obj.lora_name}' is already registered: a published " + "name has readers and must not be staged over." + ) + new_adapter, reused = ref, False + else: + new_adapter, reused = await self.lora_registry.register_or_reuse( + ref, upsert=True + ) + # No path to reload a streamed adapter from: eviction would lose the + # only engine-side copy, so the cap rejects new non-reloadable names + # instead. Reloadable names resolve over-cap by LRU eviction. + if ( + not reused + and not new_adapter.reloadable + and get_lora().max_loaded_loras is not None + and self.lora_registry.num_registered_loras + >= get_lora().max_loaded_loras + ): + raise ValueError( + f"Cannot register streamed LoRA adapter '{obj.lora_name}': " + f"{self.lora_registry.num_registered_loras} adapters already " + f"registered (--max-loaded-loras=" + f"{get_lora().max_loaded_loras}). Streamed adapters are " + "never evicted; unload one or raise the cap." + ) obj.lora_id = new_adapter.lora_id result = _merge_lora_update_results( await self.update_lora_adapter_communicator(obj) ) if result.success: - await self.lora_registry.register(new_adapter) - self.lora_ref_cache[obj.lora_name] = new_adapter - if get_lora().max_loaded_loras is not None: - while ( - self.lora_registry.num_registered_loras - > get_lora().max_loaded_loras - ): - lru_lora_name = await self.lora_registry.lru_lora_name( - exclude_pinned=True - ) - if lru_lora_name is None: - raise ValueError( - "Didn't find any LoRA adapters when trying to evict LRU LoRA adapter. " - f"LoRA registry is: {self.lora_registry._registry}" - ) - - logger.info( - f"Unloading least recently used LoRA adapter '{lru_lora_name}' " - f"(current number of adapters: {self.lora_registry.num_registered_loras}, " - f"max allowed: {get_lora().max_loaded_loras})" - ) - - unload_result = await self._unload_lora_adapter_locked( - UnloadLoRAAdapterReqInput(lora_name=lru_lora_name) - ) - if not unload_result.success: - raise ValueError( - f"Error while unloading LRU LoRA adapter '{lru_lora_name}': " - f"{unload_result.error_message}" - ) - del result.loaded_adapters[lru_lora_name] - + if obj.defer_publish: + self._pending_lora_publications[obj.lora_name] = new_adapter + result = msgspec.structs.replace(result, pending=True) + else: + if reused: + await self.lora_registry.refresh(new_adapter) + else: + await self.lora_registry.register(new_adapter) + self.lora_ref_cache[obj.lora_name] = new_adapter + await self._evict_lru_over_cap() return result except ValueError as e: - return LoadLoRAAdapterFromTensorsReqOutput( + return RegisterLoRAAdapterReqOutput( success=False, error_message=str(e), ) @@ -775,9 +922,21 @@ async def unload_lora_adapter( "Start unload Lora adapter. Lora name=%s", obj.lora_name, ) + if obj.lora_name in self._pending_lora_publications: + raise ValueError( + f"LoRA adapter '{obj.lora_name}' is awaiting publication; " + "abort its weight-update session instead of unloading it." + ) async with self.lora_update_lock: - return await self._unload_lora_adapter_locked(obj) + result = await self._unload_lora_adapter_locked(obj) + # Explicit unload is a DELETE: drop the reload-catalog entry too. + # The max_loaded_loras LRU loop calls _unload_lora_adapter_locked + # directly — an EVICT — and must keep the entry so disk-backed + # adapters can be implicitly reloaded later. + if result.success: + self.lora_ref_cache.pop(obj.lora_name, None) + return result except ValueError as e: return UnloadLoRAAdapterReqOutput(success=False, error_message=str(e)) @@ -810,6 +969,15 @@ async def resume_memory_occupation( self.auto_create_handle_loop() await self.resume_memory_occupation_communicator(obj) + async def pull_weights( + self: TokenizerManager, + obj: PullWeightsReqInput, + request: Optional[fastapi.Request] = None, + ) -> Tuple[bool, str]: + self.auto_create_handle_loop() + results = await self.pull_weights_communicator(obj) + return FanOutCommunicator.merge_results(results) + async def check_weights( self: TokenizerManager, obj: CheckWeightsReqInput, @@ -938,8 +1106,13 @@ async def update_weight_version( def _update_weight_version_if_provided( self: TokenizerManager, weight_version: Optional[str] ) -> None: - """Update weight version if provided.""" - if weight_version is not None: - self.record_config_updates( - "tokenizer.weight_version", weight_version=weight_version - ) + """Update weight version if provided; inside a weight-update session the + version is held until end_weight_update commits.""" + if weight_version is None: + return + if self._weight_update_session_open: + self._weight_update_pending_version = weight_version + return + self.record_config_updates( + "tokenizer.weight_version", weight_version=weight_version + ) diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 18fa575ad766..4007fed814b4 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -247,6 +247,18 @@ class ReqState: dispatched: bool = False abort_sent: bool = False + # An abort matched this rid while the request was still tokenizer-held + # (parked at the pause gate / model-update lock / tokenization). The + # scheduler never saw the rid, so the dispatch path must resolve the + # request as aborted instead of sending it. + abort_before_dispatch: bool = False + + # The LoRA lease for this request has already been released. Every terminal + # path funnels through TokenizerManager._finalize_lora_lease, which uses + # this flag to stay idempotent: a leaked release hangs unload's + # wait_for_zero forever, and a double release drives the ConcurrentCounter + # negative, which hangs it just the same (it waits for exactly zero). + lora_lease_released: bool = False # For streaming output last_output_offset: int = 0 @@ -289,8 +301,12 @@ def get_crash_dump_output(self) -> Dict[Any, Any]: input_token_ids_logprobs_idx: List = dataclasses.field(default_factory=list) output_token_ids_logprobs_val: List = dataclasses.field(default_factory=list) output_token_ids_logprobs_idx: List = dataclasses.field(default_factory=list) - output_token_sampling_mask: List = dataclasses.field(default_factory=list) - output_token_sampling_logprobs: List = dataclasses.field(default_factory=list) + output_token_sampling_mask: List[List[int]] = dataclasses.field( + default_factory=list + ) + output_token_sampling_logprobs: List[Union[float, List[float]]] = dataclasses.field( + default_factory=list + ) # Cached flat-format prompt top logprob fields; rebuilt only when more # prefill chunks arrive, so streaming decode chunks reuse the payload. @@ -416,10 +432,53 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): # Set by whoever owns the event loop, and left None for Engine and grpc, # which own no server. Class-level to leave the frozen __init__ alone. _server_stop_hook: Optional[Callable[[], None]] = None + _engine_state_changed_callback: Optional[Callable[[], None]] = None def set_server_stop_hook(self, hook: Callable[[], None]) -> None: self._server_stop_hook = hook + def _notify_engine_state_changed(self) -> None: + callback = self._engine_state_changed_callback + if callback is None: + return + try: + callback() + except Exception: + logger.exception("Engine-state change callback failed") + + def _set_engine_state_field(self, name: str, value: Any) -> None: + if value == getattr(self, name, None): + return + setattr(self, name, value) + self._notify_engine_state_changed() + + def set_engine_state_changed_callback(self, callback: Callable[[], None]) -> None: + self._engine_state_changed_callback = callback + + @property + def server_status(self): + return self._server_status + + @server_status.setter + def server_status(self, value) -> None: + self._set_engine_state_field("_server_status", value) + + @property + def gracefully_exit(self) -> bool: + return self._gracefully_exit + + @gracefully_exit.setter + def gracefully_exit(self, value: bool) -> None: + self._set_engine_state_field("_gracefully_exit", value) + + @property + def is_pause(self) -> bool: + return self._is_pause + + @is_pause.setter + def is_pause(self, value: bool) -> None: + self._set_engine_state_field("_is_pause", value) + @property def serving_chat_class(self): """Return the serving chat class for OpenAI API. @@ -626,6 +685,14 @@ def init_running_status(self): # Subprocess liveness watchdog — set by Engine or http_server after construction self._subprocess_watchdog = None + def is_ready(self) -> bool: + """Return whether this server should receive new requests.""" + return ( + not self.is_pause + and not self.gracefully_exit + and self.server_status == ServerStatus.Up + ) + def init_request_logging_and_dumping(self): # TODO: Refactor and organize the log export code. # Request logging @@ -669,6 +736,11 @@ def init_weight_update(self): self.model_update_tmp: List[UpdateWeightFromDiskReqOutput] = [] self.is_pause = False self.is_pause_cond = asyncio.Condition() + self._weight_update_session_open = False + # A staged session (adapter-only, deferred publications pending) runs + # under the reader lock, concurrent with generation. + self._weight_update_staged_session = False + self._weight_update_pending_version: Optional[str] = None def init_lora(self): # LoRA @@ -688,6 +760,9 @@ def init_lora(self): if get_lora().lora_paths is not None: for lora_ref in get_lora().lora_paths: self.lora_ref_cache[lora_ref.lora_name] = lora_ref + # defer_publish registrations: zeroed identities on the backends, names + # unservable until end_weight_update commits their session + self._pending_lora_publications: Dict[str, LoRARef] = {} def init_disaggregation(self, *, start_pd_bootstrap_service: bool = True): # PD Disaggregation @@ -1304,6 +1379,11 @@ def _validate_one_request( "return_sampling_mask only supports DisallowedTokensLogitsProcessor " "among custom logit processors." ) + if obj.sampling_logprobs_mode is not None and not obj.return_sampling_mask: + raise ValueError( + "sampling_logprobs_mode can only be set when " + "return_sampling_mask=true." + ) def _validate_mm_limits( self, obj: Union[GenerateReqInput, EmbeddingReqInput] @@ -1439,6 +1519,7 @@ def _create_tokenized_object( top_logprobs_num=obj.top_logprobs_num, token_ids_logprob=obj.token_ids_logprob, return_sampling_mask=obj.return_sampling_mask, + sampling_logprobs_mode=obj.sampling_logprobs_mode or "selected", return_flat_raw_top_logprobs=obj.return_flat_raw_top_logprobs, stream=obj.stream, rid=obj.rid, @@ -1607,10 +1688,38 @@ def _should_use_batch_tokenization(self, batch_size, requests) -> bool: ) ) + def _abort_instead_of_dispatch( + self, + tokenized_obj: Union[TokenizedGenerateReqInput, TokenizedEmbeddingReqInput], + ) -> bool: + """Resolve a request whose rid an abort matched while it was still + tokenizer-held (see abort_request): the scheduler never saw the rid, so + it must be finished here instead of dispatched. Returns True if the + request was aborted.""" + state = self.rid_to_state.get(tokenized_obj.rid) + if ( + state is None + or not state.abort_before_dispatch + or is_health_check_generate_req(tokenized_obj) + ): + return False + self._handle_abort_req( + AbortReq( + rid=tokenized_obj.rid, + abort_message="Aborted by AbortReq before dispatch to scheduler", + ) + ) + dispatch_ready = self.encoder_dispatch_ready.pop(tokenized_obj.rid, None) + if dispatch_ready is not None: + dispatch_ready.set() + return True + async def _send_one_request( self, tokenized_obj: Union[TokenizedGenerateReqInput, TokenizedEmbeddingReqInput], ): + if self._abort_instead_of_dispatch(tokenized_obj): + return prepared_mm_items = [] dispatched = False try: @@ -1652,6 +1761,13 @@ async def _send_batch_request( ], ): """Send a batch of tokenized requests as a single batched request to the scheduler.""" + tokenized_objs = [ + tokenized_obj + for tokenized_obj in tokenized_objs + if not self._abort_instead_of_dispatch(tokenized_obj) + ] + if not tokenized_objs: + return prepared_mm_items = [] dispatched = False try: @@ -1719,6 +1835,25 @@ def _coalesce_streaming_chunks( out["meta_info"] = meta_info return out + def _finalize_lora_lease(self, state: Optional[ReqState]) -> None: + """Release the request's LoRA lease exactly once, however it terminates: + normal finish, scheduler abort echo (queued / tokenizer-held / disagg), + status-code abort, or a failed dispatch. A request whose LoRA was never + acquired (pre-acquire failure leaves lora_id unset) has no lease to + release.""" + if state is None or state.lora_lease_released: + return + # lora_path / lora_id are declared on both input structs with default + # None (= base model / lease never acquired), so plain None checks are + # the contract. State-derived checks come first: a request without a + # lease has nothing to release, whatever the server config says. + if state.obj.lora_path is None or state.obj.lora_id is None: + return + if not self.enable_lora: + return + state.lora_lease_released = True + asyncio.create_task(self.lora_registry.release(state.obj.lora_id)) + async def _handle_abort_finish_reason( self, out: dict, @@ -1748,9 +1883,10 @@ async def _handle_abort_finish_reason( if state.obj.rid in self.rid_to_state: del self.rid_to_state[state.obj.rid] - # Mark ongoing LoRA request as finished. - if self.enable_lora and state.obj.lora_path: - await self.lora_registry.release(state.obj.lora_id) + # Status-code aborts also arrive through _handle_batch_output, which + # already released the lease and deleted the state — the finalizer's + # idempotency is what prevents the counter from going negative here. + self._finalize_lora_lease(state) if not is_stream: raise fastapi.HTTPException( status_code=finish_reason["status_code"], @@ -2039,20 +2175,39 @@ async def _stream_batch_responses(self, generators, rids): return_exceptions=True, ) - def abort_request(self, rid: str = "", abort_all: bool = False): + def abort_request( + self, rid: str = "", abort_all: bool = False, prefix: bool = False + ): # Empty rid would startswith-match every request on the scheduler. if not abort_all and not rid: logger.warning("Ignore abort_request with empty rid and abort_all=False") return state = None if abort_all else self.rid_to_state.get(rid) if not abort_all: - if state is not None: + if prefix: + if get_serving().tokenizer_worker_num == 1 and not any( + r.startswith(rid) for r in self.rid_to_state + ): + return + elif state is not None: if state.abort_sent: return state.abort_sent = True elif get_serving().tokenizer_worker_num == 1: return - req = AbortReq(rid=rid, abort_all=abort_all) + # A matching request can still be tokenizer-held: rid_to_state is + # populated in _init_req_state, several awaits (pause gate, + # model_update_lock, tokenization) before the scheduler learns the rid + # in _send_one_request. The scheduler-side abort cannot match those, so + # flag them here; the dispatch path resolves flagged requests as + # aborted instead of sending them. Already-dispatched requests are + # unaffected (the flag is only read at dispatch). + for tracked_rid, tracked_state in self.rid_to_state.items(): + if abort_all or ( + tracked_rid.startswith(rid) if prefix else tracked_rid == rid + ): + tracked_state.abort_before_dispatch = True + req = AbortReq(rid=rid, abort_all=abort_all, prefix=prefix) try: self._dispatch_to_scheduler(req) except BaseException: @@ -2326,12 +2481,11 @@ async def _handle_batch_output( ): output_sampling_mask = recv_obj.output_token_sampling_mask if output_sampling_mask is not None: - state.output_token_sampling_mask.extend(output_sampling_mask[i]) - output_sampling_logprobs = recv_obj.output_token_sampling_logprobs - if output_sampling_logprobs is not None: - state.output_token_sampling_logprobs.extend( - output_sampling_logprobs[i] - ) + masks, logprobs = output_sampling_mask[i].to_lists( + support_logprobs=state.obj.sampling_logprobs_mode == "support" + ) + state.output_token_sampling_mask.extend(masks) + state.output_token_sampling_logprobs.extend(logprobs) meta_info["output_token_sampling_mask"] = ( state.output_token_sampling_mask ) @@ -2406,6 +2560,9 @@ async def _handle_batch_output( if isinstance(val, torch.Tensor): val = pybase64.b64encode(val.numpy().tobytes()).decode("utf-8") meta_info["indexer_topk"] = val + n = getattr(recv_obj, "indexer_topk_num_layers", None) + if n is not None: + meta_info["indexer_topk_num_layers"] = n if getattr(recv_obj, "dp_ranks", None): meta_info["dp_rank"] = recv_obj.dp_ranks[i] @@ -2557,8 +2714,7 @@ async def _handle_batch_output( del self.rid_to_state[rid] # Mark ongoing LoRA request as finished. - if self.enable_lora and state.obj.lora_path: - asyncio.create_task(self.lora_registry.release(state.obj.lora_id)) + self._finalize_lora_lease(state) if out_dict is not None: state.out_list.append(out_dict) @@ -3356,6 +3512,10 @@ def _handle_abort_req(self, recv_obj: AbortReq): } if state.prompt_token_ids is not None: out["prompt_token_ids"] = state.prompt_token_ids + # This abort echo is the scheduler-side terminal ACK for queued, + # tokenizer-held, and disagg-retracted requests — none of them reach + # _handle_batch_output, so this is their only lease-release point. + self._finalize_lora_lease(self.rid_to_state.get(recv_obj.rid)) del self.rid_to_state[recv_obj.rid] state.out_list.append(out) @@ -3477,18 +3637,40 @@ async def _resolve_lora_path(self, obj: Union[GenerateReqInput, EmbeddingReqInpu unregistered_loras = await self.lora_registry.get_unregistered_loras( unique_lora_paths ) + backfill_paths = obj.lora_backfill_paths or {} for lora_path in unregistered_loras: if lora_path is None: continue - if lora_path not in self.lora_ref_cache: + if lora_path in self._pending_lora_publications: + # a staged session is streaming this name; a disk load now would give it two identities raise ValueError( - f"Got LoRA adapter that has never been loaded: {lora_path}\n" - f"All loaded adapters: {self.lora_ref_cache.keys()}." + f"LoRA adapter '{lora_path}' is awaiting publication and " + "cannot be served or backfilled yet." ) + if lora_path not in self.lora_ref_cache: + # a request-carried path lets a fresh or restarted engine serve any published version + if lora_path in backfill_paths: + self.lora_ref_cache[lora_path] = LoRARef( + lora_name=lora_path, + lora_path=backfill_paths[lora_path], + pinned=False, + ) + else: + raise ValueError( + f"Got LoRA adapter that has never been loaded: {lora_path}\n" + f"All loaded adapters: {self.lora_ref_cache.keys()}." + ) - logger.info(f"Reloading evicted adapter: {lora_path}") new_lora_ref = self.lora_ref_cache[lora_path] + if not new_lora_ref.reloadable: + raise ValueError( + f"LoRA adapter '{lora_path}' was loaded over the wire " + f"(lora_path={new_lora_ref.lora_path!r}) and has no disk " + "artifact to reload from; it must be re-pushed by the " + "trainer instead of implicitly reloaded." + ) + logger.info(f"Reloading evicted adapter: {lora_path}") load_result = await self.load_lora_adapter( LoadLoRAAdapterReqInput( lora_name=new_lora_ref.lora_name, @@ -3574,6 +3756,9 @@ def _release_req_states_on_failure(self, rids: Iterable[str]): "Failed to abort request %s during cleanup", rid ) else: + # A failed/partial dispatch never produces a scheduler + # terminal for this rid, so release the LoRA lease here. + self._finalize_lora_lease(state) del self.rid_to_state[rid] dispatch_ready = self.encoder_dispatch_ready.pop(rid, None) if dispatch_ready is not None: diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index 12b32503fdb4..15678a7593ff 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -30,14 +30,12 @@ GetWeightsByNameReqInput, InitWeightsSendGroupForRemoteInstanceReqInput, InitWeightsUpdateGroupReqInput, - LoadLoRAAdapterFromTensorsReqInput, LoadLoRAAdapterReqInput, + RegisterLoRAAdapterReqInput, SendWeightsToRemoteInstanceReqInput, UnloadLoRAAdapterReqInput, UpdateWeightFromDiskReqInput, - UpdateWeightsFromDistributedReqInput, UpdateWeightsFromIPCReqInput, - UpdateWeightsFromTensorReqInput, ) from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.scheduler import GenerationBatchResult @@ -62,14 +60,12 @@ get_spec, ) from sglang.srt.server_args import ServerArgs -from sglang.srt.utils import MultiprocessingSerializer, broadcast_pyobj, set_random_seed +from sglang.srt.utils import broadcast_pyobj, set_random_seed from sglang.srt.utils.hf_transformers_utils import ( get_processor, get_tokenizer, get_tokenizer_from_processor, ) -from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions -from sglang.srt.weight_sync.tensor_bucket import FlattenedTensorBucket if TYPE_CHECKING: from sglang.srt.managers.cache_controller import LayerDoneCounter @@ -195,36 +191,6 @@ def send_weights_to_remote_instance( ) return success, message - def update_weights_from_distributed( - self, recv_req: UpdateWeightsFromDistributedReqInput - ): - success, message = ( - self.model_runner.weight_updater.update_weights_from_distributed( - recv_req.names, - recv_req.dtypes, - recv_req.shapes, - recv_req.group_name, - recv_req.load_format, - ) - ) - return success, message - - def _deserialize_own_rank(self, serialized_named_tensors): - """Each rank deserializes only its own payload (index ps.tp_rank); - deserializing another rank's copy would break producer-side CUDA-IPC - refcounting.""" - monkey_patch_torch_reductions() - return MultiprocessingSerializer.deserialize( - serialized_named_tensors[self.ps.tp_rank] - ) - - def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput): - success, message = self.model_runner.weight_updater.update_weights_from_tensor( - named_tensors=self._deserialize_own_rank(recv_req.serialized_named_tensors), - load_format=recv_req.load_format, - ) - return success, message - def update_weights_from_ipc(self, recv_req: UpdateWeightsFromIPCReqInput): """Update weights from IPC for checkpoint-engine integration.""" success, message = self.model_runner.weight_updater.update_weights_from_ipc( @@ -246,58 +212,10 @@ def unload_lora_adapter(self, recv_req: UnloadLoRAAdapterReqInput): result = self.model_runner.unload_lora_adapter(recv_req.to_ref()) return result - def load_lora_adapter_from_tensors( - self, recv_req: LoadLoRAAdapterFromTensorsReqInput - ): - # The LoRA code handles TP sharding internally using slice_lora_a_weights - # and slice_lora_b_weights methods (see lora/layers.py and mem_pool.py). - data = self._deserialize_own_rank(recv_req.serialized_named_tensors) - if recv_req.load_format == "flattened_bucket": - bucket = FlattenedTensorBucket( - flattened_tensor=data["flattened_tensor"], - metadata=data["metadata"], - ) - tensors = dict(bucket.reconstruct_tensors()) - else: - tensors = data - if recv_req.expected_checksums is not None: - import hashlib - - exp = recv_req.expected_checksums - mismatch, missing = [], [] - for name, want in exp.items(): - if name not in tensors: - missing.append(name) - continue - got = hashlib.sha256( - tensors[name] - .detach() - .cpu() - .contiguous() - .flatten() - .view(torch.uint8) - .numpy() - .tobytes() - ).hexdigest() - if got != want: - mismatch.append(name) - extra = [n for n in tensors if n not in exp] - if mismatch or missing or extra: - raise RuntimeError( - f"[LORA-CHECK] rank{self.ps.tp_rank} adapter sync MISMATCH of {len(exp)} expected: " - f"{len(mismatch)} value-diff {mismatch[:5]}, {len(missing)} missing {missing[:5]}, " - f"{len(extra)} extra {extra[:5]}" - ) - logger.info( - f"[LORA-CHECK] rank{self.ps.tp_rank} adapter sync OK: {len(exp)}/{len(exp)} tensors match (sha256)" - ) - result = self.model_runner.load_lora_adapter_from_tensors( - recv_req.to_ref(), - tensors, - recv_req.config_dict, - recv_req.added_tokens_config, + def register_lora_adapter(self, recv_req: RegisterLoRAAdapterReqInput): + return self.model_runner.register_lora_adapter( + recv_req.to_ref(), recv_req.config_dict ) - return result def forward_batch_embedding(self, batch: ScheduleBatch): forward_batch = ForwardBatch.init_new( @@ -529,6 +447,11 @@ def _init_dllm_algorithm(self): def model_runner(self) -> ModelRunner: return self._model_runner + def iter_runners(self) -> List[Tuple[str, ModelRunner]]: + """(role, runner) pairs this worker owns for weight ops. The target worker + owns one runner and uses the empty role so its checksum keys stay unprefixed.""" + return [("", self._model_runner)] + def register_hicache_layer_transfer_counter(self, counter: LayerDoneCounter): self.hicache_layer_transfer_counter = counter diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py index be44849baa2b..4beedd4d64d9 100644 --- a/python/sglang/srt/managers/utils.py +++ b/python/sglang/srt/managers/utils.py @@ -280,6 +280,11 @@ def get_logprob_dict_from_result(result: GenerationBatchResult) -> dict: if sampling_mask_output is None else sampling_mask_output.selected_logprobs ), + "sampling_mask_support_logprobs": ( + None + if sampling_mask_output is None + else sampling_mask_output.support_logprobs + ), "sampling_mask_statuses": ( None if sampling_mask_output is None else sampling_mask_output.statuses ), @@ -301,6 +306,7 @@ def get_logprob_from_pp_outputs( token_ids=next_pp_outputs["sampling_mask_token_ids"], lengths=next_pp_outputs["sampling_mask_lengths"], selected_logprobs=next_pp_outputs["sampling_mask_selected_logprobs"], + support_logprobs=next_pp_outputs["sampling_mask_support_logprobs"], statuses=next_pp_outputs["sampling_mask_statuses"], ) logits_output = LogitsProcessorOutput( diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 04b28b8ac8cb..8de68ffbde00 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -60,6 +60,7 @@ mamba_cache_chunk_size, ) from sglang.srt.speculative.spec_info import SpecInputType +from sglang.srt.true_on_policy import is_true_on_policy_enabled from sglang.srt.utils import ( is_cpu, is_hip, @@ -903,8 +904,18 @@ def init_new( if ret.forward_mode.is_idle(): ret.positions = torch.empty((0,), dtype=torch.int64, device=device) + # Under --enable-dp-attention an IDLE rank still runs its local experts over + # the DP-GATHERED tokens of the other ranks, so MoE-LoRA batch info has to be + # re-prepared for THIS batch — sized to the gathered length and stamped via the + # idle-rank path — rather than cleared: clearing leaves those experts reading the + # previous batch's token_lora_mapping, which is undersized (OOB in the MoE-LoRA + # kernels) and points foreign tokens at the wrong adapter. Non-MoE LoRA does no + # such cross-rank work on an idle rank, so it takes the base path. if model_runner.lora_manager is not None: - model_runner.lora_manager.reset_lora_batch() + if model_runner.lora_manager.lora_backend.is_moe_lora: + model_runner.lora_manager.prepare_lora_batch(ret) + else: + model_runner.lora_manager.reset_lora_batch() return ret # Override the positions with diffusion LLM or spec_info @@ -1253,7 +1264,11 @@ def _compute_mrope_positions_decode( ) has_multimodal_input = any(mm is not None for mm in mm_inputs) - if rl_on_policy_target is not None or not has_multimodal_input: + if ( + rl_on_policy_target is not None + or is_true_on_policy_enabled() + or not has_multimodal_input + ): # Text-only positions_1d = seq_lens_int64 - 1 self.mrope_positions = positions_1d.unsqueeze(0).repeat(3, 1) @@ -1314,7 +1329,11 @@ def _compute_mrope_positions_extend( mm_input = mm_inputs[batch_idx] extend_seq_len = extend_lens[batch_idx] extend_prefix_len = prefix_lens[batch_idx] - if mm_input is None or rl_on_policy_target is not None: + if ( + mm_input is None + or rl_on_policy_target is not None + or is_true_on_policy_enabled() + ): # text only mrope_positions = ( torch.arange( diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 40a7ce104966..2d470e37db5e 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -20,7 +20,7 @@ import logging import time from dataclasses import dataclass -from typing import Callable, Optional, Union +from typing import Callable, List, Optional, Union import torch import torch.distributed as dist @@ -169,6 +169,11 @@ EagerRunner, get_batch_sizes_to_capture, ) +from sglang.srt.model_loader.loader import ( + post_load_weights, + postprocess_weight, + restore_weight, +) from sglang.srt.platforms import current_platform from sglang.srt.runtime_context import ( assert_published, @@ -211,6 +216,7 @@ get_global_experts_capturer, set_global_experts_capturer, ) +from sglang.srt.true_on_policy import is_tp_invariant_target from sglang.srt.utils import ( cpu_has_amx_support, enable_show_time_cost, @@ -274,6 +280,9 @@ def _prefill_cuda_graph_allows_context_parallel( class ModelRunnerOutput: logits_output: Union[LogitsProcessorOutput, PPProxyTensors] can_run_graph: bool + # Per-rank padded size of the graph this forward replayed (None for eager); + # the dp-gathered buffer is strided by it. + graph_num_tokens: Optional[int] = None expert_distribution_metrics: Optional[ExpertDistributionMetrics] = None routed_experts_output: Optional[TopkCaptureOutput] = None indexer_topk_output: Optional[TopkCaptureOutput] = None @@ -592,6 +601,7 @@ def init_weight_exporter(self): def init_remote_instance_weight_transporter(self): self.remote_instance_weight_transporter = RemoteInstanceWeightTransporter( + server_args=self.server_args, get_model=lambda: self.model, tp_rank=self.ps.tp_rank, gpu_id=self.gpu_id, @@ -701,6 +711,7 @@ def init_memory_saver_adapter(self): def maybe_init_remote_instance_transfer_engine(self): if remote_instance_transfer_engine_enabled(load_format=self.draft_load_format): self.remote_instance_weight_transporter.init_engine() + self.remote_instance_weight_transporter.maybe_init_parallelism_config() def maybe_init_expert_location_metadata(self): if self.is_draft_worker: @@ -796,6 +807,10 @@ def maybe_enable_batch_invariant_mode(self): from sglang.srt.batch_invariant_ops import enable_batch_invariant_mode enable_batch_invariant_mode() + if is_tp_invariant_target(): + from sglang.srt.tp_invariant_ops import enable_tp_invariant_mode + + enable_tp_invariant_mode() def get_pp_proxy_topk_size(self) -> Optional[int]: return misc_utils.resolve_pp_proxy_topk_size( @@ -931,7 +946,8 @@ def _init_post_memory_pool_components(self): self.maybe_init_hisparse_coordinator() - self.init_routed_experts_capturer() + if not self.is_draft_worker: + self.init_routed_experts_capturer() self.init_indexer_capturer() self.graph_shared_output = None @@ -1358,12 +1374,9 @@ def load_lora_adapter(self, lora_ref: LoRARef): """Load a new lora adapter from disk or huggingface.""" return self.lora_manager.load_lora_adapter(lora_ref) - def load_lora_adapter_from_tensors( - self, lora_ref: LoRARef, tensors, config_dict, added_tokens_config=None - ): - return self.lora_manager.load_lora_adapter_from_tensors( - lora_ref, tensors, config_dict, added_tokens_config - ) + def register_lora_adapter(self, lora_ref: LoRARef, config_dict): + logger.info(f"LoRA adapter registration: {lora_ref}.") + return self.lora_manager.register_lora_adapter(lora_ref, config_dict) def unload_lora_adapter(self, lora_ref: LoRARef): """Unload a lora adapter that was previously loaded during initialization or dynamic loading.""" @@ -1689,6 +1702,14 @@ def forward( output.expert_distribution_metrics = recorder_outputs.get("metrics") no_copy_to_cpu = not get_schedule().disable_overlap_schedule + # In speculative decoding more than one token is captured per request, so + # pass the actual number of tokens per DP rank in CUDA graph, not the batch + # size — the width is captured_req_width. + + # From the runner that ran this forward: the decode runner's own bs is + # another graph's stride once a prefill graph replays. + cuda_graph_num_tokens = output.graph_num_tokens + if ( not self.is_draft_worker and (experts_capturer := get_global_experts_capturer()) is not None @@ -1696,7 +1717,7 @@ def forward( output.routed_experts_output = experts_capturer.on_forward_end( forward_batch=forward_batch, can_run_graph=output.can_run_graph, - cuda_graph_batch=getattr(self.decode_cuda_graph_runner, "bs", None), + cuda_graph_batch=cuda_graph_num_tokens, no_copy_to_cpu=no_copy_to_cpu, ) @@ -1704,7 +1725,7 @@ def forward( output.indexer_topk_output = indexer_capturer.on_forward_end( forward_batch=forward_batch, can_run_graph=output.can_run_graph, - cuda_graph_batch=getattr(self.decode_cuda_graph_runner, "bs", None), + cuda_graph_batch=cuda_graph_num_tokens, no_copy_to_cpu=no_copy_to_cpu, ) @@ -1802,13 +1823,20 @@ def _forward_raw( self.hisparse_coordinator.wait_for_pending_backup() self.hisparse_coordinator.num_real_reqs.fill_(forward_batch.batch_size) + graph_num_tokens = None + # Replay cuda graph if applicable if can_run_graph: ret = self.decode_cuda_graph_runner.execute( forward_batch, pp_proxy_tensors=pp_proxy_tensors, ) - return ModelRunnerOutput(logits_output=ret, can_run_graph=can_run_graph) + return ModelRunnerOutput( + logits_output=ret, + can_run_graph=can_run_graph, + graph_num_tokens=self.decode_cuda_graph_runner.bs + * self.decode_cuda_graph_runner.captured_req_width, + ) # DP / MLP-sync padding + attn-tp normalization. Only the decode # cuda-graph path above pre-pads its static buffers and returns @@ -1857,6 +1885,7 @@ def _forward_raw( forward_batch, **kwargs ) can_run_graph = True + graph_num_tokens = self.prefill_cuda_graph_runner.last_replay_num_tokens else: # Eager: decode / extend / idle dispatched inside the runner. ret = self.eager_runner.execute( @@ -1869,7 +1898,11 @@ def _forward_raw( ): forward_batch.post_forward_mlp_sync_batch(ret) - return ModelRunnerOutput(logits_output=ret, can_run_graph=can_run_graph) + return ModelRunnerOutput( + logits_output=ret, + can_run_graph=can_run_graph, + graph_num_tokens=graph_num_tokens, + ) def _preprocess_logits( self, @@ -1986,9 +2019,29 @@ def compute_logprobs_only( forward_batch.token_ids_logprobs, ) - def check_weights(self, action: str, allow_quant_error: bool = False): + def begin_weight_update(self) -> None: + """Begin a weight-update session: restore in-place-packed weights to a + loadable state (no-op for schemes that don't repack).""" + restore_weight(self.model, torch.device(self.device)) + + def end_weight_update(self, run_post_load: bool) -> None: + """End the weight-update session: optionally run model.post_load_weights + (when load_weights was bypassed this session, e.g. P2P/RDMA), then finalize + quantized weights into kernel layout.""" + if run_post_load: + post_load_weights(self.model) + postprocess_weight(self.model, torch.device(self.device)) + + def check_weights( + self, + action: str, + allow_quant_error: bool = False, + skip_tensor_list: Optional[List[str]] = None, + ): return self._weight_checker.handle( - action=action, allow_quant_error=allow_quant_error + action=action, + allow_quant_error=allow_quant_error, + skip_tensor_list=skip_tensor_list, ) def _expand_eplb_metadata_for_scale( diff --git a/python/sglang/srt/model_executor/model_runner_components/remote_instance_weight_transporter.py b/python/sglang/srt/model_executor/model_runner_components/remote_instance_weight_transporter.py index 14690ab217a7..5698ed940478 100644 --- a/python/sglang/srt/model_executor/model_runner_components/remote_instance_weight_transporter.py +++ b/python/sglang/srt/model_executor/model_runner_components/remote_instance_weight_transporter.py @@ -6,6 +6,7 @@ import torch +from sglang.srt.distributed.parallel_state import RankParallelismConfig from sglang.srt.environ import envs from sglang.srt.model_loader.remote_instance_weight_loader_utils import ( RemoteInstanceWeightLoaderBackend, @@ -16,6 +17,7 @@ get_parallel, remote_instance_transfer_engine_enabled, ) +from sglang.srt.server_args import ServerArgs from sglang.srt.utils.network import NetworkAddress, get_local_ip_auto logger = logging.getLogger(__name__) @@ -23,12 +25,14 @@ @dataclass(slots=True, kw_only=True) class RemoteInstanceWeightTransporter: + server_args: ServerArgs get_model: Callable[[], torch.nn.Module] tp_rank: int gpu_id: int engine: Optional[Any] = None session_id: str = "" weight_info: Optional[dict[str, tuple[int, int, int]]] = None + parallelism_config: Optional[RankParallelismConfig] = None _nixl_manager: Optional[Any] = None @property @@ -55,6 +59,12 @@ def init_engine(self): local_ip, self.engine.get_rpc_port() ).to_host_port_str() + def maybe_init_parallelism_config(self) -> None: + if self.server_args.registers_parallelism_config(): + self.parallelism_config = RankParallelismConfig.from_parallel_state( + self.tp_rank + ) + def maybe_register_and_publish_weight_info(self) -> None: if ( remote_instance_transfer_engine_enabled() @@ -70,6 +80,51 @@ def maybe_register_and_publish_weight_info(self) -> None: self.weight_info = register_memory_region(self.model, self.engine) self._register_to_engine_info_bootstrap() + # The P2P weight-update client needs each rank's parallelism layout to + # map training-side parameters onto this rank's shards. + if ( + self.server_args.registers_parallelism_config() + and self.parallelism_config is not None + ): + self._register_parallelism_config_to_bootstrap() + + def _bootstrap_url(self) -> str: + if self.server_args.dist_init_addr: + bootstrap_host = ( + NetworkAddress.parse(self.server_args.dist_init_addr).resolved().host + ) + else: + bootstrap_host = "127.0.0.1" + bootstrap_port = self.server_args.engine_info_bootstrap_port + return NetworkAddress(bootstrap_host, bootstrap_port).to_url() + + def _register_parallelism_config_to_bootstrap(self) -> None: + """Register this rank's parallelism config with the EngineInfoBootstrapServer.""" + import requests as http_requests + + bootstrap_url = self._bootstrap_url() + url = f"{bootstrap_url}/register_parallelism_config" + payload = { + "tp_rank": self.tp_rank, + "parallelism_config": self.parallelism_config.to_dict(), + } + try: + resp = http_requests.put(url, json=payload, timeout=5) + if resp.status_code == 200: + logger.info( + f"Registered parallelism config for tp_rank={self.tp_rank} " + f"with bootstrap server at {bootstrap_url}" + ) + else: + logger.error( + f"Failed to register parallelism config for tp_rank={self.tp_rank}: " + f"{resp.status_code}, {resp.text}" + ) + except Exception as e: + logger.error( + f"Failed to register parallelism config for tp_rank={self.tp_rank}: {e}" + ) + def _register_to_engine_info_bootstrap(self: RemoteInstanceWeightTransporter): """Register transfer engine info with the EngineInfoBootstrapServer via HTTP PUT. diff --git a/python/sglang/srt/model_executor/model_runner_components/weight_updater.py b/python/sglang/srt/model_executor/model_runner_components/weight_updater.py index e9dd1c7023f9..1b6b4057a2b3 100644 --- a/python/sglang/srt/model_executor/model_runner_components/weight_updater.py +++ b/python/sglang/srt/model_executor/model_runner_components/weight_updater.py @@ -220,6 +220,75 @@ def model_load_weights(model, iter): logger.info("Update weights end.") return True, "Succeeded to update model weights." + def load_weights(self: WeightUpdater, weights) -> None: + """Load an in-memory list of (name, tensor) weights into this runner's model.""" + self.get_model().load_weights(weights) + + def receive_weights_from_distributed( + self: WeightUpdater, + names, + dtypes, + shapes, + group_name, + load_format: Optional[str] = None, + ): + """Receive one weight broadcast from the training engine over this runner's + `_model_update_group` and return the named tensors WITHOUT loading them. + + Only the runner that joined the group (the target / main model) can receive; + the caller loads the result into each runner it wants updated. Speculative + draft runners never join the group, so they are fed from here. + """ + assert group_name in self._model_update_group, ( + f"Group {group_name} not in {list(self._model_update_group.keys())}. " + "Please call `init_weights_update_group` first." + ) + + if load_format == "flattened_bucket": + return self._receive_bucketed_weights_from_distributed( + names, dtypes, shapes, group_name + ) + + weights = [] + handles = [] + for name, dtype, shape in zip(names, dtypes, shapes): + target_dtype = ( + dtype if isinstance(dtype, torch.dtype) else getattr(torch, dtype) + ) + weight = torch.empty(shape, dtype=target_dtype, device=self.device) + handles.append( + torch.distributed.broadcast( + weight, + src=0, + group=self._model_update_group[group_name], + async_op=True, + ) + ) + weights.append((name, weight)) + for handle in handles: + handle.wait() + return weights + + def _receive_bucketed_weights_from_distributed( + self: WeightUpdater, names, dtypes, shapes, group_name + ): + named_tensors = [] + for name, dtype, shape in zip(names, dtypes, shapes): + target_dtype = ( + dtype if isinstance(dtype, torch.dtype) else getattr(torch, dtype) + ) + named_tensors.append( + (name, torch.empty(shape, dtype=target_dtype, device=self.device)) + ) + bucket = FlattenedTensorBucket(named_tensors=named_tensors) + flattened_tensor = bucket.get_flattened_tensor() + torch.distributed.broadcast( + flattened_tensor, + src=0, + group=self._model_update_group[group_name], + ) + return bucket.reconstruct_tensors() + def update_weights_from_distributed( self: WeightUpdater, names, @@ -252,26 +321,10 @@ def update_weights_from_distributed( names, dtypes, shapes, group_name ) try: - weights = [] - handles = [] - for name, dtype, shape in zip(names, dtypes, shapes): - target_dtype = ( - dtype if isinstance(dtype, torch.dtype) else getattr(torch, dtype) - ) - weight = torch.empty(shape, dtype=target_dtype, device=self.device) - handles.append( - torch.distributed.broadcast( - weight, - src=0, - group=self._model_update_group[group_name], - async_op=True, - ) - ) - weights.append((name, weight)) - for handle in handles: - handle.wait() - - self.get_model().load_weights(weights) + weights = self.receive_weights_from_distributed( + names, dtypes, shapes, group_name, load_format + ) + self.load_weights(weights) return True, "Succeeded to update parameter online." except Exception as e: diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index d5160d43b925..fd0c49c05c6b 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py @@ -114,6 +114,9 @@ get_spec, ) from sglang.srt.speculative.ragged_verify import resolve_ragged_verify_layout +from sglang.srt.true_on_policy import ( + patch_prefill_only_deterministic_inference_for_cuda_graph, +) from sglang.srt.utils import ( empty_context, get_available_gpu_memory, @@ -1028,27 +1031,34 @@ def capture(self) -> None: # Trigger CUDA graph capture for specific shapes. # Capture the large shapes first so that the smaller shapes # can reuse the memory pool allocated for the large shapes. - with freeze_gc(get_exec().graph.enable_cudagraph_gc): - if not self.enable_pdmux: - with ( - graph_capture( - stream=get_or_create_global_graph_capture_stream() - ) as graph_capture_context, - profile_context as prof, - ): - self.stream = graph_capture_context.stream - with self.backend.capture_session(self.stream): - self._capture_one_stream() - else: - set_pdmux_status(False) - for i, sg in enumerate(self.stream_groups): + with patch_prefill_only_deterministic_inference_for_cuda_graph( + self.model_runner.server_args, + attn_backend=getattr(self.model_runner, "attn_backend", None), + dvr_target_verify_cuda_graph=getattr( + self.model_runner, "enable_dvr_target_verify_cuda_graph", False + ), + ): + with freeze_gc(get_exec().graph.enable_cudagraph_gc): + if not self.enable_pdmux: with ( - graph_capture(stream=sg[1]) as graph_capture_context, + graph_capture( + stream=get_or_create_global_graph_capture_stream() + ) as graph_capture_context, profile_context as prof, ): self.stream = graph_capture_context.stream with self.backend.capture_session(self.stream): - self._capture_one_stream(i) + self._capture_one_stream() + else: + set_pdmux_status(False) + for i, sg in enumerate(self.stream_groups): + with ( + graph_capture(stream=sg[1]) as graph_capture_context, + profile_context as prof, + ): + self.stream = graph_capture_context.stream + with self.backend.capture_session(self.stream): + self._capture_one_stream(i) if self.enable_profile_cuda_graph: self._post_process_after_profile(prof) diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index ab330617beed..b97ad99446f9 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -1922,6 +1922,9 @@ def execute( with self.backend.replay_session(): static_forward_batch = self.load_batch(forward_batch, **kwargs) static_num_tokens = len(static_forward_batch.input_ids) + # Stride of the dp-gathered buffer for this replay; surfaced through + # ModelRunnerOutput.graph_num_tokens. + self.last_replay_num_tokens = static_num_tokens raw_num_tokens = self.raw_num_tokens shape_key = self._shape_key(static_num_tokens, forward_batch) # The only variants this runner records are chunked-prefix ones. diff --git a/python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py b/python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py index 3a047c8ffa64..69015f0174f0 100644 --- a/python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py +++ b/python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py @@ -167,6 +167,10 @@ def capture_one( graph_ctx = partial( self._memory_saver_adapter.cuda_graph, tag=GPU_MEMORY_TYPE_CUDA_GRAPH, + # The pool holds capture-time-initialized device state that replays + # read but never rewrite, so pausing must preserve the pool contents + # instead of discarding them. + enable_cpu_backup=True, ) else: graph_ctx = self._device_module.graph diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 1d1a35c37d47..809fb7e3ee5a 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -304,7 +304,7 @@ def _initialize_model( return model_class(**kwargs) -def _post_load_weights(model: nn.Module) -> None: +def post_load_weights(model: nn.Module) -> None: # Loaders that bypass `model.load_weights()` (dummy / sharded state / remote instance / # remote fs) must trigger the model's post-load fixup explicitly; `model.load_weights()` # would normally do it internally. NextN subclasses override the method to fill in @@ -313,6 +313,36 @@ def _post_load_weights(model: nn.Module) -> None: model.post_load_weights() +def _apply_quant_method_hook(model: nn.Module, target_device, hook_name: str) -> None: + """Run a quant_method hook (restore/process) on every quantized module. + + LoRA wrappers forward `quant_method` for forward-path dispatch but don't own + the packed params; skip them so the inner base layer (yielded separately by + `named_modules`) handles it. + """ + from sglang.srt.lora.layers import BaseLayerWithLoRA + + for _, module in model.named_modules(): + if isinstance(module, BaseLayerWithLoRA): + continue + quant_method = getattr(module, "quant_method", None) + if quant_method is not None and hasattr(quant_method, hook_name): + with device_loading_context(module, target_device): + getattr(quant_method, hook_name)(module) + + +def restore_weight(model: nn.Module, target_device) -> None: + """Undo in-place quant packing so fresh weights can be loaded + (no-op for schemes that don't repack, e.g. plain fp8).""" + _apply_quant_method_hook(model, target_device, "restore_weights_before_loading") + + +def postprocess_weight(model: nn.Module, target_device) -> None: + """Finalize quantized weights into kernel layout (Marlin repack, UE8M0 requant, + transpose, ...).""" + _apply_quant_method_hook(model, target_device, "process_weights_after_loading") + + class BaseModelLoader(ABC): """Base class for model loaders.""" @@ -836,7 +866,7 @@ def prepare_model_for_capture( """ with set_default_torch_dtype(model_config.dtype): initialize_capture_safe_weights(model) - _post_load_weights(model) + post_load_weights(model) for _, module in model.named_modules(): quant_method = getattr(module, "quant_method", None) if quant_method is None: @@ -1643,7 +1673,7 @@ def load_model( # random values to the weights. initialize_dummy_weights(model) - _post_load_weights(model) + post_load_weights(model) for _, module in model.named_modules(): quant_method = getattr(module, "quant_method", None) @@ -1804,7 +1834,7 @@ def load_model( if state_dict: raise ValueError(f"Missing keys {tuple(state_dict)} in loaded state!") - _post_load_weights(model) + post_load_weights(model) return model.eval() @@ -3426,7 +3456,7 @@ def load_model_from_remote_instance_by_nccl( ) current_platform.synchronize() - _post_load_weights(model) + post_load_weights(model) end_get_weights_tic = time.time() logger.debug( f"finish getting all weights from remote instance, time used: {(end_get_weights_tic - start_get_weights_tic):.4f}s" @@ -3489,7 +3519,7 @@ def load_model_from_remote_instance_by_transfer_engine( logger.error(f"batch transfer failed, error: {ret}") return False - _post_load_weights(model) + post_load_weights(model) return True @@ -3579,7 +3609,7 @@ def _load_model_from_remote_kv( if state_dict: raise ValueError(f"Missing keys {tuple(state_dict)} in loaded state!") - _post_load_weights(model) + post_load_weights(model) def _load_model_from_remote_fs( self, model, client, model_config: ModelConfig, device_config: DeviceConfig diff --git a/python/sglang/srt/model_loader/parameter_mapper.py b/python/sglang/srt/model_loader/parameter_mapper.py new file mode 100644 index 000000000000..56056d570843 --- /dev/null +++ b/python/sglang/srt/model_loader/parameter_mapper.py @@ -0,0 +1,261 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Parameter mapping from HuggingFace checkpoint names to SGLang model parameters. + +This module provides utilities for translating weight names between HuggingFace +checkpoint format and SGLang's internal parameter naming, handling: + +1. Stacked Parameter Fusion + - gate_proj + up_proj → gate_up_proj (num_shards=2) + - q_proj + k_proj + v_proj → qkv_proj (num_shards=3) + - q_a_proj + kv_a_proj_with_mqa → fused_qkv_a_proj_with_mqa (DeepSeek MLA) + +2. Expert Parameter Sharding (MoE models) + - experts.{id}.gate_proj + experts.{id}.up_proj → experts.w13_weight (num_shards=2) + - experts.{id}.down_proj → experts.w2_weight (num_shards=1) + - Handles expert parallelism: num_local_experts = n_routed // ep_size + shared + +3. Scale Remapping (Quantized models) + - k_proj.k_scale → attn.k_scale + - v_proj.v_scale → attn.v_scale + - Quark-specific: output_scale → per-component scales + +Supported Models: + Dense: Llama, Qwen2, Qwen3, GLM4 + MoE: DeepSeekV2/V3/R1, Qwen3-MoE, GLM4-MoE, GLM4-MoE-Lite (GLM-4.7) + +Example: + >>> mapper = ParameterMapper.from_model(model) + >>> result = mapper.map("model.layers.0.mlp.gate_proj.weight") + >>> result.sglang_name # "model.layers.0.mlp.gate_up_proj.weight" + >>> result.shard_id # 0 + >>> result.num_shards # 2 +""" + +from dataclasses import dataclass +from typing import Callable, Dict, List, Optional, Tuple, Union + +StackedParamsEntry = Tuple[str, str, Union[int, str]] +ExpertParamsEntry = Tuple[str, str, int, Union[int, str]] + + +@dataclass +class MappingResult: + """Result of mapping a HuggingFace checkpoint weight name to SGLang parameter.""" + + sglang_name: str + shard_id: Optional[Union[int, str]] + num_shards: int + expert_id: Optional[int] + num_local_experts: Optional[int] + + +# Standard FP8 scale remapping patterns +_SCALE_REMAP_PATTERNS: List[Tuple[str, str, str]] = [ + (".k_scale", ".self_attn.k_proj.k_scale", ".self_attn.attn.k_scale"), + (".v_scale", ".self_attn.v_proj.v_scale", ".self_attn.attn.v_scale"), + (".k_scale", ".k_scale", ".attn.k_scale"), + (".v_scale", ".v_scale", ".attn.v_scale"), +] + +# Quark quantization scale remapping +_QUARK_SCALE_REMAP: Dict[str, str] = { + ".q_proj.output_scale": ".attn.q_scale", + ".k_proj.output_scale": ".attn.k_scale", + ".v_proj.output_scale": ".attn.v_scale", + "self_attn.prob_output_scale": ".attn.prob_scale", +} + + +class ParameterMapper: + """Maps HuggingFace checkpoint weight names to SGLang model parameters. + + This class pre-computes lookup tables at initialization for efficient + repeated mapping. It handles: + - Stacked/fused parameter mapping (gate_up_proj, qkv_proj, etc.) + - Expert parameter mapping with shard information + - Scale remapping for quantized models + - Model-specific weight name mutations + """ + + def __init__( + self, + stacked_params_mapping: List[StackedParamsEntry], + expert_params_mapping: List[ExpertParamsEntry], + num_local_experts: int = 0, + mutate_weight_preload: Optional[Callable[[str], str]] = None, + custom_scale_remap: Optional[Callable[[str], str]] = None, + ): + """Initialize the parameter mapper with model-specific configuration. + + Args: + stacked_params_mapping: List of (sglang_name, hf_name, shard_id) tuples. + Example: [("gate_up_proj", "gate_proj", 0), ("gate_up_proj", "up_proj", 1)] + expert_params_mapping: List of (sglang_name, hf_name, expert_id, shard_id) tuples. + Example: [("w13_weight", "experts.0.gate_proj.weight", 0, 0), ...] + num_local_experts: Number of experts in the current model rank. + For EP=1: num_local_experts = n_routed_experts + num_fused_shared_experts + For EP>1: num_local_experts = n_routed_experts // ep_size + num_fused_shared_experts + mutate_weight_preload: Optional function to transform weight names before mapping. + Used for shared expert fusion in DeepSeek (shared_experts → experts.{n_routed}). + custom_scale_remap: Optional function for model-specific scale remapping. + Used for DeepSeek k_proj/v_proj → attn_mqa scale mapping. + """ + self.num_local_experts = num_local_experts + self._mutate_weight_preload = mutate_weight_preload + self._custom_scale_remap = custom_scale_remap + + self._stacked_lookup, self._stacked_num_shards = self._build_stacked_lookup( + stacked_params_mapping + ) + self._expert_lookup, self._expert_num_shards = self._build_expert_lookup( + expert_params_mapping + ) + + @staticmethod + def _build_stacked_lookup( + mapping: List[StackedParamsEntry], + ) -> Tuple[Dict[str, Tuple[str, Union[int, str]]], Dict[str, int]]: + """Build lookup table and num_shards from stacked params mapping.""" + lookup: Dict[str, Tuple[str, Union[int, str]]] = {} + shard_counts: Dict[str, int] = {} + + for sglang_name, hf_name, shard_id in mapping: + lookup[hf_name] = (sglang_name, shard_id) + shard_counts[sglang_name] = shard_counts.get(sglang_name, 0) + 1 + + return lookup, shard_counts + + @staticmethod + def _build_expert_lookup( + mapping: List[ExpertParamsEntry], + ) -> Tuple[Dict[str, Tuple[str, int, Union[int, str]]], Dict[str, int]]: + """Build lookup table and num_shards from expert params mapping.""" + lookup: Dict[str, Tuple[str, int, Union[int, str]]] = {} + shard_counts: Dict[str, int] = {} + + for sglang_name, hf_name, expert_id, shard_id in mapping: + lookup[hf_name] = (sglang_name, expert_id, shard_id) + + for sglang_name, _, _, shard_id in mapping: + key = sglang_name + if key not in shard_counts: + unique_shards = set( + s_id for s_name, _, _, s_id in mapping if s_name == sglang_name + ) + shard_counts[key] = len(unique_shards) + + return lookup, shard_counts + + def _apply_scale_remap(self, name: str) -> str: + """Apply standard and Quark scale remapping patterns.""" + for suffix, pattern, replacement in _SCALE_REMAP_PATTERNS: + if name.endswith(suffix) and pattern in name: + return name.replace(pattern, replacement) + + for quark_suffix, replacement in _QUARK_SCALE_REMAP.items(): + if name.endswith(quark_suffix): + return name.replace(quark_suffix, replacement) + + return name + + def map(self, hf_weight_name: str) -> MappingResult: + """Map a HuggingFace checkpoint weight name to SGLang parameter info. + + Args: + hf_weight_name: The weight name from HuggingFace checkpoint. + + Returns: + MappingResult with mapped name and sharding information. + """ + name = hf_weight_name + + if self._mutate_weight_preload is not None: + name = self._mutate_weight_preload(name) + + if "scale" in name: + if self._custom_scale_remap is not None: + remapped = self._custom_scale_remap(name) + if remapped != name: + name = remapped + else: + name = self._apply_scale_remap(name) + else: + name = self._apply_scale_remap(name) + + for hf_pattern, ( + sglang_name, + expert_id, + shard_id, + ) in self._expert_lookup.items(): + if hf_pattern in name: + mapped_name = name.replace(hf_pattern, sglang_name) + return MappingResult( + sglang_name=mapped_name, + shard_id=shard_id, + num_shards=self._expert_num_shards.get(sglang_name, 1), + expert_id=expert_id, + num_local_experts=self.num_local_experts, + ) + + for hf_pattern, (sglang_name, shard_id) in self._stacked_lookup.items(): + if hf_pattern in name: + mapped_name = name.replace(hf_pattern, sglang_name) + return MappingResult( + sglang_name=mapped_name, + shard_id=shard_id, + num_shards=self._stacked_num_shards.get(sglang_name, 1), + expert_id=None, + num_local_experts=None, + ) + + return MappingResult( + sglang_name=name, + shard_id=None, + num_shards=1, + expert_id=None, + num_local_experts=None, + ) + + @classmethod + def from_model(cls, model) -> "ParameterMapper": + """Create a ParameterMapper from a model instance; currently supports + DeepseekV2ForCausalLM, Glm4ForCausalLM, Glm4MoeForCausalLM, + Glm4MoeLiteForCausalLM, LlamaForCausalLM, Qwen2ForCausalLM, + Qwen3ForCausalLM, Qwen3MoeForCausalLM.""" + stacked_mapping = list(getattr(model, "stacked_params_mapping", []) or []) + expert_mapping = list(getattr(model, "expert_params_mapping", []) or []) + + num_local_experts = 0 + if hasattr(model, "num_local_experts"): + num_local_experts = model.num_local_experts + elif expert_mapping: + expert_ids = set(entry[2] for entry in expert_mapping) + num_local_experts = len(expert_ids) + + mutate_fn = None + if hasattr(model, "mutate_weight_preload"): + mutate_fn = model.mutate_weight_preload + + scale_fn = None + if hasattr(model, "custom_scale_remap"): + scale_fn = model.custom_scale_remap + + return cls( + stacked_params_mapping=stacked_mapping, + expert_params_mapping=expert_mapping, + num_local_experts=num_local_experts, + mutate_weight_preload=mutate_fn, + custom_scale_remap=scale_fn, + ) diff --git a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py index c526df7759d1..df11ddcae002 100644 --- a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py +++ b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py @@ -25,7 +25,6 @@ from sglang.srt.distributed.parallel_state import GroupCoordinator from sglang.srt.environ import envs from sglang.srt.layers import deep_gemm_wrapper -from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.quantization.fp8_utils import ( block_quant_dequant, @@ -59,7 +58,6 @@ _use_aiter_gfx95, awq_dequantize_func, enable_nextn_moe_bf16_cast_to_fp8, - is_wint4afp8_or_wint4a16_config, ) from sglang.srt.utils import bind_or_assign, get_bool_env_var, log_info_on_rank0 @@ -210,6 +208,18 @@ class DeepseekV2WeightLoaderMixin: quant_config: Optional[QuantizationConfig] pp_group: GroupCoordinator num_fused_shared_experts: int + # Weight mapping relationships determined at model initialization time. + fuse_qkv_a_proj: bool + stacked_params_mapping: List[Tuple[str, str, int]] + expert_params_mapping: List[Tuple[str, str, int, int]] + + def mutate_weight_preload(self, name: str) -> str: + """Override in subclass for model-specific weight name mutations.""" + return name + + def custom_scale_remap(self, name: str) -> str: + """Override in subclass for model-specific scale remapping.""" + return name def do_load_weights( self, @@ -228,33 +238,7 @@ def do_load_weights( weights, NVFP4_CKPT_FP8_ATTN_QUANT_MODULES, nextn_conf ) - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("gate_up_proj", "gate_proj", 0), - ("gate_up_proj", "up_proj", 1), - ] - - # Params for weights, fp8 weight scales, fp8 activation scales - # (param_name, weight_name, expert_id, shard_id) - expert_params_mapping = FusedMoE.make_expert_params_mapping( - ckpt_gate_proj_name="gate_proj", - ckpt_down_proj_name="down_proj", - ckpt_up_proj_name="up_proj", - num_experts=self.config.n_routed_experts + self.num_fused_shared_experts, - ) - # Params for special naming rules in mixed-precision models, for example: - # model.layers.xx.mlp.experts.xx.w1.input_scale. For details, - # see https://huggingface.co/Barrrrry/DeepSeek-R1-W4AFP8/blob/main. - if is_wint4afp8_or_wint4a16_config(self.quant_config): - expert_params_mapping += FusedMoE.make_expert_input_scale_params_mapping( - num_experts=self.config.n_routed_experts - ) - - # Fuse q_a_proj and kv_a_proj_with_mqa along output dimension when q_lora_rank is not None - fuse_qkv_a_proj = hasattr(self.config, "q_lora_rank") and ( - self.config.q_lora_rank is not None - ) - cached_a_proj = {} if fuse_qkv_a_proj else None + cached_a_proj = {} if self.fuse_qkv_a_proj else None pending_indexer_wk: Dict[str, Dict[str, torch.Tensor]] = {} @@ -282,11 +266,7 @@ def do_load_weights( ) ): continue - if self.num_fused_shared_experts > 0 and "mlp.shared_experts" in name: - name = name.replace( - "mlp.shared_experts", - f"mlp.experts.{self.config.n_routed_experts}", - ) + name = self.mutate_weight_preload(name) weight_names.append(name) @@ -343,7 +323,7 @@ def do_load_weights( ): continue - for param_name, weight_name, shard_id in stacked_params_mapping: + for param_name, weight_name, shard_id in self.stacked_params_mapping: # Skip non-stacked layers and experts (experts handled below). if weight_name not in name: continue @@ -357,6 +337,10 @@ def do_load_weights( # for mlp.experts[0].gate_gate_up_proj, which breaks load. if ("mlp.experts." in name) and name not in params_dict: continue + # q_a_proj / kv_a_proj_with_mqa must bypass stacked loading + # and use the cache+concat fused A-proj path below. + if param_name == "fused_qkv_a_proj_with_mqa": + continue name = name.replace(weight_name, param_name) # Skip loading extra bias for GPTQ models. if name.endswith(".bias") and name not in params_dict: @@ -372,7 +356,7 @@ def do_load_weights( ) break else: - for mapping in expert_params_mapping: + for mapping in self.expert_params_mapping: param_name, weight_name, expert_id, shard_id = mapping if weight_name not in name: continue @@ -409,7 +393,7 @@ def do_load_weights( # Skip loading norm if not last rank in pipeline parallelism if ".norm." in name and not self.pp_group.is_last_rank: continue - if fuse_qkv_a_proj and ( + if self.fuse_qkv_a_proj and ( "q_a_proj" in name or "kv_a_proj_with_mqa" in name ): cached_a_proj[name] = _clone_if_runai_streamed_tensor( @@ -479,13 +463,7 @@ def do_load_weights( if ( "k_scale" in name or "v_scale" in name ) and name not in params_dict: - # modelopt attn kv scale is named differently - for scale in ["k_scale", "v_scale"]: - if scale in name: - name = name.replace( - f"{scale[0]}_proj", "attn_mqa" - ) - break + name = self.custom_scale_remap(name) if name not in params_dict: # modelopt ckpt contains not needed weights for MTP module: # model.decoder.self_attn.attn_mqa.v_scale and diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py index 7d61f6f7c08c..7431dddeb937 100644 --- a/python/sglang/srt/models/deepseek_nextn.py +++ b/python/sglang/srt/models/deepseek_nextn.py @@ -283,6 +283,7 @@ def __init__( # if not set, model load will be broken in DeepseekV3ForCausalLM load_weights() self.pp_group = get_pp_group() self.determine_num_fused_shared_experts() + self.init_unified_loader_mappings() nextn_quant_config = self._resolve_nextn_quant_config(config, quant_config) self.model = DeepseekModelNextN( diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index bae009ce8fb3..a2fbebe9122a 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -1842,6 +1842,13 @@ def __init__( indexer_kwargs["skip_rope"] = skip_rope self.indexer = indexer_cls(**indexer_kwargs) + # A skip-topk layer only owns an indexer when it is also nextn; the + # others reuse the previous layer's topk and build none, so there are + # no weights to exempt. + if self.skip_topk and self.indexer is not None: + for p in self.indexer.parameters(): + p._skip_weight_check = True + self.kv_b_proj = ColumnParallelLinear( self.kv_lora_rank, self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), @@ -2938,6 +2945,8 @@ def __init__( self.lm_head = PPMissingLayer() self.logits_processor = LogitsProcessor(config) + self.init_unified_loader_mappings() + self._routed_experts_weights_of_layer = LazyValue( lambda: { layer_id: self.model.layers[layer_id].mlp.get_moe_weights() @@ -2950,6 +2959,22 @@ def __init__( q_lora_rank = config.q_lora_rank if hasattr(config, "q_lora_rank") else None get_attn_tp_context().init_context(q_lora_rank, is_deepseek_dsa(config)) + def mutate_weight_preload(self, name: str) -> str: + """DeepSeek V2: shared expert fusion.""" + if self.num_fused_shared_experts > 0 and "mlp.shared_experts" in name: + return name.replace( + "mlp.shared_experts", + f"mlp.experts.{self.config.n_routed_experts}", + ) + return name + + def custom_scale_remap(self, name: str) -> str: + """DeepSeek V2: k_proj -> attn_mqa when k_scale in name, v_proj -> attn_mqa when v_scale in name.""" + for scale in ["k_scale", "v_scale"]: + if scale in name: + return name.replace(f"{scale[0]}_proj", "attn_mqa") + return name + @property def routed_experts_weights_of_layer(self): return self._routed_experts_weights_of_layer.value @@ -2959,6 +2984,45 @@ def routed_experts_weights_of_layer(self): # overrides it. fused_shared_experts_architecture = "DeepseekV3ForCausalLM" + def init_unified_loader_mappings(self): + """Mappings the unified weight loading API reads off the model. + + Call after determine_num_fused_shared_experts(); subclasses that build + their own __init__ must call this or weight loading fails on the + missing attributes. + """ + self.fuse_qkv_a_proj = ( + hasattr(self.config, "q_lora_rank") and self.config.q_lora_rank is not None + ) + self.stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + # q_a_proj + kv_a_proj_with_mqa -> fused_qkv_a_proj_with_mqa + if self.fuse_qkv_a_proj: + self.stacked_params_mapping.extend( + [ + ("fused_qkv_a_proj_with_mqa", "q_a_proj", 0), + ("fused_qkv_a_proj_with_mqa", "kv_a_proj_with_mqa", 1), + ] + ) + self.expert_params_mapping = FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.n_routed_experts + self.num_fused_shared_experts, + ) + # Params for special naming rules in mixed-precision models, for example: + # model.layers.xx.mlp.experts.xx.w1.input_scale. For details, + # see https://huggingface.co/Barrrrry/DeepSeek-R1-W4AFP8/blob/main. + if is_wint4afp8_or_wint4a16_config(self.quant_config): + self.expert_params_mapping += ( + FusedMoE.make_expert_input_scale_params_mapping( + num_experts=self.config.n_routed_experts + ) + ) + @classmethod def shared_experts_fusion_disable_reason(cls, hf_config, quant_config): """Why this checkpoint cannot fuse its shared expert, or None. diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index af76c73f7225..622b77751a5d 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -706,6 +706,7 @@ def __init__( self.attn_sink = nn.Parameter(torch.empty(self.n_heads, dtype=torch.float32)) self._attn_sink_local: Optional[torch.Tensor] = None + if fuse: self.wqkv_a = ReplicatedLinear( self.hidden_size, @@ -867,6 +868,17 @@ def _local_attn_sink(self, kernel_num_heads: Optional[int] = None) -> torch.Tens self._attn_sink_local = sink return self._attn_sink_local[:kernel_num_heads] + def refresh_attn_sink_cache(self): + # RL weight updates rewrite attn_sink in place; the padded per-rank + # slice cached by _local_attn_sink must follow it. + if self._attn_sink_local is None: + return + rank = self.attn_tp_rank + num_heads = self.n_local_heads + self._attn_sink_local[:num_heads].copy_( + self.attn_sink[rank * num_heads : (rank + 1) * num_heads] + ) + @contextmanager def maybe_use_decode_attn_tp(self, forward_batch: ForwardBatch): ctx = get_cp_decode_attn_tp_ctx() @@ -3558,6 +3570,7 @@ def post_load_weights(self, is_nextn=False, weight_names=None): ): self_attn.indexer.compressor.apply_ape_hotfix() layer.refresh_mhc_norm_weight_cache() + self_attn.refresh_attn_sink_cache() @staticmethod def remap_weight_name_to_dpsk_hf_format( @@ -4036,7 +4049,9 @@ def auto_weight_loader(module): } if unloaded_params: logger.warning( - f"Some weights are not initialized from checkpoints: {unloaded_params}" + "Some weights are not initialized from checkpoints: " + f"count={len(unloaded_params)}. " + "Ignore this message for RL weight update." ) self.post_load_weights(is_nextn=is_nextn, weight_names=weight_names) diff --git a/python/sglang/srt/models/deepseek_v4_dspark.py b/python/sglang/srt/models/deepseek_v4_dspark.py index 4994a92547b6..567361267739 100644 --- a/python/sglang/srt/models/deepseek_v4_dspark.py +++ b/python/sglang/srt/models/deepseek_v4_dspark.py @@ -146,6 +146,18 @@ def kv_proj_only(self, x: torch.Tensor) -> torch.Tensor: kv, _ = self.wkv(x) return kv + def _attn_sink_pad_width(self) -> int: + # DSpark pads q to _PAD_NUM_HEADS only when the rank has fewer heads than + # that, so the sink follows the same width rather than padded_num_heads. + return max(self.n_local_heads, _PAD_NUM_HEADS) + + def _local_attn_sink(self) -> torch.Tensor: + if self.attn_tp_size == 1: + return self.attn_sink + if self._attn_sink_local is None: + self.refresh_attn_sink_cache() + return self._attn_sink_local + def _store_block_kv( self, *, @@ -978,6 +990,14 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> None: params_dict=params_dict, loaded_params=loaded_params ) + def post_load_weights(self): + # The per-rank attn_sink slice is cached on first use; an RL weight update + # rewrites attn_sink in place, so re-slice it here or decode keeps serving + # the pre-update sink. + for module in self.modules(): + if isinstance(module, DSparkAttention): + module.refresh_attn_sink_cache() + def _assert_confidence_head_loaded( self, *, params_dict: dict, loaded_params: set ) -> None: diff --git a/python/sglang/srt/models/deepseek_v4_nextn.py b/python/sglang/srt/models/deepseek_v4_nextn.py index cecdf4583d9b..0e657340168e 100644 --- a/python/sglang/srt/models/deepseek_v4_nextn.py +++ b/python/sglang/srt/models/deepseek_v4_nextn.py @@ -235,6 +235,7 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): def post_load_weights(self, is_nextn=False, weight_names=None): super().post_load_weights(is_nextn=True, weight_names=weight_names) + self.model.decoder.self_attn.refresh_attn_sink_cache() EntryClass = [DeepseekV4ForCausalLMNextN] diff --git a/python/sglang/srt/models/glm4.py b/python/sglang/srt/models/glm4.py index 00a5057d0e3e..91f8fceaea8d 100644 --- a/python/sglang/srt/models/glm4.py +++ b/python/sglang/srt/models/glm4.py @@ -458,6 +458,17 @@ def __init__( self.logits_processor = LogitsProcessor(config) self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True) + + # Stacked params mapping for unified weight loading API + self.stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + # For EAGLE3 support self.capture_aux_hidden_states = False @@ -552,14 +563,7 @@ def end_layer(self): return self.model.end_layer def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - (".qkv_proj", ".q_proj", "q"), - (".qkv_proj", ".k_proj", "k"), - (".qkv_proj", ".v_proj", "v"), - (".gate_up_proj", ".up_proj", 1), - (".gate_up_proj", ".gate_proj", 0), - ] + stacked_params_mapping = self.stacked_params_mapping params_dict = dict(self.named_parameters()) for name, loaded_weight in weights: diff --git a/python/sglang/srt/models/glm4_moe.py b/python/sglang/srt/models/glm4_moe.py index fbda6fdc027d..29c41f8528b9 100644 --- a/python/sglang/srt/models/glm4_moe.py +++ b/python/sglang/srt/models/glm4_moe.py @@ -15,7 +15,6 @@ """Inference-only GLM-4.5, GLM-4.6 and GLM-4.7 model compatible with HuggingFace weights""" import logging -import re from typing import Any, Dict, Iterable, List, Optional, Tuple, Union import torch @@ -1165,9 +1164,34 @@ def __init__( ) self.logits_processor = LogitsProcessor(config) + # Stacked params mapping for unified weight loading API + self.stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + self.expert_params_mapping = FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.n_routed_experts + self.num_fused_shared_experts, + ) + # For EAGLE3 support self.capture_aux_hidden_states = False + def mutate_weight_preload(self, name: str) -> str: + """GLM4-MoE: shared expert fusion.""" + if self.num_fused_shared_experts > 0 and "mlp.shared_experts" in name: + return name.replace( + "mlp.shared_experts", + f"mlp.experts.{self.config.n_routed_experts}", + ) + return name + @classmethod def shared_experts_fusion_disable_reason(cls, hf_config, quant_config): """Why this checkpoint cannot fuse its shared expert, or None. Asked by @@ -1248,43 +1272,8 @@ def load_weights( else: raise ValueError("num_nextn_predict_layers is not in the config") - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("qkv_proj", "q_proj", "q"), - ("qkv_proj", "k_proj", "k"), - ("qkv_proj", "v_proj", "v"), - ("gate_up_proj", "gate_proj", 0), - ("gate_up_proj", "up_proj", 1), - ] - - if self.num_fused_shared_experts > 0: - assert self.num_fused_shared_experts == 1 - - def iter_weights_with_fused_shared_experts( - weights: Iterable[Tuple[str, torch.Tensor]], - ) -> Iterable[Tuple[str, torch.Tensor]]: - - pattern = re.compile( - r"^model\.layers\.(\d+)\.mlp\.shared_experts\.(.+)$" - ) - for name, weight in weights: - match = pattern.match(name) - if match: - layer_id = int(match.group(1)) - suffix = match.group(2) - name = f"model.layers.{layer_id}.mlp.experts.{self.config.n_routed_experts}.{suffix}" - yield name, weight - - weights = iter_weights_with_fused_shared_experts(weights) - - # Params for weights, fp8 weight scales, fp8 activation scales - # (param_name, weight_name, expert_id, shard_id) - expert_params_mapping = FusedMoE.make_expert_params_mapping( - ckpt_gate_proj_name="gate_proj", - ckpt_down_proj_name="down_proj", - ckpt_up_proj_name="up_proj", - num_experts=self.config.n_routed_experts + self.num_fused_shared_experts, - ) + stacked_params_mapping = self.stacked_params_mapping + expert_params_mapping = self.expert_params_mapping if is_nextn: nextn_layer_prefix = f"model.layers.{nextn_layer_id}" @@ -1305,6 +1294,8 @@ def iter_weights_with_fused_shared_experts( for name, loaded_weight in weights: weight_names.append(name) + name = self.mutate_weight_preload(name) + if not is_nextn: if hasattr(self.config, "num_nextn_predict_layers"): num_nextn_layers = self.config.num_nextn_predict_layers diff --git a/python/sglang/srt/models/glm4_moe_lite.py b/python/sglang/srt/models/glm4_moe_lite.py index 4d74c41dc221..98cf4b4dd925 100644 --- a/python/sglang/srt/models/glm4_moe_lite.py +++ b/python/sglang/srt/models/glm4_moe_lite.py @@ -246,6 +246,11 @@ def __init__( if config.n_shared_experts is not None and self.num_fused_shared_experts == 0: intermediate_size = config.moe_intermediate_size * config.n_shared_experts # disable tp for shared experts when enable deepep moe, or with fp4 allgather + _shared_expert_use_tp1 = ( + get_moe_a2a_backend().is_deepep() + or get_moe_a2a_backend().is_mooncake() + or should_use_flashinfer_cutlass_moe_fp4_allgather() + ) self.shared_experts = Glm4MoeLiteMLP( hidden_size=config.hidden_size, intermediate_size=intermediate_size, @@ -253,14 +258,9 @@ def __init__( quant_config=quant_config, reduce_results=False, prefix=add_prefix("shared_experts", prefix), - **( - dict(tp_rank=0, tp_size=1) - if get_moe_a2a_backend().is_deepep() - or get_moe_a2a_backend().is_mooncake() - or should_use_flashinfer_cutlass_moe_fp4_allgather() - else {} - ), + **(dict(tp_rank=0, tp_size=1) if _shared_expert_use_tp1 else {}), ) + self._shared_expert_tp1 = _shared_expert_use_tp1 is_packed_weight = hasattr( self.shared_experts.gate_up_proj.quant_method, "quant_config" ) @@ -915,6 +915,33 @@ def __init__( ) self.capture_aux_hidden_states = False + # Weight loading mappings for ParameterMapper compatibility + self.stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + # Add A-proj fusion mapping when q_lora_rank is enabled (MLA) + self.fuse_qkv_a_proj = hasattr(config, "q_lora_rank") and ( + config.q_lora_rank is not None + ) + if self.fuse_qkv_a_proj: + self.stacked_params_mapping.extend( + [ + ("fused_qkv_a_proj_with_mqa", "q_a_proj", 0), + ("fused_qkv_a_proj_with_mqa", "kv_a_proj_with_mqa", 1), + ] + ) + self.expert_params_mapping = FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=config.n_routed_experts + self.num_fused_shared_experts, + ) + @property def routed_experts_weights_of_layer(self): return self._routed_experts_weights_of_layer.value @@ -1027,6 +1054,22 @@ def set_dflash_layers_to_capture(self, layer_ids: List[int]): self.capture_aux_hidden_states = True self.model.layers_to_capture = [val + 1 for val in layer_ids] + def mutate_weight_preload(self, name: str) -> str: + """GLM4-MoE-Lite: shared expert fusion.""" + if self.num_fused_shared_experts > 0 and "mlp.shared_experts" in name: + return name.replace( + "mlp.shared_experts", + f"mlp.experts.{self.config.n_routed_experts}", + ) + return name + + def custom_scale_remap(self, name: str) -> str: + """GLM4-MoE-Lite: k_proj/v_proj -> attn_mqa for MLA kv scale.""" + for s in ["k_scale", "v_scale"]: + if s in name: + return name.replace(f"{s[0]}_proj", "attn_mqa") + return name + def load_weights( self, weights: Iterable[Tuple[str, torch.Tensor]], diff --git a/python/sglang/srt/models/glm4_moe_nextn.py b/python/sglang/srt/models/glm4_moe_nextn.py index f29928efb359..eedd8e536d29 100644 --- a/python/sglang/srt/models/glm4_moe_nextn.py +++ b/python/sglang/srt/models/glm4_moe_nextn.py @@ -25,6 +25,7 @@ from sglang.srt.layers.dp_attention import is_dp_attention_enabled from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.vocab_parallel_embedding import ( ParallelLMHead, @@ -32,7 +33,7 @@ ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.models.glm4_moe import Glm4MoeDecoderLayer, Glm4MoeForCausalLM -from sglang.srt.runtime_context import get_parallel, get_spec +from sglang.srt.runtime_context import get_exec, get_parallel, get_spec from sglang.srt.utils import add_prefix, is_npu logger = logging.getLogger(__name__) @@ -148,6 +149,26 @@ def __init__( ) self.logits_processor = LogitsProcessor(config) + self.num_fused_shared_experts = ( + 0 if get_exec().moe.disable_shared_experts_fusion else 1 + ) + + # Weight loading mappings (must match parent for load_weights to work) + self.stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + self.expert_params_mapping = FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.n_routed_experts + self.num_fused_shared_experts, + ) + @torch.no_grad() def forward( self, diff --git a/python/sglang/srt/models/inkling.py b/python/sglang/srt/models/inkling.py index ebad2504ffc1..f50260557311 100644 --- a/python/sglang/srt/models/inkling.py +++ b/python/sglang/srt/models/inkling.py @@ -1114,6 +1114,20 @@ def base_layer(module: torch.nn.Module) -> torch.nn.Module: return hidden_size, config.vocab_size raise NotImplementedError(f"get_hidden_dim not implemented for {module_name}") + @staticmethod + def get_lora_target_module_name(name: str) -> str: + parent, _, leaf = name.rpartition(".") + if parent.endswith(".attn"): + mappings = ATTENTION_PARAMS_MAPPING + elif parent.endswith((".mlp.experts", ".mlp.shared_experts")): + mappings = STACKED_DENSE_PARAMS_MAPPING + else: + return name + for module_name, weight_name, _ in mappings: + if leaf == weight_name: + return module_name + return name + def get_stacked_multiply(self, module_name: str) -> int: if module_name == "qkvr": return 4 diff --git a/python/sglang/srt/models/kimi_k25.py b/python/sglang/srt/models/kimi_k25.py index 54208e7d7132..3ed997452735 100644 --- a/python/sglang/srt/models/kimi_k25.py +++ b/python/sglang/srt/models/kimi_k25.py @@ -247,6 +247,7 @@ def __init__( .unsqueeze(1), persistent=False, ) + self.time_weight._skip_weight_check = True self.reset_parameters() diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py index b0ede48c8826..53bb54276d71 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/python/sglang/srt/models/kimi_k3.py @@ -1196,11 +1196,17 @@ def _forward_shared(self, gate_up, shared_output): shared.down_proj.weight, torch.Tensor ) assert shared is not None + shared_act = shared.act_fn(gate_up) _k3_bf16_gemm( - shared.act_fn(gate_up), + shared_act, shared.down_proj.weight, out=shared_output, ) + down_proj = shared.down_proj + if getattr(down_proj, "lora_active", False): + # the base GEMM bypassed the wrapper; LoRA B must land in the same buffer the tail add reads + lora_output = down_proj.apply_lora(shared_output, shared_act) + assert lora_output.data_ptr() == shared_output.data_ptr() def _get_fused_norm_params(self) -> tuple[torch.Tensor, float]: norm = self.routed_expert_norm @@ -1236,6 +1242,10 @@ def _forward_fused( gate_up, router_logits, routed_input = torch.split( fused, self._front_sizes, dim=-1 ) + gate_up_proj = self.shared_experts.gate_up_proj + if getattr(gate_up_proj, "lora_active", False): + # the merged front GEMM read the base shared weight through the wrapper; add its delta + gate_up = gate_up_proj.apply_lora(gate_up.contiguous(), hidden_states) if num_tokens > 1 and _is_hip and not _aiter_k3_opt: router_logits = router_logits.contiguous() if self._moe_front_needs_dense_bf16: @@ -1756,7 +1766,8 @@ def _merge_bfa_weights(self) -> None: if any(getattr(mod, "weight", None) is None for mod in mods): return self._bfa_w, sizes = _merge_weights_as_views(mods, pad_rows_to=8) - self._bfa_f_b_w = self.f_b_proj.weight + # .data: a second Parameter registration hides f_b_proj.weight from named_parameters() + self._bfa_f_b_w = self.f_b_proj.weight.data self._bfa_fa_size, self._bfa_b_size = sizes def _merge_kda_inproj_weights_hip(self) -> bool: @@ -2211,9 +2222,10 @@ def _symm_o_proj_forward(x, *args, **kwargs): if self._gate_alt_stream is not None else 0 ) - _orig_o_proj_forward = self.o_proj.forward + o_proj = self.o_proj + _orig_o_proj_forward = o_proj.forward - def _gated_o_proj_forward(x, *args, **kwargs): + def _apply_output_gate(x): gate_input = self._gate_hidden_states self._gate_hidden_states = None precomputed = self._gate_precomputed @@ -2237,9 +2249,14 @@ def _gated_o_proj_forward(x, *args, **kwargs): x = mla_output_gate.kimi_k3_mla_output_gate(x, gate) else: x = x * torch.sigmoid(gate) - return _orig_o_proj_forward(x, *args, **kwargs) + return x + + def _gated_o_proj_forward(x, *args, **kwargs): + return _orig_o_proj_forward(_apply_output_gate(x), *args, **kwargs) - self.o_proj.forward = _gated_o_proj_forward + # RowParallelLinearWithLoRA bypasses forward and applies this transform itself + o_proj.lora_input_transform = _apply_output_gate + o_proj.forward = _gated_o_proj_forward @staticmethod def _split_kv_b_weight_loader(param, loaded_weight) -> None: @@ -3383,13 +3400,15 @@ def _warm_cw(proj, norm): precompile_k3_recompute_w_u_kernel, ) - o_proj_weight = getattr(layer.self_attn.o_proj, "weight", None) + # under LoRA o_proj is a BaseLayerWithLoRA; params_dtype lives on the base layer + o_proj = layer.self_attn.o_proj + o_proj = getattr(o_proj, "base_layer", o_proj) + o_proj_weight = getattr(o_proj, "weight", None) if o_proj_weight is None: - o_proj_weight = layer.self_attn.o_proj.qweight + o_proj_weight = o_proj.qweight if precompile_k3_recompute_w_u_kernel( num_heads=layer.self_attn.local_num_heads, - dtype=getattr(layer.self_attn.o_proj, "params_dtype", None) - or o_proj_weight.dtype, + dtype=getattr(o_proj, "params_dtype", None) or o_proj_weight.dtype, device=layer.self_attn.dt_bias.device, ): rank0_log("Precompiled the Kimi-K3 KDA prefill kernel.") @@ -3438,6 +3457,15 @@ class KimiK3ForConditionalGeneration(nn.Module): }, ) + @staticmethod + def get_lora_target_module_name(name: str) -> str: + parent, _, leaf = name.rpartition(".") + if ".experts" in parent: + return {"w1": "gate_proj", "w2": "down_proj", "w3": "up_proj"}.get( + leaf, name + ) + return name + def __init__( self, config: KimiK3Config, diff --git a/python/sglang/srt/models/kimi_k3_vl.py b/python/sglang/srt/models/kimi_k3_vl.py index c42302799464..8a0fb63a10ee 100644 --- a/python/sglang/srt/models/kimi_k3_vl.py +++ b/python/sglang/srt/models/kimi_k3_vl.py @@ -228,6 +228,7 @@ def __init__( .unsqueeze(1), persistent=False, ) + self.time_weight._skip_weight_check = True def position_embeddings( self, diff --git a/python/sglang/srt/models/llama.py b/python/sglang/srt/models/llama.py index ad6aeb8bc0f8..215741e64a12 100644 --- a/python/sglang/srt/models/llama.py +++ b/python/sglang/srt/models/llama.py @@ -547,9 +547,21 @@ def __init__( (".gate_up_proj", ".gate_proj", 0), (".gate_up_proj", ".up_proj", 1), ] + # Llama-specific scale remapping patterns (suffix, pattern, replacement) + self._llama_scale_remap_patterns = [ + (".activation_scale", ".activation_scale", ".input_scale"), + (".weight_scale_inv", ".weight_scale_inv", ".weight_scale"), + ] self.capture_aux_hidden_states = False + def custom_scale_remap(self, name: str) -> str: + """Llama: activation_scale->input_scale, weight_scale_inv->weight_scale.""" + for suffix, pattern, replacement in self._llama_scale_remap_patterns: + if name.endswith(suffix) and pattern in name: + return name.replace(pattern, replacement) + return name + def _init_model( self, config: LlamaConfig, @@ -667,23 +679,11 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): return self._legacy_load_weights(weights) def _legacy_load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - (".qkv_proj", ".q_proj", "q"), - (".qkv_proj", ".k_proj", "k"), - (".qkv_proj", ".v_proj", "v"), - (".gate_up_proj", ".gate_proj", 0), - (".gate_up_proj", ".up_proj", 1), - ] params_dict = dict(self.named_parameters()) for name, loaded_weight in weights: - if name.endswith(".activation_scale"): - name = name.replace(".activation_scale", ".input_scale") - if name.endswith(".weight_scale_inv"): - name = name.replace(".weight_scale_inv", ".weight_scale") - + name = self.custom_scale_remap(name) layer_id = get_layer_id(name) if ( layer_id is not None @@ -710,7 +710,7 @@ def _legacy_load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): if name is None: continue - for param_name, weight_name, shard_id in stacked_params_mapping: + for param_name, weight_name, shard_id in self.stacked_params_mapping: if weight_name not in name: continue name = name.replace(weight_name, param_name) diff --git a/python/sglang/srt/models/minicpmv.py b/python/sglang/srt/models/minicpmv.py index b8fd89a27a29..c675d558aa52 100644 --- a/python/sglang/srt/models/minicpmv.py +++ b/python/sglang/srt/models/minicpmv.py @@ -303,6 +303,7 @@ def _set_2d_pos_cache( ) pos_embed = torch.from_numpy(pos_embed_arr).float().to(device) self.register_buffer("pos_embed", pos_embed, persistent=False) + self.pos_embed._skip_weight_check = True def _adjust_pos_cache( self, tgt_sizes: torch.Tensor, device: torch.types.Device @@ -425,6 +426,7 @@ def _set_2d_pos_cache( ) pos_embed = torch.from_numpy(pos_embed_arr).float().to(device) self.register_buffer("pos_embed", pos_embed, persistent=False) + self.pos_embed._skip_weight_check = True def _adjust_pos_cache( self, tgt_sizes: torch.Tensor, device: torch.types.Device @@ -454,6 +456,7 @@ def _set_temporal_pos_cache( .to(device) ) self.register_buffer("temporal_pos_embed", pos_embed, persistent=False) + self.temporal_pos_embed._skip_weight_check = True def _adjust_temporal_pos_cache( self, max_temporal_size: int, device: torch.types.Device = "cpu" diff --git a/python/sglang/srt/models/nemotron_h.py b/python/sglang/srt/models/nemotron_h.py index fd85df6bc178..ee5a4f74f34a 100644 --- a/python/sglang/srt/models/nemotron_h.py +++ b/python/sglang/srt/models/nemotron_h.py @@ -227,6 +227,7 @@ def __init__( ) self.topk = TopK( top_k=config.num_experts_per_tok, + layer_id=layer_idx, use_grouped_topk=True, topk_group=config.topk_group, num_expert_group=config.n_group, diff --git a/python/sglang/srt/models/qwen2.py b/python/sglang/srt/models/qwen2.py index 49eaecf8b3fd..2198948b1b18 100644 --- a/python/sglang/srt/models/qwen2.py +++ b/python/sglang/srt/models/qwen2.py @@ -50,7 +50,12 @@ kv_cache_scales_loader, ) from sglang.srt.platforms import current_platform -from sglang.srt.runtime_context import get_exec, get_parallel +from sglang.srt.runtime_context import get_parallel +from sglang.srt.true_on_policy import ( + get_on_policy_rms_norm_kwargs, + is_true_on_policy_enabled, + should_force_bfloat16_dense_tensor_math, +) from sglang.srt.utils import add_prefix, make_layers from sglang.srt.utils.hf_transformers_utils import get_rope_config @@ -95,9 +100,11 @@ def forward( x: torch.Tensor, forward_batch: ForwardBatch = None, ) -> torch.Tensor: - if get_exec().deterministic.rl_on_policy_target is not None: - x = x.bfloat16() - + if ( + should_force_bfloat16_dense_tensor_math() + or x.dtype != self.gate_up_proj.weight.dtype + ): + x = x.to(self.gate_up_proj.weight.dtype) gate_up, _ = self.gate_up_proj(x) x = self.act_fn(gate_up) x, _ = self.down_proj(x, forward_batch=forward_batch) @@ -333,11 +340,7 @@ def __init__( quant_config=quant_config, use_attn_tp_group=is_dp_attention_enabled(), prefix=add_prefix("embed_tokens", prefix), - params_dtype=( - torch.float32 - if get_exec().deterministic.rl_on_policy_target is not None - else None - ), + params_dtype=(torch.float32 if is_true_on_policy_enabled() else None), ) else: self.embed_tokens = PPMissingLayer() @@ -364,16 +367,7 @@ def __init__( prefix=add_prefix("layers", prefix), ) if self.pp_group.is_last_rank: - norm_kwargs = ( - dict( - weight_dtype=torch.float32, - cast_x_before_out_mul=True, - override_orig_dtype=torch.float32, - fp32_residual=True, - ) - if get_exec().deterministic.rl_on_policy_target is not None - else {} - ) + norm_kwargs = get_on_policy_rms_norm_kwargs() self.norm = RMSNorm( config.hidden_size, eps=config.rms_norm_eps, **norm_kwargs ) @@ -519,6 +513,17 @@ def __init__( self.logits_processor = LogitsProcessor(config) self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True) + + # Stacked params mapping for unified weight loading API + self.stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + # For EAGLE3 support self.capture_aux_hidden_states = False @@ -620,14 +625,7 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): return self._legacy_load_weights(weights) def _legacy_load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("qkv_proj", "q_proj", "q"), - ("qkv_proj", "k_proj", "k"), - ("qkv_proj", "v_proj", "v"), - ("gate_up_proj", "gate_proj", 0), - ("gate_up_proj", "up_proj", 1), - ] + stacked_params_mapping = self.stacked_params_mapping params_dict = dict(self.named_parameters()) for name, loaded_weight in weights: diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 83aad13ebbef..eb9bf7e087a7 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -97,6 +97,7 @@ get_parallel, get_stream, ) +from sglang.srt.true_on_policy import get_on_policy_rms_norm_kwargs from sglang.srt.utils import ( add_prefix, cpu_has_amx_support, @@ -1092,7 +1093,10 @@ def __init__( prefix=add_prefix("layers", prefix), ) if self.pp_group.is_last_rank: - self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + norm_kwargs = get_on_policy_rms_norm_kwargs(fp32_residual=False) + self.norm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps, **norm_kwargs + ) else: self.norm = PPMissingLayer(return_tuple=True) diff --git a/python/sglang/srt/models/qwen3.py b/python/sglang/srt/models/qwen3.py index 9fbd046709e6..ba928b3a52bf 100644 --- a/python/sglang/srt/models/qwen3.py +++ b/python/sglang/srt/models/qwen3.py @@ -33,8 +33,16 @@ from sglang.srt.models.qwen2 import Qwen2MLP as Qwen3MLP from sglang.srt.models.qwen2 import Qwen2Model from sglang.srt.models.utils import apply_qk_norm -from sglang.srt.runtime_context import get_exec, get_parallel, get_stream +from sglang.srt.runtime_context import ( + get_parallel, + get_stream, +) +from sglang.srt.true_on_policy import ( + should_disable_fused_qk_norm_mrope, + should_force_bfloat16_dense_tensor_math, +) from sglang.srt.utils import add_prefix, get_bool_env_var, is_cuda, is_hip, is_npu +from sglang.srt.utils.hf_transformers_utils import get_rope_config Qwen3Config = None @@ -106,16 +114,16 @@ def __init__( self.max_position_embeddings = max_position_embeddings self.tp_rank = get_parallel().tp_rank - norm_kwargs = ( - dict( - weight_dtype=torch.float32, - cast_x_before_out_mul=True, - ) - if get_exec().deterministic.rl_on_policy_target is not None - else {} + self.q_norm = RMSNorm( + self.head_dim, + eps=rms_norm_eps, + true_on_policy_weight_dtype=torch.float32, + ) + self.k_norm = RMSNorm( + self.head_dim, + eps=rms_norm_eps, + true_on_policy_weight_dtype=torch.float32, ) - self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps, **norm_kwargs) - self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps, **norm_kwargs) self.qkv_proj = QKVParallelLinear( hidden_size, @@ -271,14 +279,20 @@ def forward( hidden_states: torch.Tensor, forward_batch: ForwardBatch, ) -> torch.Tensor: - if get_exec().deterministic.rl_on_policy_target is not None: - hidden_states = hidden_states.bfloat16() + if ( + should_force_bfloat16_dense_tensor_math() + or hidden_states.dtype != self.qkv_proj.weight.dtype + ): + # True-on-policy RMSNorm can produce fp32 activations while dense + # projections remain bf16, including during cuda-graph capture when + # the global on-policy flag is temporarily cleared. + hidden_states = hidden_states.to(self.qkv_proj.weight.dtype) save_kv_cache = True use_aiter_fused = ( self.use_fused_qk_norm_mrope and forward_batch.forward_mode.is_decode() - and get_exec().deterministic.rl_on_policy_target is None + and not should_disable_fused_qk_norm_mrope() ) if use_aiter_fused: @@ -298,9 +312,9 @@ def forward( forward_batch=forward_batch, ) - if get_exec().deterministic.rl_on_policy_target is not None: - q = q.to(torch.bfloat16) - k = k.to(torch.bfloat16) + if should_force_bfloat16_dense_tensor_math() or q.dtype != v.dtype: + q = q.to(v.dtype) + k = k.to(v.dtype) attn_output = self.attn(q, k, v, forward_batch, save_kv_cache=save_kv_cache) output, _ = self.o_proj(attn_output) @@ -319,16 +333,7 @@ def __init__( ) -> None: super().__init__() self.hidden_size = config.hidden_size - if ( - hasattr(config, "rope_parameters") - and config.rope_parameters - and "rope_theta" in config.rope_parameters - ): - rope_theta = config.rope_parameters["rope_theta"] - rope_scaling = config.rope_parameters - else: - rope_theta = getattr(config, "rope_theta", 1000000) - rope_scaling = getattr(config, "rope_scaling", None) + rope_theta, rope_scaling = get_rope_config(config) max_position_embeddings = getattr(config, "max_position_embeddings", 32768) head_dim = getattr(config, "head_dim", None) self.self_attn = Qwen3Attention( @@ -355,21 +360,19 @@ def __init__( prefix=add_prefix("mlp", prefix), ) - norm_kwargs = ( - dict( - weight_dtype=torch.float32, - cast_x_before_out_mul=True, - override_orig_dtype=torch.float32, - fp32_residual=True, - ) - if get_exec().deterministic.rl_on_policy_target is not None - else {} - ) self.input_layernorm = RMSNorm( - config.hidden_size, eps=config.rms_norm_eps, **norm_kwargs + config.hidden_size, + eps=config.rms_norm_eps, + true_on_policy_weight_dtype=torch.float32, + true_on_policy_override_orig_dtype=torch.float32, + true_on_policy_fp32_residual=True, ) self.post_attention_layernorm = RMSNorm( - config.hidden_size, eps=config.rms_norm_eps, **norm_kwargs + config.hidden_size, + eps=config.rms_norm_eps, + true_on_policy_weight_dtype=torch.float32, + true_on_policy_override_orig_dtype=torch.float32, + true_on_policy_fp32_residual=True, ) self.layer_scatter_modes = LayerScatterModes.init_new( @@ -483,6 +486,16 @@ def __init__( config, quant_config=quant_config, prefix=add_prefix("model", prefix) ) + # Stacked params mapping for unified weight loading API + self.stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + # handle the lm head on different pp ranks if self.pp_group.is_last_rank: if self.pp_group.world_size == 1 and config.tie_word_embeddings: @@ -594,15 +607,7 @@ def end_layer(self): return self.model.end_layer def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("qkv_proj", "q_proj", "q"), - ("qkv_proj", "k_proj", "k"), - ("qkv_proj", "v_proj", "v"), - ("gate_up_proj", "gate_proj", 0), - ("gate_up_proj", "up_proj", 1), - ] - + stacked_params_mapping = self.stacked_params_mapping params_dict = dict(self.named_parameters()) for name, loaded_weight in weights: if not name.startswith("model.") and ( diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index d3691ea5ec5f..293ed6f1e596 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -464,7 +464,6 @@ def __init__( eps=self.layer_norm_epsilon, group_size=None, norm_before_gate=True, - device=torch.get_device_module().current_device(), dtype=torch.get_default_dtype(), **( {"activation": self.output_gate_type} @@ -1564,6 +1563,7 @@ class Qwen3_5ForCausalLM(nn.Module): "o_proj", "out_proj", "in_proj_qkvz", + "in_proj_ba", "gate_up_proj", "down_proj", "lm_head", @@ -1589,6 +1589,9 @@ def get_hidden_dim(self, module_name: str, layer_idx: int): key_dim = config.linear_key_head_dim * config.linear_num_key_heads value_dim = config.linear_value_head_dim * config.linear_num_value_heads return config.hidden_size, key_dim * 2 + value_dim * 2 + elif module_name == "in_proj_ba": + # b + a projections: one scalar per linear-attention value head each + return config.hidden_size, config.linear_num_value_heads * 2 elif module_name == "gate_up_proj": # MoE: shared expert uses shared_expert_intermediate_size # Dense: regular MLP uses intermediate_size diff --git a/python/sglang/srt/models/qwen3_moe.py b/python/sglang/srt/models/qwen3_moe.py index c0e93c371311..4481fdd52676 100644 --- a/python/sglang/srt/models/qwen3_moe.py +++ b/python/sglang/srt/models/qwen3_moe.py @@ -23,6 +23,7 @@ from typing import Any, Dict, Iterable, List, Optional, Tuple, TypeVar import torch +import torch.nn.functional as F from torch import nn from transformers import PretrainedConfig @@ -48,7 +49,7 @@ ) from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE -from sglang.srt.layers.moe.topk import TopK +from sglang.srt.layers.moe.topk import StandardTopKOutput, TopK from sglang.srt.layers.moe.utils import ( RoutingMethodType, filter_moe_weight_param_global_expert, @@ -68,7 +69,17 @@ create_fused_set_kv_buffer_arg, enable_fused_set_kv_buffer, ) -from sglang.srt.runtime_context import get_exec, get_forward, get_parallel, get_stream +from sglang.srt.runtime_context import ( + get_exec, + get_forward, + get_parallel, + get_stream, +) +from sglang.srt.true_on_policy import ( + get_on_policy_rms_norm_kwargs, + is_true_on_policy_enabled, + should_disable_fused_qk_norm_mrope, +) from sglang.srt.utils import ( LazyValue, add_prefix, @@ -330,7 +341,20 @@ def forward_normal( if hidden_states.shape[0] > 0: # router_logits: (num_tokens, n_experts) router_logits, _ = self.gate(hidden_states) - topk_output = self.topk(hidden_states, router_logits) + if is_true_on_policy_enabled(): + routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float) + routing_weights, selected_experts = torch.topk( + routing_weights, self.top_k, dim=-1 + ) + routing_weights /= routing_weights.sum(dim=-1, keepdim=True) + routing_weights = routing_weights.to(hidden_states.dtype) + topk_output = StandardTopKOutput( + topk_weights=routing_weights, + topk_ids=selected_experts, + router_logits=router_logits, + ) + else: + topk_output = self.topk(hidden_states, router_logits) else: topk_output = self.topk.empty_topk_output(hidden_states.device) final_hidden_states = self.experts(hidden_states, topk_output) @@ -520,7 +544,7 @@ def __init__( ) self.compatible_with_fused_kv_buffer = ( False if isinstance(self.rotary_emb, MRotaryEmbedding) else True - ) + ) and not is_true_on_policy_enabled() self.compatible_with_fused_qk_norm_rope = not isinstance( self.rotary_emb, MRotaryEmbedding ) and self.head_dim in (64, 128, 256) @@ -535,6 +559,7 @@ def __init__( torch.bfloat16, _yarn_factor != 1.0, ) + and not should_disable_fused_qk_norm_mrope() ) self.use_fused_qk_norm_rope_cpu = ( _is_cpu @@ -553,8 +578,9 @@ def __init__( prefix=add_prefix("attn", prefix), ) - self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) - self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) + norm_kwargs = get_on_policy_rms_norm_kwargs(fp32_residual=False) + self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps, **norm_kwargs) + self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps, **norm_kwargs) self.alt_stream = alt_stream def op_prepare(self, state): @@ -819,9 +845,12 @@ def __init__( quant_config=quant_config, prefix=add_prefix("mlp", prefix), ) - self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + norm_kwargs = get_on_policy_rms_norm_kwargs(fp32_residual=False) + self.input_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps, **norm_kwargs + ) self.post_attention_layernorm = RMSNorm( - config.hidden_size, eps=config.rms_norm_eps + config.hidden_size, eps=config.rms_norm_eps, **norm_kwargs ) self.layer_communicator = LayerCommunicator( @@ -998,6 +1027,23 @@ def __init__( use_attn_tp_group=get_parallel().enable_dp_lm_head, ) self.logits_processor = LogitsProcessor(config) + + # Stacked params mapping for unified weight loading API + self.stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + self.expert_params_mapping = FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.num_experts, + ) + self.capture_aux_hidden_states = False # IPC loading bypasses load_weights(), so initialize the EPLB descriptor here. self.routed_experts_weights_of_layer = LazyValue( @@ -1135,21 +1181,8 @@ def set_dflash_layers_to_capture(self, layer_ids: List[int]): def load_weights( self, weights: Iterable[Tuple[str, torch.Tensor]], is_mtp: bool = False ): - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("qkv_proj", "q_proj", "q"), - ("qkv_proj", "k_proj", "k"), - ("qkv_proj", "v_proj", "v"), - ("gate_up_proj", "gate_proj", 0), - ("gate_up_proj", "up_proj", 1), - ] - - expert_params_mapping = FusedMoE.make_expert_params_mapping( - ckpt_gate_proj_name="gate_proj", - ckpt_down_proj_name="down_proj", - ckpt_up_proj_name="up_proj", - num_experts=self.config.num_experts, - ) + stacked_params_mapping = self.stacked_params_mapping + expert_params_mapping = self.expert_params_mapping # Pre-define `params_dict` to avoid repeated expensive traversal of model parameters. params_dict = dict(self.named_parameters()) diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py index db6ba842754a..f0d90606a918 100644 --- a/python/sglang/srt/models/qwen3_vl.py +++ b/python/sglang/srt/models/qwen3_vl.py @@ -80,6 +80,7 @@ ) from sglang.srt.multimodal.vit_cuda_graph_runner import ViTCudaGraphRunner from sglang.srt.runtime_context import get_exec, get_mm, get_parallel +from sglang.srt.true_on_policy import is_true_on_policy_enabled from sglang.srt.utils import ( add_prefix, cpu_has_amx_support, @@ -490,75 +491,73 @@ def rot_pos_emb( return cos_combined, sin_combined - def _get_interpolation_indices(self, dim_size: int) -> torch.Tensor: - """ - Compute continuous interpolation indices for a single dimension. + def fast_pos_embed_interpolate(self, grid_thw): + grid_ts, grid_hs, grid_ws = grid_thw[:, 0], grid_thw[:, 1], grid_thw[:, 2] + num_grid_per_side = int(self.num_position_embeddings**0.5) + device = self.pos_embed.weight.device - Returns continuous indices. - """ - if self.align_corners: - indices = np.linspace( - 0, self.num_grid_per_side - 1, dim_size, dtype=np.float32 - ) - else: - indices = (np.arange(dim_size, dtype=np.float32) + 0.5) * ( - self.num_grid_per_side / dim_size - ) - 0.5 - indices = np.clip(indices, 0, self.num_grid_per_side - 1) - return indices + idx_list = [[] for _ in range(4)] + weight_list = [[] for _ in range(4)] - def _calculate_indices_and_weights(self, h_idxs, w_idxs): - """ - Compute bilinear interpolation indices and weights. + for t, h, w in zip(grid_ts, grid_hs, grid_ws): + h_idxs = torch.linspace(0, num_grid_per_side - 1, h) + w_idxs = torch.linspace(0, num_grid_per_side - 1, w) - Returns tuple of (indices, weights), each as 4 numpy arrays for the 4 corner points. - """ - h_f = np.floor(h_idxs).astype(np.int64) - h_c = np.clip(h_f + 1, 0, self.num_grid_per_side - 1) - dh = h_idxs - h_f + h_idxs_floor = h_idxs.int() + w_idxs_floor = w_idxs.int() + h_idxs_ceil = (h_idxs.int() + 1).clip(max=num_grid_per_side - 1) + w_idxs_ceil = (w_idxs.int() + 1).clip(max=num_grid_per_side - 1) - w_f = np.floor(w_idxs).astype(np.int64) - w_c = np.clip(w_f + 1, 0, self.num_grid_per_side - 1) - dw = w_idxs - w_f + dh = h_idxs - h_idxs_floor + dw = w_idxs - w_idxs_floor - side = self.num_grid_per_side + base_h = h_idxs_floor * num_grid_per_side + base_h_ceil = h_idxs_ceil * num_grid_per_side - indices = [ - (h_f[:, None] * side + w_f).flatten(), - (h_f[:, None] * side + w_c).flatten(), - (h_c[:, None] * side + w_f).flatten(), - (h_c[:, None] * side + w_c).flatten(), - ] - weights = [ - ((1 - dh)[:, None] * (1 - dw)).flatten(), - ((1 - dh)[:, None] * dw).flatten(), - (dh[:, None] * (1 - dw)).flatten(), - (dh[:, None] * dw).flatten(), - ] - return indices, weights + indices = [ + (base_h[None].T + w_idxs_floor[None]).flatten(), + (base_h[None].T + w_idxs_ceil[None]).flatten(), + (base_h_ceil[None].T + w_idxs_floor[None]).flatten(), + (base_h_ceil[None].T + w_idxs_ceil[None]).flatten(), + ] - def _get_position_embedding(self, patch_pos_embeds, grid_ts, grid_hs, grid_ws): - """ - Tile and reorganize position embeddings to align with the token sequence. - """ - result_parts = [] - merge_size = self.spatial_merge_size + weights = [ + ((1 - dh)[None].T * (1 - dw)[None]).flatten(), + ((1 - dh)[None].T * dw[None]).flatten(), + (dh[None].T * (1 - dw)[None]).flatten(), + (dh[None].T * dw[None]).flatten(), + ] - for pos_embed, t, h, w in zip(patch_pos_embeds, grid_ts, grid_hs, grid_ws): - pos_embed = pos_embed.repeat(t, 1) + for i in range(4): + idx_list[i].extend(indices[i].tolist()) + weight_list[i].extend(weights[i].tolist()) + + idx_tensor = torch.tensor(idx_list, dtype=torch.long, device=device) + weight_tensor = torch.tensor( + weight_list, dtype=self.pos_embed.weight.dtype, device=device + ) + pos_embeds = self.pos_embed(idx_tensor).to(device) * weight_tensor[:, :, None] + patch_pos_embeds = pos_embeds[0] + pos_embeds[1] + pos_embeds[2] + pos_embeds[3] - h_merge = h // merge_size - w_merge = w // merge_size + patch_pos_embeds = patch_pos_embeds.split( + [h * w for h, w in zip(grid_hs, grid_ws)] + ) + patch_pos_embeds_permute = [] + merge_size = self.spatial_merge_size + for pos_embed, t, h, w in zip(patch_pos_embeds, grid_ts, grid_hs, grid_ws): + pos_embed = pos_embed.repeat(t, 1) pos_embed = ( - pos_embed.view(t, h_merge, merge_size, w_merge, merge_size, -1) + pos_embed.view( + t, h // merge_size, merge_size, w // merge_size, merge_size, -1 + ) .permute(0, 1, 3, 2, 4, 5) .flatten(0, 4) ) - result_parts.append(pos_embed) + patch_pos_embeds_permute.append(pos_embed) - return torch.cat(result_parts, dim=0) + return torch.cat(patch_pos_embeds_permute) def _torch_interp_indices( self, dim_size: int, device: torch.device @@ -791,61 +790,6 @@ def bucket_flashinfer_max_seqlen(self, real_max_seqlen: int) -> int: round_up(real_max_seqlen, FLASHINFER_MAX_SEQLEN_BUCKETS[-1]), ) - def fast_pos_embed_interpolate(self, grid_thw): - """Interpolate position embeddings for (batch, 3) size input dimensions. - - Performs bilinear interpolation on spatial dimensions (height, width) and replicates - along temporal dimension. The result is reorganized according to spatial_merge_size. - - Args: - grid_thw: Tensor of shape [batch_size, 3] with (temporal, height, width) dimensions - in patches for each sample. - - Returns: - Interpolated position embeddings tensor. - """ - grid_thw_cpu = grid_thw.cpu().numpy() - - # transfer data to CPU before loop - temporal_dims = grid_thw_cpu[:, 0].tolist() - height_dims = grid_thw_cpu[:, 1].tolist() - width_dims = grid_thw_cpu[:, 2].tolist() - - device = self.pos_embed.weight.device - dtype = self.pos_embed.weight.dtype - - patches_size = [h * w for h, w in zip(height_dims, width_dims)] - total_patches = sum(patches_size) - all_indices_np = np.zeros((4, total_patches), dtype=np.int64) - all_weights_np = np.zeros((4, total_patches), dtype=np.float32) - - current_idx = 0 - - # calculate indices and weights on CPU - for t, h, w in zip(temporal_dims, height_dims, width_dims): - h_idxs = self._get_interpolation_indices(h) - w_idxs = self._get_interpolation_indices(w) - - indices, weights = self._calculate_indices_and_weights(h_idxs, w_idxs) - - end_idx = current_idx + h * w - for i in range(4): - all_indices_np[i, current_idx:end_idx] = indices[i] - all_weights_np[i, current_idx:end_idx] = weights[i] - current_idx = end_idx - - idx_tensor = torch.from_numpy(all_indices_np).to(device) - weight_tensor = torch.from_numpy(all_weights_np).to(dtype=dtype, device=device) - - # calculate interpolation - pos_embeds = self.pos_embed(idx_tensor.view(-1)) - pos_embeds = pos_embeds.view(4, total_patches, -1) - patch_pos_embeds = (pos_embeds * weight_tensor.unsqueeze(-1)).sum(dim=0) - patch_pos_embeds = patch_pos_embeds.split(patches_size) - return self._get_position_embedding( - patch_pos_embeds, temporal_dims, height_dims, width_dims - ) - def compute_flashinfer_batch_offsets_packed( self, token_cu_seqlens: np.ndarray, @@ -1163,9 +1107,7 @@ def __init__( ) # Use HF deepstack order only if rl_on_policy_target is set; # otherwise, retain original order for inference accuracy. - self.use_hf_deepstack_order = ( - get_exec().deterministic.rl_on_policy_target is not None - ) + self.use_hf_deepstack_order = is_true_on_policy_enabled() def get_deepstack_embeds( self, layer_idx: int, input_deepstack_embeds: Optional[torch.Tensor] diff --git a/python/sglang/srt/models/sdar.py b/python/sglang/srt/models/sdar.py index 3d5912d975ea..a4bfd23fcfd9 100644 --- a/python/sglang/srt/models/sdar.py +++ b/python/sglang/srt/models/sdar.py @@ -41,7 +41,15 @@ create_fused_set_kv_buffer_arg, enable_fused_set_kv_buffer, ) -from sglang.srt.runtime_context import get_exec, get_forward, get_parallel, get_stream +from sglang.srt.runtime_context import ( + get_forward, + get_parallel, + get_stream, +) +from sglang.srt.true_on_policy import ( + get_on_policy_rms_norm_kwargs, + should_force_bfloat16_dense_tensor_math, +) from sglang.srt.utils import add_prefix, is_cuda, make_layers logger = logging.getLogger(__name__) @@ -202,7 +210,7 @@ def forward( hidden_states: torch.Tensor, forward_batch: ForwardBatch, ): - if get_exec().deterministic.rl_on_policy_target is not None: + if should_force_bfloat16_dense_tensor_math(): hidden_states = hidden_states.bfloat16() qkv, _ = self.qkv_proj(hidden_states) @@ -230,7 +238,7 @@ def forward( ), ) - if get_exec().deterministic.rl_on_policy_target is not None: + if should_force_bfloat16_dense_tensor_math(): q = q.to(torch.bfloat16) k = k.to(torch.bfloat16) @@ -258,15 +266,10 @@ def __init__( self.hidden_size = config.hidden_size self.layer_id = layer_id - norm_kwargs = ( - dict( - weight_dtype=torch.float32, - cast_x_before_out_mul=True, - override_orig_dtype=torch.float32, - fp32_residual=True, - ) - if get_exec().deterministic.rl_on_policy_target is not None - else {} + norm_kwargs = get_on_policy_rms_norm_kwargs( + weight_dtype=torch.float32, + override_orig_dtype=torch.float32, + fp32_residual=True, ) self.input_layernorm = RMSNorm( self.hidden_size, eps=config.rms_norm_eps, **norm_kwargs @@ -383,15 +386,10 @@ def __init__( prefix=add_prefix("layers", prefix), ) if self.pp_group.is_last_rank: - norm_kwargs = ( - dict( - weight_dtype=torch.float32, - cast_x_before_out_mul=True, - override_orig_dtype=torch.float32, - fp32_residual=True, - ) - if get_exec().deterministic.rl_on_policy_target is not None - else {} + norm_kwargs = get_on_policy_rms_norm_kwargs( + weight_dtype=torch.float32, + override_orig_dtype=torch.float32, + fp32_residual=True, ) self.norm = RMSNorm(self.embed_dim, eps=config.rms_norm_eps, **norm_kwargs) else: diff --git a/python/sglang/srt/models/sdar_moe.py b/python/sglang/srt/models/sdar_moe.py index 9ac15f3fe7f4..2a476c7bdc3d 100644 --- a/python/sglang/srt/models/sdar_moe.py +++ b/python/sglang/srt/models/sdar_moe.py @@ -57,7 +57,16 @@ create_fused_set_kv_buffer_arg, enable_fused_set_kv_buffer, ) -from sglang.srt.runtime_context import get_exec, get_forward, get_parallel, get_stream +from sglang.srt.runtime_context import ( + get_exec, + get_forward, + get_parallel, + get_stream, +) +from sglang.srt.true_on_policy import ( + get_on_policy_rms_norm_kwargs, + should_force_bfloat16_dense_tensor_math, +) from sglang.srt.utils import LazyValue, add_prefix, is_cuda, make_layers logger = logging.getLogger(__name__) @@ -269,7 +278,7 @@ def forward( hidden_states: torch.Tensor, forward_batch: ForwardBatch, ) -> torch.Tensor: - if get_exec().deterministic.rl_on_policy_target is not None: + if should_force_bfloat16_dense_tensor_math(): hidden_states = hidden_states.bfloat16() qkv, _ = self.qkv_proj(hidden_states) @@ -297,7 +306,7 @@ def forward( ), ) - if get_exec().deterministic.rl_on_policy_target is not None: + if should_force_bfloat16_dense_tensor_math(): q = q.to(torch.bfloat16) k = k.to(torch.bfloat16) @@ -326,15 +335,10 @@ def __init__( self.hidden_size = config.hidden_size self.layer_id = layer_id - norm_kwargs = ( - dict( - weight_dtype=torch.float32, - cast_x_before_out_mul=True, - override_orig_dtype=torch.float32, - fp32_residual=True, - ) - if get_exec().deterministic.rl_on_policy_target is not None - else {} + norm_kwargs = get_on_policy_rms_norm_kwargs( + weight_dtype=torch.float32, + override_orig_dtype=torch.float32, + fp32_residual=True, ) self.input_layernorm = RMSNorm( self.hidden_size, eps=config.rms_norm_eps, **norm_kwargs @@ -466,15 +470,10 @@ def __init__( ) if self.pp_group.is_last_rank: - norm_kwargs = ( - dict( - weight_dtype=torch.float32, - cast_x_before_out_mul=True, - override_orig_dtype=torch.float32, - fp32_residual=True, - ) - if get_exec().deterministic.rl_on_policy_target is not None - else {} + norm_kwargs = get_on_policy_rms_norm_kwargs( + weight_dtype=torch.float32, + override_orig_dtype=torch.float32, + fp32_residual=True, ) self.norm = RMSNorm(self.embed_dim, eps=config.rms_norm_eps, **norm_kwargs) else: diff --git a/python/sglang/srt/models/step3p5.py b/python/sglang/srt/models/step3p5.py index c9220f2f367d..9c8a422d6c4f 100644 --- a/python/sglang/srt/models/step3p5.py +++ b/python/sglang/srt/models/step3p5.py @@ -46,7 +46,13 @@ ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_loader.weight_utils import default_weight_loader -from sglang.srt.runtime_context import get_exec, get_forward, get_parallel, get_stream +from sglang.srt.runtime_context import ( + get_exec, + get_forward, + get_parallel, + get_stream, +) +from sglang.srt.true_on_policy import is_true_on_policy_enabled from sglang.srt.utils import add_prefix, is_cuda, is_non_idle_and_non_empty, make_layers Step3p5Config = None @@ -669,11 +675,7 @@ def __init__( quant_config=quant_config, enable_tp=not is_dp_attention_enabled(), prefix=add_prefix("embed_tokens", prefix), - params_dtype=( - torch.float32 - if get_exec().deterministic.rl_on_policy_target is not None - else None - ), + params_dtype=torch.float32 if is_true_on_policy_enabled() else None, ) else: self.embed_tokens = PPMissingLayer() diff --git a/python/sglang/srt/models/utils.py b/python/sglang/srt/models/utils.py index e8a03a2f5a19..ec5d580e0c67 100644 --- a/python/sglang/srt/models/utils.py +++ b/python/sglang/srt/models/utils.py @@ -514,6 +514,8 @@ def apply_qk_norm( _is_cuda # TODO(dark): have not tested on ROCm or other backends and allow_inplace # TODO(dark): this can be relaxed if needed and (q_eps == k_eps) # TODO(dark): this can also be relaxed + and q_norm.weight.dtype == q.dtype + and k_norm.weight.dtype == k.dtype and not envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.get() and get_exec().graph.cuda_graph_config.prefill.tc_compiler != "inductor" # let inductor fuse QK norm diff --git a/python/sglang/srt/multimodal/processors/base_processor.py b/python/sglang/srt/multimodal/processors/base_processor.py index 9c6d172b0f8c..12241b2740f8 100644 --- a/python/sglang/srt/multimodal/processors/base_processor.py +++ b/python/sglang/srt/multimodal/processors/base_processor.py @@ -47,6 +47,7 @@ get_mm, get_serving, ) +from sglang.srt.true_on_policy import is_true_on_policy_enabled from sglang.srt.utils import ( CLIENT_MEDIA_EXCEPTIONS, configure_media_url_security, @@ -728,7 +729,11 @@ def _preprocessing_competes_with_the_scheduler(self) -> bool: preprocessing worker there is one more competitor for that device rather than added parallelism. """ - if _is_cpu or get_exec().deterministic.rl_on_policy_target is not None: + if ( + _is_cpu + or get_exec().deterministic.rl_on_policy_target is not None + or is_true_on_policy_enabled() + ): return False if self.disable_fast_image_processor: return False @@ -764,7 +769,11 @@ def _fast_image_processor_device(self, processor) -> Optional[str]: tokenizer process each carry their own ``base_gpu_id``. """ server_args = self.server_args - if _is_cpu or get_exec().deterministic.rl_on_policy_target is not None: + if ( + _is_cpu + or get_exec().deterministic.rl_on_policy_target is not None + or is_true_on_policy_enabled() + ): return "cpu" if _is_xpu: return "xpu" @@ -1732,7 +1741,14 @@ def resolve_image_token_counts(self, images: List) -> List[int]: """ assert images is not None - image_sizes = [(image.height, image.width) for image in images] + image_sizes = [ + ( + tuple(image.shape[-2:]) + if isinstance(image, torch.Tensor) + else (image.height, image.width) + ) + for image in images + ] num_image_tokens = self._processor._get_num_multimodal_tokens( image_sizes=image_sizes ).num_image_tokens diff --git a/python/sglang/srt/multimodal/processors/qwen_vl.py b/python/sglang/srt/multimodal/processors/qwen_vl.py index 44f179393b65..4b295a35866d 100644 --- a/python/sglang/srt/multimodal/processors/qwen_vl.py +++ b/python/sglang/srt/multimodal/processors/qwen_vl.py @@ -1011,7 +1011,7 @@ async def _process_mm_data_uncached( **kwargs, ): entry_time = time.perf_counter() - base_output = await self.load_mm_data( + base_output = await self.legacy_load_mm_data( prompt=input_text, image_data=image_data, video_data=request_obj.video_data, @@ -1180,31 +1180,11 @@ async def process_mm_data_async( *args, **kwargs, ): - if ( - not image_data - or request_obj.video_data - or request_obj.audio_data - or any(self._is_preprocessed_input(item) for item in image_data) - or ( - not self.mm_preprocess_cache.enabled - and not self.uses_media_artifacts_without_cache - ) - ): - return await self._process_mm_data_uncached( - image_data, input_text, request_obj, *args, **kwargs - ) - - prepare_artifacts = ( - self.prepare_media_artifacts - if self.mm_preprocess_cache.enabled - else self.prepare_media_artifacts_without_cache - ) - artifacts = await prepare_artifacts( - image_data, content_hashes=getattr(request_obj, "mm_content_hashes", None) - ) - composed = self.compose_image_artifacts(input_text, artifacts) - if composed is not None: - return composed + # The media-artifact pipeline rebuilds the image prompt template from + # its own grid bookkeeping, which does not agree with the token layout + # RL rollouts replay: build_input_ids then walks off the end of the grid + # list. Take the legacy loader for every request, as this processor did + # before the artifact pipeline existed. return await self._process_mm_data_uncached( image_data, input_text, request_obj, *args, **kwargs ) diff --git a/python/sglang/srt/ray/__init__.py b/python/sglang/srt/ray/__init__.py index 5927c789f926..ba619e3bc7a5 100644 --- a/python/sglang/srt/ray/__init__.py +++ b/python/sglang/srt/ray/__init__.py @@ -1,3 +1,4 @@ from sglang.srt.ray.engine import RayEngine +from sglang.srt.ray.http_server import launch_engine, launch_server, serve_http -__all__ = ["RayEngine"] +__all__ = ["RayEngine", "launch_engine", "launch_server", "serve_http"] diff --git a/python/sglang/srt/ray/engine.py b/python/sglang/srt/ray/engine.py index 41f667d9e795..368f409832cd 100644 --- a/python/sglang/srt/ray/engine.py +++ b/python/sglang/srt/ray/engine.py @@ -15,9 +15,11 @@ from __future__ import annotations +import contextvars import dataclasses import logging import threading +from contextlib import contextmanager from typing import Callable, List, Optional import ray @@ -40,6 +42,20 @@ logger = logging.getLogger(__name__) +_caller_placement_group: contextvars.ContextVar[Optional[PlacementGroup]] = ( + contextvars.ContextVar("sglang_ray_caller_placement_group", default=None) +) + + +@contextmanager +def _placement_group_context(placement_group: Optional[PlacementGroup]): + """Expose launch-only state while Engine synchronously starts schedulers.""" + token = _caller_placement_group.set(placement_group) + try: + yield + finally: + _caller_placement_group.reset(token) + @dataclasses.dataclass class RaySchedulerInitResult(SchedulerInitResult): @@ -183,6 +199,22 @@ def _validate_custom_placement_group(pg: PlacementGroup, world_size: int) -> Non ) +def get_scheduler_actor_name( + *, + rank0_node_ip: str, + dp_rank: int, + pp_rank: int, + tp_rank: int, + pg_id_hex: str, + bundle_idx: int, +) -> str: + return ( + f"sglang_scheduler_node{rank0_node_ip}" + f"_dp{dp_rank}_pp{pp_rank}_tp{tp_rank}" + f"_pg{pg_id_hex}_bundle{bundle_idx}" + ) + + def _create_scheduler_actor( pg: PlacementGroup, bundle_idx: int, @@ -206,14 +238,25 @@ def _create_scheduler_actor( """ attn_cp_rank, moe_dp_rank, moe_ep_rank = _compute_parallelism_ranks(tp_rank) + rdt = server_args.enable_rdt_weight_sync + return SchedulerActor.options( num_cpus=0, num_gpus=1, - name=( - f"sglang_scheduler_node{rank0_node_ip}" - f"_dp{dp_rank}_pp{pp_rank}_tp{tp_rank}" - f"_pg{pg.id.hex()[:8]}_bundle{bundle_idx}" + # run_event_loop() blocks one thread for the actor's lifetime; leave a spare + # for pull_weights, which the trainer calls while generation is paused. + max_concurrency=2 if rdt else 1, + name=get_scheduler_actor_name( + rank0_node_ip=rank0_node_ip, + dp_rank=dp_rank, + pp_rank=pp_rank, + tp_rank=tp_rank, + pg_id_hex=pg.id.hex()[:8], + bundle_idx=bundle_idx, ), + # SchedulerActor calls set_device() with the absolute id from + # get_accelerator_ids(), which is only valid if Ray leaves the mask alone. + runtime_env={"env_vars": {"RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1"}}, scheduling_strategy=PlacementGroupSchedulingStrategy( placement_group=pg, placement_group_bundle_index=bundle_idx, @@ -233,12 +276,13 @@ def _create_scheduler_actor( class RayEngine(Engine): - """Engine using Ray actors for scheduler processes.""" + """Engine using Ray actors for scheduler processes. + + Same constructor kwargs as :class:`Engine`, plus ``placement_group`` (a Ray + PlacementGroup handle, not a ServerArgs field). + """ - def __init__(self, **kwargs): - # Set before super().__init__(): it launches the subprocesses, which need - # the group to schedule the scheduler actors onto. - self._placement_group = kwargs.pop("placement_group", None) + def __init__(self, *, placement_group: Optional[PlacementGroup] = None, **kwargs): if "log_level" not in kwargs: kwargs["log_level"] = "error" # Schedulers are separate Ray actors; default to the Ray-backed @@ -247,7 +291,10 @@ def __init__(self, **kwargs): from sglang.srt.observability.ray_wrappers import build_ray_stat_loggers kwargs["stat_loggers"] = build_ray_stat_loggers() - super().__init__(server_args=ServerArgs(**kwargs)) + # The context is read while super().__init__() launches the subprocesses, + # which need the group to schedule the scheduler actors onto. + with _placement_group_context(placement_group): + super().__init__(server_args=ServerArgs(**kwargs)) def shutdown(self): """Shutdown the engine — kill Ray scheduler actors then local processes.""" @@ -264,8 +311,6 @@ def _launch_scheduler_processes( server_args: ServerArgs, port_args: PortArgs, run_scheduler_process_func: Callable, - *, - placement_group=None, ) -> tuple[SchedulerInitResult, None]: """Launch schedulers as Ray actors. @@ -273,6 +318,7 @@ def _launch_scheduler_processes( Tuple of (RaySchedulerInitResult, None). scheduler_procs is None since Ray uses actors instead of mp.Process. """ + placement_group = _caller_placement_group.get() pg = placement_group or ray.util.get_current_placement_group() if pg is None: from ray.util.placement_group import ( diff --git a/python/sglang/srt/ray/http_server.py b/python/sglang/srt/ray/http_server.py index 963f05d13524..6b37cd4798bd 100644 --- a/python/sglang/srt/ray/http_server.py +++ b/python/sglang/srt/ray/http_server.py @@ -15,6 +15,8 @@ from typing import Callable, Optional +from ray.util.placement_group import PlacementGroup + from sglang.srt.entrypoints.engine import ( init_tokenizer_manager, run_detokenizer_process, @@ -23,23 +25,37 @@ from sglang.srt.server_args import ServerArgs -def launch_server( +def launch_engine( server_args: ServerArgs, init_tokenizer_manager_func: Callable = init_tokenizer_manager, run_scheduler_process_func: Callable = run_scheduler_process, run_detokenizer_process_func: Callable = run_detokenizer_process, + *, + placement_group: Optional[PlacementGroup] = None, +): + """Create RayEngine subprocesses / SchedulerActors.""" + from sglang.srt.ray.engine import RayEngine, _placement_group_context + + with _placement_group_context(placement_group): + return RayEngine._launch_subprocesses( + server_args, + init_tokenizer_manager_func=init_tokenizer_manager_func, + run_scheduler_process_func=run_scheduler_process_func, + run_detokenizer_process_func=run_detokenizer_process_func, + ) + + +def serve_http( + engine, + server_args: ServerArgs, execute_warmup_func: Optional[Callable] = None, launch_callback: Optional[Callable[[], None]] = None, ): - """Launch HTTP server with Ray-based scheduler actors. - - Mirrors http_server.launch_server() but uses RayEngine for scheduler launching. - """ + """Block in uvicorn. ``engine`` is the tuple from ``launch_engine``.""" from sglang.srt.entrypoints.http_server import ( _execute_server_warmup, _setup_and_run_http_server, ) - from sglang.srt.ray.engine import RayEngine if execute_warmup_func is None: execute_warmup_func = _execute_server_warmup @@ -51,12 +67,7 @@ def launch_server( scheduler_init_result, subprocess_watchdog, _weight_cache_daemon_procs, - ) = RayEngine._launch_subprocesses( - server_args, - init_tokenizer_manager_func=init_tokenizer_manager_func, - run_scheduler_process_func=run_scheduler_process_func, - run_detokenizer_process_func=run_detokenizer_process_func, - ) + ) = engine _setup_and_run_http_server( server_args, @@ -68,3 +79,28 @@ def launch_server( execute_warmup_func=execute_warmup_func, launch_callback=launch_callback, ) + + +def launch_server( + server_args: ServerArgs, + init_tokenizer_manager_func: Callable = init_tokenizer_manager, + run_scheduler_process_func: Callable = run_scheduler_process, + run_detokenizer_process_func: Callable = run_detokenizer_process, + execute_warmup_func: Optional[Callable] = None, + launch_callback: Optional[Callable[[], None]] = None, +): + """Launch HTTP server with Ray-based scheduler actors. + + Mirrors http_server.launch_server() but uses RayEngine for scheduler launching. + """ + serve_http( + launch_engine( + server_args, + init_tokenizer_manager_func=init_tokenizer_manager_func, + run_scheduler_process_func=run_scheduler_process_func, + run_detokenizer_process_func=run_detokenizer_process_func, + ), + server_args, + execute_warmup_func=execute_warmup_func, + launch_callback=launch_callback, + ) diff --git a/python/sglang/srt/ray/scheduler_actor.py b/python/sglang/srt/ray/scheduler_actor.py index 3616d18161b4..d4b69583d903 100644 --- a/python/sglang/srt/ray/scheduler_actor.py +++ b/python/sglang/srt/ray/scheduler_actor.py @@ -16,9 +16,10 @@ from __future__ import annotations import logging -from typing import Any, Dict, Optional +from typing import Any, Dict, List, Optional import ray +from ray import ObjectRef from sglang.srt.arg_groups.overrides import declare_resolution from sglang.srt.runtime_context import publish @@ -125,6 +126,38 @@ def get_info(self) -> Dict[str, Any]: """Return scheduler initialization info for handshake.""" return self.scheduler.get_init_info() + def register_weight_for_rdt(self) -> None: + """Pin model parameters with NIXL for repeated RDT pulls.""" + if self.scheduler.server_args.enable_memory_saver: + return + + import torch + from ray.experimental import register_nixl_memory + + torch.cuda.set_device(self.scheduler.ps.gpu_id) + model = self.scheduler.tp_worker.model_runner.model + for _, param in model.named_parameters(): + register_nixl_memory(param.data) + + def pull_weights( + self, weights_refs: List[ObjectRef], param_names: List[str] + ) -> None: + """Pull a pre-sharded weight bucket from the trainer via RDT zero-copy. + + ``weights_refs`` is a list so Ray does not resolve the ref on call. + """ + import torch + from ray.experimental import set_target_for_ref + + # Runs on a different thread than run_event_loop, which owns the device binding. + torch.cuda.set_device(self.scheduler.ps.gpu_id) + + model = self.scheduler.tp_worker.model_runner.model + params_dict = dict(model.named_parameters()) + target_buffers = [params_dict[name].data for name in param_names] + set_target_for_ref(weights_refs[0], target_buffers) + ray.get(weights_refs[0]) + def run_event_loop(self) -> None: """Run the scheduler's event loop. Blocks until shutdown.""" try: diff --git a/python/sglang/srt/sampling/sampling_batch_info.py b/python/sglang/srt/sampling/sampling_batch_info.py index fad02f235047..b825a7ddcebe 100644 --- a/python/sglang/srt/sampling/sampling_batch_info.py +++ b/python/sglang/srt/sampling/sampling_batch_info.py @@ -26,6 +26,15 @@ logger = logging.getLogger(__name__) +def _rows_to_device_indices(rows: List[int], device: str) -> torch.Tensor: + """Move batch row numbers to the device without blocking on a pageable copy.""" + return torch.tensor( + rows, + dtype=torch.long, + pin_memory=is_pin_memory_available(device), + ).to(device, non_blocking=True) + + @dataclasses.dataclass class SamplingBatchInfo: # Basic batched sampling params @@ -77,7 +86,9 @@ class SamplingBatchInfo: # CPU request flags and their derived device indices. The flags make batch # filtering cheap; the indices keep sampler work limited to opted-in rows. return_sampling_masks: Optional[List[bool]] = None + return_sampling_support_logprobs: Optional[List[bool]] = None sampling_mask_batch_indices: Optional[torch.Tensor] = None + sampling_support_logprobs_capture_indices: Optional[torch.Tensor] = None # Device device: str = "cuda" @@ -147,9 +158,20 @@ def from_schedule_batch(cls, batch: ScheduleBatch, vocab_size: int): and any(r.custom_logit_processor for r in reqs) # check the flag first. ) # then check the requests. return_sampling_masks = [r.return_sampling_mask for r in reqs] + return_sampling_support_logprobs = [ + r.return_sampling_mask and r.sampling_logprobs_mode == "support" + for r in reqs + ] sampling_mask_batch_indices = cls._make_sampling_mask_batch_indices( return_sampling_masks, device ) + sampling_support_logprobs_capture_indices = ( + cls._make_sampling_support_logprobs_capture_indices( + return_sampling_masks, + return_sampling_support_logprobs, + device, + ) + ) if has_custom_logit_processor: # Merge the same type of custom logit processors together @@ -215,7 +237,11 @@ def from_schedule_batch(cls, batch: ScheduleBatch, vocab_size: int): device=device, logit_bias=logit_bias, return_sampling_masks=return_sampling_masks, + return_sampling_support_logprobs=return_sampling_support_logprobs, sampling_mask_batch_indices=sampling_mask_batch_indices, + sampling_support_logprobs_capture_indices=( + sampling_support_logprobs_capture_indices + ), ) ret.adjusted_from_schedule_batch(batch, vocab_size) return ret @@ -244,11 +270,31 @@ def _make_sampling_mask_batch_indices( ] if not indices: return None - return torch.tensor( - indices, - dtype=torch.long, - pin_memory=is_pin_memory_available(device), - ).to(device, non_blocking=True) + return _rows_to_device_indices(indices, device) + + @staticmethod + def _make_sampling_support_logprobs_capture_indices( + return_sampling_masks: List[bool], + return_sampling_support_logprobs: Optional[List[bool]], + device: str, + ) -> Optional[torch.Tensor]: + """Map support-mode requests to rows in the compact mask capture.""" + if return_sampling_support_logprobs is None: + return None + capture_index = -1 + indices = [] + for return_mask, return_support in zip( + return_sampling_masks, + return_sampling_support_logprobs, + strict=True, + ): + if return_mask: + capture_index += 1 + if return_support: + indices.append(capture_index) + if not indices: + return None + return _rows_to_device_indices(indices, device) def __len__(self): return len(self.temperatures) @@ -360,12 +406,23 @@ def filter_batch(self, keep_indices: List[int], keep_indices_device: torch.Tenso self.return_sampling_masks = [ self.return_sampling_masks[i] for i in keep_indices ] + if self.return_sampling_support_logprobs is not None: + self.return_sampling_support_logprobs = [ + self.return_sampling_support_logprobs[i] for i in keep_indices + ] if self.sampling_mask_batch_indices is not None: self.sampling_mask_batch_indices = ( self._make_sampling_mask_batch_indices( self.return_sampling_masks, self.device ) ) + self.sampling_support_logprobs_capture_indices = ( + self._make_sampling_support_logprobs_capture_indices( + self.return_sampling_masks, + self.return_sampling_support_logprobs, + self.device, + ) + ) self.adjusted_filter_batch(keep_indices, keep_indices_device) @@ -481,6 +538,16 @@ def merge_batch(self, other: SamplingBatchInfo): self.return_sampling_masks = ( self.return_sampling_masks or [False] * self_len ) + (other.return_sampling_masks or [False] * other_len) + self.return_sampling_support_logprobs = ( + self.return_sampling_support_logprobs or [False] * self_len + ) + (other.return_sampling_support_logprobs or [False] * other_len) + self.sampling_support_logprobs_capture_indices = ( + self._make_sampling_support_logprobs_capture_indices( + self.return_sampling_masks, + self.return_sampling_support_logprobs, + self.device, + ) + ) # Note: because the __len()__ operator is defined on the temperatures tensor, # please make sure any merge operation with len(self) or len(other) is done before diff --git a/python/sglang/srt/sampling/sampling_mask.py b/python/sglang/srt/sampling/sampling_mask.py new file mode 100644 index 000000000000..cb3c319b07c2 --- /dev/null +++ b/python/sglang/srt/sampling/sampling_mask.py @@ -0,0 +1,130 @@ +"""Host-side transport of output-token sampling masks.""" + +import pickle +import sys +from array import array +from typing import List, Optional, Tuple, Union + +import msgspec +import numpy as np + + +class SamplingMaskChunk( + msgspec.Struct, frozen=True, kw_only=True, array_like=True, gc=False +): + """One request's sampling masks in one output batch, one row per output token.""" + + # int32 [num_rows] + lengths: np.ndarray + # int32 [sum(lengths)], row after row + token_ids: np.ndarray + # float32 [sum(lengths)] in support mode, [num_rows] in selected mode + logprobs: np.ndarray + + def __reduce_ex__(self, protocol: int): + # Read-only buffers unpickle as bytes, which np.frombuffer wraps without the + # GC-tracked objects that numpy's reducer or a bytearray would leave per array. + if protocol < 5: + return super().__reduce_ex__(protocol) + return ( + _chunk_from_buffers, + tuple( + pickle.PickleBuffer(memoryview(values).toreadonly()) + for values in (self.lengths, self.token_ids, self.logprobs) + ), + ) + + def to_lists( + self, support_logprobs: bool + ) -> Tuple[List[List[int]], List[Union[float, List[float]]]]: + masks: List[List[int]] = [] + values: List[Union[float, List[float]]] = ( + [] if support_logprobs else self.logprobs.tolist() + ) + start = 0 + for length in self.lengths.tolist(): + end = start + length + masks.append(self.token_ids[start:end].tolist()) + if support_logprobs: + values.append(self.logprobs[start:end].tolist()) + start = end + return masks, values + + +def _chunk_from_buffers(lengths, token_ids, logprobs) -> SamplingMaskChunk: + return SamplingMaskChunk( + lengths=np.frombuffer(lengths, dtype=np.int32), + token_ids=np.frombuffer(token_ids, dtype=np.int32), + logprobs=np.frombuffer(logprobs, dtype=np.float32), + ) + + +class _GrowableArray(msgspec.Struct): + """Append-only 1-D buffer; doubling capacity copies each value O(1) times. + + ``reset`` keeps the buffer as a spare, reused once no view of it is alive, + so steady-state appends touch no new pages. + """ + + buffer: np.ndarray + size: int = 0 + spare: Optional[np.ndarray] = None + + def extend(self, values: np.ndarray) -> None: + end = self.size + len(values) + if end > len(self.buffer): + self.buffer = self._grow(end) + self.buffer[self.size : end] = values + self.size = end + + def view(self) -> np.ndarray: + return self.buffer[: self.size] + + def reset(self) -> None: + self.spare = self.buffer + self.buffer = np.empty(0, dtype=self.buffer.dtype) + self.size = 0 + + def _grow(self, end: int) -> np.ndarray: + spare, self.spare = self.spare, None + # Two references, the local and the argument, mean no view is alive. + if spare is not None and len(spare) >= end and sys.getrefcount(spare) == 2: + grown = spare + else: + grown = np.empty(max(end, 2 * len(self.buffer)), dtype=self.buffer.dtype) + grown[: self.size] = self.buffer[: self.size] + return grown + + +class SamplingMaskRows(msgspec.Struct): + """One request's SamplingMaskChunk rows, queued until they are streamed.""" + + _lengths: array = msgspec.field(default_factory=lambda: array("i")) + _token_ids: _GrowableArray = msgspec.field( + default_factory=lambda: _GrowableArray(np.empty(0, dtype=np.int32)) + ) + _logprobs: _GrowableArray = msgspec.field( + default_factory=lambda: _GrowableArray(np.empty(0, dtype=np.float32)) + ) + + def append(self, token_ids: np.ndarray, logprobs: np.ndarray) -> None: + """Queue one token's int32 support and float32 logprobs.""" + self._lengths.append(len(token_ids)) + self._token_ids.extend(token_ids) + self._logprobs.extend(logprobs) + + def view(self) -> SamplingMaskChunk: + """View the queued token IDs and logprobs without copying them.""" + return SamplingMaskChunk( + lengths=np.array(self._lengths, dtype=np.int32), + token_ids=self._token_ids.view(), + logprobs=self._logprobs.view(), + ) + + def take(self) -> SamplingMaskChunk: + """Hand off the queued rows without copying them, leaving the queue empty.""" + chunk = self.view() + del self._lengths[:] + self._token_ids.reset() + self._logprobs.reset() + return chunk diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index d79bb75d7a32..f8791c0c7d6d 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -503,6 +503,22 @@ def remote_instance_weight_loader_use_transfer_engine(self, load_format=None): ``--speculative-draft-load-format`` needs its own transfer engine.""" return remote_instance_transfer_engine_of(resolving_view(self), load_format) + def needs_engine_info_bootstrap(self) -> bool: + """Whether this node (rank 0) hosts the EngineInfoBootstrapServer.""" + return ( + resolving_view( + self + ).remote_instance_weight_loader_start_seed_via_transfer_engine + or self.enable_engine_info_bootstrap + ) + + def registers_parallelism_config(self) -> bool: + """Whether this rank publishes its parallelism config to the bootstrap server.""" + return ( + self.remote_instance_weight_loader_use_transfer_engine() + or self.enable_engine_info_bootstrap + ) + # The namespaces whose *input* fields make up the record. A namespace declares # its input and derived fields side by side; only the input half is collected, diff --git a/python/sglang/srt/session/session_controller.py b/python/sglang/srt/session/session_controller.py index c23c9653c6e3..7f28b8594a70 100644 --- a/python/sglang/srt/session/session_controller.py +++ b/python/sglang/srt/session/session_controller.py @@ -310,6 +310,7 @@ def create_req( top_logprobs_num=req.top_logprobs_num, token_ids_logprob=req.token_ids_logprob, return_sampling_mask=req.return_sampling_mask, + sampling_logprobs_mode=req.sampling_logprobs_mode, vocab_size=vocab_size, eos_token_ids=eos_token_ids, require_reasoning=req.require_reasoning, diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index 28c5e42cdec5..d4282e6c34bc 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -2,7 +2,7 @@ import math import os from dataclasses import replace -from typing import List, Optional, Tuple +from typing import TYPE_CHECKING, List, Optional, Tuple import torch import torch.distributed as dist @@ -93,6 +93,9 @@ _is_npu = is_npu() +if TYPE_CHECKING: + from sglang.srt.model_executor.model_runner import ModelRunner + logger = logging.getLogger(__name__) _FusedKVMaterializeHelper = None @@ -359,6 +362,9 @@ class DFlashWorkerV2(BaseSpecWorker): scheduler runs it synchronously when overlap is disabled. """ + def iter_runners(self) -> List[Tuple[str, "ModelRunner"]]: + return [("draft", self.draft_model_runner)] + def __init__( self, server_args: ServerArgs, diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py index 0c6d80e3dfdd..2e11f124009b 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py @@ -582,8 +582,10 @@ def execute(self, forward_batch: ForwardBatch): forward_batch.seq_lens, forward_batch.out_cache_loc, forward_batch.positions, - forward_batch.spec_info.topk_p, - forward_batch.spec_info.topk_index, + forward_batch.spec_info.topk_p.clamp(0, 1), + forward_batch.spec_info.topk_index.clamp( + 0, self.model_runner.model_config.vocab_size - 1 + ), forward_batch.req_pool_indices, ] if buffers.rids_int is not None and forward_batch.rids_int is not None: diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 539b96eb0b41..8be32f2a099a 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -2,7 +2,7 @@ import logging import time from dataclasses import replace -from typing import List, Optional +from typing import TYPE_CHECKING, List, Optional, Tuple import torch @@ -38,7 +38,6 @@ speculative_moe_a2a_backend_context, speculative_moe_backend_context, ) -from sglang.srt.managers.io_struct import UpdateWeightsFromTensorReqInput from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.scheduler import GenerationBatchResult from sglang.srt.managers.tp_worker import TpModelWorker @@ -116,7 +115,6 @@ maybe_detect_oob, ) from sglang.srt.utils.common import ( - MultiprocessingSerializer, empty_context, fast_topk, get_available_gpu_memory, @@ -128,7 +126,6 @@ is_xpu, log_info_on_rank0, ) -from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions _is_cpu = is_cpu() _is_npu = is_npu() @@ -138,6 +135,9 @@ _is_xpu = is_xpu() +if TYPE_CHECKING: + from sglang.srt.model_executor.model_runner import ModelRunner + logger = logging.getLogger(__name__) @@ -441,7 +441,11 @@ def _capture_cuda_graphs(self): self.cuda_graph_runner = None self.cuda_graph_runner_for_draft_extend = None - if _is_cpu or check_cuda_graph_backend(Phase.DECODE, Backend.DISABLED): + if ( + _is_cpu + or check_cuda_graph_backend(Phase.DECODE, Backend.DISABLED) + or self.server_args.disable_draft_cuda_graph + ): return if get_model().model_impl == "mindspore": @@ -1154,6 +1158,9 @@ def _draft_extend_for_decode( class EAGLEWorkerV2(BaseSpecWorker): + def iter_runners(self) -> List[Tuple[str, "ModelRunner"]]: + return [("draft", self.draft_worker.draft_runner)] + def __init__( self, server_args: ServerArgs, @@ -1678,25 +1685,3 @@ def verify(self, batch: ScheduleBatch, grammar_barrier=None): finalize_tree_path=True, grammar_barrier=grammar_barrier, ) - - def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput): - monkey_patch_torch_reductions() - named_tensors = MultiprocessingSerializer.deserialize( - recv_req.serialized_named_tensors[self.ps.tp_rank] - ) - success, message = ( - self.draft_worker.draft_runner.weight_updater.update_weights_from_tensor( - named_tensors=named_tensors, - load_format=recv_req.load_format, - ) - ) - if not success: - return success, message - - success, message = ( - self.target_worker.model_runner.weight_updater.update_weights_from_tensor( - named_tensors=named_tensors, - load_format=recv_req.load_format, - ) - ) - return success, message diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index b62755c53e98..26867225f5dd 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -17,7 +17,7 @@ import logging import time from dataclasses import replace -from typing import TYPE_CHECKING, List +from typing import TYPE_CHECKING, List, Tuple import torch @@ -994,6 +994,12 @@ def _draft_extend_for_decode( class MultiLayerEagleWorkerV2(BaseSpecWorker): + def iter_runners(self) -> List[Tuple[str, ModelRunner]]: + return [ + (f"draft_step_{i}", r) + for i, r in enumerate(self.draft_worker.draft_runner_list) + ] + def __init__( self, server_args: ServerArgs, diff --git a/python/sglang/srt/speculative/ngram_worker.py b/python/sglang/srt/speculative/ngram_worker.py index 826895b28eac..69a799447d9c 100644 --- a/python/sglang/srt/speculative/ngram_worker.py +++ b/python/sglang/srt/speculative/ngram_worker.py @@ -1,5 +1,5 @@ import logging -from typing import List, Optional +from typing import TYPE_CHECKING, List, Optional, Tuple import numpy as np import torch @@ -38,6 +38,9 @@ _is_cpu = is_cpu() +if TYPE_CHECKING: + from sglang.srt.model_executor.model_runner import ModelRunner + logger = logging.getLogger(__name__) @@ -73,6 +76,11 @@ def _derive_tree_links( class NGRAMWorker(BaseSpecWorker): + def iter_runners(self) -> List[Tuple[str, "ModelRunner"]]: + # NGRAM shares the target's model_runner -- no independent draft weights + # (the n-gram corpus is a CPU lookup structure built from token streams). + return [] + def alloc_memory_pool(self, **kwargs): # The target memory pool does not exist yet when __init__ runs. self.req_to_token_pool, self.token_to_kv_pool_allocator = ( @@ -152,15 +160,6 @@ def clear_cache_pool(self): self.ngram_corpus.reset() self._prev_decode_rids = set() - def update_weights_from_tensor(self, recv_req): - # NGRAM has no draft weights of its own — the n-gram corpus is a CPU - # lookup structure built from request token streams — and its - # `model_runner` is shared with the target worker. The scheduler - # mixin dispatches via `self.draft_worker or self.tp_worker`, so - # without this method any caller of `update_weights_from_tensor` - # under `--speculative-algorithm NGRAM` raises AttributeError. - return self.target_worker.update_weights_from_tensor(recv_req) - def add_external_corpus(self, corpus_id: str, token_chunks: list[list[int]]) -> int: return self.ngram_corpus.load_external_corpus_named(corpus_id, token_chunks) diff --git a/python/sglang/srt/state_capturer/base.py b/python/sglang/srt/state_capturer/base.py index 6f6b8757841b..4407561eb5ee 100644 --- a/python/sglang/srt/state_capturer/base.py +++ b/python/sglang/srt/state_capturer/base.py @@ -204,10 +204,10 @@ def on_forward_end( if no_copy_to_cpu: # Clone before the next overlapping forward reuses these buffers. return TopkCaptureOutput( - out_cache_loc=forward_batch.out_cache_loc.clone(), + out_cache_loc=forward_batch.out_cache_loc[: slice_gpu.shape[0]].clone(), topk=slice_gpu.clone(), host_cache=self.host_cache, ) - out_cache_loc_cpu = forward_batch.out_cache_loc.cpu() + out_cache_loc_cpu = forward_batch.out_cache_loc[: slice_gpu.shape[0]].cpu() self.host_cache.buffer[out_cache_loc_cpu] = slice_gpu.cpu() return None diff --git a/python/sglang/srt/tp_invariant_ops/__init__.py b/python/sglang/srt/tp_invariant_ops/__init__.py new file mode 100644 index 000000000000..2a3f7a07f4a7 --- /dev/null +++ b/python/sglang/srt/tp_invariant_ops/__init__.py @@ -0,0 +1,23 @@ +from .tp_invariant_ops import ( + disable_tp_invariant_mode, + enable_tp_invariant_mode, + is_tp_invariant_mode_enabled, + matmul_tp_inv, + matmul_tp_persistent, + moe_sum_tree_reduce, + set_tp_invariant_mode, + tree_all_reduce_sum, +) + +__version__ = "0.1.0" + +__all__ = [ + "matmul_tp_persistent", + "matmul_tp_inv", + "moe_sum_tree_reduce", + "tree_all_reduce_sum", + "set_tp_invariant_mode", + "is_tp_invariant_mode_enabled", + "disable_tp_invariant_mode", + "enable_tp_invariant_mode", +] diff --git a/python/sglang/srt/tp_invariant_ops/tp_invariant_ops.py b/python/sglang/srt/tp_invariant_ops/tp_invariant_ops.py new file mode 100644 index 000000000000..6e71bb19e29d --- /dev/null +++ b/python/sglang/srt/tp_invariant_ops/tp_invariant_ops.py @@ -0,0 +1,1939 @@ +import contextlib +import math +import os +import sys +from typing import Any, Callable, Dict + +import torch +import torch.distributed as dist +import triton +import triton.language as tl + +# Triton's constexpr tree unrolling causes deep AST recursion in the JIT +# compiler. The two-level tree (v2) bounds compilation depth to +# max(log2(SUBTREE), log2(E/SUBTREE)) instead of log2(E), but we still +# need headroom for the per-level AST visitor overhead. +if sys.getrecursionlimit() < 16384: + sys.setrecursionlimit(16384) + + +def _matmul_launch_metadata( + grid: Callable[..., Any], kernel: Any, args: Dict[str, Any] +) -> Dict[str, Any]: + ret = {} + m, n, k = args["M"], args["N"], args["K"] + ret["name"] = f"{kernel.name} [M={m}, N={n}, K={k}]" + if "tiles_per_update" in args: + ret["name"] = ( + f"{kernel.name} [M={m}, N={n}, K={k}, tiles_per_update={args['tiles_per_update']:02}]" + ) + if "c_ptr" in args: + bytes_per_elem = args["c_ptr"].element_size() + else: + bytes_per_elem = 1 if args["FP8_OUTPUT"] else 2 + ret[f"flops{bytes_per_elem * 8}"] = 2.0 * m * n * k + ret["bytes"] = bytes_per_elem * (m * k + n * k + m * n) + return ret + + +@triton.jit +def _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M, NUM_SMS): + group_id = tile_id // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (tile_id % group_size_m) + pid_n = (tile_id % num_pid_in_group) // group_size_m + return pid_m, pid_n + + +def _get_tl_dtype(dtype): + if dtype == torch.float32: + return tl.float32 + elif dtype == torch.float16: + return tl.float16 + elif dtype == torch.bfloat16: + return tl.bfloat16 + + +# ---- kernel ---- +@triton.jit(launch_metadata=_matmul_launch_metadata) +def matmul_kernel_tp_persistent( + A_ptr, + B_ptr, + C_ptr, + M, + N, + K, + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + NUM_SMS: tl.constexpr, + LEVEL_K: tl.constexpr, + TILE_K: tl.constexpr, + FIRST_LEVEL_BLOCK: tl.constexpr, + NEXT_POWER_OF_LEVEL: tl.constexpr, + NEXT_POWER_OF_REMAIN_LEVEL: tl.constexpr, + ACC_DTYPE: tl.constexpr, + OUT_DTYPE: tl.constexpr, + A_LARGE: tl.constexpr, + B_LARGE: tl.constexpr, + C_LARGE: tl.constexpr, +): + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_M) + num_pid_n = tl.cdiv(N, BLOCK_N) + num_tiles = num_pid_m * num_pid_n + + num_pid_in_group = GROUP_SIZE_M * num_pid_n + + manual_acc = 3 + acc1 = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_DTYPE) + acc2 = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_DTYPE) + acc3 = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_DTYPE) + + S = tl.zeros((NEXT_POWER_OF_REMAIN_LEVEL, BLOCK_M, BLOCK_N), dtype=ACC_DTYPE) + + S_mask = tl.arange(0, NEXT_POWER_OF_REMAIN_LEVEL)[:, None, None] + level_ids = tl.arange(0, NEXT_POWER_OF_LEVEL) + + base_offs_m = tl.arange(0, BLOCK_M) + base_offs_n = tl.arange(0, BLOCK_N) + offs_k = tl.arange(0, BLOCK_K) + offs_k = tl.max_contiguous(tl.multiple_of(offs_k, BLOCK_K), BLOCK_K) + + for tile_id in tl.range(pid, num_tiles, NUM_SMS, flatten=False): + pid_m, pid_n = _compute_pid( + tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M, NUM_SMS + ) + + start_m = pid_m * BLOCK_M + start_n = pid_n * BLOCK_N + + offs_am = start_m + base_offs_m + if A_LARGE: + offs_am = offs_am.to(tl.int64) + offs_am = tl.where(offs_am < M, offs_am, 0) + + offs_bn = start_n + base_offs_n + if B_LARGE: + offs_bn = offs_bn.to(tl.int64) + offs_bn = tl.where(offs_bn < N, offs_bn, 0) + + offs_am = tl.max_contiguous(tl.multiple_of(offs_am, BLOCK_M), BLOCK_M) + offs_bn = tl.max_contiguous(tl.multiple_of(offs_bn, BLOCK_N), BLOCK_N) + + a_ptrs = A_ptr + offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak + b_ptrs = B_ptr + offs_bn[None, :] * stride_bn + offs_k[:, None] * stride_bk + + count = tl.zeros((NEXT_POWER_OF_LEVEL,), dtype=tl.int32) + + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_DTYPE) + for s_tile_idx in range(0, TILE_K): + k0 = s_tile_idx * BLOCK_K + a = tl.load( + a_ptrs, + mask=(offs_am[:, None] < M) & ((k0 + offs_k)[None, :] < K), + other=0.0, + ) + b = tl.load( + b_ptrs, + mask=((k0 + offs_k)[:, None] < K) & (offs_bn[None, :] < N), + other=0.0, + ) + a_ptrs += BLOCK_K * stride_ak + b_ptrs += BLOCK_K * stride_bk + + acc = tl.dot(a, b).to(ACC_DTYPE) + + break_flag = 0 + for level in range(LEVEL_K): + if break_flag == 0: + idx_mask = level_ids == level + + count_value_added = tl.sum(count * idx_mask) + 1 + + table_value = FIRST_LEVEL_BLOCK if level == 0 else 2 + + carry_over = (table_value == count_value_added).to(tl.int1) + + if count_value_added > 1: + if level == 0: + acc = acc1 + acc + elif level == 1: + acc = acc2 + acc + elif level == 2: + acc = acc3 + acc + else: + tmp_acc_mask = S_mask == (level - manual_acc) + acc = ( + tl.sum(S * tmp_acc_mask, axis=0, dtype=ACC_DTYPE) + acc + ) + + count = tl.where( + idx_mask, count_value_added * (1 - carry_over), count + ) + if not carry_over: + break_flag = 1 + if level == 0: + acc1 = acc + elif level == 1: + acc2 = acc + elif level == 2: + acc3 = acc + else: + tmp_acc_mask = S_mask == (level - manual_acc) + S = tl.where(tmp_acc_mask, acc[None, :, :], S) + + c_ptr = C_ptr + (offs_am[:, None] * stride_cm + offs_bn[None, :] * stride_cn) + offs_cm = start_m + base_offs_m + offs_cn = start_n + base_offs_n + if C_LARGE: + offs_cm = offs_cm.to(tl.int64) + offs_cn = offs_cn.to(tl.int64) + offs_cm = tl.where(offs_cm < M, offs_cm, 0) + offs_cn = tl.where(offs_cn < N, offs_cn, 0) + offs_cm = tl.max_contiguous(tl.multiple_of(offs_cm, BLOCK_M), BLOCK_M) + offs_cn = tl.max_contiguous(tl.multiple_of(offs_cn, BLOCK_N), BLOCK_N) + mask_c = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptr, acc.to(OUT_DTYPE), mask=mask_c) + + +@triton.jit(launch_metadata=_matmul_launch_metadata) +def matmul_kernel_tp_persistent_optim( + A_ptr, + B_ptr, + C_ptr, + M, + N, + K, + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + NUM_SMS: tl.constexpr, + LEVEL_K: tl.constexpr, + TILE_K: tl.constexpr, + FIRST_LEVEL_BLOCK: tl.constexpr, + NEXT_POWER_OF_LEVEL: tl.constexpr, + NEXT_POWER_OF_REMAIN_LEVEL: tl.constexpr, + ACC_DTYPE: tl.constexpr, + OUT_DTYPE: tl.constexpr, + A_LARGE: tl.constexpr, + B_LARGE: tl.constexpr, + C_LARGE: tl.constexpr, +): + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_M) + num_pid_n = tl.cdiv(N, BLOCK_N) + num_tiles = num_pid_m * num_pid_n + + num_pid_in_group = GROUP_SIZE_M * num_pid_n + + # Flat loop (same structure as original) with inlined level-0 handling. + # Level-0 uses scalar counter; tree merge (levels 1+) only on carry. + # Accumulators are conditionally allocated based on LEVEL_K (constexpr) + # to minimize register pressure and maximize occupancy. + acc1 = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_DTYPE) + acc2 = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_DTYPE) + acc3 = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_DTYPE) + acc4 = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_DTYPE) + acc5 = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_DTYPE) + + S = tl.zeros((NEXT_POWER_OF_REMAIN_LEVEL, BLOCK_M, BLOCK_N), dtype=ACC_DTYPE) + S_mask = tl.arange(0, NEXT_POWER_OF_REMAIN_LEVEL)[:, None, None] + level_ids = tl.arange(0, NEXT_POWER_OF_LEVEL) + + offs_k = tl.arange(0, BLOCK_K) + offs_k = tl.max_contiguous(tl.multiple_of(offs_k, BLOCK_K), BLOCK_K) + + for tile_id in tl.range(pid, num_tiles, NUM_SMS, flatten=False): + pid_m, pid_n = _compute_pid( + tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M, NUM_SMS + ) + + start_m = pid_m * BLOCK_M + start_n = pid_n * BLOCK_N + + offs_am = start_m + tl.arange(0, BLOCK_M) + mask_m = offs_am < M + if A_LARGE: + offs_am = offs_am.to(tl.int64) + offs_am = tl.where(mask_m, offs_am, 0) + + offs_bn = start_n + tl.arange(0, BLOCK_N) + mask_n = offs_bn < N + if B_LARGE: + offs_bn = offs_bn.to(tl.int64) + offs_bn = tl.where(mask_n, offs_bn, 0) + + offs_am = tl.max_contiguous(tl.multiple_of(offs_am, BLOCK_M), BLOCK_M) + offs_bn = tl.max_contiguous(tl.multiple_of(offs_bn, BLOCK_N), BLOCK_N) + mask_m_bc = mask_m[:, None] + mask_n_bc = mask_n[None, :] + + a_ptrs = A_ptr + offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak + b_ptrs = B_ptr + offs_bn[None, :] * stride_bn + offs_k[:, None] * stride_bk + + c0 = 0 + c1 = 0 + c2 = 0 + c3 = 0 + c4 = 0 + count = tl.zeros((NEXT_POWER_OF_LEVEL,), dtype=tl.int32) + + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_DTYPE) + for _ in range(0, TILE_K): + a = tl.load(a_ptrs, mask=mask_m_bc, other=0.0) + b = tl.load(b_ptrs, mask=mask_n_bc, other=0.0) + a_ptrs += BLOCK_K * stride_ak + b_ptrs += BLOCK_K * stride_bk + + acc = tl.dot(a, b).to(ACC_DTYPE) + + c0 += 1 + if c0 > 1: + acc = acc1 + acc + if c0 < FIRST_LEVEL_BLOCK: + acc1 = acc + else: + c0 = 0 + break_flag = 0 + for level in range(1, LEVEL_K): + if break_flag == 0: + if level == 1: + c1 += 1 + if c1 > 1: + acc = acc2 + acc + if c1 == 2: + c1 = 0 + else: + acc2 = acc + break_flag = 1 + elif level == 2: + c2 += 1 + if c2 > 1: + acc = acc3 + acc + if c2 == 2: + c2 = 0 + else: + acc3 = acc + break_flag = 1 + elif level == 3: + c3 += 1 + if c3 > 1: + acc = acc4 + acc + if c3 == 2: + c3 = 0 + else: + acc4 = acc + break_flag = 1 + elif level == 4: + c4 += 1 + if c4 > 1: + acc = acc5 + acc + if c4 == 2: + c4 = 0 + else: + acc5 = acc + break_flag = 1 + else: + idx_mask = level_ids == level + count_value_added = tl.sum(count * idx_mask) + 1 + carry_over = (2 == count_value_added).to(tl.int1) + if count_value_added > 1: + tmp_acc_mask = S_mask == (level - 5) + acc = ( + tl.sum(S * tmp_acc_mask, axis=0, dtype=ACC_DTYPE) + + acc + ) + count = tl.where( + idx_mask, count_value_added * (1 - carry_over), count + ) + if not carry_over: + break_flag = 1 + tmp_acc_mask = S_mask == (level - 5) + S = tl.where(tmp_acc_mask, acc[None, :, :], S) + + offs_cm = start_m + tl.arange(0, BLOCK_M) + offs_cn = start_n + tl.arange(0, BLOCK_N) + if C_LARGE: + offs_cm = offs_cm.to(tl.int64) + offs_cn = offs_cn.to(tl.int64) + offs_cm = tl.where(mask_m, offs_cm, 0) + offs_cn = tl.where(mask_n, offs_cn, 0) + offs_cm = tl.max_contiguous(tl.multiple_of(offs_cm, BLOCK_M), BLOCK_M) + offs_cn = tl.max_contiguous(tl.multiple_of(offs_cn, BLOCK_N), BLOCK_N) + c_ptr = C_ptr + (offs_cm[:, None] * stride_cm + offs_cn[None, :] * stride_cn) + mask_c = mask_m_bc & mask_n_bc + tl.store(c_ptr, acc.to(OUT_DTYPE), mask=mask_c) + + +def _matmul_tp_persistent_impl( + A: torch.Tensor, + B: torch.Tensor, + bias: torch.Tensor = None, + fp32_accum: bool = False, + use_optim_kernel: bool = False, +): + assert A.shape[-1] == B.shape[-2], "Dim doesn't match" + + out_dtype = A.dtype + acc_dtype = torch.float32 if fp32_accum else A.dtype + + NUM_SMS = torch.cuda.get_device_properties(A.device).multi_processor_count + + # 1D launch kernel where each block gets its own program. + def grid(META): + return ( + min( + NUM_SMS, + triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]), + ), + ) + + base_configs = { + torch.bfloat16: { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_stages": 2, + "num_warps": 8, + }, + torch.float16: { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_stages": 2, + "num_warps": 8, + }, + torch.float32: { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 8, + "num_stages": 2, + "num_warps": 8, + }, + } + + configs = base_configs + BLOCK_M = configs[out_dtype]["BLOCK_SIZE_M"] + BLOCK_N = configs[out_dtype]["BLOCK_SIZE_N"] + BLOCK_K = configs[out_dtype]["BLOCK_SIZE_K"] + GROUP_SIZE_M = configs[out_dtype]["GROUP_SIZE_M"] + num_stages = configs[out_dtype]["num_stages"] + num_warps = configs[out_dtype]["num_warps"] + + M, K = A.shape + _, N = B.shape + assert K % BLOCK_K == 0, ( + f"Dimension K should be divisible by BLOCK_K. Got K={K}, BLOCK_K={BLOCK_K}." + ) + T = K // BLOCK_K + FIRST_LEVEL_BLOCK = T + + if use_optim_kernel: + num_n_tiles = triton.cdiv(N, BLOCK_N) + total_tiles = triton.cdiv(M, BLOCK_M) * num_n_tiles + + if total_tiles * 4 <= NUM_SMS: + while BLOCK_M > 16 and triton.cdiv(M, BLOCK_M) * num_n_tiles < NUM_SMS: + BLOCK_M //= 2 + elif total_tiles * 2 <= NUM_SMS: + while BLOCK_M > 32 and triton.cdiv(M, BLOCK_M) * num_n_tiles < NUM_SMS: + BLOCK_M //= 2 + + if out_dtype in (torch.bfloat16, torch.float16): + num_warps = 4 if BLOCK_M <= 16 else 8 + num_stages = 3 if K >= 1024 else 2 + + LEVEL_K = 1 + while FIRST_LEVEL_BLOCK > 2 and FIRST_LEVEL_BLOCK % 2 == 0: + FIRST_LEVEL_BLOCK //= 2 + LEVEL_K += 1 + + C = torch.empty((M, N), device=A.device, dtype=out_dtype) + + # Original kernel manually handles levels 0-2 (3 accumulators); + # optim kernel handles levels 0-4 (5 register accumulators), + # so S tensor is smaller / unused for common LEVEL_K <= 5. + manual_acc = 5 if use_optim_kernel else 3 + + NEXT_POWER_OF_LEVEL = 2 ** math.ceil(math.log2(LEVEL_K)) + NEXT_POWER_OF_REMAIN_LEVEL = ( + 2 ** math.ceil(math.log2(LEVEL_K - manual_acc)) if LEVEL_K > manual_acc else 1 + ) + + kernel = ( + matmul_kernel_tp_persistent_optim + if use_optim_kernel + else matmul_kernel_tp_persistent + ) + kernel[grid]( + A, + B, + C, + M, + N, + K, + *A.stride(), + *B.stride(), + *C.stride(), + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + BLOCK_K=BLOCK_K, + GROUP_SIZE_M=GROUP_SIZE_M, + NUM_SMS=NUM_SMS, + NEXT_POWER_OF_REMAIN_LEVEL=NEXT_POWER_OF_REMAIN_LEVEL, + LEVEL_K=LEVEL_K, + TILE_K=T, + FIRST_LEVEL_BLOCK=FIRST_LEVEL_BLOCK, + NEXT_POWER_OF_LEVEL=NEXT_POWER_OF_LEVEL, + ACC_DTYPE=_get_tl_dtype(acc_dtype), + OUT_DTYPE=_get_tl_dtype(out_dtype), + A_LARGE=A.numel() > 2**31, + B_LARGE=B.numel() > 2**31, + C_LARGE=C.numel() > 2**31, + num_warps=num_warps, + num_stages=num_stages, + ) + if bias is not None: + C += bias + return C + + +def matmul_tp_persistent( + A: torch.Tensor, + B: torch.Tensor, + bias: torch.Tensor = None, + fp32_accum: bool = False, +): + return _matmul_tp_persistent_impl( + A=A, + B=B, + bias=bias, + fp32_accum=fp32_accum, + use_optim_kernel=False, + ) + + +def matmul_tp_inv( + A: torch.Tensor, + B: torch.Tensor, + bias: torch.Tensor = None, + fp32_accum: bool = False, +): + return matmul_tp_persistent(A, B, bias=bias, fp32_accum=fp32_accum) + + +def matmul_tp_persistent_optim( + A: torch.Tensor, + B: torch.Tensor, + bias: torch.Tensor = None, + fp32_accum: bool = False, +): + return _matmul_tp_persistent_impl( + A=A, + B=B, + bias=bias, + fp32_accum=fp32_accum, + use_optim_kernel=True, + ) + + +def tree_all_reduce_sum(x: torch.Tensor, device_group=None) -> torch.Tensor: + rank = dist.get_rank(device_group) + world_size = dist.get_world_size(device_group) + + if world_size & (world_size - 1) != 0: + raise ValueError( + "world_size must be a power of 2 in order to use all_reduce_sum." + ) + + result = [torch.zeros_like(x) for _ in range(world_size)] + dist.all_gather(result, x, group=device_group) + + for level in range(1, world_size.bit_length()): + for left in range(0, world_size, 1 << level): + right = left + (1 << (level - 1)) + result[left] += result[right] + + return result[0] + + +def tree_all_reduce_sum_optim(x: torch.Tensor, device_group=None) -> torch.Tensor: + if not x.is_cuda: + raise ValueError("x must be a CUDA tensor.") + if not x.is_contiguous(): + raise ValueError( + "x must be contiguous. Call x = x.contiguous() OUTSIDE graph capture." + ) + + world_size = dist.get_world_size(device_group) + if world_size & (world_size - 1) != 0: + raise ValueError("world_size must be a power of 2.") + + # cache + if not hasattr(tree_all_reduce_sum, "_cache"): + tree_all_reduce_sum._cache = {} + + key = (id(device_group), x.device.index, tuple(x.shape), x.dtype, world_size) + st = tree_all_reduce_sum._cache.get(key) + if st is None: + gather = torch.empty( + (world_size,) + tuple(x.shape), device=x.device, dtype=x.dtype + ) + out = torch.empty_like(x) + st = tree_all_reduce_sum._cache[key] = (gather, out) + + gather, out = st + + # 1) all_gather into one contiguous buffer + dist.all_gather_into_tensor(gather, x, group=device_group) + + # 2) deterministic tree pairing EXACTLY like your original: + # for level in range(1, bit_length): + # for left in range(0, world_size, 1<> 1 + # Views only (no alloc); one add_ kernel per level + gather[0:world_size:step].add_(gather[half:world_size:step]) + + out.copy_(gather[0]) + tree_all_reduce_sum._cache.clear() + return out + + +_tp_inv_MODE = False + +try: + def_lib = torch.library.Library("tp_inv_ops", "DEF") + def_lib.define("matmul_tp_inv(Tensor a, Tensor b, Tensor? bias=None) -> Tensor") +except RuntimeError: + pass + +try: + impl = torch.library.Library("tp_inv_ops", "IMPL") + + impl.impl("matmul_tp_inv", matmul_tp_persistent, "CUDA") +except RuntimeError: + pass + + +def is_tp_invariant_mode_enabled(): + return _tp_inv_MODE + + +def enable_tp_invariant_mode(): + global _tp_inv_MODE + + if _tp_inv_MODE: + return + + _tp_inv_MODE = True + + +def disable_tp_invariant_mode(): + global _tp_inv_MODE + + _tp_inv_MODE = False + + +@contextlib.contextmanager +def set_tp_invariant_mode(enabled=True): + global _tp_inv_MODE + + old_state = _tp_inv_MODE + + if enabled: + enable_tp_invariant_mode() + else: + disable_tp_invariant_mode() + + try: + yield + finally: + _tp_inv_MODE = old_state + + +def scatter_input_by_local_expert( + topk: torch.Tensor, input: torch.Tensor, E: int +) -> torch.Tensor: + """ + Args: + topk: [M, topk], long, -1 means remote expert + input: [M, topk, hidden_size], float + E: int, number of local experts (expert ids in [0, E)) + Returns: + output: [M, E, hidden_size] + """ + M, _, hidden_size = input.shape + + # Mask out remote experts in output + valid = (topk != -1).unsqueeze(-1) # [M, topk, 1] + output_masked = input * valid.to(input.dtype) # [M, topk, hidden_size] + + # Replace -1 with 0 for safe indexing (value doesn't matter because output is zero) + topk_index = topk.clamp(min=0) # turns -1 into 0, leaves others unchanged + + # Expand index to match output + index = topk_index.unsqueeze(-1).expand( + -1, -1, hidden_size + ) # [M, topk, hidden_size] + + # Initialize result + output = torch.zeros(M, E, hidden_size, device=input.device, dtype=input.dtype) + + # Scatter add + output.scatter_add_(1, index, output_masked) + + return output + + +@triton.jit +def _load_expert_tile( + input_ptr, + input_base, + input_stride_1, + topk_ids_base, + topk_ids_stride_1, + zero_ptrs, + offs_dim, + mask, + mask_token, + e: tl.constexpr, + TOPK: tl.constexpr, + BLOCK_M: tl.constexpr, +): + """Load tile for expert e: search topk_ids for slot, redirect to zero_buf if remote.""" + found_slot = tl.full([BLOCK_M], -1, dtype=tl.int32) + for k in range(TOPK): + kid = tl.load( + topk_ids_base + k * topk_ids_stride_1, + mask=mask_token, + other=-1, + ).to(tl.int32) + match = (kid == e) & (found_slot == -1) + found_slot = tl.where(match, k, found_slot) + + is_valid = found_slot != -1 + slot_safe = tl.maximum(found_slot, 0) + input_ptrs = input_base + slot_safe[:, None] * input_stride_1 + offs_dim[None, :] + load_ptrs = tl.where(is_valid[:, None], input_ptrs, zero_ptrs) + return tl.load(load_ptrs, mask=mask, other=0.0).to(tl.float32) + + +@triton.jit +def _tree_reduce_pair( + input_ptr, + input_base, + input_stride_1, + topk_ids_base, + topk_ids_stride_1, + zero_ptrs, + offs_dim, + mask, + mask_token, + start: tl.constexpr, + size: tl.constexpr, + TOPK: tl.constexpr, + BLOCK_M: tl.constexpr, +): + """ + Recursively compute binary-tree sum over experts [start, start+size). + Returns a fp32 tile [BLOCK_M, BLOCK_DIM]. + + Tree structure (e.g. size=4, start=0): + _tree_reduce_pair(0, 4) + = _tree_reduce_pair(0, 2) + _tree_reduce_pair(2, 2) + = (_tree_reduce_pair(0,1) + _tree_reduce_pair(1,1)) + + (_tree_reduce_pair(2,1) + _tree_reduce_pair(3,1)) + = (load(e0) + load(e1)) + (load(e2) + load(e3)) + + Since size is constexpr and always a power of 2, Triton fully unrolls this + into a fixed sequence of loads and adds with no dynamic branching. + Max register depth = log2(E) tiles, e.g. E=64 -> 6 tiles. + """ + if size == 1: + return _load_expert_tile( + input_ptr, + input_base, + input_stride_1, + topk_ids_base, + topk_ids_stride_1, + zero_ptrs, + offs_dim, + mask, + mask_token, + start, + TOPK, + BLOCK_M, + ) + else: + half: tl.constexpr = size // 2 + left = _tree_reduce_pair( + input_ptr, + input_base, + input_stride_1, + topk_ids_base, + topk_ids_stride_1, + zero_ptrs, + offs_dim, + mask, + mask_token, + start, + half, + TOPK, + BLOCK_M, + ) + right = _tree_reduce_pair( + input_ptr, + input_base, + input_stride_1, + topk_ids_base, + topk_ids_stride_1, + zero_ptrs, + offs_dim, + mask, + mask_token, + start + half, + half, + TOPK, + BLOCK_M, + ) + return left + right + + +@triton.jit +def _fused_tree_reduce_kernel( + # input: [M, topk, hidden_dim], contiguous + input_ptr, + input_stride_0, + input_stride_1, + # topk_ids: [M, topk], -1 means remote + topk_ids_ptr, + topk_ids_stride_0, + topk_ids_stride_1, + # zero_buf: [hidden_dim], all zeros + zero_buf_ptr, + # output: [M, hidden_dim] + output_ptr, + output_stride_0, + # scalars + token_num, + hidden_dim, + routed_scaling_factor, + # constexpr + E: tl.constexpr, + TOPK: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_DIM: tl.constexpr, +): + """ + Fused scatter + binary-tree reduce in a single kernel. + Zero extra memory: all intermediate results live in registers. + + For each token block, loads expert tiles on-the-fly (with zero-buf redirect + for remote experts), and reduces them in binary-tree order using recursive + constexpr unrolling. The tree structure is: + result = tree_sum(0, E) + tree_sum(s, n) = tree_sum(s, n/2) + tree_sum(s+n/2, n/2) if n > 1 + tree_sum(s, 1) = load_expert(s) base case + + Register pressure: log2(E) tiles of [BLOCK_M, BLOCK_DIM] fp32. + E.g. E=64, BLOCK_M=1, BLOCK_DIM=2048 -> 6 * 8KB = 48KB, well within limits. + """ + input_stride_0 = tl.cast(input_stride_0, dtype=tl.int64) + input_stride_1 = tl.cast(input_stride_1, dtype=tl.int64) + topk_ids_stride_0 = tl.cast(topk_ids_stride_0, dtype=tl.int64) + topk_ids_stride_1 = tl.cast(topk_ids_stride_1, dtype=tl.int64) + output_stride_0 = tl.cast(output_stride_0, dtype=tl.int64) + + token_block_id = tl.program_id(0) + dim_block_id = tl.program_id(1) + + offs_token = token_block_id * BLOCK_M + tl.arange(0, BLOCK_M) + offs_dim = dim_block_id * BLOCK_DIM + tl.arange(0, BLOCK_DIM) + + mask_token = offs_token < token_num + mask_dim = offs_dim < hidden_dim + mask = mask_token[:, None] & mask_dim[None, :] + + zero_ptrs = zero_buf_ptr + offs_dim[None, :] + input_base = input_ptr + offs_token[:, None] * input_stride_0 + topk_ids_base = topk_ids_ptr + offs_token * topk_ids_stride_0 + + # Binary tree reduce over all E experts, entirely in registers. + result = _tree_reduce_pair( + input_ptr, + input_base, + input_stride_1, + topk_ids_base, + topk_ids_stride_1, + zero_ptrs, + offs_dim, + mask, + mask_token, + 0, + E, + TOPK, + BLOCK_M, + ) + + result *= routed_scaling_factor + + store_ptrs = output_ptr + offs_token[:, None] * output_stride_0 + offs_dim[None, :] + tl.store(store_ptrs, result.to(input_ptr.dtype.element_ty), mask=mask) + + +# Persistent zero buffer: only [hidden_dim] elements, negligible memory. +# Allocated once on first call, never freed. +_zero_buf_cache: torch.Tensor | None = None + + +def moe_sum_tree_reduce_v1( + input: torch.Tensor, # [M, topk, hidden_dim] + output: torch.Tensor, # [M, hidden_dim] + curr_topk_ids: torch.Tensor, # [M, topk], -1 means remote + routed_scaling_factor: float, + E: int, +): + """ + Fused MoE tree reduce: zero extra memory, CUDA Graph safe. + + Single kernel: loads expert tiles on-the-fly with zero-buf pointer redirect + for remote experts (L1 cache hit), reduces in binary-tree order entirely + in registers. No scratch buffer needed. + + Invariant guarantee: binary tree reduce order is fixed by expert id, + identical across all EP ranks regardless of which experts are local. + + Memory overhead: only a single [hidden_dim] zero buffer (~14KB for H=7168 bf16). + """ + assert input.is_contiguous() + assert output.is_contiguous() + + token_num, topk, hidden_dim = input.shape + assert output.shape[0] == token_num and output.shape[1] == hidden_dim + assert (E & (E - 1)) == 0, f"E must be power of 2, got {E}" + + # Fast path for the dominant K=8 case: avoids generic expert-tree loads/scans. + if topk == 8: + _moe_sum_tree_reduce_k8_fast_path( + input=input, + output=output, + curr_topk_ids=curr_topk_ids, + routed_scaling_factor=routed_scaling_factor, + E=E, + ) + return + + # Zero buffer: [hidden_dim], allocated once, reused forever + global _zero_buf_cache + if ( + _zero_buf_cache is None + or _zero_buf_cache.device != input.device + or _zero_buf_cache.dtype != input.dtype + or _zero_buf_cache.numel() < hidden_dim + ): + _zero_buf_cache = torch.zeros( + hidden_dim, device=input.device, dtype=input.dtype + ) + zero_buf = _zero_buf_cache + + BLOCK_M = 1 + BLOCK_DIM = 2048 + num_warps = 16 + + grid = ( + triton.cdiv(token_num, BLOCK_M), + triton.cdiv(hidden_dim, BLOCK_DIM), + ) + + _fused_tree_reduce_kernel[grid]( + input, + input.stride(0), + input.stride(1), + curr_topk_ids, + curr_topk_ids.stride(0), + curr_topk_ids.stride(1), + zero_buf, + output, + output.stride(0), + token_num=token_num, + hidden_dim=hidden_dim, + routed_scaling_factor=routed_scaling_factor, + E=E, + TOPK=topk, + BLOCK_M=BLOCK_M, + BLOCK_DIM=BLOCK_DIM, + num_warps=num_warps, + ) + return + + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _moe_sum_tree_reduce_k8_fused_kernel_opt2d( + x_ptr, + ids_ptr, + out_ptr, + sx_m: tl.constexpr, + sx_k: tl.constexpr, + sx_h: tl.constexpr, # x: [M,8,H] + sid_m: tl.constexpr, + sid_k: tl.constexpr, # ids: [M,8] + so_m: tl.constexpr, + so_h: tl.constexpr, # out: [M,H] + M, + H, # runtime OK + E_LEVEL, # runtime int (log2(E)) + routed_scaling_factor, # runtime scalar + BLOCK_M: tl.constexpr, + BLOCK_H: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_h = tl.program_id(1) + + m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + h = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) + + mask_m = m < M + mask_h = h < H + mask_mh = mask_m[:, None] & mask_h[None, :] + + # ---- load ids (int32), -1 means remote ---- + ids0 = tl.load(ids_ptr + m * sid_m + 0 * sid_k, mask=mask_m, other=-1).to(tl.int32) + ids1 = tl.load(ids_ptr + m * sid_m + 1 * sid_k, mask=mask_m, other=-1).to(tl.int32) + ids2 = tl.load(ids_ptr + m * sid_m + 2 * sid_k, mask=mask_m, other=-1).to(tl.int32) + ids3 = tl.load(ids_ptr + m * sid_m + 3 * sid_k, mask=mask_m, other=-1).to(tl.int32) + ids4 = tl.load(ids_ptr + m * sid_m + 4 * sid_k, mask=mask_m, other=-1).to(tl.int32) + ids5 = tl.load(ids_ptr + m * sid_m + 5 * sid_k, mask=mask_m, other=-1).to(tl.int32) + ids6 = tl.load(ids_ptr + m * sid_m + 6 * sid_k, mask=mask_m, other=-1).to(tl.int32) + ids7 = tl.load(ids_ptr + m * sid_m + 7 * sid_k, mask=mask_m, other=-1).to(tl.int32) + + # h comes from tl.arange and can be annotated for alignment. + tl.multiple_of(h, 8) + tl.max_contiguous(h, BLOCK_H) + + # m is also a tensor. + tl.multiple_of(m, BLOCK_M) + + # ---- load values: remote handled by mask->other=0 (NO zero_tile / NO tl.where big tiles) ---- + m0 = mask_mh & (ids0 != -1)[:, None] + m1 = mask_mh & (ids1 != -1)[:, None] + m2 = mask_mh & (ids2 != -1)[:, None] + m3 = mask_mh & (ids3 != -1)[:, None] + m4 = mask_mh & (ids4 != -1)[:, None] + m5 = mask_mh & (ids5 != -1)[:, None] + m6 = mask_mh & (ids6 != -1)[:, None] + m7 = mask_mh & (ids7 != -1)[:, None] + + v0 = tl.load( + x_ptr + m[:, None] * sx_m + 0 * sx_k + h[None, :] * sx_h, mask=m0, other=0.0 + ) + v1 = tl.load( + x_ptr + m[:, None] * sx_m + 1 * sx_k + h[None, :] * sx_h, mask=m1, other=0.0 + ) + v2 = tl.load( + x_ptr + m[:, None] * sx_m + 2 * sx_k + h[None, :] * sx_h, mask=m2, other=0.0 + ) + v3 = tl.load( + x_ptr + m[:, None] * sx_m + 3 * sx_k + h[None, :] * sx_h, mask=m3, other=0.0 + ) + v4 = tl.load( + x_ptr + m[:, None] * sx_m + 4 * sx_k + h[None, :] * sx_h, mask=m4, other=0.0 + ) + v5 = tl.load( + x_ptr + m[:, None] * sx_m + 5 * sx_k + h[None, :] * sx_h, mask=m5, other=0.0 + ) + v6 = tl.load( + x_ptr + m[:, None] * sx_m + 6 * sx_k + h[None, :] * sx_h, mask=m6, other=0.0 + ) + v7 = tl.load( + x_ptr + m[:, None] * sx_m + 7 * sx_k + h[None, :] * sx_h, mask=m7, other=0.0 + ) + + x_dtype = x_ptr.dtype.element_ty + + # ---- deterministic dense-tree-equivalent reduce (same order concept as baseline) ---- + # Remote entries have already been masked to 0.0 at load time (m0..m7). + + for bit in tl.range(0, E_LEVEL): + bitmask = 1 << bit + + # ========== lane0 as source ========== + cond = (ids0 != -1) & ((ids0 & bitmask) != 0) + target = ids0 ^ bitmask + mm1 = cond & (ids1 == target) + mm2 = cond & (ids2 == target) + mm3 = cond & (ids3 == target) + mm4 = cond & (ids4 == target) + mm5 = cond & (ids5 == target) + mm6 = cond & (ids6 == target) + mm7 = cond & (ids7 == target) + hit = mm1 | mm2 | mm3 | mm4 | mm5 | mm6 | mm7 + + src = v0.to(tl.float32) + v1 = tl.where(mm1[:, None], (v1.to(tl.float32) + src).to(x_dtype), v1) + v2 = tl.where(mm2[:, None], (v2.to(tl.float32) + src).to(x_dtype), v2) + v3 = tl.where(mm3[:, None], (v3.to(tl.float32) + src).to(x_dtype), v3) + v4 = tl.where(mm4[:, None], (v4.to(tl.float32) + src).to(x_dtype), v4) + v5 = tl.where(mm5[:, None], (v5.to(tl.float32) + src).to(x_dtype), v5) + v6 = tl.where(mm6[:, None], (v6.to(tl.float32) + src).to(x_dtype), v6) + v7 = tl.where(mm7[:, None], (v7.to(tl.float32) + src).to(x_dtype), v7) + + v0 = tl.where(hit[:, None], 0.0, v0) + ids0 = tl.where(hit, -1, ids0) + ids0 = tl.where(cond & (~hit), target, ids0) + + # ========== lane1 as source ========== + cond = (ids1 != -1) & ((ids1 & bitmask) != 0) + target = ids1 ^ bitmask + mm0 = cond & (ids0 == target) + mm2 = cond & (ids2 == target) + mm3 = cond & (ids3 == target) + mm4 = cond & (ids4 == target) + mm5 = cond & (ids5 == target) + mm6 = cond & (ids6 == target) + mm7 = cond & (ids7 == target) + hit = mm0 | mm2 | mm3 | mm4 | mm5 | mm6 | mm7 + + src = v1.to(tl.float32) + v0 = tl.where(mm0[:, None], (v0.to(tl.float32) + src).to(x_dtype), v0) + v2 = tl.where(mm2[:, None], (v2.to(tl.float32) + src).to(x_dtype), v2) + v3 = tl.where(mm3[:, None], (v3.to(tl.float32) + src).to(x_dtype), v3) + v4 = tl.where(mm4[:, None], (v4.to(tl.float32) + src).to(x_dtype), v4) + v5 = tl.where(mm5[:, None], (v5.to(tl.float32) + src).to(x_dtype), v5) + v6 = tl.where(mm6[:, None], (v6.to(tl.float32) + src).to(x_dtype), v6) + v7 = tl.where(mm7[:, None], (v7.to(tl.float32) + src).to(x_dtype), v7) + + v1 = tl.where(hit[:, None], 0.0, v1) + ids1 = tl.where(hit, -1, ids1) + ids1 = tl.where(cond & (~hit), target, ids1) + + # ========== lane2 as source ========== + cond = (ids2 != -1) & ((ids2 & bitmask) != 0) + target = ids2 ^ bitmask + mm0 = cond & (ids0 == target) + mm1 = cond & (ids1 == target) + mm3 = cond & (ids3 == target) + mm4 = cond & (ids4 == target) + mm5 = cond & (ids5 == target) + mm6 = cond & (ids6 == target) + mm7 = cond & (ids7 == target) + hit = mm0 | mm1 | mm3 | mm4 | mm5 | mm6 | mm7 + + src = v2.to(tl.float32) + v0 = tl.where(mm0[:, None], (v0.to(tl.float32) + src).to(x_dtype), v0) + v1 = tl.where(mm1[:, None], (v1.to(tl.float32) + src).to(x_dtype), v1) + v3 = tl.where(mm3[:, None], (v3.to(tl.float32) + src).to(x_dtype), v3) + v4 = tl.where(mm4[:, None], (v4.to(tl.float32) + src).to(x_dtype), v4) + v5 = tl.where(mm5[:, None], (v5.to(tl.float32) + src).to(x_dtype), v5) + v6 = tl.where(mm6[:, None], (v6.to(tl.float32) + src).to(x_dtype), v6) + v7 = tl.where(mm7[:, None], (v7.to(tl.float32) + src).to(x_dtype), v7) + + v2 = tl.where(hit[:, None], 0.0, v2) + ids2 = tl.where(hit, -1, ids2) + ids2 = tl.where(cond & (~hit), target, ids2) + + # ========== lane3 as source ========== + cond = (ids3 != -1) & ((ids3 & bitmask) != 0) + target = ids3 ^ bitmask + mm0 = cond & (ids0 == target) + mm1 = cond & (ids1 == target) + mm2 = cond & (ids2 == target) + mm4 = cond & (ids4 == target) + mm5 = cond & (ids5 == target) + mm6 = cond & (ids6 == target) + mm7 = cond & (ids7 == target) + hit = mm0 | mm1 | mm2 | mm4 | mm5 | mm6 | mm7 + + src = v3.to(tl.float32) + v0 = tl.where(mm0[:, None], (v0.to(tl.float32) + src).to(x_dtype), v0) + v1 = tl.where(mm1[:, None], (v1.to(tl.float32) + src).to(x_dtype), v1) + v2 = tl.where(mm2[:, None], (v2.to(tl.float32) + src).to(x_dtype), v2) + v4 = tl.where(mm4[:, None], (v4.to(tl.float32) + src).to(x_dtype), v4) + v5 = tl.where(mm5[:, None], (v5.to(tl.float32) + src).to(x_dtype), v5) + v6 = tl.where(mm6[:, None], (v6.to(tl.float32) + src).to(x_dtype), v6) + v7 = tl.where(mm7[:, None], (v7.to(tl.float32) + src).to(x_dtype), v7) + + v3 = tl.where(hit[:, None], 0.0, v3) + ids3 = tl.where(hit, -1, ids3) + ids3 = tl.where(cond & (~hit), target, ids3) + + # ========== lane4 as source ========== + cond = (ids4 != -1) & ((ids4 & bitmask) != 0) + target = ids4 ^ bitmask + mm0 = cond & (ids0 == target) + mm1 = cond & (ids1 == target) + mm2 = cond & (ids2 == target) + mm3 = cond & (ids3 == target) + mm5 = cond & (ids5 == target) + mm6 = cond & (ids6 == target) + mm7 = cond & (ids7 == target) + hit = mm0 | mm1 | mm2 | mm3 | mm5 | mm6 | mm7 + + src = v4.to(tl.float32) + v0 = tl.where(mm0[:, None], (v0.to(tl.float32) + src).to(x_dtype), v0) + v1 = tl.where(mm1[:, None], (v1.to(tl.float32) + src).to(x_dtype), v1) + v2 = tl.where(mm2[:, None], (v2.to(tl.float32) + src).to(x_dtype), v2) + v3 = tl.where(mm3[:, None], (v3.to(tl.float32) + src).to(x_dtype), v3) + v5 = tl.where(mm5[:, None], (v5.to(tl.float32) + src).to(x_dtype), v5) + v6 = tl.where(mm6[:, None], (v6.to(tl.float32) + src).to(x_dtype), v6) + v7 = tl.where(mm7[:, None], (v7.to(tl.float32) + src).to(x_dtype), v7) + + v4 = tl.where(hit[:, None], 0.0, v4) + ids4 = tl.where(hit, -1, ids4) + ids4 = tl.where(cond & (~hit), target, ids4) + + # ========== lane5 as source ========== + cond = (ids5 != -1) & ((ids5 & bitmask) != 0) + target = ids5 ^ bitmask + mm0 = cond & (ids0 == target) + mm1 = cond & (ids1 == target) + mm2 = cond & (ids2 == target) + mm3 = cond & (ids3 == target) + mm4 = cond & (ids4 == target) + mm6 = cond & (ids6 == target) + mm7 = cond & (ids7 == target) + hit = mm0 | mm1 | mm2 | mm3 | mm4 | mm6 | mm7 + + src = v5.to(tl.float32) + v0 = tl.where(mm0[:, None], (v0.to(tl.float32) + src).to(x_dtype), v0) + v1 = tl.where(mm1[:, None], (v1.to(tl.float32) + src).to(x_dtype), v1) + v2 = tl.where(mm2[:, None], (v2.to(tl.float32) + src).to(x_dtype), v2) + v3 = tl.where(mm3[:, None], (v3.to(tl.float32) + src).to(x_dtype), v3) + v4 = tl.where(mm4[:, None], (v4.to(tl.float32) + src).to(x_dtype), v4) + v6 = tl.where(mm6[:, None], (v6.to(tl.float32) + src).to(x_dtype), v6) + v7 = tl.where(mm7[:, None], (v7.to(tl.float32) + src).to(x_dtype), v7) + + v5 = tl.where(hit[:, None], 0.0, v5) + ids5 = tl.where(hit, -1, ids5) + ids5 = tl.where(cond & (~hit), target, ids5) + + # ========== lane6 as source ========== + cond = (ids6 != -1) & ((ids6 & bitmask) != 0) + target = ids6 ^ bitmask + mm0 = cond & (ids0 == target) + mm1 = cond & (ids1 == target) + mm2 = cond & (ids2 == target) + mm3 = cond & (ids3 == target) + mm4 = cond & (ids4 == target) + mm5 = cond & (ids5 == target) + mm7 = cond & (ids7 == target) + hit = mm0 | mm1 | mm2 | mm3 | mm4 | mm5 | mm7 + + src = v6.to(tl.float32) + v0 = tl.where(mm0[:, None], (v0.to(tl.float32) + src).to(x_dtype), v0) + v1 = tl.where(mm1[:, None], (v1.to(tl.float32) + src).to(x_dtype), v1) + v2 = tl.where(mm2[:, None], (v2.to(tl.float32) + src).to(x_dtype), v2) + v3 = tl.where(mm3[:, None], (v3.to(tl.float32) + src).to(x_dtype), v3) + v4 = tl.where(mm4[:, None], (v4.to(tl.float32) + src).to(x_dtype), v4) + v5 = tl.where(mm5[:, None], (v5.to(tl.float32) + src).to(x_dtype), v5) + v7 = tl.where(mm7[:, None], (v7.to(tl.float32) + src).to(x_dtype), v7) + + v6 = tl.where(hit[:, None], 0.0, v6) + ids6 = tl.where(hit, -1, ids6) + ids6 = tl.where(cond & (~hit), target, ids6) + + # ========== lane7 as source ========== + cond = (ids7 != -1) & ((ids7 & bitmask) != 0) + target = ids7 ^ bitmask + mm0 = cond & (ids0 == target) + mm1 = cond & (ids1 == target) + mm2 = cond & (ids2 == target) + mm3 = cond & (ids3 == target) + mm4 = cond & (ids4 == target) + mm5 = cond & (ids5 == target) + mm6 = cond & (ids6 == target) + hit = mm0 | mm1 | mm2 | mm3 | mm4 | mm5 | mm6 + + src = v7.to(tl.float32) + v0 = tl.where(mm0[:, None], (v0.to(tl.float32) + src).to(x_dtype), v0) + v1 = tl.where(mm1[:, None], (v1.to(tl.float32) + src).to(x_dtype), v1) + v2 = tl.where(mm2[:, None], (v2.to(tl.float32) + src).to(x_dtype), v2) + v3 = tl.where(mm3[:, None], (v3.to(tl.float32) + src).to(x_dtype), v3) + v4 = tl.where(mm4[:, None], (v4.to(tl.float32) + src).to(x_dtype), v4) + v5 = tl.where(mm5[:, None], (v5.to(tl.float32) + src).to(x_dtype), v5) + v6 = tl.where(mm6[:, None], (v6.to(tl.float32) + src).to(x_dtype), v6) + + v7 = tl.where(hit[:, None], 0.0, v7) + ids7 = tl.where(hit, -1, ids7) + ids7 = tl.where(cond & (~hit), target, ids7) + + # ---- final: since ids unique per token, bucket0 is unique; compute directly in fp32 (no out_tile chain) ---- + acc = tl.zeros((BLOCK_M, BLOCK_H), dtype=tl.float32) + acc += tl.where((ids0 == 0)[:, None], v0.to(tl.float32), 0.0) + acc += tl.where((ids1 == 0)[:, None], v1.to(tl.float32), 0.0) + acc += tl.where((ids2 == 0)[:, None], v2.to(tl.float32), 0.0) + acc += tl.where((ids3 == 0)[:, None], v3.to(tl.float32), 0.0) + acc += tl.where((ids4 == 0)[:, None], v4.to(tl.float32), 0.0) + acc += tl.where((ids5 == 0)[:, None], v5.to(tl.float32), 0.0) + acc += tl.where((ids6 == 0)[:, None], v6.to(tl.float32), 0.0) + acc += tl.where((ids7 == 0)[:, None], v7.to(tl.float32), 0.0) + + acc *= routed_scaling_factor + + out_ptrs = out_ptr + m[:, None] * so_m + h[None, :] * so_h + tl.store(out_ptrs, acc.to(out_ptr.dtype.element_ty), mask=mask_mh) + + +def _moe_sum_tree_reduce_k8_fast_path( + input: torch.Tensor, # [M, 8, H] + output: torch.Tensor, # [M, H] + curr_topk_ids: torch.Tensor, # [M, 8], -1 means remote + routed_scaling_factor: float, + E: int, +): + assert input.is_contiguous() + assert output.is_contiguous() + assert curr_topk_ids.is_contiguous() + M, K, H = input.shape + assert K == 8 + assert output.shape == (M, H) + assert (E & (E - 1)) == 0 + + E_LEVEL = int(math.log2(E)) + + # K=8 specialization: use wider H tiles for better memory throughput. + if H >= 4096: + BLOCK_M = 8 + BLOCK_H = 128 + num_warps = 4 + elif H >= 2048: + BLOCK_M = 8 + BLOCK_H = 256 + num_warps = 8 + else: + BLOCK_M = 8 + BLOCK_H = 128 + num_warps = 4 + + grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(H, BLOCK_H)) + _moe_sum_tree_reduce_k8_fused_kernel_opt2d[grid]( + input, + curr_topk_ids, + output, + input.stride(0), + input.stride(1), + input.stride(2), + curr_topk_ids.stride(0), + curr_topk_ids.stride(1), + output.stride(0), + output.stride(1), + M, + H, + E_LEVEL, + routed_scaling_factor, + BLOCK_M=BLOCK_M, + BLOCK_H=BLOCK_H, + num_warps=num_warps, + ) + return output + + +@triton.jit +def _moe_sum_tree_reduce_topk16_sparse_kernel( + x_ptr, + ids_ptr, + out_ptr, + sx_m, + sx_k, + sx_h, # x: [M, K, H] + sid_m, + sid_k, # ids: [M, K] + so_m, + so_h, # out: [M, H] + M, + K, + H, # runtime + E_LEVEL, # log2(E) + routed_scaling_factor, + BLOCK_H: tl.constexpr, + MAX_TOPK: tl.constexpr, # fixed compile-time capacity, e.g. 16 +): + # One program handles one token + one hidden tile. + pid_m = tl.program_id(0) + pid_h = tl.program_id(1) + + m = pid_m + h = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) + + mask_m = m < M + mask_h = h < H + mask_mh = mask_m & mask_h + + slot = tl.arange(0, MAX_TOPK) + in_k = slot < K + + ids = tl.load( + ids_ptr + m * sid_m + slot * sid_k, mask=(mask_m & in_k), other=-1 + ).to(tl.int32) + valid = (ids != -1) & in_k & mask_m + + vals = tl.zeros((MAX_TOPK, BLOCK_H), dtype=tl.float32) + for s in range(MAX_TOPK): + vmask = mask_h & valid[s] + v = tl.load( + x_ptr + m * sx_m + s * sx_k + h * sx_h, + mask=vmask, + other=0.0, + ).to(tl.float32) + vals = tl.where((slot == s)[:, None], v[None, :], vals) + + idx = tl.arange(0, MAX_TOPK) + for bit in tl.range(0, E_LEVEL): + bitmask = 1 << bit + + for s in range(MAX_TOPK): + id_s = ids[s] + cond = (id_s != -1) & ((id_s & bitmask) != 0) + target = id_s ^ bitmask + + match = (ids == target) & (idx != s) + # 1-based index; 0 means no match. + match_pos1 = tl.max(tl.where(match, idx + 1, 0)) + has = match_pos1 > 0 + j = match_pos1 - 1 + + src = vals[s, :] + hit_mask = idx == j + vals = vals + tl.where(hit_mask[:, None] & has, src[None, :], 0.0) + vals = tl.where((idx == s)[:, None] & has, 0.0, vals) + ids = tl.where((idx == s) & has, -1, ids) + + # No partner found at this level: update bucket id only. + ids = tl.where((idx == s) & cond & (~has), target, ids) + + acc = tl.zeros((BLOCK_H,), dtype=tl.float32) + for s in range(MAX_TOPK): + acc += tl.where(ids[s] == 0, vals[s, :], 0.0) + + acc *= routed_scaling_factor + tl.store( + out_ptr + m * so_m + h * so_h, acc.to(out_ptr.dtype.element_ty), mask=mask_mh + ) + + +def moe_sum_tree_reduce_v1_topk_sparse16( + input: torch.Tensor, # [M, K, H] + output: torch.Tensor, # [M, H] + curr_topk_ids: torch.Tensor, # [M, K], -1 means remote + routed_scaling_factor: float, + E: int, + max_topk: int = 16, +): + """ + Supplemental kernel path for topk != 8: + - optimized for small topk (<=16) + - keeps deterministic tree-equivalent reduction semantics + - does not replace existing v1 entrypoint automatically + """ + assert input.is_contiguous() + assert output.is_contiguous() + assert curr_topk_ids.is_contiguous() + M, K, H = input.shape + assert output.shape == (M, H) + assert K <= max_topk, f"K={K} > max_topk={max_topk}" + assert (E & (E - 1)) == 0, "E must be power of 2" + + E_LEVEL = int(math.log2(E)) + BLOCK_H = 128 if H >= 4096 else 256 + num_warps = 4 if BLOCK_H == 128 else 8 + + grid = (M, triton.cdiv(H, BLOCK_H)) + _moe_sum_tree_reduce_topk16_sparse_kernel[grid]( + input, + curr_topk_ids, + output, + input.stride(0), + input.stride(1), + input.stride(2), + curr_topk_ids.stride(0), + curr_topk_ids.stride(1), + output.stride(0), + output.stride(1), + M, + K, + H, + E_LEVEL, + routed_scaling_factor, + BLOCK_H=BLOCK_H, + MAX_TOPK=max_topk, + num_warps=num_warps, + ) + return output + + +def moe_sum_tree_reduce_v0( + input: torch.Tensor, # [M, 8, H] + output: torch.Tensor, # [M, H] + curr_topk_ids: torch.Tensor, # [M, 8], -1 means remote + routed_scaling_factor: float, + E: int, +): + # Backward-compatible alias. + return _moe_sum_tree_reduce_k8_fast_path( + input=input, + output=output, + curr_topk_ids=curr_topk_ids, + routed_scaling_factor=routed_scaling_factor, + E=E, + ) + + +@triton.jit +def _moe_tree_reduce_sparse_k8_kernel( + x_ptr, + ids_ptr, + out_ptr, + sx_m: tl.constexpr, + sx_k: tl.constexpr, + sx_h: tl.constexpr, + sid_m: tl.constexpr, + sid_k: tl.constexpr, + so_m: tl.constexpr, + so_h: tl.constexpr, + M, + H, + routed_scaling_factor, + LOGE, # runtime: tl.range does not unroll + CAST_MODE: tl.constexpr, # 0: per-level bf16 round, 1: no intermediate cast + BLOCK_M: tl.constexpr, + BLOCK_H: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_h = tl.program_id(1) + + m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + h = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) + + tl.multiple_of(m, BLOCK_M) + tl.multiple_of(h, 8) + tl.max_contiguous(h, BLOCK_H) + + mask_m = m < M + mask_h = h < H + mask_mh = mask_m[:, None] & mask_h[None, :] + + x_ty = x_ptr.dtype.element_ty + NEG2 = -2 # sentinel: topk_ids cannot be -2 + + # ---- ids ---- + ids0 = tl.load(ids_ptr + m * sid_m + 0 * sid_k, mask=mask_m, other=-1).to(tl.int32) + ids1 = tl.load(ids_ptr + m * sid_m + 1 * sid_k, mask=mask_m, other=-1).to(tl.int32) + ids2 = tl.load(ids_ptr + m * sid_m + 2 * sid_k, mask=mask_m, other=-1).to(tl.int32) + ids3 = tl.load(ids_ptr + m * sid_m + 3 * sid_k, mask=mask_m, other=-1).to(tl.int32) + ids4 = tl.load(ids_ptr + m * sid_m + 4 * sid_k, mask=mask_m, other=-1).to(tl.int32) + ids5 = tl.load(ids_ptr + m * sid_m + 5 * sid_k, mask=mask_m, other=-1).to(tl.int32) + ids6 = tl.load(ids_ptr + m * sid_m + 6 * sid_k, mask=mask_m, other=-1).to(tl.int32) + ids7 = tl.load(ids_ptr + m * sid_m + 7 * sid_k, mask=mask_m, other=-1).to(tl.int32) + + # ---- vals: load -> fp32 ---- + f0 = tl.load( + x_ptr + m[:, None] * sx_m + 0 * sx_k + h[None, :] * sx_h, + mask=mask_mh & (ids0 != -1)[:, None], + other=0.0, + ).to(tl.float32) + f1 = tl.load( + x_ptr + m[:, None] * sx_m + 1 * sx_k + h[None, :] * sx_h, + mask=mask_mh & (ids1 != -1)[:, None], + other=0.0, + ).to(tl.float32) + f2 = tl.load( + x_ptr + m[:, None] * sx_m + 2 * sx_k + h[None, :] * sx_h, + mask=mask_mh & (ids2 != -1)[:, None], + other=0.0, + ).to(tl.float32) + f3 = tl.load( + x_ptr + m[:, None] * sx_m + 3 * sx_k + h[None, :] * sx_h, + mask=mask_mh & (ids3 != -1)[:, None], + other=0.0, + ).to(tl.float32) + f4 = tl.load( + x_ptr + m[:, None] * sx_m + 4 * sx_k + h[None, :] * sx_h, + mask=mask_mh & (ids4 != -1)[:, None], + other=0.0, + ).to(tl.float32) + f5 = tl.load( + x_ptr + m[:, None] * sx_m + 5 * sx_k + h[None, :] * sx_h, + mask=mask_mh & (ids5 != -1)[:, None], + other=0.0, + ).to(tl.float32) + f6 = tl.load( + x_ptr + m[:, None] * sx_m + 6 * sx_k + h[None, :] * sx_h, + mask=mask_mh & (ids6 != -1)[:, None], + other=0.0, + ).to(tl.float32) + f7 = tl.load( + x_ptr + m[:, None] * sx_m + 7 * sx_k + h[None, :] * sx_h, + mask=mask_mh & (ids7 != -1)[:, None], + other=0.0, + ).to(tl.float32) + + # ---- main: runtime loop ---- + for bit in tl.range(0, LOGE): + bitmask = 1 << bit + + # snapshot ids for matching (critical!) + a0, a1, a2, a3, a4, a5, a6, a7 = ids0, ids1, ids2, ids3, ids4, ids5, ids6, ids7 + + # ---------------- s0 ---------------- + c0 = (a0 != -1) & ((a0 & bitmask) != 0) + t0 = tl.where(c0, a0 ^ bitmask, NEG2) + m01 = a1 == t0 + m02 = a2 == t0 + m03 = a3 == t0 + m04 = a4 == t0 + m05 = a5 == t0 + m06 = a6 == t0 + m07 = a7 == t0 + hit0 = m01 | m02 | m03 | m04 | m05 | m06 | m07 + src = f0 + f1 = tl.where(m01[:, None], f1 + src, f1) + f2 = tl.where(m02[:, None], f2 + src, f2) + f3 = tl.where(m03[:, None], f3 + src, f3) + f4 = tl.where(m04[:, None], f4 + src, f4) + f5 = tl.where(m05[:, None], f5 + src, f5) + f6 = tl.where(m06[:, None], f6 + src, f6) + f7 = tl.where(m07[:, None], f7 + src, f7) + f0 = tl.where(hit0[:, None], 0.0, f0) + ids0 = tl.where(c0, tl.where(hit0, -1, t0), ids0) + + # ---------------- s1 ---------------- + c1 = (a1 != -1) & ((a1 & bitmask) != 0) + t1 = tl.where(c1, a1 ^ bitmask, NEG2) + m10 = a0 == t1 + m12 = a2 == t1 + m13 = a3 == t1 + m14 = a4 == t1 + m15 = a5 == t1 + m16 = a6 == t1 + m17 = a7 == t1 + hit1 = m10 | m12 | m13 | m14 | m15 | m16 | m17 + src = f1 + f0 = tl.where(m10[:, None], f0 + src, f0) + f2 = tl.where(m12[:, None], f2 + src, f2) + f3 = tl.where(m13[:, None], f3 + src, f3) + f4 = tl.where(m14[:, None], f4 + src, f4) + f5 = tl.where(m15[:, None], f5 + src, f5) + f6 = tl.where(m16[:, None], f6 + src, f6) + f7 = tl.where(m17[:, None], f7 + src, f7) + f1 = tl.where(hit1[:, None], 0.0, f1) + ids1 = tl.where(c1, tl.where(hit1, -1, t1), ids1) + + # ---------------- s2 ---------------- + c2 = (a2 != -1) & ((a2 & bitmask) != 0) + t2 = tl.where(c2, a2 ^ bitmask, NEG2) + m20 = a0 == t2 + m21 = a1 == t2 + m23 = a3 == t2 + m24 = a4 == t2 + m25 = a5 == t2 + m26 = a6 == t2 + m27 = a7 == t2 + hit2 = m20 | m21 | m23 | m24 | m25 | m26 | m27 + src = f2 + f0 = tl.where(m20[:, None], f0 + src, f0) + f1 = tl.where(m21[:, None], f1 + src, f1) + f3 = tl.where(m23[:, None], f3 + src, f3) + f4 = tl.where(m24[:, None], f4 + src, f4) + f5 = tl.where(m25[:, None], f5 + src, f5) + f6 = tl.where(m26[:, None], f6 + src, f6) + f7 = tl.where(m27[:, None], f7 + src, f7) + f2 = tl.where(hit2[:, None], 0.0, f2) + ids2 = tl.where(c2, tl.where(hit2, -1, t2), ids2) + + # ---------------- s3 ---------------- + c3 = (a3 != -1) & ((a3 & bitmask) != 0) + t3 = tl.where(c3, a3 ^ bitmask, NEG2) + m30 = a0 == t3 + m31 = a1 == t3 + m32 = a2 == t3 + m34 = a4 == t3 + m35 = a5 == t3 + m36 = a6 == t3 + m37 = a7 == t3 + hit3 = m30 | m31 | m32 | m34 | m35 | m36 | m37 + src = f3 + f0 = tl.where(m30[:, None], f0 + src, f0) + f1 = tl.where(m31[:, None], f1 + src, f1) + f2 = tl.where(m32[:, None], f2 + src, f2) + f4 = tl.where(m34[:, None], f4 + src, f4) + f5 = tl.where(m35[:, None], f5 + src, f5) + f6 = tl.where(m36[:, None], f6 + src, f6) + f7 = tl.where(m37[:, None], f7 + src, f7) + f3 = tl.where(hit3[:, None], 0.0, f3) + ids3 = tl.where(c3, tl.where(hit3, -1, t3), ids3) + + # ---------------- s4 ---------------- + c4 = (a4 != -1) & ((a4 & bitmask) != 0) + t4 = tl.where(c4, a4 ^ bitmask, NEG2) + m40 = a0 == t4 + m41 = a1 == t4 + m42 = a2 == t4 + m43 = a3 == t4 + m45 = a5 == t4 + m46 = a6 == t4 + m47 = a7 == t4 + hit4 = m40 | m41 | m42 | m43 | m45 | m46 | m47 + src = f4 + f0 = tl.where(m40[:, None], f0 + src, f0) + f1 = tl.where(m41[:, None], f1 + src, f1) + f2 = tl.where(m42[:, None], f2 + src, f2) + f3 = tl.where(m43[:, None], f3 + src, f3) + f5 = tl.where(m45[:, None], f5 + src, f5) + f6 = tl.where(m46[:, None], f6 + src, f6) + f7 = tl.where(m47[:, None], f7 + src, f7) + f4 = tl.where(hit4[:, None], 0.0, f4) + ids4 = tl.where(c4, tl.where(hit4, -1, t4), ids4) + + # ---------------- s5 ---------------- + c5 = (a5 != -1) & ((a5 & bitmask) != 0) + t5 = tl.where(c5, a5 ^ bitmask, NEG2) + m50 = a0 == t5 + m51 = a1 == t5 + m52 = a2 == t5 + m53 = a3 == t5 + m54 = a4 == t5 + m56 = a6 == t5 + m57 = a7 == t5 + hit5 = m50 | m51 | m52 | m53 | m54 | m56 | m57 + src = f5 + f0 = tl.where(m50[:, None], f0 + src, f0) + f1 = tl.where(m51[:, None], f1 + src, f1) + f2 = tl.where(m52[:, None], f2 + src, f2) + f3 = tl.where(m53[:, None], f3 + src, f3) + f4 = tl.where(m54[:, None], f4 + src, f4) + f6 = tl.where(m56[:, None], f6 + src, f6) + f7 = tl.where(m57[:, None], f7 + src, f7) + f5 = tl.where(hit5[:, None], 0.0, f5) + ids5 = tl.where(c5, tl.where(hit5, -1, t5), ids5) + + # ---------------- s6 ---------------- + c6 = (a6 != -1) & ((a6 & bitmask) != 0) + t6 = tl.where(c6, a6 ^ bitmask, NEG2) + m60 = a0 == t6 + m61 = a1 == t6 + m62 = a2 == t6 + m63 = a3 == t6 + m64 = a4 == t6 + m65 = a5 == t6 + m67 = a7 == t6 + hit6 = m60 | m61 | m62 | m63 | m64 | m65 | m67 + src = f6 + f0 = tl.where(m60[:, None], f0 + src, f0) + f1 = tl.where(m61[:, None], f1 + src, f1) + f2 = tl.where(m62[:, None], f2 + src, f2) + f3 = tl.where(m63[:, None], f3 + src, f3) + f4 = tl.where(m64[:, None], f4 + src, f4) + f5 = tl.where(m65[:, None], f5 + src, f5) + f7 = tl.where(m67[:, None], f7 + src, f7) + f6 = tl.where(hit6[:, None], 0.0, f6) + ids6 = tl.where(c6, tl.where(hit6, -1, t6), ids6) + + # ---------------- s7 ---------------- + c7 = (a7 != -1) & ((a7 & bitmask) != 0) + t7 = tl.where(c7, a7 ^ bitmask, NEG2) + m70 = a0 == t7 + m71 = a1 == t7 + m72 = a2 == t7 + m73 = a3 == t7 + m74 = a4 == t7 + m75 = a5 == t7 + m76 = a6 == t7 + hit7 = m70 | m71 | m72 | m73 | m74 | m75 | m76 + src = f7 + f0 = tl.where(m70[:, None], f0 + src, f0) + f1 = tl.where(m71[:, None], f1 + src, f1) + f2 = tl.where(m72[:, None], f2 + src, f2) + f3 = tl.where(m73[:, None], f3 + src, f3) + f4 = tl.where(m74[:, None], f4 + src, f4) + f5 = tl.where(m75[:, None], f5 + src, f5) + f6 = tl.where(m76[:, None], f6 + src, f6) + f7 = tl.where(hit7[:, None], 0.0, f7) + ids7 = tl.where(c7, tl.where(hit7, -1, t7), ids7) + + # ---- per-level bf16 rounding (optional) ---- + if CAST_MODE == 0: + # emulate "store bf16 per level then reload": round once per level + f0 = f0.to(x_ty).to(tl.float32) + f1 = f1.to(x_ty).to(tl.float32) + f2 = f2.to(x_ty).to(tl.float32) + f3 = f3.to(x_ty).to(tl.float32) + f4 = f4.to(x_ty).to(tl.float32) + f5 = f5.to(x_ty).to(tl.float32) + f6 = f6.to(x_ty).to(tl.float32) + f7 = f7.to(x_ty).to(tl.float32) + + # ---- gather root (id==0) ---- + acc = tl.zeros((BLOCK_M, BLOCK_H), dtype=tl.float32) + acc += tl.where((ids0 == 0)[:, None], f0, 0.0) + acc += tl.where((ids1 == 0)[:, None], f1, 0.0) + acc += tl.where((ids2 == 0)[:, None], f2, 0.0) + acc += tl.where((ids3 == 0)[:, None], f3, 0.0) + acc += tl.where((ids4 == 0)[:, None], f4, 0.0) + acc += tl.where((ids5 == 0)[:, None], f5, 0.0) + acc += tl.where((ids6 == 0)[:, None], f6, 0.0) + acc += tl.where((ids7 == 0)[:, None], f7, 0.0) + + acc *= routed_scaling_factor + + out_ptrs = out_ptr + m[:, None] * so_m + h[None, :] * so_h + tl.store(out_ptrs, acc.to(out_ptr.dtype.element_ty), mask=mask_mh) + + +def _launch_sparse_tree_k8_k10( + input: torch.Tensor, # [M, K, H], bf16 contiguous + output: torch.Tensor, # [M, H], bf16 contiguous + curr_topk_ids: torch.Tensor, # [M, K], int32/int64 contiguous, -1 remote + routed_scaling_factor: float, + E: int, + cast_mode: int = 1, # <-- you asked for "last cast" version, so default = 1 +): + assert ( + input.is_contiguous() + and output.is_contiguous() + and curr_topk_ids.is_contiguous() + ) + M, K, H = input.shape + assert output.shape == (M, H) + assert K in (8, 10) + assert (E & (E - 1)) == 0 + LOGE = int(math.log2(E)) + + # stable params for CUDA Graph + BLOCK_M = 8 + BLOCK_H = 256 + num_warps = 8 + grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(H, BLOCK_H)) + + # NOTE: for strict CUDA Graph safety, avoid dtype conversion here. + # Please ensure curr_topk_ids is int32 upstream. + assert curr_topk_ids.dtype == torch.int32, ( + "make ids int32 upstream for CUDA Graph stability" + ) + + if K == 8: + _moe_tree_reduce_sparse_k8_kernel[grid]( + input, + curr_topk_ids, + output, + input.stride(0), + input.stride(1), + input.stride(2), + curr_topk_ids.stride(0), + curr_topk_ids.stride(1), + output.stride(0), + output.stride(1), + M, + H, + routed_scaling_factor, + LOGE=LOGE, + CAST_MODE=cast_mode, + BLOCK_M=BLOCK_M, + BLOCK_H=BLOCK_H, + num_warps=num_warps, + ) + + +def moe_sum_tree_reduce_v2( + input: torch.Tensor, + output: torch.Tensor, + curr_topk_ids: torch.Tensor, + routed_scaling_factor: float, + E: int, + *, + cast_mode: int = 0, # default to "last cast" as requested +): + assert ( + input.is_contiguous() + and output.is_contiguous() + and curr_topk_ids.is_contiguous() + ) + M, K, H = input.shape + if M == 0: + return output + if K in (8, 10): + _launch_sparse_tree_k8_k10( + input=input, + output=output, + curr_topk_ids=curr_topk_ids, + routed_scaling_factor=routed_scaling_factor, + E=E, + cast_mode=cast_mode, + ) + return output + + # fallback: keep your existing generic path here + return output + + +def moe_sum_tree_reduce( + input: torch.Tensor, + output: torch.Tensor, + curr_topk_ids: torch.Tensor, + routed_scaling_factor: float, + E: int, +): + curr_topk_ids = curr_topk_ids.to(torch.int32) + if os.environ.get("SGLANG_MOE_TREE_REDUCE_USE_V2", "0") == "1": + return moe_sum_tree_reduce_v2( + input=input, + output=output, + curr_topk_ids=curr_topk_ids, + routed_scaling_factor=routed_scaling_factor, + E=E, + ) + + if input.shape[1] == 8: + return moe_sum_tree_reduce_v0( + input=input, + output=output, + curr_topk_ids=curr_topk_ids, + routed_scaling_factor=routed_scaling_factor, + E=E, + ) + return moe_sum_tree_reduce_v1( + input=input, + output=output, + curr_topk_ids=curr_topk_ids, + routed_scaling_factor=routed_scaling_factor, + E=E, + ) + + +# if os.getenv("MOE_SUM_OPTIM", "0") == "1": +# moe_sum_tree_reduce = moe_sum_tree_reduce_optim +# print("using optimized moe_sum_tree_reduce_optim") +# else: +# print("using original moe_sum_tree_reduce_original") + +if os.getenv("TREE_ALL_REDUCE_OPTIM", "0") == "1": + tree_all_reduce_sum = tree_all_reduce_sum_optim diff --git a/python/sglang/srt/true_on_policy/__init__.py b/python/sglang/srt/true_on_policy/__init__.py new file mode 100644 index 000000000000..9fbaaeee48f3 --- /dev/null +++ b/python/sglang/srt/true_on_policy/__init__.py @@ -0,0 +1,49 @@ +"""True-on-policy runtime contract helpers.""" + +from .config import ( + ROW_LINEAR_INV_BLOCK_K, + get_on_policy_rms_norm_kwargs, + get_rl_on_policy_target, + is_tp_invariant_target, + is_true_on_policy_enabled, + patch_prefill_only_deterministic_inference_for_cuda_graph, + should_disable_flashinfer_allreduce_fusion, + should_disable_fused_qk_norm_mrope, + should_disable_mlp_allreduce_fusion_for_on_policy, + should_disable_reduce_scatter_for_on_policy, + should_force_bfloat16_dense_tensor_math, + should_force_bfloat16_lm_head, + should_use_tp_invariant_row_linear, + should_use_tp_invariant_tree_all_reduce, +) +from .contracts import ( + QWEN3_DENSE_TRUE_ON_POLICY_V1, + SGLangTrueOnPolicyContract, + SGLangTrueOnPolicyRuntimePolicy, + get_true_on_policy_contract, + resolve_true_on_policy_runtime_policy, + validate_true_on_policy_contract, +) + +__all__ = [ + "QWEN3_DENSE_TRUE_ON_POLICY_V1", + "ROW_LINEAR_INV_BLOCK_K", + "SGLangTrueOnPolicyContract", + "SGLangTrueOnPolicyRuntimePolicy", + "get_true_on_policy_contract", + "get_on_policy_rms_norm_kwargs", + "get_rl_on_policy_target", + "is_tp_invariant_target", + "is_true_on_policy_enabled", + "patch_prefill_only_deterministic_inference_for_cuda_graph", + "resolve_true_on_policy_runtime_policy", + "validate_true_on_policy_contract", + "should_disable_flashinfer_allreduce_fusion", + "should_disable_fused_qk_norm_mrope", + "should_disable_mlp_allreduce_fusion_for_on_policy", + "should_disable_reduce_scatter_for_on_policy", + "should_force_bfloat16_dense_tensor_math", + "should_force_bfloat16_lm_head", + "should_use_tp_invariant_row_linear", + "should_use_tp_invariant_tree_all_reduce", +] diff --git a/python/sglang/srt/true_on_policy/config.py b/python/sglang/srt/true_on_policy/config.py new file mode 100644 index 000000000000..82e156d6d4e8 --- /dev/null +++ b/python/sglang/srt/true_on_policy/config.py @@ -0,0 +1,150 @@ +from __future__ import annotations + +import contextlib +from typing import Any, Iterator, Optional + +import torch + +from sglang.srt.true_on_policy.contracts import resolve_true_on_policy_runtime_policy + +ROW_LINEAR_INV_BLOCK_K = 128 + + +def _get_global_server_args() -> Any: + from sglang.srt.runtime_context import get_server_args + + return get_server_args() + + +def get_rl_on_policy_target() -> Optional[str]: + return getattr(_get_global_server_args(), "rl_on_policy_target", None) + + +def is_true_on_policy_enabled() -> bool: + return resolve_true_on_policy_runtime_policy(_get_global_server_args()).enabled + + +def is_tp_invariant_target() -> bool: + return resolve_true_on_policy_runtime_policy( + _get_global_server_args() + ).tp_invariant_row_linear + + +def should_disable_reduce_scatter_for_on_policy() -> bool: + return resolve_true_on_policy_runtime_policy( + _get_global_server_args() + ).disable_reduce_scatter + + +def should_disable_mlp_allreduce_fusion_for_on_policy() -> bool: + return resolve_true_on_policy_runtime_policy( + _get_global_server_args() + ).disable_mlp_allreduce_fusion + + +def should_disable_flashinfer_allreduce_fusion() -> bool: + return resolve_true_on_policy_runtime_policy( + _get_global_server_args() + ).disable_flashinfer_allreduce_fusion + + +def should_force_bfloat16_dense_tensor_math() -> bool: + return resolve_true_on_policy_runtime_policy( + _get_global_server_args() + ).force_bfloat16_dense_tensor_math + + +def should_force_bfloat16_lm_head( + *, + use_fp32_lm_head: bool = False, +) -> bool: + return ( + resolve_true_on_policy_runtime_policy( + _get_global_server_args() + ).force_bfloat16_lm_head + and not use_fp32_lm_head + ) + + +def should_disable_fused_qk_norm_mrope() -> bool: + return resolve_true_on_policy_runtime_policy( + _get_global_server_args() + ).disable_fused_qk_norm_mrope + + +def get_on_policy_rms_norm_kwargs( + *, + weight_dtype: Optional[torch.dtype] = None, + override_orig_dtype: Optional[torch.dtype] = None, + fp32_residual: bool = False, +) -> dict[str, Any]: + if not is_true_on_policy_enabled(): + return {} + + kwargs: dict[str, Any] = { + "cast_x_before_out_mul": True, + "fp32_residual": fp32_residual, + } + if weight_dtype is not None: + kwargs["weight_dtype"] = weight_dtype + if override_orig_dtype is not None: + kwargs["override_orig_dtype"] = override_orig_dtype + return kwargs + + +def should_use_tp_invariant_row_linear( + k_size: int, + row_linear_enable_inv: Optional[bool] = None, +) -> bool: + policy_enabled = resolve_true_on_policy_runtime_policy( + _get_global_server_args() + ).tp_invariant_row_linear + if row_linear_enable_inv is not None: + policy_enabled = policy_enabled and row_linear_enable_inv + + return ( + policy_enabled + and k_size >= ROW_LINEAR_INV_BLOCK_K + and k_size % ROW_LINEAR_INV_BLOCK_K == 0 + ) + + +def should_use_tp_invariant_tree_all_reduce( + accl_binary_tree_enabled: Optional[bool] = None, +) -> bool: + policy_enabled = resolve_true_on_policy_runtime_policy( + _get_global_server_args() + ).deterministic_tree_all_reduce + if accl_binary_tree_enabled is not None: + policy_enabled = policy_enabled and not accl_binary_tree_enabled + + return policy_enabled + + +@contextlib.contextmanager +def patch_prefill_only_deterministic_inference_for_cuda_graph( + server_args: Any, + *, + attn_backend: Optional[Any] = None, + dvr_target_verify_cuda_graph: bool = False, +) -> Iterator[bool]: + enabled = ( + getattr(server_args, "enable_prefill_only_deterministic_inference", False) + and not dvr_target_verify_cuda_graph + ) + if not enabled: + yield False + return + + saved_num_splits = None + if attn_backend is not None and hasattr(attn_backend, "num_splits"): + saved_num_splits = attn_backend.num_splits + + try: + if attn_backend is not None and hasattr(attn_backend, "num_splits"): + attn_backend.num_splits = 0 + + yield True + finally: + if attn_backend is not None and hasattr(attn_backend, "num_splits"): + attn_backend.num_splits = saved_num_splits diff --git a/python/sglang/srt/true_on_policy/contracts.py b/python/sglang/srt/true_on_policy/contracts.py new file mode 100644 index 000000000000..310778c842ce --- /dev/null +++ b/python/sglang/srt/true_on_policy/contracts.py @@ -0,0 +1,111 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Optional + +from sglang.srt.true_on_policy.schema import ( + QWEN3_DENSE_TRUE_ON_POLICY_V1_SCHEMA, + TrueOnPolicyContractName, + TrueOnPolicyContractSchema, +) + +QWEN3_DENSE_TRUE_ON_POLICY_V1 = QWEN3_DENSE_TRUE_ON_POLICY_V1_SCHEMA.name + + +@dataclass(frozen=True) +class SGLangTrueOnPolicyRuntimePolicy: + """SGLang-local behavior implied by a true-on-policy parity contract.""" + + contract_name: Optional[str] + enabled: bool + force_bfloat16_dense_tensor_math: bool + force_bfloat16_lm_head: bool + disable_reduce_scatter: bool + disable_mlp_allreduce_fusion: bool + disable_flashinfer_allreduce_fusion: bool + tp_invariant_row_linear: bool + deterministic_tree_all_reduce: bool + disable_fused_qk_norm_mrope: bool + + +DEFAULT_RUNTIME_POLICY = SGLangTrueOnPolicyRuntimePolicy( + contract_name=None, + enabled=False, + force_bfloat16_dense_tensor_math=False, + force_bfloat16_lm_head=False, + disable_reduce_scatter=False, + disable_mlp_allreduce_fusion=False, + disable_flashinfer_allreduce_fusion=False, + tp_invariant_row_linear=False, + deterministic_tree_all_reduce=False, + disable_fused_qk_norm_mrope=False, +) + + +@dataclass(frozen=True) +class SGLangTrueOnPolicyContract: + """SGLang-local adapter from a shared contract schema to runtime policy.""" + + schema: TrueOnPolicyContractSchema + + @property + def name(self) -> TrueOnPolicyContractName: + return self.schema.name + + def policy_for(self, server_args: Any) -> SGLangTrueOnPolicyRuntimePolicy: + uses_tp_invariant_rollout = getattr(server_args, "tp_size", 1) > 1 + return SGLangTrueOnPolicyRuntimePolicy( + contract_name=self.name, + enabled=True, + force_bfloat16_dense_tensor_math=True, + force_bfloat16_lm_head=True, + disable_reduce_scatter=True, + disable_mlp_allreduce_fusion=True, + disable_flashinfer_allreduce_fusion=uses_tp_invariant_rollout, + tp_invariant_row_linear=uses_tp_invariant_rollout, + deterministic_tree_all_reduce=uses_tp_invariant_rollout, + disable_fused_qk_norm_mrope=True, + ) + + +QWEN3_DENSE_TRUE_ON_POLICY_CONTRACT = SGLangTrueOnPolicyContract( + schema=QWEN3_DENSE_TRUE_ON_POLICY_V1_SCHEMA, +) + + +_CONTRACT_BY_NAME = { + QWEN3_DENSE_TRUE_ON_POLICY_CONTRACT.name: QWEN3_DENSE_TRUE_ON_POLICY_CONTRACT, +} + + +def get_true_on_policy_contract(contract_name: str) -> SGLangTrueOnPolicyContract: + try: + return _CONTRACT_BY_NAME[contract_name] + except KeyError as exc: + supported = ", ".join(sorted(_CONTRACT_BY_NAME)) + raise ValueError( + f"Unsupported SGLang true-on-policy contract {contract_name!r}. " + f"Supported contracts: {supported}" + ) from exc + + +def _contract_name_for(server_args: Any) -> Optional[str]: + return getattr(server_args, "true_on_policy_contract", None) + + +def validate_true_on_policy_contract(server_args: Any) -> None: + contract_name = getattr(server_args, "true_on_policy_contract", None) + if contract_name is None: + return + get_true_on_policy_contract(contract_name) + + +def resolve_true_on_policy_runtime_policy( + server_args: Any, +) -> SGLangTrueOnPolicyRuntimePolicy: + contract_name = _contract_name_for(server_args) + if contract_name is None: + return DEFAULT_RUNTIME_POLICY + + validate_true_on_policy_contract(server_args) + return get_true_on_policy_contract(contract_name).policy_for(server_args) diff --git a/python/sglang/srt/true_on_policy/schema.py b/python/sglang/srt/true_on_policy/schema.py new file mode 100644 index 000000000000..0628573c6880 --- /dev/null +++ b/python/sglang/srt/true_on_policy/schema.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal + +TrueOnPolicyContractName = Literal["qwen3_dense_true_on_policy_v1"] +ModelFamily = Literal["qwen3_dense", "qwen3_moe", "qwen3_next"] +KernelContract = Literal["qwen3_dense_sglang_math"] +LogprobContract = Literal["sglang_prefill"] + + +@dataclass(frozen=True) +class TrueOnPolicyContractSchema: + """Declarative cross-repo identity for a true-on-policy parity contract.""" + + name: TrueOnPolicyContractName + model_family: ModelFamily + required_kernel_contracts: tuple[KernelContract, ...] + logprob_contract: LogprobContract + sglang_attention_backend: str + fsdp_attention_implementation: str + disable_megatron_sequence_parallel: bool + + +QWEN3_DENSE_TRUE_ON_POLICY_V1_SCHEMA = TrueOnPolicyContractSchema( + name="qwen3_dense_true_on_policy_v1", + model_family="qwen3_dense", + required_kernel_contracts=("qwen3_dense_sglang_math",), + logprob_contract="sglang_prefill", + sglang_attention_backend="fa3", + fsdp_attention_implementation="flash_attention_3", + disable_megatron_sequence_parallel=True, +) diff --git a/python/sglang/srt/utils/auth.py b/python/sglang/srt/utils/auth.py index e7ef0767f122..df6811179e5d 100644 --- a/python/sglang/srt/utils/auth.py +++ b/python/sglang/srt/utils/auth.py @@ -90,14 +90,19 @@ def decide_request_auth( it must be rejected (403) even if api_key is provided. NOTE : - - Health/metrics endpoints are always allowed (even when api_key/admin_api_key is set), - to support k8s/liveness/readiness and Prometheus scraping without embedding secrets. + - Health/readiness/metrics endpoints are always allowed (even when + api_key/admin_api_key is set), to support k8s probes and Prometheus + scraping without embedding secrets. - We match them by prefix to cover common variants like /health_generate. """ if method == "OPTIONS": return AuthDecision(allowed=True) - if path.startswith("/health") or path.startswith("/metrics"): + if ( + path.startswith("/health") + or path.startswith("/ready") + or path.startswith("/metrics") + ): return AuthDecision(allowed=True) def _check_bearer_token( diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 62ad5159f6ab..9be3fdbdf8f5 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -4469,6 +4469,10 @@ def mxfp8_block_convert_required(): # Inkling attention projections (merged q/k/v/r and its row-parallel output). "qkvr", "wo_ud", + # GDN (GatedDeltaNet) projections + "in_proj_qkvz", + "in_proj_ba", + "out_proj", ] LORA_TARGET_ALL_MODULES = "all" diff --git a/python/sglang/srt/utils/weight_checker.py b/python/sglang/srt/utils/weight_checker.py index 6e0b5581e2e2..433a5fe39bc0 100644 --- a/python/sglang/srt/utils/weight_checker.py +++ b/python/sglang/srt/utils/weight_checker.py @@ -1,13 +1,14 @@ import hashlib import logging import time -from typing import Any, Callable, Dict, Iterable, NamedTuple, Optional, Set +from typing import Any, Callable, Dict, Iterable, List, NamedTuple, Optional, Set import torch import torch.distributed as dist from pydantic import BaseModel, ConfigDict from sglang.srt.managers.mm_utils import tensor_hash +from sglang.srt.mem_cache.storage.mmap.mmap_allocator import alloc_mmap from sglang.srt.utils.weight_checker_comparator import ( CHUNK_NUMEL, ComparableWeight, @@ -66,51 +67,102 @@ def _is_non_persistent_buffer_name(name: str) -> bool: return any(pat in name for pat in _NON_PERSISTENT_BUFFER_PATTERNS) +def _is_skip_weight_check(name, param, skip_tensor_list=None) -> bool: + # One skip group shared by reset / compare / checksum; the _skip_weight_check + # flag is set on kv-cache k/v_scale and process_weights_after_loading placeholders. + return ( + _is_non_persistent_buffer_name(name) + or getattr(param, "_skip_weight_check", False) + or any(pat in name for pat in (skip_tensor_list or ())) + ) + + +def overall_checksum(checksums: Dict[str, str]) -> str: + h = hashlib.sha256() + for name in sorted(checksums): + h.update(name.encode()) + h.update(checksums[name].encode()) + return h.hexdigest() + + class WeightChecker: def __init__(self, *, get_model: Callable[[], Any], ps: Any): self._get_model = get_model self._ps = ps self._snapshot_tensors = None - - def handle(self, action: str, allow_quant_error: bool = False) -> Optional[Dict]: + self._snapshot_arena = None + + def handle( + self, + action: str, + allow_quant_error: bool = False, + skip_tensor_list: Optional[List[str]] = None, + ) -> Optional[Dict]: logger.info( - f"[WeightChecker] handle action={action} allow_quant_error={allow_quant_error}" + f"[WeightChecker] handle action={action} " + f"allow_quant_error={allow_quant_error} skip_tensor_list={skip_tensor_list}" ) if action == "snapshot": return self._snapshot() elif action == "reset_tensors": - return self._reset_tensors() + return self._reset_tensors(skip_tensor_list) elif action == "compare": - return self._compare(allow_quant_error=allow_quant_error) + return self._compare( + allow_quant_error=allow_quant_error, skip_tensor_list=skip_tensor_list + ) elif action == "checksum": - return self._compute_checksum() + return self._compute_checksum(skip_tensor_list) else: raise Exception(f"Unsupported {action=}") def _snapshot(self): - named_tensors = [ - (name, param.data.detach().cpu()) for name, param in self._model_state() - ] - self._snapshot_tensors = dict(named_tensors) - assert len(self._snapshot_tensors) == len(named_tensors), ( + named_params = [(name, param.data) for name, param in self._model_state()] + align = 64 # torch CPU-allocator alignment + offsets, total = [], 0 + for _, param in named_params: + offsets.append(total) + total += (param.nbytes + align - 1) // align * align + self._snapshot_arena = alloc_mmap((max(total, align),), torch.uint8) + snapshot_tensors = {} + for (name, param), offset in zip(named_params, offsets): + view = ( + self._snapshot_arena[offset : offset + param.nbytes] + .view(param.dtype) + .view(param.shape) + ) + view.copy_(param.detach()) + snapshot_tensors[name] = view + self._snapshot_tensors = snapshot_tensors + assert len(self._snapshot_tensors) == len(named_params), ( f"should not have duplicated tensor name" ) - def _reset_tensors(self): + def _skip_compare_names( + self, skip_tensor_list: Optional[List[str]] = None + ) -> Set[str]: + return { + name + for name, param in self._model_state() + if _is_skip_weight_check(name, param, skip_tensor_list) + } + + def _reset_tensors(self, skip_tensor_list: Optional[List[str]] = None): for name, param in self._model_state(): - if _is_non_persistent_buffer_name(name): + # Skip exactly what compare/checksum skip, so reset only poisons + # tensors compare will verify. + if _is_skip_weight_check(name, param, skip_tensor_list): continue param.copy_(_random_like(param)) - def _compare(self, allow_quant_error: bool = False): + def _compare( + self, + allow_quant_error: bool = False, + skip_tensor_list: Optional[List[str]] = None, + ): assert self._snapshot_tensors is not None quantized_set = _build_quantized_set(self._get_model()) - skip_compare_names = { - name - for name, param in self._model_state() - if getattr(param, "_skip_weight_check", False) - } + skip_compare_names = self._skip_compare_names(skip_tensor_list) _check_tensors( expect_tensors=_build_check_entries( self._snapshot_tensors, skip_compare_names, quantized_set @@ -120,17 +172,15 @@ def _compare(self, allow_quant_error: bool = False): ), allow_quant_error=allow_quant_error, ) + self._snapshot_tensors = None + self._snapshot_arena = None - def _compute_checksum(self) -> Dict: + def _compute_checksum(self, skip_tensor_list: Optional[List[str]] = None) -> Dict: torch.cuda.synchronize() start = time.perf_counter() quantized_set = _build_quantized_set(self._get_model()) - skip_compare_names = { - name - for name, param in self._model_state() - if getattr(param, "_skip_weight_check", False) - } + skip_compare_names = self._skip_compare_names(skip_tensor_list) # Hash the dequantized weight so two (qweight, scale) pairs with the same # bf16 hash equal. @@ -141,11 +191,7 @@ def _compute_checksum(self) -> Dict: if should_compare: checksums[name] = _hash_tensor(comparable.dequantize().data) - h = hashlib.sha256() - for name in sorted(checksums): - h.update(name.encode()) - h.update(checksums[name].encode()) - overall = h.hexdigest() + overall = overall_checksum(checksums) torch.cuda.synchronize() elapsed = time.perf_counter() - start @@ -272,7 +318,7 @@ def _build_quantized_set(model) -> Dict[str, QuantizedWeight]: prefix = f"{module_name}." if module_name else "" own = dict(module.named_parameters(recurse=False)) for name, parameter in own.items(): - scale = name.replace("weight", "weight_scale_inv") + scale = name.replace("weight", comparable_cls.SCALE_SUFFIX) if name.endswith("weight") and scale in own: quantized_set[prefix + name] = QuantizedWeight( comparable_cls, @@ -300,7 +346,7 @@ def _build_check_entries( qw = quantized_set[name] yield CheckEntry( name, - True, + name not in skip_compare_names, qw.comparable_cls( tensor, raw[qw.scale_name], is_shuffled=qw.is_shuffled ), diff --git a/python/sglang/srt/utils/weight_checker_comparator.py b/python/sglang/srt/utils/weight_checker_comparator.py index 3578e9c583b9..2747de51ee04 100644 --- a/python/sglang/srt/utils/weight_checker_comparator.py +++ b/python/sglang/srt/utils/weight_checker_comparator.py @@ -11,10 +11,16 @@ block_quant_dequant, inverse_transform_scale_ue8m0, ) +from sglang.srt.layers.quantization.marlin_utils import ( + get_scale_perms, + get_weight_perm, + marlin_permute_weights, +) from sglang.srt.layers.quantization.modelopt_quant import ( ModelOptFp4LinearMethod, ModelOptNvFp4FusedMoEMethod, ) +from sglang.srt.layers.quantization.mxfp4 import Mxfp4MoEMethod # chunk to avoid too high GPU memory peak CHUNK_NUMEL = 64 * 1024 * 1024 @@ -30,6 +36,9 @@ class CompareResult(NamedTuple): class ComparableWeight: """Base comparable-weight class; one subclass per precision or raw tensor.""" + # `weight` pairs with `SCALE_SUFFIX` on the same module + SCALE_SUFFIX = "weight_scale_inv" + @staticmethod def _quant_ulp(w_q: torch.Tensor) -> torch.Tensor: """Per-element ULP of w_q in its own dtype.""" @@ -120,6 +129,138 @@ def dequantize(self, dtype: torch.dtype = torch.bfloat16) -> torch.Tensor: return block_quant_dequant(w_q, s, block_size, dtype=dtype) +_E2M1_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) +_E2M1_SPACING = (0.5, 0.5, 0.5, 0.5, 1.0, 1.0, 2.0, 2.0) +_MXFP4_GROUP_SIZE = 32 + + +class Mxfp4MarlinComparable(ComparableWeight): + """MXFP4 MoE experts after `prepare_moe_mxfp4_layer_for_marlin`: e2m1 nibbles in + Marlin-repacked int32 `(E, K/16, 2N)`, e8m0 scales in Marlin order `(E, K/32, N)`. + + Compared in dequantized `(E, N, K)` space, so a re-quantization that picks a different + but equivalent (scale, nibble) encoding still compares equal.""" + + SCALE_SUFFIX = "weight_scale" + + # (size_k, size_n) -> (nibble source index per (k, n), scale source index per (g, n)) + _layout_cache: dict = {} + + def __init__(self, w_q: torch.Tensor, w_s: torch.Tensor, is_shuffled: bool = False): + # the checker passes every quantized weight's AITER shuffle flag; Marlin never shuffles + assert not is_shuffled, "Marlin MXFP4 experts are never AITER-shuffled" + assert w_q.dtype == torch.int32, ( + f"expected Marlin int32 weights, got {w_q.dtype}" + ) + self.w_q = w_q + self.w_s = w_s + self.size_k = w_q.shape[-2] * 16 + self.size_n = w_q.shape[-1] // 2 + assert tuple(w_s.shape[-2:]) == ( + self.size_k // _MXFP4_GROUP_SIZE, + self.size_n, + ), ( + f"scale shape {tuple(w_s.shape)} does not match weight shape {tuple(w_q.shape)}" + ) + + def __repr__(self) -> str: + return f"mxfp4_marlin(shape={tuple(self.w_q.shape)} dtype={self.w_q.dtype})" + + @classmethod + def _layout(cls, size_k: int, size_n: int, device: torch.device): + """Index maps that undo the Marlin repack, built from the Python reference permutation.""" + key = (size_k, size_n, str(device)) + if key not in cls._layout_cache: + nibble_source = marlin_permute_weights( + torch.arange(size_k * size_n).reshape(size_k, size_n), + size_k, + size_n, + get_weight_perm(num_bits=4), + ).reshape(-1) + # marlin_weights packs consecutive columns into one int32, low nibble first, so the + # flattened permuted matrix already lists nibbles in packed order + nibble_position = torch.empty_like(nibble_source) + nibble_position[nibble_source] = torch.arange(nibble_source.numel()) + + scale_perm, _ = get_scale_perms() + num_groups = size_k // _MXFP4_GROUP_SIZE + scale_source = torch.arange(num_groups * size_n) + scale_source = scale_source.reshape(-1, len(scale_perm))[:, scale_perm] + # mxfp4_marlin_process_scales swaps the middle pair of every four scales + scale_source = scale_source.reshape(-1, 4)[:, [0, 2, 1, 3]].reshape(-1) + scale_position = torch.empty_like(scale_source) + scale_position[scale_source] = torch.arange(scale_source.numel()) + + cls._layout_cache[key] = ( + nibble_position.to(device), + scale_position.to(device), + ) + return cls._layout_cache[key] + + def _dequantize_experts(self, w_q: torch.Tensor, w_s: torch.Tensor): + """`(e, K/16, 2N)` int32 + `(e, K/32, N)` e8m0 -> dequantized and ULP tensors `(e, N, K)`.""" + num_experts = w_q.shape[0] + nibble_position, scale_position = self._layout( + self.size_k, self.size_n, w_q.device + ) + shifts = torch.arange( + 0, 32, 4, device=w_q.device, dtype=torch.int32 + ) # nibble i sits at bits [4i, 4i+4) + nibbles = (w_q.reshape(num_experts, -1, 1) >> shifts) & 0xF + nibbles = nibbles.reshape(num_experts, -1)[:, nibble_position] + nibbles = nibbles.reshape(num_experts, self.size_k, self.size_n) + + exponents = w_s.view(torch.uint8).reshape(num_experts, -1)[:, scale_position] + scale = torch.exp2(exponents.to(torch.float32) - 127.0).reshape( + num_experts, self.size_k // _MXFP4_GROUP_SIZE, 1, self.size_n + ) + + magnitude = (nibbles & 0x7).long() + sign = 1.0 - 2.0 * ((nibbles >> 3) & 0x1).to(torch.float32) + values = torch.tensor(_E2M1_VALUES, device=w_q.device)[magnitude] * sign + spacing = torch.tensor(_E2M1_SPACING, device=w_q.device)[magnitude] + grouped = ( + num_experts, + self.size_k // _MXFP4_GROUP_SIZE, + _MXFP4_GROUP_SIZE, + self.size_n, + ) + dequantized = (values.reshape(grouped) * scale).reshape( + num_experts, self.size_k, self.size_n + ) + ulp = (spacing.reshape(grouped) * scale).reshape( + num_experts, self.size_k, self.size_n + ) + # the checkpoint stores experts as (N, K) + return dequantized.transpose(-1, -2), ulp.transpose(-1, -2) + + def _experts_per_chunk(self) -> int: + # the unpack holds several int32/float32 copies of every nibble + return max(1, CHUNK_NUMEL // (4 * self.size_k * self.size_n)) + + def iter_chunks(self): + q = self.w_q.reshape(-1, *self.w_q.shape[-2:]) + s = self.w_s.reshape(-1, *self.w_s.shape[-2:]) + step = self._experts_per_chunk() + for e0 in range(0, q.shape[0], step): + dequantized, ulp = self._dequantize_experts( + q[e0 : e0 + step].cuda(), s[e0 : e0 + step].cuda() + ) + yield dequantized.to(torch.bfloat16), ulp + + def dequantize(self, dtype: torch.dtype = torch.bfloat16) -> torch.Tensor: + q = self.w_q.reshape(-1, *self.w_q.shape[-2:]) + s = self.w_s.reshape(-1, *self.w_s.shape[-2:]) + step = self._experts_per_chunk() + chunks = [ + self._dequantize_experts( + q[e0 : e0 + step].cuda(), s[e0 : e0 + step].cuda() + )[0].to(dtype) + for e0 in range(0, q.shape[0], step) + ] + return torch.cat(chunks).reshape(*self.w_q.shape[:-2], self.size_n, self.size_k) + + class RawComparable(ComparableWeight): """Bitwise equal compare on raw tensor.""" @@ -189,4 +330,10 @@ def select_comparable_weight(quant_method) -> Optional[type]: raise NotImplementedError( f"weight checker has no ComparableWeight for {type(quant_method).__name__}" ) + if isinstance(quant_method, Mxfp4MoEMethod): + if quant_method.use_marlin: + return Mxfp4MarlinComparable + raise NotImplementedError( + "weight checker has no ComparableWeight for Mxfp4MoEMethod outside the Marlin backend" + ) return None diff --git a/python/sglang/srt/weight_sync/local_checkpoint.py b/python/sglang/srt/weight_sync/local_checkpoint.py new file mode 100644 index 000000000000..8f59802f5590 --- /dev/null +++ b/python/sglang/srt/weight_sync/local_checkpoint.py @@ -0,0 +1,350 @@ +"""Host-local pull of published weights (the /pull_weights endpoint). + +A trainer publishes each weight sync as a version directory ``weight_v{N:06d}/`` +under a shared ``source_dir``. Each version is a canonical HF checkpoint +directory of one of two kinds, distinguished by its index metadata: + +- **full**: an ordinary checkpoint. Pulling it copies it into the host-local + ``local_checkpoint_dir``, replacing whatever is there — no history needed. +- **delta** (index metadata carries ``delta_encoding``): safetensors files + holding zstd-compressed per-tensor diffs against version N-1, plus per-tensor + checksums of the new state. Pulling it patches the local checkpoint in place. + +Version 0 is the engine's own base checkpoint (``model_path``). Every host of a +(possibly multi-node) deployment runs the same pull; the engine then reloads the +local checkpoint through the ordinary ``update_weights_from_disk`` path. + +``pull()`` is safe to call concurrently from every scheduler rank on a host: a +per-host file lock serializes the work and an applied-version marker makes the +extra calls no-ops. +""" + +from __future__ import annotations + +import fcntl +import glob +import importlib +import json +import logging +import mmap +import os +import shutil +import struct +import threading +import zlib +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from typing import Optional + +import numpy as np +import zstandard + +logger = logging.getLogger(__name__) + +# The delta-apply phases (decompress, XOR/scatter, checksum) are memory-bandwidth +# bound and release the GIL, so a thread pool over tensors recovers the +# bandwidth one thread leaves idle. +NUM_WORKERS = min(32, (os.cpu_count() or 8)) + +# Per-checkpoint dir holding the applied-version marker and the pull lock. +SYNC_DIR = ".weight_sync" + + +def pull( + local_checkpoint_dir: str, + base_dir: str, + source_dir: str, + target_version: int, + pre_read_hook: Optional[str] = None, +) -> None: + """Bring the host-local checkpoint up to ``target_version``. + + Seeds from the newest full checkpoint at or below the target — the engine's + own base (``base_dir``) for a pure-delta stream, a published full version + otherwise — then applies the remaining deltas in order. A local checkpoint + already past the seed point just continues its delta chain. Raises on any + per-tensor checksum mismatch (fail loud, never serve bad weights). + """ + # Object-store-backed shared filesystems lack cross-host read-after-write + # consistency: the publisher's files only appear here after an explicit + # refresh, which the deployment supplies as this hook. POSIX shared + # filesystems (NFS, Lustre, ...) need none. + if target_version > 0 and pre_read_hook: + _load_hook(pre_read_hook)(source_dir, target_version) + with _pull_lock(local_checkpoint_dir): + applied = _read_applied_version(local_checkpoint_dir) # None on a fresh host + # Scan back from the target for the newest full version. Stop at the + # local state — below it a reset can never be needed (or, on a fresh + # host, at 0 = the engine's base). + floor = applied if applied is not None else 0 + start = target_version + while start > floor and _is_delta(_version_dir(source_dir, start)): + start -= 1 + if applied is None or start > applied: + seed_dir = base_dir if start == 0 else _version_dir(source_dir, start) + _reset_checkpoint(seed_dir, local_checkpoint_dir, start) + else: + start = applied + for version in range(start + 1, target_version + 1): + _apply_delta(local_checkpoint_dir, _version_dir(source_dir, version)) + + +def _load_hook(path: str): + module_path, _, name = path.rpartition(".") + return getattr(importlib.import_module(module_path), name) + + +def _version_dir(source_dir: str, version: int) -> str: + return os.path.join(source_dir, f"weight_v{version:06d}") + + +def _is_delta(version_dir: str) -> bool: + """A version is a delta iff its index metadata declares an encoding; an + ordinary HF checkpoint (with or without an index) is a full version.""" + if not os.path.isdir(version_dir): + raise FileNotFoundError(f"published weight version missing: {version_dir}") + try: + with open(os.path.join(version_dir, "model.safetensors.index.json")) as f: + return "delta_encoding" in json.load(f).get("metadata", {}) + except FileNotFoundError: + return False + + +class _Adler32: + """adler32 behind the incremental .update / .hexdigest interface the hash objects expose.""" + + def __init__(self): + self._value = 1 + + def update(self, data) -> None: + self._value = zlib.adler32(data, self._value) + + def hexdigest(self) -> str: + return f"{self._value:08x}" + + +def _new_hasher(algorithm: str): + if algorithm == "xxh3-128": + import xxhash + + return xxhash.xxh3_128() + if algorithm == "blake3": + import blake3 + + return blake3.blake3() + if algorithm == "adler32": + return _Adler32() + raise KeyError(f"unknown checksum algorithm {algorithm!r}") + + +def _checksum(algorithm: str, buf) -> str: + hasher = _new_hasher(algorithm) + hasher.update(buf) + return hasher.hexdigest() + + +@contextmanager +def _pull_lock(local_checkpoint_dir: str): + sync = os.path.join(local_checkpoint_dir, SYNC_DIR) + os.makedirs(sync, exist_ok=True) + with open(os.path.join(sync, "lock"), "w") as f: + fcntl.flock(f, fcntl.LOCK_EX) + try: + yield + finally: + fcntl.flock(f, fcntl.LOCK_UN) + + +def _read_applied_version(local_checkpoint_dir: str) -> Optional[int]: + try: + with open(os.path.join(local_checkpoint_dir, SYNC_DIR, "state.json")) as f: + return int(json.load(f)["version"]) + except FileNotFoundError: + return None + + +def _write_applied_version(local_checkpoint_dir: str, version: int) -> None: + path = os.path.join(local_checkpoint_dir, SYNC_DIR, "state.json") + tmp = path + ".tmp" + with open(tmp, "w") as f: + json.dump({"version": f"{version:06d}"}, f) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp, path) + + +def _drop_page_cache(path: str) -> None: + """Evict a file from the page cache (POSIX_FADV_DONTNEED).""" + if not hasattr(os, "posix_fadvise"): # POSIX-only (absent on macOS/Windows) + return + try: + fd = os.open(path, os.O_RDONLY) + try: + os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED) + finally: + os.close(fd) + except OSError: + pass + + +def _reset_checkpoint(src_dir: str, local_checkpoint_dir: str, version: int) -> None: + """Make local_checkpoint_dir an exact copy of the full checkpoint in src_dir + (files the new checkpoint doesn't have — e.g. differently-sharded old ones — + are pruned). Later deltas chain on top of this state.""" + logger.info( + "Pulling full checkpoint v%d %s -> %s", version, src_dir, local_checkpoint_dir + ) + os.makedirs(local_checkpoint_dir, exist_ok=True) + src_files = [entry for entry in os.scandir(src_dir) if entry.is_file()] + for entry in src_files: + shutil.copy2(entry.path, os.path.join(local_checkpoint_dir, entry.name)) + # don't let the source evict the local copy we keep resident + _drop_page_cache(entry.path) + names = {entry.name for entry in src_files} + for entry in os.scandir(local_checkpoint_dir): + if entry.is_file() and entry.name not in names: + os.remove(entry.path) + # a truncated copy (e.g. an object-store mount surfacing metadata before + # bytes) must fail loud, not serve bad weights + for entry in src_files: + copied = os.path.getsize(os.path.join(local_checkpoint_dir, entry.name)) + if copied != entry.stat().st_size: + raise RuntimeError( + f"size mismatch copying {entry.name}: src {entry.stat().st_size} != local {copied}" + ) + _write_applied_version(local_checkpoint_dir, version) + + +def _tensor_locations(ckpt_dir: str) -> dict: + """Map each tensor name to (file, byte offset, nbytes) by reading every safetensors header.""" + locations = {} + for path in glob.glob(os.path.join(ckpt_dir, "*.safetensors")): + with open(path, "rb") as f: + (header_len,) = struct.unpack(" None: + """Apply one version's delta in place: decompress + apply + checksum each tensor across a thread + pool (each writes a distinct mmap region, so the writes don't conflict). Any mismatch raises. + """ + with open(os.path.join(version_dir, "model.safetensors.index.json")) as f: + meta = json.load(f)["metadata"] + applied = _read_applied_version(local_checkpoint_dir) + if applied == int(meta["version"]): + return + if applied != int(meta["base_version"]): + raise RuntimeError( + f"out-of-order delta: local at {applied}, delta builds on {meta['base_version']}" + ) + if meta["compression_format"] != "zstd": + raise NotImplementedError( + f"compression {meta['compression_format']!r} not supported" + ) + encoding = meta["delta_encoding"] + algorithm = meta["checksum_format"] + locations = _tensor_locations(local_checkpoint_dir) + open_mmaps = {} + mismatches = [] + lock = threading.Lock() + file_bytes = [] # keep alive: items hold zero-copy views into these + items = [] # (name, compressed_view, path, offset, nbytes, want_checksum) + try: + for delta_file in sorted(glob.glob(os.path.join(version_dir, "*.safetensors"))): + with open(delta_file, "rb") as f: + blob = f.read() + file_bytes.append(blob) + (header_len,) = struct.unpack(" None: + name, compressed, path, offset, nbytes, want = item + region = np.ndarray( + (nbytes,), dtype=np.uint8, buffer=open_mmaps[path][1], offset=offset + ) + hasher = _new_hasher(algorithm) + reader = zstandard.ZstdDecompressor().stream_reader(compressed) + pos = 0 + # 2 MB chunks stay L2-resident across decompress -> XOR -> checksum + while pos < nbytes: + block = reader.read(min(2 << 20, nbytes - pos)) + if not block: + break + chunk = np.frombuffer(block, dtype=np.uint8) + region[pos : pos + chunk.size] ^= chunk + hasher.update(region[pos : pos + chunk.size]) + pos += chunk.size + if hasher.hexdigest() != want: + with lock: + mismatches.append(name) + + def apply_overwrite(item) -> None: + name, compressed, path, offset, nbytes, want = item + delta = np.frombuffer( + zstandard.ZstdDecompressor().decompress(compressed), dtype=np.uint8 + ) + region = np.ndarray( + (nbytes,), dtype=np.uint8, buffer=open_mmaps[path][1], offset=offset + ) + count = int.from_bytes(delta[:4], "little") + positions = np.frombuffer(delta[4 : 4 + 4 * count], dtype=", +} + +#[pymethods] +impl EngineStateChangedCallback { + fn __call__(&self) { + // A full snapshot is built after the notification is received, so one + // pending signal is enough to represent any number of quick changes. + let _ = self.sender.try_send(()); + } +} + fn lock_or_recover<'a, T>(mutex: &'a Mutex, name: &'static str) -> MutexGuard<'a, T> { mutex.lock().unwrap_or_else(|poisoned| { tracing::warn!(mutex = name, "Recovering from poisoned gRPC bridge mutex"); @@ -288,6 +302,25 @@ impl PyBridge { }) } + pub fn is_pause(&self) -> PyResult { + Python::attach(|py| { + let result = self.runtime_handle.call_method0(py, "is_pause")?; + result.extract::(py) + }) + } + + pub fn set_engine_state_changed_callback(&self, sender: Sender<()>) -> PyResult<()> { + Python::attach(|py| { + let callback = Py::new(py, EngineStateChangedCallback { sender })?; + self.runtime_handle.call_method1( + py, + "set_engine_state_changed_callback", + (callback,), + )?; + Ok(()) + }) + } + /// Tokenize via Python (fallback when Rust tokenizer unavailable). pub fn tokenize_py(&self, text: &str, add_special_tokens: bool) -> PyResult { Python::attach(|py| { diff --git a/rust/sglang-grpc/src/server.rs b/rust/sglang-grpc/src/server.rs index c05736ce4bed..c1f939e7e115 100644 --- a/rust/sglang-grpc/src/server.rs +++ b/rust/sglang-grpc/src/server.rs @@ -1,11 +1,12 @@ use std::collections::HashMap; use std::pin::Pin; use std::sync::Arc; +use std::time::{SystemTime, UNIX_EPOCH}; use pyo3::PyErr; use pyo3::Python; use pyo3::exceptions::{PyTypeError, PyValueError}; -use tokio::sync::{Notify, mpsc::Receiver}; +use tokio::sync::{Notify, mpsc::Receiver, watch}; use tokio::time::{Duration, timeout}; use tokio_stream::Stream; use tokio_stream::wrappers::TcpListenerStream; @@ -21,11 +22,92 @@ use crate::utils::{ pub struct SglangServiceImpl { pub bridge: Arc, pub response_timeout: Duration, + engine_state: EngineStatePublisher, + stream_shutdown: watch::Receiver, } type StreamResult = Pin> + Send + 'static>>; pub const DEFAULT_RESPONSE_TIMEOUT_SECS: u64 = 300; +#[derive(Clone)] +struct EngineStatePublisher { + bridge: Arc, + instance_id: u64, + sender: watch::Sender, +} + +impl EngineStatePublisher { + async fn new(bridge: Arc) -> Result { + let instance_id = u64::try_from( + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_err(|error| { + Status::internal(format!("system clock before Unix epoch: {error}")) + })? + .as_nanos(), + ) + .map_err(|_| Status::internal("engine instance timestamp does not fit in uint64"))?; + let snapshot = build_engine_state_snapshot(bridge.clone(), instance_id, 1).await?; + let (sender, _) = watch::channel(snapshot); + Ok(Self { + bridge, + instance_id, + sender, + }) + } + + fn subscribe(&self) -> watch::Receiver { + self.sender.subscribe() + } + + async fn publish_current(&self) -> Result<(), Status> { + let revision = self.sender.borrow().revision + 1; + let snapshot = + build_engine_state_snapshot(self.bridge.clone(), self.instance_id, revision).await?; + tracing::info!( + instance_id = snapshot.instance_id, + revision = snapshot.revision, + healthy = snapshot.healthy, + is_pause = snapshot.is_pause, + "publishing SGLang engine state" + ); + self.sender.send_replace(snapshot); + Ok(()) + } +} + +async fn build_engine_state_snapshot( + bridge: Arc, + instance_id: u64, + revision: u64, +) -> Result { + let values = tokio::task::spawn_blocking(move || { + Ok::<_, PyErr>(( + bridge.health_check()?, + bridge.is_pause()?, + bridge.get_model_info()?, + bridge.get_server_info()?, + )) + }) + .await + .map_err(|error| Status::internal(format!("engine snapshot task failed: {error}")))? + .map_err(|error| pyerr_to_status(error, "Failed to build engine state snapshot"))?; + let (healthy, is_pause, model_json, server_json) = values; + Ok(proto::EngineStateSnapshot { + instance_id, + revision, + healthy, + is_pause, + model_info: Some(proto::GetModelInfoResponse { + model_path: extract_model_path(&model_json), + json_info: model_json, + }), + server_info: Some(proto::GetServerInfoResponse { + json_info: server_json, + }), + }) +} + /// 64 MiB — leaves headroom for multimodal inputs and OpenAI JSON pass-through bodies, /// well above tonic's 4 MiB decode default. pub const DEFAULT_GRPC_MAX_MESSAGE_SIZE: usize = 64 * 1024 * 1024; @@ -494,14 +576,12 @@ impl proto::sglang_service_server::SglangService for SglangServiceImpl { } // Fallback to Python - let json_str = tokio::task::spawn_blocking({ - let bridge = self.bridge.clone(); - let text = req.text.clone(); - move || bridge.tokenize_py(&text, add_special) - }) - .await - .map_err(|e| Status::internal(format!("Task join error: {}", e)))? - .map_err(|e| pyerr_to_status(e, "Tokenize failed"))?; + let text = req.text.clone(); + let json_str = self + .blocking_bridge_call("Tokenize failed", move |bridge| { + bridge.tokenize_py(&text, add_special) + }) + .await?; let v: serde_json::Value = serde_json::from_str(&json_str) .map_err(|e| Status::internal(format!("Failed to parse JSON response: {}", e)))?; @@ -539,14 +619,11 @@ impl proto::sglang_service_server::SglangService for SglangServiceImpl { } // Fallback to Python - let json_str = tokio::task::spawn_blocking({ - let bridge = self.bridge.clone(); - let tokens = req.tokens; - move || bridge.detokenize_py(tokens) - }) - .await - .map_err(|e| Status::internal(format!("Task join error: {}", e)))? - .map_err(|e| pyerr_to_status(e, "Detokenize failed"))?; + let json_str = self + .blocking_bridge_call("Detokenize failed", move |bridge| { + bridge.detokenize_py(req.tokens) + }) + .await?; let v: serde_json::Value = serde_json::from_str(&json_str) .map_err(|e| Status::internal(format!("Failed to parse JSON response: {}", e)))?; @@ -561,28 +638,55 @@ impl proto::sglang_service_server::SglangService for SglangServiceImpl { &self, _request: Request, ) -> Result, Status> { - let healthy = tokio::task::spawn_blocking({ - let bridge = self.bridge.clone(); - move || bridge.health_check() - }) - .await - .map_err(|e| Status::internal(format!("Task join error: {}", e)))? - .map_err(|e| pyerr_to_status(e, "Health check failed"))?; + let healthy = self + .blocking_bridge_call("Health check failed", PyBridge::health_check) + .await?; Ok(Response::new(proto::HealthCheckResponse { healthy })) } + type WatchEngineStateStream = StreamResult; + + async fn watch_engine_state( + &self, + _request: Request, + ) -> Result, Status> { + let mut receiver = self.engine_state.subscribe(); + let mut shutdown = self.stream_shutdown.clone(); + let stream = async_stream::stream! { + if *shutdown.borrow_and_update() { + return; + } + let initial = receiver.borrow_and_update().clone(); + yield Ok(initial); + loop { + tokio::select! { + biased; + result = shutdown.changed() => { + if result.is_err() || *shutdown.borrow_and_update() { + break; + } + } + result = receiver.changed() => { + if result.is_err() { + break; + } + let update = receiver.borrow_and_update().clone(); + yield Ok(update); + } + } + } + }; + Ok(Response::new(Box::pin(stream))) + } + async fn get_model_info( &self, _request: Request, ) -> Result, Status> { - let json_info = tokio::task::spawn_blocking({ - let bridge = self.bridge.clone(); - move || bridge.get_model_info() - }) - .await - .map_err(|e| Status::internal(format!("Task join error: {}", e)))? - .map_err(|e| pyerr_to_status(e, "Failed to get model info"))?; + let json_info = self + .blocking_bridge_call("Failed to get model info", PyBridge::get_model_info) + .await?; Ok(Response::new(proto::GetModelInfoResponse { model_path: extract_model_path(&json_info), @@ -594,13 +698,9 @@ impl proto::sglang_service_server::SglangService for SglangServiceImpl { &self, _request: Request, ) -> Result, Status> { - let json_info = tokio::task::spawn_blocking({ - let bridge = self.bridge.clone(); - move || bridge.get_server_info() - }) - .await - .map_err(|e| Status::internal(format!("Task join error: {}", e)))? - .map_err(|e| pyerr_to_status(e, "Failed to get server info"))?; + let json_info = self + .blocking_bridge_call("Failed to get server info", PyBridge::get_server_info) + .await?; Ok(Response::new(proto::GetServerInfoResponse { json_info })) } @@ -609,13 +709,9 @@ impl proto::sglang_service_server::SglangService for SglangServiceImpl { &self, _request: Request, ) -> Result, Status> { - let json_str = tokio::task::spawn_blocking({ - let bridge = self.bridge.clone(); - move || bridge.list_models() - }) - .await - .map_err(|e| Status::internal(format!("Task join error: {}", e)))? - .map_err(|e| pyerr_to_status(e, "Failed to list models"))?; + let json_str = self + .blocking_bridge_call("Failed to list models", PyBridge::list_models) + .await?; let models_arr: Vec = serde_json::from_str(&json_str) .map_err(|e| Status::internal(format!("Failed to parse models JSON: {}", e)))?; @@ -847,8 +943,20 @@ impl proto::sglang_service_server::SglangService for SglangServiceImpl { } } -// Helper methods for OpenAI pass-through RPCs. +// Shared RPC helpers. impl SglangServiceImpl { + async fn blocking_bridge_call(&self, context: &str, call: F) -> Result + where + T: Send + 'static, + F: FnOnce(&PyBridge) -> Result + Send + 'static, + { + let bridge = self.bridge.clone(); + tokio::task::spawn_blocking(move || call(&bridge)) + .await + .map_err(|e| Status::internal(format!("Task join error: {}", e)))? + .map_err(|e| pyerr_to_status(e, context)) + } + async fn openai_streaming_rpc( &self, request: Request, @@ -984,11 +1092,26 @@ pub async fn run_grpc_server( ) -> Result<(), Box> { let addr = listener.local_addr()?; let listener = tokio::net::TcpListener::from_std(listener)?; + let (state_changed_tx, mut state_changed_rx) = tokio::sync::mpsc::channel(1); + bridge.set_engine_state_changed_callback(state_changed_tx)?; + let engine_state = EngineStatePublisher::new(bridge.clone()).await?; + let (stream_shutdown_tx, stream_shutdown_rx) = watch::channel(false); let service = SglangServiceImpl { bridge, response_timeout, + engine_state: engine_state.clone(), + stream_shutdown: stream_shutdown_rx, }; + let monitor = tokio::spawn(async move { + while state_changed_rx.recv().await.is_some() { + while state_changed_rx.try_recv().is_ok() {} + if let Err(error) = engine_state.publish_current().await { + tracing::warn!(%error, "failed to publish SGLang engine state"); + } + } + }); + let max_message_size = resolve_max_message_size(); let svc = proto::sglang_service_server::SglangServiceServer::new(service) .max_decoding_message_size(max_message_size) @@ -996,13 +1119,16 @@ pub async fn run_grpc_server( tracing::info!("gRPC server listening on {}", addr); - tonic::transport::Server::builder() + let result = tonic::transport::Server::builder() .add_service(svc) .serve_with_incoming_shutdown(TcpListenerStream::new(listener), async move { shutdown.notified().await; + stream_shutdown_tx.send_replace(true); tracing::info!("gRPC server shutting down"); }) - .await?; + .await; + monitor.abort(); + result?; Ok(()) } diff --git a/sgl-model-gateway/bindings/python/pyproject.toml b/sgl-model-gateway/bindings/python/pyproject.toml index c44b1b96abb4..98a9fbe3d9de 100644 --- a/sgl-model-gateway/bindings/python/pyproject.toml +++ b/sgl-model-gateway/bindings/python/pyproject.toml @@ -32,6 +32,7 @@ dependencies = [ "setproctitle", "aiohttp", "orjson", + "pybase64", "uvicorn", "fastapi", ] diff --git a/sgl-model-gateway/bindings/python/src/sglang_router/mini_lb.py b/sgl-model-gateway/bindings/python/src/sglang_router/mini_lb.py index cb8feefadd80..6ae3c3ddfc75 100644 --- a/sgl-model-gateway/bindings/python/src/sglang_router/mini_lb.py +++ b/sgl-model-gateway/bindings/python/src/sglang_router/mini_lb.py @@ -14,6 +14,7 @@ import aiohttp import orjson +import pybase64 import uvicorn from fastapi import FastAPI, HTTPException from fastapi.responses import ORJSONResponse, Response, StreamingResponse @@ -34,6 +35,44 @@ def maybe_wrap_ipv6_address(address: str) -> str: return address +def _merge_routed_experts(prefill: dict, decode: dict): + if "routed_experts" not in prefill or "routed_experts" not in decode: + return False + + prefill_bytes = pybase64.b64decode(prefill["routed_experts"], validate=True) + decode_bytes = pybase64.b64decode(decode["routed_experts"], validate=True) + decode["routed_experts"] = pybase64.b64encode( + prefill_bytes + decode_bytes[len(prefill_bytes) :] + ).decode("utf-8") + return True + + +def _merge_input_token_logprobs(prefill_meta: dict, decode_meta: dict): + if ( + "input_token_logprobs" not in prefill_meta + or "input_token_logprobs" not in decode_meta + ): + return + + decode_meta["input_token_logprobs"] = ( + prefill_meta["input_token_logprobs"] + decode_meta["input_token_logprobs"] + ) + + +def _merge_prefill_json(prefill_json, decode_json): + if "meta_info" in prefill_json and "meta_info" in decode_json: + prefill_meta = prefill_json["meta_info"] + decode_meta = decode_json["meta_info"] + _merge_input_token_logprobs(prefill_meta, decode_meta) + _merge_routed_experts(prefill_meta, decode_meta) + + if "sglext" not in prefill_json: + return + + if "sglext" in decode_json: + _merge_routed_experts(prefill_json["sglext"], decode_json["sglext"]) + + class MiniLoadBalancer: def __init__( self, @@ -139,17 +178,13 @@ async def generate( # Wait for both responses to complete. Prefill should end first. prefill_response, decode_response = await asyncio.gather(*tasks) - if "return_logprob" in modified_request: + if ( + "return_logprob" in modified_request + or "return_routed_experts" in modified_request + ): prefill_json = await prefill_response.json() ret_json = await decode_response.json() - - # merge `meta_info.input_token_logprobs` from prefill to decode - if "meta_info" in ret_json: - if "input_token_logprobs" in ret_json["meta_info"]: - ret_json["meta_info"]["input_token_logprobs"] = ( - prefill_json["meta_info"]["input_token_logprobs"] - + ret_json["meta_info"]["input_token_logprobs"] - ) + _merge_prefill_json(prefill_json, ret_json) else: ret_json = await decode_response.json() diff --git a/test/manual/layers/test_layernorm.py b/test/manual/layers/test_layernorm.py index 717a6c018cba..8fa69df0b0d5 100644 --- a/test/manual/layers/test_layernorm.py +++ b/test/manual/layers/test_layernorm.py @@ -61,6 +61,26 @@ def test_rms_norm(self): ): self._run_rms_norm_test(*params) + def test_rms_norm_cuda_uses_native_for_fp32_weight(self): + torch.manual_seed(0) + + hidden_size = 256 + layer = RMSNorm( + hidden_size, + cast_x_before_out_mul=True, + weight_dtype=torch.float32, + override_orig_dtype=torch.float32, + ) + layer.weight.data.normal_(mean=1.0, std=0.1) + x = torch.randn(17, hidden_size, dtype=torch.bfloat16) / hidden_size + + with torch.inference_mode(): + ref_out = layer.forward_native(x) + out = layer(x) + + self.assertEqual(out.dtype, torch.float32) + self.assertTrue(torch.equal(out, ref_out)) + class TestGemmaRMSNorm(CustomTestCase): DTYPES = [torch.half, torch.bfloat16] diff --git a/test/registered/e2e/disaggregation/test_parallelism_context_integration.py b/test/registered/e2e/disaggregation/test_parallelism_context_integration.py new file mode 100644 index 000000000000..a2479d1cb1f8 --- /dev/null +++ b/test/registered/e2e/disaggregation/test_parallelism_context_integration.py @@ -0,0 +1,276 @@ +""" +Integration tests for ParallelismContext with real sglang servers. + +Tests that ParallelismContext can instantiate models with correct tensor parallel +sharding by comparing parameter names and sizes against a running sglang server. + +Run with: + pytest test/registered/distributed/test_parallelism_context_integration.py -v + +Full test suite (non-CI): + - TP=2 small model (Qwem2.5-1.5B-Instruct) + - EP=2 small MOE model (DeepSeek-Coder-V2-Lite-Instruct) + - MLA model with hybrid dp attention (DeepSeek-Coder-V2-Lite-Instruct) + +CI test (reduced): + - TP=2 small model only +""" + +import dataclasses +import gc +import sys +from typing import Dict, List, Tuple + +import pytest +import requests +import torch + +from sglang.srt.distributed.parallel_state import RankParallelismConfig +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci +from sglang.test.test_utils import ( + DEFAULT_MLA_MODEL_NAME_FOR_TEST, + DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN, + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + is_in_ci, + popen_launch_server, +) +from sglang.utils import terminate_process + +register_cuda_ci(est_time=145, stage="extra-a", runner_config="2-gpu-large") +register_amd_ci(est_time=72, suite="stage-b-test-2-gpu-large-amd") + + +def get_transfer_engine_info(url: str, rank: int) -> Dict: + """Get transfer engine info (parameter names and sizes) for a rank.""" + response = requests.get( + f"{url}/remote_instance_transfer_engine_info", + params={"rank": rank}, + ) + response.raise_for_status() + return response.json() + + +def get_parallelism_config(url: str, rank: int) -> Dict: + """Get parallelism config for a rank.""" + response = requests.get(f"{url}/parallelism_config", params={"rank": rank}) + response.raise_for_status() + return response.json() + + +def get_server_info(url: str) -> Dict: + """Get server info.""" + response = requests.get(f"{url}/server_info") + response.raise_for_status() + return response.json() + + +def verify_model_params_match_for_rank( + url: str, + rank: int, + server_info: Dict, + test_gpu_id: int, +): + """Verify model parameters match for a specific rank by recreating a model shard.""" + transfer_info = get_transfer_engine_info(url, rank) + server_weights_info = transfer_info["remote_instance_transfer_engine_info"][1] + + # Get parallelism config from running server + parallelism_config_data = get_parallelism_config(url, rank) + parallelism_config = RankParallelismConfig.from_dict(parallelism_config_data) + # Get server args from server info + from sglang.srt.server_args import ServerArgs + + valid_fields = {f.name for f in dataclasses.fields(ServerArgs)} + filtered_info = {k: v for k, v in server_info.items() if k in valid_fields} + filtered_info.pop("model_config", None) + server_args = ServerArgs(**filtered_info) + + from sglang.srt import server_args as server_args_module + from sglang.srt.distributed.parallel_state import ParallelismContext + + original_global_server_args = server_args_module._global_server_args + + try: + # In a Mock ParallelismContext, instantiate the model for this rank. + # Use a separate GPU (test_gpu_id) to avoid memory conflicts with the running server. + server_args_module._global_server_args = server_args + with ParallelismContext(parallelism_config): + from sglang.srt.configs.device_config import DeviceConfig + from sglang.srt.configs.load_config import LoadConfig + from sglang.srt.configs.model_config import ModelConfig + from sglang.srt.model_loader import get_model + + model_config = ModelConfig.from_server_args(server_args) + load_config = LoadConfig(load_format="dummy") + device_config = DeviceConfig(device="cuda", gpu_id=test_gpu_id) + + torch.cuda.set_device(test_gpu_id) + model = get_model( + model_config=model_config, + load_config=load_config, + device_config=device_config, + ) + model_params = {} + for name, param in model.named_parameters(): + model_params[name] = param.numel() * param.element_size() + + # Verify all server parameters exist in model with same size + mismatches = [] + missing = [] + for param_name, (ptr, numel, elem_size) in server_weights_info.items(): + expected_size = numel * elem_size + if param_name not in model_params: + missing.append(param_name) + elif model_params[param_name] != expected_size: + mismatches.append( + f"{param_name}: model={model_params[param_name]}, server={expected_size}" + ) + + assert not missing, f"Rank {rank}: Missing parameters: {missing}" + assert not mismatches, f"Rank {rank}: Size mismatches: {mismatches}" + del model + torch.cuda.empty_cache() + + finally: + server_args_module._global_server_args = original_global_server_args + + +TEST_CONFIGS: List[Tuple[str, str, int, List[str], int]] = [ + # Basic TP=2 test (CI only) + ( + "tp2_small", + DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN, + 2, + [], + 2, + ), + # EP=2: MoE experts split across 2 groups, moe_tp=1 per group + ( + "mla_ep2", + DEFAULT_MLA_MODEL_NAME_FOR_TEST, + 2, + ["--ep-size", "2"], + 2, + ), + ( + "mla_dp2_tp4", + DEFAULT_MLA_MODEL_NAME_FOR_TEST, + 4, + ["--enable-dp-attention", "--dp", "2"], + 4, + ), + ( + "mla_dp2_ep2_tp4", + DEFAULT_MLA_MODEL_NAME_FOR_TEST, + 4, + ["--enable-dp-attention", "--dp", "2", "--ep-size", "2"], + 4, + ), + ( + "mla_dp2_ep4_tp4", + DEFAULT_MLA_MODEL_NAME_FOR_TEST, + 4, + ["--enable-dp-attention", "--dp", "2", "--ep-size", "4"], + 4, + ), + ( + "mla_dp4_ep2_tp4", + DEFAULT_MLA_MODEL_NAME_FOR_TEST, + 4, + ["--enable-dp-attention", "--dp", "4", "--ep-size", "2"], + 4, + ), +] + + +def get_test_configs(): + if is_in_ci(): + return [TEST_CONFIGS[0]] + else: + return TEST_CONFIGS + + +def _get_test_params(): + """Generate pytest parameters based on test configs.""" + configs = get_test_configs() + params = [] + ids = [] + for ( + test_id, + model_name, + tp_size, + extra_args, + min_gpus, + ) in configs: + params.append( + pytest.param( + (model_name, tp_size, extra_args, min_gpus), + id=test_id, + ) + ) + return params + + +class TestParallelismContextIntegration: + """ + Test that ParallelismContext can instantiate models with the same + parameter names and sizes as the sglang server engine. + """ + + @pytest.mark.parametrize("config", _get_test_params()) + def test_model_instantiation_matches_server(self, config): + """ + Test that a model instantiated with ParallelismContext has the same + parameter names and sizes as the model in the sglang server. + + This test: + 1. Starts a server with specified parallelism config + 2. Gets transfer_engine_info for all ranks (contains param names and sizes) + 3. Gets parallelism_config and server_info + 4. Uses ParallelismContext to instantiate a model for each rank + 5. Compares the parameter names and sizes + """ + model_name, tp_size, extra_args, min_gpus = config + url = DEFAULT_URL_FOR_TEST + + # Need min_gpus for server + 1 extra GPU for test model instantiation + required_gpus = min_gpus + 1 + if torch.cuda.device_count() < required_gpus: + pytest.skip( + f"Need at least {required_gpus} GPUs (server={min_gpus} + test=1), have {torch.cuda.device_count()}" + ) + test_gpu_id = min_gpus # e.g., if server uses 0-1, test uses 2 + + # Build server args + other_args = [ + "--tp-size", + str(tp_size), + "--remote-instance-weight-loader-start-seed-via-transfer-engine", + "--trust-remote-code", + ] + other_args.extend(extra_args) + + process = None + try: + process = popen_launch_server( + model_name, + url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=other_args, + ) + server_info = get_server_info(url) + + for rank in range(tp_size): + verify_model_params_match_for_rank(url, rank, server_info, test_gpu_id) + + finally: + if process is not None: + terminate_process(process) + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/e2e/lora/test_moe_lora_tail_stamp.py b/test/registered/e2e/lora/test_moe_lora_tail_stamp.py new file mode 100644 index 000000000000..b10d53389331 --- /dev/null +++ b/test/registered/e2e/lora/test_moe_lora_tail_stamp.py @@ -0,0 +1,134 @@ +"""Unit tests for the MoE-LoRA DP-gathered tail stamping in _add_moe_lora_info. + +Under --enable-dp-attention the MoE runs on DP-GATHERED tokens, so the per-token +LoRA mapping must cover [0, moe_num_tokens) while the per-rank segments only fill +[0, num_tokens). These tests pin the Tier-1 (single loaded adapter) semantics of +the gathered tail [num_tokens, moe_num_tokens): + + * active rank (local requests use the single adapter, ``_single_loaded_buffer_id`` + armed): tail stamped with that adapter's buffer id; + * true idle rank (num_tokens == 0, ``_idle_rank_active_buffer_id`` armed): the + whole mapping stamped and the adapter enabled; + * base-only local batches, multi-adapter batches, and stale idle stamps on + token-bearing batches: tail stays -1 (adapter-disabled) — foreign tokens are + never given a delta this rank cannot attribute. + +The backend object is stubbed (only the fields _add_moe_lora_info reads are +populated) so the tests run hermetically without a server or dist groups; the +gathered length is forced via forward_batch.global_dp_buffer_len. + +Usage: + python -m pytest test/registered/unit/lora/test_moe_lora_tail_stamp.py -v +""" + +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci + +register_cuda_ci(est_time=9, stage="base-b", runner_config="1-gpu-small") +register_amd_ci(est_time=9, suite="stage-b-test-1-gpu-small-amd") + +import unittest +from types import SimpleNamespace + +import torch + +from sglang.srt.lora.backend.base_backend import BaseLoRABackend +from sglang.srt.lora.utils import LoRABatchInfo + +DEVICE = "cuda" if torch.cuda.is_available() else "cpu" +GATHERED = 16 # forced gathered length (> any local num_tokens used below) +# buffer slots: 0 = the None/base uid slot (rank 0), 1/2 = adapters (rank 8), 3 = free +RANKS = [0, 8, 8, 0] + + +def _forward_batch(num_tokens: int): + mode = SimpleNamespace( + is_extend=lambda: False, + is_idle=lambda: num_tokens == 0, + is_cuda_graph=lambda: False, + ) + return SimpleNamespace( + forward_mode=mode, + batch_size=num_tokens, # decode: 1 token per seq + extend_seq_lens_cpu=None, + global_dp_buffer_len=GATHERED, + global_num_tokens_cpu=[num_tokens, GATHERED - num_tokens], + ) + + +def _batch_info(weight_indices, num_tokens): + bs = len(weight_indices) + return LoRABatchInfo( + use_cuda_graph=False, + bs=bs, + num_segments=bs, + seg_indptr=torch.arange(bs + 1, dtype=torch.int32, device=DEVICE), + weight_indices=torch.tensor(weight_indices, dtype=torch.int32, device=DEVICE), + lora_ranks=torch.tensor(RANKS, dtype=torch.int64, device=DEVICE), + scalings=torch.ones(len(RANKS), dtype=torch.float, device=DEVICE), + max_len=1, + seg_lens=torch.ones(bs, dtype=torch.int32, device=DEVICE), + permutation=None, + expected_tokens=num_tokens, + ) + + +def _run(weight_indices, num_tokens, idle_bid=None, single_bid=None): + stub = SimpleNamespace( + is_moe_lora=True, + _idle_rank_active_buffer_id=idle_bid, + _single_loaded_buffer_id=single_bid, + ) + out = BaseLoRABackend._add_moe_lora_info( + stub, _forward_batch(num_tokens), _batch_info(weight_indices, num_tokens) + ) + info = out.moe_lora_info + return info.token_lora_mapping.tolist(), info.adapter_enabled.tolist() + + +class TestMoELoRATailStamp(unittest.TestCase): + def test_base_only_local_batch_keeps_disabled_tail(self): + # All local tokens on the base (None-uid) slot, nothing armed: the + # gathered tail must stay -1 and no adapter may be enabled. + mapping, enabled = _run([0, 0, 0, 0], num_tokens=4) + self.assertTrue(all(x == -1 for x in mapping[4:]), mapping) + self.assertEqual(enabled, [0, 0, 0, 0]) + + def test_stale_idle_stamp_ignored_on_token_bearing_batch(self): + # Defense in depth: an idle stamp must only be consumed when the rank is + # truly idle (num_tokens == 0), never on a batch with local tokens. + mapping, enabled = _run([0, 0, 0, 0], num_tokens=4, idle_bid=2) + self.assertTrue(all(x == -1 for x in mapping[4:]), mapping) + self.assertEqual(enabled[2], 0) + + def test_multi_adapter_batch_keeps_disabled_tail(self): + # Two adapters used locally: the foreign tokens' adapter identity is + # unknowable host-side, so the tail must stay -1 (no mis-stamping with + # whatever the last local token happened to use). + mapping, _ = _run([1, 1, 2, 2], num_tokens=4) + self.assertTrue(all(x == -1 for x in mapping[4:]), mapping) + + def test_idle_rank_stamps_whole_mapping_and_enables_adapter(self): + # True idle rank under Tier-1: the whole gathered mapping is stamped with + # the single loaded adapter and the adapter is enabled, so foreign tokens + # routed to this rank's experts get the LoRA delta. + mapping, enabled = _run([], num_tokens=0, idle_bid=2) + self.assertTrue(all(x == 2 for x in mapping), mapping) + self.assertEqual(enabled[2], 1) + + def test_active_rank_stamps_tail_with_single_adapter(self): + # Local tokens actively use the single loaded adapter: the tail is + # stamped with its buffer id (not a copy of an arbitrary local slot). + mapping, enabled = _run([2, 2], num_tokens=2, single_bid=2) + self.assertTrue(all(x == 2 for x in mapping[2:]), mapping) + self.assertEqual(enabled[2], 1) + + def test_single_adapter_loaded_but_local_base_keeps_disabled_tail(self): + # An adapter is loaded but this rank's local tokens are all base (the + # manager arms no stamp in this case): tail stays -1. + mapping, enabled = _run([0, 0], num_tokens=2) + self.assertTrue(all(x == -1 for x in mapping[2:]), mapping) + self.assertEqual(enabled, [0, 0, 0, 0]) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/e2e/rl/test_mxfp4_marlin_inplace_reload.py b/test/registered/e2e/rl/test_mxfp4_marlin_inplace_reload.py new file mode 100644 index 000000000000..7efef96d3139 --- /dev/null +++ b/test/registered/e2e/rl/test_mxfp4_marlin_inplace_reload.py @@ -0,0 +1,84 @@ +import torch + +from sglang.srt.layers.quantization.marlin_utils_fp4 import ( + prepare_moe_mxfp4_layer_for_marlin, +) +from sglang.srt.layers.quantization.mxfp4 import Mxfp4MoEMethod +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=120, stage="base-b", runner_config="1-gpu-large") + + +def test_mxfp4_marlin_inplace_reload_matches_initial_repack(): + """A reload must land in the original Parameters, byte-identical to a fresh repack.""" + device = torch.device("cuda") + layer = torch.nn.Module() + layer.orig_dtype = torch.bfloat16 + checkpoint_tensors = { + "w13_weight": torch.randint( + 0, 256, (2, 256, 128), dtype=torch.uint8, device=device + ), + "w13_weight_scale": torch.randint( + 120, 136, (2, 256, 8), dtype=torch.uint8, device=device + ), + "w13_weight_bias": torch.randn(2, 256, dtype=torch.bfloat16, device=device), + "w2_weight": torch.randint( + 0, 256, (2, 256, 64), dtype=torch.uint8, device=device + ), + "w2_weight_scale": torch.randint( + 120, 136, (2, 256, 4), dtype=torch.uint8, device=device + ), + "w2_weight_bias": torch.randn(2, 256, dtype=torch.bfloat16, device=device), + } + source_values = { + name: tensor.clone() for name, tensor in checkpoint_tensors.items() + } + for name, tensor in checkpoint_tensors.items(): + param = torch.nn.Parameter(tensor, requires_grad=False) + param.weight_loader = object() + layer.register_parameter(name, param) + + method = Mxfp4MoEMethod.__new__(Mxfp4MoEMethod) + method.use_marlin = True + method.use_mega_moe = False + layer._mxfp4_marlin_checkpoint_parameter_specs = { + name: (tuple(tensor.shape), tensor.dtype) + for name, tensor in checkpoint_tensors.items() + } + + prepare_moe_mxfp4_layer_for_marlin(layer) + + expected = { + name: getattr(layer, name).detach().clone() + for name in method._MARLIN_CHECKPOINT_PARAMETER_NAMES + } + layer._mxfp4_marlin_runtime_parameter_specs = { + name: (tuple(tensor.shape), tensor.dtype) for name, tensor in expected.items() + } + parameter_ids = { + name: id(getattr(layer, name)) + for name in method._MARLIN_CHECKPOINT_PARAMETER_NAMES + } + storage_ptrs = { + name: getattr(layer, name).untyped_storage().data_ptr() + for name in method._MARLIN_CHECKPOINT_PARAMETER_NAMES + } + + method.restore_weights_before_loading(layer) + for name, source in source_values.items(): + param = getattr(layer, name) + assert hasattr(param, "weight_loader"), name + param.data.copy_(source) + prepare_moe_mxfp4_layer_for_marlin(layer, reuse_parameter_storage=True) + + for name, expected_tensor in expected.items(): + actual = getattr(layer, name) + assert id(actual) == parameter_ids[name] + assert actual.untyped_storage().data_ptr() == storage_ptrs[name] + assert torch.equal(actual, expected_tensor) + + +if __name__ == "__main__": + import pytest + + raise SystemExit(pytest.main([__file__, "-x"])) diff --git a/test/registered/e2e/rl/test_weight_checker_mxfp4_marlin.py b/test/registered/e2e/rl/test_weight_checker_mxfp4_marlin.py new file mode 100644 index 000000000000..d9afa8d1683a --- /dev/null +++ b/test/registered/e2e/rl/test_weight_checker_mxfp4_marlin.py @@ -0,0 +1,232 @@ +"""`Mxfp4MarlinComparable` must see through the Marlin repack and through equivalent MXFP4 encodings.""" + +import unittest + +import torch + +from sglang.srt.layers.quantization.marlin_utils_fp4 import ( + prepare_moe_mxfp4_layer_for_marlin, +) +from sglang.srt.layers.quantization.mxfp4 import Mxfp4MoEMethod +from sglang.srt.utils.weight_checker import _build_quantized_set +from sglang.srt.utils.weight_checker_comparator import ( + Mxfp4MarlinComparable, + compare_weights, + select_comparable_weight, +) +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=120, stage="base-b", runner_config="1-gpu-small") + +_E2M1_VALUES = torch.tensor((0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)) +# e2m1 code of twice the value, for codes whose value doubles within range +_DOUBLED_CODE = {0: 0, 1: 2, 2: 4, 3: 5, 4: 6, 5: 7} + +_NUM_EXPERTS, _HIDDEN, _INTERMEDIATE = 2, 256, 128 + + +def _random_checkpoint(device): + """Checkpoint-layout MXFP4 experts: packed nibbles `(E, N, K/2)` and e8m0 bytes `(E, N, K/32)`. + + Every other block keeps its values at or below 3 so it has a second, equivalent encoding. + """ + generator = torch.Generator(device="cpu").manual_seed(0) + + def packed(n, k): + magnitude = torch.randint( + 0, 8, (_NUM_EXPERTS, n, k // 32, 32), generator=generator + ) + magnitude[:, :, 0::2] %= 6 + sign = torch.randint(0, 2, magnitude.shape, generator=generator) << 3 + codes = (magnitude | sign).reshape(_NUM_EXPERTS, n, k).to(torch.uint8) + return codes[..., 0::2] | (codes[..., 1::2] << 4) + + def scales(n, k): + return torch.randint( + 118, 136, (_NUM_EXPERTS, n, k // 32), dtype=torch.uint8, generator=generator + ) + + return { + "w13_weight": packed(2 * _INTERMEDIATE, _HIDDEN).to(device), + "w13_weight_scale": scales(2 * _INTERMEDIATE, _HIDDEN).to(device), + "w2_weight": packed(_HIDDEN, _INTERMEDIATE).to(device), + "w2_weight_scale": scales(_HIDDEN, _INTERMEDIATE).to(device), + } + + +def _dequantize_checkpoint(packed, scales): + nibbles = torch.stack([packed & 0xF, packed >> 4], dim=-1).reshape( + *packed.shape[:-1], -1 + ) + values = _E2M1_VALUES.to(packed.device)[(nibbles & 0x7).long()] * ( + 1.0 - 2.0 * (nibbles >> 3).float() + ) + scale = torch.exp2(scales.float() - 127.0).repeat_interleave(32, dim=-1) + return (values * scale).to(torch.bfloat16) + + +def _repacked_layer(checkpoint): + layer = torch.nn.Module() + layer.orig_dtype = torch.bfloat16 + for name, tensor in checkpoint.items(): + layer.register_parameter( + name, torch.nn.Parameter(tensor.clone(), requires_grad=False) + ) + prepare_moe_mxfp4_layer_for_marlin(layer) + return layer + + +def _re_encode_eligible_blocks(packed, scales): + """Halve the scale and double the nibbles of the blocks that can take it: same values, new bytes.""" + nibbles = torch.stack([packed & 0xF, packed >> 4], dim=-1).reshape( + *packed.shape[:-1], -1, 32 + ) + magnitude = nibbles & 0x7 + eligible = (magnitude <= 5).all(dim=-1) + doubled = magnitude.clone() + for code, twice in _DOUBLED_CODE.items(): + doubled[magnitude == code] = twice + new_nibbles = torch.where( + eligible.unsqueeze(-1), doubled | (nibbles & 0x8), nibbles + ) + new_nibbles = new_nibbles.reshape(*packed.shape[:-1], -1) + new_packed = new_nibbles[..., 0::2] | (new_nibbles[..., 1::2] << 4) + new_scales = torch.where(eligible, scales - 1, scales) + assert eligible.any(), "the fixture must exercise at least one re-encoded block" + return new_packed.contiguous(), new_scales.contiguous() + + +class TestMxfp4MarlinComparable(CustomTestCase): + def setUp(self): + self.device = torch.device("cuda") + self.checkpoint = _random_checkpoint(self.device) + self.layer = _repacked_layer(self.checkpoint) + + def _comparable(self, layer, prefix): + return Mxfp4MarlinComparable( + getattr(layer, f"{prefix}_weight"), getattr(layer, f"{prefix}_weight_scale") + ) + + def test_dequantize_undoes_the_marlin_repack(self): + for prefix in ("w13", "w2"): + expected = _dequantize_checkpoint( + self.checkpoint[f"{prefix}_weight"], + self.checkpoint[f"{prefix}_weight_scale"], + ) + actual = self._comparable(self.layer, prefix).dequantize() + self.assertEqual(tuple(actual.shape), tuple(expected.shape), prefix) + self.assertTrue(torch.equal(actual, expected), prefix) + + def test_equivalent_encodings_compare_equal(self): + re_encoded = dict(self.checkpoint) + for prefix in ("w13", "w2"): + re_encoded[f"{prefix}_weight"], re_encoded[f"{prefix}_weight_scale"] = ( + _re_encode_eligible_blocks( + self.checkpoint[f"{prefix}_weight"], + self.checkpoint[f"{prefix}_weight_scale"], + ) + ) + self.assertFalse( + torch.equal( + re_encoded[f"{prefix}_weight"], self.checkpoint[f"{prefix}_weight"] + ) + ) + other = _repacked_layer(re_encoded) + for prefix in ("w13", "w2"): + self.assertFalse( + torch.equal( + getattr(other, f"{prefix}_weight"), + getattr(self.layer, f"{prefix}_weight"), + ) + ) + result = compare_weights( + self._comparable(self.layer, prefix), self._comparable(other, prefix) + ) + self.assertTrue(result.equal, (prefix, result)) + + def _compare_with_nibble_flipped(self, mask): + changed = dict(self.checkpoint) + changed["w2_weight"] = self.checkpoint["w2_weight"].clone() + changed["w2_weight"][0, 3, 5] ^= mask + other = _repacked_layer(changed) + return compare_weights( + self._comparable(self.layer, "w2"), self._comparable(other, "w2") + ) + + def test_a_changed_value_is_reported(self): + # flipping bit 2 of the e2m1 code moves the value by at least four ulps + result = self._compare_with_nibble_flipped(0x4) + self.assertFalse(result.equal) + self.assertGreater(result.num_exceed, 0) + self.assertGreater(result.max_abs_err, 0.0) + + def test_a_one_ulp_change_stays_within_the_quantization_tolerance(self): + result = self._compare_with_nibble_flipped(0x2) + self.assertFalse(result.equal) + self.assertEqual(result.num_exceed, 0) + + def test_accepts_the_checkers_shuffle_flag(self): + # _build_check_entries passes is_shuffled to every quantized comparable + Mxfp4MarlinComparable( + self.layer.w2_weight, self.layer.w2_weight_scale, is_shuffled=False + ) + with self.assertRaises(AssertionError): + Mxfp4MarlinComparable( + self.layer.w2_weight, self.layer.w2_weight_scale, is_shuffled=True + ) + + def test_chunking_matches_the_whole(self): + comparable = self._comparable(self.layer, "w13") + whole = comparable.dequantize() + chunks = torch.cat([dq for dq, _ in comparable.iter_chunks()]) + self.assertTrue(torch.equal(chunks, whole)) + + +class TestRouting(CustomTestCase): + def _method(self, use_marlin): + method = Mxfp4MoEMethod.__new__(Mxfp4MoEMethod) + method.use_marlin = use_marlin + return method + + def test_marlin_experts_get_the_comparable(self): + self.assertIs( + select_comparable_weight(self._method(True)), Mxfp4MarlinComparable + ) + + def test_other_mxfp4_backends_are_rejected(self): + with self.assertRaises(NotImplementedError): + select_comparable_weight(self._method(False)) + + def test_quantized_set_pairs_weight_with_weight_scale(self): + experts = torch.nn.Module() + experts.quant_method = self._method(True) + for name in ( + "w13_weight", + "w13_weight_scale", + "w2_weight", + "w2_weight_scale", + "w2_weight_bias", + ): + experts.register_parameter( + name, torch.nn.Parameter(torch.zeros(1), requires_grad=False) + ) + model = torch.nn.Module() + model.experts = experts + + quantized = _build_quantized_set(model) + + self.assertEqual( + {name: qw.scale_name for name, qw in quantized.items()}, + { + "experts.w13_weight": "experts.w13_weight_scale", + "experts.w2_weight": "experts.w2_weight_scale", + }, + ) + self.assertTrue( + all(qw.comparable_cls is Mxfp4MarlinComparable for qw in quantized.values()) + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/e2e/true_on_policy/test_tp_invariant_ops.py b/test/registered/e2e/true_on_policy/test_tp_invariant_ops.py new file mode 100644 index 000000000000..aaa4fab087c9 --- /dev/null +++ b/test/registered/e2e/true_on_policy/test_tp_invariant_ops.py @@ -0,0 +1,866 @@ +"""Tests for TP-invariant kernels (PR1). + +TP-invariance property: + Given the same TP degree and the same input data, matmul_tp_persistent + plus tree_all_reduce_sum produces bitwise identical results across runs. + When K/BLOCK_K is divisible by tp_size AND each shard yields a power-of-two + block count, TP=1 and TP=N also agree (isomorphic tree structure). + + For production K values (e.g. 3584, 5120) where block counts per shard are + not power-of-two, the tree structure differs between TP degrees. The + invariance guarantee is *determinism for a fixed TP degree*. + +All bitwise assertions use torch.equal, never approximate tolerances. +""" + +import os +import random +import unittest + +import torch +import torch.distributed as dist + +from sglang.srt.tp_invariant_ops import ( + disable_tp_invariant_mode, + enable_tp_invariant_mode, + is_tp_invariant_mode_enabled, + matmul_tp_inv, + matmul_tp_persistent, + moe_sum_tree_reduce, + set_tp_invariant_mode, + tree_all_reduce_sum, +) +from sglang.srt.tp_invariant_ops.tp_invariant_ops import ( + _MATMUL_K_BLOCK, + _fixed_tree_sum_tensors, + _is_power_of_two, +) +from sglang.test.ci.ci_register import register_cpu_ci, register_cuda_ci + +register_cpu_ci(est_time=12, suite="stage-a-test-cpu") +register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="8-gpu-h200") + +BLOCK_K = _MATMUL_K_BLOCK # 128 + + +def _simulate_tp_matmul(A, B, tp_size, **kwargs): + """Simulate what production TP does: each rank runs matmul_tp_persistent + on its K-shard, then tree_all_reduce_sum gathers and tree-sums.""" + K = A.shape[1] + shard = K // tp_size + partials = [] + for r in range(tp_size): + start = r * shard + end = start + shard + partials.append( + matmul_tp_persistent(A[:, start:end], B[start:end, :], **kwargs) + ) + return _fixed_tree_sum_tensors(partials) + + +# --------------------------------------------------------------------------- +# Mode flag +# --------------------------------------------------------------------------- +class TestTPInvariantMode(unittest.TestCase): + def tearDown(self): + disable_tp_invariant_mode() + + def test_mode_context_restores_previous_state(self): + disable_tp_invariant_mode() + self.assertFalse(is_tp_invariant_mode_enabled()) + + with set_tp_invariant_mode(True): + self.assertTrue(is_tp_invariant_mode_enabled()) + self.assertFalse(is_tp_invariant_mode_enabled()) + + enable_tp_invariant_mode() + with set_tp_invariant_mode(False): + self.assertFalse(is_tp_invariant_mode_enabled()) + self.assertTrue(is_tp_invariant_mode_enabled()) + + def test_enable_is_idempotent(self): + enable_tp_invariant_mode() + enable_tp_invariant_mode() + self.assertTrue(is_tp_invariant_mode_enabled()) + disable_tp_invariant_mode() + self.assertFalse(is_tp_invariant_mode_enabled()) + + def test_context_restores_after_exception(self): + disable_tp_invariant_mode() + try: + with set_tp_invariant_mode(True): + raise RuntimeError("deliberate") + except RuntimeError: + pass + self.assertFalse(is_tp_invariant_mode_enabled()) + + +# --------------------------------------------------------------------------- +# Reference ops: correctness +# --------------------------------------------------------------------------- +class TestTPInvariantReferenceOps(unittest.TestCase): + def test_fixed_tree_sum_order_is_stable(self): + values = [ + torch.tensor([1.0e20], dtype=torch.float32), + torch.tensor([1.0], dtype=torch.float32), + torch.tensor([-1.0e20], dtype=torch.float32), + torch.tensor([3.0], dtype=torch.float32), + ] + tree_result = _fixed_tree_sum_tensors(values) + sequential_result = values[0] + values[1] + values[2] + values[3] + + self.assertEqual(tree_result.item(), 0.0) + self.assertEqual(sequential_result.item(), 3.0) + + def test_matmul_tp_persistent_matches_torch_matmul_fp32(self): + torch.manual_seed(0) + A = torch.randn(5, 257, dtype=torch.float32) + B = torch.randn(257, 7, dtype=torch.float32) + bias = torch.randn(7, dtype=torch.float32) + + actual = matmul_tp_persistent(A, B, bias=bias) + expected = A @ B + bias + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) + + def test_matmul_tp_persistent_bf16_approximate_to_torch(self): + """BF16 block-tree matmul vs native torch matmul. Different accumulation + order means results differ by up to a few BF16 ULPs; this test only + checks that the magnitude is in the same ballpark, not bitwise equality.""" + torch.manual_seed(0) + A = torch.randn(6, 256, dtype=torch.bfloat16) + B = torch.randn(256, 10, dtype=torch.bfloat16) + bias = torch.randn(10, dtype=torch.bfloat16) + + actual = matmul_tp_persistent(A, B, bias=bias) + expected = A @ B + bias + torch.testing.assert_close(actual, expected, rtol=1e-1, atol=1e-1) + + def test_torch_custom_op_dispatches_to_matmul(self): + A = torch.arange(6, dtype=torch.float32).reshape(2, 3) + B = torch.arange(12, dtype=torch.float32).reshape(3, 4) + + actual = torch.ops.tp_inv_ops.matmul_tp_inv(A, B) + expected = A @ B + torch.testing.assert_close(actual, expected) + + def test_torch_custom_op_dispatches_with_bias(self): + torch.manual_seed(99) + A = torch.randn(4, 128, dtype=torch.float32) + B = torch.randn(128, 8, dtype=torch.float32) + bias = torch.randn(8, dtype=torch.float32) + + actual = torch.ops.tp_inv_ops.matmul_tp_inv(A, B, bias) + expected = matmul_tp_persistent(A, B, bias=bias) + self.assertTrue(torch.equal(actual, expected)) + + def test_matmul_tp_inv_matches_persistent(self): + torch.manual_seed(11) + A = torch.randn(4, 256, dtype=torch.bfloat16) + B = torch.randn(256, 16, dtype=torch.bfloat16) + self.assertTrue(torch.equal(matmul_tp_inv(A, B), matmul_tp_persistent(A, B))) + + def test_moe_sum_tree_reduce_matches_expert_order_reference(self): + input_tensor = torch.tensor( + [ + [ + [1.0e20, 1.0], + [1.0, 2.0], + [-1.0e20, 4.0], + [3.0, 8.0], + ], + [ + [5.0, 7.0], + [11.0, 13.0], + [17.0, 19.0], + [23.0, 29.0], + ], + ], + dtype=torch.float32, + ) + curr_topk_ids = torch.tensor( + [[0, 1, 2, 3], [3, -1, 1, 0]], + dtype=torch.int64, + ) + output = torch.empty(2, 2, dtype=torch.float32) + + moe_sum_tree_reduce( + input=input_tensor, + output=output, + curr_topk_ids=curr_topk_ids, + routed_scaling_factor=0.5, + E=4, + ) + + expected = torch.tensor([[0.0, 7.5], [22.5, 27.5]], dtype=torch.float32) + torch.testing.assert_close(output, expected) + + def test_moe_sum_tree_reduce_rejects_non_power_of_two_expert_count(self): + with self.assertRaisesRegex(ValueError, "power of two"): + moe_sum_tree_reduce( + input=torch.zeros(1, 1, 2), + output=torch.zeros(1, 2), + curr_topk_ids=torch.zeros(1, 1, dtype=torch.int64), + routed_scaling_factor=1.0, + E=3, + ) + + +# --------------------------------------------------------------------------- +# Matmul TP invariance: the core contract +# +# Two classes of tests: +# 1. Cross-TP: TP=1 == TP=N, requires isomorphic tree (K/BLOCK_K power-of-two +# multiple of tp_size). Uses torch.equal. +# 2. Determinism: same TP degree, two runs == bitwise identical. Works for +# any valid K alignment. +# --------------------------------------------------------------------------- +class TestTPInvarianceCrossTP(unittest.TestCase): + """TP=1 == TP=N when the binary tree structure is isomorphic.""" + + def test_fp32_cross_tp_k512(self): + torch.manual_seed(42) + A = torch.randn(8, 512, dtype=torch.float32) + B = torch.randn(512, 16, dtype=torch.float32) + + result_tp1 = matmul_tp_persistent(A, B) + result_tp2 = _simulate_tp_matmul(A, B, tp_size=2) + result_tp4 = _simulate_tp_matmul(A, B, tp_size=4) + + self.assertTrue(torch.equal(result_tp1, result_tp2)) + self.assertTrue(torch.equal(result_tp1, result_tp4)) + + def test_bf16_cross_tp_k512(self): + torch.manual_seed(42) + A = torch.randn(8, 512, dtype=torch.bfloat16) + B = torch.randn(512, 16, dtype=torch.bfloat16) + + result_tp1 = matmul_tp_persistent(A, B) + result_tp2 = _simulate_tp_matmul(A, B, tp_size=2) + result_tp4 = _simulate_tp_matmul(A, B, tp_size=4) + + self.assertTrue(torch.equal(result_tp1, result_tp2)) + self.assertTrue(torch.equal(result_tp1, result_tp4)) + + def test_bf16_cross_tp_k1024(self): + torch.manual_seed(7) + A = torch.randn(4, 1024, dtype=torch.bfloat16) + B = torch.randn(1024, 32, dtype=torch.bfloat16) + + result_tp1 = matmul_tp_persistent(A, B) + result_tp2 = _simulate_tp_matmul(A, B, tp_size=2) + result_tp4 = _simulate_tp_matmul(A, B, tp_size=4) + result_tp8 = _simulate_tp_matmul(A, B, tp_size=8) + + self.assertTrue(torch.equal(result_tp1, result_tp2)) + self.assertTrue(torch.equal(result_tp1, result_tp4)) + self.assertTrue(torch.equal(result_tp1, result_tp8)) + + def test_bf16_cross_tp_k2048_all_sizes(self): + """K=2048: 16 blocks. TP={1,2,4,8,16} all produce isomorphic trees.""" + torch.manual_seed(102) + A = torch.randn(4, 2048, dtype=torch.bfloat16) + B = torch.randn(2048, 16, dtype=torch.bfloat16) + + result_tp1 = matmul_tp_persistent(A, B) + for tp_size in [2, 4, 8, 16]: + result_tpN = _simulate_tp_matmul(A, B, tp_size=tp_size) + self.assertTrue( + torch.equal(result_tp1, result_tpN), + f"TP=1 != TP={tp_size} for K=2048", + ) + + def test_bf16_cross_tp_k4096(self): + """K=4096: 32 blocks.""" + torch.manual_seed(103) + A = torch.randn(2, 4096, dtype=torch.bfloat16) + B = torch.randn(4096, 8, dtype=torch.bfloat16) + + result_tp1 = matmul_tp_persistent(A, B) + for tp_size in [2, 4, 8]: + result_tpN = _simulate_tp_matmul(A, B, tp_size=tp_size) + self.assertTrue( + torch.equal(result_tp1, result_tpN), + f"TP=1 != TP={tp_size} for K=4096", + ) + + def test_fp16_cross_tp_k512(self): + torch.manual_seed(500) + A = torch.randn(4, 512, dtype=torch.float16) + B = torch.randn(512, 16, dtype=torch.float16) + + result_tp1 = matmul_tp_persistent(A, B) + result_tp4 = _simulate_tp_matmul(A, B, tp_size=4) + + self.assertTrue(torch.equal(result_tp1, result_tp4)) + + def test_torch_ops_dispatch_bf16_cross_tp(self): + """torch.ops.tp_inv_ops.matmul_tp_inv dispatch preserves cross-TP invariance.""" + torch.manual_seed(400) + A = torch.randn(4, 512, dtype=torch.bfloat16) + B = torch.randn(512, 16, dtype=torch.bfloat16) + + result_full = torch.ops.tp_inv_ops.matmul_tp_inv(A, B) + + K = A.shape[1] + shard = K // 4 + partials = [] + for r in range(4): + s, e = r * shard, (r + 1) * shard + partials.append(torch.ops.tp_inv_ops.matmul_tp_inv(A[:, s:e], B[s:e, :])) + result_tp4 = _fixed_tree_sum_tensors(partials) + + self.assertTrue(torch.equal(result_full, result_tp4)) + + +class TestTPInvarianceDeterminism(unittest.TestCase): + """Same TP degree, two runs -> bitwise identical. Works for all production K.""" + + def _assert_deterministic(self, K, tp_size, dtype=torch.bfloat16): + torch.manual_seed(42) + A = torch.randn(4, K, dtype=dtype) + B = torch.randn(K, 16, dtype=dtype) + + result_a = _simulate_tp_matmul(A, B, tp_size=tp_size) + result_b = _simulate_tp_matmul(A, B, tp_size=tp_size) + self.assertTrue( + torch.equal(result_a, result_b), + f"K={K} TP={tp_size} dtype={dtype} not deterministic", + ) + + def test_bf16_k3584_tp2(self): + """Qwen3-4B hidden_size. 3584/2=1792 -> 14 blocks per shard.""" + self._assert_deterministic(3584, tp_size=2) + + def test_bf16_k3584_tp4(self): + self._assert_deterministic(3584, tp_size=4) + + def test_bf16_k5120_tp4(self): + """Qwen3-30B hidden_size. 5120/4=1280 -> 10 blocks per shard.""" + self._assert_deterministic(5120, tp_size=4) + + def test_bf16_k5120_tp8(self): + self._assert_deterministic(5120, tp_size=8) + + def test_bf16_k4096_tp8(self): + self._assert_deterministic(4096, tp_size=8) + + def test_fp32_k3584_tp2(self): + self._assert_deterministic(3584, tp_size=2, dtype=torch.float32) + + def test_fp32_accum_deterministic(self): + """fp32_accum=True is deterministic for a fixed TP degree.""" + torch.manual_seed(200) + A = torch.randn(4, 512, dtype=torch.bfloat16) + B = torch.randn(512, 16, dtype=torch.bfloat16) + + result_a = _simulate_tp_matmul(A, B, tp_size=2, fp32_accum=True) + result_b = _simulate_tp_matmul(A, B, tp_size=2, fp32_accum=True) + self.assertTrue(torch.equal(result_a, result_b)) + + def test_fp32_accum_output_dtype_is_input_dtype(self): + torch.manual_seed(300) + A = torch.randn(4, 512, dtype=torch.bfloat16) + B = torch.randn(512, 16, dtype=torch.bfloat16) + + result = matmul_tp_persistent(A, B, fp32_accum=True) + self.assertEqual(result.dtype, torch.bfloat16) + + +# --------------------------------------------------------------------------- +# BFloat16 ops: approximate correctness vs torch +# --------------------------------------------------------------------------- +class TestBFloat16Ops(unittest.TestCase): + def test_moe_sum_tree_reduce_bf16(self): + input_tensor = torch.tensor( + [[[1.0, 2.0], [3.0, 4.0]]], + dtype=torch.bfloat16, + ) + curr_topk_ids = torch.tensor([[0, 1]], dtype=torch.int64) + output = torch.empty(1, 2, dtype=torch.bfloat16) + + moe_sum_tree_reduce( + input=input_tensor, + output=output, + curr_topk_ids=curr_topk_ids, + routed_scaling_factor=1.0, + E=2, + ) + + expected = torch.tensor([[4.0, 6.0]], dtype=torch.bfloat16) + torch.testing.assert_close(output, expected) + + +# --------------------------------------------------------------------------- +# MoE tree reduce: EP/slot invariance +# --------------------------------------------------------------------------- +class TestMoeReduceSlotInvariance(unittest.TestCase): + """moe_sum_tree_reduce must produce bitwise identical results regardless + of which topk slot an expert appears in. This is the EP invariance + property: different EP ranks may route the same experts to different + slot positions, but the tree-reduce result must be bitwise identical.""" + + def _run_moe_reduce(self, input_tensor, topk_ids, E, scaling=1.0): + output = torch.zeros( + input_tensor.shape[0], input_tensor.shape[2], dtype=input_tensor.dtype + ) + moe_sum_tree_reduce( + input=input_tensor, + output=output, + curr_topk_ids=topk_ids, + routed_scaling_factor=scaling, + E=E, + ) + return output + + def _permute_slots(self, values, ids_src, ids_dst, topk): + """Rearrange values so that each expert's data moves to its new slot.""" + tokens = values.shape[0] + result = torch.zeros_like(values) + for t in range(tokens): + for slot_dst in range(topk): + eid = ids_dst[t, slot_dst].item() + if eid == -1: + continue + slot_src = (ids_src[t] == eid).nonzero(as_tuple=True)[0].item() + result[t, slot_dst] = values[t, slot_src] + return result + + def test_slot_permutation_gives_identical_result_fp32(self): + torch.manual_seed(10) + H = 64 + values = torch.randn(1, 4, H, dtype=torch.float32) + + ids_a = torch.tensor([[0, 1, 2, 3]], dtype=torch.int64) + ids_b = torch.tensor([[2, 0, 3, 1]], dtype=torch.int64) + + input_b = self._permute_slots(values, ids_a, ids_b, topk=4) + + result_a = self._run_moe_reduce(values, ids_a, E=4) + result_b = self._run_moe_reduce(input_b, ids_b, E=4) + self.assertTrue(torch.equal(result_a, result_b)) + + def test_slot_permutation_gives_identical_result_bf16(self): + torch.manual_seed(20) + H = 128 + values = torch.randn(2, 8, H, dtype=torch.bfloat16) + + expert_ids = list(range(8)) + rng = random.Random(42) + shuffled = expert_ids.copy() + rng.shuffle(shuffled) + + ids_a = torch.tensor([expert_ids, expert_ids], dtype=torch.int64) + ids_b = torch.tensor([shuffled, shuffled], dtype=torch.int64) + + input_b = self._permute_slots(values, ids_a, ids_b, topk=8) + + result_a = self._run_moe_reduce(values, ids_a, E=8) + result_b = self._run_moe_reduce(input_b, ids_b, E=8) + self.assertTrue(torch.equal(result_a, result_b)) + + def test_slot_invariance_with_remote_experts(self): + """Remote experts (-1) in different positions must not change the result.""" + torch.manual_seed(30) + H = 32 + values_a = torch.randn(1, 4, H, dtype=torch.float32) + + ids_a = torch.tensor([[0, -1, 2, -1]], dtype=torch.int64) + ids_b = torch.tensor([[-1, 2, -1, 0]], dtype=torch.int64) + + values_b = torch.randn(1, 4, H, dtype=torch.float32) + for slot_b in range(4): + eid = ids_b[0, slot_b].item() + if eid == -1: + continue + slot_a = (ids_a[0] == eid).nonzero(as_tuple=True)[0].item() + values_b[0, slot_b] = values_a[0, slot_a] + + result_a = self._run_moe_reduce(values_a, ids_a, E=4) + result_b = self._run_moe_reduce(values_b, ids_b, E=4) + self.assertTrue(torch.equal(result_a, result_b)) + + def test_moe_bf16_large_hidden_invariance(self): + """Production-scale hidden dim (H=7168) with E=64 in BF16.""" + torch.manual_seed(40) + H = 7168 + topk = 8 + E = 64 + tokens = 2 + + expert_ids = torch.zeros(tokens, topk, dtype=torch.int64) + for t in range(tokens): + chosen = torch.randperm(E)[:topk] + expert_ids[t] = chosen + + values = torch.randn(tokens, topk, H, dtype=torch.bfloat16) + + output_a = torch.zeros(tokens, H, dtype=torch.bfloat16) + moe_sum_tree_reduce( + input=values, + output=output_a, + curr_topk_ids=expert_ids, + routed_scaling_factor=0.25, + E=E, + ) + + perm = torch.randperm(topk) + values_b = values[:, perm, :] + ids_b = expert_ids[:, perm] + + output_b = torch.zeros(tokens, H, dtype=torch.bfloat16) + moe_sum_tree_reduce( + input=values_b, + output=output_b, + curr_topk_ids=ids_b, + routed_scaling_factor=0.25, + E=E, + ) + + self.assertTrue(torch.equal(output_a, output_b)) + + def test_moe_deterministic_two_runs(self): + """Same input, two calls -> bitwise identical.""" + torch.manual_seed(50) + values = torch.randn(4, 4, 256, dtype=torch.bfloat16) + ids = torch.tensor( + [[0, 1, 2, 3], [3, 2, 1, 0], [0, 0, 1, 1], [2, 3, 0, 1]], + dtype=torch.int64, + ) + + out_a = torch.zeros(4, 256, dtype=torch.bfloat16) + moe_sum_tree_reduce( + input=values, + output=out_a, + curr_topk_ids=ids, + routed_scaling_factor=0.5, + E=4, + ) + + out_b = torch.zeros(4, 256, dtype=torch.bfloat16) + moe_sum_tree_reduce( + input=values, + output=out_b, + curr_topk_ids=ids, + routed_scaling_factor=0.5, + E=4, + ) + + self.assertTrue(torch.equal(out_a, out_b)) + + +# --------------------------------------------------------------------------- +# tree_all_reduce_sum: non-distributed +# --------------------------------------------------------------------------- +class TestTreeAllReduceNonDistributed(unittest.TestCase): + def test_returns_clone_when_dist_not_initialized(self): + if dist.is_initialized(): + self.skipTest("dist already initialized") + x = torch.tensor([1.0, 2.0, 3.0]) + result = tree_all_reduce_sum(x) + self.assertTrue(torch.equal(x, result)) + self.assertFalse(x.data_ptr() == result.data_ptr()) + + def test_fixed_tree_sum_is_order_deterministic(self): + torch.manual_seed(50) + world_size = 8 + shards = [torch.randn(16, dtype=torch.bfloat16) for _ in range(world_size)] + + result_a = _fixed_tree_sum_tensors(shards) + result_b = _fixed_tree_sum_tensors(list(shards)) + self.assertTrue(torch.equal(result_a, result_b)) + + def test_tree_sum_power_of_two_sizes(self): + for n in [1, 2, 4, 8, 16]: + shards = [torch.tensor([float(i + 1)]) for i in range(n)] + result = _fixed_tree_sum_tensors(shards) + self.assertAlmostEqual(result.item(), n * (n + 1) / 2, places=5) + + +# --------------------------------------------------------------------------- +# MoE reduce: order-matters proof +# --------------------------------------------------------------------------- +class TestMoeReduceOrderMatters(unittest.TestCase): + def test_expert_order_differs_from_slot_order(self): + input_tensor = torch.tensor( + [[[1.0e16, 0.0], [1.0, 0.0], [-1.0e16, 0.0], [2.0, 0.0]]], + dtype=torch.float32, + ) + curr_topk_ids = torch.tensor([[2, 0, 3, 1]], dtype=torch.int64) + output = torch.empty(1, 2, dtype=torch.float32) + + moe_sum_tree_reduce( + input=input_tensor, + output=output, + curr_topk_ids=curr_topk_ids, + routed_scaling_factor=1.0, + E=4, + ) + + slot_order_sum = ( + input_tensor[0, 0] + + input_tensor[0, 1] + + input_tensor[0, 2] + + input_tensor[0, 3] + ) + self.assertNotEqual(output[0, 0].item(), slot_order_sum[0].item()) + + +# --------------------------------------------------------------------------- +# Edge cases +# --------------------------------------------------------------------------- +class TestEdgeCases(unittest.TestCase): + def test_fixed_tree_sum_single_tensor(self): + t = torch.tensor([5.0]) + result = _fixed_tree_sum_tensors([t]) + self.assertEqual(result.item(), 5.0) + + def test_fixed_tree_sum_odd_count(self): + values = [torch.tensor([1.0]), torch.tensor([2.0]), torch.tensor([3.0])] + result = _fixed_tree_sum_tensors(values) + self.assertEqual(result.item(), 6.0) + + def test_fixed_tree_sum_empty_raises(self): + with self.assertRaises(ValueError): + _fixed_tree_sum_tensors([]) + + def test_matmul_tp_persistent_k_less_than_block(self): + torch.manual_seed(0) + A = torch.randn(3, 64, dtype=torch.float32) + B = torch.randn(64, 5, dtype=torch.float32) + actual = matmul_tp_persistent(A, B) + expected = A @ B + torch.testing.assert_close(actual, expected) + + def test_matmul_tp_persistent_k_equals_block(self): + torch.manual_seed(0) + A = torch.randn(3, 128, dtype=torch.float32) + B = torch.randn(128, 5, dtype=torch.float32) + actual = matmul_tp_persistent(A, B) + expected = A @ B + torch.testing.assert_close(actual, expected) + + def test_matmul_tp_persistent_k_multi_block_non_aligned(self): + """K not divisible by BLOCK_K: reference handles remainder gracefully.""" + torch.manual_seed(0) + A = torch.randn(3, 300, dtype=torch.float32) + B = torch.randn(300, 5, dtype=torch.float32) + actual = matmul_tp_persistent(A, B) + expected = A @ B + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) + + def test_moe_sum_tree_reduce_single_expert(self): + input_tensor = torch.tensor([[[1.0, 2.0]]], dtype=torch.float32) + curr_topk_ids = torch.tensor([[0]], dtype=torch.int64) + output = torch.empty(1, 2, dtype=torch.float32) + + moe_sum_tree_reduce( + input=input_tensor, + output=output, + curr_topk_ids=curr_topk_ids, + routed_scaling_factor=1.0, + E=1, + ) + expected = torch.tensor([[1.0, 2.0]], dtype=torch.float32) + torch.testing.assert_close(output, expected) + + def test_moe_sum_tree_reduce_single_topk(self): + input_tensor = torch.tensor( + [[[10.0, 20.0]], [[30.0, 40.0]]], + dtype=torch.float32, + ) + curr_topk_ids = torch.tensor([[1], [0]], dtype=torch.int64) + output = torch.empty(2, 2, dtype=torch.float32) + + moe_sum_tree_reduce( + input=input_tensor, + output=output, + curr_topk_ids=curr_topk_ids, + routed_scaling_factor=0.5, + E=2, + ) + expected = torch.tensor([[5.0, 10.0], [15.0, 20.0]], dtype=torch.float32) + torch.testing.assert_close(output, expected) + + def test_moe_sum_tree_reduce_all_remote(self): + input_tensor = torch.tensor( + [[[99.0, 99.0], [99.0, 99.0]]], + dtype=torch.float32, + ) + curr_topk_ids = torch.tensor([[-1, -1]], dtype=torch.int64) + output = torch.empty(1, 2, dtype=torch.float32) + + moe_sum_tree_reduce( + input=input_tensor, + output=output, + curr_topk_ids=curr_topk_ids, + routed_scaling_factor=1.0, + E=2, + ) + expected = torch.zeros(1, 2, dtype=torch.float32) + torch.testing.assert_close(output, expected) + + def test_moe_sum_tree_reduce_duplicate_expert_ids(self): + """Same expert in multiple slots: both contributions must be summed.""" + input_tensor = torch.tensor( + [[[10.0, 20.0], [30.0, 40.0]]], + dtype=torch.float32, + ) + curr_topk_ids = torch.tensor([[0, 0]], dtype=torch.int64) + output = torch.empty(1, 2, dtype=torch.float32) + + moe_sum_tree_reduce( + input=input_tensor, + output=output, + curr_topk_ids=curr_topk_ids, + routed_scaling_factor=1.0, + E=2, + ) + expected = torch.tensor([[40.0, 60.0]], dtype=torch.float32) + torch.testing.assert_close(output, expected) + + +# --------------------------------------------------------------------------- +# Input validation +# --------------------------------------------------------------------------- +class TestInputValidation(unittest.TestCase): + def test_matmul_rejects_non_2d_inputs(self): + with self.assertRaisesRegex(ValueError, "expected 2D"): + matmul_tp_persistent(torch.randn(2, 3, 4), torch.randn(4, 5)) + with self.assertRaisesRegex(ValueError, "expected 2D"): + matmul_tp_persistent(torch.randn(2, 3), torch.randn(3)) + + def test_matmul_rejects_dimension_mismatch(self): + with self.assertRaisesRegex(ValueError, "dimension mismatch"): + matmul_tp_persistent(torch.randn(2, 3), torch.randn(4, 5)) + + def test_moe_rejects_wrong_input_dims(self): + with self.assertRaisesRegex(ValueError, "tokens, topk, hidden"): + moe_sum_tree_reduce( + input=torch.zeros(4, 8), + output=torch.zeros(4, 8), + curr_topk_ids=torch.zeros(4, 2, dtype=torch.int64), + routed_scaling_factor=1.0, + E=2, + ) + + def test_moe_rejects_wrong_topk_ids_dims(self): + with self.assertRaisesRegex(ValueError, "tokens, topk"): + moe_sum_tree_reduce( + input=torch.zeros(4, 2, 8), + output=torch.zeros(4, 8), + curr_topk_ids=torch.zeros(4, dtype=torch.int64), + routed_scaling_factor=1.0, + E=2, + ) + + def test_moe_rejects_shape_mismatch_between_input_and_ids(self): + with self.assertRaisesRegex(ValueError, "must match"): + moe_sum_tree_reduce( + input=torch.zeros(4, 2, 8), + output=torch.zeros(4, 8), + curr_topk_ids=torch.zeros(4, 3, dtype=torch.int64), + routed_scaling_factor=1.0, + E=2, + ) + + def test_moe_rejects_wrong_output_shape(self): + with self.assertRaisesRegex(ValueError, "output must have shape"): + moe_sum_tree_reduce( + input=torch.zeros(4, 2, 8), + output=torch.zeros(4, 4), + curr_topk_ids=torch.zeros(4, 2, dtype=torch.int64), + routed_scaling_factor=1.0, + E=2, + ) + + def test_is_power_of_two_helper(self): + self.assertTrue(_is_power_of_two(1)) + self.assertTrue(_is_power_of_two(2)) + self.assertTrue(_is_power_of_two(64)) + self.assertFalse(_is_power_of_two(0)) + self.assertFalse(_is_power_of_two(3)) + self.assertFalse(_is_power_of_two(6)) + + +# --------------------------------------------------------------------------- +# Distributed tree all-reduce (multi-GPU only) +# --------------------------------------------------------------------------- +class TestDistributedTreeAllReduce(unittest.TestCase): + @unittest.skipUnless( + int(os.environ.get("WORLD_SIZE", "1")) > 1, + "requires torchrun with WORLD_SIZE > 1", + ) + def test_tree_all_reduce_sum_distributed(self): + own_pg = False + if not dist.is_initialized(): + backend = "nccl" if torch.cuda.is_available() else "gloo" + dist.init_process_group(backend=backend) + own_pg = True + + try: + world_size = dist.get_world_size() + if world_size & (world_size - 1) != 0: + self.skipTest("requires power-of-two world size") + + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + if torch.cuda.is_available(): + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + else: + device = torch.device("cpu") + + rank = dist.get_rank() + value = torch.full((4,), float(rank + 1), device=device) + actual = tree_all_reduce_sum(value) + expected = torch.full( + (4,), + float(world_size * (world_size + 1) // 2), + device=device, + ) + + torch.testing.assert_close(actual, expected) + dist.barrier() + finally: + if own_pg: + dist.destroy_process_group() + + @unittest.skipUnless( + int(os.environ.get("WORLD_SIZE", "1")) > 1, + "requires torchrun with WORLD_SIZE > 1", + ) + def test_tree_all_reduce_bf16_bitwise_deterministic(self): + """Run tree all-reduce twice with same inputs, verify bitwise identical.""" + own_pg = False + if not dist.is_initialized(): + backend = "nccl" if torch.cuda.is_available() else "gloo" + dist.init_process_group(backend=backend) + own_pg = True + + try: + world_size = dist.get_world_size() + if world_size & (world_size - 1) != 0: + self.skipTest("requires power-of-two world size") + + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + if torch.cuda.is_available(): + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + else: + device = torch.device("cpu") + + rank = dist.get_rank() + torch.manual_seed(rank * 1000 + 777) + value = torch.randn(256, device=device, dtype=torch.bfloat16) + + result_a = tree_all_reduce_sum(value) + result_b = tree_all_reduce_sum(value) + + self.assertTrue(torch.equal(result_a, result_b)) + dist.barrier() + finally: + if own_pg: + dist.destroy_process_group() + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/kernels/ops/moe/test_moe_wna16_marlin.py b/test/registered/kernels/ops/moe/test_moe_wna16_marlin.py index 09700073d407..419ea6b3cc8b 100644 --- a/test/registered/kernels/ops/moe/test_moe_wna16_marlin.py +++ b/test/registered/kernels/ops/moe/test_moe_wna16_marlin.py @@ -352,7 +352,8 @@ def test_moe_wna16_marlin_gemm( not (is_sm80_supported() or is_sm90_supported()), reason="Non-gated NVFP4 Marlin fallback test requires CUDA SM8X/SM9X", ) -def test_fused_marlin_moe_non_gated_relu2(): +@pytest.mark.parametrize("compiled", [False, True]) +def test_fused_marlin_moe_non_gated_relu2(compiled): torch.manual_seed(0) m = 17 @@ -376,7 +377,12 @@ def test_fused_marlin_moe_non_gated_relu2(): score_softmax = torch.softmax(router_logits, dim=-1, dtype=torch.float32) topk_weights, topk_ids = torch.topk(score_softmax, topk) - output = fused_marlin_moe( + run_moe = ( + torch.compile(fused_marlin_moe, backend="eager", fullgraph=True) + if compiled + else fused_marlin_moe + ) + output = run_moe( hidden_states=hidden_states, w1=qweight1, w2=qweight2, diff --git a/test/registered/lora/test_virtual_experts_kernels.py b/test/registered/lora/test_virtual_experts_kernels.py index e6c8d6596874..5b685aaf95a0 100644 --- a/test/registered/lora/test_virtual_experts_kernels.py +++ b/test/registered/lora/test_virtual_experts_kernels.py @@ -124,6 +124,16 @@ def test_positive_topk_remapped_correctly(self): expected = base + max(lora, 0) * num_experts self.assertEqual(virtual_ids[m, k].item(), expected) + def test_out_of_range_expert_cannot_alias_another_adapter(self): + slots = torch.arange(4, dtype=torch.int32, device=self.device) + ids = torch.tensor([[255, 256]], dtype=torch.int32, device=self.device).repeat( + 4, 1 + ) + invalid, _, _ = _fused_virtual_topk_ids(ids, slots, 256, False, 4) + self.assertTrue(torch.all(invalid[:, 1] == -1)) + valid, _, _ = _fused_virtual_topk_ids(ids, slots, 257, False, 4) + torch.testing.assert_close(valid[:, 1], slots * 257 + 256) + def test_no_lora_token_does_not_shift_base(self): """`token_lora_mapping[m] == -1` (no LoRA) keeps `safe_lora=0`, so positive bases pass through unchanged and the row mask is False.""" diff --git a/test/registered/rl/test_lora_load_from_tensor.py b/test/registered/rl/test_lora_load_from_tensor.py index bb3a308ab2d2..3d6f63edbb34 100644 --- a/test/registered/rl/test_lora_load_from_tensor.py +++ b/test/registered/rl/test_lora_load_from_tensor.py @@ -1,6 +1,7 @@ import json import os import unittest +from types import SimpleNamespace import torch from huggingface_hub import snapshot_download @@ -22,6 +23,26 @@ MAX_NEW_TOKENS = 16 +def load_lora_via_stream(engine, lora_name, tensors, config_dict): + """Wire-load an adapter over the stream path: register (control plane) + + prefixed tensors through a sync_base=False weight-update session.""" + result = engine.register_lora_adapter(lora_name=lora_name, config_dict=config_dict) + if not result.success: + return SimpleNamespace(success=False, error_message=result.error_message) + engine.begin_weight_update(selector="all", sync_base=False) + engine.update_weights_from_tensor( + named_tensors=[ + (f"{lora_name}:{key}", tensor) for key, tensor in tensors.items() + ] + ) + success, message = engine.end_weight_update() + return SimpleNamespace(success=success, error_message=message) + + +def loaded_adapter_names(engine): + return list(engine.tokenizer_manager.lora_registry._registry.keys()) + + class TestLoRALoadFromTensor(CustomTestCase): @classmethod def setUpClass(cls): @@ -45,10 +66,10 @@ def setUpClass(cls): with open(os.path.join(lora_adapter, "adapter_config.json"), "r") as f: cls.lora_config_dict = json.load(f) - def test_lora_lru_eviction(self): - print("[Test]Testing LRU LoRA eviction...") + def test_lora_stream_register_rejected_over_cap(self): + """Streamed adapters have no path to reload from, so the cap rejects new + registrations instead of evicting; same-name upserts stay allowed.""" MAX_LOADED_LORAS = 8 - print(f"[Test]Max loaded LoRAs: {MAX_LOADED_LORAS}") test_engine = sgl.Engine( model_path=MODEL_PATH, enable_lora=True, @@ -59,56 +80,61 @@ def test_lora_lru_eviction(self): max_loaded_loras=MAX_LOADED_LORAS, ) - # Load 10 LoRA adapters, max allowed is 8 - # This should trigger LRU eviction when we exceed the limit TEST_LORA_COUNT = 10 for i in range(TEST_LORA_COUNT): - print(f"[Test]Loading LoRA adapter {i + 1}/10: self_cognition_Alice_{i}") - result = test_engine.load_lora_adapter_from_tensors( - lora_name=f"self_cognition_Alice_{i}", - tensors=self.lora_tensors, - config_dict=self.lora_config_dict, + result = load_lora_via_stream( + test_engine, + f"self_cognition_Alice_{i}", + self.lora_tensors, + self.lora_config_dict, ) - self.assertTrue( - result.success, - f"Failed to load LoRA adapter {i}: {result.error_message}", - ) - print( - f"[Test]Successfully loaded LoRA {i + 1}, current loaded adapters: {list(result.loaded_adapters.keys())}" - ) - - EXPECTED_LORA_ADAPTERS = [ - "self_cognition_Alice_2", - "self_cognition_Alice_3", - "self_cognition_Alice_4", - "self_cognition_Alice_5", - "self_cognition_Alice_6", - "self_cognition_Alice_7", - "self_cognition_Alice_8", - "self_cognition_Alice_9", - ] - EXPECTED_LORA_COUNT = 8 + if i < MAX_LOADED_LORAS: + self.assertTrue( + result.success, + f"Failed to load LoRA adapter {i}: {result.error_message}", + ) + else: + self.assertFalse( + result.success, + f"Adapter {i} should have been rejected over the cap", + ) + self.assertIn("max-loaded-loras", result.error_message) + + loaded = loaded_adapter_names(test_engine) self.assertEqual( - len(result.loaded_adapters), - EXPECTED_LORA_COUNT, - f"Loaded adapters count does not match expected result: {len(result.loaded_adapters)} != {EXPECTED_LORA_COUNT}", + loaded, + [f"self_cognition_Alice_{i}" for i in range(MAX_LOADED_LORAS)], + f"Loaded adapters do not match the first {MAX_LOADED_LORAS} registrations: {loaded}", ) - self.assertEqual( - list(result.loaded_adapters.keys()), - EXPECTED_LORA_ADAPTERS, - f"Loaded adapters do not match expected result: {list(result.loaded_adapters.keys())} != {EXPECTED_LORA_ADAPTERS}", + + # Same-name upsert must still pass at the cap. + result = load_lora_via_stream( + test_engine, + "self_cognition_Alice_0", + self.lora_tensors, + self.lora_config_dict, ) - print( - f"[Test]LRU eviction test passed! Final loaded adapters: {len(result.loaded_adapters)}" + self.assertTrue(result.success, f"Upsert at cap failed: {result.error_message}") + + def test_register_rejects_name_holding_the_separator(self): + """A name holding ':' could never be matched by its own streamed tensors, + since the first ':' is what splits adapter from tensor key.""" + result = self.engine.register_lora_adapter( + lora_name="team:alice", + config_dict=self.lora_config_dict, ) + self.assertFalse(result.success, "':' in an adapter name must be rejected") + self.assertIn("must not contain ':'", result.error_message) + self.assertNotIn("team:alice", loaded_adapter_names(self.engine)) def test_lora_e2e_load_from_tensor_params(self): print("[Test]Testing LoRA load from tensor params...") - result = self.engine.load_lora_adapter_from_tensors( - lora_name="self_cognition_Alice", - tensors=self.lora_tensors, - config_dict=self.lora_config_dict, + result = load_lora_via_stream( + self.engine, + "self_cognition_Alice", + self.lora_tensors, + self.lora_config_dict, ) self.assertTrue( result.success, @@ -150,10 +176,11 @@ def test_lora_load_unload_load_from_tensor_params(self): print("[Test]Testing LoRA load, unload, load from tensor params...") # Load LoRA adapter from tensors - result = self.engine.load_lora_adapter_from_tensors( - lora_name="self_cognition_Alice_multiple", - tensors=self.lora_tensors, - config_dict=self.lora_config_dict, + result = load_lora_via_stream( + self.engine, + "self_cognition_Alice_multiple", + self.lora_tensors, + self.lora_config_dict, ) self.assertTrue( result.success, @@ -175,10 +202,11 @@ def test_lora_load_unload_load_from_tensor_params(self): lora_path=["self_cognition_Alice_multiple"], ) # Load LoRA adapter again - result_again = self.engine.load_lora_adapter_from_tensors( - lora_name="self_cognition_Alice_multiple", - tensors=self.lora_tensors, - config_dict=self.lora_config_dict, + result_again = load_lora_via_stream( + self.engine, + "self_cognition_Alice_multiple", + self.lora_tensors, + self.lora_config_dict, ) self.assertTrue( result_again.success, @@ -240,10 +268,8 @@ def test_lora_logp_diff_with_huggingface(self): "down_proj", ], ) as srt_runner: - result = srt_runner.engine.load_lora_adapter_from_tensors( - lora_name=lora_name, - tensors=self.lora_tensors, - config_dict=self.lora_config_dict, + result = load_lora_via_stream( + srt_runner.engine, lora_name, self.lora_tensors, self.lora_config_dict ) self.assertTrue( result.success, @@ -334,23 +360,27 @@ def test_lora_e2e_load_from_flattened_bucket(self): from sglang.srt.utils import MultiprocessingSerializer from sglang.srt.weight_sync.tensor_bucket import FlattenedTensorBucket - named_tensors = list(self.lora_tensors.items()) - bucket = FlattenedTensorBucket(named_tensors=[(n, t) for n, t in named_tensors]) + lora_name = "self_cognition_Alice_flattened" + named_tensors = [ + (f"{lora_name}:{key}", tensor) for key, tensor in self.lora_tensors.items() + ] + bucket = FlattenedTensorBucket(named_tensors=named_tensors) bucket_dict = { "flattened_tensor": bucket.get_flattened_tensor(), "metadata": bucket.get_metadata(), } serialized = MultiprocessingSerializer.serialize(bucket_dict, output_str=True) - # flattened_bucket callers pass one serialized copy per TP rank, same - # as Engine.update_weights_from_tensor. - result = self.engine.load_lora_adapter_from_tensors( - lora_name="self_cognition_Alice_flattened", - tensors=[serialized], - config_dict=self.lora_config_dict, - load_format="flattened_bucket", + result = self.engine.register_lora_adapter( + lora_name=lora_name, config_dict=self.lora_config_dict ) self.assertTrue(result.success, f"Failed: {result.error_message}") + self.engine.begin_weight_update(selector="all", sync_base=False) + self.engine.update_weights_from_tensor( + named_tensors=[serialized], load_format="flattened_bucket" + ) + success, message = self.engine.end_weight_update() + self.assertTrue(success, f"Failed: {message}") output = self.engine.generate( prompt=[TEST_PROMPT], diff --git a/test/registered/rl/test_weight_checker_e2e.py b/test/registered/rl/test_weight_checker_e2e.py index bf1de736422e..a508486b0eaf 100644 --- a/test/registered/rl/test_weight_checker_e2e.py +++ b/test/registered/rl/test_weight_checker_e2e.py @@ -160,7 +160,12 @@ def test_e_checksum_returns_ranks_with_hashes(self): self.assertIn("checksums", first) self.assertIn("parallelism_info", first) - info = first["parallelism_info"] + parallelism_info = first["parallelism_info"] + self.assertIsInstance(parallelism_info, list) + self.assertGreaterEqual(len(parallelism_info), 1) + + info = parallelism_info[0] + self.assertEqual(info["role"], "target") for key in ( "tp_rank", "tp_size", diff --git a/test/registered/sampling/test_sampling_mask.py b/test/registered/sampling/test_sampling_mask.py index 187e805dc912..d09631535aef 100644 --- a/test/registered/sampling/test_sampling_mask.py +++ b/test/registered/sampling/test_sampling_mask.py @@ -108,13 +108,20 @@ def _materialize(self, sampled, capture, requested_rows): output = LogitsProcessorOutput( next_token_logits=None, sampling_mask_output=self.sampler._build_sampling_mask_output( - sampled, capture + sampled, + capture, + support_capture_indices=torch.arange( + len(requested_rows), device=sampled.device + ), ), ) output.sampling_mask_output.map_device_tensors(lambda tensor: tensor.cpu()) SchedulerBatchResultProcessor.materialize_sampling_mask_output( [ - SimpleNamespace(return_sampling_mask=i in requested_rows) + SimpleNamespace( + return_sampling_mask=i in requested_rows, + sampling_logprobs_mode="support", + ) for i in range(len(sampled)) ], output, @@ -129,13 +136,21 @@ def test_min_p_capture_matches_filtered_support_and_logprob(self): sampled, capture = self._sample( probs, backend, top_k=3, top_p=1.0, min_p=0.6 ) - output = self.sampler._build_sampling_mask_output(sampled, capture) + output = self.sampler._build_sampling_mask_output( + sampled, + capture, + support_capture_indices=torch.tensor([0], device="cuda"), + ) self.assertEqual(output.statuses.tolist(), [SamplingMaskStatus.OK]) self.assertEqual(output.lengths.tolist(), [2]) - self.assertEqual(set(output.token_ids[0, :2].tolist()), {0, 1}) + support = output.token_ids[0, :2].tolist() + self.assertEqual(set(support), {0, 1}) + sampled_index = support.index(sampled.item()) expected = (0.4 if sampled.item() == 0 else 0.3) / 0.7 self.assertAlmostEqual( - output.selected_logprobs.item(), math.log(expected), places=6 + output.support_logprobs[0, sampled_index].item(), + math.log(expected), + places=6, ) def test_hard_exclusion_replay_in_mixed_batch(self): @@ -184,12 +199,15 @@ def test_hard_exclusion_replay_in_mixed_batch(self): simple_sampling_case=False, ) output = self._materialize(sampled, capture, requested_rows=[0, 1]) - support = output.next_token_sampling_mask_idx[0] + support = output.next_token_sampling_mask_idx[0].tolist() + sampling_logprobs = output.next_token_sampling_logprobs[0].tolist() self.assertEqual(set(support), {0, 1, 3}) self.assertIn(int(sampled[0]), support) expected = original[0, sampled[0]] - original[0, support].logsumexp(0) self.assertAlmostEqual( - output.next_token_sampling_logprobs[0], expected.item(), places=5 + sampling_logprobs[support.index(int(sampled[0]))], + expected.item(), + places=5, ) self.assertIn(2, output.next_token_sampling_mask_idx[1]) @@ -258,9 +276,13 @@ def test_flashinfer_capture_only_materializes_requested_rows(self): output = self._materialize(sampled, capture, requested_rows) self.assertIsNone(output.next_token_sampling_mask_idx[0]) - self.assertEqual(set(output.next_token_sampling_mask_idx[1]), {0, 1, 2}) + self.assertEqual( + set(output.next_token_sampling_mask_idx[1].tolist()), {0, 1, 2} + ) self.assertIsNone(output.next_token_sampling_mask_idx[2]) - self.assertEqual(set(output.next_token_sampling_mask_idx[3]), {0, 1, 2}) + self.assertEqual( + set(output.next_token_sampling_mask_idx[3].tolist()), {0, 1, 2} + ) self.assertIsNone(output.next_token_sampling_logprobs[0]) self.assertIsNotNone(output.next_token_sampling_logprobs[1]) self.assertIsNone(output.next_token_sampling_logprobs[2]) @@ -314,6 +336,7 @@ def _post_generate( top_logprobs_num=0, custom_logit_processor=None, stream=False, + sampling_logprobs_mode="support", ): payload = { "text": "The capital of France is", @@ -326,6 +349,8 @@ def _post_generate( "return_sampling_mask": return_sampling_mask, "stream": stream, } + if return_sampling_mask and sampling_logprobs_mode is not None: + payload["sampling_logprobs_mode"] = sampling_logprobs_mode if custom_logit_processor is not None: payload["custom_logit_processor"] = custom_logit_processor if return_logprob: @@ -337,13 +362,19 @@ def _post_generate( def _assert_sampling_masks(self, output_ids, meta_info): masks = meta_info["output_token_sampling_mask"] + sampling_logprobs = meta_info["output_token_sampling_logprobs"] self.assertEqual(len(masks), len(output_ids)) - self.assertEqual( - len(meta_info["output_token_sampling_logprobs"]), len(output_ids) - ) - for token_id, mask in zip(output_ids, masks): + self.assertEqual(len(sampling_logprobs), len(output_ids)) + for token_id, mask, logprobs in zip( + output_ids, masks, sampling_logprobs, strict=True + ): self.assertIn(token_id, mask) self.assertEqual(len(mask), len(set(mask))) + self.assertEqual(len(mask), len(logprobs)) + self.assertTrue(all(math.isfinite(logprob) for logprob in logprobs)) + self.assertAlmostEqual( + sum(math.exp(logprob) for logprob in logprobs), 1.0, delta=1e-5 + ) return masks def _generate_sampling_masks(self, sampling_params): @@ -393,6 +424,7 @@ def test_disallowed_tokens_with_replay(self): meta = output["meta_info"] token = output["output_ids"][0] mask = meta["output_token_sampling_mask"][0] + sampling_logprobs = meta["output_token_sampling_logprobs"][0] self.assertTrue(set(mask).isdisjoint(blocked)) self.assertIn(token, mask) probs = { @@ -400,7 +432,22 @@ def test_disallowed_tokens_with_replay(self): } expected = math.log(probs[token] / sum(probs[tid] for tid in mask)) self.assertAlmostEqual( - meta["output_token_sampling_logprobs"][0], expected, delta=1e-2 + sampling_logprobs[mask.index(token)], expected, delta=1e-2 + ) + + def test_sampling_logprobs_default_to_selected_token(self): + response = self._post_generate( + {"top_k": _TOP_K, "top_p": _TOP_P}, + sampling_logprobs_mode=None, + ) + self.assertEqual(response.status_code, 200, response.text) + meta = response.json()["meta_info"] + self.assertTrue(meta["output_token_sampling_logprobs"]) + self.assertTrue( + all( + isinstance(logprob, float) + for logprob in meta["output_token_sampling_logprobs"] + ) ) def test_rejected_processors_do_not_break_generation(self): @@ -434,7 +481,7 @@ def test_generate_returns_greedy_singleton_mask(self): self.assertTrue(all(len(mask) == 1 for mask in masks)) def test_sampling_mask_matches_topk_logprobs(self): - """Check the returned mask and its renormalized logprobs. + """Check the returned mask and its aligned behavior logprobs. We get a wide prefix of full-vocab logprobs via ``return_logprob`` so cutoff ties that extend beyond ``top_k`` are visible. With @@ -445,7 +492,7 @@ def test_sampling_mask_matches_topk_logprobs(self): 2. every mask token is in the returned top logprobs and at or above the top-k cutoff (ties at the cutoff survive, so the mask may exceed ``top_k``), - 3. sampling_logprob == log(p[sampled] / sum(p[t] for t in mask)). + 3. each sampling logprob equals log(p[token] / sum(p[t] for t in mask)). """ top_k, top_p = _TOP_K, _TOP_P response = self._post_generate( @@ -464,8 +511,8 @@ def test_sampling_mask_matches_topk_logprobs(self): self.assertEqual(len(top_logprobs), len(output_ids)) - for output_id, mask, mask_logprob, step_top_logprobs in zip( - output_ids, sampling_masks, sampling_logprobs, top_logprobs + for output_id, mask, mask_logprobs, step_top_logprobs in zip( + output_ids, sampling_masks, sampling_logprobs, top_logprobs, strict=True ): probs = { int(tid): math.exp(logprob) for logprob, tid, _ in step_top_logprobs @@ -480,8 +527,9 @@ def test_sampling_mask_matches_topk_logprobs(self): self.assertGreaterEqual(probs[token_id], top_k_cutoff * (1 - 1e-3)) support_mass = sum(probs[token_id] for token_id in mask_set) - expected_logprob = math.log(probs[output_id] / support_mass) - self.assertAlmostEqual(mask_logprob, expected_logprob, delta=1e-2) + for token_id, sampling_logprob in zip(mask, mask_logprobs, strict=True): + expected_logprob = math.log(probs[token_id] / support_mass) + self.assertAlmostEqual(sampling_logprob, expected_logprob, delta=1e-2) def test_chat_completions_returns_sampling_mask(self): response = requests.post( @@ -495,6 +543,7 @@ def test_chat_completions_returns_sampling_mask(self): "max_tokens": _MAX_NEW_TOKENS, "ignore_eos": True, "return_sampling_mask": True, + "sampling_logprobs_mode": "support", "return_meta_info": True, "return_token_ids": True, }, @@ -546,9 +595,78 @@ def test_selected_token_must_have_positive_captured_weight(self): selected_weight=None, ) selected = torch.tensor([2 if token_ids is None else 0]) - output = self.sampler._build_sampling_mask_output(selected, capture) + output = self.sampler._build_sampling_mask_output( + selected, capture, support_capture_indices=None + ) self.assertEqual(output.statuses.tolist(), [SamplingMaskStatus.INVALID]) + def test_selected_mode_does_not_build_support_logprobs(self): + capture = _SamplingMaskCapture( + batch_rows=torch.tensor([0]), + weights=torch.tensor([[0.6, 0.4, 0.0]]), + token_ids=None, + selected_weight=torch.tensor([0.6]), + ) + output = self.sampler._build_sampling_mask_output( + torch.tensor([0]), capture, support_capture_indices=None + ) + self.assertIsNone(output.support_logprobs) + self.assertAlmostEqual(output.selected_logprobs.item(), math.log(0.6)) + + def test_support_logprobs_align_with_packed_token_ids(self): + captures = ( + _SamplingMaskCapture( + batch_rows=torch.tensor([0]), + weights=torch.tensor([[0.1, 0.6, 0.0, 0.3]]), + token_ids=None, + selected_weight=torch.tensor([0.3]), + ), + _SamplingMaskCapture( + batch_rows=torch.tensor([0]), + weights=torch.tensor([[0.6, 0.3, 0.1, 0.0]]), + token_ids=torch.tensor([[1, 3, 0, 2]], dtype=torch.int32), + selected_weight=torch.tensor([0.3]), + ), + ) + for capture in captures: + with self.subTest(sorted_capture=capture.token_ids is not None): + output = self.sampler._build_sampling_mask_output( + torch.tensor([3]), + capture, + support_capture_indices=torch.tensor([0]), + ) + self.assertEqual(output.statuses.tolist(), [SamplingMaskStatus.OK]) + self.assertEqual(output.lengths.tolist(), [3]) + self.assertEqual(output.token_ids[0, :3].tolist(), [1, 3, 0]) + torch.testing.assert_close( + output.support_logprobs[0, :3].exp(), + torch.tensor([0.6, 0.3, 0.1]), + ) + + def test_support_logprobs_only_pack_support_mode_rows(self): + capture = _SamplingMaskCapture( + batch_rows=torch.tensor([0, 1]), + weights=torch.tensor( + [ + [0.7, 0.3, 0.0], + [0.1, 0.6, 0.3], + ] + ), + token_ids=None, + selected_weight=torch.tensor([0.7, 0.6]), + ) + output = self.sampler._build_sampling_mask_output( + torch.tensor([0, 1]), + capture, + support_capture_indices=torch.tensor([1]), + ) + self.assertEqual(tuple(output.selected_logprobs.shape), (2,)) + self.assertEqual(tuple(output.support_logprobs.shape), (1, 3)) + torch.testing.assert_close( + output.support_logprobs[0].exp(), + torch.tensor([0.6, 0.3, 0.1]), + ) + def test_synced_token_logprob_is_recomputed_from_capture(self): capture = _SamplingMaskCapture( batch_rows=torch.tensor([0]), @@ -556,9 +674,17 @@ def test_synced_token_logprob_is_recomputed_from_capture(self): token_ids=torch.tensor([[2, 1, 0]], dtype=torch.int32), selected_weight=None, ) - output = self.sampler._build_sampling_mask_output(torch.tensor([1]), capture) + output = self.sampler._build_sampling_mask_output( + torch.tensor([1]), + capture, + support_capture_indices=torch.tensor([0]), + ) self.assertEqual(output.statuses.tolist(), [SamplingMaskStatus.OK]) self.assertAlmostEqual(output.selected_logprobs.item(), math.log(0.25)) + support = output.token_ids[0, :2].tolist() + self.assertAlmostEqual( + output.support_logprobs[0, support.index(1)].item(), math.log(0.25) + ) def test_greedy_device_output_survives_async_copy(self): from sglang.srt.managers.utils import GenerationBatchResult @@ -567,7 +693,9 @@ def test_greedy_device_output_survives_async_copy(self): output = LogitsProcessorOutput( next_token_logits=None, sampling_mask_output=self.sampler._build_greedy_sampling_mask_output( - torch.tensor([0, 2], device="cuda"), tokens + torch.tensor([0, 2], device="cuda"), + tokens, + support_capture_indices=torch.tensor([0, 1], device="cuda"), ), ) result = GenerationBatchResult( @@ -578,13 +706,28 @@ def test_greedy_device_output_survives_async_copy(self): self.assertEqual(output.sampling_mask_output.token_ids.device.type, "cpu") SchedulerBatchResultProcessor.materialize_sampling_mask_output( [ - SimpleNamespace(return_sampling_mask=flag) + SimpleNamespace( + return_sampling_mask=flag, + sampling_logprobs_mode="support", + ) for flag in (True, False, True) ], output, ) - self.assertEqual(output.next_token_sampling_mask_idx, [[3], None, [5]]) - self.assertEqual(output.next_token_sampling_logprobs, [0.0, None, 0.0]) + self.assertEqual( + [ + None if row is None else row.tolist() + for row in output.next_token_sampling_mask_idx + ], + [[3], None, [5]], + ) + self.assertEqual( + [ + None if row is None else row.tolist() + for row in output.next_token_sampling_logprobs + ], + [[0.0], None, [0.0]], + ) def test_overflow_never_materializes_a_partial_mask(self): # Simulate a top-k cutoff tie: a nominal top_k below the cap can still @@ -597,7 +740,9 @@ def test_overflow_never_materializes_a_partial_mask(self): ) sampling_output = self.sampler._build_sampling_mask_output( - torch.tensor([0]), capture + torch.tensor([0]), + capture, + support_capture_indices=torch.tensor([0]), ) output = LogitsProcessorOutput( @@ -605,7 +750,13 @@ def test_overflow_never_materializes_a_partial_mask(self): sampling_mask_output=sampling_output, ) SchedulerBatchResultProcessor.materialize_sampling_mask_output( - [SimpleNamespace(return_sampling_mask=True)], output + [ + SimpleNamespace( + return_sampling_mask=True, + sampling_logprobs_mode="support", + ) + ], + output, ) self.assertEqual( output.next_token_sampling_mask_status, @@ -641,6 +792,61 @@ def test_return_sampling_mask_preserves_deterministic_sampling(self): outputs.append((output["output_ids"], output["text"])) self.assertEqual(outputs[0], outputs[1]) + def test_modes_and_streaming_return_identical_sampling_masks(self): + """For the same seeded tokens, selected and support mode report the same + supports, each selected logprob is the sampled token's support logprob, and a + stream that emits every token matches the response emitted every 50 tokens.""" + sampling_params = { + "top_k": _TOP_K, + "top_p": 1.0, + "sampling_seed": _SAMPLING_SEED, + "max_new_tokens": 64, + } + responses = {} + for mode, stream in ( + ("support", False), + ("selected", False), + ("support", True), + ): + response = self._post_generate( + sampling_params, sampling_logprobs_mode=mode, stream=stream + ) + self.assertEqual(response.status_code, 200, response.text) + if stream: + chunks = [ + json.loads(line[6:]) + for line in response.iter_lines() + if line.startswith(b"data: ") and line[6:] != b"[DONE]" + ] + responses[mode, stream] = chunks[-1] + else: + responses[mode, stream] = response.json() + + support = responses["support", False] + output_ids = support["output_ids"] + masks = self._assert_sampling_masks(output_ids, support["meta_info"]) + support_logprobs = support["meta_info"]["output_token_sampling_logprobs"] + + selected = responses["selected", False] + self.assertEqual(selected["output_ids"], output_ids) + self.assertEqual(selected["meta_info"]["output_token_sampling_mask"], masks) + self.assertEqual( + selected["meta_info"]["output_token_sampling_logprobs"], + [ + logprobs[mask.index(token)] + for token, mask, logprobs in zip(output_ids, masks, support_logprobs) + ], + ) + + streamed = responses["support", True] + self.assertEqual(streamed["output_ids"], output_ids) + for key in ( + "output_token_sampling_mask", + "output_token_sampling_logprobs", + "output_token_sampling_mask_length", + ): + self.assertEqual(streamed["meta_info"][key], support["meta_info"][key]) + class TestSamplingMaskPytorch(TestSamplingMask): _sampling_backend = "pytorch" @@ -687,12 +893,19 @@ def _check_parallel_config(self, *, tp_size, pp_size): self.assertEqual(meta["output_token_sampling_mask_length"], 4) self.assertEqual(len(masks), 4) self.assertEqual(len(logprobs), 4) - for token_id, mask, logprob in zip(token_ids, masks, logprobs): + for token_id, mask, step_logprobs in zip( + token_ids, masks, logprobs, strict=True + ): self.assertIn(token_id, mask) self.assertEqual(len(mask), len(set(mask))) self.assertLessEqual(len(mask), 64) - self.assertTrue(math.isfinite(logprob)) - self.assertLessEqual(logprob, 0.0) + self.assertEqual(len(mask), len(step_logprobs)) + self.assertTrue( + all(math.isfinite(logprob) for logprob in step_logprobs) + ) + self.assertTrue( + all(logprob <= 0.0 for logprob in step_logprobs) + ) if return_logprob: self.assertEqual(len(meta["output_token_logprobs"]), 4) @@ -705,21 +918,22 @@ def _check_parallel_config(self, *, tp_size, pp_size): process.wait(timeout=30) def _generate(self, *, return_sampling_mask, return_logprob): - response = requests.post( - DEFAULT_URL_FOR_TEST + "/generate", - json={ - "text": "The capital of France is", - "sampling_params": { - "temperature": 0.8, - "top_k": 8, - "top_p": 0.9, - "max_new_tokens": 4, - "ignore_eos": True, - }, - "return_sampling_mask": return_sampling_mask, - "return_logprob": return_logprob, + payload = { + "text": "The capital of France is", + "sampling_params": { + "temperature": 0.8, + "top_k": 8, + "top_p": 0.9, + "max_new_tokens": 4, + "ignore_eos": True, }, - timeout=120, + "return_sampling_mask": return_sampling_mask, + "return_logprob": return_logprob, + } + if return_sampling_mask: + payload["sampling_logprobs_mode"] = "support" + response = requests.post( + DEFAULT_URL_FOR_TEST + "/generate", json=payload, timeout=120 ) self.assertEqual(response.status_code, 200, response.text) return response.json() diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index 2c8709483c98..6d8d66f07828 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/test/registered/unit/disaggregation/test_disaggregation_wire.py @@ -23,6 +23,7 @@ unpack_int_lists, unpack_list_of_buffers, ) +from sglang.srt.disaggregation.decode import DecodeRequest, DecodeTransferQueue from sglang.srt.disaggregation.decode_schedule_batch_mixin import ( ScheduleBatchDisaggregationDecodeMixin, ) @@ -51,6 +52,7 @@ QSATokenToKVPool, ) from sglang.srt.runtime_context import get_context +from sglang.srt.sampling.sampling_mask import SamplingMaskRows from sglang.srt.speculative.eagle_disaggregation import ( build_eagle_disagg_draft_input, ) @@ -536,8 +538,16 @@ def _make_req( seed, metadata_buffer_index=0, sampling_mask=None, - sampling_logprob=None, + sampling_logprobs=None, + sampling_logprobs_mode="support", ): + sampling_mask_rows = None + if sampling_mask is not None: + sampling_mask_rows = SamplingMaskRows() + sampling_mask_rows.append( + np.array(sampling_mask, np.int32), + np.array(sampling_logprobs, np.float32), + ) return SimpleNamespace( metadata_buffer_index=metadata_buffer_index, output_ids=[101], @@ -548,12 +558,8 @@ def _make_req( multimodal_inputs=None, return_logprob=False, return_sampling_mask=sampling_mask is not None, - output_token_sampling_mask=( - None if sampling_mask is None else [sampling_mask] - ), - output_token_sampling_logprobs=( - None if sampling_logprob is None else [sampling_logprob] - ), + sampling_logprobs_mode=sampling_logprobs_mode, + sampling_mask_rows=sampling_mask_rows, hidden_states_tensor=torch.tensor([1.0, 2.0]), output_topk_p=torch.tensor([1.0]), output_topk_index=torch.tensor([7]), @@ -603,7 +609,7 @@ def test_sampling_mask_metadata_is_opt_in(self): self._make_req( None, sampling_mask=[7, 8, 9] if enabled else None, - sampling_logprob=-1.25 if enabled else None, + sampling_logprobs=[-1.25, -1.5, -2.0] if enabled else None, ) ) schemas.append(buffers.get_buf_infos()) @@ -611,10 +617,10 @@ def test_sampling_mask_metadata_is_opt_in(self): self.assertEqual( buffers.output_token_sampling_mask_idx.shape, (1, 3) ) - length, mask, logprob = buffers.get_buf(0)[6:9] + length, mask, logprobs = buffers.get_buf(0)[6:9] self.assertEqual(length[0].item(), 3) self.assertEqual(mask.tolist(), [7, 8, 9]) - self.assertAlmostEqual(logprob[0].item(), -1.25) + self.assertEqual(logprobs.tolist(), [-1.25, -1.5, -2.0]) else: self.assertIsNone(buffers.output_token_sampling_mask_len) self.assertIsNone(buffers.output_token_sampling_mask_idx) @@ -623,7 +629,151 @@ def test_sampling_mask_metadata_is_opt_in(self): disabled_ptrs, _, disabled_sizes = schemas[0] enabled_ptrs, _, enabled_sizes = schemas[1] self.assertEqual(len(enabled_ptrs) - len(disabled_ptrs), 3) - self.assertEqual(sum(enabled_sizes) - sum(disabled_sizes), 3 * 4 + 128) + self.assertEqual(sum(enabled_sizes) - sum(disabled_sizes), 2 * 3 * 4 + 64) + + def test_sampling_mask_selected_logprob_uses_first_metadata_slot(self): + with envs.SGLANG_ENABLE_DISAGG_SAMPLING_MASK.override(True): + buffers = MetadataBuffers( + size=1, + hidden_size=2, + hidden_states_dtype=torch.float32, + max_sampling_mask_tokens=3, + ) + buffers.set_buf( + self._make_req( + None, + sampling_mask=[7, 8, 9], + sampling_logprobs=[-0.5], + sampling_logprobs_mode="selected", + ) + ) + length, mask, logprobs = buffers.get_buf(0)[6:9] + self.assertEqual(length[0].item(), 3) + self.assertEqual(mask.tolist(), [7, 8, 9]) + self.assertEqual(logprobs[0].item(), -0.5) + + def test_sampling_mask_row_reaches_decode_unchanged(self): + """The prefill worker's first-token row is the row the decode worker streams.""" + for mode, logprobs, expected in ( + ("support", [-1.25, -1.5, -2.0], [[-1.25, -1.5, -2.0]]), + ("selected", [-0.5], [-0.5]), + ): + with ( + self.subTest(sampling_logprobs_mode=mode), + envs.SGLANG_ENABLE_DISAGG_SAMPLING_MASK.override(True), + ): + buffers = MetadataBuffers( + size=1, + hidden_size=2, + hidden_states_dtype=torch.float32, + max_sampling_mask_tokens=4, + ) + buffers.set_buf( + self._make_req( + None, + sampling_mask=[7, 8, 9], + sampling_logprobs=logprobs, + sampling_logprobs_mode=mode, + ) + ) + queue = DecodeTransferQueue.__new__(DecodeTransferQueue) + queue.scheduler = SimpleNamespace( + kv_checksum_computer=None, + batch_result_processor=SimpleNamespace( + _maybe_update_reasoning_tokens=lambda req, token_id: None + ), + ) + queue.spec_algorithm = SimpleNamespace(is_none=lambda: True) + queue.metadata_buffers = buffers + req = SimpleNamespace( + rid="r0", + bootstrap_host="127.0.0.1", + bootstrap_room=9, + output_ids=[], + return_logprob=False, + return_sampling_mask=True, + sampling_logprobs_mode=mode, + sampling_mask_rows=SamplingMaskRows(), + time_stats=SimpleNamespace(set_wait_queue_entry_time=lambda: None), + ) + + queue._commit_transfer_to_req( + DecodeRequest( + req=req, + kv_receiver=SimpleNamespace(clear=lambda: None), + metadata_buffer_index=0, + ) + ) + + self.assertEqual(req.output_ids, [101]) + self.assertEqual( + req.sampling_mask_rows.take().to_lists( + support_logprobs=mode == "support" + ), + ([[7, 8, 9]], expected), + ) + + def test_rebootstrap_replay_keeps_one_sampling_mask_row_per_token(self): + """A PD rebootstrap that replays an already-emitted token keeps that token's + sampling-mask row instead of adding the prefill worker's fresh one.""" + with envs.SGLANG_ENABLE_DISAGG_SAMPLING_MASK.override(True): + buffers = MetadataBuffers( + size=1, + hidden_size=2, + hidden_states_dtype=torch.float32, + max_sampling_mask_tokens=4, + ) + buffers.set_buf( + self._make_req( + None, + sampling_mask=[7, 8, 9], + sampling_logprobs=[-1.25, -1.5, -2.0], + ) + ) + queue = DecodeTransferQueue.__new__(DecodeTransferQueue) + queue.scheduler = SimpleNamespace( + kv_checksum_computer=None, + batch_result_processor=SimpleNamespace( + _maybe_update_reasoning_tokens=lambda req, token_id: None + ), + ) + queue.spec_algorithm = SimpleNamespace(is_none=lambda: True) + queue.metadata_buffers = buffers + # Retracting popped token 6 from output_ids; its row is still queued. + sampling_mask_rows = SamplingMaskRows() + sampling_mask_rows.append( + np.array([5, 1], np.int32), np.array([-0.5, -2.25], np.float32) + ) + sampling_mask_rows.append( + np.array([6, 2], np.int32), np.array([-0.25, -1.75], np.float32) + ) + req = SimpleNamespace( + rid="r0", + bootstrap_host="127.0.0.1", + bootstrap_room=9, + output_ids=[5], + pd_rebootstrap_forced_output_id=6, + return_logprob=False, + return_sampling_mask=True, + sampling_logprobs_mode="support", + sampling_mask_rows=sampling_mask_rows, + time_stats=SimpleNamespace(set_wait_queue_entry_time=lambda: None), + ) + + queue._commit_transfer_to_req( + DecodeRequest( + req=req, + kv_receiver=SimpleNamespace(clear=lambda: None), + metadata_buffer_index=0, + is_rebootstrap=True, + ) + ) + + self.assertEqual(req.output_ids, [5, 6]) + self.assertEqual( + req.sampling_mask_rows.take().to_lists(support_logprobs=True), + ([[5, 1], [6, 2]], [[-0.5, -2.25], [-0.25, -1.75]]), + ) def test_decode_input_requires_valid_seed_for_every_request(self): seeds = ( diff --git a/test/registered/unit/entrypoints/anthropic/test_utils.py b/test/registered/unit/entrypoints/anthropic/test_utils.py new file mode 100644 index 000000000000..593176d9460e --- /dev/null +++ b/test/registered/unit/entrypoints/anthropic/test_utils.py @@ -0,0 +1,793 @@ +"""Tests for the Anthropic conversion seams used by external frontends. + +Request and response goldens call the public runtime-independent functions in +serving.py directly. The utils-specific tests cover the composite error DTO, +eager fake-SSE synthesis, and message-ID generation. +""" + +import ast +import importlib.util +import json +import subprocess +import sys +import unittest +from pathlib import Path +from types import SimpleNamespace + +from sglang.test.test_utils import maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() # must precede imports that may pull in sgl_kernel + +from sglang.srt.entrypoints.anthropic import utils # noqa: E402 +from sglang.srt.entrypoints.anthropic.protocol import ( # noqa: E402 + AnthropicMessagesRequest, +) +from sglang.srt.entrypoints.anthropic.serving import ( # noqa: E402 + AnthropicServing, + convert_response, + convert_to_chat_completion_request, +) +from sglang.srt.entrypoints.openai.protocol import ChatCompletionResponse # noqa: E402 +from sglang.test.ci.ci_register import register_cpu_ci # noqa: E402 + +register_cpu_ci(est_time=30, suite="base-a-test-cpu") + + +class _FakeOpenAIServingChat: + """Just enough of OpenAIServingChat for ``AnthropicServing.__init__``.""" + + def __init__(self, chat_template=None): + self.tokenizer_manager = SimpleNamespace( + tokenizer=SimpleNamespace(chat_template=chat_template) + ) + + +def _real_serving(merge_inline_system: bool) -> AnthropicServing: + """A normally constructed AnthropicServing with a forced policy. + + The matrix must cover both policy values regardless of what the fake's + (absent) chat template probes to; template detection itself is covered + by ``test_serving.py``. + """ + serving = AnthropicServing(_FakeOpenAIServingChat()) + serving._merge_inline_system = merge_inline_system + return serving + + +def _payload(messages, **extra) -> dict: + return {"model": "claude-test", "max_tokens": 64, "messages": messages, **extra} + + +def _convert(payload: dict, merge_inline_system: bool = True) -> dict: + request = AnthropicMessagesRequest.model_validate(payload) + openai_request = convert_to_chat_completion_request( + request, merge_inline_system=merge_inline_system + ) + return openai_request.model_dump(mode="json", exclude_none=True, by_alias=True) + + +_TOOLS = [ + { + "name": "get_weather", + "description": "w", + "input_schema": {"type": "object", "properties": {}}, + } +] + +_REQUEST_CASES = { + "text_system_sampling": _payload( + [{"role": "user", "content": "hello"}], + system="be brief", + temperature=0.5, + top_k=20, + top_p=0.9, + stop_sequences=["END", "STOP"], + ), + "system_blocks_and_inline_system": _payload( + [ + {"role": "user", "content": "q"}, + {"role": "system", "content": [{"type": "text", "text": " inline "}]}, + ], + system=[{"type": "text", "text": "s1"}, {"type": "text", "text": "s2"}], + ), + "stream_with_options": _payload([{"role": "user", "content": "hi"}], stream=True), + "tool_roundtrip": _payload( + [ + {"role": "user", "content": "weather?"}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "checking"}, + { + "type": "tool_use", + "id": "toolu_1", + "name": "get_weather", + "input": {"city": "Paris"}, + }, + ], + }, + { + "role": "user", + "content": [ + {"type": "text", "text": "pre"}, + { + "type": "tool_result", + "tool_use_id": "toolu_1", + "content": "sunny", + "is_error": True, + }, + {"type": "text", "text": "post"}, + ], + }, + ], + tools=_TOOLS, + tool_choice={"type": "auto"}, + ), + "empty_placeholders": _payload( + [ + {"role": "assistant", "content": [{"type": "text", "text": ""}]}, + {"role": "assistant", "content": []}, + {"role": "user", "content": "q"}, + ] + ), + "images": _payload( + [ + { + "role": "user", + "content": [ + {"type": "text", "text": "look"}, + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/jpeg", + "data": "eA==", + }, + }, + { + "type": "image", + "source": {"type": "url", "url": "https://x/y.png"}, + }, + ], + } + ] + ), + "search_result_block": _payload( + [ + { + "role": "user", + "content": [ + { + "type": "search_result", + "title": "T", + "source": "https://s", + "content": [{"type": "text", "text": "C"}], + } + ], + } + ] + ), + "server_tools_skipped": _payload( + [{"role": "user", "content": "q"}], + tools=[{"type": "web_search_20250305", "name": "web_search"}, *_TOOLS], + ), +} + + +def _chat_response( + message: dict, finish_reason: str = "stop", usage: dict = None +) -> ChatCompletionResponse: + return ChatCompletionResponse.model_validate( + { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "served-alias", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", **message}, + "finish_reason": finish_reason, + } + ], + "usage": usage + or {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + ) + + +_RESPONSE_CASES = { + "reasoning_and_tools": ( + { + "content": "text", + "reasoning_content": "thought", + "tool_calls": [ + { + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": '{"a": 1}'}, + }, + { + "id": "c2", + "type": "function", + "function": {"name": "g", "arguments": "not json"}, + }, + ], + }, + "tool_calls", + None, + ), +} + + +class TestRequestGoldens(unittest.TestCase): + """Absolute value pins for the public request converter.""" + + def test_sampling_params_and_system(self): + dump = _convert(_REQUEST_CASES["text_system_sampling"]) + self.assertEqual( + (dump["temperature"], dump["top_k"], dump["top_p"], dump["stop"]), + (0.5, 20, 0.9, ["END", "STOP"]), + ) + self.assertEqual(dump["max_tokens"], 64) + self.assertIs(dump["stream"], False) + self.assertNotIn("stream_options", dump) + self.assertEqual( + dump["messages"], + [ + {"role": "system", "content": "be brief"}, + {"role": "user", "content": "hello"}, + ], + ) + + def test_stream_true_sets_stream_options(self): + dump = _convert(_REQUEST_CASES["stream_with_options"]) + self.assertIs(dump["stream"], True) + self.assertEqual( + dump["stream_options"], + {"include_usage": True, "continuous_usage_stats": True}, + ) + + def test_system_blocks_and_inline_system_merge(self): + payload = _REQUEST_CASES["system_blocks_and_inline_system"] + merged = _convert(payload, merge_inline_system=True) + self.assertEqual( + merged["messages"][0], {"role": "system", "content": "s1\ns2\ninline"} + ) + self.assertEqual([m["role"] for m in merged["messages"]], ["system", "user"]) + + unmerged = _convert(payload, merge_inline_system=False) + self.assertEqual( + unmerged["messages"][0], {"role": "system", "content": "s1\ns2"} + ) + self.assertEqual( + [m["role"] for m in unmerged["messages"]], ["system", "user", "system"] + ) + + def test_tool_use_roundtrip_and_wire_order(self): + dump = _convert(_REQUEST_CASES["tool_roundtrip"]) + self.assertEqual( + dump["messages"][1], + { + "role": "assistant", + "content": "checking", + "tool_calls": [ + { + "id": "toolu_1", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "Paris"}', + }, + } + ], + }, + ) + # Wire order preserved: user(pre) -> tool -> user(post); is_error is + # dropped from the OpenAI tool message. + self.assertEqual(dump["messages"][2], {"role": "user", "content": "pre"}) + self.assertEqual( + dump["messages"][3], + {"role": "tool", "tool_call_id": "toolu_1", "content": "sunny"}, + ) + self.assertEqual(dump["messages"][4], {"role": "user", "content": "post"}) + self.assertEqual(dump["tools"][0]["function"]["name"], "get_weather") + + def test_tool_result_variants(self): + dump = _convert( + _payload( + [ + { + "role": "user", + "content": [ + {"type": "tool_result", "id": "legacy_id", "content": None} + ], + }, + { + "role": "assistant", + "content": [ + { + "type": "tool_result", + "tool_use_id": "t2", + "content": "inner", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "t3", + "content": [ + {"type": "text", "text": "a"}, + {"type": "text", "text": "b"}, + ], + } + ], + }, + ] + ) + ) + # None content -> ""; legacy ``id`` used as tool_call_id fallback. + self.assertEqual( + dump["messages"][0], + {"role": "tool", "tool_call_id": "legacy_id", "content": ""}, + ) + # Assistant-role tool_result folds into text. + self.assertEqual( + dump["messages"][1], {"role": "assistant", "content": "Tool result: inner"} + ) + # Multi-text list keeps the parts list. + self.assertEqual( + dump["messages"][2]["content"], + [{"type": "text", "text": "a"}, {"type": "text", "text": "b"}], + ) + + def test_tool_choice_mappings(self): + msgs = [{"role": "user", "content": "q"}] + self.assertEqual( + _convert(_payload(msgs, tools=_TOOLS, tool_choice={"type": "any"}))[ + "tool_choice" + ], + "required", + ) + self.assertEqual( + _convert(_payload(msgs, tools=_TOOLS, tool_choice={"type": "none"}))[ + "tool_choice" + ], + "none", + ) + self.assertEqual( + _convert( + _payload( + msgs, + tools=_TOOLS, + tool_choice={"type": "tool", "name": "get_weather"}, + ) + )["tool_choice"], + {"type": "function", "function": {"name": "get_weather"}}, + ) + self.assertEqual(_convert(_payload(msgs, tools=_TOOLS))["tool_choice"], "auto") + with self.assertRaisesRegex(ValueError, "not in the forwarded tools list"): + _convert( + _payload( + msgs, tools=_TOOLS, tool_choice={"type": "tool", "name": "missing"} + ) + ) + with self.assertRaisesRegex(ValueError, "requires at least one custom"): + _convert(_payload(msgs, tool_choice={"type": "any"})) + + def test_empty_text_and_empty_assistant_placeholders(self): + dump = _convert(_REQUEST_CASES["empty_placeholders"]) + self.assertEqual(dump["messages"][0], {"role": "assistant", "content": ""}) + self.assertEqual(dump["messages"][1], {"role": "assistant", "content": ""}) + + +class TestResponseGoldens(unittest.TestCase): + """Absolute value pins for the public response converter.""" + + def test_default_id_factory_keeps_wire_format(self): + result = convert_response(_chat_response({"content": "hi"})) + self.assertRegex(result.id, r"^msg_[0-9a-f]{32}$") + + def test_empty_choices_response(self): + no_choices = ChatCompletionResponse.model_validate( + { + "id": "c", + "object": "chat.completion", + "created": 1, + "model": "m", + "choices": [], + "usage": { + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0, + }, + } + ) + result = convert_response(no_choices) + self.assertEqual(result.stop_reason, "end_turn") + self.assertEqual([b.type for b in result.content], ["text"]) + self.assertEqual(result.content[0].text, "") + self.assertEqual( + (result.usage.input_tokens, result.usage.output_tokens), (0, 0) + ) + + def test_block_order_and_invalid_json_tool_arguments(self): + message, finish_reason, _ = _RESPONSE_CASES["reasoning_and_tools"] + result = convert_response(_chat_response(message, finish_reason)) + types = [(b.type, getattr(b, "name", None)) for b in result.content] + self.assertEqual( + types, + [("thinking", None), ("text", None), ("tool_use", "f"), ("tool_use", "g")], + ) + self.assertEqual(result.content[2].input, {"a": 1}) + self.assertEqual(result.content[3].input, {}) + self.assertEqual(result.stop_reason, "tool_use") + + def test_stop_reason_mapping(self): + length = convert_response(_chat_response({"content": "x"}, "length")) + self.assertEqual(length.stop_reason, "max_tokens") + unmapped = convert_response(_chat_response({"content": "x"}, "content_filter")) + self.assertEqual(unmapped.stop_reason, "end_turn") + + +class TestExtendedRequestGoldens(unittest.TestCase): + def test_unknown_extra_keys_keep_pydantic_ignore_behavior(self): + payload = _payload( + [{"role": "user", "content": "hi", "unknown_msg_key": 1}], + unknown_top_key="x", + ) + request = AnthropicMessagesRequest.model_validate(payload) + self.assertEqual(request.model, "claude-test") + self.assertFalse(hasattr(request, "unknown_top_key")) + + def test_image_conversion(self): + dump = _convert(_REQUEST_CASES["images"]) + parts = dump["messages"][0]["content"] + self.assertEqual([p["type"] for p in parts], ["text", "image_url", "image_url"]) + self.assertEqual(parts[1]["image_url"]["url"], "data:image/jpeg;base64,eA==") + self.assertEqual(parts[2]["image_url"]["url"], "https://x/y.png") + + def test_output_config_effort_mapping(self): + msgs = [{"role": "user", "content": "q"}] + high = _convert(_payload(msgs, output_config={"effort": "high"})) + xhigh = _convert(_payload(msgs, output_config={"effort": "xhigh"})) + self.assertEqual(high["reasoning_effort"], "high") + self.assertEqual(xhigh["reasoning_effort"], "max") + + def test_server_tools_are_skipped_not_forwarded(self): + dump = _convert(_REQUEST_CASES["server_tools_skipped"]) + self.assertEqual( + [tool["function"]["name"] for tool in dump["tools"]], ["get_weather"] + ) + + def test_search_result_flattening(self): + dump = _convert(_REQUEST_CASES["search_result_block"]) + self.assertEqual( + dump["messages"][0], + {"role": "user", "content": "Title: T\nSource: https://s\nContent: C"}, + ) + + def test_nested_tool_reference_translation(self): + dump = _convert( + _payload( + [ + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "t", + "content": [ + { + "type": "tool_reference", + "tool_name": "deferred_fn", + } + ], + } + ], + } + ] + ) + ) + self.assertEqual( + dump["messages"][0]["content"], + [{"type": "tool_reference", "name": "deferred_fn"}], + ) + + +_ERROR_BODIES = [ + b'{"error": "boom"}', + b'{"error": {"message": "m", "type": "custom_type"}}', + b'{"message": "top-level"}', + b"gateway", + b"", +] + + +class TestErrorConversion(unittest.TestCase): + def test_matches_serving_for_shared_statuses(self): + # Statuses present in serving.py's own ERROR_TYPE_MAP; 413/422 are + # composite additions with no serving.py equivalent (below). + for status in (400, 401, 403, 404, 408, 429, 500, 502, 503, 504): + for body in _ERROR_BODIES: + with self.subTest(status=status, body=body): + serving_response = _real_serving( + True + )._convert_openai_error_response( + SimpleNamespace(status_code=status, body=body) + ) + self.assertEqual(serving_response.status_code, status) + envelope = utils.to_anthropic_error(status, body) + self.assertEqual( + envelope.model_dump(), json.loads(bytes(serving_response.body)) + ) + + def test_composite_statuses_follow_http_layer_policy(self): + self.assertEqual(utils.ERROR_TYPE_MAP[413], "request_too_large") + self.assertEqual(utils.ERROR_TYPE_MAP[422], "invalid_request_error") + env_413 = utils.to_anthropic_error(413, b'{"error": "big"}') + self.assertEqual( + (env_413.error.type, env_413.error.message), ("request_too_large", "big") + ) + env_422 = utils.to_anthropic_error(422, b'{"error": "bad"}') + self.assertEqual( + (env_422.error.type, env_422.error.message), + ("invalid_request_error", "bad"), + ) + # The composite statuses share the full message policy: custom + # error.type honored for 4xx, empty body -> "Request failed". + custom = utils.to_anthropic_error( + 413, b'{"error": {"message": "m", "type": "custom_type"}}' + ) + self.assertEqual( + (custom.error.type, custom.error.message), ("custom_type", "m") + ) + self.assertEqual( + utils.to_anthropic_error(422, b"").error.message, "Request failed" + ) + # Unlisted statuses fall through to api_error in both sources. + self.assertEqual( + utils.to_anthropic_error(409, b'{"error": "conflict"}').error.type, + "api_error", + ) + + def test_message_policy_goldens(self): + """Absolute pins for the envelope message policy (see + TestRequestGoldens for why value pins accompany the delegation + matrix).""" + # 4xx keeps the parsed upstream message; custom error.type honored. + env = utils.to_anthropic_error( + 400, b'{"error": {"message": "m", "type": "custom_type"}}' + ) + self.assertEqual((env.error.type, env.error.message), ("custom_type", "m")) + # 5xx never echoes upstream detail or type. + env = utils.to_anthropic_error( + 500, b'{"error": {"message": "secret", "type": "custom_type"}}' + ) + self.assertEqual( + (env.error.type, env.error.message), ("api_error", "Internal server error") + ) + # Empty-body fallbacks. + self.assertEqual( + utils.to_anthropic_error(400, b"").error.message, "Request failed" + ) + self.assertEqual( + utils.to_anthropic_error(502, b"").error.message, "Internal server error" + ) + # Non-JSON 4xx body passes through as a bounded hint. + self.assertEqual( + utils.to_anthropic_error(400, b"gateway").error.message, + "gateway", + ) + + def test_composite_map_matches_http_server_source(self): + """413/422 are a copy of the inline status->type dict in + ``http_server.py``'s ``/v1/messages`` exception handler, which is not + importable without loading the server app. Parse the source so a + policy change there fails loudly here instead of silently diverging + every external frontend.""" + origin = importlib.util.find_spec("sglang.srt.entrypoints.http_server").origin + tree = ast.parse(Path(origin).read_text()) + handler_maps = [] + for node in ast.walk(tree): + if not isinstance(node, ast.Dict) or not node.keys: + continue + if not all( + key is not None + and isinstance(key, ast.Constant) + and isinstance(key.value, int) + for key in node.keys + ): + continue + mapping = { + key.value: value.value + for key, value in zip(node.keys, node.values) + if isinstance(value, ast.Constant) + } + if 413 in mapping and 422 in mapping: + handler_maps.append(mapping) + self.assertTrue( + handler_maps, "no status->type dict with 413/422 found in http_server.py" + ) + for mapping in handler_maps: + for status, error_type in mapping.items(): + self.assertEqual( + utils.ERROR_TYPE_MAP.get(status, "api_error"), + error_type, + f"status {status} diverged from http_server.py policy", + ) + + def test_4xx_scrub_strips_traceback_lines_and_truncates(self): + upstream = "\n".join( + [ + "Traceback (most recent call last):", + ' File "/app/handler.py", line 3, in run', + "real cause: " + "x" * 600, + ] + ) + body = json.dumps({"error": {"message": upstream}}).encode() + message = utils.to_anthropic_error(400, body).error.message + self.assertNotIn("Traceback", message) + self.assertNotIn('File "/', message) + self.assertTrue(message.startswith("real cause: ")) + self.assertEqual(len(message), 501) + self.assertTrue(message.endswith("…")) + + +class TestFakeSse(unittest.TestCase): + def test_default_id_factory_keeps_wire_format(self): + events = utils.to_anthropic_fake_sse_events( + _chat_response({"content": "hello"}), model="claude-test" + ) + self.assertRegex(events[0].message.id, r"^msg_[0-9a-f]{32}$") + + def test_text_event_sequence_uses_request_model(self): + events = utils.to_anthropic_fake_sse_events( + _chat_response({"content": "hello"}), + model="claude-test", + id_factory=lambda: "msg_fixed", + ) + self.assertEqual( + [e.type for e in events], + [ + "message_start", + "content_block_start", + "content_block_delta", + "content_block_stop", + "message_delta", + "message_stop", + ], + ) + start = events[0].message + # Model comes from the Anthropic request, not the backend alias. + self.assertEqual(start.model, "claude-test") + self.assertEqual(start.id, "msg_fixed") + self.assertEqual(start.content, []) + self.assertEqual(start.usage.input_tokens, 10) + self.assertEqual(start.usage.output_tokens, 0) + self.assertEqual(events[2].delta.text, "hello") + self.assertEqual(events[4].delta.stop_reason, "end_turn") + self.assertIsNone(events[4].usage.input_tokens) + self.assertEqual(events[4].usage.output_tokens, 5) + + def test_multi_block_index_accounting(self): + events = utils.to_anthropic_fake_sse_events( + _chat_response( + { + "content": "txt", + "reasoning_content": "think", + "tool_calls": [ + { + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": '{"a": 1}'}, + }, + { + "id": "c2", + "type": "function", + "function": {"name": "g", "arguments": ""}, + }, + ], + }, + finish_reason="tool_calls", + ), + model="claude-test", + id_factory=lambda: "m", + ) + starts = [e for e in events if e.type == "content_block_start"] + self.assertEqual( + [(e.index, e.content_block.type) for e in starts], + [(0, "thinking"), (1, "text"), (2, "tool_use"), (3, "tool_use")], + ) + deltas = [e for e in events if e.type == "content_block_delta"] + # A zero-argument tool call emits no input_json_delta (live stream + # behavior). + self.assertEqual( + [(e.index, e.delta.type) for e in deltas], + [(0, "thinking_delta"), (1, "text_delta"), (2, "input_json_delta")], + ) + self.assertEqual(deltas[2].delta.partial_json, '{"a": 1}') + stops = [e.index for e in events if e.type == "content_block_stop"] + self.assertEqual(stops, [0, 1, 2, 3]) + self.assertEqual(events[-2].delta.stop_reason, "tool_use") + self.assertEqual(events[-1].type, "message_stop") + + def test_unmapped_finish_reason_and_cached_usage(self): + events = utils.to_anthropic_fake_sse_events( + _chat_response( + {"content": "x"}, + finish_reason="content_filter", + usage={ + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + "prompt_tokens_details": {"cached_tokens": 4}, + }, + ), + model="claude-test", + id_factory=lambda: "m", + ) + start_usage = events[0].message.usage + self.assertEqual( + ( + start_usage.input_tokens, + start_usage.cache_read_input_tokens, + start_usage.output_tokens, + ), + (6, 4, 0), + ) + message_delta = next(e for e in events if e.type == "message_delta") + # content_filter is unmapped and falls through to end_turn. + self.assertEqual(message_delta.delta.stop_reason, "end_turn") + self.assertIsNone(message_delta.usage.cache_read_input_tokens) + self.assertEqual(message_delta.usage.output_tokens, 5) + + +class TestFakeSseEdgeCases(unittest.TestCase): + def test_empty_choices_emits_bare_envelope(self): + no_choices = ChatCompletionResponse.model_validate( + { + "id": "c", + "object": "chat.completion", + "created": 1, + "model": "m", + "choices": [], + "usage": { + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0, + }, + } + ) + events = utils.to_anthropic_fake_sse_events( + no_choices, model="claude-test", id_factory=lambda: "m" + ) + self.assertEqual( + [e.type for e in events], ["message_start", "message_delta", "message_stop"] + ) + self.assertEqual(events[1].delta.stop_reason, "end_turn") + + +class TestImportHygiene(unittest.TestCase): + def test_utils_import_loads_no_serving_runtime(self): + # External frontends import this module without a serving runtime; + # it must not pull the OpenAI serving chat, tokenizer manager, or + # engine (serving.py keeps OpenAIServingChat under TYPE_CHECKING). + code = ( + "import sys\n" + "import sglang.srt.entrypoints.anthropic.utils\n" + "banned = ('sglang.srt.entrypoints.openai.serving_chat',\n" + " 'sglang.srt.managers.tokenizer_manager',\n" + " 'sglang.srt.entrypoints.engine')\n" + "loaded = [m for m in sys.modules for b in banned " + "if m == b or m.startswith(b + '.')]\n" + "assert not loaded, loaded\n" + ) + subprocess.run([sys.executable, "-c", code], check=True, timeout=300) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/test/registered/unit/entrypoints/openai/test_protocol.py b/test/registered/unit/entrypoints/openai/test_protocol.py index c580e691f409..fe312945c370 100644 --- a/test/registered/unit/entrypoints/openai/test_protocol.py +++ b/test/registered/unit/entrypoints/openai/test_protocol.py @@ -189,6 +189,7 @@ def test_basic_chat_completion_request(self): self.assertEqual(request.temperature, None) # default self.assertFalse(request.stream) # default self.assertFalse(request.return_sampling_mask) + self.assertIsNone(request.sampling_logprobs_mode) self.assertEqual(request.tool_choice, "none") # default when no tools def test_image_content_hash_validation(self): diff --git a/test/registered/unit/entrypoints/openai/test_serving_chat.py b/test/registered/unit/entrypoints/openai/test_serving_chat.py index 223314353486..2ba727ae0b5f 100644 --- a/test/registered/unit/entrypoints/openai/test_serving_chat.py +++ b/test/registered/unit/entrypoints/openai/test_serving_chat.py @@ -589,11 +589,13 @@ def test_convert_to_internal_request_single(self): ) self.basic_req.return_sampling_mask = True + self.basic_req.sampling_logprobs_mode = "support" self.basic_req.return_meta_info = True adapted, processed = self.chat._convert_to_internal_request(self.basic_req) self.assertIsInstance(adapted, GenerateReqInput) self.assertFalse(adapted.stream) self.assertTrue(adapted.return_sampling_mask) + self.assertEqual(adapted.sampling_logprobs_mode, "support") self.assertEqual(adapted.session_id, "session-1") self.assertEqual(processed, self.basic_req) diff --git a/test/registered/unit/entrypoints/test_grpc_bridge.py b/test/registered/unit/entrypoints/test_grpc_bridge.py index 85cba4d2f1cb..b48abe8b9f1c 100644 --- a/test/registered/unit/entrypoints/test_grpc_bridge.py +++ b/test/registered/unit/entrypoints/test_grpc_bridge.py @@ -4,6 +4,7 @@ from types import SimpleNamespace from sglang.srt.entrypoints.grpc_bridge import RuntimeHandle +from sglang.srt.managers.tokenizer_manager import ServerStatus, TokenizerManager from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -111,5 +112,51 @@ def test_streaming_first_finished_choice_is_not_batch_terminal(self): self.assertEqual([call[1] for call in callback.calls], [False, False, True]) +class TestEngineStateNotifications(CustomTestCase): + def setUp(self): + self.manager = TokenizerManager.__new__(TokenizerManager) + self.manager._engine_state_changed_callback = None + self.manager._server_status = ServerStatus.Starting + self.manager._gracefully_exit = False + self.manager._is_pause = False + self.notifications = 0 + self.manager.set_engine_state_changed_callback(self._notify) + + def _notify(self): + self.notifications += 1 + + def test_observable_state_notifies_only_on_changes(self): + self.manager.is_pause = False + self.manager.server_status = ServerStatus.Starting + self.manager.gracefully_exit = False + self.assertEqual(self.notifications, 0) + + self.manager.is_pause = True + self.manager.server_status = ServerStatus.Up + self.manager.gracefully_exit = True + self.assertEqual(self.notifications, 3) + + def test_runtime_handle_registers_callback_with_manager(self): + handle = RuntimeHandle.__new__(RuntimeHandle) + handle.tokenizer_manager = self.manager + callback = object() + + handle.set_engine_state_changed_callback(callback) + + self.assertIs(self.manager._engine_state_changed_callback, callback) + + def test_graceful_exit_notifies_and_changes_computed_health(self): + handle = RuntimeHandle.__new__(RuntimeHandle) + handle.tokenizer_manager = self.manager + self.manager.server_status = ServerStatus.Up + self.notifications = 0 + + self.assertTrue(handle.health_check()) + self.manager.gracefully_exit = True + + self.assertEqual(self.notifications, 1) + self.assertFalse(handle.health_check()) + + if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/test/registered/unit/entrypoints/test_ready.py b/test/registered/unit/entrypoints/test_ready.py new file mode 100644 index 000000000000..ebe960d6343a --- /dev/null +++ b/test/registered/unit/entrypoints/test_ready.py @@ -0,0 +1,56 @@ +import asyncio +import unittest +from types import SimpleNamespace + +from sglang.srt.entrypoints import http_server +from sglang.srt.managers.tokenizer_manager import ServerStatus, TokenizerManager +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +class TestReadyEndpoint(CustomTestCase): + def call_ready( + self, + *, + is_pause: bool = False, + gracefully_exit: bool = False, + server_status: ServerStatus = ServerStatus.Up, + ): + tokenizer_manager = TokenizerManager.__new__(TokenizerManager) + tokenizer_manager.is_pause = is_pause + tokenizer_manager.gracefully_exit = gracefully_exit + tokenizer_manager.server_status = server_status + + prior_state = http_server.get_global_state() + http_server.set_global_state( + SimpleNamespace(tokenizer_manager=tokenizer_manager) + ) + try: + return asyncio.run(http_server.ready()) + finally: + http_server._global_state = prior_state + + def test_ready_while_accepting_requests(self): + self.assertEqual(self.call_ready().status_code, 200) + + def test_not_ready_while_paused(self): + self.assertEqual(self.call_ready(is_pause=True).status_code, 503) + + def test_not_ready_while_starting(self): + self.assertEqual( + self.call_ready(server_status=ServerStatus.Starting).status_code, 503 + ) + + def test_not_ready_while_unhealthy(self): + self.assertEqual( + self.call_ready(server_status=ServerStatus.UnHealthy).status_code, 503 + ) + + def test_not_ready_while_gracefully_exiting(self): + self.assertEqual(self.call_ready(gracefully_exit=True).status_code, 503) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/test/registered/unit/lora/test_lora_lease.py b/test/registered/unit/lora/test_lora_lease.py new file mode 100644 index 000000000000..6bc487614184 --- /dev/null +++ b/test/registered/unit/lora/test_lora_lease.py @@ -0,0 +1,233 @@ +"""Unit tests for the exactly-once LoRA lease lifecycle. + +A leaked lease leaves the adapter's in-flight counter above zero forever, so +unload's wait_for_zero never returns; a double release drives the counter +negative, which hangs it just the same (it waits for exactly zero). Covers: + + * TokenizerManager._finalize_lora_lease is the single, idempotent release + owner for every terminal path + * pre-acquire failures (lora_id never set) release nothing + * LoRARegistry.acquire: lookup + counter increment are one atomic admission + step under the writer lock, so acquire racing an unload fails cleanly + instead of touching a deleted counter + * LoRARegistry.lru_lora_name never picks non-reloadable (wire-loaded) refs + * explicit unload drops the lora_ref_cache entry (DELETE) while the locked + helper alone (the max_loaded_loras EVICT path) keeps it +""" + +import asyncio +import unittest +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, Mock + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + +from sglang.srt.lora.lora_registry import LoRARef, LoRARegistry +from sglang.srt.managers.io_struct import UnloadLoRAAdapterReqInput +from sglang.srt.managers.tokenizer_manager import ReqState, TokenizerManager + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +def _make_tm(enable_lora: bool = True) -> TokenizerManager: + """TokenizerManager with only the fields the lease path reads.""" + tm = TokenizerManager.__new__(TokenizerManager) + tm.enable_lora = enable_lora + tm.lora_registry = MagicMock() + tm.lora_registry.release = AsyncMock() + return tm + + +def _make_state(lora_path="adapter-a", lora_id="id-1") -> ReqState: + obj = SimpleNamespace(lora_path=lora_path, lora_id=lora_id, rid="rid-1") + return ReqState([], False, asyncio.Event(), obj, Mock()) + + +class TestFinalizeLoraLease(CustomTestCase): + def _finalize(self, tm, state, times=1): + async def run(): + for _ in range(times): + tm._finalize_lora_lease(state) + # Let the create_task'd release coroutines run. + await asyncio.sleep(0) + + asyncio.run(run()) + + def test_release_exactly_once_on_double_finalize(self): + tm = _make_tm() + state = _make_state() + self._finalize(tm, state, times=3) + tm.lora_registry.release.assert_awaited_once_with("id-1") + self.assertTrue(state.lora_lease_released) + + def test_pre_acquire_failure_releases_nothing(self): + # _init_req_state runs before the LoRA acquire; a failure in between + # leaves lora_id unset and there is no lease to release. + tm = _make_tm() + state = _make_state(lora_id=None) + self._finalize(tm, state) + tm.lora_registry.release.assert_not_awaited() + self.assertFalse(state.lora_lease_released) + + def test_no_lora_path_is_noop(self): + tm = _make_tm() + state = _make_state(lora_path=None) + self._finalize(tm, state) + tm.lora_registry.release.assert_not_awaited() + + def test_lora_disabled_is_noop(self): + tm = _make_tm(enable_lora=False) + state = _make_state() + self._finalize(tm, state) + tm.lora_registry.release.assert_not_awaited() + + def test_none_state_is_noop(self): + tm = _make_tm() + + async def run(): + tm._finalize_lora_lease(None) + + asyncio.run(run()) + tm.lora_registry.release.assert_not_awaited() + + +class TestAcquireAtomicity(CustomTestCase): + def test_acquire_after_unregister_raises_value_error(self): + async def run(): + registry = LoRARegistry() + ref = LoRARef(lora_name="a", lora_path="/x") + await registry.register(ref) + await registry.unregister("a") + with self.assertRaises(ValueError): + await registry.acquire("a") + + asyncio.run(run()) + + def test_unload_waits_for_inflight_lease_then_completes(self): + async def run(): + registry = LoRARegistry() + ref = LoRARef(lora_name="a", lora_path="/x") + await registry.register(ref) + lora_id = await registry.acquire("a") + await registry.unregister("a") + + wait_task = asyncio.create_task(registry.wait_for_unload(lora_id)) + await asyncio.sleep(0.01) + self.assertFalse(wait_task.done(), "unload must wait for the lease") + + await registry.release(lora_id) + await asyncio.wait_for(wait_task, timeout=2) + + asyncio.run(run()) + + def test_batch_acquire_release_roundtrip(self): + async def run(): + registry = LoRARegistry() + for name in ("a", "b"): + await registry.register(LoRARef(lora_name=name, lora_path=f"/{name}")) + lora_ids = await registry.acquire(["a", "b", None]) + self.assertEqual(len(lora_ids), 3) + self.assertIsNone(lora_ids[2]) + await registry.release([i for i in lora_ids if i is not None]) + for name in ("a", "b"): + await registry.unregister(name) + + asyncio.run(run()) + + +class TestLruSkipsNonReloadable(CustomTestCase): + def test_lru_prefers_reloadable_even_if_older_entry_is_wire_loaded(self): + async def run(): + registry = LoRARegistry() + # Registration order = LRU order; the wire-loaded ref is oldest. + await registry.register( + LoRARef(lora_name="wire", lora_path="__stream__", reloadable=False) + ) + await registry.register(LoRARef(lora_name="disk", lora_path="/d")) + self.assertEqual(await registry.lru_lora_name(), "disk") + + asyncio.run(run()) + + def test_all_non_reloadable_yields_no_victim(self): + async def run(): + registry = LoRARegistry() + await registry.register( + LoRARef(lora_name="w1", lora_path="__stream__", reloadable=False) + ) + await registry.register( + LoRARef(lora_name="w2", lora_path="__stream__", reloadable=False) + ) + self.assertIsNone(await registry.lru_lora_name()) + self.assertIsNone(await registry.lru_lora_name(exclude_pinned=True)) + + asyncio.run(run()) + + def test_exclude_pinned_still_respected(self): + async def run(): + registry = LoRARegistry() + await registry.register( + LoRARef(lora_name="pinned-disk", lora_path="/p", pinned=True) + ) + await registry.register(LoRARef(lora_name="plain-disk", lora_path="/q")) + self.assertEqual( + await registry.lru_lora_name(exclude_pinned=True), "plain-disk" + ) + self.assertEqual(await registry.lru_lora_name(), "pinned-disk") + + asyncio.run(run()) + + +class TestUnloadRefCacheCleanup(CustomTestCase): + def _make_unload_tm(self, success: bool) -> TokenizerManager: + tm = TokenizerManager.__new__(TokenizerManager) + tm.auto_create_handle_loop = Mock() + tm.server_args = MagicMock() + tm.server_args.enable_lora = True + tm.server_args.dp_size = 1 + tm.lora_update_lock = asyncio.Lock() + tm._unload_lora_adapter_locked = AsyncMock( + return_value=SimpleNamespace(success=success) + ) + tm.lora_ref_cache = {"a": LoRARef(lora_name="a", lora_path="/x")} + return tm + + def test_explicit_unload_drops_ref_cache_entry(self): + tm = self._make_unload_tm(success=True) + asyncio.run(tm.unload_lora_adapter(UnloadLoRAAdapterReqInput(lora_name="a"))) + self.assertNotIn("a", tm.lora_ref_cache) + + def test_failed_unload_keeps_ref_cache_entry(self): + tm = self._make_unload_tm(success=False) + asyncio.run(tm.unload_lora_adapter(UnloadLoRAAdapterReqInput(lora_name="a"))) + self.assertIn("a", tm.lora_ref_cache) + + def test_locked_helper_alone_keeps_ref_cache_entry(self): + # The max_loaded_loras LRU loop calls _unload_lora_adapter_locked + # directly — an EVICT — and a disk-backed adapter must keep its + # reload-catalog entry or it can never be implicitly reloaded. + tm = TokenizerManager.__new__(TokenizerManager) + tm.lora_update_lock = asyncio.Lock() + tm.lora_registry = MagicMock() + tm.lora_registry.unregister = AsyncMock(return_value="id-a") + tm.lora_registry.wait_for_unload = AsyncMock() + tm.update_lora_adapter_communicator = AsyncMock( + return_value=[SimpleNamespace(success=True)] + ) + tm.lora_ref_cache = {"a": LoRARef(lora_name="a", lora_path="/x")} + + async def run(): + async with tm.lora_update_lock: + return await tm._unload_lora_adapter_locked( + UnloadLoRAAdapterReqInput(lora_name="a") + ) + + result = asyncio.run(run()) + self.assertTrue(result.success) + self.assertIn("a", tm.lora_ref_cache) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/lora/test_lora_staged_publication.py b/test/registered/unit/lora/test_lora_staged_publication.py new file mode 100644 index 000000000000..e23437bf0462 --- /dev/null +++ b/test/registered/unit/lora/test_lora_staged_publication.py @@ -0,0 +1,313 @@ +"""Staged LoRA publication: fresh versions stream and commit alongside +generation, while base and fixed-name sessions keep the writer lock.""" + +import asyncio +import unittest +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import torch + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + +from sglang.srt.lora.lora_registry import LoRARef, LoRARegistry +from sglang.srt.managers.io_struct import ( + BeginWeightUpdateReqInput, + EndWeightUpdateReqInput, + GenerateReqInput, + LoRAUpdateOutput, + RegisterLoRAAdapterReqInput, + UnloadLoRAAdapterReqInput, + UpdateWeightsFromDistributedReqInput, +) +from sglang.srt.managers.scheduler_components.weight_updater import ( + SchedulerWeightUpdaterManager, +) +from sglang.srt.managers.tokenizer_manager import TokenizerManager +from sglang.srt.runtime_context import get_context + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + + +def _reply(success=True): + return SimpleNamespace( + success=success, message="ok" if success else "worker failed" + ) + + +def _manager(): + tm = TokenizerManager.__new__(TokenizerManager) + tm.server_args = SimpleNamespace(max_loaded_loras=None) + tm.elastic_worker_count = 1 + tm.init_weight_update() + tm.auto_create_handle_loop = Mock() + tm._validate_lora_upsert_supported = Mock() + tm.record_config_updates = Mock() + tm.abort_request = Mock() + tm.mm_processor = None + tm.lora_update_lock = asyncio.Lock() + tm.lora_registry = LoRARegistry( + [LoRARef(lora_name=name, reloadable=False) for name in ("A@1", "B", "C")] + ) + tm.lora_ref_cache = {} + tm._pending_lora_publications = {} + tm.update_lora_adapter_communicator = AsyncMock( + side_effect=lambda obj: [LoRAUpdateOutput(success=True)] + ) + tm.begin_weight_update_communicator = AsyncMock(return_value=[_reply()]) + tm.end_weight_update_communicator = AsyncMock(return_value=[_reply()]) + tm.update_weights_from_distributed_communicator = AsyncMock(return_value=[_reply()]) + return tm + + +def _begin(**kwargs): + return BeginWeightUpdateReqInput(**(dict(sync_base=False) | kwargs)) + + +def _bucket(**kwargs): + values = dict( + names=["A@2:model.layers.0.self_attn.q_proj.lora_A.weight"], + dtypes=["float32"], + shapes=[[1]], + flush_cache=False, + ) + return UpdateWeightsFromDistributedReqInput(**(values | kwargs)) + + +def _end(**kwargs): + values = dict(expected_lora_checksums={"A@2": {"q_proj": "hash"}}) + return EndWeightUpdateReqInput(**(values | kwargs)) + + +class TestStagedPublication(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self): + config = get_context().override_server_args( + enable_lora=True, + pp_size=1, + dp_size=1, + enable_dp_attention=False, + checkpoint_engine_wait_weights_before_ready=False, + ) + config.install() + self.addCleanup(config.restore) + self.tm = _manager() + result = await self.tm.register_lora_adapter( + RegisterLoRAAdapterReqInput( + lora_name="A@2", config_dict={"r": 8}, defer_publish=True + ) + ) + self.assertTrue(result.success) + self.assertTrue(result.pending) + self.assertIn("A@2", self.tm._pending_lora_publications) + + async def test_deferred_name_is_unservable_until_commit(self): + tm = self.tm + old_ids = await tm.lora_registry.acquire(["A@1", "B", "C"]) + self.assertIsNone(await tm.lora_registry.get_lora_id("A@2")) + self.assertTrue((await tm.begin_weight_update(_begin()))[0]) + self.assertTrue((await tm.update_weights_from_distributed(_bucket()))[0]) + self.assertIsNone(await tm.lora_registry.get_lora_id("A@2")) + self.assertTrue((await tm.end_weight_update(_end()))[0]) + new_id = await tm.lora_registry.acquire("A@2") + self.assertNotIn(new_id, old_ids) + # Untouched adapters keep their identities through the publication. + self.assertEqual(await tm.lora_registry.get_lora_id("A@1"), old_ids[0]) + await tm.lora_registry.release(new_id) + await tm.lora_registry.release(old_ids) + + async def test_staged_session_completes_under_live_readers(self): + tm = self.tm + async with tm.model_update_lock.reader_lock: + self.assertTrue( + (await asyncio.wait_for(tm.begin_weight_update(_begin()), 1))[0] + ) + self.assertTrue(tm._weight_update_staged_session) + self.assertTrue( + ( + await asyncio.wait_for( + tm.update_weights_from_distributed(_bucket()), 1 + ) + )[0] + ) + self.assertTrue( + (await asyncio.wait_for(tm.end_weight_update(_end()), 1))[0] + ) + self.assertIsNotNone(await tm.lora_registry.get_lora_id("A@2")) + + async def test_base_and_fixed_name_sessions_still_wait_for_readers(self): + for sync_base in (True, False): + tm = _manager() # no pending publications: not a staged session + async with tm.model_update_lock.reader_lock: + task = asyncio.create_task( + tm.begin_weight_update( + BeginWeightUpdateReqInput(sync_base=sync_base) + ) + ) + await asyncio.sleep(0) + self.assertFalse(task.done()) + tm.begin_weight_update_communicator.assert_not_awaited() + self.assertTrue((await asyncio.wait_for(task, 1))[0]) + + async def test_base_sync_with_pending_registration_keeps_writer_lock(self): + tm = self.tm + async with tm.model_update_lock.reader_lock: + task = asyncio.create_task( + tm.begin_weight_update(BeginWeightUpdateReqInput(sync_base=True)) + ) + await asyncio.sleep(0) + self.assertFalse(task.done()) + self.assertTrue((await asyncio.wait_for(task, 1))[0]) + self.assertFalse(tm._weight_update_staged_session) + + async def test_defer_publish_requires_a_fresh_name(self): + for taken in ("A@1", "A@2"): # published name; pending name + result = await self.tm.register_lora_adapter( + RegisterLoRAAdapterReqInput( + lora_name=taken, config_dict={}, defer_publish=True + ) + ) + self.assertFalse(result.success) + + async def test_pending_name_cannot_be_reregistered_or_unloaded(self): + result = await self.tm.register_lora_adapter( + RegisterLoRAAdapterReqInput(lora_name="A@2", config_dict={}) + ) + self.assertFalse(result.success) + result = await self.tm.unload_lora_adapter( + UnloadLoRAAdapterReqInput(lora_name="A@2") + ) + self.assertFalse(result.success) + + async def test_staged_bucket_cannot_change_serving_state(self): + tm = self.tm + await tm.begin_weight_update(_begin()) + for kwargs in ( + {"flush_cache": True}, + {"abort_all_requests": True}, + {"weight_version": "2"}, + ): + self.assertFalse( + (await tm.update_weights_from_distributed(_bucket(**kwargs)))[0] + ) + tm.abort_request.assert_not_called() + tm.update_weights_from_distributed_communicator.assert_not_awaited() + + async def test_any_worker_failure_discards_the_publication(self): + tm = self.tm + await tm.begin_weight_update(_begin()) + tm.end_weight_update_communicator.return_value = [_reply(), _reply(False)] + tm.update_lora_adapter_communicator.reset_mock() + self.assertFalse((await tm.end_weight_update(_end()))[0]) + self.assertIsNone(await tm.lora_registry.get_lora_id("A@2")) + self.assertEqual(tm._pending_lora_publications, {}) + # The backends held a zeroed identity for the name; it must be dropped. + discarded = tm.update_lora_adapter_communicator.await_args.args[0] + self.assertIsInstance(discarded, UnloadLoRAAdapterReqInput) + self.assertEqual(discarded.lora_name, "A@2") + + async def test_missing_manifest_fails_closed(self): + tm = self.tm + await tm.begin_weight_update(_begin()) + success, message = await tm.end_weight_update( + _end(expected_lora_checksums=None) + ) + self.assertFalse(success) + self.assertIn("A@2", message) + sent = tm.end_weight_update_communicator.await_args.args[0] + self.assertTrue(sent.abort) + self.assertIsNone(await tm.lora_registry.get_lora_id("A@2")) + + async def test_abort_never_publishes(self): + tm = self.tm + await tm.begin_weight_update(_begin()) + self.assertTrue((await tm.end_weight_update(_end(abort=True)))[0]) + self.assertIsNone(await tm.lora_registry.get_lora_id("A@2")) + self.assertEqual(tm._pending_lora_publications, {}) + + async def test_version_is_recorded_only_at_successful_end(self): + tm = _manager() + self.assertTrue( + (await tm.begin_weight_update(BeginWeightUpdateReqInput(sync_base=True)))[0] + ) + self.assertTrue( + (await tm.update_weights_from_distributed(_bucket(weight_version="2")))[0] + ) + tm.record_config_updates.assert_not_called() + self.assertTrue((await tm.end_weight_update(EndWeightUpdateReqInput()))[0]) + tm.record_config_updates.assert_called_once_with( + "tokenizer.weight_version", weight_version="2" + ) + + +class TestRequestCarriedBackfill(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self): + config = get_context().override_server_args( + enable_lora=True, + pp_size=1, + dp_size=1, + enable_dp_attention=False, + checkpoint_engine_wait_weights_before_ready=False, + ) + config.install() + self.addCleanup(config.restore) + self.tm = _manager() + + async def load(obj): + await self.tm.lora_registry.register( + LoRARef(lora_name=obj.lora_name, lora_path=obj.lora_path) + ) + return SimpleNamespace(success=True, error_message="") + + self.tm.load_lora_adapter = AsyncMock(side_effect=load) + + async def test_backfill_path_admits_an_unknown_version(self): + obj = GenerateReqInput( + text="x", + lora_path="A@7", + lora_backfill_paths={"A@7": "/ckpt/run/sampler_weights/7"}, + ) + await self.tm._resolve_lora_path(obj) + self.assertIsNotNone(obj.lora_id) + loaded = self.tm.load_lora_adapter.await_args.args[0] + self.assertEqual( + (loaded.lora_name, loaded.lora_path), + ("A@7", "/ckpt/run/sampler_weights/7"), + ) + await self.tm.lora_registry.release(obj.lora_id) + + async def test_unknown_version_without_a_path_is_rejected(self): + with self.assertRaisesRegex(ValueError, "never been loaded"): + await self.tm._resolve_lora_path( + GenerateReqInput(text="x", lora_path="A@7") + ) + + async def test_pending_name_is_not_servable_or_backfillable(self): + self.tm._pending_lora_publications["A@7"] = LoRARef( + lora_name="A@7", lora_path="__stream__", reloadable=False + ) + with self.assertRaisesRegex(ValueError, "awaiting publication"): + await self.tm._resolve_lora_path( + GenerateReqInput( + text="x", lora_path="A@7", lora_backfill_paths={"A@7": "/ckpt/7"} + ) + ) + + +class TestStashCopiesIPCBuckets(unittest.TestCase): + def test_stash_survives_sender_reusing_the_bucket(self): + manager = SchedulerWeightUpdaterManager.__new__(SchedulerWeightUpdaterManager) + manager.tp_worker = SimpleNamespace(ps=SimpleNamespace(tp_rank=0)) + manager._lora_stash = {} + bucket = torch.ones(4) + manager._stash_lora_tensors([("A@2:q_proj", bucket)], copy_tensors=True) + bucket.zero_() + self.assertTrue( + torch.equal(manager._lora_stash["A@2"]["q_proj"], torch.ones(4)) + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/lora/test_lora_stream_update.py b/test/registered/unit/lora/test_lora_stream_update.py new file mode 100644 index 000000000000..bc94b6155451 --- /dev/null +++ b/test/registered/unit/lora/test_lora_stream_update.py @@ -0,0 +1,299 @@ +"""Streamed LoRA updates: prefix routing, the session stash, and the +whole-adapter apply at end_weight_update.""" + +import unittest +from unittest.mock import MagicMock, Mock + +import torch + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + +from sglang.srt.lora.lora_manager import LoRAManager +from sglang.srt.lora.lora_registry import LoRARef +from sglang.srt.managers.io_struct import LoRAUpdateOutput +from sglang.srt.managers.scheduler_components.weight_updater import ( + SchedulerWeightUpdaterManager, + _sha256_tensor, + _split_lora_named_tensors, +) +from sglang.srt.model_executor.model_runner_components.weight_updater import ( + LocalSerializedTensor, +) +from sglang.srt.utils import MultiprocessingSerializer + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + + +BASE_NAME = "model.layers.0.self_attn.q_proj.weight" +LORA_A = "model.layers.0.self_attn.q_proj.lora_A.weight" +LORA_B = "model.layers.0.self_attn.q_proj.lora_B.weight" + + +def _make_manager(**overrides) -> SchedulerWeightUpdaterManager: + kwargs = dict( + tp_worker=MagicMock(), + draft_worker=None, + tp_cpu_group=MagicMock(), + memory_saver_adapter=MagicMock(), + flush_cache=Mock(return_value=True), + is_fully_idle=Mock(return_value=True), + ) + kwargs.update(overrides) + return SchedulerWeightUpdaterManager(**kwargs) + + +class TestSplitLoraNamedTensors(CustomTestCase): + def test_partitions_on_prefix(self): + t = torch.zeros(1) + base, lora = _split_lora_named_tensors( + [ + (BASE_NAME, t), + (f"miles_lora:{LORA_A}", t), + (f"__miles_slot_3:{LORA_B}", t), + ] + ) + self.assertEqual([n for n, _ in base], [BASE_NAME]) + self.assertEqual( + [n for n, _ in lora], + [f"miles_lora:{LORA_A}", f"__miles_slot_3:{LORA_B}"], + ) + + def test_unprefixed_lora_tensor_is_a_sender_bug(self): + with self.assertRaisesRegex(AssertionError, "without a"): + _split_lora_named_tensors([(LORA_A, torch.zeros(1))]) + + def test_pure_base_payload_passes_through(self): + base, lora = _split_lora_named_tensors([(BASE_NAME, torch.zeros(1))]) + self.assertEqual(len(base), 1) + self.assertEqual(lora, []) + + +class TestLoraStash(CustomTestCase): + def test_stash_groups_by_adapter(self): + mgr = _make_manager() + a, b = torch.zeros(2), torch.ones(2) + mgr._stash_lora_tensors( + [(f"miles_lora:{LORA_A}", a), (f"__miles_slot_1:{LORA_A}", b)] + ) + self.assertEqual(set(mgr._lora_stash), {"miles_lora", "__miles_slot_1"}) + self.assertIs(mgr._lora_stash["miles_lora"][LORA_A], a) + self.assertIs(mgr._lora_stash["__miles_slot_1"][LORA_A], b) + + def test_split_at_first_colon_only(self): + mgr = _make_manager() + mgr._stash_lora_tensors([("name:with:colon.lora_A.weight", torch.zeros(1))]) + self.assertIn("with:colon.lora_A.weight", mgr._lora_stash["name"]) + + def test_per_rank_payload_is_unwrapped(self): + """A LoRA tensor never reaches the base loader, so the stash is the last + place a per-rank payload can be reduced to this rank's tensor.""" + tp_worker = MagicMock() + tp_worker.ps.tp_rank = 1 + mgr = _make_manager(tp_worker=tp_worker) + rank_tensors = [torch.zeros(2), torch.ones(2)] + payload = LocalSerializedTensor( + values=[MultiprocessingSerializer.serialize(t) for t in rank_tensors] + ) + mgr._stash_lora_tensors([(f"miles_lora:{LORA_A}", payload)]) + self.assertTrue( + torch.equal(mgr._lora_stash["miles_lora"][LORA_A], rank_tensors[1]) + ) + + +class TestApplyLoraStash(CustomTestCase): + def _manager_with_lora(self, apply_result=None): + mgr = _make_manager() + lora_manager = MagicMock() + lora_manager.apply_streamed_adapter.return_value = ( + apply_result if apply_result is not None else LoRAUpdateOutput(success=True) + ) + mgr.tp_worker.model_runner.lora_manager = lora_manager + return mgr, lora_manager + + def test_empty_stash_is_noop(self): + mgr = _make_manager() + mgr.tp_worker.model_runner.lora_manager = None + success, _ = mgr._apply_lora_stash(None) + self.assertTrue(success) + + def test_stash_without_lora_manager_fails(self): + mgr = _make_manager() + mgr.tp_worker.model_runner.lora_manager = None + mgr._stash_lora_tensors([(f"miles_lora:{LORA_A}", torch.zeros(1))]) + success, message = mgr._apply_lora_stash(None) + self.assertFalse(success) + self.assertIn("--enable-lora", message) + + def test_applies_each_adapter_and_clears_stash(self): + mgr, lora_manager = self._manager_with_lora() + mgr._stash_lora_tensors( + [(f"a:{LORA_A}", torch.zeros(1)), (f"b:{LORA_A}", torch.zeros(1))] + ) + success, _ = mgr._apply_lora_stash(None) + self.assertTrue(success) + self.assertEqual(lora_manager.apply_streamed_adapter.call_count, 2) + applied = [ + c.args[0] for c in lora_manager.apply_streamed_adapter.call_args_list + ] + self.assertEqual(applied, ["a", "b"]) # sorted, deterministic + self.assertEqual(mgr._lora_stash, {}) + + def test_manager_failure_propagates(self): + mgr, _ = self._manager_with_lora( + apply_result=LoRAUpdateOutput(success=False, error_message="unregistered") + ) + mgr._stash_lora_tensors([(f"a:{LORA_A}", torch.zeros(1))]) + success, message = mgr._apply_lora_stash(None) + self.assertFalse(success) + self.assertIn("unregistered", message) + + def test_checksum_match_passes(self): + mgr, _ = self._manager_with_lora() + t = torch.arange(4, dtype=torch.float32) + mgr._stash_lora_tensors([(f"a:{LORA_A}", t)]) + success, _ = mgr._apply_lora_stash({"a": {LORA_A: _sha256_tensor(t)}}) + self.assertTrue(success) + + def test_checksum_mismatch_fails(self): + mgr, lora_manager = self._manager_with_lora() + mgr._stash_lora_tensors([(f"a:{LORA_A}", torch.zeros(4))]) + success, message = mgr._apply_lora_stash( + {"a": {LORA_A: _sha256_tensor(torch.ones(4))}} + ) + self.assertFalse(success) + self.assertIn("checksum mismatch", message) + lora_manager.apply_streamed_adapter.assert_not_called() + + def test_checksum_name_set_mismatch_fails(self): + mgr, _ = self._manager_with_lora() + mgr._stash_lora_tensors([(f"a:{LORA_A}", torch.zeros(1))]) + success, message = mgr._apply_lora_stash({"a": {LORA_A: "x", LORA_B: "y"}}) + self.assertFalse(success) + self.assertIn("expected manifest", message) + + def test_manifest_adapter_set_mismatch_fails(self): + """An adapter the sender promised but streamed nothing for is a sender bug + reported to the caller, not a scheduler kill.""" + mgr, lora_manager = self._manager_with_lora() + mgr._stash_lora_tensors([(f"a:{LORA_A}", torch.zeros(1))]) + success, message = mgr._apply_lora_stash( + {"a": {LORA_A: "x"}, "b": {LORA_A: "y"}} + ) + self.assertFalse(success) + self.assertIn("expected manifest", message) + self.assertIn("'b'", message) + lora_manager.apply_streamed_adapter.assert_not_called() + + def test_forget_admits_a_different_tensor_set(self): + """Without the forget on re-register/unload, a name that legitimately + changes its tensor set is rejected as a partial stream forever.""" + mgr, _ = self._manager_with_lora() + mgr._stash_lora_tensors( + [(f"a:{LORA_A}", torch.zeros(1)), (f"a:{LORA_B}", torch.zeros(1))] + ) + self.assertTrue(mgr._apply_lora_stash(None)[0]) + mgr.forget_lora_adapter("a") + mgr._stash_lora_tensors([(f"a:{LORA_A}", torch.zeros(1))]) + self.assertTrue(mgr._apply_lora_stash(None)[0]) + + def test_shrunken_tensor_set_rejected(self): + # The adapter's tensor set is static for its lifetime; a change means a + # partial stream, which the whole-adapter replace would silently zero. + mgr, _ = self._manager_with_lora() + mgr._stash_lora_tensors( + [(f"a:{LORA_A}", torch.zeros(1)), (f"a:{LORA_B}", torch.zeros(1))] + ) + self.assertTrue(mgr._apply_lora_stash(None)[0]) + mgr._stash_lora_tensors([(f"a:{LORA_A}", torch.zeros(1))]) + success, message = mgr._apply_lora_stash(None) + self.assertFalse(success) + self.assertIn("different tensor", message) + + +class TestLoRAManagerStreamEntryPoints(CustomTestCase): + def _bare_manager(self): + manager = object.__new__(LoRAManager) + manager.lora_refs = {} + return manager + + def test_require_registered(self): + manager = self._bare_manager() + ref = LoRARef(lora_name="miles_lora", lora_path="__stream__") + manager.lora_refs[ref.lora_id] = ref + self.assertIs(manager.require_registered("miles_lora"), ref) + with self.assertRaisesRegex(ValueError, "unregistered"): + manager.require_registered("missing") + + def test_register_delegates_with_empty_tensors_and_upsert(self): + manager = self._bare_manager() + manager.load_lora_adapter_from_tensors = MagicMock( + return_value=LoRAUpdateOutput(success=True) + ) + ref = LoRARef(lora_name="miles_lora", lora_path="__stream__") + result = manager.register_lora_adapter(ref, {"r": 8}) + self.assertTrue(result.success) + kwargs = manager.load_lora_adapter_from_tensors.call_args.kwargs + self.assertEqual(kwargs["tensors"], {}) + self.assertEqual(kwargs["config_dict"], {"r": 8}) + self.assertTrue(kwargs["upsert"]) + + def test_apply_streamed_adapter_reuses_registered_config(self): + manager = self._bare_manager() + ref = LoRARef(lora_name="miles_lora", lora_path="__stream__") + manager.lora_refs[ref.lora_id] = ref + manager.load_lora_adapter_from_tensors = MagicMock( + return_value=LoRAUpdateOutput(success=True) + ) + tensors = {LORA_A: torch.zeros(1)} + result = manager.apply_streamed_adapter("miles_lora", tensors) + self.assertTrue(result.success) + args, kwargs = manager.load_lora_adapter_from_tensors.call_args + self.assertIs(args[0], ref) + self.assertIs(args[1], tensors) + self.assertIsNone(kwargs["config_dict"]) + self.assertTrue(kwargs["upsert"]) + + def test_apply_streamed_adapter_unregistered_fails(self): + manager = self._bare_manager() + result = manager.apply_streamed_adapter("missing", {LORA_A: torch.zeros(1)}) + self.assertFalse(result.success) + self.assertIn("unregistered", result.error_message) + + +if __name__ == "__main__": + unittest.main() + + +class TestWeightVersionDeferredToCommit(CustomTestCase): + """Inside a session the version reported by a bucket must not reach the + scheduler until end_weight_update commits; a failed commit drops it.""" + + def _manager(self): + mgr = _make_manager(scheduler=Mock()) + mgr._weight_update_in_progress = True + return mgr + + def test_bucket_version_is_held_while_the_session_is_open(self): + mgr = self._manager() + mgr.record_weight_version_after_update("7") + mgr.scheduler.record_weight_version_change.assert_not_called() + self.assertEqual(mgr._weight_update_pending_version, "7") + + def test_outside_a_session_the_version_is_recorded_at_once(self): + mgr = _make_manager(scheduler=Mock()) + mgr.record_weight_version_after_update("7") + mgr.scheduler.record_weight_version_change.assert_called_once_with( + new_version="7" + ) + + def test_commit_records_the_held_version(self): + mgr = self._manager() + mgr.record_weight_version_after_update("7") + mgr._weight_update_in_progress = False + mgr.record_weight_version_after_update(mgr._weight_update_pending_version) + mgr.scheduler.record_weight_version_change.assert_called_once_with( + new_version="7" + ) diff --git a/test/registered/unit/lora/test_lora_upsert.py b/test/registered/unit/lora/test_lora_upsert.py new file mode 100644 index 000000000000..6bb08e852798 --- /dev/null +++ b/test/registered/unit/lora/test_lora_upsert.py @@ -0,0 +1,425 @@ +"""Unit tests for the LoRA upsert (in-place refresh) path. + +With upsert=True, loading an adapter that is already registered refreshes +its weights in place, reusing the existing lora_id and memory-pool slot +instead of failing with a duplicate error. Covers: + + * LoRARegistry.get_lora_id / register_or_reuse / refresh + * LoRAManager.load_lora_adapter_from_tensors upsert semantics + * failed-upsert rollback (no half-updated live adapter) + * num_pinned_loras consistency across pinned flips + * LoRAManager.validate_new_adapter duplicate-name / starvation checks + * TokenizerControlMixin from_distributed id reuse; the from_tensors route + rejects upsert explicitly (only from_distributed supports in-place refresh) +""" + +import asyncio +import unittest +from unittest.mock import MagicMock, Mock + +import torch + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + +from sglang.srt.lora.lora_manager import LoRAManager +from sglang.srt.lora.lora_registry import LoRARef, LoRARegistry + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + +CONFIG_DICT = {"target_modules": ["q_proj"], "r": 8, "lora_alpha": 16} + + +class TestLoRARegistryGetLoraId(CustomTestCase): + def test_returns_id_for_registered_adapter(self): + registry = LoRARegistry() + ref = LoRARef(lora_name="a", lora_path="/x") + asyncio.run(registry.register(ref)) + + self.assertEqual(asyncio.run(registry.get_lora_id("a")), ref.lora_id) + + def test_returns_none_for_unregistered_adapter(self): + registry = LoRARegistry() + self.assertIsNone(asyncio.run(registry.get_lora_id("missing"))) + + +def _make_manager(*, lora_no_cpu_backup=False) -> LoRAManager: + """Create a LoRAManager via __new__ with only the fields the load path reads.""" + manager = LoRAManager.__new__(LoRAManager) + manager.configs = {} + manager.loras = {} + manager.lora_refs = {} + manager.num_pinned_loras = 0 + manager.lora_no_cpu_backup = lora_no_cpu_backup + manager.max_loras_per_batch = 4 + manager.base_hf_config = MagicMock(vocab_size=32000) + manager.lora_modules = [] + manager.embed_tokens_module = None + manager.lm_head_module = None + manager.device = torch.device("cpu") + manager.memory_pool = MagicMock() + manager.memory_pool.can_support.return_value = True + manager.memory_pool.uid_to_buffer_id = {} + # Weight loading needs a real base model / backend; just build a stub. + manager._create_lora_adapter_from_tensors = Mock( + side_effect=lambda ref, config, tensors: MagicMock() + ) + return manager + + +class TestLoRAManagerUpsert(CustomTestCase): + def test_fresh_load_registers_adapter(self): + manager = _make_manager() + ref = LoRARef(lora_name="a", lora_path="__stream__", pinned=True) + + result = manager.load_lora_adapter_from_tensors(ref, {}, CONFIG_DICT) + + self.assertTrue(result.success) + self.assertIn(ref.lora_id, manager.loras) + self.assertIs(manager.lora_refs[ref.lora_id], ref) + self.assertEqual(manager.num_pinned_loras, 1) + + def test_duplicate_load_without_upsert_asserts(self): + manager = _make_manager() + ref = LoRARef(lora_name="a", lora_path="__stream__") + manager.load_lora_adapter_from_tensors(ref, {}, CONFIG_DICT) + + with self.assertRaises(AssertionError): + manager.load_lora_adapter_from_tensors(ref, {}, CONFIG_DICT) + + def test_upsert_refreshes_loaded_adapter_in_place(self): + manager = _make_manager() + ref = LoRARef(lora_name="a", lora_path="__stream__", pinned=True) + manager.load_lora_adapter_from_tensors(ref, {}, CONFIG_DICT) + # Simulate the adapter occupying a memory-pool buffer slot. + manager.memory_pool.uid_to_buffer_id = {ref.lora_id: 3} + + new_config = dict(CONFIG_DICT, lora_alpha=32) + result = manager.load_lora_adapter_from_tensors( + ref, {}, new_config, upsert=True + ) + + self.assertTrue(result.success) + self.assertEqual(manager.configs[ref.lora_id].lora_alpha, 32) + # Weights are re-copied into the existing buffer slot. + manager.memory_pool.load_lora_weight_to_buffer.assert_called_once() + call = manager.memory_pool.load_lora_weight_to_buffer.call_args + self.assertEqual(call.args[0], ref.lora_id) + self.assertEqual(call.args[1], 3) + # The pinned slot is not double-counted. + self.assertEqual(manager.num_pinned_loras, 1) + + def test_upsert_skips_pool_copy_when_not_resident(self): + manager = _make_manager() + ref = LoRARef(lora_name="a", lora_path="__stream__") + manager.load_lora_adapter_from_tensors(ref, {}, CONFIG_DICT) + + result = manager.load_lora_adapter_from_tensors( + ref, {}, CONFIG_DICT, upsert=True + ) + + self.assertTrue(result.success) + manager.memory_pool.load_lora_weight_to_buffer.assert_not_called() + + def test_upsert_falls_back_to_register_when_not_loaded(self): + manager = _make_manager() + ref = LoRARef(lora_name="a", lora_path="__stream__", pinned=True) + + result = manager.load_lora_adapter_from_tensors( + ref, {}, CONFIG_DICT, upsert=True + ) + + self.assertTrue(result.success) + self.assertIn(ref.lora_id, manager.loras) + self.assertEqual(manager.num_pinned_loras, 1) + + +class TestLoRARegistryRegisterOrReuse(CustomTestCase): + def test_upsert_reuses_id_of_registered_adapter(self): + registry = LoRARegistry() + existing = LoRARef(lora_name="a", lora_path="/x", pinned=False) + asyncio.run(registry.register(existing)) + + candidate = LoRARef(lora_name="a", lora_path="__stream__", pinned=True) + resolved, reused = asyncio.run(registry.register_or_reuse(candidate, True)) + + self.assertTrue(reused) + self.assertEqual(resolved.lora_id, existing.lora_id) + self.assertEqual(resolved.lora_path, "__stream__") + self.assertTrue(resolved.pinned) + + def test_upsert_without_registered_adapter_keeps_fresh_id(self): + registry = LoRARegistry() + candidate = LoRARef(lora_name="a", lora_path="__stream__") + + resolved, reused = asyncio.run(registry.register_or_reuse(candidate, True)) + + self.assertFalse(reused) + self.assertIs(resolved, candidate) + + def test_non_upsert_never_reuses(self): + registry = LoRARegistry() + asyncio.run(registry.register(LoRARef(lora_name="a", lora_path="/x"))) + + candidate = LoRARef(lora_name="a", lora_path="/x") + resolved, reused = asyncio.run(registry.register_or_reuse(candidate, False)) + + self.assertFalse(reused) + self.assertIs(resolved, candidate) + + def test_refresh_replaces_ref_in_place(self): + registry = LoRARegistry() + existing = LoRARef(lora_name="a", lora_path="/x", pinned=False) + asyncio.run(registry.register(existing)) + + refreshed = LoRARef( + lora_id=existing.lora_id, + lora_name="a", + lora_path="__stream__", + pinned=True, + ) + asyncio.run(registry.refresh(refreshed)) + + self.assertEqual(registry.get_all_adapters()["a"].pinned, True) + self.assertEqual(asyncio.run(registry.get_lora_id("a")), existing.lora_id) + + def test_refresh_rejects_id_mismatch(self): + registry = LoRARegistry() + asyncio.run(registry.register(LoRARef(lora_name="a", lora_path="/x"))) + + with self.assertRaises(AssertionError): + asyncio.run( + registry.refresh(LoRARef(lora_name="a", lora_path="__stream__")) + ) + + +class TestUpsertRollback(CustomTestCase): + """A failed load/upsert must not leave a live adapter half-updated.""" + + def test_failed_fresh_load_leaves_no_state(self): + manager = _make_manager() + manager._create_lora_adapter_from_tensors = Mock( + side_effect=ValueError("bad tensors") + ) + ref = LoRARef(lora_name="a", lora_path="__stream__") + + result = manager.load_lora_adapter_from_tensors(ref, {}, CONFIG_DICT) + + self.assertFalse(result.success) + self.assertIn("bad tensors", result.error_message) + self.assertEqual(manager.configs, {}) + self.assertEqual(manager.loras, {}) + self.assertEqual(manager.lora_refs, {}) + + def test_failed_upsert_staging_keeps_old_adapter_serving(self): + manager = _make_manager() + ref = LoRARef(lora_name="a", lora_path="__stream__") + manager.load_lora_adapter_from_tensors(ref, {}, CONFIG_DICT) + old_config = manager.configs[ref.lora_id] + old_lora = manager.loras[ref.lora_id] + + manager._create_lora_adapter_from_tensors = Mock( + side_effect=ValueError("rank mismatch") + ) + result = manager.load_lora_adapter_from_tensors( + ref, {}, CONFIG_DICT, upsert=True + ) + + self.assertFalse(result.success) + self.assertIs(manager.configs[ref.lora_id], old_config) + self.assertIs(manager.loras[ref.lora_id], old_lora) + manager.memory_pool.load_lora_weight_to_buffer.assert_not_called() + + def test_failed_buffer_rewrite_restores_old_weights(self): + manager = _make_manager() + ref = LoRARef(lora_name="a", lora_path="__stream__") + manager.load_lora_adapter_from_tensors(ref, {}, CONFIG_DICT) + old_config = manager.configs[ref.lora_id] + old_lora = manager.loras[ref.lora_id] + manager.memory_pool.uid_to_buffer_id = {ref.lora_id: 3} + manager.memory_pool.load_lora_weight_to_buffer.side_effect = [ + RuntimeError("copy failed at layer k"), + None, # the restore pass + ] + + result = manager.load_lora_adapter_from_tensors( + ref, {}, CONFIG_DICT, upsert=True + ) + + self.assertFalse(result.success) + self.assertIn("copy failed", result.error_message) + # CPU-side state rolled back... + self.assertIs(manager.configs[ref.lora_id], old_config) + self.assertIs(manager.loras[ref.lora_id], old_lora) + # ...and the served buffer was rewritten back from the old adapter. + calls = manager.memory_pool.load_lora_weight_to_buffer.call_args_list + self.assertEqual(len(calls), 2) + self.assertIs(calls[1].args[2], old_lora) + + +class TestStreamedUpsert(CustomTestCase): + def test_failed_install_keeps_previous_metadata(self): + manager = _make_manager(lora_no_cpu_backup=True) + ref = LoRARef(lora_name="a", lora_path="__stream__", pinned=False) + manager.load_lora_adapter_from_tensors(ref, {}, CONFIG_DICT) + old_config, old_lora = manager.configs[ref.lora_id], manager.loras[ref.lora_id] + manager.memory_pool.install_streamed_adapter.side_effect = ValueError( + "bad shape" + ) + result = manager.load_lora_adapter_from_tensors( + ref, + {"weight": torch.ones(1)}, + dict(CONFIG_DICT, lora_alpha=32), + upsert=True, + ) + self.assertFalse(result.success) + self.assertIn("bad shape", result.error_message) + self.assertIs(manager.configs[ref.lora_id], old_config) + self.assertIs(manager.loras[ref.lora_id], old_lora) + self.assertIs(manager.lora_refs[ref.lora_id], ref) + self.assertEqual(manager.num_pinned_loras, 0) + + def test_failed_fresh_install_publishes_no_metadata(self): + manager = _make_manager(lora_no_cpu_backup=True) + ref = LoRARef(lora_name="a", lora_path="__stream__") + manager.memory_pool.install_streamed_adapter.side_effect = RuntimeError( + "No free slot" + ) + result = manager.load_lora_adapter_from_tensors( + ref, {"weight": torch.ones(1)}, CONFIG_DICT + ) + self.assertFalse(result.success) + self.assertEqual(manager.configs, {}) + self.assertEqual(manager.loras, {}) + self.assertEqual(manager.lora_refs, {}) + + def test_install_precedes_metadata_publication(self): + manager = _make_manager(lora_no_cpu_backup=True) + ref = LoRARef(lora_name="a", lora_path="__stream__") + manager.load_lora_adapter_from_tensors(ref, {}, CONFIG_DICT) + old_lora = manager.loras[ref.lora_id] + + def install(uid, adapter, *_args): + self.assertIs(manager.loras[uid], old_lora) + self.assertIsNot(adapter, old_lora) + + manager.memory_pool.install_streamed_adapter.side_effect = install + result = manager.load_lora_adapter_from_tensors( + ref, {"weight": torch.ones(1)}, CONFIG_DICT, upsert=True + ) + self.assertTrue(result.success) + self.assertIsNot(manager.loras[ref.lora_id], old_lora) + + def test_re_registration_installs_zero_weights_into_resident_slot(self): + manager = _make_manager(lora_no_cpu_backup=True) + ref = LoRARef(lora_name="a", lora_path="__stream__") + manager.register_lora_adapter(ref, CONFIG_DICT) + manager.memory_pool.uid_to_buffer_id = {ref.lora_id: 0} + result = manager.register_lora_adapter(ref, CONFIG_DICT) + self.assertTrue(result.success) + manager.memory_pool.install_streamed_adapter.assert_called_once() + manager.memory_pool.load_lora_weight_to_buffer.assert_not_called() + + def test_invalid_pool_blocks_batch_preparation(self): + manager = _make_manager(lora_no_cpu_backup=True) + manager.memory_pool.check_valid.side_effect = RuntimeError( + "incomplete GPU installs" + ) + with self.assertRaisesRegex(RuntimeError, "incomplete GPU installs"): + manager.prepare_lora_batch(Mock()) + + +class TestUpsertPinnedAccounting(CustomTestCase): + def test_pinned_flip_updates_counter_both_ways(self): + manager = _make_manager() + unpinned = LoRARef(lora_name="a", lora_path="__stream__", pinned=False) + manager.load_lora_adapter_from_tensors(unpinned, {}, CONFIG_DICT) + self.assertEqual(manager.num_pinned_loras, 0) + + pinned = LoRARef( + lora_id=unpinned.lora_id, + lora_name="a", + lora_path="__stream__", + pinned=True, + ) + manager.load_lora_adapter_from_tensors(pinned, {}, CONFIG_DICT, upsert=True) + self.assertEqual(manager.num_pinned_loras, 1) + + manager.load_lora_adapter_from_tensors(unpinned, {}, CONFIG_DICT, upsert=True) + self.assertEqual(manager.num_pinned_loras, 0) + + def test_unload_after_pinned_flip_keeps_counter_consistent(self): + manager = _make_manager() + unpinned = LoRARef(lora_name="a", lora_path="__stream__", pinned=False) + manager.load_lora_adapter_from_tensors(unpinned, {}, CONFIG_DICT) + pinned = LoRARef( + lora_id=unpinned.lora_id, + lora_name="a", + lora_path="__stream__", + pinned=True, + ) + manager.load_lora_adapter_from_tensors(pinned, {}, CONFIG_DICT, upsert=True) + + result = manager.unload_lora_adapter(pinned) + + self.assertTrue(result.success) + self.assertEqual(manager.num_pinned_loras, 0) + + def test_pinned_refresh_allowed_at_pin_limit(self): + """Refreshing a pinned adapter adds no pinned slot; rejecting it would + freeze RL serving on the step-1 weights.""" + manager = _make_manager() + manager.max_loras_per_batch = 2 + ref = LoRARef(lora_name="a", lora_path="__stream__", pinned=True) + manager.load_lora_adapter_from_tensors(ref, {}, CONFIG_DICT) + self.assertEqual(manager.num_pinned_loras, 1) + + result = manager.load_lora_adapter_from_tensors( + ref, {}, CONFIG_DICT, upsert=True + ) + + self.assertTrue(result.success) + self.assertEqual(manager.num_pinned_loras, 1) + + def test_fresh_pinned_load_still_rejected_at_pin_limit(self): + manager = _make_manager() + manager.max_loras_per_batch = 2 + manager.load_lora_adapter_from_tensors( + LoRARef(lora_name="a", lora_path="__stream__", pinned=True), + {}, + CONFIG_DICT, + ) + + result = manager.load_lora_adapter_from_tensors( + LoRARef(lora_name="b", lora_path="__stream__", pinned=True), + {}, + CONFIG_DICT, + ) + + self.assertFalse(result.success) + self.assertIn("not allowed to pin all slots", result.error_message) + + +class TestValidateNewAdapterDuplicates(CustomTestCase): + def test_duplicate_name_rejected_without_update(self): + manager = _make_manager() + existing = LoRARef(lora_name="a", lora_path="/x") + manager.lora_refs[existing.lora_id] = existing + config = MagicMock(lora_added_tokens_size=0, use_dora=False) + + with self.assertRaisesRegex(ValueError, "already loaded"): + manager.validate_new_adapter(config, LoRARef(lora_name="a", lora_path="/y")) + + def test_duplicate_name_allowed_for_update(self): + manager = _make_manager() + existing = LoRARef(lora_name="a", lora_path="/x") + manager.lora_refs[existing.lora_id] = existing + config = MagicMock(lora_added_tokens_size=0, use_dora=False) + + manager.validate_new_adapter(config, existing, is_update=True) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/test/registered/unit/lora/test_mem_pool_ep_unit.py b/test/registered/unit/lora/test_mem_pool_ep_unit.py index 4a3ad357e6ba..8786aa7d20d2 100644 --- a/test/registered/unit/lora/test_mem_pool_ep_unit.py +++ b/test/registered/unit/lora/test_mem_pool_ep_unit.py @@ -19,6 +19,7 @@ register_cpu_ci(est_time=8, suite="base-a-test-cpu") import ast +import itertools import types import unittest import unittest.mock as mock @@ -28,6 +29,7 @@ import torch from sglang.srt.lora.eviction_policy import get_eviction_policy +from sglang.test.test_utils import CustomTestCase REPO_ROOT = Path(__file__).resolve().parents[4] MOE_UTILS_PATH = REPO_ROOT / "python/sglang/srt/layers/moe/utils.py" @@ -36,6 +38,14 @@ ) +class _FakeFusedMoE(torch.nn.Module): + pass + + +_FUSED_MOE_STUB = types.ModuleType("sglang.srt.layers.moe.fused_moe_triton.layer") +_FUSED_MOE_STUB.FusedMoE = _FakeFusedMoE + + class _FakeBaseLayerWithLoRA: pass @@ -108,7 +118,13 @@ def slice_lora_b_weights(self, weights): def _load_lora_weight_to_buffer(pool, **kwargs): - with mock.patch.dict("sys.modules", {"sglang.srt.lora.layers": _LORA_LAYERS_STUB}): + with mock.patch.dict( + "sys.modules", + { + "sglang.srt.lora.layers": _LORA_LAYERS_STUB, + "sglang.srt.layers.moe.fused_moe_triton.layer": _FUSED_MOE_STUB, + }, + ): return pool.load_lora_weight_to_buffer(**kwargs) @@ -193,6 +209,7 @@ def _make_pool( helpers under test. """ pool = LoRAMemoryPool.__new__(LoRAMemoryPool) + pool.invalid_uids = set() pool.moe_ep_size = moe_ep_size pool.moe_ep_rank = moe_ep_rank pool.moe_use_local_expert_ids = moe_use_local_expert_ids @@ -204,6 +221,21 @@ def _make_pool( pool._num_experts_local = num_experts_global // moe_ep_size else: pool._num_experts_local = num_experts_global + pool.num_layer = 1 + pool.tp_rank = 0 + pool.max_lora_rank = 2 + pool.target_modules = {"down_proj"} + pool.experts_shared_outer_loras = False + pool.strict_loading = True + pool.lora_added_tokens_size = 0 + pool.pin_memory_available = False + pool.enable_lora_overlap_loading = False + pool.base_model = torch.nn.Identity() + pool.embedding_A_buffer = {} + pool.embedding_B_buffer = {} + pool.lm_head_A_buffer = {} + pool.lm_head_B_buffer = {} + pool.new_embeddings_buffer = {} return pool @@ -220,6 +252,8 @@ class TestDeterministicPoolSlots(unittest.TestCase): @staticmethod def _prepare(iteration_order): pool = LoRAMemoryPool.__new__(LoRAMemoryPool) + pool.invalid_uids = set() + pool.lora_no_cpu_backup = False pool.max_loras_per_batch = 4 pool.uid_to_buffer_id = {} pool.buffer_id_to_uid = [EMPTY_SLOT] * 4 @@ -261,6 +295,8 @@ def test_uid_iteration_order_does_not_change_slots_or_lru(self): @staticmethod def _evict_after_touch(iteration_order): pool = LoRAMemoryPool.__new__(LoRAMemoryPool) + pool.invalid_uids = set() + pool.lora_no_cpu_backup = False pool.max_loras_per_batch = 4 pool.uid_to_buffer_id = { None: 0, @@ -326,6 +362,7 @@ def _buffers(pool): @classmethod def _make_pool(cls): pool = LoRAMemoryPool.__new__(LoRAMemoryPool) + pool.invalid_uids = set() pool.num_layer = 2 pool.A_buffer = {"a": [torch.ones(2, 1, 2) for _ in range(2)]} pool.B_buffer = {"b": [torch.full((2, 1, 2), float("nan")) for _ in range(2)]} @@ -638,24 +675,157 @@ def test_rank_with_partial_tensor_coverage(self): list(pool._iter_local_expert_weights(weights, "weights")) -class TestSharedMoeProductionLoad(unittest.TestCase): - def test_unmarked_2d_shared_expert_uses_dense_buffers(self): - pool = _make_pool( - num_experts_global=8, +@mock.patch.dict( + "sys.modules", {"sglang.srt.layers.moe.fused_moe_triton.layer": _FUSED_MOE_STUB} +) +class TestStreamedInstall(unittest.TestCase): + def setUp(self): + self.pool = pool = _make_pool( + num_experts_global=1, moe_ep_size=1, moe_ep_rank=0, moe_use_local_expert_ids=False, ) - pool.num_layer = 1 + pool.num_layer = 2 pool.tp_rank = 0 pool.max_lora_rank = 2 + pool.max_loras_per_batch = 1 pool.target_modules = {"down_proj"} pool.experts_shared_outer_loras = False pool.strict_loading = True + pool.lora_no_cpu_backup = True pool.lora_added_tokens_size = 0 pool.pin_memory_available = False pool.enable_lora_overlap_loading = False - pool.base_model = object() + pool.base_model = torch.nn.Linear(3, 5) + pool.A_buffer = {"down_proj": [torch.full((1, 2, 3), -1.0) for _ in range(2)]} + pool.B_buffer = {"down_proj": [torch.full((1, 5, 2), -1.0) for _ in range(2)]} + pool.embedding_A_buffer = {} + pool.embedding_B_buffer = {} + pool.lm_head_A_buffer = {} + pool.lm_head_B_buffer = {} + pool.new_embeddings_buffer = {} + pool.uid_to_buffer_id = {"adapter": 0} + pool.buffer_id_to_uid = ["adapter"] + pool.eviction_policy = get_eviction_policy("lru") + pool.eviction_policy.mark_used("adapter") + self.modules = [ + {f"model.layers.{i}.mlp.down_proj": _FakeDenseLayer()} for i in range(2) + ] + self.adapter = types.SimpleNamespace( + config=types.SimpleNamespace(r=2), + layers=[ + types.SimpleNamespace( + weights={ + f"model.layers.{i}.mlp.down_proj.lora_A.weight": torch.full( + (2, 3), 2.0 + ), + f"model.layers.{i}.mlp.down_proj.lora_B.weight": torch.full( + (5, 2), 3.0 + ), + }, + pinned_weights={}, + ) + for i in range(2) + ], + embedding_layers={}, + added_tokens_embeddings={}, + release_staged_weights=mock.Mock(), + ) + + def install(self, uid="adapter"): + with mock.patch.dict( + "sys.modules", {"sglang.srt.lora.layers": _LORA_LAYERS_STUB} + ): + self.pool.install_streamed_adapter( + uid, self.adapter, self.modules, None, None + ) + + def test_later_shape_error_preserves_all_resident_weights(self): + self.adapter.layers[1].weights["model.layers.1.mlp.down_proj.lora_B.weight"] = ( + torch.zeros(4, 2) + ) + with self.assertRaisesRegex(AssertionError, "LoRA buffer shape"): + self.install() + for buffer in (self.pool.A_buffer, self.pool.B_buffer): + for tensor in buffer["down_proj"]: + self.assertTrue(torch.all(tensor == -1)) + self.assertEqual(self.pool.invalid_uids, set()) + self.assertEqual(self.pool.uid_to_buffer_id, {"adapter": 0}) + self.adapter.release_staged_weights.assert_not_called() + + def test_full_pool_preserves_resident_adapter(self): + with self.assertRaisesRegex(RuntimeError, "No free LoRA memory pool slot"): + self.install("another-adapter") + self.assertEqual(self.pool.uid_to_buffer_id, {"adapter": 0}) + self.assertEqual(self.pool.buffer_id_to_uid, ["adapter"]) + self.assertEqual(self.pool.invalid_uids, set()) + self.adapter.release_staged_weights.assert_not_called() + + def test_invalid_weights_preserve_base_placeholder(self): + self.pool.uid_to_buffer_id = {None: 0} + self.pool.buffer_id_to_uid = [None] + self.adapter.layers[1].weights["model.layers.1.mlp.down_proj.lora_B.weight"] = ( + torch.zeros(4, 2) + ) + with self.assertRaisesRegex(AssertionError, "LoRA buffer shape"): + self.install() + self.assertEqual(self.pool.uid_to_buffer_id, {None: 0}) + self.assertEqual(self.pool.buffer_id_to_uid, [None]) + + def test_failed_write_blocks_inference_until_reload(self): + copy_weight = mem_pool_module.copy_weight_into_buffer + copies = 0 + + def fail_on_third_copy(buffer, weight): + nonlocal copies + copies += 1 + if copies == 3: + raise RuntimeError("injected copy failure") + copy_weight(buffer, weight) + + with mock.patch.object( + mem_pool_module, "copy_weight_into_buffer", fail_on_third_copy + ): + with self.assertRaisesRegex(RuntimeError, "injected copy failure"): + self.install() + self.assertTrue(torch.all(self.pool.A_buffer["down_proj"][0] == 2)) + self.assertTrue(torch.all(self.pool.A_buffer["down_proj"][1] == -1)) + self.assertEqual(self.pool.invalid_uids, {"adapter"}) + self.adapter.release_staged_weights.assert_not_called() + with self.assertRaisesRegex(RuntimeError, "incomplete GPU installs"): + self.pool.prepare_lora_batch({"adapter"}, {}, [], {}, None, None) + + self.install() + self.pool.check_valid() + for a, b in zip( + self.pool.A_buffer["down_proj"], self.pool.B_buffer["down_proj"] + ): + self.assertTrue(torch.all(a == 2)) + self.assertTrue(torch.all(b == 3)) + self.adapter.release_staged_weights.assert_called_once() + + def test_unload_clears_invalid_slot(self): + self.pool.invalid_uids.add("adapter") + self.assertEqual(self.pool.remove_lora("adapter"), 0) + self.pool.check_valid() + self.assertEqual(self.pool.buffer_id_to_uid, [EMPTY_SLOT]) + for buffer in (self.pool.A_buffer, self.pool.B_buffer): + for tensor in buffer["down_proj"]: + self.assertTrue(torch.all(tensor == 0)) + + +@mock.patch.dict( + "sys.modules", {"sglang.srt.layers.moe.fused_moe_triton.layer": _FUSED_MOE_STUB} +) +class TestSharedMoeProductionLoad(CustomTestCase): + def test_unmarked_2d_shared_expert_uses_dense_buffers(self): + pool = _make_pool( + num_experts_global=8, + moe_ep_size=1, + moe_ep_rank=0, + moe_use_local_expert_ids=False, + ) pool.A_buffer = { "down_proj": [torch.full((1, 2, 3), -1.0)], "down_proj_shared_moe": [torch.full((1, 2, 2, 3), -7.0)], @@ -664,11 +834,6 @@ def test_unmarked_2d_shared_expert_uses_dense_buffers(self): "down_proj": [torch.full((1, 5, 2), -1.0)], "down_proj_shared_moe": [torch.full((1, 1, 5, 2), -7.0)], } - pool.embedding_A_buffer = {} - pool.embedding_B_buffer = {} - pool.lm_head_A_buffer = {} - pool.lm_head_B_buffer = {} - pool.new_embeddings_buffer = {} down_a = torch.arange(6, dtype=torch.float32).reshape(2, 3) down_b = torch.arange(10, dtype=torch.float32).reshape(5, 2) @@ -705,6 +870,117 @@ def test_unmarked_2d_shared_expert_uses_dense_buffers(self): self.assertTrue(torch.all(pool.A_buffer["down_proj_shared_moe"][0] == -7)) self.assertTrue(torch.all(pool.B_buffer["down_proj_shared_moe"][0] == -7)) + def test_fused_shared_weights_use_last_expert_in_each_adapter(self): + """Shared factors must reach fused slots for either spelling and packed order.""" + for (ep_size, ep_rank), layout, shared_name in itertools.product( + [(1, 0), (2, 0), (2, 1)], + ["per_expert", "packed_first", "shared_first"], + ["shared_expert", "shared_experts"], + ): + with self.subTest( + ep_size=ep_size, ep_rank=ep_rank, layout=layout, shared_name=shared_name + ): + pool = _make_pool( + num_experts_global=8, + moe_ep_size=ep_size, + moe_ep_rank=ep_rank, + moe_use_local_expert_ids=ep_size > 1, + ) + fused = _FakeFusedMoE() + fused.layer_id = 0 + fused.num_fused_shared_experts = 1 + fused.num_experts = 9 + fused._num_global_routed = 8 + fused.num_local_experts = 8 // ep_size + 1 + + def local_id(gid): + if gid == 8: + return 8 // ep_size + start = ep_rank * (8 // ep_size) + return gid - start if start <= gid < start + 8 // ep_size else -1 + + fused._map_global_expert_id_to_local_expert_id = local_id + pool.base_model = torch.nn.Sequential(fused) + n = pool._get_moe_pool_expert_dim(pool.base_model, 0) + self.assertEqual(n, 8 // ep_size + 1) + pool.A_buffer = { + "down_proj": [torch.full((2, 2, 3), -1.0)], + "down_proj_moe": [torch.full((2, n, 2, 3), -1.0)], + } + pool.B_buffer = { + "down_proj": [torch.full((2, 5, 2), -1.0)], + "down_proj_moe": [torch.full((2, n, 5, 2), -1.0)], + } + for slot in range(2): + weights = {} + for expert, value in [ + ("experts.0", 10 + slot), + (shared_name, 20 + slot), + ]: + prefix = f"model.layers.0.mlp.{expert}.down_proj" + weights[f"{prefix}.lora_A.weight"] = torch.full( + (2, 3), float(value) + ) + weights[f"{prefix}.lora_B.weight"] = torch.full( + (5, 2), float(value) + ) + if layout != "per_expert": + for factor in ("A", "B"): + prefix = "model.layers.0.mlp.experts" + key = f"down_proj.lora_{factor}.weight" + weight = weights.pop(f"{prefix}.0.{key}") + packed = weight.new_zeros((8, *weight.shape)) + packed[0] = weight + weights[f"{prefix}.{key}"] = packed + if layout == "packed_first": + weights = dict(reversed(list(weights.items()))) + adapter = types.SimpleNamespace( + config=types.SimpleNamespace(r=2), + scaling=2.5, + layers=[ + types.SimpleNamespace(weights=weights, pinned_weights={}) + ], + embedding_layers={}, + added_tokens_embeddings={}, + ) + _load_lora_weight_to_buffer( + pool, + uid=f"adapter-{slot}", + buffer_id=slot, + lora_adapter=adapter, + lora_modules=[ + {"model.layers.0.mlp.experts": _FakeRoutedMoeLayer()} + ], + lora_embed_tokens_module=None, + lora_lm_head_module=None, + ) + for slot in range(2): + a = pool.A_buffer["down_proj_moe"][0][slot] + b = pool.B_buffer["down_proj_moe"][0][slot] + self.assertTrue(torch.all(a[-1] == 20 + slot)) + self.assertTrue(torch.all(b[-1] == (20 + slot) * 2.5)) + expected = 10 + slot if ep_rank == 0 else 0 + self.assertTrue(torch.all(a[0] == expected)) + self.assertTrue(torch.all(b[0] == expected * 2.5)) + self.assertTrue(torch.all(a[1:-1] == 0)) + self.assertTrue(torch.all(b[1:-1] == 0)) + + def test_global_per_rank_shared_slots_are_rejected(self): + pool = _make_pool( + num_experts_global=8, + moe_ep_size=2, + moe_ep_rank=0, + moe_use_local_expert_ids=False, + ) + fused = _FakeFusedMoE() + fused.layer_id = 0 + fused.num_fused_shared_experts = 1 + fused.num_experts = 10 + fused._num_global_routed = 8 + model = torch.nn.Sequential(fused) + with self.assertRaisesRegex(AssertionError, "per-rank fused shared slots"): + pool._get_moe_pool_expert_dim(model, 0) + def test_rank3_loads_shared_gate_b_and_down_a(self): """Shared-sink weights stay replicated even on a nonzero EP rank.""" pool = _make_pool( @@ -713,15 +989,8 @@ def test_rank3_loads_shared_gate_b_and_down_a(self): moe_ep_rank=3, moe_use_local_expert_ids=True, ) - pool.num_layer = 1 - pool.max_lora_rank = 2 pool.target_modules = {"gate_up_proj", "down_proj"} pool.experts_shared_outer_loras = True - pool.strict_loading = True - pool.lora_added_tokens_size = 0 - pool.pin_memory_available = False - pool.enable_lora_overlap_loading = False - pool.base_model = object() pool.A_buffer = { "gate_up_proj": [torch.full((1, 4, 5), -7.0)], "down_proj": [torch.full((1, 2, 3), -7.0)], @@ -734,11 +1003,6 @@ def test_rank3_loads_shared_gate_b_and_down_a(self): "gate_up_proj_shared_moe": [torch.full((1, 2, 6, 2), -1.0)], "down_proj_shared_moe": [torch.full((1, 1, 5, 2), -1.0)], } - pool.embedding_A_buffer = {} - pool.embedding_B_buffer = {} - pool.lm_head_A_buffer = {} - pool.lm_head_B_buffer = {} - pool.new_embeddings_buffer = {} gate_b = torch.arange(2 * 6 * 2, dtype=torch.float32).reshape(2, 6, 2) down_a = torch.arange(2 * 2 * 3, dtype=torch.float32).reshape(2, 2, 3) @@ -819,15 +1083,8 @@ def test_missing_inner_factors_clear_reused_shared_and_routed_slots(self): moe_ep_rank=0, moe_use_local_expert_ids=False, ) - pool.num_layer = 1 - pool.max_lora_rank = 2 pool.target_modules = {"gate_up_proj", "down_proj"} pool.experts_shared_outer_loras = True - pool.strict_loading = True - pool.lora_added_tokens_size = 0 - pool.pin_memory_available = False - pool.enable_lora_overlap_loading = False - pool.base_model = object() pool.A_buffer = { f"gate_up_proj{suffix}": [torch.full((1, 1, 4, 5), -1.0)], f"down_proj{suffix}": [torch.full((1, 2, 2, 3), -1.0)], @@ -836,11 +1093,6 @@ def test_missing_inner_factors_clear_reused_shared_and_routed_slots(self): f"gate_up_proj{suffix}": [torch.full((1, 2, 6, 2), -1.0)], f"down_proj{suffix}": [torch.full((1, 1, 5, 2), -1.0)], } - pool.embedding_A_buffer = {} - pool.embedding_B_buffer = {} - pool.lm_head_A_buffer = {} - pool.lm_head_B_buffer = {} - pool.new_embeddings_buffer = {} adapter = types.SimpleNamespace( config=types.SimpleNamespace(r=2), @@ -1102,6 +1354,9 @@ def get_hidden_dim(self, module_name: str, layer_idx: int): return _Model() +@mock.patch.dict( + "sys.modules", {"sglang.srt.layers.moe.fused_moe_triton.layer": _FUSED_MOE_STUB} +) class TestMoeBufferShardsByMoeTp(unittest.TestCase): """Regression: per-expert MoE LoRA buffers must shard by `moe_tp_size`, not the outer attention `tp_size`. @@ -1124,6 +1379,7 @@ def _pool( ep_rank: int = 0, ) -> LoRAMemoryPool: pool = LoRAMemoryPool.__new__(LoRAMemoryPool) + pool.invalid_uids = set() pool.max_loras_per_batch = 2 pool.tp_size = tp_size pool.tp_rank = 0 @@ -1243,6 +1499,7 @@ class TestAttnModulesShardByAttnTp(unittest.TestCase): def _pool(self, *, tp_size: int, attn_tp_size: int) -> LoRAMemoryPool: pool = LoRAMemoryPool.__new__(LoRAMemoryPool) + pool.invalid_uids = set() pool.max_loras_per_batch = 2 pool.tp_size = tp_size pool.tp_rank = 0 @@ -1328,6 +1585,8 @@ def test_moe_tp_rank_used_for_slicing_when_ep_lt_tp(self): # tp=4 ep=2 → moe_tp_size=2. Pick OUTER rank 3 so moe_tp_rank=1. # The two values differ; the bug would surface on this exact rank. pool = LoRAMemoryPool.__new__(LoRAMemoryPool) + pool.base_model = torch.nn.Identity() + pool.invalid_uids = set() pool.tp_size = 4 pool.tp_rank = 3 pool.moe_tp_size = 2 diff --git a/test/registered/unit/lora/test_uno_inactive_lora_batch.py b/test/registered/unit/lora/test_uno_inactive_lora_batch.py index 5b47794beeac..42c3ac5faa80 100644 --- a/test/registered/unit/lora/test_uno_inactive_lora_batch.py +++ b/test/registered/unit/lora/test_uno_inactive_lora_batch.py @@ -2,6 +2,7 @@ import unittest from types import SimpleNamespace +from unittest.mock import Mock from sglang.srt.lora.lora_manager import LoRAManager from sglang.test.ci.ci_register import register_cpu_ci @@ -29,6 +30,7 @@ def test_all_base_batch_clears_stale_routing_before_graph_metadata(self): backend = _InactiveSkippingBackend() manager = LoRAManager.__new__(LoRAManager) manager.lora_backend = backend + manager.memory_pool = SimpleNamespace(check_valid=Mock()) forward_batch = SimpleNamespace(lora_ids=[None], batch_size=1) manager.prepare_lora_batch(forward_batch) diff --git a/test/registered/unit/managers/test_abort_request_prefix.py b/test/registered/unit/managers/test_abort_request_prefix.py new file mode 100644 index 000000000000..36e83d046150 --- /dev/null +++ b/test/registered/unit/managers/test_abort_request_prefix.py @@ -0,0 +1,450 @@ +"""Unit tests for abort-by-rid-prefix. + +Tokenizer side: with prefix=True, the rid is treated as a prefix — the +early-return gate matches any tracked rid starting with it, matching +tokenizer-held (not yet dispatched) requests are resolved locally, and the +AbortReq is forwarded with prefix=True. + +Scheduler side: matching is prefix-based (``rid.startswith``) regardless of +the flag, because batch requests derive child rids as ``f"{rid}_{i}"``. +""" + +import asyncio +import unittest +from types import SimpleNamespace +from unittest.mock import MagicMock, Mock, patch + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + +from sglang.srt.disaggregation.utils import DisaggregationMode +from sglang.srt.managers.io_struct import AbortReq +from sglang.srt.managers.schedule_batch import FINISH_ABORT +from sglang.srt.managers.scheduler import Scheduler +from sglang.srt.managers.tokenizer_manager import ReqState, TokenizerManager + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +def _make_tokenizer_manager(rids=(), tokenizer_worker_num=1) -> TokenizerManager: + """Create a TokenizerManager with mocked dependencies, bypassing __init__.""" + tm = TokenizerManager.__new__(TokenizerManager) + tm.server_args = MagicMock() + tm.server_args.tokenizer_worker_num = tokenizer_worker_num + tm.enable_metrics = False + tm.enable_lora = False + tm.rid_to_state = {rid: Mock() for rid in rids} + tm.send_to_scheduler = MagicMock() + tm.tokenizer_ipc_name = None + # The IPC boundary: sock_send's wire format varies (pickle/msgpack), so + # tests observe dispatched objects here instead of on the zmq socket. + tm._dispatch_to_scheduler = MagicMock() + return tm + + +def _sent_req(tm) -> AbortReq: + tm._dispatch_to_scheduler.assert_called_once() + return tm._dispatch_to_scheduler.call_args.args[0] + + +class TestAbortRequestPrefix(CustomTestCase): + def test_prefix_match_sends_abort(self): + tm = _make_tokenizer_manager(rids=["job-1-seq-0", "job-1-seq-1", "other"]) + tm.abort_request(rid="job-1", prefix=True) + + req = _sent_req(tm) + self.assertEqual(req.rid, "job-1") + self.assertTrue(req.prefix) + self.assertFalse(req.abort_all) + + def test_prefix_without_match_is_ignored(self): + tm = _make_tokenizer_manager(rids=["other-1", "other-2"]) + tm.abort_request(rid="job-1", prefix=True) + + tm._dispatch_to_scheduler.assert_not_called() + + def test_prefix_requires_full_prefix_not_substring(self): + tm = _make_tokenizer_manager(rids=["seq-job-1"]) + tm.abort_request(rid="job-1", prefix=True) + + tm._dispatch_to_scheduler.assert_not_called() + + def test_exact_match_still_works_without_prefix(self): + tm = _make_tokenizer_manager(rids=["job-1"]) + tm.abort_request(rid="job-1") + + req = _sent_req(tm) + self.assertEqual(req.rid, "job-1") + self.assertFalse(req.prefix) + + def test_exact_mode_does_not_prefix_match(self): + # rid is only a prefix of a tracked request; exact mode must ignore it. + tm = _make_tokenizer_manager(rids=["job-1-seq-0"]) + tm.abort_request(rid="job-1") + + tm._dispatch_to_scheduler.assert_not_called() + + def test_empty_rid_is_ignored(self): + # An empty rid would prefix-match every request on the scheduler. + tm = _make_tokenizer_manager(rids=["job-1"]) + tm.abort_request(rid="", prefix=True) + + tm._dispatch_to_scheduler.assert_not_called() + + def test_multi_tokenizer_worker_skips_local_check(self): + # With >1 tokenizer workers, rid_to_state is not authoritative; the + # abort must be forwarded even if this worker tracks no matching rid. + tm = _make_tokenizer_manager(rids=[], tokenizer_worker_num=2) + tm.abort_request(rid="job-1", prefix=True) + + req = _sent_req(tm) + self.assertEqual(req.rid, "job-1") + self.assertTrue(req.prefix) + + +def _make_state(rid: str) -> ReqState: + # lora_path / lora_id model the real input-struct contract: always present, + # None = base model (the lease finalizer does plain None checks on them). + obj = SimpleNamespace( + rid=rid, stream=False, return_logprob=False, lora_path=None, lora_id=None + ) + return ReqState([], False, asyncio.Event(), obj, MagicMock()) + + +def _make_tokenizer_manager_with_states(rids) -> TokenizerManager: + tm = _make_tokenizer_manager() + tm.rid_to_state = {rid: _make_state(rid) for rid in rids} + return tm + + +class TestAbortTokenizerHeldRequests(CustomTestCase): + """An abort must also cover requests admitted to the tokenizer (rid_to_state + populated in _init_req_state) but not yet dispatched to the scheduler — + e.g. parked at the pause gate during a weight update. The scheduler cannot + match those rids, so abort_request flags them and the dispatch path + resolves them as aborted instead of sending them.""" + + def test_prefix_abort_flags_matching_states(self): + tm = _make_tokenizer_manager_with_states( + ["job-1-seq-0", "job-1-seq-1", "other"] + ) + tm.abort_request(rid="job-1", prefix=True) + + self.assertTrue(tm.rid_to_state["job-1-seq-0"].abort_before_dispatch) + self.assertTrue(tm.rid_to_state["job-1-seq-1"].abort_before_dispatch) + self.assertFalse(tm.rid_to_state["other"].abort_before_dispatch) + # The AbortReq is still forwarded for already-dispatched requests. + self.assertEqual(_sent_req(tm).rid, "job-1") + + def test_exact_abort_flags_only_exact_state(self): + tm = _make_tokenizer_manager_with_states(["job-1", "job-1-seq-0"]) + tm.abort_request(rid="job-1") + + self.assertTrue(tm.rid_to_state["job-1"].abort_before_dispatch) + self.assertFalse(tm.rid_to_state["job-1-seq-0"].abort_before_dispatch) + + def test_abort_all_flags_every_state(self): + tm = _make_tokenizer_manager_with_states(["a", "b"]) + tm.abort_request(abort_all=True) + + self.assertTrue(tm.rid_to_state["a"].abort_before_dispatch) + self.assertTrue(tm.rid_to_state["b"].abort_before_dispatch) + + def test_dispatch_resolves_flagged_request_locally(self): + tm = _make_tokenizer_manager_with_states(["job-1-seq-0"]) + tm.server_args.weight_version = "v0" + state = tm.rid_to_state["job-1-seq-0"] + tm.abort_request(rid="job-1", prefix=True) + tm._dispatch_to_scheduler.reset_mock() + + tm._send_one_request(SimpleNamespace(rid="job-1-seq-0")) + + # Never dispatched; resolved as aborted so _wait_one_response returns. + tm._dispatch_to_scheduler.assert_not_called() + self.assertTrue(state.finished) + self.assertTrue(state.event.is_set()) + self.assertNotIn("job-1-seq-0", tm.rid_to_state) + finish_reason = state.out_list[-1]["meta_info"]["finish_reason"] + self.assertEqual(finish_reason["type"], "abort") + + def test_dispatch_sends_unflagged_request(self): + tm = _make_tokenizer_manager_with_states(["job-1", "other"]) + tm.abort_request(rid="job-1") + tm._dispatch_to_scheduler.reset_mock() + + tokenized_obj = MagicMock() + tokenized_obj.rid = "other" + with unittest.mock.patch( + "sglang.srt.managers.tokenizer_manager.wrap_shm_features", + side_effect=lambda obj: obj, + ): + tm._send_one_request(tokenized_obj) + + tm._dispatch_to_scheduler.assert_called_once_with(tokenized_obj) + self.assertIn("other", tm.rid_to_state) + + def test_batch_dispatch_filters_flagged_requests(self): + tm = _make_tokenizer_manager_with_states(["job-1-seq-0", "job-1-seq-1"]) + tm.server_args.weight_version = "v0" + tm.abort_request(rid="job-1", prefix=True) + tm._dispatch_to_scheduler.reset_mock() + + tm._send_batch_request( + [SimpleNamespace(rid="job-1-seq-0"), SimpleNamespace(rid="job-1-seq-1")] + ) + + tm._dispatch_to_scheduler.assert_not_called() + self.assertEqual(tm.rid_to_state, {}) + + +class FakeReq: + def __init__(self, rid: str): + self.rid = rid + self.mamba_pool_idx = None + self.to_finish = None + + def finished(self) -> bool: + return False + + +def _make_scheduler(waiting_rids=(), running_rids=(), chunked_rid=None): + sched = SimpleNamespace() + sched.chunked_req = FakeReq(chunked_rid) if chunked_rid is not None else None + sched.waiting_queue = [FakeReq(rid) for rid in waiting_rids] + sched.enable_hicache_storage = False + sched.dllm_config = None # abort_request reads it since the dLLM rework (#27877) + sched.disaggregation_mode = DisaggregationMode.NULL + sched.grammar_manager = MagicMock() + sched.running_batch = SimpleNamespace(reqs=[FakeReq(rid) for rid in running_rids]) + sched.cur_batch = None + # abort_request collects in-flight reqs from (running_batch, last_batch) on + # the pp_size == 1 path since the v0.5.16 rebase. + sched.ps = SimpleNamespace(pp_size=1) + sched.last_batch = None + sched.ipc_channels = MagicMock() + return sched + + +class TestSchedulerAbortMatching(CustomTestCase): + """Scheduler-side matching semantics for AbortReq (see io_struct.AbortReq: + always ``rid.startswith``, so batch children ``f"{rid}_{i}"`` are covered).""" + + def test_prefix_abort_isolates_namespaces(self): + sched = _make_scheduler( + waiting_rids=["A::1", "A::2", "B::1"], + running_rids=["A::3", "B::2"], + ) + Scheduler.abort_request(sched, AbortReq(rid="A::", prefix=True)) + + self.assertEqual([req.rid for req in sched.waiting_queue], ["B::1"]) + running = {req.rid: req for req in sched.running_batch.reqs} + self.assertIsInstance(running["A::3"].to_finish, FINISH_ABORT) + self.assertIsNone(running["B::2"].to_finish) + # Every waiting-queue abort echoes back to the tokenizer for cleanup. + aborted_rids = { + call.args[0].rid + for call in sched.ipc_channels.send_to_tokenizer.send_output.call_args_list + } + self.assertEqual(aborted_rids, {"A::1", "A::2"}) + + def test_matching_is_prefix_based_even_without_prefix_flag(self): + # Pre-existing scheduler semantics: batch requests derive child rids as + # f"{rid}_{i}", so an exact-mode abort for the parent must cover them. + sched = _make_scheduler(waiting_rids=["job-1_0", "job-1_1", "job-2_0"]) + Scheduler.abort_request(sched, AbortReq(rid="job-1", prefix=False)) + + self.assertEqual([req.rid for req in sched.waiting_queue], ["job-2_0"]) + + def test_chunked_request_is_prefix_matched(self): + sched = _make_scheduler(chunked_rid="A::9") + Scheduler.abort_request(sched, AbortReq(rid="A::", prefix=True)) + + self.assertIs(sched._pending_chunked_abort_req, sched.chunked_req) + + def test_abort_all_covers_everything(self): + sched = _make_scheduler(waiting_rids=["A::1"], running_rids=["B::1"]) + Scheduler.abort_request(sched, AbortReq(abort_all=True)) + + self.assertEqual(sched.waiting_queue, []) + self.assertIsInstance(sched.running_batch.reqs[0].to_finish, FINISH_ABORT) + + +def _make_disagg_req(rid: str, pending_bootstrap: bool = False) -> FakeReq: + req = FakeReq(rid) + req.pending_bootstrap = pending_bootstrap + req.disagg_kv_sender = Mock() + return req + + +def _make_prefill_scheduler(waiting_rids=(), bootstrap_rids=(), inflight_rids=()): + sched = _make_scheduler() + sched.disaggregation_mode = DisaggregationMode.PREFILL + sched.waiting_queue = [ + _make_disagg_req(rid, pending_bootstrap=True) for rid in waiting_rids + ] + sched.req_to_metadata_buffer_idx_allocator = MagicMock() + sched.disagg_prefill_bootstrap_queue = SimpleNamespace( + queue=[_make_disagg_req(rid) for rid in bootstrap_rids] + ) + sched.disagg_prefill_inflight_queue = [ + _make_disagg_req(rid) for rid in inflight_rids + ] + return sched + + +def _make_decode_req(rid: str) -> SimpleNamespace: + return SimpleNamespace(req=FakeReq(rid), kv_receiver=Mock()) + + +def _make_retracted_req(rid: str) -> SimpleNamespace: + return SimpleNamespace(rid=rid, kv_cache_cpu=object()) + + +def _make_decode_scheduler( + waiting_rids=(), prealloc_rids=(), transfer_rids=(), retracted_rids=() +): + sched = _make_scheduler(waiting_rids=waiting_rids) + sched.disaggregation_mode = DisaggregationMode.DECODE + sched.tree_cache = MagicMock() + sched.disagg_decode_prealloc_queue = SimpleNamespace( + queue=[_make_decode_req(rid) for rid in prealloc_rids], + retracted_queue=[_make_retracted_req(rid) for rid in retracted_rids], + ) + sched.disagg_decode_transfer_queue = SimpleNamespace( + queue=[_make_decode_req(rid) for rid in transfer_rids] + ) + return sched + + +def _echoed_rids(sched) -> set: + return { + call.args[0].rid + for call in sched.ipc_channels.send_to_tokenizer.send_output.call_args_list + } + + +class TestSchedulerDisaggPrefillAbort(CustomTestCase): + """PREFILL-side disaggregation abort matching: the bootstrap and in-flight + queues hold requests the waiting queue no longer tracks, and the waiting + queue itself must release the metadata buffer slot and abort a + still-bootstrapping KV sender.""" + + def test_bootstrap_and_inflight_queues_prefix_matched(self): + sched = _make_prefill_scheduler( + bootstrap_rids=["A::1", "B::1"], inflight_rids=["A::2", "B::2"] + ) + Scheduler.abort_request(sched, AbortReq(rid="A::", prefix=True)) + + bootstrap = {r.rid: r for r in sched.disagg_prefill_bootstrap_queue.queue} + bootstrap["A::1"].disagg_kv_sender.abort.assert_called_once() + bootstrap["B::1"].disagg_kv_sender.abort.assert_not_called() + inflight = {r.rid: r for r in sched.disagg_prefill_inflight_queue} + inflight["A::2"].disagg_kv_sender.abort.assert_called_once() + inflight["B::2"].disagg_kv_sender.abort.assert_not_called() + + def test_waiting_queue_releases_metadata_buffer(self): + sched = _make_prefill_scheduler(waiting_rids=["A::1", "B::1"]) + with patch( + "sglang.srt.managers.scheduler.maybe_release_metadata_buffer" + ) as release: + Scheduler.abort_request(sched, AbortReq(rid="A::", prefix=True)) + + self.assertEqual([r.rid for r in sched.waiting_queue], ["B::1"]) + release.assert_called_once() + released_req, allocator = release.call_args.args + self.assertEqual(released_req.rid, "A::1") + self.assertIs(allocator, sched.req_to_metadata_buffer_idx_allocator) + self.assertEqual(_echoed_rids(sched), {"A::1"}) + + def test_waiting_queue_pending_bootstrap_gates_sender_abort(self): + pending = _make_disagg_req("A::1", pending_bootstrap=True) + bootstrapped = _make_disagg_req("A::2", pending_bootstrap=False) + sched = _make_prefill_scheduler() + sched.waiting_queue = [pending, bootstrapped] + with patch("sglang.srt.managers.scheduler.maybe_release_metadata_buffer"): + Scheduler.abort_request(sched, AbortReq(rid="A::", prefix=True)) + + pending.disagg_kv_sender.abort.assert_called_once() + bootstrapped.disagg_kv_sender.abort.assert_not_called() + + def test_abort_all_covers_prefill_queues(self): + sched = _make_prefill_scheduler(bootstrap_rids=["A::1"], inflight_rids=["B::1"]) + Scheduler.abort_request(sched, AbortReq(abort_all=True)) + + sched.disagg_prefill_bootstrap_queue.queue[ + 0 + ].disagg_kv_sender.abort.assert_called_once() + sched.disagg_prefill_inflight_queue[ + 0 + ].disagg_kv_sender.abort.assert_called_once() + + +class TestSchedulerDisaggDecodeAbort(CustomTestCase): + """DECODE-side disaggregation abort matching: prealloc/transfer queues + abort their KV receivers, the retracted queue frees CPU KV cache and + echoes the abort back to the tokenizer, and waiting-queue requests + release their preallocated KV cache.""" + + def test_prealloc_and_transfer_queues_prefix_matched(self): + sched = _make_decode_scheduler( + prealloc_rids=["A::1", "B::1"], transfer_rids=["A::2", "B::2"] + ) + Scheduler.abort_request(sched, AbortReq(rid="A::", prefix=True)) + + prealloc = {d.req.rid: d for d in sched.disagg_decode_prealloc_queue.queue} + prealloc["A::1"].kv_receiver.abort.assert_called_once() + prealloc["B::1"].kv_receiver.abort.assert_not_called() + transfer = {d.req.rid: d for d in sched.disagg_decode_transfer_queue.queue} + transfer["A::2"].kv_receiver.abort.assert_called_once() + transfer["B::2"].kv_receiver.abort.assert_not_called() + + def test_retracted_queue_frees_cpu_cache_and_echoes(self): + sched = _make_decode_scheduler(retracted_rids=["A::1", "B::1"]) + aborted = sched.disagg_decode_prealloc_queue.retracted_queue[0] + Scheduler.abort_request(sched, AbortReq(rid="A::", prefix=True)) + + self.assertEqual( + [d.rid for d in sched.disagg_decode_prealloc_queue.retracted_queue], + ["B::1"], + ) + self.assertFalse(hasattr(aborted, "kv_cache_cpu")) + self.assertEqual(_echoed_rids(sched), {"A::1"}) + + def test_waiting_queue_releases_kv_cache(self): + sched = _make_decode_scheduler(waiting_rids=["A::1", "B::1"]) + with patch("sglang.srt.managers.scheduler.release_kv_cache") as release: + Scheduler.abort_request(sched, AbortReq(rid="A::", prefix=True)) + + self.assertEqual([r.rid for r in sched.waiting_queue], ["B::1"]) + release.assert_called_once() + self.assertEqual(release.call_args.args[0].rid, "A::1") + + def test_exact_mode_is_still_prefix_matched_in_disagg_queues(self): + # Same load-bearing semantics as the non-disagg queues: batch children + # derive rids as f"{rid}_{i}", so exact-mode must cover them here too. + sched = _make_decode_scheduler(prealloc_rids=["job-1_0"]) + Scheduler.abort_request(sched, AbortReq(rid="job-1", prefix=False)) + + sched.disagg_decode_prealloc_queue.queue[ + 0 + ].kv_receiver.abort.assert_called_once() + + def test_abort_all_covers_decode_queues(self): + sched = _make_decode_scheduler( + prealloc_rids=["A::1"], transfer_rids=["B::1"], retracted_rids=["C::1"] + ) + Scheduler.abort_request(sched, AbortReq(abort_all=True)) + + sched.disagg_decode_prealloc_queue.queue[ + 0 + ].kv_receiver.abort.assert_called_once() + sched.disagg_decode_transfer_queue.queue[ + 0 + ].kv_receiver.abort.assert_called_once() + self.assertEqual(sched.disagg_decode_prealloc_queue.retracted_queue, []) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/test/registered/unit/managers/test_auxiliary_output.py b/test/registered/unit/managers/test_auxiliary_output.py index 2afa825d2506..fc895c12bb46 100644 --- a/test/registered/unit/managers/test_auxiliary_output.py +++ b/test/registered/unit/managers/test_auxiliary_output.py @@ -129,6 +129,7 @@ def test_sampling_mask_output_uses_generation_result_copy_path(): token_ids=torch.tensor([[3, 5]], dtype=torch.int32), lengths=torch.tensor([2], dtype=torch.int32), selected_logprobs=torch.tensor([-0.5]), + support_logprobs=torch.tensor([[-0.5, -1.0]]), statuses=torch.tensor([SamplingMaskStatus.OK], dtype=torch.int32), ) result = GenerationBatchResult( @@ -146,8 +147,10 @@ def test_sampling_mask_output_uses_generation_result_copy_path(): ) as copy_tensor: result.copy_to_cpu(return_logprob=False) - assert copy_tensor.call_count == 5 + assert copy_tensor.call_count == 6 assert sampling_output.token_ids.tolist() == [[3, 5]] + assert sampling_output.selected_logprobs.tolist() == [-0.5] + assert sampling_output.support_logprobs.tolist() == [[-0.5, -1.0]] assert sampling_output.lengths.tolist() == [2] assert sampling_output.statuses.tolist() == [SamplingMaskStatus.OK] assert result.copy_done.record_count == 1 @@ -158,6 +161,7 @@ def test_pipeline_sampling_mask_round_trip_without_logprobs(): token_ids=torch.tensor([[3, 5]], dtype=torch.int32), lengths=torch.tensor([2], dtype=torch.int32), selected_logprobs=torch.tensor([-0.5]), + support_logprobs=torch.tensor([[-0.5, -1.0]]), statuses=torch.tensor([SamplingMaskStatus.OK], dtype=torch.int32), ) result = GenerationBatchResult( @@ -170,7 +174,13 @@ def test_pipeline_sampling_mask_round_trip_without_logprobs(): SimpleNamespace(), result, SimpleNamespace(return_logprob=False) ) output, _, _ = get_logprob_from_pp_outputs(PPProxyTensors(payload)) - for name in ("token_ids", "lengths", "selected_logprobs", "statuses"): + for name in ( + "token_ids", + "lengths", + "selected_logprobs", + "support_logprobs", + "statuses", + ): torch.testing.assert_close( getattr(output.sampling_mask_output, name), getattr(sampling_output, name) ) diff --git a/test/registered/unit/managers/test_batch_result_processor_hidden_states.py b/test/registered/unit/managers/test_batch_result_processor_hidden_states.py index 363515e10343..103432afbd62 100644 --- a/test/registered/unit/managers/test_batch_result_processor_hidden_states.py +++ b/test/registered/unit/managers/test_batch_result_processor_hidden_states.py @@ -52,6 +52,61 @@ def _make_processor(case, server_mode: str = "full") -> SchedulerBatchResultProc class TestSamplingMaskMaterialization(CustomTestCase): + def test_selected_and_support_modes_share_one_batch(self): + output = LogitsProcessorOutput( + next_token_logits=None, + sampling_mask_output=SimpleNamespace( + token_ids=torch.tensor([[7, 8], [9, 10]], dtype=torch.int32), + lengths=torch.tensor([2, 2]), + selected_logprobs=torch.tensor([-0.5, -0.25]), + support_logprobs=torch.tensor([[-0.25, -1.5]]), + statuses=torch.tensor([SamplingMaskStatus.OK, SamplingMaskStatus.OK]), + ), + ) + SchedulerBatchResultProcessor.materialize_sampling_mask_output( + reqs=[ + SimpleNamespace( + return_sampling_mask=True, + sampling_logprobs_mode="selected", + ), + SimpleNamespace( + return_sampling_mask=True, + sampling_logprobs_mode="support", + ), + ], + output=output, + ) + self.assertEqual( + [row.tolist() for row in output.next_token_sampling_mask_idx], + [[7, 8], [9, 10]], + ) + self.assertEqual( + [row.tolist() for row in output.next_token_sampling_logprobs], + [[-0.5], [-0.25, -1.5]], + ) + + def test_selected_mode_does_not_require_support_tensor(self): + output = LogitsProcessorOutput( + next_token_logits=None, + sampling_mask_output=SimpleNamespace( + token_ids=torch.tensor([[7, 8]], dtype=torch.int32), + lengths=torch.tensor([2]), + selected_logprobs=torch.tensor([-0.5]), + support_logprobs=None, + statuses=torch.tensor([SamplingMaskStatus.OK]), + ), + ) + SchedulerBatchResultProcessor.materialize_sampling_mask_output( + reqs=[ + SimpleNamespace( + return_sampling_mask=True, + sampling_logprobs_mode="selected", + ) + ], + output=output, + ) + self.assertEqual(output.next_token_sampling_logprobs[0].tolist(), [-0.5]) + def test_packed_ids_are_copied_before_per_request_slicing(self): """Non-overlap capture must not perform one device copy per request.""" packed_ids = Mock() @@ -63,16 +118,35 @@ def test_packed_ids_are_copied_before_per_request_slicing(self): token_ids=packed_ids, lengths=torch.tensor([2, 1]), selected_logprobs=torch.tensor([-0.5, -0.25]), + support_logprobs=torch.tensor([[-0.5, -0.75, 0.0], [-0.25, 0.0, 0.0]]), statuses=torch.tensor([SamplingMaskStatus.OK, SamplingMaskStatus.OK]), ), ) SchedulerBatchResultProcessor.materialize_sampling_mask_output( - reqs=[SimpleNamespace(return_sampling_mask=x) for x in (True, False, True)], + reqs=[ + SimpleNamespace( + return_sampling_mask=x, + sampling_logprobs_mode="support", + ) + for x in (True, False, True) + ], output=output, ) packed_ids.cpu.assert_called_once_with() - self.assertEqual(output.next_token_sampling_mask_idx, [[7, 8], None, [9]]) - self.assertEqual(output.next_token_sampling_logprobs, [-0.5, None, -0.25]) + self.assertEqual( + [ + None if row is None else row.tolist() + for row in output.next_token_sampling_mask_idx + ], + [[7, 8], None, [9]], + ) + self.assertEqual( + [ + None if row is None else row.tolist() + for row in output.next_token_sampling_logprobs + ], + [[-0.5, -0.75], None, [-0.25]], + ) self.assertIsNone(output.sampling_mask_output) diff --git a/test/registered/unit/managers/test_flat_raw_top_logprobs.py b/test/registered/unit/managers/test_flat_raw_top_logprobs.py index f7e4e09e64a6..a9b89239bf0e 100644 --- a/test/registered/unit/managers/test_flat_raw_top_logprobs.py +++ b/test/registered/unit/managers/test_flat_raw_top_logprobs.py @@ -601,7 +601,6 @@ def _make_batch_token_id_output(**overrides) -> BatchTokenIDOutput: output_token_ids_logprobs_idx=[[], []], output_token_entropy_val=None, output_token_sampling_mask=None, - output_token_sampling_logprobs=None, output_hidden_states=None, routed_experts=None, indexer_topk=None, diff --git a/test/registered/unit/managers/test_io_struct.py b/test/registered/unit/managers/test_io_struct.py index c35ba141acfd..9d13c113feeb 100644 --- a/test/registered/unit/managers/test_io_struct.py +++ b/test/registered/unit/managers/test_io_struct.py @@ -407,6 +407,24 @@ def setUp(self): rid=["id1", "id2"], ) + def test_sampling_logprobs_mode_normalization(self): + default_req = copy.deepcopy(self.base_req) + default_req.normalize_batch_and_arguments() + self.assertEqual(default_req.sampling_logprobs_mode, [None, None]) + + req = copy.deepcopy(self.base_req) + req.return_sampling_mask = [True, True] + req.sampling_logprobs_mode = ["selected", "support"] + + req.normalize_batch_and_arguments() + + self.assertEqual( + req.sampling_logprobs_mode, + ["selected", "support"], + ) + self.assertEqual(req[0].sampling_logprobs_mode, "selected") + self.assertEqual(req[1].sampling_logprobs_mode, "support") + def test_single_image_to_list_of_lists(self): """Test that a single image is converted to a list of single-image lists.""" req = copy.deepcopy(self.base_req) diff --git a/test/registered/unit/managers/test_msgpack_ipc_roundtrip.py b/test/registered/unit/managers/test_msgpack_ipc_roundtrip.py index a206c27fbc9d..3ea4b2d76d40 100644 --- a/test/registered/unit/managers/test_msgpack_ipc_roundtrip.py +++ b/test/registered/unit/managers/test_msgpack_ipc_roundtrip.py @@ -24,7 +24,6 @@ ExpertWeightPointer, GetInternalStateReqOutput, GetWeightsByNameReqOutput, - LoadLoRAAdapterFromTensorsReqInput, ParallelismInfo, RpcReqInput, SetInternalStateReq, @@ -65,9 +64,17 @@ def _contains_dataclass(obj) -> bool: return False -def _parallelism_info() -> ParallelismInfo: +def _parallelism_info(role: str = "target") -> ParallelismInfo: return ParallelismInfo( - tp_rank=0, tp_size=2, dp_rank=0, dp_size=1, pp_rank=0, pp_size=1, rank=0, size=2 + role=role, + tp_rank=0, + tp_size=2, + dp_rank=0, + dp_size=1, + pp_rank=0, + pp_size=1, + rank=0, + size=2, ) @@ -75,7 +82,8 @@ def _checksum_info(tag: str) -> ChecksumInfo: return ChecksumInfo( checksums={f"model.layers.{tag}": "deadbeef"}, per_gpu_checksum="cafef00d", - parallelism_info=_parallelism_info(), + # One entry per role: the target model plus each speculative draft runner. + parallelism_info=[_parallelism_info("target"), _parallelism_info("draft")], ) @@ -125,12 +133,6 @@ def _checksum_info(tag: str) -> ChecksumInfo: parameters={"flag": True, "n": 1, "ratio": 2.0, "name": "x", "opt": None}, ), "RpcReqInput/none": RpcReqInput(method="collective_rpc", parameters=None), - "LoadLoRAAdapterFromTensorsReqInput": LoadLoRAAdapterFromTensorsReqInput( - lora_name="adapter", - config_dict={"r": 8, "lora_alpha": 16, "target_modules": ["q_proj", "v_proj"]}, - serialized_named_tensors=[b"tp0-bytes", b"tp1-bytes"], - added_tokens_config={"": 32000}, - ), "DumperControlReqInput": DumperControlReqInput(method="start", body={"k": "v"}), "DumperControlReqOutput": DumperControlReqOutput( success=True, response=[{"worker": 0, "ok": True}] @@ -194,29 +196,46 @@ def test_check_weights_multi_rank_payload(self): self.assertEqual(len(decoded.payload), 2) as_dict = msgspec_to_builtins(decoded.payload[0]) self.assertEqual(as_dict["per_gpu_checksum"], "cafef00d") - self.assertIn("tp_rank", as_dict["parallelism_info"]) + self.assertEqual( + [pi["role"] for pi in as_dict["parallelism_info"]], ["target", "draft"] + ) + self.assertIn("tp_rank", as_dict["parallelism_info"][0]) def test_check_weights_producer_conversion(self): - # Mirrors weight_updater.check_weights: WeightChecker returns - # ChecksumInfo.model_dump() (a dict), converted to the msgspec struct via - # msgspec.convert, and the result round-trips as the payload. - pydantic_checksum = PydanticChecksumInfo( - checksums={"model.layers.0": "deadbeef"}, - per_gpu_checksum="cafef00d", - parallelism_info=PydanticParallelismInfo( - tp_rank=0, - tp_size=2, - dp_rank=0, - dp_size=1, - pp_rank=0, - pp_size=1, - rank=0, - size=2, - ), + # Mirrors weight_updater.check_weights end to end: each runner's + # WeightChecker returns a single-role ChecksumInfo.model_dump(), + # _merge_checksum_payloads folds the roles into one per-role payload, and + # that is what msgspec.convert has to accept. + from sglang.srt.managers.scheduler_components.weight_updater import ( + _merge_checksum_payloads, + ) + + def _dump(role_tag): + return PydanticChecksumInfo( + checksums={f"model.layers.0{role_tag}": "deadbeef"}, + per_gpu_checksum="cafef00d", + parallelism_info=PydanticParallelismInfo( + tp_rank=0, + tp_size=2, + dp_rank=0, + dp_size=1, + pp_rank=0, + pp_size=1, + rank=0, + size=2, + ), + ).model_dump() + + merged = _merge_checksum_payloads([("", _dump("")), ("draft", _dump(""))]) + converted = msgspec.convert(merged, ChecksumInfo) + # Draft keys are role-prefixed so they never collide with the target's. + self.assertEqual( + sorted(converted.checksums), ["draft.model.layers.0", "model.layers.0"] + ) + self.assertEqual( + [pi.role for pi in converted.parallelism_info], ["target", "draft"] ) - converted = msgspec.convert(pydantic_checksum.model_dump(), ChecksumInfo) - self.assertEqual(converted.per_gpu_checksum, "cafef00d") - self.assertEqual(converted.parallelism_info.tp_rank, 0) + self.assertEqual(converted.parallelism_info[0].tp_rank, 0) output = CheckWeightsReqOutput(success=True, message="ok", payload=[converted]) self.assertEqual(_round_trip(output), output) diff --git a/test/registered/unit/managers/test_multi_tokenizer_mixin.py b/test/registered/unit/managers/test_multi_tokenizer_mixin.py index 7468d1fc5042..02a05ca209e5 100644 --- a/test/registered/unit/managers/test_multi_tokenizer_mixin.py +++ b/test/registered/unit/managers/test_multi_tokenizer_mixin.py @@ -1,5 +1,8 @@ import unittest +import numpy as np + +from sglang.srt.sampling.sampling_mask import SamplingMaskChunk from sglang.srt.utils.weight_versions import WeightVersionSpan from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import maybe_stub_sgl_kernel @@ -69,8 +72,18 @@ def _make_batch_str_output() -> BatchStrOutput: output_token_ids_logprobs_val=[[], []], output_token_ids_logprobs_idx=[[], []], output_token_entropy_val=[0.0, 0.0], - output_token_sampling_mask=[[], []], - output_token_sampling_logprobs=[[], []], + output_token_sampling_mask=[ + SamplingMaskChunk( + lengths=np.array([2], np.int32), + token_ids=np.array([7, 8], np.int32), + logprobs=np.array([-0.5, -1.0], np.float32), + ), + SamplingMaskChunk( + lengths=np.array([1], np.int32), + token_ids=np.array([9], np.int32), + logprobs=np.array([0.0], np.float32), + ), + ], output_hidden_states=[None, None], routed_experts=[None, None], indexer_topk=[None, None], @@ -118,6 +131,16 @@ def test_batch_str_output_keeps_weight_versions_nested_per_request(self): [[WeightVersionSpan(version="v2", start=0, end=2)]], ) + def test_batch_str_output_keeps_sampling_distribution_aligned(self): + output = _make_batch_str_output() + + single_output = _handle_output_by_index(output, 0) + + (chunk,) = single_output.output_token_sampling_mask + self.assertEqual( + chunk.to_lists(support_logprobs=True), ([[7, 8]], [[-0.5, -1.0]]) + ) + def test_batch_str_output_without_weight_versions_stays_none(self): """An output from an older server without the field splits into None.""" output = _make_batch_str_output() diff --git a/test/registered/unit/managers/test_output_streamer_customized_info.py b/test/registered/unit/managers/test_output_streamer_customized_info.py index cd103f7f799f..1139dcf7de35 100644 --- a/test/registered/unit/managers/test_output_streamer_customized_info.py +++ b/test/registered/unit/managers/test_output_streamer_customized_info.py @@ -1,13 +1,29 @@ +import dataclasses +import pickle import unittest +from array import array from types import SimpleNamespace from unittest.mock import Mock, patch +import numpy as np +import torch + from sglang.srt.disaggregation.utils import DisaggregationMode -from sglang.srt.managers.io_struct import unwrap_from_pickle +from sglang.srt.layers.logits_processor import LogitsProcessorOutput, SamplingMaskStatus +from sglang.srt.managers import io_struct +from sglang.srt.managers.io_struct import ( + msgpack_decode, + msgpack_encode, + unwrap_from_pickle, +) +from sglang.srt.managers.scheduler_components.batch_result_processor import ( + SchedulerBatchResultProcessor, +) from sglang.srt.managers.scheduler_components.output_streamer import ( SchedulerOutputStreamer, _GenerationStreamAccumulator, ) +from sglang.srt.sampling.sampling_mask import SamplingMaskRows from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.utils.weight_versions import ( WeightVersionSpan, @@ -26,6 +42,7 @@ def __init__( customized_info=None, *, finished=False, + sampling_mask_rows=None, ): self.rid = rid self.http_worker_ipc = None @@ -57,7 +74,8 @@ def __init__( self.return_hidden_states = False self.return_routed_experts = False self.return_indexer_topk = False - self.return_sampling_mask = False + self.return_sampling_mask = sampling_mask_rows is not None + self.sampling_mask_rows = sampling_mask_rows self.mm_image_tokens = 0 self.mm_audio_tokens = 0 self.mm_video_tokens = 0 @@ -75,12 +93,13 @@ def check_match_stop_str_prefix(self): return False -def _accumulator(current_weight_version="default"): +def _accumulator(current_weight_version="default", return_sampling_mask=False): return _GenerationStreamAccumulator( return_logprob=False, return_hidden_states=False, return_routed_experts=False, return_indexer_topk=False, + return_sampling_mask=return_sampling_mask, spec_algorithm=SpeculativeAlgorithm.NONE, disaggregation_mode=DisaggregationMode.NULL, default_stream_interval=1, @@ -374,5 +393,154 @@ def test_payload_omits_spans_while_all_requests_stream(self): self.assertIsNone(payload.weight_versions) +_IPC_ROUND_TRIPS = { + "pickle": lambda payload: pickle.loads( + pickle.dumps(payload, protocol=pickle.HIGHEST_PROTOCOL) + ), + "msgpack": lambda payload: msgpack_decode(msgpack_encode(payload)), +} + + +class TestOutputStreamerSamplingMasks(unittest.TestCase): + def test_rows_stream_once_and_expand_to_response_lists(self): + """Queued rows cross either IPC codec once, in batch order, as per-token lists; + rows queued later never overwrite a chunk that is still referenced.""" + for codec, round_trip in _IPC_ROUND_TRIPS.items(): + with ( + self.subTest(codec=codec), + patch.object(io_struct, "_USE_PICKLE_IPC", codec == "pickle"), + ): + selected_rows = SamplingMaskRows() + selected_rows.append( + np.array([5, 3], np.int32), np.array([-0.5], np.float32) + ) + selected_rows.append( + np.array([7], np.int32), np.array([0.0], np.float32) + ) + support_rows = SamplingMaskRows() + support_rows.append( + np.array([9, 8], np.int32), np.array([-0.25, -1.5], np.float32) + ) + support_rows.append( + np.array([6], np.int32), np.array([0.0], np.float32) + ) + selected = _FakeReq( + "selected", array("q", [5, 7]), sampling_mask_rows=selected_rows + ) + support = _FakeReq( + "support", array("q", [9, 6]), sampling_mask_rows=support_rows + ) + accumulator = _accumulator(return_sampling_mask=True) + for req in (selected, _FakeReq("plain", array("q", [1])), support): + accumulator.accept(req=req) + payload = accumulator.to_payload(dp_rank=0, is_idle_batch=False) + masks = round_trip(payload).output_token_sampling_mask + self.assertEqual( + masks[0].to_lists(support_logprobs=False), + ([[5, 3], [7]], [-0.5, 0.0]), + ) + self.assertIsNone(masks[1]) + self.assertEqual( + masks[2].to_lists(support_logprobs=True), + ([[9, 8], [6]], [[-0.25, -1.5], [0.0]]), + ) + + selected.output_ids.append(4) + selected.sampling_mask_rows.append( + np.array([4], np.int32), np.array([-2.0], np.float32) + ) + self.assertEqual( + payload.output_token_sampling_mask[0].to_lists( + support_logprobs=False + ), + ([[5, 3], [7]], [-0.5, 0.0]), + ) + accumulator = _accumulator(return_sampling_mask=True) + accumulator.accept(req=selected) + masks = round_trip( + accumulator.to_payload(dp_rank=0, is_idle_batch=False) + ).output_token_sampling_mask + self.assertEqual( + masks[0].to_lists(support_logprobs=False), ([[4]], [-2.0]) + ) + + def test_clients_receive_the_sampler_rows_unchanged(self): + """Across decode steps, stream emissions and either IPC codec, each request gets + the sampler's rows exactly as the tensor-to-list conversion defines them.""" + processor = SchedulerBatchResultProcessor( + **{f.name: None for f in dataclasses.fields(SchedulerBatchResultProcessor)} + ) + modes = ["selected", "support", None, "support", "selected"] + mask_reqs = [i for i, mode in enumerate(modes) if mode is not None] + support_reqs = [i for i in mask_reqs if modes[i] == "support"] + for codec, round_trip in _IPC_ROUND_TRIPS.items(): + with ( + self.subTest(codec=codec), + patch.object(io_struct, "_USE_PICKLE_IPC", codec == "pickle"), + ): + generator = torch.Generator().manual_seed(0) + reqs = [] + for i, mode in enumerate(modes): + rows = None if mode is None else SamplingMaskRows() + req = _FakeReq(f"r{i}", array("q"), sampling_mask_rows=rows) + req.sampling_logprobs_mode = mode + reqs.append(req) + expected = [([], []) for _ in modes] + received = [([], []) for _ in modes] + for step in range(12): + width = (1, 7, 64)[step % 3] + lengths = torch.randint( + 1, width + 1, (len(mask_reqs),), generator=generator + ) + token_ids = torch.randint( + 0, 1 << 20, (len(mask_reqs), width), generator=generator + ).int() + selected = torch.randn(len(mask_reqs), generator=generator) + support = torch.randn(len(support_reqs), width, generator=generator) + output = LogitsProcessorOutput( + next_token_logits=None, + sampling_mask_output=SimpleNamespace( + token_ids=token_ids, + lengths=lengths, + selected_logprobs=selected, + support_logprobs=support, + statuses=torch.full( + (len(mask_reqs),), SamplingMaskStatus.OK + ), + ), + ) + SchedulerBatchResultProcessor.materialize_sampling_mask_output( + reqs, output + ) + for row, i in enumerate(mask_reqs): + processor.add_sampling_mask_return_values(i, reqs[i], output) + length = int(lengths[row]) + expected[i][0].append(token_ids[row, :length].tolist()) + if modes[i] == "support": + support_row = support_reqs.index(i) + expected[i][1].append( + support[support_row, :length].tolist() + ) + else: + expected[i][1].append(float(selected[row])) + if step in (0, 1, 5, 6, 11): + accumulator = _accumulator(return_sampling_mask=True) + for req in reqs: + accumulator.accept(req=req) + chunks = round_trip( + accumulator.to_payload(dp_rank=0, is_idle_batch=False) + ).output_token_sampling_mask + for i, chunk in enumerate(chunks): + if modes[i] is None: + self.assertIsNone(chunk) + continue + masks, logprobs = chunk.to_lists( + support_logprobs=modes[i] == "support" + ) + received[i][0].extend(masks) + received[i][1].extend(logprobs) + self.assertEqual(received, expected) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/managers/test_release_memory_flush_order.py b/test/registered/unit/managers/test_release_memory_flush_order.py new file mode 100644 index 000000000000..71bc719d8ec1 --- /dev/null +++ b/test/registered/unit/managers/test_release_memory_flush_order.py @@ -0,0 +1,96 @@ +import unittest +from types import SimpleNamespace +from unittest.mock import Mock, call, patch + +from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE +from sglang.srt.disaggregation.utils import DisaggregationMode +from sglang.srt.managers.io_struct import ( + ReleaseMemoryOccupationReqInput, + ResumeMemoryOccupationReqInput, +) +from sglang.srt.managers.scheduler import Scheduler +from sglang.srt.managers.scheduler_components.weight_updater import ( + SchedulerWeightUpdaterManager, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=11, suite="base-a-test-cpu") + + +class TestReleaseMemoryFlushOrder(CustomTestCase): + """Releasing the KV region must flush the pools exactly once, before the region is tagged as offloaded.""" + + def _harness(self): + scheduler = SimpleNamespace( + _engine_paused=False, + is_fully_idle=lambda ignore_waiting=False: True, + cur_batch_for_debug=None, + last_batch=None, + tree_cache=Mock(), + req_to_token_pool=Mock(), + token_to_kv_pool_allocator=Mock(), + grammar_manager=Mock(), + metrics_reporter=SimpleNamespace( + reset_metrics=lambda: None, is_stats_logging_rank=False + ), + draft_worker=None, + waiting_queue=[], + running_batch=SimpleNamespace(reqs=[]), + disaggregation_mode=DisaggregationMode.NULL, + ) + manager = SchedulerWeightUpdaterManager( + tp_worker=None, + draft_worker=None, + tp_cpu_group=None, + memory_saver_adapter=Mock(), + flush_cache=lambda empty_cache=False: Scheduler.flush_cache( + scheduler, empty_cache=False + ), + is_fully_idle=scheduler.is_fully_idle, + scheduler=scheduler, + ) + scheduler.weight_updater = manager + return scheduler, manager + + def test_release_flushes_once_before_the_pause_and_later_flushes_are_skipped(self): + scheduler, manager = self._harness() + calls = Mock() + manager.flush_cache = Mock(wraps=manager.flush_cache) + calls.attach_mock(manager.flush_cache, "flush") + calls.attach_mock(manager.memory_saver_adapter.pause, "pause") + with patch( + "sglang.srt.managers.scheduler_components.weight_updater.torch.get_device_module", + return_value=SimpleNamespace(synchronize=calls.synchronize), + ): + manager.release_memory_occupation( + ReleaseMemoryOccupationReqInput(tags=[GPU_MEMORY_TYPE_KV_CACHE]) + ) + + self.assertEqual( + calls.mock_calls, + [ + call.flush(), + call.synchronize(), + call.pause(GPU_MEMORY_TYPE_KV_CACHE), + call.synchronize(), + ], + ) + self.assertEqual(scheduler.tree_cache.reset.call_count, 1) + self.assertEqual(manager.memory_saver_adapter.pause.call_count, 1) + self.assertIn(GPU_MEMORY_TYPE_KV_CACHE, manager.offload_tags) + + # the trainer's /flush_cache while the KV region is unmapped must not touch the pools + self.assertTrue(Scheduler.flush_cache(scheduler, empty_cache=False)) + self.assertEqual(scheduler.tree_cache.reset.call_count, 1) + + manager.resume_memory_occupation( + ResumeMemoryOccupationReqInput(tags=[GPU_MEMORY_TYPE_KV_CACHE]) + ) + self.assertNotIn(GPU_MEMORY_TYPE_KV_CACHE, manager.offload_tags) + self.assertTrue(Scheduler.flush_cache(scheduler, empty_cache=False)) + self.assertEqual(scheduler.tree_cache.reset.call_count, 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/managers/test_sampling_mask_validation.py b/test/registered/unit/managers/test_sampling_mask_validation.py index a41dfe548595..963bd94d2327 100644 --- a/test/registered/unit/managers/test_sampling_mask_validation.py +++ b/test/registered/unit/managers/test_sampling_mask_validation.py @@ -29,12 +29,18 @@ def setUp(self): self.manager.validate_total_tokens = False self.manager.is_generation = True - def _validate(self, processor, return_sampling_mask=True): + def _validate( + self, + processor, + return_sampling_mask=True, + sampling_logprobs_mode=None, + ): req = GenerateReqInput( input_ids=[1, 2, 3], sampling_params={"top_k": 10}, custom_logit_processor=processor, return_sampling_mask=return_sampling_mask, + sampling_logprobs_mode=sampling_logprobs_mode, ) self.manager._validate_one_request(req, req.input_ids) @@ -78,6 +84,23 @@ def test_other_processors_still_work_without_sampling_masks(self): Qwen3ThinkingBudgetLogitProcessor.to_str(), return_sampling_mask=False ) + def test_sampling_logprobs_mode_requires_sampling_mask(self): + self._validate(None, return_sampling_mask=False) + for mode in ("selected", "support"): + with self.subTest(mode=mode): + with self.assertRaisesRegex( + ValueError, "can only be set when return_sampling_mask=true" + ): + self._validate( + None, + return_sampling_mask=False, + sampling_logprobs_mode=mode, + ) + + for mode in (None, "selected", "support"): + with self.subTest(mode=mode): + self._validate(None, sampling_logprobs_mode=mode) + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/managers/test_scheduler_flush_cache_after_retract.py b/test/registered/unit/managers/test_scheduler_flush_cache_after_retract.py new file mode 100644 index 000000000000..ff3f771eb9da --- /dev/null +++ b/test/registered/unit/managers/test_scheduler_flush_cache_after_retract.py @@ -0,0 +1,109 @@ +import unittest +from collections import deque +from unittest.mock import MagicMock + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + +from sglang.srt.disaggregation.utils import DisaggregationMode +from sglang.srt.managers.scheduler import Scheduler + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") +register_cpu_ci(est_time=6, suite="base-c-test-cpu") + + +def _make_req(is_retracted: bool) -> MagicMock: + req = MagicMock() + req.is_retracted = is_retracted + return req + + +class TestSchedulerFlushCacheAfterRetract(unittest.TestCase): + """Regression coverage for flush_cache() no-oping while the engine is paused. + + With generation paused and running_batch drained, nothing in waiting_queue + holds KV/pool state: neither requests pause_generation(mode="retract") just + moved there, nor requests that were never scheduled at all. The latter are + always present under high concurrency (in-flight targets exceed engine + parallelism), where gating the flush on them starves weight updates until + the caller's retry budget times out (reproduced in RL training at 512 + in-flight requests). See RETRACT_MODE_FLUSH_CACHE_FIX.md in miles.""" + + def _new_scheduler(self) -> Scheduler: + scheduler = Scheduler.__new__(Scheduler) + scheduler._engine_paused = True + scheduler.enable_overlap = False + scheduler.last_batch = None + scheduler.cur_batch = None + scheduler.chunked_req = None + scheduler.disaggregation_mode = DisaggregationMode.NULL + scheduler.running_batch = MagicMock() + scheduler.running_batch.is_empty.return_value = True + scheduler.result_queue = deque() + scheduler.waiting_queue = [] + scheduler.dllm_manager = MagicMock() + scheduler.dllm_manager.any_staging_reqs.return_value = False + scheduler._pp_microbatches_drained = MagicMock(return_value=True) + scheduler.grammar_manager = MagicMock() + scheduler.grammar_manager.grammar_queue = [] + scheduler.enable_hierarchical_cache = False + scheduler.enable_hisparse = False + scheduler.tree_cache = MagicMock() + scheduler.req_to_token_pool = MagicMock() + scheduler.token_to_kv_pool_allocator = MagicMock() + scheduler.draft_worker = None + scheduler.metrics_reporter = MagicMock() + return scheduler + + def test_flush_cache_succeeds_when_only_retracted_reqs_are_waiting(self): + scheduler = self._new_scheduler() + scheduler.waiting_queue = [_make_req(is_retracted=True)] + + self.assertTrue(scheduler.flush_cache()) + scheduler.tree_cache.reset.assert_called_once() + scheduler.req_to_token_pool.clear.assert_called_once() + scheduler.token_to_kv_pool_allocator.clear.assert_called_once() + + def test_flush_cache_succeeds_with_never_scheduled_reqs_while_paused(self): + """A request that was queued but never scheduled holds no KV/pool state + either; a paused-engine flush must not starve on it.""" + scheduler = self._new_scheduler() + scheduler.waiting_queue = [_make_req(is_retracted=False)] + + self.assertTrue(scheduler.flush_cache()) + scheduler.tree_cache.reset.assert_called_once() + + def test_flush_cache_succeeds_on_mixed_queue_while_paused(self): + scheduler = self._new_scheduler() + scheduler.waiting_queue = [ + _make_req(is_retracted=True), + _make_req(is_retracted=False), + ] + + self.assertTrue(scheduler.flush_cache()) + scheduler.tree_cache.reset.assert_called_once() + + def test_is_fully_idle_default_ignores_nothing(self): + """Regression guard: every other is_fully_idle() caller passes no args, + so a non-empty waiting_queue must still report not-idle by default.""" + scheduler = self._new_scheduler() + scheduler.waiting_queue = [_make_req(is_retracted=True)] + + self.assertFalse(scheduler.is_fully_idle()) + self.assertTrue(scheduler.is_fully_idle(ignore_waiting=True)) + + def test_flush_cache_does_not_ignore_waiting_reqs_while_unpaused(self): + """ignore_waiting is gated on _engine_paused inside flush_cache — a + flush issued without pausing first must not be silently permissive.""" + scheduler = self._new_scheduler() + scheduler._engine_paused = False + scheduler.waiting_queue = [_make_req(is_retracted=True)] + + self.assertFalse(scheduler.flush_cache()) + scheduler.tree_cache.reset.assert_not_called() + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/managers/test_scheduler_pause_generation.py b/test/registered/unit/managers/test_scheduler_pause_generation.py index 716c59e57859..f0551d6ba6af 100644 --- a/test/registered/unit/managers/test_scheduler_pause_generation.py +++ b/test/registered/unit/managers/test_scheduler_pause_generation.py @@ -44,6 +44,7 @@ def _new_scheduler(self) -> Scheduler: scheduler.last_batch = None scheduler.cur_batch_for_debug = None scheduler.chunked_req = None + scheduler.disaggregation_mode = DisaggregationMode.NULL scheduler.running_batch = MagicMock() scheduler.running_batch.reqs = [] scheduler.running_batch.is_empty.return_value = True diff --git a/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py b/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py index 464fa5024126..f1ebea5e6902 100644 --- a/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py +++ b/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py @@ -201,6 +201,8 @@ def _make_batch_str_output(rid: str, finished_reason=None) -> BatchStrOutput: kwargs[f.name] = [[]] elif f.name in _PER_REQUEST_OPTIONAL_FIELDS: kwargs[f.name] = [None] + elif f.name == "output_token_sampling_mask": + kwargs[f.name] = None # Fields with class defaults — skip, let the default be used elif ( f.default is not msgspec.NODEFAULT diff --git a/test/registered/unit/mem_cache/test_session_token_share_unit.py b/test/registered/unit/mem_cache/test_session_token_share_unit.py index cf12e063c689..b762f5026a2e 100644 --- a/test/registered/unit/mem_cache/test_session_token_share_unit.py +++ b/test/registered/unit/mem_cache/test_session_token_share_unit.py @@ -39,6 +39,7 @@ def _recv(rid, input_ids, max_new_tokens=8): top_logprobs_num=0, token_ids_logprob=None, return_sampling_mask=False, + sampling_logprobs_mode="selected", require_reasoning=False, return_hidden_states=False, return_routed_experts=False, diff --git a/test/registered/unit/model_executor/test_mlp_sync_pad_unpad.py b/test/registered/unit/model_executor/test_mlp_sync_pad_unpad.py index 376982a86b83..a7c9565e19fd 100644 --- a/test/registered/unit/model_executor/test_mlp_sync_pad_unpad.py +++ b/test/registered/unit/model_executor/test_mlp_sync_pad_unpad.py @@ -15,6 +15,7 @@ import torch +from sglang.srt.layers.logits_processor import LogitsMetadata, LogitsProcessor from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.model_executor.runner.decode_cuda_graph_runner import ( DecodeCudaGraphRunner, @@ -40,6 +41,31 @@ def _logits_output(num_rows: int) -> SimpleNamespace: class TestMlpSyncPadUnpad(CustomTestCase): + def test_idle_rank_does_not_index_dummy_last_token(self): + # MLP-sync turns an idle rank into a dummy zero-token EXTEND batch. + empty = torch.empty(0, dtype=torch.int64) + batch = ForwardBatch( + forward_mode=ForwardMode.EXTEND, + batch_size=1, + input_ids=empty, + req_pool_indices=torch.tensor([0]), + seq_lens=torch.tensor([0]), + out_cache_loc=empty, + seq_lens_sum=0, + positions=empty, + extend_seq_lens=torch.tensor([0]), + extend_seq_lens_cpu=[0], + _original_forward_mode=ForwardMode.IDLE, + _original_batch_size=0, + ) + hidden = torch.empty(0, 4) + pruned, *_ = LogitsProcessor._get_pruned_states( + None, hidden, None, None, LogitsMetadata.from_forward_batch(batch) + ) + self.assertEqual(pruned.shape, (0, 4)) + # Attention and MLP execution still use the padded mode. + self.assertEqual(batch.forward_mode, ForwardMode.EXTEND) + def test_init_mlp_sync_metadata_scales_speculative_request_width(self): spec_info = SimpleNamespace( num_tokens_per_req=4, diff --git a/test/registered/unit/rl/test_distributed_weight_update_spec_worker.py b/test/registered/unit/rl/test_distributed_weight_update_spec_worker.py new file mode 100644 index 000000000000..e16c0330ce4a --- /dev/null +++ b/test/registered/unit/rl/test_distributed_weight_update_spec_worker.py @@ -0,0 +1,227 @@ +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import pytest + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +from sglang.srt.managers.io_struct import ( + BeginWeightUpdateReqInput, + EndWeightUpdateReqInput, + UpdateWeightsFromDistributedReqInput, +) +from sglang.srt.managers.scheduler_components.weight_updater import ( + SchedulerWeightUpdaterManager, +) + + +def _distributed_req(selector="all"): + return UpdateWeightsFromDistributedReqInput( + names=["model.layers.0.weight"], + dtypes=["float32"], + shapes=[[1]], + group_name="weight_update_group", + flush_cache=False, + selector=selector, + ) + + +def _manager(tp_worker, draft_worker): + # metrics_collector defaults to None, so _observe_weight_load is a no-op; the + # reqs below set flush_cache=False so flush_cache is never called either. + manager = SchedulerWeightUpdaterManager( + tp_worker=tp_worker, + draft_worker=draft_worker, + tp_cpu_group=object(), + memory_saver_adapter=Mock(), + flush_cache=Mock(return_value=True), + is_fully_idle=Mock(return_value=True), + ) + # update_weights_from_* assert an open begin_weight_update session. + manager._weight_update_in_progress = True + return manager + + +def test_scheduler_distributed_update_receives_once_on_target_loads_into_each(): + # Default selector ("all"): only the target (main model) owns the update group, + # so it receives the broadcast once; that single weights object is then loaded + # into every selected runner — receive once on the target, load into each. + weights = object() + target_runner = Mock() + target_runner.weight_updater.receive_weights_from_distributed.return_value = weights + draft_runner = Mock() + manager = _manager( + tp_worker=SimpleNamespace( + model_runner=target_runner, + iter_runners=lambda: [("", target_runner)], + ), + draft_worker=SimpleNamespace(iter_runners=lambda: [("draft", draft_runner)]), + ) + + output = manager.update_weights_from_distributed(_distributed_req()) + + assert output.success is True + target_runner.weight_updater.receive_weights_from_distributed.assert_called_once_with( + ["model.layers.0.weight"], + ["float32"], + [[1]], + "weight_update_group", + None, + ) + # The single received weights object is loaded into every selected runner. + target_runner.weight_updater.load_weights.assert_called_once_with(weights) + draft_runner.weight_updater.load_weights.assert_called_once_with(weights) + + +def test_scheduler_distributed_update_target_only_selector_skips_draft(): + # selector="target": the target still receives once, but the draft worker is + # never enumerated and no draft runner is loaded. + weights = object() + target_runner = Mock() + target_runner.weight_updater.receive_weights_from_distributed.return_value = weights + draft_worker = Mock() + manager = _manager( + tp_worker=SimpleNamespace( + model_runner=target_runner, + iter_runners=lambda: [("", target_runner)], + ), + draft_worker=draft_worker, + ) + + output = manager.update_weights_from_distributed( + _distributed_req(selector="target") + ) + + assert output.success is True + target_runner.weight_updater.receive_weights_from_distributed.assert_called_once() + target_runner.weight_updater.load_weights.assert_called_once_with(weights) + draft_worker.iter_runners.assert_not_called() + + +def _session_manager(target_runner, draft_runner): + return _manager( + tp_worker=SimpleNamespace(iter_runners=lambda: [("", target_runner)]), + draft_worker=SimpleNamespace(iter_runners=lambda: [("draft", draft_runner)]), + ) + + +def test_begin_weight_update_restores_target_and_draft(): + # The session begins on every runner (target + draft): the draft model is + # restored to a loadable state identically to the target. + target_runner = Mock() + draft_runner = Mock() + manager = _session_manager(target_runner, draft_runner) + manager._weight_update_in_progress = False + + with patch("torch.distributed.barrier"): + output = manager.begin_weight_update(BeginWeightUpdateReqInput()) + + assert output.success is True + target_runner.begin_weight_update.assert_called_once_with() + draft_runner.begin_weight_update.assert_called_once_with() + assert manager._weight_update_in_progress is True + assert manager._weight_update_loaded is False + + +def test_end_weight_update_runs_post_load_on_both_when_load_was_bypassed(): + # No load_weights happened this session (e.g. P2P/RDMA), so end runs + # post_load_weights then quant finalize on BOTH target and draft. + target_runner = Mock() + draft_runner = Mock() + manager = _session_manager(target_runner, draft_runner) + manager._weight_update_loaded = False + + with patch("torch.distributed.barrier"): + output = manager.end_weight_update(EndWeightUpdateReqInput()) + + assert output.success is True + target_runner.end_weight_update.assert_called_once_with(run_post_load=True) + draft_runner.end_weight_update.assert_called_once_with(run_post_load=True) + assert manager._weight_update_in_progress is False + + +def test_end_weight_update_skips_post_load_on_both_when_weights_loaded(): + # A distributed/tensor load happened this session, so post_load is skipped on + # both runners; only quant finalize runs. + target_runner = Mock() + draft_runner = Mock() + manager = _session_manager(target_runner, draft_runner) + manager._weight_update_loaded = True + + with patch("torch.distributed.barrier"): + manager.end_weight_update(EndWeightUpdateReqInput()) + + target_runner.end_weight_update.assert_called_once_with(run_post_load=False) + draft_runner.end_weight_update.assert_called_once_with(run_post_load=False) + + +def test_model_runner_begin_end_wire_to_loader_hooks(): + # ModelRunner.begin/end delegate to the loader: begin restores; end runs + # post_load only when requested, always finalizes quant layout. + import sglang.srt.model_executor.model_runner as mr + + runner = SimpleNamespace(model=object(), device="cpu") + + with patch.object(mr, "restore_weight") as restore: + mr.ModelRunner.begin_weight_update(runner) + restore.assert_called_once() + + with ( + patch.object(mr, "post_load_weights") as post_load, + patch.object(mr, "postprocess_weight") as postprocess, + ): + mr.ModelRunner.end_weight_update(runner, run_post_load=True) + post_load.assert_called_once() + postprocess.assert_called_once() + + with ( + patch.object(mr, "post_load_weights") as post_load, + patch.object(mr, "postprocess_weight") as postprocess, + ): + mr.ModelRunner.end_weight_update(runner, run_post_load=False) + post_load.assert_not_called() + postprocess.assert_called_once() + + +def test_begin_weight_update_selector_restores_only_selected_and_is_recorded(): + # begin(selector="draft") opens the session on the draft only; the target is + # untouched, and the selector is recorded for end to reuse. + target_runner = Mock() + draft_runner = Mock() + manager = _session_manager(target_runner, draft_runner) + manager._weight_update_in_progress = False + + with patch("torch.distributed.barrier"): + manager.begin_weight_update(BeginWeightUpdateReqInput(selector="draft")) + + target_runner.begin_weight_update.assert_not_called() + draft_runner.begin_weight_update.assert_called_once_with() + assert manager._weight_update_selector == "draft" + + +def test_end_weight_update_reuses_session_selector_from_begin(): + # end has no selector of its own; it finalizes exactly the set begin opened. + target_runner = Mock() + draft_runner = Mock() + manager = _session_manager(target_runner, draft_runner) + manager._weight_update_in_progress = False + + with patch("torch.distributed.barrier"): + manager.begin_weight_update(BeginWeightUpdateReqInput(selector="draft")) + manager.end_weight_update(EndWeightUpdateReqInput()) + + target_runner.end_weight_update.assert_not_called() + draft_runner.end_weight_update.assert_called_once() + + +def test_begin_weight_update_rejects_reentry(): + # A second begin while a session is open would leave the first session's + # restored runners unfinalized — reject it loudly. + manager = _session_manager(Mock(), Mock()) + manager._weight_update_in_progress = True + + with patch("torch.distributed.barrier"): + with pytest.raises(AssertionError, match="already open"): + manager.begin_weight_update(BeginWeightUpdateReqInput()) diff --git a/test/registered/unit/sampling/test_sampling_batch_info.py b/test/registered/unit/sampling/test_sampling_batch_info.py index 91cefe30ac70..f9146a5e333d 100644 --- a/test/registered/unit/sampling/test_sampling_batch_info.py +++ b/test/registered/unit/sampling/test_sampling_batch_info.py @@ -101,6 +101,32 @@ def test_len_matches_batch_size(self): class TestSamplingMaskBatchIndices(CustomTestCase): + def test_filter_and_merge_preserve_support_logprob_modes(self): + info = _make_info( + batch_size=3, + return_sampling_masks=[True, True, False], + return_sampling_support_logprobs=[False, True, False], + sampling_mask_batch_indices=torch.tensor([0, 1]), + ) + info.filter_batch([1, 2], torch.tensor([1, 2])) + self.assertEqual(info.return_sampling_support_logprobs, [True, False]) + self.assertEqual(info.sampling_support_logprobs_capture_indices.tolist(), [0]) + + other = _make_info( + batch_size=2, + return_sampling_masks=[True, True], + return_sampling_support_logprobs=[False, True], + sampling_mask_batch_indices=torch.tensor([0, 1]), + ) + info.merge_batch(other) + self.assertEqual( + info.return_sampling_support_logprobs, + [True, False, False, True], + ) + self.assertEqual( + info.sampling_support_logprobs_capture_indices.tolist(), [0, 2] + ) + def test_filter_removes_last_opted_in_row_then_merge_restores_capture(self): info = _make_info( batch_size=2, diff --git a/test/registered/unit/true_on_policy/test_dense_deterministic_math.py b/test/registered/unit/true_on_policy/test_dense_deterministic_math.py new file mode 100644 index 000000000000..a0aafcd5f658 --- /dev/null +++ b/test/registered/unit/true_on_policy/test_dense_deterministic_math.py @@ -0,0 +1,344 @@ +import json +import os +import subprocess +import textwrap +import unittest +from types import SimpleNamespace + +import torch + +from sglang.srt.true_on_policy import ( + QWEN3_DENSE_TRUE_ON_POLICY_V1, + get_on_policy_rms_norm_kwargs, + should_force_bfloat16_dense_tensor_math, + should_force_bfloat16_lm_head, +) +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=12, suite="stage-a-test-cpu") + + +def _run_dense_math_script(script_body: str) -> dict[str, object]: + stubbed_imports = textwrap.dedent(""" + import importlib.machinery + import json + import sys + import types + from pydantic import BaseModel + + def install_openai_stubs(): + openai_mod = types.ModuleType("openai") + openai_types_mod = types.ModuleType("openai.types") + openai_responses_mod = types.ModuleType("openai.types.responses") + openai_response_mod = types.ModuleType("openai.types.responses.response") + openai_tool_mod = types.ModuleType("openai.types.responses.tool") + + openai_mod.__spec__ = importlib.machinery.ModuleSpec("openai", loader=None) + openai_types_mod.__spec__ = importlib.machinery.ModuleSpec("openai.types", loader=None) + openai_responses_mod.__spec__ = importlib.machinery.ModuleSpec( + "openai.types.responses", loader=None + ) + openai_response_mod.__spec__ = importlib.machinery.ModuleSpec( + "openai.types.responses.response", loader=None + ) + openai_tool_mod.__spec__ = importlib.machinery.ModuleSpec( + "openai.types.responses.tool", loader=None + ) + + for name in [ + "ResponseFunctionToolCall", + "ResponseInputItemParam", + "ResponseOutputItem", + "ResponseOutputMessage", + "ResponseOutputText", + "ResponseReasoningItem", + ]: + setattr(openai_responses_mod, name, type(name, (BaseModel,), {})) + + openai_response_mod.ToolChoice = type("ToolChoice", (BaseModel,), {}) + openai_tool_mod.Tool = type("Tool", (BaseModel,), {}) + + sys.modules.setdefault("openai", openai_mod) + sys.modules.setdefault("openai.types", openai_types_mod) + sys.modules.setdefault("openai.types.responses", openai_responses_mod) + sys.modules.setdefault("openai.types.responses.response", openai_response_mod) + sys.modules.setdefault("openai.types.responses.tool", openai_tool_mod) + + install_openai_stubs() + + hf_utils_mod = types.ModuleType("sglang.srt.utils.hf_transformers_utils") + hf_utils_mod.__spec__ = importlib.machinery.ModuleSpec( + "sglang.srt.utils.hf_transformers_utils", loader=None + ) + hf_utils_mod.check_gguf_file = lambda *args, **kwargs: False + hf_utils_mod.get_rope_config = lambda config: ( + getattr(config, "rope_theta", 1000000), + getattr(config, "rope_scaling", None), + ) + sys.modules.setdefault("sglang.srt.utils.hf_transformers_utils", hf_utils_mod) + + gguf_mod = types.ModuleType("gguf") + gguf_mod.__spec__ = importlib.machinery.ModuleSpec("gguf", loader=None) + gguf_mod.GGMLQuantizationType = type( + "GGMLQuantizationType", + (), + { + "F32": 0, + "F16": 1, + "BF16": 2, + "Q4_0": 3, + "Q4_1": 4, + "Q5_0": 5, + "Q5_1": 6, + "Q8_0": 7, + "Q8_1": 8, + "Q2_K": 9, + "Q3_K": 10, + "Q4_K": 11, + "Q5_K": 12, + "Q6_K": 13, + "IQ1_S": 14, + "IQ1_M": 15, + "IQ2_XXS": 16, + "IQ2_XS": 17, + "IQ2_S": 18, + "IQ3_XXS": 19, + "IQ3_S": 20, + "IQ4_NL": 21, + "IQ4_XS": 22, + }, + ) + sys.modules.setdefault("gguf", gguf_mod) + """) + + env = dict(os.environ) + pythonpath = env.get("PYTHONPATH") + repo_python = "python" + env["PYTHONPATH"] = ( + f"{repo_python}{os.pathsep}{pythonpath}" if pythonpath else repo_python + ) + script = f"{stubbed_imports}\n{script_body}" + completed = subprocess.run( + ["python", "-c", script], + check=True, + capture_output=True, + text=True, + env=env, + ) + return json.loads(completed.stdout) + + +class TestDenseOnPolicyHelpers(unittest.TestCase): + def test_default_dense_math_helpers_are_inactive(self): + server_args = SimpleNamespace( + true_on_policy_contract=None, + tp_size=1, + ) + + self.assertFalse(should_force_bfloat16_dense_tensor_math(server_args)) + self.assertFalse( + should_force_bfloat16_lm_head( + server_args=server_args, + use_fp32_lm_head=False, + ) + ) + self.assertEqual(get_on_policy_rms_norm_kwargs(server_args), {}) + + def test_on_policy_dense_math_helpers_enable_bfloat16_and_rms_norm_kwargs(self): + server_args = SimpleNamespace( + true_on_policy_contract=QWEN3_DENSE_TRUE_ON_POLICY_V1, + tp_size=1, + ) + + kwargs = get_on_policy_rms_norm_kwargs( + server_args, + weight_dtype=torch.float32, + override_orig_dtype=torch.float32, + fp32_residual=True, + ) + + self.assertTrue(should_force_bfloat16_dense_tensor_math(server_args)) + self.assertTrue( + should_force_bfloat16_lm_head( + server_args=server_args, + use_fp32_lm_head=False, + ) + ) + self.assertFalse( + should_force_bfloat16_lm_head( + server_args=server_args, + use_fp32_lm_head=True, + ) + ) + self.assertEqual(kwargs["weight_dtype"], torch.float32) + self.assertEqual(kwargs["override_orig_dtype"], torch.float32) + self.assertTrue(kwargs["cast_x_before_out_mul"]) + self.assertTrue(kwargs["fp32_residual"]) + + +class TestDenseOnPolicyContracts(unittest.TestCase): + def test_qwen3_style_rms_norm_keeps_fp32_weight_output_and_residual(self): + result = _run_dense_math_script( + textwrap.dedent(""" + import json + from types import SimpleNamespace + + import torch + + from sglang.srt.layers.layernorm import RMSNorm + from sglang.srt.true_on_policy import get_on_policy_rms_norm_kwargs + + from sglang.srt.true_on_policy import QWEN3_DENSE_TRUE_ON_POLICY_V1 + + server_args = SimpleNamespace( + true_on_policy_contract=QWEN3_DENSE_TRUE_ON_POLICY_V1, + tp_size=1, + ) + norm = RMSNorm( + 4, + eps=1e-6, + **get_on_policy_rms_norm_kwargs( + server_args, + weight_dtype=torch.float32, + override_orig_dtype=torch.float32, + fp32_residual=True, + ), + ) + x = torch.randn(2, 4, dtype=torch.bfloat16) + residual = torch.randn(2, 4, dtype=torch.bfloat16) + out, residual_out = norm.forward_native(x, residual) + print( + json.dumps( + { + "weight_dtype": str(norm.weight.dtype), + "out_dtype": str(out.dtype), + "residual_dtype": str(residual_out.dtype), + } + ) + ) + """) + ) + + self.assertEqual(result["weight_dtype"], "torch.float32") + self.assertEqual(result["out_dtype"], "torch.float32") + self.assertEqual(result["residual_dtype"], "torch.float32") + + def test_rms_norm_can_self_configure_from_true_on_policy_role_hints(self): + result = _run_dense_math_script( + textwrap.dedent(""" + import json + + import torch + + from sglang.srt.layers.layernorm import RMSNorm + from sglang.srt.server_args import ( + ServerArgs, + get_global_server_args, + set_global_server_args_for_scheduler, + ) + from sglang.srt.true_on_policy import QWEN3_DENSE_TRUE_ON_POLICY_V1 + + set_global_server_args_for_scheduler(ServerArgs(model_path="dummy")) + server_args = get_global_server_args() + server_args.true_on_policy_contract = QWEN3_DENSE_TRUE_ON_POLICY_V1 + server_args.tp_size = 1 + norm = RMSNorm( + 4, + eps=1e-6, + true_on_policy_weight_dtype=torch.float32, + true_on_policy_override_orig_dtype=torch.float32, + true_on_policy_fp32_residual=True, + ) + print( + json.dumps( + { + "weight_dtype": str(norm.weight.dtype), + "cast_x_before_out_mul": norm.cast_x_before_out_mul, + "fp32_residual": norm.fp32_residual, + "override_orig_dtype": str(norm.override_orig_dtype), + } + ) + ) + """) + ) + + self.assertEqual(result["weight_dtype"], "torch.float32") + self.assertTrue(result["cast_x_before_out_mul"]) + self.assertTrue(result["fp32_residual"]) + self.assertEqual(result["override_orig_dtype"], "torch.float32") + + def test_on_policy_lm_head_forces_bfloat16_matmul_inputs(self): + result = _run_dense_math_script( + textwrap.dedent(""" + import json + from types import SimpleNamespace + from unittest.mock import patch + + import torch + import torch.nn as nn + + from sglang.srt.layers.logits_processor import LogitsProcessor + from sglang.srt.server_args import ( + ServerArgs, + get_global_server_args, + set_global_server_args_for_scheduler, + ) + from sglang.srt.true_on_policy import QWEN3_DENSE_TRUE_ON_POLICY_V1 + + class DummyMeta: + gathered_buffer = None + next_token_logits_buffer = None + + def compute_dp_attention_metadata(self): + return None + + class LMHeadStub(nn.Module): + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.randn(8, 4, dtype=torch.float32)) + + set_global_server_args_for_scheduler(ServerArgs(model_path="dummy")) + get_global_server_args().enable_dp_lm_head = False + get_global_server_args().enable_fp32_lm_head = False + get_global_server_args().true_on_policy_contract = QWEN3_DENSE_TRUE_ON_POLICY_V1 + get_global_server_args().tp_size = 1 + + processor = LogitsProcessor( + SimpleNamespace(vocab_size=8, final_logit_softcapping=None), + skip_all_gather=True, + logit_scale=None, + ) + hidden_states = torch.randn(2, 4, dtype=torch.float32) + head = LMHeadStub() + captured = {} + + original_matmul = torch.matmul + + def probe_matmul(a, b, *args, **kwargs): + if not captured: + captured["a_dtype"] = str(a.dtype) + captured["b_dtype"] = str(b.dtype) + return original_matmul(a, b, *args, **kwargs) + + with patch("torch.matmul", new=probe_matmul): + logits = processor._get_logits(hidden_states, head, DummyMeta()) + + print( + json.dumps( + { + "a_dtype": captured["a_dtype"], + "b_dtype": captured["b_dtype"], + "logits_dtype": str(logits.dtype), + } + ) + ) + """) + ) + + self.assertEqual(result["a_dtype"], "torch.bfloat16") + self.assertEqual(result["b_dtype"], "torch.bfloat16") + self.assertEqual(result["logits_dtype"], "torch.bfloat16") + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/test/registered/unit/true_on_policy/test_on_policy_wiring.py b/test/registered/unit/true_on_policy/test_on_policy_wiring.py new file mode 100644 index 000000000000..946f1922a9e0 --- /dev/null +++ b/test/registered/unit/true_on_policy/test_on_policy_wiring.py @@ -0,0 +1,606 @@ +import json +import os +import subprocess +import textwrap +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.srt.true_on_policy import ( + QWEN3_DENSE_TRUE_ON_POLICY_V1, + get_rl_on_policy_target, + get_true_on_policy_contract, + is_tp_invariant_target, + is_true_on_policy_enabled, + patch_prefill_only_deterministic_inference_for_cuda_graph, + resolve_true_on_policy_runtime_policy, + should_disable_flashinfer_allreduce_fusion, + should_disable_fused_qk_norm_mrope, + should_disable_mlp_allreduce_fusion_for_on_policy, + should_disable_reduce_scatter_for_on_policy, + should_use_tp_invariant_row_linear, + should_use_tp_invariant_tree_all_reduce, +) +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=12, suite="stage-a-test-cpu") + +_PATCH_TARGET = "sglang.srt.server_args.get_global_server_args" + + +def _run_server_args_script(argv: list[str]) -> dict[str, object]: + stubbed_imports = textwrap.dedent(""" + import argparse + import importlib.machinery + import json + import sys + import types + from types import SimpleNamespace + from unittest.mock import patch + + from pydantic import BaseModel + + def install_openai_stubs(): + openai_mod = types.ModuleType("openai") + openai_types_mod = types.ModuleType("openai.types") + openai_responses_mod = types.ModuleType("openai.types.responses") + openai_response_mod = types.ModuleType("openai.types.responses.response") + openai_tool_mod = types.ModuleType("openai.types.responses.tool") + + openai_mod.__spec__ = importlib.machinery.ModuleSpec("openai", loader=None) + openai_types_mod.__spec__ = importlib.machinery.ModuleSpec("openai.types", loader=None) + openai_responses_mod.__spec__ = importlib.machinery.ModuleSpec( + "openai.types.responses", loader=None + ) + openai_response_mod.__spec__ = importlib.machinery.ModuleSpec( + "openai.types.responses.response", loader=None + ) + openai_tool_mod.__spec__ = importlib.machinery.ModuleSpec( + "openai.types.responses.tool", loader=None + ) + + for name in [ + "ResponseFunctionToolCall", + "ResponseInputItemParam", + "ResponseOutputItem", + "ResponseOutputMessage", + "ResponseOutputText", + "ResponseReasoningItem", + ]: + setattr(openai_responses_mod, name, type(name, (BaseModel,), {})) + + openai_response_mod.ToolChoice = type("ToolChoice", (BaseModel,), {}) + openai_tool_mod.Tool = type("Tool", (BaseModel,), {}) + + sys.modules.setdefault("openai", openai_mod) + sys.modules.setdefault("openai.types", openai_types_mod) + sys.modules.setdefault("openai.types.responses", openai_responses_mod) + sys.modules.setdefault("openai.types.responses.response", openai_response_mod) + sys.modules.setdefault("openai.types.responses.tool", openai_tool_mod) + + install_openai_stubs() + + hf_utils_mod = types.ModuleType("sglang.srt.utils.hf_transformers_utils") + hf_utils_mod.__spec__ = importlib.machinery.ModuleSpec( + "sglang.srt.utils.hf_transformers_utils", loader=None + ) + hf_utils_mod.check_gguf_file = lambda *args, **kwargs: False + sys.modules.setdefault("sglang.srt.utils.hf_transformers_utils", hf_utils_mod) + + from sglang.srt.server_args import ServerArgs + + def _mock_model_config(): + return SimpleNamespace( + hf_config=SimpleNamespace(architectures=["Qwen2ForCausalLM"]) + ) + + parser = argparse.ArgumentParser() + ServerArgs.add_cli_args(parser) + cli_args = parser.parse_args(ARGV) + + with patch("sglang.srt.server_args.get_device", return_value="cuda"), patch.object( + ServerArgs, "get_model_config", return_value=_mock_model_config() + ): + server_args = ServerArgs.from_cli_args(cli_args) + server_args._handle_deterministic_inference() + + print( + json.dumps( + { + "enable_deterministic_inference": server_args.enable_deterministic_inference, + "enable_prefill_only_deterministic_inference": server_args.enable_prefill_only_deterministic_inference, + "enable_flashinfer_allreduce_fusion": server_args.enable_flashinfer_allreduce_fusion, + "rl_on_policy_target": server_args.rl_on_policy_target, + "true_on_policy_contract": server_args.true_on_policy_contract, + "sampling_backend": server_args.sampling_backend, + } + ) + ) + """) + + env = dict(os.environ) + pythonpath = env.get("PYTHONPATH") + repo_python = "python" + env["PYTHONPATH"] = ( + f"{repo_python}{os.pathsep}{pythonpath}" if pythonpath else repo_python + ) + script = f"ARGV = {argv!r}\n{stubbed_imports}" + completed = subprocess.run( + ["python", "-c", script], + check=True, + capture_output=True, + text=True, + env=env, + ) + return json.loads(completed.stdout) + + +class TestOnPolicyServerArgs(unittest.TestCase): + def test_cli_parses_prefill_only_deterministic_flag(self): + result = _run_server_args_script( + [ + "--model-path", + "dummy", + "--attention-backend", + "triton", + "--enable-prefill-only-deterministic-inference", + ] + ) + + self.assertTrue(result["enable_prefill_only_deterministic_inference"]) + self.assertTrue(result["enable_deterministic_inference"]) + self.assertIsNone(result["rl_on_policy_target"]) + self.assertEqual(result["sampling_backend"], "pytorch") + + def test_cli_accepts_fsdp_and_fsdp_tp_targets(self): + fsdp_tp_result = _run_server_args_script( + [ + "--model-path", + "dummy", + "--attention-backend", + "triton", + "--rl-on-policy-target", + "fsdp_tp", + ] + ) + self.assertEqual(fsdp_tp_result["rl_on_policy_target"], "fsdp_tp") + self.assertIsNone(fsdp_tp_result["true_on_policy_contract"]) + self.assertTrue(fsdp_tp_result["enable_deterministic_inference"]) + + fsdp_result = _run_server_args_script( + [ + "--model-path", + "dummy", + "--attention-backend", + "triton", + "--rl-on-policy-target", + "fsdp", + ] + ) + self.assertEqual(fsdp_result["rl_on_policy_target"], "fsdp") + self.assertIsNone(fsdp_result["true_on_policy_contract"]) + self.assertTrue(fsdp_result["enable_deterministic_inference"]) + + def test_cli_accepts_explicit_true_on_policy_contract(self): + result = _run_server_args_script( + [ + "--model-path", + "dummy", + "--attention-backend", + "triton", + "--true-on-policy-contract", + QWEN3_DENSE_TRUE_ON_POLICY_V1, + ] + ) + + self.assertIsNone(result["rl_on_policy_target"]) + self.assertEqual( + result["true_on_policy_contract"], QWEN3_DENSE_TRUE_ON_POLICY_V1 + ) + self.assertTrue(result["enable_deterministic_inference"]) + + def test_contract_tp_rollout_disables_flashinfer_allreduce_fusion(self): + result = _run_server_args_script( + [ + "--model-path", + "dummy", + "--attention-backend", + "triton", + "--tensor-parallel-size", + "2", + "--true-on-policy-contract", + QWEN3_DENSE_TRUE_ON_POLICY_V1, + "--enable-flashinfer-allreduce-fusion", + ] + ) + self.assertFalse(result["enable_flashinfer_allreduce_fusion"]) + + def test_legacy_target_keeps_flashinfer_allreduce_fusion_available(self): + result = _run_server_args_script( + [ + "--model-path", + "dummy", + "--attention-backend", + "triton", + "--rl-on-policy-target", + "fsdp_tp", + "--enable-flashinfer-allreduce-fusion", + ] + ) + self.assertTrue(result["enable_flashinfer_allreduce_fusion"]) + + +def _mock_args(**kwargs): + defaults = dict( + rl_on_policy_target=None, + true_on_policy_contract=None, + tp_size=1, + ) + defaults.update(kwargs) + return SimpleNamespace(**defaults) + + +def _contract_args(*, tp_size: int = 1): + return _mock_args( + true_on_policy_contract=QWEN3_DENSE_TRUE_ON_POLICY_V1, + tp_size=tp_size, + ) + + +class TestDefaultPathUnchanged(unittest.TestCase): + """Default serving must not enter true-on-policy policy paths.""" + + def setUp(self): + self.default_args = _mock_args() + + def test_default_args_no_on_policy(self): + with patch(_PATCH_TARGET, return_value=self.default_args): + self.assertIsNone(get_rl_on_policy_target()) + self.assertFalse(is_true_on_policy_enabled()) + self.assertFalse(is_tp_invariant_target()) + + def test_default_args_row_linear_uses_quant_method(self): + with patch(_PATCH_TARGET, return_value=self.default_args): + self.assertFalse( + should_use_tp_invariant_row_linear( + 256, + row_linear_enable_inv=True, + ) + ) + + def test_default_args_tree_allreduce_not_selected(self): + with patch(_PATCH_TARGET, return_value=self.default_args): + self.assertFalse( + should_use_tp_invariant_tree_all_reduce( + accl_binary_tree_enabled=False, + ) + ) + + def test_default_args_reduce_scatter_available(self): + with patch(_PATCH_TARGET, return_value=self.default_args): + self.assertFalse(should_disable_reduce_scatter_for_on_policy()) + + def test_default_args_mlp_fusion_available(self): + with patch(_PATCH_TARGET, return_value=self.default_args): + self.assertFalse(should_disable_mlp_allreduce_fusion_for_on_policy()) + + def test_default_args_flashinfer_fusion_available(self): + with patch(_PATCH_TARGET, return_value=self.default_args): + self.assertFalse(should_disable_flashinfer_allreduce_fusion()) + + def test_default_server_args_cli_no_on_policy_flags(self): + result = _run_server_args_script( + ["--model-path", "dummy", "--attention-backend", "triton"] + ) + self.assertIsNone(result["rl_on_policy_target"]) + self.assertFalse(result["enable_deterministic_inference"]) + self.assertFalse(result["enable_prefill_only_deterministic_inference"]) + + +class TestOnPolicyHelpers(unittest.TestCase): + def test_tp_invariant_row_linear_selection(self): + with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=2)): + self.assertTrue( + should_use_tp_invariant_row_linear( + 256, + row_linear_enable_inv=True, + ) + ) + + def test_tp_invariant_row_linear_selection_is_contract_owned(self): + with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=2)): + with patch.dict(os.environ, {"ROW_LINEAR_ENABLE_INV": "0"}): + self.assertTrue(should_use_tp_invariant_row_linear(256)) + + def test_contract_resolver_ignores_legacy_target_without_contract(self): + policy = resolve_true_on_policy_runtime_policy( + _mock_args(rl_on_policy_target="fsdp_tp", tp_size=2) + ) + + self.assertIsNone(policy.contract_name) + self.assertFalse(policy.enabled) + self.assertFalse(policy.tp_invariant_row_linear) + self.assertFalse(policy.deterministic_tree_all_reduce) + + def test_contract_resolver_accepts_explicit_qwen3_dense_contract(self): + args_tp1 = _contract_args(tp_size=1) + policy = resolve_true_on_policy_runtime_policy(args_tp1) + + self.assertTrue(policy.enabled) + self.assertTrue(policy.force_bfloat16_dense_tensor_math) + self.assertFalse(policy.tp_invariant_row_linear) + with patch(_PATCH_TARGET, return_value=args_tp1): + self.assertFalse( + should_use_tp_invariant_row_linear( + 96, + row_linear_enable_inv=True, + ) + ) + self.assertFalse( + should_use_tp_invariant_row_linear( + 256, + row_linear_enable_inv=True, + ) + ) + + def test_contract_object_owns_sglang_runtime_policy_values(self): + contract = get_true_on_policy_contract(QWEN3_DENSE_TRUE_ON_POLICY_V1) + + policy = contract.policy_for(_contract_args(tp_size=2)) + + self.assertEqual(contract.schema.name, QWEN3_DENSE_TRUE_ON_POLICY_V1) + self.assertEqual(contract.schema.model_family, "qwen3_dense") + self.assertEqual(policy.contract_name, QWEN3_DENSE_TRUE_ON_POLICY_V1) + self.assertTrue(policy.enabled) + self.assertTrue(policy.force_bfloat16_dense_tensor_math) + self.assertTrue(policy.force_bfloat16_lm_head) + self.assertTrue(policy.disable_reduce_scatter) + self.assertTrue(policy.disable_mlp_allreduce_fusion) + self.assertTrue(policy.disable_flashinfer_allreduce_fusion) + self.assertTrue(policy.tp_invariant_row_linear) + self.assertTrue(policy.deterministic_tree_all_reduce) + self.assertTrue(policy.disable_fused_qk_norm_mrope) + + def test_reduce_scatter_and_fusion_are_disabled_for_contract(self): + with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=1)): + self.assertTrue(should_disable_reduce_scatter_for_on_policy()) + with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=2)): + self.assertTrue(should_disable_mlp_allreduce_fusion_for_on_policy()) + with patch(_PATCH_TARGET, return_value=_mock_args()): + self.assertFalse(should_disable_reduce_scatter_for_on_policy()) + + def test_tree_all_reduce_selection_requires_tp_rollout_and_no_accl(self): + with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=2)): + self.assertTrue( + should_use_tp_invariant_tree_all_reduce( + accl_binary_tree_enabled=False, + ) + ) + self.assertFalse( + should_use_tp_invariant_tree_all_reduce( + accl_binary_tree_enabled=True, + ) + ) + with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=1)): + self.assertFalse( + should_use_tp_invariant_tree_all_reduce( + accl_binary_tree_enabled=False, + ) + ) + + def test_tree_all_reduce_selection_is_contract_owned(self): + with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=2)): + with patch.dict(os.environ, {"ACCL_BINARY_TREE_ENABLE": "1"}): + self.assertTrue(should_use_tp_invariant_tree_all_reduce()) + + def test_attention_handoff_tree_reduce_uses_attention_tp_group(self): + from sglang.srt.layers.communicator import ( + CommunicateWithAllReduceAndLayerNormFn, + ) + + hidden_states = torch.ones(2, 4) + residual = torch.full((2, 4), 3.0) + + class FakeNorm: + def __call__(self, x, residual): + return x + residual, residual + + with ( + patch( + "sglang.srt.layers.communicator.get_attn_tp_context", + return_value=SimpleNamespace(input_scattered=False), + ), + patch( + "sglang.srt.layers.communicator.apply_aiter_all_reduce_fusion", + return_value=False, + ), + patch( + "sglang.srt.layers.communicator.apply_flashinfer_allreduce_fusion", + return_value=False, + ), + patch( + "sglang.srt.layers.communicator.attention_tensor_model_parallel_all_reduce", + side_effect=lambda x: x + 10.0, + ) as attn_tree_reduce, + ): + output, output_residual = ( + CommunicateWithAllReduceAndLayerNormFn._gather_hidden_states_and_residual( + hidden_states, + residual, + forward_batch=None, + layernorm=FakeNorm(), + context=SimpleNamespace(attn_dp_size=1, cache=None), + residual_input_mode=None, + ) + ) + + attn_tree_reduce.assert_called_once() + torch.testing.assert_close(output, hidden_states + 10.0 + residual) + torch.testing.assert_close(output_residual, residual) + + def test_prefill_only_cuda_graph_patch_only_scopes_attention_splits(self): + server_args = SimpleNamespace( + enable_prefill_only_deterministic_inference=True, + enable_deterministic_inference=True, + enable_flashinfer_allreduce_fusion=False, + rl_on_policy_target="fsdp_tp", + true_on_policy_contract=QWEN3_DENSE_TRUE_ON_POLICY_V1, + disable_custom_all_reduce=True, + ) + attn_backend = SimpleNamespace(num_splits=1) + + with patch.dict( + os.environ, + { + "SGLANG_ENABLE_DETERMINISTIC_INFERENCE": "1", + "SGLANG_DISABLE_CUSTOM_ALL_REDUCE": "1", + "NCCL_ALGO": "allreduce:tree", + }, + clear=False, + ): + with patch_prefill_only_deterministic_inference_for_cuda_graph( + server_args, + attn_backend=attn_backend, + ) as patched: + self.assertTrue(patched) + self.assertTrue(server_args.enable_deterministic_inference) + self.assertFalse(server_args.enable_flashinfer_allreduce_fusion) + self.assertEqual(server_args.rl_on_policy_target, "fsdp_tp") + self.assertEqual( + server_args.true_on_policy_contract, + QWEN3_DENSE_TRUE_ON_POLICY_V1, + ) + self.assertEqual(attn_backend.num_splits, 0) + self.assertEqual( + os.environ["SGLANG_ENABLE_DETERMINISTIC_INFERENCE"], "1" + ) + self.assertEqual(os.environ["SGLANG_DISABLE_CUSTOM_ALL_REDUCE"], "1") + self.assertEqual(os.environ["NCCL_ALGO"], "allreduce:tree") + + self.assertTrue(server_args.enable_deterministic_inference) + self.assertFalse(server_args.enable_flashinfer_allreduce_fusion) + self.assertEqual(server_args.rl_on_policy_target, "fsdp_tp") + self.assertEqual( + server_args.true_on_policy_contract, + QWEN3_DENSE_TRUE_ON_POLICY_V1, + ) + self.assertTrue(server_args.disable_custom_all_reduce) + self.assertEqual(attn_backend.num_splits, 1) + self.assertEqual(os.environ["SGLANG_ENABLE_DETERMINISTIC_INFERENCE"], "1") + self.assertEqual(os.environ["SGLANG_DISABLE_CUSTOM_ALL_REDUCE"], "1") + self.assertEqual(os.environ["NCCL_ALGO"], "allreduce:tree") + + def test_row_linear_k_alignment_edge_cases(self): + with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=2)): + self.assertFalse( + should_use_tp_invariant_row_linear(64, row_linear_enable_inv=True), + ) + self.assertTrue( + should_use_tp_invariant_row_linear(128, row_linear_enable_inv=True), + ) + self.assertFalse( + should_use_tp_invariant_row_linear(300, row_linear_enable_inv=True), + ) + self.assertTrue( + should_use_tp_invariant_row_linear(3584, row_linear_enable_inv=True), + ) + + def test_row_linear_explicit_override_can_disable(self): + with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=2)): + self.assertFalse( + should_use_tp_invariant_row_linear(256, row_linear_enable_inv=False) + ) + + def test_flashinfer_allreduce_fusion_helpers(self): + with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=2)): + self.assertTrue(should_disable_flashinfer_allreduce_fusion()) + with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=1)): + self.assertFalse(should_disable_flashinfer_allreduce_fusion()) + with patch(_PATCH_TARGET, return_value=_mock_args()): + self.assertFalse(should_disable_flashinfer_allreduce_fusion()) + + def test_fused_qk_norm_mrope_helper_follows_true_on_policy_contract(self): + with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=1)): + self.assertTrue(should_disable_fused_qk_norm_mrope()) + with patch(_PATCH_TARGET, return_value=_mock_args()): + self.assertFalse(should_disable_fused_qk_norm_mrope()) + + def test_get_rl_on_policy_target_returns_correct_value(self): + with patch( + _PATCH_TARGET, return_value=_mock_args(rl_on_policy_target="fsdp_tp") + ): + self.assertEqual(get_rl_on_policy_target(), "fsdp_tp") + with patch(_PATCH_TARGET, return_value=_mock_args(rl_on_policy_target="fsdp")): + self.assertEqual(get_rl_on_policy_target(), "fsdp") + with patch(_PATCH_TARGET, return_value=_mock_args()): + self.assertIsNone(get_rl_on_policy_target()) + + def test_is_true_on_policy_enabled_for_both_targets(self): + with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=1)): + self.assertTrue(is_true_on_policy_enabled()) + with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=2)): + self.assertTrue(is_true_on_policy_enabled()) + with patch(_PATCH_TARGET, return_value=_mock_args()): + self.assertFalse(is_true_on_policy_enabled()) + + def test_is_tp_invariant_target_only_fsdp_tp(self): + with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=2)): + self.assertTrue(is_tp_invariant_target()) + with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=1)): + self.assertFalse(is_tp_invariant_target()) + with patch(_PATCH_TARGET, return_value=_mock_args()): + self.assertFalse(is_tp_invariant_target()) + + def test_cuda_graph_patch_noop_when_disabled(self): + server_args = SimpleNamespace( + enable_prefill_only_deterministic_inference=False, + enable_deterministic_inference=True, + rl_on_policy_target="fsdp_tp", + ) + with patch_prefill_only_deterministic_inference_for_cuda_graph( + server_args, + ) as patched: + self.assertFalse(patched) + self.assertTrue(server_args.enable_deterministic_inference) + self.assertEqual(server_args.rl_on_policy_target, "fsdp_tp") + + def test_cuda_graph_patch_noop_when_dvr_verify(self): + server_args = SimpleNamespace( + enable_prefill_only_deterministic_inference=True, + enable_deterministic_inference=True, + enable_flashinfer_allreduce_fusion=False, + rl_on_policy_target="fsdp_tp", + disable_custom_all_reduce=True, + ) + with patch_prefill_only_deterministic_inference_for_cuda_graph( + server_args, + dvr_target_verify_cuda_graph=True, + ) as patched: + self.assertFalse(patched) + self.assertTrue(server_args.enable_deterministic_inference) + self.assertEqual(server_args.rl_on_policy_target, "fsdp_tp") + + def test_tp_invariant_ops_import_is_available(self): + import sglang.srt.tp_invariant_ops as tp_invariant_ops + + self.assertTrue(hasattr(tp_invariant_ops, "matmul_tp_inv")) + + def test_legacy_on_policy_utils_import_matches_true_on_policy_namespace(self): + from sglang.srt import true_on_policy + from sglang.srt.layers import on_policy_utils as legacy + + self.assertIs( + legacy.should_use_tp_invariant_row_linear, + true_on_policy.should_use_tp_invariant_row_linear, + ) + self.assertIs( + legacy.patch_prefill_only_deterministic_inference_for_cuda_graph, + true_on_policy.patch_prefill_only_deterministic_inference_for_cuda_graph, + ) + self.assertTrue(hasattr(torch.ops, "tp_inv_ops")) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/utils/test_auth.py b/test/registered/unit/utils/test_auth.py index cfecfdb0d321..0473081c2c9c 100644 --- a/test/registered/unit/utils/test_auth.py +++ b/test/registered/unit/utils/test_auth.py @@ -80,6 +80,17 @@ def test_health_subpath_always_allowed(self): ) self.assertTrue(decision.allowed) + def test_ready_path_always_allowed(self): + decision = decide_request_auth( + method="GET", + path="/ready", + authorization_header=None, + api_key="secret", + admin_api_key=None, + auth_level=AuthLevel.NORMAL, + ) + self.assertTrue(decision.allowed) + def test_metrics_path_always_allowed(self): decision = decide_request_auth( method="GET", diff --git a/test/registered/unit/utils/test_http_server_auth.py b/test/registered/unit/utils/test_http_server_auth.py index 87e1f400e76e..fd3c8c3298c0 100644 --- a/test/registered/unit/utils/test_http_server_auth.py +++ b/test/registered/unit/utils/test_http_server_auth.py @@ -274,6 +274,7 @@ def test_health_and_metrics_are_always_allowed(self): paths_allowed = [ "/health", "/health_generate", + "/ready", "/metrics", "/metrics/", "/metrics/prometheus", diff --git a/test/registered/unit/utils/test_weight_checker.py b/test/registered/unit/utils/test_weight_checker.py index 702923f510e2..f1eb92d00afc 100644 --- a/test/registered/unit/utils/test_weight_checker.py +++ b/test/registered/unit/utils/test_weight_checker.py @@ -306,6 +306,18 @@ def test_substring_match_not_endswith(self): [("weird.cos_sin_cache.foo.bar", False, RawComparable(t))], ) + def test_skip_set_marks_quantized_entry_not_compared(self): + qweight, sf_fp32, _ = _build_fp8_quant_pair() + raw = {"x.weight": qweight, "x.weight_scale_inv": sf_fp32} + quantized_set = { + "x.weight": QuantizedWeight(Fp8BlockComparable, "x.weight_scale_inv") + } + ref = Fp8BlockComparable(qweight, sf_fp32) + _assert_entries_close( + _build_check_entries(raw, {"x.weight"}, quantized_set), + [("x.weight", False, ref)], + ) + # --- fp8 quant pair (real dequant on real fp8 tensors) --- @unittest.skipIf(is_hip(), "DeepGEMM is not supported on ROCm") @@ -627,6 +639,21 @@ def test_passes_when_unchanged(self): self.checker._snapshot() self.checker._compare() # no exception + def test_success_releases_snapshot(self): + self.checker._snapshot() + self.checker._compare() + self.assertIsNone(self.checker._snapshot_tensors) + self.assertIsNone(self.checker._snapshot_arena) + with self.assertRaises(AssertionError): + self.checker._compare() + + def test_failure_keeps_snapshot(self): + self.checker._snapshot() + self.checker._reset_tensors() + with self.assertRaises(Exception): + self.checker._compare() + self.assertIsNotNone(self.checker._snapshot_tensors) + def test_fails_after_reset_on_normal_param(self): self.checker._snapshot() self.checker._reset_tensors() diff --git a/test/registered/vlm/test_token_id_retokenize_e2e.py b/test/registered/vlm/test_token_id_retokenize_e2e.py index af37df453218..419d09a7a1a5 100644 --- a/test/registered/vlm/test_token_id_retokenize_e2e.py +++ b/test/registered/vlm/test_token_id_retokenize_e2e.py @@ -9,6 +9,9 @@ original tokens verbatim and only expands the image placeholder, so prompt_tokens stays faithful to what the client sent. +The test uses JPEG so CUDA decoding returns a CHW tensor, covering the same +exact-token path as PIL-backed images. + For each model we launch a real server twice with the same predefined, non-canonical prompt ("Describe" split into "D"+"escribe") plus one image: @@ -42,8 +45,8 @@ def _data_uri(): img = Image.new("RGB", (64, 64), (128, 128, 128)) buf = io.BytesIO() - img.save(buf, format="PNG") - return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode() + img.save(buf, format="JPEG") + return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode() def _build_drift_prompt(model, image_token): diff --git a/test/srt/models/test_params_mapping.py b/test/srt/models/test_params_mapping.py new file mode 100644 index 000000000000..e064266cab6a --- /dev/null +++ b/test/srt/models/test_params_mapping.py @@ -0,0 +1,292 @@ +"""Unit tests for ParameterMapper.""" + +from types import SimpleNamespace + +import pytest + +from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE +from sglang.srt.model_loader.parameter_mapper import ParameterMapper + +_DEEPSEEK_N_ROUTED = 4 +_DEEPSEEK_N_LOCAL = _DEEPSEEK_N_ROUTED + 1 # +1 fused shared expert +_QWEN3MOE_N = 4 +_GLM4LITE_N_ROUTED = 4 +_GLM4LITE_N_LOCAL = _GLM4LITE_N_ROUTED + 1 # +1 fused shared expert + + +def _make_model(**kwargs): + """Create a stub model object for ParameterMapper.from_model().""" + return SimpleNamespace(**kwargs) + + +def _deepseek_mutate(name): + if "mlp.shared_experts" in name: + return name.replace("mlp.shared_experts", f"mlp.experts.{_DEEPSEEK_N_ROUTED}") + return name + + +def _deepseek_scale_remap(name): + for s in ["k_scale", "v_scale"]: + if s in name: + return name.replace(f"{s[0]}_proj", "attn_mqa") + return name + + +def _glm4lite_mutate(name): + if "mlp.shared_experts" in name: + return name.replace("mlp.shared_experts", f"mlp.experts.{_GLM4LITE_N_ROUTED}") + return name + + +_LLAMA_SCALE_PATTERNS = [ + (".activation_scale", ".activation_scale", ".input_scale"), + (".weight_scale_inv", ".weight_scale_inv", ".weight_scale"), +] + + +def _llama_scale_remap(name): + for suffix, pattern, replacement in _LLAMA_SCALE_PATTERNS: + if name.endswith(suffix) and pattern in name: + return name.replace(pattern, replacement) + return name + + +@pytest.fixture +def qwen_mapper(): + """Qwen2/Qwen3 (dense): QKV fusion, gate/up fusion, no experts.""" + return ParameterMapper.from_model( + _make_model( + stacked_params_mapping=[ + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ], + ) + ) + + +@pytest.fixture +def llama_mapper(): + """Llama/GLM4 (dense): dot-prefixed stacked params, custom scale remap.""" + return ParameterMapper.from_model( + _make_model( + stacked_params_mapping=[ + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ], + custom_scale_remap=_llama_scale_remap, + ) + ) + + +@pytest.fixture +def qwen3moe_mapper(): + """Qwen3-MoE: QKV fusion + experts, no shared expert fusion.""" + return ParameterMapper.from_model( + _make_model( + stacked_params_mapping=[ + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ], + expert_params_mapping=FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=_QWEN3MOE_N, + ), + ) + ) + + +@pytest.fixture +def deepseek_mapper(): + """DeepSeek V2/V3: MLA A-proj fusion, shared expert fusion, custom scale remap.""" + return ParameterMapper.from_model( + _make_model( + stacked_params_mapping=[ + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ("fused_qkv_a_proj_with_mqa", "q_a_proj", 0), + ("fused_qkv_a_proj_with_mqa", "kv_a_proj_with_mqa", 1), + ], + expert_params_mapping=FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=_DEEPSEEK_N_LOCAL, + ), + mutate_weight_preload=_deepseek_mutate, + custom_scale_remap=_deepseek_scale_remap, + ) + ) + + +@pytest.fixture +def glm4lite_mapper(): + """GLM4-MoE-Lite (GLM-4.7): QKV fusion, MLA A-proj fusion, shared expert fusion, custom scale remap.""" + return ParameterMapper.from_model( + _make_model( + stacked_params_mapping=[ + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ("fused_qkv_a_proj_with_mqa", "q_a_proj", 0), + ("fused_qkv_a_proj_with_mqa", "kv_a_proj_with_mqa", 1), + ], + expert_params_mapping=FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=_GLM4LITE_N_LOCAL, + ), + mutate_weight_preload=_glm4lite_mutate, + custom_scale_remap=_deepseek_scale_remap, + ) + ) + + +# ── Helpers ────────────────────────────────────────────────────────────────── + + +def to_expect(name, shard=None, n=1, expert=None, n_exp=None): + """Shorthand for expected MappingResult fields.""" + return (name, shard, n, expert, n_exp) + + +def _assert(mapper, ckpt, expected): + r = mapper.map(ckpt) + name, shard, n, expert, n_exp = expected + assert ( + r.sglang_name, + r.shard_id, + r.num_shards, + r.expert_id, + r.num_local_experts, + ) == (name, shard, n, expert, n_exp), f"map({ckpt!r}) = {r}" + + +# ── Tests ──────────────────────────────────────────────────────────────────── + +# fmt: off +_QWEN_CASES = [ + # QKV fusion (Qwen2, Qwen3, GLM4-MoE) + ("layers.0.attn.q_proj.weight", to_expect("layers.0.attn.qkv_proj.weight", "q", 3)), + ("layers.0.attn.k_proj.weight", to_expect("layers.0.attn.qkv_proj.weight", "k", 3)), + ("layers.0.attn.v_proj.weight", to_expect("layers.0.attn.qkv_proj.weight", "v", 3)), + # Gate/Up fusion + ("layers.0.mlp.gate_proj.weight", to_expect("layers.0.mlp.gate_up_proj.weight", 0, 2)), + ("layers.0.mlp.up_proj.weight", to_expect("layers.0.mlp.gate_up_proj.weight", 1, 2)), + # Pass-through + ("layers.0.mlp.down_proj.weight", to_expect("layers.0.mlp.down_proj.weight")), + ("embed_tokens.weight", to_expect("embed_tokens.weight")), + # Standard scale remap (no custom_scale_remap) + ("model.layers.0.self_attn.k_scale", to_expect("model.layers.0.self_attn.attn.k_scale")), + ("model.layers.0.self_attn.v_scale", to_expect("model.layers.0.self_attn.attn.v_scale")), +] + +_LLAMA_CASES = [ + # Dot-prefixed QKV fusion (Llama, GLM4) + ("model.layers.0.self_attn.q_proj.weight", to_expect("model.layers.0.self_attn.qkv_proj.weight", "q", 3)), + ("model.layers.0.self_attn.k_proj.weight", to_expect("model.layers.0.self_attn.qkv_proj.weight", "k", 3)), + # Dot-prefixed gate/up + ("model.layers.0.mlp.gate_proj.weight", to_expect("model.layers.0.mlp.gate_up_proj.weight", 0, 2)), + # Llama-specific scale remap + stacked (scales follow their weights) + ("model.layers.0.mlp.gate_proj.activation_scale", to_expect("model.layers.0.mlp.gate_up_proj.input_scale", 0, 2)), + ("model.layers.0.mlp.gate_proj.weight_scale_inv", to_expect("model.layers.0.mlp.gate_up_proj.weight_scale", 0, 2)), + # Pass-through + ("model.layers.0.mlp.down_proj.weight", to_expect("model.layers.0.mlp.down_proj.weight")), +] + +_QWEN3MOE_CASES = [ + # QKV fusion + ("model.layers.0.self_attn.q_proj.weight", to_expect("model.layers.0.self_attn.qkv_proj.weight", "q", 3)), + # Expert mapping (no shared expert fusion) + ("model.layers.0.mlp.experts.0.gate_proj.weight", to_expect("model.layers.0.mlp.experts.w13_weight", "w1", 2, 0, _QWEN3MOE_N)), + ("model.layers.0.mlp.experts.3.down_proj.weight", to_expect("model.layers.0.mlp.experts.w2_weight", "w2", 1, 3, _QWEN3MOE_N)), + # shared_experts falls through to stacked mapping (no mutate_weight_preload) + ("model.layers.0.mlp.shared_experts.gate_proj.weight", to_expect("model.layers.0.mlp.shared_experts.gate_up_proj.weight", 0, 2)), +] + +_DEEPSEEK_CASES = [ + # MLA A-proj fusion + ("model.layers.0.self_attn.q_a_proj.weight", to_expect("model.layers.0.self_attn.fused_qkv_a_proj_with_mqa.weight", 0, 2)), + ("model.layers.0.self_attn.kv_a_proj_with_mqa.weight", to_expect("model.layers.0.self_attn.fused_qkv_a_proj_with_mqa.weight", 1, 2)), + # Shared expert fusion via mutate_weight_preload + ("model.layers.0.mlp.shared_experts.gate_proj.weight", to_expect("model.layers.0.mlp.experts.w13_weight", "w1", 2, _DEEPSEEK_N_ROUTED, _DEEPSEEK_N_LOCAL)), + ("model.layers.0.mlp.shared_experts.down_proj.weight", to_expect("model.layers.0.mlp.experts.w2_weight", "w2", 1, _DEEPSEEK_N_ROUTED, _DEEPSEEK_N_LOCAL)), + # Custom scale remap (k_proj/v_proj -> attn_mqa, NOT double-remapped) + ("model.layers.0.self_attn.k_proj.k_scale", to_expect("model.layers.0.self_attn.attn_mqa.k_scale")), + ("model.layers.0.self_attn.v_proj.v_scale", to_expect("model.layers.0.self_attn.attn_mqa.v_scale")), + # kv_b_proj pass-through (decomposed in post_load_weights) + ("model.layers.0.self_attn.kv_b_proj.weight", to_expect("model.layers.0.self_attn.kv_b_proj.weight")), +] + +_GLM4LITE_CASES = [ + # QKV fusion (GLM-4.7 uses standard QKV unlike DeepSeek which uses MLA-only) + ("model.layers.0.self_attn.q_proj.weight", to_expect("model.layers.0.self_attn.qkv_proj.weight", "q", 3)), + ("model.layers.0.self_attn.k_proj.weight", to_expect("model.layers.0.self_attn.qkv_proj.weight", "k", 3)), + ("model.layers.0.self_attn.v_proj.weight", to_expect("model.layers.0.self_attn.qkv_proj.weight", "v", 3)), + # MLA A-proj fusion (GLM-4.7 also uses MLA with q_lora_rank) + ("model.layers.0.self_attn.q_a_proj.weight", to_expect("model.layers.0.self_attn.fused_qkv_a_proj_with_mqa.weight", 0, 2)), + ("model.layers.0.self_attn.kv_a_proj_with_mqa.weight", to_expect("model.layers.0.self_attn.fused_qkv_a_proj_with_mqa.weight", 1, 2)), + # Gate/Up fusion (non-expert layers) + ("model.layers.0.mlp.gate_proj.weight", to_expect("model.layers.0.mlp.gate_up_proj.weight", 0, 2)), + ("model.layers.0.mlp.up_proj.weight", to_expect("model.layers.0.mlp.gate_up_proj.weight", 1, 2)), + # Expert mapping + ("model.layers.0.mlp.experts.0.gate_proj.weight", to_expect("model.layers.0.mlp.experts.w13_weight", "w1", 2, 0, _GLM4LITE_N_LOCAL)), + ("model.layers.0.mlp.experts.0.up_proj.weight", to_expect("model.layers.0.mlp.experts.w13_weight", "w3", 2, 0, _GLM4LITE_N_LOCAL)), + ("model.layers.0.mlp.experts.3.down_proj.weight", to_expect("model.layers.0.mlp.experts.w2_weight", "w2", 1, 3, _GLM4LITE_N_LOCAL)), + # Shared expert fusion via mutate_weight_preload + ("model.layers.0.mlp.shared_experts.gate_proj.weight", to_expect("model.layers.0.mlp.experts.w13_weight", "w1", 2, _GLM4LITE_N_ROUTED, _GLM4LITE_N_LOCAL)), + ("model.layers.0.mlp.shared_experts.down_proj.weight", to_expect("model.layers.0.mlp.experts.w2_weight", "w2", 1, _GLM4LITE_N_ROUTED, _GLM4LITE_N_LOCAL)), + # Custom scale remap (same as DeepSeek: k_proj/v_proj -> attn_mqa) + ("model.layers.0.self_attn.k_proj.k_scale", to_expect("model.layers.0.self_attn.attn_mqa.k_scale")), + ("model.layers.0.self_attn.v_proj.v_scale", to_expect("model.layers.0.self_attn.attn_mqa.v_scale")), + # Pass-through + ("model.layers.0.mlp.down_proj.weight", to_expect("model.layers.0.mlp.down_proj.weight")), + ("model.layers.0.self_attn.kv_b_proj.weight", to_expect("model.layers.0.self_attn.kv_b_proj.weight")), +] +# fmt: on + + +@pytest.mark.parametrize("ckpt,expected", _QWEN_CASES, ids=[c[0] for c in _QWEN_CASES]) +def test_qwen(qwen_mapper, ckpt, expected): + _assert(qwen_mapper, ckpt, expected) + + +@pytest.mark.parametrize( + "ckpt,expected", _LLAMA_CASES, ids=[c[0] for c in _LLAMA_CASES] +) +def test_llama(llama_mapper, ckpt, expected): + _assert(llama_mapper, ckpt, expected) + + +@pytest.mark.parametrize( + "ckpt,expected", _QWEN3MOE_CASES, ids=[c[0] for c in _QWEN3MOE_CASES] +) +def test_qwen3moe(qwen3moe_mapper, ckpt, expected): + _assert(qwen3moe_mapper, ckpt, expected) + + +@pytest.mark.parametrize( + "ckpt,expected", _DEEPSEEK_CASES, ids=[c[0] for c in _DEEPSEEK_CASES] +) +def test_deepseek(deepseek_mapper, ckpt, expected): + _assert(deepseek_mapper, ckpt, expected) + + +@pytest.mark.parametrize( + "ckpt,expected", _GLM4LITE_CASES, ids=[c[0] for c in _GLM4LITE_CASES] +) +def test_glm4lite(glm4lite_mapper, ckpt, expected): + _assert(glm4lite_mapper, ckpt, expected)