From f81b6df3a3fbff7203e604e935091e26e29662b1 Mon Sep 17 00:00:00 2001 From: Kurt Shuster Date: Sun, 12 Apr 2026 19:25:02 -0400 Subject: [PATCH 01/52] [lora] Fix partial MoE rank loading, VL lm_head, strict loading, deepseek on-demand (#21864) Co-authored-by: Yusheng Su --- python/sglang/srt/lora/lora_manager.py | 35 +++-- python/sglang/srt/lora/mem_pool.py | 206 +++++++++++++++++++------ python/sglang/srt/server_args.py | 12 +- 3 files changed, 187 insertions(+), 66 deletions(-) diff --git a/python/sglang/srt/lora/lora_manager.py b/python/sglang/srt/lora/lora_manager.py index 06f3b09dde48..9f3ef601ed5c 100644 --- a/python/sglang/srt/lora/lora_manager.py +++ b/python/sglang/srt/lora/lora_manager.py @@ -35,6 +35,7 @@ from sglang.srt.lora.lora_registry import LoRARef from sglang.srt.lora.mem_pool import LoRAMemoryPool from sglang.srt.lora.utils import ( + EMBEDDING_NAMES, LoRAType, auto_detect_lora_target_modules, get_normalized_target_modules, @@ -85,6 +86,9 @@ def __init__( self._experts_shared_outer_override: Optional[bool] = ( server_args.experts_shared_outer_loras ) + self.lora_strict_loading: bool = getattr( + server_args, "lora_strict_loading", False + ) # LoRA backend for running sgemm kernels logger.info(f"Using {lora_backend} as backend of LoRA kernels.") @@ -507,9 +511,19 @@ def init_lora_shapes( ): """Infer LoRA target modules and max_lora_rank from loaded adapters if not provided.""" - self.target_modules = ( - get_normalized_target_modules(target_modules) if target_modules else set() - ) + if target_modules and target_modules == {"all"}: + self.target_modules = auto_detect_lora_target_modules(self.base_model) + self.target_modules.update(EMBEDDING_NAMES) + logger.info( + "CLI --lora-target-modules='all' resolved to %s " + "by inspecting the base model.", + sorted(self.target_modules), + ) + target_modules = self.target_modules + elif target_modules: + self.target_modules = get_normalized_target_modules(target_modules) + else: + self.target_modules = set() for lora_id, config in self.configs.items(): # Handle PEFT shorthand strings like "all-linear" or "all". @@ -682,6 +696,7 @@ def init_memory_pool(self): eviction_policy=self.eviction_policy, lora_added_tokens_size=self.lora_added_tokens_size, experts_shared_outer_loras=self.experts_shared_outer_loras, + strict_loading=self.lora_strict_loading, ) # Initializing memory pool with base model @@ -737,17 +752,9 @@ def init_lora_modules(self): self.base_model.lm_head = untied_lm_head for module_name, module in self.base_model.named_modules(): - # TODO (lifuhuang): in the future, we should consider generalizing the - # should_apply_lora function to support mapping by full module name instead - # of just the last part (e.g., "qkv_proj") to support scenarios with multiple - # attention stacks (e.g., multimodal models). - # See: https://github.com/sgl-project/sglang/issues/6608 - if getattr( - self.base_model, "should_apply_lora", None - ) and not self.base_model.should_apply_lora(module_name): - continue - - # Check if module should be wrapped with LoRA + # Handle embed_tokens and lm_head before the should_apply_lora gate, + # since VL models' should_apply_lora patterns only match language + # model layers and would incorrectly skip these. # Handle embed_tokens if "embed_tokens" in module_name and "embed_tokens" in self.target_modules: if isinstance(module, VocabParallelEmbedding) and not isinstance( diff --git a/python/sglang/srt/lora/mem_pool.py b/python/sglang/srt/lora/mem_pool.py index 6c923afc5f7d..34444edfc415 100644 --- a/python/sglang/srt/lora/mem_pool.py +++ b/python/sglang/srt/lora/mem_pool.py @@ -4,7 +4,7 @@ import torch -from sglang.srt.distributed import divide +from sglang.srt.distributed import divide, get_pp_group from sglang.srt.lora.eviction_policy import get_eviction_policy from sglang.srt.lora.layers import BaseLayerWithLoRA from sglang.srt.lora.lora import LoRAAdapter @@ -62,6 +62,7 @@ def __init__( eviction_policy: str, lora_added_tokens_size: int, experts_shared_outer_loras: bool = False, + strict_loading: bool = False, ): self.base_hf_config: AutoConfig = base_hf_config self.num_layer: int = base_hf_config.num_hidden_layers @@ -73,6 +74,7 @@ def __init__( self.max_lora_rank: int = max_lora_rank self.target_modules: Set[str] = target_modules self.experts_shared_outer_loras: bool = experts_shared_outer_loras + self.strict_loading: bool = strict_loading # Initialize eviction policy self.eviction_policy = get_eviction_policy(eviction_policy) @@ -533,6 +535,43 @@ def load_lora_weight_tensor( assert lora_adapter is not None lora_rank = lora_adapter.config.r + + # Pre-validate weight names against target modules across all layers + # and embedding weights. This catches mismatches before any GPU + # buffers are mutated. + skipped_weight_names: set = set() + matched_modules: set = set() + all_weight_names: list = [] + for layer in lora_adapter.layers: + all_weight_names.extend(layer.weights.keys()) + if lora_adapter.embedding_layers: + all_weight_names.extend(lora_adapter.embedding_layers.keys()) + for name in all_weight_names: + try: + target_module = get_target_module_name(name, self.target_modules) + matched_modules.add(target_module) + except ValueError: + skipped_weight_names.add(name) + if matched_modules: + logger.info( + "LoRA adapter '%s': loaded weights for target modules %s.", + uid, + sorted(matched_modules), + ) + if skipped_weight_names: + msg = ( + f"LoRA adapter '{uid}': {len(skipped_weight_names)} weight(s) " + f"skipped because they did not match any target module in " + f"{sorted(self.target_modules)}. Skipped weights: " + f"{sorted(skipped_weight_names)}. This likely indicates a " + f"mismatch between the adapter's target modules and the base " + f"model architecture." + ) + if self.strict_loading: + raise ValueError(msg) + else: + logger.warning(msg) + for layer_id in range(self.num_layer): layer_weights = lora_adapter.layers[layer_id].weights # - Standard: module_name -> torch.Tensor @@ -576,59 +615,58 @@ def load_lora_weight_tensor( else: temp_B_buffer[target_module] = weights - if self.tp_size > 1: - cur_layer_modules = lora_modules[layer_id] - for module_name, module in cur_layer_modules.items(): - # TODO (Jonahcb): check if the code can be refactored to avoid the special handling for FusedMoEWithLoRA - # Handle FusedMoEWithLoRA specially - it contains multiple target modules - from sglang.srt.lora.layers import FusedMoEWithLoRA - - if isinstance(module, FusedMoEWithLoRA): - moe_target_modules = ["gate_up_proj_moe", "down_proj_moe"] - for target_module in moe_target_modules: - if temp_A_buffer.get(target_module) is not None: - temp_A_buffer[target_module] = ( - module.slice_moe_lora_a_weights( - temp_A_buffer[target_module], - self.tp_rank, - target_module, - ) + cur_layer_modules = lora_modules[layer_id] + for module_name, module in cur_layer_modules.items(): + # TODO (Jonahcb): check if the code can be refactored to avoid the special handling for FusedMoEWithLoRA + # Handle FusedMoEWithLoRA specially - it contains multiple target modules + from sglang.srt.lora.layers import FusedMoEWithLoRA + + if isinstance(module, FusedMoEWithLoRA): + moe_target_modules = ["gate_up_proj_moe", "down_proj_moe"] + for target_module in moe_target_modules: + if temp_A_buffer.get(target_module) is not None: + temp_A_buffer[target_module] = ( + module.slice_moe_lora_a_weights( + temp_A_buffer[target_module], + self.tp_rank, + target_module, ) - if temp_B_buffer.get(target_module) is not None: - temp_B_buffer[target_module] = ( - module.slice_moe_lora_b_weights( - temp_B_buffer[target_module], - self.tp_rank, - target_module, - ) + ) + if temp_B_buffer.get(target_module) is not None: + temp_B_buffer[target_module] = ( + module.slice_moe_lora_b_weights( + temp_B_buffer[target_module], + self.tp_rank, + target_module, ) + ) - continue + continue - # Handle regular modules - target_module = get_target_module_name( - module_name, self.target_modules - ) + # Handle regular modules + target_module = get_target_module_name(module_name, self.target_modules) - if temp_A_buffer[target_module] is None: - # Skip weight slicing if the weight is not present in the adapter - continue + if temp_A_buffer[target_module] is None: + # Skip weight slicing if the weight is not present in the adapter + continue - # Handle standard modules - temp_A_buffer[target_module] = module.slice_lora_a_weights( - temp_A_buffer[target_module], self.tp_rank - ) - temp_B_buffer[target_module] = module.slice_lora_b_weights( - temp_B_buffer[target_module], self.tp_rank - ) + # Handle standard modules + temp_A_buffer[target_module] = module.slice_lora_a_weights( + temp_A_buffer[target_module], self.tp_rank + ) + temp_B_buffer[target_module] = module.slice_lora_b_weights( + temp_B_buffer[target_module], self.tp_rank + ) for name, weights in temp_A_buffer.items(): c = get_stacked_multiply(name) + max_r = self.max_lora_rank target_buffer = self.A_buffer[name][layer_id] if name in ["gate_up_proj_moe", "down_proj_moe"]: if self.experts_shared_outer_loras and name == "gate_up_proj_moe": if weights is None: + representative_weight = None buffer_view = target_buffer[ buffer_id, 0, : lora_rank * c, : ] @@ -640,6 +678,7 @@ def load_lora_weight_tensor( f"gate_up_proj_moe lora_A has expert_dim=" f"{weights.shape[0]} (expected 1)." ) + representative_weight = weights[0] buffer_view = target_buffer[ buffer_id, 0, : lora_rank * c, : ] @@ -652,6 +691,7 @@ def load_lora_weight_tensor( f"{len(weights)} entries (expected 1)." ) rep = next(iter(weights.values())) + representative_weight = rep buffer_view = target_buffer[ buffer_id, 0, : lora_rank * c, : ] @@ -662,18 +702,65 @@ def load_lora_weight_tensor( f"type={type(weights)}, " f"shape={weights.shape if isinstance(weights, torch.Tensor) else 'N/A'}" ) + # 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_() + if representative_weight is not None: + for ci in range(c): + buffer_view = target_buffer[ + buffer_id, 0, ci * max_r : ci * max_r + lora_rank, : + ] + load_lora_weight_tensor( + buffer_view, + representative_weight[ + ci * lora_rank : (ci + 1) * lora_rank, : + ], + ) elif isinstance(weights, torch.Tensor) and weights.dim() == 3: for eid in range(weights.shape[0]): - buffer_view = target_buffer[ - buffer_id, eid, : lora_rank * c, : - ] - load_lora_weight_tensor(buffer_view, weights[eid]) + # Place each component at max_rank-spaced positions + # and zero gaps so the MoE kernel (which processes + # the full max_rank) sees correct data. + target_buffer[buffer_id, eid].zero_() + expert_weight = weights[eid] + if expert_weight is not None: + for ci in range(c): + buffer_view = target_buffer[ + buffer_id, + eid, + ci * max_r : ci * max_r + lora_rank, + :, + ] + load_lora_weight_tensor( + buffer_view, + expert_weight[ + ci * lora_rank : (ci + 1) * lora_rank, : + ], + ) elif isinstance(weights, dict): - for expert_id, expert_weight in weights.items(): - buffer_view = target_buffer[ - buffer_id, expert_id, : lora_rank * c, : - ] - load_lora_weight_tensor(buffer_view, expert_weight) + if weights is not None: + for expert_id, expert_weight in weights.items(): + # Place each component at max_rank-spaced positions + # and zero gaps so the MoE kernel (which processes + # the full max_rank) sees correct data. + target_buffer[buffer_id, expert_id].zero_() + if expert_weight is not None: + for ci in range(c): + buffer_view = target_buffer[ + buffer_id, + expert_id, + ci * max_r : ci * max_r + lora_rank, + :, + ] + load_lora_weight_tensor( + buffer_view, + expert_weight[ + ci * lora_rank : (ci + 1) * lora_rank, : + ], + ) + else: + target_buffer[buffer_id].zero_() else: buffer_view = target_buffer[buffer_id, : lora_rank * c, :] load_lora_weight_tensor(buffer_view, weights) @@ -698,6 +785,8 @@ def load_lora_weight_tensor( if w is not None: w = w * lora_adapter.scaling load_lora_weight_tensor(buffer_view, w) + # Zero beyond loaded rank — MoE kernel reads full max_rank + target_buffer[buffer_id, 0, :, lora_rank:].zero_() elif isinstance(weights, dict) and len(weights) > 0: if len(weights) != 1: raise ValueError( @@ -710,6 +799,8 @@ def load_lora_weight_tensor( if rep is not None: rep = rep * lora_adapter.scaling load_lora_weight_tensor(buffer_view, rep) + # Zero beyond loaded rank — MoE kernel reads full max_rank + target_buffer[buffer_id, 0, :, lora_rank:].zero_() else: raise ValueError( f"Unexpected weight format for shared outer down_proj_moe lora_B: " @@ -723,6 +814,8 @@ def load_lora_weight_tensor( if w is not None: w = w * lora_adapter.scaling load_lora_weight_tensor(buffer_view, w) + # Zero beyond loaded rank — MoE kernel reads full max_rank + target_buffer[buffer_id, eid, :, lora_rank:].zero_() elif isinstance(weights, dict): for expert_id, expert_weight in weights.items(): buffer_view = target_buffer[ @@ -732,6 +825,8 @@ def load_lora_weight_tensor( if w is not None: w = w * lora_adapter.scaling load_lora_weight_tensor(buffer_view, w) + # Zero beyond loaded rank — MoE kernel reads full max_rank + target_buffer[buffer_id, expert_id, :, lora_rank:].zero_() else: buffer_view = target_buffer[buffer_id, :, :lora_rank] load_lora_weight_tensor(buffer_view, weights) @@ -779,6 +874,7 @@ def load_lora_weight_tensor( elif ( target_module == "lm_head" + and lora_lm_head_module is not None and "lm_head" in name and ("lora_embedding_A" in name or "lora_A" in name) ): @@ -791,12 +887,14 @@ def load_lora_weight_tensor( load_lora_weight_tensor(buffer_view, weights) elif ( target_module == "lm_head" + and lora_lm_head_module is not None and "lm_head" in name and ("lora_embedding_B" in name or "lora_B" in name) ): + assert lora_lm_head_module is not None lora_b_weights = weights # Slice B along vocab dimension for this TP rank - if self.tp_size > 1 and lora_lm_head_module is not None: + if self.tp_size > 1: lora_b_weights = lora_lm_head_module.slice_lora_b_weights( lora_b_weights, self.tp_rank ) @@ -807,6 +905,14 @@ def load_lora_weight_tensor( :lora_rank, ] load_lora_weight_tensor(buffer_view, lora_b_weights) + elif target_module == "lm_head" and "lm_head" in name: + # Non-last PP stages do not own lm_head, so adapters can + # legitimately contain lm_head LoRA weights with no local + # module to load them into, otherwise we should have been able to load this weight. + assert ( + not get_pp_group().is_last_rank + ), f"Failed to load lm_head LoRA weight: {name}, this is only expected to happen on non-last PP stages." + continue else: # Zero out embedding/lm_head buffers for adapters without embedding LoRA # to avoid using garbage values from uninitialized memory diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index cb549bfcda1f..9e66c282801f 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -475,6 +475,7 @@ class ServerArgs: lora_backend: str = "csgmv" max_lora_chunk_size: Optional[int] = 16 experts_shared_outer_loras: Optional[bool] = None + lora_strict_loading: bool = False # Kernel backend attention_backend: Optional[str] = None @@ -4961,6 +4962,13 @@ def add_cli_args(parser: argparse.ArgumentParser): "(expert_dim=1). Use --no-experts-shared-outer-loras to force disable. " "By default this is auto-detected from adapter weights.", ) + parser.add_argument( + "--lora-strict-loading", + default=ServerArgs.lora_strict_loading, + action=argparse.BooleanOptionalAction, + help="Enable strict loading for LoRA adapters. " + "When set, mismatched or missing keys in the adapter weights will raise an error.", + ) # Kernel backend parser.add_argument( @@ -6669,14 +6677,14 @@ def check_lora_server_args(self): "Expected a list or a dictionary." ) - # Expand target modules + # Normalize target modules to a set; keep {"all"} as a sentinel + # that gets resolved model-awarely in lora_manager.init_lora_shapes(). if self.lora_target_modules: self.lora_target_modules = set(self.lora_target_modules) if "all" in self.lora_target_modules: assert ( len(self.lora_target_modules) == 1 ), "If 'all' is specified in --lora-target-modules, it should be the only module specified." - self.lora_target_modules = set(SUPPORTED_LORA_TARGET_MODULES) # Ensure sufficient information is provided for LoRA initialization. assert self.lora_paths or ( From 03a1a7b81cb966a70af2a1ced20295dea6c74042 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com> Date: Mon, 13 Apr 2026 07:57:41 +0800 Subject: [PATCH 02/52] [Diffusion] Add FLUX.1-dev ModelOpt NVFP4 support (#22574) --- docs/diffusion/quantization.md | 22 +- python/sglang/jit_kernel/nvfp4.py | 13 + .../sglang-diffusion-modelopt-quant/SKILL.md | 55 +-- .../sglang-diffusion-performance/SKILL.md | 2 +- .../layers/quantization/modelopt_quant.py | 23 +- .../runtime/loader/fsdp_load.py | 11 +- .../runtime/models/dits/flux.py | 9 + .../pipelines_core/stages/denoising.py | 18 + .../test/unit/test_transformer_quant.py | 138 +++++- ...t.py => build_modelopt_fp8_transformer.py} | 252 +++++++++-- .../tools/build_modelopt_nvfp4_transformer.py | 402 ++++++++++++++++++ 11 files changed, 871 insertions(+), 74 deletions(-) rename python/sglang/multimodal_gen/tools/{convert_modelopt_fp8_checkpoint.py => build_modelopt_fp8_transformer.py} (66%) create mode 100644 python/sglang/multimodal_gen/tools/build_modelopt_nvfp4_transformer.py diff --git a/docs/diffusion/quantization.md b/docs/diffusion/quantization.md index 8714e5ed4e56..96c872d4f56d 100644 --- a/docs/diffusion/quantization.md +++ b/docs/diffusion/quantization.md @@ -48,6 +48,26 @@ backend. | `nunchaku-svdq` | Pre-quantized Nunchaku transformer weights, usually named `svdq-{int4\|fp4}_r{rank}-...` | `--transformer-weights-path` | Model-specific support such as Qwen-Image, FLUX, and Z-Image | `nunchaku` | SGLang can infer precision and rank from the filename and supports both `int4` and `nvfp4` | | `msmodelslim` | Pre-quantized msmodelslim transformer weights | `--model-path` | Wan2.2 family | None | Currently only compatible with the Ascend NPU family and supports both `w8a8` and `w4a4` | +## Validated ModelOpt Checkpoints + +This section is the canonical support matrix for diffusion ModelOpt checkpoints +that have been brought up and verified in SGLang. + +### FP8 + +| Base Model | Validated Scope | HF DiT Weights | Notes | +| --- | --- | --- | --- | +| `black-forest-labs/FLUX.1-dev` | single-transformer override, deterministic latent/image comparison, H100 benchmark, torch-profiler trace | `BBuf/flux1-dev-modelopt-fp8-sglang-transformer` | SGLang converter keeps a validated BF16 fallback set for modulation and FF projection layers; use `--model-id FLUX.1-dev` for local mirrors | +| `black-forest-labs/FLUX.2-dev` | single-transformer override load and generation path | `BBuf/flux2-dev-modelopt-fp8-sglang-transformer` | published SGLang-ready transformer override | +| `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | primary `transformer` quantized, `transformer_2` kept BF16 | `BBuf/wan22-t2v-a14b-modelopt-fp8-sglang-transformer` | do not describe this as dual-transformer full-model FP8 unless that path is validated separately | + +### NVFP4 + +| Base Model | Validated Scope | HF DiT Weights | Notes | +| --- | --- | --- | --- | +| `black-forest-labs/FLUX.1-dev` | mixed BF16+NVFP4 transformer override, correctness validation, 4x RTX 5090 benchmark, torch-profiler trace | `unpublished` | use `build_modelopt_nvfp4_transformer.py`; validated builder keeps selected FLUX.1 modules in BF16 and sets `swap_weight_nibbles=false` | +| `black-forest-labs/FLUX.2-dev` | packed-QKV load path | `black-forest-labs/FLUX.2-dev-NVFP4` | validated packed export detection and runtime layout handling | + ## ModelOpt FP8 ### Usage Examples @@ -82,7 +102,7 @@ sglang generate \ ModelOpt FP8 checkpoints because the runtime expects the transformed FP8 weights to remain GPU-resident in their column-major layout. - To build the converted checkpoint yourself from a ModelOpt diffusers export, - use `python -m sglang.multimodal_gen.tools.convert_modelopt_fp8_checkpoint`. + use `python -m sglang.multimodal_gen.tools.build_modelopt_fp8_transformer`. ## NVFP4 diff --git a/python/sglang/jit_kernel/nvfp4.py b/python/sglang/jit_kernel/nvfp4.py index f54e6414b9e7..ff3c20072668 100644 --- a/python/sglang/jit_kernel/nvfp4.py +++ b/python/sglang/jit_kernel/nvfp4.py @@ -47,6 +47,19 @@ def _nvfp4_arch_env(): return override_jit_cuda_arch(major, minor, suffix="a") +@torch.compiler.disable +def prewarm_nvfp4_jit_modules( + *, include_expert_quant: bool = False, include_blockwise_moe: bool = False +) -> None: + """Materialize NVFP4 JIT modules before torch.compile traces the model.""" + _jit_nvfp4_quant_module() + _jit_nvfp4_scaled_mm_module() + if include_expert_quant: + _jit_nvfp4_expert_quant_module() + if include_blockwise_moe: + _jit_nvfp4_blockwise_moe_module() + + @cache_once def _jit_nvfp4_quant_module() -> Module: with _nvfp4_arch_env(): diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-modelopt-quant/SKILL.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-modelopt-quant/SKILL.md index 3379baaae755..46dd8137cdd0 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-modelopt-quant/SKILL.md +++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-modelopt-quant/SKILL.md @@ -23,8 +23,8 @@ This skill owns the ModelOpt-to-SGLang bridge. It is not a generic kernel-tuning - Benchmark only when BF16 and quantized commands are identical except for the checkpoint override being tested. - For diffusion FP8, pin `dit_cpu_offload=false` and `dit_layerwise_offload=false`. - For multi-transformer pipelines, use per-component overrides when different components need different checkpoints. -- When a branch is missing the validated helper tools, refresh `python/sglang/multimodal_gen/tools/convert_modelopt_fp8_checkpoint.py` and `python/sglang/multimodal_gen/tools/compare_diffusion_trajectory_similarity.py` instead of inventing one-off scripts elsewhere. -- After validating a new ModelOpt quant path, update both the FP8 and NVFP4 support tables in this skill before closing the task. +- When a branch is missing the validated helper tools, refresh `python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py`, `python/sglang/multimodal_gen/tools/build_modelopt_nvfp4_transformer.py`, and `python/sglang/multimodal_gen/tools/compare_diffusion_trajectory_similarity.py` instead of inventing one-off scripts elsewhere. +- After validating a new ModelOpt quant path, update the ModelOpt support matrix in `docs/diffusion/quantization.md` before closing the task. ## Read First @@ -38,7 +38,8 @@ Read these sources before changing code: - `python/sglang/multimodal_gen/runtime/utils/quantization_utils.py` - `python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py` - Helper tools in this repo: - - [`python/sglang/multimodal_gen/tools/convert_modelopt_fp8_checkpoint.py`](../../../tools/convert_modelopt_fp8_checkpoint.py) + - [`python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py`](../../../tools/build_modelopt_fp8_transformer.py) + - [`python/sglang/multimodal_gen/tools/build_modelopt_nvfp4_transformer.py`](../../../tools/build_modelopt_nvfp4_transformer.py) - [`python/sglang/multimodal_gen/tools/compare_diffusion_trajectory_similarity.py`](../../../tools/compare_diffusion_trajectory_similarity.py) If you are working on a new model family, inspect the transformer's config and tensor naming before changing the generic converter. @@ -52,32 +53,17 @@ This repo now contains: - diffusion-side NVFP4 loading from ModelOpt exports - FLUX.2 packed-QKV detection that distinguishes packed NVFP4 checkpoints from standard diffusers exports - automatic protection against incompatible FP8 offload modes -- FP8 export conversion: - [`python/sglang/multimodal_gen/tools/convert_modelopt_fp8_checkpoint.py`](../../../tools/convert_modelopt_fp8_checkpoint.py) +- FP8 transformer build: + [`python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py`](../../../tools/build_modelopt_fp8_transformer.py) - trajectory similarity validation: [`python/sglang/multimodal_gen/tools/compare_diffusion_trajectory_similarity.py`](../../../tools/compare_diffusion_trajectory_similarity.py) ## Documentation Maintenance -- Keep two separate support tables in this skill: one for FP8 and one for NVFP4. -- After finishing a new quant support path, update both tables in every mirrored copy of this skill. -- Each row must record the validated scope, the Hugging Face repo or path for the quantized DiT weights, and the key caveats. +- Keep the validated ModelOpt support matrix in `docs/diffusion/quantization.md`. +- Each row should record the validated scope, the Hugging Face repo or path for the quantized DiT weights, and the key caveats. - If the quantized DiT weights are not published yet, write `unpublished` explicitly instead of leaving the field blank. -## FP8 Supported Models - -| Base Model | Validated Scope | HF DiT Weights | Notes | -| --- | --- | --- | --- | -| `black-forest-labs/FLUX.1-dev` | single-transformer override, deterministic latent/image comparison, H100 benchmark, torch-profiler trace | `BBuf/flux1-dev-modelopt-fp8-sglang-transformer` | SGLang converter keeps a validated BF16 fallback set for modulation and FF projection layers; use `--model-id FLUX.1-dev` for local mirrors | -| `black-forest-labs/FLUX.2-dev` | single-transformer override load and generation path | `BBuf/flux2-dev-modelopt-fp8-sglang-transformer` | published SGLang-ready transformer override | -| `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | primary `transformer` quantized, `transformer_2` kept BF16 | `BBuf/wan22-t2v-a14b-modelopt-fp8-sglang-transformer` | do not describe this as dual-transformer full-model FP8 unless that path is validated separately | - -## NVFP4 Supported Models - -| Base Model | Validated Scope | HF DiT Weights | Notes | -| --- | --- | --- | --- | -| `black-forest-labs/FLUX.2-dev` | packed-QKV load path | `black-forest-labs/FLUX.2-dev-NVFP4` | validated packed export detection and runtime layout handling | - ## FP8 Vs NVFP4 FP8 and NVFP4 are not wired into SGLang in exactly the same way. @@ -120,7 +106,7 @@ python quantize.py \ --model \ --override-model-path \ --model-dtype \ - --format \ + --format \ --batch-size 1 \ --calib-size \ --n-steps \ @@ -130,6 +116,9 @@ python quantize.py \ --hf-ckpt-dir /hf ``` +For current ModelOpt diffusion examples, use `--format fp4` for NVFP4 exports. +Do not assume the checked-out ModelOpt version accepts a literal `nvfp4` format string unless you verified it locally. + For multi-transformer models: - quantize each backbone deliberately @@ -141,7 +130,7 @@ For multi-transformer models: FP8 requires an extra conversion step: ```bash -PYTHONPATH=python python3 -m sglang.multimodal_gen.tools.convert_modelopt_fp8_checkpoint \ +PYTHONPATH=python python3 -m sglang.multimodal_gen.tools.build_modelopt_fp8_transformer \ --modelopt-hf-dir /hf \ --modelopt-backbone-ckpt /ckpt/backbone.pt \ --base-transformer-dir \ @@ -166,9 +155,25 @@ For `FLUX.1-dev`, the validated fallback set currently keeps these modules in BF - `transformer_blocks.*.ff_context.net.0.proj` - `transformer_blocks.*.ff_context.net.2` - `single_transformer_blocks.*.norm.linear` +- `single_transformer_blocks.*.proj_mlp` Use `--model-type flux1` to force that profile, or rely on `--model-type auto` when the export config identifies `FluxTransformer2DModel`. +For FLUX.1-dev NVFP4 model families that need a mixed BF16+NVFP4 checkpoint, build the merged transformer explicitly: + +```bash +PYTHONPATH=python python3 -m sglang.multimodal_gen.tools.build_modelopt_nvfp4_transformer \ + --base-transformer-dir \ + --modelopt-hf-dir /hf/transformer \ + --output-dir /transformer-mixed \ + --pattern-preset flux1-nvfp4 +``` + +The validated FLUX.1-dev mixed builder also needs to preserve: + +- `quant_type: NVFP4` in `config.json` +- `swap_weight_nibbles: false` for the validated diffusers export + ### 4. Load The Quantized Checkpoint In SGLang Single-transformer example: @@ -304,5 +309,5 @@ When documenting results: | `runtime/utils/quantization_utils.py` | resolves flat ModelOpt configs and reconstructs NVFP4 config from metadata | | `runtime/loader/transformer_load_utils.py` | guards incompatible FP8 offload modes | | `runtime/models/dits/flux_2.py` | packed-QKV handling for the packed FLUX.2 NVFP4 family | -| `tools/convert_modelopt_fp8_checkpoint.py` | FP8 offline conversion into SGLang-native layout | +| `tools/build_modelopt_fp8_transformer.py` | Build an SGLang-loadable FP8 transformer from a ModelOpt export | | `tools/compare_diffusion_trajectory_similarity.py` | reduced deterministic BF16-vs-quantized validation | diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md index e81449004096..4fe50c2b056a 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md +++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md @@ -47,7 +47,7 @@ These options **trade output quality** for speed or VRAM savings. Results will d | **Approximate Attention** | `--attention-backend sage_attn` / `sage_attn_3` / `sliding_tile_attn` / `video_sparse_attn` / `sparse_video_gen_2_attn` / `vmoba_attn` / `sla_attn` / `sage_sla_attn` | Replaces exact attention with approximate or sparse variants. `sage_attn`: INT8/FP8 quantized Q·K; `sliding_tile_attn`: spatial-temporal tile skipping; others: model-specific sparse patterns. | ~1.5–2x on attention (varies by backend) | Quality degradation varies by backend and model. `sage_attn` is the most general; sparse backends (`sliding_tile_attn`, `video_sparse_attn`, etc.) are video-model-specific and may require config files (e.g. `--mask-strategy-file-path` for STA). Requires corresponding packages installed. | | **Cache-DiT** | `SGLANG_CACHE_DIT_ENABLED=true` + `--cache-dit-config ` | Caches intermediate residuals across denoising steps and skips redundant computations via a Selective Computation Mask (SCM). | ~1.5–2x on supported models | Quality depends on SCM config. Incompatible with `--dit-layerwise-offload`. Requires correct per-model config YAML. | | **Quantized Models (Nunchaku / SVDQuant)** | `--enable-svdquant --transformer-weights-path ` + optional `--quantization-precision int4\|nvfp4`, `--quantization-rank 32` | W4A4-style quantization via [Nunchaku](https://nunchaku.tech). Reduces DiT weight memory by ~4x. Precision/rank can be auto-inferred from weight filename or set explicitly. | ~1.5–2x compute speedup | Lossy quantization; quality depends on rank and precision. Requires pre-quantized weights. Ampere (SM8x) or SM12x only (no Hopper SM90). Higher rank = better quality but more memory. | -| **Pre-quantized Weights** | `--transformer-weights-path ` | Load any pre-quantized transformer weights (FP8, INT8, etc.) from a single `.safetensors` file, a directory, or a HuggingFace repo ID. | ~1.3–1.5x compute (dtype dependent) | Requires pre-converted weights (e.g. via `tools/convert_hf_to_fp8.py` for FP8). Quality slightly worse than BF16; varies by quantization format. | +| **Pre-quantized Weights** | `--transformer-weights-path ` | Load any pre-quantized transformer weights (FP8, INT8, etc.) from a single `.safetensors` file, a directory, or a HuggingFace repo ID. | ~1.3–1.5x compute (dtype dependent) | Requires a validated quantized transformer override, such as one produced by `tools/build_modelopt_fp8_transformer.py` for ModelOpt FP8. Quality slightly worse than BF16; varies by quantization format. | | **Component Precision Override** | `--dit-precision fp16`, `--vae-precision fp16\|bf16` | On-the-fly dtype conversion for individual components. E.g. convert a BF16 model to FP16 at load time, or run VAE in BF16 instead of FP32. | Reduces memory; FP16 can be faster on some GPUs | May affect numerical stability. VAE is FP32 by default for accuracy; lowering it is lossy. DiT defaults to BF16. | | **Fewer Inference Steps** | `--num-inference-steps N` (sampling param) | Reduces the number of denoising steps. Fewer steps = faster. | Linear speedup | Quality degrades with too few steps. Model-dependent optimal range. | diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py b/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py index 627edb93edeb..d7ec21a3570b 100755 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py @@ -52,6 +52,15 @@ def _get_fp4_gemm_op(): return current_platform.get_modelopt_fp4_gemm_op() +def _prepare_nvfp4_weight_bytes( + weight: torch.Tensor, *, swap_weight_nibbles: bool +) -> torch.Tensor: + """Normalize serialized NVFP4 bytes before padding for the runtime kernel.""" + if not swap_weight_nibbles: + return weight.contiguous() + return ((weight >> 4) | (weight << 4)).contiguous() + + class ModelOptQuantConfig(QuantizationConfig): def __init__( self, @@ -180,6 +189,7 @@ def __init__( exclude_modules: List[str] = None, packed_modules_mapping: Optional[Dict[str, List[str]]] = None, checkpoint_uses_packed_qkv: bool = False, + swap_weight_nibbles: bool = True, ) -> None: super().__init__(exclude_modules, packed_modules_mapping) self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized @@ -190,6 +200,7 @@ def __init__( ) self.group_size = group_size self.checkpoint_uses_packed_qkv = checkpoint_uses_packed_qkv + self.swap_weight_nibbles = swap_weight_nibbles @classmethod def get_name(cls) -> str: @@ -237,6 +248,7 @@ def _add_group_size_from_dict(config: dict): def from_config(cls, config: Dict[str, Any]) -> ModelOptFp4Config: group_size = None exclude_modules = [] + swap_weight_nibbles = True # Flat format (config.json quantization_config) quant_method = config.get("quant_algo") @@ -248,6 +260,7 @@ def from_config(cls, config: Dict[str, Any]) -> ModelOptFp4Config: first_group = next(iter(config_groups.values()), {}) group_size = first_group.get("weights", {}).get("group_size") exclude_modules = config.get("ignore", []) + swap_weight_nibbles = config.get("swap_weight_nibbles", True) else: # Nested format (hf_quant_config.json) try: @@ -255,6 +268,10 @@ def from_config(cls, config: Dict[str, Any]) -> ModelOptFp4Config: quant_method = quant_config["quant_algo"] group_size = ModelOptFp4Config.common_group_size(config) exclude_modules = quant_config.get("exclude_modules", []) + swap_weight_nibbles = quant_config.get( + "swap_weight_nibbles", + config.get("swap_weight_nibbles", True), + ) except (ValueError, KeyError): raise ValueError("Cannot find 'quant_algo' in quantization config.") @@ -274,6 +291,7 @@ def from_config(cls, config: Dict[str, Any]) -> ModelOptFp4Config: exclude_modules=exclude_modules, packed_modules_mapping=config.get("packed_modules_mapping"), checkpoint_uses_packed_qkv=config.get("checkpoint_uses_packed_qkv", False), + swap_weight_nibbles=swap_weight_nibbles, ) def get_quant_method(self, layer: torch.nn.Module, prefix: str): @@ -459,9 +477,10 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.output_size_per_partition = layer.weight.shape[0] - # Swap nibbles: (byte >> 4) | (byte << 4). w = layer.weight.data - w_swapped = ((w >> 4) | (w << 4)).contiguous() + w_swapped = _prepare_nvfp4_weight_bytes( + w, swap_weight_nibbles=self.quant_config.swap_weight_nibbles + ) weight, weights_padding_cols = pad_nvfp4_weight(w_swapped) layer.weights_padding_cols = weights_padding_cols copy_or_rebind_param(layer, "weight", weight) diff --git a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py index 2941f5e24f95..c0375be9b385 100644 --- a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py +++ b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py @@ -410,7 +410,16 @@ def load_model_from_full_model_state_dict( ): requires_grad = False temp_param.requires_grad = requires_grad - weight_loader(temp_param, full_tensor) + try: + weight_loader(temp_param, full_tensor) + except AssertionError as exc: + raise AssertionError( + "Failed to shard/load parameter " + f"{target_param_name}: full_tensor.shape={tuple(full_tensor.shape)}, " + f"meta_sharded_param.shape={tuple(meta_sharded_param.shape)}, " + f"temp_param.shape={tuple(temp_param.shape)}, " + f"param_cls={type(actual_param).__name__}" + ) from exc sharded_tensor = temp_param.data else: # In cases where parts of the model aren't sharded, some parameters will be plain tensors diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux.py b/python/sglang/multimodal_gen/runtime/models/dits/flux.py index 6ecfc22503fd..481296a2da99 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/flux.py @@ -261,6 +261,7 @@ def __init__( bias=bias, gather_output=True, quant_config=quant_config, + prefix=f"{prefix}.to_q" if prefix else "to_q", ) self.to_k = ColumnParallelLinear( query_dim, @@ -268,6 +269,7 @@ def __init__( bias=bias, gather_output=True, quant_config=quant_config, + prefix=f"{prefix}.to_k" if prefix else "to_k", ) self.to_v = ColumnParallelLinear( query_dim, @@ -275,6 +277,7 @@ def __init__( bias=bias, gather_output=True, quant_config=quant_config, + prefix=f"{prefix}.to_v" if prefix else "to_v", ) if not self.pre_only: self.to_out = torch.nn.ModuleList([]) @@ -310,6 +313,7 @@ def __init__( bias=added_proj_bias, gather_output=True, quant_config=quant_config, + prefix=f"{prefix}.add_q_proj" if prefix else "add_q_proj", ) self.add_k_proj = ColumnParallelLinear( added_kv_proj_dim, @@ -317,6 +321,7 @@ def __init__( bias=added_proj_bias, gather_output=True, quant_config=quant_config, + prefix=f"{prefix}.add_k_proj" if prefix else "add_k_proj", ) self.add_v_proj = ColumnParallelLinear( added_kv_proj_dim, @@ -324,6 +329,7 @@ def __init__( bias=added_proj_bias, gather_output=True, quant_config=quant_config, + prefix=f"{prefix}.add_v_proj" if prefix else "add_v_proj", ) self.to_add_out = ColumnParallelLinear( self.inner_dim, @@ -495,6 +501,7 @@ def __init__( bias=True, gather_output=True, quant_config=quant_config, + prefix=f"{prefix}.proj_mlp" if prefix else "proj_mlp", ) self.act_mlp = nn.GELU(approximate="tanh") self.proj_out = ColumnParallelLinear( @@ -503,6 +510,7 @@ def __init__( bias=True, gather_output=True, quant_config=quant_config, + prefix=f"{prefix}.proj_out" if prefix else "proj_out", ) self.attn = FluxAttention( query_dim=dim, @@ -513,6 +521,7 @@ def __init__( eps=1e-6, pre_only=True, quant_config=quant_config, + prefix=f"{prefix}.attn" if prefix else "attn", ) def forward( 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 7f385c967d69..4cd004784567 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -19,6 +19,7 @@ from einops import rearrange from tqdm.auto import tqdm +from sglang.jit_kernel.nvfp4 import prewarm_nvfp4_jit_modules from sglang.multimodal_gen import envs from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType, STA_Mode from sglang.multimodal_gen.configs.pipeline_configs.flux import ( @@ -206,9 +207,26 @@ def _maybe_enable_torch_compile(self, module: object) -> None: compile_kwargs["mode"] = mode logger.info(f"Compiling transformer with mode: {mode}") + if self._needs_nvfp4_jit_prewarm(module): + logger.info( + "Prewarming NVFP4 JIT modules before torch.compile to avoid " + "Dynamo tracing JIT initialization." + ) + prewarm_nvfp4_jit_modules() + # TODO(triple-mu): support customized fullgraph and dynamic in the future module.compile(**compile_kwargs) + @staticmethod + def _needs_nvfp4_jit_prewarm(module: nn.Module) -> bool: + for submodule in module.modules(): + quant_method = getattr(submodule, "quant_method", None) + if quant_method is None: + continue + if type(quant_method).__name__ == "ModelOptFp4LinearMethod": + return True + return False + def _maybe_enable_cache_dit( self, num_inference_steps: int | tuple[int, int], batch: Req ) -> None: diff --git a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py index 4d37b5bc7d02..3a7340d4c082 100644 --- a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py +++ b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py @@ -10,6 +10,8 @@ from types import SimpleNamespace from unittest.mock import patch +import torch + partial_json_parser = types.ModuleType("partial_json_parser") partial_json_parser_core = types.ModuleType("partial_json_parser.core") partial_json_parser_exceptions = types.ModuleType("partial_json_parser.core.exceptions") @@ -41,15 +43,24 @@ def _loads(input_str, _flags=None): ) sys.modules.setdefault("partial_json_parser.core.options", partial_json_parser_options) +from sglang.multimodal_gen.runtime.layers.linear import UnquantizedLinearMethod from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import ( NunchakuConfig, ) +from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import ( + ModelOptFp4Config, + _prepare_nvfp4_weight_bytes, +) from sglang.multimodal_gen.runtime.loader.transformer_load_utils import ( _filter_duplicate_precision_variant_safetensors, _Flux2Nvfp4FallbackAdapter, resolve_transformer_quant_load_spec, resolve_transformer_safetensors_to_load, ) +from sglang.multimodal_gen.runtime.models.dits.flux import FluxSingleTransformerBlock +from sglang.multimodal_gen.tools.build_modelopt_nvfp4_transformer import ( + _updated_quant_config, +) class _FakeFluxTransformer: @@ -148,22 +159,21 @@ def test_resolve_transformer_quant_load_spec_keeps_nunchaku_hook( mock_metadata.return_value = { "config": json.dumps({"_class_name": _FakeFluxTransformer.__name__}) } - nunchaku_config = NunchakuConfig( - transformer_weights_path="/tmp/svdq-int4_r32.safetensors" - ) - server_args = self._make_server_args( - transformer_weights_path=nunchaku_config.transformer_weights_path, - nunchaku_config=nunchaku_config, - ) + with tempfile.NamedTemporaryFile(suffix=".safetensors") as f: + nunchaku_config = NunchakuConfig(transformer_weights_path=f.name) + server_args = self._make_server_args( + transformer_weights_path=nunchaku_config.transformer_weights_path, + nunchaku_config=nunchaku_config, + ) - spec = resolve_transformer_quant_load_spec( - hf_config={}, - server_args=server_args, - safetensors_list=[nunchaku_config.transformer_weights_path], - component_model_path="/unused/component/path", - model_cls=_FakeFluxTransformer, - cls_name=_FakeFluxTransformer.__name__, - ) + spec = resolve_transformer_quant_load_spec( + hf_config={}, + server_args=server_args, + safetensors_list=[nunchaku_config.transformer_weights_path], + component_model_path="/unused/component/path", + model_cls=_FakeFluxTransformer, + cls_name=_FakeFluxTransformer.__name__, + ) self.assertIsNone(spec.quant_config) self.assertIs(spec.nunchaku_config, nunchaku_config) @@ -189,6 +199,104 @@ def test_flux2_mixed_nvfp4_fallback_disables_conflicting_offloads(self): self.assertFalse(server_args.dit_cpu_offload) self.assertFalse(server_args.text_encoder_cpu_offload) + def test_prepare_nvfp4_weight_bytes_swaps_nibbles(self): + weight = torch.tensor([[0xAB, 0x10]], dtype=torch.uint8) + + prepared = _prepare_nvfp4_weight_bytes(weight, swap_weight_nibbles=True) + + self.assertEqual(prepared.tolist(), [[0xBA, 0x01]]) + + def test_prepare_nvfp4_weight_bytes_can_skip_nibble_swap(self): + weight = torch.tensor([[0xAB, 0x10]], dtype=torch.uint8) + + prepared = _prepare_nvfp4_weight_bytes(weight, swap_weight_nibbles=False) + + self.assertEqual(prepared.tolist(), [[0xAB, 0x10]]) + + def test_modelopt_fp4_config_reads_swap_weight_nibbles_from_flat_config(self): + config = ModelOptFp4Config.from_config( + { + "quant_algo": "NVFP4", + "group_size": 16, + "ignore": [], + "swap_weight_nibbles": False, + } + ) + + self.assertFalse(config.swap_weight_nibbles) + + def test_modelopt_fp4_config_reads_swap_weight_nibbles_from_nested_config(self): + config = ModelOptFp4Config.from_config( + { + "quantization": { + "quant_algo": "NVFP4", + "exclude_modules": [], + "swap_weight_nibbles": False, + }, + "config_groups": {"default": {"weights": {"group_size": 16}}}, + } + ) + + self.assertFalse(config.swap_weight_nibbles) + + def test_builder_adds_diffusers_quant_type_for_nvfp4(self): + updated = _updated_quant_config( + { + "quantization_config": { + "quant_method": "modelopt", + "quant_algo": "NVFP4", + "ignore": [], + } + }, + fallback_patterns=["single_transformer_blocks.*.proj_mlp*"], + swap_weight_nibbles=False, + ) + + self.assertEqual(updated["quantization_config"]["quant_type"], "NVFP4") + self.assertEqual( + updated["quantization_config"]["ignore"], + ["single_transformer_blocks.*.proj_mlp*"], + ) + + @patch("sglang.multimodal_gen.runtime.layers.linear.get_group_rank", return_value=0) + @patch("sglang.multimodal_gen.runtime.layers.linear.get_group_size", return_value=1) + @patch( + "sglang.multimodal_gen.runtime.layers.linear.get_tp_group", return_value=None + ) + def test_flux_single_transformer_block_modelopt_excludes_use_full_prefix( + self, + _mock_tp_group, + _mock_group_size, + _mock_group_rank, + ): + quant_config = ModelOptFp4Config( + is_checkpoint_nvfp4_serialized=True, + group_size=16, + exclude_modules=[ + "single_transformer_blocks.*.proj_mlp*", + "single_transformer_blocks.*.proj_out*", + "single_transformer_blocks.*.attn.to_q", + ], + ) + + block = FluxSingleTransformerBlock( + dim=64, + num_attention_heads=4, + attention_head_dim=16, + mlp_ratio=2.0, + quant_config=quant_config, + prefix="single_transformer_blocks.0", + ) + + self.assertEqual(block.proj_mlp.prefix, "single_transformer_blocks.0.proj_mlp") + self.assertEqual(block.proj_out.prefix, "single_transformer_blocks.0.proj_out") + self.assertEqual( + block.attn.to_q.prefix, "single_transformer_blocks.0.attn.to_q" + ) + self.assertIsInstance(block.proj_mlp.quant_method, UnquantizedLinearMethod) + self.assertIsInstance(block.proj_out.quant_method, UnquantizedLinearMethod) + self.assertIsInstance(block.attn.to_q.quant_method, UnquantizedLinearMethod) + if __name__ == "__main__": unittest.main() diff --git a/python/sglang/multimodal_gen/tools/convert_modelopt_fp8_checkpoint.py b/python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py similarity index 66% rename from python/sglang/multimodal_gen/tools/convert_modelopt_fp8_checkpoint.py rename to python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py index 89a731fba863..5d87f5ff1bf2 100644 --- a/python/sglang/multimodal_gen/tools/convert_modelopt_fp8_checkpoint.py +++ b/python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py @@ -1,4 +1,4 @@ -"""Convert a ModelOpt diffusion FP8 export into an SGLang-loadable checkpoint. +"""Build an SGLang-loadable ModelOpt FP8 diffusion transformer. The core conversion path is model-agnostic: - read the ModelOpt diffusers transformer export @@ -12,7 +12,7 @@ Example: - python -m sglang.multimodal_gen.tools.convert_modelopt_fp8_checkpoint \ + python -m sglang.multimodal_gen.tools.build_modelopt_fp8_transformer \ --modelopt-hf-dir /tmp/modelopt_flux2_fp8/hf \ --modelopt-backbone-ckpt /tmp/modelopt_flux2_fp8/ckpt/backbone.pt \ --base-transformer-dir /path/to/FLUX.2-dev/transformer \ @@ -57,6 +57,21 @@ r"^transformer_blocks\.\d+\.ff_context\.net\.2$", r"^single_transformer_blocks\.\d+\.norm\.linear$", ] +DEFAULT_LTX2_KEEP_BF16_PATTERNS = [ + r"^(audio_)?adaln_single\.emb\.timestep_embedder\.linear_[12]$", + r"^(audio_)?adaln_single\.linear$", + r"^audio_caption_projection\.linear_[12]$", + r"^audio_patchify_proj$", + r"^audio_proj_out$", + r"^av_ca_(a2v_gate|audio_scale_shift|v2a_gate|video_scale_shift)_adaln_single\.emb\.timestep_embedder\.linear_[12]$", + r"^av_ca_(a2v_gate|audio_scale_shift|v2a_gate|video_scale_shift)_adaln_single\.linear$", + r"^caption_projection\.linear_[12]$", + r"^patchify_proj$", + r"^proj_out$", + r"^transformer_blocks\.(0|43|44|45|46|47)\.(attn1|attn2|audio_attn1|audio_attn2|audio_to_video_attn|video_to_audio_attn)\.to_(q|k|v)$", + r"^transformer_blocks\.(0|43|44|45|46|47)\.(attn1|attn2|audio_attn1|audio_attn2|audio_to_video_attn|video_to_audio_attn)\.to_out\.0$", + r"^transformer_blocks\.(0|43|44|45|46|47)\.(ff|audio_ff)\.proj_(in|out)$", +] def _resolve_transformer_dir(path: str) -> str: @@ -126,9 +141,116 @@ def _load_config(model_dir: str) -> dict: return json.load(f) +def _load_first_shard_metadata( + model_dir: str, weight_map: Mapping[str, str] +) -> dict[str, str]: + if not weight_map: + return {} + first_shard = next(iter(weight_map.values())) + with safe_open( + os.path.join(model_dir, first_shard), framework="pt", device="cpu" + ) as f: + return dict(f.metadata() or {}) + + +def _module_name_variants(weight_name: str) -> list[str]: + module_name = weight_name[:-7] if weight_name.endswith(".weight") else weight_name + variants = [module_name] + + for prefix in ("model.diffusion_model.", "velocity_model."): + if module_name.startswith(prefix): + variants.append(module_name[len(prefix) :]) + + canonicalized: list[str] = [] + for variant in variants: + canonicalized.append( + re.sub(r"(\.audio_ff|\.ff)\.net\.0\.proj$", r"\1.proj_in", variant) + ) + canonicalized.append( + re.sub(r"(\.audio_ff|\.ff)\.net\.2$", r"\1.proj_out", variant) + ) + variants.extend(canonicalized) + + deduped: list[str] = [] + for variant in variants: + if variant not in deduped: + deduped.append(variant) + return deduped + + +def _preferred_module_name(weight_name: str) -> str: + return _module_name_variants(weight_name)[-1] + + +def _scale_key_candidates(weight_name: str) -> list[str]: + candidates = [weight_name] + if weight_name.startswith("model.diffusion_model."): + candidates.append( + "velocity_model." + weight_name[len("model.diffusion_model.") :] + ) + return candidates + + +def _resolve_scale_key( + weight_name: str, + scale_map: Mapping[str, Mapping[str, torch.Tensor]], +) -> str | None: + for candidate in _scale_key_candidates(weight_name): + if candidate in scale_map: + return candidate + return None + + +def _is_ltx2_x0_export( + *, + config: Mapping[str, object], + source_metadata: Mapping[str, str], + source_weight_map: Mapping[str, str], +) -> bool: + if config.get("_class_name") != "X0Model": + return False + if not any(name.startswith("model.diffusion_model.") for name in source_weight_map): + return False + try: + metadata_config = json.loads(str(source_metadata.get("config", ""))) + except json.JSONDecodeError: + return False + return isinstance(metadata_config.get("transformer"), dict) + + +def _build_output_config( + *, + source_config: Mapping[str, object], + source_metadata: Mapping[str, str], + quant_config: Mapping[str, object], + is_ltx2_x0_export: bool, +) -> dict[str, object]: + if is_ltx2_x0_export: + metadata_config = json.loads(str(source_metadata["config"])) + output_config = dict(metadata_config["transformer"]) + output_config["_class_name"] = "LTX2VideoTransformer3DModel" + else: + output_config = dict(source_config) + + output_config["quantization_config"] = dict(quant_config) + return output_config + + +def _should_keep_ltx2_transformer_key(weight_name: str) -> bool: + if not weight_name.startswith("model.diffusion_model."): + return False + connector_prefixes = ( + "model.diffusion_model.audio_embeddings_connector.", + "model.diffusion_model.video_embeddings_connector.", + ) + return not weight_name.startswith(connector_prefixes) + + def get_default_keep_bf16_patterns( *, model_type: str, class_name: str | None ) -> list[str]: + if model_type == "ltx2": + return list(DEFAULT_LTX2_KEEP_BF16_PATTERNS) if model_type == "flux1": return list(DEFAULT_FLUX1_KEEP_BF16_PATTERNS) if model_type == "flux2": @@ -149,8 +271,11 @@ def should_keep_bf16( if not keep_bf16_patterns: return False - module_name = weight_name[:-7] if weight_name.endswith(".weight") else weight_name - return any(re.search(pattern, module_name) for pattern in keep_bf16_patterns) + return any( + re.search(pattern, module_name) + for pattern in keep_bf16_patterns + for module_name in _module_name_variants(weight_name) + ) def is_ignored_by_modelopt( @@ -160,10 +285,12 @@ def is_ignored_by_modelopt( if not ignore_patterns: return False - module_name = weight_name[:-7] if weight_name.endswith(".weight") else weight_name for pattern in ignore_patterns: regex_str = pattern.replace(".", r"\.").replace("*", r".*") - if re.fullmatch(regex_str, module_name): + if any( + re.fullmatch(regex_str, module_name) + for module_name in _module_name_variants(weight_name) + ): return True return False @@ -242,7 +369,7 @@ def _load_selected_tensors( return tensors -def convert_modelopt_fp8_checkpoint( +def build_modelopt_fp8_transformer( *, modelopt_hf_dir: str, modelopt_backbone_ckpt: str, @@ -265,23 +392,29 @@ def convert_modelopt_fp8_checkpoint( raise ValueError( "Expected a flat quantization_config dict in the ModelOpt export." ) - if ( - quant_config.get("quant_method") != "modelopt" - or "FP8" not in str(quant_config.get("quant_algo", "")).upper() - ): + if quant_config.get("quant_method") != "modelopt": raise ValueError( "This tool only supports ModelOpt diffusers FP8 exports " - "(quant_method=modelopt, quant_algo=FP8)." + "(quant_method=modelopt)." ) + source_weight_map_all, index_filename = _load_weight_map(source_dir) + source_metadata = _load_first_shard_metadata(source_dir, source_weight_map_all) + is_ltx2_export = _is_ltx2_x0_export( + config=config, + source_metadata=source_metadata, + source_weight_map=source_weight_map_all, + ) class_name = config.get("_class_name") ignore_patterns = list(quant_config.get("ignore", []) or []) patterns = list( get_default_keep_bf16_patterns(model_type=model_type, class_name=class_name) ) + if is_ltx2_export and model_type == "auto": + patterns.extend(DEFAULT_LTX2_KEEP_BF16_PATTERNS) if keep_bf16_patterns: patterns.extend(keep_bf16_patterns) - if patterns and base_dir is None: + if patterns and base_dir is None and not is_ltx2_export: raise ValueError( "BF16 fallback patterns are enabled, but --base-transformer-dir was not provided." ) @@ -298,25 +431,73 @@ def convert_modelopt_fp8_checkpoint( _copy_non_shard_files(source_dir, str(output_path)) - source_weight_map, index_filename = _load_weight_map(source_dir) + if is_ltx2_export: + source_weight_map = { + name: filename + for name, filename in source_weight_map_all.items() + if _should_keep_ltx2_transformer_key(name) + } + else: + source_weight_map = source_weight_map_all base_weight_map: dict[str, str] = {} if base_dir is not None: base_weight_map, _ = _load_weight_map(base_dir) + fallback_weight_names = sorted( + weight_name + for weight_name in source_weight_map + if weight_name.endswith(".weight") and should_keep_bf16(weight_name, patterns) + ) + fallback_weight_names_set = set(fallback_weight_names) backbone_state = torch.load(backbone_ckpt_path, map_location="cpu")[ "model_state_dict" ] fp8_scale_map = build_fp8_scale_map(backbone_state, maxbound=maxbound) - serialized_quant_config = json.dumps(quant_config, sort_keys=True) - - fallback_weight_names = sorted( - weight_name - for weight_name in source_weight_map - if weight_name.endswith(".weight") and should_keep_bf16(weight_name, patterns) + quant_algo = str(quant_config.get("quant_algo", "")).upper() + if quant_algo and "FP8" not in quant_algo: + raise ValueError( + "This tool only supports ModelOpt diffusers FP8 exports, " + f"got quant_algo={quant_config.get('quant_algo')!r}." + ) + if not quant_algo and not fp8_scale_map: + raise ValueError( + "Could not infer an FP8 ModelOpt export: quantization_config.quant_algo " + "is missing and backbone.pt does not contain FP8 scale tensors." + ) + effective_quant_config = json.loads(json.dumps(quant_config)) + if not quant_algo: + effective_quant_config["quant_algo"] = "FP8" + + auto_ignore_modules = sorted( + { + _preferred_module_name(weight_name) + for weight_name in source_weight_map + if weight_name.endswith(".weight") + and _resolve_scale_key(weight_name, fp8_scale_map) is None + } + ) + fallback_ignore_modules = sorted( + {_preferred_module_name(weight_name) for weight_name in fallback_weight_names} ) + ignore_patterns = sorted( + { + *ignore_patterns, + *auto_ignore_modules, + *fallback_ignore_modules, + } + ) + effective_quant_config["ignore"] = ignore_patterns + serialized_quant_config = json.dumps(effective_quant_config, sort_keys=True) + output_config = _build_output_config( + source_config=config, + source_metadata=source_metadata, + quant_config=effective_quant_config, + is_ltx2_x0_export=is_ltx2_export, + ) + fallback_tensors = ( _load_selected_tensors(base_dir, base_weight_map, fallback_weight_names) - if fallback_weight_names + if fallback_weight_names and base_dir is not None else {} ) fallback_scale_names = { @@ -340,15 +521,23 @@ def convert_modelopt_fp8_checkpoint( for filename, names in sorted(weights_by_file.items()): shard_path = os.path.join(source_dir, filename) shard_tensors = load_file(shard_path, device="cpu") + selected_names = set(names) with safe_open(shard_path, framework="pt", device="cpu") as f: metadata = dict(f.metadata() or {}) metadata.setdefault("format", "pt") + metadata["_class_name"] = str( + output_config.get("_class_name", metadata.get("_class_name", "")) + ) + metadata["config"] = json.dumps(output_config, sort_keys=True) metadata["quantization_config"] = serialized_quant_config metadata["_quantization_metadata"] = serialized_quant_config for name in list(shard_tensors.keys()): + if name not in selected_names: + del shard_tensors[name] + continue if "_quantizer." in name: del shard_tensors[name] continue @@ -362,12 +551,14 @@ def convert_modelopt_fp8_checkpoint( continue if name in fallback_tensors: shard_tensors[name] = fallback_tensors[name] + scale_key = _resolve_scale_key(name, fp8_scale_map) if ( name.endswith(".weight") - and name in fp8_scale_map + and scale_key is not None and name not in fallback_tensors + and name not in fallback_weight_names_set ): - scale_tensors = fp8_scale_map[name] + scale_tensors = fp8_scale_map[scale_key] shard_tensors[name] = quantize_fp8_weight( shard_tensors[name], scale_tensors["weight_scale"] ) @@ -397,12 +588,15 @@ def convert_modelopt_fp8_checkpoint( sort_keys=True, ) + with open(output_path / "config.json", "w", encoding="utf-8") as f: + json.dump(output_config, f, indent=2, sort_keys=True) + return { "quantized_weights": sum( 1 for name in source_weight_map if name.endswith(".weight") - and name in fp8_scale_map + and _resolve_scale_key(name, fp8_scale_map) is not None and not is_ignored_by_modelopt(name, ignore_patterns) ), "bf16_fallback_weights": len(fallback_weight_names), @@ -415,8 +609,8 @@ def convert_modelopt_fp8_checkpoint( def _parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description=( - "Inject FP8 scales from ModelOpt backbone.pt into a diffusers export so " - "SGLang diffusion can load it natively." + "Build an SGLang-loadable ModelOpt FP8 diffusion transformer from a " + "ModelOpt diffusers export." ) ) parser.add_argument( @@ -443,11 +637,11 @@ def _parse_args() -> argparse.Namespace: ) parser.add_argument( "--model-type", - choices=["auto", "flux1", "flux2", "none"], + choices=["auto", "flux1", "flux2", "ltx2", "none"], default="auto", help=( "Optional model-family BF16 fallback profile. 'none' uses the generic " - "conversion path. 'auto' enables the validated FLUX.1 / FLUX.2 " + "conversion path. 'auto' enables the validated FLUX.1 / FLUX.2 / LTX-2 " "fallback set when the export config matches those transformer classes." ), ) @@ -477,7 +671,7 @@ def _parse_args() -> argparse.Namespace: def main() -> None: args = _parse_args() - stats = convert_modelopt_fp8_checkpoint( + stats = build_modelopt_fp8_transformer( modelopt_hf_dir=args.modelopt_hf_dir, modelopt_backbone_ckpt=args.modelopt_backbone_ckpt, output_dir=args.output_dir, diff --git a/python/sglang/multimodal_gen/tools/build_modelopt_nvfp4_transformer.py b/python/sglang/multimodal_gen/tools/build_modelopt_nvfp4_transformer.py new file mode 100644 index 000000000000..e6f13b306a41 --- /dev/null +++ b/python/sglang/multimodal_gen/tools/build_modelopt_nvfp4_transformer.py @@ -0,0 +1,402 @@ +"""Build an SGLang-loadable ModelOpt NVFP4 diffusion transformer. + +This tool keeps the ModelOpt-exported NVFP4 tensors for most transformer +modules, but can replace a validated subset of numerically sensitive modules +with their original BF16 tensors from the base transformer checkpoint. + +It is primarily intended for FLUX.1-dev style ModelOpt NVFP4 exports where: +- the base pipeline should remain separate from the quantized transformer +- fallback BF16 modules are model-family specific +- the serialized FP4 weight byte order may already match the runtime kernel +""" + +from __future__ import annotations + +import argparse +import json +import os +import re +import shutil +from collections import defaultdict +from pathlib import Path +from typing import Iterable, Mapping, Sequence + +from safetensors import safe_open +from safetensors.torch import load_file, save_file + +INDEX_FILENAMES = [ + "model.safetensors.index.json", + "diffusion_pytorch_model.safetensors.index.json", +] + +DEFAULT_FLUX1_NVFP4_FALLBACK_PATTERNS = [ + "transformer_blocks.*.norm1.linear*", + "transformer_blocks.*.norm1_context.linear*", + "transformer_blocks.*.ff.net.0.proj*", + "transformer_blocks.*.ff.net.2*", + "transformer_blocks.*.ff_context.net.0.proj*", + "transformer_blocks.*.ff_context.net.2*", + "single_transformer_blocks.*.norm.linear*", + "single_transformer_blocks.*.proj_mlp*", +] + +_TENSOR_MODULE_SUFFIXES = ( + ".weight_scale_2", + ".weight_scale", + ".input_scale", + ".weight", + ".bias", +) + + +def _resolve_transformer_dir(path: str) -> str: + candidate = Path(path).expanduser().resolve() + if (candidate / "config.json").is_file(): + return str(candidate) + transformer_dir = candidate / "transformer" + if (transformer_dir / "config.json").is_file(): + return str(transformer_dir) + raise FileNotFoundError(f"Could not resolve a transformer directory from: {path}") + + +def _find_index_file(model_dir: str) -> str | None: + for filename in INDEX_FILENAMES: + candidate = os.path.join(model_dir, filename) + if os.path.isfile(candidate): + return filename + + matches = sorted( + filename + for filename in os.listdir(model_dir) + if filename.endswith(".safetensors.index.json") + ) + return matches[0] if matches else None + + +def _load_weight_map(model_dir: str) -> tuple[dict[str, str], str | None]: + index_filename = _find_index_file(model_dir) + if index_filename is not None: + with open(os.path.join(model_dir, index_filename), encoding="utf-8") as f: + index_data = json.load(f) + return dict(index_data["weight_map"]), index_filename + + safetensors_files = sorted( + filename + for filename in os.listdir(model_dir) + if filename.endswith(".safetensors") + ) + if len(safetensors_files) != 1: + raise ValueError( + f"Expected an index file or a single safetensors shard in {model_dir}, " + f"found {len(safetensors_files)} shard(s)." + ) + + shard_name = safetensors_files[0] + with safe_open( + os.path.join(model_dir, shard_name), framework="pt", device="cpu" + ) as f: + weight_map = {key: shard_name for key in f.keys()} + index_filename = f"{Path(shard_name).stem}.safetensors.index.json" + return weight_map, index_filename + + +def _load_config(model_dir: str) -> dict: + config_path = os.path.join(model_dir, "config.json") + with open(config_path, encoding="utf-8") as f: + return json.load(f) + + +def _write_config(model_dir: Path, config: Mapping[str, object]) -> None: + with open(model_dir / "config.json", "w", encoding="utf-8") as f: + json.dump(config, f, indent=2, sort_keys=True) + f.write("\n") + + +def _copy_non_shard_files(source_dir: str, output_dir: str) -> None: + ignored = set(INDEX_FILENAMES) + for entry in os.listdir(source_dir): + if entry.endswith(".safetensors") or entry in ignored: + continue + source_path = os.path.join(source_dir, entry) + output_path = os.path.join(output_dir, entry) + if os.path.isdir(source_path): + shutil.copytree(source_path, output_path, dirs_exist_ok=True) + else: + shutil.copy2(source_path, output_path) + + +def _load_selected_tensors( + model_dir: str, + weight_map: Mapping[str, str], + tensor_names: Iterable[str], +): + tensors = {} + names_by_file: dict[str, list[str]] = defaultdict(list) + for name in tensor_names: + names_by_file[weight_map[name]].append(name) + + for filename, names in names_by_file.items(): + shard_path = os.path.join(model_dir, filename) + with safe_open(shard_path, framework="pt", device="cpu") as f: + for name in names: + tensors[name] = f.get_tensor(name).contiguous() + return tensors + + +def _module_name_for_tensor(tensor_name: str) -> str: + for suffix in _TENSOR_MODULE_SUFFIXES: + if tensor_name.endswith(suffix): + return tensor_name[: -len(suffix)] + return tensor_name + + +def _matches_any_pattern(module_name: str, patterns: Sequence[str]) -> bool: + if not patterns: + return False + for pattern in patterns: + regex_str = pattern.replace(".", r"\.").replace("*", r".*") + if re.fullmatch(regex_str, module_name): + return True + return False + + +def _preset_patterns(pattern_preset: str) -> list[str]: + if pattern_preset == "none": + return [] + if pattern_preset == "flux1-nvfp4": + return list(DEFAULT_FLUX1_NVFP4_FALLBACK_PATTERNS) + raise ValueError(f"Unsupported pattern preset: {pattern_preset}") + + +def _updated_quant_config( + source_config: Mapping[str, object], + *, + fallback_patterns: Sequence[str], + swap_weight_nibbles: bool, +) -> dict[str, object]: + output_config = json.loads(json.dumps(source_config)) + quant_config = output_config.get("quantization_config") + if not isinstance(quant_config, dict): + raise ValueError("Expected a flat quantization_config dict in config.json.") + if ( + quant_config.get("quant_method") != "modelopt" + or "FP4" not in str(quant_config.get("quant_algo", "")).upper() + ): + raise ValueError( + "This tool only supports ModelOpt diffusion NVFP4 exports " + "(quant_method=modelopt, quant_algo=FP4/NVFP4)." + ) + + ignore_patterns = list(quant_config.get("ignore", []) or []) + for pattern in fallback_patterns: + if pattern not in ignore_patterns: + ignore_patterns.append(pattern) + + quant_config["ignore"] = ignore_patterns + quant_config.setdefault( + "quant_type", str(quant_config.get("quant_algo", "")).upper() + ) + quant_config["swap_weight_nibbles"] = swap_weight_nibbles + return output_config + + +def build_modelopt_nvfp4_transformer( + *, + base_transformer_dir: str, + modelopt_hf_dir: str, + output_dir: str, + pattern_preset: str = "none", + keep_bf16_patterns: Sequence[str] | None = None, + swap_weight_nibbles: bool | None = None, + overwrite: bool = False, +) -> dict[str, int | bool]: + source_dir = _resolve_transformer_dir(modelopt_hf_dir) + base_dir = _resolve_transformer_dir(base_transformer_dir) + + patterns = _preset_patterns(pattern_preset) + if keep_bf16_patterns: + patterns.extend(keep_bf16_patterns) + + resolved_swap_weight_nibbles = ( + swap_weight_nibbles + if swap_weight_nibbles is not None + else (False if pattern_preset == "flux1-nvfp4" else True) + ) + output_config = _updated_quant_config( + _load_config(source_dir), + fallback_patterns=patterns, + swap_weight_nibbles=resolved_swap_weight_nibbles, + ) + quant_config = output_config["quantization_config"] + serialized_quant_config = json.dumps(quant_config, sort_keys=True) + + output_path = Path(output_dir).expanduser().resolve() + if output_path.exists(): + if not overwrite: + raise FileExistsError( + f"Output directory already exists: {output_path}. " + "Use --overwrite to replace it." + ) + shutil.rmtree(output_path) + output_path.mkdir(parents=True, exist_ok=True) + + _copy_non_shard_files(source_dir, str(output_path)) + _write_config(output_path, output_config) + + source_weight_map, index_filename = _load_weight_map(source_dir) + base_weight_map, _ = _load_weight_map(base_dir) + + fallback_tensor_names = sorted( + name + for name in base_weight_map + if name in source_weight_map + and _matches_any_pattern(_module_name_for_tensor(name), patterns) + ) + fallback_tensors = _load_selected_tensors( + base_dir, + base_weight_map, + fallback_tensor_names, + ) + fallback_modules = { + _module_name_for_tensor(tensor_name) for tensor_name in fallback_tensor_names + } + + weights_by_file: dict[str, list[str]] = defaultdict(list) + for tensor_name, filename in source_weight_map.items(): + weights_by_file[filename].append(tensor_name) + + updated_weight_map: dict[str, str] = {} + total_size = 0 + replaced_tensor_count = 0 + removed_aux_tensor_count = 0 + + for filename, tensor_names in sorted(weights_by_file.items()): + shard_path = os.path.join(source_dir, filename) + shard_tensors = load_file(shard_path, device="cpu") + + with safe_open(shard_path, framework="pt", device="cpu") as f: + metadata = dict(f.metadata() or {}) + + metadata.setdefault("format", "pt") + metadata["quantization_config"] = serialized_quant_config + metadata["_quantization_metadata"] = serialized_quant_config + + for name in list(shard_tensors.keys()): + if "_quantizer." in name: + del shard_tensors[name] + removed_aux_tensor_count += 1 + continue + + module_name = _module_name_for_tensor(name) + if module_name not in fallback_modules: + continue + + if name in fallback_tensors: + shard_tensors[name] = fallback_tensors[name] + replaced_tensor_count += 1 + else: + del shard_tensors[name] + removed_aux_tensor_count += 1 + + save_file(shard_tensors, os.path.join(output_path, filename), metadata=metadata) + + for name, tensor in shard_tensors.items(): + updated_weight_map[name] = filename + total_size += tensor.element_size() * tensor.numel() + + if index_filename is None: + raise ValueError( + "Expected a sharded or indexed ModelOpt HF export, but no index file was found." + ) + + with open(output_path / index_filename, "w", encoding="utf-8") as f: + json.dump( + { + "metadata": {"total_size": total_size}, + "weight_map": updated_weight_map, + }, + f, + indent=2, + sort_keys=True, + ) + f.write("\n") + + return { + "fallback_modules": len(fallback_modules), + "replaced_tensors": replaced_tensor_count, + "removed_aux_tensors": removed_aux_tensor_count, + "output_shards": len(weights_by_file), + "swap_weight_nibbles": resolved_swap_weight_nibbles, + } + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Build an SGLang-loadable ModelOpt NVFP4 diffusion transformer and " + "optionally keep selected modules in BF16." + ) + ) + parser.add_argument( + "--base-transformer-dir", + required=True, + help="Original BF16 transformer directory, or a parent model directory.", + ) + parser.add_argument( + "--modelopt-hf-dir", + required=True, + help="ModelOpt --hf-ckpt-dir output, or its transformer subdirectory.", + ) + parser.add_argument( + "--output-dir", + required=True, + help="Directory to write the mixed transformer checkpoint.", + ) + parser.add_argument( + "--pattern-preset", + choices=["none", "flux1-nvfp4"], + default="none", + help="Optional model-family BF16 fallback preset.", + ) + parser.add_argument( + "--keep-bf16-pattern", + action="append", + default=[], + help=( + "Glob-style pattern matched against module names without trailing tensor " + "suffixes such as .weight or .bias." + ), + ) + parser.add_argument( + "--swap-weight-nibbles", + action=argparse.BooleanOptionalAction, + default=None, + help=( + "Whether the runtime should swap packed FP4 nibbles before padding. " + "Defaults to false for --pattern-preset flux1-nvfp4 and true otherwise." + ), + ) + parser.add_argument( + "--overwrite", + action="store_true", + help="Replace --output-dir if it already exists.", + ) + return parser.parse_args() + + +def main() -> None: + args = _parse_args() + stats = build_modelopt_nvfp4_transformer( + base_transformer_dir=args.base_transformer_dir, + modelopt_hf_dir=args.modelopt_hf_dir, + output_dir=args.output_dir, + pattern_preset=args.pattern_preset, + keep_bf16_patterns=args.keep_bf16_pattern, + swap_weight_nibbles=args.swap_weight_nibbles, + overwrite=args.overwrite, + ) + print(json.dumps(stats, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() From 37fc47c64559e54ab131ada6ea63be539f0277e6 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com> Date: Mon, 13 Apr 2026 08:01:54 +0800 Subject: [PATCH 03/52] diffusion: fix layerwise offload for ModelOpt quantized DiTs (#22594) --- docs/diffusion/quantization.md | 9 +- .../sglang-diffusion-modelopt-quant/SKILL.md | 14 +- .../runtime/loader/transformer_load_utils.py | 23 +-- .../runtime/utils/layerwise_offload.py | 135 +++++++++++--- .../test/unit/test_layerwise_offload.py | 166 ++++++++++++++++++ 5 files changed, 297 insertions(+), 50 deletions(-) create mode 100644 python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py diff --git a/docs/diffusion/quantization.md b/docs/diffusion/quantization.md index 96c872d4f56d..5a33a354412a 100644 --- a/docs/diffusion/quantization.md +++ b/docs/diffusion/quantization.md @@ -43,7 +43,7 @@ backend. | quant_family | checkpoint form | canonical CLI | supported models | extra dependency | platform / notes | |------------------|--------------------------------------------------------------------------------------------|------------------------------------------------------|--------------------------------------------------------------|---------------------------------------|-----------------------------------------------------------------------------------------------------------------------| | `fp8` | Quantized transformer component folder, or safetensors with `quantization_config` metadata | `--transformer-path` or `--transformer-weights-path` | ALL | None | Component-folder and single-file flows are both supported | -| `modelopt-fp8` | Converted ModelOpt FP8 transformer directory or repo with `config.json` | `--transformer-weights-path` | FLUX.2, Wan2.2 | None | Override config is read from the quantized transformer repo; `dit_cpu_offload` and `dit_layerwise_offload` are auto-disabled | +| `modelopt-fp8` | Converted ModelOpt FP8 transformer directory or repo with `config.json` | `--transformer-weights-path` | FLUX.2, Wan2.2 | None | Override config is read from the quantized transformer repo; `dit_layerwise_offload` is supported and `dit_cpu_offload` stays disabled | | `nvfp4-modelopt` | NVFP4 safetensors file, sharded directory, or repo providing transformer weights | `--transformer-weights-path` | FLUX.2 | `comfy-kitchen` optional on Blackwell | Blackwell can use a best-performance kit when available; otherwise SGLang falls back to the generic ModelOpt FP4 path | | `nunchaku-svdq` | Pre-quantized Nunchaku transformer weights, usually named `svdq-{int4\|fp4}_r{rank}-...` | `--transformer-weights-path` | Model-specific support such as Qwen-Image, FLUX, and Z-Image | `nunchaku` | SGLang can infer precision and rank from the filename and supports both `int4` and `nvfp4` | | `msmodelslim` | Pre-quantized msmodelslim transformer weights | `--model-path` | Wan2.2 family | None | Currently only compatible with the Ascend NPU family and supports both `w8a8` and `w4a4` | @@ -98,9 +98,10 @@ sglang generate \ - If the override repo or local directory contains its own `config.json`, SGLang reads the quantization config from that override instead of relying on the base model config. -- `dit_cpu_offload` and `dit_layerwise_offload` are automatically disabled for - ModelOpt FP8 checkpoints because the runtime expects the transformed FP8 - weights to remain GPU-resident in their column-major layout. +- `dit_layerwise_offload` is supported for ModelOpt FP8 checkpoints. +- `dit_cpu_offload` still stays disabled for ModelOpt FP8 checkpoints. +- The layerwise offload path now preserves the non-contiguous FP8 weight stride + expected by the runtime FP8 GEMM path. - To build the converted checkpoint yourself from a ModelOpt diffusers export, use `python -m sglang.multimodal_gen.tools.build_modelopt_fp8_transformer`. diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-modelopt-quant/SKILL.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-modelopt-quant/SKILL.md index 46dd8137cdd0..4e81ef4eb2f7 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-modelopt-quant/SKILL.md +++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-modelopt-quant/SKILL.md @@ -21,7 +21,7 @@ This skill owns the ModelOpt-to-SGLang bridge. It is not a generic kernel-tuning - Use ModelOpt's official `quantize.py` as the PTQ source of truth. - Keep the workflow generic. Put model-specific fallback logic in small isolated branches, not in the main conversion path. - Benchmark only when BF16 and quantized commands are identical except for the checkpoint override being tested. -- For diffusion FP8, pin `dit_cpu_offload=false` and `dit_layerwise_offload=false`. +- For diffusion FP8, keep `dit_cpu_offload=false`. `dit_layerwise_offload=true` is valid on the fixed path when you want lower DiT residency. - For multi-transformer pipelines, use per-component overrides when different components need different checkpoints. - When a branch is missing the validated helper tools, refresh `python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py`, `python/sglang/multimodal_gen/tools/build_modelopt_nvfp4_transformer.py`, and `python/sglang/multimodal_gen/tools/compare_diffusion_trajectory_similarity.py` instead of inventing one-off scripts elsewhere. - After validating a new ModelOpt quant path, update the ModelOpt support matrix in `docs/diffusion/quantization.md` before closing the task. @@ -52,7 +52,7 @@ This repo now contains: - diffusion-side ModelOpt FP8 linear loading - diffusion-side NVFP4 loading from ModelOpt exports - FLUX.2 packed-QKV detection that distinguishes packed NVFP4 checkpoints from standard diffusers exports -- automatic protection against incompatible FP8 offload modes +- automatic protection against incompatible FP8 CPU offload while keeping layerwise DiT offload available - FP8 transformer build: [`python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py`](../../../tools/build_modelopt_fp8_transformer.py) - trajectory similarity validation: @@ -279,18 +279,18 @@ Do not turn one validated model quirk into a generic rule unless another family Current diffusion ModelOpt FP8 support requires: - `dit_cpu_offload=false` -- `dit_layerwise_offload=false` +- `dit_layerwise_offload` may be enabled when you want lower DiT residency Reason: - the FP8 linear path depends on a CUTLASS-compatible weight layout after loading -- the offload and restore path does not preserve that layout -- in particular, layerwise offload can flatten and rebuild FP8 weights in a way that breaks the column-major requirement used by the FP8 GEMM path +- `dit_cpu_offload` is still treated conservatively +- the fixed layerwise offload path now preserves non-contiguous tensor strides instead of flattening and rebuilding FP8 weights into a contiguous layout Runtime behavior: -- SGLang currently force-disables these two flags when it detects `modelopt_fp8` -- benchmark commands should still pin them explicitly so the command line itself makes the comparison rule obvious +- SGLang still force-disables `dit_cpu_offload` when it detects `modelopt_fp8` +- benchmark commands should still pin offload flags explicitly so the command line itself makes the comparison rule obvious ## Claim Discipline diff --git a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py index 5ecffdfd0217..5885fa03b920 100644 --- a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py @@ -186,26 +186,13 @@ def _maybe_disable_incompatible_dit_offload_modes( if quant_name != "modelopt_fp8": return - disabled_args: list[str] = [] - if server_args.dit_cpu_offload: server_args.dit_cpu_offload = False - disabled_args.append("dit_cpu_offload") - - if server_args.dit_layerwise_offload: - server_args.dit_layerwise_offload = False - disabled_args.append("dit_layerwise_offload") - - if not disabled_args: - return - - logger.warning( - "ModelOpt FP8 diffusion checkpoints currently require the transformer " - "FP8 weights to stay GPU-resident in their column-major layout; " - "disabling %s for this run. Text encoder / VAE offload settings are " - "left unchanged.", - ", ".join(disabled_args), - ) + logger.warning( + "ModelOpt FP8 diffusion checkpoints currently keep dit_cpu_offload " + "disabled. Layerwise DiT offload stays enabled because the runtime " + "now preserves the restored FP8 tensor strides.", + ) def prepare(self) -> None: _ModelOptFp8OffloadAdapter._maybe_disable_incompatible_dit_offload_modes( diff --git a/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py b/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py index 05a3d676d8d5..29d601845246 100644 --- a/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py +++ b/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py @@ -56,6 +56,9 @@ def __init__( # layer_idx -> {dtype: consolidated_pinned_cpu_tensor} # stores the consolidated weight from a same layer, of same dtype self._consolidated_cpu_weights: Dict[int, Dict[torch.dtype, torch.Tensor]] = {} + # layer_idx -> {name: pinned_cpu_tensor_with_original_stride} + # stores tensors whose original non-contiguous stride/layout must be preserved + self._strided_cpu_weights: Dict[int, Dict[str, torch.Tensor]] = {} # layer_idx -> {name: {dtype, offset, numel, shape}} # stores the offset and numel of each weight from a same layer, of same dtype self._weight_metadata: Dict[int, Dict[str, Dict[str, Any]]] = {} @@ -88,6 +91,21 @@ def _get_shared_empty_tensor(self, dtype: torch.dtype) -> torch.Tensor: self._offload_placeholders[dtype] = placeholder return placeholder + @staticmethod + def _get_alignment_numel(dtype: torch.dtype, alignment_bytes: int = 32) -> int: + element_size = torch.empty((), dtype=dtype).element_size() + return max(1, alignment_bytes // element_size) + + @classmethod + def _align_numel_offset( + cls, offset: int, dtype: torch.dtype, alignment_bytes: int = 32 + ) -> int: + alignment_numel = cls._get_alignment_numel(dtype, alignment_bytes) + remainder = offset % alignment_numel + if remainder == 0: + return offset + return offset + alignment_numel - remainder + @torch.compiler.disable def _initialize(self) -> None: if not self.enabled: @@ -110,10 +128,49 @@ def _initialize(self) -> None: # 2. concat and offload (in pinned memory) for layer_idx, dtype_to_params in layer_groups.items(): self._consolidated_cpu_weights[layer_idx] = {} + self._strided_cpu_weights[layer_idx] = {} self._weight_metadata[layer_idx] = {} for dtype, weights in dtype_to_params.items(): - total_numel = sum(t.numel() for _, t in weights) + contiguous_weights: List[Tuple[str, torch.Tensor]] = [] + for name, weight in weights: + if weight.is_contiguous(): + contiguous_weights.append((name, weight)) + continue + + # Preserve non-contiguous layouts such as the transposed FP8 + # weight views expected by CUTLASS kernels. + cpu_tensor = torch.empty_strided( + size=weight.shape, + stride=weight.stride(), + dtype=dtype, + pin_memory=self.pin_cpu_memory, + ) + cpu_tensor.copy_(weight) + self._strided_cpu_weights[layer_idx][name] = cpu_tensor + self._weight_metadata[layer_idx][name] = { + "dtype": dtype, + "shape": weight.shape, + "stride": weight.stride(), + "preserve_strides": True, + } + weight.data = self._get_shared_empty_tensor(dtype) + + if not contiguous_weights: + continue + + current_offset = 0 + aligned_offsets: Dict[str, int] = {} + for name, weight in contiguous_weights: + # Some fused diffusion kernels require tensor base pointers to + # satisfy a 32-byte alignment contract. Reusing one flat buffer + # is still fine, but each logical tensor slice must start on an + # aligned offset inside that buffer. + current_offset = self._align_numel_offset(current_offset, dtype) + aligned_offsets[name] = current_offset + current_offset += weight.numel() + + total_numel = current_offset # create concatenated CPU buffer (in pinned memory) cpu_buffer = torch.empty( @@ -121,8 +178,8 @@ def _initialize(self) -> None: ) # offload weights to the buffer - current_offset = 0 - for name, weight in weights: + for name, weight in contiguous_weights: + current_offset = aligned_offsets[name] numel = weight.numel() cpu_buffer[current_offset : current_offset + numel].copy_( weight.flatten() @@ -132,6 +189,8 @@ def _initialize(self) -> None: "offset": current_offset, "numel": numel, "shape": weight.shape, + "stride": weight.stride(), + "preserve_strides": False, } weight.data = self._get_shared_empty_tensor(dtype) @@ -190,22 +249,39 @@ def prefetch_layer(self, layer_idx: int, non_blocking: bool = True) -> None: gpu_buffer.copy_(cpu_buffer, non_blocking=non_blocking) gpu_buffers[dtype] = gpu_buffer - # record the prefetch event of this layer + # restore model's weights by their metadata using the same copy stream + # so the recorded event covers both flat-buffer and stride-preserving copies. + for name, meta in self._weight_metadata[layer_idx].items(): + target = self.get_target_with_name(name) + if meta.get("preserve_strides", False): + # Recreate the original view layout instead of flatten+view. + # ModelOpt FP8 relies on a transposed runtime weight layout, + # so preserving stride is part of correctness, not just an + # optimization detail. + cpu_tensor = self._strided_cpu_weights[layer_idx][name] + gpu_tensor = torch.empty_strided( + size=meta["shape"], + stride=meta["stride"], + dtype=meta["dtype"], + device=self.device, + ) + gpu_tensor.copy_(cpu_tensor, non_blocking=non_blocking) + target.data = gpu_tensor + continue + + dtype = meta["dtype"] + gpu_buffer = gpu_buffers[dtype] + + # map the parameter's data to the correct slice of the GPU buffer + target.data = gpu_buffer[ + meta["offset"] : meta["offset"] + meta["numel"] + ].view(meta["shape"]) + + # record the prefetch event of this layer after all copies are enqueued event = torch.get_device_module().Event() event.record(self.copy_stream) self._prefetch_events[layer_idx] = event - # restore model's weights by their metadata using gpu buffer - for name, meta in self._weight_metadata[layer_idx].items(): - dtype = meta["dtype"] - gpu_buffer = gpu_buffers[dtype] - - # map the parameter's data to the correct slice of the GPU buffer - target = self.get_target_with_name(name) - target.data = gpu_buffer[ - meta["offset"] : meta["offset"] + meta["numel"] - ].view(meta["shape"]) - self._gpu_layers.add(layer_idx) @torch.compiler.disable @@ -266,6 +342,10 @@ def sync_layer_to_cpu(self, layer_idx: int) -> None: # Collect current GPU weights and write back to CPU buffer for name, meta in self._weight_metadata.get(layer_idx, {}).items(): target = self.get_target_with_name(name) + if meta.get("preserve_strides", False): + self._strided_cpu_weights[layer_idx][name].copy_(target.data.cpu()) + continue + gpu_weight = target.data.flatten().cpu() dtype = meta["dtype"] @@ -331,12 +411,17 @@ def update_cpu_weights( ) dtype = meta["dtype"] - offset = meta["offset"] - numel = meta["numel"] - cpu_buffer = self._consolidated_cpu_weights[layer_idx][dtype] - cpu_buffer[offset : offset + numel].copy_( - loaded_weight.to(dtype=dtype).flatten() - ) + if meta.get("preserve_strides", False): + self._strided_cpu_weights[layer_idx][name].copy_( + loaded_weight.to(dtype=dtype) + ) + else: + offset = meta["offset"] + numel = meta["numel"] + cpu_buffer = self._consolidated_cpu_weights[layer_idx][dtype] + cpu_buffer[offset : offset + numel].copy_( + loaded_weight.to(dtype=dtype).flatten() + ) # If this layer is currently on GPU, update the live parameter. if layer_idx in self._gpu_layers: @@ -358,6 +443,14 @@ def iter_cpu_weights(self): """ for layer_idx in sorted(self._weight_metadata): for name, meta in self._weight_metadata[layer_idx].items(): + if meta.get("preserve_strides", False): + # Some quantized weights rely on a non-contiguous layout. + # Yield the strided tensor directly instead of rebuilding it + # from the flat buffer, which would silently lose the + # original stride information. + yield name, self._strided_cpu_weights[layer_idx][name] + continue + dtype = meta["dtype"] offset = meta["offset"] numel = meta["numel"] diff --git a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py new file mode 100644 index 000000000000..73b256683a61 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py @@ -0,0 +1,166 @@ +from contextlib import nullcontext +from types import SimpleNamespace + +import torch + +from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import ( + ModelOptFp8Config, +) +from sglang.multimodal_gen.runtime.loader.transformer_load_utils import ( + _ModelOptFp8OffloadAdapter, +) +from sglang.multimodal_gen.runtime.utils import ( + layerwise_offload as layerwise_offload_mod, +) +from sglang.multimodal_gen.runtime.utils.layerwise_offload import ( + LayerwiseOffloadManager, +) + + +class _FakeStream: + def wait_stream(self, _stream) -> None: + return None + + def wait_event(self, _event) -> None: + return None + + +class _FakeEvent: + def record(self, _stream) -> None: + return None + + +class _FakeDeviceModule: + Stream = _FakeStream + Event = _FakeEvent + + @staticmethod + def is_available() -> bool: + return True + + @staticmethod + def current_device() -> int: + return 0 + + @staticmethod + def current_stream() -> _FakeStream: + return _FakeStream() + + @staticmethod + def stream(_stream): + return nullcontext() + + +class _DummyBlock(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + base = torch.arange(12, dtype=torch.float32).reshape(3, 4) + self.weight = torch.nn.Parameter(base.t()) + self.bias = torch.nn.Parameter(torch.arange(3, dtype=torch.float32)) + + +class _DummyModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.blocks = torch.nn.ModuleList([_DummyBlock()]) + + +def test_layerwise_offload_preserves_non_contiguous_stride(monkeypatch): + monkeypatch.setattr( + layerwise_offload_mod.torch, "get_device_module", lambda: _FakeDeviceModule + ) + monkeypatch.setattr(layerwise_offload_mod.current_platform, "device_type", "cpu") + + model = _DummyModel() + original_weight = model.blocks[0].weight.detach().clone() + original_stride = model.blocks[0].weight.stride() + assert not model.blocks[0].weight.is_contiguous() + + manager = LayerwiseOffloadManager( + model=model, + layers_attr_str="blocks", + num_layers=1, + enabled=True, + pin_cpu_memory=False, + prefetch_size=1, + ) + + meta = manager._weight_metadata[0]["blocks.0.weight"] + assert meta["preserve_strides"] is True + + restored_weight = model.blocks[0].weight.data + assert restored_weight.shape == original_weight.shape + assert restored_weight.stride() == original_stride + assert not restored_weight.is_contiguous() + assert torch.equal(restored_weight, original_weight) + + manager.release_layer(0) + manager.prefetch_layer(0, non_blocking=False) + + reloaded_weight = model.blocks[0].weight.data + assert reloaded_weight.stride() == original_stride + assert not reloaded_weight.is_contiguous() + assert torch.equal(reloaded_weight, original_weight) + + +def test_modelopt_fp8_adapter_keeps_layerwise_offload_enabled(): + server_args = SimpleNamespace( + dit_cpu_offload=True, + dit_layerwise_offload=True, + ) + quant_config = ModelOptFp8Config(is_checkpoint_fp8_serialized=True) + + _ModelOptFp8OffloadAdapter._maybe_disable_incompatible_dit_offload_modes( + server_args=server_args, + quant_config=quant_config, + ) + + assert server_args.dit_cpu_offload is False + assert server_args.dit_layerwise_offload is True + + +def test_layerwise_offload_aligns_contiguous_tensor_offsets(monkeypatch): + monkeypatch.setattr( + layerwise_offload_mod.torch, "get_device_module", lambda: _FakeDeviceModule + ) + monkeypatch.setattr(layerwise_offload_mod.current_platform, "device_type", "cpu") + + class _AlignedDummyBlock(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = torch.nn.Parameter( + torch.arange(9, dtype=torch.float32).reshape(3, 3) + ) + self.bias = torch.nn.Parameter(torch.arange(3, dtype=torch.float32)) + + class _AlignedDummyModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.blocks = torch.nn.ModuleList([_AlignedDummyBlock()]) + + model = _AlignedDummyModel() + original_weight = model.blocks[0].weight.detach().clone() + original_bias = model.blocks[0].bias.detach().clone() + + manager = LayerwiseOffloadManager( + model=model, + layers_attr_str="blocks", + num_layers=1, + enabled=True, + pin_cpu_memory=False, + prefetch_size=1, + ) + + weight_meta = manager._weight_metadata[0]["blocks.0.weight"] + bias_meta = manager._weight_metadata[0]["blocks.0.bias"] + assert weight_meta["preserve_strides"] is False + assert bias_meta["preserve_strides"] is False + assert weight_meta["offset"] == 0 + assert bias_meta["offset"] % 8 == 0 + + restored_weight = model.blocks[0].weight.data + restored_bias = model.blocks[0].bias.data + assert restored_weight.data_ptr() % 32 == 0 + assert restored_bias.data_ptr() % 32 == 0 + assert torch.equal(restored_weight, original_weight) + assert torch.equal(restored_bias, original_bias) From c1ab68b45e7c6b920793550ee1ac946e60b49233 Mon Sep 17 00:00:00 2001 From: ishandhanani <82981111+ishandhanani@users.noreply.github.com> Date: Sun, 12 Apr 2026 20:05:23 -0500 Subject: [PATCH 04/52] fix: streaming session race condition + some metrics (#21875) Co-authored-by: Claude Opus 4.6 (1M context) Co-authored-by: hnyls2002 Co-authored-by: Liangsheng Yin --- python/sglang/srt/managers/scheduler.py | 24 +- .../scheduler_runtime_checker_mixin.py | 9 + .../sglang/srt/managers/session_controller.py | 57 ++- .../managers/tokenizer_communicator_mixin.py | 10 +- .../sglang/srt/managers/tokenizer_manager.py | 12 +- python/sglang/srt/mem_cache/common.py | 48 ++- .../srt/mem_cache/session_aware_cache.py | 143 ++++++- .../srt/observability/metrics_collector.py | 28 ++ .../observability/scheduler_metrics_mixin.py | 3 + .../sessions/test_streaming_session.py | 300 +++++++++++++++ .../mem_cache/test_streaming_session_unit.py | 358 ++++++++++++++++++ 11 files changed, 966 insertions(+), 26 deletions(-) create mode 100644 test/registered/unit/mem_cache/test_streaming_session_unit.py diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index b4118777ec83..1f95448002dc 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1848,8 +1848,11 @@ def handle_generate_request( self.stream_output([req], req.return_logprob) return - elif session_id in self.session_controller: - # Session exists: create request from session + elif ( + session_id in self.session_controller + and not self.session_controller.get(session_id).close_on_finish + ): + # Session exists and is not closing: create request from session session = self.session_controller.get(session_id) req = session.create_req( recv_req, @@ -1866,7 +1869,13 @@ def handle_generate_request( return else: - # Session ID provided but session not found + # Session not found, or session is closing + if session_id in self.session_controller: + error_msg = ( + f"Invalid request: close was requested for session {session_id}" + ) + else: + error_msg = f"Invalid request: session id {session_id} does not exist" req = Req( recv_req.rid, recv_req.input_text, @@ -1875,9 +1884,7 @@ def handle_generate_request( vocab_size=self.model_config.vocab_size, ) req.tokenizer = self.tokenizer - req.set_finish_with_abort( - f"Invalid request: session id {session_id} does not exist" - ) + req.set_finish_with_abort(error_msg) self.init_req_max_new_tokens(req) self._add_request_to_queue(req) return @@ -3461,7 +3468,10 @@ def expert_distribution_handle(self, recv_req: ExpertDistributionReq): return ExpertDistributionReqOutput() def open_session(self, recv_req: OpenSessionReqInput): - return self.session_controller.open(recv_req) + output = self.session_controller.open(recv_req) + if self.pp_rank == 0 and self.tp_rank == 0 and self.attn_cp_rank == 0: + return output + return None def close_session(self, recv_req: CloseSessionReqInput): self.session_controller.close(recv_req) diff --git a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py index 94aa0da06a4a..189afae5b51b 100644 --- a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py +++ b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py @@ -117,6 +117,13 @@ def update_scheduler_stats(self, stats: SchedulerStats) -> None: class SchedulerRuntimeCheckerMixin: + def _alive_streaming_session_count(self: Scheduler) -> int: + return sum( + 1 + for session in self.session_controller.sessions.values() + if session.streaming + ) + def _session_held_tokens(self: Scheduler) -> int: if isinstance(self.tree_cache, SessionAwareCache): return self.tree_cache.session_held_tokens() @@ -451,6 +458,8 @@ def _maybe_log_idle_metrics(self: Scheduler): return self.get_pool_stats().update_scheduler_stats(self.stats) + self.stats.num_streaming_sessions = self._alive_streaming_session_count() + self.stats.streaming_session_held_tokens = self._session_held_tokens() priority_enabled = self.enable_priority_scheduling self.stats.num_running_reqs = QueueCount.from_reqs( diff --git a/python/sglang/srt/managers/session_controller.py b/python/sglang/srt/managers/session_controller.py index caf165c32165..9bb763fcebed 100644 --- a/python/sglang/srt/managers/session_controller.py +++ b/python/sglang/srt/managers/session_controller.py @@ -25,6 +25,7 @@ ) from sglang.srt.managers.schedule_batch import FINISH_ABORT, Req from sglang.srt.mem_cache.session_aware_cache import SessionAwareCache +from sglang.srt.utils.common import log_info_on_rank0 if TYPE_CHECKING: from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache @@ -92,6 +93,7 @@ def __init__( self.timeout = timeout self.last_active_time: float = time.monotonic() self.req_nodes: Dict[str, SessionReqNode] = {} + self.close_on_finish: bool = False def is_timed_out(self) -> bool: if self.timeout is None: @@ -275,6 +277,9 @@ def open(self, recv_req: OpenSessionReqInput) -> OpenSessionReqOutput: streaming=bool(recv_req.streaming), timeout=recv_req.timeout, ) + log_info_on_rank0( + logger, f"Session opened: {session_id} (active={len(self.sessions)})" + ) return OpenSessionReqOutput(session_id, True) def close(self, recv_req: CloseSessionReqInput): @@ -286,11 +291,31 @@ def close(self, recv_req: CloseSessionReqInput): def _close(self, session_id: str): session = self.sessions[session_id] + req = None + has_unfinished_request = False if session.streaming and session.req_nodes: assert len(session.req_nodes) == 1 req = next(iter(session.req_nodes.values())).req if not req.finished(): - req.session = None + has_unfinished_request = True + + if has_unfinished_request: + # An in-flight request is still decoding on this session's KV + # memory. Freeing now would corrupt the scheduler. Mark the + # session for deferred cleanup: the request keeps its session + # reference so cache_finished_req takes the streaming path, + # and we schedule release_session for after it completes. + session.close_on_finish = True + logger.info( + "Deferring session close for %s (unfinished request)", + session_id, + ) + return + + # No active request -- safe to release immediately. + if session.streaming and session.req_nodes: + req = next(iter(session.req_nodes.values())).req + req.session = None # Release multimodal features held by session requests. # Session reqs skip the normal mm cleanup path (scheduler and @@ -304,20 +329,46 @@ def _close(self, session_id: str): node.req.multimodal_inputs = None if isinstance(self.tree_cache, SessionAwareCache): - self.tree_cache.release_session(session_id) + self.tree_cache.release_session( + session_id, req if session.streaming else None + ) del self.sessions[session_id] + log_info_on_rank0( + logger, f"Session closed: {session_id} (active={len(self.sessions)})" + ) def maybe_reap(self, now: float, interval: float = 1.0): # reap sessions every second if now - self._last_reap_time > interval: self._last_reap_time = now + + # Finish deferred closes for sessions whose requests completed. + pending = [ + sid + for sid, session in self.sessions.items() + if session.close_on_finish and self._all_requests_finished(session) + ] + for sid in pending: + log_info_on_rank0( + logger, f"Deferred close ready for session {sid}, releasing." + ) + # Reset close_on_finish so _close proceeds with the release. + self.sessions[sid].close_on_finish = False + self._close(sid) + timed_out = [ sid for sid, session in self.sessions.items() if session.is_timed_out() ] for sid in timed_out: - logger.info(f"Session {sid} timed out, closing.") + log_info_on_rank0(logger, f"Session {sid} timed out, closing.") self._close(sid) + @staticmethod + def _all_requests_finished(session: "Session") -> bool: + if not session.req_nodes: + return True + return all(node.req.finished() for node in session.req_nodes.values()) + @staticmethod def adjust_mm_offsets(recv_req: TokenizedGenerateReqInput, req: Req, image_inputs): # For session requests, adjust mm_inputs offsets by the prefix length. diff --git a/python/sglang/srt/managers/tokenizer_communicator_mixin.py b/python/sglang/srt/managers/tokenizer_communicator_mixin.py index b035cac37e04..8ec7ee4f67cc 100644 --- a/python/sglang/srt/managers/tokenizer_communicator_mixin.py +++ b/python/sglang/srt/managers/tokenizer_communicator_mixin.py @@ -1090,12 +1090,14 @@ async def open_session( elif obj.session_id in self.session_futures: return None + future = asyncio.Future() + self.session_futures[obj.session_id] = future self.send_to_scheduler.send_pyobj(obj) - self.session_futures[obj.session_id] = asyncio.Future() - session_id = await self.session_futures[obj.session_id] - del self.session_futures[obj.session_id] - return session_id + try: + return await future + finally: + self.session_futures.pop(obj.session_id, None) async def close_session( self: TokenizerManager, diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 30e4ab9f9def..bfc5bef63aa5 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -2365,9 +2365,15 @@ def update_active_ranks(self, ranks: ActiveRanksOutput): self.send_to_scheduler.send_pyobj(ranks) def _handle_open_session_req_output(self, recv_obj): - self.session_futures[recv_obj.session_id].set_result( - recv_obj.session_id if recv_obj.success else None - ) + future = self.session_futures.get(recv_obj.session_id) + if future is None: + logger.warning( + "Open session response arrived after waiter cleanup: %s", + recv_obj.session_id, + ) + return + if not future.done(): + future.set_result(recv_obj.session_id if recv_obj.success else None) def _handle_update_weights_from_disk_req_output(self, recv_obj): if self.server_args.dp_size == 1: diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 5a759ed11bcd..391eea077780 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -9,6 +9,7 @@ from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, EvictParams from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool +from sglang.srt.mem_cache.session_aware_cache import SessionAwareCache from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import support_triton @@ -476,13 +477,52 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr req.mamba_pool_idx = None return + # Streaming sessions transfer req_pool ownership into SessionSlot objects. + # Trim any speculative tail before that transfer, otherwise later turns + # restore only the committed prefix and can strand unreachable KV pages. + # + # Aborted streaming-session requests (e.g. input too long) skip the + # streaming path entirely. match_prefix did not restore the slot's KV + # state, so the request has a fresh pool slot that should be freed by + # cache_finished_req below (which also sets req_pool_idx = None). + from sglang.srt.managers.schedule_batch import FINISH_ABORT + + is_streaming_session = ( + isinstance(tree_cache, SessionAwareCache) + and getattr(req, "session", None) is not None + and req.session.streaming + ) + is_aborted_streaming = is_streaming_session and isinstance( + getattr(req, "finished_reason", None), FINISH_ABORT + ) + if is_streaming_session and not is_aborted_streaming: + start_p, end_p = req.pop_overallocated_kv_cache() + page_size = get_global_server_args().page_size + if page_size > 1: + start_p = ceil_align(start_p, page_size) + if start_p < end_p: + indices_to_free = tree_cache.req_to_token_pool.req_to_token[ + req.req_pool_idx + ][start_p:end_p] + tree_cache.token_to_kv_pool_allocator.free(indices_to_free) + req.kv_allocated_len = req.kv_committed_len + tree_cache.cache_finished_req(req, is_insert=is_insert) - # FIXME: SessionAwareCache.cache_finished_req sets req_pool_idx = None to - # transfer KV ownership to the SessionSlot, so we skip the remaining - # cleanup (overalloc free + pool slot free). This means over-allocated - # tokens from speculative decoding are NOT freed between turns. + # SessionAwareCache.cache_finished_req sets req_pool_idx = None to transfer + # KV ownership to the SessionSlot, so the remaining cleanup is skipped. + # Streaming-session specific overalloc trimming must therefore happen + # before cache_finished_req above. if req.req_pool_idx is None: + if is_streaming_session: + # The request no longer owns any KV once SessionAwareCache either + # transfers it into the session slot or frees it on abort. Mark + # both bookkeeping flags so busy-time memory checks do not keep + # counting this finished request as uncached KV. + if not req.kv_committed_freed: + req.pop_committed_kv_cache() + if not req.kv_overallocated_freed: + req.pop_overallocated_kv_cache() return start_p, end_p = req.pop_overallocated_kv_cache() diff --git a/python/sglang/srt/mem_cache/session_aware_cache.py b/python/sglang/srt/mem_cache/session_aware_cache.py index 983395ec6c29..cd5f01c706de 100644 --- a/python/sglang/srt/mem_cache/session_aware_cache.py +++ b/python/sglang/srt/mem_cache/session_aware_cache.py @@ -1,5 +1,6 @@ from __future__ import annotations +import logging from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Dict, Optional @@ -22,6 +23,9 @@ from sglang.srt.managers.schedule_batch import Req +logger = logging.getLogger(__name__) + + class _VirtualNode: """Sentinel node for streaming session requests. @@ -187,6 +191,16 @@ def match_prefix(self, params: MatchPrefixParams) -> MatchResult: if slot is None or slot.req_pool_idx is None: return self.inner.match_prefix(params) + # If the request is destined for abort (e.g. input too long), + # do NOT restore the slot's KV state. set_finish_with_abort + # truncates origin_input_ids to [0], so alloc_for_extend would + # overwrite the slot's req_to_token row with a 1-token prefix, + # destroying the session's accumulated KV mapping. By skipping + # restore, the request gets a fresh pool slot from alloc_for_extend + # and the session slot remains untouched. + if req.to_finish is not None: + return self.inner.match_prefix(params) + slot.restore_to_req(req) # logprob_start_len is already forced to -1 for streaming sessions @@ -208,13 +222,56 @@ def cache_finished_req(self, req: Req, is_insert: bool = True, **kwargs): if not _is_streaming(req): return self.inner.cache_finished_req(req, is_insert=is_insert, **kwargs) + from sglang.srt.managers.schedule_batch import FINISH_ABORT + session_id = req.session.session_id slot = self.slots.get(session_id) is_first = slot is None + + # When an aborted streaming-session request was scheduled (e.g. + # input too long), match_prefix skipped restore_to_req so the + # request got a fresh pool slot from alloc_for_extend. Don't + # overwrite the session slot -- free the transient KV and pool slot. + if not is_first and isinstance(req.finished_reason, FINISH_ABORT): + if req.req_pool_idx is not None: + # Free all KV pages allocated for this aborted request. + end = req.kv_allocated_len + if end > 0: + kv_indices = self.req_to_token_pool.req_to_token[ + req.req_pool_idx, :end + ] + self.token_to_kv_pool_allocator.free(kv_indices) + self.req_to_token_pool.free_slots.append(req.req_pool_idx) + req.req_pool_idx = None + return + if is_first: slot = SessionSlot() self.slots[session_id] = slot + # If the session's KV is shrinking (e.g. client sent a shorter + # prompt after an abort), free the orphaned tail pages before + # save_from_req overwrites the slot's committed length. + # Never free tree-protected tokens — those are managed by the tree. + if ( + not is_first + and slot.is_holding_kv + and req.kv_committed_len < slot.kv_committed_len + ): + old_end = slot.kv_allocated_len + new_end = req.kv_committed_len + if self.page_size > 1: + new_end = ceil_align(new_end, self.page_size) + new_end = max(new_end, slot.cache_protected_len) + if new_end < old_end: + kv_indices = self.req_to_token_pool.req_to_token[ + slot.req_pool_idx, new_end:old_end + ] + self.token_to_kv_pool_allocator.free(kv_indices) + slot.cache_protected_len = min( + slot.cache_protected_len, req.kv_committed_len + ) + slot.save_from_req(req, is_first=is_first) def cache_unfinished_req(self, req: Req, **kwargs): @@ -251,23 +308,99 @@ def dec_lock_ref( # -- Session lifecycle -- - def release_session(self, session_id: str): + def _resolve_release_state( + self, slot: SessionSlot, req: Optional[Req] + ) -> tuple[int, Any]: + """Resolve the currently tree-owned prefix for a session slot. + + A long-lived session can outlive radix-tree splits caused by unrelated + traffic. In that case, the saved `last_node` may no longer represent the + full protected prefix even though the slot's req_to_token row still + contains tree-owned indices at the front. Re-match the current request + text, then intersect the returned tree indices with the slot's row so + release uses the prefix that is still actually backed by the tree. + """ + protected_len = slot.cache_protected_len + lock_node = slot.last_node + + # TODO: re-match logic disabled — match_prefix has side effects + # (splits) that disturb tree accounting. Directly using + # slot.last_node + cache_protected_len is safe after split analysis. + return protected_len, lock_node + + if ( + req is None + or not slot.is_holding_kv + or slot.req_pool_idx is None + or protected_len <= 0 + ): + return protected_len, lock_node + + from sglang.srt.mem_cache.radix_cache import RadixKey + + token_ids = (req.origin_input_ids + req.output_ids)[: slot.kv_committed_len] + if not token_ids: + return 0, None + + match = self.inner.match_prefix( + MatchPrefixParams( + key=RadixKey(token_ids=token_ids, extra_key=req.extra_key), + req=None, + ) + ) + if len(match.device_indices) == 0: + return 0, None + + max_protected_len = min(len(match.device_indices), protected_len) + row_indices = self.req_to_token_pool.req_to_token[ + slot.req_pool_idx, :max_protected_len + ].to(dtype=torch.int64) + match_indices = match.device_indices[:max_protected_len] + mismatches = (match_indices != row_indices).nonzero(as_tuple=False) + if mismatches.numel() == 0 and max_protected_len == len(match.device_indices): + common_len = max_protected_len + return common_len, match.last_device_node + + common_len = ( + int(mismatches[0].item()) if mismatches.numel() > 0 else max_protected_len + ) + if self.page_size > 1: + common_len = (common_len // self.page_size) * self.page_size + if common_len <= 0: + return 0, None + + rematch = self.inner.match_prefix( + MatchPrefixParams( + key=RadixKey(token_ids=token_ids[:common_len], extra_key=req.extra_key), + req=None, + ) + ) + return len(rematch.device_indices), rematch.last_device_node + + def release_session(self, session_id: str, req: Optional[Req] = None): """Release all KV resources held by a streaming session.""" slot = self.slots.pop(session_id, None) if slot is None: return + protected_len, lock_node = self._resolve_release_state(slot, req) + tokens_freed = ( + max(0, slot.kv_allocated_len - protected_len) if slot.is_holding_kv else 0 + ) + logger.info( + "Session KV released: %s (%d tokens freed)", session_id, tokens_freed + ) - if slot.last_node is not None: + if lock_node is not None: if slot.swa_uuid_for_lock is not None: self.inner.dec_lock_ref( - slot.last_node, + lock_node, DecLockRefParams(swa_uuid_for_lock=slot.swa_uuid_for_lock), ) else: - self.inner.dec_lock_ref(slot.last_node) + self.inner.dec_lock_ref(lock_node) if slot.is_holding_kv: - start = slot.cache_protected_len + start = protected_len end = slot.kv_allocated_len if start < end: kv_indices = self.req_to_token_pool.req_to_token[ diff --git a/python/sglang/srt/observability/metrics_collector.py b/python/sglang/srt/observability/metrics_collector.py index 9a81ee78f182..3c34302161dd 100644 --- a/python/sglang/srt/observability/metrics_collector.py +++ b/python/sglang/srt/observability/metrics_collector.py @@ -132,6 +132,10 @@ class SchedulerStats: hicache_host_used_tokens: int = 0 hicache_host_total_tokens: int = 0 + # Streaming session metrics + num_streaming_sessions: int = 0 + streaming_session_held_tokens: int = 0 + # Routing key metrics num_unique_running_routing_keys: int = 0 routing_key_running_req_counts: List[int] = field(default_factory=list) @@ -176,6 +180,7 @@ def __init__( labels: Dict[str, str], enable_lora: bool = False, enable_hierarchical_cache: bool = False, + enable_streaming_session: bool = False, server_args: Optional["ServerArgs"] = None, ) -> None: # We need to import prometheus_client after setting the env variable `PROMETHEUS_MULTIPROC_DIR` @@ -184,6 +189,7 @@ def __init__( self.labels = labels self.enable_lora = enable_lora self.enable_hierarchical_cache = enable_hierarchical_cache + self.enable_streaming_session = enable_streaming_session self.last_log_time = time.perf_counter() self._known_priorities: Set[int] = set() @@ -654,6 +660,21 @@ def __init__( multiprocess_mode="mostrecent", ) + # Streaming session metrics (only created when streaming sessions are enabled) + if self.enable_streaming_session: + self.num_streaming_sessions = Gauge( + name="sglang:num_streaming_sessions", + documentation="The number of active streaming sessions.", + labelnames=labels.keys(), + multiprocess_mode="mostrecent", + ) + self.streaming_session_held_tokens = Gauge( + name="sglang:streaming_session_held_tokens", + documentation="The number of KV tokens currently held by streaming session slots.", + labelnames=labels.keys(), + multiprocess_mode="mostrecent", + ) + self.num_unique_running_routing_keys = Gauge( name="sglang:num_unique_running_routing_keys", documentation="Number of unique routing keys in running batch.", @@ -1049,6 +1070,13 @@ def log_stats(self, stats: SchedulerStats) -> None: self.hicache_host_total_tokens, stats.hicache_host_total_tokens ) + # Streaming session metrics (only logged if streaming sessions are enabled) + if self.enable_streaming_session: + self._log_gauge(self.num_streaming_sessions, stats.num_streaming_sessions) + self._log_gauge( + self.streaming_session_held_tokens, stats.streaming_session_held_tokens + ) + self._log_gauge( self.num_unique_running_routing_keys, stats.num_unique_running_routing_keys ) diff --git a/python/sglang/srt/observability/scheduler_metrics_mixin.py b/python/sglang/srt/observability/scheduler_metrics_mixin.py index 593d6dd58680..796c312c364b 100644 --- a/python/sglang/srt/observability/scheduler_metrics_mixin.py +++ b/python/sglang/srt/observability/scheduler_metrics_mixin.py @@ -148,6 +148,7 @@ def init_metrics( labels=labels, enable_lora=self.enable_lora, enable_hierarchical_cache=self.enable_hierarchical_cache, + enable_streaming_session=self.server_args.enable_streaming_session, server_args=self.server_args, ) self.enable_mfu_metrics = bool(self.server_args.enable_mfu_metrics) @@ -602,6 +603,8 @@ def report_decode_stats( self.stats.cache_hit_rate = cache_hit_rate self.stats.max_total_num_tokens = self.max_total_num_tokens + self.stats.num_streaming_sessions = self._alive_streaming_session_count() + self.stats.streaming_session_held_tokens = self._session_held_tokens() # Speculative decoding self.stats.spec_accept_rate = spec_accept_rate diff --git a/test/registered/sessions/test_streaming_session.py b/test/registered/sessions/test_streaming_session.py index 89c1a1fd8814..59bb76981971 100644 --- a/test/registered/sessions/test_streaming_session.py +++ b/test/registered/sessions/test_streaming_session.py @@ -10,8 +10,12 @@ """ import asyncio +import json +import os +import tempfile import time import unittest +from pathlib import Path from typing import Any, Optional import aiohttp @@ -66,6 +70,21 @@ "We promptly judged antique ivory buckles for the next prize. " ) * 20 +# --------------------------------------------------------------------------- +# Abort-heavy chunked prefill leak repro constants +# --------------------------------------------------------------------------- + +ABORT_REPRO_CONTEXT_LEN = 512 +ABORT_REPRO_PAGE_SIZE = 16 +ABORT_REPRO_GEN_LEN = 8 +ABORT_REPRO_SESSIONS = 4 +ABORT_REPRO_WARMUP_TURNS = 2 +ABORT_REPRO_ROUNDS = 8 +ABORT_REPRO_STREAM_TOKENS = 150 +ABORT_REPRO_ABORT_TOKENS = 320 +ABORT_REPRO_NON_STREAMING_TOKENS = 96 +ABORT_REPRO_CHUNKED_PREFILL_SIZE = 128 + # --------------------------------------------------------------------------- # Logprob leak helpers @@ -206,6 +225,198 @@ async def _leak_run_all(base_url: str, tokenizer: Any) -> None: assert resp.status == 200 +def _make_token_sized_ids( + tokenizer: Any, prefix: str, min_tokens: int, max_tokens: Optional[int] = None +) -> list[int]: + text = prefix + chunk = " pack quartz wizard sphinx zebra fox " * 16 + token_ids = tokenizer.encode(text) + while len(token_ids) < min_tokens: + text += chunk + token_ids = tokenizer.encode(text) + if max_tokens is not None: + token_ids = token_ids[:max_tokens] + return token_ids + + +async def _abort_repro_generate( + base_url: str, + session: aiohttp.ClientSession, + input_ids: list[int], + max_new_tokens: int, + session_params: Optional[dict[str, Any]] = None, + expect_abort: bool = False, +) -> Optional[dict[str, Any]]: + payload: dict[str, Any] = { + "input_ids": input_ids, + "sampling_params": { + "temperature": 0, + "max_new_tokens": max_new_tokens, + "no_stop_trim": True, + "skip_special_tokens": False, + }, + } + if session_params: + payload["session_params"] = session_params + + async with session.post(base_url + "/generate", json=payload) as resp: + text = await resp.text() + if expect_abort: + if resp.status == 200: + data = json.loads(text) + finish_reason = data.get("meta_info", {}).get("finish_reason", {}) + assert finish_reason.get("type") == "abort", text + assert "maximum allowed length" in finish_reason.get( + "message", "" + ), text + return data + assert resp.status == 400, text + assert "maximum allowed length" in text, text + return None + + assert resp.status == 200, text + data = json.loads(text) + finish_reason = data.get("meta_info", {}).get("finish_reason", {}) + assert finish_reason.get("type") != "abort", text + return data + + +def _read_tail(path: str, num_lines: int = 120) -> str: + if not path or not os.path.exists(path): + return "" + lines = Path(path).read_text(errors="replace").splitlines() + return "\n".join(lines[-num_lines:]) + + +async def _abort_repro_run_all(base_url: str, tokenizer: Any) -> None: + timeout = aiohttp.ClientTimeout(total=300) + async with aiohttp.ClientSession(timeout=timeout) as http: + session_ids = [] + for _ in range(ABORT_REPRO_SESSIONS): + async with http.post( + base_url + "/open_session", + json={"capacity_of_str_len": 50000, "streaming": True}, + ) as resp: + assert resp.status == 200, await resp.text() + session_ids.append(await resp.json()) + + try: + for warmup_turn in range(ABORT_REPRO_WARMUP_TURNS): + warmup_tasks = [] + for session_idx, session_id in enumerate(session_ids): + input_ids = _make_token_sized_ids( + tokenizer, + prefix=f"[warmup={warmup_turn} session={session_idx}]", + min_tokens=ABORT_REPRO_STREAM_TOKENS, + max_tokens=ABORT_REPRO_STREAM_TOKENS + 8, + ) + warmup_tasks.append( + _abort_repro_generate( + base_url, + http, + input_ids, + ABORT_REPRO_GEN_LEN, + session_params={"id": session_id, "rid": None}, + ) + ) + await asyncio.gather(*warmup_tasks) + + for round_idx in range(ABORT_REPRO_ROUNDS): + mixed_tasks = [] + for session_idx, session_id in enumerate(session_ids): + input_ids = _make_token_sized_ids( + tokenizer, + prefix=f"[round={round_idx} ok session={session_idx}]", + min_tokens=ABORT_REPRO_STREAM_TOKENS, + max_tokens=ABORT_REPRO_STREAM_TOKENS + 8, + ) + mixed_tasks.append( + _abort_repro_generate( + base_url, + http, + input_ids, + ABORT_REPRO_GEN_LEN, + session_params={"id": session_id, "rid": None}, + ) + ) + + for ns_idx in range(2): + input_ids = _make_token_sized_ids( + tokenizer, + prefix=f"[round={round_idx} ns={ns_idx}]", + min_tokens=ABORT_REPRO_NON_STREAMING_TOKENS, + max_tokens=ABORT_REPRO_NON_STREAMING_TOKENS + 8, + ) + mixed_tasks.append( + _abort_repro_generate( + base_url, + http, + input_ids, + ABORT_REPRO_GEN_LEN, + ) + ) + await asyncio.gather(*mixed_tasks) + + abort_tasks = [] + for session_idx, session_id in enumerate(session_ids): + input_ids = _make_token_sized_ids( + tokenizer, + prefix=f"[round={round_idx} abort session={session_idx}]", + min_tokens=ABORT_REPRO_ABORT_TOKENS, + ) + abort_tasks.append( + _abort_repro_generate( + base_url, + http, + input_ids, + ABORT_REPRO_GEN_LEN, + session_params={"id": session_id, "rid": None}, + expect_abort=True, + ) + ) + await asyncio.gather(*abort_tasks) + + recovery_tasks = [] + for session_idx, session_id in enumerate(session_ids): + input_ids = _make_token_sized_ids( + tokenizer, + prefix=f"[round={round_idx} recover session={session_idx}]", + min_tokens=ABORT_REPRO_NON_STREAMING_TOKENS, + max_tokens=ABORT_REPRO_NON_STREAMING_TOKENS + 8, + ) + recovery_tasks.append( + _abort_repro_generate( + base_url, + http, + input_ids, + ABORT_REPRO_GEN_LEN, + session_params={"id": session_id, "rid": None}, + ) + ) + recovery_results = await asyncio.gather(*recovery_tasks) + for result in recovery_results: + assert result is not None + assert result["meta_info"]["cached_tokens"] > 0, result + + health = requests.get(base_url + "/health", timeout=10) + if health.status_code != 200: + raise RuntimeError( + f"server unhealthy after round={round_idx}: " + f"{health.status_code} {health.text}" + ) + finally: + for session_id in session_ids: + async with http.post( + base_url + "/close_session", json={"session_id": session_id} + ) as resp: + assert resp.status == 200, await resp.text() + + +# =================================================================== +# Test class +# =================================================================== + + class TestStreamingSession(CustomTestCase): @classmethod def setUpClass(cls): @@ -467,5 +678,94 @@ def tearDownClass(cls): kill_process_tree(cls.process.pid) +class TestStreamingSessionAbortLeakRepro(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST + cls.base_url = DEFAULT_URL_FOR_TEST + cls.stdout = tempfile.NamedTemporaryFile( + prefix="streaming-session-abort-repro.", + suffix=".stdout.log", + delete=False, + mode="w+", + encoding="utf-8", + ) + cls.stderr = tempfile.NamedTemporaryFile( + prefix="streaming-session-abort-repro.", + suffix=".stderr.log", + delete=False, + mode="w+", + encoding="utf-8", + ) + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--enable-streaming-session", + "--chunked-prefill-size", + str(ABORT_REPRO_CHUNKED_PREFILL_SIZE), + "--context-length", + str(ABORT_REPRO_CONTEXT_LEN), + "--page-size", + str(ABORT_REPRO_PAGE_SIZE), + "--max-running-requests", + "32", + "--log-level", + "info", + ], + return_stdout_stderr=(cls.stdout, cls.stderr), + ) + cls.tokenizer = get_tokenizer(cls.model) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + for handle in (cls.stdout, cls.stderr): + path = handle.name + handle.close() + if os.path.exists(path): + os.remove(path) + + def test_abort_heavy_chunked_prefill_does_not_leak(self) -> None: + requests.post(self.base_url + "/flush_cache") + + asyncio.run(_abort_repro_run_all(self.base_url, self.tokenizer)) + + for i in range(3): + ids = self.tokenizer.encode(f"Post-session cleanup request {i}.") + response = requests.post( + self.base_url + "/generate", + json={ + "input_ids": ids, + "sampling_params": {"temperature": 0, "max_new_tokens": 4}, + }, + timeout=30, + ) + self.assertEqual(response.status_code, 200, response.text) + + time.sleep(5) + self.assertIsNone( + self.process.poll(), + "Server crashed during abort-heavy streaming session repro.\n" + f"---- stderr tail ----\n{_read_tail(self.stderr.name)}", + ) + + health = requests.get(self.base_url + "/health", timeout=10) + self.assertEqual( + health.status_code, + 200, + "Server unhealthy after abort-heavy streaming session cleanup.\n" + f"---- stderr tail ----\n{_read_tail(self.stderr.name)}", + ) + + stderr_tail = _read_tail(self.stderr.name) + self.assertNotIn( + "token_to_kv_pool_allocator memory leak detected", + stderr_tail, + stderr_tail, + ) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_streaming_session_unit.py b/test/registered/unit/mem_cache/test_streaming_session_unit.py new file mode 100644 index 000000000000..759cff86221a --- /dev/null +++ b/test/registered/unit/mem_cache/test_streaming_session_unit.py @@ -0,0 +1,358 @@ +from types import SimpleNamespace + +import torch + +from sglang.srt.managers.schedule_batch import FINISH_ABORT +from sglang.srt.mem_cache.base_prefix_cache import MatchResult +from sglang.srt.mem_cache.common import release_kv_cache +from sglang.srt.mem_cache.session_aware_cache import SessionAwareCache, SessionSlot +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=8, suite="stage-a-test-cpu") + + +class _FakeAllocator: + def __init__(self): + self.freed = [] + + def free(self, free_index: torch.Tensor): + self.freed.append(free_index.clone()) + + +class _FakeInnerCache: + def __init__(self, req_to_token_pool, allocator, page_size, match_results=None): + self.req_to_token_pool = req_to_token_pool + self.token_to_kv_pool_allocator = allocator + self.page_size = page_size + self.match_results = list(match_results or []) + self.dec_lock_ref_calls = [] + + def cache_finished_req(self, *args, **kwargs): + raise AssertionError("Streaming requests should not delegate to inner cache") + + def match_prefix(self, *args, **kwargs): + if not self.match_results: + raise AssertionError("Unexpected match_prefix call") + return self.match_results.pop(0) + + def dec_lock_ref(self, node, *args, **kwargs): + self.dec_lock_ref_calls.append(node) + + def supports_mamba(self): + return False + + def sanity_check(self): + return None + + +class _FakeReq: + def __init__( + self, session_id: str, req_pool_idx: int, committed: int, allocated: int + ): + self.session = SimpleNamespace(session_id=session_id, streaming=True) + self.req_pool_idx = req_pool_idx + self.kv_committed_len = committed + self.kv_allocated_len = allocated + self.kv_committed_freed = False + self.kv_overallocated_freed = False + self.origin_input_ids = list(range(committed)) + self.output_ids = [] + self.extra_key = None + self.swa_evicted_seqlen = 0 + self.last_node = None + self.cache_protected_len = 0 + self.swa_uuid_for_lock = None + self.mamba_pool_idx = None + self.mamba_ping_pong_track_buffer = None + self.mamba_next_track_idx = None + self.mamba_last_track_seqlen = None + self.mamba_branching_seqlen = None + self.pop_overallocated_calls = 0 + self.to_finish = None + self.finished_reason = None + + def pop_committed_kv_cache(self): + assert not self.kv_committed_freed + self.kv_committed_freed = True + return self.kv_committed_len + + def pop_overallocated_kv_cache(self): + assert not self.kv_overallocated_freed + self.pop_overallocated_calls += 1 + self.kv_overallocated_freed = True + return self.kv_committed_len, self.kv_allocated_len + + +def test_streaming_release_kv_cache_trims_overallocated_tail(monkeypatch): + page_size = 16 + req_to_token = torch.arange(128, dtype=torch.int32).reshape(1, 128) + req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[]) + allocator = _FakeAllocator() + tree_cache = SessionAwareCache( + _FakeInnerCache(req_to_token_pool, allocator, page_size) + ) + req = _FakeReq("session-a", req_pool_idx=0, committed=17, allocated=40) + + monkeypatch.setattr( + "sglang.srt.mem_cache.common.get_global_server_args", + lambda: SimpleNamespace(page_size=page_size, speculative_algorithm="eagle"), + ) + + release_kv_cache(req, tree_cache) + + slot = tree_cache.slots["session-a"] + assert req.pop_overallocated_calls == 1 + assert req.kv_committed_freed is True + assert req.kv_overallocated_freed is True + assert req.req_pool_idx is None + assert slot.kv_committed_len == 17 + assert slot.kv_allocated_len == 17 + assert len(allocator.freed) == 1 + assert allocator.freed[0].tolist() == list(range(32, 40)) + + +def test_release_session_recomputes_current_tree_owned_prefix(): + page_size = 16 + req_to_token = torch.arange(128, dtype=torch.int32).reshape(1, 128) + req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[]) + allocator = _FakeAllocator() + + full_match = MatchResult( + device_indices=torch.tensor(list(range(16)) + list(range(64, 96))), + last_device_node="stale-expanded", + last_host_node="stale-expanded", + ) + protected_match = MatchResult( + device_indices=torch.tensor(list(range(16))), + last_device_node="current-protected", + last_host_node="current-protected", + ) + inner = _FakeInnerCache( + req_to_token_pool, + allocator, + page_size, + match_results=[full_match, protected_match], + ) + tree_cache = SessionAwareCache(inner) + + tree_cache.slots["session-a"] = SessionSlot( + req_pool_idx=0, + kv_committed_len=48, + kv_allocated_len=48, + last_node="outdated-node", + cache_protected_len=32, + ) + req = _FakeReq("session-a", req_pool_idx=0, committed=48, allocated=48) + + tree_cache.release_session("session-a", req) + + assert inner.dec_lock_ref_calls == ["current-protected"] + assert req_to_token_pool.free_slots == [0] + assert len(allocator.freed) == 1 + assert allocator.freed[0].tolist() == list(range(16, 48)) + + +def test_release_session_never_grows_tree_owned_prefix(): + page_size = 16 + req_to_token = torch.arange(128, dtype=torch.int32).reshape(1, 128) + req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[]) + allocator = _FakeAllocator() + + overmatched = MatchResult( + device_indices=torch.tensor(list(range(48))), + last_device_node="overmatched-node", + last_host_node="overmatched-node", + ) + capped_match = MatchResult( + device_indices=torch.tensor(list(range(16))), + last_device_node="original-lock-node", + last_host_node="original-lock-node", + ) + inner = _FakeInnerCache( + req_to_token_pool, + allocator, + page_size, + match_results=[overmatched, capped_match], + ) + tree_cache = SessionAwareCache(inner) + + tree_cache.slots["session-a"] = SessionSlot( + req_pool_idx=0, + kv_committed_len=48, + kv_allocated_len=48, + last_node="outdated-node", + cache_protected_len=16, + ) + req = _FakeReq("session-a", req_pool_idx=0, committed=48, allocated=48) + + tree_cache.release_session("session-a", req) + + assert inner.dec_lock_ref_calls == ["original-lock-node"] + assert req_to_token_pool.free_slots == [0] + assert len(allocator.freed) == 1 + assert allocator.freed[0].tolist() == list(range(16, 48)) + + +def test_match_prefix_abort_does_not_restore_live_session_slot(): + req_to_token = torch.arange(256, dtype=torch.int32).reshape(2, 128) + req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[]) + allocator = _FakeAllocator() + inner = _FakeInnerCache( + req_to_token_pool, + allocator, + page_size=16, + match_results=[ + MatchResult( + device_indices=torch.tensor([], dtype=torch.int64), + last_device_node=None, + last_host_node=None, + ) + ], + ) + tree_cache = SessionAwareCache(inner) + tree_cache.slots["session-a"] = SessionSlot( + req_pool_idx=0, + kv_committed_len=48, + kv_allocated_len=48, + cache_protected_len=16, + ) + + req = _FakeReq("session-a", req_pool_idx=1, committed=1, allocated=1) + req.to_finish = FINISH_ABORT("too long") + + result = tree_cache.match_prefix( + SimpleNamespace( + req=req, + key=SimpleNamespace(token_ids=list(range(64))), + ) + ) + + slot = tree_cache.slots["session-a"] + assert req.req_pool_idx == 1 + assert req.kv_committed_len == 1 + assert req.kv_allocated_len == 1 + assert slot.req_pool_idx == 0 + assert slot.kv_committed_len == 48 + assert slot.kv_allocated_len == 48 + assert len(result.device_indices) == 0 + + +def test_aborted_streaming_turn_preserves_slot_and_accounting(monkeypatch): + page_size = 16 + req_to_token = torch.arange(256, dtype=torch.int32).reshape(2, 128) + req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[]) + allocator = _FakeAllocator() + tree_cache = SessionAwareCache( + _FakeInnerCache(req_to_token_pool, allocator, page_size) + ) + tree_cache.slots["session-a"] = SessionSlot( + req_pool_idx=0, + kv_committed_len=48, + kv_allocated_len=48, + cache_protected_len=16, + swa_evicted_seqlen=8, + last_node="lock-node", + ) + + req = _FakeReq("session-a", req_pool_idx=1, committed=5, allocated=23) + req.finished_reason = FINISH_ABORT("too long") + + monkeypatch.setattr( + "sglang.srt.mem_cache.common.get_global_server_args", + lambda: SimpleNamespace(page_size=page_size, speculative_algorithm="eagle"), + ) + + release_kv_cache(req, tree_cache) + + slot = tree_cache.slots["session-a"] + assert slot.req_pool_idx == 0 + assert slot.kv_committed_len == 48 + assert slot.kv_allocated_len == 48 + assert req.kv_committed_freed is True + assert req.kv_overallocated_freed is True + assert req.req_pool_idx is None + assert req.pop_overallocated_calls == 1 + assert tree_cache.session_held_tokens() == 32 + assert tree_cache.session_held_full_tokens() == 32 + assert tree_cache.session_held_swa_tokens() == 32 + assert tree_cache.session_held_req_count() == 1 + assert req_to_token_pool.free_slots == [1] + assert len(allocator.freed) == 1 + assert allocator.freed[0].tolist() == list(range(128, 151)) + + tree_cache.release_session("session-a") + + assert tree_cache.session_held_tokens() == 0 + assert tree_cache.session_held_swa_tokens() == 0 + assert tree_cache.session_held_req_count() == 0 + assert req_to_token_pool.free_slots == [1, 0] + assert len(allocator.freed) == 2 + assert allocator.freed[1].tolist() == list(range(16, 48)) + + +def test_session_shrink_frees_orphaned_tail(): + """When a session's KV shrinks (client retried with shorter prompt), + the orphaned tail pages must be freed before save_from_req overwrites + the slot.""" + page_size = 16 + pool_size = 256 + req_to_token = torch.arange(pool_size, dtype=torch.int32).reshape(1, pool_size) + req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[]) + allocator = _FakeAllocator() + inner = _FakeInnerCache(req_to_token_pool, allocator, page_size) + tree_cache = SessionAwareCache(inner) + + # Session slot has 128 tokens committed + tree_cache.slots["session-a"] = SessionSlot( + req_pool_idx=0, + kv_committed_len=128, + kv_allocated_len=128, + last_node="lock-node", + cache_protected_len=16, + ) + + # New request finished with only 48 tokens (client truncated) + req = _FakeReq("session-a", req_pool_idx=0, committed=48, allocated=48) + + tree_cache.cache_finished_req(req) + + slot = tree_cache.slots["session-a"] + # Slot should now reflect the shrunk state + assert slot.kv_committed_len == 48 + assert slot.kv_allocated_len == 48 + # The tail [48:128] should have been freed (page-aligned: [48:128]) + assert len(allocator.freed) == 1 + assert allocator.freed[0].tolist() == list(range(48, 128)) + + +def test_session_shrink_page_aligns_free_start(): + """The shrink free should page-align the start to avoid freeing + tokens that are still part of the new committed prefix.""" + page_size = 16 + pool_size = 256 + req_to_token = torch.arange(pool_size, dtype=torch.int32).reshape(1, pool_size) + req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[]) + allocator = _FakeAllocator() + inner = _FakeInnerCache(req_to_token_pool, allocator, page_size) + tree_cache = SessionAwareCache(inner) + + # Session slot has 128 tokens + tree_cache.slots["session-a"] = SessionSlot( + req_pool_idx=0, + kv_committed_len=128, + kv_allocated_len=128, + last_node="lock-node", + cache_protected_len=16, + ) + + # New request committed 50 tokens (not page-aligned) + req = _FakeReq("session-a", req_pool_idx=0, committed=50, allocated=50) + + tree_cache.cache_finished_req(req) + + slot = tree_cache.slots["session-a"] + assert slot.kv_committed_len == 50 + # Free start should be ceil_align(50, 16) = 64, not 50 + # So freed range is [64:128] + assert len(allocator.freed) == 1 + assert allocator.freed[0].tolist() == list(range(64, 128)) From 28e40d873ca1289f402b595e58e4770cd4a8dd90 Mon Sep 17 00:00:00 2001 From: Lawrence Wu Date: Sun, 12 Apr 2026 18:07:51 -0700 Subject: [PATCH 05/52] fix(PD): respect pause_generation in disagg event loops (#20908) --- python/sglang/srt/disaggregation/decode.py | 6 +- python/sglang/srt/disaggregation/prefill.py | 6 ++ .../test_disaggregation_basic.py | 97 +++++++++++++++++++ 3 files changed, 107 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index e3ee0ae1c3ae..875900ad8579 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -1178,8 +1178,9 @@ def event_loop_normal_disagg_decode(self: Scheduler): # Receive requests recv_reqs = self.recv_requests() self.process_input_requests(recv_reqs) - # polling and allocating kv cache self.process_decode_queue() + if self._engine_paused: + continue # Get the next batch to run batch = self.get_next_disagg_decode_batch_to_run() @@ -1205,8 +1206,9 @@ def event_loop_overlap_disagg_decode(self: Scheduler): # Receive requests recv_reqs = self.recv_requests() self.process_input_requests(recv_reqs) - # polling and allocating kv cache self.process_decode_queue() + if self._engine_paused: + continue # Get the next batch to run batch = self.get_next_disagg_decode_batch_to_run() diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 0e0a653fdaa0..d9016d0f94d3 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -397,6 +397,9 @@ def event_loop_normal_disagg_prefill(self: Scheduler) -> None: self.waiting_queue.extend( self.disagg_prefill_bootstrap_queue.pop_bootstrapped() ) + if self._engine_paused: + self.process_disagg_prefill_inflight_queue() + continue # Get the next batch to run batch = self.get_next_disagg_prefill_batch_to_run() @@ -428,6 +431,9 @@ def event_loop_overlap_disagg_prefill(self: Scheduler) -> None: self.waiting_queue.extend( self.disagg_prefill_bootstrap_queue.pop_bootstrapped() ) + if self._engine_paused: + self.process_disagg_prefill_inflight_queue() + continue # Get the next batch to run batch = self.get_next_disagg_prefill_batch_to_run() diff --git a/test/registered/disaggregation/test_disaggregation_basic.py b/test/registered/disaggregation/test_disaggregation_basic.py index fd63195b35fa..deeda4a8a92a 100644 --- a/test/registered/disaggregation/test_disaggregation_basic.py +++ b/test/registered/disaggregation/test_disaggregation_basic.py @@ -1,7 +1,9 @@ import asyncio import json import os +import time import unittest +from concurrent.futures import ThreadPoolExecutor, as_completed from types import SimpleNamespace import aiohttp @@ -97,6 +99,101 @@ def test_structured_output(self): # ensure the output is a valid JSON json.loads(output) + def test_pause_resume_in_place(self): + """Send requests, pause mid-generation, verify no progress during pause, resume.""" + NUM_REQUESTS = 32 + MAX_NEW_TOKENS = 512 + REQUEST_TIMEOUT = 180 + PAUSE_DURATION = 5 + + def _generate(prompt_id): + return requests.post( + self.lb_url + "/generate", + json={ + "text": f"Question {prompt_id}: Write a short essay about the number {prompt_id}.", + "sampling_params": { + "temperature": 0.8, + "max_new_tokens": MAX_NEW_TOKENS, + }, + }, + timeout=REQUEST_TIMEOUT, + ) + + with ThreadPoolExecutor(max_workers=NUM_REQUESTS) as executor: + futures = {executor.submit(_generate, i): i for i in range(NUM_REQUESTS)} + + time.sleep(1) + + requests.post( + self.prefill_url + "/pause_generation", + json={"mode": "in_place"}, + timeout=30, + ).raise_for_status() + requests.post( + self.decode_url + "/pause_generation", + json={"mode": "in_place"}, + timeout=30, + ).raise_for_status() + + time.sleep(0.5) + done_before = sum(1 for f in futures if f.done()) + + time.sleep(PAUSE_DURATION) + done_after = sum(1 for f in futures if f.done()) + + self.assertLess( + done_before, + NUM_REQUESTS, + "All requests completed before pause took effect — " + "increase MAX_NEW_TOKENS to make the test meaningful.", + ) + + self.assertEqual( + done_after - done_before, + 0, + f"{done_after - done_before} requests completed during pause " + f"({done_before} before, {done_after} after) — " + f"pause_generation was not respected by the disagg scheduler.", + ) + + requests.post( + self.decode_url + "/continue_generation", + json={}, + timeout=30, + ).raise_for_status() + requests.post( + self.prefill_url + "/continue_generation", + json={}, + timeout=30, + ).raise_for_status() + + completed = 0 + errors = [] + for future in as_completed(futures, timeout=REQUEST_TIMEOUT): + prompt_id = futures[future] + try: + resp = future.result() + if resp.status_code == 200: + body = resp.json() + self.assertIn("text", body) + self.assertGreater(len(body["text"]), 0) + completed += 1 + else: + errors.append(f"Request {prompt_id}: status={resp.status_code}") + except Exception as e: + errors.append(f"Request {prompt_id}: exception={e}") + + self.assertEqual( + completed + len(errors), + NUM_REQUESTS, + "Some requests did not resolve within the timeout — likely hung during pause.", + ) + self.assertEqual( + completed, + NUM_REQUESTS, + f"Some requests failed: {completed}/{NUM_REQUESTS} succeeded. Errors: {errors}", + ) + def test_first_token_finish(self): client = openai.Client(api_key="empty", base_url=f"{self.lb_url}/v1") tokenizer = AutoTokenizer.from_pretrained(self.model) From da6b8e14483cae7ab3fd2ac902472f3491481688 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sun, 12 Apr 2026 18:49:37 -0700 Subject: [PATCH 06/52] Extract pause_resume_in_place kit; rename test_abort to test_scheduler_control (#22647) --- .../sglang/test/kits/pause_generation_kit.py | 111 ++++++++++++++++++ .../test_disaggregation_basic.py | 102 +--------------- ...est_abort.py => test_scheduler_control.py} | 5 +- 3 files changed, 118 insertions(+), 100 deletions(-) create mode 100644 python/sglang/test/kits/pause_generation_kit.py rename test/registered/scheduler/{test_abort.py => test_scheduler_control.py} (98%) diff --git a/python/sglang/test/kits/pause_generation_kit.py b/python/sglang/test/kits/pause_generation_kit.py new file mode 100644 index 000000000000..00c1ef4cf0d9 --- /dev/null +++ b/python/sglang/test/kits/pause_generation_kit.py @@ -0,0 +1,111 @@ +import time +from concurrent.futures import ThreadPoolExecutor, as_completed + +import requests + +_REQUEST_TIMEOUT = 180 + + +class PauseResumeInPlaceMixin: + """Test pause/resume with in_place mode. + + Sends concurrent requests, pauses mid-generation, verifies no progress + during the pause window, then resumes and verifies all requests complete. + + Subclass must set: + - pause_generate_url: URL to send /generate requests (or falls back to self.base_url) + - pause_target_urls: list of URLs to send /pause_generation and /continue_generation + """ + + pause_num_requests: int = 32 + pause_max_new_tokens: int = 512 + pause_duration: float = 5 + pause_generate_url: str = "" + pause_target_urls: list = [] + + def test_pause_resume_in_place(self): + generate_url = self.pause_generate_url or self.base_url + target_urls = self.pause_target_urls or [self.base_url] + num_requests = self.pause_num_requests + + def _generate(prompt_id): + return requests.post( + generate_url + "/generate", + json={ + "text": f"Question {prompt_id}: Write a short essay about the number {prompt_id}.", + "sampling_params": { + "temperature": 0.8, + "max_new_tokens": self.pause_max_new_tokens, + }, + }, + timeout=_REQUEST_TIMEOUT, + ) + + with ThreadPoolExecutor(max_workers=num_requests) as executor: + futures = {executor.submit(_generate, i): i for i in range(num_requests)} + + time.sleep(1) + + # Pause all targets + for url in target_urls: + requests.post( + url + "/pause_generation", + json={"mode": "in_place"}, + timeout=30, + ).raise_for_status() + + time.sleep(0.5) + done_before = sum(1 for f in futures if f.done()) + + time.sleep(self.pause_duration) + done_after = sum(1 for f in futures if f.done()) + + self.assertLess( + done_before, + num_requests, + "All requests completed before pause took effect -- " + "increase pause_max_new_tokens to make the test meaningful.", + ) + + self.assertEqual( + done_after - done_before, + 0, + f"{done_after - done_before} requests completed during pause " + f"({done_before} before, {done_after} after) -- " + f"pause_generation was not respected by the scheduler.", + ) + + # Resume all targets (reverse order to unblock downstream first) + for url in reversed(target_urls): + requests.post( + url + "/continue_generation", + json={}, + timeout=30, + ).raise_for_status() + + completed = 0 + errors = [] + for future in as_completed(futures, timeout=_REQUEST_TIMEOUT): + prompt_id = futures[future] + try: + resp = future.result() + if resp.status_code == 200: + body = resp.json() + self.assertIn("text", body) + self.assertGreater(len(body["text"]), 0) + completed += 1 + else: + errors.append(f"Request {prompt_id}: status={resp.status_code}") + except Exception as e: + errors.append(f"Request {prompt_id}: exception={e}") + + self.assertEqual( + completed + len(errors), + num_requests, + "Some requests did not resolve within the timeout -- likely hung during pause.", + ) + self.assertEqual( + completed, + num_requests, + f"Some requests failed: {completed}/{num_requests} succeeded. Errors: {errors}", + ) diff --git a/test/registered/disaggregation/test_disaggregation_basic.py b/test/registered/disaggregation/test_disaggregation_basic.py index deeda4a8a92a..e20d59ea818a 100644 --- a/test/registered/disaggregation/test_disaggregation_basic.py +++ b/test/registered/disaggregation/test_disaggregation_basic.py @@ -1,9 +1,7 @@ import asyncio import json import os -import time import unittest -from concurrent.futures import ThreadPoolExecutor, as_completed from types import SimpleNamespace import aiohttp @@ -12,6 +10,7 @@ from transformers import AutoTokenizer from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kits.pause_generation_kit import PauseResumeInPlaceMixin from sglang.test.run_eval import run_eval from sglang.test.server_fixtures.disaggregation_fixture import ( PDDisaggregationServerBase, @@ -25,11 +24,13 @@ register_cuda_ci(est_time=394, suite="stage-b-test-2-gpu-large") -class TestDisaggregationAccuracy(PDDisaggregationServerBase): +class TestDisaggregationAccuracy(PauseResumeInPlaceMixin, PDDisaggregationServerBase): @classmethod def setUpClass(cls): super().setUpClass() cls.model = DEFAULT_MODEL_NAME_FOR_TEST + cls.pause_generate_url = cls.lb_url + cls.pause_target_urls = [cls.prefill_url, cls.decode_url] cls.launch_all() def test_gsm8k(self): @@ -99,101 +100,6 @@ def test_structured_output(self): # ensure the output is a valid JSON json.loads(output) - def test_pause_resume_in_place(self): - """Send requests, pause mid-generation, verify no progress during pause, resume.""" - NUM_REQUESTS = 32 - MAX_NEW_TOKENS = 512 - REQUEST_TIMEOUT = 180 - PAUSE_DURATION = 5 - - def _generate(prompt_id): - return requests.post( - self.lb_url + "/generate", - json={ - "text": f"Question {prompt_id}: Write a short essay about the number {prompt_id}.", - "sampling_params": { - "temperature": 0.8, - "max_new_tokens": MAX_NEW_TOKENS, - }, - }, - timeout=REQUEST_TIMEOUT, - ) - - with ThreadPoolExecutor(max_workers=NUM_REQUESTS) as executor: - futures = {executor.submit(_generate, i): i for i in range(NUM_REQUESTS)} - - time.sleep(1) - - requests.post( - self.prefill_url + "/pause_generation", - json={"mode": "in_place"}, - timeout=30, - ).raise_for_status() - requests.post( - self.decode_url + "/pause_generation", - json={"mode": "in_place"}, - timeout=30, - ).raise_for_status() - - time.sleep(0.5) - done_before = sum(1 for f in futures if f.done()) - - time.sleep(PAUSE_DURATION) - done_after = sum(1 for f in futures if f.done()) - - self.assertLess( - done_before, - NUM_REQUESTS, - "All requests completed before pause took effect — " - "increase MAX_NEW_TOKENS to make the test meaningful.", - ) - - self.assertEqual( - done_after - done_before, - 0, - f"{done_after - done_before} requests completed during pause " - f"({done_before} before, {done_after} after) — " - f"pause_generation was not respected by the disagg scheduler.", - ) - - requests.post( - self.decode_url + "/continue_generation", - json={}, - timeout=30, - ).raise_for_status() - requests.post( - self.prefill_url + "/continue_generation", - json={}, - timeout=30, - ).raise_for_status() - - completed = 0 - errors = [] - for future in as_completed(futures, timeout=REQUEST_TIMEOUT): - prompt_id = futures[future] - try: - resp = future.result() - if resp.status_code == 200: - body = resp.json() - self.assertIn("text", body) - self.assertGreater(len(body["text"]), 0) - completed += 1 - else: - errors.append(f"Request {prompt_id}: status={resp.status_code}") - except Exception as e: - errors.append(f"Request {prompt_id}: exception={e}") - - self.assertEqual( - completed + len(errors), - NUM_REQUESTS, - "Some requests did not resolve within the timeout — likely hung during pause.", - ) - self.assertEqual( - completed, - NUM_REQUESTS, - f"Some requests failed: {completed}/{NUM_REQUESTS} succeeded. Errors: {errors}", - ) - def test_first_token_finish(self): client = openai.Client(api_key="empty", base_url=f"{self.lb_url}/v1") tokenizer = AutoTokenizer.from_pretrained(self.model) diff --git a/test/registered/scheduler/test_abort.py b/test/registered/scheduler/test_scheduler_control.py similarity index 98% rename from test/registered/scheduler/test_abort.py rename to test/registered/scheduler/test_scheduler_control.py index b82a22e9ad91..2042ff00febf 100644 --- a/test/registered/scheduler/test_abort.py +++ b/test/registered/scheduler/test_scheduler_control.py @@ -10,6 +10,7 @@ from sglang.srt.utils import kill_process_tree from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.test.kits.abort_timeout_kit import AbortAllMixin, WaitingTimeoutMixin +from sglang.test.kits.pause_generation_kit import PauseResumeInPlaceMixin from sglang.test.test_utils import ( DEFAULT_MODEL_NAME_FOR_TEST, DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, @@ -19,7 +20,7 @@ run_and_check_memory_leak, ) -register_cuda_ci(est_time=328, suite="stage-b-test-1-gpu-small") +register_cuda_ci(est_time=360, suite="stage-b-test-1-gpu-small") register_amd_ci(est_time=300, suite="stage-b-test-1-gpu-small-amd") @@ -110,7 +111,7 @@ def test_memory_leak_with_api_key(self): ) -class TestAbortAll(AbortAllMixin, CustomTestCase): +class TestSchedulerControl(AbortAllMixin, PauseResumeInPlaceMixin, CustomTestCase): @classmethod def setUpClass(cls): cls.model = DEFAULT_MODEL_NAME_FOR_TEST From 934e19a6104b86d563a7f43a65e72be81b1349b7 Mon Sep 17 00:00:00 2001 From: blzheng Date: Mon, 13 Apr 2026 09:59:13 +0800 Subject: [PATCH 07/52] =?UTF-8?q?[CPU]=20Fix=20argument=20issues=20in=20qk?= =?UTF-8?q?v=5Fproj=5Fwith=5Frope=5Ffused=5Fweight=20and=20bmm=E2=80=A6=20?= =?UTF-8?q?(#21367)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Ma Mingfei --- python/sglang/srt/model_executor/cpu_graph_runner.py | 1 + .../attention_forward_methods/forward_mla_fused_rope_cpu.py | 4 ++-- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/model_executor/cpu_graph_runner.py b/python/sglang/srt/model_executor/cpu_graph_runner.py index 8db80fc6623c..28b808fd0c3d 100644 --- a/python/sglang/srt/model_executor/cpu_graph_runner.py +++ b/python/sglang/srt/model_executor/cpu_graph_runner.py @@ -222,6 +222,7 @@ def _( use_fp8_w8a16, qkv_a_proj_scale, q_b_proj_scale, + w_scale, is_vnni, block_size, q_lora_rank, diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_cpu.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_cpu.py index b6c76df09d2b..fa781b39c7fe 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_cpu.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_cpu.py @@ -100,7 +100,7 @@ def forward_absorb_fused_mla_rope_cpu_prepare( else None ) ), - self.w_scale, + self.w_scale if self.qkv_proj_with_rope_is_fp8 else None, True, # is_vnni self.weight_block_size, self.q_lora_rank, @@ -145,7 +145,7 @@ def forward_absorb_fused_mla_rope_cpu_core( attn_output.transpose(0, 1), self.w_vc, True, # is_vnni - self.w_scale, # scale + self.w_scale if self.qkv_proj_with_rope_is_fp8 else None, # scale ) attn_output = output output, _ = self.o_proj(attn_output) From 55935399426690573ad2e8fe040cb61dde837d82 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Sun, 12 Apr 2026 19:08:45 -0700 Subject: [PATCH 08/52] [RL] Refactor NVFP4 shuffling/swizzling to in-place replacement (#22204) --- .../moe/moe_runner/flashinfer_trtllm.py | 51 +++------ .../compressed_tensors_w4a4_nvfp4_moe.py | 35 ++---- .../srt/layers/quantization/modelopt_quant.py | 13 +-- python/sglang/srt/server_args.py | 3 +- .../test_flashinfer_trtllm_gen_moe_backend.py | 55 ++++++++++ ...est_update_weights_from_disk_blackwell.py} | 103 ++++++++++++------ 6 files changed, 157 insertions(+), 103 deletions(-) rename test/registered/rl/{test_update_weights_from_disk_mxfp8.py => test_update_weights_from_disk_blackwell.py} (69%) diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py index f4add35b391d..6a637805f4ad 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py @@ -275,15 +275,15 @@ def align_fp4_moe_weights_for_flashinfer_trtllm(layer: Module) -> None: w13_weight.size(0), # num_experts ) - # Set flashinfer parameters + # Set flashinfer parameters in-place + copy_or_rebind_param(layer, "w13_weight", gemm1_weights_fp4_shuffled.contiguous()) + copy_or_rebind_param(layer, "w2_weight", gemm2_weights_fp4_shuffled.contiguous()) copy_or_rebind_param( - layer, "gemm1_weights_fp4_shuffled", gemm1_weights_fp4_shuffled + layer, "w13_weight_scale", gemm1_scales_fp4_shuffled.contiguous() ) copy_or_rebind_param( - layer, "gemm2_weights_fp4_shuffled", gemm2_weights_fp4_shuffled + layer, "w2_weight_scale", gemm2_scales_fp4_shuffled.contiguous() ) - copy_or_rebind_param(layer, "gemm1_scales_fp4_shuffled", gemm1_scales_fp4_shuffled) - copy_or_rebind_param(layer, "gemm2_scales_fp4_shuffled", gemm2_scales_fp4_shuffled) # Compute additional scaling factor needed for TRT-LLM w2_input_scale_quant = cast(torch.Tensor, layer.w2_input_scale_quant) @@ -294,14 +294,6 @@ def align_fp4_moe_weights_for_flashinfer_trtllm(layer: Module) -> None: (w2_input_scale_quant * g1_alphas).to(torch.float32), ) - # Clean up weights that won't be used by TRT-LLM - del ( - layer.w2_weight, - layer.w2_weight_scale, - layer.w13_weight, - layer.w13_weight_scale, - ) - @dataclass class FlashInferTrtllmFp8MoeQuantInfo(MoeQuantInfo): @@ -560,11 +552,10 @@ def fused_experts_none_to_flashinfer_trtllm_fp8( class FlashInferTrtllmFp4MoeQuantInfo(MoeQuantInfo): """Quantization payload consumed by FlashInfer TRT-LLM FP4 MoE kernels.""" - # Shuffled FP4 weights (processed by align_fp4_moe_weights_for_flashinfer_trtllm) - gemm1_weights_fp4_shuffled: torch.Tensor - gemm2_weights_fp4_shuffled: torch.Tensor - gemm1_scales_fp4_shuffled: torch.Tensor - gemm2_scales_fp4_shuffled: torch.Tensor + w13_weight: torch.Tensor + w2_weight: torch.Tensor + w13_weight_scale: torch.Tensor + w2_weight_scale: torch.Tensor # Scaling factors g1_scale_c: torch.Tensor @@ -666,18 +657,14 @@ def fused_experts_none_to_flashinfer_trtllm_fp4( routing_bias=None, hidden_states=hs_fp4, hidden_states_scale=hs_scale, - gemm1_weights=quant_info.gemm1_weights_fp4_shuffled, - gemm1_weights_scale=quant_info.gemm1_scales_fp4_shuffled.view( - torch.float8_e4m3fn - ), + gemm1_weights=quant_info.w13_weight, + gemm1_weights_scale=quant_info.w13_weight_scale.view(torch.float8_e4m3fn), gemm1_bias=None, gemm1_alpha=None, gemm1_beta=None, gemm1_clamp_limit=None, - gemm2_weights=quant_info.gemm2_weights_fp4_shuffled, - gemm2_weights_scale=quant_info.gemm2_scales_fp4_shuffled.view( - torch.float8_e4m3fn - ), + gemm2_weights=quant_info.w2_weight, + gemm2_weights_scale=quant_info.w2_weight_scale.view(torch.float8_e4m3fn), gemm2_bias=None, output1_scale_scalar=quant_info.g1_scale_c, output1_scale_gate_scalar=quant_info.g1_alphas, @@ -716,18 +703,14 @@ def fused_experts_none_to_flashinfer_trtllm_fp4( routing_bias=correction_bias, hidden_states=hs_fp4, hidden_states_scale=hs_scale, - gemm1_weights=quant_info.gemm1_weights_fp4_shuffled, - gemm1_weights_scale=quant_info.gemm1_scales_fp4_shuffled.view( - torch.float8_e4m3fn - ), + gemm1_weights=quant_info.w13_weight, + gemm1_weights_scale=quant_info.w13_weight_scale.view(torch.float8_e4m3fn), gemm1_bias=None, gemm1_alpha=None, gemm1_beta=None, gemm1_clamp_limit=None, - gemm2_weights=quant_info.gemm2_weights_fp4_shuffled, - gemm2_weights_scale=quant_info.gemm2_scales_fp4_shuffled.view( - torch.float8_e4m3fn - ), + gemm2_weights=quant_info.w2_weight, + gemm2_weights_scale=quant_info.w2_weight_scale.view(torch.float8_e4m3fn), gemm2_bias=None, output1_scale_scalar=quant_info.g1_scale_c, output1_scale_gate_scalar=quant_info.g1_alphas, diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4_moe.py index 5898a078dbba..6b285809ba16 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4_moe.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4_moe.py @@ -20,6 +20,7 @@ from sglang.srt.layers.quantization.utils import ( prepare_static_weights_for_trtllm_fp4_moe, reorder_w1w3_to_w3w1, + replace_parameter, swizzle_blockscale, ) from sglang.srt.utils import next_power_of_2, set_weight_attrs @@ -257,30 +258,16 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: ) logger.debug("Finished shuffling weights for TRT-LLM MOE") - layer.gemm1_weights_fp4_shuffled = torch.nn.Parameter( - gemm1_weights_fp4_shuffled, requires_grad=False - ) - layer.gemm2_weights_fp4_shuffled = torch.nn.Parameter( - gemm2_weights_fp4_shuffled, requires_grad=False - ) - layer.gemm1_scales_fp4_shuffled = torch.nn.Parameter( - gemm1_scales_fp4_shuffled, requires_grad=False - ) - layer.gemm2_scales_fp4_shuffled = torch.nn.Parameter( - gemm2_scales_fp4_shuffled, requires_grad=False - ) + replace_parameter(layer, "w13_weight", gemm1_weights_fp4_shuffled) + replace_parameter(layer, "w2_weight", gemm2_weights_fp4_shuffled) + replace_parameter(layer, "w13_weight_scale", gemm1_scales_fp4_shuffled) + replace_parameter(layer, "w2_weight_scale", gemm2_scales_fp4_shuffled) # Additional parameter needed for TRT-LLM layer.g1_scale_c = torch.nn.Parameter( (layer.w2_input_scale_quant * layer.g1_alphas).to(torch.float32), requires_grad=False, ) - - # Clean up weights that won't be used by TRT-LLM - del layer.w2_weight - del layer.w2_weight_scale - del layer.w13_weight - del layer.w13_weight_scale else: # swizzle weight scales layer.w13_weight_scale = torch.nn.Parameter( @@ -370,18 +357,14 @@ def apply_weights( routing_bias=correction_bias, hidden_states=hs_fp4, hidden_states_scale=hs_scale, - gemm1_weights=layer.gemm1_weights_fp4_shuffled, - gemm1_weights_scale=layer.gemm1_scales_fp4_shuffled.view( - torch.float8_e4m3fn - ), + gemm1_weights=layer.w13_weight, + gemm1_weights_scale=layer.w13_weight_scale.view(torch.float8_e4m3fn), gemm1_bias=None, gemm1_alpha=None, gemm1_beta=None, gemm1_clamp_limit=None, - gemm2_weights=layer.gemm2_weights_fp4_shuffled, - gemm2_weights_scale=layer.gemm2_scales_fp4_shuffled.view( - torch.float8_e4m3fn - ), + gemm2_weights=layer.w2_weight, + gemm2_weights_scale=layer.w2_weight_scale.view(torch.float8_e4m3fn), gemm2_bias=None, output1_scale_scalar=layer.g1_scale_c, output1_scale_gate_scalar=layer.g1_alphas, diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index 3115bf0a5b6b..c24185bdf633 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -1980,9 +1980,8 @@ def apply( ), f"{activation=} missing from {ACT_STR_TO_TYPE_MAP.keys()=}" moe_runner_config = self.moe_runner_config - # FlashInfer TRTLLM FP4 path - layer has shuffled weights only when - # backend is flashinfer_trtllm - if hasattr(layer, "gemm1_weights_fp4_shuffled"): + # FlashInfer TRTLLM FP4 path + if self.enable_flashinfer_trtllm_moe and hasattr(layer, "g1_scale_c"): from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( FlashInferTrtllmFp4MoeQuantInfo, ) @@ -1994,10 +1993,10 @@ def apply( ) quant_info = FlashInferTrtllmFp4MoeQuantInfo( - gemm1_weights_fp4_shuffled=layer.gemm1_weights_fp4_shuffled.data, - gemm2_weights_fp4_shuffled=layer.gemm2_weights_fp4_shuffled.data, - gemm1_scales_fp4_shuffled=layer.gemm1_scales_fp4_shuffled.data, - gemm2_scales_fp4_shuffled=layer.gemm2_scales_fp4_shuffled.data, + w13_weight=layer.w13_weight.data, + w2_weight=layer.w2_weight.data, + w13_weight_scale=layer.w13_weight_scale.data, + w2_weight_scale=layer.w2_weight_scale.data, g1_scale_c=layer.g1_scale_c.data, g1_alphas=layer.g1_alphas.data, g2_alphas=layer.g2_alphas.data, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 9e66c282801f..64ea2c1856fd 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2824,8 +2824,9 @@ def _handle_moe_kernel_config(self): assert self.quantization in [ "fp8", "mxfp8", + "modelopt_fp4", None, - ], f"Invalid quantization '{self.quantization}'. \nFlashInfer TRTLLM routed MOE supports only: 'fp8', 'mxfp8', or bfloat16 (None)." + ], f"Invalid quantization '{self.quantization}'. \nFlashInfer TRTLLM routed MOE supports only: 'fp8', 'mxfp8', 'modelopt_fp4', or bfloat16 (None)." self.disable_shared_experts_fusion = True logger.warning( "FlashInfer TRTLLM routed MoE is enabled. --disable-shared-experts-fusion is automatically set." diff --git a/test/registered/backends/test_flashinfer_trtllm_gen_moe_backend.py b/test/registered/backends/test_flashinfer_trtllm_gen_moe_backend.py index b63447a60cd5..9c7260b19966 100644 --- a/test/registered/backends/test_flashinfer_trtllm_gen_moe_backend.py +++ b/test/registered/backends/test_flashinfer_trtllm_gen_moe_backend.py @@ -157,6 +157,49 @@ def test_gsm8k(self): self.assertGreater(metrics["score"], 0.93) +class FlashinferTrtllmGenMoeBackendNVFP4Base: + backend = None + + @classmethod + def setUpClass(cls): + cls.model = "nvidia/Qwen3-30B-A3B-NVFP4" + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + env={**os.environ, "SGLANG_ENABLE_JIT_DEEPGEMM": "False"}, + other_args=[ + "--moe-runner-backend", + cls.backend, + "--tp-size", + "4", + "--ep-size", + "4", + "--mem-fraction-static", + "0.7", + ], + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="gsm8k", + api="completion", + max_tokens=512, + num_examples=200, + num_threads=128, + ) + metrics = run_eval(args) + print(f"{metrics=}") + self.assertGreater(metrics["score"], 0.89) + + class TestFlashinferTrtllmGenMoeBackendFP8( FlashinferTrtllmGenMoeBackendFP8Base, CustomTestCase ): @@ -175,6 +218,12 @@ class TestFlashinferTrtllmGenMoeBackendBF16( backend = "flashinfer_trtllm" +class TestFlashinferTrtllmGenMoeBackendNVFP4( + FlashinferTrtllmGenMoeBackendNVFP4Base, CustomTestCase +): + backend = "flashinfer_trtllm" + + class TestFlashinferTrtllmGenMoeBackendFP8Routed( FlashinferTrtllmGenMoeBackendFP8Base, CustomTestCase ): @@ -193,5 +242,11 @@ class TestFlashinferTrtllmGenMoeBackendBF16Routed( backend = "flashinfer_trtllm_routed" +class TestFlashinferTrtllmGenMoeBackendNVFP4Routed( + FlashinferTrtllmGenMoeBackendNVFP4Base, CustomTestCase +): + backend = "flashinfer_trtllm_routed" + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/rl/test_update_weights_from_disk_mxfp8.py b/test/registered/rl/test_update_weights_from_disk_blackwell.py similarity index 69% rename from test/registered/rl/test_update_weights_from_disk_mxfp8.py rename to test/registered/rl/test_update_weights_from_disk_blackwell.py index 12ac0758cdb6..58d79c375699 100644 --- a/test/registered/rl/test_update_weights_from_disk_mxfp8.py +++ b/test/registered/rl/test_update_weights_from_disk_blackwell.py @@ -15,37 +15,43 @@ ) -class TestServerUpdateWeightsFromDiskMXFP8(CustomTestCase): - model = "zianglih/Qwen3-30B-A3B-Instruct-2507-MXFP8-last-8-BF16" +class UpdateWeightsFromDiskBase: + model = None base_url = DEFAULT_URL_FOR_TEST request_timeout = 120 update_timeout = 240 + launch_env = None decode_payload = { "text": "The capital of France is", "sampling_params": {"temperature": 0, "max_new_tokens": 16}, } - backend_test_suites = ( - { - "fp8_gemm_backend": "flashinfer_trtllm", - "moe_runner_backend": "flashinfer_trtllm_routed", - }, + backend_test_suites = () + update_test_suites = ( + {"flush_cache": True, "abort_all_requests": False}, + {"flush_cache": False, "abort_all_requests": False}, ) - def _launch_server(self, fp8_gemm_backend, moe_runner_backend): + @classmethod + def setUpClass(cls): + super().setUpClass() + if cls.model is None: + raise NotImplementedError("Subclass must set 'model' attribute") + if not cls.backend_test_suites: + raise NotImplementedError( + "Subclass must set non-empty 'backend_test_suites'" + ) + + def _launch_server(self, backend_test_suite): + launch_kwargs = {} + if self.launch_env is not None: + launch_kwargs["env"] = self.launch_env + other_args = backend_test_suite.get("other_args") return popen_launch_server( self.model, self.base_url, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=[ - "--base-gpu-id", - "0", - "--tp-size", - "4", - "--fp8-gemm-backend", - fp8_gemm_backend, - "--moe-runner-backend", - moe_runner_backend, - ], + other_args=other_args, + **launch_kwargs, ) def _get_json(self, endpoint, timeout=None): @@ -119,30 +125,19 @@ def _run_update_weights( timeout=self.update_timeout, ) - def test_parameterized_update_weights_mxfp8(self): - update_test_suites = ( - {"flush_cache": True, "abort_all_requests": False}, - {"flush_cache": False, "abort_all_requests": False}, - ) + def test_parameterized_update_weights_from_disk(self): for backend_test_suite in self.backend_test_suites: - with self.subTest(**backend_test_suite): - process = self._launch_server( - backend_test_suite["fp8_gemm_backend"], - backend_test_suite["moe_runner_backend"], - ) + case_name = backend_test_suite.get("name", "default") + with self.subTest(model=self.model, case_name=case_name): + process = self._launch_server(backend_test_suite) try: origin_model_path = self._get_model_info() self.assertEqual(origin_model_path, self.model) self._assert_non_empty_decode() baseline_sig = self._get_decode_logprob_signature() - for update_test_suite in update_test_suites: - with self.subTest( - fp8_gemm_backend=backend_test_suite["fp8_gemm_backend"], - moe_runner_backend=backend_test_suite["moe_runner_backend"], - flush_cache=update_test_suite["flush_cache"], - abort_all_requests=update_test_suite["abort_all_requests"], - ): + for update_test_suite in self.update_test_suites: + with self.subTest(case_name=case_name, **update_test_suite): ret = self._run_update_weights( self.model, flush_cache=update_test_suite["flush_cache"], @@ -161,5 +156,43 @@ def test_parameterized_update_weights_mxfp8(self): kill_process_tree(process.pid) +class TestServerUpdateWeightsFromDiskMXFP8(UpdateWeightsFromDiskBase, CustomTestCase): + model = "zianglih/Qwen3-30B-A3B-Instruct-2507-MXFP8-last-8-BF16" + backend_test_suites = ( + { + "name": "flashinfer_trtllm_routed_mxfp8", + "other_args": ( + "--base-gpu-id", + "0", + "--tp-size", + "4", + "--fp8-gemm-backend", + "flashinfer_trtllm", + "--moe-runner-backend", + "flashinfer_trtllm_routed", + ), + }, + ) + + +class TestServerUpdateWeightsFromDiskNVFP4(UpdateWeightsFromDiskBase, CustomTestCase): + model = "nvidia/Qwen3-30B-A3B-NVFP4" + backend_test_suites = ( + { + "name": "flashinfer_trtllm_nvfp4", + "other_args": ( + "--base-gpu-id", + "0", + "--tp-size", + "4", + "--fp4-gemm-backend", + "flashinfer_trtllm", + "--moe-runner-backend", + "flashinfer_trtllm_routed", + ), + }, + ) + + if __name__ == "__main__": unittest.main() From bc59cc0f9682079f30318fb1c658fc9ecc67422b Mon Sep 17 00:00:00 2001 From: Zhangheng Date: Mon, 13 Apr 2026 10:28:22 +0800 Subject: [PATCH 09/52] [RaidxTree Refactor]: Support Unified HybridRadixTree V2 (#21206) Co-authored-by: ispobock Co-authored-by: pansicheng Co-authored-by: yizhang2077 <1109276519@qq.com> Co-authored-by: xiezhq-hermann --- python/sglang/srt/environ.py | 3 + python/sglang/srt/managers/scheduler.py | 15 + .../sglang/srt/mem_cache/base_prefix_cache.py | 2 +- .../sglang/srt/mem_cache/cache_init_params.py | 3 + .../unified_cache_components/README.md | 329 ++++++ .../unified_cache_components/__init__.py | 25 + .../full_component.py | 102 ++ .../mamba_component.py | 270 +++++ .../unified_cache_components/swa_component.py | 297 ++++++ .../tree_component.py | 292 ++++++ .../srt/mem_cache/unified_radix_cache.py | 986 ++++++++++++++++++ python/sglang/test/kl_multiturn_utils.py | 460 ++++++++ .../test_unified_radix_cache_kl.py | 266 +++++ .../test_unified_radix_cache_bench.py | 775 ++++++++++++++ .../test_unified_radix_cache_unittest.py | 883 ++++++++++++++++ 15 files changed, 4707 insertions(+), 1 deletion(-) create mode 100644 python/sglang/srt/mem_cache/unified_cache_components/README.md create mode 100644 python/sglang/srt/mem_cache/unified_cache_components/__init__.py create mode 100644 python/sglang/srt/mem_cache/unified_cache_components/full_component.py create mode 100644 python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py create mode 100644 python/sglang/srt/mem_cache/unified_cache_components/swa_component.py create mode 100644 python/sglang/srt/mem_cache/unified_cache_components/tree_component.py create mode 100644 python/sglang/srt/mem_cache/unified_radix_cache.py create mode 100644 python/sglang/test/kl_multiturn_utils.py create mode 100644 test/registered/radix_cache/test_unified_radix_cache_kl.py create mode 100644 test/registered/unit/mem_cache/test_unified_radix_cache_bench.py create mode 100644 test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index d5afb072dfb7..42309d9039fe 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -472,6 +472,9 @@ class Envs: SGLANG_MAMBA_CONV_DTYPE = EnvStr("bfloat16") SGLANG_MAMBA_SSM_DTYPE = EnvStr(None) + # Unified Radix Tree + SGLANG_ENABLE_UNIFIED_RADIX_TREE = EnvBool(False) + # Breakable CUDA Graph SGLANG_USE_BREAKABLE_CUDA_GRAPH = EnvBool(False) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 1f95448002dc..30c242daa966 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -835,6 +835,21 @@ def init_cache_with_memory_pool(self): self.tp_worker.register_hicache_layer_transfer_counter( self.tree_cache.cache_controller.layer_done_counter ) + elif envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.get(): + from sglang.srt.mem_cache.unified_cache_components import ( + ComponentType, + ) + from sglang.srt.mem_cache.unified_radix_cache import ( + UnifiedRadixCache, + ) + + tree_components = [ComponentType.FULL] + if self.is_hybrid_swa or self.is_hybrid_ssm: + tree_components.append( + ComponentType.SWA if self.is_hybrid_swa else ComponentType.MAMBA + ) + params.tree_components = tuple(tree_components) + self.tree_cache = UnifiedRadixCache(params) elif self.is_hybrid_swa: from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index be219339cef6..d19be907f26f 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -47,7 +47,7 @@ class MatchPrefixParams: class InsertParams: """Unified parameters for insert across different cache types""" - key: RadixKey + key: Optional[RadixKey] = None value: Optional[torch.Tensor] = None # Mamba specific diff --git a/python/sglang/srt/mem_cache/cache_init_params.py b/python/sglang/srt/mem_cache/cache_init_params.py index d8731160c842..6f6fafae05aa 100644 --- a/python/sglang/srt/mem_cache/cache_init_params.py +++ b/python/sglang/srt/mem_cache/cache_init_params.py @@ -8,6 +8,7 @@ if TYPE_CHECKING: from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.memory_pool import ReqToTokenPool + from sglang.srt.mem_cache.unified_cache_components import ComponentType @dataclasses.dataclass @@ -39,3 +40,5 @@ class CacheInitParams: # Time-to-live for cache entries in seconds. If None, TTL is disabled. cache_ttl_seconds: Optional[float] = None + + tree_components: Optional[tuple[ComponentType, ...]] = None diff --git a/python/sglang/srt/mem_cache/unified_cache_components/README.md b/python/sglang/srt/mem_cache/unified_cache_components/README.md new file mode 100644 index 000000000000..c698b542ad72 --- /dev/null +++ b/python/sglang/srt/mem_cache/unified_cache_components/README.md @@ -0,0 +1,329 @@ +# Unified Radix Cache + +A component-based, pluggable prefix cache framework for SGLang that unifies Full-attention, Sliding-Window-Attention (SWA), and Mamba/SSM caching into a single radix tree. + +## Design Goals + +1. **Unified tree structure** — One radix tree manages all KV cache types instead of separate specialized implementations (`SWARadixCache`, `MambaRadixCache`, etc.). +2. **Pluggable components** — Each attention/state type (Full, SWA, Mamba) is a `TreeComponent` that implements hook interfaces. Adding a new cache type only requires adding a new component. +3. **Per-component resource isolation** — Each component has its own LRU list, lock reference counting, evictable/protected size tracking, and eviction driver. +4. **Cascade eviction with priority** — When a component evicts a node, lower-or-equal-priority components on the same node are evicted together, maintaining cross-component consistency. +5. **Zero special-casing in the main tree** — The tree operates purely on keys (logical). All physical resource management (allocation, freeing, copy-on-write) is handled by components through hooks. + +## Architecture + +``` +┌───────────────────────────────────────────────┐ +│ UnifiedRadixCache │ +│ (unified_radix_cache.py) │ +│ │ +│ root_node ──► UnifiedTreeNode (radix tree) │ +│ components ► {name → TreeComponent} │ +│ lru_lists ─► {name → UnifiedLRUList} │ +└──────────┬───────────┬───────────┬────────────┘ + │ │ │ + ▼ ▼ ▼ + ┌────────────┐ ┌──────────┐ ┌─────────────┐ + │ Full │ │ SWA │ │ Mamba │ + │ Component │ │Component │ │ Component │ + └─────┬──────┘ └────┬─────┘ └──────┬──────┘ + │ │ │ + └─────────────┼──────────────┘ + ▼ + ┌──────────────┐ + │TreeComponent │ + │ (ABC) │ + └──────────────┘ +``` + +### Key Data Structures + +**`UnifiedTreeNode`** — Each node stores per-component data independently: + +```python +node.component_data = { + "full": ComponentData(value=Tensor|None, lock_ref=int, metadata={}), + "swa": ComponentData(value=Tensor|None, lock_ref=int, metadata={}), + "mamba": ComponentData(value=Tensor|None, lock_ref=int, metadata={}), +} +``` + +**`UnifiedLRUList`** — One doubly-linked list per component, threaded through the same tree nodes via `lru_prev[name]`/`lru_next[name]`. Supports O(1) insert/remove/promote and O(L) scan for eviction (L = locked nodes skipped). + +**`ComponentData`** — Per-component data stored on each node: +- `value: Tensor | None` — Device indices into the component's memory pool (`TokenToKVPool` for Full, `SWAKVPool` for SWA, `MambaPool` for Mamba). `None` means tombstone (data evicted but node structure retained). +- `lock_ref: int` — Reference count of active requests using this node's component data. `lock_ref > 0` protects the node from eviction. +- `metadata: dict` — Component-specific state (e.g., SWA stores `component_uuid` for window-lock boundary tracking). + +--- + +## File Layout + +| File | Contents | +|------|----------| +| `../unified_radix_cache.py` | `UnifiedRadixCache`, `UnifiedTreeNode`, `UnifiedLRUList`, factory `create_unified_radix_cache` | +| `tree_component.py` | `TreeComponent` ABC, `ComponentType`, `ComponentData`, `get_and_increase_time_counter`, `next_component_uuid` | +| `full_component.py` | `FullComponent` — standard full-attention KV cache component | +| `swa_component.py` | `SWAComponent` — sliding-window attention component with tombstone/window tracking | +| `mamba_component.py` | `MambaComponent` — Mamba/SSM state component with copy-on-write | +| `hybrid_cache_controller.py` | `HybridCacheController` — HiCache 3-tier storage controller (L1 GPU → L2 CPU → L3 Disk) | +| `__init__.py` | Re-exports: `ComponentName`, `ComponentData`, `TreeComponent`, `FullComponent`, `SWAComponent`, `MambaComponent` | + +--- + +## Public API Reference + +All public APIs are on `UnifiedRadixCache`, which implements `BasePrefixCache`. + +**Notation**: K = key length (tokens), D = matched path depth in tree (D ≤ K/P), P = page_size, C = number of components (≤ 3, treated as constant). + +All tree traversal operations have two cost components: **O(K)** for data operations (key comparison, tensor clone/concat) + **O(D·C)** for component overhead (C hooks per node). Since D ≤ K/P and C is constant, overall **O(K)**. + +### `match_prefix(params: MatchPrefixParams) → MatchResult` + +Find the longest cached prefix for a token sequence. + +| Aspect | Detail | +|--------|--------| +| **Purpose** | Walk the radix tree to find the longest prefix where **all** component validators pass | +| **Inputs** | `params.key: RadixKey` — token IDs + optional extra key for namespace isolation | +| **Output** | `MatchResult(device_indices, last_device_node, last_host_node, mamba_branching_seqlen, ...)` | +| **Mutation** | Updates `last_access_time` on matched path; promotes matched nodes to MRU in all component LRU lists; may trigger `_split_node` if match ends mid-node | +| **Complexity** | **O(K + D·C)** | + +**Algorithm detail:** +1. Calls `create_match_validator()` once per component — returns a stateful closure (e.g., SWA tracks accumulated window length) +2. Walks tree edges via `key_match_fn`; at each node, calls all validator closures — the match boundary is only advanced when **all** validators return `True` +3. If match ends mid-node, calls `_split_node` → triggers `redistribute_on_node_split()` per component +4. Post-match (`_match_post_processor`): + - Promotes matched path to MRU in each component's LRU via `node_has_component_data()` as filter + - Updates `last_access_time` with decreasing timestamps up the path (parent < child) + - Concatenates matched device indices via `torch.cat` (concat length ≤ K, subsumed by O(K)) + - Calls `finalize_match_result()` per component (Mamba performs copy-on-write: allocates new pool slot, copies SSM state) + +--- + +### `insert(params: InsertParams) → InsertResult` + +Insert a key-value pair into the tree. + +| Aspect | Detail | +|--------|--------| +| **Purpose** | Insert token sequence + KV indices, reusing existing prefix and freeing duplicate KV slots | +| **Inputs** | `params.key: RadixKey`, `params.value: Tensor` (KV pool indices), plus component-specific fields (`mamba_value`, `swa_evicted_seqlen`, `prev_prefix_len`) | +| **Output** | `InsertResult(prefix_len, mamba_exist)` — `prefix_len` = length of reused prefix | +| **Mutation** | Creates new leaf nodes; updates component data on overlapping nodes; frees duplicate KV indices; may split nodes; updates LRU lists and evictable sizes | +| **Complexity** | **O(K + D·C)** | + +**Algorithm detail** (`_insert_helper`): +1. At each existing node, calls `_touch_node` → promotes to MRU via `node_has_component_data()` +2. If key diverges mid-node, calls `_split_node` → `redistribute_on_node_split()` per component +3. For each overlapping node, calls `update_component_on_insert_overlap()` per component — returns `consumed_from` index; the tree frees `value[dup_start:consumed_from]` as duplicate pool indices + - Full: returns `prefix_len` (no consumption, default behavior) + - SWA: checks if the overlapping node is a tombstone (SWA value = None) within the SWA window boundary (`swa_evicted_seqlen`): + - If entirely within window: **recovers tombstone** — frees old `full_value`, clones `value_slice`, translates to SWA indices, inserts into SWA LRU (returns `0` = all consumed) + - If partially within window: **splits node** at boundary, recovers SWA on the window portion (returns `start_idx`) + - If entirely outside window: returns `prefix_len` (no consumption) + - Mamba: returns `prefix_len` (no consumption, default behavior) +4. Before creating a new leaf, checks `should_skip_leaf_creation()` per component — any veto aborts leaf creation and frees remaining value +5. Creates leaf via `_add_new_node` (clones value tensor, inserts into Full LRU) +6. Calls `commit_insert_component_data()` per component on the final target node (SWA may trigger a secondary split for window boundary; Mamba sets mamba pool indices and inserts into Mamba LRU) + +--- + +### `evict(params: EvictParams) → EvictResult` + +Free cached tokens to reclaim memory. + +| Aspect | Detail | +|--------|--------| +| **Purpose** | Each component drives eviction from its own LRU list until its target is met | +| **Inputs** | `params.num_tokens` (full), `params.swa_num_tokens` (SWA), `params.mamba_num` (Mamba) | +| **Output** | `EvictResult(num_tokens_evicted, swa_num_tokens_evicted, mamba_num_evicted)` | +| **Mutation** | Frees pool indices; removes nodes from LRU lists; deletes leaf nodes from tree; cascades to lower-priority components; walks up parent chain to delete tombstone ancestors | +| **Complexity** | **O(E·H + L)** — E = nodes evicted, H = tombstone chain height, L = locked nodes skipped in LRU scan. | + +**Algorithm detail:** +1. Calls `drive_eviction()` for each component: + - Full: scans Full LRU from tail, only evicts **leaf** nodes (`get_leaf_lru_no_lock` — **O(L)**); calls `evict_component()` to free pool indices + - SWA: scans SWA LRU from tail; **internal** nodes are tombstoned (evict SWA data, keep node), **leaf** nodes are fully deleted; both trigger cascade + - Mamba: scans Mamba LRU from tail; **internal** nodes are tombstoned, **leaf** nodes are fully deleted; both trigger cascade +2. After each node eviction, calls `_cascade_evict`: + - Queries `eviction_priority()` per component; evicts all with priority ≤ trigger's + - Calls `evict_component()` + `node_has_component_data()` for cascaded components + - For leaf: removes from parent, then `_iteratively_delete_tombstone_leaf` walks up **O(H)** ancestors + +**Cascade eviction rules:** +- **Leaf nodes**: all priorities = 0 → evicting any cascades to all (node deleted) +- **Internal nodes**: Full(2) > SWA(1) > Mamba(0) + - Evicting Mamba: no cascade + - Evicting SWA: cascades to Mamba + - Evicting Full: cascades to SWA + Mamba + +--- + +### `inc_lock_ref(node: UnifiedTreeNode) → IncLockRefResult` + +Lock a node to protect it (and its ancestors) from eviction. + +| Aspect | Detail | +|--------|--------| +| **Purpose** | Called when a request begins using a cached prefix — prevents eviction of nodes it depends on | +| **Inputs** | `node` — the last matched node (deepest) | +| **Output** | `IncLockRefResult(swa_uuid_for_lock)` | +| **Mutation** | Increments `lock_ref` per component along the path; moves tokens from evictable to protected size counters | +| **Complexity** | **O(D)** — Full: node to root; SWA: up to window boundary O(min(D, W)); Mamba: O(1).| + +**Algorithm detail:** Calls `acquire_component_lock()` for each component. + +| Component | Strategy | +|-----------|----------| +| Full | **Path-lock**: walks from node to root, `lock_ref += 1` on every ancestor. On first lock (`lock_ref: 0→1`), moves tokens from `component_evictable_size_` to `component_protected_size_`. | +| SWA | **Window-lock**: walks upward, accumulating SWA value lengths until `sliding_window_size` is filled. Records a `component_uuid` at the boundary node for `dec_lock_ref` to know where to stop. | +| Mamba | **Single-node lock**: only `lock_ref += 1` on the node itself (mamba state is per-leaf, not per-path). | + +--- + +### `dec_lock_ref(node, params?) → DecLockRefResult` + +Unlock a previously locked node path. + +| Aspect | Detail | +|--------|--------| +| **Purpose** | Called when a request finishes — releases eviction protection | +| **Inputs** | `node`, optional `params.swa_uuid_for_lock` for SWA boundary detection | +| **Output** | `DecLockRefResult()` | +| **Mutation** | Decrements `lock_ref` per component; moves tokens from protected back to evictable when `lock_ref` reaches 0 | +| **Complexity** | **O(D)** — symmetric to `inc_lock_ref` | + +**Algorithm detail:** Calls `release_component_lock()` for each component. Full walks to root; SWA walks up until matching `component_uuid`; Mamba decrements single node. + +--- + +### `cache_finished_req(req: Req, is_insert: bool = True)` + +Cache a completed request's KV data into the tree. + +| Aspect | Detail | +|--------|--------| +| **Purpose** | After a request finishes, insert its token/KV data into the tree for future reuse | +| **Inputs** | `req` — the finished request; `is_insert` — whether to insert (True) or just release locks (False) | +| **Output** | `None` | +| **Mutation** | Calls component hooks → `insert` → `dec_lock_ref` → component cleanup. Frees unaligned tail KV indices; frees non-inserted KV indices when `is_insert=False`. | +| **Complexity** | **O(K + D·C)** — insert O(K + D·C) + lock release O(D). Simplifies to **O(K)**. | + +**Algorithm detail:** +1. `prepare_for_caching_req()` per component — sets component-specific insert params, returns effective cache length (SWA: sets `swa_evicted_seqlen`; Mamba: prepares `mamba_value` from ping-pong buffer, returns `mamba_last_track_seqlen` as truncation hint) +2. Truncates if `effective_cache_len < len(token_ids)`: frees excess pool indices +3. Converts token IDs (bigram if EAGLE), page-aligns keys, then calls `insert()` +4. Frees unaligned tail KV indices beyond page boundary +5. Calls `dec_lock_ref()` on the previous `req.last_node` +6. `cleanup_after_caching_req()` per component (Mamba: frees forked mamba_value based on `mamba_exist`, handles ping-pong buffer cleanup) + +--- + +### `cache_unfinished_req(req: Req, chunked=False)` + +Cache an in-progress request's partial KV data (chunked prefill). + +| Aspect | Detail | +|--------|--------| +| **Purpose** | During chunked prefill, insert partial results so the next chunk can match the prefix | +| **Inputs** | `req` — the in-progress request | +| **Output** | `None` | +| **Mutation** | Inserts partial KV → re-matches prefix → updates `req.prefix_indices`, `req.cache_protected_len`, `req.last_node`; transfers lock from old node to new node | +| **Complexity** | **O(K + D·C)** — two tree traversals: insert O(K + D·C) + re-match O(K + D·C) + lock transfer O(D). Simplifies to **O(K)**. | + +**Algorithm detail:** +1. `prepare_for_caching_req()` per component +2. `insert()` — first tree traversal +3. `match_prefix()` — **second** tree traversal to get updated indices +4. Writes new prefix indices into `req_to_token_pool` +5. `dec_lock_ref()` on old `req.last_node` +6. `inc_lock_ref()` on new matched node +7. Updates `req.prefix_indices`, `req.cache_protected_len`, `req.last_node` +8. `cleanup_after_caching_req()` per component + +--- + +## TreeComponent Hook Reference + +Each component implements these hooks. See `tree_component.py` for the ABC and docstrings. + +### Match Phase + +| Hook | Purpose | Called By | Default | +|------|---------|-----------|----------| +| `create_match_validator()` | Return a per-match stateful predicate that decides whether a node is a valid match boundary. Full: always True. SWA: tracks accumulated window length, True when contiguous window ≥ `sliding_window_size`. Mamba: True iff node has mamba data. | `_match_prefix_helper` | *abstract* | +| `finalize_match_result()` | Post-process the match result after prefix matching completes. Full/SWA: pass-through. Mamba: copy-on-write — allocates a new mamba pool slot, copies SSM state into the request pool, records `branching_seqlen`. | `_match_post_processor` | pass-through | + +### Insert Phase + +| Hook | Purpose | Called By | Default | +|------|---------|-----------|----------| +| `update_component_on_insert_overlap()` | Handle key overlap with an existing node during insert. Returns the index within `value_slice` from which this component consumed (took ownership of) pool slots. Full/Mamba: no consumption (`prefix_len`). SWA: may recover tombstoned nodes within the sliding window boundary. | `_insert_helper` | returns `prefix_len` | +| `should_skip_leaf_creation()` | Veto leaf creation when the entire new leaf would be a tombstone for this component. SWA: vetoes if `swa_evicted_seqlen ≥ total_prefix_len + key_len`. | `_insert_helper` | `False` | +| `commit_insert_component_data()` | Finalize component data on the target node after the insert walk completes. Full: no-op (handled by `_add_new_node`). SWA: checks window boundary, may split node — parent becomes tombstone, child gets SWA data. Mamba: sets mamba pool indices and inserts into Mamba LRU. | `_insert_helper` | no-op | + +### Node Split + +| Hook | Purpose | Called By | Default | +|------|---------|-----------|----------| +| `redistribute_on_node_split()` | Redistribute component data between new parent (prefix) and child (suffix) when a node is split. Full: copies `lock_ref` to parent. SWA: slices SWA value, copies `lock_ref` and `component_uuid`. Mamba: parent gets `None`/`lock_ref=0` (mamba stays on leaf). | `_split_node` | *abstract* | + +### Eviction Phase + +| Hook | Purpose | Called By | Default | +|------|---------|-----------|----------| +| `evict_component()` | Free this component's KV resources on a node being evicted. Internal nodes: free memory and tombstone (`value = None`). Leaf nodes: free memory, node will be deleted. Returns number of tokens freed. | `_evict_component_and_detach_lru` | *abstract* | +| `eviction_priority()` | Return cascade eviction priority (higher = evicted later). Leaf: all 0. Internal: Full(2) > SWA(1) > Mamba(0). When evicting, all components with ≤ priority on the same node are cascade-evicted. | `_cascade_evict` | `0` | +| `drive_eviction()` | Drive eviction from this component's LRU list until the target amount is freed. Full: leaf-only from Full LRU. SWA: both internal (tombstone) and leaf from SWA LRU. Mamba: both internal (tombstone) and leaf from Mamba LRU. | `evict` | *abstract* | + +### Lock Phase + +| Hook | Purpose | Called By | Default | +|------|---------|-----------|----------| +| `acquire_component_lock()` | Increment `lock_ref` to protect nodes from eviction; moves tokens from evictable to protected. Full: path-lock to root. SWA: window-lock with UUID boundary. Mamba: single-node lock. | `inc_lock_ref` | *abstract* | +| `release_component_lock()` | Decrement `lock_ref` to un-protect nodes; moves tokens from protected to evictable when `lock_ref` → 0. Full: path-unlock to root. SWA: walks up to UUID boundary. Mamba: single-node unlock. | `dec_lock_ref` | *abstract* | + +### Caching Phase + +| Hook | Purpose | Called By | Default | +|------|---------|-----------|----------| +| `prepare_for_caching_req()` | Prepare component-specific data before insert, fill fields in `InsertParams`, return effective cache length. Full: no-op. SWA: sets `swa_evicted_seqlen`. Mamba: prepares `mamba_value` from ping-pong buffer, returns `mamba_last_track_seqlen`. | `cache_finished/unfinished_req` | returns `None` | +| `cleanup_after_caching_req()` | Post-cache cleanup. Full/SWA: no-op. Mamba: frees forked `mamba_value` based on `mamba_exist`, handles ping-pong buffer `keep_idx`, resets `mamba_last_track_seqlen` on unfinished. | `cache_finished/unfinished_req` | no-op | + +### Utility + +| Hook | Purpose | Called By | Default | +|------|---------|-----------|----------| +| `node_has_component_data()` | Check if a node has this component's data. Used as filter for LRU operations and cascade checks. Full overrides to check `full_value` directly. | multiple | `value is not None` | + +### Component Behavior Summary + +| Behavior | FullComponent | SWAComponent | MambaComponent | +|----------|--------------|-------------|----------------| +| **Validator** | Always `True` | Tracks accumulated window; `True` when ≥ `sliding_window_size` | `True` iff node has mamba data | +| **Lock strategy** | Path-lock (root → node) | Window-lock (up to window boundary, UUID-tagged) | Single-node lock | +| **Internal eviction priority** | 2 (last) | 1 (middle) | 0 (first) | +| **Split behavior** | Copy `lock_ref` to parent | Slice SWA value + copy UUID | Parent gets `None` (mamba stays on leaf) | +| **Match finalize** | No-op | No-op | Copy-on-write: allocate new mamba slot, copy state | +| **Drive eviction** | Full LRU (leaf-only) → cascade all | SWA LRU → tombstone internal, cascade leaf | Mamba LRU → tombstone internal, cascade leaf | + +--- + +## Factory Function + +```python +def create_unified_radix_cache( + params: CacheInitParams, + component_names: Optional[tuple[ComponentName, ...]] = None, +) -> UnifiedRadixCache +``` + +Auto-detects component configuration from `params` if `component_names` is not specified: +- `SWATokenToKVPoolAllocator` → `(SWA,)` → `UnifiedSWARadixCache` +- `HybridReqToTokenPool` → `(MAMBA,)` → `UnifiedMambaRadixCache` +- Explicit tuple → `UnifiedRadixCache` with specified components + +Enable via `--enable-unified-radix-tree` server flag. diff --git a/python/sglang/srt/mem_cache/unified_cache_components/__init__.py b/python/sglang/srt/mem_cache/unified_cache_components/__init__.py new file mode 100644 index 000000000000..d0fde2786daf --- /dev/null +++ b/python/sglang/srt/mem_cache/unified_cache_components/__init__.py @@ -0,0 +1,25 @@ +from sglang.srt.mem_cache.unified_cache_components.full_component import FullComponent +from sglang.srt.mem_cache.unified_cache_components.mamba_component import MambaComponent +from sglang.srt.mem_cache.unified_cache_components.swa_component import SWAComponent +from sglang.srt.mem_cache.unified_cache_components.tree_component import ( + _NUM_COMPONENT_TYPES, + BASE_COMPONENT_TYPE, + ComponentData, + ComponentType, + TreeComponent, + get_and_increase_time_counter, + next_component_uuid, +) + +__all__ = [ + "BASE_COMPONENT_TYPE", + "ComponentData", + "ComponentType", + "FullComponent", + "MambaComponent", + "SWAComponent", + "TreeComponent", + "_NUM_COMPONENT_TYPES", + "next_component_uuid", + "get_and_increase_time_counter", +] diff --git a/python/sglang/srt/mem_cache/unified_cache_components/full_component.py b/python/sglang/srt/mem_cache/unified_cache_components/full_component.py new file mode 100644 index 000000000000..d963775d12e8 --- /dev/null +++ b/python/sglang/srt/mem_cache/unified_cache_components/full_component.py @@ -0,0 +1,102 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Callable, Optional + +from sglang.srt.mem_cache.base_prefix_cache import ( + DecLockRefParams, + EvictParams, + IncLockRefResult, +) +from sglang.srt.mem_cache.unified_cache_components.tree_component import ( + ComponentType, + TreeComponent, +) + +if TYPE_CHECKING: + from sglang.srt.mem_cache.unified_radix_cache import ( + UnifiedTreeNode, + ) + + +class FullComponent(TreeComponent): + component_type = ComponentType.FULL + + def __init__(self, cache, params): + super().__init__(cache, params) + allocator = cache.token_to_kv_pool_allocator + # When SWA is present, only free full-attention KV here; + # SWA KV will be freed by cascade via SWAComponent.evict_component. + if ComponentType.SWA in cache.tree_components: + self._free_full = allocator.full_attn_allocator.free + else: + self._free_full = allocator.free + + def node_has_component_data(self, node: UnifiedTreeNode) -> bool: + # Override so _for_each_component_lru includes Full in LRU operations + return node.component_data[self.component_type].value is not None + + def create_match_validator(self) -> Callable[[UnifiedTreeNode], bool]: + return lambda node: True + + def redistribute_on_node_split( + self, new_parent: UnifiedTreeNode, child: UnifiedTreeNode + ): + new_parent.component_data[self.component_type].lock_ref = child.component_data[ + self.component_type + ].lock_ref + + def evict_component(self, node: UnifiedTreeNode, is_leaf: bool) -> int: + cd = node.component_data[self.component_type] + self._free_full(cd.value) + freed = len(cd.value) + self.cache.component_evictable_size_[self.component_type] -= freed + return freed + + def eviction_priority(self, is_leaf: bool) -> int: + return 0 if is_leaf else 2 + + def drive_eviction( + self, params: EvictParams, tracker: dict[ComponentType, int] + ) -> None: + request = params.num_tokens + lru = self.cache.lru_lists[self.component_type] + while tracker[self.component_type] < request: + x = lru.get_leaf_lru_no_lock() + if x is None: + break + self.cache._evict_component_and_detach_lru( + x, self, is_leaf=True, tracker=tracker + ) + self.cache._cascade_evict(x, self, tracker) + + def acquire_component_lock( + self, node: UnifiedTreeNode, result: IncLockRefResult + ) -> IncLockRefResult: + ct = self.component_type + root = self.cache.root_node + cur = node + while cur != root: + cd = cur.component_data[ct] + if cd.lock_ref == 0: + key_len = len(cd.value) + self.cache.component_evictable_size_[ct] -= key_len + self.cache.component_protected_size_[ct] += key_len + cd.lock_ref += 1 + cur = cur.parent + return result + + def release_component_lock( + self, node: UnifiedTreeNode, params: Optional[DecLockRefParams] + ) -> None: + ct = self.component_type + root = self.cache.root_node + cur = node + while cur != root: + cd = cur.component_data[ct] + assert cd.lock_ref > 0 + if cd.lock_ref == 1: + key_len = len(cd.value) + self.cache.component_evictable_size_[ct] += key_len + self.cache.component_protected_size_[ct] -= key_len + cd.lock_ref -= 1 + cur = cur.parent diff --git a/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py b/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py new file mode 100644 index 000000000000..032e905dd62b --- /dev/null +++ b/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py @@ -0,0 +1,270 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Callable, Optional + +import torch + +from sglang.srt.mem_cache.base_prefix_cache import ( + DecLockRefParams, + EvictParams, + IncLockRefResult, + InsertParams, + InsertResult, + MatchPrefixParams, + MatchResult, +) +from sglang.srt.mem_cache.unified_cache_components.tree_component import ( + ComponentType, + TreeComponent, + get_and_increase_time_counter, +) +from sglang.srt.server_args import get_global_server_args + +if TYPE_CHECKING: + from sglang.srt.managers.schedule_batch import Req + from sglang.srt.mem_cache.cache_init_params import CacheInitParams + from sglang.srt.mem_cache.unified_radix_cache import ( + UnifiedRadixCache, + UnifiedTreeNode, + ) + + +class MambaComponent(TreeComponent): + component_type = ComponentType.MAMBA + + def __init__(self, cache: UnifiedRadixCache, params: CacheInitParams): + from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool + + assert isinstance( + cache.req_to_token_pool, HybridReqToTokenPool + ), f"MambaComponent requires HybridReqToTokenPool, got {type(cache.req_to_token_pool)}" + if not params.enable_mamba_extra_buffer: + assert ( + cache.page_size == 1 + ), f"MambaComponent requires page_size=1 when mamba_extra_buffer is disabled, got {cache.page_size}" + super().__init__(cache, params) + self.enable_mamba_extra_buffer = params.enable_mamba_extra_buffer + + def create_match_validator(self) -> Callable[[UnifiedTreeNode], bool]: + ct = self.component_type + return lambda node: node.component_data[ct].value is not None + + def finalize_match_result( + self, + result: MatchResult, + params: MatchPrefixParams, + value_chunks: list[torch.Tensor], + best_value_len: int, + ) -> MatchResult: + cow_mamba = params.cow_mamba + req = params.req + last_node = result.last_device_node + + if len(value_chunks) > best_value_len: + chunk_size = get_global_server_args().mamba_cache_chunk_size + aligned_seqlen = ( + sum(len(v) for v in value_chunks) // chunk_size + ) * chunk_size + branching_seqlen = aligned_seqlen if aligned_seqlen > 0 else None + else: + branching_seqlen = None + + mamba_value = last_node.component_data[self.component_type].value + if cow_mamba and mamba_value is not None: + assert req is not None + if req.mamba_pool_idx is None: + dst_index = self.cache.req_to_token_pool.mamba_pool.alloc(1) + if dst_index is None: + self.cache.inc_lock_ref(last_node) + self.cache.evict(EvictParams(num_tokens=0, mamba_num=1)) + dst_index = self.cache.req_to_token_pool.mamba_pool.alloc(1) + self.cache.dec_lock_ref(last_node) + assert dst_index is not None, "Can not alloc mamba cache" + self.cache.req_to_token_pool.mamba_pool.copy_from( + mamba_value, dst_index + ) + req.mamba_pool_idx = dst_index[0] + else: + dst_index = req.mamba_pool_idx.unsqueeze(0) + self.cache.req_to_token_pool.mamba_pool.copy_from( + mamba_value, dst_index + ) + + return result._replace(mamba_branching_seqlen=branching_seqlen) + + def commit_insert_component_data( + self, + node: UnifiedTreeNode, + is_new_leaf: bool, + params: InsertParams, + result: InsertResult, + ) -> None: + assert params.mamba_value is not None + if is_new_leaf: + node.component_data[self.component_type].value = params.mamba_value + self.cache.lru_lists[self.component_type].insert_mru(node) + self.cache.component_evictable_size_[self.component_type] += len( + params.mamba_value + ) + return + if node.component_data[self.component_type].value is None: + node.component_data[self.component_type].value = params.mamba_value + self.cache.lru_lists[self.component_type].insert_mru(node) + self.cache.component_evictable_size_[self.component_type] += len( + params.mamba_value + ) + node.last_access_time = get_and_increase_time_counter() + return + self.cache.lru_lists[self.component_type].reset_node_mru(node) + node.last_access_time = get_and_increase_time_counter() + result.mamba_exist = True + + def redistribute_on_node_split( + self, new_parent: UnifiedTreeNode, child: UnifiedTreeNode + ): + new_parent.component_data[self.component_type].value = None + new_parent.component_data[self.component_type].lock_ref = 0 + + def evict_component(self, node: UnifiedTreeNode, is_leaf: bool) -> int: + value = node.component_data[self.component_type].value + self.cache.req_to_token_pool.mamba_pool.free(value) + freed = len(value) + self.cache.component_evictable_size_[self.component_type] -= freed + if not is_leaf: + node.component_data[self.component_type].value = None + return freed + + def drive_eviction( + self, params: EvictParams, tracker: dict[ComponentType, int] + ) -> None: + request = params.mamba_num + lru = self.cache.lru_lists[self.component_type] + x = lru.get_lru_no_lock() + while ( + tracker[self.component_type] < request and x is not None and lru.in_list(x) + ): + assert x.component_data[self.component_type].value is not None + if len(x.children) > 0: + x_next = lru.get_prev_no_lock(x) + self.cache._evict_component_and_detach_lru( + x, self, is_leaf=False, tracker=tracker + ) + self.cache._cascade_evict(x, self, tracker) + x = x_next + else: + self.cache._evict_component_and_detach_lru( + x, self, is_leaf=True, tracker=tracker + ) + self.cache._cascade_evict(x, self, tracker) + x = lru.get_lru_no_lock() + + def acquire_component_lock( + self, node: UnifiedTreeNode, result: IncLockRefResult + ) -> IncLockRefResult: + ct = self.component_type + cd = node.component_data[ct] + value = cd.value + if value is not None: + if cd.lock_ref == 0: + vlen = len(value) + self.cache.component_evictable_size_[ct] -= vlen + self.cache.component_protected_size_[ct] += vlen + cd.lock_ref += 1 + return result + + def release_component_lock( + self, node: UnifiedTreeNode, params: Optional[DecLockRefParams] + ) -> None: + ct = self.component_type + cd = node.component_data[ct] + value = cd.value + if value is not None: + assert cd.lock_ref > 0 + if cd.lock_ref == 1: + vlen = len(value) + self.cache.component_evictable_size_[ct] += vlen + self.cache.component_protected_size_[ct] -= vlen + cd.lock_ref -= 1 + + def prepare_for_caching_req( + self, + req: Req, + insert_params: InsertParams, + token_ids_len: int, + is_finished: bool, + ) -> Optional[int]: + cache_len = ( + req.mamba_last_track_seqlen + if self.enable_mamba_extra_buffer + else token_ids_len + ) + if is_finished: + if cache_len is None: + cache_len = 0 + if self.enable_mamba_extra_buffer: + keep_idx = self.cache.req_to_token_pool.get_mamba_ping_pong_other_idx( + req.mamba_next_track_idx + ) + mamba_value = ( + req.mamba_ping_pong_track_buffer[keep_idx].unsqueeze(-1).clone() + ) + else: + mamba_value = req.mamba_pool_idx.unsqueeze(-1).clone() + insert_params.mamba_value = mamba_value + return cache_len + else: + if cache_len is None: + return 0 + if self.enable_mamba_extra_buffer: + keep_idx = self.cache.req_to_token_pool.get_mamba_ping_pong_other_idx( + req.mamba_next_track_idx + ) + mamba_value = ( + req.mamba_ping_pong_track_buffer[keep_idx].unsqueeze(-1).clone() + ) + else: + mamba_value = self.cache.req_to_token_pool.get_mamba_indices( + req.req_pool_idx + ).unsqueeze(-1) + mamba_value_forked = self.cache.req_to_token_pool.mamba_pool.fork_from( + mamba_value + ) + if mamba_value_forked is None: + self.cache.evict(EvictParams(num_tokens=0, mamba_num=1)) + mamba_value_forked = self.cache.req_to_token_pool.mamba_pool.fork_from( + mamba_value + ) + assert mamba_value_forked is not None, "Can not alloc mamba cache" + insert_params.mamba_value = mamba_value_forked + return cache_len + + def cleanup_after_caching_req( + self, + req: Req, + is_finished: bool, + insert_result: Optional[InsertResult] = None, + insert_params: Optional[InsertParams] = None, + ) -> None: + if is_finished: + mamba_exist = ( + insert_result.mamba_exist if insert_result is not None else True + ) + if self.enable_mamba_extra_buffer: + keep_idx = self.cache.req_to_token_pool.get_mamba_ping_pong_other_idx( + req.mamba_next_track_idx + ) + else: + keep_idx = None + if mamba_exist: + keep_idx = None + free_mamba_cache = True if self.enable_mamba_extra_buffer else mamba_exist + if free_mamba_cache: + self.cache.req_to_token_pool.free_mamba_cache( + req, mamba_ping_pong_track_buffer_to_keep=keep_idx + ) + else: + if insert_params.mamba_value is not None and ( + insert_result is None or insert_result.mamba_exist + ): + self.cache.req_to_token_pool.mamba_pool.free(insert_params.mamba_value) + req.mamba_last_track_seqlen = None diff --git a/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py b/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py new file mode 100644 index 000000000000..b38b8842f6f0 --- /dev/null +++ b/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py @@ -0,0 +1,297 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Callable, Optional + +import torch + +from sglang.srt.mem_cache.base_prefix_cache import ( + DecLockRefParams, + EvictParams, + IncLockRefResult, + InsertParams, + InsertResult, +) +from sglang.srt.mem_cache.unified_cache_components.tree_component import ( + BASE_COMPONENT_TYPE, + ComponentType, + TreeComponent, + next_component_uuid, +) + +if TYPE_CHECKING: + from sglang.srt.managers.schedule_batch import Req + from sglang.srt.mem_cache.cache_init_params import CacheInitParams + from sglang.srt.mem_cache.unified_radix_cache import ( + UnifiedRadixCache, + UnifiedTreeNode, + ) + + +class SWAComponent(TreeComponent): + """Sliding window attention component. + + Each SWA node stores translated SWA pool indices as its component + value, independent of the full attention indices on the same tree node. + When SWA data is evicted from an internal node the node is tombstoned + — its SWA component value becomes None while the full attention + value stays intact. + """ + + def __init__(self, cache: UnifiedRadixCache, params: CacheInitParams): + from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator + + assert isinstance( + cache.token_to_kv_pool_allocator, SWATokenToKVPoolAllocator + ), f"SWAComponent requires SWATokenToKVPoolAllocator, got {type(cache.token_to_kv_pool_allocator)}" + super().__init__(cache, params) + self.sliding_window_size = params.sliding_window_size + + component_type = ComponentType.SWA + + def _translate_full_to_swa(self, full_indices: torch.Tensor) -> torch.Tensor: + return self.cache.token_to_kv_pool_allocator.translate_loc_from_full_to_swa( + full_indices + ) + + def create_match_validator(self) -> Callable[[UnifiedTreeNode], bool]: + sliding_window_size = self.sliding_window_size + ct = self.component_type + state = {"len": float("inf")} + + def validator(node: UnifiedTreeNode) -> bool: + if node.component_data[ct].value is None: + state["len"] = 0 + return False + state["len"] += len(node.key) + return state["len"] >= sliding_window_size + + return validator + + def update_component_on_insert_overlap( + self, + node: UnifiedTreeNode, + prefix_len: int, + total_prefix_len: int, + value_slice: torch.Tensor, + params: InsertParams, + ) -> int: + if params.prev_prefix_len >= total_prefix_len + prefix_len: + return prefix_len + + is_tombstone = node.component_data[self.component_type].value is None + if not is_tombstone: + return prefix_len + + swa_evicted_seqlen = params.swa_evicted_seqlen + assert ( + node.component_data[self.component_type].lock_ref == 0 + ), f"tombstone {self.component_type} lock_ref should be 0, node {node.id}" + assert ( + swa_evicted_seqlen % self.cache.page_size == 0 + ), f"{self.component_type}: swa_evicted_seqlen must be page-aligned, {swa_evicted_seqlen=}" + + if swa_evicted_seqlen <= total_prefix_len: + # Branch 1: entire value_slice is within SWA window — recover + self.cache.token_to_kv_pool_allocator.free( + node.component_data[BASE_COMPONENT_TYPE].value + ) + node.component_data[BASE_COMPONENT_TYPE].value = value_slice.clone() + swa_value = self._translate_full_to_swa( + node.component_data[BASE_COMPONENT_TYPE].value + ) + node.component_data[self.component_type].value = swa_value + self.cache.lru_lists[self.component_type].insert_mru(node) + self.cache.component_evictable_size_[self.component_type] += len(swa_value) + return 0 + elif swa_evicted_seqlen < total_prefix_len + prefix_len: + # Branch 2: value_slice[start_idx:] is within SWA window — partial recover + start_idx = swa_evicted_seqlen - total_prefix_len + self.cache.token_to_kv_pool_allocator.free( + node.component_data[BASE_COMPONENT_TYPE].value[start_idx:] + ) + self.cache._split_node(node.key, node, start_idx) + node.component_data[BASE_COMPONENT_TYPE].value = value_slice[ + start_idx: + ].clone() + swa_value = self._translate_full_to_swa( + node.component_data[BASE_COMPONENT_TYPE].value + ) + node.component_data[self.component_type].value = swa_value + self.cache.lru_lists[self.component_type].insert_mru(node) + self.cache.component_evictable_size_[self.component_type] += len(swa_value) + return start_idx + else: + # Branch 3: entire value_slice is outside SWA window — not consumed + return prefix_len + + def should_skip_leaf_creation( + self, total_prefix_len: int, key_len: int, params: InsertParams + ) -> bool: + return params.swa_evicted_seqlen >= total_prefix_len + key_len + + def commit_insert_component_data( + self, + node: UnifiedTreeNode, + is_new_leaf: bool, + params: InsertParams, + result: InsertResult, + ) -> None: + if not is_new_leaf: + return + + node_start = result.prefix_len + split_pos = params.swa_evicted_seqlen - node_start + + if split_pos <= 0: + swa_value = self._translate_full_to_swa( + node.component_data[BASE_COMPONENT_TYPE].value + ) + node.component_data[self.component_type].value = swa_value + self.cache.lru_lists[self.component_type].insert_mru(node) + self.cache.component_evictable_size_[self.component_type] += len(swa_value) + elif split_pos < len(node.key): + # Node straddles the SWA eviction boundary + # Split into parent (tombstone, no SWA) and child (with SWA) + # After _split_node, `node` becomes the child + self.cache._split_node(node.key, node, split_pos) + swa_value = self._translate_full_to_swa( + node.component_data[BASE_COMPONENT_TYPE].value + ) + node.component_data[self.component_type].value = swa_value + self.cache.lru_lists[self.component_type].insert_mru(node) + self.cache.component_evictable_size_[self.component_type] += len(swa_value) + + def redistribute_on_node_split( + self, new_parent: UnifiedTreeNode, child: UnifiedTreeNode + ): + new_parent.component_data[self.component_type].lock_ref = child.component_data[ + self.component_type + ].lock_ref + + child_swa_value = child.component_data[self.component_type].value + if child_swa_value is not None: + split_len = len(new_parent.key) + new_parent.component_data[self.component_type].value = child_swa_value[ + :split_len + ].clone() + child.component_data[self.component_type].value = child_swa_value[ + split_len: + ].clone() + else: + new_parent.component_data[self.component_type].value = None + + # parent inherits the swa_uuid from child for swa lock ref + new_parent.component_data[self.component_type].metadata["uuid"] = ( + child.component_data[self.component_type].metadata.get("uuid") + ) + child.component_data[self.component_type].metadata.pop("uuid", None) + + def evict_component(self, node: UnifiedTreeNode, is_leaf: bool) -> int: + swa_value = node.component_data[self.component_type].value + if swa_value is None: + return 0 + # Direct swa_attn_allocator.free(swa_value) would double-free + # free_swa(full_value) has the mapping guard to avoid double-free + # TODO: decoupling full and swa free, need further discussion on mapping necessity + self.cache.token_to_kv_pool_allocator.free_swa( + node.component_data[BASE_COMPONENT_TYPE].value + ) + freed = len(swa_value) + self.cache.component_evictable_size_[self.component_type] -= freed + if not is_leaf: + node.component_data[self.component_type].value = None + return freed + + def eviction_priority(self, is_leaf: bool) -> int: + return 0 if is_leaf else 1 + + def drive_eviction( + self, params: EvictParams, tracker: dict[ComponentType, int] + ) -> None: + request = params.swa_num_tokens + lru = self.cache.lru_lists[self.component_type] + x = lru.get_lru_no_lock() + while ( + tracker[self.component_type] < request and x is not None and lru.in_list(x) + ): + assert x.component_data[self.component_type].value is not None + if len(x.children) > 0: + x_next = lru.get_prev_no_lock(x) + self.cache._evict_component_and_detach_lru( + x, self, is_leaf=False, tracker=tracker + ) + self.cache._cascade_evict(x, self, tracker) + x = x_next + else: + self.cache._evict_component_and_detach_lru( + x, self, is_leaf=True, tracker=tracker + ) + self.cache._cascade_evict(x, self, tracker) + x = lru.get_lru_no_lock() + + def acquire_component_lock( + self, node: UnifiedTreeNode, result: IncLockRefResult + ) -> IncLockRefResult: + ct = self.component_type + root = self.cache.root_node + sliding_window_size = self.sliding_window_size + swa_lock_size = 0 + swa_uuid_for_lock = None + + cur = node + while cur != root and swa_lock_size < sliding_window_size: + assert ( + cur.component_data[ct].value is not None + ), f"acquire_component_lock({ct}) on tombstoned node {cur.id}" + comp = cur.component_data[ct] + if comp.lock_ref == 0: + key_len = len(cur.key) + self.cache.component_evictable_size_[ct] -= key_len + self.cache.component_protected_size_[ct] += key_len + comp.lock_ref += 1 + swa_lock_size += len(cur.key) + if swa_lock_size >= sliding_window_size: + if comp.metadata.get("uuid") is None: + comp.metadata["uuid"] = next_component_uuid() + swa_uuid_for_lock = comp.metadata["uuid"] + cur = cur.parent + + result.swa_uuid_for_lock = swa_uuid_for_lock + return result + + def release_component_lock( + self, node: UnifiedTreeNode, params: Optional[DecLockRefParams] + ) -> None: + ct = self.component_type + root = self.cache.root_node + swa_uuid_for_lock = params.swa_uuid_for_lock if params else None + dec_swa = True + + cur = node + while cur != root and dec_swa: + assert ( + cur.component_data[ct].value is not None + ), f"release_component_lock({ct}) on tombstoned node {cur.id}" + comp = cur.component_data[ct] + assert ( + comp.lock_ref > 0 + ), f"release_component_lock({ct}) on node with lock_ref=0, node {cur.id}" + if comp.lock_ref == 1: + key_len = len(cur.key) + self.cache.component_evictable_size_[ct] += key_len + self.cache.component_protected_size_[ct] -= key_len + comp.lock_ref -= 1 + if swa_uuid_for_lock and comp.metadata.get("uuid") == swa_uuid_for_lock: + dec_swa = False + cur = cur.parent + + def prepare_for_caching_req( + self, + req: Req, + insert_params: InsertParams, + token_ids_len: int, + is_finished: bool, + ) -> Optional[int]: + if is_finished: + insert_params.swa_evicted_seqlen = req.swa_evicted_seqlen + return None diff --git a/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py b/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py new file mode 100644 index 000000000000..739db523b223 --- /dev/null +++ b/python/sglang/srt/mem_cache/unified_cache_components/tree_component.py @@ -0,0 +1,292 @@ +from __future__ import annotations + +import dataclasses +from abc import ABC, abstractmethod +from enum import Enum +from typing import TYPE_CHECKING, Any, Callable, Optional + +import torch +from numpy import float64 + +from sglang.srt.mem_cache.base_prefix_cache import ( + DecLockRefParams, + EvictParams, + IncLockRefResult, + InsertParams, + InsertResult, + MatchPrefixParams, + MatchResult, +) + +if TYPE_CHECKING: + from sglang.srt.managers.schedule_batch import Req + from sglang.srt.mem_cache.cache_init_params import CacheInitParams + from sglang.srt.mem_cache.unified_radix_cache import ( + UnifiedRadixCache, + UnifiedTreeNode, + ) + + +class ComponentType(int, Enum): + """Integer enum so that per-node list/tuple storage can be indexed directly.""" + + FULL = 0 + SWA = 1 + MAMBA = 2 + + def __str__(self) -> str: # keep human-readable logging + return self.name.lower() + + @property + def is_full(self) -> bool: + return self == ComponentType.FULL + + @property + def is_swa(self) -> bool: + return self == ComponentType.SWA + + @property + def is_mamba(self) -> bool: + return self == ComponentType.MAMBA + + +BASE_COMPONENT_TYPE = ComponentType.FULL +_NUM_COMPONENT_TYPES = len(ComponentType) + +_LAST_ACCESS_TIME_COUNTER_FLOAT = float64(1.0) +_COMPONENT_UUID_COUNTER = 1 + + +@dataclasses.dataclass +class ComponentData: + value: Optional[torch.Tensor] = None + lock_ref: int = 0 + metadata: dict[str, Any] = dataclasses.field(default_factory=dict) + + +def get_and_increase_time_counter() -> float64: + global _LAST_ACCESS_TIME_COUNTER_FLOAT + ret = _LAST_ACCESS_TIME_COUNTER_FLOAT + _LAST_ACCESS_TIME_COUNTER_FLOAT += 1.0 + return ret + + +def next_component_uuid() -> int: + global _COMPONENT_UUID_COUNTER + _COMPONENT_UUID_COUNTER += 1 + return _COMPONENT_UUID_COUNTER + + +class TreeComponent(ABC): + def __init__(self, cache: UnifiedRadixCache, params: CacheInitParams): + self.cache = cache + + # Subclasses MUST set this as a class attribute (not @property) + component_type: ComponentType + + def node_has_component_data(self, node: UnifiedTreeNode) -> bool: + return node.component_data[self.component_type].value is not None + + def value_len(self, node: UnifiedTreeNode) -> int: + value = node.component_data[self.component_type].value + return len(value) if value is not None else 0 + + @abstractmethod + def create_match_validator(self) -> Callable[[UnifiedTreeNode], bool]: + """Return a per-match stateful predicate that decides whether a node + is a valid match boundary for this component. + Called once per match_prefix; the returned closure may carry state. + - Full: always True (every node is valid). + - SWA: tracks accumulated length since last gap; returns True only + when the contiguous window reaches swa_sliding_window_size. + - Mamba: returns True iff the node has mamba component data.""" + ... + + def finalize_match_result( + self, + result: MatchResult, + params: MatchPrefixParams, + value_chunks: list[torch.Tensor], + best_value_len: int, + ) -> MatchResult: + """Post-process the match result after prefix matching completes. + - Full & SWA: pass through unchanged. + - Mamba: performs copy-on-write — allocates a new mamba slot, copies + the matched node's mamba state into the request pool, and records + branching_seqlen in result.""" + return result + + def update_component_on_insert_overlap( + self, + node: UnifiedTreeNode, + prefix_len: int, + total_prefix_len: int, + value_slice: torch.Tensor, + params: InsertParams, + ) -> int: + """Called per-node when an insert's key overlaps an existing node. + Returns the index within value_slice from which this component + consumed (took ownership of) the underlying KV pool slots. + Returns prefix_len if nothing was consumed (default). + _insert_helper uses this to free only the non-consumed duplicate + portion: value_slice[dup_start:consumed_from].""" + return prefix_len + + def should_skip_leaf_creation( + self, total_prefix_len: int, key_len: int, params: InsertParams + ) -> bool: + """Return True to veto leaf creation when the entire new leaf would + be a tombstone for this component.""" + return False + + def commit_insert_component_data( + self, + node: UnifiedTreeNode, + is_new_leaf: bool, + params: InsertParams, + result: InsertResult, + ) -> None: + """Finalize component data on the target (leaf) node after the insert + walk completes. Called once per insert. + - Full: no-op (full data is handled by _add_new_node). + - SWA: for new leaves, checks whether the node straddles the SWA + eviction boundary (swa_evicted_seqlen). If so, splits the node + via _split_node — the parent becomes a tombstone (no SWA) and the + child (the deeper portion) receives SWA data. If the entire node + is within the window, sets SWA directly. If entirely outside, + leaves SWA as None (tombstone). + - Mamba: sets the mamba component value from params, inserts into + mamba LRU list, and increments evictable size. If the node already + has mamba data, resets its LRU position instead.""" + pass + + @abstractmethod + def redistribute_on_node_split( + self, new_parent: UnifiedTreeNode, child: UnifiedTreeNode + ): + """Redistribute component data between new_parent and child when a + node is split. new_parent is the newly created prefix node. + - Full: copies child's lock_ref to new_parent. + - SWA: slices (or clones) the swa value for new_parent, copies + lock_ref and component_uuid metadata, then syncs child's swa + value with its (now-trimmed) full_value. + - Mamba: sets new_parent's mamba value to None and lock_ref to 0 + (mamba data stays on the original leaf, not on prefix nodes).""" + ... + + @abstractmethod + def evict_component(self, node: UnifiedTreeNode, is_leaf: bool) -> int: + """Free this component's KV resources on a node being evicted. + For internal (non-leaf) nodes: free memory and tombstone the value + (set to None); the node structure is kept. + For leaf nodes: free memory; the node will be deleted by caller. + Returns the number of tokens/slots freed. + - Full: frees full_value via token_to_kv_pool_allocator. + - SWA: frees swa value via swa_token_to_kv_pool_allocator; + only tombstones on internal nodes. + - Mamba: frees mamba value via mamba_token_to_kv_pool_allocator; + only tombstones on internal nodes.""" + ... + + def eviction_priority(self, is_leaf: bool) -> int: + """Eviction priority on this node type. Higher = evicted later. + When a component is evicted, all other components with equal or + lower priority on the same node are also cascade-evicted. + + Leaf: all components equal (0) — evicting any cascades to all, + because the node will be deleted. + + Internal: full=2 > swa=1 > mamba=0. + Why swa > mamba: SWA data on internal nodes is *path data* — + the sliding window needs continuous SWA coverage along the path + from root to the match boundary. E.g. A->B->C->D->E where C + and E both have mamba and the window covers C->E: if C's mamba + is evicted, C's SWA must stay so E remains reachable. + Mamba data, by contrast, is only meaningful at the match + boundary node; on internal nodes it + contributes nothing to the path. So SWA is more valuable to + keep and should be evicted later. + + Cascade consequences: + - Mamba evict internal: no cascade. + - SWA evict internal: cascades to Mamba. SWA gone -> SWA + validator fails -> mamba data is useless (match requires all + validators to pass). + - Full evict internal: cascades to SWA + Mamba.""" + return 0 + + @abstractmethod + def drive_eviction( + self, params: EvictParams, tracker: dict[ComponentType, int] + ) -> None: + """Drive eviction from this component's LRU list. + Each component extracts its own request from params, walks its own + LRU, evicts, and calls cache._cascade_evict for priority cascade. + Updates the shared tracker with freed amounts for all components. + - Full: walks leaf LRU, evicts full then cascades entire leaf. + - Mamba: walks full LRU; tombstones internal nodes (with cascade + to equal-priority components like swa), cascades leaves to all.""" + ... + + @abstractmethod + def acquire_component_lock( + self, node: UnifiedTreeNode, result: IncLockRefResult + ) -> IncLockRefResult: + """Increment lock_ref for this component, protecting nodes from + eviction. Updates evictable → protected size on first lock. + - Full: path-lock — walks from node up to root, incrementing + lock_ref on every ancestor. + - SWA: path-lock — walks upward collecting swa values until the + sliding window is filled; records a component_uuid at the + boundary for release_component_lock to know where to stop. + - Mamba: single-node lock — only increments lock_ref on the + node itself (mamba state is per-leaf, not per-path).""" + ... + + @abstractmethod + def release_component_lock( + self, node: UnifiedTreeNode, params: Optional[DecLockRefParams] + ) -> None: + """Decrement lock_ref for this component, un-protecting nodes. + Updates protected → evictable size when lock_ref drops to 0. + - Full: path-unlock — walks from node up to root, decrementing + lock_ref on every ancestor. + - SWA: path-unlock — walks upward, stopping at the node whose + component_uuid matches the one recorded during acquire. + - Mamba: single-node unlock — only decrements lock_ref on the + node itself.""" + ... + + def prepare_for_caching_req( + self, + req: Req, + insert_params: InsertParams, + token_ids_len: int, + is_finished: bool, + ) -> Optional[int]: + """Prepare component-specific data before insert, fill component + fields in insert_params, return effective cache_len. + Return None for no truncation opinion (use full length); + return int >= 0 for effective cache length. + - Full: no-op, returns None. + - SWA: sets insert_params.swa_evicted_seqlen on finished; returns None. + - Mamba: prepares mamba_value (finished from ping-pong buffer, + unfinished fork from req); returns mamba_last_track_seqlen.""" + return None + + def cleanup_after_caching_req( + self, + req: Req, + is_finished: bool, + insert_result: Optional[InsertResult] = None, + insert_params: Optional[InsertParams] = None, + ) -> None: + """Post-cache cleanup for component-specific resources. + + ``is_finished`` — whether the request has finished generation. + True means the request is complete and its resources can be released; + ``insert_result`` is None when insert was skipped (cache disabled + or effective_cache_len <= 0); treat as "no insert happened". + ``insert_params`` is None only on the disabled path; on early-return + paths it is still provided so components can free their resources.""" + pass diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py new file mode 100644 index 000000000000..73b645a3b216 --- /dev/null +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -0,0 +1,986 @@ +from __future__ import annotations + +import logging +import time +from collections import defaultdict +from functools import partial +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.srt.mem_cache.base_prefix_cache import ( + BasePrefixCache, + DecLockRefParams, + DecLockRefResult, + EvictParams, + EvictResult, + IncLockRefResult, + InsertParams, + InsertResult, + MatchPrefixParams, + MatchResult, +) +from sglang.srt.mem_cache.radix_cache import ( + RadixKey, + _key_match_page_size1, + _key_match_paged, + get_child_key, + maybe_bigram_convert, + page_align_keys, +) +from sglang.srt.mem_cache.unified_cache_components import ( + _NUM_COMPONENT_TYPES, + BASE_COMPONENT_TYPE, + ComponentData, + ComponentType, + FullComponent, + MambaComponent, + SWAComponent, + TreeComponent, + get_and_increase_time_counter, +) +from sglang.srt.mem_cache.utils import convert_to_bigram_key + +if TYPE_CHECKING: + from sglang.srt.managers.schedule_batch import Req + from sglang.srt.mem_cache.cache_init_params import CacheInitParams + + +class UnifiedTreeNode: + counter = 0 + + def __init__(self, tree_components: tuple[ComponentType, ...]): + self.children = defaultdict(partial(UnifiedTreeNode, tree_components)) + self.parent: UnifiedTreeNode | None = None + self.key: Optional[RadixKey] = None + self.tree_components = tree_components + # list indexed by ComponentType (int enum 0..N-1) + self.component_data: list[ComponentData] = [ + ComponentData() for _ in range(_NUM_COMPONENT_TYPES) + ] + self.last_access_time = get_and_increase_time_counter() + self.host_value = None + self.hit_count = 0 + self.lru_prev: list[UnifiedTreeNode | None] = [None] * _NUM_COMPONENT_TYPES + self.lru_next: list[UnifiedTreeNode | None] = [None] * _NUM_COMPONENT_TYPES + self.id = UnifiedTreeNode.counter + UnifiedTreeNode.counter += 1 + + def component(self, component_type: ComponentType) -> ComponentData: + return self.component_data[component_type] + + def __lt__(self, other: UnifiedTreeNode): + return self.last_access_time < other.last_access_time + + +class UnifiedLRUList: + def __init__( + self, component_type: ComponentType, tree_components: tuple[ComponentType, ...] + ): + self.component_type = component_type + self.head = UnifiedTreeNode(tree_components) + self.tail = UnifiedTreeNode(tree_components) + self.head.lru_next[component_type] = self.tail + self.tail.lru_prev[component_type] = self.head + self.cache: dict[int, UnifiedTreeNode] = {} + + def _add_node_after(self, prev_node: UnifiedTreeNode, new_node: UnifiedTreeNode): + ct = self.component_type + new_node.lru_prev[ct] = prev_node + new_node.lru_next[ct] = prev_node.lru_next[ct] + prev_node.lru_next[ct].lru_prev[ct] = new_node + prev_node.lru_next[ct] = new_node + + def _add_node(self, node: UnifiedTreeNode): + self._add_node_after(self.head, node) + + def _remove_node(self, node: UnifiedTreeNode): + ct = self.component_type + node.lru_prev[ct].lru_next[ct] = node.lru_next[ct] + node.lru_next[ct].lru_prev[ct] = node.lru_prev[ct] + + def insert_mru(self, node: UnifiedTreeNode): + assert node.id not in self.cache + self.cache[node.id] = node + self._add_node(node) + + def remove_node(self, node: UnifiedTreeNode): + assert node.id in self.cache + del self.cache[node.id] + self._remove_node(node) + + def reset_node_mru(self, node: UnifiedTreeNode): + assert node.id in self.cache + self._remove_node(node) + self._add_node(node) + + def reset_node_and_parents_mru( + self, + node: UnifiedTreeNode, + root_node: UnifiedTreeNode, + should_include, + ): + prev_node = self.head + while node != root_node: + if should_include(node): + assert node.id in self.cache + self._remove_node(node) + self._add_node_after(prev_node, node) + prev_node = node + node = node.parent + + def in_list(self, node: Optional[UnifiedTreeNode]): + return node is not None and node.id in self.cache + + def get_prev_no_lock(self, node: UnifiedTreeNode, check_id: bool = True): + if check_id: + assert node.id in self.cache + ct = self.component_type + x = node.lru_prev[ct] + while x.component_data[ct].lock_ref > 0: + x = x.lru_prev[ct] + if x == self.head: + return None + return x + + def get_prev_leaf_no_lock(self, node: UnifiedTreeNode, check_id: bool = True): + if check_id: + assert node.id in self.cache + ct = self.component_type + x = node.lru_prev[ct] + while x.component_data[ct].lock_ref > 0 or len(x.children) > 0: + x = x.lru_prev[ct] + if x == self.head: + return None + return x + + def get_lru_no_lock(self): + return self.get_prev_no_lock(self.tail, check_id=False) + + def get_leaf_lru_no_lock(self): + return self.get_prev_leaf_no_lock(self.tail, check_id=False) + + +COMPONENT_REGISTRY: dict[ComponentType, type[TreeComponent]] = { + ComponentType.FULL: FullComponent, + ComponentType.MAMBA: MambaComponent, + ComponentType.SWA: SWAComponent, +} + +logger = logging.getLogger(__name__) + + +class UnifiedRadixCache(BasePrefixCache): + def __init__( + self, + params: CacheInitParams, + ): + self.req_to_token_pool = params.req_to_token_pool + self.token_to_kv_pool_allocator = params.token_to_kv_pool_allocator + self.page_size = params.page_size + self.disable = params.disable + self.is_eagle = params.is_eagle + + if self.token_to_kv_pool_allocator: + self.device = self.token_to_kv_pool_allocator.device + else: + self.device = torch.device("cpu") + + if params.enable_metrics: + self.init_metrics_collector() + + if self.page_size == 1: + self.key_match_fn = _key_match_page_size1 + self.get_child_key_fn = get_child_key + else: + self.key_match_fn = partial(_key_match_paged, page_size=self.page_size) + self.get_child_key_fn = partial(get_child_key, page_size=self.page_size) + + assert params.tree_components is not None + self.tree_components = tuple(params.tree_components) + self.components: dict[ComponentType, TreeComponent] = { + ct: COMPONENT_REGISTRY[ct](self, params) for ct in self.tree_components + } + self._components_tuple: tuple[TreeComponent, ...] = tuple( + self.components.values() + ) + if self.is_eagle: + self.key_convert_fn = convert_to_bigram_key + else: + self.key_convert_fn = lambda key: key + self.reset() + logger.info(f"Init Unified RadixTree with components {self.tree_components}") + + def reset(self) -> None: + self.root_node = UnifiedTreeNode(self.tree_components) + self.root_node.key = RadixKey([], None) + self.root_node.component_data[BASE_COMPONENT_TYPE].value = [] + for ct in self.tree_components: + self.root_node.component_data[ct].lock_ref = 1 + self.component_evictable_size_ = {ct: 0 for ct in self.tree_components} + self.component_protected_size_ = {ct: 0 for ct in self.tree_components} + self.lru_lists = { + ct: UnifiedLRUList(ct, self.tree_components) for ct in self.tree_components + } + + def match_prefix(self, params: MatchPrefixParams) -> MatchResult: + key = params.key + key, _ = maybe_bigram_convert(self.is_eagle, key) + if self.disable or len(key) == 0: + return MatchResult( + device_indices=torch.empty( + (0,), + dtype=torch.int64, + device=self.device, + ), + last_device_node=self.root_node, + last_host_node=self.root_node, + ) + if self.page_size != 1: + page_aligned_len = len(key) // self.page_size * self.page_size + key = key[:page_aligned_len] + + value, last_node, best_value_len = self._match_prefix_helper(key) + return self._match_post_processor(params, value, last_node, best_value_len) + + def insert(self, params: InsertParams) -> InsertResult: + if self.disable: + return InsertResult(prefix_len=0) + + key = params.key + value = params.value + if value is None: + value = torch.tensor([x for x in key.token_ids], dtype=torch.int64) + + key, value = maybe_bigram_convert(self.is_eagle, key, value) + result = self._insert_helper(self.root_node, key, value, params) + return result + + def evict(self, params: EvictParams) -> EvictResult: + if self.disable: + return EvictResult() + start_time = time.perf_counter() + tracker = {ct: 0 for ct in self.tree_components} + + for component in self._components_tuple: + component.drive_eviction(params=params, tracker=tracker) + + self.update_eviction_metrics(sum(tracker.values()), start_time) + return EvictResult( + num_tokens_evicted=tracker[BASE_COMPONENT_TYPE], + swa_num_tokens_evicted=tracker.get(ComponentType.SWA, 0), + mamba_num_evicted=tracker.get(ComponentType.MAMBA, 0), + ) + + def inc_lock_ref(self, node: UnifiedTreeNode) -> IncLockRefResult: + if self.disable: + return IncLockRefResult() + result = IncLockRefResult() + for component in self._components_tuple: + result = component.acquire_component_lock(node=node, result=result) + return result + + def dec_lock_ref( + self, node: UnifiedTreeNode, params: Optional[DecLockRefParams] = None + ) -> DecLockRefResult: + if self.disable: + return DecLockRefResult() + for component in self._components_tuple: + component.release_component_lock(node=node, params=params) + # TODO: delta is not aggregated from components; no caller uses it yet. + return DecLockRefResult() + + def cache_finished_req(self, req: Req, is_insert: bool = True) -> None: + kv_committed_len = req.pop_committed_kv_cache() + + if self.disable: + kv_indices = self.req_to_token_pool.req_to_token[ + req.req_pool_idx, :kv_committed_len + ] + self.token_to_kv_pool_allocator.free(kv_indices) + for comp in self._components_tuple: + comp.cleanup_after_caching_req(req, is_finished=True) + return + + token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len] + kv_indices = self.req_to_token_pool.req_to_token[ + req.req_pool_idx, :kv_committed_len + ] + + result = None + insert_params = None + + if is_insert: + insert_params = InsertParams(prev_prefix_len=req.cache_protected_len) + + # components prepare insert data + return effective cache_len + effective_cache_len = len(token_ids) + for comp in self._components_tuple: + cl = comp.prepare_for_caching_req( + req=req, + insert_params=insert_params, + token_ids_len=len(token_ids), + is_finished=True, + ) + if cl is not None: + effective_cache_len = min(effective_cache_len, cl) + + # Truncate if needed + if effective_cache_len < len(token_ids): + free_start = max(effective_cache_len, req.cache_protected_len) + self.token_to_kv_pool_allocator.free(kv_indices[free_start:]) + token_ids = token_ids[:effective_cache_len] + kv_indices = kv_indices[:effective_cache_len] + + # Key convert + page align + keys = self.key_convert_fn(token_ids) + keys = page_align_keys(keys, self.page_size) + page_aligned_len = len(keys) + values = kv_indices[:page_aligned_len].to(dtype=torch.int64, copy=True) + radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle) + + insert_params.key = radix_key + insert_params.value = values + result = self.insert(insert_params) + + # Free unaligned tail + self.token_to_kv_pool_allocator.free(kv_indices[page_aligned_len:]) + else: + self.token_to_kv_pool_allocator.free(kv_indices[req.cache_protected_len :]) + + self.dec_lock_ref( + req.last_node, + DecLockRefParams(swa_uuid_for_lock=getattr(req, "swa_uuid_for_lock", None)), + ) + + # cleanup + for comp in self._components_tuple: + comp.cleanup_after_caching_req( + req, is_finished=True, insert_result=result, insert_params=insert_params + ) + + def cache_unfinished_req(self, req: Req, chunked=False) -> None: + token_ids = req.fill_ids + + if self.disable: + kv_indices = self.req_to_token_pool.req_to_token[ + req.req_pool_idx, : len(token_ids) + ] + req.prefix_indices = kv_indices + return + + kv_indices_orig = self.req_to_token_pool.req_to_token[ + req.req_pool_idx, : len(token_ids) + ] + + # components prepare insert data + return effective cache_len + insert_params = InsertParams(prev_prefix_len=req.cache_protected_len) + effective_cache_len = len(token_ids) + for comp in self._components_tuple: + cl = comp.prepare_for_caching_req( + req=req, + insert_params=insert_params, + token_ids_len=len(token_ids), + is_finished=False, + ) + if cl is not None: + effective_cache_len = min(effective_cache_len, cl) + + if effective_cache_len <= 0: + req.prefix_indices = kv_indices_orig.to(dtype=torch.int64, copy=True) + for comp in self._components_tuple: + comp.cleanup_after_caching_req( + req, is_finished=False, insert_params=insert_params + ) + return + + kv_indices = kv_indices_orig[:effective_cache_len] + + # Key convert + page align + keys = self.key_convert_fn(token_ids[:effective_cache_len]) + keys = page_align_keys(keys, self.page_size) + page_aligned_len = len(keys) + values = kv_indices[:page_aligned_len].to(dtype=torch.int64, copy=True) + radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle) + + insert_params.key = radix_key + insert_params.value = values + result = self.insert(insert_params) + + # Match prefix + match_result = self.match_prefix(MatchPrefixParams(key=radix_key)) + new_indices = match_result.device_indices + new_last_node = match_result.last_device_node + new_prefix_len = result.prefix_len + assert ( + req.cache_protected_len <= len(new_indices) + self.page_size - 1 + ), f"{req.cache_protected_len=}, {len(new_indices)=}, {page_aligned_len=}" + assert new_prefix_len <= len( + new_indices + ), f"{new_prefix_len=}, {len(new_indices)=}" + self.req_to_token_pool.write( + (req.req_pool_idx, slice(req.cache_protected_len, len(new_indices))), + new_indices[req.cache_protected_len :], + ) + + self.dec_lock_ref( + req.last_node, + DecLockRefParams(swa_uuid_for_lock=getattr(req, "swa_uuid_for_lock", None)), + ) + lock_result = self.inc_lock_ref(new_last_node) + + # Update req fields + if len(new_indices) < len(kv_indices_orig): + req.prefix_indices = torch.cat( + [new_indices, kv_indices_orig[len(new_indices) :]] + ) + else: + req.prefix_indices = new_indices + req.cache_protected_len = len(new_indices) + req.last_node = new_last_node + req.swa_uuid_for_lock = lock_result.swa_uuid_for_lock + + # cleanup + for comp in self._components_tuple: + comp.cleanup_after_caching_req( + req, + is_finished=False, + insert_result=result, + insert_params=insert_params, + ) + + # ---- Internal Helpers ---- + + def _match_prefix_helper_readonly( + self, key: RadixKey + ) -> tuple[list[torch.Tensor], UnifiedTreeNode, int]: + """Read-only version of _match_prefix_helper that does not split nodes. + Only considers fully matched nodes, ignores partial matches. + + Not used yet; reserved for future read-only match operations.""" + node = self.root_node + child_key = self.get_child_key_fn(key) + value: list[torch.Tensor] = [] + best_value_len = 0 + best_node = node + validators = tuple( + comp.create_match_validator() for comp in self._components_tuple + ) + + def _update_best_if_valid(node): + nonlocal best_value_len, best_node + if all(v(node) for v in validators): + best_value_len = len(value) + best_node = node + + while len(key) > 0 and child_key in node.children: + child = node.children[child_key] + prefix_len = self.key_match_fn(child.key, key) + if prefix_len < len(child.key): + # Read-only: do not split, ignore partial match and stop + break + value.append(child.component_data[BASE_COMPONENT_TYPE].value) + node = child + _update_best_if_valid(node) + key = key[prefix_len:] + if len(key): + child_key = self.get_child_key_fn(key) + return value, best_node, best_value_len + + def _match_prefix_helper( + self, key: RadixKey + ) -> tuple[list[torch.Tensor], UnifiedTreeNode, int]: + node = self.root_node + child_key = self.get_child_key_fn(key) + value: list[torch.Tensor] = [] + best_value_len = 0 + best_node = node + validators = tuple( + comp.create_match_validator() for comp in self._components_tuple + ) + + def _update_best_if_valid(node): + nonlocal best_value_len, best_node + if all(v(node) for v in validators): + best_value_len = len(value) + best_node = node + + while len(key) > 0 and child_key in node.children: + child = node.children[child_key] + prefix_len = self.key_match_fn(child.key, key) + if prefix_len < len(child.key): + node = self._split_node(child.key, child, prefix_len) + value.append(node.component_data[BASE_COMPONENT_TYPE].value) + _update_best_if_valid(node) + break + value.append(child.component_data[BASE_COMPONENT_TYPE].value) + node = child + _update_best_if_valid(node) + key = key[prefix_len:] + if len(key): + child_key = self.get_child_key_fn(key) + return value, best_node, best_value_len + + def _match_post_processor( + self, + params: MatchPrefixParams, + value: list[torch.Tensor], + last_node: UnifiedTreeNode, + best_value_len: int, + ) -> MatchResult: + node_update = last_node + for comp in self._components_tuple: + self.lru_lists[comp.component_type].reset_node_and_parents_mru( + node_update, self.root_node, comp.node_has_component_data + ) + cur_time = get_and_increase_time_counter() + while node_update: + node_update.last_access_time = cur_time + cur_time -= 0.00001 + node_update = node_update.parent + + if best_value_len > 0: + device_indices = torch.cat(value[:best_value_len]) + else: + device_indices = torch.empty((0,), dtype=torch.int64, device=self.device) + result = MatchResult( + device_indices=device_indices, + last_device_node=last_node, + last_host_node=last_node, + ) + + for component in self._components_tuple: + result = component.finalize_match_result( + result=result, + params=params, + value_chunks=value, + best_value_len=best_value_len, + ) + return result + + def _split_node( + self, key: RadixKey, child: UnifiedTreeNode, split_len: int + ) -> UnifiedTreeNode: + new_node = UnifiedTreeNode(self.tree_components) + new_node.children = {self.get_child_key_fn(key[split_len:]): child} + new_node.parent = child.parent + new_node.key = child.key[:split_len] + new_node.component_data[BASE_COMPONENT_TYPE].value = ( + child.component_data[BASE_COMPONENT_TYPE].value[:split_len].clone() + ) + + self._for_each_component_lru(child, UnifiedLRUList.remove_node) + + child.parent = new_node + child.key = child.key[split_len:] + child.component_data[BASE_COMPONENT_TYPE].value = ( + child.component_data[BASE_COMPONENT_TYPE].value[split_len:].clone() + ) + + for component in self._components_tuple: + component.redistribute_on_node_split(new_parent=new_node, child=child) + new_node.parent.children[self.get_child_key_fn(key)] = new_node + + self._for_each_component_lru(new_node, UnifiedLRUList.insert_mru) + self._for_each_component_lru(child, UnifiedLRUList.insert_mru) + child.last_access_time = get_and_increase_time_counter() + return new_node + + def _touch_node(self, node: UnifiedTreeNode): + node.last_access_time = get_and_increase_time_counter() + if node != self.root_node: + self._for_each_component_lru(node, UnifiedLRUList.reset_node_mru) + + def _add_new_node( + self, + parent: UnifiedTreeNode, + key: RadixKey, + value: torch.Tensor, + ) -> UnifiedTreeNode: + new_node = UnifiedTreeNode(self.tree_components) + new_node.parent = parent + new_node.key = key + new_node.component_data[BASE_COMPONENT_TYPE].value = value.clone() + parent.children[self.get_child_key_fn(key)] = new_node + self.lru_lists[BASE_COMPONENT_TYPE].insert_mru(new_node) + self.component_evictable_size_[BASE_COMPONENT_TYPE] += len(value) + return new_node + + def _insert_helper( + self, + node: UnifiedTreeNode, + key: RadixKey, + value: torch.Tensor, + params: InsertParams, + ) -> InsertResult: + self._touch_node(node) + if len(key) == 0: + return InsertResult(prefix_len=0, mamba_exist=True) + + child_key = self.get_child_key_fn(key) + total_prefix_length = 0 + while len(key) > 0 and child_key in node.children: + node = node.children[child_key] + self._touch_node(node) + prefix_len = self.key_match_fn(node.key, key) + if prefix_len < len(node.key): + node = self._split_node(node.key, node, prefix_len) + + value_slice = value[:prefix_len] + consumed_from = prefix_len + # Let each component claim ownership of overlapping KV slots + for component in self._components_tuple: + comp_consumed_from = component.update_component_on_insert_overlap( + node=node, + prefix_len=prefix_len, + total_prefix_len=total_prefix_length, + value_slice=value_slice, + params=params, + ) + consumed_from = min(consumed_from, comp_consumed_from) + + dup_start = max(0, params.prev_prefix_len - total_prefix_length) + if dup_start < consumed_from: + self.token_to_kv_pool_allocator.free( + value_slice[dup_start:consumed_from] + ) + + total_prefix_length += prefix_len + key = key[prefix_len:] + value = value[prefix_len:] + if len(key): + child_key = self.get_child_key_fn(key) + + is_new_leaf = False + # Create new leaf for remaining suffix + if len(key): + if any( + comp.should_skip_leaf_creation( + total_prefix_len=total_prefix_length, + key_len=len(key), + params=params, + ) + for comp in self._components_tuple + ): + # TODO: When leaf creation is skipped, We should release all component + # resources here or propagate a flag so that + # cleanup_after_caching_req can free them properly. + self.token_to_kv_pool_allocator.free(value) + return InsertResult(prefix_len=total_prefix_length) + target_node = self._add_new_node(node, key, value) + is_new_leaf = True + else: + target_node = node + + # Finalize: let each component attach its data to the target node. + # e.g. Mamba attaches mamba_value to the leaf node + result = InsertResult(prefix_len=total_prefix_length) + for component in self._components_tuple: + component.commit_insert_component_data( + node=target_node, + is_new_leaf=is_new_leaf, + params=params, + result=result, + ) + return result + + def _cascade_evict( + self, + node: UnifiedTreeNode, + trigger: TreeComponent, + tracker: dict[ComponentType, int], + ): + """Cascade eviction from trigger to lower-or-equal priority components. + + When a component evicts a node, all other components with equal or + lower eviction_priority on the same node are also evicted. + If the node is a leaf, it is removed from the tree and any + resulting tombstone ancestors are cleaned up recursively.""" + is_leaf = len(node.children) == 0 + trigger_priority = trigger.eviction_priority(is_leaf) + + for comp in self._components_tuple: + if comp.eviction_priority(is_leaf) <= trigger_priority: + if comp is not trigger and comp.node_has_component_data(node): + assert node.component_data[comp.component_type].lock_ref == 0 + self._evict_component_and_detach_lru( + node, comp, is_leaf=is_leaf, tracker=tracker + ) + + if is_leaf: + self._remove_leaf_from_parent(node) + self._iteratively_delete_tombstone_leaf(node, tracker) + + def _remove_leaf_from_parent(self, node: UnifiedTreeNode): + key = self.get_child_key_fn(node.key) + v = node.parent.children.pop(key, None) + assert v == node + + def _evict_component_and_detach_lru( + self, + node: UnifiedTreeNode, + comp: TreeComponent, + is_leaf: bool, + tracker: dict[ComponentType, int], + ) -> int: + freed = comp.evict_component(node, is_leaf=is_leaf) + tracker[comp.component_type] += freed + lru = self.lru_lists[comp.component_type] + if lru.in_list(node): + lru.remove_node(node) + return freed + + def _iteratively_delete_tombstone_leaf( + self, deleted_node: UnifiedTreeNode, tracker: dict[ComponentType, int] + ): + """After a leaf is removed, walk up the parent chain and delete + any ancestor that is leaf node and has lost any component data (tombstoned).""" + cur = deleted_node.parent + while cur != self.root_node and len(cur.children) == 0: + has_tombstone = any( + not comp.node_has_component_data(cur) + for comp in self.components.values() + ) + if not has_tombstone: + break + + if any( + cur.component_data[comp.component_type].lock_ref > 0 + for comp in self.components.values() + if comp.node_has_component_data(cur) + ): + break + + for comp in self.components.values(): + if comp.node_has_component_data(cur): + self._evict_component_and_detach_lru( + cur, comp, is_leaf=True, tracker=tracker + ) + self._remove_leaf_from_parent(cur) + cur = cur.parent + + def _for_each_component_lru(self, node: UnifiedTreeNode, lru_op): + for ct in self.tree_components: + if node.component_data[ct].value is not None: + lru_op(self.lru_lists[ct], node) + + # ---- Query / Inspection APIs ---- + # These APIs exist for compatibility with other RadixTree implementations. + # TODO: simplify and consolidate in a future refactor. + + @property + def sliding_window_size(self): + swa = self.components.get(ComponentType.SWA) + return swa.sliding_window_size if swa else None + + def supports_swa(self) -> bool: + return ComponentType.SWA in self.components + + def supports_mamba(self) -> bool: + return ComponentType.MAMBA in self.components + + def evictable_size(self) -> int: + return self.component_evictable_size_.get(BASE_COMPONENT_TYPE, 0) + + def protected_size(self) -> int: + return self.component_protected_size_.get(BASE_COMPONENT_TYPE, 0) + + def full_evictable_size(self) -> int: + return self.evictable_size() + + def full_protected_size(self) -> int: + return self.protected_size() + + def swa_evictable_size(self) -> int: + return self.component_evictable_size_.get(ComponentType.SWA, 0) + + def mamba_evictable_size(self) -> int: + return self.component_evictable_size_.get(ComponentType.MAMBA, 0) + + def swa_protected_size(self) -> int: + return self.component_protected_size_.get(ComponentType.SWA, 0) + + def mamba_protected_size(self) -> int: + return self.component_protected_size_.get(ComponentType.MAMBA, 0) + + def total_size(self): + total_size = 0 + total_aux_size = 0 + stack = [self.root_node] + while stack: + node = stack.pop() + total_size += len(node.component_data[BASE_COMPONENT_TYPE].value) + for ct in self.tree_components: + if ct == BASE_COMPONENT_TYPE: + continue + value = node.component_data[ct].value + if value is not None: + total_aux_size += len(value) + for child in node.children.values(): + stack.append(child) + return total_size, total_aux_size + + def all_values_flatten(self) -> torch.Tensor: + values = [] + + def _dfs(node: UnifiedTreeNode): + for child in node.children.values(): + values.append(child.component_data[BASE_COMPONENT_TYPE].value) + _dfs(child) + + _dfs(self.root_node) + if values: + return torch.cat(values) + return torch.tensor([], dtype=torch.int64, device=self.device) + + def _all_component_values_flatten( + self, component_type: ComponentType + ) -> torch.Tensor: + if component_type not in self.components: + return torch.tensor([], dtype=torch.int64, device=self.device) + + values = [] + + def _dfs(node: UnifiedTreeNode): + value = node.component_data[component_type].value + if value is not None: + values.append(value) + for child in node.children.values(): + _dfs(child) + + _dfs(self.root_node) + if values: + return torch.cat(values) + return torch.tensor([], dtype=torch.int64, device=self.device) + + def all_mamba_values_flatten(self) -> torch.Tensor: + return self._all_component_values_flatten(ComponentType.MAMBA) + + def all_swa_values_flatten(self) -> torch.Tensor: + return self._all_component_values_flatten(ComponentType.SWA) + + def available_and_evictable_str(self) -> str: + if self.supports_swa(): + full_available_size = self.token_to_kv_pool_allocator.full_available_size() + else: + full_available_size = self.token_to_kv_pool_allocator.available_size() + full_evictable = self.component_evictable_size_[BASE_COMPONENT_TYPE] + lines = [ + f"Available full tokens: {full_available_size + full_evictable} " + f"(full_available_size={full_available_size} + full_evictable_size_={full_evictable})" + ] + for ct in self.tree_components: + if ct == BASE_COMPONENT_TYPE: + continue + if ct.is_swa: + available_size = self.token_to_kv_pool_allocator.swa_available_size() + elif ct.is_mamba: + available_size = self.req_to_token_pool.mamba_pool.available_size() + else: + continue + + lines.append( + f"Available {ct}: {available_size + self.component_evictable_size_[ct]} " + f"(available_size={available_size} + component_evictable_size_={self.component_evictable_size_[ct]})" + ) + return "\n".join(lines) + "\n" + + def _collect_all_nodes(self) -> list[UnifiedTreeNode]: + nodes = [] + stack = [self.root_node] + while stack: + node = stack.pop() + nodes.append(node) + stack.extend(node.children.values()) + return nodes + + def sanity_check(self): + """Thorough sanity check: verify LRU membership, lock state, linked-list + integrity, and evictable sizes for every component. + Expensive — use only in tests or idle checks.""" + try: + # 1. Collect all nodes from tree + all_nodes = self._collect_all_nodes() + + for ct in self.tree_components: + # 2. Basic size invariants + assert ( + self.component_evictable_size_[ct] >= 0 + ), f"component_evictable_size_[{ct}] = {self.component_evictable_size_[ct]} < 0" + assert ( + self.component_protected_size_[ct] >= 0 + ), f"component_protected_size_[{ct}] = {self.component_protected_size_[ct]} < 0" + + # 3. Verify LRU membership: tree nodes with data == LRU cache entries + lru = self.lru_lists[ct] + tree_ids = { + n.id + for n in all_nodes + if n != self.root_node and n.component_data[ct].value is not None + } + lru_ids = set(lru.cache.keys()) + assert tree_ids == lru_ids, ( + f"[{ct}] LRU membership mismatch: " + f"in_tree_not_lru={tree_ids - lru_ids}, " + f"in_lru_not_tree={lru_ids - tree_ids}" + ) + + # 4. Walk LRU doubly-linked list: verify structural integrity + # and that all nodes are unlocked (idle check) + visited = set() + x = lru.head.lru_next[ct] + prev = lru.head + while x != lru.tail: + assert ( + x.lru_prev[ct] == prev + ), f"[{ct}] broken prev link at node {x.id}" + assert ( + x.id in lru.cache + ), f"[{ct}] node {x.id} in linked list but not in cache dict" + assert x.id not in visited, f"[{ct}] cycle detected at node {x.id}" + assert x.component_data[ct].lock_ref == 0, ( + f"[{ct}] node {x.id} should not be locked when idle, " + f"lock_ref={x.component_data[ct].lock_ref}" + ) + visited.add(x.id) + prev = x + x = x.lru_next[ct] + assert len(visited) == len(lru.cache), ( + f"[{ct}] linked list has {len(visited)} nodes, " + f"cache dict has {len(lru.cache)}" + ) + + # 5. Verify evictable size by walking unlocked LRU nodes + recomputed = 0 + x = lru.get_lru_no_lock() + while lru.in_list(x): + v = x.component_data[ct].value + recomputed += len(v) if v is not None else 0 + x = lru.get_prev_no_lock(x) + assert self.component_evictable_size_[ct] == recomputed, ( + f"[{ct}] evictable_size_={self.component_evictable_size_[ct]} " + f"!= recomputed={recomputed}" + ) + + except Exception as e: + logger.error(f"Unified RadixTree sanity check failed: {e}") + self.pretty_print() + raise + + def pretty_print(self) -> None: + stack = [(self.root_node, 0)] + while stack: + node, indent = stack.pop() + component_str = " ".join( + f"{ct}={'yes' if node.component_data[ct].value is not None else 'no'}" + for ct in self.tree_components + ) + print( + " " * indent, + f"[{node.id}]", + len(node.key), + f"full_lock={node.component_data[BASE_COMPONENT_TYPE].lock_ref}", + component_str, + ) + for child in node.children.values(): + stack.append((child, indent + 2)) diff --git a/python/sglang/test/kl_multiturn_utils.py b/python/sglang/test/kl_multiturn_utils.py new file mode 100644 index 000000000000..af11edc8600c --- /dev/null +++ b/python/sglang/test/kl_multiturn_utils.py @@ -0,0 +1,460 @@ +"""Enhanced multi-turn KL divergence test helpers.""" + +from __future__ import annotations + +from typing import Callable + +from sglang.test.kl_test_utils import ( + _extract_output_logprobs, + _flush_cache, + _generate, + _get_input_logprobs, + compare_kl_divergence, + get_input_ids, +) + +__all__ = [ + # Cache assertion callbacks + "default_prefill_cache_assert", + "default_decode_cache_assert", + "make_mamba_prefill_assert", + "make_mamba_decode_assert", + # Enhanced test helpers + "test_input_output_logprobs_match_helper", + "test_input_output_logprobs_match_prefill_cache_hit_helper", + "test_input_output_logprobs_match_decode_cache_hit_helper", + # Internal helpers (for custom inline tests) + "_replay_and_compare_kl", + # Re-exports from kl_test_utils + "get_input_ids", + "_generate", + "_flush_cache", + "_extract_output_logprobs", +] + + +# ============================================================================= +# Cache assertion callbacks +# ============================================================================= +# Prefill signature: (result, prefix_len, label) -> None +# Decode signature: (result, history_len, output_len, label) -> None + + +def default_prefill_cache_assert(result: dict, prefix_len: int, label: str): + """Standard radix cache: cached_tokens == prefix_len.""" + actual = result["meta_info"]["cached_tokens"] + assert ( + actual == prefix_len + ), f"{label}: expected cached_tokens={prefix_len}, got {actual}" + + +def default_decode_cache_assert( + result: dict, history_len: int, output_len: int, label: str +): + """Standard radix cache: cached_tokens == history_len + output_len.""" + expected = history_len + output_len + actual = result["meta_info"]["cached_tokens"] + assert ( + actual == expected + ), f"{label}: expected cached_tokens={expected}, got {actual}" + + +def make_mamba_prefill_assert(chunk_size: int = 64) -> Callable: + """Mamba: cached_tokens in [rounded_down - chunk_size, rounded_down].""" + + def _check(result: dict, prefix_len: int, label: str): + actual = result["meta_info"]["cached_tokens"] + upper = (prefix_len // chunk_size) * chunk_size + lower = max(0, upper - chunk_size) + assert ( + lower <= actual <= upper + ), f"{label}: expected cached_tokens in [{lower}, {upper}], got {actual}" + + return _check + + +def make_mamba_decode_assert(track_interval: int = 16) -> Callable: + """Mamba: cached_tokens = floor((history+output-1)/interval)*interval.""" + + def _check(result: dict, history_len: int, output_len: int, label: str): + actual = result["meta_info"]["cached_tokens"] + if output_len <= 0: + expected = history_len + else: + expected = ( + (history_len + output_len - 1) // track_interval + ) * track_interval + assert ( + actual == expected + ), f"{label}: expected cached_tokens={expected}, got {actual}" + + return _check + + +# ============================================================================= +# Internal helpers +# ============================================================================= + + +def _replay_and_compare_kl( + base_url: str, + model_name: str, + kl_threshold: float, + replay_input_ids: list[list[int]], + output_logprobs: list[list[float]], + label: str, + batch_size: int = 1, +): + """Flush cache, run replay prefill in batches, compare KL divergence.""" + all_input_logprobs = [] + for start in range(0, len(replay_input_ids), batch_size): + end = start + batch_size + all_input_logprobs.extend( + _get_input_logprobs( + base_url, + replay_input_ids[start:end], + output_logprobs[start:end], + ) + ) + acc = {model_name: {"kl_div": kl_threshold}} + compare_kl_divergence(all_input_logprobs, output_logprobs, acc, model_name, label) + + +def _interleave_order(n: int, branches_per_group: int) -> list[int] | None: + """Build interleaved submission order for branch stress testing. + + Given n items grouped into groups of branches_per_group, returns indices + that interleave branches across groups: [g0b0, g1b0, ..., g0b1, g1b1, ...]. + + Returns None if no interleaving is needed. + """ + if branches_per_group <= 0 or branches_per_group >= n: + return None + num_groups = n // branches_per_group + order = [ + g * branches_per_group + b + for b in range(branches_per_group) + for g in range(num_groups) + ] + # Append remainder indices not covered by complete groups + for i in range(num_groups * branches_per_group, n): + order.append(i) + return order + + +def _generate_maybe_interleaved(base_url, inputs, max_new_tokens, order=None): + """Generate with optional interleaved submission order. + + Submits inputs reordered by ``order``, then maps results back to the + original order so the caller always sees results[i] corresponds to + inputs[i]. + """ + if order is None: + return _generate(base_url, inputs, max_new_tokens, return_logprob=True) + ordered = [inputs[i] for i in order] + results = _generate(base_url, ordered, max_new_tokens, return_logprob=True) + unordered = [None] * len(results) + for idx, orig in enumerate(order): + unordered[orig] = results[idx] + return unordered + + +# ============================================================================= +# Helper 1: test_input_output_logprobs_match_helper +# ============================================================================= + + +def test_input_output_logprobs_match_helper( + base_url: str, + model_name: str, + kl_threshold: float, + input_ids: list[list[int]], + *, + label: str = "logprobs_match", + max_new_tokens: int = 256, + # --- Multi-turn --- + # turn_suffixes[t][i] = suffix tokens for sample i at turn t+1 + turn_suffixes: list[list[list[int]]] | None = None, + # --- Cache assertion (for turns > 0) --- + assert_decode_cached_tokens: Callable | None = None, + replay_batch_size: int = 1, +): + """Verify decode logprobs match prefill replay. + + Single-turn (turn_suffixes=None): + flush -> generate(input_ids) -> replay -> KL + + Multi-turn (turn_suffixes provided): + flush -> generate turn 0 -> + for t in range(len(turn_suffixes)): + input = accumulated + output + suffix[t] -> generate -> + assert_decode_cached_tokens (optional) -> + replay last turn -> KL + + Multi-branch: caller passes input_ids where multiple entries share + a prefix. + """ + n = len(input_ids) + num_turns = 1 + (len(turn_suffixes) if turn_suffixes else 0) + print(f"[{label}] {n} samples, {num_turns} turns, max_new_tokens={max_new_tokens}") + + _flush_cache(base_url) + + current_input = list(input_ids) + last_outputs = None + prev_input_lens = [0] * n + prev_output_lens = [0] * n + + for turn in range(num_turns): + if turn > 0: + suffixes = turn_suffixes[turn - 1] + current_input = [ + current_input[i] + last_outputs[i] + suffixes[i] for i in range(n) + ] + + results = _generate( + base_url, current_input, max_new_tokens, return_logprob=True + ) + assert len(results) == n + + if turn > 0 and assert_decode_cached_tokens: + for i, result in enumerate(results): + assert_decode_cached_tokens( + result, + prev_input_lens[i], + prev_output_lens[i], + f"{label}[turn{turn}][{i}]", + ) + + last_outputs = [r["output_ids"] for r in results] + prev_input_lens = [len(current_input[i]) for i in range(n)] + prev_output_lens = [len(last_outputs[i]) for i in range(n)] + + # Replay last turn + replay_ids = [current_input[i] + results[i]["output_ids"] for i in range(n)] + output_lps = [_extract_output_logprobs(r) for r in results] + + _replay_and_compare_kl( + base_url, + model_name, + kl_threshold, + replay_ids, + output_lps, + label=label, + batch_size=replay_batch_size, + ) + + +# ============================================================================= +# Helper 2: test_input_output_logprobs_match_prefill_cache_hit_helper +# ============================================================================= + + +def test_input_output_logprobs_match_prefill_cache_hit_helper( + base_url: str, + model_name: str, + kl_threshold: float, + input_ids: list[list[int]] | None = None, + *, + # --- Multi-branch: explicit prefix/full split --- + prefix_input_ids: list[list[int]] | None = None, + full_input_ids: list[list[int]] | None = None, + label: str = "prefill_cache_hit", + max_new_tokens: int = 256, + # --- Multi-turn: additional turns after the cache-hit generation --- + turn_suffixes: list[list[list[int]]] | None = None, + # --- Cache assertions --- + assert_prefill_cached_tokens: Callable | None = None, # turn 0 + assert_decode_cached_tokens: Callable | None = None, # turns > 0 + # --- Interleaving for branch stress --- + branches_per_group: int = 0, + replay_batch_size: int = 1, +): + """Verify logprobs when prefill cache is hit. + + Original (input_ids only, backward compat): + flush -> seed(input_ids) -> generate(input_ids, cache hit) -> replay -> KL + + Multi-branch (prefix_input_ids + full_input_ids): + flush -> seed(prefixes) -> generate(fulls, prefix cache hit) -> + assert_prefill_cached_tokens -> replay -> KL + + Multi-turn (+ turn_suffixes): + ... after prefill cache-hit turn, additional turns: + input = accumulated + output + suffix -> generate -> + assert_decode_cached_tokens -> replay last turn -> KL + + Interleaving (branches_per_group > 0): + Reorders submission for decode-cache-hit turns to interleave branches + across groups, stressing the radix tree with competing branches. + """ + # Resolve inputs: backward compat with input_ids-only + if input_ids is not None and prefix_input_ids is None: + prefix_input_ids = input_ids + full_input_ids = input_ids + assert prefix_input_ids is not None and full_input_ids is not None + assert len(prefix_input_ids) == len(full_input_ids) + + if assert_prefill_cached_tokens is None: + assert_prefill_cached_tokens = default_prefill_cache_assert + + n = len(full_input_ids) + num_turns = 1 + (len(turn_suffixes) if turn_suffixes else 0) + order = _interleave_order(n, branches_per_group) + print(f"[{label}] {n} samples, {num_turns} turns, max_new_tokens={max_new_tokens}") + + # Seed cache with prefixes + _flush_cache(base_url) + _generate(base_url, prefix_input_ids, max_new_tokens=0) + + # Turn 0: prefill cache hit (NOT interleaved, matching original behavior) + results = _generate(base_url, full_input_ids, max_new_tokens, return_logprob=True) + assert len(results) == n + + for i, result in enumerate(results): + assert_prefill_cached_tokens( + result, len(prefix_input_ids[i]), f"{label}[turn0][{i}]" + ) + + current_input = list(full_input_ids) + last_outputs = [r["output_ids"] for r in results] + prev_input_lens = [len(full_input_ids[i]) for i in range(n)] + prev_output_lens = [len(last_outputs[i]) for i in range(n)] + + # Additional turns: decode cache hits (interleaved if order is set) + if turn_suffixes: + if assert_decode_cached_tokens is None: + assert_decode_cached_tokens = default_decode_cache_assert + + for t, suffixes in enumerate(turn_suffixes): + current_input = [ + current_input[i] + last_outputs[i] + suffixes[i] for i in range(n) + ] + results = _generate_maybe_interleaved( + base_url, current_input, max_new_tokens, order + ) + assert len(results) == n + + for i, result in enumerate(results): + assert_decode_cached_tokens( + result, + prev_input_lens[i], + prev_output_lens[i], + f"{label}[turn{t + 1}][{i}]", + ) + + last_outputs = [r["output_ids"] for r in results] + prev_input_lens = [len(current_input[i]) for i in range(n)] + prev_output_lens = [len(last_outputs[i]) for i in range(n)] + + # Replay last turn + replay_ids = [current_input[i] + results[i]["output_ids"] for i in range(n)] + output_lps = [_extract_output_logprobs(r) for r in results] + + _replay_and_compare_kl( + base_url, + model_name, + kl_threshold, + replay_ids, + output_lps, + label=label, + batch_size=replay_batch_size, + ) + + +# ============================================================================= +# Helper 3: test_input_output_logprobs_match_decode_cache_hit_helper +# ============================================================================= + + +def test_input_output_logprobs_match_decode_cache_hit_helper( + base_url: str, + model_name: str, + kl_threshold: float, + first_turn_input_ids: list[list[int]], + *, + # --- Multi-turn --- + # turn_suffixes[t][i] = suffix for sample i at turn t+2 + turn_suffixes: list[list[list[int]]], + label: str = "decode_cache_hit", + max_new_tokens: int = 256, + # --- Cache assertion --- + assert_decode_cached_tokens: Callable | None = None, + # --- Interleaving --- + branches_per_group: int = 0, + replay_batch_size: int = 1, +): + """Verify logprobs when decode cache is hit. + + 2-turn (turn_suffixes has 1 entry): + flush -> generate turn 1 -> + turn 2: input = turn1 + output + suffix -> generate -> + assert_decode_cached_tokens -> replay -> KL + + Multi-turn (turn_suffixes has N entries): + flush -> generate turn 1 -> + for each turn t: input = accumulated + output + suffix[t] -> generate -> + assert_decode_cached_tokens -> replay last turn -> KL + + Multi-branch: caller duplicates first_turn_input_ids entries and provides + different suffixes per branch. Use branches_per_group for interleaved + submission to stress the radix tree. + """ + assert ( + len(turn_suffixes) >= 1 + ), "turn_suffixes must have at least 1 entry (for turn 2)" + if assert_decode_cached_tokens is None: + assert_decode_cached_tokens = default_decode_cache_assert + + n = len(first_turn_input_ids) + num_turns = 1 + len(turn_suffixes) + order = _interleave_order(n, branches_per_group) + print(f"[{label}] {n} samples, {num_turns} turns, max_new_tokens={max_new_tokens}") + + # Turn 1: populate cache, no assertion, no interleaving + _flush_cache(base_url) + results = _generate( + base_url, first_turn_input_ids, max_new_tokens, return_logprob=True + ) + assert len(results) == n + + current_input = list(first_turn_input_ids) + last_outputs = [r["output_ids"] for r in results] + prev_input_lens = [len(first_turn_input_ids[i]) for i in range(n)] + prev_output_lens = [len(last_outputs[i]) for i in range(n)] + + # Turns 2..N: decode cache hits (interleaved if order is set) + for t, suffixes in enumerate(turn_suffixes): + current_input = [ + current_input[i] + last_outputs[i] + suffixes[i] for i in range(n) + ] + results = _generate_maybe_interleaved( + base_url, current_input, max_new_tokens, order + ) + assert len(results) == n + + for i, result in enumerate(results): + assert_decode_cached_tokens( + result, + prev_input_lens[i], + prev_output_lens[i], + f"{label}[turn{t + 1}][{i}]", + ) + + last_outputs = [r["output_ids"] for r in results] + prev_input_lens = [len(current_input[i]) for i in range(n)] + prev_output_lens = [len(last_outputs[i]) for i in range(n)] + + # Replay last turn + replay_ids = [current_input[i] + results[i]["output_ids"] for i in range(n)] + output_lps = [_extract_output_logprobs(r) for r in results] + + _replay_and_compare_kl( + base_url, + model_name, + kl_threshold, + replay_ids, + output_lps, + label=label, + batch_size=replay_batch_size, + ) diff --git a/test/registered/radix_cache/test_unified_radix_cache_kl.py b/test/registered/radix_cache/test_unified_radix_cache_kl.py new file mode 100644 index 000000000000..abf1324c3b81 --- /dev/null +++ b/test/registered/radix_cache/test_unified_radix_cache_kl.py @@ -0,0 +1,266 @@ +import random +import unittest +from types import SimpleNamespace +from urllib.parse import urlparse + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kl_multiturn_utils import ( + get_input_ids, + make_mamba_decode_assert, + make_mamba_prefill_assert, + test_input_output_logprobs_match_decode_cache_hit_helper, + test_input_output_logprobs_match_helper, + test_input_output_logprobs_match_prefill_cache_hit_helper, +) +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + + +def _random_suffixes(n, length, seed): + """Generate n random token-id lists of the given length.""" + rng = random.Random(seed) + return [[rng.randint(1, 30000) for _ in range(length)] for _ in range(n)] + + +MAMBA_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct" +MAMBA_CHUNK_SIZE = 64 +MAMBA_TRACK_INTERVAL = 128 + +SWA_MODEL = "openai/gpt-oss-20b" +FULL_MODEL = "Qwen/Qwen3-32B" + +register_cuda_ci(est_time=1200, suite="stage-c-test-4-gpu-h100") + + +class UnifiedRadixTreeTestMixin: + """Mixin: gsm8k、mmlu and multi-turn KL tests with multi-branch interleaving.""" + + kl_threshold: float = 0.003 + max_new_tokens: int = 512 + num_groups: int = 3 + branches_per_group: int = 3 + prefix_len: int = 512 + prefill_cache_assert = None + decode_cache_assert = None + + gsm8k_threshold: float = 0.93 + mmlu_threshold: float = 0.8 + num_gsm8k_questions: int = 200 + + def test_gsm8k(self): + """Few-shot GSM8K math reasoning accuracy.""" + from sglang.test.few_shot_gsm8k import run_eval as run_few_shot_gsm8k + + url = urlparse(self.base_url) + args = SimpleNamespace( + num_shots=10, + data_path=None, + num_questions=self.num_gsm8k_questions, + max_new_tokens=16000, + parallel=128, + host=f"http://{url.hostname}", + port=int(url.port), + ) + metrics = run_few_shot_gsm8k(args) + print( + f"[{self.__class__.__name__}] GSM8K accuracy: {metrics['accuracy']:.3f} " + f"(threshold: {self.gsm8k_threshold})" + ) + self.assertGreaterEqual(metrics["accuracy"], self.gsm8k_threshold) + + def test_mmlu(self): + """Simple-evals MMLU multi-task accuracy.""" + from sglang.test.run_eval import run_eval as run_simple_eval + + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="mmlu", + num_examples=64, + num_threads=32, + ) + metrics = run_simple_eval(args) + print( + f"[{self.__class__.__name__}] MMLU score: {metrics['score']:.3f} " + f"(threshold: {self.mmlu_threshold})" + ) + self.assertGreaterEqual(metrics["score"], self.mmlu_threshold) + + def test_multiturn_logprobs_match(self): + """Helper 1: 3-turn, no explicit cache seeding.""" + ids = self.input_ids[:4] + n = len(ids) + t2 = _random_suffixes(n, 512, seed=100) + t3 = _random_suffixes(n, 256, seed=200) + test_input_output_logprobs_match_helper( + self.base_url, + self.model, + self.kl_threshold, + ids, + turn_suffixes=[t2, t3], + assert_decode_cached_tokens=self.decode_cache_assert, + max_new_tokens=self.max_new_tokens, + ) + + def test_multiturn_prefill_cache_hit_branching(self): + """Helper 2: prefill hit + 2 decode-hit turns, multi-branch interleaved.""" + num_groups = self.num_groups + branches = self.branches_per_group + n = num_groups * branches + rng = random.Random(456) + prefix_ids, full_ids = [], [] + for g in range(num_groups): + prefix = self.input_ids[g][: self.prefix_len] + for b in range(branches): + suffix = [rng.randint(1, 30000) for _ in range(256 + b * 64)] + prefix_ids.append(list(prefix)) + full_ids.append(prefix + suffix) + + t2 = _random_suffixes(n, 512, seed=789) + t3 = _random_suffixes(n, 256, seed=890) + test_input_output_logprobs_match_prefill_cache_hit_helper( + self.base_url, + self.model, + self.kl_threshold, + prefix_input_ids=prefix_ids, + full_input_ids=full_ids, + turn_suffixes=[t2, t3], + assert_prefill_cached_tokens=self.prefill_cache_assert, + assert_decode_cached_tokens=self.decode_cache_assert, + branches_per_group=branches, + max_new_tokens=self.max_new_tokens, + ) + + def test_multiturn_decode_cache_hit_branching(self): + """Helper 3: 3-turn decode hit, multi-branch interleaved.""" + num_groups = self.num_groups + branches = self.branches_per_group + n = num_groups * branches + first_turn = [] + for g in range(num_groups): + base = self.input_ids[g][: self.prefix_len] + for _ in range(branches): + first_turn.append(list(base)) + + t2 = _random_suffixes(n, 512, seed=300) + t3 = _random_suffixes(n, 256, seed=400) + test_input_output_logprobs_match_decode_cache_hit_helper( + self.base_url, + self.model, + self.kl_threshold, + first_turn, + turn_suffixes=[t2, t3], + assert_decode_cached_tokens=self.decode_cache_assert, + branches_per_group=branches, + max_new_tokens=self.max_new_tokens, + ) + + +class TestUnifiedFullRadixCache(UnifiedRadixTreeTestMixin, CustomTestCase): + """Full attention.""" + + kl_threshold = 0.0025 + + @classmethod + def setUpClass(cls): + cls.model = FULL_MODEL + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--tp-size", + "4", + "--mem-fraction-static", + "0.80", + "--page-size", + "64", + ], + env={"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"}, + ) + cls.input_ids = get_input_ids(cls.model, num_samples=18) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + +class TestUnifiedMambaRadixCache(UnifiedRadixTreeTestMixin, CustomTestCase): + """Mamba hybrid + UnifiedRadixCache.""" + + kl_threshold = 0.003 + prefill_cache_assert = staticmethod( + make_mamba_prefill_assert(chunk_size=MAMBA_CHUNK_SIZE) + ) + decode_cache_assert = staticmethod( + make_mamba_decode_assert(track_interval=MAMBA_TRACK_INTERVAL) + ) + + @classmethod + def setUpClass(cls): + cls.model = MAMBA_MODEL + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--tp-size", + "4", + "--chunked-prefill-size", + "2048", + "--mem-fraction-static", + "0.85", + "--mamba-scheduler-strategy", + "extra_buffer", + "--mamba-track-interval", + str(MAMBA_TRACK_INTERVAL), + ], + env={"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"}, + ) + cls.input_ids = get_input_ids(cls.model, num_samples=18) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + +class TestUnifiedSWARadixCache(UnifiedRadixTreeTestMixin, CustomTestCase): + """SWA hybrid + UnifiedRadixCache.""" + + kl_threshold = 0.03 + gsm8k_threshold = 0.75 + mmlu_threshold = 0.75 + + @classmethod + def setUpClass(cls): + cls.model = SWA_MODEL + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--tp-size", + "4", + "--mem-fraction-static", + "0.7", + "--disable-piecewise-cuda-graph", + ], + env={"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "0"}, + ) + cls.input_ids = get_input_ids(cls.model, num_samples=18) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py new file mode 100644 index 000000000000..a2e607dabca7 --- /dev/null +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py @@ -0,0 +1,775 @@ +"""Large-scale benchmark + fuzz correctness tests for UnifiedRadixCache. + +Usage (standalone): + bench: python3 test/registered/unit/mem_cache/test_unified_radix_cache_bench.py --num-seqs 5000 --verify --components mamba legacy-mamba swa legacy-swa + CI Test: python -m pytest test/registered/unit/mem_cache/test_unified_radix_cache_bench.py -v -s +""" + +import argparse +import gc +import logging +import random +import statistics +import time +import unittest +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Callable + +import torch + +from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape +from sglang.srt.environ import envs +from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator +from sglang.srt.mem_cache.base_prefix_cache import ( + DecLockRefParams, + EvictParams, + InsertParams, + MatchPrefixParams, +) +from sglang.srt.mem_cache.cache_init_params import CacheInitParams +from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache +from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, HybridReqToTokenPool +from sglang.srt.mem_cache.radix_cache import RadixKey +from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache +from sglang.srt.mem_cache.unified_cache_components.tree_component import ComponentType +from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache +from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler +from sglang.srt.utils import get_device +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=60, suite="stage-b-test-1-gpu-small") + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- +_PAGE_SIZE = 1 +_HEAD_NUM = 2 +_HEAD_DIM = 16 +_NUM_LAYERS = 8 +_GLOBAL_INTERVAL = 4 +_DTYPE = torch.bfloat16 +_SWA_WINDOW_SIZE = 128 + +_BENCH_NUM_SEQS = 5000 +_BENCH_KV_SIZE = 500_000 +_BENCH_CHUNK_LEN = 256 + +_DEFAULT_COMPONENTS = (ComponentType.FULL, ComponentType.MAMBA) + + +@contextmanager +def _suppress_logs(): + root = logging.getLogger() + prev = root.level + root.setLevel(logging.WARNING) + try: + yield + finally: + root.setLevel(prev) + + +def _full_attention_layer_ids(): + return list(range(_GLOBAL_INTERVAL - 1, _NUM_LAYERS, _GLOBAL_INTERVAL)) + + +def _non_full_layer_ids(): + full = set(_full_attention_layer_ids()) + return [i for i in range(_NUM_LAYERS) if i not in full] + + +# =================================================================== +# Sequence generator +# =================================================================== +def gen_random_sequences( + num_seqs: int = 2000, + chunk_len: int = 256, + vocab_size: int = 32000, + seed: int = 42, +) -> list[list[int]]: + """Generate *num_seqs* token sequences with tree-like prefix sharing. + + Phase 1 (50%): chain growth — each new seq extends a random existing one. + Phase 2 (50%): fan-out burst — multiple children from the same parent. + """ + rng = random.Random(seed) + root_prefix = [rng.randint(1, vocab_size) for _ in range(max(1, chunk_len // 4))] + sequences: list[list[int]] = [root_prefix[:]] + + # Phase 1: chain growth + for _ in range(num_seqs // 2): + parent = rng.choice(sequences) + sequences.append( + parent + [rng.randint(1, vocab_size)] * rng.randint(1, chunk_len) + ) + + # Phase 2: fan-out burst + remaining = num_seqs - num_seqs // 2 + while remaining > 0: + fan = min(rng.randint(2, 10), remaining) + parent = rng.choice(sequences) + for _ in range(fan): + sequences.append( + parent + [rng.randint(1, vocab_size)] * rng.randint(1, chunk_len) + ) + remaining -= fan + + rng.shuffle(sequences) + return sequences + + +# =================================================================== +# Cache factory +# =================================================================== +def create_bench_cache( + kv_size, + max_num_reqs, + max_context_len, + components, + page_size=_PAGE_SIZE, + tree_cls=None, +): + """Create cache. Returns (tree, allocator, req_to_token_pool, make_req).""" + device = get_device() + has_mamba = ComponentType.MAMBA in components + has_swa = ComponentType.SWA in components + + mamba2_cache_params = None + if has_mamba: + with envs.SGLANG_MAMBA_SSM_DTYPE.override("bfloat16"): + shape = Mamba2StateShape.create( + tp_world_size=1, + intermediate_size=256, + n_groups=1, + num_heads=2, + head_dim=16, + state_size=16, + conv_kernel=4, + ) + mamba2_cache_params = Mamba2CacheParams( + shape=shape, layers=_non_full_layer_ids() + ) + + # --- req_to_token pool --- + if has_mamba: + req_to_token_pool = HybridReqToTokenPool( + size=max_num_reqs, + mamba_size=max(max_num_reqs * 2, 200), + mamba_spec_state_size=max_num_reqs, + max_context_len=max_context_len, + device=device, + enable_memory_saver=False, + cache_params=mamba2_cache_params, + mamba_layer_ids=_non_full_layer_ids(), + enable_mamba_extra_buffer=False, + speculative_num_draft_tokens=3, + ) + else: + from sglang.srt.mem_cache.memory_pool import ReqToTokenPool + + req_to_token_pool = ReqToTokenPool( + size=max_num_reqs, + max_context_len=max_context_len, + device=device, + enable_memory_saver=False, + ) + + # --- KV pool + allocator --- + if has_swa: + from sglang.srt.mem_cache.swa_memory_pool import ( + SWAKVPool, + SWATokenToKVPoolAllocator, + ) + + pool = SWAKVPool( + size=kv_size, + size_swa=kv_size, + page_size=page_size, + dtype=_DTYPE, + head_num=_HEAD_NUM, + head_dim=_HEAD_DIM, + swa_attention_layer_ids=_non_full_layer_ids(), + full_attention_layer_ids=_full_attention_layer_ids(), + enable_kvcache_transpose=False, + device=device, + ) + allocator = SWATokenToKVPoolAllocator( + size=kv_size, + size_swa=kv_size, + page_size=page_size, + dtype=_DTYPE, + device=device, + kvcache=pool, + need_sort=False, + ) + else: + pool = HybridLinearKVPool( + size=kv_size, + dtype=_DTYPE, + page_size=page_size, + head_num=_HEAD_NUM, + head_dim=_HEAD_DIM, + full_attention_layer_ids=_full_attention_layer_ids(), + enable_kvcache_transpose=False, + device=device, + enable_memory_saver=False, + mamba_pool=req_to_token_pool.mamba_pool if has_mamba else None, + ) + allocator = TokenToKVPoolAllocator( + size=kv_size, + dtype=_DTYPE, + device=device, + kvcache=pool, + need_sort=False, + ) + + # --- tree --- + if tree_cls is None: + tree_cls = UnifiedRadixCache + tree = tree_cls( + params=CacheInitParams( + req_to_token_pool=req_to_token_pool, + token_to_kv_pool_allocator=allocator, + page_size=page_size, + disable=False, + tree_components=components if tree_cls is UnifiedRadixCache else None, + sliding_window_size=_SWA_WINDOW_SIZE if has_swa else None, + ) + ) + + _rid = [0] + + def make_req(): + from sglang.srt.managers.schedule_batch import Req + from sglang.srt.sampling.sampling_params import SamplingParams + + req = Req( + rid=_rid[0], + origin_input_text="", + origin_input_ids=[], + sampling_params=SamplingParams(temperature=0, max_new_tokens=1), + ) + _rid[0] += 1 + req_to_token_pool.alloc([req]) + return req + + return tree, allocator, req_to_token_pool, make_req + + +# =================================================================== +# Shared bench environment + helpers +# =================================================================== +@dataclass +class _Env: + tree: object + alloc: object + rtp: object + make_req: Callable + seqs: list + has_mamba: bool + avg_tokens: int + + +def _make_env(num_seqs, chunk_len, kv_size, components, tree_cls=None): + """Create sequences + cache, return shared _Env.""" + if components is None: + components = _DEFAULT_COMPONENTS + seqs = gen_random_sequences(num_seqs=num_seqs, chunk_len=chunk_len) + max_seq_len = max(len(s) for s in seqs) + avg_tokens = sum(len(s) for s in seqs) // len(seqs) + with _suppress_logs(): + tree, alloc, rtp, make_req = create_bench_cache( + kv_size=kv_size, + max_num_reqs=num_seqs + 100, + max_context_len=max_seq_len + 10, + components=components, + tree_cls=tree_cls, + ) + return _Env( + tree, alloc, rtp, make_req, seqs, ComponentType.MAMBA in components, avg_tokens + ) + + +def _alloc_with_evict(env, n): + """Alloc *n* tokens, evicting if necessary. Returns tensor or None.""" + v = env.alloc.alloc(n) + if v is None: + env.tree.evict(EvictParams(num_tokens=n * 2, mamba_num=2)) + v = env.alloc.alloc(n) + return v + + +def _insert_seq(env, seq): + """Insert one sequence (alloc + evict-fallback). Returns True on success.""" + v = _alloc_with_evict(env, len(seq)) + if v is None: + return False + mamba_val = None + if env.has_mamba: + req = env.make_req() + mamba_val = req.mamba_pool_idx.unsqueeze(0) + env.tree.insert(InsertParams(key=RadixKey(seq), value=v, mamba_value=mamba_val)) + return True + + +def _populate(env, count): + """Insert first *count* sequences (with evict-fallback).""" + for seq in env.seqs[:count]: + _insert_seq(env, seq) + + +def _fill_no_evict(env): + """Insert sequences until pool exhausted (no eviction). Returns count.""" + inserted = 0 + for seq in env.seqs: + v = env.alloc.alloc(len(seq)) + if v is None: + break + mamba_val = None + if env.has_mamba: + req = env.make_req() + mamba_val = req.mamba_pool_idx.unsqueeze(0) + env.tree.insert(InsertParams(key=RadixKey(seq), value=v, mamba_value=mamba_val)) + inserted += 1 + return inserted + + +# =================================================================== +# Benchmark result + runner +# =================================================================== +@dataclass +class BenchResult: + name: str + num_ops: int + total_tokens: int + elapsed_s: float + latencies_us: list[float] + + @property + def ops_per_sec(self): + return self.num_ops / self.elapsed_s if self.elapsed_s > 0 else 0 + + @property + def tokens_per_sec(self): + return self.total_tokens / self.elapsed_s if self.elapsed_s > 0 else 0 + + @property + def p50_us(self): + return statistics.median(self.latencies_us) if self.latencies_us else 0 + + @property + def p99_us(self): + if not self.latencies_us: + return 0 + idx = int(len(self.latencies_us) * 0.99) + return sorted(self.latencies_us)[min(idx, len(self.latencies_us) - 1)] + + def report(self): + tok = ( + f"{self.tokens_per_sec:>12,.0f} tok/s" + if self.total_tokens > 0 + else f"{'N/A':>12s} tok/s" + ) + return ( + f" {self.name:<18s} | {tok} | {self.ops_per_sec:>10,.0f} ops/s | " + f"p50={self.p50_us:>8,.0f}us p99={self.p99_us:>8,.0f}us" + ) + + +def bench_api( + name, setup_fn, op_fn, num_ops, tokens_per_op=0, warmup=10, verify_fn=None +): + """Time *op_fn(item)* for each item from *setup_fn()*. + + *verify_fn*, if provided, runs during warmup and once after timing + (excluded from latency measurement). + """ + items = setup_fn() + assert ( + len(items) >= num_ops + warmup + ), f"need {num_ops + warmup} items, got {len(items)}" + + for i in range(warmup): + op_fn(items[i]) + if verify_fn: + verify_fn(items[i]) + + gc.collect() + gc_was = gc.isenabled() + gc.disable() + + latencies: list[float] = [] + t0 = time.perf_counter() + for i in range(warmup, warmup + num_ops): + ts = time.perf_counter() + op_fn(items[i]) + latencies.append((time.perf_counter() - ts) * 1e6) + elapsed = time.perf_counter() - t0 + + if gc_was: + gc.enable() + if verify_fn: + verify_fn(items[warmup + num_ops - 1]) + + return BenchResult( + name, + num_ops, + tokens_per_op * num_ops if tokens_per_op > 0 else 0, + elapsed, + latencies, + ) + + +# =================================================================== +# Five benchmark scenarios +# =================================================================== +def bench_insert( + num_seqs=5000, + chunk_len=256, + kv_size=500_000, + components=None, + verify=False, + tree_cls=None, +): + """Insert throughput (alloc + evict-fallback + insert).""" + env = _make_env(num_seqs, chunk_len, kv_size, components, tree_cls) + warmup = min(20, num_seqs // 10) + + return bench_api( + "insert", + lambda: list(range(len(env.seqs))), + lambda idx: _insert_seq(env, env.seqs[idx]), + num_seqs - warmup, + env.avg_tokens, + warmup, + (lambda _: env.tree.sanity_check()) if verify else None, + ) + + +def bench_match_prefix( + num_seqs=5000, + chunk_len=256, + kv_size=500_000, + components=None, + verify=False, + tree_cls=None, +): + """Prefix matching throughput (hit / partial / miss mix).""" + env = _make_env(num_seqs, chunk_len, kv_size, components, tree_cls) + _populate(env, num_seqs // 2) + + rng = random.Random(123) + pop = num_seqs // 2 + queries: list[list[int]] = [] + for _ in env.seqs: + roll = rng.random() + if roll < 0.33: + queries.append(env.seqs[rng.randint(0, pop - 1)]) + elif roll < 0.66: + base = env.seqs[rng.randint(0, pop - 1)] + queries.append(base + [rng.randint(1, 32000)] * rng.randint(10, 100)) + else: + queries.append([rng.randint(1, 32000)] * rng.randint(50, 300)) + + def verify_fn(q): + r1 = env.tree.match_prefix(MatchPrefixParams(key=RadixKey(q))) + r2 = env.tree.match_prefix(MatchPrefixParams(key=RadixKey(q))) + assert len(r1.device_indices) == len(r2.device_indices), "match not idempotent" + + warmup = min(20, len(queries) // 10) + return bench_api( + "match_prefix", + lambda: queries, + lambda q: env.tree.match_prefix(MatchPrefixParams(key=RadixKey(q))), + min(len(queries) - warmup, num_seqs), + env.avg_tokens, + warmup, + verify_fn if verify else None, + ) + + +def bench_evict( + num_seqs=5000, + chunk_len=256, + kv_size=500_000, + components=None, + verify=False, + tree_cls=None, +): + """Eviction throughput — fill pool then repeatedly evict batches.""" + env = _make_env(num_seqs, chunk_len, kv_size, components, tree_cls) + inserted = _fill_no_evict(env) + + evict_batch = max(100, kv_size // 200) + num_evictions = max(inserted // 5, 100) + items = [(evict_batch,)] * (num_evictions + 50) + warmup = min(20, num_evictions // 10) + + return bench_api( + "evict", + lambda: items, + lambda item: env.tree.evict(EvictParams(num_tokens=item[0], mamba_num=2)), + num_evictions - warmup, + evict_batch, + warmup, + (lambda _: env.tree.sanity_check()) if verify else None, + ) + + +def bench_lock_unlock( + num_seqs=5000, + chunk_len=256, + kv_size=500_000, + components=None, + verify=False, + tree_cls=None, +): + """Lock/unlock throughput — match nodes then cycle lock/unlock.""" + env = _make_env(num_seqs, chunk_len, kv_size, components, tree_cls) + _populate(env, num_seqs // 2) + + nodes = [] + for seq in env.seqs[: num_seqs // 2]: + r = env.tree.match_prefix(MatchPrefixParams(key=RadixKey(seq))) + if r.last_device_node != env.tree.root_node: + nodes.append(r.last_device_node) + if not nodes: + return BenchResult("lock_unlock", 0, 0, 0, []) + + rng = random.Random(99) + num_pairs = min(len(nodes) * 2, num_seqs) + items = [rng.choice(nodes) for _ in range(num_pairs + 50)] + + def op_fn(node): + lr = env.tree.inc_lock_ref(node) + env.tree.dec_lock_ref( + node, + DecLockRefParams(swa_uuid_for_lock=getattr(lr, "swa_uuid_for_lock", None)), + ) + + warmup = min(20, num_pairs // 10) + return bench_api( + "lock_unlock", + lambda: items, + op_fn, + num_pairs - warmup, + 0, + warmup, + (lambda _: env.tree.sanity_check()) if verify else None, + ) + + +def bench_cache_finished( + num_seqs=5000, + chunk_len=256, + kv_size=500_000, + components=None, + verify=False, + tree_cls=None, +): + """cache_finished_req throughput — full request lifecycle. + + Simulates: match_prefix → inc_lock_ref → alloc → fill req_to_token → cache_finished_req. + """ + env = _make_env(num_seqs, chunk_len, kv_size, components, tree_cls) + + # Pre-build Req objects with token IDs filled into req_to_token + req_items: list = [] + for seq in env.seqs: + key = RadixKey(seq) + mr = env.tree.match_prefix(MatchPrefixParams(key=key)) + matched_len = len(mr.device_indices) + node = mr.last_device_node + lr = env.tree.inc_lock_ref(node) + + remaining = len(seq) - matched_len + if remaining > 0: + v = _alloc_with_evict(env, remaining) + if v is None: + env.tree.dec_lock_ref( + node, + DecLockRefParams( + swa_uuid_for_lock=getattr(lr, "swa_uuid_for_lock", None) + ), + ) + continue + kv_indices = torch.cat([mr.device_indices, v]) + else: + kv_indices = mr.device_indices + + req = env.make_req() + req.origin_input_ids = list(seq) + req.output_ids = [] + req.fill_ids = list(seq) + req.last_node = node + req.cache_protected_len = matched_len + req.kv_committed_len = len(seq) + req.kv_committed_freed = False + if hasattr(lr, "swa_uuid_for_lock"): + req.swa_uuid_for_lock = lr.swa_uuid_for_lock + env.rtp.req_to_token[req.req_pool_idx, : len(kv_indices)] = kv_indices + req_items.append(req) + + if not req_items: + return BenchResult("cache_finished", 0, 0, 0, []) + + warmup = min(20, len(req_items) // 10) + return bench_api( + "cache_finished", + lambda: req_items, + lambda req: env.tree.cache_finished_req(req, is_insert=True), + len(req_items) - warmup, + env.avg_tokens, + warmup, + # Pool math doesn't hold here (many reqs still hold allocated tokens). + (lambda _: env.tree.sanity_check()) if verify else None, + ) + + +# =================================================================== +# Runner +# =================================================================== +ALL_BENCHMARKS = { + "insert": bench_insert, + "match": bench_match_prefix, + "evict": bench_evict, + "lock": bench_lock_unlock, + "cache_finished": bench_cache_finished, +} + + +def run_all_benchmarks( + num_seqs=5000, + chunk_len=256, + kv_size=500_000, + components=None, + verify=False, + benchmarks=None, + tree_cls=None, +): + if components is None: + components = _DEFAULT_COMPONENTS + if benchmarks is None or "all" in benchmarks: + benchmarks = list(ALL_BENCHMARKS.keys()) + + set_global_server_args_for_scheduler( + ServerArgs(model_path="dummy", page_size=_PAGE_SIZE) + ) + + impl_name = (tree_cls or UnifiedRadixCache).__name__ + results = [] + for name in benchmarks: + if name not in ALL_BENCHMARKS: + print(f"[WARN] Unknown benchmark: {name}, skipping") + continue + results.append( + ALL_BENCHMARKS[name]( + num_seqs=num_seqs, + chunk_len=chunk_len, + kv_size=kv_size, + components=components, + verify=verify, + tree_cls=tree_cls, + ) + ) + + print("=" * 100) + print( + f"{impl_name} Benchmark | " + f"num_seqs={num_seqs} chunk_len={chunk_len} kv_size={kv_size} " + f"components={[c.value for c in components]} verify={verify}" + ) + print("-" * 100) + for r in results: + print(r.report()) + print("=" * 100) + return results + + +# =================================================================== +# pytest wrapper +# =================================================================== +class TestUnifiedRadixCacheBench(unittest.TestCase): + + @classmethod + def setUpClass(cls): + set_global_server_args_for_scheduler( + ServerArgs(model_path="dummy", page_size=_PAGE_SIZE) + ) + + def test_bench_insert(self): + r = bench_insert(_BENCH_NUM_SEQS, _BENCH_CHUNK_LEN, _BENCH_KV_SIZE, verify=True) + self.assertGreater(r.num_ops, 0) + self.assertGreater(r.ops_per_sec, 0) + + def test_bench_match_prefix(self): + r = bench_match_prefix( + _BENCH_NUM_SEQS, _BENCH_CHUNK_LEN, _BENCH_KV_SIZE, verify=True + ) + self.assertGreater(r.num_ops, 0) + self.assertGreater(r.ops_per_sec, 0) + + def test_bench_evict(self): + r = bench_evict(_BENCH_NUM_SEQS, _BENCH_CHUNK_LEN, _BENCH_KV_SIZE, verify=True) + self.assertGreater(r.num_ops, 0) + + def test_bench_lock_unlock(self): + r = bench_lock_unlock( + _BENCH_NUM_SEQS, _BENCH_CHUNK_LEN, _BENCH_KV_SIZE, verify=True + ) + self.assertGreater(r.num_ops, 0) + + def test_bench_cache_finished(self): + r = bench_cache_finished( + _BENCH_NUM_SEQS, _BENCH_CHUNK_LEN, _BENCH_KV_SIZE, verify=True + ) + self.assertGreater(r.num_ops, 0) + self.assertGreater(r.ops_per_sec, 0) + + +# =================================================================== +# CLI +# =================================================================== +_TREE_CONFIGS = { + "full": ((ComponentType.FULL,), None), + "mamba": ((ComponentType.FULL, ComponentType.MAMBA), None), + "swa": ((ComponentType.FULL, ComponentType.SWA), None), + "all": ((ComponentType.FULL, ComponentType.SWA, ComponentType.MAMBA), None), + "legacy-mamba": ((ComponentType.FULL, ComponentType.MAMBA), MambaRadixCache), + "legacy-swa": ((ComponentType.FULL, ComponentType.SWA), SWARadixCache), +} + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="UnifiedRadixCache benchmark") + parser.add_argument("--num-seqs", type=int, default=5000) + parser.add_argument("--chunk-len", type=int, default=256) + parser.add_argument("--kv-size", type=int, default=500_000) + parser.add_argument( + "--components", + nargs="+", + choices=list(_TREE_CONFIGS.keys()), + default=["mamba", "legacy-mamba"], + help="Component configs to benchmark", + ) + parser.add_argument( + "--verify", action="store_true", help="Enable correctness assertions" + ) + parser.add_argument( + "--benchmarks", + nargs="+", + default=["all"], + help="insert match evict lock cache_finished all", + ) + args, _ = parser.parse_known_args() + + for comp_name in args.components: + components, tree_cls = _TREE_CONFIGS[comp_name] + run_all_benchmarks( + num_seqs=args.num_seqs, + chunk_len=args.chunk_len, + kv_size=args.kv_size, + components=components, + verify=args.verify, + benchmarks=args.benchmarks, + tree_cls=tree_cls, + ) diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py new file mode 100644 index 000000000000..b8ce9c48ba67 --- /dev/null +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -0,0 +1,883 @@ +"""Unit tests for UnifiedRadixCache""" + +import unittest + +import torch + +from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape +from sglang.srt.environ import envs +from sglang.srt.managers.schedule_batch import Req +from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator +from sglang.srt.mem_cache.base_prefix_cache import ( + DecLockRefParams, + EvictParams, + EvictResult, + InsertParams, + MatchPrefixParams, +) +from sglang.srt.mem_cache.cache_init_params import CacheInitParams +from sglang.srt.mem_cache.common import available_and_evictable_str +from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, HybridReqToTokenPool +from sglang.srt.mem_cache.radix_cache import RadixKey +from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool, SWATokenToKVPoolAllocator +from sglang.srt.mem_cache.unified_cache_components.tree_component import ComponentType +from sglang.srt.mem_cache.unified_radix_cache import ( + UnifiedRadixCache, +) +from sglang.srt.sampling.sampling_params import SamplingParams +from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler +from sglang.srt.utils import get_device +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=10, suite="stage-b-test-1-gpu-small") + +# --------------------------------------------------------------------------- +# Shared constants +# --------------------------------------------------------------------------- +_PAGE_SIZE = 1 +_HEAD_NUM = 2 +_HEAD_DIM = 128 +_NUM_LAYERS = 24 +_GLOBAL_INTERVAL = 4 +_DTYPE = torch.bfloat16 + + +def _full_attention_layer_ids(): + return [i for i in range(_GLOBAL_INTERVAL - 1, _NUM_LAYERS, _GLOBAL_INTERVAL)] + + +def _mamba_layer_ids(): + full_set = set(_full_attention_layer_ids()) + return [i for i in range(_NUM_LAYERS) if i not in full_set] + + +def _swa_attention_layer_ids(): + full_set = set(_full_attention_layer_ids()) + return [i for i in range(_NUM_LAYERS) if i not in full_set] + + +# =================================================================== +# Test: Full + Mamba components (no SWA) +# =================================================================== +class TestUnifiedRadixCacheMamba(unittest.TestCase): + """UnifiedRadixCache with (Full, Mamba) components.""" + + @classmethod + def setUpClass(cls): + set_global_server_args_for_scheduler( + ServerArgs(model_path="dummy", page_size=_PAGE_SIZE) + ) + + def _build_tree( + self, + kv_size: int = 128, + max_num_reqs: int = 10, + mamba_cache_size: int = 20, + max_context_len: int = 128, + ): + device = get_device() + with envs.SGLANG_MAMBA_SSM_DTYPE.override("bfloat16"): + shape = Mamba2StateShape.create( + tp_world_size=1, + intermediate_size=4096, + n_groups=16, + num_heads=32, + head_dim=128, + state_size=128, + conv_kernel=4, + ) + mamba2_cache_params = Mamba2CacheParams( + shape=shape, layers=_mamba_layer_ids() + ) + + req_to_token_pool = HybridReqToTokenPool( + size=max_num_reqs, + mamba_size=mamba_cache_size, + mamba_spec_state_size=max_num_reqs, + max_context_len=max_context_len, + device=device, + enable_memory_saver=False, + cache_params=mamba2_cache_params, + mamba_layer_ids=_mamba_layer_ids(), + enable_mamba_extra_buffer=False, + speculative_num_draft_tokens=3, + ) + pool = HybridLinearKVPool( + size=kv_size, + dtype=_DTYPE, + page_size=_PAGE_SIZE, + head_num=_HEAD_NUM, + head_dim=_HEAD_DIM, + full_attention_layer_ids=_full_attention_layer_ids(), + enable_kvcache_transpose=False, + device=device, + enable_memory_saver=False, + mamba_pool=req_to_token_pool.mamba_pool, + ) + allocator = TokenToKVPoolAllocator( + size=kv_size, + dtype=_DTYPE, + device=device, + kvcache=pool, + need_sort=False, + ) + tree = UnifiedRadixCache( + params=CacheInitParams( + req_to_token_pool=req_to_token_pool, + token_to_kv_pool_allocator=allocator, + page_size=_PAGE_SIZE, + disable=False, + tree_components=(ComponentType.FULL, ComponentType.MAMBA), + ), + ) + + def make_req(): + sp = SamplingParams(temperature=0, max_new_tokens=1) + req = Req( + rid=0, + origin_input_text="", + origin_input_ids=[], + sampling_params=sp, + ) + req_to_token_pool.alloc([req]) + return req + + return tree, allocator, req_to_token_pool, make_req + + # ------- insert + match ------- + def test_insert_and_match_basic(self): + tree, alloc, _, make_req = self._build_tree() + + # Insert [1,2,3] + req1 = make_req() + v1 = alloc.alloc(3) + result = tree.insert( + InsertParams( + key=RadixKey([1, 2, 3]), + value=v1, + mamba_value=req1.mamba_pool_idx.unsqueeze(0), + ) + ) + self.assertEqual(result.prefix_len, 0) + + # Insert [1,2,3,4,5] — shares prefix [1,2,3] + req2 = make_req() + v2 = alloc.alloc(5) + result = tree.insert( + InsertParams( + key=RadixKey([1, 2, 3, 4, 5]), + value=v2, + mamba_value=req2.mamba_pool_idx.unsqueeze(0), + ) + ) + self.assertEqual(result.prefix_len, 3) + + # Match [1,2,3,4,5] — full hit + m = tree.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3, 4, 5]))) + self.assertEqual(len(m.device_indices), 5) + + # Match [1,2,3,4,5,6] — partial hit (5 tokens) + m = tree.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3, 4, 5, 6]))) + self.assertEqual(len(m.device_indices), 5) + + # Match [10,11] — no hit + m = tree.match_prefix(MatchPrefixParams(key=RadixKey([10, 11]))) + self.assertEqual(len(m.device_indices), 0) + + tree.sanity_check() + + # ------- evict: full-only ------- + def test_evict_full_tokens(self): + tree, alloc, _, make_req = self._build_tree() + + # Insert two disjoint sequences + req1 = make_req() + tree.insert( + InsertParams( + key=RadixKey([1, 2, 3]), + value=alloc.alloc(3), + mamba_value=req1.mamba_pool_idx.unsqueeze(0), + ) + ) + req2 = make_req() + tree.insert( + InsertParams( + key=RadixKey([10, 11, 12]), + value=alloc.alloc(3), + mamba_value=req2.mamba_pool_idx.unsqueeze(0), + ) + ) + self.assertEqual(tree.full_evictable_size(), 6) + + # Evict 3 full tokens — should remove one leaf + result = tree.evict(EvictParams(num_tokens=3)) + self.assertIsInstance(result, EvictResult) + self.assertGreaterEqual(result.num_tokens_evicted, 3) + self.assertTrue(tree.full_evictable_size() <= 3) + tree.sanity_check() + + # ------- evict: mamba-only ------- + def test_evict_mamba_only(self): + tree, alloc, rtp, make_req = self._build_tree() + mamba_pool = rtp.mamba_pool + + req1 = make_req() + tree.insert( + InsertParams( + key=RadixKey([1, 2, 3]), + value=alloc.alloc(3), + mamba_value=req1.mamba_pool_idx.unsqueeze(0), + ) + ) + req2 = make_req() + tree.insert( + InsertParams( + key=RadixKey([1, 2, 3, 4, 5, 6, 7]), + value=alloc.alloc(7), + mamba_value=req2.mamba_pool_idx.unsqueeze(0), + ) + ) + self.assertEqual(tree.mamba_evictable_size(), 2) + + # Evict 1 mamba state + result = tree.evict(EvictParams(num_tokens=0, mamba_num=1)) + self.assertGreaterEqual(result.mamba_num_evicted, 1) + # After mamba eviction on an internal node, full tokens remain + self.assertGreaterEqual(tree.full_evictable_size(), 0) + tree.sanity_check() + + # ------- evict: mamba → match stops at tombstone ------- + def test_evict_mamba_breaks_match(self): + tree, alloc, _, make_req = self._build_tree() + + req1 = make_req() + tree.insert( + InsertParams( + key=RadixKey([1, 2, 3]), + value=alloc.alloc(3), + mamba_value=req1.mamba_pool_idx.unsqueeze(0), + ) + ) + req2 = make_req() + tree.insert( + InsertParams( + key=RadixKey([1, 2, 3, 4, 5]), + value=alloc.alloc(5), + mamba_value=req2.mamba_pool_idx.unsqueeze(0), + ) + ) + + # Evict all mamba (2 states) + tree.evict(EvictParams(num_tokens=0, mamba_num=2)) + self.assertEqual(tree.mamba_evictable_size(), 0) + + # Now match should return 0 because mamba validator fails + m = tree.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3, 4, 5]))) + self.assertEqual(len(m.device_indices), 0) + tree.sanity_check() + + # ------- evict: lock_ref protection ------- + def test_evict_respects_lock_ref(self): + tree, alloc, _, make_req = self._build_tree() + + req1 = make_req() + tree.insert( + InsertParams( + key=RadixKey([1, 2, 3]), + value=alloc.alloc(3), + mamba_value=req1.mamba_pool_idx.unsqueeze(0), + ) + ) + req2 = make_req() + tree.insert( + InsertParams( + key=RadixKey([10, 11, 12]), + value=alloc.alloc(3), + mamba_value=req2.mamba_pool_idx.unsqueeze(0), + ) + ) + + # Lock the first leaf + m = tree.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3]))) + locked_node = m.last_device_node + tree.inc_lock_ref(locked_node) + + # Evict all full tokens — only unlocked leaf should be evicted + result = tree.evict(EvictParams(num_tokens=6)) + self.assertGreaterEqual(result.num_tokens_evicted, 3) + + # [1,2,3] is still matchable because it was locked + m = tree.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3]))) + self.assertEqual(len(m.device_indices), 3) + + # Unlock and verify we can now evict it + tree.dec_lock_ref(locked_node) + result = tree.evict(EvictParams(num_tokens=3)) + self.assertGreaterEqual(result.num_tokens_evicted, 3) + tree.sanity_check() + + # ------- evict: verify EvictResult accounting ------- + def test_evict_result_accounting(self): + tree, alloc, _, make_req = self._build_tree() + + req1 = make_req() + tree.insert( + InsertParams( + key=RadixKey([1, 2, 3]), + value=alloc.alloc(3), + mamba_value=req1.mamba_pool_idx.unsqueeze(0), + ) + ) + + # Request 0 mamba + 3 full → full evicted, mamba cascaded + result = tree.evict(EvictParams(num_tokens=3)) + self.assertGreaterEqual(result.num_tokens_evicted, 3) + # Leaf eviction cascades all components; mamba also freed + self.assertGreaterEqual(result.mamba_num_evicted, 1) + tree.sanity_check() + + # ------- insert: prev_prefix_len controls overlap free ------- + def test_insert_prev_prefix_len(self): + tree, alloc, _, make_req = self._build_tree() + initial_avail = alloc.available_size() + + # Step 1: Insert [1,2,3] + req1 = make_req() + tree.insert( + InsertParams( + key=RadixKey([1, 2, 3]), + value=alloc.alloc(3), + mamba_value=req1.mamba_pool_idx.unsqueeze(0), + ) + ) + self.assertEqual(alloc.available_size(), initial_avail - 3) + + # Step 2: Insert [1,2,3,4,5] with prev_prefix_len=0 → frees overlap [0:3] + req2 = make_req() + result = tree.insert( + InsertParams( + key=RadixKey([1, 2, 3, 4, 5]), + value=alloc.alloc(5), + mamba_value=req2.mamba_pool_idx.unsqueeze(0), + prev_prefix_len=0, + ) + ) + self.assertEqual(result.prefix_len, 3) + # alloc 5, freed 3 overlap, stored 2 new → net -2 + self.assertEqual(alloc.available_size(), initial_avail - 3 - 2) + + # Step 3: Insert [1,2,3,4,5,6] with prev_prefix_len=5 → nothing freed + req3 = make_req() + avail_before = alloc.available_size() + result = tree.insert( + InsertParams( + key=RadixKey([1, 2, 3, 4, 5, 6]), + value=alloc.alloc(6), + mamba_value=req3.mamba_pool_idx.unsqueeze(0), + prev_prefix_len=5, + ) + ) + self.assertEqual(result.prefix_len, 5) + # alloc 6, freed 0, stored 1 → net -6 + self.assertEqual(alloc.available_size(), avail_before - 6) + tree.sanity_check() + + # ------- available_and_evictable_str + pretty_print ------- + def test_diagnostics(self): + tree, alloc, _, make_req = self._build_tree() + + req1 = make_req() + tree.insert( + InsertParams( + key=RadixKey([1, 2, 3]), + value=alloc.alloc(3), + mamba_value=req1.mamba_pool_idx.unsqueeze(0), + ) + ) + + diag = tree.available_and_evictable_str() + self.assertIn("Available full tokens", diag) + self.assertIn("mamba", diag.lower()) + + diag2 = available_and_evictable_str(tree) + self.assertIn("Available full tokens", diag2) + + tree.pretty_print() + tree.sanity_check() + + +# =================================================================== +# Test: Full + SWA + Mamba components +# =================================================================== +class TestUnifiedRadixCacheSWAMamba(unittest.TestCase): + """UnifiedRadixCache with (Full, SWA, Mamba) components — the most complex config.""" + + @classmethod + def setUpClass(cls): + set_global_server_args_for_scheduler( + ServerArgs(model_path="dummy", page_size=_PAGE_SIZE) + ) + + def _build_tree( + self, + kv_size: int = 128, + kv_size_swa: int = 64, + max_num_reqs: int = 10, + mamba_cache_size: int = 20, + max_context_len: int = 128, + sliding_window_size: int = 4, + ): + device = get_device() + with envs.SGLANG_MAMBA_SSM_DTYPE.override("bfloat16"): + shape = Mamba2StateShape.create( + tp_world_size=1, + intermediate_size=4096, + n_groups=16, + num_heads=32, + head_dim=128, + state_size=128, + conv_kernel=4, + ) + mamba2_cache_params = Mamba2CacheParams( + shape=shape, layers=_mamba_layer_ids() + ) + + req_to_token_pool = HybridReqToTokenPool( + size=max_num_reqs, + mamba_size=mamba_cache_size, + mamba_spec_state_size=max_num_reqs, + max_context_len=max_context_len, + device=device, + enable_memory_saver=False, + cache_params=mamba2_cache_params, + mamba_layer_ids=_mamba_layer_ids(), + enable_mamba_extra_buffer=False, + speculative_num_draft_tokens=3, + ) + + kv_pool = SWAKVPool( + size=kv_size, + size_swa=kv_size_swa, + page_size=_PAGE_SIZE, + dtype=_DTYPE, + head_num=_HEAD_NUM, + head_dim=_HEAD_DIM, + swa_attention_layer_ids=_swa_attention_layer_ids(), + full_attention_layer_ids=_full_attention_layer_ids(), + enable_kvcache_transpose=False, + device=device, + ) + allocator = SWATokenToKVPoolAllocator( + size=kv_size, + size_swa=kv_size_swa, + page_size=_PAGE_SIZE, + dtype=_DTYPE, + device=device, + kvcache=kv_pool, + need_sort=False, + ) + + tree = UnifiedRadixCache( + params=CacheInitParams( + req_to_token_pool=req_to_token_pool, + token_to_kv_pool_allocator=allocator, + page_size=_PAGE_SIZE, + disable=False, + sliding_window_size=sliding_window_size, + tree_components=( + ComponentType.FULL, + ComponentType.SWA, + ComponentType.MAMBA, + ), + ), + ) + + def make_req(): + sp = SamplingParams(temperature=0, max_new_tokens=1) + req = Req( + rid=0, + origin_input_text="", + origin_input_ids=[], + sampling_params=sp, + ) + req_to_token_pool.alloc([req]) + return req + + return tree, allocator, req_to_token_pool, make_req + + # ------- basic insert + match with SWA ------- + def test_insert_and_match_with_swa(self): + tree, alloc, _, make_req = self._build_tree(sliding_window_size=4) + + req1 = make_req() + tree.insert( + InsertParams( + key=RadixKey([1, 2, 3, 4, 5]), + value=alloc.alloc(5), + mamba_value=req1.mamba_pool_idx.unsqueeze(0), + ) + ) + + # Match: SWA validator requires contiguous window >= sliding_window_size + m = tree.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3, 4, 5]))) + # With sliding_window_size=4 and 5 tokens on single node → should match + self.assertEqual(len(m.device_indices), 5) + tree.sanity_check() + + # ------- evict SWA → cascade Mamba ------- + def test_evict_swa_cascades_mamba(self): + tree, alloc, _, make_req = self._build_tree(sliding_window_size=4) + + # Build tree: [1,2,3] → [4,5,6,7] + req1 = make_req() + tree.insert( + InsertParams( + key=RadixKey([1, 2, 3]), + value=alloc.alloc(3), + mamba_value=req1.mamba_pool_idx.unsqueeze(0), + ) + ) + req2 = make_req() + tree.insert( + InsertParams( + key=RadixKey([1, 2, 3, 4, 5, 6, 7]), + value=alloc.alloc(7), + mamba_value=req2.mamba_pool_idx.unsqueeze(0), + ) + ) + initial_mamba = tree.mamba_evictable_size() + + # Evict SWA — on internal node, SWA eviction cascades to Mamba (priority: swa=1 > mamba=0) + result = tree.evict(EvictParams(num_tokens=0, swa_num_tokens=3)) + self.assertGreaterEqual(result.swa_num_tokens_evicted, 0) + + tree.sanity_check() + + # ------- evict full leaf ------- + def test_evict_full_leaf_cascades_all(self): + tree, alloc, _, make_req = self._build_tree(sliding_window_size=4) + + req1 = make_req() + tree.insert( + InsertParams( + key=RadixKey([1, 2, 3, 4, 5]), + value=alloc.alloc(5), + mamba_value=req1.mamba_pool_idx.unsqueeze(0), + ) + ) + req2 = make_req() + tree.insert( + InsertParams( + key=RadixKey([10, 11, 12, 13, 14]), + value=alloc.alloc(5), + mamba_value=req2.mamba_pool_idx.unsqueeze(0), + ) + ) + self.assertEqual(tree.full_evictable_size(), 10) + + # Evict one leaf (5 full tokens) → also cascades SWA + Mamba + result = tree.evict(EvictParams(num_tokens=5)) + self.assertGreaterEqual(result.num_tokens_evicted, 5) + # Leaf eviction should cascade all components + self.assertGreaterEqual(result.mamba_num_evicted, 1) + self.assertGreaterEqual(result.swa_num_tokens_evicted, 0) + tree.sanity_check() + + # ------- evict with SWA lock ------- + def test_swa_lock_protects_from_eviction(self): + tree, alloc, _, make_req = self._build_tree(sliding_window_size=4) + + req1 = make_req() + tree.insert( + InsertParams( + key=RadixKey([1, 2, 3, 4, 5]), + value=alloc.alloc(5), + mamba_value=req1.mamba_pool_idx.unsqueeze(0), + ) + ) + req2 = make_req() + tree.insert( + InsertParams( + key=RadixKey([10, 11, 12, 13, 14]), + value=alloc.alloc(5), + mamba_value=req2.mamba_pool_idx.unsqueeze(0), + ) + ) + + # Lock the first entry + m = tree.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3, 4, 5]))) + lock_result = tree.inc_lock_ref(m.last_device_node) + + # Try to evict all full tokens + result = tree.evict(EvictParams(num_tokens=10)) + # Only the unlocked one (5 tokens) should be evictable + self.assertGreaterEqual(result.num_tokens_evicted, 5) + + # Locked one is still matchable + m = tree.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3, 4, 5]))) + self.assertEqual(len(m.device_indices), 5) + + # Unlock + tree.dec_lock_ref( + m.last_device_node, + DecLockRefParams(swa_uuid_for_lock=lock_result.swa_uuid_for_lock), + ) + tree.sanity_check() + + # ------- cache_finished_req (with insert) ------- + def test_cache_finished_req_insert(self): + tree, alloc, rtp, make_req = self._build_tree() + + req = make_req() + req.origin_input_ids = [1, 2, 3, 4, 5] + req.output_ids = [6, 7] + kv_len = len(req.origin_input_ids) + len(req.output_ids) + kv_indices = alloc.alloc(kv_len) + rtp.write((req.req_pool_idx, slice(0, kv_len)), kv_indices) + req.kv_committed_len = kv_len + req.last_node = tree.root_node + req.cache_protected_len = 0 + req.swa_uuid_for_lock = None + req.extra_key = None + req.mamba_last_track_seqlen = kv_len + req.fill_ids = req.origin_input_ids + req.output_ids + + tree.cache_finished_req(req, is_insert=True) + + # Verify the tokens are in the tree + m = tree.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3, 4, 5, 6, 7]))) + self.assertEqual(len(m.device_indices), 7) + tree.sanity_check() + + # ------- cache_finished_req (no insert) ------- + def test_cache_finished_req_no_insert(self): + tree, alloc, rtp, make_req = self._build_tree() + + req = make_req() + req.origin_input_ids = [1, 2, 3] + req.output_ids = [] + kv_len = 3 + kv_indices = alloc.alloc(kv_len) + rtp.write((req.req_pool_idx, slice(0, kv_len)), kv_indices) + req.kv_committed_len = kv_len + req.last_node = tree.root_node + req.cache_protected_len = 0 + req.swa_uuid_for_lock = None + req.extra_key = None + req.fill_ids = req.origin_input_ids + + avail_before = alloc.available_size() + tree.cache_finished_req(req, is_insert=False) + + # KV indices should be freed back + self.assertEqual(alloc.available_size(), avail_before + kv_len) + + # Nothing in tree + m = tree.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3]))) + self.assertEqual(len(m.device_indices), 0) + tree.sanity_check() + + # ------- cache_unfinished_req ------- + def test_cache_unfinished_req(self): + tree, alloc, rtp, make_req = self._build_tree() + + req = make_req() + req.origin_input_ids = [1, 2, 3, 4, 5] + req.output_ids = [] + req.fill_ids = req.origin_input_ids[:] + kv_len = len(req.fill_ids) + kv_indices = alloc.alloc(kv_len) + rtp.write((req.req_pool_idx, slice(0, kv_len)), kv_indices) + req.kv_committed_len = kv_len + req.last_node = tree.root_node + req.cache_protected_len = 0 + req.swa_uuid_for_lock = None + req.extra_key = None + req.mamba_last_track_seqlen = kv_len + + tree.cache_unfinished_req(req) + + # After caching, prefix_indices should be set + self.assertGreater(len(req.prefix_indices), 0) + self.assertEqual(req.cache_protected_len, len(req.prefix_indices)) + self.assertIsNotNone(req.last_node) + + # Release the lock acquired by cache_unfinished_req before idle check + tree.dec_lock_ref( + req.last_node, + DecLockRefParams(swa_uuid_for_lock=getattr(req, "swa_uuid_for_lock", None)), + ) + tree.sanity_check() + + # ------- evict empty tree → no crash ------- + def test_evict_empty_tree(self): + tree, alloc, _, _ = self._build_tree() + result = tree.evict(EvictParams(num_tokens=10, mamba_num=5)) + self.assertEqual(result.num_tokens_evicted, 0) + self.assertEqual(result.mamba_num_evicted, 0) + tree.sanity_check() + + # ------- multiple evictions until empty ------- + def test_evict_until_empty(self): + tree, alloc, _, make_req = self._build_tree() + + for i in range(5): + req = make_req() + tokens = list(range(i * 10, i * 10 + 5)) + tree.insert( + InsertParams( + key=RadixKey(tokens), + value=alloc.alloc(5), + mamba_value=req.mamba_pool_idx.unsqueeze(0), + ) + ) + self.assertEqual(tree.full_evictable_size(), 25) + + # Evict all + result = tree.evict(EvictParams(num_tokens=100)) + self.assertGreaterEqual(result.num_tokens_evicted, 25) + self.assertEqual(tree.full_evictable_size(), 0) + self.assertEqual(tree.mamba_evictable_size(), 0) + + # Verify tree is empty (no matches) + m = tree.match_prefix(MatchPrefixParams(key=RadixKey([0, 1, 2, 3, 4]))) + self.assertEqual(len(m.device_indices), 0) + tree.sanity_check() + + # ------- cow mamba on match ------- + def test_match_cow_mamba(self): + tree, alloc, rtp, make_req = self._build_tree() + mamba_pool = rtp.mamba_pool + + req1 = make_req() + tree.insert( + InsertParams( + key=RadixKey([1, 2, 3, 4, 5]), + value=alloc.alloc(5), + mamba_value=req1.mamba_pool_idx.unsqueeze(0), + ) + ) + + # Match with cow_mamba + req2 = make_req() + m = tree.match_prefix( + MatchPrefixParams(key=RadixKey([1, 2, 3, 4, 5]), cow_mamba=True, req=req2) + ) + self.assertEqual(len(m.device_indices), 5) + # req2 should now have its own mamba state (copied) + self.assertIsNotNone(req2.mamba_pool_idx) + + # Verify the copy matches + src_value = m.last_device_node.component_data[ComponentType.MAMBA].value + self.assertTrue( + torch.all( + mamba_pool.mamba_cache.conv[0][:, req2.mamba_pool_idx] + == mamba_pool.mamba_cache.conv[0][:, src_value] + ) + ) + tree.sanity_check() + + +# =================================================================== +# Test: Helper functions +# =================================================================== +class TestUnifiedRadixCacheHelpers(unittest.TestCase): + """Tests for internal helper functions of UnifiedRadixCache.""" + + @classmethod + def setUpClass(cls): + set_global_server_args_for_scheduler( + ServerArgs(model_path="dummy", page_size=_PAGE_SIZE) + ) + + def _build_tree( + self, + kv_size: int = 128, + max_num_reqs: int = 10, + max_context_len: int = 128, + ): + from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, ReqToTokenPool + + device = get_device() + req_to_token_pool = ReqToTokenPool( + size=max_num_reqs, + max_context_len=max_context_len, + device=device, + enable_memory_saver=False, + ) + kv_pool = MHATokenToKVPool( + size=kv_size, + page_size=_PAGE_SIZE, + dtype=_DTYPE, + head_num=_HEAD_NUM, + head_dim=_HEAD_DIM, + layer_num=_NUM_LAYERS, + device=device, + enable_memory_saver=False, + ) + allocator = TokenToKVPoolAllocator( + size=kv_size, + dtype=_DTYPE, + device=device, + kvcache=kv_pool, + need_sort=False, + ) + tree = UnifiedRadixCache( + params=CacheInitParams( + req_to_token_pool=req_to_token_pool, + token_to_kv_pool_allocator=allocator, + page_size=_PAGE_SIZE, + disable=False, + tree_components=( + ComponentType.FULL, + ), # Full attention only, no mamba/swa + ), + ) + return tree, allocator + + def test_readonly_does_not_modify_tree(self): + """Verify readonly match does not modify tree structure (no split).""" + tree, alloc = self._build_tree() + + # Insert [1, 2, 3, 4, 5] + tree.insert( + InsertParams( + key=RadixKey([1, 2, 3, 4, 5]), + value=alloc.alloc(5), + ) + ) + + def count_nodes(node): + count = 1 + for child in node.children.values(): + count += count_nodes(child) + return count + + node_count_before = count_nodes(tree.root_node) + self.assertEqual(node_count_before, 2) # root_node and [1, 2, 3, 4, 5] + + # Regular match with partial key [1, 2] creates a split + value, best_node, best_value_len = tree._match_prefix_helper(RadixKey([1, 2])) + # Regular match with partial key [1, 2, 3, 4] creates a split + value, best_node, best_value_len = tree._match_prefix_helper( + RadixKey([1, 2, 3, 4]) + ) + self.assertEqual(best_value_len, 2) + self.assertEqual(best_node.key.token_ids, [3, 4]) + node_count_after_regular = count_nodes(tree.root_node) + self.assertEqual(node_count_after_regular, node_count_before + 2) + + # Readonly match with partial key [1, 2, 3] should NOT create a split + value, best_node, best_value_len = tree._match_prefix_helper_readonly( + RadixKey([1, 2, 3]) + ) + self.assertEqual(best_value_len, 1) + self.assertEqual(best_node.key.token_ids, [1, 2]) + node_count_after_readonly = count_nodes(tree.root_node) + self.assertEqual(node_count_after_readonly, node_count_after_regular) + + tree.sanity_check() + + +if __name__ == "__main__": + unittest.main() From bf022e177cf6345c77989bdbab4dd92fa79092b1 Mon Sep 17 00:00:00 2001 From: Mick Date: Mon, 13 Apr 2026 11:17:32 +0800 Subject: [PATCH 10/52] Revert "[Diffusion] Add FLUX.1-dev ModelOpt NVFP4 support (#22574)" (#22649) --- docs/diffusion/quantization.md | 22 +- python/sglang/jit_kernel/nvfp4.py | 13 - .../sglang-diffusion-modelopt-quant/SKILL.md | 55 ++- .../sglang-diffusion-performance/SKILL.md | 2 +- .../layers/quantization/modelopt_quant.py | 23 +- .../runtime/loader/fsdp_load.py | 11 +- .../runtime/models/dits/flux.py | 9 - .../pipelines_core/stages/denoising.py | 18 - .../test/unit/test_transformer_quant.py | 138 +----- .../tools/build_modelopt_nvfp4_transformer.py | 402 ------------------ ....py => convert_modelopt_fp8_checkpoint.py} | 252 ++--------- 11 files changed, 74 insertions(+), 871 deletions(-) delete mode 100644 python/sglang/multimodal_gen/tools/build_modelopt_nvfp4_transformer.py rename python/sglang/multimodal_gen/tools/{build_modelopt_fp8_transformer.py => convert_modelopt_fp8_checkpoint.py} (66%) diff --git a/docs/diffusion/quantization.md b/docs/diffusion/quantization.md index 5a33a354412a..b325458b9b42 100644 --- a/docs/diffusion/quantization.md +++ b/docs/diffusion/quantization.md @@ -48,26 +48,6 @@ backend. | `nunchaku-svdq` | Pre-quantized Nunchaku transformer weights, usually named `svdq-{int4\|fp4}_r{rank}-...` | `--transformer-weights-path` | Model-specific support such as Qwen-Image, FLUX, and Z-Image | `nunchaku` | SGLang can infer precision and rank from the filename and supports both `int4` and `nvfp4` | | `msmodelslim` | Pre-quantized msmodelslim transformer weights | `--model-path` | Wan2.2 family | None | Currently only compatible with the Ascend NPU family and supports both `w8a8` and `w4a4` | -## Validated ModelOpt Checkpoints - -This section is the canonical support matrix for diffusion ModelOpt checkpoints -that have been brought up and verified in SGLang. - -### FP8 - -| Base Model | Validated Scope | HF DiT Weights | Notes | -| --- | --- | --- | --- | -| `black-forest-labs/FLUX.1-dev` | single-transformer override, deterministic latent/image comparison, H100 benchmark, torch-profiler trace | `BBuf/flux1-dev-modelopt-fp8-sglang-transformer` | SGLang converter keeps a validated BF16 fallback set for modulation and FF projection layers; use `--model-id FLUX.1-dev` for local mirrors | -| `black-forest-labs/FLUX.2-dev` | single-transformer override load and generation path | `BBuf/flux2-dev-modelopt-fp8-sglang-transformer` | published SGLang-ready transformer override | -| `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | primary `transformer` quantized, `transformer_2` kept BF16 | `BBuf/wan22-t2v-a14b-modelopt-fp8-sglang-transformer` | do not describe this as dual-transformer full-model FP8 unless that path is validated separately | - -### NVFP4 - -| Base Model | Validated Scope | HF DiT Weights | Notes | -| --- | --- | --- | --- | -| `black-forest-labs/FLUX.1-dev` | mixed BF16+NVFP4 transformer override, correctness validation, 4x RTX 5090 benchmark, torch-profiler trace | `unpublished` | use `build_modelopt_nvfp4_transformer.py`; validated builder keeps selected FLUX.1 modules in BF16 and sets `swap_weight_nibbles=false` | -| `black-forest-labs/FLUX.2-dev` | packed-QKV load path | `black-forest-labs/FLUX.2-dev-NVFP4` | validated packed export detection and runtime layout handling | - ## ModelOpt FP8 ### Usage Examples @@ -103,7 +83,7 @@ sglang generate \ - The layerwise offload path now preserves the non-contiguous FP8 weight stride expected by the runtime FP8 GEMM path. - To build the converted checkpoint yourself from a ModelOpt diffusers export, - use `python -m sglang.multimodal_gen.tools.build_modelopt_fp8_transformer`. + use `python -m sglang.multimodal_gen.tools.convert_modelopt_fp8_checkpoint`. ## NVFP4 diff --git a/python/sglang/jit_kernel/nvfp4.py b/python/sglang/jit_kernel/nvfp4.py index ff3c20072668..f54e6414b9e7 100644 --- a/python/sglang/jit_kernel/nvfp4.py +++ b/python/sglang/jit_kernel/nvfp4.py @@ -47,19 +47,6 @@ def _nvfp4_arch_env(): return override_jit_cuda_arch(major, minor, suffix="a") -@torch.compiler.disable -def prewarm_nvfp4_jit_modules( - *, include_expert_quant: bool = False, include_blockwise_moe: bool = False -) -> None: - """Materialize NVFP4 JIT modules before torch.compile traces the model.""" - _jit_nvfp4_quant_module() - _jit_nvfp4_scaled_mm_module() - if include_expert_quant: - _jit_nvfp4_expert_quant_module() - if include_blockwise_moe: - _jit_nvfp4_blockwise_moe_module() - - @cache_once def _jit_nvfp4_quant_module() -> Module: with _nvfp4_arch_env(): diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-modelopt-quant/SKILL.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-modelopt-quant/SKILL.md index 4e81ef4eb2f7..6336984fd0d3 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-modelopt-quant/SKILL.md +++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-modelopt-quant/SKILL.md @@ -23,8 +23,8 @@ This skill owns the ModelOpt-to-SGLang bridge. It is not a generic kernel-tuning - Benchmark only when BF16 and quantized commands are identical except for the checkpoint override being tested. - For diffusion FP8, keep `dit_cpu_offload=false`. `dit_layerwise_offload=true` is valid on the fixed path when you want lower DiT residency. - For multi-transformer pipelines, use per-component overrides when different components need different checkpoints. -- When a branch is missing the validated helper tools, refresh `python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py`, `python/sglang/multimodal_gen/tools/build_modelopt_nvfp4_transformer.py`, and `python/sglang/multimodal_gen/tools/compare_diffusion_trajectory_similarity.py` instead of inventing one-off scripts elsewhere. -- After validating a new ModelOpt quant path, update the ModelOpt support matrix in `docs/diffusion/quantization.md` before closing the task. +- When a branch is missing the validated helper tools, refresh `python/sglang/multimodal_gen/tools/convert_modelopt_fp8_checkpoint.py` and `python/sglang/multimodal_gen/tools/compare_diffusion_trajectory_similarity.py` instead of inventing one-off scripts elsewhere. +- After validating a new ModelOpt quant path, update both the FP8 and NVFP4 support tables in this skill before closing the task. ## Read First @@ -38,8 +38,7 @@ Read these sources before changing code: - `python/sglang/multimodal_gen/runtime/utils/quantization_utils.py` - `python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py` - Helper tools in this repo: - - [`python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py`](../../../tools/build_modelopt_fp8_transformer.py) - - [`python/sglang/multimodal_gen/tools/build_modelopt_nvfp4_transformer.py`](../../../tools/build_modelopt_nvfp4_transformer.py) + - [`python/sglang/multimodal_gen/tools/convert_modelopt_fp8_checkpoint.py`](../../../tools/convert_modelopt_fp8_checkpoint.py) - [`python/sglang/multimodal_gen/tools/compare_diffusion_trajectory_similarity.py`](../../../tools/compare_diffusion_trajectory_similarity.py) If you are working on a new model family, inspect the transformer's config and tensor naming before changing the generic converter. @@ -53,17 +52,32 @@ This repo now contains: - diffusion-side NVFP4 loading from ModelOpt exports - FLUX.2 packed-QKV detection that distinguishes packed NVFP4 checkpoints from standard diffusers exports - automatic protection against incompatible FP8 CPU offload while keeping layerwise DiT offload available -- FP8 transformer build: - [`python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py`](../../../tools/build_modelopt_fp8_transformer.py) +- FP8 export conversion: + [`python/sglang/multimodal_gen/tools/convert_modelopt_fp8_checkpoint.py`](../../../tools/convert_modelopt_fp8_checkpoint.py) - trajectory similarity validation: [`python/sglang/multimodal_gen/tools/compare_diffusion_trajectory_similarity.py`](../../../tools/compare_diffusion_trajectory_similarity.py) ## Documentation Maintenance -- Keep the validated ModelOpt support matrix in `docs/diffusion/quantization.md`. -- Each row should record the validated scope, the Hugging Face repo or path for the quantized DiT weights, and the key caveats. +- Keep two separate support tables in this skill: one for FP8 and one for NVFP4. +- After finishing a new quant support path, update both tables in every mirrored copy of this skill. +- Each row must record the validated scope, the Hugging Face repo or path for the quantized DiT weights, and the key caveats. - If the quantized DiT weights are not published yet, write `unpublished` explicitly instead of leaving the field blank. +## FP8 Supported Models + +| Base Model | Validated Scope | HF DiT Weights | Notes | +| --- | --- | --- | --- | +| `black-forest-labs/FLUX.1-dev` | single-transformer override, deterministic latent/image comparison, H100 benchmark, torch-profiler trace | `BBuf/flux1-dev-modelopt-fp8-sglang-transformer` | SGLang converter keeps a validated BF16 fallback set for modulation and FF projection layers; use `--model-id FLUX.1-dev` for local mirrors | +| `black-forest-labs/FLUX.2-dev` | single-transformer override load and generation path | `BBuf/flux2-dev-modelopt-fp8-sglang-transformer` | published SGLang-ready transformer override | +| `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | primary `transformer` quantized, `transformer_2` kept BF16 | `BBuf/wan22-t2v-a14b-modelopt-fp8-sglang-transformer` | do not describe this as dual-transformer full-model FP8 unless that path is validated separately | + +## NVFP4 Supported Models + +| Base Model | Validated Scope | HF DiT Weights | Notes | +| --- | --- | --- | --- | +| `black-forest-labs/FLUX.2-dev` | packed-QKV load path | `black-forest-labs/FLUX.2-dev-NVFP4` | validated packed export detection and runtime layout handling | + ## FP8 Vs NVFP4 FP8 and NVFP4 are not wired into SGLang in exactly the same way. @@ -106,7 +120,7 @@ python quantize.py \ --model \ --override-model-path \ --model-dtype \ - --format \ + --format \ --batch-size 1 \ --calib-size \ --n-steps \ @@ -116,9 +130,6 @@ python quantize.py \ --hf-ckpt-dir /hf ``` -For current ModelOpt diffusion examples, use `--format fp4` for NVFP4 exports. -Do not assume the checked-out ModelOpt version accepts a literal `nvfp4` format string unless you verified it locally. - For multi-transformer models: - quantize each backbone deliberately @@ -130,7 +141,7 @@ For multi-transformer models: FP8 requires an extra conversion step: ```bash -PYTHONPATH=python python3 -m sglang.multimodal_gen.tools.build_modelopt_fp8_transformer \ +PYTHONPATH=python python3 -m sglang.multimodal_gen.tools.convert_modelopt_fp8_checkpoint \ --modelopt-hf-dir /hf \ --modelopt-backbone-ckpt /ckpt/backbone.pt \ --base-transformer-dir \ @@ -155,25 +166,9 @@ For `FLUX.1-dev`, the validated fallback set currently keeps these modules in BF - `transformer_blocks.*.ff_context.net.0.proj` - `transformer_blocks.*.ff_context.net.2` - `single_transformer_blocks.*.norm.linear` -- `single_transformer_blocks.*.proj_mlp` Use `--model-type flux1` to force that profile, or rely on `--model-type auto` when the export config identifies `FluxTransformer2DModel`. -For FLUX.1-dev NVFP4 model families that need a mixed BF16+NVFP4 checkpoint, build the merged transformer explicitly: - -```bash -PYTHONPATH=python python3 -m sglang.multimodal_gen.tools.build_modelopt_nvfp4_transformer \ - --base-transformer-dir \ - --modelopt-hf-dir /hf/transformer \ - --output-dir /transformer-mixed \ - --pattern-preset flux1-nvfp4 -``` - -The validated FLUX.1-dev mixed builder also needs to preserve: - -- `quant_type: NVFP4` in `config.json` -- `swap_weight_nibbles: false` for the validated diffusers export - ### 4. Load The Quantized Checkpoint In SGLang Single-transformer example: @@ -309,5 +304,5 @@ When documenting results: | `runtime/utils/quantization_utils.py` | resolves flat ModelOpt configs and reconstructs NVFP4 config from metadata | | `runtime/loader/transformer_load_utils.py` | guards incompatible FP8 offload modes | | `runtime/models/dits/flux_2.py` | packed-QKV handling for the packed FLUX.2 NVFP4 family | -| `tools/build_modelopt_fp8_transformer.py` | Build an SGLang-loadable FP8 transformer from a ModelOpt export | +| `tools/convert_modelopt_fp8_checkpoint.py` | FP8 offline conversion into SGLang-native layout | | `tools/compare_diffusion_trajectory_similarity.py` | reduced deterministic BF16-vs-quantized validation | diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md index 4fe50c2b056a..e81449004096 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md +++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md @@ -47,7 +47,7 @@ These options **trade output quality** for speed or VRAM savings. Results will d | **Approximate Attention** | `--attention-backend sage_attn` / `sage_attn_3` / `sliding_tile_attn` / `video_sparse_attn` / `sparse_video_gen_2_attn` / `vmoba_attn` / `sla_attn` / `sage_sla_attn` | Replaces exact attention with approximate or sparse variants. `sage_attn`: INT8/FP8 quantized Q·K; `sliding_tile_attn`: spatial-temporal tile skipping; others: model-specific sparse patterns. | ~1.5–2x on attention (varies by backend) | Quality degradation varies by backend and model. `sage_attn` is the most general; sparse backends (`sliding_tile_attn`, `video_sparse_attn`, etc.) are video-model-specific and may require config files (e.g. `--mask-strategy-file-path` for STA). Requires corresponding packages installed. | | **Cache-DiT** | `SGLANG_CACHE_DIT_ENABLED=true` + `--cache-dit-config ` | Caches intermediate residuals across denoising steps and skips redundant computations via a Selective Computation Mask (SCM). | ~1.5–2x on supported models | Quality depends on SCM config. Incompatible with `--dit-layerwise-offload`. Requires correct per-model config YAML. | | **Quantized Models (Nunchaku / SVDQuant)** | `--enable-svdquant --transformer-weights-path ` + optional `--quantization-precision int4\|nvfp4`, `--quantization-rank 32` | W4A4-style quantization via [Nunchaku](https://nunchaku.tech). Reduces DiT weight memory by ~4x. Precision/rank can be auto-inferred from weight filename or set explicitly. | ~1.5–2x compute speedup | Lossy quantization; quality depends on rank and precision. Requires pre-quantized weights. Ampere (SM8x) or SM12x only (no Hopper SM90). Higher rank = better quality but more memory. | -| **Pre-quantized Weights** | `--transformer-weights-path ` | Load any pre-quantized transformer weights (FP8, INT8, etc.) from a single `.safetensors` file, a directory, or a HuggingFace repo ID. | ~1.3–1.5x compute (dtype dependent) | Requires a validated quantized transformer override, such as one produced by `tools/build_modelopt_fp8_transformer.py` for ModelOpt FP8. Quality slightly worse than BF16; varies by quantization format. | +| **Pre-quantized Weights** | `--transformer-weights-path ` | Load any pre-quantized transformer weights (FP8, INT8, etc.) from a single `.safetensors` file, a directory, or a HuggingFace repo ID. | ~1.3–1.5x compute (dtype dependent) | Requires pre-converted weights (e.g. via `tools/convert_hf_to_fp8.py` for FP8). Quality slightly worse than BF16; varies by quantization format. | | **Component Precision Override** | `--dit-precision fp16`, `--vae-precision fp16\|bf16` | On-the-fly dtype conversion for individual components. E.g. convert a BF16 model to FP16 at load time, or run VAE in BF16 instead of FP32. | Reduces memory; FP16 can be faster on some GPUs | May affect numerical stability. VAE is FP32 by default for accuracy; lowering it is lossy. DiT defaults to BF16. | | **Fewer Inference Steps** | `--num-inference-steps N` (sampling param) | Reduces the number of denoising steps. Fewer steps = faster. | Linear speedup | Quality degrades with too few steps. Model-dependent optimal range. | diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py b/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py index d7ec21a3570b..627edb93edeb 100755 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py @@ -52,15 +52,6 @@ def _get_fp4_gemm_op(): return current_platform.get_modelopt_fp4_gemm_op() -def _prepare_nvfp4_weight_bytes( - weight: torch.Tensor, *, swap_weight_nibbles: bool -) -> torch.Tensor: - """Normalize serialized NVFP4 bytes before padding for the runtime kernel.""" - if not swap_weight_nibbles: - return weight.contiguous() - return ((weight >> 4) | (weight << 4)).contiguous() - - class ModelOptQuantConfig(QuantizationConfig): def __init__( self, @@ -189,7 +180,6 @@ def __init__( exclude_modules: List[str] = None, packed_modules_mapping: Optional[Dict[str, List[str]]] = None, checkpoint_uses_packed_qkv: bool = False, - swap_weight_nibbles: bool = True, ) -> None: super().__init__(exclude_modules, packed_modules_mapping) self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized @@ -200,7 +190,6 @@ def __init__( ) self.group_size = group_size self.checkpoint_uses_packed_qkv = checkpoint_uses_packed_qkv - self.swap_weight_nibbles = swap_weight_nibbles @classmethod def get_name(cls) -> str: @@ -248,7 +237,6 @@ def _add_group_size_from_dict(config: dict): def from_config(cls, config: Dict[str, Any]) -> ModelOptFp4Config: group_size = None exclude_modules = [] - swap_weight_nibbles = True # Flat format (config.json quantization_config) quant_method = config.get("quant_algo") @@ -260,7 +248,6 @@ def from_config(cls, config: Dict[str, Any]) -> ModelOptFp4Config: first_group = next(iter(config_groups.values()), {}) group_size = first_group.get("weights", {}).get("group_size") exclude_modules = config.get("ignore", []) - swap_weight_nibbles = config.get("swap_weight_nibbles", True) else: # Nested format (hf_quant_config.json) try: @@ -268,10 +255,6 @@ def from_config(cls, config: Dict[str, Any]) -> ModelOptFp4Config: quant_method = quant_config["quant_algo"] group_size = ModelOptFp4Config.common_group_size(config) exclude_modules = quant_config.get("exclude_modules", []) - swap_weight_nibbles = quant_config.get( - "swap_weight_nibbles", - config.get("swap_weight_nibbles", True), - ) except (ValueError, KeyError): raise ValueError("Cannot find 'quant_algo' in quantization config.") @@ -291,7 +274,6 @@ def from_config(cls, config: Dict[str, Any]) -> ModelOptFp4Config: exclude_modules=exclude_modules, packed_modules_mapping=config.get("packed_modules_mapping"), checkpoint_uses_packed_qkv=config.get("checkpoint_uses_packed_qkv", False), - swap_weight_nibbles=swap_weight_nibbles, ) def get_quant_method(self, layer: torch.nn.Module, prefix: str): @@ -477,10 +459,9 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.output_size_per_partition = layer.weight.shape[0] + # Swap nibbles: (byte >> 4) | (byte << 4). w = layer.weight.data - w_swapped = _prepare_nvfp4_weight_bytes( - w, swap_weight_nibbles=self.quant_config.swap_weight_nibbles - ) + w_swapped = ((w >> 4) | (w << 4)).contiguous() weight, weights_padding_cols = pad_nvfp4_weight(w_swapped) layer.weights_padding_cols = weights_padding_cols copy_or_rebind_param(layer, "weight", weight) diff --git a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py index c0375be9b385..2941f5e24f95 100644 --- a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py +++ b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py @@ -410,16 +410,7 @@ def load_model_from_full_model_state_dict( ): requires_grad = False temp_param.requires_grad = requires_grad - try: - weight_loader(temp_param, full_tensor) - except AssertionError as exc: - raise AssertionError( - "Failed to shard/load parameter " - f"{target_param_name}: full_tensor.shape={tuple(full_tensor.shape)}, " - f"meta_sharded_param.shape={tuple(meta_sharded_param.shape)}, " - f"temp_param.shape={tuple(temp_param.shape)}, " - f"param_cls={type(actual_param).__name__}" - ) from exc + weight_loader(temp_param, full_tensor) sharded_tensor = temp_param.data else: # In cases where parts of the model aren't sharded, some parameters will be plain tensors diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux.py b/python/sglang/multimodal_gen/runtime/models/dits/flux.py index 481296a2da99..6ecfc22503fd 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/flux.py @@ -261,7 +261,6 @@ def __init__( bias=bias, gather_output=True, quant_config=quant_config, - prefix=f"{prefix}.to_q" if prefix else "to_q", ) self.to_k = ColumnParallelLinear( query_dim, @@ -269,7 +268,6 @@ def __init__( bias=bias, gather_output=True, quant_config=quant_config, - prefix=f"{prefix}.to_k" if prefix else "to_k", ) self.to_v = ColumnParallelLinear( query_dim, @@ -277,7 +275,6 @@ def __init__( bias=bias, gather_output=True, quant_config=quant_config, - prefix=f"{prefix}.to_v" if prefix else "to_v", ) if not self.pre_only: self.to_out = torch.nn.ModuleList([]) @@ -313,7 +310,6 @@ def __init__( bias=added_proj_bias, gather_output=True, quant_config=quant_config, - prefix=f"{prefix}.add_q_proj" if prefix else "add_q_proj", ) self.add_k_proj = ColumnParallelLinear( added_kv_proj_dim, @@ -321,7 +317,6 @@ def __init__( bias=added_proj_bias, gather_output=True, quant_config=quant_config, - prefix=f"{prefix}.add_k_proj" if prefix else "add_k_proj", ) self.add_v_proj = ColumnParallelLinear( added_kv_proj_dim, @@ -329,7 +324,6 @@ def __init__( bias=added_proj_bias, gather_output=True, quant_config=quant_config, - prefix=f"{prefix}.add_v_proj" if prefix else "add_v_proj", ) self.to_add_out = ColumnParallelLinear( self.inner_dim, @@ -501,7 +495,6 @@ def __init__( bias=True, gather_output=True, quant_config=quant_config, - prefix=f"{prefix}.proj_mlp" if prefix else "proj_mlp", ) self.act_mlp = nn.GELU(approximate="tanh") self.proj_out = ColumnParallelLinear( @@ -510,7 +503,6 @@ def __init__( bias=True, gather_output=True, quant_config=quant_config, - prefix=f"{prefix}.proj_out" if prefix else "proj_out", ) self.attn = FluxAttention( query_dim=dim, @@ -521,7 +513,6 @@ def __init__( eps=1e-6, pre_only=True, quant_config=quant_config, - prefix=f"{prefix}.attn" if prefix else "attn", ) def forward( 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 4cd004784567..7f385c967d69 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -19,7 +19,6 @@ from einops import rearrange from tqdm.auto import tqdm -from sglang.jit_kernel.nvfp4 import prewarm_nvfp4_jit_modules from sglang.multimodal_gen import envs from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType, STA_Mode from sglang.multimodal_gen.configs.pipeline_configs.flux import ( @@ -207,26 +206,9 @@ def _maybe_enable_torch_compile(self, module: object) -> None: compile_kwargs["mode"] = mode logger.info(f"Compiling transformer with mode: {mode}") - if self._needs_nvfp4_jit_prewarm(module): - logger.info( - "Prewarming NVFP4 JIT modules before torch.compile to avoid " - "Dynamo tracing JIT initialization." - ) - prewarm_nvfp4_jit_modules() - # TODO(triple-mu): support customized fullgraph and dynamic in the future module.compile(**compile_kwargs) - @staticmethod - def _needs_nvfp4_jit_prewarm(module: nn.Module) -> bool: - for submodule in module.modules(): - quant_method = getattr(submodule, "quant_method", None) - if quant_method is None: - continue - if type(quant_method).__name__ == "ModelOptFp4LinearMethod": - return True - return False - def _maybe_enable_cache_dit( self, num_inference_steps: int | tuple[int, int], batch: Req ) -> None: diff --git a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py index 3a7340d4c082..4d37b5bc7d02 100644 --- a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py +++ b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py @@ -10,8 +10,6 @@ from types import SimpleNamespace from unittest.mock import patch -import torch - partial_json_parser = types.ModuleType("partial_json_parser") partial_json_parser_core = types.ModuleType("partial_json_parser.core") partial_json_parser_exceptions = types.ModuleType("partial_json_parser.core.exceptions") @@ -43,24 +41,15 @@ def _loads(input_str, _flags=None): ) sys.modules.setdefault("partial_json_parser.core.options", partial_json_parser_options) -from sglang.multimodal_gen.runtime.layers.linear import UnquantizedLinearMethod from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import ( NunchakuConfig, ) -from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import ( - ModelOptFp4Config, - _prepare_nvfp4_weight_bytes, -) from sglang.multimodal_gen.runtime.loader.transformer_load_utils import ( _filter_duplicate_precision_variant_safetensors, _Flux2Nvfp4FallbackAdapter, resolve_transformer_quant_load_spec, resolve_transformer_safetensors_to_load, ) -from sglang.multimodal_gen.runtime.models.dits.flux import FluxSingleTransformerBlock -from sglang.multimodal_gen.tools.build_modelopt_nvfp4_transformer import ( - _updated_quant_config, -) class _FakeFluxTransformer: @@ -159,21 +148,22 @@ def test_resolve_transformer_quant_load_spec_keeps_nunchaku_hook( mock_metadata.return_value = { "config": json.dumps({"_class_name": _FakeFluxTransformer.__name__}) } - with tempfile.NamedTemporaryFile(suffix=".safetensors") as f: - nunchaku_config = NunchakuConfig(transformer_weights_path=f.name) - server_args = self._make_server_args( - transformer_weights_path=nunchaku_config.transformer_weights_path, - nunchaku_config=nunchaku_config, - ) + nunchaku_config = NunchakuConfig( + transformer_weights_path="/tmp/svdq-int4_r32.safetensors" + ) + server_args = self._make_server_args( + transformer_weights_path=nunchaku_config.transformer_weights_path, + nunchaku_config=nunchaku_config, + ) - spec = resolve_transformer_quant_load_spec( - hf_config={}, - server_args=server_args, - safetensors_list=[nunchaku_config.transformer_weights_path], - component_model_path="/unused/component/path", - model_cls=_FakeFluxTransformer, - cls_name=_FakeFluxTransformer.__name__, - ) + spec = resolve_transformer_quant_load_spec( + hf_config={}, + server_args=server_args, + safetensors_list=[nunchaku_config.transformer_weights_path], + component_model_path="/unused/component/path", + model_cls=_FakeFluxTransformer, + cls_name=_FakeFluxTransformer.__name__, + ) self.assertIsNone(spec.quant_config) self.assertIs(spec.nunchaku_config, nunchaku_config) @@ -199,104 +189,6 @@ def test_flux2_mixed_nvfp4_fallback_disables_conflicting_offloads(self): self.assertFalse(server_args.dit_cpu_offload) self.assertFalse(server_args.text_encoder_cpu_offload) - def test_prepare_nvfp4_weight_bytes_swaps_nibbles(self): - weight = torch.tensor([[0xAB, 0x10]], dtype=torch.uint8) - - prepared = _prepare_nvfp4_weight_bytes(weight, swap_weight_nibbles=True) - - self.assertEqual(prepared.tolist(), [[0xBA, 0x01]]) - - def test_prepare_nvfp4_weight_bytes_can_skip_nibble_swap(self): - weight = torch.tensor([[0xAB, 0x10]], dtype=torch.uint8) - - prepared = _prepare_nvfp4_weight_bytes(weight, swap_weight_nibbles=False) - - self.assertEqual(prepared.tolist(), [[0xAB, 0x10]]) - - def test_modelopt_fp4_config_reads_swap_weight_nibbles_from_flat_config(self): - config = ModelOptFp4Config.from_config( - { - "quant_algo": "NVFP4", - "group_size": 16, - "ignore": [], - "swap_weight_nibbles": False, - } - ) - - self.assertFalse(config.swap_weight_nibbles) - - def test_modelopt_fp4_config_reads_swap_weight_nibbles_from_nested_config(self): - config = ModelOptFp4Config.from_config( - { - "quantization": { - "quant_algo": "NVFP4", - "exclude_modules": [], - "swap_weight_nibbles": False, - }, - "config_groups": {"default": {"weights": {"group_size": 16}}}, - } - ) - - self.assertFalse(config.swap_weight_nibbles) - - def test_builder_adds_diffusers_quant_type_for_nvfp4(self): - updated = _updated_quant_config( - { - "quantization_config": { - "quant_method": "modelopt", - "quant_algo": "NVFP4", - "ignore": [], - } - }, - fallback_patterns=["single_transformer_blocks.*.proj_mlp*"], - swap_weight_nibbles=False, - ) - - self.assertEqual(updated["quantization_config"]["quant_type"], "NVFP4") - self.assertEqual( - updated["quantization_config"]["ignore"], - ["single_transformer_blocks.*.proj_mlp*"], - ) - - @patch("sglang.multimodal_gen.runtime.layers.linear.get_group_rank", return_value=0) - @patch("sglang.multimodal_gen.runtime.layers.linear.get_group_size", return_value=1) - @patch( - "sglang.multimodal_gen.runtime.layers.linear.get_tp_group", return_value=None - ) - def test_flux_single_transformer_block_modelopt_excludes_use_full_prefix( - self, - _mock_tp_group, - _mock_group_size, - _mock_group_rank, - ): - quant_config = ModelOptFp4Config( - is_checkpoint_nvfp4_serialized=True, - group_size=16, - exclude_modules=[ - "single_transformer_blocks.*.proj_mlp*", - "single_transformer_blocks.*.proj_out*", - "single_transformer_blocks.*.attn.to_q", - ], - ) - - block = FluxSingleTransformerBlock( - dim=64, - num_attention_heads=4, - attention_head_dim=16, - mlp_ratio=2.0, - quant_config=quant_config, - prefix="single_transformer_blocks.0", - ) - - self.assertEqual(block.proj_mlp.prefix, "single_transformer_blocks.0.proj_mlp") - self.assertEqual(block.proj_out.prefix, "single_transformer_blocks.0.proj_out") - self.assertEqual( - block.attn.to_q.prefix, "single_transformer_blocks.0.attn.to_q" - ) - self.assertIsInstance(block.proj_mlp.quant_method, UnquantizedLinearMethod) - self.assertIsInstance(block.proj_out.quant_method, UnquantizedLinearMethod) - self.assertIsInstance(block.attn.to_q.quant_method, UnquantizedLinearMethod) - if __name__ == "__main__": unittest.main() diff --git a/python/sglang/multimodal_gen/tools/build_modelopt_nvfp4_transformer.py b/python/sglang/multimodal_gen/tools/build_modelopt_nvfp4_transformer.py deleted file mode 100644 index e6f13b306a41..000000000000 --- a/python/sglang/multimodal_gen/tools/build_modelopt_nvfp4_transformer.py +++ /dev/null @@ -1,402 +0,0 @@ -"""Build an SGLang-loadable ModelOpt NVFP4 diffusion transformer. - -This tool keeps the ModelOpt-exported NVFP4 tensors for most transformer -modules, but can replace a validated subset of numerically sensitive modules -with their original BF16 tensors from the base transformer checkpoint. - -It is primarily intended for FLUX.1-dev style ModelOpt NVFP4 exports where: -- the base pipeline should remain separate from the quantized transformer -- fallback BF16 modules are model-family specific -- the serialized FP4 weight byte order may already match the runtime kernel -""" - -from __future__ import annotations - -import argparse -import json -import os -import re -import shutil -from collections import defaultdict -from pathlib import Path -from typing import Iterable, Mapping, Sequence - -from safetensors import safe_open -from safetensors.torch import load_file, save_file - -INDEX_FILENAMES = [ - "model.safetensors.index.json", - "diffusion_pytorch_model.safetensors.index.json", -] - -DEFAULT_FLUX1_NVFP4_FALLBACK_PATTERNS = [ - "transformer_blocks.*.norm1.linear*", - "transformer_blocks.*.norm1_context.linear*", - "transformer_blocks.*.ff.net.0.proj*", - "transformer_blocks.*.ff.net.2*", - "transformer_blocks.*.ff_context.net.0.proj*", - "transformer_blocks.*.ff_context.net.2*", - "single_transformer_blocks.*.norm.linear*", - "single_transformer_blocks.*.proj_mlp*", -] - -_TENSOR_MODULE_SUFFIXES = ( - ".weight_scale_2", - ".weight_scale", - ".input_scale", - ".weight", - ".bias", -) - - -def _resolve_transformer_dir(path: str) -> str: - candidate = Path(path).expanduser().resolve() - if (candidate / "config.json").is_file(): - return str(candidate) - transformer_dir = candidate / "transformer" - if (transformer_dir / "config.json").is_file(): - return str(transformer_dir) - raise FileNotFoundError(f"Could not resolve a transformer directory from: {path}") - - -def _find_index_file(model_dir: str) -> str | None: - for filename in INDEX_FILENAMES: - candidate = os.path.join(model_dir, filename) - if os.path.isfile(candidate): - return filename - - matches = sorted( - filename - for filename in os.listdir(model_dir) - if filename.endswith(".safetensors.index.json") - ) - return matches[0] if matches else None - - -def _load_weight_map(model_dir: str) -> tuple[dict[str, str], str | None]: - index_filename = _find_index_file(model_dir) - if index_filename is not None: - with open(os.path.join(model_dir, index_filename), encoding="utf-8") as f: - index_data = json.load(f) - return dict(index_data["weight_map"]), index_filename - - safetensors_files = sorted( - filename - for filename in os.listdir(model_dir) - if filename.endswith(".safetensors") - ) - if len(safetensors_files) != 1: - raise ValueError( - f"Expected an index file or a single safetensors shard in {model_dir}, " - f"found {len(safetensors_files)} shard(s)." - ) - - shard_name = safetensors_files[0] - with safe_open( - os.path.join(model_dir, shard_name), framework="pt", device="cpu" - ) as f: - weight_map = {key: shard_name for key in f.keys()} - index_filename = f"{Path(shard_name).stem}.safetensors.index.json" - return weight_map, index_filename - - -def _load_config(model_dir: str) -> dict: - config_path = os.path.join(model_dir, "config.json") - with open(config_path, encoding="utf-8") as f: - return json.load(f) - - -def _write_config(model_dir: Path, config: Mapping[str, object]) -> None: - with open(model_dir / "config.json", "w", encoding="utf-8") as f: - json.dump(config, f, indent=2, sort_keys=True) - f.write("\n") - - -def _copy_non_shard_files(source_dir: str, output_dir: str) -> None: - ignored = set(INDEX_FILENAMES) - for entry in os.listdir(source_dir): - if entry.endswith(".safetensors") or entry in ignored: - continue - source_path = os.path.join(source_dir, entry) - output_path = os.path.join(output_dir, entry) - if os.path.isdir(source_path): - shutil.copytree(source_path, output_path, dirs_exist_ok=True) - else: - shutil.copy2(source_path, output_path) - - -def _load_selected_tensors( - model_dir: str, - weight_map: Mapping[str, str], - tensor_names: Iterable[str], -): - tensors = {} - names_by_file: dict[str, list[str]] = defaultdict(list) - for name in tensor_names: - names_by_file[weight_map[name]].append(name) - - for filename, names in names_by_file.items(): - shard_path = os.path.join(model_dir, filename) - with safe_open(shard_path, framework="pt", device="cpu") as f: - for name in names: - tensors[name] = f.get_tensor(name).contiguous() - return tensors - - -def _module_name_for_tensor(tensor_name: str) -> str: - for suffix in _TENSOR_MODULE_SUFFIXES: - if tensor_name.endswith(suffix): - return tensor_name[: -len(suffix)] - return tensor_name - - -def _matches_any_pattern(module_name: str, patterns: Sequence[str]) -> bool: - if not patterns: - return False - for pattern in patterns: - regex_str = pattern.replace(".", r"\.").replace("*", r".*") - if re.fullmatch(regex_str, module_name): - return True - return False - - -def _preset_patterns(pattern_preset: str) -> list[str]: - if pattern_preset == "none": - return [] - if pattern_preset == "flux1-nvfp4": - return list(DEFAULT_FLUX1_NVFP4_FALLBACK_PATTERNS) - raise ValueError(f"Unsupported pattern preset: {pattern_preset}") - - -def _updated_quant_config( - source_config: Mapping[str, object], - *, - fallback_patterns: Sequence[str], - swap_weight_nibbles: bool, -) -> dict[str, object]: - output_config = json.loads(json.dumps(source_config)) - quant_config = output_config.get("quantization_config") - if not isinstance(quant_config, dict): - raise ValueError("Expected a flat quantization_config dict in config.json.") - if ( - quant_config.get("quant_method") != "modelopt" - or "FP4" not in str(quant_config.get("quant_algo", "")).upper() - ): - raise ValueError( - "This tool only supports ModelOpt diffusion NVFP4 exports " - "(quant_method=modelopt, quant_algo=FP4/NVFP4)." - ) - - ignore_patterns = list(quant_config.get("ignore", []) or []) - for pattern in fallback_patterns: - if pattern not in ignore_patterns: - ignore_patterns.append(pattern) - - quant_config["ignore"] = ignore_patterns - quant_config.setdefault( - "quant_type", str(quant_config.get("quant_algo", "")).upper() - ) - quant_config["swap_weight_nibbles"] = swap_weight_nibbles - return output_config - - -def build_modelopt_nvfp4_transformer( - *, - base_transformer_dir: str, - modelopt_hf_dir: str, - output_dir: str, - pattern_preset: str = "none", - keep_bf16_patterns: Sequence[str] | None = None, - swap_weight_nibbles: bool | None = None, - overwrite: bool = False, -) -> dict[str, int | bool]: - source_dir = _resolve_transformer_dir(modelopt_hf_dir) - base_dir = _resolve_transformer_dir(base_transformer_dir) - - patterns = _preset_patterns(pattern_preset) - if keep_bf16_patterns: - patterns.extend(keep_bf16_patterns) - - resolved_swap_weight_nibbles = ( - swap_weight_nibbles - if swap_weight_nibbles is not None - else (False if pattern_preset == "flux1-nvfp4" else True) - ) - output_config = _updated_quant_config( - _load_config(source_dir), - fallback_patterns=patterns, - swap_weight_nibbles=resolved_swap_weight_nibbles, - ) - quant_config = output_config["quantization_config"] - serialized_quant_config = json.dumps(quant_config, sort_keys=True) - - output_path = Path(output_dir).expanduser().resolve() - if output_path.exists(): - if not overwrite: - raise FileExistsError( - f"Output directory already exists: {output_path}. " - "Use --overwrite to replace it." - ) - shutil.rmtree(output_path) - output_path.mkdir(parents=True, exist_ok=True) - - _copy_non_shard_files(source_dir, str(output_path)) - _write_config(output_path, output_config) - - source_weight_map, index_filename = _load_weight_map(source_dir) - base_weight_map, _ = _load_weight_map(base_dir) - - fallback_tensor_names = sorted( - name - for name in base_weight_map - if name in source_weight_map - and _matches_any_pattern(_module_name_for_tensor(name), patterns) - ) - fallback_tensors = _load_selected_tensors( - base_dir, - base_weight_map, - fallback_tensor_names, - ) - fallback_modules = { - _module_name_for_tensor(tensor_name) for tensor_name in fallback_tensor_names - } - - weights_by_file: dict[str, list[str]] = defaultdict(list) - for tensor_name, filename in source_weight_map.items(): - weights_by_file[filename].append(tensor_name) - - updated_weight_map: dict[str, str] = {} - total_size = 0 - replaced_tensor_count = 0 - removed_aux_tensor_count = 0 - - for filename, tensor_names in sorted(weights_by_file.items()): - shard_path = os.path.join(source_dir, filename) - shard_tensors = load_file(shard_path, device="cpu") - - with safe_open(shard_path, framework="pt", device="cpu") as f: - metadata = dict(f.metadata() or {}) - - metadata.setdefault("format", "pt") - metadata["quantization_config"] = serialized_quant_config - metadata["_quantization_metadata"] = serialized_quant_config - - for name in list(shard_tensors.keys()): - if "_quantizer." in name: - del shard_tensors[name] - removed_aux_tensor_count += 1 - continue - - module_name = _module_name_for_tensor(name) - if module_name not in fallback_modules: - continue - - if name in fallback_tensors: - shard_tensors[name] = fallback_tensors[name] - replaced_tensor_count += 1 - else: - del shard_tensors[name] - removed_aux_tensor_count += 1 - - save_file(shard_tensors, os.path.join(output_path, filename), metadata=metadata) - - for name, tensor in shard_tensors.items(): - updated_weight_map[name] = filename - total_size += tensor.element_size() * tensor.numel() - - if index_filename is None: - raise ValueError( - "Expected a sharded or indexed ModelOpt HF export, but no index file was found." - ) - - with open(output_path / index_filename, "w", encoding="utf-8") as f: - json.dump( - { - "metadata": {"total_size": total_size}, - "weight_map": updated_weight_map, - }, - f, - indent=2, - sort_keys=True, - ) - f.write("\n") - - return { - "fallback_modules": len(fallback_modules), - "replaced_tensors": replaced_tensor_count, - "removed_aux_tensors": removed_aux_tensor_count, - "output_shards": len(weights_by_file), - "swap_weight_nibbles": resolved_swap_weight_nibbles, - } - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser( - description=( - "Build an SGLang-loadable ModelOpt NVFP4 diffusion transformer and " - "optionally keep selected modules in BF16." - ) - ) - parser.add_argument( - "--base-transformer-dir", - required=True, - help="Original BF16 transformer directory, or a parent model directory.", - ) - parser.add_argument( - "--modelopt-hf-dir", - required=True, - help="ModelOpt --hf-ckpt-dir output, or its transformer subdirectory.", - ) - parser.add_argument( - "--output-dir", - required=True, - help="Directory to write the mixed transformer checkpoint.", - ) - parser.add_argument( - "--pattern-preset", - choices=["none", "flux1-nvfp4"], - default="none", - help="Optional model-family BF16 fallback preset.", - ) - parser.add_argument( - "--keep-bf16-pattern", - action="append", - default=[], - help=( - "Glob-style pattern matched against module names without trailing tensor " - "suffixes such as .weight or .bias." - ), - ) - parser.add_argument( - "--swap-weight-nibbles", - action=argparse.BooleanOptionalAction, - default=None, - help=( - "Whether the runtime should swap packed FP4 nibbles before padding. " - "Defaults to false for --pattern-preset flux1-nvfp4 and true otherwise." - ), - ) - parser.add_argument( - "--overwrite", - action="store_true", - help="Replace --output-dir if it already exists.", - ) - return parser.parse_args() - - -def main() -> None: - args = _parse_args() - stats = build_modelopt_nvfp4_transformer( - base_transformer_dir=args.base_transformer_dir, - modelopt_hf_dir=args.modelopt_hf_dir, - output_dir=args.output_dir, - pattern_preset=args.pattern_preset, - keep_bf16_patterns=args.keep_bf16_pattern, - swap_weight_nibbles=args.swap_weight_nibbles, - overwrite=args.overwrite, - ) - print(json.dumps(stats, indent=2, sort_keys=True)) - - -if __name__ == "__main__": - main() diff --git a/python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py b/python/sglang/multimodal_gen/tools/convert_modelopt_fp8_checkpoint.py similarity index 66% rename from python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py rename to python/sglang/multimodal_gen/tools/convert_modelopt_fp8_checkpoint.py index 5d87f5ff1bf2..89a731fba863 100644 --- a/python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py +++ b/python/sglang/multimodal_gen/tools/convert_modelopt_fp8_checkpoint.py @@ -1,4 +1,4 @@ -"""Build an SGLang-loadable ModelOpt FP8 diffusion transformer. +"""Convert a ModelOpt diffusion FP8 export into an SGLang-loadable checkpoint. The core conversion path is model-agnostic: - read the ModelOpt diffusers transformer export @@ -12,7 +12,7 @@ Example: - python -m sglang.multimodal_gen.tools.build_modelopt_fp8_transformer \ + python -m sglang.multimodal_gen.tools.convert_modelopt_fp8_checkpoint \ --modelopt-hf-dir /tmp/modelopt_flux2_fp8/hf \ --modelopt-backbone-ckpt /tmp/modelopt_flux2_fp8/ckpt/backbone.pt \ --base-transformer-dir /path/to/FLUX.2-dev/transformer \ @@ -57,21 +57,6 @@ r"^transformer_blocks\.\d+\.ff_context\.net\.2$", r"^single_transformer_blocks\.\d+\.norm\.linear$", ] -DEFAULT_LTX2_KEEP_BF16_PATTERNS = [ - r"^(audio_)?adaln_single\.emb\.timestep_embedder\.linear_[12]$", - r"^(audio_)?adaln_single\.linear$", - r"^audio_caption_projection\.linear_[12]$", - r"^audio_patchify_proj$", - r"^audio_proj_out$", - r"^av_ca_(a2v_gate|audio_scale_shift|v2a_gate|video_scale_shift)_adaln_single\.emb\.timestep_embedder\.linear_[12]$", - r"^av_ca_(a2v_gate|audio_scale_shift|v2a_gate|video_scale_shift)_adaln_single\.linear$", - r"^caption_projection\.linear_[12]$", - r"^patchify_proj$", - r"^proj_out$", - r"^transformer_blocks\.(0|43|44|45|46|47)\.(attn1|attn2|audio_attn1|audio_attn2|audio_to_video_attn|video_to_audio_attn)\.to_(q|k|v)$", - r"^transformer_blocks\.(0|43|44|45|46|47)\.(attn1|attn2|audio_attn1|audio_attn2|audio_to_video_attn|video_to_audio_attn)\.to_out\.0$", - r"^transformer_blocks\.(0|43|44|45|46|47)\.(ff|audio_ff)\.proj_(in|out)$", -] def _resolve_transformer_dir(path: str) -> str: @@ -141,116 +126,9 @@ def _load_config(model_dir: str) -> dict: return json.load(f) -def _load_first_shard_metadata( - model_dir: str, weight_map: Mapping[str, str] -) -> dict[str, str]: - if not weight_map: - return {} - first_shard = next(iter(weight_map.values())) - with safe_open( - os.path.join(model_dir, first_shard), framework="pt", device="cpu" - ) as f: - return dict(f.metadata() or {}) - - -def _module_name_variants(weight_name: str) -> list[str]: - module_name = weight_name[:-7] if weight_name.endswith(".weight") else weight_name - variants = [module_name] - - for prefix in ("model.diffusion_model.", "velocity_model."): - if module_name.startswith(prefix): - variants.append(module_name[len(prefix) :]) - - canonicalized: list[str] = [] - for variant in variants: - canonicalized.append( - re.sub(r"(\.audio_ff|\.ff)\.net\.0\.proj$", r"\1.proj_in", variant) - ) - canonicalized.append( - re.sub(r"(\.audio_ff|\.ff)\.net\.2$", r"\1.proj_out", variant) - ) - variants.extend(canonicalized) - - deduped: list[str] = [] - for variant in variants: - if variant not in deduped: - deduped.append(variant) - return deduped - - -def _preferred_module_name(weight_name: str) -> str: - return _module_name_variants(weight_name)[-1] - - -def _scale_key_candidates(weight_name: str) -> list[str]: - candidates = [weight_name] - if weight_name.startswith("model.diffusion_model."): - candidates.append( - "velocity_model." + weight_name[len("model.diffusion_model.") :] - ) - return candidates - - -def _resolve_scale_key( - weight_name: str, - scale_map: Mapping[str, Mapping[str, torch.Tensor]], -) -> str | None: - for candidate in _scale_key_candidates(weight_name): - if candidate in scale_map: - return candidate - return None - - -def _is_ltx2_x0_export( - *, - config: Mapping[str, object], - source_metadata: Mapping[str, str], - source_weight_map: Mapping[str, str], -) -> bool: - if config.get("_class_name") != "X0Model": - return False - if not any(name.startswith("model.diffusion_model.") for name in source_weight_map): - return False - try: - metadata_config = json.loads(str(source_metadata.get("config", ""))) - except json.JSONDecodeError: - return False - return isinstance(metadata_config.get("transformer"), dict) - - -def _build_output_config( - *, - source_config: Mapping[str, object], - source_metadata: Mapping[str, str], - quant_config: Mapping[str, object], - is_ltx2_x0_export: bool, -) -> dict[str, object]: - if is_ltx2_x0_export: - metadata_config = json.loads(str(source_metadata["config"])) - output_config = dict(metadata_config["transformer"]) - output_config["_class_name"] = "LTX2VideoTransformer3DModel" - else: - output_config = dict(source_config) - - output_config["quantization_config"] = dict(quant_config) - return output_config - - -def _should_keep_ltx2_transformer_key(weight_name: str) -> bool: - if not weight_name.startswith("model.diffusion_model."): - return False - connector_prefixes = ( - "model.diffusion_model.audio_embeddings_connector.", - "model.diffusion_model.video_embeddings_connector.", - ) - return not weight_name.startswith(connector_prefixes) - - def get_default_keep_bf16_patterns( *, model_type: str, class_name: str | None ) -> list[str]: - if model_type == "ltx2": - return list(DEFAULT_LTX2_KEEP_BF16_PATTERNS) if model_type == "flux1": return list(DEFAULT_FLUX1_KEEP_BF16_PATTERNS) if model_type == "flux2": @@ -271,11 +149,8 @@ def should_keep_bf16( if not keep_bf16_patterns: return False - return any( - re.search(pattern, module_name) - for pattern in keep_bf16_patterns - for module_name in _module_name_variants(weight_name) - ) + module_name = weight_name[:-7] if weight_name.endswith(".weight") else weight_name + return any(re.search(pattern, module_name) for pattern in keep_bf16_patterns) def is_ignored_by_modelopt( @@ -285,12 +160,10 @@ def is_ignored_by_modelopt( if not ignore_patterns: return False + module_name = weight_name[:-7] if weight_name.endswith(".weight") else weight_name for pattern in ignore_patterns: regex_str = pattern.replace(".", r"\.").replace("*", r".*") - if any( - re.fullmatch(regex_str, module_name) - for module_name in _module_name_variants(weight_name) - ): + if re.fullmatch(regex_str, module_name): return True return False @@ -369,7 +242,7 @@ def _load_selected_tensors( return tensors -def build_modelopt_fp8_transformer( +def convert_modelopt_fp8_checkpoint( *, modelopt_hf_dir: str, modelopt_backbone_ckpt: str, @@ -392,29 +265,23 @@ def build_modelopt_fp8_transformer( raise ValueError( "Expected a flat quantization_config dict in the ModelOpt export." ) - if quant_config.get("quant_method") != "modelopt": + if ( + quant_config.get("quant_method") != "modelopt" + or "FP8" not in str(quant_config.get("quant_algo", "")).upper() + ): raise ValueError( "This tool only supports ModelOpt diffusers FP8 exports " - "(quant_method=modelopt)." + "(quant_method=modelopt, quant_algo=FP8)." ) - source_weight_map_all, index_filename = _load_weight_map(source_dir) - source_metadata = _load_first_shard_metadata(source_dir, source_weight_map_all) - is_ltx2_export = _is_ltx2_x0_export( - config=config, - source_metadata=source_metadata, - source_weight_map=source_weight_map_all, - ) class_name = config.get("_class_name") ignore_patterns = list(quant_config.get("ignore", []) or []) patterns = list( get_default_keep_bf16_patterns(model_type=model_type, class_name=class_name) ) - if is_ltx2_export and model_type == "auto": - patterns.extend(DEFAULT_LTX2_KEEP_BF16_PATTERNS) if keep_bf16_patterns: patterns.extend(keep_bf16_patterns) - if patterns and base_dir is None and not is_ltx2_export: + if patterns and base_dir is None: raise ValueError( "BF16 fallback patterns are enabled, but --base-transformer-dir was not provided." ) @@ -431,73 +298,25 @@ def build_modelopt_fp8_transformer( _copy_non_shard_files(source_dir, str(output_path)) - if is_ltx2_export: - source_weight_map = { - name: filename - for name, filename in source_weight_map_all.items() - if _should_keep_ltx2_transformer_key(name) - } - else: - source_weight_map = source_weight_map_all + source_weight_map, index_filename = _load_weight_map(source_dir) base_weight_map: dict[str, str] = {} if base_dir is not None: base_weight_map, _ = _load_weight_map(base_dir) - fallback_weight_names = sorted( - weight_name - for weight_name in source_weight_map - if weight_name.endswith(".weight") and should_keep_bf16(weight_name, patterns) - ) - fallback_weight_names_set = set(fallback_weight_names) backbone_state = torch.load(backbone_ckpt_path, map_location="cpu")[ "model_state_dict" ] fp8_scale_map = build_fp8_scale_map(backbone_state, maxbound=maxbound) - quant_algo = str(quant_config.get("quant_algo", "")).upper() - if quant_algo and "FP8" not in quant_algo: - raise ValueError( - "This tool only supports ModelOpt diffusers FP8 exports, " - f"got quant_algo={quant_config.get('quant_algo')!r}." - ) - if not quant_algo and not fp8_scale_map: - raise ValueError( - "Could not infer an FP8 ModelOpt export: quantization_config.quant_algo " - "is missing and backbone.pt does not contain FP8 scale tensors." - ) - effective_quant_config = json.loads(json.dumps(quant_config)) - if not quant_algo: - effective_quant_config["quant_algo"] = "FP8" - - auto_ignore_modules = sorted( - { - _preferred_module_name(weight_name) - for weight_name in source_weight_map - if weight_name.endswith(".weight") - and _resolve_scale_key(weight_name, fp8_scale_map) is None - } - ) - fallback_ignore_modules = sorted( - {_preferred_module_name(weight_name) for weight_name in fallback_weight_names} - ) - ignore_patterns = sorted( - { - *ignore_patterns, - *auto_ignore_modules, - *fallback_ignore_modules, - } - ) - effective_quant_config["ignore"] = ignore_patterns - serialized_quant_config = json.dumps(effective_quant_config, sort_keys=True) - output_config = _build_output_config( - source_config=config, - source_metadata=source_metadata, - quant_config=effective_quant_config, - is_ltx2_x0_export=is_ltx2_export, - ) + serialized_quant_config = json.dumps(quant_config, sort_keys=True) + fallback_weight_names = sorted( + weight_name + for weight_name in source_weight_map + if weight_name.endswith(".weight") and should_keep_bf16(weight_name, patterns) + ) fallback_tensors = ( _load_selected_tensors(base_dir, base_weight_map, fallback_weight_names) - if fallback_weight_names and base_dir is not None + if fallback_weight_names else {} ) fallback_scale_names = { @@ -521,23 +340,15 @@ def build_modelopt_fp8_transformer( for filename, names in sorted(weights_by_file.items()): shard_path = os.path.join(source_dir, filename) shard_tensors = load_file(shard_path, device="cpu") - selected_names = set(names) with safe_open(shard_path, framework="pt", device="cpu") as f: metadata = dict(f.metadata() or {}) metadata.setdefault("format", "pt") - metadata["_class_name"] = str( - output_config.get("_class_name", metadata.get("_class_name", "")) - ) - metadata["config"] = json.dumps(output_config, sort_keys=True) metadata["quantization_config"] = serialized_quant_config metadata["_quantization_metadata"] = serialized_quant_config for name in list(shard_tensors.keys()): - if name not in selected_names: - del shard_tensors[name] - continue if "_quantizer." in name: del shard_tensors[name] continue @@ -551,14 +362,12 @@ def build_modelopt_fp8_transformer( continue if name in fallback_tensors: shard_tensors[name] = fallback_tensors[name] - scale_key = _resolve_scale_key(name, fp8_scale_map) if ( name.endswith(".weight") - and scale_key is not None + and name in fp8_scale_map and name not in fallback_tensors - and name not in fallback_weight_names_set ): - scale_tensors = fp8_scale_map[scale_key] + scale_tensors = fp8_scale_map[name] shard_tensors[name] = quantize_fp8_weight( shard_tensors[name], scale_tensors["weight_scale"] ) @@ -588,15 +397,12 @@ def build_modelopt_fp8_transformer( sort_keys=True, ) - with open(output_path / "config.json", "w", encoding="utf-8") as f: - json.dump(output_config, f, indent=2, sort_keys=True) - return { "quantized_weights": sum( 1 for name in source_weight_map if name.endswith(".weight") - and _resolve_scale_key(name, fp8_scale_map) is not None + and name in fp8_scale_map and not is_ignored_by_modelopt(name, ignore_patterns) ), "bf16_fallback_weights": len(fallback_weight_names), @@ -609,8 +415,8 @@ def build_modelopt_fp8_transformer( def _parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description=( - "Build an SGLang-loadable ModelOpt FP8 diffusion transformer from a " - "ModelOpt diffusers export." + "Inject FP8 scales from ModelOpt backbone.pt into a diffusers export so " + "SGLang diffusion can load it natively." ) ) parser.add_argument( @@ -637,11 +443,11 @@ def _parse_args() -> argparse.Namespace: ) parser.add_argument( "--model-type", - choices=["auto", "flux1", "flux2", "ltx2", "none"], + choices=["auto", "flux1", "flux2", "none"], default="auto", help=( "Optional model-family BF16 fallback profile. 'none' uses the generic " - "conversion path. 'auto' enables the validated FLUX.1 / FLUX.2 / LTX-2 " + "conversion path. 'auto' enables the validated FLUX.1 / FLUX.2 " "fallback set when the export config matches those transformer classes." ), ) @@ -671,7 +477,7 @@ def _parse_args() -> argparse.Namespace: def main() -> None: args = _parse_args() - stats = build_modelopt_fp8_transformer( + stats = convert_modelopt_fp8_checkpoint( modelopt_hf_dir=args.modelopt_hf_dir, modelopt_backbone_ckpt=args.modelopt_backbone_ckpt, output_dir=args.output_dir, From c775c31cb30f0f6c32e742fb4d5fbed5be384573 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sun, 12 Apr 2026 20:26:54 -0700 Subject: [PATCH 11/52] Simplify test_chunked_prefill; remove redundant tests (#22652) --- .../scheduler/test_chunked_prefill.py | 35 ------------ .../scheduler/test_mixed_chunked_prefill.py | 55 +++++++++++++++++++ 2 files changed, 55 insertions(+), 35 deletions(-) delete mode 100644 test/registered/scheduler/test_chunked_prefill.py create mode 100644 test/registered/scheduler/test_mixed_chunked_prefill.py diff --git a/test/registered/scheduler/test_chunked_prefill.py b/test/registered/scheduler/test_chunked_prefill.py deleted file mode 100644 index a4257dedb1b3..000000000000 --- a/test/registered/scheduler/test_chunked_prefill.py +++ /dev/null @@ -1,35 +0,0 @@ -""" -python3 -m unittest test_chunked_prefill.TestChunkedPrefill.test_mixed_chunked_prefill_without_radix_cache -""" - -import unittest - -from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci -from sglang.test.test_utils import CustomTestCase, run_mmlu_test, run_mulit_request_test - -register_cuda_ci(est_time=360, suite="stage-b-test-1-gpu-small") -register_amd_ci(est_time=312, suite="stage-b-test-1-gpu-small-amd") - - -class TestChunkedPrefill(CustomTestCase): - def test_chunked_prefill(self): - run_mmlu_test(disable_radix_cache=False, enable_mixed_chunk=False) - - def test_mixed_chunked_prefill(self): - run_mmlu_test(disable_radix_cache=False, enable_mixed_chunk=True) - - def test_chunked_prefill_without_radix_cache(self): - run_mmlu_test(disable_radix_cache=True, enable_mixed_chunk=False) - - def test_mixed_chunked_prefill_without_radix_cache(self): - run_mmlu_test(disable_radix_cache=True, enable_mixed_chunk=True) - - def test_mixed_chunked_prefill_multi_requests(self): - run_mulit_request_test( - enable_mixed_chunk=True, - chunked_prefill_size=2048, - ) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/scheduler/test_mixed_chunked_prefill.py b/test/registered/scheduler/test_mixed_chunked_prefill.py new file mode 100644 index 000000000000..895943394a9f --- /dev/null +++ b/test/registered/scheduler/test_mixed_chunked_prefill.py @@ -0,0 +1,55 @@ +import unittest + +from sglang.srt.environ import envs +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci +from sglang.test.kits.eval_accuracy_kit import GSM8KMixin +from sglang.test.test_utils import ( + DEFAULT_MODEL_NAME_FOR_TEST, + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +register_cuda_ci(est_time=180, suite="stage-b-test-1-gpu-small") +register_amd_ci(est_time=180, suite="stage-b-test-1-gpu-small-amd") + + +class TestMixedChunkedPrefill(GSM8KMixin, CustomTestCase): + model = DEFAULT_MODEL_NAME_FOR_TEST + base_url = DEFAULT_URL_FOR_TEST + gsm8k_accuracy_thres = 0.62 + + extra_args = [ + "--enable-mixed-chunk", + "--chunked-prefill-size", + "32", + ] + + @classmethod + def setUpClass(cls): + with envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.override(2): + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=cls.extra_args, + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + +class TestMixedChunkedPrefillNoRadixCache(TestMixedChunkedPrefill): + extra_args = [ + "--enable-mixed-chunk", + "--chunked-prefill-size", + "32", + "--disable-radix-cache", + ] + + +if __name__ == "__main__": + unittest.main() From fae0a2fc3c51ce383759d22ade6fd638a9716356 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com> Date: Mon, 13 Apr 2026 12:23:32 +0800 Subject: [PATCH 12/52] [codex] Add LTX-2.3 benchmark skill recipes (#22631) --- .../SKILL.md | 2 +- .../benchmark-and-profile.md | 57 +++++++++++++++++++ .../scripts/bench_diffusion_denoise.py | 47 +++++++++++++-- .../sglang-diffusion-performance/SKILL.md | 31 ++++++++++ 4 files changed, 130 insertions(+), 7 deletions(-) diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/SKILL.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/SKILL.md index d91f2c2a21e2..29b12dab0249 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/SKILL.md +++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/SKILL.md @@ -27,7 +27,7 @@ Before running any benchmark, profiler, or kernel-validation command: ## Main Reference -- [benchmark-and-profile.md](benchmark-and-profile.md) — canonical denoise benchmark, perf dump, and `torch.profiler` workflow; uses the checked-in nightly-aligned presets, including `LTX-2` two-stage +- [benchmark-and-profile.md](benchmark-and-profile.md) — canonical denoise benchmark, perf dump, and `torch.profiler` workflow; uses the checked-in nightly-aligned presets, plus `LTX-2`, `LTX-2.3` one-stage, and `LTX-2.3` two-stage benchmark recipes - [existing-fast-paths.md](existing-fast-paths.md) — map bottlenecks to existing fused kernels, packed QKV paths, fused `QK norm + RoPE`, and distributed overlap patterns before proposing new code - [scripts/diffusion_skill_env.py](scripts/diffusion_skill_env.py) — preflight helper: repo root discovery via `sglang.__file__`, write-access probe, benchmark/profile output directories, idle GPU selection - [scripts/bench_diffusion_denoise.py](scripts/bench_diffusion_denoise.py) — end-to-end denoise benchmark preset runner via `sglang generate`; use `--list-models` to inspect preset order, then save perf dumps by label and compare them with `compare_perf.py` diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/benchmark-and-profile.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/benchmark-and-profile.md index be4aac2158e4..4ea949fb7c3a 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/benchmark-and-profile.md +++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/benchmark-and-profile.md @@ -100,6 +100,26 @@ PYTHONPATH=python python3 \ --output-dir "${BENCH_DIR}" ``` +Run the `LTX-2.3` one-stage skill preset: + +```bash +PYTHONPATH=python python3 \ + python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py \ + --model ltx23-one-stage \ + --label baseline \ + --output-dir "${BENCH_DIR}" +``` + +Run the `LTX-2.3` two-stage skill preset: + +```bash +PYTHONPATH=python python3 \ + python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py \ + --model ltx23-two-stage \ + --label baseline \ + --output-dir "${BENCH_DIR}" +``` + Run the full preset sweep: ```bash @@ -123,6 +143,8 @@ Nightly-aligned presets come first; skill-only presets stay available after them | `wan-ti2v` | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | Yes: `wan22_ti2v_5b_720p` | Uses the nightly cat image and motion prompt | | `ltx2` | `Lightricks/LTX-2` | Yes: `ltx2_twostage_t2v` | Uses `LTX2TwoStagePipeline`; nightly-aligned prompt, negative prompt, 1536x1024, 121 frames, fps 24, seed 1234 | | `wan-i2v` | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | Yes: `wan22_i2v_a14b_720p` | Aligned to nightly CFG-parallel 4-GPU launch | +| `ltx23-one-stage` | `Lightricks/LTX-2.3` | No | Skill-only extra preset for the native `LTX-2.3` one-stage baseline; 2 GPUs, 768x512, 121 frames, fps 24, 30 steps, guidance 3.0, seed 1234 | +| `ltx23-two-stage` | `Lightricks/LTX-2.3` | No | Skill-only extra preset for the native `LTX-2.3` two-stage path; uses `LTX2TwoStagePipeline`, 2 GPUs, 1536x1024, 121 frames, fps 24, 30 steps, guidance 3.0, seed 1234 | | `hunyuanvideo` | `hunyuanvideo-community/HunyuanVideo` | No | Skill-only extra preset | | `mova-720p` | `OpenMOSS-Team/MOVA-720p` | No | Skill-only extra preset | | `helios` | `BestWishYsh/Helios-Base` | No | Skill-only extra preset | @@ -151,6 +173,41 @@ After [PR #20707](https://github.com/sgl-project/sglang/pull/20707), `LTX2TwoStagePipeline` is a native path. The spatial upsampler and distilled LoRA are auto-resolved from the same model snapshot unless you override them. +### Manual command example: LTX-2.3 One-Stage + +```bash +sglang generate \ + --model-path=Lightricks/LTX-2.3 \ + --prompt="A beautiful sunset over the ocean" \ + --negative-prompt="shaky, glitchy, low quality, worst quality, deformed, distorted, disfigured, motion smear, motion artifacts, fused fingers, bad anatomy, weird hand, ugly, transition, static." \ + --width=768 --height=512 \ + --num-frames=121 --fps=24 \ + --num-inference-steps=30 --guidance-scale=3.0 \ + --seed=1234 --num-gpus=2 \ + --save-output --enable-torch-compile --warmup +``` + +Use this when you want the native `LTX2Pipeline` baseline for `LTX-2.3` at the +validated one-stage resolution. + +### Manual command example: LTX-2.3 Two-Stage + +```bash +sglang generate \ + --model-path=Lightricks/LTX-2.3 \ + --pipeline-class-name=LTX2TwoStagePipeline \ + --prompt="A beautiful sunset over the ocean" \ + --negative-prompt="shaky, glitchy, low quality, worst quality, deformed, distorted, disfigured, motion smear, motion artifacts, fused fingers, bad anatomy, weird hand, ugly, transition, static." \ + --width=1536 --height=1024 \ + --num-frames=121 --fps=24 \ + --num-inference-steps=30 --guidance-scale=3.0 \ + --seed=1234 --num-gpus=2 \ + --save-output --enable-torch-compile --warmup +``` + +This matches the new `ltx23-two-stage` skill preset and is a good benchmark target for +the recently merged `LTX-2.3` two-stage path. + ### Manual command example: Wan2.2-I2V-A14B 720P ```bash diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py index 002c13429e3a..224462896d2d 100755 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py +++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py @@ -12,7 +12,7 @@ # Tag the run for later compare_perf.py usage python3 python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py --model flux --label tuned - # All 12 preset models + # All 14 preset models python3 python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py --all # Show preset order, model path, and nightly mapping @@ -193,6 +193,39 @@ ], }, # 10. Skill-only extra preset + "ltx23-one-stage": { + "path": "Lightricks/LTX-2.3", + "prompt": "A beautiful sunset over the ocean", + "negative_prompt": "shaky, glitchy, low quality, worst quality, deformed, distorted, disfigured, motion smear, motion artifacts, fused fingers, bad anatomy, weird hand, ugly, transition, static.", + "seed": 1234, + "extra_args": [ + "--width=768", + "--height=512", + "--num-frames=121", + "--fps=24", + "--num-inference-steps=30", + "--guidance-scale=3.0", + "--num-gpus=2", + ], + }, + # 11. Skill-only extra preset + "ltx23-two-stage": { + "path": "Lightricks/LTX-2.3", + "prompt": "A beautiful sunset over the ocean", + "negative_prompt": "shaky, glitchy, low quality, worst quality, deformed, distorted, disfigured, motion smear, motion artifacts, fused fingers, bad anatomy, weird hand, ugly, transition, static.", + "seed": 1234, + "extra_args": [ + "--pipeline-class-name=LTX2TwoStagePipeline", + "--width=1536", + "--height=1024", + "--num-frames=121", + "--fps=24", + "--num-inference-steps=30", + "--guidance-scale=3.0", + "--num-gpus=2", + ], + }, + # 12. Skill-only extra preset "hunyuanvideo": { "path": "hunyuanvideo-community/HunyuanVideo", "prompt": "A cat and a dog baking a cake together in a kitchen. The cat is carefully measuring flour, while the dog is stirring the batter with a wooden spoon. The kitchen is cozy, with sunlight streaming through the window.", @@ -205,7 +238,7 @@ "--num-inference-steps=30", ], }, - # 11. Skill-only extra preset + # 13. Skill-only extra preset # Requires: /inputs/diffusion_benchmark/figs/mova_single_person.jpg "mova-720p": { "path": "OpenMOSS-Team/MOVA-720p", @@ -221,7 +254,7 @@ "--num-inference-steps=2", ], }, - # 12. Skill-only extra preset + # 14. Skill-only extra preset "helios": { "path": "BestWishYsh/Helios-Base", "prompt": "A curious raccoon", @@ -247,6 +280,8 @@ def required_gpus_for_model(model_key: str) -> int: return 4 if model_key == "mova-720p": return 4 + if model_key in {"ltx23-one-stage", "ltx23-two-stage"}: + return 2 return 1 @@ -260,11 +295,11 @@ def print_model_catalog(): print("=" * 95) print("MODEL PRESETS — Nightly-aligned first, skill-only extras after") print("=" * 95) - print(f"{'Preset':<14} {'Nightly':<28} {'Model Path':<46} {'GPUs':>4}") + print(f"{'Preset':<17} {'Nightly':<28} {'Model Path':<46} {'GPUs':>4}") print("-" * 95) for model_key, cfg in MODELS.items(): print( - f"{model_key:<14} {model_nightly_case_id(model_key):<28} {cfg['path']:<46} {required_gpus_for_model(model_key):>4}" + f"{model_key:<17} {model_nightly_case_id(model_key):<28} {cfg['path']:<46} {required_gpus_for_model(model_key):>4}" ) print("-" * 112) print( @@ -474,7 +509,7 @@ def main(): choices=list(MODELS.keys()), help="Model to benchmark (default: flux)", ) - parser.add_argument("--all", action="store_true", help="Benchmark all 12 models") + parser.add_argument("--all", action="store_true", help="Benchmark all 14 models") parser.add_argument( "--list-models", action="store_true", diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md index e81449004096..5c8d91e01f1f 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md +++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md @@ -88,6 +88,37 @@ sglang generate --model-path Lightricks/LTX-2 \ Note: this generate recipe is aligned with the nightly comparison case `ltx2_twostage_t2v`. After [PR #20707](https://github.com/sgl-project/sglang/pull/20707), `LTX2TwoStagePipeline` is a native path and auto-resolves the spatial upsampler plus distilled LoRA from the same model snapshot unless you override them. +### Native baseline, 2 GPUs: LTX-2.3 one-stage + +```bash +sglang generate --model-path Lightricks/LTX-2.3 \ + --prompt "A beautiful sunset over the ocean" \ + --negative-prompt "shaky, glitchy, low quality, worst quality, deformed, distorted, disfigured, motion smear, motion artifacts, fused fingers, bad anatomy, weird hand, ugly, transition, static." \ + --width 768 --height 512 \ + --num-frames 121 --fps 24 \ + --num-inference-steps 30 --guidance-scale 3.0 \ + --seed 1234 --num-gpus 2 \ + --enable-torch-compile --warmup --save-output +``` + +Note: use this as the native `LTX2Pipeline` baseline for `LTX-2.3`. It keeps the validated one-stage resolution and explicit `LTX-2.3` sampling defaults, and matches the `ltx23-one-stage` benchmark preset in `sglang-diffusion-benchmark-profile`. + +### Benchmark target, 2 GPUs: LTX-2.3 two-stage + +```bash +sglang generate --model-path Lightricks/LTX-2.3 \ + --pipeline-class-name LTX2TwoStagePipeline \ + --prompt "A beautiful sunset over the ocean" \ + --negative-prompt "shaky, glitchy, low quality, worst quality, deformed, distorted, disfigured, motion smear, motion artifacts, fused fingers, bad anatomy, weird hand, ugly, transition, static." \ + --width 1536 --height 1024 \ + --num-frames 121 --fps 24 \ + --num-inference-steps 30 --guidance-scale 3.0 \ + --seed 1234 --num-gpus 2 \ + --enable-torch-compile --warmup --save-output +``` + +Note: this is the recommended benchmark command for the new `LTX-2.3` two-stage path. It uses the native `LTX2TwoStagePipeline` and matches the `ltx23-two-stage` benchmark preset in `sglang-diffusion-benchmark-profile`. + ### Maximum speed, image model, single GPU, lossless ```bash From 4dbd59850b8bad54598d164ebce9d17ce378ecc9 Mon Sep 17 00:00:00 2001 From: Mohammad Miadh Angkad Date: Mon, 13 Apr 2026 12:41:42 +0800 Subject: [PATCH 13/52] Add bfloat16 KV cache validation for HiSparse (#22505) --- python/sglang/srt/server_args.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 64ea2c1856fd..81debbb9be68 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -6565,6 +6565,12 @@ def check_server_args(self): f"Please use --nsa-{label}-backend=flashmla_sparse or omit it." ) + if self.kv_cache_dtype != "bfloat16": + raise ValueError( + f"HiSparse requires bfloat16 KV cache, but got --kv-cache-dtype={self.kv_cache_dtype}. " + f"Please use --kv-cache-dtype=bfloat16." + ) + assert ( self.schedule_conservativeness >= 0 ), "schedule_conservativeness must be non-negative" From 3f4fbc165d83e17ce3732471465121a326d9db90 Mon Sep 17 00:00:00 2001 From: Alison Shao <54658187+alisonshao@users.noreply.github.com> Date: Sun, 12 Apr 2026 23:48:40 -0500 Subject: [PATCH 14/52] Upgrade CI default CUDA version from 12.9 to 13.0 (#21441) --- .github/workflows/pr-test.yml | 31 +++++++++---------- python/pyproject.toml | 8 ++--- .../cuda/ci_download_flashinfer_jit_cache.sh | 2 +- scripts/ci/cuda/ci_install_deepep.sh | 9 ++---- scripts/ci/cuda/ci_install_dependency.sh | 16 +++++++--- 5 files changed, 33 insertions(+), 33 deletions(-) diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 8feb364bd10a..7c50a46d53f2 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -429,10 +429,7 @@ jobs: matrix: include: - python-version: "3.10" - cuda-version: "12.9" - # Add back when CUDA 13.0 is supported on CI - # - python-version: "3.10" - # cuda-version: "13.0" + cuda-version: "13.0" name: Build Wheel steps: - name: Cleanup @@ -480,7 +477,7 @@ jobs: matrix: include: - python-version: "3.10" - cuda-version: "12.9" + cuda-version: "13.0" name: Build Wheel Arm steps: - name: Cleanup @@ -587,7 +584,7 @@ jobs: with: path: sgl-kernel/dist/ merge-multiple: true - pattern: wheel-python3.10-cuda12.9 + pattern: wheel-python3.10-cuda13.0 - name: Install dependencies timeout-minutes: 20 @@ -694,7 +691,7 @@ jobs: with: path: sgl-kernel/dist/ merge-multiple: true - pattern: wheel-python3.10-cuda12.9 + pattern: wheel-python3.10-cuda13.0 - name: Install dependencies timeout-minutes: 20 @@ -752,7 +749,7 @@ jobs: with: path: sgl-kernel/dist/ merge-multiple: true - pattern: wheel-python3.10-cuda12.9 + pattern: wheel-python3.10-cuda13.0 - name: Install dependencies timeout-minutes: 20 @@ -807,7 +804,7 @@ jobs: with: path: sgl-kernel/dist/ merge-multiple: true - pattern: wheel-python3.10-cuda12.9 + pattern: wheel-python3.10-cuda13.0 - name: Install dependencies timeout-minutes: 20 @@ -860,7 +857,7 @@ jobs: with: path: sgl-kernel/dist/ merge-multiple: true - pattern: wheel-python3.10-cuda12.9 + pattern: wheel-python3.10-cuda13.0 - name: Install dependencies timeout-minutes: 20 @@ -949,7 +946,7 @@ jobs: with: path: sgl-kernel/dist/ merge-multiple: true - pattern: wheel-python3.10-cuda12.9 + pattern: wheel-python3.10-cuda13.0 - name: Install dependencies timeout-minutes: 20 @@ -1003,7 +1000,7 @@ jobs: with: path: sgl-kernel/dist/ merge-multiple: true - pattern: wheel-python3.10-cuda12.9 + pattern: wheel-python3.10-cuda13.0 - name: Install dependencies timeout-minutes: 20 @@ -1073,7 +1070,7 @@ jobs: with: path: sgl-kernel/dist/ merge-multiple: true - pattern: wheel-python3.10-cuda12.9 + pattern: wheel-python3.10-cuda13.0 - name: Install dependencies timeout-minutes: 20 @@ -1123,7 +1120,7 @@ jobs: with: path: sgl-kernel/dist/ merge-multiple: true - pattern: wheel-python3.10-cuda12.9 + pattern: wheel-python3.10-cuda13.0 - name: Install dependencies timeout-minutes: 20 @@ -1183,7 +1180,7 @@ jobs: with: path: sgl-kernel/dist/ merge-multiple: true - pattern: wheel-python3.10-cuda12.9 + pattern: wheel-python3.10-cuda13.0 - name: Install dependencies timeout-minutes: 20 @@ -1249,7 +1246,7 @@ jobs: with: path: sgl-kernel/dist/ merge-multiple: true - pattern: wheel-python3.10-cuda12.9 + pattern: wheel-python3.10-cuda13.0 - name: Install dependencies timeout-minutes: 20 @@ -1303,7 +1300,7 @@ jobs: # with: # path: sgl-kernel/dist/ # merge-multiple: true - # pattern: wheel-python3.10-cuda12.9-aarch64 + # pattern: wheel-python3.10-cuda13.0-aarch64 # # - name: Install dependencies # timeout-minutes: 20 diff --git a/python/pyproject.toml b/python/pyproject.toml index 538ad4908d7f..7eeb44f59783 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -22,7 +22,7 @@ dependencies = [ "blobfile==3.0.0", "build", "compressed-tensors", - "cuda-python==12.9", + "cuda-python>=13.0", "decord2 ; sys_platform == 'linux' and (platform_machine == 'aarch64' or platform_machine == 'arm64' or platform_machine == 'armv7l')", "datasets", "einops", @@ -86,14 +86,14 @@ url = "https://pypi.org/simple" default = true [[tool.uv.index]] -name = "torch-cu129" -url = "https://download.pytorch.org/whl/cu129" +name = "torch-cu130" +url = "https://download.pytorch.org/whl/cu130" explicit = true [tool.uv.sources] torch = [ { index = "pypi", marker = "platform_machine == 'x86_64'"}, - { index = "torch-cu129", marker = "platform_machine == 'aarch64'"}, + { index = "torch-cu130", marker = "platform_machine == 'aarch64'"}, ] [project.optional-dependencies] diff --git a/scripts/ci/cuda/ci_download_flashinfer_jit_cache.sh b/scripts/ci/cuda/ci_download_flashinfer_jit_cache.sh index ab41dcef22c0..6358b5caf9b8 100755 --- a/scripts/ci/cuda/ci_download_flashinfer_jit_cache.sh +++ b/scripts/ci/cuda/ci_download_flashinfer_jit_cache.sh @@ -5,7 +5,7 @@ # Required environment (caller must export or set): # UNINSTALL_JIT_CACHE — literal true/false (skip download when false) # FLASHINFER_PYTHON_REQUIRED — e.g. from python/pyproject.toml (flashinfer_python) -# CU_VERSION — e.g. cu129 +# CU_VERSION — e.g. cu130 # PIP_CMD — e.g. "pip" or "uv pip" # PIP_INSTALL_SUFFIX — extra pip args for this runner set -euxo pipefail diff --git a/scripts/ci/cuda/ci_install_deepep.sh b/scripts/ci/cuda/ci_install_deepep.sh index bb4185799571..b6e1887514be 100755 --- a/scripts/ci/cuda/ci_install_deepep.sh +++ b/scripts/ci/cuda/ci_install_deepep.sh @@ -100,13 +100,8 @@ if [ "$GRACE_BLACKWELL" = "1" ]; then if [ "$CUDA_VERSION" = "12.8" ]; then CHOSEN_TORCH_CUDA_ARCH_LIST='10.0' elif awk -v ver="$CUDA_VERSION" 'BEGIN {exit !(ver > 12.8)}'; then - # With cuda > 12.8, the compiler supports 10.3, so we should use - # CHOSEN_TORCH_CUDA_ARCH_LIST='10.0;10.3' - # - # However, our CI machine has a weird setup and nvidia-smi reports wrong CUDA version in the container. - # The container is actually cuda 12.8, but nvidia-smi reports 13.0, leading to compilation errors. so we - # drop 10.3. - CHOSEN_TORCH_CUDA_ARCH_LIST='10.0' + # CUDA > 12.8 supports sm_103 (Blackwell) + CHOSEN_TORCH_CUDA_ARCH_LIST='10.0;10.3' else echo "Unsupported CUDA version for Grace Blackwell: $CUDA_VERSION" && exit 1 fi && \ diff --git a/scripts/ci/cuda/ci_install_dependency.sh b/scripts/ci/cuda/ci_install_dependency.sh index 7735977ef278..9ac79327c30e 100755 --- a/scripts/ci/cuda/ci_install_dependency.sh +++ b/scripts/ci/cuda/ci_install_dependency.sh @@ -23,7 +23,7 @@ set -euxo pipefail # Configuration & timing # ------------------------------------------------------------------------------ # Set up environment variables -CU_VERSION="cu129" +CU_VERSION="cu130" # Nvidia package versions we override (torch pins older versions). # Used both as pip constraints during install and for post-install verification. @@ -244,7 +244,15 @@ if [ "${CUSTOM_BUILD_SGL_KERNEL:-}" = "true" ] && [ -d "sgl-kernel/dist" ]; then else WHEEL_ARCH="x86_64" fi - $PIP_CMD install sgl-kernel/dist/sglang_kernel-${SGL_KERNEL_VERSION_FROM_KERNEL}-cp310-abi3-manylinux2014_${WHEEL_ARCH}.whl --force-reinstall $PIP_INSTALL_SUFFIX + # Wheel may have +cuXYZ suffix (e.g. sglang_kernel-0.4.0+cu130-...) depending on CUDA version + KERNEL_WHL=$(ls sgl-kernel/dist/sglang_kernel-${SGL_KERNEL_VERSION_FROM_KERNEL}*-cp310-abi3-manylinux2014_${WHEEL_ARCH}.whl 2>/dev/null | head -1) + if [ -z "$KERNEL_WHL" ]; then + echo "ERROR: No matching sgl-kernel wheel found in sgl-kernel/dist/ for version ${SGL_KERNEL_VERSION_FROM_KERNEL} arch ${WHEEL_ARCH}" + ls -alh sgl-kernel/dist/ + exit 1 + fi + echo "Installing sgl-kernel wheel: $KERNEL_WHL" + $PIP_CMD install "$KERNEL_WHL" --force-reinstall $PIP_INSTALL_SUFFIX elif [ "${CUSTOM_BUILD_SGL_KERNEL:-}" = "true" ] && [ ! -d "sgl-kernel/dist" ]; then # CUSTOM_BUILD_SGL_KERNEL was set but artifacts not available (e.g., stage rerun without wheel build) # Fail instead of falling back to PyPI - we need to test the built kernel, not PyPI version @@ -318,10 +326,10 @@ mark_step_done "Install extra dependency" # Fix other dependencies # ------------------------------------------------------------------------------ # Fix CUDA version mismatch between torch and torchaudio. -# PyPI's torch 2.9.1 bundles cu128 but torchaudio from pytorch.org/cu129 uses cu129. +# PyPI's torch bundles a specific CUDA version but torchaudio from pytorch.org/cu130 may use a different one. # This mismatch causes torchaudio's C extension to fail loading, producing: # "partially initialized module 'torchaudio' has no attribute 'lib'" -# We cannot replace torch with cu129 (breaks sgl_kernel ABI), so instead we reinstall +# We cannot replace torch with a different CUDA version (breaks sgl_kernel ABI), so instead we reinstall # torchaudio/torchvision from an index matching torch's CUDA version. TORCH_CUDA_VER=$(python3 -c "import torch; v=torch.version.cuda; parts=v.split('.'); print(f'cu{parts[0]}{parts[1]}')") echo "Detected torch CUDA version: ${TORCH_CUDA_VER}" From 305b42935afbb69c88e19116427ce9372ace4c4b Mon Sep 17 00:00:00 2001 From: Zhangheng Date: Mon, 13 Apr 2026 12:49:18 +0800 Subject: [PATCH 15/52] [HiSparse]: Add benchmark for hisparse kernel (#22187) --- .../jit_kernel/benchmark/bench_hisparse.py | 190 ++++++++++++++++++ 1 file changed, 190 insertions(+) create mode 100644 python/sglang/jit_kernel/benchmark/bench_hisparse.py diff --git a/python/sglang/jit_kernel/benchmark/bench_hisparse.py b/python/sglang/jit_kernel/benchmark/bench_hisparse.py new file mode 100644 index 000000000000..6eb5fb5f0596 --- /dev/null +++ b/python/sglang/jit_kernel/benchmark/bench_hisparse.py @@ -0,0 +1,190 @@ +import itertools +from typing import Dict, Tuple + +import torch +import triton +import triton.testing + +from sglang.jit_kernel.benchmark.utils import DEFAULT_DEVICE, DEFAULT_DTYPE +from sglang.jit_kernel.hisparse import load_cache_to_device_buffer_mla +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=12, suite="stage-b-kernel-benchmark-1-gpu-large") + +DEVICE = DEFAULT_DEVICE +DTYPE = DEFAULT_DTYPE +TOP_K = 2048 +ITEM_SIZE_BYTES = 512 +MISS_RATES = [0.2, 0.001] +ROUNDS = 5 +WARMUP_ROUNDS = 5 +BATCH_SIZES = [1, 10, 100] +HOT_BUFFER_SIZES = [4096, 8192] +CONFIGS = [ + ( + batch_size, + hot_buffer_size, + miss_rate, + batch_size * round(TOP_K * miss_rate), + ) + for batch_size, hot_buffer_size, miss_rate in itertools.product( + BATCH_SIZES, HOT_BUFFER_SIZES, MISS_RATES + ) +] + +LINE_VALS = ["jit"] +LINE_NAMES = ["SGL JIT Kernel"] +STYLES = [("blue", "--")] + + +def _make_top_k_tokens( + num_hits: int, num_misses: int, hot_buffer_size: int +) -> torch.Tensor: + hit_tokens = torch.arange(num_hits, dtype=torch.int32, device=DEVICE) + miss_tokens = hot_buffer_size + torch.arange( + num_misses, dtype=torch.int32, device=DEVICE + ) + return torch.cat([hit_tokens, miss_tokens]) + + +def _miss_tokens_per_req(miss_rate: float) -> int: + return round(TOP_K * miss_rate) + + +def _build_inputs( + batch_size: int, hot_buffer_size: int, miss_rate: float +) -> Dict[str, torch.Tensor | int]: + dtype_bytes = torch.empty((), dtype=DTYPE).element_size() + kv_dim = ITEM_SIZE_BYTES // dtype_bytes + padded_buffer_size = hot_buffer_size + 1 + seq_len = hot_buffer_size + TOP_K + 1 + num_misses = _miss_tokens_per_req(miss_rate) + num_hits = TOP_K - num_misses + + top_k_row = _make_top_k_tokens(num_hits, num_misses, hot_buffer_size) + top_k_tokens = top_k_row.view(1, -1).repeat(batch_size, 1).contiguous() + + host_stride = seq_len + total_host_tokens = batch_size * host_stride + host_cache = torch.empty( + (total_host_tokens, 1, kv_dim), dtype=DTYPE, device="cpu", pin_memory=True + ) + host_cache.copy_(torch.randn_like(host_cache)) + + total_device_tokens = batch_size * padded_buffer_size + device_buffer = torch.empty( + (total_device_tokens, 1, kv_dim), dtype=DTYPE, device=DEVICE + ) + device_buffer.normal_() + + device_buffer_locs = torch.arange( + total_device_tokens, dtype=torch.int32, device=DEVICE + ).view(batch_size, padded_buffer_size) + device_buffer_tokens = torch.full( + (batch_size, padded_buffer_size), -1, dtype=torch.int32, device=DEVICE + ) + device_buffer_tokens[:, :hot_buffer_size] = torch.arange( + hot_buffer_size, dtype=torch.int32, device=DEVICE + ) + + lru_slots = ( + torch.arange(hot_buffer_size, dtype=torch.int16, device=DEVICE) + .view(1, -1) + .repeat(batch_size, 1) + ) + + return { + "top_k_tokens": top_k_tokens, + "device_buffer_tokens": device_buffer_tokens, + "initial_device_buffer_tokens": device_buffer_tokens.clone(), + "host_cache_locs": torch.arange( + total_host_tokens, dtype=torch.int64, device=DEVICE + ).view(batch_size, host_stride), + "device_buffer_locs": device_buffer_locs, + "host_cache": host_cache, + "device_buffer": device_buffer, + "top_k_device_locs": torch.empty( + (batch_size, TOP_K), dtype=torch.int32, device=DEVICE + ), + "req_pool_indices": torch.arange(batch_size, dtype=torch.int64, device=DEVICE), + "seq_lens": torch.full( + (batch_size,), seq_len, dtype=torch.int32, device=DEVICE + ), + "lru_slots": lru_slots, + "initial_lru_slots": lru_slots.clone(), + "num_real_reqs": torch.tensor([batch_size], dtype=torch.int32, device=DEVICE), + } + + +def _time_kernel(batch_size: int, hot_buffer_size: int, miss_rate: float) -> float: + state = _build_inputs(batch_size, hot_buffer_size, miss_rate) + + def run_once(): + state["device_buffer_tokens"].copy_(state["initial_device_buffer_tokens"]) + state["lru_slots"].copy_(state["initial_lru_slots"]) + state["top_k_device_locs"].fill_(-1) + load_cache_to_device_buffer_mla( + top_k_tokens=state["top_k_tokens"], + device_buffer_tokens=state["device_buffer_tokens"], + host_cache_locs=state["host_cache_locs"], + device_buffer_locs=state["device_buffer_locs"], + host_cache=state["host_cache"], + device_buffer=state["device_buffer"], + top_k_device_locs=state["top_k_device_locs"], + req_pool_indices=state["req_pool_indices"], + seq_lens=state["seq_lens"], + lru_slots=state["lru_slots"], + item_size_bytes=ITEM_SIZE_BYTES, + num_top_k=TOP_K, + hot_buffer_size=hot_buffer_size, + block_size=1024, + num_real_reqs=state["num_real_reqs"], + ) + + run_once() + torch.cuda.synchronize() + for _ in range(WARMUP_ROUNDS): + run_once() + torch.cuda.synchronize() + + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(ROUNDS): + run_once() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) * 1000.0 / ROUNDS + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["batch_size", "hot_buffer_size", "miss_rate", "miss_tokens_cnt"], + x_vals=CONFIGS, + line_arg="provider", + line_vals=LINE_VALS, + line_names=LINE_NAMES, + styles=STYLES, + ylabel="us", + plot_name="hisparse-latency", + args={}, + ) +) +def benchmark_latency( + batch_size: int, + hot_buffer_size: int, + miss_rate: float, + miss_tokens_cnt: int, + provider: str, +) -> Tuple[float, float, float]: + assert provider == "jit" + batch_size = int(batch_size) + hot_buffer_size = int(hot_buffer_size) + miss_rate = float(miss_rate) + assert miss_tokens_cnt == batch_size * _miss_tokens_per_req(miss_rate) + avg_us = _time_kernel(batch_size, hot_buffer_size, miss_rate) + return avg_us, avg_us, avg_us + + +if __name__ == "__main__": + benchmark_latency.run(print_data=True) From 5549d910c66fc8b2bfc72ce0c8c83e6797b001e5 Mon Sep 17 00:00:00 2001 From: Zhangheng Date: Mon, 13 Apr 2026 12:50:29 +0800 Subject: [PATCH 16/52] [hisparse]: Adding ci for hisparse kvcache-swap-in jit-kernel (#22155) --- .../sglang/jit_kernel/tests/test_hisparse.py | 299 ++++++++++++++++++ 1 file changed, 299 insertions(+) create mode 100644 python/sglang/jit_kernel/tests/test_hisparse.py diff --git a/python/sglang/jit_kernel/tests/test_hisparse.py b/python/sglang/jit_kernel/tests/test_hisparse.py new file mode 100644 index 000000000000..190f22355ff5 --- /dev/null +++ b/python/sglang/jit_kernel/tests/test_hisparse.py @@ -0,0 +1,299 @@ +import sys + +import pytest +import torch + +from sglang.jit_kernel.hisparse import load_cache_to_device_buffer_mla +from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=10, suite="stage-b-kernel-unit-1-gpu-large") +register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() + or is_npu() + or is_xpu() + or not (is_cuda() or is_hip()), + reason="HiSparse JIT tests require CUDA/ROCm.", +) + +DEVICE = "cuda" +DTYPE = torch.float32 +KV_DIM = 8 +HOT_BUFFER_SIZE = 4 +PADDED_BUFFER_SIZE = HOT_BUFFER_SIZE + 1 +HOST_CACHE_SIZE = 16 +DEVICE_CACHE_SIZE = 16 +ITEM_SIZE_BYTES = KV_DIM * torch.empty((), dtype=DTYPE).element_size() + + +def _host_cache() -> torch.Tensor: + host_cache = torch.empty( + (HOST_CACHE_SIZE, 1, KV_DIM), dtype=DTYPE, device="cpu", pin_memory=True + ) + host_cache.copy_(torch.arange(host_cache.numel(), dtype=DTYPE).view_as(host_cache)) + return host_cache + + +def _run_kernel( + *, + top_k_tokens: torch.Tensor, + device_buffer_tokens: torch.Tensor, + host_cache_locs: torch.Tensor, + device_buffer_locs: torch.Tensor, + host_cache: torch.Tensor, + device_buffer: torch.Tensor, + lru_slots: torch.Tensor, + seq_len: int | None = None, + seq_lens: torch.Tensor | None = None, + seq_lens_dtype: torch.dtype = torch.int32, + req_pool_indices: torch.Tensor | None = None, + num_real_reqs: int | None = None, +) -> torch.Tensor: + batch_size = top_k_tokens.shape[0] + if req_pool_indices is None: + req_pool_indices = torch.arange(batch_size, dtype=torch.int64, device=DEVICE) + if seq_lens is None: + seq_lens = torch.full( + (batch_size,), seq_len, dtype=seq_lens_dtype, device=DEVICE + ) + if num_real_reqs is None: + num_real_reqs = batch_size + + out = torch.full_like(top_k_tokens, -1) + load_cache_to_device_buffer_mla( + top_k_tokens=top_k_tokens, + device_buffer_tokens=device_buffer_tokens, + host_cache_locs=host_cache_locs, + device_buffer_locs=device_buffer_locs, + host_cache=host_cache, + device_buffer=device_buffer, + top_k_device_locs=out, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + lru_slots=lru_slots, + item_size_bytes=ITEM_SIZE_BYTES, + num_top_k=top_k_tokens.shape[1], + hot_buffer_size=HOT_BUFFER_SIZE, + page_size=1, + block_size=256, + num_real_reqs=torch.tensor([num_real_reqs], dtype=torch.int32, device=DEVICE), + ) + torch.cuda.synchronize() + return out + + +def _make_state( + device_buffer_locs_rows: list[list[int]], + device_buffer_tokens_rows: list[list[int]], + newest_tokens: list[int], +): + host_cache = _host_cache() + device_buffer = torch.full( + (DEVICE_CACHE_SIZE, 1, KV_DIM), -1, dtype=DTYPE, device=DEVICE + ) + device_buffer_locs = torch.tensor( + device_buffer_locs_rows, dtype=torch.int32, device=DEVICE + ) + device_buffer_tokens = torch.tensor( + device_buffer_tokens_rows, dtype=torch.int32, device=DEVICE + ) + lru_slots = ( + torch.arange(HOT_BUFFER_SIZE, dtype=torch.int16, device=DEVICE) + .view(1, -1) + .repeat(device_buffer_locs.shape[0], 1) + ) + host_cache_locs = ( + torch.arange(HOST_CACHE_SIZE, dtype=torch.int64, device=DEVICE) + .view(1, -1) + .repeat(device_buffer_locs.shape[0], 1) + ) + + # Slots 0..3 participate in LRU; slot 4 is the reserved newest slot. + for rid, newest_token in enumerate(newest_tokens): + for slot, token in enumerate(device_buffer_tokens_rows[rid][:HOT_BUFFER_SIZE]): + if token >= 0: + device_buffer[device_buffer_locs[rid, slot]].copy_( + host_cache[token].to(DEVICE, non_blocking=True) + ) + device_buffer[device_buffer_locs[rid, HOT_BUFFER_SIZE]].copy_( + host_cache[newest_token].to(DEVICE, non_blocking=True) + ) + torch.cuda.synchronize() + + return { + "host_cache": host_cache, + "device_buffer": device_buffer, + "device_buffer_locs": device_buffer_locs, + "device_buffer_tokens": device_buffer_tokens, + "lru_slots": lru_slots, + "host_cache_locs": host_cache_locs, + } + + +def _long_case(): + # One-request baseline used by the stateful cases below: + # req 0 LRU slots : [0, 1, 2, 3] + # req 0 cached tokens : slot0->1, slot1->4, slot2->2, slot3->5 + # req 0 physical locs : slot0->9, slot1->7, slot2->3, slot3->5 + # req 0 newest slot : slot4/newest -> token 7 at physical loc 11 + return _make_state([[9, 7, 3, 5, 11]], [[1, 4, 2, 5, -1]], [7]) + + +@pytest.mark.parametrize("seq_lens_dtype", [torch.int32, torch.int64]) +def test_load_cache_to_device_buffer_fast_path(seq_lens_dtype: torch.dtype) -> None: + host_cache = _host_cache() + device_buffer = torch.arange( + DEVICE_CACHE_SIZE * KV_DIM, dtype=DTYPE, device=DEVICE + ).view(DEVICE_CACHE_SIZE, 1, KV_DIM) + device_buffer_before = device_buffer.clone() + device_buffer_locs = torch.tensor( + [[13, 9, 5, 1, 15]], dtype=torch.int32, device=DEVICE + ) + device_buffer_tokens = torch.tensor( + [[10, 11, 12, 13, -1]], dtype=torch.int32, device=DEVICE + ) + device_buffer_tokens_before = device_buffer_tokens.clone() + lru_slots = torch.tensor([[0, 1, 2, 3]], dtype=torch.int16, device=DEVICE) + lru_slots_before = lru_slots.clone() + + # Short-sequence layout: + # token position 0 -> physical loc 13 + # token position 1 -> physical loc 9 + # token position 2 -> physical loc 5 + # + # seq_len <= HOT_BUFFER_SIZE should skip host loads and LRU mutations, + # so top_k_tokens acts like direct indexing into device_buffer_locs. + out = _run_kernel( + top_k_tokens=torch.tensor([[2, 0, 1]], dtype=torch.int32, device=DEVICE), + device_buffer_tokens=device_buffer_tokens, + host_cache_locs=torch.arange( + HOST_CACHE_SIZE, dtype=torch.int64, device=DEVICE + ).view(1, -1), + device_buffer_locs=device_buffer_locs, + host_cache=host_cache, + device_buffer=device_buffer, + lru_slots=lru_slots, + seq_len=3, + seq_lens_dtype=seq_lens_dtype, + ) + + assert torch.equal(out.cpu(), torch.tensor([[5, 13, 9]], dtype=torch.int32)) + assert torch.equal(device_buffer_tokens.cpu(), device_buffer_tokens_before.cpu()) + assert torch.equal(lru_slots.cpu(), lru_slots_before.cpu()) + assert torch.equal(device_buffer.cpu(), device_buffer_before.cpu()) + + +def test_load_cache_to_device_buffer_hits_newest_and_updates_lru() -> None: + state = _long_case() + + # Query [4, 2, 7]: + # 4 hits slot1 -> loc 7 + # 2 hits slot2 -> loc 3 + # 7 is the newest token -> reserved newest loc 11 + # + # Hits move to the MRU tail, so [0, 1, 2, 3] becomes [0, 3, 1, 2]. + out = _run_kernel( + top_k_tokens=torch.tensor([[4, 2, 7]], dtype=torch.int32, device=DEVICE), + seq_len=8, + **state, + ) + + assert torch.equal(out.cpu(), torch.tensor([[7, 3, 11]], dtype=torch.int32)) + assert torch.equal( + state["device_buffer_tokens"].cpu(), + torch.tensor([[1, 4, 2, 5, -1]], dtype=torch.int32), + ) + assert torch.equal( + state["lru_slots"].cpu(), torch.tensor([[0, 3, 1, 2]], dtype=torch.int16) + ) + + +def test_load_cache_to_device_buffer_miss_uses_updated_lru_slot() -> None: + state = _long_case() + + # Step 1: touch tokens [4, 2], so LRU becomes [0, 3, 1, 2]. + # Step 2: query token 6, which is a miss. + # The kernel should reuse the new LRU head slot0, whose physical loc is 9. + # This round has no regular hits, so the freshly loaded miss slot ends up at the tail. + _run_kernel( + top_k_tokens=torch.tensor([[4, 2]], dtype=torch.int32, device=DEVICE), + seq_len=8, + **state, + ) + out = _run_kernel( + top_k_tokens=torch.tensor([[6]], dtype=torch.int32, device=DEVICE), + seq_len=8, + **state, + ) + + assert torch.equal(out.cpu(), torch.tensor([[9]], dtype=torch.int32)) + assert torch.equal( + state["device_buffer_tokens"].cpu(), + torch.tensor([[6, 4, 2, 5, -1]], dtype=torch.int32), + ) + assert torch.equal( + state["lru_slots"].cpu(), torch.tensor([[3, 1, 2, 0]], dtype=torch.int16) + ) + assert torch.equal(state["device_buffer"][9].cpu(), state["host_cache"][6]) + + +def test_load_cache_to_device_buffer_batched_with_padding() -> None: + state = _make_state( + [ + [9, 7, 3, 5, 11], + [12, 10, 8, 6, 14], + [15, 4, 2, 1, 13], + ], + [ + [1, 4, 2, 5, -1], + [0, 1, 2, 3, -1], + [9, 8, 7, 6, -1], + ], + [7, 4, 5], + ) + padded_tokens_before = state["device_buffer_tokens"][2].clone() + padded_lru_before = state["lru_slots"][2].clone() + + # req 0: long path + # cached tokens/locs : 1@9, 4@7, 2@3, 5@5, newest 7@11 + # query [4, 6, 7] : hit loc 7, miss into slot0/loc 9, newest loc 11 + # LRU update : remaining evictables [2, 3], then miss [0], then hit [1] + # : [0, 1, 2, 3] -> [2, 3, 0, 1] + # + # req 1: fast path + # seq_len = 3 <= HOT_BUFFER_SIZE, so [2, 1, 0] maps directly to locs [8, 10, 12] + # + # req 2: padded block + # num_real_reqs = 2 means this row must be ignored entirely. + out = _run_kernel( + top_k_tokens=torch.tensor( + [[4, 6, 7], [2, 1, 0], [9, 8, 7]], dtype=torch.int32, device=DEVICE + ), + seq_lens=torch.tensor([8, 3, 8], dtype=torch.int32, device=DEVICE), + num_real_reqs=2, + **state, + ) + + assert torch.equal( + out.cpu(), + torch.tensor([[7, 9, 11], [8, 10, 12], [-1, -1, -1]], dtype=torch.int32), + ) + assert torch.equal( + state["device_buffer_tokens"][:2].cpu(), + torch.tensor([[6, 4, 2, 5, -1], [0, 1, 2, 3, -1]], dtype=torch.int32), + ) + assert torch.equal( + state["lru_slots"][:2].cpu(), + torch.tensor([[2, 3, 0, 1], [0, 1, 2, 3]], dtype=torch.int16), + ) + assert torch.equal( + state["device_buffer_tokens"][2].cpu(), padded_tokens_before.cpu() + ) + assert torch.equal(state["lru_slots"][2].cpu(), padded_lru_before.cpu()) + assert torch.equal(state["device_buffer"][9].cpu(), state["host_cache"][6]) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v", "-s"])) From 7d2c11970ce626109685b5d86285031b851e7074 Mon Sep 17 00:00:00 2001 From: Polisetty V R K Jyothendra Varma Date: Mon, 13 Apr 2026 10:46:24 +0530 Subject: [PATCH 17/52] [Intel GPU] Upgrade pytorch xpu version to 2.11 (#21908) Signed-off-by: P V R K Jyothendra Varma Co-authored-by: Ma Mingfei --- docker/xpu.Dockerfile | 4 ++-- docs/platforms/xpu.md | 2 +- python/pyproject_xpu.toml | 8 ++++---- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/docker/xpu.Dockerfile b/docker/xpu.Dockerfile index 90a2414da3f9..fff8cba56f38 100644 --- a/docker/xpu.Dockerfile +++ b/docker/xpu.Dockerfile @@ -48,7 +48,7 @@ RUN --mount=type=secret,id=github_token \ cd /home/sdp && \ . /home/sdp/miniforge3/bin/activate && \ conda activate py${PYTHON_VERSION} && \ - pip3 install torch==2.10.0+xpu torchao torchvision torchaudio==2.10.0+xpu triton-xpu==3.6.0 --index-url https://download.pytorch.org/whl/xpu + pip3 install torch==2.11.0+xpu torchao torchvision torchaudio==2.11.0+xpu --index-url https://download.pytorch.org/whl/xpu RUN --mount=type=secret,id=github_token \ cd /home/sdp && \ @@ -58,7 +58,7 @@ RUN --mount=type=secret,id=github_token \ git clone --branch ${SG_LANG_BRANCH} --single-branch ${SG_LANG_REPO} && \ cd sglang && cd python && \ cp pyproject_xpu.toml pyproject.toml && \ - pip install . && \ + pip install . --extra-index-url https://download.pytorch.org/whl/xpu && \ pip install xgrammar --no-deps && \ pip install msgspec blake3 py-cpuinfo compressed_tensors gguf partial_json_parser einops tabulate --root-user-action=ignore && \ conda install libsqlite=3.48.0 -y && \ diff --git a/docs/platforms/xpu.md b/docs/platforms/xpu.md index 259f3081c47f..4ffa39fe5015 100644 --- a/docs/platforms/xpu.md +++ b/docs/platforms/xpu.md @@ -30,7 +30,7 @@ conda create -n sgl-xpu python=3.12 -y conda activate sgl-xpu # Set PyTorch XPU as primary pip install channel to avoid installing the larger CUDA-enabled version and prevent potential runtime issues. -pip3 install torch==2.10.0+xpu torchao torchvision torchaudio triton-xpu==3.6.0 --index-url https://download.pytorch.org/whl/xpu +pip3 install torch==2.11.0+xpu torchao torchvision torchaudio==2.11.0+xpu --index-url https://download.pytorch.org/whl/xpu pip3 install xgrammar --no-deps # xgrammar will introduce CUDA-enabled triton which might conflict with XPU # Clone the SGLang code diff --git a/python/pyproject_xpu.toml b/python/pyproject_xpu.toml index 315d9e8509a0..821dc88e5a75 100644 --- a/python/pyproject_xpu.toml +++ b/python/pyproject_xpu.toml @@ -15,10 +15,10 @@ classifiers = [ ] dependencies = [ - "torch==2.10.0+xpu", - "torchcodec==0.10.0 ; sys_platform != 'linux' or (sys_platform == 'linux' and platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')", # torchcodec does not exist in those systems. torch==2.10.0 on XPU uses 0.10.0 + "torch==2.11.0+xpu", + "torchcodec==0.11.0 ; sys_platform != 'linux' or (sys_platform == 'linux' and platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')", # torchcodec does not exist in those systems. torch==2.11.0 on XPU uses 0.11.0 "av ; sys_platform == 'linux' and (platform_machine == 'aarch64' or platform_machine == 'arm64' or platform_machine == 'armv7l')", - "torchaudio==2.10.0+xpu", + "torchaudio==2.11.0+xpu", "torchvision", "sgl-kernel @ git+https://github.com/sgl-project/sgl-kernel-xpu.git", "IPython", @@ -58,7 +58,7 @@ dependencies = [ "soundfile==0.13.1", "tiktoken", "timm==1.0.16", - "torchao==0.9.0", + "torchao==0.9.0+xpu", "tqdm", "mistral_common>=1.9.0", "transformers==5.3.0", From d524f110ac677d8bb138bff0e2f306e0ba7b62be Mon Sep 17 00:00:00 2001 From: Mick Date: Mon, 13 Apr 2026 13:34:37 +0800 Subject: [PATCH 18/52] [diffusion] refactor: streamline denoising stages (#22633) --- .../configs/pipeline_configs/base.py | 34 +- .../configs/pipeline_configs/ltx_2.py | 18 +- .../runtime/pipelines_core/stages/__init__.py | 4 + .../pipelines_core/stages/denoising.py | 632 +++--- .../pipelines_core/stages/denoising_av.py | 1971 +---------------- .../pipelines_core/stages/ltx_2_denoising.py | 1240 +++++++++++ .../stages/model_specific_stages/__init__.py | 1 + .../stages/model_specific_stages/wan_ti2v.py | 174 ++ .../utils/diffusion/comparison_configs.json | 17 + 9 files changed, 1785 insertions(+), 2306 deletions(-) create mode 100644 python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py create mode 100644 python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/__init__.py create mode 100644 python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/wan_ti2v.py diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index b149b9cfb490..6406af765ab6 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py @@ -151,7 +151,6 @@ def maybe_unpad_latents(latents, batch): return latents -# config for a single pipeline @dataclass class PipelineConfig: """The base configuration class for a generation pipeline.""" @@ -360,10 +359,28 @@ def normalize_vae_encode(self, image_latents, vae): def preprocess_decoding(self, latents, server_args=None, vae=None): return latents + @staticmethod + def _gather_sp_tensor(tensor: torch.Tensor, *, dim: int) -> torch.Tensor: + """All-gather an SP-sharded tensor along the specified logical dimension.""" + return sequence_model_parallel_all_gather(tensor.contiguous(), dim=dim) + + @staticmethod + def _trim_sp_gather_padding( + tensor: torch.Tensor, *, orig_len: int | None, dim: int + ) -> torch.Tensor: + """Trim padding introduced before SP sharding back to the original length.""" + if orig_len is None: + return tensor + orig_len = int(orig_len) + if orig_len <= 0 or tensor.shape[dim] <= orig_len: + return tensor + slices = [slice(None)] * tensor.ndim + slices[dim] = slice(orig_len) + return tensor[tuple(slices)] + def gather_latents_for_sp(self, latents, batch=None): # For video latents [B, C, T_local, H, W], gather along time dim=2 - latents = sequence_model_parallel_all_gather(latents, dim=2) - return latents + return self._gather_sp_tensor(latents, dim=2) def can_shard_audio_latents_for_sp(self, audio_latents) -> bool: """Return whether this pipeline uses packed audio latents that can be SP-sharded.""" @@ -405,9 +422,9 @@ def gather_noise_pred_for_sp(self, batch, noise_pred): noise_pred = self.gather_latents_for_sp(noise_pred) raw_latent_shape = getattr(batch, "raw_latent_shape", None) if raw_latent_shape is not None and noise_pred.dim() == 3: - orig_s = raw_latent_shape[1] - if noise_pred.shape[1] > orig_s: - noise_pred = noise_pred[:, :orig_s, :] + noise_pred = self._trim_sp_gather_padding( + noise_pred, orig_len=raw_latent_shape[1], dim=1 + ) return noise_pred def preprocess_vae_image(self, batch, vae_image_processor): @@ -850,8 +867,7 @@ def shard_latents_for_sp(self, batch, latents): def gather_latents_for_sp(self, latents, batch=None): # For image latents [B, S_local, D], gather along sequence dim=1 - latents = sequence_model_parallel_all_gather(latents, dim=1) - return latents + return self._gather_sp_tensor(latents, dim=1) def _unpad_and_unpack_latents(self, latents, batch): vae_scale_factor = self.vae_config.arch_config.vae_scale_factor @@ -908,7 +924,7 @@ def gather_latents_for_sp(self, latents, batch=None): if latents.dim() != 4: return super().gather_latents_for_sp(latents, batch=batch) # Gather along dim=2 (H') to match shard_latents_for_sp - return sequence_model_parallel_all_gather(latents, dim=2) + return self._gather_sp_tensor(latents, dim=2) @dataclass diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py b/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py index 126cf1a547c1..00b5dc80cbb3 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py @@ -19,7 +19,6 @@ from sglang.multimodal_gen.runtime.distributed import ( get_sp_parallel_rank, get_sp_world_size, - sequence_model_parallel_all_gather, ) @@ -383,7 +382,7 @@ def gather_latents_for_sp(self, latents, batch=None): if get_sp_world_size() <= 1: return latents if isinstance(latents, torch.Tensor) and latents.ndim == 3: - return sequence_model_parallel_all_gather(latents.contiguous(), dim=1) + return self._gather_sp_tensor(latents, dim=1) return super().gather_latents_for_sp(latents, batch=batch) def shard_audio_latents_for_sp(self, batch, audio_latents): @@ -432,18 +431,21 @@ def can_shard_audio_latents_for_sp(self, audio_latents) -> bool: ) def gather_audio_latents_for_sp(self, audio_latents, batch): + """Gather packed audio latents after SP and trim any pad-only tail tokens.""" if get_sp_world_size() <= 1: return audio_latents if not (isinstance(audio_latents, torch.Tensor) and audio_latents.ndim == 3): return audio_latents - audio_latents = sequence_model_parallel_all_gather( - audio_latents.contiguous(), dim=1 + audio_latents = self._gather_sp_tensor( + audio_latents, + dim=1, + ) + return self._trim_sp_gather_padding( + audio_latents, + orig_len=getattr(batch, "sp_audio_orig_num_frames", None), + dim=1, ) - orig_num_frames = int(batch.sp_audio_orig_num_frames) - if orig_num_frames > 0: - audio_latents = audio_latents[:, :orig_num_frames, :] - return audio_latents def prepare_video_rope_coords_for_sp( self, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py index ba6fc5f46689..208791f2e6b5 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py @@ -56,6 +56,9 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.latent_preparation_av import ( LTX2AVLatentPreparationStage, ) +from sglang.multimodal_gen.runtime.pipelines_core.stages.ltx_2_denoising import ( + LTX2DenoisingStage, +) from sglang.multimodal_gen.runtime.pipelines_core.stages.text_connector import ( LTX2TextConnectorStage, ) @@ -80,6 +83,7 @@ "LTX2AVLatentPreparationStage", "DenoisingStage", "DmdDenoisingStage", + "LTX2DenoisingStage", "LTX2AVDenoisingStage", "CausalDMDDenoisingStage", "EncodingStage", 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 7f385c967d69..7b868be9faf6 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -10,13 +10,13 @@ import os import time import weakref -from collections.abc import Iterable +from collections.abc import Callable, Iterable +from dataclasses import dataclass, field, fields from functools import lru_cache from typing import Any import torch import torch.nn as nn -from einops import rearrange from tqdm.auto import tqdm from sglang.multimodal_gen import envs @@ -25,9 +25,6 @@ Flux2PipelineConfig, FluxPipelineConfig, ) -from sglang.multimodal_gen.configs.pipeline_configs.wan import ( - Wan2_2_TI2V_5B_Config, -) from sglang.multimodal_gen.configs.pipeline_configs.zimage import ZImagePipelineConfig from sglang.multimodal_gen.runtime.cache.cache_dit_integration import ( CacheDitConfig, @@ -41,7 +38,6 @@ cfg_model_parallel_all_reduce, get_local_torch_device, get_sp_group, - get_sp_parallel_rank, get_sp_world_size, get_tp_group, get_world_group, @@ -68,6 +64,13 @@ PipelineStage, StageParallelismType, ) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.wan_ti2v import ( + blend_wan_ti2v_latents, + expand_wan_ti2v_timestep, + prepare_wan_ti2v_latents, + prepare_wan_ti2v_sp_inputs, + should_apply_wan_ti2v, +) from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( StageValidators as V, ) @@ -89,12 +92,60 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler from sglang.multimodal_gen.runtime.utils.profiler import SGLDiffusionProfiler -from sglang.multimodal_gen.utils import dict_to_3d_list, masks_like +from sglang.multimodal_gen.utils import dict_to_3d_list from sglang.srt.utils.common import get_compiler_backend logger = init_logger(__name__) +@dataclass(slots=True) +class DenoisingContext: + """Loop-scoped state shared across the denoising skeleton and its hooks.""" + + extra_step_kwargs: dict[str, Any] + target_dtype: torch.dtype + autocast_enabled: bool + timesteps: torch.Tensor + num_inference_steps: int + num_warmup_steps: int + image_kwargs: dict[str, Any] + pos_cond_kwargs: dict[str, Any] + neg_cond_kwargs: dict[str, Any] + latents: torch.Tensor + boundary_timestep: float | None + z: torch.Tensor | None + reserved_frames_mask: torch.Tensor | None + seq_len: int | None + guidance: torch.Tensor + is_warmup: bool + trajectory_timesteps: list[torch.Tensor] = field(default_factory=list) + trajectory_latents: list[torch.Tensor] = field(default_factory=list) + extra: dict[str, Any] = field(default_factory=dict) + + def __getitem__(self, key: str) -> Any: + return getattr(self, key) + + def get(self, key: str, default: Any = None) -> Any: + return getattr(self, key, default) + + def to_kwargs(self) -> dict[str, Any]: + """Return a shallow field mapping for derived context construction.""" + return {item.name: getattr(self, item.name) for item in fields(self)} + + +@dataclass(slots=True) +class DenoisingStepState: + """Per-step hot-path state computed once and reused within a denoising step.""" + + step_index: int + t_host: torch.Tensor + t_device: torch.Tensor + t_int: int + current_model: Any + current_guidance_scale: Any + attn_metadata: Any | None + + class DenoisingStage(PipelineStage): """ Stage for running the denoising loop in diffusion pipelines. @@ -431,118 +482,6 @@ def parallelism_type(self) -> StageParallelismType: # return StageParallelismType.CFG_PARALLEL if get_global_server_args().enable_cfg_parallel else StageParallelismType.REPLICATED return StageParallelismType.REPLICATED - def _preprocess_latents_for_ti2v( - self, latents, target_dtype, batch, server_args: ServerArgs - ): - # FIXME: should probably move to latent preparation stage, to handle with offload - # Wan2.2 TI2V directly replaces the first frame of the latent with - # the image latent instead of appending along the channel dim - assert batch.image_latent is None, "TI2V task should not have image latents" - assert self.vae is not None, "VAE is not provided for TI2V task" - self.vae = self.vae.to(batch.condition_image.device) - z = self.vae.encode(batch.condition_image).mean.float() - if self.vae.device != "cpu" and server_args.vae_cpu_offload: - self.vae = self.vae.to("cpu") - if hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None: - if isinstance(self.vae.shift_factor, torch.Tensor): - z -= self.vae.shift_factor.to(z.device, z.dtype) - else: - z -= self.vae.shift_factor - - if isinstance(self.vae.scaling_factor, torch.Tensor): - z = z * self.vae.scaling_factor.to(z.device, z.dtype) - else: - z = z * self.vae.scaling_factor - # z: [B, C, 1, H, W] - latent_model_input = latents.to(target_dtype) - # Keep as [B, C, T, H, W] for proper broadcasting - assert latent_model_input.ndim == 5 - - # Create mask with proper shape [B, C, T, H, W] - latent_for_mask = latent_model_input.squeeze(0) # [C, T, H, W] - _, reserved_frames_masks = masks_like([latent_for_mask], zero=True) - reserved_frames_mask = reserved_frames_masks[0].unsqueeze(0) # [1, C, T, H, W] - - # replace GLOBAL first frame with image - proper broadcasting - # z: [B, C, 1, H, W], reserved_frames_mask: [1, C, T, H, W] - # Both will broadcast correctly - latents = ( - 1.0 - reserved_frames_mask - ) * z + reserved_frames_mask * latent_model_input - assert latents.ndim == 5 - latents = latents.to(get_local_torch_device()) - batch.latents = latents - - F = batch.num_frames - temporal_scale = ( - server_args.pipeline_config.vae_config.arch_config.scale_factor_temporal - ) - spatial_scale = ( - server_args.pipeline_config.vae_config.arch_config.scale_factor_spatial - ) - patch_size = server_args.pipeline_config.dit_config.arch_config.patch_size - seq_len = ( - ((F - 1) // temporal_scale + 1) - * (batch.height // spatial_scale) - * (batch.width // spatial_scale) - // (patch_size[1] * patch_size[2]) - ) - seq_len = int(math.ceil(seq_len / get_sp_world_size())) * get_sp_world_size() - return seq_len, z, reserved_frames_masks - - def _postprocess_latents_for_ti2v(self, z, reserved_frames_masks, batch): - rank_in_sp_group = get_sp_parallel_rank() - sp_world_size = get_sp_world_size() - - if getattr(batch, "did_sp_shard_latents", False): - # Shard z (image latent) along time dimension - # z shape: [1, C, 1, H, W] - only first frame - # Only rank 0 has the first frame after sharding - if z.shape[2] == 1: - # z is single frame, only rank 0 needs it - if rank_in_sp_group == 0: - z_sp = z - else: - # Other ranks don't have the first frame - z_sp = None - else: - # Should not happen for TI2V - z_sp = z - - # Shard reserved_frames_mask along time dimension to match sharded latents - # reserved_frames_mask is a list from masks_like, extract reserved_frames_mask[0] first - # reserved_frames_mask[0] shape: [C, T, H, W] - # All ranks need their portion of reserved_frames_mask for timestep calculation - if reserved_frames_masks is not None: - reserved_frames_mask = reserved_frames_masks[ - 0 - ] # Extract tensor from list - time_dim = reserved_frames_mask.shape[1] # [C, T, H, W] - if time_dim > 0 and time_dim % sp_world_size == 0: - reserved_frames_mask_sp_tensor = rearrange( - reserved_frames_mask, - "c (n t) h w -> c n t h w", - n=sp_world_size, - ).contiguous() - reserved_frames_mask_sp_tensor = reserved_frames_mask_sp_tensor[ - :, rank_in_sp_group, :, :, : - ] - reserved_frames_mask_sp = ( - reserved_frames_mask_sp_tensor # Store as tensor, not list - ) - else: - reserved_frames_mask_sp = reserved_frames_mask - else: - reserved_frames_mask_sp = None - else: - # SP not enabled or latents not sharded - z_sp = z - reserved_frames_mask_sp = ( - reserved_frames_masks[0] if reserved_frames_masks is not None else None - ) # Extract tensor - - return reserved_frames_mask_sp, z_sp - def _handle_boundary_ratio( self, server_args, @@ -572,7 +511,7 @@ def _prepare_denoising_loop(self, batch: Req, server_args: ServerArgs): Prepare all necessary invariant variables for the denoising loop. Returns: - A dictionary containing all the prepared variables for the denoising loop. + A context object containing the invariant state for the denoising loop. """ assert self.transformer is not None pipeline = self.pipeline() if self.pipeline else None @@ -642,17 +581,16 @@ def _prepare_denoising_loop(self, batch: Req, server_args: ServerArgs): assert neg_prompt_embeds is not None # Removed Tensor truthiness assert to avoid GPU sync - # specifically for Wan2_2_TI2V_5B_Config, not applicable for FastWan2_2_TI2V_5B_Config - should_preprocess_for_wan_ti2v = ( - server_args.pipeline_config.task_type == ModelTaskType.TI2V - and batch.condition_image is not None - and type(server_args.pipeline_config) is Wan2_2_TI2V_5B_Config - ) + should_preprocess_for_wan_ti2v = should_apply_wan_ti2v(batch, server_args) # TI2V specific preparations - before SP sharding if should_preprocess_for_wan_ti2v: - seq_len, z, reserved_frames_masks = self._preprocess_latents_for_ti2v( - latents, target_dtype, batch, server_args + seq_len, z, reserved_frames_masks = prepare_wan_ti2v_latents( + self.vae, + latents, + target_dtype, + batch, + server_args, ) else: seq_len, z, reserved_frames_masks = ( @@ -667,7 +605,7 @@ def _prepare_denoising_loop(self, batch: Req, server_args: ServerArgs): # Shard z and reserved_frames_mask for TI2V if SP is enabled if should_preprocess_for_wan_ti2v: - reserved_frames_mask_sp, z_sp = self._postprocess_latents_for_ti2v( + reserved_frames_mask_sp, z_sp = prepare_wan_ti2v_sp_inputs( z, reserved_frames_masks, batch ) else: @@ -732,26 +670,199 @@ def _prepare_denoising_loop(self, batch: Req, server_args: ServerArgs): else: neg_cond_kwargs = {} - return { - "extra_step_kwargs": extra_step_kwargs, - "target_dtype": target_dtype, - "autocast_enabled": autocast_enabled, - "timesteps": timesteps, - "num_inference_steps": num_inference_steps, - "num_warmup_steps": num_warmup_steps, - "image_kwargs": image_kwargs, - "pos_cond_kwargs": pos_cond_kwargs, - "neg_cond_kwargs": neg_cond_kwargs, - "latents": latents, - "prompt_embeds": prompt_embeds, - "neg_prompt_embeds": neg_prompt_embeds, - "boundary_timestep": boundary_timestep, - "z": z_sp, # Use SP-sharded version - # ndim == 5 - "reserved_frames_mask": reserved_frames_mask_sp, # Use SP-sharded version - "seq_len": seq_len, - "guidance": guidance, - } + return DenoisingContext( + extra_step_kwargs=extra_step_kwargs, + target_dtype=target_dtype, + autocast_enabled=autocast_enabled, + timesteps=timesteps, + num_inference_steps=num_inference_steps, + num_warmup_steps=num_warmup_steps, + image_kwargs=image_kwargs, + pos_cond_kwargs=pos_cond_kwargs, + neg_cond_kwargs=neg_cond_kwargs, + latents=latents, + boundary_timestep=boundary_timestep, + z=z_sp, + reserved_frames_mask=reserved_frames_mask_sp, + seq_len=seq_len, + guidance=guidance, + is_warmup=batch.is_warmup, + ) + + def _before_denoising_loop( + self, ctx: DenoisingContext, batch: Req, server_args: ServerArgs + ) -> None: + """Prepare scheduler state before entering the shared denoising loop.""" + self.scheduler.set_begin_index(0) + + def _prepare_step_state( + self, + ctx: DenoisingContext, + batch: Req, + server_args: ServerArgs, + step_index: int, + t_host: torch.Tensor, + timesteps_cpu: torch.Tensor, + ) -> DenoisingStepState: + """Build the per-step state shared by the loop and model-specific hooks.""" + t_int = int(t_host.item()) + t_device = ctx.timesteps[step_index] + current_model, current_guidance_scale = self._select_and_manage_model( + t_int=t_int, + boundary_timestep=ctx.boundary_timestep, + server_args=server_args, + batch=batch, + ) + attn_metadata = self._prepare_step_attn_metadata( + ctx=ctx, + batch=batch, + server_args=server_args, + step_index=step_index, + t_int=t_int, + timesteps_cpu=timesteps_cpu, + ) + return DenoisingStepState( + step_index=step_index, + t_host=t_host, + t_device=t_device, + t_int=t_int, + current_model=current_model, + current_guidance_scale=current_guidance_scale, + attn_metadata=attn_metadata, + ) + + def _prepare_step_attn_metadata( + self, + ctx: DenoisingContext, + batch: Req, + server_args: ServerArgs, + step_index: int, + t_int: int, + timesteps_cpu: torch.Tensor, + ) -> Any | None: + """Build attention metadata for the current denoising step.""" + # Keep attention metadata preparation overridable so model-specific stages + # can preserve their original semantics without duplicating step state setup. + return self._build_attn_metadata( + step_index, + batch, + server_args, + timestep_value=t_int, + timesteps=timesteps_cpu, + ) + + def _get_prompt_embeds_validator( + self, batch: Req + ) -> Callable[[Any], bool] | list[Callable[[Any], bool]]: + """Return the prompt-embedding validator used by verify_input.""" + del batch + return V.list_not_empty + + def _get_negative_prompt_embeds_validator( + self, batch: Req + ) -> Callable[[Any], bool] | list[Callable[[Any], bool]]: + """Return the negative-prompt validator used by verify_input.""" + return lambda x: not batch.do_classifier_free_guidance or V.list_not_empty(x) + + def _run_denoising_step( + self, + ctx: DenoisingContext, + step: DenoisingStepState, + batch: Req, + server_args: ServerArgs, + ) -> None: + """Run one scheduler-backed denoising step in the shared base path. + + Model-specific stages should override this instead of the whole loop whenever possible to achieve better performance + """ + # 1. Prepare latent inputs in the model's compute dtype. + latent_model_input = ctx.latents.to(ctx.target_dtype) + if batch.image_latent is not None: + assert ( + not server_args.pipeline_config.task_type == ModelTaskType.TI2V + ), "image latents should not be provided for TI2V task" + latent_model_input = torch.cat( + [latent_model_input, batch.image_latent], dim=1 + ).to(ctx.target_dtype) + + # 2. Expand the timestep to the shape expected by the current model. + timestep = self.expand_timestep_before_forward( + batch, + server_args, + step.t_device, + ctx.target_dtype, + ctx.seq_len, + ctx.reserved_frames_mask, + ) + + # 3. Apply scheduler-side input scaling before the model forward. + latent_model_input = self.scheduler.scale_model_input( + latent_model_input, step.t_device + ) + + # 4. Run the model prediction path, including CFG when enabled. + noise_pred = self._predict_noise_with_cfg( + current_model=step.current_model, + latent_model_input=latent_model_input, + timestep=timestep, + batch=batch, + timestep_index=step.step_index, + attn_metadata=step.attn_metadata, + target_dtype=ctx.target_dtype, + current_guidance_scale=step.current_guidance_scale, + image_kwargs=ctx.image_kwargs, + pos_cond_kwargs=ctx.pos_cond_kwargs, + neg_cond_kwargs=ctx.neg_cond_kwargs, + server_args=server_args, + guidance=ctx.guidance, + latents=ctx.latents, + ) + if server_args.comfyui_mode: + batch.noise_pred = noise_pred + + # 5. Advance the scheduler state with the predicted noise. + ctx.latents = self.scheduler.step( + model_output=noise_pred, + timestep=step.t_device, + sample=ctx.latents, + **ctx.extra_step_kwargs, + return_dict=False, + )[0] + + # 6. Re-apply any model-specific latent constraints after the update. + ctx.latents = self.post_forward_for_ti2v_task( + batch, + server_args, + ctx.reserved_frames_mask, + ctx.latents, + ctx.z, + ) + + def _record_trajectory( + self, + ctx: DenoisingContext, + step: DenoisingStepState, + batch: Req, + server_args: ServerArgs, + ) -> None: + """Append the current step to the returned latent trajectory, if requested.""" + if not batch.return_trajectory_latents: + return + ctx.trajectory_timesteps.append(step.t_host) + ctx.trajectory_latents.append(ctx.latents) + + def _finalize_denoising_loop( + self, ctx: DenoisingContext, batch: Req, server_args: ServerArgs + ) -> None: + """Finalize the shared loop by handing state to post-denoising processing.""" + self._post_denoising_loop( + batch=batch, + latents=ctx.latents, + trajectory_latents=ctx.trajectory_latents, + trajectory_timesteps=ctx.trajectory_timesteps, + server_args=server_args, + is_warmup=ctx.is_warmup, + ) def _post_denoising_loop( self, @@ -972,45 +1083,18 @@ def expand_timestep_before_forward( reserved_frames_mask, ): bsz = batch.raw_latent_shape[0] - should_preprocess_for_wan_ti2v = ( - server_args.pipeline_config.task_type == ModelTaskType.TI2V - and batch.condition_image is not None - and type(server_args.pipeline_config) is Wan2_2_TI2V_5B_Config - ) + should_preprocess_for_wan_ti2v = should_apply_wan_ti2v(batch, server_args) # expand timestep if should_preprocess_for_wan_ti2v: - # Explicitly cast t_device to the target float type at the beginning. - # This ensures any precision-based rounding (e.g., float32(999.0) -> bfloat16(1000.0)) - # is applied consistently *before* it's used by any rank. - t_device_rounded = t_device.to(target_dtype) - - local_seq_len = seq_len - if get_sp_world_size() > 1 and getattr( - batch, "did_sp_shard_latents", False - ): - local_seq_len = seq_len // get_sp_world_size() - - if get_sp_parallel_rank() == 0 and reserved_frames_mask is not None: - # Rank 0 has the first frame, create a special timestep tensor - # NOTE: The spatial downsampling in the next line is suspicious but kept - # to match original model's potential training configuration. - temp_ts = ( - reserved_frames_mask[0][:, ::2, ::2] * t_device_rounded - ).flatten() - - # Pad to full local sequence length - temp_ts = torch.cat( - [ - temp_ts, - temp_ts.new_ones(local_seq_len - temp_ts.size(0)) - * t_device_rounded, - ] - ) - timestep = temp_ts.unsqueeze(0).repeat(bsz, 1) - else: - # Other ranks get a uniform timestep tensor of the correct shape [B, local_seq_len] - timestep = t_device.repeat(bsz, local_seq_len) + assert seq_len is not None, "Wan TI2V requires a token sequence length." + timestep = expand_wan_ti2v_timestep( + batch, + t_device, + target_dtype, + seq_len, + reserved_frames_mask, + ) else: timestep = t_device.repeat(bsz) return timestep @@ -1018,27 +1102,10 @@ def expand_timestep_before_forward( def post_forward_for_ti2v_task( self, batch: Req, server_args: ServerArgs, reserved_frames_mask, latents, z ): - """ - For Wan2.2 ti2v task, global first frame should be replaced with encoded image after each timestep - """ - should_preprocess_for_wan_ti2v = ( - server_args.pipeline_config.task_type == ModelTaskType.TI2V - and batch.condition_image is not None - and type(server_args.pipeline_config) is Wan2_2_TI2V_5B_Config - ) + """Re-apply Wan TI2V first-frame conditioning after each denoising step.""" + should_preprocess_for_wan_ti2v = should_apply_wan_ti2v(batch, server_args) if should_preprocess_for_wan_ti2v: - # Apply TI2V mask blending with SP-aware z and reserved_frames_mask. - # This ensures the first frame is always the condition image after each step. - # This is only applied on rank 0, where z is not None. - if z is not None and reserved_frames_mask is not None: - # z: [1, C, 1, H, W] - # latents: [1, C, T_local, H, W] - # reserved_frames_mask: [C, T_local, H, W] - # Unsqueeze mask to [1, C, T_local, H, W] for broadcasting. - # z will broadcast along the time dimension. - latents = ( - 1.0 - reserved_frames_mask.unsqueeze(0) - ) * z + reserved_frames_mask.unsqueeze(0) * latents + latents = blend_wan_ti2v_latents(latents, reserved_frames_mask, z) return latents @@ -1051,159 +1118,56 @@ def forward( """ Run the denoising loop. """ - # Prepare variables for the denoising loop - - prepared_vars = self._prepare_denoising_loop(batch, server_args) - extra_step_kwargs = prepared_vars["extra_step_kwargs"] - target_dtype = prepared_vars["target_dtype"] - autocast_enabled = prepared_vars["autocast_enabled"] - timesteps = prepared_vars["timesteps"] - num_inference_steps = prepared_vars["num_inference_steps"] - num_warmup_steps = prepared_vars["num_warmup_steps"] - image_kwargs = prepared_vars["image_kwargs"] - pos_cond_kwargs = prepared_vars["pos_cond_kwargs"] - neg_cond_kwargs = prepared_vars["neg_cond_kwargs"] - latents = prepared_vars["latents"] - boundary_timestep = prepared_vars["boundary_timestep"] - z = prepared_vars["z"] - reserved_frames_mask = prepared_vars["reserved_frames_mask"] - seq_len = prepared_vars["seq_len"] - guidance = prepared_vars["guidance"] - - # Initialize lists for ODE trajectory - trajectory_timesteps: list[torch.Tensor] = [] - trajectory_latents: list[torch.Tensor] = [] - - # Run denoising loop + ctx = self._prepare_denoising_loop(batch, server_args) denoising_start_time = time.time() - + self._before_denoising_loop(ctx, batch, server_args) # to avoid device-sync caused by timestep comparison - is_warmup = batch.is_warmup - self.scheduler.set_begin_index(0) - timesteps_cpu = timesteps.cpu() + timesteps_cpu = ctx.timesteps.cpu() num_timesteps = timesteps_cpu.shape[0] with torch.autocast( device_type=current_platform.device_type, - dtype=target_dtype, - enabled=autocast_enabled, + dtype=ctx.target_dtype, + enabled=ctx.autocast_enabled, ): - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t_host in enumerate(timesteps_cpu): + with self.progress_bar(total=ctx.num_inference_steps) as progress_bar: + for step_index, t_host in enumerate(timesteps_cpu): with StageProfiler( - f"denoising_step_{i}", + f"denoising_step_{step_index}", logger=logger, metrics=batch.metrics, perf_dump_path_provided=batch.perf_dump_path is not None, record_as_step=True, ): - t_int = int(t_host.item()) - t_device = timesteps[i] - current_model, current_guidance_scale = ( - self._select_and_manage_model( - t_int=t_int, - boundary_timestep=boundary_timestep, - server_args=server_args, - batch=batch, - ) - ) - - # Expand latents for I2V - latent_model_input = latents.to(target_dtype) - if batch.image_latent is not None: - assert ( - not server_args.pipeline_config.task_type - == ModelTaskType.TI2V - ), "image latents should not be provided for TI2V task" - latent_model_input = torch.cat( - [latent_model_input, batch.image_latent], dim=1 - ).to(target_dtype) - - timestep = self.expand_timestep_before_forward( + step = self._prepare_step_state( + ctx, batch, server_args, - t_device, - target_dtype, - seq_len, - reserved_frames_mask, + step_index, + t_host, + timesteps_cpu, ) + self._run_denoising_step(ctx, step, batch, server_args) + self._record_trajectory(ctx, step, batch, server_args) - latent_model_input = self.scheduler.scale_model_input( - latent_model_input, t_device - ) - - # Predict noise residual - attn_metadata = self._build_attn_metadata( - i, - batch, - server_args, - timestep_value=t_int, - timesteps=timesteps_cpu, - ) - noise_pred = self._predict_noise_with_cfg( - current_model=current_model, - latent_model_input=latent_model_input, - timestep=timestep, - batch=batch, - timestep_index=i, - attn_metadata=attn_metadata, - target_dtype=target_dtype, - current_guidance_scale=current_guidance_scale, - image_kwargs=image_kwargs, - pos_cond_kwargs=pos_cond_kwargs, - neg_cond_kwargs=neg_cond_kwargs, - server_args=server_args, - guidance=guidance, - latents=latents, - ) - # Save noise_pred to batch for external access (e.g., ComfyUI) - if server_args.comfyui_mode: - batch.noise_pred = noise_pred - - # Compute the previous noisy sample - latents = self.scheduler.step( - model_output=noise_pred, - timestep=t_device, - sample=latents, - **extra_step_kwargs, - return_dict=False, - )[0] - - latents = self.post_forward_for_ti2v_task( - batch, server_args, reserved_frames_mask, latents, z - ) - - # save trajectory latents if needed - if batch.return_trajectory_latents: - trajectory_timesteps.append(t_host) - trajectory_latents.append(latents) - - # Update progress bar - if i == num_timesteps - 1 or ( - (i + 1) > num_warmup_steps - and (i + 1) % self.scheduler.order == 0 + if step_index == num_timesteps - 1 or ( + (step_index + 1) > ctx.num_warmup_steps + and (step_index + 1) % self.scheduler.order == 0 and progress_bar is not None ): progress_bar.update() - if not is_warmup: + if not ctx.is_warmup: self.step_profile() denoising_end_time = time.time() - if num_timesteps > 0 and not is_warmup: + if num_timesteps > 0 and not ctx.is_warmup: self.log_info( "average time per step: %.4f seconds", - (denoising_end_time - denoising_start_time) / len(timesteps), + (denoising_end_time - denoising_start_time) / len(ctx.timesteps), ) - self._post_denoising_loop( - batch=batch, - latents=latents, - trajectory_latents=trajectory_latents, - trajectory_timesteps=trajectory_timesteps, - server_args=server_args, - is_warmup=is_warmup, - ) + self._finalize_denoising_loop(ctx, batch, server_args) return batch # TODO: this will extends the preparation stage, should let subclass/passed-in variables decide which to prepare @@ -1846,7 +1810,11 @@ def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResul result.add_check("timesteps", batch.timesteps, [V.is_tensor, V.min_dims(1)]) # disable temporarily for image-generation models # result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)]) - result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty) + result.add_check( + "prompt_embeds", + batch.prompt_embeds, + self._get_prompt_embeds_validator(batch), + ) result.add_check("image_embeds", batch.image_embeds, V.is_list) # result.add_check( # "image_latent", batch.image_latent, V.none_or_tensor_with_dims(5) @@ -1865,7 +1833,7 @@ def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResul result.add_check( "negative_prompt_embeds", batch.negative_prompt_embeds, - lambda x: not batch.do_classifier_free_guidance or V.list_not_empty(x), + self._get_negative_prompt_embeds_validator(batch), ) return result diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py index a688d8f7e397..e735960bd788 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py @@ -1,54 +1,26 @@ import copy -import json -import math -import os -import time -from io import BytesIO -import av -import numpy as np -import PIL.Image import torch -from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution -from diffusers.models.modeling_outputs import AutoencoderKLOutput from diffusers.utils.torch_utils import randn_tensor -from safetensors.torch import load_file as safetensors_load_file from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import ( is_ltx23_native_variant, ) -from sglang.multimodal_gen.runtime.distributed import get_sp_world_size -from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context -from sglang.multimodal_gen.runtime.models.vaes.ltx_2_3_condition_encoder import ( - LTX23VideoConditionEncoder, -) -from sglang.multimodal_gen.runtime.models.vision_utils import ( - load_image, - normalize, - numpy_to_pt, - pil_to_numpy, -) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req -from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import DenoisingStage -from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( - StageValidators as V, -) -from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( - VerificationResult, +from sglang.multimodal_gen.runtime.pipelines_core.stages.ltx_2_denoising import ( + LTX2DenoisingStage, ) -from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger -from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler -from sglang.multimodal_gen.utils import PRECISION_TO_TYPE logger = init_logger(__name__) -class LTX2AVDenoisingStage(DenoisingStage): +class LTX2AVDenoisingStage(LTX2DenoisingStage): """ - LTX-2 specific denoising stage that handles joint video and audio generation. + Thin AV layer that adds audio trajectory gathering and final unpacking on top of + the LTX-2 denoising semantics. """ def __init__(self, transformer, scheduler, vae=None, audio_vae=None, **kwargs): @@ -56,1862 +28,6 @@ def __init__(self, transformer, scheduler, vae=None, audio_vae=None, **kwargs): transformer=transformer, scheduler=scheduler, vae=vae, **kwargs ) self.audio_vae = audio_vae - self._condition_image_encoder = None - self._condition_image_encoder_dir = None - - @staticmethod - def _get_video_latent_num_frames_for_model( - batch: Req, server_args: ServerArgs, latents: torch.Tensor - ) -> int: - """Return the latent-frame length the DiT model should see. - - - If video latents were time-sharded for SP and are packed as token latents - ([B, S, D]), the model only sees the local shard and must use the local - latent-frame count (stored on the batch during SP sharding). - - Otherwise, fall back to the global latent-frame count inferred from the - requested output frames and the VAE temporal compression ratio. - """ - did_sp_shard = bool(getattr(batch, "did_sp_shard_latents", False)) - is_token_latents = isinstance(latents, torch.Tensor) and latents.ndim == 3 - - if did_sp_shard and is_token_latents: - if not hasattr(batch, "sp_video_latent_num_frames"): - raise ValueError( - "SP-sharded LTX2 token latents require `batch.sp_video_latent_num_frames` " - "to be set by `LTX2PipelineConfig.shard_latents_for_sp()`." - ) - return int(batch.sp_video_latent_num_frames) - - pc = server_args.pipeline_config - return int( - (batch.num_frames - 1) - // int(pc.vae_config.arch_config.temporal_compression_ratio) - + 1 - ) - - @staticmethod - def _truncate_sp_padded_token_latents( - batch: Req, latents: torch.Tensor - ) -> torch.Tensor: - """Remove token padding introduced by SP time-sharding (if applicable).""" - did_sp_shard = bool(getattr(batch, "did_sp_shard_latents", False)) - if not did_sp_shard or not ( - isinstance(latents, torch.Tensor) and latents.ndim == 3 - ): - return latents - - raw_shape = getattr(batch, "raw_latent_shape", None) - if not (isinstance(raw_shape, tuple) and len(raw_shape) == 3): - return latents - - orig_s = int(raw_shape[1]) - cur_s = int(latents.shape[1]) - if cur_s == orig_s: - return latents - if cur_s < orig_s: - raise ValueError( - f"Unexpected gathered token-latents seq_len {cur_s} < original seq_len {orig_s}." - ) - return latents[:, :orig_s, :].contiguous() - - def _maybe_enable_cache_dit(self, num_inference_steps: int, batch: Req) -> None: - """Disable cache-dit for TI2V-style requests (image-conditioned), to avoid stale activations. - - NOTE: base denoising stage calls this hook with (num_inference_steps, batch). - """ - if getattr(self, "_disable_cache_dit_for_request", False): - return - return super()._maybe_enable_cache_dit(num_inference_steps, batch) - - def _get_ltx2_stage1_guider_params( - self, batch: Req, server_args: ServerArgs, stage: str - ) -> dict[str, object] | None: - if stage != "stage1": - return None - return batch.extra.get("ltx2_stage1_guider_params") - - @staticmethod - def _ltx2_should_skip_step(step_index: int, skip_step: int) -> bool: - if skip_step == 0: - return False - return step_index % (skip_step + 1) != 0 - - @staticmethod - def _ltx2_apply_rescale( - cond: torch.Tensor, pred: torch.Tensor, rescale_scale: float - ) -> torch.Tensor: - if rescale_scale == 0.0: - return pred - factor = cond.std() / pred.std() - factor = rescale_scale * factor + (1.0 - rescale_scale) - return pred * factor - - @staticmethod - def _prepare_ltx2_ti2v_clean_state( - latents: torch.Tensor, - image_latent: torch.Tensor, - num_img_tokens: int, - zero_clean_latent: bool, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - latents = latents.clone() - conditioned = image_latent[:, :num_img_tokens, :].to( - device=latents.device, dtype=latents.dtype - ) - latents[:, :num_img_tokens, :] = conditioned - denoise_mask = torch.ones( - (latents.shape[0], latents.shape[1], 1), - device=latents.device, - dtype=torch.float32, - ) - denoise_mask[:, :num_img_tokens, :] = 0.0 - if zero_clean_latent: - clean_latent = torch.zeros_like(latents) - else: - clean_latent = latents.detach().clone() - clean_latent[:, :num_img_tokens, :] = conditioned - return latents, denoise_mask, clean_latent - - @staticmethod - def _ltx2_velocity_to_x0( - sample: torch.Tensor, - velocity: torch.Tensor, - sigma: float | torch.Tensor, - ) -> torch.Tensor: - if isinstance(sigma, torch.Tensor): - sigma = sigma.to(device=sample.device, dtype=torch.float32) - while sigma.ndim < sample.ndim: - sigma = sigma.unsqueeze(-1) - return (sample.float() - sigma * velocity.float()).to(sample.dtype) - return (sample.float() - float(sigma) * velocity.float()).to(sample.dtype) - - @staticmethod - def _repeat_batch_dim(tensor: torch.Tensor, target_batch_size: int) -> torch.Tensor: - """Repeat along batch dim while preserving any tokenwise timestep layout.""" - if tensor.shape[0] == int(target_batch_size): - return tensor - if tensor.shape[0] <= 0 or int(target_batch_size) % int(tensor.shape[0]) != 0: - raise ValueError( - f"Cannot repeat tensor with batch={tensor.shape[0]} to target_batch_size={target_batch_size}" - ) - repeat_factor = int(target_batch_size) // int(tensor.shape[0]) - return tensor.repeat(repeat_factor, *([1] * (tensor.ndim - 1))) - - @staticmethod - def _build_ltx2_sp_padding_mask( - batch: Req, - *, - seq_len: int, - batch_size: int, - key: str, - device: torch.device, - ) -> torch.Tensor | None: - valid = getattr(batch, key, None) - if valid is None: - return None - valid = int(valid) - if valid <= 0 or valid >= int(seq_len): - return None - mask = torch.ones( - (batch_size, int(seq_len)), device=device, dtype=torch.float32 - ) - mask[:, valid:] = 0.0 - return mask - - @staticmethod - def _get_ltx_prompt_attention_mask( - batch: Req, - *, - is_ltx23_variant: bool, - negative: bool = False, - ) -> torch.Tensor | None: - if is_ltx23_variant: - return None - return ( - batch.negative_attention_mask if negative else batch.prompt_attention_mask - ) - - @classmethod - def _should_use_ltx23_legacy_one_stage( - cls, - server_args: ServerArgs, - pipeline_name: str | None, - ) -> bool: - if not is_ltx23_native_variant( - server_args.pipeline_config.vae_config.arch_config - ): - return False - if server_args.pipeline_class_name == "LTX2TwoStagePipeline": - return False - return pipeline_name != "LTX2TwoStagePipeline" - - @classmethod - def _should_shard_ltx23_legacy_one_stage_audio_latents( - cls, - batch: Req, - server_args: ServerArgs, - ) -> bool: - return bool( - get_sp_world_size() > 1 - and is_ltx23_native_variant( - server_args.pipeline_config.vae_config.arch_config - ) - and cls._should_use_ltx23_legacy_one_stage(server_args, None) - and server_args.pipeline_config.can_shard_audio_latents_for_sp( - batch.audio_latents - ) - ) - - @classmethod - def _ltx2_calculate_guided_x0( - cls, - *, - cond: torch.Tensor, - uncond_text: torch.Tensor | float, - uncond_perturbed: torch.Tensor | float, - uncond_modality: torch.Tensor | float, - cfg_scale: float, - stg_scale: float, - rescale_scale: float, - modality_scale: float, - ) -> torch.Tensor: - pred = ( - cond - + (cfg_scale - 1.0) * (cond - uncond_text) - + stg_scale * (cond - uncond_perturbed) - + (modality_scale - 1.0) * (cond - uncond_modality) - ) - return cls._ltx2_apply_rescale(cond, pred, rescale_scale) - - @staticmethod - def _resize_center_crop( - img: PIL.Image.Image, *, width: int, height: int - ) -> PIL.Image.Image: - return img.resize((width, height), resample=PIL.Image.Resampling.BILINEAR) - - @staticmethod - def _apply_video_codec_compression( - img_array: np.ndarray, crf: int = 33 - ) -> np.ndarray: - """Encode as a single H.264 frame and decode back to simulate compression artifacts.""" - if crf == 0: - return img_array - height, width = img_array.shape[0] // 2 * 2, img_array.shape[1] // 2 * 2 - img_array = img_array[:height, :width] - buffer = BytesIO() - container = av.open(buffer, mode="w", format="mp4") - stream = container.add_stream( - "libx264", rate=1, options={"crf": str(crf), "preset": "veryfast"} - ) - stream.height, stream.width = height, width - frame = av.VideoFrame.from_ndarray(img_array, format="rgb24").reformat( - format="yuv420p" - ) - container.mux(stream.encode(frame)) - container.mux(stream.encode()) - container.close() - buffer.seek(0) - container = av.open(buffer) - decoded = next(container.decode(container.streams.video[0])) - container.close() - return decoded.to_ndarray(format="rgb24") - - @staticmethod - def _resize_center_crop_tensor( - img: PIL.Image.Image, - *, - width: int, - height: int, - device: torch.device, - dtype: torch.dtype, - apply_codec_compression: bool = True, - codec_crf: int = 33, - ) -> torch.Tensor: - """Resize, center-crop, and normalize to [1, C, 1, H, W] tensor in [-1, 1].""" - img_array = np.array(img).astype(np.uint8)[..., :3] - if apply_codec_compression: - img_array = LTX2AVDenoisingStage._apply_video_codec_compression( - img_array, crf=codec_crf - ) - tensor = ( - torch.from_numpy(img_array.astype(np.float32)) - .permute(2, 0, 1) - .unsqueeze(0) - .to(device=device) - ) - src_h, src_w = tensor.shape[2], tensor.shape[3] - scale = max(height / src_h, width / src_w) - new_h, new_w = math.ceil(src_h * scale), math.ceil(src_w * scale) - tensor = torch.nn.functional.interpolate( - tensor, size=(new_h, new_w), mode="bilinear", align_corners=False - ) - top, left = (new_h - height) // 2, (new_w - width) // 2 - tensor = tensor[:, :, top : top + height, left : left + width] - return ((tensor / 127.5 - 1.0).to(dtype=dtype)).unsqueeze(2) - - @staticmethod - def _pil_to_normed_tensor(img: PIL.Image.Image) -> torch.Tensor: - # PIL -> numpy [0,1] -> torch [B,C,H,W], then [-1,1] - arr = pil_to_numpy(img) - t = numpy_to_pt(arr) - return normalize(t) - - @staticmethod - def _should_apply_ltx2_ti2v(batch: Req) -> bool: - """True if we have an image-latent token prefix to condition with. - - SP note: when token latents are time-sharded, only the rank that owns the - *global* first latent frame should apply TI2V conditioning (rank with start_frame==0). - """ - if ( - batch.image_latent is None - or int(getattr(batch, "ltx2_num_image_tokens", 0)) <= 0 - ): - return False - did_sp_shard = bool(getattr(batch, "did_sp_shard_latents", False)) - if not did_sp_shard: - return True - return int(getattr(batch, "sp_video_start_frame", 0)) == 0 - - @staticmethod - def _should_replicate_ltx23_audio_for_sp( - batch: Req, - server_args: ServerArgs, - *, - is_ltx23_variant: bool, - ) -> bool: - return False - - def _get_condition_image_encoder( - self, - server_args: ServerArgs, - *, - device: torch.device, - dtype: torch.dtype, - ) -> LTX23VideoConditionEncoder | None: - arch_config = server_args.pipeline_config.vae_config.arch_config - encoder_subdir = str(getattr(arch_config, "condition_encoder_subdir", "")) - if not encoder_subdir: - return None - - vae_model_path = server_args.model_paths["vae"] - encoder_dir = os.path.join(vae_model_path, encoder_subdir) - config_path = os.path.join(encoder_dir, "config.json") - weights_path = os.path.join(encoder_dir, "model.safetensors") - if not os.path.exists(config_path) or not os.path.exists(weights_path): - raise ValueError( - f"LTX-2 condition encoder files not found under {encoder_dir}" - ) - - cached_dir = self._condition_image_encoder_dir - encoder = self._condition_image_encoder - if encoder is None or cached_dir != encoder_dir: - with open(config_path, encoding="utf-8") as f: - config = json.load(f) - encoder = LTX23VideoConditionEncoder(config) - encoder.load_state_dict(safetensors_load_file(weights_path), strict=True) - self._condition_image_encoder = encoder - self._condition_image_encoder_dir = encoder_dir - - encoder = encoder.to(device=device, dtype=dtype) - return encoder - - def _prepare_ltx2_image_latent(self, batch: Req, server_args: ServerArgs) -> None: - """Encode `batch.image_path` into packed token latents for LTX-2 TI2V.""" - if ( - batch.image_latent is not None - and int(getattr(batch, "ltx2_num_image_tokens", 0)) > 0 - ): - return - batch.ltx2_num_image_tokens = 0 - batch.image_latent = None - - if batch.image_path is None: - return - if batch.width is None or batch.height is None: - raise ValueError("width/height must be provided for LTX-2 TI2V.") - if self.vae is None: - raise ValueError("VAE must be provided for LTX-2 TI2V.") - - image_path = ( - batch.image_path[0] - if isinstance(batch.image_path, list) - else batch.image_path - ) - - img = load_image(image_path) - img_array = np.array(img).astype(np.uint8)[..., :3] - img_array = self._apply_video_codec_compression(img_array, crf=33) - conditioned_img = PIL.Image.fromarray(img_array) - batch.condition_image = self._resize_center_crop( - conditioned_img, width=int(batch.width), height=int(batch.height) - ) - - latents_device = ( - batch.latents.device - if isinstance(batch.latents, torch.Tensor) - else torch.device("cpu") - ) - encode_dtype = batch.latents.dtype - original_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] - vae_autocast_enabled = ( - original_dtype != torch.float32 - ) and not server_args.disable_autocast - condition_image_encoder = self._get_condition_image_encoder( - server_args, device=latents_device, dtype=encode_dtype - ) - if condition_image_encoder is None: - self.vae = self.vae.to(device=latents_device, dtype=encode_dtype) - - video_condition = self._resize_center_crop_tensor( - conditioned_img, - width=int(batch.width), - height=int(batch.height), - device=latents_device, - dtype=encode_dtype, - apply_codec_compression=False, - ) - - with torch.autocast( - device_type=current_platform.device_type, - dtype=original_dtype, - enabled=vae_autocast_enabled, - ): - try: - if ( - condition_image_encoder is None - and server_args.pipeline_config.vae_tiling - ): - self.vae.enable_tiling() - except Exception: - pass - if not vae_autocast_enabled: - video_condition = video_condition.to(encode_dtype) - - if condition_image_encoder is not None: - latent = condition_image_encoder(video_condition) - else: - latent_dist: DiagonalGaussianDistribution = self.vae.encode( - video_condition - ) - if isinstance(latent_dist, AutoencoderKLOutput): - latent_dist = latent_dist.latent_dist - - if condition_image_encoder is None: - mode = server_args.pipeline_config.vae_config.encode_sample_mode() - if mode == "argmax": - latent = latent_dist.mode() - elif mode == "sample": - if batch.generator is None: - raise ValueError("Generator must be provided for VAE sampling.") - latent = latent_dist.sample(batch.generator) - else: - raise ValueError(f"Unsupported encode_sample_mode: {mode}") - - # Per-channel normalization: normalized = (x - mean) / std - mean = self.vae.latents_mean.view(1, -1, 1, 1, 1).to(latent) - std = self.vae.latents_std.view(1, -1, 1, 1, 1).to(latent) - latent = (latent - mean) / std - else: - latent = latent.to(dtype=encode_dtype) - - packed = server_args.pipeline_config.maybe_pack_latents( - latent, latent.shape[0], batch - ) - if not (isinstance(packed, torch.Tensor) and packed.ndim == 3): - raise ValueError("Expected packed image latents [B, S0, D].") - - # Fail-fast token count: must match one latent frame's tokens. - vae_sf = int(server_args.pipeline_config.vae_scale_factor) - patch = int(server_args.pipeline_config.patch_size) - latent_h = int(batch.height) // vae_sf - latent_w = int(batch.width) // vae_sf - expected_tokens = (latent_h // patch) * (latent_w // patch) - if int(packed.shape[1]) != int(expected_tokens): - raise ValueError( - "LTX-2 conditioning token count mismatch: " - f"{int(packed.shape[1])=} {int(expected_tokens)=}." - ) - - batch.image_latent = packed - batch.ltx2_num_image_tokens = int(packed.shape[1]) - - if batch.debug: - logger.info( - "LTX2 TI2V conditioning prepared: %d tokens (shape=%s) for %sx%s", - batch.ltx2_num_image_tokens, - tuple(batch.image_latent.shape), - batch.width, - batch.height, - ) - - if condition_image_encoder is None: - self.vae.to(original_dtype) - if server_args.vae_cpu_offload: - self.vae = self.vae.to("cpu") - if condition_image_encoder is not None: - self._condition_image_encoder = condition_image_encoder.to("cpu") - - @torch.no_grad() - def _forward_ltx23_legacy_one_stage( - self, - batch: Req, - server_args: ServerArgs, - prepared_vars: dict[str, object], - ) -> Req: - target_dtype = prepared_vars["target_dtype"] - autocast_enabled = prepared_vars["autocast_enabled"] - timesteps = prepared_vars["timesteps"] - num_inference_steps = prepared_vars["num_inference_steps"] - num_warmup_steps = prepared_vars["num_warmup_steps"] - latents = prepared_vars["latents"] - boundary_timestep = prepared_vars["boundary_timestep"] - z = prepared_vars["z"] - reserved_frames_mask = prepared_vars["reserved_frames_mask"] - stage = "stage1" - audio_latents = batch.audio_latents - audio_scheduler = copy.deepcopy(self.scheduler) - batch.ltx23_audio_replicated_for_sp = False - batch.did_sp_shard_audio_latents = False - - latent_num_frames_for_model = self._get_video_latent_num_frames_for_model( - batch=batch, server_args=server_args, latents=latents - ) - latent_height = ( - batch.height - // server_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio - ) - latent_width = ( - batch.width - // server_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio - ) - - trajectory_timesteps: list[torch.Tensor] = [] - trajectory_latents: list[torch.Tensor] = [] - trajectory_audio_latents: list[torch.Tensor] = [] - - denoising_start_time = time.time() - - is_warmup = batch.is_warmup - self.scheduler.set_begin_index(0) - audio_scheduler.set_begin_index(0) - timesteps_cpu = timesteps.cpu() - num_timesteps = timesteps_cpu.shape[0] - - do_ti2v = self._should_apply_ltx2_ti2v(batch) - num_img_tokens = int(getattr(batch, "ltx2_num_image_tokens", 0)) - denoise_mask = None - clean_latent = None - if do_ti2v: - if not (isinstance(latents, torch.Tensor) and latents.ndim == 3): - raise ValueError("LTX-2 TI2V expects packed token latents [B, S, D].") - latents, denoise_mask, clean_latent = self._prepare_ltx2_ti2v_clean_state( - latents=latents, - image_latent=batch.image_latent, - num_img_tokens=num_img_tokens, - zero_clean_latent=True, - ) - - with torch.autocast( - device_type=current_platform.device_type, - dtype=target_dtype, - enabled=autocast_enabled, - ): - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t_host in enumerate(timesteps_cpu): - with StageProfiler( - f"denoising_step_{i}", - logger=logger, - metrics=batch.metrics, - perf_dump_path_provided=batch.perf_dump_path is not None, - record_as_step=True, - ): - t_int = int(t_host.item()) - t_device = timesteps[i] - ( - current_model, - current_guidance_scale, - ) = self._select_and_manage_model( - t_int=t_int, - boundary_timestep=boundary_timestep, - server_args=server_args, - batch=batch, - ) - - attn_metadata = self._build_attn_metadata(i, batch, server_args) - - sigmas = getattr(self.scheduler, "sigmas", None) - if sigmas is None or not isinstance(sigmas, torch.Tensor): - raise ValueError( - "Expected scheduler.sigmas to be a tensor for LTX-2." - ) - sigma = sigmas[i].to(device=latents.device, dtype=torch.float32) - sigma_next = sigmas[i + 1].to( - device=latents.device, dtype=torch.float32 - ) - dt = sigma_next - sigma - - latent_model_input = latents.to(target_dtype) - audio_latent_model_input = audio_latents.to(target_dtype) - stage1_guider_params = self._get_ltx2_stage1_guider_params( - batch, server_args, stage - ) - latent_num_frames = latent_num_frames_for_model - - if audio_latent_model_input.ndim == 3: - audio_num_frames_latent = int( - audio_latent_model_input.shape[1] - ) - elif audio_latent_model_input.ndim == 4: - audio_num_frames_latent = int( - audio_latent_model_input.shape[2] - ) - else: - raise ValueError( - f"Unexpected audio latents rank: {audio_latent_model_input.ndim}, shape={tuple(audio_latent_model_input.shape)}" - ) - - video_coords = None - audio_coords = None - - timestep = t_device.expand(int(latent_model_input.shape[0])) - if do_ti2v and denoise_mask is not None: - timestep_video = timestep.unsqueeze( - -1 - ) * denoise_mask.squeeze(-1) - else: - timestep_video = timestep - timestep_audio = timestep - - use_official_cfg_path = stage1_guider_params is None - if use_official_cfg_path: - encoder_hidden_states = batch.prompt_embeds[0] - audio_encoder_hidden_states = batch.audio_prompt_embeds[0] - encoder_attention_mask = batch.prompt_attention_mask - if batch.do_classifier_free_guidance: - latent_model_input = torch.cat( - [latent_model_input] * 2, dim=0 - ) - audio_latent_model_input = torch.cat( - [audio_latent_model_input] * 2, dim=0 - ) - encoder_hidden_states = torch.cat( - [ - batch.negative_prompt_embeds[0], - encoder_hidden_states, - ], - dim=0, - ) - audio_encoder_hidden_states = torch.cat( - [ - batch.negative_audio_prompt_embeds[0], - audio_encoder_hidden_states, - ], - dim=0, - ) - encoder_attention_mask = torch.cat( - [ - batch.negative_attention_mask, - encoder_attention_mask, - ], - dim=0, - ) - cfg_batch_size = int(latent_model_input.shape[0]) - timestep_video = self._repeat_batch_dim( - timestep_video, cfg_batch_size - ) - timestep_audio = self._repeat_batch_dim( - timestep_audio, cfg_batch_size - ) - - with set_forward_context( - current_timestep=i, attn_metadata=attn_metadata - ): - model_video, model_audio = current_model( - hidden_states=latent_model_input, - audio_hidden_states=audio_latent_model_input, - encoder_hidden_states=encoder_hidden_states, - audio_encoder_hidden_states=audio_encoder_hidden_states, - timestep=timestep_video, - audio_timestep=timestep_audio, - encoder_attention_mask=encoder_attention_mask, - audio_encoder_attention_mask=encoder_attention_mask, - num_frames=latent_num_frames, - height=latent_height, - width=latent_width, - fps=batch.fps, - audio_num_frames=audio_num_frames_latent, - video_coords=video_coords, - audio_coords=audio_coords, - return_latents=False, - return_dict=False, - ) - - model_video = model_video.float() - model_audio = model_audio.float() - if batch.do_classifier_free_guidance: - ( - model_video_uncond, - model_video_text, - ) = model_video.chunk(2) - ( - model_audio_uncond, - model_audio_text, - ) = model_audio.chunk(2) - model_video = model_video_uncond + ( - batch.guidance_scale - * (model_video_text - model_video_uncond) - ) - model_audio = model_audio_uncond + ( - batch.guidance_scale - * (model_audio_text - model_audio_uncond) - ) - v_pos = model_video - a_v_pos = model_audio - - latents = self.scheduler.step( - v_pos, t_device, latents, return_dict=False - )[0] - audio_latents = audio_scheduler.step( - a_v_pos, t_device, audio_latents, return_dict=False - )[0] - latents = self.post_forward_for_ti2v_task( - batch, server_args, reserved_frames_mask, latents, z - ) - - if batch.return_trajectory_latents: - trajectory_timesteps.append(t_host) - trajectory_latents.append(latents) - if audio_latents is not None: - trajectory_audio_latents.append(audio_latents) - - if i == num_timesteps - 1 or ( - (i + 1) > num_warmup_steps - and (i + 1) % self.scheduler.order == 0 - and progress_bar is not None - ): - progress_bar.update() - - if not is_warmup: - self.step_profile() - continue - - encoder_hidden_states = batch.prompt_embeds[0] - audio_encoder_hidden_states = batch.audio_prompt_embeds[0] - encoder_attention_mask = batch.prompt_attention_mask - with set_forward_context( - current_timestep=i, attn_metadata=attn_metadata - ): - v_pos, a_v_pos = current_model( - hidden_states=latent_model_input, - audio_hidden_states=audio_latent_model_input, - encoder_hidden_states=encoder_hidden_states, - audio_encoder_hidden_states=audio_encoder_hidden_states, - timestep=timestep_video, - audio_timestep=timestep_audio, - encoder_attention_mask=encoder_attention_mask, - audio_encoder_attention_mask=encoder_attention_mask, - num_frames=latent_num_frames, - height=latent_height, - width=latent_width, - fps=batch.fps, - audio_num_frames=audio_num_frames_latent, - video_coords=video_coords, - audio_coords=audio_coords, - return_latents=False, - return_dict=False, - ) - - if ( - stage1_guider_params is not None - or batch.do_classifier_free_guidance - ): - neg_encoder_hidden_states = ( - batch.negative_prompt_embeds[0] - ) - neg_audio_encoder_hidden_states = ( - batch.negative_audio_prompt_embeds[0] - ) - neg_encoder_attention_mask = ( - batch.negative_attention_mask - ) - - v_neg, a_v_neg = current_model( - hidden_states=latent_model_input, - audio_hidden_states=audio_latent_model_input, - encoder_hidden_states=neg_encoder_hidden_states, - audio_encoder_hidden_states=neg_audio_encoder_hidden_states, - timestep=timestep_video, - audio_timestep=timestep_audio, - encoder_attention_mask=neg_encoder_attention_mask, - audio_encoder_attention_mask=neg_encoder_attention_mask, - num_frames=latent_num_frames, - height=latent_height, - width=latent_width, - fps=batch.fps, - audio_num_frames=audio_num_frames_latent, - video_coords=video_coords, - audio_coords=audio_coords, - return_latents=False, - return_dict=False, - ) - else: - v_neg = None - a_v_neg = None - - v_pos = v_pos.float() - a_v_pos = a_v_pos.float() - if v_neg is not None: - v_neg = v_neg.float() - if a_v_neg is not None: - a_v_neg = a_v_neg.float() - - sigma_val = float(sigma.item()) - video_sigma_for_x0: float | torch.Tensor = sigma_val - if do_ti2v and denoise_mask is not None: - video_sigma_for_x0 = sigma.to( - device=latents.device, dtype=torch.float32 - ) * denoise_mask.squeeze(-1) - denoised_video = self._ltx2_velocity_to_x0( - latents, v_pos, video_sigma_for_x0 - ) - denoised_audio = self._ltx2_velocity_to_x0( - audio_latents, a_v_pos, sigma_val - ) - denoised_video_neg = None - denoised_audio_neg = None - denoised_video_perturbed = None - denoised_audio_perturbed = None - denoised_video_modality = None - denoised_audio_modality = None - - if ( - ( - stage1_guider_params is not None - or batch.do_classifier_free_guidance - ) - and v_neg is not None - and a_v_neg is not None - ): - denoised_video_neg = self._ltx2_velocity_to_x0( - latents, v_neg, video_sigma_for_x0 - ) - denoised_audio_neg = self._ltx2_velocity_to_x0( - audio_latents, a_v_neg, sigma_val - ) - if stage1_guider_params is not None: - video_skip = self._ltx2_should_skip_step( - i, int(stage1_guider_params["video_skip_step"]) - ) - audio_skip = self._ltx2_should_skip_step( - i, int(stage1_guider_params["audio_skip_step"]) - ) - - need_perturbed = ( - float(stage1_guider_params["video_stg_scale"]) != 0.0 - or float(stage1_guider_params["audio_stg_scale"]) != 0.0 - ) - if need_perturbed: - with set_forward_context( - current_timestep=i, attn_metadata=attn_metadata - ): - v_ptb, a_v_ptb = current_model( - hidden_states=latent_model_input, - audio_hidden_states=audio_latent_model_input, - encoder_hidden_states=encoder_hidden_states, - audio_encoder_hidden_states=audio_encoder_hidden_states, - timestep=timestep_video, - audio_timestep=timestep_audio, - encoder_attention_mask=encoder_attention_mask, - audio_encoder_attention_mask=encoder_attention_mask, - num_frames=latent_num_frames, - height=latent_height, - width=latent_width, - fps=batch.fps, - audio_num_frames=audio_num_frames_latent, - video_coords=video_coords, - audio_coords=audio_coords, - return_latents=False, - return_dict=False, - skip_video_self_attn_blocks=tuple( - stage1_guider_params["video_stg_blocks"] - ), - skip_audio_self_attn_blocks=tuple( - stage1_guider_params["audio_stg_blocks"] - ), - ) - denoised_video_perturbed = self._ltx2_velocity_to_x0( - latents, v_ptb.float(), video_sigma_for_x0 - ) - denoised_audio_perturbed = self._ltx2_velocity_to_x0( - audio_latents, a_v_ptb.float(), sigma_val - ) - - need_modality = ( - float(stage1_guider_params["video_modality_scale"]) - != 1.0 - or float(stage1_guider_params["audio_modality_scale"]) - != 1.0 - ) - if need_modality: - with set_forward_context( - current_timestep=i, attn_metadata=attn_metadata - ): - v_mod, a_v_mod = current_model( - hidden_states=latent_model_input, - audio_hidden_states=audio_latent_model_input, - encoder_hidden_states=encoder_hidden_states, - audio_encoder_hidden_states=audio_encoder_hidden_states, - timestep=timestep_video, - audio_timestep=timestep_audio, - encoder_attention_mask=encoder_attention_mask, - audio_encoder_attention_mask=encoder_attention_mask, - num_frames=latent_num_frames, - height=latent_height, - width=latent_width, - fps=batch.fps, - audio_num_frames=audio_num_frames_latent, - video_coords=video_coords, - audio_coords=audio_coords, - return_latents=False, - return_dict=False, - disable_a2v_cross_attn=True, - disable_v2a_cross_attn=True, - ) - denoised_video_modality = self._ltx2_velocity_to_x0( - latents, v_mod.float(), video_sigma_for_x0 - ) - denoised_audio_modality = self._ltx2_velocity_to_x0( - audio_latents, a_v_mod.float(), sigma_val - ) - - if not video_skip: - denoised_video = self._ltx2_calculate_guided_x0( - cond=denoised_video, - uncond_text=( - denoised_video_neg - if denoised_video_neg is not None - else denoised_video - ), - uncond_perturbed=( - denoised_video_perturbed - if denoised_video_perturbed is not None - else 0.0 - ), - uncond_modality=( - denoised_video_modality - if denoised_video_modality is not None - else 0.0 - ), - cfg_scale=float( - stage1_guider_params["video_cfg_scale"] - ), - stg_scale=float( - stage1_guider_params["video_stg_scale"] - ), - rescale_scale=float( - stage1_guider_params["video_rescale_scale"] - ), - modality_scale=float( - stage1_guider_params["video_modality_scale"] - ), - ) - if not audio_skip: - denoised_audio = self._ltx2_calculate_guided_x0( - cond=denoised_audio, - uncond_text=( - denoised_audio_neg - if denoised_audio_neg is not None - else denoised_audio - ), - uncond_perturbed=( - denoised_audio_perturbed - if denoised_audio_perturbed is not None - else 0.0 - ), - uncond_modality=( - denoised_audio_modality - if denoised_audio_modality is not None - else 0.0 - ), - cfg_scale=float( - stage1_guider_params["audio_cfg_scale"] - ), - stg_scale=float( - stage1_guider_params["audio_stg_scale"] - ), - rescale_scale=float( - stage1_guider_params["audio_rescale_scale"] - ), - modality_scale=float( - stage1_guider_params["audio_modality_scale"] - ), - ) - elif ( - batch.do_classifier_free_guidance - and denoised_video_neg is not None - and denoised_audio_neg is not None - ): - denoised_video = denoised_video + ( - batch.guidance_scale - 1.0 - ) * (denoised_video - denoised_video_neg) - denoised_audio = denoised_audio + ( - batch.guidance_scale - 1.0 - ) * (denoised_audio - denoised_audio_neg) - - if ( - do_ti2v - and denoise_mask is not None - and clean_latent is not None - ): - denoised_video = ( - denoised_video * denoise_mask - + clean_latent.float() * (1.0 - denoise_mask) - ) - if sigma_val == 0.0: - v_video = torch.zeros_like(denoised_video) - v_audio = torch.zeros_like(denoised_audio) - else: - v_video = ( - (latents.float() - denoised_video.float()) / sigma_val - ).to(latents.dtype) - v_audio = ( - (audio_latents.float() - denoised_audio.float()) - / sigma_val - ).to(audio_latents.dtype) - - latents = (latents.float() + v_video.float() * dt).to( - dtype=latents.dtype - ) - audio_latents = ( - audio_latents.float() + v_audio.float() * dt - ).to(dtype=audio_latents.dtype) - - latents = self.post_forward_for_ti2v_task( - batch, server_args, reserved_frames_mask, latents, z - ) - - if batch.return_trajectory_latents: - trajectory_timesteps.append(t_host) - trajectory_latents.append(latents) - if audio_latents is not None: - trajectory_audio_latents.append(audio_latents) - - if i == num_timesteps - 1 or ( - (i + 1) > num_warmup_steps - and (i + 1) % self.scheduler.order == 0 - and progress_bar is not None - ): - progress_bar.update() - - if not is_warmup: - self.step_profile() - - denoising_end_time = time.time() - - if num_timesteps > 0 and not is_warmup: - self.log_info( - "average time per step: %.4f seconds", - (denoising_end_time - denoising_start_time) / len(timesteps), - ) - - batch.audio_latents = audio_latents - self._post_denoising_loop( - batch=batch, - latents=latents, - trajectory_latents=trajectory_latents, - trajectory_timesteps=trajectory_timesteps, - trajectory_audio_latents=trajectory_audio_latents, - server_args=server_args, - is_warmup=is_warmup, - ) - - return batch - - @torch.no_grad() - def forward(self, batch: Req, server_args: ServerArgs) -> Req: - """ - Run the denoising loop. - - Args: - batch: The current batch information. - server_args: The inference arguments. - - Returns: - The batch with denoised latents. - """ - # Disable cache-dit for image-conditioned requests (TI2V-style) for correctness/debuggability. - self._disable_cache_dit_for_request = batch.image_path is not None - - # Prepare variables for the denoising loop - prepared_vars = self._prepare_denoising_loop(batch, server_args) - target_dtype = prepared_vars["target_dtype"] - autocast_enabled = prepared_vars["autocast_enabled"] - timesteps = prepared_vars["timesteps"] - num_inference_steps = prepared_vars["num_inference_steps"] - num_warmup_steps = prepared_vars["num_warmup_steps"] - latents = prepared_vars["latents"] - boundary_timestep = prepared_vars["boundary_timestep"] - z = prepared_vars["z"] - reserved_frames_mask = prepared_vars["reserved_frames_mask"] - is_ltx23_variant = is_ltx23_native_variant( - server_args.pipeline_config.vae_config.arch_config - ) - phase = batch.extra.get("ltx2_phase") - pipeline = self.pipeline() if self.pipeline else None - pipeline_name = pipeline.pipeline_name if pipeline is not None else None - use_ltx23_legacy_one_stage = self._should_use_ltx23_legacy_one_stage( - server_args, pipeline_name - ) - stage = ( - phase - if phase is not None - else ("stage1" if use_ltx23_legacy_one_stage else "one_stage") - ) - audio_latents = batch.audio_latents - audio_scheduler = copy.deepcopy(self.scheduler) - - self._prepare_ltx2_image_latent(batch, server_args) - if use_ltx23_legacy_one_stage: - return self._forward_ltx23_legacy_one_stage( - batch, server_args, prepared_vars - ) - do_ti2v = self._should_apply_ltx2_ti2v(batch) - replicate_audio_for_sp = self._should_replicate_ltx23_audio_for_sp( - batch, - server_args, - is_ltx23_variant=is_ltx23_variant, - ) - batch.ltx23_audio_replicated_for_sp = bool(replicate_audio_for_sp) - - if ( - is_ltx23_variant - and get_sp_world_size() > 1 - and server_args.pipeline_config.can_shard_audio_latents_for_sp( - batch.audio_latents - ) - and not replicate_audio_for_sp - and not use_ltx23_legacy_one_stage - ): - ( - batch.audio_latents, - batch.did_sp_shard_audio_latents, - ) = server_args.pipeline_config.shard_audio_latents_for_sp( - batch, batch.audio_latents - ) - audio_latents = batch.audio_latents - else: - batch.did_sp_shard_audio_latents = False - - # For LTX-2 packed token latents, SP sharding happens on the time dimension - # (frames). The model must see local latent frames (RoPE offset is applied - # inside the model using SP rank). - latent_num_frames_for_model = self._get_video_latent_num_frames_for_model( - batch=batch, server_args=server_args, latents=latents - ) - latent_height = ( - batch.height - // server_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio - ) - latent_width = ( - batch.width - // server_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio - ) - - # Initialize lists for ODE trajectory - trajectory_timesteps: list[torch.Tensor] = [] - trajectory_latents: list[torch.Tensor] = [] - trajectory_audio_latents: list[torch.Tensor] = [] - - # Run denoising loop - denoising_start_time = time.time() - - # to avoid device-sync caused by timestep comparison - is_warmup = batch.is_warmup - self.scheduler.set_begin_index(0) - audio_scheduler.set_begin_index(0) - timesteps_cpu = timesteps.cpu() - num_timesteps = timesteps_cpu.shape[0] - - num_img_tokens = int(getattr(batch, "ltx2_num_image_tokens", 0)) - denoise_mask = None - clean_latent = None - if do_ti2v: - if not (isinstance(latents, torch.Tensor) and latents.ndim == 3): - raise ValueError("LTX-2 TI2V expects packed token latents [B, S, D].") - use_zero_clean_latent = is_ltx23_native_variant( - server_args.pipeline_config.vae_config.arch_config - ) - latents, denoise_mask, clean_latent = self._prepare_ltx2_ti2v_clean_state( - latents=latents, - image_latent=batch.image_latent, - num_img_tokens=num_img_tokens, - zero_clean_latent=use_zero_clean_latent, - ) - with torch.autocast( - device_type=current_platform.device_type, - dtype=target_dtype, - enabled=autocast_enabled, - ): - with self.progress_bar(total=num_inference_steps) as progress_bar: - for i, t_host in enumerate(timesteps_cpu): - with StageProfiler( - f"denoising_step_{i}", - logger=logger, - metrics=batch.metrics, - perf_dump_path_provided=batch.perf_dump_path is not None, - record_as_step=True, - ): - t_int = int(t_host.item()) - t_device = timesteps[i] - ( - current_model, - current_guidance_scale, - ) = self._select_and_manage_model( - t_int=t_int, - boundary_timestep=boundary_timestep, - server_args=server_args, - batch=batch, - ) - - # Predict noise residual - attn_metadata = self._build_attn_metadata(i, batch, server_args) - - # === LTX-2 sigma-space Euler step (flow matching) === - # Use scheduler-generated sigmas (includes terminal sigma=0). - sigmas = getattr(self.scheduler, "sigmas", None) - if sigmas is None or not isinstance(sigmas, torch.Tensor): - raise ValueError( - "Expected scheduler.sigmas to be a tensor for LTX-2." - ) - sigma = sigmas[i].to(device=latents.device, dtype=torch.float32) - sigma_next = sigmas[i + 1].to( - device=latents.device, dtype=torch.float32 - ) - dt = sigma_next - sigma - - latent_model_input = latents.to(target_dtype) - audio_latent_model_input = audio_latents.to(target_dtype) - stage1_guider_params = self._get_ltx2_stage1_guider_params( - batch, server_args, stage - ) - latent_num_frames = latent_num_frames_for_model - - # Audio latent dims - if audio_latent_model_input.ndim == 3: - audio_num_frames_latent = int( - audio_latent_model_input.shape[1] - ) - elif audio_latent_model_input.ndim == 4: - audio_num_frames_latent = int( - audio_latent_model_input.shape[2] - ) - else: - raise ValueError( - f"Unexpected audio latents rank: {audio_latent_model_input.ndim}, shape={tuple(audio_latent_model_input.shape)}" - ) - - video_coords = None - audio_coords = None - if not use_ltx23_legacy_one_stage: - video_coords = server_args.pipeline_config.prepare_video_rope_coords_for_sp( - current_model, - batch, - latent_model_input, - num_frames=latent_num_frames, - height=latent_height, - width=latent_width, - ) - audio_coords = server_args.pipeline_config.prepare_audio_rope_coords_for_sp( - current_model, - batch, - audio_latent_model_input, - num_frames=audio_num_frames_latent, - ) - - batch_size = int(latent_model_input.shape[0]) - video_num_tokens = int(latent_model_input.shape[1]) - is_ltx23_variant = is_ltx23_native_variant( - server_args.pipeline_config.vae_config.arch_config - ) - timestep = t_device.expand(batch_size) - if do_ti2v and denoise_mask is not None: - timestep_video = timestep.unsqueeze( - -1 - ) * denoise_mask.squeeze(-1) - elif is_ltx23_variant and not use_ltx23_legacy_one_stage: - timestep_video = timestep.view(batch_size, 1).expand( - batch_size, video_num_tokens - ) - else: - timestep_video = timestep - - if ( - is_ltx23_variant - and not use_ltx23_legacy_one_stage - and audio_latent_model_input.ndim == 3 - ): - audio_num_tokens = int(audio_latent_model_input.shape[1]) - timestep_audio = timestep.view(batch_size, 1).expand( - batch_size, audio_num_tokens - ) - else: - timestep_audio = timestep - prompt_timestep_video = None - prompt_timestep_audio = None - if is_ltx23_variant and not use_ltx23_legacy_one_stage: - timestep_scale_multiplier = float( - getattr( - current_model, "timestep_scale_multiplier", 1000 - ) - ) - prompt_timestep_video = ( - sigma.to( - device=latent_model_input.device, - dtype=torch.float32, - ) - * timestep_scale_multiplier - ).expand(batch_size) - prompt_timestep_audio = ( - sigma.to( - device=audio_latent_model_input.device, - dtype=torch.float32, - ) - * timestep_scale_multiplier - ).expand(batch_size) - - if use_ltx23_legacy_one_stage: - video_self_attention_mask = None - audio_self_attention_mask = None - a2v_cross_attention_mask = None - v2a_cross_attention_mask = None - else: - video_self_attention_mask = ( - self._build_ltx2_sp_padding_mask( - batch, - seq_len=video_num_tokens, - batch_size=batch_size, - key="sp_video_valid_token_count", - device=latent_model_input.device, - ) - ) - audio_self_attention_mask = ( - self._build_ltx2_sp_padding_mask( - batch, - seq_len=audio_num_frames_latent, - batch_size=batch_size, - key="sp_audio_valid_token_count", - device=audio_latent_model_input.device, - ) - ) - a2v_cross_attention_mask = audio_self_attention_mask - v2a_cross_attention_mask = video_self_attention_mask - - def build_model_kwargs( - *, - encoder_hidden_states: torch.Tensor, - audio_encoder_hidden_states: torch.Tensor, - encoder_attention_mask: torch.Tensor | None, - skip_video_self_attn_blocks: tuple[int, ...] | None = None, - skip_audio_self_attn_blocks: tuple[int, ...] | None = None, - disable_a2v_cross_attn: bool = False, - disable_v2a_cross_attn: bool = False, - ) -> dict[str, object]: - kwargs: dict[str, object] = { - "hidden_states": latent_model_input, - "audio_hidden_states": audio_latent_model_input, - "encoder_hidden_states": encoder_hidden_states, - "audio_encoder_hidden_states": audio_encoder_hidden_states, - "timestep": timestep_video, - "audio_timestep": timestep_audio, - "encoder_attention_mask": encoder_attention_mask, - "audio_encoder_attention_mask": encoder_attention_mask, - "num_frames": latent_num_frames, - "height": latent_height, - "width": latent_width, - "fps": batch.fps, - "audio_num_frames": audio_num_frames_latent, - "video_coords": video_coords, - "audio_coords": audio_coords, - "return_latents": False, - "return_dict": False, - } - if not use_ltx23_legacy_one_stage: - kwargs.update( - { - "prompt_timestep": prompt_timestep_video, - "audio_prompt_timestep": prompt_timestep_audio, - "video_self_attention_mask": video_self_attention_mask, - "audio_self_attention_mask": audio_self_attention_mask, - "a2v_cross_attention_mask": a2v_cross_attention_mask, - "v2a_cross_attention_mask": v2a_cross_attention_mask, - "audio_replicated_for_sp": replicate_audio_for_sp, - "legacy_ltx23_one_stage_semantics": False, - } - ) - if skip_video_self_attn_blocks is not None: - kwargs["skip_video_self_attn_blocks"] = ( - skip_video_self_attn_blocks - ) - if skip_audio_self_attn_blocks is not None: - kwargs["skip_audio_self_attn_blocks"] = ( - skip_audio_self_attn_blocks - ) - if disable_a2v_cross_attn: - kwargs["disable_a2v_cross_attn"] = True - if disable_v2a_cross_attn: - kwargs["disable_v2a_cross_attn"] = True - return kwargs - - use_official_cfg_path = stage1_guider_params is None - if use_official_cfg_path: - encoder_hidden_states = batch.prompt_embeds[0] - audio_encoder_hidden_states = batch.audio_prompt_embeds[0] - encoder_attention_mask = ( - self._get_ltx_prompt_attention_mask( - batch, - is_ltx23_variant=( - is_ltx23_variant - and not use_ltx23_legacy_one_stage - ), - ) - ) - if batch.do_classifier_free_guidance: - latent_model_input = torch.cat( - [latent_model_input] * 2, dim=0 - ) - audio_latent_model_input = torch.cat( - [audio_latent_model_input] * 2, dim=0 - ) - encoder_hidden_states = torch.cat( - [ - batch.negative_prompt_embeds[0], - encoder_hidden_states, - ], - dim=0, - ) - audio_encoder_hidden_states = torch.cat( - [ - batch.negative_audio_prompt_embeds[0], - audio_encoder_hidden_states, - ], - dim=0, - ) - if encoder_attention_mask is not None: - encoder_attention_mask = torch.cat( - [ - self._get_ltx_prompt_attention_mask( - batch, - is_ltx23_variant=( - is_ltx23_variant - and not use_ltx23_legacy_one_stage - ), - negative=True, - ), - encoder_attention_mask, - ], - dim=0, - ) - cfg_batch_size = int(latent_model_input.shape[0]) - timestep_video = self._repeat_batch_dim( - timestep_video, cfg_batch_size - ) - timestep_audio = self._repeat_batch_dim( - timestep_audio, cfg_batch_size - ) - if prompt_timestep_video is not None: - prompt_timestep_video = self._repeat_batch_dim( - prompt_timestep_video, cfg_batch_size - ) - if prompt_timestep_audio is not None: - prompt_timestep_audio = self._repeat_batch_dim( - prompt_timestep_audio, cfg_batch_size - ) - if video_self_attention_mask is not None: - video_self_attention_mask = self._repeat_batch_dim( - video_self_attention_mask, cfg_batch_size - ) - if audio_self_attention_mask is not None: - audio_self_attention_mask = self._repeat_batch_dim( - audio_self_attention_mask, cfg_batch_size - ) - if a2v_cross_attention_mask is not None: - a2v_cross_attention_mask = self._repeat_batch_dim( - a2v_cross_attention_mask, cfg_batch_size - ) - if v2a_cross_attention_mask is not None: - v2a_cross_attention_mask = self._repeat_batch_dim( - v2a_cross_attention_mask, cfg_batch_size - ) - - with set_forward_context( - current_timestep=i, attn_metadata=attn_metadata - ): - model_video, model_audio = current_model( - **build_model_kwargs( - encoder_hidden_states=encoder_hidden_states, - audio_encoder_hidden_states=audio_encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - ) - ) - - model_video = model_video.float() - model_audio = model_audio.float() - if batch.do_classifier_free_guidance: - ( - model_video_uncond, - model_video_text, - ) = model_video.chunk(2) - ( - model_audio_uncond, - model_audio_text, - ) = model_audio.chunk(2) - model_video = model_video_uncond + ( - batch.guidance_scale - * (model_video_text - model_video_uncond) - ) - model_audio = model_audio_uncond + ( - batch.guidance_scale - * (model_audio_text - model_audio_uncond) - ) - v_pos = model_video - a_v_pos = model_audio - v_neg = None - a_v_neg = None - - latents = self.scheduler.step( - v_pos, t_device, latents, return_dict=False - )[0] - audio_latents = audio_scheduler.step( - a_v_pos, t_device, audio_latents, return_dict=False - )[0] - latents = self.post_forward_for_ti2v_task( - batch, server_args, reserved_frames_mask, latents, z - ) - - if batch.return_trajectory_latents: - trajectory_timesteps.append(t_host) - trajectory_latents.append(latents) - if audio_latents is not None: - trajectory_audio_latents.append(audio_latents) - - if i == num_timesteps - 1 or ( - (i + 1) > num_warmup_steps - and (i + 1) % self.scheduler.order == 0 - and progress_bar is not None - ): - progress_bar.update() - - if not is_warmup: - self.step_profile() - continue - else: - # Follow ltx-pipelines structure: separate pos/neg forward passes, - # then apply CFG on denoised (x0) predictions. - encoder_hidden_states = batch.prompt_embeds[0] - audio_encoder_hidden_states = batch.audio_prompt_embeds[0] - encoder_attention_mask = ( - self._get_ltx_prompt_attention_mask( - batch, - is_ltx23_variant=( - is_ltx23_variant - and not use_ltx23_legacy_one_stage - ), - ) - ) - with set_forward_context( - current_timestep=i, attn_metadata=attn_metadata - ): - v_pos, a_v_pos = current_model( - **build_model_kwargs( - encoder_hidden_states=encoder_hidden_states, - audio_encoder_hidden_states=audio_encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - ) - ) - - if ( - stage1_guider_params is not None - or batch.do_classifier_free_guidance - ): - neg_encoder_hidden_states = ( - batch.negative_prompt_embeds[0] - ) - neg_audio_encoder_hidden_states = ( - batch.negative_audio_prompt_embeds[0] - ) - neg_encoder_attention_mask = ( - self._get_ltx_prompt_attention_mask( - batch, - is_ltx23_variant=( - is_ltx23_variant - and not use_ltx23_legacy_one_stage - ), - negative=True, - ) - ) - - v_neg, a_v_neg = current_model( - **build_model_kwargs( - encoder_hidden_states=neg_encoder_hidden_states, - audio_encoder_hidden_states=neg_audio_encoder_hidden_states, - encoder_attention_mask=neg_encoder_attention_mask, - ) - ) - else: - v_neg = None - a_v_neg = None - - v_pos = v_pos.float() - a_v_pos = a_v_pos.float() - if v_neg is not None: - v_neg = v_neg.float() - if a_v_neg is not None: - a_v_neg = a_v_neg.float() - - sigma_val = float(sigma.item()) - video_sigma_for_x0: float | torch.Tensor = sigma_val - if do_ti2v and denoise_mask is not None: - video_sigma_for_x0 = sigma.to( - device=latents.device, dtype=torch.float32 - ) * denoise_mask.squeeze(-1) - denoised_video = self._ltx2_velocity_to_x0( - latents, v_pos, video_sigma_for_x0 - ) - denoised_audio = self._ltx2_velocity_to_x0( - audio_latents, a_v_pos, sigma_val - ) - denoised_video_cond = denoised_video - denoised_audio_cond = denoised_audio - denoised_video_neg = None - denoised_audio_neg = None - denoised_video_perturbed = None - denoised_audio_perturbed = None - denoised_video_modality = None - denoised_audio_modality = None - - if ( - ( - stage1_guider_params is not None - or batch.do_classifier_free_guidance - ) - and v_neg is not None - and a_v_neg is not None - ): - denoised_video_neg = self._ltx2_velocity_to_x0( - latents, v_neg, video_sigma_for_x0 - ) - denoised_audio_neg = self._ltx2_velocity_to_x0( - audio_latents, a_v_neg, sigma_val - ) - if stage1_guider_params is not None: - video_skip = self._ltx2_should_skip_step( - i, int(stage1_guider_params["video_skip_step"]) - ) - audio_skip = self._ltx2_should_skip_step( - i, int(stage1_guider_params["audio_skip_step"]) - ) - - need_perturbed = ( - float(stage1_guider_params["video_stg_scale"]) != 0.0 - or float(stage1_guider_params["audio_stg_scale"]) != 0.0 - ) - if need_perturbed: - with set_forward_context( - current_timestep=i, attn_metadata=attn_metadata - ): - v_ptb, a_v_ptb = current_model( - **build_model_kwargs( - encoder_hidden_states=encoder_hidden_states, - audio_encoder_hidden_states=audio_encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - skip_video_self_attn_blocks=tuple( - stage1_guider_params["video_stg_blocks"] - ), - skip_audio_self_attn_blocks=tuple( - stage1_guider_params["audio_stg_blocks"] - ), - ) - ) - denoised_video_perturbed = self._ltx2_velocity_to_x0( - latents, v_ptb.float(), video_sigma_for_x0 - ) - denoised_audio_perturbed = self._ltx2_velocity_to_x0( - audio_latents, a_v_ptb.float(), sigma_val - ) - - need_modality = ( - float(stage1_guider_params["video_modality_scale"]) - != 1.0 - or float(stage1_guider_params["audio_modality_scale"]) - != 1.0 - ) - if need_modality: - with set_forward_context( - current_timestep=i, attn_metadata=attn_metadata - ): - v_mod, a_v_mod = current_model( - **build_model_kwargs( - encoder_hidden_states=encoder_hidden_states, - audio_encoder_hidden_states=audio_encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - disable_a2v_cross_attn=True, - disable_v2a_cross_attn=True, - ) - ) - denoised_video_modality = self._ltx2_velocity_to_x0( - latents, v_mod.float(), video_sigma_for_x0 - ) - denoised_audio_modality = self._ltx2_velocity_to_x0( - audio_latents, a_v_mod.float(), sigma_val - ) - - if not video_skip: - denoised_video = self._ltx2_calculate_guided_x0( - cond=denoised_video, - uncond_text=( - denoised_video_neg - if denoised_video_neg is not None - else denoised_video - ), - uncond_perturbed=( - denoised_video_perturbed - if denoised_video_perturbed is not None - else 0.0 - ), - uncond_modality=( - denoised_video_modality - if denoised_video_modality is not None - else 0.0 - ), - cfg_scale=float( - stage1_guider_params["video_cfg_scale"] - ), - stg_scale=float( - stage1_guider_params["video_stg_scale"] - ), - rescale_scale=float( - stage1_guider_params["video_rescale_scale"] - ), - modality_scale=float( - stage1_guider_params["video_modality_scale"] - ), - ) - if not audio_skip: - denoised_audio = self._ltx2_calculate_guided_x0( - cond=denoised_audio, - uncond_text=( - denoised_audio_neg - if denoised_audio_neg is not None - else denoised_audio - ), - uncond_perturbed=( - denoised_audio_perturbed - if denoised_audio_perturbed is not None - else 0.0 - ), - uncond_modality=( - denoised_audio_modality - if denoised_audio_modality is not None - else 0.0 - ), - cfg_scale=float( - stage1_guider_params["audio_cfg_scale"] - ), - stg_scale=float( - stage1_guider_params["audio_stg_scale"] - ), - rescale_scale=float( - stage1_guider_params["audio_rescale_scale"] - ), - modality_scale=float( - stage1_guider_params["audio_modality_scale"] - ), - ) - elif ( - batch.do_classifier_free_guidance - and denoised_video_neg is not None - and denoised_audio_neg is not None - ): - denoised_video = denoised_video + ( - batch.guidance_scale - 1.0 - ) * (denoised_video - denoised_video_neg) - denoised_audio = denoised_audio + ( - batch.guidance_scale - 1.0 - ) * (denoised_audio - denoised_audio_neg) - - # Apply conditioning mask (keep conditioned tokens clean). - if ( - do_ti2v - and denoise_mask is not None - and clean_latent is not None - ): - denoised_video = ( - denoised_video * denoise_mask - + clean_latent.float() * (1.0 - denoise_mask) - ) - # Euler step in sigma space: x_next = x + (sigma_next - sigma) * v, - # where v = (x - x0) / sigma. - if sigma_val == 0.0: - v_video = torch.zeros_like(denoised_video) - v_audio = torch.zeros_like(denoised_audio) - else: - v_video = ( - (latents.float() - denoised_video.float()) / sigma_val - ).to(latents.dtype) - v_audio = ( - (audio_latents.float() - denoised_audio.float()) - / sigma_val - ).to(audio_latents.dtype) - - latents = (latents.float() + v_video.float() * dt).to( - dtype=latents.dtype - ) - audio_latents = ( - audio_latents.float() + v_audio.float() * dt - ).to(dtype=audio_latents.dtype) - - latents = self.post_forward_for_ti2v_task( - batch, server_args, reserved_frames_mask, latents, z - ) - - # save trajectory latents if needed - if batch.return_trajectory_latents: - trajectory_timesteps.append(t_host) - trajectory_latents.append(latents) - if audio_latents is not None: - trajectory_audio_latents.append(audio_latents) - - # Update progress bar - if i == num_timesteps - 1 or ( - (i + 1) > num_warmup_steps - and (i + 1) % self.scheduler.order == 0 - and progress_bar is not None - ): - progress_bar.update() - - if not is_warmup: - self.step_profile() - - denoising_end_time = time.time() - - if num_timesteps > 0 and not is_warmup: - self.log_info( - "average time per step: %.4f seconds", - (denoising_end_time - denoising_start_time) / len(timesteps), - ) - - batch.audio_latents = audio_latents - self._post_denoising_loop( - batch=batch, - latents=latents, - trajectory_latents=trajectory_latents, - trajectory_timesteps=trajectory_timesteps, - trajectory_audio_latents=trajectory_audio_latents, - server_args=server_args, - is_warmup=is_warmup, - ) - - return batch def _post_denoising_loop( self, @@ -1925,7 +41,7 @@ def _post_denoising_loop( *args, **kwargs, ): - # 1. Handle Trajectory (Video) - Copy from base + """Finalize AV requests by gathering audio latents and unpacking both streams.""" if trajectory_latents: trajectory_tensor = torch.stack(trajectory_latents, dim=1) trajectory_timesteps_tensor = torch.stack(trajectory_timesteps, dim=0) @@ -1936,26 +52,16 @@ def _post_denoising_loop( latents, trajectory_tensor = self._postprocess_sp_latents( batch, latents, trajectory_tensor ) - - # If SP time-sharding padded whole frames worth of tokens, remove padding - # after gather and before unpacking. latents = self._truncate_sp_padded_token_latents(batch, latents) if trajectory_tensor is not None and trajectory_timesteps_tensor is not None: batch.trajectory_timesteps = trajectory_timesteps_tensor.cpu() batch.trajectory_latents = trajectory_tensor.cpu() - # 2. Handle Trajectory (Audio) - LTX-2 specific if trajectory_audio_latents: trajectory_audio_tensor = torch.stack(trajectory_audio_latents, dim=1) - # We don't have SP support for audio latents yet (or needed?) batch.trajectory_audio_latents = trajectory_audio_tensor.cpu() - # 3. Unpack and Denormalize - # Call pipeline_config._unpad_and_unpack_latents - # latents is video latents. - # batch.audio_latents is audio latents. - audio_latents = batch.audio_latents if batch.did_sp_shard_audio_latents and isinstance(audio_latents, torch.Tensor): audio_latents = server_args.pipeline_config.gather_audio_latents_for_sp( @@ -1963,7 +69,6 @@ def _post_denoising_loop( ) batch.audio_latents = audio_latents - # NOTE: self.vae and self.audio_vae should be populated via __init__ or manual setting if self.vae is None or self.audio_vae is None: logger.warning( "VAE or Audio VAE not found in DenoisingStage. Skipping unpack and denormalize." @@ -1971,66 +76,22 @@ def _post_denoising_loop( batch.latents = latents batch.audio_latents = audio_latents else: - ( - latents, - audio_latents, - ) = server_args.pipeline_config._unpad_and_unpack_latents( - latents, audio_latents, batch, self.vae, self.audio_vae + latents, audio_latents = ( + server_args.pipeline_config._unpad_and_unpack_latents( + latents, audio_latents, batch, self.vae, self.audio_vae + ) ) - batch.latents = latents batch.audio_latents = audio_latents - # 4. Cleanup - # TODO: make this a general denoising-stage hook + if isinstance(self.transformer, OffloadableDiTMixin): for manager in self.transformer.layerwise_offload_managers: manager.release_all() - def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: - """Verify denoising stage inputs. - - Note: LTX-2 connector stage converts `prompt_embeds`/`negative_prompt_embeds` - from list-of-tensors to a single tensor (video context) and stores audio - context separately. - """ - - result = VerificationResult() - result.add_check("timesteps", batch.timesteps, [V.is_tensor, V.min_dims(1)]) - - # LTX-2 may carry prompt embeddings as either a tensor (preferred) or legacy list. - result.add_check( - "prompt_embeds", - batch.prompt_embeds, - lambda x: V.is_tensor(x) or V.list_not_empty(x), - ) - - # Keep base expectation: image_embeds is always a list (may be empty). - result.add_check("image_embeds", batch.image_embeds, V.is_list) - - result.add_check( - "num_inference_steps", batch.num_inference_steps, V.positive_int - ) - result.add_check("guidance_scale", batch.guidance_scale, V.non_negative_float) - result.add_check("eta", batch.eta, V.non_negative_float) - result.add_check("generator", batch.generator, V.generator_or_list_generators) - result.add_check( - "do_classifier_free_guidance", - batch.do_classifier_free_guidance, - V.bool_value, - ) - - # When CFG is enabled, negative prompt embeddings must exist (tensor or legacy list). - result.add_check( - "negative_prompt_embeds", - batch.negative_prompt_embeds, - lambda x: (not batch.do_classifier_free_guidance) - or V.is_tensor(x) - or V.list_not_empty(x), - ) - return result - class LTX2RefinementStage(LTX2AVDenoisingStage): + """Stage-2 refinement wrapper that re-noises distilled LTX latents once.""" + def __init__( self, transformer, scheduler, distilled_sigmas, vae=None, audio_vae=None ): @@ -2085,9 +146,6 @@ def _reset_stage2_generators(batch: Req) -> None: @staticmethod def _should_reset_stage2_generators(server_args: ServerArgs) -> bool: - # Official LTX-2.3 two-stage refinement continues from the generator state - # after stage 1. Resetting back to the request seed changes the distilled - # noise injection immediately at stage 2 step 0. arch_config = getattr( server_args.pipeline_config.vae_config, "arch_config", None ) @@ -2096,6 +154,7 @@ def _should_reset_stage2_generators(server_args: ServerArgs) -> bool: return "LTX-2.3" not in str(getattr(server_args, "model_path", "")) def forward(self, batch: Req, server_args: ServerArgs) -> Req: + """Run the distilled refinement schedule on top of the shared AV denoiser.""" batch.extra["ltx2_phase"] = "stage2" if self._should_reset_stage2_generators(server_args): self._reset_stage2_generators(batch) @@ -2125,11 +184,9 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req: device=batch.audio_latents.device, dtype=torch.float32 ) - # Stage 2 runs at full resolution, so Stage 1 TI2V conditioning is invalid. batch.image_latent = None batch.ltx2_num_image_tokens = 0 - # Use a private scheduler copy to avoid mutating shared state. original_scheduler = self.scheduler original_batch_timesteps = batch.timesteps original_batch_num_inference_steps = batch.num_inference_steps diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py new file mode 100644 index 000000000000..a644fc693f88 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py @@ -0,0 +1,1240 @@ +import copy +import json +import math +import os +from dataclasses import dataclass, field +from io import BytesIO + +import av +import numpy as np +import PIL.Image +import torch +from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution +from diffusers.models.modeling_outputs import AutoencoderKLOutput +from safetensors.torch import load_file as safetensors_load_file + +from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import ( + is_ltx23_native_variant, +) +from sglang.multimodal_gen.runtime.distributed import get_sp_world_size +from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context +from sglang.multimodal_gen.runtime.models.vaes.ltx_2_3_condition_encoder import ( + LTX23VideoConditionEncoder, +) +from sglang.multimodal_gen.runtime.models.vision_utils import ( + load_image, + normalize, + numpy_to_pt, + pil_to_numpy, +) +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import ( + DenoisingContext, + DenoisingStage, + DenoisingStepState, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( + StageValidators as V, +) +from sglang.multimodal_gen.runtime.platforms import current_platform +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.multimodal_gen.utils import PRECISION_TO_TYPE + +logger = init_logger(__name__) + + +@dataclass(slots=True) +class LTX2DenoisingContext(DenoisingContext): + """Loop-scoped denoising state for joint LTX-2 video and audio generation.""" + + audio_latents: torch.Tensor | None = None + audio_scheduler: object | None = None + is_ltx23_variant: bool = False + use_ltx23_legacy_one_stage: bool = False + replicate_audio_for_sp: bool = False + stage: str = "one_stage" + latent_num_frames_for_model: int = 0 + latent_height: int = 0 + latent_width: int = 0 + denoise_mask: torch.Tensor | None = None + clean_latent: torch.Tensor | None = None + trajectory_audio_latents: list[torch.Tensor] = field(default_factory=list) + + +class LTX2DenoisingStage(DenoisingStage): + """ + LTX-2 specific denoising stage that handles joint video and audio generation. + """ + + def __init__(self, transformer, scheduler, vae=None, **kwargs): + super().__init__( + transformer=transformer, scheduler=scheduler, vae=vae, **kwargs + ) + self._condition_image_encoder = None + self._condition_image_encoder_dir = None + + @staticmethod + def _get_video_latent_num_frames_for_model( + batch: Req, server_args: ServerArgs, latents: torch.Tensor + ) -> int: + """Return the latent-frame length the DiT model should see. + + - If video latents were time-sharded for SP and are packed as token latents + ([B, S, D]), the model only sees the local shard and must use the local + latent-frame count (stored on the batch during SP sharding). + - Otherwise, fall back to the global latent-frame count inferred from the + requested output frames and the VAE temporal compression ratio. + """ + did_sp_shard = bool(getattr(batch, "did_sp_shard_latents", False)) + is_token_latents = isinstance(latents, torch.Tensor) and latents.ndim == 3 + + if did_sp_shard and is_token_latents: + if not hasattr(batch, "sp_video_latent_num_frames"): + raise ValueError( + "SP-sharded LTX2 token latents require `batch.sp_video_latent_num_frames` " + "to be set by `LTX2PipelineConfig.shard_latents_for_sp()`." + ) + return int(batch.sp_video_latent_num_frames) + + pc = server_args.pipeline_config + return int( + (batch.num_frames - 1) + // int(pc.vae_config.arch_config.temporal_compression_ratio) + + 1 + ) + + @staticmethod + def _truncate_sp_padded_token_latents( + batch: Req, latents: torch.Tensor + ) -> torch.Tensor: + """Remove token padding introduced by SP time-sharding (if applicable).""" + did_sp_shard = bool(getattr(batch, "did_sp_shard_latents", False)) + if not did_sp_shard or not ( + isinstance(latents, torch.Tensor) and latents.ndim == 3 + ): + return latents + + raw_shape = getattr(batch, "raw_latent_shape", None) + if not (isinstance(raw_shape, tuple) and len(raw_shape) == 3): + return latents + + orig_s = int(raw_shape[1]) + cur_s = int(latents.shape[1]) + if cur_s == orig_s: + return latents + if cur_s < orig_s: + raise ValueError( + f"Unexpected gathered token-latents seq_len {cur_s} < original seq_len {orig_s}." + ) + return latents[:, :orig_s, :].contiguous() + + def _maybe_enable_cache_dit(self, num_inference_steps: int, batch: Req) -> None: + """Disable cache-dit for TI2V-style requests (image-conditioned), to avoid stale activations. + + NOTE: base denoising stage calls this hook with (num_inference_steps, batch). + """ + if getattr(self, "_disable_cache_dit_for_request", False): + return + return super()._maybe_enable_cache_dit(num_inference_steps, batch) + + def _get_ltx2_stage1_guider_params( + self, batch: Req, server_args: ServerArgs, stage: str + ) -> dict[str, object] | None: + if stage != "stage1": + return None + return batch.extra.get("ltx2_stage1_guider_params") + + @staticmethod + def _ltx2_should_skip_step(step_index: int, skip_step: int) -> bool: + if skip_step == 0: + return False + return step_index % (skip_step + 1) != 0 + + @staticmethod + def _ltx2_apply_rescale( + cond: torch.Tensor, pred: torch.Tensor, rescale_scale: float + ) -> torch.Tensor: + if rescale_scale == 0.0: + return pred + factor = cond.std() / pred.std() + factor = rescale_scale * factor + (1.0 - rescale_scale) + return pred * factor + + @staticmethod + def _prepare_ltx2_ti2v_clean_state( + latents: torch.Tensor, + image_latent: torch.Tensor, + num_img_tokens: int, + zero_clean_latent: bool, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + latents = latents.clone() + conditioned = image_latent[:, :num_img_tokens, :].to( + device=latents.device, dtype=latents.dtype + ) + latents[:, :num_img_tokens, :] = conditioned + denoise_mask = torch.ones( + (latents.shape[0], latents.shape[1], 1), + device=latents.device, + dtype=torch.float32, + ) + denoise_mask[:, :num_img_tokens, :] = 0.0 + if zero_clean_latent: + clean_latent = torch.zeros_like(latents) + else: + clean_latent = latents.detach().clone() + clean_latent[:, :num_img_tokens, :] = conditioned + return latents, denoise_mask, clean_latent + + @staticmethod + def _ltx2_velocity_to_x0( + sample: torch.Tensor, + velocity: torch.Tensor, + sigma: float | torch.Tensor, + ) -> torch.Tensor: + if isinstance(sigma, torch.Tensor): + sigma = sigma.to(device=sample.device, dtype=torch.float32) + while sigma.ndim < sample.ndim: + sigma = sigma.unsqueeze(-1) + return (sample.float() - sigma * velocity.float()).to(sample.dtype) + return (sample.float() - float(sigma) * velocity.float()).to(sample.dtype) + + @staticmethod + def _repeat_batch_dim(tensor: torch.Tensor, target_batch_size: int) -> torch.Tensor: + """Repeat along batch dim while preserving any tokenwise timestep layout.""" + if tensor.shape[0] == int(target_batch_size): + return tensor + if tensor.shape[0] <= 0 or int(target_batch_size) % int(tensor.shape[0]) != 0: + raise ValueError( + f"Cannot repeat tensor with batch={tensor.shape[0]} to target_batch_size={target_batch_size}" + ) + repeat_factor = int(target_batch_size) // int(tensor.shape[0]) + return tensor.repeat(repeat_factor, *([1] * (tensor.ndim - 1))) + + @staticmethod + def _build_ltx2_sp_padding_mask( + batch: Req, + *, + seq_len: int, + batch_size: int, + key: str, + device: torch.device, + ) -> torch.Tensor | None: + valid = getattr(batch, key, None) + if valid is None: + return None + valid = int(valid) + if valid <= 0 or valid >= int(seq_len): + return None + mask = torch.ones( + (batch_size, int(seq_len)), device=device, dtype=torch.float32 + ) + mask[:, valid:] = 0.0 + return mask + + @staticmethod + def _get_ltx_prompt_attention_mask( + batch: Req, + *, + is_ltx23_variant: bool, + negative: bool = False, + ) -> torch.Tensor | None: + if is_ltx23_variant: + return None + return ( + batch.negative_attention_mask if negative else batch.prompt_attention_mask + ) + + @classmethod + def _should_use_ltx23_legacy_one_stage( + cls, + server_args: ServerArgs, + pipeline_name: str | None, + ) -> bool: + if not is_ltx23_native_variant( + server_args.pipeline_config.vae_config.arch_config + ): + return False + if server_args.pipeline_class_name == "LTX2TwoStagePipeline": + return False + return pipeline_name != "LTX2TwoStagePipeline" + + @classmethod + def _should_shard_ltx23_legacy_one_stage_audio_latents( + cls, + batch: Req, + server_args: ServerArgs, + ) -> bool: + return bool( + get_sp_world_size() > 1 + and is_ltx23_native_variant( + server_args.pipeline_config.vae_config.arch_config + ) + and cls._should_use_ltx23_legacy_one_stage(server_args, None) + and server_args.pipeline_config.can_shard_audio_latents_for_sp( + batch.audio_latents + ) + ) + + @classmethod + def _ltx2_calculate_guided_x0( + cls, + *, + cond: torch.Tensor, + uncond_text: torch.Tensor | float, + uncond_perturbed: torch.Tensor | float, + uncond_modality: torch.Tensor | float, + cfg_scale: float, + stg_scale: float, + rescale_scale: float, + modality_scale: float, + ) -> torch.Tensor: + pred = ( + cond + + (cfg_scale - 1.0) * (cond - uncond_text) + + stg_scale * (cond - uncond_perturbed) + + (modality_scale - 1.0) * (cond - uncond_modality) + ) + return cls._ltx2_apply_rescale(cond, pred, rescale_scale) + + @staticmethod + def _resize_center_crop( + img: PIL.Image.Image, *, width: int, height: int + ) -> PIL.Image.Image: + return img.resize((width, height), resample=PIL.Image.Resampling.BILINEAR) + + @staticmethod + def _apply_video_codec_compression( + img_array: np.ndarray, crf: int = 33 + ) -> np.ndarray: + """Encode as a single H.264 frame and decode back to simulate compression artifacts.""" + if crf == 0: + return img_array + height, width = img_array.shape[0] // 2 * 2, img_array.shape[1] // 2 * 2 + img_array = img_array[:height, :width] + buffer = BytesIO() + container = av.open(buffer, mode="w", format="mp4") + stream = container.add_stream( + "libx264", rate=1, options={"crf": str(crf), "preset": "veryfast"} + ) + stream.height, stream.width = height, width + frame = av.VideoFrame.from_ndarray(img_array, format="rgb24").reformat( + format="yuv420p" + ) + container.mux(stream.encode(frame)) + container.mux(stream.encode()) + container.close() + buffer.seek(0) + container = av.open(buffer) + decoded = next(container.decode(container.streams.video[0])) + container.close() + return decoded.to_ndarray(format="rgb24") + + @staticmethod + def _resize_center_crop_tensor( + img: PIL.Image.Image, + *, + width: int, + height: int, + device: torch.device, + dtype: torch.dtype, + apply_codec_compression: bool = True, + codec_crf: int = 33, + ) -> torch.Tensor: + """Resize, center-crop, and normalize to [1, C, 1, H, W] tensor in [-1, 1].""" + img_array = np.array(img).astype(np.uint8)[..., :3] + if apply_codec_compression: + img_array = LTX2DenoisingStage._apply_video_codec_compression( + img_array, crf=codec_crf + ) + tensor = ( + torch.from_numpy(img_array.astype(np.float32)) + .permute(2, 0, 1) + .unsqueeze(0) + .to(device=device) + ) + src_h, src_w = tensor.shape[2], tensor.shape[3] + scale = max(height / src_h, width / src_w) + new_h, new_w = math.ceil(src_h * scale), math.ceil(src_w * scale) + tensor = torch.nn.functional.interpolate( + tensor, size=(new_h, new_w), mode="bilinear", align_corners=False + ) + top, left = (new_h - height) // 2, (new_w - width) // 2 + tensor = tensor[:, :, top : top + height, left : left + width] + return ((tensor / 127.5 - 1.0).to(dtype=dtype)).unsqueeze(2) + + @staticmethod + def _pil_to_normed_tensor(img: PIL.Image.Image) -> torch.Tensor: + # PIL -> numpy [0,1] -> torch [B,C,H,W], then [-1,1] + arr = pil_to_numpy(img) + t = numpy_to_pt(arr) + return normalize(t) + + @staticmethod + def _should_apply_ltx2_ti2v(batch: Req) -> bool: + """True if we have an image-latent token prefix to condition with. + + SP note: when token latents are time-sharded, only the rank that owns the + *global* first latent frame should apply TI2V conditioning (rank with start_frame==0). + """ + if ( + batch.image_latent is None + or int(getattr(batch, "ltx2_num_image_tokens", 0)) <= 0 + ): + return False + did_sp_shard = bool(getattr(batch, "did_sp_shard_latents", False)) + if not did_sp_shard: + return True + return int(getattr(batch, "sp_video_start_frame", 0)) == 0 + + @staticmethod + def _should_replicate_ltx23_audio_for_sp( + batch: Req, + server_args: ServerArgs, + *, + is_ltx23_variant: bool, + ) -> bool: + return False + + def _get_condition_image_encoder( + self, + server_args: ServerArgs, + *, + device: torch.device, + dtype: torch.dtype, + ) -> LTX23VideoConditionEncoder | None: + arch_config = server_args.pipeline_config.vae_config.arch_config + encoder_subdir = str(getattr(arch_config, "condition_encoder_subdir", "")) + if not encoder_subdir: + return None + + vae_model_path = server_args.model_paths["vae"] + encoder_dir = os.path.join(vae_model_path, encoder_subdir) + config_path = os.path.join(encoder_dir, "config.json") + weights_path = os.path.join(encoder_dir, "model.safetensors") + if not os.path.exists(config_path) or not os.path.exists(weights_path): + raise ValueError( + f"LTX-2 condition encoder files not found under {encoder_dir}" + ) + + cached_dir = self._condition_image_encoder_dir + encoder = self._condition_image_encoder + if encoder is None or cached_dir != encoder_dir: + with open(config_path, encoding="utf-8") as f: + config = json.load(f) + encoder = LTX23VideoConditionEncoder(config) + encoder.load_state_dict(safetensors_load_file(weights_path), strict=True) + self._condition_image_encoder = encoder + self._condition_image_encoder_dir = encoder_dir + + encoder = encoder.to(device=device, dtype=dtype) + return encoder + + def _prepare_ltx2_image_latent(self, batch: Req, server_args: ServerArgs) -> None: + """Encode `batch.image_path` into packed token latents for LTX-2 TI2V.""" + if ( + batch.image_latent is not None + and int(getattr(batch, "ltx2_num_image_tokens", 0)) > 0 + ): + return + batch.ltx2_num_image_tokens = 0 + batch.image_latent = None + + if batch.image_path is None: + return + if batch.width is None or batch.height is None: + raise ValueError("width/height must be provided for LTX-2 TI2V.") + if self.vae is None: + raise ValueError("VAE must be provided for LTX-2 TI2V.") + + image_path = ( + batch.image_path[0] + if isinstance(batch.image_path, list) + else batch.image_path + ) + + img = load_image(image_path) + img_array = np.array(img).astype(np.uint8)[..., :3] + img_array = self._apply_video_codec_compression(img_array, crf=33) + conditioned_img = PIL.Image.fromarray(img_array) + batch.condition_image = self._resize_center_crop( + conditioned_img, width=int(batch.width), height=int(batch.height) + ) + + latents_device = ( + batch.latents.device + if isinstance(batch.latents, torch.Tensor) + else torch.device("cpu") + ) + encode_dtype = batch.latents.dtype + original_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] + vae_autocast_enabled = ( + original_dtype != torch.float32 + ) and not server_args.disable_autocast + condition_image_encoder = self._get_condition_image_encoder( + server_args, device=latents_device, dtype=encode_dtype + ) + if condition_image_encoder is None: + self.vae = self.vae.to(device=latents_device, dtype=encode_dtype) + + video_condition = self._resize_center_crop_tensor( + conditioned_img, + width=int(batch.width), + height=int(batch.height), + device=latents_device, + dtype=encode_dtype, + apply_codec_compression=False, + ) + + with torch.autocast( + device_type=current_platform.device_type, + dtype=original_dtype, + enabled=vae_autocast_enabled, + ): + try: + if ( + condition_image_encoder is None + and server_args.pipeline_config.vae_tiling + ): + self.vae.enable_tiling() + except Exception: + pass + if not vae_autocast_enabled: + video_condition = video_condition.to(encode_dtype) + + if condition_image_encoder is not None: + latent = condition_image_encoder(video_condition) + else: + latent_dist: DiagonalGaussianDistribution = self.vae.encode( + video_condition + ) + if isinstance(latent_dist, AutoencoderKLOutput): + latent_dist = latent_dist.latent_dist + + if condition_image_encoder is None: + mode = server_args.pipeline_config.vae_config.encode_sample_mode() + if mode == "argmax": + latent = latent_dist.mode() + elif mode == "sample": + if batch.generator is None: + raise ValueError("Generator must be provided for VAE sampling.") + latent = latent_dist.sample(batch.generator) + else: + raise ValueError(f"Unsupported encode_sample_mode: {mode}") + + # Per-channel normalization: normalized = (x - mean) / std + mean = self.vae.latents_mean.view(1, -1, 1, 1, 1).to(latent) + std = self.vae.latents_std.view(1, -1, 1, 1, 1).to(latent) + latent = (latent - mean) / std + else: + latent = latent.to(dtype=encode_dtype) + + packed = server_args.pipeline_config.maybe_pack_latents( + latent, latent.shape[0], batch + ) + if not (isinstance(packed, torch.Tensor) and packed.ndim == 3): + raise ValueError("Expected packed image latents [B, S0, D].") + + # Fail-fast token count: must match one latent frame's tokens. + vae_sf = int(server_args.pipeline_config.vae_scale_factor) + patch = int(server_args.pipeline_config.patch_size) + latent_h = int(batch.height) // vae_sf + latent_w = int(batch.width) // vae_sf + expected_tokens = (latent_h // patch) * (latent_w // patch) + if int(packed.shape[1]) != int(expected_tokens): + raise ValueError( + "LTX-2 conditioning token count mismatch: " + f"{int(packed.shape[1])=} {int(expected_tokens)=}." + ) + + batch.image_latent = packed + batch.ltx2_num_image_tokens = int(packed.shape[1]) + + if batch.debug: + logger.info( + "LTX2 TI2V conditioning prepared: %d tokens (shape=%s) for %sx%s", + batch.ltx2_num_image_tokens, + tuple(batch.image_latent.shape), + batch.width, + batch.height, + ) + + if condition_image_encoder is None: + self.vae.to(original_dtype) + if server_args.vae_cpu_offload: + self.vae = self.vae.to("cpu") + if condition_image_encoder is not None: + self._condition_image_encoder = condition_image_encoder.to("cpu") + + def _prepare_denoising_loop( + self, + batch: Req, + server_args: ServerArgs, + ) -> LTX2DenoisingContext: + """Extend the base context with LTX-2 audio, SP, and TI2V state.""" + self._disable_cache_dit_for_request = batch.image_path is not None + base_ctx = super()._prepare_denoising_loop(batch, server_args) + ctx = LTX2DenoisingContext(**base_ctx.to_kwargs()) + ctx.is_ltx23_variant = is_ltx23_native_variant( + server_args.pipeline_config.vae_config.arch_config + ) + phase = batch.extra.get("ltx2_phase") + pipeline = self.pipeline() if self.pipeline else None + pipeline_name = pipeline.pipeline_name if pipeline is not None else None + ctx.use_ltx23_legacy_one_stage = self._should_use_ltx23_legacy_one_stage( + server_args, pipeline_name + ) + ctx.stage = ( + phase + if phase is not None + else ("stage1" if ctx.use_ltx23_legacy_one_stage else "one_stage") + ) + ctx.audio_latents = batch.audio_latents + # Video and audio keep separate scheduler state throughout the denoising loop. + ctx.audio_scheduler = copy.deepcopy(self.scheduler) + + # Prepare image latents and embeddings for LTX-2 TI2V generation. + self._prepare_ltx2_image_latent(batch, server_args) + do_ti2v = self._should_apply_ltx2_ti2v(batch) + + if ctx.use_ltx23_legacy_one_stage: + batch.ltx23_audio_replicated_for_sp = False + batch.did_sp_shard_audio_latents = False + else: + ctx.replicate_audio_for_sp = self._should_replicate_ltx23_audio_for_sp( + batch, + server_args, + is_ltx23_variant=ctx.is_ltx23_variant, + ) + batch.ltx23_audio_replicated_for_sp = bool(ctx.replicate_audio_for_sp) + if ( + ctx.is_ltx23_variant + and get_sp_world_size() > 1 + and server_args.pipeline_config.can_shard_audio_latents_for_sp( + batch.audio_latents + ) + and not ctx.replicate_audio_for_sp + ): + ( + batch.audio_latents, + batch.did_sp_shard_audio_latents, + ) = server_args.pipeline_config.shard_audio_latents_for_sp( + batch, batch.audio_latents + ) + ctx.audio_latents = batch.audio_latents + else: + batch.did_sp_shard_audio_latents = False + + # For LTX-2 packed token latents, SP sharding happens on the time dimension + # (frames). The model must see local latent frames (RoPE offset is applied + # inside the model using SP rank). + ctx.latent_num_frames_for_model = self._get_video_latent_num_frames_for_model( + batch=batch, server_args=server_args, latents=ctx.latents + ) + ctx.latent_height = ( + batch.height + // server_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio + ) + ctx.latent_width = ( + batch.width + // server_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio + ) + if do_ti2v: + if not (isinstance(ctx.latents, torch.Tensor) and ctx.latents.ndim == 3): + raise ValueError("LTX-2 TI2V expects packed token latents [B, S, D].") + # Keep conditioned tokens clean and reuse the mask during every step update. + ctx.latents, ctx.denoise_mask, ctx.clean_latent = ( + self._prepare_ltx2_ti2v_clean_state( + latents=ctx.latents, + image_latent=batch.image_latent, + num_img_tokens=int(getattr(batch, "ltx2_num_image_tokens", 0)), + zero_clean_latent=ctx.is_ltx23_variant, + ) + ) + return ctx + + def _before_denoising_loop( + self, ctx: LTX2DenoisingContext, batch: Req, server_args: ServerArgs + ) -> None: + """Reset the mirrored audio scheduler before the shared loop begins.""" + super()._before_denoising_loop(ctx, batch, server_args) + if ctx.audio_scheduler is None: + raise ValueError("LTX-2 audio scheduler was not prepared.") + ctx.audio_scheduler.set_begin_index(0) + + def _prepare_step_attn_metadata( + self, + ctx: LTX2DenoisingContext, + batch: Req, + server_args: ServerArgs, + step_index: int, + t_int: int, + timesteps_cpu: torch.Tensor, + ): + """Preserve the legacy LTX-2 attention-metadata contract.""" + # Legacy LTX-2 paths used the plain attention-metadata builder call here. + del ctx, t_int, timesteps_cpu + return self._build_attn_metadata(step_index, batch, server_args) + + def _run_denoising_step( + self, + ctx: LTX2DenoisingContext, + step: DenoisingStepState, + batch: Req, + server_args: ServerArgs, + ) -> None: + """Run one joint video/audio denoising step with LTX-2-specific guidance.""" + if ctx.audio_latents is None: + raise ValueError("LTX-2 requires audio latents for denoising.") + if ctx.audio_scheduler is None: + raise ValueError("LTX-2 audio scheduler was not prepared.") + + # 1. Read the scheduler sigma pair and derive the Euler delta. + sigmas = getattr(self.scheduler, "sigmas", None) + if sigmas is None or not isinstance(sigmas, torch.Tensor): + raise ValueError("Expected scheduler.sigmas to be a tensor for LTX-2.") + sigma = sigmas[step.step_index].to( + device=ctx.latents.device, dtype=torch.float32 + ) + sigma_next = sigmas[step.step_index + 1].to( + device=ctx.latents.device, dtype=torch.float32 + ) + dt = sigma_next - sigma + + # 2. Materialize the current video/audio latent inputs in the compute dtype. + latent_model_input = ctx.latents.to(ctx.target_dtype) + audio_latent_model_input = ctx.audio_latents.to(ctx.target_dtype) + stage1_guider_params = self._get_ltx2_stage1_guider_params( + batch, server_args, ctx.stage + ) + + if audio_latent_model_input.ndim == 3: + audio_num_frames_latent = int(audio_latent_model_input.shape[1]) + elif audio_latent_model_input.ndim == 4: + audio_num_frames_latent = int(audio_latent_model_input.shape[2]) + else: + raise ValueError( + f"Unexpected audio latents rank: {audio_latent_model_input.ndim}, shape={tuple(audio_latent_model_input.shape)}" + ) + + # 3. Prepare any LTX-specific RoPE coordinates and timestep layouts. + video_coords = None + audio_coords = None + if not ctx.use_ltx23_legacy_one_stage: + video_coords = server_args.pipeline_config.prepare_video_rope_coords_for_sp( + step.current_model, + batch, + latent_model_input, + num_frames=ctx.latent_num_frames_for_model, + height=ctx.latent_height, + width=ctx.latent_width, + ) + audio_coords = server_args.pipeline_config.prepare_audio_rope_coords_for_sp( + step.current_model, + batch, + audio_latent_model_input, + num_frames=audio_num_frames_latent, + ) + + batch_size = int(latent_model_input.shape[0]) + timestep = step.t_device.expand(batch_size) + if ctx.denoise_mask is not None: + timestep_video = timestep.unsqueeze(-1) * ctx.denoise_mask.squeeze(-1) + elif ctx.is_ltx23_variant and not ctx.use_ltx23_legacy_one_stage: + timestep_video = timestep.view(batch_size, 1).expand( + batch_size, int(latent_model_input.shape[1]) + ) + else: + timestep_video = timestep + + if ( + ctx.is_ltx23_variant + and not ctx.use_ltx23_legacy_one_stage + and audio_latent_model_input.ndim == 3 + ): + timestep_audio = timestep.view(batch_size, 1).expand( + batch_size, int(audio_latent_model_input.shape[1]) + ) + else: + timestep_audio = timestep + + prompt_timestep_video = None + prompt_timestep_audio = None + if ctx.is_ltx23_variant and not ctx.use_ltx23_legacy_one_stage: + timestep_scale_multiplier = float( + getattr(step.current_model, "timestep_scale_multiplier", 1000) + ) + prompt_timestep_video = ( + sigma.to(device=latent_model_input.device, dtype=torch.float32) + * timestep_scale_multiplier + ).expand(batch_size) + prompt_timestep_audio = ( + sigma.to(device=audio_latent_model_input.device, dtype=torch.float32) + * timestep_scale_multiplier + ).expand(batch_size) + + # 4. Build attention masks that account for SP padding and replicated audio. + if ctx.use_ltx23_legacy_one_stage: + video_self_attention_mask = None + audio_self_attention_mask = None + a2v_cross_attention_mask = None + v2a_cross_attention_mask = None + else: + video_self_attention_mask = self._build_ltx2_sp_padding_mask( + batch, + seq_len=int(latent_model_input.shape[1]), + batch_size=batch_size, + key="sp_video_valid_token_count", + device=latent_model_input.device, + ) + audio_self_attention_mask = self._build_ltx2_sp_padding_mask( + batch, + seq_len=audio_num_frames_latent, + batch_size=batch_size, + key="sp_audio_valid_token_count", + device=audio_latent_model_input.device, + ) + a2v_cross_attention_mask = audio_self_attention_mask + v2a_cross_attention_mask = video_self_attention_mask + + def build_model_kwargs( + *, + encoder_hidden_states: torch.Tensor, + audio_encoder_hidden_states: torch.Tensor, + encoder_attention_mask: torch.Tensor | None, + skip_video_self_attn_blocks: tuple[int, ...] | None = None, + skip_audio_self_attn_blocks: tuple[int, ...] | None = None, + disable_a2v_cross_attn: bool = False, + disable_v2a_cross_attn: bool = False, + ) -> dict[str, object]: + kwargs: dict[str, object] = { + "hidden_states": latent_model_input, + "audio_hidden_states": audio_latent_model_input, + "encoder_hidden_states": encoder_hidden_states, + "audio_encoder_hidden_states": audio_encoder_hidden_states, + "timestep": timestep_video, + "audio_timestep": timestep_audio, + "encoder_attention_mask": encoder_attention_mask, + "audio_encoder_attention_mask": encoder_attention_mask, + "num_frames": ctx.latent_num_frames_for_model, + "height": ctx.latent_height, + "width": ctx.latent_width, + "fps": batch.fps, + "audio_num_frames": audio_num_frames_latent, + "video_coords": video_coords, + "audio_coords": audio_coords, + "return_latents": False, + "return_dict": False, + } + if not ctx.use_ltx23_legacy_one_stage: + kwargs.update( + { + "prompt_timestep": prompt_timestep_video, + "audio_prompt_timestep": prompt_timestep_audio, + "video_self_attention_mask": video_self_attention_mask, + "audio_self_attention_mask": audio_self_attention_mask, + "a2v_cross_attention_mask": a2v_cross_attention_mask, + "v2a_cross_attention_mask": v2a_cross_attention_mask, + "audio_replicated_for_sp": ctx.replicate_audio_for_sp, + "legacy_ltx23_one_stage_semantics": False, + } + ) + if skip_video_self_attn_blocks is not None: + kwargs["skip_video_self_attn_blocks"] = skip_video_self_attn_blocks + if skip_audio_self_attn_blocks is not None: + kwargs["skip_audio_self_attn_blocks"] = skip_audio_self_attn_blocks + if disable_a2v_cross_attn: + kwargs["disable_a2v_cross_attn"] = True + if disable_v2a_cross_attn: + kwargs["disable_v2a_cross_attn"] = True + return kwargs + + # 5. Run the branch-specific LTX forward path and apply CFG/guider logic. + prompt_attention_mask = self._get_ltx_prompt_attention_mask( + batch, + is_ltx23_variant=( + ctx.is_ltx23_variant and not ctx.use_ltx23_legacy_one_stage + ), + ) + use_official_cfg_path = stage1_guider_params is None + if use_official_cfg_path: + encoder_hidden_states = batch.prompt_embeds[0] + audio_encoder_hidden_states = batch.audio_prompt_embeds[0] + encoder_attention_mask = prompt_attention_mask + if batch.do_classifier_free_guidance: + latent_model_input = torch.cat([latent_model_input] * 2, dim=0) + audio_latent_model_input = torch.cat( + [audio_latent_model_input] * 2, dim=0 + ) + encoder_hidden_states = torch.cat( + [batch.negative_prompt_embeds[0], encoder_hidden_states], dim=0 + ) + audio_encoder_hidden_states = torch.cat( + [ + batch.negative_audio_prompt_embeds[0], + audio_encoder_hidden_states, + ], + dim=0, + ) + if encoder_attention_mask is not None: + encoder_attention_mask = torch.cat( + [ + self._get_ltx_prompt_attention_mask( + batch, + is_ltx23_variant=( + ctx.is_ltx23_variant + and not ctx.use_ltx23_legacy_one_stage + ), + negative=True, + ), + encoder_attention_mask, + ], + dim=0, + ) + cfg_batch_size = int(latent_model_input.shape[0]) + timestep_video = self._repeat_batch_dim(timestep_video, cfg_batch_size) + timestep_audio = self._repeat_batch_dim(timestep_audio, cfg_batch_size) + if prompt_timestep_video is not None: + prompt_timestep_video = self._repeat_batch_dim( + prompt_timestep_video, cfg_batch_size + ) + if prompt_timestep_audio is not None: + prompt_timestep_audio = self._repeat_batch_dim( + prompt_timestep_audio, cfg_batch_size + ) + if video_self_attention_mask is not None: + video_self_attention_mask = self._repeat_batch_dim( + video_self_attention_mask, cfg_batch_size + ) + if audio_self_attention_mask is not None: + audio_self_attention_mask = self._repeat_batch_dim( + audio_self_attention_mask, cfg_batch_size + ) + if a2v_cross_attention_mask is not None: + a2v_cross_attention_mask = self._repeat_batch_dim( + a2v_cross_attention_mask, cfg_batch_size + ) + if v2a_cross_attention_mask is not None: + v2a_cross_attention_mask = self._repeat_batch_dim( + v2a_cross_attention_mask, cfg_batch_size + ) + + with set_forward_context( + current_timestep=step.step_index, attn_metadata=step.attn_metadata + ): + model_video, model_audio = step.current_model( + **build_model_kwargs( + encoder_hidden_states=encoder_hidden_states, + audio_encoder_hidden_states=audio_encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + ) + ) + + model_video = model_video.float() + model_audio = model_audio.float() + if batch.do_classifier_free_guidance: + model_video_uncond, model_video_text = model_video.chunk(2) + model_audio_uncond, model_audio_text = model_audio.chunk(2) + model_video = model_video_uncond + ( + batch.guidance_scale * (model_video_text - model_video_uncond) + ) + model_audio = model_audio_uncond + ( + batch.guidance_scale * (model_audio_text - model_audio_uncond) + ) + + ctx.latents = self.scheduler.step( + model_video, step.t_device, ctx.latents, return_dict=False + )[0] + ctx.audio_latents = ctx.audio_scheduler.step( + model_audio, step.t_device, ctx.audio_latents, return_dict=False + )[0] + ctx.latents = self.post_forward_for_ti2v_task( + batch, server_args, ctx.reserved_frames_mask, ctx.latents, ctx.z + ) + return + + encoder_hidden_states = batch.prompt_embeds[0] + audio_encoder_hidden_states = batch.audio_prompt_embeds[0] + encoder_attention_mask = prompt_attention_mask + with set_forward_context( + current_timestep=step.step_index, attn_metadata=step.attn_metadata + ): + v_pos, a_v_pos = step.current_model( + **build_model_kwargs( + encoder_hidden_states=encoder_hidden_states, + audio_encoder_hidden_states=audio_encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + ) + ) + + if stage1_guider_params is not None or batch.do_classifier_free_guidance: + v_neg, a_v_neg = step.current_model( + **build_model_kwargs( + encoder_hidden_states=batch.negative_prompt_embeds[0], + audio_encoder_hidden_states=batch.negative_audio_prompt_embeds[ + 0 + ], + encoder_attention_mask=self._get_ltx_prompt_attention_mask( + batch, + is_ltx23_variant=( + ctx.is_ltx23_variant + and not ctx.use_ltx23_legacy_one_stage + ), + negative=True, + ), + ) + ) + else: + v_neg = None + a_v_neg = None + + v_pos = v_pos.float() + a_v_pos = a_v_pos.float() + if v_neg is not None: + v_neg = v_neg.float() + if a_v_neg is not None: + a_v_neg = a_v_neg.float() + + sigma_val = float(sigma.item()) + video_sigma_for_x0: float | torch.Tensor = sigma_val + if ctx.denoise_mask is not None: + video_sigma_for_x0 = sigma.to( + device=ctx.latents.device, dtype=torch.float32 + ) * ctx.denoise_mask.squeeze(-1) + + denoised_video = self._ltx2_velocity_to_x0( + ctx.latents, v_pos, video_sigma_for_x0 + ) + denoised_audio = self._ltx2_velocity_to_x0( + ctx.audio_latents, a_v_pos, sigma_val + ) + denoised_video_neg = None + denoised_audio_neg = None + denoised_video_perturbed = None + denoised_audio_perturbed = None + denoised_video_modality = None + denoised_audio_modality = None + + if ( + (stage1_guider_params is not None or batch.do_classifier_free_guidance) + and v_neg is not None + and a_v_neg is not None + ): + denoised_video_neg = self._ltx2_velocity_to_x0( + ctx.latents, v_neg, video_sigma_for_x0 + ) + denoised_audio_neg = self._ltx2_velocity_to_x0( + ctx.audio_latents, a_v_neg, sigma_val + ) + if stage1_guider_params is not None: + video_skip = self._ltx2_should_skip_step( + step.step_index, int(stage1_guider_params["video_skip_step"]) + ) + audio_skip = self._ltx2_should_skip_step( + step.step_index, int(stage1_guider_params["audio_skip_step"]) + ) + + need_perturbed = ( + float(stage1_guider_params["video_stg_scale"]) != 0.0 + or float(stage1_guider_params["audio_stg_scale"]) != 0.0 + ) + if need_perturbed: + with set_forward_context( + current_timestep=step.step_index, attn_metadata=step.attn_metadata + ): + v_ptb, a_v_ptb = step.current_model( + **build_model_kwargs( + encoder_hidden_states=encoder_hidden_states, + audio_encoder_hidden_states=audio_encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + skip_video_self_attn_blocks=tuple( + stage1_guider_params["video_stg_blocks"] + ), + skip_audio_self_attn_blocks=tuple( + stage1_guider_params["audio_stg_blocks"] + ), + ) + ) + denoised_video_perturbed = self._ltx2_velocity_to_x0( + ctx.latents, v_ptb.float(), video_sigma_for_x0 + ) + denoised_audio_perturbed = self._ltx2_velocity_to_x0( + ctx.audio_latents, a_v_ptb.float(), sigma_val + ) + + need_modality = ( + float(stage1_guider_params["video_modality_scale"]) != 1.0 + or float(stage1_guider_params["audio_modality_scale"]) != 1.0 + ) + if need_modality: + with set_forward_context( + current_timestep=step.step_index, attn_metadata=step.attn_metadata + ): + v_mod, a_v_mod = step.current_model( + **build_model_kwargs( + encoder_hidden_states=encoder_hidden_states, + audio_encoder_hidden_states=audio_encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + disable_a2v_cross_attn=True, + disable_v2a_cross_attn=True, + ) + ) + denoised_video_modality = self._ltx2_velocity_to_x0( + ctx.latents, v_mod.float(), video_sigma_for_x0 + ) + denoised_audio_modality = self._ltx2_velocity_to_x0( + ctx.audio_latents, a_v_mod.float(), sigma_val + ) + + if not video_skip: + denoised_video = self._ltx2_calculate_guided_x0( + cond=denoised_video, + uncond_text=( + denoised_video_neg + if denoised_video_neg is not None + else denoised_video + ), + uncond_perturbed=( + denoised_video_perturbed + if denoised_video_perturbed is not None + else 0.0 + ), + uncond_modality=( + denoised_video_modality + if denoised_video_modality is not None + else 0.0 + ), + cfg_scale=float(stage1_guider_params["video_cfg_scale"]), + stg_scale=float(stage1_guider_params["video_stg_scale"]), + rescale_scale=float(stage1_guider_params["video_rescale_scale"]), + modality_scale=float(stage1_guider_params["video_modality_scale"]), + ) + if not audio_skip: + denoised_audio = self._ltx2_calculate_guided_x0( + cond=denoised_audio, + uncond_text=( + denoised_audio_neg + if denoised_audio_neg is not None + else denoised_audio + ), + uncond_perturbed=( + denoised_audio_perturbed + if denoised_audio_perturbed is not None + else 0.0 + ), + uncond_modality=( + denoised_audio_modality + if denoised_audio_modality is not None + else 0.0 + ), + cfg_scale=float(stage1_guider_params["audio_cfg_scale"]), + stg_scale=float(stage1_guider_params["audio_stg_scale"]), + rescale_scale=float(stage1_guider_params["audio_rescale_scale"]), + modality_scale=float(stage1_guider_params["audio_modality_scale"]), + ) + elif ( + batch.do_classifier_free_guidance + and denoised_video_neg is not None + and denoised_audio_neg is not None + ): + denoised_video = denoised_video + (batch.guidance_scale - 1.0) * ( + denoised_video - denoised_video_neg + ) + denoised_audio = denoised_audio + (batch.guidance_scale - 1.0) * ( + denoised_audio - denoised_audio_neg + ) + + if ctx.denoise_mask is not None and ctx.clean_latent is not None: + denoised_video = ( + denoised_video * ctx.denoise_mask + + ctx.clean_latent.float() * (1.0 - ctx.denoise_mask) + ) + + # 6. Convert x0 predictions back to velocity and update both latent streams. + if sigma_val == 0.0: + v_video = torch.zeros_like(denoised_video) + v_audio = torch.zeros_like(denoised_audio) + else: + v_video = ((ctx.latents.float() - denoised_video.float()) / sigma_val).to( + ctx.latents.dtype + ) + v_audio = ( + (ctx.audio_latents.float() - denoised_audio.float()) / sigma_val + ).to(ctx.audio_latents.dtype) + + ctx.latents = (ctx.latents.float() + v_video.float() * dt).to( + dtype=ctx.latents.dtype + ) + ctx.audio_latents = (ctx.audio_latents.float() + v_audio.float() * dt).to( + dtype=ctx.audio_latents.dtype + ) + ctx.latents = self.post_forward_for_ti2v_task( + batch, server_args, ctx.reserved_frames_mask, ctx.latents, ctx.z + ) + + def _record_trajectory( + self, + ctx: LTX2DenoisingContext, + step: DenoisingStepState, + batch: Req, + server_args: ServerArgs, + ) -> None: + """Record audio trajectory alongside the base video trajectory.""" + super()._record_trajectory(ctx, step, batch, server_args) + if batch.return_trajectory_latents and ctx.audio_latents is not None: + ctx.trajectory_audio_latents.append(ctx.audio_latents) + + def _finalize_denoising_loop( + self, ctx: LTX2DenoisingContext, batch: Req, server_args: ServerArgs + ) -> None: + """Expose audio latents before delegating to AV-aware postprocessing.""" + batch.audio_latents = ctx.audio_latents + self._post_denoising_loop( + batch=batch, + latents=ctx.latents, + trajectory_latents=ctx.trajectory_latents, + trajectory_timesteps=ctx.trajectory_timesteps, + trajectory_audio_latents=ctx.trajectory_audio_latents, + server_args=server_args, + is_warmup=ctx.is_warmup, + ) + + def _post_denoising_loop( + self, + batch: Req, + latents: torch.Tensor, + trajectory_latents: list, + trajectory_timesteps: list, + server_args: ServerArgs, + trajectory_audio_latents: list | None = None, + is_warmup: bool = False, + *args, + **kwargs, + ): + """Trim SP token padding before delegating to the base finalizer.""" + if trajectory_audio_latents: + batch.trajectory_audio_latents = torch.stack( + trajectory_audio_latents, dim=1 + ).cpu() + latents = self._truncate_sp_padded_token_latents(batch, latents) + super()._post_denoising_loop( + batch=batch, + latents=latents, + trajectory_latents=trajectory_latents, + trajectory_timesteps=trajectory_timesteps, + server_args=server_args, + is_warmup=is_warmup, + ) + + def _get_prompt_embeds_validator(self, batch: Req): + """Allow either tensor or list prompt embeddings for LTX-2 prompts.""" + del batch + return lambda x: V.is_tensor(x) or V.list_not_empty(x) + + def _get_negative_prompt_embeds_validator(self, batch: Req): + """Allow either tensor or list negative prompt embeddings for LTX-2 CFG.""" + return ( + lambda x: (not batch.do_classifier_free_guidance) + or V.is_tensor(x) + or V.list_not_empty(x) + ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/__init__.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/__init__.py new file mode 100644 index 000000000000..7ac2d6cbec57 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/__init__.py @@ -0,0 +1 @@ +"""Model-specific helpers and stages for diffusion pipeline components.""" diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/wan_ti2v.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/wan_ti2v.py new file mode 100644 index 000000000000..b26d00fa6852 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/wan_ti2v.py @@ -0,0 +1,174 @@ +"""WAN TI2V-specific helpers shared by the generic denoising stage.""" + +import math + +import torch +from einops import rearrange + +from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType +from sglang.multimodal_gen.configs.pipeline_configs.wan import ( + Wan2_2_TI2V_5B_Config, +) +from sglang.multimodal_gen.runtime.distributed import ( + get_local_torch_device, + get_sp_parallel_rank, + get_sp_world_size, +) +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.utils import masks_like + + +def should_apply_wan_ti2v(batch: Req, server_args: ServerArgs) -> bool: + """Return whether the request should use the Wan2.2 TI2V latent path.""" + + return bool( + server_args.pipeline_config.task_type == ModelTaskType.TI2V + and batch.condition_image is not None + and type(server_args.pipeline_config) is Wan2_2_TI2V_5B_Config + ) + + +def prepare_wan_ti2v_latents( + vae: object, + latents: torch.Tensor, + target_dtype: torch.dtype, + batch: Req, + server_args: ServerArgs, +) -> tuple[int, torch.Tensor, list[torch.Tensor]]: + """Encode the conditioning image and splice it into Wan TI2V latents.""" + + # Wan2.2 TI2V directly replaces the first frame of the latent with + # the image latent instead of appending along the channel dim. + assert batch.image_latent is None, "TI2V task should not have image latents" + assert vae is not None, "VAE is not provided for TI2V task" + + vae = vae.to(batch.condition_image.device) + z = vae.encode(batch.condition_image).mean.float() + if getattr(vae, "device", None) != "cpu" and server_args.vae_cpu_offload: + vae = vae.to("cpu") + + if hasattr(vae, "shift_factor") and vae.shift_factor is not None: + if isinstance(vae.shift_factor, torch.Tensor): + z -= vae.shift_factor.to(z.device, z.dtype) + else: + z -= vae.shift_factor + + if isinstance(vae.scaling_factor, torch.Tensor): + z = z * vae.scaling_factor.to(z.device, z.dtype) + else: + z = z * vae.scaling_factor + + latent_model_input = latents.to(target_dtype) + assert latent_model_input.ndim == 5 + + latent_for_mask = latent_model_input.squeeze(0) + _, reserved_frames_masks = masks_like([latent_for_mask], zero=True) + reserved_frames_mask = reserved_frames_masks[0].unsqueeze(0) + + latents = ( + 1.0 - reserved_frames_mask + ) * z + reserved_frames_mask * latent_model_input + assert latents.ndim == 5 + batch.latents = latents.to(get_local_torch_device()) + + num_frames = batch.num_frames + temporal_scale = ( + server_args.pipeline_config.vae_config.arch_config.scale_factor_temporal + ) + spatial_scale = ( + server_args.pipeline_config.vae_config.arch_config.scale_factor_spatial + ) + patch_size = server_args.pipeline_config.dit_config.arch_config.patch_size + seq_len = ( + ((num_frames - 1) // temporal_scale + 1) + * (batch.height // spatial_scale) + * (batch.width // spatial_scale) + // (patch_size[1] * patch_size[2]) + ) + seq_len = int(math.ceil(seq_len / get_sp_world_size())) * get_sp_world_size() + return seq_len, z, reserved_frames_masks + + +def prepare_wan_ti2v_sp_inputs( + z: torch.Tensor | None, + reserved_frames_masks: list[torch.Tensor] | None, + batch: Req, +) -> tuple[torch.Tensor | None, torch.Tensor | None]: + """Shard Wan TI2V image-conditioning state to match SP-sharded video latents.""" + + rank_in_sp_group = get_sp_parallel_rank() + sp_world_size = get_sp_world_size() + + if getattr(batch, "did_sp_shard_latents", False): + if z is not None and z.shape[2] == 1: + z_sp = z if rank_in_sp_group == 0 else None + else: + z_sp = z + + if reserved_frames_masks is not None: + reserved_frames_mask = reserved_frames_masks[0] + time_dim = reserved_frames_mask.shape[1] + if time_dim > 0 and time_dim % sp_world_size == 0: + reserved_frames_mask_sp_tensor = rearrange( + reserved_frames_mask, + "c (n t) h w -> c n t h w", + n=sp_world_size, + ).contiguous() + reserved_frames_mask_sp = reserved_frames_mask_sp_tensor[ + :, rank_in_sp_group, :, :, : + ] + else: + reserved_frames_mask_sp = reserved_frames_mask + else: + reserved_frames_mask_sp = None + else: + z_sp = z + reserved_frames_mask_sp = ( + reserved_frames_masks[0] if reserved_frames_masks is not None else None + ) + + return reserved_frames_mask_sp, z_sp + + +def expand_wan_ti2v_timestep( + batch: Req, + t_device: torch.Tensor, + target_dtype: torch.dtype, + seq_len: int, + reserved_frames_mask: torch.Tensor | None, +) -> torch.Tensor: + """Expand the timestep tensor for Wan TI2V's first-frame masking semantics.""" + + batch_size = batch.raw_latent_shape[0] + t_device_rounded = t_device.to(target_dtype) + + local_seq_len = seq_len + if get_sp_world_size() > 1 and getattr(batch, "did_sp_shard_latents", False): + local_seq_len = seq_len // get_sp_world_size() + + if get_sp_parallel_rank() == 0 and reserved_frames_mask is not None: + temp_ts = (reserved_frames_mask[0][:, ::2, ::2] * t_device_rounded).flatten() + temp_ts = torch.cat( + [ + temp_ts, + temp_ts.new_ones(local_seq_len - temp_ts.size(0)) * t_device_rounded, + ] + ) + return temp_ts.unsqueeze(0).repeat(batch_size, 1) + + return t_device.repeat(batch_size, local_seq_len) + + +def blend_wan_ti2v_latents( + latents: torch.Tensor, + reserved_frames_mask: torch.Tensor | None, + z: torch.Tensor | None, +) -> torch.Tensor: + """Restore Wan TI2V's conditioned first frame after each denoising step.""" + + if z is None or reserved_frames_mask is None: + return latents + return ( + 1.0 - reserved_frames_mask.unsqueeze(0) + ) * z + reserved_frames_mask.unsqueeze(0) * latents diff --git a/scripts/ci/utils/diffusion/comparison_configs.json b/scripts/ci/utils/diffusion/comparison_configs.json index b1b766591c4f..db51d53fa518 100644 --- a/scripts/ci/utils/diffusion/comparison_configs.json +++ b/scripts/ci/utils/diffusion/comparison_configs.json @@ -135,6 +135,23 @@ } } }, + { + "id": "ltx2.3_twostage_t2v_2gpus", + "model": "Lightricks/LTX-2.3", + "task": "text-to-video", + "prompt": "A cat and a dog baking a cake together in a kitchen.", + "width": 768, + "height": 512, + "num_frames": 121, + "seed": 42, + "num_gpus": 2, + "frameworks": { + "sglang": { + "serve_args": "--enable-torch-compile --warmup --enable-cfg-parallel --pipeline-class-name LTX2TwoStagePipeline", + "extra_env": {} + } + } + }, { "id": "wan22_i2v_a14b_720p", "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers", From 1e9eecfa36ce899d27cb2d794a083049ba82d937 Mon Sep 17 00:00:00 2001 From: Chandrakant Khandelwal Date: Mon, 13 Apr 2026 11:15:48 +0530 Subject: [PATCH 19/52] [Intel GPU] Enable sgl-kernel-xpu fused_experts MoE kernel path for GPT-OSS bf16 models. (#22417) --- python/sglang/srt/layers/quantization/unquant.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index 7339640c4abc..3cc28e2aef14 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/python/sglang/srt/layers/quantization/unquant.py @@ -593,6 +593,8 @@ def forward_xpu( b1=getattr(layer, "w13_weight_bias", None), b2=getattr(layer, "w2_weight_bias", None), activation=moe_runner_config.activation, + gemm1_alpha=moe_runner_config.gemm1_alpha, + gemm1_limit=moe_runner_config.gemm1_clamp_limit, ) return StandardCombineInput(hidden_states=output) else: From 4a746ea462882106c85724f8228eae1079414f05 Mon Sep 17 00:00:00 2001 From: Thomas Wang Date: Mon, 13 Apr 2026 15:01:37 +0800 Subject: [PATCH 20/52] [AMD] Remove aiter hotfixes in Dockerfile covered by aiter v0.1.12.post1 (#22657) --- docker/rocm.Dockerfile | 19 ------------------- 1 file changed, 19 deletions(-) diff --git a/docker/rocm.Dockerfile b/docker/rocm.Dockerfile index 23fdee68e24c..abc2f8615b58 100644 --- a/docker/rocm.Dockerfile +++ b/docker/rocm.Dockerfile @@ -180,27 +180,8 @@ RUN git clone ${AITER_REPO} \ && git submodule update --init --recursive \ && pip install -r requirements.txt -# Hot patches for AITER in v0.1.10.post3 -# This is for ROCm 7.2 only, because of the image rebase from vllm -# to rocm/pytorch. -RUN set -eux; \ - case "${GPU_ARCH}" in \ - *rocm720*) \ - echo "ROCm 7.2 flavor detected from GPU_ARCH=${GPU_ARCH}"; \ - cd aiter \ - && sed -i '459 s/if.*:/if False:/' aiter/ops/triton/attention/pa_mqa_logits.py; \ - ;; \ - *) \ - echo "Not rocm720 (GPU_ARCH=${GPU_ARCH}), skip patch"; \ - ;; \ - esac -# [WA] from kk-huang -# add sed -i '/c1 = torch.empty((M, D, S1 + S3) for aiter triton gemm config issue -# the corresponding pr is https://github.com/ROCm/aiter/pull/2173 -# it will be removed when server launched issue is fixed by aiter RUN cd aiter \ && echo "[AITER] GPU_ARCH=${GPU_ARCH}" \ - && sed -i '/c1 = torch.empty((M, D, S1 + S3), dtype=dtype, device=x.device)/i\ config = dict(config)' aiter/ops/triton/gemm/fused/fused_gemm_afp4wfp4_split_cat.py \ && if [ "$BUILD_AITER_ALL" = "1" ] && [ "$BUILD_LLVM" = "1" ]; then \ sh -c "HIP_CLANG_PATH=/sgl-workspace/llvm-project/build/bin/ PREBUILD_KERNELS=1 GPU_ARCHS=$GPU_ARCH_LIST python setup.py build_ext --inplace" \ && sh -c "HIP_CLANG_PATH=/sgl-workspace/llvm-project/build/bin/ GPU_ARCHS=$GPU_ARCH_LIST pip install -e ."; \ From 9e6d1c066ec3445b534eef6ccdcd752f1bddae1f Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Mon, 13 Apr 2026 00:37:29 -0700 Subject: [PATCH 21/52] TestStreamingSessionAbortLeakRepro: inherit stdout/stderr instead of tempfile (#22668) --- .../sessions/test_streaming_session.py | 43 +------------------ 1 file changed, 2 insertions(+), 41 deletions(-) diff --git a/test/registered/sessions/test_streaming_session.py b/test/registered/sessions/test_streaming_session.py index 59bb76981971..8c8f82702b4e 100644 --- a/test/registered/sessions/test_streaming_session.py +++ b/test/registered/sessions/test_streaming_session.py @@ -11,11 +11,8 @@ import asyncio import json -import os -import tempfile import time import unittest -from pathlib import Path from typing import Any, Optional import aiohttp @@ -281,13 +278,6 @@ async def _abort_repro_generate( return data -def _read_tail(path: str, num_lines: int = 120) -> str: - if not path or not os.path.exists(path): - return "" - lines = Path(path).read_text(errors="replace").splitlines() - return "\n".join(lines[-num_lines:]) - - async def _abort_repro_run_all(base_url: str, tokenizer: Any) -> None: timeout = aiohttp.ClientTimeout(total=300) async with aiohttp.ClientSession(timeout=timeout) as http: @@ -683,20 +673,6 @@ class TestStreamingSessionAbortLeakRepro(CustomTestCase): def setUpClass(cls): cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST cls.base_url = DEFAULT_URL_FOR_TEST - cls.stdout = tempfile.NamedTemporaryFile( - prefix="streaming-session-abort-repro.", - suffix=".stdout.log", - delete=False, - mode="w+", - encoding="utf-8", - ) - cls.stderr = tempfile.NamedTemporaryFile( - prefix="streaming-session-abort-repro.", - suffix=".stderr.log", - delete=False, - mode="w+", - encoding="utf-8", - ) cls.process = popen_launch_server( cls.model, cls.base_url, @@ -714,18 +690,12 @@ def setUpClass(cls): "--log-level", "info", ], - return_stdout_stderr=(cls.stdout, cls.stderr), ) cls.tokenizer = get_tokenizer(cls.model) @classmethod def tearDownClass(cls): kill_process_tree(cls.process.pid) - for handle in (cls.stdout, cls.stderr): - path = handle.name - handle.close() - if os.path.exists(path): - os.remove(path) def test_abort_heavy_chunked_prefill_does_not_leak(self) -> None: requests.post(self.base_url + "/flush_cache") @@ -747,23 +717,14 @@ def test_abort_heavy_chunked_prefill_does_not_leak(self) -> None: time.sleep(5) self.assertIsNone( self.process.poll(), - "Server crashed during abort-heavy streaming session repro.\n" - f"---- stderr tail ----\n{_read_tail(self.stderr.name)}", + "Server crashed during abort-heavy streaming session repro.", ) health = requests.get(self.base_url + "/health", timeout=10) self.assertEqual( health.status_code, 200, - "Server unhealthy after abort-heavy streaming session cleanup.\n" - f"---- stderr tail ----\n{_read_tail(self.stderr.name)}", - ) - - stderr_tail = _read_tail(self.stderr.name) - self.assertNotIn( - "token_to_kv_pool_allocator memory leak detected", - stderr_tail, - stderr_tail, + "Server unhealthy after abort-heavy streaming session cleanup.", ) From 4df60434d7e15284f56f72b881db7c218eac7563 Mon Sep 17 00:00:00 2001 From: Xiaole Guo <119391905+gxlvera@users.noreply.github.com> Date: Mon, 13 Apr 2026 01:07:06 -0700 Subject: [PATCH 22/52] [diffusion] model: support stable-diffusion-3-medium-diffusers (#19225) Co-authored-by: zhaochenyang20 Co-authored-by: Kangrui Du Co-authored-by: Xiaole Guo --- .../configs/models/dits/__init__.py | 4 + .../configs/models/dits/stablediffusion3.py | 37 ++++ .../configs/models/vaes/__init__.py | 4 + .../configs/models/vaes/stablediffusion3.py | 72 +++++++ .../configs/pipeline_configs/__init__.py | 4 + .../configs/pipeline_configs/base.py | 27 +++ .../configs/pipeline_configs/flux.py | 15 ++ .../pipeline_configs/stablediffusion3.py | 202 ++++++++++++++++++ .../configs/sample/stablediffusion3.py | 20 ++ python/sglang/multimodal_gen/registry.py | 26 +++ .../component_loaders/text_encoder_loader.py | 47 ++-- .../loader/component_loaders/vae_loader.py | 7 + .../multimodal_gen/runtime/loader/utils.py | 4 +- .../runtime/models/dits/stablediffusion3.py | 180 ++++++++++++++++ .../runtime/models/encoders/clip.py | 44 +++- .../runtime/pipelines/stable_diffusion_3.py | 111 ++++++++++ .../pipelines_core/stages/text_encoding.py | 25 ++- 17 files changed, 800 insertions(+), 29 deletions(-) create mode 100644 python/sglang/multimodal_gen/configs/models/dits/stablediffusion3.py create mode 100644 python/sglang/multimodal_gen/configs/models/vaes/stablediffusion3.py create mode 100644 python/sglang/multimodal_gen/configs/pipeline_configs/stablediffusion3.py create mode 100644 python/sglang/multimodal_gen/configs/sample/stablediffusion3.py create mode 100644 python/sglang/multimodal_gen/runtime/models/dits/stablediffusion3.py create mode 100644 python/sglang/multimodal_gen/runtime/pipelines/stable_diffusion_3.py diff --git a/python/sglang/multimodal_gen/configs/models/dits/__init__.py b/python/sglang/multimodal_gen/configs/models/dits/__init__.py index ba114c18fa89..63ac893ddd42 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/__init__.py +++ b/python/sglang/multimodal_gen/configs/models/dits/__init__.py @@ -5,6 +5,9 @@ from sglang.multimodal_gen.configs.models.dits.hunyuanvideo import HunyuanVideoConfig from sglang.multimodal_gen.configs.models.dits.mova_audio import MOVAAudioConfig from sglang.multimodal_gen.configs.models.dits.mova_video import MOVAVideoConfig +from sglang.multimodal_gen.configs.models.dits.stablediffusion3 import ( + StableDiffusion3TransformerConfig, +) from sglang.multimodal_gen.configs.models.dits.wanvideo import WanVideoConfig __all__ = [ @@ -14,4 +17,5 @@ "Hunyuan3DDiTConfig", "MOVAAudioConfig", "MOVAVideoConfig", + "StableDiffusion3TransformerConfig", ] diff --git a/python/sglang/multimodal_gen/configs/models/dits/stablediffusion3.py b/python/sglang/multimodal_gen/configs/models/dits/stablediffusion3.py new file mode 100644 index 000000000000..cbe3141b7092 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/models/dits/stablediffusion3.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +"""StableDiffusion3 Transformer model configuration""" + +from dataclasses import dataclass, field + +from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig + + +@dataclass +class StableDiffusion3TransformerArchConfig(DiTArchConfig): + """Architecture configuration for StableDiffusion3 Transformer, applicable to SD3-medium, SD3.5-medium, SD3.5-large.""" + + sample_size: int = 128 + patch_size: int = 2 + in_channels: int = 16 + out_channels: int = 16 + num_layers: int = 18 + attention_head_dim: int = 64 + num_attention_heads: int = 18 + cross_attention_dim: int = 4096 + joint_attention_dim: int = 4096 + caption_projection_dim: int = 1152 + pooled_projection_dim: int = 2048 + pos_embed_max_size: int = 96 + dual_attention_layers: tuple[int, ...] = () + qk_norm: str | None = None + + _class_name: str = "SD3Transformer2DModel" + + +@dataclass +class StableDiffusion3TransformerConfig(DiTConfig): + """Configuration for StableDiffusion3 Transformer model.""" + + arch_config: StableDiffusion3TransformerArchConfig = field( + default_factory=StableDiffusion3TransformerArchConfig + ) diff --git a/python/sglang/multimodal_gen/configs/models/vaes/__init__.py b/python/sglang/multimodal_gen/configs/models/vaes/__init__.py index 8b9dd60fe5ef..3438b1b8937f 100644 --- a/python/sglang/multimodal_gen/configs/models/vaes/__init__.py +++ b/python/sglang/multimodal_gen/configs/models/vaes/__init__.py @@ -3,11 +3,15 @@ from sglang.multimodal_gen.configs.models.vaes.dac import DacVAEConfig from sglang.multimodal_gen.configs.models.vaes.hunyuan3d import Hunyuan3DVAEConfig from sglang.multimodal_gen.configs.models.vaes.hunyuanvae import HunyuanVAEConfig +from sglang.multimodal_gen.configs.models.vaes.stablediffusion3 import ( + StableDiffusion3VAEConfig, +) from sglang.multimodal_gen.configs.models.vaes.wanvae import WanVAEConfig __all__ = [ "DacVAEConfig", "HunyuanVAEConfig", + "StableDiffusion3VAEConfig", "WanVAEConfig", "Hunyuan3DVAEConfig", ] diff --git a/python/sglang/multimodal_gen/configs/models/vaes/stablediffusion3.py b/python/sglang/multimodal_gen/configs/models/vaes/stablediffusion3.py new file mode 100644 index 000000000000..fe6351b2e639 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/models/vaes/stablediffusion3.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +"""StableDiffusion3 VAE configuration.""" + +from dataclasses import dataclass, field + +from sglang.multimodal_gen.configs.models.vaes.base import VAEArchConfig, VAEConfig + + +@dataclass +class StableDiffusion3VAEArchConfig(VAEArchConfig): + """Architecture configuration for StableDiffusion3 VAE.""" + + scaling_factor: float = 1.5305 + shift_factor: float = 0.0609 + + spatial_compression_ratio: int = 8 + temporal_compression_ratio: int = 1 + + in_channels: int = 3 + out_channels: int = 3 + latent_channels: int = 16 + sample_size: int = 128 + + block_out_channels: tuple[int, ...] = (128, 256, 512, 512) + layers_per_block: int = 2 + act_fn: str = "silu" + norm_num_groups: int = 32 + + down_block_types: tuple[str, ...] = ( + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D", + ) + up_block_types: tuple[str, ...] = ( + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D", + ) + + attention_head_dim: int = 8 + mid_block_add_attention: bool = True + use_quant_conv: bool = False + use_post_quant_conv: bool = False + + +@dataclass +class StableDiffusion3VAEConfig(VAEConfig): + """Configuration for StableDiffusion3 VAE.""" + + arch_config: StableDiffusion3VAEArchConfig = field( + default_factory=StableDiffusion3VAEArchConfig + ) + + tile_sample_min_height: int = 512 + tile_sample_min_width: int = 512 + tile_sample_min_num_frames: int = 1 + tile_sample_stride_height: int = 448 + tile_sample_stride_width: int = 448 + tile_sample_stride_num_frames: int = 1 + + use_tiling: bool = True + use_temporal_tiling: bool = False + use_parallel_tiling: bool = True + use_temporal_scaling_frames: bool = False + + def __post_init__(self) -> None: + """Post initialization for SD3 VAE specific setup.""" + super().__post_init__() + self.update_model_arch({"_class_name": "AutoencoderKL"}) + self.blend_num_frames = 0 diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py index e2e7c36dc6a1..005bc862c558 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py @@ -30,6 +30,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import LTX2PipelineConfig from sglang.multimodal_gen.configs.pipeline_configs.mova import MOVAPipelineConfig from sglang.multimodal_gen.configs.pipeline_configs.sana import SanaPipelineConfig +from sglang.multimodal_gen.configs.pipeline_configs.stablediffusion3 import ( + StableDiffusion3PipelineConfig, +) from sglang.multimodal_gen.configs.pipeline_configs.wan import ( SelfForcingWanT2V480PConfig, WanI2V480PConfig, @@ -55,6 +58,7 @@ "SanaPipelineConfig", "SlidingTileAttnConfig", "MOVAPipelineConfig", + "StableDiffusion3PipelineConfig", "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig", diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index 6406af765ab6..46086cba7dd3 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py @@ -459,6 +459,33 @@ def shard_latents_for_sp(self, batch, latents): sharded_tensor = sharded_tensor[:, :, rank_in_sp_group, :, :, :] return sharded_tensor, True + def get_text_encoder_attention_mask( + self, text_inputs: dict, encoder_index: int + ) -> "torch.Tensor | None": + """Return the attention mask for the given text encoder. + + Override to suppress (return None) or modify the mask per model. + """ + return text_inputs.get("attention_mask") + + def get_text_encoder_pooler_output( + self, outputs: "BaseEncoderOutput", encoder_index: int + ) -> "torch.Tensor | None": + """Return the pooler output for the given text encoder, or None to skip. + + Override for models that need pooled embeddings (e.g. FLUX v1, SD3). + """ + return None + + def select_vae_weight_files( + self, + safetensors_list: list[str], + component_model_path: str, + component_name: str, + vae_precision: str, + ) -> list[str]: + return safetensors_list + def get_pos_prompt_embeds(self, batch): return batch.prompt_embeds diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py b/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py index 76c7ddea6ab3..38829098ff8a 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py @@ -81,6 +81,13 @@ class FluxPipelineConfig(ImagePipelineConfig): ] ) + def get_text_encoder_attention_mask(self, text_inputs, encoder_index): + # Flux v1 does not use attention masks for text encoders. + return None + + def get_text_encoder_pooler_output(self, outputs, encoder_index): + return outputs.pooler_output + def prepare_sigmas(self, sigmas, num_inference_steps): return self._prepare_sigmas(sigmas, num_inference_steps) @@ -391,6 +398,14 @@ class Flux2PipelineConfig(FluxPipelineConfig): ] ) + def get_text_encoder_attention_mask(self, text_inputs, encoder_index): + # Flux2 uses standard attention masks (unlike Flux v1). + return text_inputs.get("attention_mask") + + def get_text_encoder_pooler_output(self, outputs, encoder_index): + # Flux2 does not use pooler output. + return None + def tokenize_prompt(self, prompts: list[str], tokenizer, tok_kwargs) -> dict: messages = build_flux2_text_messages(prompts) inputs = tokenizer.apply_chat_template( diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/stablediffusion3.py b/python/sglang/multimodal_gen/configs/pipeline_configs/stablediffusion3.py new file mode 100644 index 000000000000..beea2109717e --- /dev/null +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/stablediffusion3.py @@ -0,0 +1,202 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Stable Diffusion 3 pipeline configuration.""" + +import os +from dataclasses import dataclass, field +from typing import Callable + +import torch + +from sglang.multimodal_gen.configs.models import DiTConfig, EncoderConfig, VAEConfig +from sglang.multimodal_gen.configs.models.dits import StableDiffusion3TransformerConfig +from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput +from sglang.multimodal_gen.configs.models.encoders.base import TextEncoderArchConfig +from sglang.multimodal_gen.configs.models.encoders.clip import ( + CLIPTextArchConfig, + CLIPTextConfig, +) +from sglang.multimodal_gen.configs.models.encoders.t5 import ( + T5ArchConfig, + T5Config, +) +from sglang.multimodal_gen.configs.models.vaes.stablediffusion3 import ( + StableDiffusion3VAEConfig, +) +from sglang.multimodal_gen.configs.pipeline_configs.base import ( + ModelTaskType, + SpatialImagePipelineConfig, +) + + +def sd3_clip_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tensor: + """Extract pre-final hidden state for SD3 CLIP encoders.""" + if outputs.hidden_states is None: + raise ValueError( + "SD3 CLIP postprocessing requires hidden_states from encoder output." + ) + return outputs.hidden_states[-2] + + +def t5_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tensor: + return outputs.last_hidden_state + + +def select_sd3_vae_weight_files( + safetensors_list: list[str], + component_model_path: str, + component_name: str, + vae_precision: str, +) -> list[str]: + """Select SD3 VAE checkpoint file candidates with minimal policy.""" + if component_name not in ("vae", "video_vae"): + return safetensors_list + + base_name = "diffusion_pytorch_model" + if vae_precision == "fp16": + fp16_path = os.path.join(component_model_path, f"{base_name}.fp16.safetensors") + if os.path.exists(fp16_path): + return [fp16_path] + + full_path = os.path.join(component_model_path, f"{base_name}.safetensors") + if os.path.exists(full_path): + return [full_path] + return safetensors_list + + +@dataclass +class SD3CLIPTextArchConfig(CLIPTextArchConfig): + def __post_init__(self) -> None: + super().__post_init__() + self.tokenizer_kwargs.update( + { + "max_length": self.text_len, + "padding": "max_length", + } + ) + + +@dataclass +class SD3CLIPTextConfig(CLIPTextConfig): + arch_config: TextEncoderArchConfig = field(default_factory=SD3CLIPTextArchConfig) + + +@dataclass +class SD3T5ArchConfig(T5ArchConfig): + def __post_init__(self) -> None: + super().__post_init__() + self.tokenizer_kwargs.update({"max_length": 256}) + + +@dataclass +class SD3T5Config(T5Config): + arch_config: TextEncoderArchConfig = field(default_factory=SD3T5ArchConfig) + + +@dataclass +class StableDiffusion3PipelineConfig(SpatialImagePipelineConfig): + """Configuration for SD3 image generation pipeline. + + This config intentionally relies on SD3-specific encoder configs to provide + tokenizer kwargs, instead of stage-level tokenizer overrides. + """ + + task_type: ModelTaskType = ModelTaskType.T2I + + dit_config: DiTConfig = field(default_factory=StableDiffusion3TransformerConfig) + vae_config: VAEConfig = field(default_factory=StableDiffusion3VAEConfig) + + text_encoder_configs: tuple[EncoderConfig, ...] = field( + default_factory=lambda: ( + SD3CLIPTextConfig(), + SD3CLIPTextConfig(), + SD3T5Config(), + ) + ) + + text_encoder_precisions: tuple[str, ...] = field( + default_factory=lambda: ("fp16", "fp16", "fp32") + ) + + preprocess_text_funcs: tuple[Callable[[str], str] | None, ...] = field( + default_factory=lambda: ( + None, + None, + None, + ) + ) + + postprocess_text_funcs: tuple[ + Callable[[BaseEncoderOutput, dict], torch.Tensor], ... + ] = field( + default_factory=lambda: ( + sd3_clip_postprocess_text, + sd3_clip_postprocess_text, + t5_postprocess_text, + ) + ) + + should_use_guidance: bool = False + guidance_scale: float = 7.0 + + def __post_init__(self) -> None: + configs = list(self.text_encoder_configs) + configs[0].update_model_arch({"_class_name": "CLIPTextModelWithProjection"}) + configs[1].update_model_arch({"_class_name": "CLIPTextModelWithProjection"}) + configs[2].update_model_arch({"_class_name": "T5EncoderModel"}) + self.text_encoder_configs = tuple(configs) + + def get_text_encoder_pooler_output(self, outputs, encoder_index): + # SD3 uses pooled embeddings only from the two CLIP encoders (indices 0 and 1). + if encoder_index <= 1: + return outputs.pooler_output + return None + + def select_vae_weight_files( + self, + safetensors_list: list[str], + component_model_path: str, + component_name: str, + vae_precision: str, + ) -> list[str]: + return select_sd3_vae_weight_files( + safetensors_list=safetensors_list, + component_model_path=component_model_path, + component_name=component_name, + vae_precision=vae_precision, + ) + + def tokenize_prompt(self, prompt: list[str], tokenizer, tok_kwargs) -> dict: + text_inputs = tokenizer(prompt, **tok_kwargs) + text_inputs["attention_mask"] = None + return text_inputs + + def get_pos_prompt_embeds(self, batch): + return batch.prompt_embeds[0] + + def get_neg_prompt_embeds(self, batch): + return batch.negative_prompt_embeds[0] + + def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype): + return { + "pooled_projections": ( + batch.pooled_embeds[0] if batch.pooled_embeds else None + ) + } + + def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype): + return { + "pooled_projections": ( + batch.neg_pooled_embeds[0] if batch.neg_pooled_embeds else None + ) + } + + # SD3 image latents are spatial (B, C, H, W), not video-like (B, C, T, H, W). + def prepare_latent_shape(self, batch, batch_size, num_frames): # noqa: ARG002 + spatial_ratio = self.vae_config.arch_config.spatial_compression_ratio + in_channels = self.dit_config.arch_config.in_channels + return ( + batch_size, + in_channels, + batch.height // spatial_ratio, + batch.width // spatial_ratio, + ) diff --git a/python/sglang/multimodal_gen/configs/sample/stablediffusion3.py b/python/sglang/multimodal_gen/configs/sample/stablediffusion3.py new file mode 100644 index 000000000000..f586ee073c36 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/sample/stablediffusion3.py @@ -0,0 +1,20 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +# SPDX-License-Identifier: Apache-2.0 +"""StableDiffusion3 sampling parameters configuration.""" + +from dataclasses import dataclass + +from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams + + +@dataclass +class StableDiffusion3SamplingParams(SamplingParams): + """Sampling parameters for StableDiffusion3.""" + + # A single space ensures tokenizers produce valid (non-empty) input for CFG. + negative_prompt: str = " " + num_frames: int = 1 + num_inference_steps: int = 50 + guidance_scale: float = 7.0 + guidance_rescale: float = 0.0 diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index 93bbd9cc2c67..aaa49b08e10f 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -68,6 +68,9 @@ QwenImagePipelineConfig, ) from sglang.multimodal_gen.configs.pipeline_configs.sana import SanaPipelineConfig +from sglang.multimodal_gen.configs.pipeline_configs.stablediffusion3 import ( + StableDiffusion3PipelineConfig, +) from sglang.multimodal_gen.configs.pipeline_configs.wan import ( FastWan2_1_T2V_480P_Config, FastWan2_2_TI2V_5B_Config, @@ -109,6 +112,9 @@ QwenImageSamplingParams, ) from sglang.multimodal_gen.configs.sample.sana import SanaSamplingParams +from sglang.multimodal_gen.configs.sample.stablediffusion3 import ( + StableDiffusion3SamplingParams, +) from sglang.multimodal_gen.configs.sample.wan import ( FastWanT2V480PConfig, Turbo_Wan2_2_I2V_A14B_SamplingParam, @@ -849,6 +855,26 @@ def _register_configs(): hf_model_paths=["Qwen/Qwen-Image-Layered"], model_detectors=[lambda hf_id: "qwen-image-layered" in hf_id.lower()], ) + register_configs( + sampling_param_cls=StableDiffusion3SamplingParams, + pipeline_config_cls=StableDiffusion3PipelineConfig, + hf_model_paths=[ + "stabilityai/stable-diffusion-3-medium", + "stabilityai/stable-diffusion-3-medium-diffusers", + "stabilityai/stable-diffusion-3.5-medium", + "stabilityai/stable-diffusion-3.5-medium-diffusers", + "stabilityai/stable-diffusion-3.5-large", + "stabilityai/stable-diffusion-3.5-large-diffusers", + ], + model_detectors=[ + lambda hf_id: "stable-diffusion-3-medium" in hf_id.lower() + or "stable-diffusion-3.5-medium" in hf_id.lower() + or "stable-diffusion-3.5-large" in hf_id.lower() + or "sd3-medium" in hf_id.lower() + or "sd3.5-medium" in hf_id.lower() + or "sd3.5-large" in hf_id.lower() + ], + ) register_configs( sampling_param_cls=GlmImageSamplingParams, diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py index 01be87142c20..d7f9b0aeceb2 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py @@ -1,6 +1,7 @@ import dataclasses import glob import os +import re from collections.abc import Generator, Iterable from typing import cast @@ -218,21 +219,21 @@ def load_customized( component_path=component_model_path ) - def is_not_first_encoder(module_name): - return "2" in module_name - # TODO(mick): had to throw an exception for different text-encoder arch - if not is_not_first_encoder(component_name): - encoder_config = server_args.pipeline_config.text_encoder_configs[0] - encoder_config.update_model_arch(model_config) + encoder_index = self._extract_encoder_index(component_name) + assert encoder_index < len( + server_args.pipeline_config.text_encoder_configs + ) and encoder_index < len(server_args.pipeline_config.text_encoder_precisions) + + encoder_config = server_args.pipeline_config.text_encoder_configs[encoder_index] + encoder_config.update_model_arch(model_config) + + if encoder_index == 0: for key, value in diffusers_pretrained_config.__dict__.items(): setattr(encoder_config.arch_config, key, value) - encoder_dtype = server_args.pipeline_config.text_encoder_precisions[0] - else: - assert len(server_args.pipeline_config.text_encoder_configs) == 2 - encoder_config = server_args.pipeline_config.text_encoder_configs[1] - encoder_config.update_model_arch(model_config) - encoder_dtype = server_args.pipeline_config.text_encoder_precisions[1] + encoder_dtype = server_args.pipeline_config.text_encoder_precisions[ + encoder_index + ] # TODO(will): add support for other dtypes return self.load_model( component_model_path, @@ -242,6 +243,28 @@ def is_not_first_encoder(module_name): cpu_offload_flag=cpu_offload_flag, ) + @staticmethod + def _extract_encoder_index(component_name: str) -> int: + """ + Map text encoder component names to zero-based indices. + + Examples: + - text_encoder -> 0 + - text_encoder_2 -> 1 + - text_encoder_3 -> 2 + """ + match = re.search(r"_(\d+)$", component_name) + if match is None: + return 0 + + suffix_num = int(match.group(1)) + if suffix_num <= 0: + raise ValueError( + f"Invalid text encoder component name '{component_name}': " + "numeric suffix must be >= 1." + ) + return suffix_num - 1 + def load_model( self, model_path: str, diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py index 52c998ac8b33..c49aa4d6780e 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py @@ -129,6 +129,13 @@ def load_customized( vae = vae_cls(vae_config).to(target_device) safetensors_list = _list_safetensors_files(component_model_path) + safetensors_list = server_args.pipeline_config.select_vae_weight_files( + safetensors_list=safetensors_list, + component_model_path=component_model_path, + component_name=component_name, + vae_precision=vae_precision, + ) + assert ( len(safetensors_list) >= 1 ), f"Found no safetensors files in {component_model_path}" diff --git a/python/sglang/multimodal_gen/runtime/loader/utils.py b/python/sglang/multimodal_gen/runtime/loader/utils.py index d8ed09c2e850..edbcf65fb0c8 100644 --- a/python/sglang/multimodal_gen/runtime/loader/utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/utils.py @@ -181,9 +181,7 @@ def __exit__(self, exc_type, exc_value, traceback): def _normalize_component_type(module_type: str) -> str: """Normalize module types like 'text_encoder_2' -> 'text_encoder'.""" - if module_type.endswith("_2"): - return module_type[:-2] - return module_type + return re.sub(r"_\d+$", "", module_type) def _clean_hf_config_inplace(model_config: dict) -> None: diff --git a/python/sglang/multimodal_gen/runtime/models/dits/stablediffusion3.py b/python/sglang/multimodal_gen/runtime/models/dits/stablediffusion3.py new file mode 100644 index 000000000000..1ceeec2d9092 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/dits/stablediffusion3.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: Apache-2.0 +"""StableDiffusion3 Transformer model implementation. + +NOTE: This initial implementation uses diffusers' JointTransformerBlock directly. +A native SGLang attention implementation is needed for FlashAttention, TP/SP, +quantization, and LoRA support. +""" + +from typing import Any + +import torch +import torch.nn as nn +from diffusers.models.attention import JointTransformerBlock +from diffusers.models.embeddings import CombinedTimestepTextProjEmbeddings, PatchEmbed +from diffusers.models.normalization import AdaLayerNormContinuous + +from sglang.multimodal_gen.configs.models.dits.stablediffusion3 import ( + StableDiffusion3TransformerConfig, +) +from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + + +class SD3Transformer2DModel(CachableDiT): + _supports_gradient_checkpointing = True + _no_split_modules = ["JointTransformerBlock"] + _skip_layerwise_casting_patterns = ["pos_embed", "norm"] + + def __init__( + self, + config: StableDiffusion3TransformerConfig, + hf_config: dict[str, Any] | None = None, + quant_config=None, + ): + super().__init__(config=config, hf_config=hf_config) + self.config = config + arch_config = config.arch_config + sample_size = arch_config.sample_size + patch_size = arch_config.patch_size + in_channels = arch_config.in_channels + num_layers = arch_config.num_layers + attention_head_dim = arch_config.attention_head_dim + num_attention_heads = arch_config.num_attention_heads + joint_attention_dim = arch_config.joint_attention_dim + caption_projection_dim = arch_config.caption_projection_dim + pooled_projection_dim = arch_config.pooled_projection_dim + out_channels = arch_config.out_channels + pos_embed_max_size = arch_config.pos_embed_max_size + dual_attention_layers = arch_config.dual_attention_layers + qk_norm = arch_config.qk_norm + + self.out_channels = out_channels if out_channels is not None else in_channels + self.inner_dim = num_attention_heads * attention_head_dim + self.patch_size = patch_size + + self.pos_embed = PatchEmbed( + height=sample_size, + width=sample_size, + patch_size=patch_size, + in_channels=in_channels, + embed_dim=self.inner_dim, + pos_embed_max_size=pos_embed_max_size, + ) + self.time_text_embed = CombinedTimestepTextProjEmbeddings( + embedding_dim=self.inner_dim, pooled_projection_dim=pooled_projection_dim + ) + self.context_embedder = nn.Linear(joint_attention_dim, caption_projection_dim) + + self.transformer_blocks = nn.ModuleList( + [ + JointTransformerBlock( + dim=self.inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + context_pre_only=i == num_layers - 1, + qk_norm=qk_norm, + use_dual_attention=i in dual_attention_layers, + ) + for i in range(num_layers) + ] + ) + + self.norm_out = AdaLayerNormContinuous( + self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6 + ) + self.proj_out = nn.Linear( + self.inner_dim, patch_size * patch_size * self.out_channels, bias=True + ) + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + pooled_projections: torch.Tensor | None = None, + timestep: torch.LongTensor | None = None, + block_controlnet_hidden_states: list | None = None, + guidance: torch.Tensor | None = None, + joint_attention_kwargs: dict[str, Any] | None = None, + skip_layers: list[int] | None = None, + ) -> torch.Tensor: + if encoder_hidden_states is None: + raise ValueError("encoder_hidden_states must be provided.") + if pooled_projections is None: + raise ValueError("pooled_projections must be provided.") + + encoder_embeddings = encoder_hidden_states + + height, width = hidden_states.shape[-2:] + + hidden_states = self.pos_embed(hidden_states) + temb = self.time_text_embed(timestep, pooled_projections) + encoder_embeddings = self.context_embedder(encoder_embeddings) + + skip_layer_set = set(skip_layers) if skip_layers else set() + + if block_controlnet_hidden_states is not None: + interval_control = len(self.transformer_blocks) / len( + block_controlnet_hidden_states + ) + else: + interval_control = 0 + + for index_block, block in enumerate(self.transformer_blocks): + if index_block not in skip_layer_set: + encoder_embeddings, hidden_states = block( + hidden_states=hidden_states, + encoder_hidden_states=encoder_embeddings, + temb=temb, + joint_attention_kwargs=joint_attention_kwargs, + ) + + # controlnet residual + if ( + block_controlnet_hidden_states is not None + and block.context_pre_only is False + ): + hidden_states = ( + hidden_states + + block_controlnet_hidden_states[ + int(index_block / interval_control) + ] + ) + + hidden_states = self.norm_out(hidden_states, temb) + hidden_states = self.proj_out(hidden_states) + + # unpatchify + patch_size = self.patch_size + height = height // patch_size + width = width // patch_size + + hidden_states = hidden_states.reshape( + shape=( + hidden_states.shape[0], + height, + width, + patch_size, + patch_size, + self.out_channels, + ) + ) + hidden_states = hidden_states.permute(0, 5, 1, 3, 2, 4) + output = hidden_states.reshape( + shape=( + hidden_states.shape[0], + self.out_channels, + height * patch_size, + width * patch_size, + ) + ) + + return output + + +# Entry class for registry +EntryClass = SD3Transformer2DModel diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/clip.py b/python/sglang/multimodal_gen/runtime/models/encoders/clip.py index 56b4932e58a7..cda4f90ce8af 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/clip.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/clip.py @@ -594,6 +594,48 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loaded_params +class CLIPTextModelWithProjection(CLIPTextModel): + """ + CLIP text encoder with projection head for pooled_output. + """ + + def __init__( + self, + config: CLIPTextConfig, + ) -> None: + super().__init__(config) + self.text_projection = nn.Linear( + config.hidden_size, config.projection_dim, bias=False + ) + + def forward( + self, + input_ids: torch.Tensor | None, + position_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + **kwargs, + ) -> BaseEncoderOutput: + outputs: BaseEncoderOutput = self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + output_hidden_states=output_hidden_states, + ) + + pooled_output = outputs.pooler_output + if pooled_output is not None: + pooled_output = self.text_projection(pooled_output) + + return BaseEncoderOutput( + last_hidden_state=outputs.last_hidden_state, + pooler_output=pooled_output, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + class CLIPVisionTransformer(nn.Module): def __init__( @@ -759,4 +801,4 @@ class BertModel(CLIPTextModel): pass -EntryClass = [CLIPTextModel, CLIPVisionModel] +EntryClass = [CLIPTextModel, CLIPTextModelWithProjection, CLIPVisionModel] diff --git a/python/sglang/multimodal_gen/runtime/pipelines/stable_diffusion_3.py b/python/sglang/multimodal_gen/runtime/pipelines/stable_diffusion_3.py new file mode 100644 index 000000000000..f0b8c0b9fb71 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines/stable_diffusion_3.py @@ -0,0 +1,111 @@ +# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo + +# SPDX-License-Identifier: Apache-2.0 +"""StableDiffusion3 pipeline implementation.""" + +import torch + +from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( + ComposedPipelineBase, +) +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.pipelines_core.stages import ( + InputValidationStage, + PipelineStage, + TextEncodingStage, +) +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + + +class SD3ConditioningStage(PipelineStage): + """Merge CLIP-T, CLIP-G and T5 embeddings into unified prompt/pooled tensors.""" + + @torch.no_grad() + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + batch.prompt_embeds, batch.pooled_embeds = self._merge( + batch.prompt_embeds, batch.pooled_embeds + ) + if batch.do_classifier_free_guidance: + batch.negative_prompt_embeds, batch.neg_pooled_embeds = self._merge( + batch.negative_prompt_embeds, batch.neg_pooled_embeds + ) + return batch + + @staticmethod + def _merge( + embeds_list: list[torch.Tensor], + pooled_list: list[torch.Tensor], + ) -> tuple[list[torch.Tensor], list[torch.Tensor]]: + """Merge 3 encoder outputs into unified prompt/pooled tensors. + + SD3-medium uses exactly 3 text encoders (CLIP-L, CLIP-G, T5). + Returns single-element lists to match the batch field format expected + by downstream stages (get_pos_prompt_embeds accesses index [0]). + """ + if len(embeds_list) != 3: + raise ValueError( + f"SD3 requires exactly 3 prompt embedding tensors, got {len(embeds_list)}." + ) + if len(pooled_list) < 2: + raise ValueError( + f"SD3 requires at least 2 pooled embedding tensors, got {len(pooled_list)}." + ) + + clipt, clipg, t5 = embeds_list + clip_merged = torch.cat([clipt, clipg], dim=-1) + clip_merged = torch.nn.functional.pad( + clip_merged, (0, t5.shape[-1] - clip_merged.shape[-1]) + ) + merged_embeds = [torch.cat([clip_merged, t5], dim=-2)] + merged_pooled = [torch.cat([pooled_list[0], pooled_list[1]], dim=-1)] + return merged_embeds, merged_pooled + + +class StableDiffusion3Pipeline(ComposedPipelineBase): + """StableDiffusion3 pipeline implementation.""" + + pipeline_name = "StableDiffusion3Pipeline" + + _required_config_modules = [ + "text_encoder", + "text_encoder_2", + "text_encoder_3", + "tokenizer", + "tokenizer_2", + "tokenizer_3", + "vae", + "transformer", + "scheduler", + ] + + def create_pipeline_stages(self, server_args: ServerArgs): + self.add_stage(InputValidationStage()) + + self.add_stage( + TextEncodingStage( + text_encoders=[ + self.get_module("text_encoder"), + self.get_module("text_encoder_2"), + self.get_module("text_encoder_3"), + ], + tokenizers=[ + self.get_module("tokenizer"), + self.get_module("tokenizer_2"), + self.get_module("tokenizer_3"), + ], + ), + "prompt_encoding_stage_primary", + ) + + self.add_stage(SD3ConditioningStage()) + + self.add_standard_timestep_preparation_stage() + self.add_standard_latent_preparation_stage() + self.add_standard_denoising_stage() + self.add_standard_decoding_stage() + + +EntryClass = StableDiffusion3Pipeline diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py index b5d51ab6ecac..efb409e50be1 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py @@ -12,8 +12,6 @@ import torch from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput -from sglang.multimodal_gen.configs.pipeline_configs import FluxPipelineConfig -from sglang.multimodal_gen.configs.pipeline_configs.flux import Flux2PipelineConfig from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req @@ -263,11 +261,11 @@ def encode_text( ).to(target_device) input_ids = text_inputs["input_ids"] - is_flux_v1 = isinstance( - server_args.pipeline_config, FluxPipelineConfig - ) and not isinstance(server_args.pipeline_config, Flux2PipelineConfig) - - attention_mask = None if is_flux_v1 else text_inputs["attention_mask"] + attention_mask = ( + server_args.pipeline_config.get_text_encoder_attention_mask( + text_inputs, i + ) + ) encoder_forward_kwargs = { "input_ids": input_ids, "output_hidden_states": True, @@ -292,12 +290,13 @@ def encode_text( prompt_embeds = prompt_embeds.to(device=target_device) embeds_list.append(prompt_embeds) - if is_flux_v1 and outputs.pooler_output is not None: - # FLUX.1 only consumes the pooled CLIP projection. The T5 - # encoder in the same pipeline has no pooler output. - pooled_embeds_list.append( - outputs.pooler_output.to(device=target_device) - ) + + pooled_output = server_args.pipeline_config.get_text_encoder_pooler_output( + outputs, i + ) + if pooled_output is not None: + pooled_embeds_list.append(pooled_output.to(device=target_device)) + if return_attention_mask: mask_to_store = ( attention_mask.to(device=target_device) From 314d6ecf0880eefbd9d56cdc21e1b57bb390df80 Mon Sep 17 00:00:00 2001 From: DarkSharpness <76582120+DarkSharpness@users.noreply.github.com> Date: Mon, 13 Apr 2026 20:29:47 +0800 Subject: [PATCH 23/52] [Feature][JIT Kernel] Fused TP QK norm For Minimax (#20673) Co-authored-by: Mingyang Jiang <13463932+jmydurant@users.noreply.github.com> --- python/sglang/jit_kernel/all_reduce.py | 56 ++- .../jit_kernel/benchmark/bench_tp_qknorm.py | 170 +++++++++ .../jit_kernel/csrc/distributed/tp_qknorm.cuh | 325 ++++++++++++++++++ .../include/sgl_kernel/distributed/common.cuh | 20 +- .../distributed/custom_all_reduce.cuh | 23 +- .../jit_kernel/include/sgl_kernel/warp.cuh | 6 +- .../tests/test_custom_all_reduce.py | 31 +- .../sglang/jit_kernel/tests/test_tp_qknorm.py | 168 +++++++++ python/sglang/jit_kernel/tests/utils.py | 41 +++ .../custom_all_reduce_v2.py | 31 +- python/sglang/srt/models/minimax_m2.py | 134 ++++++-- 11 files changed, 923 insertions(+), 82 deletions(-) create mode 100644 python/sglang/jit_kernel/benchmark/bench_tp_qknorm.py create mode 100644 python/sglang/jit_kernel/csrc/distributed/tp_qknorm.cuh create mode 100644 python/sglang/jit_kernel/tests/test_tp_qknorm.py create mode 100644 python/sglang/jit_kernel/tests/utils.py diff --git a/python/sglang/jit_kernel/all_reduce.py b/python/sglang/jit_kernel/all_reduce.py index 48f763259beb..26167c698f10 100644 --- a/python/sglang/jit_kernel/all_reduce.py +++ b/python/sglang/jit_kernel/all_reduce.py @@ -5,6 +5,7 @@ import torch import tvm_ffi +from tvm_ffi import Module from sglang.jit_kernel.utils import ( cache_once, @@ -92,7 +93,7 @@ def config_pull( @cache_once -def _jit_custom_all_reduce_pull_module(dtype: torch.dtype, world_size: int): +def _jit_custom_all_reduce_pull_module(dtype: torch.dtype, world_size: int) -> Module: args = make_cpp_args(dtype, world_size, is_arch_support_pdl()) return load_jit( "custom_all_reduce_pull", @@ -104,7 +105,7 @@ def _jit_custom_all_reduce_pull_module(dtype: torch.dtype, world_size: int): @cache_once -def _jit_custom_all_reduce_push_module(dtype: torch.dtype, world_size: int): +def _jit_custom_all_reduce_push_module(dtype: torch.dtype, world_size: int) -> Module: args = make_cpp_args(dtype, world_size, is_arch_support_pdl()) return load_jit( "custom_all_reduce_push", @@ -115,6 +116,24 @@ def _jit_custom_all_reduce_push_module(dtype: torch.dtype, world_size: int): ) +@cache_once +def _jit_fused_parallel_qknorm_module( + dtype: torch.dtype, world_size: int, q_dim: int, k_dim: int +) -> Module: + args = make_cpp_args(dtype, world_size, q_dim, k_dim, is_arch_support_pdl()) + cls_name = f"FusedParallelQKNormAcrossHead<{args}>" + return load_jit( + "tp_qknorm", + *args, + extra_ldflags=["-lcuda"], + cuda_files=["distributed/tp_qknorm.cuh"], + cuda_wrappers=[ + ("fused_parallel_qknorm", f"{cls_name}::run"), + ("get_max_occupancy", f"{cls_name}::get_max_occupancy"), + ], + ) + + @cache_once def get_custom_all_reduce_cls() -> type[CustomAllReduceObj]: module = load_jit( @@ -144,18 +163,21 @@ def __init__( max_pull_blocks: Optional[int] = None, max_push_blocks: Optional[int] = None, ) -> None: + max_pull_blocks = NUM_CTA if max_pull_blocks is None else max_pull_blocks + max_push_blocks = NUM_CTA if max_push_blocks is None else max_push_blocks self.__ffi_init__( rank, world_size, - NUM_CTA if max_pull_blocks is None else max_pull_blocks, - NUM_CTA if max_push_blocks is None else max_push_blocks, + max_pull_blocks, + max_push_blocks, pull_buffer_bytes, push_buffer_bytes, graph_input_count, ) self._world_size = world_size - self._pull_config = ConfigResult(NUM_CTA, MAX_THREADS) - self.configure_pull(*self._pull_config) # type: ignore + self._pull_config = ConfigResult(min(NUM_CTA, max_pull_blocks), MAX_THREADS) + if max_pull_blocks > 0: # special case: cannot configure 0 blocks + self.configure_pull(*self._pull_config) # type: ignore @property def world_size(self) -> int: @@ -194,3 +216,25 @@ def free(self, tp_cpu_group: torch.distributed.ProcessGroup) -> None: self.free_storage() # type: ignore return cast(type["CustomAllReduceObj"], CustomAllReduceObjReal) + + +def get_fused_parallel_qknorm_max_occupancy( + dtype: torch.dtype, world_size: int, q_dim: int, k_dim: int +) -> int: + module = _jit_fused_parallel_qknorm_module(dtype, world_size, q_dim, k_dim) + return module.get_max_occupancy() + + +def fused_parallel_qknorm( + custom_ar: CustomAllReduceObj, + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + eps: float = 1e-6, +) -> None: + world_size = custom_ar.world_size + q_dim = q.shape[-1] * world_size + k_dim = k.shape[-1] * world_size + module = _jit_fused_parallel_qknorm_module(q.dtype, world_size, q_dim, k_dim) + module.fused_parallel_qknorm(custom_ar, q, k, q_weight, k_weight, eps) diff --git a/python/sglang/jit_kernel/benchmark/bench_tp_qknorm.py b/python/sglang/jit_kernel/benchmark/bench_tp_qknorm.py new file mode 100644 index 000000000000..d63545b10f48 --- /dev/null +++ b/python/sglang/jit_kernel/benchmark/bench_tp_qknorm.py @@ -0,0 +1,170 @@ +from __future__ import annotations + +import argparse +import os + +import torch +import torch.distributed as dist + +import sglang.srt.distributed.parallel_state as ps +from sglang.jit_kernel.all_reduce import ( + fused_parallel_qknorm, + get_fused_parallel_qknorm_max_occupancy, +) +from sglang.jit_kernel.utils import get_ci_test_range +from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import ( + CustomAllReduceV2, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci( + est_time=120, + suite="stage-b-kernel-benchmark-1-gpu-large", + disabled="requires multi-GPU, self-skips in CI", +) + +Q_K_DIMS = [(6144, 1024)] +DTYPE = torch.bfloat16 +EPS = 1e-6 +BATCH_SIZES = get_ci_test_range([2**i for i in range(15)], [1, 64, 1024]) +NUM_LAYERS = 8 + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--iters", type=int, default=100) + return parser.parse_args() + + +def init_distributed(): + local_rank = int(os.environ["LOCAL_RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + rank = local_rank + device = torch.device(f"cuda:{rank}") + torch.cuda.set_device(device) + + dist.init_process_group(backend="gloo") + ps._WORLD = coord = ps.init_world_group( + ranks=list(range(world_size)), + local_rank=local_rank, + backend="nccl", + ) + + cpu_group = coord.cpu_group + max_occupancy = get_fused_parallel_qknorm_max_occupancy( + DTYPE, world_size, Q_K_DIMS[0][0], Q_K_DIMS[0][1] + ) + if rank == 0: + print(f"Max occupancy for fused_parallel_qknorm: {max_occupancy} blocks/SM") + + props = torch.cuda.get_device_properties(device) + comm = CustomAllReduceV2( + cpu_group, + device, + max_pull_size=0, + max_push_size=8 * max(BATCH_SIZES), + max_push_blocks=props.multi_processor_count * max_occupancy, + ) + comm_ = CustomAllReduceV2(cpu_group, device) + if comm.disabled or comm_.disabled: + raise RuntimeError("JIT CustomAllReduceV2 is disabled on this system") + return rank, world_size, device, cpu_group, comm, comm_ + + +@torch.inference_mode() +def bench_one(fn, warmup: int, iters: int) -> float: + for _ in range(warmup): + fn(0) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for i in range(NUM_LAYERS): + fn(i) + + graph.replay() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + graph.replay() + start.record() + for i in range(iters): + graph.replay() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) * 1000.0 / (iters * NUM_LAYERS) + + +def rmsnorm_baseline( + comm_, + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + world_size: int, +) -> None: + from sglang.srt.models.minimax_m2 import rms_apply_serial, rms_sumsq_serial + + sum_sq = rms_sumsq_serial(q, k) + sum_sq = comm_.custom_all_reduce(sum_sq) + rms_apply_serial(q, k, q_weight, k_weight, sum_sq, world_size, EPS) + + +def main(): + args = parse_args() + rank, world_size, device, _, comm, comm_ = init_distributed() + torch.cuda.set_stream(torch.cuda.Stream()) + + if rank == 0: + print( + f"{'q_dim':>8} {'k_dim':>8} {'batch':>8} {'fused_us':>12} {'baseline_us':>12}" + ) + + for q_dim, k_dim in Q_K_DIMS: + local_q_dim = q_dim // world_size + local_k_dim = k_dim // world_size + for batch_size in BATCH_SIZES: + q = torch.randn( + NUM_LAYERS, batch_size, local_q_dim, device=device, dtype=DTYPE + ) + k = torch.randn( + NUM_LAYERS, batch_size, local_k_dim, device=device, dtype=DTYPE + ) + q_weight = torch.randn(NUM_LAYERS, local_q_dim, device=device, dtype=DTYPE) + k_weight = torch.randn(NUM_LAYERS, local_k_dim, device=device, dtype=DTYPE) + + def run_fused(i: int): + fused_parallel_qknorm( + comm.obj, + q[i], + k[i], + q_weight[i], + k_weight[i], + EPS, + ) + + def run_baseline(i: int): + rmsnorm_baseline( + comm_, + q[i], + k[i], + q_weight[i], + k_weight[i], + world_size, + ) + + fused_us = bench_one(run_fused, args.warmup, args.iters) + baseline_us = bench_one(run_baseline, args.warmup, args.iters) + + if rank == 0: + print( + f"{q_dim:8d} {k_dim:8d} {batch_size:8d} " + f"{fused_us:12.1f} {baseline_us:12.1f}" + ) + + comm.close() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/python/sglang/jit_kernel/csrc/distributed/tp_qknorm.cuh b/python/sglang/jit_kernel/csrc/distributed/tp_qknorm.cuh new file mode 100644 index 000000000000..ca80e1efcdf1 --- /dev/null +++ b/python/sglang/jit_kernel/csrc/distributed/tp_qknorm.cuh @@ -0,0 +1,325 @@ +// Adapted from https://github.com/NVIDIA/TensorRT-LLM/pull/12163 +// We reuse the custom all reduce push buffer in SGLang +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include + +namespace { + +using device::distributed::PushController; +using host::distributed::CustomAllReduceBase, host::distributed::CustomAllReduceRef; + +struct ParallelQKNormParams { + void* __restrict__ buffer[device::distributed::kMaxNumGPU]; + void* q_ptr; + void* k_ptr; + const void* __restrict__ q_weight; + const void* __restrict__ k_weight; + int64_t q_stride_bytes; + int64_t k_stride_bytes; + float eps; + uint32_t rank; + uint32_t num_tokens; + uint32_t epoch_bytes; + uint32_t num_clean_up_count = 0; +}; + +template +SGL_DEVICE void ld_global_volatile_8B(T& x, const void* addr, int64_t offset) { + static_assert(alignof(T) == 8 && sizeof(T) == 8); + addr = device::pointer::offset(addr, offset); + uint2 val; + asm volatile("ld.volatile.global.v2.b32 {%0, %1}, [%2];" : "=r"(val.x), "=r"(val.y) : "l"(addr)); + x = *reinterpret_cast(&val); +} + +template +SGL_DEVICE void st_global_volatile_8B(const T& x, void* addr, int64_t offset) { + static_assert(alignof(T) == 8 && sizeof(T) == 8); + const uint2 val = *reinterpret_cast(&x); + addr = device::pointer::offset(addr, offset); + asm volatile("st.volatile.global.v2.b32 [%2], {%0, %1};" ::"r"(val.x), "r"(val.y), "l"(addr)); +} + +[[maybe_unused]] +SGL_DEVICE float sync_float(float x) { + return __shfl_sync(0xffffffffu, x, 0); +} + +[[maybe_unused]] +constexpr auto next_pow_of_2(uint32_t x) { + uint32_t y = 1; + while (y < x) + y *= 2; + return y; +} + +template +struct KernelTrait { + // rename the arguments to avoid confusion with the template parameters + using DType = DType_; + static constexpr uint32_t kNumGPU = kNumGPU_; + static constexpr int64_t kQDim = kQDim_; + static constexpr int64_t kKDim = kKDim_; + static constexpr bool kUsePDL = kUsePDL_; + + static constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2); + static constexpr int64_t kLocalQDim = kQDim / kNumGPU; + static constexpr int64_t kLocalKDim = kKDim / kNumGPU; + static constexpr uint32_t kNumQThreads = kLocalQDim / (kVecSize * 2); + static constexpr uint32_t kNumKThreads = kLocalKDim / (kVecSize * 2); + static constexpr uint32_t kNumQWarps = kNumQThreads / device::kWarpThreads; + static constexpr uint32_t kNumKWarps = host::div_ceil(kNumKThreads, device::kWarpThreads); + static constexpr uint32_t kBlockSize = (kNumQWarps + kNumKWarps) * device::kWarpThreads; + static constexpr uint32_t kOccupancy = 2048 / kBlockSize; + + using DType2 = packed_t; + using Storage = device::AlignedVector; + + static_assert(std::has_single_bit(kNumGPU), "must be pow of 2"); + static_assert(kQDim % kNumGPU == 0); + static_assert(kKDim % kNumGPU == 0); + static_assert(kLocalQDim % (kVecSize * 2) == 0); + static_assert(kLocalKDim % (kVecSize * 2) == 0); + static_assert(kNumQThreads % device::kWarpThreads == 0); + static_assert(kBlockSize <= 1024); + static_assert(sizeof(Storage) == 16 && alignof(Storage) == 16); + static_assert(kOccupancy * kBlockSize <= 2048); +}; + +template +__global__ __launch_bounds__(Trait::kBlockSize, Trait::kOccupancy) void parallel_qknorm_across_head( + const ParallelQKNormParams __grid_constant__ params, const PushController __grid_constant__ ctrl) { + using namespace device; + + // each cta will handle exactly 1 token + using Storage = typename Trait::Storage; + using DType2 = typename Trait::DType2; + const auto &[ + buffer, q_ptr, k_ptr, q_weight, k_weight, q_stride_bytes, k_stride_bytes, // + eps, rank, num_tokens, epoch_bytes, num_clean_up_count + ] = params; + + using Package = AlignedVector; + constexpr uint32_t kNumGPU = Trait::kNumGPU; + constexpr uint32_t kNumQReduce = next_pow_of_2(Trait::kNumQWarps); + constexpr uint32_t kNumKReduce = next_pow_of_2(Trait::kNumKWarps); + __shared__ float smem_qk[Trait::kNumQWarps + Trait::kNumKWarps]; + __shared__ float scale_q; + __shared__ float scale_k; + const auto tx = threadIdx.x; + const auto bx = blockIdx.x; + /// NOTE: this can hint compiler to optimize `is_valid` out when not needed + constexpr uint32_t kActiveThreads = Trait::kNumQThreads + Trait::kNumKThreads; + const auto is_valid = Trait::kBlockSize == kActiveThreads || tx < kActiveThreads; + const auto smem_q = smem_qk + 0; + const auto smem_k = smem_qk + Trait::kNumQWarps; + const auto load_q = tx < Trait::kNumQThreads; + const auto offset = load_q ? tx : tx - Trait::kNumQThreads; + const auto input_ptr = load_q ? q_ptr : k_ptr; + const auto weight_ptr = load_q ? q_weight : k_weight; + const auto input_stride_bytes = load_q ? q_stride_bytes : k_stride_bytes; + PDLWaitPrimary(); + PDLTriggerSecondary(); + if (bx >= num_tokens) { + [[unlikely]]; + // In this case, we use the last few blocks to clean up other controllers + const auto start = (bx - num_tokens) * blockDim.x + threadIdx.x; + const auto stride = (gridDim.x - num_tokens) * blockDim.x; + for (uint32_t i = start; i < num_clean_up_count; i += stride) + ctrl.exit_unsafe(num_tokens + i); + return; + } + const auto epoch_offset = ctrl.epoch() * epoch_bytes; // only for comm + + __builtin_assume(bx < num_tokens); // since we have `bx >= num_tokens` + Storage next_input; + void* input_i_ptr = pointer::offset(input_ptr, bx * input_stride_bytes); + if (is_valid) next_input.load(input_i_ptr, offset); + + for (uint32_t i = bx; i < num_tokens; i += gridDim.x) { + // Stage 1. local reduce (warp-level) + Storage local_input; + { + float local_sum = 0.0; + if (is_valid) { + local_input = next_input; +#pragma unroll + for (uint32_t j = 0; j < Trait::kVecSize; ++j) { + const auto [x, y] = cast(local_input[j]); + local_sum += x * x + y * y; + } + } + smem_qk[threadIdx.x / kWarpThreads] = warp::reduce_sum(local_sum); + } + + // Stage 2. block reduce + push to peer ranks + poll from local rank + __syncthreads(); + + Storage local_weight; + const auto input_next_ptr = pointer::offset(input_i_ptr, gridDim.x * input_stride_bytes); + /** + * NOTE: Prefetch to hide the latency. + * This brings around 20% of performance gain in large batches + * The P2P communication is mainly latency bound, so during this waiting period, + * We can let some data loading transparently in the background. + */ + if (is_valid) { + local_weight.load(weight_ptr, offset); + if (i + gridDim.x < num_tokens) next_input.load(input_next_ptr, offset); + } + + if (tx < kWarpThreads) { + const auto local_sum_q = tx < Trait::kNumQWarps ? smem_q[tx] : 0.0f; + const auto local_sum_k = tx < Trait::kNumKWarps ? smem_k[tx] : 0.0f; + const auto sum_q = sync_float(warp::reduce_sum(local_sum_q)); + const auto sum_k = sync_float(warp::reduce_sum(local_sum_k)); + if (tx < kNumGPU) { // push a float2 pack to the peer + Package sum_q_k; + /// NOTE: eps should be scaled down by kNumGPU from host side + /// we add here to ensure that the sum is never zero + sum_q_k[0] = sum_q + eps; + sum_q_k[1] = sum_k + eps; + const auto push_ptr = pointer::offset(buffer[tx], epoch_offset); + st_global_volatile_8B(sum_q_k, push_ptr, i * kNumGPU + rank); + const auto poll_ptr = pointer::offset(buffer[rank], epoch_offset); + while (true) { + ld_global_volatile_8B(sum_q_k, poll_ptr, i * kNumGPU + tx); + if (sum_q_k[0] != 0.0f && sum_q_k[1] != 0.0f) break; + } + constexpr uint32_t kActiveMask = (1 << kNumGPU) - 1; + const auto global_sum_q = warp::reduce_sum(sum_q_k[0], kActiveMask); + const auto global_sum_k = warp::reduce_sum(sum_q_k[1], kActiveMask); + scale_q = math::rsqrt(global_sum_q / static_cast(Trait::kQDim)); + scale_k = math::rsqrt(global_sum_k / static_cast(Trait::kKDim)); + Package zeros; + zeros.fill(0.0f); + zeros.store(poll_ptr, i * kNumGPU + tx); + } + } + + __syncthreads(); + const auto scale = load_q ? scale_q : scale_k; + if (is_valid) { +#pragma unroll + for (uint32_t j = 0; j < Trait::kVecSize; ++j) { + const auto fp32_input = cast(local_input[j]); + const auto fp32_weight = cast(local_weight[j]); + const auto scaled_x = fp32_input.x * scale * fp32_weight.x; + const auto scaled_y = fp32_input.y * scale * fp32_weight.y; + local_input[j] = cast(fp32x2_t{scaled_x, scaled_y}); + } + local_input.store(input_i_ptr, offset); + } + input_i_ptr = input_next_ptr; + } + ctrl.exit(); +} + +template +struct FusedParallelQKNormAcrossHead : public CustomAllReduceBase { + using Trait = KernelTrait; + static constexpr auto kernel = parallel_qknorm_across_head; + static_assert(kNumGPU <= device::distributed::kMaxNumGPU, "kNumGPU exceeds the maximum supported GPUs"); + + void _run( + const tvm::ffi::Tensor q, + const tvm::ffi::Tensor k, + const tvm::ffi::Tensor q_weight, + const tvm::ffi::Tensor k_weight, + const float eps // passed in unscaled + ) { + using namespace host; + constexpr auto Q = Trait::kLocalQDim; + constexpr auto K = Trait::kLocalKDim; + auto N = SymbolicSize{"num_tokens"}; + auto device_ = SymbolicDevice{}; + device_.set_options(); + TensorMatcher({N, Q}) // q + .with_strides({-1, 1}) + .with_dtype() + .with_device(device_) + .verify(q); + TensorMatcher({N, K}) // k + .with_strides({-1, 1}) + .with_dtype() + .with_device(device_) + .verify(k); + TensorMatcher({Q}) // q_weight + .with_dtype() + .with_device(device_) + .verify(q_weight); + TensorMatcher({K}) // k_weight + .with_dtype() + .with_device(device_) + .verify(k_weight); + const auto device = device_.unwrap(); + const auto num_tokens = static_cast(N.unwrap()); + // use at most `world_size` blocks to clean up, + // this is based on the observation that occupancy is usually linear + // with respect to the world size + const bool need_clean = num_tokens < m_max_num_cta_push; + const auto num_clean = need_clean ? (m_max_num_cta_push - num_tokens) : 0; + const auto num_blocks = need_clean ? num_tokens + div_ceil(num_clean, Trait::kBlockSize) // + : m_max_num_cta_push; // + const auto num_threads = Trait::kBlockSize; + RuntimeCheck(num_blocks <= m_max_num_cta_push, "internal error"); + ParallelQKNormParams params; + for (uint32_t i = 0; i < kNumGPU; ++i) { + params.buffer[i] = get_push_buffer(m_peer_storage[i]); + } + params.q_ptr = q.data_ptr(); + params.k_ptr = k.data_ptr(); + params.q_weight = q_weight.data_ptr(); + params.k_weight = k_weight.data_ptr(); + params.q_stride_bytes = q.stride(0) * sizeof(DType); + params.k_stride_bytes = k.stride(0) * sizeof(DType); + params.eps = eps / kNumGPU; // scale down eps by number of GPUs + params.rank = m_rank; + params.num_tokens = num_tokens; + params.epoch_bytes = m_push_buffer_bytes; + params.num_clean_up_count = num_clean; + + const auto needed_buffer_bytes = static_cast(num_tokens) * 2 * sizeof(float); + RuntimeCheck(m_num_gpu == kNumGPU, "Number of GPUs mismatch"); + RuntimeCheck(m_push_ctrl.has_value(), "Controller is not initialized"); + RuntimeCheck(std::bit_cast(params.q_ptr) % 16 == 0, "q pointer is not properly aligned"); + RuntimeCheck(std::bit_cast(params.k_ptr) % 16 == 0, "k pointer is not properly aligned"); + RuntimeCheck(std::bit_cast(params.q_weight) % 16 == 0, "q_weight pointer is not properly aligned"); + RuntimeCheck(std::bit_cast(params.k_weight) % 16 == 0, "k_weight pointer is not properly aligned"); + RuntimeCheck(needed_buffer_bytes <= m_push_buffer_bytes, "Push buffer is too small"); + + LaunchKernel(num_blocks, num_threads, device) // + .enable_pdl(kUsePDL)(kernel, params, *m_push_ctrl); + } + + static uint32_t get_max_occupancy() { + return host::runtime::get_blocks_per_sm(kernel, Trait::kBlockSize); + } + + static void + run(CustomAllReduceRef obj, + const tvm::ffi::Tensor q, + const tvm::ffi::Tensor k, + const tvm::ffi::Tensor q_weight, + const tvm::ffi::Tensor k_weight, + const float eps) { + using Self = FusedParallelQKNormAcrossHead; + return static_cast(obj.get())->_run(q, k, q_weight, k_weight, eps); + } +}; + +} // namespace diff --git a/python/sglang/jit_kernel/include/sgl_kernel/distributed/common.cuh b/python/sglang/jit_kernel/include/sgl_kernel/distributed/common.cuh index c9711c5b6557..e0ce2dc086c1 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/distributed/common.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/distributed/common.cuh @@ -48,6 +48,8 @@ struct alignas(128) Semaphore { struct PullController { public: + using SignalType = Semaphore; + PullController(void** signals, uint32_t num_gpu) { for (uint32_t i = 0; i < num_gpu; ++i) { m_signals[i] = static_cast(signals[i]); @@ -90,25 +92,29 @@ struct PullController { struct PushController { public: + using SignalType = uint32_t; static constexpr int64_t kNumStages = 2; - PushController(void* ptr) : m_local_signal(static_cast(ptr)) {} + PushController(void* ptr) : m_local_signal(static_cast(ptr)) {} - SGL_DEVICE uint32_t epoch() const { - return m_local_signal[blockIdx.x].get_counter(); + SGL_DEVICE SignalType epoch() const { + return m_local_signal[blockIdx.x]; } SGL_DEVICE void exit() const { __syncthreads(); if (threadIdx.x == 0) { - auto& signal = m_local_signal[blockIdx.x]; - const auto epoch = signal.get_counter(); - signal.set_counter((epoch + 1) % kNumStages); + this->exit_unsafe(blockIdx.x); } } + SGL_DEVICE void exit_unsafe(uint32_t which) const { + auto& signal = m_local_signal[which]; + signal = (signal + 1) % kNumStages; + } + private: - Semaphore* m_local_signal; + SignalType* m_local_signal; }; } // namespace device::distributed diff --git a/python/sglang/jit_kernel/include/sgl_kernel/distributed/custom_all_reduce.cuh b/python/sglang/jit_kernel/include/sgl_kernel/distributed/custom_all_reduce.cuh index 7b627a9e9568..239fac71a198 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/distributed/custom_all_reduce.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/distributed/custom_all_reduce.cuh @@ -93,12 +93,14 @@ struct CustomAllReduceBase : public tvm::ffi::Object { // default config for pull kernel, can be updated by `configure()` m_num_cta(max_num_cta_pull), m_cta_size(256) { - RuntimeDeviceCheck(cudaMalloc(&m_storage, storage_bytes())); + RuntimeCheck(pull_buffer_size % 128 == 0, "Pull buffer size should be aligned to 128 bytes"); + RuntimeCheck(push_buffer_size % 128 == 0, "Push buffer size should be aligned to 128 bytes"); RuntimeCheck(rank < num_gpu, "Invalid rank: ", rank); const int64_t kU32Max = static_cast(std::numeric_limits::max()); const int64_t push_buffer_size_all = push_all_ranks_bytes(); - RuntimeCheck(pull_buffer_size <= kU32Max, "Buffer size is too large: ", pull_buffer_size); + RuntimeCheck(pull_buffer_size <= kU32Max, "Pull buffer size is too large: ", pull_buffer_size); RuntimeCheck(push_buffer_size_all <= kU32Max, "Push buffer size is too large: ", push_buffer_size_all); + RuntimeDeviceCheck(cudaMalloc(&m_storage, storage_bytes())); } ExternHandle share_storage() { @@ -252,19 +254,18 @@ struct CustomAllReduceBase : public tvm::ffi::Object { return static_cast(m_graph_capture_inputs.size()); } int64_t pull_signal_bytes() const { - return sizeof(device::distributed::Semaphore) * m_max_num_cta_pull; + return _align_bytes(sizeof(PullController::SignalType) * m_max_num_cta_pull); } int64_t push_signal_bytes() const { - return sizeof(device::distributed::Semaphore) * m_max_num_cta_push; + return _align_bytes(sizeof(PushController::SignalType) * m_max_num_cta_push); } - int64_t params_bytes() const { - return sizeof(AllReduceData) * (1 + m_graph_buffer_count); // 1 for default + int64_t graph_param_bytes() const { + return _align_bytes(sizeof(AllReduceData) * (1 + m_graph_buffer_count)); // 1 for default } int64_t push_all_ranks_bytes() const { - return PushController::kNumStages * m_num_gpu * m_push_buffer_bytes; + return _align_bytes(PushController::kNumStages * m_num_gpu * m_push_buffer_bytes); } int64_t storage_bytes() const { - // | SignalArray (pull + push) | GraphBuffers (pull params) | Buffers (pull + push) | return _get_offset_impl(5); } void* get_pull_signal(void* ptr) const { @@ -283,16 +284,20 @@ struct CustomAllReduceBase : public tvm::ffi::Object { return pointer::offset(ptr, _get_offset_impl(4)); } int64_t _get_offset_impl(int64_t which) const { + // | SignalArray (pull + push) | GraphBuffers (pull params) | Buffers (pull + push) | const int64_t offset_map[5] = { /*[0]=*/pull_signal_bytes(), /*[1]=*/push_signal_bytes(), - /*[2]=*/params_bytes(), + /*[2]=*/graph_param_bytes(), /*[3]=*/m_pull_buffer_bytes, /*[4]=*/push_all_ranks_bytes(), }; RuntimeCheck(which >= 0 && which <= 5, "Invalid offset index: ", which); return std::accumulate(offset_map, offset_map + which, int64_t(0)); } + static int64_t _align_bytes(int64_t size) { + return div_ceil(size, 128) * 128; + } const int64_t m_pull_buffer_bytes; const int64_t m_push_buffer_bytes; diff --git a/python/sglang/jit_kernel/include/sgl_kernel/warp.cuh b/python/sglang/jit_kernel/include/sgl_kernel/warp.cuh index 52952ff81406..769c16c08603 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/warp.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/warp.cuh @@ -21,10 +21,12 @@ static constexpr uint32_t kFullMask = 0xffffffffu; * \param active_mask Bitmask of participating lanes (default: all 32). * \return The sum across all active lanes. */ -template +template SGL_DEVICE T reduce_sum(T value, uint32_t active_mask = kFullMask) { + static_assert(kNumThreads >= 1 && kNumThreads <= kWarpThreads); + static_assert(std::has_single_bit(kNumThreads), "must be pow of 2"); #pragma unroll - for (int mask = 16; mask > 0; mask >>= 1) + for (int mask = kNumThreads / 2; mask > 0; mask >>= 1) value = value + __shfl_xor_sync(active_mask, value, mask, 32); return value; } diff --git a/python/sglang/jit_kernel/tests/test_custom_all_reduce.py b/python/sglang/jit_kernel/tests/test_custom_all_reduce.py index 2d7e0253eb43..96a61bfbf785 100644 --- a/python/sglang/jit_kernel/tests/test_custom_all_reduce.py +++ b/python/sglang/jit_kernel/tests/test_custom_all_reduce.py @@ -18,8 +18,6 @@ import logging import multiprocessing as mp import os -import subprocess -import sys from typing import Dict, Optional, Tuple import pytest @@ -32,6 +30,7 @@ _jit_custom_all_reduce_pull_module, _jit_custom_all_reduce_push_module, ) +from sglang.jit_kernel.tests.utils import multiprocess_main, multiprocess_test from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import ( CustomAllReduceV2, ) @@ -79,26 +78,6 @@ # --------------------------------------------------------------------------- -def _run_torchrun(nproc: int, timeout: int = 300) -> None: - """Launch this script as a torchrun worker and assert success.""" - cmd = [ - "torchrun", - f"--nproc_per_node={nproc}", - __file__, - ] - result = subprocess.run( - cmd, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - text=True, - timeout=timeout, - ) - assert result.returncode == 0, ( - f"torchrun (nproc={nproc}) failed with rc={result.returncode}\n" - f"{result.stdout}" - ) - - def _compile_one(dtype: torch.dtype, world_size: int): _jit_custom_all_reduce_push_module(dtype, world_size) _jit_custom_all_reduce_pull_module(dtype, world_size) @@ -129,7 +108,7 @@ def test_custom_allreduce(nproc: int) -> None: pytest.skip( f"Requires at least {nproc} GPUs, but only {device_count} available" ) - _run_torchrun(nproc) + multiprocess_test(__file__, nproc) # --------------------------------------------------------------------------- @@ -229,7 +208,6 @@ def run_eager(x: torch.Tensor) -> torch.Tensor: def worker_main() -> None: """Entry point for each torchrun worker process.""" rank, device, cpu_group, nccl_group, comm = init_distributed() - world_size = dist.get_world_size() torch.cuda.set_stream(torch.cuda.Stream()) @@ -258,7 +236,4 @@ def worker_main() -> None: if __name__ == "__main__": - if "LOCAL_RANK" in os.environ: - worker_main() - else: - sys.exit(pytest.main([__file__, "-x", "-vv", "-s"])) + multiprocess_main(__file__, worker_main) diff --git a/python/sglang/jit_kernel/tests/test_tp_qknorm.py b/python/sglang/jit_kernel/tests/test_tp_qknorm.py new file mode 100644 index 000000000000..98a803ca06dc --- /dev/null +++ b/python/sglang/jit_kernel/tests/test_tp_qknorm.py @@ -0,0 +1,168 @@ +from __future__ import annotations + +import itertools +import os +from typing import Optional + +import pytest +import torch +import torch.distributed as dist +import triton + +from sglang.jit_kernel.all_reduce import fused_parallel_qknorm +from sglang.jit_kernel.tests.test_custom_all_reduce import multiprocess_test +from sglang.jit_kernel.tests.utils import multiprocess_main +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci( + est_time=300, + suite="stage-b-kernel-unit-8-gpu-h200", +) +register_cuda_ci( + est_time=300, + suite="nightly-kernel-8-gpu-h200", + nightly=True, +) + + +Q_K_DIMS = [(6144, 1024)] +EPS = 1e-6 +BATCH_SIZES = [2**n for n in range(0, 14)] +DTYPES = [torch.float16, torch.bfloat16, torch.float32] +TEST_CONFIG = list(itertools.product(Q_K_DIMS, BATCH_SIZES, DTYPES)) + + +@pytest.mark.parametrize("nproc", [2, 4, 8]) +def test_tp_qknorm(nproc: int) -> None: + device_count = torch.cuda.device_count() + if device_count < nproc: + pytest.skip( + f"Requires at least {nproc} GPUs, but only {device_count} available" + ) + multiprocess_test(__file__, nproc) + + +def init_distributed(): + import sglang.srt.distributed.parallel_state as ps + from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import ( + CustomAllReduceV2, + ) + + local_rank = int(os.environ["LOCAL_RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + rank = local_rank + device = torch.device(f"cuda:{rank}") + torch.cuda.set_device(device) + + dist.init_process_group(backend="gloo") + ps._WORLD = coord = ps.init_world_group( + ranks=list(range(world_size)), + local_rank=local_rank, + backend="nccl", + ) + + cpu_group = coord.cpu_group + nccl_group = coord.device_group + assert nccl_group is not None + + max_pull_size = 0 + max_push_size = 8 * max(BATCH_SIZES) + comm = CustomAllReduceV2(cpu_group, device, max_pull_size, max_push_size) + if comm.disabled: + raise RuntimeError("JIT CustomAllReduceV2 is disabled on this system") + + return rank, world_size, device, cpu_group, nccl_group, comm + + +def _all_gather_cat(x: torch.Tensor, group: dist.ProcessGroup) -> torch.Tensor: + gathered = [torch.empty_like(x) for _ in range(dist.get_world_size(group=group))] + dist.all_gather(gathered, x, group=group) + return torch.cat(gathered, dim=-1) + + +def _rmsnorm_ref(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + x_fp32 = x.float() + scale = (x_fp32.pow(2).mean(dim=-1, keepdim=True) + eps).rsqrt() + return (x_fp32 * scale * weight.float()).to(x.dtype) + + +@torch.inference_mode() +def worker_test( + rank: int, + world_size: int, + device: torch.device, + nccl_group: dist.ProcessGroup, + comm, + q_k_dim: tuple[int, int], + batch_size: int, + dtype: torch.dtype, +) -> Optional[RuntimeError]: + q_dim, k_dim = q_k_dim + local_q_dim = q_dim // world_size + local_k_dim = k_dim // world_size + + q = torch.randn(batch_size, local_q_dim, device=device, dtype=dtype) + k = torch.randn(batch_size, local_k_dim, device=device, dtype=dtype) + q_weight = torch.randn(local_q_dim, device=device, dtype=dtype) + k_weight = torch.randn(local_k_dim, device=device, dtype=dtype) + + q_ref = _all_gather_cat(q, nccl_group) + k_ref = _all_gather_cat(k, nccl_group) + q_weight_ref = _all_gather_cat(q_weight.unsqueeze(0), nccl_group).squeeze(0) + k_weight_ref = _all_gather_cat(k_weight.unsqueeze(0), nccl_group).squeeze(0) + + q_expected = _rmsnorm_ref(q_ref, q_weight_ref, EPS) + k_expected = _rmsnorm_ref(k_ref, k_weight_ref, EPS) + q_expected = q_expected[:, rank * local_q_dim : (rank + 1) * local_q_dim] + k_expected = k_expected[:, rank * local_k_dim : (rank + 1) * local_k_dim] + + fused_parallel_qknorm( + comm.obj, + q, + k, + q_weight, + k_weight, + EPS, + ) + + try: + triton.testing.assert_close(q, q_expected, atol=1e-2, rtol=1e-2) + triton.testing.assert_close(k, k_expected, atol=1e-2, rtol=1e-2) + except AssertionError as err: + return RuntimeError( + f"TP QKNorm mismatch for {batch_size=}, {dtype=}, {world_size=}, {rank=}: {err}" + ) + return None + + +def worker_main() -> None: + rank, world_size, device, cpu_group, nccl_group, comm = init_distributed() + torch.cuda.set_stream(torch.cuda.Stream()) + + for q_k_dim, batch_size, dtype in TEST_CONFIG: + error = worker_test( + rank, + world_size, + device, + nccl_group, + comm, + q_k_dim, + batch_size, + dtype, + ) + result = torch.tensor([int(error is not None)]) + dist.all_reduce(result, group=cpu_group) + if error is not None: + print(str(error)) + if bool(result.item()): + raise RuntimeError( + f"TP QKNorm test failed for {q_k_dim=}, {batch_size=}, {dtype=}, {world_size=}" + ) + + print(f"Rank {rank} passed all tests.") + comm.close() + dist.destroy_process_group() + + +if __name__ == "__main__": + multiprocess_main(__file__, worker_main) diff --git a/python/sglang/jit_kernel/tests/utils.py b/python/sglang/jit_kernel/tests/utils.py new file mode 100644 index 000000000000..56f3511ab8c7 --- /dev/null +++ b/python/sglang/jit_kernel/tests/utils.py @@ -0,0 +1,41 @@ +import os +import subprocess +import sys +from typing import Callable + +import pytest + + +def multiprocess_test(file: str, nproc: int, timeout: int = 90) -> None: + """Launch this script as a torchrun worker and assert success.""" + cmd = [ + "torchrun", + f"--nproc_per_node={nproc}", + file, + ] + try: + result = subprocess.run( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + timeout=timeout, + ) + except subprocess.TimeoutExpired as e: + raise RuntimeError( + f"torchrun (nproc={nproc}) timed out after {timeout} seconds\n" + f"{e.stdout}" + ) from e + + assert result.returncode == 0, ( + f"torchrun (nproc={nproc}) failed with rc={result.returncode}\n" + f"{result.stdout}" + ) + + +def multiprocess_main(file: str, main: Callable[[], None]) -> None: + """Helper to run a function in a multiprocess torchrun context.""" + if "LOCAL_RANK" in os.environ: + main() + else: + sys.exit(pytest.main([file, "-v", "-s"])) diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py index 2006554f9fd6..11599552fdf9 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py +++ b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py @@ -35,6 +35,8 @@ def __init__( device: torch.device, max_pull_size: Optional[int] = None, max_push_size: Optional[int] = None, + max_pull_blocks: Optional[int] = None, + max_push_blocks: Optional[int] = None, ) -> None: _init_config() self.disabled = True @@ -50,14 +52,15 @@ def __init__( self.group = group self.rank = dist.get_rank(group=self.group) self.world_size = dist.get_world_size(group=self.group) - self.override_shot(None) - if max_pull_size is None: - max_pull_size = 16 * 1024 * 1024 # default to 16MB - if max_push_size is None: - max_push_size = self.config.one_shot_push_threshold - max_push_size = min(max_push_size, max_pull_size) + if max_pull_size is None: # default to 16MB + max_pull_size = 16 * 1024 * 1024 + if max_push_size is None: # default to recommended size + config = THRESHOLD_2_SHOT_MAP[self.world_size] + max_push_size = config.one_shot_push_threshold self.max_pull_size = max_pull_size self.max_push_size = max_push_size + self.max_size = max(max_pull_size, max_push_size) + self.override_shot(None) # set default config based on world size self.override_algo: Optional[AllReduceAlgo] = None self.obj = get_custom_all_reduce_cls()( rank=self.rank, @@ -65,6 +68,8 @@ def __init__( pull_buffer_bytes=self.max_pull_size, push_buffer_bytes=self.max_push_size, graph_input_count=131072, + max_pull_blocks=max_pull_blocks, + max_push_blocks=max_push_blocks, ) self._post_init_obj() self.disabled = False @@ -72,11 +77,19 @@ def __init__( def override_shot(self, shot: int | None): if shot is None: - self.config = THRESHOLD_2_SHOT_MAP[self.world_size] + config = THRESHOLD_2_SHOT_MAP[self.world_size] else: assert shot in (1, 2) threshold = INF if shot == 1 else 0 - self.config = replace(self.config, one_shot_pull_threshold=threshold) + config = replace(self.config, one_shot_pull_threshold=threshold) + # need to clip the config thresholds to max sizes to avoid invalid config + push_threshold = min(config.one_shot_push_threshold, self.max_push_size) + pull_threshold = min(config.one_shot_pull_threshold, self.max_pull_size) + self.config: ModeConfig = replace( + config, + one_shot_push_threshold=push_threshold, + one_shot_pull_threshold=pull_threshold, + ) @contextmanager def capture(self): @@ -109,7 +122,7 @@ def should_custom_ar(self, inp: torch.Tensor) -> bool: return False if not is_weak_contiguous(inp): return False - return inp_size <= self.max_pull_size + return inp_size <= self.max_size def custom_all_reduce(self, input: torch.Tensor) -> torch.Tensor: if is_in_piecewise_cuda_graph(): # disable inplace optimization diff --git a/python/sglang/srt/models/minimax_m2.py b/python/sglang/srt/models/minimax_m2.py index 1904d046a2a5..eea60e5868cc 100644 --- a/python/sglang/srt/models/minimax_m2.py +++ b/python/sglang/srt/models/minimax_m2.py @@ -17,7 +17,8 @@ import logging from contextlib import nullcontext -from typing import Iterable, Optional, Set, Tuple, Union +from functools import lru_cache +from typing import Any, Dict, Iterable, Optional, Set, Tuple, Union import torch import triton @@ -25,9 +26,14 @@ from torch import nn from transformers import PretrainedConfig +from sglang.jit_kernel.all_reduce import ( + fused_parallel_qknorm, + get_fused_parallel_qknorm_max_occupancy, +) from sglang.kernel_api_logging import debug_kernel_api from sglang.srt.batch_overlap.two_batch_overlap import model_forward_maybe_tbo from sglang.srt.distributed import ( + get_bool_env_var, get_moe_expert_parallel_world_size, get_pp_group, get_tensor_model_parallel_world_size, @@ -42,6 +48,7 @@ ) from sglang.srt.layers.dp_attention import ( attn_tp_all_reduce, + get_attention_tp_group, get_attention_tp_rank, get_attention_tp_size, is_dp_attention_enabled, @@ -78,12 +85,15 @@ BumpAllocator, add_prefix, get_compiler_backend, + is_cuda, is_non_idle_and_non_empty, make_layers, ) +from sglang.srt.utils.custom_op import register_custom_op from sglang.srt.utils.hf_transformers_utils import get_rope_config logger = logging.getLogger(__name__) +_is_cuda = is_cuda() @triton.jit @@ -324,27 +334,114 @@ def forward( return x + +@register_custom_op(mutates_args=["q", "k"]) +def fused_tp_qknorm( + counter: int, + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + eps: float, +) -> None: + return fused_parallel_qknorm( + MiniMaxM2QKRMSNorm.COMM_MAP[counter].obj, + q, + k, + q_weight, + k_weight, + eps=eps, + ) + + +class MiniMaxM2QKRMSNorm: + COUNTER = 0 + COMM_MAP: Dict[int, Any] = {} + + def __init__( + self, + q_norm: MiniMaxM2RMSNormTP, + k_norm: MiniMaxM2RMSNormTP, + ) -> None: + assert q_norm.variance_epsilon == k_norm.variance_epsilon + self._q_norm = q_norm + self._k_norm = k_norm + self._world_size = self._q_norm.attn_tp_size + self._eps = q_norm.variance_epsilon + use_fused_norm = get_bool_env_var("SGLANG_USE_FUSED_PARALLEL_QKNORM") + + self._forward_impl = self._forward_naive + if self._world_size > 1 and _is_cuda and use_fused_norm: + occupancy = get_fused_parallel_qknorm_max_occupancy( + q_norm.weight.dtype, + self._world_size, + # NOTE: we need full dimension + q_dim=q_norm.weight.shape[0] * self._world_size, + k_dim=k_norm.weight.shape[0] * self._world_size, + ) + counter = MiniMaxM2QKRMSNorm._get_comm(q_norm.weight.device, occupancy) + if counter is not None: + self._counter = counter + self._forward_impl = self._forward_fused + + @lru_cache @staticmethod - def forward_qk( - q_norm: "MiniMaxM2RMSNormTP", - k_norm: "MiniMaxM2RMSNormTP", - q: torch.Tensor, - k: torch.Tensor, - ) -> torch.Tensor: + def _get_comm(device: torch.device, occupancy: int): + from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import ( + CustomAllReduceV2, + ) + + props = torch.cuda.get_device_properties(device) + # probe the maximum tokens for one prefill + server_args = get_global_server_args() + max_tokens = server_args.chunked_prefill_size + if max_tokens is None: + max_tokens = server_args.model_config.context_len + max_tokens = max(max_tokens, server_args.max_prefill_tokens) + logger.info(f"[AR] Using CustomAllReduceV2 for MiniMaxM2 with {max_tokens = }") + ALIGN = 512 + # typically, this should not exceed 1M, since max_tokens is usually less than 16384 + max_size = ((8 * max_tokens + ALIGN - 1) // ALIGN) * ALIGN + comm = CustomAllReduceV2( + group=get_attention_tp_group().cpu_group, + device=device, + max_pull_size=0, + max_pull_blocks=0, + max_push_size=max_size, + max_push_blocks=props.multi_processor_count * occupancy, + ) + counter = MiniMaxM2QKRMSNorm.COUNTER + MiniMaxM2QKRMSNorm.COUNTER += 1 + MiniMaxM2QKRMSNorm.COMM_MAP[counter] = comm + return counter if not comm.disabled else None + + def forward(self, q: torch.Tensor, k: torch.Tensor): + return self._forward_impl(q, k) + + def _forward_naive(self, q: torch.Tensor, k: torch.Tensor): + q, k = q.contiguous(), k.contiguous() sum_sq = rms_sumsq_serial(q, k) - if q_norm.attn_tp_size > 1: + if self._world_size > 1: sum_sq = attn_tp_all_reduce(sum_sq) - - q, k = rms_apply_serial( + return rms_apply_serial( q, k, - q_norm.weight, - k_norm.weight, + self._q_norm.weight, + self._k_norm.weight, sum_sq, - q_norm.attn_tp_size, - q_norm.variance_epsilon, + self._world_size, + self._eps, ) + def _forward_fused(self, q: torch.Tensor, k: torch.Tensor): + fused_tp_qknorm( + self._counter, + q, + k, + self._q_norm.weight, + self._k_norm.weight, + self._eps, + ) return q, k @@ -681,6 +778,7 @@ def __init__( num_heads=self.total_num_kv_heads, eps=config.rms_norm_eps, ) + self.qk_norm_impl = MiniMaxM2QKRMSNorm(self.q_norm, self.k_norm) else: raise ValueError(f"Unsupported qk_norm_type: {self.qk_norm_type}") @@ -708,13 +806,7 @@ def forward_prepare( qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) if self.use_qk_norm: - # q = self.q_norm(q.contiguous()) - # k = self.k_norm(k.contiguous()) - q, k = MiniMaxM2RMSNormTP.forward_qk( - self.q_norm, self.k_norm, q.contiguous(), k.contiguous() - ) - else: - q, k = q.contiguous(), k.contiguous() + q, k = self.qk_norm_impl.forward(q, k) q, k = self.rotary_emb(positions, q, k) inner_state = q, k, v, forward_batch return None, forward_batch, inner_state From d01b2bf257e7d12bc7513ed69e8296be963b4b90 Mon Sep 17 00:00:00 2001 From: xdtbynd <231153518+xdtbynd@users.noreply.github.com> Date: Mon, 13 Apr 2026 21:22:37 +0800 Subject: [PATCH 24/52] [Docs] Fix default values and options in Ascend server arguments documentation (#22698) Co-authored-by: xdtbynd --- .../ascend/ascend_npu_support_features.md | 64 +++++++++---------- 1 file changed, 32 insertions(+), 32 deletions(-) diff --git a/docs/platforms/ascend/ascend_npu_support_features.md b/docs/platforms/ascend/ascend_npu_support_features.md index 3d36ca76f9d9..c5311c83ee9b 100644 --- a/docs/platforms/ascend/ascend_npu_support_features.md +++ b/docs/platforms/ascend/ascend_npu_support_features.md @@ -34,7 +34,7 @@ click [Server Arguments](https://docs.sglang.io/advanced_features/server_argumen | `--warmups` | `None` | Type: str | A2, A3 | | `--nccl-port` | `None` | Type: int | A2, A3 | | `--fastapi-root-path` | `None` | Type: str | A2, A3 | -| `--grpc-mode` | `False` | bool flag (set to enable) | A2, A3 | +| `--grpc-mode` | `False` | `False` | Planned | ## Quantization and data type @@ -143,19 +143,19 @@ click [Server Arguments](https://docs.sglang.io/advanced_features/server_argumen ## API related -| Argument | Defaults | Options | Server supported | -|-------------------------|-----------|--------------------------------|:----------------:| -| `--api-key` | `None` | Type: str | A2, A3 | -| `--admin-api-key` | `None` | Type: str | A2, A3 | -| `--served-model-name` | `None` | Type: str | A2, A3 | -| `--weight-version` | `default` | Type: str | A2, A3 | -| `--chat-template` | `None` | Type: str | A2, A3 | -| `--hf-chat-template-name` | `None` | Type: str | A2, A3 | -| `--completion-template` | `None` | Type: str | A2, A3 | -| `--enable-cache-report` | `False` | bool flag
(set to enable) | A2, A3 | -| `--reasoning-parser` | `None` | `deepseek-r1`
`deepseek-v3`
`glm45`
`gpt-oss`
`kimi`
`qwen3`
`qwen3-thinking`
`step3` | A2, A3 | -| `--tool-call-parser` | `None` | `deepseekv3`
`deepseekv31`
`glm`
`glm45`
`glm47`
`gpt-oss`
`kimi_k2`
`llama3`
`mistral`
`pythonic`
`qwen`
`qwen25`
`qwen3_coder`
`step3`
`gigachat3` | A2, A3 | -| `--sampling-defaults` | `model` | `openai`, `model` | A2, A3 | +| Argument | Defaults | Options | Server supported | +|---------------------------|-----------|-------------------------------------------------------------------------------------------------------------------|:----------------:| +| `--api-key` | `None` | Type: str | A2, A3 | +| `--admin-api-key` | `None` | Type: str | A2, A3 | +| `--served-model-name` | `None` | Type: str | A2, A3 | +| `--weight-version` | `default` | Type: str | A2, A3 | +| `--chat-template` | `None` | Type: str | A2, A3 | +| `--hf-chat-template-name` | `None` | Type: str | A2, A3 | +| `--completion-template` | `None` | Type: str | A2, A3 | +| `--enable-cache-report` | `False` | bool flag
(set to enable) | A2, A3 | +| `--reasoning-parser` | `None` | `deepseek-r1`
`deepseek-v3`
`glm45`
`gpt-oss`
`kimi`
`qwen3`
`qwen3-thinking`
`step3` | A2, A3 | +| `--tool-call-parser` | `None` | `llama3`
``pythonic`
`qwen`
`qwen3_coder` | A2, A3 | +| `--sampling-defaults` | `model` | `openai`, `model` | A2, A3 | ## Data parallelism @@ -211,23 +211,23 @@ click [Server Arguments](https://docs.sglang.io/advanced_features/server_argumen ## Speculative decoding -| Argument | Defaults | Options | Server supported | -|------------------------------------------------------------------|-----------|--------------------------|:----------------:| -| `--speculative-algorithm` | `None` | `EAGLE3`,
`NEXTN` | A2, A3 | -| `--speculative-draft-model-path`
`--speculative-draft-model` | `None` | Type: str | A2, A3 | -| `--speculative-draft-model-`
`revision` | `None` | Type: str | A2, A3 | -| `--speculative-draft-load-format` | `None` | `auto` | A2, A3 | -| `--speculative-num-steps` | `None` | Type: int | A2, A3 | -| `--speculative-eagle-topk` | `None` | Type: int | A2, A3 | -| `--speculative-num-draft-tokens` | `None` | Type: int | A2, A3 | -| `--speculative-accept-`
`threshold-single` | `1.0` | Type: float | Special for GPU | -| `--speculative-accept-`
`threshold-acc` | `1.0` | Type: float | Special for GPU | -| `--speculative-token-map` | `None` | Type: str | A2, A3 | -| `--speculative-attention-`
`mode` | `prefill` | `prefill`,
`decode` | A2, A3 | -| `--speculative-moe-runner-`
`backend` | `None` | `auto` | A2, A3 | -| `--speculative-moe-a2a-`
`backend` | `None` | `ascend_fuseep` | A2, A3 | -| `--speculative-draft-attention-backend` | `None` | `ascend` | A2, A3 | -| `--speculative-draft-model-quantization` | `None` | `unquant` | A2, A3 | +| Argument | Defaults | Options | Server supported | +|------------------------------------------------------------------|-----------|------------------------------------------------------------------|:----------------:| +| `--speculative-algorithm` | `None` | `EAGLE3`,
`NEXTN` | A2, A3 | +| `--speculative-draft-model-path`
`--speculative-draft-model` | `None` | Type: str | A2, A3 | +| `--speculative-draft-model-`
`revision` | `None` | Type: str,
`branch name`,
`tag name`,
`commit id` | A2, A3 | +| `--speculative-draft-load-format` | `auto` | `auto`,
`dummy` | A2, A3 | +| `--speculative-num-steps` | `None` | Type: int | A2, A3 | +| `--speculative-eagle-topk` | `None` | Type: int | A2, A3 | +| `--speculative-num-draft-tokens` | `None` | Type: int | A2, A3 | +| `--speculative-accept-`
`threshold-single` | `1.0` | Type: float | Special for GPU | +| `--speculative-accept-`
`threshold-acc` | `1.0` | Type: float | Special for GPU | +| `--speculative-token-map` | `None` | Type: str | A2, A3 | +| `--speculative-attention-`
`mode` | `prefill` | `prefill`,
`decode` | A2, A3 | +| `--speculative-moe-runner-`
`backend` | `None` | `auto` | A2, A3 | +| `--speculative-moe-a2a-`
`backend` | `None` | `ascend_fuseep` | A2, A3 | +| `--speculative-draft-attention-backend` | `None` | `ascend` | A2, A3 | +| `--speculative-draft-model-quantization` | `None` | `unquant` | A2, A3 | ## Ngram speculative decoding @@ -398,7 +398,7 @@ click [Server Arguments](https://docs.sglang.io/advanced_features/server_argumen | `--disaggregation-transfer-backend` | `mooncake` | `ascend` | A2, A3 | | `--disaggregation-bootstrap-port` | `8998` | Type: int | A2, A3 | | `--disaggregation-ib-device` | `None` | Type: str | Special for GPU | -| `--disaggregation-decode-`
`enable-offload-kvcache` | `False` | bool flag
(set to enable) | A2, A3 | +| `--disaggregation-decode-`
`enable-offload-kvcache` | `False` | `False` | Planned | | `--num-reserved-decode-tokens` | `512` | Type: int | A2, A3 | | `--disaggregation-decode-`
`polling-interval` | `1` | Type: int | A2, A3 | From 56c97c7738c471f0c11691166b5c6518a75de17b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=9C=8B=E6=B5=B7=E7=9A=84=E4=BA=BA?= <37328560+zhsurpass@users.noreply.github.com> Date: Mon, 13 Apr 2026 21:55:38 +0800 Subject: [PATCH 25/52] [NPU] update npu doc (#22697) Co-authored-by: zhsurpass --- .../platforms/ascend/ascend_npu_support_features.md | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/docs/platforms/ascend/ascend_npu_support_features.md b/docs/platforms/ascend/ascend_npu_support_features.md index c5311c83ee9b..54bdf7c979dc 100644 --- a/docs/platforms/ascend/ascend_npu_support_features.md +++ b/docs/platforms/ascend/ascend_npu_support_features.md @@ -404,12 +404,13 @@ click [Server Arguments](https://docs.sglang.io/advanced_features/server_argumen ## Encode prefill disaggregation -| Argument | Defaults | Options | Server supported | -|------------------------------|--------------------|----------------------------------------------------------------|:----------------:| -| `--encoder-only` | `False` | bool flag
(set to enable) | A2, A3 | -| `--language-only` | `False` | bool flag
(set to enable) | A2, A3 | -| `--encoder-transfer-backend` | `zmq_to_scheduler` | `zmq_to_scheduler`,
`zmq_to_tokenizer`,
`mooncake` | A2, A3 | -| `--encoder-urls` | `[]` | List[str] | A2, A3 | +| Argument | Defaults | Options | Server supported | +| --------------------------------------- | ------------------ | ---------------------------------------------------------------- |:----------------:| +| `--enable-adaptive-dispatch-to-encoder` | `False` | bool flag
(set to enable adaptively dispatch) | A2, A3 | +| `--encoder-only` | `False` | bool flag
(set to launch an encoder-only server) | A2, A3 | +| `--language-only` | `False` | bool flag
(set to load weights for the language model only) | A2, A3 | +| `--encoder-transfer-backend` | `zmq_to_scheduler` | `zmq_to_scheduler`,
`zmq_to_tokenizer`,
`mooncake` | A2, A3 | +| `--encoder-urls` | `[]` | List[str]
(List of encoder server urls) | A2, A3 | ## Custom weight loader From 2089ac86a7cd325a1a4150508e41f7d06e2ddb51 Mon Sep 17 00:00:00 2001 From: loading66 <102725781+loading66@users.noreply.github.com> Date: Mon, 13 Apr 2026 22:02:56 +0800 Subject: [PATCH 26/52] Improve parameters usage constraints for npu deployment (#22700) Co-authored-by: h30064329 --- .../ascend/ascend_npu_support_features.md | 48 +++++++++---------- 1 file changed, 23 insertions(+), 25 deletions(-) diff --git a/docs/platforms/ascend/ascend_npu_support_features.md b/docs/platforms/ascend/ascend_npu_support_features.md index 54bdf7c979dc..7f0127a76bb0 100644 --- a/docs/platforms/ascend/ascend_npu_support_features.md +++ b/docs/platforms/ascend/ascend_npu_support_features.md @@ -243,29 +243,29 @@ click [Server Arguments](https://docs.sglang.io/advanced_features/server_argumen ## Expert parallelism -| Argument | Defaults | Options | Server supported | -|-------------------------------------------------------|-----------|---------------------------------------------|:----------------:| -| `--expert-parallel-size`
`--ep-size`
`--ep` | `1` | Type: int | A2, A3 | -| `--moe-a2a-backend` | `none` | `none`,
`deepep`,
`ascend_fuseep` | A2, A3 | -| `--moe-runner-backend` | `auto` | `auto`, `triton` | A2, A3 | -| `--flashinfer-mxfp4-`
`moe-precision` | `default` | `default`,
`bf16` | Special for GPU | -| `--enable-flashinfer-`
`allreduce-fusion` | `False` | bool flag
(set to enable) | Special for GPU | -| `--deepep-mode` | `auto` | `normal`,
`low_latency`,
`auto` | A2, A3 | -| `--deepep-config` | `None` | Type: str | Special for GPU | -| `--ep-num-redundant-experts` | `0` | Type: int | A2, A3 | -| `--ep-dispatch-algorithm` | `None` | `static`,
`dynamic`,
`fake` | A2, A3 | -| `--init-expert-location` | `trivial` | Type: str | A2, A3 | -| `--enable-eplb` | `False` | bool flag
(set to enable) | A2, A3 | -| `--eplb-algorithm` | `deepseek`| Type: str | A2, A3 | -| `--eplb-rebalance-num-iterations` | `1000` | Type: int | A2, A3 | -| `--eplb-rebalance-layers-`
`per-chunk` | `None` | Type: int | A2, A3 | -| `--eplb-min-rebalancing-`
`utilization-threshold` | `1.0` | Type: float | A2, A3 | -| `--expert-distribution-`
`recorder-mode` | `None` | `stat`,
`stat_approx`,
`per_pass`,
`per_token` | A2, A3 | -| `--expert-distribution-`
`recorder-buffer-size` | `None` | Type: int | A2, A3 | -| `--enable-expert-distribution-`
`metrics` | `False` | bool flag (set to enable) | A2, A3 | -| `--moe-dense-tp-size` | `None` | `1` | A2, A3 | -| `--elastic-ep-backend` | `None` | `none`, `mooncake` | Special for GPU | -| `--mooncake-ib-device` | `None` | Type: str | Special for GPU | +| Argument | Defaults | Options | Server supported | +|-------------------------------------------------------|-----------|---------------------------------------------------------------------------|:----------------:| +| `--expert-parallel-size`
`--ep-size`
`--ep` | `1` | Type: int | A2, A3 | +| `--moe-a2a-backend` | `none` | `none`,
`deepep`,
`ascend_fuseep`(It is incompatible with eplb) | A2, A3 | +| `--moe-runner-backend` | `auto` | `auto`, `triton` | A2, A3 | +| `--flashinfer-mxfp4-`
`moe-precision` | `default` | `default`,
`bf16` | Special for GPU | +| `--enable-flashinfer-`
`allreduce-fusion` | `False` | bool flag
(set to enable) | Special for GPU | +| `--deepep-mode` | `auto` | `normal`,
`low_latency`,
`auto` | A2, A3 | +| `--deepep-config` | `None` | Type: str | Special for GPU | +| `--ep-num-redundant-experts` | `0` | Type: int | A2, A3 | +| `--ep-dispatch-algorithm` | `None` | `static`,
`dynamic`,
`fake` | A2, A3 | +| `--init-expert-location` | `trivial` | Type: str | A2, A3 | +| `--enable-eplb` | `False` | bool flag
(set to enable) | A2, A3 | +| `--eplb-algorithm` | `deepseek`| Type: str | A2, A3 | +| `--eplb-rebalance-num-iterations` | `1000` | Type: int | A2, A3 | +| `--eplb-rebalance-layers-`
`per-chunk` | `None` | Type: int | A2, A3 | +| `--eplb-min-rebalancing-`
`utilization-threshold` | `1.0` | Type: float | A2, A3 | +| `--expert-distribution-`
`recorder-mode` | `None` | `stat`,
`stat_approx`,
`per_pass`,
`per_token` | A2, A3 | +| `--expert-distribution-`
`recorder-buffer-size` | `None` | Type: int | A2, A3 | +| `--enable-expert-distribution-`
`metrics` | `False` | bool flag (set to enable) | A2, A3 | +| `--moe-dense-tp-size` | `None` | `1` | A2, A3 | +| `--elastic-ep-backend` | `None` | `none`, `mooncake` | Special for GPU | +| `--mooncake-ib-device` | `None` | Type: str | Special for GPU | ## Mamba Cache @@ -436,8 +436,6 @@ click [Server Arguments](https://docs.sglang.io/advanced_features/server_argumen | Argument | Defaults | Options | Server supported | |-----------------------------------------------|----------|--------------------------------|:----------------:| -| `--mm-max-concurrent-calls` | `32` | Type: int | A2, A3 | -| `--mm-per-request-timeout` | `10.0` | Type: float | A2, A3 | | `--enable-broadcast-mm-`
`inputs-process` | `False` | bool flag
(set to enable) | A2, A3 | | `--mm-process-config` | `None` | Type: JSON / Dict | A2, A3 | | `--mm-enable-dp-encoder` | `False` | bool flag
(set to enable) | A2, A3 | From 8d904e50f2d937991c424f69fa1c052836d11bd7 Mon Sep 17 00:00:00 2001 From: Liwansi <62291011+Liwansi@users.noreply.github.com> Date: Mon, 13 Apr 2026 22:20:17 +0800 Subject: [PATCH 27/52] [NPU]qwen3-8b and 32b md bugfix (#22687) --- docs/platforms/ascend/ascend_npu_best_practice.md | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/docs/platforms/ascend/ascend_npu_best_practice.md b/docs/platforms/ascend/ascend_npu_best_practice.md index 7e395aa6c3cf..8e011e6d5ad1 100644 --- a/docs/platforms/ascend/ascend_npu_best_practice.md +++ b/docs/platforms/ascend/ascend_npu_best_practice.md @@ -2397,7 +2397,6 @@ LOCAL_HOST2=`hostname -I|awk -F " " '{print$2}'` echo "${LOCAL_HOST1}" echo "${LOCAL_HOST2}" -export HCCL_BUFFSIZE=400 export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo export HCCL_OP_EXPANSION_MODE="AIV" @@ -2410,7 +2409,7 @@ python -m sglang.launch_server --model-path $MODEL_PATH \ --max-running-requests 16 \ --disable-radix-cache \ --speculative-draft-model-quantization unquant \ - --speculative-algorithm EAGLE3 --speculative-draft-model-path xxx --speculative-draft-model-quantization unquant \ + --speculative-algorithm EAGLE3 --speculative-draft-model-path xxx \ --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 \ --chunked-prefill-size -1 --max-prefill-tokens 16384 \ --tp-size 4 --mem-fraction-static 0.843 --cuda-graph-bs 1 4 8 16 --dtype bfloat16 @@ -2470,7 +2469,6 @@ LOCAL_HOST2=`hostname -I|awk -F " " '{print$2}'` echo "${LOCAL_HOST1}" echo "${LOCAL_HOST2}" -export HCCL_BUFFSIZE=400 export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo export HCCL_OP_EXPANSION_MODE="AIV" @@ -2483,7 +2481,7 @@ python -m sglang.launch_server --model-path $MODEL_PATH \ --max-running-requests 16 \ --disable-radix-cache \ --speculative-draft-model-quantization unquant \ - --speculative-algorithm EAGLE3 --speculative-draft-model-path xxx --speculative-draft-model-quantization unquant \ + --speculative-algorithm EAGLE3 --speculative-draft-model-path xxx \ --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 \ --chunked-prefill-size -1 --max-prefill-tokens 16384 \ --tp-size 4 --mem-fraction-static 0.843 --cuda-graph-bs 1 4 10 15 16 --dtype bfloat16 @@ -2543,7 +2541,6 @@ LOCAL_HOST2=`hostname -I|awk -F " " '{print$2}'` echo "${LOCAL_HOST1}" echo "${LOCAL_HOST2}" -export HCCL_BUFFSIZE=400 export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo export HCCL_OP_EXPANSION_MODE="AIV" @@ -2556,7 +2553,7 @@ python -m sglang.launch_server --model-path $MODEL_PATH \ --max-running-requests 16 \ --disable-radix-cache \ --speculative-draft-model-quantization unquant \ - --speculative-algorithm EAGLE3 --speculative-draft-model-path xxx --speculative-draft-model-quantization unquant \ + --speculative-algorithm EAGLE3 --speculative-draft-model-path xxx \ --speculative-num-steps 4 --speculative-eagle-topk 1 --speculative-num-draft-tokens 5 \ --chunked-prefill-size -1 --max-prefill-tokens 16384 \ --tp-size 2 --mem-fraction-static 0.894 --cuda-graph-bs 1 2 4 6 9 10 15 16 --dtype bfloat16 @@ -2616,7 +2613,6 @@ LOCAL_HOST2=`hostname -I|awk -F " " '{print$2}'` echo "${LOCAL_HOST1}" echo "${LOCAL_HOST2}" -export HCCL_BUFFSIZE=400 export HCCL_SOCKET_IFNAME=lo export GLOO_SOCKET_IFNAME=lo export HCCL_OP_EXPANSION_MODE="AIV" @@ -2629,7 +2625,7 @@ python -m sglang.launch_server --model-path $MODEL_PATH \ --max-running-requests 16 \ --disable-radix-cache \ --speculative-draft-model-quantization unquant \ - --speculative-algorithm EAGLE3 --speculative-draft-model-path xxx --speculative-draft-model-quantization unquant \ + --speculative-algorithm EAGLE3 --speculative-draft-model-path xxx \ --speculative-num-steps 4 --speculative-eagle-topk 1 --speculative-num-draft-tokens 5 \ --chunked-prefill-size -1 --max-prefill-tokens 16384 \ --tp-size 2 --mem-fraction-static 0.894 --cuda-graph-bs 1 5 15 16 --dtype bfloat16 From b6a91b1afefc3caac396f720aa9beeb798aabf75 Mon Sep 17 00:00:00 2001 From: jianzhao-xu <978716854@qq.com> Date: Mon, 13 Apr 2026 22:42:34 +0800 Subject: [PATCH 28/52] [NPU] --attn-cp-size --init-expert-location --eplb-algorithm parameter docs update (#22704) Co-authored-by: Jianzhao Xu --- docs/platforms/ascend/ascend_npu_support_features.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/platforms/ascend/ascend_npu_support_features.md b/docs/platforms/ascend/ascend_npu_support_features.md index 7f0127a76bb0..a46f4854a0a5 100644 --- a/docs/platforms/ascend/ascend_npu_support_features.md +++ b/docs/platforms/ascend/ascend_npu_support_features.md @@ -87,7 +87,7 @@ click [Server Arguments](https://docs.sglang.io/advanced_features/server_argumen | `--device` | `None` | Type: str | A2, A3 | | `--tensor-parallel-size`
`--tp-size` | `1` | Type: int | A2, A3 | | `--pipeline-parallel-size`
`--pp-size` | `1` | Type: int; Currently `2` not supported | A2, A3 | -| `--attention-context-parallel-size`
`--attn-cp-size` | `1` | Type: int | A2, A3 | +| `--attention-context-parallel-size`
`--attn-cp-size` | `1` | Type: int; must be equal to --tp-size | A2, A3 | | `--moe-data-parallel-size`
`--moe-dp-size` | `1` | Type: int | A2, A3 | | `--pp-max-micro-batch-size` | `None` | Type: int | A2, A3 | | `--pp-async-batch-depth` | `None` | Type: int | A2, A3 | @@ -254,9 +254,9 @@ click [Server Arguments](https://docs.sglang.io/advanced_features/server_argumen | `--deepep-config` | `None` | Type: str | Special for GPU | | `--ep-num-redundant-experts` | `0` | Type: int | A2, A3 | | `--ep-dispatch-algorithm` | `None` | `static`,
`dynamic`,
`fake` | A2, A3 | -| `--init-expert-location` | `trivial` | Type: str | A2, A3 | +| `--init-expert-location` | `trivial` | `trivial`,
``,
``,
`` | A2, A3 | | `--enable-eplb` | `False` | bool flag
(set to enable) | A2, A3 | -| `--eplb-algorithm` | `deepseek`| Type: str | A2, A3 | +| `--eplb-algorithm` | `deepseek`| `auto`,
`deepseek` | A2, A3 | | `--eplb-rebalance-num-iterations` | `1000` | Type: int | A2, A3 | | `--eplb-rebalance-layers-`
`per-chunk` | `None` | Type: int | A2, A3 | | `--eplb-min-rebalancing-`
`utilization-threshold` | `1.0` | Type: float | A2, A3 | From c6403a11cbfc76c245ede585bfd66a89890f52e0 Mon Sep 17 00:00:00 2001 From: chx96642264 Date: Mon, 13 Apr 2026 22:50:48 +0800 Subject: [PATCH 29/52] Modify the optional values and constraints of parameter. (#22705) --- .../ascend/ascend_npu_support_features.md | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/docs/platforms/ascend/ascend_npu_support_features.md b/docs/platforms/ascend/ascend_npu_support_features.md index a46f4854a0a5..9b17a090becf 100644 --- a/docs/platforms/ascend/ascend_npu_support_features.md +++ b/docs/platforms/ascend/ascend_npu_support_features.md @@ -274,22 +274,22 @@ click [Server Arguments](https://docs.sglang.io/advanced_features/server_argumen | `--max-mamba-cache-size` | `None` | Type: int | A2, A3 | | `--mamba-ssm-dtype` | `float32` | `float32`,
`bfloat16`,
`float16` | A2, A3 | | `--mamba-full-memory-ratio` | `0.9` | Type: float | A2, A3 | -| `--mamba-scheduler-strategy` | `auto` | Only `auto`, `no_buffer` supported | A2, A3 | +| `--mamba-scheduler-strategy` | `auto` | `auto`,
`no_buffer`,
`extra_buffer` | A2, A3 | | `--mamba-track-interval` | `256` | Type: int | A2, A3 | ## Hierarchical cache -| Argument | Defaults | Options | Server supported | -|-------------------------------------------------|-----------------|---------------------------------------------------------------------|:----------------:| -| `--enable-hierarchical-`
`cache` | `False` | bool flag
(set to enable) | A2, A3 | -| `--hicache-ratio` | `2.0` | Type: float | A2, A3 | -| `--hicache-size` | `0` | Type: int | A2, A3 | -| `--hicache-write-policy` | `write_through` | Currently only `write_back` supported | A2, A3 | -| `--hicache-io-backend` | `kernel` | `kernel_ascend`,
`direct` | A2, A3 | -| `--hicache-mem-layout` | `layer_first` | `page_first_direct`,
`page_first_kv_split` | A2, A3 | -| `--hicache-storage-`
`backend` | `None` | `file` | A2, A3 | -| `--hicache-storage-`
`prefetch-policy` | `best_effort` | `best_effort`,
`wait_complete`,
`timeout` | Special for GPU | -| `--hicache-storage-`
`backend-extra-config` | `None` | Type: str | Special for GPU | +| Argument | Defaults | Options | Server supported | +|-------------------------------------------------|-----------------|-------------------------------------------------------------------------------|:----------------:| +| `--enable-hierarchical-`
`cache` | `False` | bool flag
(set to enable).
Currently, mamba cache is not supported. | A2, A3 | +| `--hicache-ratio` | `2.0` | Type: float | A2, A3 | +| `--hicache-size` | `0` | Type: int | A2, A3 | +| `--hicache-write-policy` | `write_through` | Currently only `write_back` supported | A2, A3 | +| `--hicache-io-backend` | `kernel` | `kernel_ascend`,
`direct` | A2, A3 | +| `--hicache-mem-layout` | `layer_first` | `page_first_direct`,
`page_first_kv_split` | A2, A3 | +| `--hicache-storage-`
`backend` | `None` | `file` | A2, A3 | +| `--hicache-storage-`
`prefetch-policy` | `best_effort` | `best_effort`,
`wait_complete`,
`timeout` | Special for GPU | +| `--hicache-storage-`
`backend-extra-config` | `None` | Type: str | Special for GPU | ## LMCache From 13a4aafdbe693da7a7bbc80952383c1bcabf39a1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=9C=8B=E6=B5=B7=E7=9A=84=E4=BA=BA?= <37328560+zhsurpass@users.noreply.github.com> Date: Mon, 13 Apr 2026 22:53:24 +0800 Subject: [PATCH 30/52] [NPU] update glm5 running guide (#22712) --- docs/platforms/ascend/ascend_npu_glm5_examples.md | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/docs/platforms/ascend/ascend_npu_glm5_examples.md b/docs/platforms/ascend/ascend_npu_glm5_examples.md index f613e8956583..6c3e1500cbd2 100644 --- a/docs/platforms/ascend/ascend_npu_glm5_examples.md +++ b/docs/platforms/ascend/ascend_npu_glm5_examples.md @@ -53,10 +53,16 @@ docker run -itd --shm-size=16g --privileged=true --name ${NAME} \ swr.cn-southwest-2.myhuaweicloud.com/base_image/dockerhub/lmsysorg/sglang:${TAG} ``` -Note: Using this image, you need to update transformers to main branch +### Best Practices +Note: Using this image for **best practices**, you need to update transformers to version 5.3.0 ``` shell # reinstall transformers -pip install git+https://github.com/huggingface/transformers.git + +# Install transformers version 5.3.0 from PyPI +pip install transformers==5.3.0 + +# Install from GitHub v5.3.0 tag from GitHub +pip install git+https://github.com/huggingface/transformers.git@v5.3.0 ``` ## Deployment From 39810762d2cd4ad6a513566b54b5f50fc6be328d Mon Sep 17 00:00:00 2001 From: Asish Kumar <87874775+officialasishkumar@users.noreply.github.com> Date: Mon, 13 Apr 2026 22:15:45 +0530 Subject: [PATCH 31/52] fix: use describe mode for SGLang version detection (#22600) Signed-off-by: Asish Kumar --- 3rdparty/amd/wheel/sglang/pyproject.toml | 4 +- python/pyproject.toml | 2 +- python/pyproject_cpu.toml | 4 +- python/pyproject_npu.toml | 4 +- python/pyproject_other.toml | 4 +- python/pyproject_xpu.toml | 2 +- .../unit/tools/test_get_version_tag.py | 87 +++++++++++++++++++ 7 files changed, 101 insertions(+), 6 deletions(-) create mode 100644 test/registered/unit/tools/test_get_version_tag.py diff --git a/3rdparty/amd/wheel/sglang/pyproject.toml b/3rdparty/amd/wheel/sglang/pyproject.toml index 9b9c24fd22b8..99a99c0b5050 100644 --- a/3rdparty/amd/wheel/sglang/pyproject.toml +++ b/3rdparty/amd/wheel/sglang/pyproject.toml @@ -210,4 +210,6 @@ exclude = [ [tool.setuptools_scm] root = ".." version_file = "sglang/_version.py" -git_describe_command = ["python3", "python/tools/get_version_tag.py", "--tag-only"] +git_describe_command = ["python3", "python/tools/get_version_tag.py"] +# Allow editable installs even when .git metadata is not available. +fallback_version = "0.0.0.dev0" diff --git a/python/pyproject.toml b/python/pyproject.toml index 7eeb44f59783..80ba96cc6cdf 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -205,7 +205,7 @@ exclude = [ [tool.setuptools_scm] root = ".." version_file = "sglang/_version.py" -git_describe_command = ["python3", "python/tools/get_version_tag.py", "--tag-only"] +git_describe_command = ["python3", "python/tools/get_version_tag.py"] # Allow editable installs even when .git metadata is not available. fallback_version = "0.0.0.dev0" diff --git a/python/pyproject_cpu.toml b/python/pyproject_cpu.toml index 550780ffc727..1597212e99f5 100644 --- a/python/pyproject_cpu.toml +++ b/python/pyproject_cpu.toml @@ -127,4 +127,6 @@ exclude = [ [tool.setuptools_scm] root = ".." version_file = "sglang/_version.py" -git_describe_command = ["python3", "python/tools/get_version_tag.py", "--tag-only"] +git_describe_command = ["python3", "python/tools/get_version_tag.py"] +# Allow editable installs even when .git metadata is not available. +fallback_version = "0.0.0.dev0" diff --git a/python/pyproject_npu.toml b/python/pyproject_npu.toml index 704df6e3e908..e2bbfd798840 100644 --- a/python/pyproject_npu.toml +++ b/python/pyproject_npu.toml @@ -146,4 +146,6 @@ exclude = [ [tool.setuptools_scm] root = ".." version_file = "sglang/_version.py" -git_describe_command = ["python3", "python/tools/get_version_tag.py", "--tag-only"] +git_describe_command = ["python3", "python/tools/get_version_tag.py"] +# Allow editable installs even when .git metadata is not available. +fallback_version = "0.0.0.dev0" diff --git a/python/pyproject_other.toml b/python/pyproject_other.toml index b7f4411cffc1..0e6da96c18e9 100755 --- a/python/pyproject_other.toml +++ b/python/pyproject_other.toml @@ -214,4 +214,6 @@ exclude = [ [tool.setuptools_scm] root = ".." version_file = "sglang/_version.py" -git_describe_command = ["python3", "python/tools/get_version_tag.py", "--tag-only"] +git_describe_command = ["python3", "python/tools/get_version_tag.py"] +# Allow editable installs even when .git metadata is not available. +fallback_version = "0.0.0.dev0" diff --git a/python/pyproject_xpu.toml b/python/pyproject_xpu.toml index 821dc88e5a75..8e15e14c8008 100644 --- a/python/pyproject_xpu.toml +++ b/python/pyproject_xpu.toml @@ -154,6 +154,6 @@ exclude = [ [tool.setuptools_scm] root = ".." version_file = "sglang/_version.py" -git_describe_command = ["python3", "python/tools/get_version_tag.py", "--tag-only"] +git_describe_command = ["python3", "python/tools/get_version_tag.py"] # Allow editable installs even when .git metadata is not available. fallback_version = "0.0.0.dev0" diff --git a/test/registered/unit/tools/test_get_version_tag.py b/test/registered/unit/tools/test_get_version_tag.py new file mode 100644 index 000000000000..d64dac347386 --- /dev/null +++ b/test/registered/unit/tools/test_get_version_tag.py @@ -0,0 +1,87 @@ +import importlib.util +import sys +import unittest +from pathlib import Path +from unittest.mock import patch + +REPO_ROOT = Path(__file__).resolve().parents[4] +CI_REGISTER_PATH = REPO_ROOT / "python" / "sglang" / "test" / "ci" / "ci_register.py" +VERSION_HELPER_PATH = REPO_ROOT / "python" / "tools" / "get_version_tag.py" +PYPROJECT_PATHS = [ + REPO_ROOT / "python" / "pyproject.toml", + REPO_ROOT / "python" / "pyproject_cpu.toml", + REPO_ROOT / "python" / "pyproject_npu.toml", + REPO_ROOT / "python" / "pyproject_other.toml", + REPO_ROOT / "python" / "pyproject_xpu.toml", + REPO_ROOT / "3rdparty" / "amd" / "wheel" / "sglang" / "pyproject.toml", +] +DESCRIBE_COMMAND = ( + 'git_describe_command = ["python3", "python/tools/get_version_tag.py"]' +) +TAG_ONLY_DESCRIBE_COMMAND = ( + 'git_describe_command = ["python3", "python/tools/get_version_tag.py", ' + '"--tag-only"]' +) +FALLBACK_VERSION = 'fallback_version = "0.0.0.dev0"' + + +def _load_module(name, path): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +register_cpu_ci = _load_module("ci_register", CI_REGISTER_PATH).register_cpu_ci +register_cpu_ci(est_time=2, suite="stage-a-test-cpu") + + +class TestGetVersionTag(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.version_helper = _load_module("get_version_tag", VERSION_HELPER_PATH) + + def test_parse_version_tuple_sorts_stable_above_rc_and_post_above_stable(self): + tags = ["v0.5.10rc0", "v0.5.9", "v0.5.10.post1", "v0.5.10"] + + self.assertEqual( + sorted(tags, key=self.version_helper.parse_version_tuple, reverse=True), + ["v0.5.10.post1", "v0.5.10", "v0.5.10rc0", "v0.5.9"], + ) + + def test_exact_version_tag_takes_precedence_over_latest_tag(self): + with patch.object( + self.version_helper, "get_exact_version_tag", return_value="v0.5.9" + ), patch.object( + self.version_helper, "get_latest_version_tag_describe" + ) as latest_describe: + self.assertEqual(self.version_helper.get_version_describe(), "v0.5.9") + + latest_describe.assert_not_called() + + def test_pyprojects_use_describe_mode_for_setuptools_scm(self): + for path in PYPROJECT_PATHS: + with self.subTest(path=path): + content = path.read_text() + self.assertIn(DESCRIBE_COMMAND, content) + self.assertNotIn(TAG_ONLY_DESCRIBE_COMMAND, content) + self.assertIn(FALLBACK_VERSION, content) + + def test_tag_only_cli_mode_remains_available_for_callers_that_need_latest_tag(self): + with patch.object( + sys, "argv", ["get_version_tag.py", "--tag-only"] + ), patch.object( + self.version_helper, "get_latest_version_tag", return_value="v0.5.10" + ), patch.object( + self.version_helper, "get_version_describe" + ) as version_describe, patch( + "builtins.print" + ) as print_mock: + self.version_helper.main() + + version_describe.assert_not_called() + print_mock.assert_called_once_with("v0.5.10") + + +if __name__ == "__main__": + unittest.main() From 90ef8ce54de7dc4eb41f4ad1b6ba5abbad583769 Mon Sep 17 00:00:00 2001 From: Mohammad Miadh Angkad Date: Tue, 14 Apr 2026 00:48:22 +0800 Subject: [PATCH 32/52] [Docker] Remove flashinfer cache copy (#22653) --- docker/Dockerfile | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index 487594aa21d5..0b79785c98ea 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -428,8 +428,7 @@ COPY --from=deepep_builder /build/DeepEP /sgl-workspace/DeepEP RUN --mount=type=cache,target=/root/.cache/pip \ pip install /tmp/wheels/deepep/*.whl && rm -rf /tmp/wheels/deepep -# Copy flashinfer cache and jit-cache package (if installed) -COPY --from=flashinfer_cache /root/.cache/flashinfer /root/.cache/flashinfer +# Copy flashinfer jit-cache package (if installed) COPY --from=flashinfer_cache /flashinfer_jit_output/ /usr/local/lib/python3.12/dist-packages/ # Copy dev tools From f51ce2c92f9ced73db0678f607919ed6f12c6576 Mon Sep 17 00:00:00 2001 From: R0CKSTAR Date: Tue, 14 Apr 2026 03:41:07 +0800 Subject: [PATCH 33/52] Update CODEOWNERS for musa/mlx (#22593) Signed-off-by: Xiaodong Ye --- .github/CODEOWNERS | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 83bcad1b476c..3ad22947b6e1 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -23,6 +23,8 @@ /python/sglang/srt/eplb @fzyzcjy @ch-wan /python/sglang/srt/function_call @CatherineSue @JustinTong0323 /python/sglang/srt/grpc @CatherineSue @slin1237 +/python/sglang/srt/hardware_backend/mlx @yeahdongcn +/python/sglang/srt/hardware_backend/musa @yeahdongcn /python/sglang/srt/hardware_backend/npu @ping1jing2 @iforgetmyname /python/sglang/srt/hardware_backend/npu/quantization @OrangeRedeng @TamirBaydasov @iforgetmyname /python/sglang/srt/layers @merrymercy @Ying1123 @Fridge003 @ispobock @HaiShaw @ch-wan @BBuf @Edwardf0t1 From 6b2bf66cd9cd0448b0e9f3af8a54e9e10686fdf2 Mon Sep 17 00:00:00 2001 From: ishandhanani <82981111+ishandhanani@users.noreply.github.com> Date: Mon, 13 Apr 2026 15:10:07 -0500 Subject: [PATCH 34/52] fix[glm4.7 flash]: properly detect `gfx95_quant_format` (#22720) --- python/sglang/srt/models/glm4_moe_lite.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/sglang/srt/models/glm4_moe_lite.py b/python/sglang/srt/models/glm4_moe_lite.py index 03ff3d201d45..5b08f1f4581d 100644 --- a/python/sglang/srt/models/glm4_moe_lite.py +++ b/python/sglang/srt/models/glm4_moe_lite.py @@ -403,6 +403,8 @@ def __init__( config.hidden_size, eps=config.rms_norm_eps ) + self._gfx95_quant_format = self._detect_gfx95_quant_format() + self.layer_communicator = LayerCommunicator( layer_scatter_modes=self.layer_scatter_modes, input_layernorm=self.input_layernorm, From ff13dfee45df4be0f85ea0da3cd630c5a4912eaf Mon Sep 17 00:00:00 2001 From: Kurt Shuster Date: Mon, 13 Apr 2026 17:19:30 -0400 Subject: [PATCH 35/52] [lora][moe] Virtual experts for LoRA MoE (#22122) Co-authored-by: Yusheng Su --- .../fused_moe_triton_kernels.py | 49 +- .../srt/layers/moe/moe_runner/runner.py | 25 +- python/sglang/srt/lora/layers.py | 2 + python/sglang/srt/lora/lora_manager.py | 3 +- python/sglang/srt/lora/lora_moe_runners.py | 343 +++++++-- python/sglang/srt/lora/triton_ops/__init__.py | 2 + .../srt/lora/triton_ops/virtual_experts.py | 662 ++++++++++++++++++ python/sglang/srt/server_args.py | 14 + test/registered/lora/test_lora_moe_runner.py | 153 ++++ .../lora/test_marlin_lora_correctness.py | 1 + 10 files changed, 1148 insertions(+), 106 deletions(-) create mode 100644 python/sglang/srt/lora/triton_ops/virtual_experts.py diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py index b80e1827baea..6a2b46d99786 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py @@ -335,6 +335,7 @@ def fused_moe_kernel( sorted_token_ids_ptr, expert_ids_ptr, num_tokens_post_padded_ptr, + add_mask_ptr, # Matrix dimensions N, K, @@ -377,6 +378,7 @@ def fused_moe_kernel( c_sorted: tl.constexpr, filter_expert: tl.constexpr, swap_ab: tl.constexpr, + FUSE_ADD_TO_OUTPUT: tl.constexpr, FUSE_SUM_ALL_REDUCE: tl.constexpr, ROUTER_TOPK: tl.constexpr, ): @@ -441,18 +443,20 @@ def fused_moe_kernel( # ----------------------------------------------------------- # Write back zeros to the output when the expert is not # in the current expert parallel rank. - write_zeros_to_output( - c_ptr, - stride_cm, - stride_cn, - pid_n, - N, - offs_token, - token_mask, - BLOCK_SIZE_M, - BLOCK_SIZE_N, - compute_type, - ) + if not FUSE_ADD_TO_OUTPUT: + # skip the zero-write to preserve existing values. + write_zeros_to_output( + c_ptr, + stride_cm, + stride_cn, + pid_n, + N, + offs_token, + token_mask, + BLOCK_SIZE_M, + BLOCK_SIZE_N, + compute_type, + ) return offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N @@ -604,7 +608,15 @@ def fused_moe_kernel( # Write back the block of the output offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) - if FUSE_SUM_ALL_REDUCE: + if FUSE_ADD_TO_OUTPUT: + # Accumulate into existing output with per-token mask. + offs_token_out = offs_token // ROUTER_TOPK + add_mask = tl.load(add_mask_ptr + offs_token_out, mask=token_mask, other=False) + c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :] + c_mask = token_mask[:, None] & add_mask[:, None] & (offs_cn[None, :] < N) + existing = tl.load(c_ptrs, mask=c_mask, other=0.0) + tl.store(c_ptrs, existing + accumulator, mask=c_mask) + elif FUSE_SUM_ALL_REDUCE: offs_token_out = offs_token // ROUTER_TOPK c_ptrs = ( c_ptr + stride_cm * offs_token_out[:, None] + stride_cn * offs_cn[None, :] @@ -717,6 +729,8 @@ def invoke_fused_moe_kernel( filter_expert: bool = True, fuse_sum_all_reduce: bool = False, router_topk: int = 1, + fuse_add_to_output: bool = False, + add_output_mask: Optional[torch.Tensor] = None, ) -> None: assert topk_weights.stride(1) == 1 assert sorted_token_ids.stride(0) == 1 @@ -786,6 +800,13 @@ def invoke_fused_moe_kernel( if fuse_sum_all_reduce: assert not c_sorted, "fuse_sum_all_reduce only supports c_sorted=False" + if fuse_add_to_output: + assert ( + not fuse_sum_all_reduce + ), "fuse_add_to_output and fuse_sum_all_reduce are mutually exclusive" + assert ( + add_output_mask is not None + ), "add_output_mask required when fuse_add_to_output=True" if ( (use_int8_w8a16 or use_int4_w4a16) @@ -870,6 +891,7 @@ def invoke_fused_moe_kernel( sorted_token_ids, expert_ids, num_tokens_post_padded, + add_output_mask, B.shape[1], B.shape[2] - padded_size, sorted_token_ids.shape[0], @@ -901,6 +923,7 @@ def invoke_fused_moe_kernel( c_sorted=c_sorted, filter_expert=filter_expert, swap_ab=swap_ab, + FUSE_ADD_TO_OUTPUT=fuse_add_to_output, FUSE_SUM_ALL_REDUCE=fuse_sum_all_reduce, ROUTER_TOPK=router_topk, **config, diff --git a/python/sglang/srt/layers/moe/moe_runner/runner.py b/python/sglang/srt/layers/moe/moe_runner/runner.py index 9dcee54da7d2..09b04c25080d 100644 --- a/python/sglang/srt/layers/moe/moe_runner/runner.py +++ b/python/sglang/srt/layers/moe/moe_runner/runner.py @@ -21,7 +21,6 @@ from sglang.srt.layers.moe.utils import MoeRunnerBackend from sglang.srt.lora.lora_moe_runners import LoRAHooks - logger = logging.getLogger(__name__) @@ -98,9 +97,6 @@ def run( assert self.runner_core is not None def _maybe_build_lora_hooks(_runner_input: Any) -> LoRAHooks: - if not self.lora_enabled or lora_info is None: - return None - from sglang.srt.layers.moe.token_dispatcher.base import DispatchOutput from sglang.srt.lora.lora_moe_runners import build_lora_hooks @@ -109,19 +105,16 @@ def _maybe_build_lora_hooks(_runner_input: Any) -> LoRAHooks: _runner_input.hidden_states, _runner_input.topk_output.topk_ids, ) - elif hasattr(_runner_input, "topk_ids"): - hidden_states, topk_ids = ( - _runner_input.hidden_states, - _runner_input.topk_ids, - ) else: - return None - - return build_lora_hooks( - hidden_states, - lora_info, - topk_ids, - ) + hidden_states = _runner_input.hidden_states + topk_ids = getattr(_runner_input, "topk_ids", None) + if self.lora_enabled and lora_info is not None: + return build_lora_hooks( + hidden_states, + lora_info, + topk_ids, + ) + return None # Runners that handle dispatch_output directly (e.g., MarlinRunnerCore) # bypass the pre-permute step and do their own alignment internally. diff --git a/python/sglang/srt/lora/layers.py b/python/sglang/srt/lora/layers.py index 6c79e9063f4e..7b184e71a187 100644 --- a/python/sglang/srt/lora/layers.py +++ b/python/sglang/srt/lora/layers.py @@ -797,6 +797,7 @@ def __init__( super().__init__(base_layer, lora_backend) self.experts_shared_outer_loras: bool = False + self.lora_use_virtual_experts: bool = False self.quant_method = base_layer.quant_method self.tp_size = getattr(base_layer, "moe_tp_size", 1) @@ -903,6 +904,7 @@ def _get_lora_info(self): tp_size=self.tp_size, tp_rank=self.tp_rank, hidden_size=getattr(self.base_layer, "hidden_size", 0), + lora_use_virtual_experts=self.lora_use_virtual_experts, ) def forward(self, hidden_states: torch.Tensor, topk_output: TopKOutput, **kwargs): diff --git a/python/sglang/srt/lora/lora_manager.py b/python/sglang/srt/lora/lora_manager.py index 9f3ef601ed5c..58d00551a283 100644 --- a/python/sglang/srt/lora/lora_manager.py +++ b/python/sglang/srt/lora/lora_manager.py @@ -86,6 +86,7 @@ def __init__( self._experts_shared_outer_override: Optional[bool] = ( server_args.experts_shared_outer_loras ) + self.lora_use_virtual_experts: bool = server_args.lora_use_virtual_experts self.lora_strict_loading: bool = getattr( server_args, "lora_strict_loading", False ) @@ -763,7 +764,6 @@ def init_lora_modules(self): lora_module = self.set_lora_module(module_name, module) self.embed_tokens_module = lora_module continue - # Handle lm_head if "lm_head" in module_name and "lm_head" in self.target_modules: if isinstance(module, ParallelLMHead) and not isinstance( @@ -808,4 +808,5 @@ def init_lora_modules(self): layer_id = get_layer_id(module_name) lora_module = self.set_lora_module(module_name, module) lora_module.experts_shared_outer_loras = self.experts_shared_outer_loras + lora_module.lora_use_virtual_experts = self.lora_use_virtual_experts self.lora_modules[layer_id][module_name] = lora_module diff --git a/python/sglang/srt/lora/lora_moe_runners.py b/python/sglang/srt/lora/lora_moe_runners.py index f07d762eef81..983064154ccf 100644 --- a/python/sglang/srt/lora/lora_moe_runners.py +++ b/python/sglang/srt/lora/lora_moe_runners.py @@ -34,6 +34,7 @@ _is_cuda = is_cuda() _is_hip = is_hip() +_is_hip = is_hip() _is_xpu = is_xpu() if _is_cuda or _is_hip or _is_xpu: @@ -63,6 +64,112 @@ def _get_moe_lora_block_config(max_lora_rank: int) -> dict: _SPARSITY_FACTOR = 8 +def _naive_moe_lora_align_block_size( + topk_ids: torch.Tensor, + seg_indptr: torch.Tensor, + req_to_lora: torch.Tensor, + num_experts: int, + block_size_m: int, + max_loras: int, + max_num_tokens_padded: int, + max_num_m_blocks: int, + adapter_enabled: torch.Tensor, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Construct LoRA token-expert alignment on CPU for small batches. + + When the number of tokens is very small, the overhead of launching the + CUDA-based moe_lora_align_block_size kernel exceeds the actual + computation. This function builds the same data structures using simple + Python loops on CPU and transfers the result to GPU in one shot. + """ + M, top_k = topk_ids.shape + num_valid_tokens = M * top_k + + sorted_token_ids = torch.full( + (max_loras * max_num_tokens_padded,), + num_valid_tokens, + dtype=torch.int32, + ) + expert_ids_out = torch.full((max_loras * max_num_m_blocks,), -1, dtype=torch.int32) + num_tokens_post_padded = torch.zeros(max_loras, dtype=torch.int32) + + seg_indptr_list = seg_indptr.cpu().tolist() + req_to_lora_list = req_to_lora.cpu().tolist() + topk_ids_list = topk_ids.cpu().tolist() + adapter_enabled_list = adapter_enabled.cpu().tolist() + + for lora_id in range(max_loras): + if not adapter_enabled_list[lora_id]: + continue + + pairs: list[tuple[int, int]] = [] + for seg_idx in range(len(seg_indptr_list) - 1): + if req_to_lora_list[seg_idx] != lora_id: + continue + start = seg_indptr_list[seg_idx] + end = seg_indptr_list[seg_idx + 1] + for m in range(start, end): + for k in range(top_k): + pairs.append((topk_ids_list[m][k], m * top_k + k)) + + if not pairs: + continue + + pairs.sort() + + base_t = lora_id * max_num_tokens_padded + base_e = lora_id * max_num_m_blocks + pos = 0 + block_idx = 0 + i = 0 + while i < len(pairs): + cur_expert = pairs[i][0] + group_start = pos + while i < len(pairs) and pairs[i][0] == cur_expert: + sorted_token_ids[base_t + pos] = pairs[i][1] + pos += 1 + i += 1 + group_len = pos - group_start + padded_len = ((group_len + block_size_m - 1) // block_size_m) * block_size_m + num_blocks = padded_len // block_size_m + for b in range(num_blocks): + expert_ids_out[base_e + block_idx + b] = cur_expert + block_idx += num_blocks + pos = group_start + padded_len + + num_tokens_post_padded[lora_id] = pos + + return ( + sorted_token_ids.to(device), + expert_ids_out.to(device), + num_tokens_post_padded.to(device), + ) + + +def _get_moe_lora_block_config(max_lora_rank: int) -> dict: + """Compute rank-aware block sizes for MoE LoRA kernels. + + Shrink: output dim is the rank -> cap BLOCK_SIZE_N to avoid waste. + Expand: input dim is the rank -> cap BLOCK_SIZE_K similarly. + """ + if max_lora_rank <= 0: + rank_pow2 = 64 + else: + rank_pow2 = next_power_of_2(max_lora_rank) + + shrink_n = min(64, rank_pow2) + expand_k = max(16, min(64, rank_pow2)) + + return { + "shrink_block_size_n": shrink_n, + "expand_block_size_k": expand_k, + } + + +_SPARSITY_FACTOR = 8 + + def _naive_moe_lora_align_block_size( topk_ids: torch.Tensor, seg_indptr: torch.Tensor, @@ -181,11 +288,13 @@ class LoRAInfo: num_experts: int experts_shared_outer_loras: bool = False cg_buffers: dict | None = None + cg_buffers: dict | None = None fully_sharded: bool = False tp_size: int = 1 tp_rank: int = 0 hidden_size: int = 0 + lora_use_virtual_experts: bool = False @dataclass @@ -200,11 +309,27 @@ class LoRAHooks: ) = None +def _compute_token_lora_mapping( + hidden_states: torch.Tensor, + lora_info: LoRAInfo, +) -> torch.Tensor: + """Map each token to its LoRA adapter index (-1 for no LoRA).""" + token_positions = torch.arange( + hidden_states.shape[0], device=hidden_states.device, dtype=torch.int32 + ) + req_indices = torch.searchsorted( + lora_info.seg_indptr[1:].to(torch.int32), + token_positions, + right=True, + ) + return lora_info.req_to_lora.to(torch.int32)[req_indices] + + def _compute_lora_alignment( topk_ids: torch.Tensor, lora_info: LoRAInfo, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Compute LoRA alignment tensors for MoE LoRA computation. + """Compute LoRA alignment tensors for the non-virtual-expert (classic) path. Returns: (sorted_token_ids_reshaped, expert_ids_reshaped, num_tokens_post_padded_lora, lora_ids) """ @@ -305,13 +430,18 @@ def _add_lora_gate_up_delta( topk_weights: torch.Tensor, topk_ids: torch.Tensor, lora_info: LoRAInfo, + token_lora_mapping: torch.Tensor | None, sorted_token_ids_reshaped: torch.Tensor | None, expert_ids_reshaped: torch.Tensor | None, num_tokens_post_padded_lora: torch.Tensor | None, lora_ids: torch.Tensor | None, + routing_cache: dict | None = None, ) -> None: """Add LoRA gate_up delta to intermediate_cache in-place.""" - from sglang.srt.lora.triton_ops import fused_moe_lora + from sglang.srt.lora.triton_ops import ( + fused_moe_lora, + merged_experts_fused_moe_lora_add, + ) if get_is_capture_mode(): # During CUDA graph capture, always enter the LoRA path so that @@ -338,43 +468,63 @@ def _add_lora_gate_up_delta( gate_up_a = lora_info.gate_up_lora_a_weights gate_up_b = lora_info.gate_up_lora_b_weights inter_size = gate_up_b.shape[2] // 2 + M, top_k, gate_up_dim = intermediate_cache.shape + 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 + inter_size = gate_up_b.shape[2] // 2 - if lora_info.experts_shared_outer_loras: + 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) inter_size = gate_up_b.shape[2] // 2 lora_a_stacked = [gate_up_a[:, :, :r, :], gate_up_a[:, :, r : 2 * r, :]] lora_b_stacked = [gate_up_b[:, :, :inter_size, :], gate_up_b[:, :, inter_size:, :]] - blk = _get_moe_lora_block_config(r) - fused_moe_lora( - output=intermediate_cache, - qcurr_hidden_states=hidden_states, - lora_a_stacked=lora_a_stacked, - lora_b_stacked=lora_b_stacked, - topk_weights=topk_weights, - sorted_token_ids=sorted_token_ids_reshaped, - expert_ids=expert_ids_reshaped, - num_tokens_post_padded=num_tokens_post_padded_lora, - max_lora_rank=r, - top_k_num=top_k, - lora_ids=lora_ids, - adapter_enabled=lora_info.adapter_enabled, - shrink_block_size_m=64, - shrink_block_size_n=blk["shrink_block_size_n"], - shrink_block_size_k=64, - shrink_group_size_m=8, - shrink_num_warps=4, - shrink_num_stages=2, - shrink_split_k=1, - expand_block_size_m=64, - expand_block_size_n=64, - expand_block_size_k=blk["expand_block_size_k"], - expand_group_size_m=8, - expand_num_warps=4, - expand_num_stages=2, - expand_split_k=1, - fully_sharded=lora_info.fully_sharded, - ) + if lora_info.lora_use_virtual_experts: + merged_experts_fused_moe_lora_add( + output=intermediate_cache, + hidden_states=hidden_states, + lora_a=gate_up_a, + lora_b=gate_up_b, + topk_ids=topk_ids, + topk_weights=topk_weights, + token_lora_mapping=token_lora_mapping, + mul_routed_weight=False, + experts_shared_outer_loras_a=lora_info.experts_shared_outer_loras, + experts_shared_outer_loras_b=False, + routing_cache=routing_cache, + ) + else: + blk = _get_moe_lora_block_config(r) + fused_moe_lora( + output=intermediate_cache, + qcurr_hidden_states=hidden_states, + lora_a_stacked=lora_a_stacked, + lora_b_stacked=lora_b_stacked, + topk_weights=topk_weights, + sorted_token_ids=sorted_token_ids_reshaped, + expert_ids=expert_ids_reshaped, + num_tokens_post_padded=num_tokens_post_padded_lora, + max_lora_rank=r, + top_k_num=top_k, + lora_ids=lora_ids, + adapter_enabled=lora_info.adapter_enabled, + shrink_block_size_m=64, + shrink_block_size_n=blk["shrink_block_size_n"], + shrink_block_size_k=64, + shrink_group_size_m=8, + shrink_num_warps=4, + shrink_num_stages=2, + shrink_split_k=1, + expand_block_size_m=64, + expand_block_size_n=64, + expand_block_size_k=blk["expand_block_size_k"], + expand_group_size_m=8, + expand_num_warps=4, + expand_num_stages=2, + expand_split_k=1, + fully_sharded=lora_info.fully_sharded, + ) def _add_lora_down_delta( @@ -383,13 +533,18 @@ def _add_lora_down_delta( topk_weights: torch.Tensor, topk_ids: torch.Tensor, lora_info: LoRAInfo, + token_lora_mapping: torch.Tensor | None, sorted_token_ids_reshaped: torch.Tensor | None, expert_ids_reshaped: torch.Tensor | None, num_tokens_post_padded_lora: torch.Tensor | None, lora_ids: torch.Tensor | None, + routing_cache: dict | None = None, ) -> None: """Add LoRA down delta to intermediate_cache in-place.""" - from sglang.srt.lora.triton_ops import fused_moe_lora + from sglang.srt.lora.triton_ops import ( + fused_moe_lora, + merged_experts_fused_moe_lora_add, + ) if lora_info.max_lora_rank == 0: return @@ -398,47 +553,67 @@ 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: + if lora_info.experts_shared_outer_loras and not lora_info.lora_use_virtual_experts: 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: + shard_size = lora_info.hidden_size // lora_info.tp_size + offset = shard_size * lora_info.tp_rank + else: + offset = 0 if lora_info.fully_sharded and lora_info.tp_size > 1: shard_size = lora_info.hidden_size // lora_info.tp_size offset = shard_size * lora_info.tp_rank else: offset = 0 - blk = _get_moe_lora_block_config(lora_info.max_lora_rank) - fused_moe_lora( - output=intermediate_cache, - qcurr_hidden_states=intermediate_input, - lora_a_stacked=[down_lora_a], - lora_b_stacked=[down_lora_b], - topk_weights=topk_weights, - sorted_token_ids=sorted_token_ids_reshaped, - expert_ids=expert_ids_reshaped, - num_tokens_post_padded=num_tokens_post_padded_lora, - max_lora_rank=lora_info.max_lora_rank, - top_k_num=top_k, - lora_ids=lora_ids, - adapter_enabled=lora_info.adapter_enabled, - shrink_block_size_m=64, - shrink_block_size_n=blk["shrink_block_size_n"], - shrink_block_size_k=64, - shrink_group_size_m=8, - shrink_num_warps=4, - shrink_num_stages=2, - shrink_split_k=1, - expand_block_size_m=64, - expand_block_size_n=64, - expand_block_size_k=blk["expand_block_size_k"], - expand_group_size_m=8, - expand_num_warps=4, - expand_num_stages=2, - expand_split_k=1, - mul_routed_weight=True, - fully_sharded=lora_info.fully_sharded, - offset=offset, - ) + if lora_info.lora_use_virtual_experts: + merged_experts_fused_moe_lora_add( + output=intermediate_cache, + hidden_states=intermediate_input, + lora_a=down_lora_a, + lora_b=down_lora_b, + topk_ids=topk_ids, + topk_weights=topk_weights, + token_lora_mapping=token_lora_mapping, + mul_routed_weight=True, + experts_shared_outer_loras_a=False, + experts_shared_outer_loras_b=lora_info.experts_shared_outer_loras, + routing_cache=routing_cache, + ) + else: + blk = _get_moe_lora_block_config(lora_info.max_lora_rank) + fused_moe_lora( + output=intermediate_cache, + qcurr_hidden_states=intermediate_input, + lora_a_stacked=[down_lora_a], + lora_b_stacked=[down_lora_b], + topk_weights=topk_weights, + sorted_token_ids=sorted_token_ids_reshaped, + expert_ids=expert_ids_reshaped, + num_tokens_post_padded=num_tokens_post_padded_lora, + max_lora_rank=lora_info.max_lora_rank, + top_k_num=top_k, + lora_ids=lora_ids, + adapter_enabled=lora_info.adapter_enabled, + shrink_block_size_m=64, + shrink_block_size_n=blk["shrink_block_size_n"], + shrink_block_size_k=64, + shrink_group_size_m=8, + shrink_num_warps=4, + shrink_num_stages=2, + shrink_split_k=1, + expand_block_size_m=64, + expand_block_size_n=64, + expand_block_size_k=blk["expand_block_size_k"], + expand_group_size_m=8, + expand_num_warps=4, + expand_num_stages=2, + expand_split_k=1, + mul_routed_weight=True, + fully_sharded=lora_info.fully_sharded, + offset=offset, + ) def build_lora_hooks( @@ -448,19 +623,31 @@ def build_lora_hooks( ) -> LoRAHooks: """Build LoRA hook closures for injection into any MoE runner. - Computes alignment tensors once, then returns closures that capture - them for the two injection points. + Computes token_lora_mapping and alignment tensors once, then returns + closures that capture them for the two injection points. """ if lora_info is None or lora_info.max_lora_rank == 0: return LoRAHooks() - # Compute alignment tensors (once, shared by both hooks) - ( - sorted_token_ids_reshaped, - expert_ids_reshaped, - num_tokens_post_padded_lora, - lora_ids, - ) = _compute_lora_alignment(topk_ids, lora_info) + # Compute alignment / mapping (once, shared by both hooks) + token_lora_mapping: torch.Tensor | None = None + sorted_token_ids_reshaped: torch.Tensor | None = None + expert_ids_reshaped: torch.Tensor | None = None + num_tokens_post_padded_lora: torch.Tensor | None = None + lora_ids: torch.Tensor | None = None + + if lora_info.lora_use_virtual_experts: + token_lora_mapping = _compute_token_lora_mapping(hidden_states, lora_info) + else: + ( + sorted_token_ids_reshaped, + expert_ids_reshaped, + num_tokens_post_padded_lora, + lora_ids, + ) = _compute_lora_alignment(topk_ids, lora_info) + + # Shared routing cache: gate_up and down reuse routing for same (num_experts, shared_outer, block_size) + routing_cache: dict = {} def after_gate_up( hidden_states: torch.Tensor, @@ -474,10 +661,12 @@ def after_gate_up( topk_weights, topk_ids, lora_info, + token_lora_mapping, sorted_token_ids_reshaped, expert_ids_reshaped, num_tokens_post_padded_lora, lora_ids, + routing_cache=routing_cache, ) def after_down( @@ -492,10 +681,12 @@ def after_down( topk_weights, topk_ids, lora_info, + token_lora_mapping, sorted_token_ids_reshaped, expert_ids_reshaped, num_tokens_post_padded_lora, lora_ids, + routing_cache=routing_cache, ) return LoRAHooks(after_gate_up=after_gate_up, after_down=after_down) diff --git a/python/sglang/srt/lora/triton_ops/__init__.py b/python/sglang/srt/lora/triton_ops/__init__.py index 7f96837cd367..276e879c9bb1 100644 --- a/python/sglang/srt/lora/triton_ops/__init__.py +++ b/python/sglang/srt/lora/triton_ops/__init__.py @@ -7,6 +7,7 @@ from .qkv_lora_b import qkv_lora_b_fwd from .sgemm_lora_a import sgemm_lora_a_fwd from .sgemm_lora_b import sgemm_lora_b_fwd +from .virtual_experts import merged_experts_fused_moe_lora_add __all__ = [ "gate_up_lora_b_fwd", @@ -18,4 +19,5 @@ "fused_moe_lora", "chunked_embedding_lora_a_forward", "embedding_lora_a_fwd", + "merged_experts_fused_moe_lora_add", ] diff --git a/python/sglang/srt/lora/triton_ops/virtual_experts.py b/python/sglang/srt/lora/triton_ops/virtual_experts.py new file mode 100644 index 000000000000..22b5b55aeb7f --- /dev/null +++ b/python/sglang/srt/lora/triton_ops/virtual_experts.py @@ -0,0 +1,662 @@ +""" +LoRA Virtual Experts Triton Ops. +""" + +import functools +from typing import Any + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _fused_virtual_topk_ids_kernel( + topk_ids_ptr, + token_lora_mapping_ptr, + virtual_topk_ids_ptr, + token_lora_mask_ptr, + num_experts_for_weight: tl.constexpr, + M, + top_k: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + """ + Fuses _get_virtual_topk_ids: comparison + clamp + arithmetic into one kernel. + + For each (m, k): + lora_id = token_lora_mapping[m] + mask[m] = (lora_id >= 0) + safe_lora = max(lora_id, 0) + if shared_outer: (handled by num_experts_for_weight == 0 sentinel) + virtual_topk_ids[m, k] = safe_lora * 1 (= safe_lora) + else: + virtual_topk_ids[m, k] = topk_ids[m, k] + safe_lora * num_experts_for_weight + """ + pid = tl.program_id(0) + offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + total = M * top_k + valid = offs < total + + m = offs // top_k + # k = offs % top_k # not needed directly + + lora_id = tl.load(token_lora_mapping_ptr + m, mask=valid, other=0) + mask_val = lora_id >= 0 + safe_lora = tl.maximum(lora_id, 0) + + base = tl.load(topk_ids_ptr + offs, mask=valid, other=0) + result = base + safe_lora * num_experts_for_weight + tl.store(virtual_topk_ids_ptr + offs, result, mask=valid) + + # Write mask once per row (at first k position) + k = offs % top_k + is_first_k = k == 0 + tl.store(token_lora_mask_ptr + m, mask_val, mask=valid & is_first_k) + + +def _fused_virtual_topk_ids( + topk_ids: torch.Tensor, + token_lora_mapping: torch.Tensor, + num_experts: int, + shared_outer: bool, + max_loras: int, +) -> tuple[torch.Tensor, torch.Tensor, int]: + """ + Returns virtual topk_ids, token_lora_mask, and virtual_num_experts. + """ + M, top_k = topk_ids.shape + device = topk_ids.device + + if shared_outer: + num_experts_for_weight = 1 + # For shared_outer, we need topk_ids to be zeros + zero_topk = torch.zeros_like(topk_ids) + input_topk = zero_topk + else: + num_experts_for_weight = num_experts + input_topk = topk_ids + + virtual_topk_ids = torch.empty_like(topk_ids) + token_lora_mask = torch.empty(M, dtype=torch.bool, device=device) + + BLOCK_SIZE = 1024 + grid = ((M * top_k + BLOCK_SIZE - 1) // BLOCK_SIZE,) + + _fused_virtual_topk_ids_kernel[grid]( + input_topk, + token_lora_mapping, + virtual_topk_ids, + token_lora_mask, + num_experts_for_weight, + M, + top_k, + BLOCK_SIZE, + ) + + virtual_num_experts = num_experts_for_weight * max_loras + return virtual_topk_ids, token_lora_mask, virtual_num_experts + + +@triton.jit +def _fused_sanitize_expert_ids_kernel( + expert_ids_ptr, + output_ptr, + num_virtual_experts, + N, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + valid = offs < N + + eid = tl.load(expert_ids_ptr + offs, mask=valid, other=0) + result = tl.where(eid < num_virtual_experts, eid, -1) + tl.store(output_ptr + offs, result, mask=valid) + + +def fused_sanitize_expert_ids( + expert_ids: torch.Tensor, + num_virtual_experts: int, +) -> torch.Tensor: + """ + Sanitize expert_ids by replacing values >= num_virtual_experts with -1. + + Returns a new tensor with expert_ids >= num_virtual_experts replaced by -1. + """ + N = expert_ids.numel() + output = torch.empty_like(expert_ids) + + BLOCK_SIZE = 1024 + grid = ((N + BLOCK_SIZE - 1) // BLOCK_SIZE,) + + _fused_sanitize_expert_ids_kernel[grid]( + expert_ids, + output, + num_virtual_experts, + N, + BLOCK_SIZE, + ) + return output + + +@triton.jit +def _moe_lora_shrink_splitk_kernel( + # Pointers + a_ptr, # type: ignore # [num_tokens, K] + b_ptr, # type: ignore # [num_virtual_experts, N, K] + c_ptr, # type: ignore # [num_tokens * top_k, N] (pre-zeroed when SPLIT_K > 1) + sorted_token_ids_ptr, # type: ignore + expert_ids_ptr, # type: ignore + num_tokens_post_padded_ptr, # type: ignore + # Dimensions + N, # type: ignore + K, # type: ignore + num_valid_tokens, # type: ignore + # Strides + stride_am, # type: ignore + stride_ak, # type: ignore + stride_be, # type: ignore + stride_bn, # type: ignore + stride_bk, # type: ignore + stride_cm, # type: ignore + stride_cn, # type: ignore + # Constexprs + top_k: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + SPLIT_K: tl.constexpr, +): + """Split-K grouped GEMM for the LoRA A (shrink) stage with few virtual experts.""" + pid = tl.program_id(0) + pid_sk = pid % SPLIT_K + pid_mn = pid // SPLIT_K + + num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr) + num_pid_m = tl.cdiv(num_tokens_post_padded, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid_mn // 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 + ((pid_mn % num_pid_in_group) % group_size_m) + pid_n = (pid_mn % num_pid_in_group) // group_size_m + + if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded: + return + + # Token routing (same pattern as fused_moe_triton_kernels) + offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64) + offs_token = tl.load(sorted_token_ids_ptr + offs_token_id).to(tl.int64) + token_mask = offs_token < num_valid_tokens + + off_expert = tl.load(expert_ids_ptr + pid_m).to(tl.int64) + if off_expert == -1: + return + + # Pointers + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N + offs_k = pid_sk * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + + a_ptrs = a_ptr + ( + offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak + ) + b_ptrs = ( + b_ptr + + off_expert * stride_be + + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + ) + + # Accumulate + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + grid_k = tl.cdiv(K, BLOCK_SIZE_K * SPLIT_K) + for k in range(0, grid_k): + k_remaining = K - k * (BLOCK_SIZE_K * SPLIT_K) + k_mask = offs_k[:, None] < k_remaining + a = tl.load( + a_ptrs, + mask=token_mask[:, None] & (offs_k[None, :] < k_remaining), + other=0.0, + ) + b = tl.load(b_ptrs, mask=k_mask, other=0.0) + accumulator += tl.dot(a, b.to(a.dtype)) + a_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_ak + b_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_bk + + accumulator = accumulator.to(c_ptr.dtype.element_ty) + + # Write output + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :] + c_mask = token_mask[:, None] & (offs_cn[None, :] < N) + if SPLIT_K == 1: + tl.store(c_ptrs, accumulator, mask=c_mask) + else: + tl.atomic_add(c_ptrs, accumulator, mask=c_mask, sem="relaxed") + + +def _invoke_moe_lora_shrink_splitk( + hidden_states: torch.Tensor, + weight: torch.Tensor, + output: torch.Tensor, + topk_ids: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + top_k: int, + config: dict[str, Any], +) -> None: + """Launch split-K shrink kernel for LoRA A with few virtual experts.""" + N = weight.shape[1] + K = weight.shape[2] + BLOCK_SIZE_M = config["BLOCK_SIZE_M"] + BLOCK_SIZE_N = min(config.get("BLOCK_SIZE_N", 64), max(16, N)) + BLOCK_SIZE_K = config.get("BLOCK_SIZE_K", 64) + GROUP_SIZE_M = config.get("GROUP_SIZE_M", 1) + + num_m_blocks = triton.cdiv(sorted_token_ids.shape[0], BLOCK_SIZE_M) + num_n_blocks = triton.cdiv(N, BLOCK_SIZE_N) + base_grid = num_m_blocks * num_n_blocks + max_split_k = max(1, K // BLOCK_SIZE_K) + SPLIT_K = min(max_split_k, max(1, 128 // base_grid)) if base_grid < 128 else 1 + + grid = (SPLIT_K * base_grid,) + + _moe_lora_shrink_splitk_kernel[grid]( + hidden_states, + weight, + output, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + N, + K, + topk_ids.numel(), + hidden_states.stride(0), + hidden_states.stride(1), + weight.stride(0), + weight.stride(1), + weight.stride(2), + output.stride(0), + output.stride(1), + top_k=top_k, + BLOCK_SIZE_M=BLOCK_SIZE_M, + BLOCK_SIZE_N=BLOCK_SIZE_N, + BLOCK_SIZE_K=BLOCK_SIZE_K, + GROUP_SIZE_M=GROUP_SIZE_M, + SPLIT_K=SPLIT_K, + num_warps=config.get("num_warps", 4), + num_stages=config.get("num_stages", 4), + ) + + +@torch.compile(dynamic=True) +def _align_block_size_torch( + topk_ids: torch.Tensor, + block_size: int, + num_experts: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Pure-PyTorch align_block_size for num_experts > 1024, compiled via torch.compile.""" + device = topk_ids.device + flat_topk_ids = topk_ids.reshape(-1).to(torch.int64) + num_valid_tokens = flat_topk_ids.numel() + max_total_padded_tokens = ( + (num_valid_tokens + num_experts * (block_size - 1) + block_size - 1) + // block_size + ) * block_size + max_num_blocks = max_total_padded_tokens // block_size + + sorted_token_ids = torch.full( + (max_total_padded_tokens,), + num_valid_tokens, + dtype=torch.int32, + device=device, + ) + expert_ids = torch.full( + (max_num_blocks,), + -1, + dtype=torch.int32, + device=device, + ) + + if num_valid_tokens == 0: + num_tokens_post_padded = torch.zeros((1,), dtype=torch.int32, device=device) + return sorted_token_ids, expert_ids, num_tokens_post_padded + + sorted_order = torch.argsort(flat_topk_ids) + sorted_expert_ids = flat_topk_ids[sorted_order] + expert_range = torch.arange(num_experts, device=device, dtype=torch.int64) + counts_offsets = torch.searchsorted(sorted_expert_ids, expert_range, right=False) + counts_end = torch.searchsorted(sorted_expert_ids, expert_range, right=True) + counts = counts_end - counts_offsets + padded_counts = ((counts + block_size - 1) // block_size) * block_size + total_padded_tokens = padded_counts.sum().to(torch.int32).reshape(1) + padded_offsets = torch.cumsum(padded_counts, dim=0) - padded_counts + + token_ranks = ( + torch.arange(num_valid_tokens, device=device, dtype=torch.int64) + - counts_offsets[sorted_expert_ids] + ) + output_positions = padded_offsets[sorted_expert_ids] + token_ranks + sorted_token_ids.scatter_( + 0, + output_positions.to(torch.int64), + sorted_order.to(torch.int32), + ) + + block_counts = padded_counts // block_size + actual_num_blocks = block_counts.sum() + + if max_num_blocks <= 0: + return sorted_token_ids, expert_ids, total_padded_tokens + + block_offsets = torch.cumsum(block_counts, dim=0) + all_block_positions = torch.arange(max_num_blocks, device=device, dtype=torch.int64) + assigned_experts = torch.searchsorted( + block_offsets, all_block_positions, right=True + ).to(torch.int32) + expert_ids.copy_( + torch.where( + all_block_positions < actual_num_blocks, + assigned_experts, + torch.full_like(assigned_experts, -1), + ) + ) + + return sorted_token_ids, expert_ids, total_padded_tokens + + +_align_block_size_large = _align_block_size_torch + + +def _merged_experts_fused_moe_lora_add_fake( + output: torch.Tensor, + hidden_states: torch.Tensor, + lora_a: torch.Tensor, + lora_b: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + token_lora_mapping: torch.Tensor, + mul_routed_weight: bool, + experts_shared_outer_loras_a: bool, + experts_shared_outer_loras_b: bool, +) -> None: + return + + +def _merged_experts_fused_moe_lora_add_impl( + output: torch.Tensor, + hidden_states: torch.Tensor, + lora_a: torch.Tensor, + lora_b: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + token_lora_mapping: torch.Tensor, + mul_routed_weight: bool, + experts_shared_outer_loras_a: bool, + experts_shared_outer_loras_b: bool, + routing_cache: dict | None = None, +) -> None: + """ + 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. + """ + max_loras, _, max_lora_rank, _ = lora_a.shape + 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: + # [max_loras, num_experts, x, y] -> [max_loras * num_experts, x, y] + return t.reshape(t.shape[0] * t.shape[1], t.shape[2], t.shape[3]) + + def _get_stage_config( + weight: torch.Tensor, + stage_top_k: int, + ) -> dict[str, Any]: + from sglang.srt.layers.moe.fused_moe_triton.fused_moe_triton_config import ( + get_config_dtype_str, + try_get_optimal_moe_config, + ) + + config_dtype = get_config_dtype_str(dtype=hidden_states.dtype) + get_config_func = functools.partial( + try_get_optimal_moe_config, + weight.shape, + weight.shape, + stage_top_k, + config_dtype, + ) + try: + cfg = get_config_func(token_lora_mapping.shape[0]) + except ValueError: + K_dim = weight.shape[2] + N_dim = weight.shape[1] + if K_dim >= 1024: + default_block_k = 256 + elif K_dim >= 64: + default_block_k = 64 + else: + default_block_k = max(16, K_dim) + cfg = { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": min(64, max(16, N_dim)), + "BLOCK_SIZE_K": min(default_block_k, max(16, K_dim)), + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4, + } + return cfg + + def _align_block_size( + topk_ids: torch.Tensor, + block_size: int, + num_experts: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + # The native align kernel consumes num_experts + 1 internally for its + # sentinel bucket, so the 1024-expert boundary must use the fallback path. + if num_experts < 1024: + from sglang.srt.layers.moe.fused_moe_triton.moe_align_block_size import ( + moe_align_block_size as native_moe_align_block_size, + ) + + return native_moe_align_block_size(topk_ids, block_size, num_experts) + return _align_block_size_large(topk_ids, block_size, num_experts) + + def _get_routing( + topk_ids: torch.Tensor, + token_lora_mapping: torch.Tensor, + num_experts: int, + shared_outer: bool, + block_size: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + # Check routing_cache for cross-call reuse (gate_up and down share routing) + cache_key = (num_experts, shared_outer, block_size) + if routing_cache is not None: + cached = routing_cache.get(cache_key) + if cached is not None: + return cached + + virtual_topk_ids, token_lora_mask, virtual_num_experts = ( + _fused_virtual_topk_ids( + topk_ids, token_lora_mapping, num_experts, shared_outer, max_loras + ) + ) + sorted_token_ids, expert_ids, num_tokens_post_padded = _align_block_size( + virtual_topk_ids, + 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] + expert_ids = fused_sanitize_expert_ids(expert_ids, virtual_num_experts) + result = ( + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + token_lora_mask, + ) + + if routing_cache is not None: + routing_cache[cache_key] = result + + return result + + from sglang.srt.layers.moe.fused_moe_triton.fused_moe_triton_kernels import ( + invoke_fused_moe_kernel, + ) + + lora_a_virtual = _merge_lora_expert_weight(lora_a) + lora_b_virtual = _merge_lora_expert_weight(lora_b) + num_experts_a = lora_a.shape[1] + num_experts_b = lora_b.shape[1] + + intermediate = torch.zeros( + [token_lora_mapping.shape[0], topk_ids.shape[1], max_lora_rank], + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + + a_stage_config = _get_stage_config(lora_a_virtual, input_top_k) + ( + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + token_lora_mask, + ) = _get_routing( + topk_ids, + token_lora_mapping, + num_experts_a, + experts_shared_outer_loras_a, + a_stage_config["BLOCK_SIZE_M"], + ) + + _invoke_moe_lora_shrink_splitk( + hidden_states, + lora_a_virtual, + intermediate.view(-1, max_lora_rank), + topk_ids, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + input_top_k, + a_stage_config, + ) + + b_stage_config = _get_stage_config(lora_b_virtual, 1) + ( + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + token_lora_mask, + ) = _get_routing( + topk_ids, + token_lora_mapping, + num_experts_b, + experts_shared_outer_loras_b, + 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], + ) + + +def _merged_experts_fused_moe_lora_add_op( + output: torch.Tensor, + hidden_states: torch.Tensor, + lora_a: torch.Tensor, + lora_b: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + token_lora_mapping: torch.Tensor, + mul_routed_weight: bool, + experts_shared_outer_loras_a: bool, + experts_shared_outer_loras_b: bool, +) -> None: + _merged_experts_fused_moe_lora_add_impl( + output, + hidden_states, + lora_a, + lora_b, + topk_ids, + topk_weights, + token_lora_mapping, + mul_routed_weight, + experts_shared_outer_loras_a, + experts_shared_outer_loras_b, + ) + + +from sglang.srt.utils.common import direct_register_custom_op + +direct_register_custom_op( + op_name="merged_experts_fused_moe_lora_add", + op_func=_merged_experts_fused_moe_lora_add_op, + mutates_args=["output"], + fake_impl=_merged_experts_fused_moe_lora_add_fake, +) + + +def merged_experts_fused_moe_lora_add( + output: torch.Tensor, + hidden_states: torch.Tensor, + lora_a: torch.Tensor, + lora_b: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + token_lora_mapping: torch.Tensor, + mul_routed_weight: bool, + experts_shared_outer_loras_a: bool, + experts_shared_outer_loras_b: bool, + routing_cache: dict | None = None, +) -> None: + """Public API: wraps the registered op with routing_cache support.""" + _merged_experts_fused_moe_lora_add_impl( + output, + hidden_states, + lora_a, + lora_b, + topk_ids, + topk_weights, + token_lora_mapping, + mul_routed_weight, + experts_shared_outer_loras_a, + experts_shared_outer_loras_b, + routing_cache, + ) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 81debbb9be68..8c7cdb12a38b 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -475,6 +475,7 @@ class ServerArgs: lora_backend: str = "csgmv" max_lora_chunk_size: Optional[int] = 16 experts_shared_outer_loras: Optional[bool] = None + lora_use_virtual_experts: bool = False lora_strict_loading: bool = False # Kernel backend @@ -4963,6 +4964,12 @@ def add_cli_args(parser: argparse.ArgumentParser): "(expert_dim=1). Use --no-experts-shared-outer-loras to force disable. " "By default this is auto-detected from adapter weights.", ) + parser.add_argument( + "--lora-use-virtual-experts", + default=ServerArgs.lora_use_virtual_experts, + action="store_true", + help="Enable virtual expert computation for MoE models. When set, the model will use virtual expert computation.", + ) parser.add_argument( "--lora-strict-loading", default=ServerArgs.lora_strict_loading, @@ -6715,6 +6722,13 @@ def check_lora_server_args(self): and (self.max_lora_chunk_size & (self.max_lora_chunk_size - 1)) == 0 ), "--max-lora-chunk-size must be a power of 2 between 16 and 128." + if self.lora_use_virtual_experts: + assert self.lora_backend == "triton", ( + "--lora-use-virtual-experts requires --lora-backend triton. " + f"Got: {self.lora_backend}" + ) + logger.info("Virtual expert computation enabled.") + def validate_buckets_rule(self, arg_name: str, buckets_rule: List[str]): if not buckets_rule: return diff --git a/test/registered/lora/test_lora_moe_runner.py b/test/registered/lora/test_lora_moe_runner.py index 0aad023957ed..eab94a4e6be6 100644 --- a/test/registered/lora/test_lora_moe_runner.py +++ b/test/registered/lora/test_lora_moe_runner.py @@ -120,6 +120,7 @@ def create_lora_info( gate_up_dim, dtype, device, + lora_use_virtual_experts=False, ): # ------------------------------------------------------------------------- # 1. Deterministic LoRA A Initialization @@ -186,6 +187,7 @@ def create_lora_info( adapter_enabled=adapter_enabled, max_lora_rank=max_lora_rank, num_experts=num_experts, + lora_use_virtual_experts=lora_use_virtual_experts, ) @@ -461,6 +463,155 @@ class MockServerArgs: torch.testing.assert_close(sglang_delta, torch_delta, atol=tol, rtol=tol) +@pytest.mark.parametrize("num_tokens", [32, 64]) +@pytest.mark.parametrize("top_k_num", [1, 2]) +@pytest.mark.parametrize("num_experts", [8, 20]) +@pytest.mark.parametrize("max_lora_rank", [8, 16]) +def test_lora_moe_runner_virtual_experts( + num_tokens, top_k_num, num_experts, max_lora_rank +): + # Fixed parameters + max_loras = 2 + hidden_dim = 512 + intermediate_dim = 1024 + + dtype = torch.float32 + device = "cuda:0" + seed = 42 + + torch.set_default_device(device) + set_random_seed(seed) + + num_sequences = 4 + + # Generate Data using the new Request-Based generator + topk_ids, topk_weights, seg_indptr, req_to_lora, token_lora_mapping = sample_data( + num_tokens, num_sequences, max_loras, num_experts, top_k_num, dtype, device + ) + + gate_up_dim = intermediate_dim * 2 + + # Initialize experts + w13 = torch.randn(num_experts, gate_up_dim, hidden_dim, dtype=dtype) * 0.1 + w2 = torch.randn(num_experts, hidden_dim, intermediate_dim, dtype=dtype) * 0.1 + b13 = torch.randn(num_experts, gate_up_dim, dtype=dtype) * 0.1 + b2 = torch.randn(num_experts, hidden_dim, dtype=dtype) * 0.1 + + hidden_states = torch.randn(num_tokens, hidden_dim, dtype=dtype) + + # Create LoRA Info with virtual experts enabled + lora_info_delta = create_lora_info( + seg_indptr=seg_indptr, + weight_indices=req_to_lora, + topk_ids=topk_ids, + max_loras=max_loras, + num_experts=num_experts, + max_lora_rank=max_lora_rank, + hidden_dim=hidden_dim, + intermediate_dim=intermediate_dim, + gate_up_dim=gate_up_dim, + dtype=dtype, + device=device, + lora_use_virtual_experts=True, + ) + + lora_info_baseline = create_lora_info( + seg_indptr=seg_indptr, + weight_indices=req_to_lora, + topk_ids=topk_ids, + max_loras=max_loras, + num_experts=num_experts, + max_lora_rank=0, + hidden_dim=hidden_dim, + intermediate_dim=intermediate_dim, + gate_up_dim=gate_up_dim, + dtype=dtype, + device=device, + lora_use_virtual_experts=True, + ) + + quant_info = TritonMoeQuantInfo( + w13_weight=w13, + w2_weight=w2, + b13=b13, + b2=b2, + ) + + config = MoeRunnerConfig( + activation="silu", + is_gated=True, + inplace=False, + no_combine=False, + gemm1_alpha=None, + gemm1_clamp_limit=None, + routed_scaling_factor=1.0, + apply_router_weight_on_input=False, + num_local_experts=num_experts, + ) + + router_logits = torch.randn(num_tokens, num_experts, dtype=dtype, device=device) + topk_output = StandardTopKOutput( + topk_weights=topk_weights, + topk_ids=topk_ids, + router_logits=router_logits, + ) + + dispatch_output = StandardDispatchOutput( + hidden_states=hidden_states, + hidden_states_scale=None, + topk_output=topk_output, + ) + + class MockServerArgs: + enable_deterministic_inference = False + + with patch( + "sglang.srt.layers.moe.fused_moe_triton.fused_moe_triton_config.get_global_server_args", + return_value=MockServerArgs(), + ): + runner = MoeRunner(MoeRunnerBackend.TRITON, config, lora_enabled=True) + + output_with_lora = runner.run( + dispatch_output, quant_info, lora_info_delta + ).hidden_states + output_baseline = runner.run( + dispatch_output, quant_info, lora_info_baseline + ).hidden_states + + # Run Naive Torch Implementation (Uses dense mapping for verification) + torch_output_lora = torch_naive_moe_with_lora( + hidden_states, + w13, + w2, + b13, + b2, + topk_weights, + topk_ids, + lora_info_delta, + token_lora_mapping, + ) + + torch_output_base = torch_naive_moe_with_lora( + hidden_states, + w13, + w2, + b13, + b2, + topk_weights, + topk_ids, + lora_info_baseline, + token_lora_mapping, + ) + + # The actual "Delta" (LoRA effect) for both + sglang_delta = output_with_lora - output_baseline + torch_delta = torch_output_lora - torch_output_base + + # Larger expert counts accumulate more numerical drift in Triton kernels on GB300 + tol = 0.15 if num_experts >= 20 else 5e-2 + torch.testing.assert_close(sglang_delta, torch_delta, atol=tol, rtol=tol) + + def _setup_marlin_moe_weights(num_experts, n, k, dtype): """Quantize float weights into AWQ Marlin format for testing.""" from sgl_kernel.scalar_type import scalar_types @@ -547,6 +698,7 @@ def test_lora_moe_runner_marlin(num_tokens, top_k_num, num_experts, max_lora_ran gate_up_dim=gate_up_dim, dtype=dtype, device=device, + lora_use_virtual_experts=True, ) lora_info_baseline = create_lora_info( @@ -561,6 +713,7 @@ def test_lora_moe_runner_marlin(num_tokens, top_k_num, num_experts, max_lora_ran gate_up_dim=gate_up_dim, dtype=dtype, device=device, + lora_use_virtual_experts=True, ) quant_info = MarlinMoeQuantInfo( diff --git a/test/registered/lora/test_marlin_lora_correctness.py b/test/registered/lora/test_marlin_lora_correctness.py index bc593debc03d..696fb5f8d648 100644 --- a/test/registered/lora/test_marlin_lora_correctness.py +++ b/test/registered/lora/test_marlin_lora_correctness.py @@ -204,6 +204,7 @@ def _make_lora_info(rank): adapter_enabled=torch.ones(num_loras + 1, dtype=torch.int32, device=device), max_lora_rank=rank, num_experts=num_experts, + lora_use_virtual_experts=True, experts_shared_outer_loras=True, ) From ba7bcca6b36a663f45f573537d9c4a0707065661 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Mon, 13 Apr 2026 14:29:12 -0700 Subject: [PATCH 36/52] Use reshape instead of contiguous().view() in TRTLLMHAAttnBackend (#22517) --- python/sglang/srt/layers/attention/trtllm_mha_backend.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index 722c598b6926..0fe16379e866 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mha_backend.py @@ -728,7 +728,7 @@ def forward_decode( # For XQA, q_dtype should be bf16 if self.data_type == torch.float8_e4m3fn and (not self.is_xqa_impl): q = q.to(torch.float8_e4m3fn) - q = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim) + q = q.reshape(-1, layer.tp_q_head_num, layer.head_dim) k_cache, v_cache = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id) # shape conversion: # [num_pages, page_size, num_kv_heads, head_dim] -> [num_pages, num_kv_heads, page_size, head_dim] @@ -813,7 +813,7 @@ def forward_extend( if self.data_type == torch.float8_e4m3fn: q = q.to(torch.float8_e4m3fn) - q = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim) + q = q.reshape(-1, layer.tp_q_head_num, layer.head_dim) # [num_pages, page_size, num_kv_heads, head_dim] -> [num_pages, num_kv_heads, page_size, head_dim] k_cache, v_cache = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id) k_cache = k_cache.view( From b441317aa43084d9c43709ef00ec692b8f3ff20e Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Mon, 13 Apr 2026 14:39:24 -0700 Subject: [PATCH 37/52] Revert "Upgrade CI default CUDA version from 12.9 to 13.0" (#22727) --- .github/workflows/pr-test.yml | 31 ++++++++++--------- python/pyproject.toml | 8 ++--- .../cuda/ci_download_flashinfer_jit_cache.sh | 2 +- scripts/ci/cuda/ci_install_deepep.sh | 9 ++++-- scripts/ci/cuda/ci_install_dependency.sh | 16 +++------- 5 files changed, 33 insertions(+), 33 deletions(-) diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 7c50a46d53f2..8feb364bd10a 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -429,7 +429,10 @@ jobs: matrix: include: - python-version: "3.10" - cuda-version: "13.0" + cuda-version: "12.9" + # Add back when CUDA 13.0 is supported on CI + # - python-version: "3.10" + # cuda-version: "13.0" name: Build Wheel steps: - name: Cleanup @@ -477,7 +480,7 @@ jobs: matrix: include: - python-version: "3.10" - cuda-version: "13.0" + cuda-version: "12.9" name: Build Wheel Arm steps: - name: Cleanup @@ -584,7 +587,7 @@ jobs: with: path: sgl-kernel/dist/ merge-multiple: true - pattern: wheel-python3.10-cuda13.0 + pattern: wheel-python3.10-cuda12.9 - name: Install dependencies timeout-minutes: 20 @@ -691,7 +694,7 @@ jobs: with: path: sgl-kernel/dist/ merge-multiple: true - pattern: wheel-python3.10-cuda13.0 + pattern: wheel-python3.10-cuda12.9 - name: Install dependencies timeout-minutes: 20 @@ -749,7 +752,7 @@ jobs: with: path: sgl-kernel/dist/ merge-multiple: true - pattern: wheel-python3.10-cuda13.0 + pattern: wheel-python3.10-cuda12.9 - name: Install dependencies timeout-minutes: 20 @@ -804,7 +807,7 @@ jobs: with: path: sgl-kernel/dist/ merge-multiple: true - pattern: wheel-python3.10-cuda13.0 + pattern: wheel-python3.10-cuda12.9 - name: Install dependencies timeout-minutes: 20 @@ -857,7 +860,7 @@ jobs: with: path: sgl-kernel/dist/ merge-multiple: true - pattern: wheel-python3.10-cuda13.0 + pattern: wheel-python3.10-cuda12.9 - name: Install dependencies timeout-minutes: 20 @@ -946,7 +949,7 @@ jobs: with: path: sgl-kernel/dist/ merge-multiple: true - pattern: wheel-python3.10-cuda13.0 + pattern: wheel-python3.10-cuda12.9 - name: Install dependencies timeout-minutes: 20 @@ -1000,7 +1003,7 @@ jobs: with: path: sgl-kernel/dist/ merge-multiple: true - pattern: wheel-python3.10-cuda13.0 + pattern: wheel-python3.10-cuda12.9 - name: Install dependencies timeout-minutes: 20 @@ -1070,7 +1073,7 @@ jobs: with: path: sgl-kernel/dist/ merge-multiple: true - pattern: wheel-python3.10-cuda13.0 + pattern: wheel-python3.10-cuda12.9 - name: Install dependencies timeout-minutes: 20 @@ -1120,7 +1123,7 @@ jobs: with: path: sgl-kernel/dist/ merge-multiple: true - pattern: wheel-python3.10-cuda13.0 + pattern: wheel-python3.10-cuda12.9 - name: Install dependencies timeout-minutes: 20 @@ -1180,7 +1183,7 @@ jobs: with: path: sgl-kernel/dist/ merge-multiple: true - pattern: wheel-python3.10-cuda13.0 + pattern: wheel-python3.10-cuda12.9 - name: Install dependencies timeout-minutes: 20 @@ -1246,7 +1249,7 @@ jobs: with: path: sgl-kernel/dist/ merge-multiple: true - pattern: wheel-python3.10-cuda13.0 + pattern: wheel-python3.10-cuda12.9 - name: Install dependencies timeout-minutes: 20 @@ -1300,7 +1303,7 @@ jobs: # with: # path: sgl-kernel/dist/ # merge-multiple: true - # pattern: wheel-python3.10-cuda13.0-aarch64 + # pattern: wheel-python3.10-cuda12.9-aarch64 # # - name: Install dependencies # timeout-minutes: 20 diff --git a/python/pyproject.toml b/python/pyproject.toml index 80ba96cc6cdf..491958878f39 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -22,7 +22,7 @@ dependencies = [ "blobfile==3.0.0", "build", "compressed-tensors", - "cuda-python>=13.0", + "cuda-python==12.9", "decord2 ; sys_platform == 'linux' and (platform_machine == 'aarch64' or platform_machine == 'arm64' or platform_machine == 'armv7l')", "datasets", "einops", @@ -86,14 +86,14 @@ url = "https://pypi.org/simple" default = true [[tool.uv.index]] -name = "torch-cu130" -url = "https://download.pytorch.org/whl/cu130" +name = "torch-cu129" +url = "https://download.pytorch.org/whl/cu129" explicit = true [tool.uv.sources] torch = [ { index = "pypi", marker = "platform_machine == 'x86_64'"}, - { index = "torch-cu130", marker = "platform_machine == 'aarch64'"}, + { index = "torch-cu129", marker = "platform_machine == 'aarch64'"}, ] [project.optional-dependencies] diff --git a/scripts/ci/cuda/ci_download_flashinfer_jit_cache.sh b/scripts/ci/cuda/ci_download_flashinfer_jit_cache.sh index 6358b5caf9b8..ab41dcef22c0 100755 --- a/scripts/ci/cuda/ci_download_flashinfer_jit_cache.sh +++ b/scripts/ci/cuda/ci_download_flashinfer_jit_cache.sh @@ -5,7 +5,7 @@ # Required environment (caller must export or set): # UNINSTALL_JIT_CACHE — literal true/false (skip download when false) # FLASHINFER_PYTHON_REQUIRED — e.g. from python/pyproject.toml (flashinfer_python) -# CU_VERSION — e.g. cu130 +# CU_VERSION — e.g. cu129 # PIP_CMD — e.g. "pip" or "uv pip" # PIP_INSTALL_SUFFIX — extra pip args for this runner set -euxo pipefail diff --git a/scripts/ci/cuda/ci_install_deepep.sh b/scripts/ci/cuda/ci_install_deepep.sh index b6e1887514be..bb4185799571 100755 --- a/scripts/ci/cuda/ci_install_deepep.sh +++ b/scripts/ci/cuda/ci_install_deepep.sh @@ -100,8 +100,13 @@ if [ "$GRACE_BLACKWELL" = "1" ]; then if [ "$CUDA_VERSION" = "12.8" ]; then CHOSEN_TORCH_CUDA_ARCH_LIST='10.0' elif awk -v ver="$CUDA_VERSION" 'BEGIN {exit !(ver > 12.8)}'; then - # CUDA > 12.8 supports sm_103 (Blackwell) - CHOSEN_TORCH_CUDA_ARCH_LIST='10.0;10.3' + # With cuda > 12.8, the compiler supports 10.3, so we should use + # CHOSEN_TORCH_CUDA_ARCH_LIST='10.0;10.3' + # + # However, our CI machine has a weird setup and nvidia-smi reports wrong CUDA version in the container. + # The container is actually cuda 12.8, but nvidia-smi reports 13.0, leading to compilation errors. so we + # drop 10.3. + CHOSEN_TORCH_CUDA_ARCH_LIST='10.0' else echo "Unsupported CUDA version for Grace Blackwell: $CUDA_VERSION" && exit 1 fi && \ diff --git a/scripts/ci/cuda/ci_install_dependency.sh b/scripts/ci/cuda/ci_install_dependency.sh index 9ac79327c30e..7735977ef278 100755 --- a/scripts/ci/cuda/ci_install_dependency.sh +++ b/scripts/ci/cuda/ci_install_dependency.sh @@ -23,7 +23,7 @@ set -euxo pipefail # Configuration & timing # ------------------------------------------------------------------------------ # Set up environment variables -CU_VERSION="cu130" +CU_VERSION="cu129" # Nvidia package versions we override (torch pins older versions). # Used both as pip constraints during install and for post-install verification. @@ -244,15 +244,7 @@ if [ "${CUSTOM_BUILD_SGL_KERNEL:-}" = "true" ] && [ -d "sgl-kernel/dist" ]; then else WHEEL_ARCH="x86_64" fi - # Wheel may have +cuXYZ suffix (e.g. sglang_kernel-0.4.0+cu130-...) depending on CUDA version - KERNEL_WHL=$(ls sgl-kernel/dist/sglang_kernel-${SGL_KERNEL_VERSION_FROM_KERNEL}*-cp310-abi3-manylinux2014_${WHEEL_ARCH}.whl 2>/dev/null | head -1) - if [ -z "$KERNEL_WHL" ]; then - echo "ERROR: No matching sgl-kernel wheel found in sgl-kernel/dist/ for version ${SGL_KERNEL_VERSION_FROM_KERNEL} arch ${WHEEL_ARCH}" - ls -alh sgl-kernel/dist/ - exit 1 - fi - echo "Installing sgl-kernel wheel: $KERNEL_WHL" - $PIP_CMD install "$KERNEL_WHL" --force-reinstall $PIP_INSTALL_SUFFIX + $PIP_CMD install sgl-kernel/dist/sglang_kernel-${SGL_KERNEL_VERSION_FROM_KERNEL}-cp310-abi3-manylinux2014_${WHEEL_ARCH}.whl --force-reinstall $PIP_INSTALL_SUFFIX elif [ "${CUSTOM_BUILD_SGL_KERNEL:-}" = "true" ] && [ ! -d "sgl-kernel/dist" ]; then # CUSTOM_BUILD_SGL_KERNEL was set but artifacts not available (e.g., stage rerun without wheel build) # Fail instead of falling back to PyPI - we need to test the built kernel, not PyPI version @@ -326,10 +318,10 @@ mark_step_done "Install extra dependency" # Fix other dependencies # ------------------------------------------------------------------------------ # Fix CUDA version mismatch between torch and torchaudio. -# PyPI's torch bundles a specific CUDA version but torchaudio from pytorch.org/cu130 may use a different one. +# PyPI's torch 2.9.1 bundles cu128 but torchaudio from pytorch.org/cu129 uses cu129. # This mismatch causes torchaudio's C extension to fail loading, producing: # "partially initialized module 'torchaudio' has no attribute 'lib'" -# We cannot replace torch with a different CUDA version (breaks sgl_kernel ABI), so instead we reinstall +# We cannot replace torch with cu129 (breaks sgl_kernel ABI), so instead we reinstall # torchaudio/torchvision from an index matching torch's CUDA version. TORCH_CUDA_VER=$(python3 -c "import torch; v=torch.version.cuda; parts=v.split('.'); print(f'cu{parts[0]}{parts[1]}')") echo "Detected torch CUDA version: ${TORCH_CUDA_VER}" From f81b6e8f51705e58dcc589259277d4f732b53887 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Mon, 13 Apr 2026 14:42:49 -0700 Subject: [PATCH 38/52] [Misc] Add @cache_once to is_arch_support_pdl in jit_kernel (#22724) --- python/sglang/jit_kernel/utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/sglang/jit_kernel/utils.py b/python/sglang/jit_kernel/utils.py index ec3ebd5abc80..59ced5e57e68 100644 --- a/python/sglang/jit_kernel/utils.py +++ b/python/sglang/jit_kernel/utils.py @@ -52,7 +52,7 @@ def cache_once(fn: F) -> F: @functools.wraps(fn) def wrapper(*args, **kwargs): - key = (args, tuple(sorted(kwargs.items(), key=lambda x: x[0]))) + key = (args, tuple(sorted(kwargs.items()))) if key not in result_map: result_map[key] = fn(*args, **kwargs) return result_map[key] @@ -305,6 +305,7 @@ def get_jit_cuda_arch() -> ArchInfo: return _CUDA_ARCH +@cache_once def is_arch_support_pdl() -> bool: if is_hip_runtime(): return False From b1efce342ce3946c40c27e00812f0468d735c1e1 Mon Sep 17 00:00:00 2001 From: Yilong Zhao <74357408+happierpig@users.noreply.github.com> Date: Mon, 13 Apr 2026 15:37:59 -0700 Subject: [PATCH 39/52] env: add knob to control SWA eviction interval (#22645) --- python/sglang/srt/environ.py | 1 + python/sglang/srt/managers/schedule_batch.py | 14 ++++++++++++-- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 42309d9039fe..42b340a3dd70 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -263,6 +263,7 @@ class Envs: SGLANG_NCCL_ALL_GATHER_IN_OVERLAP_SCHEDULER_SYNC_BATCH = EnvBool(False) SGLANG_REQ_RUNNING_TIMEOUT = EnvFloat(-1) # in seconds SGLANG_DISAGGREGATION_BOOTSTRAP_ENTRY_CLEANUP_INTERVAL = EnvInt(120) + SGLANG_SWA_EVICTION_INTERVAL_MULTIPLIER = EnvFloat(1.0) # Test: pd-disaggregation SGLANG_TEST_PD_DISAGG_BACKEND = EnvStr("mooncake") diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index ddddbcef1f62..a6172890fcc8 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -2482,12 +2482,22 @@ def maybe_evict_swa(self): sliding_window_size = self.tree_cache.sliding_window_size server_args = get_global_server_args() + # Eviction_interval: trade-off between SWA token waste and eviction overhead + page_size = self.tree_cache.page_size + eviction_interval = max( + page_size, + int( + sliding_window_size + * envs.SGLANG_SWA_EVICTION_INTERVAL_MULTIPLIER.get() + ), + ) + eviction_interval = (eviction_interval // page_size) * page_size for idx, req in enumerate(self.reqs): if self.forward_mode.is_decode(): # We set evict_swa condition here with two reasons: # 1. In overlap scheduler, we cannot evict swa when req.decode_batch_idx == 0 since the prev extend batch is still running. - # 2. Evict swa every window_size tokens to reduce the overhead. - if req.decode_batch_idx % sliding_window_size == 1: + # 2. Evict swa every eviction_interval tokens to reduce the overhead. + if req.decode_batch_idx % eviction_interval == 1: self._evict_swa(req, req.seqlen - 1) elif self.forward_mode.is_extend() and self.tree_cache.is_chunk_cache(): pre_len = self.prefix_lens[idx] From f4f9e681891673ecb2d92dca1fe00e65c171eff8 Mon Sep 17 00:00:00 2001 From: mqhc2020 Date: Tue, 14 Apr 2026 06:50:15 +0800 Subject: [PATCH 40/52] [AMD] Add MoE weights and scales padding (#21097) Co-authored-by: HAI --- .../layers/moe/fused_moe_triton/fused_moe.py | 4 +- .../fused_moe_triton_kernels.py | 4 +- .../srt/layers/moe/fused_moe_triton/layer.py | 30 ++++++---- python/sglang/srt/layers/moe/utils.py | 51 +++++++++++++++++ .../compressed_tensors_w8a8_fp8_moe.py | 22 +++++-- python/sglang/srt/layers/quantization/fp8.py | 57 ++++++++++++------- .../quark/schemes/quark_w4a4_mxfp4_moe.py | 25 ++++++-- .../sglang/srt/model_executor/model_runner.py | 6 +- 8 files changed, 153 insertions(+), 46 deletions(-) diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py index f2621473cb49..7a7d9fd0b4b0 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py @@ -6,7 +6,6 @@ from __future__ import annotations import functools -import os from typing import TYPE_CHECKING, List, Optional import torch @@ -14,6 +13,7 @@ import triton.language as tl from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig +from sglang.srt.layers.moe.utils import get_moe_padding_size from sglang.srt.utils import ( cpu_has_amx_support, get_bool_env_var, @@ -75,7 +75,7 @@ # Fallback: vllm not available, will use native PyTorch implementations _has_vllm_ops = False -padding_size = 128 if bool(int(os.getenv("SGLANG_MOE_PADDING", "0"))) else 0 +padding_size = get_moe_padding_size(_use_aiter) @register_custom_op(mutates_args=["hidden_states"]) diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py index 6a2b46d99786..2b22f16fec23 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py @@ -1,7 +1,6 @@ from __future__ import annotations import functools -import os from collections import OrderedDict from typing import Any, Dict, List, Optional @@ -10,6 +9,7 @@ import triton.language as tl from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled +from sglang.srt.layers.moe.utils import get_moe_padding_size from sglang.srt.layers.quantization.fp8_kernel import ( per_token_group_quant_fp8, scaled_fp8_quant, @@ -49,7 +49,7 @@ elif _is_hip: pass -padding_size = 128 if bool(int(os.getenv("SGLANG_MOE_PADDING", "0"))) else 0 +padding_size = get_moe_padding_size(_use_aiter) def support_tensor_descriptor(): diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index 01e9c16f60b4..6bb195758656 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -2,6 +2,7 @@ import logging from enum import Enum +from functools import cached_property from typing import List, Optional, Tuple import torch @@ -325,6 +326,21 @@ def __init__( if self.quant_method is not None and hasattr(self.quant_method, "runner"): self.runner = self.quant_method.runner + @cached_property + def use_padded_loading(self) -> bool: + # This handles the case where the loaded weights are smaller than the padded expert_data + # Use narrow_padded_param_and_loaded_weight for: + # 1. CPU (always) + # 2. GPU with flashinfer_trtllm padding (when intermediate_size is padded to 128) + # 3. GPU with Aiter padding + aiter_padded = ( + _use_aiter + and hasattr(self, "w2_weight") + and getattr(self.w2_weight, "weight_padded", False) + ) + + return _is_cpu or self.use_flashinfer_trtllm_moe or aiter_padded + def _load_per_tensor_weight_scale( self, shard_id: str, @@ -436,12 +452,7 @@ def _load_w13( else: start = 0 - # Use narrow_padded_param_and_loaded_weight for: - # 1. CPU (always) - # 2. GPU with flashinfer_trtllm padding (when intermediate_size is padded to 128) - # This handles the case where the loaded weights are smaller than the padded expert_data - use_padded_loading = _is_cpu or self.use_flashinfer_trtllm_moe - if use_padded_loading: + if self.use_padded_loading: expert_data, loaded_weight = narrow_padded_param_and_loaded_weight( expert_data, loaded_weight, @@ -510,12 +521,7 @@ def _load_w2( # for w2 in TP, it shards the input_features, i.e., shard_dim=2 shard_size = expert_data.shape[shard_dim] - # Use narrow_padded_param_and_loaded_weight for: - # 1. CPU (always) - # 2. GPU with flashinfer_trtllm padding (when intermediate_size is padded to 128) - # This handles the case where the loaded weights are smaller than the padded expert_data - use_padded_loading = _is_cpu or self.use_flashinfer_trtllm_moe - if use_padded_loading: + if self.use_padded_loading: expert_data, loaded_weight = narrow_padded_param_and_loaded_weight( expert_data, loaded_weight, diff --git a/python/sglang/srt/layers/moe/utils.py b/python/sglang/srt/layers/moe/utils.py index 8eeacde05e99..279945acdf1c 100644 --- a/python/sglang/srt/layers/moe/utils.py +++ b/python/sglang/srt/layers/moe/utils.py @@ -1,10 +1,13 @@ from __future__ import annotations import logging +import os from contextlib import contextmanager from enum import Enum, IntEnum from typing import TYPE_CHECKING, Optional +import torch + from sglang.srt.distributed.parallel_state import get_moe_expert_parallel_world_size from sglang.srt.layers.dp_attention import ( get_attention_dp_size, @@ -351,3 +354,51 @@ class RoutingMethodType(IntEnum): TopK = (5,) # Unspecified Unspecified = 6 + + +AITER_PADDING_SIZE = 128 +TRITON_PADDING_SIZE = 128 + + +# Unit of padding - context dependent +def get_moe_padding_size(is_aiter_moe): + if is_aiter_moe: + return AITER_PADDING_SIZE + else: + return ( + TRITON_PADDING_SIZE + if bool(int(os.getenv("SGLANG_MOE_PADDING", "0"))) + else 0 + ) + + +def get_moe_weight_sizes(inter_dim, is_concat, is_packed, is_aiter_moe): + """ + Calculate dimensions for MoE weight tensors. + + Args: + inter_dim: Base intermediate dimension. + is_concat: If True, fusions W1 (gate) and W3 (up) projections. + is_packed: If True, uses 4-bit quantization (two FP4 elements per byte). + is_aiter_moe: If True, applies Aiter-specific kernel padding alignment. + """ + # w2_down_dim is the packing rank, but w13_up_dim not (of matrix to matmul) + w13_up_dim = 2 * inter_dim if is_concat else inter_dim + w2_down_dim = inter_dim // 2 if is_packed else inter_dim + + if is_aiter_moe: + padding_size = get_moe_padding_size(True) + align_aiter = lambda n: ((n + padding_size - 1) // padding_size) * padding_size + is_padded = (w2_down_dim % padding_size) > 0 + if is_padded: + # w2_down_dim, padding & aligned, unit: parameter dtype + w2_down_dim = align_aiter(w2_down_dim) + # up proj + gate fusion : 2x + if is_concat: + w13_up_dim = w2_down_dim * 2 + # packed + if hasattr(torch, "float4_e2m1fn_x2") and is_packed: + # w13_up_dim (row rank of matmul matrix) is not packing dim, *2 to recover + w13_up_dim *= 2 + + return (w13_up_dim, w2_down_dim, False if not is_aiter_moe else is_padded) diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8_moe.py index fdd3c9622ad2..c4139f54aab1 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8_moe.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8_moe.py @@ -12,7 +12,10 @@ FlashInferTrtllmFp8MoeQuantInfo, ) from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo -from sglang.srt.layers.moe.utils import get_moe_runner_backend +from sglang.srt.layers.moe.utils import ( + get_moe_runner_backend, + get_moe_weight_sizes, +) from sglang.srt.layers.quantization.compressed_tensors.schemes import ( CompressedTensorsMoEScheme, ) @@ -120,11 +123,22 @@ def create_weights( f"weight quantization block_k = {block_k}." ) + w13_up_dim, w2_down_dim, weight_padded = get_moe_weight_sizes( + intermediate_size_per_partition, + is_aiter_moe=_use_aiter, + is_concat=True, + is_packed=False, + ) + + extra_weight_attrs.update( + {"weight_padded": weight_padded}, + ) + # WEIGHTS w13_weight = torch.nn.Parameter( torch.empty( num_experts, - 2 * intermediate_size_per_partition, + w13_up_dim, hidden_size, dtype=params_dtype, ), @@ -137,7 +151,7 @@ def create_weights( torch.empty( num_experts, hidden_size, - intermediate_size_per_partition, + w2_down_dim, dtype=params_dtype, ), requires_grad=False, @@ -161,7 +175,7 @@ def create_weights( w13_weight_scale = torch.nn.Parameter( torch.ones( num_experts, - 2 * intermediate_size_per_partition, + w13_up_dim, 1, dtype=torch.float32, ), diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 4aaa910effc9..1ca9deab7203 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -26,7 +26,12 @@ FlashInferTrtllmFp8MoeQuantInfo, ) from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo -from sglang.srt.layers.moe.utils import RoutingMethodType, get_moe_runner_backend +from sglang.srt.layers.moe.utils import ( + RoutingMethodType, + get_moe_padding_size, + get_moe_runner_backend, + get_moe_weight_sizes, +) from sglang.srt.layers.parameter import ( BlockQuantScaleParameter, ModelWeightParameter, @@ -814,27 +819,38 @@ def create_weights( if self.quant_config.is_checkpoint_fp8_serialized: params_dtype = torch.uint32 if _use_hip_int4 else torch.float8_e4m3fn tp_size = get_tensor_model_parallel_world_size() + + w13_up_dim, w2_up_dim, weight_padded = get_moe_weight_sizes( + intermediate_size_per_partition, + is_aiter_moe=_use_aiter, + is_concat=True, + is_packed=False, + ) + if self.block_quant: block_n, block_k = ( self.quant_config.weight_block_size[0], self.quant_config.weight_block_size[1], ) - # NOTE(HandH1998): To ensure proper alignment of the block-wise quantization scales, the output_size of the weights for both the gate and up layers must be divisible by block_n. - # Required by column parallel or enabling merged weights - if intermediate_size_per_partition % block_n != 0: - raise ValueError( - f"The output_size of gate's and up's weight = " - f"{intermediate_size_per_partition} is not divisible by " - f"weight quantization block_n = {block_n}." - ) - if tp_size > 1: - # Required by row parallel - if intermediate_size_per_partition % block_k != 0: + + padding_size = get_moe_padding_size(_use_aiter) + if not (_use_aiter and padding_size == block_n == block_k): + # NOTE(HandH1998): To ensure proper alignment of the block-wise quantization scales, the output_size of the weights for both the gate and up layers must be divisible by block_n. + # Required by column parallel or enabling merged weights + if intermediate_size_per_partition % block_n != 0: raise ValueError( - f"The input_size of down's weight = " + f"The output_size of gate's and up's weight = " f"{intermediate_size_per_partition} is not divisible by " - f"weight quantization block_k = {block_k}." + f"weight quantization block_n = {block_n}." ) + if tp_size > 1: + # Required by row parallel + if intermediate_size_per_partition % block_k != 0: + raise ValueError( + f"The input_size of down's weight = " + f"{intermediate_size_per_partition} is not divisible by " + f"weight quantization block_k = {block_k}." + ) # WEIGHTS if _is_hip and _use_hip_int4: @@ -861,7 +877,7 @@ def create_weights( w13_weight = torch.nn.Parameter( torch.empty( num_experts, - 2 * intermediate_size_per_partition, + w13_up_dim, hidden_size, dtype=params_dtype, ), @@ -871,12 +887,16 @@ def create_weights( torch.empty( num_experts, hidden_size, - intermediate_size_per_partition, + w2_up_dim, dtype=params_dtype, ), requires_grad=False, ) + extra_weight_attrs.update( + {"weight_padded": weight_padded}, + ) + layer.register_parameter("w13_weight", w13_weight) set_weight_attrs(w13_weight, extra_weight_attrs) @@ -1439,10 +1459,7 @@ def process_weights_hip_int4(self, layer: Module): layer.w2_weight_scale1[expert_id] *= layer.w2_weight_scale[expert_id] def process_weights_hip_scale_padding(self, layer: Module): - from sglang.srt.layers.moe.fused_moe_triton.fused_moe import ( - padding_size, # Avoid circular import - ) - + padding_size = get_moe_padding_size(_use_aiter) if _use_aiter: layer.w13_weight = torch.nn.Parameter( shuffle_weight(layer.w13_weight.data, (16, 16)), diff --git a/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py b/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py index ae9743577fb8..68386e3b9de9 100644 --- a/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py +++ b/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py @@ -8,6 +8,7 @@ import torch from sglang.srt.layers.moe import MoeRunnerConfig +from sglang.srt.layers.moe.utils import get_moe_weight_sizes from sglang.srt.layers.quantization.quark.schemes import QuarkMoEScheme from sglang.srt.utils import ( get_bool_env_var, @@ -73,10 +74,20 @@ def create_weights( from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported + w13_up_dim, w2_down_dim, weight_padded = get_moe_weight_sizes( + intermediate_size_per_partition, + is_aiter_moe=_use_aiter, + is_concat=True, + is_packed=True, + ) + # Add the quantization method used (per tensor/grouped/channel) # to ensure the weight scales are loaded in properly extra_weight_attrs.update( - {"quant_method": FusedMoeWeightScaleSupported.BLOCK.value} + { + "quant_method": FusedMoeWeightScaleSupported.BLOCK.value, + "weight_padded": weight_padded, + }, ) params_dtype = torch.uint8 @@ -85,7 +96,7 @@ def create_weights( w13_weight = torch.nn.Parameter( torch.empty( num_experts, - 2 * intermediate_size_per_partition, + w13_up_dim, hidden_size // 2, dtype=params_dtype, ), @@ -99,7 +110,7 @@ def create_weights( torch.empty( num_experts, hidden_size, - intermediate_size_per_partition // 2, + w2_down_dim, dtype=params_dtype, ), requires_grad=False, @@ -112,17 +123,21 @@ def create_weights( w13_weight_scale = torch.nn.Parameter( torch.ones( num_experts, - 2 * intermediate_size_per_partition, + w13_up_dim, hidden_size // OCP_MX_BLOCK_SIZE, dtype=params_dtype, ), requires_grad=False, ) + + # 1. w2 scale is floor division of inter_dim by blockscale. + # 2. w2 scale needs to scale up just as w2. + # We combine 1. and 2. to keep the integer precision. w2_weight_scale = torch.nn.Parameter( torch.ones( num_experts, hidden_size, - intermediate_size_per_partition // OCP_MX_BLOCK_SIZE, + (w2_down_dim * 2) // OCP_MX_BLOCK_SIZE, dtype=params_dtype, ), requires_grad=False, diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 7dd70af63938..75eee24c1d62 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -165,6 +165,7 @@ empty_context, enable_show_time_cost, get_available_gpu_memory, + get_bool_env_var, get_cpu_ids_by_node, init_custom_process_group, is_hip, @@ -201,6 +202,7 @@ _is_npu = is_npu() _is_cpu_amx_available = cpu_has_amx_support() _is_cpu_arm64 = is_host_cpu_arm64() +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip if _is_npu: from sglang.srt.hardware_backend.npu.utils import init_npu_backend @@ -955,7 +957,9 @@ def check_quantized_moe_compatibility(self): f"moe_intermediate_size {moe_intermediate_size} must be divisible by moe_tp_size ({moe_tp_size}) which is tp_size ({self.tp_size}) divided by moe_ep_size ({self.moe_ep_size})." ) - if (moe_intermediate_size // moe_tp_size) % weight_block_size_n != 0: + if ( + moe_intermediate_size // moe_tp_size + ) % weight_block_size_n != 0 and not _use_aiter: raise ValueError( f"For quantized MoE models, please make sure ({moe_intermediate_size=} / {moe_tp_size=}) % {weight_block_size_n=} == 0 " f"where moe_tp_size is equal to tp_size ({self.tp_size}) divided by ep_size ({self.moe_ep_size}). " From 8f9553bccbf98e7c25dd9e8a0c152cc7722d8ce5 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Mon, 13 Apr 2026 16:10:31 -0700 Subject: [PATCH 41/52] [Misc] Migrate SGLANG_SET_CPU_AFFINITY to envs and refactor model config building (#22730) --- python/sglang/srt/managers/scheduler.py | 7 +------ python/sglang/srt/model_executor/model_runner.py | 14 ++++++++++++-- .../srt/observability/scheduler_metrics_mixin.py | 5 +++-- 3 files changed, 16 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 30c242daa966..5297f070b09f 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -379,10 +379,6 @@ def __init__( ) ) - self.enable_kv_cache_events = bool( - server_args.kv_events_config and self.attn_tp_rank == 0 - ) - # Init model configs self.init_model_config() @@ -3688,12 +3684,11 @@ def run_scheduler_process( dp_rank = configure_scheduler( server_args, tp_rank, attn_cp_rank, moe_dp_rank, moe_ep_rank, pp_rank, dp_rank ) - kill_itself_when_parent_died() parent_process = psutil.Process().parent() # Set cpu affinity to this gpu process - if get_bool_env_var("SGLANG_SET_CPU_AFFINITY"): + if envs.SGLANG_SET_CPU_AFFINITY.get(): set_gpu_proc_affinity( server_args.pp_size, server_args.tp_size, server_args.nnodes, gpu_id ) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 75eee24c1d62..74cec2565838 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -362,7 +362,7 @@ def __init__( self.dflash_draft_num_layers = None if self.spec_algorithm.is_eagle3() and not self.is_draft_worker: # load draft config - draft_model_config = ModelConfig.from_server_args( + draft_model_config = self._build_model_config( server_args, model_path=(server_args.speculative_draft_model_path), model_revision=server_args.speculative_draft_model_revision, @@ -391,7 +391,7 @@ def __init__( ) # Select target layers to capture for building DFlash context features. - draft_model_config = ModelConfig.from_server_args( + draft_model_config = self._build_model_config( server_args, model_path=(server_args.speculative_draft_model_path), model_revision=server_args.speculative_draft_model_revision, @@ -497,6 +497,16 @@ def __init__( self._model_update_group = {} self._weights_send_group = {} + def _build_model_config( + self, server_args, model_path=None, model_revision=None, is_draft_model=False + ): + return ModelConfig.from_server_args( + server_args, + model_path=model_path, + model_revision=model_revision, + is_draft_model=is_draft_model, + ) + def init_mindspore_runner(self): # Init the mindspore runner # for now, there is only some communication initialization work diff --git a/python/sglang/srt/observability/scheduler_metrics_mixin.py b/python/sglang/srt/observability/scheduler_metrics_mixin.py index 796c312c364b..6a313958f0b0 100644 --- a/python/sglang/srt/observability/scheduler_metrics_mixin.py +++ b/python/sglang/srt/observability/scheduler_metrics_mixin.py @@ -166,14 +166,15 @@ def init_metrics( reporter=self.metrics_collector.increment_gpu_overlap_wait_seconds, ) - if self.enable_kv_cache_events: - self.init_kv_events(self.server_args.kv_events_config) + self.init_kv_events(self.server_args.kv_events_config) self.scheduler_status_logger = SchedulerStatusLogger.maybe_create( enable_metrics=self.enable_metrics ) def init_kv_events(self: Scheduler, kv_events_config: Optional[str]): + self.enable_kv_cache_events = bool(kv_events_config and self.attn_tp_rank == 0) + if self.enable_kv_cache_events: self.kv_event_publisher = EventPublisherFactory.create( kv_events_config, self.attn_dp_rank From a2b51119621a3f599c99adac5dd45fab17fd8f18 Mon Sep 17 00:00:00 2001 From: Jia Guo Date: Mon, 13 Apr 2026 16:27:52 -0700 Subject: [PATCH 42/52] perf: skip KV cache in FA backend for embedding mode (#21971) Co-authored-by: Claude Opus 4.6 --- .../attention/flashattention_backend.py | 37 ++++++++++++++++++- 1 file changed, 36 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index ab4cf841ff97..7680f2c6a4b2 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -214,6 +214,16 @@ def __init__( else 0 ) + # In embedding mode with no chunked prefill and radix cache disabled, + # skip KV cache write and use flash_attn_varlen_func with raw K/V + # instead of flash_attn_with_kvcache, bypassing paged KV cache entirely. + server_args = model_runner.server_args + self.fa_skip_kv_cache = ( + server_args.is_embedding + and server_args.chunked_prefill_size == -1 + and server_args.disable_radix_cache + ) + def _compute_scheduler_metadata( self, batch_size, max_seq_len_k, cache_seqlens, cu_seqlens_q ): @@ -606,7 +616,7 @@ def forward_extend( and self.attn_cp_size > 1 ) - if save_kv_cache and not is_cp_mode: + if save_kv_cache and not is_cp_mode and not self.fa_skip_kv_cache: cache_loc = ( forward_batch.out_cache_loc if not layer.is_cross_attention @@ -764,6 +774,31 @@ def _fa_cp_attn( self.device, _fa_cp_attn, ) + elif self.fa_skip_kv_cache: + # Embedding mode: skip KV cache read and use raw K/V tensors + # directly via flash_attn_varlen_func. The KV cache write is + # also skipped (guarded above). This eliminates store_kvcache + # and prepare_varlen_num_blocks overhead per layer. + assert k is not None, "fa_skip_kv_cache requires k to be provided" + assert k_descale is None and v_descale is None, ( + "fa_skip_kv_cache uses raw K/V tensors, " + "FP8 KV cache descaling is not supported in this mode" + ) + result = flash_attn_varlen_func( + q=q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim), + k=k.view(-1, layer.tp_k_head_num, layer.head_dim), + v=v.view(-1, layer.tp_v_head_num, layer.v_head_dim), + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_q, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_q, + softmax_scale=layer.scaling, + causal=causal, + window_size=window_size, + softcap=layer.logit_cap, + num_splits=self.num_splits, + **kwargs, + ) else: result = flash_attn_with_kvcache( q=q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim), From 9fb00ede15d5297f08aaa61485b08ae4ddf959f3 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Mon, 13 Apr 2026 16:47:32 -0700 Subject: [PATCH 43/52] Clean up TokenizerManager and req_time_stats: reduce overhead and simplify (#21646) --- .../sglang/srt/managers/tokenizer_manager.py | 279 ++++++++---------- 1 file changed, 130 insertions(+), 149 deletions(-) diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index bfc5bef63aa5..09a9df885260 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -84,7 +84,6 @@ from sglang.srt.observability.metrics_collector import TokenizerMetricsCollector from sglang.srt.observability.req_time_stats import ( APIServerReqTimeStats, - calibrate_time_diff, convert_time_to_realtime, real_time, set_time_batch, @@ -205,22 +204,6 @@ def _slice_streaming_output_meta_info( meta_info[key] = meta_info[key][last_output_offset:] -def _merge_incremental_stream_meta_info( - out_list: list[dict[str, Any]], -) -> dict[str, Any]: - """Preserve delta-style output metadata when queued chunks are coalesced.""" - meta_info_list = [chunk["meta_info"] for chunk in out_list] - meta_info = dict(meta_info_list[-1]) - for key in _INCREMENTAL_STREAMING_META_INFO_KEYS: - if any(key in chunk_meta_info for chunk_meta_info in meta_info_list): - meta_info[key] = [ - item - for chunk_meta_info in meta_info_list - for item in chunk_meta_info.get(key, []) - ] - return meta_info - - class InputFormat(Enum): """Input format types for tokenization handling.""" @@ -268,9 +251,6 @@ def __init__( # Init PD disaggregation and encoder disaggregation self.init_disaggregation() - # Subprocess liveness watchdog — set by Engine or http_server after construction - self._subprocess_watchdog = None - # Init metric collector and watchdog self.init_metric_collector_watchdog() @@ -395,6 +375,9 @@ def init_running_status(self): # Session self.session_futures = {} # session_id -> asyncio event + # Subprocess liveness watchdog — set by Engine or http_server after construction + self._subprocess_watchdog = None + def init_request_logging_and_dumping(self): # TODO: Refactor and organize the log export code. # Request logging @@ -539,7 +522,6 @@ async def generate_request( # Normalize the request obj.normalize_batch_and_arguments() self._set_default_priority(obj) - self._validate_rid_not_in_flight(obj) if isinstance(obj, GenerateReqInput) and obj.routed_dp_rank is not None: dp_size = self.server_args.dp_size @@ -552,7 +534,7 @@ async def generate_request( f"routed_dp_rank={obj.routed_dp_rank} out of range [0, {dp_size})" ) - self._req_stats_init(obj, request) + self._init_req_state(obj, request) if self.server_args.language_only: self._handle_epd_disaggregation_encode_request(obj) if self.server_args.tokenizer_worker_num > 1: @@ -570,9 +552,8 @@ async def generate_request( # Tokenize the request and send it to the scheduler if obj.is_single: tokenized_obj = await self._tokenize_one_request(obj) - state = self.rid_to_state[obj.rid] self._send_one_request(tokenized_obj) - async for response in self._wait_one_response(obj, state, request): + async for response in self._wait_one_response(obj, request): yield response else: async for response in self._handle_batch_request(obj, request): @@ -825,17 +806,6 @@ async def _tokenize_one_request( obj, input_text, input_ids, input_embeds, mm_inputs, token_type_ids ) - def _validate_rid_not_in_flight( - self, obj: Union[GenerateReqInput, EmbeddingReqInput] - ) -> None: - """Validate that request IDs are not already in flight.""" - if obj.rid is None: - return - rids = obj.rid if isinstance(obj.rid, list) else [obj.rid] - conflicts = set(rids) & self.rid_to_state.keys() - if conflicts: - raise ValueError(f"Duplicate request IDs detected: {list(conflicts)}") - def _validate_one_request( self, obj: Union[GenerateReqInput, EmbeddingReqInput], input_ids: List[int] ) -> None: @@ -1204,13 +1174,90 @@ def _send_batch_request( self.send_to_scheduler.send_pyobj(batch_req) set_time_batch(tokenized_objs, "set_api_server_dispatch_finish_time") + def _coalesce_streaming_chunks( + self, + out_list: list, + rid: str, + ) -> dict: + """Coalesce multiple incremental streaming chunks into one. + + Both text and output_ids are incremental deltas, so we concatenate them; + all other fields (meta_info, etc.) are taken from the last chunk. + """ + if len(out_list) >= 20: + logger.warning( + "Streaming backlog: rid=%s, coalescing %d queued chunks into one. " + "This may inflate P99 ITL for affected requests.", + rid, + len(out_list), + ) + out = dict(out_list[-1]) + if "output_ids" in out: + out["output_ids"] = [id for chunk in out_list for id in chunk["output_ids"]] + if "text" in out: + out["text"] = "".join(chunk["text"] for chunk in out_list) + if "meta_info" in out: + meta_info_list = [chunk["meta_info"] for chunk in out_list] + meta_info = dict(meta_info_list[-1]) + for key in _INCREMENTAL_STREAMING_META_INFO_KEYS: + if any(key in m for m in meta_info_list): + meta_info[key] = [ + item for m in meta_info_list for item in m.get(key, []) + ] + out["meta_info"] = meta_info + return out + + async def _handle_abort_finish_reason( + self, + out: dict, + state: ReqState, + is_stream: bool, + ) -> Optional[dict]: + """Handle abort/error finish reasons from the scheduler. + + Returns the output dict if it should be yielded (stream abort), or None + for normal flow. Raises ValueError or HTTPException for non-stream aborts. + """ + finish_reason = out["meta_info"]["finish_reason"] + + if ( + finish_reason.get("type") == "abort" + and finish_reason.get("status_code") == HTTPStatus.BAD_REQUEST + ): + if not is_stream: + raise ValueError(finish_reason["message"]) + return out + + if finish_reason.get("type") == "abort" and finish_reason.get( + "status_code" + ) in ( + HTTPStatus.SERVICE_UNAVAILABLE, + HTTPStatus.INTERNAL_SERVER_ERROR, + ): + # Delete the key to prevent resending abort request to the scheduler and + # to ensure aborted request state is cleaned up. + 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.server_args.enable_lora and state.obj.lora_path: + await self.lora_registry.release(state.obj.lora_id) + if not is_stream: + raise fastapi.HTTPException( + status_code=finish_reason["status_code"], + detail=finish_reason["message"], + ) + return out + + return None + async def _wait_one_response( self, obj: Union[GenerateReqInput, EmbeddingReqInput], - state: ReqState, request: Optional[fastapi.Request] = None, ): """Wait for the response of one request.""" + state = self.rid_to_state[obj.rid] # Not all request types have `stream` (e.g., EmbeddingReqInput). Default to non-streaming. is_stream = getattr(obj, "stream", False) while True: @@ -1233,38 +1280,18 @@ async def _wait_one_response( continue # Drain all pending outputs atomically. - # With incremental streaming output, each chunk carries only a - # delta, so every queued chunk must be yielded to avoid dropping - # token ids. Without it, outputs are cumulative and only the - # latest chunk contains the full result, so we can safely skip - # intermediate ones. - incremental_stream = ( - is_stream and self.server_args.incremental_streaming_output - ) out_list = state.out_list state.out_list = [] finished = state.finished state.event.clear() + # With incremental streaming, each chunk is a delta — coalesce + # multiple queued chunks to avoid dropping token ids. + incremental_stream = ( + is_stream and self.server_args.incremental_streaming_output + ) if incremental_stream and len(out_list) > 1: - if len(out_list) >= 20: - logger.warning( - "Streaming backlog: rid=%s, coalescing %d queued chunks into one. " - "This may inflate P99 ITL for affected requests.", - obj.rid, - len(out_list), - ) - # Coalesce all deltas into a single chunk. Text, output_ids, - # and output-side incremental metadata all need to be merged. - out = dict(out_list[-1]) - if "output_ids" in out: - out["output_ids"] = [ - id for chunk in out_list for id in chunk["output_ids"] - ] - if "text" in out: - out["text"] = "".join(chunk["text"] for chunk in out_list) - if "meta_info" in out: - out["meta_info"] = _merge_incremental_stream_meta_info(out_list) + out = self._coalesce_streaming_chunks(out_list, obj.rid) else: out = out_list[-1] @@ -1280,7 +1307,6 @@ async def _wait_one_response( out["text"] = state.get_text() if finished: - # For non-streaming cases, response has not been sent yet (`response_sent_to_client_time` has not been set yet). # Record response sent time right before we log finished results and metrics. if not state.time_stats.response_sent_to_client_time: state.time_stats.set_response_sent_to_client_time() @@ -1294,47 +1320,19 @@ async def _wait_one_response( ) if self.request_metrics_exporter_manager.exporter_enabled(): - # Asynchronously write metrics for this request using the exporter manager. asyncio.create_task( self.request_metrics_exporter_manager.write_record(obj, out) ) # Check if this was an abort/error created by scheduler if isinstance(out["meta_info"].get("finish_reason"), dict): - finish_reason = out["meta_info"]["finish_reason"] - if ( - finish_reason.get("type") == "abort" - and finish_reason.get("status_code") == HTTPStatus.BAD_REQUEST - ): - if not is_stream: - raise ValueError(finish_reason["message"]) - else: - yield out - break - - if finish_reason.get("type") == "abort" and finish_reason.get( - "status_code" - ) in ( - HTTPStatus.SERVICE_UNAVAILABLE, - HTTPStatus.INTERNAL_SERVER_ERROR, - ): - # This is an abort request initiated by scheduler. - # Delete the key to prevent resending abort request to the scheduler and - # to ensure aborted request state is cleaned up. - 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.server_args.enable_lora and state.obj.lora_path: - await self.lora_registry.release(state.obj.lora_id) - if not is_stream: - raise fastapi.HTTPException( - status_code=finish_reason["status_code"], - detail=finish_reason["message"], - ) - else: - yield out - break + abort_out = await self._handle_abort_finish_reason( + out, state, is_stream + ) + if abort_out is not None: + yield abort_out + break + yield out break @@ -1346,8 +1344,7 @@ async def _wait_one_response( "response_sent_to_client_ts" ] = state.time_stats.get_response_sent_to_client_realtime() yield out - - if not is_stream: + else: if ( request is not None and not obj.background @@ -1377,9 +1374,7 @@ async def _handle_batch_request( # Set up generators for each request in the batch for i in range(batch_size): tmp_obj = obj[i] - state = self.rid_to_state[tmp_obj.rid] - state.obj = tmp_obj - generators.append(self._wait_one_response(tmp_obj, state, request)) + generators.append(self._wait_one_response(tmp_obj, request)) rids.append(tmp_obj.rid) else: # Sequential tokenization and processing @@ -1391,12 +1386,8 @@ async def _handle_batch_request( for i in range(batch_size): tmp_obj = obj[i] tokenized_obj = await self._tokenize_one_request(tmp_obj) - state = self.rid_to_state[tmp_obj.rid] - state.obj = tmp_obj self._send_one_request(tokenized_obj) - generators.append( - self._wait_one_response(tmp_obj, state, request) - ) + generators.append(self._wait_one_response(tmp_obj, request)) rids.append(tmp_obj.rid) else: # FIXME: When using batch and parallel_sample_num together, the perf is not optimal. @@ -1421,11 +1412,9 @@ async def _handle_batch_request( tokenized_obj.sampling_params = copy.copy(tokenized_obj.sampling_params) tokenized_obj.sampling_params.max_new_tokens = 0 tokenized_obj.stream = False - self._req_stats_init(tmp_obj) - state = self.rid_to_state[tmp_obj.rid] - tokenized_obj.time_stats = state.time_stats + self._init_req_state(tmp_obj) self._send_one_request(tokenized_obj) - await self._wait_one_response(tmp_obj, state, request).__anext__() + await self._wait_one_response(tmp_obj, request).__anext__() # Expand requests, assign new rids for them, and send them for i in range(batch_size): @@ -1433,11 +1422,10 @@ async def _handle_batch_request( tmp_obj = copy.copy(objs[i]) tokenized_obj = copy.copy(tokenized_objs[i]) tokenized_obj.rid = tmp_obj.regenerate_rid() - self._req_stats_init(tmp_obj) - state = self.rid_to_state[tmp_obj.rid] - tokenized_obj.time_stats = state.time_stats + self._init_req_state(tmp_obj) + tokenized_obj.time_stats = self.rid_to_state[tmp_obj.rid].time_stats self._send_one_request(tokenized_obj) - generators.append(self._wait_one_response(tmp_obj, state, request)) + generators.append(self._wait_one_response(tmp_obj, request)) rids.append(tmp_obj.rid) self.rid_to_state[objs[i].rid].time_stats.set_finished_time() @@ -1795,9 +1783,10 @@ def _handle_batch_output( state.time_stats.set_first_token_time() if state.finished: - state.time_stats.trace_ctx.trace_set_root_attrs( - self.convert_to_span_attrs(state, recv_obj, i) - ) + if state.time_stats.trace_ctx.tracing_enable: + state.time_stats.trace_ctx.trace_set_root_attrs( + self.convert_to_span_attrs(state, recv_obj, i) + ) state.time_stats.set_finished_time() meta_info["e2e_latency"] = state.time_stats.get_e2e_latency() @@ -2459,12 +2448,11 @@ async def _resolve_lora_path(self, obj: Union[GenerateReqInput, EmbeddingReqInpu obj.lora_id[i] if isinstance(obj.lora_id, list) else obj.lora_id ) - def _req_stats_init( + def _init_req_state( self, obj: Union[GenerateReqInput, EmbeddingReqInput], request: Optional[fastapi.Request] = None, ): - calibrate_time_diff() created_time = obj.received_time external_trace_header = None @@ -2478,39 +2466,32 @@ def _req_stats_init( external_trace_header = extract_trace_headers(request.headers) obj.external_trace_header = external_trace_header + # Normalize single/batch into a uniform list of (rid, sub_obj, bootstrap_room) if not hasattr(obj, "is_single") or obj.is_single: - time_stats = APIServerReqTimeStats(disagg_mode=self.disaggregation_mode) - state = ReqState([], False, asyncio.Event(), obj, time_stats) - self.rid_to_state[obj.rid] = state - - if self.server_args.enable_trace: - bootstrap_room = ( - obj.bootstrap_room if hasattr(obj, "bootstrap_room") else None - ) - time_stats.init_trace_ctx( - obj.rid, - bootstrap_room, - external_trace_header, - ) - time_stats.set_created_time(created_time) + items = [(obj.rid, obj, getattr(obj, "bootstrap_room", None))] else: - for i in range(len(obj.rid)): - time_stats = APIServerReqTimeStats(disagg_mode=self.disaggregation_mode) - state = ReqState([], False, asyncio.Event(), obj[i], time_stats) - self.rid_to_state[obj.rid[i]] = state - - if self.server_args.enable_trace: - bootstrap_room = ( + items = [ + ( + obj.rid[i], + obj[i], + ( obj.bootstrap_room[i] if hasattr(obj, "bootstrap_room") and obj.bootstrap_room else None - ) - time_stats.init_trace_ctx( - obj.rid[i], - bootstrap_room, - external_trace_header, - ) - time_stats.set_created_time(created_time) + ), + ) + for i in range(len(obj.rid)) + ] + + for rid, sub_obj, bootstrap_room in items: + if rid in self.rid_to_state: + raise ValueError(f"Duplicate request ID detected: {rid}") + time_stats = APIServerReqTimeStats(disagg_mode=self.disaggregation_mode) + state = ReqState([], False, asyncio.Event(), sub_obj, time_stats) + self.rid_to_state[rid] = state + if self.server_args.enable_trace: + time_stats.init_trace_ctx(rid, bootstrap_room, external_trace_header) + time_stats.set_created_time(created_time) def _should_dispatch_to_encoder( self, obj: Union[GenerateReqInput, EmbeddingReqInput] From 33a3ba256f25d0fd6a5ad3c2e6361feb69adeb50 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Mon, 13 Apr 2026 17:02:40 -0700 Subject: [PATCH 44/52] Delete dead rematch path in SessionAwareCache.release_session (#22735) --- .../sglang/srt/managers/session_controller.py | 4 +- .../srt/mem_cache/session_aware_cache.py | 81 +++--------------- .../mem_cache/test_streaming_session_unit.py | 82 ------------------- 3 files changed, 11 insertions(+), 156 deletions(-) diff --git a/python/sglang/srt/managers/session_controller.py b/python/sglang/srt/managers/session_controller.py index 9bb763fcebed..063a6192de98 100644 --- a/python/sglang/srt/managers/session_controller.py +++ b/python/sglang/srt/managers/session_controller.py @@ -329,9 +329,7 @@ def _close(self, session_id: str): node.req.multimodal_inputs = None if isinstance(self.tree_cache, SessionAwareCache): - self.tree_cache.release_session( - session_id, req if session.streaming else None - ) + self.tree_cache.release_session(session_id) del self.sessions[session_id] log_info_on_rank0( logger, f"Session closed: {session_id} (active={len(self.sessions)})" diff --git a/python/sglang/srt/mem_cache/session_aware_cache.py b/python/sglang/srt/mem_cache/session_aware_cache.py index cd5f01c706de..c453d0ebc8e5 100644 --- a/python/sglang/srt/mem_cache/session_aware_cache.py +++ b/python/sglang/srt/mem_cache/session_aware_cache.py @@ -308,81 +308,20 @@ def dec_lock_ref( # -- Session lifecycle -- - def _resolve_release_state( - self, slot: SessionSlot, req: Optional[Req] - ) -> tuple[int, Any]: - """Resolve the currently tree-owned prefix for a session slot. - - A long-lived session can outlive radix-tree splits caused by unrelated - traffic. In that case, the saved `last_node` may no longer represent the - full protected prefix even though the slot's req_to_token row still - contains tree-owned indices at the front. Re-match the current request - text, then intersect the returned tree indices with the slot's row so - release uses the prefix that is still actually backed by the tree. + def release_session(self, session_id: str): + """Release all KV resources held by a streaming session. + + `slot.last_node` + `slot.cache_protected_len` are trusted directly: radix + tree splits mutate TreeNode objects in place (see RadixCache._split_node), + so a saved TreeNode reference remains valid and the locked prefix length + is unchanged. No rematch needed -- and `match_prefix` here would cause + further splits that disturb accounting. """ - protected_len = slot.cache_protected_len - lock_node = slot.last_node - - # TODO: re-match logic disabled — match_prefix has side effects - # (splits) that disturb tree accounting. Directly using - # slot.last_node + cache_protected_len is safe after split analysis. - return protected_len, lock_node - - if ( - req is None - or not slot.is_holding_kv - or slot.req_pool_idx is None - or protected_len <= 0 - ): - return protected_len, lock_node - - from sglang.srt.mem_cache.radix_cache import RadixKey - - token_ids = (req.origin_input_ids + req.output_ids)[: slot.kv_committed_len] - if not token_ids: - return 0, None - - match = self.inner.match_prefix( - MatchPrefixParams( - key=RadixKey(token_ids=token_ids, extra_key=req.extra_key), - req=None, - ) - ) - if len(match.device_indices) == 0: - return 0, None - - max_protected_len = min(len(match.device_indices), protected_len) - row_indices = self.req_to_token_pool.req_to_token[ - slot.req_pool_idx, :max_protected_len - ].to(dtype=torch.int64) - match_indices = match.device_indices[:max_protected_len] - mismatches = (match_indices != row_indices).nonzero(as_tuple=False) - if mismatches.numel() == 0 and max_protected_len == len(match.device_indices): - common_len = max_protected_len - return common_len, match.last_device_node - - common_len = ( - int(mismatches[0].item()) if mismatches.numel() > 0 else max_protected_len - ) - if self.page_size > 1: - common_len = (common_len // self.page_size) * self.page_size - if common_len <= 0: - return 0, None - - rematch = self.inner.match_prefix( - MatchPrefixParams( - key=RadixKey(token_ids=token_ids[:common_len], extra_key=req.extra_key), - req=None, - ) - ) - return len(rematch.device_indices), rematch.last_device_node - - def release_session(self, session_id: str, req: Optional[Req] = None): - """Release all KV resources held by a streaming session.""" slot = self.slots.pop(session_id, None) if slot is None: return - protected_len, lock_node = self._resolve_release_state(slot, req) + protected_len = slot.cache_protected_len + lock_node = slot.last_node tokens_freed = ( max(0, slot.kv_allocated_len - protected_len) if slot.is_holding_kv else 0 ) diff --git a/test/registered/unit/mem_cache/test_streaming_session_unit.py b/test/registered/unit/mem_cache/test_streaming_session_unit.py index 759cff86221a..ef0b786fc231 100644 --- a/test/registered/unit/mem_cache/test_streaming_session_unit.py +++ b/test/registered/unit/mem_cache/test_streaming_session_unit.py @@ -111,88 +111,6 @@ def test_streaming_release_kv_cache_trims_overallocated_tail(monkeypatch): assert allocator.freed[0].tolist() == list(range(32, 40)) -def test_release_session_recomputes_current_tree_owned_prefix(): - page_size = 16 - req_to_token = torch.arange(128, dtype=torch.int32).reshape(1, 128) - req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[]) - allocator = _FakeAllocator() - - full_match = MatchResult( - device_indices=torch.tensor(list(range(16)) + list(range(64, 96))), - last_device_node="stale-expanded", - last_host_node="stale-expanded", - ) - protected_match = MatchResult( - device_indices=torch.tensor(list(range(16))), - last_device_node="current-protected", - last_host_node="current-protected", - ) - inner = _FakeInnerCache( - req_to_token_pool, - allocator, - page_size, - match_results=[full_match, protected_match], - ) - tree_cache = SessionAwareCache(inner) - - tree_cache.slots["session-a"] = SessionSlot( - req_pool_idx=0, - kv_committed_len=48, - kv_allocated_len=48, - last_node="outdated-node", - cache_protected_len=32, - ) - req = _FakeReq("session-a", req_pool_idx=0, committed=48, allocated=48) - - tree_cache.release_session("session-a", req) - - assert inner.dec_lock_ref_calls == ["current-protected"] - assert req_to_token_pool.free_slots == [0] - assert len(allocator.freed) == 1 - assert allocator.freed[0].tolist() == list(range(16, 48)) - - -def test_release_session_never_grows_tree_owned_prefix(): - page_size = 16 - req_to_token = torch.arange(128, dtype=torch.int32).reshape(1, 128) - req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[]) - allocator = _FakeAllocator() - - overmatched = MatchResult( - device_indices=torch.tensor(list(range(48))), - last_device_node="overmatched-node", - last_host_node="overmatched-node", - ) - capped_match = MatchResult( - device_indices=torch.tensor(list(range(16))), - last_device_node="original-lock-node", - last_host_node="original-lock-node", - ) - inner = _FakeInnerCache( - req_to_token_pool, - allocator, - page_size, - match_results=[overmatched, capped_match], - ) - tree_cache = SessionAwareCache(inner) - - tree_cache.slots["session-a"] = SessionSlot( - req_pool_idx=0, - kv_committed_len=48, - kv_allocated_len=48, - last_node="outdated-node", - cache_protected_len=16, - ) - req = _FakeReq("session-a", req_pool_idx=0, committed=48, allocated=48) - - tree_cache.release_session("session-a", req) - - assert inner.dec_lock_ref_calls == ["original-lock-node"] - assert req_to_token_pool.free_slots == [0] - assert len(allocator.freed) == 1 - assert allocator.freed[0].tolist() == list(range(16, 48)) - - def test_match_prefix_abort_does_not_restore_live_session_slot(): req_to_token = torch.arange(256, dtype=torch.int32).reshape(2, 128) req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[]) From ff61b2e47015a0ce61acac8f4983987359067662 Mon Sep 17 00:00:00 2001 From: Sahithi Chigurupati <58884509+csahithi@users.noreply.github.com> Date: Mon, 13 Apr 2026 17:08:18 -0700 Subject: [PATCH 45/52] [CI] Add workflow_dispatch and environment gate to GB200 nightly pipeline (#22733) --- .github/workflows/nightly-72-gpu-gb200.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/nightly-72-gpu-gb200.yml b/.github/workflows/nightly-72-gpu-gb200.yml index f9b3a093d58c..4d0a56419ca4 100644 --- a/.github/workflows/nightly-72-gpu-gb200.yml +++ b/.github/workflows/nightly-72-gpu-gb200.yml @@ -6,6 +6,7 @@ name: Nightly Test (GB200 72GPU) on: schedule: - cron: '0 2 * * *' # 2 AM UTC daily (offset from other nightly runs) + workflow_dispatch: # allow manual trigger; gated by gb200-ci environment concurrency: group: nightly-test-gb200 @@ -47,6 +48,7 @@ jobs: prepare-image: needs: setup if: github.repository == 'sgl-project/sglang' + environment: gb200-ci runs-on: 72-gpu-gb200 outputs: squash_file: ${{ steps.import.outputs.squash_file }} From 1ec018f27a4ac99ebfe41d1a5ebe92d11137026c Mon Sep 17 00:00:00 2001 From: yuki-brook Date: Tue, 14 Apr 2026 08:12:37 +0800 Subject: [PATCH 46/52] [Feature] Add SiMM as sglang HiCache Storage backend (#18016) --- .../sglang/srt/managers/cache_controller.py | 5 +- .../srt/mem_cache/storage/backend_factory.py | 8 + .../srt/mem_cache/storage/simm/README.md | 121 ++++ .../mem_cache/storage/simm/hicache_simm.py | 544 ++++++++++++++++++ .../srt/mem_cache/storage/simm/test_simm.py | 181 ++++++ python/sglang/srt/server_args.py | 11 +- 6 files changed, 868 insertions(+), 2 deletions(-) create mode 100644 python/sglang/srt/mem_cache/storage/simm/README.md create mode 100644 python/sglang/srt/mem_cache/storage/simm/hicache_simm.py create mode 100644 python/sglang/srt/mem_cache/storage/simm/test_simm.py diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index ddd34fd1688f..a78e4059d6ab 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -483,7 +483,10 @@ def attach_storage_backend( self.page_get_func = self._generic_page_get self.page_set_func = self._generic_page_set - if (self.storage_backend_type in ["hf3fs", "mooncake", "eic", "nixl"]) or ( + if ( + self.storage_backend_type + in ["hf3fs", "mooncake", "eic", "nixl", "simm"] + ) or ( self.storage_backend_type == "dynamic" and bool(self.storage_config.extra_config.get("interface_v1", 0)) ): diff --git a/python/sglang/srt/mem_cache/storage/backend_factory.py b/python/sglang/srt/mem_cache/storage/backend_factory.py index eaa5a3e18c64..1fe731520508 100644 --- a/python/sglang/srt/mem_cache/storage/backend_factory.py +++ b/python/sglang/srt/mem_cache/storage/backend_factory.py @@ -183,6 +183,8 @@ def _create_builtin_backend( return backend_class.from_env_config(bytes_per_page, dtype, storage_config) elif backend_name == "eic": return backend_class(storage_config, mem_pool_host) + elif backend_name == "simm": + return backend_class(storage_config, mem_pool_host) else: raise ValueError(f"Unknown built-in backend: {backend_name}") @@ -221,3 +223,9 @@ def _create_builtin_backend( "sglang.srt.mem_cache.storage.eic.eic_storage", "EICStorage", ) + +StorageBackendFactory.register_backend( + "simm", + "sglang.srt.mem_cache.storage.simm.hicache_simm", + "HiCacheSiMM", +) diff --git a/python/sglang/srt/mem_cache/storage/simm/README.md b/python/sglang/srt/mem_cache/storage/simm/README.md new file mode 100644 index 000000000000..67418cdac517 --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/simm/README.md @@ -0,0 +1,121 @@ +# SiMM as L3 KV Cache + +This document describes how to use SiMM as the L3 KV cache for SGLang. + +## About SiMM + +SiMM(Scalable In-Memory Middleware) is a distributed, high-performance, elastic cache acceleration layer for all AI workloads. + +For more details about SiMM, please refer to [SiMM project](https://github.com/scitix/SiMM) and [SiMM documents](https://github.com/scitix/SiMM/tree/main/docs). + +### SiMM & SGLang HiCache + +SiMM serves as a high-performance L3 storage backend for SGLang HiCache, enabling distributed KV cache storage across multiple servers with RDMA-baed transport. This integration addresses the capacity limitations of traditional GPU-only or GPU+CPU caching by providing virtually unlimited cache storage through a distributed memory pool. + +When a cache miss occurs in L1 and L2, HiCache automatically fetches the required KV cache from SiMM's distributed memory pool. The system uses intelligent prefetching strategies to minimize latency, and utilize RDMA technology and zero-copy technique to ensure high-bandwidth, low-latency data transfer between SGLang instances and SiMM data servers. + +## Install SiMM + +**from source** + +Clone SiMM project: + +```bash +git clone https://github.com/scitix/SiMM --recursive +``` + +Install dependencies: + +```bash +cd SiMM +bash configure.sh +``` + +Build and install SiMM: + +```bash +bash build.sh --mode=release --clean +``` + +For more details, please refer to [SiMM official installation guide](https://github.com/scitix/SiMM/blob/main/README.md). + +## Deployment + +**SiMM** + +Before launch `SGLang server` with SiMM, you should launch SiMM `cluster manager service` and `data server service`. + +You can visit [SiMM official deploy guide](https://github.com/scitix/SiMM/blob/main/docs/deploy_guide.md) and deploy SiMM on your K8S cluster with RDMA network. + +**Start the `SGLang server` with SiMM enabled:** + +There are three ways to configure SiMM: + +1. Via extra configuration passed through sglang parameters +2. Using JSON configuration files +3. Using environment variables + +SiMM loads configuration in the following priority order: + +1. If SiMM-specific options are provided in `--hicache-storage-backend-extra-config`, they are used first. +2. If not, SiMM checks whether the environment variable `DEFAULT_SIMM_CONFIG_PATH_ENV` is set, and loads the JSON config file from that path. +3. If neither of the above is provided, SiMM falls back to environment variables. + +**HiCache Related Parameters for SGLang Server** + +For a comprehensive overview of HiCache-related parameters, please refer to [this document](https://docs.sglang.io/advanced_features/hicache_design.html#related-parameters). + + +Note that, for `--hicache-mem-layout {layer_first,page_first,page_first_direct}`, which specifies the memory layout for the host memory pool, `page_first` or `page_first_direct` are required if use SiMM backend. + +### Distributed Deployment + +**Using extra-config of sglang arguments to configure SiMM** + +```bash +python -m sglang.launch_server \ + --enable-hierarchical-cache \ + --hicache-storage-backend simm \ + --model-path [model_path] \ + --hicache-storage-backend-extra-config '{"manager_address": "127.0.0.1:30001"}' +``` + +**Using JSON file to configure SiMM** + +SGLang server can load SiMM config from `SGLANG_HICACHE_SIMM_CONFIG_PATH`. + +```bash +export SGLANG_HICACHE_SIMM_CONFIG_PATH=/sgl-workspace/sglang/benchmark/hicache/simm_config.json + +echo '{ + "manager_address": "127.0.0.1:30001" +}' > ${SGLANG_HICACHE_SIMM_CONFIG_PATH} + +python -m sglang.launch_server \ + --enable-hierarchical-cache \ + --hicache-storage-backend simm \ + --model-path [model_path] +``` + +**Using env variables to configure SiMM** + +```bash +SIMM_CLUSTER_MANAGER="127.0.0.1:30001" +python -m sglang.launch_server \ + --enable-hierarchical-cache \ + --hicache-storage-backend simm \ + --model-path [model_path] +``` + +## Test SiMM + +This test is intended for developers to quickly verify that the SiMM class interfaces are functioning correctly. + +First, start the `cluster manager service` and `data server service`. Then run the `test_hicache_simm.py`. + +```bash +SIMM_CLUSTER_MANAGER="127.0.0.1:30001" \ +python3 [path of test_hicache_simm.py] +``` + +If all tests pass, the message "✅ All tests passed" will be printed at the end. diff --git a/python/sglang/srt/mem_cache/storage/simm/hicache_simm.py b/python/sglang/srt/mem_cache/storage/simm/hicache_simm.py new file mode 100644 index 000000000000..8b29fd2d21d6 --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/simm/hicache_simm.py @@ -0,0 +1,544 @@ +import json +import logging +import os +import re +import time +import uuid +from collections import defaultdict +from dataclasses import dataclass +from datetime import datetime +from typing import Any, Dict, List, Optional + +import torch + +from sglang.srt.mem_cache.hicache_storage import ( + HiCacheStorage, + HiCacheStorageConfig, + HiCacheStorageExtraInfo, +) +from sglang.srt.mem_cache.memory_pool_host import HostKVCache + +# Third Party +try: + from simm.kv import BlockView, Store, register_mr, set_flag +except ImportError as e: + raise ImportError( + "Please install simm by following the instructions at https://github.com/scitix/SiMM " + "to run SGLang with SimmConnector." + ) from e + +SGLANG_HICACHE_SIMM_JSON_ENV_VAR = "SGLANG_HICACHE_SIMM_CONFIG_PATH" + +logger = logging.getLogger(__name__) + + +@dataclass +class SiMMConfig: + manager_address: str + clnt_threadpool_size: int + enable_profile: bool + + @staticmethod + def from_file() -> "SiMMConfig": + """Load the config from a JSON file.""" + if os.environ.get(SGLANG_HICACHE_SIMM_JSON_ENV_VAR) is None: + raise RuntimeError( + f"Config file path not set. Please set {SGLANG_HICACHE_SIMM_JSON_ENV_VAR}" + ) + file_path = os.environ.get(SGLANG_HICACHE_SIMM_JSON_ENV_VAR) + try: + with open(file_path) as fin: + config = json.load(fin) + except Exception as e: + raise RuntimeError(f"Failed to load config from {file_path}: {str(e)}") + + if "manager_address" not in config: + raise ValueError("Manager_address is required in config file") + + return SiMMConfig( + manager_address=config.get("manager_address"), + clnt_threadpool_size=config.get("clnt_threadpool_size", 10), + enable_profile=config.get("enable_profile", False), + ) + + @staticmethod + def load_from_extra_config(extra_config: dict) -> "SiMMConfig": + """Load config from extra_config dictionary.""" + if "manager_address" not in extra_config: + raise ValueError("manager_address is required in extra_config") + + return SiMMConfig( + manager_address=extra_config.get("manager_address"), + clnt_threadpool_size=extra_config.get("clnt_threadpool_size", 10), + enable_profile=extra_config.get("enable_profile", False), + ) + + +def get_current_process_numa() -> int: + """ + Return value: numa_node of current process, failed return -1 + """ + try: + # get current cpu + with open("/proc/self/stat", "r") as f: + stat_data = f.read() + + # the 39th field is processor + fields = stat_data.split() + if len(fields) < 39: + return -1 + current_cpu = int(fields[38]) + numa_path = f"/sys/devices/system/cpu/cpu{current_cpu}/node0" + if os.path.exists(numa_path) and os.path.islink(numa_path): + link_target = os.readlink(numa_path) + # parse numa node from path + match = re.search(r"node(\d+)$", link_target) + if match: + return int(match.group(1)) + + return -1 + except Exception: + return -1 + + +def get_numa_nic_mapping() -> Dict[int, List[str]]: + """ + Return value: Dict[numa_node, List(rdma_device_name)] + """ + ib_root = "/sys/class/infiniband" + device_map = defaultdict(list) + + if not os.path.exists(ib_root): + logger.error(f"SiMM ERROR: {ib_root} not found. Are RDMA drivers loaded?") + return [] + + for device_name in os.listdir(ib_root): + numa_path = os.path.join(ib_root, device_name, "device", "numa_node") + numa_node = -1 # default value, if system is UMA. + + try: + if os.path.exists(numa_path): + with open(numa_path, "r") as f: + content = f.read().strip() + numa_node = int(content) + except (IOError, ValueError): + pass + device_map[numa_node].append(device_name) + + return device_map + + +class HiCacheSiMM(HiCacheStorage): + + def __init__( + self, storage_config: HiCacheStorageConfig = None, mem_pool: HostKVCache = None + ): + try: + extra_config = ( + getattr(storage_config, "extra_config", None) + if storage_config + else None + ) + # Load configuration with manager_address prioritized from extra_config if available + if ( + extra_config is not None + and extra_config.get("manager_address") is not None + ): + # Load from extra_config + self.config = SiMMConfig.load_from_extra_config(extra_config) + logger.info("SiMM Configuration loaded from extra_config successfully.") + else: + # Load from config file + self.config = SiMMConfig.from_file() + logger.info("SiMM Configuration loaded from file successfully.") + + # Check if extra_backend_tag should be passed to SiMM data server + self.extra_backend_tag = None + if extra_config and "extra_backend_tag" in extra_config: + self.extra_backend_tag = extra_config["extra_backend_tag"] + logger.info(f"Using extra_backend_tag: {self.extra_backend_tag}") + + # Set nic device according to current process numa node + nic_mapping = get_numa_nic_mapping() + logger.info(f"SiMM NUMA-awared allocation: {nic_mapping}") + current_numa = get_current_process_numa() + if current_numa >= 0: + rdma_devices = nic_mapping.get(current_numa) + if rdma_devices is not None and len(rdma_devices) > 0: + rdma_device_str = ",".join(rdma_devices) + os.environ["SICL_NET_DEVICES"] = rdma_device_str + logger.info(f"SiMM using rdma {rdma_device_str}") + + # Set simm log path: /var/log/simm/{filename_ts}-{pid}/simm_clnt.log + filename_ts = datetime.now().strftime("%Y%m%d-%H%M%S") + log_file_path: str = ( + f"/var/log/simm/{filename_ts}-{os.getpid()}/simm_clnt.log" + ) + + cm_ip = self.config.manager_address.split(":")[0] + cm_port = self.config.manager_address.split(":")[1] + set_flag("cm_primary_node_ip", cm_ip) + set_flag("cm_primary_node_port", cm_port) + set_flag("clnt_log_file", log_file_path) + set_flag("clnt_thread_pool_size", str(self.config.clnt_threadpool_size)) + + self.store = Store() + logger.info("SiMM store setup successfully.") + self.mr_ext = None + + self.warmup() + logger.info("SiMM store warmup successfully.") + + if storage_config is not None: + self.model_name = storage_config.model_name + self.is_mla_backend = storage_config.is_mla_model + self.local_rank = storage_config.tp_rank + self.pp_rank = storage_config.pp_rank + self.pp_size = storage_config.pp_size + else: + self.model_name = "" + self.is_mla_backend = False + self.local_rank = 0 + self.pp_rank = 0 + self.pp_size = 1 + + self.enable_pp = self.pp_size > 1 + if self.enable_pp: + self.mha_suffix = f"{self.local_rank}_{self.pp_rank}" + self.mla_suffix = f"{self.pp_rank}" + else: + self.mha_suffix = f"{self.local_rank}" + self.mla_suffix = "" + + except ValueError as e: + logger.error("Configuration loading failed: %s", e) + raise + except Exception as exc: + logger.error("An error occurred while loading the configuration: %s", exc) + raise + + def warmup(self): + """Dryrun a key to warmup SiMM client""" + logger.info("begin warm up SiMM client") + start_time = time.perf_counter_ns() + warmup_key = "sglang_simm_warmup_key" + uuid.uuid4().hex + warmup_tensor = torch.frombuffer( + bytearray(warmup_key.encode()), dtype=torch.uint8 + ) + warmup_size = 4 * 1024 # 4 KB + block = self.store.allocate(warmup_size) + block_ = block.as_ref() + block_[: len(warmup_key)] = warmup_tensor + if self.store.put(warmup_key, block.view()) != 0: + logger.warning(f"SiMM client warmup put key {warmup_key} failed") + if not self.store.exists(warmup_key): + logger.warning(f"SiMM client warmup key {warmup_key} not exists") + got_block = self.store.allocate(warmup_size) + if self.store.get(warmup_key, got_block.view()) < 0: + logger.warning(f"SiMM client warmup get key {warmup_key} failed") + if not all(got_block.as_ref()[: len(warmup_key)] == warmup_tensor): + logger.warning(f"SiMM client warmup key {warmup_key} data wrong") + logger.info( + f"finish SiMM client warm up, cost {(time.perf_counter_ns() - start_time)/1000:.2f} us" + ) + + def register_mem_pool_host(self, mem_pool_host: HostKVCache): + super().register_mem_pool_host(mem_pool_host) + assert self.mem_pool_host.layout in [ + "page_first", + "page_first_direct", + ], "simm storage backend only support page first or page first direct layout" + buffer = self.mem_pool_host.kv_buffer + try: + self.mr_ext = register_mr(buffer) + if self.mr_ext is None: + logger.error( + f"Failed to register buffer, {buffer=}, please check buffer and RDMA network" + ) + raise RuntimeError(f"Failed to register buffer to SiMM") + except TypeError as err: + logger.error("Failed to register buffer to SiMM: %s", err) + raise TypeError("SiMM Register Buffer Error.") from err + + def _get_mha_buffer_meta(self, keys, indices): + ptr_list, element_size_list = self.mem_pool_host.get_page_buffer_meta(indices) + key_list = [] + for key_ in keys: + key_list.append(f"{key_}_{self.mha_suffix}_k") + key_list.append(f"{key_}_{self.mha_suffix}_v") + if len(key_list) != len(ptr_list): + logger.error( + f"key size {len(key_list)} not equal with incides ptr size {len(ptr_list)}" + ) + assert len(key_list) == len(ptr_list) + return key_list, ptr_list, element_size_list + + def _get_mla_buffer_meta(self, keys, indices): + ptr_list, element_size_list = self.mem_pool_host.get_page_buffer_meta(indices) + key_list = [] + for key_ in keys: + key_list.append(f"{key_}_{self.mla_suffix}_k") + if len(key_list) != len(ptr_list): + logger.error( + f"key size {len(key_list)} not equal with incides ptr size {len(ptr_list)}" + ) + assert len(key_list) == len(ptr_list) + return key_list, ptr_list, element_size_list + + def _batch_preprocess(self, keys, host_indices): + assert len(keys) > 0 + assert len(keys) == len(host_indices) // self.mem_pool_host.page_size + if self.is_mla_backend: + return self._get_mla_buffer_meta(keys, host_indices) + else: + return self._get_mha_buffer_meta(keys, host_indices) + + def _batch_postprocess(self, results: List[int], is_set_operate=False): + """ + for batch_get_into, results is Vector of integers, + where each element is the number of bytes read on success, or a negative value on error + for batch_put_from, results is Vector of integers, + where each element is 0 on success, or a negative value on error + """ + if self.is_mla_backend: + return [k_res == 0 if is_set_operate else k_res > 0 for k_res in results] + else: + kv_pairs = zip(results[::2], results[1::2]) + return [ + ( + (k_res == 0 and v_res == 0) + if is_set_operate + else (k_res > 0 and v_res > 0) + ) + for k_res, v_res in kv_pairs + ] + + def batch_get_v1( + self, + keys: List[str], + host_indices: torch.Tensor, + extra_info: Optional[HiCacheStorageExtraInfo] = None, + ) -> List[bool]: + # Apply extra_backend_tag prefix if available + if self.extra_backend_tag is not None: + prefix = self.extra_backend_tag + keys = [f"{prefix}_{key}" for key in keys] + + t1 = time.perf_counter_ns() + key_strs, buffer_ptrs, buffer_sizes = self._batch_preprocess(keys, host_indices) + get_results = self._get_batch_zero_copy_impl( + key_strs, buffer_ptrs, buffer_sizes + ) + t2 = time.perf_counter_ns() + total_size = sum([k_res if k_res > 0 else 0 for k_res in get_results]) + if self.config.enable_profile: + logger.info( + f"SiMM batch_get_v1 {len(keys)} keys, total size: {total_size / 1024**2} MiB, \ + using {(t2 - t1)/1000} us, Throughput: {total_size / 1024**3 / ((t2 - t1) / 1000**3):.2f} GiB/s" + ) + return self._batch_postprocess(get_results, is_set_operate=False) + + def batch_set_v1( + self, + keys: List[str], + host_indices: torch.Tensor, + extra_info: Optional[HiCacheStorageExtraInfo] = None, + ) -> List[bool]: + # Apply extra_backend_tag prefix if available + if self.extra_backend_tag is not None: + prefix = self.extra_backend_tag + keys = [f"{prefix}_{key}" for key in keys] + + t1 = time.perf_counter_ns() + key_strs, buffer_ptrs, buffer_sizes = self._batch_preprocess(keys, host_indices) + exist_result = self._batch_exist_impl(key_strs) + t2 = time.perf_counter_ns() + if self.config.enable_profile: + logger.info( + f"SiMM batch exists {len(keys)} keys, using {(t2 - t1)/1000} us" + ) + + set_keys = [] + set_buffer_ptrs = [] + set_buffer_sizes = [] + set_indices = [] + set_results = [-1] * len(key_strs) + total_size = 0 + for i in range(len(key_strs)): + if not exist_result[i]: + set_keys.append(key_strs[i]) + set_buffer_ptrs.append(buffer_ptrs[i]) + set_buffer_sizes.append(buffer_sizes[i]) + set_indices.append(i) + total_size += buffer_sizes[i] + else: + set_results[i] = 0 + + # Only set non-existing keys to storage + if len(set_keys) > 0: + put_results = self._put_batch_zero_copy_impl( + set_keys, set_buffer_ptrs, set_buffer_sizes + ) + for i in range(len(set_indices)): + set_results[set_indices[i]] = put_results[i] + t3 = time.perf_counter_ns() + if self.config.enable_profile: + logger.info( + f"SiMM batch_put_v1 {len(keys)} keys, total size: {total_size / 1024**2} MiB, \ + using {(t3 - t2)/1000} us, Throughput: {total_size / 1024**3 / ((t3 - t2) / 1000**3):.2f} GiB/s" + ) + + return self._batch_postprocess(set_results, is_set_operate=True) + + def set( + self, + key, + value: Optional[Any] = None, + target_location: Optional[List[int]] = None, + target_sizes: Optional[List[int]] = None, + ) -> bool: + # Only support zero copy set for now + assert target_location is not None and target_sizes is not None + exist_result = self._batch_exist_impl([key]) + if exist_result[0]: + return True + put_result = self._put_batch_zero_copy_impl( + [key], [target_location], [target_sizes] + ) + return put_result[0] == 0 + + def batch_set( + self, + keys: List[str], + values: Optional[List[torch.Tensor]] = None, + target_locations: Optional[List[int]] = None, + target_sizes: Optional[List[int]] = None, + ) -> bool: + # Only support zero copy set for now + assert target_locations is not None and target_sizes is not None + assert len(keys) == len(target_locations) == len(target_sizes) + + if len(keys) == 0: + return False + + for i in range(len(keys)): + if ( + keys[i] is None + or target_locations[i] is None + or target_sizes[i] is None + ): + return False + + exist_result = self._batch_exist_impl(keys) + set_keys = [] + set_target_locations = [] + set_target_sizes = [] + set_indices = [] + for i in range(len(keys)): + if not exist_result[i]: + set_keys.append(keys[i]) + set_target_locations.append(target_locations[i]) + set_target_sizes.append(target_sizes[i]) + set_indices.append(i) + # Only set non-existing keys to storage + put_result = self._put_batch_zero_copy_impl( + set_keys, set_target_locations, set_target_sizes + ) + for i in range(len(set_indices)): + if put_result[i] == 0: + exist_result[set_indices[i]] = 1 + + # return the number of consecutive successful operations from the start. + success_count = 0 + for i in range(len(keys)): + if exist_result[i] == 0: + break + success_count += 1 + return success_count == len(keys) + + def get( + self, + key, + target_location: Optional[Any] = None, + target_sizes: Optional[Any] = None, + ) -> bool: + assert target_location is not None and target_sizes is not None + get_result = self._get_batch_zero_copy_impl( + [key], [target_location], [target_sizes] + ) + return get_result[0] >= 0 + + def batch_get( + self, + keys: List[str], + target_locations: Optional[Any] = None, + target_sizes: Optional[Any] = None, + ) -> int: + assert len(keys) == len(target_locations) == len(target_sizes) + if len(keys) == 0: + return 0 + get_result = self._get_batch_zero_copy_impl( + keys, target_locations, target_sizes + ) + if self.is_mla_backend: + key_multiplier = 1 + else: + key_multiplier = 2 + for i in range(len(keys)): + if get_result[i] < 0: + return i // key_multiplier + return len(keys) // key_multiplier + + def exists(self, key) -> bool: + exist_result = self._batch_exist_impl([key]) + return exist_result[0] + + def batch_exists( + self, keys, extra_info: Optional[HiCacheStorageExtraInfo] = None + ) -> int: + if self.is_mla_backend: + query_keys = [f"{key}_{self.mla_suffix}_k" for key in keys] + key_multiplier = 1 + else: + query_keys = [] + for key in keys: + query_keys.append(f"{key}_{self.mha_suffix}_k") + query_keys.append(f"{key}_{self.mha_suffix}_v") + key_multiplier = 2 + + t1 = time.perf_counter_ns() + exist_result = self._batch_exist_impl(query_keys) + t2 = time.perf_counter_ns() + if self.config.enable_profile: + logger.info( + f"SiMM batch exists {len(keys)} keys, using {(t2 - t1)/1000} us" + ) + for i in range(len(query_keys)): + if not exist_result[i]: + return i // key_multiplier + return len(query_keys) // key_multiplier + + def _put_batch_zero_copy_impl( + self, key_strs: List[str], buffer_ptrs: List[int], buffer_sizes: List[int] + ) -> List[int]: + block_views = [] + for i in range(len(buffer_ptrs)): + block_view = BlockView.from_buffer( + buffer_ptrs[i], buffer_sizes[i], self.mr_ext + ) + block_views.append(block_view) + return self.store.mput(key_strs, block_views) + + def _get_batch_zero_copy_impl( + self, key_strs: List[str], buffer_ptrs: List[int], buffer_sizes: List[int] + ) -> List[int]: + block_views = [] + for i in range(len(buffer_ptrs)): + block_view = BlockView.from_buffer( + buffer_ptrs[i], buffer_sizes[i], self.mr_ext + ) + block_views.append(block_view) + return self.store.mget(key_strs, block_views) + + def _batch_exist_impl(self, key_strs: List[str]) -> List[bool]: + return self.store.mexists(key_strs) diff --git a/python/sglang/srt/mem_cache/storage/simm/test_simm.py b/python/sglang/srt/mem_cache/storage/simm/test_simm.py new file mode 100644 index 000000000000..74e4ec65b4cd --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/simm/test_simm.py @@ -0,0 +1,181 @@ +import logging +import uuid + +import torch + +from python.sglang.srt.mem_cache.storage.simm.hicache_simm import HiCacheSiMM +from sglang.srt.mem_cache.hicache_storage import HiCacheStorageConfig + +logging.basicConfig( + level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" +) +logger = logging.getLogger(__name__) + + +def generate_batch_query_keys(kv_num: int, config: HiCacheStorageConfig): + keys = ["test_" + str(uuid.uuid4()) for _ in range(kv_num)] + set_keys = [] + for key in keys: + if config.is_mla_model: + set_keys.append(key + "_k") + else: + set_keys.append(key + f"_{config.tp_rank}_k") + set_keys.append(key + f"_{config.tp_rank}_v") + get_keys = set_keys + exist_keys = keys + return set_keys, get_keys, exist_keys + + +def create_mock_host_kv_cache(buffer_size, dtype=torch.float32): + """Create a mock HostKVCache-like object for testing.""" + buffer = torch.randn(buffer_size, dtype=dtype) + + class MockHostKVCache: + def __init__(self, buffer): + self.kv_buffer = buffer + self.layout = "page_first" + self.page_size = 1 # Simple page size for testing + + def get_page_buffer_meta(self, indices): + """Mock implementation of get_page_buffer_meta.""" + ptr_list = [] + element_size_list = [] + for idx in indices: + # Create mock pointers and sizes for each page + ptr_list.append(idx * self.page_size * self.kv_buffer.element_size()) + element_size_list.append(self.page_size * self.kv_buffer.element_size()) + return ptr_list, element_size_list + + return MockHostKVCache(buffer), buffer + + +def test_single_operation(): + """Test the set API with a single key-value pair.""" + print("=" * 100) + print("Testing single operation") + + buffer_size = 1024 * 1024 * 16 # 16MB + value_elements = 1024 + store = HiCacheSiMM() + mock_host_kv_cache, buffer = create_mock_host_kv_cache(buffer_size) + + # Register the memory pool host - this is the proper workflow + store.register_mem_pool_host(mock_host_kv_cache) + + value_size = value_elements * buffer.element_size() + + key = str(uuid.uuid4()) + set_slice = buffer[:value_elements] + get_slice = buffer[value_elements : 2 * value_elements] + set_location = set_slice.data_ptr() + get_location = get_slice.data_ptr() + + # Test set operation + result = store.set(key, target_location=set_location, target_sizes=value_size) + assert result is True, f"❌set operation failed for key: {key}" + + # Test exists operation + assert store.exists(key), f"❌key {key} should exist after set operation" + + # Test get operation + result = store.get(key, target_location=get_location, target_sizes=value_size) + assert result is True, f"❌get operation failed for key: {key}" + + # Compare the data using proper tensor indices + assert torch.allclose( + set_slice, get_slice, atol=1e-6 + ), f"❌get operation failed for key: {key}" + + logger.info(f"✅ Single operation passed") + + +def test_batch_operation(config: HiCacheStorageConfig): + """Test the batch set/get APIs with multiple key-value pairs.""" + print("=" * 100) + print(f"Testing batch operation with config: {config}") + + buffer_size = 1024 * 1024 * 16 # 16MB + value_elements = 256 + kv_num = 13 + store = HiCacheSiMM(config) + mock_host_kv_cache, buffer = create_mock_host_kv_cache(buffer_size) + + store.register_mem_pool_host(mock_host_kv_cache) + + value_size = value_elements * buffer.element_size() + + set_keys, get_keys, exist_keys = generate_batch_query_keys(kv_num, config) + set_slices = [ + buffer[i * value_elements : (i + 1) * value_elements] + for i in range(len(set_keys)) + ] + set_indices = torch.cat(set_slices) + + # Test batch set operation + result = store.batch_set_v1(set_keys, set_indices) + assert all(result), f"❌batch set operation failed" + + # Test batch exists operation + assert store.batch_exists( + exist_keys + ), f"❌keys should exist after batch set operation" + + # Test batch get operation + get_slices = [ + buffer[ + (len(set_keys) + i) + * value_elements : (len(set_keys) + i + 1) + * value_elements + ] + for i in range(len(get_keys)) + ] + get_indices = torch.cat(get_slices) + result = store.batch_get_v1(get_keys, get_indices) + assert all(result), f"❌batch get operation failed" + for i in range(len(get_keys)): + assert torch.allclose( + set_slices[i], get_slices[i], atol=1e-6 + ), f"❌batch get operation failed for key: {get_keys[i]}" + + logger.info(f"✅ Batch operation passed") + + +if __name__ == "__main__": + test_single_operation() + test_batch_operation( + HiCacheStorageConfig( + is_mla_model=False, + tp_rank=0, + tp_size=1, + model_name=None, + is_page_first_layout=True, + ) + ) + test_batch_operation( + HiCacheStorageConfig( + is_mla_model=True, + tp_rank=0, + tp_size=1, + model_name=None, + is_page_first_layout=True, + ) + ) + test_batch_operation( + HiCacheStorageConfig( + is_mla_model=False, + tp_rank=1, + tp_size=4, + model_name=None, + is_page_first_layout=True, + ) + ) + test_batch_operation( + HiCacheStorageConfig( + is_mla_model=True, + tp_rank=3, + tp_size=8, + model_name=None, + is_page_first_layout=True, + ) + ) + logger.info(f"✅ All tests passed") diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 8c7cdb12a38b..8cea237cc2d5 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -5519,7 +5519,16 @@ def add_cli_args(parser: argparse.ArgumentParser): parser.add_argument( "--hicache-storage-backend", type=str, - choices=["file", "mooncake", "hf3fs", "nixl", "aibrix", "dynamic", "eic"], + choices=[ + "file", + "mooncake", + "hf3fs", + "nixl", + "aibrix", + "dynamic", + "eic", + "simm", + ], default=ServerArgs.hicache_storage_backend, help="The storage backend for hierarchical KV cache. " "Built-in backends: file, mooncake, hf3fs, nixl, aibrix. " From c456cba7fdec0ff0f7a3d7b0395d22029f33ccde Mon Sep 17 00:00:00 2001 From: Zhai Feiyue <80079571+ZhaiFeiyue@users.noreply.github.com> Date: Tue, 14 Apr 2026 08:28:00 +0800 Subject: [PATCH 47/52] [gateway] Support SGLANG_LOG_MS for millisecond precision in router logs (#22506) --- sgl-model-gateway/src/observability/logging.rs | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/sgl-model-gateway/src/observability/logging.rs b/sgl-model-gateway/src/observability/logging.rs index 188d9b1b1584..a85fbdd7f73b 100644 --- a/sgl-model-gateway/src/observability/logging.rs +++ b/sgl-model-gateway/src/observability/logging.rs @@ -16,8 +16,16 @@ use super::otel_trace::get_otel_layer; use crate::config::TraceConfig; const TIME_FORMAT: &str = "%Y-%m-%d %H:%M:%S"; +const TIME_FORMAT_MS: &str = "%Y-%m-%d %H:%M:%S%.3f"; const DEFAULT_LOG_TARGET: &str = "smg"; +fn get_time_format() -> &'static str { + match std::env::var("SGLANG_LOG_MS") { + Ok(v) if matches!(v.trim().to_lowercase().as_str(), "true" | "1") => TIME_FORMAT_MS, + _ => TIME_FORMAT, + } +} + #[derive(Debug, Clone)] pub struct LoggingConfig { pub level: Level, @@ -102,11 +110,13 @@ pub fn init_logging(config: LoggingConfig, otel_layer_config: Option Date: Tue, 14 Apr 2026 09:03:25 +0800 Subject: [PATCH 48/52] [HiSparse] Clarify decode token usage logs (#22331) --- .../srt/managers/hisparse_coordinator.py | 24 +++++++++++++++ .../scheduler_runtime_checker_mixin.py | 30 +++++++++++++++++++ 2 files changed, 54 insertions(+) diff --git a/python/sglang/srt/managers/hisparse_coordinator.py b/python/sglang/srt/managers/hisparse_coordinator.py index 9336571976ba..07c3aa7aea24 100644 --- a/python/sglang/srt/managers/hisparse_coordinator.py +++ b/python/sglang/srt/managers/hisparse_coordinator.py @@ -27,6 +27,13 @@ class HiSparseAct(NamedTuple): req: Req +class HiSparseTokenStats(NamedTuple): + device_tokens: int + device_token_usage: float + host_tokens: int + host_token_usage: float + + class HiSparseCoordinator: def __init__( self, @@ -125,6 +132,23 @@ def __init__( def set_decode_producer_stream(self, stream) -> None: self.decode_producer_stream = stream + def get_token_stats(self) -> HiSparseTokenStats: + device_allocator = self.token_to_kv_pool_allocator.hisparse_attn_allocator + device_capacity = device_allocator.size + device_tokens = device_capacity - device_allocator.available_size() + host_capacity = self.mem_pool_host.size + host_tokens = host_capacity - self.mem_pool_host.available_size() + return HiSparseTokenStats( + device_tokens=device_tokens, + device_token_usage=( + device_tokens / device_capacity if device_capacity > 0 else 0.0 + ), + host_tokens=host_tokens, + host_token_usage=( + host_tokens / host_capacity if host_capacity > 0 else 0.0 + ), + ) + def admit_request_into_staging(self, req: Req) -> None: req.hisparse_staging = True logical_indices = self.req_to_token_pool.req_to_token[ diff --git a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py index 189afae5b51b..ef8e1f3fda2c 100644 --- a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py +++ b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py @@ -31,6 +31,7 @@ class PoolStats: is_hybrid_swa: bool = False is_hybrid_ssm: bool = False + is_hisparse: bool = False # For hybrid-swa pools swa_num_used: Optional[int] = None @@ -44,6 +45,12 @@ class PoolStats: mamba_available_size: Optional[int] = None mamba_evictable_size: Optional[int] = None + # HiSparse device/host breakdown for decode logs (plain KV pool only) + hisparse_device_tokens: Optional[int] = None + hisparse_device_token_usage: Optional[float] = None + hisparse_host_tokens: Optional[int] = None + hisparse_host_token_usage: Optional[float] = None + def get_kv_token_stats(self) -> Tuple[int, float]: # NOTE: mamba pool is not included in the "token usage" calculation. if self.is_hybrid_swa: @@ -98,6 +105,13 @@ def get_decode_usage_msg_parts(self) -> List[str]: f"mamba num: {self.mamba_num_used}", f"mamba usage: {self.mamba_usage:.2f}", ] + if self.is_hisparse: + parts += [ + f"#gpu token: {self.hisparse_device_tokens}", + f"gpu token usage: {self.hisparse_device_token_usage:.2f}", + f"#cpu token: {self.hisparse_host_tokens}", + f"cpu token usage: {self.hisparse_host_token_usage:.2f}", + ] if not parts: parts.append( f"#token: {self.full_num_used}, token usage: {self.full_token_usage:.2f}" @@ -149,6 +163,8 @@ def get_pool_stats(self: Scheduler) -> PoolStats: pool_stats = self._get_swa_token_info() elif self.is_hybrid_ssm: return self._get_mamba_token_info() + elif self.enable_hisparse: + return self._get_hisparse_token_info() else: return self._get_token_info() @@ -175,6 +191,20 @@ def _get_token_info(self: Scheduler) -> PoolStats: full_evictable_size=evictable_size, ) + def _get_hisparse_token_info(self: Scheduler) -> PoolStats: + pool_stats = self._get_token_info() + if self.enable_hisparse and self.hisparse_coordinator is not None: + h = self.hisparse_coordinator.get_token_stats() + return dataclasses.replace( + pool_stats, + is_hisparse=True, + hisparse_device_tokens=h.device_tokens, + hisparse_device_token_usage=h.device_token_usage, + hisparse_host_tokens=h.host_tokens, + hisparse_host_token_usage=h.host_token_usage, + ) + return pool_stats + def _get_mamba_token_info(self: Scheduler): is_mamba_radix_cache = ( self.tree_cache.supports_mamba() and self.tree_cache.is_tree_cache() From cc449ac4e59658f92f8ef5cc641de2bc83002d60 Mon Sep 17 00:00:00 2001 From: ishandhanani <82981111+ishandhanani@users.noreply.github.com> Date: Mon, 13 Apr 2026 20:30:36 -0500 Subject: [PATCH 49/52] feat(metrics): expose raw KV cache pool token counts as prometheus gauges (#22726) --- .../scheduler_runtime_checker_mixin.py | 3 +++ .../srt/observability/metrics_collector.py | 25 +++++++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py index ef8e1f3fda2c..d35efdb5356a 100644 --- a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py +++ b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py @@ -128,6 +128,9 @@ def update_scheduler_stats(self, stats: SchedulerStats) -> None: stats.swa_token_usage = self.swa_token_usage if self.is_hybrid_ssm: stats.mamba_usage = self.mamba_usage + stats.kv_available_tokens = self.full_available_size + stats.kv_evictable_tokens = self.full_evictable_size + stats.kv_used_tokens = self.full_num_used class SchedulerRuntimeCheckerMixin: diff --git a/python/sglang/srt/observability/metrics_collector.py b/python/sglang/srt/observability/metrics_collector.py index 3c34302161dd..f02ac70eee74 100644 --- a/python/sglang/srt/observability/metrics_collector.py +++ b/python/sglang/srt/observability/metrics_collector.py @@ -94,6 +94,9 @@ class SchedulerStats: cache_hit_rate: float = 0.0 max_total_num_tokens: int = 0 + kv_available_tokens: int = 0 + kv_evictable_tokens: int = 0 + kv_used_tokens: int = 0 # Speculative decoding spec_accept_length: float = 0.0 @@ -279,6 +282,25 @@ def __init__( multiprocess_mode="mostrecent", ) + self.kv_available_tokens = Gauge( + name="sglang:kv_available_tokens", + documentation="Number of free token slots in the KV cache pool.", + labelnames=labels.keys(), + multiprocess_mode="mostrecent", + ) + self.kv_evictable_tokens = Gauge( + name="sglang:kv_evictable_tokens", + documentation="Number of evictable (radix-cached) token slots in the KV cache pool.", + labelnames=labels.keys(), + multiprocess_mode="mostrecent", + ) + self.kv_used_tokens = Gauge( + name="sglang:kv_used_tokens", + documentation="Number of actively used token slots in the KV cache pool.", + labelnames=labels.keys(), + multiprocess_mode="mostrecent", + ) + # Speculative decoding self.spec_accept_length = Gauge( name="sglang:spec_accept_length", @@ -1014,6 +1036,9 @@ def log_stats(self, stats: SchedulerStats) -> None: self._log_gauge(self.cache_hit_rate, stats.cache_hit_rate) self._log_gauge(self.max_total_num_tokens, stats.max_total_num_tokens) + self._log_gauge(self.kv_available_tokens, stats.kv_available_tokens) + self._log_gauge(self.kv_evictable_tokens, stats.kv_evictable_tokens) + self._log_gauge(self.kv_used_tokens, stats.kv_used_tokens) # Speculative decoding self._log_gauge(self.spec_accept_length, stats.spec_accept_length) From e9d6b9eb2d676f245863609ea611bee78ab5ce70 Mon Sep 17 00:00:00 2001 From: huangtingwei <141888744+huangtingwei9988@users.noreply.github.com> Date: Tue, 14 Apr 2026 09:47:36 +0800 Subject: [PATCH 50/52] [HiCache & HybridModel] mooncake backend support DSA & mamba model (#21259) Signed-off-by: Vladislav Nosivskoy Co-authored-by: hzh0425 Co-authored-by: pansicheng Co-authored-by: ispobock Co-authored-by: Vladislav Nosivskoy --- .../srt/mem_cache/hi_mamba_radix_cache.py | 97 +----- .../sglang/srt/mem_cache/hicache_storage.py | 3 +- python/sglang/srt/mem_cache/hiradix_cache.py | 115 +++++-- .../hybrid_cache/hybrid_cache_controller.py | 27 +- .../hybrid_cache/hybrid_pool_assembler.py | 212 +++++++++++++ .../sglang/srt/mem_cache/memory_pool_host.py | 298 ++++++++++++++---- .../storage/mooncake_store/mooncake_store.py | 216 ++++++++++--- .../unit/mem_cache/test_nsa_pool_host_unit.py | 24 +- 8 files changed, 760 insertions(+), 232 deletions(-) create mode 100644 python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py index 10e63c7f8087..0b9c8f680cb6 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -24,9 +24,11 @@ ) from sglang.srt.mem_cache.hicache_storage import PoolHitPolicy, PoolName, PoolTransfer from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( - HybridCacheController, PrefetchOperation, ) +from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import ( + build_mamba_hybrid_stack, +) from sglang.srt.mem_cache.mamba_radix_cache import ( LRUList, MambaRadixCache, @@ -34,13 +36,6 @@ get_last_access_time, ) from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, HybridReqToTokenPool -from sglang.srt.mem_cache.memory_pool_host import ( - HostPoolGroup, - MambaPoolHost, - MHATokenToKVPoolHost, - MLATokenToKVPoolHost, - PoolEntry, -) from sglang.srt.mem_cache.radix_cache import ( RadixKey, compute_node_hash_values, @@ -115,66 +110,6 @@ def __init__(self, params: CacheInitParams, server_args: ServerArgs): ) self.kvcache = self.hybrid_kv_cache.full_kv_pool - kv_host_pool_cls = ( - MLATokenToKVPoolHost - if self.hybrid_kv_cache.use_mla - else MHATokenToKVPoolHost - ) - self.full_kv_pool_host = kv_host_pool_cls( - self.kvcache, - server_args.hicache_ratio, - server_args.hicache_size, - params.page_size, - server_args.hicache_mem_layout, - allocator_type=server_args.hicache_storage_backend, - ) - self.mamba_pool_host = MambaPoolHost( - params.req_to_token_pool.mamba_pool, - server_args.hicache_ratio, - server_args.hicache_size, - allocator_type=server_args.hicache_storage_backend, - layout=server_args.hicache_mem_layout, - ) - - full_layer_ids = sorted( - self.hybrid_kv_cache.full_attention_layer_id_mapping.keys() - ) - mamba_layer_ids = sorted(params.req_to_token_pool.mamba_map.keys()) - self.transfer_layer_num = len(set(full_layer_ids) | set(mamba_layer_ids)) - - full_layer_mapping = dict(self.hybrid_kv_cache.full_attention_layer_id_mapping) - mamba_layer_mapping = dict(params.req_to_token_pool.mamba_map) - transfer_layer_num = self.transfer_layer_num - - def kv_layer_mapper(layer_id: int) -> Optional[int]: - if not 0 <= layer_id < transfer_layer_num: - return None - return full_layer_mapping.get(layer_id) - - def mamba_layer_mapper(layer_id: int) -> Optional[int]: - if not 0 <= layer_id < transfer_layer_num: - return None - return mamba_layer_mapping.get(layer_id) - - self.host_pool_group = HostPoolGroup( - [ - PoolEntry( - name=PoolName.KV, - host_pool=self.full_kv_pool_host, - device_pool=self.kvcache, - layer_mapper=kv_layer_mapper, - is_primary_index_anchor=True, - ), - PoolEntry( - name=PoolName.MAMBA, - host_pool=self.mamba_pool_host, - device_pool=params.req_to_token_pool.mamba_pool, - layer_mapper=mamba_layer_mapper, - host_evict_fn=self.evict_mamba_host, - device_evict_fn=self.evict_mamba, - ), - ] - ) self.tp_group = params.tp_cache_group self.tp_world_size = ( @@ -200,27 +135,13 @@ def mamba_layer_mapper(layer_id: int) -> Optional[int]: self.prefetch_stop_policy = server_args.hicache_storage_prefetch_policy self.load_cache_event = threading.Event() - self.cache_controller = HybridCacheController( - params.token_to_kv_pool_allocator, - self.host_pool_group, - params.page_size, - self.tp_group, - load_cache_event=self.load_cache_event, - write_policy=server_args.hicache_write_policy, - io_backend=server_args.hicache_io_backend, - storage_backend=server_args.hicache_storage_backend, + build_mamba_hybrid_stack( + self, + params, + server_args, + extra_config=extra_config, prefetch_threshold=prefetch_threshold, - model_name=server_args.served_model_name, - storage_backend_extra_config=extra_config, - pp_rank=params.pp_rank, - pp_size=params.pp_size, - transfer_layer_num=self.transfer_layer_num, - ) - params.req_to_token_pool.register_layer_transfer_counter( - self.cache_controller.layer_done_counter - ) - self.hybrid_kv_cache.register_layer_transfer_counter( - self.cache_controller.layer_done_counter + load_cache_event=self.load_cache_event, ) self._apply_storage_runtime_config( storage_backend=server_args.hicache_storage_backend, diff --git a/python/sglang/srt/mem_cache/hicache_storage.py b/python/sglang/srt/mem_cache/hicache_storage.py index b2195c4c39ab..fa0a5769b6b8 100644 --- a/python/sglang/srt/mem_cache/hicache_storage.py +++ b/python/sglang/srt/mem_cache/hicache_storage.py @@ -41,12 +41,13 @@ class PoolName(str, Enum): KV = "kv" MAMBA = "mamba" + INDEXER = "indexer" class PoolHitPolicy(str, Enum): """Hit policy for batch_exists_v2 per-pool prefix matching. - ALL_PAGES : every page in [0, kv_hit) must exist (default). + ALL_PAGES : every page in [0, kv_hit) must exist (e.g. DSA). TRAILING_PAGES : only the last N pages must exist (e.g. Mamba/SWA states). """ diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index b4a5eb5b8eef..f9ecc0e6909b 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -25,6 +25,17 @@ MatchPrefixParams, MatchResult, ) +from sglang.srt.mem_cache.hicache_storage import ( + PoolHitPolicy, + PoolName, + PoolTransfer, +) +from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( + HybridCacheController, +) +from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import ( + build_nsa_hybrid_stack, +) from sglang.srt.mem_cache.memory_pool import ( MHATokenToKVPool, MLATokenToKVPool, @@ -33,7 +44,6 @@ from sglang.srt.mem_cache.memory_pool_host import ( MHATokenToKVPoolHost, MLATokenToKVPoolHost, - NSATokenToKVPoolHost, ) from sglang.srt.mem_cache.radix_cache import ( RadixCache, @@ -70,14 +80,8 @@ def __init__(self, params: CacheInitParams, server_args: ServerArgs): allocator_type=server_args.hicache_storage_backend, ) elif isinstance(self.kv_cache, NSATokenToKVPool): - self.token_to_kv_pool_host = NSATokenToKVPoolHost( - self.kv_cache, - server_args.hicache_ratio, - server_args.hicache_size, - self.page_size, - server_args.hicache_mem_layout, - allocator_type=server_args.hicache_storage_backend, - ) + # Filled by build_nsa_hybrid_stack after storage extra_config is parsed. + self.token_to_kv_pool_host = None elif isinstance(self.kv_cache, MLATokenToKVPool): self.token_to_kv_pool_host = MLATokenToKVPoolHost( self.kv_cache, @@ -88,7 +92,9 @@ def __init__(self, params: CacheInitParams, server_args: ServerArgs): allocator_type=server_args.hicache_storage_backend, ) else: - raise ValueError(f"HiRadixCache only supports MHA and MLA yet") + raise ValueError( + "HiRadixCache only supports MHA, MLA, and NSA (DSA) models" + ) self.tp_group = params.tp_cache_group self.tp_world_size = torch.distributed.get_world_size(group=self.tp_group) @@ -114,24 +120,35 @@ def __init__(self, params: CacheInitParams, server_args: ServerArgs): self.prefetch_stop_policy = server_args.hicache_storage_prefetch_policy self.load_cache_event = threading.Event() - self.cache_controller = HiCacheController( - params.token_to_kv_pool_allocator, - self.token_to_kv_pool_host, - self.page_size, - self.tp_group, - load_cache_event=self.load_cache_event, - write_policy=server_args.hicache_write_policy, - io_backend=server_args.hicache_io_backend, - storage_backend=server_args.hicache_storage_backend, - prefetch_threshold=prefetch_threshold, - model_name=server_args.served_model_name, - storage_backend_extra_config=extra_config, - pp_rank=self.pp_rank, - pp_size=self.pp_size, - attn_cp_rank=self.attn_cp_rank, - attn_cp_size=self.attn_cp_size, - enable_storage_metrics=self.enable_storage_metrics, - ) + if isinstance(self.kv_cache, NSATokenToKVPool): + build_nsa_hybrid_stack( + self, + params, + server_args, + extra_config=extra_config, + prefetch_threshold=prefetch_threshold, + enable_storage_metrics=self.enable_storage_metrics, + load_cache_event=self.load_cache_event, + ) + else: + self.cache_controller = HiCacheController( + params.token_to_kv_pool_allocator, + self.token_to_kv_pool_host, + self.page_size, + self.tp_group, + load_cache_event=self.load_cache_event, + write_policy=server_args.hicache_write_policy, + io_backend=server_args.hicache_io_backend, + storage_backend=server_args.hicache_storage_backend, + prefetch_threshold=prefetch_threshold, + model_name=server_args.served_model_name, + storage_backend_extra_config=extra_config, + pp_rank=self.pp_rank, + pp_size=self.pp_size, + attn_cp_rank=self.attn_cp_rank, + attn_cp_size=self.attn_cp_size, + enable_storage_metrics=self.enable_storage_metrics, + ) self._apply_storage_runtime_config( storage_backend=server_args.hicache_storage_backend, prefetch_threshold=prefetch_threshold, @@ -326,6 +343,7 @@ def attach_storage_backend( prefetch_threshold=prefetch_threshold, model_name=served_model_name, storage_backend_extra_config=extra_config, + **self._get_hybrid_storage_attach_kwargs(), ) except Exception as e: logger.exception( @@ -591,6 +609,24 @@ def get_height(self, node: TreeNode): height += 1 return height + def _get_extra_pools(self) -> dict: + if not isinstance(self.cache_controller, HybridCacheController): + return {} + if isinstance(self.kv_cache, NSATokenToKVPool): + pool = PoolTransfer( + name=PoolName.INDEXER, + hit_policy=PoolHitPolicy.ALL_PAGES, + ) + return {"extra_pools": [pool]} + else: + return {} + + def _get_hybrid_storage_attach_kwargs(self) -> dict: + """Extra kwargs for attach_storage_backend when controller is HybridCacheController.""" + if isinstance(self.cache_controller, HybridCacheController): + return {"host_pools": self.cache_controller.mem_pool_host.entries} + return {} + def clear_storage_backend(self) -> bool: if self.enable_storage: try: @@ -625,12 +661,14 @@ def write_backup(self, node: TreeNode, write_back=False) -> int: host_indices = self.cache_controller.write( device_indices=node.value, node_id=node.id, + **self._get_extra_pools(), ) if host_indices is None: self.evict_host(len(node.value)) host_indices = self.cache_controller.write( device_indices=node.value, node_id=node.id, + **self._get_extra_pools(), ) if host_indices is not None: node.host_value = host_indices.clone() @@ -652,7 +690,11 @@ def write_backup_storage(self, node: TreeNode): ) operation_id = self.cache_controller.write_storage( - node.host_value, node.key, node.hash_value, prefix_keys + node.host_value, + node.key, + node.hash_value, + prefix_keys, + **self._get_extra_pools(), ) self.ongoing_backup[operation_id] = node node.protect_host() @@ -925,12 +967,16 @@ def load_back( return None device_indices = self.cache_controller.load( - host_indices=host_indices, node_id=last_hit_node.id + host_indices=host_indices, + node_id=last_hit_node.id, + **self._get_extra_pools(), ) if device_indices is None: self.evict(EvictParams(num_tokens=len(host_indices))) device_indices = self.cache_controller.load( - host_indices=host_indices, node_id=last_hit_node.id + host_indices=host_indices, + node_id=last_hit_node.id, + **self._get_extra_pools(), ) self.dec_lock_ref(ancester_node) if device_indices is None: @@ -1231,7 +1277,12 @@ def prefetch_from_storage( # no sufficient host memory for prefetch return operation = self.cache_controller.prefetch( - req_id, host_indices, new_input_tokens, last_hash, prefix_keys + req_id, + host_indices, + new_input_tokens, + last_hash, + prefix_keys, + **self._get_extra_pools(), ) self.ongoing_prefetch[req_id] = ( last_host_node, diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py index 7d6abdab41ef..7fd9580b26fd 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py @@ -163,6 +163,7 @@ def __init__( pp_rank: int = 0, pp_size: int = 1, transfer_layer_num: Optional[int] = None, + enable_storage_metrics: bool = False, ): startup_storage_backend = storage_backend super().__init__( @@ -179,6 +180,7 @@ def __init__( storage_backend_extra_config=storage_backend_extra_config, pp_rank=pp_rank, pp_size=pp_size, + enable_storage_metrics=enable_storage_metrics, ) # Override layer_num: hybrid models transfer all layers (For example, Linear Model (KV + Mamba)), # not just the full attention layers reported by full_kv_pool. @@ -230,7 +232,10 @@ def write( if host_indices is None: return None pool_transfers = self._resolve_pool_transfers_allocation( - extra_pools, alloc_host=True + extra_pools, + alloc_host=True, + kv_device_indices=device_indices, + kv_host_indices=host_indices, ) if pool_transfers is None and extra_pools: self.mem_pool_host.free(host_indices) @@ -289,7 +294,10 @@ def load( device_indices = torch.empty((0,), dtype=torch.int64, device=self.device) pool_transfers = self._resolve_pool_transfers_allocation( - extra_pools, alloc_host=False + extra_pools, + alloc_host=False, + kv_device_indices=device_indices, + kv_host_indices=host_indices, ) if pool_transfers is None and extra_pools: if need_load_kv: @@ -415,6 +423,7 @@ def _storage_hit_query(self, operation) -> tuple[list[str], int]: def _page_transfer(self, operation): # Transfer extra pools if operation.pool_transfers and not operation.is_terminated(): + self._resolve_shared_pool_transfers(operation) results = self.storage_backend.batch_get_v2(operation.pool_transfers) operation.pool_storage_result.update_extra_pool_hit_pages(results) @@ -424,12 +433,20 @@ def _page_transfer(self, operation): def _page_backup(self, operation): # Backup extra pools if operation.pool_transfers: + self._resolve_shared_pool_transfers(operation) results = self.storage_backend.batch_set_v2(operation.pool_transfers) operation.pool_storage_result.update_extra_pool_hit_pages(results) # Backup kv pools super()._page_backup(operation) + def _resolve_shared_pool_transfers(self, operation): + for transfer in operation.pool_transfers: + entry = self.mem_pool_host.entry_map.get(transfer.name) + if entry.share_indices_with_anchor: + transfer.keys = operation.hash_value + transfer.host_indices = operation.host_indices + def _sync_trailing_keys( self, pool_transfers: list[PoolTransfer], @@ -454,6 +471,8 @@ def _resolve_pool_transfers_allocation( self, extra_pools: Optional[list[PoolTransfer]], alloc_host: bool, + kv_device_indices: Optional[torch.Tensor] = None, + kv_host_indices: Optional[torch.Tensor] = None, ) -> Optional[list[PoolTransfer]]: """Auto-alloc host or device indices for PoolTransfers where they are None.""" if not extra_pools: @@ -463,6 +482,10 @@ def _resolve_pool_transfers_allocation( entry = self.mem_pool_host.entry_map.get(pool.name) if entry is None: continue + if entry.share_indices_with_anchor: + pool.device_indices = kv_device_indices + pool.host_indices = kv_host_indices + continue if alloc_host: if pool.host_indices is not None or pool.device_indices is None: continue diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py new file mode 100644 index 000000000000..99198dfece94 --- /dev/null +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py @@ -0,0 +1,212 @@ +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Optional + +from sglang.srt.mem_cache.hicache_storage import PoolName +from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( + HybridCacheController, +) +from sglang.srt.mem_cache.memory_pool_host import ( + HostPoolGroup, + MambaPoolHost, + MHATokenToKVPoolHost, + MLATokenToKVPoolHost, + NSAIndexerPoolHost, + PoolEntry, +) + +if TYPE_CHECKING: + from sglang.srt.mem_cache.cache_init_params import CacheInitParams + from sglang.srt.mem_cache.hi_mamba_radix_cache import HiMambaRadixCache + from sglang.srt.mem_cache.hiradix_cache import HiRadixCache + from sglang.srt.server_args import ServerArgs + +logger = logging.getLogger(__name__) + + +def build_nsa_hybrid_stack( + radix_cache: "HiRadixCache", + params: "CacheInitParams", + server_args: "ServerArgs", + *, + extra_config: dict, + prefetch_threshold: int, + enable_storage_metrics: bool, + load_cache_event, +) -> None: + """HostPoolGroup (KV + indexer) + HybridCacheController for NSA (DSA).""" + try: + kv = radix_cache.kv_cache + mla_host = MLATokenToKVPoolHost( + kv, + server_args.hicache_ratio, + server_args.hicache_size, + radix_cache.page_size, + server_args.hicache_mem_layout, + allocator_type=server_args.hicache_storage_backend, + override_kv_cache_dim=kv.kv_cache_dim, + ) + indexer_host = NSAIndexerPoolHost( + kv, + mla_host, + server_args.hicache_mem_layout, + allocator_type=server_args.hicache_storage_backend, + ) + layer_num = kv.layer_num + + def layer_mapper(layer_id: int): + if 0 <= layer_id < layer_num: + return layer_id + return None + + host_pool_group = HostPoolGroup( + [ + PoolEntry( + name=PoolName.KV, + host_pool=mla_host, + device_pool=kv, + layer_mapper=layer_mapper, + is_primary_index_anchor=True, + ), + PoolEntry( + name=PoolName.INDEXER, + host_pool=indexer_host, + device_pool=kv, + layer_mapper=layer_mapper, + share_indices_with_anchor=True, + ), + ] + ) + cache_controller = HybridCacheController( + params.token_to_kv_pool_allocator, + host_pool_group, + radix_cache.page_size, + radix_cache.tp_group, + load_cache_event=load_cache_event, + write_policy=server_args.hicache_write_policy, + io_backend=server_args.hicache_io_backend, + storage_backend=server_args.hicache_storage_backend, + prefetch_threshold=prefetch_threshold, + model_name=server_args.served_model_name, + storage_backend_extra_config=extra_config, + pp_rank=radix_cache.pp_rank, + pp_size=radix_cache.pp_size, + transfer_layer_num=layer_num, + enable_storage_metrics=enable_storage_metrics, + ) + radix_cache.full_kv_pool_host = mla_host + radix_cache.token_to_kv_pool_host = host_pool_group + radix_cache.cache_controller = cache_controller + logger.info( + "Hybrid hierarchical cache: HostPoolGroup(KV + INDEXER), HybridCacheController, " + "transfer_layer_num=%s", + layer_num, + ) + except Exception: + logger.exception("build_nsa_hybrid_stack failed") + raise + + +def build_mamba_hybrid_stack( + mamba_cache: "HiMambaRadixCache", + params: "CacheInitParams", + server_args: "ServerArgs", + *, + extra_config: dict, + prefetch_threshold: int, + load_cache_event, + enable_storage_metrics: bool = False, +) -> None: + """HostPoolGroup (KV + Mamba) + HybridCacheController for hybrid SSM models.""" + try: + hybrid_kv = mamba_cache.hybrid_kv_cache + kvcache = mamba_cache.kvcache + kv_host_pool_cls = ( + MLATokenToKVPoolHost if hybrid_kv.use_mla else MHATokenToKVPoolHost + ) + full_kv_pool_host = kv_host_pool_cls( + kvcache, + server_args.hicache_ratio, + server_args.hicache_size, + params.page_size, + server_args.hicache_mem_layout, + allocator_type=server_args.hicache_storage_backend, + ) + mamba_pool_host = MambaPoolHost( + params.req_to_token_pool.mamba_pool, + server_args.hicache_ratio, + server_args.hicache_size, + allocator_type=server_args.hicache_storage_backend, + layout=server_args.hicache_mem_layout, + ) + + full_layer_ids = sorted(hybrid_kv.full_attention_layer_id_mapping.keys()) + mamba_layer_ids = sorted(params.req_to_token_pool.mamba_map.keys()) + transfer_layer_num = len(set(full_layer_ids) | set(mamba_layer_ids)) + full_layer_mapping = dict(hybrid_kv.full_attention_layer_id_mapping) + mamba_layer_mapping = dict(params.req_to_token_pool.mamba_map) + + def kv_layer_mapper(layer_id: int) -> Optional[int]: + if not 0 <= layer_id < transfer_layer_num: + return None + return full_layer_mapping.get(layer_id) + + def mamba_layer_mapper(layer_id: int) -> Optional[int]: + if not 0 <= layer_id < transfer_layer_num: + return None + return mamba_layer_mapping.get(layer_id) + + host_pool_group = HostPoolGroup( + [ + PoolEntry( + name=PoolName.KV, + host_pool=full_kv_pool_host, + device_pool=kvcache, + layer_mapper=kv_layer_mapper, + is_primary_index_anchor=True, + ), + PoolEntry( + name=PoolName.MAMBA, + host_pool=mamba_pool_host, + device_pool=params.req_to_token_pool.mamba_pool, + layer_mapper=mamba_layer_mapper, + host_evict_fn=mamba_cache.evict_mamba_host, + device_evict_fn=mamba_cache.evict_mamba, + ), + ] + ) + cache_controller = HybridCacheController( + params.token_to_kv_pool_allocator, + host_pool_group, + params.page_size, + params.tp_cache_group, + load_cache_event=load_cache_event, + write_policy=server_args.hicache_write_policy, + io_backend=server_args.hicache_io_backend, + storage_backend=server_args.hicache_storage_backend, + prefetch_threshold=prefetch_threshold, + model_name=server_args.served_model_name, + storage_backend_extra_config=extra_config, + pp_rank=params.pp_rank, + pp_size=params.pp_size, + transfer_layer_num=transfer_layer_num, + enable_storage_metrics=enable_storage_metrics, + ) + mamba_cache.full_kv_pool_host = full_kv_pool_host + mamba_cache.mamba_pool_host = mamba_pool_host + mamba_cache.transfer_layer_num = transfer_layer_num + mamba_cache.host_pool_group = host_pool_group + mamba_cache.cache_controller = cache_controller + params.req_to_token_pool.register_layer_transfer_counter( + cache_controller.layer_done_counter + ) + hybrid_kv.register_layer_transfer_counter(cache_controller.layer_done_counter) + logger.info( + "Hybrid hierarchical cache: HostPoolGroup(KV + MAMBA), HybridCacheController, " + "transfer_layer_num=%s", + transfer_layer_num, + ) + except Exception: + logger.exception("build_mamba_hybrid_stack failed") + raise diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 8410f8471aeb..e5fd2da76f97 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -64,6 +64,9 @@ logger = logging.getLogger(__name__) +# Host RAM to leave free when sizing HiCache pools (OS, other processes). +HICACHE_HOST_MEMORY_RESERVE_BYTES: int = 10 * (1024**3) + def synchronized(func): @wraps(func) @@ -187,9 +190,7 @@ def __init__( # Verify there is enough available host memory. host_mem = psutil.virtual_memory() requested_bytes = self.size * self.size_per_token - # preserve at least 10GB for other usage - ten_gb = 10 * (1024**3) - available_bytes = host_mem.available - ten_gb + available_bytes = host_mem.available - HICACHE_HOST_MEMORY_RESERVE_BYTES if requested_bytes > available_bytes: raise ValueError( f"Not enough host memory available. Requesting " @@ -1213,6 +1214,10 @@ def __init__( conv_state.shape[2:] for conv_state in device_pool.mamba_cache.conv ] self.temporal_state_shape = device_pool.mamba_cache.temporal.shape[2:] + self.temporal_state_elem_size = int(np.prod(self.temporal_state_shape)) + self.conv_state_elem_sizes = [ + int(np.prod(conv_shape)) for conv_shape in self.conv_state_shapes + ] self.conv_dtype = device_pool.mamba_cache.conv[0].dtype self.temporal_dtype = device_pool.mamba_cache.temporal.dtype self.dtype = self.conv_dtype @@ -1232,8 +1237,7 @@ def __init__( host_mem = psutil.virtual_memory() requested_bytes = self.size * self.size_per_token - ten_gb = 10 * (1024**3) - available_bytes = host_mem.available - ten_gb + available_bytes = host_mem.available - HICACHE_HOST_MEMORY_RESERVE_BYTES if requested_bytes > available_bytes: raise ValueError( f"Not enough host memory available. Requesting " @@ -1306,6 +1310,10 @@ def init_kv_buffer(self): ) ) + def get_hybrid_pool_buffer(self): + # Expose all mamba host tensors that need Mooncake buffer registration. + return [self.temporal_buffer, *self.conv_buffer] + def _iter_page_tensors(self, index: int): if self.layout in ["page_first", "page_first_direct"]: yield self.temporal_buffer[index] @@ -1347,12 +1355,11 @@ def free(self, indices: torch.Tensor) -> int: return len(indices) def get_size_per_token(self): - conv_total_size = 0 - for conv_shape in self.conv_state_shapes: - conv_total_size += int(np.prod(conv_shape)) * self.conv_dtype.itemsize - temporal_size = ( - int(np.prod(self.temporal_state_shape)) * self.temporal_dtype.itemsize + conv_total_size = sum( + conv_elem_size * self.conv_dtype.itemsize + for conv_elem_size in self.conv_state_elem_sizes ) + temporal_size = self.temporal_state_elem_size * self.temporal_dtype.itemsize return (conv_total_size + temporal_size) * self.num_mamba_layers def get_ksize_per_token(self): @@ -1588,6 +1595,65 @@ def set_from_flat_data_page( restored = tensor_bytes.view(dtype=tensor.dtype).reshape(tensor.shape) tensor.copy_(restored) + def get_page_buffer_meta(self, indices): + """Meta data for zero-copy storage I/O. + + Only page-first layouts are supported for mamba storage zero-copy because + each page slot in temporal/conv buffers is directly addressable. + """ + assert len(indices) % self.page_size == 0 + if self.layout not in ["page_first", "page_first_direct"]: + raise ValueError( + f"Mamba storage zero-copy requires page_first layout, got {self.layout}" + ) + indices = indices.tolist() + ptr_list = [] + element_size_list = [] + + # Compute base pointers once; each page pointer is offset from these bases. + temporal_base_ptr = self.temporal_buffer.data_ptr() + conv_base_ptrs = [buf.data_ptr() for buf in self.conv_buffer] + # Component sizes are constant across pages, so precompute once as well. + temporal_element_size = ( + self.page_size + * self.num_mamba_layers + * self.temporal_dtype.itemsize + * self.temporal_state_elem_size + ) + conv_element_sizes = [ + ( + self.page_size + * self.num_mamba_layers + * self.conv_dtype.itemsize + * self.conv_state_elem_sizes[i] + ) + for i in range(len(self.conv_state_shapes)) + ] + + for i in range(0, len(indices), self.page_size): + # Emit component pointers in stable order: + # temporal first, then conv_0..conv_n for this page. + temporal_ptr = ( + temporal_base_ptr + + indices[i] + * self.num_mamba_layers + * self.temporal_state_elem_size + * self.temporal_dtype.itemsize + ) + ptr_list.append(temporal_ptr) + element_size_list.append(temporal_element_size) + for j in range(len(self.conv_buffer)): + conv_ptr = ( + conv_base_ptrs[j] + + indices[i] + * self.num_mamba_layers + * self.conv_state_elem_sizes[j] + * self.conv_dtype.itemsize + ) + ptr_list.append(conv_ptr) + element_size_list.append(conv_element_sizes[j]) + return ptr_list, element_size_list + @dataclass class PoolEntry: @@ -1596,6 +1662,9 @@ class PoolEntry: device_pool: Any layer_mapper: Callable[[int], Optional[int]] is_primary_index_anchor: bool = False + # When True, host_pool uses the same logical slot indices as the anchor pool + # (e.g. DSA indexer); HostPoolGroup.free mirrors frees to this pool. + share_indices_with_anchor: bool = False # Optional eviction callbacks for auto-alloc in HybridCacheController. # host_evict_fn(n): evict n slots from the host pool (used by write()). # device_evict_fn(n): evict n slots from the device pool (used by load()). @@ -1619,15 +1688,54 @@ def __init__(self, entries: list[PoolEntry]): self.device = self.anchor_entry.host_pool.device self.size = self.anchor_entry.host_pool.size + @property + def kv_buffer(self): + return self.anchor_entry.host_pool.kv_buffer + + @property + def size_per_token(self): + return self.anchor_entry.host_pool.size_per_token + + @property + def allocator(self): + return self.anchor_entry.host_pool.allocator + + @property + def dtype(self): + return self.anchor_entry.host_pool.dtype + + @property + def start_layer(self): + return self.anchor_entry.host_pool.start_layer + + @property + def end_layer(self): + return self.anchor_entry.host_pool.end_layer + + def get_ksize_per_token(self): + return self.anchor_entry.host_pool.get_ksize_per_token() + + def get_page_buffer_meta(self, indices): + return self.anchor_entry.host_pool.get_page_buffer_meta(indices) + def clear(self) -> None: for entry in self.entries: entry.host_pool.clear() + def available_size(self): + return self.anchor_entry.host_pool.available_size() + def alloc(self, need_size: int) -> Optional[torch.Tensor]: return self.anchor_entry.host_pool.alloc(need_size) def free(self, indices: torch.Tensor) -> int: - return self.anchor_entry.host_pool.free(indices) + n = self.anchor_entry.host_pool.free(indices) + for entry in self.entries: + if entry is self.anchor_entry: + continue + if getattr(entry, "share_indices_with_anchor", False): + entry.host_pool.free(indices) + return n def get_data_page(self, index, flat: bool = True): return self.anchor_entry.host_pool.get_data_page(index, flat) @@ -1703,21 +1811,31 @@ def backup_from_device_all_layer( ) -class NSATokenToKVPoolHost(MLATokenToKVPoolHost): +class NSAIndexerPoolHost(HostKVCache): + """Host-side NSA index buffers only. Slot layout matches the anchor MLA host pool.""" + device_pool: NSATokenToKVPool def __init__( self, device_pool: NSATokenToKVPool, - host_to_device_ratio: float, - host_size: int, - page_size: int, + anchor_host: MLATokenToKVPoolHost, layout: str, pin_memory: bool = True, device: str = "cpu", allocator_type: str = "default", ): - # Initialize indexer metadata before HostKVCache.__init__ calls get_size_per_token. + self.device_pool = device_pool + self.page_size = anchor_host.page_size + self.layout = layout + self.pin_memory = pin_memory + self.device = device + self.allocator = get_allocator_from_storage(allocator_type) + self.dtype = device_pool.store_dtype + self.start_layer = device_pool.start_layer + self.end_layer = device_pool.end_layer + self.layer_num = device_pool.layer_num + self.index_head_dim = device_pool.index_head_dim self.indexer_quant_block_size = device_pool.quant_block_size self.indexer_dtype = NSATokenToKVPool.index_k_with_scale_buffer_dtype @@ -1725,35 +1843,46 @@ def __init__( self.index_head_dim + self.index_head_dim // self.indexer_quant_block_size * 4 ) - super().__init__( - device_pool, - host_to_device_ratio, - host_size, - page_size, - layout, - pin_memory, - device, - allocator_type, - override_kv_cache_dim=device_pool.kv_cache_dim, - ) + self.size = anchor_host.size + self.page_num = anchor_host.page_num + self.indexer_page_stride_size = ( self.indexer_size_per_token * self.page_size * self.indexer_dtype.itemsize ) self.indexer_layout_dim = self.indexer_page_stride_size * self.layer_num self.indexer_page_num = (self.size + self.page_size + 1) // self.page_size - self._init_indexer_buffers() + self.size_per_token = ( + self.indexer_size_per_token * self.layer_num * self.indexer_dtype.itemsize + ) + + buf_elem_size = self.page_num * self.layer_num * self.indexer_page_stride_size + requested_bytes = buf_elem_size * self.indexer_dtype.itemsize + host_mem = psutil.virtual_memory() + available_bytes = host_mem.available - HICACHE_HOST_MEMORY_RESERVE_BYTES + if requested_bytes > available_bytes: + raise ValueError( + f"Not enough host memory for NSA indexer hierarchical cache. " + f"Requesting {requested_bytes / 1e9:.2f} GB but only have " + f"{available_bytes / 1e9:.2f} GB free." + ) logger.info( - f"NSATokenToKVPoolHost initialized with indexer page stride size: {self.indexer_page_stride_size}, page num: {self.indexer_page_num}" + "Allocating %.2f GB host memory for NSA indexer (layout=%s).", + requested_bytes / 1e9, + layout, ) + self.init_kv_buffer() + self.lock = threading.RLock() + self.clear() def get_size_per_token(self): - base = super().get_size_per_token() return ( - base - + self.indexer_size_per_token * self.layer_num * self.indexer_dtype.itemsize + self.indexer_size_per_token * self.layer_num * self.indexer_dtype.itemsize ) - def _init_indexer_buffers(self): + def get_ksize_per_token(self): + return self.get_size_per_token() + + def init_kv_buffer(self): alloc_func = ALLOC_MEMORY_FUNCS[self.device_pool.device] self.index_k_device_ptrs = torch.tensor( [x.data_ptr() for x in self.device_pool.index_k_with_scale_buffer], @@ -1761,16 +1890,13 @@ def _init_indexer_buffers(self): device=self.device_pool.device, ) if self.layout == "layer_first": - self.index_k_with_scale_buffer = [ - alloc_func( - (self.indexer_page_num, self.indexer_page_stride_size), - dtype=self.indexer_dtype, - device=self.device, - pin_memory=self.pin_memory, - allocator=self.allocator, - ) - for _ in range(self.layer_num) - ] + self.index_k_with_scale_buffer = alloc_func( + (self.layer_num, self.indexer_page_num, self.indexer_page_stride_size), + dtype=self.indexer_dtype, + device=self.device, + pin_memory=self.pin_memory, + allocator=self.allocator, + ) self.index_k_data_refs = [ self.index_k_with_scale_buffer[i] for i in range(self.layer_num) ] @@ -1795,6 +1921,9 @@ def _init_indexer_buffers(self): else: raise ValueError(f"Unsupported layout: {self.layout}") + def get_hybrid_pool_buffer(self): + return [self.index_k_with_scale_buffer] + def _get_indexer_page_indices(self, host_indices, device_indices): if host_indices.numel() == 0: return host_indices, device_indices @@ -1810,7 +1939,7 @@ def _get_indexer_page_indices(self, host_indices, device_indices): ) return host_page_indices, device_page_indices - def _load_indexer_to_device_per_layer( + def load_to_device_per_layer( self, device_pool, host_indices, device_indices, layer_id, io_backend ): host_page_indices, device_page_indices = self._get_indexer_page_indices( @@ -1861,7 +1990,7 @@ def _load_indexer_to_device_per_layer( else: raise ValueError(f"Unsupported IO backend: {io_backend}") - def _backup_indexer_from_device_all_layer( + def backup_from_device_all_layer( self, device_pool, host_indices, device_indices, io_backend ): host_page_indices, device_page_indices = self._get_indexer_page_indices( @@ -1894,7 +2023,7 @@ def _backup_indexer_from_device_all_layer( if self.layout == "layer_first": transfer_kv_direct( src_layers=device_pool.index_k_with_scale_buffer, - dst_layers=self.index_k_with_scale_buffer, + dst_layers=self.index_k_data_refs, src_indices=device_page_indices, dst_indices=host_page_indices, page_size=1, @@ -1912,27 +2041,60 @@ def _backup_indexer_from_device_all_layer( else: raise ValueError(f"Unsupported IO backend: {io_backend}") - def load_to_device_per_layer( - self, - device_pool, - host_indices, - device_indices, - layer_id, - io_backend, - ): - super().load_to_device_per_layer( - device_pool, host_indices, device_indices, layer_id, io_backend - ) - self._load_indexer_to_device_per_layer( - device_pool, host_indices, device_indices, layer_id, io_backend - ) + def get_data_page(self, index, flat: bool = True) -> torch.Tensor: + page_idx = int(index) // self.page_size + if self.layout == "layer_first": + data_page = self.index_k_with_scale_buffer[:, page_idx : page_idx + 1, :] + elif self.layout in ["page_first", "page_first_direct"]: + data_page = self.index_k_with_scale_buffer[page_idx : page_idx + 1, :, :, :] + else: + raise ValueError(f"Unsupported layout: {self.layout}") + if flat: + data_page = data_page.flatten() + return data_page - def backup_from_device_all_layer( - self, device_pool, host_indices, device_indices, io_backend - ): - super().backup_from_device_all_layer( - device_pool, host_indices, device_indices, io_backend - ) - self._backup_indexer_from_device_all_layer( - device_pool, host_indices, device_indices, io_backend + def get_dummy_flat_data_page(self) -> torch.Tensor: + return torch.zeros( + (self.layer_num, self.indexer_page_stride_size), + dtype=self.indexer_dtype, + device=self.device, + pin_memory=self.pin_memory, + ).flatten() + + def set_from_flat_data_page(self, index: int, data_page: torch.Tensor) -> None: + page_idx = int(index) // self.page_size + if self.layout == "layer_first": + self.index_k_with_scale_buffer[:, page_idx : page_idx + 1, :] = ( + data_page.reshape( + self.layer_num, + 1, + self.indexer_page_stride_size, + ) + ) + elif self.layout in ["page_first", "page_first_direct"]: + self.index_k_with_scale_buffer[page_idx : page_idx + 1, :, :, :] = ( + data_page.reshape( + 1, + self.layer_num, + 1, + self.indexer_page_stride_size, + ) + ) + else: + raise ValueError(f"Unsupported layout: {self.layout}") + + def get_page_buffer_meta(self, indices): + """Meta data for zero-copy storage I/O.""" + assert len(indices) % self.page_size == 0 + if self.layout not in ["page_first", "page_first_direct"]: + raise ValueError(f"Unsupported layout: {self.layout}") + ptr_list = [] + indices = indices.tolist() + page_stride_bytes = ( + self.layer_num * self.indexer_page_stride_size * self.indexer_dtype.itemsize ) + base_ptr = self.index_k_with_scale_buffer.data_ptr() + for i in range(0, len(indices), self.page_size): + page_index = int(indices[i]) // self.page_size + ptr_list.append(base_ptr + page_index * page_stride_bytes) + return ptr_list, [page_stride_bytes] * len(ptr_list) diff --git a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py index 1923d5baf789..a8e4b08878ea 100644 --- a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py +++ b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py @@ -5,7 +5,7 @@ import time import uuid from dataclasses import dataclass -from typing import Any, List, Optional +from typing import Any, List, Optional, Tuple import requests import torch @@ -15,6 +15,10 @@ HiCacheStorage, HiCacheStorageConfig, HiCacheStorageExtraInfo, + PoolHitPolicy, + PoolName, + PoolTransfer, + PoolTransferResult, ) from sglang.srt.mem_cache.memory_pool_host import HostKVCache, HostTensorAllocator from sglang.srt.observability.metrics_collector import StorageMetrics @@ -432,6 +436,8 @@ def __init__( else: self.mha_suffix = [f"{rank}" for rank in target_ranks] + self.registered_pools = {} + self.gb_per_page = None self.prefetch_pgs = [] self.backup_pgs = [] @@ -502,6 +508,154 @@ def register_mem_pool_host(self, mem_pool_host: HostKVCache): bytes_per_page = mem_pool_host.get_ksize_per_token() * mem_pool_host.page_size self.gb_per_page = bytes_per_page / (1 << 30) + def register_mem_host_pool_v2(self, host_pool: HostKVCache, host_pool_name): + # KV anchor memory is already registered via register_mem_pool_host(). + # v2 here only registers additional hybrid pools. + if host_pool_name == PoolName.KV: + return + # Keep a name->pool mapping so batch v2 can resolve PoolTransfer.name to + # the corresponding host pool implementation at runtime. + self.registered_pools[host_pool_name] = host_pool + + # Hybrid pools expose the tensors that Mooncake needs for zero-copy I/O. + # The storage backend only depends on this accessor, not concrete fields. + buf_list = host_pool.get_hybrid_pool_buffer() + for buf in buf_list: + super().register_buffer(buf) + + def _tag_keys(self, keys: List[str]) -> List[str]: + if self.extra_backend_tag is None: + return keys + return [f"{ self.extra_backend_tag}_{key}" for key in keys] + + def _get_hybrid_page_component_keys( + self, page_keys: List[str], transfer: PoolTransfer + ) -> Tuple[List[str], int]: + # A logical "page" may map to multiple physical objects in storage. + # - INDEXER: one key per page + # - MAMBA : one temporal key + N conv keys per page + # key_multiplier records how many component keys are generated per page. + name = transfer.name + suffixes = [] + if name == PoolName.INDEXER: + suffixes = [f"_{self.mla_suffix}_{PoolName.INDEXER}"] + elif name == PoolName.MAMBA: + pools = getattr(self, "registered_pools", {}) + mamba_pool = pools.get(PoolName.MAMBA) + conv_num = len(getattr(mamba_pool, "conv_buffer", None) or []) + base_suffix = f"_{self.mha_suffix}" + suffixes = [f"{base_suffix}_temporal"] + [ + f"{base_suffix}_conv_{i}" for i in range(conv_num) + ] + key_multiplier = len(suffixes) + component_keys = [ + f"{page_key}{suffix}" for page_key in page_keys for suffix in suffixes + ] + return component_keys, key_multiplier + + def batch_exists_v2( + self, + keys: List[str], + pool_transfers: Optional[List[PoolTransfer]] = None, + extra_info: Optional[HiCacheStorageExtraInfo] = None, + ) -> PoolTransferResult: + qkeys = self._tag_keys(keys) + kv_pages = self.batch_exists(qkeys, extra_info) + + hit_count: dict = {PoolName.KV: kv_pages} if kv_pages else {} + final_pages = kv_pages + + for transfer in pool_transfers or []: + if final_pages == 0: + break + component_keys, key_multiplier = self._get_hybrid_page_component_keys( + qkeys, transfer + ) + ex = self._batch_exist(component_keys) + if key_multiplier > 0: + page_exists = [ + all( + r == 1 + for r in ex[i * key_multiplier : (i + 1) * key_multiplier] + ) + for i in range(kv_pages) + ] + else: + page_exists = [False] * kv_pages + boundary = 0 + if transfer.hit_policy == PoolHitPolicy.ALL_PAGES: + try: + boundary = page_exists.index(False) + except ValueError: + boundary = kv_pages + elif transfer.hit_policy == PoolHitPolicy.TRAILING_PAGES: + trailing = max(1, len(transfer.keys) if transfer.keys else 1) + for prefix_len in range(kv_pages, 0, -1): + if all( + page_exists[i] + for i in range(max(0, prefix_len - trailing), prefix_len) + ): + boundary = prefix_len + break + if boundary: + hit_count[transfer.name] = boundary + final_pages = min(final_pages, boundary) + + return PoolTransferResult(final_pages, hit_count) + + def _batch_io_v2(self, transfers: List[PoolTransfer], is_set: bool): + # Unified v2 I/O path: each PoolTransfer can expand to one or more + # storage objects per logical page, but API still reports page-level result. + results: dict = {} + for transfer in transfers: + host_pool = getattr(self, "registered_pools", {}).get(transfer.name) + keys = transfer.keys + page_size = getattr(host_pool, "page_size", 1) or 1 + host_indices = transfer.host_indices + assert len(keys) > 0 + assert len(keys) == len(host_indices) // page_size + + ptr_list, element_size_list = host_pool.get_page_buffer_meta(host_indices) + key_strs, key_multiplier = self._get_hybrid_page_component_keys( + keys, transfer + ) + key_strs = self._tag_keys(key_strs) + + if is_set: + exist_result = self._batch_exist(key_strs) + io_results = [0 if state == 1 else -1 for state in exist_result] + missing_idx = [i for i, state in enumerate(exist_result) if state != 1] + if missing_idx: + put_results = self._put_batch_zero_copy_impl( + [key_strs[i] for i in missing_idx], + [ptr_list[i] for i in missing_idx], + [element_size_list[i] for i in missing_idx], + ) + for i, res in zip(missing_idx, put_results): + io_results[i] = res + else: + io_results = self._get_batch_zero_copy_impl( + key_strs, ptr_list, element_size_list + ) + results[transfer.name] = self._batch_postprocess( + io_results, is_set_operate=is_set, key_multiplier=key_multiplier + ) + return results + + def batch_get_v2( + self, + transfers: List[PoolTransfer], + extra_info: Optional[HiCacheStorageExtraInfo] = None, + ) -> dict: + return self._batch_io_v2(transfers, is_set=False) + + def batch_set_v2( + self, + transfers: List[PoolTransfer], + extra_info: Optional[HiCacheStorageExtraInfo] = None, + ) -> dict: + return self._batch_io_v2(transfers, is_set=True) + def _get_mha_split_heads_buffer_meta(self, keys, indices): ptr_list, element_size_list = ( self.mem_pool_host.get_split_heads_page_buffer_meta( @@ -544,7 +698,9 @@ def _batch_preprocess(self, keys, host_indices): else: return self._get_mha_buffer_meta(keys, host_indices) - def _batch_postprocess(self, results: List[int], is_set_operate=False): + def _batch_postprocess( + self, results: List[int], is_set_operate=False, key_multiplier=None + ): """ refer to https://github.com/kvcache-ai/Mooncake/blob/main/mooncake-store/include/pybind_client.h for batch_get_into, results is Vector of integers, @@ -552,32 +708,26 @@ def _batch_postprocess(self, results: List[int], is_set_operate=False): for batch_put_from, results is Vector of integers, where each element is 0 on success, or a negative value on error """ - if self.is_mla_backend: - return [k_res == 0 if is_set_operate else k_res > 0 for k_res in results] - else: - if self.storage_config.should_split_heads: - kv_groups = [ - results[i : i + self.split_factor * 2] - for i in range(0, len(results), self.split_factor * 2) - ] - return [ - ( - all(res == 0 for res in kv_group) - if is_set_operate - else all(res > 0 for res in kv_group) - ) - for kv_group in kv_groups - ] + if key_multiplier is None: + if self.is_mla_backend: + key_multiplier = 1 else: - kv_pairs = zip(results[::2], results[1::2]) - return [ - ( - (k_res == 0 and v_res == 0) - if is_set_operate - else (k_res > 0 and v_res > 0) - ) - for k_res, v_res in kv_pairs - ] + key_multiplier = 2 + if self.storage_config.should_split_heads: + key_multiplier *= self.split_factor + + result_groups = [ + results[i : i + key_multiplier] + for i in range(0, len(results), key_multiplier) + ] + return [ + ( + all(res == 0 for res in group) + if is_set_operate + else all(res > 0 for res in group) + ) + for group in result_groups + ] def batch_get_v1( self, @@ -586,9 +736,7 @@ def batch_get_v1( extra_info: Optional[HiCacheStorageExtraInfo] = None, ) -> List[bool]: # Apply extra_backend_tag prefix if available - if self.extra_backend_tag is not None: - prefix = self.extra_backend_tag - keys = [f"{prefix}_{key}" for key in keys] + keys = self._tag_keys(keys) key_strs, buffer_ptrs, buffer_sizes = self._batch_preprocess(keys, host_indices) @@ -613,9 +761,7 @@ def batch_set_v1( extra_info: Optional[HiCacheStorageExtraInfo] = None, ) -> List[bool]: # Apply extra_backend_tag prefix if available - if self.extra_backend_tag is not None: - prefix = self.extra_backend_tag - keys = [f"{prefix}_{key}" for key in keys] + keys = self._tag_keys(keys) key_strs, buffer_ptrs, buffer_sizes = self._batch_preprocess(keys, host_indices) exist_result = self._batch_exist(key_strs) @@ -780,9 +926,7 @@ def batch_exists( self, keys, extra_info: Optional[HiCacheStorageExtraInfo] = None ) -> int: # Apply extra_backend_tag prefix if available - if self.extra_backend_tag is not None: - prefix = self.extra_backend_tag - keys = [f"{prefix}_{key}" for key in keys] + keys = self._tag_keys(keys) if self.is_mla_backend: query_keys = [f"{key}_{self.mla_suffix}_k" for key in keys] diff --git a/test/registered/unit/mem_cache/test_nsa_pool_host_unit.py b/test/registered/unit/mem_cache/test_nsa_pool_host_unit.py index f0462e267a87..2602e196a195 100644 --- a/test/registered/unit/mem_cache/test_nsa_pool_host_unit.py +++ b/test/registered/unit/mem_cache/test_nsa_pool_host_unit.py @@ -5,7 +5,8 @@ from sglang.srt.mem_cache.memory_pool import NSATokenToKVPool from sglang.srt.mem_cache.memory_pool_host import ( ALLOC_MEMORY_FUNCS, - NSATokenToKVPoolHost, + MLATokenToKVPoolHost, + NSAIndexerPoolHost, alloc_with_pin_memory, ) from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu @@ -58,7 +59,7 @@ def _run_device_to_host_indexer_copy(self, io_backend: str): if pin_memory: ALLOC_MEMORY_FUNCS["cuda"] = alloc_with_pin_memory try: - host_pool = NSATokenToKVPoolHost( + mla_host = MLATokenToKVPoolHost( device_pool=device_pool, host_to_device_ratio=2.0, host_size=0, @@ -66,6 +67,16 @@ def _run_device_to_host_indexer_copy(self, io_backend: str): layout="layer_first", pin_memory=pin_memory, device="cpu", + allocator_type="default", + override_kv_cache_dim=device_pool.kv_cache_dim, + ) + indexer_host = NSAIndexerPoolHost( + device_pool=device_pool, + anchor_host=mla_host, + layout="layer_first", + pin_memory=pin_memory, + device="cpu", + allocator_type="default", ) finally: ALLOC_MEMORY_FUNCS["cuda"] = original_alloc @@ -97,7 +108,10 @@ def _run_device_to_host_indexer_copy(self, io_backend: str): device="cuda" if io_backend == "kernel" else "cpu", ) - host_pool.backup_from_device_all_layer( + mla_host.backup_from_device_all_layer( + device_pool, host_indices, device_indices, io_backend + ) + indexer_host.backup_from_device_all_layer( device_pool, host_indices, device_indices, io_backend ) @@ -105,14 +119,14 @@ def _run_device_to_host_indexer_copy(self, io_backend: str): for host_page, device_page in zip( host_pages.tolist(), device_pages.tolist() ): - got = host_pool.index_k_with_scale_buffer[layer_id][host_page].cpu() + got = indexer_host.index_k_with_scale_buffer[layer_id][host_page].cpu() expected = device_pool.index_k_with_scale_buffer[layer_id][ device_page ].cpu() self.assertTrue(torch.equal(got, expected)) host_start = host_page * page_size device_start = device_page * page_size - got_kv = host_pool.kv_buffer[layer_id][ + got_kv = mla_host.kv_buffer[layer_id][ host_start : host_start + page_size ].cpu() expected_kv = device_pool.kv_buffer[layer_id][ From 25300540afe8811dc2e6f6cebd4a73a9261398aa Mon Sep 17 00:00:00 2001 From: David <12414531+DavidBellamy@users.noreply.github.com> Date: Fri, 17 Apr 2026 18:59:36 -0700 Subject: [PATCH 51/52] fix(mooncake): retry transient batch_transfer_sync failures Wrap engine.batch_transfer_sync in a small retry loop so a single transient transfer failure does not permakill a PD session. Background: the disaggregation polling loop increments session_failures on any non-zero return from batch_transfer_sync, and the default kill threshold is low. Mooncake has internal per-slice retries for WC errors but not for handshake RPC failures or endpoint resets under load, so a single transient error from those paths takes the session out for the rest of the run. Observed under heavy concurrent transfer on an H200 cluster running PD disagg with multiple decode pools. Behavior: - Retry up to N attempts with linear backoff between them. - Return 0 on first success and log a recovery warning if it took more than one attempt. - On persistent failure, log session id, slice count, total bytes, first src/dst/len, then return the last non-zero ret unchanged so the existing failure-accounting path handles it. Knobs (env, defaults preserve existing behavior at N=1): SGLANG_DISAGG_TRANSFER_RETRIES default 3 SGLANG_DISAGG_TRANSFER_RETRY_BACKOFF_MS default 50 --- .../srt/disaggregation/mooncake/conn.py | 42 +++++++++++++++++-- 1 file changed, 39 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 64d97f5c6966..c913cd4962c2 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -572,9 +572,45 @@ def _transfer_data(self, mooncake_session_id, transfer_blocks): return 0 src_addrs, dst_addrs, lengths = zip(*transfer_blocks) - return self.engine.batch_transfer_sync( - mooncake_session_id, list(src_addrs), list(dst_addrs), list(lengths) - ) + # Retry transient transfer failures before returning non-zero. Upstream + # sglang increments session_failures on any ret != 0 and the default + # kill threshold is low, so a single transient error permakills the + # session. Mooncake itself has per-slice retries for WC errors but not + # for handshake RPC failures / endpoint resets under load; this catches + # those. Defaults: 3 attempts, 50ms backoff. Knobs: + # SGLANG_DISAGG_TRANSFER_RETRIES (int, default 3) + # SGLANG_DISAGG_TRANSFER_RETRY_BACKOFF_MS (int, default 50) + max_retries = int(os.environ.get("SGLANG_DISAGG_TRANSFER_RETRIES", "3")) + backoff_ms = int(os.environ.get("SGLANG_DISAGG_TRANSFER_RETRY_BACKOFF_MS", "50")) + import time as _time + ret = 0 + for attempt in range(max_retries): + ret = self.engine.batch_transfer_sync( + mooncake_session_id, list(src_addrs), list(dst_addrs), list(lengths) + ) + if ret == 0: + if attempt > 0: + logger.warning( + f"batch_transfer_sync recovered on retry {attempt} for session {mooncake_session_id}" + ) + return 0 + if attempt < max_retries - 1: + _time.sleep(backoff_ms / 1000.0) + # On persistent failure, log enough context to diagnose. + try: + _first_src = list(src_addrs)[0] + _first_dst = list(dst_addrs)[0] + _total_bytes = sum(list(lengths)) + _n_slices = len(list(lengths)) + logger.error( + f"PD_TRANSFER_FAIL session={mooncake_session_id} ret={ret} " + f"n_slices={_n_slices} total_bytes={_total_bytes} " + f"first_src=0x{_first_src:x} first_dst=0x{_first_dst:x} " + f"first_len={list(lengths)[0]}" + ) + except Exception as _e: + logger.error(f"PD_TRANSFER_FAIL session={mooncake_session_id} ret={ret} (instrumentation error: {_e})") + return ret def _send_kvcache_generic( self, From 3397cfc2582f688578b5b74f4a1f11ba03d1ff81 Mon Sep 17 00:00:00 2001 From: David <12414531+DavidBellamy@users.noreply.github.com> Date: Fri, 17 Apr 2026 19:00:10 -0700 Subject: [PATCH 52/52] feat(mooncake): make session-failure kill threshold configurable The disaggregation polling loop currently marks a session as failed the first time session_failures hits 1. Combined with the batch_transfer_sync retry loop in the previous commit, this is too aggressive: even after the transfer-level retries recover, a future transient on the same session can still tip it into the failed set because the threshold is 1. Make the threshold configurable via SGLANG_DISAGG_SESSION_FAILURE_THRESHOLD and raise the default to 10. This trades a small amount of latency on truly dead sessions for resilience against intermittent fabric hiccups, which in our experience are by far the more common case at scale on PD disagg with mooncake. Setting the env var to 1 restores the prior fail-fast behavior for anyone who wants it. --- python/sglang/srt/disaggregation/mooncake/conn.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index c913cd4962c2..e3c92ed8ffbd 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -1302,8 +1302,12 @@ def transfer_worker( if ret != 0: with self.session_lock: self.session_failures[req.mooncake_session_id] += 1 - # Failures should never happen if the session is not dead, if the session fails once, mark it as failed - if self.session_failures[req.mooncake_session_id] >= 1: + # Mark the session dead once it accumulates enough failures. + # Threshold is configurable via SGLANG_DISAGG_SESSION_FAILURE_THRESHOLD + # (default 10) so transient failures retried by the transfer loop + # above do not immediately permakill the session. + _threshold = int(os.environ.get("SGLANG_DISAGG_SESSION_FAILURE_THRESHOLD", "10")) + if self.session_failures[req.mooncake_session_id] >= _threshold: self.failed_sessions.add(req.mooncake_session_id) logger.error( f"Session {req.mooncake_session_id} failed."