chore: import upstream snapshot with attribution

python/sglang/jit_kernel/diffusion/triton/causal_conv3d_pad.py

@@ -0,0 +1,122 @@
+from __future__ import annotations
+
+import torch
+import triton  # type: ignore
+import triton.language as tl  # type: ignore
+
+
+@triton.jit
+def _fused_cat_pad_5d_kernel(
+    x_ptr,
+    cache_ptr,
+    out_ptr,
+    total,
+    channels,
+    t_size,
+    h_size,
+    w_size,
+    cache_t,
+    out_t,
+    out_h,
+    out_w,
+    pad_d_left,
+    pad_h_top,
+    pad_w_left,
+    block_size: tl.constexpr,
+):
+    offsets = tl.program_id(0) * block_size + tl.arange(0, block_size)
+    mask = offsets < total
+
+    ow = offsets % out_w
+    tmp = offsets // out_w
+    oh = tmp % out_h
+    tmp = tmp // out_h
+    out = tmp % out_t
+    tmp = tmp // out_t
+    oc = tmp % channels
+    ob = tmp // channels
+
+    iw = ow - pad_w_left
+    ih = oh - pad_h_top
+    src_t = out - pad_d_left
+
+    valid = (
+        mask
+        & (iw >= 0)
+        & (iw < w_size)
+        & (ih >= 0)
+        & (ih < h_size)
+        & (src_t >= 0)
+        & (src_t < cache_t + t_size)
+    )
+    from_cache = src_t < cache_t
+
+    x_t = src_t - cache_t
+    clamped_iw = tl.minimum(tl.maximum(iw, 0), w_size - 1)
+    clamped_ih = tl.minimum(tl.maximum(ih, 0), h_size - 1)
+    clamped_x_t = tl.minimum(tl.maximum(x_t, 0), t_size - 1)
+    clamped_src_t = tl.minimum(tl.maximum(src_t, 0), cache_t - 1)
+
+    x_offsets = (
+        ((ob * channels + oc) * t_size + clamped_x_t) * h_size + clamped_ih
+    ) * w_size + clamped_iw
+    cache_offsets = (
+        ((ob * channels + oc) * cache_t + clamped_src_t) * h_size + clamped_ih
+    ) * w_size + clamped_iw
+
+    x_vals = tl.load(x_ptr + x_offsets, mask=valid & ~from_cache, other=0.0)
+    cache_vals = tl.load(cache_ptr + cache_offsets, mask=valid & from_cache, other=0.0)
+    vals = tl.where(from_cache, cache_vals, x_vals)
+    tl.store(out_ptr + offsets, vals, mask=mask)
+
+
+def fused_causal_conv3d_cat_pad(
+    x: torch.Tensor,
+    cache_x: torch.Tensor,
+    padding: list[int] | tuple[int, ...],
+) -> torch.Tensor:
+    width_left, width_right, height_top, height_bottom, depth_left, depth_right = (
+        padding
+    )
+    depth_left -= cache_x.shape[2]
+    assert depth_left >= 0
+    assert depth_right == 0
+    assert width_left == width_right
+    assert height_top == height_bottom
+
+    bsz, channels, t_size, h_size, w_size = x.shape
+    cache_t = cache_x.shape[2]
+    out = torch.empty(
+        (
+            bsz,
+            channels,
+            t_size + cache_t + depth_left + depth_right,
+            h_size + height_top + height_bottom,
+            w_size + width_left + width_right,
+        ),
+        device=x.device,
+        dtype=x.dtype,
+    )
+    block_size = 256
+    total = out.numel()
+    grid = (triton.cdiv(total, block_size),)
+    with torch.get_device_module().device(x.device):
+        _fused_cat_pad_5d_kernel[grid](
+            x,
+            cache_x,
+            out,
+            total,
+            channels,
+            t_size,
+            h_size,
+            w_size,
+            cache_t,
+            out.shape[2],
+            out.shape[3],
+            out.shape[4],
+            depth_left,
+            height_top,
+            width_left,
+            block_size,
+        )
+    return out

python/sglang/jit_kernel/diffusion/triton/group_norm_silu.py

@@ -0,0 +1,412 @@
+import math
+
+import torch
+import torch.nn.functional as F
+import triton  # type: ignore
+import triton.language as tl  # type: ignore
+
+from sglang.srt.utils.custom_op import register_custom_op
+
+_SUPPORTED_DTYPES = {torch.float16, torch.bfloat16, torch.float32}
+_LARGE_GROUP_THRESHOLD = 1 << 18
+_BLOCK_SIZE = 4096
+_BLOCKS_PER_PROGRAM = 2
+_CHUNK_SIZE = _BLOCK_SIZE * _BLOCKS_PER_PROGRAM
+
+
+@triton.jit
+def _group_norm_silu_contiguous_kernel(
+    input_ptr,
+    weight_ptr,
+    bias_ptr,
+    output_ptr,
+    channels,
+    spatial_size,
+    channels_per_group,
+    group_size,
+    eps,
+    BLOCK_SIZE: tl.constexpr,
+):
+    group_id = tl.program_id(0).to(tl.int64)
+    batch_id = tl.program_id(1).to(tl.int64)
+
+    group_base = batch_id * channels * spatial_size + group_id * group_size
+    offsets = tl.arange(0, BLOCK_SIZE)
+
+    sum_val = tl.zeros((), dtype=tl.float32)
+    sum_sq = tl.zeros((), dtype=tl.float32)
+    for off in range(0, group_size, BLOCK_SIZE):
+        idx = off + offsets
+        mask = idx < group_size
+        x = tl.load(input_ptr + group_base + idx, mask=mask, other=0.0).to(tl.float32)
+        sum_val += tl.sum(x, axis=0)
+        sum_sq += tl.sum(x * x, axis=0)
+
+    inv_group = 1.0 / group_size
+    mean = sum_val * inv_group
+    var = sum_sq * inv_group - mean * mean
+    rstd = tl.rsqrt(var + eps)
+
+    weight_group_offset = group_id * channels_per_group
+    for off in range(0, group_size, BLOCK_SIZE):
+        idx = off + offsets
+        mask = idx < group_size
+        x = tl.load(input_ptr + group_base + idx, mask=mask, other=0.0).to(tl.float32)
+        channel_offsets = weight_group_offset + idx // spatial_size
+        weight = tl.load(weight_ptr + channel_offsets, mask=mask, other=1.0).to(
+            tl.float32
+        )
+        bias = tl.load(bias_ptr + channel_offsets, mask=mask, other=0.0).to(tl.float32)
+        y = (x - mean) * rstd
+        y = y * weight + bias
+        y = y * tl.sigmoid(y)
+        tl.store(output_ptr + group_base + idx, y, mask=mask)
+
+
+@triton.jit
+def _group_norm_stats_kernel(
+    input_ptr,
+    partial_sum_ptr,
+    partial_sq_ptr,
+    channels,
+    spatial_size,
+    num_groups,
+    channels_per_group,
+    group_size,
+    chunks_per_row,
+    BLOCK_SIZE: tl.constexpr,
+    BLOCKS_PER_PROGRAM: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    chunk_id = tl.program_id(1).to(tl.int64)
+
+    batch_id = row // num_groups
+    group_id = row - batch_id * num_groups
+    chunk_start = chunk_id * BLOCK_SIZE * BLOCKS_PER_PROGRAM
+    group_base = batch_id * channels * spatial_size + group_id * group_size
+
+    sum_val = tl.zeros((), dtype=tl.float32)
+    sum_sq = tl.zeros((), dtype=tl.float32)
+    offsets = tl.arange(0, BLOCK_SIZE)
+
+    for block_id in range(BLOCKS_PER_PROGRAM):
+        idx = chunk_start + block_id * BLOCK_SIZE + offsets
+        mask = idx < group_size
+        x = tl.load(input_ptr + group_base + idx, mask=mask, other=0.0).to(tl.float32)
+        sum_val += tl.sum(x, axis=0)
+        sum_sq += tl.sum(x * x, axis=0)
+
+    partial_index = row * chunks_per_row + chunk_id
+    tl.store(partial_sum_ptr + partial_index, sum_val)
+    tl.store(partial_sq_ptr + partial_index, sum_sq)
+
+
+@triton.jit
+def _group_norm_finalize_stats_kernel(
+    partial_sum_ptr,
+    partial_sq_ptr,
+    stats_ptr,
+    chunks_per_row,
+    group_size,
+    eps,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    offsets = tl.arange(0, BLOCK_SIZE)
+
+    sum_val = tl.zeros((), dtype=tl.float32)
+    sum_sq = tl.zeros((), dtype=tl.float32)
+    base = row * chunks_per_row
+    for off in range(0, chunks_per_row, BLOCK_SIZE):
+        idx = off + offsets
+        mask = idx < chunks_per_row
+        sum_val += tl.sum(
+            tl.load(partial_sum_ptr + base + idx, mask=mask, other=0.0), axis=0
+        )
+        sum_sq += tl.sum(
+            tl.load(partial_sq_ptr + base + idx, mask=mask, other=0.0), axis=0
+        )
+
+    inv_group = 1.0 / group_size
+    mean = sum_val * inv_group
+    var = sum_sq * inv_group - mean * mean
+    rstd = tl.rsqrt(var + eps)
+    tl.store(stats_ptr + row * 2, mean)
+    tl.store(stats_ptr + row * 2 + 1, rstd)
+
+
+@triton.jit
+def _group_norm_apply_kernel(
+    input_ptr,
+    weight_ptr,
+    bias_ptr,
+    output_ptr,
+    stats_ptr,
+    channels,
+    spatial_size,
+    num_groups,
+    channels_per_group,
+    group_size,
+    chunks_per_row,
+    BLOCK_SIZE: tl.constexpr,
+    BLOCKS_PER_PROGRAM: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    chunk_id = tl.program_id(1).to(tl.int64)
+
+    batch_id = row // num_groups
+    group_id = row - batch_id * num_groups
+    chunk_start = chunk_id * BLOCK_SIZE * BLOCKS_PER_PROGRAM
+    group_base = batch_id * channels * spatial_size + group_id * group_size
+    weight_group_offset = group_id * channels_per_group
+
+    mean = tl.load(stats_ptr + row * 2)
+    rstd = tl.load(stats_ptr + row * 2 + 1)
+    offsets = tl.arange(0, BLOCK_SIZE)
+
+    for block_id in range(BLOCKS_PER_PROGRAM):
+        idx = chunk_start + block_id * BLOCK_SIZE + offsets
+        mask = idx < group_size
+        x = tl.load(input_ptr + group_base + idx, mask=mask, other=0.0).to(tl.float32)
+        channel_offsets = weight_group_offset + idx // spatial_size
+        weight = tl.load(weight_ptr + channel_offsets, mask=mask, other=1.0).to(
+            tl.float32
+        )
+        bias = tl.load(bias_ptr + channel_offsets, mask=mask, other=0.0).to(tl.float32)
+        y = (x - mean) * rstd
+        y = y * weight + bias
+        y = y * tl.sigmoid(y)
+        tl.store(output_ptr + group_base + idx, y, mask=mask)
+
+
+@triton.jit
+def _group_norm_apply_scalar_affine_kernel(
+    input_ptr,
+    weight_ptr,
+    bias_ptr,
+    output_ptr,
+    stats_ptr,
+    channels,
+    spatial_size,
+    num_groups,
+    channels_per_group,
+    group_size,
+    chunks_per_row,
+    BLOCK_SIZE: tl.constexpr,
+    BLOCKS_PER_PROGRAM: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    chunk_id = tl.program_id(1).to(tl.int64)
+
+    batch_id = row // num_groups
+    group_id = row - batch_id * num_groups
+    chunk_start = chunk_id * BLOCK_SIZE * BLOCKS_PER_PROGRAM
+    group_base = batch_id * channels * spatial_size + group_id * group_size
+
+    channel_id = chunk_start // spatial_size
+    affine_offset = group_id * channels_per_group + channel_id
+    weight = tl.load(weight_ptr + affine_offset).to(tl.float32)
+    bias = tl.load(bias_ptr + affine_offset).to(tl.float32)
+
+    mean = tl.load(stats_ptr + row * 2)
+    rstd = tl.load(stats_ptr + row * 2 + 1)
+    offsets = tl.arange(0, BLOCK_SIZE)
+
+    for block_id in range(BLOCKS_PER_PROGRAM):
+        idx = chunk_start + block_id * BLOCK_SIZE + offsets
+        mask = idx < group_size
+        x = tl.load(input_ptr + group_base + idx, mask=mask, other=0.0).to(tl.float32)
+        y = (x - mean) * rstd
+        y = y * weight + bias
+        y = y * tl.sigmoid(y)
+        tl.store(output_ptr + group_base + idx, y, mask=mask)
+
+
+def _group_norm_silu_native(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    num_groups: int,
+    eps: float,
+) -> torch.Tensor:
+    return F.silu(F.group_norm(x, num_groups, weight=weight, bias=bias, eps=eps))
+
+
+def _can_use_triton_group_norm_silu(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    num_groups: int,
+) -> bool:
+    return (
+        x.is_cuda
+        and not torch.is_grad_enabled()
+        and not x.requires_grad
+        and x.dtype in _SUPPORTED_DTYPES
+        and x.ndim in (2, 3, 4, 5)
+        and x.shape[1] % num_groups == 0
+        and weight.is_cuda
+        and bias.is_cuda
+        and weight.dtype == x.dtype
+        and bias.dtype == x.dtype
+        and weight.ndim == 1
+        and bias.ndim == 1
+        and weight.shape == bias.shape == (x.shape[1],)
+    )
+
+
+def _launch_one_pass(
+    x_contiguous: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    num_groups: int,
+    eps: float,
+) -> torch.Tensor:
+    batch_size, channels = x_contiguous.shape[:2]
+    spatial_size = math.prod(x_contiguous.shape[2:]) if x_contiguous.ndim > 2 else 1
+    channels_per_group = channels // num_groups
+    group_size = channels_per_group * spatial_size
+
+    x_flat = x_contiguous.reshape(batch_size, channels, spatial_size, 1)
+    y_flat = torch.empty_like(x_flat)
+    block_size = min(4096, triton.next_power_of_2(max(1, min(group_size, 4096))))
+
+    _group_norm_silu_contiguous_kernel[(num_groups, batch_size)](
+        x_flat,
+        weight,
+        bias,
+        y_flat,
+        channels,
+        spatial_size,
+        channels_per_group,
+        group_size,
+        eps,
+        BLOCK_SIZE=block_size,
+    )
+    return y_flat.reshape_as(x_contiguous)
+
+
+def _launch_chunked(
+    x_contiguous: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    num_groups: int,
+    eps: float,
+) -> torch.Tensor:
+    batch_size, channels = x_contiguous.shape[:2]
+    spatial_size = math.prod(x_contiguous.shape[2:]) if x_contiguous.ndim > 2 else 1
+    channels_per_group = channels // num_groups
+    group_size = channels_per_group * spatial_size
+    rows = batch_size * num_groups
+    chunks_per_row = triton.cdiv(group_size, _CHUNK_SIZE)
+
+    x_flat = x_contiguous.reshape(-1)
+    y = torch.empty_like(x_contiguous)
+    y_flat = y.reshape(-1)
+    partial_sum = torch.empty(
+        (rows, chunks_per_row), device=x_contiguous.device, dtype=torch.float32
+    )
+    partial_sq = torch.empty_like(partial_sum)
+    stats = torch.empty((rows, 2), device=x_contiguous.device, dtype=torch.float32)
+
+    _group_norm_stats_kernel[(rows, chunks_per_row)](
+        x_flat,
+        partial_sum,
+        partial_sq,
+        channels,
+        spatial_size,
+        num_groups,
+        channels_per_group,
+        group_size,
+        chunks_per_row,
+        BLOCK_SIZE=_BLOCK_SIZE,
+        BLOCKS_PER_PROGRAM=_BLOCKS_PER_PROGRAM,
+        num_warps=8,
+        num_stages=3,
+    )
+
+    reduce_block = min(1024, triton.next_power_of_2(max(1, chunks_per_row)))
+    _group_norm_finalize_stats_kernel[(rows,)](
+        partial_sum,
+        partial_sq,
+        stats,
+        chunks_per_row,
+        group_size,
+        eps,
+        BLOCK_SIZE=reduce_block,
+        num_warps=4,
+        num_stages=2,
+    )
+
+    if spatial_size % _CHUNK_SIZE == 0 and chunks_per_row >= 64:
+        _group_norm_apply_scalar_affine_kernel[(rows, chunks_per_row)](
+            x_flat,
+            weight,
+            bias,
+            y_flat,
+            stats,
+            channels,
+            spatial_size,
+            num_groups,
+            channels_per_group,
+            group_size,
+            chunks_per_row,
+            BLOCK_SIZE=_BLOCK_SIZE,
+            BLOCKS_PER_PROGRAM=_BLOCKS_PER_PROGRAM,
+            num_warps=4,
+            num_stages=3,
+        )
+    else:
+        _group_norm_apply_kernel[(rows, chunks_per_row)](
+            x_flat,
+            weight,
+            bias,
+            y_flat,
+            stats,
+            channels,
+            spatial_size,
+            num_groups,
+            channels_per_group,
+            group_size,
+            chunks_per_row,
+            BLOCK_SIZE=_BLOCK_SIZE,
+            BLOCKS_PER_PROGRAM=_BLOCKS_PER_PROGRAM,
+            num_warps=8,
+            num_stages=3,
+        )
+    return y
+
+
+@register_custom_op(op_name="triton_group_norm_silu_cuda", out_shape="x")
+def _triton_group_norm_silu_cuda(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    num_groups: int,
+    eps: float = 1e-5,
+) -> torch.Tensor:
+    if not _can_use_triton_group_norm_silu(x, weight, bias, num_groups):
+        return _group_norm_silu_native(x, weight, bias, num_groups, eps)
+
+    x_contiguous = x.contiguous()
+    spatial_size = math.prod(x_contiguous.shape[2:]) if x_contiguous.ndim > 2 else 1
+    channels_per_group = x_contiguous.shape[1] // num_groups
+    group_size = channels_per_group * spatial_size
+
+    with torch.cuda.device(x.device):
+        if group_size >= _LARGE_GROUP_THRESHOLD:
+            return _launch_chunked(x_contiguous, weight, bias, num_groups, eps)
+        return _launch_one_pass(x_contiguous, weight, bias, num_groups, eps)
+
+
+def triton_group_norm_silu(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    num_groups: int,
+    eps: float = 1e-5,
+) -> torch.Tensor:
+    return _triton_group_norm_silu_cuda(x, weight, bias, num_groups, eps)
+
+
+__all__ = ["triton_group_norm_silu"]

python/sglang/jit_kernel/diffusion/triton/ltx2_ada_values.py

@@ -0,0 +1,184 @@
+# Adapted from NVlabs/Sana sol-engine LTX2 Ada-value fusion.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _ltx2_ada_values9_kernel(
+    temb_ptr,
+    table_ptr,
+    out0_ptr,
+    out1_ptr,
+    out2_ptr,
+    out3_ptr,
+    out4_ptr,
+    out5_ptr,
+    out6_ptr,
+    out7_ptr,
+    out8_ptr,
+    rows: tl.constexpr,
+    hidden: tl.constexpr,
+    total_params: tl.constexpr,
+    table_stride_p: tl.constexpr,
+    table_stride_d: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    cols = tl.arange(0, BLOCK_N)
+    mask = cols < hidden
+    temb_row = temb_ptr + row * total_params * hidden
+    base = row * hidden + cols
+
+    table0 = tl.load(
+        table_ptr + 0 * table_stride_p + cols * table_stride_d,
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    temb0 = tl.load(
+        temb_row + (0 * hidden + cols),
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    table1 = tl.load(
+        table_ptr + 1 * table_stride_p + cols * table_stride_d,
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    temb1 = tl.load(
+        temb_row + (1 * hidden + cols),
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    table2 = tl.load(
+        table_ptr + 2 * table_stride_p + cols * table_stride_d,
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    temb2 = tl.load(
+        temb_row + (2 * hidden + cols),
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    table3 = tl.load(
+        table_ptr + 3 * table_stride_p + cols * table_stride_d,
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    temb3 = tl.load(
+        temb_row + (3 * hidden + cols),
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    table4 = tl.load(
+        table_ptr + 4 * table_stride_p + cols * table_stride_d,
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    temb4 = tl.load(
+        temb_row + (4 * hidden + cols),
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    table5 = tl.load(
+        table_ptr + 5 * table_stride_p + cols * table_stride_d,
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    temb5 = tl.load(
+        temb_row + (5 * hidden + cols),
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    table6 = tl.load(
+        table_ptr + 6 * table_stride_p + cols * table_stride_d,
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    temb6 = tl.load(
+        temb_row + (6 * hidden + cols),
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    table7 = tl.load(
+        table_ptr + 7 * table_stride_p + cols * table_stride_d,
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    temb7 = tl.load(
+        temb_row + (7 * hidden + cols),
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    table8 = tl.load(
+        table_ptr + 8 * table_stride_p + cols * table_stride_d,
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    temb8 = tl.load(
+        temb_row + (8 * hidden + cols),
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+
+    tl.store(out0_ptr + base, (table0 + temb0).to(tl.bfloat16), mask=mask)
+    tl.store(out1_ptr + base, (table1 + temb1).to(tl.bfloat16), mask=mask)
+    tl.store(out2_ptr + base, (table2 + temb2).to(tl.bfloat16), mask=mask)
+    tl.store(out3_ptr + base, (table3 + temb3).to(tl.bfloat16), mask=mask)
+    tl.store(out4_ptr + base, (table4 + temb4).to(tl.bfloat16), mask=mask)
+    tl.store(out5_ptr + base, (table5 + temb5).to(tl.bfloat16), mask=mask)
+    tl.store(out6_ptr + base, (table6 + temb6).to(tl.bfloat16), mask=mask)
+    tl.store(out7_ptr + base, (table7 + temb7).to(tl.bfloat16), mask=mask)
+    tl.store(out8_ptr + base, (table8 + temb8).to(tl.bfloat16), mask=mask)
+
+
+def ltx2_ada_values9(
+    scale_shift_table: torch.Tensor,
+    timestep: torch.Tensor,
+) -> tuple[torch.Tensor, ...]:
+    if timestep.ndim != 3:
+        raise ValueError("timestep must have shape [B, S, 9 * D]")
+    if not timestep.is_cuda or timestep.dtype != torch.bfloat16:
+        raise ValueError("timestep must be a CUDA bfloat16 tensor")
+    if not timestep.is_contiguous():
+        raise ValueError("timestep must be contiguous")
+    if scale_shift_table.ndim != 2 or scale_shift_table.shape[0] != 9:
+        raise ValueError("scale_shift_table must have shape [9, D]")
+    if (
+        not scale_shift_table.is_cuda
+        or scale_shift_table.dtype not in (torch.bfloat16, torch.float32)
+        or scale_shift_table.stride(-1) != 1
+    ):
+        raise ValueError(
+            "scale_shift_table must be CUDA, bf16/fp32, last-dim contiguous"
+        )
+
+    total_params = int(scale_shift_table.shape[0])
+    hidden = int(scale_shift_table.shape[1])
+    if hidden <= 0 or timestep.shape[-1] != total_params * hidden:
+        raise ValueError("timestep last dim must equal 9 * hidden")
+    if hidden % 256 != 0 or hidden > 8192:
+        raise ValueError("hidden size is outside the supported LTX2 fast-path range")
+
+    batch, seq, _ = timestep.shape
+    rows = int(batch * seq)
+    outs = tuple(
+        torch.empty((batch, seq, hidden), device=timestep.device, dtype=timestep.dtype)
+        for _ in range(9)
+    )
+    _ltx2_ada_values9_kernel[(rows,)](
+        timestep,
+        scale_shift_table,
+        *outs,
+        rows,
+        hidden,
+        total_params,
+        scale_shift_table.stride(0),
+        scale_shift_table.stride(1),
+        BLOCK_N=triton.next_power_of_2(hidden),
+        num_warps=4 if hidden >= 4096 else 8,
+    )
+    return outs

python/sglang/jit_kernel/diffusion/triton/ltx2_rotary.py

@@ -0,0 +1,102 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _ltx2_split_rotary_kernel(
+    out_ptr,
+    x_ptr,
+    cos_ptr,
+    sin_ptr,
+    seq_len: tl.constexpr,
+    num_heads: tl.constexpr,
+    head_dim: tl.constexpr,
+    half_dim: tl.constexpr,
+    stride_cos_b: tl.constexpr,
+    stride_cos_h: tl.constexpr,
+    stride_cos_t: tl.constexpr,
+    stride_sin_b: tl.constexpr,
+    stride_sin_h: tl.constexpr,
+    stride_sin_t: tl.constexpr,
+    BLOCK_HEADS: tl.constexpr,
+    BLOCK_HALF: tl.constexpr,
+):
+    pid_bt = tl.program_id(0)
+    head_block = tl.program_id(1)
+    batch = pid_bt // seq_len
+    token = pid_bt - batch * seq_len
+    heads = head_block * BLOCK_HEADS + tl.arange(0, BLOCK_HEADS)
+    offsets = tl.arange(0, BLOCK_HALF)
+    mask = (heads[:, None] < num_heads) & (offsets[None, :] < half_dim)
+
+    x_base = ((batch * seq_len + token) * num_heads + heads[:, None]) * head_dim
+    cos_base = (
+        batch * stride_cos_b + heads[:, None] * stride_cos_h + token * stride_cos_t
+    )
+    sin_base = (
+        batch * stride_sin_b + heads[:, None] * stride_sin_h + token * stride_sin_t
+    )
+
+    x_first = tl.load(x_ptr + x_base + offsets[None, :], mask=mask, other=0.0)
+    x_second = tl.load(
+        x_ptr + x_base + half_dim + offsets[None, :], mask=mask, other=0.0
+    )
+    cos = tl.load(cos_ptr + cos_base + offsets[None, :], mask=mask, other=0.0)
+    sin = tl.load(sin_ptr + sin_base + offsets[None, :], mask=mask, other=0.0)
+
+    # Match the original PyTorch order: x * cos is written as BF16 first, then
+    # addcmul_ computes the sine product in FP32 before the final BF16 store.
+    out_first = (x_first * cos).to(tl.bfloat16).to(tl.float32) + (
+        -x_second.to(tl.float32) * sin.to(tl.float32)
+    )
+    out_second = (x_second * cos).to(tl.bfloat16).to(tl.float32) + (
+        x_first.to(tl.float32) * sin.to(tl.float32)
+    )
+
+    tl.store(out_ptr + x_base + offsets[None, :], out_first, mask=mask)
+    tl.store(out_ptr + x_base + half_dim + offsets[None, :], out_second, mask=mask)
+
+
+def apply_ltx2_split_rotary_emb(
+    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> torch.Tensor:
+    batch, seq_len, inner_dim = x.shape
+    cos_batch, num_heads, cos_seq_len, half_dim = cos.shape
+    head_dim = half_dim * 2
+    if (
+        cos_batch != batch
+        or cos_seq_len != seq_len
+        or inner_dim != num_heads * head_dim
+        or sin.shape != cos.shape
+    ):
+        raise ValueError(
+            "LTX2 split RoPE shape mismatch: "
+            f"x={tuple(x.shape)}, cos={tuple(cos.shape)}, sin={tuple(sin.shape)}"
+        )
+
+    out = torch.empty_like(x)
+    block_half = triton.next_power_of_2(half_dim)
+    block_heads = min(16, triton.next_power_of_2(num_heads))
+    num_warps = min(8, max(1, block_heads))
+    grid = (batch * seq_len, triton.cdiv(num_heads, block_heads))
+    _ltx2_split_rotary_kernel[grid](
+        out,
+        x,
+        cos,
+        sin,
+        seq_len,
+        num_heads,
+        head_dim,
+        half_dim,
+        cos.stride(0),
+        cos.stride(1),
+        cos.stride(2),
+        sin.stride(0),
+        sin.stride(1),
+        sin.stride(2),
+        BLOCK_HEADS=block_heads,
+        BLOCK_HALF=block_half,
+        num_warps=num_warps,
+    )
+    return out

python/sglang/jit_kernel/diffusion/triton/mps_fallback.py

@@ -0,0 +1,125 @@
+"""MPS (Apple Silicon) fallbacks for Triton diffusion kernels.
+
+Triton is not available on macOS / Metal, so these pure-PyTorch (and
+optionally MLX-accelerated) implementations replace the Triton kernels
+at import time when ``current_platform.is_mps()`` is True.
+
+MLX acceleration (opt-in via ``SGLANG_USE_MLX=1``):
+    Norm ops use ``mx.fast.rms_norm`` / ``mx.fast.layer_norm`` — single fused
+    Metal kernels that are 1.4x–2.9x faster than the multi-step PyTorch MPS
+    decomposition for medium-to-large tensors.
+"""
+
+from typing import Optional
+
+import torch
+from torch import Tensor
+
+from sglang.srt.utils.tensor_bridge import mlx_to_torch, torch_to_mlx, use_mlx
+
+from .torch_fallback import (
+    apply_rotary_embedding_native,
+    fuse_scale_shift_kernel_native,
+    norm_infer_native,
+    rms_norm_fn_native,
+    triton_one_pass_rms_norm_native,
+)
+
+_use_mlx = use_mlx()
+
+if _use_mlx:
+    import mlx.core as mx
+
+# use the common torch native version form torch_fallback
+fuse_scale_shift_kernel_native = fuse_scale_shift_kernel_native
+apply_rotary_embedding_native = apply_rotary_embedding_native
+norm_infer_native = norm_infer_native
+triton_one_pass_rms_norm_native = triton_one_pass_rms_norm_native
+rms_norm_fn_native = rms_norm_fn_native
+
+# MLX-accelerated norm ops (1.4x–2.9x faster than torch native on MPS)
+# Uses mx.fast.rms_norm / mx.fast.layer_norm — single fused Metal kernels
+# instead of 7+ separate PyTorch MPS kernel launches.
+
+if _use_mlx:
+
+    def norm_infer_native(  # noqa: F811
+        x: Tensor,
+        weight: Optional[Tensor],
+        bias: Optional[Tensor],
+        eps: float,
+        is_rms_norm: bool = False,
+        out: Optional[Tensor] = None,
+    ) -> Tensor:
+        """MLX-accelerated norm_infer (layer norm / rms norm inference)."""
+        device = x.device
+        orig_dtype = x.dtype
+        x_mx = torch_to_mlx(x)
+        if is_rms_norm:
+            w_mx = (
+                torch_to_mlx(weight) if weight is not None else mx.ones(x_mx.shape[-1])
+            )
+            result_mx = mx.fast.rms_norm(x_mx, w_mx, eps)
+        else:
+            w_mx = torch_to_mlx(weight) if weight is not None else None
+            b_mx = torch_to_mlx(bias) if bias is not None else None
+            result_mx = mx.fast.layer_norm(x_mx, w_mx, b_mx, eps)
+        result = mlx_to_torch(result_mx, device).to(orig_dtype)
+        if out is not None:
+            out.copy_(result)
+            return out
+        return result
+
+    def triton_one_pass_rms_norm_native(  # noqa: F811
+        x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6
+    ) -> torch.Tensor:
+        """MLX-accelerated triton_one_pass_rms_norm."""
+        device = x.device
+        orig_dtype = x.dtype
+        x_mx = torch_to_mlx(x)
+        w_mx = torch_to_mlx(w)
+        result_mx = mx.fast.rms_norm(x_mx, w_mx, eps)
+        return mlx_to_torch(result_mx, device).to(orig_dtype)
+
+    def rms_norm_fn_native(  # noqa: F811
+        x,
+        weight,
+        bias,
+        residual=None,
+        x1=None,
+        weight1=None,
+        bias1=None,
+        eps=1e-6,
+        dropout_p=0.0,
+        rowscale=None,
+        prenorm=False,
+        residual_in_fp32=False,
+        zero_centered_weight=False,
+        return_dropout_mask=False,
+        out_dtype=None,
+        out=None,
+        residual_out=None,
+    ):
+        """MLX-accelerated rms_norm_fn (inference only, no dropout/x1 support)."""
+        device = x.device
+        orig_dtype = x.dtype
+        if residual is not None:
+            x = x.float() + residual.float()
+            residual_out_val = x.to(torch.float32 if residual_in_fp32 else orig_dtype)
+        else:
+            residual_out_val = None
+        if weight is not None and zero_centered_weight:
+            w = weight.float() + 1.0
+        else:
+            w = weight
+        x_mx = torch_to_mlx(x)
+        w_mx = torch_to_mlx(w) if w is not None else mx.ones(x_mx.shape[-1])
+        result_mx = mx.fast.rms_norm(x_mx, w_mx, eps)
+        x_hat = mlx_to_torch(result_mx, device)
+        if bias is not None:
+            x_hat = x_hat + bias.to(x_hat.device, x_hat.dtype)
+        final_dtype = out_dtype if out_dtype is not None else orig_dtype
+        y = x_hat.to(final_dtype)
+        if residual is not None and residual_out_val is not None:
+            return y, residual_out_val
+        return y

python/sglang/jit_kernel/diffusion/triton/norm.py

@@ -0,0 +1,660 @@
+from typing import Optional, Tuple
+
+import torch
+import triton  # type: ignore
+import triton.language as tl  # type: ignore
+from torch import Tensor
+
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.srt.utils.custom_op import register_custom_op
+
+
+# RMSNorm-fp32
+def maybe_contiguous_lastdim(x):
+    return x.contiguous() if x is not None and x.stride(-1) != 1 else x
+
+
+def maybe_contiguous(x):
+    return x.contiguous() if x is not None else None
+
+
+def triton_autotune_configs():
+    # Return configs with a valid warp count for the current device
+    # Maximum threads per block is architecture-dependent in theory, but in reality all are 1024
+    max_threads_per_block = 1024
+    # Default to warp size 32 if not defined by device
+    warp_size = getattr(
+        torch.get_device_module().get_device_properties(
+            torch.get_device_module().current_device()
+        ),
+        "warp_size",
+        32,
+    )
+    if warp_size is None:
+        warp_size = 32
+    # Autotune for warp counts which are powers of 2 and do not exceed thread per block limit
+    return [
+        triton.Config({}, num_warps=warp_count)
+        for warp_count in [1, 2, 4, 8, 16, 32]
+        if warp_count * warp_size <= max_threads_per_block
+    ]
+
+
+# Copied from flash-attn
+@triton.autotune(
+    configs=triton_autotune_configs(),
+    key=[
+        "N",
+        "HAS_RESIDUAL",
+        "STORE_RESIDUAL_OUT",
+        "IS_RMS_NORM",
+        "HAS_BIAS",
+        "HAS_WEIGHT",
+        "HAS_X1",
+        "HAS_W1",
+        "HAS_B1",
+    ],
+)
+# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel
+# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None})
+# @triton.heuristics({"HAS_RESIDUAL": lambda args: args["RESIDUAL"] is not None})
+# @triton.heuristics({"HAS_X1": lambda args: args["X1"] is not None})
+# @triton.heuristics({"HAS_W1": lambda args: args["W1"] is not None})
+# @triton.heuristics({"HAS_B1": lambda args: args["B1"] is not None})
+@triton.jit
+def _layer_norm_fwd_1pass_kernel(
+    X,  # pointer to the input
+    Y,  # pointer to the output
+    W,  # pointer to the weights
+    B,  # pointer to the biases
+    RESIDUAL,  # pointer to the residual
+    X1,
+    W1,
+    B1,
+    Y1,
+    RESIDUAL_OUT,  # pointer to the residual
+    ROWSCALE,
+    SEEDS,  # Dropout seeds for each row
+    DROPOUT_MASK,
+    DROPOUT_MASK1,
+    Mean,  # pointer to the mean
+    Rstd,  # pointer to the 1/std
+    stride_x_row,  # how much to increase the pointer when moving by 1 row
+    stride_y_row,
+    stride_res_row,
+    stride_res_out_row,
+    stride_x1_row,
+    stride_y1_row,
+    M,  # number of rows in X
+    N,  # number of columns in X
+    eps,  # epsilon to avoid division by zero
+    dropout_p,  # Dropout probability
+    zero_centered_weight,  # If true, add 1.0 to the weight
+    IS_RMS_NORM: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    HAS_RESIDUAL: tl.constexpr,
+    STORE_RESIDUAL_OUT: tl.constexpr,
+    HAS_WEIGHT: tl.constexpr,
+    HAS_BIAS: tl.constexpr,
+    HAS_DROPOUT: tl.constexpr,
+    STORE_DROPOUT_MASK: tl.constexpr,
+    HAS_ROWSCALE: tl.constexpr,
+    HAS_X1: tl.constexpr,
+    HAS_W1: tl.constexpr,
+    HAS_B1: tl.constexpr,
+):
+    # Map the program id to the row of X and Y it should compute.
+    row = tl.program_id(0)
+    X += row * stride_x_row
+    Y += row * stride_y_row
+    if HAS_RESIDUAL:
+        RESIDUAL += row * stride_res_row
+    if STORE_RESIDUAL_OUT:
+        RESIDUAL_OUT += row * stride_res_out_row
+    if HAS_X1:
+        X1 += row * stride_x1_row
+    if HAS_W1:
+        Y1 += row * stride_y1_row
+    # Compute mean and variance
+    cols = tl.arange(0, BLOCK_N)
+    x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32)
+    if HAS_ROWSCALE:
+        rowscale = tl.load(ROWSCALE + row).to(tl.float32)
+        x *= rowscale
+    if HAS_DROPOUT:
+        # Compute dropout mask
+        # 7 rounds is good enough, and reduces register pressure
+        keep_mask = (
+            tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p
+        )
+        x = tl.where(keep_mask, x / (1.0 - dropout_p), 0.0)
+        if STORE_DROPOUT_MASK:
+            tl.store(DROPOUT_MASK + row * N + cols, keep_mask, mask=cols < N)
+    if HAS_X1:
+        x1 = tl.load(X1 + cols, mask=cols < N, other=0.0).to(tl.float32)
+        if HAS_ROWSCALE:
+            rowscale = tl.load(ROWSCALE + M + row).to(tl.float32)
+            x1 *= rowscale
+        if HAS_DROPOUT:
+            # Compute dropout mask
+            # 7 rounds is good enough, and reduces register pressure
+            keep_mask = (
+                tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7)
+                > dropout_p
+            )
+            x1 = tl.where(keep_mask, x1 / (1.0 - dropout_p), 0.0)
+            if STORE_DROPOUT_MASK:
+                tl.store(DROPOUT_MASK1 + row * N + cols, keep_mask, mask=cols < N)
+        x += x1
+    if HAS_RESIDUAL:
+        residual = tl.load(RESIDUAL + cols, mask=cols < N, other=0.0).to(tl.float32)
+        x += residual
+    if STORE_RESIDUAL_OUT:
+        tl.store(RESIDUAL_OUT + cols, x, mask=cols < N)
+    if not IS_RMS_NORM:
+        mean = tl.sum(x, axis=0) / N
+        tl.store(Mean + row, mean)
+        xbar = tl.where(cols < N, x - mean, 0.0)
+        var = tl.sum(xbar * xbar, axis=0) / N
+    else:
+        xbar = tl.where(cols < N, x, 0.0)
+        var = tl.sum(xbar * xbar, axis=0) / N
+    rstd = 1 / tl.sqrt(var + eps)
+    tl.store(Rstd + row, rstd)
+    # Normalize and apply linear transformation
+    mask = cols < N
+    if HAS_WEIGHT:
+        w = tl.load(W + cols, mask=mask).to(tl.float32)
+        if zero_centered_weight:
+            w += 1.0
+    if HAS_BIAS:
+        b = tl.load(B + cols, mask=mask).to(tl.float32)
+    x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd
+    if HAS_WEIGHT:
+        y = x_hat * w + b if HAS_BIAS else x_hat * w
+    else:
+        y = x_hat + b if HAS_BIAS else x_hat
+    # Write output
+    tl.store(Y + cols, y, mask=mask)
+    if HAS_W1:
+        w1 = tl.load(W1 + cols, mask=mask).to(tl.float32)
+        if zero_centered_weight:
+            w1 += 1.0
+        if HAS_B1:
+            b1 = tl.load(B1 + cols, mask=mask).to(tl.float32)
+        y1 = x_hat * w1 + b1 if HAS_B1 else x_hat * w1
+        tl.store(Y1 + cols, y1, mask=mask)
+
+
+def _layer_norm_fwd(
+    x: Tensor,
+    weight: Optional[Tensor],
+    bias: Optional[Tensor],
+    eps: float,
+    residual: Optional[Tensor] = None,
+    x1: Optional[Tensor] = None,
+    weight1: Optional[Tensor] = None,
+    bias1: Optional[Tensor] = None,
+    dropout_p: float = 0.0,
+    rowscale: Optional[Tensor] = None,
+    out_dtype: Optional[torch.dtype] = None,
+    residual_dtype: Optional[torch.dtype] = None,
+    zero_centered_weight: bool = False,
+    is_rms_norm: bool = False,
+    return_dropout_mask: bool = False,
+    out: Optional[Tensor] = None,
+    residual_out: Optional[Tensor] = None,
+) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor):
+    # Allocate aliases upfront so the custom op only mutates explicit outputs.
+    if out is None:
+        out = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype)
+    if residual is not None:
+        residual_dtype = residual.dtype
+    if residual_out is None and (
+        residual is not None
+        or (residual_dtype is not None and residual_dtype != x.dtype)
+        or dropout_p > 0.0
+        or rowscale is not None
+        or x1 is not None
+    ):
+        residual_out = torch.empty_like(
+            x, dtype=residual_dtype if residual_dtype is not None else x.dtype
+        )
+    else:
+        residual_out = None
+    y1, mean, rstd, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd_impl(
+        x,
+        weight,
+        bias,
+        eps,
+        out,
+        residual=residual,
+        x1=x1,
+        weight1=weight1,
+        bias1=bias1,
+        dropout_p=dropout_p,
+        rowscale=rowscale,
+        zero_centered_weight=zero_centered_weight,
+        is_rms_norm=is_rms_norm,
+        return_dropout_mask=return_dropout_mask,
+        residual_out=residual_out,
+    )
+    # residual_out is None if residual is None and residual_dtype == input_dtype and dropout_p == 0.0
+    if residual_out is None:
+        residual_out = x
+    return out, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1
+
+
+@register_custom_op(
+    op_name="diffusion_layer_norm_fwd_impl_cuda",
+    mutates_args=[
+        "out",
+        "y1",
+        "mean",
+        "rstd",
+        "residual_out",
+        "dropout_mask",
+        "dropout_mask1",
+    ],
+)
+def _layer_norm_fwd_impl_cuda(
+    x: Tensor,
+    weight: Optional[Tensor],
+    bias: Optional[Tensor],
+    eps: float,
+    out: Tensor,
+    y1: Optional[Tensor],
+    mean: Optional[Tensor],
+    rstd: Tensor,
+    residual: Optional[Tensor] = None,
+    x1: Optional[Tensor] = None,
+    weight1: Optional[Tensor] = None,
+    bias1: Optional[Tensor] = None,
+    residual_out: Optional[Tensor] = None,
+    rowscale: Optional[Tensor] = None,
+    seeds: Optional[Tensor] = None,
+    dropout_mask: Optional[Tensor] = None,
+    dropout_mask1: Optional[Tensor] = None,
+    dropout_p: float = 0.0,
+    zero_centered_weight: bool = False,
+    is_rms_norm: bool = False,
+) -> None:
+    M, N = x.shape
+    assert x.stride(-1) == 1
+    if residual is not None:
+        assert residual.stride(-1) == 1
+        assert residual.shape == (M, N)
+    if weight is not None:
+        assert weight.shape == (N,)
+        assert weight.stride(-1) == 1
+    if bias is not None:
+        assert bias.stride(-1) == 1
+        assert bias.shape == (N,)
+    if x1 is not None:
+        assert x1.shape == x.shape
+        assert rowscale is None
+        assert x1.stride(-1) == 1
+    if weight1 is not None:
+        assert weight1.shape == (N,)
+        assert weight1.stride(-1) == 1
+    if bias1 is not None:
+        assert bias1.shape == (N,)
+        assert bias1.stride(-1) == 1
+    if rowscale is not None:
+        assert rowscale.is_contiguous()
+        assert rowscale.shape == (M,)
+    assert out.shape == x.shape
+    assert out.stride(-1) == 1
+    if residual_out is not None:
+        assert residual_out.shape == x.shape
+        assert residual_out.stride(-1) == 1
+    if y1 is not None:
+        assert y1.shape == x.shape
+        assert y1.stride(-1) == 1
+    if mean is not None:
+        assert mean.shape == (M,)
+    assert rstd.shape == (M,)
+    if seeds is not None:
+        assert seeds.shape == (M if x1 is None else 2 * M,)
+    if dropout_mask is not None:
+        assert dropout_mask.shape == (M, N)
+    if dropout_mask1 is not None:
+        assert dropout_mask1.shape == (M, N)
+    # Less than 64KB per feature: enqueue fused kernel
+    MAX_FUSED_SIZE = 65536 // x.element_size()
+    BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N))
+    if N > BLOCK_N:
+        raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
+    with torch.get_device_module().device(x.device):
+        _layer_norm_fwd_1pass_kernel[(M,)](
+            x,
+            out,
+            weight if weight is not None else x,  # unused when HAS_WEIGHT == False
+            bias,
+            residual,
+            x1,
+            weight1,
+            bias1,
+            y1,
+            residual_out,
+            rowscale,
+            seeds,
+            dropout_mask,
+            dropout_mask1,
+            mean,
+            rstd,
+            x.stride(0),
+            out.stride(0),
+            residual.stride(0) if residual is not None else 0,
+            residual_out.stride(0) if residual_out is not None else 0,
+            x1.stride(0) if x1 is not None else 0,
+            y1.stride(0) if y1 is not None else 0,
+            M,
+            N,
+            eps,
+            dropout_p,
+            # Passing bool make torch inductor very unhappy since it then tries to compare to int_max
+            int(zero_centered_weight),
+            is_rms_norm,
+            BLOCK_N,
+            residual is not None,
+            residual_out is not None,
+            weight is not None,
+            bias is not None,
+            dropout_p > 0.0,
+            dropout_mask is not None,
+            rowscale is not None,
+            HAS_X1=x1 is not None,
+            HAS_W1=weight1 is not None,
+            HAS_B1=bias1 is not None,
+        )
+    return None
+
+
+def _layer_norm_fwd_impl(
+    x: Tensor,
+    weight: Optional[Tensor],
+    bias: Optional[Tensor],
+    eps: float,
+    out: Tensor,
+    residual: Optional[Tensor] = None,
+    x1: Optional[Tensor] = None,
+    weight1: Optional[Tensor] = None,
+    bias1: Optional[Tensor] = None,
+    dropout_p: float = 0.0,
+    rowscale: Optional[Tensor] = None,
+    zero_centered_weight: bool = False,
+    is_rms_norm: bool = False,
+    return_dropout_mask: bool = False,
+    residual_out: Optional[Tensor] = None,
+) -> Tuple[
+    Optional[Tensor],
+    Optional[Tensor],
+    Tensor,
+    Optional[Tensor],
+    Optional[Tensor],
+    Optional[Tensor],
+]:
+    M, N = x.shape
+    y1 = torch.empty_like(out) if weight1 is not None else None
+    mean = (
+        torch.empty((M,), dtype=torch.float32, device=x.device)
+        if not is_rms_norm
+        else None
+    )
+    rstd = torch.empty((M,), dtype=torch.float32, device=x.device)
+    seeds = (
+        torch.randint(
+            2**32, (M if x1 is None else 2 * M), device=x.device, dtype=torch.int64
+        )
+        if dropout_p > 0.0
+        else None
+    )
+    if return_dropout_mask and dropout_p > 0.0:
+        dropout_mask = torch.empty((M, N), dtype=torch.bool, device=x.device)
+        dropout_mask1 = (
+            torch.empty((M, N), dtype=torch.bool, device=x.device)
+            if x1 is not None
+            else None
+        )
+    else:
+        dropout_mask = dropout_mask1 = None
+    _layer_norm_fwd_impl_cuda(
+        x,
+        weight,
+        bias,
+        eps,
+        out,
+        y1,
+        mean,
+        rstd,
+        residual=residual,
+        x1=x1,
+        weight1=weight1,
+        bias1=bias1,
+        residual_out=residual_out,
+        rowscale=rowscale,
+        seeds=seeds,
+        dropout_mask=dropout_mask,
+        dropout_mask1=dropout_mask1,
+        dropout_p=dropout_p,
+        zero_centered_weight=zero_centered_weight,
+        is_rms_norm=is_rms_norm,
+    )
+    return y1, mean, rstd, seeds, dropout_mask, dropout_mask1
+
+
+def _norm_forward(
+    x,
+    weight,
+    bias,
+    residual=None,
+    x1=None,
+    weight1=None,
+    bias1=None,
+    eps=1e-6,
+    dropout_p=0.0,
+    rowscale=None,
+    prenorm=False,
+    residual_in_fp32=False,
+    zero_centered_weight=False,
+    is_rms_norm=False,
+    return_dropout_mask=False,
+    out_dtype=None,
+    out=None,
+    residual_out=None,
+):
+    x_shape_og = x.shape
+    # reshape input data into 2D tensor
+    x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1]))
+    if residual is not None:
+        assert residual.shape == x_shape_og
+        residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1]))
+    if x1 is not None:
+        assert x1.shape == x_shape_og
+        assert rowscale is None, "rowscale is not supported with parallel LayerNorm"
+        x1 = maybe_contiguous_lastdim(x1.reshape(-1, x1.shape[-1]))
+    # weight can be None when elementwise_affine=False for LayerNorm
+    if weight is not None:
+        weight = weight.contiguous()
+    bias = maybe_contiguous(bias)
+    weight1 = maybe_contiguous(weight1)
+    bias1 = maybe_contiguous(bias1)
+    if rowscale is not None:
+        rowscale = rowscale.reshape(-1).contiguous()
+    residual_dtype = (
+        residual.dtype
+        if residual is not None
+        else (torch.float32 if residual_in_fp32 else None)
+    )
+    if out is not None:
+        out = out.reshape(-1, out.shape[-1])
+    if residual_out is not None:
+        residual_out = residual_out.reshape(-1, residual_out.shape[-1])
+    y, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 = (
+        _layer_norm_fwd(
+            x,
+            weight,
+            bias,
+            eps,
+            residual,
+            x1,
+            weight1,
+            bias1,
+            dropout_p=dropout_p,
+            rowscale=rowscale,
+            out_dtype=out_dtype,
+            residual_dtype=residual_dtype,
+            zero_centered_weight=zero_centered_weight,
+            is_rms_norm=is_rms_norm,
+            return_dropout_mask=return_dropout_mask,
+            out=out,
+            residual_out=residual_out,
+        )
+    )
+    y = y.reshape(x_shape_og)
+    if residual is not None:
+        residual_out = residual_out.reshape(x_shape_og)
+        return y, residual_out
+    return y
+
+
+def rms_norm_fn(
+    x,
+    weight,
+    bias,
+    residual=None,
+    x1=None,
+    weight1=None,
+    bias1=None,
+    eps=1e-6,
+    dropout_p=0.0,
+    rowscale=None,
+    prenorm=False,
+    residual_in_fp32=False,
+    zero_centered_weight=False,
+    return_dropout_mask=False,
+    out_dtype=None,
+    out=None,
+    residual_out=None,
+):
+    return _norm_forward(
+        x,
+        weight,
+        bias,
+        residual,
+        x1,
+        weight1,
+        bias1,
+        eps,
+        dropout_p,
+        rowscale,
+        prenorm,
+        residual_in_fp32,
+        zero_centered_weight,
+        True,
+        return_dropout_mask,
+        out_dtype,
+        out,
+        residual_out,
+    )
+
+
+@triton.jit
+def _norm_infer_kernel(
+    X,
+    Y,
+    W,
+    B,
+    stride_x_row,
+    stride_y_row,
+    M,
+    N,
+    eps,
+    IS_RMS_NORM: tl.constexpr,
+    HAS_WEIGHT: tl.constexpr,
+    HAS_BIAS: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    row = tl.program_id(0)
+    X += row * stride_x_row
+    Y += row * stride_y_row
+    if HAS_WEIGHT:
+        W += 0
+    if HAS_BIAS:
+        B += 0
+    cols = tl.arange(0, BLOCK_N)
+    x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32)
+    if not IS_RMS_NORM:
+        mean = tl.sum(x, axis=0) / N
+        xbar = tl.where(cols < N, x - mean, 0.0)
+        var = tl.sum(xbar * xbar, axis=0) / N
+    else:
+        xbar = tl.where(cols < N, x, 0.0)
+        var = tl.sum(xbar * xbar, axis=0) / N
+    rstd = 1 / tl.sqrt(var + eps)
+    x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd
+    if HAS_WEIGHT:
+        w = tl.load(W + cols, mask=cols < N, other=1.0).to(tl.float32)
+        y = x_hat * w
+    else:
+        y = x_hat
+    if HAS_BIAS:
+        b = tl.load(B + cols, mask=cols < N, other=0.0).to(tl.float32)
+        y += b
+    tl.store(Y + cols, y, mask=cols < N)
+
+
+def norm_infer(
+    x: Tensor,
+    weight: Optional[Tensor],
+    bias: Optional[Tensor],
+    eps: float,
+    is_rms_norm: bool = False,
+    out: Optional[Tensor] = None,
+):
+    M, N = x.shape
+    x = x.contiguous()
+    if weight is not None:
+        assert weight.shape == (N,)
+        assert weight.stride(-1) == 1
+    if bias is not None:
+        assert bias.shape == (N,)
+        assert bias.stride(-1) == 1
+    if out is None:
+        out = torch.empty_like(x)
+    MAX_FUSED_SIZE = 65536 // x.element_size()
+    BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N))
+    if N > BLOCK_N:
+        raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
+    num_warps = min(max(BLOCK_N // 256, 1), 8)
+    _norm_infer_kernel[(M,)](
+        x,
+        out,
+        weight if weight is not None else x,  # dummy when HAS_WEIGHT=False
+        bias if bias is not None else x,  # dummy when HAS_BIAS=False
+        x.stride(0),
+        out.stride(0),
+        M,
+        N,
+        eps,
+        IS_RMS_NORM=is_rms_norm,
+        HAS_WEIGHT=weight is not None,
+        HAS_BIAS=bias is not None,
+        BLOCK_N=BLOCK_N,
+        num_warps=num_warps,
+    )
+    return out
+
+
+if current_platform.is_mps():
+    from .mps_fallback import norm_infer_native, rms_norm_fn_native
+
+    norm_infer = norm_infer_native
+    rms_norm_fn = rms_norm_fn_native
+
+if current_platform.is_cpu():
+    from .torch_fallback import norm_infer_native, rms_norm_fn_native
+
+    norm_infer = norm_infer_native
+    rms_norm_fn = rms_norm_fn_native

python/sglang/jit_kernel/diffusion/triton/npu_fallback.py

@@ -0,0 +1,50 @@
+import torch
+import torch_npu
+
+NPU_ROTARY_MUL_MAX_NUM_HEADS = 1000
+NPU_ROTARY_MUL_MAX_HEAD_SIZE = 896
+
+
+# TODO: remove this when triton ascend bug is fixed
+def fuse_scale_shift_native(
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    block_l: int = 128,
+    block_c: int = 128,
+):
+    return x * (1 + scale) + shift
+
+
+# TODO: remove this when triton ascend bug is fixed
+def apply_rotary_embedding_native(
+    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, interleaved: bool = False
+) -> torch.Tensor:
+    if interleaved and cos.shape[-1] == x.shape[-1]:
+        cos = cos[..., ::2]
+        sin = sin[..., ::2]
+    cos = cos.unsqueeze(-2).to(x.dtype)
+    sin = sin.unsqueeze(-2).to(x.dtype)
+
+    if (
+        cos.dim() == 3
+        and x.dim() == 3
+        and x.shape[1] < NPU_ROTARY_MUL_MAX_NUM_HEADS
+        and x.shape[2] < NPU_ROTARY_MUL_MAX_HEAD_SIZE
+        and not interleaved
+    ):
+        if cos.size(-1) * 2 == x.size(-1):
+            cos = torch.cat([cos, cos], dim=-1)
+            sin = torch.cat([sin, sin], dim=-1)
+        cos = cos.unsqueeze(0)
+        sin = sin.unsqueeze(0)
+        x = x.unsqueeze(0)
+        x_embed = torch_npu.npu_rotary_mul(x, cos, sin)
+        x_embed = x_embed.squeeze(0)
+        return x_embed
+
+    x1 = x[..., ::2]
+    x2 = x[..., 1::2]
+    o1 = x1 * cos - x2 * sin
+    o2 = x2 * cos + x1 * sin
+    return torch.stack((o1, o2), dim=-1).flatten(-2)

python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py

@@ -0,0 +1,83 @@
+import torch
+import triton  # type: ignore
+import triton.language as tl  # type: ignore
+
+from sglang.kernel_api_logging import debug_kernel_api
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.srt.utils.custom_op import register_custom_op
+
+
+# Adapted from https://github.com/ModelTC/LightX2V/blob/main/lightx2v/common/ops/norm/triton_ops.py#L905-L956
+@triton.jit
+def _rms_norm_tiled_onepass(
+    y_ptr,
+    x_ptr,
+    w_ptr,
+    SEQ: tl.constexpr,
+    DIM: tl.constexpr,
+    EPS: tl.constexpr,
+    BLOCK_SIZE_SEQ: tl.constexpr,
+    BLOCK_SIZE_DIM: tl.constexpr,
+):
+    seq_blk_id = tl.program_id(0)
+    seq_id = seq_blk_id * BLOCK_SIZE_SEQ
+
+    seq_offset = seq_id + tl.arange(0, BLOCK_SIZE_SEQ)[:, None]
+    s_mask = seq_offset < SEQ
+    d_offset = tl.arange(0, BLOCK_SIZE_DIM)[None, :]
+    d_mask = d_offset < DIM
+    y_blk = y_ptr + seq_offset * DIM + d_offset
+    x_blk = x_ptr + seq_offset * DIM + d_offset
+    mask = s_mask & d_mask
+
+    x = tl.load(x_blk, mask=mask, other=0.0).to(tl.float32)
+    mean_square = tl.sum(x * x, axis=1, keep_dims=True) / DIM
+    rstd = tl.math.rsqrt(mean_square + EPS)
+    w = tl.load(w_ptr + d_offset, mask=d_mask)
+    tl.store(y_blk, x * rstd * w, mask=mask)
+
+
+@register_custom_op(op_name="triton_one_pass_rms_norm_cuda", out_shape="x")
+def _triton_one_pass_rms_norm_cuda(
+    x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6
+) -> torch.Tensor:
+    shape = x.shape
+    x = x.contiguous()
+    y = torch.empty_like(x)
+    x_view = x.reshape(-1, shape[-1])
+    y_view = y.reshape(-1, shape[-1])
+    S, D = x_view.shape
+
+    block_size_seq = min(16, triton.next_power_of_2(max(1, S // 512)))
+    grid = (triton.cdiv(S, block_size_seq),)
+
+    with torch.get_device_module().device(x.device):
+        _rms_norm_tiled_onepass[grid](
+            y_view,
+            x_view,
+            w,
+            S,
+            D,
+            eps,
+            BLOCK_SIZE_DIM=triton.next_power_of_2(D),
+            BLOCK_SIZE_SEQ=block_size_seq,
+        )
+    return y
+
+
+def triton_one_pass_rms_norm(x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6):
+    return _triton_one_pass_rms_norm_cuda(x, w, eps)
+
+
+if current_platform.is_mps():
+    from .mps_fallback import triton_one_pass_rms_norm_native
+
+    @debug_kernel_api
+    def triton_one_pass_rms_norm(x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6):
+        return triton_one_pass_rms_norm_native(x, w, eps)
+
+
+if current_platform.is_cpu():
+    from .torch_fallback import triton_one_pass_rms_norm_native
+
+    triton_one_pass_rms_norm = triton_one_pass_rms_norm_native

python/sglang/jit_kernel/diffusion/triton/rotary.py

@@ -0,0 +1,141 @@
+import torch
+import triton  # type: ignore
+import triton.language as tl  # type: ignore
+
+from sglang.multimodal_gen.runtime.platforms import current_platform
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({"BLOCK_HEADS": 1, "BLOCK_HS_HALF": 32}, num_warps=2),
+        triton.Config({"BLOCK_HEADS": 2, "BLOCK_HS_HALF": 32}, num_warps=2),
+        triton.Config({"BLOCK_HEADS": 4, "BLOCK_HS_HALF": 32}, num_warps=4),
+        triton.Config({"BLOCK_HEADS": 4, "BLOCK_HS_HALF": 64}, num_warps=4),
+        triton.Config({"BLOCK_HEADS": 8, "BLOCK_HS_HALF": 64}, num_warps=8),
+    ],
+    key=["num_heads", "head_size"],
+)
+@triton.jit
+def _rotary_embedding_kernel(
+    output_ptr,
+    x_ptr,
+    cos_ptr,
+    sin_ptr,
+    num_heads,
+    head_size,
+    num_tokens,
+    stride_out_bt,
+    stride_out_head,
+    stride_x_bt,
+    stride_x_head,
+    stride_cos_row,
+    stride_sin_row,
+    BLOCK_HEADS: tl.constexpr,
+    BLOCK_HS_HALF: tl.constexpr,
+):
+    bt_idx = tl.program_id(0)
+    head_block_idx = tl.program_id(1)
+    token_idx = bt_idx % num_tokens
+
+    cos_row_ptr = cos_ptr + token_idx * stride_cos_row
+    sin_row_ptr = sin_ptr + token_idx * stride_sin_row
+    head_offsets = head_block_idx * BLOCK_HEADS + tl.arange(0, BLOCK_HEADS)
+    head_mask = head_offsets < num_heads
+
+    head_size_half = head_size // 2
+    x_row_ptrs = x_ptr + bt_idx * stride_x_bt + head_offsets[:, None] * stride_x_head
+    output_row_ptrs = (
+        output_ptr + bt_idx * stride_out_bt + head_offsets[:, None] * stride_out_head
+    )
+
+    for block_start in range(0, head_size_half, BLOCK_HS_HALF):
+        offsets_half = block_start + tl.arange(0, BLOCK_HS_HALF)
+        half_mask = offsets_half < head_size_half
+        mask = head_mask[:, None] & half_mask[None, :]
+
+        cos_vals = tl.load(cos_row_ptr + offsets_half, mask=half_mask, other=0.0)
+        sin_vals = tl.load(sin_row_ptr + offsets_half, mask=half_mask, other=0.0)
+
+        offsets_x1 = 2 * offsets_half
+        offsets_x2 = 2 * offsets_half + 1
+
+        x1_vals = tl.load(x_row_ptrs + offsets_x1[None, :], mask=mask, other=0.0)
+        x2_vals = tl.load(x_row_ptrs + offsets_x2[None, :], mask=mask, other=0.0)
+
+        x1_fp32 = x1_vals.to(tl.float32)
+        x2_fp32 = x2_vals.to(tl.float32)
+        cos_fp32 = cos_vals.to(tl.float32)[None, :]
+        sin_fp32 = sin_vals.to(tl.float32)[None, :]
+        o1_vals = tl.fma(-x2_fp32, sin_fp32, x1_fp32 * cos_fp32)
+        o2_vals = tl.fma(x1_fp32, sin_fp32, x2_fp32 * cos_fp32)
+
+        tl.store(
+            output_row_ptrs + offsets_x1[None, :],
+            o1_vals.to(x1_vals.dtype),
+            mask=mask,
+        )
+        tl.store(
+            output_row_ptrs + offsets_x2[None, :],
+            o2_vals.to(x2_vals.dtype),
+            mask=mask,
+        )
+
+
+def apply_rotary_embedding(
+    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, interleaved: bool = False
+) -> torch.Tensor:
+    output = torch.empty_like(x)
+
+    if x.dim() > 3:
+        bsz, num_tokens, num_heads, head_size = x.shape
+    else:
+        num_tokens, num_heads, head_size = x.shape
+        bsz = 1
+
+    assert head_size % 2 == 0, "head_size must be divisible by 2"
+
+    x_reshaped = x.view(bsz * num_tokens, num_heads, head_size)
+    output_reshaped = output.view(bsz * num_tokens, num_heads, head_size)
+
+    if interleaved and cos.shape[-1] == head_size:
+        cos = cos[..., ::2].contiguous()
+        sin = sin[..., ::2].contiguous()
+    else:
+        cos = cos.contiguous()
+        sin = sin.contiguous()
+
+    _rotary_embedding_kernel[
+        lambda META: (bsz * num_tokens, triton.cdiv(num_heads, META["BLOCK_HEADS"]))
+    ](
+        output_reshaped,
+        x_reshaped,
+        cos,
+        sin,
+        num_heads,
+        head_size,
+        num_tokens,
+        output_reshaped.stride(0),
+        output_reshaped.stride(1),
+        x_reshaped.stride(0),
+        x_reshaped.stride(1),
+        cos.stride(0),
+        sin.stride(0),
+    )
+
+    return output
+
+
+if current_platform.is_npu():
+    from .npu_fallback import apply_rotary_embedding_native
+
+    apply_rotary_embedding = apply_rotary_embedding_native
+
+if current_platform.is_mps():
+    from .mps_fallback import apply_rotary_embedding_native
+
+    apply_rotary_embedding = apply_rotary_embedding_native
+
+if current_platform.is_cpu():
+    from .torch_fallback import apply_rotary_embedding_native
+
+    apply_rotary_embedding = apply_rotary_embedding_native

python/sglang/jit_kernel/diffusion/triton/sana_wm_gdn.py

@@ -0,0 +1,241 @@
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# 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.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Inference-side helpers for the bidirectional fused GDN path.
+
+Precision knob: env var ``FUSED_GDN_PRECISION`` or ``PRECISION_OVERRIDE``:
+  0=IEEE fp32 dots, 1=TF32, 2=bf16 TC + fp32 state [default], 3=bf16 TC + bf16 state.
+"""
+
+# ruff: noqa: E501
+
+from __future__ import annotations
+
+import os
+
+import torch
+import triton
+import triton.language as tl
+
+# =====================================================================
+#  GPU-adaptive kernel config
+# =====================================================================
+
+
+def _get_kernel_config() -> dict:
+    """Return optimal kernel parameters for the current GPU.
+
+    STATE_FP32 (fp32 state_prev) needs ~128KB SRAM (H100 228KB), vs ~96KB for
+    bf16 state_prev (fits GB10's 101KB).
+    """
+    if not torch.cuda.is_available():
+        return {"BLOCK_S": 64, "num_stages": 1, "num_warps": 4, "STATE_FP32": False}
+    smem = torch.cuda.get_device_properties(0).shared_memory_per_multiprocessor
+    state_fp32 = smem >= 150 * 1024  # H100 (228KB) yes, GB10 (101KB) no
+    return {"BLOCK_S": 64, "num_stages": 1, "num_warps": 8, "STATE_FP32": state_fp32}
+
+
+_KCFG = None
+
+
+def _kcfg():
+    global _KCFG
+    if _KCFG is None:
+        _KCFG = _get_kernel_config()
+    return _KCFG
+
+
+# precision=0 → IEEE fp32 dots + fp32 state  (DOT_PRECISION=2, STATE_FP32=1)
+# precision=1 → TF32  dots   + fp32 state    (DOT_PRECISION=1, STATE_FP32=1)
+# precision=2 → bf16  dots   + fp32 state    (DOT_PRECISION=0, STATE_FP32=1) [default]
+# precision=3 → bf16  dots   + bf16 state    (DOT_PRECISION=0, STATE_FP32=0)
+def _precision_params(precision: int) -> tuple:
+    if precision == 0:
+        return 2, True
+    elif precision == 1:
+        return 1, True
+    elif precision == 3:
+        return 0, False
+    else:  # default
+        return 0, True
+
+
+_env_prec = os.environ.get("FUSED_GDN_PRECISION", None)
+PRECISION_OVERRIDE: int | None = int(_env_prec) if _env_prec is not None else None
+
+
+def _resolve_launch_config() -> tuple:
+    """Returns (prec, dot_prec, state_fp32, num_warps).
+
+    Uses ``PRECISION_OVERRIDE`` when set, else ``_kcfg()`` (per-GPU SRAM).
+    num_warps clamped to 4 when dots run on fp32 operands (more registers).
+    """
+    cfg = _kcfg()
+    prec = PRECISION_OVERRIDE if PRECISION_OVERRIDE is not None else 2
+    dot_prec, state_fp32 = _precision_params(prec)
+    if PRECISION_OVERRIDE is None:
+        state_fp32 = cfg["STATE_FP32"]
+    nw = cfg["num_warps"]
+    if dot_prec >= 1:
+        nw = min(nw, 4)
+    return prec, dot_prec, state_fp32, nw
+
+
+def prepare_rope_tables(
+    rotary_emb, N: int, D: int, device
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Complex rotary_emb `(1, 1, N, D//2)` → expanded (N, D) cos/sin tables.
+
+    Encodes the interleaved-pair rotation
+        y[2i]   = x[2i]*cos[i] - x[2i+1]*sin[i]
+        y[2i+1] = x[2i]*sin[i] + x[2i+1]*cos[i]
+    as  y[d] = x[d]*cos_exp[d] + x[d^1]*sin_exp[d]
+    where sin_exp[2i] = -sin[i], sin_exp[2i+1] = +sin[i].
+
+    Returns (cos_exp, sin_exp) both (N, D) float32, contiguous.
+    """
+    if rotary_emb is None:
+        return (
+            torch.ones(N, D, device=device, dtype=torch.float32),
+            torch.zeros(N, D, device=device, dtype=torch.float32),
+        )
+    freqs = rotary_emb.squeeze(0).squeeze(0)  # (N, D//2) complex
+    cos_half = freqs.real.float()
+    sin_half = freqs.imag.float()
+    rope_cos = cos_half.repeat_interleave(2, dim=-1)
+    rope_sin = torch.stack([-sin_half, sin_half], dim=-1).reshape(N, D)
+    return rope_cos.contiguous(), rope_sin.contiguous()
+
+
+# =====================================================================
+#  Fused single-pass Q+K inverse-RMS Triton kernel
+# =====================================================================
+# Single Triton launch that reads each `(b, n)` row of `qkv` once and emits
+# both `q_inv_rms[b, n]` and `k_inv_rms[b, n]`. Replaces two separate PyTorch
+# scans (cast→square→sum→rsqrt) over `qkv[:, :, 0]` and `qkv[:, :, 1]`.
+#
+# Layout assumed: `qkv` is (B, N, 3, H, D) contiguous, so the C = H*D channels
+# for a given (b, n, qkv_idx) live in a contiguous memory span.
+
+
+@triton.jit
+def _fused_qk_inv_rms_kernel(
+    qkv_ptr,  # *T_in     (B, N, 3, H, D), contiguous
+    q_inv_rms_ptr,  # *float32  (B, N)
+    k_inv_rms_ptr,  # *float32  (B, N)
+    N: tl.constexpr,
+    C: tl.constexpr,  # H * D
+    eps,
+    BLOCK_C: tl.constexpr,
+):
+    bn_id = tl.program_id(0)
+    qkv_row_stride = 3 * C
+    row_base = bn_id * qkv_row_stride
+    q_base = row_base
+    k_base = row_base + C
+
+    offs = tl.arange(0, BLOCK_C)
+    mask = offs < C
+
+    q_vals = tl.load(qkv_ptr + q_base + offs, mask=mask, other=0.0).to(tl.float32)
+    k_vals = tl.load(qkv_ptr + k_base + offs, mask=mask, other=0.0).to(tl.float32)
+
+    q_sq = tl.sum(q_vals * q_vals, axis=0)
+    k_sq = tl.sum(k_vals * k_vals, axis=0)
+
+    inv_c = 1.0 / C
+    q_inv = tl.rsqrt(q_sq * inv_c + eps)
+    k_inv = tl.rsqrt(k_sq * inv_c + eps)
+
+    tl.store(q_inv_rms_ptr + bn_id, q_inv)
+    tl.store(k_inv_rms_ptr + bn_id, k_inv)
+
+
+def fused_qk_inv_rms(
+    qkv: torch.Tensor,
+    eps: float = 1e-5,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Single-pass Triton fused Q+K inverse-RMS.
+
+    Replaces two separate PyTorch RMS scans with one launch that reads each
+    ``(b, n)`` row of ``qkv`` exactly once.
+    qkv: (B, N, 3, H, D) contiguous. Returns (q_inv_rms, k_inv_rms), each (B, N) float32.
+    """
+    assert qkv.is_contiguous(), "qkv must be contiguous (B, N, 3, H, D)"
+    assert (
+        qkv.dim() == 5 and qkv.shape[2] == 3
+    ), f"expected (B, N, 3, H, D), got {tuple(qkv.shape)}"
+    B, N, _, H, D = qkv.shape
+    C = H * D
+    q_inv_rms = torch.empty((B, N), dtype=torch.float32, device=qkv.device)
+    k_inv_rms = torch.empty((B, N), dtype=torch.float32, device=qkv.device)
+    BLOCK_C = triton.next_power_of_2(C)
+    _fused_qk_inv_rms_kernel[(B * N,)](
+        qkv,
+        q_inv_rms,
+        k_inv_rms,
+        N=N,
+        C=C,
+        eps=eps,
+        BLOCK_C=BLOCK_C,
+    )
+    return q_inv_rms, k_inv_rms
+
+
+# =====================================================================
+#  Bidirectional GDN entry point (delegates to chunkwise)
+# =====================================================================
+
+
+def fused_bigdn_func(
+    qkv: torch.Tensor,  # (B, N, 3, H, D)
+    q_inv_rms: torch.Tensor,  # (B, N) float32
+    k_inv_rms: torch.Tensor,  # (B, N) float32
+    q_norm_weight: torch.Tensor,  # (C,) float32
+    k_norm_weight: torch.Tensor,  # (C,) float32
+    rope_cos: torch.Tensor,  # (N, D) float32
+    rope_sin: torch.Tensor,  # (N, D) float32
+    beta: torch.Tensor,  # (B, H, F, S)
+    decay: torch.Tensor,  # (B, H, F)
+    F: int,
+    S: int,
+    k_scale: float,
+    eps: float = 1e-6,
+) -> torch.Tensor:
+    """Bidirectional fused GDN. Returns ``(B, N, H, D)``.
+
+    Thin entry point kept for call-site stability; delegates to
+    :func:`fused_bigdn_bidi_chunkwise` from ``sana_wm_gdn_chunkwise``.
+    """
+    from sglang.jit_kernel.diffusion.triton.sana_wm_gdn_chunkwise import (
+        fused_bigdn_bidi_chunkwise,
+    )
+
+    return fused_bigdn_bidi_chunkwise(
+        qkv,
+        q_inv_rms,
+        k_inv_rms,
+        q_norm_weight,
+        k_norm_weight,
+        rope_cos,
+        rope_sin,
+        beta,
+        decay,
+        F=F,
+        S=S,
+        k_scale=k_scale,
+        eps=eps,
+    )

python/sglang/jit_kernel/diffusion/triton/scale_shift.py

@@ -0,0 +1,691 @@
+import torch
+import triton  # type: ignore
+import triton.language as tl  # type: ignore
+
+from sglang.multimodal_gen.runtime.platforms import current_platform
+
+
+@triton.jit
+def _fused_layernorm_scale_shift_gate_select01_kernel(
+    output_ptr,
+    gate_out_ptr,
+    x_ptr,
+    weight_ptr,
+    bias_ptr,
+    scale0_ptr,
+    shift0_ptr,
+    gate0_ptr,
+    scale1_ptr,
+    shift1_ptr,
+    gate1_ptr,
+    index_ptr,
+    inner_dim,
+    seq_len,
+    stride_x_row,
+    stride_out_row,
+    stride_go_row,
+    stride_w,
+    stride_b,
+    stride_s0_b,
+    stride_s0_c,
+    stride_sh0_b,
+    stride_sh0_c,
+    stride_g0_b,
+    stride_g0_c,
+    stride_s1_b,
+    stride_s1_c,
+    stride_sh1_b,
+    stride_sh1_c,
+    stride_g1_b,
+    stride_g1_c,
+    stride_i_b,
+    stride_i_l,
+    eps,
+    HAS_WEIGHT: tl.constexpr,
+    HAS_BIAS: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    row = tl.program_id(0)
+    cols = tl.arange(0, BLOCK_N)
+    mask = cols < inner_dim
+
+    x_row_ptr = x_ptr + row * stride_x_row
+    out_row_ptr = output_ptr + row * stride_out_row
+    gate_row_ptr = gate_out_ptr + row * stride_go_row
+
+    x = tl.load(x_row_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+    mean = tl.sum(x, axis=0) / inner_dim
+    xbar = tl.where(mask, x - mean, 0.0)
+    var = tl.sum(xbar * xbar, axis=0) / inner_dim
+    rstd = tl.rsqrt(var + eps)
+    x_hat = (x - mean) * rstd
+
+    if HAS_WEIGHT:
+        w = tl.load(weight_ptr + cols * stride_w, mask=mask, other=1.0).to(tl.float32)
+        x_hat = x_hat * w
+    if HAS_BIAS:
+        b = tl.load(bias_ptr + cols * stride_b, mask=mask, other=0.0).to(tl.float32)
+        x_hat = x_hat + b
+
+    batch_idx = row // seq_len
+    seq_idx = row % seq_len
+    idx = tl.load(index_ptr + batch_idx * stride_i_b + seq_idx * stride_i_l).to(tl.int1)
+
+    scale0_ptrs = scale0_ptr + batch_idx * stride_s0_b + cols * stride_s0_c
+    shift0_ptrs = shift0_ptr + batch_idx * stride_sh0_b + cols * stride_sh0_c
+    gate0_ptrs = gate0_ptr + batch_idx * stride_g0_b + cols * stride_g0_c
+
+    scale1_ptrs = scale1_ptr + batch_idx * stride_s1_b + cols * stride_s1_c
+    shift1_ptrs = shift1_ptr + batch_idx * stride_sh1_b + cols * stride_sh1_c
+    gate1_ptrs = gate1_ptr + batch_idx * stride_g1_b + cols * stride_g1_c
+
+    # Branch on scalar idx instead of using tl.where on pointers.
+    # tl.where on pointers triggers an assertion in AMD Triton's
+    # CanonicalizePointers pass (ConvertArithSelectOp) on gfx950.
+    # This keeps it at 3 loads (not 6), avoids the pointer-level
+    # tl.where entirely, and since idx is uniform across all threads
+    # the branch has no divergence cost.
+    if idx:
+        scale = tl.load(scale1_ptrs, mask=mask, other=0.0).to(tl.float32)
+        shift = tl.load(shift1_ptrs, mask=mask, other=0.0).to(tl.float32)
+        gate = tl.load(gate1_ptrs, mask=mask, other=0.0)
+    else:
+        scale = tl.load(scale0_ptrs, mask=mask, other=0.0).to(tl.float32)
+        shift = tl.load(shift0_ptrs, mask=mask, other=0.0).to(tl.float32)
+        gate = tl.load(gate0_ptrs, mask=mask, other=0.0)
+    y = x_hat * (1.0 + scale) + shift
+
+    tl.store(out_row_ptr + cols, y, mask=mask)
+    tl.store(gate_row_ptr + cols, gate, mask=mask)
+
+
+@triton.jit
+def _fused_residual_layernorm_scale_shift_gate_select01_kernel(
+    output_ptr,
+    residual_out_ptr,
+    gate_out_ptr,
+    x_ptr,
+    residual_ptr,
+    residual_gate_ptr,
+    weight_ptr,
+    bias_ptr,
+    scale0_ptr,
+    shift0_ptr,
+    gate0_ptr,
+    scale1_ptr,
+    shift1_ptr,
+    gate1_ptr,
+    index_ptr,
+    inner_dim,
+    seq_len,
+    stride_x_row,
+    stride_res_row,
+    stride_rg_row,
+    stride_out_row,
+    stride_res_out_row,
+    stride_go_row,
+    stride_w,
+    stride_b,
+    stride_s0_b,
+    stride_s0_c,
+    stride_sh0_b,
+    stride_sh0_c,
+    stride_g0_b,
+    stride_g0_c,
+    stride_s1_b,
+    stride_s1_c,
+    stride_sh1_b,
+    stride_sh1_c,
+    stride_g1_b,
+    stride_g1_c,
+    stride_i_b,
+    stride_i_l,
+    eps,
+    HAS_WEIGHT: tl.constexpr,
+    HAS_BIAS: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    row = tl.program_id(0)
+    cols = tl.arange(0, BLOCK_N)
+    mask = cols < inner_dim
+
+    x_row_ptr = x_ptr + row * stride_x_row
+    res_row_ptr = residual_ptr + row * stride_res_row
+    rg_row_ptr = residual_gate_ptr + row * stride_rg_row
+    out_row_ptr = output_ptr + row * stride_out_row
+    res_out_row_ptr = residual_out_ptr + row * stride_res_out_row
+    gate_row_ptr = gate_out_ptr + row * stride_go_row
+
+    x = tl.load(x_row_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+    residual = tl.load(res_row_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+    residual_gate = tl.load(rg_row_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+    residual_out = residual + residual_gate * x
+    tl.store(res_out_row_ptr + cols, residual_out, mask=mask)
+
+    mean = tl.sum(residual_out, axis=0) / inner_dim
+    xbar = tl.where(mask, residual_out - mean, 0.0)
+    var = tl.sum(xbar * xbar, axis=0) / inner_dim
+    rstd = tl.rsqrt(var + eps)
+    x_hat = (residual_out - mean) * rstd
+
+    if HAS_WEIGHT:
+        w = tl.load(weight_ptr + cols * stride_w, mask=mask, other=1.0).to(tl.float32)
+        x_hat = x_hat * w
+    if HAS_BIAS:
+        b = tl.load(bias_ptr + cols * stride_b, mask=mask, other=0.0).to(tl.float32)
+        x_hat = x_hat + b
+
+    batch_idx = row // seq_len
+    seq_idx = row % seq_len
+    idx = tl.load(index_ptr + batch_idx * stride_i_b + seq_idx * stride_i_l).to(tl.int1)
+
+    scale0_ptrs = scale0_ptr + batch_idx * stride_s0_b + cols * stride_s0_c
+    shift0_ptrs = shift0_ptr + batch_idx * stride_sh0_b + cols * stride_sh0_c
+    gate0_ptrs = gate0_ptr + batch_idx * stride_g0_b + cols * stride_g0_c
+
+    scale1_ptrs = scale1_ptr + batch_idx * stride_s1_b + cols * stride_s1_c
+    shift1_ptrs = shift1_ptr + batch_idx * stride_sh1_b + cols * stride_sh1_c
+    gate1_ptrs = gate1_ptr + batch_idx * stride_g1_b + cols * stride_g1_c
+
+    # Branch on scalar idx instead of using tl.where on pointers.
+    # tl.where on pointers triggers an assertion in AMD Triton's
+    # CanonicalizePointers pass (ConvertArithSelectOp) on gfx950.
+    # This keeps it at 3 loads (not 6), avoids the pointer-level
+    # tl.where entirely, and since idx is uniform across all threads
+    # the branch has no divergence cost.
+    if idx:
+        scale = tl.load(scale1_ptrs, mask=mask, other=0.0).to(tl.float32)
+        shift = tl.load(shift1_ptrs, mask=mask, other=0.0).to(tl.float32)
+        gate = tl.load(gate1_ptrs, mask=mask, other=0.0)
+    else:
+        scale = tl.load(scale0_ptrs, mask=mask, other=0.0).to(tl.float32)
+        shift = tl.load(shift0_ptrs, mask=mask, other=0.0).to(tl.float32)
+        gate = tl.load(gate0_ptrs, mask=mask, other=0.0)
+    y = x_hat * (1.0 + scale) + shift
+
+    tl.store(out_row_ptr + cols, y, mask=mask)
+    tl.store(gate_row_ptr + cols, gate, mask=mask)
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({"BLOCK_N": 64}, num_warps=2),
+        triton.Config({"BLOCK_N": 128}, num_warps=4),
+        triton.Config({"BLOCK_N": 256}, num_warps=4),
+        triton.Config({"BLOCK_N": 512}, num_warps=4),
+        triton.Config({"BLOCK_N": 1024}, num_warps=8),
+    ],
+    key=["inner_dim"],
+)
+@triton.jit
+def _fused_scale_shift_4d_kernel(
+    output_ptr,
+    normalized_ptr,
+    scale_ptr,
+    shift_ptr,
+    scale_constant: tl.constexpr,  # scale_constant is either 0 or 1.
+    inner_dim,
+    seq_len,
+    num_frames,
+    frame_seqlen,
+    BLOCK_N: tl.constexpr,
+):
+    pid_row = tl.program_id(0)
+    pid_col = tl.program_id(1)
+
+    col_offsets = pid_col * BLOCK_N + tl.arange(0, BLOCK_N)
+    mask = col_offsets < inner_dim
+
+    # Pointers for normalized and output
+    row_base = pid_row * inner_dim
+    norm_ptrs = normalized_ptr + row_base + col_offsets
+    out_ptrs = output_ptr + row_base + col_offsets
+
+    # Pointers for scale (per-frame) and shift (per-token)
+    b_idx = pid_row // seq_len
+    t_idx = pid_row % seq_len
+    frame_idx_in_batch = t_idx // frame_seqlen
+
+    scale_row_idx = b_idx * num_frames + frame_idx_in_batch
+    scale_ptrs = scale_ptr + scale_row_idx * inner_dim + col_offsets
+    # shift is per-token [B*L, C], indexed by pid_row directly
+    shift_ptrs = shift_ptr + pid_row * inner_dim + col_offsets
+
+    normalized = tl.load(norm_ptrs, mask=mask, other=0.0)
+    scale = tl.load(scale_ptrs, mask=mask, other=0.0)
+    shift = tl.load(shift_ptrs, mask=mask, other=0.0)
+
+    scale_const_tensor = tl.full([BLOCK_N], scale_constant, dtype=scale.dtype)
+    output = normalized * (scale_const_tensor + scale) + shift
+
+    tl.store(out_ptrs, output, mask=mask)
+
+
+@triton.jit
+def fuse_scale_shift_kernel_blc_opt(
+    x_ptr,
+    shift_ptr,
+    scale_ptr,
+    scale_constant: tl.constexpr,  # scale_constant is either 0 or 1.,
+    y_ptr,
+    B,
+    L,
+    C,
+    stride_x_b,
+    stride_x_l,
+    stride_x_c,
+    stride_s_b,
+    stride_s_l,
+    stride_s_c,
+    stride_sc_b,
+    stride_sc_l,
+    stride_sc_c,
+    SCALE_IS_SCALAR: tl.constexpr,
+    SHIFT_IS_SCALAR: tl.constexpr,
+    BLOCK_L: tl.constexpr,
+    BLOCK_C: tl.constexpr,
+):
+    pid_l = tl.program_id(0)
+    pid_c = tl.program_id(1)
+    pid_b = tl.program_id(2)
+
+    l_offsets = pid_l * BLOCK_L + tl.arange(0, BLOCK_L)
+    c_offsets = pid_c * BLOCK_C + tl.arange(0, BLOCK_C)
+
+    mask_l = l_offsets < L
+    mask_c = c_offsets < C
+    mask = mask_l[:, None] & mask_c[None, :]
+
+    x_off = (
+        pid_b * stride_x_b
+        + l_offsets[:, None] * stride_x_l
+        + c_offsets[None, :] * stride_x_c
+    )
+    x = tl.load(x_ptr + x_off, mask=mask, other=0)
+
+    if SHIFT_IS_SCALAR:
+        shift_val = tl.load(shift_ptr)
+        shift = tl.full((BLOCK_L, BLOCK_C), shift_val, dtype=shift_val.dtype)
+    else:
+        s_off = (
+            pid_b * stride_s_b
+            + l_offsets[:, None] * stride_s_l
+            + c_offsets[None, :] * stride_s_c
+        )
+        shift = tl.load(shift_ptr + s_off, mask=mask, other=0)
+
+    if SCALE_IS_SCALAR:
+        scale_val = tl.load(scale_ptr)
+        scale = tl.full((BLOCK_L, BLOCK_C), scale_val, dtype=scale_val.dtype)
+    else:
+        sc_off = (
+            pid_b * stride_sc_b
+            + l_offsets[:, None] * stride_sc_l
+            + c_offsets[None, :] * stride_sc_c
+        )
+        scale = tl.load(scale_ptr + sc_off, mask=mask, other=0)
+
+    y = x * (scale_constant + scale) + shift
+    tl.store(y_ptr + x_off, y, mask=mask)
+
+
+def fuse_scale_shift_kernel(
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    scale_constant: float = 1.0,
+    block_l: int = 128,
+    block_c: int = 128,
+):
+    assert (x.is_cuda and scale.is_cuda) or (x.is_xpu and scale.is_xpu)
+    assert x.is_contiguous()
+
+    B, L, C = x.shape
+    output = torch.empty_like(x)
+    if x.numel() == 0:
+        return output
+
+    if scale.dim() == 4:
+        # scale/shift: [B, F, 1, C]
+        rows = B * L
+        x_2d = x.view(rows, C)
+        output_2d = output.view(rows, C)
+
+        def grid(meta):
+            return (rows, triton.cdiv(C, meta["BLOCK_N"]))
+
+        num_frames = scale.shape[1]
+        assert (
+            L % num_frames == 0
+        ), "seq_len must be divisible by num_frames for 4D scale/shift"
+        frame_seqlen = L // num_frames
+
+        # Compact scale [B, F, 1, C] -> [B*F, C] (per-frame)
+        scale_reshaped = scale.squeeze(2).reshape(-1, C).contiguous()
+        if shift.dim() == 4 and current_platform.is_hip():
+            # ROCm has no fused CUTLASS scale-shift kernel, so this native path
+            # handles the causal Wan / LingBot output AdaLN, which passes a
+            # per-frame shift [B, F, 1, C]. Broadcast it across each frame's
+            # tokens to per-token [B, L, C] before flattening to [B*L, C],
+            # matching the per-token indexing in _fused_scale_shift_4d_kernel
+            # (the CUDA fused path accepts [B, F, 1, C] shift and broadcasts it
+            # per-frame).
+            shift_reshaped = (
+                shift.expand(B, num_frames, frame_seqlen, C)
+                .reshape(rows, C)
+                .contiguous()
+            )
+        else:
+            # shift is per-token [B, L, C] -> [B*L, C]
+            shift_reshaped = shift.reshape(rows, C).contiguous()
+
+        _fused_scale_shift_4d_kernel[grid](
+            output_2d,
+            x_2d,
+            scale_reshaped,
+            shift_reshaped,
+            scale_constant,
+            C,
+            L,
+            num_frames,
+            frame_seqlen,
+        )
+    else:
+        # 2D: [B, C] or [1, C]  -> treat as [B, 1, C] and broadcast over L
+        # 3D: [B, L, C] (or broadcastable variants like [B, 1, C], [1, L, C], [1, 1, C])
+        # Also support scalar (0D or 1-element)
+        if scale.dim() == 0 or (scale.dim() == 1 and scale.numel() == 1):
+            scale_blc = scale.reshape(1)
+        elif scale.dim() == 2:
+            scale_blc = scale[:, None, :]
+        elif scale.dim() == 3:
+            scale_blc = scale
+        else:
+            raise ValueError("scale must be 0D/1D(1)/2D/3D or 4D")
+
+        if shift.dim() == 0 or (shift.dim() == 1 and shift.numel() == 1):
+            shift_blc = shift.reshape(1)
+        elif shift.dim() == 2:
+            shift_blc = shift[:, None, :]
+        elif shift.dim() == 3:
+            shift_blc = shift
+        else:
+            # broadcast later via expand if possible
+            shift_blc = shift
+
+        need_scale_scalar = scale_blc.dim() == 1 and scale_blc.numel() == 1
+        need_shift_scalar = shift_blc.dim() == 1 and shift_blc.numel() == 1
+
+        if not need_scale_scalar:
+            scale_exp = scale_blc.expand(B, L, C)
+            s_sb, s_sl, s_sc = scale_exp.stride()
+        else:
+            s_sb = s_sl = s_sc = 0
+
+        if not need_shift_scalar:
+            shift_exp = shift_blc.expand(B, L, C)
+            sh_sb, sh_sl, sh_sc = shift_exp.stride()
+        else:
+            sh_sb = sh_sl = sh_sc = 0
+
+        # If both scalars and both zero, copy fast-path
+        if need_scale_scalar and need_shift_scalar:
+            if not (
+                scale_blc.any().to("cpu", non_blocking=True)
+                or shift_blc.any().to("cpu", non_blocking=True)
+            ):
+                output.copy_(x)
+                return output
+
+        grid = (triton.cdiv(L, block_l), triton.cdiv(C, block_c), B)
+        fuse_scale_shift_kernel_blc_opt[grid](
+            x,
+            shift_blc if need_shift_scalar else shift_exp,
+            scale_blc if need_scale_scalar else scale_exp,
+            scale_constant,
+            output,
+            B,
+            L,
+            C,
+            x.stride(0),
+            x.stride(1),
+            x.stride(2),
+            sh_sb,
+            sh_sl,
+            sh_sc,
+            s_sb,
+            s_sl,
+            s_sc,
+            SCALE_IS_SCALAR=need_scale_scalar,
+            SHIFT_IS_SCALAR=need_shift_scalar,
+            BLOCK_L=block_l,
+            BLOCK_C=block_c,
+            num_warps=4,
+            num_stages=2,
+        )
+    return output
+
+
+def fuse_layernorm_scale_shift_gate_select01_kernel(
+    x: torch.Tensor,
+    weight: torch.Tensor | None,
+    bias: torch.Tensor | None,
+    scale0: torch.Tensor,
+    shift0: torch.Tensor,
+    gate0: torch.Tensor,
+    scale1: torch.Tensor,
+    shift1: torch.Tensor,
+    gate1: torch.Tensor,
+    index: torch.Tensor,
+    eps: float,
+):
+    assert x.is_cuda
+    assert x.is_contiguous()
+    B, L, C = x.shape
+    output = torch.empty_like(x)
+    gate_out = torch.empty_like(x)
+
+    if (
+        scale0.dim() != 2
+        or shift0.dim() != 2
+        or gate0.dim() != 2
+        or scale1.dim() != 2
+        or shift1.dim() != 2
+        or gate1.dim() != 2
+    ):
+        raise ValueError("scale0/shift0/gate0/scale1/shift1/gate1 must be 2D [B, C]")
+    if index.dim() != 2:
+        raise ValueError("index must be 2D [B, L]")
+    if weight is not None and (weight.dim() != 1 or weight.shape[0] != C):
+        raise ValueError("weight must be 1D [C]")
+    if bias is not None and (bias.dim() != 1 or bias.shape[0] != C):
+        raise ValueError("bias must be 1D [C]")
+
+    x_2d = x.view(B * L, C)
+    output_2d = output.view(B * L, C)
+    gate_out_2d = gate_out.view(B * L, C)
+    weight = weight.contiguous() if weight is not None else x_2d
+    bias = bias.contiguous() if bias is not None else x_2d
+
+    MAX_FUSED_SIZE = 65536 // x_2d.element_size()
+    BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(C))
+    if C > BLOCK_N:
+        raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
+    num_warps, num_stages = 4, 4
+
+    grid = (B * L,)
+    _fused_layernorm_scale_shift_gate_select01_kernel[grid](
+        output_2d,
+        gate_out_2d,
+        x_2d,
+        weight,
+        bias,
+        scale0.contiguous(),
+        shift0.contiguous(),
+        gate0.contiguous(),
+        scale1.contiguous(),
+        shift1.contiguous(),
+        gate1.contiguous(),
+        index.contiguous(),
+        C,
+        L,
+        x_2d.stride(0),
+        output_2d.stride(0),
+        gate_out_2d.stride(0),
+        weight.stride(0) if weight.dim() == 1 else 0,
+        bias.stride(0) if bias.dim() == 1 else 0,
+        scale0.stride(0),
+        scale0.stride(1),
+        shift0.stride(0),
+        shift0.stride(1),
+        gate0.stride(0),
+        gate0.stride(1),
+        scale1.stride(0),
+        scale1.stride(1),
+        shift1.stride(0),
+        shift1.stride(1),
+        gate1.stride(0),
+        gate1.stride(1),
+        index.stride(0),
+        index.stride(1),
+        eps,
+        HAS_WEIGHT=weight is not x_2d,
+        HAS_BIAS=bias is not x_2d,
+        BLOCK_N=BLOCK_N,
+        num_warps=num_warps,
+        num_stages=num_stages,
+    )
+    return output, gate_out
+
+
+def fuse_residual_layernorm_scale_shift_gate_select01_kernel(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    residual_gate: torch.Tensor,
+    weight: torch.Tensor | None,
+    bias: torch.Tensor | None,
+    scale0: torch.Tensor,
+    shift0: torch.Tensor,
+    gate0: torch.Tensor,
+    scale1: torch.Tensor,
+    shift1: torch.Tensor,
+    gate1: torch.Tensor,
+    index: torch.Tensor,
+    eps: float,
+):
+    assert x.is_cuda
+    assert x.is_contiguous()
+    assert residual.is_contiguous()
+    assert residual_gate.is_contiguous()
+    B, L, C = x.shape
+    output = torch.empty_like(x)
+    residual_out = torch.empty_like(x)
+    gate_out = torch.empty_like(x)
+
+    if residual.shape != x.shape:
+        raise ValueError("residual must have the same shape as x")
+    if residual_gate.shape != x.shape:
+        raise ValueError("residual_gate must have the same shape as x")
+    if (
+        scale0.dim() != 2
+        or shift0.dim() != 2
+        or gate0.dim() != 2
+        or scale1.dim() != 2
+        or shift1.dim() != 2
+        or gate1.dim() != 2
+    ):
+        raise ValueError("scale0/shift0/gate0/scale1/shift1/gate1 must be 2D [B, C]")
+    if index.dim() != 2:
+        raise ValueError("index must be 2D [B, L]")
+    if weight is not None and (weight.dim() != 1 or weight.shape[0] != C):
+        raise ValueError("weight must be 1D [C]")
+    if bias is not None and (bias.dim() != 1 or bias.shape[0] != C):
+        raise ValueError("bias must be 1D [C]")
+
+    x_2d = x.view(B * L, C)
+    residual_2d = residual.view(B * L, C)
+    residual_gate_2d = residual_gate.view(B * L, C)
+    output_2d = output.view(B * L, C)
+    residual_out_2d = residual_out.view(B * L, C)
+    gate_out_2d = gate_out.view(B * L, C)
+    weight = weight.contiguous() if weight is not None else x_2d
+    bias = bias.contiguous() if bias is not None else x_2d
+
+    MAX_FUSED_SIZE = 65536 // x_2d.element_size()
+    BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(C))
+    if C > BLOCK_N:
+        raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
+    num_warps, num_stages = 4, 4
+
+    grid = (B * L,)
+    _fused_residual_layernorm_scale_shift_gate_select01_kernel[grid](
+        output_2d,
+        residual_out_2d,
+        gate_out_2d,
+        x_2d,
+        residual_2d,
+        residual_gate_2d,
+        weight,
+        bias,
+        scale0.contiguous(),
+        shift0.contiguous(),
+        gate0.contiguous(),
+        scale1.contiguous(),
+        shift1.contiguous(),
+        gate1.contiguous(),
+        index.contiguous(),
+        C,
+        L,
+        x_2d.stride(0),
+        residual_2d.stride(0),
+        residual_gate_2d.stride(0),
+        output_2d.stride(0),
+        residual_out_2d.stride(0),
+        gate_out_2d.stride(0),
+        weight.stride(0) if weight.dim() == 1 else 0,
+        bias.stride(0) if bias.dim() == 1 else 0,
+        scale0.stride(0),
+        scale0.stride(1),
+        shift0.stride(0),
+        shift0.stride(1),
+        gate0.stride(0),
+        gate0.stride(1),
+        scale1.stride(0),
+        scale1.stride(1),
+        shift1.stride(0),
+        shift1.stride(1),
+        gate1.stride(0),
+        gate1.stride(1),
+        index.stride(0),
+        index.stride(1),
+        eps,
+        HAS_WEIGHT=weight is not x_2d,
+        HAS_BIAS=bias is not x_2d,
+        BLOCK_N=BLOCK_N,
+        num_warps=num_warps,
+        num_stages=num_stages,
+    )
+    return output, residual_out, gate_out
+
+
+if current_platform.is_npu():
+    from .npu_fallback import fuse_scale_shift_native
+
+    fuse_scale_shift_kernel = fuse_scale_shift_native
+
+if current_platform.is_mps():
+    from .mps_fallback import fuse_scale_shift_kernel_native
+
+    fuse_scale_shift_kernel = fuse_scale_shift_kernel_native
+
+if current_platform.is_musa():
+    from .torch_fallback import fuse_scale_shift_kernel_native
+
+    fuse_scale_shift_kernel = fuse_scale_shift_kernel_native
+
+if current_platform.is_cpu():
+    from .torch_fallback import (
+        fuse_scale_shift_kernel_native,
+    )
+
+    fuse_scale_shift_kernel = fuse_scale_shift_kernel_native

python/sglang/jit_kernel/diffusion/triton/torch_fallback.py

@@ -0,0 +1,146 @@
+"""Pytorch native based fallbacks for Triton diffusion kernels.
+
+Triton is not available on some platforms, so these pure-PyTorch
+implementations replace the Triton kernels
+
+"""
+
+from typing import Optional
+
+import torch
+from torch import Tensor
+
+
+def fuse_scale_shift_kernel_native(
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    scale_constant: float = 1.0,
+    block_l: int = 128,
+    block_c: int = 128,
+):
+    """Native fallback for fuse_scale_shift_kernel with scale_constant support."""
+    B, L, C = x.shape
+
+    def _expand(t: torch.Tensor) -> torch.Tensor:
+        if t.dim() == 4:
+            # [B, F, 1, C] -> [B, L, C]
+            num_frames = t.shape[1]
+            frame_seqlen = L // num_frames
+            return (
+                t.squeeze(2)
+                .unsqueeze(2)
+                .expand(-1, -1, frame_seqlen, -1)
+                .reshape(B, L, C)
+            )
+        elif t.dim() == 2:
+            # [B, C] -> [B, 1, C]
+            return t.unsqueeze(1)
+        return t
+
+    scale = _expand(scale)
+    shift = _expand(shift)
+
+    return x * (scale_constant + scale) + shift
+
+
+def apply_rotary_embedding_native(
+    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, interleaved: bool = False
+) -> torch.Tensor:
+    """Native fallback for rotary embedding (shared with NPU implementation)."""
+    if interleaved and cos.shape[-1] == x.shape[-1]:
+        cos = cos[..., ::2]
+        sin = sin[..., ::2]
+    cos = cos.unsqueeze(-2).to(x.dtype)
+    sin = sin.unsqueeze(-2).to(x.dtype)
+    x1 = x[..., ::2]
+    x2 = x[..., 1::2]
+    o1 = x1 * cos - x2 * sin
+    o2 = x2 * cos + x1 * sin
+    return torch.stack((o1, o2), dim=-1).flatten(-2)
+
+
+def norm_infer_native(
+    x: Tensor,
+    weight: Optional[Tensor],
+    bias: Optional[Tensor],
+    eps: float,
+    is_rms_norm: bool = False,
+    out: Optional[Tensor] = None,
+) -> Tensor:
+    """Native fallback for norm_infer (layer norm / rms norm inference)."""
+    orig_dtype = x.dtype
+    x = x.contiguous().float()
+    if is_rms_norm:
+        variance = x.pow(2).mean(dim=-1, keepdim=True)
+        x_hat = x * torch.rsqrt(variance + eps)
+    else:
+        mean = x.mean(dim=-1, keepdim=True)
+        variance = (x - mean).pow(2).mean(dim=-1, keepdim=True)
+        x_hat = (x - mean) * torch.rsqrt(variance + eps)
+    if weight is not None:
+        x_hat = x_hat * weight.float()
+    if bias is not None:
+        x_hat = x_hat + bias.float()
+    result = x_hat.to(orig_dtype)
+    if out is not None:
+        out.copy_(result)
+        return out
+    return result
+
+
+def triton_one_pass_rms_norm_native(
+    x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6
+) -> torch.Tensor:
+    """Native fallback for triton_one_pass_rms_norm."""
+    shape = x.shape
+    orig_dtype = x.dtype
+    x = x.contiguous().float()
+    variance = x.pow(2).mean(dim=-1, keepdim=True)
+    x_hat = x * torch.rsqrt(variance + eps)
+    return (x_hat * w.float()).to(orig_dtype).view(shape)
+
+
+def rms_norm_fn_native(
+    x,
+    weight,
+    bias,
+    residual=None,
+    x1=None,
+    weight1=None,
+    bias1=None,
+    eps=1e-6,
+    dropout_p=0.0,
+    rowscale=None,
+    prenorm=False,
+    residual_in_fp32=False,
+    zero_centered_weight=False,
+    return_dropout_mask=False,
+    out_dtype=None,
+    out=None,
+    residual_out=None,
+):
+    """Native fallback for rms_norm_fn (inference only, no dropout/x1 support)."""
+    x_shape_og = x.shape
+    orig_dtype = x.dtype
+    x = x.reshape(-1, x.shape[-1]).float()
+    if residual is not None:
+        residual = residual.reshape(-1, residual.shape[-1]).float()
+        x = x + residual
+        residual_out_val = x.to(torch.float32 if residual_in_fp32 else orig_dtype)
+    else:
+        residual_out_val = None
+    variance = x.pow(2).mean(dim=-1, keepdim=True)
+    x_hat = x * torch.rsqrt(variance + eps)
+    if weight is not None:
+        w = weight.float()
+        if zero_centered_weight:
+            w = w + 1.0
+        x_hat = x_hat * w
+    if bias is not None:
+        x_hat = x_hat + bias.float()
+    final_dtype = out_dtype if out_dtype is not None else orig_dtype
+    y = x_hat.to(final_dtype).reshape(x_shape_og)
+    if residual is not None and residual_out_val is not None:
+        return y, residual_out_val.reshape(x_shape_og)
+    return y

python/sglang/jit_kernel/diffusion/triton/varlen_pack_pad.py

@@ -0,0 +1,191 @@
+"""Fused Triton pack/scatter kernels for the varlen mask path.
+
+Used by ``USPAttention.forward`` masked branch to gather Q/K/V at valid
+positions and scatter the FA output back to the dense ``[B, S, H, D]`` layout.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton  # type: ignore
+import triton.language as tl  # type: ignore
+
+# ---------------------------------------------------------------------------
+# Pack (unpad) — gather Q/K/V at indices into packed [total_valid, H, D]
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _fused_pack_qkv_kernel(
+    Q_ptr,
+    K_ptr,
+    V_ptr,
+    Q_unpad_ptr,
+    K_unpad_ptr,
+    V_unpad_ptr,
+    indices_ptr,
+    HD,  # H * D, flattened feature dim
+    src_row_stride,  # stride between rows in Q/K/V (B*S row -> next row)
+    dst_row_stride,  # stride in Q_unpad/K_unpad/V_unpad
+    BLOCK_HD: tl.constexpr,
+):
+    """One program per packed row; copies Q[src], K[src], V[src] to dst row."""
+    out_row = tl.program_id(0)
+    src_row = tl.load(indices_ptr + out_row).to(tl.int64)
+
+    cols = tl.arange(0, BLOCK_HD)
+    col_mask = cols < HD
+
+    src_offset = src_row * src_row_stride + cols
+    dst_offset = out_row * dst_row_stride + cols
+
+    q_val = tl.load(Q_ptr + src_offset, mask=col_mask)
+    k_val = tl.load(K_ptr + src_offset, mask=col_mask)
+    v_val = tl.load(V_ptr + src_offset, mask=col_mask)
+
+    tl.store(Q_unpad_ptr + dst_offset, q_val, mask=col_mask)
+    tl.store(K_unpad_ptr + dst_offset, k_val, mask=col_mask)
+    tl.store(V_unpad_ptr + dst_offset, v_val, mask=col_mask)
+
+
+def fused_pack_qkv(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    indices: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Pack ``[B, S, H, D]`` Q/K/V at ``indices`` into ``[total_valid, H, D]``.
+
+    ``indices`` is the int64 flat ``B*S`` position for each kept token.
+    Non-contiguous inputs are made contiguous internally.
+    """
+    assert q.shape == k.shape == v.shape, "Q/K/V must share shape"
+    assert q.dtype == k.dtype == v.dtype, "Q/K/V must share dtype"
+    assert q.dim() == 4, "Q/K/V must be [B, S, H, D]"
+    assert indices.dtype in (torch.int32, torch.int64)
+    q = q.contiguous()
+    k = k.contiguous()
+    v = v.contiguous()
+    bs, seq, num_heads, head_dim = q.shape
+    hd = num_heads * head_dim
+    n_valid = indices.shape[0]
+
+    if n_valid == 0:
+        return (
+            q.new_empty(0, num_heads, head_dim),
+            k.new_empty(0, num_heads, head_dim),
+            v.new_empty(0, num_heads, head_dim),
+        )
+
+    block_hd = triton.next_power_of_2(hd)
+    q_flat = q.view(bs * seq, hd)
+    k_flat = k.view(bs * seq, hd)
+    v_flat = v.view(bs * seq, hd)
+    q_unpad = torch.empty(n_valid, hd, dtype=q.dtype, device=q.device)
+    k_unpad = torch.empty(n_valid, hd, dtype=k.dtype, device=k.device)
+    v_unpad = torch.empty(n_valid, hd, dtype=v.dtype, device=v.device)
+
+    with torch.get_device_module().device(q.device):
+        _fused_pack_qkv_kernel[(n_valid,)](
+            q_flat,
+            k_flat,
+            v_flat,
+            q_unpad,
+            k_unpad,
+            v_unpad,
+            indices,
+            hd,
+            q_flat.stride(0),
+            q_unpad.stride(0),
+            BLOCK_HD=block_hd,
+        )
+
+    return (
+        q_unpad.view(n_valid, num_heads, head_dim),
+        k_unpad.view(n_valid, num_heads, head_dim),
+        v_unpad.view(n_valid, num_heads, head_dim),
+    )
+
+
+# ---------------------------------------------------------------------------
+# Scatter (pad) — write packed output to [B, S, H, D] with zeros at invalid
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _fused_scatter_to_padded_kernel(
+    Out_unpad_ptr,
+    Out_padded_ptr,
+    inv_indices_ptr,  # [B*S]: pack idx for valid row, -1 for invalid
+    HD,
+    src_row_stride,
+    dst_row_stride,
+    BLOCK_HD: tl.constexpr,
+):
+    """One program per padded row; writes from pack or zeros."""
+    pad_row = tl.program_id(0)
+    inv_idx = tl.load(inv_indices_ptr + pad_row).to(tl.int64)
+
+    cols = tl.arange(0, BLOCK_HD)
+    col_mask = cols < HD
+    valid = inv_idx >= 0
+
+    safe_idx = tl.where(valid, inv_idx, 0)
+    src_offset = safe_idx * src_row_stride + cols
+    dst_offset = pad_row * dst_row_stride + cols
+
+    val = tl.load(Out_unpad_ptr + src_offset, mask=col_mask & valid, other=0.0)
+    tl.store(Out_padded_ptr + dst_offset, val, mask=col_mask)
+
+
+def fused_scatter_to_padded(
+    out_unpad: torch.Tensor,
+    inv_indices: torch.Tensor,
+    batch_size: int,
+    seqlen: int,
+) -> torch.Tensor:
+    """Scatter packed varlen output back to ``[B, S, H, D]`` with zero padding.
+
+    ``inv_indices`` is ``[B*S]`` giving the pack row index for each padded
+    position (``-1`` for padding). Non-contiguous ``out_unpad`` is made contiguous.
+    """
+    assert out_unpad.dim() == 3, "out_unpad must be [total_valid, H, D]"
+    assert inv_indices.shape == (batch_size * seqlen,)
+    assert inv_indices.dtype in (torch.int32, torch.int64)
+    out_unpad = out_unpad.contiguous()
+    _, num_heads, head_dim = out_unpad.shape
+    hd = num_heads * head_dim
+    block_hd = triton.next_power_of_2(hd)
+
+    out_padded = torch.empty(
+        batch_size * seqlen, hd, dtype=out_unpad.dtype, device=out_unpad.device
+    )
+    out_unpad_flat = out_unpad.view(-1, hd)
+
+    with torch.get_device_module().device(out_unpad.device):
+        _fused_scatter_to_padded_kernel[(batch_size * seqlen,)](
+            out_unpad_flat,
+            out_padded,
+            inv_indices,
+            hd,
+            out_unpad_flat.stride(0),
+            out_padded.stride(0),
+            BLOCK_HD=block_hd,
+        )
+
+    return out_padded.view(batch_size, seqlen, num_heads, head_dim)
+
+
+# ---------------------------------------------------------------------------
+# Inverse-index builder (called once per request alongside indices)
+# ---------------------------------------------------------------------------
+
+
+def build_inv_indices(indices: torch.Tensor, total_rows: int) -> torch.Tensor:
+    """For each padded row in ``[B*S]``, return its pack index or ``-1``."""
+    n_valid = indices.shape[0]
+    inv = torch.full((total_rows,), -1, dtype=torch.int32, device=indices.device)
+    inv[indices.long()] = torch.arange(
+        n_valid, dtype=torch.int32, device=indices.device
+    )
+    return inv

python/sglang/jit_kernel/diffusion/triton/zimage_native_norm.py

@@ -0,0 +1,182 @@
+import torch
+import triton  # type: ignore
+import triton.language as tl  # type: ignore
+
+
+@triton.jit
+def _tanh(x):
+    return 2.0 / (1.0 + tl.exp(-2.0 * x)) - 1.0
+
+
+@triton.jit
+def _rmsnorm_scale_kernel(
+    y_ptr,
+    x_ptr,
+    weight_ptr,
+    scale_ptr,
+    x_row_stride,
+    scale_row_stride,
+    seq_len,
+    dim: tl.constexpr,
+    eps: tl.constexpr,
+    block_dim: tl.constexpr,
+):
+    row = tl.program_id(0)
+    offsets = tl.arange(0, block_dim)
+    mask = offsets < dim
+
+    x = tl.load(x_ptr + row * x_row_stride + offsets, mask=mask, other=0.0)
+    square = (x * x).to(tl.bfloat16)
+    mean_square = (tl.sum(square, axis=0) / dim).to(tl.bfloat16)
+    rstd = tl.rsqrt((mean_square + eps).to(tl.bfloat16).to(tl.float32)).to(tl.bfloat16)
+
+    batch = row // seq_len
+    weight = tl.load(weight_ptr + offsets, mask=mask, other=0.0)
+    scale = tl.load(
+        scale_ptr + batch * scale_row_stride + offsets, mask=mask, other=0.0
+    )
+    y = (((x * rstd).to(tl.bfloat16) * weight).to(tl.bfloat16) * scale).to(tl.bfloat16)
+    tl.store(y_ptr + row * dim + offsets, y, mask=mask)
+
+
+@triton.jit
+def _rmsnorm_tanh_residual_kernel(
+    y_ptr,
+    x_ptr,
+    gate_ptr,
+    residual_ptr,
+    weight_ptr,
+    x_row_stride,
+    gate_row_stride,
+    residual_row_stride,
+    seq_len,
+    dim: tl.constexpr,
+    eps: tl.constexpr,
+    block_dim: tl.constexpr,
+):
+    row = tl.program_id(0)
+    offsets = tl.arange(0, block_dim)
+    mask = offsets < dim
+
+    x = tl.load(x_ptr + row * x_row_stride + offsets, mask=mask, other=0.0)
+    square = (x * x).to(tl.bfloat16)
+    mean_square = (tl.sum(square, axis=0) / dim).to(tl.bfloat16)
+    rstd = tl.rsqrt((mean_square + eps).to(tl.bfloat16).to(tl.float32)).to(tl.bfloat16)
+
+    batch = row // seq_len
+    gate = tl.load(gate_ptr + batch * gate_row_stride + offsets, mask=mask, other=0.0)
+    residual = tl.load(
+        residual_ptr + row * residual_row_stride + offsets, mask=mask, other=0.0
+    )
+    weight = tl.load(weight_ptr + offsets, mask=mask, other=0.0)
+    norm = ((x * rstd).to(tl.bfloat16) * weight).to(tl.bfloat16)
+    gated = (_tanh(gate.to(tl.float32)).to(tl.bfloat16) * norm).to(tl.bfloat16)
+    y = (residual + gated).to(tl.bfloat16)
+    tl.store(y_ptr + row * dim + offsets, y, mask=mask)
+
+
+def _flat_row_stride(x: torch.Tensor) -> int | None:
+    if x.dim() < 2 or x.stride(-1) != 1:
+        return None
+    row_stride = x.stride(-2)
+    expected_stride = row_stride * x.shape[-2]
+    for dim in range(x.dim() - 3, -1, -1):
+        if x.stride(dim) != expected_stride:
+            return None
+        expected_stride *= x.shape[dim]
+    return row_stride
+
+
+def _can_use(x: torch.Tensor, weight: torch.Tensor, other: torch.Tensor) -> bool:
+    return (
+        x.is_cuda
+        and weight.is_cuda
+        and other.is_cuda
+        and x.dtype == torch.bfloat16
+        and weight.dtype == torch.bfloat16
+        and other.dtype == torch.bfloat16
+        and weight.is_contiguous()
+        and x.shape[-1] <= 8192
+        and _flat_row_stride(x) is not None
+        and _flat_row_stride(other) is not None
+    )
+
+
+def zimage_rmsnorm_scale(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    scale: torch.Tensor,
+    eps: float,
+) -> torch.Tensor | None:
+    if not _can_use(x, weight, scale):
+        return None
+    shape = x.shape
+    dim = shape[-1]
+    x_rows = x.numel() // dim
+    scale_rows = scale.numel() // dim
+    if x_rows % scale_rows != 0:
+        return None
+    seq_len = x_rows // scale_rows
+    x_row_stride = _flat_row_stride(x)
+    scale_row_stride = _flat_row_stride(scale)
+    if x_row_stride is None or scale_row_stride is None:
+        return None
+    y = torch.empty_like(x, memory_format=torch.contiguous_format)
+    with torch.get_device_module().device(x.device):
+        _rmsnorm_scale_kernel[(x_rows,)](
+            y.reshape(-1, dim),
+            x,
+            weight,
+            scale,
+            x_row_stride,
+            scale_row_stride,
+            seq_len,
+            dim,
+            eps,
+            block_dim=triton.next_power_of_2(dim),
+            num_warps=8,
+        )
+    return y
+
+
+def zimage_rmsnorm_tanh_residual(
+    x: torch.Tensor,
+    gate: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+) -> torch.Tensor | None:
+    if not (_can_use(x, weight, gate) and residual.is_cuda):
+        return None
+    if residual.dtype != x.dtype or _flat_row_stride(residual) is None:
+        return None
+    shape = x.shape
+    dim = shape[-1]
+    x_rows = x.numel() // dim
+    gate_rows = gate.numel() // dim
+    if x_rows % gate_rows != 0:
+        return None
+    seq_len = x_rows // gate_rows
+    x_row_stride = _flat_row_stride(x)
+    gate_row_stride = _flat_row_stride(gate)
+    residual_row_stride = _flat_row_stride(residual)
+    if x_row_stride is None or gate_row_stride is None or residual_row_stride is None:
+        return None
+    y = torch.empty_like(x, memory_format=torch.contiguous_format)
+    with torch.get_device_module().device(x.device):
+        _rmsnorm_tanh_residual_kernel[(x_rows,)](
+            y.reshape(-1, dim),
+            x,
+            gate,
+            residual,
+            weight,
+            x_row_stride,
+            gate_row_stride,
+            residual_row_stride,
+            seq_len,
+            dim,
+            eps,
+            block_dim=triton.next_power_of_2(dim),
+            num_warps=8,
+        )
+    return y