[Kernel] Migrate linear-attention, MiniMax-sparse and diffusion kernels to sglang.kernels (RFC #29630, Phase 2.5, 6/7) (#30793)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

python/sglang/kernels/ops/attention/__init__.py

@@ -43,6 +43,32 @@ del _mod, _fn
 __all__ = []
 
 
+# Linear-attention / MiniMax-sparse / diffusion kernels migrated in Phase 2.5
+# (RFC #29630); registered for inventory.
+for _grp, _mod, _fn in [
+    ("attention", "linear.seg_la", "seg_la_fwd"),
+    ("attention", "linear.lightning_attn", "lightning_attention"),
+    ("attention", "linear.lightning_attn", "linear_decode_forward_triton"),
+    (
+        "attention",
+        "minimax_sparse.decode.flash_with_topk_idx",
+        "flash_decode_with_topk_idx",
+    ),
+    (
+        "attention",
+        "minimax_sparse.prefill.flash_with_topk_idx",
+        "flash_prefill_with_topk_index",
+    ),
+]:
+    register_kernel(
+        KernelSpec(
+            op=f"{_grp}.{_fn}",
+            backend=KernelBackend.TRITON,
+            target=f"sglang.kernels.ops.{_grp}.{_mod}:{_fn}",
+        )
+    )
+del _grp, _mod, _fn
+
 # DeepSeek DSA / DSV4 kernels migrated in Phase 2.5 (RFC #29630);
 # registered for inventory. Import them from their modules.
 for _mod, _fn in [

python/sglang/kernels/ops/attention/linear/__init__.py

@@ -0,0 +1 @@
+"""Linear-attention kernels (RFC #29630, Phase 2.5)."""

python/sglang/kernels/ops/attention/minimax_sparse/__init__.py

@@ -0,0 +1 @@
+"""MiniMax sparse-attention kernels (RFC #29630, Phase 2.5)."""

python/sglang/kernels/ops/diffusion/__init__.py

@@ -107,3 +107,13 @@ __all__ = [
     "residual_gate_add",
     "fused_inplace_qknorm_rope",
 ]
+
+
+# Migrated from multimodal_gen (RFC #29630, Phase 2.5).
+register_kernel(
+    KernelSpec(
+        op="diffusion.sparse_linear_attn_fwd",
+        backend=KernelBackend.TRITON,
+        target="sglang.kernels.ops.diffusion.sparse_linear_attn_kernels:get_block_map",
+    )
+)

python/sglang/kernels/ops/diffusion/render/hunyuan3d_rasterizer/__init__.py

@@ -12,7 +12,8 @@ import os
 from typing import Tuple
 
 import torch
-from sglang.multimodal_gen.csrc.render import load_extension_with_recovery
+
+from sglang.kernels.ops.diffusion.render import load_extension_with_recovery
 
 _abs_path = os.path.dirname(os.path.abspath(__file__))
 _custom_rasterizer_kernel = None
@@ -26,14 +27,15 @@ def _load_custom_rasterizer(
 
     if _custom_rasterizer_kernel is not None:
         return _custom_rasterizer_kernel
-    
+
     cuda_enabled_flag = ["-DCUDA_ENABLED"] if is_cuda else []
-    
+
     _custom_rasterizer_kernel = load_extension_with_recovery(
         name="custom_rasterizer_kernel",
         sources=[
             f"{_abs_path}/rasterizer.cpp",
-        ] + ([f"{_abs_path}/rasterizer_gpu.cu"] if is_cuda else []),
+        ]
+        + ([f"{_abs_path}/rasterizer_gpu.cu"] if is_cuda else []),
         extra_cflags=["-O3"] + cuda_enabled_flag,
         extra_cuda_cflags=["-O3", "--use_fast_math"] + cuda_enabled_flag,
         verbose=False,
@@ -60,7 +62,13 @@ def rasterize(
         pos = pos[0]
 
     findices, barycentric = kernel.rasterize_image(
-        pos.to(device), tri.to(device), clamp_depth.to(device), resolution[1], resolution[0], 1e-6, use_depth_prior
+        pos.to(device),
+        tri.to(device),
+        clamp_depth.to(device),
+        resolution[1],
+        resolution[0],
+        1e-6,
+        use_depth_prior,
     )
 
     findices = findices.to(pos.device)

python/sglang/kernels/ops/diffusion/render/mesh_processor/__init__.py

@@ -12,7 +12,8 @@ import os
 from typing import Tuple
 
 import numpy as np
-from sglang.multimodal_gen.csrc.render import load_extension_with_recovery
+
+from sglang.kernels.ops.diffusion.render import load_extension_with_recovery
 
 _abs_path = os.path.dirname(os.path.abspath(__file__))
 _mesh_processor_kernel = None
@@ -55,7 +56,9 @@ def meshVerticeInpaint(
     pos_idx = np.ascontiguousarray(pos_idx, dtype=np.int32)
     uv_idx = np.ascontiguousarray(uv_idx, dtype=np.int32)
 
-    return kernel.meshVerticeInpaint(texture, mask, vtx_pos, vtx_uv, pos_idx, uv_idx, method)
+    return kernel.meshVerticeInpaint(
+        texture, mask, vtx_pos, vtx_uv, pos_idx, uv_idx, method
+    )
 
 
 __all__ = ["meshVerticeInpaint"]

python/sglang/kernels/ops/diffusion/sparse_linear_attn_kernels.py

@@ -0,0 +1,130 @@
+"""Sparse linear-attention block-map and fwd kernels, migrated from
+``sglang.multimodal_gen.runtime.layers.attention.backends.sparse_linear_attn``
+(RFC #29630, Phase 2.5).
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+
+def get_block_map(q, k, topk_ratio, BLKQ=64, BLKK=64):
+    arg_k = k - torch.mean(
+        k, dim=-2, keepdim=True
+    )  # smooth-k technique in SageAttention
+    pooled_qblocks = mean_pool(q, BLKQ)
+    pooled_kblocks = mean_pool(arg_k, BLKK)
+    pooled_score = pooled_qblocks @ pooled_kblocks.transpose(-1, -2)
+
+    K = pooled_score.shape[-1]
+    topk = min(K, int(topk_ratio * K))
+    lut = torch.topk(pooled_score, topk, dim=-1, sorted=False).indices
+
+    sparse_map = torch.zeros_like(pooled_score, dtype=torch.int8)
+    sparse_map.scatter_(-1, lut, 1)
+    return sparse_map, lut, topk
+
+
+def mean_pool(x, BLK):
+    assert x.is_contiguous()
+
+    B, H, L, D = x.shape
+    L_BLOCKS = (L + BLK - 1) // BLK
+    x_mean = torch.empty((B, H, L_BLOCKS, D), device=x.device, dtype=x.dtype)
+
+    grid = (L_BLOCKS, B * H)
+    compress_kernel[grid](x, x_mean, L, D, BLK)
+    return x_mean
+
+
+@triton.jit
+def compress_kernel(
+    X,
+    XM,
+    L: tl.constexpr,
+    D: tl.constexpr,
+    BLOCK_L: tl.constexpr,
+):
+    idx_l = tl.program_id(0)
+    idx_bh = tl.program_id(1)
+
+    offs_l = idx_l * BLOCK_L + tl.arange(0, BLOCK_L)
+    offs_d = tl.arange(0, D)
+
+    x_offset = idx_bh * L * D
+    xm_offset = idx_bh * ((L + BLOCK_L - 1) // BLOCK_L) * D
+    x = tl.load(
+        X + x_offset + offs_l[:, None] * D + offs_d[None, :], mask=offs_l[:, None] < L
+    )
+
+    nx = min(BLOCK_L, L - idx_l * BLOCK_L)
+    x_mean = tl.sum(x, axis=0, dtype=tl.float32) / nx
+    tl.store(XM + xm_offset + idx_l * D + offs_d, x_mean.to(XM.dtype.element_ty))
+
+
+@triton.jit
+def _attn_fwd(
+    Q,
+    K,
+    V,
+    qk_scale: tl.constexpr,
+    topk: tl.constexpr,
+    LUT,
+    LSE,
+    OS,
+    L: tl.constexpr,
+    M_BLOCKS: tl.constexpr,
+    D: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    idx_m = tl.program_id(0).to(tl.int64)
+    idx_bh = tl.program_id(1).to(tl.int64)
+
+    qkv_offset = idx_bh * L * D
+    lut_offset = (idx_bh * M_BLOCKS + idx_m) * topk
+    lse_offset = idx_bh * L
+    offs_m = idx_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = tl.arange(0, BLOCK_N)
+    offs_d = tl.arange(0, D)
+
+    Q_ptrs = Q + qkv_offset + offs_m[:, None] * D + offs_d[None, :]
+    K_ptrs = K + qkv_offset + offs_n[None, :] * D + offs_d[:, None]
+    V_ptrs = V + qkv_offset + offs_n[:, None] * D + offs_d[None, :]
+    OS_ptrs = OS + qkv_offset + offs_m[:, None] * D + offs_d[None, :]
+    LUT_ptr = LUT + lut_offset
+    LSE_ptrs = LSE + lse_offset + offs_m
+
+    m_i = tl.full([BLOCK_M], -float("inf"), dtype=tl.float32)
+    l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
+    o_s = tl.zeros([BLOCK_M, D], dtype=tl.float32)
+
+    q = tl.load(Q_ptrs, mask=offs_m[:, None] < L)
+    for block_idx in tl.range(topk):
+        idx_n = tl.load(LUT_ptr + block_idx)
+        n_mask = offs_n < L - idx_n * BLOCK_N
+
+        k = tl.load(K_ptrs + idx_n * BLOCK_N * D, mask=n_mask[None, :])
+        qk = tl.dot(q, k) * (qk_scale * 1.4426950408889634)  # = 1 / ln(2)
+        if L - idx_n * BLOCK_N < BLOCK_N:
+            qk = tl.where(n_mask[None, :], qk, float("-inf"))
+
+        v = tl.load(V_ptrs + idx_n * BLOCK_N * D, mask=n_mask[:, None])
+        local_m = tl.max(qk, 1)
+        new_m = tl.maximum(m_i, local_m)
+        qk = qk - new_m[:, None]
+
+        p = tl.math.exp2(qk)
+        l_ij = tl.sum(p, 1)
+        alpha = tl.math.exp2(m_i - new_m)
+        o_s = o_s * alpha[:, None]
+        o_s += tl.dot(p.to(v.dtype), v)
+
+        l_i = l_i * alpha + l_ij
+        m_i = new_m
+
+    o_s = o_s / l_i[:, None]
+    tl.store(OS_ptrs, o_s.to(OS.type.element_ty), mask=offs_m[:, None] < L)
+
+    m_i += tl.math.log2(l_i)
+    tl.store(LSE_ptrs, m_i, mask=offs_m < L)

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

@@ -22,7 +22,6 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import triton
-import triton.language as tl
 
 from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
     AttentionBackend,
@@ -37,126 +36,10 @@ logger = init_logger(__name__)
 
 
 # ==================================SLA Functions===================================
-def get_block_map(q, k, topk_ratio, BLKQ=64, BLKK=64):
-    arg_k = k - torch.mean(
-        k, dim=-2, keepdim=True
-    )  # smooth-k technique in SageAttention
-    pooled_qblocks = mean_pool(q, BLKQ)
-    pooled_kblocks = mean_pool(arg_k, BLKK)
-    pooled_score = pooled_qblocks @ pooled_kblocks.transpose(-1, -2)
-
-    K = pooled_score.shape[-1]
-    topk = min(K, int(topk_ratio * K))
-    lut = torch.topk(pooled_score, topk, dim=-1, sorted=False).indices
-
-    sparse_map = torch.zeros_like(pooled_score, dtype=torch.int8)
-    sparse_map.scatter_(-1, lut, 1)
-    return sparse_map, lut, topk
-
-
-def mean_pool(x, BLK):
-    assert x.is_contiguous()
-
-    B, H, L, D = x.shape
-    L_BLOCKS = (L + BLK - 1) // BLK
-    x_mean = torch.empty((B, H, L_BLOCKS, D), device=x.device, dtype=x.dtype)
-
-    grid = (L_BLOCKS, B * H)
-    compress_kernel[grid](x, x_mean, L, D, BLK)
-    return x_mean
-
-
-@triton.jit
-def compress_kernel(
-    X,
-    XM,
-    L: tl.constexpr,
-    D: tl.constexpr,
-    BLOCK_L: tl.constexpr,
-):
-    idx_l = tl.program_id(0)
-    idx_bh = tl.program_id(1)
-
-    offs_l = idx_l * BLOCK_L + tl.arange(0, BLOCK_L)
-    offs_d = tl.arange(0, D)
-
-    x_offset = idx_bh * L * D
-    xm_offset = idx_bh * ((L + BLOCK_L - 1) // BLOCK_L) * D
-    x = tl.load(
-        X + x_offset + offs_l[:, None] * D + offs_d[None, :], mask=offs_l[:, None] < L
-    )
-
-    nx = min(BLOCK_L, L - idx_l * BLOCK_L)
-    x_mean = tl.sum(x, axis=0, dtype=tl.float32) / nx
-    tl.store(XM + xm_offset + idx_l * D + offs_d, x_mean.to(XM.dtype.element_ty))
-
-
-@triton.jit
-def _attn_fwd(
-    Q,
-    K,
-    V,
-    qk_scale: tl.constexpr,
-    topk: tl.constexpr,
-    LUT,
-    LSE,
-    OS,
-    L: tl.constexpr,
-    M_BLOCKS: tl.constexpr,
-    D: tl.constexpr,
-    BLOCK_M: tl.constexpr,
-    BLOCK_N: tl.constexpr,
-):
-    idx_m = tl.program_id(0).to(tl.int64)
-    idx_bh = tl.program_id(1).to(tl.int64)
-
-    qkv_offset = idx_bh * L * D
-    lut_offset = (idx_bh * M_BLOCKS + idx_m) * topk
-    lse_offset = idx_bh * L
-    offs_m = idx_m * BLOCK_M + tl.arange(0, BLOCK_M)
-    offs_n = tl.arange(0, BLOCK_N)
-    offs_d = tl.arange(0, D)
-
-    Q_ptrs = Q + qkv_offset + offs_m[:, None] * D + offs_d[None, :]
-    K_ptrs = K + qkv_offset + offs_n[None, :] * D + offs_d[:, None]
-    V_ptrs = V + qkv_offset + offs_n[:, None] * D + offs_d[None, :]
-    OS_ptrs = OS + qkv_offset + offs_m[:, None] * D + offs_d[None, :]
-    LUT_ptr = LUT + lut_offset
-    LSE_ptrs = LSE + lse_offset + offs_m
-
-    m_i = tl.full([BLOCK_M], -float("inf"), dtype=tl.float32)
-    l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
-    o_s = tl.zeros([BLOCK_M, D], dtype=tl.float32)
-
-    q = tl.load(Q_ptrs, mask=offs_m[:, None] < L)
-    for block_idx in tl.range(topk):
-        idx_n = tl.load(LUT_ptr + block_idx)
-        n_mask = offs_n < L - idx_n * BLOCK_N
-
-        k = tl.load(K_ptrs + idx_n * BLOCK_N * D, mask=n_mask[None, :])
-        qk = tl.dot(q, k) * (qk_scale * 1.4426950408889634)  # = 1 / ln(2)
-        if L - idx_n * BLOCK_N < BLOCK_N:
-            qk = tl.where(n_mask[None, :], qk, float("-inf"))
-
-        v = tl.load(V_ptrs + idx_n * BLOCK_N * D, mask=n_mask[:, None])
-        local_m = tl.max(qk, 1)
-        new_m = tl.maximum(m_i, local_m)
-        qk = qk - new_m[:, None]
-
-        p = tl.math.exp2(qk)
-        l_ij = tl.sum(p, 1)
-        alpha = tl.math.exp2(m_i - new_m)
-        o_s = o_s * alpha[:, None]
-        o_s += tl.dot(p.to(v.dtype), v)
-
-        l_i = l_i * alpha + l_ij
-        m_i = new_m
-
-    o_s = o_s / l_i[:, None]
-    tl.store(OS_ptrs, o_s.to(OS.type.element_ty), mask=offs_m[:, None] < L)
-
-    m_i += tl.math.log2(l_i)
-    tl.store(LSE_ptrs, m_i, mask=offs_m < L)
+from sglang.kernels.ops.diffusion.sparse_linear_attn_kernels import (
+    _attn_fwd,
+    get_block_map,
+)
 
 
 def _get_cuda_arch(device_index: int) -> str:

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

@@ -18,7 +18,7 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 logger = init_logger(__name__)
 
 # Import C++ mesh processor extension
-from sglang.multimodal_gen.csrc.render.mesh_processor import meshVerticeInpaint
+from sglang.kernels.ops.diffusion.render.mesh_processor import meshVerticeInpaint
 
 
 def transform_pos(
@@ -363,7 +363,7 @@ class MeshRender:
         resolution: Tuple[int, int],
     ) -> torch.Tensor:
         """Rasterize using CUDA rasterizer."""
-        from sglang.multimodal_gen.csrc.render.hunyuan3d_rasterizer import rasterize
+        from sglang.kernels.ops.diffusion.render.hunyuan3d_rasterizer import rasterize
 
         if pos_clip.dim() == 2:
             pos_clip = pos_clip.unsqueeze(0)
@@ -380,7 +380,7 @@ class MeshRender:
         tri: torch.Tensor,
     ) -> torch.Tensor:
         """Interpolate vertex attributes."""
-        from sglang.multimodal_gen.csrc.render.hunyuan3d_rasterizer import interpolate
+        from sglang.kernels.ops.diffusion.render.hunyuan3d_rasterizer import interpolate
 
         barycentric = rast_out[0, ..., :-1]
         findices = rast_out[0, ..., -1].int()

python/sglang/srt/layers/attention/linear/kernels/gdn_cutedsl.py

@@ -65,11 +65,11 @@ class CuteDSLGDNKernel(LinearAttnKernelBase):
             raise RuntimeError(
                 f"CuTe DSL GDN prefill requires head_k_dim=128, got {head_k_dim}."
             )
-        from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd
-        from sglang.srt.layers.attention.linear.kernels.gdn_blackwell import (
+        from sglang.kernels.ops.attention.linear.gdn_blackwell import (
             chunk_gated_delta_rule_cutedsl,
             prepare_metadata_cutedsl,
         )
+        from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd
 
         self._extend_fn = chunk_gated_delta_rule_cutedsl
         self._prepare_meta_fn = prepare_metadata_cutedsl

python/sglang/srt/layers/attention/linear/kernels/kda_cutedsl.py

@@ -49,10 +49,10 @@ class CuteDSLKDAKernel(LinearAttnKernelBase):
             raise RuntimeError(
                 f"CuTe DSL KDA prefill requires head_k_dim=128, got {head_k_dim}."
             )
-        from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd
-        from sglang.srt.layers.attention.linear.kernels.kda_blackwell import (
+        from sglang.kernels.ops.attention.linear.kda_blackwell import (
             chunk_kda_cutedsl,
         )
+        from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd
 
         self._extend_fn = chunk_kda_cutedsl
         self._l2norm_fn = l2norm_fwd

python/sglang/srt/layers/attention/linear/lightning_backend.py

@@ -3,15 +3,15 @@ import math
 
 import torch
 
-from sglang.srt.layers.attention.hybrid_linear_attn_backend import MambaAttnBackendBase
-from sglang.srt.layers.attention.linear.lightning_attn import (
+from sglang.kernels.ops.attention.linear.lightning_attn import (
     BailingLinearKernel,
     linear_decode_forward_triton,
 )
+from sglang.kernels.ops.attention.linear.seg_la import SegLaMeta, seg_la_fwd
+from sglang.srt.layers.attention.hybrid_linear_attn_backend import MambaAttnBackendBase
 from sglang.srt.layers.attention.linear.linear_metadata import (
     BailingLinearMetadata,
 )
-from sglang.srt.layers.attention.linear.seg_la import SegLaMeta, seg_la_fwd
 from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_executor.model_runner import ModelRunner

python/sglang/srt/layers/attention/minimax_sparse_ops/minimax_sparse.py

@@ -5,12 +5,20 @@ from typing import Callable, List, Optional, Tuple
 
 import torch
 
-from .common.index import topk_index_reduce
-from .common.utils import get_cu_seqblocks
-from .decode.flash_with_topk_idx import flash_decode_with_topk_idx
-from .decode.topk_sparse import flash_decode_with_gqa_share_sparse
-from .prefill.flash_with_topk_idx import flash_prefill_with_topk_index
-from .prefill.topk_sparse import flash_prefill_with_gqa_share_sparse
+from sglang.kernels.ops.attention.minimax_sparse.common.index import topk_index_reduce
+from sglang.kernels.ops.attention.minimax_sparse.common.utils import get_cu_seqblocks
+from sglang.kernels.ops.attention.minimax_sparse.decode.flash_with_topk_idx import (
+    flash_decode_with_topk_idx,
+)
+from sglang.kernels.ops.attention.minimax_sparse.decode.topk_sparse import (
+    flash_decode_with_gqa_share_sparse,
+)
+from sglang.kernels.ops.attention.minimax_sparse.prefill.flash_with_topk_idx import (
+    flash_prefill_with_topk_index,
+)
+from sglang.kernels.ops.attention.minimax_sparse.prefill.topk_sparse import (
+    flash_prefill_with_gqa_share_sparse,
+)
 
 logger = logging.getLogger(__name__)
 _msa_fallback_warned = False

python/sglang/srt/layers/attention/minimax_sparse_ops/tests/test_flash_with_topk_idx.py

@@ -3,10 +3,10 @@ import sys
 import pytest
 import torch
 
-from sglang.srt.environ import envs
-from sglang.srt.layers.attention.minimax_sparse_ops.decode.flash_with_topk_idx import (
+from sglang.kernels.ops.attention.minimax_sparse.decode.flash_with_topk_idx import (
     flash_decode_with_topk_idx,
 )
+from sglang.srt.environ import envs
 
 DEVICE = "cuda"
 RTOL_VS_REF = 5e-3

python/sglang/srt/layers/attention/minimax_sparse_ops/tests/test_sparse_gqa.py

@@ -10,7 +10,7 @@ import sys
 import pytest
 import torch
 
-from sglang.srt.layers.attention.minimax_sparse_ops.decode.topk_sparse import (
+from sglang.kernels.ops.attention.minimax_sparse.decode.topk_sparse import (
     flash_decode_with_gqa_share_sparse,
 )