From d7ff74b9847ab51c98b6c34a73defe5b3e5ca16c Mon Sep 17 00:00:00 2001 From: Kangyan-Zhou Date: Tue, 29 Sep 2026 17:30:38 -0700 Subject: [PATCH 01/56] [Cherry-pick to release/v0.5.21] [Fix] Restore deferred layer dumps and pin SentencePiece for InternVL (#41777) (#41788) Co-authored-by: Baizhou Zhang --- python/pyproject.toml | 2 +- python/pyproject_cpu.toml | 2 +- python/pyproject_npu.toml | 2 +- python/pyproject_other.toml | 2 +- python/pyproject_xpu.toml | 2 +- python/sglang/srt/debug_utils/dumper.py | 12 ++++++++++++ .../accuracy/models/test_vlms_mmmu_eval.py | 2 +- 7 files changed, 18 insertions(+), 6 deletions(-) diff --git a/python/pyproject.toml b/python/pyproject.toml index 0619b5899577..50394749d90b 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -74,7 +74,7 @@ dependencies = [ "regex", "requests", "scipy", - "sentencepiece", + "sentencepiece==0.2.1", # 0.2.2 rejects null pieces in InternVL tokenizers. "setproctitle", "sgl-deep-ep==0.1.2", "sgl-deep-gemm==0.2.0", diff --git a/python/pyproject_cpu.toml b/python/pyproject_cpu.toml index 6a9902a01b68..01b8f50d9472 100644 --- a/python/pyproject_cpu.toml +++ b/python/pyproject_cpu.toml @@ -54,7 +54,7 @@ dependencies = [ "regex", "requests", "scipy", - "sentencepiece", + "sentencepiece==0.2.1", # 0.2.2 rejects null pieces in InternVL tokenizers. "setproctitle", "smg-grpc-servicer>=0.9.0", "soundfile==0.13.1", diff --git a/python/pyproject_npu.toml b/python/pyproject_npu.toml index 8d097df8dc8f..47f76d6f3069 100644 --- a/python/pyproject_npu.toml +++ b/python/pyproject_npu.toml @@ -55,7 +55,7 @@ dependencies = [ "regex", "requests", "scipy", - "sentencepiece", + "sentencepiece==0.2.1", # 0.2.2 rejects null pieces in InternVL tokenizers. "setproctitle", "smg-grpc-servicer>=0.9.0", "soundfile==0.13.1", diff --git a/python/pyproject_other.toml b/python/pyproject_other.toml index 2eb25a40cdfe..4ca21b90a8bb 100755 --- a/python/pyproject_other.toml +++ b/python/pyproject_other.toml @@ -56,7 +56,7 @@ runtime_base = [ "regex", "requests", "scipy", - "sentencepiece", + "sentencepiece==0.2.1", # 0.2.2 rejects null pieces in InternVL tokenizers. "setproctitle", "smg-grpc-servicer>=0.9.0", "soundfile==0.13.1", diff --git a/python/pyproject_xpu.toml b/python/pyproject_xpu.toml index 4239621ab214..79fc47541c85 100644 --- a/python/pyproject_xpu.toml +++ b/python/pyproject_xpu.toml @@ -54,7 +54,7 @@ dependencies = [ "regex", "requests", "scipy", - "sentencepiece", + "sentencepiece==0.2.1", # 0.2.2 rejects null pieces in InternVL tokenizers. "setproctitle", "sglang-kernel-xpu @ https://github.com/sgl-project/whl/releases/download/v0.3.0+xpu/sglang_kernel_xpu-0.3.0-cp310-abi3-manylinux_2_38_x86_64.whl", "smg-grpc-servicer>=0.9.0", diff --git a/python/sglang/srt/debug_utils/dumper.py b/python/sglang/srt/debug_utils/dumper.py index 4d788b50596a..6bffc95e6ed8 100644 --- a/python/sglang/srt/debug_utils/dumper.py +++ b/python/sglang/srt/debug_utils/dumper.py @@ -1780,6 +1780,18 @@ def convert_value( if isinstance(value, self.PPProxyTensors): return {k: v for k, v in value.tensors.items()} + from sglang.srt.layers.layer_boundary.output import UnreducedOutput + from sglang.srt.layers.layer_boundary.residual.stream import OwedOutput + + # Observe rank-local storage without completing the deferred reduction or + # consuming the residual stream. Opaque handoffs and consumed handles have + # no tensor to dump. + if isinstance(value, UnreducedOutput): + return {"": value.partial} + if isinstance(value, OwedOutput): + tensor = value.contribution.value + return {"": tensor} if tensor is not None else {} + return None def detect_layer_id(self, module: "torch.nn.Module") -> Optional[int]: diff --git a/test/registered/accuracy/models/test_vlms_mmmu_eval.py b/test/registered/accuracy/models/test_vlms_mmmu_eval.py index 48dd51745751..ef845e6a4fae 100644 --- a/test/registered/accuracy/models/test_vlms_mmmu_eval.py +++ b/test/registered/accuracy/models/test_vlms_mmmu_eval.py @@ -18,7 +18,7 @@ # Use a longer timeout than the default 600s. NIGHTLY_EVAL_SERVER_TIMEOUT = 1800 -register_cuda_ci(est_time=7200, stage="nightly", runner_config="2-gpu-large") +register_cuda_ci(est_time=3600, stage="nightly", runner_config="2-gpu-large") MODEL_THRESHOLDS = { # Conservative thresholds on 100 MMMU samples. Latency baselines account for From e00930c5489053f26d86b179cee0d087f846acbb Mon Sep 17 00:00:00 2001 From: Kangyan-Zhou Date: Tue, 29 Sep 2026 19:44:12 -0700 Subject: [PATCH 02/56] [Cherry-pick to release/v0.5.21] [Test] Demote PD test RDMA openability check to a warning (#41681) (#41802) Co-authored-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Co-authored-by: hnyls2002 --- .../server_fixtures/disaggregation_fixture.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/python/sglang/test/server_fixtures/disaggregation_fixture.py b/python/sglang/test/server_fixtures/disaggregation_fixture.py index 2eab0c3e7b82..16b543b7d369 100644 --- a/python/sglang/test/server_fixtures/disaggregation_fixture.py +++ b/python/sglang/test/server_fixtures/disaggregation_fixture.py @@ -97,7 +97,7 @@ def setUpClass(cls): cls._ucx_net_devices_set = False if is_in_ci(): cls.transfer_backend = ["--disaggregation-transfer-backend", "mooncake"] - _check_rdma_openable() + _warn_if_rdma_unopenable() ib_devices = get_rdma_devices_args() cls.rdma_devices = ["--disaggregation-ib-device", ib_devices] cls._mc_gid_index_set = _maybe_set_roce_gid_index(ib_devices) @@ -428,15 +428,16 @@ def _ibverbs_device_names() -> Optional[set]: # Without an openable device the PD transfer backend silently falls back to TCP, -# and transfers then fail under load instead of at setup. -def _check_rdma_openable() -> None: +# and transfers may then fail under load; warn so such failures are traceable. +def _warn_if_rdma_unopenable() -> None: if not _get_available_ib_devices(): return if _ibverbs_device_names() == set(): - raise RuntimeError( - f"Active RDMA devices are listed in {_IB_SYSFS} but ibverbs can open " - "none; the PD transfer backend would lose RDMA. " - "Check /dev/infiniband in the container." + logger.warning( + "Active RDMA devices are listed in %s but ibverbs can open none; " + "the PD transfer backend will fall back to TCP. " + "Check /dev/infiniband in the container.", + _IB_SYSFS, ) From 2beb575082d41030c8a9e7552b1bbc3d472e487c Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Tue, 10 Feb 2026 20:28:00 +0000 Subject: [PATCH 03/56] [1/32] [sglang-miles] True on-policy training and tuple-aware row-parallel math (#18639, #30742) Co-authored-by: Xinyu Jiang --- .../sglang/srt/arg_groups/attention_hook.py | 45 +- python/sglang/srt/arg_groups/choices.py | 2 +- python/sglang/srt/arg_groups/fields/exec_.py | 9 + python/sglang/srt/arg_groups/overrides.py | 12 +- .../srt/distributed/communication_op.py | 8 + .../sglang/srt/distributed/parallel_state.py | 5 +- python/sglang/srt/layers/activation.py | 8 +- python/sglang/srt/layers/attention/vision.py | 9 +- .../sglang/srt/layers/layer_boundary/exit.py | 6 +- .../srt/layers/layer_boundary/factories.py | 14 +- python/sglang/srt/layers/layernorm.py | 62 +- python/sglang/srt/layers/linear.py | 23 +- python/sglang/srt/layers/logits_processor.py | 13 +- .../moe/moe_runner/triton_utils/fused_moe.py | 6 +- python/sglang/srt/layers/on_policy_utils.py | 6 + .../srt/layers/rotary_embedding/base.py | 12 +- python/sglang/srt/layers/sampler.py | 25 +- .../srt/model_executor/forward_batch_info.py | 13 +- .../sglang/srt/model_executor/model_runner.py | 5 + .../runner/decode_cuda_graph_runner.py | 40 +- python/sglang/srt/models/qwen2.py | 32 +- python/sglang/srt/models/qwen2_moe.py | 6 +- python/sglang/srt/models/qwen3.py | 64 +- python/sglang/srt/models/qwen3_moe.py | 38 +- python/sglang/srt/models/qwen3_vl.py | 164 +- python/sglang/srt/models/sdar.py | 36 +- python/sglang/srt/models/sdar_moe.py | 34 +- python/sglang/srt/models/step3p5.py | 14 +- python/sglang/srt/models/utils.py | 2 + .../multimodal/processors/base_processor.py | 13 +- .../sglang/srt/tp_invariant_ops/__init__.py | 23 + .../srt/tp_invariant_ops/tp_invariant_ops.py | 1939 +++++++++++++++++ python/sglang/srt/true_on_policy/__init__.py | 49 + python/sglang/srt/true_on_policy/config.py | 150 ++ python/sglang/srt/true_on_policy/contracts.py | 111 + python/sglang/srt/true_on_policy/schema.py | 33 + test/manual/layers/test_layernorm.py | 20 + .../true_on_policy/test_tp_invariant_ops.py | 866 ++++++++ .../test_dense_deterministic_math.py | 344 +++ .../true_on_policy/test_on_policy_wiring.py | 589 +++++ 40 files changed, 4550 insertions(+), 300 deletions(-) create mode 100644 python/sglang/srt/layers/on_policy_utils.py create mode 100644 python/sglang/srt/tp_invariant_ops/__init__.py create mode 100644 python/sglang/srt/tp_invariant_ops/tp_invariant_ops.py create mode 100644 python/sglang/srt/true_on_policy/__init__.py create mode 100644 python/sglang/srt/true_on_policy/config.py create mode 100644 python/sglang/srt/true_on_policy/contracts.py create mode 100644 python/sglang/srt/true_on_policy/schema.py create mode 100644 test/registered/e2e/true_on_policy/test_tp_invariant_ops.py create mode 100644 test/registered/unit/true_on_policy/test_dense_deterministic_math.py create mode 100644 test/registered/unit/true_on_policy/test_on_policy_wiring.py diff --git a/python/sglang/srt/arg_groups/attention_hook.py b/python/sglang/srt/arg_groups/attention_hook.py index b5ef4ed97f1e..5f09cf97021b 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, ) @@ -494,20 +498,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 809e33ff193b..e9a818642298 100644 --- a/python/sglang/srt/arg_groups/choices.py +++ b/python/sglang/srt/arg_groups/choices.py @@ -183,7 +183,7 @@ RADIX_EVICTION_POLICY_CHOICES = ["lru", "lfu", "slru", "priority", "tlru"] -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/fields/exec_.py b/python/sglang/srt/arg_groups/fields/exec_.py index 21afe7f02b3c..c680bf55cbc0 100644 --- a/python/sglang/srt/arg_groups/fields/exec_.py +++ b/python/sglang/srt/arg_groups/fields/exec_.py @@ -1010,6 +1010,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/overrides.py b/python/sglang/srt/arg_groups/overrides.py index a2de0bc2937b..dd41e5863563 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -1305,7 +1305,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 {} @@ -1362,8 +1367,9 @@ def _tp_lm_head_all_to_all_default(view: Any) -> dict: def _dp_lm_head_validation(view: Any) -> dict: """Require DP attention for DP LM head and TP LM-head all-to-all.""" 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/distributed/communication_op.py b/python/sglang/srt/distributed/communication_op.py index df7170c40c61..49950986acde 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_) @@ -119,6 +123,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/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 37c6732744d8..a014e28a14e4 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -3176,7 +3176,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 diff --git a/python/sglang/srt/layers/activation.py b/python/sglang/srt/layers/activation.py index 7c1bf82fed1f..d31e50b71052 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/vision.py b/python/sglang/srt/layers/attention/vision.py index 693963080ea5..e37f9abdd216 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, @@ -1253,7 +1253,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( @@ -1395,10 +1395,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/layer_boundary/exit.py b/python/sglang/srt/layers/layer_boundary/exit.py index 9528802f2f69..99be211205d2 100644 --- a/python/sglang/srt/layers/layer_boundary/exit.py +++ b/python/sglang/srt/layers/layer_boundary/exit.py @@ -65,6 +65,9 @@ get_lora, get_parallel, ) +from sglang.srt.true_on_policy import ( + should_disable_mlp_allreduce_fusion_for_on_policy, +) def _reduce_and_redistribute_output_step( @@ -188,7 +191,8 @@ def _select_ffn_completion( steps=steps, ) defer_moe_finalize = ( - self.plan.fusions is not None + not should_disable_mlp_allreduce_fusion_for_on_policy() + and self.plan.fusions is not None and self.plan.fusions.can_defer_finalize(self.plan, forward_batch) ) if not steps.output.leaves_for_next_layer and not ( diff --git a/python/sglang/srt/layers/layer_boundary/factories.py b/python/sglang/srt/layers/layer_boundary/factories.py index 5efd7ce9dff0..cbedb6ff3ecd 100644 --- a/python/sglang/srt/layers/layer_boundary/factories.py +++ b/python/sglang/srt/layers/layer_boundary/factories.py @@ -49,6 +49,10 @@ ) from sglang.srt.layers.moe import is_moe_input_scattered_across_dp_ranks from sglang.srt.runtime_context import get_exec, get_parallel +from sglang.srt.true_on_policy import ( + should_disable_mlp_allreduce_fusion_for_on_policy, + should_disable_reduce_scatter_for_on_policy, +) def _variants(): @@ -92,8 +96,13 @@ def _ffn_decl( "boundary_reduction must be resolved before model construction" ) can_move_output = output_transform is None or output_transform.before_reduce_scatter - use_reduce_scatter = strategy in ("rs", "rs+rsv") and can_move_output - use_reduce_scatterv = strategy in ("rsv", "rs+rsv") and can_move_output + allow_reduce_scatter = not should_disable_reduce_scatter_for_on_policy() + use_reduce_scatter = ( + strategy in ("rs", "rs+rsv") and can_move_output and allow_reduce_scatter + ) + use_reduce_scatterv = ( + strategy in ("rsv", "rs+rsv") and can_move_output and allow_reduce_scatter + ) cp_shards = _generic_prefill_cp_shards_tokens() axes, attention, local, full = _rows(variant) on_local = ( @@ -165,6 +174,7 @@ def _ffn_decl( rows, group=group, leaves_for_next_layer=may_leave + and not should_disable_mlp_allreduce_fusion_for_on_policy() and not terminal and not update.at_producer and update.can_defer_across_layers diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index 6a4d049b653e..2d3367d1d4ff 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, @@ -228,7 +232,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: @@ -455,14 +459,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 @@ -514,11 +537,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 @@ -721,10 +750,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, @@ -771,10 +800,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 ) @@ -817,15 +846,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: @@ -1136,7 +1168,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 4504bbcb74a5..611aad4b27ed 100644 --- a/python/sglang/srt/layers/linear.py +++ b/python/sglang/srt/layers/linear.py @@ -40,6 +40,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: @@ -1650,7 +1651,27 @@ def forward( get_parallel().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 77b47df9c9d3..70b28101b46b 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -63,6 +63,7 @@ get_forward, get_parallel, ) +from sglang.srt.true_on_policy import should_force_bfloat16_lm_head from sglang.srt.utils.common import ( is_cpu, is_npu, @@ -986,6 +987,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), @@ -993,11 +999,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 @@ -1154,6 +1155,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/moe_runner/triton_utils/fused_moe.py b/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py index d4baddfe6877..1c7163643d18 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 @@ -126,7 +126,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/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/rotary_embedding/base.py b/python/sglang/srt/layers/rotary_embedding/base.py index 486a9d606236..ee699dd7709f 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, @@ -104,10 +105,7 @@ def __init__( self.is_neox_style = is_neox_style self.dtype = dtype self.position_start = position_start - 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. @@ -145,10 +143,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/sampler.py b/python/sglang/srt/layers/sampler.py index c1dae6ef24c7..e77b853de019 100644 --- a/python/sglang/srt/layers/sampler.py +++ b/python/sglang/srt/layers/sampler.py @@ -22,6 +22,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, @@ -117,13 +118,13 @@ def __init__(self): if get_parallel().attn_cp_size > 1: 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 @@ -209,17 +210,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. @@ -245,8 +242,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 @@ -658,13 +653,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/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index ee2e5d2bf056..5f70842508f8 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -61,6 +61,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, @@ -1451,7 +1452,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) @@ -1512,7 +1517,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 604a72f4ed7f..9033087aada9 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -217,6 +217,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, @@ -773,6 +774,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_dspark_hidden_size(self) -> int: return misc_utils.resolve_pp_proxy_dspark_hidden_size( 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 a7b73754114c..2640b0257c23 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 @@ -115,6 +115,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, @@ -1061,26 +1064,33 @@ 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: - 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: + 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/models/qwen2.py b/python/sglang/srt/models/qwen2.py index 588aecf7939a..5a37f29672d8 100644 --- a/python/sglang/srt/models/qwen2.py +++ b/python/sglang/srt/models/qwen2.py @@ -54,7 +54,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 @@ -99,9 +104,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) @@ -337,11 +344,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() @@ -368,16 +371,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 ) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 8755df797ef8..4b3301512b41 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -99,6 +99,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, @@ -1152,7 +1153,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 2939fc608d82..cf2bc13901b6 100644 --- a/python/sglang/srt/models/qwen3.py +++ b/python/sglang/srt/models/qwen3.py @@ -37,7 +37,11 @@ from sglang.srt.models.qwen2 import Qwen2Model from sglang.srt.models.utils import apply_qk_norm from sglang.srt.platforms import current_platform -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 Qwen3Config = None @@ -110,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, @@ -275,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: @@ -302,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) @@ -359,21 +369,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.attn_boundary, self.ffn_boundary = make_stages( diff --git a/python/sglang/srt/models/qwen3_moe.py b/python/sglang/srt/models/qwen3_moe.py index d99f61687b5a..e00e9363e3d9 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 @@ -49,7 +50,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, TopKOutputChecker +from sglang.srt.layers.moe.topk import StandardTopKOutput, TopK, TopKOutputChecker from sglang.srt.layers.moe.utils import ( RoutingMethodType, filter_moe_weight_param_global_expert, @@ -70,6 +71,11 @@ enable_fused_set_kv_buffer, ) from sglang.srt.runtime_context import get_exec, 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, @@ -338,7 +344,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) @@ -570,7 +589,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) @@ -585,6 +604,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 @@ -603,8 +623,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): @@ -861,9 +882,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.attn_boundary, self.ffn_boundary = make_stages( diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py index 6f9faffd5684..54183b082e9b 100644 --- a/python/sglang/srt/models/qwen3_vl.py +++ b/python/sglang/srt/models/qwen3_vl.py @@ -82,6 +82,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, @@ -502,75 +503,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 @@ -803,61 +802,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, @@ -1175,9 +1119,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 6e8d510e0e32..0059e51f87df 100644 --- a/python/sglang/srt/models/sdar.py +++ b/python/sglang/srt/models/sdar.py @@ -45,7 +45,11 @@ create_fused_set_kv_buffer_arg, enable_fused_set_kv_buffer, ) -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 ( + 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__) @@ -206,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) @@ -234,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) @@ -262,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 @@ -372,15 +371,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 f4cbd6a29bfd..fe92fc168f4d 100644 --- a/python/sglang/srt/models/sdar_moe.py +++ b/python/sglang/srt/models/sdar_moe.py @@ -59,6 +59,10 @@ enable_fused_set_kv_buffer, ) from sglang.srt.runtime_context import get_exec, 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__) @@ -267,7 +271,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) @@ -295,7 +299,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) @@ -324,15 +328,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 @@ -438,15 +437,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 c640e094dd65..475ffa2a2433 100644 --- a/python/sglang/srt/models/step3p5.py +++ b/python/sglang/srt/models/step3p5.py @@ -50,7 +50,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 @@ -647,11 +653,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 7b9e979f5ff2..6ad95674bc34 100644 --- a/python/sglang/srt/multimodal/processors/base_processor.py +++ b/python/sglang/srt/multimodal/processors/base_processor.py @@ -49,6 +49,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, @@ -731,7 +732,11 @@ def _preprocessing_competes_with_the_scheduler(self) -> bool: the same accelerator as the scheduler. A second preprocessing worker there adds device contention rather than CPU preprocessing 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 @@ -772,7 +777,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" 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/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/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/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..1aa6de2f425a --- /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 +from unittest.mock import patch + +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, + ) + + with patch( + "sglang.srt.true_on_policy.config._get_global_server_args", + return_value=server_args, + ): + self.assertFalse(should_force_bfloat16_dense_tensor_math()) + self.assertFalse( + should_force_bfloat16_lm_head( + use_fp32_lm_head=False, + ) + ) + self.assertEqual(get_on_policy_rms_norm_kwargs(), {}) + + 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, + ) + + with patch( + "sglang.srt.true_on_policy.config._get_global_server_args", + return_value=server_args, + ): + kwargs = get_on_policy_rms_norm_kwargs( + weight_dtype=torch.float32, + override_orig_dtype=torch.float32, + fp32_residual=True, + ) + + self.assertTrue(should_force_bfloat16_dense_tensor_math()) + self.assertTrue( + should_force_bfloat16_lm_head( + use_fp32_lm_head=False, + ) + ) + self.assertFalse( + should_force_bfloat16_lm_head( + 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 + from unittest.mock import patch + + 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, + ) + patch("sglang.srt.true_on_policy.config._get_global_server_args", return_value=server_args).start() + norm = RMSNorm( + 4, + eps=1e-6, + **get_on_policy_rms_norm_kwargs( + 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 + from sglang.srt.runtime_context import get_server_args, publish + from sglang.srt.true_on_policy import QWEN3_DENSE_TRUE_ON_POLICY_V1 + + publish(ServerArgs(model_path="dummy"), role="scheduler") + server_args = get_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 + from sglang.srt.runtime_context import get_server_args, publish + 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)) + + publish(ServerArgs(model_path="dummy"), role="scheduler") + get_server_args().enable_dp_lm_head = False + get_server_args().enable_fp32_lm_head = False + get_server_args().true_on_policy_contract = QWEN3_DENSE_TRUE_ON_POLICY_V1 + get_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..def4e1bcb827 --- /dev/null +++ b/test/registered/unit/true_on_policy/test_on_policy_wiring.py @@ -0,0 +1,589 @@ +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.true_on_policy.config._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.layer_boundary.prepare import _mlp_input_without_dp + from sglang.srt.model_executor.forward_batch_info import ForwardMode + + 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.layer_boundary.ops.attention_tensor_model_parallel_all_reduce", + side_effect=lambda x: x + 10.0, + ) as attn_tree_reduce: + output, output_residual = _mlp_input_without_dp( + hidden_states, + residual, + forward_batch=SimpleNamespace(forward_mode=ForwardMode.DECODE), + layernorm=FakeNorm(), + gathers_residual=False, + fusions=(), + ) + + 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() From 31d10bfd1040835978b92d751a031d2582b9a401 Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Wed, 11 Feb 2026 23:14:29 +0000 Subject: [PATCH 04/56] [2/32] [sglang-miles] R3 draft isolation and graph capture layout (#18642, #36404) Co-authored-by: Claude Opus 4.6 --- python/sglang/srt/layers/dp_attention.py | 13 ++++--- python/sglang/srt/layers/moe/topk.py | 24 +++++++++---- .../sglang/srt/model_executor/model_runner.py | 34 ++++++++++++++++--- .../runner/prefill_cuda_graph_runner.py | 3 ++ python/sglang/srt/state_capturer/base.py | 1 + 5 files changed, 60 insertions(+), 15 deletions(-) diff --git a/python/sglang/srt/layers/dp_attention.py b/python/sglang/srt/layers/dp_attention.py index bf6c1e8f6743..d2329e699aec 100644 --- a/python/sglang/srt/layers/dp_attention.py +++ b/python/sglang/srt/layers/dp_attention.py @@ -494,16 +494,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 = dp_slot_in(global_num_tokens) - 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/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index a4b6541f6441..60767eb65bbb 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -2208,19 +2208,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( @@ -2240,7 +2250,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 recorder_was_fused = False diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 9033087aada9..2f4b097d8943 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -280,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 @@ -924,7 +927,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 @@ -1772,6 +1776,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 @@ -1779,7 +1791,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, ) @@ -1787,7 +1799,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, ) @@ -1889,13 +1901,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 @@ -1945,6 +1964,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( @@ -1957,7 +1977,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, 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 1b506b826c7e..d429e95bf30c 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 @@ -2103,6 +2103,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/state_capturer/base.py b/python/sglang/srt/state_capturer/base.py index f451f7e6f5d2..dcfd7fd42e17 100644 --- a/python/sglang/srt/state_capturer/base.py +++ b/python/sglang/srt/state_capturer/base.py @@ -209,6 +209,7 @@ def on_forward_end( else: # Kernel batches may be padded; only real request tokens are stored. slice_gpu = slice_gpu[: out_cache_loc.shape[0]] + out_cache_loc = out_cache_loc[: slice_gpu.shape[0]] if no_copy_to_cpu: # Clone before the next overlapping forward reuses these buffers. return TopkCaptureOutput( From ea1fa7ea4b85674a41bb1ed46df9163d5039c056 Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Wed, 11 Feb 2026 23:58:12 +0000 Subject: [PATCH 05/56] [3/32] [sglang-miles] PD disaggregation for RL (#18646) Co-authored-by: Claude Opus 4.6 --- python/sglang/srt/managers/schedule_batch.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 52c72a0e24b2..a94a1df7ae7a 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -3305,8 +3305,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 From 6a24ca3b768c39961abcc9bbef5d652f55c3f83a Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Thu, 12 Feb 2026 00:15:26 +0000 Subject: [PATCH 06/56] [4/32] [sglang-miles] MTP related fix (#18647) Co-authored-by: Claude Opus 4.6 --- python/sglang/srt/arg_groups/fields/exec_.py | 3 +++ .../sglang/srt/speculative/eagle_draft_cuda_graph_runner.py | 6 ++++-- python/sglang/srt/speculative/eagle_worker_v2.py | 6 +++++- 3 files changed, 12 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/arg_groups/fields/exec_.py b/python/sglang/srt/arg_groups/fields/exec_.py index c680bf55cbc0..159d4b0a0b1c 100644 --- a/python/sglang/srt/arg_groups/fields/exec_.py +++ b/python/sglang/srt/arg_groups/fields/exec_.py @@ -543,6 +543,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.", 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 543690c11321..af198d2d5d37 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py @@ -594,8 +594,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 self.model_runner.model_config.model_is_mrope: diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 993b53f637c1..da7a0f32d741 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -458,7 +458,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 get_exec().graph.disable_draft_cuda_graph + ): return if get_model().model_impl == "mindspore": From 1b05a82e9990d80886dfdf6ddc46b87799655390 Mon Sep 17 00:00:00 2001 From: yueming-yuan Date: Fri, 13 Feb 2026 06:29:45 +0000 Subject: [PATCH 07/56] [5/32] [sglang-miles] VLM training multimodal fallback fixes (#18781) --- .../srt/multimodal/processors/qwen_vl.py | 32 ++++--------------- 1 file changed, 6 insertions(+), 26 deletions(-) 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 ) From 4a4cbe7aa79b6ffa236c631cb69ae2dd95d1679b Mon Sep 17 00:00:00 2001 From: maocheng23 Date: Sun, 12 Apr 2026 07:06:14 +0000 Subject: [PATCH 08/56] [6/32] [sglang-miles] Allow a ragged final merged FP8 matrix (#22623) Co-authored-by: Claude Opus 4.6 (1M context) --- python/sglang/srt/layers/quantization/fp8.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 25b39ae95e6a..5634a2617a6a 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -579,11 +579,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 = " From 01957d9104373ffb16bdedfc2b90ce96542b4252 Mon Sep 17 00:00:00 2001 From: zyzshishui Date: Thu, 16 Apr 2026 00:04:51 +0000 Subject: [PATCH 09/56] [7/32] [sglang-miles] R3 support on PD disaggregation mini_lb (#22916) --- python/sglang/srt/disaggregation/prefill.py | 4 +- .../bindings/python/pyproject.toml | 1 + .../python/src/sglang_router/mini_lb.py | 53 +++++++++++++++---- 3 files changed, 48 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index b72dc2a29622..9383cfb679e9 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -1056,9 +1056,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/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 3cb49fa53bf0..7734f8830209 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, @@ -175,17 +214,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() From 91f1fbed04692b10f0dd362c226a2d54612fcd34 Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Wed, 13 May 2026 18:09:12 +0000 Subject: [PATCH 10/56] [8/32] [sglang-miles] Add KimiK2 raw tool call id parser (#25196) --- .../srt/entrypoints/openai/serving_chat.py | 9 +++++++ python/sglang/srt/function_call/core_types.py | 6 +++++ .../srt/function_call/function_call_parser.py | 6 ++++- .../srt/function_call/kimik2_detector.py | 24 +++++++++++++++++++ 4 files changed, 44 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index 6845aa4936ef..5917a93bfa47 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -2537,6 +2537,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/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 3184ac7dc3e3..c2cf71c6be4c 100644 --- a/python/sglang/srt/function_call/function_call_parser.py +++ b/python/sglang/srt/function_call/function_call_parser.py @@ -37,7 +37,10 @@ from sglang.srt.function_call.internlm_detector import InternlmDetector from sglang.srt.function_call.iquest_q1_detector import IQuestQ1Detector 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 @@ -88,6 +91,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 7263893f1c3e..b26b95b47ece 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: @@ -473,3 +477,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 From f6df9fcb9105a3147134e25cfaadf8ec27fe45bc Mon Sep 17 00:00:00 2001 From: yueming-yuan Date: Sat, 22 Aug 2026 12:41:21 -0700 Subject: [PATCH 11/56] [9/32] [sglang-miles] Fix GLM4 MoE Lite shared expert TP flag --- python/sglang/srt/models/glm4_moe_lite.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/python/sglang/srt/models/glm4_moe_lite.py b/python/sglang/srt/models/glm4_moe_lite.py index e910e693da03..853115b7c391 100644 --- a/python/sglang/srt/models/glm4_moe_lite.py +++ b/python/sglang/srt/models/glm4_moe_lite.py @@ -244,6 +244,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, @@ -251,14 +256,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" ) From 0584e0eeb4ef29a1500e3e8f8083f63fa7b94897 Mon Sep 17 00:00:00 2001 From: Nan Date: Sat, 23 May 2026 21:53:23 +0000 Subject: [PATCH 12/56] [10/32] [sglang-miles] MoE LoRA support and expert-target validation (#25141, #29874, #31251, #32376) Co-authored-by: Yusheng Su --- .../sglang/kernels/ops/moe/virtual_experts.py | 154 ++++++++++++------ python/sglang/srt/configs/model_config.py | 5 +- .../sglang/srt/lora/backend/base_backend.py | 2 +- python/sglang/srt/lora/layers.py | 19 +++ python/sglang/srt/lora/lora_manager.py | 35 +++- python/sglang/srt/lora/lora_moe_runners.py | 7 +- python/sglang/srt/lora/mem_pool.py | 40 ++++- 7 files changed, 198 insertions(+), 64 deletions(-) diff --git a/python/sglang/kernels/ops/moe/virtual_experts.py b/python/sglang/kernels/ops/moe/virtual_experts.py index e8b35e4465d5..6267c6db4695 100644 --- a/python/sglang/kernels/ops/moe/virtual_experts.py +++ b/python/sglang/kernels/ops/moe/virtual_experts.py @@ -517,7 +517,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 +526,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 +633,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,12 +663,22 @@ 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)" + ) intermediate = torch.zeros( - [token_lora_mapping.shape[0], topk_ids.shape[1], max_lora_rank], + [topk_ids.shape[0], topk_ids.shape[1], max_lora_rank], dtype=hidden_states.dtype, device=hidden_states.device, ) @@ -680,7 +709,7 @@ def _get_routing( a_stage_config, ) - b_stage_config = _get_stage_config(lora_b_virtual, 1) + b_stage_config = _get_stage_config(lora_b_virtuals[0], 1) ( sorted_token_ids, expert_ids, @@ -694,33 +723,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 +812,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 +821,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/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index b7d31afaef77..77ab4a863fd3 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -391,7 +391,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/lora/backend/base_backend.py b/python/sglang/srt/lora/backend/base_backend.py index 455e9f3fe4a9..911ac199394d 100644 --- a/python/sglang/srt/lora/backend/base_backend.py +++ b/python/sglang/srt/lora/backend/base_backend.py @@ -266,7 +266,7 @@ def init_cuda_graph_moe_buffers( base = moe_layer.base_layer top_k = base.top_k - device = moe_layer._quant_info.w13_weight.device + device = next(base.parameters()).device num_experts = base.num_experts block_size_m = 64 diff --git a/python/sglang/srt/lora/layers.py b/python/sglang/srt/lora/layers.py index 6ce0205d4c9e..3ebb9b550a9c 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) @@ -990,6 +997,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,6 +1016,8 @@ 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 @@ -1193,6 +1209,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_manager.py b/python/sglang/srt/lora/lora_manager.py index dae118d8082b..7a00431ea0fa 100644 --- a/python/sglang/srt/lora/lora_manager.py +++ b/python/sglang/srt/lora/lora_manager.py @@ -1039,6 +1039,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, @@ -1130,9 +1132,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_runners.py b/python/sglang/srt/lora/lora_moe_runners.py index 8b1d3e5e4533..1c8b383807bd 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) @@ -338,6 +337,8 @@ def _add_lora_gate_up_delta( if is_gated: inter_size = gate_up_b.shape[2] // 2 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 +352,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, @@ -418,6 +419,8 @@ def _add_lora_down_delta( down_lora_a = lora_info.down_lora_a_weights down_lora_b = lora_info.down_lora_b_weights 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/mem_pool.py b/python/sglang/srt/lora/mem_pool.py index 506a469c3498..f556f57faa38 100644 --- a/python/sglang/srt/lora/mem_pool.py +++ b/python/sglang/srt/lora/mem_pool.py @@ -1018,19 +1018,45 @@ def load_lora_weight_tensor( temp_B_buffer[target_module] = weights temp_B_cache_keys[target_module] = 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 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 elif "experts" in name and weights.dim() == 3: From d8744423fd751b22d73c8cfdf890c1209e7bea40 Mon Sep 17 00:00:00 2001 From: yueming-yuan Date: Fri, 29 May 2026 19:38:35 +0000 Subject: [PATCH 13/56] [11/32] [sglang-miles] DeepSeek V4 RL fixes (#27131, #27603, #27604, #27728) --- .../layers/deep_gemm_wrapper/compile_utils.py | 8 ++++- python/sglang/srt/layers/linear.py | 32 +++++++++++++++++++ python/sglang/srt/models/deepseek_v4.py | 4 ++- 3 files changed, 42 insertions(+), 2 deletions(-) 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 b1f282d90ab9..6f3d03f6f427 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/linear.py b/python/sglang/srt/layers/linear.py index 611aad4b27ed..6186005f1707 100644 --- a/python/sglang/srt/layers/linear.py +++ b/python/sglang/srt/layers/linear.py @@ -1587,6 +1587,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 +1618,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_, diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index c28eb599fe04..e4d06703faf3 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -5742,7 +5742,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) From 18ce4146b61b7b0d17cadb1bc81ec18c6201744c Mon Sep 17 00:00:00 2001 From: gongyisheng Date: Thu, 4 Jun 2026 07:05:15 +0000 Subject: [PATCH 14/56] [12/32] [sglang-miles] Streamed LoRA updates, upsert and deferred publication (#27268, #31759, #30913, #30912, #33299, #36384, #38165) Register adapter identities, stream their tensors within weight-update sessions, and publish checksum-verified fresh versions alongside generation. Preserve per-rank failure cleanup and request-prefix cancellation for adapter lifecycle management. Co-authored-by: Mathew Han Co-authored-by: Ethan (Yusheng) Su Co-authored-by: Yueming Yuan --- python/sglang/srt/entrypoints/EngineBase.py | 4 +- python/sglang/srt/entrypoints/engine.py | 53 +-- python/sglang/srt/entrypoints/http_server.py | 17 +- .../srt/entrypoints/http_server_engine.py | 14 +- python/sglang/srt/lora/lora_manager.py | 195 ++++++-- python/sglang/srt/lora/lora_registry.py | 93 +++- python/sglang/srt/managers/io_struct.py | 63 ++- python/sglang/srt/managers/scheduler.py | 27 +- .../scheduler_components/weight_updater.py | 202 +++++++- .../srt/managers/tokenizer_control_mixin.py | 359 +++++++++----- .../sglang/srt/managers/tokenizer_manager.py | 146 +++++- python/sglang/srt/managers/tp_worker.py | 57 +-- .../sglang/srt/model_executor/model_runner.py | 9 +- .../rl/test_lora_load_from_tensor.py | 156 +++--- test/registered/unit/lora/test_lora_lease.py | 233 +++++++++ .../unit/lora/test_lora_staged_publication.py | 314 ++++++++++++ .../unit/lora/test_lora_stream_update.py | 301 ++++++++++++ test/registered/unit/lora/test_lora_upsert.py | 352 ++++++++++++++ .../managers/test_abort_request_prefix.py | 450 ++++++++++++++++++ .../managers/test_lora_update_result_merge.py | 19 +- .../managers/test_msgpack_ipc_roundtrip.py | 7 - 21 files changed, 2655 insertions(+), 416 deletions(-) create mode 100644 test/registered/unit/lora/test_lora_lease.py create mode 100644 test/registered/unit/lora/test_lora_staged_publication.py create mode 100644 test/registered/unit/lora/test_lora_stream_update.py create mode 100644 test/registered/unit/lora/test_lora_upsert.py create mode 100644 test/registered/unit/managers/test_abort_request_prefix.py diff --git a/python/sglang/srt/entrypoints/EngineBase.py b/python/sglang/srt/entrypoints/EngineBase.py index f4af4df5eae8..653f0832f65a 100644 --- a/python/sglang/srt/entrypoints/EngineBase.py +++ b/python/sglang/srt/entrypoints/EngineBase.py @@ -46,12 +46,12 @@ def flush_cache(self): pass @abstractmethod - def begin_weight_update(self, selector: str = "all"): + def begin_weight_update(self, selector: str = "all", sync_base: bool = True): """Open a weight-update session; update_weights_from_* must run inside one.""" pass @abstractmethod - def end_weight_update(self): + def end_weight_update(self, expected_lora_checksums: Optional[Dict] = None): """Close the weight-update session and finalize quantized weights.""" pass diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index 4df81823346e..231c480b1794 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -77,12 +77,12 @@ GenerateReqInput, GetWeightsByNameReqInput, InitWeightsUpdateGroupReqInput, - LoadLoRAAdapterFromTensorsReqInput, LoadLoRAAdapterReqInput, MultimodalDataInputFormat, OpenSessionReqInput, ProfileReq, ProfileReqType, + RegisterLoRAAdapterReqInput, ReleaseMemoryOccupationReqInput, ResumeMemoryOccupationReqInput, ReturnHiddenStatesMode, @@ -1553,20 +1553,39 @@ def destroy_weights_update_group( self.tokenizer_manager.destroy_weights_update_group(obj, None) ) - def begin_weight_update(self, selector: str = "all"): - """Open a weight-update session; close it with end_weight_update().""" - obj = BeginWeightUpdateReqInput(selector=selector) + 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): - """Close the session and finalize quantized weights into kernel layout.""" - obj = EndWeightUpdateReqInput() + 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], @@ -1671,26 +1690,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/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 3258cfb9e974..72651ba43b51 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -135,13 +135,13 @@ GetWeightsByNameReqInput, InitWeightsSendGroupForRemoteInstanceReqInput, InitWeightsUpdateGroupReqInput, - LoadLoRAAdapterFromTensorsReqInput, LoadLoRAAdapterReqInput, OpenSessionReqInput, ParseFunctionCallReq, PauseGenerationReqInput, PdRoleSwitchReqInput, ProfileReq, + RegisterLoRAAdapterReqInput, ReleaseMemoryOccupationReqInput, ResumeMemoryOccupationReqInput, SendWeightsToRemoteInstanceReqInput, @@ -1602,14 +1602,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) @@ -1665,7 +1664,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/http_server_engine.py b/python/sglang/srt/entrypoints/http_server_engine.py index 1d4334451353..0974440fc9d0 100644 --- a/python/sglang/srt/entrypoints/http_server_engine.py +++ b/python/sglang/srt/entrypoints/http_server_engine.py @@ -1,6 +1,6 @@ import multiprocessing import time -from typing import List, Optional, Tuple +from typing import Dict, List, Optional, Tuple import requests import torch @@ -86,11 +86,15 @@ def _make_request(self, endpoint: str, payload: Optional[dict] = None): response.raise_for_status() return response.json() - def begin_weight_update(self, selector: str = "all"): - return self._make_request("begin_weight_update", {"selector": selector}) + def begin_weight_update(self, selector: str = "all", sync_base: bool = True): + return self._make_request( + "begin_weight_update", {"selector": selector, "sync_base": sync_base} + ) - def end_weight_update(self): - return self._make_request("end_weight_update") + def end_weight_update(self, expected_lora_checksums: Optional[Dict] = None): + return self._make_request( + "end_weight_update", {"expected_lora_checksums": expected_lora_checksums} + ) def update_weights_from_tensor( self, diff --git a/python/sglang/srt/lora/lora_manager.py b/python/sglang/srt/lora/lora_manager.py index 7a00431ea0fa..626d537e11a5 100644 --- a/python/sglang/srt/lora/lora_manager.py +++ b/python/sglang/srt/lora/lora_manager.py @@ -329,9 +329,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. """ assert not self.enable_dp_attention or lora_ref.pinned, ( "DP-attention LoRA requires pinned adapters" @@ -346,18 +356,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) @@ -370,7 +381,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 " @@ -934,9 +951,18 @@ 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, @@ -949,18 +975,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 @@ -971,36 +1029,111 @@ 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), ) + self.configs[uid] = new_config + self.loras[uid] = new_lora + + if ( + 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.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): diff --git a/python/sglang/srt/lora/lora_registry.py b/python/sglang/srt/lora/lora_registry.py index ed7101b9553c..88f2309d9b29 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/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 9629b7aab621..6be95c8bb1a6 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -274,6 +274,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 @@ -1217,6 +1219,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 @@ -2097,6 +2101,10 @@ class BeginWeightUpdateReqInput(BaseReq, kw_only=True): """Open a weight-update session: restore in-place-packed weights so new ones can load.""" 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): @@ -2105,7 +2113,15 @@ class BeginWeightUpdateReqOutput(BaseReq, kw_only=True): class EndWeightUpdateReqInput(BaseReq, kw_only=True): - """Close the weight-update session opened by BeginWeightUpdateReqInput.""" + """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): @@ -2186,6 +2202,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 @@ -2428,27 +2451,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, ) @@ -2456,11 +2487,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/scheduler.py b/python/sglang/srt/managers/scheduler.py index a7d8baeecc37..a0b4e0437c3e 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -154,8 +154,6 @@ InitWeightsUpdateGroupReqInput, ListExternalCorporaReqInput, ListExternalCorporaReqOutput, - LoadLoRAAdapterFromTensorsReqInput, - LoadLoRAAdapterFromTensorsReqOutput, LoadLoRAAdapterReqInput, LoadLoRAAdapterReqOutput, MMInputsProcessError, @@ -163,6 +161,8 @@ PauseGenerationReqInput, PdRoleSwitchReqInput, ProfileReq, + RegisterLoRAAdapterReqInput, + RegisterLoRAAdapterReqOutput, ReleaseMemoryOccupationReqInput, RemoveExternalCorpusReqInput, RemoveExternalCorpusReqOutput, @@ -1802,10 +1802,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), @@ -5788,20 +5785,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/weight_updater.py b/python/sglang/srt/managers/scheduler_components/weight_updater.py index d9a11737328e..9a4237c908bb 100644 --- a/python/sglang/srt/managers/scheduler_components/weight_updater.py +++ b/python/sglang/srt/managers/scheduler_components/weight_updater.py @@ -1,5 +1,6 @@ from __future__ import annotations +import hashlib import logging import time import traceback @@ -44,12 +45,40 @@ 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.weight_checker import overall_checksum +from sglang.srt.weight_sync.tensor_bucket import FlattenedTensorBucket logger = logging.getLogger(__name__) +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 = [] @@ -82,6 +111,8 @@ class _WeightUpdateSession(msgspec.Struct, frozen=True): # recorded at begin so end finalizes the same runners selector: str loaded_weights: bool = False + sync_base: bool = True + pending_version: Optional[str] = None @dataclass(kw_only=True, slots=True) @@ -98,6 +129,8 @@ class SchedulerWeightUpdaterManager: stashed_model_static_state: Any = None # replicated on every TP rank, so a rejected call returns on all ranks before any barrier _session: Optional[_WeightUpdateSession] = None + _lora_stash: Dict[str, Dict[str, torch.Tensor]] = field(default_factory=dict) + _lora_applied_names: Dict[str, frozenset] = field(default_factory=dict) @contextmanager def _observe_weight_load(self, source: str) -> Iterator[None]: @@ -122,6 +155,12 @@ 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._session is not None: + if weight_version is not None: + self._session = msgspec.structs.replace( + self._session, pending_version=weight_version + ) + return self.scheduler.record_weight_version_change(new_version=weight_version) def update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput): @@ -193,21 +232,31 @@ def update_weights_from_distributed( group_name=recv_req.group_name, load_format=recv_req.load_format, ) + base_tensors, lora_tensors = _split_lora_named_tensors(weights) + if base_tensors: + assert self._session.sync_base, ( + "base tensors arrived in a sync_base=False weight-update session" + ) + self._stash_lora_tensors(lora_tensors) except Exception as e: success, message = False, f"Failed to receive weights: {e}" logger.error(message) else: success, message = True, "Succeeded to update parameter online." - for _, runner in self._select_runners(recv_req.selector): - success, message = ( - runner.weight_updater.load_weights_from_distributed(weights) - ) - if not success: - break + if base_tensors: + for _, runner in self._select_runners(recv_req.selector): + success, message = ( + runner.weight_updater.load_weights_from_distributed( + base_tensors + ) + ) + if not success: + break if success: - self._session = msgspec.structs.replace( - self._session, loaded_weights=True - ) + if base_tensors: + self._session = msgspec.structs.replace( + self._session, loaded_weights=True + ) self.flush_cache_after_weight_update(recv_req) self.record_weight_version_after_update(recv_req.weight_version) return UpdateWeightsFromDistributedReqOutput( @@ -226,18 +275,39 @@ def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput): named_tensors = self.tp_worker.deserialize_own_rank( recv_req.serialized_named_tensors ) + 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: + 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" - for _, runner in self._select_runners(recv_req.selector): - success, message = runner.weight_updater.update_weights_from_tensor( - named_tensors=named_tensors, - load_format=recv_req.load_format, + if base_tensors: + assert self._session.sync_base, ( + "base tensors arrived in a sync_base=False weight-update session" ) - if not success: - break + for _, runner in self._select_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._session = msgspec.structs.replace( + self._session, loaded_weights=True + ) if success: - self._session = msgspec.structs.replace( - self._session, loaded_weights=True - ) self.flush_cache_after_weight_update(recv_req) self.record_weight_version_after_update(recv_req.weight_version) else: @@ -295,9 +365,13 @@ def begin_weight_update(self, recv_req: BeginWeightUpdateReqInput): message="a weight-update session is already open; " "call end_weight_update() first", ) - for _, runner in self._select_runners(recv_req.selector): - runner.weight_updater.begin_weight_update() - self._session = _WeightUpdateSession(selector=recv_req.selector) + self._lora_stash = {} + if recv_req.sync_base: + for _, runner in self._select_runners(recv_req.selector): + runner.weight_updater.begin_weight_update() + self._session = _WeightUpdateSession( + selector=recv_req.selector, sync_base=recv_req.sync_base + ) torch.distributed.barrier(group=self.tp_cpu_group) return BeginWeightUpdateReqOutput(success=True, message="Success") @@ -308,12 +382,90 @@ def end_weight_update(self, recv_req: EndWeightUpdateReqInput): success=False, message="no weight-update session is open; call begin_weight_update() first", ) - run_post_load = not self._session.loaded_weights - for _, runner in self._select_runners(self._session.selector): - runner.weight_updater.end_weight_update(run_post_load=run_post_load) + if self._session.sync_base: + run_post_load = not self._session.loaded_weights + for _, runner in self._select_runners(self._session.selector): + runner.weight_updater.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) + pending_version = self._session.pending_version self._session = None + if success and not recv_req.abort: + self.record_weight_version_after_update(pending_version) torch.distributed.barrier(group=self.tp_cpu_group) - return EndWeightUpdateReqOutput(success=True, message="Success") + 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.model_runner.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 diff --git a/python/sglang/srt/managers/tokenizer_control_mixin.py b/python/sglang/srt/managers/tokenizer_control_mixin.py index f5051295bd33..d41724b451c8 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.disaggregation.utils import DisaggregationMode from sglang.srt.managers.communicator import FanOutCommunicator @@ -47,8 +48,6 @@ InitWeightsUpdateGroupReqOutput, ListExternalCorporaReqInput, ListExternalCorporaReqOutput, - LoadLoRAAdapterFromTensorsReqInput, - LoadLoRAAdapterFromTensorsReqOutput, LoadLoRAAdapterReqInput, LoadLoRAAdapterReqOutput, LoRAUpdateOutput, @@ -58,6 +57,8 @@ ProfileReq, ProfileReqOutput, ProfileReqType, + RegisterLoRAAdapterReqInput, + RegisterLoRAAdapterReqOutput, ReleaseMemoryOccupationReqInput, ReleaseMemoryOccupationReqOutput, RemoveExternalCorpusReqInput, @@ -465,8 +466,12 @@ async def destroy_weights_update_group( return FanOutCommunicator.merge_results(results) async def _weight_update_session_call( - self: TokenizerManager, communicator, obj + 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 @@ -474,27 +479,108 @@ async def _weight_update_session_call( if is_paused: results = await communicator(obj) if not is_paused: - async with self.model_update_lock.writer_lock: + 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]: - return await self._weight_update_session_call( - self.begin_weight_update_communicator, obj + 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]: - return await self._weight_update_session_call( - self.end_weight_update_communicator, obj + 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(): + self.pending_lora_unloads[ref.lora_name] = ref.lora_id + result = _merge_lora_update_results( + await self.update_lora_adapter_communicator( + UnloadLoRAAdapterReqInput( + lora_name=ref.lora_name, lora_id=ref.lora_id + ) + ) + ) + if result.success: + self.pending_lora_unloads.pop(ref.lora_name) + else: + 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, @@ -506,20 +592,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: @@ -566,6 +648,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) @@ -573,16 +657,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: @@ -701,35 +780,8 @@ async def load_lora_adapter( ): self.pending_lora_unloads[obj.lora_name] = new_adapter.lora_id - 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: @@ -738,11 +790,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: @@ -750,76 +850,84 @@ 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: if obj.lora_name in self.pending_lora_unloads: raise ValueError( f"LoRA adapter '{obj.lora_name}' has an incomplete unload. " "Retry the unload before loading it again." ) - - 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 rank_results = await self.update_lora_adapter_communicator(obj) result = _merge_lora_update_results(rank_results) + if result.success: - await self.lora_registry.register(new_adapter) - self.lora_ref_cache[obj.lora_name] = new_adapter - elif ( - obj.lora_name not in self.lora_registry.get_all_adapters() - and _lora_load_needs_cleanup(rank_results, obj.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() + elif not reused and _lora_load_needs_cleanup( + rank_results, obj.lora_name ): self.pending_lora_unloads[obj.lora_name] = new_adapter.lora_id - - 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] - return result except ValueError as e: - return LoadLoRAAdapterFromTensorsReqOutput( + return RegisterLoRAAdapterReqOutput( success=False, error_message=str(e), ) @@ -848,9 +956,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)) @@ -1047,8 +1167,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 eb74ca022316..918dd731c22f 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 @@ -724,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 @@ -745,6 +762,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 @@ -1675,10 +1695,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: @@ -1720,6 +1768,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: @@ -1787,6 +1842,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, @@ -1816,9 +1890,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"], @@ -2111,20 +2186,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: @@ -2628,8 +2722,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) @@ -3427,6 +3520,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) @@ -3548,18 +3645,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, @@ -3645,6 +3764,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 6adfbe043792..f38f6e6b57d6 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -28,8 +28,8 @@ GetWeightsByNameReqInput, InitWeightsSendGroupForRemoteInstanceReqInput, InitWeightsUpdateGroupReqInput, - LoadLoRAAdapterFromTensorsReqInput, LoadLoRAAdapterReqInput, + RegisterLoRAAdapterReqInput, SendWeightsToRemoteInstanceReqInput, UnloadLoRAAdapterReqInput, ) @@ -65,7 +65,6 @@ 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 @@ -212,58 +211,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.model_runner.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.model_runner.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( diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 2f4b097d8943..cf4a62cc3169 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1389,12 +1389,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.""" 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/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..18bd94a7ca50 --- /dev/null +++ b/test/registered/unit/lora/test_lora_staged_publication.py @@ -0,0 +1,314 @@ +"""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.pending_lora_unloads = {} + 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..bf026ceeaab2 --- /dev/null +++ b/test/registered/unit/lora/test_lora_stream_update.py @@ -0,0 +1,301 @@ +"""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, + _WeightUpdateSession, +) +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.model_runner.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._session = _WeightUpdateSession(selector="all") + 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._session.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") + pending_version = mgr._session.pending_version + mgr._session = None + mgr.record_weight_version_after_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..f9e7b805c971 --- /dev/null +++ b/test/registered/unit/lora/test_lora_upsert.py @@ -0,0 +1,352 @@ +"""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() -> 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.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 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/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_lora_update_result_merge.py b/test/registered/unit/managers/test_lora_update_result_merge.py index 22fe49b50e45..5165af0d6962 100644 --- a/test/registered/unit/managers/test_lora_update_result_merge.py +++ b/test/registered/unit/managers/test_lora_update_result_merge.py @@ -20,9 +20,9 @@ from sglang.srt.lora.lora_registry import LoRARef, LoRARegistry from sglang.srt.managers.io_struct import ( - LoadLoRAAdapterFromTensorsReqInput, LoadLoRAAdapterReqInput, LoRAUpdateOutput, + RegisterLoRAAdapterReqInput, UnloadLoRAAdapterReqInput, ) from sglang.srt.managers.tokenizer_control_mixin import ( @@ -89,13 +89,14 @@ def test_failed_load_does_not_redirect_existing_adapter_unload(self): """Rejected duplicates must keep the original adapter's unload ID; fresh partial loads must remain cleanable via their new ID.""" - async def scenario(from_tensors, duplicate): + async def scenario(from_stream, duplicate): original = LoRARef(lora_name="a", lora_path="path", pinned=False) manager = TokenizerControlMixin() manager.auto_create_handle_loop = Mock() manager.lora_update_lock = asyncio.Lock() manager.lora_registry = LoRARegistry([original] if duplicate else []) manager.pending_lora_unloads = {} + manager._pending_lora_publications = {} manager.lora_ref_cache = {} replies = ( [_err("already loaded", {"a": "path"})] @@ -105,11 +106,9 @@ async def scenario(from_tensors, duplicate): manager.update_lora_adapter_communicator = AsyncMock( side_effect=[replies, [_ok()]] ) - if from_tensors: - request = LoadLoRAAdapterFromTensorsReqInput( - lora_name="a", config_dict={}, serialized_named_tensors=[b""] - ) - result = await manager.load_lora_adapter_from_tensors(request) + if from_stream: + request = RegisterLoRAAdapterReqInput(lora_name="a", config_dict={}) + result = await manager.load_lora_adapter_from_stream(request) else: request = LoadLoRAAdapterReqInput(lora_name="a", lora_path="path") result = await manager.load_lora_adapter(request) @@ -136,10 +135,10 @@ async def scenario(from_tensors, duplicate): self.assertEqual(manager.pending_lora_unloads, {}) with get_context().override_server_args(enable_lora=True, dp_size=1): - for from_tensors in (False, True): + for from_stream in (False, True): for duplicate in (True, False): - with self.subTest(from_tensors=from_tensors, duplicate=duplicate): - asyncio.run(scenario(from_tensors, duplicate)) + with self.subTest(from_stream=from_stream, duplicate=duplicate): + asyncio.run(scenario(from_stream, duplicate)) if __name__ == "__main__": diff --git a/test/registered/unit/managers/test_msgpack_ipc_roundtrip.py b/test/registered/unit/managers/test_msgpack_ipc_roundtrip.py index f718711a4be6..65a28ebb9660 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, @@ -136,12 +135,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}] From 2b83564e74e8f77a88241b8f6de944a6cd53b6aa Mon Sep 17 00:00:00 2001 From: JD-ETH Date: Tue, 24 Mar 2026 05:05:43 +0000 Subject: [PATCH 15/56] [13/32] [sglang-miles] P2P weight mapping and scoped replica parallelism (#21278, #22663, #33478) Expose rank topology for P2P weight mapping and create replicas using scoped runtime-context overrides. Restore the previous configuration on exit, including failed context entry. Co-authored-by: Claude Opus 4.6 Co-authored-by: yueming-yuan --- .../sglang/srt/distributed/parallel_state.py | 164 +++++++++- .../engine_info_bootstrap_server.py | 41 ++- python/sglang/srt/entrypoints/http_server.py | 25 ++ .../remote_instance_weight_transporter.py | 50 +++ .../srt/model_loader/parameter_mapper.py | 261 ++++++++++++++++ .../deepseek_common/deepseek_weight_loader.py | 66 ++-- python/sglang/srt/models/deepseek_nextn.py | 1 + python/sglang/srt/models/deepseek_v2.py | 57 ++++ python/sglang/srt/models/glm4.py | 20 +- python/sglang/srt/models/glm4_moe.py | 66 ++-- python/sglang/srt/models/glm4_moe_lite.py | 43 +++ python/sglang/srt/models/glm4_moe_nextn.py | 23 +- python/sglang/srt/models/llama.py | 28 +- python/sglang/srt/models/qwen2.py | 20 +- python/sglang/srt/models/qwen3.py | 32 +- python/sglang/srt/models/qwen3_moe.py | 34 +- .../test_parallelism_context_integration.py | 276 +++++++++++++++++ test/srt/models/test_params_mapping.py | 292 ++++++++++++++++++ 18 files changed, 1350 insertions(+), 149 deletions(-) create mode 100644 python/sglang/srt/model_loader/parameter_mapper.py create mode 100644 test/registered/e2e/disaggregation/test_parallelism_context_integration.py create mode 100644 test/srt/models/test_params_mapping.py diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index a014e28a14e4..683d837d25f2 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -34,11 +34,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 @@ -3467,6 +3467,166 @@ def monkey_patch_vllm_parallel_state(reverse: bool = False): 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.""" + from sglang.srt.runtime_context import get_parallel + + ps = get_parallel() + + return cls( + tp_size=ps.tp_size, + tp_rank=ps.tp_rank, + pp_size=ps.pp_size, + pp_rank=ps.pp_rank, + ep_size=ps.moe_ep_size, + ep_rank=ps.moe_ep_rank, + moe_tp_size=ps.moe_tp_size, + moe_tp_rank=ps.moe_tp_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=ps.moe_dp_size, + moe_dp_rank=ps.moe_dp_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, + ) + + +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._scope = None + + 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): + from sglang.srt.distributed import parallel_state + from sglang.srt.runtime_context import get_flags, get_parallel + + conf = self.config + groups = { + "tp": ("_TP", conf.tp_size, conf.tp_rank), + "pp": ("_PP", conf.pp_size, conf.pp_rank), + "moe_ep": ("_MOE_EP", conf.ep_size, conf.ep_rank), + "moe_tp": ("_MOE_TP", conf.moe_tp_size, conf.moe_tp_rank), + "attn_tp": ("_ATTN_TP", conf.attn_tp_size, conf.attn_tp_rank), + "attn_cp": ("_ATTN_CP", conf.attn_cp_size, conf.attn_cp_rank), + "moe_dp": ("_MOE_DP", conf.moe_dp_size, conf.moe_dp_rank), + } + values = { + "attn_dp_size": conf.attn_dp_size, + "attn_dp_rank": conf.attn_dp_rank, + "enable_dp_attention": conf.attn_dp_size > 1, + "launch_world_size": conf.world_size, + "launch_world_rank": conf.global_rank, + } + with contextlib.ExitStack() as scope: + for prefix, (name, size, rank) in groups.items(): + group = self._create_mock_group(size, rank) + scope.enter_context(patch.object(parallel_state, name, group)) + values.update( + { + f"{prefix}_size": size, + f"{prefix}_rank": rank, + f"{prefix}_group": group, + } + ) + scope.enter_context(get_flags().dp.override(enabled=conf.attn_dp_size > 1)) + # Model layers read the runtime context; changing legacy groups alone + # would leave them using the previously published topology. + scope.enter_context(get_parallel().override(**values)) + self._scope = scope.pop_all() + logger.info(f"[ParallelismContext] Activated: {conf}") + return self + + def __exit__(self, *args): + self._scope.__exit__(*args) + logger.info("[ParallelismContext] Deactivated") + return False + + # Use `get_parallel()` outside this package. Warn once per deprecated getter. _EXEMPT_CALLERS = ("sglang.srt.distributed.",) 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/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 72651ba43b51..e050865472f4 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -1343,6 +1343,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( 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 8d10417696ba..98392e3a266a 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, @@ -30,6 +31,7 @@ class RemoteInstanceWeightTransporter: 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 +57,9 @@ def init_engine(self): self.session_id = NetworkAddress( local_ip, self.engine.get_rpc_port() ).to_host_port_str() + self.parallelism_config = RankParallelismConfig.from_parallel_state( + self.tp_rank + ) def maybe_register_and_publish_weight_info(self) -> None: if ( @@ -71,6 +76,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.remote_instance_weight_loader_use_transfer_engine() + 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_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 774f325b7c75..5d0003272306 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 7bc345ac2ca1..d5f89429b3a0 100644 --- a/python/sglang/srt/models/deepseek_nextn.py +++ b/python/sglang/srt/models/deepseek_nextn.py @@ -282,6 +282,7 @@ def __init__( # if not set, model load will be broken in DeepseekV3ForCausalLM load_weights() self.pp_group = get_parallel().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 0181eb7b91d2..440601eebb66 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -3220,6 +3220,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() @@ -3232,6 +3234,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 @@ -3241,6 +3259,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/glm4.py b/python/sglang/srt/models/glm4.py index f1feda9e9ec3..2e654b37343e 100644 --- a/python/sglang/srt/models/glm4.py +++ b/python/sglang/srt/models/glm4.py @@ -455,6 +455,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 @@ -549,14 +560,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 e2a2875f80a1..14f74e6b844b 100644 --- a/python/sglang/srt/models/glm4_moe.py +++ b/python/sglang/srt/models/glm4_moe.py @@ -1157,9 +1157,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 @@ -1240,43 +1265,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}" @@ -1297,6 +1287,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 853115b7c391..78a683c6afc8 100644 --- a/python/sglang/srt/models/glm4_moe_lite.py +++ b/python/sglang/srt/models/glm4_moe_lite.py @@ -877,6 +877,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 @@ -989,6 +1016,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 83750d24328d..b8ceb8c7e786 100644 --- a/python/sglang/srt/models/glm4_moe_nextn.py +++ b/python/sglang/srt/models/glm4_moe_nextn.py @@ -26,6 +26,7 @@ from sglang.srt.layers.layer_boundary.residual import batch as residual_batch 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, @@ -33,7 +34,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__) @@ -147,6 +148,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/llama.py b/python/sglang/srt/models/llama.py index 0da5cb20a5fe..0f2b2e8026db 100644 --- a/python/sglang/srt/models/llama.py +++ b/python/sglang/srt/models/llama.py @@ -546,9 +546,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, @@ -666,23 +678,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 @@ -709,7 +709,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/qwen2.py b/python/sglang/srt/models/qwen2.py index 5a37f29672d8..7eb7f4d00e0c 100644 --- a/python/sglang/srt/models/qwen2.py +++ b/python/sglang/srt/models/qwen2.py @@ -518,6 +518,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 @@ -619,14 +630,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/qwen3.py b/python/sglang/srt/models/qwen3.py index cf2bc13901b6..ee00b8d9d7d5 100644 --- a/python/sglang/srt/models/qwen3.py +++ b/python/sglang/srt/models/qwen3.py @@ -43,6 +43,7 @@ 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 @@ -333,16 +334,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( @@ -527,6 +519,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: @@ -636,15 +638,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_moe.py b/python/sglang/srt/models/qwen3_moe.py index e00e9363e3d9..eedee0906cd1 100644 --- a/python/sglang/srt/models/qwen3_moe.py +++ b/python/sglang/srt/models/qwen3_moe.py @@ -1045,6 +1045,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( @@ -1183,21 +1200,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/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..a91173c231f1 --- /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 gc +import sys +from typing import Dict, List, Tuple + +import msgspec +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 msgspec.structs.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.distributed.parallel_state import ParallelismContext + from sglang.srt.runtime_context import publish, restore_context, snapshot_context + + original_context = snapshot_context() + + 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. + publish(server_args, role="model_runner") + 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: + restore_context(original_context) + + +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/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) From 4a7bc4a36b219af6ac78e5ea18b070a765fe1cbf Mon Sep 17 00:00:00 2001 From: yueming-yuan Date: Sat, 22 Aug 2026 12:43:45 -0700 Subject: [PATCH 16/56] [14/32] [sglang-miles] rollout indexer replay: raw seq-relative topk capture + meta_info threading --- .../srt/layers/attention/dsa/dsa_indexer.py | 103 +++++++++++++----- .../attention/dsa/dsa_indexer_metadata.py | 2 + .../layers/attention/dsa/dsa_topk_backend.py | 48 +++++--- .../srt/managers/detokenizer_manager.py | 1 + python/sglang/srt/managers/io_struct.py | 7 ++ .../srt/managers/multi_tokenizer_mixin.py | 2 + .../scheduler_components/output_streamer.py | 10 ++ .../sglang/srt/managers/tokenizer_manager.py | 3 + 8 files changed, 134 insertions(+), 42 deletions(-) diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py index f4e23f74b6a7..fd7ffcf4bc27 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py @@ -61,6 +61,7 @@ get_parallel, ) from sglang.srt.state_capturer.indexer_topk import ( + get_global_indexer_capturer, maybe_capture_indexer_topk, ) from sglang.srt.utils import ( @@ -1139,7 +1140,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 @@ -1150,7 +1158,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: cached_budget = self._mqa_logits_budget_bytes.get(device_index) @@ -1229,8 +1239,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() @@ -1309,9 +1321,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))) @@ -1386,19 +1406,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, @@ -1476,7 +1509,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,...] @@ -1487,13 +1520,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, @@ -1595,6 +1634,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, @@ -1659,7 +1707,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, @@ -1669,7 +1717,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. @@ -1861,6 +1909,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 @@ -1879,14 +1928,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: @@ -1897,7 +1946,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, @@ -1908,4 +1957,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 7802f95904a8..ec1586b5419b 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py @@ -105,9 +105,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 @@ -130,7 +142,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 @@ -147,9 +160,10 @@ def topk_transform( and 0 < topk <= 2048 and lengths.shape[0] == logits.shape[0] == topk_indices_offset.shape[0] ): - return _topk_transform_v2_ragged( + result = _topk_transform_v2_ragged( logits, lengths, topk, topk_indices_offset, row_starts ) + return (result, raw_indices) if return_raw_indices else result # Packed PAGED extend (GLM DSA prefill), ROCm-only: CUDA gets the same # fusion from RAGGED above. Unsupported shapes fall back, not raise. @@ -175,7 +189,7 @@ def topk_transform( and row_to_batch is not None and row_to_batch.shape[0] == logits.shape[0] ): - return _topk_transform_v2_packed( + result = _topk_transform_v2_packed( logits, lengths, topk, @@ -183,6 +197,7 @@ def topk_transform( row_starts=row_starts, row_to_batch=row_to_batch, ) + return (result, raw_indices) if return_raw_indices else result # The legacy transforms below read attn_metadata.page_table_1 (page_size=1), # which is always present here: the fold only drops it for the decode case @@ -201,7 +216,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, @@ -209,22 +224,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: @@ -236,7 +251,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(), @@ -248,13 +263,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(), @@ -264,9 +279,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/managers/detokenizer_manager.py b/python/sglang/srt/managers/detokenizer_manager.py index 00236933cc37..ef253a939838 100644 --- a/python/sglang/srt/managers/detokenizer_manager.py +++ b/python/sglang/srt/managers/detokenizer_manager.py @@ -498,6 +498,7 @@ def handle_batch_token_id_out(self, recv_obj: BatchTokenIDOutput): 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 6be95c8bb1a6..8e875932c5e9 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -1603,6 +1603,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 @@ -1697,6 +1701,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 diff --git a/python/sglang/srt/managers/multi_tokenizer_mixin.py b/python/sglang/srt/managers/multi_tokenizer_mixin.py index 48cb012c13ec..150ee49055b5 100644 --- a/python/sglang/srt/managers/multi_tokenizer_mixin.py +++ b/python/sglang/srt/managers/multi_tokenizer_mixin.py @@ -256,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, @@ -374,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/scheduler_components/output_streamer.py b/python/sglang/srt/managers/scheduler_components/output_streamer.py index 752d456b9b98..0bb77a8c6885 100644 --- a/python/sglang/srt/managers/scheduler_components/output_streamer.py +++ b/python/sglang/srt/managers/scheduler_components/output_streamer.py @@ -679,6 +679,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, @@ -739,6 +748,7 @@ def to_payload( 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/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 918dd731c22f..110c1d9fd573 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -2571,6 +2571,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] From 5e5043295ee5778d0a42946514a073ccf47f6f38 Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Thu, 25 Jun 2026 22:54:58 +0000 Subject: [PATCH 17/56] [15/32] [sglang-miles] Exclude shared skip-topk indexer weights from RL weight checks (#29339) Co-authored-by: yueming-yuan --- python/sglang/srt/models/deepseek_v2.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 440601eebb66..d1d50ed9b3b7 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -2092,6 +2092,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), From 47fbef9d2b90d4f10c339db116945ffcda8a1861 Mon Sep 17 00:00:00 2001 From: yueming-yuan Date: Thu, 2 Jul 2026 08:26:50 +0000 Subject: [PATCH 18/56] [16/32] [sglang-miles] Fix stale _attn_sink_local cache after RL weight updates (#30421) --- python/sglang/srt/models/deepseek_v4.py | 13 +++++++++++++ python/sglang/srt/models/deepseek_v4_dspark.py | 13 +++++++++++++ python/sglang/srt/models/deepseek_v4_nextn.py | 1 + 3 files changed, 27 insertions(+) diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index e4d06703faf3..f9f75894a83b 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -836,6 +836,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, @@ -1005,6 +1006,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() @@ -5218,6 +5230,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( diff --git a/python/sglang/srt/models/deepseek_v4_dspark.py b/python/sglang/srt/models/deepseek_v4_dspark.py index baebc2de6fcc..5fe64615c858 100644 --- a/python/sglang/srt/models/deepseek_v4_dspark.py +++ b/python/sglang/srt/models/deepseek_v4_dspark.py @@ -149,6 +149,10 @@ def kv_proj_only(self, x: torch.Tensor) -> torch.Tensor: kv, _ = self.wkv(x) return kv + def _local_attn_sink(self) -> torch.Tensor: + # DSpark pads Q to this width; keep the base helper's stable allocation. + return super()._local_attn_sink(max(self.n_local_heads, _PAD_NUM_HEADS)) + def _store_block_kv( self, *, @@ -1146,6 +1150,15 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> None: self._assert_confidence_head_loaded( params_dict=params_dict, loaded_params=loaded_params ) + self.post_load_weights() + + 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 diff --git a/python/sglang/srt/models/deepseek_v4_nextn.py b/python/sglang/srt/models/deepseek_v4_nextn.py index 6694abb603a8..0948c5fe84e2 100644 --- a/python/sglang/srt/models/deepseek_v4_nextn.py +++ b/python/sglang/srt/models/deepseek_v4_nextn.py @@ -256,6 +256,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] From ebee922531e0708d189ee3cf2592713190e47323 Mon Sep 17 00:00:00 2001 From: Yusheng Su Date: Wed, 8 Jul 2026 08:07:23 +0000 Subject: [PATCH 19/56] [17/32] [sglang-miles] Support Qwen3.5 GDN LoRA targets (#30499, #37466, #37887) Co-authored-by: fzyzcjy Co-authored-by: yueming-yuan --- python/sglang/srt/lora/lora.py | 60 +++++++++++++++++++++++++++++ python/sglang/srt/lora/utils.py | 4 ++ python/sglang/srt/models/qwen3_5.py | 4 ++ python/sglang/srt/utils/common.py | 4 ++ 4 files changed, 72 insertions(+) diff --git a/python/sglang/srt/lora/lora.py b/python/sglang/srt/lora/lora.py index c154e1c5581b..99848d822e5f 100644 --- a/python/sglang/srt/lora/lora.py +++ b/python/sglang/srt/lora/lora.py @@ -210,6 +210,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 +452,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 +480,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/utils.py b/python/sglang/srt/lora/utils.py index f4310e1f0020..5c1b27305cc6 100644 --- a/python/sglang/srt/lora/utils.py +++ b/python/sglang/srt/lora/utils.py @@ -285,6 +285,8 @@ def get_normalized_target_modules( "v_proj": "qkv_proj", "gate_proj": "gate_up_proj", "up_proj": "gate_up_proj", + "in_proj_b": "in_proj_ba", + "in_proj_a": "in_proj_ba", "out_proj": "out_proj", "embed_tokens": "embed_tokens", "vocab_emb": "embed_tokens", @@ -323,6 +325,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 +405,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", diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index fac16f30a38a..1da735faef4a 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -1600,6 +1600,7 @@ class Qwen3_5ForCausalLM(nn.Module): "o_proj", "out_proj", "in_proj_qkvz", + "in_proj_ba", "gate_up_proj", "down_proj", "lm_head", @@ -1625,6 +1626,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/utils/common.py b/python/sglang/srt/utils/common.py index 59572210b982..4bc416920edc 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -4600,6 +4600,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" From 0e6fcdc0ae71fc8aba81f07b45c1da37f5a565b6 Mon Sep 17 00:00:00 2001 From: Nan Date: Tue, 7 Jul 2026 08:13:19 +0000 Subject: [PATCH 20/56] [18/32] [sglang-miles] Add /pull_weights: engine-side pull of published weights into a host-local checkpoint (#30366, #28524) --- python/sglang/srt/arg_groups/fields/model.py | 4 + python/sglang/srt/entrypoints/http_server.py | 14 + python/sglang/srt/managers/io_struct.py | 16 + python/sglang/srt/managers/scheduler.py | 5 + .../scheduler_components/weight_updater.py | 40 ++ .../srt/managers/tokenizer_control_mixin.py | 12 + .../srt/weight_sync/local_checkpoint.py | 350 ++++++++++++++++++ 7 files changed, 441 insertions(+) create mode 100644 python/sglang/srt/weight_sync/local_checkpoint.py diff --git a/python/sglang/srt/arg_groups/fields/model.py b/python/sglang/srt/arg_groups/fields/model.py index b6ca07469190..acc91da017d1 100644 --- a/python/sglang/srt/arg_groups/fields/model.py +++ b/python/sglang/srt/arg_groups/fields/model.py @@ -277,6 +277,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.", diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index e050865472f4..bb76f303d128 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -141,6 +141,7 @@ PauseGenerationReqInput, PdRoleSwitchReqInput, ProfileReq, + PullWeightsReqInput, RegisterLoRAAdapterReqInput, ReleaseMemoryOccupationReqInput, ResumeMemoryOccupationReqInput, @@ -1268,6 +1269,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( diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 8e875932c5e9..4ebce9cb53d7 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -1909,6 +1909,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] diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index a0b4e0437c3e..f24f4984c47b 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -161,6 +161,7 @@ PauseGenerationReqInput, PdRoleSwitchReqInput, ProfileReq, + PullWeightsReqInput, RegisterLoRAAdapterReqInput, RegisterLoRAAdapterReqOutput, ReleaseMemoryOccupationReqInput, @@ -1789,6 +1790,10 @@ def init_request_dispatcher(self): CheckWeightsReqInput, self.weight_updater.check_weights, ), + ( + PullWeightsReqInput, + self.weight_updater.pull_weights, + ), (SlowDownReqInput, self.slow_down), (PdRoleSwitchReqInput, self.handle_pd_role_switch), ( diff --git a/python/sglang/srt/managers/scheduler_components/weight_updater.py b/python/sglang/srt/managers/scheduler_components/weight_updater.py index 9a4237c908bb..f4f68eff5a59 100644 --- a/python/sglang/srt/managers/scheduler_components/weight_updater.py +++ b/python/sglang/srt/managers/scheduler_components/weight_updater.py @@ -32,6 +32,8 @@ GetWeightsByNameReqOutput, InitWeightsUpdateGroupReqInput, InitWeightsUpdateGroupReqOutput, + PullWeightsReqInput, + PullWeightsReqOutput, ReleaseMemoryOccupationReqInput, ReleaseMemoryOccupationReqOutput, ResumeMemoryOccupationReqInput, @@ -188,6 +190,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) diff --git a/python/sglang/srt/managers/tokenizer_control_mixin.py b/python/sglang/srt/managers/tokenizer_control_mixin.py index d41724b451c8..4abd64d477d9 100644 --- a/python/sglang/srt/managers/tokenizer_control_mixin.py +++ b/python/sglang/srt/managers/tokenizer_control_mixin.py @@ -57,6 +57,8 @@ ProfileReq, ProfileReqOutput, ProfileReqType, + PullWeightsReqInput, + PullWeightsReqOutput, RegisterLoRAAdapterReqInput, RegisterLoRAAdapterReqOutput, ReleaseMemoryOccupationReqInput, @@ -123,6 +125,7 @@ ("release_memory_occupation", ReleaseMemoryOccupationReqOutput), ("resume_memory_occupation", ResumeMemoryOccupationReqOutput), ("check_weights", CheckWeightsReqOutput), + ("pull_weights", PullWeightsReqOutput), ("slow_down", SlowDownReqOutput), ("pd_role_switch", PdRoleSwitchReqOutput), ("flush_cache", FlushCacheReqOutput), @@ -1003,6 +1006,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, 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=" Date: Tue, 4 Aug 2026 05:05:18 +0000 Subject: [PATCH 21/56] [19/32] [sglang-miles] Nemotron hybrid metadata and routing replay (#27110, #38718) Co-authored-by: Shi Dong --- .../srt/layers/attention/mamba/mamba2_metadata.py | 11 ++++++++++- python/sglang/srt/models/nemotron_h.py | 1 + 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py b/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py index b2843a536e5b..d4ac22143177 100644 --- a/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py +++ b/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py @@ -250,7 +250,16 @@ def prepare_mixed( num_prefill_tokens = int(sum(extend_seq_lens_cpu)) else: num_prefill_tokens = int(forward_batch.extend_num_tokens) - num_decodes = len(forward_batch.seq_lens) - 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/models/nemotron_h.py b/python/sglang/srt/models/nemotron_h.py index 4a0cc531c65c..f361910552ee 100644 --- a/python/sglang/srt/models/nemotron_h.py +++ b/python/sglang/srt/models/nemotron_h.py @@ -216,6 +216,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, From d2fc909aae8d9b0f8bc866e89f588cdf3cbdabb5 Mon Sep 17 00:00:00 2001 From: xyuzh Date: Tue, 9 Jun 2026 23:16:35 +0000 Subject: [PATCH 22/56] [20/32] [sglang-miles] RDT/NIXL weight sync for Ray scheduler actors (#27723) Co-authored-by: Claude Opus 4.8 (1M context) Co-authored-by: yueming-yuan --- python/pyproject.toml | 2 +- python/sglang/srt/arg_groups/fields/model.py | 8 +++ python/sglang/srt/arg_groups/pipeline.py | 2 + python/sglang/srt/arg_groups/serving_hook.py | 14 +++++ python/sglang/srt/entrypoints/engine.py | 2 +- .../sglang/srt/model_executor/model_runner.py | 2 + .../remote_instance_weight_transporter.py | 19 +++--- python/sglang/srt/ray/__init__.py | 3 +- python/sglang/srt/ray/engine.py | 35 +++++++++-- python/sglang/srt/ray/http_server.py | 60 +++++++++++++++---- python/sglang/srt/ray/scheduler_actor.py | 43 ++++++++++++- python/sglang/srt/server_args.py | 16 +++++ 12 files changed, 178 insertions(+), 28 deletions(-) diff --git a/python/pyproject.toml b/python/pyproject.toml index 50394749d90b..c9a945554572 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -142,7 +142,7 @@ diffusion-qvg = [ ] ray = [ - "ray[default]>=2.55.1", + "ray[default]>=2.56.0", ] tracing = [ diff --git a/python/sglang/srt/arg_groups/fields/model.py b/python/sglang/srt/arg_groups/fields/model.py index acc91da017d1..4ee0d4bde0d6 100644 --- a/python/sglang/srt/arg_groups/fields/model.py +++ b/python/sglang/srt/arg_groups/fields/model.py @@ -318,6 +318,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/pipeline.py b/python/sglang/srt/arg_groups/pipeline.py index 3963f548c128..616f755187b4 100644 --- a/python/sglang/srt/arg_groups/pipeline.py +++ b/python/sglang/srt/arg_groups/pipeline.py @@ -54,6 +54,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, @@ -70,6 +71,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 e18fc6648aab..bd4c52a20823 100644 --- a/python/sglang/srt/arg_groups/serving_hook.py +++ b/python/sglang/srt/arg_groups/serving_hook.py @@ -391,6 +391,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/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index 231c480b1794..851439bf172c 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -1125,7 +1125,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 diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index cf4a62cc3169..939314f13469 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -585,6 +585,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, gpu_id=self.gpu_id, ) @@ -685,6 +686,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: 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 98392e3a266a..f41c26622310 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 @@ -17,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__) @@ -24,6 +25,7 @@ @dataclass(slots=True, kw_only=True) class RemoteInstanceWeightTransporter: + server_args: ServerArgs get_model: Callable[[], torch.nn.Module] # Registration may run after the runner's construction scope has exited. tp_rank: int = field(init=False, default_factory=lambda: get_parallel().tp_rank) @@ -57,9 +59,12 @@ def init_engine(self): self.session_id = NetworkAddress( local_ip, self.engine.get_rpc_port() ).to_host_port_str() - self.parallelism_config = RankParallelismConfig.from_parallel_state( - self.tp_rank - ) + + 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 ( @@ -79,19 +84,19 @@ def maybe_register_and_publish_weight_info(self) -> None: # 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.remote_instance_weight_loader_use_transfer_engine() + 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: + if get_parallel().dist_init_addr: bootstrap_host = ( - NetworkAddress.parse(self.server_args.dist_init_addr).resolved().host + NetworkAddress.parse(get_parallel().dist_init_addr).resolved().host ) else: bootstrap_host = "127.0.0.1" - bootstrap_port = self.server_args.engine_info_bootstrap_port + bootstrap_port = get_model().engine_info_bootstrap_port return NetworkAddress(bootstrap_host, bootstrap_port).to_url() def _register_parallelism_config_to_bootstrap(self) -> None: 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 0d6b8e3d53a2..07b540d84623 100644 --- a/python/sglang/srt/ray/engine.py +++ b/python/sglang/srt/ray/engine.py @@ -182,6 +182,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, @@ -203,14 +219,25 @@ def _create_scheduler_actor( rank0_node_ip: IP of rank-0's node, used for NCCL rendezvous. dist_init_addr: Distributed init address (tcp://rank0_node_ip:nccl_port). """ + 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, diff --git a/python/sglang/srt/ray/http_server.py b/python/sglang/srt/ray/http_server.py index 963f05d13524..8b9389641b5b 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 + + 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, + placement_group=placement_group, + ) + + +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 8ec28b8d1db4..37f604397a9b 100644 --- a/python/sglang/srt/ray/scheduler_actor.py +++ b/python/sglang/srt/ray/scheduler_actor.py @@ -16,12 +16,19 @@ 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 SpawnRanks, get_device, publish, spawn_world_rank +from sglang.srt.runtime_context import ( + SpawnRanks, + get_device, + get_exec, + publish, + spawn_world_rank, +) from sglang.srt.server_args import PortArgs, ServerArgs logger = logging.getLogger(__name__) @@ -122,6 +129,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 get_exec().features.enable_memory_saver: + return + + import torch + from ray.experimental import register_nixl_memory + + torch.cuda.set_device(get_device().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(get_device().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/server_args.py b/python/sglang/srt/server_args.py index a49b2300990c..bb551ac50c75 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -425,6 +425,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 + ) + # Collect input fields only; field_order.py preserves positional argument order. From d1b0ae27571d67476f06cf4ffba16ea319fb4158 Mon Sep 17 00:00:00 2001 From: yueming-yuan Date: Wed, 26 Aug 2026 18:08:33 -0700 Subject: [PATCH 23/56] [21/32] [sglang-miles] args: drop DeepseekV4 from the FlashInfer allreduce-fusion auto-enable --- python/sglang/srt/arg_groups/overrides.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index dd41e5863563..298a232e628a 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -887,7 +887,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", From 0946f89221260cf91cd75967a465545b3d7eb781 Mon Sep 17 00:00:00 2001 From: yueming-yuan Date: Sun, 23 Aug 2026 21:25:49 +0000 Subject: [PATCH 24/56] [22/32] [sglang-miles] Add SGLANG_DISABLE_MULTIMEM_AG to force the NCCL all-gather path (#36110) --- .../srt/distributed/device_communicators/triton_symm_mem_ag.py | 3 +++ python/sglang/srt/environ.py | 1 + 2 files changed, 4 insertions(+) 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 f60b65ed7e7d..23278d9f1934 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/environ.py b/python/sglang/srt/environ.py index 787234e08ec4..17e77045e89e 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1795,6 +1795,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) From 4312f178c1ca6e6034397de993d11d7f537fab58 Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Tue, 11 Aug 2026 19:49:39 +0000 Subject: [PATCH 25/56] [23/32] [sglang-miles] Extract Anthropic conversion into standalone utils (#35127) Co-authored-by: Claude Fable 5 Co-authored-by: Shi Dong --- .../srt/entrypoints/anthropic/serving.py | 1181 +++++++++-------- .../sglang/srt/entrypoints/anthropic/utils.py | 219 +++ .../unit/entrypoints/anthropic/test_utils.py | 793 +++++++++++ 3 files changed, 1616 insertions(+), 577 deletions(-) create mode 100644 python/sglang/srt/entrypoints/anthropic/utils.py create mode 100644 test/registered/unit/entrypoints/anthropic/test_utils.py 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/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) From 9e1f55033abd381aa7e6abfe9be3a68cc1eabef7 Mon Sep 17 00:00:00 2001 From: yueming-yuan Date: Wed, 2 Sep 2026 20:04:49 +0000 Subject: [PATCH 26/56] [24/32] [sglang-miles] Exclude derived sincos and MRoPE buffers from weight checks (#37649) Co-authored-by: Zhichenzzz --- python/sglang/srt/layers/rotary_embedding/mrope.py | 4 ++++ python/sglang/srt/models/kimi_k25.py | 1 + python/sglang/srt/models/kimi_k3_vl.py | 1 + python/sglang/srt/models/minicpmv.py | 3 +++ 4 files changed, 9 insertions(+) diff --git a/python/sglang/srt/layers/rotary_embedding/mrope.py b/python/sglang/srt/layers/rotary_embedding/mrope.py index 5f42f89fceef..15386e359832 100644 --- a/python/sglang/srt/layers/rotary_embedding/mrope.py +++ b/python/sglang/srt/layers/rotary_embedding/mrope.py @@ -105,6 +105,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/models/kimi_k25.py b/python/sglang/srt/models/kimi_k25.py index 09cc7952c47f..57c071099709 100644 --- a/python/sglang/srt/models/kimi_k25.py +++ b/python/sglang/srt/models/kimi_k25.py @@ -248,6 +248,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_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/minicpmv.py b/python/sglang/srt/models/minicpmv.py index 4fa7b66a6e43..a88e9912c215 100644 --- a/python/sglang/srt/models/minicpmv.py +++ b/python/sglang/srt/models/minicpmv.py @@ -304,6 +304,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 @@ -426,6 +427,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 @@ -455,6 +457,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" From 4f88adb3578d9dc1d9115d7d043c7102c59b79fd Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Fri, 4 Sep 2026 07:27:39 +0000 Subject: [PATCH 27/56] [25/32] [sglang-miles] fix(qwen3_5): let the GDN gated norm follow the model's device (#37954) --- python/sglang/srt/models/qwen3_5.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 1da735faef4a..982afe72cbf4 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -440,7 +440,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} From 2ac868be27b133e47f1c6a2a661167d1b496c11a Mon Sep 17 00:00:00 2001 From: yueming-yuan Date: Thu, 17 Sep 2026 21:19:48 +0000 Subject: [PATCH 28/56] [26/32] [sglang-miles] Compare MXFP4 Marlin experts in dequantized space in the weight checker (#40028) --- .../srt/layers/quantization/marlin_utils.py | 45 ++++ python/sglang/srt/utils/weight_checker.py | 2 +- .../srt/utils/weight_checker_comparator.py | 145 ++++++++++++ python/sglang/test/test_marlin_utils.py | 48 +--- .../rl/test_weight_checker_mxfp4_marlin.py | 222 ++++++++++++++++++ 5 files changed, 415 insertions(+), 47 deletions(-) create mode 100644 test/registered/e2e/rl/test_weight_checker_mxfp4_marlin.py diff --git a/python/sglang/srt/layers/quantization/marlin_utils.py b/python/sglang/srt/layers/quantization/marlin_utils.py index e972584d5419..e2f63be384b7 100644 --- a/python/sglang/srt/layers/quantization/marlin_utils.py +++ b/python/sglang/srt/layers/quantization/marlin_utils.py @@ -298,6 +298,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/utils/weight_checker.py b/python/sglang/srt/utils/weight_checker.py index 0fdacbca0a0e..2598a84a25cb 100644 --- a/python/sglang/srt/utils/weight_checker.py +++ b/python/sglang/srt/utils/weight_checker.py @@ -343,7 +343,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, diff --git a/python/sglang/srt/utils/weight_checker_comparator.py b/python/sglang/srt/utils/weight_checker_comparator.py index 3578e9c583b9..a7f7191cb6e0 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,136 @@ 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): + 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 +328,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/test/test_marlin_utils.py b/python/sglang/test/test_marlin_utils.py index b87958775bcd..11ce1f436eee 100644 --- a/python/sglang/test/test_marlin_utils.py +++ b/python/sglang/test/test_marlin_utils.py @@ -14,8 +14,9 @@ from sgl_kernel.scalar_type import ScalarType from sglang.srt.layers.quantization.marlin_utils import ( - GPTQ_MARLIN_TILE, + get_weight_perm, marlin_permute_scales, + marlin_permute_weights, marlin_zero_points, ) from sglang.srt.layers.quantization.utils import ( @@ -61,21 +62,6 @@ def __init__(self, out_features, min_thread_n, max_parallel): self.scratch = torch.zeros(max_workspace_size, dtype=torch.int, device="cuda") -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 marlin_weights(q_w, size_k, size_n, num_bits, perm): # Permute q_w = marlin_permute_weights(q_w, size_k, size_n, perm) @@ -95,36 +81,6 @@ def marlin_weights(q_w, size_k, size_n, num_bits, perm): return q_packed -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 = np.array(perm_list) - - if num_bits == 4: - interleave = np.array([0, 2, 4, 6, 1, 3, 5, 7]) - elif num_bits == 8: - interleave = np.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 marlin_quantize( w: torch.Tensor, quant_type: ScalarType, 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..d1986da0f294 --- /dev/null +++ b/test/registered/e2e/rl/test_weight_checker_mxfp4_marlin.py @@ -0,0 +1,222 @@ +"""`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_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() From 1139a6d77a85caaac29615936e1d303d16980623 Mon Sep 17 00:00:00 2001 From: yueming-yuan Date: Thu, 3 Sep 2026 00:12:11 +0000 Subject: [PATCH 29/56] [27/32] [sglang-miles] Kimi K3 LoRA, fused shared experts and in-place MXFP4 reload (#37704, #41074) --- .../kernels/ops/moe/fused_moe_lora_kernel.py | 5 +- .../sglang/kernels/ops/moe/virtual_experts.py | 4 +- python/sglang/srt/arg_groups/field_order.py | 1 + python/sglang/srt/arg_groups/fields/lora.py | 4 + .../moe/fused_moe_triton/fused_marlin_moe.py | 90 +++-- .../layers/quantization/marlin_utils_fp4.py | 179 ++++++++-- .../sglang/srt/layers/quantization/mxfp4.py | 79 ++++- python/sglang/srt/lora/layers.py | 19 +- python/sglang/srt/lora/lora.py | 24 +- python/sglang/srt/lora/lora_manager.py | 28 +- .../sglang/srt/lora/lora_moe_runner_marlin.py | 30 +- python/sglang/srt/lora/lora_moe_runners.py | 16 + python/sglang/srt/lora/mem_pool.py | 243 ++++++++++--- python/sglang/srt/lora/utils.py | 22 +- python/sglang/srt/models/kimi_k3.py | 39 ++- .../srt/utils/weight_checker_comparator.py | 4 +- .../rl/test_mxfp4_marlin_inplace_reload.py | 84 +++++ .../rl/test_weight_checker_mxfp4_marlin.py | 10 + .../kernels/ops/moe/test_moe_wna16_marlin.py | 10 +- .../lora/test_virtual_experts_kernels.py | 10 + test/registered/unit/lora/test_lora_upsert.py | 75 +++- .../unit/lora/test_mem_pool_ep_unit.py | 331 ++++++++++++++++-- .../unit/lora/test_uno_inactive_lora_batch.py | 2 + 23 files changed, 1124 insertions(+), 185 deletions(-) create mode 100644 test/registered/e2e/rl/test_mxfp4_marlin_inplace_reload.py 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 6267c6db4695..f37851f46f3f 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) diff --git a/python/sglang/srt/arg_groups/field_order.py b/python/sglang/srt/arg_groups/field_order.py index 20c35da68093..fd126f2a4061 100644 --- a/python/sglang/srt/arg_groups/field_order.py +++ b/python/sglang/srt/arg_groups/field_order.py @@ -399,6 +399,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/lora.py b/python/sglang/srt/arg_groups/fields/lora.py index c7ea9ea8f033..6857110dc9a3 100644 --- a/python/sglang/srt/arg_groups/fields/lora.py +++ b/python/sglang/srt/arg_groups/fields/lora.py @@ -106,6 +106,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/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/quantization/marlin_utils_fp4.py b/python/sglang/srt/layers/quantization/marlin_utils_fp4.py index 1a4b5d52699e..8cba3ca42868 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 1d25a81febc6..f0e580357252 100644 --- a/python/sglang/srt/layers/quantization/mxfp4.py +++ b/python/sglang/srt/layers/quantization/mxfp4.py @@ -371,6 +371,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, @@ -618,6 +627,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 ( @@ -639,11 +687,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 @@ -1490,6 +1556,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/lora/layers.py b/python/sglang/srt/lora/layers.py index 3ebb9b550a9c..8d406aa2fb49 100644 --- a/python/sglang/srt/lora/layers.py +++ b/python/sglang/srt/lora/layers.py @@ -820,6 +820,10 @@ def forward(self, input_: torch.Tensor, skip_all_reduce=False, forward_batch=Non ) input_parallel = splitted_input[self.base_layer.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) @@ -1024,9 +1028,10 @@ def __init__( 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 @@ -1080,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__}" ) diff --git a/python/sglang/srt/lora/lora.py b/python/sglang/srt/lora/lora.py index 99848d822e5f..4fa687027cd0 100644 --- a/python/sglang/srt/lora/lora.py +++ b/python/sglang/srt/lora/lora.py @@ -60,9 +60,11 @@ 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 assert self.config.hf_config["peft_type"].lower() == "lora" self.base_hf_config: AutoConfig = base_hf_config @@ -161,6 +163,22 @@ def initialize_weights_from_tensors(self, tensors: Dict[str, torch.Tensor]): self._normalize_weights() + 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() + + 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): from sglang.srt.lora.utils import get_normalized_target_modules @@ -175,7 +193,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 @@ -186,14 +204,14 @@ def _process_weight(self, name: str, loaded_weight: torch.Tensor): "all" in normalized_target_modules or module_name in 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" diff --git a/python/sglang/srt/lora/lora_manager.py b/python/sglang/srt/lora/lora_manager.py index 626d537e11a5..e02d130f00fe 100644 --- a/python/sglang/srt/lora/lora_manager.py +++ b/python/sglang/srt/lora/lora_manager.py @@ -128,6 +128,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 @@ -504,6 +505,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 @@ -967,6 +969,7 @@ def _create_lora_adapter_from_tensors( 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) @@ -1079,10 +1082,26 @@ def _load_lora_adapter_from_tensors( error_message=str(e), ) - self.configs[uid] = new_config - self.loras[uid] = new_lora - 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 @@ -1128,6 +1147,8 @@ def _load_lora_adapter_from_tensors( 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. @@ -1153,6 +1174,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 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 1c8b383807bd..a7950694caa0 100644 --- a/python/sglang/srt/lora/lora_moe_runners.py +++ b/python/sglang/srt/lora/lora_moe_runners.py @@ -336,6 +336,15 @@ 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). @@ -418,6 +427,13 @@ 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. diff --git a/python/sglang/srt/lora/mem_pool.py b/python/sglang/srt/lora/mem_pool.py index f556f57faa38..50e49a4ab857 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 @@ -149,7 +150,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 @@ -351,6 +358,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 @@ -375,6 +414,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. @@ -384,7 +424,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 @@ -460,7 +505,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", @@ -566,7 +611,7 @@ def get_lora_B_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 ( "down_proj_moe", "down_proj_shared_moe", @@ -616,9 +661,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) @@ -642,16 +686,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): @@ -790,6 +842,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 "")) @@ -835,6 +888,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] @@ -864,6 +923,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.""" @@ -890,6 +1003,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( @@ -900,21 +1014,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 @@ -983,12 +1101,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 ( @@ -1007,16 +1135,16 @@ 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. # Init A and B INDEPENDENTLY (both buffer and cache_keys). Under @@ -1044,7 +1172,7 @@ def load_lora_weight_tensor( 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 @@ -1058,24 +1186,41 @@ def load_lora_weight_tensor( 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 @@ -1233,7 +1378,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[ @@ -1246,14 +1391,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, @@ -1263,6 +1408,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 @@ -1362,14 +1508,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, @@ -1379,6 +1527,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 @@ -1404,7 +1553,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 @@ -1535,18 +1686,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 5c1b27305cc6..5451fa57d4b8 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. diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py index b04ac8b40cac..c02ed419aa62 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/python/sglang/srt/models/kimi_k3.py @@ -1299,11 +1299,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 @@ -1339,6 +1345,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: @@ -1846,7 +1856,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: @@ -2283,9 +2294,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 if gate_input is not None and not isinstance(x, tuple): @@ -2302,9 +2314,14 @@ def _gated_o_proj_forward(x, *args, **kwargs): # Even a skipped gate must close its capture branch. torch.cuda.current_stream().wait_stream(self._gate_pending_stream) self._gate_pending_stream = None - 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: @@ -3481,13 +3498,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.") diff --git a/python/sglang/srt/utils/weight_checker_comparator.py b/python/sglang/srt/utils/weight_checker_comparator.py index a7f7191cb6e0..2747de51ee04 100644 --- a/python/sglang/srt/utils/weight_checker_comparator.py +++ b/python/sglang/srt/utils/weight_checker_comparator.py @@ -146,7 +146,9 @@ class Mxfp4MarlinComparable(ComparableWeight): # (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): + 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}" ) 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 index d1986da0f294..d9afa8d1683a 100644 --- a/test/registered/e2e/rl/test_weight_checker_mxfp4_marlin.py +++ b/test/registered/e2e/rl/test_weight_checker_mxfp4_marlin.py @@ -166,6 +166,16 @@ def test_a_one_ulp_change_stays_within_the_quantization_tolerance(self): 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() 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/unit/lora/test_lora_upsert.py b/test/registered/unit/lora/test_lora_upsert.py index f9e7b805c971..6bb08e852798 100644 --- a/test/registered/unit/lora/test_lora_upsert.py +++ b/test/registered/unit/lora/test_lora_upsert.py @@ -45,13 +45,14 @@ def test_returns_none_for_unregistered_adapter(self): self.assertIsNone(asyncio.run(registry.get_lora_id("missing"))) -def _make_manager() -> LoRAManager: +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 = [] @@ -258,6 +259,78 @@ def test_failed_buffer_rewrite_restores_old_weights(self): 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() 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 7f73273bc908..211f81207138 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 @@ -111,7 +121,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) @@ -196,6 +212,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 @@ -207,6 +224,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 @@ -223,6 +255,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 @@ -264,6 +298,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, @@ -329,6 +365,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)]} @@ -637,24 +674,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)], @@ -663,11 +833,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) @@ -704,6 +869,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( @@ -712,15 +988,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)], @@ -733,11 +1002,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) @@ -818,15 +1082,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)], @@ -835,11 +1092,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), @@ -1101,6 +1353,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`. @@ -1123,6 +1378,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 @@ -1242,6 +1498,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 @@ -1327,6 +1584,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) From 97586dc5938a4ae06e4b506483b84f00a5d64fda Mon Sep 17 00:00:00 2001 From: yueming-yuan Date: Fri, 18 Sep 2026 21:56:09 +0000 Subject: [PATCH 30/56] [28/32] [sglang-miles] Resolve HF LoRA targets through model-aware normalization (#40242) --- python/sglang/kernels/ops/moe/virtual_experts.py | 4 +++- python/sglang/srt/arg_groups/fields/lora.py | 11 +++++------ .../srt/arg_groups/model_overrides/qwen3_moe.py | 10 ++++++++-- python/sglang/srt/lora/lora.py | 14 ++++++-------- python/sglang/srt/lora/lora_manager.py | 13 +++++++++++-- python/sglang/srt/lora/mem_pool.py | 4 +++- python/sglang/srt/lora/utils.py | 13 ++++++++++++- python/sglang/srt/models/inkling.py | 14 ++++++++++++++ python/sglang/srt/models/kimi_k3.py | 9 +++++++++ 9 files changed, 71 insertions(+), 21 deletions(-) diff --git a/python/sglang/kernels/ops/moe/virtual_experts.py b/python/sglang/kernels/ops/moe/virtual_experts.py index f37851f46f3f..74e22ef02549 100644 --- a/python/sglang/kernels/ops/moe/virtual_experts.py +++ b/python/sglang/kernels/ops/moe/virtual_experts.py @@ -679,9 +679,10 @@ def _get_routing( 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( [topk_ids.shape[0], topk_ids.shape[1], max_lora_rank], - dtype=hidden_states.dtype, + dtype=torch.float32, device=hidden_states.device, ) @@ -711,6 +712,7 @@ def _get_routing( a_stage_config, ) + intermediate = intermediate.to(hidden_states.dtype) b_stage_config = _get_stage_config(lora_b_virtuals[0], 1) ( sorted_token_ids, diff --git a/python/sglang/srt/arg_groups/fields/lora.py b/python/sglang/srt/arg_groups/fields/lora.py index 6857110dc9a3..d8ddc88f7c4d 100644 --- a/python/sglang/srt/arg_groups/fields/lora.py +++ b/python/sglang/srt/arg_groups/fields/lora.py @@ -17,10 +17,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): @@ -46,9 +42,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[ 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 f7d3ca012874..82dc9906f1e8 100644 --- a/python/sglang/srt/arg_groups/model_overrides/qwen3_moe.py +++ b/python/sglang/srt/arg_groups/model_overrides/qwen3_moe.py @@ -83,9 +83,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/lora/lora.py b/python/sglang/srt/lora/lora.py index 4fa687027cd0..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 @@ -66,6 +67,9 @@ def __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 @@ -180,12 +184,6 @@ def release_staged_weights(self): self.pinned_added_tokens_embeddings.clear() def _process_weight(self, name: str, loaded_weight: torch.Tensor): - from sglang.srt.lora.utils import get_normalized_target_modules - - normalized_target_modules = get_normalized_target_modules( - self.config.target_modules - ) - # Remap PEFT "unembed_tokens" key to "lm_head" so the weight is # recognized and loaded into the correct buffer. if "unembed_tokens" in name: @@ -201,8 +199,8 @@ 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] = self._stage_weight(loaded_weight) else: diff --git a/python/sglang/srt/lora/lora_manager.py b/python/sglang/srt/lora/lora_manager.py index e02d130f00fe..40d8c586bcb3 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, @@ -822,7 +823,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() @@ -866,7 +869,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: @@ -884,6 +887,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 diff --git a/python/sglang/srt/lora/mem_pool.py b/python/sglang/srt/lora/mem_pool.py index 50e49a4ab857..f1d61a7b521b 100644 --- a/python/sglang/srt/lora/mem_pool.py +++ b/python/sglang/srt/lora/mem_pool.py @@ -257,7 +257,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) diff --git a/python/sglang/srt/lora/utils.py b/python/sglang/srt/lora/utils.py index 5451fa57d4b8..e982d1d01123 100644 --- a/python/sglang/srt/lora/utils.py +++ b/python/sglang/srt/lora/utils.py @@ -274,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. @@ -299,6 +300,11 @@ 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", @@ -321,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) @@ -476,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 diff --git a/python/sglang/srt/models/inkling.py b/python/sglang/srt/models/inkling.py index e9eeb9e616e4..675f89c7adeb 100644 --- a/python/sglang/srt/models/inkling.py +++ b/python/sglang/srt/models/inkling.py @@ -1111,6 +1111,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_k3.py b/python/sglang/srt/models/kimi_k3.py index c02ed419aa62..dc9e9a123d1f 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/python/sglang/srt/models/kimi_k3.py @@ -3555,6 +3555,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, From 9a459031679238b053feb1be78d86eb3d06f2295 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Wed, 23 Sep 2026 01:45:37 +0000 Subject: [PATCH 31/56] [29/32] [sglang-miles] Flush and synchronize hybrid caches before offload; restore bf16 MoE layout (#40821) Co-authored-by: Yueming Yuan --- python/sglang/srt/layers/attention/qsa/mqa.py | 6 +- .../sglang/srt/layers/quantization/unquant.py | 2 +- python/sglang/srt/managers/scheduler.py | 4 + .../scheduler_components/weight_updater.py | 13 ++- .../test_release_memory_flush_order.py | 96 +++++++++++++++++++ 5 files changed, 112 insertions(+), 9 deletions(-) create mode 100644 test/registered/unit/managers/test_release_memory_flush_order.py diff --git a/python/sglang/srt/layers/attention/qsa/mqa.py b/python/sglang/srt/layers/attention/qsa/mqa.py index 88e201ee060a..78e27a33cc6f 100644 --- a/python/sglang/srt/layers/attention/qsa/mqa.py +++ b/python/sglang/srt/layers/attention/qsa/mqa.py @@ -122,7 +122,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") @@ -152,9 +151,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/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index 8edd7e72405f..31f63142197f 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/python/sglang/srt/layers/quantization/unquant.py @@ -918,7 +918,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/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index f24f4984c47b..fe1cd0584fbf 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -74,6 +74,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 import role_switch @@ -5188,6 +5189,9 @@ 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 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 diff --git a/python/sglang/srt/managers/scheduler_components/weight_updater.py b/python/sglang/srt/managers/scheduler_components/weight_updater.py index f4f68eff5a59..8e702a226c13 100644 --- a/python/sglang/srt/managers/scheduler_components/weight_updater.py +++ b/python/sglang/srt/managers/scheduler_components/weight_updater.py @@ -518,9 +518,6 @@ def release_memory_occupation(self, recv_req: ReleaseMemoryOccupationReqInput): 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: if scheduler is not None: if scheduler.disaggregation_mode == DisaggregationMode.DECODE: @@ -535,9 +532,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( 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() From aea827bec6be81a7d510bdc5b8105b0221b2a085 Mon Sep 17 00:00:00 2001 From: Kangrui Du <89372739+Rockdu@users.noreply.github.com> Date: Thu, 20 Aug 2026 02:20:03 +0000 Subject: [PATCH 32/56] [30/32] [sglang-miles] H3 rollout trajectories, response timing and IPC LoRA updates (#40786) Keep H3 rollout and streamed-weight behavior on the release implementation. The release already contains the LoRA wrapper type guards. Co-authored-by: WenhaoZhang <42087078+niehen6174@users.noreply.github.com> Co-authored-by: niehen6174 --- .../configs/sample/minimax_h3.py | 13 +- .../entrypoints/post_training/io_struct.py | 2 + .../post_training/request_timing.py | 83 ++++++ .../entrypoints/post_training/rollout_api.py | 73 +++++- .../entrypoints/post_training/utils.py | 53 ++++ .../runtime/pipelines_core/lora/pipeline.py | 8 +- .../minimax_h3/denoise_loop.py | 53 +++- .../minimax_h3/minimax_h3_rollout.py | 242 ++++++++++++++++++ .../minimax_h3/stages/denoising.py | 53 ++++ .../runtime/post_training/rl_dataclasses.py | 5 +- .../post_training/rollout_denoising_mixin.py | 7 +- .../runtime/post_training/weights_updater.py | 30 ++- .../test/unit/test_fused_lora_compose.py | 43 ++++ .../test/unit/test_msgpack_splice.py | 67 +++++ .../test/unit/test_rollout_api.py | 77 +++++- python/sglang/srt/managers/scheduler.py | 1 + .../scheduler_components/output_streamer.py | 3 + 17 files changed, 772 insertions(+), 41 deletions(-) create mode 100644 python/sglang/multimodal_gen/runtime/entrypoints/post_training/request_timing.py create mode 100644 python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/minimax_h3_rollout.py create mode 100644 python/sglang/multimodal_gen/test/unit/test_msgpack_splice.py diff --git a/python/sglang/multimodal_gen/configs/sample/minimax_h3.py b/python/sglang/multimodal_gen/configs/sample/minimax_h3.py index a6c158c73866..855c04002edd 100644 --- a/python/sglang/multimodal_gen/configs/sample/minimax_h3.py +++ b/python/sglang/multimodal_gen/configs/sample/minimax_h3.py @@ -262,14 +262,15 @@ def _validate(self) -> None: self.spectrum_params = SpectrumParams() apply_h3_spectrum_param_defaults(self.spectrum_params) 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/pipelines_core/lora/pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines_core/lora/pipeline.py index 9a12ef0799e7..a61b0306238b 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/lora/pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/lora/pipeline.py @@ -58,12 +58,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 @@ -1016,7 +1014,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 1b7b65c8467c..2203461e1e36 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 @@ -783,6 +783,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( @@ -821,7 +869,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 012c69c1fc9f..59858a27a7c6 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 e6ef5094ce19..f8317a106f66 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py +++ b/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py @@ -68,6 +68,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 @@ -266,12 +267,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 @@ -712,7 +719,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: @@ -724,6 +731,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/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index fe1cd0584fbf..3eafec682506 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2467,6 +2467,7 @@ def init_output_streamer(self) -> None: tree_cache=self.tree_cache, 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: ( diff --git a/python/sglang/srt/managers/scheduler_components/output_streamer.py b/python/sglang/srt/managers/scheduler_components/output_streamer.py index 0bb77a8c6885..c394a9f64979 100644 --- a/python/sglang/srt/managers/scheduler_components/output_streamer.py +++ b/python/sglang/srt/managers/scheduler_components/output_streamer.py @@ -62,6 +62,7 @@ class SchedulerOutputStreamer: # Rust egress ring via `rust_server.push_generation` instead of the zmq # detokenizer. None otherwise. (Rust-specific state lives in RustServer.) rust_server: Optional[RustServer] = None + is_multimodal_gen: bool = False _test_stream_output_count: int = 0 def __post_init__(self) -> None: @@ -205,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=get_parallel().dp_rank, is_idle_batch=is_idle_batch, From 8035002c4ed0360579004791378e5b7d1740503b Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Mon, 21 Sep 2026 22:40:43 +0000 Subject: [PATCH 33/56] [31/32] [sglang-miles] MiMo colocated RL loaders, pausable caches and graph registration (#40648, #40978, #40979, #40980, #40977, #41972) Retain colocated training fixes above the release MiMo implementation, including address-stable Marlin reload and mandatory TopK layer IDs. Disable NCCL graph registration for TP LM-head all-to-all and pausable graph pools. Co-authored-by: Claude Fable 5.1 --- .../benchmark_sglang_fused_moe_triton.py | 1 + python/sglang/srt/arg_groups/parallel_hook.py | 41 +++- python/sglang/srt/configs/model_config.py | 8 +- python/sglang/srt/layers/moe/topk.py | 4 +- .../layers/quantization/mxfp4_marlin_moe.py | 77 +++++++- .../sglang/srt/mem_cache/index_key_cache.py | 28 +-- .../srt/mem_cache/kv_cache_configurator.py | 3 +- .../sglang/srt/mem_cache/swa_memory_pool.py | 3 +- python/sglang/srt/models/afmoe.py | 3 + python/sglang/srt/models/bailing_moe.py | 1 + .../sglang/srt/models/bailing_moe_linear.py | 1 + python/sglang/srt/models/dbrx.py | 4 +- python/sglang/srt/models/deepseek.py | 3 + python/sglang/srt/models/ernie45_moe_vl.py | 2 + python/sglang/srt/models/exaone_moe.py | 1 + python/sglang/srt/models/granitemoe.py | 1 + python/sglang/srt/models/hunyuan_v3.py | 1 + python/sglang/srt/models/kimi_k3.py | 1 + python/sglang/srt/models/kimi_linear.py | 1 + python/sglang/srt/models/llada2.py | 1 + python/sglang/srt/models/llama4.py | 1 + python/sglang/srt/models/mimo_v2.py | 36 +++- python/sglang/srt/models/minimax_m2.py | 1 + python/sglang/srt/models/mixtral.py | 1 + python/sglang/srt/models/olmoe.py | 1 + python/sglang/srt/models/phimoe.py | 1 + python/sglang/srt/models/xverse_moe.py | 3 + test/registered/moe/test_triton_fused_moe.py | 1 + .../models/test_mimo_v2_fused_qkv_loader.py | 177 ++++++++++++++++++ 29 files changed, 371 insertions(+), 36 deletions(-) create mode 100644 test/registered/unit/models/test_mimo_v2_fused_qkv_loader.py diff --git a/benchmark/kernels/fused_moe_triton/benchmark_sglang_fused_moe_triton.py b/benchmark/kernels/fused_moe_triton/benchmark_sglang_fused_moe_triton.py index 6934650adb92..1dabf2f3582c 100644 --- a/benchmark/kernels/fused_moe_triton/benchmark_sglang_fused_moe_triton.py +++ b/benchmark/kernels/fused_moe_triton/benchmark_sglang_fused_moe_triton.py @@ -38,6 +38,7 @@ def fused_moe_triton_api( ): topk_op = TopK( top_k=topk, + layer_id=0, renormalize=False, use_grouped_topk=False, output_format=TopKOutputFormat.TRITON_KERNEL, diff --git a/python/sglang/srt/arg_groups/parallel_hook.py b/python/sglang/srt/arg_groups/parallel_hook.py index 1a396935fbf1..a63ce623290c 100644 --- a/python/sglang/srt/arg_groups/parallel_hook.py +++ b/python/sglang/srt/arg_groups/parallel_hook.py @@ -302,12 +302,29 @@ def handle_data_parallelism(server_args: Any): run_post_process_pass(server_args, _tp_lm_head_all_to_all_default) run_post_process_pass(server_args, _dp_lm_head_validation) if resolving_view(server_args).enable_tp_lm_head_all_to_all: - _disable_nccl_graph_buffer_registration() + _disable_nccl_graph_buffer_registration( + "the graph-captured TP LM-head all-to-all can deadlock with " + "registered buffers" + ) + if _graph_pool_is_pausable(server_args): + _disable_nccl_graph_buffer_registration( + "graph replay can hang once torch_memory_saver resume remaps the " + "graph pool, leaving the capture-time registrations on released pages" + ) -def _disable_nccl_graph_buffer_registration() -> None: - """Keep NCCL from registering the buffers of the graph-captured PyNccl - all-to-all. +def _graph_pool_is_pausable(server_args: Any) -> bool: + """Whether CUDA graphs are captured into torch_memory_saver memory, which + release/resume of the `cuda_graph` tag remaps to new physical pages at the + same virtual addresses.""" + return bool( + resolving_view(server_args).enable_memory_saver + and envs.SGLANG_MEMORY_SAVER_CUDA_GRAPH.get() + ) + + +def _disable_nccl_graph_buffer_registration(reason: str) -> None: + """Keep NCCL from registering the buffers of graph-captured collectives. NCCL_GRAPH_REGISTER (default on) registers the send/recv buffers of every collective captured in a CUDA graph for the lifetime of the graph, and @@ -319,15 +336,21 @@ def _disable_nccl_graph_buffer_registration() -> None: spin in ncclDevKernel_SendRecv forever, and every DP rank hangs. Reproduced on tp4/dp4/ep4 and on a multi-node tp16/dp16/ep16 PD decode deployment; disabling the registration removes the hang while dedicated - all-to-all buffers alone do not. Must run before the schedulers create - their NCCL communicators, which inherit this environment. An explicit - setting wins. + all-to-all buffers alone do not. + + A pausable graph pool breaks the registrations too: torch_memory_saver + resume keeps the pool's virtual addresses but maps new physical pages, + while NCCL's capture-time registrations still point at the released pages, + so graph replay can hang (fzyzcjy/torch_memory_saver#88). + + Must run before the schedulers create their NCCL communicators, which + inherit this environment. An explicit setting wins. """ if os.environ.setdefault("NCCL_GRAPH_REGISTER", "0") != "0": logger.warning( - "NCCL_GRAPH_REGISTER=%s was set explicitly; the graph-captured TP " - "LM-head all-to-all can deadlock with registered buffers.", + "NCCL_GRAPH_REGISTER=%s was set explicitly; %s.", os.environ["NCCL_GRAPH_REGISTER"], + reason, ) diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 77ab4a863fd3..acd086833d7c 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -121,13 +121,15 @@ def unwrap_modelopt_quantization_config(quant_config: dict) -> dict: def get_mimo_v2_fused_qkv_expected_tp_size(hf_config): layout = getattr(hf_config, "attention_projection_layout", None) - if layout is None: + # MiMoV2Config fills in "split" when the checkpoint omits the field; split + # q/k/v projections load through the regular stacked-params path. + if layout in (None, "split"): return None if layout != "fused_qkv": raise ValueError( "MiMoV2 hf_config has unsupported " - f"attention_projection_layout={layout!r}; expected 'fused_qkv' " - "or unset." + f"attention_projection_layout={layout!r}; expected 'fused_qkv', " + "'split' or unset." ) num_key_value_heads = getattr(hf_config, "num_key_value_heads", None) diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 60767eb65bbb..d806c3d8f635 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -537,6 +537,8 @@ class TopK(BaseFusedOp): """ Parameters: --top_k: The all number of top experts selected per token, including the fused shared expert(s). + --layer_id: global index of the owning layer. Required: routed-experts capture and + per-layer EPLB state are keyed by it. --num_fused_shared_experts: num of shared experts, can be activate both in TP or EP mode. --routed_scaling_factor: the scaling factor for routed experts in topk_weights. --fused_shared_experts_scaling_factor: scaling factor applied to the fused shared experts' @@ -548,7 +550,7 @@ def __init__( self, top_k: int, *, - layer_id: Optional[int] = None, + layer_id: int, use_grouped_topk: bool = False, topk_group: Optional[int] = None, num_expert_group: Optional[int] = None, diff --git a/python/sglang/srt/layers/quantization/mxfp4_marlin_moe.py b/python/sglang/srt/layers/quantization/mxfp4_marlin_moe.py index 52503455e076..8557ab5947e1 100644 --- a/python/sglang/srt/layers/quantization/mxfp4_marlin_moe.py +++ b/python/sglang/srt/layers/quantization/mxfp4_marlin_moe.py @@ -46,6 +46,14 @@ class Mxfp4MarlinMoEMethod: """MXFP4 (E8M0 scales) MoE quantization method using the Marlin backend.""" fuse_routed_scaling_factor_in_topk = True + # The loader writes the scales as *_weight_scale_inv (as for Fp8MoEMethod); Marlin keeps + # the repacked scales as *_weight_scale. + _LOADER_TO_RUNTIME_NAMES = { + "w13_weight": "w13_weight", + "w13_weight_scale_inv": "w13_weight_scale", + "w2_weight": "w2_weight", + "w2_weight_scale_inv": "w2_weight_scale", + } def __init__(self, fp8_method, prefix: str): self._fp8 = fp8_method @@ -127,6 +135,46 @@ def _e8m0_ones(*shape: int) -> torch.Tensor: layer.register_parameter("w2_weight_scale_inv", w2_weight_scale) set_weight_attrs(w2_weight_scale, scale_attrs) + layer._mxfp4_marlin_scale_attrs = scale_attrs + layer._mxfp4_marlin_checkpoint_parameter_specs = { + name: (tuple(getattr(layer, name).shape), getattr(layer, name).dtype) + for name in self._LOADER_TO_RUNTIME_NAMES + } + layer._mxfp4_marlin_runtime_parameter_specs = None + layer._mxfp4_marlin_reload_in_place = False + + def restore_weights_before_loading(self, layer: Module) -> None: + from sglang.srt.layers.quantization.marlin_utils_fp4 import ( + _view_tensor_storage, + ) + + # A weight update loads the checkpoint layout into the Marlin storage and repacks it there: + # CUDA graphs and the memory saver hold these parameters' addresses. + runtime_specs = layer._mxfp4_marlin_runtime_parameter_specs + if runtime_specs is None: + raise RuntimeError( + f"MXFP4 Marlin weights of {self.prefix} were never post-processed." + ) + for loader_name, runtime_name in self._LOADER_TO_RUNTIME_NAMES.items(): + shape, dtype = layer._mxfp4_marlin_checkpoint_parameter_specs[loader_name] + param = getattr(layer, runtime_name) + current_spec = (tuple(param.shape), param.dtype) + if current_spec not in ((shape, dtype), runtime_specs[runtime_name]): + raise ValueError( + f"Unexpected MXFP4 Marlin parameter state for {runtime_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_() + if loader_name != runtime_name: + # the loader looks the scales up by their checkpoint name + alias = torch.nn.Parameter(param.data, requires_grad=False) + set_weight_attrs(alias, layer._mxfp4_marlin_scale_attrs) + alias.format_ue8m0 = False + layer.register_parameter(loader_name, alias) + layer._mxfp4_marlin_reload_in_place = True + def process_weights_after_loading(self, layer: Module) -> None: from sglang.srt.layers.quantization.marlin_utils import ( check_moe_marlin_supports_layer, @@ -157,13 +205,38 @@ def process_weights_after_loading(self, layer: Module) -> None: # Unlike the flashinfer trtllm_fp4 kernel (which wants [w3, w1]), # we must *not* call ``reorder_w1w3_to_w3w1`` here. + reuse_parameter_storage = layer._mxfp4_marlin_reload_in_place + if reuse_parameter_storage: + # the repack works on the runtime names, which share storage with these aliases + for loader_name, runtime_name in self._LOADER_TO_RUNTIME_NAMES.items(): + if loader_name != runtime_name: + delattr(layer, loader_name) + log_info_on_rank0( logger, f"Preparing MXFP4 experts for Marlin backend (layer: {self.prefix})...", ) if self.runner.config.gemm1_alpha is not None: - 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._LOADER_TO_RUNTIME_NAMES.values() + } + 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._dsv4_mxfp4_backend = "marlin" def apply( diff --git a/python/sglang/srt/mem_cache/index_key_cache.py b/python/sglang/srt/mem_cache/index_key_cache.py index 0975e177a343..a1036472fdc6 100644 --- a/python/sglang/srt/mem_cache/index_key_cache.py +++ b/python/sglang/srt/mem_cache/index_key_cache.py @@ -6,6 +6,7 @@ import torch from sglang.kernels.ops.attention.dsa import index_buf_accessor +from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE if TYPE_CHECKING: from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool @@ -15,19 +16,20 @@ class IndexKeyCache: def __init__(self, pool: DSATokenToKVPool, index_buf_size: int): self.pool = pool num_pages = (index_buf_size + pool.page_size + 1) // pool.page_size - with ( - torch.cuda.use_mem_pool(pool.custom_mem_pool) - if pool.custom_mem_pool - else nullcontext() - ): - self.buffer = [ - torch.zeros( - self._buffer_shape(self._layer_num_pages(i, num_pages)), - dtype=pool.index_k_with_scale_buffer_dtype, - device=pool.device, - ) - for i in range(pool.layer_num) - ] + with pool.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE): + with ( + torch.cuda.use_mem_pool(pool.custom_mem_pool) + if pool.custom_mem_pool + else nullcontext() + ): + self.buffer = [ + torch.zeros( + self._buffer_shape(self._layer_num_pages(i, num_pages)), + dtype=pool.index_k_with_scale_buffer_dtype, + device=pool.device, + ) + for i in range(pool.layer_num) + ] def _buffer_shape(self, num_pages: int) -> tuple[int, int]: pool = self.pool diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 13210daa44d3..2094584b1736 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -1695,7 +1695,7 @@ def _build_hybrid_mla_swa_kv_pool( common = { "page_size": get_schedule().page_size, "device": self.device, - "enable_memory_saver": False, + "enable_memory_saver": get_exec().features.enable_memory_saver, } full_pool_kwargs = { **common, @@ -1808,6 +1808,7 @@ def _build_hybrid_swa_kv_pool( device=self.device, enable_kv_cache_copy=(get_spec().speculative_algorithm is not None), token_to_kv_pool_class=swa_pool_class, + enable_memory_saver=get_exec().features.enable_memory_saver, **kwargs, ) return token_to_kv_pool diff --git a/python/sglang/srt/mem_cache/swa_memory_pool.py b/python/sglang/srt/mem_cache/swa_memory_pool.py index 93db9fe06457..bf037bdad4c4 100644 --- a/python/sglang/srt/mem_cache/swa_memory_pool.py +++ b/python/sglang/srt/mem_cache/swa_memory_pool.py @@ -42,6 +42,7 @@ def __init__( swa_kv_pool_class: Optional[type] = None, full_kv_pool_kwargs: Optional[dict] = None, swa_kv_pool_kwargs: Optional[dict] = None, + enable_memory_saver: bool = False, **kwargs, ): self.size = size @@ -66,7 +67,7 @@ def __init__( swa_kv_pool_class = swa_kv_pool_class or token_to_kv_pool_class common_kwargs = { "page_size": page_size, - "enable_memory_saver": False, + "enable_memory_saver": enable_memory_saver, "device": device, } if full_kv_pool_kwargs is None: diff --git a/python/sglang/srt/models/afmoe.py b/python/sglang/srt/models/afmoe.py index ed5521f77174..e21b3e47ce72 100644 --- a/python/sglang/srt/models/afmoe.py +++ b/python/sglang/srt/models/afmoe.py @@ -152,6 +152,7 @@ def _custom_routing_function( def __init__( self, config: PretrainedConfig, + layer_id: int, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): @@ -234,6 +235,7 @@ def __init__( renormalize = self.route_norm if self.score_func == "sigmoid" else False self.topk = TopK( top_k=self.top_k, + layer_id=layer_id, renormalize=renormalize, use_grouped_topk=self.use_grouped_topk, num_expert_group=self.n_group if self.use_grouped_topk else None, @@ -455,6 +457,7 @@ def __init__( if use_moe: self.mlp = AfmoeMoE( config=config, + layer_id=layer_id, quant_config=quant_config, prefix=add_prefix("mlp", prefix), ) diff --git a/python/sglang/srt/models/bailing_moe.py b/python/sglang/srt/models/bailing_moe.py index 4b77564048d2..b485de6bd156 100644 --- a/python/sglang/srt/models/bailing_moe.py +++ b/python/sglang/srt/models/bailing_moe.py @@ -285,6 +285,7 @@ def __init__( self.topk = TopK( top_k=self.top_k, + layer_id=self.layer_id, renormalize=self.norm_topk_prob, use_grouped_topk=self.use_grouped_topk, num_expert_group=self.num_expert_group, diff --git a/python/sglang/srt/models/bailing_moe_linear.py b/python/sglang/srt/models/bailing_moe_linear.py index 7c6ee7a1f47b..097b42586a85 100644 --- a/python/sglang/srt/models/bailing_moe_linear.py +++ b/python/sglang/srt/models/bailing_moe_linear.py @@ -300,6 +300,7 @@ def __init__( self.topk = TopK( top_k=self.top_k, + layer_id=self.layer_id, use_grouped_topk=self.use_grouped_topk, renormalize=self.norm_expert_prob, num_expert_group=self.num_expert_group, diff --git a/python/sglang/srt/models/dbrx.py b/python/sglang/srt/models/dbrx.py index d0425e603f8f..18037900beae 100644 --- a/python/sglang/srt/models/dbrx.py +++ b/python/sglang/srt/models/dbrx.py @@ -97,6 +97,7 @@ class DbrxExperts(nn.Module): def __init__( self, config: DbrxConfig, + layer_id: int, quant_config: Optional[QuantizationConfig] = None, params_dtype: Optional[torch.dtype] = None, prefix: str = "", @@ -115,6 +116,7 @@ def __init__( self.router = DbrxRouter(config, self.params_dtype) self.topk = TopK( self.top_k, + layer_id=layer_id, renormalize=True, ) self.moe_runner_config = MoeRunnerConfig(inplace=True) @@ -336,7 +338,7 @@ def __init__( quant_config=quant_config, prefix=add_prefix("norm_attn_norm", prefix), ) - self.ffn = DbrxExperts(config, quant_config=quant_config) + self.ffn = DbrxExperts(config, layer_id, quant_config=quant_config) def forward( self, diff --git a/python/sglang/srt/models/deepseek.py b/python/sglang/srt/models/deepseek.py index 6bb1996b60e4..a28f2a0a4db6 100644 --- a/python/sglang/srt/models/deepseek.py +++ b/python/sglang/srt/models/deepseek.py @@ -109,6 +109,7 @@ class DeepseekMoE(nn.Module): def __init__( self, config: PretrainedConfig, + layer_id: int, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): @@ -125,6 +126,7 @@ def __init__( ) self.topk = TopK( top_k=self.top_k, + layer_id=layer_id, renormalize=config.norm_topk_prob, ) self.experts = nn.ModuleList( @@ -339,6 +341,7 @@ def __init__( ): self.mlp = DeepseekMoE( config=config, + layer_id=layer_id, quant_config=quant_config, prefix=add_prefix("mlp", prefix), ) diff --git a/python/sglang/srt/models/ernie45_moe_vl.py b/python/sglang/srt/models/ernie45_moe_vl.py index f15071207bc0..cdadc9314f45 100644 --- a/python/sglang/srt/models/ernie45_moe_vl.py +++ b/python/sglang/srt/models/ernie45_moe_vl.py @@ -206,6 +206,7 @@ def __init__( self.text_experts_topk = TopK( top_k=config.moe_k, + layer_id=self.layer_id, renormalize=True, use_grouped_topk=False, correction_bias=self.e_score_correction_bias[0], @@ -237,6 +238,7 @@ def __init__( self.vision_experts_topk = TopK( top_k=config.moe_k, + layer_id=self.layer_id, renormalize=True, use_grouped_topk=False, correction_bias=self.e_score_correction_bias[1], diff --git a/python/sglang/srt/models/exaone_moe.py b/python/sglang/srt/models/exaone_moe.py index 211a12de392c..01a488b7e84e 100755 --- a/python/sglang/srt/models/exaone_moe.py +++ b/python/sglang/srt/models/exaone_moe.py @@ -172,6 +172,7 @@ def __init__( self.topk = TopK( top_k=config.num_experts_per_tok, + layer_id=self.layer_id, renormalize=config.norm_topk_prob, use_grouped_topk=True, num_expert_group=config.n_group, diff --git a/python/sglang/srt/models/granitemoe.py b/python/sglang/srt/models/granitemoe.py index 5c6d4d91c797..0b5090a951f5 100644 --- a/python/sglang/srt/models/granitemoe.py +++ b/python/sglang/srt/models/granitemoe.py @@ -134,6 +134,7 @@ def __init__( self.topk = TopK( top_k=top_k, + layer_id=layer_id, renormalize=True, ) diff --git a/python/sglang/srt/models/hunyuan_v3.py b/python/sglang/srt/models/hunyuan_v3.py index f95840dd9976..7312e386258c 100644 --- a/python/sglang/srt/models/hunyuan_v3.py +++ b/python/sglang/srt/models/hunyuan_v3.py @@ -120,6 +120,7 @@ def __init__( ) self.topk = TopK( top_k=config.num_experts_per_tok, + layer_id=layer_id, use_grouped_topk=True, num_expert_group=1, topk_group=1, diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py index dc9e9a123d1f..e732cd1fbcdb 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/python/sglang/srt/models/kimi_k3.py @@ -489,6 +489,7 @@ def __init__( self.topk = TopK( top_k=config.num_experts_per_token, + layer_id=self.layer_idx, renormalize=moe_renormalize, use_grouped_topk=True, num_expert_group=config.num_expert_group, diff --git a/python/sglang/srt/models/kimi_linear.py b/python/sglang/srt/models/kimi_linear.py index dc4d6f890ab6..115327c057b6 100644 --- a/python/sglang/srt/models/kimi_linear.py +++ b/python/sglang/srt/models/kimi_linear.py @@ -120,6 +120,7 @@ def __init__( self.topk = TopK( top_k=config.num_experts_per_token, + layer_id=self.layer_idx, renormalize=moe_renormalize, use_grouped_topk=True, num_expert_group=config.num_expert_group, diff --git a/python/sglang/srt/models/llada2.py b/python/sglang/srt/models/llada2.py index 129e3b9525aa..5e10a9e2956a 100644 --- a/python/sglang/srt/models/llada2.py +++ b/python/sglang/srt/models/llada2.py @@ -259,6 +259,7 @@ def __init__( self.topk = TopK( top_k=self.top_k, + layer_id=self.layer_id, renormalize=self.norm_topk_prob, use_grouped_topk=self.use_grouped_topk, num_expert_group=self.num_expert_group, diff --git a/python/sglang/srt/models/llama4.py b/python/sglang/srt/models/llama4.py index 81aeb414a28e..48d1e007db35 100644 --- a/python/sglang/srt/models/llama4.py +++ b/python/sglang/srt/models/llama4.py @@ -113,6 +113,7 @@ def __init__( self.topk = TopK( top_k=self.top_k, + layer_id=layer_id, renormalize=False, custom_routing_function=Llama4MoE.custom_routing_function, ) diff --git a/python/sglang/srt/models/mimo_v2.py b/python/sglang/srt/models/mimo_v2.py index de9bc0bef35c..4c28970bf0f7 100644 --- a/python/sglang/srt/models/mimo_v2.py +++ b/python/sglang/srt/models/mimo_v2.py @@ -100,6 +100,7 @@ def load_mimo_v2_qkv_proj_weight( loaded_weight, expected_fused_tp_size: Optional[int] = None, deferred_scale_inv: Optional[Dict[str, torch.Tensor]] = None, + config=None, ): tp_size = get_parallel().attn_tp_size tp_rank = get_parallel().attn_tp_rank @@ -144,11 +145,34 @@ def load_mimo_v2_qkv_proj_weight( default_weight_loader(param, loaded_weight.chunk(tp_size, dim=0)[tp_rank]) else: shards_per_rank = ckpt_tp // tp_size - shards = loaded_weight.chunk(ckpt_tp, dim=0) - merged = torch.cat( - shards[tp_rank * shards_per_rank : (tp_rank + 1) * shards_per_rank], - dim=0, - ) + shards = loaded_weight.chunk(ckpt_tp, dim=0)[ + tp_rank * shards_per_rank : (tp_rank + 1) * shards_per_rank + ] + if loaded_weight.dtype in (torch.float8_e4m3fn, torch.float8_e4m3fnuz): + # Block-FP8 codes keep the checkpoint shard order; + # _resolve_deferred_qkv_scale_inv regroups them with their scales. + merged = torch.cat(shards, dim=0) + else: + qkv_sizes = ( + _get_ckpt_qkv_shard_sizes( + config=config, layer_name=name, ckpt_tp=ckpt_tp + ) + if config is not None + else None + ) + if qkv_sizes is None: + raise ValueError( + f"qkv_proj weight {name}: attention TP {tp_size} below the " + f"checkpoint's {ckpt_tp} kv-head shards needs the decoder layer " + f"config to regroup q|k|v; MTP layers are not supported here" + ) + q_per_shard, k_per_shard, v_per_shard = qkv_sizes + merged = _deinterleave_qkv_shards( + shards, + q_per_shard=q_per_shard, + k_per_shard=k_per_shard, + v_per_shard=v_per_shard, + ) default_weight_loader(param, merged) @@ -427,6 +451,7 @@ def __init__( self.topk = TopK( top_k=config.num_experts_per_tok, + layer_id=self.layer_id, renormalize=config.norm_topk_prob, use_grouped_topk=True, num_expert_group=config.n_group, @@ -1572,6 +1597,7 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): loaded_weight, expected_fused_tp_size, deferred_scale_inv=deferred_qkv_scale_inv, + config=self.config, ) continue diff --git a/python/sglang/srt/models/minimax_m2.py b/python/sglang/srt/models/minimax_m2.py index 0beecf8eef2a..39a3539cecad 100644 --- a/python/sglang/srt/models/minimax_m2.py +++ b/python/sglang/srt/models/minimax_m2.py @@ -549,6 +549,7 @@ def __init__( ) self.topk = TopK( top_k=config.num_experts_per_tok, + layer_id=layer_id, renormalize=True, scoring_func=config.scoring_func, correction_bias=self.e_score_correction_bias, diff --git a/python/sglang/srt/models/mixtral.py b/python/sglang/srt/models/mixtral.py index 4b2b5c90807d..5eb781d76233 100644 --- a/python/sglang/srt/models/mixtral.py +++ b/python/sglang/srt/models/mixtral.py @@ -90,6 +90,7 @@ def __init__( self.topk = TopK( top_k=top_k, + layer_id=layer_id, renormalize=True, ) diff --git a/python/sglang/srt/models/olmoe.py b/python/sglang/srt/models/olmoe.py index c66156b33b46..ec0c2ece48d4 100644 --- a/python/sglang/srt/models/olmoe.py +++ b/python/sglang/srt/models/olmoe.py @@ -82,6 +82,7 @@ def __init__( self.topk = TopK( top_k=top_k, + layer_id=layer_id, renormalize=False, ) diff --git a/python/sglang/srt/models/phimoe.py b/python/sglang/srt/models/phimoe.py index f9f2ab06fef4..e7e5f4342685 100644 --- a/python/sglang/srt/models/phimoe.py +++ b/python/sglang/srt/models/phimoe.py @@ -200,6 +200,7 @@ def __init__( self.topk = TopK( top_k=top_k, + layer_id=layer_id, renormalize=False, custom_routing_function=phimoe_routing_function, ) diff --git a/python/sglang/srt/models/xverse_moe.py b/python/sglang/srt/models/xverse_moe.py index 898cf3abc42d..6dcb4e132078 100644 --- a/python/sglang/srt/models/xverse_moe.py +++ b/python/sglang/srt/models/xverse_moe.py @@ -94,6 +94,7 @@ class XverseMoE(nn.Module): def __init__( self, config: PretrainedConfig, + layer_id: int, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): @@ -134,6 +135,7 @@ def __init__( ) self.topk = TopK( top_k=self.top_k, + layer_id=layer_id, renormalize=getattr(self.config, "norm_topk_prob", False), ) @@ -305,6 +307,7 @@ def __init__( if config.num_experts is not None: self.mlp = XverseMoE( config=config, + layer_id=layer_id, quant_config=quant_config, prefix=add_prefix("mlp", prefix), ) diff --git a/test/registered/moe/test_triton_fused_moe.py b/test/registered/moe/test_triton_fused_moe.py index 3591507b091d..9dcb63eedb7b 100644 --- a/test/registered/moe/test_triton_fused_moe.py +++ b/test/registered/moe/test_triton_fused_moe.py @@ -106,6 +106,7 @@ def _test_case(self, m, n, k, e, topk, dtype): topk_op = TopK( top_k=topk, + layer_id=0, renormalize=False, use_grouped_topk=False, ) diff --git a/test/registered/unit/models/test_mimo_v2_fused_qkv_loader.py b/test/registered/unit/models/test_mimo_v2_fused_qkv_loader.py new file mode 100644 index 000000000000..07a3e83aa9d6 --- /dev/null +++ b/test/registered/unit/models/test_mimo_v2_fused_qkv_loader.py @@ -0,0 +1,177 @@ +"""Below the fused kv-head shard count, each MiMo-V2 rank holds its q, k and v heads in +that order, for BF16 weights and for block-FP8 codes regrouped with their scales.""" + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +import sglang.srt.models.mimo_v2 as mimo_v2 +from sglang.test.test_utils import CustomTestCase + +SHARDS = 4 +HIDDEN = 8 +FP8_MAX = torch.finfo(torch.float8_e4m3fn).max +# layer 0: global attention, one kv head per shard; layer 1: SWA, two +CONFIG = SimpleNamespace( + hybrid_layer_pattern=[0, 1], + num_attention_heads=8, + num_key_value_heads=4, + head_dim=3, + v_head_dim=2, + swa_num_attention_heads=8, + swa_num_key_value_heads=8, + swa_head_dim=3, + swa_v_head_dim=2, +) + + +def _qkv(layer): + prefix = "swa_" if CONFIG.hybrid_layer_pattern[layer] else "" + heads, kv_heads, head_dim, v_head_dim = ( + getattr(CONFIG, prefix + key) + for key in ( + "num_attention_heads", + "num_key_value_heads", + "head_dim", + "v_head_dim", + ) + ) + return ( + torch.randn(heads * head_dim, HIDDEN), + torch.randn(kv_heads * head_dim, HIDDEN), + torch.randn(kv_heads * v_head_dim, HIDDEN), + ) + + +def _interleave(q, k, v): + return torch.cat( + [ + torch.cat(parts) + for parts in zip( + q.chunk(SHARDS), k.chunk(SHARDS), v.chunk(SHARDS), strict=True + ) + ] + ) + + +def _rank_rows(q, k, v, tp_size, tp_rank): + return torch.cat([t.chunk(tp_size)[tp_rank] for t in (q, k, v)]) + + +def _parallel(tp_size, tp_rank): + parallel = SimpleNamespace(attn_tp_size=tp_size, attn_tp_rank=tp_rank) + return patch.object(mimo_v2, "get_parallel", return_value=parallel) + + +def _param(rows, dtype): + return torch.nn.Parameter( + torch.empty(rows, HIDDEN, dtype=dtype), requires_grad=False + ) + + +class TestMiMoV2FusedQKVLoader(CustomTestCase): + def test_bf16_rank_holds_its_q_k_v_heads_at_every_tp(self): + for layer in (0, 1): + q, k, v = _qkv(layer) + fused = _interleave(q, k, v).to(torch.bfloat16) + for tp_size in (1, 2, 4): + for tp_rank in range(tp_size): + param = _param(fused.shape[0] // tp_size, torch.bfloat16) + with _parallel(tp_size, tp_rank): + mimo_v2.load_mimo_v2_qkv_proj_weight( + f"model.layers.{layer}.self_attn.qkv_proj.weight", + param, + fused, + expected_fused_tp_size=SHARDS, + deferred_scale_inv={}, + config=CONFIG, + ) + want = _rank_rows(q, k, v, tp_size, tp_rank).to(torch.bfloat16) + self.assertTrue( + torch.equal(param.data, want), + f"layer {layer} tp {tp_size} rank {tp_rank}", + ) + + def test_block_fp8_regroups_with_its_deferred_scales(self): + # every shard fits one 128x128 block, so each carries its own scale row + for layer in (0, 1): + q, k, v = _qkv(layer) + shards = _interleave(q, k, v).chunk(SHARDS) + scales = torch.stack([s.abs().amax() / FP8_MAX for s in shards]).view(-1, 1) + codes = torch.cat( + [ + (s / scale).to(torch.float8_e4m3fn) + for s, scale in zip(shards, scales) + ] + ) + decoded = torch.cat( + [c.float() * scale for c, scale in zip(codes.chunk(SHARDS), scales)] + ) + sizes = [t.shape[0] // SHARDS for t in (q, k, v)] + parts = [torch.split(shard, sizes) for shard in decoded.chunk(SHARDS)] + dq, dk, dv = (torch.cat(group) for group in zip(*parts)) + for tp_size in (1, 2): + for tp_rank in range(tp_size): + name = f"model.layers.{layer}.self_attn.qkv_proj.weight" + params = { + name: _param(codes.shape[0] // tp_size, torch.float8_e4m3fn), + name + "_scale_inv": torch.nn.Parameter( + torch.empty(1, 1), requires_grad=False + ), + } + deferred = {} + with _parallel(tp_size, tp_rank): + for tensor_name, tensor in ( + (name, codes), + (name + "_scale_inv", scales), + ): + mimo_v2.load_mimo_v2_qkv_proj_weight( + tensor_name, + params[tensor_name], + tensor, + expected_fused_tp_size=SHARDS, + deferred_scale_inv=deferred, + config=CONFIG, + ) + mimo_v2._resolve_deferred_qkv_scale_inv( + params, + deferred, + expected_fused_tp_size=SHARDS, + config=CONFIG, + ) + got = params[name].data.float() * params[name + "_scale_inv"].data + want = _rank_rows(dq, dk, dv, tp_size, tp_rank) + # the resolver re-quantizes the rank's rows as one block + tolerance = 0.07 * want.abs().max() + self.assertLessEqual( + (got - want).abs().max().item(), + tolerance.item(), + f"layer {layer} tp {tp_size} rank {tp_rank}", + ) + + def test_bf16_without_layer_sizes_raises(self): + q, k, v = _qkv(1) + fused = _interleave(q, k, v).to(torch.bfloat16) + for name, config in ( + ("model.layers.1.self_attn.qkv_proj.weight", None), + ("model.mtp_block.self_attn.qkv_proj.weight", CONFIG), + ): + with _parallel(2, 0), self.assertRaises(ValueError): + mimo_v2.load_mimo_v2_qkv_proj_weight( + name, + _param(fused.shape[0] // 2, torch.bfloat16), + fused, + expected_fused_tp_size=SHARDS, + deferred_scale_inv={}, + config=config, + ) + + +if __name__ == "__main__": + unittest.main() From 6c94e57f9ff5b84ba9746f5983ff13ce7e84c68b Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Sun, 30 Aug 2026 07:14:35 +0000 Subject: [PATCH 34/56] [32/32] [sglang-miles] fix: initialize Inkling safely on ROCm (#37294) --- python/sglang/srt/mem_cache/swa_memory_pool.py | 8 +------- python/sglang/srt/models/inkling_common/attn.py | 8 +++++--- 2 files changed, 6 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/mem_cache/swa_memory_pool.py b/python/sglang/srt/mem_cache/swa_memory_pool.py index bf037bdad4c4..ca732eab549a 100644 --- a/python/sglang/srt/mem_cache/swa_memory_pool.py +++ b/python/sglang/srt/mem_cache/swa_memory_pool.py @@ -211,13 +211,7 @@ def get_value_buffer(self, layer_id: int): return self.full_kv_pool.get_value_buffer(layer_id_pool) def get_v_head_dim(self): - # The FULL side's dim, as HybridLinearKVPool.get_v_head_dim(): a caller - # asking a pool for "the" v_head_dim wants the full-attention geometry. - # `start_layer`, not 0, so pipeline parallelism (start_layer > 0) works, - # and because layer 0 need not be a full-attention layer. - return self.full_kv_pool.get_value_buffer(self.full_kv_pool.start_layer).shape[ - -1 - ] + return self.full_kv_pool.get_v_head_dim() def get_kv_buffer(self, layer_id: int): self._wait_for_layer(layer_id) diff --git a/python/sglang/srt/models/inkling_common/attn.py b/python/sglang/srt/models/inkling_common/attn.py index 5659b3efc35f..3db9755eabf5 100644 --- a/python/sglang/srt/models/inkling_common/attn.py +++ b/python/sglang/srt/models/inkling_common/attn.py @@ -6,9 +6,6 @@ import torch from torch import nn -from sglang.kernels.ops.attention.flash_attn.cute.batch_invariance import ( - is_batch_invariant, -) from sglang.kernels.ops.attention.log_scaling_tau import ( apply_log_scaling_tau as _apply_log_scaling_tau, ) @@ -45,10 +42,14 @@ import cutlass.cute as cute from cutlass.cute import Float32 + from sglang.kernels.ops.attention.flash_attn.cute.batch_invariance import ( + is_batch_invariant, + ) from sglang.kernels.ops.attention.flash_attn.cute.seqlen_info import SeqlenInfoQK except Exception as _import_error: cute = None Float32 = None + is_batch_invariant = None SeqlenInfoQK = None _cute_import_error = _import_error else: @@ -953,6 +954,7 @@ def forward( if ( envs.SGLANG_OPT_USE_INKLING_SHEARED_BIAS.get() and fa4 + and is_batch_invariant is not None and not is_batch_invariant() ): # FA4 sheared-bias kernel: pass rel_logits directly; the kernel shears From 2b7313dc21a57251f568553247d36d78e585a614 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Fri, 2 Oct 2026 16:25:36 -0700 Subject: [PATCH 35/56] [sglang-miles] Allocate packed weight receive buffers directly from metadata (#39265) Co-authored-by: yueming-yuan --- .../model_runner_components/weight_updater.py | 10 +-- .../sglang/srt/weight_sync/tensor_bucket.py | 28 ++++++++ test/registered/rl/test_tensor_bucket.py | 67 +++++++++++++++++++ 3 files changed, 96 insertions(+), 9 deletions(-) create mode 100644 test/registered/rl/test_tensor_bucket.py 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 5d852a613675..4cb5af6354aa 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 @@ -288,15 +288,7 @@ def receive_weights_from_distributed( 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) + bucket = FlattenedTensorBucket.empty(names, dtypes, shapes, self.device) flattened_tensor = bucket.get_flattened_tensor() torch.distributed.broadcast( flattened_tensor, diff --git a/python/sglang/srt/weight_sync/tensor_bucket.py b/python/sglang/srt/weight_sync/tensor_bucket.py index c1d592ddbb0d..a4ac966cd9b3 100644 --- a/python/sglang/srt/weight_sync/tensor_bucket.py +++ b/python/sglang/srt/weight_sync/tensor_bucket.py @@ -83,6 +83,34 @@ def get_flattened_tensor(self) -> torch.Tensor: """Get the flattened tensor containing all bucket tensors""" return self.flattened_tensor + @classmethod + def empty(cls, names, dtypes, shapes, device): + """Allocate a receive bucket without constructing or copying its tensors.""" + metadata = [] + offset = 0 + for name, dtype, shape in zip(names, dtypes, shapes, strict=True): + dtype = dtype if isinstance(dtype, torch.dtype) else getattr(torch, dtype) + shape = torch.Size(shape) + if any(dim < 0 for dim in shape): + raise ValueError( + f"Negative dimension in received tensor {name}: {shape}" + ) + if offset % dtype.itemsize: + raise ValueError( + f"Unaligned received tensor {name}: byte offset {offset}, dtype {dtype}" + ) + size = shape.numel() * dtype.itemsize + metadata.append( + FlattenedTensorMetadata(name, shape, dtype, offset, offset + size, size) + ) + offset += size + if not metadata: + raise ValueError("Cannot create empty tensor bucket") + return cls( + flattened_tensor=torch.empty(offset, dtype=torch.uint8, device=device), + metadata=metadata, + ) + def get_metadata(self) -> List[FlattenedTensorMetadata]: """Get metadata for all tensors in the bucket""" return self.metadata diff --git a/test/registered/rl/test_tensor_bucket.py b/test/registered/rl/test_tensor_bucket.py new file mode 100644 index 000000000000..63bdbc99f5ab --- /dev/null +++ b/test/registered/rl/test_tensor_bucket.py @@ -0,0 +1,67 @@ +import sys + +import pytest +import torch + +from sglang.srt.weight_sync.tensor_bucket import FlattenedTensorBucket +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=3, suite="base-a-test-cpu") + + +@pytest.mark.parametrize("dtype_strings", [False, True]) +def test_receive_bucket_preserves_mixed_dtype_wire_bytes(dtype_strings): + scale_bits = torch.tensor( + [0, -2147483648, 2143289345, 2139095040], dtype=torch.int32 + ) + values = [ + ("packed_weight", torch.arange(16, dtype=torch.uint8).reshape(2, 8)), + ("block_scales", torch.arange(8, dtype=torch.uint8).view(torch.float8_e4m3fn)), + ("global_scale", scale_bits.view(torch.float32)[2].reshape(())), + ("other_scales", scale_bits.view(torch.float32)), + ("bf16", torch.tensor([1.5, -2.0], dtype=torch.bfloat16)), + ("empty", torch.empty((0, 4), dtype=torch.float32)), + ] + sent = FlattenedTensorBucket(named_tensors=values) + dtypes = [value.dtype for _, value in values] + if dtype_strings: + dtypes = [str(dtype).removeprefix("torch.") for dtype in dtypes] + received = FlattenedTensorBucket.empty( + [name for name, _ in values], + dtypes, + [value.shape for _, value in values], + "cpu", + ) + assert received.metadata == sent.metadata + received.flattened_tensor.copy_(sent.flattened_tensor) + for (name, original), (received_name, actual), meta in zip( + values, received.reconstruct_tensors(), sent.metadata, strict=True + ): + assert name == received_name + assert original.shape == actual.shape and original.dtype == actual.dtype + assert torch.equal( + original.reshape(-1).view(torch.uint8), actual.reshape(-1).view(torch.uint8) + ) + if actual.numel(): + assert ( + actual.data_ptr() + == received.flattened_tensor.data_ptr() + meta.start_idx + ) + + +@pytest.mark.parametrize( + "names,dtypes,shapes", + [ + ([], [], []), + (["a"], [], [[4]]), + (["a"], [torch.float32], [[-1, -1]]), + (["a", "b"], [torch.uint8, torch.float32], [[3], []]), + ], +) +def test_receive_bucket_rejects_invalid_metadata(names, dtypes, shapes): + with pytest.raises(ValueError): + FlattenedTensorBucket.empty(names, dtypes, shapes, "cpu") + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__])) From 258bc1a8d928925dc82654470240417e9ab09d77 Mon Sep 17 00:00:00 2001 From: Xinyu Jiang Date: Sat, 3 Oct 2026 13:27:01 -0400 Subject: [PATCH 36/56] [sglang-miles] Fix Kimi-K3 weight reloads after graph capture (#41925) Co-authored-by: Zhiyao Jiang --- python/sglang/srt/models/kimi_k3.py | 41 +++++++++++++++++++++++------ 1 file changed, 33 insertions(+), 8 deletions(-) diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py index e732cd1fbcdb..6569eb053fe7 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/python/sglang/srt/models/kimi_k3.py @@ -222,10 +222,11 @@ def _k3_bf16_gemm( def _merge_weights_as_views( - mods: list, pad_rows_to: int = 1 + mods: list, pad_rows_to: int = 1, merged: Optional[torch.Tensor] = None ) -> tuple[torch.Tensor, list[int]]: """Cat module weights along dim 0; re-point each module's weight to a view of the merged buffer so the original storage is freed (net extra memory ~0). + Reuse an existing merged buffer during reload to preserve graph addresses. With pad_rows_to > 1 the merged buffer gets zero rows appended up to the next multiple, so every row of the fused GEMM output stays 16-byte aligned @@ -235,7 +236,18 @@ def _merge_weights_as_views( pad = (-sum(sizes)) % pad_rows_to if pad: ws = ws + [ws[0].new_zeros((pad, ws[0].shape[1]))] - merged = torch.cat(ws, dim=0).contiguous() + if merged is None: + merged = torch.cat(ws, dim=0).contiguous() + else: + # Keep the buffer address stable for graph replay after weight reloads + assert merged.shape == (sum(w.shape[0] for w in ws), ws[0].shape[1]) + assert merged.dtype == ws[0].dtype and merged.device == ws[0].device + off = 0 + for w in ws: + destination = merged[off : off + w.shape[0]] + if destination.data_ptr() != w.data_ptr(): + destination.copy_(w) + off += w.shape[0] off = 0 for m, n in zip(mods, sizes): m.weight.data = merged[off : off + n] @@ -715,7 +727,9 @@ def _merge_front_weights(self) -> None: dtypes = {m.weight.dtype for m in mods} if len(dtypes) != 1 or dtypes.pop() not in (torch.bfloat16, torch.float16): return - self._front_w, self._front_sizes = _merge_weights_as_views(mods) + self._front_w, self._front_sizes = _merge_weights_as_views( + mods, merged=self._front_w + ) self._front_is_ep_pair = len(mods) == 2 # Invalidate the cached properties. for prop in ( @@ -1842,7 +1856,7 @@ def _merge_bfa_weights(self) -> None: if _is_hip and self._merge_kda_inproj_weights_hip(): # Split-path f_b GEMM still uses this when the fused in-proj # is above the token threshold. - self._bfa_f_b_w = self.f_b_proj.weight + self._bfa_f_b_w = self.f_b_proj.weight.data return mods = [self.f_a_proj, self.b_proj] if self._bfa_uses_block_fp8: @@ -1851,12 +1865,20 @@ def _merge_bfa_weights(self) -> None: pad = (-sum(sizes)) % 8 if pad: weights.append(weights[0].new_zeros((pad, weights[0].shape[1]))) - self._bfa_w = torch.cat(weights, dim=0).contiguous() - self._bfa_f_b_w = _get_k3_dense_weight(self.f_b_proj).contiguous() + bfa_w = torch.cat(weights, dim=0).contiguous() + f_b_w = _get_k3_dense_weight(self.f_b_proj).contiguous() + if self._bfa_w is None: + self._bfa_w, self._bfa_f_b_w = bfa_w, f_b_w + else: + # Keep the buffer addresses stable for graph replay after weight reloads + self._bfa_w.copy_(bfa_w) + self._bfa_f_b_w.copy_(f_b_w) else: 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_w, sizes = _merge_weights_as_views( + mods, pad_rows_to=8, merged=self._bfa_w + ) # .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 @@ -1875,8 +1897,11 @@ def _merge_kda_inproj_weights_hip(self) -> bool: # [q,k,v,g | f_a | b | pad]; f_a/b keep the same relative order and the # same pad (both widths are 4 short of a multiple of 8), so the tail # view is byte-identical to the wide-only merge. + previous = getattr(self, "_qkvgbfa_layer", None) merged, sizes = _merge_weights_as_views( - [self.fused_qkvg_proj, self.f_a_proj, self.b_proj], pad_rows_to=8 + [self.fused_qkvg_proj, self.f_a_proj, self.b_proj], + pad_rows_to=8, + merged=previous.weight if previous is not None else None, ) self._bfa_fa_size, self._bfa_b_size = sizes[-2:] self._bfa_w = merged[sizes[0] :] From 48214466d9297be79771bb02e5353a26634b35ff Mon Sep 17 00:00:00 2001 From: Kangrui Du <89372739+Rockdu@users.noreply.github.com> Date: Sun, 4 Oct 2026 13:05:34 -0700 Subject: [PATCH 37/56] [cherrypick from #42332] [diffusion] fix: report the port each scheduler actually bound to its clients (#42470) Signed-off-by: rockdu --- python/sglang/multimodal_gen/runtime/launch_server.py | 6 +++++- python/sglang/multimodal_gen/runtime/managers/gpu_worker.py | 2 ++ python/sglang/multimodal_gen/runtime/managers/scheduler.py | 5 +++++ .../multimodal_gen/runtime/server_args/server_args.py | 3 +++ 4 files changed, 15 insertions(+), 1 deletion(-) diff --git a/python/sglang/multimodal_gen/runtime/launch_server.py b/python/sglang/multimodal_gen/runtime/launch_server.py index d3fa0a1deac2..5567803d61ec 100644 --- a/python/sglang/multimodal_gen/runtime/launch_server.py +++ b/python/sglang/multimodal_gen/runtime/launch_server.py @@ -209,6 +209,10 @@ def launch_server(server_args: ServerArgs, launch_http_server: bool = True): ) scheduler_infos.append(data) reader.close() + if data["actual_scheduler_port"] is not None: + server_args.scheduler_ports[data["dp_replica"]] = data[ + "actual_scheduler_port" + ] logger.debug("All workers are ready") logger.info("[server-load] workers_ready_monotonic_ns=%d", time.monotonic_ns()) @@ -241,7 +245,7 @@ def launch_server(server_args: ServerArgs, launch_http_server: bool = True): logger.info("Launch FastAPI server in another process because of webui.") http_server_process = worker_context.Process( target=bootstrap_http_server_process, - args=(server_args_payload,), + args=(ServerArgsPayload.capture(server_args),), name="sglang-diffusion-webui", daemon=True, ) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 0ce719e3abf3..43a7d5e0565f 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -1614,6 +1614,8 @@ def run_scheduler_process( pipe_writer.send( { "status": "ready", + "dp_replica": scheduler.dp_replica, + "actual_scheduler_port": scheduler.actual_scheduler_port, } ) scheduler.event_loop() diff --git a/python/sglang/multimodal_gen/runtime/managers/scheduler.py b/python/sglang/multimodal_gen/runtime/managers/scheduler.py index f2eb77b07cb4..9539532d80dc 100644 --- a/python/sglang/multimodal_gen/runtime/managers/scheduler.py +++ b/python/sglang/multimodal_gen/runtime/managers/scheduler.py @@ -121,12 +121,17 @@ def __init__( gpus_per_replica = max(1, server_args.num_gpus // server_args.dp_size) self.dp_replica = gpu_id // gpus_per_replica self.context = zmq.Context(io_threads=2) + # Port actually bound by this replica's driver rank; reported to the launcher. + self.actual_scheduler_port = None if gpu_id % gpus_per_replica == 0: endpoint = server_args.scheduler_endpoint_for(self.dp_replica) # router allocates identify (envelope) for each connection self.receiver, actual_endpoint = get_zmq_socket( self.context, zmq.ROUTER, endpoint, True ) + self.actual_scheduler_port = int(actual_endpoint.rsplit(":", 1)[1]) + # Readers in this process, e.g. the metrics label, see the bound port. + server_args.scheduler_ports[self.dp_replica] = self.actual_scheduler_port logger.info( f"Scheduler (dp replica {self.dp_replica}) bind at endpoint: " f"{actual_endpoint}" diff --git a/python/sglang/multimodal_gen/runtime/server_args/server_args.py b/python/sglang/multimodal_gen/runtime/server_args/server_args.py index 75b683f989df..0e927e8b566e 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args/server_args.py @@ -1434,6 +1434,9 @@ def _adjust_network_ports(self): ) seen_ports[port] = name self._require_port(port, name) + self.scheduler_ports = [ + self.scheduler_port + replica for replica in range(self.dp_size or 1) + ] else: settled_ports: set[int] = set() if needs_http: From f19dcfb124ec33442029eb5fe8c4c6cf5950af30 Mon Sep 17 00:00:00 2001 From: Kangrui Du <89372739+Rockdu@users.noreply.github.com> Date: Sun, 4 Oct 2026 13:06:38 -0700 Subject: [PATCH 38/56] [sglang-miles] [diffusion] fix: give the scheduler its own host (#42436) Signed-off-by: rockdu --- .../entrypoints/diffusion_generator.py | 3 +- .../multimodal_gen/runtime/launch_server.py | 4 +-- .../runtime/managers/scheduler.py | 18 +++++----- .../runtime/server_args/server_args.py | 17 ++++----- .../test/unit/test_ipc_array.py | 35 +++++++++++++++++++ 5 files changed, 58 insertions(+), 19 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py index 9fce9bb852f9..bba0c5768f84 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py @@ -214,7 +214,8 @@ def _check_remote_scheduler(self): raise ConnectionError( f"Could not connect to remote scheduler at " f"{self.server_args.scheduler_endpoint} with `local mode` as False. " - "Please ensure the server is running." + "Please ensure the server is running, binds a reachable " + "--scheduler-host, and uses the same --scheduler-port as this client." ) logger.info( f"Successfully connected to remote scheduler at " diff --git a/python/sglang/multimodal_gen/runtime/launch_server.py b/python/sglang/multimodal_gen/runtime/launch_server.py index 5567803d61ec..14ccf191d647 100644 --- a/python/sglang/multimodal_gen/runtime/launch_server.py +++ b/python/sglang/multimodal_gen/runtime/launch_server.py @@ -446,7 +446,7 @@ def find_port(start): logger.info("All pool role instances ready") # Start DiffusionServer - frontend_endpoint = f"tcp://{host}:{server_args.scheduler_port}" + frontend_endpoint = server_args.scheduler_endpoint diffusion_server = DiffusionServer( frontend_endpoint=frontend_endpoint, @@ -552,7 +552,7 @@ def launch_disagg_server(server_args: ServerArgs): denoiser_result_ep = f"tcp://{host}:{base_port + 2}" decoder_result_ep = f"tcp://{host}:{base_port + 3}" - frontend_endpoint = f"tcp://{host}:{base_port}" + frontend_endpoint = server_args.scheduler_endpoint logger.info( "Starting DiffusionServer: %d encoder(s), %d denoiser(s), %d decoder(s)", diff --git a/python/sglang/multimodal_gen/runtime/managers/scheduler.py b/python/sglang/multimodal_gen/runtime/managers/scheduler.py index 9539532d80dc..b524a4f4f55b 100644 --- a/python/sglang/multimodal_gen/runtime/managers/scheduler.py +++ b/python/sglang/multimodal_gen/runtime/managers/scheduler.py @@ -773,7 +773,9 @@ def return_result( with maybe_record_function("REPLY spill+pickle+send"): # if the server is local, use temp file to spill the frame array # instead of leaving it in OutputBatch to be pickled later - if is_local_endpoint(self.server_args.scheduler_endpoint): + if is_local_endpoint( + self.server_args.scheduler_endpoint_for(self.dp_replica) + ): with self._record_return_stage( output_batch, "Scheduler.return_result.spill_arrays" ): @@ -781,13 +783,13 @@ def return_result( output_batch.output ) - with self._record_return_stage( - output_batch, "Scheduler.return_result.spill_cuda" - ): - # The previous reply has already been mapped: the client - # only sends the next hop after materializing the last one. - release_retained_producer_tensors() - spill_cuda_tensors(output_batch, in_place=True) + with self._record_return_stage( + output_batch, "Scheduler.return_result.spill_cuda" + ): + # The previous reply has already been mapped: the client + # only sends the next hop after materializing the last one. + release_retained_producer_tensors() + spill_cuda_tensors(output_batch, in_place=True) with self._record_return_stage( output_batch, "Scheduler.return_result.pickle" ): diff --git a/python/sglang/multimodal_gen/runtime/server_args/server_args.py b/python/sglang/multimodal_gen/runtime/server_args/server_args.py index 0e927e8b566e..20743c7e4fc1 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args/server_args.py @@ -536,6 +536,7 @@ class ServerArgs(DisaggServerArgsMixin): webui: bool = False webui_port: int | None = 12312 + scheduler_host: str = "127.0.0.1" scheduler_port: int = 5555 # settled ingress ports, one per DP replica; None until ports are settled scheduler_ports: list[int] | None = None @@ -2821,6 +2822,12 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: default=ServerArgs.master_port, help="Master port for distributed inference. If not set, a random free port will be used.", ) + parser.add_argument( + "--scheduler-host", + type=str, + default=ServerArgs.scheduler_host, + help="Host the scheduler binds for client requests.", + ) parser.add_argument( "--scheduler-port", type=int, @@ -3107,22 +3114,16 @@ def url(self): @property def scheduler_endpoint(self): - """ - Internal endpoint for scheduler. - Prefers the configured host but normalizes localhost -> 127.0.0.1 to avoid ZMQ issues. - """ + """Internal endpoint for scheduler.""" return self.scheduler_endpoint_for(0) def scheduler_endpoint_for(self, replica: int) -> str: """Ingress endpoint of one DP replica's driver rank.""" - scheduler_host = self.host - if scheduler_host is None or scheduler_host == "localhost": - scheduler_host = "127.0.0.1" if self.scheduler_ports is not None: port = self.scheduler_ports[replica] else: port = self.scheduler_port + replica - return f"tcp://{scheduler_host}:{port}" + return f"tcp://{self.scheduler_host}:{port}" @property def scheduler_endpoints(self) -> list[str]: diff --git a/python/sglang/multimodal_gen/test/unit/test_ipc_array.py b/python/sglang/multimodal_gen/test/unit/test_ipc_array.py index cff1731e1a89..62e6f7861a0c 100644 --- a/python/sglang/multimodal_gen/test/unit/test_ipc_array.py +++ b/python/sglang/multimodal_gen/test/unit/test_ipc_array.py @@ -1,7 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 +import pickle import tempfile from pathlib import Path +from unittest.mock import Mock import numpy as np import pytest @@ -83,6 +85,39 @@ def test_local_endpoint_detection(): assert not is_local_endpoint("tcp://10.0.0.2:30000") +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA IPC needs a GPU") +@pytest.mark.parametrize( + "scheduler_host, wire_type", [("127.0.0.1", "CudaIpcRef"), ("10.0.0.2", "Tensor")] +) +def test_reply_reaches_the_caller_as_a_tensor(scheduler_host, wire_type): + from sglang.multimodal_gen.runtime.managers.scheduler import Scheduler + from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch + from sglang.multimodal_gen.runtime.scheduler_client import ( + _materialize_local_cuda_refs, + _materialize_output_batch_file_refs, + ) + from sglang.multimodal_gen.runtime.server_args.server_args import ServerArgs + + server_args = object.__new__(ServerArgs) + server_args.host = "172.18.0.2" + server_args.scheduler_host = scheduler_host + server_args.scheduler_port = 5555 + server_args.scheduler_ports = None + scheduler = Scheduler.__new__(Scheduler) + scheduler.server_args = server_args + scheduler.dp_replica = 0 + scheduler.receiver = Mock() + output = torch.arange(24, device="cuda", dtype=torch.float32) + + scheduler.return_result(OutputBatch(output=output.clone()), identity=b"client") + received = pickle.loads(scheduler.receiver.send_multipart.call_args.args[0][2]) + assert type(received.output).__name__ == wire_type + + _materialize_output_batch_file_refs(received) + _materialize_local_cuda_refs(server_args.scheduler_endpoint, received) + assert torch.equal(received.output, output) + + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.uint8]) def test_spill_large_tensors_round_trips(monkeypatch, tmp_path, dtype): monkeypatch.setattr(ipc_array, "_array_ipc_dir", lambda: str(tmp_path)) From b84bec5761a6cdc8b1df86b9b89ead80a7a9bb02 Mon Sep 17 00:00:00 2001 From: Zhichen Zeng Date: Wed, 7 Oct 2026 15:10:31 -0700 Subject: [PATCH 39/56] [sglang-miles] Fix the bf16 triton MoE path and mhc_pre splitk for DeepSeek-V4.1 (#42408) --- python/sglang/kernels/ops/layernorm/mhc.py | 15 +++++++++------ .../moe/moe_runner/triton_utils/fused_moe.py | 4 ++-- 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/python/sglang/kernels/ops/layernorm/mhc.py b/python/sglang/kernels/ops/layernorm/mhc.py index 78c163f0c373..0c237e8a1d75 100644 --- a/python/sglang/kernels/ops/layernorm/mhc.py +++ b/python/sglang/kernels/ops/layernorm/mhc.py @@ -1178,14 +1178,17 @@ def mhc_pre( else: if num_tokens <= 2048: assert n_splits == 1 - if hc_hidden_size == 16384: - hidden_block = 256 - elif hc_hidden_size == 28672: - hidden_block = 128 + split_size = hc_hidden_size // n_splits_pre + for hidden_block in (256, 128): + if ( + hc_hidden_size % hidden_block == 0 + and split_size % hidden_block == 0 + ): + break else: raise NotImplementedError( - f"mhc_pre splitk kernel only supports hc_hidden_size in {{16384, 28672}}, " - f"got {hc_hidden_size}" + f"mhc_pre splitk kernel needs a hidden_block in {{256, 128}} that divides " + f"hc_hidden_size // split_k = {split_size}, got hc_hidden_size={hc_hidden_size}" ) kernel_0, _ = mhc_pre_gemm_sqrsum_splitk_kernel( hc_mult3, 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 1c7163643d18..e9595bdf2ca6 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 @@ -295,7 +295,7 @@ def fused_experts( a1_q: Optional[torch.Tensor] = None, fuse_swiglu_interleaved: bool = False, ): - topk_weights, topk_ids, _ = topk_output + topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids filter_expert = ( moe_runner_config.num_experts is None or moe_runner_config.num_experts != moe_runner_config.num_local_experts @@ -1160,7 +1160,7 @@ def fused_moe( - torch.Tensor: The output tensor after applying the MoE layer. """ if _is_xpu and not get_moe_runner_backend().is_triton(): - topk_weight, topk_ids, _ = topk_output + topk_weight, topk_ids = topk_output.topk_weights, topk_output.topk_ids from sgl_kernel import fused_experts as sgl_fused_experts return sgl_fused_experts( From d1faf8244e6d60e8491d7e81d2e465e3f9a2fd94 Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Sat, 10 Oct 2026 08:30:26 +0800 Subject: [PATCH 40/56] fix(dsv4): support R3 capture and MXFP4 online updates (#39392) Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- .../srt/layers/moe/flashinfer_megamoe.py | 19 +- python/sglang/srt/layers/moe/hash_topk.py | 28 +- python/sglang/srt/layers/moe/topk.py | 4 +- .../srt/layers/quantization/base_config.py | 2 +- python/sglang/srt/layers/quantization/fp8.py | 38 ++- .../srt/layers/quantization/fp8_utils.py | 3 +- .../mxfp4_flashinfer_trtllm_moe.py | 256 +++++++++------ .../model_runner_components/weight_updater.py | 42 ++- python/sglang/srt/model_loader/loader.py | 13 +- python/sglang/srt/server_args.py | 7 +- .../srt/state_capturer/routed_experts.py | 23 +- ...test_mxfp4_flashinfer_trtllm_hot_reload.py | 297 ++++++++++++++++++ test/registered/moe/test_hash_topk.py | 140 ++++++++- .../layers/moe/test_flashinfer_megamoe.py | 109 +++++++ .../test_deepgemm_ue8m0_requant.py | 41 ++- .../test_fp8_moe_runner_fallback.py | 46 +++ .../test_fp8_moe_runner_ownership.py | 56 +++- .../model_loader/test_default_model_loader.py | 51 +++ .../unit/server_args/test_server_args.py | 27 ++ .../test_routed_experts_scattered_a2a.py | 50 +++ 20 files changed, 1097 insertions(+), 155 deletions(-) create mode 100644 test/registered/kernels/ops/quantization/test_mxfp4_flashinfer_trtllm_hot_reload.py diff --git a/python/sglang/srt/layers/moe/flashinfer_megamoe.py b/python/sglang/srt/layers/moe/flashinfer_megamoe.py index 1e336f576896..7cc61209d480 100644 --- a/python/sglang/srt/layers/moe/flashinfer_megamoe.py +++ b/python/sglang/srt/layers/moe/flashinfer_megamoe.py @@ -140,7 +140,10 @@ def _forward_megamoe_legacy(mega: Any, tensors: Any) -> torch.Tensor: def _select_megamoe_forward(mega: Any) -> Callable[[Any, Any], torch.Tensor]: import inspect - if "return_workspace_view" in inspect.signature(mega.forward).parameters: + if ( + getattr(mega, "supports_output_view", False) + and "return_workspace_view" in inspect.signature(mega.forward).parameters + ): return _forward_megamoe_with_workspace_view return _forward_megamoe_legacy @@ -324,8 +327,8 @@ def ensure_fp4_moe_layer_for_flashinfer_megamoe(layer: FusedMoE) -> Any: top_k=layer.top_k, activation_clamp=layer.moe_runner_config.swiglu_limit, ), - w13_scale=layer.w13_weight_scale_inv, - w2_scale=layer.w2_weight_scale_inv, + w13_scale=layer.w13_weight_scale_inv_shuffled, + w2_scale=layer.w2_weight_scale_inv_shuffled, ) @@ -402,7 +405,9 @@ def prepare_fp4_moe_weights_for_flashinfer_megamoe( current moe_ep API owns backend-specific weight preprocessing, including DeepGEMM scale layout transforms. """ - _init_flashinfer_megamoe_layer_state(layer) + # Captured graphs retain the runtime's workspace and transformed tensors. + # Same-geometry reloads refresh their storage without rebuilding the runtime. + _get_or_init_flashinfer_megamoe_layer_state(layer) from flashinfer.moe_ep import ( MoEWeightPack, @@ -420,11 +425,13 @@ def prepare_fp4_moe_weights_for_flashinfer_megamoe( intermediate_size=layer.intermediate_size_per_partition, hidden_size=layer.hidden_size, ) + # Kernel scales pack four UE8M0 columns per INT32; loaders still need the + # original FP32 scale shape for the next checkpoint or tensor update. _bind_transformed_weights( layer, transformed_weights, - w13_scale_name="w13_weight_scale_inv", - w2_scale_name="w2_weight_scale_inv", + w13_scale_name="w13_weight_scale_inv_shuffled", + w2_scale_name="w2_weight_scale_inv_shuffled", ) diff --git a/python/sglang/srt/layers/moe/hash_topk.py b/python/sglang/srt/layers/moe/hash_topk.py index d32c600ff200..5bc7869af84d 100644 --- a/python/sglang/srt/layers/moe/hash_topk.py +++ b/python/sglang/srt/layers/moe/hash_topk.py @@ -20,6 +20,7 @@ TopKConfig, _mask_topk_ids_padded_region, _zero_topk_weights_padded_region, + capture_routed_experts_if_allowed, remap_topk_for_per_rank_shared_slots, ) from sglang.srt.layers.moe.utils import has_per_rank_fused_shared_slots @@ -54,12 +55,26 @@ def __init__( self.waterfill_balancer = None if self.enable_waterfill: + # Routed-experts capture sizes its buffer by the model's fused + # shared-expert count, but under waterfill the ids captured here + # carry routed slots only, so the two widths disagree. + if get_exec().features.enable_return_routed_experts: + raise ValueError( + "HashTopK cannot capture routed experts with waterfill enabled: " + "disable --enable-waterfill or --enable-return-routed-experts." + ) # Waterfill appends the shared expert after EPLB maps routed IDs. topk -= num_fused_shared_experts num_fused_shared_experts = 0 self.num_experts = num_experts self.topk = topk + self.topk_config = TopKConfig( + top_k=topk, + num_fused_shared_experts=num_fused_shared_experts, + routed_scaling_factor=routed_scaling_factor, + scoring_func=scoring_func, + ) self.routed_scaling_factor = routed_scaling_factor self.num_fused_shared_experts = num_fused_shared_experts self.score_func = scoring_func @@ -250,6 +265,13 @@ def forward( if self.apply_routed_scaling_factor_on_output: topk_weights = topk_weights * self.routed_scaling_factor + capture_routed_experts_if_allowed( + self.topk_config, + self.layer_id, + topk_ids, + num_token_non_padded, + ) + num_fused_shared_experts = self.num_fused_shared_experts log2phy_prob = None if ( @@ -284,11 +306,7 @@ def forward( topk_weights, num_fused_shared_experts, num_physical_routed_experts, - TopKConfig( - top_k=self.topk, - num_fused_shared_experts=num_fused_shared_experts, - routed_scaling_factor=self.routed_scaling_factor, - ), + self.topk_config, ) else: topk_ids = topk_ids_logical_to_physical( diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index d806c3d8f635..1f43a194b3b1 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -1666,7 +1666,7 @@ def _mask_topk_ids_padded_region( mask_topk_ids(topk_ids, num_token_non_padded) elif _is_npu: return - elif _can_fuse_padded_region(topk_ids): + elif not topk_ids.is_cpu and _can_fuse_padded_region(topk_ids): _fill_padded_rows(topk_ids, num_token_non_padded, fill_value) else: indices = torch.arange(0, topk_ids.shape[0], device=topk_ids.device) @@ -1679,7 +1679,7 @@ def _zero_topk_weights_padded_region( ): if num_token_non_padded is None: return - if _can_fuse_padded_region(topk_weights): + if not topk_weights.is_cpu and _can_fuse_padded_region(topk_weights): _fill_padded_rows(topk_weights, num_token_non_padded, 0.0) return indices = torch.arange(0, topk_weights.shape[0], device=topk_weights.device) diff --git a/python/sglang/srt/layers/quantization/base_config.py b/python/sglang/srt/layers/quantization/base_config.py index 6ab29d05d6d1..8ded3b128523 100644 --- a/python/sglang/srt/layers/quantization/base_config.py +++ b/python/sglang/srt/layers/quantization/base_config.py @@ -49,7 +49,7 @@ def process_weights_after_loading(self, layer: nn.Module) -> None: def restore_weights_before_loading(self, layer: nn.Module) -> None: """Undo an in-place repack so checkpoint-format weights can be loaded again. - Needed only when `process_weights_after_loading` changes parameter shapes. + Needed when `process_weights_after_loading` changes parameter shapes or layout. """ return diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 5634a2617a6a..3bbc31665f1e 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -445,8 +445,12 @@ def get_quant_method( return Mxfp4HummingMoEMethod(fp8_method, prefix=prefix) - if self.is_fp4_experts and get_moe_runner_backend().is_flashinfer_mxfp4(): - # SM100 uses TRT-LLM; SM90 uses W4A16 and SM120 uses MXFP8xMXFP4. + if self.is_fp4_experts and ( + get_moe_runner_backend().is_flashinfer_mxfp4() + or get_moe_runner_backend().is_flashinfer_trtllm_routed() + ): + # flashinfer_mxfp4 also supports CUTLASS on SM90 and SM120. + # The routed TRT-LLM backend is validated as SM100 at model load. if get_platform().is_sm90 or get_platform().is_sm120: from sglang.srt.layers.quantization.mxfp4_flashinfer_cutlass_moe import ( Mxfp4FlashinferCutlassMoEMethod, @@ -2590,17 +2594,27 @@ def _prepare_flashinfer_trtllm_activation_params(self, layer: Module) -> None: ("gemm1_beta", self.moe_runner_config.gemm1_beta), ("gemm1_clamp_limit", clamp_limit), ): - tensor = ( - None - if value is None - else torch.full( - (num_experts,), - float(value), - dtype=torch.float32, - device=device, + buffer_name = f"_flashinfer_trtllm_{name}" + if buffer_name in layer._buffers: + existing = layer._buffers[buffer_name] + if (existing is None) != (value is None): + raise ValueError( + f"{name} cannot be enabled or disabled after initialization" + ) + # CUDA graphs retain this storage across online weight updates. + if existing is not None: + existing.fill_(float(value)) + else: + tensor = ( + None + if value is None + else torch.full( + (num_experts,), float(value), dtype=torch.float32, device=device + ) ) - ) - setattr(layer, f"_flashinfer_trtllm_{name}", tensor) + # Memory-saver weight offload preserves registered buffers, not + # arbitrary tensor attributes. These have no checkpoint payload. + layer.register_buffer(buffer_name, tensor, persistent=False) def _prepare_hpc_ops_weights(self, layer: Module) -> None: """Precompute the scale layouts consumed by the HPC-Ops fused MoE kernels. diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index fd86e97283fd..8344cb9a220c 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -1828,8 +1828,9 @@ def quant_weight_ue8m0( weight_dequant: torch.Tensor, weight_block_size: List[int], ): + """Quantize BF16 or FP32 weights with 128x128 power-of-two block scales.""" assert weight_block_size == [128, 128] - assert weight_dequant.dtype == torch.bfloat16, ( + assert weight_dequant.dtype in (torch.bfloat16, torch.float32), ( f"{weight_dequant.dtype=} {weight_dequant.shape=}" ) diff --git a/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py b/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py index 7e5f844b7da5..d5a284b76be5 100644 --- a/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py +++ b/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py @@ -12,6 +12,8 @@ ) from sglang.srt.layers.dp_attention import is_allocation_symmetric from sglang.srt.layers.moe.utils import RoutingMethodType +from sglang.srt.layers.quantization.base_config import QuantizeMethodBase +from sglang.srt.layers.utils.common import copy_or_rebind_param from sglang.srt.runtime_context import ( get_exec, get_parallel, @@ -27,13 +29,16 @@ _MXFP8_QUANTIZE_BACKEND = "cute-dsl" if get_platform().is_sm100 else "cuda" if is_flashinfer_available(): - from flashinfer import shuffle_matrix_a, shuffle_matrix_sf_a from flashinfer.fp4_quantization import block_scale_interleave from flashinfer.fused_moe import trtllm_fp4_block_scale_routed_moe from flashinfer.fused_moe.core import ( _maybe_get_cached_w3_w1_permute_indices, get_w2_permute_indices_with_cache, ) + from flashinfer.utils import ( + get_shuffle_matrix_a_row_indices, + get_shuffle_matrix_sf_a_row_indices, + ) logger = logging.getLogger(__name__) @@ -47,12 +52,22 @@ ) -def _pad_intermediate_size(layer: Module) -> None: - intermediate_size = layer.w13_weight.shape[1] // 2 +def _pad_intermediate_size(layer: Module, intermediate_size: int) -> None: padded_size = (intermediate_size + 127) // 128 * 128 if padded_size == intermediate_size: return + if layer.w13_weight.shape[1] == 2 * padded_size: + # Partial checkpoint loads leave the padding untouched. After a kernel + # row permutation those slots can contain data, so clear them on reload. + for name, fill_value in (("w13_weight", 0), ("w13_weight_scale_inv", 1)): + param = getattr(layer, name).data + param[:, intermediate_size:padded_size].fill_(fill_value) + param[:, padded_size + intermediate_size :].fill_(fill_value) + layer.w2_weight.data[:, :, intermediate_size // 2 :].zero_() + layer.w2_weight_scale_inv.data[:, :, intermediate_size // 32 :].fill_(1) + return + # Gate and up occupy separate halves; each needs its own zero tail. for name, fill_value in ( ("w13_weight", 0), @@ -98,9 +113,9 @@ def _pad_intermediate_size(layer: Module) -> None: def routed_hidden_size(layer: Module) -> int: - """Hidden size the routed GEMM1 expects (uint8 weights hold two fp4/row).""" + """Hidden size the routed GEMM1 expects (packed bytes hold two FP4 values).""" w13 = layer.w13_weight - return w13.shape[2] * 2 if w13.dtype == torch.uint8 else w13.shape[2] + return w13.shape[2] * 2 if w13.dtype in (torch.int8, torch.uint8) else w13.shape[2] class Mxfp8RoutedInputPreQuant(NamedTuple): @@ -112,7 +127,12 @@ class Mxfp8RoutedInputPreQuant(NamedTuple): ready: Optional[torch.cuda.Event] -class Mxfp4FlashinferTrtllmMoEMethod: +def _shuffled_scale_name(name: str) -> str: + """Name of the kernel-layout copy kept next to a checkpoint-layout scale.""" + return f"{name}_shuffled" + + +class Mxfp4FlashinferTrtllmMoEMethod(QuantizeMethodBase): fuse_routed_scaling_factor_in_topk = True def __init__(self, fp8_method, prefix: str): @@ -133,7 +153,7 @@ def create_moe_runner(self, layer, moe_runner_config): self.runner = None swiglu_limit = moe_runner_config.swiglu_limit - self._gemm1_clamp_limit_tensor = ( + clamp_limit_tensor = ( torch.full( (layer.num_local_experts,), swiglu_limit, @@ -143,6 +163,11 @@ def create_moe_runner(self, layer, moe_runner_config): if swiglu_limit is not None else None ) + layer.register_buffer( + "_gemm1_clamp_limit_tensor", + clamp_limit_tensor, + persistent=False, + ) def create_weights( self, @@ -156,6 +181,10 @@ def create_weights( from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported fp4_block_k = 32 + self._intermediate_size_unpadded = intermediate_size_per_partition + self._weights_in_kernel_layout = False + layer.register_buffer("_mxfp4_w13_restore_rows", None, persistent=False) + layer.register_buffer("_mxfp4_w2_restore_rows", None, persistent=False) w13_weight = Parameter( torch.empty( @@ -207,22 +236,40 @@ def create_weights( layer.register_parameter("w2_weight_scale_inv", w2_weight_scale) set_weight_attrs(w2_weight_scale, scale_attrs) - def process_weights_after_loading(self, layer: Module) -> None: - from sglang.srt.layers.quantization.utils import reorder_w1w3_to_w3w1 + def restore_weights_before_loading(self, layer: Module) -> None: + """Restore checkpoint row order without replacing captured weight storage.""" + if not self._weights_in_kernel_layout: + return + for weight, rows in ( + (layer.w13_weight, layer._mxfp4_w13_restore_rows), + (layer.w2_weight, layer._mxfp4_w2_restore_rows), + ): + for expert in weight.data: + expert.copy_(expert[rows]) + self._weights_in_kernel_layout = False + def process_weights_after_loading(self, layer: Module) -> None: + """Turn the freshly loaded checkpoint layout into what the kernel reads. + + The kernel wants each expert's rows permuted, with ``w13`` reordered + from ``[w1, w3]`` to ``[w3, w1]`` first. For the weights that is a pure + row permutation of a tensor with the same shape and dtype, so it is + applied in place, one expert at a time; the parameters keep their + identity and storage, which is what ``load_weights`` (via + ``weight_loader``) and a captured CUDA graph both depend on. The + scales additionally change dtype, interleaving and extent, so their + kernel layout lives in separate ``*_shuffled`` parameters rebuilt from + the checkpoint-layout scales here; the checkpoint-layout scales are + never modified and stay loadable. + """ + if self._weights_in_kernel_layout: + return self._fp8.process_weights_after_loading(layer) if getattr(layer, "_mega_moe_weights_built", False): return - _pad_intermediate_size(layer) - - w13_w, w13_s = reorder_w1w3_to_w3w1( - layer.w13_weight.data, layer.w13_weight_scale_inv.data - ) - layer.w13_weight = Parameter(w13_w, requires_grad=False) - layer.w13_weight_scale_inv = Parameter(w13_s, requires_grad=False) - + _pad_intermediate_size(layer, self._intermediate_size_unpadded) log_info_on_rank0( logger, f"Shuffling FP4 expert weights for TRT-LLM MxFP4 kernel " @@ -234,92 +281,95 @@ def process_weights_after_loading(self, layer: Module) -> None: w13_scale = layer.w13_weight_scale_inv.data w2_scale = layer.w2_weight_scale_inv.data num_experts = w13.shape[0] + device = w13.device if w13_scale.dtype == torch.float32: w13_scale = w13_scale.to(torch.float8_e8m0fnu) w2_scale = w2_scale.to(torch.float8_e8m0fnu) + w13_u8 = w13.view(torch.uint8) + w13_s_u8 = w13_scale.view(torch.uint8) + w2_u8 = w2.view(torch.uint8) + w2_s_u8 = w2_scale.view(torch.uint8) + epilogue_tile_m = 128 - g1_w, g1_s, g2_w, g2_s = [], [], [], [] if _USE_OFFICIAL_SHUFFLE: cache: dict = {} - for i in range(num_experts): - w13_u8 = w13[i].view(torch.uint8) - w13_s_u8 = w13_scale[i].view(torch.uint8) - w2_u8 = w2[i].view(torch.uint8) - w2_s_u8 = w2_scale[i].view(torch.uint8) - - perm = _maybe_get_cached_w3_w1_permute_indices( - cache, - w13_u8, - epilogue_tile_m, - ) - g1_w.append(w13_u8[perm.to(w13_u8.device)].contiguous()) - perm_sf = _maybe_get_cached_w3_w1_permute_indices( - cache, - w13_s_u8, - epilogue_tile_m, - num_elts_per_sf=16, - ) - g1_s.append( - block_scale_interleave( - w13_s_u8[perm_sf.to(w13_s_u8.device)].contiguous() - ) - ) - - perm = get_w2_permute_indices_with_cache( - cache, - w2_u8, - epilogue_tile_m, - ) - g2_w.append(w2_u8[perm.to(w2_u8.device)].contiguous()) - perm_sf = get_w2_permute_indices_with_cache( - cache, - w2_s_u8, - epilogue_tile_m, - num_elts_per_sf=16, - ) - g2_s.append( - block_scale_interleave( - w2_s_u8[perm_sf.to(w2_s_u8.device)].contiguous() - ) - ) + w13_rows = _maybe_get_cached_w3_w1_permute_indices( + cache, w13_u8[0], epilogue_tile_m + ) + w13_sf_rows = _maybe_get_cached_w3_w1_permute_indices( + cache, w13_s_u8[0], epilogue_tile_m, num_elts_per_sf=16 + ) + w2_rows = get_w2_permute_indices_with_cache( + cache, w2_u8[0], epilogue_tile_m + ) + w2_sf_rows = get_w2_permute_indices_with_cache( + cache, w2_s_u8[0], epilogue_tile_m, num_elts_per_sf=16 + ) else: - for i in range(num_experts): - g1_w.append(shuffle_matrix_a(w13[i].view(torch.uint8), epilogue_tile_m)) - g1_s.append( - shuffle_matrix_sf_a(w13_scale[i].view(torch.uint8), epilogue_tile_m) - ) - g2_w.append(shuffle_matrix_a(w2[i].view(torch.uint8), epilogue_tile_m)) - g2_s.append( - shuffle_matrix_sf_a(w2_scale[i].view(torch.uint8), epilogue_tile_m) - ) - - layer.w13_weight = Parameter(torch.stack(g1_w), requires_grad=False) - layer.w13_weight_scale_inv = Parameter( - torch.stack(g1_s) - .view(torch.float8_e4m3fn) - .reshape(num_experts, w13.shape[1], -1), - requires_grad=False, + w13_rows = get_shuffle_matrix_a_row_indices(w13_u8[0], epilogue_tile_m) + w13_sf_rows = get_shuffle_matrix_sf_a_row_indices( + w13_s_u8[0], epilogue_tile_m + ) + w2_rows = get_shuffle_matrix_a_row_indices(w2_u8[0], epilogue_tile_m) + w2_sf_rows = get_shuffle_matrix_sf_a_row_indices( + w2_s_u8[0], epilogue_tile_m + ) + # [w1, w3] -> [w3, w1] is a row permutation as well; fold it into the + # gather so each expert is rewritten exactly once. + half = w13.shape[1] // 2 + w3_w1_rows = torch.cat( + [ + torch.arange(half, 2 * half, device=device), + torch.arange(0, half, device=device), + ] ) - layer.w2_weight = Parameter(torch.stack(g2_w), requires_grad=False) - layer.w2_weight_scale_inv = Parameter( - torch.stack(g2_s) - .view(torch.float8_e4m3fn) - .reshape(num_experts, w2.shape[1], -1), - requires_grad=False, + w13_rows = w3_w1_rows[w13_rows.to(device)] + w13_sf_rows = w3_w1_rows[w13_sf_rows.to(device)] + w2_rows = w2_rows.to(device) + w2_sf_rows = w2_sf_rows.to(device) + first_build = layer._mxfp4_w13_restore_rows is None + if first_build: + layer._mxfp4_w13_restore_rows = torch.argsort(w13_rows) + layer._mxfp4_w2_restore_rows = torch.argsort(w2_rows) + + g1_s, g2_s = [], [] + for i in range(num_experts): + w13_u8[i].copy_(w13_u8[i][w13_rows]) + w2_u8[i].copy_(w2_u8[i][w2_rows]) + g1_s.append(block_scale_interleave(w13_s_u8[i][w13_sf_rows].contiguous())) + g2_s.append(block_scale_interleave(w2_s_u8[i][w2_sf_rows].contiguous())) + + # Keep scales as Parameters: EPLB moves expert-indexed parameters, not buffers. + # FP32 load scales stay resident alongside shuffled scales (1/4 of packed weight bytes). + copy_or_rebind_param( + layer, + _shuffled_scale_name("w13_weight_scale_inv"), + torch.stack(g1_s).reshape(num_experts, w13.shape[1], -1), + ) + copy_or_rebind_param( + layer, + _shuffled_scale_name("w2_weight_scale_inv"), + torch.stack(g2_s).reshape(num_experts, w2.shape[1], -1), ) self._register_static_scale_ones(layer) - torch.cuda.empty_cache() + self._weights_in_kernel_layout = True + if first_build: + torch.cuda.empty_cache() def _register_static_scale_ones(self, layer: Module) -> None: + # Constant across reloads; created once so their addresses stay valid + # for a captured CUDA graph. device = layer.w13_weight.device for name in ( "output1_scale_scalar", "output1_scale_gate_scalar", "output2_scale_scalar", ): + if getattr(layer, name, None) is not None: + continue layer.register_buffer( name, torch.ones(layer.num_local_experts, device=device, dtype=torch.float32), @@ -349,19 +399,39 @@ def apply( layer: Module, dispatch_output: DispatchOutput, ) -> CombineInput: - from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput + from sglang.srt.layers.moe.token_dispatcher import ( + DispatchOutputChecker, + StandardCombineInput, + ) from sglang.srt.layers.moe.topk import TopKOutputChecker hidden_states = dispatch_output.hidden_states topk_output = dispatch_output.topk_output pre_quant = getattr(dispatch_output, "hidden_states_pre_quant", None) + is_flashinfer_dispatch = DispatchOutputChecker.format_is_flashinfer( + dispatch_output + ) + if is_flashinfer_dispatch and ( + hidden_states.dtype != torch.bfloat16 + or dispatch_output.hidden_states_scale is not None + ): + raise ValueError( + "FlashInfer A2A + TRT-LLM MXFP4 MoE requires an unquantized BF16 " + "dispatch payload; activations are quantized locally before GEMM." + ) - w13 = layer.w13_weight - w2 = layer.w2_weight - w13_scale = layer.w13_weight_scale_inv - w2_scale = layer.w2_weight_scale_inv + # Loaders keep signed checkpoint bytes; FlashInfer identifies packed FP4 + # and its logical dimensions through unsigned, zero-copy byte views. + w13 = layer.w13_weight.view(torch.uint8) + w2 = layer.w2_weight.view(torch.uint8) + w13_scale = getattr(layer, _shuffled_scale_name("w13_weight_scale_inv")).view( + torch.float8_e4m3fn + ) + w2_scale = getattr(layer, _shuffled_scale_name("w2_weight_scale_inv")).view( + torch.float8_e4m3fn + ) - intermediate_size = w2.shape[2] * 2 if w2.dtype == torch.uint8 else w2.shape[2] + intermediate_size = w2.shape[2] * 2 hidden_size = routed_hidden_size(layer) num_local_experts = layer.num_local_experts @@ -449,7 +519,7 @@ def apply( gemm1_bias=None, gemm1_alpha=None, gemm1_beta=None, - gemm1_clamp_limit=self._gemm1_clamp_limit_tensor, + gemm1_clamp_limit=layer._gemm1_clamp_limit_tensor, gemm2_weights=w2, gemm2_weights_scale=w2_scale, gemm2_bias=None, @@ -475,6 +545,14 @@ def apply( else: output = result[0] + if is_flashinfer_dispatch: + from sglang.srt.layers.moe.token_dispatcher.flashinfer import ( + FlashinferCombineInput, + ) + + return FlashinferCombineInput(hidden_states=output) + # Prefill uses all-gather dispatch and reduce-scatter combine rather than + # the decode A2A workspace, even when FlashInfer A2A is selected globally. return StandardCombineInput(hidden_states=output) 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 4cb5af6354aa..0230690544dc 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 @@ -339,21 +339,30 @@ def update_weights_from_tensor( monkey_patch_torch_reductions() self._assert_weight_cache_inactive("update_weights_from_tensor") - if load_format == "flattened_bucket": - # Handle flattened bucket format - return self._update_weights_from_flattened_bucket( - flattened_tensor_bucket_dict=named_tensors - ) - # We need to get device after patch otherwise the device would be wrong device_module = torch.get_device_module(self.device) infered_device = device_module.current_device() - named_tensors = [ - (name, _unwrap_tensor(tensor, tp_rank=self.tp_rank, device=infered_device)) - for name, tensor in named_tensors - ] - if load_format == "direct": + # Two input shapes reach this point. The per-tensor formats hand over a + # list of (name, tensor) pairs, where a tensor may still be a + # LocalSerializedTensor carrying one CUDA IPC handle per TP rank: open + # this rank's handle and move the result to this device before loading. + # A flattened bucket is instead a dict holding one big device tensor and + # the metadata that says how to slice it; its loader below does that. + if load_format != "flattened_bucket": + named_tensors = [ + ( + name, + _unwrap_tensor(tensor, tp_rank=self.tp_rank, device=infered_device), + ) + for name, tensor in named_tensors + ] + + if load_format == "flattened_bucket": + self._update_weights_from_flattened_bucket( + flattened_tensor_bucket_dict=named_tensors + ) + elif load_format == "direct": _model_load_weights_direct(self.get_model(), named_tensors) elif load_format in self.custom_weight_loaders: custom_loader = dynamic_import(load_format) @@ -362,13 +371,18 @@ def update_weights_from_tensor( self.get_model().load_weights(named_tensors) else: raise NotImplementedError(f"Unknown load_format={load_format}") + # Tensors deserialized from CUDA IPC handles alias storage owned by the + # sender, who may free or reuse it as soon as this call returns. The + # loads above only enqueue device-to-device copies, so wait for them + # before handing control back. + device_module.synchronize() return True, "Success" def _update_weights_from_flattened_bucket( self: WeightUpdater, flattened_tensor_bucket_dict, - ): - """Handle flattened bucket format for weight updates""" + ) -> None: + """Load a flattened bucket, raising if reconstruction or loading fails.""" flattened_tensor = flattened_tensor_bucket_dict["flattened_tensor"] metadata = flattened_tensor_bucket_dict["metadata"] @@ -394,8 +408,6 @@ def _update_weights_from_flattened_bucket( # Load the reconstructed tensors using the standard method self.get_model().load_weights(reconstructed_tensors) - return True, "Success" - def update_weights_from_ipc(self: WeightUpdater, recv_req): """Update weights from IPC for checkpoint-engine integration.""" self._assert_weight_cache_inactive("update_weights_from_ipc") diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 788ac19b8cde..658121ffa3c3 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -42,6 +42,7 @@ import torch from sglang.srt.constants import GIB_BYTES +from sglang.srt.layers.moe import get_moe_runner_backend from sglang.srt.model_loader.post_load import stage_module_for_post_load from sglang.srt.model_loader.remote_instance_weight_loader_utils import ( RemoteInstanceWeightLoaderBackend, @@ -52,6 +53,7 @@ get_exec, get_model, get_parallel, + get_platform, get_server_args, ) from sglang.srt.utils import get_available_gpu_memory @@ -221,6 +223,15 @@ def _get_quantization_config( quant_config.is_dsv4_fp4_experts = is_deepseek_v4(model_config.hf_config) quant_config.dequant_fp4_to_fp8 = envs.SGLANG_DSV4_FP4_DEQUANT.get() + if ( + quant_config.is_fp4_experts + and not quant_config.dequant_fp4_to_fp8 + and get_moe_runner_backend().is_flashinfer_trtllm_routed() + and not get_platform().is_sm100 + ): + raise ValueError( + "flashinfer_trtllm_routed with MXFP4 experts requires SM100." + ) # Handle hybrid NVFP4 moe (nvidia/DeepSeek-V4-Pro-NVFP4) nvfp4_meta = model_config.nvfp4_moe_meta if nvfp4_meta is not None: @@ -1084,7 +1095,7 @@ def postprocess_weights(model, target_device): def restore_weights_before_loading(model, target_device): """Undo in-place quant packing so fresh weights can be loaded.""" for module, quant_method in _modules_with_quant_method(model): - # AMX packing and the MXFP4 backend wrappers are duck-typed and cannot restore + # Duck-typed methods without the quantization hook contract cannot restore. if isinstance(quant_method, QuantizeMethodBase): with device_loading_context(module, target_device): quant_method.restore_weights_before_loading(module) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index bb551ac50c75..ad631afe0e53 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -779,6 +779,9 @@ def init_new( scheduler_input_port = worker_ports[dp_rank] is_joiner = ep_joiner_of(resolving_view(server_args)) + # Ordinary peer nodes connect to the head's TCP endpoints. Checking + # those port numbers on the peer can conflict with unrelated services. + check_local_ports = server_args.node_rank == 0 or is_joiner # Under SGLANG_DISTRIBUTED_INIT_METHOD_OVERRIDE, SGLang never binds # dist_init_port / nccl_port (rendezvous uses the externally-managed # store; see distributed/bootstrap.py:_resolve_dist_init_method), so @@ -787,7 +790,7 @@ def init_new( envs.SGLANG_DISTRIBUTED_INIT_METHOD_OVERRIDE.get() ) try: - if dp_rank is None: + if dp_rank is None and check_local_ports: if not (is_joiner or dist_init_overridden): wait_port_available(dist_init_port, "dist_init_port") wait_port_available(port_base, "port_base") @@ -800,7 +803,7 @@ def init_new( wait_port_available(load_collector_port, "load_collector_port") # Check scheduler_input_port only for dp. # Skip check when using worker_ports since the port is already bound by our ZMQ socket - if dp_rank is None or worker_ports is None: + if check_local_ports and (dp_rank is None or worker_ports is None): wait_port_available(scheduler_input_port, "scheduler_input_port") except ValueError: logger.exception( diff --git a/python/sglang/srt/state_capturer/routed_experts.py b/python/sglang/srt/state_capturer/routed_experts.py index fab5c4d9207e..309bd4853ae9 100644 --- a/python/sglang/srt/state_capturer/routed_experts.py +++ b/python/sglang/srt/state_capturer/routed_experts.py @@ -11,6 +11,7 @@ is_dp_attention_enabled, ) from sglang.srt.layers.moe import get_moe_a2a_backend +from sglang.srt.layers.moe.utils import is_moe_input_scattered_across_dp_ranks from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.runtime_context import ( get_exec, @@ -30,10 +31,9 @@ def _is_scattered_a2a_backend() -> bool: class RoutedExpertsCapturer(BaseTopkCapturer): """Capturer for routed experts with host buffer. - Routed experts share a global device buffer across DP ranks (indexed by - dp_rank), so `_get_local_slice` overrides the default to apply DP-rank-aware - slicing. The device cache also holds extra columns for any fused shared - experts; the host cache and user-facing return drop them via the + DP-concatenated routing needs a DP-rank offset at readback; DP-local routing + already starts at buffer offset zero. The device cache holds extra columns + for fused shared experts; the host cache and user-facing return drop them via the [:topk_size] truncation. """ @@ -120,8 +120,8 @@ def _get_local_slice( can_run_graph: bool, cuda_graph_batch: Optional[int], ) -> torch.Tensor: - # Gathered rows start at buffer offset zero on every DP rank. - if is_dp_attention_enabled() and not _is_scattered_a2a_backend(): + # DP-local routing and attention-TP gathered rows both start at zero. + if is_dp_attention_enabled() and not is_moe_input_scattered_across_dp_ranks(): # GPU->CPU sync would break overlap; operate on CPU directly. local_start_pos, local_num_tokens = get_dp_local_slice_cpu( forward_batch, can_run_graph, cuda_graph_batch @@ -164,14 +164,13 @@ def extract_routed_experts_from_meta_info(data): def disable_routed_experts_capture_for_draft(model: Any) -> None: """Opt every draft MoE ``TopK`` out of routed-experts (R3) capture. - Capture is target-only; a draft ``TopK`` must never write the target's - process-global buffer. ``HashTopK`` has no ``topk_config`` and never - captures, so it is left untouched. + Capture is target-only; draft ``TopK`` and ``HashTopK`` modules must never + write the target's process-global buffer. """ - # Lazy import: ``layers.moe.topk`` imports ``get_global_experts_capturer`` - # from this module, so a top-level import here would be circular. + # Lazy imports avoid circular dependencies with the capture call sites. + from sglang.srt.layers.moe.hash_topk import HashTopK from sglang.srt.layers.moe.topk import TopK for module in model.modules(): - if isinstance(module, TopK): + if isinstance(module, (TopK, HashTopK)): module.topk_config.allow_routed_experts_capture = False diff --git a/test/registered/kernels/ops/quantization/test_mxfp4_flashinfer_trtllm_hot_reload.py b/test/registered/kernels/ops/quantization/test_mxfp4_flashinfer_trtllm_hot_reload.py new file mode 100644 index 000000000000..3d7ac2c49a9d --- /dev/null +++ b/test/registered/kernels/ops/quantization/test_mxfp4_flashinfer_trtllm_hot_reload.py @@ -0,0 +1,297 @@ +"""Online reload contract for the FlashInfer TRT-LLM MXFP4 MoE layout.""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch +from torch.nn import Module + +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="4-gpu-b200") + +if not torch.cuda.is_available(): + pytest.skip("CUDA required", allow_module_level=True) + +from sglang.srt.layers.quantization.mxfp4_flashinfer_trtllm_moe import ( + Mxfp4FlashinferTrtllmMoEMethod, +) +from sglang.srt.model_loader.loader import DefaultModelLoader +from sglang.srt.runtime_context import get_parallel +from sglang.srt.utils import is_sm100_supported +from sglang.test.test_utils import CustomTestCase + +if not is_sm100_supported(): + pytest.skip("FlashInfer TRT-LLM MXFP4 MoE requires SM100+", allow_module_level=True) + + +LOAD_PARAMS = ( + "w13_weight", + "w2_weight", + "w13_weight_scale_inv", + "w2_weight_scale_inv", +) +KERNEL_TENSORS = ( + "w13_weight", + "w2_weight", + "w13_weight_scale_inv_shuffled", + "w2_weight_scale_inv_shuffled", +) +NUM_EXPERTS = 2 +HIDDEN_SIZE = 128 +INTERMEDIATE_SIZE = 128 + + +class _NoOpFp8: + """Stand-in for the wrapped ``Fp8MoEMethod``. + + On the FlashInfer MXFP4 backend the real ``process_weights_after_loading`` + only re-views the FP4 payloads as ``int8`` and returns before any scale + transform, so running it a second time is a no-op and this test can leave + it out. If that branch ever starts transforming weights or scales, the + reload contract below must be checked with the real method instead. + """ + + def process_weights_after_loading(self, layer: Module) -> None: + pass + + +def _build_layer( + *, + num_experts=NUM_EXPERTS, + hidden_size=HIDDEN_SIZE, + intermediate_size=INTERMEDIATE_SIZE, +) -> Module: + layer = Module().to("cuda") + layer.num_local_experts = num_experts + layer.num_experts = num_experts + layer.moe_ep_rank = 0 + method = Mxfp4FlashinferTrtllmMoEMethod.__new__(Mxfp4FlashinferTrtllmMoEMethod) + method._fp8 = _NoOpFp8() + method.prefix = "test" + method.flashinfer_mxfp4_moe_precision = "default" + layer.quant_method = method + with torch.device("cuda"): + method.create_weights( + layer, + num_experts, + hidden_size, + intermediate_size, + torch.bfloat16, + weight_loader=lambda *args, **kwargs: None, + ) + method.create_moe_runner(layer, SimpleNamespace(swiglu_limit=10.0)) + return layer + + +def _load_weights(layer: Module, seed: int, names=LOAD_PARAMS) -> None: + generator = torch.Generator(device="cuda").manual_seed(seed) + intermediate_size = layer.quant_method._intermediate_size_unpadded + for name in LOAD_PARAMS: + param = getattr(layer, name) + shape = list(param.shape) + if name.startswith("w13"): + shape[1] = 2 * intermediate_size + else: + shape[2] = intermediate_size // (32 if "scale" in name else 2) + if "scale" in name: + encoded = torch.randint( + 100, + 140, + shape, + generator=generator, + device="cuda", + dtype=torch.uint8, + ) + loaded = encoded.view(torch.float8_e8m0fnu) + else: + loaded = torch.randint( + -128, + 127, + shape, + generator=generator, + device="cuda", + dtype=torch.int8, + ) + if name not in names: + continue + if name.startswith("w13"): + half = param.shape[1] // 2 + param.data[:, :intermediate_size].copy_(loaded[:, :intermediate_size]) + param.data[:, half : half + intermediate_size].copy_( + loaded[:, intermediate_size:] + ) + else: + param.data[:, :, : loaded.shape[2]].copy_(loaded) + + +def _layout(layer: Module) -> dict[str, tuple[tuple[int, ...], torch.dtype]]: + return { + name: (tuple(getattr(layer, name).shape), getattr(layer, name).dtype) + for name in LOAD_PARAMS + } + + +def _kernel_tensors(layer: Module) -> dict[str, torch.Tensor]: + return {name: getattr(layer, name) for name in KERNEL_TENSORS} + + +@pytest.mark.parametrize("intermediate_size", [64, 128]) +def test_hot_reload_matches_fresh_load_and_preserves_kernel_addresses( + intermediate_size, +): + layer = _build_layer(intermediate_size=intermediate_size) + load_params = {name: getattr(layer, name) for name in LOAD_PARAMS} + + _load_weights(layer, seed=0) + DefaultModelLoader.postprocess_weights(layer, torch.device("cuda")) + load_layout = _layout(layer) + initial = { + name: tensor.detach().clone() for name, tensor in _kernel_tensors(layer).items() + } + addresses = { + name: tensor.data_ptr() for name, tensor in _kernel_tensors(layer).items() + } + + reference = _build_layer(intermediate_size=intermediate_size) + _load_weights(reference, seed=1) + DefaultModelLoader.postprocess_weights(reference, torch.device("cuda")) + expected = { + name: tensor.detach().clone() + for name, tensor in _kernel_tensors(reference).items() + } + + # The loader hooks bracket every update session; neither may disturb the + # checkpoint-layout parameters ``load_weights`` writes into. + DefaultModelLoader.restore_weights_before_loading(layer, torch.device("cuda")) + assert _layout(layer) == load_layout + for name in LOAD_PARAMS: + assert getattr(layer, name) is load_params[name] + assert hasattr(getattr(layer, name), "weight_loader") + + _load_weights(layer, seed=1) + DefaultModelLoader.postprocess_weights(layer, torch.device("cuda")) + + for name in LOAD_PARAMS: + assert getattr(layer, name) is load_params[name] + for name, tensor in _kernel_tensors(layer).items(): + assert tensor.data_ptr() == addresses[name] + assert not torch.equal( + initial[name].view(torch.uint8), tensor.view(torch.uint8) + ) + assert torch.equal(expected[name].view(torch.uint8), tensor.view(torch.uint8)) + + +@pytest.mark.parametrize("intermediate_size", [64, 128]) +@pytest.mark.parametrize("names", [(), ("w2_weight",), ("w13_weight_scale_inv",)]) +def test_partial_update_preserves_untouched_experts(intermediate_size, names): + """A session must not reshuffle untouched expert weights or discard scale-only updates.""" + layer = _build_layer(intermediate_size=intermediate_size) + _load_weights(layer, seed=0) + DefaultModelLoader.postprocess_weights(layer, torch.device("cuda")) + addresses = { + name: tensor.data_ptr() for name, tensor in _kernel_tensors(layer).items() + } + + reference = _build_layer(intermediate_size=intermediate_size) + _load_weights(reference, seed=0) + _load_weights(reference, seed=1, names=names) + DefaultModelLoader.postprocess_weights(reference, torch.device("cuda")) + + DefaultModelLoader.restore_weights_before_loading(layer, torch.device("cuda")) + DefaultModelLoader.restore_weights_before_loading(layer, torch.device("cuda")) + _load_weights(layer, seed=1, names=names) + DefaultModelLoader.postprocess_weights(layer, torch.device("cuda")) + DefaultModelLoader.postprocess_weights(layer, torch.device("cuda")) + for name, tensor in _kernel_tensors(layer).items(): + assert tensor.data_ptr() == addresses[name] + assert torch.equal( + tensor.view(torch.uint8), getattr(reference, name).view(torch.uint8) + ) + + +class TestMxfp4KernelForward(CustomTestCase): + @torch.inference_mode() + def test_mxfp8_forward_accepts_int8_load_parameters(self): + """Signed checkpoint bytes must remain loadable across real MoE forwards.""" + from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput + from sglang.srt.layers.moe.token_dispatcher.flashinfer import ( + FlashinferCombineInput, + FlashinferDispatchOutput, + ) + from sglang.srt.layers.moe.topk import StandardTopKOutput + from sglang.srt.layers.quantization import mxfp4_flashinfer_trtllm_moe + + layer = _build_layer(num_experts=32, hidden_size=4096, intermediate_size=2048) + # Module device moves can replace buffers; the kernel must read the live copy. + old_clamp = layer._gemm1_clamp_limit_tensor + layer._apply(lambda tensor: tensor.clone()) + old_clamp.zero_() + generator = torch.Generator(device="cuda").manual_seed(0) + hidden_states = torch.randn( + 8, 4096, generator=generator, device="cuda", dtype=torch.bfloat16 + ) + dispatch = StandardDispatchOutput( + hidden_states=hidden_states, + hidden_states_scale=None, + topk_output=StandardTopKOutput( + topk_weights=torch.full( + (8, 6), 1 / 6, device="cuda", dtype=torch.float32 + ), + topk_ids=torch.arange(6, device="cuda", dtype=torch.int32).repeat(8, 1), + router_logits=None, + ), + ) + # A2A reserves padding slots for unequal source token counts. + dispatch.topk_output.topk_ids[-1].fill_(-1) + dispatch.topk_output.topk_weights[-1].zero_() + a2a_dispatch = FlashinferDispatchOutput( + hidden_states=dispatch.hidden_states, + hidden_states_scale=None, + topk_output=dispatch.topk_output, + ) + outputs = [] + # A local kernel call needs no distributed symmetric-memory allocator. + with ( + get_parallel().override(tp_group=None), + patch.object( + mxfp4_flashinfer_trtllm_moe, + "is_allocation_symmetric", + return_value=False, + ), + ): + for packed_byte in (-86, 34): + if outputs: + DefaultModelLoader.restore_weights_before_loading( + layer, torch.device("cuda") + ) + for name in LOAD_PARAMS: + getattr(layer, name).fill_( + 1 / 32 if "scale" in name else packed_byte + ) + DefaultModelLoader.postprocess_weights(layer, torch.device("cuda")) + output = layer.quant_method.apply(layer, dispatch).hidden_states + torch.cuda.synchronize() + self.assertEqual(output.shape, hidden_states.shape) + self.assertEqual(output.dtype, torch.bfloat16) + self.assertTrue(torch.isfinite(output).all().item()) + self.assertGreater(torch.count_nonzero(output).item(), 0) + self.assertEqual(torch.count_nonzero(output[-1]).item(), 0) + a2a_result = layer.quant_method.apply(layer, a2a_dispatch) + self.assertIsInstance(a2a_result, FlashinferCombineInput) + torch.testing.assert_close( + a2a_result.hidden_states, output, rtol=0, atol=0 + ) + self.assertEqual(layer.w13_weight.dtype, torch.int8) + self.assertEqual(layer.w2_weight.dtype, torch.int8) + outputs.append(output.clone()) + self.assertFalse(torch.equal(*outputs)) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__])) diff --git a/test/registered/moe/test_hash_topk.py b/test/registered/moe/test_hash_topk.py index a4546ca0d5be..4d0f54ded2a2 100644 --- a/test/registered/moe/test_hash_topk.py +++ b/test/registered/moe/test_hash_topk.py @@ -6,6 +6,7 @@ from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo from sglang.srt.layers.moe import hash_topk as hash_topk_module +from sglang.srt.layers.moe import topk as topk_module from sglang.srt.layers.moe.hash_topk import HashTopK from sglang.srt.layers.moe.topk import ( StandardTopKOutput, @@ -13,9 +14,12 @@ from sglang.srt.models.deepseek_v2 import DeepseekV2MoE from sglang.srt.runtime_context import get_parallel from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler +from sglang.srt.state_capturer.routed_experts import ( + disable_routed_experts_capture_for_draft, +) from sglang.test.ci.ci_register import register_cpu_ci -register_cpu_ci(est_time=6, suite="stage-a-test-cpu-intel") +register_cpu_ci(est_time=30, suite="stage-a-test-cpu-intel") @pytest.fixture(autouse=True) @@ -83,6 +87,140 @@ def on_select_experts(self, *, topk_ids): assert recorded["topk_ids"].tolist() == [[0, 65], [63, 127]] +def test_hash_topk_captures_logical_expert_ids(monkeypatch): + captured = {} + + class FakeCapturer: + def capture(self, *, layer_id, topk_indices): + captured["layer_id"] = layer_id + captured["topk_ids"] = topk_indices.clone() + + monkeypatch.setattr( + hash_topk_module, + "capture_routed_experts_if_allowed", + lambda config, layer_id, topk_ids, _num_token_non_padded=None: ( + FakeCapturer().capture(layer_id=layer_id, topk_indices=topk_ids) + if config.allow_routed_experts_capture + else None + ), + ) + + topk = HashTopK( + topk=2, + num_experts=8, + num_fused_shared_experts=0, + vocab_size=2, + scoring_func="sqrtsoftplus", + layer_id=3, + ) + with torch.no_grad(): + topk.tid2eid.copy_(torch.tensor([[1, 5], [2, 7]], dtype=torch.int32)) + + with hash_topk_module.envs.SGLANG_OPT_USE_FUSED_HASH_TOPK.override(False): + output = topk( + hidden_states=torch.empty(2, 4), + router_logits=torch.ones(2, 8), + input_ids=torch.tensor([0, 1], dtype=torch.int64), + ) + + assert output.topk_ids.tolist() == [[1, 5], [2, 7]] + assert captured["layer_id"] == 3 + assert torch.equal( + captured["topk_ids"], + torch.tensor([[1, 5], [2, 7]], dtype=torch.int32), + ) + + +def test_hash_topk_capture_masks_padded_tokens(monkeypatch): + captured = {} + + class FakeCapturer: + def capture(self, *, layer_id, topk_indices): + captured["layer_id"] = layer_id + captured["topk_ids"] = topk_indices.clone() + + monkeypatch.setattr( + topk_module, "get_global_experts_capturer", lambda: FakeCapturer() + ) + + def mask_padded_rows(topk_ids, num_token_non_padded, fill_value=-1): + topk_ids[int(num_token_non_padded.item()) :].fill_(fill_value) + + # ``hash_topk`` imported the helper by name, so it holds its own reference. + monkeypatch.setattr(topk_module, "_mask_topk_ids_padded_region", mask_padded_rows) + monkeypatch.setattr( + hash_topk_module, "_mask_topk_ids_padded_region", mask_padded_rows + ) + + topk = HashTopK( + topk=2, + num_experts=8, + num_fused_shared_experts=0, + vocab_size=2, + scoring_func="sqrtsoftplus", + layer_id=4, + ) + with torch.no_grad(): + topk.tid2eid.copy_(torch.tensor([[1, 5], [2, 7]], dtype=torch.int32)) + + with hash_topk_module.envs.SGLANG_OPT_USE_FUSED_HASH_TOPK.override(False): + topk( + hidden_states=torch.empty(2, 4), + router_logits=torch.ones(2, 8), + input_ids=torch.tensor([0, 1], dtype=torch.int64), + num_token_non_padded=torch.tensor(1), + ) + + assert captured["layer_id"] == 4 + assert torch.equal( + captured["topk_ids"], + torch.tensor([[1, 5], [-1, -1]], dtype=torch.int32), + ) + + +def test_hash_topk_capture_can_be_disabled(monkeypatch): + captured = [] + monkeypatch.setattr( + hash_topk_module, + "capture_routed_experts_if_allowed", + lambda config, *_args: ( + captured.append(True) if config.allow_routed_experts_capture else None + ), + ) + + topk = HashTopK( + topk=2, + num_experts=8, + num_fused_shared_experts=0, + vocab_size=2, + layer_id=0, + ) + topk.topk_config.allow_routed_experts_capture = False + with hash_topk_module.envs.SGLANG_OPT_USE_FUSED_HASH_TOPK.override(False): + topk( + hidden_states=torch.empty(1, 4), + router_logits=torch.ones(1, 8), + input_ids=torch.tensor([0], dtype=torch.int64), + ) + + assert captured == [] + + +def test_draft_hash_topk_capture_is_disabled(): + topk = HashTopK( + topk=2, + num_experts=8, + num_fused_shared_experts=0, + vocab_size=2, + layer_id=0, + ) + model = torch.nn.Sequential(topk) + + disable_routed_experts_capture_for_draft(model) + + assert not topk.topk_config.allow_routed_experts_capture + + def test_hash_topk_empty_output_keeps_per_rank_shared_slot(monkeypatch): monkeypatch.setattr( hash_topk_module, "has_per_rank_fused_shared_slots", lambda *_args: True diff --git a/test/registered/unit/layers/moe/test_flashinfer_megamoe.py b/test/registered/unit/layers/moe/test_flashinfer_megamoe.py index cd73104ed9c9..37c08635ee19 100644 --- a/test/registered/unit/layers/moe/test_flashinfer_megamoe.py +++ b/test/registered/unit/layers/moe/test_flashinfer_megamoe.py @@ -6,6 +6,7 @@ import torch from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase register_cpu_ci(est_time=5, suite="base-a-test-cpu") @@ -214,6 +215,114 @@ def original(value): assert dgm.pack_ue8m0_to_int is original +class TestMegaMoeOutputViewCapability(CustomTestCase): + def test_unsupported_workspace_view_uses_owning_output(self): + """A shared forward signature does not imply backend view support.""" + import pytest + + with pytest.MonkeyPatch.context() as monkeypatch: + module = _load_megamoe_module(monkeypatch) + output = torch.arange(4, dtype=torch.bfloat16) + + class Mega: + supports_output_view = False + + def forward(self, tensors, *, return_workspace_view=False): + if return_workspace_view: + raise ValueError("workspace views are unsupported") + return output + + mega = Mega() + result = module._select_megamoe_forward(mega)(mega, object()) + self.assertIs(result, output) + + +class TestMegaMoeFp4Reload(CustomTestCase): + def test_reload_preserves_load_scales_and_cached_kernel_storage(self): + """Reloads must accept raw scales and refresh the live kernel's tensors.""" + import pytest + + from sglang.srt.layers.utils.common import copy_or_rebind_param + + with pytest.MonkeyPatch.context() as monkeypatch: + module = _load_megamoe_module(monkeypatch) + fake_common = types.ModuleType("sglang.srt.layers.utils.common") + fake_common.copy_or_rebind_param = copy_or_rebind_param + monkeypatch.setitem(sys.modules, fake_common.__name__, fake_common) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) + + class Mega(torch.nn.Module): + def __init__(self, *, bootstrap, fleet_params, weights, backend): + super().__init__() + self.transformed = backend.transformed_weights + + def forward(self, tensors): + return tensors + + def preprocess(weights, **kwargs): + return ( + (weights.w13.flip(1), weights.w13_scale[..., ::4].to(torch.int32)), + (weights.w2.flip(1), weights.w2_scale[..., ::4].to(torch.int32)), + ) + + fake_ep = types.ModuleType("flashinfer.moe_ep") + for name in ( + "MoEWeightPack", + "BootstrapConfig", + "FleetParams", + "MegaConfig", + "DeepGemmMegaMoeConfig", + ): + setattr(fake_ep, name, types.SimpleNamespace) + fake_ep.MoEEpMegaLayer = Mega + fake_ep.preprocess_mega_weights = preprocess + monkeypatch.setitem(sys.modules, "flashinfer.moe_ep", fake_ep) + + layer = torch.nn.Module() + layer.hidden_size = layer.intermediate_size_per_partition = 128 + layer.moe_ep_size = layer.num_experts = 8 + layer.moe_ep_rank = layer.layer_id = 0 + layer.top_k = 1 + layer.moe_runner_config = types.SimpleNamespace(swiglu_limit=10.0) + for name, shape, dtype in ( + ("w13_weight", (1, 256, 64), torch.int8), + ("w2_weight", (1, 128, 64), torch.int8), + ("w13_weight_scale_inv", (1, 256, 4), torch.float32), + ("w2_weight_scale_inv", (1, 128, 4), torch.float32), + ): + layer.register_parameter( + name, + torch.nn.Parameter( + torch.ones(shape, dtype=dtype), requires_grad=False + ), + ) + raw_scales = (layer.w13_weight_scale_inv, layer.w2_weight_scale_inv) + raw_ptrs = [scale.data_ptr() for scale in raw_scales] + cached = None + for value in (1, 2, 4): + for scale in raw_scales: + scale.data.copy_(torch.full_like(scale, value, dtype=torch.float32)) + module.prepare_fp4_moe_weights_for_flashinfer_megamoe(layer) + for scale, ptr in zip(raw_scales, raw_ptrs): + self.assertEqual(scale.dtype, torch.float32) + self.assertEqual(scale.shape[-1], 4) + self.assertEqual(scale.data_ptr(), ptr) + mega = module.ensure_fp4_moe_layer_for_flashinfer_megamoe(layer) + tensors = [tensor for pair in mega.transformed for tensor in pair] + if cached is None: + cached = mega + kernel_ptrs = [tensor.data_ptr() for tensor in tensors] + self.assertIs(mega, cached) + self.assertEqual([tensor.data_ptr() for tensor in tensors], kernel_ptrs) + for _, scale in cached.transformed: + torch.testing.assert_close( + scale, + torch.full_like(scale, value, dtype=torch.int32), + rtol=0, + atol=0, + ) + + if __name__ == "__main__": import pytest diff --git a/test/registered/unit/layers/quantization/test_deepgemm_ue8m0_requant.py b/test/registered/unit/layers/quantization/test_deepgemm_ue8m0_requant.py index ddc039fcd450..36f64120e506 100644 --- a/test/registered/unit/layers/quantization/test_deepgemm_ue8m0_requant.py +++ b/test/registered/unit/layers/quantization/test_deepgemm_ue8m0_requant.py @@ -1,4 +1,4 @@ -"""CPU guards for DeepGEMM UE8M0 weight-scale requantization decisions.""" +"""CPU guards for UE8M0 weight quantization and DeepGEMM requantization.""" from sglang.test.ci.ci_register import register_cpu_ci @@ -28,6 +28,45 @@ def _make_params(n: int = 64, k: int = 128): return weight, weight_scale +class TestUE8M0WeightQuantization(CustomTestCase): + def test_fp32_exports_match_bf16_quantization(self): + """Upcasting BF16 weights for export must preserve their FP8 encoding.""" + generator = torch.Generator().manual_seed(0) + for shape in ((256, 512), (2, 128, 512)): + with self.subTest(shape=shape): + weight = torch.randn(shape, generator=generator, dtype=torch.bfloat16) + exported = weight.float() + expected_weight, expected_scale = fp8_utils.quant_weight_ue8m0( + weight, BLOCK_SIZE + ) + actual_weight, actual_scale = fp8_utils.quant_weight_ue8m0( + exported, BLOCK_SIZE + ) + self.assertTrue( + torch.equal( + actual_weight.view(torch.uint8), + expected_weight.view(torch.uint8), + ) + ) + torch.testing.assert_close(actual_scale, expected_scale, rtol=0, atol=0) + self.assertEqual(exported.dtype, torch.float32) + torch.testing.assert_close(exported, weight.float(), rtol=0, atol=0) + + def test_fp32_values_are_not_rounded_through_bf16(self): + """FP32 values above an FP8 midpoint must not become BF16 ties first.""" + weight = torch.zeros((128, 128), dtype=torch.float32) + weight[0, 0] = 448.0 + # With scale 1, the E4M3 midpoint is 1.0625, exactly representable in BF16. + weight[0, 1] = 1.063 + + quantized, scale = fp8_utils.quant_weight_ue8m0(weight, BLOCK_SIZE) + rounded, _ = fp8_utils.quant_weight_ue8m0(weight.bfloat16(), BLOCK_SIZE) + + self.assertEqual(scale.item(), 1.0) + self.assertEqual(quantized[0, 1].float().item(), 1.125) + self.assertEqual(rounded[0, 1].float().item(), 1.0) + + class TestDeepGemmUE8M0Requant(CustomTestCase): def _enabled_deepgemm_ue8m0(self): return patch.multiple( diff --git a/test/registered/unit/layers/quantization/test_fp8_moe_runner_fallback.py b/test/registered/unit/layers/quantization/test_fp8_moe_runner_fallback.py index 4b3927e32c20..f5376b8f3ef4 100644 --- a/test/registered/unit/layers/quantization/test_fp8_moe_runner_fallback.py +++ b/test/registered/unit/layers/quantization/test_fp8_moe_runner_fallback.py @@ -6,14 +6,60 @@ register_cpu_ci(est_time=5, suite="base-a-test-cpu") import unittest +from types import SimpleNamespace from unittest.mock import patch +import torch + import sglang.srt.layers.quantization.fp8 as fp8 from sglang.srt.layers.moe import MoeRunnerBackend, MoeRunnerConfig +from sglang.srt.layers.moe.fused_moe_triton import FusedMoE +from sglang.srt.layers.quantization.mxfp4_flashinfer_trtllm_moe import ( + Mxfp4FlashinferTrtllmMoEMethod, +) from sglang.test.test_utils import CustomTestCase class TestFp8MoeRunnerFallback(CustomTestCase): + def test_routed_trtllm_selects_the_checkpoint_expert_dtype(self): + """Packed MXFP4 must not enter the routed runner's FP8 GEMM branch.""" + layer = FusedMoE.__new__(FusedMoE) + torch.nn.Module.__init__(layer) + with ( + patch.object( + fp8, + "get_moe_runner_backend", + return_value=MoeRunnerBackend.FLASHINFER_TRTLLM_ROUTED, + ), + patch.object( + fp8, + "get_platform", + return_value=SimpleNamespace( + is_sm100=True, is_sm90=False, is_sm120=False + ), + ), + patch( + "sglang.srt.layers.quantization.mxfp4_flashinfer_trtllm_moe.get_exec", + return_value=SimpleNamespace( + moe=SimpleNamespace(flashinfer_mxfp4_moe_precision="default") + ), + ), + ): + for is_fp4 in (True, False): + with self.subTest(is_fp4=is_fp4): + config = fp8.Fp8Config( + is_checkpoint_fp8_serialized=True, + weight_block_size=[128, 128], + is_fp4_experts=is_fp4, + ) + method = config.get_quant_method( + layer, "model.layers.0.mlp.experts" + ) + expected = ( + Mxfp4FlashinferTrtllmMoEMethod if is_fp4 else fp8.Fp8MoEMethod + ) + self.assertIsInstance(method, expected) + def _runner_backend_for(self, global_backend): method = fp8.Fp8MoEMethod.__new__(fp8.Fp8MoEMethod) with patch.object(fp8, "get_moe_runner_backend", return_value=global_backend): diff --git a/test/registered/unit/layers/quantization/test_fp8_moe_runner_ownership.py b/test/registered/unit/layers/quantization/test_fp8_moe_runner_ownership.py index 55f679941035..ec1d3013d35a 100644 --- a/test/registered/unit/layers/quantization/test_fp8_moe_runner_ownership.py +++ b/test/registered/unit/layers/quantization/test_fp8_moe_runner_ownership.py @@ -20,6 +20,10 @@ from sglang.srt.layers.moe.utils import MoeRunnerBackend from sglang.srt.layers.quantization import fp8 as fp8_module from sglang.srt.layers.quantization.fp8 import Fp8Config, Fp8MoEMethod +from sglang.srt.managers.scheduler_components.weight_updater import ( + _export_static_state, + _import_static_state, +) from sglang.srt.runtime_context import get_flags from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -50,18 +54,18 @@ def _make_block_fp8_method() -> Fp8MoEMethod: ) @staticmethod - def _make_layer(num_local_experts: int = 2) -> SimpleNamespace: - return SimpleNamespace( - num_local_experts=num_local_experts, - w13_weight=torch.empty(num_local_experts, 4), - ) + def _make_layer(num_local_experts: int = 2) -> torch.nn.Module: + layer = torch.nn.Module() + layer.num_local_experts = num_local_experts + layer.w13_weight = torch.nn.Parameter(torch.empty(num_local_experts, 4)) + return layer - def _run_post_load(self, method: Fp8MoEMethod, layer: SimpleNamespace) -> None: + def _run_post_load(self, method: Fp8MoEMethod, layer: torch.nn.Module) -> None: with patch.object(method, "process_weights_after_loading_block_quant") as work: method.process_weights_after_loading(layer) work.assert_called_once_with(layer) - def _assert_activation_params_absent(self, layer: SimpleNamespace) -> None: + def _assert_activation_params_absent(self, layer: torch.nn.Module) -> None: for name in _ACTIVATION_PARAMS: self.assertFalse(hasattr(layer, f"_flashinfer_trtllm_{name}")) @@ -117,6 +121,44 @@ def test_owning_method_skips_params_on_non_trtllm_backend(self): self._assert_activation_params_absent(layer) + def test_activation_params_survive_memory_release_and_reload(self): + """Offload must restore activation constants at the addresses retained by CUDA graphs.""" + method = self._make_block_fp8_method() + layer = self._make_layer() + method.create_moe_runner( + layer=layer, + moe_runner_config=MoeRunnerConfig( + gemm1_alpha=1.5, gemm1_beta=0.25, gemm1_clamp_limit=10.0 + ), + ) + self._run_post_load(method, layer) + retained = { + f"_flashinfer_trtllm_{name}": getattr(layer, f"_flashinfer_trtllm_{name}") + for name in _ACTIVATION_PARAMS + } + expected = {name: tensor.clone() for name, tensor in retained.items()} + backup = _export_static_state(layer) + # Discarded weight-region pages do not retain non-checkpoint constants. + for tensor in retained.values(): + tensor.zero_() + _import_static_state(layer, backup) + for name, tensor in retained.items(): + torch.testing.assert_close(tensor, expected[name], rtol=0, atol=0) + + self._run_post_load(method, layer) + for name, tensor in retained.items(): + self.assertIs(getattr(layer, name), tensor) + torch.testing.assert_close(tensor, expected[name], rtol=0, atol=0) + self.assertNotIn(name, layer.state_dict()) + + method.moe_runner_config.gemm1_alpha = None + with self.assertRaisesRegex(ValueError, "cannot be enabled or disabled"): + self._run_post_load(method, layer) + self.assertIs( + layer._flashinfer_trtllm_gemm1_alpha, + retained["_flashinfer_trtllm_gemm1_alpha"], + ) + class TestFp8MoEAiterQuantInfo(CustomTestCase): """maybe_get_hip_aiter_quant_info assembles what the AITER runner consumes. diff --git a/test/registered/unit/model_loader/test_default_model_loader.py b/test/registered/unit/model_loader/test_default_model_loader.py index da5b3676aeda..f2cb180b1003 100644 --- a/test/registered/unit/model_loader/test_default_model_loader.py +++ b/test/registered/unit/model_loader/test_default_model_loader.py @@ -6,6 +6,8 @@ import torch import sglang.srt.model_loader.loader as loader_mod +from sglang.srt.layers.moe import MoeRunnerBackend +from sglang.srt.layers.quantization.fp8 import Fp8Config from sglang.srt.model_loader.loader import DefaultModelLoader from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -14,6 +16,55 @@ class TestDefaultModelLoader(CustomTestCase): + def test_routed_mxfp4_hardware_is_validated_before_model_construction(self): + for is_fp4, is_sm100 in ((True, False), (True, True), (False, False)): + with self.subTest(is_fp4=is_fp4, is_sm100=is_sm100): + config = Fp8Config( + is_checkpoint_fp8_serialized=True, + weight_block_size=[128, 128], + ) + model_config = SimpleNamespace( + quantization="fp8", + dtype=torch.bfloat16, + is_fp4_experts=is_fp4, + hf_config=SimpleNamespace(architectures=["DeepseekV4ForCausalLM"]), + nvfp4_moe_meta=None, + ) + with ( + patch.object( + loader_mod, + "get_model_architecture", + return_value=(type("Model", (), {}), None), + ), + patch.object(loader_mod, "get_quant_config", return_value=config), + patch.object( + loader_mod, + "get_moe_runner_backend", + return_value=MoeRunnerBackend.FLASHINFER_TRTLLM_ROUTED, + ), + patch.object( + loader_mod, + "get_platform", + return_value=SimpleNamespace(is_sm100=is_sm100), + ), + patch.object( + loader_mod, + "get_device_capability", + return_value=(10, 0) if is_sm100 else (9, 0), + ), + loader_mod.envs.SGLANG_DSV4_FP4_DEQUANT.override(False), + ): + if is_fp4 and not is_sm100: + with self.assertRaisesRegex( + ValueError, "MXFP4 experts requires SM100" + ): + loader_mod._get_quantization_config(model_config, object()) + else: + self.assertIs( + loader_mod._get_quantization_config(model_config, object()), + config, + ) + def test_load_weights_only_precedes_postprocessing(self): events = [] model = Mock() diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 67010b13066f..8aff0ab199c2 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -1758,6 +1758,33 @@ def test_megamoe_combine_dtype_conflicts_with_ikr(self, _): handle_a2a_moe(self._make_args()) +class TestPeerPortArgs(CustomTestCase): + def test_dp_peer_does_not_check_head_ports_on_local_host(self): + """A local listener must not block connecting to the head's TCP endpoints.""" + with ( + socket.socket() as listener, + patch.dict(os.environ, {"SGLANG_WAIT_PORT_TIMEOUT": "1"}), + ): + listener.bind(("127.0.0.1", 0)) + listener.listen() + port = listener.getsockname()[1] + args = ServerArgs(model_path="dummy") + args.enable_dp_attention = True + args.nnodes = 2 + args.node_rank = 1 + args.dist_init_addr = f"192.0.2.1:{port}" + args.nccl_port = port + ports = PortArgs.init_new(args) + self.assertTrue(ports.tokenizer_ipc_name.startswith("tcp://192.0.2.1:")) + self.assertTrue( + ports.scheduler_input_ipc_name.startswith("tcp://192.0.2.1:") + ) + + args.node_rank = 0 + with self.assertRaisesRegex(ValueError, "dist_init_port"): + PortArgs.init_new(args) + + class TestPortArgs(unittest.TestCase): @patch("sglang.srt.server_args.tempfile.NamedTemporaryFile") def test_init_new_standard_case(self, mock_temp_file): diff --git a/test/registered/unit/state_capturer/test_routed_experts_scattered_a2a.py b/test/registered/unit/state_capturer/test_routed_experts_scattered_a2a.py index c8e3654dfba0..3c5b506f7721 100644 --- a/test/registered/unit/state_capturer/test_routed_experts_scattered_a2a.py +++ b/test/registered/unit/state_capturer/test_routed_experts_scattered_a2a.py @@ -6,6 +6,7 @@ import torch +from sglang.srt.layers.moe import utils as moe_utils from sglang.srt.layers.moe.utils import MoeA2ABackend from sglang.srt.state_capturer import routed_experts as re_mod from sglang.test.ci.ci_register import register_cpu_ci @@ -34,6 +35,27 @@ def test_classification(self): class TestGetLocalSliceBackendBranch(CustomTestCase): T, L, K = 16, 3, 4 + def setUp(self): + super().setUp() + patches = ( + mock.patch.object( + moe_utils, + "get_moe_a2a_backend", + side_effect=lambda: re_mod.get_moe_a2a_backend(), + ), + mock.patch.object( + moe_utils, + "should_use_flashinfer_cutlass_moe_fp4_allgather", + return_value=False, + ), + mock.patch.object( + moe_utils, "get_parallel", return_value=SimpleNamespace(dwdp_size=1) + ), + ) + for patcher in patches: + patcher.start() + self.addCleanup(patcher.stop) + def _capturer(self): cap = object.__new__(re_mod.RoutedExpertsCapturer) buf = torch.arange(self.T * self.L * self.K, dtype=torch.int32).reshape( @@ -83,6 +105,34 @@ def test_tp_moe_reads_global_offset(self): out = self._slice(cap, n_local=999) self.assertTrue(torch.equal(out, buf[6:10, :, : self.K])) + def test_dp_local_a2a_preserves_routes_on_nonzero_dp_rank(self): + """DP-local routing must not read the zero-filled global-offset region.""" + for backend in ("flashinfer", "flashinfer_megamoe", "megamoe"): + for can_run_graph in (False, True): + with self.subTest(backend=backend, cuda_graph=can_run_graph): + cap, buf = self._capturer() + expected = buf[:3].clone() + buf[3:].zero_() + with ( + mock.patch.object( + re_mod, "is_dp_attention_enabled", return_value=True + ), + mock.patch.object( + re_mod, + "get_moe_a2a_backend", + return_value=MoeA2ABackend(backend), + ), + mock.patch.object( + re_mod, "get_dp_local_slice_cpu", return_value=(8, 3) + ), + ): + out = cap._get_local_slice( + SimpleNamespace(out_cache_loc=torch.empty(3)), + can_run_graph=can_run_graph, + cuda_graph_batch=4 if can_run_graph else None, + ) + torch.testing.assert_close(out, expected, rtol=0, atol=0) + if __name__ == "__main__": unittest.main() From 6b5f3dfc5a395d6fd56d105c34c8e3542846ca4c Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Sat, 10 Oct 2026 02:49:44 -0700 Subject: [PATCH 41/56] [sglang-miles] [Diffusion] Support distributed weight updates (#42854) --- .../entrypoints/post_training/io_struct.py | 27 +++ .../entrypoints/post_training/weights_api.py | 39 +++ .../runtime/managers/gpu_worker.py | 1 + .../runtime/managers/scheduler.py | 6 + .../gpu_worker_post_training_mixin.py | 76 +++++- .../scheduler_post_training_mixin.py | 42 +++- .../runtime/scheduler_client.py | 19 +- .../multimodal_gen/test/server/gpu_cases.py | 1 + .../test_update_weights_from_distributed.py | 226 ++++++++++++++++++ .../unit/test_distributed_weight_update.py | 82 +++++++ python/sglang/srt/utils/common.py | 42 +++- 11 files changed, 547 insertions(+), 14 deletions(-) create mode 100644 python/sglang/multimodal_gen/test/single_test_file/test_update_weights_from_distributed.py create mode 100644 python/sglang/multimodal_gen/test/unit/test_distributed_weight_update.py 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 de01bc49283b..0c993d7dc3ee 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 @@ -115,3 +115,30 @@ class RolloutResponse(BaseModel): inference_time_s: Optional[float] = None peak_memory_mb: Optional[float] = None + + +@dataclass +class InitWeightsUpdateGroupReqInput: + master_address: str + master_port: int + rank_offset: int + world_size: int + group_name: str + backend: str = "nccl" + + +@dataclass +class UpdateWeightsFromDistributedReqInput: + names: list[str] + dtypes: list[str] + shapes: list[list[int]] + group_name: str + target_modules: list[str] | None = None + weight_update_mode: str | None = None + lora_alpha: int | None = None + lora_rank: int | None = None + + +@dataclass +class DestroyWeightsUpdateGroupReqInput: + group_name: str diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/weights_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/weights_api.py index 2e9a0956998e..216cf63b9209 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/weights_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/weights_api.py @@ -1,14 +1,18 @@ """Weight update API for the diffusion engine.""" +import msgspec from fastapi import APIRouter, Request from sglang.multimodal_gen.runtime.entrypoints.post_training.io_struct import ( + DestroyWeightsUpdateGroupReqInput, GetWeightsChecksumReqInput, + InitWeightsUpdateGroupReqInput, ReleaseMemoryOccupationReqInput, ResumeMemoryOccupationReqInput, UpdateWeightFromDiskReqInput, UpdateWeightFromTensorCheckerReqInput, UpdateWeightFromTensorReqInput, + UpdateWeightsFromDistributedReqInput, ) from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client from sglang.srt.utils.json_response import orjson_response @@ -201,3 +205,38 @@ async def resume_memory_occupation(): payload = response.output success = bool(payload["success"]) return orjson_response(payload, status_code=200 if success else 400) + + +async def _parse_and_forward(request: Request, request_type: type): + try: + req = msgspec.convert(await request.json(), type=request_type) + except (ValueError, TypeError) as e: + return orjson_response({"success": False, "message": str(e)}, status_code=400) + + try: + response = await async_scheduler_client.forward(req) + except Exception as e: + return orjson_response({"success": False, "message": str(e)}, status_code=500) + + if response.output is None: + return orjson_response( + {"success": False, "message": response.error}, status_code=500 + ) + + result = response.output + return orjson_response(result, status_code=200 if result["success"] else 400) + + +@router.post("/init_weights_update_group") +async def init_weights_update_group(request: Request): + return await _parse_and_forward(request, InitWeightsUpdateGroupReqInput) + + +@router.post("/update_weights_from_distributed") +async def update_weights_from_distributed(request: Request): + return await _parse_and_forward(request, UpdateWeightsFromDistributedReqInput) + + +@router.post("/destroy_weights_update_group") +async def destroy_weights_update_group(request: Request): + return await _parse_and_forward(request, DestroyWeightsUpdateGroupReqInput) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 43a7d5e0565f..d7113fe2ea9e 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -248,6 +248,7 @@ def __init__( self.server_args = server_args self.metrics = init_metrics(server_args, rank) self.pipeline: ComposedPipelineBase = None + self._weights_update_groups: dict[str, torch.distributed.ProcessGroup] = {} self.init_device_and_model() load_snapshot = None if current_platform.is_cpu() else capture_memory_snapshot() diff --git a/python/sglang/multimodal_gen/runtime/managers/scheduler.py b/python/sglang/multimodal_gen/runtime/managers/scheduler.py index b524a4f4f55b..487795444449 100644 --- a/python/sglang/multimodal_gen/runtime/managers/scheduler.py +++ b/python/sglang/multimodal_gen/runtime/managers/scheduler.py @@ -32,12 +32,15 @@ UnmergeLoraWeightsReq, ) from sglang.multimodal_gen.runtime.entrypoints.post_training.io_struct import ( + DestroyWeightsUpdateGroupReqInput, GetWeightsChecksumReqInput, + InitWeightsUpdateGroupReqInput, ReleaseMemoryOccupationReqInput, ResumeMemoryOccupationReqInput, UpdateWeightFromDiskReqInput, UpdateWeightFromTensorCheckerReqInput, UpdateWeightFromTensorReqInput, + UpdateWeightsFromDistributedReqInput, ) from sglang.multimodal_gen.runtime.ipc_array import ( is_local_endpoint, @@ -168,6 +171,9 @@ def __init__( self._handle_update_weights_from_tensor_checker ), GetWeightsChecksumReqInput: self._handle_get_weights_checksum, + InitWeightsUpdateGroupReqInput: self._handle_init_weights_update_group, + UpdateWeightsFromDistributedReqInput: self._handle_update_weights_from_distributed, + DestroyWeightsUpdateGroupReqInput: self._handle_destroy_weights_update_group, ReleaseMemoryOccupationReqInput: self._handle_release_memory_occupation, ResumeMemoryOccupationReqInput: self._handle_resume_memory_occupation, } diff --git a/python/sglang/multimodal_gen/runtime/post_training/gpu_worker_post_training_mixin.py b/python/sglang/multimodal_gen/runtime/post_training/gpu_worker_post_training_mixin.py index e93cb3ae7f19..ed0abef0677e 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/gpu_worker_post_training_mixin.py +++ b/python/sglang/multimodal_gen/runtime/post_training/gpu_worker_post_training_mixin.py @@ -1,7 +1,11 @@ from __future__ import annotations +from datetime import timedelta from typing import TYPE_CHECKING +import torch +import torch.distributed as dist + from sglang.multimodal_gen.runtime.distributed import ( get_tp_rank, get_tp_world_size, @@ -19,22 +23,92 @@ get_updatable_modules, ) from sglang.srt.platforms import current_platform -from sglang.srt.utils import MultiprocessingSerializer +from sglang.srt.utils import MultiprocessingSerializer, init_custom_process_group +from sglang.srt.utils.network import NetworkAddress from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions if TYPE_CHECKING: from sglang.multimodal_gen.runtime.entrypoints.post_training.io_struct import ( + DestroyWeightsUpdateGroupReqInput, + InitWeightsUpdateGroupReqInput, UpdateWeightFromTensorCheckerReqInput, UpdateWeightFromTensorReqInput, + UpdateWeightsFromDistributedReqInput, ) +_SENDER_RANK = 0 + + def _normalize_gpu_uuid(uuid: str) -> str: # NVML prefixes uuids with "GPU-"/"MIG-"; torch device properties do not. return uuid.removeprefix("MIG-").removeprefix("GPU-").lower() +def _receive_broadcast_weights( + req: UpdateWeightsFromDistributedReqInput, group: dist.ProcessGroup +) -> list[tuple[str, torch.Tensor]]: + named_tensors = [ + ( + name, + torch.empty( + shape, dtype=getattr(torch, dtype), device=torch.cuda.current_device() + ), + ) + for name, dtype, shape in zip(req.names, req.dtypes, req.shapes, strict=True) + ] + handles = [ + dist.broadcast(tensor, src=_SENDER_RANK, group=group, async_op=True) + for _, tensor in named_tensors + ] + for handle in handles: + handle.wait() + return named_tensors + + class GPUWorkerPostTrainingMixin: + def init_weights_update_group( + self, req: InitWeightsUpdateGroupReqInput + ) -> tuple[bool, str]: + if req.group_name in self._weights_update_groups: + return False, f"Group {req.group_name} already exists" + engine_world = get_world_group() + if req.rank_offset + engine_world.world_size > req.world_size: + return False, "Engine ranks exceed the update group size" + dist_timeout = self.server_args.dist_timeout + self._weights_update_groups[req.group_name] = init_custom_process_group( + backend=req.backend, + init_method=NetworkAddress(req.master_address, req.master_port).to_tcp(), + world_size=req.world_size, + rank=req.rank_offset + engine_world.rank_in_group, + group_name=req.group_name, + timeout=None if dist_timeout is None else timedelta(seconds=dist_timeout), + ) + return True, "Initialized weight update group" + + def update_weights_from_distributed( + self, req: UpdateWeightsFromDistributedReqInput + ) -> tuple[bool, str]: + group = self._weights_update_groups.get(req.group_name) + if group is None: + return False, f"Unknown weight update group {req.group_name}" + named_tensors = _receive_broadcast_weights(req, group) + return WeightsUpdater(self.pipeline).update_weights_from_tensor( + named_tensors=named_tensors, + target_modules=req.target_modules, + weight_update_mode=req.weight_update_mode, + lora_alpha=req.lora_alpha, + lora_rank=req.lora_rank, + ) + + def destroy_weights_update_group( + self, req: DestroyWeightsUpdateGroupReqInput + ) -> tuple[bool, str]: + group = self._weights_update_groups.pop(req.group_name, None) + if group is not None: + dist.destroy_process_group(group) + return True, "Destroyed weight update group" + def update_weights_from_disk( self, model_path: str, diff --git a/python/sglang/multimodal_gen/runtime/post_training/scheduler_post_training_mixin.py b/python/sglang/multimodal_gen/runtime/post_training/scheduler_post_training_mixin.py index 025643b09cc0..ad18bbf78630 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/scheduler_post_training_mixin.py +++ b/python/sglang/multimodal_gen/runtime/post_training/scheduler_post_training_mixin.py @@ -1,11 +1,51 @@ from __future__ import annotations -from typing import Any, List +from typing import Any, Callable, List +import torch.distributed as dist + +from sglang.multimodal_gen.runtime.distributed import get_world_group from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch class SchedulerPostTrainingMixin: + def _run_on_all_ranks( + self, worker_op: Callable[[Any], tuple[bool, str]], req: Any + ) -> OutputBatch: + """Every rank joins the all-gather, even one whose op raised, and the + request succeeds only if every rank succeeded.""" + try: + own_result = worker_op(req) + except Exception as e: + own_result = False, f"{type(e).__name__}: {e}" + + world = get_world_group() + all_results = [None] * world.world_size + dist.all_gather_object(all_results, own_result, group=world.cpu_group) + + failures = [ + f"Rank {rank}: {message}" + for rank, (success, message) in enumerate(all_results) + if not success + ] + if failures: + return OutputBatch( + output={"success": False, "message": failures[0]}, error=failures[0] + ) + _, message = own_result + return OutputBatch(output={"success": True, "message": message}) + + def _handle_init_weights_update_group(self, reqs: List[Any]) -> OutputBatch: + return self._run_on_all_ranks(self.worker.init_weights_update_group, reqs[0]) + + def _handle_update_weights_from_distributed(self, reqs: List[Any]) -> OutputBatch: + return self._run_on_all_ranks( + self.worker.update_weights_from_distributed, reqs[0] + ) + + def _handle_destroy_weights_update_group(self, reqs: List[Any]) -> OutputBatch: + return self._run_on_all_ranks(self.worker.destroy_weights_update_group, reqs[0]) + def _handle_update_weights_from_disk(self, reqs: List[Any]) -> OutputBatch: req = reqs[0] success, message = self.worker.update_weights_from_disk( diff --git a/python/sglang/multimodal_gen/runtime/scheduler_client.py b/python/sglang/multimodal_gen/runtime/scheduler_client.py index df6ab5a20576..6eba9998664f 100644 --- a/python/sglang/multimodal_gen/runtime/scheduler_client.py +++ b/python/sglang/multimodal_gen/runtime/scheduler_client.py @@ -1,3 +1,4 @@ +import asyncio import itertools import pickle import time @@ -20,12 +21,15 @@ UnmergeLoraWeightsReq, ) from sglang.multimodal_gen.runtime.entrypoints.post_training.io_struct import ( + DestroyWeightsUpdateGroupReqInput, GetWeightsChecksumReqInput, + InitWeightsUpdateGroupReqInput, ReleaseMemoryOccupationReqInput, ResumeMemoryOccupationReqInput, UpdateWeightFromDiskReqInput, UpdateWeightFromTensorCheckerReqInput, UpdateWeightFromTensorReqInput, + UpdateWeightsFromDistributedReqInput, ) from sglang.multimodal_gen.runtime.ipc_array import ( is_local_endpoint, @@ -44,6 +48,14 @@ logger = init_logger(__name__) +# All replicas join one process group for these ops, so they must reach every +# replica at once; a sequential fan-out would deadlock. +_COLLECTIVE_REQ_TYPES = ( + InitWeightsUpdateGroupReqInput, + UpdateWeightsFromDistributedReqInput, + DestroyWeightsUpdateGroupReqInput, +) + # Control ops mutate replica state (weights, LoRA, memory, shutdown), so with # DP they must reach every replica rather than one. _CONTROL_REQ_TYPES = ( @@ -270,7 +282,12 @@ async def forward(self, batch: Any, timeout_ms: int | None = None) -> Any: ) endpoints = self.server_args.scheduler_endpoints - if isinstance(batch, _CONTROL_REQ_TYPES): + if isinstance(batch, _COLLECTIVE_REQ_TYPES): + results = await asyncio.gather( + *(self._forward_one(ep, batch, timeout_ms) for ep in endpoints) + ) + output_batch = _merge_fanout_results(results) + elif isinstance(batch, _CONTROL_REQ_TYPES): # replica state (weights, LoRA, memory) must change everywhere results = [ await self._forward_one(ep, batch, timeout_ms) for ep in endpoints diff --git a/python/sglang/multimodal_gen/test/server/gpu_cases.py b/python/sglang/multimodal_gen/test/server/gpu_cases.py index 83c62dcd36fc..5e1008249bc2 100644 --- a/python/sglang/multimodal_gen/test/server/gpu_cases.py +++ b/python/sglang/multimodal_gen/test/server/gpu_cases.py @@ -1516,6 +1516,7 @@ def _discover_unit_tests() -> list[str]: "../single_test_file/test_disagg_server.py", "../single_test_file/test_ar_models.py", "../single_test_file/test_ipc_a2a_2_gpu.py", + "../single_test_file/test_update_weights_from_distributed.py", "../single_test_file/test_encoder_fold_srt_linear_2_gpu.py", "../single_test_file/test_encoder_fold_srt_2_gpu.py", "../single_test_file/test_diffusion_bcg_tp2_zimage_turbo.py", diff --git a/python/sglang/multimodal_gen/test/single_test_file/test_update_weights_from_distributed.py b/python/sglang/multimodal_gen/test/single_test_file/test_update_weights_from_distributed.py new file mode 100644 index 000000000000..4e286c4c403f --- /dev/null +++ b/python/sglang/multimodal_gen/test/single_test_file/test_update_weights_from_distributed.py @@ -0,0 +1,226 @@ +import base64 +import os +import pickle +import socket +import sys +import time +from concurrent.futures import ThreadPoolExecutor +from datetime import timedelta +from pathlib import Path + +import pytest +import requests +import torch +import torch.distributed as dist +from huggingface_hub import snapshot_download +from safetensors import safe_open + +from sglang.multimodal_gen.test.server.test_server_utils import ServerManager +from sglang.multimodal_gen.test.test_utils import get_dynamic_server_port +from sglang.srt.utils import init_custom_process_group + +GROUP_NAME = "test-update" +LORA_RANK = 4 +LORA_ALPHA = 4 + + +def post(base_url, endpoint, body): + response = requests.post(f"{base_url}/{endpoint}", json=body, timeout=180) + response.raise_for_status() + return response.json() + + +def run_while_posting(sender_side, base_url, endpoint, body): + with ThreadPoolExecutor(max_workers=1) as executor: + response = executor.submit(post, base_url, endpoint, body) + sender_result = sender_side() + assert response.result()["success"] + return sender_result + + +def free_port(): + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +def init_device_bound_default_group(device, rendezvous_file): + dist.init_process_group( + "nccl", + init_method=f"file://{rendezvous_file}", + rank=0, + world_size=1, + device_id=device, + ) + + +@pytest.fixture(scope="module") +def engine(tmp_path_factory): + num_gpus = int(os.environ.get("SGLANG_TEST_WEIGHT_ENGINE_GPUS", "1")) + if torch.cuda.device_count() <= num_gpus: + pytest.skip("Requires the engine GPUs plus a separate sender GPU") + model = os.environ.get( + "SGLANG_TEST_WEIGHT_MODEL", "black-forest-labs/FLUX.2-klein-base-4B" + ) + visible_gpus = os.environ.get( + "CUDA_VISIBLE_DEVICES", + ",".join(str(i) for i in range(torch.cuda.device_count())), + ).split(",") + engine_gpus = visible_gpus[:num_gpus] + port = get_dynamic_server_port() + server = ServerManager( + model=model, + port=port, + extra_args=f"--num-gpus {num_gpus} --enable-cfg-parallel false --warmup-mode off --dit-precision bf16 " + + os.environ.get("SGLANG_TEST_WEIGHT_SERVER_ARGS", ""), + env_vars={"CUDA_VISIBLE_DEVICES": ",".join(engine_gpus)}, + ).start() + sender_device = torch.device("cuda", len(engine_gpus)) + torch.cuda.set_device(sender_device) + init_device_bound_default_group( + sender_device, tmp_path_factory.mktemp("weights") / "rendezvous" + ) + try: + yield f"http://127.0.0.1:{port}", model, num_gpus + finally: + dist.destroy_process_group() + server.cleanup() + + +def init_update_group(base_url, num_gpus): + port = free_port() + world_size = 1 + num_gpus + + def join_as_sender(): + return init_custom_process_group( + backend="nccl", + init_method=f"tcp://127.0.0.1:{port}", + rank=0, + world_size=world_size, + group_name=GROUP_NAME, + timeout=timedelta(seconds=120), + ) + + return run_while_posting( + join_as_sender, + base_url, + "init_weights_update_group", + dict( + master_address="127.0.0.1", + master_port=port, + rank_offset=1, + world_size=world_size, + group_name=GROUP_NAME, + ), + ) + + +def destroy_update_group(base_url, group): + run_while_posting( + lambda: dist.destroy_process_group(group), + base_url, + "destroy_weights_update_group", + {"group_name": GROUP_NAME}, + ) + + +def update_options(mode): + return dict( + target_modules=["transformer"], + weight_update_mode=mode, + lora_alpha=LORA_ALPHA, + lora_rank=LORA_RANK, + ) + + +def update_weights_from_tensor(base_url, weights, mode): + payload = {"transformer": [(name, tensor.cpu()) for name, tensor in weights]} + response = post( + base_url, + "update_weights_from_tensor", + dict( + serialized_named_tensors=[base64.b64encode(pickle.dumps(payload)).decode()], + **update_options(mode), + ), + ) + assert response["success"] + + +def update_weights_from_distributed(base_url, group, weights, mode): + def broadcast(): + tensors = [tensor.contiguous() for _, tensor in weights] + handles = [ + dist.broadcast(tensor, src=0, group=group, async_op=True) + for tensor in tensors + ] + for handle in handles: + handle.wait() + + run_while_posting( + broadcast, + base_url, + "update_weights_from_distributed", + dict( + names=[name for name, _ in weights], + dtypes=[str(tensor.dtype).removeprefix("torch.") for _, tensor in weights], + shapes=[list(tensor.shape) for _, tensor in weights], + group_name=GROUP_NAME, + **update_options(mode), + ), + ) + + +def transformer_checksum(base_url): + return post(base_url, "get_weights_checksum", {"module_names": ["transformer"]}) + + +def load_attention_weight(model): + model_dir = Path(model) if Path(model).is_dir() else Path(snapshot_download(model)) + for path in sorted((model_dir / "transformer").glob("*.safetensors")): + with safe_open(str(path), framework="pt") as reader: + for name in reader.keys(): + if name.endswith(".attn.to_q.weight"): + return name, reader.get_tensor(name).cuda() + pytest.fail("No attention projection found") + + +def make_weights(name, original, mode): + if mode != "lora_merge": + return [(name, original + 0.01)] + prefix = name.removesuffix(".weight") + out_features, in_features = original.shape + lora_a = torch.randn(in_features, LORA_RANK, device="cuda").t() * 0.1 + lora_b = torch.randn(LORA_RANK, out_features, device="cuda").t() * 0.1 + return [(f"{prefix}.lora_A", lora_a), (f"{prefix}.lora_B", lora_b)] + + +@pytest.mark.parametrize("mode", [None, "lora_merge"]) +def test_distributed_matches_tensor_updates_and_reconnects(engine, mode): + base_url, model, num_gpus = engine + name, original = load_attention_weight(model) + torch.manual_seed(123) + weights = make_weights(name, original, mode) + other_weights = [(key, tensor + 0.02) for key, tensor in weights] + + update_weights_from_tensor(base_url, weights, mode) + expected = transformer_checksum(base_url) + + group = init_update_group(base_url, num_gpus) + durations = [] + try: + update_weights_from_distributed(base_url, group, other_weights, mode) + assert transformer_checksum(base_url) != expected + for _ in range(2): + start = time.perf_counter() + update_weights_from_distributed(base_url, group, weights, mode) + durations.append(time.perf_counter() - start) + assert transformer_checksum(base_url) == expected + finally: + destroy_update_group(base_url, group) + print( + f"weight_update mode={mode} bytes={sum(t.numel() * t.element_size() for _, t in weights)} seconds={durations}" + ) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v", "-s"])) diff --git a/python/sglang/multimodal_gen/test/unit/test_distributed_weight_update.py b/python/sglang/multimodal_gen/test/unit/test_distributed_weight_update.py new file mode 100644 index 000000000000..4f5f63d5aaa2 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_distributed_weight_update.py @@ -0,0 +1,82 @@ +import asyncio +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from sglang.multimodal_gen.runtime.entrypoints.post_training.io_struct import ( + InitWeightsUpdateGroupReqInput, +) +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch +from sglang.multimodal_gen.runtime.post_training import scheduler_post_training_mixin +from sglang.multimodal_gen.runtime.scheduler_client import AsyncSchedulerClient + + +def test_collective_request_reaches_dp_replicas_concurrently(): + async def run(): + client = AsyncSchedulerClient() + client.context = object() + client.request_logger = MagicMock() + client.server_args = SimpleNamespace(scheduler_endpoints=["a", "b"]) + entered = [] + ready = asyncio.Event() + + async def forward(endpoint, batch, timeout): + entered.append(endpoint) + if len(entered) == 2: + ready.set() + await ready.wait() + return OutputBatch( + output={"success": endpoint == "a"}, + error="rank failure" if endpoint == "b" else None, + ) + + client._forward_one = forward + req = InitWeightsUpdateGroupReqInput( + master_address="127.0.0.1", + master_port=12345, + rank_offset=1, + world_size=3, + group_name="update", + ) + result = await asyncio.wait_for(client.forward(req), timeout=2) + assert set(entered) == {"a", "b"} + assert result.error == "rank failure" + + asyncio.run(run()) + + +@pytest.fixture +def two_rank_scheduler(monkeypatch): + module = scheduler_post_training_mixin + monkeypatch.setattr( + module, + "get_world_group", + lambda: SimpleNamespace(world_size=2, cpu_group="world"), + ) + gathered = [] + + def gather(results, result, group): + gathered.append(result) + results[:] = [result, (False, "load failed")] + + monkeypatch.setattr(module.dist, "all_gather_object", gather) + return module.SchedulerPostTrainingMixin(), gathered + + +def test_non_primary_worker_failure_reaches_http_result(two_rank_scheduler): + scheduler, _ = two_rank_scheduler + result = scheduler._run_on_all_ranks(lambda req: (True, "loaded"), None) + assert not result.output["success"] + assert result.error == "Rank 1: load failed" + + +def test_raising_rank_still_joins_the_gather(two_rank_scheduler): + scheduler, gathered = two_rank_scheduler + + def op(req): + raise RuntimeError("out of memory") + + result = scheduler._run_on_all_ranks(op, None) + assert gathered == [(False, "RuntimeError: out of memory")] + assert result.error == "Rank 0: RuntimeError: out of memory" diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 4bc416920edc..927d1003c0ed 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -2791,6 +2791,25 @@ def _get_fastapi_request_path(request) -> Tuple[str, bool]: return request.url.path, False +@contextmanager +def _no_split_from_default_group(): + """torch splits every new group from a device-bound default group. A group + that spans processes outside the default world must not be split, or each + process silently ends up in a group of its own.""" + from torch.distributed.distributed_c10d import _world + + default_pg = _world.default_pg + bound_device_id = None if default_pg is None else default_pg.bound_device_id + if bound_device_id is None: + yield + return + default_pg.bound_device_id = None + try: + yield + finally: + default_pg.bound_device_id = bound_device_id + + # Copy from pytorch and OpenRLHF to allow creating multiple main groups. # https://github.com/pytorch/pytorch/blob/main/torch/distributed/distributed_c10d.py # https://github.com/OpenRLHF/OpenRLHF/blob/main/openrlhf/utils/distributed_util.py @@ -2848,17 +2867,18 @@ def init_custom_process_group( pg_options_param_name = ( "backend_options" if torch_release >= (2, 6) else "pg_options" ) - pg, _ = _new_process_group_helper( - world_size, - rank, - [], - backend, - store, - group_name=group_name, - **{pg_options_param_name: pg_options}, - timeout=timeout, - device_id=device_id, - ) + with _no_split_from_default_group(): + pg, _ = _new_process_group_helper( + world_size, + rank, + [], + backend, + store, + group_name=group_name, + **{pg_options_param_name: pg_options}, + timeout=timeout, + device_id=device_id, + ) _world.pg_group_ranks[pg] = {i: i for i in range(world_size)} From e5117d3df3816f648d3b6b6bf687102327a66de1 Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Mon, 12 Oct 2026 16:55:38 -0700 Subject: [PATCH 42/56] [sglang-miles] Support P2P replica loading and target/draft weight updates (#42820) --- .../engine_info_bootstrap_server.py | 48 +-- python/sglang/srt/entrypoints/http_server.py | 26 +- .../sglang/srt/layers/quantization/unquant.py | 10 + python/sglang/srt/managers/scheduler.py | 5 + .../sglang/srt/model_executor/model_runner.py | 6 +- .../remote_instance_weight_transporter.py | 16 +- .../model_runner_components/weight_updater.py | 29 ++ python/sglang/srt/model_loader/loader.py | 50 ++- .../srt/model_loader/parameter_mapper.py | 261 ---------------- .../remote_instance_weight_loader_utils.py | 8 +- .../deepseek_common/deepseek_weight_loader.py | 66 ++-- python/sglang/srt/models/deepseek_nextn.py | 1 - python/sglang/srt/models/deepseek_v2.py | 57 ---- python/sglang/srt/models/glm4.py | 20 +- python/sglang/srt/models/glm4_moe.py | 66 ++-- python/sglang/srt/models/glm4_moe_lite.py | 43 --- python/sglang/srt/models/glm4_moe_nextn.py | 23 +- python/sglang/srt/models/llama.py | 28 +- python/sglang/srt/models/qwen2.py | 20 +- python/sglang/srt/models/qwen3.py | 32 +- python/sglang/srt/models/qwen3_moe.py | 34 +- .../quantization/test_unquant_moe_restore.py | 72 +++++ .../test_remote_instance_weight_roles.py | 146 +++++++++ .../test_weight_updater_published_storage.py | 100 ++++++ .../model_loader/test_default_model_loader.py | 61 ++++ .../test_remote_instance_loader.py | 43 +++ test/srt/models/test_params_mapping.py | 292 ------------------ 27 files changed, 726 insertions(+), 837 deletions(-) delete mode 100644 python/sglang/srt/model_loader/parameter_mapper.py create mode 100644 test/registered/unit/layers/quantization/test_unquant_moe_restore.py create mode 100644 test/registered/unit/model_executor/model_runner_components/test_remote_instance_weight_roles.py create mode 100644 test/registered/unit/model_executor/model_runner_components/test_weight_updater_published_storage.py delete mode 100644 test/srt/models/test_params_mapping.py diff --git a/python/sglang/srt/entrypoints/engine_info_bootstrap_server.py b/python/sglang/srt/entrypoints/engine_info_bootstrap_server.py index 88e48075a644..96de6cf6ec70 100644 --- a/python/sglang/srt/entrypoints/engine_info_bootstrap_server.py +++ b/python/sglang/srt/entrypoints/engine_info_bootstrap_server.py @@ -32,17 +32,19 @@ class EngineInfoBootstrapServer: query via HTTP GET. Currently supports transfer engine memory registration info and - per-rank parallelism configuration. + per-rank parallelism configuration. Each model runner of a rank registers + under its weight-update role (``target``, or a draft role such as ``draft``), + so a speculative draft never replaces its target's entry. """ def __init__(self, host: str, port: int): self.host = host self.port = port - # 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] = {} + # Storage: {(role, tp_rank): (session_id, weights_info_dict)} + self.transfer_engine_info: Dict[Tuple[str, int], Tuple] = {} + # Storage: {(role, tp_rank): parallelism_config_dict} + self.parallelism_config: Dict[Tuple[str, int], dict] = {} self.lock = threading.Lock() app = FastAPI() @@ -54,19 +56,20 @@ def health(): @app.put("/register_transfer_engine_info") def register_transfer_engine_info(data: dict): try: + role = data["role"] tp_rank = data["tp_rank"] info = data["transfer_engine_info"] session_id = info["session_id"] weights_info_dict = info["weights_info_dict"] with self.lock: - self.transfer_engine_info[tp_rank] = ( + self.transfer_engine_info[(role, tp_rank)] = ( session_id, weights_info_dict, ) logger.info( - f"Registered transfer engine info for tp_rank={tp_rank}, " + f"Registered transfer engine info for role={role} tp_rank={tp_rank}, " f"session_id={session_id}" ) return PlainTextResponse("OK") @@ -75,17 +78,17 @@ def register_transfer_engine_info(data: dict): raise HTTPException(status_code=400, detail=str(e)) @app.get("/get_transfer_engine_info") - def get_transfer_engine_info(rank: int): + def get_transfer_engine_info(rank: int, role: str = "target"): if rank < 0: raise HTTPException(status_code=400, detail="Invalid rank parameter") with self.lock: - info = self.transfer_engine_info.get(rank) + info = self.transfer_engine_info.get((role, rank)) if info is None: raise HTTPException( status_code=404, - detail=f"No transfer engine info for rank {rank}", + detail=f"No transfer engine info for role {role} rank {rank}", ) return {"rank": rank, "remote_instance_transfer_engine_info": list(info)} @@ -96,30 +99,33 @@ def get_transfer_engine_info(rank: int): @app.put("/register_parallelism_config") def register_parallelism_config(data: dict): try: + role = data["role"] tp_rank = data["tp_rank"] config = data["parallelism_config"] with self.lock: - self.parallelism_config[tp_rank] = config + self.parallelism_config[(role, tp_rank)] = config - logger.info(f"Registered parallelism config for tp_rank={tp_rank}") + logger.info( + f"Registered parallelism config for role={role} 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): + def get_parallelism_config(rank: int, role: str = "target"): if rank < 0: raise HTTPException(status_code=400, detail="Invalid rank parameter") with self.lock: - config = self.parallelism_config.get(rank) + config = self.parallelism_config.get((role, rank)) if config is None: raise HTTPException( status_code=404, - detail=f"No parallelism config for rank {rank}", + detail=f"No parallelism config for role {role} rank {rank}", ) return config @@ -135,10 +141,14 @@ def close(self): self._server.should_exit = True self._thread.join(timeout=5) - def get_transfer_engine_info(self, rank: int) -> Optional[Tuple]: + def get_transfer_engine_info( + self, rank: int, role: str = "target" + ) -> Optional[Tuple]: """Direct in-process access for co-located HTTP server (no HTTP round-trip).""" - return self.transfer_engine_info.get(rank) + return self.transfer_engine_info.get((role, rank)) - def get_parallelism_config_info(self, rank: int) -> Optional[dict]: + def get_parallelism_config_info( + self, rank: int, role: str = "target" + ) -> Optional[dict]: """Direct in-process access for parallelism config (no HTTP round-trip).""" - return self.parallelism_config.get(rank) + return self.parallelism_config.get((role, rank)) diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index bb76f303d128..85412a1b644a 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -1321,18 +1321,20 @@ async def send_weights_to_remote_instance( @app.get("/get_remote_instance_transfer_engine_info") @auth_level(AuthLevel.ADMIN_OPTIONAL) -async def get_remote_instance_transfer_engine_info(rank: int = None): +async def get_remote_instance_transfer_engine_info( + rank: int = None, role: str = "target" +): """Get the server information (deprecated - use /remote_instance_transfer_engine_info instead).""" logger.warning( "Endpoint '/get_remote_instance_transfer_engine_info' is deprecated and will be removed in a future version. " "Please use '/remote_instance_transfer_engine_info' instead." ) - return await remote_instance_transfer_engine_info(rank=rank) + return await remote_instance_transfer_engine_info(rank=rank, role=role) @app.get("/remote_instance_transfer_engine_info") @auth_level(AuthLevel.ADMIN_OPTIONAL) -async def remote_instance_transfer_engine_info(rank: int = None): +async def remote_instance_transfer_engine_info(rank: int = None, role: str = "target"): if rank is None or rank < 0: return ORJSONResponse( {"error": {"message": "Missing or invalid rank parameter"}}, @@ -1343,24 +1345,30 @@ async def remote_instance_transfer_engine_info(rank: int = None): try: resp = requests.get( f"{server_args.engine_info_bootstrap_url}/get_transfer_engine_info", - params={"rank": rank}, + params={"rank": rank, "role": role}, timeout=5, ) if resp.status_code == 200: return resp.json() except (requests.exceptions.RequestException, ValueError) as e: - logger.warning(f"Failed to get transfer engine info for rank {rank}: {e}") + logger.warning( + f"Failed to get transfer engine info for role {role} rank {rank}: {e}" + ) return ORJSONResponse( - {"error": {"message": f"Failed to get transfer engine info for rank {rank}"}}, + { + "error": { + "message": f"Failed to get transfer engine info for role {role} rank {rank}" + } + }, status_code=HTTPStatus.BAD_REQUEST, ) @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.""" +async def parallelism_config(rank: int = None, role: str = "target"): + """Get per-rank parallelism config of the `role` runner (target or a draft) from the bootstrap server.""" if rank is None or rank < 0: return ORJSONResponse( {"error": {"message": "Missing or invalid rank parameter"}}, @@ -1371,7 +1379,7 @@ async def parallelism_config(rank: int = None): try: resp = requests.get( f"{server_args.engine_info_bootstrap_url}/get_parallelism_config", - params={"rank": rank}, + params={"rank": rank, "role": role}, timeout=5, ) if resp.status_code == 200: diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index 31f63142197f..02d4bf67d9e2 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/python/sglang/srt/layers/quantization/unquant.py @@ -906,6 +906,16 @@ def _maybe_interleave_w13_for_fused_swiglu(self, layer: torch.nn.Module) -> None "Interleaved w13 gate/up: the SwiGLU is applied by the MoE up-GEMM epilogue." ) + def restore_weights_before_loading(self, layer: torch.nn.Module) -> None: + """Return the expert weights to the canonical layout loads expect, which the TRT-LLM BF16 postprocess + left in block layout. A no-op on other MoE backends.""" + for weight_name in ("w13_weight", "w2_weight"): + self.maybe_restore_flashinfer_trtllm_bf16_weight_shape_for_load( + layer=layer, + param=getattr(layer, weight_name), + weight_name=f".experts.{weight_name}", + ) + def maybe_restore_flashinfer_trtllm_bf16_weight_shape_for_load( self, layer: torch.nn.Module, diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 3eafec682506..8bd6402da8e0 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1076,6 +1076,11 @@ def init_memory_pools(self): token_to_kv_pool_allocator=allocator, ) self.draft_worker.init_hicache_draft_plan() + # only now: allocating the draft's pools binds the target's embed and head into it + for role, runner in self.draft_worker.weight_update_runners(): + runner.remote_instance_weight_transporter.maybe_register_and_publish_weight_info( + role=role + ) def init_all_attention_backends(self): """Initialize attention backends for all workers.""" diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 939314f13469..9f7f8636d226 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -666,7 +666,11 @@ def initialize(self): if self.is_draft_worker: disable_routed_experts_capture_for_draft(self.model) self.maybe_init_expert_backup_client() - self.remote_instance_weight_transporter.maybe_register_and_publish_weight_info() + # drafts publish once the spec worker binds the target's embed and head (Scheduler.init_memory_pools) + if not self.is_draft_worker: + self.remote_instance_weight_transporter.maybe_register_and_publish_weight_info( + role="target" + ) self.layer_info: ModelLayerInfo = resolve_layer_indices( model=self.model, model_config=self.model_config, 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 f41c26622310..46e3f721ab79 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 @@ -66,7 +66,9 @@ def maybe_init_parallelism_config(self) -> None: self.tp_rank ) - def maybe_register_and_publish_weight_info(self) -> None: + def maybe_register_and_publish_weight_info(self, role: str) -> None: + """Registers this runner's weights with its transfer engine and publishes them, with its parallelism + config, under `role`: ``target``, or the draft role `weight_update_runners()` names.""" if ( remote_instance_transfer_engine_enabled() # ModelExpress owns TransferEngine memory registration and metadata @@ -79,7 +81,7 @@ def maybe_register_and_publish_weight_info(self) -> None: ): # Register memory and upstream the transfer engine info to the bootstrap server self.weight_info = register_memory_region(self.model, self.engine) - self._register_to_engine_info_bootstrap() + self._register_to_engine_info_bootstrap(role) # The P2P weight-update client needs each rank's parallelism layout to # map training-side parameters onto this rank's shards. @@ -87,7 +89,7 @@ def maybe_register_and_publish_weight_info(self) -> None: self.server_args.registers_parallelism_config() and self.parallelism_config is not None ): - self._register_parallelism_config_to_bootstrap() + self._register_parallelism_config_to_bootstrap(role) def _bootstrap_url(self) -> str: if get_parallel().dist_init_addr: @@ -99,13 +101,14 @@ def _bootstrap_url(self) -> str: bootstrap_port = get_model().engine_info_bootstrap_port return NetworkAddress(bootstrap_host, bootstrap_port).to_url() - def _register_parallelism_config_to_bootstrap(self) -> None: + def _register_parallelism_config_to_bootstrap(self, role: str) -> 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 = { + "role": role, "tp_rank": self.tp_rank, "parallelism_config": self.parallelism_config.to_dict(), } @@ -126,7 +129,9 @@ def _register_parallelism_config_to_bootstrap(self) -> None: f"Failed to register parallelism config for tp_rank={self.tp_rank}: {e}" ) - def _register_to_engine_info_bootstrap(self: RemoteInstanceWeightTransporter): + def _register_to_engine_info_bootstrap( + self: RemoteInstanceWeightTransporter, role: str + ): """Register transfer engine info with the EngineInfoBootstrapServer via HTTP PUT. The bootstrap server runs on node_rank==0. For multi-node setups, the @@ -148,6 +153,7 @@ def _register_to_engine_info_bootstrap(self: RemoteInstanceWeightTransporter): url = f"{bootstrap_na.to_url()}/register_transfer_engine_info" payload = { + "role": role, "tp_rank": self.tp_rank, "transfer_engine_info": { "session_id": self.session_id, 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 0230690544dc..267f2c7e665c 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 @@ -301,6 +301,7 @@ def begin_weight_update(self: WeightUpdater) -> None: DefaultModelLoader.restore_weights_before_loading( self.get_model(), torch.device(self.device) ) + self._assert_published_storage_unmoved() def end_weight_update(self: WeightUpdater, *, run_post_load: bool) -> None: if run_post_load: @@ -308,6 +309,34 @@ def end_weight_update(self: WeightUpdater, *, run_post_load: bool) -> None: DefaultModelLoader.postprocess_weights( self.get_model(), torch.device(self.device) ) + self._assert_published_storage_unmoved() + + def _assert_published_storage_unmoved(self: WeightUpdater) -> None: + """Peers write weight updates straight into the addresses this rank published at startup (p2p), so a + weight-update session must leave every published parameter's storage where it was.""" + published_weight_info = ( + self.get_model_runner().remote_instance_weight_transporter.weight_info + ) + # nothing published: nobody writes by address + if published_weight_info is None: + return + params = dict(self.get_model().named_parameters()) + moved = [] + for name, (address, numel, element_size) in published_weight_info.items(): + param = params.get(name) + now = ( + None + if param is None + else (param.data_ptr(), param.numel() * param.element_size()) + ) + if now != (address, numel * element_size): + moved.append( + f"{name} published at {(address, numel * element_size)}, now {now}" + ) + assert not moved, ( + "a weight-update session moved published parameters, so peer writes would land in freed memory: " + f"{', '.join(moved[:5])} ({len(moved)} in all)" + ) def load_weights_from_distributed( self: WeightUpdater, named_tensors: List[Tuple[str, torch.Tensor]] diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 658121ffa3c3..cd13ea193a84 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -838,6 +838,33 @@ def initialize_model_for_startup( ) return model + def initialize_model_without_storage( + self, + *, + model_config: ModelConfig, + device: torch.device, + ) -> Tuple[nn.Module, Dict[str, torch.Size]]: + """Build the model of `initialize_model_for_startup` without allocating parameter storage. + + Returns the model, with every parameter 0-size on `device`, and the shape each parameter was built with, + by name. Parameters keep their subclass, attributes and ties; other tensors stay on meta, so the model + cannot run forward. For a trainer that binds its own buffers before loading weights in this rank's layout + (p2p weight updates). + """ + quant_config = _get_quantization_config(model_config, self.load_config) + with set_default_torch_dtype(model_config.dtype): + with torch.device("meta"): + model = _initialize_model( + model_config, + self.load_config, + quant_config, + ) + built_shapes_by_name = { + name: param.shape for name, param in model.named_parameters() + } + _replace_parameters_with_empty(model, device) + return model, built_shapes_by_name + def prepare_model_for_capture( self, *, @@ -1101,6 +1128,24 @@ def restore_weights_before_loading(model, target_device): quant_method.restore_weights_before_loading(module) +def _replace_parameters_with_empty(model: nn.Module, device: torch.device) -> None: + # by the original's id, so a parameter shared by several modules stays shared + replacements: Dict[int, nn.Parameter] = {} + for module in model.modules(): + for name, param in list(module._parameters.items()): + if param is None: + continue + if id(param) not in replacements: + replacement = torch.Tensor._make_subclass( + type(param), + torch.empty(0, dtype=param.dtype, device=device), + param.requires_grad, + ) + replacement.__dict__.update(param.__dict__) + replacements[id(param)] = replacement + module._parameters[name] = replacements[id(param)] + + class LayeredModelLoader(DefaultModelLoader): """Model loader that loads weights layer by layer so that one can quantize a layer before loading another to make the peak memory envelope smaller.""" @@ -3393,6 +3438,7 @@ def load_model( load_config.remote_instance_weight_loader_transfer_engine, f"http://{load_config.remote_instance_weight_loader_seed_instance_ip}:{load_config.remote_instance_weight_loader_seed_instance_service_port}", load_config.tp_rank, + role="draft" if model_config.is_draft_model else "target", ) if not success: raise RuntimeError( @@ -3471,11 +3517,11 @@ def load_model_from_remote_instance_by_nccl( current_platform.empty_cache() def load_model_from_remote_instance_by_transfer_engine( - self, model, transfer_engine, seed_url, tp_rank + self, model, transfer_engine, seed_url, tp_rank, role ) -> bool: # get remote weights metadata from source instance seed_transfer_engine_session_id, seed_transfer_engine_weight_info = ( - get_remote_instance_transfer_engine_info_per_rank(seed_url, tp_rank) + get_remote_instance_transfer_engine_info_per_rank(seed_url, tp_rank, role) ) if ( seed_transfer_engine_session_id is None diff --git a/python/sglang/srt/model_loader/parameter_mapper.py b/python/sglang/srt/model_loader/parameter_mapper.py deleted file mode 100644 index 56056d570843..000000000000 --- a/python/sglang/srt/model_loader/parameter_mapper.py +++ /dev/null @@ -1,261 +0,0 @@ -# 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/model_loader/remote_instance_weight_loader_utils.py b/python/sglang/srt/model_loader/remote_instance_weight_loader_utils.py index 2a0aeb047ed6..5b5c7f9c6985 100644 --- a/python/sglang/srt/model_loader/remote_instance_weight_loader_utils.py +++ b/python/sglang/srt/model_loader/remote_instance_weight_loader_utils.py @@ -79,13 +79,13 @@ def trigger_transferring_weights_request( raise -def get_remote_instance_transfer_engine_info_per_rank(seed_url: str, rank: int): +def get_remote_instance_transfer_engine_info_per_rank( + seed_url: str, rank: int, role: str +): try: response = requests.get( f"{seed_url}/get_remote_instance_transfer_engine_info", - params={ - "rank": rank, - }, + params={"rank": rank, "role": role}, ) if response.status_code == 200: 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 5d0003272306..774f325b7c75 100644 --- a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py +++ b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py @@ -25,6 +25,7 @@ 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, @@ -58,6 +59,7 @@ _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 @@ -208,18 +210,6 @@ 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, @@ -238,7 +228,33 @@ def do_load_weights( weights, NVFP4_CKPT_FP8_ATTN_QUANT_MODULES, nextn_conf ) - cached_a_proj = {} if self.fuse_qkv_a_proj else None + 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 pending_indexer_wk: Dict[str, Dict[str, torch.Tensor]] = {} @@ -266,7 +282,11 @@ def do_load_weights( ) ): continue - name = self.mutate_weight_preload(name) + 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}", + ) weight_names.append(name) @@ -323,7 +343,7 @@ def do_load_weights( ): continue - for param_name, weight_name, shard_id in self.stacked_params_mapping: + for param_name, weight_name, shard_id in stacked_params_mapping: # Skip non-stacked layers and experts (experts handled below). if weight_name not in name: continue @@ -337,10 +357,6 @@ 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: @@ -356,7 +372,7 @@ def do_load_weights( ) break else: - for mapping in self.expert_params_mapping: + for mapping in expert_params_mapping: param_name, weight_name, expert_id, shard_id = mapping if weight_name not in name: continue @@ -393,7 +409,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 self.fuse_qkv_a_proj and ( + if 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( @@ -463,7 +479,13 @@ def do_load_weights( if ( "k_scale" in name or "v_scale" in name ) and name not in params_dict: - name = self.custom_scale_remap(name) + # 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 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 d5f89429b3a0..7bc345ac2ca1 100644 --- a/python/sglang/srt/models/deepseek_nextn.py +++ b/python/sglang/srt/models/deepseek_nextn.py @@ -282,7 +282,6 @@ def __init__( # if not set, model load will be broken in DeepseekV3ForCausalLM load_weights() self.pp_group = get_parallel().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 d1d50ed9b3b7..eeb76d58cc09 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -3227,8 +3227,6 @@ 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() @@ -3241,22 +3239,6 @@ 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 @@ -3266,45 +3248,6 @@ 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/glm4.py b/python/sglang/srt/models/glm4.py index 2e654b37343e..f1feda9e9ec3 100644 --- a/python/sglang/srt/models/glm4.py +++ b/python/sglang/srt/models/glm4.py @@ -455,17 +455,6 @@ 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 @@ -560,7 +549,14 @@ def end_layer(self): return self.model.end_layer def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): - stacked_params_mapping = self.stacked_params_mapping + 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), + ] 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 14f74e6b844b..e2a2875f80a1 100644 --- a/python/sglang/srt/models/glm4_moe.py +++ b/python/sglang/srt/models/glm4_moe.py @@ -1157,34 +1157,9 @@ 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 @@ -1265,8 +1240,43 @@ def load_weights( else: raise ValueError("num_nextn_predict_layers is not in the config") - stacked_params_mapping = self.stacked_params_mapping - expert_params_mapping = self.expert_params_mapping + 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, + ) if is_nextn: nextn_layer_prefix = f"model.layers.{nextn_layer_id}" @@ -1287,8 +1297,6 @@ def load_weights( 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 78a683c6afc8..853115b7c391 100644 --- a/python/sglang/srt/models/glm4_moe_lite.py +++ b/python/sglang/srt/models/glm4_moe_lite.py @@ -877,33 +877,6 @@ 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 @@ -1016,22 +989,6 @@ 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 b8ceb8c7e786..83750d24328d 100644 --- a/python/sglang/srt/models/glm4_moe_nextn.py +++ b/python/sglang/srt/models/glm4_moe_nextn.py @@ -26,7 +26,6 @@ from sglang.srt.layers.layer_boundary.residual import batch as residual_batch 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, @@ -34,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_exec, get_parallel, get_spec +from sglang.srt.runtime_context import get_parallel, get_spec from sglang.srt.utils import add_prefix, is_npu logger = logging.getLogger(__name__) @@ -148,26 +147,6 @@ 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/llama.py b/python/sglang/srt/models/llama.py index 0f2b2e8026db..0da5cb20a5fe 100644 --- a/python/sglang/srt/models/llama.py +++ b/python/sglang/srt/models/llama.py @@ -546,21 +546,9 @@ 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, @@ -678,11 +666,23 @@ 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: - name = self.custom_scale_remap(name) + 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") + layer_id = get_layer_id(name) if ( layer_id is not None @@ -709,7 +709,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 self.stacked_params_mapping: + for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue name = name.replace(weight_name, param_name) diff --git a/python/sglang/srt/models/qwen2.py b/python/sglang/srt/models/qwen2.py index 7eb7f4d00e0c..5a37f29672d8 100644 --- a/python/sglang/srt/models/qwen2.py +++ b/python/sglang/srt/models/qwen2.py @@ -518,17 +518,6 @@ 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 @@ -630,7 +619,14 @@ 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 = self.stacked_params_mapping + 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: diff --git a/python/sglang/srt/models/qwen3.py b/python/sglang/srt/models/qwen3.py index ee00b8d9d7d5..cf2bc13901b6 100644 --- a/python/sglang/srt/models/qwen3.py +++ b/python/sglang/srt/models/qwen3.py @@ -43,7 +43,6 @@ 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 @@ -334,7 +333,16 @@ def __init__( ) -> None: super().__init__() self.hidden_size = config.hidden_size - rope_theta, rope_scaling = get_rope_config(config) + 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) max_position_embeddings = getattr(config, "max_position_embeddings", 32768) head_dim = getattr(config, "head_dim", None) self.self_attn = Qwen3Attention( @@ -519,16 +527,6 @@ 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: @@ -638,7 +636,15 @@ def end_layer(self): return self.model.end_layer def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): - stacked_params_mapping = self.stacked_params_mapping + 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 not name.startswith("model.") and ( diff --git a/python/sglang/srt/models/qwen3_moe.py b/python/sglang/srt/models/qwen3_moe.py index eedee0906cd1..e00e9363e3d9 100644 --- a/python/sglang/srt/models/qwen3_moe.py +++ b/python/sglang/srt/models/qwen3_moe.py @@ -1045,23 +1045,6 @@ 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( @@ -1200,8 +1183,21 @@ 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 = self.stacked_params_mapping - expert_params_mapping = self.expert_params_mapping + 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, + ) # Pre-define `params_dict` to avoid repeated expensive traversal of model parameters. params_dict = dict(self.named_parameters()) diff --git a/test/registered/unit/layers/quantization/test_unquant_moe_restore.py b/test/registered/unit/layers/quantization/test_unquant_moe_restore.py new file mode 100644 index 000000000000..5bde5508326d --- /dev/null +++ b/test/registered/unit/layers/quantization/test_unquant_moe_restore.py @@ -0,0 +1,72 @@ +"""Restoring a TRT-LLM BF16 MoE before a weight update must hand loads the canonical expert layout, in place.""" + +import unittest +from types import SimpleNamespace + +import torch + +from sglang.srt.layers.quantization.unquant import UnquantizedFusedMoEMethod +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") + +EXPERTS, HIDDEN, INTERMEDIATE = 2, 4, 3 + + +def _layer_in_block_layout(): + # the TRT-LLM postprocess keeps each expert's bytes in another shape + def expert_weight(numel): + return torch.nn.Parameter( + torch.arange(numel, dtype=torch.bfloat16).reshape(EXPERTS, -1, 2), + requires_grad=False, + ) + + return SimpleNamespace( + num_local_experts=EXPERTS, + hidden_size=HIDDEN, + intermediate_size_per_partition=INTERMEDIATE, + moe_runner_config=SimpleNamespace(is_gated=True), + w13_weight=expert_weight(EXPERTS * 2 * INTERMEDIATE * HIDDEN), + w2_weight=expert_weight(EXPERTS * HIDDEN * INTERMEDIATE), + ) + + +def _moe_method(use_flashinfer_trtllm_moe): + method = object.__new__(UnquantizedFusedMoEMethod) + method.use_flashinfer_trtllm_moe = use_flashinfer_trtllm_moe + return method + + +class TestRestoreWeightsBeforeLoading(CustomTestCase): + def test_trtllm_expert_weights_return_to_the_canonical_layout_in_place(self): + """Loads copy canonical expert tensors into storage peers write by address, which must not move.""" + layer = _layer_in_block_layout() + data_ptrs = (layer.w13_weight.data_ptr(), layer.w2_weight.data_ptr()) + + _moe_method(use_flashinfer_trtllm_moe=True).restore_weights_before_loading( + layer + ) + + self.assertEqual( + tuple(layer.w13_weight.shape), (EXPERTS, 2 * INTERMEDIATE, HIDDEN) + ) + self.assertEqual(tuple(layer.w2_weight.shape), (EXPERTS, HIDDEN, INTERMEDIATE)) + self.assertEqual( + (layer.w13_weight.data_ptr(), layer.w2_weight.data_ptr()), data_ptrs + ) + + def test_other_moe_backends_are_left_as_they_are(self): + """Only the TRT-LLM BF16 postprocess changes expert shapes.""" + layer = _layer_in_block_layout() + shapes_before = (layer.w13_weight.shape, layer.w2_weight.shape) + + _moe_method(use_flashinfer_trtllm_moe=False).restore_weights_before_loading( + layer + ) + + self.assertEqual((layer.w13_weight.shape, layer.w2_weight.shape), shapes_before) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/model_executor/model_runner_components/test_remote_instance_weight_roles.py b/test/registered/unit/model_executor/model_runner_components/test_remote_instance_weight_roles.py new file mode 100644 index 000000000000..599bd2d99c3e --- /dev/null +++ b/test/registered/unit/model_executor/model_runner_components/test_remote_instance_weight_roles.py @@ -0,0 +1,146 @@ +"""A speculative draft publishes its weights under its own role, after the target's embed and head are bound.""" + +import socket +import time +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import requests +import torch + +import sglang.srt.managers.scheduler as scheduler_mod +import sglang.srt.model_executor.model_runner_components.remote_instance_weight_transporter as transporter_mod +from sglang.srt.entrypoints.engine_info_bootstrap_server import ( + EngineInfoBootstrapServer, +) +from sglang.srt.managers.scheduler import Scheduler +from sglang.srt.model_executor.model_runner_components.remote_instance_weight_transporter import ( + RemoteInstanceWeightTransporter, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +def _free_port(): + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +def _start_bootstrap_server(): + port = _free_port() + server = EngineInfoBootstrapServer("127.0.0.1", port) + url = f"http://127.0.0.1:{port}" + for _ in range(100): + try: + if requests.get(f"{url}/health", timeout=1).status_code == 200: + return server, url + except requests.exceptions.ConnectionError: + time.sleep(0.05) + raise TimeoutError(f"the bootstrap server on {url} did not start") + + +def _publish(role, session_id, port): + model_config = SimpleNamespace( + remote_instance_weight_loader_backend=None, engine_info_bootstrap_port=port + ) + with ( + patch.object( + transporter_mod, + "get_parallel", + lambda: SimpleNamespace(tp_rank=0, dist_init_addr=None), + ), + patch.object(transporter_mod, "get_model", lambda: model_config), + patch.object( + transporter_mod, "remote_instance_transfer_engine_enabled", lambda: True + ), + patch.object( + transporter_mod, + "register_memory_region", + lambda model, engine: {f"{role}.weight": (1, 4, 2)}, + ), + ): + transporter = RemoteInstanceWeightTransporter( + server_args=SimpleNamespace(registers_parallelism_config=lambda: True), + get_model=lambda: torch.nn.Module(), + gpu_id=0, + ) + transporter.engine = object() + transporter.session_id = session_id + transporter.parallelism_config = SimpleNamespace( + to_dict=lambda: {"published_by": role} + ) + transporter.maybe_register_and_publish_weight_info(role=role) + + +class TestPublicationRoles(CustomTestCase): + def test_a_draft_publishing_for_the_same_rank_keeps_its_target(self): + """Keyed by rank alone, the draft replaced the target, and p2p wrote target weights by the draft's table.""" + server, url = _start_bootstrap_server() + self.addCleanup(server.close) + _publish("target", "target-session", server.port) + _publish("draft", "draft-session", server.port) + + def get(path, **params): + return requests.get(f"{url}/{path}", params={"rank": 0, **params}).json() + + self.assertEqual( + get("get_transfer_engine_info")["remote_instance_transfer_engine_info"], + ["target-session", {"target.weight": [1, 4, 2]}], + ) + self.assertEqual( + get("get_transfer_engine_info", role="draft")[ + "remote_instance_transfer_engine_info" + ], + ["draft-session", {"draft.weight": [1, 4, 2]}], + ) + self.assertEqual(get("get_parallelism_config"), {"published_by": "target"}) + self.assertEqual( + get("get_parallelism_config", role="draft"), {"published_by": "draft"} + ) + + def test_drafts_publish_once_their_pools_bind_the_targets_embed_and_head(self): + """EAGLE binds the target's embed and head into its draft while allocating the draft's pools; a table + published before that holds the draft's own embed and head, which are freed.""" + events = [] + + class _Transporter: + def maybe_register_and_publish_weight_info(self, role): + events.append(("published", role)) + + runner = SimpleNamespace(remote_instance_weight_transporter=_Transporter()) + + class _DraftWorker: + def alloc_memory_pool(self, **kwargs): + events.append("bound the target's embed and head") + + def init_hicache_draft_plan(self): + pass + + def weight_update_runners(self): + return [("draft", runner)] + + scheduler = Scheduler.__new__(Scheduler) + scheduler.init_target_memory_pool = lambda: None + scheduler.tp_worker = SimpleNamespace( + get_memory_pool=lambda: (None, None), + model_runner=SimpleNamespace(memory_pool_config=None), + ) + scheduler.draft_worker = _DraftWorker() + with patch.object( + scheduler_mod.kv_cache_builder, + "resolve_decode_retraction_backup", + lambda tp_worker: None, + ): + scheduler.init_memory_pools() + + self.assertEqual( + events, ["bound the target's embed and head", ("published", "draft")] + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/model_executor/model_runner_components/test_weight_updater_published_storage.py b/test/registered/unit/model_executor/model_runner_components/test_weight_updater_published_storage.py new file mode 100644 index 000000000000..b2de83768ef8 --- /dev/null +++ b/test/registered/unit/model_executor/model_runner_components/test_weight_updater_published_storage.py @@ -0,0 +1,100 @@ +"""A weight-update session must leave every parameter that peers write by address where it was published.""" + +import unittest +from contextlib import contextmanager +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +import sglang.srt.model_executor.model_runner_components.weight_updater as weight_updater_mod +from sglang.srt.model_executor.model_runner_components.weight_updater import ( + WeightUpdater, +) +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 _Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.zeros(4), requires_grad=False) + + +def _published(model): + return { + name: (param.data_ptr(), param.numel(), param.element_size()) + for name, param in model.named_parameters() + } + + +def _weight_updater(model, published_weight_info): + model_runner = SimpleNamespace( + remote_instance_weight_transporter=SimpleNamespace( + weight_info=published_weight_info + ) + ) + return WeightUpdater( + tp_rank=0, + device="cpu", + gpu_id=0, + model_config=None, + custom_weight_loaders={}, + get_model=lambda: model, + update_model_fields=None, + recapture_cuda_graph=None, + get_model_runner=lambda: model_runner, + ) + + +@contextmanager +def _postprocess(postprocess_weights): + loader = weight_updater_mod.DefaultModelLoader + with ( + patch.object( + loader, "restore_weights_before_loading", lambda model, device: None + ), + patch.object(loader, "postprocess_weights", postprocess_weights), + ): + yield + + +def _rebind_weight(model, device): + model.weight.data = torch.zeros(4) + + +class TestPublishedStorage(CustomTestCase): + def test_a_session_that_keeps_published_storage_passes(self): + model = _Model() + weight_updater = _weight_updater(model, _published(model)) + + with _postprocess(lambda model, device: None): + weight_updater.begin_weight_update() + weight_updater.end_weight_update(run_post_load=False) + + def test_a_postprocess_that_moves_a_published_parameter_fails_the_session(self): + """Peer writes into the old address would land in freed memory.""" + model = _Model() + weight_updater = _weight_updater(model, _published(model)) + + with _postprocess(_rebind_weight): + weight_updater.begin_weight_update() + with self.assertRaisesRegex( + AssertionError, "moved published parameters.*weight published at" + ): + weight_updater.end_weight_update(run_post_load=False) + + def test_a_rank_that_published_nothing_is_not_checked(self): + """Without published addresses no peer writes by address, so storage may move.""" + model = _Model() + weight_updater = _weight_updater(model, None) + + with _postprocess(_rebind_weight): + weight_updater.begin_weight_update() + weight_updater.end_weight_update(run_post_load=False) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/model_loader/test_default_model_loader.py b/test/registered/unit/model_loader/test_default_model_loader.py index f2cb180b1003..2ed200d3a9e1 100644 --- a/test/registered/unit/model_loader/test_default_model_loader.py +++ b/test/registered/unit/model_loader/test_default_model_loader.py @@ -167,5 +167,66 @@ def load_weights_and_postprocess(self, model, weights, target_device): ) +class _ShardedParameter(torch.nn.Parameter): + pass + + +class _TiedModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.embed = torch.nn.Embedding(8, 4) + self.lm_head = torch.nn.Linear(4, 8, bias=False) + self.lm_head.weight = self.embed.weight + self.proj = torch.nn.Module() + self.proj.weight = _ShardedParameter(torch.empty(3, 4), requires_grad=False) + self.proj.weight.weight_loader = "sharded loader" + self.register_buffer("cos_sin_cache", torch.empty(5)) + + +class TestInitializeModelWithoutStorage(CustomTestCase): + def _initialize(self): + built_on = [] + + def initialize_model(model_config, load_config, quant_config): + built_on.append(torch.empty(0).device.type) + return _TiedModel() + + loader = object.__new__(DefaultModelLoader) + loader.load_config = object() + with ( + patch.object(loader_mod, "_get_quantization_config", return_value=None), + patch.object(loader_mod, "_initialize_model", side_effect=initialize_model), + ): + model, built_shapes_by_name = loader.initialize_model_without_storage( + model_config=SimpleNamespace(dtype=torch.float32), + device=torch.device("cpu"), + ) + return model, built_shapes_by_name, built_on + + def test_the_model_is_built_on_meta_and_keeps_no_parameter_storage(self): + """A trainer holding a replica per engine rank cannot afford each one's startup allocation.""" + model, built_shapes_by_name, built_on = self._initialize() + + self.assertEqual(built_on, ["meta"]) + self.assertEqual( + {name: param.numel() for name, param in model.named_parameters()}, + {"embed.weight": 0, "proj.weight": 0}, + ) + self.assertEqual( + built_shapes_by_name, + {"embed.weight": torch.Size([8, 4]), "proj.weight": torch.Size([3, 4])}, + ) + self.assertEqual(model.cos_sin_cache.device.type, "meta") + + def test_parameters_keep_their_subclass_attributes_and_ties(self): + """Loaders dispatch on the parameter class and its attributes, and a tied weight must stay one parameter.""" + model, _, _ = self._initialize() + + self.assertIsInstance(model.proj.weight, _ShardedParameter) + self.assertEqual(model.proj.weight.weight_loader, "sharded loader") + self.assertFalse(model.proj.weight.requires_grad) + self.assertIs(model.lm_head.weight, model.embed.weight) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/model_loader/test_remote_instance_loader.py b/test/registered/unit/model_loader/test_remote_instance_loader.py index 5962ccb994c4..829f02e6b0e9 100644 --- a/test/registered/unit/model_loader/test_remote_instance_loader.py +++ b/test/registered/unit/model_loader/test_remote_instance_loader.py @@ -1,6 +1,10 @@ """Unit tests for RemoteInstanceModelLoader construction - no server, no weights.""" import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch import sglang.srt.model_loader.loader as loader_mod from sglang.srt.configs.load_config import LoadConfig, LoadFormat @@ -73,5 +77,44 @@ def test_all_backends_construct_without_extra_config(self): self.assertFalse(loader.load_config.model_loader_extra_config) +class TestRemoteInstanceTransferEngineRole(CustomTestCase): + def test_each_runner_loads_the_seed_weights_of_its_own_role(self): + loader = loader_mod.RemoteInstanceModelLoader( + LoadConfig( + load_format=LoadFormat.REMOTE_INSTANCE, + remote_instance_weight_loader_backend=RemoteInstanceWeightLoaderBackend.TRANSFER_ENGINE, + remote_instance_weight_loader_transfer_engine=object(), + remote_instance_weight_loader_seed_instance_ip="127.0.0.1", + remote_instance_weight_loader_seed_instance_service_port=30000, + tp_rank=1, + ) + ) + for is_draft_model, role in ((False, "target"), (True, "draft")): + with self.subTest(role=role): + queried = [] + with ( + patch.object(loader_mod, "_get_quantization_config"), + patch.object( + loader_mod, + "_initialize_model", + return_value=torch.nn.Linear(1, 1), + ), + patch.object(loader_mod, "register_memory_region"), + patch.object( + loader_mod, + "get_remote_instance_transfer_engine_info_per_rank", + side_effect=lambda *args: queried.append(args) or (None, None), + ), + ): + with self.assertRaises(RuntimeError): + loader.load_model( + model_config=SimpleNamespace( + is_draft_model=is_draft_model, dtype=torch.float32 + ), + device_config=SimpleNamespace(device="cpu"), + ) + self.assertEqual(queried, [("http://127.0.0.1:30000", 1, role)]) + + if __name__ == "__main__": unittest.main() diff --git a/test/srt/models/test_params_mapping.py b/test/srt/models/test_params_mapping.py deleted file mode 100644 index e064266cab6a..000000000000 --- a/test/srt/models/test_params_mapping.py +++ /dev/null @@ -1,292 +0,0 @@ -"""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) From 7f79d11d948aa37b49821475e075369bb0e801f3 Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Mon, 12 Oct 2026 16:55:38 -0700 Subject: [PATCH 43/56] Quantized weight-update fixes: Qwen3.5 draft KV budget, UE8M0 partial blocks, mxfp8 fused indexer, NVFP4 input scales (#43251) --- python/sglang/srt/layers/layernorm.py | 5 +- .../srt/layers/quantization/fp8_utils.py | 16 ++-- .../srt/layers/quantization/modelopt_quant.py | 3 + .../spec_aux_hidden_state.py | 5 ++ .../model_runner_components/weight_updater.py | 3 + python/sglang/srt/model_loader/loader.py | 9 +++ .../deepseek_common/deepseek_weight_loader.py | 24 ++++-- .../sglang/srt/models/glm4_moe_lite_nextn.py | 4 + .../layers/quantization/test_fp8_utils.py | 17 ++-- .../test_modelopt_nvfp4_moe_weight_check.py | 81 +++++++++++++++++++ .../test_spec_aux_hidden_state.py | 33 ++++++++ .../test_weight_updater_published_storage.py | 21 ++++- .../model_loader/test_default_model_loader.py | 19 ++++- .../test_deepseek_fused_indexer_mxfp8.py | 53 ++++++++++++ .../test_glm4_moe_lite_nextn_post_load.py | 39 +++++++++ 15 files changed, 309 insertions(+), 23 deletions(-) create mode 100644 test/registered/unit/layers/quantization/test_modelopt_nvfp4_moe_weight_check.py create mode 100644 test/registered/unit/models/test_deepseek_fused_indexer_mxfp8.py create mode 100644 test/registered/unit/models/test_glm4_moe_lite_nextn_post_load.py diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index 2d3367d1d4ff..d5954344ea61 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/python/sglang/srt/layers/layernorm.py @@ -1134,8 +1134,11 @@ def __init__( def _weight_loader(self, param: torch.Tensor, loaded_weight: torch.Tensor) -> None: assert param.size() == loaded_weight.size() param.data.copy_(loaded_weight) + self.refresh_gemma_weight() + + def refresh_gemma_weight(self) -> None: # Keep storage stable for CUDA graphs or fused paths that capture this buffer. - torch.add(param.data, 1.0, out=self.gemma_weight) + torch.add(self.weight.data, 1.0, out=self.gemma_weight) def _forward_impl( self, diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index 8344cb9a220c..0ffee501442c 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -1937,24 +1937,24 @@ def _inverse_transform_scale_ue8m0_impl(sf_packed): assert len(sf_packed.shape) == 2, f"{sf_packed.shape=}" assert sf_packed.dtype == torch.int32 - mn_repeat_128, k_div_4 = sf_packed.shape - mn = mn_repeat_128 // block_size + # one row per weight row, so the last block is partial when the rows are not a multiple of 128 + num_rows, k_div_4 = sf_packed.shape + mn = -(-num_rows // block_size) k = k_div_4 * 4 # packed u8 -> fp32 - sf_u8 = sf_packed.contiguous().flatten().view(torch.uint8).view(mn_repeat_128, k) + sf_u8 = sf_packed.contiguous().flatten().view(torch.uint8).view(num_rows, k) sf_fp32 = (sf_u8.to(torch.int32) << 23).view(torch.float32) # remove repeat - sf_reshaped = sf_fp32.view(mn, block_size, k) - sf_unrepeated = sf_reshaped[:, 0:1, :] - if not torch.all(sf_unrepeated == sf_reshaped): + sf_unrepeated = sf_fp32[::block_size].contiguous() + sf_repeated = sf_unrepeated.repeat_interleave(block_size, dim=0)[:num_rows] + if not torch.all(sf_repeated == sf_fp32): from sglang.srt.debug_utils.dumper import get_tensor_info raise AssertionError( - f"sf_unrepeated != sf_reshaped ({get_tensor_info(sf_unrepeated)=} {get_tensor_info(sf_reshaped)=})" + f"sf_repeated != sf_fp32 ({get_tensor_info(sf_repeated)=} {get_tensor_info(sf_fp32)=})" ) - sf_unrepeated = sf_unrepeated.squeeze(1).contiguous() assert sf_unrepeated.shape == (mn, k) return sf_unrepeated diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index 806c52bc25b8..7daad0580d31 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -2574,6 +2574,9 @@ def create_weights( ) w2_input_scale._sglang_require_global_experts = True layer.register_parameter("w2_input_scale", w2_input_scale) + # calibrated activation scales, as the KV cache scales: no weight update carries them + w13_input_scale._skip_weight_check = True + w2_input_scale._skip_weight_check = True def _build_mega_moe_weights(self, layer: torch.nn.Module) -> None: # Activations are quantized per token at dispatch, so w13_input_scale diff --git a/python/sglang/srt/model_executor/model_runner_components/spec_aux_hidden_state.py b/python/sglang/srt/model_executor/model_runner_components/spec_aux_hidden_state.py index 950761105fae..a5d38150828e 100644 --- a/python/sglang/srt/model_executor/model_runner_components/spec_aux_hidden_state.py +++ b/python/sglang/srt/model_executor/model_runner_components/spec_aux_hidden_state.py @@ -94,6 +94,11 @@ def _resolve_eagle_aux_hidden_state( is_draft_model=True, ) num_nextn_predict_layers = draft_model_config.num_nextn_predict_layers + if num_nextn_predict_layers is None: + # Qwen3.5 names its MTP depth mtp_num_hidden_layers; models without MTP have neither + num_nextn_predict_layers = getattr( + draft_model_config.hf_text_config, "mtp_num_hidden_layers", None + ) if num_nextn_predict_layers is not None: config.eagle_draft_num_layers = int(num_nextn_predict_layers) elif get_spec().speculative_draft_model_path: 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 267f2c7e665c..072c97212c68 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 @@ -12,6 +12,7 @@ DefaultModelLoader, get_model_loader, post_load_weights, + refresh_gemma_weights, ) from sglang.srt.model_loader.utils import set_default_torch_dtype from sglang.srt.model_loader.weight_utils import default_weight_loader @@ -306,6 +307,8 @@ def begin_weight_update(self: WeightUpdater) -> None: def end_weight_update(self: WeightUpdater, *, run_post_load: bool) -> None: if run_post_load: post_load_weights(self.get_model()) + else: + refresh_gemma_weights(self.get_model()) DefaultModelLoader.postprocess_weights( self.get_model(), torch.device(self.device) ) diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index cd13ea193a84..ae8266aede23 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -85,6 +85,7 @@ from sglang.srt.distributed import ( model_parallel_is_initialized, ) +from sglang.srt.layers.layernorm import GemmaRMSNorm from sglang.srt.layers.modelopt_utils import QUANT_CFG_CHOICES from sglang.srt.layers.moe.utils import ( install_shared_experts_fusion_decision, @@ -325,6 +326,14 @@ def post_load_weights(model: nn.Module) -> None: # `is_nextn=True`, so the loader doesn't need to know. if hasattr(model, "post_load_weights"): model.post_load_weights() + refresh_gemma_weights(model) + + +def refresh_gemma_weights(model: nn.Module) -> None: + # its weight loader derives the buffer, so writes that bypass the loader leave it stale + for module in model.modules(): + if isinstance(module, GemmaRMSNorm): + module.refresh_gemma_weight() def _modules_with_quant_method(model: nn.Module): 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 774f325b7c75..3b4607b8bfd0 100644 --- a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py +++ b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py @@ -39,6 +39,7 @@ from sglang.srt.layers.quantization.int8_utils import ( block_dequant as int8_block_dequant, ) +from sglang.srt.layers.quantization.mxfp8_block_convert import dequant_mxfp8_2d_to_bf16 from sglang.srt.layers.utils import get_layer_id from sglang.srt.model_loader.utils import ( maybe_executor_submit, @@ -86,6 +87,19 @@ def _get_indexer_weight_block_size( return [128, 128] +def _dequant_fused_indexer_shard( + weight: torch.Tensor, + scale: torch.Tensor, + quant_config: Optional[QuantizationConfig], +) -> torch.Tensor: + # mxfp8 scales are UE8M0 exponents stored as uint8, not fp32 multipliers + if scale.dtype == torch.uint8: + return dequant_mxfp8_2d_to_bf16(weight, scale) + return block_quant_dequant( + weight, scale, _get_indexer_weight_block_size(quant_config), torch.bfloat16 + ) + + def _load_fused_indexer_wk( name: str, loaded_weight: torch.Tensor, @@ -119,9 +133,8 @@ def _load_fused_indexer_wk( ) if "weight" in entry and "scale" in entry: pending.pop(fused_name + ".weights_proj") - block_size = _get_indexer_weight_block_size(quant_config) - weights_bf16 = block_quant_dequant( - entry["weight"], entry["scale"], block_size, torch.bfloat16 + weights_bf16 = _dequant_fused_indexer_shard( + entry["weight"], entry["scale"], quant_config ) fused_param.data[-weights_bf16.shape[0] :].copy_(weights_bf16) return True @@ -139,9 +152,8 @@ def _load_fused_indexer_wk( ) if "weight" in entry and "scale" in entry: pending.pop(fused_name) - block_size = _get_indexer_weight_block_size(quant_config) - wk_bf16 = block_quant_dequant( - entry["weight"], entry["scale"], block_size, torch.bfloat16 + wk_bf16 = _dequant_fused_indexer_shard( + entry["weight"], entry["scale"], quant_config ) fused_param.data[: wk_bf16.shape[0]].copy_(wk_bf16) return True diff --git a/python/sglang/srt/models/glm4_moe_lite_nextn.py b/python/sglang/srt/models/glm4_moe_lite_nextn.py index 25e0b7c93a86..86f6b9f50668 100644 --- a/python/sglang/srt/models/glm4_moe_lite_nextn.py +++ b/python/sglang/srt/models/glm4_moe_lite_nextn.py @@ -179,5 +179,9 @@ def forward( def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): super().load_weights(weights, is_nextn=True) + def post_load_weights(self, is_nextn=True, weight_names=None): + # loaders and weight-update sessions that bypass load_weights call it with no arguments + super().post_load_weights(is_nextn=True, weight_names=weight_names) + EntryClass = [Glm4MoeLiteForCausalLMNextN] diff --git a/test/registered/unit/layers/quantization/test_fp8_utils.py b/test/registered/unit/layers/quantization/test_fp8_utils.py index d0ed39ca3a1c..10626babd195 100644 --- a/test/registered/unit/layers/quantization/test_fp8_utils.py +++ b/test/registered/unit/layers/quantization/test_fp8_utils.py @@ -82,13 +82,18 @@ def is_deep_gemm(self): class TestInverseTransformScaleUe8m0(CustomTestCase): def test_round_trip(self): + for shape in ( + # DeepSeek V3 kv_b_proj + (32768, 512), + # GLM-5 fused q_a + kv_a: 2048 + 576 rows, the last block partial + (2624, 6144), + ): + with self.subTest(shape=shape): + self._check_round_trip(shape) + + def _check_round_trip(self, shape): for _ in range(100): - weight_bf16 = torch.randn( - # DeepSeek V3 kv_b_proj - (32768, 512), - dtype=torch.bfloat16, - device="cuda", - ) + weight_bf16 = torch.randn(shape, dtype=torch.bfloat16, device="cuda") weight_block_size = [128, 128] diff --git a/test/registered/unit/layers/quantization/test_modelopt_nvfp4_moe_weight_check.py b/test/registered/unit/layers/quantization/test_modelopt_nvfp4_moe_weight_check.py new file mode 100644 index 000000000000..df3d941fb93b --- /dev/null +++ b/test/registered/unit/layers/quantization/test_modelopt_nvfp4_moe_weight_check.py @@ -0,0 +1,81 @@ +"""The weight checker must leave the NVFP4 MoE activation input scales alone: no weight update carries them, so +a reset before an update leaves them random and the comparison after it fails whatever the update wrote. + +The platform check is stubbed so the layer is built on CPU, as in test_modelopt_nvfp4_moe_dispatch.py. +""" + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=6, suite="base-a-test-cpu") + +import unittest +from types import SimpleNamespace +from unittest import mock + +import torch + +# Import modelopt_quant first (see test_modelopt_nvfp4_moe_scales.py for the circular-import reason). +# isort: off +from sglang.srt.layers.quantization import modelopt_quant +from sglang.srt.layers.quantization.modelopt_quant import ModelOptNvFp4FusedMoEMethod + +# isort: on +from sglang.srt.layers.moe.utils import MoeRunnerBackend +from sglang.srt.utils.weight_checker import _is_skip_weight_check +from sglang.test.test_utils import CustomTestCase + +NUM_EXPERTS = 8 +HIDDEN = 64 +INTERMEDIATE = 32 + + +def _nvfp4_moe_layer() -> torch.nn.Module: + with ( + mock.patch.object( + modelopt_quant, + "get_moe_runner_backend", + return_value=MoeRunnerBackend.FLASHINFER_TRTLLM_ROUTED, + ), + mock.patch.object( + modelopt_quant, + "get_platform", + return_value=SimpleNamespace(is_blackwell=True), + ), + mock.patch.object(modelopt_quant, "is_cuda", return_value=False), + ): + method = ModelOptNvFp4FusedMoEMethod( + SimpleNamespace( + group_size=16, + is_checkpoint_nvfp4_serialized=True, + use_per_token_activation=False, + get_name=lambda: "modelopt_fp4", + ) + ) + layer = torch.nn.Module() + layer.num_experts = NUM_EXPERTS + layer.num_local_experts = NUM_EXPERTS + layer.moe_ep_rank = 0 + layer.moe_runner_config = SimpleNamespace(is_gated=True) + method.create_weights( + layer, + num_experts=NUM_EXPERTS, + hidden_size=HIDDEN, + intermediate_size_per_partition=INTERMEDIATE, + params_dtype=torch.bfloat16, + weight_loader=lambda *args, **kwargs: None, + ) + return layer + + +class TestNvFp4MoeWeightCheck(CustomTestCase): + def test_the_checker_skips_the_input_scales_and_checks_the_weight_scales(self): + params = dict(_nvfp4_moe_layer().named_parameters()) + + for name in ("w13_input_scale", "w2_input_scale"): + self.assertTrue(_is_skip_weight_check(name, params[name]), name) + for name in ("w13_weight", "w13_weight_scale", "w13_weight_scale_2"): + self.assertFalse(_is_skip_weight_check(name, params[name]), name) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/model_executor/model_runner_components/test_spec_aux_hidden_state.py b/test/registered/unit/model_executor/model_runner_components/test_spec_aux_hidden_state.py index a48604188f99..0b28912dd11e 100644 --- a/test/registered/unit/model_executor/model_runner_components/test_spec_aux_hidden_state.py +++ b/test/registered/unit/model_executor/model_runner_components/test_spec_aux_hidden_state.py @@ -3,9 +3,13 @@ import pytest +from sglang.srt.model_executor.model_runner_components import spec_aux_hidden_state from sglang.srt.model_executor.model_runner_components.spec_aux_hidden_state import ( + SpecAuxHiddenStateConfig, _map_muse_target_layer_ids, + _resolve_eagle_aux_hidden_state, ) +from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=11, suite="base-a-test-cpu") @@ -34,5 +38,34 @@ def test_muse_target_layer_id_mapping(target_model_type, draft_architecture, exp ) +def test_an_mtp_depth_named_mtp_num_hidden_layers_counts_as_the_eagle_draft( + monkeypatch, +): + """Qwen3.5 names its MTP depth mtp_num_hidden_layers. Read as no draft, the target's + KV budget leaves the draft's pool, of the target's token count, outside + mem_fraction_static.""" + monkeypatch.setattr( + spec_aux_hidden_state, + "get_spec", + lambda: SimpleNamespace(speculative_draft_model_path=None), + ) + config = SpecAuxHiddenStateConfig() + + _resolve_eagle_aux_hidden_state( + config=config, + server_args=None, + model_config=SimpleNamespace( + num_nextn_predict_layers=None, + hf_text_config=SimpleNamespace(mtp_num_hidden_layers=1), + is_hybrid_swa=False, + is_deepseek_v4_arch=False, + ), + spec_algorithm=SpeculativeAlgorithm.EAGLE, + is_draft_worker=False, + ) + + assert config.eagle_draft_num_layers == 1 + + if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/unit/model_executor/model_runner_components/test_weight_updater_published_storage.py b/test/registered/unit/model_executor/model_runner_components/test_weight_updater_published_storage.py index b2de83768ef8..4cb81a6fb3b4 100644 --- a/test/registered/unit/model_executor/model_runner_components/test_weight_updater_published_storage.py +++ b/test/registered/unit/model_executor/model_runner_components/test_weight_updater_published_storage.py @@ -1,4 +1,5 @@ -"""A weight-update session must leave every parameter that peers write by address where it was published.""" +"""A weight-update session must leave every parameter that peers write by address where it was published, and +rebuild what the loaders derive from the weights it wrote.""" import unittest from contextlib import contextmanager @@ -8,6 +9,7 @@ import torch import sglang.srt.model_executor.model_runner_components.weight_updater as weight_updater_mod +from sglang.srt.layers.layernorm import GemmaRMSNorm from sglang.srt.model_executor.model_runner_components.weight_updater import ( WeightUpdater, ) @@ -96,5 +98,22 @@ def test_a_rank_that_published_nothing_is_not_checked(self): weight_updater.end_weight_update(run_post_load=False) +class TestDerivedState(CustomTestCase): + def test_gemma_weight_follows_a_norm_weight_loaded_without_its_loader(self): + """`direct` and custom loads write params without GemmaRMSNorm's loader yet mark the session loaded, so + fused allreduce kernels would read the old `gemma_weight`.""" + model = torch.nn.Module() + model.norm = GemmaRMSNorm(4) + weight_updater = _weight_updater(model, published_weight_info=None) + new_weight = torch.tensor([0.5, -0.25, 0.0, 2.0]) + + with _postprocess(lambda model, device: None): + weight_updater.begin_weight_update() + model.norm.weight.data.copy_(new_weight) + weight_updater.end_weight_update(run_post_load=False) + + torch.testing.assert_close(model.norm.gemma_weight, new_weight + 1) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/model_loader/test_default_model_loader.py b/test/registered/unit/model_loader/test_default_model_loader.py index 2ed200d3a9e1..5cfb2eeda675 100644 --- a/test/registered/unit/model_loader/test_default_model_loader.py +++ b/test/registered/unit/model_loader/test_default_model_loader.py @@ -6,9 +6,10 @@ import torch import sglang.srt.model_loader.loader as loader_mod +from sglang.srt.layers.layernorm import GemmaRMSNorm from sglang.srt.layers.moe import MoeRunnerBackend from sglang.srt.layers.quantization.fp8 import Fp8Config -from sglang.srt.model_loader.loader import DefaultModelLoader +from sglang.srt.model_loader.loader import DefaultModelLoader, post_load_weights from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -228,5 +229,21 @@ def test_parameters_keep_their_subclass_attributes_and_ties(self): self.assertIs(model.lm_head.weight, model.embed.weight) +class TestPostLoadWeights(CustomTestCase): + def test_gemma_weight_follows_a_weight_written_without_its_loader(self): + """p2p weight updates and the remote-instance loader write params by address; fused allreduce kernels read + `gemma_weight`, so a stale one would run the old norm.""" + model = torch.nn.Module() + model.norm = GemmaRMSNorm(4) + gemma_weight_address = model.norm.gemma_weight.data_ptr() + new_weight = torch.tensor([0.5, -0.25, 0.0, 2.0]) + + model.norm.weight.data.copy_(new_weight) + post_load_weights(model) + + torch.testing.assert_close(model.norm.gemma_weight, new_weight + 1) + self.assertEqual(model.norm.gemma_weight.data_ptr(), gemma_weight_address) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/models/test_deepseek_fused_indexer_mxfp8.py b/test/registered/unit/models/test_deepseek_fused_indexer_mxfp8.py new file mode 100644 index 000000000000..860314fd1f57 --- /dev/null +++ b/test/registered/unit/models/test_deepseek_fused_indexer_mxfp8.py @@ -0,0 +1,53 @@ +"""An mxfp8 indexer wk / weights_proj must reach the fused bf16 wk_weights_proj dequantized with its UE8M0 scales.""" + +import unittest +from types import SimpleNamespace + +import torch + +from sglang.srt.layers.quantization.mxfp8_block_convert import ( + dequant_mxfp8_2d_to_bf16, +) +from sglang.srt.models.deepseek_common.deepseek_weight_loader import ( + _load_fused_indexer_wk, +) +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") + +HEAD_DIM, N_HEADS, HIDDEN = 128, 32, 256 +PREFIX = "model.layers.0.self_attn.indexer" + + +class TestFusedIndexerMxfp8(CustomTestCase): + def test_mxfp8_shards_are_dequantized_with_their_ue8m0_scales(self): + """The scales are biased exponents: used as multipliers, the indexer keys + are off by orders of magnitude until a bf16 update overwrites them.""" + torch.manual_seed(0) + fused = torch.zeros(HEAD_DIM + N_HEADS, HIDDEN, dtype=torch.bfloat16) + params_dict = {f"{PREFIX}.wk_weights_proj.weight": fused} + quant_config = SimpleNamespace(weight_block_size=[1, 32]) + pending = {} + expected = {} + for shard, rows in (("wk", HEAD_DIM), ("weights_proj", N_HEADS)): + weight = torch.randn(rows, HIDDEN).to(torch.float8_e4m3fn) + scale = torch.randint(120, 130, (rows, HIDDEN // 32), dtype=torch.uint8) + expected[shard] = dequant_mxfp8_2d_to_bf16(weight, scale) + for name, tensor in (("weight", weight), ("weight_scale_inv", scale)): + self.assertTrue( + _load_fused_indexer_wk( + f"{PREFIX}.{shard}.{name}", + tensor, + params_dict, + pending, + quant_config, + ) + ) + + self.assertTrue(torch.equal(fused[:HEAD_DIM], expected["wk"])) + self.assertTrue(torch.equal(fused[-N_HEADS:], expected["weights_proj"])) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/models/test_glm4_moe_lite_nextn_post_load.py b/test/registered/unit/models/test_glm4_moe_lite_nextn_post_load.py new file mode 100644 index 000000000000..a48382fe3965 --- /dev/null +++ b/test/registered/unit/models/test_glm4_moe_lite_nextn_post_load.py @@ -0,0 +1,39 @@ +"""The GLM-lite NextN draft must take the NextN branch of post_load_weights when called without arguments.""" + +import unittest +from unittest.mock import patch + +import torch + +from sglang.srt.model_loader.loader import post_load_weights +from sglang.srt.models.glm4_moe_lite import Glm4MoeLiteForCausalLM +from sglang.srt.models.glm4_moe_lite_nextn import Glm4MoeLiteForCausalLMNextN +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +class TestGlm4MoeLiteNextNPostLoad(CustomTestCase): + def test_a_bare_post_load_runs_the_nextn_branch(self): + """A weight-update session that writes weights without load_weights (p2p) finishes with a bare + post_load_weights; on the draft the main model's branch walks a layer range it does not have.""" + nextn_flags = [] + draft = Glm4MoeLiteForCausalLMNextN.__new__(Glm4MoeLiteForCausalLMNextN) + # an empty module, without the layers its __init__ would build + torch.nn.Module.__init__(draft) + + with patch.object( + Glm4MoeLiteForCausalLM, + "post_load_weights", + lambda self, is_nextn=False, weight_names=None: nextn_flags.append( + is_nextn + ), + ): + post_load_weights(draft) + + self.assertEqual(nextn_flags, [True]) + + +if __name__ == "__main__": + unittest.main() From 3042bf6aebf0ec2b46e66c9350bcdc4141804ed7 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Wed, 14 Oct 2026 10:36:03 +0800 Subject: [PATCH 44/56] Register the batch-invariant matmul fallback flag in Envs (#44166) --- python/sglang/srt/environ.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 17e77045e89e..e27bd2d17288 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1170,6 +1170,7 @@ class Envs: # DeepGEMM # =================================================================== SGLANG_ENABLE_JIT_DEEPGEMM = EnvBool(True) + SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_FALLBACK_VARIANT = EnvBool(False) # Enable the allowlisted low-M BF16 Split-K GEMM path on Blackwell. Shapes # outside the measured allowlist continue to use CuTe DSL/cuBLAS. SGLANG_ENABLE_BF16_SPLITK_GEMM = EnvBool(True) From f31278164939dd0746037abbb7566de4164e6767 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Wed, 14 Oct 2026 11:02:31 +0800 Subject: [PATCH 45/56] Report prompt KV weight versions in generation metadata (#44164) --- docs/docs/advanced_features/sglang_for_rl.mdx | 1 + python/sglang/srt/arg_groups/field_order.py | 1 + python/sglang/srt/arg_groups/fields/exec_.py | 4 + python/sglang/srt/arg_groups/kv_cache_hook.py | 14 + python/sglang/srt/disaggregation/prefill.py | 3 + .../srt/managers/detokenizer_manager.py | 1 + python/sglang/srt/managers/io_struct.py | 3 + .../srt/managers/multi_tokenizer_mixin.py | 6 + python/sglang/srt/managers/schedule_batch.py | 3 + python/sglang/srt/managers/scheduler.py | 13 + .../batch_result_processor.py | 12 + .../scheduler_components/output_streamer.py | 8 + .../sglang/srt/managers/tokenizer_manager.py | 15 +- python/sglang/srt/managers/tp_worker.py | 1 + python/sglang/srt/managers/utils.py | 3 + .../mem_cache/kv_weight_version_tracker.py | 159 +++++++++ .../sglang/srt/model_executor/model_runner.py | 6 + .../srt/speculative/dflash_worker_v2.py | 1 + .../srt/speculative/eagle_worker_common.py | 1 + python/sglang/srt/utils/weight_versions.py | 23 +- .../rl/test_weight_version_spans.py | 303 ++++++++++++++-- .../mlx/test_attention_patching.py | 1 + ...st_batch_result_processor_hidden_states.py | 3 + ...t_batch_result_processor_mamba_boundary.py | 1 + ...est_batch_result_processor_spec_grammar.py | 37 ++ .../managers/test_msgpack_ipc_roundtrip.py | 27 +- .../managers/test_multi_tokenizer_mixin.py | 23 ++ .../test_output_streamer_customized_info.py | 50 +++ .../test_scheduler_weight_version_tracking.py | 256 ++++++++++++- .../test_kv_weight_version_tracker.py | 336 ++++++++++++++++++ .../test_model_runner_kv_weight_versions.py | 103 ++++++ .../unit/server_args/test_server_args.py | 26 ++ .../unit/utils/test_weight_versions.py | 46 +++ 33 files changed, 1441 insertions(+), 49 deletions(-) create mode 100644 python/sglang/srt/mem_cache/kv_weight_version_tracker.py create mode 100644 test/registered/unit/mem_cache/test_kv_weight_version_tracker.py create mode 100644 test/registered/unit/model_executor/test_model_runner_kv_weight_versions.py diff --git a/docs/docs/advanced_features/sglang_for_rl.mdx b/docs/docs/advanced_features/sglang_for_rl.mdx index 1c2c663dd6a3..e86875a1ff64 100644 --- a/docs/docs/advanced_features/sglang_for_rl.mdx +++ b/docs/docs/advanced_features/sglang_for_rl.mdx @@ -608,6 +608,7 @@ A request can outlive a version change: the RL flow retracts it with `pause_gene ``` - Ranges are half-open `[start, end)` over output-token indices, prompt excluded. The usual case is a single span. +- `--enable-prefill-weight-versions` adds `prefill_weight_versions`, the same span list over the prompt tokens, so KV reused from the prefix cache under older weights is visible. ## Easy To Postpone Generation diff --git a/python/sglang/srt/arg_groups/field_order.py b/python/sglang/srt/arg_groups/field_order.py index fd126f2a4061..464fd4ef5223 100644 --- a/python/sglang/srt/arg_groups/field_order.py +++ b/python/sglang/srt/arg_groups/field_order.py @@ -520,5 +520,6 @@ "enable_decoder_swa_bounded_replay", "enable_response_store", "disaggregation_decode_host_receive_threshold", + "enable_prefill_weight_versions", ) # fmt: on diff --git a/python/sglang/srt/arg_groups/fields/exec_.py b/python/sglang/srt/arg_groups/fields/exec_.py index 159d4b0a0b1c..bb3db33894d0 100644 --- a/python/sglang/srt/arg_groups/fields/exec_.py +++ b/python/sglang/srt/arg_groups/fields/exec_.py @@ -98,6 +98,10 @@ class ExecFeatures(msgspec.Struct): bool, "DeepSeek-V4.1 decoder SWA bounded replay: after the last kv_source layer, run the remaining layers over only the last window_size tokens of a prefill. Main and indexer KV stay exact; nothing is replayed. Deterministic for a fixed prompt and chunk size.", ] = False + enable_prefill_weight_versions: A[ + bool, + "Enable returning the weight version that computed each prompt token's KV cache with responses, as meta_info['prefill_weight_versions'].", + ] = False sampling_mask_max_tokens: A[ int, "The maximum number of token IDs in a returned sampling mask. Requests " diff --git a/python/sglang/srt/arg_groups/kv_cache_hook.py b/python/sglang/srt/arg_groups/kv_cache_hook.py index 5db8194b9dda..94d47365322a 100644 --- a/python/sglang/srt/arg_groups/kv_cache_hook.py +++ b/python/sglang/srt/arg_groups/kv_cache_hook.py @@ -396,6 +396,20 @@ def handle_cache_compatibility(server_args: Any) -> None: "and cannot be used at the same time. Please use only one of them." ) + if cfg.enable_prefill_weight_versions: + if cfg.enable_hierarchical_cache: + raise ValueError( + "The argument enable-prefill-weight-versions cannot be combined with " + "enable-hierarchical-cache: KV restored from the host or storage tiers carries " + "no weight version, so the reported prompt versions would be wrong." + ) + if cfg.disaggregation_mode != "null": + raise ValueError( + "The argument enable-prefill-weight-versions is not supported in disaggregation " + f"mode {cfg.disaggregation_mode!r}: prefill workers do not report prompt weight " + "versions yet." + ) + if cfg.disaggregation_decode_enable_offload_kvcache: if cfg.disaggregation_mode != "decode": raise ValueError( diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 9383cfb679e9..64b23b65e939 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -795,6 +795,9 @@ def process_batch_result_disagg_prefill( if result.indexer_topk_output is not None: result.indexer_topk_output.finalize() result.indexer_topk_output = None + if (record := result.kv_weight_version_record) is not None: + record.finalize(tracker=self.kv_weight_version_tracker) + result.kv_weight_version_record = None logprob_pt = 0 aborted_reqs: List[Req] = [] diff --git a/python/sglang/srt/managers/detokenizer_manager.py b/python/sglang/srt/managers/detokenizer_manager.py index ef253a939838..01058f95a3de 100644 --- a/python/sglang/srt/managers/detokenizer_manager.py +++ b/python/sglang/srt/managers/detokenizer_manager.py @@ -504,6 +504,7 @@ def handle_batch_token_id_out(self, recv_obj: BatchTokenIDOutput): placeholder_tokens_val=None, retraction_counts=recv_obj.retraction_counts, weight_versions=recv_obj.weight_versions, + prefill_weight_versions=recv_obj.prefill_weight_versions, token_steps=recv_obj.token_steps, beam_search_output=recv_obj.beam_search_output, dp_ranks=recv_obj.dp_ranks, diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 4ebce9cb53d7..a5044cc114e0 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -1588,6 +1588,7 @@ class BatchTokenIDOutput(BaseBatchReq, kw_only=True): beam_search_output: Optional[List[Optional[BeamSearchOutput]]] = None weight_versions: Optional[List[Optional[WeightVersionSpans]]] = None + prefill_weight_versions: Optional[List[Optional[WeightVersionSpans]]] = None # The trainer step id. Used to know which step's weights are used for sampling. token_steps: Optional[List[List[int]]] = None @@ -1686,6 +1687,7 @@ class BatchStrOutput(BaseBatchReq, kw_only=True): beam_search_output: Optional[List[Optional[BeamSearchOutput]]] = None weight_versions: Optional[List[Optional[WeightVersionSpans]]] = None + prefill_weight_versions: Optional[List[Optional[WeightVersionSpans]]] = None # The trainer step id. Used to know which step's weights are used for sampling. token_steps: Optional[List[List[int]]] = None @@ -2232,6 +2234,7 @@ class AbortReq(BaseReq, kw_only=True): # because batch requests derive child rids as ``f"{rid}_{i}"`` and an # abort for the parent rid must cover them. prefix: bool = False + prefill_weight_versions: Optional[WeightVersionSpans] = None def __post_init__(self): # FIXME: This is a hack to keep the same with the old code diff --git a/python/sglang/srt/managers/multi_tokenizer_mixin.py b/python/sglang/srt/managers/multi_tokenizer_mixin.py index 150ee49055b5..0021979fa948 100644 --- a/python/sglang/srt/managers/multi_tokenizer_mixin.py +++ b/python/sglang/srt/managers/multi_tokenizer_mixin.py @@ -259,6 +259,9 @@ def _handle_output_by_index(output, i): 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), + prefill_weight_versions=_extract_field_by_index( + output, "prefill_weight_versions", i + ), placeholder_tokens_idx=None, placeholder_tokens_val=None, beam_search_output=_extract_field_by_index( @@ -387,6 +390,9 @@ def _handle_output_by_index(output, i): output, "beam_search_output", i, check_length=True ), weight_versions=_extract_field_by_index(output, "weight_versions", i), + prefill_weight_versions=_extract_field_by_index( + output, "prefill_weight_versions", i + ), token_steps=_extract_field_by_index( output, "token_steps", i, check_length=False ), diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index a94a1df7ae7a..74bed1fcb048 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -22,6 +22,7 @@ ) from sglang.srt.utils.weight_versions import ( WeightVersionEvent, + WeightVersionSpans, truncate_weight_version_events, ) @@ -1230,6 +1231,7 @@ def __init__( self.retracted_stain = False self.weight_version_events: List[WeightVersionEvent] = [] + self.prefill_weight_versions: Optional[WeightVersionSpans] = None # Incremental streamining self.send_token_offset: int = 0 @@ -1994,6 +1996,7 @@ def reset_for_retract(self): self.already_computed = 0 assert not self.kv.holds_kv, "expect it is already released" self.kv.kv_committed_len = 0 + self.prefill_weight_versions = None self.extend_batch_idx = 0 self.decode_batch_idx = 0 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 8bd6402da8e0..a74a095e3d0c 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -297,6 +297,7 @@ maybe_cache_unfinished_req, release_kv_cache, ) +from sglang.srt.mem_cache.kv_weight_version_tracker import KvWeightVersionTracker from sglang.srt.model_executor.forward_batch_info import PPProxyTensors from sglang.srt.model_executor.runner_utils.pool import prewarm_graph_pool_borrow from sglang.srt.model_loader.utils import get_resolved_model_impl @@ -710,6 +711,8 @@ def __init__( self.init_load_inquirer() + self.init_kv_weight_version_tracker() + self.init_output_streamer() self.init_batch_result_processor() @@ -2466,6 +2469,14 @@ def init_load_inquirer(self) -> None: get_decode_moment_totals=lambda: self.decode_moment_totals, ) + def init_kv_weight_version_tracker(self) -> None: + self.kv_weight_version_tracker = KvWeightVersionTracker.maybe_create( + server_args=self.server_args, + model_config=self.model_config, + allocator=self.token_to_kv_pool_allocator, + req_to_token_pool=self.req_to_token_pool, + ) + def init_output_streamer(self) -> None: self.output_streamer = self.get_output_streamer_class()( send_to_detokenizer=self.ipc_channels.send_to_detokenizer, @@ -2508,6 +2519,7 @@ def init_batch_result_processor(self) -> None: tree_cache=self.tree_cache, hisparse_coordinator=self.hisparse_coordinator, req_to_token_pool=self.req_to_token_pool, + kv_weight_version_tracker=self.kv_weight_version_tracker, decode_offload_manager=self.decode_offload_manager, metrics_collector=self.metrics_collector, metrics_reporter=self.metrics_reporter, @@ -6136,4 +6148,5 @@ def _make_abort_req( current_version=get_serving().weight_version, num_output_tokens=len(req.output_ids), ), + prefill_weight_versions=req.prefill_weight_versions, ) 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 e0f922fe8df6..9ebcfbcd064f 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py @@ -76,6 +76,7 @@ ) from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache + from sglang.srt.mem_cache.kv_weight_version_tracker import KvWeightVersionTracker from sglang.srt.mem_cache.memory_pool import ReqToTokenPool from sglang.srt.observability.metrics_collector import SchedulerMetricsCollector @@ -103,6 +104,7 @@ class SchedulerBatchResultProcessor: tree_cache: BasePrefixCache hisparse_coordinator: Optional[HiSparseCoordinator] req_to_token_pool: ReqToTokenPool + kv_weight_version_tracker: Optional[KvWeightVersionTracker] decode_offload_manager: Optional[DecodeKVCacheOffloadManager] metrics_collector: SchedulerMetricsCollector metrics_reporter: SchedulerMetricsReporter @@ -275,6 +277,9 @@ def process_batch_result_prefill( if result.indexer_topk_output is not None: result.indexer_topk_output.finalize() result.indexer_topk_output = None + if (record := result.kv_weight_version_record) is not None: + record.finalize(tracker=self.kv_weight_version_tracker) + result.kv_weight_version_record = None ( logits_output, @@ -343,6 +348,8 @@ def process_batch_result_prefill( if req.inflight_middle_chunks <= 0: req.time_stats.set_prefill_finished_time() + if (x := self.kv_weight_version_tracker) is not None: + x.fill_req_prefill_weight_versions(req) if sampling_mask_finish_reason is not None: req.to_finish = sampling_mask_finish_reason @@ -463,6 +470,8 @@ def process_batch_result_prefill( req.pooled_hidden_state = phs[i] if req.inflight_middle_chunks <= 0: req.time_stats.set_prefill_finished_time() + if (x := self.kv_weight_version_tracker) is not None: + x.fill_req_prefill_weight_versions(req) # Dummy output token for embedding models req.output_ids.append(0) req.update_finish_state() @@ -933,6 +942,9 @@ def process_batch_result_decode( if result.indexer_topk_output is not None: result.indexer_topk_output.finalize() result.indexer_topk_output = None + if (record := result.kv_weight_version_record) is not None: + record.finalize(tracker=self.kv_weight_version_tracker) + result.kv_weight_version_record = None logits_output, next_token_ids, can_run_cuda_graph = ( result.logits_output, diff --git a/python/sglang/srt/managers/scheduler_components/output_streamer.py b/python/sglang/srt/managers/scheduler_components/output_streamer.py index c394a9f64979..9594e3c2230f 100644 --- a/python/sglang/srt/managers/scheduler_components/output_streamer.py +++ b/python/sglang/srt/managers/scheduler_components/output_streamer.py @@ -356,6 +356,7 @@ class _GenerationStreamAccumulator: spec_cap_lens_histogram: list = field(default_factory=list) retraction_counts: list = field(default_factory=list) weight_versions: list = field(default_factory=list) + prefill_weight_versions: list = field(default_factory=list) output_hidden_states: Optional[list] = None routed_experts: Optional[list] = None indexer_topk: Optional[list] = None @@ -524,8 +525,10 @@ def accept(self, *, req: Req) -> None: num_output_tokens=len(output_ids_), ) ) + self.prefill_weight_versions.append(req.prefill_weight_versions) else: self.weight_versions.append(None) + self.prefill_weight_versions.append(None) self.time_stats.append(req.time_stats) @@ -768,5 +771,10 @@ def to_payload( weight_versions=( self.weight_versions if any(self.weight_versions) else None ), + prefill_weight_versions=( + self.prefill_weight_versions + if any(self.prefill_weight_versions) + else None + ), dp_ranks=dp_ranks, ) diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 110c1d9fd573..81cd975306d9 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -182,7 +182,10 @@ from sglang.srt.utils.network import get_zmq_socket from sglang.srt.utils.request_logger import RequestLogger from sglang.srt.utils.watchdog import Watchdog -from sglang.srt.utils.weight_versions import add_weight_versions_to_meta_info +from sglang.srt.utils.weight_versions import ( + add_weight_versions_to_meta_info, + weight_version_spans_to_json, +) from sglang.utils import TypeBasedDispatcher, get_exception_traceback asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) @@ -2524,6 +2527,12 @@ async def _handle_batch_output( spans, num_output_tokens=recv_obj.completion_tokens[i], ) + if recv_obj.prefill_weight_versions and ( + prefill_spans := recv_obj.prefill_weight_versions[i] + ): + meta_info["prefill_weight_versions"] = weight_version_spans_to_json( + prefill_spans + ) # Add detailed cache breakdown if available if ( hasattr(recv_obj, "cached_tokens_details") @@ -3501,6 +3510,10 @@ def _handle_abort_req(self, recv_obj: AbortReq): recv_obj.weight_versions, num_output_tokens=len(state.output_ids), ) + if prefill_spans := recv_obj.prefill_weight_versions: + meta_info["prefill_weight_versions"] = weight_version_spans_to_json( + prefill_spans + ) is_stream = getattr(state.obj, "stream", False) if getattr(state.obj, "return_logprob", False): self.add_logprob_to_meta_info( diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index f38f6e6b57d6..90c1b0ce6653 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -600,6 +600,7 @@ def forward_batch_generation( expert_distribution_metrics=out.expert_distribution_metrics, routed_experts_output=out.routed_experts_output, indexer_topk_output=out.indexer_topk_output, + kv_weight_version_record=out.kv_weight_version_record, ) capture_pre_sample_logits(batch, forward_batch, logits_output) diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py index ed40ce3694a5..7a6a49f0f7c8 100644 --- a/python/sglang/srt/managers/utils.py +++ b/python/sglang/srt/managers/utils.py @@ -17,6 +17,7 @@ ) from sglang.srt.managers import io_struct from sglang.srt.managers.schedule_batch import Req +from sglang.srt.mem_cache.kv_weight_version_tracker import KvWeightVersionRecord from sglang.srt.model_executor.forward_batch_info import PPProxyTensors from sglang.srt.runtime_context import get_spec, max_speculative_num_draft_tokens from sglang.srt.speculative.spec_info import SpeculativeAlgorithm @@ -131,6 +132,7 @@ class GenerationBatchResult: # Routed experts: pending async D2H for overlap scheduling routed_experts_output: Optional[TopkCaptureOutput] = None indexer_topk_output: Optional[TopkCaptureOutput] = None + kv_weight_version_record: Optional[KvWeightVersionRecord] = None # metrics expert_distribution_metrics: Optional[ExpertDistributionMetrics] = None @@ -207,6 +209,7 @@ def copy_to_cpu(self, return_logprob: bool, return_hidden_states: bool = True): sampling_mask_output, self.routed_experts_output, self.indexer_topk_output, + self.kv_weight_version_record, self.expert_distribution_metrics, ): if holder is not None: diff --git a/python/sglang/srt/mem_cache/kv_weight_version_tracker.py b/python/sglang/srt/mem_cache/kv_weight_version_tracker.py new file mode 100644 index 000000000000..3633f48e4214 --- /dev/null +++ b/python/sglang/srt/mem_cache/kv_weight_version_tracker.py @@ -0,0 +1,159 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Callable, Dict, List, Optional + +import msgspec +import torch + +from sglang.srt.runtime_context import get_serving +from sglang.srt.utils.weight_versions import WeightVersionSpan, WeightVersionSpans + +if TYPE_CHECKING: + from sglang.srt.configs.model_config import ModelConfig + from sglang.srt.managers.schedule_batch import Req + from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator + from sglang.srt.mem_cache.memory_pool import ReqToTokenPool + from sglang.srt.model_executor.forward_batch_info import ForwardBatch + from sglang.srt.model_executor.model_runner import ModelRunner + from sglang.srt.server_args import ServerArgs + + +_UNWRITTEN_VERSION_ID = -1 + + +class KvWeightVersionRecord(msgspec.Struct): + slot_indices: torch.Tensor + version: str + + @classmethod + def maybe_capture( + cls, model_runner: ModelRunner, forward_batch: ForwardBatch + ) -> Optional[KvWeightVersionRecord]: + if ( + not model_runner.is_draft_worker + and model_runner.server_args.enable_prefill_weight_versions + and forward_batch.out_cache_loc is not None + ): + return cls.capture( + slot_indices=forward_batch.out_cache_loc, + version=get_serving().weight_version, + ) + return None + + @classmethod + def capture( + cls, *, slot_indices: torch.Tensor, version: Optional[str] + ) -> KvWeightVersionRecord: + assert version is not None + return cls(slot_indices=slot_indices.clone(), version=version) + + def map_device_tensors(self, fn: Callable[[torch.Tensor], torch.Tensor]) -> None: + self.slot_indices = fn(self.slot_indices) + + def finalize(self, *, tracker: Optional[KvWeightVersionTracker]) -> None: + assert tracker is not None + tracker.record(slot_indices=self.slot_indices, version=self.version) + + +class KvWeightVersionTracker: + @classmethod + def maybe_create( + cls, + *, + server_args: ServerArgs, + model_config: ModelConfig, + allocator: BaseTokenToKVPoolAllocator, + req_to_token_pool: ReqToTokenPool, + ) -> Optional[KvWeightVersionTracker]: + if not server_args.enable_prefill_weight_versions: + return None + + assert not model_config.is_encoder_decoder, ( + "--enable-prefill-weight-versions does not support encoder-decoder models" + ) + assert model_config.is_generation, ( + "--enable-prefill-weight-versions does not support embedding or reward models" + ) + assert server_args.pp_size == 1, ( + "--enable-prefill-weight-versions does not support pipeline parallelism" + ) + + from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator + + if isinstance(allocator, SWATokenToKVPoolAllocator): + allocator = allocator.full_attn_allocator + + return cls( + num_slots=allocator.size_full + allocator.page_size, + device=allocator.device, + req_to_token_pool=req_to_token_pool, + ) + + def __init__( + self, *, num_slots: int, device: str, req_to_token_pool: ReqToTokenPool + ): + self._slot_version_ids = torch.full( + (num_slots,), _UNWRITTEN_VERSION_ID, dtype=torch.int32, device=device + ) + self._req_to_token_pool = req_to_token_pool + self._versions = _StringInterner() + + def record(self, *, slot_indices: torch.Tensor, version: Optional[str]) -> None: + assert version is not None + self._slot_version_ids[slot_indices] = self._versions.intern(version) + + def fill_req_prefill_weight_versions(self, req: Req) -> None: + num_prompt_tokens = len(req.origin_input_ids) + assert req.kv.kv_committed_len >= num_prompt_tokens, ( + f"prefill finished with {req.kv.kv_committed_len} committed KV tokens " + f"for a {num_prompt_tokens}-token prompt" + ) + req.prefill_weight_versions = self._lookup_spans( + self._req_to_token_pool.req_to_token[ + req.kv.req_pool_idx, :num_prompt_tokens + ] + ) + + def _lookup_spans(self, slot_indices: torch.Tensor) -> WeightVersionSpans: + version_ids = self._slot_version_ids[slot_indices] + if len(version_ids) == 0: + return [] + if (is_unwritten := version_ids == _UNWRITTEN_VERSION_ID).any(): + raise ValueError( + "KV slots without a recorded weight version were looked up: " + f"{slot_indices[is_unwritten].tolist()}" + ) + + version_changes_at: List[int] = ( + (version_ids[1:] != version_ids[:-1]).nonzero().flatten().tolist() + ) + run_starts = [0] + [position + 1 for position in version_changes_at] + run_ends = run_starts[1:] + [len(version_ids)] + run_version_ids: List[int] = version_ids[run_starts].tolist() + + return [ + WeightVersionSpan( + version=self._versions.lookup(version_id), start=start, end=end + ) + for version_id, start, end in zip( + run_version_ids, run_starts, run_ends, strict=True + ) + ] + + +class _StringInterner: + def __init__(self): + self._str_by_id: List[str] = [] + self._id_by_str: Dict[str, int] = {} + + def intern(self, value: str) -> int: + if (value_id := self._id_by_str.get(value)) is not None: + return value_id + + value_id = len(self._str_by_id) + self._str_by_id.append(value) + self._id_by_str[value] = value_id + return value_id + + def lookup(self, value_id: int) -> str: + return self._str_by_id[value_id] diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 9f7f8636d226..437a632441ab 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -87,6 +87,7 @@ KVCacheConfigurator, ) from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator +from sglang.srt.mem_cache.kv_weight_version_tracker import KvWeightVersionRecord from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool from sglang.srt.model_executor.cuda_graph_config import ( Backend, @@ -286,6 +287,7 @@ class ModelRunnerOutput: expert_distribution_metrics: Optional[ExpertDistributionMetrics] = None routed_experts_output: Optional[TopkCaptureOutput] = None indexer_topk_output: Optional[TopkCaptureOutput] = None + kv_weight_version_record: Optional[KvWeightVersionRecord] = None def resolve_draft_attention_backend( @@ -1806,6 +1808,10 @@ def forward( no_copy_to_cpu=no_copy_to_cpu, ) + output.kv_weight_version_record = KvWeightVersionRecord.maybe_capture( + model_runner=self, forward_batch=forward_batch + ) + # should_advance_eplb_counter is always True on non-hip platform if self.eplb_manager is not None and should_advance_eplb_counter(forward_batch): self.eplb_manager.on_forward_pass_end() diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index 35361a58a518..f570abe4b789 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -2896,4 +2896,5 @@ def forward_batch_generation( new_seq_lens=new_seq_lens, routed_experts_output=target_out.routed_experts_output, indexer_topk_output=target_out.indexer_topk_output, + kv_weight_version_record=target_out.kv_weight_version_record, ) diff --git a/python/sglang/srt/speculative/eagle_worker_common.py b/python/sglang/srt/speculative/eagle_worker_common.py index 0c2c53086020..f390e7d6cf4b 100644 --- a/python/sglang/srt/speculative/eagle_worker_common.py +++ b/python/sglang/srt/speculative/eagle_worker_common.py @@ -671,5 +671,6 @@ def run_eagle_verify( new_seq_lens=new_seq_lens, routed_experts_output=forward_batch_output.routed_experts_output, indexer_topk_output=forward_batch_output.indexer_topk_output, + kv_weight_version_record=forward_batch_output.kv_weight_version_record, extra_keep_alive_refs=[verify_forward_batch], ) diff --git a/python/sglang/srt/utils/weight_versions.py b/python/sglang/srt/utils/weight_versions.py index 6ea1db20a91b..c64cd3479b2c 100644 --- a/python/sglang/srt/utils/weight_versions.py +++ b/python/sglang/srt/utils/weight_versions.py @@ -1,6 +1,7 @@ from __future__ import annotations import dataclasses +import sys from typing import TYPE_CHECKING, Any, Dict, Iterable, List import msgspec @@ -96,17 +97,21 @@ def add_weight_versions_to_meta_info( span for span in spans if span.start < num_output_tokens or span.start == 0 ] - meta_info["weight_versions"] = [ - { - "version": span.version, - "start": span.start, - "end": min(span.end, num_output_tokens), - } - for span in visible - ] + meta_info["weight_versions"] = weight_version_spans_to_json( + visible, end_limit=num_output_tokens + ) meta_info["weight_version"] = visible[-1].version +def weight_version_spans_to_json( + spans: WeightVersionSpans, end_limit: int = sys.maxsize +) -> List[Dict[str, Any]]: + return [ + {"version": span.version, "start": span.start, "end": min(span.end, end_limit)} + for span in spans + ] + + # ====================================================================== # OpenAI-compatible endpoints # ====================================================================== @@ -114,4 +119,6 @@ def build_endpoint_weight_version_metadata(meta_info: Dict[str, Any]) -> Dict[st metadata = {"weight_version": meta_info["weight_version"]} if "weight_versions" in meta_info: metadata["weight_versions"] = meta_info["weight_versions"] + if "prefill_weight_versions" in meta_info: + metadata["prefill_weight_versions"] = meta_info["prefill_weight_versions"] return metadata diff --git a/test/registered/rl/test_weight_version_spans.py b/test/registered/rl/test_weight_version_spans.py index f5fa8d4fe3f8..4ded05a71b6e 100644 --- a/test/registered/rl/test_weight_version_spans.py +++ b/test/registered/rl/test_weight_version_spans.py @@ -2,6 +2,7 @@ import time import unittest from concurrent.futures import ThreadPoolExecutor +from typing import Optional import requests @@ -15,21 +16,37 @@ popen_launch_server, ) -register_cuda_ci( - est_time=180, stage="nightly", runner_config="2-gpu-large", nightly=True -) +register_cuda_ci(est_time=180, stage="nightly", runner_config="2-gpu-large") _REQUEST_TIMEOUT = 180 +_LONG_GENERATION_TOKENS = 3072 + +_SHARED_PREFIX = ( + "You are a meticulous assistant. Follow every instruction to the letter, " + "explain your reasoning step by step, and never skip an intermediate result. " + "Here is the background material you must rely on for the whole conversation. " +) def _assert_spans_contiguous(test, meta_info): - spans = meta_info["weight_versions"] + spans = _assert_contiguous(test, meta_info["weight_versions"]) + test.assertEqual(meta_info["weight_version"], spans[-1]["version"]) + return spans + + +def _assert_prefill_spans_contiguous(test, meta_info, prompt_tokens=None): + spans = _assert_contiguous(test, meta_info["prefill_weight_versions"]) + if prompt_tokens is not None: + test.assertEqual(spans[-1]["end"], prompt_tokens) + return spans + + +def _assert_contiguous(test, spans): test.assertGreater(len(spans), 0) test.assertEqual(spans[0]["start"], 0) for prev, cur in zip(spans, spans[1:]): test.assertEqual(prev["end"], cur["start"]) test.assertNotEqual(prev["version"], cur["version"]) - test.assertEqual(meta_info["weight_version"], spans[-1]["version"]) return spans @@ -65,6 +82,7 @@ def setUpClass(cls): "32", "--max-running-requests", "8", + "--enable-prefill-weight-versions", ], ) @@ -72,7 +90,13 @@ def setUpClass(cls): def tearDownClass(cls): kill_process_tree(cls.process.pid) - def _generate(self, max_new_tokens: int, prompt: str = "The capital of France is"): + def _generate( + self, + max_new_tokens: int, + prompt: str = "The capital of France is", + dp_rank: Optional[int] = None, + extra_key: Optional[str] = None, + ): response = requests.post( f"{self.base_url}/generate", json={ @@ -82,6 +106,8 @@ def _generate(self, max_new_tokens: int, prompt: str = "The capital of France is "max_new_tokens": max_new_tokens, "ignore_eos": True, }, + "routed_dp_rank": dp_rank, + "extra_key": extra_key, }, timeout=_REQUEST_TIMEOUT, ) @@ -107,6 +133,14 @@ def _continue(self): timeout=30, ).raise_for_status() + def _abort_all(self): + requests.post( + f"{self.base_url}/abort_request", json={"abort_all": True}, timeout=30 + ).raise_for_status() + + def _flush_cache(self): + requests.post(f"{self.base_url}/flush_cache", timeout=30).raise_for_status() + def _set_weight_version(self, new_version: str, abort_all_requests: bool = False): response = requests.post( f"{self.base_url}/update_weight_version", @@ -133,7 +167,7 @@ def _run_while_paused( num_requests: int, while_paused, mode: str = "retract", - max_new_tokens: int = 1024, + max_new_tokens: int = _LONG_GENERATION_TOKENS, ): with ThreadPoolExecutor(max_workers=num_requests) as executor: futures = [ @@ -164,6 +198,11 @@ def test_01_single_span_without_update(self): self.assertEqual(spans[0]["version"], "base-v0") self.assertEqual(spans[0]["end"], meta_info["completion_tokens"]) self.assertEqual(meta_info["weight_version"], "base-v0") + prefill_spans = _assert_prefill_spans_contiguous( + self, meta_info, prompt_tokens=meta_info["prompt_tokens"] + ) + self.assertEqual(len(prefill_spans), 1) + self.assertEqual(prefill_spans[0]["version"], "base-v0") def test_02_update_weight_version_endpoint_applies_to_new_requests(self): """The endpoint returns only once every scheduler stamps new requests with the new version.""" @@ -181,6 +220,10 @@ def test_02_update_weight_version_endpoint_applies_to_new_requests(self): spans = _assert_spans_contiguous(self, meta_info) self.assertEqual(len(spans), 1) self.assertEqual(spans[0]["version"], "endpoint-v1") + for prefill_span in _assert_prefill_spans_contiguous( + self, meta_info, prompt_tokens=meta_info["prompt_tokens"] + ): + self.assertIn(prefill_span["version"], ("base-v0", "endpoint-v1")) def test_03_spans_split_across_pause_update_continue(self): """Requests spanning pause -> update_weights_from_disk -> continue report one span per version.""" @@ -201,6 +244,9 @@ def test_03_spans_split_across_pause_update_continue(self): versions = [span["version"] for span in spans] self.assertEqual(versions[0], base_version) self.assertIn(versions[-1], (base_version, "disk-v2")) + _assert_prefill_spans_contiguous( + self, meta_info, prompt_tokens=meta_info["prompt_tokens"] + ) if len(spans) > 1: multi_span_count += 1 self.assertEqual(versions, [base_version, "disk-v2"]) @@ -235,6 +281,9 @@ def test_04_openai_metadata_contains_weight_versions(self): self.assertEqual(metadata["weight_version"], self._current_version()) self.assertEqual(spans[0]["start"], 0) self.assertEqual(spans[0]["end"], data["usage"]["completion_tokens"]) + _assert_prefill_spans_contiguous( + self, metadata, prompt_tokens=data["usage"]["prompt_tokens"] + ) def test_04b_openai_metadata_reports_the_first_choice_when_n_is_greater_than_one( self, @@ -261,18 +310,14 @@ def test_04b_openai_metadata_reports_the_first_choice_when_n_is_greater_than_one self.assertEqual(spans[0]["version"], metadata["weight_version"]) self.assertEqual(metadata["weight_version"], self._current_version()) self.assertEqual(spans[0]["start"], 0) + _assert_prefill_spans_contiguous( + self, metadata, prompt_tokens=data["usage"]["prompt_tokens"] + ) def test_05_aborted_retracted_requests_report_spans(self): """Requests aborted while retracted in the waiting queue still report their spans.""" - def abort_all(): - requests.post( - f"{self.base_url}/abort_request", - json={"abort_all": True}, - timeout=30, - ).raise_for_status() - - results = self._run_while_paused(num_requests=4, while_paused=abort_all) + results = self._run_while_paused(num_requests=4, while_paused=self._abort_all) aborted_with_spans = 0 for data in results: @@ -299,7 +344,6 @@ def test_06_running_requests_split_without_abort(self): num_requests=4, while_paused=lambda: self._set_weight_version("inplace-v3"), mode="in_place", - max_new_tokens=1024, ) self.assertEqual(self._current_version(), "inplace-v3") @@ -310,6 +354,10 @@ def test_06_running_requests_split_without_abort(self): spans = _assert_spans_contiguous(self, meta_info) self.assertEqual(spans[-1]["end"], meta_info["completion_tokens"]) self.assertEqual(spans[0]["version"], previous_version) + for prefill_span in _assert_prefill_spans_contiguous( + self, meta_info, prompt_tokens=meta_info["prompt_tokens"] + ): + self.assertNotEqual(prefill_span["version"], "inplace-v3") if len(spans) > 1: split_count += 1 self.assertEqual( @@ -339,9 +387,13 @@ def test_07_update_without_weight_version_does_not_split(self): self.assertEqual(len(spans), 1) self.assertEqual(spans[0]["version"], version) self.assertEqual(spans[0]["end"], meta_info["completion_tokens"]) + _assert_prefill_spans_contiguous( + self, meta_info, prompt_tokens=meta_info["prompt_tokens"] + ) def test_08_reannouncing_current_version_is_a_noop(self): """Re-announcing the version the server already has must not split anything.""" + self._flush_cache() version = self._current_version() results = self._run_while_paused( @@ -353,9 +405,14 @@ def test_08_reannouncing_current_version_is_a_noop(self): self.assertEqual(self._current_version(), version) for data in results: - spans = _assert_spans_contiguous(self, data["meta_info"]) + meta_info = data["meta_info"] + spans = _assert_spans_contiguous(self, meta_info) self.assertEqual(len(spans), 1) self.assertEqual(spans[0]["version"], version) + for prefill_span in _assert_prefill_spans_contiguous( + self, meta_info, prompt_tokens=meta_info["prompt_tokens"] + ): + self.assertEqual(prefill_span["version"], version) def test_09_three_spans_across_two_updates(self): """Two updates during one request produce three ordered, non-empty spans.""" @@ -365,7 +422,7 @@ def test_09_three_spans_across_two_updates(self): futures = [ executor.submit( self._generate, - max_new_tokens=2048, + max_new_tokens=_LONG_GENERATION_TOKENS, prompt=f"Write a long story about the number {i}.", ) for i in range(4) @@ -391,6 +448,10 @@ def test_09_three_spans_across_two_updates(self): for span in spans: self.assertGreater(span["end"], span["start"]) self.assertEqual(spans[-1]["end"], meta_info["completion_tokens"]) + for prefill_span in _assert_prefill_spans_contiguous( + self, meta_info, prompt_tokens=meta_info["prompt_tokens"] + ): + self.assertNotIn(prefill_span["version"], ("multi-a", "multi-b")) if len(spans) == 3: three_span_count += 1 @@ -404,16 +465,9 @@ def test_10_abort_before_first_token_reports_empty_span(self): """A request aborted before producing a token still reports a well-formed span.""" version = self._current_version() - def abort_all(): - requests.post( - f"{self.base_url}/abort_request", - json={"abort_all": True}, - timeout=30, - ).raise_for_status() - results = self._run_while_paused( num_requests=16, - while_paused=abort_all, + while_paused=self._abort_all, ) empty_aborts = [ @@ -429,6 +483,7 @@ def abort_all(): [{"version": version, "start": 0, "end": 0}], ) self.assertEqual(meta_info["weight_version"], version) + self.assertNotIn("prefill_weight_versions", meta_info) def test_11_streaming_reports_spans_only_on_the_final_chunk(self): """Intermediate stream chunks carry no spans; the finishing chunk carries them all.""" @@ -462,12 +517,16 @@ def test_11_streaming_reports_spans_only_on_the_final_chunk(self): self.assertGreater(len(chunks), 1) for chunk in chunks[:-1]: self.assertNotIn("weight_versions", chunk["meta_info"]) + self.assertNotIn("prefill_weight_versions", chunk["meta_info"]) self.assertEqual(chunk["meta_info"]["weight_version"], version) meta_info = chunks[-1]["meta_info"] spans = _assert_spans_contiguous(self, meta_info) self.assertEqual(spans[-1]["end"], meta_info["completion_tokens"]) self.assertEqual(spans[-1]["end"], max_new_tokens) + _assert_prefill_spans_contiguous( + self, meta_info, prompt_tokens=meta_info["prompt_tokens"] + ) def test_12_completions_endpoint_reports_metadata(self): """/v1/completions surfaces the spans the same way /v1/chat/completions does.""" @@ -491,6 +550,9 @@ def test_12_completions_endpoint_reports_metadata(self): self.assertEqual(spans[0]["version"], metadata["weight_version"]) self.assertEqual(spans[0]["start"], 0) self.assertEqual(spans[0]["end"], data["usage"]["completion_tokens"]) + _assert_prefill_spans_contiguous( + self, metadata, prompt_tokens=data["usage"]["prompt_tokens"] + ) def test_13_new_requests_after_the_endpoint_returns_see_the_new_version(self): """Once the endpoint returns, every concurrently admitted request stamps the new version.""" @@ -570,6 +632,9 @@ def test_16_retract_update_continue_keeps_exact_boundaries(self): spans = _assert_spans_contiguous(self, meta_info) self.assertEqual(spans[-1]["end"], meta_info["completion_tokens"]) self.assertEqual(spans[0]["version"], previous_version) + _assert_prefill_spans_contiguous( + self, meta_info, prompt_tokens=meta_info["prompt_tokens"] + ) if len(spans) > 1: split_count += 1 self.assertEqual( @@ -585,6 +650,192 @@ def test_16_retract_update_continue_keeps_exact_boundaries(self): "speculative decoding was not recorded.", ) + def test_17_reused_prefix_keeps_the_version_that_computed_it(self): + """A radix hit on KV computed before a relabel shows up as a stale leading prefill span.""" + self._flush_cache() + self._generate( + max_new_tokens=8, prompt=_SHARED_PREFIX + "Name three colors.", dp_rank=0 + ) + version_before = self._current_version() + + self._set_weight_version("reuse-v4") + data = self._generate( + max_new_tokens=8, + prompt=_SHARED_PREFIX + "Name three colors. Then name three shapes.", + dp_rank=0, + ) + + meta_info = data["meta_info"] + self.assertGreater(meta_info["cached_tokens"], 0) + prefill_spans = _assert_prefill_spans_contiguous( + self, meta_info, prompt_tokens=meta_info["prompt_tokens"] + ) + self.assertEqual( + [span["version"] for span in prefill_spans], [version_before, "reuse-v4"] + ) + self.assertEqual(prefill_spans[0]["end"], meta_info["cached_tokens"]) + spans = _assert_spans_contiguous(self, meta_info) + self.assertEqual([span["version"] for span in spans], ["reuse-v4"]) + + def test_18_flushed_cache_recomputes_the_prefix(self): + """After a flush the whole prompt is recomputed under the current version.""" + self._flush_cache() + data = self._generate(max_new_tokens=8, prompt=_SHARED_PREFIX + "Say hello.") + + meta_info = data["meta_info"] + self.assertEqual(meta_info["cached_tokens"], 0) + prefill_spans = _assert_prefill_spans_contiguous( + self, meta_info, prompt_tokens=meta_info["prompt_tokens"] + ) + self.assertEqual(len(prefill_spans), 1) + self.assertEqual(prefill_spans[0]["version"], self._current_version()) + + def test_19_aborted_after_prefill_reports_both_spans(self): + """A request aborted mid-decode keeps its prompt spans and reports output spans.""" + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit( + self._generate, + max_new_tokens=_LONG_GENERATION_TOKENS, + prompt=_SHARED_PREFIX + "Write an even longer story.", + ) + time.sleep(1) + self._pause("in_place") + try: + self._abort_all() + finally: + self._continue() + data = future.result() + + meta_info = data["meta_info"] + self.assertEqual(meta_info["finish_reason"]["type"], "abort") + self.assertGreater(meta_info["completion_tokens"], 0) + spans = _assert_spans_contiguous(self, meta_info) + self.assertEqual(spans[-1]["end"], meta_info["completion_tokens"]) + _assert_prefill_spans_contiguous(self, meta_info) + + def test_20_prefix_grown_across_versions_reports_every_version(self): + """A prompt whose prefix was extended under three successive versions reports three prefill spans.""" + self._flush_cache() + self._set_weight_version("grow-v1") + self._generate(max_new_tokens=8, prompt=_SHARED_PREFIX + "Step one.", dp_rank=0) + self._set_weight_version("grow-v2") + self._generate( + max_new_tokens=8, + prompt=_SHARED_PREFIX + "Step one. Step two follows the first one.", + dp_rank=0, + ) + self._set_weight_version("grow-v3") + + data = self._generate( + max_new_tokens=8, + prompt=_SHARED_PREFIX + + "Step one. Step two follows the first one. Step three ends it.", + dp_rank=0, + ) + + meta_info = data["meta_info"] + prefill_spans = _assert_prefill_spans_contiguous( + self, meta_info, prompt_tokens=meta_info["prompt_tokens"] + ) + self.assertEqual( + [span["version"] for span in prefill_spans], + ["grow-v1", "grow-v2", "grow-v3"], + ) + self.assertEqual(prefill_spans[1]["end"], meta_info["cached_tokens"]) + self.assertEqual( + [span["version"] for span in _assert_spans_contiguous(self, meta_info)], + ["grow-v3"], + ) + + def test_21_multi_turn_reuses_the_previous_turn_under_its_version(self): + """The second turn of a conversation reports the first turn's prompt and reply as a stale span.""" + self._flush_cache() + self._set_weight_version("turn-v1") + first_prompt = _SHARED_PREFIX + "User: Name a planet.\nAssistant:" + first = self._generate(max_new_tokens=16, prompt=first_prompt, dp_rank=0) + self.assertGreater(first["meta_info"]["spec_verify_ct"], 0) + self.assertEqual( + [span["version"] for span in first["meta_info"]["prefill_weight_versions"]], + ["turn-v1"], + ) + + self._set_weight_version("turn-v2") + second = self._generate( + max_new_tokens=16, + prompt=first_prompt + + first["text"] + + "\nUser: Name another one.\nAssistant:", + dp_rank=0, + ) + + meta_info = second["meta_info"] + self.assertGreater(meta_info["spec_verify_ct"], 0) + self.assertGreater( + meta_info["cached_tokens"], first["meta_info"]["prompt_tokens"] + ) + prefill_spans = _assert_prefill_spans_contiguous( + self, meta_info, prompt_tokens=meta_info["prompt_tokens"] + ) + self.assertEqual( + [span["version"] for span in prefill_spans], ["turn-v1", "turn-v2"] + ) + self.assertEqual(prefill_spans[0]["end"], meta_info["cached_tokens"]) + self.assertEqual( + [span["version"] for span in _assert_spans_contiguous(self, meta_info)], + ["turn-v2"], + ) + + def test_22_the_other_dp_rank_recomputes_the_prefix(self): + """A prefix cached on one DP rank is not shared, so the other rank prefills it fresh.""" + self._flush_cache() + self._set_weight_version("rank-v1") + prompt = _SHARED_PREFIX + "Which rank served me?" + self._generate(max_new_tokens=8, prompt=prompt, dp_rank=0) + self._set_weight_version("rank-v2") + + cached = self._generate(max_new_tokens=8, prompt=prompt, dp_rank=0) + fresh = self._generate(max_new_tokens=8, prompt=prompt, dp_rank=1) + + self.assertEqual( + [ + span["version"] + for span in cached["meta_info"]["prefill_weight_versions"] + ], + ["rank-v1", "rank-v2"], + ) + self.assertEqual(fresh["meta_info"]["cached_tokens"], 0) + self.assertEqual( + [span["version"] for span in fresh["meta_info"]["prefill_weight_versions"]], + ["rank-v2"], + ) + + def test_23_extra_key_partitions_stay_apart(self): + """A prefix cached under one extra_key is invisible to another, so only the same key sees the stale span.""" + self._flush_cache() + self._set_weight_version("key-v1") + prompt = _SHARED_PREFIX + "Which partition served me?" + self._generate(max_new_tokens=8, prompt=prompt, dp_rank=0, extra_key="alpha") + self._set_weight_version("key-v2") + + same_key = self._generate( + max_new_tokens=8, prompt=prompt, dp_rank=0, extra_key="alpha" + ) + other_key = self._generate( + max_new_tokens=8, prompt=prompt, dp_rank=0, extra_key="beta" + ) + no_key = self._generate(max_new_tokens=8, prompt=prompt, dp_rank=0) + + self.assertEqual( + [s["version"] for s in same_key["meta_info"]["prefill_weight_versions"]], + ["key-v1", "key-v2"], + ) + for data in (other_key, no_key): + self.assertEqual(data["meta_info"]["cached_tokens"], 0) + self.assertEqual( + [s["version"] for s in data["meta_info"]["prefill_weight_versions"]], + ["key-v2"], + ) + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/hardware_backend/mlx/test_attention_patching.py b/test/registered/unit/hardware_backend/mlx/test_attention_patching.py index 38964a655d9d..1a673cb62849 100644 --- a/test/registered/unit/hardware_backend/mlx/test_attention_patching.py +++ b/test/registered/unit/hardware_backend/mlx/test_attention_patching.py @@ -1203,6 +1203,7 @@ def test_finished_request_snapshots_before_release(self): tree_cache=tree_cache, hisparse_coordinator=SimpleNamespace(request_finished=lambda req: None), req_to_token_pool=None, + kv_weight_version_tracker=None, decode_offload_manager=None, metrics_collector=None, metrics_reporter=None, 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 103432afbd62..e5d159b02b4d 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 @@ -39,6 +39,7 @@ def _make_processor(case, server_mode: str = "full") -> SchedulerBatchResultProc tree_cache=None, hisparse_coordinator=None, req_to_token_pool=None, + kv_weight_version_tracker=None, decode_offload_manager=None, metrics_collector=None, metrics_reporter=metrics_reporter, @@ -236,6 +237,7 @@ def test_active_middle_chunk_advances_before_new_last_request(self): auxiliary_host_output=None, routed_experts_output=None, indexer_topk_output=None, + kv_weight_version_record=None, logits_output=SimpleNamespace( hidden_states=hidden_states, customized_info=None, @@ -291,6 +293,7 @@ def test_sampling_mask_middle_chunk_does_not_require_logits_output(self): routed_experts_output=None, indexer_topk_output=None, logits_output=None, + kv_weight_version_record=None, next_token_ids=torch.zeros(1, dtype=torch.int64), extend_input_len_per_req=None, extend_logprob_start_len_per_req=None, diff --git a/test/registered/unit/managers/test_batch_result_processor_mamba_boundary.py b/test/registered/unit/managers/test_batch_result_processor_mamba_boundary.py index 97eb071fb1f0..ffce27861120 100644 --- a/test/registered/unit/managers/test_batch_result_processor_mamba_boundary.py +++ b/test/registered/unit/managers/test_batch_result_processor_mamba_boundary.py @@ -67,6 +67,7 @@ def _make_processor() -> SchedulerBatchResultProcessor: hisparse_coordinator=None, beam_coordinator=MagicMock(), req_to_token_pool=None, + kv_weight_version_tracker=None, decode_offload_manager=None, metrics_collector=None, metrics_reporter=metrics_reporter, diff --git a/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py b/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py index 6fc1a88c424e..0d69add93a1f 100644 --- a/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py +++ b/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py @@ -5,6 +5,7 @@ """ import unittest +from dataclasses import replace from types import SimpleNamespace import torch @@ -15,6 +16,7 @@ SchedulerBatchResultProcessor, ) from sglang.srt.managers.utils import GenerationBatchResult +from sglang.srt.mem_cache.kv_weight_version_tracker import KvWeightVersionTracker from sglang.srt.sampling.sampling_params import ( REQUEST_REASONING_END_TOKEN_IDS_KEY, SamplingParams, @@ -75,6 +77,7 @@ def _make_processor() -> SchedulerBatchResultProcessor: tree_cache=None, hisparse_coordinator=None, req_to_token_pool=None, + kv_weight_version_tracker=None, decode_offload_manager=None, metrics_collector=None, metrics_reporter=SimpleNamespace(), @@ -154,6 +157,40 @@ def _commit_disagg_handoff( class TestSpecV2GrammarTruncation(CustomTestCase): + def test_grammar_commit_does_not_relabel_forward_recorded_kv(self) -> None: + """Grammar commitment leaves forward-recorded provenance unchanged.""" + for overlap in (False, True): + with self.subTest(overlap=overlap): + pool = SimpleNamespace(req_to_token=torch.tensor([[1, 2, 3, 7, 5, 6]])) + tracker = KvWeightVersionTracker( + num_slots=10, device="cpu", req_to_token_pool=pool + ) + tracker.record( + slot_indices=torch.tensor([1, 2, 3, 7, 5, 6]), version="v0" + ) + tracker.record(slot_indices=torch.tensor([7, 5, 6]), version="v1") + req = _make_req(terminate_after=2) + req.kv.req_pool_idx = 0 + req.kv.kv_committed_len = 3 + processor = replace( + _make_processor(), + enable_overlap=overlap, + req_to_token_pool=pool, + kv_weight_version_tracker=tracker, + ) + batch = _FakeBatch([req]) + batch.out_cache_loc = torch.tensor([1, 2, 3]) + result = _make_result(4, [3], [101, 102, 103, 0]) + + processor._resolve_spec_v2_tokens(result, batch) + + spans = tracker._lookup_spans(torch.tensor([1, 2, 3, 7, 5, 6])) + self.assertEqual( + [(span.version, span.start, span.end) for span in spans], + [("v0", 0, 3), ("v1", 3, 6)], + ) + self.assertEqual(req.kv.kv_committed_len, 5) + def test_resolve_truncates_after_grammar_completion(self): req = _make_req(terminate_after=2) proc = _make_processor() diff --git a/test/registered/unit/managers/test_msgpack_ipc_roundtrip.py b/test/registered/unit/managers/test_msgpack_ipc_roundtrip.py index 65a28ebb9660..99f5ebd8304d 100644 --- a/test/registered/unit/managers/test_msgpack_ipc_roundtrip.py +++ b/test/registered/unit/managers/test_msgpack_ipc_roundtrip.py @@ -178,12 +178,13 @@ def test_backup_dram_req_is_narrowed(self): pointer = decoded.weight_pointer_map["experts.0.gate_proj"] self.assertEqual((pointer.weight_ptr, pointer.byte_size), (8, 4)) - def test_check_weights_mirrors_match_pydantic_models(self): + def test_check_weights_mirrors_match_pydantic_models(self) -> None: + """Wire fields preserve checker metadata and add the merged runner role.""" # Field-parity guard: the msgspec wire structs must not drift from the # pydantic source of truth in weight_checker. self.assertEqual( {f.name for f in msgspec.structs.fields(ParallelismInfo)}, - set(PydanticParallelismInfo.model_fields), + set(PydanticParallelismInfo.model_fields) | {"role"}, ) self.assertEqual( {f.name for f in msgspec.structs.fields(ChecksumInfo)}, @@ -282,6 +283,28 @@ def test_abort_req_carries_spans(self): def test_abort_req_defaults_to_no_spans(self): """The field is optional on the wire, so an abort without spans decodes to None.""" self.assertIsNone(_round_trip(AbortReq(rid="r0")).weight_versions) + self.assertIsNone(_round_trip(AbortReq(rid="r0")).prefill_weight_versions) + + def test_abort_req_carries_prefill_spans(self): + """Prompt-side spans survive the wire independently of the sampling spans.""" + obj = AbortReq( + rid="r0", + prefill_weight_versions=[ + WeightVersionSpan(version="unknown", start=0, end=2), + WeightVersionSpan(version="v1", start=2, end=9), + ], + ) + + decoded = _double_hop(obj) + + self.assertIsNone(decoded.weight_versions) + self.assertEqual( + decoded.prefill_weight_versions, + [ + WeightVersionSpan(version="unknown", start=0, end=2), + WeightVersionSpan(version="v1", start=2, end=9), + ], + ) if __name__ == "__main__": diff --git a/test/registered/unit/managers/test_multi_tokenizer_mixin.py b/test/registered/unit/managers/test_multi_tokenizer_mixin.py index 02a05ca209e5..b73251bc9a09 100644 --- a/test/registered/unit/managers/test_multi_tokenizer_mixin.py +++ b/test/registered/unit/managers/test_multi_tokenizer_mixin.py @@ -97,6 +97,10 @@ def _make_batch_str_output() -> BatchStrOutput: ], [WeightVersionSpan(version="v2", start=0, end=2)], ], + prefill_weight_versions=[ + [WeightVersionSpan(version="v0", start=0, end=7)], + None, + ], ) @@ -148,6 +152,25 @@ def test_batch_str_output_without_weight_versions_stays_none(self): self.assertIsNone(_handle_output_by_index(output, 0).weight_versions) + def test_batch_str_output_keeps_prefill_weight_versions_nested_per_request(self): + """Prefill spans split per request the same way the sampling spans do.""" + output = _make_batch_str_output() + + self.assertEqual( + _handle_output_by_index(output, 0).prefill_weight_versions, + [[WeightVersionSpan(version="v0", start=0, end=7)]], + ) + self.assertEqual( + _handle_output_by_index(output, 1).prefill_weight_versions, [None] + ) + + def test_batch_str_output_without_prefill_weight_versions_stays_none(self): + """An output produced without the flag splits into None.""" + output = _make_batch_str_output() + output.prefill_weight_versions = None + + self.assertIsNone(_handle_output_by_index(output, 0).prefill_weight_versions) + def test_get_tokenizer_worker_class_uses_default(self): self.assertIs(get_tokenizer_worker_class(DefaultServerArgs()), TokenizerWorker) 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 8817dc762561..d02e58dd8e1c 100644 --- a/test/registered/unit/managers/test_output_streamer_customized_info.py +++ b/test/registered/unit/managers/test_output_streamer_customized_info.py @@ -83,6 +83,7 @@ def __init__( self.multimodal_inputs = None self.customized_info = customized_info self.weight_version_events = [] + self.prefill_weight_versions = None def finished(self): return self._finished @@ -537,5 +538,54 @@ def test_clients_receive_the_sampler_rows_unchanged(self): self.assertEqual(received, expected) +class TestOutputStreamerPrefillWeightVersions(unittest.TestCase): + def test_payload_carries_prefill_spans_for_finished_requests(self): + """Prefill spans computed before the KV release ride out with the finished request.""" + streaming_req = _FakeReq("r0", [10, 11]) + finished_req = _FakeReq("r1", [20], finished=True) + finished_req.prefill_weight_versions = [ + WeightVersionSpan(version="v0", start=0, end=4), + WeightVersionSpan(version="v1", start=4, end=6), + ] + + accumulator = _accumulator(current_weight_version="v1") + accumulator.accept(req=streaming_req) + accumulator.accept(req=finished_req) + payload = accumulator.to_payload(dp_rank=0, is_idle_batch=False) + + self.assertEqual( + payload.prefill_weight_versions, + [ + None, + [ + WeightVersionSpan(version="v0", start=0, end=4), + WeightVersionSpan(version="v1", start=4, end=6), + ], + ], + ) + + def test_an_unfinished_request_withholds_its_prefill_spans(self): + """Spans of a request still generating must not ship before its finishing chunk.""" + streaming_req = _FakeReq("r0", [10, 11]) + streaming_req.prefill_weight_versions = [ + WeightVersionSpan(version="v0", start=0, end=4) + ] + + accumulator = _accumulator(current_weight_version="v1") + accumulator.accept(req=streaming_req) + payload = accumulator.to_payload(dp_rank=0, is_idle_batch=False) + + self.assertIsNone(payload.prefill_weight_versions) + + def test_payload_omits_prefill_spans_when_the_flag_is_off(self): + """With tracking disabled every request contributes None, so nothing goes on the wire.""" + accumulator = _accumulator(current_weight_version="v1") + accumulator.accept(req=_FakeReq("r0", [10], finished=True)) + payload = accumulator.to_payload(dp_rank=0, is_idle_batch=False) + + self.assertIsNone(payload.prefill_weight_versions) + self.assertIsNotNone(payload.weight_versions) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/managers/test_scheduler_weight_version_tracking.py b/test/registered/unit/managers/test_scheduler_weight_version_tracking.py index 8eb58463dd79..7e3c79651eb2 100644 --- a/test/registered/unit/managers/test_scheduler_weight_version_tracking.py +++ b/test/registered/unit/managers/test_scheduler_weight_version_tracking.py @@ -1,16 +1,31 @@ import unittest from types import SimpleNamespace -from unittest.mock import Mock, patch +from unittest.mock import MagicMock, Mock, patch +import torch + +from sglang.srt.disaggregation.utils import DisaggregationMode from sglang.srt.managers.io_struct import ( BeginWeightUpdateReqInput, EndWeightUpdateReqInput, ) +from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.scheduler import Scheduler +from sglang.srt.managers.scheduler_components.batch_result_processor import ( + SchedulerBatchResultProcessor, +) from sglang.srt.managers.scheduler_components.weight_updater import ( SchedulerWeightUpdaterManager, _WeightUpdateSession, ) +from sglang.srt.managers.utils import GenerationBatchResult +from sglang.srt.mem_cache.kv_weight_version_tracker import ( + KvWeightVersionRecord, + KvWeightVersionTracker, +) +from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode, ForwardMode +from sglang.srt.model_executor.model_runner import ModelRunner +from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -101,6 +116,203 @@ def test_every_source_of_live_requests_is_stamped(self): ) +class TestSchedulerBatchWeightVersion(CustomTestCase): + def test_embedding_rejects_enabled_prefill_tracking(self) -> None: + """Embedding and reward models reject enabled prefill tracking before allocation.""" + scheduler = Scheduler.__new__(Scheduler) + scheduler.token_to_kv_pool_allocator = None + scheduler.req_to_token_pool = None + scheduler.server_args = SimpleNamespace(enable_prefill_weight_versions=True) + scheduler.model_config = SimpleNamespace( + is_encoder_decoder=False, is_generation=False + ) + + with self.assertRaisesRegex( + AssertionError, "does not support embedding or reward" + ): + scheduler.init_kv_weight_version_tracker() + + def test_embedding_without_prefill_tracking_remains_supported(self) -> None: + """Embedding initialization remains unchanged when prefill tracking is disabled.""" + scheduler = Scheduler.__new__(Scheduler) + scheduler.token_to_kv_pool_allocator = None + scheduler.req_to_token_pool = None + scheduler.server_args = SimpleNamespace(enable_prefill_weight_versions=False) + scheduler.model_config = SimpleNamespace( + is_encoder_decoder=False, is_generation=False + ) + + scheduler.init_kv_weight_version_tracker() + + self.assertIsNone(scheduler.kv_weight_version_tracker) + + def test_encoder_decoder_rejects_enabled_prefill_tracking(self) -> None: + """Encoder-decoder models reject prefill tracking before allocating a tracker.""" + scheduler = Scheduler.__new__(Scheduler) + scheduler.token_to_kv_pool_allocator = None + scheduler.req_to_token_pool = None + scheduler.server_args = SimpleNamespace(enable_prefill_weight_versions=True) + scheduler.model_config = SimpleNamespace(is_encoder_decoder=True) + + with self.assertRaisesRegex(AssertionError, "does not support encoder-decoder"): + scheduler.init_kv_weight_version_tracker() + + def test_encoder_decoder_without_prefill_tracking_remains_supported(self) -> None: + """Disabled tracking leaves encoder-decoder initialization unchanged.""" + scheduler = Scheduler.__new__(Scheduler) + scheduler.token_to_kv_pool_allocator = None + scheduler.req_to_token_pool = None + scheduler.server_args = SimpleNamespace(enable_prefill_weight_versions=False) + scheduler.model_config = SimpleNamespace(is_encoder_decoder=True) + + scheduler.init_kv_weight_version_tracker() + + self.assertIsNone(scheduler.kv_weight_version_tracker) + + def _scheduler(self) -> Scheduler: + scheduler = Scheduler.__new__(Scheduler) + scheduler.token_to_kv_pool_allocator = None + scheduler.req_to_token_pool = None + scheduler.publish_load_snapshot = MagicMock() + scheduler.scheduler_stage_metrics = None + scheduler.load_publisher = MagicMock() + scheduler.load_inquirer = MagicMock() + scheduler.tree_cache = SimpleNamespace(flush_pending_backups=MagicMock()) + scheduler.kv_weight_version_tracker = KvWeightVersionTracker( + num_slots=8, + device="cpu", + req_to_token_pool=SimpleNamespace(req_to_token=torch.tensor([[3]])), + ) + scheduler.batch_result_processor = SchedulerBatchResultProcessor( + is_generation=True, + disaggregation_mode=DisaggregationMode.NULL, + enable_overlap=False, + enable_overlap_mlx=False, + model_config=SimpleNamespace(think_end_ids=None), + token_to_kv_pool_allocator=MagicMock(), + tree_cache=scheduler.tree_cache, + hisparse_coordinator=None, + req_to_token_pool=None, + kv_weight_version_tracker=scheduler.kv_weight_version_tracker, + decode_offload_manager=None, + metrics_collector=None, + metrics_reporter=MagicMock(), + draft_worker=None, + model_worker=MagicMock(), + logprob_result_processor=None, + output_streamer=MagicMock(), + beam_coordinator=MagicMock(), + abort_request=lambda *args, **kwargs: None, + ) + for name, value in ( + ("get_observability", SimpleNamespace(enable_metrics=False)), + ("get_server_return_hidden_states_mode", CaptureHiddenMode.NULL), + ): + patcher = patch( + f"sglang.srt.managers.scheduler_components.batch_result_processor.{name}", + return_value=value, + ) + patcher.start() + self.addCleanup(patcher.stop) + scheduler.disaggregation_mode = DisaggregationMode.NULL + scheduler._record_step_counters = MagicMock() + scheduler.metrics_reporter = MagicMock() + scheduler.enable_fpm = False + scheduler._maybe_clear_mm_inputs = MagicMock() + scheduler.maybe_send_health_check_signal = MagicMock() + return scheduler + + def _batch(self) -> ScheduleBatch: + return ScheduleBatch( + reqs=[], + forward_mode=ForwardMode.EXTEND, + out_cache_loc=torch.tensor([3], dtype=torch.int64), + ) + + def test_pending_result_uses_version_captured_when_batch_ran(self) -> None: + """A pending result keeps the forward-time version across a serving update.""" + serving = _ServingStub("v0") + scheduler = self._scheduler() + batch = self._batch() + + runner = ModelRunner.__new__(ModelRunner) + runner.is_draft_worker = False + runner.server_args = SimpleNamespace(enable_prefill_weight_versions=True) + with patch( + "sglang.srt.mem_cache.kv_weight_version_tracker.get_serving", + return_value=serving, + ): + result = GenerationBatchResult( + next_token_ids=torch.tensor([], dtype=torch.int64), + kv_weight_version_record=KvWeightVersionRecord.maybe_capture( + model_runner=runner, forward_batch=batch + ), + ) + queued_batch = batch.copy() + serving.weight_version = "v1" + Scheduler.process_batch_result(scheduler, queued_batch, result) + + spans = scheduler.kv_weight_version_tracker._lookup_spans( + queued_batch.out_cache_loc + ) + self.assertEqual( + [(span.version, span.start, span.end) for span in spans], [("v0", 0, 1)] + ) + + def test_spec_decode_does_not_restamp_restored_prefill_slots(self) -> None: + """Restoring the batch after speculative forward must not relabel old prompt KV.""" + for overlap in (False, True): + with self.subTest(overlap=overlap): + scheduler = self._scheduler() + scheduler.record_batch_in_overlap = lambda batch: None + batch = self._batch() + batch.forward_mode = ForwardMode.DECODE + batch.spec_algorithm = SpeculativeAlgorithm.EAGLE + tracker = scheduler.kv_weight_version_tracker + tracker.record(slot_indices=batch.out_cache_loc, version="v0") + + with scheduler._forward_isolation(batch, overlap=overlap): + forward_slots = torch.tensor([4, 5, 6]) + batch.out_cache_loc = forward_slots + result = GenerationBatchResult( + next_token_ids=torch.tensor([], dtype=torch.int64), + accept_lens=torch.tensor([], dtype=torch.int32), + speculative_num_draft_tokens=3, + kv_weight_version_record=KvWeightVersionRecord.capture( + slot_indices=forward_slots, version="v1" + ), + ) + forward_slots.fill_(7) + scheduler.process_batch_result(batch.copy(), result) + + spans = tracker._lookup_spans(torch.tensor([3, 4, 5, 6])) + self.assertEqual( + [(span.version, span.start, span.end) for span in spans], + [("v0", 0, 1), ("v1", 1, 4)], + ) + + def test_result_copy_finishes_before_slot_indices_are_consumed(self) -> None: + """Pending slot copies complete before the tracker reads their contents.""" + scheduler = self._scheduler() + batch = self._batch() + record = KvWeightVersionRecord.capture( + slot_indices=torch.tensor([0]), version="v0" + ) + synchronize = MagicMock(side_effect=lambda: record.slot_indices.fill_(3)) + result = GenerationBatchResult( + next_token_ids=torch.tensor([], dtype=torch.int64), + kv_weight_version_record=record, + copy_done=SimpleNamespace(synchronize=synchronize), + ) + + scheduler.process_batch_result(batch, result) + + spans = scheduler.kv_weight_version_tracker._lookup_spans(torch.tensor([3])) + self.assertEqual([span.version for span in spans], ["v0"]) + self.assertIsNone(result.kv_weight_version_record) + synchronize.assert_called_once_with() + + def _runner(result=(True, "ok")): runner = Mock() for method in ( @@ -110,6 +322,9 @@ def _runner(result=(True, "ok")): "load_weights_from_distributed", ): getattr(runner.weight_updater, method).return_value = result + runner.weight_updater.receive_weights_from_distributed.return_value = [ + ("weight", torch.tensor([1.0])) + ] return runner @@ -171,23 +386,27 @@ def _manager(self, target, draft=None, *, session=True): class TestRecordWeightVersionAfterUpdate(_WeightUpdaterManagerTestBase): def test_successful_update_records_the_version(self): """A refit that reports success advances the scheduler-side version.""" - output = self._manager(_runner()).update_weights_from_disk(_request()) + output = self._manager(_runner(), session=False).update_weights_from_disk( + _request() + ) self.assertTrue(output.success) self.assertEqual(self.recorded, ["v2"]) def test_failed_update_does_not_record_the_version(self): """A refit that fails must leave the version alone, or later tokens are mislabelled.""" - output = self._manager(_runner((False, "boom"))).update_weights_from_disk( - _request() - ) + output = self._manager( + _runner((False, "boom")), session=False + ).update_weights_from_disk(_request()) self.assertFalse(output.success) self.assertEqual(self.recorded, []) def test_draft_failure_does_not_record_the_version(self): """The target succeeding is not enough: a failed draft refit leaves the engine mixed.""" - manager = self._manager(_runner(), draft=_runner((False, "draft boom"))) + manager = self._manager( + _runner(), draft=_runner((False, "draft boom")), session=False + ) output = manager.update_weights_from_disk(_request()) @@ -196,7 +415,12 @@ def test_draft_failure_does_not_record_the_version(self): def test_successful_distributed_update_records_the_version(self): """The distributed refit is the path an RL trainer actually drives, so it must record too.""" - output = self._manager(_runner()).update_weights_from_distributed(_request()) + manager = self._manager(_runner()) + output = manager.update_weights_from_distributed(_request()) + + self.assertTrue(output.success) + self.assertEqual(self.recorded, []) + output = manager.end_weight_update(EndWeightUpdateReqInput()) self.assertTrue(output.success) self.assertEqual(self.recorded, ["v2"]) @@ -212,23 +436,31 @@ def test_failed_distributed_update_does_not_record_the_version(self): def test_successful_tensor_update_records_the_version(self): """The tensor refit records the version once the load reports success.""" - output = self._manager(_runner()).update_weights_from_tensor(_request()) + updater = self._manager(_runner()) + output = updater.update_weights_from_tensor(_request()) + + self.assertTrue(output.success) + self.assertEqual(self.recorded, []) + with patch("torch.distributed.barrier"): + output = updater.end_weight_update(EndWeightUpdateReqInput()) self.assertTrue(output.success) self.assertEqual(self.recorded, ["v2"]) def test_successful_ipc_update_records_the_version(self): """The checkpoint-engine IPC refit records the version like every other path.""" - output = self._manager(_runner()).update_weights_from_ipc(_request()) + output = self._manager(_runner(), session=False).update_weights_from_ipc( + _request() + ) self.assertTrue(output.success) self.assertEqual(self.recorded, ["v2"]) def test_failed_ipc_update_does_not_record_the_version(self): """The IPC path branches on success separately from the cache flush, so failure must record nothing.""" - output = self._manager(_runner((False, "boom"))).update_weights_from_ipc( - _request() - ) + output = self._manager( + _runner((False, "boom")), session=False + ).update_weights_from_ipc(_request()) self.assertFalse(output.success) self.assertEqual(self.recorded, []) diff --git a/test/registered/unit/mem_cache/test_kv_weight_version_tracker.py b/test/registered/unit/mem_cache/test_kv_weight_version_tracker.py new file mode 100644 index 000000000000..4b90a25865e7 --- /dev/null +++ b/test/registered/unit/mem_cache/test_kv_weight_version_tracker.py @@ -0,0 +1,336 @@ +import unittest +from types import SimpleNamespace +from typing import List + +import torch + +from sglang.srt.mem_cache.kv_weight_version_tracker import ( + KvWeightVersionRecord, + KvWeightVersionTracker, + _StringInterner, +) +from sglang.srt.utils.weight_versions import WeightVersionSpan +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=4, suite="base-a-test-cpu") + + +class _ReqStub: + def __init__(self, num_prompt_tokens: int, kv_committed_len: int): + self.origin_input_ids = [0] * num_prompt_tokens + self.kv = SimpleNamespace(kv_committed_len=kv_committed_len, req_pool_idx=1) + self.prefill_weight_versions = None + + +class _ReqToTokenPoolStub: + def __init__(self, slots_of_req: List[int]): + self.req_to_token = torch.zeros((2, 32), dtype=torch.int32) + self.req_to_token[1, : len(slots_of_req)] = torch.tensor( + slots_of_req, dtype=torch.int32 + ) + + +def _table(slots_of_req: List[int] = ()) -> KvWeightVersionTracker: + return KvWeightVersionTracker( + num_slots=16, + device="cpu", + req_to_token_pool=_ReqToTokenPoolStub(list(slots_of_req)), + ) + + +def _slots(*indices: int) -> torch.Tensor: + return torch.tensor(indices, dtype=torch.int64) + + +class TestKvWeightVersionRecord(CustomTestCase): + def test_capture_without_a_version_fails(self) -> None: + """A forward cannot publish KV provenance without a weight version.""" + with self.assertRaises(AssertionError): + KvWeightVersionRecord.capture(slot_indices=_slots(1), version=None) + + def test_capture_survives_source_reuse_and_tensor_mapping(self) -> None: + """Reusing forward buffers cannot change a pending record's slots.""" + slots = _slots(4, 5, 6) + record = KvWeightVersionRecord.capture(slot_indices=slots, version="v0") + slots.fill_(9) + record.map_device_tensors(lambda tensor: tensor.to(dtype=torch.int32)) + table = _table() + + table.record(slot_indices=record.slot_indices, version=record.version) + + self.assertEqual( + table._lookup_spans(_slots(4, 5, 6)), + [WeightVersionSpan(version="v0", start=0, end=3)], + ) + self.assertEqual(record.slot_indices.dtype, torch.int32) + + +class TestKvWeightVersionTracker(CustomTestCase): + def test_pipeline_parallelism_is_rejected_only_when_tracking_is_enabled( + self, + ) -> None: + """The optional tracker rejects PP before allocation without restricting disabled runs.""" + for enabled in (False, True): + with self.subTest(enabled=enabled): + kwargs = dict( + server_args=SimpleNamespace( + enable_prefill_weight_versions=enabled, pp_size=2 + ), + model_config=SimpleNamespace( + is_encoder_decoder=False, is_generation=True + ), + allocator=None, + req_to_token_pool=None, + ) + if enabled: + with self.assertRaisesRegex(AssertionError, "pipeline parallelism"): + KvWeightVersionTracker.maybe_create(**kwargs) + else: + self.assertIsNone(KvWeightVersionTracker.maybe_create(**kwargs)) + + def test_factory_includes_allocator_padding_and_uses_the_request_pool(self) -> None: + """Factory sizing includes the allocator's padding page and preserves its mapping.""" + tracker = KvWeightVersionTracker.maybe_create( + server_args=SimpleNamespace(enable_prefill_weight_versions=True, pp_size=1), + model_config=SimpleNamespace(is_encoder_decoder=False, is_generation=True), + allocator=SimpleNamespace(size_full=12, page_size=4, device="cpu"), + req_to_token_pool=_ReqToTokenPoolStub([15]), + ) + req = _ReqStub(num_prompt_tokens=1, kv_committed_len=1) + + tracker.record(slot_indices=_slots(15), version="v0") + tracker.fill_req_prefill_weight_versions(req) + + self.assertEqual( + req.prefill_weight_versions, + [WeightVersionSpan(version="v0", start=0, end=1)], + ) + + def test_unified_swa_virtual_slots_above_the_static_quota_keep_their_versions( + self, + ) -> None: + """Legal virtual slots above the static token quota remain addressable by the tracker.""" + from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( + UnifiedSWATokenToKVPoolAllocator, + ) + + for page_size in (1, 4): + with self.subTest(page_size=page_size): + allocator = UnifiedSWATokenToKVPoolAllocator.__new__( + UnifiedSWATokenToKVPoolAllocator + ) + allocator._size_full = 8 + allocator.page_size = page_size + allocator.device = "cpu" + allocator.full_attn_allocator = SimpleNamespace( + size_full=32, page_size=page_size, device="cpu" + ) + tracker = KvWeightVersionTracker.maybe_create( + server_args=SimpleNamespace( + enable_prefill_weight_versions=True, pp_size=1 + ), + model_config=SimpleNamespace( + is_encoder_decoder=False, is_generation=True + ), + allocator=allocator, + req_to_token_pool=_ReqToTokenPoolStub([31]), + ) + req = _ReqStub(num_prompt_tokens=1, kv_committed_len=1) + + tracker.record(slot_indices=_slots(31), version="v2") + tracker.fill_req_prefill_weight_versions(req) + + self.assertEqual( + req.prefill_weight_versions, + [WeightVersionSpan(version="v2", start=0, end=1)], + ) + + def test_equal_neighbours_merge_while_a_returning_version_starts_a_new_span(self): + """Run-length compression merges adjacent equal versions and never merges across a change.""" + table = _table() + table.record(slot_indices=_slots(4, 5), version="v0") + table.record(slot_indices=_slots(6), version="v1") + table.record(slot_indices=_slots(7), version="v0") + + self.assertEqual( + table._lookup_spans(_slots(4, 5, 6, 7)), + [ + WeightVersionSpan(version="v0", start=0, end=2), + WeightVersionSpan(version="v1", start=2, end=3), + WeightVersionSpan(version="v0", start=3, end=4), + ], + ) + + def test_a_single_slot_lookup_yields_one_unit_span(self): + """A one-token prompt maps to exactly one [0, 1) span.""" + table = _table() + table.record(slot_indices=_slots(7), version="v0") + + self.assertEqual( + table._lookup_spans(_slots(7)), + [WeightVersionSpan(version="v0", start=0, end=1)], + ) + + def test_an_empty_lookup_returns_no_spans(self): + """Looking up an empty prompt yields an empty span list.""" + self.assertEqual(_table()._lookup_spans(_slots()), []) + + def test_unwritten_slots_are_reported_by_index_in_lookup_order(self): + """The error names exactly the never-stamped slots, ordered by their place in the lookup.""" + table = _table() + table.record(slot_indices=_slots(1, 3), version="v0") + + with self.assertRaisesRegex(ValueError, r": \[2, 0\]$"): + table._lookup_spans(_slots(2, 3, 0, 1)) + + def test_a_reused_slot_reports_the_version_that_last_wrote_it(self): + """The table is keyed by slot, so a slot refilled under a new version forgets the old one.""" + table = _table() + table.record(slot_indices=_slots(3), version="v0") + table.record(slot_indices=_slots(3), version="v1") + + self.assertEqual( + table._lookup_spans(_slots(3)), + [WeightVersionSpan(version="v1", start=0, end=1)], + ) + + def test_int32_slot_indices_are_accepted(self): + """out_cache_loc arrives as int32 and must index the table like int64 does.""" + table = _table() + table.record(slot_indices=torch.tensor([2, 3], dtype=torch.int32), version="v0") + + self.assertEqual( + table._lookup_spans(torch.tensor([2, 3], dtype=torch.int32)), + [WeightVersionSpan(version="v0", start=0, end=2)], + ) + + def test_the_last_slot_is_addressable(self): + """The table covers num_slots entries, so index num_slots - 1 is valid.""" + table = _table() + table.record(slot_indices=_slots(15), version="v0") + + self.assertEqual( + table._lookup_spans(_slots(15)), + [WeightVersionSpan(version="v0", start=0, end=1)], + ) + + def test_version_ids_are_dense_and_reused(self): + """Every new version gets the next id and already seen ones are reused.""" + table = _table() + for index, version in enumerate(["v0", "v1", "v2", "v1", "v0"]): + table.record(slot_indices=_slots(index), version=version) + + self.assertEqual(table._slot_version_ids[:5].tolist(), [0, 1, 2, 1, 0]) + + +class TestStringInterner(CustomTestCase): + def test_first_value_gets_id_zero(self): + """A fresh interner starts numbering at zero.""" + self.assertEqual(_StringInterner().intern("v0"), 0) + + def test_new_values_get_consecutive_ids_in_first_seen_order(self): + """Ids are dense and follow the order in which values were first interned.""" + interner = _StringInterner() + + self.assertEqual([interner.intern(v) for v in ["b", "a", "c"]], [0, 1, 2]) + + def test_a_seen_value_reuses_its_id(self): + """Interning a value again returns the id it already has instead of a new one.""" + interner = _StringInterner() + first = interner.intern("v0") + interner.intern("v1") + + self.assertEqual(interner.intern("v0"), first) + self.assertEqual(interner.intern("v2"), 2) + + def test_lookup_returns_the_interned_value(self): + """lookup is the inverse of intern for every id handed out.""" + interner = _StringInterner() + values = ["v0", "v1", "v2"] + ids = [interner.intern(v) for v in values] + + self.assertEqual([interner.lookup(i) for i in ids], values) + + def test_lookup_of_an_unknown_id_fails(self): + """An id that was never handed out cannot be resolved.""" + interner = _StringInterner() + interner.intern("v0") + + with self.assertRaises(IndexError): + interner.lookup(1) + + def test_distinct_strings_never_share_an_id(self): + """Values that differ only slightly are still distinct entries.""" + interner = _StringInterner() + + ids = {interner.intern(v) for v in ["1", "01", "1 ", "default", ""]} + + self.assertEqual(len(ids), 5) + + +class TestFillReqPrefillWeightVersions(CustomTestCase): + def test_prompt_slots_resolve_to_the_versions_that_computed_them(self): + """Only the prompt's own slots are read, even when more KV is committed than there are prompt tokens.""" + table = _table(slots_of_req=[4, 5, 6, 7, 8]) + table.record(slot_indices=_slots(4, 5), version="v0") + table.record(slot_indices=_slots(6), version="v1") + req = _ReqStub(num_prompt_tokens=3, kv_committed_len=5) + + table.fill_req_prefill_weight_versions(req) + + self.assertEqual( + req.prefill_weight_versions, + [ + WeightVersionSpan(version="v0", start=0, end=2), + WeightVersionSpan(version="v1", start=2, end=3), + ], + ) + + def test_a_partially_committed_prompt_fails_the_fill(self): + """Filling before the whole prompt is committed is a scheduling bug and must not yield short spans.""" + table = _table(slots_of_req=[4, 5]) + table.record(slot_indices=_slots(4, 5), version="v0") + req = _ReqStub(num_prompt_tokens=5, kv_committed_len=2) + + with self.assertRaisesRegex( + AssertionError, "2 committed KV tokens for a 5-token prompt" + ): + table.fill_req_prefill_weight_versions(req) + + def test_an_empty_prompt_yields_no_spans(self): + """A request without prompt tokens reports an empty span list, not an error.""" + table = _table(slots_of_req=[]) + req = _ReqStub(num_prompt_tokens=0, kv_committed_len=0) + + table.fill_req_prefill_weight_versions(req) + + self.assertEqual(req.prefill_weight_versions, []) + + def test_an_unstamped_prompt_slot_fails_the_fill(self): + """A prompt slot the tracker never saw surfaces as an error on the request path too.""" + table = _table(slots_of_req=[4, 5, 6]) + table.record(slot_indices=_slots(4, 6), version="v0") + req = _ReqStub(num_prompt_tokens=3, kv_committed_len=3) + + with self.assertRaisesRegex(ValueError, r"\[5\]"): + table.fill_req_prefill_weight_versions(req) + + def test_a_second_fill_replaces_the_previous_spans(self): + """Filling again after a re-prefill replaces the stale spans instead of appending.""" + table = _table(slots_of_req=[4, 5]) + table.record(slot_indices=_slots(4, 5), version="v0") + req = _ReqStub(num_prompt_tokens=2, kv_committed_len=2) + table.fill_req_prefill_weight_versions(req) + table.record(slot_indices=_slots(4, 5), version="v1") + + table.fill_req_prefill_weight_versions(req) + + self.assertEqual( + req.prefill_weight_versions, + [WeightVersionSpan(version="v1", start=0, end=2)], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/model_executor/test_model_runner_kv_weight_versions.py b/test/registered/unit/model_executor/test_model_runner_kv_weight_versions.py new file mode 100644 index 000000000000..8caa2294bc57 --- /dev/null +++ b/test/registered/unit/model_executor/test_model_runner_kv_weight_versions.py @@ -0,0 +1,103 @@ +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.srt.managers.tp_worker import TpModelWorker +from sglang.srt.managers.utils import GenerationBatchResult +from sglang.srt.mem_cache.kv_weight_version_tracker import KvWeightVersionRecord +from sglang.srt.model_executor.model_runner import ModelRunner, ModelRunnerOutput +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=3, suite="base-a-test-cpu") + + +class TestModelRunnerKvWeightVersions(CustomTestCase): + def test_only_enabled_target_forwards_capture_slots(self) -> None: + """Draft forwards and disabled tracking never publish target provenance.""" + for enabled, draft, has_slots in ( + (False, False, True), + (True, True, True), + (True, False, False), + (True, False, True), + ): + with self.subTest(enabled=enabled, draft=draft, has_slots=has_slots): + runner = ModelRunner.__new__(ModelRunner) + runner.is_draft_worker = draft + runner.server_args = SimpleNamespace( + enable_prefill_weight_versions=enabled + ) + batch = SimpleNamespace( + out_cache_loc=torch.tensor([4, 5]) if has_slots else None + ) + with patch( + "sglang.srt.mem_cache.kv_weight_version_tracker.get_serving", + return_value=SimpleNamespace(weight_version="v0"), + ): + record = KvWeightVersionRecord.maybe_capture( + model_runner=runner, forward_batch=batch + ) + + if enabled and not draft and has_slots: + batch.out_cache_loc.fill_(9) + self.assertEqual(record.slot_indices.tolist(), [4, 5]) + self.assertEqual(record.version, "v0") + else: + self.assertIsNone(record) + + def test_target_worker_preserves_forward_provenance(self) -> None: + """TP result construction retains the target forward record.""" + record = KvWeightVersionRecord.capture( + slot_indices=torch.tensor([4, 5]), version="v0" + ) + output = ModelRunnerOutput( + logits_output=None, + can_run_graph=False, + kv_weight_version_record=record, + ) + worker = SimpleNamespace( + is_dllm=lambda: False, + pp_group=SimpleNamespace(is_last_rank=True), + model_runner=SimpleNamespace(forward=lambda *args, **kwargs: output), + ) + batch = SimpleNamespace( + apply_deprecated_skip_attn_backend_init=lambda value: None + ) + + result = TpModelWorker.forward_batch_generation( + worker, batch=None, forward_batch=batch, is_verify=True + ) + + self.assertIs(result.kv_weight_version_record, record) + + def test_result_copies_record_before_publishing_completion(self) -> None: + """The shared result-copy path includes KV slots before its completion event.""" + record = KvWeightVersionRecord.capture( + slot_indices=torch.tensor([4, 5]), version="v0" + ) + source = record.slot_indices + completed: list[bool] = [] + + def publish_completion() -> None: + self.assertIsNot(record.slot_indices, source) + completed.append(True) + + result = GenerationBatchResult( + next_token_ids=torch.tensor([1]), + kv_weight_version_record=record, + copy_done=SimpleNamespace(record=publish_completion), + ) + with patch( + "sglang.srt.managers.utils._async_d2h", + side_effect=lambda tensor: tensor.clone(), + ): + result.copy_to_cpu(return_logprob=False, return_hidden_states=False) + + self.assertEqual(record.slot_indices.tolist(), [4, 5]) + self.assertEqual(completed, [True]) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 8aff0ab199c2..e379b04869bb 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -2244,6 +2244,32 @@ def test_hicache_kernel_keeps_implicit_fa3_decode_backend(self): self.assertEqual(resolution_result(args, "hicache_mem_layout"), "page_first") self.assertIsNone(resolution_result(args, "decode_attention_backend")) + def test_prefill_weight_versions_rejects_hierarchical_cache(self): + """KV restored from a host or storage tier carries no weight version, so the two cannot combine.""" + args = self._make_args( + enable_prefill_weight_versions=True, + enable_hierarchical_cache=True, + ) + + with self.assertRaisesRegex(ValueError, "enable-hierarchical-cache"): + handle_cache_compatibility(args) + + def test_prefill_weight_versions_rejects_disaggregation(self): + """PD prefill workers do not stamp prompt tokens yet, so the flag is refused there.""" + args = self._make_args( + enable_prefill_weight_versions=True, + disaggregation_mode="prefill", + ) + + with self.assertRaisesRegex(ValueError, "disaggregation"): + handle_cache_compatibility(args) + + def test_prefill_weight_versions_alone_is_accepted(self): + """A plain engine may report prompt weight versions.""" + args = self._make_args(enable_prefill_weight_versions=True) + + handle_cache_compatibility(args) + def test_decode_offload_rejects_host_pool_retraction(self): args = self._make_args( disaggregation_mode="decode", diff --git a/test/registered/unit/utils/test_weight_versions.py b/test/registered/unit/utils/test_weight_versions.py index 7c2508b4b799..556407028b70 100644 --- a/test/registered/unit/utils/test_weight_versions.py +++ b/test/registered/unit/utils/test_weight_versions.py @@ -13,6 +13,7 @@ compute_weight_version_spans, record_weight_version_events, truncate_weight_version_events, + weight_version_spans_to_json, ) from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import ( @@ -28,6 +29,7 @@ class _ReqStub: def __init__(self, output_len: int): self.output_ids = [0] * output_len self.weight_version_events = [] + self.prefill_weight_versions = None def record_weight_version_change(self, old_version): record_weight_version_events([self], old_version=old_version) @@ -540,6 +542,50 @@ def test_metadata_omits_spans_when_absent(self): metadata = build_endpoint_weight_version_metadata({"weight_version": "v2"}) self.assertEqual(metadata, {"weight_version": "v2"}) + def test_metadata_carries_prefill_spans_when_present(self): + """prefill_weight_versions rides along into endpoint metadata next to the output spans.""" + prefill_spans = [{"version": "v1", "start": 0, "end": 5}] + metadata = build_endpoint_weight_version_metadata( + { + "weight_version": "v2", + "weight_versions": [{"version": "v2", "start": 0, "end": 3}], + "prefill_weight_versions": prefill_spans, + "prompt_tokens": 5, + } + ) + self.assertEqual(metadata["prefill_weight_versions"], prefill_spans) + self.assertNotIn("prompt_tokens", metadata) + + +class TestWeightVersionSpansToJson(CustomTestCase): + def test_spans_pass_through_unclamped_without_a_limit(self): + """Prefill spans keep their exact ends because no output-token limit applies.""" + spans = [ + WeightVersionSpan(version="v1", start=0, end=5), + WeightVersionSpan(version="v2", start=5, end=9), + ] + self.assertEqual( + weight_version_spans_to_json(spans), + [ + {"version": "v1", "start": 0, "end": 5}, + {"version": "v2", "start": 5, "end": 9}, + ], + ) + + def test_ends_are_clamped_to_the_limit(self): + """With a limit every end is capped, matching the output-span clamp.""" + spans = [ + WeightVersionSpan(version="v1", start=0, end=5), + WeightVersionSpan(version="v2", start=5, end=9), + ] + self.assertEqual( + weight_version_spans_to_json(spans, end_limit=7), + [ + {"version": "v1", "start": 0, "end": 5}, + {"version": "v2", "start": 5, "end": 7}, + ], + ) + class TestAddWeightVersionsToMetaInfo(CustomTestCase): def test_meta_info_gets_dict_spans_and_legacy_field(self): From a5dc4b4f716ccefa98ddc5d717df175672ea0327 Mon Sep 17 00:00:00 2001 From: rockdu Date: Wed, 14 Oct 2026 12:13:38 -0700 Subject: [PATCH 46/56] refactor(diffusion): rename scheduler_rl_mixin.py to flow_rollout_controller.py Pure file rename (git R100); the proof tool has no file-rename primitive. Refactor-Kind: non_mechanical_provable Signed-off-by: rockdu --- .../schedulers/scheduling_flow_match_euler_discrete.py | 2 +- .../{scheduler_rl_mixin.py => flow_rollout_controller.py} | 0 .../runtime/post_training/rollout_denoising_mixin.py | 6 +++--- .../multimodal_gen/test/unit/test_scheduler_rollout_unit.py | 4 ++-- 4 files changed, 6 insertions(+), 6 deletions(-) rename python/sglang/multimodal_gen/runtime/post_training/{scheduler_rl_mixin.py => flow_rollout_controller.py} (100%) diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py index e38d4d792eb5..08aeb7366d24 100644 --- a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py +++ b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -32,7 +32,7 @@ from diffusers.utils import BaseOutput from sglang.multimodal_gen.runtime.models.schedulers.base import BaseScheduler -from sglang.multimodal_gen.runtime.post_training.scheduler_rl_mixin import ( +from sglang.multimodal_gen.runtime.post_training.flow_rollout_controller import ( SchedulerRLMixin, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger diff --git a/python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_mixin.py b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py similarity index 100% rename from python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_mixin.py rename to python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py 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 fde2d199a5e8..5bd235283747 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 @@ -10,14 +10,14 @@ import torch from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.post_training.flow_rollout_controller import ( + SchedulerRLMixin, +) from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import ( RolloutDenoisingEnv, RolloutDitTrajectory, RolloutTrajectoryData, ) -from sglang.multimodal_gen.runtime.post_training.scheduler_rl_mixin import ( - SchedulerRLMixin, -) from sglang.multimodal_gen.runtime.post_training.sp_utils import ( gather_stacked_latents_for_sp, ) diff --git a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py index 21cdb2136468..943ae2d81297 100644 --- a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py +++ b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py @@ -4,8 +4,8 @@ import torch -import sglang.multimodal_gen.runtime.post_training.scheduler_rl_mixin as rl_mixin_module -from sglang.multimodal_gen.runtime.post_training.scheduler_rl_mixin import ( +import sglang.multimodal_gen.runtime.post_training.flow_rollout_controller as rl_mixin_module +from sglang.multimodal_gen.runtime.post_training.flow_rollout_controller import ( SchedulerRLMixin, ) From afe76ad0df1a370295ffae23d22fd5d6cb6bf2a3 Mon Sep 17 00:00:00 2001 From: rockdu Date: Wed, 14 Oct 2026 12:13:38 -0700 Subject: [PATCH 47/56] refactor(move): move SchedulerRLDebugMixin into flow_rollout_controller.py Refactor-Kind: mechanical_provable Signed-off-by: rockdu --- .../post_training/flow_rollout_controller.py | 107 +++++++++++++++- .../post_training/scheduler_rl_debug_mixin.py | 116 ------------------ 2 files changed, 104 insertions(+), 119 deletions(-) delete mode 100644 python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_debug_mixin.py diff --git a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py index aa52a25e6542..b0298a3d7a93 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py +++ b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py @@ -7,19 +7,120 @@ import torch from sglang.multimodal_gen.runtime.distributed import ( + get_local_torch_device, get_sp_world_size, ) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import ( + RolloutDebugTensors, RolloutSessionData, ) -from sglang.multimodal_gen.runtime.post_training.scheduler_rl_debug_mixin import ( - SchedulerRLDebugMixin, -) _LOG_SQRT_2PI = math.log(math.sqrt(2 * math.pi)) +class SchedulerRLDebugMixin: + @staticmethod + def _reset_rollout_debug_tensors(rollout_session_data: RolloutSessionData) -> None: + rollout_session_data.local_variance_noises = [] + rollout_session_data.local_prev_sample_means = [] + rollout_session_data.local_noise_std_devs = [] + rollout_session_data.local_model_outputs = [] + + def append_local_rollout_debug_tensors( + self, + batch, + *, + variance_noise: torch.Tensor, + prev_sample_mean: torch.Tensor, + noise_std_dev: torch.Tensor, + model_output: torch.Tensor, + ) -> None: + rollout_session_data = batch._rollout_session_data + batch_size = variance_noise.shape[0] + # the underlying noise buffer in batch._rollout_session_data.noise_buffer is reused. + rollout_session_data.local_variance_noises.append(variance_noise.clone()) + rollout_session_data.local_prev_sample_means.append(prev_sample_mean) + rollout_session_data.local_noise_std_devs.append( + noise_std_dev.expand((batch_size, 1)) + ) + rollout_session_data.local_model_outputs.append(model_output) + + def consume_local_rollout_debug_tensors( + self, + batch, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + rollout_session_data = batch._rollout_session_data + variance_noises = torch.stack(rollout_session_data.local_variance_noises, dim=1) + prev_sample_means = torch.stack( + rollout_session_data.local_prev_sample_means, dim=1 + ) + noise_std_devs = torch.stack(rollout_session_data.local_noise_std_devs, dim=1) + model_outputs = torch.stack(rollout_session_data.local_model_outputs, dim=1) + self._reset_rollout_debug_tensors(rollout_session_data) + return variance_noises, prev_sample_means, noise_std_devs, model_outputs + + def collect_rollout_debug_tensors(self, batch: Req) -> RolloutDebugTensors: + """ + Consume rollout debug tensors and merge for all SP ranks. + + Returns rollout debug tensors with shape [B, T, ...]. + """ + rollout_session_data = batch._rollout_session_data + variance_noises, prev_sample_means, noise_std_devs, model_outputs = ( + self.consume_local_rollout_debug_tensors(batch) + ) + + if get_sp_world_size() > 1 and getattr(batch, "did_sp_shard_latents", False): + variance_noises = variance_noises.to(get_local_torch_device()) + prev_sample_means = prev_sample_means.to(get_local_torch_device()) + noise_std_devs = noise_std_devs.to(get_local_torch_device()) + model_outputs = model_outputs.to(get_local_torch_device()) + pipeline_config = rollout_session_data.pipeline_config + bsz, num_steps = variance_noises.shape[0], variance_noises.shape[1] + + # [B, T, ...] -> [B*T, ...] + variance_noises_packed = variance_noises.contiguous().reshape( + bsz * num_steps, *variance_noises.shape[2:] + ) + prev_sample_means_packed = prev_sample_means.contiguous().reshape( + bsz * num_steps, *prev_sample_means.shape[2:] + ) + model_outputs_packed = model_outputs.contiguous().reshape( + bsz * num_steps, *model_outputs.shape[2:] + ) + + # Gather on packed tensors first. + variance_noises_packed = pipeline_config.gather_latents_for_sp( + variance_noises_packed, batch=batch + ) + prev_sample_means_packed = pipeline_config.gather_latents_for_sp( + prev_sample_means_packed, batch=batch + ) + model_outputs_packed = pipeline_config.gather_latents_for_sp( + model_outputs_packed, batch=batch + ) + + # Unpack back to [B, T, ...]. + variance_noises = variance_noises_packed.reshape( + bsz, num_steps, *variance_noises_packed.shape[1:] + ) + prev_sample_means = prev_sample_means_packed.reshape( + bsz, num_steps, *prev_sample_means_packed.shape[1:] + ) + model_outputs = model_outputs_packed.reshape( + bsz, num_steps, *model_outputs_packed.shape[1:] + ) + # noise_std_devs is same on every device, not a sharded latent tensor. + + return RolloutDebugTensors( + rollout_variance_noises=variance_noises.cpu(), + rollout_prev_sample_means=prev_sample_means.cpu(), + rollout_noise_std_devs=noise_std_devs.cpu(), + rollout_model_outputs=model_outputs.cpu(), + ) + + class SchedulerRLMixin(SchedulerRLDebugMixin): @staticmethod def _get_rollout_session_data(batch) -> RolloutSessionData: diff --git a/python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_debug_mixin.py b/python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_debug_mixin.py deleted file mode 100644 index e6113876ae38..000000000000 --- a/python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_debug_mixin.py +++ /dev/null @@ -1,116 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Debug tensor helpers for rollout-enabled schedulers.""" - -import torch - -from sglang.multimodal_gen.runtime.distributed import ( - get_local_torch_device, - get_sp_world_size, -) -from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req -from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import ( - RolloutDebugTensors, - RolloutSessionData, -) - - -class SchedulerRLDebugMixin: - @staticmethod - def _reset_rollout_debug_tensors(rollout_session_data: RolloutSessionData) -> None: - rollout_session_data.local_variance_noises = [] - rollout_session_data.local_prev_sample_means = [] - rollout_session_data.local_noise_std_devs = [] - rollout_session_data.local_model_outputs = [] - - def append_local_rollout_debug_tensors( - self, - batch, - *, - variance_noise: torch.Tensor, - prev_sample_mean: torch.Tensor, - noise_std_dev: torch.Tensor, - model_output: torch.Tensor, - ) -> None: - rollout_session_data = batch._rollout_session_data - batch_size = variance_noise.shape[0] - # the underlying noise buffer in batch._rollout_session_data.noise_buffer is reused. - rollout_session_data.local_variance_noises.append(variance_noise.clone()) - rollout_session_data.local_prev_sample_means.append(prev_sample_mean) - rollout_session_data.local_noise_std_devs.append( - noise_std_dev.expand((batch_size, 1)) - ) - rollout_session_data.local_model_outputs.append(model_output) - - def consume_local_rollout_debug_tensors( - self, - batch, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - rollout_session_data = batch._rollout_session_data - variance_noises = torch.stack(rollout_session_data.local_variance_noises, dim=1) - prev_sample_means = torch.stack( - rollout_session_data.local_prev_sample_means, dim=1 - ) - noise_std_devs = torch.stack(rollout_session_data.local_noise_std_devs, dim=1) - model_outputs = torch.stack(rollout_session_data.local_model_outputs, dim=1) - self._reset_rollout_debug_tensors(rollout_session_data) - return variance_noises, prev_sample_means, noise_std_devs, model_outputs - - def collect_rollout_debug_tensors(self, batch: Req) -> RolloutDebugTensors: - """ - Consume rollout debug tensors and merge for all SP ranks. - - Returns rollout debug tensors with shape [B, T, ...]. - """ - rollout_session_data = batch._rollout_session_data - variance_noises, prev_sample_means, noise_std_devs, model_outputs = ( - self.consume_local_rollout_debug_tensors(batch) - ) - - if get_sp_world_size() > 1 and getattr(batch, "did_sp_shard_latents", False): - variance_noises = variance_noises.to(get_local_torch_device()) - prev_sample_means = prev_sample_means.to(get_local_torch_device()) - noise_std_devs = noise_std_devs.to(get_local_torch_device()) - model_outputs = model_outputs.to(get_local_torch_device()) - pipeline_config = rollout_session_data.pipeline_config - bsz, num_steps = variance_noises.shape[0], variance_noises.shape[1] - - # [B, T, ...] -> [B*T, ...] - variance_noises_packed = variance_noises.contiguous().reshape( - bsz * num_steps, *variance_noises.shape[2:] - ) - prev_sample_means_packed = prev_sample_means.contiguous().reshape( - bsz * num_steps, *prev_sample_means.shape[2:] - ) - model_outputs_packed = model_outputs.contiguous().reshape( - bsz * num_steps, *model_outputs.shape[2:] - ) - - # Gather on packed tensors first. - variance_noises_packed = pipeline_config.gather_latents_for_sp( - variance_noises_packed, batch=batch - ) - prev_sample_means_packed = pipeline_config.gather_latents_for_sp( - prev_sample_means_packed, batch=batch - ) - model_outputs_packed = pipeline_config.gather_latents_for_sp( - model_outputs_packed, batch=batch - ) - - # Unpack back to [B, T, ...]. - variance_noises = variance_noises_packed.reshape( - bsz, num_steps, *variance_noises_packed.shape[1:] - ) - prev_sample_means = prev_sample_means_packed.reshape( - bsz, num_steps, *prev_sample_means_packed.shape[1:] - ) - model_outputs = model_outputs_packed.reshape( - bsz, num_steps, *model_outputs_packed.shape[1:] - ) - # noise_std_devs is same on every device, not a sharded latent tensor. - - return RolloutDebugTensors( - rollout_variance_noises=variance_noises.cpu(), - rollout_prev_sample_means=prev_sample_means.cpu(), - rollout_noise_std_devs=noise_std_devs.cpu(), - rollout_model_outputs=model_outputs.cpu(), - ) From 2d8004670a75f9311ca6a09d6e62be19ac9391bc Mon Sep 17 00:00:00 2001 From: rockdu Date: Wed, 14 Oct 2026 12:13:38 -0700 Subject: [PATCH 48/56] refactor(move): extract the FlowMatch Euler update into _transition Refactor-Kind: mechanical_provable Signed-off-by: rockdu --- .../scheduling_flow_match_euler_discrete.py | 46 ++++++++++++++----- 1 file changed, 35 insertions(+), 11 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py index 08aeb7366d24..b16e0a1d810b 100644 --- a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py +++ b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -531,17 +531,16 @@ def step( batch, model_output, sample, current_sigma, next_sigma, generator ) else: - if self.config.stochastic_sampling: - x0 = sample - current_sigma * model_output - noise = torch.randn( - sample.shape, - generator=generator, - device=sample.device, - dtype=sample.dtype, - ) - prev_sample = (1.0 - next_sigma) * x0 + next_sigma * noise - else: - prev_sample = sample + dt * model_output + prev_sample = self._transition( + model_output=model_output, + sample=sample, + current_sigma=current_sigma, + next_sigma=next_sigma, + dt=dt, + generator=generator, + step_index=self.step_index, + timestep=timestep, + ) # upon completion increase step index by one assert self._step_index is not None, "_step_index should not be None" @@ -645,6 +644,31 @@ def _convert_to_beta( ) return sigmas + def _transition( + self, + *, + model_output: torch.FloatTensor, + sample: torch.FloatTensor, + current_sigma: torch.FloatTensor, + next_sigma: torch.FloatTensor, + dt: torch.FloatTensor, + generator: torch.Generator | None, + step_index: int, + timestep: int | torch.Tensor, + ) -> torch.FloatTensor: + if self.config.stochastic_sampling: + x0 = sample - current_sigma * model_output + noise = torch.randn( + sample.shape, + generator=generator, + device=sample.device, + dtype=sample.dtype, + ) + prev_sample = (1.0 - next_sigma) * x0 + next_sigma * noise + else: + prev_sample = sample + dt * model_output + return prev_sample + def _time_shift_exponential( self, mu: float, sigma: float, t: torch.Tensor | np.ndarray ) -> torch.Tensor | np.ndarray: From adb178bbefaf5e06f9eb2b66031ce04f760334fb Mon Sep 17 00:00:00 2001 From: rockdu Date: Wed, 14 Oct 2026 12:13:38 -0700 Subject: [PATCH 49/56] refactor(move): move RolloutSessionData into flow_rollout_controller.py Refactor-Kind: mechanical_provable Signed-off-by: rockdu --- .../post_training/flow_rollout_controller.py | 23 ++++++++++++++++++- .../runtime/post_training/rl_dataclasses.py | 23 +------------------ 2 files changed, 23 insertions(+), 23 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py index b0298a3d7a93..e17161399b99 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py +++ b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py @@ -2,6 +2,7 @@ """Flow-matching rollout step utilities for log-prob computation.""" import math +from dataclasses import dataclass, field from typing import Any, Union import torch @@ -13,12 +14,32 @@ from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import ( RolloutDebugTensors, - RolloutSessionData, ) _LOG_SQRT_2PI = math.log(math.sqrt(2 * math.pi)) +@dataclass +class RolloutSessionData: + """Per-batch rollout state created by prepare_rollout(), lives on the batch object. + + Cleared by setting ``batch._rollout_session_data = None``. + """ + + pipeline_config: Any = None + sigma_max: float = 0.0 + latents_shape: tuple | None = None + noise_buffer: torch.Tensor | None = None + + local_log_prob_sum: list[torch.Tensor] = field(default_factory=list) + local_log_prob_count: list[torch.Tensor] = field(default_factory=list) + + local_variance_noises: list[torch.Tensor] = field(default_factory=list) + local_prev_sample_means: list[torch.Tensor] = field(default_factory=list) + local_noise_std_devs: list[torch.Tensor] = field(default_factory=list) + local_model_outputs: list[torch.Tensor] = field(default_factory=list) + + class SchedulerRLDebugMixin: @staticmethod def _reset_rollout_debug_tensors(rollout_session_data: RolloutSessionData) -> None: 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 59858a27a7c6..e9756e0747ff 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/rl_dataclasses.py +++ b/python/sglang/multimodal_gen/runtime/post_training/rl_dataclasses.py @@ -1,33 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 """RL-specific dataclasses used by post-training and rollout paths.""" -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import Any import torch -@dataclass -class RolloutSessionData: - """Per-batch rollout state created by prepare_rollout(), lives on the batch object. - - Cleared by setting ``batch._rollout_session_data = None``. - """ - - pipeline_config: Any = None - sigma_max: float = 0.0 - latents_shape: tuple | None = None - noise_buffer: torch.Tensor | None = None - - local_log_prob_sum: list[torch.Tensor] = field(default_factory=list) - local_log_prob_count: list[torch.Tensor] = field(default_factory=list) - - local_variance_noises: list[torch.Tensor] = field(default_factory=list) - local_prev_sample_means: list[torch.Tensor] = field(default_factory=list) - local_noise_std_devs: list[torch.Tensor] = field(default_factory=list) - local_model_outputs: list[torch.Tensor] = field(default_factory=list) - - @dataclass class RolloutDebugTensors: """Container for rollout debug tensors collected during denoising.""" From 74c54346547bb3cab665e472ed1cdfa6e08028ed Mon Sep 17 00:00:00 2001 From: rockdu Date: Wed, 14 Oct 2026 12:13:38 -0700 Subject: [PATCH 50/56] refactor(diffusion): prepare the rollout session methods to move onto RolloutSessionData Refactor-Kind: non_mechanical_provable Signed-off-by: rockdu --- .../scheduling_flow_match_euler_discrete.py | 11 +- .../pipelines_core/stages/denoising.py | 1 - .../stages/model_specific_stages/cosmos3.py | 1 - .../post_training/flow_rollout_controller.py | 201 ++++++++---------- .../post_training/rollout_denoising_mixin.py | 8 +- .../test/unit/test_scheduler_rollout_unit.py | 113 +++++----- 6 files changed, 161 insertions(+), 174 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py index b16e0a1d810b..4dcf81e3285d 100644 --- a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py +++ b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -527,8 +527,15 @@ def step( if batch is not None and batch.rollout: if not self.already_prepared_rollout(batch): raise RuntimeError("Rollout not prepared before step") - prev_sample = self.flow_sde_sampling( - batch, model_output, sample, current_sigma, next_sigma, generator + prev_sample = SchedulerRLMixin.flow_sde_sampling( + self._get_rollout_session_data(batch), + batch, + model_output, + sample, + current_sigma, + next_sigma, + generator, + step_index=self.step_index, ) else: prev_sample = self._transition( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index bd6a0886ce2f..b675abfd5715 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -2103,7 +2103,6 @@ def _denoise( # pre-step value. Gated on batch.rollout to keep the # non-rollout path strictly untouched. if batch.rollout: - batch._rollout_loop_step_index = step_index self._maybe_append_dit_trajectory_step( batch=batch, latents=ctx.latents, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py index a00599e69b77..3e0b7f43e2b6 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py @@ -1781,7 +1781,6 @@ def _denoise_once( if batch.rollout: # Capture the pre-step x_{t_i} before the scheduler advances it. - batch._rollout_loop_step_index = i self._maybe_append_dit_trajectory_step( batch=batch, latents=latents, diff --git a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py index e17161399b99..737466aea69d 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py +++ b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py @@ -8,13 +8,15 @@ import torch from sglang.multimodal_gen.runtime.distributed import ( - get_local_torch_device, get_sp_world_size, ) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import ( RolloutDebugTensors, ) +from sglang.multimodal_gen.runtime.post_training.sp_utils import ( + gather_stacked_latents_for_sp, +) _LOG_SQRT_2PI = math.log(math.sqrt(2 * math.pi)) @@ -40,99 +42,71 @@ class RolloutSessionData: local_model_outputs: list[torch.Tensor] = field(default_factory=list) -class SchedulerRLDebugMixin: +class SchedulerRLMixin: @staticmethod - def _reset_rollout_debug_tensors(rollout_session_data: RolloutSessionData) -> None: - rollout_session_data.local_variance_noises = [] - rollout_session_data.local_prev_sample_means = [] - rollout_session_data.local_noise_std_devs = [] - rollout_session_data.local_model_outputs = [] + def _reset_rollout_debug_tensors(self: RolloutSessionData) -> None: + self.local_variance_noises = [] + self.local_prev_sample_means = [] + self.local_noise_std_devs = [] + self.local_model_outputs = [] + @staticmethod def append_local_rollout_debug_tensors( - self, - batch, + self: RolloutSessionData, *, variance_noise: torch.Tensor, prev_sample_mean: torch.Tensor, noise_std_dev: torch.Tensor, model_output: torch.Tensor, ) -> None: - rollout_session_data = batch._rollout_session_data batch_size = variance_noise.shape[0] - # the underlying noise buffer in batch._rollout_session_data.noise_buffer is reused. - rollout_session_data.local_variance_noises.append(variance_noise.clone()) - rollout_session_data.local_prev_sample_means.append(prev_sample_mean) - rollout_session_data.local_noise_std_devs.append( - noise_std_dev.expand((batch_size, 1)) - ) - rollout_session_data.local_model_outputs.append(model_output) + # the underlying noise buffer in self.noise_buffer is reused. + self.local_variance_noises.append(variance_noise.clone()) + self.local_prev_sample_means.append(prev_sample_mean) + self.local_noise_std_devs.append(noise_std_dev.expand((batch_size, 1))) + self.local_model_outputs.append(model_output) + @staticmethod def consume_local_rollout_debug_tensors( - self, - batch, + self: RolloutSessionData, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - rollout_session_data = batch._rollout_session_data - variance_noises = torch.stack(rollout_session_data.local_variance_noises, dim=1) - prev_sample_means = torch.stack( - rollout_session_data.local_prev_sample_means, dim=1 - ) - noise_std_devs = torch.stack(rollout_session_data.local_noise_std_devs, dim=1) - model_outputs = torch.stack(rollout_session_data.local_model_outputs, dim=1) - self._reset_rollout_debug_tensors(rollout_session_data) + variance_noises = torch.stack(self.local_variance_noises, dim=1) + prev_sample_means = torch.stack(self.local_prev_sample_means, dim=1) + noise_std_devs = torch.stack(self.local_noise_std_devs, dim=1) + model_outputs = torch.stack(self.local_model_outputs, dim=1) + SchedulerRLMixin._reset_rollout_debug_tensors(self) return variance_noises, prev_sample_means, noise_std_devs, model_outputs - def collect_rollout_debug_tensors(self, batch: Req) -> RolloutDebugTensors: + @staticmethod + def collect_rollout_debug_tensors( + self: RolloutSessionData, batch: Req + ) -> RolloutDebugTensors: """ Consume rollout debug tensors and merge for all SP ranks. Returns rollout debug tensors with shape [B, T, ...]. """ - rollout_session_data = batch._rollout_session_data variance_noises, prev_sample_means, noise_std_devs, model_outputs = ( - self.consume_local_rollout_debug_tensors(batch) + SchedulerRLMixin.consume_local_rollout_debug_tensors(self) ) - if get_sp_world_size() > 1 and getattr(batch, "did_sp_shard_latents", False): - variance_noises = variance_noises.to(get_local_torch_device()) - prev_sample_means = prev_sample_means.to(get_local_torch_device()) - noise_std_devs = noise_std_devs.to(get_local_torch_device()) - model_outputs = model_outputs.to(get_local_torch_device()) - pipeline_config = rollout_session_data.pipeline_config - bsz, num_steps = variance_noises.shape[0], variance_noises.shape[1] - - # [B, T, ...] -> [B*T, ...] - variance_noises_packed = variance_noises.contiguous().reshape( - bsz * num_steps, *variance_noises.shape[2:] - ) - prev_sample_means_packed = prev_sample_means.contiguous().reshape( - bsz * num_steps, *prev_sample_means.shape[2:] - ) - model_outputs_packed = model_outputs.contiguous().reshape( - bsz * num_steps, *model_outputs.shape[2:] - ) - - # Gather on packed tensors first. - variance_noises_packed = pipeline_config.gather_latents_for_sp( - variance_noises_packed, batch=batch - ) - prev_sample_means_packed = pipeline_config.gather_latents_for_sp( - prev_sample_means_packed, batch=batch - ) - model_outputs_packed = pipeline_config.gather_latents_for_sp( - model_outputs_packed, batch=batch - ) - - # Unpack back to [B, T, ...]. - variance_noises = variance_noises_packed.reshape( - bsz, num_steps, *variance_noises_packed.shape[1:] - ) - prev_sample_means = prev_sample_means_packed.reshape( - bsz, num_steps, *prev_sample_means_packed.shape[1:] - ) - model_outputs = model_outputs_packed.reshape( - bsz, num_steps, *model_outputs_packed.shape[1:] - ) - # noise_std_devs is same on every device, not a sharded latent tensor. + pipeline_config = self.pipeline_config + variance_noises = gather_stacked_latents_for_sp( + pipeline_config=pipeline_config, + batch=batch, + stacked_latents=variance_noises, + ) + prev_sample_means = gather_stacked_latents_for_sp( + pipeline_config=pipeline_config, + batch=batch, + stacked_latents=prev_sample_means, + ) + model_outputs = gather_stacked_latents_for_sp( + pipeline_config=pipeline_config, + batch=batch, + stacked_latents=model_outputs, + ) + # noise_std_devs is same on every device, not a sharded latent tensor. return RolloutDebugTensors( rollout_variance_noises=variance_noises.cpu(), @@ -141,8 +115,6 @@ def collect_rollout_debug_tensors(self, batch: Req) -> RolloutDebugTensors: rollout_model_outputs=model_outputs.cpu(), ) - -class SchedulerRLMixin(SchedulerRLDebugMixin): @staticmethod def _get_rollout_session_data(batch) -> RolloutSessionData: """Return the RolloutSessionData attached to *batch*, or raise if not prepared.""" @@ -172,15 +144,15 @@ def prepare_rollout(self, batch: Req, pipeline_config: Any = None) -> None: def already_prepared_rollout(self, batch) -> bool: return getattr(batch, "_rollout_session_data", None) is not None + @staticmethod def _get_or_create_rollout_noise_buffer( - self, - rollout_session_data: RolloutSessionData, + self: RolloutSessionData, full_shape: tuple, device: torch.device, dtype: torch.dtype, ) -> torch.Tensor: """Get or create the reusable noise buffer (local or full shape) for rollout.""" - buffer = rollout_session_data.noise_buffer + buffer = self.noise_buffer if ( buffer is None or buffer.shape != full_shape @@ -188,11 +160,12 @@ def _get_or_create_rollout_noise_buffer( or buffer.device != device ): buffer = torch.empty(full_shape, device=device, dtype=dtype) - rollout_session_data.noise_buffer = buffer + self.noise_buffer = buffer return buffer + @staticmethod def _rollout_variance_noise( - self, + self: RolloutSessionData, batch, model_output: torch.FloatTensor, generator: Union[torch.Generator, list[torch.Generator]], @@ -200,7 +173,6 @@ def _rollout_variance_noise( """Generate variance noise for rollout. If generator is a list, use generator[i] for the i-th batch item.""" assert generator is not None, "Generator must be provided" - rollout_session_data = self._get_rollout_session_data(batch) device = model_output.device dtype = model_output.dtype local_shape = tuple(model_output.shape) @@ -214,17 +186,17 @@ def _rollout_variance_noise( "Generator list must have the same length as batch size" ) - buffer = self._get_or_create_rollout_noise_buffer( - rollout_session_data, rollout_session_data.latents_shape, device, dtype + buffer = SchedulerRLMixin._get_or_create_rollout_noise_buffer( + self, self.latents_shape, device, dtype ) for i in range(B): torch.randn( - rollout_session_data.latents_shape, + self.latents_shape, out=buffer[i : i + 1], generator=generator[i], ) - sharded_noise, _ = rollout_session_data.pipeline_config.shard_latents_for_sp( + sharded_noise, _ = self.pipeline_config.shard_latents_for_sp( batch=batch, latents=buffer ) if tuple(sharded_noise.shape) != local_shape: @@ -234,14 +206,16 @@ def _rollout_variance_noise( ) return sharded_noise + @staticmethod def flow_sde_sampling( - self, + self: RolloutSessionData, batch, model_output: torch.FloatTensor, sample: torch.FloatTensor, current_sigma: torch.FloatTensor, next_sigma: torch.FloatTensor, generator: torch.Generator, + step_index: int, ) -> torch.Tensor: """Flow rollout step for log-prob / sampling (see FlowGRPO-style references). @@ -251,7 +225,6 @@ def flow_sde_sampling( 2. ``"cps"``: Coupled Particle Sampling. 3. ``"ode"``: Deterministic ODE step (no diffusion noise). """ - rollout_session_data = self._get_rollout_session_data(batch) sde_type = batch.rollout_sde_type noise_level = float(batch.rollout_noise_level) log_prob_no_const = batch.rollout_log_prob_no_const @@ -264,16 +237,11 @@ def flow_sde_sampling( dt = next_sigma - current_sigma - # step_index comes from the denoising-loop counter stashed by - # DenoisingStage — scheduler._step_index would differ when - # _begin_index != 0 (e.g. partial denoising). sde_step_indices = getattr(batch, "rollout_sde_step_indices", None) - loop_step_index = getattr(batch, "_rollout_loop_step_index", None) if ( sde_type != "ode" and sde_step_indices is not None - and loop_step_index is not None - and loop_step_index not in sde_step_indices + and step_index not in sde_step_indices ): effective_sde_type = "ode" else: @@ -288,10 +256,10 @@ def flow_sde_sampling( if effective_sde_type == "sde": model_output = model_output.float() sample = sample.float() - variance_noise = self._rollout_variance_noise( - batch, model_output, generator + variance_noise = SchedulerRLMixin._rollout_variance_noise( + self, batch, model_output, generator ) - full_variance_noise = rollout_session_data.noise_buffer + full_variance_noise = self.noise_buffer std_dev_t = ( torch.sqrt( current_sigma @@ -299,7 +267,7 @@ def flow_sde_sampling( 1 - torch.where( torch.isclose(current_sigma, current_sigma.new_tensor(1.0)), - rollout_session_data.sigma_max, + self.sigma_max, current_sigma, ) ) @@ -321,10 +289,10 @@ def flow_sde_sampling( elif effective_sde_type == "cps": model_output = model_output.float() sample = sample.float() - variance_noise = self._rollout_variance_noise( - batch, model_output, generator + variance_noise = SchedulerRLMixin._rollout_variance_noise( + self, batch, model_output, generator ) - full_variance_noise = rollout_session_data.noise_buffer + full_variance_noise = self.noise_buffer std_dev_t = next_sigma * math.sin(noise_level * math.pi / 2) noise_std_dev = std_dev_t pred_original_sample = sample - current_sigma * model_output @@ -345,7 +313,7 @@ def flow_sde_sampling( (), device=model_output.device, dtype=model_output.dtype ) log_prob_no_const_val = torch.zeros( - rollout_session_data.latents_shape, + self.latents_shape, device=model_output.device, dtype=torch.float32, ) @@ -375,43 +343,48 @@ def flow_sde_sampling( ).sum(dim=list(range(1, len(log_prob_no_const_val.shape)))) if debug_mode: - self.append_local_rollout_debug_tensors( - batch, + SchedulerRLMixin.append_local_rollout_debug_tensors( + self, variance_noise=variance_noise, prev_sample_mean=prev_sample_mean, noise_std_dev=noise_std_dev, model_output=model_output, ) - self.append_local_rollout_log_probs(batch, log_prob_local_sum, local_elem_count) + SchedulerRLMixin.append_local_rollout_log_probs( + self, log_prob_local_sum, local_elem_count + ) return prev_sample + @staticmethod def append_local_rollout_log_probs( - self, batch, log_prob_sum: torch.Tensor, log_prob_count: torch.Tensor + self: RolloutSessionData, + log_prob_sum: torch.Tensor, + log_prob_count: torch.Tensor, ) -> None: - rollout_session_data = self._get_rollout_session_data(batch) - rollout_session_data.local_log_prob_sum.append(log_prob_sum) - rollout_session_data.local_log_prob_count.append(log_prob_count) + self.local_log_prob_sum.append(log_prob_sum) + self.local_log_prob_count.append(log_prob_count) + @staticmethod def consume_local_rollout_log_probs( - self, batch + self: RolloutSessionData, ) -> tuple[torch.Tensor, torch.Tensor]: - rollout_session_data = self._get_rollout_session_data(batch) # [B, T]: batch dim 0, denoising step dim 1 - values_sum = torch.stack(rollout_session_data.local_log_prob_sum, dim=1) - values_count = torch.stack(rollout_session_data.local_log_prob_count, dim=1) - rollout_session_data.local_log_prob_sum = [] - rollout_session_data.local_log_prob_count = [] + values_sum = torch.stack(self.local_log_prob_sum, dim=1) + values_count = torch.stack(self.local_log_prob_count, dim=1) + self.local_log_prob_sum = [] + self.local_log_prob_count = [] return values_sum, values_count - def collect_rollout_log_probs(self, batch: Req) -> torch.Tensor | None: + @staticmethod + def collect_rollout_log_probs(self: RolloutSessionData) -> torch.Tensor | None: """Per-step sums are already computed on the full pre-shard noise buffer inside flow_sde_sampling, so every SP rank holds identical values here and no all-reduce is needed.""" trajectory_log_prob_sum, trajectory_log_prob_count = ( - self.consume_local_rollout_log_probs(batch) + SchedulerRLMixin.consume_local_rollout_log_probs(self) ) rollout_log_probs_tensor = trajectory_log_prob_sum / trajectory_log_prob_count return rollout_log_probs_tensor.cpu() 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 5bd235283747..908a36f78f92 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 @@ -67,11 +67,15 @@ def _maybe_collect_rollout_log_probs(self, batch: Req): if batch.rollout_trajectory_data is None: batch.rollout_trajectory_data = RolloutTrajectoryData() batch.rollout_trajectory_data.rollout_log_probs = ( - scheduler.collect_rollout_log_probs(batch) + SchedulerRLMixin.collect_rollout_log_probs( + scheduler._get_rollout_session_data(batch) + ) ) if batch.rollout_debug_mode: batch.rollout_trajectory_data.rollout_debug_tensors = ( - scheduler.collect_rollout_debug_tensors(batch) + SchedulerRLMixin.collect_rollout_debug_tensors( + scheduler._get_rollout_session_data(batch), batch + ) ) scheduler.release_rollout_resources(batch) diff --git a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py index 943ae2d81297..667b764ce335 100644 --- a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py +++ b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py @@ -38,26 +38,30 @@ def test_ode_step_does_not_call_variance_noise_sampler(self): batch = self._build_batch(debug_mode=False) scheduler.prepare_rollout(batch) - def _raise_if_called(*args, **kwargs): - raise AssertionError("ODE path should not call _rollout_variance_noise") - - scheduler._rollout_variance_noise = _raise_if_called # type: ignore[method-assign] + generator = torch.Generator().manual_seed(1) + generator_state = generator.get_state() sample = torch.randn(2, 4, 8, 8, dtype=torch.float32) model_output = torch.randn_like(sample) current_sigma = torch.tensor(0.6, dtype=torch.float32) next_sigma = torch.tensor(0.4, dtype=torch.float32) - prev_sample = scheduler.flow_sde_sampling( + prev_sample = SchedulerRLMixin.flow_sde_sampling( + batch._rollout_session_data, batch, model_output=model_output, sample=sample, current_sigma=current_sigma, next_sigma=next_sigma, - generator=torch.Generator(device=sample.device).manual_seed(1), + generator=generator, + step_index=2, ) + # The ODE path must not draw variance noise from the request generator. + self.assertTrue(torch.equal(generator.get_state(), generator_state)) log_prob_local_sum, local_elem_count = ( - scheduler.consume_local_rollout_log_probs(batch) + SchedulerRLMixin.consume_local_rollout_log_probs( + batch._rollout_session_data + ) ) log_prob_local_sum = log_prob_local_sum.squeeze(-1) local_elem_count = local_elem_count.squeeze(-1) @@ -89,13 +93,15 @@ def test_ode_bit_exact_with_non_rollout_path(self): next_sigma = torch.tensor(0.4, dtype=torch.float32) dt = next_sigma - current_sigma - rollout_prev = scheduler.flow_sde_sampling( + rollout_prev = SchedulerRLMixin.flow_sde_sampling( + batch._rollout_session_data, batch, model_output=model_output, sample=sample, current_sigma=current_sigma, next_sigma=next_sigma, generator=torch.Generator(device=sample.device).manual_seed(0), + step_index=2, ) # Exact expression used by the non-rollout branch at # scheduling_flow_match_euler_discrete.py `prev_sample = sample + @@ -132,13 +138,15 @@ def test_ode_debug_tensors_have_shape_safe_noise_std(self): current_sigma = torch.tensor(0.6, dtype=torch.float32) next_sigma = torch.tensor(0.4, dtype=torch.float32) - scheduler.flow_sde_sampling( + SchedulerRLMixin.flow_sde_sampling( + batch._rollout_session_data, batch, model_output=model_output, sample=sample, current_sigma=current_sigma, next_sigma=next_sigma, generator=torch.Generator(device=sample.device).manual_seed(2), + step_index=2, ) ( @@ -146,7 +154,9 @@ def test_ode_debug_tensors_have_shape_safe_noise_std(self): prev_sample_means, noise_std_devs, model_outputs, - ) = scheduler.consume_local_rollout_debug_tensors(batch) + ) = SchedulerRLMixin.consume_local_rollout_debug_tensors( + batch._rollout_session_data + ) # [B, T, ...] with one step in this test. self.assertEqual(tuple(variance_noises.shape), (2, 1, 4, 8, 8)) @@ -253,7 +263,7 @@ def test_single_step_matches_flowgrpo_reference(self): shape = (1, 16, 1, 32, 32) atol = 1e-6 pipeline_config = types.SimpleNamespace( - shard_latents_for_sp=lambda _batch, latents: (latents, False) + shard_latents_for_sp=lambda batch, latents: (latents, False) ) for sde_type in ("sde", "cps"): @@ -265,32 +275,29 @@ def test_single_step_matches_flowgrpo_reference(self): g = torch.Generator(device="cpu").manual_seed(seed) model_output = torch.randn(shape, generator=g, dtype=torch.float32) sample = torch.randn(shape, generator=g, dtype=torch.float32) - variance_noise = torch.randn(shape, generator=g, dtype=torch.float32) - - def _mock_rollout_variance_noise(_batch, *_args, **_kwargs): - # flow_sde_sampling reads the full pre-shard noise from - # rollout_session_data.noise_buffer to compute log_prob, so - # the mock must populate it alongside returning the - # (single-GPU trivially-sharded) noise. - scheduler._get_rollout_session_data( # type: ignore[attr-defined] - _batch - ).noise_buffer = variance_noise - return variance_noise - - scheduler._rollout_variance_noise = ( # type: ignore[method-assign] - _mock_rollout_variance_noise + # The controller draws its variance noise from the request + # generator; a twin generator reproduces it for the reference. + noise_generator = torch.Generator(device="cpu").manual_seed(seed + 100) + variance_noise = torch.randn( + shape, + generator=torch.Generator(device="cpu").manual_seed(seed + 100), + dtype=torch.float32, ) - prev_sgl = scheduler.flow_sde_sampling( + prev_sgl = SchedulerRLMixin.flow_sde_sampling( + batch._rollout_session_data, batch, model_output=model_output, sample=sample, current_sigma=current_sigma, next_sigma=next_sigma, - generator=g, + generator=noise_generator, + step_index=2, ) - log_prob_sum, elem_count = scheduler.consume_local_rollout_log_probs( - batch + log_prob_sum, elem_count = ( + SchedulerRLMixin.consume_local_rollout_log_probs( + batch._rollout_session_data + ) ) log_prob_sum = log_prob_sum.squeeze(-1) elem_count = elem_count.squeeze(-1) @@ -299,7 +306,9 @@ def _mock_rollout_variance_noise(_batch, *_args, **_kwargs): prev_sample_means, noise_std_devs, _model_outputs, - ) = scheduler.consume_local_rollout_debug_tensors(batch) + ) = SchedulerRLMixin.consume_local_rollout_debug_tensors( + batch._rollout_session_data + ) prev_ref, log_prob_ref, prev_mean_ref, noise_std_ref = ( _flowgrpo_sde_step_with_logprob( @@ -369,21 +378,25 @@ def test_sde_cps_force_fp32_with_bf16_model_output(self): # propagates from the (original) model_output.dtype into the # noise buffer. If flow_sde_sampling fails to cast to fp32 at # entry, the buffer is bf16 → log_prob becomes bf16. - scheduler.flow_sde_sampling( + SchedulerRLMixin.flow_sde_sampling( + batch._rollout_session_data, batch, model_output=model_output, sample=sample, current_sigma=current_sigma, next_sigma=next_sigma, generator=g, + step_index=2, + ) + log_prob_sum, _count = SchedulerRLMixin.consume_local_rollout_log_probs( + batch._rollout_session_data ) - log_prob_sum, _count = scheduler.consume_local_rollout_log_probs(batch) self.assertEqual( log_prob_sum.dtype, torch.float32, msg=f"{sde_type}: log_prob_sum must be fp32 with bf16 model_output", ) - noise_buffer = scheduler._get_rollout_session_data(batch).noise_buffer + noise_buffer = batch._rollout_session_data.noise_buffer self.assertEqual( noise_buffer.dtype, torch.float32, @@ -403,7 +416,7 @@ def test_timestep_filters_gate_sde_and_trajectory(self): scheduler = _DummyScheduler() shape = (1, 4, 8, 8) pipeline_config = types.SimpleNamespace( - shard_latents_for_sp=lambda _batch, latents: (latents, False) + shard_latents_for_sp=lambda batch, latents: (latents, False) ) batch = types.SimpleNamespace( rollout_log_prob_no_const=False, @@ -422,48 +435,40 @@ def test_timestep_filters_gate_sde_and_trajectory(self): current_sigma = torch.tensor(0.6, dtype=torch.float32) next_sigma = torch.tensor(0.4, dtype=torch.float32) - variance_noise_ref = torch.randn(shape, generator=g, dtype=torch.float32) - variance_noise_call_count = {"n": 0} - - def _mock_variance_noise(_batch, *_args, **_kwargs): - variance_noise_call_count["n"] += 1 - scheduler._get_rollout_session_data( - _batch - ).noise_buffer = variance_noise_ref - return variance_noise_ref - - scheduler._rollout_variance_noise = ( # type: ignore[method-assign] - _mock_variance_noise - ) + generator_state = g.get_state() # Step 0: not in filter → deterministic ODE transition, no noise draw. - batch._rollout_loop_step_index = 0 - prev_0 = scheduler.flow_sde_sampling( + prev_0 = SchedulerRLMixin.flow_sde_sampling( + batch._rollout_session_data, batch, model_output=model_output, sample=sample, current_sigma=current_sigma, next_sigma=next_sigma, generator=g, + step_index=0, ) - self.assertEqual(variance_noise_call_count["n"], 0) + self.assertTrue(torch.equal(g.get_state(), generator_state)) expected_ode = sample + (next_sigma - current_sigma) * model_output self.assertTrue(torch.allclose(prev_0, expected_ode, atol=1e-6)) # Step 1: in filter → real SDE, noise drawn, prev differs from ODE form. - batch._rollout_loop_step_index = 1 - prev_1 = scheduler.flow_sde_sampling( + prev_1 = SchedulerRLMixin.flow_sde_sampling( + batch._rollout_session_data, batch, model_output=model_output, sample=sample, current_sigma=current_sigma, next_sigma=next_sigma, generator=g, + step_index=1, ) - self.assertEqual(variance_noise_call_count["n"], 1) + self.assertFalse(torch.equal(g.get_state(), generator_state)) self.assertFalse(torch.allclose(prev_1, expected_ode, atol=1e-3)) - log_prob_sum, elem_count = scheduler.consume_local_rollout_log_probs(batch) + log_prob_sum, elem_count = SchedulerRLMixin.consume_local_rollout_log_probs( + batch._rollout_session_data + ) self.assertEqual(tuple(log_prob_sum.shape), (shape[0], 2)) # Filtered step contributes zero log-prob; real SDE step does not. self.assertTrue( From c3f77b16024495f37c995da278040ab4369e9e83 Mon Sep 17 00:00:00 2001 From: rockdu Date: Wed, 14 Oct 2026 12:13:38 -0700 Subject: [PATCH 51/56] refactor(move): move the rollout session methods into RolloutSessionData Refactor-Kind: mechanical_provable Signed-off-by: rockdu --- .../scheduling_flow_match_euler_discrete.py | 3 +- .../post_training/flow_rollout_controller.py | 119 ++++++++---------- .../post_training/rollout_denoising_mixin.py | 10 +- .../test/unit/test_scheduler_rollout_unit.py | 45 +++---- 4 files changed, 72 insertions(+), 105 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py index 4dcf81e3285d..116f8cead96f 100644 --- a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py +++ b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -527,8 +527,7 @@ def step( if batch is not None and batch.rollout: if not self.already_prepared_rollout(batch): raise RuntimeError("Rollout not prepared before step") - prev_sample = SchedulerRLMixin.flow_sde_sampling( - self._get_rollout_session_data(batch), + prev_sample = self._get_rollout_session_data(batch).flow_sde_sampling( batch, model_output, sample, diff --git a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py index 737466aea69d..ae000b1c5fb8 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py +++ b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py @@ -41,18 +41,14 @@ class RolloutSessionData: local_noise_std_devs: list[torch.Tensor] = field(default_factory=list) local_model_outputs: list[torch.Tensor] = field(default_factory=list) - -class SchedulerRLMixin: - @staticmethod - def _reset_rollout_debug_tensors(self: RolloutSessionData) -> None: + def _reset_rollout_debug_tensors(self) -> None: self.local_variance_noises = [] self.local_prev_sample_means = [] self.local_noise_std_devs = [] self.local_model_outputs = [] - @staticmethod def append_local_rollout_debug_tensors( - self: RolloutSessionData, + self, *, variance_noise: torch.Tensor, prev_sample_mean: torch.Tensor, @@ -66,28 +62,24 @@ def append_local_rollout_debug_tensors( self.local_noise_std_devs.append(noise_std_dev.expand((batch_size, 1))) self.local_model_outputs.append(model_output) - @staticmethod def consume_local_rollout_debug_tensors( - self: RolloutSessionData, + self, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: variance_noises = torch.stack(self.local_variance_noises, dim=1) prev_sample_means = torch.stack(self.local_prev_sample_means, dim=1) noise_std_devs = torch.stack(self.local_noise_std_devs, dim=1) model_outputs = torch.stack(self.local_model_outputs, dim=1) - SchedulerRLMixin._reset_rollout_debug_tensors(self) + self._reset_rollout_debug_tensors() return variance_noises, prev_sample_means, noise_std_devs, model_outputs - @staticmethod - def collect_rollout_debug_tensors( - self: RolloutSessionData, batch: Req - ) -> RolloutDebugTensors: + def collect_rollout_debug_tensors(self, batch: Req) -> RolloutDebugTensors: """ Consume rollout debug tensors and merge for all SP ranks. Returns rollout debug tensors with shape [B, T, ...]. """ variance_noises, prev_sample_means, noise_std_devs, model_outputs = ( - SchedulerRLMixin.consume_local_rollout_debug_tensors(self) + self.consume_local_rollout_debug_tensors() ) pipeline_config = self.pipeline_config @@ -115,38 +107,8 @@ def collect_rollout_debug_tensors( rollout_model_outputs=model_outputs.cpu(), ) - @staticmethod - def _get_rollout_session_data(batch) -> RolloutSessionData: - """Return the RolloutSessionData attached to *batch*, or raise if not prepared.""" - rollout_session_data = getattr(batch, "_rollout_session_data", None) - if rollout_session_data is None: - raise RuntimeError("prepare_rollout() not called before rollout") - return rollout_session_data - - def release_rollout_resources(self, batch) -> None: - """Release rollout-owned resources. Call when denoising ends or before a new rollout.""" - batch._rollout_session_data = None - - def prepare_rollout(self, batch: Req, pipeline_config: Any = None) -> None: - """Enable rollout and set SDE/CPS params. Call once before the denoising loop.""" - if get_sp_world_size() > 1 and pipeline_config is None: - raise RuntimeError( - "SP rollout requires pipeline_config to be passed to prepare_rollout()." - ) - batch._rollout_session_data = RolloutSessionData( - pipeline_config=pipeline_config, - sigma_max=self.sigmas[min(1, len(self.sigmas) - 1)].item(), - latents_shape=( - tuple(batch.latents.shape) if batch.latents is not None else None - ), - ) - - def already_prepared_rollout(self, batch) -> bool: - return getattr(batch, "_rollout_session_data", None) is not None - - @staticmethod def _get_or_create_rollout_noise_buffer( - self: RolloutSessionData, + self, full_shape: tuple, device: torch.device, dtype: torch.dtype, @@ -163,9 +125,8 @@ def _get_or_create_rollout_noise_buffer( self.noise_buffer = buffer return buffer - @staticmethod def _rollout_variance_noise( - self: RolloutSessionData, + self, batch, model_output: torch.FloatTensor, generator: Union[torch.Generator, list[torch.Generator]], @@ -186,8 +147,8 @@ def _rollout_variance_noise( "Generator list must have the same length as batch size" ) - buffer = SchedulerRLMixin._get_or_create_rollout_noise_buffer( - self, self.latents_shape, device, dtype + buffer = self._get_or_create_rollout_noise_buffer( + self.latents_shape, device, dtype ) for i in range(B): torch.randn( @@ -206,9 +167,8 @@ def _rollout_variance_noise( ) return sharded_noise - @staticmethod def flow_sde_sampling( - self: RolloutSessionData, + self, batch, model_output: torch.FloatTensor, sample: torch.FloatTensor, @@ -256,8 +216,8 @@ def flow_sde_sampling( if effective_sde_type == "sde": model_output = model_output.float() sample = sample.float() - variance_noise = SchedulerRLMixin._rollout_variance_noise( - self, batch, model_output, generator + variance_noise = self._rollout_variance_noise( + batch, model_output, generator ) full_variance_noise = self.noise_buffer std_dev_t = ( @@ -289,8 +249,8 @@ def flow_sde_sampling( elif effective_sde_type == "cps": model_output = model_output.float() sample = sample.float() - variance_noise = SchedulerRLMixin._rollout_variance_noise( - self, batch, model_output, generator + variance_noise = self._rollout_variance_noise( + batch, model_output, generator ) full_variance_noise = self.noise_buffer std_dev_t = next_sigma * math.sin(noise_level * math.pi / 2) @@ -343,32 +303,27 @@ def flow_sde_sampling( ).sum(dim=list(range(1, len(log_prob_no_const_val.shape)))) if debug_mode: - SchedulerRLMixin.append_local_rollout_debug_tensors( - self, + self.append_local_rollout_debug_tensors( variance_noise=variance_noise, prev_sample_mean=prev_sample_mean, noise_std_dev=noise_std_dev, model_output=model_output, ) - SchedulerRLMixin.append_local_rollout_log_probs( - self, log_prob_local_sum, local_elem_count - ) + self.append_local_rollout_log_probs(log_prob_local_sum, local_elem_count) return prev_sample - @staticmethod def append_local_rollout_log_probs( - self: RolloutSessionData, + self, log_prob_sum: torch.Tensor, log_prob_count: torch.Tensor, ) -> None: self.local_log_prob_sum.append(log_prob_sum) self.local_log_prob_count.append(log_prob_count) - @staticmethod def consume_local_rollout_log_probs( - self: RolloutSessionData, + self, ) -> tuple[torch.Tensor, torch.Tensor]: # [B, T]: batch dim 0, denoising step dim 1 values_sum = torch.stack(self.local_log_prob_sum, dim=1) @@ -377,14 +332,44 @@ def consume_local_rollout_log_probs( self.local_log_prob_count = [] return values_sum, values_count - @staticmethod - def collect_rollout_log_probs(self: RolloutSessionData) -> torch.Tensor | None: + def collect_rollout_log_probs(self) -> torch.Tensor | None: """Per-step sums are already computed on the full pre-shard noise buffer inside flow_sde_sampling, so every SP rank holds identical values here and no all-reduce is needed.""" trajectory_log_prob_sum, trajectory_log_prob_count = ( - SchedulerRLMixin.consume_local_rollout_log_probs(self) + self.consume_local_rollout_log_probs() ) rollout_log_probs_tensor = trajectory_log_prob_sum / trajectory_log_prob_count return rollout_log_probs_tensor.cpu() + + +class SchedulerRLMixin: + @staticmethod + def _get_rollout_session_data(batch) -> RolloutSessionData: + """Return the RolloutSessionData attached to *batch*, or raise if not prepared.""" + rollout_session_data = getattr(batch, "_rollout_session_data", None) + if rollout_session_data is None: + raise RuntimeError("prepare_rollout() not called before rollout") + return rollout_session_data + + def release_rollout_resources(self, batch) -> None: + """Release rollout-owned resources. Call when denoising ends or before a new rollout.""" + batch._rollout_session_data = None + + def prepare_rollout(self, batch: Req, pipeline_config: Any = None) -> None: + """Enable rollout and set SDE/CPS params. Call once before the denoising loop.""" + if get_sp_world_size() > 1 and pipeline_config is None: + raise RuntimeError( + "SP rollout requires pipeline_config to be passed to prepare_rollout()." + ) + batch._rollout_session_data = RolloutSessionData( + pipeline_config=pipeline_config, + sigma_max=self.sigmas[min(1, len(self.sigmas) - 1)].item(), + latents_shape=( + tuple(batch.latents.shape) if batch.latents is not None else None + ), + ) + + def already_prepared_rollout(self, batch) -> bool: + return getattr(batch, "_rollout_session_data", None) is not None 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 908a36f78f92..fc9a0f8df244 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 @@ -67,15 +67,13 @@ def _maybe_collect_rollout_log_probs(self, batch: Req): if batch.rollout_trajectory_data is None: batch.rollout_trajectory_data = RolloutTrajectoryData() batch.rollout_trajectory_data.rollout_log_probs = ( - SchedulerRLMixin.collect_rollout_log_probs( - scheduler._get_rollout_session_data(batch) - ) + scheduler._get_rollout_session_data(batch).collect_rollout_log_probs() ) if batch.rollout_debug_mode: batch.rollout_trajectory_data.rollout_debug_tensors = ( - SchedulerRLMixin.collect_rollout_debug_tensors( - scheduler._get_rollout_session_data(batch), batch - ) + scheduler._get_rollout_session_data( + batch + ).collect_rollout_debug_tensors(batch) ) scheduler.release_rollout_resources(batch) diff --git a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py index 667b764ce335..5de2b942a0bc 100644 --- a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py +++ b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py @@ -46,8 +46,7 @@ def test_ode_step_does_not_call_variance_noise_sampler(self): current_sigma = torch.tensor(0.6, dtype=torch.float32) next_sigma = torch.tensor(0.4, dtype=torch.float32) - prev_sample = SchedulerRLMixin.flow_sde_sampling( - batch._rollout_session_data, + prev_sample = batch._rollout_session_data.flow_sde_sampling( batch, model_output=model_output, sample=sample, @@ -59,9 +58,7 @@ def test_ode_step_does_not_call_variance_noise_sampler(self): # The ODE path must not draw variance noise from the request generator. self.assertTrue(torch.equal(generator.get_state(), generator_state)) log_prob_local_sum, local_elem_count = ( - SchedulerRLMixin.consume_local_rollout_log_probs( - batch._rollout_session_data - ) + batch._rollout_session_data.consume_local_rollout_log_probs() ) log_prob_local_sum = log_prob_local_sum.squeeze(-1) local_elem_count = local_elem_count.squeeze(-1) @@ -93,8 +90,7 @@ def test_ode_bit_exact_with_non_rollout_path(self): next_sigma = torch.tensor(0.4, dtype=torch.float32) dt = next_sigma - current_sigma - rollout_prev = SchedulerRLMixin.flow_sde_sampling( - batch._rollout_session_data, + rollout_prev = batch._rollout_session_data.flow_sde_sampling( batch, model_output=model_output, sample=sample, @@ -138,8 +134,7 @@ def test_ode_debug_tensors_have_shape_safe_noise_std(self): current_sigma = torch.tensor(0.6, dtype=torch.float32) next_sigma = torch.tensor(0.4, dtype=torch.float32) - SchedulerRLMixin.flow_sde_sampling( - batch._rollout_session_data, + batch._rollout_session_data.flow_sde_sampling( batch, model_output=model_output, sample=sample, @@ -154,9 +149,7 @@ def test_ode_debug_tensors_have_shape_safe_noise_std(self): prev_sample_means, noise_std_devs, model_outputs, - ) = SchedulerRLMixin.consume_local_rollout_debug_tensors( - batch._rollout_session_data - ) + ) = batch._rollout_session_data.consume_local_rollout_debug_tensors() # [B, T, ...] with one step in this test. self.assertEqual(tuple(variance_noises.shape), (2, 1, 4, 8, 8)) @@ -284,8 +277,7 @@ def test_single_step_matches_flowgrpo_reference(self): dtype=torch.float32, ) - prev_sgl = SchedulerRLMixin.flow_sde_sampling( - batch._rollout_session_data, + prev_sgl = batch._rollout_session_data.flow_sde_sampling( batch, model_output=model_output, sample=sample, @@ -295,9 +287,7 @@ def test_single_step_matches_flowgrpo_reference(self): step_index=2, ) log_prob_sum, elem_count = ( - SchedulerRLMixin.consume_local_rollout_log_probs( - batch._rollout_session_data - ) + batch._rollout_session_data.consume_local_rollout_log_probs() ) log_prob_sum = log_prob_sum.squeeze(-1) elem_count = elem_count.squeeze(-1) @@ -306,9 +296,7 @@ def test_single_step_matches_flowgrpo_reference(self): prev_sample_means, noise_std_devs, _model_outputs, - ) = SchedulerRLMixin.consume_local_rollout_debug_tensors( - batch._rollout_session_data - ) + ) = batch._rollout_session_data.consume_local_rollout_debug_tensors() prev_ref, log_prob_ref, prev_mean_ref, noise_std_ref = ( _flowgrpo_sde_step_with_logprob( @@ -378,8 +366,7 @@ def test_sde_cps_force_fp32_with_bf16_model_output(self): # propagates from the (original) model_output.dtype into the # noise buffer. If flow_sde_sampling fails to cast to fp32 at # entry, the buffer is bf16 → log_prob becomes bf16. - SchedulerRLMixin.flow_sde_sampling( - batch._rollout_session_data, + batch._rollout_session_data.flow_sde_sampling( batch, model_output=model_output, sample=sample, @@ -388,8 +375,8 @@ def test_sde_cps_force_fp32_with_bf16_model_output(self): generator=g, step_index=2, ) - log_prob_sum, _count = SchedulerRLMixin.consume_local_rollout_log_probs( - batch._rollout_session_data + log_prob_sum, _count = ( + batch._rollout_session_data.consume_local_rollout_log_probs() ) self.assertEqual( log_prob_sum.dtype, @@ -438,8 +425,7 @@ def test_timestep_filters_gate_sde_and_trajectory(self): generator_state = g.get_state() # Step 0: not in filter → deterministic ODE transition, no noise draw. - prev_0 = SchedulerRLMixin.flow_sde_sampling( - batch._rollout_session_data, + prev_0 = batch._rollout_session_data.flow_sde_sampling( batch, model_output=model_output, sample=sample, @@ -453,8 +439,7 @@ def test_timestep_filters_gate_sde_and_trajectory(self): self.assertTrue(torch.allclose(prev_0, expected_ode, atol=1e-6)) # Step 1: in filter → real SDE, noise drawn, prev differs from ODE form. - prev_1 = SchedulerRLMixin.flow_sde_sampling( - batch._rollout_session_data, + prev_1 = batch._rollout_session_data.flow_sde_sampling( batch, model_output=model_output, sample=sample, @@ -466,8 +451,8 @@ def test_timestep_filters_gate_sde_and_trajectory(self): self.assertFalse(torch.equal(g.get_state(), generator_state)) self.assertFalse(torch.allclose(prev_1, expected_ode, atol=1e-3)) - log_prob_sum, elem_count = SchedulerRLMixin.consume_local_rollout_log_probs( - batch._rollout_session_data + log_prob_sum, elem_count = ( + batch._rollout_session_data.consume_local_rollout_log_probs() ) self.assertEqual(tuple(log_prob_sum.shape), (shape[0], 2)) # Filtered step contributes zero log-prob; real SDE step does not. From ba6b224517cf0beab8c54aaa3ae2ff4f03d7ef09 Mon Sep 17 00:00:00 2001 From: rockdu Date: Wed, 14 Oct 2026 12:13:38 -0700 Subject: [PATCH 52/56] refactor(diffusion): turn RolloutSessionData into a per-request FlowRolloutController and drop SchedulerRLMixin Refactor-Kind: non_mechanical_provable Signed-off-by: rockdu --- .../scheduling_flow_match_euler_discrete.py | 12 +- .../runtime/pipelines_core/schedule_batch.py | 9 +- .../post_training/flow_rollout_controller.py | 146 +++++++----------- .../post_training/rollout_denoising_mixin.py | 76 +++++---- .../test/unit/test_scheduler_rollout_unit.py | 137 ++++++++-------- 5 files changed, 172 insertions(+), 208 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py index 116f8cead96f..03407987d946 100644 --- a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py +++ b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -32,9 +32,6 @@ from diffusers.utils import BaseOutput from sglang.multimodal_gen.runtime.models.schedulers.base import BaseScheduler -from sglang.multimodal_gen.runtime.post_training.flow_rollout_controller import ( - SchedulerRLMixin, -) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) @@ -54,9 +51,7 @@ class FlowMatchEulerDiscreteSchedulerOutput(BaseOutput): prev_sample: torch.FloatTensor -class FlowMatchEulerDiscreteScheduler( - SchedulerMixin, ConfigMixin, BaseScheduler, SchedulerRLMixin -): +class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): """ Euler scheduler. @@ -525,10 +520,9 @@ def step( dt = sigma_next - sigma if batch is not None and batch.rollout: - if not self.already_prepared_rollout(batch): + if "visual" not in batch.rollout_controllers: raise RuntimeError("Rollout not prepared before step") - prev_sample = self._get_rollout_session_data(batch).flow_sde_sampling( - batch, + prev_sample = batch.rollout_controllers["visual"].flow_sde_sampling( model_output, sample, current_sigma, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py index 5e07f55b23ef..b7a1b04da621 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py @@ -18,7 +18,7 @@ from collections import Counter from copy import deepcopy from dataclasses import MISSING, asdict, dataclass, field, fields -from typing import Any, Optional, Sequence, Union +from typing import TYPE_CHECKING, Any, Optional, Sequence, Union import PIL.Image import torch @@ -41,6 +41,11 @@ from sglang.multimodal_gen.runtime.utils.perf_logger import RequestMetrics from sglang.srt.observability.trace import TraceNullContext, TraceReqContext +if TYPE_CHECKING: + from sglang.multimodal_gen.runtime.post_training.flow_rollout_controller import ( + FlowRolloutController, + ) + logger = init_logger(__name__) SAMPLING_PARAMS_FIELDS = {f.name for f in fields(SamplingParams)} @@ -194,6 +199,8 @@ class Req: trajectory_timesteps: torch.Tensor | None = None trajectory_latents: torch.Tensor | None = None rollout_trajectory_data: RolloutTrajectoryData | None = None + # Rollout controllers of this request, keyed by latent stream. + rollout_controllers: dict[str, FlowRolloutController] = field(default_factory=dict) trajectory_audio_latents: torch.Tensor | None = None # Extra parameters that might be needed by specific pipeline implementations (e.g., LTX2.3 DenoisingAVStage) diff --git a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py index ae000b1c5fb8..3e0e7c66fc4b 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py +++ b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py @@ -1,45 +1,61 @@ # SPDX-License-Identifier: Apache-2.0 """Flow-matching rollout step utilities for log-prob computation.""" +from __future__ import annotations + import math -from dataclasses import dataclass, field -from typing import Any, Union +from typing import Callable, Union import torch -from sglang.multimodal_gen.runtime.distributed import ( - get_sp_world_size, -) -from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import ( RolloutDebugTensors, ) -from sglang.multimodal_gen.runtime.post_training.sp_utils import ( - gather_stacked_latents_for_sp, -) _LOG_SQRT_2PI = math.log(math.sqrt(2 * math.pi)) -@dataclass -class RolloutSessionData: - """Per-batch rollout state created by prepare_rollout(), lives on the batch object. - - Cleared by setting ``batch._rollout_session_data = None``. +class FlowRolloutController: + """Rollout controller of one latent stream of one request, shared by every + flow-matching model: its SDE/CPS/ODE transition and the per-step log-probs. + The denoising stage creates it in ``batch.rollout_controllers`` and clears + that dict when the loop ends. """ - pipeline_config: Any = None - sigma_max: float = 0.0 - latents_shape: tuple | None = None - noise_buffer: torch.Tensor | None = None - - local_log_prob_sum: list[torch.Tensor] = field(default_factory=list) - local_log_prob_count: list[torch.Tensor] = field(default_factory=list) - - local_variance_noises: list[torch.Tensor] = field(default_factory=list) - local_prev_sample_means: list[torch.Tensor] = field(default_factory=list) - local_noise_std_devs: list[torch.Tensor] = field(default_factory=list) - local_model_outputs: list[torch.Tensor] = field(default_factory=list) + def __init__( + self, + *, + sde_type: str, + noise_level: float, + log_prob_no_const: bool, + debug_mode: bool, + sde_step_indices: list[int] | None, + shard_for_sp: Callable[[torch.Tensor], torch.Tensor], + gather_for_sp: Callable[[torch.Tensor], torch.Tensor], + sigmas: torch.Tensor, + latents_shape: tuple | None, + ) -> None: + self.sde_type = sde_type + self.noise_level = noise_level + self.log_prob_no_const = log_prob_no_const + self.debug_mode = debug_mode + self.sde_step_indices = sde_step_indices + # Shard a full-shape tensor to this rank / gather [B, T, ...] stacks back. + self.shard_for_sp = shard_for_sp + self.gather_for_sp = gather_for_sp + # [T+1] sigma grid of the loop, terminal 0 included. + self.sigmas = sigmas + self.sigma_max = sigmas[min(1, len(sigmas) - 1)].item() + self.latents_shape = latents_shape + self.noise_buffer: torch.Tensor | None = None + + self.local_log_prob_sum: list[torch.Tensor] = [] + self.local_log_prob_count: list[torch.Tensor] = [] + + self.local_variance_noises: list[torch.Tensor] = [] + self.local_prev_sample_means: list[torch.Tensor] = [] + self.local_noise_std_devs: list[torch.Tensor] = [] + self.local_model_outputs: list[torch.Tensor] = [] def _reset_rollout_debug_tensors(self) -> None: self.local_variance_noises = [] @@ -72,7 +88,7 @@ def consume_local_rollout_debug_tensors( self._reset_rollout_debug_tensors() return variance_noises, prev_sample_means, noise_std_devs, model_outputs - def collect_rollout_debug_tensors(self, batch: Req) -> RolloutDebugTensors: + def collect_rollout_debug_tensors(self) -> RolloutDebugTensors: """ Consume rollout debug tensors and merge for all SP ranks. @@ -82,22 +98,9 @@ def collect_rollout_debug_tensors(self, batch: Req) -> RolloutDebugTensors: self.consume_local_rollout_debug_tensors() ) - pipeline_config = self.pipeline_config - variance_noises = gather_stacked_latents_for_sp( - pipeline_config=pipeline_config, - batch=batch, - stacked_latents=variance_noises, - ) - prev_sample_means = gather_stacked_latents_for_sp( - pipeline_config=pipeline_config, - batch=batch, - stacked_latents=prev_sample_means, - ) - model_outputs = gather_stacked_latents_for_sp( - pipeline_config=pipeline_config, - batch=batch, - stacked_latents=model_outputs, - ) + variance_noises = self.gather_for_sp(variance_noises) + prev_sample_means = self.gather_for_sp(prev_sample_means) + model_outputs = self.gather_for_sp(model_outputs) # noise_std_devs is same on every device, not a sharded latent tensor. return RolloutDebugTensors( @@ -127,7 +130,6 @@ def _get_or_create_rollout_noise_buffer( def _rollout_variance_noise( self, - batch, model_output: torch.FloatTensor, generator: Union[torch.Generator, list[torch.Generator]], ) -> torch.FloatTensor: @@ -157,9 +159,7 @@ def _rollout_variance_noise( generator=generator[i], ) - sharded_noise, _ = self.pipeline_config.shard_latents_for_sp( - batch=batch, latents=buffer - ) + sharded_noise = self.shard_for_sp(buffer) if tuple(sharded_noise.shape) != local_shape: raise ValueError( "Rollout SP noise shape mismatch after shard. " @@ -169,7 +169,6 @@ def _rollout_variance_noise( def flow_sde_sampling( self, - batch, model_output: torch.FloatTensor, sample: torch.FloatTensor, current_sigma: torch.FloatTensor, @@ -179,16 +178,16 @@ def flow_sde_sampling( ) -> torch.Tensor: """Flow rollout step for log-prob / sampling (see FlowGRPO-style references). - ``rollout_sde_type`` (from batch SamplingParams): + ``sde_type`` (from the request's ``rollout_sde_type``): 1. ``"sde"``: Standard stochastic differential equation transition (Gaussian). 2. ``"cps"``: Coupled Particle Sampling. 3. ``"ode"``: Deterministic ODE step (no diffusion noise). """ - sde_type = batch.rollout_sde_type - noise_level = float(batch.rollout_noise_level) - log_prob_no_const = batch.rollout_log_prob_no_const - debug_mode = bool(getattr(batch, "rollout_debug_mode", False)) + sde_type = self.sde_type + noise_level = self.noise_level + log_prob_no_const = self.log_prob_no_const + debug_mode = self.debug_mode if not log_prob_no_const and sde_type != "ode": assert noise_level > 0, ( @@ -197,7 +196,7 @@ def flow_sde_sampling( dt = next_sigma - current_sigma - sde_step_indices = getattr(batch, "rollout_sde_step_indices", None) + sde_step_indices = self.sde_step_indices if ( sde_type != "ode" and sde_step_indices is not None @@ -216,9 +215,7 @@ def flow_sde_sampling( if effective_sde_type == "sde": model_output = model_output.float() sample = sample.float() - variance_noise = self._rollout_variance_noise( - batch, model_output, generator - ) + variance_noise = self._rollout_variance_noise(model_output, generator) full_variance_noise = self.noise_buffer std_dev_t = ( torch.sqrt( @@ -249,9 +246,7 @@ def flow_sde_sampling( elif effective_sde_type == "cps": model_output = model_output.float() sample = sample.float() - variance_noise = self._rollout_variance_noise( - batch, model_output, generator - ) + variance_noise = self._rollout_variance_noise(model_output, generator) full_variance_noise = self.noise_buffer std_dev_t = next_sigma * math.sin(noise_level * math.pi / 2) noise_std_dev = std_dev_t @@ -342,34 +337,3 @@ def collect_rollout_log_probs(self) -> torch.Tensor | None: ) rollout_log_probs_tensor = trajectory_log_prob_sum / trajectory_log_prob_count return rollout_log_probs_tensor.cpu() - - -class SchedulerRLMixin: - @staticmethod - def _get_rollout_session_data(batch) -> RolloutSessionData: - """Return the RolloutSessionData attached to *batch*, or raise if not prepared.""" - rollout_session_data = getattr(batch, "_rollout_session_data", None) - if rollout_session_data is None: - raise RuntimeError("prepare_rollout() not called before rollout") - return rollout_session_data - - def release_rollout_resources(self, batch) -> None: - """Release rollout-owned resources. Call when denoising ends or before a new rollout.""" - batch._rollout_session_data = None - - def prepare_rollout(self, batch: Req, pipeline_config: Any = None) -> None: - """Enable rollout and set SDE/CPS params. Call once before the denoising loop.""" - if get_sp_world_size() > 1 and pipeline_config is None: - raise RuntimeError( - "SP rollout requires pipeline_config to be passed to prepare_rollout()." - ) - batch._rollout_session_data = RolloutSessionData( - pipeline_config=pipeline_config, - sigma_max=self.sigmas[min(1, len(self.sigmas) - 1)].item(), - latents_shape=( - tuple(batch.latents.shape) if batch.latents is not None else None - ), - ) - - def already_prepared_rollout(self, batch) -> bool: - return getattr(batch, "_rollout_session_data", None) is not None 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 fc9a0f8df244..47573bfce89a 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 @@ -9,9 +9,12 @@ import torch +from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler, +) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.post_training.flow_rollout_controller import ( - SchedulerRLMixin, + FlowRolloutController, ) from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import ( RolloutDenoisingEnv, @@ -38,44 +41,49 @@ def _kwargs_to_cpu(d: Any) -> Any: class RolloutDenoisingMixin: def _maybe_prepare_rollout(self, batch: Req): - """Prepare denoising loop for rollout.""" - scheduler = batch.scheduler - if not isinstance(scheduler, SchedulerRLMixin): - if batch.rollout: - raise ValueError( - f"Scheduler {type(scheduler)} does not support rollout" - ) - return + """Create the rollout controller of the visual stream of this request. - scheduler.release_rollout_resources(batch) - if batch.rollout: - scheduler.prepare_rollout( - batch=batch, - pipeline_config=self.server_args.pipeline_config, + Must run before SP sharding so the controller records the full latent shape. + """ + scheduler = batch.scheduler + if not isinstance(scheduler, FlowMatchEulerDiscreteScheduler): + raise ValueError(f"Scheduler {type(scheduler)} does not support rollout") + + pipeline_config = self.server_args.pipeline_config + batch.rollout_controllers = { + "visual": FlowRolloutController( + sde_type=batch.rollout_sde_type, + noise_level=float(batch.rollout_noise_level), + log_prob_no_const=batch.rollout_log_prob_no_const, + debug_mode=batch.rollout_debug_mode, + sde_step_indices=batch.rollout_sde_step_indices, + shard_for_sp=lambda latents: pipeline_config.shard_latents_for_sp( + batch=batch, latents=latents + )[0], + gather_for_sp=lambda stacked_latents: gather_stacked_latents_for_sp( + pipeline_config=pipeline_config, + batch=batch, + stacked_latents=stacked_latents, + ), + sigmas=scheduler.sigmas, + latents_shape=( + tuple(batch.latents.shape) if batch.latents is not None else None + ), ) + } def _maybe_collect_rollout_log_probs(self, batch: Req): - scheduler = batch.scheduler - if not isinstance(scheduler, SchedulerRLMixin): - if batch.rollout: - raise ValueError( - f"Scheduler {type(scheduler)} does not support rollout" - ) - return - - if batch.rollout: - if batch.rollout_trajectory_data is None: - batch.rollout_trajectory_data = RolloutTrajectoryData() - batch.rollout_trajectory_data.rollout_log_probs = ( - scheduler._get_rollout_session_data(batch).collect_rollout_log_probs() + rollout = batch.rollout_controllers["visual"] + if batch.rollout_trajectory_data is None: + batch.rollout_trajectory_data = RolloutTrajectoryData() + batch.rollout_trajectory_data.rollout_log_probs = ( + rollout.collect_rollout_log_probs() + ) + if batch.rollout_debug_mode: + batch.rollout_trajectory_data.rollout_debug_tensors = ( + rollout.collect_rollout_debug_tensors() ) - if batch.rollout_debug_mode: - batch.rollout_trajectory_data.rollout_debug_tensors = ( - scheduler._get_rollout_session_data( - batch - ).collect_rollout_debug_tensors(batch) - ) - scheduler.release_rollout_resources(batch) + batch.rollout_controllers = {} def _postprocess_rollout_outputs( self, diff --git a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py index 5de2b942a0bc..a8d175e66823 100644 --- a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py +++ b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py @@ -4,39 +4,44 @@ import torch -import sglang.multimodal_gen.runtime.post_training.flow_rollout_controller as rl_mixin_module from sglang.multimodal_gen.runtime.post_training.flow_rollout_controller import ( - SchedulerRLMixin, + FlowRolloutController, ) +_SIGMAS = torch.tensor([1.0, 0.8, 0.6, 0.4, 0.2, 0.0], dtype=torch.float32) -class _DummyScheduler(SchedulerRLMixin): - def __init__(self): - self.sigmas = torch.tensor([1.0, 0.8, 0.6, 0.4, 0.2, 0.0], dtype=torch.float32) +def _make_rollout(batch, pipeline_config=None) -> FlowRolloutController: + """The controller the denoising stage creates for ``batch`` (single GPU).""" + return FlowRolloutController( + sde_type=batch.rollout_sde_type, + noise_level=float(batch.rollout_noise_level), + log_prob_no_const=batch.rollout_log_prob_no_const, + debug_mode=batch.rollout_debug_mode, + sde_step_indices=batch.rollout_sde_step_indices, + shard_for_sp=lambda latents: pipeline_config.shard_latents_for_sp( + batch=batch, latents=latents + )[0], + gather_for_sp=lambda stacked_latents: stacked_latents, + sigmas=_SIGMAS, + latents_shape=tuple(batch.latents.shape), + ) -class TestSchedulerRolloutOdeUnit(unittest.TestCase): - def setUp(self): - self._orig_get_sp_world_size = rl_mixin_module.get_sp_world_size - rl_mixin_module.get_sp_world_size = lambda: 1 - - def tearDown(self): - rl_mixin_module.get_sp_world_size = self._orig_get_sp_world_size +class TestSchedulerRolloutOdeUnit(unittest.TestCase): def _build_batch(self, *, debug_mode: bool) -> types.SimpleNamespace: return types.SimpleNamespace( rollout_log_prob_no_const=True, rollout_noise_level=0.5, rollout_sde_type="ode", rollout_debug_mode=debug_mode, + rollout_sde_step_indices=None, latents=torch.zeros(2, 4, 8, 8, dtype=torch.float32), - _rollout_session_data=None, ) def test_ode_step_does_not_call_variance_noise_sampler(self): - scheduler = _DummyScheduler() batch = self._build_batch(debug_mode=False) - scheduler.prepare_rollout(batch) + rollout = _make_rollout(batch) generator = torch.Generator().manual_seed(1) generator_state = generator.get_state() @@ -46,8 +51,7 @@ def test_ode_step_does_not_call_variance_noise_sampler(self): current_sigma = torch.tensor(0.6, dtype=torch.float32) next_sigma = torch.tensor(0.4, dtype=torch.float32) - prev_sample = batch._rollout_session_data.flow_sde_sampling( - batch, + prev_sample = rollout.flow_sde_sampling( model_output=model_output, sample=sample, current_sigma=current_sigma, @@ -57,9 +61,7 @@ def test_ode_step_does_not_call_variance_noise_sampler(self): ) # The ODE path must not draw variance noise from the request generator. self.assertTrue(torch.equal(generator.get_state(), generator_state)) - log_prob_local_sum, local_elem_count = ( - batch._rollout_session_data.consume_local_rollout_log_probs() - ) + log_prob_local_sum, local_elem_count = rollout.consume_local_rollout_log_probs() log_prob_local_sum = log_prob_local_sum.squeeze(-1) local_elem_count = local_elem_count.squeeze(-1) @@ -80,9 +82,8 @@ def test_ode_bit_exact_with_non_rollout_path(self): wrapped-scalar promotion difference that a spurious ``model_output.float()`` in the ODE branch would introduce is most visible at bf16 precision.""" - scheduler = _DummyScheduler() batch = self._build_batch(debug_mode=False) - scheduler.prepare_rollout(batch) + rollout = _make_rollout(batch) sample = torch.randn(2, 4, 8, 8, dtype=torch.float32) model_output = torch.randn_like(sample).to(torch.bfloat16) @@ -90,8 +91,7 @@ def test_ode_bit_exact_with_non_rollout_path(self): next_sigma = torch.tensor(0.4, dtype=torch.float32) dt = next_sigma - current_sigma - rollout_prev = batch._rollout_session_data.flow_sde_sampling( - batch, + rollout_prev = rollout.flow_sde_sampling( model_output=model_output, sample=sample, current_sigma=current_sigma, @@ -125,17 +125,15 @@ def test_ode_bit_exact_with_non_rollout_path(self): self.assertEqual(post_cast_max_abs_diff, 0.0) def test_ode_debug_tensors_have_shape_safe_noise_std(self): - scheduler = _DummyScheduler() batch = self._build_batch(debug_mode=True) - scheduler.prepare_rollout(batch) + rollout = _make_rollout(batch) sample = torch.randn(2, 4, 8, 8, dtype=torch.float32) model_output = torch.randn_like(sample) current_sigma = torch.tensor(0.6, dtype=torch.float32) next_sigma = torch.tensor(0.4, dtype=torch.float32) - batch._rollout_session_data.flow_sde_sampling( - batch, + rollout.flow_sde_sampling( model_output=model_output, sample=sample, current_sigma=current_sigma, @@ -149,7 +147,7 @@ def test_ode_debug_tensors_have_shape_safe_noise_std(self): prev_sample_means, noise_std_devs, model_outputs, - ) = batch._rollout_session_data.consume_local_rollout_debug_tensors() + ) = rollout.consume_local_rollout_debug_tensors() # [B, T, ...] with one step in this test. self.assertEqual(tuple(variance_noises.shape), (2, 1, 4, 8, 8)) @@ -228,13 +226,6 @@ def _flowgrpo_sde_step_with_logprob( class TestSchedulerFlowGRPOStepAlignmentUnit(unittest.TestCase): - def setUp(self): - self._orig_get_sp_world_size = rl_mixin_module.get_sp_world_size - rl_mixin_module.get_sp_world_size = lambda: 1 - - def tearDown(self): - rl_mixin_module.get_sp_world_size = self._orig_get_sp_world_size - def _build_batch( self, *, sde_type: str, shape: tuple[int, ...] ) -> types.SimpleNamespace: @@ -243,14 +234,13 @@ def _build_batch( rollout_noise_level=0.5, rollout_sde_type=sde_type, rollout_debug_mode=True, + rollout_sde_step_indices=None, latents=torch.empty(shape, dtype=torch.float32), - _rollout_session_data=None, ) def test_single_step_matches_flowgrpo_reference(self): """Verify prev_sample, prev_sample_mean, noise_std_dev, and log_prob all match FlowGRPO's ``sde_step_with_logprob`` for SDE and CPS.""" - scheduler = _DummyScheduler() current_sigma = torch.tensor(0.5, dtype=torch.float32) next_sigma = torch.tensor(0.3, dtype=torch.float32) shape = (1, 16, 1, 32, 32) @@ -262,8 +252,7 @@ def test_single_step_matches_flowgrpo_reference(self): for sde_type in ("sde", "cps"): for seed in (0, 1, 2, 3): batch = self._build_batch(sde_type=sde_type, shape=shape) - scheduler.release_rollout_resources(batch) - scheduler.prepare_rollout(batch=batch, pipeline_config=pipeline_config) + rollout = _make_rollout(batch, pipeline_config) g = torch.Generator(device="cpu").manual_seed(seed) model_output = torch.randn(shape, generator=g, dtype=torch.float32) @@ -277,8 +266,7 @@ def test_single_step_matches_flowgrpo_reference(self): dtype=torch.float32, ) - prev_sgl = batch._rollout_session_data.flow_sde_sampling( - batch, + prev_sgl = rollout.flow_sde_sampling( model_output=model_output, sample=sample, current_sigma=current_sigma, @@ -286,9 +274,7 @@ def test_single_step_matches_flowgrpo_reference(self): generator=noise_generator, step_index=2, ) - log_prob_sum, elem_count = ( - batch._rollout_session_data.consume_local_rollout_log_probs() - ) + log_prob_sum, elem_count = rollout.consume_local_rollout_log_probs() log_prob_sum = log_prob_sum.squeeze(-1) elem_count = elem_count.squeeze(-1) ( @@ -296,7 +282,7 @@ def test_single_step_matches_flowgrpo_reference(self): prev_sample_means, noise_std_devs, _model_outputs, - ) = batch._rollout_session_data.consume_local_rollout_debug_tensors() + ) = rollout.consume_local_rollout_debug_tensors() prev_ref, log_prob_ref, prev_mean_ref, noise_std_ref = ( _flowgrpo_sde_step_with_logprob( @@ -305,7 +291,7 @@ def test_single_step_matches_flowgrpo_reference(self): variance_noise=variance_noise, sigma=current_sigma, sigma_prev=next_sigma, - sigma_max=scheduler.sigmas[1].item(), + sigma_max=_SIGMAS[1].item(), noise_level=0.5, sde_type=sde_type, ) @@ -343,7 +329,6 @@ def test_sde_cps_force_fp32_with_bf16_model_output(self): branches therefore cast ``model_output.float()`` at entry. Passing bf16 ``model_output`` must still yield an fp32 noise buffer and an fp32 log-prob sum.""" - scheduler = _DummyScheduler() current_sigma = torch.tensor(0.5, dtype=torch.float32) next_sigma = torch.tensor(0.3, dtype=torch.float32) shape = (1, 16, 1, 32, 32) @@ -353,8 +338,7 @@ def test_sde_cps_force_fp32_with_bf16_model_output(self): for sde_type in ("sde", "cps"): batch = self._build_batch(sde_type=sde_type, shape=shape) - scheduler.release_rollout_resources(batch) - scheduler.prepare_rollout(batch=batch, pipeline_config=pipeline_config) + rollout = _make_rollout(batch, pipeline_config) g = torch.Generator(device="cpu").manual_seed(0) model_output = torch.randn(shape, generator=g, dtype=torch.float32).to( @@ -366,8 +350,7 @@ def test_sde_cps_force_fp32_with_bf16_model_output(self): # propagates from the (original) model_output.dtype into the # noise buffer. If flow_sde_sampling fails to cast to fp32 at # entry, the buffer is bf16 → log_prob becomes bf16. - batch._rollout_session_data.flow_sde_sampling( - batch, + rollout.flow_sde_sampling( model_output=model_output, sample=sample, current_sigma=current_sigma, @@ -375,15 +358,13 @@ def test_sde_cps_force_fp32_with_bf16_model_output(self): generator=g, step_index=2, ) - log_prob_sum, _count = ( - batch._rollout_session_data.consume_local_rollout_log_probs() - ) + log_prob_sum, _count = rollout.consume_local_rollout_log_probs() self.assertEqual( log_prob_sum.dtype, torch.float32, msg=f"{sde_type}: log_prob_sum must be fp32 with bf16 model_output", ) - noise_buffer = batch._rollout_session_data.noise_buffer + noise_buffer = rollout.noise_buffer self.assertEqual( noise_buffer.dtype, torch.float32, @@ -400,7 +381,6 @@ def test_timestep_filters_gate_sde_and_trajectory(self): ) # --- Part 1: rollout_sde_step_indices gates SDE noise injection --- - scheduler = _DummyScheduler() shape = (1, 4, 8, 8) pipeline_config = types.SimpleNamespace( shard_latents_for_sp=lambda batch, latents: (latents, False) @@ -412,9 +392,8 @@ def test_timestep_filters_gate_sde_and_trajectory(self): rollout_debug_mode=False, rollout_sde_step_indices=[1], # only step 1 is stochastic latents=torch.empty(shape, dtype=torch.float32), - _rollout_session_data=None, ) - scheduler.prepare_rollout(batch=batch, pipeline_config=pipeline_config) + rollout = _make_rollout(batch, pipeline_config) g = torch.Generator(device="cpu").manual_seed(0) sample = torch.randn(shape, generator=g, dtype=torch.float32) @@ -425,8 +404,7 @@ def test_timestep_filters_gate_sde_and_trajectory(self): generator_state = g.get_state() # Step 0: not in filter → deterministic ODE transition, no noise draw. - prev_0 = batch._rollout_session_data.flow_sde_sampling( - batch, + prev_0 = rollout.flow_sde_sampling( model_output=model_output, sample=sample, current_sigma=current_sigma, @@ -439,8 +417,7 @@ def test_timestep_filters_gate_sde_and_trajectory(self): self.assertTrue(torch.allclose(prev_0, expected_ode, atol=1e-6)) # Step 1: in filter → real SDE, noise drawn, prev differs from ODE form. - prev_1 = batch._rollout_session_data.flow_sde_sampling( - batch, + prev_1 = rollout.flow_sde_sampling( model_output=model_output, sample=sample, current_sigma=current_sigma, @@ -451,9 +428,7 @@ def test_timestep_filters_gate_sde_and_trajectory(self): self.assertFalse(torch.equal(g.get_state(), generator_state)) self.assertFalse(torch.allclose(prev_1, expected_ode, atol=1e-3)) - log_prob_sum, elem_count = ( - batch._rollout_session_data.consume_local_rollout_log_probs() - ) + log_prob_sum, elem_count = rollout.consume_local_rollout_log_probs() self.assertEqual(tuple(log_prob_sum.shape), (shape[0], 2)) # Filtered step contributes zero log-prob; real SDE step does not. self.assertTrue( @@ -479,7 +454,11 @@ class _DummyDit(RolloutDenoisingMixin): rollout=True, rollout_return_dit_trajectory=True, rollout_return_step_indices=[0, 2], - _rollout_denoising_env_state={"step_latents": [], "step_timesteps": []}, + _rollout_denoising_env_state={ + "step_latents": [], + "step_latent_indices": [], + "step_timesteps": [], + }, ) for i in range(3): dit._maybe_append_dit_trajectory_step( @@ -492,7 +471,7 @@ class _DummyDit(RolloutDenoisingMixin): len(traj_filtered._rollout_denoising_env_state["step_latents"]), 2 ) self.assertEqual( - len(traj_filtered._rollout_denoising_env_state["step_timesteps"]), 2 + len(traj_filtered._rollout_denoising_env_state["step_timesteps"]), 3 ) # None (default) → all steps appended (back-compat). @@ -500,7 +479,11 @@ class _DummyDit(RolloutDenoisingMixin): rollout=True, rollout_return_dit_trajectory=True, rollout_return_step_indices=None, - _rollout_denoising_env_state={"step_latents": [], "step_timesteps": []}, + _rollout_denoising_env_state={ + "step_latents": [], + "step_latent_indices": [], + "step_timesteps": [], + }, ) for i in range(3): dit._maybe_append_dit_trajectory_step( @@ -517,7 +500,11 @@ class _DummyDit(RolloutDenoisingMixin): rollout=True, rollout_return_dit_trajectory=True, rollout_return_step_indices=[0, 1, 2], # excludes T=3 - _rollout_denoising_env_state={"step_latents": [], "step_timesteps": []}, + _rollout_denoising_env_state={ + "step_latents": [], + "step_latent_indices": [], + "step_timesteps": [], + }, ) for i in range(3): dit._maybe_append_dit_trajectory_step( @@ -537,7 +524,7 @@ class _DummyDit(RolloutDenoisingMixin): len(traj_exclude_final._rollout_denoising_env_state["step_latents"]), 3 ) self.assertEqual( - len(traj_exclude_final._rollout_denoising_env_state["step_timesteps"]), 3 + len(traj_exclude_final._rollout_denoising_env_state["step_timesteps"]), 4 ) # Filter includes only step_index=T → only the final latent survives. @@ -545,7 +532,11 @@ class _DummyDit(RolloutDenoisingMixin): rollout=True, rollout_return_dit_trajectory=True, rollout_return_step_indices=[3], - _rollout_denoising_env_state={"step_latents": [], "step_timesteps": []}, + _rollout_denoising_env_state={ + "step_latents": [], + "step_latent_indices": [], + "step_timesteps": [], + }, ) for i in range(3): dit._maybe_append_dit_trajectory_step( @@ -564,7 +555,7 @@ class _DummyDit(RolloutDenoisingMixin): len(traj_only_final._rollout_denoising_env_state["step_latents"]), 1 ) self.assertEqual( - len(traj_only_final._rollout_denoising_env_state["step_timesteps"]), 1 + len(traj_only_final._rollout_denoising_env_state["step_timesteps"]), 4 ) From 118bf21f06444f660c9d4b3e5ce9bfc0d3ece500 Mon Sep 17 00:00:00 2001 From: rockdu Date: Wed, 14 Oct 2026 12:13:38 -0700 Subject: [PATCH 53/56] refactor(diffusion): pass the rollout transition to scheduler step instead of branching on batch Refactor-Kind: non_mechanical_provable Co-authored-by: WenhaoZhang <42087078+niehen6174@users.noreply.github.com> Signed-off-by: rockdu --- .../scheduling_flow_match_euler_discrete.py | 42 ++++++++---------- .../pipelines_core/stages/denoising.py | 9 +++- .../stages/model_specific_stages/cosmos3.py | 28 ++++++------ .../model_specific_stages/ltx_2/denoising.py | 29 +++++-------- .../post_training/flow_rollout_controller.py | 43 ++++++++++++++++++- .../post_training/rollout_denoising_mixin.py | 1 + .../test/unit/test_scheduler_rollout_unit.py | 34 +++++++++++++++ 7 files changed, 125 insertions(+), 61 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py index 03407987d946..02c6ed174158 100644 --- a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py +++ b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -22,7 +22,7 @@ # ============================================================================== import math from dataclasses import dataclass -from typing import Any +from typing import Any, Callable import numpy as np import scipy.stats @@ -449,7 +449,7 @@ def step( s_noise: float = 1.0, generator: torch.Generator | None = None, per_token_timesteps: torch.Tensor | None = None, - batch=None, + transition: Callable[..., torch.FloatTensor] | None = None, return_dict: bool = True, ) -> FlowMatchEulerDiscreteSchedulerOutput | tuple[torch.FloatTensor, ...]: """ @@ -472,6 +472,10 @@ def step( A random number generator. per_token_timesteps (`torch.Tensor`, *optional*): The timesteps for each token in the sample. + transition (`Callable`, *optional*): + Replaces this scheduler's own `_transition` for this call, with the + same keyword arguments; rollout passes + `FlowRolloutController.transition` here. return_dict (`bool`): Whether or not to return a [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or tuple. @@ -519,28 +523,18 @@ def step( next_sigma = sigma_next dt = sigma_next - sigma - if batch is not None and batch.rollout: - if "visual" not in batch.rollout_controllers: - raise RuntimeError("Rollout not prepared before step") - prev_sample = batch.rollout_controllers["visual"].flow_sde_sampling( - model_output, - sample, - current_sigma, - next_sigma, - generator, - step_index=self.step_index, - ) - else: - prev_sample = self._transition( - model_output=model_output, - sample=sample, - current_sigma=current_sigma, - next_sigma=next_sigma, - dt=dt, - generator=generator, - step_index=self.step_index, - timestep=timestep, - ) + if transition is None: + transition = self._transition + prev_sample = transition( + model_output=model_output, + sample=sample, + current_sigma=current_sigma, + next_sigma=next_sigma, + dt=dt, + generator=generator, + step_index=self.step_index, + timestep=timestep, + ) # upon completion increase step index by one assert self._step_index is not None, "_step_index should not be None" diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index b675abfd5715..a60bcc84f025 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -153,6 +153,9 @@ AttentionBackendEnum, current_platform, ) +from sglang.multimodal_gen.runtime.post_training.flow_rollout_controller import ( + rollout_transition, +) from sglang.multimodal_gen.runtime.post_training.rollout_denoising_mixin import ( RolloutDenoisingMixin, ) @@ -1283,7 +1286,11 @@ def _prepare_denoising_loop(self, batch: Req, server_args: ServerArgs): # Prepare extra step kwargs for scheduler extra_step_kwargs = self.prepare_extra_func_kwargs( scheduler.step, - {"generator": batch.generator, "eta": batch.eta, "batch": batch}, + { + "generator": batch.generator, + "eta": batch.eta, + "transition": rollout_transition(batch), + }, ) # Setup precision and autocast settings diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py index 3e0b7f43e2b6..dff950bb0a61 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py @@ -59,6 +59,9 @@ VerificationResult, ) from sglang.multimodal_gen.runtime.platforms import current_platform +from sglang.multimodal_gen.runtime.post_training.flow_rollout_controller import ( + rollout_transition, +) from sglang.multimodal_gen.runtime.post_training.rollout_denoising_mixin import ( RolloutDenoisingMixin, ) @@ -1519,6 +1522,7 @@ def _denoise_once( }, guidance=None, ) + transition = rollout_transition(batch) do_cfg = guidance_scale > 1.0 # Control-CFG runs even when text guidance is off (its own extra @@ -1787,22 +1791,14 @@ def _denoise_once( timestep_value=t, step_index=i, ) - latents = scheduler.step( - noise_pred, - t, - latents, - generator=batch.generator, - batch=batch, - return_dict=False, - )[0] - else: - latents = scheduler.step( - noise_pred, - t, - latents, - generator=generator, - return_dict=False, - )[0] + latents = scheduler.step( + noise_pred, + t, + latents, + generator=generator, + transition=transition, + return_dict=False, + )[0] if action_noise_pred is not None: # Zero the velocity at conditioned (clean) action tokens and at diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py index f6594bc7afee..631dfb7bf485 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py @@ -38,6 +38,9 @@ from sglang.multimodal_gen.runtime.platforms import ( current_platform, ) +from sglang.multimodal_gen.runtime.post_training.flow_rollout_controller import ( + rollout_transition, +) from sglang.multimodal_gen.runtime.server_args import ( ServerArgs, is_ltx2_two_stage_pipeline_name, @@ -180,12 +183,6 @@ def __init__( # set per request by _prepare_denoising_loop before the cache-dit hook self._disable_cache_dit_for_request = False - def _scheduler_step_kwargs(self, batch: Req, scheduler) -> dict: - return self.prepare_extra_func_kwargs( - scheduler.step, - {"generator": batch.generator, "eta": batch.eta, "batch": batch}, - ) - @staticmethod def _randn_like_with_batch_generators( reference_tensor: torch.Tensor, batch: Req @@ -1825,19 +1822,13 @@ def _stage2_midpoint_model_call( midpoint_model_call=_stage2_midpoint_model_call, ) else: - if batch.rollout: - ctx.scheduler._step_index = step.step_index - ctx.latents = ctx.scheduler.step( - model_video, - step.t_device, - ctx.latents, - return_dict=False, - **self._scheduler_step_kwargs(batch, ctx.scheduler), - )[0] - else: - ctx.latents = ctx.scheduler.step( - model_video, step.t_device, ctx.latents, return_dict=False - )[0] + ctx.latents = ctx.scheduler.step( + model_video, + step.t_device, + ctx.latents, + transition=rollout_transition(batch), + return_dict=False, + )[0] ctx.audio_latents = ctx.audio_scheduler.step( model_audio, step.t_device, ctx.audio_latents, return_dict=False )[0] diff --git a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py index 3e0e7c66fc4b..f4d4de163d4c 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py +++ b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py @@ -4,7 +4,7 @@ from __future__ import annotations import math -from typing import Callable, Union +from typing import TYPE_CHECKING, Any, Callable, Union import torch @@ -12,9 +12,19 @@ RolloutDebugTensors, ) +if TYPE_CHECKING: + from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req + _LOG_SQRT_2PI = math.log(math.sqrt(2 * math.pi)) +def rollout_transition(batch: Req) -> Callable[..., torch.FloatTensor] | None: + """The ``transition`` to pass to ``scheduler.step`` for this request: the + visual-stream rollout controller, or None to keep the scheduler's own update.""" + rollout = batch.rollout_controllers.get("visual") + return None if rollout is None else rollout.transition + + class FlowRolloutController: """Rollout controller of one latent stream of one request, shared by every flow-matching model: its SDE/CPS/ODE transition and the per-step log-probs. @@ -30,6 +40,7 @@ def __init__( log_prob_no_const: bool, debug_mode: bool, sde_step_indices: list[int] | None, + generator: Union[torch.Generator, list[torch.Generator]], shard_for_sp: Callable[[torch.Tensor], torch.Tensor], gather_for_sp: Callable[[torch.Tensor], torch.Tensor], sigmas: torch.Tensor, @@ -40,6 +51,7 @@ def __init__( self.log_prob_no_const = log_prob_no_const self.debug_mode = debug_mode self.sde_step_indices = sde_step_indices + self.generator = generator # Shard a full-shape tensor to this rank / gather [B, T, ...] stacks back. self.shard_for_sp = shard_for_sp self.gather_for_sp = gather_for_sp @@ -110,6 +122,35 @@ def collect_rollout_debug_tensors(self) -> RolloutDebugTensors: rollout_model_outputs=model_outputs.cpu(), ) + def transition( + self, + *, + model_output: torch.FloatTensor, + sample: torch.FloatTensor, + current_sigma: torch.FloatTensor, + next_sigma: torch.FloatTensor, + dt: torch.FloatTensor, + generator: Any, + step_index: int, + timestep: Any, + ) -> torch.FloatTensor: + """Stand-in for ``scheduler._transition``, passed as ``step(transition=...)``; + ``sample`` arrives already upcast to fp32 by ``step()``. + + Draws variance noise from the request generator captured at creation, so + the caller's ``generator`` argument is not used. + """ + if current_sigma.dim() != 0: + raise ValueError("Rollout does not support per-token timesteps") + return self.flow_sde_sampling( + model_output=model_output, + sample=sample, + current_sigma=current_sigma, + next_sigma=next_sigma, + generator=self.generator, + step_index=step_index, + ) + def _get_or_create_rollout_noise_buffer( self, full_shape: tuple, 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 47573bfce89a..4c4429b15f46 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 @@ -57,6 +57,7 @@ def _maybe_prepare_rollout(self, batch: Req): log_prob_no_const=batch.rollout_log_prob_no_const, debug_mode=batch.rollout_debug_mode, sde_step_indices=batch.rollout_sde_step_indices, + generator=batch.generator, shard_for_sp=lambda latents: pipeline_config.shard_latents_for_sp( batch=batch, latents=latents )[0], diff --git a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py index a8d175e66823..1e2c0a481e83 100644 --- a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py +++ b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py @@ -4,6 +4,9 @@ import torch +from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler, +) from sglang.multimodal_gen.runtime.post_training.flow_rollout_controller import ( FlowRolloutController, ) @@ -19,6 +22,7 @@ def _make_rollout(batch, pipeline_config=None) -> FlowRolloutController: log_prob_no_const=batch.rollout_log_prob_no_const, debug_mode=batch.rollout_debug_mode, sde_step_indices=batch.rollout_sde_step_indices, + generator=torch.Generator().manual_seed(0), shard_for_sp=lambda latents: pipeline_config.shard_latents_for_sp( batch=batch, latents=latents )[0], @@ -124,6 +128,36 @@ def test_ode_bit_exact_with_non_rollout_path(self): self.assertEqual(pre_cast_max_abs_diff, 0.0) self.assertEqual(post_cast_max_abs_diff, 0.0) + def test_ode_transition_matches_serving_step(self): + """Passing the rollout transition to a real FlowMatchEulerDiscreteScheduler + must reproduce rollout=False bit for bit in ODE mode: step() keeps its + index, sigma lookup and dtype casts around either transition.""" + serving, rollout_scheduler = ( + FlowMatchEulerDiscreteScheduler(shift=3.0), + FlowMatchEulerDiscreteScheduler(shift=3.0), + ) + for scheduler in (serving, rollout_scheduler): + scheduler.set_timesteps(5) + scheduler.set_begin_index(0) + rollout = _make_rollout(self._build_batch(debug_mode=False)) + + g = torch.Generator().manual_seed(0) + sample = torch.randn(2, 4, 8, 8) + expected, actual = sample.clone(), sample.clone() + for t in serving.timesteps: + model_output = torch.randn(sample.shape, generator=g).to(torch.bfloat16) + expected = serving.step(model_output, t, expected, return_dict=False)[0] + actual = rollout_scheduler.step( + model_output, + t, + actual, + transition=rollout.transition, + return_dict=False, + )[0] + self.assertTrue(torch.equal(actual, expected)) + self.assertEqual(rollout_scheduler.step_index, serving.step_index) + self.assertEqual(len(rollout.local_log_prob_sum), len(serving.timesteps)) + def test_ode_debug_tensors_have_shape_safe_noise_std(self): batch = self._build_batch(debug_mode=True) rollout = _make_rollout(batch) From 88adb216b3b15cd3c5f82cc3e18670d462eb0fdb Mon Sep 17 00:00:00 2001 From: rockdu Date: Wed, 14 Oct 2026 12:13:38 -0700 Subject: [PATCH 54/56] refactor(diffusion): record rollout latents in the transition instead of the denoise loops Refactor-Kind: non_mechanical_provable Signed-off-by: rockdu --- .../pipelines_core/stages/denoising.py | 12 -- .../stages/model_specific_stages/cosmos3.py | 8 - .../post_training/flow_rollout_controller.py | 31 ++- .../post_training/rollout_denoising_mixin.py | 47 +---- .../test/unit/test_scheduler_rollout_unit.py | 199 ++++++++---------- 5 files changed, 128 insertions(+), 169 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index a60bcc84f025..0645104e9f86 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -2104,18 +2104,6 @@ def _denoise( t_host, timesteps_cpu, ) - # Capture the raw (pre-scale, pre-I2V-concat) noisy latent - # x_{t_i} for rollout trajectory collection. Must run - # BEFORE _run_denoising_step so ctx.latents is still the - # pre-step value. Gated on batch.rollout to keep the - # non-rollout path strictly untouched. - if batch.rollout: - self._maybe_append_dit_trajectory_step( - batch=batch, - latents=ctx.latents, - timestep_value=step.t_host, - step_index=step_index, - ) self._run_denoising_step(ctx, step, batch, server_args) self._record_trajectory(ctx, step, batch, server_args) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py index dff950bb0a61..57e6721c35d7 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py @@ -1783,14 +1783,6 @@ def _denoise_once( if velocity_mask is not None: noise_pred = noise_pred * velocity_mask - if batch.rollout: - # Capture the pre-step x_{t_i} before the scheduler advances it. - self._maybe_append_dit_trajectory_step( - batch=batch, - latents=latents, - timestep_value=t, - step_index=i, - ) latents = scheduler.step( noise_pred, t, diff --git a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py index f4d4de163d4c..0173822563a8 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py +++ b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py @@ -41,6 +41,8 @@ def __init__( debug_mode: bool, sde_step_indices: list[int] | None, generator: Union[torch.Generator, list[torch.Generator]], + return_dit_trajectory: bool, + return_step_indices: list[int] | None, shard_for_sp: Callable[[torch.Tensor], torch.Tensor], gather_for_sp: Callable[[torch.Tensor], torch.Tensor], sigmas: torch.Tensor, @@ -52,6 +54,8 @@ def __init__( self.debug_mode = debug_mode self.sde_step_indices = sde_step_indices self.generator = generator + self.return_dit_trajectory = return_dit_trajectory + self.return_step_indices = return_step_indices # Shard a full-shape tensor to this rank / gather [B, T, ...] stacks back. self.shard_for_sp = shard_for_sp self.gather_for_sp = gather_for_sp @@ -69,6 +73,10 @@ def __init__( self.local_noise_std_devs: list[torch.Tensor] = [] self.local_model_outputs: list[torch.Tensor] = [] + self.step_latents: list[torch.Tensor] = [] + self.step_latent_indices: list[int] = [] + self.step_timesteps: list[torch.Tensor] = [] + def _reset_rollout_debug_tensors(self) -> None: self.local_variance_noises = [] self.local_prev_sample_means = [] @@ -138,10 +146,14 @@ def transition( ``sample`` arrives already upcast to fp32 by ``step()``. Draws variance noise from the request generator captured at creation, so - the caller's ``generator`` argument is not used. + the caller's ``generator`` argument is not used. Records ``sample`` as the + step's x_t before advancing it. """ if current_sigma.dim() != 0: raise ValueError("Rollout does not support per-token timesteps") + self._maybe_append_dit_trajectory_step( + latents=sample, timestep_value=timestep, step_index=step_index + ) return self.flow_sde_sampling( model_output=model_output, sample=sample, @@ -378,3 +390,20 @@ def collect_rollout_log_probs(self) -> torch.Tensor | None: ) rollout_log_probs_tensor = trajectory_log_prob_sum / trajectory_log_prob_count return rollout_log_probs_tensor.cpu() + + def _maybe_append_dit_trajectory_step( + self, + latents: torch.Tensor, + timestep_value: torch.Tensor, + step_index: int, + ) -> None: + if not self.return_dit_trajectory: + return + + self.step_timesteps.append(timestep_value.detach().cpu()) + return_step_indices = self.return_step_indices + if return_step_indices is not None and step_index not in return_step_indices: + return + + self.step_latents.append(latents.detach()) + self.step_latent_indices.append(step_index) 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 4c4429b15f46..a77d0d8ee591 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 @@ -58,6 +58,8 @@ def _maybe_prepare_rollout(self, batch: Req): debug_mode=batch.rollout_debug_mode, sde_step_indices=batch.rollout_sde_step_indices, generator=batch.generator, + return_dit_trajectory=batch.rollout_return_dit_trajectory, + return_step_indices=batch.rollout_return_step_indices, shard_for_sp=lambda latents: pipeline_config.shard_latents_for_sp( batch=batch, latents=latents )[0], @@ -84,7 +86,6 @@ def _maybe_collect_rollout_log_probs(self, batch: Req): batch.rollout_trajectory_data.rollout_debug_tensors = ( rollout.collect_rollout_debug_tensors() ) - batch.rollout_controllers = {} def _postprocess_rollout_outputs( self, @@ -104,8 +105,7 @@ def _postprocess_rollout_outputs( # Append final denoised latent as the (T+1)-th entry (step_index=T), # routed through the same filter so rollout_return_step_indices can # include/exclude it. - self._maybe_append_dit_trajectory_step( - batch=batch, + batch.rollout_controllers["visual"]._maybe_append_dit_trajectory_step( latents=latents, timestep_value=final_timestep, step_index=num_inference_steps, @@ -114,6 +114,7 @@ def _postprocess_rollout_outputs( batch=batch, pipeline_config=server_args.pipeline_config, ) + batch.rollout_controllers = {} def _maybe_init_denoising_env_collection( self, @@ -148,59 +149,33 @@ 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, } - def _maybe_append_dit_trajectory_step( - self, - batch, - latents: torch.Tensor, - timestep_value: torch.Tensor, - step_index: int, - ) -> None: - if not batch.rollout or not batch.rollout_return_dit_trajectory: - return - state = getattr(batch, "_rollout_denoising_env_state", None) - 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_latent_indices"].append(step_index) - def _maybe_finalize_denoising_env_collection(self, batch, pipeline_config) -> None: state = getattr(batch, "_rollout_denoising_env_state", None) if state is None: return env: RolloutDenoisingEnv | None = state["env"] - step_latents: list[torch.Tensor] = state["step_latents"] - step_timesteps: list[torch.Tensor] = state["step_timesteps"] + rollout = batch.rollout_controllers["visual"] + step_latents: list[torch.Tensor] = rollout.step_latents + step_timesteps: list[torch.Tensor] = rollout.step_timesteps if batch.rollout_trajectory_data is None: batch.rollout_trajectory_data = RolloutTrajectoryData() if step_latents and batch.rollout_return_dit_trajectory: - step_latents_tensor = torch.stack(step_latents, dim=1) - step_latents_tensor = gather_stacked_latents_for_sp( - pipeline_config=pipeline_config, - batch=batch, - stacked_latents=step_latents_tensor, + step_latents_tensor = rollout.gather_for_sp( + torch.stack(step_latents, dim=1) ) batch.rollout_trajectory_data.dit_trajectory = RolloutDitTrajectory( latents=step_latents_tensor.cpu(), timesteps=torch.stack(step_timesteps, dim=0).cpu(), - sigmas=batch.scheduler.sigmas.detach().cpu().clone(), + sigmas=rollout.sigmas.detach().cpu().clone(), latent_step_indices=torch.tensor( - state["step_latent_indices"], dtype=torch.long + rollout.step_latent_indices, dtype=torch.long ), ) diff --git a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py index 1e2c0a481e83..58400b97d8c3 100644 --- a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py +++ b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py @@ -1,3 +1,28 @@ +"""Unit tests for FlowRolloutController, whose transition rollout passes to +scheduler.step(). + + scheduler.step(model_output, t, x_t, transition=rollout.transition) + | upcast x_t to fp32, look up sigma_t / sigma_next, dt = sigma_next - sigma_t + v + rollout.transition(...) + |-- record x_t into the trajectory (filtered by return_step_indices) + +-- flow_sde_sampling(...) + |-- ode : x_next = x_t + dt * v (no noise, log_prob 0) + |-- sde / cps: x_next = mean + std * noise (fp32, log_prob summed) + +-- step outside sde_step_indices -> ode form + v + x_next (cast back by step()) + +TestSchedulerRolloutOdeUnit + ODE rollout == serving step, bit for bit (bare update and through step()); + ODE draws no noise; the trajectory records x_t, not x_next; + debug tensors keep [B, T, ...] shapes. +TestSchedulerFlowGRPOStepAlignmentUnit + SDE / CPS match the verbatim FlowGRPO reference; bf16 model_output still + yields fp32 noise and log-probs; sde_step_indices gates the noise draw; + return_step_indices filters latents while timesteps cover every step. +""" + import math import types import unittest @@ -23,6 +48,8 @@ def _make_rollout(batch, pipeline_config=None) -> FlowRolloutController: debug_mode=batch.rollout_debug_mode, sde_step_indices=batch.rollout_sde_step_indices, generator=torch.Generator().manual_seed(0), + return_dit_trajectory=batch.rollout_return_dit_trajectory, + return_step_indices=batch.rollout_return_step_indices, shard_for_sp=lambda latents: pipeline_config.shard_latents_for_sp( batch=batch, latents=latents )[0], @@ -40,6 +67,8 @@ def _build_batch(self, *, debug_mode: bool) -> types.SimpleNamespace: rollout_sde_type="ode", rollout_debug_mode=debug_mode, rollout_sde_step_indices=None, + rollout_return_dit_trajectory=False, + rollout_return_step_indices=None, latents=torch.zeros(2, 4, 8, 8, dtype=torch.float32), ) @@ -158,6 +187,27 @@ def test_ode_transition_matches_serving_step(self): self.assertEqual(rollout_scheduler.step_index, serving.step_index) self.assertEqual(len(rollout.local_log_prob_sum), len(serving.timesteps)) + def test_transition_records_pre_step_latents(self): + """The rollout transition records x_t (the latents entering the step), + not the updated x_{t-1} it returns.""" + scheduler = FlowMatchEulerDiscreteScheduler(shift=3.0) + scheduler.set_timesteps(5) + scheduler.set_begin_index(0) + batch = self._build_batch(debug_mode=False) + batch.rollout_return_dit_trajectory = True + rollout = _make_rollout(batch) + sample = torch.randn(2, 4, 8, 8) + model_output = torch.randn(2, 4, 8, 8) + prev_sample = scheduler.step( + model_output, + scheduler.timesteps[0], + sample, + transition=rollout.transition, + return_dict=False, + )[0] + self.assertTrue(torch.equal(rollout.step_latents[0], sample)) + self.assertFalse(torch.equal(rollout.step_latents[0], prev_sample)) + def test_ode_debug_tensors_have_shape_safe_noise_std(self): batch = self._build_batch(debug_mode=True) rollout = _make_rollout(batch) @@ -269,6 +319,8 @@ def _build_batch( rollout_sde_type=sde_type, rollout_debug_mode=True, rollout_sde_step_indices=None, + rollout_return_dit_trajectory=False, + rollout_return_step_indices=None, latents=torch.empty(shape, dtype=torch.float32), ) @@ -410,9 +462,6 @@ def test_timestep_filters_gate_sde_and_trajectory(self): injection (excluded steps = ODE transition + zero log-prob); independently, rollout_return_step_indices gates the dit_trajectory append. Both features are exercised here because they share the same step_index predicate.""" - from sglang.multimodal_gen.runtime.post_training.rollout_denoising_mixin import ( - RolloutDenoisingMixin, - ) # --- Part 1: rollout_sde_step_indices gates SDE noise injection --- shape = (1, 4, 8, 8) @@ -425,6 +474,8 @@ def test_timestep_filters_gate_sde_and_trajectory(self): rollout_sde_type="sde", rollout_debug_mode=False, rollout_sde_step_indices=[1], # only step 1 is stochastic + rollout_return_dit_trajectory=False, + rollout_return_step_indices=None, latents=torch.empty(shape, dtype=torch.float32), ) rollout = _make_rollout(batch, pipeline_config) @@ -476,121 +527,45 @@ def test_timestep_filters_gate_sde_and_trajectory(self): self.assertTrue(torch.all(elem_count > 0)) # --- Part 2: rollout_return_step_indices gates dit trajectory append --- - class _DummyDit(RolloutDenoisingMixin): - pass - - dit = _DummyDit() lat = torch.zeros(1, 4, 8, 8) ts = torch.tensor(0.5) - # Filter [0, 2] over steps 0,1,2 → steps 0 and 2 appended, step 1 skipped. - traj_filtered = types.SimpleNamespace( - rollout=True, - rollout_return_dit_trajectory=True, - rollout_return_step_indices=[0, 2], - _rollout_denoising_env_state={ - "step_latents": [], - "step_latent_indices": [], - "step_timesteps": [], - }, - ) - for i in range(3): - dit._maybe_append_dit_trajectory_step( - batch=traj_filtered, - latents=lat, - timestep_value=ts, - step_index=i, - ) - self.assertEqual( - len(traj_filtered._rollout_denoising_env_state["step_latents"]), 2 - ) - self.assertEqual( - len(traj_filtered._rollout_denoising_env_state["step_timesteps"]), 3 - ) - - # None (default) → all steps appended (back-compat). - traj_all = types.SimpleNamespace( - rollout=True, - rollout_return_dit_trajectory=True, - rollout_return_step_indices=None, - _rollout_denoising_env_state={ - "step_latents": [], - "step_latent_indices": [], - "step_timesteps": [], - }, - ) - for i in range(3): - dit._maybe_append_dit_trajectory_step( - batch=traj_all, - latents=lat, - timestep_value=ts, - step_index=i, + def _recorded_steps(return_step_indices: list[int] | None) -> list[int]: + # Three loop steps, then the final latent appended at step_index=T=3. + rollout = _make_rollout( + types.SimpleNamespace( + rollout_log_prob_no_const=False, + rollout_noise_level=0.5, + rollout_sde_type="sde", + rollout_debug_mode=False, + rollout_sde_step_indices=None, + rollout_return_dit_trajectory=True, + rollout_return_step_indices=return_step_indices, + latents=lat, + ) ) - self.assertEqual(len(traj_all._rollout_denoising_env_state["step_latents"]), 3) - - # Filter excludes step_index=T (the final/(T+1)-th latent appended by - # _postprocess_rollout_outputs). Simulate T=3 loop steps + final append. - traj_exclude_final = types.SimpleNamespace( - rollout=True, - rollout_return_dit_trajectory=True, - rollout_return_step_indices=[0, 1, 2], # excludes T=3 - _rollout_denoising_env_state={ - "step_latents": [], - "step_latent_indices": [], - "step_timesteps": [], - }, - ) - for i in range(3): - dit._maybe_append_dit_trajectory_step( - batch=traj_exclude_final, - latents=lat, - timestep_value=ts, - step_index=i, + for i in range(3): + rollout._maybe_append_dit_trajectory_step( + latents=lat, timestep_value=ts, step_index=i + ) + rollout._maybe_append_dit_trajectory_step( + latents=lat, timestep_value=torch.zeros(()), step_index=3 ) - # Mimic the final append routed through the same filter. - dit._maybe_append_dit_trajectory_step( - batch=traj_exclude_final, - latents=lat, - timestep_value=torch.zeros(()), - step_index=3, - ) - self.assertEqual( - len(traj_exclude_final._rollout_denoising_env_state["step_latents"]), 3 - ) - self.assertEqual( - len(traj_exclude_final._rollout_denoising_env_state["step_timesteps"]), 4 - ) - - # Filter includes only step_index=T → only the final latent survives. - traj_only_final = types.SimpleNamespace( - rollout=True, - rollout_return_dit_trajectory=True, - rollout_return_step_indices=[3], - _rollout_denoising_env_state={ - "step_latents": [], - "step_latent_indices": [], - "step_timesteps": [], - }, - ) - for i in range(3): - dit._maybe_append_dit_trajectory_step( - batch=traj_only_final, - latents=lat, - timestep_value=ts, - step_index=i, + # Timesteps cover every step; only latents follow the filter. + self.assertEqual(len(rollout.step_timesteps), 4) + self.assertEqual( + len(rollout.step_latents), len(rollout.step_latent_indices) ) - dit._maybe_append_dit_trajectory_step( - batch=traj_only_final, - latents=lat, - timestep_value=torch.zeros(()), - step_index=3, - ) - self.assertEqual( - len(traj_only_final._rollout_denoising_env_state["step_latents"]), 1 - ) - self.assertEqual( - len(traj_only_final._rollout_denoising_env_state["step_timesteps"]), 4 - ) + return rollout.step_latent_indices + + # Filter [0, 2] over steps 0,1,2 -> steps 0 and 2 appended, step 1 skipped. + self.assertEqual(_recorded_steps([0, 2]), [0, 2]) + # None (default) -> all steps appended, including the final latent. + self.assertEqual(_recorded_steps(None), [0, 1, 2, 3]) + # Filter excludes step_index=T (the final/(T+1)-th latent). + self.assertEqual(_recorded_steps([0, 1, 2]), [0, 1, 2]) + # Filter includes only step_index=T -> only the final latent survives. + self.assertEqual(_recorded_steps([3]), [3]) if __name__ == "__main__": From 6286565d03340fdacc789d4fe749279b22e94696 Mon Sep 17 00:00:00 2001 From: rockdu Date: Wed, 14 Oct 2026 12:13:38 -0700 Subject: [PATCH 55/56] refactor(diffusion): replace the rollout denoising mixin with rollout hook functions Refactor-Kind: non_mechanical_provable Co-authored-by: WenhaoZhang <42087078+niehen6174@users.noreply.github.com> Signed-off-by: rockdu --- .../pipelines_core/stages/denoising.py | 44 ++-- .../stages/model_specific_stages/cosmos3.py | 58 ++--- .../post_training/rollout_denoising_mixin.py | 199 ---------------- .../runtime/post_training/rollout_hooks.py | 220 ++++++++++++++++++ 4 files changed, 272 insertions(+), 249 deletions(-) delete mode 100644 python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py create mode 100644 python/sglang/multimodal_gen/runtime/post_training/rollout_hooks.py diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 0645104e9f86..542fe742f608 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -156,8 +156,10 @@ from sglang.multimodal_gen.runtime.post_training.flow_rollout_controller import ( rollout_transition, ) -from sglang.multimodal_gen.runtime.post_training.rollout_denoising_mixin import ( - RolloutDenoisingMixin, +from sglang.multimodal_gen.runtime.post_training.rollout_hooks import ( + finalize_rollout_if_enabled, + init_rollout_env_if_enabled, + prepare_rollout_if_enabled, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger @@ -331,7 +333,7 @@ class DualTransformerExecutionMode(str, Enum): PAIRED_PER_STEP = "paired_per_step" -class DenoisingStage(PipelineStage, RolloutDenoisingMixin): +class DenoisingStage(PipelineStage): """ Stage for running the denoising loop in diffusion pipelines. @@ -1280,8 +1282,9 @@ def _prepare_denoising_loop(self, batch: Req, server_args: ServerArgs): if freshly_loaded: register_loaded_transformer(self, server_args, pipeline) - if batch.rollout: - self._maybe_prepare_rollout(batch) + prepare_rollout_if_enabled( + batch=batch, pipeline_config=server_args.pipeline_config + ) # Prepare extra step kwargs for scheduler extra_step_kwargs = self.prepare_extra_func_kwargs( @@ -2052,15 +2055,13 @@ def _denoise( Run the denoising loop. """ ctx = self._prepare_denoising_loop(batch, server_args) - if batch.rollout: - self._maybe_init_denoising_env_collection( - batch=batch, - pipeline_config=server_args.pipeline_config, - image_kwargs=ctx.image_kwargs, - pos_cond_kwargs=ctx.pos_cond_kwargs, - neg_cond_kwargs=ctx.neg_cond_kwargs, - guidance=ctx.guidance, - ) + init_rollout_env_if_enabled( + batch=batch, + image_kwargs=ctx.image_kwargs, + pos_cond_kwargs=ctx.pos_cond_kwargs, + neg_cond_kwargs=ctx.neg_cond_kwargs, + guidance=ctx.guidance, + ) denoising_start_time = time.time() self._before_denoising_loop(ctx, batch, server_args) # to avoid device-sync caused by timestep comparison @@ -2133,14 +2134,13 @@ def _denoise( # the final scheduler.step output (ctx.latents) is still SP-sharded and # can be gathered uniformly alongside the per-step dit_trajectory via # gather_stacked_latents_for_sp. - if batch.rollout: - self._postprocess_rollout_outputs( - batch=batch, - latents=ctx.latents, - num_inference_steps=num_timesteps, - final_timestep=timesteps_cpu.new_zeros(()), - server_args=server_args, - ) + finalize_rollout_if_enabled( + batch=batch, + latents=ctx.latents, + num_inference_steps=num_timesteps, + final_timestep=timesteps_cpu.new_zeros(()), + pipeline_config=server_args.pipeline_config, + ) self._finalize_denoising_loop(ctx, batch, server_args) return batch diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py index 57e6721c35d7..9e6c983110e1 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py @@ -62,8 +62,10 @@ from sglang.multimodal_gen.runtime.post_training.flow_rollout_controller import ( rollout_transition, ) -from sglang.multimodal_gen.runtime.post_training.rollout_denoising_mixin import ( - RolloutDenoisingMixin, +from sglang.multimodal_gen.runtime.post_training.rollout_hooks import ( + finalize_rollout_if_enabled, + init_rollout_env_if_enabled, + prepare_rollout_if_enabled, ) from sglang.multimodal_gen.runtime.post_training.rollout_scheduler import ( prepare_rollout_request_scheduler, @@ -1065,7 +1067,7 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req: return batch -class Cosmos3DenoisingStage(PipelineStage, RolloutDenoisingMixin): +class Cosmos3DenoisingStage(PipelineStage): def default_workload_iterations( self, batch: Req, num_inference_steps: int ) -> int | None: @@ -1505,23 +1507,24 @@ def _denoise_once( raise ValueError( "Cosmos3 rollout does not support action/sound modalities." ) - self._maybe_prepare_rollout(batch) - self._maybe_init_denoising_env_collection( - batch=batch, - pipeline_config=server_args.pipeline_config, - image_kwargs={}, - pos_cond_kwargs={ - "text_ids": cond_text_ids, - "text_mask": cond_text_mask, - "fps": fps, - }, - neg_cond_kwargs={ - "text_ids": uncond_text_ids, - "text_mask": uncond_text_mask, - "fps": fps, - }, - guidance=None, - ) + prepare_rollout_if_enabled( + batch=batch, pipeline_config=server_args.pipeline_config + ) + init_rollout_env_if_enabled( + batch=batch, + image_kwargs={}, + pos_cond_kwargs={ + "text_ids": cond_text_ids, + "text_mask": cond_text_mask, + "fps": fps, + }, + neg_cond_kwargs={ + "text_ids": uncond_text_ids, + "text_mask": uncond_text_mask, + "fps": fps, + }, + guidance=None, + ) transition = rollout_transition(batch) do_cfg = guidance_scale > 1.0 @@ -1839,14 +1842,13 @@ def _denoise_once( # next request's steps. self.transformer.reset_denoising_step() - if batch.rollout: - self._postprocess_rollout_outputs( - batch=batch, - latents=latents, - num_inference_steps=len(timesteps), - final_timestep=timesteps.new_zeros(()).cpu(), - server_args=server_args, - ) + finalize_rollout_if_enabled( + batch=batch, + latents=latents, + num_inference_steps=len(timesteps), + final_timestep=timesteps.new_zeros(()).cpu(), + pipeline_config=server_args.pipeline_config, + ) batch.latents = latents if action_latents is not None: 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 deleted file mode 100644 index a77d0d8ee591..000000000000 --- a/python/sglang/multimodal_gen/runtime/post_training/rollout_denoising_mixin.py +++ /dev/null @@ -1,199 +0,0 @@ -"""Mixin for rollout-related denoising hooks. - -Moved out of DenoisingStage to keep the core stage lean. -""" - -from __future__ import annotations - -from typing import Any - -import torch - -from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import ( - FlowMatchEulerDiscreteScheduler, -) -from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req -from sglang.multimodal_gen.runtime.post_training.flow_rollout_controller import ( - FlowRolloutController, -) -from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import ( - RolloutDenoisingEnv, - RolloutDitTrajectory, - RolloutTrajectoryData, -) -from sglang.multimodal_gen.runtime.post_training.sp_utils import ( - gather_stacked_latents_for_sp, -) -from sglang.multimodal_gen.runtime.server_args import ServerArgs - - -def _kwargs_to_cpu(d: Any) -> Any: - if isinstance(d, torch.Tensor): - return d.detach().cpu() - if isinstance(d, dict): - return {k: _kwargs_to_cpu(v) for k, v in d.items()} - if isinstance(d, list): - return [_kwargs_to_cpu(v) for v in d] - if isinstance(d, tuple): - return tuple(_kwargs_to_cpu(v) for v in d) - return d - - -class RolloutDenoisingMixin: - def _maybe_prepare_rollout(self, batch: Req): - """Create the rollout controller of the visual stream of this request. - - Must run before SP sharding so the controller records the full latent shape. - """ - scheduler = batch.scheduler - if not isinstance(scheduler, FlowMatchEulerDiscreteScheduler): - raise ValueError(f"Scheduler {type(scheduler)} does not support rollout") - - pipeline_config = self.server_args.pipeline_config - batch.rollout_controllers = { - "visual": FlowRolloutController( - sde_type=batch.rollout_sde_type, - noise_level=float(batch.rollout_noise_level), - log_prob_no_const=batch.rollout_log_prob_no_const, - debug_mode=batch.rollout_debug_mode, - sde_step_indices=batch.rollout_sde_step_indices, - generator=batch.generator, - return_dit_trajectory=batch.rollout_return_dit_trajectory, - return_step_indices=batch.rollout_return_step_indices, - shard_for_sp=lambda latents: pipeline_config.shard_latents_for_sp( - batch=batch, latents=latents - )[0], - gather_for_sp=lambda stacked_latents: gather_stacked_latents_for_sp( - pipeline_config=pipeline_config, - batch=batch, - stacked_latents=stacked_latents, - ), - sigmas=scheduler.sigmas, - latents_shape=( - tuple(batch.latents.shape) if batch.latents is not None else None - ), - ) - } - - def _maybe_collect_rollout_log_probs(self, batch: Req): - rollout = batch.rollout_controllers["visual"] - if batch.rollout_trajectory_data is None: - batch.rollout_trajectory_data = RolloutTrajectoryData() - batch.rollout_trajectory_data.rollout_log_probs = ( - rollout.collect_rollout_log_probs() - ) - if batch.rollout_debug_mode: - batch.rollout_trajectory_data.rollout_debug_tensors = ( - rollout.collect_rollout_debug_tensors() - ) - - def _postprocess_rollout_outputs( - self, - batch: Req, - latents: torch.Tensor, - num_inference_steps: int, - final_timestep: torch.Tensor, - server_args: ServerArgs, - ) -> None: - """Finalize rollout-only outputs. - - Must be called before ``_post_denoising_loop`` so that ``latents`` (the - last ``scheduler.step`` output) is still SP-sharded and can be gathered - uniformly with the per-step trajectory latents. - """ - self._maybe_collect_rollout_log_probs(batch) - # Append final denoised latent as the (T+1)-th entry (step_index=T), - # routed through the same filter so rollout_return_step_indices can - # include/exclude it. - batch.rollout_controllers["visual"]._maybe_append_dit_trajectory_step( - latents=latents, - timestep_value=final_timestep, - step_index=num_inference_steps, - ) - self._maybe_finalize_denoising_env_collection( - batch=batch, - pipeline_config=server_args.pipeline_config, - ) - batch.rollout_controllers = {} - - def _maybe_init_denoising_env_collection( - self, - batch, - pipeline_config, - image_kwargs: dict[str, Any], - pos_cond_kwargs: dict[str, Any], - neg_cond_kwargs: dict[str, Any], - guidance: torch.Tensor | None, - ) -> None: - collect_env = batch.rollout_return_denoising_env - collect_traj = batch.rollout_return_dit_trajectory - if not (collect_env or collect_traj): - batch._rollout_denoising_env_state = None - return - - if collect_env: - env = RolloutDenoisingEnv( - image_kwargs=_kwargs_to_cpu(image_kwargs), - pos_cond_kwargs=_kwargs_to_cpu(pos_cond_kwargs), - neg_cond_kwargs=( - _kwargs_to_cpu(neg_cond_kwargs) if neg_cond_kwargs else None - ), - guidance=guidance.detach().cpu() if guidance is not None else None, - ) - pos_src = pos_cond_kwargs - neg_src = neg_cond_kwargs - else: - env = None - pos_src = None - neg_src = None - - batch._rollout_denoising_env_state = { - "env": env, - "pos_cond_kwargs_src": pos_src, - "neg_cond_kwargs_src": neg_src, - } - - def _maybe_finalize_denoising_env_collection(self, batch, pipeline_config) -> None: - state = getattr(batch, "_rollout_denoising_env_state", None) - if state is None: - return - - env: RolloutDenoisingEnv | None = state["env"] - rollout = batch.rollout_controllers["visual"] - step_latents: list[torch.Tensor] = rollout.step_latents - step_timesteps: list[torch.Tensor] = rollout.step_timesteps - - if batch.rollout_trajectory_data is None: - batch.rollout_trajectory_data = RolloutTrajectoryData() - - if step_latents and batch.rollout_return_dit_trajectory: - step_latents_tensor = rollout.gather_for_sp( - torch.stack(step_latents, dim=1) - ) - batch.rollout_trajectory_data.dit_trajectory = RolloutDitTrajectory( - latents=step_latents_tensor.cpu(), - timesteps=torch.stack(step_timesteps, dim=0).cpu(), - sigmas=rollout.sigmas.detach().cpu().clone(), - latent_step_indices=torch.tensor( - rollout.step_latent_indices, dtype=torch.long - ), - ) - - if env is not None and batch.rollout_return_denoising_env: - gather_fn = getattr( - pipeline_config, "gather_denoising_env_static_for_sp", None - ) - - pos_src = state.get("pos_cond_kwargs_src") - if pos_src is not None and env.pos_cond_kwargs is not None: - gathered_pos = gather_fn(batch, pos_src) if gather_fn else pos_src - env.pos_cond_kwargs = _kwargs_to_cpu(gathered_pos) - - neg_src = state.get("neg_cond_kwargs_src") - if neg_src is not None and env.neg_cond_kwargs is not None: - gathered_neg = gather_fn(batch, neg_src) if gather_fn else neg_src - env.neg_cond_kwargs = _kwargs_to_cpu(gathered_neg) - - batch.rollout_trajectory_data.denoising_env = env - - batch._rollout_denoising_env_state = None diff --git a/python/sglang/multimodal_gen/runtime/post_training/rollout_hooks.py b/python/sglang/multimodal_gen/runtime/post_training/rollout_hooks.py new file mode 100644 index 000000000000..91b7f759db60 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/post_training/rollout_hooks.py @@ -0,0 +1,220 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Rollout hooks for denoise loops. + +A loop calls ``prepare_rollout_if_enabled`` before SP sharding, +``init_rollout_env_if_enabled`` at the loop start and +``finalize_rollout_if_enabled`` at the loop end, and passes +``rollout_transition(batch)`` to ``scheduler.step``. Every hook is a no-op for +serving requests. +""" + +from __future__ import annotations + +from typing import Any, Callable + +import torch + +from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler, +) +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.post_training.flow_rollout_controller import ( + FlowRolloutController, +) +from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import ( + RolloutDenoisingEnv, + RolloutDitTrajectory, + RolloutTrajectoryData, +) +from sglang.multimodal_gen.runtime.post_training.sp_utils import ( + gather_stacked_latents_for_sp, +) + + +def _kwargs_to_cpu(d: Any) -> Any: + if isinstance(d, torch.Tensor): + return d.detach().cpu() + if isinstance(d, dict): + return {k: _kwargs_to_cpu(v) for k, v in d.items()} + if isinstance(d, list): + return [_kwargs_to_cpu(v) for v in d] + if isinstance(d, tuple): + return tuple(_kwargs_to_cpu(v) for v in d) + return d + + +def _build_scheduler_rollout_controller( + *, batch: Req, pipeline_config: Any +) -> FlowRolloutController: + """The visual-stream controller of a loop that steps ``batch.scheduler``.""" + scheduler = batch.scheduler + if not isinstance(scheduler, FlowMatchEulerDiscreteScheduler): + raise ValueError(f"Scheduler {type(scheduler)} does not support rollout") + + controller = FlowRolloutController( + sde_type=batch.rollout_sde_type, + noise_level=float(batch.rollout_noise_level), + log_prob_no_const=batch.rollout_log_prob_no_const, + debug_mode=batch.rollout_debug_mode, + sde_step_indices=batch.rollout_sde_step_indices, + generator=batch.generator, + return_dit_trajectory=batch.rollout_return_dit_trajectory, + return_step_indices=batch.rollout_return_step_indices, + shard_for_sp=lambda latents: pipeline_config.shard_latents_for_sp( + batch=batch, latents=latents + )[0], + gather_for_sp=lambda stacked_latents: gather_stacked_latents_for_sp( + pipeline_config=pipeline_config, + batch=batch, + stacked_latents=stacked_latents, + ), + sigmas=scheduler.sigmas, + latents_shape=( + tuple(batch.latents.shape) if batch.latents is not None else None + ), + ) + return controller + + +def prepare_rollout_if_enabled( + *, + batch: Req, + pipeline_config: Any, + build_controller: Callable[..., FlowRolloutController] = ( + _build_scheduler_rollout_controller + ), +) -> None: + """Create the rollout controller of the visual stream of this request. + + Must run before SP sharding so the controller records the full latent shape. + ``build_controller(batch=, pipeline_config=)`` returns it; a loop without a + scheduler passes its own. + """ + if not batch.rollout: + return + controller = build_controller(batch=batch, pipeline_config=pipeline_config) + batch.rollout_controllers = {"visual": controller} + + +def init_rollout_env_if_enabled( + *, + batch: Req, + image_kwargs: dict[str, Any], + pos_cond_kwargs: dict[str, Any], + neg_cond_kwargs: dict[str, Any], + guidance: torch.Tensor | None, +) -> None: + """Snapshot the DiT conditioning that the trainer replays.""" + if not batch.rollout: + return + collect_env = batch.rollout_return_denoising_env + collect_traj = batch.rollout_return_dit_trajectory + if not (collect_env or collect_traj): + batch._rollout_denoising_env_state = None + return + + if collect_env: + env = RolloutDenoisingEnv( + image_kwargs=_kwargs_to_cpu(image_kwargs), + pos_cond_kwargs=_kwargs_to_cpu(pos_cond_kwargs), + neg_cond_kwargs=( + _kwargs_to_cpu(neg_cond_kwargs) if neg_cond_kwargs else None + ), + guidance=guidance.detach().cpu() if guidance is not None else None, + ) + pos_src = pos_cond_kwargs + neg_src = neg_cond_kwargs + else: + env = None + pos_src = None + neg_src = None + + batch._rollout_denoising_env_state = { + "env": env, + "pos_cond_kwargs_src": pos_src, + "neg_cond_kwargs_src": neg_src, + } + + +def finalize_rollout_if_enabled( + *, + batch: Req, + latents: torch.Tensor, + num_inference_steps: int, + final_timestep: torch.Tensor, + pipeline_config: Any, +) -> None: + """Finalize rollout-only outputs. + + Must be called before ``_post_denoising_loop`` so that ``latents`` (the + last ``scheduler.step`` output) is still SP-sharded and can be gathered + uniformly with the per-step trajectory latents. + """ + if not batch.rollout: + return + _collect_rollout_log_probs(batch) + # Append final denoised latent as the (T+1)-th entry (step_index=T), + # routed through the same filter so rollout_return_step_indices can + # include/exclude it. + batch.rollout_controllers["visual"]._maybe_append_dit_trajectory_step( + latents=latents, + timestep_value=final_timestep, + step_index=num_inference_steps, + ) + _finalize_rollout_env(batch=batch, pipeline_config=pipeline_config) + batch.rollout_controllers = {} + + +def _collect_rollout_log_probs(batch: Req) -> None: + rollout = batch.rollout_controllers["visual"] + if batch.rollout_trajectory_data is None: + batch.rollout_trajectory_data = RolloutTrajectoryData() + batch.rollout_trajectory_data.rollout_log_probs = ( + rollout.collect_rollout_log_probs() + ) + if batch.rollout_debug_mode: + batch.rollout_trajectory_data.rollout_debug_tensors = ( + rollout.collect_rollout_debug_tensors() + ) + + +def _finalize_rollout_env(*, batch: Req, pipeline_config: Any) -> None: + state = getattr(batch, "_rollout_denoising_env_state", None) + if state is None: + return + + env: RolloutDenoisingEnv | None = state["env"] + rollout = batch.rollout_controllers["visual"] + step_latents: list[torch.Tensor] = rollout.step_latents + step_timesteps: list[torch.Tensor] = rollout.step_timesteps + + if batch.rollout_trajectory_data is None: + batch.rollout_trajectory_data = RolloutTrajectoryData() + + if step_latents and batch.rollout_return_dit_trajectory: + step_latents_tensor = rollout.gather_for_sp(torch.stack(step_latents, dim=1)) + batch.rollout_trajectory_data.dit_trajectory = RolloutDitTrajectory( + latents=step_latents_tensor.cpu(), + timesteps=torch.stack(step_timesteps, dim=0).cpu(), + sigmas=rollout.sigmas.detach().cpu().clone(), + latent_step_indices=torch.tensor( + rollout.step_latent_indices, dtype=torch.long + ), + ) + + if env is not None and batch.rollout_return_denoising_env: + gather_fn = getattr(pipeline_config, "gather_denoising_env_static_for_sp", None) + + pos_src = state.get("pos_cond_kwargs_src") + if pos_src is not None and env.pos_cond_kwargs is not None: + gathered_pos = gather_fn(batch, pos_src) if gather_fn else pos_src + env.pos_cond_kwargs = _kwargs_to_cpu(gathered_pos) + + neg_src = state.get("neg_cond_kwargs_src") + if neg_src is not None and env.neg_cond_kwargs is not None: + gathered_neg = gather_fn(batch, neg_src) if gather_fn else neg_src + env.neg_cond_kwargs = _kwargs_to_cpu(gathered_neg) + + batch.rollout_trajectory_data.denoising_env = env + + batch._rollout_denoising_env_state = None From 8d835f684c6ac72a043c17b0fae8c30187813082 Mon Sep 17 00:00:00 2001 From: rockdu Date: Wed, 14 Oct 2026 15:48:20 -0700 Subject: [PATCH 56/56] feat(diffusion): let a model pass its own ODE update to the rollout controller Refactor-Kind: non_mechanical_provable Signed-off-by: rockdu --- .../post_training/flow_rollout_controller.py | 12 ++++- .../test/unit/test_scheduler_rollout_unit.py | 49 ++++++++++++++++++- 2 files changed, 59 insertions(+), 2 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py index 0173822563a8..0ec23e6a5a7f 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py +++ b/python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py @@ -47,6 +47,7 @@ def __init__( gather_for_sp: Callable[[torch.Tensor], torch.Tensor], sigmas: torch.Tensor, latents_shape: tuple | None, + ode_transition: Callable[..., torch.Tensor] | None = None, ) -> None: self.sde_type = sde_type self.noise_level = noise_level @@ -63,6 +64,10 @@ def __init__( self.sigmas = sigmas self.sigma_max = sigmas[min(1, len(sigmas) - 1)].item() self.latents_shape = latents_shape + # The model's own deterministic update, called as + # ode_transition(model_output=, sample=, step_index=) on ODE steps; + # None means the flow-matching Euler update sample + dt * model_output. + self.ode_transition = ode_transition self.noise_buffer: torch.Tensor | None = None self.local_log_prob_sum: list[torch.Tensor] = [] @@ -314,7 +319,12 @@ def flow_sde_sampling( log_prob_no_const_val = -((full_variance_noise * noise_std_dev) ** 2) elif effective_sde_type == "ode": - prev_sample = sample + dt * model_output + if self.ode_transition is None: + prev_sample = sample + dt * model_output + else: + prev_sample = self.ode_transition( + model_output=model_output, sample=sample, step_index=step_index + ) prev_sample_mean = prev_sample variance_noise = torch.zeros_like(model_output) noise_std_dev = torch.zeros( diff --git a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py index 58400b97d8c3..f2d18e074bf7 100644 --- a/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py +++ b/python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py @@ -8,6 +8,7 @@ |-- record x_t into the trajectory (filtered by return_step_indices) +-- flow_sde_sampling(...) |-- ode : x_next = x_t + dt * v (no noise, log_prob 0) + | or the model's own ode_transition |-- sde / cps: x_next = mean + std * noise (fp32, log_prob summed) +-- step outside sde_step_indices -> ode form v @@ -16,6 +17,7 @@ TestSchedulerRolloutOdeUnit ODE rollout == serving step, bit for bit (bare update and through step()); ODE draws no noise; the trajectory records x_t, not x_next; + a model's own ode_transition replaces the Euler update on ODE steps only; debug tensors keep [B, T, ...] shapes. TestSchedulerFlowGRPOStepAlignmentUnit SDE / CPS match the verbatim FlowGRPO reference; bf16 model_output still @@ -39,7 +41,9 @@ _SIGMAS = torch.tensor([1.0, 0.8, 0.6, 0.4, 0.2, 0.0], dtype=torch.float32) -def _make_rollout(batch, pipeline_config=None) -> FlowRolloutController: +def _make_rollout( + batch, pipeline_config=None, ode_transition=None +) -> FlowRolloutController: """The controller the denoising stage creates for ``batch`` (single GPU).""" return FlowRolloutController( sde_type=batch.rollout_sde_type, @@ -56,6 +60,7 @@ def _make_rollout(batch, pipeline_config=None) -> FlowRolloutController: gather_for_sp=lambda stacked_latents: stacked_latents, sigmas=_SIGMAS, latents_shape=tuple(batch.latents.shape), + ode_transition=ode_transition, ) @@ -208,6 +213,48 @@ def test_transition_records_pre_step_latents(self): self.assertTrue(torch.equal(rollout.step_latents[0], sample)) self.assertFalse(torch.equal(rollout.step_latents[0], prev_sample)) + def test_ode_transition_replaces_euler_on_ode_steps_only(self): + """A model's own ode_transition runs on the steps outside the SDE window + and nowhere else; SDE steps keep the shared stochastic update.""" + batch = self._build_batch(debug_mode=False) + batch.rollout_sde_type = "sde" + batch.rollout_sde_step_indices = [1] + pipeline_config = types.SimpleNamespace( + shard_latents_for_sp=lambda batch, latents: (latents, False) + ) + ode_calls = [] + + def ode_transition(*, model_output, sample, step_index): + ode_calls.append(step_index) + return sample * 2 + + rollout = _make_rollout(batch, pipeline_config, ode_transition=ode_transition) + sample = torch.randn(2, 4, 8, 8) + model_output = torch.randn_like(sample) + + # step 0 is outside the window: the model's update, zero log-prob + prev_0 = rollout.flow_sde_sampling( + model_output=model_output, + sample=sample, + current_sigma=torch.tensor(0.6), + next_sigma=torch.tensor(0.4), + generator=None, + step_index=0, + ) + self.assertTrue(torch.equal(prev_0, sample * 2)) + # step 1 is inside the window: stochastic update, ode_transition not called + rollout.flow_sde_sampling( + model_output=model_output, + sample=sample, + current_sigma=torch.tensor(0.6), + next_sigma=torch.tensor(0.4), + generator=[torch.Generator().manual_seed(i) for i in range(2)], + step_index=1, + ) + self.assertEqual(ode_calls, [0]) + log_prob_sum, _ = rollout.consume_local_rollout_log_probs() + self.assertTrue(torch.equal(log_prob_sum[:, 0], torch.zeros(2))) + def test_ode_debug_tensors_have_shape_safe_noise_std(self): batch = self._build_batch(debug_mode=True) rollout = _make_rollout(batch)