From c00131ebaaebca044129bf946bada3ab01a08f28 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Wed, 15 Jul 2026 11:21:36 +0800 Subject: [PATCH] [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 --- .pre-commit-config.yaml | 2 +- .../bench_gdn_prefill_cutedsl.py | 8 +- .../bench_kda_prefill_cutedsl.py | 12 +- .../sglang/kernels/ops/attention/__init__.py | 26 ++++ .../kernels/ops/attention/linear/__init__.py | 1 + .../linear}/gdn_blackwell/__init__.py | 0 .../linear}/gdn_blackwell/kernel_h.py | 0 .../gdn_blackwell/kernel_kkt_inv_uw.py | 0 .../linear}/gdn_blackwell/kernel_o.py | 0 .../linear}/kda_blackwell/__init__.py | 0 .../linear}/kda_blackwell/kernel_h.py | 0 .../kda_blackwell/kernel_kkt_inv_uw.py | 0 .../linear}/kda_blackwell/kernel_o.py | 0 .../linear}/kda_blackwell/prologue.py | 0 .../ops}/attention/linear/lightning_attn.py | 0 .../ops}/attention/linear/seg_la.py | 0 .../ops/attention/minimax_sparse/__init__.py | 1 + .../attention/minimax_sparse}/common/index.py | 0 .../attention/minimax_sparse}/common/utils.py | 0 .../decode/flash_with_topk_idx.py | 0 .../minimax_sparse}/decode/topk_sparse.py | 0 .../prefill/flash_with_topk_idx.py | 0 .../minimax_sparse}/prefill/topk_sparse.py | 0 .../sglang/kernels/ops/diffusion/__init__.py | 10 ++ .../ops/diffusion}/render/__init__.py | 0 .../render/hunyuan3d_rasterizer/__init__.py | 18 ++- .../hunyuan3d_rasterizer/rasterizer.cpp | 0 .../render/hunyuan3d_rasterizer/rasterizer.h | 0 .../hunyuan3d_rasterizer/rasterizer_gpu.cu | 0 .../render/mesh_processor/__init__.py | 7 +- .../render/mesh_processor/mesh_processor.cpp | 0 .../diffusion/sparse_linear_attn_kernels.py | 130 ++++++++++++++++++ .../attention/backends/sparse_linear_attn.py | 125 +---------------- .../runtime/utils/mesh3d_utils.py | 6 +- .../attention/linear/kernels/gdn_cutedsl.py | 4 +- .../attention/linear/kernels/kda_cutedsl.py | 4 +- .../attention/linear/lightning_backend.py | 6 +- .../minimax_sparse_ops/minimax_sparse.py | 20 ++- .../tests/test_flash_with_topk_idx.py | 4 +- .../tests/test_sparse_gqa.py | 2 +- .../attention/test_gdn_prefill_cutedsl.py | 8 +- .../attention/test_kda_prefill_cutedsl.py | 8 +- .../minimax/bench_minimax_decode_topk.py | 2 +- .../test_minimax_decode_topk_page_table.py | 2 +- 44 files changed, 238 insertions(+), 168 deletions(-) create mode 100644 python/sglang/kernels/ops/attention/linear/__init__.py rename python/sglang/{srt/layers/attention/linear/kernels => kernels/ops/attention/linear}/gdn_blackwell/__init__.py (100%) rename python/sglang/{srt/layers/attention/linear/kernels => kernels/ops/attention/linear}/gdn_blackwell/kernel_h.py (100%) rename python/sglang/{srt/layers/attention/linear/kernels => kernels/ops/attention/linear}/gdn_blackwell/kernel_kkt_inv_uw.py (100%) rename python/sglang/{srt/layers/attention/linear/kernels => kernels/ops/attention/linear}/gdn_blackwell/kernel_o.py (100%) rename python/sglang/{srt/layers/attention/linear/kernels => kernels/ops/attention/linear}/kda_blackwell/__init__.py (100%) rename python/sglang/{srt/layers/attention/linear/kernels => kernels/ops/attention/linear}/kda_blackwell/kernel_h.py (100%) rename python/sglang/{srt/layers/attention/linear/kernels => kernels/ops/attention/linear}/kda_blackwell/kernel_kkt_inv_uw.py (100%) rename python/sglang/{srt/layers/attention/linear/kernels => kernels/ops/attention/linear}/kda_blackwell/kernel_o.py (100%) rename python/sglang/{srt/layers/attention/linear/kernels => kernels/ops/attention/linear}/kda_blackwell/prologue.py (100%) rename python/sglang/{srt/layers => kernels/ops}/attention/linear/lightning_attn.py (100%) rename python/sglang/{srt/layers => kernels/ops}/attention/linear/seg_la.py (100%) create mode 100644 python/sglang/kernels/ops/attention/minimax_sparse/__init__.py rename python/sglang/{srt/layers/attention/minimax_sparse_ops => kernels/ops/attention/minimax_sparse}/common/index.py (100%) rename python/sglang/{srt/layers/attention/minimax_sparse_ops => kernels/ops/attention/minimax_sparse}/common/utils.py (100%) rename python/sglang/{srt/layers/attention/minimax_sparse_ops => kernels/ops/attention/minimax_sparse}/decode/flash_with_topk_idx.py (100%) rename python/sglang/{srt/layers/attention/minimax_sparse_ops => kernels/ops/attention/minimax_sparse}/decode/topk_sparse.py (100%) rename python/sglang/{srt/layers/attention/minimax_sparse_ops => kernels/ops/attention/minimax_sparse}/prefill/flash_with_topk_idx.py (100%) rename python/sglang/{srt/layers/attention/minimax_sparse_ops => kernels/ops/attention/minimax_sparse}/prefill/topk_sparse.py (100%) rename python/sglang/{multimodal_gen/csrc => kernels/ops/diffusion}/render/__init__.py (100%) rename python/sglang/{multimodal_gen/csrc => kernels/ops/diffusion}/render/hunyuan3d_rasterizer/__init__.py (88%) rename python/sglang/{multimodal_gen/csrc => kernels/ops/diffusion}/render/hunyuan3d_rasterizer/rasterizer.cpp (100%) rename python/sglang/{multimodal_gen/csrc => kernels/ops/diffusion}/render/hunyuan3d_rasterizer/rasterizer.h (100%) rename python/sglang/{multimodal_gen/csrc => kernels/ops/diffusion}/render/hunyuan3d_rasterizer/rasterizer_gpu.cu (100%) rename python/sglang/{multimodal_gen/csrc => kernels/ops/diffusion}/render/mesh_processor/__init__.py (89%) rename python/sglang/{multimodal_gen/csrc => kernels/ops/diffusion}/render/mesh_processor/mesh_processor.cpp (100%) create mode 100644 python/sglang/kernels/ops/diffusion/sparse_linear_attn_kernels.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index bd5365fc0..a4064717d 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,5 +1,5 @@ default_stages: [pre-commit, pre-push, manual] -exclude: ^(python/sglang/multimodal_gen/csrc|python/sglang/jit_kernel/flash_attention/cute) +exclude: ^(python/sglang/multimodal_gen/csrc|python/sglang/kernels/ops/diffusion/render|python/sglang/jit_kernel/flash_attention/cute) repos: - repo: https://github.com/pre-commit/pre-commit-hooks diff --git a/benchmark/bench_linear_attention/bench_gdn_prefill_cutedsl.py b/benchmark/bench_linear_attention/bench_gdn_prefill_cutedsl.py index acc9fd57f..bdcbfc6b5 100644 --- a/benchmark/bench_linear_attention/bench_gdn_prefill_cutedsl.py +++ b/benchmark/bench_linear_attention/bench_gdn_prefill_cutedsl.py @@ -26,14 +26,14 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "python") import torch +from sglang.kernels.ops.attention.linear.gdn_blackwell import ( + chunk_gated_delta_rule_cutedsl, + prepare_metadata_cutedsl, +) from sglang.srt.layers.attention.fla.chunk import ( chunk_gated_delta_rule as triton_chunk_gated_delta_rule, ) from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd -from sglang.srt.layers.attention.linear.kernels.gdn_blackwell import ( - chunk_gated_delta_rule_cutedsl, - prepare_metadata_cutedsl, -) # --------------------------------------------------------------------------- # Helpers (shared shape: pool layout [N, H, K, V] with K-last stride) diff --git a/benchmark/bench_linear_attention/bench_kda_prefill_cutedsl.py b/benchmark/bench_linear_attention/bench_kda_prefill_cutedsl.py index 08fe40003..e899c777d 100644 --- a/benchmark/bench_linear_attention/bench_kda_prefill_cutedsl.py +++ b/benchmark/bench_linear_attention/bench_kda_prefill_cutedsl.py @@ -29,20 +29,20 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "python") import torch import torch.nn.functional as F -from sglang.srt.layers.attention.fla.kda import chunk_kda, fused_recurrent_kda -from sglang.srt.layers.attention.linear.kernels.kda_blackwell import prepare_metadata -from sglang.srt.layers.attention.linear.kernels.kda_blackwell.kernel_h import ( +from sglang.kernels.ops.attention.linear.kda_blackwell import prepare_metadata +from sglang.kernels.ops.attention.linear.kda_blackwell.kernel_h import ( kda_h_cutedsl, ) -from sglang.srt.layers.attention.linear.kernels.kda_blackwell.kernel_kkt_inv_uw import ( +from sglang.kernels.ops.attention.linear.kda_blackwell.kernel_kkt_inv_uw import ( kkt_inv_uw_cutedsl, ) -from sglang.srt.layers.attention.linear.kernels.kda_blackwell.kernel_o import ( +from sglang.kernels.ops.attention.linear.kda_blackwell.kernel_o import ( kda_o_cutedsl, ) -from sglang.srt.layers.attention.linear.kernels.kda_blackwell.prologue import ( +from sglang.kernels.ops.attention.linear.kda_blackwell.prologue import ( kda_prologue, ) +from sglang.srt.layers.attention.fla.kda import chunk_kda, fused_recurrent_kda BT = 64 # chunk size diff --git a/python/sglang/kernels/ops/attention/__init__.py b/python/sglang/kernels/ops/attention/__init__.py index dc9dcf962..557e41f06 100644 --- a/python/sglang/kernels/ops/attention/__init__.py +++ b/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 [ diff --git a/python/sglang/kernels/ops/attention/linear/__init__.py b/python/sglang/kernels/ops/attention/linear/__init__.py new file mode 100644 index 000000000..f259447d5 --- /dev/null +++ b/python/sglang/kernels/ops/attention/linear/__init__.py @@ -0,0 +1 @@ +"""Linear-attention kernels (RFC #29630, Phase 2.5).""" diff --git a/python/sglang/srt/layers/attention/linear/kernels/gdn_blackwell/__init__.py b/python/sglang/kernels/ops/attention/linear/gdn_blackwell/__init__.py similarity index 100% rename from python/sglang/srt/layers/attention/linear/kernels/gdn_blackwell/__init__.py rename to python/sglang/kernels/ops/attention/linear/gdn_blackwell/__init__.py diff --git a/python/sglang/srt/layers/attention/linear/kernels/gdn_blackwell/kernel_h.py b/python/sglang/kernels/ops/attention/linear/gdn_blackwell/kernel_h.py similarity index 100% rename from python/sglang/srt/layers/attention/linear/kernels/gdn_blackwell/kernel_h.py rename to python/sglang/kernels/ops/attention/linear/gdn_blackwell/kernel_h.py diff --git a/python/sglang/srt/layers/attention/linear/kernels/gdn_blackwell/kernel_kkt_inv_uw.py b/python/sglang/kernels/ops/attention/linear/gdn_blackwell/kernel_kkt_inv_uw.py similarity index 100% rename from python/sglang/srt/layers/attention/linear/kernels/gdn_blackwell/kernel_kkt_inv_uw.py rename to python/sglang/kernels/ops/attention/linear/gdn_blackwell/kernel_kkt_inv_uw.py diff --git a/python/sglang/srt/layers/attention/linear/kernels/gdn_blackwell/kernel_o.py b/python/sglang/kernels/ops/attention/linear/gdn_blackwell/kernel_o.py similarity index 100% rename from python/sglang/srt/layers/attention/linear/kernels/gdn_blackwell/kernel_o.py rename to python/sglang/kernels/ops/attention/linear/gdn_blackwell/kernel_o.py diff --git a/python/sglang/srt/layers/attention/linear/kernels/kda_blackwell/__init__.py b/python/sglang/kernels/ops/attention/linear/kda_blackwell/__init__.py similarity index 100% rename from python/sglang/srt/layers/attention/linear/kernels/kda_blackwell/__init__.py rename to python/sglang/kernels/ops/attention/linear/kda_blackwell/__init__.py diff --git a/python/sglang/srt/layers/attention/linear/kernels/kda_blackwell/kernel_h.py b/python/sglang/kernels/ops/attention/linear/kda_blackwell/kernel_h.py similarity index 100% rename from python/sglang/srt/layers/attention/linear/kernels/kda_blackwell/kernel_h.py rename to python/sglang/kernels/ops/attention/linear/kda_blackwell/kernel_h.py diff --git a/python/sglang/srt/layers/attention/linear/kernels/kda_blackwell/kernel_kkt_inv_uw.py b/python/sglang/kernels/ops/attention/linear/kda_blackwell/kernel_kkt_inv_uw.py similarity index 100% rename from python/sglang/srt/layers/attention/linear/kernels/kda_blackwell/kernel_kkt_inv_uw.py rename to python/sglang/kernels/ops/attention/linear/kda_blackwell/kernel_kkt_inv_uw.py diff --git a/python/sglang/srt/layers/attention/linear/kernels/kda_blackwell/kernel_o.py b/python/sglang/kernels/ops/attention/linear/kda_blackwell/kernel_o.py similarity index 100% rename from python/sglang/srt/layers/attention/linear/kernels/kda_blackwell/kernel_o.py rename to python/sglang/kernels/ops/attention/linear/kda_blackwell/kernel_o.py diff --git a/python/sglang/srt/layers/attention/linear/kernels/kda_blackwell/prologue.py b/python/sglang/kernels/ops/attention/linear/kda_blackwell/prologue.py similarity index 100% rename from python/sglang/srt/layers/attention/linear/kernels/kda_blackwell/prologue.py rename to python/sglang/kernels/ops/attention/linear/kda_blackwell/prologue.py diff --git a/python/sglang/srt/layers/attention/linear/lightning_attn.py b/python/sglang/kernels/ops/attention/linear/lightning_attn.py similarity index 100% rename from python/sglang/srt/layers/attention/linear/lightning_attn.py rename to python/sglang/kernels/ops/attention/linear/lightning_attn.py diff --git a/python/sglang/srt/layers/attention/linear/seg_la.py b/python/sglang/kernels/ops/attention/linear/seg_la.py similarity index 100% rename from python/sglang/srt/layers/attention/linear/seg_la.py rename to python/sglang/kernels/ops/attention/linear/seg_la.py diff --git a/python/sglang/kernels/ops/attention/minimax_sparse/__init__.py b/python/sglang/kernels/ops/attention/minimax_sparse/__init__.py new file mode 100644 index 000000000..300109248 --- /dev/null +++ b/python/sglang/kernels/ops/attention/minimax_sparse/__init__.py @@ -0,0 +1 @@ +"""MiniMax sparse-attention kernels (RFC #29630, Phase 2.5).""" diff --git a/python/sglang/srt/layers/attention/minimax_sparse_ops/common/index.py b/python/sglang/kernels/ops/attention/minimax_sparse/common/index.py similarity index 100% rename from python/sglang/srt/layers/attention/minimax_sparse_ops/common/index.py rename to python/sglang/kernels/ops/attention/minimax_sparse/common/index.py diff --git a/python/sglang/srt/layers/attention/minimax_sparse_ops/common/utils.py b/python/sglang/kernels/ops/attention/minimax_sparse/common/utils.py similarity index 100% rename from python/sglang/srt/layers/attention/minimax_sparse_ops/common/utils.py rename to python/sglang/kernels/ops/attention/minimax_sparse/common/utils.py diff --git a/python/sglang/srt/layers/attention/minimax_sparse_ops/decode/flash_with_topk_idx.py b/python/sglang/kernels/ops/attention/minimax_sparse/decode/flash_with_topk_idx.py similarity index 100% rename from python/sglang/srt/layers/attention/minimax_sparse_ops/decode/flash_with_topk_idx.py rename to python/sglang/kernels/ops/attention/minimax_sparse/decode/flash_with_topk_idx.py diff --git a/python/sglang/srt/layers/attention/minimax_sparse_ops/decode/topk_sparse.py b/python/sglang/kernels/ops/attention/minimax_sparse/decode/topk_sparse.py similarity index 100% rename from python/sglang/srt/layers/attention/minimax_sparse_ops/decode/topk_sparse.py rename to python/sglang/kernels/ops/attention/minimax_sparse/decode/topk_sparse.py diff --git a/python/sglang/srt/layers/attention/minimax_sparse_ops/prefill/flash_with_topk_idx.py b/python/sglang/kernels/ops/attention/minimax_sparse/prefill/flash_with_topk_idx.py similarity index 100% rename from python/sglang/srt/layers/attention/minimax_sparse_ops/prefill/flash_with_topk_idx.py rename to python/sglang/kernels/ops/attention/minimax_sparse/prefill/flash_with_topk_idx.py diff --git a/python/sglang/srt/layers/attention/minimax_sparse_ops/prefill/topk_sparse.py b/python/sglang/kernels/ops/attention/minimax_sparse/prefill/topk_sparse.py similarity index 100% rename from python/sglang/srt/layers/attention/minimax_sparse_ops/prefill/topk_sparse.py rename to python/sglang/kernels/ops/attention/minimax_sparse/prefill/topk_sparse.py diff --git a/python/sglang/kernels/ops/diffusion/__init__.py b/python/sglang/kernels/ops/diffusion/__init__.py index 12c2d42bc..2ae0c0ee0 100644 --- a/python/sglang/kernels/ops/diffusion/__init__.py +++ b/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", + ) +) diff --git a/python/sglang/multimodal_gen/csrc/render/__init__.py b/python/sglang/kernels/ops/diffusion/render/__init__.py similarity index 100% rename from python/sglang/multimodal_gen/csrc/render/__init__.py rename to python/sglang/kernels/ops/diffusion/render/__init__.py diff --git a/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/__init__.py b/python/sglang/kernels/ops/diffusion/render/hunyuan3d_rasterizer/__init__.py similarity index 88% rename from python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/__init__.py rename to python/sglang/kernels/ops/diffusion/render/hunyuan3d_rasterizer/__init__.py index abb38de74..74ae6c262 100644 --- a/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/__init__.py +++ b/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) diff --git a/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/rasterizer.cpp b/python/sglang/kernels/ops/diffusion/render/hunyuan3d_rasterizer/rasterizer.cpp similarity index 100% rename from python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/rasterizer.cpp rename to python/sglang/kernels/ops/diffusion/render/hunyuan3d_rasterizer/rasterizer.cpp diff --git a/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/rasterizer.h b/python/sglang/kernels/ops/diffusion/render/hunyuan3d_rasterizer/rasterizer.h similarity index 100% rename from python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/rasterizer.h rename to python/sglang/kernels/ops/diffusion/render/hunyuan3d_rasterizer/rasterizer.h diff --git a/python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/rasterizer_gpu.cu b/python/sglang/kernels/ops/diffusion/render/hunyuan3d_rasterizer/rasterizer_gpu.cu similarity index 100% rename from python/sglang/multimodal_gen/csrc/render/hunyuan3d_rasterizer/rasterizer_gpu.cu rename to python/sglang/kernels/ops/diffusion/render/hunyuan3d_rasterizer/rasterizer_gpu.cu diff --git a/python/sglang/multimodal_gen/csrc/render/mesh_processor/__init__.py b/python/sglang/kernels/ops/diffusion/render/mesh_processor/__init__.py similarity index 89% rename from python/sglang/multimodal_gen/csrc/render/mesh_processor/__init__.py rename to python/sglang/kernels/ops/diffusion/render/mesh_processor/__init__.py index 03fb86410..9c6a9d67b 100644 --- a/python/sglang/multimodal_gen/csrc/render/mesh_processor/__init__.py +++ b/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"] diff --git a/python/sglang/multimodal_gen/csrc/render/mesh_processor/mesh_processor.cpp b/python/sglang/kernels/ops/diffusion/render/mesh_processor/mesh_processor.cpp similarity index 100% rename from python/sglang/multimodal_gen/csrc/render/mesh_processor/mesh_processor.cpp rename to python/sglang/kernels/ops/diffusion/render/mesh_processor/mesh_processor.cpp diff --git a/python/sglang/kernels/ops/diffusion/sparse_linear_attn_kernels.py b/python/sglang/kernels/ops/diffusion/sparse_linear_attn_kernels.py new file mode 100644 index 000000000..e7f41d894 --- /dev/null +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_linear_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_linear_attn.py index 793e2c52b..a55363908 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sparse_linear_attn.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/utils/mesh3d_utils.py b/python/sglang/multimodal_gen/runtime/utils/mesh3d_utils.py index 8d17eaf9d..3ee1e49c7 100644 --- a/python/sglang/multimodal_gen/runtime/utils/mesh3d_utils.py +++ b/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() diff --git a/python/sglang/srt/layers/attention/linear/kernels/gdn_cutedsl.py b/python/sglang/srt/layers/attention/linear/kernels/gdn_cutedsl.py index 311f40093..80e271207 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/gdn_cutedsl.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/linear/kernels/kda_cutedsl.py b/python/sglang/srt/layers/attention/linear/kernels/kda_cutedsl.py index b8583ada8..ed3df69d1 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/kda_cutedsl.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/linear/lightning_backend.py b/python/sglang/srt/layers/attention/linear/lightning_backend.py index ce343ba25..8013cbb2b 100644 --- a/python/sglang/srt/layers/attention/linear/lightning_backend.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/minimax_sparse_ops/minimax_sparse.py b/python/sglang/srt/layers/attention/minimax_sparse_ops/minimax_sparse.py index 046d7ff7c..ab8231a5d 100644 --- a/python/sglang/srt/layers/attention/minimax_sparse_ops/minimax_sparse.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/minimax_sparse_ops/tests/test_flash_with_topk_idx.py b/python/sglang/srt/layers/attention/minimax_sparse_ops/tests/test_flash_with_topk_idx.py index 4438a5632..837aa221d 100644 --- a/python/sglang/srt/layers/attention/minimax_sparse_ops/tests/test_flash_with_topk_idx.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/minimax_sparse_ops/tests/test_sparse_gqa.py b/python/sglang/srt/layers/attention/minimax_sparse_ops/tests/test_sparse_gqa.py index 0b2ec7e76..92d6dd21a 100644 --- a/python/sglang/srt/layers/attention/minimax_sparse_ops/tests/test_sparse_gqa.py +++ b/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, ) diff --git a/test/registered/attention/test_gdn_prefill_cutedsl.py b/test/registered/attention/test_gdn_prefill_cutedsl.py index cefaa78e5..975b09e29 100644 --- a/test/registered/attention/test_gdn_prefill_cutedsl.py +++ b/test/registered/attention/test_gdn_prefill_cutedsl.py @@ -23,6 +23,10 @@ if not (torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 1 allow_module_level=True, ) +from sglang.kernels.ops.attention.linear.gdn_blackwell import ( # noqa: E402 + chunk_gated_delta_rule_cutedsl, + prepare_metadata_cutedsl, +) from sglang.srt.layers.attention.fla.fused_recurrent import ( # noqa: E402 fused_recurrent_gated_delta_rule, ) @@ -30,10 +34,6 @@ from sglang.srt.layers.attention.fla.index import ( # noqa: E402 prepare_chunk_indices, prepare_chunk_offsets, ) -from sglang.srt.layers.attention.linear.kernels.gdn_blackwell import ( # noqa: E402 - chunk_gated_delta_rule_cutedsl, - prepare_metadata_cutedsl, -) @pytest.mark.parametrize("num_seqs", [1, 5, 257]) diff --git a/test/registered/attention/test_kda_prefill_cutedsl.py b/test/registered/attention/test_kda_prefill_cutedsl.py index 96ffc08bb..83fa7eae6 100644 --- a/test/registered/attention/test_kda_prefill_cutedsl.py +++ b/test/registered/attention/test_kda_prefill_cutedsl.py @@ -25,15 +25,15 @@ if not (torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 1 allow_module_level=True, ) +from sglang.kernels.ops.attention.linear.kda_blackwell import ( # noqa: E402 + chunk_kda_cutedsl, + prepare_metadata, +) from sglang.srt.layers.attention.fla.index import ( # noqa: E402 prepare_chunk_indices, prepare_chunk_offsets, ) from sglang.srt.layers.attention.fla.kda import fused_recurrent_kda # noqa: E402 -from sglang.srt.layers.attention.linear.kernels.kda_blackwell import ( # noqa: E402 - chunk_kda_cutedsl, - prepare_metadata, -) def _l2norm(x: torch.Tensor) -> torch.Tensor: diff --git a/test/registered/jit/benchmark/minimax/bench_minimax_decode_topk.py b/test/registered/jit/benchmark/minimax/bench_minimax_decode_topk.py index ec98c465c..9ac8b9d18 100644 --- a/test/registered/jit/benchmark/minimax/bench_minimax_decode_topk.py +++ b/test/registered/jit/benchmark/minimax/bench_minimax_decode_topk.py @@ -11,7 +11,7 @@ import triton from sglang.jit_kernel.benchmark import marker from sglang.jit_kernel.minimax_decode_topk import minimax_decode_topk -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 ( _topk_index_merge_kernel, _topk_index_partial_kernel, ) diff --git a/test/registered/jit/minimax/test_minimax_decode_topk_page_table.py b/test/registered/jit/minimax/test_minimax_decode_topk_page_table.py index 55a825e1a..48a1039ec 100644 --- a/test/registered/jit/minimax/test_minimax_decode_topk_page_table.py +++ b/test/registered/jit/minimax/test_minimax_decode_topk_page_table.py @@ -20,7 +20,7 @@ from sglang.jit_kernel.minimax_decode_topk import ( minimax_decode_topk, minimax_decode_topk_page_table, ) -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, ) from sglang.test.ci.ci_register import register_cuda_ci