Kernscope

A small LLM inference operator library with PyTorch, Triton CUDA, and TileLang-Ascend backends.

Kernscope provides directly callable inference operators behind a small Python API. Operators have explicit numerical contracts, reference implementations, and correctness tests; accelerator kernels are available where listed below.

Current operators:

RMSNorm
fused residual-add RMSNorm
SwiGLU activation (silu_and_mul)
GEMV
Softmax

Current backends:

PyTorch — portable reference implementation
Triton — CUDA kernels
TileLang-Ascend — Ascend NPU kernels

Backend support

Operator	PyTorch	Triton CUDA	TileLang-Ascend
`rms_norm`	CPU / CUDA	CUDA	Ascend NPU
`fused_add_rms_norm`	CPU / CUDA	CUDA	—
`silu_and_mul`	CPU / CUDA	CUDA	—
`gemv`	CPU / CUDA	CUDA	—
`softmax`	CPU / CUDA	CUDA	—

Backend selection is explicit:

from kernscope import rms_norm

y = rms_norm(x, weight, backend="torch")
y = rms_norm(x, weight, backend="triton")
y = rms_norm(x, weight, backend="tilelang_ascend")

The public API does not require callers to interact with Triton or TileLang directly.

Install

Kernscope requires Python 3.10 or newer.

Clone the repository:

git clone https://github.com/luicarus/kernscope.git
cd kernscope

For the PyTorch reference implementations and tests:

python -m pip install -e ".[test]"

For the Triton CUDA backend:

python -m pip install -e ".[gpu,test]"

A CUDA-enabled PyTorch installation is required before using backend="triton".

The TileLang-Ascend backend requires an Ascend PyTorch environment with torch_npu and a compatible TileLang-Ascend installation. These dependencies are imported only when backend="tilelang_ascend" is selected, so CPU and CUDA users do not need the Ascend stack installed.

TileLang-Ascend RMSNorm is experimental and has not yet been compiled or validated on Ascend hardware.

Quick start

RMSNorm

import torch

from kernscope import rms_norm

x = torch.randn(
    512,
    1024,
    device="cuda",
    dtype=torch.bfloat16,
)
weight = torch.ones(
    1024,
    device="cuda",
    dtype=torch.bfloat16,
)

y = rms_norm(
    x,
    weight,
    eps=1e-6,
    backend="triton",
)

For the PyTorch reference:

y = rms_norm(
    x,
    weight,
    eps=1e-6,
    backend="torch",
)

On Ascend:

import torch
import torch_npu

from kernscope import rms_norm

x = torch.randn(
    512,
    1024,
    device="npu",
    dtype=torch.bfloat16,
)
weight = torch.ones(
    1024,
    device="npu",
    dtype=torch.bfloat16,
)

y = rms_norm(
    x,
    weight,
    eps=1e-6,
    backend="tilelang_ascend",
)

Fused residual-add RMSNorm

import torch

from kernscope import fused_add_rms_norm

x = torch.randn(
    512,
    1024,
    device="cuda",
    dtype=torch.bfloat16,
)
residual = torch.randn_like(x)
weight = torch.ones(
    1024,
    device="cuda",
    dtype=torch.bfloat16,
)

fused_add_rms_norm(
    x,
    residual,
    weight,
    eps=1e-6,
    backend="triton",
)

The fused operator updates both x and residual in place and returns None.

weight is read-only.

SwiGLU activation

import torch

from kernscope import silu_and_mul

x = torch.randn(8, 2 * 4096, device="cuda", dtype=torch.bfloat16)
y = silu_and_mul(x, backend="triton")
# y has shape (8, 4096); x is unchanged.

GEMV

import torch

from kernscope import gemv

x = torch.randn(5, dtype=torch.bfloat16)
weight = torch.randn(3, 5, dtype=x.dtype)
y = gemv(x, weight)
# y has shape (3,); x and weight are unchanged.

Softmax

import torch

from kernscope import softmax

x = torch.randn(2, 5, dtype=torch.bfloat16)
y = softmax(x)
# y has shape (2, 5); x is unchanged.

Numerical contract

RMSNorm

For input x and weight weight, RMSNorm computes:

x_fp32 = fp32(x)

scale =
    rsqrt(
        mean(x_fp32², axis=-1)
        + eps
    )

output =
    cast_to_input_dtype(
        x_fp32
        * scale
        * fp32(weight)
    )

Reduction and scaling use FP32 intermediates.

Fused residual-add RMSNorm

For input snapshots x_in and residual_in:

s =
    fp32(x_in)
    + fp32(residual_in)

residual =
    cast_to_input_dtype(s)

x =
    cast_to_input_dtype(
        s
        * rsqrt(
            mean(s², axis=-1)
            + eps
        )
        * fp32(weight)
    )

Normalization uses the FP32 residual sum before it is rounded for storage in residual.

SwiGLU activation

D      = x.shape[-1] // 2
gate   = fp32(x[..., :D])
up     = fp32(x[..., D:])
output = cast_to_input_dtype(gate * sigmoid(gate) * up)

All intermediate calculations use FP32; only the final output is cast to the input dtype.

GEMV

output = cast_to_input_dtype(fp32(weight) @ fp32(x))

GEMV computes a matrix-vector product without bias, using FP32 multiplication and accumulation. The reference disables autocast locally to preserve this precision.

Softmax

z      = fp32(x) - max(fp32(x), dim=-1, keepdim=True)
output = cast_to_input_dtype(exp(z) / sum(exp(z), dim=-1, keepdim=True))

The reference uses torch.softmax(x.float(), dim=-1) with autocast disabled locally. Negative infinity masks individual entries when the row contains finite values. Rows containing NaN or positive infinity, and rows entirely equal to negative infinity, produce NaN outputs as in PyTorch.

API contract

The RMSNorm operators accept eps=1e-6. All operators use keyword-only backend="torch" by default.

The PyTorch implementation is the default backend:

rms_norm(
    x,
    weight,
    eps=1e-6,
    backend="torch",
)

fused_add_rms_norm(
    x,
    residual,
    weight,
    eps=1e-6,
    backend="torch",
)

rms_norm currently supports:

torch
triton
tilelang_ascend

fused_add_rms_norm currently supports:

torch
triton

silu_and_mul supports the same two backends and has no eps parameter.

gemv(x, weight, *, backend="torch") supports torch and triton, with no bias or eps parameter.

softmax(x, *, backend="torch") supports torch and triton, and always operates on the last dimension.

Inputs

Supported dtypes:

FP16
BF16
FP32

For the RMSNorm operators, x has shape:

(..., hidden_size)

and weight has shape:

(hidden_size,)

Requirements:

the final dimension must be nonempty;
tensors must be contiguous;
participating tensors must share dtype and device;
eps must be finite and positive.

For fused_add_rms_norm:

residual must have the same shape as x;
x, residual, and weight storage regions must not overlap;
x and residual are modified in place.
both its PyTorch and Triton implementations are inference-only.

For silu_and_mul, x has shape (..., 2D) with a nonempty, even final dimension. It must be contiguous and use a supported dtype. D may be odd, and batch dimensions may be zero. The result is a new tensor of shape (..., D), preserving dtype and device; x is read-only.

For gemv, x has shape (H,) with H > 0, and weight has shape (N, H). Both must be contiguous and share a supported dtype and device. N=0 returns an empty output. The result is a new tensor of shape (N,), preserving dtype and device. Both inputs are read-only and may share storage.

For softmax, x has shape (..., H) with H > 0, is contiguous, and uses a supported dtype. Leading dimensions may be zero. The result is a new tensor with the same shape, dtype and device; x is read-only.

The PyTorch references for rms_norm, silu_and_mul, gemv, and softmax support autograd.

Triton and TileLang-Ascend kernels are inference-only and reject tensors requiring gradients.

Triton CUDA backend

The CUDA kernels are implemented in:

src/kernscope/backends/triton.py
src/kernscope/backends/triton_gemv.py
src/kernscope/backends/triton_softmax.py

RMSNorm

The current RMSNorm implementation processes multiple rows per Triton program.

The kernel:

loads the input row into FP32;
computes the sum of squares;
evaluates the RMS normalization scale;
applies the learned weight;
writes the result in the original input dtype.

The launch configuration used by the current measured implementation is:

4 rows / program
8 warps / program

Fused residual-add RMSNorm

One Triton program handles one row.

The kernel loads:

x
residual
weight

and keeps the residual sum in FP32 through normalization before writing both outputs.

Conceptually:

x ─────────┐
           ├─ FP32 add ───────────────┐
residual ──┘                          │
                                      ├─ RMSNorm ─→ x
                                      │
                                      └────────────→ residual
weight ───────────────────────────────┘

Fusing residual addition with RMSNorm removes a separate addition kernel launch and avoids reloading the residual sum from global memory.

The current launch configuration is:

1 row / program
1 warp / program

SwiGLU activation

One Triton program handles 1024 flattened output elements with four warps. It loads matching gate and up values, computes SiLU and multiplication in FP32, and writes the input-dtype output. Masks handle partial blocks; the input remains unchanged.

GEMV

One Triton program computes one output row with four warps. It uses 1024-element blocks when N>=512 and H is 1024 or 2048, and 256-element blocks otherwise. It loads x and weight into FP32, accumulates products, then reduces and stores the input-dtype result. Masks cover partial blocks, and empty outputs skip the launch. Both inputs remain unchanged; this backend is inference-only.

Softmax

One Triton program normalizes two rows when the flattened row count N>=512 and H<=1024, excluding H=1023; other cases use one row. It uses one warp when H<=256 and four warps otherwise. The block size is the next power of two at least as large as H; masks cover partial rows and columns, with masked loads using negative infinity. Max, exponential, sum and division use FP32, followed by an input-dtype store. Empty leading dimensions skip the launch; inputs requiring gradients are rejected.

TileLang-Ascend backend

The Ascend RMSNorm implementation is in:

src/kernscope/backends/tilelang_ascend.py

It is compiled through TileLang with:

@tilelang.jit(
    out_idx=[2],
    target="ascendc",
    pass_configs={
        tilelang.PassConfigKey.TL_ASCEND_AUTO_SYNC: True,
        tilelang.PassConfigKey.TL_ASCEND_AUTO_CV_COMBINE: True,
        tilelang.PassConfigKey.TL_ASCEND_MEMORY_PLANNING: False,
    },
)

The input is flattened into rows:

(..., hidden_size)
        │
        ▼
(rows, hidden_size)

and the kernel maps one NPU Vector block to each row.

The main data path is:

Global Memory
     │
     ▼
     UB
     │
     ▼
    FP32
     │
     ├─ square
     ├─ reduce_sum
     ├─ / hidden_size
     ├─ + eps
     └─ rsqrt
     │
     ▼
normalize
     │
     ├─ multiply weight
     ├─ cast to input dtype
     │
     ▼
Global Memory

FP16 and BF16 inputs are staged through UB buffers and converted to FP32 before reduction.

The working width is padded to a multiple of 64:

block_size = (
    (hidden_size + 63) // 64
) * 64

Padding is zero-filled before reduction, while the reduction itself uses the real hidden size.

The current implementation keeps full-row intermediate buffers in UB. This avoids an additional hidden-dimension tiling path but also makes UB capacity a constraint for sufficiently large hidden sizes.

CUDA performance

The following compares earlier and current Triton implementations on an RTX 3050 Ti Laptop GPU (4 GB) at 512 × 1024 BF16.

Values are medians of three kernel launches captured by Nsight Compute.

Operator	Before	Current	Latency reduction	Current configuration
RMSNorm	12.544 µs	10.176 µs	18.9%	4 rows/program, 8 warps
Fused add RMSNorm	20.064 µs	19.136 µs	4.6%	1 row/program, 1 warp

Recorded environment:

Python        3.12
PyTorch       2.8.0+cu128
CUDA          12.8
Triton        3.4.0
GPU           RTX 3050 Ti Laptop GPU
VRAM          4 GB
Driver        610.47

These measurements apply to the recorded workload. Launch settings are currently fixed, so other shapes and dtypes may produce different tradeoffs.

For fused residual-add RMSNorm, the current configuration increases registers per thread from:

29 → 84

and lowers achieved occupancy from:

84.93% → 25.95%

while still reducing measured kernel latency.

This is also a useful example of why higher occupancy does not necessarily imply lower kernel latency.

The complete profiler data is available in:

docs/rms_norm_nsight_results.md

CUDA Graph benchmark results and Nsight Compute kernel measurements use different timing methods and should only be compared within the same measurement method.

SwiGLU baseline

At BF16 input 512×8192 and output 512×4096, CUDA Graph latency is 68.634 µs for Triton versus 437.280 µs for this project's PyTorch eager reference with FP32 intermediates (6.37×). Values are medians of five repeat medians. The separate Nsight Compute median of three launches is 63.648 µs, with DRAM throughput at 95.27% of peak sustained throughput.

The SwiGLU results table records the selected initial configuration: 1024 output elements/program and four warps. Subsequent parameter, indexing, cache, and load-order experiments did not establish sufficient benefit to replace it. These results apply to the measured reference and workload.

GEMV measurements

At BF16 weights 4096×4096 and a single input vector of length 4096, CUDA Graph medians are 180.332 µs for Triton and 181.499 µs for native torch.mv. The FP32 reference is 926.073 µs, including weight conversion. These are separate paths; the native and Triton timings are close. Nsight records 181.760 µs kernel latency and 96.70% DRAM throughput. See the GEMV metrics table and profiling guide for scope and raw evidence.

The original baseline above uses 1024 columns per block. The block-256 stage uses 256 columns and 4 warps. In a separate paired Nsight comparison at BF16 weights 4096×8192, latency falls from 392.416 to 360.544 µs (8.1%), registers from 64 to 39, and compiler-reported spills from 2 to 0. At 4096×4096, paired latency remains 181.376 µs. CUDA Graph follow-up at 4096×8192 improves 390.071→359.431 µs (7.9%); 4096×4096 stays near 180.33 µs, while 512×1024 regresses 13.2%. An independent block-256 run records 359.497 µs for Triton versus 360.838 µs for native torch.mv; see the profiling guide for raw samples and measurement scope.

The current short-H rule selects 1024 columns at N>=512, H∈{1024,2048}. Paired BF16 CUDA Graph latency at 512×1024 improves 3.167→2.824 µs (10.8%), restoring the original block-1024 performance level; 512×2048 improves 4.438→4.282 µs (3.5%). Wide-H controls remain essentially unchanged. Cold-cache Nsight latency at 512×1024 is essentially unchanged (8.000→8.064 µs), with registers 22→21. See the profiling guide for the paired and independent benchmark samples.

Softmax measurements

At BF16 input/output 512×1024, baseline CUDA Graph medians are 6.096 µs for Triton and 8.472 µs for native torch.softmax (1.39×); the conversion-inclusive FP32 reference is 47.102 µs. The four-warp baseline is about 75% slower at 512×128, so the gain is shape-specific. The separate Nsight Compute median of three launches is 10.464 µs, with 23 registers/thread, 77.17% achieved occupancy and 86.57% DRAM throughput. See the Softmax metrics table and profiling guide for sampling, source hashes and raw evidence.

The short-H warp stage selects one warp for H<=256 and keeps four warps for wider rows. Paired BF16 CUDA Graph latency at 512×128 falls 3.882→2.233 µs (42.5%), and at 512×256 5.889→3.748 µs (36.3%). Each paired repeat improves at both shapes; the 512×1024/4096/8192 controls remain near their earlier latency. Separate paired Nsight latency falls 6.304→4.160 µs (34.0%) and 6.656→4.672 µs (29.8%) at the two short-row shapes. At H=128, registers stay at 16 while occupancy falls 82.49%→22.11%; at H=256, registers rise 17→20. Results apply to these shapes and the recorded GPU state.

The current row-grouping stage uses two rows when N>=512 and H<=1024, excluding H=1023 after measured regressions. Paired BF16 Nsight medians of five launches fall 10.912→10.592 µs (2.9%) at 512×1024 and 4.096→3.552 µs (13.3%) at 512×128. At the main shape, registers rise 23→34 and occupancy falls 76.40%→64.42%, with no compiler-reported spills. CUDA Graph medians independently fall 6.090→5.483 µs (10.0%) and 2.238→1.625 µs (27.4%). The main Nsight gain is modest and individual samples vary. Performance claims here cover BF16; all three dtypes pass correctness checks. See the profiling guide for raw evidence and sampling scope.

Correctness

Run the complete test suite with:

python -m pytest -q

CUDA cases automatically skip when CUDA is unavailable.

Ascend hardware-dependent cases automatically skip when torch_npu, TileLang, or an available NPU runtime is missing.

Numerical tolerances

dtype	`rtol`	`atol`
FP16	`1e-3`	`1e-3`
BF16	`1e-2`	`1e-2`
FP32	`1e-6`	`1e-6`

CPU tests are checked against high-precision reference computation.

Accelerator tests check each implementation against its corresponding public numerical contract.

Ascend RMSNorm tests currently cover shapes including:

(5,)
(2, 3, 5)
(3, 1000)
(512, 1024)

as well as zero-valued and empty inputs.

Backend dispatch, device validation, autograd rejection, and dependency isolation are tested separately.

SwiGLU checks cover all three dtypes, CPU FP64 reference math, CUDA/PyTorch comparisons, autograd, odd D, empty batches, contiguous offset views, and extreme gate values.

GEMV CPU checks cover all three dtypes, FP64 reference math, FP32 accumulation, input gradients, empty outputs, contiguous offset views, and read-only input aliasing. CUDA tests compare Triton with both the PyTorch reference and FP64 math, including partial blocks, offset views, aliasing, and inference-only validation. Autocast checks preserve the FP32 contract.

Softmax CPU/CUDA checks cover all three dtypes, FP64 reference math, widths 1–16384, large logits, nonfinite rows, empty batches, contiguous offset views and autocast state preservation. A 65536-column FP16 case checks FP32 denominator accumulation. The PyTorch reference supports input gradients; Triton rejects them. CPU reference use remains independent of Triton installation.

Benchmarks

Run the CUDA RMSNorm benchmark:

python benchmarks/bench_rms_norm.py \
    --operator rms_norm \
    --rows 512 \
    --runs 5

Run fused residual-add RMSNorm:

python benchmarks/bench_rms_norm.py \
    --operator fused_add_rms_norm \
    --rows 512 \
    --runs 5

Run SwiGLU activation:

python benchmarks/bench_silu_and_mul.py \
    --dtype bf16 \
    --runs 5

Run GEMV:

python benchmarks/bench_gemv.py \
    --dtype bf16 \
    --runs 5

Run Softmax:

python benchmarks/bench_softmax.py \
    --dtype bf16 \
    --runs 5

Benchmarks use CUDA Graph replay. RMSNorm benchmarks sweep:

FP16
BF16
FP32

Timestamped CSV files record metadata including:

device
software versions
shape
dtype
backend
GPU state

Fused timing uses zero inputs so repeated in-place calls remain idempotent. Correctness tests use nonzero inputs.

SwiGLU defaults to BF16 with rows 1, 16, 128, 512 and output widths 1024, 4096, 8192. It warms up each backend for 20 calls, alternates backend order across five repeats, and records the Triton source hash alongside timing and GPU state metadata.

See docs/profiling.md for Nsight Compute commands, archived reports, and the source revisions corresponding to recorded measurements.

Repository layout

src/kernscope/
├── ops/
│   ├── rms_norm.py
│   ├── fused_add_rms_norm.py
│   ├── silu_and_mul.py
│   ├── gemv.py
│   ├── softmax.py
│   └── ...
│
└── backends/
    ├── triton.py
    ├── triton_gemv.py
    ├── triton_softmax.py
    └── tilelang_ascend.py

benchmarks/
└── CUDA Graph benchmarks
    and profiling entry points

docs/
├── profiling.md
├── rms_norm_nsight_results.md
├── silu_and_mul_baseline.md
├── gemv_nsight_results.md
├── softmax_nsight_results.md
└── ...

tests/ops/
├── PyTorch reference tests
├── Triton CUDA tests
├── TileLang-Ascend API tests
└── TileLang-Ascend NPU tests

The split between ops and backends is intentional:

ops/
    public API
    validation
    numerical contract
    PyTorch reference

backends/
    hardware-specific kernel implementations

This keeps backend dependencies outside the public operator definitions until the corresponding backend is selected.

Scope

Kernscope is a small operator library rather than a serving framework.

Serving-specific integration belongs in consumer projects such as QuantAssay. Kernscope can be installed and imported without SGLang or QuantAssay.

Current coverage is:

RMSNorm
├── PyTorch
├── Triton CUDA
└── TileLang-Ascend

Fused residual-add RMSNorm
├── PyTorch
└── Triton CUDA

SwiGLU activation
├── PyTorch
└── Triton CUDA

GEMV
├── PyTorch
└── Triton CUDA

Softmax
├── PyTorch
└── Triton CUDA

An earlier SGLang serving comparison predates the current Triton optimizations and did not establish an end-to-end serving speedup. The optimized kernels have not yet been re-evaluated in that workload.

Additional operators and backend implementations are added around concrete inference workloads rather than maintained as standalone kernel demos.

Name		Name	Last commit message	Last commit date
benchmarks		benchmarks
docs		docs
src/kernscope		src/kernscope
tests/ops		tests/ops
.gitignore		.gitignore
AGENTS.md		AGENTS.md
README.md		README.md
pyproject.toml		pyproject.toml

Latest commit

History

Folders and files

Repository files navigation

Kernscope

Backend support

Install

Quick start

RMSNorm

Fused residual-add RMSNorm

SwiGLU activation

GEMV

Softmax

Numerical contract

RMSNorm

Fused residual-add RMSNorm

SwiGLU activation

GEMV

Softmax

API contract

Inputs

Triton CUDA backend

RMSNorm

Fused residual-add RMSNorm

SwiGLU activation

GEMV

Softmax

TileLang-Ascend backend

CUDA performance

SwiGLU baseline

GEMV measurements

Softmax measurements

Correctness

Numerical tolerances

Benchmarks

Repository layout

Scope

About

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages