[KDA-Pilot] Add diffusion causal Conv3D cat-pad CUDA fast path for Cosmos3 (#29281)
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Xiaoyu Zhang
co-authored by
Claude Opus 4.8
Claude Opus 4.8
parent
a10a24e9a7
commit
5996b54bd3
@@ -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 <sgl_kernel/tensor.h> // For TensorMatcher, SymbolicSize, SymbolicDevice
|
||||
#include <sgl_kernel/utils.h> // For RuntimeCheck, div_ceil
|
||||
|
||||
#include <sgl_kernel/utils.cuh> // For LaunchKernel
|
||||
|
||||
#include <cstdint>
|
||||
|
||||
namespace sglang_causal_conv3d_cat_pad {
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr int kBlockSize = 256;
|
||||
|
||||
template <typename ET, int kVec>
|
||||
__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<int64_t>(gridDim.x) * blockDim.x;
|
||||
for (int64_t vid = static_cast<int64_t>(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<uint4*>(out)[vid] = pack.raw;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename ET, int kVec>
|
||||
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<uint32_t>(host::div_ceil(total_vecs, static_cast<int64_t>(kBlockSize)));
|
||||
host::LaunchKernel(grid, kBlockSize, device)(
|
||||
cat_pad_flat_kernel<ET, kVec>,
|
||||
static_cast<const ET*>(x),
|
||||
static_cast<const ET*>(cache),
|
||||
static_cast<ET*>(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 <typename T>
|
||||
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<kDLCUDA>();
|
||||
|
||||
TensorMatcher({bsz, channels, t_size, h_size, w_size})
|
||||
.with_dtype<T>()
|
||||
.template with_device<kDLCUDA>(device)
|
||||
.verify(x);
|
||||
TensorMatcher({bsz, channels, cache_t, h_size, w_size})
|
||||
.with_dtype<T>()
|
||||
.template with_device<kDLCUDA>(device)
|
||||
.verify(cache);
|
||||
TensorMatcher({bsz, channels, out_t, out_h, out_w})
|
||||
.with_dtype<T>()
|
||||
.template with_device<kDLCUDA>(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<uintptr_t>(out.data_ptr()) % 16 == 0, "output pointer must be 16-byte aligned");
|
||||
|
||||
if constexpr (sizeof(T) == 2) {
|
||||
launch_cat_pad_flat<uint16_t, kVec>(
|
||||
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<uint32_t, kVec>(
|
||||
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
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
+1
-1
@@ -464,7 +464,7 @@ class Cosmos3DenoisingStage(PipelineStage):
|
||||
compile_kwargs = {
|
||||
"mode": "default",
|
||||
"fullgraph": False,
|
||||
"dynamic": True,
|
||||
"dynamic": False,
|
||||
}
|
||||
|
||||
gen_layers = getattr(transformer, "gen_layers", None)
|
||||
|
||||
Reference in New Issue
Block a user