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

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

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

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

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

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)

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