chore: import upstream snapshot with attribution

python/sglang/multimodal_gen/runtime/layers/__init__.py

@@ -0,0 +1 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo

python/sglang/multimodal_gen/runtime/layers/activation.py

@@ -0,0 +1,178 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/activation.py
+"""Custom activation functions."""
+
+import math
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from sglang.multimodal_gen.runtime.platforms import current_platform
+
+_is_cuda = current_platform.is_cuda()
+_is_hip = current_platform.is_hip()
+_is_npu = current_platform.is_npu()
+_is_xpu = current_platform.is_xpu()
+
+if _is_cuda:
+    from sglang.jit_kernel.activation import silu_and_mul
+elif _is_hip or _is_xpu:
+    from sgl_kernel import silu_and_mul
+
+
+if _is_npu:
+    import torch_npu
+# TODO (will): remove this dependency
+from sglang.multimodal_gen.runtime.layers.custom_op import CustomOp
+
+
+@CustomOp.register("silu_and_mul")
+class SiluAndMul(CustomOp):
+    """An activation function for SwiGLU.
+
+    The function computes x -> silu(x[:d]) * x[d:] where d = x.shape[-1] // 2.
+
+    Shapes:
+        x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d)
+        return: (num_tokens, d) or (batch_size, seq_len, d)
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
+        d = x.shape[-1] // 2
+        output_shape = x.shape[:-1] + (d,)
+        out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
+        silu_and_mul(x, out)
+        return out
+
+    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
+        """PyTorch-native implementation equivalent to forward()."""
+        d = x.shape[-1] // 2
+        return F.silu(x[..., :d]) * x[..., d:]
+
+    def forward_npu(self, x: torch.Tensor) -> torch.Tensor:
+        out = torch_npu.npu_swiglu(x)
+        return out
+
+    def forward_musa(self, x: torch.Tensor) -> torch.Tensor:
+        return nn.SwishGLU()(x)
+
+    def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
+        d = x.shape[-1] // 2
+        output_shape = x.shape[:-1] + (d,)
+        out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
+        silu_and_mul(x, out)
+        return out
+
+
+@CustomOp.register("gelu_and_mul")
+class GeluAndMul(CustomOp):
+    """An activation function for GeGLU.
+
+    The function computes x -> GELU(x[:d]) * x[d:] where d = x.shape[-1] // 2.
+
+    Shapes:
+        x: (batch_size, seq_len, 2 * d) or (num_tokens, 2 * d)
+        return: (batch_size, seq_len, d) or (num_tokens, d)
+    """
+
+    def __init__(self, approximate: str = "none"):
+        super().__init__()
+        self.approximate = approximate
+        if approximate not in ("none", "tanh"):
+            raise ValueError(f"Unknown approximate mode: {approximate}")
+
+    def forward_cuda(self, *args, **kwargs) -> Any:
+        return self.forward_native(*args, **kwargs)
+
+    def forward_npu(self, x: torch.Tensor) -> torch.Tensor:
+        y_npu, _ = torch_npu.npu_geglu(
+            x,
+            dim=-1,
+            approximate=1 if self.approximate == "tanh" else 0,
+            activate_left=True,
+        )
+        return y_npu
+
+    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
+        """PyTorch-native implementation equivalent to forward()."""
+        d = x.shape[-1] // 2
+        return F.gelu(x[..., :d], approximate=self.approximate) * x[..., d:]
+
+    def extra_repr(self) -> str:
+        return f"approximate={repr(self.approximate)}"
+
+
+@CustomOp.register("gelu_new")
+class NewGELU(CustomOp):
+
+    def __init__(self):
+        super().__init__()
+
+    def forward_cuda(self, *args, **kwargs) -> Any:
+        return self.forward_native(*args, **kwargs)
+
+    def forward_xpu(self, *args, **kwargs) -> Any:
+        return self.forward_native(*args, **kwargs)
+
+    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
+        """PyTorch-native implementation equivalent to forward()."""
+        c = math.sqrt(2.0 / math.pi)
+        return 0.5 * x * (1.0 + torch.tanh(c * (x + 0.044715 * torch.pow(x, 3.0))))
+
+
+@CustomOp.register("quick_gelu")
+class QuickGELU(CustomOp):
+    # https://github.com/huggingface/transformers/blob/main/src/transformers/activations.py#L90
+    def __init__(self):
+        super().__init__()
+
+    def forward_cuda(self, *args, **kwargs) -> Any:
+        return self.forward_native(*args, **kwargs)
+
+    def forward_xpu(self, *args, **kwargs) -> Any:
+        return self.forward_native(*args, **kwargs)
+
+    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
+        """PyTorch-native implementation equivalent to forward()."""
+        return x * torch.sigmoid(1.702 * x)
+
+
+_ACTIVATION_REGISTRY = {
+    "gelu": nn.GELU,
+    "gelu_new": NewGELU,
+    "gelu_pytorch_tanh": lambda: nn.GELU(approximate="tanh"),
+    "relu": nn.ReLU,
+    "silu": nn.SiLU,
+    "quick_gelu": QuickGELU,
+}
+
+
+def get_act_fn(act_fn_name: str) -> nn.Module:
+    """Get an activation function by name."""
+    act_fn_name = act_fn_name.lower()
+    if act_fn_name not in _ACTIVATION_REGISTRY:
+        raise ValueError(f"Activation function {act_fn_name!r} is not supported.")
+
+    return _ACTIVATION_REGISTRY[act_fn_name]()
+
+
+_ACTIVATION_AND_MUL_REGISTRY = {
+    "gelu": GeluAndMul,
+    "silu": SiluAndMul,
+}
+
+
+def get_act_and_mul_fn(act_fn_name: str) -> nn.Module:
+    """Get an activation-and-mul (i.e. SiluAndMul) function by name."""
+    act_fn_name = act_fn_name.lower()
+    if act_fn_name not in _ACTIVATION_AND_MUL_REGISTRY:
+        raise ValueError(f"Activation function {act_fn_name!r} is not supported.")
+
+    return _ACTIVATION_AND_MUL_REGISTRY[act_fn_name]()

python/sglang/multimodal_gen/runtime/layers/attention/STA_configuration.py

@@ -0,0 +1,414 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+import json
+import os
+from collections import defaultdict
+from typing import Any
+
+import numpy as np
+
+from sglang.multimodal_gen.utils import dict_to_3d_list
+
+
+def configure_sta(
+    mode: str = "STA_searching",
+    layer_num: int = 40,
+    time_step_num: int = 50,
+    head_num: int = 40,
+    **kwargs,
+) -> list[list[list[Any]]]:
+    """
+    Configure Sliding Tile Attention (STA) parameters based on the specified mode.
+
+    Parameters:
+    ----------
+    mode : str
+        The STA mode to use. Options are:
+        - 'STA_searching': Generate a set of mask candidates for initial search
+        - 'STA_tuning': Select best mask strategy based on previously saved results
+        - 'STA_inference': Load and use a previously tuned mask strategy
+    layer_num: int, number of layers
+    time_step_num: int, number of timesteps
+    head_num: int, number of heads
+
+    **kwargs : dict
+        Mode-specific parameters:
+
+        For 'STA_searching':
+        - mask_candidates: list of str, optional, mask candidates to use
+        - mask_selected: list of int, optional, indices of selected masks
+
+        For 'STA_tuning':
+        - mask_search_files_path: str, required, path to mask search results
+        - mask_candidates: list of str, optional, mask candidates to use
+        - mask_selected: list of int, optional, indices of selected masks
+        - skip_time_steps: int, optional, number of time steps to use full attention (default 12)
+        - save_dir: str, optional, directory to save mask strategy (default "mask_candidates")
+
+        For 'STA_inference':
+        - load_path: str, optional, path to load mask strategy (default "mask_candidates/mask_strategy.json")
+    """
+    valid_modes = ["STA_searching", "STA_tuning", "STA_inference", "STA_tuning_cfg"]
+    if mode not in valid_modes:
+        raise ValueError(f"Mode must be one of {valid_modes}, got {mode}")
+
+    if mode == "STA_searching":
+        # Get parameters with defaults
+        mask_candidates: list[str] | None = kwargs.get("mask_candidates")
+        if mask_candidates is None:
+            raise ValueError("mask_candidates is required for STA_searching mode")
+        mask_selected: list[int] = kwargs.get(
+            "mask_selected", list(range(len(mask_candidates)))
+        )
+
+        # Parse selected masks
+        selected_masks: list[list[int]] = []
+        for index in mask_selected:
+            mask = mask_candidates[index]
+            masks_list = [int(x) for x in mask.split(",")]
+            selected_masks.append(masks_list)
+
+        # Create 3D mask structure with fixed dimensions (t=50, l=60)
+        masks_3d: list[list[list[list[int]]]] = []
+        for i in range(time_step_num):  # Fixed t dimension = 50
+            row = []
+            for j in range(layer_num):  # Fixed l dimension = 60
+                row.append(selected_masks)  # Add all masks at each position
+            masks_3d.append(row)
+
+        return masks_3d
+
+    elif mode == "STA_tuning":
+        # Get required parameters
+        mask_search_files_path: str | None = kwargs.get("mask_search_files_path")
+        if not mask_search_files_path:
+            raise ValueError("mask_search_files_path is required for STA_tuning mode")
+
+        # Get optional parameters with defaults
+        mask_candidates_tuning: list[str] | None = kwargs.get("mask_candidates")
+        if mask_candidates_tuning is None:
+            raise ValueError("mask_candidates is required for STA_tuning mode")
+        mask_selected_tuning: list[int] = kwargs.get(
+            "mask_selected", list(range(len(mask_candidates_tuning)))
+        )
+        skip_time_steps_tuning: int | None = kwargs.get("skip_time_steps")
+        save_dir_tuning: str | None = kwargs.get("save_dir", "mask_candidates")
+
+        # Parse selected masks
+        selected_masks_tuning: list[list[int]] = []
+        for index in mask_selected_tuning:
+            mask = mask_candidates_tuning[index]
+            masks_list = [int(x) for x in mask.split(",")]
+            selected_masks_tuning.append(masks_list)
+
+        # Read JSON results
+        results = read_specific_json_files(mask_search_files_path)
+        averaged_results = average_head_losses(results, selected_masks_tuning)
+
+        # Add full attention mask for specific cases
+        full_attention_mask_tuning: list[int] | None = kwargs.get("full_attention_mask")
+        if full_attention_mask_tuning is not None:
+            selected_masks_tuning.append(full_attention_mask_tuning)
+
+        # Select best mask strategy
+        timesteps_tuning: int = kwargs.get("timesteps", time_step_num)
+        if skip_time_steps_tuning is None:
+            skip_time_steps_tuning = 12
+        mask_strategy, sparsity, strategy_counts = select_best_mask_strategy(
+            averaged_results,
+            selected_masks_tuning,
+            skip_time_steps_tuning,
+            timesteps_tuning,
+            head_num,
+        )
+
+        # Save mask strategy
+        if save_dir_tuning is not None:
+            os.makedirs(save_dir_tuning, exist_ok=True)
+            file_path = os.path.join(
+                save_dir_tuning, f"mask_strategy_s{skip_time_steps_tuning}.json"
+            )
+            with open(file_path, "w") as f:
+                json.dump(mask_strategy, f, indent=4)
+            print(f"Successfully saved mask_strategy to {file_path}")
+
+        # Print sparsity and strategy counts for information
+        print(f"Overall sparsity: {sparsity:.4f}")
+        print("\nStrategy usage counts:")
+        total_heads = time_step_num * layer_num * head_num  # Fixed dimensions
+        for strategy, count in strategy_counts.items():
+            print(f"Strategy {strategy}: {count} heads ({count/total_heads*100:.2f}%)")
+
+        # Convert dictionary to 3D list with fixed dimensions
+        mask_strategy_3d = dict_to_3d_list(
+            mask_strategy, t_max=time_step_num, l_max=layer_num, h_max=head_num
+        )
+
+        return mask_strategy_3d
+    elif mode == "STA_tuning_cfg":
+        # Get required parameters for both positive and negative paths
+        mask_search_files_path_pos: str | None = kwargs.get(
+            "mask_search_files_path_pos"
+        )
+        mask_search_files_path_neg: str | None = kwargs.get(
+            "mask_search_files_path_neg"
+        )
+        save_dir_cfg: str | None = kwargs.get("save_dir")
+
+        if (
+            not mask_search_files_path_pos
+            or not mask_search_files_path_neg
+            or not save_dir_cfg
+        ):
+            raise ValueError(
+                "mask_search_files_path_pos, mask_search_files_path_neg, and save_dir are required for STA_tuning_cfg mode"
+            )
+
+        # Get optional parameters with defaults
+        mask_candidates_cfg: list[str] | None = kwargs.get("mask_candidates")
+        if mask_candidates_cfg is None:
+            raise ValueError("mask_candidates is required for STA_tuning_cfg mode")
+        mask_selected_cfg: list[int] = kwargs.get(
+            "mask_selected", list(range(len(mask_candidates_cfg)))
+        )
+        skip_time_steps_cfg: int | None = kwargs.get("skip_time_steps")
+
+        # Parse selected masks
+        selected_masks_cfg: list[list[int]] = []
+        for index in mask_selected_cfg:
+            mask = mask_candidates_cfg[index]
+            masks_list = [int(x) for x in mask.split(",")]
+            selected_masks_cfg.append(masks_list)
+
+        # Read JSON results for both positive and negative paths
+        pos_results = read_specific_json_files(mask_search_files_path_pos)
+        neg_results = read_specific_json_files(mask_search_files_path_neg)
+        # Combine positive and negative results into one list
+        combined_results = pos_results + neg_results
+
+        # Average the combined results
+        averaged_results = average_head_losses(combined_results, selected_masks_cfg)
+
+        # Add full attention mask for specific cases
+        full_attention_mask_cfg: list[int] | None = kwargs.get("full_attention_mask")
+        if full_attention_mask_cfg is not None:
+            selected_masks_cfg.append(full_attention_mask_cfg)
+
+        timesteps_cfg: int = kwargs.get("timesteps", time_step_num)
+        if skip_time_steps_cfg is None:
+            skip_time_steps_cfg = 12
+        # Select best mask strategy using combined results
+        mask_strategy, sparsity, strategy_counts = select_best_mask_strategy(
+            averaged_results,
+            selected_masks_cfg,
+            skip_time_steps_cfg,
+            timesteps_cfg,
+            head_num,
+        )
+
+        # Save mask strategy
+        os.makedirs(save_dir_cfg, exist_ok=True)
+        file_path = os.path.join(
+            save_dir_cfg, f"mask_strategy_s{skip_time_steps_cfg}.json"
+        )
+        with open(file_path, "w") as f:
+            json.dump(mask_strategy, f, indent=4)
+        print(f"Successfully saved mask_strategy to {file_path}")
+
+        # Print sparsity and strategy counts for information
+        print(f"Overall sparsity: {sparsity:.4f}")
+        print("\nStrategy usage counts:")
+        total_heads = time_step_num * layer_num * head_num  # Fixed dimensions
+        for strategy, count in strategy_counts.items():
+            print(f"Strategy {strategy}: {count} heads ({count/total_heads*100:.2f}%)")
+
+        # Convert dictionary to 3D list with fixed dimensions
+        mask_strategy_3d = dict_to_3d_list(
+            mask_strategy, t_max=time_step_num, l_max=layer_num, h_max=head_num
+        )
+
+        return mask_strategy_3d
+
+    else:  # STA_inference
+        # Get parameters with defaults
+        load_path: str | None = kwargs.get(
+            "load_path", "mask_candidates/mask_strategy.json"
+        )
+        if load_path is None:
+            raise ValueError("load_path is required for STA_inference mode")
+
+        # Load previously saved mask strategy
+        with open(load_path) as f:
+            mask_strategy = json.load(f)
+
+        # Convert dictionary to 3D list with fixed dimensions
+        mask_strategy_3d = dict_to_3d_list(
+            mask_strategy, t_max=time_step_num, l_max=layer_num, h_max=head_num
+        )
+
+        return mask_strategy_3d
+
+
+# Helper functions
+
+
+def read_specific_json_files(folder_path: str) -> list[dict[str, Any]]:
+    """Read and parse JSON files containing mask search results."""
+    json_contents: list[dict[str, Any]] = []
+
+    # List files only in the current directory (no walk)
+    files = os.listdir(folder_path)
+    # Filter files
+    matching_files = [f for f in files if "mask" in f and f.endswith(".json")]
+    print(f"Found {len(matching_files)} matching files: {matching_files}")
+
+    for file_name in matching_files:
+        file_path = os.path.join(folder_path, file_name)
+        with open(file_path) as file:
+            data = json.load(file)
+            json_contents.append(data)
+
+    return json_contents
+
+
+def average_head_losses(
+    results: list[dict[str, Any]], selected_masks: list[list[int]]
+) -> dict[str, dict[str, np.ndarray]]:
+    """Average losses across all prompts for each mask strategy."""
+    # Initialize a dictionary to store the averaged results
+    averaged_losses: dict[str, dict[str, np.ndarray]] = {}
+    loss_type = "L2_loss"
+    # Get all loss types (e.g., 'L2_loss')
+    averaged_losses[loss_type] = {}
+
+    for mask in selected_masks:
+        mask_str = str(mask)
+        data_shape = np.array(results[0][loss_type][mask_str]).shape
+        accumulated_data = np.zeros(data_shape)
+
+        # Sum across all prompts
+        for prompt_result in results:
+            accumulated_data += np.array(prompt_result[loss_type][mask_str])
+
+        # Average by dividing by number of prompts
+        averaged_data = accumulated_data / len(results)
+        averaged_losses[loss_type][mask_str] = averaged_data
+
+    return averaged_losses
+
+
+def select_best_mask_strategy(
+    averaged_results: dict[str, dict[str, np.ndarray]],
+    selected_masks: list[list[int]],
+    skip_time_steps: int = 12,
+    timesteps: int = 50,
+    head_num: int = 40,
+) -> tuple[dict[str, list[int]], float, dict[str, int]]:
+    """Select the best mask strategy for each head based on loss minimization."""
+    best_mask_strategy: dict[str, list[int]] = {}
+    loss_type = "L2_loss"
+    # Get the shape of time steps and layers
+    layers = len(averaged_results[loss_type][str(selected_masks[0])][0])
+
+    # Counter for sparsity calculation
+    total_tokens = 0  # total number of masked tokens
+    total_length = 0  # total sequence length
+
+    strategy_counts: dict[str, int] = {str(strategy): 0 for strategy in selected_masks}
+    full_attn_strategy = selected_masks[-1]  # Last strategy is full attention
+    print(f"Strategy {full_attn_strategy}, skip first {skip_time_steps} steps ")
+
+    for t in range(timesteps):
+        for layer_idx in range(layers):
+            for h in range(head_num):
+                if t < skip_time_steps:  # First steps use full attention
+                    strategy = full_attn_strategy
+                else:
+                    # Get losses for this head across all strategies
+                    head_losses = []
+                    for strategy in selected_masks[:-1]:  # Exclude full attention
+                        head_losses.append(
+                            averaged_results[loss_type][str(strategy)][t][layer_idx][h]
+                        )
+
+                    # Find which strategy gives minimum loss
+                    best_strategy_idx = np.argmin(head_losses)
+                    strategy = selected_masks[best_strategy_idx]
+
+                best_mask_strategy[f"{t}_{layer_idx}_{h}"] = strategy
+
+                # Calculate sparsity
+                nums = strategy  # strategy is already a list of numbers
+                total_tokens += (
+                    nums[0] * nums[1] * nums[2]
+                )  # masked tokens for chosen strategy
+                total_length += (
+                    full_attn_strategy[0]
+                    * full_attn_strategy[1]
+                    * full_attn_strategy[2]
+                )
+
+                # Count strategy usage
+                strategy_counts[str(strategy)] += 1
+
+    overall_sparsity = 1 - total_tokens / total_length
+
+    return best_mask_strategy, overall_sparsity, strategy_counts
+
+
+def save_mask_search_results(
+    mask_search_final_result: list[dict[str, list[float]]],
+    prompt: str,
+    mask_strategies: list[str],
+    output_dir: str = "output/mask_search_result/",
+) -> str | None:
+    if not mask_search_final_result:
+        print("No mask search results to save")
+        return None
+
+    # Create result dictionary with defaultdict for nested lists
+    mask_search_dict: dict[str, dict[str, list[list[float]]]] = {
+        "L2_loss": defaultdict(list),
+        "L1_loss": defaultdict(list),
+    }
+
+    mask_selected = list(range(len(mask_strategies)))
+    selected_masks: list[list[int]] = []
+    for index in mask_selected:
+        mask = mask_strategies[index]
+        masks_list = [int(x) for x in mask.split(",")]
+        selected_masks.append(masks_list)
+
+    # Process each mask strategy
+    for i, mask_strategy in enumerate(selected_masks):
+        mask_strategy_str = str(mask_strategy)
+        # Process L2 loss
+        step_results: list[list[float]] = []
+        for step_data in mask_search_final_result:
+            if isinstance(step_data, dict) and "L2_loss" in step_data:
+                layer_losses = [float(loss) for loss in step_data["L2_loss"]]
+                step_results.append(layer_losses)
+        mask_search_dict["L2_loss"][mask_strategy_str] = step_results
+
+        step_results = []
+        for step_data in mask_search_final_result:
+            if isinstance(step_data, dict) and "L1_loss" in step_data:
+                layer_losses = [float(loss) for loss in step_data["L1_loss"]]
+                step_results.append(layer_losses)
+        mask_search_dict["L1_loss"][mask_strategy_str] = step_results
+
+    # Create the output directory if it doesn't exist
+    os.makedirs(output_dir, exist_ok=True)
+
+    # Create a filename based on the first 20 characters of the prompt
+    filename = prompt[:50].replace(" ", "_")
+    filepath = os.path.join(output_dir, f"mask_search_{filename}.json")
+
+    # Save the results to a JSON file
+    with open(filepath, "w") as f:
+        json.dump(mask_search_dict, f, indent=4)
+
+    print(f"Successfully saved mask research results to {filepath}")
+
+    return filepath

python/sglang/multimodal_gen/runtime/layers/attention/__init__.py

@@ -0,0 +1,38 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.layers.attention.layer import (
+    DynamicVarlenMaskMeta,
+    LocalAttention,
+    UlyssesAttention,
+    UlyssesAttention_VSA,
+    USPAttention,
+    build_varlen_mask_meta,
+    build_varlen_mask_meta_from_lengths,
+    build_varlen_mask_meta_from_ranges,
+)
+from sglang.multimodal_gen.runtime.layers.attention.selector import get_attn_backend
+from sglang.multimodal_gen.runtime.layers.attention.turbo_layer import MinimalA2AAttnOp
+
+__all__ = [
+    "USPAttention",
+    "LocalAttention",
+    "DynamicVarlenMaskMeta",
+    "UlyssesAttention",
+    "UlyssesAttention_VSA",
+    "MinimalA2AAttnOp",
+    "AttentionBackend",
+    "AttentionMetadata",
+    "AttentionMetadataBuilder",
+    # "AttentionState",
+    "get_attn_backend",
+    "build_varlen_mask_meta",
+    "build_varlen_mask_meta_from_lengths",
+    "build_varlen_mask_meta_from_ranges",
+]

python/sglang/multimodal_gen/runtime/layers/attention/backends/__init__.py

@@ -0,0 +1 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo

python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter.py

@@ -0,0 +1,207 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+import os
+
+import aiter
+import torch
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.platforms.aiter import USE_AITER_GFX95
+
+logger = logging.getLogger(__name__)
+
+_use_fp8_attn = os.environ.get("SGLANG_DIFFUSION_AITER_FP8_ATTN", "0") == "1"
+_fp8_dtype = torch.float8_e4m3fn
+
+# fmha_fwd_hd128_fp8_gfx950 ASM kernel. Support full MHA with q/k/v head_dim == 128 -- e.g., Wan 2.2 self- and cross-attention.
+_FMHA_FP8_HEAD_DIM = 128
+
+
+if _use_fp8_attn:
+    logger.info("DiT FP8 attention enabled via SGLANG_DIFFUSION_AITER_FP8_ATTN=1")
+
+
+def _can_use_fmha_fp8_prefill(
+    q_head_dim: int,
+    k_head_dim: int,
+    v_head_dim: int,
+    num_heads: int,
+    num_kv_heads: int,
+) -> bool:
+    """True if MHA q/k/v head_dim==128 on a gfx950-class arch."""
+    if not USE_AITER_GFX95:
+        return False
+    if num_kv_heads != num_heads:
+        return False
+    return q_head_dim == k_head_dim == v_head_dim == _FMHA_FP8_HEAD_DIM
+
+
+def _fmha_fp8_prefill_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    softmax_scale: float,
+    is_causal: bool,
+    q_scale: torch.Tensor,
+    k_scale: torch.Tensor,
+    v_scale: torch.Tensor,
+) -> torch.Tensor:
+    """
+    FP8 FMHA prefill via aiter.flash_attn_fp8_pertensor_func.
+
+    Expects q, k, v as (batch, seqlen, nheads, 128) FP8, contiguous.
+    """
+
+    def _ensure_fp8_descale(scale: torch.Tensor) -> torch.Tensor:
+        """Per-tensor descale as shape (1,) float32 for flash_attn_fp8_pertensor_func."""
+        return scale.to(dtype=torch.float32).reshape(1).contiguous()
+
+    q = q.contiguous()
+    k = k.contiguous()
+    v = v.contiguous()
+    q_descale = _ensure_fp8_descale(q_scale)
+    k_descale = _ensure_fp8_descale(k_scale)
+    v_descale = _ensure_fp8_descale(v_scale)
+
+    return aiter.flash_attn_fp8_pertensor_func(
+        q,
+        k,
+        v,
+        q_descale,
+        k_descale,
+        v_descale,
+        causal=is_causal,
+        softmax_scale=softmax_scale,
+        window_size=(-1, -1, 0),
+    )
+
+
+class AITerBackend(AttentionBackend):
+    """
+    Backend for AITemplate attention implementation.
+    """
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.AITER
+
+    @staticmethod
+    def get_impl_cls() -> type["AITerImpl"]:
+        return AITerImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        # AITer backend does not require special metadata.
+        return AttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
+        raise NotImplementedError("AITer backend does not have a metadata builder.")
+
+
+class AITerImpl(AttentionImpl):
+    """
+    Implementation of attention using AITemplate.
+    """
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        dropout_p: float = 0.0,
+        **extra_impl_args,
+    ) -> None:
+        if num_kv_heads is not None and num_kv_heads != num_heads:
+            raise NotImplementedError(
+                "AITer backend does not support Grouped Query Attention yet."
+            )
+        self.causal = causal
+        self.dropout_p = dropout_p
+        self.softmax_scale = softmax_scale
+
+    @torch.compiler.disable
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        """
+        Performs attention using one of:
+          - _fmha_fp8_prefill_attention (FP8, SGLANG_DIFFUSION_AITER_FP8_ATTN=1 when eligible)
+          - flash_attn_func (BF16, default or FP8 fallback for unsupported shapes)
+
+        Args:
+            query: Query tensor of shape [batch_size, seq_len, num_heads, head_dim]
+            key: Key tensor of shape [batch_size, seq_len, num_heads, head_dim]
+            value: Value tensor of shape [batch_size, seq_len, num_heads, head_dim]
+            attn_metadata: Metadata for the attention operation (unused).
+
+        Returns:
+            Output tensor of shape [batch_size, seq_len, num_heads, head_dim]
+        """
+        if _use_fp8_attn:
+            if query.dtype != _fp8_dtype:
+                q_fp8, q_scale = aiter.per_tensor_quant(query, quant_dtype=_fp8_dtype)
+                k_fp8, k_scale = aiter.per_tensor_quant(key, quant_dtype=_fp8_dtype)
+                v_fp8, v_scale = aiter.per_tensor_quant(value, quant_dtype=_fp8_dtype)
+            else:
+                q_fp8, k_fp8, v_fp8 = query, key, value
+                one = torch.tensor(1.0, dtype=torch.float32, device=query.device)
+                q_scale = k_scale = v_scale = one
+
+            d_q = q_fp8.shape[-1]
+            d_k = k_fp8.shape[-1]
+            d_v = v_fp8.shape[-1]
+            h_q = q_fp8.shape[2]
+            h_kv = k_fp8.shape[2]
+
+            if _can_use_fmha_fp8_prefill(d_q, d_k, d_v, h_q, h_kv):
+                return _fmha_fp8_prefill_attention(
+                    q_fp8,
+                    k_fp8,
+                    v_fp8,
+                    softmax_scale=self.softmax_scale,
+                    is_causal=self.causal,
+                    q_scale=q_scale,
+                    k_scale=k_scale,
+                    v_scale=v_scale,
+                )
+
+            logger.warning_once(
+                "FP8 FMHA prefill unsupported for this shape (need gfx950-class AITER, "
+                "full MHA, q/k/v head_dim=%d; got q=%d, k=%d, v=%d, num_heads=%d, "
+                "num_kv_heads=%d). Falling back to BF16.",
+                _FMHA_FP8_HEAD_DIM,
+                d_q,
+                d_k,
+                d_v,
+                h_q,
+                h_kv,
+            )
+
+        # BF16 path
+        output, _ = aiter.flash_attn_func(
+            query,
+            key,
+            value,
+            dropout_p=self.dropout_p,
+            causal=self.causal,
+            return_attn_probs=False,
+            return_lse=True,
+        )
+        return output

python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter_sage.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+
+
+import torch
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+
+
+class AITERSageBackend(AttentionBackend):
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.AITER_SAGE
+
+    @staticmethod
+    def get_impl_cls() -> type["AITERSageImpl"]:
+        return AITERSageImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        # AITER Sage backend does not require special metadata.
+        return AttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
+        raise NotImplementedError(
+            "AITER Sage backend does not have a metadata builder."
+        )
+
+
+class AITERSageImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        dropout_p: float = 0.0,
+        **extra_impl_args,
+    ) -> None:
+
+        try:
+            from aiter.ops.triton.attention.fav3_sage import fav3_sage_wrapper_func
+
+            self.aiter_sage_attn_fn = fav3_sage_wrapper_func
+        except ImportError:
+            raise ImportError(
+                "AITER Sage attention is not available, please update AITER version."
+            )
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        """
+        Performs attention using aiter sage backend.
+
+        Args:
+            query: Query tensor of shape [batch_size, seq_len, head_num, head_dim]
+            key: Key tensor of shape [batch_size, seq_len, head_num, head_dim]
+            value: Value tensor of shape [batch_size, seq_len, head_num, head_dim]
+            attn_metadata: Metadata for the attention operation (unused).
+
+        Returns:
+            Output tensor of shape [batch_size, seq_len, head_num, head_dim]
+        """
+
+        output = self.aiter_sage_attn_fn(query, key, value)
+        return output

python/sglang/multimodal_gen/runtime/layers/attention/backends/ascend_fa.py

@@ -0,0 +1,104 @@
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class AscendFAMetadata:
+    pass
+
+
+class AscendFAMetadataBuilder(AttentionMetadataBuilder):
+    def __init__(self) -> None:
+        pass
+
+    def prepare(self) -> None:
+        pass
+
+    def build(
+        self,
+        **kwargs: dict[str, Any],
+    ) -> AttentionMetadata:
+        return AscendFAMetadata()
+
+
+class AscendFABackend(AttentionBackend):
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.FA
+
+    @staticmethod
+    def get_impl_cls() -> type["AscendFAImpl"]:
+        return AscendFAImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        raise NotImplementedError
+
+    @staticmethod
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
+        return AscendFAMetadataBuilder
+
+
+class AscendFAImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+        return_softmax_lse: bool = False,
+    ) -> torch.Tensor:
+        mask = None
+        num_heads, num_key_value_heads = query.shape[2], key.shape[2]
+        if self.causal:
+            seq_len = query.shape[1]
+            mask = torch.triu(
+                torch.ones(seq_len, seq_len, device=query.device), diagonal=1
+            ).bool()
+        # transpose to bs, heads, seq_len, head_dim
+        query = query.transpose(1, 2)
+        key = key.transpose(1, 2)
+        value = value.transpose(1, 2)
+        output, lse = torch.ops.npu.npu_fused_infer_attention_score(
+            query,
+            key,
+            value,
+            num_heads=num_heads,
+            num_key_value_heads=num_key_value_heads,
+            scale=self.softmax_scale,
+            input_layout="BNSD",
+            softmax_lse_flag=return_softmax_lse,
+            atten_mask=mask,
+        )
+        output = output.transpose(1, 2)
+        if return_softmax_lse:
+            return output, lse
+        return output

python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py

@@ -0,0 +1,179 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/attention/backends/abstract.py
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, fields
+from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar
+
+if TYPE_CHECKING:
+    pass
+
+import torch
+
+from sglang.kernel_api_logging import wrap_method_with_debug_kernel_once
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+
+
+class AttentionBackend(ABC):
+    """Abstract class for attention backends."""
+
+    # For some attention backends, we allocate an output tensor before
+    # calling the custom op. When piecewise cudagraph is enabled, this
+    # makes sure the output tensor is allocated inside the cudagraph.
+    accept_output_buffer: bool = False
+
+    @staticmethod
+    @abstractmethod
+    def get_enum() -> AttentionBackendEnum:
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def get_impl_cls() -> type["AttentionImpl"]:
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        raise NotImplementedError
+
+    # @staticmethod
+    # @abstractmethod
+    # def get_state_cls() -> Type["AttentionState"]:
+    #     raise NotImplementedError
+
+    # @classmethod
+    # def make_metadata(cls, *args, **kwargs) -> "AttentionMetadata":
+    #     return cls.get_metadata_cls()(*args, **kwargs)
+
+    @staticmethod
+    @abstractmethod
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
+        return None
+
+
+@dataclass
+class AttentionMetadata:
+    """Attention metadata for prefill and decode batched together."""
+
+    # Current step of diffusion process
+    current_timestep: int
+
+    def asdict_zerocopy(self, skip_fields: set[str] | None = None) -> dict[str, Any]:
+        """Similar to dataclasses.asdict, but avoids deepcopying."""
+        if skip_fields is None:
+            skip_fields = set()
+        # Note that if we add dataclasses as fields, they will need
+        # similar handling.
+        return {
+            field.name: getattr(self, field.name)
+            for field in fields(self)
+            if field.name not in skip_fields
+        }
+
+
+T = TypeVar("T", bound=AttentionMetadata)
+
+
+class AttentionMetadataBuilder(ABC, Generic[T]):
+    """Abstract class for attention metadata builders."""
+
+    @abstractmethod
+    def __init__(self) -> None:
+        """Create the builder, remember some configuration and parameters."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def prepare(self) -> None:
+        """Prepare for one batch."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def build(
+        self,
+        **kwargs: dict[str, Any],
+    ) -> AttentionMetadata:
+        """Build attention metadata with on-device tensors."""
+        raise NotImplementedError
+
+
+class AttentionLayer(Protocol):
+
+    _k_scale: torch.Tensor
+    _v_scale: torch.Tensor
+    _k_scale_float: float
+    _v_scale_float: float
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        kv_cache: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor: ...
+
+
+class AttentionImpl(ABC, Generic[T]):
+
+    @abstractmethod
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        raise NotImplementedError
+
+    def preprocess_qkv(self, qkv: torch.Tensor, attn_metadata: T) -> torch.Tensor:
+        """Preprocess QKV tensor before performing attention operation.
+
+        Default implementation returns the tensor unchanged.
+        Subclasses can override this to implement custom preprocessing
+        like reshaping, tiling, scaling, or other transformations.
+
+        Called AFTER all_to_all for distributed attention
+
+        """
+        return qkv
+
+    def postprocess_output(
+        self,
+        output: torch.Tensor,
+        attn_metadata: T,
+    ) -> torch.Tensor:
+        """Postprocess the output tensor after the attention operation.
+
+        Default implementation returns the tensor unchanged.
+        Subclasses can override this to implement custom postprocessing
+        like untiling, scaling, or other transformations.
+
+        Called BEFORE all_to_all for distributed attention
+
+        """
+
+        return output
+
+    @abstractmethod
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: T,
+    ) -> torch.Tensor:
+        raise NotImplementedError
+
+
+def wrap_attention_impl_forward(attn_impl: AttentionImpl) -> AttentionImpl:
+    return wrap_method_with_debug_kernel_once(
+        attn_impl,
+        "forward",
+        op_name=f"diffusion.attn_impl.{attn_impl.__class__.__name__}.forward",
+    )

python/sglang/multimodal_gen/runtime/layers/attention/backends/block_sparse_attn.py

@@ -0,0 +1,279 @@
+from dataclasses import dataclass
+from typing import Any
+
+import attentions  # noqa: F401
+import torch
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.layers.attention.backends.laser_attn import (
+    LaserAttentionBackend,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+BSA_BLOCK_SIZE = 128
+
+
+class BlockSparseAttentionBackend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [32, 64, 96, 128]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.BLOCK_SPARSE_ATTN
+
+    @staticmethod
+    def get_impl_cls() -> type["BlockSparseAttentionImpl"]:
+        return BlockSparseAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["BlockSparseAttentionMetadata"]:
+        return BlockSparseAttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["BlockSparseAttentionMetadataBuilder"]:
+        return BlockSparseAttentionMetadataBuilder
+
+
+@dataclass
+class BlockSparseAttentionMetadata(AttentionMetadata):
+    current_timestep: int
+    skip_first_steps: int
+    sparsity: float
+    block_frame_stride: int
+
+
+class BlockSparseAttentionMetadataBuilder(AttentionMetadataBuilder):
+    def __init__(self) -> None:
+        pass
+
+    def prepare(self) -> None:
+        pass
+
+    def build(
+        self,
+        current_timestep: int,
+        skip_first_steps: int,
+        sparsity: float,
+        raw_latent_shape: list[int],
+        patch_size: tuple[int, int, int],
+        **kwargs: dict[str, Any],
+    ) -> BlockSparseAttentionMetadata:
+        """
+        Builds BlockSparseAttention metadata.
+
+        Args:
+            current_timestep: The current diffusion timestep.
+            skip_first_steps: Number of initial timesteps to skip before applying
+                sparsity. Must be non‑negative.
+            sparsity: Fraction of tokens to drop (block‑wise) in the block sparse
+                attention mechanism. Must be in the range [0.0, 1.0).
+            raw_latent_shape: Shape of the latent tensor before patching.
+            patch_size: Patch size as (T, height, width). Only the height
+                and width components are used to divide the latent dimensions.
+            **kwargs: Additional keyword arguments (ignored, but accepted for
+                compatibility with base class or calling conventions).
+
+        Returns:
+            BlockSparseAttentionMetadata
+        Note:
+            The `block_frame_stride` is needed to set the first blocks to be non‑sparse.
+        """
+        if not (skip_first_steps >= 0 and 0.0 <= sparsity < 1.0):
+            raise ValueError(
+                (
+                    "Invalid attention metadata values."
+                    f"Sparsity should be in [0, 1), skip_first_steps should be non-negative."
+                    f"Got sparsity={sparsity}, skip_first_steps={skip_first_steps}"
+                )
+            )
+
+        if sparsity == 0.0:
+            logger.warning(
+                (
+                    "Sparsity is set to 0.0, which means no tokens will be dropped."
+                    "For better performance use Laser Attention or increase sparsity."
+                )
+            )
+
+        if len(raw_latent_shape) >= 5:
+            latent_height, latent_width = raw_latent_shape[3:5]
+        else:
+            latent_height, latent_width = raw_latent_shape[-2:]
+
+        latent_height //= patch_size[1]
+        latent_width //= patch_size[2]
+
+        frame_stride = latent_height * latent_width
+        block_frame_stride = (frame_stride + BSA_BLOCK_SIZE - 1) // BSA_BLOCK_SIZE
+
+        return BlockSparseAttentionMetadata(
+            current_timestep=current_timestep,
+            skip_first_steps=skip_first_steps,
+            sparsity=sparsity,
+            block_frame_stride=block_frame_stride,
+        )
+
+
+class BlockSparseAttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads or num_heads
+        self.block_size = BSA_BLOCK_SIZE
+        self.stride = 8
+        self.default_tokens = 214748647
+
+        self.laser_attn_impl = LaserAttentionBackend.get_impl_cls()(
+            num_heads,
+            head_size,
+            causal,
+            softmax_scale,
+            num_kv_heads,
+            prefix,
+            **extra_impl_args,
+        )
+
+    def _get_estimate_mask(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        sparsity: float,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return torch.ops.attentions.sparse_block_estimate(
+            query=query,
+            key=key,
+            actual_seq_lengths=None,
+            actual_seq_lengths_kv=None,
+            input_layout="BNSD",
+            stride=self.stride,
+            sparse_size=self.block_size,
+            num_heads=query.shape[1],
+            num_key_value_heads=key.shape[1],
+            scale_value=self.softmax_scale / self.stride,
+            threshold=1.0,
+            causal=self.causal,
+            keep_sink=True,
+            keep_recent=True,
+            row_sparse=1.0 - sparsity,
+        )
+
+    def _block_sparse_attention(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        smask: torch.Tensor,
+        sct: torch.Tensor,
+    ) -> torch.Tensor:
+        return torch.ops.attentions.block_sparse_attention(
+            query=query,
+            key=key,
+            value=value,
+            sparse_mask=smask,
+            sparse_count_table=sct,
+            input_layout="BNSD",
+            sparse_size=self.block_size,
+            num_heads=query.shape[1],
+            num_key_value_heads=key.shape[1],
+            scale_value=self.softmax_scale,
+            causal=self.causal,
+            inner_precise=1,
+            pre_tokens=self.default_tokens,
+            next_tokens=self.default_tokens,
+            actual_seq_lengths=None,
+            actual_seq_lengths_kv=None,
+        )
+
+    def _get_smask(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        block_frame_stride: int,
+        sparsity: float,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        smask, sct = self._get_estimate_mask(
+            query,
+            key,
+            sparsity,
+        )
+
+        seq_len = smask.shape[2]
+
+        # Set the first blocks to be non-sparse to ensure the quality of the first few steps
+        smask[:, :, :block_frame_stride, :seq_len] = 1
+        smask[:, :, :seq_len, :block_frame_stride] = 1
+        smask = smask.to(torch.int8)
+        sct = smask.sum(dim=-1, dtype=torch.int32)
+        return smask, sct
+
+    def _adaptive_block_sparse_attention(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        block_frame_stride: int,
+        sparsity: float,
+    ) -> torch.Tensor:
+        # TODO Currently implementation for BSND input layout has quality issues
+        # When the implementation is improved, transposes can be removed
+        q = query.permute(0, 2, 1, 3).contiguous()
+        k = key.permute(0, 2, 1, 3).contiguous()
+        v = value.permute(0, 2, 1, 3).contiguous()
+
+        smask, sct = self._get_smask(
+            q,
+            k,
+            block_frame_stride,
+            sparsity,
+        )
+        output = self._block_sparse_attention(q, k, v, smask, sct)
+        output = output.permute(0, 2, 1, 3).contiguous()
+
+        return output
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        if attn_metadata.current_timestep < attn_metadata.skip_first_steps:
+            output = self.laser_attn_impl.forward(
+                query,
+                key,
+                value,
+                attn_metadata,
+            )
+        else:
+            output = self._adaptive_block_sparse_attention(
+                query,
+                key,
+                value,
+                attn_metadata.block_frame_stride,
+                attn_metadata.sparsity,
+            )
+
+        return output

python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py

@@ -0,0 +1,445 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass
+from typing import Any, List, Optional, Tuple
+
+import torch
+
+from sglang.jit_kernel.flash_attention import flash_attn_varlen_func
+from sglang.multimodal_gen.runtime.layers.utils import register_custom_op
+from sglang.multimodal_gen.runtime.platforms import (
+    AttentionBackendEnum,
+)
+
+
+def maybe_contiguous(x: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
+    return x.contiguous() if x is not None and x.stride(-1) != 1 else x
+
+
+# -----------------------------
+# Fake implementations for schema / tracing
+# custom op schema requires FIXED return structure.
+# We provide TWO ops:
+# 1) out-only op: always returns Tensor
+# 2) out+lse op: always returns Tuple[Tensor, Tensor]
+# -----------------------------
+def flash_attn_varlen_func_fake_out(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    max_seqlen_q: Optional[int] = None,
+    max_seqlen_k: Optional[int] = None,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    page_table: Optional[torch.Tensor] = None,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    qv: Optional[torch.Tensor] = None,
+    q_descale: Optional[torch.Tensor] = None,
+    k_descale: Optional[torch.Tensor] = None,
+    v_descale: Optional[torch.Tensor] = None,
+    window_size: Optional[List[int]] = None,
+    attention_chunk: int = 0,
+    softcap: float = 0.0,
+    num_splits: int = 1,
+    pack_gqa: Optional[bool] = None,
+    sm_margin: int = 0,
+    return_softmax_lse: bool = False,
+    sinks: Optional[torch.Tensor] = None,
+    ver: int = 4,
+) -> torch.Tensor:
+    assert ver == 4, "only support flash attention v4"
+    q, k, v = [maybe_contiguous(t) for t in (q, k, v)]
+    num_head, head_dim = q.shape[-2:]
+    if cu_seqlens_q is None:
+        batch_size, seqlen_q = q.shape[:2]
+    else:
+        batch_size = cu_seqlens_q.shape[0] - 1
+        seqlen_q = None
+    head_dim_v = v.shape[-1]
+
+    if cu_seqlens_q is not None:
+        assert cu_seqlens_q.shape == (
+            batch_size + 1,
+        ), "cu_seqlens_q must have shape (batch_size + 1,)"
+        assert cu_seqlens_q.dtype == torch.int32, "cu_seqlens_q must be int32"
+        assert cu_seqlens_q.stride(0) == 1, "cu_seqlens_q must be contiguous"
+
+    assert q.dtype in [
+        torch.float16,
+        torch.bfloat16,
+    ], "inputs must be float16 or bfloat16"
+    assert q.dtype == k.dtype == v.dtype, "inputs must have the same dtype"
+    assert head_dim <= 256, "head_dim must be less than or equal to 256"
+    alignment = 16 // q.element_size()
+    assert head_dim_v % alignment == 0, f"head_dim_v must be divisible by {alignment}"
+
+    q_batch_seqlen_shape = (
+        (batch_size, seqlen_q) if cu_seqlens_q is None else (q.shape[0],)
+    )
+    out = q.new_empty(*q_batch_seqlen_shape, num_head, head_dim_v)
+    return out
+
+
+def flash_attn_varlen_func_fake_out_lse(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    max_seqlen_q: Optional[int] = None,
+    max_seqlen_k: Optional[int] = None,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    page_table: Optional[torch.Tensor] = None,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    qv: Optional[torch.Tensor] = None,
+    q_descale: Optional[torch.Tensor] = None,
+    k_descale: Optional[torch.Tensor] = None,
+    v_descale: Optional[torch.Tensor] = None,
+    window_size: Optional[List[int]] = None,
+    attention_chunk: int = 0,
+    softcap: float = 0.0,
+    num_splits: int = 1,
+    pack_gqa: Optional[bool] = None,
+    sm_margin: int = 0,
+    return_softmax_lse: bool = True,
+    sinks: Optional[torch.Tensor] = None,
+    ver: int = 4,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    assert ver == 4, "only support flash attention v4"
+    q, k, v = [maybe_contiguous(t) for t in (q, k, v)]
+    num_head, head_dim = q.shape[-2:]
+    if cu_seqlens_q is None:
+        batch_size, seqlen_q = q.shape[:2]
+        total_q = batch_size * seqlen_q
+    else:
+        batch_size = cu_seqlens_q.shape[0] - 1
+        seqlen_q = None
+        total_q = q.shape[0]
+    head_dim_v = v.shape[-1]
+
+    if cu_seqlens_q is not None:
+        assert cu_seqlens_q.shape == (
+            batch_size + 1,
+        ), "cu_seqlens_q must have shape (batch_size + 1,)"
+        assert cu_seqlens_q.dtype == torch.int32, "cu_seqlens_q must be int32"
+        assert cu_seqlens_q.stride(0) == 1, "cu_seqlens_q must be contiguous"
+
+    assert q.dtype in [
+        torch.float16,
+        torch.bfloat16,
+    ], "inputs must be float16 or bfloat16"
+    assert q.dtype == k.dtype == v.dtype, "inputs must have the same dtype"
+    assert head_dim <= 256, "head_dim must be less than or equal to 256"
+    alignment = 16 // q.element_size()
+    assert head_dim_v % alignment == 0, f"head_dim_v must be divisible by {alignment}"
+
+    q_batch_seqlen_shape = (
+        (batch_size, seqlen_q) if cu_seqlens_q is None else (total_q,)
+    )
+    lse_shape = (
+        (batch_size, num_head, seqlen_q)
+        if cu_seqlens_q is None
+        else (num_head, total_q)
+    )
+
+    out = q.new_empty(*q_batch_seqlen_shape, num_head, head_dim_v)
+    lse = q.new_empty(lse_shape, dtype=torch.float32)
+    return out, lse
+
+
+# -----------------------------
+# Registered custom ops
+# NOTE: fixed return schemas to avoid:
+# "Object of type 'Tensor' is not an instance of 'sequence'"
+# -----------------------------
+@register_custom_op(fake_impl=flash_attn_varlen_func_fake_out)
+def flash_attn_varlen_func_op(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    max_seqlen_q: Optional[int] = None,
+    max_seqlen_k: Optional[int] = None,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    page_table: Optional[torch.Tensor] = None,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    qv: Optional[torch.Tensor] = None,
+    q_descale: Optional[torch.Tensor] = None,
+    k_descale: Optional[torch.Tensor] = None,
+    v_descale: Optional[torch.Tensor] = None,
+    window_size: Optional[List[int]] = None,
+    attention_chunk: int = 0,
+    softcap: float = 0.0,
+    num_splits: int = 1,
+    pack_gqa: Optional[bool] = None,
+    sm_margin: int = 0,
+    return_softmax_lse: bool = False,
+    sinks: Optional[torch.Tensor] = None,
+    ver: int = 4,
+) -> torch.Tensor:
+    if window_size is None:
+        window_size = [-1, -1]
+    if return_softmax_lse:
+        raise ValueError(
+            "flash_attn_varlen_func_op is out-only op; return_softmax_lse must be False. "
+            "Use flash_attn_varlen_func_op_lse for (out, lse)."
+        )
+    return flash_attn_varlen_func(
+        q,
+        k,
+        v,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_q=max_seqlen_q,
+        max_seqlen_k=max_seqlen_k,
+        seqused_q=seqused_q,
+        seqused_k=seqused_k,
+        page_table=page_table,
+        softmax_scale=softmax_scale,
+        causal=causal,
+        qv=qv,
+        q_descale=q_descale,
+        k_descale=k_descale,
+        v_descale=v_descale,
+        window_size=tuple(window_size),
+        attention_chunk=attention_chunk,
+        softcap=softcap,
+        num_splits=num_splits,
+        pack_gqa=pack_gqa,
+        sm_margin=sm_margin,
+        return_softmax_lse=False,
+        sinks=sinks,
+        ver=ver,
+    )
+
+
+@register_custom_op(fake_impl=flash_attn_varlen_func_fake_out_lse)
+def flash_attn_varlen_func_op_lse(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    max_seqlen_q: Optional[int] = None,
+    max_seqlen_k: Optional[int] = None,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    page_table: Optional[torch.Tensor] = None,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    qv: Optional[torch.Tensor] = None,
+    q_descale: Optional[torch.Tensor] = None,
+    k_descale: Optional[torch.Tensor] = None,
+    v_descale: Optional[torch.Tensor] = None,
+    window_size: Optional[List[int]] = None,
+    attention_chunk: int = 0,
+    softcap: float = 0.0,
+    num_splits: int = 1,
+    pack_gqa: Optional[bool] = None,
+    sm_margin: int = 0,
+    return_softmax_lse: bool = True,
+    sinks: Optional[torch.Tensor] = None,
+    ver: int = 4,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    if window_size is None:
+        window_size = [-1, -1]
+    if not return_softmax_lse:
+        raise ValueError(
+            "flash_attn_varlen_func_op_lse is out+lse op; return_softmax_lse must be True. "
+            "Use flash_attn_varlen_func_op for out-only."
+        )
+    return flash_attn_varlen_func(
+        q,
+        k,
+        v,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_q=max_seqlen_q,
+        max_seqlen_k=max_seqlen_k,
+        seqused_q=seqused_q,
+        seqused_k=seqused_k,
+        page_table=page_table,
+        softmax_scale=softmax_scale,
+        causal=causal,
+        qv=qv,
+        q_descale=q_descale,
+        k_descale=k_descale,
+        v_descale=v_descale,
+        window_size=tuple(window_size),
+        attention_chunk=attention_chunk,
+        softcap=softcap,
+        num_splits=num_splits,
+        pack_gqa=pack_gqa,
+        sm_margin=sm_margin,
+        return_softmax_lse=True,
+        sinks=sinks,
+        ver=ver,
+    )
+
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+
+fa_ver = 3
+
+
+def set_fa_ver(ver: int) -> None:
+    global fa_ver
+    fa_ver = ver
+
+
+@dataclass
+class FlashAttentionMetadata:
+    # Sequence lengths for the forward batch
+    # Maximum sequence length for query
+    max_seqlen_q: int = 1
+    # Maximum sequence length for key
+    max_seqlen_k: int = 0
+    # Cumulative sequence lengths for query
+    cu_seqlens_q: torch.Tensor = None
+    # Cumulative sequence lengths for key
+    cu_seqlens_k: torch.Tensor = None
+
+
+class FlashAttentionMetadataBuilder(AttentionMetadataBuilder):
+    def __init__(self) -> None:
+        pass
+
+    def prepare(self) -> None:
+        pass
+
+    def build(  # type: ignore
+        self,
+        raw_latent_shape=list,
+        **kwargs: dict[str, Any],
+    ) -> FlashAttentionMetadata:
+        # TODO: put empty values here to be set at first-run, since the q_len calculation can be complicated
+        return FlashAttentionMetadata(max_seqlen_q=None, max_seqlen_k=None)
+
+
+class FlashAttentionBackend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [32, 64, 96, 128, 160, 192, 224, 256]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.FA
+
+    @staticmethod
+    def get_impl_cls() -> type["FlashAttentionImpl"]:
+        return FlashAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        raise NotImplementedError
+
+    @staticmethod
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
+        return FlashAttentionMetadataBuilder
+
+
+class FlashAttentionImpl(AttentionImpl):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.head_size = head_size
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.attention_metadata = FlashAttentionMetadata()
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata = None,
+        *,
+        return_softmax_lse: bool = False,
+    ):
+        if attn_metadata is not None:
+            if attn_metadata.max_seqlen_q is None:
+                attn_metadata.max_seqlen_q = query.shape[1]
+            if attn_metadata.max_seqlen_k is None:
+                attn_metadata.max_seqlen_k = key.shape[1]
+            max_seqlen_q = attn_metadata.max_seqlen_q
+            max_seqlen_k = attn_metadata.max_seqlen_k
+        else:
+            max_seqlen_q = query.shape[1]
+            max_seqlen_k = key.shape[1]
+
+        # FA version selection:
+        # - fa_ver == 3: call python function (can return Tensor or (Tensor, Tensor) depending on flag)
+        # - fa_ver == 4: call custom ops with FIXED return schema
+        if fa_ver == 3:
+            flash_attn_op = flash_attn_varlen_func
+            output = flash_attn_op(
+                q=query,
+                k=key,
+                v=value,
+                cu_seqlens_q=None,
+                cu_seqlens_k=None,
+                max_seqlen_q=max_seqlen_q,
+                max_seqlen_k=max_seqlen_k,
+                softmax_scale=self.softmax_scale,
+                causal=self.causal,
+                return_softmax_lse=return_softmax_lse,
+                ver=fa_ver,
+            )
+            return output
+
+        if fa_ver == 4:
+            if return_softmax_lse:
+                out_tensor, softmax_lse = flash_attn_varlen_func_op_lse(
+                    q=query,
+                    k=key,
+                    v=value,
+                    cu_seqlens_q=None,
+                    cu_seqlens_k=None,
+                    max_seqlen_q=max_seqlen_q,
+                    max_seqlen_k=max_seqlen_k,
+                    softmax_scale=self.softmax_scale,
+                    causal=self.causal,
+                    return_softmax_lse=True,
+                    ver=fa_ver,
+                )
+                return out_tensor, softmax_lse
+            out_tensor = flash_attn_varlen_func_op(
+                q=query,
+                k=key,
+                v=value,
+                cu_seqlens_q=None,
+                cu_seqlens_k=None,
+                max_seqlen_q=max_seqlen_q,
+                max_seqlen_k=max_seqlen_k,
+                softmax_scale=self.softmax_scale,
+                causal=self.causal,
+                return_softmax_lse=False,
+                ver=fa_ver,
+            )
+            return out_tensor
+
+        raise ValueError(f"flash attention version {fa_ver} is not supported.")

python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn_2.py

@@ -0,0 +1,79 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import (
+    flash_attn_func,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class FlashAttention2Backend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [32, 64, 96, 128, 160, 192, 224, 256]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.FA2
+
+    @staticmethod
+    def get_impl_cls() -> type["FlashAttention2Impl"]:
+        return FlashAttention2Impl
+
+    @staticmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        raise NotImplementedError
+
+    @staticmethod
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
+        raise NotImplementedError
+
+
+class FlashAttention2Impl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ):
+        output = flash_attn_func(
+            q=query,  # type: ignore[no-untyped-call]
+            k=key,
+            v=value,
+            cu_seqlens_q=None,
+            cu_seqlens_k=None,
+            max_seqlen_q=None,
+            max_seqlen_k=None,
+            softmax_scale=self.softmax_scale,
+            causal=self.causal,
+        )
+        return output

python/sglang/multimodal_gen/runtime/layers/attention/backends/laser_attn.py

@@ -0,0 +1,191 @@
+import torch
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+from sglang.multimodal_gen.runtime.layers.attention.backends.sdpa import SDPABackend
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+# Import to use torch.ops.attentions, install package with sgl_kernel_npu
+try:
+    import attentions  # noqa: F401
+except ImportError as e:
+    raise ImportError(
+        (
+            "The required 'attentions' package is not installed."
+            "The package can be installed with sgl_kernel_npu"
+        )
+    ) from e
+
+logger = init_logger(__name__)
+
+
+class LaserAttentionBackend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [32, 64, 96, 128]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.LASER_ATTN
+
+    @staticmethod
+    def get_impl_cls() -> type["LaserAttentionImpl"]:
+        return LaserAttentionImpl
+
+
+class LaserAttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.softmax_scale = softmax_scale
+
+        # After preprocess input layout should be BNSD.
+        self.seqlen_base = 256
+        self.seqlen_index = 2
+        self.dim_index = 3
+        self.dim_base = 128
+        self.max_token = 2**31 - 1
+        self.seq_len_pad_base = 256
+
+        # the laser attention operator has issues with small seq_len
+        self.min_seqlen = 2048
+        self.sdpa_impl = SDPABackend.get_impl_cls()(
+            num_heads,
+            head_size,
+            causal,
+            softmax_scale,
+            num_kv_heads,
+            prefix,
+            **extra_impl_args,
+        )
+
+    def _pad(self, input_tensor: torch.Tensor) -> torch.Tensor:
+        """
+        Pad the input tensor along the sequence length and head dimension.
+        to multiples of base values. self.seqlen_index and self.dim_index should be positive integers.
+        """
+
+        seq_len = input_tensor.size(self.seqlen_index)
+        head_dim = input_tensor.size(self.dim_index)
+
+        pad_seq = 0
+        if seq_len % self.seqlen_base != 0:
+            pad_seq = ((seq_len // self.seqlen_base) + 1) * self.seqlen_base - seq_len
+
+        pad_dim = 0
+        if head_dim % self.dim_base != 0:
+            pad_dim = ((head_dim // self.dim_base) + 1) * self.dim_base - head_dim
+
+        if pad_seq == 0 and pad_dim == 0:
+            return input_tensor
+
+        pad_list = [0] * (2 * input_tensor.ndim)
+
+        pad_list[len(pad_list) - 2 * self.seqlen_index - 1] = pad_seq
+        pad_list[len(pad_list) - 2 * self.dim_index - 1] = pad_dim
+
+        return torch.nn.functional.pad(input_tensor, pad_list)
+
+    def _la_preprocess_input(
+        self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        # Currently BSND input layout is not supported
+        q = query.transpose(1, 2)
+        k = key.transpose(1, 2)
+        v = value.transpose(1, 2)
+
+        if q.dtype != torch.float16:
+            q = q.to(torch.float16)
+            k = k.to(torch.float16)
+            v = v.to(torch.float16)
+
+        q = self._pad(q)
+        k = self._pad(k)
+        v = self._pad(v)
+
+        return q, k, v
+
+    def _la_postprocess_output(
+        self,
+        attention_out: torch.Tensor,
+        dtype: torch.dtype,
+        qseqlen: int,
+        head_dim: int,
+    ) -> torch.Tensor:
+        if dtype != attention_out.dtype:
+            attention_out = attention_out.to(dtype)
+
+        attention_out = attention_out[:, :, :qseqlen, :head_dim]
+        attention_out = attention_out.transpose(1, 2).contiguous()
+        return attention_out
+
+    def _laser_attention(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        head_num: int,
+        pre_tokens: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return torch.ops.attentions.la(
+            query=query,
+            key=key,
+            value=value,
+            atten_mask=None,
+            alibi_mask=None,
+            drop_mask=None,
+            scale_value=self.softmax_scale,
+            head_num=head_num,
+            input_layout="BNSD",
+            keep_prob=1.0,
+            pre_tokens=pre_tokens,
+            next_tokens=1,
+            is_highPrecision=True,
+        )
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        q_seqlen, head_dim = query.shape[1], query.shape[3]
+        kv_seqlen = key.shape[1]
+
+        if q_seqlen < self.min_seqlen or kv_seqlen != q_seqlen:
+            output = self.sdpa_impl.forward(
+                query,
+                key,
+                value,
+                attn_metadata,
+            )
+        else:
+            pre_tokens = self.max_token
+            if kv_seqlen % self.seq_len_pad_base != 0:
+                pre_tokens = (
+                    kv_seqlen // self.seq_len_pad_base + 1
+                ) * self.seq_len_pad_base - kv_seqlen
+
+            q, k, v = self._la_preprocess_input(query, key, value)
+            _, la_output = self._laser_attention(q, k, v, q.shape[1], pre_tokens)
+            output = self._la_postprocess_output(
+                la_output, query.dtype, q_seqlen, head_dim
+            )
+
+        return output

python/sglang/multimodal_gen/runtime/layers/attention/backends/rain_fusion_attn.py

@@ -0,0 +1,414 @@
+import math
+from dataclasses import dataclass
+from typing import Any, List, Optional
+
+import attentions  # noqa: F401
+import torch
+from einops import rearrange
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.layers.attention.backends.laser_attn import (
+    LaserAttentionBackend,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class RainFusionAttentionBackend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [32, 64, 96, 128]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.RAIN_FUSION_ATTN
+
+    @staticmethod
+    def get_impl_cls() -> type["RainFusionAttentionImpl"]:
+        return RainFusionAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["RainFusionAttentionMetadata"]:
+        return RainFusionAttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["RainFusionAttentionMetadataBuilder"]:
+        return RainFusionAttentionMetadataBuilder
+
+
+@dataclass
+class RainFusionAttentionMetadata(AttentionMetadata):
+    current_timestep: int
+    skip_first_steps: int
+    sparsity: float
+    latent_shape: list[int]
+
+
+class RainFusionAttentionMetadataBuilder(AttentionMetadataBuilder):
+    def __init__(self) -> None:
+        pass
+
+    def prepare(self) -> None:
+        pass
+
+    def build(
+        self,
+        current_timestep: int,
+        skip_first_steps: int,
+        sparsity: float,
+        raw_latent_shape: list[int],
+        patch_size: tuple[int, int, int],
+        **kwargs: dict[str, Any],
+    ) -> RainFusionAttentionMetadata:
+        if not (skip_first_steps >= 0 and 0.0 <= sparsity < 1.0):
+            raise ValueError(
+                (
+                    "Invalid attention metadata values."
+                    f"Sparsity should be in [0, 1), skip_first_steps should be non-negative."
+                    f"Got sparsity={sparsity}, skip_first_steps={skip_first_steps}"
+                )
+            )
+
+        if sparsity == 0.0:
+            logger.warning(
+                (
+                    "Sparsity is set to 0.0, which means no tokens will be dropped."
+                    "For better performance use Laser Attention or increase sparsity."
+                )
+            )
+
+        latent_shape = raw_latent_shape[-3:]
+        latent_shape = [latent_shape[i] // patch_size[i] for i in range(3)]
+
+        return RainFusionAttentionMetadata(
+            current_timestep=current_timestep,
+            skip_first_steps=skip_first_steps,
+            sparsity=sparsity,
+            latent_shape=latent_shape,
+        )
+
+
+class RainFusionAttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.block_size = 128
+        self.inner_precise = 0
+
+        self.laser_attn_impl = LaserAttentionBackend.get_impl_cls()(
+            num_heads,
+            head_size,
+            causal,
+            softmax_scale,
+            num_kv_heads,
+            prefix,
+            **extra_impl_args,
+        )
+
+    def _avgpool(
+        self, input_tensor: torch.Tensor, pool_size: int = 128
+    ) -> torch.Tensor:
+        batch, seqlen, heads, dim = input_tensor.shape
+        x = input_tensor.permute(0, 2, 3, 1).reshape(batch * heads, dim, seqlen)
+
+        pooled = torch.nn.functional.avg_pool1d(
+            x, kernel_size=pool_size, stride=pool_size, ceil_mode=True
+        )
+        out = pooled.reshape(batch, heads, dim, -1).permute(0, 3, 1, 2).contiguous()
+
+        return out
+
+    def _get_mask_index(self, mask: torch.Tensor) -> torch.Tensor:
+        batch_size, num_heads, seq_len, _ = mask.shape
+
+        mask_reshaped = mask.reshape(-1, seq_len)
+        row_indices = torch.arange(
+            seq_len, device=mask.device, dtype=torch.float32
+        ).unsqueeze(0)
+
+        sorted_vals = torch.where(mask_reshaped, row_indices, seq_len)
+        sorted_vals, _ = torch.sort(sorted_vals, dim=-1)
+        valid_count = mask_reshaped.sum(dim=-1, keepdim=True)
+        keep_mask = row_indices < valid_count
+        result = torch.where(keep_mask, sorted_vals, -1)
+
+        pos_matrix = result.reshape(batch_size, num_heads, seq_len, seq_len).to(
+            torch.int64
+        )
+        return pos_matrix
+
+    def _get_blockwise_mask(
+        self,
+        qkv_pool: torch.Tensor,
+        sparsity: float,
+        scale: float,
+        pool_size: int,
+        latent_shape: tuple,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        first_frame_len = latent_shape[1] * latent_shape[2]
+
+        query_pool, key_pool, value_pool = torch.chunk(qkv_pool, 3, dim=0)
+        attn_scores = (
+            query_pool.permute(0, 2, 1, 3) @ key_pool.permute(0, 2, 3, 1) * scale
+        )
+
+        keep_len = math.ceil(attn_scores.shape[-1] * (1 - sparsity))
+
+        topk_values, _ = torch.topk(attn_scores, k=keep_len, dim=-1)
+        mask = attn_scores >= topk_values[..., -1:]
+
+        firstframe_block_num = (first_frame_len + pool_size - 1) // pool_size
+        if firstframe_block_num > 0:
+            mask[:, :, :firstframe_block_num, :] = True
+            mask[:, :, :, :firstframe_block_num] = True
+
+        select_idx = self._get_mask_index(mask)
+        select_idx = select_idx[0].transpose(0, 1)
+        select_num_idx = mask[0].transpose(0, 1).sum(dim=-1)
+        return select_idx, select_num_idx
+
+    def _rearrange_with_remaining(
+        self, tensor: torch.Tensor, latent_shape: tuple[int, int, int]
+    ) -> torch.Tensor:
+        """
+        b (f hn hb wn wb) n d -> b (f hn wn hb wb) n d
+        or
+        b n (f hn hb wn wb) d -> b n (f hn wn hb wb) d
+        """
+        tq, hq, wq = latent_shape
+        first_frame_len, frame_num = hq * wq, tq
+
+        b, s, n, d = tensor.shape
+
+        if (hq % 8 != 0) or (wq % 8 != 0):
+            tensor_first = tensor[:, :first_frame_len, :, :]
+            tensor = tensor[:, first_frame_len:, :, :]
+            tensor_hwt = rearrange(
+                tensor, "b (f h w) n d -> b f h w n d", f=frame_num - 1, h=hq, w=wq
+            )
+            if hq % 8 != 0:
+                tensor_hwt, tensor_h_r = torch.split(tensor_hwt, hq - (hq % 8), dim=2)
+                tensor_h_r = tensor_h_r.reshape(b, frame_num - 1, -1, n, d)
+            if wq % 8 != 0:
+                tensor_hwt, tensor_w_r = torch.split(tensor_hwt, wq - (wq % 8), dim=3)
+                tensor_w_r = tensor_w_r.reshape(b, frame_num - 1, -1, n, d)
+            tensor_hwt = rearrange(
+                tensor_hwt,
+                "b f (hn hb) (wn wb) n d -> b f (hn wn hb wb) n d",
+                f=frame_num - 1,
+                hb=8,
+                wb=8,
+                hn=hq // 8,
+                wn=wq // 8,
+            )
+            if hq % 8 != 0:
+                tensor_hwt = torch.cat((tensor_hwt, tensor_h_r), dim=2)
+            if wq % 8 != 0:
+                tensor_hwt = torch.cat((tensor_hwt, tensor_w_r), dim=2)
+            tensor_hwt = tensor_hwt.reshape(b, -1, n, d)
+            tensor_hwt = torch.cat([tensor_first, tensor_hwt], dim=1)
+        else:
+            tensor_hwt = rearrange(
+                tensor,
+                "b (f hn hb wn wb) n d -> b (f hn wn hb wb) n d",
+                f=frame_num,
+                hb=8,
+                wb=8,
+                hn=hq // 8,
+                wn=wq // 8,
+            )
+
+        return tensor_hwt
+
+    def _inv_rearrange_with_remaining(
+        self, tensor: torch.Tensor, latent_shape: tuple[int, int, int]
+    ) -> torch.Tensor:
+        tq, hq, wq = latent_shape
+        first_frame_len, frame_num = hq * wq, tq
+
+        b, s, n, d = tensor.shape
+
+        if (hq % 8 != 0) or (wq % 8 != 0):
+            tensor_first = tensor[:, :first_frame_len, :, :]
+            tensor = tensor[:, first_frame_len:, :, :]
+            tensor_hwt = rearrange(
+                tensor, "b (f h w) n d -> b f h w n d", f=frame_num - 1, h=hq, w=wq
+            )
+            if hq % 8 != 0:
+                tensor_hwt, tensor_h_r = torch.split(tensor_hwt, hq - (hq % 8), dim=2)
+            if wq % 8 != 0:
+                tensor_hwt, tensor_w_r = torch.split(tensor_hwt, wq - (wq % 8), dim=3)
+            tensor_hwt = tensor_hwt.reshape(b, frame_num - 1, -1, n, d)
+            tensor_hwt = rearrange(
+                tensor_hwt,
+                "b f (hn wn hb wb) n d -> b f (hn hb) (wn wb) n d",
+                f=frame_num - 1,
+                hb=8,
+                wb=8,
+                hn=hq // 8,
+                wn=wq // 8,
+            )
+            if wq % 8 != 0:
+                tensor_hwt = torch.cat((tensor_hwt, tensor_w_r), dim=3)
+            if hq % 8 != 0:
+                tensor_hwt = torch.cat((tensor_hwt, tensor_h_r), dim=2)
+            tensor_hwt = tensor_hwt.reshape(b, -1, n, d)
+            tensor_hwt = torch.cat([tensor_first, tensor_hwt], dim=1)
+        else:
+            tensor_hwt = rearrange(
+                tensor,
+                "b (f hn wn hb wb) n h -> b (f hn hb wn wb) n h",
+                f=frame_num,
+                hb=8,
+                wb=8,
+                hn=hq // 8,
+                wn=wq // 8,
+            )
+
+        return tensor_hwt
+
+    def _do_tensor_rearrange_pooling(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        pool_size: int,
+        latent_shape: tuple[int, int, int],
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Tensor block rearrangement + pooling operation
+        """
+        tensor = torch.cat((query, key, value), dim=0)
+
+        tensor = self._rearrange_with_remaining(tensor, latent_shape)
+        tensor_pool = self._avgpool(tensor, pool_size)
+
+        query_, key_, value_ = torch.chunk(tensor, 3, dim=0)
+        return query_, key_, value_, tensor_pool
+
+    def _rain_fusion_attention(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        select_idx: torch.Tensor,
+        select_num_idx: torch.Tensor,
+        blockshape: List[int],
+        scale: float = 1.0,
+        head_num: int = 1,
+        input_layout: str = "TND",
+        actual_seq_lengths=Optional[torch.Tensor],
+        actual_seq_lengths_kv=Optional[torch.Tensor],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return torch.ops.attentions.rainfusionattention(
+            query=query,
+            key=key,
+            value=value,
+            select_idx=select_idx,
+            select_num_idx=select_num_idx,
+            blockshape=blockshape,
+            attn_mask=None,
+            actual_seq_qlen=actual_seq_lengths,
+            actual_seq_kvlen=actual_seq_lengths_kv,
+            block_table=None,
+            q_input_layout=input_layout,
+            kv_input_layout=input_layout,
+            head_num=head_num,
+            mask_type=0,
+            scale=scale,
+            inner_precise=self.inner_precise,
+            block_size=0,
+        )
+
+    def _rain_fusion_sparse_attention(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        latent_shape: tuple[int, int, int],
+        sparsity: float,
+    ):
+        q, k, v, qkv_pool = self._do_tensor_rearrange_pooling(
+            query, key, value, self.block_size, latent_shape
+        )
+
+        select_idx, select_num_idx = self._get_blockwise_mask(
+            qkv_pool,
+            sparsity,
+            self.softmax_scale,
+            self.block_size,
+            latent_shape,
+        )
+
+        batch_size, seqlen_q, head_num, head_dim = q.shape
+        seqlen_kv = k.shape[1]
+
+        layout = "TND"
+        q = q.reshape(-1, head_num, head_dim)
+        k = k.reshape(-1, head_num, head_dim)
+        v = v.reshape(-1, head_num, head_dim)
+
+        actual_seq_lengths = [seqlen_q] * batch_size
+        actual_seq_lengths_kv = [seqlen_kv] * batch_size
+
+        out, _ = self._rain_fusion_attention(
+            q,
+            k,
+            v,
+            scale=self.softmax_scale,
+            head_num=head_num,
+            input_layout=layout,
+            select_idx=select_idx,
+            select_num_idx=select_num_idx,
+            blockshape=[self.block_size, self.block_size],
+            actual_seq_lengths=actual_seq_lengths,
+            actual_seq_lengths_kv=actual_seq_lengths_kv,
+        )
+
+        out = out.reshape(batch_size, seqlen_q, head_num, head_dim)
+        out = self._inv_rearrange_with_remaining(out, latent_shape)
+        return out
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        if attn_metadata.current_timestep < attn_metadata.skip_first_steps:
+            output = self.laser_attn_impl.forward(
+                query,
+                key,
+                value,
+                attn_metadata,
+            )
+        else:
+            output = self._rain_fusion_sparse_attention(
+                query,
+                key,
+                value,
+                attn_metadata.latent_shape,
+                attn_metadata.sparsity,
+            )
+        return output

python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn.py

@@ -0,0 +1,74 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+
+import torch
+from sageattention import sageattn
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (  # FlashAttentionMetadata,
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class SageAttentionBackend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [32, 64, 96, 128, 160, 192, 224, 256]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.SAGE_ATTN
+
+    @staticmethod
+    def get_impl_cls() -> type["SageAttentionImpl"]:
+        return SageAttentionImpl
+
+
+class SageAttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.dropout = extra_impl_args.get("dropout_p", 0.0)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+        *,
+        return_softmax_lse: bool = False,
+    ) -> torch.Tensor:
+        output = sageattn(
+            query,
+            key,
+            value,
+            # since input is (batch_size, seq_len, head_num, head_dim)
+            tensor_layout="NHD",
+            is_causal=self.causal,
+            sm_scale=self.softmax_scale,
+            return_lse=return_softmax_lse,
+        )
+        if return_softmax_lse:
+            output, softmax_lse = output
+            return output, softmax_lse
+        return output

python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn3.py

@@ -0,0 +1,92 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import torch.nn.functional as F
+from sageattn3 import sageattn3_blackwell
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class SageAttention3Backend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 128, 256]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.SAGE_ATTN_3
+
+    @staticmethod
+    def get_impl_cls() -> type["SageAttention3Impl"]:
+        return SageAttention3Impl
+
+    @staticmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        raise NotImplementedError
+
+
+class SageAttention3Impl(AttentionImpl):
+    _warned_gqa_fallback_global: bool = False
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.dropout = extra_impl_args.get("dropout_p", 0.0)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        query = query.transpose(1, 2)
+        key = key.transpose(1, 2)
+        value = value.transpose(1, 2)
+        # SageAttention3's Blackwell kernel assumes MHA (Hq == Hkv). For GQA/MQA
+        # (Hq != Hkv), fall back to torch SDPA which supports GQA.
+        if key.shape[1] != query.shape[1]:
+            if query.shape[1] % key.shape[1] != 0:
+                raise ValueError(
+                    "GQA/MQA requires query heads to be a multiple of KV heads, "
+                    f"got q_heads={query.shape[1]} and kv_heads={key.shape[1]}"
+                )
+            if not type(self)._warned_gqa_fallback_global:
+                logger.warning(
+                    "SageAttention3 does not support GQA/MQA (Hq != Hkv); falling back to torch SDPA."
+                )
+                type(self)._warned_gqa_fallback_global = True
+            output = F.scaled_dot_product_attention(
+                query,
+                key,
+                value,
+                is_causal=self.causal,
+                dropout_p=self.dropout,
+                scale=self.softmax_scale,
+                enable_gqa=True,
+            )
+        else:
+            output = sageattn3_blackwell(query, key, value, is_causal=self.causal)
+        output = output.transpose(1, 2)
+        return output

python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py

@@ -0,0 +1,95 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+from contextlib import nullcontext
+
+import torch
+from torch.nn.attention import SDPBackend, sdpa_kernel
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (  # FlashAttentionMetadata,
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS = [
+    SDPBackend.CUDNN_ATTENTION,
+    SDPBackend.FLASH_ATTENTION,
+    SDPBackend.EFFICIENT_ATTENTION,
+    SDPBackend.MATH,
+]
+
+
+class SDPABackend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [32, 64, 96, 128, 160, 192, 224, 256]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.TORCH_SDPA
+
+    @staticmethod
+    def get_impl_cls() -> type["SDPAImpl"]:
+        return SDPAImpl
+
+    # @staticmethod
+    # def get_metadata_cls() -> Type["AttentionMetadata"]:
+    #     return FlashAttentionMetadata
+
+
+class SDPAImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.dropout = extra_impl_args.get("dropout_p", 0.0)
+        self.allow_cudnn_sdp = bool(extra_impl_args.get("allow_cudnn_sdp", False))
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        # transpose to bs, heads, seq_len, head_dim
+        query = query.transpose(1, 2)
+        key = key.transpose(1, 2)
+        value = value.transpose(1, 2)
+        attn_kwargs = {
+            "attn_mask": None,
+            "dropout_p": self.dropout,
+            "is_causal": self.causal,
+            "scale": self.softmax_scale,
+        }
+        if query.shape[1] != key.shape[1]:
+            attn_kwargs["enable_gqa"] = True
+        sdpa_context = (
+            sdpa_kernel(_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS)
+            if self.allow_cudnn_sdp and query.device.type == "cuda"
+            else nullcontext()
+        )
+        with sdpa_context:
+            output = torch.nn.functional.scaled_dot_product_attention(
+                query, key, value, **attn_kwargs
+            )
+        output = output.transpose(1, 2)
+        return output

python/sglang/multimodal_gen/runtime/layers/attention/backends/sliding_tile_attn.py

@@ -0,0 +1,316 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+import json
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from einops import rearrange
+
+from sglang.multimodal_gen.runtime.distributed import get_sp_group
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.managers.forward_context import (
+    ForwardContext,
+    get_forward_context,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.server_args import get_global_server_args
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.utils import dict_to_3d_list
+
+try:
+    from st_attn import sliding_tile_attention
+
+    st_attn_backend_available = True
+except Exception:
+    st_attn_backend_available = False
+
+logger = init_logger(__name__)
+
+
+class RangeDict(dict):
+
+    def __getitem__(self, item: int) -> str:
+        for key in self.keys():
+            if isinstance(key, tuple):
+                low, high = key
+                if low <= item <= high:
+                    return str(super().__getitem__(key))
+            elif key == item:
+                return str(super().__getitem__(key))
+        raise KeyError(f"seq_len {item} not supported for STA")
+
+
+class SlidingTileAttentionBackend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        # TODO(will-refactor): check this
+        return [32, 64, 96, 128, 160, 192, 224, 256]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.SLIDING_TILE_ATTN
+
+    @staticmethod
+    def get_impl_cls() -> type["SlidingTileAttentionImpl"]:
+        return SlidingTileAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["SlidingTileAttentionMetadata"]:
+        return SlidingTileAttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["SlidingTileAttentionMetadataBuilder"]:
+        return SlidingTileAttentionMetadataBuilder
+
+
+@dataclass
+class SlidingTileAttentionMetadata(AttentionMetadata):
+    current_timestep: int
+    STA_param: list[
+        list[Any]
+    ]  # each timestep with one metadata, shape [num_layers, num_heads]
+
+
+class SlidingTileAttentionMetadataBuilder(AttentionMetadataBuilder):
+
+    def __init__(self):
+        pass
+
+    def prepare(self):
+        pass
+
+    def build(  # type: ignore
+        self,
+        STA_param: list[list[Any]],
+        current_timestep: int,
+        **kwargs: dict[str, Any],
+    ) -> SlidingTileAttentionMetadata:
+        param = STA_param
+        if param is None:
+            return SlidingTileAttentionMetadata(
+                current_timestep=current_timestep, STA_param=[]
+            )
+        return SlidingTileAttentionMetadata(
+            current_timestep=current_timestep, STA_param=param[current_timestep]
+        )
+
+
+class SlidingTileAttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        if not st_attn_backend_available:
+            raise ValueError("st attn not supported")
+        # TODO(will-refactor): for now this is the mask strategy, but maybe we should
+        # have a more general config for STA?
+        mask_strategy_file_path = (
+            get_global_server_args().attention_backend_config.mask_strategy_file_path
+        )
+        if mask_strategy_file_path is None:
+            raise ValueError("SGLANG_DIFFUSION_ATTENTION_CONFIG is not set")
+
+        # TODO(kevin): get mask strategy for different STA modes
+        with open(mask_strategy_file_path) as f:
+            mask_strategy = json.load(f)
+        self.mask_strategy = dict_to_3d_list(mask_strategy)
+
+        self.prefix = prefix
+        sp_group = get_sp_group()
+        self.sp_size = sp_group.world_size
+        # STA config
+        self.STA_base_tile_size = [6, 8, 8]
+        self.dit_seq_shape_mapping = RangeDict(
+            {
+                (115200, 115456): "30x48x80",
+                82944: "36x48x48",
+                69120: "18x48x80",
+            }
+        )
+        self.full_window_mapping = {
+            "30x48x80": [5, 6, 10],
+            "36x48x48": [6, 6, 6],
+            "18x48x80": [3, 6, 10],
+        }
+
+    def tile(self, x: torch.Tensor) -> torch.Tensor:
+        return rearrange(
+            x,
+            "b (n_t ts_t n_h ts_h n_w ts_w) h d -> b (n_t n_h n_w ts_t ts_h ts_w) h d",
+            n_t=self.full_window_size[0],
+            n_h=self.full_window_size[1],
+            n_w=self.full_window_size[2],
+            ts_t=self.STA_base_tile_size[0],
+            ts_h=self.STA_base_tile_size[1],
+            ts_w=self.STA_base_tile_size[2],
+        )
+
+    def untile(self, x: torch.Tensor) -> torch.Tensor:
+        x = rearrange(
+            x,
+            "b (n_t n_h n_w ts_t ts_h ts_w) h d -> b (n_t ts_t n_h ts_h n_w ts_w) h d",
+            n_t=self.full_window_size[0],
+            n_h=self.full_window_size[1],
+            n_w=self.full_window_size[2],
+            ts_t=self.STA_base_tile_size[0],
+            ts_h=self.STA_base_tile_size[1],
+            ts_w=self.STA_base_tile_size[2],
+        )
+        return x
+
+    def preprocess_qkv(
+        self,
+        qkv: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        img_sequence_length = qkv.shape[1]
+        self.dit_seq_shape_str = self.dit_seq_shape_mapping[img_sequence_length]
+        self.full_window_size = self.full_window_mapping[self.dit_seq_shape_str]
+        self.dit_seq_shape_int = list(map(int, self.dit_seq_shape_str.split("x")))
+        self.img_seq_length = (
+            self.dit_seq_shape_int[0]
+            * self.dit_seq_shape_int[1]
+            * self.dit_seq_shape_int[2]
+        )
+        return self.tile(qkv)
+
+    def postprocess_output(
+        self,
+        output: torch.Tensor,
+        attn_metadata: SlidingTileAttentionMetadata,
+    ) -> torch.Tensor:
+        return self.untile(output)
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        attn_metadata: SlidingTileAttentionMetadata,
+    ) -> torch.Tensor:
+        if self.mask_strategy is None:
+            raise ValueError("mask_strategy cannot be None for SlidingTileAttention")
+        if self.mask_strategy[0] is None:
+            raise ValueError("mask_strategy[0] cannot be None for SlidingTileAttention")
+
+        timestep = attn_metadata.current_timestep
+        forward_context: ForwardContext = get_forward_context()
+        forward_batch = forward_context.forward_batch
+        if forward_batch is None:
+            raise ValueError("forward_batch cannot be None")
+        # pattern:'.double_blocks.0.attn.impl' or '.single_blocks.0.attn.impl'
+        layer_idx = int(self.prefix.split(".")[-3])
+        if attn_metadata.STA_param is None or len(attn_metadata.STA_param) <= layer_idx:
+            raise ValueError("Invalid STA_param")
+        STA_param = attn_metadata.STA_param[layer_idx]
+
+        text_length = q.shape[1] - self.img_seq_length
+        has_text = text_length > 0
+
+        query = q.transpose(1, 2).contiguous()
+        key = k.transpose(1, 2).contiguous()
+        value = v.transpose(1, 2).contiguous()
+
+        head_num = query.size(1)
+        sp_group = get_sp_group()
+        current_rank = sp_group.rank_in_group
+        start_head = current_rank * head_num
+
+        # searching or tuning mode
+        if len(STA_param) < head_num * sp_group.world_size:
+            sparse_attn_hidden_states_all = []
+            full_mask_window = STA_param[-1]
+            for window_size in STA_param[:-1]:
+                sparse_hidden_states = sliding_tile_attention(
+                    query,
+                    key,
+                    value,
+                    [window_size] * head_num,
+                    text_length,
+                    has_text,
+                    self.dit_seq_shape_str,
+                ).transpose(1, 2)
+                sparse_attn_hidden_states_all.append(sparse_hidden_states)
+
+            hidden_states = sliding_tile_attention(
+                query,
+                key,
+                value,
+                [full_mask_window] * head_num,
+                text_length,
+                has_text,
+                self.dit_seq_shape_str,
+            ).transpose(1, 2)
+
+            attn_L2_loss = []
+            attn_L1_loss = []
+            # average loss across all heads
+            for sparse_attn_hidden_states in sparse_attn_hidden_states_all:
+                # L2 loss
+                attn_L2_loss_ = (
+                    torch.mean(
+                        (sparse_attn_hidden_states.float() - hidden_states.float())
+                        ** 2,
+                        dim=[0, 1, 3],
+                    )
+                    .cpu()
+                    .numpy()
+                )
+                attn_L2_loss_ = [round(float(x), 6) for x in attn_L2_loss_]
+                attn_L2_loss.append(attn_L2_loss_)
+                # L1 loss
+                attn_L1_loss_ = (
+                    torch.mean(
+                        torch.abs(
+                            sparse_attn_hidden_states.float() - hidden_states.float()
+                        ),
+                        dim=[0, 1, 3],
+                    )
+                    .cpu()
+                    .numpy()
+                )
+                attn_L1_loss_ = [round(float(x), 6) for x in attn_L1_loss_]
+                attn_L1_loss.append(attn_L1_loss_)
+
+            layer_loss_save = {"L2_loss": attn_L2_loss, "L1_loss": attn_L1_loss}
+
+            if forward_batch.is_cfg_negative:
+                if forward_batch.mask_search_final_result_neg is not None:
+                    forward_batch.mask_search_final_result_neg[timestep].append(
+                        layer_loss_save
+                    )
+            else:
+                if forward_batch.mask_search_final_result_pos is not None:
+                    forward_batch.mask_search_final_result_pos[timestep].append(
+                        layer_loss_save
+                    )
+        else:
+            windows = [STA_param[head_idx + start_head] for head_idx in range(head_num)]
+
+            hidden_states = sliding_tile_attention(
+                query,
+                key,
+                value,
+                windows,
+                text_length,
+                has_text,
+                self.dit_seq_shape_str,
+            ).transpose(1, 2)
+
+        return hidden_states

python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_linear_attn.py

@@ -0,0 +1,695 @@
+"""
+Copyright (c) 2025 by SLA team.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+
+This implementation is adapted from: from https://github.com/thu-ml/TurboDiffusion/blob/main/turbodiffusion/SLA/core.py and https://github.com/thu-ml/SLA/blob/main/SageSLA/core.py
+Citation (please cite if you use this code):
+
+@article{zhang2025sla,
+  title={SLA: Beyond Sparsity in Diffusion Transformers via Fine-Tunable Sparse-Linear Attention},
+  author={Jintao Zhang and Haoxu Wang and Kai Jiang and Shuo Yang and Kaiwen Zheng and Haocheng Xi and Ziteng Wang and Hongzhou Zhu and Min Zhao and Ion Stoica and Joseph E. Gonzalez and Jun Zhu and Jianfei Chen},
+  journal={arXiv preprint arXiv:2509.24006},
+  year={2025}
+}
+"""
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import triton
+import triton.language as tl
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+# ==================================SLA Functions===================================
+def get_block_map(q, k, topk_ratio, BLKQ=64, BLKK=64):
+    arg_k = k - torch.mean(
+        k, dim=-2, keepdim=True
+    )  # smooth-k technique in SageAttention
+    pooled_qblocks = mean_pool(q, BLKQ)
+    pooled_kblocks = mean_pool(arg_k, BLKK)
+    pooled_score = pooled_qblocks @ pooled_kblocks.transpose(-1, -2)
+
+    K = pooled_score.shape[-1]
+    topk = min(K, int(topk_ratio * K))
+    lut = torch.topk(pooled_score, topk, dim=-1, sorted=False).indices
+
+    sparse_map = torch.zeros_like(pooled_score, dtype=torch.int8)
+    sparse_map.scatter_(-1, lut, 1)
+    return sparse_map, lut, topk
+
+
+def mean_pool(x, BLK):
+    assert x.is_contiguous()
+
+    B, H, L, D = x.shape
+    L_BLOCKS = (L + BLK - 1) // BLK
+    x_mean = torch.empty((B, H, L_BLOCKS, D), device=x.device, dtype=x.dtype)
+
+    grid = (L_BLOCKS, B * H)
+    compress_kernel[grid](x, x_mean, L, D, BLK)
+    return x_mean
+
+
+@triton.jit
+def compress_kernel(
+    X,
+    XM,
+    L: tl.constexpr,
+    D: tl.constexpr,
+    BLOCK_L: tl.constexpr,
+):
+    idx_l = tl.program_id(0)
+    idx_bh = tl.program_id(1)
+
+    offs_l = idx_l * BLOCK_L + tl.arange(0, BLOCK_L)
+    offs_d = tl.arange(0, D)
+
+    x_offset = idx_bh * L * D
+    xm_offset = idx_bh * ((L + BLOCK_L - 1) // BLOCK_L) * D
+    x = tl.load(
+        X + x_offset + offs_l[:, None] * D + offs_d[None, :], mask=offs_l[:, None] < L
+    )
+
+    nx = min(BLOCK_L, L - idx_l * BLOCK_L)
+    x_mean = tl.sum(x, axis=0, dtype=tl.float32) / nx
+    tl.store(XM + xm_offset + idx_l * D + offs_d, x_mean.to(XM.dtype.element_ty))
+
+
+@triton.jit
+def _attn_fwd(
+    Q,
+    K,
+    V,
+    qk_scale: tl.constexpr,
+    topk: tl.constexpr,
+    LUT,
+    LSE,
+    OS,
+    L: tl.constexpr,
+    M_BLOCKS: tl.constexpr,
+    D: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    idx_m = tl.program_id(0).to(tl.int64)
+    idx_bh = tl.program_id(1).to(tl.int64)
+
+    qkv_offset = idx_bh * L * D
+    lut_offset = (idx_bh * M_BLOCKS + idx_m) * topk
+    lse_offset = idx_bh * L
+    offs_m = idx_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = tl.arange(0, BLOCK_N)
+    offs_d = tl.arange(0, D)
+
+    Q_ptrs = Q + qkv_offset + offs_m[:, None] * D + offs_d[None, :]
+    K_ptrs = K + qkv_offset + offs_n[None, :] * D + offs_d[:, None]
+    V_ptrs = V + qkv_offset + offs_n[:, None] * D + offs_d[None, :]
+    OS_ptrs = OS + qkv_offset + offs_m[:, None] * D + offs_d[None, :]
+    LUT_ptr = LUT + lut_offset
+    LSE_ptrs = LSE + lse_offset + offs_m
+
+    m_i = tl.full([BLOCK_M], -float("inf"), dtype=tl.float32)
+    l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
+    o_s = tl.zeros([BLOCK_M, D], dtype=tl.float32)
+
+    q = tl.load(Q_ptrs, mask=offs_m[:, None] < L)
+    for block_idx in tl.range(topk):
+        idx_n = tl.load(LUT_ptr + block_idx)
+        n_mask = offs_n < L - idx_n * BLOCK_N
+
+        k = tl.load(K_ptrs + idx_n * BLOCK_N * D, mask=n_mask[None, :])
+        qk = tl.dot(q, k) * (qk_scale * 1.4426950408889634)  # = 1 / ln(2)
+        if L - idx_n * BLOCK_N < BLOCK_N:
+            qk = tl.where(n_mask[None, :], qk, float("-inf"))
+
+        v = tl.load(V_ptrs + idx_n * BLOCK_N * D, mask=n_mask[:, None])
+        local_m = tl.max(qk, 1)
+        new_m = tl.maximum(m_i, local_m)
+        qk = qk - new_m[:, None]
+
+        p = tl.math.exp2(qk)
+        l_ij = tl.sum(p, 1)
+        alpha = tl.math.exp2(m_i - new_m)
+        o_s = o_s * alpha[:, None]
+        o_s += tl.dot(p.to(v.dtype), v)
+
+        l_i = l_i * alpha + l_ij
+        m_i = new_m
+
+    o_s = o_s / l_i[:, None]
+    tl.store(OS_ptrs, o_s.to(OS.type.element_ty), mask=offs_m[:, None] < L)
+
+    m_i += tl.math.log2(l_i)
+    tl.store(LSE_ptrs, m_i, mask=offs_m < L)
+
+
+def _get_cuda_arch(device_index: int) -> str:
+    """Get CUDA architecture string for the given device."""
+    major, minor = torch.cuda.get_device_capability(device_index)
+    return f"sm{major}{minor}"
+
+
+# ==================================SLA Class===================================
+class SparseLinearAttentionBackend(AttentionBackend):
+    """Sparse Linear Attention Backend for efficient attention computation."""
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 128]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.SLA_ATTN
+
+    @staticmethod
+    def get_impl_cls() -> type["SparseLinearAttentionImpl"]:
+        return SparseLinearAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["SparseLinearAttentionMetadata"]:
+        return SparseLinearAttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["SparseLinearAttentionMetadataBuilder"]:
+        return SparseLinearAttentionMetadataBuilder
+
+
+@dataclass
+class SparseLinearAttentionMetadata(AttentionMetadata):
+    """Metadata for Sparse Linear Attention computation."""
+
+    # Basic attention parameters
+    current_timestep: int
+
+    # Sparse attention configuration
+    topk_ratio: float = 0.1
+
+
+class SparseLinearAttentionMetadataBuilder(AttentionMetadataBuilder):
+    """Builder for SparseLinearAttentionMetadata."""
+
+    def __init__(self) -> None:
+        pass
+
+    def prepare(self) -> None:
+        pass
+
+    def build(
+        self,
+        current_timestep: int,
+        topk_ratio: float = 0.1,
+        **kwargs: dict[str, Any],
+    ) -> SparseLinearAttentionMetadata:
+        return SparseLinearAttentionMetadata(
+            current_timestep=current_timestep,
+            topk_ratio=topk_ratio,
+        )
+
+
+class SparseLinearAttentionImpl(AttentionImpl, nn.Module):
+    """Implementation of sparse linear attention for the backend."""
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool = False,
+        softmax_scale: float | None = None,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        # SLA-specific parameters - matched to TurboDiffusion defaults
+        topk_ratio: float = 0.1,  # TurboDiffusion uses topk=0.1
+        feature_map: str = "softmax",
+        BLKQ: int = 128,  # TurboDiffusion uses BLKQ=128
+        BLKK: int = 64,  # TurboDiffusion uses BLKK=64
+        use_bf16: bool = True,
+        **extra_impl_args,
+    ) -> None:
+        nn.Module.__init__(self)
+
+        # SLA-specific config
+        self.topk_ratio = topk_ratio
+        self.BLKQ = BLKQ
+        self.BLKK = BLKK
+        self.dtype = torch.bfloat16 if use_bf16 else torch.float16
+
+        # Learnable linear projection for combining sparse + linear attention
+        self.proj_l = nn.Linear(head_size, head_size, dtype=torch.float32)
+
+        # Feature map for linear attention
+        # Type annotation for callables
+        self.feature_map_q: Callable[[torch.Tensor], torch.Tensor]
+        self.feature_map_k: Callable[[torch.Tensor], torch.Tensor]
+        if feature_map == "elu":
+            self.feature_map_q = lambda x: F.elu(x) + 1
+            self.feature_map_k = lambda x: F.elu(x) + 1
+        elif feature_map == "relu":
+            self.feature_map_q = F.relu
+            self.feature_map_k = F.relu
+        elif feature_map == "softmax":
+            self.feature_map_q = lambda x: F.softmax(x, dim=-1)
+            self.feature_map_k = lambda x: F.softmax(x, dim=-1)
+        else:
+            raise ValueError(f"Unknown feature map: {feature_map}")
+
+        self._init_weights()
+
+    def _init_weights(self) -> None:
+        """Initialize projection weights to zero for residual-like behavior."""
+        with torch.no_grad():
+            nn.init.zeros_(self.proj_l.weight)
+            nn.init.zeros_(self.proj_l.bias)  # type: ignore[arg-type]
+
+    def _calc_linear_attention_with_torch(self, q, k, v):
+        kv = torch.matmul(k.transpose(-1, -2), v)
+        k_sum = torch.sum(k, dim=-2, keepdim=True)
+        return torch.matmul(q, kv) / (1e-5 + torch.matmul(q, k_sum.transpose(-1, -2)))
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: SparseLinearAttentionMetadata = None,
+    ) -> torch.Tensor:
+        """Forward pass for sparse linear attention.
+
+        Args:
+            query: query tensor of shape (B, H, L, D)
+            key: key tensor of shape (B, H, L, D)
+            value: value tensor of shape (B, H, L, D)
+            attn_metadata: attention metadata containing configuration
+        Returns:
+            output tensor of shape (B, H, L, D)
+        """
+        dtype = query.dtype
+
+        # Transpose for computation
+        query = query.transpose(1, 2).contiguous()
+        key = key.transpose(1, 2).contiguous()
+        value = value.transpose(1, 2).contiguous()
+
+        # Get sparse attention map
+        sparse_map, lut, real_topk = get_block_map(
+            query, key, topk_ratio=self.topk_ratio, BLKQ=self.BLKQ, BLKK=self.BLKK
+        )
+
+        # Convert to computation dtype
+        query = query.to(self.dtype)
+        key = key.to(self.dtype)
+        value = value.to(self.dtype)
+
+        # Sparse attention computation
+        o_s = _attention.apply(
+            query, key, value, sparse_map, lut, real_topk, self.BLKQ, self.BLKK
+        )
+
+        # Apply feature maps
+        query = self.feature_map_q(query).to(self.dtype)  # c_q
+        key = self.feature_map_k(key).to(self.dtype)  # c_k
+        # Linear attention computation
+        o_l = self._calc_linear_attention_with_torch(query, key, value)
+
+        # Apply projection and combine results
+        with torch.amp.autocast("cuda", dtype=self.dtype):
+            o_l = self.proj_l(o_l)
+
+        # Combine sparse and linear attention
+        output = (o_s + o_l).to(dtype).transpose(1, 2)
+
+        return output
+
+
+class _attention(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, q, k, v, k_block_id, lut, topk, BLOCK_M, BLOCK_N, qk_scale=None):
+        assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous()
+        assert k_block_id.is_contiguous() and lut.is_contiguous()
+
+        # We recommend the following two settings
+        assert BLOCK_M == 64 or BLOCK_M == 128
+        assert BLOCK_N == 64
+
+        B, H, L, D = q.shape
+        if qk_scale is None:
+            qk_scale = D**-0.5
+
+        M_BLOCKS = triton.cdiv(L, BLOCK_M)
+
+        o_s = torch.empty_like(v)
+        lse = torch.empty(q.shape[:-1], device=q.device, dtype=torch.float32)
+
+        grid = (M_BLOCKS, B * H)
+        _attn_fwd[grid](
+            q,
+            k,
+            v,
+            qk_scale,
+            topk,
+            lut,
+            lse,
+            o_s,
+            L,
+            M_BLOCKS,
+            D,
+            BLOCK_M,
+            BLOCK_N,
+            num_warps=4 if q.shape[-1] == 64 else 8,
+            num_stages=3,
+        )
+
+        ctx.save_for_backward(q, k, v, k_block_id, lut, lse, o_s)
+        ctx.qk_scale = qk_scale
+        ctx.topk = topk
+        ctx.BLOCK_M = BLOCK_M
+        ctx.BLOCK_N = BLOCK_N
+        return o_s
+
+
+# ==================================SageSLA Class===================================
+SAGESLA_ENABLED = True
+try:
+    import spas_sage_attn._fused as fused
+    import spas_sage_attn._qattn as qattn
+    from spas_sage_attn.utils import block_map_lut_triton, get_vanilla_qk_quant
+except ImportError:
+    SAGESLA_ENABLED = False
+
+SAGE2PP_ENABLED = True
+try:
+    from spas_sage_attn._qattn import (
+        qk_int8_sv_f8_accum_f16_block_sparse_attn_inst_buf_fuse_v_scale_with_pv_threshold,
+    )
+except ImportError:
+    SAGE2PP_ENABLED = False
+
+
+class SageSparseLinearAttentionBackend(AttentionBackend):
+    """Quantized Sparse-Linear Attention backend using SageAttention kernels."""
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 128]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.SAGE_SLA_ATTN
+
+    @staticmethod
+    def get_impl_cls() -> type["SageSparseLinearAttentionImpl"]:
+        return SageSparseLinearAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["SageSparseLinearAttentionMetadata"]:
+        return SageSparseLinearAttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["SageSparseLinearAttentionMetadataBuilder"]:
+        return SageSparseLinearAttentionMetadataBuilder
+
+
+@dataclass
+class SageSparseLinearAttentionMetadata(AttentionMetadata):
+    """Metadata for Sage Sparse Linear Attention computation."""
+
+    # Basic attention parameters
+    current_timestep: int
+
+    # Sparse attention configuration
+    topk_ratio: float = 0.1
+
+
+class SageSparseLinearAttentionMetadataBuilder(AttentionMetadataBuilder):
+    """Builder for SageSparseLinearAttentionMetadata."""
+
+    def __init__(self) -> None:
+        pass
+
+    def prepare(self) -> None:
+        pass
+
+    def build(
+        self,
+        current_timestep: int,
+        topk_ratio: float = 0.1,
+        **kwargs: dict[str, Any],
+    ) -> SageSparseLinearAttentionMetadata:
+        return SageSparseLinearAttentionMetadata(
+            current_timestep=current_timestep,
+            topk_ratio=topk_ratio,
+        )
+
+
+class SageSparseLinearAttentionImpl(AttentionImpl, nn.Module):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool = False,
+        softmax_scale: float | None = None,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        topk_ratio: float = 0.5,
+        feature_map: str = "softmax",
+        use_bf16: bool = True,
+        **extra_impl_args,
+    ) -> None:
+        nn.Module.__init__(self)
+
+        assert (
+            SAGESLA_ENABLED
+        ), "Install spas_sage_attn(pip install git+https://github.com/thu-ml/SpargeAttn.git --no-build-isolation) first to enable SageSLA."
+
+        self.num_heads = num_heads
+        self.head_size = head_size
+        self.softmax_scale = softmax_scale if softmax_scale else head_size**-0.5
+        self.causal = causal
+        self.prefix = prefix
+
+        self.topk_ratio = topk_ratio
+        self.dtype = torch.bfloat16 if use_bf16 else torch.float16
+
+        # Learnable linear projection for combining sparse + linear attention
+        self.proj_l = nn.Linear(head_size, head_size, dtype=torch.float32)
+
+        # Feature map for linear attention
+        # Type annotation for callables
+        self.feature_map_q: Callable[[torch.Tensor], torch.Tensor]
+        self.feature_map_k: Callable[[torch.Tensor], torch.Tensor]
+        if feature_map == "elu":
+            self.feature_map_q = lambda x: F.elu(x) + 1
+            self.feature_map_k = lambda x: F.elu(x) + 1
+        elif feature_map == "relu":
+            self.feature_map_q = F.relu
+            self.feature_map_k = F.relu
+        elif feature_map == "softmax":
+            self.feature_map_q = lambda x: F.softmax(x, dim=-1)
+            self.feature_map_k = lambda x: F.softmax(x, dim=-1)
+        else:
+            raise ValueError(f"Unknown feature map: {feature_map}")
+
+        self._init_weights()
+
+    def _init_weights(self) -> None:
+        """Initialize projection weights to zero for residual-like behavior."""
+        with torch.no_grad():
+            nn.init.zeros_(self.proj_l.weight)
+            nn.init.zeros_(self.proj_l.bias)  # type: ignore[arg-type]
+
+    def _calc_linear_attention_with_torch(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+    ):
+        kv = torch.matmul(k.transpose(-1, -2), v)
+        k_sum = torch.sum(k, dim=-2, keepdim=True)
+        return torch.matmul(q, kv) / (1e-5 + torch.matmul(q, k_sum.transpose(-1, -2)))
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        """Forward pass for Sage Sparse Linear attention with quantized kernels.
+        Args:
+            query: query tensor of shape (B, L, H, D)
+            key: key tensor of shape (B, L, H, D)
+            value: value tensor of shape (B, L, H, D)
+            attn_metadata: attention metadata containing configuration
+        Returns:
+            output tensor of shape (B, L, H, D)
+        """
+        dtype = query.dtype
+
+        # Transpose from (B, L, H, D) to SLA format (B, H, L, D)
+        q = query.transpose(1, 2).contiguous()
+        k = key.transpose(1, 2).contiguous()
+        v = value.transpose(1, 2).contiguous()
+
+        # Determine block sizes based on GPU architecture
+        arch = _get_cuda_arch(q.device.index)
+
+        if arch == "sm90":
+            BLKQ = 64
+            BLKK = 128
+        else:
+            BLKQ = 128
+            BLKK = 64
+        # Compute block-sparse attention pattern
+        sparse_map, lut, real_topk = get_block_map(
+            q, k, topk_ratio=self.topk_ratio, BLKQ=BLKQ, BLKK=BLKK
+        )
+
+        # Convert to compute dtype
+        q = q.to(self.dtype)
+        k = k.to(self.dtype)
+        v = v.to(self.dtype)
+
+        ########## SPARGE BEGIN ##########
+        km = k.mean(dim=-2, keepdim=True)
+        headdim = q.size(-1)
+        assert headdim in [
+            64,
+            128,
+        ], "headdim should be in [64, 128]. For other headdim, you can use padding and specify the softmax scale."
+
+        # Quantize Q, K to INT8
+        q_int8, q_scale, k_int8, k_scale = get_vanilla_qk_quant(q, k, km, BLKQ, BLKK)
+        lut, valid_block_num = block_map_lut_triton(sparse_map)
+        scale = 1.0 / (headdim**0.5)
+
+        o_s = torch.empty_like(q)
+
+        if arch in ("sm80", "sm86", "sm87"):
+            pvthreshold = torch.full(
+                (q.shape[-3],), 1e6, dtype=torch.float32, device=q.device
+            )
+            v_fp16 = v.to(torch.float16)
+            qattn.qk_int8_sv_f16_accum_f16_block_sparse_attn_inst_buf_with_pv_threshold(
+                q_int8,
+                k_int8,
+                v_fp16,
+                o_s,
+                lut,
+                valid_block_num,
+                pvthreshold,
+                q_scale,
+                k_scale,
+                1,
+                False,
+                1,
+                scale,
+                0,
+            )
+        else:
+            b, h_kv, kv_len, head_dim = v.shape
+            padded_len = (kv_len + 127) // 128 * 128
+            v_transposed_permutted = torch.empty(
+                (b, h_kv, head_dim, padded_len), dtype=v.dtype, device=v.device
+            )
+            fused.transpose_pad_permute_cuda(v, v_transposed_permutted, 1)
+            v_fp8 = torch.empty(
+                v_transposed_permutted.shape, dtype=torch.float8_e4m3fn, device=v.device
+            )
+            v_scale = torch.empty(
+                (b, h_kv, head_dim), dtype=torch.float32, device=v.device
+            )
+            fused.scale_fuse_quant_cuda(
+                v_transposed_permutted, v_fp8, v_scale, kv_len, 2.25, 1
+            )
+
+            if arch == "sm90":
+                qattn.qk_int8_sv_f8_accum_f32_block_sparse_attn_inst_buf_fuse_v_scale_sm90(
+                    q_int8,
+                    k_int8,
+                    v_fp8,
+                    o_s,
+                    lut,
+                    valid_block_num,
+                    q_scale,
+                    k_scale,
+                    v_scale,
+                    1,
+                    False,
+                    1,
+                    scale,
+                )
+            else:
+                pvthreshold = torch.full(
+                    (q.shape[-3],), 1e6, dtype=torch.float32, device=q.device
+                )
+                if SAGE2PP_ENABLED:
+                    qk_int8_sv_f8_accum_f16_block_sparse_attn_inst_buf_fuse_v_scale_with_pv_threshold(
+                        q_int8,
+                        k_int8,
+                        v_fp8,
+                        o_s,
+                        lut,
+                        valid_block_num,
+                        pvthreshold,
+                        q_scale,
+                        k_scale,
+                        v_scale,
+                        1,
+                        False,
+                        1,
+                        scale,
+                        0,
+                    )
+                else:
+                    qattn.qk_int8_sv_f8_accum_f32_block_sparse_attn_inst_buf_fuse_v_scale_with_pv_threshold(
+                        q_int8,
+                        k_int8,
+                        v_fp8,
+                        o_s,
+                        lut,
+                        valid_block_num,
+                        pvthreshold,
+                        q_scale,
+                        k_scale,
+                        v_scale,
+                        1,
+                        False,
+                        1,
+                        scale,
+                        0,
+                    )
+
+        ########## SPARGE END ##########
+
+        # Linear attention with feature maps
+        q_linear = self.feature_map_q(q).to(self.dtype)
+        k_linear = self.feature_map_k(k).to(self.dtype)
+        o_l = self._calc_linear_attention_with_torch(q_linear, k_linear, v)
+
+        # Project linear attention output and combine
+        with torch.amp.autocast("cuda", dtype=self.dtype):
+            o_l = self.proj_l(o_l)
+
+        # Combine sparse and linear outputs
+        output = (o_s + o_l).to(dtype).transpose(1, 2)
+
+        return output

python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_video_gen_2_attn.py

@@ -0,0 +1,562 @@
+"""
+Sparse Video Gen 2 (SAP) attention backend.
+
+This is a baseline integration that wires the backend into the
+attention framework.
+
+Adapted from https://github.com/svg-project/Sparse-VideoGen/blob/main/svg/models/wan/attention.py
+"""
+
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from torch.nn.attention import SDPBackend, sdpa_kernel
+
+try:
+    from svg.kernels.triton.permute import (
+        apply_inverse_permutation_triton,
+        permute_tensor_by_labels_triton,
+    )
+    from svg.kmeans_utils import (
+        batch_kmeans_Euclid,
+        dynamic_block_sparse_fwd_flashinfer,
+        identify_dynamic_map,
+    )
+
+    svg2_available = True
+except ImportError:
+    svg2_available = False
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class SparseVideoGen2AttentionBackend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 128, 256]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.SPARSE_VIDEO_GEN_2_ATTN
+
+    @staticmethod
+    def get_impl_cls() -> type["SparseVideoGen2AttentionImpl"]:
+        return SparseVideoGen2AttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["SparseVideoGen2AttentionMetadata"]:
+        return SparseVideoGen2AttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["SparseVideoGen2AttentionMetadataBuilder"]:
+        return SparseVideoGen2AttentionMetadataBuilder
+
+
+@dataclass
+class Svg2LayerCache:
+    # centroids for kmeans clustering
+    q_centroids: torch.Tensor | None = None
+    k_centroids: torch.Tensor | None = None
+    centroids_initialized: bool = False
+
+
+@dataclass
+class Svg2Cache:
+    layers: dict[int, Svg2LayerCache] = field(default_factory=dict)
+
+    def get_layer(self, layer_idx: int) -> Svg2LayerCache:
+        layer_cache = self.layers.get(layer_idx)
+        if layer_cache is None:
+            layer_cache = Svg2LayerCache()
+            self.layers[layer_idx] = layer_cache
+        return layer_cache
+
+
+@dataclass
+class SparseVideoGen2AttentionMetadata(AttentionMetadata):
+    current_timestep: int
+    num_q_centroids: int
+    num_k_centroids: int
+    top_p_kmeans: float
+    min_kc_ratio: float
+    kmeans_iter_init: int
+    kmeans_iter_step: int
+    zero_step_kmeans_init: bool
+    first_layers_fp: float
+    first_times_fp: float
+    context_length: int
+    num_frame: int
+    frame_size: int
+    cache: Svg2Cache
+    prompt_length: int | None = None
+    max_seqlen_q: int | None = None
+    max_seqlen_k: int | None = None
+
+
+def _require_kwarg(kwargs: dict[str, Any], name: str) -> Any:
+    if name not in kwargs:
+        raise ValueError(
+            f"Missing required argument for SparseVideoGen2Attention: {name}"
+        )
+    return kwargs[name]
+
+
+class SparseVideoGen2AttentionMetadataBuilder(AttentionMetadataBuilder):
+
+    def __init__(self) -> None:
+        pass
+
+    def prepare(self) -> None:
+        pass
+
+    def build(  # type: ignore[override]
+        self,
+        current_timestep: int,
+        raw_latent_shape: tuple[int, ...],
+        patch_size: tuple[int, int, int],
+        cache: Svg2Cache,
+        num_q_centroids: int,
+        num_k_centroids: int,
+        top_p_kmeans: float,
+        min_kc_ratio: float,
+        kmeans_iter_init: int,
+        kmeans_iter_step: int,
+        zero_step_kmeans_init: bool,
+        first_layers_fp: float,
+        first_times_fp: float,
+        context_length: int = 0,
+        prompt_length: int | None = None,
+        **kwargs: dict[str, Any],
+    ) -> SparseVideoGen2AttentionMetadata:
+        raw_shape = tuple(raw_latent_shape)
+        if len(raw_shape) == 5:
+            t, h, w = raw_shape[2:5]
+        elif len(raw_shape) == 3:
+            t, h, w = raw_shape
+        else:
+            raise ValueError(
+                "raw_latent_shape must be (T, H, W) or (B, C, T, H, W) for SAP attention"
+            )
+        pt, ph, pw = patch_size
+        if t % pt != 0 or h % ph != 0 or w % pw != 0:
+            raise ValueError(
+                "raw_latent_shape must be divisible by patch_size for SAP attention"
+            )
+
+        num_frame = t // pt
+        frame_size = (h // ph) * (w // pw)
+
+        return SparseVideoGen2AttentionMetadata(
+            current_timestep=current_timestep,
+            num_q_centroids=num_q_centroids,
+            num_k_centroids=num_k_centroids,
+            top_p_kmeans=top_p_kmeans,
+            min_kc_ratio=min_kc_ratio,
+            kmeans_iter_init=kmeans_iter_init,
+            kmeans_iter_step=kmeans_iter_step,
+            zero_step_kmeans_init=zero_step_kmeans_init,
+            first_layers_fp=first_layers_fp,
+            first_times_fp=first_times_fp,
+            context_length=context_length,
+            prompt_length=prompt_length,
+            num_frame=num_frame,
+            frame_size=frame_size,
+            cache=cache,
+        )
+
+
+class SparseVideoGen2AttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        if causal:
+            raise ValueError(
+                "Sparse Video Gen 2 attention does not support causal attention"
+            )
+        if not svg2_available:
+            raise ImportError(
+                "Sparse Video Gen 2 attention backend requires svg package to be installed"
+                "Please install it by following the instructions at "
+                "https://github.com/svg-project/Sparse-VideoGen"
+            )
+        self.prefix = prefix
+        self.layer_idx = self._get_layer_idx(prefix)
+
+    def _get_layer_idx(self, prefix: str) -> int:
+        parts = prefix.split(".")
+        if len(parts) < 3:
+            raise ValueError(
+                f"Invalid prefix for SparseVideoGen2AttentionImpl: {prefix}"
+            )
+        return int(parts[-3])
+
+    def kmeans_init(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        attn_metadata: SparseVideoGen2AttentionMetadata,
+    ):
+        cfg, num_heads, seq_len, dim = query.size()
+        qlabels, qcentroids, qcluster_sizes, qiter = batch_kmeans_Euclid(
+            query.reshape(cfg * num_heads, seq_len, dim),
+            n_clusters=attn_metadata.num_q_centroids,
+            max_iters=attn_metadata.kmeans_iter_init,
+        )
+        klabels, kcentroids, kcluster_sizes, kiter = batch_kmeans_Euclid(
+            key.reshape(cfg * num_heads, seq_len, dim),
+            n_clusters=attn_metadata.num_k_centroids,
+            max_iters=attn_metadata.kmeans_iter_init,
+        )
+
+        layer_cache = attn_metadata.cache.get_layer(self.layer_idx)
+        layer_cache.q_centroids = qcentroids
+        layer_cache.k_centroids = kcentroids
+
+        return (
+            qlabels,
+            qcentroids,
+            qcluster_sizes,
+            qiter,
+            klabels,
+            kcentroids,
+            kcluster_sizes,
+            kiter,
+        )
+
+    def kmeans_step(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        attn_metadata: SparseVideoGen2AttentionMetadata,
+    ):
+        cfg, num_heads, seq_len, dim = query.size()
+        layer_cache = attn_metadata.cache.get_layer(self.layer_idx)
+        qlabels, qcentroids, qcluster_sizes, qiter = batch_kmeans_Euclid(
+            query.reshape(cfg * num_heads, seq_len, dim),
+            n_clusters=attn_metadata.num_q_centroids,
+            max_iters=attn_metadata.kmeans_iter_step,
+            init_centroids=layer_cache.q_centroids,
+        )
+        klabels, kcentroids, kcluster_sizes, kiter = batch_kmeans_Euclid(
+            key.reshape(cfg * num_heads, seq_len, dim),
+            n_clusters=attn_metadata.num_k_centroids,
+            max_iters=attn_metadata.kmeans_iter_step,
+            init_centroids=layer_cache.k_centroids,
+        )
+
+        layer_cache.q_centroids = qcentroids
+        layer_cache.k_centroids = kcentroids
+
+        return (
+            qlabels,
+            qcentroids,
+            qcluster_sizes,
+            qiter,
+            klabels,
+            kcentroids,
+            kcluster_sizes,
+            kiter,
+        )
+
+    def kmeans_clustering(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        attn_metadata: SparseVideoGen2AttentionMetadata,
+    ):
+        layer_cache = attn_metadata.cache.get_layer(self.layer_idx)
+        if not layer_cache.centroids_initialized:
+            (
+                qlabels,
+                qcentroids,
+                qcluster_sizes,
+                qiter,
+                klabels,
+                kcentroids,
+                kcluster_sizes,
+                kiter,
+            ) = self.kmeans_init(query, key, attn_metadata)
+            layer_cache.centroids_initialized = True
+            logger.debug(
+                "Centroids initialized at layer %s (init iters: %s).",
+                self.layer_idx,
+                attn_metadata.kmeans_iter_init,
+            )
+        else:
+            (
+                qlabels,
+                qcentroids,
+                qcluster_sizes,
+                qiter,
+                klabels,
+                kcentroids,
+                kcluster_sizes,
+                kiter,
+            ) = self.kmeans_step(query, key, attn_metadata)
+
+        return (
+            qlabels,
+            qcentroids,
+            qcluster_sizes,
+            qiter,
+            klabels,
+            kcentroids,
+            kcluster_sizes,
+            kiter,
+        )
+
+    def semantic_aware_permutation(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: SparseVideoGen2AttentionMetadata,
+    ):
+        cfg, num_heads, seq_len, dim = query.size()
+
+        # 1. Kmeans clustering
+        (
+            qlabels,
+            qcentroids,
+            qcluster_sizes,
+            qiter,
+            klabels,
+            kcentroids,
+            kcluster_sizes,
+            kiter,
+        ) = self.kmeans_clustering(query, key, attn_metadata)
+
+        # 2. Identify dynamic map
+        q_cluster_sizes = qcluster_sizes.view(
+            cfg, num_heads, attn_metadata.num_q_centroids
+        )
+        k_cluster_sizes = kcluster_sizes.view(
+            cfg, num_heads, attn_metadata.num_k_centroids
+        )
+
+        dynamic_map = identify_dynamic_map(
+            qcentroids.view(cfg, num_heads, attn_metadata.num_q_centroids, dim),
+            kcentroids.view(cfg, num_heads, attn_metadata.num_k_centroids, dim),
+            q_cluster_sizes,
+            k_cluster_sizes,
+            attn_metadata.top_p_kmeans,
+            attn_metadata.min_kc_ratio,
+        )
+
+        # 3. Permute the query, key, value
+        q_permuted, q_sorted_indices = permute_tensor_by_labels_triton(
+            query, qlabels, dim=2
+        )
+        k_permuted, k_sorted_indices = permute_tensor_by_labels_triton(
+            key, klabels, dim=2
+        )
+        v_permuted, v_sorted_indices = permute_tensor_by_labels_triton(
+            value, klabels, dim=2, sorted_indices=k_sorted_indices
+        )
+
+        return (
+            q_permuted,
+            k_permuted,
+            v_permuted,
+            dynamic_map,
+            q_cluster_sizes,
+            k_cluster_sizes,
+            q_sorted_indices,
+        )
+
+    def _hunyuan_dynamic_map_post_processing(
+        self,
+        q_perm: torch.Tensor,
+        k_perm: torch.Tensor,
+        v_perm: torch.Tensor,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        dyn_map: torch.Tensor,
+        qc_sz_s: torch.Tensor,
+        kc_sz_s: torch.Tensor,
+        q_sorted_indices: torch.Tensor,
+        video_length: int,
+        context_length: int,
+        prompt_length: int,
+        unprompt_length: int,
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+    ]:
+        # Place the permuted video tokens back and keep text tokens at the tail.
+        query[:, :, :-context_length, :] = q_perm
+        key[:, :, :-context_length, :] = k_perm
+        value[:, :, :-context_length, :] = v_perm
+
+        # Add prompt/unprompt clusters to the dynamic map.
+        dyn_map = F.pad(dyn_map, (0, 2, 0, 2), value=0)
+        dyn_map[:, :, -2, :-1] = True
+        dyn_map[:, :, :-1, -2] = True
+        dyn_map[:, :, -1, -1] = True
+
+        qc_sz_s = F.pad(qc_sz_s, (0, 2), value=0)
+        qc_sz_s[:, :, -2] = prompt_length
+        qc_sz_s[:, :, -1] = unprompt_length
+        kc_sz_s = F.pad(kc_sz_s, (0, 2), value=0)
+        kc_sz_s[:, :, -2] = prompt_length
+        kc_sz_s[:, :, -1] = unprompt_length
+
+        q_sorted_indices = F.pad(q_sorted_indices, (0, context_length), value=0)
+        q_sorted_indices[:, video_length:] = torch.arange(
+            video_length,
+            video_length + context_length,
+            device=q_sorted_indices.device,
+        )
+        return query, key, value, dyn_map, qc_sz_s, kc_sz_s, q_sorted_indices
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: SparseVideoGen2AttentionMetadata,
+    ) -> torch.Tensor:
+        torch.backends.cuda.preferred_linalg_library(backend="magma")
+        res = None
+        # bshd -> bhsd
+        query = query.transpose(1, 2).contiguous()
+        key = key.transpose(1, 2).contiguous()
+        value = value.transpose(1, 2).contiguous()
+        batch_size, num_heads, seq_len, dim = query.size()
+
+        context_length, num_frame, frame_size = (
+            attn_metadata.context_length,
+            attn_metadata.num_frame,
+            attn_metadata.frame_size,
+        )
+        prompt_length = attn_metadata.prompt_length
+        if prompt_length is None:
+            prompt_length = context_length
+
+        assert (
+            seq_len == context_length + num_frame * frame_size
+        ), f"Query Shape: {seq_len} is not equivalent to {context_length} + {num_frame} * {frame_size}"
+
+        # Determine if we use Full Attention to calculate
+        full_attention_flag = False
+
+        if self.layer_idx < attn_metadata.first_layers_fp:
+            full_attention_flag = True
+        if attn_metadata.current_timestep > attn_metadata.first_times_fp:
+            full_attention_flag = True
+
+        if full_attention_flag:
+            if attn_metadata.zero_step_kmeans_init:
+                video_length = attn_metadata.num_frame * attn_metadata.frame_size
+                query_video = query[:, :, :video_length, :].contiguous()
+                key_video = key[:, :, :video_length, :].contiguous()
+                self.kmeans_clustering(query_video, key_video, attn_metadata)
+
+            with sdpa_kernel(
+                SDPBackend.CUDNN_ATTENTION
+            ):  # not sure why we need to force cudnn here, but it's faster than flash attention
+                output_hidden_states = torch.nn.functional.scaled_dot_product_attention(
+                    query, key, value, dropout_p=0.0, is_causal=False
+                )
+
+            res = output_hidden_states.reshape(
+                batch_size, num_heads, seq_len, dim
+            ).transpose(1, 2)
+        else:
+            if context_length > 0:
+                video_length = num_frame * frame_size
+                unprompt_length = max(context_length - prompt_length, 0)
+                query_video = query[:, :, :video_length, :].contiguous()
+                key_video = key[:, :, :video_length, :].contiguous()
+                value_video = value[:, :, :video_length, :].contiguous()
+
+                (
+                    q_perm,
+                    k_perm,
+                    v_perm,
+                    dyn_map,
+                    qc_sz_s,
+                    kc_sz_s,
+                    q_sorted_indices,
+                ) = self.semantic_aware_permutation(
+                    query_video, key_video, value_video, attn_metadata
+                )
+                (
+                    q_perm,
+                    k_perm,
+                    v_perm,
+                    dyn_map,
+                    qc_sz_s,
+                    kc_sz_s,
+                    q_sorted_indices,
+                ) = self._hunyuan_dynamic_map_post_processing(
+                    q_perm,
+                    k_perm,
+                    v_perm,
+                    query,
+                    key,
+                    value,
+                    dyn_map,
+                    qc_sz_s,
+                    kc_sz_s,
+                    q_sorted_indices,
+                    video_length,
+                    context_length,
+                    prompt_length,
+                    unprompt_length,
+                )
+            else:
+                (
+                    q_perm,
+                    k_perm,
+                    v_perm,
+                    dyn_map,
+                    qc_sz_s,
+                    kc_sz_s,
+                    q_sorted_indices,
+                ) = self.semantic_aware_permutation(query, key, value, attn_metadata)
+
+            output_permuted = dynamic_block_sparse_fwd_flashinfer(
+                q_perm, k_perm, v_perm, dyn_map, qc_sz_s, kc_sz_s, is_cpu=False
+            )
+
+            attn_output = apply_inverse_permutation_triton(
+                output_permuted, q_sorted_indices, dim=2
+            )
+
+            res = attn_output.reshape(batch_size, num_heads, seq_len, dim).transpose(
+                1, 2
+            )
+
+        torch.backends.cuda.preferred_linalg_library(
+            backend="default"
+        )  # reset to default
+        return res.contiguous()

python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn.py

@@ -0,0 +1,330 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+import functools
+import math
+from dataclasses import dataclass
+
+import torch
+
+try:
+    from vsa import video_sparse_attn
+except ImportError:
+    video_sparse_attn = None
+
+from typing import Any
+
+from sglang.multimodal_gen.runtime.distributed import get_sp_group
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+VSA_TILE_SIZE = (4, 4, 4)
+
+
+@functools.lru_cache(maxsize=10)
+def get_tile_partition_indices(
+    dit_seq_shape: tuple[int, int, int],
+    tile_size: tuple[int, int, int],
+    device: torch.device,
+) -> torch.LongTensor:
+    T, H, W = dit_seq_shape
+    ts, hs, ws = tile_size
+    indices = torch.arange(T * H * W, device=device, dtype=torch.long).reshape(T, H, W)
+    ls = []
+    for t in range(math.ceil(T / ts)):
+        for h in range(math.ceil(H / hs)):
+            for w in range(math.ceil(W / ws)):
+                ls.append(
+                    indices[
+                        t * ts : min(t * ts + ts, T),
+                        h * hs : min(h * hs + hs, H),
+                        w * ws : min(w * ws + ws, W),
+                    ].flatten()
+                )
+    index = torch.cat(ls, dim=0)
+    return index
+
+
+@functools.lru_cache(maxsize=10)
+def get_reverse_tile_partition_indices(
+    dit_seq_shape: tuple[int, int, int],
+    tile_size: tuple[int, int, int],
+    device: torch.device,
+) -> torch.LongTensor:
+    return torch.argsort(get_tile_partition_indices(dit_seq_shape, tile_size, device))
+
+
+@functools.lru_cache(maxsize=10)
+def construct_variable_block_sizes(
+    dit_seq_shape: tuple[int, int, int],
+    num_tiles: tuple[int, int, int],
+    device: torch.device,
+) -> torch.LongTensor:
+    """
+    Compute the number of valid (non‑padded) tokens inside every
+    (ts_t × ts_h × ts_w) tile after padding ‑‑ flattened in the order
+    (t‑tile, h‑tile, w‑tile) that `rearrange` uses.
+
+    Returns
+    -------
+    torch.LongTensor  # shape: [∏ full_window_size]
+    """
+    # unpack
+    t, h, w = dit_seq_shape
+    ts_t, ts_h, ts_w = VSA_TILE_SIZE
+    n_t, n_h, n_w = num_tiles
+
+    def _sizes(dim_len: int, tile: int, n_tiles: int) -> torch.LongTensor:
+        """Vector with the size of each tile along one dimension."""
+        sizes = torch.full((n_tiles,), tile, dtype=torch.int, device=device)
+        # size of last (possibly partial) tile
+        remainder = dim_len - (n_tiles - 1) * tile
+        sizes[-1] = remainder if remainder > 0 else tile
+        return sizes
+
+    t_sizes = _sizes(t, ts_t, n_t)  # [n_t]
+    h_sizes = _sizes(h, ts_h, n_h)  # [n_h]
+    w_sizes = _sizes(w, ts_w, n_w)  # [n_w]
+
+    # broadcast‑multiply to get voxels per tile, then flatten
+    block_sizes = (
+        t_sizes[:, None, None]  # [n_t, 1,   1]
+        * h_sizes[None, :, None]  # [1,   n_h, 1]
+        * w_sizes[None, None, :]  # [1,   1,   n_w]
+    ).reshape(
+        -1
+    )  # [n_t * n_h * n_w]
+
+    return block_sizes
+
+
+@functools.lru_cache(maxsize=10)
+def get_non_pad_index(
+    variable_block_sizes: torch.LongTensor,
+    max_block_size: int,
+):
+    n_win = variable_block_sizes.shape[0]
+    device = variable_block_sizes.device
+    starts_pad = torch.arange(n_win, device=device) * max_block_size
+    index_pad = (
+        starts_pad[:, None] + torch.arange(max_block_size, device=device)[None, :]
+    )
+    index_mask = (
+        torch.arange(max_block_size, device=device)[None, :]
+        < variable_block_sizes[:, None]
+    )
+    return index_pad[index_mask]
+
+
+class VideoSparseAttentionBackend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 128]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.VIDEO_SPARSE_ATTN
+
+    @staticmethod
+    def get_impl_cls() -> type["VideoSparseAttentionImpl"]:
+        return VideoSparseAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["VideoSparseAttentionMetadata"]:
+        return VideoSparseAttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["VideoSparseAttentionMetadataBuilder"]:
+        return VideoSparseAttentionMetadataBuilder
+
+
+@dataclass
+class VideoSparseAttentionMetadata(AttentionMetadata):
+    current_timestep: int
+    dit_seq_shape: list[int]
+    VSA_sparsity: float
+    num_tiles: list[int]
+    total_seq_length: int
+    tile_partition_indices: torch.LongTensor
+    reverse_tile_partition_indices: torch.LongTensor
+    variable_block_sizes: torch.LongTensor
+    non_pad_index: torch.LongTensor
+    untile_combined_index: torch.LongTensor
+    tile_buf: torch.Tensor | None = None
+
+    # adaption for FastWan2.1-T2V-1.3B-Diffusers
+    # Sequence lengths for the forward batch
+    # Maximum sequence length for query
+    max_seqlen_q: int = 1
+    # Maximum sequence length for key
+    max_seqlen_k: int = 0
+
+
+class VideoSparseAttentionMetadataBuilder(AttentionMetadataBuilder):
+
+    def __init__(self):
+        pass
+
+    def prepare(self):
+        pass
+
+    def build(  # type: ignore
+        self,
+        current_timestep: int,
+        raw_latent_shape: tuple[int, int, int],
+        patch_size: tuple[int, int, int],
+        VSA_sparsity: float,
+        device: torch.device,
+        **kwargs: dict[str, Any],
+    ) -> VideoSparseAttentionMetadata:
+        patch_size = patch_size
+        dit_seq_shape = (
+            raw_latent_shape[0] // patch_size[0],
+            raw_latent_shape[1] // patch_size[1],
+            raw_latent_shape[2] // patch_size[2],
+        )
+
+        num_tiles = (
+            math.ceil(dit_seq_shape[0] / VSA_TILE_SIZE[0]),
+            math.ceil(dit_seq_shape[1] / VSA_TILE_SIZE[1]),
+            math.ceil(dit_seq_shape[2] / VSA_TILE_SIZE[2]),
+        )
+        total_seq_length = math.prod(dit_seq_shape)
+
+        tile_partition_indices = get_tile_partition_indices(
+            dit_seq_shape, VSA_TILE_SIZE, device
+        )
+        reverse_tile_partition_indices = get_reverse_tile_partition_indices(
+            dit_seq_shape, VSA_TILE_SIZE, device
+        )
+        variable_block_sizes = construct_variable_block_sizes(
+            dit_seq_shape, num_tiles, device
+        )
+        non_pad_index = get_non_pad_index(
+            variable_block_sizes, math.prod(VSA_TILE_SIZE)
+        )
+        untile_combined_index = non_pad_index[reverse_tile_partition_indices]
+
+        return VideoSparseAttentionMetadata(
+            current_timestep=current_timestep,
+            dit_seq_shape=dit_seq_shape,  # type: ignore
+            VSA_sparsity=VSA_sparsity,  # type: ignore
+            num_tiles=num_tiles,  # type: ignore
+            total_seq_length=total_seq_length,  # type: ignore
+            tile_partition_indices=tile_partition_indices,  # type: ignore
+            reverse_tile_partition_indices=reverse_tile_partition_indices,
+            variable_block_sizes=variable_block_sizes,
+            non_pad_index=non_pad_index,
+            untile_combined_index=untile_combined_index,
+        )
+
+
+class VideoSparseAttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.prefix = prefix
+        sp_group = get_sp_group()
+        self.sp_size = sp_group.world_size
+
+    def tile(
+        self,
+        x: torch.Tensor,
+        attn_metadata: VideoSparseAttentionMetadata,
+    ) -> torch.Tensor:
+        num_tiles = attn_metadata.num_tiles
+        t_padded_size = num_tiles[0] * VSA_TILE_SIZE[0]
+        h_padded_size = num_tiles[1] * VSA_TILE_SIZE[1]
+        w_padded_size = num_tiles[2] * VSA_TILE_SIZE[2]
+        target_shape = (
+            x.shape[0],
+            t_padded_size * h_padded_size * w_padded_size,
+            x.shape[-2],
+            x.shape[-1],
+        )
+
+        buf = attn_metadata.tile_buf
+        if (
+            buf is None
+            or buf.shape != target_shape
+            or buf.dtype != x.dtype
+            or buf.device != x.device
+        ):
+            buf = torch.zeros(target_shape, device=x.device, dtype=x.dtype)
+            attn_metadata.tile_buf = buf
+
+        buf[:, attn_metadata.non_pad_index] = x[:, attn_metadata.tile_partition_indices]
+        return buf
+
+    def untile(
+        self,
+        x: torch.Tensor,
+        untile_combined_index: torch.LongTensor,
+    ) -> torch.Tensor:
+        return x[:, untile_combined_index]
+
+    def preprocess_qkv(
+        self,
+        qkv: torch.Tensor,
+        attn_metadata: VideoSparseAttentionMetadata,
+    ) -> torch.Tensor:
+        return self.tile(qkv, attn_metadata)
+
+    def postprocess_output(
+        self,
+        output: torch.Tensor,
+        attn_metadata: VideoSparseAttentionMetadata,
+    ) -> torch.Tensor:
+        return self.untile(output, attn_metadata.untile_combined_index)
+
+    def forward(  # type: ignore[override]
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        gate_compress: torch.Tensor,
+        attn_metadata: VideoSparseAttentionMetadata,
+    ) -> torch.Tensor:
+        query = query.transpose(1, 2).contiguous()
+        key = key.transpose(1, 2).contiguous()
+        value = value.transpose(1, 2).contiguous()
+        gate_compress = gate_compress.transpose(1, 2).contiguous()
+
+        VSA_sparsity = attn_metadata.VSA_sparsity
+
+        cur_topk = math.ceil(
+            (1 - VSA_sparsity)
+            * (attn_metadata.total_seq_length / math.prod(VSA_TILE_SIZE))
+        )
+
+        if video_sparse_attn is None:
+            raise NotImplementedError("video_sparse_attn is not installed")
+        hidden_states = video_sparse_attn(
+            query,
+            key,
+            value,
+            variable_block_sizes=attn_metadata.variable_block_sizes,
+            topk=cur_topk,
+            block_size=VSA_TILE_SIZE,
+            compress_attn_weight=gate_compress,
+        ).transpose(1, 2)
+
+        return hidden_states

python/sglang/multimodal_gen/runtime/layers/attention/backends/vmoba.py

@@ -0,0 +1,259 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+import re
+from dataclasses import dataclass
+
+import torch
+from einops import rearrange
+from kernel.attn.vmoba_attn.vmoba import (
+    moba_attn_varlen,
+    process_moba_input,
+    process_moba_output,
+)
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class VMOBAAttentionBackend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.VMOBA_ATTN
+
+    @staticmethod
+    def get_impl_cls() -> type["VMOBAAttentionImpl"]:
+        return VMOBAAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["VideoMobaAttentionMetadata"]:
+        return VideoMobaAttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["VideoMobaAttentionMetadataBuilder"]:
+        return VideoMobaAttentionMetadataBuilder
+
+
+@dataclass
+class VideoMobaAttentionMetadata(AttentionMetadata):
+    current_timestep: int
+
+    temporal_chunk_size: int
+    temporal_topk: int
+    spatial_chunk_size: tuple[int, int]
+    spatial_topk: int
+    st_chunk_size: tuple[int, int, int]
+    st_topk: int
+
+    moba_select_mode: str
+    moba_threshold: float
+    moba_threshold_type: str
+    patch_resolution: list[int]
+
+    first_full_step: int = 12
+    first_full_layer: int = 0
+    # temporal_layer -> spatial_layer -> st_layer
+    temporal_layer: int = 1
+    spatial_layer: int = 1
+    st_layer: int = 1
+
+
+def pad_input(hidden_states, indices, batch, seqlen):
+    """
+    Arguments:
+        hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask.
+        indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence.
+        batch: int, batch size for the padded sequence.
+        seqlen: int, maximum sequence length for the padded sequence.
+    Return:
+        hidden_states: (batch, seqlen, ...)
+    """
+    dim = hidden_states.shape[1:]
+    output = torch.zeros(
+        (batch * seqlen), *dim, device=hidden_states.device, dtype=hidden_states.dtype
+    )
+    output[indices] = hidden_states
+    return rearrange(output, "(b s) ... -> b s ...", b=batch)
+
+
+class VideoMobaAttentionMetadataBuilder(AttentionMetadataBuilder):
+
+    def __init__(self):
+        pass
+
+    def prepare(self):
+        pass
+
+    def build(  # type: ignore
+        self,
+        current_timestep: int,
+        raw_latent_shape: tuple[int, int, int],
+        patch_size: tuple[int, int, int],
+        temporal_chunk_size: int,
+        temporal_topk: int,
+        spatial_chunk_size: tuple[int, int],
+        spatial_topk: int,
+        st_chunk_size: tuple[int, int, int],
+        st_topk: int,
+        moba_select_mode: str = "threshold",
+        moba_threshold: float = 0.25,
+        moba_threshold_type: str = "query_head",
+        device: torch.device = None,
+        first_full_layer: int = 0,
+        first_full_step: int = 12,
+        temporal_layer: int = 1,
+        spatial_layer: int = 1,
+        st_layer: int = 1,
+        **kwargs,
+    ) -> VideoMobaAttentionMetadata:
+        if device is None:
+            device = torch.device("cpu")
+        assert (
+            raw_latent_shape[0] % patch_size[0] == 0
+            and raw_latent_shape[1] % patch_size[1] == 0
+            and raw_latent_shape[2] % patch_size[2] == 0
+        ), f"spatial patch_resolution {raw_latent_shape} should be divisible by patch_size {patch_size}"
+        patch_resolution = [
+            t // pt for t, pt in zip(raw_latent_shape, patch_size, strict=False)
+        ]
+
+        return VideoMobaAttentionMetadata(
+            current_timestep=current_timestep,
+            temporal_chunk_size=temporal_chunk_size,
+            temporal_topk=temporal_topk,
+            spatial_chunk_size=spatial_chunk_size,
+            spatial_topk=spatial_topk,
+            st_chunk_size=st_chunk_size,
+            st_topk=st_topk,
+            moba_select_mode=moba_select_mode,
+            moba_threshold=moba_threshold,
+            moba_threshold_type=moba_threshold_type,
+            patch_resolution=patch_resolution,
+            first_full_layer=first_full_layer,
+            first_full_step=first_full_step,
+            temporal_layer=temporal_layer,
+            spatial_layer=spatial_layer,
+            st_layer=st_layer,
+        )
+
+
+class VMOBAAttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads,
+        head_size,
+        softmax_scale,
+        causal=False,
+        num_kv_heads=None,
+        prefix="",
+        **extra_impl_args,
+    ) -> None:
+        self.prefix = prefix
+        self.layer_idx = self._get_layer_idx(prefix)
+
+        self.pad_input = pad_input
+
+    def _get_layer_idx(self, prefix: str) -> int | None:
+        match = re.search(r"blocks\.(\d+)", prefix)
+        if not match:
+            raise ValueError(f"Invalid prefix: {prefix}")
+        return int(match.group(1))
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        """
+        query: [B, L, H, D]
+        key:   [B, L, H, D]
+        value: [B, L, H, D]
+        attn_metadata: AttentionMetadata
+        """
+        batch_size, sequence_length, num_heads, head_dim = query.shape
+
+        # select chunk type according to layer idx:
+        loop_layer_num = (
+            attn_metadata.temporal_layer
+            + attn_metadata.spatial_layer
+            + attn_metadata.st_layer
+        )
+        moba_layer = self.layer_idx - attn_metadata.first_full_layer
+        if moba_layer % loop_layer_num < attn_metadata.temporal_layer:
+            moba_chunk_size = attn_metadata.temporal_chunk_size
+            moba_topk = attn_metadata.temporal_topk
+        elif (
+            moba_layer % loop_layer_num
+            < attn_metadata.temporal_layer + attn_metadata.spatial_layer
+        ):
+            moba_chunk_size = attn_metadata.spatial_chunk_size
+            moba_topk = attn_metadata.spatial_topk
+        elif (
+            moba_layer % loop_layer_num
+            < attn_metadata.temporal_layer
+            + attn_metadata.spatial_layer
+            + attn_metadata.st_layer
+        ):
+            moba_chunk_size = attn_metadata.st_chunk_size
+            moba_topk = attn_metadata.st_topk
+
+        query, chunk_size = process_moba_input(
+            query, attn_metadata.patch_resolution, moba_chunk_size
+        )
+        key, chunk_size = process_moba_input(
+            key, attn_metadata.patch_resolution, moba_chunk_size
+        )
+        value, chunk_size = process_moba_input(
+            value, attn_metadata.patch_resolution, moba_chunk_size
+        )
+        max_seqlen = query.shape[1]
+        indices_q = torch.arange(
+            0, query.shape[0] * query.shape[1], device=query.device
+        )
+        cu_seqlens = torch.arange(
+            0,
+            query.shape[0] * query.shape[1] + 1,
+            query.shape[1],
+            dtype=torch.int32,
+            device=query.device,
+        )
+        query = rearrange(query, "b s ... -> (b s) ...")
+        key = rearrange(key, "b s ... -> (b s) ...")
+        value = rearrange(value, "b s ... -> (b s) ...")
+
+        # current_timestep=attn_metadata.current_timestep
+        hidden_states = moba_attn_varlen(
+            query,
+            key,
+            value,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
+            moba_chunk_size=chunk_size,
+            moba_topk=moba_topk,
+            select_mode=attn_metadata.moba_select_mode,
+            simsum_threshold=attn_metadata.moba_threshold,
+            threshold_type=attn_metadata.moba_threshold_type,
+        )
+        hidden_states = self.pad_input(
+            hidden_states, indices_q, batch_size, sequence_length
+        )
+        hidden_states = process_moba_output(
+            hidden_states, attn_metadata.patch_resolution, moba_chunk_size
+        )
+
+        return hidden_states

python/sglang/multimodal_gen/runtime/layers/attention/backends/xpu_backend.py

@@ -0,0 +1,122 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+from functools import lru_cache
+
+import torch
+
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+
+try:
+    from sgl_kernel.flash_attn import flash_attn_varlen_func
+
+    flash_attn_func = flash_attn_varlen_func
+except ImportError as e:
+    raise e
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import (
+    FlashAttentionMetadataBuilder,
+)
+
+
+class XPUAttentionBackend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 96, 128, 192, 256]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.FA
+
+    @staticmethod
+    def get_impl_cls() -> type["XPUAttentionImpl"]:
+        return XPUAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        """XPU backend does not require special metadata."""
+        return AttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
+        return FlashAttentionMetadataBuilder
+
+
+@lru_cache(maxsize=128)
+def _get_cu_seqlens(device_index: int, bsz: int, seqlen: int) -> torch.Tensor:
+    return torch.arange(
+        0,
+        (bsz + 1) * seqlen,
+        step=seqlen,
+        device=torch.device("xpu", device_index),
+        dtype=torch.int32,
+    )
+
+
+class XPUAttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.head_size = head_size
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata = None,
+        *,
+        return_softmax_lse: bool = False,
+    ):
+        bsz, seqlen_q, nheads_q, d = tuple(query.shape)
+        _, seqlen_k, nheads_k, _ = tuple(key.shape)
+
+        max_seqlen_q = seqlen_q
+        max_seqlen_k = seqlen_k
+
+        q_ = query.contiguous().reshape(bsz * seqlen_q, nheads_q, d)
+        k_ = key.contiguous().reshape(bsz * seqlen_k, nheads_k, d)
+        v_ = value.contiguous().reshape(bsz * seqlen_k, nheads_k, d)
+        cu_q = _get_cu_seqlens(q_.device.index, bsz, seqlen_q)
+        cu_k = _get_cu_seqlens(q_.device.index, bsz, seqlen_k)
+
+        out = flash_attn_func(
+            q=q_,
+            k=k_,
+            v=v_,
+            cu_seqlens_q=cu_q,
+            cu_seqlens_k=cu_k,
+            max_seqlen_q=max_seqlen_q,
+            max_seqlen_k=max_seqlen_k,
+            softmax_scale=self.softmax_scale,
+            causal=self.causal,
+            return_softmax_lse=return_softmax_lse,
+        )
+
+        if return_softmax_lse:
+            out_tensor, softmax_lse = out[:2]
+            result = out_tensor.reshape(bsz, seqlen_q, nheads_q, d)
+            return result, softmax_lse
+
+        result = out.reshape(bsz, seqlen_q, nheads_q, d)
+        return result

python/sglang/multimodal_gen/runtime/layers/attention/selector.py

@@ -0,0 +1,288 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/attention/selector.py
+
+import os
+from collections.abc import Generator
+from contextlib import contextmanager
+from contextvars import ContextVar
+from functools import cache
+from typing import NamedTuple, cast
+
+import torch
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.server_args import get_global_server_args
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.utils import STR_BACKEND_ENV_VAR, resolve_obj_by_qualname
+
+logger = init_logger(__name__)
+
+
+def backend_name_to_enum(backend_name: str) -> AttentionBackendEnum | None:
+    """
+    Convert a string backend name to a _Backend enum value.
+
+    Returns:
+    * _Backend: enum value if backend_name is a valid in-tree type
+    * None: otherwise it's an invalid in-tree type or an out-of-tree platform is
+            loaded.
+    """
+    assert backend_name is not None
+    return (
+        AttentionBackendEnum[backend_name]
+        if backend_name in AttentionBackendEnum.__members__
+        else None
+    )
+
+
+def get_env_variable_attn_backend() -> AttentionBackendEnum | None:
+    """
+    Get the backend override specified by the sglang-diffusion attention
+    backend environment variable, if one is specified.
+
+    Returns:
+
+    * _Backend enum value if an override is specified
+    * None otherwise
+    """
+    backend_name = os.environ.get(STR_BACKEND_ENV_VAR)
+    return None if backend_name is None else backend_name_to_enum(backend_name)
+
+
+# Global state allows a particular choice of backend
+# to be forced, overriding the logic which auto-selects
+# a backend based on system & workload configuration
+# (default behavior if this variable is None)
+#
+# THIS SELECTION TAKES PRECEDENCE OVER THE
+# FASTVIDEO ATTENTION BACKEND ENVIRONMENT VARIABLE
+forced_attn_backend: AttentionBackendEnum | None = None
+
+
+class ComponentAttnBackendContext(NamedTuple):
+    backend: AttentionBackendEnum | None
+    component_name: str | None
+    selected_backends: dict[str, str | None]
+
+
+component_attn_backend_context: ContextVar[ComponentAttnBackendContext | None] = (
+    ContextVar("component_attn_backend_context", default=None)
+)
+
+
+def global_force_attn_backend(attn_backend: AttentionBackendEnum | None) -> None:
+    """
+    Force all attention operations to use a specified backend.
+
+    Passing `None` for the argument re-enables automatic
+    backend selection.,
+
+    Arguments:
+
+    * attn_backend: backend selection (None to revert to auto)
+    """
+    global forced_attn_backend
+    forced_attn_backend = attn_backend
+
+
+def get_global_forced_attn_backend() -> AttentionBackendEnum | None:
+    """
+    Get the currently-forced choice of attention backend,
+    or None if auto-selection is currently enabled.
+    """
+    return forced_attn_backend
+
+
+def get_component_attn_backend_context() -> ComponentAttnBackendContext | None:
+    return component_attn_backend_context.get()
+
+
+def get_component_forced_attn_backend() -> AttentionBackendEnum | None:
+    context = get_component_attn_backend_context()
+    return context.backend if context is not None else None
+
+
+def get_component_attn_backend_name() -> str | None:
+    context = get_component_attn_backend_context()
+    return context.component_name if context is not None else None
+
+
+def _record_component_attn_backend(backend_name: str, reason: str | None) -> bool:
+    context = get_component_attn_backend_context()
+    if context is None or context.component_name is None:
+        return False
+
+    existing_reason = context.selected_backends.get(backend_name)
+    if backend_name not in context.selected_backends or existing_reason is None:
+        context.selected_backends[backend_name] = reason
+    return True
+
+
+def _log_component_attn_backend_summary(
+    context: ComponentAttnBackendContext | None,
+) -> None:
+    if (
+        context is None
+        or context.component_name is None
+        or not context.selected_backends
+    ):
+        return
+
+    backend_parts = []
+    for backend_name, reason in context.selected_backends.items():
+        if reason:
+            backend_parts.append(f"{backend_name} ({reason})")
+        else:
+            backend_parts.append(backend_name)
+
+    logger.info_once(
+        f"Attention backends for {context.component_name}: "
+        f"{', '.join(backend_parts)}"
+    )
+
+
+def get_attn_backend(
+    head_size: int,
+    dtype: torch.dtype,
+    supported_attention_backends: set[AttentionBackendEnum] | None = None,
+    selected_attention_backend: AttentionBackendEnum | None = None,
+) -> type[AttentionBackend]:
+    if supported_attention_backends is None:
+        be_tuple = tuple()
+    else:
+        # Sort the backend names to ensure consistent cache key
+        be_tuple = tuple(
+            sorted(list(supported_attention_backends), key=lambda b: b.name)
+        )
+
+    selected_backend = selected_attention_backend or get_global_forced_attn_backend()
+    if selected_backend is None:
+        selected_backend = get_component_forced_attn_backend()
+    if selected_backend is None:
+        server_args = get_global_server_args()
+        if server_args.attention_backend is not None:
+            try:
+                selected_backend = AttentionBackendEnum[
+                    server_args.attention_backend.upper()
+                ]
+            except KeyError:
+                raise ValueError(
+                    f"Invalid attention backend '{server_args.attention_backend}' specified via command line. "
+                    f"Available options are: {[e.name.lower() for e in AttentionBackendEnum]}"
+                )
+
+    constraint_backend = None
+    if selected_backend is None and len(be_tuple) == 1:
+        constraint_backend = be_tuple[0].name.lower()
+
+    attention_backend_cls = _cached_get_attn_backend(
+        head_size,
+        dtype,
+        be_tuple,
+        selected_backend,
+    )
+
+    backend_name = attention_backend_cls.get_enum().name.lower()
+    reason = "component constraint" if backend_name == constraint_backend else None
+    if not _record_component_attn_backend(backend_name, reason):
+        logger.info_once(f"Using {backend_name} attention backend")
+    return attention_backend_cls
+
+
+@cache
+def _cached_get_attn_backend(
+    head_size: int,
+    dtype: torch.dtype,
+    supported_attention_backends: tuple[AttentionBackendEnum],
+    selected_backend: AttentionBackendEnum | None,
+) -> type[AttentionBackend]:
+    from sglang.multimodal_gen.runtime.platforms import current_platform
+
+    supported_attention_backends = set(supported_attention_backends)
+
+    # get device-specific attn_backend
+    if len(supported_attention_backends) == 0:
+        # all attention backends are allowed
+        pass
+    elif selected_backend is None and len(supported_attention_backends) == 1:
+        selected_backend = next(iter(supported_attention_backends))
+    elif (
+        selected_backend is not None
+        and selected_backend not in supported_attention_backends
+    ):
+        supported_attention_backends_str = [
+            supported_attention_backend.__str__()
+            for supported_attention_backend in supported_attention_backends
+        ]
+        logger.debug(
+            "Selected attention backend: '%s' not in supported attention backends: %s",
+            selected_backend,
+            supported_attention_backends_str,
+        )
+        selected_backend = None
+
+    attention_cls = current_platform.get_attn_backend_cls_str(
+        selected_backend, head_size, dtype
+    )
+    if not attention_cls:
+        raise ValueError(
+            f"Invalid attention backend for {current_platform.device_name}"
+        )
+    return cast(type[AttentionBackend], resolve_obj_by_qualname(attention_cls))
+
+
+@contextmanager
+def component_attn_backend_context_manager(
+    attn_backend: AttentionBackendEnum | None,
+    component_name: str | None = None,
+) -> Generator[None, None, None]:
+    if attn_backend is None and component_name is None:
+        yield
+        return
+
+    token = component_attn_backend_context.set(
+        ComponentAttnBackendContext(attn_backend, component_name, {})
+    )
+    try:
+        yield
+    finally:
+        context = component_attn_backend_context.get()
+        _log_component_attn_backend_summary(context)
+        component_attn_backend_context.reset(token)
+
+
+@contextmanager
+def global_force_attn_backend_context_manager(
+    attn_backend: AttentionBackendEnum,
+) -> Generator[None, None, None]:
+    """
+    Globally force a sglang-diffusion attention backend override within a
+    context manager, reverting the global attention backend
+    override to its prior state upon exiting the context
+    manager.
+
+    Arguments:
+    * attn_backend: attention backend to force
+
+    Returns:
+
+    * Generator
+    """
+
+    # Save the current state of the global backend override (if any)
+    original_value = get_global_forced_attn_backend()
+
+    # Globally force the new backend override
+    global_force_attn_backend(attn_backend)
+
+    # Yield control back to the enclosed code block
+    try:
+        yield
+    finally:
+        # Revert the original global backend override, if any
+        global_force_attn_backend(original_value)

python/sglang/multimodal_gen/runtime/layers/attention/turbo_layer.py

@@ -0,0 +1,325 @@
+# copy and modify from https://github.com/thu-ml/TurboDiffusion/blob/main/turbodiffusion/rcm/utils/a2a_cp.py and https://github.com/thu-ml/TurboDiffusion/blob/main/turbodiffusion/SLA/core.py
+
+from typing import Any, Callable, List, Tuple, Type, Union
+
+import torch
+import torch.distributed as dist
+from einops import rearrange
+from torch import Tensor
+from torch.distributed import ProcessGroup
+from torch.nn import Module
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionImpl,
+)
+from sglang.multimodal_gen.runtime.layers.attention.selector import get_attn_backend
+from sglang.multimodal_gen.runtime.managers.forward_context import (
+    ForwardContext,
+    get_forward_context,
+)
+from sglang.multimodal_gen.runtime.platforms.interface import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.server_args import get_global_server_args
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.utils import get_compute_dtype
+
+logger = init_logger(__name__)
+
+_TURBO_WAN_SPARSE_BACKENDS = {
+    AttentionBackendEnum.SLA_ATTN,
+    AttentionBackendEnum.SAGE_SLA_ATTN,
+}
+
+
+def post_all2all(local_seq_2_local_head, seq_world_size):
+    def post_func(input):
+        # b, s, n, h
+        if local_seq_2_local_head:
+            output = rearrange(input, "w bs seq h d -> bs (w seq) h d")
+        else:
+            output = rearrange(input, "w bs s h d -> bs s (w h) d", w=seq_world_size)
+
+        return output
+
+    return post_func
+
+
+def single_all_to_all(input, local_seq_2_local_head, group, async_op=False):
+    seq_world_size = dist.get_world_size(group)
+
+    # b, s, n, h
+    if local_seq_2_local_head:
+        bs, local_seq_len, num_total_head, head_dim = input.shape
+        assert (
+            num_total_head % seq_world_size == 0
+        ), f"Number of heads ({num_total_head}) must be divisible by the sequence parallel size ({seq_world_size})!"
+        input_t = rearrange(
+            input,
+            "bs seq_len (w h) d -> w bs seq_len h d",
+            w=seq_world_size,
+            h=num_total_head // seq_world_size,
+        ).contiguous()
+        post_all2all_fun = post_all2all(local_seq_2_local_head, seq_world_size)
+    else:
+        bs, global_seq_len, num_local_head, head_dim = input.shape
+        input_t = rearrange(
+            input,
+            "bs (w s) h d -> w bs s h d",
+            w=seq_world_size,
+            s=global_seq_len // seq_world_size,
+        ).contiguous()
+        post_all2all_fun = post_all2all(local_seq_2_local_head, seq_world_size)
+
+    output = torch.empty_like(input_t)
+    dist.all_to_all_single(output, input_t, group=group, async_op=async_op)
+
+    res = post_all2all_fun(output)
+    return res
+
+
+def _attention_backend_from_name(
+    backend_name: str | None,
+) -> AttentionBackendEnum | None:
+    if backend_name is None:
+        return None
+    try:
+        return AttentionBackendEnum[backend_name.upper()]
+    except KeyError:
+        return None
+
+
+def _resolve_turbo_wan_sparse_backend(
+    attention_type: str,
+    requested_attention_backend: str | None = None,
+    supported_attention_backends: set[AttentionBackendEnum] | None = None,
+) -> tuple[AttentionBackendEnum, str | None]:
+    available_backends = _TURBO_WAN_SPARSE_BACKENDS
+    if supported_attention_backends is not None:
+        available_backends = _TURBO_WAN_SPARSE_BACKENDS & supported_attention_backends
+        if not available_backends:
+            available_backends = _TURBO_WAN_SPARSE_BACKENDS
+
+    preferred_backend = (
+        AttentionBackendEnum.SAGE_SLA_ATTN
+        if attention_type == "sagesla"
+        else AttentionBackendEnum.SLA_ATTN
+    )
+    if preferred_backend not in available_backends:
+        preferred_backend = sorted(available_backends, key=lambda b: b.name)[0]
+
+    requested_backend = _attention_backend_from_name(requested_attention_backend)
+
+    if requested_backend in available_backends:
+        return requested_backend, None
+    if requested_attention_backend is None:
+        return preferred_backend, None
+
+    return (
+        preferred_backend,
+        "TurboWan only supports `sla_attn` or `sage_sla_attn`; "
+        f"got attention_backend={requested_attention_backend!r}. "
+        f"Using `{preferred_backend.name.lower()}` from "
+        f"attention_type={attention_type!r}.",
+    )
+
+
+def async_a2a_communicate(
+    a2a_inputs: Union[torch.Tensor, List[torch.Tensor]],
+    cp_size: int,
+    cp_group: ProcessGroup,
+    cp_stream: torch.get_device_module().Stream,
+    local_seq_2_local_head: bool,
+) -> Union[torch.Tensor, List[torch.Tensor]]:
+    """
+    A2A communication for context parallelism. best used in communicate qkv
+    Modified from Nvidia Transformer Engine.
+    """
+    a2a_inputs = [a2a_inputs] if not isinstance(a2a_inputs, list) else a2a_inputs
+    a2a_outputs, a2a_reqs = [None] * len(a2a_inputs), [None] * len(a2a_inputs)
+    a2a_post_fns = [None] * len(a2a_inputs)
+    if local_seq_2_local_head:
+        for i in range(len(a2a_inputs) + 2):
+            if 0 < i < len(a2a_inputs) + 1:
+                a2a_outputs[i - 1] = torch.empty_like(a2a_inputs[i - 1])
+                a2a_reqs[i - 1] = torch.distributed.all_to_all_single(
+                    a2a_outputs[i - 1], a2a_inputs[i - 1], group=cp_group, async_op=True
+                )
+                a2a_post_fns[i - 1] = post_all2all(local_seq_2_local_head, cp_size)
+            if i > 1:
+                with torch.get_device_module().stream(cp_stream):
+                    a2a_reqs[i - 2].wait()
+                    a2a_outputs[i - 2] = a2a_post_fns[i - 2](a2a_outputs[i - 2])
+            if i < len(a2a_inputs):
+                a2a_inputs[i] = rearrange(
+                    a2a_inputs[i], "bs seq_len (w h) d -> w bs seq_len h d", w=cp_size
+                ).contiguous()
+    else:
+        for i in range(len(a2a_inputs) + 2):
+            if 0 < i < len(a2a_inputs) + 1:
+                a2a_outputs[i - 1] = torch.empty_like(a2a_inputs[i - 1])
+                a2a_reqs[i - 1] = torch.distributed.all_to_all_single(
+                    a2a_outputs[i - 1], a2a_inputs[i - 1], group=cp_group, async_op=True
+                )
+                a2a_post_fns[i - 1] = post_all2all(local_seq_2_local_head, cp_size)
+            if i < len(a2a_inputs):
+                a2a_inputs[i] = rearrange(
+                    a2a_inputs[i], "bs (w s) h d -> w bs s h d", w=cp_size
+                ).contiguous()
+            if i > 1:
+                with torch.get_device_module().stream(cp_stream):
+                    a2a_reqs[i - 2].wait()
+                    a2a_outputs[i - 2] = a2a_post_fns[i - 2](a2a_outputs[i - 2])
+    torch.get_device_module().current_stream().wait_stream(cp_stream)
+    return a2a_outputs[0] if len(a2a_inputs) == 1 else a2a_outputs
+
+
+class _SeqAllToAll(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx: Any, group: dist.ProcessGroup, input: Tensor, local_seq_2_local_head: bool
+    ) -> Tensor:
+        ctx.group = group
+        res = single_all_to_all(input, local_seq_2_local_head, group, False)
+        ctx.local_seq_2_local_head = local_seq_2_local_head
+        return res
+
+    @staticmethod
+    def backward(ctx: Any, *grad_output: Tensor) -> Tuple[None, Tensor, None]:
+        return (
+            None,
+            _SeqAllToAll.apply(ctx.group, *grad_output, not ctx.local_seq_2_local_head),
+            None,
+        )
+
+
+class _SeqAllToAllQKV(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx: Any,
+        group: dist.ProcessGroup,
+        q: Tensor,
+        k: Tensor,
+        v: Tensor,
+        cp_size: int,
+        cp_stream: torch.get_device_module().Stream,
+        local_seq_2_local_head: bool,
+    ) -> Tuple[Tensor, Tensor, Tensor]:
+        ctx.group = group
+        ctx.cp_size = cp_size
+        ctx.cp_stream = cp_stream
+        ctx.local_seq_2_local_head = local_seq_2_local_head
+        q, k, v = async_a2a_communicate(
+            [q, k, v], cp_size, group, cp_stream, local_seq_2_local_head
+        )
+        return q, k, v
+
+    @staticmethod
+    def backward(
+        ctx: Any, *grad_output: Tensor
+    ) -> Tuple[None, Tensor, Tensor, Tensor, None, None, None]:
+        q_grad, k_grad, v_grad = _SeqAllToAllQKV.apply(
+            ctx.group,
+            *grad_output,
+            ctx.cp_size,
+            ctx.cp_stream,
+            not ctx.local_seq_2_local_head,
+        )
+        return (None, q_grad, k_grad, v_grad, None, None, None)
+
+
+class DistributedAttention(torch.nn.Module):
+    """Initialization.
+
+    Arguments:
+        local_attention (Module): local attention with q,k,v
+        sequence_process_group (ProcessGroup): sequence parallel process group
+    """
+
+    def __init__(self, local_attention: Union[Module, Callable]) -> None:
+        super(DistributedAttention, self).__init__()
+        self.local_attn = local_attention
+        self.pg = None
+        self.stream = None
+
+    def forward(
+        self, query: Tensor, key: Tensor, value: Tensor, ctx_attn_metadata
+    ) -> Tensor:
+        """forward
+
+        Arguments:
+            query (Tensor): query input to the layer
+            key (Tensor): key input to the layer
+            value (Tensor): value input to the layer
+
+        Returns:
+            * output (Tensor): context output
+        """
+        if self.pg is None:
+            return self.local_attn(query, key, value, ctx_attn_metadata)
+        pg_size = dist.get_world_size(self.pg)
+        if pg_size < 2:
+            return self.local_attn(query, key, value, ctx_attn_metadata)
+
+        query_layer, key_layer, value_layer = _SeqAllToAllQKV.apply(
+            self.pg, query, key, value, pg_size, self.stream, True
+        )
+        context_layer = self.local_attn(
+            query_layer, key_layer, value_layer, ctx_attn_metadata
+        )
+
+        output = _SeqAllToAll.apply(self.pg, context_layer, False)
+        return output
+
+    def set_context_parallel_group(self, group, stream):
+        self.pg = group
+        self.stream = stream
+
+
+class MinimalA2AAttnOp(DistributedAttention):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        attention_type: str,
+        topk: float,
+        supported_attention_backends: set[AttentionBackendEnum] | None = None,
+        prefix: str = "",
+    ):
+        dtype = get_compute_dtype()
+        try:
+            requested_attention_backend = get_global_server_args().attention_backend
+        except ValueError:
+            requested_attention_backend = None
+        selected_attention_backend, warning_message = _resolve_turbo_wan_sparse_backend(
+            attention_type,
+            requested_attention_backend,
+            supported_attention_backends,
+        )
+        if warning_message is not None:
+            logger.warning_once(warning_message)
+
+        attn_backend = get_attn_backend(
+            head_size,
+            dtype,
+            supported_attention_backends={selected_attention_backend},
+            selected_attention_backend=selected_attention_backend,
+        )
+        impl_cls: Type[AttentionImpl] = attn_backend.get_impl_cls()
+        local_attn = impl_cls(
+            num_heads=num_heads,
+            head_size=head_size,
+            topk_ratio=topk,
+            prefix=f"{prefix}.impl",
+        )
+        super(MinimalA2AAttnOp, self).__init__(local_attn)
+
+    def set_context_parallel_group(self, process_group, ranks, stream):
+        del ranks
+        super().set_context_parallel_group(process_group, stream)
+
+    def forward(
+        self, query: Tensor, key: Tensor, value: Tensor, *args: Any, **kwargs
+    ) -> Tensor:
+        forward_context: ForwardContext = get_forward_context()
+        ctx_attn_metadata = forward_context.attn_metadata
+        results = super().forward(query, key, value, ctx_attn_metadata)
+        return rearrange(results, "b ... h l -> b ... (h l)")

python/sglang/multimodal_gen/runtime/layers/custom_op.py

@@ -0,0 +1,115 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/custom_op.py
+
+from collections.abc import Callable
+from typing import Any
+
+import torch.nn as nn
+
+from sglang.kernel_api_logging import debug_kernel_api
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+_is_cuda = current_platform.is_cuda()
+
+
+class CustomOp(nn.Module):
+    """
+    Base class for custom ops.
+    Dispatches the forward method to the appropriate backend.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._forward_method = self.dispatch_forward()
+
+    @debug_kernel_api
+    def forward(self, *args, **kwargs) -> Any:
+        return self._forward_method(*args, **kwargs)
+
+    def forward_native(self, *args, **kwargs) -> Any:
+        """PyTorch-native implementation of the forward method.
+        This method is optional. If implemented, it can be used with compilers
+        such as torch.compile or PyTorch XLA. Also, it can be used for testing
+        purposes.
+        """
+        raise NotImplementedError
+
+    def forward_cuda(self, *args, **kwargs) -> Any:
+        raise NotImplementedError
+
+    def forward_hip(self, *args, **kwargs) -> Any:
+        # ROCm kernels follow the CUDA path by default.
+        return self.forward_cuda(*args, **kwargs)
+
+    def forward_cpu(self, *args, **kwargs) -> Any:
+        # By default, we assume that CPU ops are compatible with CUDA ops.
+        return self.forward_cuda(*args, **kwargs)
+
+    def forward_tpu(self, *args, **kwargs) -> Any:
+        # By default, we assume that TPU ops are compatible with the
+        # PyTorch-native implementation.
+        return self.forward_native(*args, **kwargs)
+
+    def forward_musa(self, *args, **kwargs) -> Any:
+        # MUSA kernels follow the CUDA path by default.
+        return self.forward_cuda(*args, **kwargs)
+
+    def forward_oot(self, *args, **kwargs) -> Any:
+        # By default, we assume that OOT ops are compatible with the
+        # PyTorch-native implementation.
+        return self.forward_native(*args, **kwargs)
+
+    def forward_npu(self, *args, **kwargs) -> Any:
+        # By default, we assume that NPU ops are compatible with the
+        # PyTorch-native implementation.
+        return self.forward_native(*args, **kwargs)
+
+    def dispatch_forward(self) -> Callable:
+        if _is_cuda:
+            return self.forward_cuda
+        elif current_platform.is_hip():
+            return self.forward_hip
+        elif current_platform.is_npu():
+            return self.forward_npu
+        elif current_platform.is_xpu():
+            return self.forward_xpu
+        elif current_platform.is_musa():
+            return self.forward_musa
+        else:
+            return self.forward_native
+
+    @classmethod
+    def enabled(cls) -> bool:
+        # since we are not using Inductor, we always return True
+        return True
+
+    @staticmethod
+    def default_on() -> bool:
+        """
+        On by default if level < CompilationLevel.PIECEWISE
+        Specifying 'all' or 'none' in custom_op takes precedence.
+        """
+        raise NotImplementedError
+
+    # Dictionary of all custom ops (classes, indexed by registered name).
+    # To check if an op with a name is enabled, call .enabled() on the class.
+    # Examples:
+    # - MyOp.enabled()
+    # - op_registry["my_op"].enabled()
+    op_registry: dict[str, type["CustomOp"]] = {}
+
+    # Decorator to register custom ops.
+    @classmethod
+    def register(cls, name: str) -> Callable:
+
+        def decorator(op_cls):
+            assert name not in cls.op_registry, f"Duplicate op name: {name}"
+            op_cls.name = name
+            cls.op_registry[name] = op_cls
+            return op_cls
+
+        return decorator

python/sglang/multimodal_gen/runtime/layers/elementwise.py

@@ -0,0 +1,53 @@
+import torch
+
+from sglang.jit_kernel.diffusion.triton.scale_shift import fuse_scale_shift_kernel
+from sglang.multimodal_gen.runtime.layers.custom_op import CustomOp
+
+
+class MulAdd(CustomOp):
+    """
+    Fuse elementwise mul and add
+    Input: a, b, c, OptionalInt[k]
+    Output: a * (k + b) + c
+    """
+
+    def __init__(self, prefix: str = ""):
+        super().__init__()
+
+    def forward_native(
+        self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, k: int = 0
+    ) -> torch.Tensor:
+        # a.shape: [batch_size, seq_len, inner_dim]
+        if b.dim() == 4:
+            # b.shape: [batch_size, num_frames, 1, inner_dim]
+            num_frames = b.shape[1]
+            frame_seqlen = a.shape[1] // num_frames
+            return c + (
+                a.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (k + b)
+            ).flatten(1, 2)
+        else:
+            # b.shape: [batch_size, 1, inner_dim]
+            return c + a * (k + b)
+
+    def forward_cuda(
+        self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, k: int = 0
+    ):
+        return fuse_scale_shift_kernel(a, b, c, scale_constant=k)
+
+    def forward_xpu(
+        self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, k: int = 0
+    ):
+        return self.forward_native(a, b, c, k=k)
+
+    @torch.compile
+    def forward_musa(
+        self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, k: int = 0
+    ):
+        return self.forward_native(a, b, c, k=k)
+
+    def forward_npu(
+        self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, k: int = 0
+    ):
+        from sgl_kernel_npu.norm.scale_shift import fused_scale_shift
+
+        return fused_scale_shift(a, b, c, scale_constant=k)

python/sglang/multimodal_gen/runtime/layers/fused_scale_shift_gate.py

@@ -0,0 +1,159 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+
+from sglang.multimodal_gen.runtime.layers.custom_op import CustomOp
+from sglang.multimodal_gen.runtime.platforms import current_platform
+
+_is_cuda = current_platform.is_cuda()
+if _is_cuda:
+    from sglang.jit_kernel.diffusion.triton.scale_shift import (
+        fuse_layernorm_scale_shift_gate_select01_kernel,
+        fuse_residual_layernorm_scale_shift_gate_select01_kernel,
+    )
+
+
+@CustomOp.register("fuse_layernorm_scale_shift_gate_select01")
+class FusedLayerNormScaleShiftGateSelect01(CustomOp):
+    """Fused layernorm + scale/shift + gate with binary index selection.
+
+    CUDA path uses a Triton kernel; other platforms fall back to PyTorch ops.
+    """
+
+    def forward_cuda(
+        self,
+        x: torch.Tensor,
+        weight: Optional[torch.Tensor],
+        bias: Optional[torch.Tensor],
+        scale0: torch.Tensor,
+        shift0: torch.Tensor,
+        gate0: torch.Tensor,
+        scale1: torch.Tensor,
+        shift1: torch.Tensor,
+        gate1: torch.Tensor,
+        index: torch.Tensor,
+        eps: float,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if not x.is_contiguous():
+            x = x.contiguous()
+        if not index.is_contiguous():
+            index = index.contiguous()
+        return fuse_layernorm_scale_shift_gate_select01_kernel(
+            x,
+            weight=weight,
+            bias=bias,
+            scale0=scale0.contiguous(),
+            shift0=shift0.contiguous(),
+            gate0=gate0.contiguous(),
+            scale1=scale1.contiguous(),
+            shift1=shift1.contiguous(),
+            gate1=gate1.contiguous(),
+            index=index,
+            eps=eps,
+        )
+
+    def forward_hip(self, *args, **kwargs):
+        return self.forward_native(*args, **kwargs)
+
+    def forward_native(
+        self,
+        x: torch.Tensor,
+        weight: Optional[torch.Tensor],
+        bias: Optional[torch.Tensor],
+        scale0: torch.Tensor,
+        shift0: torch.Tensor,
+        gate0: torch.Tensor,
+        scale1: torch.Tensor,
+        shift1: torch.Tensor,
+        gate1: torch.Tensor,
+        index: torch.Tensor,
+        eps: float,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        idx = index.to(dtype=torch.bool).unsqueeze(-1)
+        shift = torch.where(idx, shift1.unsqueeze(1), shift0.unsqueeze(1))
+        scale = torch.where(idx, scale1.unsqueeze(1), scale0.unsqueeze(1))
+        gate = torch.where(idx, gate1.unsqueeze(1), gate0.unsqueeze(1))
+        x = F.layer_norm(x, (x.shape[-1],), weight=weight, bias=bias, eps=eps)
+        x = x * (1 + scale) + shift
+        return x, gate
+
+
+@CustomOp.register("fuse_residual_layernorm_scale_shift_gate_select01")
+class FusedResidualLayerNormScaleShiftGateSelect01(CustomOp):
+    """Fused residual + layernorm + scale/shift + gate with binary index selection.
+
+    CUDA path uses a Triton kernel; other platforms fall back to PyTorch ops.
+    """
+
+    def forward_cuda(
+        self,
+        x: torch.Tensor,
+        residual: torch.Tensor,
+        residual_gate: torch.Tensor,
+        weight: Optional[torch.Tensor],
+        bias: Optional[torch.Tensor],
+        scale0: torch.Tensor,
+        shift0: torch.Tensor,
+        gate0: torch.Tensor,
+        scale1: torch.Tensor,
+        shift1: torch.Tensor,
+        gate1: torch.Tensor,
+        index: torch.Tensor,
+        eps: float,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        if not x.is_contiguous():
+            x = x.contiguous()
+        if not index.is_contiguous():
+            index = index.contiguous()
+        if not residual.is_contiguous():
+            residual = residual.contiguous()
+        if not residual_gate.is_contiguous():
+            residual_gate = residual_gate.contiguous()
+        return fuse_residual_layernorm_scale_shift_gate_select01_kernel(
+            x,
+            residual=residual,
+            residual_gate=residual_gate,
+            weight=weight,
+            bias=bias,
+            scale0=scale0.contiguous(),
+            shift0=shift0.contiguous(),
+            gate0=gate0.contiguous(),
+            scale1=scale1.contiguous(),
+            shift1=shift1.contiguous(),
+            gate1=gate1.contiguous(),
+            index=index,
+            eps=eps,
+        )
+
+    def forward_hip(self, *args, **kwargs):
+        return self.forward_native(*args, **kwargs)
+
+    def forward_native(
+        self,
+        x: torch.Tensor,
+        residual: torch.Tensor,
+        residual_gate: torch.Tensor,
+        weight: Optional[torch.Tensor],
+        bias: Optional[torch.Tensor],
+        scale0: torch.Tensor,
+        shift0: torch.Tensor,
+        gate0: torch.Tensor,
+        scale1: torch.Tensor,
+        shift1: torch.Tensor,
+        gate1: torch.Tensor,
+        index: torch.Tensor,
+        eps: float,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        idx = index.to(dtype=torch.bool).unsqueeze(-1)
+        shift = torch.where(idx, shift1.unsqueeze(1), shift0.unsqueeze(1))
+        scale = torch.where(idx, scale1.unsqueeze(1), scale0.unsqueeze(1))
+        gate = torch.where(idx, gate1.unsqueeze(1), gate0.unsqueeze(1))
+        residual_out = residual_gate * x + residual
+        x = F.layer_norm(
+            residual_out, (residual_out.shape[-1],), weight=weight, bias=bias, eps=eps
+        )
+        x = x * (1 + scale) + shift
+        return x, residual_out, gate

python/sglang/multimodal_gen/runtime/layers/kvcache/__init__.py

@@ -0,0 +1,13 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from sglang.multimodal_gen.runtime.layers.kvcache.causal_attention_cache import (
+    CausalAttentionKVView,
+    CausalSelfAttentionKVCache,
+    CrossAttentionKVCache,
+)
+
+__all__ = [
+    "CausalAttentionKVView",
+    "CausalSelfAttentionKVCache",
+    "CrossAttentionKVCache",
+]

python/sglang/multimodal_gen/runtime/layers/kvcache/causal_attention_cache.py

@@ -0,0 +1,413 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass(slots=True)
+class CausalAttentionKVView:
+    k: torch.Tensor
+    v: torch.Tensor
+    local_start_index: int
+    local_end_index: int
+    visible_local_end: int
+    visible_global_end: int
+
+
+@dataclass(slots=True)
+class CausalSelfAttentionKVCache:
+    """one transformer block's causal self-attn K/V cache and write cursors"""
+
+    k: torch.Tensor
+    v: torch.Tensor
+    # the right bound of the valid global token range
+    # e.g., 12000 means [0, 12000) has been generated and cached
+    global_end_index: torch.Tensor
+    # the right bound of the valid local token range within the buffer (when cache is unfilled)
+    local_end_index: torch.Tensor
+    global_end_index_int: int | None = None
+    local_end_index_int: int | None = None
+    cache_size: int = 0
+    sink_tokens: int = 0
+    attention_window_size: int = 0
+    allow_growth: bool = False
+
+    def __post_init__(self) -> None:
+        if self.cache_size == 0:
+            self.cache_size = self.k.shape[1]
+        if self.attention_window_size == 0:
+            self.attention_window_size = self.cache_size
+
+    def reset_indices(self) -> None:
+        self.global_end_index.zero_()
+        self.local_end_index.zero_()
+        if self.global_end_index_int is not None:
+            self.global_end_index_int = 0
+        if self.local_end_index_int is not None:
+            self.local_end_index_int = 0
+
+    def _read_indices(self) -> tuple[int, int]:
+        global_end_index = self.global_end_index_int
+        local_end_index = self.local_end_index_int
+        if global_end_index is None or local_end_index is None:
+            global_end_index = int(self.global_end_index.item())
+            local_end_index = int(self.local_end_index.item())
+            self.global_end_index_int = global_end_index
+            self.local_end_index_int = local_end_index
+        return global_end_index, local_end_index
+
+    def _write_indices(self, *, global_end_index: int, local_end_index: int) -> None:
+        if (
+            self.global_end_index_int == global_end_index
+            and self.local_end_index_int == local_end_index
+        ):
+            return
+        if self.global_end_index_int is not None:
+            self.global_end_index_int = global_end_index
+        if self.local_end_index_int is not None:
+            self.local_end_index_int = local_end_index
+        self.global_end_index.fill_(global_end_index)
+        self.local_end_index.fill_(local_end_index)
+
+    def _grow_to_fit(self, required_tokens: int) -> None:
+        if required_tokens <= self.cache_size:
+            return
+        old_cache_size = self.cache_size
+        new_cache_size = max(required_tokens, old_cache_size * 2)
+
+        new_k = self.k.new_zeros(
+            self.k.shape[0],
+            new_cache_size,
+            self.k.shape[2],
+            self.k.shape[3],
+        )
+        new_v = self.v.new_zeros(
+            self.v.shape[0],
+            new_cache_size,
+            self.v.shape[2],
+            self.v.shape[3],
+        )
+        new_k[:, :old_cache_size] = self.k
+        new_v[:, :old_cache_size] = self.v
+        self.k = new_k
+        self.v = new_v
+        self.cache_size = new_cache_size
+        if self.attention_window_size == old_cache_size:
+            self.attention_window_size = new_cache_size
+
+    def can_direct_current_attention(self, num_new_tokens: int) -> bool:
+        return (
+            self.sink_tokens == 0
+            and self.cache_size == num_new_tokens
+            and self.attention_window_size == num_new_tokens
+        )
+
+    def update_and_get_attention_kv(
+        self,
+        *,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        current_chunk_start: int,
+        cache_head_start: int | None = None,
+        recent_window_tokens: int | None = None,
+        debug_name: str = "causal KV cache",
+    ) -> CausalAttentionKVView:
+        """write fresh kv into the cache, returns the part of view visible to the current chunk
+
+        Args:
+            current_chunk_start: the global position of the start of the chunk
+            cache_head_start: first cache head for key/value when they only
+                carry a local slice of the cache heads; other heads are left untouched
+            recent_window_tokens: recent-window attention size. ``None``
+                returns the full visible attention window. ``0`` keeps only sink
+                tokens plus the current chunk. A positive value keeps sink tokens,
+                up to that many tokens before the current chunk, and the current
+                chunk. Negative values are invalid.
+
+        """
+        num_new_tokens = key.shape[1]
+        num_input_heads = key.shape[2]
+        num_cache_heads = self.k.shape[2]
+        cache_head_slice = None
+        if num_cache_heads != num_input_heads:
+            if cache_head_start is None:
+                raise ValueError(
+                    f"{debug_name} requires cache_head_start when cache heads "
+                    f"({num_cache_heads}) differ from input heads ({num_input_heads})."
+                )
+            cache_head_slice = slice(
+                cache_head_start, cache_head_start + num_input_heads
+            )
+        current_chunk_end = current_chunk_start + num_new_tokens
+        kv_cache_size = self.cache_size
+        sink_tokens = self.sink_tokens
+        global_end_index, local_end_index_prev = self._read_indices()
+
+        # local_start(/end)_index: the local position of the start/end of current chunk
+        # updated_local_end: the updated local end
+        # updated_global_end: the updated global end
+
+        # the global position of the start of the buffer
+        window_start = global_end_index - local_end_index_prev
+
+        if current_chunk_end <= global_end_index:
+            # the window stays as previous
+            # cache layout:
+            # [sink tokens, recent window tokens, current chunk tokens, uninitialized tokens (optional)]
+            local_start_index = current_chunk_start - window_start
+            local_end_index = local_start_index + num_new_tokens
+
+            # the local end and global end remains unchanged (since the chunk hasn't proceed)
+            updated_local_end = local_end_index_prev
+            updated_global_end = global_end_index
+        else:
+            # the chunk window has proceed, append new tokens, and evict earliest (if have to)
+            appended_tokens = current_chunk_end - global_end_index
+            if self.allow_growth:
+                self._grow_to_fit(local_end_index_prev + appended_tokens)
+                kv_cache_size = self.cache_size
+            if local_end_index_prev + appended_tokens > kv_cache_size:
+                # the new tokens can't fit in the remaining space (after local_end_index_prev), start evicting:
+                # before:
+                # [sink tokens, evicted tokens, rolled tokens, remaining space]
+                #                                            ^ end of previous chunk
+                # after:
+                # [sink tokens, rolled tokens,           remaining space      ]
+
+                # 1. keep sink tokens ([0: sink_tokens]) untouched
+                # 2. evict obsolete tokens in: [sink_tokens:sink_tokens + num_evicted_tokens]
+                num_evicted_tokens = (
+                    local_end_index_prev + appended_tokens - kv_cache_size
+                )
+
+                # number of tokens to move
+                num_rolled_tokens = max(
+                    0,
+                    local_end_index_prev - num_evicted_tokens - sink_tokens,
+                )
+                if num_rolled_tokens > 0:
+                    if cache_head_slice is None:
+                        self.k[:, sink_tokens : sink_tokens + num_rolled_tokens] = (
+                            self.k[
+                                :,
+                                sink_tokens
+                                + num_evicted_tokens : sink_tokens
+                                + num_evicted_tokens
+                                + num_rolled_tokens,
+                            ].clone()
+                        )
+                        self.v[:, sink_tokens : sink_tokens + num_rolled_tokens] = (
+                            self.v[
+                                :,
+                                sink_tokens
+                                + num_evicted_tokens : sink_tokens
+                                + num_evicted_tokens
+                                + num_rolled_tokens,
+                            ].clone()
+                        )
+                    else:
+                        self.k[
+                            :,
+                            sink_tokens : sink_tokens + num_rolled_tokens,
+                            cache_head_slice,
+                            :,
+                        ] = self.k[
+                            :,
+                            sink_tokens
+                            + num_evicted_tokens : sink_tokens
+                            + num_evicted_tokens
+                            + num_rolled_tokens,
+                            cache_head_slice,
+                            :,
+                        ].clone()
+                        self.v[
+                            :,
+                            sink_tokens : sink_tokens + num_rolled_tokens,
+                            cache_head_slice,
+                            :,
+                        ] = self.v[
+                            :,
+                            sink_tokens
+                            + num_evicted_tokens : sink_tokens
+                            + num_evicted_tokens
+                            + num_rolled_tokens,
+                            cache_head_slice,
+                            :,
+                        ].clone()
+
+                # if we move the minimum number of tokens, the right bound of the append token would be aligned with end of the buffer
+                local_end_index = kv_cache_size
+            else:
+                # enough space, directly append new tokens after end of previous chunk
+                local_end_index = local_end_index_prev + appended_tokens
+            local_start_index = local_end_index - num_new_tokens
+            updated_local_end = local_end_index
+            # after filling in the proceeded new chunk, the global end aligns with the global end of the current chunk
+            updated_global_end = current_chunk_end
+
+        if (
+            local_start_index < 0
+            or local_end_index > kv_cache_size
+            or local_end_index - local_start_index != num_new_tokens
+        ):
+            raise RuntimeError(
+                f"Invalid {debug_name} write range: "
+                f"local=[{local_start_index}, {local_end_index}), "
+                f"global_end={global_end_index}, "
+                f"prev_local_end={local_end_index_prev}, "
+                f"kv_cache_size={kv_cache_size}, "
+                f"num_new_tokens={num_new_tokens}, "
+                f"current_start={current_chunk_start}, current_end={current_chunk_end}"
+            )
+
+        if self.k.requires_grad:
+            self.k = self.k.detach()
+        if self.v.requires_grad:
+            self.v = self.v.detach()
+        attn_start_index = max(0, updated_local_end - self.attention_window_size)
+
+        # write fresh kv and return visible view
+        if cache_head_slice is None:
+            self.k[:, local_start_index:local_end_index] = key
+            self.v[:, local_start_index:local_end_index] = value
+            visible_k, visible_v = self._visible_attention_kv(
+                local_start_index=local_start_index,
+                updated_local_end=updated_local_end,
+                attn_start_index=attn_start_index,
+                recent_window_tokens=recent_window_tokens,
+            )
+        else:
+            self.k[:, local_start_index:local_end_index, cache_head_slice, :] = key
+            self.v[:, local_start_index:local_end_index, cache_head_slice, :] = value
+            visible_k, visible_v = self._visible_attention_kv(
+                local_start_index=local_start_index,
+                updated_local_end=updated_local_end,
+                attn_start_index=attn_start_index,
+                recent_window_tokens=recent_window_tokens,
+                cache_head_slice=cache_head_slice,
+            )
+
+        self._write_indices(
+            global_end_index=updated_global_end,
+            local_end_index=updated_local_end,
+        )
+        return CausalAttentionKVView(
+            k=visible_k,
+            v=visible_v,
+            local_start_index=local_start_index,
+            local_end_index=local_end_index,
+            visible_local_end=updated_local_end,
+            visible_global_end=updated_global_end,
+        )
+
+    def _visible_attention_kv(
+        self,
+        *,
+        local_start_index: int,
+        updated_local_end: int,
+        attn_start_index: int,
+        recent_window_tokens: int | None,
+        cache_head_slice: slice | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Return the visible KV slice for the current attention call.
+
+        When ``recent_window_tokens`` is ``None``, the returned token range is
+        the standard sliding window::
+
+            [attn_start_index, updated_local_end)
+
+        When recent-window selection is enabled, ``recent_window_tokens`` must be
+        non-negative and the returned token ranges are::
+
+            sink_end = min(self.sink_tokens, updated_local_end)
+            recent_start = max(sink_end, local_start_index - recent_window_tokens)
+            [0, sink_end) + [recent_start, updated_local_end)
+
+        Thus ``0`` keeps only sink tokens plus the current chunk.
+        ``cache_head_slice`` applies the same token ranges to a subset of KV
+        heads.
+        """
+        if recent_window_tokens is None:
+            if cache_head_slice is None:
+                return (
+                    self.k[:, attn_start_index:updated_local_end],
+                    self.v[:, attn_start_index:updated_local_end],
+                )
+            return (
+                self.k[:, attn_start_index:updated_local_end, cache_head_slice, :],
+                self.v[:, attn_start_index:updated_local_end, cache_head_slice, :],
+            )
+        if recent_window_tokens < 0:
+            raise ValueError("recent_window_tokens must be non-negative or None")
+
+        sink_end = min(self.sink_tokens, updated_local_end)
+        recent_start = max(sink_end, local_start_index - recent_window_tokens)
+        if recent_start <= sink_end:
+            if cache_head_slice is None:
+                return self.k[:, :updated_local_end], self.v[:, :updated_local_end]
+            return (
+                self.k[:, :updated_local_end, cache_head_slice, :],
+                self.v[:, :updated_local_end, cache_head_slice, :],
+            )
+        if sink_end <= 0:
+            if cache_head_slice is None:
+                return (
+                    self.k[:, recent_start:updated_local_end],
+                    self.v[:, recent_start:updated_local_end],
+                )
+            return (
+                self.k[:, recent_start:updated_local_end, cache_head_slice, :],
+                self.v[:, recent_start:updated_local_end, cache_head_slice, :],
+            )
+
+        if cache_head_slice is None:
+            return (
+                torch.cat(
+                    [
+                        self.k[:, :sink_end],
+                        self.k[:, recent_start:updated_local_end],
+                    ],
+                    dim=1,
+                ),
+                torch.cat(
+                    [
+                        self.v[:, :sink_end],
+                        self.v[:, recent_start:updated_local_end],
+                    ],
+                    dim=1,
+                ),
+            )
+        return (
+            torch.cat(
+                [
+                    self.k[:, :sink_end, cache_head_slice, :],
+                    self.k[:, recent_start:updated_local_end, cache_head_slice, :],
+                ],
+                dim=1,
+            ),
+            torch.cat(
+                [
+                    self.v[:, :sink_end, cache_head_slice, :],
+                    self.v[:, recent_start:updated_local_end, cache_head_slice, :],
+                ],
+                dim=1,
+            ),
+        )
+
+
+@dataclass(slots=True)
+class CrossAttentionKVCache:
+    """one transformer block's cross-attn condition K/V cache"""
+
+    k: torch.Tensor
+    v: torch.Tensor
+    is_init: bool = False
+
+    def store(self, k: torch.Tensor, v: torch.Tensor) -> None:
+        self.k = k.detach()
+        self.v = v.detach()
+        self.is_init = True
+
+    def reset(self) -> None:
+        self.is_init = False

python/sglang/multimodal_gen/runtime/layers/lora/linear.py

@@ -0,0 +1,730 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# Code adapted from SGLang https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/lora/layers.py
+import os
+
+import torch
+from torch import nn
+from torch.distributed._composable.fsdp import (
+    CPUOffloadPolicy,
+    OffloadPolicy,
+    fully_shard,
+)
+from torch.distributed.tensor import DTensor
+
+from sglang.multimodal_gen.runtime.distributed import (
+    get_local_torch_device,
+    get_tp_rank,
+    split_tensor_along_last_dim,
+    tensor_model_parallel_all_gather,
+    tensor_model_parallel_all_reduce,
+)
+from sglang.multimodal_gen.runtime.layers.linear import (
+    ColumnParallelLinear,
+    LinearBase,
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    ReplicatedLinear,
+    RowParallelLinear,
+)
+from sglang.multimodal_gen.runtime.layers.vocab_parallel_embedding import (
+    VocabParallelEmbedding,
+)
+from sglang.multimodal_gen.utils import get_mixed_precision_state
+
+torch._dynamo.config.recompile_limit = 64
+
+
+LORA_MERGE_CHUNK_BYTES = 32 * 1024 * 1024
+LoRAWeightEntry = tuple[
+    torch.nn.Parameter,
+    torch.nn.Parameter,
+    str | None,
+    float,
+    int | None,
+    int | None,
+]
+
+
+class BaseLayerWithLoRA(nn.Module):
+    def __init__(
+        self,
+        base_layer: nn.Module,
+        lora_rank: int | None = None,
+        lora_alpha: int | None = None,
+    ):
+        super().__init__()
+        self.base_layer: nn.Module = base_layer
+
+        self.merged: bool = False
+        # Immutable base-weight snapshot; `to("cpu")` may alias CPU storage.
+        # Use `clone()` so merge updates cannot mutate this backup tensor.
+        self.cpu_weight = base_layer.weight.detach().to("cpu").clone()
+        # indicates adapter weights don't contain this layer
+        # (which shouldn't normally happen, but we want to separate it from the case of erroneous merging)
+        # Default to True to prevent using uninitialized weights; set to False when weights are loaded
+        self.disable_lora: bool = True
+        self.lora_rank = lora_rank
+        self.lora_alpha = lora_alpha
+        self.lora_weights_list: list[LoRAWeightEntry] = []
+        self.lora_path: str | None = None
+        self.strength: float = 1.0
+
+        self.lora_A = None
+        self.lora_B = None
+
+    @property
+    def weight(self):
+        return self.base_layer.weight
+
+    @property
+    def bias(self):
+        return getattr(self.base_layer, "bias", None)
+
+    @torch.compile()
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        lora_A = self.lora_A
+        lora_B = self.lora_B
+        if isinstance(self.lora_B, DTensor):
+            lora_B = self.lora_B.to_local()
+            lora_A = self.lora_A.to_local()
+
+        # TODO: Support multiple LoRA adapters when use not merged mode
+        if not self.merged and not self.disable_lora:
+            lora_dtype = lora_A.dtype
+            x_lora = x.to(dtype=lora_dtype)
+            lora_A_sliced = self.slice_lora_a_weights(
+                lora_A.to(device=x.device, non_blocking=True)
+            )
+            lora_B_sliced = self.slice_lora_b_weights(
+                lora_B.to(device=x.device, non_blocking=True)
+            )
+            delta = x_lora @ lora_A_sliced.T @ lora_B_sliced.T
+            if self.lora_alpha != self.lora_rank:
+                delta = delta * (
+                    self.lora_alpha / self.lora_rank  # type: ignore
+                )  # type: ignore
+            delta = delta * self.strength
+            out, output_bias = self.base_layer(x)
+            return out + delta.to(dtype=out.dtype), output_bias
+        else:
+            out, output_bias = self.base_layer(x)
+            return out, output_bias
+
+    def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor:
+        return A
+
+    def slice_lora_b_weights(self, B: torch.Tensor) -> torch.Tensor:
+        return B
+
+    @staticmethod
+    def _as_mutable_tensor(tensor: torch.Tensor) -> torch.Tensor:
+        # lora can be reconfigured after executor forwards create inference tensors
+        if tensor.is_inference():
+            with torch.inference_mode(False):
+                return tensor.detach().clone()
+        return tensor
+
+    def set_lora_weights(
+        self,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        lora_path: str | None = None,
+        strength: float = 1.0,
+        clear_existing: bool = False,
+        merge_weights: bool = True,
+    ) -> None:
+        """
+        Set LoRA weights. Supports multiple LoRA adapters.
+
+        Args:
+            A: LoRA A weight tensor
+            B: LoRA B weight tensor
+            lora_path: Path to the LoRA adapter (for logging)
+            strength: LoRA strength
+            clear_existing: If True, clear existing LoRA weights before adding new one.
+                          If False, append to existing list (for multi-LoRA support).
+        """
+        lora_A_param = torch.nn.Parameter(
+            A
+        )  # share storage with weights in the pipeline
+        lora_B_param = torch.nn.Parameter(B)
+
+        if clear_existing:
+            self.lora_weights_list.clear()
+            # Also clear backward compatibility attributes
+            self.lora_A = None
+            self.lora_B = None
+            self.lora_path = None
+            self.strength = 1.0
+
+        # Add to list for multi-LoRA support
+        self.lora_weights_list.append(
+            (
+                lora_A_param,
+                lora_B_param,
+                lora_path,
+                strength,
+                self.lora_rank,
+                self.lora_alpha,
+            )
+        )
+
+        # Set backward compatibility attributes to point to the last LoRA (for single LoRA case)
+        # This ensures backward compatibility while supporting multiple LoRA
+        self.lora_A = lora_A_param
+        self.lora_B = lora_B_param
+        self.lora_path = lora_path
+        self.strength = strength
+
+        self.disable_lora = False
+        if merge_weights:
+            self.merge_lora_weights()
+        elif self.merged:
+            self.unmerge_lora_weights()
+
+    @torch.no_grad()
+    def _merge_lora_into_data(
+        self,
+        data: torch.Tensor,
+        lora_list: list[LoRAWeightEntry],
+    ) -> None:
+        """
+        Merge all LoRA adapters into the data tensor in-place.
+
+        Args:
+            data: The base weight tensor to merge LoRA into (modified in-place)
+            lora_list: List of (lora_A, lora_B, lora_path, lora_strength, rank, alpha) tuples
+        """
+        # Merge all LoRA adapters in order
+        for lora_A, lora_B, _, lora_strength, lora_rank, lora_alpha in lora_list:
+            lora_A_sliced = self.slice_lora_a_weights(lora_A.to(data))
+            lora_B_sliced = self.slice_lora_b_weights(lora_B.to(data))
+
+            scale = lora_strength
+            if (
+                lora_alpha is not None
+                and lora_rank is not None
+                and lora_alpha != lora_rank
+            ):
+                scale *= lora_alpha / lora_rank
+
+            if not isinstance(lora_B_sliced, torch.Tensor):
+                lora_delta = lora_B_sliced @ lora_A_sliced
+                if isinstance(lora_delta, torch.Tensor) and lora_delta.dim() > 2:
+                    lora_delta = lora_delta.reshape(-1, lora_delta.shape[-1])
+                data.add_(lora_delta, alpha=scale)
+                continue
+
+            if lora_A_sliced.dim() > 2 or lora_B_sliced.dim() > 2:
+                lora_delta = lora_B_sliced @ lora_A_sliced
+                if lora_delta.dim() > 2:
+                    lora_delta = lora_delta.reshape(-1, lora_delta.shape[-1])
+                data_2d = data.reshape(-1, data.shape[-1]) if data.dim() > 2 else data
+                data_2d.add_(lora_delta, alpha=scale)
+                continue
+
+            data_2d = data.reshape(-1, data.shape[-1]) if data.dim() > 2 else data
+            lora_B_2d = (
+                lora_B_sliced.reshape(-1, lora_B_sliced.shape[-1])
+                if lora_B_sliced.dim() > 2
+                else lora_B_sliced
+            )
+
+            chunk_rows = max(
+                1,
+                LORA_MERGE_CHUNK_BYTES
+                // (data_2d.shape[-1] * max(1, data_2d.element_size())),
+            )
+            for start in range(0, lora_B_2d.shape[0], chunk_rows):
+                end = min(start + chunk_rows, lora_B_2d.shape[0])
+                chunk_delta = lora_B_2d[start:end] @ lora_A_sliced
+                data_2d[start:end].add_(chunk_delta, alpha=scale)
+
+    def _should_merge_in_fp32(
+        self,
+        lora_list: list[LoRAWeightEntry],
+    ) -> bool:
+        if os.getenv("SGLANG_DIFFUSION_LORA_MERGE_FP32", "1") != "1":
+            return False
+        for _, _, lora_path, _, _, _ in lora_list:
+            if lora_path and "distilled-lora" in lora_path.lower():
+                return False
+        return True
+
+    @torch.no_grad()
+    def merge_lora_weights(self, strength: float | None = None) -> None:
+        if strength is not None:
+            self.strength = strength
+            if self.lora_weights_list:
+                self.lora_weights_list = [
+                    (lora_A, lora_B, lora_path, strength, lora_rank, lora_alpha)
+                    for (
+                        lora_A,
+                        lora_B,
+                        lora_path,
+                        _,
+                        lora_rank,
+                        lora_alpha,
+                    ) in self.lora_weights_list
+                ]
+
+        if self.disable_lora:
+            return
+
+        if self.merged:
+            self.unmerge_lora_weights()
+
+        # Use lora_weights_list if available, otherwise fall back to single LoRA for backward compatibility
+        lora_list = self.lora_weights_list if self.lora_weights_list else []
+        if not lora_list and self.lora_A is not None and self.lora_B is not None:
+            lora_list = [
+                (
+                    self.lora_A,
+                    self.lora_B,
+                    self.lora_path,
+                    self.strength,
+                    self.lora_rank,
+                    self.lora_alpha,
+                )
+            ]
+
+        if not lora_list:
+            raise ValueError("LoRA weights not set. Please set them first.")
+
+        merge_in_fp32 = self._should_merge_in_fp32(lora_list)
+
+        if isinstance(self.base_layer.weight, DTensor):
+            mesh = self.base_layer.weight.data.device_mesh
+            unsharded_base_layer = ReplicatedLinear(
+                input_size=self.base_layer.input_size,
+                output_size=self.base_layer.output_size,
+                bias=getattr(self.base_layer, "bias", None) is not None,
+                skip_bias_add=self.base_layer.skip_bias_add,
+                params_dtype=self.base_layer.params_dtype,
+                quant_config=self.base_layer.quant_config,
+                prefix=self.base_layer.prefix,
+            )
+            # Using offload param is on CPU, so current_device is for "CPU -> GPU -> merge -> CPU"
+            current_device = self.base_layer.weight.data.device
+            data = self.base_layer.weight.data.to(
+                get_local_torch_device()
+            ).full_tensor()
+            data = self._as_mutable_tensor(data)
+            target_dtype = data.dtype
+            if (
+                merge_in_fp32
+                and data.is_floating_point()
+                and data.dtype != torch.float32
+            ):
+                data = data.to(torch.float32)
+
+            self._merge_lora_into_data(data, lora_list)
+
+            unsharded_base_layer.weight = nn.Parameter(
+                self._as_mutable_tensor(data.to(current_device, dtype=target_dtype))
+            )
+            if isinstance(getattr(self.base_layer, "bias", None), DTensor):
+                bias_data = (
+                    self.base_layer.bias.to(get_local_torch_device(), non_blocking=True)
+                    .full_tensor()
+                    .to(current_device)
+                )
+                unsharded_base_layer.bias = nn.Parameter(
+                    self._as_mutable_tensor(bias_data)
+                )
+
+            offload_policy = (
+                CPUOffloadPolicy() if "cpu" in str(current_device) else OffloadPolicy()
+            )
+            mp_policy = get_mixed_precision_state().mp_policy
+
+            self.base_layer = fully_shard(
+                unsharded_base_layer,
+                mesh=mesh,
+                mp_policy=mp_policy,
+                offload_policy=offload_policy,
+            )
+        else:
+            current_device = self.base_layer.weight.data.device
+            data = self.base_layer.weight.data.to(get_local_torch_device())
+            data = self._as_mutable_tensor(data)
+            target_dtype = data.dtype
+            if (
+                merge_in_fp32
+                and data.is_floating_point()
+                and data.dtype != torch.float32
+            ):
+                data = data.to(torch.float32)
+
+            self._merge_lora_into_data(data, lora_list)
+
+            self.base_layer.weight.data = self._as_mutable_tensor(
+                data.to(current_device, dtype=target_dtype, non_blocking=True)
+            )
+
+        self.merged = True
+
+    @torch.no_grad()
+    # @torch.compile(dynamic=True)
+    def unmerge_lora_weights(self) -> None:
+        if self.disable_lora:
+            return
+
+        if not self.merged:
+            raise ValueError(
+                "LoRA weights not merged. Please merge them first before unmerging."
+            )
+
+        # avoid precision loss
+        if isinstance(self.base_layer.weight, DTensor):
+            device = self.base_layer.weight.data.device
+            old_weight = self.base_layer.weight
+            new_weight_data = self._as_mutable_tensor(
+                self.cpu_weight.to(device, non_blocking=True)
+            )
+            self.base_layer.weight = nn.Parameter(new_weight_data)
+            del old_weight
+        else:
+            current_device = self.base_layer.weight.data.device
+            cpu_weight_on_device = self.cpu_weight.to(current_device, non_blocking=True)
+            if self.base_layer.weight.data.is_inference():
+                self.base_layer.weight.data = self._as_mutable_tensor(
+                    cpu_weight_on_device
+                )
+            else:
+                self.base_layer.weight.data.copy_(cpu_weight_on_device)
+            if (
+                cpu_weight_on_device.data_ptr()
+                != self.base_layer.weight.data.data_ptr()
+            ):
+                del cpu_weight_on_device
+
+        self.merged = False
+
+    @torch.no_grad()
+    def commit_merged_as_base(self) -> None:
+        """Promote the currently merged weights to the permanent base.
+
+        Re-snapshots ``cpu_weight`` so the merged weights become the restore
+        target and resets adapter bookkeeping (``merged=False``). A later dynamic
+        ``set_lora_weights`` then adds its delta on top of the merged base instead
+        of unmerging it.
+        """
+        if not self.merged:
+            return
+        weight = self.base_layer.weight
+        if isinstance(weight, DTensor):
+            weight = weight.to_local()
+        # clone(): to("cpu") may alias storage; we must not mutate this backup.
+        self.cpu_weight = weight.detach().to("cpu").clone()
+        self.merged = False
+        self.disable_lora = True
+        self.lora_weights_list = []
+        self.lora_A = None
+        self.lora_B = None
+        self.lora_path = None
+        self.strength = 1.0
+
+
+class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):
+    """
+    Vocab parallel embedding layer with support for LoRA (Low-Rank Adaptation).
+
+    Note: The current version does not yet implement the LoRA functionality.
+    This class behaves exactly the same as the base VocabParallelEmbedding.
+    Future versions will integrate LoRA functionality to support efficient parameter fine-tuning.
+    """
+
+    def __init__(
+        self,
+        base_layer: VocabParallelEmbedding,
+    ) -> None:
+        super().__init__(base_layer)
+
+    def forward(self, input_: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError(
+            "We don't support VocabParallelEmbeddingWithLoRA yet."
+        )
+
+
+class ColumnParallelLinearWithLoRA(BaseLayerWithLoRA):
+    def __init__(
+        self,
+        base_layer: ColumnParallelLinear,
+        lora_rank: int | None = None,
+        lora_alpha: int | None = None,
+    ) -> None:
+        super().__init__(base_layer, lora_rank, lora_alpha)
+
+    def forward(self, input_: torch.Tensor) -> torch.Tensor:
+        if self.merged or self.disable_lora:
+            return self.base_layer(input_)
+
+        lora_A = self.lora_A
+        lora_B = self.lora_B
+        if isinstance(self.lora_B, DTensor):
+            lora_B = self.lora_B.to_local()
+            lora_A = self.lora_A.to_local()
+
+        bias = self.base_layer.bias if not self.base_layer.skip_bias_add else None
+        output_parallel = self.base_layer.quant_method.apply(
+            self.base_layer, input_, bias
+        )
+        if not self.merged and not self.disable_lora:
+            lora_dtype = lora_A.dtype
+            input_lora = input_.to(dtype=lora_dtype)
+            lora_A_sliced = self.slice_lora_a_weights(
+                lora_A.to(device=input_.device, non_blocking=True)
+            )
+            lora_B_sliced = self.slice_lora_b_weights(
+                lora_B.to(device=input_.device, non_blocking=True)
+            )
+            delta_parallel = input_lora @ lora_A_sliced.T @ lora_B_sliced.T
+            if self.lora_alpha != self.lora_rank:
+                delta_parallel = delta_parallel * (
+                    self.lora_alpha / self.lora_rank  # type: ignore
+                )  # type: ignore
+            delta_parallel = delta_parallel * self.strength
+            output_parallel = output_parallel + delta_parallel.to(
+                dtype=output_parallel.dtype
+            )
+        if self.base_layer.gather_output:
+            output = tensor_model_parallel_all_gather(output_parallel)
+        else:
+            output = output_parallel
+        output_bias = self.base_layer.bias if self.base_layer.skip_bias_add else None
+        return output, output_bias
+
+    def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor:
+        return A
+
+    def slice_lora_b_weights(self, B: torch.Tensor) -> torch.Tensor:
+        tp_rank = get_tp_rank()
+        shard_size = self.base_layer.output_partition_sizes[0]
+        start_idx = tp_rank * shard_size
+        end_idx = (tp_rank + 1) * shard_size
+        B = B[start_idx:end_idx, :]
+        return B
+
+
+class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
+    def __init__(
+        self,
+        base_layer: MergedColumnParallelLinear,
+        lora_rank: int | None = None,
+        lora_alpha: int | None = None,
+    ) -> None:
+        super().__init__(base_layer, lora_rank, lora_alpha)
+
+    def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor:
+        return A
+
+    def slice_lora_b_weights(self, B: torch.Tensor) -> torch.Tensor:
+        tp_rank = get_tp_rank()
+        # Since the outputs for both gate and up are identical, we use a random one.
+        shard_size = self.base_layer.output_partition_sizes[0]
+        start_idx = tp_rank * shard_size
+        end_idx = (tp_rank + 1) * shard_size
+        return B[:, start_idx:end_idx, :]
+
+
+class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
+    def __init__(
+        self,
+        base_layer: QKVParallelLinear,
+        lora_rank: int | None = None,
+        lora_alpha: int | None = None,
+    ) -> None:
+        super().__init__(base_layer, lora_rank, lora_alpha)
+
+    def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor:
+        return A
+
+    def slice_lora_b_weights(
+        self, B: list[torch.Tensor]
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        tp_rank = get_tp_rank()
+        B_q, B_kv = B
+        base_layer = self.base_layer
+        q_proj_shard_size = base_layer.q_proj_shard_size
+        kv_proj_shard_size = base_layer.kv_proj_shard_size
+        num_kv_head_replicas = base_layer.num_kv_head_replicas
+
+        q_start_idx = q_proj_shard_size * tp_rank
+        q_end_idx = q_start_idx + q_proj_shard_size
+
+        kv_shard_id = tp_rank // num_kv_head_replicas
+        kv_start_idx = kv_proj_shard_size * kv_shard_id
+        kv_end_idx = kv_start_idx + kv_proj_shard_size
+
+        return B_q[q_start_idx:q_end_idx, :], B_kv[:, kv_start_idx:kv_end_idx, :]
+
+
+class RowParallelLinearWithLoRA(BaseLayerWithLoRA):
+    def __init__(
+        self,
+        base_layer: RowParallelLinear,
+        lora_rank: int | None = None,
+        lora_alpha: int | None = None,
+    ) -> None:
+        super().__init__(base_layer, lora_rank, lora_alpha)
+
+    def forward(self, input_: torch.Tensor):
+        if self.merged or self.disable_lora:
+            return self.base_layer(input_)
+
+        lora_A = self.lora_A
+        lora_B = self.lora_B
+        if isinstance(self.lora_B, DTensor):
+            lora_B = self.lora_B.to_local()
+            lora_A = self.lora_A.to_local()
+
+        if self.base_layer.input_is_parallel:
+            input_parallel = input_
+        else:
+            tp_rank = get_tp_rank()
+            splitted_input = split_tensor_along_last_dim(
+                input_, num_partitions=self.base_layer.tp_size
+            )
+            input_parallel = splitted_input[tp_rank].contiguous()
+        output_parallel = self.base_layer.quant_method.apply(
+            self.base_layer, input_parallel
+        )
+        if not self.merged and not self.disable_lora:
+            lora_dtype = lora_A.dtype
+            input_parallel_lora = input_parallel.to(dtype=lora_dtype)
+            lora_A_sliced = self.slice_lora_a_weights(
+                lora_A.to(device=input_parallel.device, non_blocking=True)
+            )
+            lora_B_sliced = self.slice_lora_b_weights(
+                lora_B.to(device=input_parallel.device, non_blocking=True)
+            )
+            delta_parallel = input_parallel_lora @ lora_A_sliced.T @ lora_B_sliced.T
+            if self.lora_alpha != self.lora_rank:
+                delta_parallel = delta_parallel * (
+                    self.lora_alpha / self.lora_rank  # type: ignore
+                )  # type: ignore
+            delta_parallel = delta_parallel * self.strength
+            output_parallel = output_parallel + delta_parallel.to(
+                dtype=output_parallel.dtype
+            )
+
+        if self.base_layer.reduce_results and self.base_layer.tp_size > 1:
+            output_ = tensor_model_parallel_all_reduce(output_parallel)
+        else:
+            output_ = output_parallel
+
+        if not self.base_layer.skip_bias_add:
+            output = (
+                output_ + self.base_layer.bias
+                if self.base_layer.bias is not None
+                else output_
+            )
+            output_bias = None
+        else:
+            output = output_
+            output_bias = self.base_layer.bias
+        return output, output_bias
+
+    def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor:
+        tp_rank = get_tp_rank()
+        shard_size = self.base_layer.input_size_per_partition
+        start_idx = tp_rank * shard_size
+        end_idx = (tp_rank + 1) * shard_size
+        A = A[:, start_idx:end_idx].contiguous()
+        return A
+
+    def slice_lora_b_weights(self, B: torch.Tensor) -> torch.Tensor:
+        return B
+
+
+class LinearWithLoRA(BaseLayerWithLoRA):
+    """
+    Wrapper for standard torch.nn.Linear to support LoRA.
+    Unlike custom LinearBase classes, nn.Linear.forward() returns a single tensor,
+    not a tuple of (output, bias).
+    """
+
+    def __init__(
+        self,
+        base_layer: nn.Linear,
+        lora_rank: int | None = None,
+        lora_alpha: int | None = None,
+    ) -> None:
+        super().__init__(base_layer, lora_rank, lora_alpha)
+
+    @torch.compile()
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        lora_A = self.lora_A
+        lora_B = self.lora_B
+        if isinstance(self.lora_B, DTensor):
+            lora_B = self.lora_B.to_local()
+            lora_A = self.lora_A.to_local()
+
+        # TODO: Support multiple LoRA adapters when use not merged mode
+        if not self.merged and not self.disable_lora:
+            lora_dtype = lora_A.dtype
+            x_lora = x.to(dtype=lora_dtype)
+            lora_A_sliced = self.slice_lora_a_weights(
+                lora_A.to(device=x.device, non_blocking=True)
+            )
+            lora_B_sliced = self.slice_lora_b_weights(
+                lora_B.to(device=x.device, non_blocking=True)
+            )
+            delta = x_lora @ lora_A_sliced.T @ lora_B_sliced.T
+            if self.lora_alpha != self.lora_rank:
+                delta = delta * (
+                    self.lora_alpha / self.lora_rank  # type: ignore
+                )  # type: ignore
+            delta = delta * self.strength
+            # nn.Linear.forward() returns a single tensor, not a tuple
+            out = self.base_layer(x)
+            return out + delta.to(dtype=out.dtype)
+        else:
+            # nn.Linear.forward() returns a single tensor
+            out = self.base_layer(x)
+            return out
+
+
+def wrap_with_lora_layer(
+    layer: nn.Module,
+    lora_rank: int | None = None,
+    lora_alpha: int | None = None,
+) -> BaseLayerWithLoRA | None:
+    """
+    transform the given layer to its corresponding LoRA layer
+    """
+    supported_layer_types: dict[
+        type[LinearBase] | type[nn.Linear], type[BaseLayerWithLoRA]
+    ] = {
+        # the order matters
+        # VocabParallelEmbedding: VocabParallelEmbeddingWithLoRA,
+        QKVParallelLinear: QKVParallelLinearWithLoRA,
+        MergedColumnParallelLinear: MergedColumnParallelLinearWithLoRA,
+        ColumnParallelLinear: ColumnParallelLinearWithLoRA,
+        RowParallelLinear: RowParallelLinearWithLoRA,
+        ReplicatedLinear: BaseLayerWithLoRA,
+        nn.Linear: LinearWithLoRA,
+    }
+    for src_layer_type, lora_layer_type in supported_layer_types.items():
+        if isinstance(layer, src_layer_type):  # type: ignore[arg-type]
+            ret = lora_layer_type(
+                layer,
+                lora_rank=lora_rank,
+                lora_alpha=lora_alpha,
+            )
+            return ret
+    return None
+
+
+# source: https://github.com/vllm-project/vllm/blob/93b38bea5dd03e1b140ca997dfaadef86f8f1855/vllm/lora/utils.py#L9
+def replace_submodule(
+    model: nn.Module, module_name: str, new_module: nn.Module
+) -> nn.Module:
+    """Replace a submodule in a model with a new module."""
+    parent = model.get_submodule(".".join(module_name.split(".")[:-1]))
+    target_name = module_name.split(".")[-1]
+    setattr(parent, target_name, new_module)
+    return new_module

python/sglang/multimodal_gen/runtime/layers/mlp.py

@@ -0,0 +1,121 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from diffusers.models.activations import (
+    GEGLU,
+    GELU,
+    ApproximateGELU,
+    LinearActivation,
+    SwiGLU,
+)
+
+from sglang.multimodal_gen.runtime.layers.activation import get_act_fn
+from sglang.multimodal_gen.runtime.layers.linear import (
+    ColumnParallelLinear,
+    RowParallelLinear,
+)
+from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig
+from sglang.srt.utils import add_prefix
+
+
+class MLP(nn.Module):
+    """
+    MLP for DiT blocks, NO gated linear units
+    """
+
+    def __init__(
+        self,
+        input_dim: int,
+        mlp_hidden_dim: int,
+        output_dim: int | None = None,
+        bias: bool = True,
+        act_type: str = "gelu_pytorch_tanh",
+        dtype: torch.dtype | None = None,
+        prefix: str = "",
+        quant_config: QuantizationConfig = None,
+    ):
+        super().__init__()
+        self.fc_in = ColumnParallelLinear(
+            input_dim,
+            mlp_hidden_dim,
+            bias=True,
+            gather_output=False,
+            quant_config=quant_config,
+            prefix=add_prefix("fc_in", prefix),
+        )
+
+        self.act = get_act_fn(act_type)
+        if output_dim is None:
+            output_dim = input_dim
+        self.fc_out = RowParallelLinear(
+            mlp_hidden_dim,
+            output_dim,
+            bias=True,
+            input_is_parallel=True,
+            quant_config=quant_config,
+            prefix=add_prefix("fc_out", prefix),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x, _ = self.fc_in(x)
+        x = self.act(x)
+        x, _ = self.fc_out(x)
+        return x
+
+
+class FeedForward(nn.Module):
+    r"""
+    A feed-forward layer.
+
+    Parameters:
+        dim (`int`): The number of channels in the input.
+        dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`.
+        mult (`int`, *optional*, defaults to 4): The multiplier to use for the hidden dimension.
+        activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
+        bias (`bool`, defaults to True): Whether to use a bias in the linear layer.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        dim_out: Optional[int] = None,
+        mult: int = 4,
+        activation_fn: str = "geglu",
+        inner_dim=None,
+        bias: bool = True,
+    ):
+        super().__init__()
+        if inner_dim is None:
+            inner_dim = int(dim * mult)
+        dim_out = dim_out if dim_out is not None else dim
+
+        if activation_fn == "gelu":
+            act_fn = GELU(dim, inner_dim, bias=bias)
+        if activation_fn == "gelu-approximate":
+            act_fn = GELU(dim, inner_dim, approximate="tanh", bias=bias)
+        elif activation_fn == "geglu":
+            act_fn = GEGLU(dim, inner_dim, bias=bias)
+        elif activation_fn == "geglu-approximate":
+            act_fn = ApproximateGELU(dim, inner_dim, bias=bias)
+        elif activation_fn == "swiglu":
+            act_fn = SwiGLU(dim, inner_dim, bias=bias)
+        elif activation_fn == "linear-silu":
+            act_fn = LinearActivation(dim, inner_dim, bias=bias, activation="silu")
+
+        self.net = nn.ModuleList([])
+        # project in
+        self.net.append(act_fn)
+        # dummy dropout layer to match with checkpoints compatible with diffusers
+        self.net.append(nn.Dropout(0.0))
+        # project out
+        self.net.append(nn.Linear(inner_dim, dim_out, bias=bias))
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        for module in self.net:
+            hidden_states = module(hidden_states)
+        return hidden_states

python/sglang/multimodal_gen/runtime/layers/parallel_conv.py

@@ -0,0 +1,801 @@
+import contextvars
+import math
+from contextlib import contextmanager
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+import torch.nn.functional as F
+
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    get_decode_parallel_group_coordinator,
+    get_decode_parallel_rank,
+    get_decode_parallel_world_size,
+)
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+if current_platform.is_cuda():
+    from sglang.jit_kernel.diffusion.causal_conv3d_cat_pad import (
+        can_use_fused_causal_conv3d_cat_pad_cuda,
+        fused_causal_conv3d_cat_pad_cuda,
+    )
+    from sglang.jit_kernel.diffusion.triton.causal_conv3d_pad import (
+        fused_causal_conv3d_cat_pad as fused_causal_conv3d_cat_pad_triton,
+    )
+else:
+    can_use_fused_causal_conv3d_cat_pad_cuda = None
+    fused_causal_conv3d_cat_pad_cuda = None
+    fused_causal_conv3d_cat_pad_triton = None
+
+
+_causal_conv3d_cat_pad_cuda_failed = False
+
+
+def fused_causal_conv3d_cat_pad(
+    x: torch.Tensor,
+    cache_x: torch.Tensor,
+    padding: list[int],
+) -> torch.Tensor:
+    global _causal_conv3d_cat_pad_cuda_failed
+    if (
+        fused_causal_conv3d_cat_pad_cuda is not None
+        and can_use_fused_causal_conv3d_cat_pad_cuda(x, cache_x, padding)
+        and not _causal_conv3d_cat_pad_cuda_failed
+    ):
+        try:
+            return fused_causal_conv3d_cat_pad_cuda(x, cache_x, padding)
+        except Exception:
+            logger.warning(
+                "fused_causal_conv3d_cat_pad_cuda failed, falling back to Triton",
+                exc_info=True,
+            )
+            _causal_conv3d_cat_pad_cuda_failed = True
+    if fused_causal_conv3d_cat_pad_triton is None:
+        raise RuntimeError("causal Conv3D cat/pad fusion is only available on CUDA")
+    return fused_causal_conv3d_cat_pad_triton(x, cache_x, padding)
+
+
+_SPATIAL_PARALLEL_DECODE_DISABLED = contextvars.ContextVar(
+    "spatial_parallel_decode_disabled", default=False
+)
+
+
+@contextmanager
+def disable_spatial_parallel_decode():
+    token = _SPATIAL_PARALLEL_DECODE_DISABLED.set(True)
+    try:
+        yield
+    finally:
+        _SPATIAL_PARALLEL_DECODE_DISABLED.reset(token)
+
+
+def spatial_parallel_decode_disabled() -> bool:
+    return _SPATIAL_PARALLEL_DECODE_DISABLED.get()
+
+
+def _tensor_pad(x: torch.Tensor, len_to_pad: int, dim: int = -2):
+    return torch.cat(
+        [
+            x,
+            torch.zeros(
+                *x.shape[:dim],
+                len_to_pad,
+                *x.shape[dim + 1 :],
+                dtype=x.dtype,
+                device=x.device,
+            ),
+        ],
+        dim=dim,
+    )
+
+
+def _tensor_chunk(x: torch.Tensor, dim: int = -2, world_size: int = 1, rank: int = 0):
+    if x is None:
+        return x
+    if world_size <= 1:
+        return x
+    return torch.tensor_split(x, world_size, dim=dim)[rank].contiguous(
+        memory_format=_halo_memory_format(x)
+    )
+
+
+def _can_fuse_causal_conv3d_cat_pad(
+    x: torch.Tensor,
+    cache_x: torch.Tensor | None,
+    padding: list[int],
+) -> bool:
+    if cache_x is None or fused_causal_conv3d_cat_pad is None:
+        return False
+    if not current_platform.is_cuda():
+        return False
+    if not x.is_cuda or not x.is_contiguous() or not cache_x.is_contiguous():
+        return False
+    if x.dim() != 5 or cache_x.dim() != 5 or x.dtype != cache_x.dtype:
+        return False
+    if x.shape[0] != cache_x.shape[0] or x.shape[1] != cache_x.shape[1]:
+        return False
+    if x.shape[3:] != cache_x.shape[3:]:
+        return False
+
+    width_left, width_right, height_top, height_bottom, depth_left, depth_right = (
+        padding
+    )
+    if width_left != width_right or height_top != height_bottom or depth_right != 0:
+        return False
+    if depth_left < cache_x.shape[2]:
+        return False
+    return bool(width_left or height_top)
+
+
+def causal_conv3d_cat_pad(
+    x: torch.Tensor,
+    cache_x: torch.Tensor | None,
+    padding: list[int],
+) -> torch.Tensor:
+    if cache_x is not None and padding[4] > 0:
+        if cache_x.device != x.device:
+            cache_x = cache_x.to(x.device)
+        if _can_fuse_causal_conv3d_cat_pad(x, cache_x, padding):
+            return fused_causal_conv3d_cat_pad(x, cache_x, padding)
+        x = torch.cat([cache_x, x], dim=2)
+        padding[4] -= cache_x.shape[2]
+    if any(padding):
+        x = F.pad(x, padding)
+    return x
+
+
+def split_for_parallel_decode(
+    x: torch.Tensor, upsample_count: int, world_size: int, rank: int
+):
+    return split_height_for_parallel_decode(
+        x,
+        expected_height=x.shape[-2] * (2**upsample_count),
+        world_size=world_size,
+        rank=rank,
+    )
+
+
+def split_height_for_parallel_decode(
+    x: torch.Tensor, expected_height: int, world_size: int, rank: int
+):
+    if spatial_parallel_decode_disabled():
+        return x, None
+    x = _tensor_chunk(x, dim=-2, world_size=world_size, rank=rank)
+    return x, expected_height
+
+
+def _maybe_contiguous_for_sp_gather(x: torch.Tensor) -> torch.Tensor:
+    if (
+        x.dim() == 5
+        and hasattr(torch, "channels_last_3d")
+        and x.is_contiguous(memory_format=torch.channels_last_3d)
+        and not x.is_contiguous()
+    ):
+        return x.contiguous()
+    if (
+        x.dim() == 4
+        and x.is_contiguous(memory_format=torch.channels_last)
+        and not x.is_contiguous()
+    ):
+        return x.contiguous()
+    return x
+
+
+def gather_and_trim_height(x: torch.Tensor, expected_height: int | None):
+    if spatial_parallel_decode_disabled():
+        return x
+    if expected_height is None:
+        return x
+    x, _ = gather_variable_height(x)
+    if x.shape[-2] != expected_height:
+        x = x[..., :expected_height, :].contiguous()
+    return x
+
+
+def gather_height_for_global_op(x: torch.Tensor) -> torch.Tensor:
+    if spatial_parallel_decode_disabled():
+        return x
+    return gather_variable_height(x)[0]
+
+
+def chunk_height_for_parallel_decode(x: torch.Tensor) -> torch.Tensor:
+    if spatial_parallel_decode_disabled():
+        return x
+    return _tensor_chunk(
+        x,
+        dim=-2,
+        world_size=get_decode_parallel_world_size(),
+        rank=get_decode_parallel_rank(),
+    )
+
+
+def chunk_height_by_sizes(x: torch.Tensor, heights: list[int]) -> torch.Tensor:
+    if spatial_parallel_decode_disabled():
+        return x
+    rank = get_decode_parallel_rank()
+    start = sum(heights[:rank])
+    return x[..., start : start + heights[rank], :].contiguous(
+        memory_format=_halo_memory_format(x)
+    )
+
+
+def gather_height_sizes(x: torch.Tensor) -> list[int]:
+    """gather heights of sharded feature_maps from peers"""
+    if spatial_parallel_decode_disabled():
+        return [x.shape[-2]]
+    world_size = get_decode_parallel_world_size()
+    if world_size <= 1:
+        return [x.shape[-2]]
+    local_height = torch.tensor([x.shape[-2]], device=x.device, dtype=torch.int64)
+    gathered = [torch.empty_like(local_height) for _ in range(world_size)]
+    dist.all_gather(
+        gathered,
+        local_height,
+        group=get_decode_parallel_group_coordinator().device_group,
+    )
+    return [int(height.item()) for height in gathered]
+
+
+def gather_variable_height(x: torch.Tensor) -> tuple[torch.Tensor, list[int]]:
+    if spatial_parallel_decode_disabled():
+        return x, [x.shape[-2]]
+    world_size = get_decode_parallel_world_size()
+    if world_size <= 1:
+        return x, [x.shape[-2]]
+
+    heights = gather_height_sizes(x)
+    max_height = max(heights)
+    if x.shape[-2] < max_height:
+        x = _tensor_pad(x, max_height - x.shape[-2], dim=-2)
+
+    gathered = get_decode_parallel_group_coordinator().all_gather(
+        _maybe_contiguous_for_sp_gather(x), dim=-2
+    )
+    chunks = torch.split(gathered, max_height, dim=-2)
+    return (
+        torch.cat(
+            [chunk[..., :height, :] for chunk, height in zip(chunks, heights)], dim=-2
+        ),
+        heights,
+    )
+
+
+def _halo_memory_format(reference: torch.Tensor) -> torch.memory_format:
+    if reference.dim() > 1 and reference.stride(1) == 1:
+        if reference.dim() == 5 and hasattr(torch, "channels_last_3d"):
+            return torch.channels_last_3d
+        if reference.dim() == 4:
+            return torch.channels_last
+    return torch.contiguous_format
+
+
+def _ensure_recv_buf(
+    recv_buf: torch.Tensor | None, reference: torch.Tensor
+) -> torch.Tensor:
+    memory_format = _halo_memory_format(reference)
+    if (
+        recv_buf is None
+        or recv_buf.shape != reference.shape
+        or recv_buf.dtype != reference.dtype
+        or recv_buf.device != reference.device
+        or not recv_buf.is_contiguous(memory_format=memory_format)
+    ):
+        return torch.empty(
+            reference.shape,
+            dtype=reference.dtype,
+            device=reference.device,
+            memory_format=memory_format,
+        )
+    return recv_buf
+
+
+def halo_exchange(
+    x: torch.Tensor,
+    height_halo_size: int = 1,
+    recv_top_buf: torch.Tensor | None = None,
+    recv_bottom_buf: torch.Tensor | None = None,
+    height_pad_mode: str = "zeros",
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """exchange(send and recv) top/bottom conv-input halos with adjacent spatial ranks"""
+    if spatial_parallel_decode_disabled():
+        return x, recv_top_buf, recv_bottom_buf
+    if height_halo_size == 0:
+        return x, recv_top_buf, recv_bottom_buf
+
+    decode_group = get_decode_parallel_group_coordinator()
+    rank = get_decode_parallel_rank()
+    world_size = get_decode_parallel_world_size()
+    group = decode_group.device_group
+    group_ranks = decode_group.ranks
+
+    top_row_ref = x[..., :height_halo_size, :]
+    bottom_row_ref = x[..., -height_halo_size:, :]
+
+    recv_top_buf = _ensure_recv_buf(recv_top_buf, top_row_ref)
+    recv_bottom_buf = _ensure_recv_buf(recv_bottom_buf, bottom_row_ref)
+    p2p_ops = []
+
+    if rank > 0:
+        prev_rank = group_ranks[rank - 1]
+        top_row = top_row_ref.contiguous(memory_format=_halo_memory_format(top_row_ref))
+        p2p_ops.append(dist.P2POp(dist.irecv, recv_top_buf, prev_rank, group))
+        p2p_ops.append(dist.P2POp(dist.isend, top_row, prev_rank, group))
+    if rank < world_size - 1:
+        next_rank = group_ranks[rank + 1]
+        bottom_row = bottom_row_ref.contiguous(
+            memory_format=_halo_memory_format(bottom_row_ref)
+        )
+        p2p_ops.append(dist.P2POp(dist.isend, bottom_row, next_rank, group))
+        p2p_ops.append(dist.P2POp(dist.irecv, recv_bottom_buf, next_rank, group))
+
+    if p2p_ops:
+        reqs = dist.batch_isend_irecv(p2p_ops)
+        for req in reqs:
+            req.wait()
+
+    if rank == 0:
+        recv_top_buf.copy_(
+            _make_boundary_halo(
+                x,
+                recv_bottom_buf if world_size > 1 else None,
+                height_halo_size=height_halo_size,
+                is_top=True,
+                mode=height_pad_mode,
+            )
+        )
+    if rank == world_size - 1:
+        recv_bottom_buf.copy_(
+            _make_boundary_halo(
+                x,
+                recv_top_buf if world_size > 1 else None,
+                height_halo_size=height_halo_size,
+                is_top=False,
+                mode=height_pad_mode,
+            )
+        )
+
+    return (
+        torch.concat([recv_top_buf, x, recv_bottom_buf], dim=-2),
+        recv_top_buf,
+        recv_bottom_buf,
+    )
+
+
+def _make_boundary_halo(
+    x: torch.Tensor,
+    neighbor: torch.Tensor | None,
+    *,
+    height_halo_size: int,
+    is_top: bool,
+    mode: str,
+) -> torch.Tensor:
+    if mode == "zeros":
+        shape = list(x.shape)
+        shape[-2] = height_halo_size
+        return torch.zeros(shape, dtype=x.dtype, device=x.device)
+    if mode == "replicate":
+        edge = x[..., :1, :] if is_top else x[..., -1:, :]
+        return edge.expand(*edge.shape[:-2], height_halo_size, edge.shape[-1])
+    if mode == "reflect":
+        source = x
+        if is_top and neighbor is not None:
+            source = torch.cat([x, neighbor], dim=-2)
+        elif not is_top and neighbor is not None:
+            source = torch.cat([neighbor, x], dim=-2)
+        if is_top:
+            index = torch.arange(
+                height_halo_size, 0, -1, device=x.device, dtype=torch.long
+            )
+        else:
+            index = torch.arange(
+                source.shape[-2] - 2,
+                source.shape[-2] - 2 - height_halo_size,
+                -1,
+                device=x.device,
+                dtype=torch.long,
+            )
+        return source.index_select(-2, index)
+    raise ValueError(f"Unsupported spatial padding mode for parallel decode: {mode}")
+
+
+def _pad_with_mode(
+    x: torch.Tensor, padding: tuple[int, ...], mode: str
+) -> torch.Tensor:
+    if mode == "zeros":
+        return F.pad(x, padding)
+    return F.pad(x, padding, mode=mode)
+
+
+def _set_conv_padding(module: nn.Module, padding: tuple[int, ...]) -> None:
+    module.padding = padding
+    module._reversed_padding_repeated_twice = tuple(
+        value for pad in reversed(padding) for value in (pad, pad)
+    )
+
+
+def _conv_preserves_local_height(
+    *,
+    height_halo_size: int,
+    height_pad_top: int,
+    height_pad_bottom: int,
+    kernel_height: int,
+    dilation_height: int,
+    stride_height: int,
+) -> bool:
+    kernel_span = dilation_height * (kernel_height - 1)
+    return (
+        stride_height == 1
+        and 2 * height_halo_size == kernel_span
+        and height_pad_top == height_halo_size
+        and height_pad_bottom == height_halo_size
+    )
+
+
+def _conv3d_weight_is_channels_last_3d(weight: torch.Tensor) -> bool:
+    return (
+        weight.dim() == 5
+        and hasattr(torch, "channels_last_3d")
+        and (current_platform.is_cuda() or current_platform.is_rocm())
+        and weight.is_contiguous(memory_format=torch.channels_last_3d)
+    )
+
+
+def _match_conv3d_input_format(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+    if x.dim() == 5 and _conv3d_weight_is_channels_last_3d(weight):
+        return x.contiguous(memory_format=torch.channels_last_3d)
+    return x
+
+
+def _spatial_parallel_conv_forward(
+    module: nn.Module,
+    x: torch.Tensor,
+    conv_forward,
+    *,
+    height_pad_mode: str,
+    match_conv3d_format: bool = False,
+) -> torch.Tensor:
+    # send and recv halo
+    # x_padded: concatenated input
+    x_padded, module._halo_recv_top_buf, module._halo_recv_bottom_buf = halo_exchange(
+        x,
+        height_halo_size=module.height_halo_size,
+        recv_top_buf=module._halo_recv_top_buf,
+        recv_bottom_buf=module._halo_recv_bottom_buf,
+        height_pad_mode=height_pad_mode,
+    )
+    if match_conv3d_format:
+        x_padded = _match_conv3d_input_format(x_padded, module.weight)
+    if module.height_halo_size == 0:
+        return conv_forward(x_padded)
+
+    stride = module.stride[-2]
+    if _conv_preserves_local_height(
+        height_halo_size=module.height_halo_size,
+        height_pad_top=module.height_pad_top,
+        height_pad_bottom=module.height_pad_bottom,
+        kernel_height=module.kernel_size[-2],
+        dilation_height=module.dilation[-2],
+        stride_height=stride,
+    ):
+        return conv_forward(x_padded)
+
+    heights = gather_height_sizes(x)
+    global_start = sum(heights[: module.rank])
+    global_height = sum(heights)
+    if stride > 1:
+        shift = (
+            global_start - module.height_halo_size + module.height_pad_top
+        ) % stride
+        if shift:
+            x_padded = x_padded[..., shift:, :]
+            global_start += shift
+        if match_conv3d_format:
+            x_padded = _match_conv3d_input_format(x_padded, module.weight)
+
+    out = conv_forward(x_padded)
+
+    # trim the output to original shape
+    return _trim_conv_output_height(
+        out,
+        local_height=x.shape[-2],
+        global_height=global_height,
+        global_start=global_start,
+        height_halo_size=module.height_halo_size,
+        height_pad_top=module.height_pad_top,
+        height_pad_bottom=module.height_pad_bottom,
+        kernel_height=module.kernel_size[-2],
+        dilation_height=module.dilation[-2],
+        stride_height=stride,
+    )
+
+
+class SpatialParallelConv2d(nn.Conv2d):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int | tuple[int, int],
+        stride: int | tuple[int, int] = 1,
+        padding: int | tuple[int, int] = 0,
+        dilation: int | tuple[int, int] = 1,
+        groups: int = 1,
+        bias: bool = True,
+        padding_mode: str = "zeros",
+        height_padding: tuple[int, int] | None = None,
+    ):
+        super().__init__(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+            bias=bias,
+            padding_mode=padding_mode,
+        )
+        self.height_halo_size = (self.dilation[-2] * (self.kernel_size[-2] - 1)) // 2
+        if height_padding is None:
+            height_padding = (self.padding[-2], self.padding[-2])
+        self.height_pad_top, self.height_pad_bottom = height_padding
+
+        self.padding: tuple[int, int]
+        if self.height_halo_size > 0:
+            self._padding = (0, 0, 0, 0)
+        else:
+            self._padding = (0, 0, self.padding[0], self.padding[0])
+
+        _set_conv_padding(self, (0, self.padding[1]))
+        self._halo_recv_top_buf: torch.Tensor | None = None
+        self._halo_recv_bottom_buf: torch.Tensor | None = None
+        self.rank = get_decode_parallel_rank()
+        self.world_size = get_decode_parallel_world_size()
+
+    def forward(self, x):
+        if spatial_parallel_decode_disabled():
+            return self._direct_forward(x)
+
+        if any(self._padding):
+            x = _pad_with_mode(x, self._padding, self.padding_mode)
+
+        return _spatial_parallel_conv_forward(
+            self,
+            x,
+            super().forward,
+            height_pad_mode=self.padding_mode,
+        )
+
+    def _direct_forward(self, x):
+        width_pad = self.padding[-1]
+        padding = (
+            width_pad,
+            width_pad,
+            self.height_pad_top,
+            self.height_pad_bottom,
+        )
+        if any(padding):
+            x = _pad_with_mode(x, padding, self.padding_mode)
+        return F.conv2d(
+            x,
+            self.weight,
+            self.bias,
+            self.stride,
+            (0, 0),
+            self.dilation,
+            self.groups,
+        )
+
+
+class SpatialParallelCausalConv3d(nn.Conv3d):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int | tuple[int, int, int],
+        stride: int | tuple[int, int, int] = 1,
+        padding: int | tuple[int, int, int] = 0,
+    ):
+        super().__init__(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+        )
+
+        self.height_pad_top = self.padding[1]
+        self.height_pad_bottom = self.padding[1]
+        self.height_halo_size = (self.kernel_size[-2] - 1) // 2
+
+        self.padding: tuple[int, int, int]
+        if self.height_halo_size > 0:
+            self._padding = (
+                self.padding[2],
+                self.padding[2],
+                0,
+                0,
+                2 * self.padding[0],
+                0,
+            )
+        else:
+            self._padding = (
+                self.padding[2],
+                self.padding[2],
+                self.padding[1],
+                self.padding[1],
+                2 * self.padding[0],
+                0,
+            )
+        self.padding = (0, 0, 0)
+        self._halo_recv_top_buf: torch.Tensor | None = None
+        self._halo_recv_bottom_buf: torch.Tensor | None = None
+        self.rank = get_decode_parallel_rank()
+        self.world_size = get_decode_parallel_world_size()
+
+    def forward(self, x, cache_x=None):
+        padding = list(self._padding)
+        if spatial_parallel_decode_disabled():
+            padding[2] = self.height_pad_top
+            padding[3] = self.height_pad_bottom
+        x = causal_conv3d_cat_pad(x, cache_x, padding)
+        x = x if current_platform.is_amp_supported() else x.to(self.weight.dtype)
+
+        if spatial_parallel_decode_disabled():
+            x = _match_conv3d_input_format(x, self.weight)
+            return F.conv3d(
+                x,
+                self.weight,
+                self.bias,
+                self.stride,
+                (0, 0, 0),
+                self.dilation,
+                self.groups,
+            )
+
+        return _spatial_parallel_conv_forward(
+            self,
+            x,
+            super().forward,
+            height_pad_mode="zeros",
+            match_conv3d_format=True,
+        )
+
+
+class SpatialParallelConv3d(nn.Conv3d):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int | tuple[int, int, int],
+        stride: int | tuple[int, int, int] = 1,
+        padding: int | tuple[int, int, int] = 0,
+        dilation: int | tuple[int, int, int] = 1,
+        groups: int = 1,
+        bias: bool = True,
+        padding_mode: str = "zeros",
+        height_padding: tuple[int, int] | None = None,
+    ):
+        super().__init__(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+            bias=bias,
+            padding_mode=padding_mode,
+        )
+        self.height_halo_size = (self.dilation[-2] * (self.kernel_size[-2] - 1)) // 2
+        if height_padding is None:
+            height_padding = (self.padding[-2], self.padding[-2])
+        self.height_pad_top, self.height_pad_bottom = height_padding
+
+        self.padding: tuple[int, int, int]
+        if self.height_halo_size > 0:
+            self._padding = (0, 0, 0, 0, 0, 0)
+        else:
+            self._padding = (
+                0,
+                0,
+                self.padding[1],
+                self.padding[1],
+                0,
+                0,
+            )
+
+        _set_conv_padding(self, (self.padding[0], 0, self.padding[2]))
+        self._halo_recv_top_buf: torch.Tensor | None = None
+        self._halo_recv_bottom_buf: torch.Tensor | None = None
+        self.rank = get_decode_parallel_rank()
+        self.world_size = get_decode_parallel_world_size()
+
+    def forward(self, x):
+        if spatial_parallel_decode_disabled():
+            return self._direct_forward(x)
+
+        if any(self._padding):
+            x = _pad_with_mode(x, self._padding, self.padding_mode)
+
+        return _spatial_parallel_conv_forward(
+            self,
+            x,
+            super().forward,
+            height_pad_mode=self.padding_mode,
+            match_conv3d_format=True,
+        )
+
+    def _direct_forward(self, x):
+        time_pad = self.padding[0]
+        width_pad = self.padding[-1]
+        padding = (
+            width_pad,
+            width_pad,
+            self.height_pad_top,
+            self.height_pad_bottom,
+            time_pad,
+            time_pad,
+        )
+        if any(padding):
+            x = _pad_with_mode(x, padding, self.padding_mode)
+        x = _match_conv3d_input_format(x, self.weight)
+        return F.conv3d(
+            x,
+            self.weight,
+            self.bias,
+            self.stride,
+            (0, 0, 0),
+            self.dilation,
+            self.groups,
+        )
+
+
+class SpatialParallelZeroPad2d(nn.Module):
+    def __init__(self, padding: tuple[int, int, int, int]) -> None:
+        super().__init__()
+        self.padding = padding
+        self.rank = get_decode_parallel_rank()
+        self.world_size = get_decode_parallel_world_size()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if spatial_parallel_decode_disabled():
+            return F.pad(x, self.padding)
+        left, right, top, bottom = self.padding
+        top = top if self.rank == 0 else 0
+        bottom = bottom if self.rank == self.world_size - 1 else 0
+        return F.pad(x, (left, right, top, bottom))
+
+
+def _trim_conv_output_height(
+    out: torch.Tensor,
+    *,
+    local_height: int,
+    global_height: int,
+    global_start: int,
+    height_halo_size: int,
+    height_pad_top: int,
+    height_pad_bottom: int,
+    kernel_height: int,
+    dilation_height: int,
+    stride_height: int,
+) -> torch.Tensor:
+    kernel_span = dilation_height * (kernel_height - 1)
+    min_i = math.ceil(
+        ((-height_pad_top) - (global_start - height_halo_size)) / stride_height
+    )
+    max_i = math.floor(
+        (
+            (global_height - 1 + height_pad_bottom)
+            - kernel_span
+            - (global_start - height_halo_size)
+        )
+        / stride_height
+    )
+    start = max(min_i, 0)
+    end = min(max_i + 1, out.shape[-2])
+    if start != 0 or end != out.shape[-2]:
+        out = out[..., start:end, :]
+    return out

python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py

@@ -0,0 +1,98 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+from typing import Literal, get_args
+
+from sglang.multimodal_gen.runtime.layers.quantization.bitsandbytes import (
+    BitsAndBytesConfig,
+)
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+)
+from sglang.multimodal_gen.runtime.layers.quantization.fp8 import Fp8Config
+from sglang.multimodal_gen.runtime.layers.quantization.modelopt_fp8 import (
+    ModelOptFp8Config as ModelOptFp8DiffusionConfig,
+)
+from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
+    ModelOptFp4Config,
+    ModelOptFp8Config,
+)
+from sglang.multimodal_gen.runtime.layers.quantization.modelslim import ModelSlimConfig
+from sglang.multimodal_gen.runtime.layers.quantization.mxfp4 import Mxfp4Config
+from sglang.multimodal_gen.runtime.layers.quantization.mxfp4_npu import (
+    NPUMXFP4Config,
+)
+from sglang.multimodal_gen.runtime.layers.quantization.mxfp8_npu import MXFP8Config
+
+QuantizationMethods = Literal[
+    "fp8",
+    "modelopt",
+    "modelopt_fp8",
+    "modelopt_fp4",
+    "bitsandbytes",
+    "modelslim",
+    "mxfp8",
+    "mxfp4",
+    "mxfp4_npu",
+]
+
+QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods))
+
+# The customized quantization methods which will be added to this dict.
+_CUSTOMIZED_METHOD_TO_QUANT_CONFIG = {
+    "modelopt": ModelOptFp8DiffusionConfig,
+    "modelopt_fp8": ModelOptFp8Config,
+    "modelopt_fp4": ModelOptFp4Config,
+    "bitsandbytes": BitsAndBytesConfig,
+    "modelslim": ModelSlimConfig,
+    "fp8": Fp8Config,
+    "mxfp4": Mxfp4Config,
+    "mxfp8": MXFP8Config,
+    "mxfp4_npu": NPUMXFP4Config,
+}
+
+
+def register_quantization_config(quantization: str):
+    """Register a customized vllm quantization config.
+
+    When a quantization method is not supported by vllm, you can register a customized
+    quantization config to support it.
+
+    Args:
+        quantization (str): The quantization method name.
+
+
+    """  # noqa: E501
+
+    def _wrapper(quant_config_cls):
+        if quantization in QUANTIZATION_METHODS:
+            raise ValueError(
+                f"The quantization method `{quantization}` is already exists."
+            )
+        if not issubclass(quant_config_cls, QuantizationConfig):
+            raise ValueError(
+                "The quantization config must be a subclass of " "`QuantizationConfig`."
+            )
+        _CUSTOMIZED_METHOD_TO_QUANT_CONFIG[quantization] = quant_config_cls
+        QUANTIZATION_METHODS.append(quantization)
+        return quant_config_cls
+
+    return _wrapper
+
+
+def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
+    if quantization not in QUANTIZATION_METHODS:
+        raise ValueError(f"Invalid quantization method: {quantization}")
+
+    method_to_config: dict[str, type[QuantizationConfig]] = {}
+    # Update the `method_to_config` with customized quantization methods.
+    method_to_config.update(_CUSTOMIZED_METHOD_TO_QUANT_CONFIG)
+
+    return method_to_config[quantization]
+
+
+__all__ = [
+    "QuantizationMethods",
+    "QuantizationConfig",
+    "get_quantization_config",
+    "QUANTIZATION_METHODS",
+]

python/sglang/multimodal_gen/runtime/layers/quantization/bitsandbytes.py

@@ -0,0 +1,437 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+import torch
+import torch.nn as nn
+from packaging import version
+
+from sglang.multimodal_gen.runtime.layers.linear import (
+    LinearBase,
+    LinearMethodBase,
+    UnquantizedLinearMethod,
+)
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+    QuantizeMethodBase,
+)
+from sglang.multimodal_gen.runtime.utils.weight_attrs import set_weight_attrs
+
+
+def _require_bitsandbytes() -> None:
+    try:
+        import bitsandbytes
+
+        if version.parse(bitsandbytes.__version__) < version.parse("0.46.1"):
+            raise ImportError(
+                "bitsandbytes version is wrong. Please install bitsandbytes>=0.46.1."
+            )
+    except ImportError as err:
+        raise ImportError(
+            "Please install bitsandbytes>=0.46.1 via "
+            "`pip install bitsandbytes>=0.46.1` to use bitsandbytes quantizer."
+        ) from err
+
+
+def _calculate_quant_ratio(dtype: torch.dtype) -> int:
+    if dtype.is_floating_point:
+        return torch.finfo(dtype).bits // torch.iinfo(torch.uint8).bits
+    return torch.iinfo(dtype).bits // torch.iinfo(torch.uint8).bits
+
+
+def _is_layer_skipped(prefix: str, skipped_modules: list[str]) -> bool:
+    components = prefix.split(".")
+    if any(module_name in components for module_name in skipped_modules):
+        return True
+
+    prefixes = {".".join(components[: i + 1]) for i in range(len(components))}
+    return bool(set(skipped_modules) & prefixes)
+
+
+class BitsAndBytesConfig(QuantizationConfig):
+    """Config class for pre-quantized bitsandbytes 4-bit checkpoints."""
+
+    def __init__(
+        self,
+        load_in_8bit: bool = False,
+        load_in_4bit: bool = True,
+        bnb_4bit_compute_dtype: str = "float32",
+        bnb_4bit_quant_storage: str = "uint8",
+        bnb_4bit_quant_type: str = "fp4",
+        bnb_4bit_use_double_quant: bool = False,
+        llm_int8_enable_fp32_cpu_offload: bool = False,
+        llm_int8_has_fp16_weight: bool = False,
+        llm_int8_skip_modules: list[str] | None = None,
+        llm_int8_threshold: float = 6.0,
+    ) -> None:
+        super().__init__()
+        self.load_in_8bit = load_in_8bit
+        self.load_in_4bit = load_in_4bit
+        self.bnb_4bit_compute_dtype = bnb_4bit_compute_dtype
+        self.bnb_4bit_quant_storage = bnb_4bit_quant_storage
+        self.bnb_4bit_quant_type = bnb_4bit_quant_type
+        self.bnb_4bit_use_double_quant = bnb_4bit_use_double_quant
+        self.llm_int8_enable_fp32_cpu_offload = llm_int8_enable_fp32_cpu_offload
+        self.llm_int8_has_fp16_weight = llm_int8_has_fp16_weight
+        self.llm_int8_skip_modules = llm_int8_skip_modules or []
+        self.llm_int8_threshold = llm_int8_threshold
+
+        if self.load_in_8bit or not self.load_in_4bit:
+            raise ValueError("SGLang diffusion only supports bitsandbytes 4-bit.")
+        if self.bnb_4bit_quant_storage != "uint8":
+            raise ValueError(
+                f"Unsupported bnb_4bit_quant_storage: {self.bnb_4bit_quant_storage}"
+            )
+
+    @classmethod
+    def get_name(cls) -> str:
+        return "bitsandbytes"
+
+    def get_scaled_act_names(self) -> list[str]:
+        return []
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> list[torch.dtype]:
+        return [torch.float32, torch.float16, torch.bfloat16]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 70
+
+    @staticmethod
+    def get_config_filenames() -> list[str]:
+        return []
+
+    @classmethod
+    def from_config(cls, config: dict[str, Any]) -> BitsAndBytesConfig:
+        def get_safe_value(keys, default_value=None):
+            try:
+                value = QuantizationConfig.get_from_keys(config, keys)
+                return value if value is not None else default_value
+            except ValueError:
+                return default_value
+
+        return cls(
+            load_in_8bit=get_safe_value(["load_in_8bit"], False),
+            load_in_4bit=get_safe_value(["load_in_4bit"], True),
+            bnb_4bit_compute_dtype=get_safe_value(
+                ["bnb_4bit_compute_dtype"], "float32"
+            ),
+            bnb_4bit_quant_storage=get_safe_value(["bnb_4bit_quant_storage"], "uint8"),
+            bnb_4bit_quant_type=get_safe_value(["bnb_4bit_quant_type"], "fp4"),
+            bnb_4bit_use_double_quant=get_safe_value(
+                ["bnb_4bit_use_double_quant"], False
+            ),
+            llm_int8_enable_fp32_cpu_offload=get_safe_value(
+                ["llm_int8_enable_fp32_cpu_offload"], False
+            ),
+            llm_int8_has_fp16_weight=get_safe_value(
+                ["llm_int8_has_fp16_weight"], False
+            ),
+            llm_int8_skip_modules=get_safe_value(["llm_int8_skip_modules"], []),
+            llm_int8_threshold=get_safe_value(["llm_int8_threshold"], 6.0),
+        )
+
+    def get_quant_method(
+        self, layer: torch.nn.Module, prefix: str
+    ) -> Optional[QuantizeMethodBase]:
+        if isinstance(layer, LinearBase):
+            if _is_layer_skipped(prefix, self.llm_int8_skip_modules):
+                return UnquantizedLinearMethod()
+            return BitsAndBytesLinearMethod(self)
+        return None
+
+
+class BitsAndBytesLinearMethod(LinearMethodBase):
+    """Linear method for pre-quantized bitsandbytes 4-bit weights."""
+
+    def __init__(self, quant_config: BitsAndBytesConfig):
+        _require_bitsandbytes()
+        self.quant_config = quant_config
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes: list[int],
+        input_size: int,
+        output_size: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ) -> None:
+        quant_ratio = _calculate_quant_ratio(params_dtype)
+        output_size_per_partition = sum(output_partition_sizes)
+        total_size = input_size_per_partition * output_size_per_partition
+        if total_size % quant_ratio != 0:
+            raise ValueError(
+                "The input size is not aligned with the quantized weight shape."
+            )
+
+        qweight = nn.Parameter(
+            torch.empty(total_size // quant_ratio, 1, dtype=torch.uint8),
+            requires_grad=False,
+        )
+        set_weight_attrs(
+            qweight,
+            {
+                "input_dim": 0,
+                "output_dim": 0,
+                "pack_factor": quant_ratio,
+                "use_bitsandbytes_4bit": True,
+                "bnb_full_shape": (output_size, input_size),
+                "bnb_local_shape": (
+                    output_size_per_partition,
+                    input_size_per_partition,
+                ),
+                "bnb_output_shard_start": getattr(layer, "tp_rank", 0)
+                * output_size_per_partition,
+                "bnb_input_shard_start": (
+                    0
+                    if input_size_per_partition == input_size
+                    else getattr(layer, "tp_rank", 0) * input_size_per_partition
+                ),
+            },
+        )
+        layer.register_parameter("weight", qweight)
+        set_weight_attrs(qweight, extra_weight_attrs)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        original_type = x.dtype
+        original_shape = x.shape
+        if x.ndim > 2:
+            x = x.reshape(-1, x.size(-1))
+
+        out_dim = sum(
+            quant_state.shape[0]
+            for quant_state in layer.weight.bnb_quant_state.values()
+        )
+        out = torch.empty(x.shape[0], out_dim, dtype=torch.bfloat16, device=x.device)
+        apply_bnb_4bit(x.to(torch.bfloat16), layer.weight, out)
+        out = out.to(original_type)
+
+        if len(original_shape) > 2:
+            out = out.view(*original_shape[:-1], out.size(-1))
+
+        if bias is not None:
+            out = out + bias
+        return out
+
+
+def apply_bnb_4bit(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    from bitsandbytes import matmul_4bit
+
+    offsets = weight.bnb_shard_offsets
+    quant_states = weight.bnb_quant_state
+    current_index = 0
+    for i in range(len(quant_states)):
+        output_size = quant_states[i].shape[0]
+        out[:, current_index : current_index + output_size] = matmul_4bit(
+            x,
+            weight[offsets[i] : offsets[i + 1]].t(),
+            quant_states[i],
+        )
+        current_index += output_size
+
+
+class BitsAndBytes4BitLinear(nn.Module):
+    """Storage-only bitsandbytes 4-bit linear for nn.Linear-based encoders."""
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        compute_dtype: torch.dtype | None = None,
+    ) -> None:
+        super().__init__()
+        _require_bitsandbytes()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.compute_dtype = compute_dtype
+        quant_ratio = _calculate_quant_ratio(compute_dtype or torch.get_default_dtype())
+        total_size = in_features * out_features
+        if total_size % quant_ratio != 0:
+            raise ValueError(
+                "The input size is not aligned with the quantized weight shape."
+            )
+
+        self.weight = nn.Parameter(
+            torch.empty(total_size // quant_ratio, 1, dtype=torch.uint8),
+            requires_grad=False,
+        )
+        set_weight_attrs(
+            self.weight,
+            {
+                "pack_factor": quant_ratio,
+                "use_bitsandbytes_4bit": True,
+            },
+        )
+        if bias:
+            self.bias = nn.Parameter(
+                torch.empty(
+                    out_features, dtype=compute_dtype or torch.get_default_dtype()
+                ),
+                requires_grad=False,
+            )
+        else:
+            self.register_parameter("bias", None)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        original_type = x.dtype
+        original_shape = x.shape
+        if x.ndim > 2:
+            x = x.reshape(-1, x.size(-1))
+
+        out = torch.empty(
+            x.shape[0], self.out_features, dtype=torch.bfloat16, device=x.device
+        )
+        apply_bnb_4bit(x.to(torch.bfloat16), self.weight, out)
+        out = out.to(original_type)
+
+        if len(original_shape) > 2:
+            out = out.view(*original_shape[:-1], out.size(-1))
+
+        if self.bias is not None:
+            out = out + self.bias
+        return out
+
+
+def swap_linears_to_bitsandbytes_4bit(module: nn.Module) -> None:
+    for name, child in list(module.named_children()):
+        if isinstance(child, nn.Linear):
+            replacement = BitsAndBytes4BitLinear(
+                child.in_features,
+                child.out_features,
+                bias=child.bias is not None,
+                compute_dtype=child.weight.dtype,
+            )
+            setattr(module, name, replacement)
+        else:
+            swap_linears_to_bitsandbytes_4bit(child)
+
+
+_BNB_4BIT_STATE_SUFFIXES = {
+    "absmax",
+    "quant_map",
+    "nested_absmax",
+    "nested_quant_map",
+    "bitsandbytes",
+}
+
+
+def is_bitsandbytes_4bit_state_name(weight_name: str) -> bool:
+    suffix = weight_name.split(".")[-1]
+    return any(state_suffix in suffix for state_suffix in _BNB_4BIT_STATE_SUFFIXES)
+
+
+def split_bitsandbytes_4bit_state(
+    weights: Any,
+) -> tuple[list[tuple[str, torch.Tensor]], dict[str, torch.Tensor]]:
+    normal_weights: list[tuple[str, torch.Tensor]] = []
+    quant_state_dict: dict[str, torch.Tensor] = {}
+    for name, tensor in weights:
+        if is_bitsandbytes_4bit_state_name(name):
+            if "quant_state.bitsandbytes" in name:
+                tensor = tensor.cpu().data
+            quant_state_dict[name] = tensor
+            continue
+        normal_weights.append((name, tensor))
+    return normal_weights, quant_state_dict
+
+
+def build_bitsandbytes_4bit_quant_states(
+    normal_weight_names: list[str],
+    quant_state_dict: dict[str, torch.Tensor],
+    device: torch.device,
+    param_names_mapping=None,
+) -> dict[str, Any]:
+    from bitsandbytes.functional import QuantState
+
+    quant_states: dict[str, Any] = {}
+    device_str = str(device)
+    for source_name in normal_weight_names:
+        if (
+            f"{source_name}.quant_state.bitsandbytes__nf4" not in quant_state_dict
+            and f"{source_name}.quant_state.bitsandbytes__fp4" not in quant_state_dict
+        ):
+            continue
+        target_name = source_name
+        if param_names_mapping is not None:
+            target_name, _, _ = param_names_mapping(source_name)
+        state_tensors = {
+            name: tensor
+            for name, tensor in quant_state_dict.items()
+            if name.startswith(f"{source_name}.")
+        }
+        quant_states[target_name] = QuantState.from_dict(
+            state_tensors, device=device_str
+        )
+    return quant_states
+
+
+def attach_bitsandbytes_4bit_quant_states(
+    params_dict: dict[str, torch.nn.Parameter],
+    quant_states: dict[str, Any],
+) -> None:
+    for param_name, quant_state in quant_states.items():
+        param = params_dict.get(param_name)
+        if param is None:
+            raise ValueError(f"Parameter {param_name} not found in the model.")
+
+        quant_state = _maybe_shard_bitsandbytes_4bit_quant_state(param, quant_state)
+        state_by_shard = {0: quant_state}
+        set_weight_attrs(param, {"bnb_quant_state": state_by_shard})
+        offsets = torch.tensor([0, param.numel()]).cpu()
+        set_weight_attrs(param, {"bnb_shard_offsets": offsets})
+
+
+def _maybe_shard_bitsandbytes_4bit_quant_state(
+    param: torch.nn.Parameter,
+    quant_state: Any,
+) -> Any:
+    full_shape = tuple(getattr(param, "bnb_full_shape", tuple(quant_state.shape or ())))
+    local_shape = tuple(getattr(param, "bnb_local_shape", full_shape))
+    if not full_shape or local_shape == full_shape:
+        return quant_state
+
+    output_start = getattr(param, "bnb_output_shard_start", 0)
+    input_start = getattr(param, "bnb_input_shard_start", 0)
+    if input_start != 0 or local_shape[1] != full_shape[1]:
+        raise NotImplementedError(
+            "bitsandbytes 4-bit TP only supports column-parallel output shards."
+        )
+    if getattr(quant_state, "nested", False):
+        raise NotImplementedError(
+            "bitsandbytes 4-bit TP does not support nested quant states."
+        )
+
+    blocksize = quant_state.blocksize
+    start_elem = output_start * full_shape[1]
+    local_numel = local_shape[0] * local_shape[1]
+    if start_elem % blocksize != 0 or local_numel % blocksize != 0:
+        raise ValueError(
+            "bitsandbytes 4-bit TP shard is not aligned to quantization blocks."
+        )
+    start_block = start_elem // blocksize
+    num_blocks = local_numel // blocksize
+    return type(quant_state)(
+        absmax=quant_state.absmax.narrow(0, start_block, num_blocks).contiguous(),
+        shape=torch.Size(local_shape),
+        code=quant_state.code,
+        blocksize=quant_state.blocksize,
+        quant_type=quant_state.quant_type,
+        dtype=quant_state.dtype,
+        offset=None,
+        state2=None,
+    )

python/sglang/multimodal_gen/runtime/layers/quantization/configs/base_config.py

@@ -0,0 +1,155 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/quantization/base_config.py
+
+import inspect
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any
+
+import torch
+from torch import nn
+
+if TYPE_CHECKING:
+    from sglang.multimodal_gen.runtime.layers.quantization import QuantizationMethods
+else:
+    QuantizationMethods = str
+
+
+class QuantizeMethodBase(ABC):
+    """Base class for different quantized methods."""
+
+    @abstractmethod
+    def create_weights(
+        self, layer: torch.nn.Module, *weight_args, **extra_weight_attrs
+    ):
+        """Create weights for a layer.
+
+        The weights will be set as attributes of the layer."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def apply(self, layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor:
+        """Apply the weights in layer to the input tensor.
+
+        Expects create_weights to have been called before on the layer."""
+        raise NotImplementedError
+
+    # Not required functions
+    def embedding(self, layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor:
+        """Gather embeddings in the layer based on indices in the input tensor.
+
+        Expects create_weights to have been called before on the layer."""
+        raise NotImplementedError
+
+    def process_weights_after_loading(self, layer: nn.Module) -> None:
+        """Process the weight after loading.
+
+        This can be used for example, to transpose weights for computation.
+        """
+        return
+
+
+def method_has_implemented_embedding(method_class: type[QuantizeMethodBase]) -> bool:
+    """
+    Not all quant methods have embedding implemented, so we need to check that
+    it exists for our given method. We check this by making sure the function
+    has been changed from the base implementation.
+    """
+    base_embedding = inspect.getattr_static(QuantizeMethodBase, "embedding", None)
+    class_embedding = inspect.getattr_static(method_class, "embedding", None)
+
+    return class_embedding is not None and class_embedding is not base_embedding
+
+
+class QuantizationConfig(ABC):
+    """Base class for quantization configs."""
+
+    # for quantization frameworks with a separate quantized model provided, e.g. Nunchaku
+    quantized_model_path: str | None = None
+
+    def __init__(self):
+        super().__init__()
+        # mapping is updated by models as they initialize
+        self.packed_modules_mapping: dict[str, list[str]] = dict()
+
+    @abstractmethod
+    def get_name(self) -> QuantizationMethods:
+        """Name of the quantization method."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_supported_act_dtypes(self) -> list[torch.dtype]:
+        """List of supported activation dtypes."""
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def get_min_capability(cls) -> int:
+        """Minimum GPU capability to support the quantization method.
+
+        E.g., 70 for Volta, 75 for Turing, 80 for Ampere.
+        This requirement is due to the custom CUDA kernels used by the
+        quantization method.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def get_config_filenames() -> list[str]:
+        """List of filenames to search for in the model directory."""
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def from_config(cls, config: dict[str, Any]) -> "QuantizationConfig":
+        """Create a config class from the model's quantization config."""
+        raise NotImplementedError
+
+    @classmethod
+    def override_quantization_method(
+        cls, hf_quant_cfg, user_quant
+    ) -> QuantizationMethods | None:
+        """
+        Detects if this quantization method can support a given checkpoint
+        format by overriding the user specified quantization method --
+        this method should only be overwritten by subclasses in exceptional
+        circumstances
+        """
+        return None
+
+    @staticmethod
+    def get_from_keys(config: dict[str, Any], keys: list[str]) -> Any:
+        """Get a value from the model's quantization config."""
+        for key in keys:
+            if key in config:
+                return config[key]
+        raise ValueError(
+            f"Cannot find any of {keys} in the model's " "quantization config."
+        )
+
+    @staticmethod
+    def get_from_keys_or(config: dict[str, Any], keys: list[str], default: Any) -> Any:
+        """Get a optional value from the model's quantization config."""
+        try:
+            return QuantizationConfig.get_from_keys(config, keys)
+        except ValueError:
+            return default
+
+    @abstractmethod
+    def get_quant_method(
+        self, layer: torch.nn.Module, prefix: str
+    ) -> QuantizeMethodBase | None:
+        """Get the quantize method to use for the quantized layer.
+
+        Args:
+            layer: The layer for the quant method.
+            prefix: The full name of the layer in the state dict
+        Returns:
+            The quantize method. None if the given layer doesn't support quant
+            method.
+        """
+        raise NotImplementedError
+
+    def get_cache_scale(self, name: str) -> str | None:
+        return None

python/sglang/multimodal_gen/runtime/layers/quantization/configs/nunchaku_config.py

@@ -0,0 +1,283 @@
+# SPDX-License-Identifier: Apache-2.0
+import json
+import os
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import Any, Optional
+
+import torch
+from safetensors.torch import load_file as safetensors_load_file
+from torch import nn
+
+from sglang.multimodal_gen.runtime.layers.linear import LinearBase
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+from .base_config import QuantizationConfig, QuantizeMethodBase
+
+logger = init_logger(__name__)
+
+
+@lru_cache(maxsize=1)
+def is_nunchaku_available() -> bool:
+    try:
+        import nunchaku  # noqa
+
+        logger.debug("Nunchaku package detected")
+        return True
+    except Exception:
+        return False
+
+
+@dataclass
+class NunchakuConfig(QuantizationConfig):
+    """
+    Configuration for Nunchaku (SVDQuant) W4A4-style quantization.
+
+    Attributes:
+        precision: Quantization precision type. Options:
+            - "int4": Standard INT4 quantization
+            - "nvfp4": FP4 quantization
+        rank: SVD low-rank dimension for absorbing outliers
+        group_size: Quantization group size (automatically set based on precision)
+        act_unsigned: Use unsigned activation quantization
+        transformer_weights_path: Path to pre-quantized transformer weights (.safetensors)
+        model_cls: DiT model class that provides quantization rules via get_nunchaku_quant_rules()
+    """
+
+    precision: str = "int4"
+    rank: int = 32
+    group_size: Optional[int] = None
+    act_unsigned: bool = False
+    transformer_weights_path: Optional[str] = None
+    model_cls: Optional[type] = None
+
+    @classmethod
+    def get_name(cls) -> str:
+        return "svdquant"
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> list[torch.dtype]:
+        return [torch.bfloat16, torch.float16]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 70
+
+    @staticmethod
+    def get_config_filenames() -> list[str]:
+        return ["quantization_config.json", "quant_config.json"]
+
+    @classmethod
+    def from_config(cls, config: dict[str, Any]) -> "NunchakuConfig":
+
+        return cls(
+            precision=config.get("precision", "int4"),
+            rank=int(config.get("rank", 32)),
+            group_size=config.get("group_size"),
+            act_unsigned=bool(config.get("act_unsigned", False)),
+            transformer_weights_path=config.get("transformer_weights_path"),
+        )
+
+    def get_quant_method(
+        self, layer: torch.nn.Module, prefix: str
+    ) -> Optional[QuantizeMethodBase]:
+        if not isinstance(layer, LinearBase):
+            return None
+
+        # get quantization rules from model class
+        quant_rules = self._get_quant_rules()
+
+        # priority: skip > awq_w4a16 > svdq_w4a4 > default
+        skip_patterns = quant_rules.get("skip", [])
+        for pattern in skip_patterns:
+            if pattern in prefix.lower():
+                return None
+
+        awq_patterns = quant_rules.get("awq_w4a16", [])
+        for pattern in awq_patterns:
+            if pattern in prefix:
+                from ..nunchaku_linear import NunchakuAWQLinearMethod
+
+                return NunchakuAWQLinearMethod(group_size=64)
+
+        svdq_patterns = quant_rules.get("svdq_w4a4", [])
+        for pattern in svdq_patterns:
+            if pattern in prefix:
+                from ..nunchaku_linear import NunchakuSVDQLinearMethod
+
+                return NunchakuSVDQLinearMethod(
+                    precision=self.precision,
+                    rank=self.rank,
+                    act_unsigned=self.act_unsigned,
+                )
+
+        # default: apply svdq_w4a4 to all remaining linear layers
+        from ..nunchaku_linear import NunchakuSVDQLinearMethod
+
+        return NunchakuSVDQLinearMethod(
+            precision=self.precision,
+            rank=self.rank,
+            act_unsigned=self.act_unsigned,
+        )
+
+    def _get_quant_rules(self) -> dict[str, list[str]]:
+        if self.model_cls is not None and hasattr(
+            self.model_cls, "get_nunchaku_quant_rules"
+        ):
+            return self.model_cls.get_nunchaku_quant_rules()
+        return {}
+
+    def __post_init__(self):
+        if self.group_size is None:
+            if self.precision == "nvfp4":
+                self.group_size = 16
+            elif self.precision == "int4":
+                self.group_size = 64
+            else:
+                raise ValueError(
+                    f"Invalid precision: {self.precision}. Must be 'int4' or 'nvfp4'"
+                )
+
+        if self.precision not in ["int4", "nvfp4"]:
+            raise ValueError(
+                f"Invalid precision: {self.precision}. Must be 'int4' or 'nvfp4'"
+            )
+
+        if self.rank <= 0:
+            raise ValueError(f"Rank must be positive, got {self.rank}")
+
+    @classmethod
+    def from_dict(cls, config_dict: dict) -> "NunchakuConfig":
+        """Create configuration from dictionary."""
+        return cls(**config_dict)
+
+    def to_dict(self) -> dict:
+        """Convert configuration to dictionary."""
+        return {
+            "precision": self.precision,
+            "rank": self.rank,
+            "group_size": self.group_size,
+            "act_unsigned": self.act_unsigned,
+            "transformer_weights_path": self.transformer_weights_path,
+        }
+
+    @classmethod
+    def from_pretrained(cls, model_path: str) -> Optional["NunchakuConfig"]:
+        for filename in cls.get_config_filenames():
+            config_path = os.path.join(model_path, filename)
+            if os.path.exists(config_path):
+                with open(config_path, "r") as f:
+                    config_dict = json.load(f)
+                if config_dict.get("quant_method") == cls.get_name():
+                    return cls.from_config(config_dict)
+        return None
+
+
+def _patch_native_svdq_linear(
+    module: nn.Module, tensor: Any, svdq_linear_cls: type
+) -> bool:
+    if (
+        isinstance(module, svdq_linear_cls)
+        and getattr(module, "wtscale", None) is not None
+    ):
+        module.wtscale = tensor
+        return True
+    return False
+
+
+def _patch_sglang_svdq_linear(
+    module: nn.Module, tensor: Any, svdq_method_cls: type
+) -> bool:
+    quant_method = getattr(module, "quant_method", None)
+    if not isinstance(quant_method, svdq_method_cls):
+        return False
+
+    existing = getattr(module, "wtscale", None)
+    if isinstance(existing, nn.Parameter):
+        with torch.no_grad():
+            existing.data.copy_(tensor.to(existing.data.dtype))
+    else:
+        module.wtscale = tensor
+
+    # Keep alpha in sync (kernel reads `layer._nunchaku_alpha`)
+    try:
+        module._nunchaku_alpha = float(tensor.detach().cpu().item())
+    except Exception:
+        module._nunchaku_alpha = None
+    return True
+
+
+def _patch_sglang_svdq_wcscales(
+    module: nn.Module, tensor: Any, svdq_method_cls: type
+) -> bool:
+    quant_method = getattr(module, "quant_method", None)
+    if not isinstance(quant_method, svdq_method_cls):
+        return False
+
+    existing = getattr(module, "wcscales", None)
+    if isinstance(existing, nn.Parameter):
+        with torch.no_grad():
+            existing.data.copy_(tensor.to(existing.data.dtype))
+    else:
+        module.wcscales = tensor
+    return True
+
+
+def _patch_nunchaku_scales(
+    model: nn.Module,
+    safetensors_list: list[str],
+) -> None:
+    """Patch transformer module with Nunchaku scale tensors from safetensors weights.
+
+    For NVFP4 checkpoints, correctness depends on `wtscale` and attention
+    `wcscales`. The FSDP loader may skip some of these metadata tensors.
+    """
+
+    if not safetensors_list:
+        return
+
+    if len(safetensors_list) != 1:
+        logger.warning(
+            "Nunchaku scale patch expects a single safetensors file, "
+            "but got %d files. Skipping.",
+            len(safetensors_list),
+        )
+        return
+
+    from nunchaku.models.linear import SVDQW4A4Linear  # type: ignore[import]
+
+    state_dict = safetensors_load_file(safetensors_list[0])
+    if state_dict is None:
+        return
+
+    num_wtscale = 0
+    num_wcscales = 0
+
+    from ..nunchaku_linear import NunchakuSVDQLinearMethod
+
+    for name, module in model.named_modules():
+        wt = state_dict.get(f"{name}.wtscale")
+        if wt is not None:
+            if _patch_native_svdq_linear(module, wt, SVDQW4A4Linear):
+                num_wtscale += 1
+            elif _patch_sglang_svdq_linear(module, wt, NunchakuSVDQLinearMethod):
+                num_wtscale += 1
+
+        wc = state_dict.get(f"{name}.wcscales")
+        if wc is not None:
+            # Some modules may have wcscales as a direct attribute/Parameter.
+            existing = getattr(module, "wcscales", None)
+            if isinstance(existing, nn.Parameter):
+                with torch.no_grad():
+                    existing.data.copy_(wc.to(existing.data.dtype))
+                num_wcscales += 1
+            elif existing is not None:
+                setattr(module, "wcscales", wc)
+                num_wcscales += 1
+            elif _patch_sglang_svdq_wcscales(module, wc, NunchakuSVDQLinearMethod):
+                num_wcscales += 1
+
+    if num_wtscale > 0:
+        logger.info("Patched wtscale for %d layers", num_wtscale)
+    if num_wcscales > 0:
+        logger.info("Patched wcscales for %d layers", num_wcscales)

python/sglang/multimodal_gen/runtime/layers/quantization/fp8.py

@@ -0,0 +1,508 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+
+import torch
+from torch.nn import Module
+from torch.nn.parameter import Parameter
+
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    get_tensor_model_parallel_world_size,
+)
+from sglang.multimodal_gen.runtime.layers.linear import (
+    LinearMethodBase,
+    UnquantizedLinearMethod,
+)
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+    QuantizeMethodBase,
+)
+from sglang.multimodal_gen.runtime.models.parameter import (
+    BlockQuantScaleParameter,
+    ModelWeightParameter,
+    PerTensorScaleParameter,
+)
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.multimodal_gen.runtime.platforms.aiter import USE_AITER
+from sglang.multimodal_gen.runtime.utils.common import (
+    cpu_has_amx_support,
+    get_bool_env_var,
+    use_intel_amx_backend,
+)
+from sglang.srt.layers.amx_utils import _amx_process_weight_after_loading
+from sglang.srt.layers.quantization.fp8_kernel import (
+    is_fp8_fnuz,
+    per_token_group_quant_fp8,
+)
+from sglang.srt.layers.quantization.fp8_utils import (
+    apply_fp8_linear,
+    can_auto_enable_marlin_fp8,
+    cutlass_fp8_supported,
+    dispatch_w8a8_block_fp8_linear,
+    input_to_float8,
+    normalize_e4m3fn_to_e4m3fnuz,
+    requant_weight_ue8m0_inplace,
+)
+from sglang.srt.layers.quantization.marlin_utils_fp8 import (
+    apply_fp8_marlin_linear,
+    prepare_fp8_layer_for_marlin,
+)
+from sglang.srt.layers.quantization.utils import (
+    convert_to_channelwise,
+    is_layer_skipped,
+    requantize_with_max_scale,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.quantization.w4afp8 import W4AFp8Config
+
+_is_hip = current_platform.is_hip()
+_is_cuda = current_platform.is_cuda()
+_is_npu = current_platform.is_npu()
+_is_cpu_amx_available = cpu_has_amx_support()
+_is_cpu = current_platform.is_cpu()
+_is_fp8_fnuz = is_fp8_fnuz()
+_use_hip_int4 = get_bool_env_var("SGLANG_INT4_WEIGHT") and _is_hip
+
+if USE_AITER or _use_hip_int4:
+    pass
+
+
+ACTIVATION_SCHEMES = ["static", "dynamic"]
+
+logger = logging.getLogger(__name__)
+
+
+class Fp8Config(QuantizationConfig):
+    """Config class for FP8.
+
+    No-arg ``Fp8Config()`` selects online (post-load) weight quantization:
+    ``is_checkpoint_fp8_serialized=False`` with ``activation_scheme="dynamic"``.
+    """
+
+    def __init__(
+        self,
+        is_checkpoint_fp8_serialized: bool = False,
+        activation_scheme: str = "dynamic",
+        ignored_layers: Optional[List[str]] = None,
+        weight_block_size: List[int] = None,
+        packed_modules_mapping: Optional[Dict[str, List[str]]] = None,
+    ) -> None:
+        self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized
+        if is_checkpoint_fp8_serialized:
+            logger.info("Detected fp8 checkpoint.")
+        if activation_scheme not in ACTIVATION_SCHEMES:
+            raise ValueError(f"Unsupported activation scheme {activation_scheme}")
+        self.activation_scheme = activation_scheme
+        self.ignored_layers = ignored_layers or []
+        self.packed_modules_mapping = packed_modules_mapping or {}
+        if weight_block_size is not None:
+            if not is_checkpoint_fp8_serialized:
+                raise ValueError(
+                    "The block-wise quantization only supports fp8-serialized checkpoint for now."
+                )
+            if len(weight_block_size) != 2:
+                raise ValueError(
+                    f"The quantization block size of weight must have 2 dimensions, but got {len(weight_block_size)} dimensions."
+                )
+            if activation_scheme != "dynamic":
+                raise ValueError(
+                    f"The block-wise quantization only supports dynamic activation scheme for now, but got {activation_scheme} activation scheme."
+                )
+        self.weight_block_size = weight_block_size
+
+    @classmethod
+    def get_name(cls) -> str:
+        return "fp8"
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> List[torch.dtype]:
+        return [torch.bfloat16, torch.half]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 80
+
+    @classmethod
+    def get_config_filenames(cls) -> List[str]:
+        return []
+
+    @classmethod
+    def from_config(cls, config: Dict[str, Any]) -> Fp8Config:
+        quant_method = cls.get_from_keys(config, ["quant_method"])
+        is_checkpoint_fp8_serialized = "fp8" in quant_method
+        activation_scheme = cls.get_from_keys(config, ["activation_scheme"])
+        ignored_layers = cls.get_from_keys_or(
+            config, ["ignored_layers", "modules_to_not_convert"], None
+        )
+        if ignored_layers:
+            # hacking ministral
+            ignored_layers = [layer.replace("model.", "") for layer in ignored_layers]
+        weight_block_size = cls.get_from_keys_or(config, ["weight_block_size"], None)
+        return cls(
+            is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized,
+            activation_scheme=activation_scheme,
+            ignored_layers=ignored_layers,
+            weight_block_size=weight_block_size,
+        )
+
+    def get_quant_method(
+        self, layer: torch.nn.Module, prefix: str
+    ) -> Optional[QuantizeMethodBase]:
+        from sglang.multimodal_gen.runtime.layers.linear import LinearBase
+
+        if isinstance(layer, LinearBase):
+            if is_layer_skipped(
+                prefix,
+                self.ignored_layers,
+                fused_mapping=self.packed_modules_mapping,
+            ):
+                return UnquantizedLinearMethod()
+            return Fp8LinearMethod(self)
+        return None
+
+    def get_scaled_act_names(self) -> List[str]:
+        return []
+
+
+class Fp8LinearMethod(LinearMethodBase):
+    """Linear method for FP8.
+    Supports loading FP8 checkpoints with static weight scale and
+    dynamic/static activation scale.
+
+    Also supports loading quantized FP16/BF16 model checkpoints with dynamic
+    activation scaling. The weight scaling factor will be initialized after
+    the model weights are loaded.
+
+    Limitations:
+    1. Only support per-tensor quantization due to torch._scaled_mm support.
+    2. Only support float8_e4m3fn data type due to the limitation of
+       torch._scaled_mm (https://github.com/pytorch/pytorch/blob/2e48b39603411a41c5025efbe52f89560b827825/aten/src/ATen/native/cuda/Blas.cpp#L854-L856)
+
+    Args:
+        quant_config: The quantization config.
+    """
+
+    def __init__(self, quant_config: Union[Fp8Config, W4AFp8Config]):
+        self.quant_config = quant_config
+        self.cutlass_fp8_supported = cutlass_fp8_supported()
+
+        # For GPUs that lack FP8 hardware support, we can leverage the Marlin
+        # kernel for fast weight-only FP8 quantization
+        self.use_marlin = False
+        if _is_cuda:
+            force_marlin = get_bool_env_var("SGLANG_FORCE_FP8_MARLIN")
+            auto_enable = can_auto_enable_marlin_fp8()
+            self.use_marlin = force_marlin or auto_enable
+
+        self.block_quant = self.quant_config.weight_block_size is not None
+
+        self.w8a8_block_fp8_linear = dispatch_w8a8_block_fp8_linear()
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes: List[int],
+        input_size: int,
+        output_size: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        output_size_per_partition = sum(output_partition_sizes)
+        weight_loader = extra_weight_attrs.get("weight_loader")
+
+        tp_size = get_tensor_model_parallel_world_size()
+        if self.block_quant:
+            block_n, block_k = (
+                self.quant_config.weight_block_size[0],
+                self.quant_config.weight_block_size[1],
+            )
+            # Required by row parallel
+            if tp_size > 1 and input_size // input_size_per_partition == tp_size:
+                if input_size_per_partition % block_k != 0:
+                    raise ValueError(
+                        f"Weight input_size_per_partition = "
+                        f"{input_size_per_partition} is not divisible by "
+                        f"weight quantization block_k = {block_k}."
+                    )
+            # Required by column parallel or enabling merged weights
+            if (
+                tp_size > 1 and output_size // output_size_per_partition == tp_size
+            ) or len(output_partition_sizes) > 1:
+                for output_partition_size in output_partition_sizes:
+                    if output_partition_size % block_n != 0:
+                        raise ValueError(
+                            f"Weight output_partition_size = "
+                            f"{output_partition_size} is not divisible by "
+                            f"weight quantization block_n = {block_n}."
+                        )
+
+        layer.logical_widths = output_partition_sizes
+        layer.input_size_per_partition = input_size_per_partition
+        layer.output_size_per_partition = output_size_per_partition
+        layer.orig_dtype = params_dtype
+
+        # WEIGHT
+        weight_dtype = (
+            torch.float8_e4m3fn
+            if self.quant_config.is_checkpoint_fp8_serialized
+            else params_dtype
+        )
+
+        weight = ModelWeightParameter(
+            data=torch.empty(
+                output_size_per_partition, input_size_per_partition, dtype=weight_dtype
+            ),
+            input_dim=1,
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("weight", weight)
+
+        # If checkpoint is serialized fp8, load them.
+        # Otherwise, wait until process_weights_after_loading.
+        if self.quant_config.is_checkpoint_fp8_serialized:
+            # WEIGHT SCALE
+            if self.block_quant:
+                if hasattr(self.quant_config, "activation_scheme"):
+                    assert self.quant_config.activation_scheme == "dynamic"
+                elif hasattr(self.quant_config, "linear_activation_scheme"):
+                    assert self.quant_config.linear_activation_scheme == "dynamic"
+                scale = BlockQuantScaleParameter(
+                    data=torch.empty(
+                        (output_size_per_partition + block_n - 1) // block_n,
+                        (input_size_per_partition + block_k - 1) // block_k,
+                        dtype=torch.float32,
+                    ),
+                    input_dim=1,
+                    output_dim=0,
+                    weight_loader=weight_loader,
+                )
+                scale.format_ue8m0 = False
+                scale[:] = torch.finfo(torch.float32).min
+                layer.register_parameter("weight_scale_inv", scale)
+            else:
+                scale = PerTensorScaleParameter(
+                    data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
+                    weight_loader=weight_loader,
+                )
+                scale[:] = torch.finfo(torch.float32).min
+                layer.register_parameter("weight_scale", scale)
+
+            # INPUT ACTIVATION SCALE
+            if (
+                hasattr(self.quant_config, "activation_scheme")
+                and self.quant_config.activation_scheme == "static"
+            ) or (
+                hasattr(self.quant_config, "linear_activation_scheme")
+                and self.quant_config.linear_activation_scheme == "static"
+            ):
+                scale = PerTensorScaleParameter(
+                    data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
+                    weight_loader=weight_loader,
+                )
+
+                scale[:] = torch.finfo(torch.float32).min
+                layer.register_parameter("input_scale", scale)
+            else:
+                layer.register_parameter("input_scale", None)
+
+    def process_weights_after_loading(self, layer: Module) -> None:
+        if self.block_quant:
+            # If ROCm, normalize the weights and scales to e4m3fnuz
+            if _is_fp8_fnuz:
+                # activation_scheme: dynamic
+                weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz(
+                    weight=layer.weight,
+                    weight_scale=layer.weight_scale_inv,
+                    input_scale=None,
+                )
+                layer.input_scale = None
+            elif _is_cpu:
+                assert (
+                    _is_cpu_amx_available
+                ), "Fp8LinearMethod on CPU requires that CPU has AMX support"
+                _amx_process_weight_after_loading(layer, ["weight"])
+                layer.weight_scale_inv = torch.nn.Parameter(
+                    layer.weight_scale_inv.data, requires_grad=False
+                )
+                return
+            else:
+                # For fp8 linear weights run with deepgemm, the weights and scales need be requantized to ue8m0
+                from sglang.srt.layers.quantization.fp8_utils import (
+                    deepgemm_w8a8_block_fp8_linear_with_fallback,
+                )
+                from sglang.srt.model_loader.utils import (
+                    should_deepgemm_weight_requant_ue8m0,
+                )
+
+                if (
+                    should_deepgemm_weight_requant_ue8m0(
+                        weight_block_size=getattr(
+                            self.quant_config, "weight_block_size", None
+                        ),
+                    )
+                    and (
+                        self.w8a8_block_fp8_linear
+                        is deepgemm_w8a8_block_fp8_linear_with_fallback
+                    )
+                    and (not layer.weight_scale_inv.format_ue8m0)
+                ):
+                    requant_weight_ue8m0_inplace(
+                        layer.weight,
+                        layer.weight_scale_inv,
+                        self.quant_config.weight_block_size,
+                    )
+                    layer.weight_scale_inv.format_ue8m0 = True
+                weight, weight_scale = layer.weight.data, layer.weight_scale_inv.data
+
+            layer.weight.data = weight.data
+            layer.weight_scale_inv.data = weight_scale.data
+        else:
+            layer.weight = Parameter(layer.weight.data, requires_grad=False)
+
+            # If checkpoint not serialized fp8, quantize the weights.
+            if not self.quant_config.is_checkpoint_fp8_serialized:
+                if self.cutlass_fp8_supported or self.use_marlin:
+                    # apply per-channel quantization default as
+                    # cutlass sgl-kernel and marlin only support per-channel scale
+                    qweight, weight_scale = per_token_group_quant_fp8(
+                        layer.weight, layer.weight.shape[-1]
+                    )
+                    weight_scale = weight_scale.t().contiguous()
+                else:
+                    # per-tensor quantization
+                    qweight, weight_scale = input_to_float8(layer.weight)
+
+                # Update the layer with the new values.
+                layer.weight = Parameter(qweight.t(), requires_grad=False)
+                layer.weight_scale = Parameter(weight_scale, requires_grad=False)
+                layer.input_scale = None
+
+            # If checkpoint is fp8, handle that there are N scales for N
+            # shards in a fused module
+            else:
+                layer.weight_scale = Parameter(
+                    layer.weight_scale.data, requires_grad=False
+                )
+                if (
+                    hasattr(self.quant_config, "activation_scheme")
+                    and self.quant_config.activation_scheme == "static"
+                ) or (
+                    hasattr(self.quant_config, "linear_activation_scheme")
+                    and self.quant_config.linear_activation_scheme == "static"
+                ):
+                    layer.input_scale = Parameter(
+                        layer.input_scale.data, requires_grad=False
+                    )
+
+                # cutlass sgl-kernel and marlin only support per-channel scale
+                if self.cutlass_fp8_supported or self.use_marlin:
+                    weight = layer.weight
+                    weight_scale = convert_to_channelwise(
+                        layer.weight_scale, layer.logical_widths
+                    )
+                else:
+                    # Dequant -> Quant with max scale so we can run per tensor.
+                    weight = layer.weight
+                    weight_scale = layer.weight_scale
+                    # If ROCm, normalize the weights and scales to e4m3fnuz
+                    if _is_fp8_fnuz:
+                        weight, weight_scale, input_scale = (
+                            normalize_e4m3fn_to_e4m3fnuz(
+                                weight=weight,
+                                weight_scale=weight_scale,
+                                input_scale=layer.input_scale,
+                            )
+                        )
+                        if input_scale is not None:
+                            layer.input_scale = Parameter(
+                                input_scale, requires_grad=False
+                            )
+
+                    weight_scale, weight = requantize_with_max_scale(
+                        weight=weight,
+                        weight_scale=weight_scale,
+                        logical_widths=layer.logical_widths,
+                    )
+
+                # Update layer with new values.
+                layer.weight = Parameter(weight.t(), requires_grad=False)
+                layer.weight_scale = Parameter(weight_scale, requires_grad=False)
+                if (
+                    hasattr(self.quant_config, "activation_scheme")
+                    and self.quant_config.activation_scheme == "static"
+                ) or (
+                    hasattr(self.quant_config, "linear_activation_scheme")
+                    and self.quant_config.linear_activation_scheme == "static"
+                ):
+                    layer.input_scale = Parameter(
+                        layer.input_scale.max(), requires_grad=False
+                    )
+
+        if self.use_marlin:
+            if self.block_quant:
+                layer.weight_block_size = self.quant_config.weight_block_size
+            prepare_fp8_layer_for_marlin(layer, not self.block_quant)
+            # Activations not quantized for marlin.
+            del layer.input_scale
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if self.use_marlin:
+            return apply_fp8_marlin_linear(
+                input=x,
+                weight=layer.weight,
+                weight_scale=layer.weight_scale,
+                workspace=layer.workspace,
+                size_n=layer.output_size_per_partition,
+                size_k=layer.input_size_per_partition,
+                bias=bias,
+            )
+
+        if self.block_quant:
+            if use_intel_amx_backend(layer):
+                return torch.ops.sgl_kernel.fp8_scaled_mm_cpu(
+                    x,
+                    layer.weight,
+                    layer.weight_scale_inv,
+                    self.quant_config.weight_block_size,
+                    bias,
+                    x.dtype,
+                    True,  # is_vnni
+                )
+
+            if isinstance(x, tuple):
+                return self.w8a8_block_fp8_linear(
+                    input=x[0],
+                    weight=layer.weight,
+                    block_size=self.quant_config.weight_block_size,
+                    weight_scale=layer.weight_scale_inv,
+                    input_scale=x[1],
+                    bias=bias,
+                )
+
+            return self.w8a8_block_fp8_linear(
+                input=x,
+                weight=layer.weight,
+                block_size=self.quant_config.weight_block_size,
+                weight_scale=layer.weight_scale_inv,
+                input_scale=None,
+                bias=bias,
+            )
+
+        return apply_fp8_linear(
+            input=x,
+            weight=layer.weight,
+            weight_scale=layer.weight_scale,
+            input_scale=layer.input_scale,
+            bias=bias,
+            cutlass_fp8_supported=self.cutlass_fp8_supported,
+            use_per_token_if_dynamic=False,
+        )

python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_fp8.py

@@ -0,0 +1,210 @@
+"""ModelOpt FP8 quantization support for diffusion models.
+
+Handles checkpoints produced by NVIDIA Model Optimizer (ModelOpt) with
+``quant_algo: "FP8"`` and ``quant_method: "modelopt"``.
+
+Per quantized linear layer the checkpoint contains:
+    .weight         float8_e4m3fn  [out, in]   FP8 quantized weight
+    .weight_scale   float32        scalar       per-tensor weight scale
+    .input_scale    float32        scalar       per-tensor static activation scale
+    .bias           bfloat16       [out]        bias (unquantized)
+    ._amax          (ignored)                   calibration artifact
+
+Layers listed in the ``ignore`` field of the quantization config remain in
+bfloat16 and use the standard unquantized linear method.
+"""
+
+from __future__ import annotations
+
+import fnmatch
+import logging
+from typing import Any, Dict, List, Optional
+
+import torch
+
+from sglang.multimodal_gen.runtime.layers.linear import (
+    LinearMethodBase,
+    UnquantizedLinearMethod,
+)
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+    QuantizeMethodBase,
+)
+from sglang.multimodal_gen.runtime.models.parameter import (
+    ModelWeightParameter,
+    PerTensorScaleParameter,
+)
+from sglang.srt.layers.quantization.fp8_utils import (
+    apply_fp8_linear,
+    cutlass_fp8_supported,
+)
+from sglang.srt.layers.quantization.utils import convert_to_channelwise
+
+logger = logging.getLogger(__name__)
+
+
+class ModelOptFp8Config(QuantizationConfig):
+    """Config for ModelOpt static per-tensor FP8 quantization."""
+
+    def __init__(
+        self,
+        is_checkpoint_fp8_serialized: bool = True,
+        ignore: Optional[List[str]] = None,
+    ) -> None:
+        super().__init__()
+        self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized
+        self.ignore = ignore or []
+
+    # -- QuantizationConfig interface ----------------------------------------
+
+    @classmethod
+    def get_name(cls) -> str:
+        return "modelopt"
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> list[torch.dtype]:
+        return [torch.bfloat16, torch.half]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 89
+
+    @staticmethod
+    def get_config_filenames() -> list[str]:
+        return []
+
+    @classmethod
+    def from_config(
+        cls,
+        config: Dict[str, Any],
+        ignore_remap: Optional[Dict[str, str]] = None,
+    ) -> ModelOptFp8Config:
+        quant_algo = config.get("quant_algo")
+        if quant_algo is None:
+            raise ValueError(
+                "ModelOptFp8Config requires 'quant_algo' in the quantization config."
+            )
+        if "FP8" not in quant_algo:
+            raise ValueError(
+                f"ModelOptFp8Config only supports FP8, got quant_algo={quant_algo!r}."
+            )
+        ignore = config.get("ignore", [])
+        if ignore_remap and ignore:
+            ignore = [ignore_remap.get(pattern, pattern) for pattern in ignore]
+        return cls(is_checkpoint_fp8_serialized=True, ignore=ignore)
+
+    def _is_layer_ignored(self, prefix: str) -> bool:
+        """Check whether *prefix* matches any pattern in the ignore list.
+
+        ModelOpt ignore patterns are matched against the full prefix as a glob
+        (e.g. ``"norm_out*"`` matches ``"norm_out.linear"``) **and** against the
+        first path component (e.g. ``"proj_out"`` matches only the top-level
+        ``proj_out``, not ``single_transformer_blocks.0.proj_out``).
+        """
+        first_component = prefix.split(".")[0]
+        for pattern in self.ignore:
+            if fnmatch.fnmatch(prefix, pattern):
+                return True
+            if fnmatch.fnmatch(first_component, pattern):
+                return True
+        return False
+
+    def get_quant_method(
+        self, layer: torch.nn.Module, prefix: str
+    ) -> Optional[QuantizeMethodBase]:
+        from sglang.multimodal_gen.runtime.layers.linear import LinearBase
+
+        if isinstance(layer, LinearBase):
+            if self._is_layer_ignored(prefix):
+                return UnquantizedLinearMethod()
+            return ModelOptFp8LinearMethod(self)
+        return None
+
+    def get_scaled_act_names(self) -> list[str]:
+        return []
+
+
+class ModelOptFp8LinearMethod(LinearMethodBase):
+    """Linear method for ModelOpt static per-tensor FP8 quantization.
+
+    Uses ``torch._scaled_mm`` (or CUTLASS FP8 GEMM when available) for
+    the FP8 matrix multiply - the same kernels used by the LLM runtime.
+    """
+
+    def __init__(self, quant_config: ModelOptFp8Config):
+        self.quant_config = quant_config
+        self.cutlass_fp8_supported = cutlass_fp8_supported()
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes: List[int],
+        input_size: int,
+        output_size: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ) -> None:
+        output_size_per_partition = sum(output_partition_sizes)
+        weight_loader = extra_weight_attrs.get("weight_loader")
+
+        layer.logical_widths = output_partition_sizes
+        layer.input_size_per_partition = input_size_per_partition
+        layer.output_size_per_partition = output_size_per_partition
+
+        weight = ModelWeightParameter(
+            data=torch.empty(
+                output_size_per_partition,
+                input_size_per_partition,
+                dtype=torch.float8_e4m3fn,
+            ),
+            input_dim=1,
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("weight", weight)
+
+        for scale_name in ("weight_scale", "input_scale"):
+            scale = PerTensorScaleParameter(
+                data=torch.full(
+                    (len(output_partition_sizes),),
+                    torch.finfo(torch.float32).min,
+                    dtype=torch.float32,
+                ),
+                weight_loader=weight_loader,
+            )
+            layer.register_parameter(scale_name, scale)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        # Diffusion models use single-partition layers (no TP, no fused QKV),
+        # so we just take the max scale directly without the
+        # dequantize-requantize round-trip that the LLM path does (which
+        # requires CUDA kernels that are unavailable during CPU-phase loading).
+        max_w_scale = layer.weight_scale.max()
+
+        # Transpose weight to [in, out] column-major layout for
+        # apply_fp8_linear / CUTLASS fp8_scaled_mm. Do not call .contiguous();
+        # the kernel requires column-major stride.
+        layer.weight = torch.nn.Parameter(layer.weight.data.t(), requires_grad=False)
+
+        if self.cutlass_fp8_supported:
+            max_w_scale = convert_to_channelwise(max_w_scale, layer.logical_widths)
+        layer.weight_scale = torch.nn.Parameter(max_w_scale, requires_grad=False)
+        layer.input_scale = torch.nn.Parameter(
+            layer.input_scale.max(), requires_grad=False
+        )
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        return apply_fp8_linear(
+            input=x,
+            weight=layer.weight,
+            weight_scale=layer.weight_scale,
+            input_scale=layer.input_scale,
+            bias=bias,
+            cutlass_fp8_supported=self.cutlass_fp8_supported,
+        )

python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py

@@ -0,0 +1,683 @@
+# Adapted from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/quantization/modelopt_quant.py
+from __future__ import annotations
+
+import logging
+import re
+from functools import lru_cache
+from typing import Any, Dict, List, Optional
+
+import torch
+
+from sglang.multimodal_gen.runtime.layers.linear import (
+    LinearMethodBase,
+    UnquantizedLinearMethod,
+)
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+    QuantizeMethodBase,
+)
+from sglang.multimodal_gen.runtime.models.parameter import (
+    ModelWeightParameter,
+    PerTensorScaleParameter,
+)
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.multimodal_gen.runtime.utils.weight_attrs import set_weight_attrs
+from sglang.srt.layers.quantization.fp8_utils import (
+    apply_fp8_linear,
+    cutlass_fp8_supported,
+)
+from sglang.srt.layers.quantization.modelopt_quant import (
+    pad_nvfp4_activation_for_cutlass,
+    pad_nvfp4_weight,
+    slice_nvfp4_output,
+)
+from sglang.srt.layers.quantization.utils import (
+    convert_to_channelwise,
+    is_layer_skipped,
+    requantize_with_max_scale,
+)
+from sglang.srt.layers.utils.common import copy_or_rebind_param
+from sglang.srt.utils.common import is_flashinfer_available, round_up
+
+logger = logging.getLogger(__name__)
+
+if is_flashinfer_available():
+    import flashinfer
+else:
+    flashinfer = None
+
+
+@lru_cache(maxsize=1)
+def _get_fp4_quantize_op():
+    return current_platform.get_modelopt_fp4_quantize_op()
+
+
+@lru_cache(maxsize=1)
+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()
+
+
+def _swizzled_nvfp4_scales_to_linear(scales: torch.Tensor) -> torch.Tensor:
+    """Convert FlashInfer/CUTLASS-swizzled FP4 scales back to row-major layout."""
+    scale_ndim = scales.ndim
+    if scale_ndim == 2:
+        scales = scales.unsqueeze(0)
+    assert scales.ndim == 3
+
+    B, M, K = scales.shape
+    M_padded = round_up(M, 128)
+    K_padded = round_up(K, 4)
+    if M != M_padded or K != K_padded:
+        padded = torch.zeros(
+            (B, M_padded, K_padded), dtype=scales.dtype, device=scales.device
+        )
+        padded[:B, :M, :K] = scales
+        scales = padded
+
+    linear = scales.reshape(B, M_padded // 128, K_padded // 4, 32, 4, 4)
+    linear = linear.permute(0, 1, 4, 3, 2, 5).contiguous()
+    linear = linear.reshape(B, M_padded, K_padded)[:, :M, :K]
+    return linear.squeeze(0) if scale_ndim == 2 else linear
+
+
+def _require_flashinfer():
+    if flashinfer is None:
+        raise RuntimeError(
+            "flashinfer is required for the diffusion NVFP4 FlashInfer path."
+        )
+    return flashinfer
+
+
+class ModelOptQuantConfig(QuantizationConfig):
+    def __init__(
+        self,
+        exclude_modules: Optional[List[str]],
+        packed_modules_mapping: Optional[Dict[str, List[str]]],
+    ):
+        super().__init__()
+        self.packed_modules_mapping = packed_modules_mapping or {}
+        self.exclude_modules = exclude_modules or []
+
+    def _get_quant_method(
+        self,
+        layer: torch.nn.Module,
+        prefix: str,
+        *,
+        Linear: type[LinearMethodBase],
+    ) -> Optional[QuantizeMethodBase]:
+        from sglang.multimodal_gen.runtime.layers.linear import LinearBase
+
+        if isinstance(layer, LinearBase):
+            if self.is_layer_excluded(prefix) or (
+                self.packed_modules_mapping
+                and is_layer_skipped(prefix, [], self.packed_modules_mapping)
+            ):
+                return UnquantizedLinearMethod()
+            return Linear(self)
+        return None
+
+    @classmethod
+    def get_config_filenames(cls) -> List[str]:
+        return ["hf_quant_config.json"]
+
+    def get_scaled_act_names(self) -> List[str]:
+        return []
+
+    @classmethod
+    def override_quantization_method(cls, hf_quant_config, user_quant) -> Optional[str]:
+        if hf_quant_config is None:
+            return None
+
+        quant_algo = (
+            hf_quant_config.get("quant_algo")
+            or hf_quant_config.get("quantization", {}).get("quant_algo")
+            or ""
+        ).upper()
+        if user_quant in {"modelopt", "modelopt_fp8"} and "FP8" in quant_algo:
+            return "modelopt_fp8"
+        if user_quant in {"modelopt", "modelopt_fp4"} and (
+            "NVFP4" in quant_algo or "FP4" in quant_algo
+        ):
+            return "modelopt_fp4"
+        return None
+
+    def is_layer_excluded(self, prefix: str) -> bool:
+        for pattern in self.exclude_modules:
+            regex_str = re.escape(pattern).replace(r"\*", r".*")
+            if re.fullmatch(regex_str, prefix):
+                return True
+        return False
+
+
+class ModelOptFp8Config(ModelOptQuantConfig):
+    """Config class for ModelOpt FP8 diffusion checkpoints."""
+
+    def __init__(
+        self,
+        is_checkpoint_fp8_serialized: bool = False,
+        exclude_modules: Optional[List[str]] = None,
+        packed_modules_mapping: Optional[Dict[str, List[str]]] = None,
+    ) -> None:
+        super().__init__(exclude_modules, packed_modules_mapping)
+        self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized
+        if is_checkpoint_fp8_serialized:
+            logger.warning(
+                "Detected ModelOpt FP8 checkpoint. The format is experimental and subject to change."
+            )
+
+    @classmethod
+    def get_name(cls) -> str:
+        return "modelopt_fp8"
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> List[torch.dtype]:
+        return [torch.bfloat16, torch.half]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 89
+
+    @classmethod
+    def from_config(
+        cls,
+        config: Dict[str, Any],
+        ignore_remap: Optional[Dict[str, str]] = None,
+    ) -> ModelOptFp8Config:
+        quant_method = config.get("quant_algo")
+        exclude_modules = config.get("ignore")
+        if quant_method is None:
+            try:
+                quantization_section = cls.get_from_keys(config, ["quantization"])
+                quant_method = quantization_section.get("quant_algo")
+                exclude_modules = quantization_section.get("exclude_modules")
+            except ValueError as exc:
+                raise ValueError(
+                    "Cannot find 'quant_algo' in the model's quantization config."
+                ) from exc
+
+        if quant_method is None or "FP8" not in quant_method:
+            raise ValueError(
+                "ModelOptFp8Config only supports static FP8 quantization in SGLang diffusion."
+            )
+
+        if ignore_remap and exclude_modules:
+            exclude_modules = [ignore_remap.get(p, p) for p in exclude_modules]
+
+        return cls(
+            is_checkpoint_fp8_serialized=True,
+            exclude_modules=exclude_modules,
+            packed_modules_mapping=config.get("packed_modules_mapping"),
+        )
+
+    def get_quant_method(self, layer: torch.nn.Module, prefix: str):
+        return self._get_quant_method(layer, prefix, Linear=ModelOptFp8LinearMethod)
+
+
+class ModelOptFp4Config(ModelOptQuantConfig):
+    """Config class for NVFP4."""
+
+    def __init__(
+        self,
+        is_checkpoint_nvfp4_serialized: bool = False,
+        group_size: int = None,
+        exclude_modules: List[str] = None,
+        packed_modules_mapping: Optional[Dict[str, List[str]]] = None,
+        checkpoint_uses_packed_qkv: bool = False,
+        swap_weight_nibbles: bool = False,
+        checkpoint_weight_scale_layout: str = "linear",
+    ) -> None:
+        super().__init__(exclude_modules, packed_modules_mapping)
+        self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized
+        if is_checkpoint_nvfp4_serialized:
+            logger.warning(
+                "Detected nvfp4 checkpoint. Please note that the "
+                "format is experimental and subject to change."
+            )
+        self.group_size = group_size
+        self.checkpoint_uses_packed_qkv = checkpoint_uses_packed_qkv
+        self.swap_weight_nibbles = swap_weight_nibbles
+        self.checkpoint_weight_scale_layout = checkpoint_weight_scale_layout
+
+    @classmethod
+    def get_name(cls) -> str:
+        return "modelopt_fp4"
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> List[torch.dtype]:
+        return [torch.bfloat16, torch.half, torch.float8_e4m3fn]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 100
+
+    @staticmethod
+    def common_group_size(cfg: dict) -> int:
+        """Return the unique group_size across the config; raise if missing/mismatched."""
+        sizes = set()
+
+        def _add_group_size_from_dict(config: dict):
+            group_size = config.get("group_size")
+            if isinstance(group_size, int):
+                sizes.add(group_size)
+
+        # Top-level and 'quantization' block
+        _add_group_size_from_dict(cfg)
+        quantization = cfg.get("quantization")
+        if isinstance(quantization, dict):
+            _add_group_size_from_dict(quantization)
+
+        # config_groups: accept group-level or nested dicts (e.g., weights/input_activations)
+        for config_groups in (cfg.get("config_groups") or {}).values():
+            if isinstance(config_groups, dict):
+                _add_group_size_from_dict(config_groups)
+                for config_group in config_groups.values():
+                    if isinstance(config_group, dict):
+                        _add_group_size_from_dict(config_group)
+
+        if not sizes:
+            raise ValueError("No group_size found in config.")
+        if len(sizes) > 1:
+            raise ValueError(f"Inconsistent group_size values: {sorted(sizes)}")
+        return next(iter(sizes))
+
+    @classmethod
+    def from_config(cls, config: Dict[str, Any]) -> ModelOptFp4Config:
+        group_size = None
+        exclude_modules = []
+        swap_weight_nibbles = False
+
+        # Flat format (config.json quantization_config)
+        quant_method = config.get("quant_algo")
+        if quant_method is not None:
+            group_size = config.get("group_size")
+            if group_size is None:
+                config_groups = config.get("config_groups", {})
+                if config_groups:
+                    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",
+                config.get("checkpoint_uses_packed_qkv", False),
+            )
+        else:
+            # Nested format (hf_quant_config.json)
+            try:
+                quant_config = cls.get_from_keys(config, ["quantization"])
+                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",
+                        config.get("checkpoint_uses_packed_qkv", False),
+                    ),
+                )
+            except (ValueError, KeyError):
+                raise ValueError("Cannot find 'quant_algo' in quantization config.")
+
+        if quant_method not in ["NVFP4"]:
+            raise ValueError(
+                f"Only NVFP4 quantization is supported for diffusion, got '{quant_method}'."
+            )
+
+        if group_size is None or exclude_modules is None:
+            raise ValueError(
+                "NVFP4 quantization requires group_size and exclude_modules "
+                "in the quantization config"
+            )
+        return cls(
+            is_checkpoint_nvfp4_serialized=True,
+            group_size=group_size,
+            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,
+            checkpoint_weight_scale_layout=config.get(
+                "checkpoint_weight_scale_layout", "linear"
+            ),
+        )
+
+    def get_quant_method(self, layer: torch.nn.Module, prefix: str):
+        return self._get_quant_method(layer, prefix, Linear=ModelOptFp4LinearMethod)
+
+
+class ModelOptFp8LinearMethod(LinearMethodBase):
+    """Linear method for ModelOpt static FP8 checkpoints."""
+
+    def __init__(self, quant_config: ModelOptFp8Config):
+        self.quant_config = quant_config
+        self.cutlass_fp8_supported = cutlass_fp8_supported()
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes: List[int],
+        input_size: int,
+        output_size: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        del input_size, output_size
+        output_size_per_partition = sum(output_partition_sizes)
+        weight_loader = extra_weight_attrs.get("weight_loader")
+
+        layer.logical_widths = output_partition_sizes
+        layer.input_size_per_partition = input_size_per_partition
+        layer.output_size_per_partition = output_size_per_partition
+
+        weight_dtype = (
+            torch.float8_e4m3fn
+            if self.quant_config.is_checkpoint_fp8_serialized
+            else params_dtype
+        )
+        layer.register_parameter(
+            "weight",
+            ModelWeightParameter(
+                data=torch.empty(
+                    output_size_per_partition,
+                    input_size_per_partition,
+                    dtype=weight_dtype,
+                ),
+                input_dim=1,
+                output_dim=0,
+                weight_loader=weight_loader,
+            ),
+        )
+
+        if self.quant_config.is_checkpoint_fp8_serialized:
+            for scale_name in ["weight_scale", "input_scale"]:
+                layer.register_parameter(
+                    scale_name,
+                    PerTensorScaleParameter(
+                        data=torch.full(
+                            (len(output_partition_sizes),),
+                            torch.finfo(torch.float32).min,
+                            dtype=torch.float32,
+                        ),
+                        weight_loader=weight_loader,
+                    ),
+                )
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        max_w_scale, quantized_weight = requantize_with_max_scale(
+            layer.weight, layer.weight_scale, layer.logical_widths
+        )
+        # Preserve the parameter subclass metadata while rebinding to the
+        # transposed FP8 view expected by the runtime.
+        layer.weight.data = quantized_weight.t().detach()
+        layer.weight.requires_grad_(False)
+        if self.cutlass_fp8_supported:
+            max_w_scale = convert_to_channelwise(max_w_scale, layer.logical_widths)
+        copy_or_rebind_param(layer, "weight_scale", max_w_scale)
+        copy_or_rebind_param(layer, "input_scale", layer.input_scale.max())
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        return apply_fp8_linear(
+            input=x,
+            weight=layer.weight,
+            weight_scale=layer.weight_scale,
+            input_scale=layer.input_scale,
+            bias=bias,
+            cutlass_fp8_supported=self.cutlass_fp8_supported,
+        )
+
+
+class ModelOptFp4LinearMethod(LinearMethodBase):
+    """NVFP4 linear method using the selected FP4 GEMM backend."""
+
+    def __init__(self, quant_config: ModelOptFp4Config):
+        self.quant_config = quant_config
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes: List[int],
+        input_size: int,
+        output_size: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        del input_size, output_size
+        if not self.quant_config.is_checkpoint_nvfp4_serialized:
+            raise ValueError(
+                "NVFP4 quantization was selected, "
+                " dynamic quantization is not supported."
+            )
+        if input_size_per_partition % 16 != 0:
+            raise ValueError(
+                f"Unsupported model when input features size is {input_size_per_partition}, not multiple of 16, for NVFP4 quantization."
+            )
+
+        output_size_per_partition = sum(output_partition_sizes)
+        weight_loader = extra_weight_attrs.get("weight_loader")
+
+        layer.logical_widths = output_partition_sizes
+
+        layer.input_size_per_partition = input_size_per_partition
+        layer.output_size_per_partition = output_size_per_partition
+
+        weight_dtype = (
+            torch.float8_e4m3fn
+            if self.quant_config.is_checkpoint_nvfp4_serialized
+            else params_dtype
+        )
+
+        weight = ModelWeightParameter(
+            data=torch.empty(
+                output_size_per_partition,
+                input_size_per_partition // 2,
+                dtype=torch.uint8,
+            ),
+            input_dim=1,
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("weight", weight)
+
+        input_scale = PerTensorScaleParameter(
+            data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
+            weight_loader=weight_loader,
+        )
+        set_weight_attrs(input_scale, {"missing_param_init": "ones"})
+        layer.register_parameter("input_scale", input_scale)
+
+        weight_scale_2 = PerTensorScaleParameter(
+            data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
+            weight_loader=weight_loader,
+        )
+        set_weight_attrs(weight_scale_2, {"missing_param_init": "ones"})
+        layer.register_parameter("weight_scale_2", weight_scale_2)
+
+        weight_scale = ModelWeightParameter(
+            data=torch.empty(
+                output_size_per_partition,
+                input_size_per_partition // self.quant_config.group_size,
+                dtype=weight_dtype,
+            ),
+            input_dim=1,
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        set_weight_attrs(weight_scale, {"missing_param_init": "ones"})
+        layer.register_parameter("weight_scale", weight_scale)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        input_scale_2 = layer.input_scale.max().to(torch.float32)
+        weight_scale_2 = layer.weight_scale_2.max().to(torch.float32)
+
+        copy_or_rebind_param(
+            layer, "alpha", (input_scale_2 * weight_scale_2).to(torch.float32)
+        )
+        copy_or_rebind_param(
+            layer, "input_scale_inv", (1 / input_scale_2).to(torch.float32)
+        )
+
+        layer.output_size_per_partition = layer.weight.shape[0]
+
+        w = layer.weight.data
+        w_swapped = _prepare_nvfp4_weight_bytes(
+            w,
+            swap_weight_nibbles=getattr(
+                self.quant_config, "swap_weight_nibbles", False
+            ),
+        )
+        scales = layer.weight_scale
+        if (
+            getattr(self.quant_config, "checkpoint_weight_scale_layout", "linear")
+            == "swizzled"
+        ):
+            scales = _swizzled_nvfp4_scales_to_linear(scales)
+
+        _, flashinfer_backend = _get_fp4_gemm_op()
+        if flashinfer_backend == "trtllm":
+            flashinfer_ops = _require_flashinfer()
+
+            weight, _ = pad_nvfp4_weight(w_swapped, n_alignment=128, k_alignment=0)
+            if scales.shape[0] != weight.shape[0]:
+                pad_n = weight.shape[0] - scales.shape[0]
+                scales = torch.nn.functional.pad(scales, (0, 0, 0, pad_n))
+
+            scale_k = scales.shape[1]
+            weights_padding_cols = 0
+            if scale_k % 4 != 0:
+                padded_scale_k = round_up(scale_k, 4)
+                pad_scale_k = padded_scale_k - scale_k
+                scales = torch.nn.functional.pad(scales, (0, pad_scale_k, 0, 0))
+                pad_weight_k = pad_scale_k * 8
+                weight = torch.nn.functional.pad(weight, (0, pad_weight_k, 0, 0))
+                weights_padding_cols = pad_weight_k
+
+            epilogue_tile_m = 128
+            shuffled_scale_shape = scales.shape
+            if not weight.is_cuda:
+                weight = weight.cuda()
+            if scales.device != weight.device:
+                scales = scales.to(device=weight.device)
+            weight = flashinfer_ops.shuffle_matrix_a(
+                weight.view(torch.uint8), epilogue_tile_m
+            )
+            scales = (
+                flashinfer_ops.shuffle_matrix_sf_a(
+                    scales.view(torch.uint8), epilogue_tile_m
+                )
+                .reshape(shuffled_scale_shape)
+                .view(torch.float8_e4m3fn)
+            )
+
+            layer.weights_padding_cols = weights_padding_cols
+            copy_or_rebind_param(layer, "weight", weight)
+            copy_or_rebind_param(layer, "weight_scale_interleaved", scales)
+            return
+        weight, weights_padding_cols = pad_nvfp4_weight(w_swapped)
+        layer.weights_padding_cols = weights_padding_cols
+        copy_or_rebind_param(layer, "weight", weight)
+
+        scale_ndim = scales.ndim
+        if scale_ndim == 2:
+            scales = scales.unsqueeze(0)
+        assert scales.ndim == 3
+        B, M, K = scales.shape
+        M_padded = round_up(M, 128)
+        K_padded = round_up(K, 4)
+        padded_scales = torch.zeros((B, M_padded, K_padded), dtype=scales.dtype)
+        padded_scales[:B, :M, :K] = scales
+
+        _, flashinfer_backend = _get_fp4_gemm_op()
+        uses_flux1_scale_layout = not getattr(
+            self.quant_config, "checkpoint_uses_packed_qkv", False
+        ) and getattr(layer, "prefix", "").startswith(
+            ("transformer_blocks.", "single_transformer_blocks.")
+        )
+        if flashinfer_backend is None or uses_flux1_scale_layout:
+            # CUTLASS and FLUX.1 CUDNN paths need the TMA scale layout.
+            padded_scales = padded_scales.reshape(
+                B, M_padded // 128, 4, 32, K_padded // 4, 4
+            )
+            padded_scales = padded_scales.permute(0, 1, 4, 3, 2, 5)
+
+        padded_scales = padded_scales.contiguous().cuda()
+        padded_scales = (
+            padded_scales.reshape(M_padded, K_padded)
+            if scale_ndim == 2
+            else padded_scales.reshape(B, M_padded, K_padded)
+        )
+        copy_or_rebind_param(layer, "weight_scale_interleaved", padded_scales)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        output_dtype = x.dtype
+        input_shape = x.shape
+        x = x.view(-1, input_shape[-1])
+
+        output_size = layer.output_size_per_partition
+        output_shape = list(input_shape[:-1]) + [output_size]
+
+        fp4_quantize = _get_fp4_quantize_op()
+        if fp4_quantize is None:
+            raise RuntimeError(
+                "No FP4 quantization kernel available. Install flashinfer or sgl_kernel."
+            )
+
+        x_fp4, x_scale_interleaved = fp4_quantize(x, layer.input_scale_inv)
+        weights_padding_cols = getattr(layer, "weights_padding_cols", 0)
+        x_fp4 = pad_nvfp4_activation_for_cutlass(x_fp4, weights_padding_cols)
+
+        w = layer.weight
+        w_scale_interleaved = layer.weight_scale_interleaved
+
+        if x_scale_interleaved.dtype == torch.uint8:
+            x_scale_interleaved = x_scale_interleaved.view(torch.float8_e4m3fn)
+        if w_scale_interleaved.dtype == torch.uint8:
+            w_scale_interleaved = w_scale_interleaved.view(torch.float8_e4m3fn)
+        fp4_gemm, flashinfer_backend = _get_fp4_gemm_op()
+        if flashinfer_backend is not None:
+            out = fp4_gemm(
+                x_fp4,
+                w.T,
+                x_scale_interleaved,
+                w_scale_interleaved.T,
+                layer.alpha,
+                output_dtype,
+                backend=flashinfer_backend,
+            )
+        elif fp4_gemm is not None:
+            out = fp4_gemm(
+                x_fp4,
+                w,
+                x_scale_interleaved,
+                w_scale_interleaved,
+                layer.alpha,
+                output_dtype,
+            )
+        else:
+            raise RuntimeError(
+                "No FP4 GEMM kernel available. Install flashinfer or sgl_kernel."
+            )
+
+        out = slice_nvfp4_output(out, output_size)
+
+        if bias is not None:
+            out = out + bias
+        return out.view(*output_shape)

python/sglang/multimodal_gen/runtime/layers/quantization/modelslim.py

@@ -0,0 +1,253 @@
+from __future__ import annotations
+
+import logging
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, cast
+
+import torch
+
+from sglang.multimodal_gen.runtime.layers.linear import (
+    LinearMethodBase,
+    UnquantizedLinearMethod,
+)
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+    QuantizeMethodBase,
+)
+from sglang.srt.layers.quantization.compressed_tensors.utils import should_ignore_layer
+from sglang.srt.layers.quantization.modelslim.schemes import (
+    ModelSlimW4A4Int4,
+    ModelSlimW8A8Int8,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.quantization.modelslim.schemes import (
+        ModelSlimLinearScheme,
+    )
+
+from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping
+
+logger = logging.getLogger(__name__)
+
+
+class ModelSlimConfig(QuantizationConfig):
+    """
+    Config class for ModelSlim Quantization of Diffusion models https://gitcode.com/Ascend/msmodelslim, a NPU-specific quantization type.
+    The quantization method (W8A8, W4A4, etc.) will be automatically parsed from the `quant_model_description.json` config.
+
+    ModelSlim for Diffusion models includes support for various quantization schemes, such as:
+    - W4A4 dynamic linear
+    - W8A8 static linear
+    - W8A8 dynamic linear
+    """
+
+    def __init__(
+        self,
+        quant_config: Dict[str, Any] = {},
+        reverse_param_names_mapping: dict = None,
+    ):
+        super().__init__()
+        self.quant_description = quant_config
+        ignore = cast(List[str], quant_config.get("ignore", []))
+        self.ignore = ignore
+        packed_modules_mapping = quant_config.get("packed_modules_mapping", {})
+        self.packed_modules_mapping = (
+            packed_modules_mapping if packed_modules_mapping is not None else {}
+        )
+        self._name_mapper = (
+            get_param_names_mapping(reverse_param_names_mapping)
+            if reverse_param_names_mapping is not None
+            else None
+        )
+
+    def get_linear_method(self) -> ModelSlimLinearMethod:
+        return ModelSlimLinearMethod(self)
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> List[torch.dtype]:
+        return [torch.int8, torch.float16, torch.bfloat16]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 0
+
+    @classmethod
+    def get_name(cls) -> str:
+        return "modelslim"
+
+    @classmethod
+    def get_config_filenames(cls) -> List[str]:
+        filenames = ["quant_model_description.json"]
+        return filenames
+
+    @classmethod
+    def from_config(
+        cls, config: Dict[str, Any], reverse_param_names_mapping: dict = None
+    ) -> ModelSlimConfig:
+        return cls(config, reverse_param_names_mapping)
+
+    def get_quant_method(
+        self,
+        layer: torch.nn.Module,
+        prefix: str,
+    ) -> Optional[QuantizeMethodBase]:
+        from sglang.multimodal_gen.runtime.layers.linear import LinearBase
+
+        if isinstance(layer, LinearBase):
+            if should_ignore_layer(
+                prefix,
+                ignore=self.ignore,
+                fused_mapping=self.packed_modules_mapping,
+            ):
+                return UnquantizedLinearMethod()
+            key = "model"
+            packed_modules_mapping_subset = self.packed_modules_mapping.get(key, {})
+            prefix_in_quant_config = prefix
+            proj_name = prefix.split(".")[-1]
+            if proj_name in packed_modules_mapping_subset:
+                prefix_in_quant_config = prefix.replace(
+                    proj_name, packed_modules_mapping_subset[proj_name][0]
+                )
+
+            if self.is_layer_skipped(prefix, packed_modules_mapping_subset):
+                return UnquantizedLinearMethod()
+            scheme = self.get_scheme(layer=layer, layer_name=prefix_in_quant_config)
+            layer.scheme = scheme
+            return ModelSlimLinearMethod(self)
+        else:
+            return None
+
+    def _get_scheme_from_parts(
+        self,
+        layer_name: str,
+    ) -> ModelSlimLinearScheme:
+        full_weight_name = layer_name + ".weight"
+        if self._name_mapper is not None:
+            mapped_name, _, _ = self._name_mapper(full_weight_name)
+        else:
+            mapped_name = full_weight_name
+
+        quant_type = self.quant_description.get(mapped_name, "")
+        prefix = mapped_name.removesuffix(".weight")
+        if quant_type == "W8A8_DYNAMIC" or quant_type == "W8A8":
+            return ModelSlimW8A8Int8(quant_config=self.quant_description, prefix=prefix)
+        elif quant_type == "W4A4_DYNAMIC":
+            return ModelSlimW4A4Int4(quant_config=self.quant_description, prefix=prefix)
+        elif quant_type == "W8A8_MXFP8":
+            from sglang.multimodal_gen.runtime.layers.quantization.modelslim_mxfp8_scheme import (
+                ModelSlimMXFP8Scheme,
+            )
+
+            return ModelSlimMXFP8Scheme()
+        elif quant_type in ("W4A4_MXFP4", "W4A4_MXFP4_DUALSCALE"):
+            from sglang.multimodal_gen.runtime.layers.quantization.modelslim_mxfp4_scheme import (
+                ModelSlimMXFP4Scheme,
+            )
+
+            return ModelSlimMXFP4Scheme()
+        raise NotImplementedError(
+            f"No modelslim compatible scheme was found for layer '{layer_name}'. "
+            f"quant_description['{layer_name}.weight'] = '{quant_type}'"
+        )
+
+    def get_scheme(
+        self, layer: torch.nn.Module, layer_name: Optional[str] = None
+    ) -> Optional[ModelSlimLinearScheme]:
+        """
+        get_scheme method adjusted for modelslim, taken from
+        python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py
+        """
+        scheme = self._get_scheme_from_parts(
+            layer_name=layer_name,
+        )
+
+        # Ascend doesn't support device capability
+        logger.debug("Using scheme: %s for %s", scheme.__class__.__name__, layer_name)
+        return scheme
+
+    def is_layer_skipped(
+        self, prefix: str, fused_mapping: Mapping[str, List[str]] = MappingProxyType({})
+    ):
+        # adapted from vllm.model_executor.layers.quantization.utils.quant_utils.is_layer_skipped
+        proj_name = prefix.split(".")[-1]
+        if proj_name in fused_mapping:
+            shard_prefixes = [
+                prefix.replace(proj_name, shard_proj_name)
+                for shard_proj_name in fused_mapping[proj_name]
+            ]
+
+            is_skipped = None
+            for shard_prefix in shard_prefixes:
+                is_shard_skipped = (
+                    self.quant_description.get(shard_prefix + ".weight", "") == "FLOAT"
+                )
+
+                if is_skipped is None:
+                    is_skipped = is_shard_skipped
+                elif is_shard_skipped != is_skipped:
+                    raise ValueError(
+                        f"Detected some but not all shards of {prefix} "
+                        "are quantized. All shards of fused layers "
+                        "to have the same precision."
+                    )
+        else:
+            is_skipped = self.quant_description.get(prefix + ".weight", "") == "FLOAT"
+
+        assert is_skipped is not None
+        return is_skipped
+
+    def get_scaled_act_names(self) -> List[str]:
+        return []
+
+
+class ModelSlimLinearMethod(LinearMethodBase):
+
+    def __init__(self, quantization_config: ModelSlimConfig):
+        self.quantization_config = quantization_config
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        layer.scheme.process_weights_after_loading(layer)
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes: List[int],
+        input_size: int,
+        output_size: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        """
+        Use the ModelSlimLinearScheme associated with each layer to create
+        the necessary parameters for the layer. See LinearMethodBase for param
+        details
+        """
+        weight_loader = extra_weight_attrs.get("weight_loader")
+        layer.scheme.create_weights(
+            layer=layer,
+            input_size=input_size,
+            input_size_per_partition=input_size_per_partition,
+            output_partition_sizes=output_partition_sizes,
+            output_size=output_size,
+            params_dtype=params_dtype,
+            weight_loader=weight_loader,
+        )
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ):
+        """
+        Use the output of create_weights and the CompressedTensorsScheme
+        associated with the layer to apply the forward pass with the
+        layer input.  See LinearMethodBase for param details
+
+        """
+
+        scheme = layer.scheme
+        if scheme is None:
+            raise ValueError("A scheme must be defined for each layer")
+        return scheme.apply_weights(layer, x, bias=bias)

python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp4_scheme.py

@@ -0,0 +1,197 @@
+"""ModelSlim MXFP4 scheme for pre-quantized weight inference on Ascend NPU.
+
+Loads weights pre-quantized by msmodelslim and runs MXFP4 dual-level
+matmul at inference via npu_dual_level_quant_matmul.
+
+Checkpoint tensor formats (verified from msmodelslim export):
+  weight:           [out, in]           float8_e4m3fn  (FP4 data in fp8 container)
+  weight_scale:     [out, in/32]        uint8          (L1 block scales, e8m0+127)
+  weight_dual_scale:[out, in/512, 1]    float32        (L0 coarse scales)
+  mul_scale:        [in]                float32        (smooth quant activation scale)
+
+Reference: MindIE-SD W4A4MXFP4DualQuantLinear
+(MindIE-SD/mindiesd/quantization/layer.py)
+"""
+
+from typing import List, Optional
+
+import torch
+
+from sglang.multimodal_gen.runtime.platforms import current_platform
+
+_is_npu = current_platform.is_npu()
+
+if _is_npu:
+    import torch_npu
+
+from sglang.multimodal_gen.runtime.models.parameter import (
+    BasevLLMParameter,
+    GroupQuantScaleParameter,
+    ModelWeightParameter,
+)
+from sglang.srt.layers.quantization.modelslim.schemes import ModelSlimLinearScheme
+
+MXFP4_BLOCK_SIZE = 32
+# L1 (dual) scale groups this many L0 blocks together.
+# L1 block covers 16 * 32 = 512 elements.
+MXFP4_DUAL_LEVEL_RATIO = 16
+
+
+class ModelSlimMXFP4Scheme(ModelSlimLinearScheme):
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes: List[int],
+        input_size: int,
+        output_size: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        weight_loader = extra_weight_attrs.get("weight_loader")
+        output_size_per_partition = sum(output_partition_sizes)
+
+        # msmodelslim exports weight as float8_e4m3fn, shape [out, in].
+        # Each byte is a float8 container for FP4 data; the actual FP4 packing
+        # (npu_dtype_cast → float4_e2m1fn_x2) happens in process_weights_after_loading.
+        weight = ModelWeightParameter(
+            data=torch.empty(
+                (output_size_per_partition, input_size_per_partition),
+                dtype=torch.float8_e4m3fn,
+            ),
+            input_dim=1,
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("weight", weight)
+
+        # L1 block scale: uint8 [out, in/32], e8m0 scale with +127 offset.
+        scale_dim = input_size_per_partition // MXFP4_BLOCK_SIZE
+        weight_scale = GroupQuantScaleParameter(
+            data=torch.empty(
+                (output_size_per_partition, scale_dim),
+                dtype=torch.uint8,
+            ),
+            input_dim=1,
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("weight_scale", weight_scale)
+
+        # L0 (coarse) scale for dual-level quantization matmul.
+        # Each L0 block covers MXFP4_DUAL_LEVEL_RATIO L1 blocks = 16 * 32 = 512 elements.
+        dual_scale_dim = scale_dim // MXFP4_DUAL_LEVEL_RATIO  # in/32 / 16 = in/512
+        weight_dual_scale = GroupQuantScaleParameter(
+            data=torch.empty(
+                (output_size_per_partition, dual_scale_dim, 1),
+                dtype=torch.float32,
+            ),
+            input_dim=1,
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("weight_dual_scale", weight_dual_scale)
+
+        # Smooth quant activation scale (mul_scale) from NonFusionSmoothQuantWrapper.
+        # msmodelslim exports this as `<prefix>.div.mul_scale` with shape [in].
+        # After repack, it becomes `<prefix>.mul_scale`.
+        # This is CRITICAL: the offline-quantized weights were calibrated with
+        # x * mul_scale applied to the activation. Omitting it causes mosaic output.
+        # Ref: MindIE-SD W4A4MXFP4DualQuantLinear.quant_matmul lines 385-386.
+        mul_scale = BasevLLMParameter(
+            data=torch.empty(
+                (input_size_per_partition,),
+                dtype=torch.float32,
+            ),
+            weight_loader=weight_loader,
+        )
+        # If mul_scale is not in the checkpoint (e.g. non-smooth-quant model
+        # or old repack without .div. handling), initialize to ones so that
+        # x * 1.0 = x (no-op). fsdp_load.py checks this attribute.
+        mul_scale.missing_param_init = "ones"
+        layer.register_parameter("mul_scale", mul_scale)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module):
+        # Cast weight from fp8 container to FP4 packed format
+        weight = layer.weight.data
+        if not weight.is_npu:
+            weight = weight.to(f"npu:{torch.npu.current_device()}")
+        weight = torch_npu.npu_dtype_cast(weight, torch_npu.float4_e2m1fn_x2)
+        # npu_dual_level_quant_matmul requires x2 in FRACTAL_NZ format (format 29).
+        # Reference: MindIE-SD W4A4MXFP4DualQuantLinear._init_dynamic_quant_param
+        weight = torch_npu.npu_format_cast(
+            weight.view(torch.int8), 29, customize_dtype=torch.int8
+        )
+        layer.weight = torch.nn.Parameter(weight, requires_grad=False)
+
+        # Reshape weight_scale: [out, in/32] -> [out, in/64, 2]
+        # The dual-level matmul API expects L1 scales in this 3D format
+        weight_scale = layer.weight_scale.data
+        if not weight_scale.is_npu:
+            weight_scale = weight_scale.to(f"npu:{torch.npu.current_device()}")
+        weight_scale = weight_scale.reshape(weight_scale.shape[0], -1, 2)
+        layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
+
+        # Transform weight_dual_scale: [out, in/512, 1] -> [in/512, out]
+        weight_dual_scale = layer.weight_dual_scale.data
+        if not weight_dual_scale.is_npu:
+            weight_dual_scale = weight_dual_scale.to(
+                f"npu:{torch.npu.current_device()}"
+            )
+        weight_dual_scale = weight_dual_scale.squeeze(-1).transpose(0, 1).contiguous()
+        layer.weight_dual_scale = torch.nn.Parameter(
+            weight_dual_scale, requires_grad=False
+        )
+
+        # Move mul_scale to NPU if present and not already there
+        mul_scale = layer.mul_scale.data
+        if not mul_scale.is_npu:
+            mul_scale = mul_scale.to(f"npu:{torch.npu.current_device()}")
+        layer.mul_scale = torch.nn.Parameter(mul_scale, requires_grad=False)
+        layer.use_mul_scale = not torch.all(mul_scale == 1.0).item()
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+
+        original_dtype = x.dtype
+        if original_dtype not in (torch.float16, torch.bfloat16):
+            x = x.to(torch.bfloat16)
+            original_dtype = torch.bfloat16
+
+        # Flatten to 2D for npu_dynamic_dual_level_mx_quant
+        input_shape = x.shape
+        x_2d = x.reshape(-1, x.shape[-1])
+
+        # Apply smooth quant scale before activation quantization.
+        # The offline-quantized weights were calibrated under x * mul_scale,
+        # so we MUST apply it here for scale alignment.
+        # Reference: MindIE-SD W4A4MXFP4DualQuantLinear.quant_matmul
+        mul_scale = layer.mul_scale
+        if getattr(layer, "use_mul_scale", True):
+            x_2d = x_2d * mul_scale.to(x_2d.dtype)
+
+        # Dual-level MXFP4 activation quantization
+        x1, l0_scale, l1_scale = torch_npu.npu_dynamic_dual_level_mx_quant(
+            x_2d, smooth_scale=None
+        )
+
+        # Dual-level MXFP4 matmul
+        output = torch_npu.npu_dual_level_quant_matmul(
+            x1,
+            layer.weight,
+            l0_scale,
+            layer.weight_dual_scale,
+            l1_scale,
+            layer.weight_scale,
+            bias=bias.to(torch.float32) if bias is not None else None,
+            output_dtype=original_dtype,
+        )
+
+        # Restore original shape
+        output_shape = list(input_shape[:-1]) + [output.shape[-1]]
+        return output.reshape(output_shape)

python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp8_scheme.py

@@ -0,0 +1,124 @@
+"""ModelSlim MXFP8 scheme for pre-quantized weight inference on Ascend NPU.
+
+Loads weights pre-quantized by msmodelslim (float8_e4m3fn weights,
+uint8 scales) and runs MXFP8 matmul at inference.
+"""
+
+from typing import List, Optional
+
+import torch
+
+from sglang.multimodal_gen.runtime.platforms import current_platform
+
+_is_npu = current_platform.is_npu()
+
+if _is_npu:
+    import torch_npu
+
+from sglang.multimodal_gen.runtime.models.parameter import (
+    GroupQuantScaleParameter,
+    ModelWeightParameter,
+)
+from sglang.srt.layers.quantization.modelslim.schemes import ModelSlimLinearScheme
+
+MXFP8_BLOCK_SIZE = 32
+
+
+class ModelSlimMXFP8Scheme(ModelSlimLinearScheme):
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes: List[int],
+        input_size: int,
+        output_size: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        weight_loader = extra_weight_attrs.get("weight_loader")
+        output_size_per_partition = sum(output_partition_sizes)
+
+        # msmodelslim exports weight as float8_e4m3fn, shape [out, in]
+        weight = ModelWeightParameter(
+            data=torch.empty(
+                (output_size_per_partition, input_size_per_partition),
+                dtype=torch.float8_e4m3fn,
+            ),
+            input_dim=1,
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("weight", weight)
+
+        # msmodelslim exports weight_scale as uint8, shape [out, in/32].
+        # NOTE: This parameter is intentionally named "weight_scale" (not
+        # "weight_scale_inv" as used in mxfp8_npu.py) because the weight loader
+        # matches parameter names to checkpoint keys, and msmodelslim checkpoints
+        # store this tensor under the key "<layer>.weight_scale".
+        scale_dim = input_size_per_partition // MXFP8_BLOCK_SIZE
+        weight_scale = GroupQuantScaleParameter(
+            data=torch.empty(
+                (output_size_per_partition, scale_dim),
+                dtype=torch.uint8,
+            ),
+            input_dim=1,
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("weight_scale", weight_scale)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module):
+        # weight is already float8_e4m3fn, no cast needed
+        weight = layer.weight.data
+        layer.weight = torch.nn.Parameter(weight, requires_grad=False)
+
+        # Reshape weight_scale: [out, in/32] -> [out, in/32//2, 2]
+        weight_scale = layer.weight_scale.data
+        weight_scale = weight_scale.reshape(weight_scale.shape[0], -1, 2)
+        layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+
+        original_dtype = x.dtype
+        if original_dtype not in (torch.float16, torch.bfloat16):
+            # npu_dynamic_mx_quant only accepts fp16/bf16 activations
+            x = x.to(torch.bfloat16)
+            original_dtype = torch.bfloat16
+
+        # npu_dynamic_mx_quant requires a 2D input [tokens, hidden_size].
+        # Diffusion transformer inputs are typically 3D [batch, seq, hidden] or
+        # higher. Flattening to 2D merges all leading dimensions into a single
+        # token axis so the NPU kernel can compute per-token MXFP8 scales, then
+        # we restore the original shape from the output.
+        input_shape = x.shape
+        x_2d = x.reshape(-1, x.shape[-1])
+
+        # Dynamic MXFP8 activation quantisation
+        qx, input_scale = torch_npu.npu_dynamic_mx_quant(
+            x_2d, dst_type=torch_npu.float8_e4m3fn
+        )
+
+        # MXFP8 matmul
+        output = torch_npu.npu_quant_matmul(
+            qx,
+            layer.weight.transpose(0, 1),
+            layer.weight_scale.transpose(0, 1),
+            scale_dtype=torch_npu.float8_e8m0fnu,
+            pertoken_scale=input_scale,
+            pertoken_scale_dtype=torch_npu.float8_e8m0fnu,
+            bias=bias.to(torch.float32) if bias is not None else None,
+            output_dtype=original_dtype,
+            group_sizes=[1, 1, MXFP8_BLOCK_SIZE],
+        )
+
+        # Restore original shape
+        output_shape = list(input_shape[:-1]) + [output.shape[-1]]
+        output = output.reshape(output_shape)
+
+        return output

python/sglang/multimodal_gen/runtime/layers/quantization/mxfp4.py

@@ -0,0 +1,238 @@
+import logging
+from typing import Dict, List, Optional
+
+import torch
+from torch.nn.parameter import Parameter
+
+from sglang.multimodal_gen.runtime.layers.linear import (
+    LinearMethodBase,
+    UnquantizedLinearMethod,
+)
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+)
+from sglang.multimodal_gen.runtime.models.parameter import (
+    ModelWeightParameter,
+    PerTensorScaleParameter,
+)
+from sglang.srt.layers.quantization.utils import is_layer_skipped
+from sglang.srt.utils import is_hip, mxfp_supported
+
+logger = logging.getLogger(__name__)
+_is_hip = is_hip()
+
+if _is_hip:
+    try:
+        import aiter
+        from aiter.ops.gemm_op_a4w4 import gemm_a4w4
+        from aiter.ops.shuffle import shuffle_weight
+        from aiter.utility.fp4_utils import dynamic_mxfp4_quant
+    except ImportError as e:
+        logger.warning(f"aiter MXFP4 kernels not available: {e}")
+        aiter = None
+        shuffle_weight = None
+        dynamic_mxfp4_quant = None
+        gemm_a4w4 = None
+
+# The gemm_a4w4 ASM kernel has degraded precision when the output
+# dimension (N) is smaller than its minimum tile size.
+# Layers with output_size falls below this threshold will stay unquantized
+_MXFP4_MIN_OUTPUT_DIM = 256
+
+
+class Mxfp4Config(QuantizationConfig):
+    """
+    MXFP4 quantization config for diffusion models.
+
+    Supports online quantization from unquantized BF16/FP16 checkpoints;
+    no-arg ``Mxfp4Config()`` selects that online (post-load) path.
+    Note: MXFP4 requires ROCm and MI350+ (gfx95x).
+    """
+
+    def __init__(
+        self,
+        is_checkpoint_mxfp4_serialized: bool = False,
+        ignored_layers: Optional[List[str]] = None,
+        packed_modules_mapping: Optional[Dict[str, List[str]]] = None,
+    ):
+        super().__init__()
+        self.is_checkpoint_mxfp4_serialized = is_checkpoint_mxfp4_serialized
+        self.ignored_layers = ignored_layers or []
+        self.packed_modules_mapping = packed_modules_mapping or {}
+
+    @classmethod
+    def get_name(cls) -> str:
+        return "mxfp4"
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> list[torch.dtype]:
+        return [torch.bfloat16, torch.float16]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 95  # gfx95x, Note: mxfp_supported() is a better check
+
+    @classmethod
+    def get_config_filenames(cls) -> list[str]:
+        return []  # No config file needed for online quantization
+
+    @classmethod
+    def from_config(cls, config: dict) -> "Mxfp4Config":
+        """Create from model config (for pre-quantized checkpoints)."""
+        is_serialized = config.get("quant_method") == "mxfp4"
+        return cls(is_checkpoint_mxfp4_serialized=is_serialized)
+
+    def get_quant_method(self, layer, prefix: str):
+        from sglang.multimodal_gen.runtime.layers.linear import LinearBase
+
+        if isinstance(layer, LinearBase):
+            if is_layer_skipped(
+                prefix,
+                self.ignored_layers,
+                fused_mapping=self.packed_modules_mapping,
+            ):
+                logger.debug(
+                    f"MXFP4: Keeping layer {prefix} unquantized (in ignored_layers)"
+                )
+                return UnquantizedLinearMethod()
+            # Skip layers whose output dims are too small, see ASM kernel comment above
+            output_size = getattr(layer, "output_size", None)
+            if output_size is not None and output_size < _MXFP4_MIN_OUTPUT_DIM:
+                logger.info(
+                    f"MXFP4: Keeping layer {prefix} unquantized "
+                    f"(output_size={output_size} < {_MXFP4_MIN_OUTPUT_DIM})"
+                )
+                return UnquantizedLinearMethod()
+            logger.debug(f"MXFP4: Replacing layer {prefix} with MXFP4 linear method")
+            return Mxfp4LinearMethod(self)
+        else:
+            logger.debug(f"MXFP4: Skipping layer {prefix} (not a LinearBase)")
+        return None
+
+
+class Mxfp4LinearMethod(LinearMethodBase):
+    """
+    MXFP4 online quantization method for linear layers.
+
+    Quantizes unquantized BF16/FP16 weights to MXFP4 format during
+    process_weights_after_loading().
+    """
+
+    def __init__(self, quant_config: Mxfp4Config):
+        self.quant_config = quant_config
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes: list[int],
+        input_size: int,
+        output_size: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        """
+        Creates BF16/FP16 parameters that will be
+        quantized to MXFP4 in process_weights_after_loading().
+        """
+        output_size_per_partition = sum(output_partition_sizes)
+        weight_loader = extra_weight_attrs.get("weight_loader")
+
+        weight = ModelWeightParameter(
+            data=torch.empty(
+                output_size_per_partition,
+                input_size_per_partition,
+                dtype=params_dtype,
+            ),
+            weight_loader=weight_loader,
+            input_dim=1,
+            output_dim=0,
+        )
+        layer.register_parameter("weight", weight)
+
+        # Placeholder scale (will be created during quantization)
+        weight_scale = PerTensorScaleParameter(
+            data=torch.empty(1, dtype=torch.float32),
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("weight_scale", weight_scale)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module):
+        """
+        Quantize BF16/FP16 weights to MXFP4 after loading from checkpoint.
+
+        Converts weights from unquantized format to:
+        - Packed uint8 (2 FP4 values per byte)
+        - E8M0 scales (one per 32-element block)
+        """
+        if not mxfp_supported():
+            platform = "unknown"
+            if _is_hip:
+                try:
+                    platform = torch.cuda.get_device_properties(0).gcnArchName
+                except:
+                    platform = "ROCm (unknown arch)"
+            raise RuntimeError(
+                f"MXFP4 quantization requires ROCm and MI350+ (gfx95x). "
+                f"Current platform: {platform}."
+            )
+
+        # Check if weights are already quantized
+        if layer.weight.dtype not in [torch.bfloat16, torch.float16]:
+            # Already quantized or unexpected dtype
+            logger.info("Weights are quantized or unexpected dtype")
+            return
+
+        if any(fn is None for fn in (dynamic_mxfp4_quant, shuffle_weight, gemm_a4w4)):
+            raise RuntimeError(
+                "aiter MXFP4 kernels not available. "
+                "Install aiter with MXFP4 support."
+            )
+
+        weight_data = layer.weight.data
+        was_on_cpu = weight_data.device.type == "cpu"
+        if was_on_cpu:
+            weight_data = weight_data.cuda()
+
+        w_quant, mx_scales = dynamic_mxfp4_quant(weight_data, shuffle=True)
+
+        w_quant_shuffled = shuffle_weight(w_quant)
+
+        if was_on_cpu:
+            w_quant_shuffled = w_quant_shuffled.cpu()
+            mx_scales = mx_scales.cpu()
+
+        layer.weight = Parameter(w_quant_shuffled, requires_grad=False)
+        layer.weight_scale = Parameter(mx_scales, requires_grad=False)
+
+        logger.debug(
+            f"MXFP4: Quantized layer weights - weight {layer.weight.shape} {layer.weight.dtype}, "
+            f"scale {layer.weight_scale.shape}"
+        )
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+
+        if not mxfp_supported():
+            raise RuntimeError(
+                "MXFP4 inference requires ROCm and MI350+ (gfx95x). "
+                "Current platform not supported."
+            )
+
+        # Handle 3D input tensors [batch, seq, hidden]
+        original_shape = x.shape
+        if x.dim() == 3:
+            x = x.view(-1, x.shape[-1])
+
+        x_fp4, x_scale = dynamic_mxfp4_quant(x, shuffle=True)
+
+        y = gemm_a4w4(x_fp4, layer.weight, x_scale, layer.weight_scale)
+
+        if bias is not None:
+            y = y + bias
+
+        return y.view(*original_shape[:-1], layer.weight.shape[0])

python/sglang/multimodal_gen/runtime/layers/quantization/mxfp4_npu.py

@@ -0,0 +1,201 @@
+"""Online MXFP4 quantization for Diffusion models on Ascend NPU.
+
+Provides ``NPUMXFP4Config`` (registered as ``"mxfp4_npu"``) and
+``NPUMXFP4DiffusionLinearMethod`` which quantises FP16/BF16 weights to MXFP4
+at load time using dual-level MX quantization and uses
+``npu_dynamic_dual_level_mx_quant`` + ``npu_dual_level_quant_matmul`` for
+inference.
+
+The ``"mxfp4_npu"`` key is distinct from upstream's ROCm ``"mxfp4"``
+(``Mxfp4Config`` in ``mxfp4.py``) which targets AMD MI350+ via aiter kernels.
+
+NOTE: Online weight quantization via ``npu_dynamic_dual_level_mx_quant`` is
+experimental. MindIE-SD only uses an offline (pre-quantized) path for MXFP4
+weights. The online path quantizes FP16/BF16 weights at load time, which may
+produce different numerical results than the offline calibrated path.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+import torch
+from torch.nn.parameter import Parameter
+
+from sglang.multimodal_gen.runtime.platforms import current_platform
+
+_is_npu = current_platform.is_npu()
+
+if _is_npu:
+    import torch_npu
+
+from sglang.multimodal_gen.runtime.layers.linear import LinearBase, LinearMethodBase
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+    QuantizeMethodBase,
+)
+from sglang.multimodal_gen.runtime.models.parameter import ModelWeightParameter
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class NPUMXFP4Config(QuantizationConfig):
+    """Config for online MXFP4 quantization on Ascend NPU (Diffusion)."""
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    @classmethod
+    def get_name(cls) -> str:
+        return "mxfp4_npu"
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> List[torch.dtype]:
+        return [torch.bfloat16, torch.float16]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 0  # NPU, not CUDA
+
+    @classmethod
+    def get_config_filenames(cls) -> List[str]:
+        return []
+
+    @classmethod
+    def from_config(cls, config: Dict[str, Any]) -> NPUMXFP4Config:
+        return cls()
+
+    def get_quant_method(
+        self, layer: torch.nn.Module, prefix: str
+    ) -> Optional[QuantizeMethodBase]:
+        if isinstance(layer, LinearBase):
+            return NPUMXFP4DiffusionLinearMethod(self)
+        return None
+
+    def get_scaled_act_names(self) -> List[str]:
+        return []
+
+
+class NPUMXFP4DiffusionLinearMethod(LinearMethodBase):
+    """Ascend NPU MXFP4 linear method for Diffusion models (dual-level).
+
+    Online mode: loads FP16/BF16 weights → quantises to MXFP4 at load time
+    via ``npu_dynamic_dual_level_mx_quant``.
+    Inference: dynamic dual-level MXFP4 activation quant + dual-level matmul.
+
+    Reference: MindIE-SD ``W4A4MXFP4DualQuantLinear`` (offline path only).
+    """
+
+    def __init__(self, quant_config: NPUMXFP4Config):
+        self.quant_config = quant_config
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes: List[int],
+        input_size: int,
+        output_size: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        output_size_per_partition = sum(output_partition_sizes)
+        weight_loader = extra_weight_attrs.get("weight_loader")
+
+        layer.logical_widths = output_partition_sizes
+        layer.input_size_per_partition = input_size_per_partition
+        layer.output_size_per_partition = output_size_per_partition
+        layer.orig_dtype = params_dtype
+
+        # Load weights in original dtype; quantise later in process_weights_after_loading
+        weight = ModelWeightParameter(
+            data=torch.empty(
+                output_size_per_partition,
+                input_size_per_partition,
+                dtype=params_dtype,
+            ),
+            input_dim=1,
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("weight", weight)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        weight_fp = layer.weight.data
+        if weight_fp.dtype not in (torch.float16, torch.bfloat16):
+            weight_fp = weight_fp.to(torch.bfloat16)
+
+        # Move weight to NPU if needed. dit_cpu_offload defaults to True in
+        # ServerArgs, which causes fsdp_load to move parameters back to CPU
+        # after loading. npu_dynamic_dual_level_mx_quant requires an NPU tensor.
+        if not weight_fp.is_npu:
+            weight_fp = weight_fp.to(f"npu:{torch.npu.current_device()}")
+
+        # Online dual-level MXFP4 weight quantisation.
+        # NOTE: This is experimental — MindIE-SD only has an offline path for
+        # MXFP4 weights. We assume npu_dynamic_dual_level_mx_quant can also
+        # quantise weights (not just activations).
+        # Returns: (qw, w_dual_scale, w_scale)
+        #   qw          — quantized weight in float4_e2m1fn_x2 (2 FP4 packed/byte)
+        #   w_dual_scale — L0-level scale (goes to pos 3 in npu_dual_level_quant_matmul)
+        #   w_scale      — L1-level scale (goes to pos 5 in npu_dual_level_quant_matmul)
+        qw, w_dual_scale, w_scale = torch_npu.npu_dynamic_dual_level_mx_quant(
+            weight_fp, smooth_scale=None
+        )
+
+        # npu_dual_level_quant_matmul requires x2 (weight) in FRACTAL_NZ format.
+        # Reference: MindIE-SD W4A4MXFP4DualQuantLinear._init_dynamic_quant_param
+        qw = torch_npu.npu_format_cast(
+            qw.view(torch.int8), 29, customize_dtype=torch.int8
+        )
+
+        # x2Level0Scale must be [in/level0_block_size, out] — transpose from
+        # the [out, in/level0_block_size] shape returned by the quant op.
+        # Reference: MindIE-SD layer.py:409
+        w_dual_scale = w_dual_scale.squeeze(-1).transpose(0, 1).contiguous()
+
+        layer.weight = Parameter(qw, requires_grad=False)
+        layer.weight_dual_scale = Parameter(w_dual_scale, requires_grad=False)
+        layer.weight_scale = Parameter(w_scale, requires_grad=False)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        original_dtype = x.dtype
+        if original_dtype not in (torch.float16, torch.bfloat16):
+            x = x.to(torch.bfloat16)
+            original_dtype = torch.bfloat16
+
+        # Flatten to 2D [tokens, hidden] for the quantization operators
+        input_shape = x.shape
+        x_2d = x.reshape(-1, x.shape[-1])
+
+        # Dynamic dual-level MXFP4 activation quantisation
+        qx, act_l0_scale, act_l1_scale = torch_npu.npu_dynamic_dual_level_mx_quant(
+            x_2d, smooth_scale=None
+        )
+
+        # Dual-level MXFP4 matmul
+        # Arg order: act_quant, weight_quant, act_l0_scale, weight_dual_scale,
+        #            act_l1_scale, weight_scale, bias=, output_dtype=
+        # NOTE: weight is NOT transposed (unlike MXFP8's npu_quant_matmul).
+        output = torch_npu.npu_dual_level_quant_matmul(
+            qx,
+            layer.weight,
+            act_l0_scale,
+            layer.weight_dual_scale,
+            act_l1_scale,
+            layer.weight_scale,
+            bias=bias.to(torch.float32) if bias is not None else None,
+            output_dtype=original_dtype,
+        )
+
+        # Restore original shape (replace last dim with output features)
+        output_shape = list(input_shape[:-1]) + [output.shape[-1]]
+        output = output.reshape(output_shape)
+
+        return output

python/sglang/multimodal_gen/runtime/layers/quantization/mxfp8_npu.py

@@ -0,0 +1,176 @@
+"""Online MXFP8 quantization for Diffusion models on Ascend NPU.
+
+Provides ``MXFP8Config`` (registered as ``"mxfp8"``) and
+``NPUMXFP8DiffusionLinearMethod`` which quantise FP16/BF16 weights to MXFP8
+at load time and use ``npu_dynamic_mx_quant`` + ``npu_quant_matmul`` for
+inference, mirroring the LLM-side ``NPUMXFP8LinearMethod``.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+import torch
+from torch.nn.parameter import Parameter
+
+from sglang.multimodal_gen.runtime.platforms import current_platform
+
+_is_npu = current_platform.is_npu()
+
+if _is_npu:
+    import torch_npu
+
+from sglang.multimodal_gen.runtime.layers.linear import LinearBase, LinearMethodBase
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+    QuantizeMethodBase,
+)
+from sglang.multimodal_gen.runtime.models.parameter import ModelWeightParameter
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+MXFP8_BLOCK_SIZE = 32
+
+
+class MXFP8Config(QuantizationConfig):
+    """Config for online MXFP8 quantization on Ascend NPU (Diffusion)."""
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    @classmethod
+    def get_name(cls) -> str:
+        return "mxfp8"
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> List[torch.dtype]:
+        return [torch.bfloat16, torch.float16]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 0  # NPU, not CUDA
+
+    @classmethod
+    def get_config_filenames(cls) -> List[str]:
+        return []
+
+    @classmethod
+    def from_config(cls, config: Dict[str, Any]) -> MXFP8Config:
+        return cls()
+
+    def get_quant_method(
+        self, layer: torch.nn.Module, prefix: str
+    ) -> Optional[QuantizeMethodBase]:
+        if isinstance(layer, LinearBase):
+            return NPUMXFP8DiffusionLinearMethod(self)
+        return None
+
+    def get_scaled_act_names(self) -> List[str]:
+        return []
+
+
+class NPUMXFP8DiffusionLinearMethod(LinearMethodBase):
+    """Ascend NPU MXFP8 linear method for Diffusion models.
+
+    Online mode: loads FP16/BF16 weights → quantises to MXFP8 at load time.
+    Inference: dynamic MXFP8 activation quant + MXFP8 matmul (block_size=32).
+    """
+
+    def __init__(self, quant_config: MXFP8Config):
+        self.quant_config = quant_config
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes: List[int],
+        input_size: int,
+        output_size: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        output_size_per_partition = sum(output_partition_sizes)
+        weight_loader = extra_weight_attrs.get("weight_loader")
+
+        layer.logical_widths = output_partition_sizes
+        layer.input_size_per_partition = input_size_per_partition
+        layer.output_size_per_partition = output_size_per_partition
+        layer.orig_dtype = params_dtype
+
+        # Load weights in original dtype; quantise later in process_weights_after_loading
+        weight = ModelWeightParameter(
+            data=torch.empty(
+                output_size_per_partition,
+                input_size_per_partition,
+                dtype=params_dtype,
+            ),
+            input_dim=1,
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("weight", weight)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+
+        weight_fp = layer.weight.data
+        if weight_fp.dtype not in (torch.float16, torch.bfloat16):
+            weight_fp = weight_fp.to(torch.bfloat16)
+
+        # Move weight to NPU if needed. We intentionally use a conditional
+        # move rather than an assert because `dit_cpu_offload` defaults to
+        # True in ServerArgs, which causes fsdp_load to move every parameter
+        # back to CPU after loading (even when the target device is NPU).
+        # npu_dynamic_mx_quant requires an NPU tensor, so we must transfer
+        # here. The quantized fp8 weights produced below will remain on NPU
+        # for inference; if the model still needs to be offloaded after
+        # quantization (e.g. very large model on a small NPU), a higher-level
+        # offload pass can move them back afterwards.
+        if not weight_fp.is_npu:
+            weight_fp = weight_fp.to(f"npu:{torch.npu.current_device()}")
+
+        # Online MXFP8 quantisation of weights (block_size=32)
+        qw, w_scale = torch_npu.npu_dynamic_mx_quant(
+            weight_fp, dst_type=torch_npu.float8_e4m3fn
+        )
+        layer.weight = Parameter(qw, requires_grad=False)
+        layer.weight_scale_inv = Parameter(w_scale, requires_grad=False)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        original_dtype = x.dtype
+        if original_dtype not in (torch.float16, torch.bfloat16):
+            x = x.to(torch.bfloat16)
+            original_dtype = torch.bfloat16
+
+        # Flatten to 2D [tokens, hidden] so npu_dynamic_mx_quant returns 3D scale
+        input_shape = x.shape
+        x_2d = x.reshape(-1, x.shape[-1])
+
+        # Dynamic MXFP8 activation quantisation
+        qx, input_scale = torch_npu.npu_dynamic_mx_quant(
+            x_2d, dst_type=torch_npu.float8_e4m3fn
+        )
+
+        # MXFP8 matmul
+        output = torch_npu.npu_quant_matmul(
+            qx,
+            layer.weight.transpose(0, 1),
+            layer.weight_scale_inv.transpose(0, 1),
+            scale_dtype=torch_npu.float8_e8m0fnu,
+            pertoken_scale=input_scale,
+            pertoken_scale_dtype=torch_npu.float8_e8m0fnu,
+            bias=bias.to(torch.float32) if bias is not None else None,
+            output_dtype=original_dtype,
+            group_sizes=[1, 1, MXFP8_BLOCK_SIZE],
+        )
+
+        # Restore original shape (replace last dim with output features)
+        output_shape = list(input_shape[:-1]) + [output.shape[-1]]
+        output = output.reshape(output_shape)
+
+        return output

python/sglang/multimodal_gen/runtime/layers/quantization/nunchaku_linear.py

@@ -0,0 +1,291 @@
+# SPDX-License-Identifier: Apache-2.0
+from typing import List, Optional
+
+import torch
+import torch.nn as nn
+from torch.nn.parameter import Parameter
+
+from sglang.multimodal_gen.runtime.layers.linear import LinearMethodBase
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.runtime.utils.weight_attrs import set_weight_attrs
+
+logger = init_logger(__name__)
+
+try:
+    from nunchaku.ops.gemm import svdq_gemm_w4a4_cuda
+    from nunchaku.ops.gemv import awq_gemv_w4a16_cuda
+    from nunchaku.ops.quantize import svdq_quantize_w4a4_act_fuse_lora_cuda
+except ImportError:
+    svdq_gemm_w4a4_cuda = None
+    awq_gemv_w4a16_cuda = None
+    svdq_quantize_w4a4_act_fuse_lora_cuda = None
+
+
+class NunchakuSVDQLinearMethod(LinearMethodBase):
+    def __init__(
+        self,
+        precision: str = "int4",
+        rank: int = 32,
+        act_unsigned: bool = False,
+    ):
+        self.precision = precision
+        self.rank = rank
+        self.act_unsigned = act_unsigned
+
+        if precision == "nvfp4":
+            self.group_size = 16
+        else:
+            self.group_size = 64
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes: List[int],
+        input_size: int,
+        output_size: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ) -> None:
+        output_size_per_partition = sum(output_partition_sizes)
+
+        qweight = Parameter(
+            torch.empty(
+                output_size_per_partition,
+                input_size_per_partition // 2,
+                dtype=torch.int8,
+            ),
+            requires_grad=False,
+        )
+        set_weight_attrs(qweight, {"input_dim": 1, "output_dim": 0})
+
+        num_groups = input_size_per_partition // self.group_size
+        if self.precision == "nvfp4":
+            scale_dtype = torch.float8_e4m3fn
+        else:
+            scale_dtype = params_dtype
+        wscales = Parameter(
+            torch.empty(num_groups, output_size_per_partition, dtype=scale_dtype),
+            requires_grad=False,
+        )
+
+        smooth_factor = Parameter(
+            torch.empty(input_size_per_partition, dtype=params_dtype),
+            requires_grad=False,
+        )
+
+        smooth_factor_orig = Parameter(
+            torch.empty(input_size_per_partition, dtype=params_dtype),
+            requires_grad=False,
+        )
+
+        proj_down = Parameter(
+            torch.empty(input_size_per_partition, self.rank, dtype=params_dtype),
+            requires_grad=False,
+        )
+        proj_up = Parameter(
+            torch.empty(output_size_per_partition, self.rank, dtype=params_dtype),
+            requires_grad=False,
+        )
+
+        if self.precision == "nvfp4":
+            wcscales = Parameter(
+                torch.empty(
+                    output_size_per_partition,
+                    dtype=params_dtype,
+                ),
+                requires_grad=False,
+            )
+            wtscale = Parameter(
+                torch.empty(1, dtype=params_dtype),
+                requires_grad=False,
+            )
+        else:
+            wcscales = None
+            wtscale = None
+
+        layer.register_parameter("qweight", qweight)
+        layer.register_parameter("wscales", wscales)
+        layer.register_parameter("smooth_factor", smooth_factor)
+        layer.register_parameter("smooth_factor_orig", smooth_factor_orig)
+        layer.register_parameter("proj_down", proj_down)
+        layer.register_parameter("proj_up", proj_up)
+        if wcscales is not None:
+            layer.register_parameter("wcscales", wcscales)
+        if wtscale is not None:
+            layer.register_parameter("wtscale", wtscale)
+
+        layer.input_size_per_partition = input_size_per_partition
+        layer.output_size_per_partition = output_size_per_partition
+        layer.precision = self.precision
+        layer.rank = self.rank
+        layer.group_size = self.group_size
+        layer.act_unsigned = self.act_unsigned
+
+        weight_loader = extra_weight_attrs.get("weight_loader")
+        if weight_loader is not None:
+            set_weight_attrs(qweight, {"weight_loader": weight_loader})
+            set_weight_attrs(wscales, {"weight_loader": weight_loader})
+            set_weight_attrs(smooth_factor, {"weight_loader": weight_loader})
+            set_weight_attrs(smooth_factor_orig, {"weight_loader": weight_loader})
+            set_weight_attrs(proj_down, {"weight_loader": weight_loader})
+            set_weight_attrs(proj_up, {"weight_loader": weight_loader})
+            if wcscales is not None:
+                set_weight_attrs(wcscales, {"weight_loader": weight_loader})
+            if wtscale is not None:
+                set_weight_attrs(wtscale, {"weight_loader": weight_loader})
+
+    def process_weights_after_loading(self, layer: nn.Module) -> None:
+        layer.qweight = Parameter(layer.qweight.data, requires_grad=False)
+        layer.wscales = Parameter(layer.wscales.data, requires_grad=False)
+        layer.smooth_factor = Parameter(layer.smooth_factor.data, requires_grad=False)
+        layer.smooth_factor_orig = Parameter(
+            layer.smooth_factor_orig.data, requires_grad=False
+        )
+        layer.proj_down = Parameter(layer.proj_down.data, requires_grad=False)
+        layer.proj_up = Parameter(layer.proj_up.data, requires_grad=False)
+        if hasattr(layer, "wcscales") and layer.wcscales is not None:
+            layer.wcscales = Parameter(layer.wcscales.data, requires_grad=False)
+        if hasattr(layer, "wtscale") and layer.wtscale is not None:
+            layer.wtscale = Parameter(layer.wtscale.data, requires_grad=False)
+
+        alpha: float | None = None
+        wtscale = getattr(layer, "wtscale", None)
+        if wtscale is not None:
+            if isinstance(wtscale, Parameter):
+                wtscale = wtscale.data
+            if isinstance(wtscale, torch.Tensor):
+                alpha = float(wtscale.detach().cpu().item())
+            else:
+                alpha = float(wtscale)
+        layer._nunchaku_alpha = alpha
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+
+        orig_shape = x.shape
+        x_2d = x.reshape(-1, orig_shape[-1])
+        quantized_x, ascales, lora_act_out = svdq_quantize_w4a4_act_fuse_lora_cuda(
+            x_2d,
+            lora_down=layer.proj_down,
+            smooth=layer.smooth_factor,
+            fp4=layer.precision == "nvfp4",
+            pad_size=256,
+        )
+        out_2d = torch.empty(
+            x_2d.shape[0],
+            layer.output_size_per_partition,
+            dtype=x_2d.dtype,
+            device=x_2d.device,
+        )
+        alpha: float | None = getattr(layer, "_nunchaku_alpha", None)
+        wcscales = getattr(layer, "wcscales", None)
+
+        svdq_gemm_w4a4_cuda(
+            act=quantized_x,
+            wgt=layer.qweight,
+            out=out_2d,
+            ascales=ascales,
+            wscales=layer.wscales,
+            lora_act_in=lora_act_out,
+            lora_up=layer.proj_up,
+            bias=bias,
+            fp4=layer.precision == "nvfp4",
+            alpha=alpha,
+            wcscales=wcscales,
+            act_unsigned=getattr(layer, "act_unsigned", False),
+        )
+        out = out_2d.reshape(*orig_shape[:-1], layer.output_size_per_partition)
+        return out
+
+
+class NunchakuAWQLinearMethod(LinearMethodBase):
+    def __init__(self, group_size: int = 64):
+        self.group_size = group_size
+        self.pack_factor = 8
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes: List[int],
+        input_size: int,
+        output_size: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ) -> None:
+        output_size_per_partition = sum(output_partition_sizes)
+
+        qweight = Parameter(
+            torch.empty(
+                output_size_per_partition // 4,
+                input_size_per_partition // 2,
+                dtype=torch.int32,
+            ),
+            requires_grad=False,
+        )
+        set_weight_attrs(qweight, {"input_dim": 1, "output_dim": 0})
+
+        num_groups = input_size_per_partition // self.group_size
+        wscales = Parameter(
+            torch.empty(num_groups, output_size_per_partition, dtype=params_dtype),
+            requires_grad=False,
+        )
+
+        wzeros = Parameter(
+            torch.empty(num_groups, output_size_per_partition, dtype=params_dtype),
+            requires_grad=False,
+        )
+
+        layer.register_parameter("qweight", qweight)
+        layer.register_parameter("wscales", wscales)
+        layer.register_parameter("wzeros", wzeros)
+
+        layer.input_size_per_partition = input_size_per_partition
+        layer.output_size_per_partition = output_size_per_partition
+        layer.group_size = self.group_size
+        layer.pack_factor = self.pack_factor
+
+        weight_loader = extra_weight_attrs.get("weight_loader")
+        if weight_loader is not None:
+            set_weight_attrs(qweight, {"weight_loader": weight_loader})
+            set_weight_attrs(wscales, {"weight_loader": weight_loader})
+            set_weight_attrs(wzeros, {"weight_loader": weight_loader})
+
+    def process_weights_after_loading(self, layer: nn.Module) -> None:
+        layer.qweight = Parameter(layer.qweight.data, requires_grad=False)
+        layer.wscales = Parameter(layer.wscales.data, requires_grad=False)
+        layer.wzeros = Parameter(layer.wzeros.data, requires_grad=False)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+
+        orig_shape = x.shape
+        x_2d = x.reshape(-1, orig_shape[-1])
+
+        in_features = layer.input_size_per_partition
+        out_features = layer.output_size_per_partition
+        out_2d = awq_gemv_w4a16_cuda(
+            in_feats=x_2d,
+            kernel=layer.qweight,
+            scaling_factors=layer.wscales,
+            zeros=layer.wzeros,
+            m=x_2d.shape[0],
+            n=out_features,
+            k=in_features,
+            group_size=layer.group_size,
+        )
+        if bias is not None:
+            view_shape = [1] * (out_2d.ndim - 1) + [-1]
+            out_2d.add_(bias.view(view_shape))
+
+        out = out_2d.reshape(*orig_shape[:-1], out_features)
+        return out

python/sglang/multimodal_gen/runtime/layers/quantization/weight_only_fp8.py

@@ -0,0 +1,437 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import sglang.multimodal_gen.envs as envs
+from sglang.multimodal_gen.runtime.distributed import (
+    divide,
+    get_tp_group,
+    split_tensor_along_last_dim,
+    tensor_model_parallel_all_gather,
+    tensor_model_parallel_all_reduce,
+)
+from sglang.multimodal_gen.runtime.layers.utils import get_group_rank, get_group_size
+from sglang.multimodal_gen.runtime.utils.weight_attrs import set_weight_attrs
+
+FP8_WEIGHT_DTYPE = torch.float8_e4m3fn
+W8A8_FP8_GEMM_ENV = "SGLANG_DIFFUSION_ENABLE_W8A8_FP8_GEMM"
+
+logger = logging.getLogger(__name__)
+_w8a8_fp8_gemm_warning_logged = False
+
+
+def _can_apply_fused_w8a8_fp8_linear(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    compute_dtype: torch.dtype,
+) -> bool:
+    return (
+        x.device.type == "cuda"
+        and weight.device.type == "cuda"
+        and weight_scale.device.type == "cuda"
+        and not x.is_meta
+        and not weight.is_meta
+        and not weight_scale.is_meta
+        and compute_dtype in (torch.float16, torch.bfloat16)
+    )
+
+
+def dequantize_rowwise_fp8_weight(
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    if weight.ndim != 2:
+        raise ValueError(f"FP8 linear weight must be 2-D, got shape {weight.shape}")
+    if weight_scale.ndim != 1 or weight_scale.shape[0] != weight.shape[0]:
+        raise ValueError(
+            "FP8 row-wise scale must have shape (out_features,), "
+            f"got weight={tuple(weight.shape)} scale={tuple(weight_scale.shape)}"
+        )
+    return weight.to(dtype) * weight_scale.to(dtype).unsqueeze(1)
+
+
+def _apply_srt_w8a8_fp8_linear(*args, **kwargs) -> torch.Tensor:
+    from sglang.srt.layers.quantization.fp8_utils import apply_fp8_linear
+
+    return apply_fp8_linear(*args, **kwargs)
+
+
+def _is_cutlass_fp8_supported() -> bool:
+    from sglang.srt.layers.quantization.fp8_utils import cutlass_fp8_supported
+
+    return cutlass_fp8_supported()
+
+
+def _apply_weight_only_fp8_linear(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    bias: torch.Tensor | None,
+    compute_dtype: torch.dtype,
+    enable_fused_w8a8: bool,
+) -> torch.Tensor:
+    x = x.to(compute_dtype)
+    bias = bias.to(compute_dtype) if bias is not None else None
+    if enable_fused_w8a8 and _can_apply_fused_w8a8_fp8_linear(
+        x, weight, weight_scale, compute_dtype
+    ):
+        try:
+            # The fused kernel uses W8A8 compute; fallback keeps BF16/FP16
+            # activations after dequantizing the FP8 weights.
+            output = _apply_srt_w8a8_fp8_linear(
+                input=x,
+                weight=weight.t(),
+                weight_scale=weight_scale,
+                input_scale=None,
+                bias=bias,
+                cutlass_fp8_supported=_is_cutlass_fp8_supported(),
+            )
+            _log_w8a8_fp8_gemm_warning_once()
+            return output
+        except (ImportError, NotImplementedError):
+            pass
+
+    dequant_weight = dequantize_rowwise_fp8_weight(weight, weight_scale, compute_dtype)
+    return F.linear(x, dequant_weight, bias)
+
+
+class WeightOnlyFP8Linear(nn.Module):
+    """Storage-only e4m3 FP8 linear with row-wise weight scales."""
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        compute_dtype: torch.dtype | None = None,
+        enable_fused_w8a8: bool | None = None,
+    ) -> None:
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.compute_dtype = compute_dtype
+        self.enable_fused_w8a8 = _resolve_enable_fused_w8a8(enable_fused_w8a8)
+        self.weight = nn.Parameter(
+            torch.empty(out_features, in_features, dtype=FP8_WEIGHT_DTYPE),
+            requires_grad=False,
+        )
+        self.weight_scale = nn.Parameter(
+            torch.empty(out_features, dtype=torch.float32),
+            requires_grad=False,
+        )
+        set_weight_attrs(self.weight_scale, {"missing_param_init": "error"})
+        if bias:
+            self.bias = nn.Parameter(
+                torch.empty(
+                    out_features, dtype=compute_dtype or torch.get_default_dtype()
+                ),
+                requires_grad=False,
+            )
+        else:
+            self.register_parameter("bias", None)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        compute_dtype = self.compute_dtype or x.dtype
+        return _apply_weight_only_fp8_linear(
+            x,
+            self.weight,
+            self.weight_scale,
+            self.bias,
+            compute_dtype,
+            self.enable_fused_w8a8,
+        )
+
+
+class WeightOnlyFP8ColumnParallelLinear(nn.Module):
+    """Column-parallel storage-only e4m3 FP8 linear."""
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        compute_dtype: torch.dtype | None = None,
+        gather_output: bool = True,
+        tp_group=None,
+        enable_fused_w8a8: bool | None = None,
+    ) -> None:
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.compute_dtype = compute_dtype
+        self.gather_output = gather_output
+        self.enable_fused_w8a8 = _resolve_enable_fused_w8a8(enable_fused_w8a8)
+        self.tp_group = tp_group or get_tp_group()
+        self.tp_size = get_group_size(self.tp_group)
+        self.tp_rank = get_group_rank(self.tp_group)
+        self.out_features_per_partition = divide(out_features, self.tp_size)
+        self.weight = nn.Parameter(
+            torch.empty(
+                self.out_features_per_partition,
+                in_features,
+                dtype=FP8_WEIGHT_DTYPE,
+            ),
+            requires_grad=False,
+        )
+        set_weight_attrs(
+            self.weight,
+            {
+                "output_dim": 0,
+                "weight_loader": self.weight_loader,
+            },
+        )
+        self.weight_scale = nn.Parameter(
+            torch.empty(self.out_features_per_partition, dtype=torch.float32),
+            requires_grad=False,
+        )
+        set_weight_attrs(
+            self.weight_scale,
+            {
+                "missing_param_init": "error",
+                "output_dim": 0,
+                "weight_loader": self.weight_loader,
+            },
+        )
+        if bias:
+            self.bias = nn.Parameter(
+                torch.empty(
+                    self.out_features_per_partition,
+                    dtype=compute_dtype or torch.get_default_dtype(),
+                ),
+                requires_grad=False,
+            )
+            set_weight_attrs(
+                self.bias,
+                {
+                    "output_dim": 0,
+                    "weight_loader": self.weight_loader,
+                },
+            )
+        else:
+            self.register_parameter("bias", None)
+
+    def weight_loader(
+        self, param: torch.nn.Parameter, loaded_weight: torch.Tensor
+    ) -> None:
+        output_dim = getattr(param, "output_dim", None)
+        if output_dim is not None:
+            shard_size = param.data.shape[output_dim]
+            loaded_weight = loaded_weight.narrow(
+                output_dim, self.tp_rank * shard_size, shard_size
+            )
+        if len(loaded_weight.shape) == 0:
+            loaded_weight = loaded_weight.reshape(1)
+        assert param.data.shape == loaded_weight.shape
+        param.data.copy_(loaded_weight)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        compute_dtype = self.compute_dtype or x.dtype
+        output_parallel = _apply_weight_only_fp8_linear(
+            x,
+            self.weight,
+            self.weight_scale,
+            self.bias,
+            compute_dtype,
+            self.enable_fused_w8a8,
+        )
+        if self.gather_output:
+            return tensor_model_parallel_all_gather(
+                output_parallel, tp_group=self.tp_group
+            )
+        return output_parallel
+
+
+class WeightOnlyFP8MergedColumnParallelLinear(WeightOnlyFP8ColumnParallelLinear):
+    """Column-parallel storage-only FP8 packed linear."""
+
+    def __init__(
+        self,
+        in_features: int,
+        output_sizes: list[int],
+        bias: bool = True,
+        compute_dtype: torch.dtype | None = None,
+        gather_output: bool = False,
+        tp_group=None,
+        enable_fused_w8a8: bool | None = None,
+    ) -> None:
+        self.output_sizes = output_sizes
+        super().__init__(
+            in_features,
+            sum(output_sizes),
+            bias=bias,
+            compute_dtype=compute_dtype,
+            gather_output=gather_output,
+            tp_group=tp_group,
+            enable_fused_w8a8=enable_fused_w8a8,
+        )
+        assert all(output_size % self.tp_size == 0 for output_size in output_sizes)
+
+    def weight_loader(
+        self, param: torch.nn.Parameter, loaded_weight: torch.Tensor
+    ) -> None:
+        output_dim = getattr(param, "output_dim", None)
+        if output_dim is not None:
+            shards = []
+            current_offset = 0
+            for output_size in self.output_sizes:
+                loaded_shard = loaded_weight.narrow(
+                    output_dim, current_offset, output_size
+                )
+                shard_size = output_size // self.tp_size
+                loaded_shard = loaded_shard.narrow(
+                    output_dim, self.tp_rank * shard_size, shard_size
+                )
+                shards.append(loaded_shard)
+                current_offset += output_size
+            loaded_weight = torch.cat(shards, dim=output_dim)
+        if len(loaded_weight.shape) == 0:
+            loaded_weight = loaded_weight.reshape(1)
+        assert param.data.shape == loaded_weight.shape
+        param.data.copy_(loaded_weight)
+
+
+class WeightOnlyFP8RowParallelLinear(nn.Module):
+    """Row-parallel storage-only e4m3 FP8 linear."""
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        compute_dtype: torch.dtype | None = None,
+        input_is_parallel: bool = True,
+        reduce_results: bool = True,
+        tp_group=None,
+        enable_fused_w8a8: bool | None = None,
+    ) -> None:
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.compute_dtype = compute_dtype
+        self.input_is_parallel = input_is_parallel
+        self.reduce_results = reduce_results
+        self.enable_fused_w8a8 = _resolve_enable_fused_w8a8(enable_fused_w8a8)
+        self.tp_group = tp_group or get_tp_group()
+        self.tp_size = get_group_size(self.tp_group)
+        self.tp_rank = get_group_rank(self.tp_group)
+        self.in_features_per_partition = divide(in_features, self.tp_size)
+        self.weight = nn.Parameter(
+            torch.empty(
+                out_features,
+                self.in_features_per_partition,
+                dtype=FP8_WEIGHT_DTYPE,
+            ),
+            requires_grad=False,
+        )
+        set_weight_attrs(
+            self.weight,
+            {
+                "input_dim": 1,
+                "weight_loader": self.weight_loader,
+            },
+        )
+        self.weight_scale = nn.Parameter(
+            torch.empty(out_features, dtype=torch.float32),
+            requires_grad=False,
+        )
+        set_weight_attrs(
+            self.weight_scale,
+            {
+                "missing_param_init": "error",
+                "weight_loader": self.weight_loader,
+            },
+        )
+        if bias:
+            self.bias = nn.Parameter(
+                torch.empty(
+                    out_features, dtype=compute_dtype or torch.get_default_dtype()
+                ),
+                requires_grad=False,
+            )
+            set_weight_attrs(self.bias, {"weight_loader": self.weight_loader})
+        else:
+            self.register_parameter("bias", None)
+
+    def weight_loader(
+        self, param: torch.nn.Parameter, loaded_weight: torch.Tensor
+    ) -> None:
+        input_dim = getattr(param, "input_dim", None)
+        if input_dim is not None:
+            shard_size = param.data.shape[input_dim]
+            loaded_weight = loaded_weight.narrow(
+                input_dim, self.tp_rank * shard_size, shard_size
+            )
+        if len(loaded_weight.shape) == 0:
+            loaded_weight = loaded_weight.reshape(1)
+        assert param.data.shape == loaded_weight.shape
+        param.data.copy_(loaded_weight)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.input_is_parallel:
+            input_parallel = x
+        else:
+            input_parallel = split_tensor_along_last_dim(
+                x, num_partitions=self.tp_size
+            )[self.tp_rank].contiguous()
+
+        compute_dtype = self.compute_dtype or x.dtype
+        bias = None if self.tp_rank > 0 else self.bias
+        output_parallel = _apply_weight_only_fp8_linear(
+            input_parallel,
+            self.weight,
+            self.weight_scale,
+            bias,
+            compute_dtype,
+            self.enable_fused_w8a8,
+        )
+        if self.reduce_results and self.tp_size > 1:
+            return tensor_model_parallel_all_reduce(
+                output_parallel, tp_group=self.tp_group
+            )
+        return output_parallel
+
+
+def _resolve_enable_fused_w8a8(value: bool | None) -> bool:
+    if value is not None:
+        return value
+    return envs.SGLANG_DIFFUSION_ENABLE_W8A8_FP8_GEMM
+
+
+def _log_w8a8_fp8_gemm_warning_once() -> None:
+    global _w8a8_fp8_gemm_warning_logged
+    if _w8a8_fp8_gemm_warning_logged:
+        return
+    logger.warning(
+        "%s=1 enables W8A8 FP8 GEMM for weight-only FP8 linears; activations "
+        "are dynamically quantized to FP8 and outputs may differ from the "
+        "official weight-only FP8 path.",
+        W8A8_FP8_GEMM_ENV,
+    )
+    _w8a8_fp8_gemm_warning_logged = True
+
+
+def swap_linears_to_weight_only_fp8(module: nn.Module) -> None:
+    """Recursively replace nn.Linear with WeightOnlyFP8Linear.
+
+    Ideogram FP8 checkpoints provide ``<linear>.weight_scale`` for every
+    quantized linear. Swapping before load lets strict state-dict checks verify
+    both the FP8 weight and its row-wise scale.
+    """
+
+    for name, child in list(module.named_children()):
+        if isinstance(child, nn.Linear):
+            replacement = WeightOnlyFP8Linear(
+                child.in_features,
+                child.out_features,
+                bias=child.bias is not None,
+                compute_dtype=child.weight.dtype,
+            )
+            setattr(module, name, replacement)
+        else:
+            swap_linears_to_weight_only_fp8(child)

python/sglang/multimodal_gen/runtime/layers/rotary_embedding/__init__.py

@@ -0,0 +1,54 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/rotary_embedding.py
+
+# Adapted from
+# https://github.com/huggingface/transformers/blob/v4.33.2/src/transformers/models/llama/modeling_llama.py
+# Copyright 2023 The vLLM team.
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Rotary Positional Embeddings — unified public API (drop-in replacement)."""
+
+from .base import RotaryEmbedding
+from .factory import get_rope, get_rotary_pos_embed
+from .mrope import (
+    NDRotaryEmbedding,
+    Qwen3VLTextRotaryEmbedding,
+    qwen3_apply_rotary_pos_emb,
+)
+from .utils import (
+    _apply_rotary_emb,
+    apply_flashinfer_rope_qk_inplace,
+)
+
+__all__ = [
+    # _utils
+    "_apply_rotary_emb",
+    "apply_flashinfer_rope_qk_inplace",
+    # _base
+    "RotaryEmbedding",
+    # _mrope
+    "NDRotaryEmbedding",
+    "Qwen3VLTextRotaryEmbedding",
+    "qwen3_apply_rotary_pos_emb",
+    # _factory
+    "get_rope",
+    "get_rotary_pos_embed",
+]

python/sglang/multimodal_gen/runtime/layers/rotary_embedding/base.py

@@ -0,0 +1,133 @@
+"""RotaryEmbedding base class and LinearScalingRotaryEmbedding variant."""
+
+import torch
+
+from sglang.multimodal_gen.runtime.layers.custom_op import CustomOp
+
+from .utils import _apply_rotary_emb
+
+
+@CustomOp.register("rotary_embedding")
+class RotaryEmbedding(CustomOp):
+    """Original rotary positional embedding."""
+
+    def __init__(
+        self,
+        head_size: int,
+        rotary_dim: int,
+        max_position_embeddings: int,
+        base: int | float,
+        is_neox_style: bool,
+        dtype: torch.dtype,
+    ) -> None:
+        super().__init__()
+        self.head_size = head_size
+        self.rotary_dim = rotary_dim
+        self.max_position_embeddings = max_position_embeddings
+        self.base = base
+        self.is_neox_style = is_neox_style
+        self.dtype = dtype
+
+        cache = self._compute_cos_sin_cache()
+        cache = cache.to(dtype)
+        self.cos_sin_cache: torch.Tensor
+        self.register_buffer("cos_sin_cache", cache, persistent=False)
+
+    def _compute_inv_freq(self, base: int | float) -> torch.Tensor:
+        """Compute the inverse frequency."""
+        # NOTE(woosuk): To exactly match the HF implementation, we need to
+        # use CPU to compute the cache and then move it to GPU. However, we
+        # create the cache on GPU for faster initialization. This may cause
+        # a slight numerical difference between the HF implementation and ours.
+        inv_freq = 1.0 / (
+            base
+            ** (
+                torch.arange(0, self.rotary_dim, 2, dtype=torch.float) / self.rotary_dim
+            )
+        )
+        return inv_freq
+
+    def _compute_cos_sin_cache(self) -> torch.Tensor:
+        """Compute the cos and sin cache."""
+        inv_freq = self._compute_inv_freq(self.base)
+        t = torch.arange(self.max_position_embeddings, dtype=torch.float)
+
+        freqs = torch.einsum("i,j -> ij", t, inv_freq)
+        cos = freqs.cos()
+        sin = freqs.sin()
+        cache = torch.cat((cos, sin), dim=-1)
+        return cache
+
+    def forward_cuda(self, *args, **kwargs):
+        return self.forward_native(*args, **kwargs)
+
+    def forward_xpu(self, *args, **kwargs):
+        return self.forward_native(*args, **kwargs)
+
+    def forward_native(
+        self,
+        positions: torch.Tensor,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        offsets: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """A PyTorch-native implementation of forward()."""
+        if offsets is not None:
+            positions = positions + offsets
+        positions = positions.flatten()
+        num_tokens = positions.shape[0]
+        cos_sin = self.cos_sin_cache.index_select(0, positions)
+        cos, sin = cos_sin.chunk(2, dim=-1)
+
+        query_shape = query.shape
+        query = query.reshape(num_tokens, -1, self.head_size)
+        query_rot = query[..., : self.rotary_dim]
+        query_pass = query[..., self.rotary_dim :]
+        query_rot = _apply_rotary_emb(query_rot, cos, sin, self.is_neox_style)
+        query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
+
+        key_shape = key.shape
+        key = key.reshape(num_tokens, -1, self.head_size)
+        key_rot = key[..., : self.rotary_dim]
+        key_pass = key[..., self.rotary_dim :]
+        key_rot = _apply_rotary_emb(key_rot, cos, sin, self.is_neox_style)
+        key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
+        return query, key
+
+    def extra_repr(self) -> str:
+        s = f"head_size={self.head_size}, rotary_dim={self.rotary_dim}"
+        s += f", max_position_embeddings={self.max_position_embeddings}"
+        s += f", base={self.base}, is_neox_style={self.is_neox_style}"
+        return s
+
+
+class LinearScalingRotaryEmbedding(RotaryEmbedding):
+    def __init__(
+        self,
+        head_size: int,
+        rotary_dim: int,
+        max_position_embeddings: int,
+        base: int | float,
+        is_neox_style: bool,
+        dtype: torch.dtype,
+        scaling_factor: float,
+    ) -> None:
+        self.scaling_factor = float(scaling_factor)
+        super().__init__(
+            head_size=head_size,
+            rotary_dim=rotary_dim,
+            max_position_embeddings=max_position_embeddings,
+            base=base,
+            is_neox_style=is_neox_style,
+            dtype=dtype,
+        )
+
+    def _compute_cos_sin_cache(self) -> torch.Tensor:
+        inv_freq = self._compute_inv_freq(self.base)
+        t = torch.arange(self.max_position_embeddings, dtype=torch.float)
+        t = t / self.scaling_factor
+        freqs = torch.einsum("i,j -> ij", t, inv_freq)
+        cos = freqs.cos()
+        sin = freqs.sin()
+        cache = torch.cat((cos, sin), dim=-1)
+        return cache

python/sglang/multimodal_gen/runtime/layers/rotary_embedding/factory.py

@@ -0,0 +1,171 @@
+"""get_rope / get_rotary_pos_embed factory functions and module-level caches."""
+
+from collections import OrderedDict
+from typing import Any
+
+import torch
+
+from .base import LinearScalingRotaryEmbedding, RotaryEmbedding
+from .mrope import NDRotaryEmbedding, _to_tuple
+
+_ROPE_DICT: dict[tuple, RotaryEmbedding] = {}
+_ND_ROPE_CACHE: "OrderedDict[tuple, NDRotaryEmbedding]" = OrderedDict()
+_ROPE_3D_CACHE: "OrderedDict[tuple, tuple[torch.Tensor, torch.Tensor]]" = OrderedDict()
+
+
+def get_rope(
+    head_size: int,
+    rotary_dim: int,
+    max_position: int,
+    base: int | float,
+    is_neox_style: bool = True,
+    rope_scaling: dict[str, Any] | None = None,
+    dtype: torch.dtype | None = None,
+    partial_rotary_factor: float = 1.0,
+) -> RotaryEmbedding:
+    if dtype is None:
+        dtype = torch.get_default_dtype()
+    if rope_scaling is not None:
+        # Transforms every value that is a list into a tuple for caching calls
+        rope_scaling_tuple = {
+            k: tuple(v) if isinstance(v, list) else v for k, v in rope_scaling.items()
+        }
+        rope_scaling_args = tuple(rope_scaling_tuple.items())
+    else:
+        rope_scaling_args = None
+    if partial_rotary_factor < 1.0:
+        rotary_dim = int(rotary_dim * partial_rotary_factor)
+    max_position_embeddings = max_position
+    rope_type = None
+    if rope_scaling is not None:
+        rope_type = rope_scaling.get("rope_type", rope_scaling.get("type", None))
+        if rope_type in (None, "default"):
+            rope_scaling = None
+        elif rope_type == "linear":
+            factor = float(rope_scaling.get("factor", 1.0))
+            original_max = rope_scaling.get("original_max_position_embeddings", None)
+            if original_max is not None:
+                max_position_embeddings = max(
+                    max_position_embeddings, int(float(original_max) * factor)
+                )
+    key = (
+        head_size,
+        rotary_dim,
+        max_position_embeddings,
+        base,
+        is_neox_style,
+        rope_scaling_args,
+        dtype,
+    )
+    if key in _ROPE_DICT:
+        return _ROPE_DICT[key]
+
+    if rope_scaling is None:
+        rotary_emb = RotaryEmbedding(
+            head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype
+        )
+    else:
+        if rope_type == "linear":
+            factor = float(rope_scaling.get("factor", 1.0))
+            rotary_emb = LinearScalingRotaryEmbedding(
+                head_size=head_size,
+                rotary_dim=rotary_dim,
+                max_position_embeddings=max_position_embeddings,
+                base=base,
+                is_neox_style=is_neox_style,
+                dtype=dtype,
+                scaling_factor=factor,
+            )
+        else:
+            raise ValueError(f"Unknown RoPE scaling {rope_scaling}")
+    _ROPE_DICT[key] = rotary_emb
+    return rotary_emb
+
+
+def get_rotary_pos_embed(
+    rope_sizes,
+    hidden_size,
+    heads_num,
+    rope_dim_list,
+    rope_theta,
+    theta_rescale_factor=1.0,
+    interpolation_factor=1.0,
+    shard_dim: int = 0,
+    dtype: torch.dtype = torch.float32,
+    start_frame: int = 0,
+    device: torch.device | str | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Generate rotary positional embeddings for the given sizes.
+
+    Args:
+        rope_sizes: Tuple of dimensions (t, h, w)
+        hidden_size: Hidden dimension size
+        heads_num: Number of attention heads
+        rope_dim_list: List of dimensions for each axis, or None
+        rope_theta: Base for frequency calculations
+        theta_rescale_factor: Rescale factor for theta. Defaults to 1.0
+        interpolation_factor: Factor to scale positions. Defaults to 1.0
+        shard_dim: Which dimension to shard for sequence parallelism. Defaults to 0.
+
+    Returns:
+        Tuple of (cos, sin) tensors for rotary embeddings
+    """
+
+    target_ndim = 3
+    head_dim = hidden_size // heads_num
+
+    if rope_dim_list is None:
+        rope_dim_list = [head_dim // target_ndim for _ in range(target_ndim)]
+
+    assert (
+        sum(rope_dim_list) == head_dim
+    ), "sum(rope_dim_list) should equal to head_dim of attention layer"
+
+    # Get SP info - now handled within NDRotaryEmbedding
+    # sp_group = get_sp_group()
+    # sp_rank = sp_group.rank_in_group
+    # sp_world_size = sp_group.world_size
+
+    # Simple LRU cache keyed by parameters
+    global _ND_ROPE_CACHE
+    key = (
+        tuple(rope_dim_list),
+        float(rope_theta),
+        (
+            tuple(theta_rescale_factor)
+            if isinstance(theta_rescale_factor, list)
+            else float(theta_rescale_factor)
+        ),
+        (
+            tuple(interpolation_factor)
+            if isinstance(interpolation_factor, list)
+            else float(interpolation_factor)
+        ),
+        dtype,
+    )
+
+    cache_hit = key in _ND_ROPE_CACHE
+    if cache_hit:
+        rope_emb = _ND_ROPE_CACHE.pop(key)
+        _ND_ROPE_CACHE[key] = rope_emb  # move to end (most-recent)
+    else:
+        rope_emb = NDRotaryEmbedding(
+            rope_dim_list=rope_dim_list,
+            rope_theta=rope_theta,
+            theta_rescale_factor=theta_rescale_factor,
+            interpolation_factor=interpolation_factor,
+            dtype=dtype,
+        )
+        _ND_ROPE_CACHE[key] = rope_emb
+        if len(_ND_ROPE_CACHE) > 16:
+            # pop least-recently-used
+            _ND_ROPE_CACHE.pop(next(iter(_ND_ROPE_CACHE)))
+
+    freqs_cos, freqs_sin = rope_emb.forward_from_grid(
+        grid_size=_to_tuple(rope_sizes, dim=3),
+        shard_dim=shard_dim,
+        start_frame=start_frame,
+        device=device,
+    )
+    return freqs_cos, freqs_sin

python/sglang/multimodal_gen/runtime/layers/rotary_embedding/mrope.py

@@ -0,0 +1,502 @@
+"""MRotaryEmbedding, YaRNScalingMRotaryEmbedding, NDRotaryEmbedding, OneDRotaryEmbedding."""
+
+import functools
+
+import torch
+
+from sglang.multimodal_gen.runtime.distributed.parallel_state import get_sp_group
+
+
+def _to_tuple(x: int | tuple[int, ...], dim: int = 2) -> tuple[int, ...]:
+    if isinstance(x, int):
+        return (x,) * dim
+    elif len(x) == dim:
+        return x
+    else:
+        raise ValueError(f"Expected length {dim} or int, but got {x}")
+
+
+def get_1d_rotary_pos_embed(
+    dim: int,
+    pos: torch.FloatTensor | int,
+    theta: float = 10000.0,
+    theta_rescale_factor: float = 1.0,
+    interpolation_factor: float = 1.0,
+    dtype: torch.dtype = torch.float32,
+    device: torch.device | str | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Precompute the frequency tensor for complex exponential (cis) with given dimensions.
+    (Note: `cis` means `cos + i * sin`, where i is the imaginary unit.)
+
+    This function calculates a frequency tensor with complex exponential using the given dimension 'dim'
+    and the end index 'end'. The 'theta' parameter scales the frequencies.
+
+    Args:
+        dim (int): Dimension of the frequency tensor.
+        pos (int or torch.FloatTensor): Position indices for the frequency tensor. [S] or scalar
+        theta (float, optional): Scaling factor for frequency computation. Defaults to 10000.0.
+        theta_rescale_factor (float, optional): Rescale factor for theta. Defaults to 1.0.
+        interpolation_factor (float, optional): Factor to scale positions. Defaults to 1.0.
+
+    Returns:
+        freqs_cos, freqs_sin: Precomputed frequency tensor with real and imaginary parts separately. [S, D]
+    """
+    if isinstance(pos, int):
+        pos = torch.arange(pos, dtype=dtype, device=device)
+    elif (
+        isinstance(pos, torch.Tensor)
+        and device is not None
+        and pos.device != torch.device(device)
+    ):
+        # Ensure positions are on the requested device to avoid implicit CPU ops.
+        pos = pos.to(device)
+
+    # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning
+    # has some connection to NTK literature
+    if theta_rescale_factor != 1.0:
+        theta *= theta_rescale_factor ** (dim / (dim - 2))
+
+    freqs = 1.0 / (
+        theta
+        ** (torch.arange(0, dim, 2, device=device)[: (dim // 2)].to(dtype) / dim).to(
+            device=device
+        )
+    )  # [D/2]
+    freqs = torch.outer(pos * interpolation_factor, freqs)  # [S, D/2]
+    freqs_cos = freqs.cos()  # [S, D/2]
+    freqs_sin = freqs.sin()  # [S, D/2]
+    return freqs_cos, freqs_sin
+
+
+def qwen3_apply_rotary_pos_emb(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Apply Qwen3-style RoPE to q/k tensors shaped [B, S, H, D]."""
+    half = q.shape[-1] // 2
+    q1 = q[..., :half]
+    q2 = q[..., half:]
+    q_embed = torch.empty_like(q)
+    q_embed[..., :half] = q1 * cos[..., :half] - q2 * sin[..., :half]
+    q_embed[..., half:] = q2 * cos[..., half:] + q1 * sin[..., half:]
+
+    half = k.shape[-1] // 2
+    k1 = k[..., :half]
+    k2 = k[..., half:]
+    k_embed = torch.empty_like(k)
+    k_embed[..., :half] = k1 * cos[..., :half] - k2 * sin[..., :half]
+    k_embed[..., half:] = k2 * cos[..., half:] + k1 * sin[..., half:]
+    return q_embed, k_embed
+
+
+class Qwen3VLTextRotaryEmbedding(torch.nn.Module):
+    """Qwen3-VL multi-dimensional rotary embedding with interleaved mRoPE."""
+
+    def __init__(
+        self,
+        head_dim: int = 128,
+        rope_theta: float = 5_000_000.0,
+        mrope_section: tuple[int, int, int] | list[int] = (24, 20, 20),
+    ):
+        super().__init__()
+        self.rope_type = "default"
+        self.max_seq_len_cached = 262144
+        self.mrope_section = list(mrope_section)
+        self.head_dim = head_dim
+
+        inv_freq = 1.0 / (
+            rope_theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)
+        )
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.attention_scaling = 1.0
+
+    def apply_interleaved_mrope(
+        self, freqs: torch.Tensor, mrope_section: list[int]
+    ) -> torch.Tensor:
+        freqs_t = freqs[0].clone()
+        for dim, offset in enumerate((1, 2), start=1):
+            length = mrope_section[dim] * 3
+            idx = slice(offset, length, 3)
+            freqs_t[..., idx] = freqs[dim, ..., idx]
+        return freqs_t
+
+    def _normalize_position_ids(self, position_ids: torch.Tensor) -> torch.Tensor:
+        if position_ids.ndim == 3 and position_ids.shape[-1] == 3:
+            position_ids = position_ids.permute(2, 0, 1)
+        elif position_ids.ndim == 2:
+            position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
+        elif position_ids.ndim != 3 or position_ids.shape[0] != 3:
+            raise ValueError(
+                "Qwen3 mRoPE position_ids must have shape [3, B, S], [B, S, 3], "
+                f"or [B, S], got {tuple(position_ids.shape)}"
+            )
+        return position_ids
+
+    def _compute_interleaved_freqs(self, position_ids: torch.Tensor) -> torch.Tensor:
+        position_ids = self._normalize_position_ids(position_ids)
+
+        inv_freq_expanded = (
+            self.inv_freq[None, None, :, None]
+            .float()
+            .expand(3, position_ids.shape[1], -1, 1)
+            .to(position_ids.device)
+        )
+        position_ids_expanded = position_ids[:, :, None, :].float()
+
+        freqs = (inv_freq_expanded @ position_ids_expanded).transpose(2, 3)
+        return self.apply_interleaved_mrope(freqs, self.mrope_section)
+
+    @torch.no_grad()
+    def build_rope_cache_inputs(
+        self, position_ids: torch.Tensor, *, cache_dtype: torch.dtype | None = None
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        freqs = self._compute_interleaved_freqs(position_ids)
+        cos = freqs.cos() * self.attention_scaling
+        sin = freqs.sin() * self.attention_scaling
+        if cache_dtype is not None and cache_dtype != torch.float32:
+            cos = cos.to(cache_dtype).float()
+            sin = sin.to(cache_dtype).float()
+        cos_sin_cache = torch.cat((cos, sin), dim=-1).reshape(-1, self.head_dim)
+        cos_sin_cache = cos_sin_cache.contiguous()
+        cache_positions = torch.arange(
+            cos_sin_cache.shape[0], device=cos_sin_cache.device, dtype=torch.long
+        )
+        return cos_sin_cache, cache_positions
+
+    @torch.no_grad()
+    def forward(
+        self, x: torch.Tensor, position_ids: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Return cos/sin for position IDs shaped [3, B, S], [B, S, 3], or [B, S]."""
+        freqs = self._compute_interleaved_freqs(position_ids)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        cos = emb.cos() * self.attention_scaling
+        sin = emb.sin() * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+class OneDRotaryEmbedding(torch.nn.Module):
+    """1D rotary positional embedding with caching."""
+
+    def __init__(
+        self,
+        dim: int,
+        theta: float = 10000.0,
+        theta_rescale_factor: float = 1.0,
+        interpolation_factor: float = 1.0,
+        dtype: torch.dtype = torch.float32,
+        use_real: bool = False,
+        repeat_interleave_real: bool = False,
+    ):
+        super().__init__()
+        assert dim % 2 == 0
+        self.dim = dim
+        self.theta = theta
+        self.theta_rescale_factor = theta_rescale_factor
+        self.interpolation_factor = interpolation_factor
+        # dtype of freqs
+        self.dtype = dtype
+        self.use_real = use_real
+        self.repeat_interleave_real = repeat_interleave_real
+
+    def build_freqs(self, device):
+        freqs = 1.0 / (
+            self.theta
+            ** (
+                torch.arange(0, self.dim, 2, dtype=self.dtype, device=device)[
+                    : (self.dim // 2)
+                ]
+                / self.dim
+            ).to(device=device)
+        )
+        return freqs
+
+    def build_freqs_outer(self, pos: torch.Tensor, device):
+        theta = self.theta
+        # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning
+        # has some connection to NTK literature
+        if self.theta_rescale_factor != 1.0:
+            theta *= self.theta_rescale_factor ** (self.dim / (self.dim - 2))
+
+        freqs = self.build_freqs(device)
+
+        freqs = torch.outer(pos * self.interpolation_factor, freqs)
+        freqs_cos = freqs.cos()
+        freqs_sin = freqs.sin()
+
+        if self.use_real and self.repeat_interleave_real:
+            freqs_cos = freqs_cos.repeat_interleave(2, dim=1)
+            freqs_sin = freqs_sin.repeat_interleave(2, dim=1)
+
+        return freqs_cos.float(), freqs_sin.float()
+
+    @functools.lru_cache(maxsize=16)
+    def forward_from_grid(
+        self, seq_len: int, start_pos: int, device_str: str
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        device = torch.device(device_str)
+        pos = torch.arange(
+            start_pos, start_pos + seq_len, dtype=self.dtype, device=device
+        )
+
+        freqs_cos, freqs_sin = self.build_freqs_outer(pos, device)
+        return freqs_cos, freqs_sin
+
+    def forward(self, pos: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Calculates 1D rotary embeddings for the given positions.
+
+        This method converts the input tensor to a hashable representation
+        and calls a cached helper method to perform the computation.
+        """
+        pos_tuple = tuple(pos.tolist())
+        device_str = str(pos.device)
+        return self._forward_cached(pos_tuple, device_str)
+
+    @functools.lru_cache(maxsize=16)
+    def _forward_cached(
+        self, pos_tuple: tuple, device_str: str
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        The core implementation that computes 1D rotary embeddings.
+        This method is wrapped by an LRU cache.
+        """
+        device = torch.device(device_str)
+        pos = torch.as_tensor(pos_tuple, dtype=self.dtype, device=device)
+        freqs_cos, freqs_sin = self.build_freqs_outer(pos, device)
+        return freqs_cos, freqs_sin
+
+
+class NDRotaryEmbedding(torch.nn.Module):
+    """N-dimensional rotary positional embedding."""
+
+    def __init__(
+        self,
+        rope_dim_list: list[int],
+        rope_theta: float,
+        theta_rescale_factor: float | list[float] = 1.0,
+        interpolation_factor: float | list[float] = 1.0,
+        use_real: bool = False,
+        repeat_interleave_real: bool = False,
+        dtype: torch.dtype = torch.float32,
+    ):
+        super().__init__()
+        self.rope_dim_list = rope_dim_list
+        self.ndim = len(rope_dim_list)
+        self.rope_theta = rope_theta
+        # dtype of freqs
+        # does not control the output dtype
+        self.dtype = dtype
+
+        if isinstance(theta_rescale_factor, (int, float)):
+            self.theta_rescale_factor = [theta_rescale_factor] * self.ndim
+        elif isinstance(theta_rescale_factor, list) and len(theta_rescale_factor) == 1:
+            self.theta_rescale_factor = [theta_rescale_factor[0]] * self.ndim
+        else:
+            self.theta_rescale_factor = theta_rescale_factor
+        assert (
+            len(self.theta_rescale_factor) == self.ndim
+        ), "len(theta_rescale_factor) should equal to len(rope_dim_list)"
+
+        if isinstance(interpolation_factor, (int, float)):
+            self.interpolation_factor = [interpolation_factor] * self.ndim
+        elif isinstance(interpolation_factor, list) and len(interpolation_factor) == 1:
+            self.interpolation_factor = [interpolation_factor[0]] * self.ndim
+        else:
+            self.interpolation_factor = interpolation_factor
+        assert (
+            len(self.interpolation_factor) == self.ndim
+        ), "len(interpolation_factor) should equal to len(rope_dim_list)"
+
+        self.rope_generators: list[OneDRotaryEmbedding] = torch.nn.ModuleList()
+        _config_to_gen_idx: dict[tuple, int] = {}
+        self.dim_idx_to_gen_idx: list[int] = []
+
+        for i in range(self.ndim):
+            dim = self.rope_dim_list[i]
+            rescale = self.theta_rescale_factor[i]
+            interp = self.interpolation_factor[i]
+
+            config_key = (dim, rescale, interp, use_real, repeat_interleave_real)
+            if config_key not in _config_to_gen_idx:
+                generator = OneDRotaryEmbedding(
+                    dim=dim,
+                    theta=self.rope_theta,
+                    theta_rescale_factor=rescale,
+                    interpolation_factor=interp,
+                    dtype=self.dtype,
+                    use_real=use_real,
+                    repeat_interleave_real=repeat_interleave_real,
+                )
+                _config_to_gen_idx[config_key] = len(self.rope_generators)
+                self.rope_generators.append(generator)
+
+            gen_idx = _config_to_gen_idx[config_key]
+            self.dim_idx_to_gen_idx.append(gen_idx)
+
+    def forward(self, positions: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Calculates n-d rotary embeddings for given absolute positions.
+
+        Args:
+            positions (torch.Tensor): A tensor of shape `[num_tokens, ndim]`
+                containing the integer coordinates for each token.
+
+        Returns:
+            A tuple of (cos, sin) tensors.
+        """
+        # Caching wrapper: convert tensor to a hashable tuple of tuples.
+        pos_tuple = tuple(map(tuple, positions.tolist()))
+        device_str = str(positions.device)
+        return self._forward_cached(pos_tuple, device_str)
+
+    @functools.lru_cache(maxsize=16)
+    def _forward_cached(
+        self, pos_tuple: tuple[tuple[int, ...], ...], device_str: str
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        The core implementation that computes embeddings from a position tensor.
+        This method is wrapped by an LRU cache.
+        """
+        device = torch.device(device_str)
+        positions = torch.tensor(pos_tuple, dtype=torch.long, device=device)
+        return self.forward_uncached(pos=positions)
+
+    def forward_uncached(self, pos: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        The core implementation that computes embeddings from a position tensor.
+        This method is wrapped by an LRU cache.
+        """
+        device = pos.device
+
+        # Pre-allocate the final tensors for efficiency.
+        num_tokens = pos.shape[0]
+        first_generator = self.rope_generators[0]
+        if first_generator.use_real and first_generator.repeat_interleave_real:
+            head_dim = sum(self.rope_dim_list)
+        else:
+            head_dim = sum(self.rope_dim_list) // 2
+
+        cos = torch.empty((num_tokens, head_dim), device=device, dtype=self.dtype)
+        sin = torch.empty((num_tokens, head_dim), device=device, dtype=self.dtype)
+
+        col_offset = 0
+        for i in range(self.ndim):
+            # Extract position coordinates for the current dimension for all tokens.
+            pos_i = pos[:, i].to(self.dtype)
+
+            # Get the appropriate 1D generator.
+            gen_idx = self.dim_idx_to_gen_idx[i]
+            generator = self.rope_generators[gen_idx]
+
+            # Calculate 1D embeddings.
+            cos_1d, sin_1d = generator(pos_i)
+
+            slice_width = cos_1d.shape[1]
+            cos[:, col_offset : col_offset + slice_width] = cos_1d
+            sin[:, col_offset : col_offset + slice_width] = sin_1d
+            col_offset += slice_width
+
+        return cos.float(), sin.float()
+
+    def forward_from_grid(
+        self,
+        grid_size: tuple[int, ...],
+        shard_dim: int = 0,
+        start_frame: int = 0,
+        device: torch.device | str | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Handles sp internally
+        """
+        # Caching wrapper: use grid parameters directly as the key.
+        # grid_tuple = _to_tuple(grid_size, dim=self.ndim)
+        device_str = str(device) if device is not None else "cpu"
+        return self._forward_cached_from_grid(
+            grid_size, shard_dim, start_frame, device_str
+        )
+
+    @functools.lru_cache(maxsize=16)
+    def _forward_cached_from_grid(
+        self,
+        grid_size: tuple[int, ...],
+        shard_dim: int,
+        start_frame: int,
+        device_str: str,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Computes embeddings for a structured grid, using a highly efficient
+        implementation that avoids materializing the full position tensor.
+        This method is wrapped by an LRU cache.
+        """
+        device = torch.device(device_str)
+        sp_group = get_sp_group()
+        sp_rank = sp_group.rank_in_group
+        sp_world_size = sp_group.world_size
+
+        sizes = _to_tuple(grid_size, dim=self.ndim)
+        starts = (0,) * self.ndim
+
+        # Apply sequence parallel sharding to the sizes and compute shard offset
+        shard_sizes = list(sizes)
+        shard_offsets = [0] * self.ndim
+        if sp_world_size > 1:
+            assert sizes[shard_dim] % sp_world_size == 0, (
+                f"Dimension {shard_dim} with size {sizes[shard_dim]} is not divisible "
+                f"by sequence parallel world size {sp_world_size}"
+            )
+            shard_size = sizes[shard_dim] // sp_world_size
+            shard_offsets[shard_dim] = sp_rank * shard_size
+            shard_sizes[shard_dim] = shard_size
+
+        # Pre-allocate outputs on the requested device to avoid CPU ops and extra cats
+        num_tokens = 1
+        for s in shard_sizes:
+            num_tokens *= int(s)
+        head_dim_half = sum(self.rope_dim_list) // 2
+        cos = torch.empty((num_tokens, head_dim_half), device=device, dtype=self.dtype)
+        sin = torch.empty((num_tokens, head_dim_half), device=device, dtype=self.dtype)
+
+        # Compute per-axis 1D embeddings once and expand via repeats to [N, d_i/2]
+        col_offset = 0
+        for i in range(self.ndim):
+            dim_i = self.rope_dim_list[i]
+            dim_i_half = dim_i // 2
+            size_i = int(shard_sizes[i])
+
+            # Starting position for this axis, with optional frame offset for time axis (i==0)
+            base_offset = starts[i]
+            if i == 0 and start_frame > 0:
+                base_offset += start_frame
+            if sp_world_size > 1 and i == shard_dim:
+                base_offset += shard_offsets[i]
+
+            gen_idx = self.dim_idx_to_gen_idx[i]
+            generator = self.rope_generators[gen_idx]
+            cos_1d, sin_1d = generator.forward_from_grid(
+                size_i, base_offset, device_str
+            )
+
+            # Expand to [num_tokens, dim_i/2] matching flatten order (last dims vary fastest)
+            repeats_per_entry = 1
+            for j in range(i + 1, self.ndim):
+                repeats_per_entry *= int(shard_sizes[j])
+            tile_count = 1
+            for j in range(0, i):
+                tile_count *= int(shard_sizes[j])
+
+            cos_expanded = cos_1d.repeat_interleave(repeats_per_entry, dim=0)
+            sin_expanded = sin_1d.repeat_interleave(repeats_per_entry, dim=0)
+            if tile_count > 1:
+                cos_expanded = cos_expanded.repeat(tile_count, 1)
+                sin_expanded = sin_expanded.repeat(tile_count, 1)
+
+            cos[:, col_offset : col_offset + dim_i_half] = cos_expanded
+            sin[:, col_offset : col_offset + dim_i_half] = sin_expanded
+            col_offset += dim_i_half
+
+        return cos.float(), sin.float()

python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py

@@ -0,0 +1,196 @@
+"""Primitive RoPE ops: rotate helpers and apply_rotary_emb utilities."""
+
+from typing import Optional, Tuple
+
+import torch
+
+from sglang.jit_kernel.diffusion.triton.rotary import apply_rotary_embedding
+from sglang.kernel_api_logging import debug_kernel_api
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.srt.utils.custom_op import register_custom_op_from_extern
+
+logger = init_logger(__name__)
+
+_is_cuda = current_platform.is_cuda()
+if _is_cuda:
+    try:
+        from flashinfer.rope import (
+            apply_rope_with_cos_sin_cache_inplace as _flashinfer_apply_rope_inplace,
+        )
+    except Exception:
+        _flashinfer_apply_rope_inplace = None
+else:
+    _flashinfer_apply_rope_inplace = None
+
+if _flashinfer_apply_rope_inplace is not None:
+    flashinfer_apply_rope_inplace = register_custom_op_from_extern(
+        _flashinfer_apply_rope_inplace,
+        op_name="flashinfer_apply_rope_with_cos_sin_cache_inplace",
+        mutates_args=["query", "key"],
+    )
+else:
+    flashinfer_apply_rope_inplace = None
+
+
+def _apply_rotary_emb(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    is_neox_style: bool,
+    interleaved: bool = False,
+) -> torch.Tensor:
+    """
+    Args:
+        x: [num_tokens, num_heads, head_size] or [num_tokens, head_size]
+        cos: [num_tokens, head_size // 2]
+        sin: [num_tokens, head_size // 2]
+        is_neox_style: Whether to use the Neox-style or GPT-J-style rotary
+            positional embeddings.
+    """
+    # cos = cos.unsqueeze(-2).to(x.dtype)
+    # sin = sin.unsqueeze(-2).to(x.dtype)
+    if is_neox_style:
+        cos = cos.unsqueeze(-2)
+        sin = sin.unsqueeze(-2)
+        if is_neox_style:
+            x1, x2 = torch.chunk(x, 2, dim=-1)
+        else:
+            x1 = x[..., ::2]
+            x2 = x[..., 1::2]
+        o1 = (x1.float() * cos - x2.float() * sin).type_as(x)
+        o2 = (x2.float() * cos + x1.float() * sin).type_as(x)
+        return torch.cat((o1, o2), dim=-1)
+    else:
+        return apply_rotary_embedding(x, cos, sin, interleaved)
+
+
+@debug_kernel_api
+def apply_flashinfer_rope_qk_inplace(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    *,
+    head_size: Optional[int] = None,
+    is_neox: bool = False,
+    positions: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    if q.dim() != 4 or k.dim() != 4:
+        raise ValueError(
+            f"Expected q/k to be 4D [bsz, seqlen, nheads, head_size], "
+            f"got q:{tuple(q.shape)} k:{tuple(k.shape)}"
+        )
+    if q.shape[:2] != k.shape[:2] or q.shape[-1] != k.shape[-1]:
+        raise ValueError(
+            f"q and k must share batch, sequence, and head size, got {q.shape} vs {k.shape}"
+        )
+
+    if not (isinstance(cos_sin_cache, torch.Tensor) and cos_sin_cache.dim() == 2):
+        raise ValueError("cos_sin_cache must be a 2D torch.Tensor")
+
+    bsz, seqlen, q_heads, d = q.shape
+    k_heads = k.shape[2]
+    rope_dim = cos_sin_cache.shape[-1]
+    if k.device != q.device or cos_sin_cache.device != q.device:
+        raise ValueError(
+            "q, k, and cos_sin_cache must be on the same device, "
+            f"got q={q.device}, k={k.device}, cos_sin_cache={cos_sin_cache.device}"
+        )
+    if rope_dim % 2 != 0 or rope_dim > d:
+        raise ValueError(
+            f"cos_sin_cache width must be even and <= head_size, got {rope_dim} vs {d}"
+        )
+    if head_size is None:
+        head_size = d
+    if head_size != d:
+        raise ValueError(f"head_size mismatch: inferred {d}, but head_size={head_size}")
+
+    use_flashinfer = (
+        flashinfer_apply_rope_inplace is not None
+        and q.is_cuda
+        and k.is_cuda
+        and cos_sin_cache.is_cuda
+        and q_heads == k_heads
+    )
+
+    if not use_flashinfer:
+        if flashinfer_apply_rope_inplace is None:
+            _warn_about_missing_flashinfer()
+
+        half_size = rope_dim // 2
+        if positions is None:
+            cos = cos_sin_cache[:seqlen, :half_size].to(q.dtype)
+            sin = cos_sin_cache[:seqlen, half_size:].to(q.dtype)
+            cos = cos.unsqueeze(0).expand(bsz, -1, -1).reshape(bsz * seqlen, -1)
+            sin = sin.unsqueeze(0).expand(bsz, -1, -1).reshape(bsz * seqlen, -1)
+        else:
+            positions = positions.to(device=q.device, dtype=torch.long).view(-1)
+            cos = cos_sin_cache[positions, :half_size].to(q.dtype)
+            sin = cos_sin_cache[positions, half_size:].to(q.dtype)
+
+        if current_platform.is_npu():
+            q_flat = q.reshape(bsz * seqlen, q_heads, d)
+            k_flat = k.reshape(bsz * seqlen, k_heads, d)
+            q_rot = apply_rotary_embedding(q_flat, cos, sin, interleaved=not is_neox)
+            k_rot = apply_rotary_embedding(k_flat, cos, sin, interleaved=not is_neox)
+            return q_rot.view(bsz, seqlen, q_heads, d), k_rot.view(
+                bsz, seqlen, k_heads, d
+            )
+
+        def apply_rope_prefix(x: torch.Tensor, num_heads: int) -> torch.Tensor:
+            x_flat = x.reshape(bsz * seqlen, num_heads, d)
+            x_rot = x_flat[..., :rope_dim]
+            out_rot = torch.empty_like(x_rot)
+            cos_b = cos.unsqueeze(-2)
+            sin_b = sin.unsqueeze(-2)
+            if is_neox:
+                x1, x2 = torch.chunk(x_rot, 2, dim=-1)
+                out_rot[..., :half_size] = x1 * cos_b - x2 * sin_b
+                out_rot[..., half_size:] = x2 * cos_b + x1 * sin_b
+            else:
+                x1 = x_rot[..., ::2]
+                x2 = x_rot[..., 1::2]
+                out_rot[..., ::2] = x1 * cos_b - x2 * sin_b
+                out_rot[..., 1::2] = x2 * cos_b + x1 * sin_b
+            if rope_dim == d:
+                return out_rot.view(bsz, seqlen, num_heads, d)
+            out = x_flat.clone()
+            out[..., :rope_dim] = out_rot
+            return out.view(bsz, seqlen, num_heads, d)
+
+        return apply_rope_prefix(q, q_heads), apply_rope_prefix(k, k_heads)
+
+    if positions is None:
+        pos_1d = torch.arange(seqlen, device=q.device, dtype=torch.long)
+        positions = pos_1d if bsz == 1 else pos_1d.repeat(bsz)
+    else:
+        if not (isinstance(positions, torch.Tensor) and positions.dim() == 1):
+            raise ValueError("positions must be a 1D Tensor")
+        if positions.numel() != bsz * seqlen:
+            raise ValueError(
+                f"positions length must be bsz*seqlen={bsz*seqlen}, got {positions.numel()}"
+            )
+        positions = positions.to(device=q.device, dtype=torch.long)
+
+    q_flat = q.reshape(bsz * seqlen, q_heads * d).contiguous()
+    k_flat = k.reshape(bsz * seqlen, k_heads * d).contiguous()
+    flashinfer_apply_rope_inplace(
+        positions=positions,
+        query=q_flat,
+        key=k_flat,
+        head_size=d,
+        cos_sin_cache=cos_sin_cache,
+        is_neox=is_neox,
+    )
+    return q_flat.view(bsz, seqlen, q_heads, d), k_flat.view(bsz, seqlen, k_heads, d)
+
+
+@torch.compiler.assume_constant_result
+def _warn_about_missing_flashinfer():
+    """
+    Function to warn about the missing FlashInfer.
+    Exists to not cause a graph break during the compilation.
+    """
+    logger.warning_once(
+        "FlashInfer not available, using Triton fallback for RoPE",
+    )

python/sglang/multimodal_gen/runtime/layers/usp.py

@@ -0,0 +1,429 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+import logging
+from typing import TYPE_CHECKING
+
+import torch
+import torch.distributed as dist
+import torch.distributed._functional_collectives as ft_c
+from torch.distributed.tensor.experimental._attention import _cp_options
+
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    get_sp_group,
+    get_ulysses_parallel_rank,
+    get_ulysses_parallel_world_size,
+)
+from sglang.srt.utils.common import torch_release
+
+_cp_options.enable_load_balance = False
+
+if TYPE_CHECKING:
+    from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+        AttentionImpl,
+    )
+
+logger = logging.getLogger(__name__)
+
+
+def _maybe_wait(tensor: torch.Tensor) -> torch.Tensor:
+    """
+    When tracing the code, the result tensor is not an AsyncCollectiveTensor,
+    so we cannot call ``wait()``.
+    """
+    if isinstance(tensor, ft_c.AsyncCollectiveTensor):
+        return tensor.wait()
+    return tensor
+
+
+def _usp_all_to_all_single(x: torch.Tensor) -> torch.Tensor:
+    ulysses_pg = get_sp_group().ulysses_group
+    assert ulysses_pg is not None, "Ulysses process group is not initialized."
+    x_shape = x.shape
+    x = x.flatten().contiguous()
+    output = torch.empty_like(x)
+    # USP calls this collective many times per denoising step and waits
+    # immediately, so avoid the extra wrapper overhead of functional collectives.
+    torch.distributed.all_to_all_single(output, x, group=ulysses_pg)
+    return output.reshape(x_shape)
+
+
+def _usp_all_to_all_single_varlen(
+    x: torch.Tensor,
+    output_split_sizes: list[int],
+    input_split_sizes: list[int],
+) -> torch.Tensor:
+    ulysses_pg = get_sp_group().ulysses_group
+    assert ulysses_pg is not None, "Ulysses process group is not initialized."
+    x = x.flatten().contiguous()
+    output = torch.empty(sum(output_split_sizes), dtype=x.dtype, device=x.device)
+    dist.all_to_all_single(
+        output,
+        x,
+        output_split_sizes=output_split_sizes,
+        input_split_sizes=input_split_sizes,
+        group=ulysses_pg,
+    )
+    return output
+
+
+def _usp_input_all_to_all(x: torch.Tensor, head_dim: int = 1) -> torch.Tensor:
+    """
+    Perform Ulysses-style input all-to-all over the head dimension.
+
+    Default layout expects heads at dim=1 and sequence at dim=2:
+        [b, h, s_local, d] -> [b, h_local, s_global, d]
+
+    If heads are at dim=2 (input is [b, s_local, h, d]), set head_dim=2, and the
+    function returns [b, s_global, h_local, d], preserving the original
+    head/sequence dim ordering.
+
+    Args:
+        x: A 4D tensor with layout [b, *, *, d] where '*' are sequence and heads
+        head_dim: Which dimension index corresponds to heads (1 or 2)
+
+    Returns:
+        Tensor with the same dim order as input, with heads sharded and sequence gathered.
+    """
+    world_size = get_ulysses_parallel_world_size()
+    if world_size <= 1:
+        return x
+
+    assert x.ndim == 4, f"x must have 4 dimensions, got {x.ndim}"
+    assert head_dim in (1, 2), f"head_dim must be 1 or 2, got {head_dim}"
+
+    # Move the dimension to be split (h_global) to dim 0 for all_to_all_single
+    if head_dim == 1:
+        b, h_global, s_local, d = x.shape
+        # Shape transition: [b, h_global, s_local, d] -> [h_global, b, s_local, d]
+        permute_order = (1, 0, 2, 3)
+    else:  # head_dim == 2
+        b, s_local, h_global, d = x.shape
+        # Shape transition: [b, s_local, h_global, d] -> [h_global, b, s_local, d]
+        permute_order = (2, 0, 1, 3)
+
+    assert (
+        h_global % world_size == 0
+    ), f"h_global ({h_global}) must be divisible by world_size ({world_size})"
+
+    h_local, s_global = h_global // world_size, s_local * world_size
+
+    x = x.permute(permute_order).contiguous()
+    x = _usp_all_to_all_single(x)
+    x = x.reshape(world_size, h_local, b, s_local, d)
+
+    # Reorder dims to place 'world_size' adjacent to 's_local' to merge them into 's_global'
+    if head_dim == 1:
+        # Shape transition: [world_size, h_local, b, s_local, d] -> [b, h_local, world_size, s_local, d]
+        x = x.permute(2, 1, 0, 3, 4).contiguous().reshape(b, h_local, s_global, d)
+    else:  # head_dim == 2
+        # Shape transition: [world_size, h_local, b, s_local, d] -> [b, world_size, s_local, h_local, d]
+        x = x.permute(2, 0, 3, 1, 4).contiguous().reshape(b, s_global, h_local, d)
+
+    return x
+
+
+def _usp_input_all_to_all_varlen(
+    x: torch.Tensor, seq_lens: list[int], head_dim: int = 1
+) -> torch.Tensor:
+    """
+    Perform Ulysses-style input all-to-all over the head dimension with variable
+    local sequence lengths.
+
+    Default layout expects heads at dim=1 and sequence at dim=2:
+        [b, h, s_local, d] -> [b, h_local, s_global, d]
+
+    If heads are at dim=2 (input is [b, s_local, h, d]), set head_dim=2, and the
+    function returns [b, s_global, h_local, d], preserving the original
+    head/sequence dim ordering.
+
+    Args:
+        x: A 4D tensor with layout [b, *, *, d] where '*' are sequence and heads
+        seq_lens: Local sequence lengths for each rank in the Ulysses group
+        head_dim: Which dimension index corresponds to heads (1 or 2)
+
+    Returns:
+        Tensor with the same dim order as input, with heads sharded and sequence gathered.
+    """
+    world_size = get_ulysses_parallel_world_size()
+    if world_size <= 1:
+        return x
+
+    assert x.ndim == 4, f"x must have 4 dimensions, got {x.ndim}"
+    assert head_dim in (1, 2), f"head_dim must be 1 or 2, got {head_dim}"
+    assert (
+        len(seq_lens) == world_size
+    ), f"seq_lens must have length {world_size}, got {len(seq_lens)}"
+
+    rank = get_ulysses_parallel_rank()
+
+    # Move the dimension to be split (h_global) to dim 0 for all_to_all_single
+    if head_dim == 1:
+        b, h_global, s_local, d = x.shape
+        # Shape transition: [b, h_global, s_local, d] -> [h_global, b, s_local, d]
+        permute_order = (1, 0, 2, 3)
+    else:  # head_dim == 2
+        b, s_local, h_global, d = x.shape
+        # Shape transition: [b, s_local, h_global, d] -> [h_global, b, s_local, d]
+        permute_order = (2, 0, 1, 3)
+
+    assert (
+        s_local == seq_lens[rank]
+    ), f"s_local ({s_local}) must equal seq_lens[{rank}] ({seq_lens[rank]})"
+    assert (
+        h_global % world_size == 0
+    ), f"h_global ({h_global}) must be divisible by world_size ({world_size})"
+
+    h_local = h_global // world_size
+
+    x = x.permute(permute_order).contiguous()
+    x = x.reshape(world_size, h_local, b, s_local, d)
+    input_split_sizes = [h_local * b * s_local * d] * world_size
+    output_split_sizes = [h_local * b * seq_len * d for seq_len in seq_lens]
+    x = _usp_all_to_all_single_varlen(x, output_split_sizes, input_split_sizes)
+
+    chunks = []
+    offset = 0
+    for seq_len, split_size in zip(seq_lens, output_split_sizes):
+        chunk = x[offset : offset + split_size].reshape(h_local, b, seq_len, d)
+        chunks.append(chunk)
+        offset += split_size
+    x = torch.cat(chunks, dim=2)
+
+    if head_dim == 1:
+        # Shape transition: [h_local, b, s_global, d] -> [b, h_local, s_global, d]
+        x = x.permute(1, 0, 2, 3).contiguous()
+    else:  # head_dim == 2
+        # Shape transition: [h_local, b, s_global, d] -> [b, s_global, h_local, d]
+        x = x.permute(1, 2, 0, 3).contiguous()
+
+    return x
+
+
+def _usp_output_all_to_all(x: torch.Tensor, head_dim: int = 1) -> torch.Tensor:
+    """
+    Perform Ulysses-style output all-to-all over the head dimension (inverse of input).
+
+    Default layout expects heads at dim=1 and sequence at dim=2:
+        [b, h_local, s, d] -> [b, h, s_local, d]
+
+    If heads are at dim=2 (input is [b, s_global, h // world_size, d]), set head_dim=2,
+    and the function returns [b, s_local, h, d], preserving the original head/sequence
+    dim ordering.
+
+    Args:
+        x: A 4D tensor with layout [b, *, *, d] where '*' are sequence and heads
+        head_dim: Which dimension index corresponds to heads (1 or 2)
+
+    Returns:
+        Tensor with the same dim order as input, with heads gathered and sequence sharded.
+    """
+    world_size = get_ulysses_parallel_world_size()
+    if world_size <= 1:
+        return x
+
+    assert x.ndim == 4, f"x must have 4 dimensions, got {x.ndim}"
+    assert head_dim in (1, 2), f"head_dim must be 1 or 2, got {head_dim}"
+
+    # Move the dimension to be split (s_global) to dim 0 for all_to_all_single
+    if head_dim == 1:
+        b, h_local, s_global, d = x.shape
+        # Shape transition: [b, h_local, s_global, d] -> [s_global, b, h_local, d]
+        permute_order = (2, 0, 1, 3)
+    else:  # head_dim == 2
+        b, s_global, h_local, d = x.shape
+        # Shape transition: [b, s_global, h_local, d] -> [s_global, b, h_local, d]
+        permute_order = (1, 0, 2, 3)
+
+    assert (
+        s_global % world_size == 0
+    ), f"s_global ({s_global}) must be divisible by world_size ({world_size})"
+
+    s_local, h_global = s_global // world_size, h_local * world_size
+
+    x = x.permute(permute_order).contiguous()
+    x = _usp_all_to_all_single(x)
+    x = x.reshape(world_size, s_local, b, h_local, d)
+
+    # Reorder dims to place 'world_size' adjacent to 'h_local' to merge them into 'h_global'
+    if head_dim == 1:
+        # Shape transition: [world_size, s_local, b, h_local, d] -> [b, world_size, h_local, s_local, d]
+        x = x.permute(2, 0, 3, 1, 4).contiguous().reshape(b, h_global, s_local, d)
+    else:  # head_dim == 2
+        # Shape transition: [world_size, s_local, b, h_local, d] -> [b, s_local, world_size, h_local, d]
+        x = x.permute(2, 1, 0, 3, 4).contiguous().reshape(b, s_local, h_global, d)
+
+    return x
+
+
+def _usp_output_all_to_all_varlen(
+    x: torch.Tensor, seq_lens: list[int], head_dim: int = 1
+) -> torch.Tensor:
+    """
+    Perform Ulysses-style output all-to-all over the head dimension (inverse of input)
+    with variable local sequence lengths.
+
+    Default layout expects heads at dim=1 and sequence at dim=2:
+        [b, h_local, s, d] -> [b, h, s_local, d]
+
+    If heads are at dim=2 (input is [b, s_global, h // world_size, d]), set head_dim=2,
+    and the function returns [b, s_local, h, d], preserving the original head/sequence
+    dim ordering.
+
+    Args:
+        x: A 4D tensor with layout [b, *, *, d] where '*' are sequence and heads
+        seq_lens: Local sequence lengths for each rank in the Ulysses group
+        head_dim: Which dimension index corresponds to heads (1 or 2)
+
+    Returns:
+        Tensor with the same dim order as input, with heads gathered and sequence sharded.
+    """
+    world_size = get_ulysses_parallel_world_size()
+    if world_size <= 1:
+        return x
+
+    assert x.ndim == 4, f"x must have 4 dimensions, got {x.ndim}"
+    assert head_dim in (1, 2), f"head_dim must be 1 or 2, got {head_dim}"
+    assert (
+        len(seq_lens) == world_size
+    ), f"seq_lens must have length {world_size}, got {len(seq_lens)}"
+
+    rank = get_ulysses_parallel_rank()
+
+    # Move the sequence dimension to dim 2 for splitting across seq_lens
+    if head_dim == 1:
+        b, h_local, s_global, d = x.shape
+        # Shape transition: [b, h_local, s_global, d] -> [h_local, b, s_global, d]
+        permute_order = (1, 0, 2, 3)
+    else:  # head_dim == 2
+        b, s_global, h_local, d = x.shape
+        # Shape transition: [b, s_global, h_local, d] -> [h_local, b, s_global, d]
+        permute_order = (2, 0, 1, 3)
+
+    assert s_global == sum(
+        seq_lens
+    ), f"s_global ({s_global}) must equal sum(seq_lens) ({sum(seq_lens)})"
+
+    s_local = seq_lens[rank]
+
+    x = x.permute(permute_order).contiguous()
+    input_chunks = []
+    start = 0
+    for seq_len in seq_lens:
+        end = start + seq_len
+        input_chunks.append(x[:, :, start:end, :].contiguous().reshape(-1))
+        start = end
+    x = torch.cat(input_chunks, dim=0)
+    input_split_sizes = [h_local * b * seq_len * d for seq_len in seq_lens]
+    output_split_sizes = [h_local * b * s_local * d] * world_size
+    x = _usp_all_to_all_single_varlen(x, output_split_sizes, input_split_sizes)
+
+    chunks = []
+    offset = 0
+    for split_size in output_split_sizes:
+        chunk = x[offset : offset + split_size].reshape(h_local, b, s_local, d)
+        chunks.append(chunk)
+        offset += split_size
+    x = torch.cat(chunks, dim=0)
+
+    if head_dim == 1:
+        # Shape transition: [h_global, b, s_local, d] -> [b, h_global, s_local, d]
+        x = x.permute(1, 0, 2, 3).contiguous()
+    else:  # head_dim == 2
+        # Shape transition: [h_global, b, s_local, d] -> [b, s_local, h_global, d]
+        x = x.permute(1, 2, 0, 3).contiguous()
+
+    return x
+
+
+def ring_attn(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attn_impl: "AttentionImpl",
+    is_causal: bool = False,
+    dropout_p: float = 0.0,
+):
+    """
+    Ring Attention implementation.
+
+    This function implements Ring Attention, a strategy for distributed attention
+    computation that reduces peak memory usage. It accepts a generic attention
+    implementation (`attn_impl`) which is called by the underlying PyTorch
+    distributed attention primitive.
+
+    Args:
+        query, key, value: The input tensors for attention.
+        attn_impl: An instance of an attention implementation backend
+                   (e.g., FlashAttentionImpl) whose `forward` method will be
+                   used as the computational kernel.
+        is_causal: Whether to apply causal masking.
+        dropout_p: Dropout probability.
+    """
+    # torch.distributed.tensor.experimental._attention is not a public API,
+    from torch.distributed.tensor.experimental._attention import (
+        _templated_ring_attention,
+    )
+
+    ring_pg = get_sp_group().ring_group
+    assert ring_pg is not None, "Ring process group is not initialized."
+
+    # Ring attention primitives expect tensors in [B, H, S, D] layout.
+    # We permute the inputs here.
+    query = torch.permute(query, [0, 2, 1, 3]).contiguous()
+    key = torch.permute(key, [0, 2, 1, 3]).contiguous()
+    value = torch.permute(value, [0, 2, 1, 3]).contiguous()
+
+    # Create an adapter function that matches the signature expected by
+    # _templated_ring_attention. The `attn_impl` already has dropout and
+    # causal settings configured during its initialization.
+
+    # Note: Please be aware that Attention Backend and Ring Attention may require different QKV tensor shapes.
+    # For example, FlashAttention expects the format to be BSHD.
+    def attn_callable_adapter(q, k, v, *args, **kwargs):
+        # We ignore the dropout_p and is_causal passed by _templated_ring_attention
+        # and rely on the pre-configured attn_impl.
+        # The `attn_metadata` is not available here, so we pass None.
+        # This is a limitation we must accept when using this experimental API.
+        q = torch.permute(q, [0, 2, 1, 3])
+        k = torch.permute(k, [0, 2, 1, 3])
+        v = torch.permute(v, [0, 2, 1, 3])
+        # logger.warning(f"Warning: return_softmax_lse is only supported for FlashAttentionImpl")
+        output, softmax_lse, *rest = attn_impl.forward(
+            q,
+            k,
+            v,
+            attn_metadata=None,
+            return_softmax_lse=True,
+        )
+        output = torch.permute(output, [0, 2, 1, 3])
+        return output, softmax_lse, *rest
+
+    # Starting from torch 2.6.0, _templated_ring_attention expects an integer
+    # segment_id for the attention function.
+    use_segment_id = torch_release >= (2, 6)
+
+    attn_kwargs = dict(
+        op=attn_callable_adapter,
+        dropout_p=dropout_p,
+        is_causal=is_causal,
+        query=query,
+        key=key,
+        value=value,
+        group=ring_pg,  # https://github.com/pytorch/pytorch/blob/c907c778f42ba2fdaf25b733dd25baf9779c6a12/torch/distributed/tensor/experimental/_context_parallel/_attention.py#L309
+    )
+
+    if use_segment_id:
+        # For torch >= 2.6, segment_id is required. The value '1' is a placeholder
+        # as we are not using complex segmentation features.
+        out, *_ = _templated_ring_attention(
+            seq_dim=1,  # segment_id
+            **attn_kwargs,
+        )
+    else:
+        out, *_ = _templated_ring_attention(
+            **attn_kwargs,
+        )
+
+    # Permute the output back to [B, S, H, D] layout.
+    output = torch.permute(out, [0, 2, 1, 3])
+    return output

python/sglang/multimodal_gen/runtime/layers/utils.py

@@ -0,0 +1,267 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/utils.py
+"""Utility methods for model layers."""
+
+import inspect
+from typing import Any, Callable, List, Optional
+
+import torch
+from torch.library import Library
+
+from sglang.kernel_api_logging import debug_torch_op
+from sglang.multimodal_gen.runtime.platforms import current_platform
+
+
+def get_group_size(group) -> int:
+    if hasattr(group, "world_size"):
+        return group.world_size  # GroupCoordinator
+    elif hasattr(group, "size") and callable(getattr(group, "size", None)):
+        return group.size()  # ProcessGroup
+    else:
+        raise ValueError(f"Unsupported group type: {type(group)}")
+
+
+def get_group_rank(group) -> int:
+    if hasattr(group, "rank_in_group"):
+        return group.rank_in_group  # GroupCoordinator
+    elif hasattr(group, "rank") and callable(getattr(group, "rank", None)):
+        return group.rank()  # ProcessGroup
+    else:
+        raise ValueError(f"Unsupported group type: {type(group)}")
+
+
+def get_token_bin_counts_and_mask(
+    tokens: torch.Tensor,
+    vocab_size: int,
+    num_seqs: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    # Compute the bin counts for the tokens.
+    # vocab_size + 1 for padding.
+    bin_counts = torch.zeros(
+        (num_seqs, vocab_size + 1), dtype=torch.long, device=tokens.device
+    )
+    bin_counts.scatter_add_(1, tokens, torch.ones_like(tokens))
+    bin_counts = bin_counts[:, :vocab_size]
+    mask = bin_counts > 0
+
+    return bin_counts, mask
+
+
+sglang_lib = Library("sglang", "FRAGMENT")  # noqa
+
+
+def direct_register_custom_op(
+    op_name: str,
+    op_func: Callable,
+    mutates_args: List[str],
+    fake_impl: Optional[Callable] = None,
+    target_lib: Optional[Library] = None,
+):
+    """
+    `torch.library.custom_op` can have significant overhead because it
+    needs to consider complicated dispatching logic. This function
+    directly registers a custom op and dispatches it to the CUDA backend.
+    See https://gist.github.com/youkaichao/ecbea9ec9fc79a45d2adce1784d7a9a5
+    for more details.
+
+    By default, the custom op is registered to the vLLM library. If you
+    want to register it to a different library, you can pass the library
+    object to the `target_lib` argument.
+
+    IMPORTANT: the lifetime of the operator is tied to the lifetime of the
+    library object. If you want to bind the operator to a different library,
+    make sure the library object is alive when the operator is used.
+
+    Note: This function will silently skip registration if the operator
+    with the same name is already registered to avoid RuntimeError in
+    multi-engine scenarios (e.g., VERL framework).
+    """
+    import torch.library
+
+    my_lib = target_lib or sglang_lib
+
+    # Check if operator is already registered to avoid duplicate registration
+    # This is important for scenarios where multiple SGLang engines run in the same process
+    try:
+        # Try to access the operator to see if it's already registered
+        lib_name = my_lib.m.name if hasattr(my_lib.m, "name") else "sglang"
+        if hasattr(torch.ops, lib_name) and hasattr(
+            getattr(torch.ops, lib_name), op_name
+        ):
+            # Operator already exists, skip registration
+            return
+    except (AttributeError, RuntimeError):
+        # Operator doesn't exist, proceed with registration
+        pass
+
+    if hasattr(torch.library, "infer_schema"):
+        schema_str = torch.library.infer_schema(op_func, mutates_args=mutates_args)
+    else:
+        # for pytorch 2.4
+        import torch._custom_op.impl
+
+        schema_str = torch._custom_op.impl.infer_schema(op_func, mutates_args)
+
+    try:
+        my_lib.define(op_name + schema_str)
+        my_lib.impl(
+            op_name, op_func, "CUDA" if not current_platform.is_npu() else "PrivateUse1"
+        )
+        if fake_impl is not None:
+            my_lib._register_fake(op_name, fake_impl)
+    except RuntimeError as error:
+        if "Tried to register an operator" in str(error) and "multiple times" in str(
+            error
+        ):
+            # Silently ignore duplicate registration errors
+            # This can happen in multi-engine scenarios
+            pass
+        else:
+            # Re-raise other RuntimeErrors
+            raise error
+    except AttributeError as error:
+        # Always re-raise AttributeError as it indicates missing dependencies
+        raise error
+
+
+class CustomOpWrapper:
+    def __init__(
+        self,
+        op_name: str,
+        op_func: Callable,
+        mutates_args: List[str],
+        **extra_kwargs,
+    ):
+        self.op_name = op_name
+        self.op_func = op_func
+        self.mutates_args = mutates_args
+        self.extra_kwargs = extra_kwargs
+        self._impl: Optional[Callable] = None
+
+    def __call__(self, *args, **kwargs):
+        return self.real_impl(*args, **kwargs)
+
+    @property
+    def real_impl(self) -> Callable:
+        if self._impl is None:
+            if not hasattr(torch.ops.sglang, self.op_name):
+
+                # NOTE(dark): if torch compile fail here, mark the decorator as eager
+                # lazy registration does not work with torch compile
+                direct_register_custom_op(
+                    op_name=self.op_name,
+                    op_func=self.op_func,
+                    mutates_args=self.mutates_args,
+                    fake_impl=self.fake_impl,
+                )
+            self._impl = debug_torch_op(self.op_func, self.op_name)
+            assert self._impl is not None
+        return self._impl
+
+    @property
+    def fake_impl(self) -> Callable:
+        if "fake_impl" in self.extra_kwargs:
+            return self.extra_kwargs["fake_impl"]
+        assert "out_shape" in self.extra_kwargs
+        signature = inspect.signature(self.op_func)
+        out_shape = self.extra_kwargs["out_shape"]
+
+        # check out_shape in signature
+
+        def fake_impl(*args, **kwargs):
+            if out_shape is None:
+                return None
+            bound = signature.bind(*args, **kwargs)
+            bound.apply_defaults()
+            try:
+                return torch.empty_like(
+                    bound.args[out_shape]
+                    if isinstance(out_shape, int)
+                    else bound.arguments[out_shape]
+                )
+            except (IndexError, KeyError):
+                raise RuntimeError(
+                    f"Cannot find output argument at position `{out_shape}` for "
+                    f"custom operator `{self.op_name}` with signature `{signature}`."
+                )
+
+        return fake_impl
+
+
+# Real implementation
+def register_custom_op(
+    fn: Optional[Callable] = None,
+    *,
+    op_name: Optional[str] = None,
+    mutates_args: Optional[List[str]] = None,
+    eager: bool = True,
+    **extra_kwargs,
+) -> Any:
+    """
+    A decorator to register a custom operator.
+
+    Example usage:
+    ```python
+    # inplace operator, out_shape is None by default
+    @register_custom_op(mutates_args=["x"])
+    def add_1_(x: torch.Tensor) -> None:
+        x.add_(1)
+
+    # operator with output, out_shape indicates the position of output
+    @register_custom_op(mutates_args=["x"], out_shape=0)
+    def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        return x.add_(y)
+    ```
+
+    :param fn: The function to be registered as a custom operator.
+               If None, return a decorator.
+    :type fn: Callable
+    :param op_name: The name of the operator. If None, use the function name
+    :type op_name: Optional[str]
+    :param mutates_args: A list of argument names that are mutated in-place.
+    :type mutates_args: List[str]
+    :param out_shape: The position (int for positional, str for keyword) of the output-shape tensor.
+                      It is used to generate a fake implementation for torch.compile compatibility.
+                      If the operator is inplace and has no output, set to None.
+    :type out_shape: Optional[List[Union[int, str]]]
+    :param fake_impl: A fake implementation for the operator.
+                      Only one of `out_shape` or `fake_impl` should be provided.
+    :type fake_impl: Optional[Callable]
+    :param eager: Whether to register the operator eagerly.
+                  If False, the registration will be deferred until the first call.
+                  If you met any issue with torch.compile, try to set eager=True.
+                  Currently, to avoid misuse, we set eager=True by default.
+    :type eager: bool
+    :return: The registered JIT custom operator, or a decorator.
+             NOTE: the real register will occur at the first call of the function.
+    :rtype: Callable
+    """
+    extra_kwarg_keys = set(extra_kwargs.keys())
+    expected_kwarg_keys = set({"out_shape", "fake_impl"})
+    assert (
+        expected_kwarg_keys >= extra_kwarg_keys
+    ), f"Unexpected extra kwargs: {extra_kwarg_keys - expected_kwarg_keys}"
+
+    has_out_shape = "out_shape" in extra_kwargs
+    has_fake_impl = "fake_impl" in extra_kwargs
+    assert not (
+        has_out_shape and has_fake_impl
+    ), "Only one of `out_shape` or `fake_impl` should be provided."
+    # Assume inplace if neither out_shape nor fake_impl is provided
+    if not (has_out_shape or has_fake_impl):
+        extra_kwargs["out_shape"] = None
+
+    def decorator(op_func: Callable) -> Callable:
+        wrapper = CustomOpWrapper(
+            op_name=op_name or op_func.__name__,
+            op_func=op_func,
+            mutates_args=mutates_args or [],
+            **extra_kwargs,
+        )
+        return wrapper.real_impl if eager else wrapper
+
+    if fn is not None:
+        return decorator(fn)
+    return decorator

python/sglang/multimodal_gen/runtime/layers/visual_embedding.py

@@ -0,0 +1,353 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.models.embeddings import (
+    CombinedTimestepGuidanceTextProjEmbeddings as _CombinedTimestepGuidanceTextProjEmbeddings,
+)
+from diffusers.models.embeddings import (
+    CombinedTimestepTextProjEmbeddings as _CombinedTimestepTextProjEmbeddings,
+)
+from diffusers.models.embeddings import (
+    PixArtAlphaTextProjection,
+    TimestepEmbedding,
+)
+from diffusers.models.embeddings import Timesteps as _Timesteps
+from diffusers.models.embeddings import (
+    get_timestep_embedding as timestep_embedding_diffusers,
+)
+
+from sglang.jit_kernel.timestep_embedding import (
+    timestep_embedding as timestep_embedding_cuda,
+)
+from sglang.multimodal_gen.runtime.layers.activation import get_act_fn
+from sglang.multimodal_gen.runtime.layers.linear import ColumnParallelLinear
+from sglang.multimodal_gen.runtime.layers.mlp import MLP
+from sglang.multimodal_gen.runtime.platforms import current_platform
+
+_is_cuda = current_platform.is_cuda()
+
+
+class PatchEmbed(nn.Module):
+    """2D Image to Patch Embedding
+
+    Image to Patch Embedding using Conv2d
+
+    A convolution based approach to patchifying a 2D image w/ embedding projection.
+
+    Based on the impl in https://github.com/google-research/vision_transformer
+
+    Hacked together by / Copyright 2020 Ross Wightman
+
+    Remove the _assert function in forward function to be compatible with multi-resolution images.
+    """
+
+    def __init__(
+        self,
+        patch_size=16,
+        in_chans=3,
+        embed_dim=768,
+        norm_layer=None,
+        flatten=True,
+        bias=True,
+        dtype=None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        if isinstance(patch_size, list | tuple):
+            if len(patch_size) == 1:
+                patch_size = (1, patch_size[0], patch_size[0])
+            elif len(patch_size) == 2:
+                patch_size = (1, patch_size[0], patch_size[1])
+        else:
+            patch_size = (1, patch_size, patch_size)
+
+        self.patch_size = patch_size
+        self.flatten = flatten
+
+        self.proj = nn.Conv3d(
+            in_chans,
+            embed_dim,
+            kernel_size=patch_size,
+            stride=patch_size,
+            bias=bias,
+            dtype=dtype,
+        )
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x):
+        if x.dim() == 5:
+            B, C, T, H, W = x.shape
+            pt, ph, pw = self.patch_size
+
+            if T % pt == 0 and H % ph == 0 and W % pw == 0:
+                T_ = T // pt
+                H_ = H // ph
+                W_ = W // pw
+
+                x = x.reshape(B, C, T_, pt, H_, ph, W_, pw)
+                x = x.permute(0, 2, 4, 6, 1, 3, 5, 7).contiguous()
+                x = x.reshape(B, T_ * H_ * W_, C * pt * ph * pw)
+
+                w = self.proj.weight.reshape(self.proj.weight.shape[0], -1)
+                x = F.linear(x, w, self.proj.bias)  # [B, T'*H'*W', embed_dim]
+
+                if not self.flatten:
+                    x = x.reshape(B, T_, H_, W_, -1).permute(0, 4, 1, 2, 3).contiguous()
+
+                x = self.norm(x)
+                return x
+
+        # Fallback to Conv3d for non-5D input or indivisible spatial dims.
+        x = self.proj(x)
+        if self.flatten:
+            x = x.flatten(2).transpose(1, 2)
+        x = self.norm(x)
+        return x
+
+
+class WanCamControlPatchEmbedding(nn.Module):
+    """Patch embedding used by LingBotWorld camera/plucker controls."""
+
+    def __init__(
+        self,
+        patch_size=(1, 2, 2),
+        in_chans=384,
+        embed_dim=2048,
+        bias=True,
+        dtype=None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        del prefix
+        if isinstance(patch_size, list | tuple):
+            if len(patch_size) != 3:
+                raise ValueError(
+                    f"patch_size must have length 3, got {len(patch_size)}"
+                )
+            patch_size = tuple(patch_size)
+        else:
+            raise ValueError(f"Unsupported patch_size type: {type(patch_size)}")
+
+        self.patch_size = patch_size
+        pt, ph, pw = self.patch_size
+        self.in_features = in_chans * pt * ph * pw
+        self.proj = nn.Linear(self.in_features, embed_dim, bias=bias, dtype=dtype)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if x.dim() != 5:
+            raise ValueError(
+                f"Expected camera embedding shape [B, C, F, H, W], got {tuple(x.shape)}"
+            )
+
+        bsz, channels, frames, height, width = x.shape
+        pt, ph, pw = self.patch_size
+        if (frames % pt) != 0 or (height % ph) != 0 or (width % pw) != 0:
+            raise ValueError(
+                f"Input shape {tuple(x.shape)} must be divisible by patch_size {self.patch_size}"
+            )
+
+        x = x.view(
+            bsz,
+            channels,
+            frames // pt,
+            pt,
+            height // ph,
+            ph,
+            width // pw,
+            pw,
+        )
+        x = x.permute(0, 2, 4, 6, 1, 3, 5, 7).reshape(bsz, -1, self.in_features)
+        return self.proj(x)
+
+
+class Timesteps(_Timesteps):
+    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+        if _is_cuda:
+            return timestep_embedding_cuda(
+                timesteps,
+                self.num_channels,
+                flip_sin_to_cos=self.flip_sin_to_cos,
+                downscale_freq_shift=self.downscale_freq_shift,
+                scale=self.scale,
+            )
+        else:
+            return timestep_embedding_diffusers(
+                timesteps,
+                self.num_channels,
+                flip_sin_to_cos=self.flip_sin_to_cos,
+                downscale_freq_shift=self.downscale_freq_shift,
+                scale=self.scale,
+            )
+
+
+class CombinedTimestepGuidanceTextProjEmbeddings(
+    _CombinedTimestepGuidanceTextProjEmbeddings
+):
+    def __init__(self, embedding_dim, pooled_projection_dim):
+        nn.Module.__init__(self)
+
+        # use sgld op
+        self.time_proj = Timesteps(
+            num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0
+        )
+        # use diffusers op
+        self.timestep_embedder = TimestepEmbedding(
+            in_channels=256, time_embed_dim=embedding_dim
+        )
+        self.guidance_embedder = TimestepEmbedding(
+            in_channels=256, time_embed_dim=embedding_dim
+        )
+        self.text_embedder = PixArtAlphaTextProjection(
+            pooled_projection_dim, embedding_dim, act_fn="silu"
+        )
+
+
+class CombinedTimestepTextProjEmbeddings(_CombinedTimestepTextProjEmbeddings):
+    def __init__(self, embedding_dim, pooled_projection_dim):
+        nn.Module.__init__(self)
+
+        # use sgld op
+        self.time_proj = Timesteps(
+            num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0
+        )
+        # use diffusers op
+        self.timestep_embedder = TimestepEmbedding(
+            in_channels=256, time_embed_dim=embedding_dim
+        )
+        self.text_embedder = PixArtAlphaTextProjection(
+            pooled_projection_dim, embedding_dim, act_fn="silu"
+        )
+
+
+class TimestepEmbedder(nn.Module):
+    """
+    Embeds scalar timesteps into vector representations.
+    """
+
+    def __init__(
+        self,
+        hidden_size,
+        act_layer="silu",
+        frequency_embedding_size=256,
+        max_period=10000,
+        dtype=None,
+        freq_dtype=torch.float32,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.frequency_embedding_size = frequency_embedding_size
+        self.max_period = max_period
+
+        self.mlp = MLP(
+            frequency_embedding_size,
+            hidden_size,
+            hidden_size,
+            act_type=act_layer,
+            dtype=dtype,
+        )
+        self.freq_dtype = freq_dtype
+
+    def forward(
+        self, t: torch.Tensor, timestep_seq_len: int | None = None
+    ) -> torch.Tensor:
+        t_freq = timestep_embedding(
+            t, self.frequency_embedding_size, self.max_period, dtype=self.freq_dtype
+        ).to(self.mlp.fc_in.weight.dtype)
+        if timestep_seq_len is not None:
+            assert (
+                t_freq.shape[0] % timestep_seq_len == 0
+            ), "timestep length is not divisible by timestep_seq_len"
+            batch_size = t_freq.shape[0] // timestep_seq_len
+            t_freq = t_freq.unflatten(0, (batch_size, timestep_seq_len))
+        # t_freq = t_freq.to(self.mlp.fc_in.weight.dtype)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+def timestep_embedding(
+    t: torch.Tensor,
+    dim: int,
+    max_period: int = 10000,
+    dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    """
+    Create sinusoidal timestep embeddings.
+
+    Args:
+        t: Tensor of shape [B] with timesteps
+        dim: Embedding dimension
+        max_period: Controls the minimum frequency of the embeddings
+
+    Returns:
+        Tensor of shape [B, dim] with embeddings
+    """
+    half = dim // 2
+    freqs = torch.exp(
+        -math.log(max_period)
+        * torch.arange(start=0, end=half, dtype=dtype, device=t.device)
+        / half
+    )
+    args = t[:, None].float() * freqs[None]
+    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    return embedding
+
+
+class ModulateProjection(nn.Module):
+    """Modulation layer for DiT blocks."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        factor: int = 2,
+        act_layer: str = "silu",
+        dtype: torch.dtype | None = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.factor = factor
+        self.hidden_size = hidden_size
+        self.linear = ColumnParallelLinear(
+            hidden_size,
+            hidden_size * factor,
+            bias=True,
+            gather_output=True,
+            params_dtype=dtype,
+        )
+        self.act = get_act_fn(act_layer)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.act(x)
+        x, _ = self.linear(x)
+        return x
+
+
+def unpatchify(x, t, h, w, patch_size, channels) -> torch.Tensor:
+    """
+    Convert patched representation back to image space.
+
+    Args:
+        x: Tensor of shape [B, T*H*W, C*P_t*P_h*P_w]
+        t, h, w: Temporal and spatial dimensions
+
+    Returns:
+        Unpatchified tensor of shape [B, C, T*P_t, H*P_h, W*P_w]
+    """
+    assert x.ndim == 3, f"x.ndim: {x.ndim}"
+    assert len(patch_size) == 3, f"patch_size: {patch_size}"
+    assert t * h * w == x.shape[1], f"t * h * w: {t * h * w}, x.shape[1]: {x.shape[1]}"
+    c = channels
+    pt, ph, pw = patch_size
+
+    x = x.reshape(shape=(x.shape[0], t, h, w, c, pt, ph, pw))
+    x = torch.einsum("nthwcopq->nctohpwq", x)
+    imgs = x.reshape(shape=(x.shape[0], c, t * pt, h * ph, w * pw))
+
+    return imgs

python/sglang/multimodal_gen/runtime/layers/vocab_parallel_embedding.py

@@ -0,0 +1,490 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from torch.nn.parameter import Parameter, UninitializedParameter
+
+from sglang.multimodal_gen.runtime.distributed import (
+    divide,
+    get_tp_group,
+    tensor_model_parallel_all_reduce,
+)
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+    QuantizeMethodBase,
+    method_has_implemented_embedding,
+)
+from sglang.multimodal_gen.runtime.layers.utils import get_group_rank, get_group_size
+from sglang.multimodal_gen.runtime.models.parameter import BasevLLMParameter
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.multimodal_gen.runtime.utils.weight_attrs import set_weight_attrs
+
+DEFAULT_VOCAB_PADDING_SIZE = 64
+
+
+class UnquantizedEmbeddingMethod(QuantizeMethodBase):
+    """Unquantized method for embeddings."""
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes: list[int],
+        input_size: int,
+        output_size: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        """Create weights for embedding layer."""
+
+        weight = Parameter(
+            torch.empty(
+                sum(output_partition_sizes),
+                input_size_per_partition,
+                dtype=params_dtype,
+            ),
+            requires_grad=False,
+        )
+        set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0})
+        layer.register_parameter("weight", weight)
+        set_weight_attrs(weight, extra_weight_attrs)
+
+    def apply(
+        self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        return F.linear(x, layer.weight, bias)
+
+    def embedding(self, layer: torch.nn.Module, input_: torch.Tensor) -> torch.Tensor:
+        return F.embedding(input_, layer.weight)
+
+
+def pad_vocab_size(vocab_size: int, pad_to: int = DEFAULT_VOCAB_PADDING_SIZE) -> int:
+    """Pad the vocab size to the given value."""
+    return ((vocab_size + pad_to - 1) // pad_to) * pad_to
+
+
+def vocab_range_from_per_partition_vocab_size(
+    per_partition_vocab_size: int, rank: int, offset: int = 0
+) -> Sequence[int]:
+    index_f = rank * per_partition_vocab_size
+    index_l = index_f + per_partition_vocab_size
+    return index_f + offset, index_l + offset
+
+
+def vocab_range_from_global_vocab_size(
+    global_vocab_size: int, rank: int, world_size: int, offset: int = 0
+) -> Sequence[int]:
+    per_partition_vocab_size = divide(global_vocab_size, world_size)
+    return vocab_range_from_per_partition_vocab_size(
+        per_partition_vocab_size, rank, offset=offset
+    )
+
+
+@dataclass
+class VocabParallelEmbeddingShardIndices:
+    """Indices for a shard of a vocab parallel embedding."""
+
+    padded_org_vocab_start_index: int
+    padded_org_vocab_end_index: int
+    padded_added_vocab_start_index: int
+    padded_added_vocab_end_index: int
+
+    org_vocab_start_index: int
+    org_vocab_end_index: int
+    added_vocab_start_index: int
+    added_vocab_end_index: int
+
+    @property
+    def num_org_elements(self) -> int:
+        return self.org_vocab_end_index - self.org_vocab_start_index
+
+    @property
+    def num_added_elements(self) -> int:
+        return self.added_vocab_end_index - self.added_vocab_start_index
+
+    @property
+    def num_org_elements_padded(self) -> int:
+        return self.padded_org_vocab_end_index - self.padded_org_vocab_start_index
+
+    @property
+    def num_added_elements_padded(self) -> int:
+        return self.padded_added_vocab_end_index - self.padded_added_vocab_start_index
+
+    @property
+    def num_org_vocab_padding(self) -> int:
+        return self.num_org_elements_padded - self.num_org_elements
+
+    @property
+    def num_added_vocab_padding(self) -> int:
+        return self.num_added_elements_padded - self.num_added_elements
+
+    @property
+    def num_elements_padded(self) -> int:
+        return self.num_org_elements_padded + self.num_added_elements_padded
+
+    def __post_init__(self):
+        # sanity checks
+        assert self.padded_org_vocab_start_index <= self.padded_org_vocab_end_index
+        assert self.padded_added_vocab_start_index <= self.padded_added_vocab_end_index
+
+        assert self.org_vocab_start_index <= self.org_vocab_end_index
+        assert self.added_vocab_start_index <= self.added_vocab_end_index
+
+        assert self.org_vocab_start_index <= self.padded_org_vocab_start_index
+        assert self.added_vocab_start_index <= self.padded_added_vocab_start_index
+        assert self.org_vocab_end_index <= self.padded_org_vocab_end_index
+        assert self.added_vocab_end_index <= self.padded_added_vocab_end_index
+
+        assert self.num_org_elements <= self.num_org_elements_padded
+        assert self.num_added_elements <= self.num_added_elements_padded
+
+
+@torch.compile(
+    dynamic=True,
+    backend=current_platform.simple_compile_backend,
+    disable=current_platform.is_npu(),
+)
+def get_masked_input_and_mask(
+    input_: torch.Tensor,
+    org_vocab_start_index: int,
+    org_vocab_end_index: int,
+    num_org_vocab_padding: int,
+    added_vocab_start_index: int,
+    added_vocab_end_index: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    # torch.compile will fuse all of the pointwise ops below
+    # into a single kernel, making it very fast
+    org_vocab_mask = (input_ >= org_vocab_start_index) & (input_ < org_vocab_end_index)
+    added_vocab_mask = (input_ >= added_vocab_start_index) & (
+        input_ < added_vocab_end_index
+    )
+    added_offset = (
+        added_vocab_start_index
+        - (org_vocab_end_index - org_vocab_start_index)
+        - num_org_vocab_padding
+    )
+    valid_offset = (org_vocab_start_index * org_vocab_mask) + (
+        added_offset * added_vocab_mask
+    )
+    vocab_mask = org_vocab_mask | added_vocab_mask
+    input_ = vocab_mask * (input_ - valid_offset)
+    return input_, ~vocab_mask
+
+
+class VocabParallelEmbedding(torch.nn.Module):
+    """Embedding parallelized in the vocabulary dimension.
+
+    Adapted from torch.nn.Embedding, note that we pad the vocabulary size to
+    make sure it is divisible by the number of model parallel GPUs.
+
+    In order to support various loading methods, we ensure that LoRA-added
+    embeddings are always at the end of TP-sharded tensors. In other words,
+    we shard base embeddings and LoRA embeddings separately (both padded),
+    and place them in the same tensor.
+    In this example, we will have the original vocab size = 1010,
+    added vocab size = 16 and padding to 64. Therefore, the total
+    vocab size with padding will be 1088 (because we first pad 1010 to
+    1024, add 16, and then pad to 1088).
+    Therefore, the tensor format looks like the following:
+    TP1, rank 0 (no sharding):
+                            |< --------BASE-------- >|< -BASE PADDING-- >|< -----LORA------ >|< -LORA PADDING-- >|
+    corresponding token_id: |  0  |  1  | ... | 1009 |  -1  | ... |  -1  | 1010 | ... | 1015 |  -1  | ... |  -1  |
+                     index: |  0  |  1  | ... | 1009 | 1010 | ... | 1023 | 1024 | ... | 1039 | 1040 | ... | 1087 |
+
+    TP2, rank 0:
+                            |< --------------------BASE--------------------- >|< -----LORA------ >|< -LORA PADDING- >|
+    corresponding token_id: |  0  |  1  |  2  | ... | 497  | 498 | ...  | 511 | 1000 | ... | 1015 |  -1  | ... |  -1 |
+                     index: |  0  |  1  |  2  | ... | 497  | 498 | ...  | 511 | 512  | ... | 527  |  520 | ... | 543 |
+    TP2, rank 1:
+                            |< -----------BASE----------- >|< -BASE PADDING- >|< -----------LORA PADDING----------- >|
+    corresponding token_id: | 512 | 513 | 514 | ... | 1009 | -1  | ...  | -1  |  -1  | ... |  -1  | -1  | ... |   -1 |
+                     index: |  0  |  1  |  2  | ... | 497  | 498 | ...  | 511 | 512  | ... | 519  | 520 | ... |  543 |
+
+    Args:
+        num_embeddings: vocabulary size.
+        embedding_dim: size of hidden state.
+        params_dtype: type of the parameters.
+        org_num_embeddings: original vocabulary size (without LoRA).
+        padding_size: padding size for the vocabulary.
+        quant_config: quant config for the layer
+        prefix: full name of the layer in the state dict
+    """  # noqa: E501
+
+    def __init__(
+        self,
+        num_embeddings: int,
+        embedding_dim: int,
+        params_dtype: torch.dtype | None = None,
+        org_num_embeddings: int | None = None,
+        padding_size: int = DEFAULT_VOCAB_PADDING_SIZE,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+        tp_group: dist.ProcessGroup = None,
+    ):
+        super().__init__()
+
+        # Keep the input dimensions.
+        tp_group = tp_group or get_tp_group()
+        tp_rank = get_group_rank(tp_group)
+        self.tp_size = get_group_size(tp_group)
+        self.tp_group = tp_group
+        self.num_embeddings = num_embeddings
+        self.padding_size = padding_size
+        self.org_vocab_size = org_num_embeddings or num_embeddings
+        num_added_embeddings = num_embeddings - self.org_vocab_size
+        self.org_vocab_size_padded = pad_vocab_size(
+            self.org_vocab_size, self.padding_size
+        )
+        self.num_embeddings_padded = pad_vocab_size(
+            self.org_vocab_size_padded + num_added_embeddings, self.padding_size
+        )
+        assert self.org_vocab_size_padded <= self.num_embeddings_padded
+
+        self.shard_indices = self._get_indices(
+            self.num_embeddings_padded,
+            self.org_vocab_size_padded,
+            self.num_embeddings,
+            self.org_vocab_size,
+            tp_rank,
+            self.tp_size,
+        )
+        self.embedding_dim = embedding_dim
+
+        quant_method = None
+        if quant_config is not None:
+            quant_method = quant_config.get_quant_method(self, prefix=prefix)
+        if quant_method is None:
+            quant_method = UnquantizedEmbeddingMethod()
+
+        # If we are making an embedding layer, then our quantization linear
+        # method must implement the embedding operation. If we are another
+        # layer type like ParallelLMHead, this is not important.
+        is_embedding_layer = type(self.__class__) is VocabParallelEmbedding
+        quant_method_implements_embedding = method_has_implemented_embedding(
+            type(quant_method)
+        )
+        if is_embedding_layer and not quant_method_implements_embedding:
+            raise NotImplementedError(
+                f"The class {type(quant_method).__name__} must implement "
+                "the 'embedding' method, see UnquantizedEmbeddingMethod."
+            )
+
+        self.quant_method: QuantizeMethodBase = quant_method
+
+        if params_dtype is None:
+            params_dtype = torch.get_default_dtype()
+        # Divide the weight matrix along the vocaburaly dimension.
+        self.num_added_embeddings = self.num_embeddings - self.org_vocab_size
+        self.num_embeddings_per_partition = divide(
+            self.num_embeddings_padded, self.tp_size
+        )
+        assert (
+            self.shard_indices.num_elements_padded == self.num_embeddings_per_partition
+        )
+        self.num_org_embeddings_per_partition = (
+            self.shard_indices.org_vocab_end_index
+            - self.shard_indices.org_vocab_start_index
+        )
+        self.num_added_embeddings_per_partition = (
+            self.shard_indices.added_vocab_end_index
+            - self.shard_indices.added_vocab_start_index
+        )
+
+        self.quant_method.create_weights(
+            self,
+            self.embedding_dim,
+            [self.num_embeddings_per_partition],
+            self.embedding_dim,
+            self.num_embeddings_padded,
+            params_dtype=params_dtype,
+            weight_loader=self.weight_loader,
+        )
+
+    @classmethod
+    def _get_indices(
+        cls,
+        vocab_size_padded: int,
+        org_vocab_size_padded: int,
+        vocab_size: int,
+        org_vocab_size: int,
+        tp_rank: int,
+        tp_size: int,
+    ) -> VocabParallelEmbeddingShardIndices:
+        """Get start and end indices for vocab parallel embedding, following the
+        layout outlined in the class docstring, based on the given tp_rank and
+        tp_size."""
+        num_added_embeddings_padded = vocab_size_padded - org_vocab_size_padded
+        padded_org_vocab_start_index, padded_org_vocab_end_index = (
+            vocab_range_from_global_vocab_size(org_vocab_size_padded, tp_rank, tp_size)
+        )
+        padded_added_vocab_start_index, padded_added_vocab_end_index = (
+            vocab_range_from_global_vocab_size(
+                num_added_embeddings_padded, tp_rank, tp_size, offset=org_vocab_size
+            )
+        )
+        # remove padding
+        org_vocab_start_index = min(padded_org_vocab_start_index, org_vocab_size)
+        org_vocab_end_index = min(padded_org_vocab_end_index, org_vocab_size)
+        added_vocab_start_index = min(padded_added_vocab_start_index, vocab_size)
+        added_vocab_end_index = min(padded_added_vocab_end_index, vocab_size)
+        return VocabParallelEmbeddingShardIndices(
+            padded_org_vocab_start_index,
+            padded_org_vocab_end_index,
+            padded_added_vocab_start_index,
+            padded_added_vocab_end_index,
+            org_vocab_start_index,
+            org_vocab_end_index,
+            added_vocab_start_index,
+            added_vocab_end_index,
+        )
+
+    def get_sharded_to_full_mapping(self) -> list[int] | None:
+        """Get a mapping that can be used to reindex the gathered
+        logits for sampling.
+
+        During sampling, we gather logits from all ranks. The relationship
+        of index->token_id will follow the same format as outlined in the class
+        docstring. However, after the gather, we want to reindex the final
+        logits tensor to map index->token_id one-to-one (the index is always
+        equal the token_id it corresponds to). The indices returned by this
+        method allow us to do that.
+        """
+        if self.tp_size < 2:
+            return None
+
+        base_embeddings: list[int] = []
+        added_embeddings: list[int] = []
+        padding: list[int] = []
+        for tp_rank in range(self.tp_size):
+            shard_indices = self._get_indices(
+                self.num_embeddings_padded,
+                self.org_vocab_size_padded,
+                self.num_embeddings,
+                self.org_vocab_size,
+                tp_rank,
+                self.tp_size,
+            )
+            range_start = self.num_embeddings_per_partition * tp_rank
+            range_end = self.num_embeddings_per_partition * (tp_rank + 1)
+            base_embeddings.extend(
+                range(range_start, range_start + shard_indices.num_org_elements)
+            )
+            padding.extend(
+                range(
+                    range_start + shard_indices.num_org_elements,
+                    range_start + shard_indices.num_org_elements_padded,
+                )
+            )
+            added_embeddings.extend(
+                range(
+                    range_start + shard_indices.num_org_elements_padded,
+                    range_start
+                    + shard_indices.num_org_elements_padded
+                    + shard_indices.num_added_elements,
+                )
+            )
+            padding.extend(
+                range(
+                    range_start
+                    + shard_indices.num_org_elements_padded
+                    + shard_indices.num_added_elements,
+                    range_start
+                    + shard_indices.num_org_elements_padded
+                    + shard_indices.num_added_elements_padded,
+                )
+            )
+            assert (
+                range_start
+                + shard_indices.num_org_elements_padded
+                + shard_indices.num_added_elements_padded
+                == range_end
+            )
+        ret = base_embeddings + added_embeddings + padding
+        assert len(ret) == self.num_embeddings_padded
+        return ret
+
+    def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor):
+        output_dim = getattr(param, "output_dim", None)
+        packed_dim = getattr(param, "packed_dim", None)
+
+        # If the parameter is a gguf weight, then load it directly.
+        if getattr(param, "is_gguf_weight_type", None):
+            param.data.copy_(loaded_weight)
+            param.weight_type = loaded_weight.item()
+            return
+        elif isinstance(param, UninitializedParameter):
+            shape = list(loaded_weight.shape)
+            if output_dim is not None:
+                shape[output_dim] = self.num_embeddings_per_partition
+            param.materialize(tuple(shape), dtype=loaded_weight.dtype)
+
+        # If parameter does not have output dim, then it should
+        # be copied onto all gpus (e.g. g_idx for act_order gptq).
+        if output_dim is None:
+            assert param.data.shape == loaded_weight.shape
+            param.data.copy_(loaded_weight)
+            return
+
+        # Shard indexes for loading the weight
+        start_idx = self.shard_indices.org_vocab_start_index
+        shard_size = self.shard_indices.org_vocab_end_index - start_idx
+
+        # If param packed on the same dim we are sharding on, then
+        # need to adjust offsets of loaded weight by pack_factor.
+        if packed_dim is not None and packed_dim == output_dim:
+            packed_factor = (
+                param.packed_factor
+                if isinstance(param, BasevLLMParameter)
+                else param.pack_factor
+            )
+            assert loaded_weight.shape[output_dim] == (
+                self.org_vocab_size // param.packed_factor
+            )
+            start_idx = start_idx // packed_factor
+            shard_size = shard_size // packed_factor
+        else:
+            assert loaded_weight.shape[output_dim] == self.org_vocab_size
+
+        # Copy the data. Select chunk corresponding to current shard.
+        loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size)
+
+        param[: loaded_weight.shape[0]].data.copy_(loaded_weight)
+        param[loaded_weight.shape[0] :].data.fill_(0)
+
+    def forward(self, input_):
+        if self.tp_size > 1:
+            # Build the mask.
+            masked_input, input_mask = get_masked_input_and_mask(
+                input_,
+                self.shard_indices.org_vocab_start_index,
+                self.shard_indices.org_vocab_end_index,
+                self.shard_indices.num_org_vocab_padding,
+                self.shard_indices.added_vocab_start_index,
+                self.shard_indices.added_vocab_end_index,
+            )
+        else:
+            masked_input = input_
+        # Get the embeddings.
+        output_parallel = self.quant_method.embedding(self, masked_input.long())
+        # Mask the output embedding.
+        if self.tp_size > 1:
+            output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0)
+        # Reduce across all the model parallel GPUs.
+        output = tensor_model_parallel_all_reduce(
+            output_parallel, tp_group=self.tp_group
+        )
+        return output
+
+    def extra_repr(self) -> str:
+        s = f"num_embeddings={self.num_embeddings_per_partition}"
+        s += f", embedding_dim={self.embedding_dim}"
+        s += f", org_vocab_size={self.org_vocab_size}"
+        s += f", num_embeddings_padded={self.num_embeddings_padded}"
+        s += f", tp_size={self.tp_size}"
+        return s