add-cuda-kernel

name	add-cuda-kernel
description	Step-by-step tutorial for adding new CUDA kernels to Oasr

Tutorial: Adding a New Kernel to Oasr

This tutorial walks through adding a simple element-wise scale operation to Oasr. We'll implement scale(x, factor) = x * factor to demonstrate the complete workflow, with references to real kernels (norm, activation, conv, gemm) throughout.

Goal

Add a new operation that scales each element of a tensor by a scalar factor:

• Input: tensor x and scalar factor
• Output: x * factor (element-wise)
• Support multiple dtypes (FP16, BF16, FP32)

Step 1: Define CUDA Kernel in include/

Create include/oasr/scale.cuh:

#pragma once
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cuda_bf16.h>

namespace oasr {

/*!
 * \brief Element-wise scale kernel
 * \tparam T Data type (half, __nv_bfloat16, float)
 * \param input Input tensor
 * \param output Output tensor
 * \param factor Scale factor
 * \param n Number of elements
 */
template <typename T>
__global__ void ScaleKernel(const T* input, T* output, T factor, int n) {
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx < n) {
    output[idx] = input[idx] * factor;
  }
}

/*!
 * \brief Launch scale kernel
 * \tparam T Data type
 * \param input Input pointer
 * \param output Output pointer
 * \param factor Scale factor
 * \param n Number of elements
 * \param stream CUDA stream
 */
template <typename T>
cudaError_t ScaleLauncher(const T* input, T* output, T factor, int n,
                          cudaStream_t stream = nullptr) {
  const int threads = 256;
  const int blocks = (n + threads - 1) / threads;

  ScaleKernel<T><<<blocks, threads, 0, stream>>>(input, output, factor, n);

  return cudaGetLastError();
}

}  // namespace oasr

Key points:

• Framework-agnostic (no Torch headers)
• Uses raw pointers
• Template-based for dtype flexibility
• Only includes what's needed (cuda_runtime, cuda_fp16, cuda_bf16)
• Namespace follows oasr:: pattern (or oasr::<family>:: for larger families, e.g. oasr::norm::, oasr::activation::)

Real examples:

• include/oasr/activation.cuh -- oasr::activation::GLU<T>(), oasr::activation::Swish<T>()
• include/oasr/norm.cuh -- oasr::norm::LayerNorm<T>(), oasr::norm::RMSNorm<T>()
• include/oasr/conv/conv1d.cuh -- Depthwise/pointwise conv1d kernels
• include/oasr/gemm/gemm.cuh -- CUTLASS GEMM kernels

Step 2: Create Launcher in csrc/

Create csrc/scale.cu:

#include <oasr/scale.cuh>

#include "tvm_ffi_utils.h"

using namespace oasr;

void scale_run(TensorView output, TensorView input, double factor) {
  CHECK_INPUT(input);
  CHECK_INPUT(output);

  int n = 1;
  for (int i = 0; i < input.ndim(); ++i) {
    n *= input.size(i);
  }

  cudaStream_t stream = get_stream(input.device());

  DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP16(input.dtype(), c_type, [&] {
    cudaError_t status = ScaleLauncher<c_type>(
      static_cast<const c_type*>(input.data_ptr()),
      static_cast<c_type*>(output.data_ptr()),
      static_cast<c_type>(factor),
      n,
      stream
    );
    TVM_FFI_ICHECK(status == cudaSuccess)
        << "Failed to run ScaleLauncher: " << cudaGetErrorString(status);
    return true;
  });
}

Key points:

• Includes kernel header from include/oasr/ and "tvm_ffi_utils.h" (TVM-FFI utils only in csrc/)
• Uses TensorView (alias for tvm::ffi::TensorView) as tensor type
• Uses Optional (alias for tvm::ffi::Optional<TensorView>) for optional tensors
• Output tensor is the first parameter in C++ launchers (destination-passing convention)
• Uses CHECK_INPUT(x) macro to verify tensor is on CUDA
• Gets CUDA stream via get_stream(device)
• Dispatches on dtype with DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP16 (handles FP32/FP16/BF16)
• Converts TensorView to raw pointers with static_cast<T*>(x.data_ptr())
• Checks kernel result with TVM_FFI_ICHECK
• Add descriptive error messages with << operator

Available validation macros (from csrc/tvm_ffi_utils.h):

Macro	Purpose
CHECK_INPUT(x)	Verify tensor is on CUDA
CHECK_DIM(expected, x)	Verify dimensionality
CHECK_DEVICE(x, y)	Same-device check
CHECK_LAST_DIM_CONTIGUOUS_INPUT(x)	Contiguity check

Available dispatch macros:

Macro	Dtypes
DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP16(dtype, c_type, ...)	FP32, FP16, BF16
DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP32_FP16(dtype, c_type, ...)	Alias for the above
DISPATCH_DLPACK_HALF_DTYPE(dtype, c_type, ...)	FP16, BF16 only (for GEMM/conv)

TVM-FFI Error Handling:

• TVM_FFI_ICHECK(condition) << "message" -- Assert with error message (used inside dispatch macros or when you need a simple assertion)
• TVM_FFI_THROW(ValueError) << "message" -- Throw ValueError with custom message (standard runtime error handling)
• TVM_FFI_THROW(TypeError) << "message" -- Throw TypeError
• Use << to chain multiple values in the error message
• Errors are properly propagated back to Python

When to use TVM_FFI_THROW vs TVM_FFI_LOG_AND_THROW:

• TVM_FFI_THROW: Use for normal runtime error handling. This is the standard way to report errors that will be caught and propagated to Python.

  void scale_run(TensorView output, TensorView input, double factor) {
    if (!input.device().device_type == kDLCUDA) {
      TVM_FFI_THROW(ValueError) << "Input must be a CUDA tensor";
    }
  }
  

• TVM_FFI_LOG_AND_THROW: Use only in cases where:

  1. The function may be called during object construction time (e.g., validation in constructors or setup methods)
  2. The exception may not be caught properly (e.g., during module initialization)
  3. The error condition almost never fails in practice (e.g., internal errors, unsupported dtype combinations in dispatch macros)

  This variant logs the error message before throwing, ensuring visibility even if the exception doesn't propagate correctly.

  void check_weights_shape(std::string which_weights) const {
    if (which_weights != "gemm1" && which_weights != "gemm2") {
      // Internal error that should never happen - use LOG_AND_THROW
      TVM_FFI_LOG_AND_THROW(InternalError)
          << "Internal error: which_weights = " << which_weights;
    }
  }
  

Real example (from csrc/activation.cu):

void glu(TensorView output, TensorView input) {
    CHECK_INPUT(input);
    CHECK_INPUT(output);
    CHECK_LAST_DIM_CONTIGUOUS_INPUT(input);
    CHECK_LAST_DIM_CONTIGUOUS_INPUT(output);

    unsigned int batch_size = input.size(0);
    unsigned int seq_len = input.size(1);
    unsigned int channels = input.size(2) / 2;

    cudaStream_t stream = get_stream(input.device());

    DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP16(input.dtype(), c_type, [&] {
        cudaError_t status = activation::GLU<c_type>(
            static_cast<const c_type*>(input.data_ptr()),
            static_cast<c_type*>(output.data_ptr()),
            batch_size, seq_len, channels, stream);
        TVM_FFI_ICHECK(status == cudaSuccess)
            << "GLU kernel failed: " << cudaGetErrorString(status);
        return true;
    });
}

Step 3: Create TVM-FFI Binding in csrc/

Create csrc/scale_jit_binding.cu:

#include "tvm_ffi_utils.h"

// Forward declaration
void scale_run(TensorView output, TensorView input, double factor);

// Export to TVM-FFI
TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, scale_run);

Key points:

• Include "tvm_ffi_utils.h" for TVM-FFI macros and type aliases
• Forward declare the launcher function(s)
• Export with TVM_FFI_DLL_EXPORT_TYPED_FUNC(exported_name, function) -- the exported name is how Python accesses it

Real example (from csrc/activation_jit_binding.cu):

#include "tvm_ffi_utils.h"

// Forward declarations of launcher functions
void glu(TensorView output, TensorView input);
void swish(TensorView output, TensorView input);

// TVM-FFI symbol exports
TVM_FFI_DLL_EXPORT_TYPED_FUNC(glu, glu);
TVM_FFI_DLL_EXPORT_TYPED_FUNC(swish, swish);

Real example (from csrc/norm_jit_binding.cu):

#include "tvm_ffi_utils.h"

void layernorm(TensorView output, TensorView input, TensorView weight,
               Optional bias_opt, double eps);
void rmsnorm(TensorView output, TensorView input, TensorView weight,
             Optional bias_opt, double eps);
// ... more forward declarations ...

TVM_FFI_DLL_EXPORT_TYPED_FUNC(layernorm, layernorm);
TVM_FFI_DLL_EXPORT_TYPED_FUNC(rmsnorm, rmsnorm);
// ... more exports ...

Note: Multiple launcher functions can be exported from a single binding file, as the norm and activation families do.

Step 4: Create JIT Generator in oasr/jit/

Create oasr/jit/scale.py:

from .core import gen_jit_spec, JitSpec
from . import env


def gen_scale_module() -> JitSpec:
    """Generate JIT spec for scale kernel."""
    return gen_jit_spec(
        "scale",
        [
            env.OASR_CSRC_DIR / "scale.cu",
            env.OASR_CSRC_DIR / "scale_jit_binding.cu",
        ],
    )

Key points:

• Import gen_jit_spec and JitSpec from .core, paths from .env
• gen_jit_spec() auto-detects GPU architecture and sets default NVCC flags
• Source files are the launcher .cu + binding .cu from csrc/
• NEVER write to package directories -- JIT cache goes to ~/.cache/oasr/jit/

Real examples:

# oasr/jit/activation.py
def gen_activation_module() -> JitSpec:
    return gen_jit_spec(
        "activation",
        [env.OASR_CSRC_DIR / "activation.cu",
         env.OASR_CSRC_DIR / "activation_jit_binding.cu"],
    )

# oasr/jit/norm.py
def gen_norm_module() -> JitSpec:
    return gen_jit_spec(
        "norm",
        [env.OASR_CSRC_DIR / "norm.cu",
         env.OASR_CSRC_DIR / "norm_jit_binding.cu"],
    )

(Optional) Specifying Supported CUDA Architectures

Oasr uses CompilationContext to manage CUDA architecture targets. This is critical because some kernels only work on specific GPU architectures (e.g., Hopper SM90, Blackwell SM100).

How CompilationContext Works

Automatic Detection (default):

from oasr.compilation_context import CompilationContext

ctx = CompilationContext()
# Automatically detects all GPUs in the system
# For SM90+, adds 'a' suffix (e.g., 9.0a for Hopper)
# Result: ctx.TARGET_CUDA_ARCHS = {(9, '0a'), (10, '0a'), ...}

Manual Override (via environment variable):

export OASR_CUDA_ARCH_LIST="8.0 9.0a 10.0a"
# Now only these architectures will be compiled

Specifying Architectures in Your JIT Module

When creating a JIT module, specify which major SM versions are supported:

from oasr.jit.core import gen_jit_spec
from oasr.jit import current_compilation_context

def gen_my_hopper_only_module():
    """Example: Kernel works on SM90 and later supported architectures."""
    nvcc_flags = current_compilation_context.get_nvcc_flags_list(
        # Explicitly list supported SM versions -- no automatic future compatibility
        supported_major_versions=[9, 10, 11, 12]  # SM90, SM100, SM110, SM120
    )

    return gen_jit_spec(
        name="my_hopper_kernel",
        sources=sources,
        extra_cuda_cflags=nvcc_flags,
    )

def gen_my_blackwell_only_module():
    """Example: Kernel only works on SM100 (Blackwell)."""
    nvcc_flags = current_compilation_context.get_nvcc_flags_list(
        supported_major_versions=[10]  # SM100 only
    )

    return gen_jit_spec(
        name="my_blackwell_kernel",
        sources=sources,
        extra_cuda_cflags=nvcc_flags,
    )

def gen_my_universal_module():
    """Example: Kernel works on all architectures (default)."""
    # No need to call get_nvcc_flags_list -- gen_jit_spec auto-detects
    return gen_jit_spec(
        name="my_universal_kernel",
        sources=sources,
    )

What Happens:

• If user's GPU is SM90 and they call a Hopper-only module -> Compiles and runs
• If user's GPU is SM80 and they call a Hopper-only module -> RuntimeError: No supported CUDA architectures found for major versions [9, 10, 11, 12]

Common Architecture Specifications

Supported Versions	Architectures	Use Case
None	All available GPUs	Universal kernels (default)
[9, 10, 11, 12]	SM90, SM100, SM110, SM120	Hopper, Blackwell
[10, 11, 12]	SM100, SM110, SM120	Blackwell only
[12]	SM120	Specific architecture only
[8, 9, 10, 11, 12]	SM80, SM90, SM100, SM110, SM120	Ampere, Hopper, Blackwell

Step 5: Create Python API in oasr/

Create oasr/scale.py:

import functools
from typing import Optional

import torch

from .api_logging import oasr_api


@functools.cache
def _get_scale_module():
    """Get or compile scale module (cached)."""
    from oasr.jit.scale import gen_scale_module

    return gen_scale_module().build_and_load()


@oasr_api
def scale(input: torch.Tensor, factor: float,
          out: Optional[torch.Tensor] = None) -> torch.Tensor:
    """Element-wise scale operation.

    Parameters
    ----------
    input : torch.Tensor
        Input tensor (CUDA).
    factor : float
        Scale factor.
    out : Optional[torch.Tensor]
        Output tensor (if None, allocate new tensor).

    Returns
    -------
    output : torch.Tensor
        Scaled tensor (input * factor).

    Examples
    --------
    >>> import torch
    >>> import oasr
    >>> x = torch.randn(1024, dtype=torch.float16, device="cuda")
    >>> y = oasr.scale(x, 2.0)
    >>> torch.allclose(y, x * 2.0)
    True
    """
    if out is None:
        out = torch.empty_like(input)

    # Call TVM-FFI function (output first in C++ convention)
    _get_scale_module().run(out, input, float(factor))

    return out

Key points:

• Uses @functools.cache to cache the compiled module (compile once per process)
• @oasr_api decorator (from oasr.api_logging) enables debug logging
• Destination-passing style: Output tensor is an optional Python parameter (out=None) but passed first to the C++ TVM-FFI function
• Import the JIT module lazily inside the cached function to avoid import-time compilation

Real example (from oasr/activation.py):

@functools.cache
def _get_activation_module():
    from oasr.jit.activation import gen_activation_module
    return gen_activation_module().build_and_load()


@oasr_api
def glu(input: torch.Tensor, out: Optional[torch.Tensor] = None) -> torch.Tensor:
    """Gated Linear Unit activation."""
    if out is None:
        out = torch.empty(
            input.shape[:-1] + (input.shape[-1] // 2,),
            device=input.device, dtype=input.dtype,
        )
    _get_activation_module().glu(out, input)  # output first!
    return out


@oasr_api
def swish(input: torch.Tensor, out: Optional[torch.Tensor] = None) -> torch.Tensor:
    """Swish (SiLU) activation: x * sigmoid(x)."""
    if out is None:
        out = torch.empty_like(input)
    _get_activation_module().swish(out, input)
    return out

(Advanced) Using @backend_requirement and @supported_compute_capability Decorators

For kernels with compute capability requirements or multiple backend choices, Oasr provides two decorators (in oasr.utils):

@supported_compute_capability Decorator

Marks a function with its supported CUDA compute capabilities:

from oasr.utils import supported_compute_capability

@supported_compute_capability([80, 86, 89, 90, 100, 103, 110, 120])
def _my_check_function(input, output):
    """Supports SM80 (Ampere) through SM120 (Blackwell)."""
    # Validation logic here
    return True

@backend_requirement Decorator

Enforces backend and problem size requirements at runtime. There are three usage patterns:

Pattern 1: Single Backend (No Backend Choices)

For kernels with only one implementation:

from oasr.utils import backend_requirement, supported_compute_capability

@supported_compute_capability([80, 86, 89, 90, 100, 103, 110, 120])
def _check_my_kernel(input, output):
    """Validate inputs. Must return True if valid."""
    if input.shape[-1] > 256:
        raise ValueError("Head dimension must be <= 256")
    return True

@backend_requirement(
    backend_checks={},  # Empty dict = no backend parameter
    common_check=_check_my_kernel,
)
def my_kernel(input, output):
    # Kernel implementation
    pass

Pattern 2: Multiple Backends

For kernels with multiple implementation backends (e.g., CUTLASS, cuDNN):

@supported_compute_capability([80, 86, 89, 90])
def _cutlass_check(q, k, v, backend):
    """CUTLASS backend: Ampere through Hopper."""
    if q.shape[-1] > 256:
        raise ValueError("CUTLASS: head_dim must be <= 256")
    return True

@supported_compute_capability([75, 80, 86, 89, 90, 100])
def _cudnn_check(q, k, v, backend):
    """cuDNN backend: Turing through Blackwell."""
    return True

@backend_requirement(
    backend_checks={
        "cutlass": _cutlass_check,
        "cudnn": _cudnn_check,
    },
    common_check=None,  # Optional: shared validation for all backends
)
def attention(q, k, v, backend="cutlass"):
    if backend == "cutlass":
        # CUTLASS implementation
        pass
    elif backend == "cudnn":
        # cuDNN implementation
        pass

Pattern 3: Auto Backend Selection

For kernels that can automatically select the best backend:

def _heuristic_func(suitable_backends, q, k, v, backend):
    """Return backends in order of preference."""
    if q.shape[-1] <= 128:
        preferred = ["cutlass", "cudnn"]
    else:
        preferred = ["cudnn", "cutlass"]
    return [b for b in preferred if b in suitable_backends]

@backend_requirement(
    backend_checks={
        "cutlass": _cutlass_check,
        "cudnn": _cudnn_check,
    },
    common_check=_common_validation,
    heuristic_func=_heuristic_func,  # Required when backend="auto" is used
)
def attention(q, k, v, backend="auto"):
    if backend == "auto":
        backend = attention.suitable_auto_backends[0]
    # ... rest of implementation

Features Added by @backend_requirement

The decorator adds these methods to the wrapped function:

# Check if a backend is supported (optionally for a specific CC)
scale.is_backend_supported("cutlass")           # True/False
scale.is_backend_supported("cutlass", cc=90)    # True/False for Hopper

# Check if any backend supports this compute capability
scale.is_compute_capability_supported(90)       # True/False

# Check if a backend exists
scale.has_backend("cutlass")                    # True/False

# Check if there are multiple backend choices
scale.has_backend_choices()                     # True/False

skip_check Keyword Argument

The decorator adds a skip_check keyword argument to bypass validation for performance-critical code paths:

# Normal call with validation
result = scale(x, 2.0)

# Skip validation for performance (use with caution!)
result = scale(x, 2.0, skip_check=True)

Check Function Requirements

Check functions must:

1. Accept the same arguments as the decorated function
2. Return True if validation passes
3. Raise ValueError with descriptive message if validation fails
4. Be decorated with @supported_compute_capability to specify supported architectures

Step 6: Write Tests in tests/

Create tests following the flat tests/test_<kernel>.py layout. The conftest.py provides: device (CUDA, skips if unavailable), dtype/dtype_all fixtures, batch_seq_hidden common shapes, and get_rtol_atol(dtype) helper.

Create tests/test_scale.py:

import pytest
import torch

import oasr


class TestScale:
    """Tests for oasr.scale() functional API."""

    @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
    @pytest.mark.parametrize("size", [128, 1024, 4096])
    def test_scale_correctness(self, dtype, size):
        """Test scale operation correctness."""
        x = torch.randn(size, dtype=dtype, device="cuda")
        factor = 3.14

        y = oasr.scale(x, factor)

        expected = x * factor
        if dtype == torch.float32:
            rtol, atol = 1e-5, 1e-6
        else:
            rtol, atol = 1e-3, 1e-3

        torch.testing.assert_close(y, expected, rtol=rtol, atol=atol)

    @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
    def test_scale_destination_passing(self, dtype):
        """Test scale with pre-allocated output."""
        x = torch.randn(1024, dtype=dtype, device="cuda")
        out = torch.empty_like(x)
        factor = 2.0

        result = oasr.scale(x, factor, out=out)

        # Should return the same tensor
        assert result.data_ptr() == out.data_ptr()

        expected = x * factor
        torch.testing.assert_close(result, expected, rtol=1e-3, atol=1e-3)

    def test_scale_cpu_error(self):
        """Test that CPU tensors raise an error."""
        x = torch.randn(128, dtype=torch.float32)

        with pytest.raises(Exception):
            oasr.scale(x, 2.0)

Key points:

• Use pytest.mark.parametrize for multiple configurations
• Compare against reference implementation
• Set appropriate tolerances per dtype (use get_rtol_atol() from conftest)
• Test destination-passing style: verify result.data_ptr() == out.data_ptr()
• Test error cases

Real example (from tests/test_activation.py):

class TestGLU:
    @pytest.mark.parametrize(
        "batch_size,seq_len,channels",
        [(2, 128, 256), (4, 256, 512)],
    )
    @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
    def test_glu(self, batch_size, seq_len, channels, dtype):
        x = torch.randn(batch_size, seq_len, 2 * channels, device="cuda", dtype=dtype)
        output = oasr.glu(x)
        expected = F.glu(x, dim=-1).to(dtype)
        torch.testing.assert_close(output, expected, rtol=1e-2, atol=1e-2)

    def test_glu_destination_passing(self, dtype):
        x = torch.randn(2, 128, 512, device="cuda", dtype=torch.float16)
        out = torch.empty(2, 128, 256, device="cuda", dtype=torch.float16)
        result = oasr.glu(x, out=out)
        assert result.data_ptr() == out.data_ptr()

Testing with Architecture Requirements

When your kernel has architecture requirements, add skip checks:

import pytest
import torch
from oasr.utils import is_sm90a_supported

def test_hopper_kernel():
    if not is_sm90a_supported(torch.device("cuda")):
        pytest.skip("SM90a is not supported on this GPU")
    # Test code here
    ...

Step 7: Register in AOT

Register your kernel in oasr/aot.py so users with oasr-jit-cache can skip JIT compilation.

Edit oasr/aot.py:

def gen_all_modules() -> List:
    from oasr.jit.activation import gen_activation_module
    from oasr.jit.norm import gen_norm_module
    from oasr.jit.conv import gen_conv_module, gen_conv2d_module, gen_cudnn_conv2d_module
    from oasr.jit.gemm import gen_gemm_module, gen_bmm_module, gen_group_gemm_module
    from oasr.jit.scale import gen_scale_module  # NEW

    return [
        gen_activation_module(),
        gen_norm_module(),
        gen_conv_module(),
        gen_conv2d_module(),
        gen_cudnn_conv2d_module(),
        gen_gemm_module(),
        gen_bmm_module(),
        gen_group_gemm_module(),
        gen_scale_module(),  # NEW
    ]

Step 8: Export API

Edit oasr/__init__.py:

from .scale import scale as scale

Add "scale" to the __all__ list.

Step 9: Run and Test

# The kernel compiles automatically on first use
pytest tests/test_scale.py -v

# Run a single test
pytest tests/test_scale.py::TestScale::test_scale_correctness -v

Step 10: Add Benchmark

All new kernels should have benchmarks. This helps track performance regressions and allows users to compare against other implementations.

Benchmarks follow the routines + thin-wrapper pattern: a single routine module in benchmarks/routines/<family>.py exposes the kernel(s) to both the unified CLI (oasr_benchmark.py) and to per-kernel bench_*.py scripts that act as thin wrappers around run_standalone(). Reference: benchmarks/routines/activation.py + benchmarks/bench_glu.py.

Step 10a: Create the routine module

Create benchmarks/routines/scale.py:

"""Scale family benchmark routines."""

from __future__ import annotations

import argparse
from typing import Any

import torch

import oasr
from benchmarks.routines.bench_utils import (
    BenchResult,
    OutputWriter,
    bench_fn,
    check_close,
    compute_bandwidth_tb_s,
    dtype_size,
    parse_dtype,
    run_main,
)

SUBROUTINES = ["scale"]

# ---------------------------------------------------------------------------
# Default configs
# ---------------------------------------------------------------------------

DEFAULT_CONFIGS: dict[str, list[dict[str, Any]]] = {
    "scale": [
        {"size": 1024},
        {"size": 4096},
        {"size": 16384},
        {"size": 65536},
        {"size": 262144},
    ],
}

PROFILE_CONFIGS: dict[str, tuple] = {
    "scale": (65536,),
}


def get_default_configs() -> dict[str, list[dict[str, Any]]]:
    return DEFAULT_CONFIGS


# ---------------------------------------------------------------------------
# Setup functions  -- return (oasr_fn, pytorch_fn) closures
# ---------------------------------------------------------------------------


def setup_scale(size, dtype=torch.float16):
    x = torch.randn(size, device="cuda", dtype=dtype)
    factor = 2.0

    def oasr_fn():
        return oasr.scale(x, factor)

    def pytorch_fn():
        return x * factor

    return oasr_fn, pytorch_fn


# ---------------------------------------------------------------------------
# CLI args (used by oasr_benchmark.py)
# ---------------------------------------------------------------------------


def parse_args(parser: argparse.ArgumentParser) -> None:
    parser.add_argument("--size", type=int, default=None, help="Number of elements")


# ---------------------------------------------------------------------------
# run_test  -- entry point invoked by oasr_benchmark.py
# ---------------------------------------------------------------------------


def _scale_bytes(size, dtype):
    """Bytes accessed: read input + write output."""
    return 2 * size * dtype_size(dtype)


def run_test(args: argparse.Namespace, output: OutputWriter) -> None:
    subroutine = getattr(args, "subroutine", "scale")
    dtype_str = getattr(args, "dtype", "float16")
    dtype = parse_dtype(dtype_str)
    do_check = getattr(args, "refcheck", False)
    allow_mismatch = getattr(args, "allow_output_mismatch", False)
    dry_run_iters = getattr(args, "dry_run_iters", 5)
    num_iters = getattr(args, "num_iters", 30)
    use_cuda_events = getattr(args, "use_cuda_events", False)

    configs = _resolve_configs(args, subroutine)

    for cfg in configs:
        oasr_fn, pytorch_fn = setup_scale(cfg["size"], dtype)
        fn_map = get_fn_map(subroutine, oasr_fn, pytorch_fn)
        backends = getattr(args, "backends", None) or list(fn_map.keys())

        bytes_accessed = _scale_bytes(cfg["size"], dtype)
        shape_str = f"[{cfg['size']}]"

        if do_check and "torch" in backends and any(b in fn_map and b != "torch" for b in backends):
            passed, max_diff = check_close(oasr_fn(), pytorch_fn())
            if not passed:
                print(f"  [ERROR] Output mismatch for {shape_str} (max_diff={max_diff:.6f})")
                if not allow_mismatch:
                    continue

        for backend in backends:
            if backend not in fn_map:
                print(f"  [WARNING] Unknown backend '{backend}', skipping")
                continue
            median_ms, std_ms = bench_fn(
                fn_map[backend],
                dry_run_iters=dry_run_iters,
                num_iters=num_iters,
                use_cuda_events=use_cuda_events,
            )
            bw = compute_bandwidth_tb_s(bytes_accessed, median_ms)
            output.write_result(BenchResult(
                routine="scale",
                subroutine=subroutine,
                backend=backend,
                shape=shape_str,
                dtype=dtype_str,
                median_ms=median_ms,
                std_ms=std_ms,
                bandwidth_tb_s=bw,
            ))


def _resolve_configs(args, subroutine):
    size = getattr(args, "size", None)
    if size is not None:
        return [{"size": size}]
    return DEFAULT_CONFIGS.get(subroutine, DEFAULT_CONFIGS["scale"])


def get_fn_map(subroutine, cuda_fn, torch_fn):
    """Return {backend_name: fn} -- backend names match what users pass to --backends."""
    return {"cuda": cuda_fn, "torch": torch_fn}


# ---------------------------------------------------------------------------
# Standalone entry  -- used by bench_scale.py thin wrapper
# ---------------------------------------------------------------------------


def run_standalone(variant: str = "scale") -> None:
    subs = [variant]
    pcfg = {k: PROFILE_CONFIGS[k] for k in subs if k in PROFILE_CONFIGS}
    setup_funcs = {sub: _make_profile_setup(sub) for sub in subs if sub in PROFILE_CONFIGS}

    def benchmark():
        output = OutputWriter()
        for sub in subs:
            output.write_header(f"{sub.upper()} Kernel Benchmark")
            for cfg in DEFAULT_CONFIGS.get(sub, []):
                oasr_fn, pytorch_fn = setup_scale(cfg["size"], torch.float16)
                bytes_accessed = _scale_bytes(cfg["size"], torch.float16)
                shape_str = f"[{cfg['size']}]"
                for backend, fn in get_fn_map(sub, oasr_fn, pytorch_fn).items():
                    median_ms, std_ms = bench_fn(fn)
                    bw = compute_bandwidth_tb_s(bytes_accessed, median_ms)
                    output.write_result(BenchResult(
                        routine="scale", subroutine=sub, backend=backend,
                        shape=shape_str, dtype="float16",
                        median_ms=median_ms, std_ms=std_ms,
                        bandwidth_tb_s=bw,
                    ))
        output.finalize()

    run_main(f"{variant.upper()} Kernel", pcfg, setup_funcs, benchmark)


def _make_profile_setup(subroutine):
    cfg_tuple = PROFILE_CONFIGS[subroutine]

    def _setup():
        return setup_scale(*cfg_tuple)

    return _setup

Key points:

• SUBROUTINES, parse_args, run_test, get_default_configs, run_standalone are the contract the routine registry expects (see benchmarks/routines/__init__.py).
• setup_*() functions return two closures (oasr_fn, pytorch_fn) that take no arguments -- this is what bench_fn consumes and what the profile path replays.
• Backend names ("cuda" / "torch") are family-conventional. Norm/Conv1D/Activation use cuda/torch; GEMM/Conv2D use cutlass/torch. Match the family your kernel belongs to.
• Use compute_bandwidth_tb_s for memory-bound kernels and compute_gemm_tflops / compute_bmm_tflops for compute-bound ones.

Step 10b: Register the routine

Edit benchmarks/routines/__init__.py and add the routine to ROUTINE_REGISTRY:

ROUTINE_REGISTRY: dict[str, str] = {
    "gemm": "benchmarks.routines.gemm",
    "norm": "benchmarks.routines.norm",
    # ...
    "scale": "benchmarks.routines.scale",  # NEW
}

This makes python benchmarks/oasr_benchmark.py --routine scale --subroutine scale ... work.

Step 10c: Create the thin wrapper

Create benchmarks/bench_scale.py:

#!/usr/bin/env python3
"""OASR Scale Benchmark -- CUDA vs PyTorch."""

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from benchmarks.routines.scale import run_standalone


if __name__ == "__main__":
    run_standalone("scale")

For families with multiple subroutines (e.g. activation), create one thin wrapper per subroutine (bench_glu.py, bench_swish.py) -- each calls run_standalone("<subroutine>"). See benchmarks/bench_glu.py and benchmarks/bench_swish.py.

Step 10d: Run it

# Standalone (thin wrapper)
python benchmarks/bench_scale.py

# Unified CLI
python benchmarks/oasr_benchmark.py --routine scale --subroutine scale \
    --backends cuda torch --size 4096 --dtype float16 --refcheck -vv

# Profiling mode (NVTX markers for Nsight Compute)
ncu --set full -o scale_profile python benchmarks/bench_scale.py --profile --target oasr

Benchmark utilities (from benchmarks/routines/bench_utils.py):

Function	Purpose
bench_fn(fn, ...)	Time a function, returns (median_ms, std_ms)
profile_kernel(name, fn, ...)	Run with NVTX markers for Nsight Compute
check_close(actual, expected, ...)	Compare tensors, returns (passed, max_diff)
compute_bandwidth_tb_s(bytes, ms)	Memory bandwidth for memory-bound kernels
compute_gemm_tflops(M, N, K, ms)	TFLOPS for compute-bound GEMM-like kernels
BenchResult(...)	Structured result dataclass
OutputWriter()	Manages terminal [PERF] lines + CSV output
run_main(title, pcfg, setup, fn)	Standard standalone main with --profile support

-> For complete benchmarking guide, see .claude/skills/benchmark-kernel/SKILL.md

Existing OASR Kernel Families

When adding a new kernel, look at these existing families as references:

Family	Kernel Header	Launcher	Binding	JIT Generator	Python API
Activation	include/oasr/activation.cuh	csrc/activation.cu	csrc/activation_jit_binding.cu	jit/activation.py	activation.py
Norm	include/oasr/norm.cuh	csrc/norm.cu	csrc/norm_jit_binding.cu	jit/norm.py	norm.py
Conv1D	include/oasr/conv/conv1d.cuh	csrc/conv.cu	csrc/conv_jit_binding.cu	jit/conv.py	conv.py
Conv2D	include/oasr/conv/conv2d.cuh	csrc/conv2d.cu	csrc/conv2d_jit_binding.cu	jit/conv.py	conv.py
GEMM	include/oasr/gemm/gemm.cuh	csrc/gemm.cu	csrc/gemm_jit_binding.cu	jit/gemm.py	gemm.py
BMM	include/oasr/gemm/bmm.cuh	csrc/bmm.cu	csrc/bmm_jit_binding.cu	jit/gemm.py	gemm.py
Group GEMM	include/oasr/gemm/group_gemm.cuh	csrc/group_gemm.cu	csrc/group_gemm_jit_binding.cu	jit/gemm.py	gemm.py

Summary of Files Created/Modified

include/oasr/scale.cuh                  # NEW: CUDA kernel definition
csrc/scale.cu                           # NEW: TVM-FFI launcher
csrc/scale_jit_binding.cu               # NEW: TVM-FFI binding
oasr/jit/scale.py                       # NEW: JIT generator
oasr/scale.py                           # NEW: Python API
oasr/__init__.py                        # MODIFIED: Export API
oasr/aot.py                             # MODIFIED: Register AOT
tests/test_scale.py                     # NEW: Unit tests
benchmarks/routines/scale.py            # NEW: Benchmark routine module
benchmarks/routines/__init__.py         # MODIFIED: Register routine
benchmarks/bench_scale.py               # NEW: Standalone benchmark wrapper