[KDA-Pilot] Add diffusion causal Conv3D cat-pad CUDA fast path for Cosmos3 (#29281)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

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()

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__]))