From 5996b54bd3e8aca60bb417ba05f6266d1a9851a1 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Fri, 26 Jun 2026 15:06:49 +0800 Subject: [PATCH] [KDA-Pilot] Add diffusion causal Conv3D cat-pad CUDA fast path for Cosmos3 (#29281) Co-authored-by: Claude Opus 4.8 --- .../csrc/diffusion/causal_conv3d_cat_pad.cuh | 255 ++++++++++++++++++ .../diffusion/causal_conv3d_cat_pad.py | 141 ++++++++++ .../runtime/layers/parallel_conv.py | 41 ++- .../stages/model_specific_stages/cosmos3.py | 2 +- .../diffusion/bench_causal_conv3d_cat_pad.py | 93 +++++++ .../diffusion/test_causal_conv3d_cat_pad.py | 88 ++++++ 6 files changed, 617 insertions(+), 3 deletions(-) create mode 100644 python/sglang/jit_kernel/csrc/diffusion/causal_conv3d_cat_pad.cuh create mode 100644 python/sglang/jit_kernel/diffusion/causal_conv3d_cat_pad.py create mode 100644 test/registered/jit/benchmark/diffusion/bench_causal_conv3d_cat_pad.py create mode 100644 test/registered/jit/diffusion/test_causal_conv3d_cat_pad.py diff --git a/python/sglang/jit_kernel/csrc/diffusion/causal_conv3d_cat_pad.cuh b/python/sglang/jit_kernel/csrc/diffusion/causal_conv3d_cat_pad.cuh new file mode 100644 index 000000000..e5b447318 --- /dev/null +++ b/python/sglang/jit_kernel/csrc/diffusion/causal_conv3d_cat_pad.cuh @@ -0,0 +1,255 @@ +// Native CUDA fast path for Cosmos3 VAE causal-Conv3D cat/pad copy. +// +// The op writes the output of: +// pad(cat(cache_x, x, dim=T), (Wl, Wr, Ht, Hb, Dl - cache_t, Dr)) +// for 5D NCTHW tensors. It is a memory-bound copy/zero-fill kernel and is only +// entered for contiguous CUDA tensors; unsupported cases fall back to Triton in +// the Python caller. +// +// Developed with MIT HAN Lab Kernel Design Agents: +// https://github.com/mit-han-lab/kernel-design-agents + +#pragma once + +#include // For TensorMatcher, SymbolicSize, SymbolicDevice +#include // For RuntimeCheck, div_ceil + +#include // For LaunchKernel + +#include + +namespace sglang_causal_conv3d_cat_pad { + +namespace { + +constexpr int kBlockSize = 256; + +template +__global__ void __launch_bounds__(kBlockSize) cat_pad_flat_kernel( + const ET* __restrict__ x, + const ET* __restrict__ cache, + ET* __restrict__ out, + int64_t total_vecs, + int64_t channels, + int64_t t_size, + int64_t h_size, + int64_t w_size, + int64_t cache_t, + int64_t out_t, + int64_t out_h, + int64_t out_w, + int64_t pad_d_left, + int64_t pad_h_top, + int64_t pad_w_left) { + union Pack { + ET elem[kVec]; + uint4 raw; + }; + + const int64_t nthreads = static_cast(gridDim.x) * blockDim.x; + for (int64_t vid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; vid < total_vecs; vid += nthreads) { + int64_t base = vid * kVec; + int64_t ow = base % out_w; + int64_t tmp = base / out_w; + int64_t oh = tmp % out_h; + tmp /= out_h; + int64_t od = tmp % out_t; + tmp /= out_t; + int64_t oc = tmp % channels; + int64_t ob = tmp / channels; + + int64_t ih = oh - pad_h_top; + int64_t src_t = od - pad_d_left; + bool interior = ih >= 0 && ih < h_size && src_t >= 0 && src_t < cache_t + t_size; + + const ET* src = nullptr; + if (interior) { + if (src_t < cache_t) { + src = cache + (((ob * channels + oc) * cache_t + src_t) * h_size + ih) * w_size; + } else { + src = x + (((ob * channels + oc) * t_size + (src_t - cache_t)) * h_size + ih) * w_size; + } + } + + Pack pack; +#pragma unroll + for (int i = 0; i < kVec; ++i) { + ET value = ET(0); + if (interior) { + const int64_t iw = ow - pad_w_left; + if (iw >= 0 && iw < w_size) { + value = SGLANG_LDG(src + iw); + } + } + pack.elem[i] = value; + + if (++ow == out_w) { + ow = 0; + if (++oh == out_h) { + oh = 0; + if (++od == out_t) { + od = 0; + if (++oc == channels) { + oc = 0; + ++ob; + } + } + } + ih = oh - pad_h_top; + src_t = od - pad_d_left; + interior = ih >= 0 && ih < h_size && src_t >= 0 && src_t < cache_t + t_size; + if (interior) { + if (src_t < cache_t) { + src = cache + (((ob * channels + oc) * cache_t + src_t) * h_size + ih) * w_size; + } else { + src = x + (((ob * channels + oc) * t_size + (src_t - cache_t)) * h_size + ih) * w_size; + } + } else { + src = nullptr; + } + } + } + + reinterpret_cast(out)[vid] = pack.raw; + } +} + +template +void launch_cat_pad_flat( + const void* x, + const void* cache, + void* out, + int64_t total, + int64_t channels, + int64_t t_size, + int64_t h_size, + int64_t w_size, + int64_t cache_t, + int64_t out_t, + int64_t out_h, + int64_t out_w, + int64_t depth_left, + int64_t pad_h_top, + int64_t pad_w_left, + DLDevice device) { + const int64_t total_vecs = total / kVec; + const uint32_t grid = static_cast(host::div_ceil(total_vecs, static_cast(kBlockSize))); + host::LaunchKernel(grid, kBlockSize, device)( + cat_pad_flat_kernel, + static_cast(x), + static_cast(cache), + static_cast(out), + total_vecs, + channels, + t_size, + h_size, + w_size, + cache_t, + out_t, + out_h, + out_w, + depth_left, + pad_h_top, + pad_w_left); +} + +} // namespace + +template +struct CausalConv3dCatPadKernel { + static void + run(tvm::ffi::TensorView out, + tvm::ffi::TensorView x, + tvm::ffi::TensorView cache, + int64_t pad_w_left, + int64_t pad_w_right, + int64_t pad_h_top, + int64_t pad_h_bottom, + int64_t pad_d_left, + int64_t pad_d_right) { + using namespace host; + + auto bsz = SymbolicSize{"batch"}; + auto channels = SymbolicSize{"channels"}; + auto t_size = SymbolicSize{"t_size"}; + auto h_size = SymbolicSize{"h_size"}; + auto w_size = SymbolicSize{"w_size"}; + auto cache_t = SymbolicSize{"cache_t"}; + auto out_t = SymbolicSize{"out_t"}; + auto out_h = SymbolicSize{"out_h"}; + auto out_w = SymbolicSize{"out_w"}; + auto device = SymbolicDevice{}; + device.set_options(); + + TensorMatcher({bsz, channels, t_size, h_size, w_size}) + .with_dtype() + .template with_device(device) + .verify(x); + TensorMatcher({bsz, channels, cache_t, h_size, w_size}) + .with_dtype() + .template with_device(device) + .verify(cache); + TensorMatcher({bsz, channels, out_t, out_h, out_w}) + .with_dtype() + .template with_device(device) + .verify(out); + + const int64_t depth_left = pad_d_left - cache_t.unwrap(); + RuntimeCheck(depth_left >= 0, "pad_d_left must be >= cache_t"); + RuntimeCheck(pad_d_right == 0, "pad_d_right must be 0"); + RuntimeCheck(pad_w_left == pad_w_right, "width padding must be symmetric"); + RuntimeCheck(pad_h_top == pad_h_bottom, "height padding must be symmetric"); + RuntimeCheck(out_t.unwrap() == t_size.unwrap() + cache_t.unwrap() + depth_left + pad_d_right, "out_t mismatch"); + RuntimeCheck(out_h.unwrap() == h_size.unwrap() + pad_h_top + pad_h_bottom, "out_h mismatch"); + RuntimeCheck(out_w.unwrap() == w_size.unwrap() + pad_w_left + pad_w_right, "out_w mismatch"); + + const int64_t total = bsz.unwrap() * channels.unwrap() * out_t.unwrap() * out_h.unwrap() * out_w.unwrap(); + if (total == 0) { + return; + } + + constexpr int kVec = 16 / sizeof(T); + RuntimeCheck(total % kVec == 0, "output element count must be divisible by vector width"); + RuntimeCheck(reinterpret_cast(out.data_ptr()) % 16 == 0, "output pointer must be 16-byte aligned"); + + if constexpr (sizeof(T) == 2) { + launch_cat_pad_flat( + x.data_ptr(), + cache.data_ptr(), + out.data_ptr(), + total, + channels.unwrap(), + t_size.unwrap(), + h_size.unwrap(), + w_size.unwrap(), + cache_t.unwrap(), + out_t.unwrap(), + out_h.unwrap(), + out_w.unwrap(), + depth_left, + pad_h_top, + pad_w_left, + device.unwrap()); + } else { + launch_cat_pad_flat( + x.data_ptr(), + cache.data_ptr(), + out.data_ptr(), + total, + channels.unwrap(), + t_size.unwrap(), + h_size.unwrap(), + w_size.unwrap(), + cache_t.unwrap(), + out_t.unwrap(), + out_h.unwrap(), + out_w.unwrap(), + depth_left, + pad_h_top, + pad_w_left, + device.unwrap()); + } + } +}; + +} // namespace sglang_causal_conv3d_cat_pad diff --git a/python/sglang/jit_kernel/diffusion/causal_conv3d_cat_pad.py b/python/sglang/jit_kernel/diffusion/causal_conv3d_cat_pad.py new file mode 100644 index 000000000..60ff5626d --- /dev/null +++ b/python/sglang/jit_kernel/diffusion/causal_conv3d_cat_pad.py @@ -0,0 +1,141 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args +from sglang.srt.utils.custom_op import register_custom_op + +if TYPE_CHECKING: + from tvm_ffi.module import Module + + +_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32) + + +@cache_once +def _jit_causal_conv3d_cat_pad_module(dtype: torch.dtype) -> Module: + args = make_cpp_args(dtype) + return load_jit( + "diffusion_causal_conv3d_cat_pad", + *args, + cuda_files=["diffusion/causal_conv3d_cat_pad.cuh"], + cuda_wrappers=[ + ( + "causal_conv3d_cat_pad", + "sglang_causal_conv3d_cat_pad::" + f"CausalConv3dCatPadKernel<{args}>::run", + ) + ], + ) + + +def _causal_conv3d_cat_pad_fake_impl( + x: torch.Tensor, + cache_x: torch.Tensor, + pad_w_left: int, + pad_w_right: int, + pad_h_top: int, + pad_h_bottom: int, + pad_d_left: int, + pad_d_right: int, +) -> torch.Tensor: + cache_t = cache_x.shape[2] + depth_left = pad_d_left - cache_t + return torch.empty( + ( + x.shape[0], + x.shape[1], + x.shape[2] + cache_t + depth_left + pad_d_right, + x.shape[3] + pad_h_top + pad_h_bottom, + x.shape[4] + pad_w_left + pad_w_right, + ), + device=x.device, + dtype=x.dtype, + ) + + +@register_custom_op( + op_name="diffusion_causal_conv3d_cat_pad", + mutates_args=[], + fake_impl=_causal_conv3d_cat_pad_fake_impl, +) +def _causal_conv3d_cat_pad_custom_op( + x: torch.Tensor, + cache_x: torch.Tensor, + pad_w_left: int, + pad_w_right: int, + pad_h_top: int, + pad_h_bottom: int, + pad_d_left: int, + pad_d_right: int, +) -> torch.Tensor: + out = _causal_conv3d_cat_pad_fake_impl( + x, + cache_x, + pad_w_left, + pad_w_right, + pad_h_top, + pad_h_bottom, + pad_d_left, + pad_d_right, + ) + module = _jit_causal_conv3d_cat_pad_module(x.dtype) + module.causal_conv3d_cat_pad( + out, + x, + cache_x, + pad_w_left, + pad_w_right, + pad_h_top, + pad_h_bottom, + pad_d_left, + pad_d_right, + ) + return out + + +def fused_causal_conv3d_cat_pad_cuda( + x: torch.Tensor, + cache_x: torch.Tensor, + padding: list[int] | tuple[int, ...], +) -> torch.Tensor: + if x.dtype not in _SUPPORTED_DTYPES: + raise RuntimeError(f"unsupported dtype for causal Conv3D cat/pad: {x.dtype}") + if not torch.compiler.is_compiling(): + if ( + not x.is_cuda + or not cache_x.is_cuda + or x.dim() != 5 + or cache_x.dim() != 5 + or not x.is_contiguous() + or not cache_x.is_contiguous() + or not can_use_fused_causal_conv3d_cat_pad_cuda(x, cache_x, padding) + ): + raise RuntimeError("unsupported input for causal Conv3D cat/pad CUDA") + return _causal_conv3d_cat_pad_custom_op(x, cache_x, *padding) + + +def can_use_fused_causal_conv3d_cat_pad_cuda( + x: torch.Tensor, + cache_x: torch.Tensor, + padding: list[int] | tuple[int, ...], +) -> bool: + if x.dtype not in _SUPPORTED_DTYPES: + return False + pad_w_left, pad_w_right, pad_h_top, pad_h_bottom, pad_d_left, pad_d_right = padding + cache_t = cache_x.shape[2] + depth_left = pad_d_left - cache_t + if depth_left < 0 or pad_d_right != 0: + return False + out_numel = ( + x.shape[0] + * x.shape[1] + * (x.shape[2] + cache_t + depth_left + pad_d_right) + * (x.shape[3] + pad_h_top + pad_h_bottom) + * (x.shape[4] + pad_w_left + pad_w_right) + ) + elem_size = 4 if x.dtype == torch.float32 else 2 + vec_elems = 16 // elem_size + return out_numel % vec_elems == 0 diff --git a/python/sglang/multimodal_gen/runtime/layers/parallel_conv.py b/python/sglang/multimodal_gen/runtime/layers/parallel_conv.py index 780d354c9..b35f181db 100644 --- a/python/sglang/multimodal_gen/runtime/layers/parallel_conv.py +++ b/python/sglang/multimodal_gen/runtime/layers/parallel_conv.py @@ -13,13 +13,50 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_decode_parallel_world_size, ) from sglang.multimodal_gen.runtime.platforms import current_platform +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) if current_platform.is_cuda(): + from sglang.jit_kernel.diffusion.causal_conv3d_cat_pad import ( + can_use_fused_causal_conv3d_cat_pad_cuda, + fused_causal_conv3d_cat_pad_cuda, + ) from sglang.jit_kernel.diffusion.triton.causal_conv3d_pad import ( - fused_causal_conv3d_cat_pad, + fused_causal_conv3d_cat_pad as fused_causal_conv3d_cat_pad_triton, ) else: - fused_causal_conv3d_cat_pad = None + can_use_fused_causal_conv3d_cat_pad_cuda = None + fused_causal_conv3d_cat_pad_cuda = None + fused_causal_conv3d_cat_pad_triton = None + + +_causal_conv3d_cat_pad_cuda_failed = False + + +def fused_causal_conv3d_cat_pad( + x: torch.Tensor, + cache_x: torch.Tensor, + padding: list[int], +) -> torch.Tensor: + global _causal_conv3d_cat_pad_cuda_failed + if ( + fused_causal_conv3d_cat_pad_cuda is not None + and can_use_fused_causal_conv3d_cat_pad_cuda(x, cache_x, padding) + and not _causal_conv3d_cat_pad_cuda_failed + ): + try: + return fused_causal_conv3d_cat_pad_cuda(x, cache_x, padding) + except Exception: + logger.warning( + "fused_causal_conv3d_cat_pad_cuda failed, falling back to Triton", + exc_info=True, + ) + _causal_conv3d_cat_pad_cuda_failed = True + if fused_causal_conv3d_cat_pad_triton is None: + raise RuntimeError("causal Conv3D cat/pad fusion is only available on CUDA") + return fused_causal_conv3d_cat_pad_triton(x, cache_x, padding) + _SPATIAL_PARALLEL_DECODE_DISABLED = contextvars.ContextVar( "spatial_parallel_decode_disabled", default=False diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py index 0a16be28e..1d21f5e26 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py @@ -464,7 +464,7 @@ class Cosmos3DenoisingStage(PipelineStage): compile_kwargs = { "mode": "default", "fullgraph": False, - "dynamic": True, + "dynamic": False, } gen_layers = getattr(transformer, "gen_layers", None) diff --git a/test/registered/jit/benchmark/diffusion/bench_causal_conv3d_cat_pad.py b/test/registered/jit/benchmark/diffusion/bench_causal_conv3d_cat_pad.py new file mode 100644 index 000000000..229eeeffa --- /dev/null +++ b/test/registered/jit/benchmark/diffusion/bench_causal_conv3d_cat_pad.py @@ -0,0 +1,93 @@ +from dataclasses import dataclass + +import torch + +from sglang.jit_kernel.benchmark import marker +from sglang.jit_kernel.diffusion.causal_conv3d_cat_pad import ( + fused_causal_conv3d_cat_pad_cuda, +) +from sglang.jit_kernel.diffusion.triton.causal_conv3d_pad import ( + fused_causal_conv3d_cat_pad as fused_causal_conv3d_cat_pad_triton, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci( + est_time=20, + suite="base-b-kernel-benchmark-1-gpu-large", + disabled="standalone benchmark", +) + +DEVICE = "cuda" +DTYPE = torch.bfloat16 + + +@dataclass(frozen=True) +class Case: + name: str + channels: int + t_size: int + h_size: int + w_size: int + cache_t: int + trace_count: int + + +CASES = [ + Case("c1024_t1_h30_w52_cache1", 1024, 1, 30, 52, 1, 8), + Case("c1024_t1_h30_w52_cache2", 1024, 1, 30, 52, 2, 8), + Case("c1024_t2_h60_w104_cache1", 1024, 2, 60, 104, 1, 5), + Case("c1024_t2_h60_w104_cache2", 1024, 2, 60, 104, 2, 5), + Case("c512_t4_h120_w208_cache1", 512, 4, 120, 208, 1, 5), + Case("c512_t4_h120_w208_cache2", 512, 4, 120, 208, 2, 5), + Case("c256_t4_h240_w416_cache1", 256, 4, 240, 416, 1, 6), + Case("c256_t4_h240_w416_cache2", 256, 4, 240, 416, 2, 6), +] +CASE_BY_NAME = {case.name: case for case in CASES} +CASE_NAMES = [case.name for case in CASES] + + +def make_inputs(case: Case) -> tuple[torch.Tensor, torch.Tensor, tuple[int, ...]]: + generator = torch.Generator(device=DEVICE) + generator.manual_seed(case.channels * 1009 + case.t_size * 251 + case.cache_t) + x = torch.randn( + (1, case.channels, case.t_size, case.h_size, case.w_size), + device=DEVICE, + dtype=DTYPE, + generator=generator, + ) + cache_x = torch.randn( + (1, case.channels, case.cache_t, case.h_size, case.w_size), + device=DEVICE, + dtype=DTYPE, + generator=generator, + ) + padding = (1, 1, 1, 1, case.cache_t, 0) + return x, cache_x, padding + + +@marker.parametrize("case_name", CASE_NAMES, ci_vals=CASE_NAMES[:2]) +@marker.benchmark("provider", ["triton", "cuda"]) +def benchmark(case_name: str, provider: str) -> marker.BenchResult: + case = CASE_BY_NAME[case_name] + x, cache_x, padding = make_inputs(case) + fn = ( + fused_causal_conv3d_cat_pad_triton + if provider == "triton" + else fused_causal_conv3d_cat_pad_cuda + ) + actual = fused_causal_conv3d_cat_pad_cuda(x, cache_x, padding) + expected = fused_causal_conv3d_cat_pad_triton(x, cache_x, padding) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + return marker.do_bench( + fn, + input_args=(x, cache_x, padding), + use_cuda_graph=False, + replay_iters=200, + graph_clone_args=(0, 1), + memory_args=(x, cache_x), + memory_output="out", + ) + + +if __name__ == "__main__": + benchmark.run() diff --git a/test/registered/jit/diffusion/test_causal_conv3d_cat_pad.py b/test/registered/jit/diffusion/test_causal_conv3d_cat_pad.py new file mode 100644 index 000000000..68afd07ac --- /dev/null +++ b/test/registered/jit/diffusion/test_causal_conv3d_cat_pad.py @@ -0,0 +1,88 @@ +import sys + +import pytest +import torch + +from sglang.jit_kernel.diffusion.causal_conv3d_cat_pad import ( + fused_causal_conv3d_cat_pad_cuda, +) +from sglang.jit_kernel.diffusion.triton.causal_conv3d_pad import ( + fused_causal_conv3d_cat_pad as fused_causal_conv3d_cat_pad_triton, +) +from sglang.jit_kernel.utils import get_ci_test_range +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=45, suite="base-b-kernel-unit-1-gpu-large") +register_cuda_ci(est_time=45, suite="base-b-kernel-unit-1-gpu-b200") + +DEVICE = "cuda" +DTYPE = torch.bfloat16 + +COSMOS3_CASES = get_ci_test_range( + [ + (1024, 1, 30, 52, 1), + (1024, 1, 30, 52, 2), + (1024, 2, 60, 104, 1), + (1024, 2, 60, 104, 2), + (512, 4, 120, 208, 1), + (512, 4, 120, 208, 2), + (256, 4, 240, 416, 1), + (256, 4, 240, 416, 2), + ], + [(1024, 1, 30, 52, 1), (512, 4, 120, 208, 2)], +) + + +def _make_inputs( + channels: int, + t_size: int, + h_size: int, + w_size: int, + cache_t: int, +) -> tuple[torch.Tensor, torch.Tensor, tuple[int, ...]]: + generator = torch.Generator(device=DEVICE) + generator.manual_seed(channels * 1009 + t_size * 251 + h_size + cache_t) + x = torch.randn( + (1, channels, t_size, h_size, w_size), + device=DEVICE, + dtype=DTYPE, + generator=generator, + ) + cache_x = torch.randn( + (1, channels, cache_t, h_size, w_size), + device=DEVICE, + dtype=DTYPE, + generator=generator, + ) + padding = (1, 1, 1, 1, cache_t, 0) + return x, cache_x, padding + + +@pytest.mark.parametrize("channels,t_size,h_size,w_size,cache_t", COSMOS3_CASES) +def test_causal_conv3d_cat_pad( + channels: int, + t_size: int, + h_size: int, + w_size: int, + cache_t: int, +) -> None: + x, cache_x, padding = _make_inputs(channels, t_size, h_size, w_size, cache_t) + actual = fused_causal_conv3d_cat_pad_cuda(x, cache_x, padding) + expected = fused_causal_conv3d_cat_pad_triton(x, cache_x, padding) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +def test_causal_conv3d_cat_pad_torch_compile() -> None: + x, cache_x, padding = _make_inputs(1024, 1, 30, 52, 1) + + @torch.compile(fullgraph=True) + def fn(x: torch.Tensor, cache_x: torch.Tensor) -> torch.Tensor: + return fused_causal_conv3d_cat_pad_cuda(x, cache_x, padding) + + actual = fn(x, cache_x) + expected = fused_causal_conv3d_cat_pad_triton(x, cache_x, padding) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__]))