From 27cff30833ee145e97a6c84dc976b3def08eb2da Mon Sep 17 00:00:00 2001
From: Laura Wang <3700467+Laurawly@users.noreply.github.com>
Date: Tue, 6 Jan 2026 10:58:01 -0800
Subject: [PATCH 1/8] Provide a vLLM general plugin that registers
oink::rmsnorm and oink::fused_add_rms_norm backed by an SM100 CuTeDSL
RMSNorm kernel.
The ops are torch.compile-friendly (stride-preserving for padded-row inputs)
and the fused op matches vLLM's in-place residual-add RMSNorm semantics.
---
oink/README.md | 57 +
oink/pyproject.toml | 29 +
oink/src/kernelagent_oink/__init__.py | 95 +
.../kernelagent_oink/blackwell/__init__.py | 3 +
.../kernelagent_oink/blackwell/lite_quack.py | 350 +++
.../blackwell/oink_custom_ops.py | 224 ++
.../src/kernelagent_oink/blackwell/rmsnorm.py | 2660 +++++++++++++++++
7 files changed, 3418 insertions(+)
create mode 100644 oink/README.md
create mode 100644 oink/pyproject.toml
create mode 100644 oink/src/kernelagent_oink/__init__.py
create mode 100644 oink/src/kernelagent_oink/blackwell/__init__.py
create mode 100644 oink/src/kernelagent_oink/blackwell/lite_quack.py
create mode 100644 oink/src/kernelagent_oink/blackwell/oink_custom_ops.py
create mode 100644 oink/src/kernelagent_oink/blackwell/rmsnorm.py
diff --git a/oink/README.md b/oink/README.md
new file mode 100644
index 0000000..427f69f
--- /dev/null
+++ b/oink/README.md
@@ -0,0 +1,57 @@
+# KernelAgent Oink (vLLM plugin)
+
+This subproject provides an **out-of-tree vLLM plugin** that registers
+`torch.library.custom_op` entrypoints under the `oink::` namespace:
+
+- `torch.ops.oink.rmsnorm`
+- `torch.ops.oink.fused_add_rms_norm`
+
+The implementation is backed by a CuTeDSL (CUTLASS) RMSNorm kernel tuned for
+**NVIDIA Blackwell (SM100)**.
+
+## Install (editable)
+
+From the `KernelAgent` repo root:
+
+```bash
+pip install -e ./oink
+```
+
+This plugin requires the CuTeDSL stack:
+
+```bash
+pip install nvidia-cutlass-dsl cuda-python
+```
+
+## Use with vLLM
+
+1. Enable the vLLM integration:
+
+```bash
+export VLLM_USE_OINK_RMSNORM=1
+```
+
+2. Ensure vLLM keeps `rms_norm` as a custom op when using `torch.compile` /
+CUDA graphs. In Python:
+
+```python
+from vllm import LLM
+
+llm = LLM(
+ model=...,
+ tensor_parallel_size=...,
+ enforce_eager=False,
+ compilation_config={"custom_ops": ["none", "+rms_norm"]},
+)
+```
+
+Without `+rms_norm`, Inductor may fuse RMSNorm into larger Triton kernels and
+neither vLLM's CUDA RMSNorm nor Oink will run.
+
+## Notes
+
+- This plugin is designed to be **safe to import even when disabled**; it only
+ registers ops when `VLLM_USE_OINK_RMSNORM` is truthy (`"1"` / `"true"`).
+- The ops preserve **padded-row layouts** for 2D tensors (shape `[M, N]`,
+ `stride(1) == 1`, and potentially `stride(0) > N`), which is required for
+ `torch.compile` stride verification on some models (e.g., MLA padded inputs).
diff --git a/oink/pyproject.toml b/oink/pyproject.toml
new file mode 100644
index 0000000..a9ec306
--- /dev/null
+++ b/oink/pyproject.toml
@@ -0,0 +1,29 @@
+[build-system]
+requires = ["setuptools>=61.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "kernelagent-oink"
+version = "0.1.0"
+description = "vLLM plugin that registers Oink Blackwell RMSNorm custom ops"
+readme = "README.md"
+requires-python = ">=3.10"
+license = {text = "Apache-2.0"}
+authors = [{name = "PyTorch Labs"}]
+
+# Keep dependencies minimal, but include the CuTeDSL stack required by the
+# Blackwell RMSNorm implementation.
+#
+# We intentionally do NOT depend on `torch` here because vLLM already pins and
+# provides a compatible PyTorch build.
+dependencies = [
+ "nvidia-cutlass-dsl",
+ "cuda-python",
+]
+
+[project.entry-points."vllm.general_plugins"]
+oink = "kernelagent_oink:register"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+include = ["kernelagent_oink*"]
diff --git a/oink/src/kernelagent_oink/__init__.py b/oink/src/kernelagent_oink/__init__.py
new file mode 100644
index 0000000..542e59e
--- /dev/null
+++ b/oink/src/kernelagent_oink/__init__.py
@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+import logging
+import os
+
+logger = logging.getLogger(__name__)
+
+_OPS_REGISTERED = False
+
+
+def _env_truthy(name: str) -> bool:
+ val = os.environ.get(name)
+ if val is None:
+ return False
+ return val.strip().lower() in ("1", "true", "yes", "on")
+
+
+def _infer_cuda_device_index() -> int:
+ local_rank = os.environ.get("LOCAL_RANK")
+ if local_rank is not None:
+ try:
+ return int(local_rank)
+ except ValueError:
+ pass
+ return 0
+
+
+def _compute_cutedsl_arch(major: int, minor: int) -> str:
+ # CuTeDSL uses an "a" suffix for >= Hopper.
+ suffix = "a" if major >= 9 else ""
+ # Match cutlass/base_dsl/env_manager.py: map sm_110 -> sm_101.
+ if major == 11 and minor == 0:
+ major, minor = 10, 1
+ return f"sm_{major}{minor}{suffix}"
+
+
+def register() -> None:
+ """vLLM plugin entrypoint.
+
+ This function must be safe to call multiple times and must not raise.
+ vLLM executes it in multiple processes (engine + workers).
+ """
+ global _OPS_REGISTERED
+
+ if _OPS_REGISTERED:
+ return
+
+ # Gate on the vLLM integration flag so installing the package does not
+ # change behavior unless explicitly enabled.
+ if not _env_truthy("VLLM_USE_OINK_RMSNORM"):
+ return
+
+ try:
+ import torch
+ except Exception as e: # pragma: no cover
+ logger.debug("Oink plugin: torch import failed: %s", e)
+ return
+
+ try:
+ if not torch.cuda.is_available():
+ return
+ device_index = _infer_cuda_device_index()
+ major, minor = torch.cuda.get_device_capability(device_index)
+ sm = 10 * int(major) + int(minor)
+ if sm < 100:
+ return
+
+ # Ensure required deps are importable before registering ops so that vLLM
+ # doesn't detect ops that would later fail at first use.
+ try:
+ import cutlass # noqa: F401
+ import cuda.bindings.driver as _cuda # noqa: F401
+ except Exception as e:
+ logger.warning(
+ "Oink plugin: CuTeDSL deps missing; skipping op registration. "
+ "Install `nvidia-cutlass-dsl` + `cuda-python`. Error: %s",
+ e,
+ )
+ return
+
+ # Ensure CuTeDSL sees a target arch early. If the user has already set it,
+ # respect their choice.
+ os.environ.setdefault("CUTE_DSL_ARCH", _compute_cutedsl_arch(int(major), int(minor)))
+
+ # Import registers the ops via torch.library.custom_op decorators.
+ from .blackwell import oink_custom_ops # noqa: F401
+ except Exception as e: # pragma: no cover
+ # Do not raise: vLLM plugin loader does not guard plugin execution.
+ logger.exception("Oink plugin: failed to register ops: %s", e)
+ return
+
+ _OPS_REGISTERED = True
+
+
+__all__ = ["register"]
diff --git a/oink/src/kernelagent_oink/blackwell/__init__.py b/oink/src/kernelagent_oink/blackwell/__init__.py
new file mode 100644
index 0000000..4d21ee8
--- /dev/null
+++ b/oink/src/kernelagent_oink/blackwell/__init__.py
@@ -0,0 +1,3 @@
+from __future__ import annotations
+
+__all__ = []
diff --git a/oink/src/kernelagent_oink/blackwell/lite_quack.py b/oink/src/kernelagent_oink/blackwell/lite_quack.py
new file mode 100644
index 0000000..3c3f750
--- /dev/null
+++ b/oink/src/kernelagent_oink/blackwell/lite_quack.py
@@ -0,0 +1,350 @@
+"""
+Lightweight local clone of the small subset of helpers that the SM100
+RMSNorm CuteDSL kernels depend on.
+
+This module intentionally avoids importing the `quack` package so that
+Oink Blackwell kernels can run without Quack installed, while keeping
+numerical behaviour and performance close to the original reference
+implementations.
+"""
+
+from __future__ import annotations
+
+import math
+import operator
+from typing import Callable, Optional, Tuple
+
+import cuda.bindings.driver as cuda # type: ignore
+import torch
+from torch import Tensor
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Float32, Int32, const_expr
+from cutlass.cute.runtime import from_dlpack
+from cutlass.cutlass_dsl import T, dsl_user_op
+from cutlass._mlir.dialects import llvm, nvvm, vector
+
+
+# -------------------------
+# Dtype mapping
+# -------------------------
+
+TORCH2CUTE_DTYPE = {
+ torch.float16: cutlass.Float16,
+ torch.bfloat16: cutlass.BFloat16,
+ torch.float32: cutlass.Float32,
+}
+
+
+# -------------------------
+# Tensor conversion helpers
+# -------------------------
+
+def convert_from_dlpack(
+ x: Tensor,
+ leading_dim: int,
+ alignment: int = 16,
+ divisibility: int = 1,
+) -> cute.Tensor:
+ """
+ Wrap a torch.Tensor in a CuteDSL tensor with layout metadata that
+ matches the logical leading dimension and alignment/divisibility
+ constraints expected by SM100 kernels.
+ """
+ return (
+ from_dlpack(x, assumed_align=alignment)
+ .mark_layout_dynamic(leading_dim=leading_dim)
+ .mark_compact_shape_dynamic(
+ mode=leading_dim,
+ stride_order=x.dim_order(),
+ divisibility=divisibility,
+ )
+ )
+
+
+# -------------------------
+# SM90/SM100 cluster helpers
+# -------------------------
+
+
+@dsl_user_op
+def elem_pointer(x: cute.Tensor, coord: cute.Coord, *, loc=None, ip=None) -> cute.Pointer:
+ return x.iterator + cute.crd2idx(coord, x.layout, loc=loc, ip=ip)
+
+
+@dsl_user_op
+def set_block_rank(
+ smem_ptr: cute.Pointer,
+ peer_cta_rank_in_cluster: cute.Int32,
+ *,
+ loc=None,
+ ip=None,
+) -> cutlass.Int32:
+ """Map the given smem pointer to the address at another CTA rank in the cluster."""
+ smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value()
+ return cutlass.Int32(
+ llvm.inline_asm(
+ T.i32(),
+ [smem_ptr_i32, peer_cta_rank_in_cluster.ir_value()],
+ "mapa.shared::cluster.u32 $0, $1, $2;",
+ "=r,r,r",
+ has_side_effects=False,
+ is_align_stack=False,
+ asm_dialect=llvm.AsmDialect.AD_ATT,
+ )
+ )
+
+
+@dsl_user_op
+def store_shared_remote(
+ val: float | Float32 | Int32 | cutlass.Int64,
+ smem_ptr: cute.Pointer,
+ mbar_ptr: cute.Pointer,
+ peer_cta_rank_in_cluster: cute.typing.Int,
+ *,
+ loc=None,
+ ip=None,
+) -> None:
+ remote_smem_ptr_i32 = set_block_rank(
+ smem_ptr,
+ peer_cta_rank_in_cluster,
+ loc=loc,
+ ip=ip,
+ ).ir_value()
+ remote_mbar_ptr_i32 = set_block_rank(
+ mbar_ptr,
+ peer_cta_rank_in_cluster,
+ loc=loc,
+ ip=ip,
+ ).ir_value()
+ if const_expr(isinstance(val, float)):
+ val = Float32(val)
+ assert isinstance(val, (Float32, Int32, cutlass.Int64)), "val must be Float32, Int32, or Int64"
+ suffix = {Float32: "f32", Int32: "s32", cutlass.Int64: "s64"}[type(val)]
+ constraint = {Float32: "f", Int32: "r", cutlass.Int64: "l"}[type(val)]
+ llvm.inline_asm(
+ None,
+ [remote_smem_ptr_i32, val.ir_value(loc=loc, ip=ip), remote_mbar_ptr_i32],
+ f"st.async.shared::cluster.mbarrier::complete_tx::bytes.{suffix} [$0], $1, [$2];",
+ f"r,{constraint},r",
+ has_side_effects=True,
+ is_align_stack=False,
+ asm_dialect=llvm.AsmDialect.AD_ATT,
+ )
+
+
+@cute.jit
+def predicate_k(tAcA: cute.Tensor, limit: cutlass.Int32) -> cute.Tensor:
+ """
+ Build a predicate tensor for the K dimension only. Values beyond
+ `limit` are masked out.
+ """
+ tApA = cute.make_fragment(
+ cute.make_layout(
+ (cute.size(tAcA, mode=[0, 1]), cute.size(tAcA, mode=[1]), cute.size(tAcA, mode=[2])),
+ stride=(cute.size(tAcA, mode=[2]), 0, 1),
+ ),
+ cutlass.Boolean,
+ )
+ for rest_v in cutlass.range_constexpr(tApA.shape[0]):
+ for rest_k in cutlass.range_constexpr(tApA.shape[2]):
+ tApA[rest_v, 0, rest_k] = cute.elem_less(tAcA[(0, rest_v), 0, rest_k][1], limit)
+ return tApA
+
+
+@dsl_user_op
+def domain_offset_i64(coord: cute.Coord, tensor: cute.Tensor, *, loc=None, ip=None) -> cute.Tensor:
+ """
+ Return a tensor whose iterator is offset by an Int64 byte offset
+ computed from `coord` and the tensor's strides.
+ """
+ flat_coord_i64 = tuple(cutlass.Int64(c) for c in cute.flatten(coord))
+ flat_stride = cute.flatten_to_tuple(tensor.stride)
+ assert len(flat_coord_i64) == len(flat_stride), (
+ "Coordinate and stride must have the same length"
+ )
+ offset = sum(c * s for c, s in zip(flat_coord_i64, flat_stride))
+ assert isinstance(tensor.iterator, cute.Pointer)
+ new_ptr = cute.make_ptr(
+ tensor.element_type,
+ tensor.iterator.toint() + offset * tensor.element_type.width // 8,
+ tensor.memspace,
+ assumed_align=tensor.iterator.max_alignment,
+ )
+ return cute.make_tensor(new_ptr, tensor.layout)
+
+
+# -------------------------
+# Reduction helpers
+# -------------------------
+
+
+@cute.jit
+def warp_reduce(
+ val: cute.TensorSSA | cute.Numeric,
+ op: Callable,
+ width: cutlass.Constexpr[int] = cute.arch.WARP_SIZE,
+) -> cute.TensorSSA | cute.Numeric:
+ """
+ Warp-level reduction for either scalar values or small TensorSSA
+ fragments.
+ """
+ if cutlass.const_expr(isinstance(val, cute.TensorSSA)):
+ res = cute.make_fragment(val.shape, val.dtype)
+ res.store(val)
+ for i in cutlass.range_constexpr(cute.size(val.shape)):
+ res[i] = warp_reduce(res[i], op, width)
+ return res.load()
+ for i in cutlass.range_constexpr(int(math.log2(width))):
+ val = op(val, cute.arch.shuffle_sync_bfly(val, offset=1 << i))
+ return val
+
+
+@cute.jit
+def block_reduce(
+ val: cute.Numeric,
+ op: Callable,
+ reduction_buffer: cute.Tensor,
+ init_val: cute.Numeric = 0.0,
+) -> cute.Numeric:
+ """Block-level reduction across warps."""
+ lane_idx, warp_idx = cute.arch.lane_idx(), cute.arch.warp_idx()
+ warps_per_row = cute.size(reduction_buffer.shape[1])
+ row_idx, col_idx = warp_idx // warps_per_row, warp_idx % warps_per_row
+ if lane_idx == 0:
+ reduction_buffer[row_idx, col_idx] = val
+ cute.arch.barrier()
+ block_reduce_val = init_val
+ if lane_idx < warps_per_row:
+ block_reduce_val = reduction_buffer[row_idx, lane_idx]
+ return warp_reduce(block_reduce_val, op)
+
+
+@cute.jit
+def cluster_reduce(
+ val: cute.Numeric,
+ op: Callable,
+ reduction_buffer: cute.Tensor,
+ mbar_ptr: cute.Pointer,
+ init_val: cute.Numeric = 0.0,
+ phase: Optional[cutlass.Int32] = None,
+) -> cute.Numeric:
+ """
+ Cluster-wide reduction using shared memory and mbarrier. The
+ reduction_buffer has shape (rows_per_block, (warps_per_row, cluster_n)).
+ """
+ cta_rank_in_cluster = cute.arch.block_idx_in_cluster()
+ lane_idx, warp_idx = cute.arch.lane_idx(), cute.arch.warp_idx()
+ rows_per_block, (warps_per_row, cluster_n) = reduction_buffer.shape
+ row_idx, col_idx = warp_idx // warps_per_row, warp_idx % warps_per_row
+ if warp_idx == 0:
+ with cute.arch.elect_one():
+ num_warps = rows_per_block * warps_per_row
+ cute.arch.mbarrier_arrive_and_expect_tx(
+ mbar_ptr,
+ num_warps * cluster_n * reduction_buffer.element_type.width // 8,
+ )
+ if lane_idx < cluster_n:
+ store_shared_remote(
+ val,
+ elem_pointer(reduction_buffer, (row_idx, (col_idx, cta_rank_in_cluster))),
+ mbar_ptr,
+ peer_cta_rank_in_cluster=lane_idx,
+ )
+ cute.arch.mbarrier_wait(mbar_ptr, phase=phase if phase is not None else 0)
+ block_reduce_val = init_val
+ num_iter = cute.ceil_div(warps_per_row * cluster_n, cute.arch.WARP_SIZE)
+ for i in cutlass.range_constexpr(num_iter):
+ idx = lane_idx + i * cute.arch.WARP_SIZE
+ if idx < cute.size(reduction_buffer, mode=[1]):
+ block_reduce_val = op(block_reduce_val, reduction_buffer[row_idx, idx])
+ return warp_reduce(block_reduce_val, op)
+
+
+@cute.jit
+def block_or_cluster_reduce(
+ val: cute.Numeric,
+ op: Callable,
+ reduction_buffer: cute.Tensor,
+ mbar_ptr: Optional[cute.Pointer],
+ phase: Optional[cutlass.Int32] = None,
+ init_val: cute.Numeric = 0.0,
+) -> cute.Numeric:
+ """Dispatch between block or cluster reduction depending on mbar_ptr."""
+ if cutlass.const_expr(mbar_ptr is None):
+ return block_reduce(val, op, reduction_buffer, init_val=init_val)
+ return cluster_reduce(val, op, reduction_buffer, mbar_ptr, init_val=init_val, phase=phase)
+
+
+@cute.jit
+def row_reduce(
+ x: cute.TensorSSA | cute.Numeric,
+ op: cute.ReductionOp,
+ threads_per_row: cutlass.Constexpr[int],
+ reduction_buffer: Optional[cute.Tensor] = None,
+ mbar_ptr: Optional[cute.Pointer] = None,
+ phase: Optional[cutlass.Int32] = None,
+ init_val: cute.Numeric = 0.0,
+ hook_fn: Optional[Callable] = None,
+) -> cute.Numeric:
+ """
+ Row-wise reduction used by RMSNorm and similar kernels.
+
+ reduction_buffer must have shape
+ (num_warps / warps_per_row, (warps_per_row, cluster_n)).
+ """
+ if cutlass.const_expr(isinstance(x, cute.TensorSSA)):
+ val = x.reduce(op, init_val=init_val, reduction_profile=0)
+ else:
+ val = x
+ warp_op = {
+ cute.ReductionOp.ADD: operator.add,
+ cute.ReductionOp.MAX: cute.arch.fmax if cutlass.const_expr(x.dtype == Float32) else max,
+ cute.ReductionOp.MIN: min,
+ cute.ReductionOp.MUL: operator.mul,
+ }[op]
+ val = warp_reduce(
+ val,
+ warp_op,
+ width=min(threads_per_row, cute.arch.WARP_SIZE),
+ )
+ if cutlass.const_expr(hook_fn is not None):
+ hook_fn()
+ if cutlass.const_expr(reduction_buffer is not None):
+ warps_per_row, cluster_n = reduction_buffer.shape[1]
+ assert cluster_n == 1 or mbar_ptr is not None, (
+ "mbar_ptr must be provided for cluster reduction"
+ )
+ if cutlass.const_expr(warps_per_row > 1 or cluster_n > 1):
+ val = block_or_cluster_reduce(
+ val,
+ warp_op,
+ reduction_buffer,
+ mbar_ptr,
+ phase=phase,
+ init_val=init_val,
+ )
+ return val
+
+
+# -------------------------
+# SM count helper
+# -------------------------
+
+
+def get_sm_count(N: int, device: torch.device) -> int:
+ """
+ Heuristic for the number of persistent CTAs (sm_count) based on N and
+ the GPU's SM count. This mirrors the behaviour used in Quack's
+ RMSNorm kernels but lives entirely in this local module.
+ """
+ sm_count_multiple = (
+ 16 if N <= 256 else (8 if N <= 1024 else (4 if N <= 2048 else (2 if N <= 4096 else 1)))
+ )
+ sm_count = torch.cuda.get_device_properties(device).multi_processor_count
+ sm_count = (
+ sm_count * sm_count_multiple if N <= 8192 else sm_count // 2 if N <= 16384 else sm_count * 2
+ )
+ return sm_count
+
diff --git a/oink/src/kernelagent_oink/blackwell/oink_custom_ops.py b/oink/src/kernelagent_oink/blackwell/oink_custom_ops.py
new file mode 100644
index 0000000..8225025
--- /dev/null
+++ b/oink/src/kernelagent_oink/blackwell/oink_custom_ops.py
@@ -0,0 +1,224 @@
+from __future__ import annotations
+
+"""
+Torch custom ops wrapping Oink's Blackwell RMSNorm kernels.
+
+These ops are designed to be:
+- Architecture-aware (use CuTeDSL SM100 kernels when available, fall back
+ to a safe reference elsewhere).
+- Layout-preserving for 2D row-major inputs, including padded MLA-style
+ layouts where stride(0) > N and stride(1) == 1.
+- torch.compile-friendly via proper fake implementations that mirror
+ runtime shapes and strides.
+
+Public ops (Python signatures):
+
+ torch.ops.oink.rmsnorm(x: Tensor, weight: Tensor, eps: float) -> Tensor
+ Functional RMSNorm. Returns a new tensor with the same shape and
+ stride as x when using the fast CuTeDSL path.
+
+ torch.ops.oink.fused_add_rms_norm(
+ x: Tensor, residual: Tensor, weight: Tensor, eps: float
+ ) -> None
+ In-place fused residual-add + RMSNorm matching vLLM semantics:
+ residual = x + residual (stored into `residual`)
+ x = RMSNorm(residual, w) (stored into `x`)
+ Mutates `x` and `residual` in-place and returns None.
+"""
+
+import importlib
+import threading
+
+import torch
+from torch.library import custom_op
+
+_RMSNORM_MOD: object | None = None
+_RMSNORM_MOD_LOCK = threading.Lock()
+
+
+def _get_rmsnorm_mod():
+ """Lazy import to keep plugin registration lightweight.
+
+ Importing the CuTeDSL kernel stack can be expensive and may require a CUDA
+ context. We defer it until the first actual execution of the custom op.
+ """
+ global _RMSNORM_MOD
+
+ cached = _RMSNORM_MOD
+ if cached is not None:
+ return cached
+
+ with _RMSNORM_MOD_LOCK:
+ if _RMSNORM_MOD is None:
+ _RMSNORM_MOD = importlib.import_module("kernelagent_oink.blackwell.rmsnorm")
+ return _RMSNORM_MOD
+
+
+def _get_sm(device: torch.device | None = None) -> int:
+ """Return SM version as an int (e.g., 100 for SM100 / Blackwell)."""
+ if device is None:
+ device = torch.device("cuda")
+ major, minor = torch.cuda.get_device_capability(device)
+ return 10 * major + minor
+
+
+#
+# RMSNorm (functional)
+#
+
+@custom_op("oink::rmsnorm", mutates_args=())
+def oink_rmsnorm(
+ x: torch.Tensor,
+ weight: torch.Tensor,
+ eps: float,
+) -> torch.Tensor:
+ """
+ Functional RMSNorm entrypoint.
+
+ This op is model-agnostic. It expects a 2D [M, N] view of the input
+ where the last dimension is contiguous (stride(1) == 1). The leading
+ dimension stride(0) may be larger than N (padded-row layouts), and
+ will be preserved on the fast CuTeDSL path.
+
+ On SM100 (and newer), this dispatches to the tuned CuTeDSL Blackwell
+ RMSNorm kernel in rmsnorm.rmsnorm_forward, which in turn selects the
+ best internal schedule (including DSv3-specific stage-2 kernels where
+ applicable) and preserves the input's 2D stride when using the
+ pointer-based path.
+
+ On older architectures it falls back to a safe PyTorch reference
+ implementation for correctness.
+ """
+ assert x.is_cuda, "oink::rmsnorm requires CUDA tensors"
+ assert x.dim() == 2, "oink::rmsnorm expects a 2D [M, N] tensor view"
+ assert weight.dim() == 1, "weight must be 1D [N]"
+
+ sm = _get_sm(x.device)
+ if sm >= 100:
+ # Use the tuned CuTeDSL SM100 kernel. The public API already
+ # contains all necessary gating and layout checks internally.
+ _rms = _get_rmsnorm_mod()
+ y, _rstd, _res = _rms.rmsnorm_forward(
+ x,
+ weight=weight,
+ bias=None,
+ residual=None,
+ eps=eps,
+ store_rstd=False,
+ )
+ return y
+
+ # Fallback: reference implementation (correctness-first).
+ _rms = _get_rmsnorm_mod()
+ return _rms.rmsnorm_ref(
+ x,
+ w=weight,
+ b=None,
+ residual=None,
+ eps=eps,
+ )
+
+
+@oink_rmsnorm.register_fake
+def oink_rmsnorm_fake(
+ x: torch.Tensor,
+ weight: torch.Tensor,
+ eps: float,
+) -> torch.Tensor:
+ """
+ Fake (meta) implementation for oink::rmsnorm.
+
+ We must preserve x's logical layout (shape + stride) so that Inductor's
+ CUDA graph capture sees the same stride contract as the real kernel.
+ """
+ # x is a FakeTensor here; x.shape/x.stride()/x.device/x.dtype are defined.
+ return torch.empty_strided(
+ x.shape,
+ x.stride(),
+ device=x.device,
+ dtype=x.dtype,
+ )
+
+
+#
+# Fused residual-add + RMSNorm (in-place, vLLM semantics)
+#
+
+@custom_op("oink::fused_add_rms_norm", mutates_args=("x", "residual"))
+def oink_fused_add_rms_norm(
+ x: torch.Tensor,
+ residual: torch.Tensor,
+ weight: torch.Tensor,
+ eps: float,
+) -> None:
+ """
+ In-place fused residual-add + RMSNorm:
+
+ residual <- x + residual
+ x <- RMSNorm(residual, weight, eps)
+
+ Returns:
+ None (mutates `x` and `residual` in-place).
+ """
+ assert x.is_cuda and residual.is_cuda, "oink::fused_add_rms_norm requires CUDA tensors"
+ assert x.shape == residual.shape, "x and residual must have the same shape"
+ assert x.dtype == residual.dtype, "x and residual must have the same dtype"
+ assert weight.dim() == 1, "weight must be 1D [N]"
+
+ sm = _get_sm(x.device)
+ if sm >= 100:
+ _rms = _get_rmsnorm_mod()
+ # Prefer the lowest-overhead in-place entrypoint (returns None).
+ if hasattr(_rms, "fused_add_rmsnorm_inplace_"):
+ _rms.fused_add_rmsnorm_inplace_( # type: ignore[misc]
+ x,
+ residual,
+ weight,
+ eps=eps,
+ )
+ return None
+ # Backward-compatible wrapper (returns (x, residual)).
+ if hasattr(_rms, "fused_add_rmsnorm_forward_inplace"):
+ _rms.fused_add_rmsnorm_forward_inplace( # type: ignore[misc]
+ x,
+ residual,
+ weight,
+ eps=eps,
+ )
+ return None
+
+ # Extremely defensive fallback if the Oink module doesn't provide
+ # the in-place entrypoint.
+ y, z = _rms.fused_add_rmsnorm_forward(x, residual, weight, eps=eps)
+ x.copy_(y)
+ residual.copy_(z)
+ return None
+
+ # Non-SM100 fallback: keep semantics in-place (correctness-first).
+ residual.add_(x)
+ _rms = _get_rmsnorm_mod()
+ y = _rms.rmsnorm_ref(residual, w=weight, b=None, residual=None, eps=eps)
+ x.copy_(y)
+ return None
+
+
+@oink_fused_add_rms_norm.register_fake
+def oink_fused_add_rms_norm_fake(
+ x: torch.Tensor,
+ residual: torch.Tensor,
+ weight: torch.Tensor,
+ eps: float,
+) -> None:
+ """
+ Fake (meta) implementation for oink::fused_add_rms_norm.
+
+ Because this op mutates its inputs in-place, the outputs alias the input
+ buffers and therefore have the same shapes and strides.
+ """
+ return None
+
+
+__all__ = [
+ "oink_rmsnorm",
+ "oink_fused_add_rms_norm",
+]
diff --git a/oink/src/kernelagent_oink/blackwell/rmsnorm.py b/oink/src/kernelagent_oink/blackwell/rmsnorm.py
new file mode 100644
index 0000000..d6c2c20
--- /dev/null
+++ b/oink/src/kernelagent_oink/blackwell/rmsnorm.py
@@ -0,0 +1,2660 @@
+"""
+RMSNorm kernel for SM100 (Blackwell) in CuteDSL.
+
+This implementation targets Blackwell with:
+- A stride-preserving pointer path for padded-row layouts (e.g. MLA stride0> N).
+- A one-pass fused-add RMSNorm schedule for bf16/fp16 (DSv3 N=7168) that keeps
+ `x + residual` in registers (avoids re-reading gmem) and uses FP32 accumulation.
+- Optional experimental schedule knobs (env vars) to explore copy widths and
+ stage-2 cp.async variants.
+
+Note: This file expects the local CuTeDSL (cutlass) and SM100 helper modules
+to be available in the Python environment (e.g., `nvidia-cutlass-dsl` and
+`cuda-python`). It is shipped as part of the KernelAgent Oink vLLM plugin.
+"""
+
+from __future__ import annotations
+
+import ctypes
+import importlib.metadata
+import os
+import re
+import subprocess
+import sys
+import threading
+from typing import Optional, Tuple
+
+_HERE = os.path.dirname(__file__)
+
+# CuTeDSL caches generated MLIR into a tempdir under a global default
+# (`/tmp/$USER/cutlass_python_cache`). The cache bytecode format can differ across
+# `nvidia-cutlass-dsl` versions (e.g. 4.3.2 vs 4.3.4), and cross-version cache
+# sharing causes noisy "invalid section ID" warnings (and disables cache reuse).
+#
+# If the user has not pinned `CUTE_DSL_CACHE_DIR`, isolate by version so multiple
+# CuTeDSL envs can coexist on the same machine without stepping on each other.
+if "CUTE_DSL_CACHE_DIR" not in os.environ:
+ try:
+ _dsl_ver = importlib.metadata.version("nvidia-cutlass-dsl")
+ except Exception:
+ _dsl_ver = "unknown"
+ _dsl_ver = re.sub(r"[^0-9A-Za-z]+", "_", _dsl_ver)
+ _user = os.environ.get("USER") or os.environ.get("USERNAME") or "user"
+ _tmp = os.environ.get("TMPDIR") or "/tmp"
+ os.environ["CUTE_DSL_CACHE_DIR"] = os.path.join(
+ _tmp, _user, f"cutlass_python_cache_{_dsl_ver}"
+ )
+
+try:
+ import cutlass # type: ignore # noqa: F401
+except Exception as e:
+ raise ImportError(
+ "kernelagent_oink.blackwell.rmsnorm requires CuTeDSL's Python package "
+ "(`cutlass`, typically provided by `nvidia-cutlass-dsl`)."
+ ) from e
+
+import torch
+from torch import Tensor
+
+import cuda.bindings.driver as cuda # provided by NVIDIA cuda-python
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Float32, Int32, const_expr
+from cutlass.cute import runtime as rt
+
+# Simple compile cache declared early so direct execution works
+_PTR_COMPILE_CACHE = {}
+
+# Thread-local cache for the fast-launch path. We keep per-thread packed args and
+# pointer/scalar storage so concurrent callers don't race on in-place updates.
+_PTR_FAST_LAUNCH_TLS = threading.local()
+
+def _env_flag(name: str, default: bool) -> bool:
+ val = os.environ.get(name)
+ if val is None:
+ return default
+ return val.strip().lower() not in {"0", "false", "no", "off", ""}
+
+
+# Fast-launch uses a few private-ish CuTeDSL internals (packed args plumbing and
+# runtime pointer descriptors). Keep it enabled by default for our pinned CuTeDSL
+# environment, but allow disabling it via env var and auto-disable it if those
+# internals are not present in a future upgrade.
+_ENABLE_FAST_LAUNCH = _env_flag("OINK_CUTEDSL_FAST_LAUNCH", default=True)
+_FAST_LAUNCH_SUPPORTED = True
+
+# Fused-add RMSNorm schedule knobs (read once at import time; set env vars before
+# importing this module if you want to override).
+_DIRECT_GMEM_POLICY = (os.environ.get("OINK_RMSNORM_DIRECT_GMEM", "auto").strip().lower() or "auto")
+_COPY_BITS_POLICY = (os.environ.get("OINK_RMSNORM_COPY_BITS", "auto").strip().lower() or "auto")
+_ENABLE_CLUSTER_ILP = _env_flag("OINK_RMSNORM_ENABLE_CLUSTER_ILP", default=False)
+_ENABLE_CLUSTER_ILP_UNSAFE = _env_flag("OINK_RMSNORM_ENABLE_CLUSTER_ILP_UNSAFE", default=False)
+_ENABLE_TPR256 = _env_flag("OINK_RMSNORM_ENABLE_TPR256", default=False)
+_ENABLE_STAGE2 = _env_flag("OINK_RMSNORM_ENABLE_STAGE2", default=False)
+
+# CuTeDSL stability probe for the experimental cluster_n>1 + direct-GMEM schedule.
+#
+# Some CuTeDSL builds segfault during JIT compilation when combining:
+# - cluster launches (cluster_n>1) and
+# - direct-GMEM loads/stores (no staging SMEM tiles).
+#
+# We keep the schedule gated behind `OINK_RMSNORM_ENABLE_CLUSTER_ILP=1` +
+# `OINK_RMSNORM_ENABLE_CLUSTER_ILP_UNSAFE=1`, and additionally run a one-time
+# out-of-process compile probe so we can safely fall back to the staged SMEM
+# path instead of crashing the parent process.
+#
+# This is (currently) sensitive to the vector width: we have observed
+# reproducible segfaults for the 256b universal-copy path, while the 128b path
+# can succeed. Cache the maximum supported copy width (0 = unsupported).
+_CLUSTER_DIRECT_GMEM_MAX_COPY_BITS: Optional[int] = None
+_CLUSTER_DIRECT_GMEM_PROBE_LOCK = threading.Lock()
+_CLUSTER_DIRECT_GMEM_PROBE_WARNED = False
+
+
+def _probe_cluster_direct_gmem_max_copy_bits() -> int:
+ global _CLUSTER_DIRECT_GMEM_MAX_COPY_BITS
+ global _CLUSTER_DIRECT_GMEM_PROBE_WARNED
+
+ override = os.environ.get("OINK_RMSNORM_CLUSTER_DIRECT_GMEM_MAX_COPY_BITS")
+ if override is not None and override.strip() != "":
+ try:
+ value = int(override)
+ except ValueError:
+ value = 0
+ value = 256 if value >= 256 else 128 if value >= 128 else 0
+ _CLUSTER_DIRECT_GMEM_MAX_COPY_BITS = value
+ return value
+
+ if _CLUSTER_DIRECT_GMEM_MAX_COPY_BITS is not None:
+ return _CLUSTER_DIRECT_GMEM_MAX_COPY_BITS
+
+ with _CLUSTER_DIRECT_GMEM_PROBE_LOCK:
+ if _CLUSTER_DIRECT_GMEM_MAX_COPY_BITS is not None:
+ return _CLUSTER_DIRECT_GMEM_MAX_COPY_BITS
+
+ script_template = r"""
+import os
+
+os.environ["OINK_CUTEDSL_FAST_LAUNCH"] = "0"
+
+import cutlass
+import cutlass.cute as cute
+import cuda.bindings.driver as cuda
+from cutlass import Float32, Int32
+from cutlass.cute import runtime as rt
+
+from kernelagent_oink.blackwell import rmsnorm
+
+N = 7168
+dtype = cutlass.BFloat16
+
+copy_bits = int(os.environ["OINK_PROBE_COPY_BITS"])
+assumed_align = int(os.environ["OINK_PROBE_ASSUMED_ALIGN"])
+
+op = rmsnorm.RMSNormSM100(
+ N,
+ dtype,
+ stage=1,
+ copy_bits=copy_bits,
+ use_async=False,
+ direct_gmem=True,
+)
+op._cluster_n_override = 2 # 2 CTAs per row
+
+ptr_x = rt.make_ptr(dtype, 0, mem_space=rt.AddressSpace.gmem, assumed_align=assumed_align)
+ptr_res = rt.make_ptr(dtype, 0, mem_space=rt.AddressSpace.gmem, assumed_align=assumed_align)
+ptr_w = rt.make_ptr(dtype, 0, mem_space=rt.AddressSpace.gmem, assumed_align=assumed_align)
+
+_ = cute.compile(
+ op.launch_from_ptrs_fused_add_inplace,
+ ptr_x,
+ ptr_w,
+ ptr_res,
+ Int32(4096),
+ Int32(N),
+ Int32(N),
+ cuda.CUstream(0),
+ Float32(1e-6),
+)
+print(f"ok {copy_bits}")
+"""
+
+ env = os.environ.copy()
+ env["PYTHONNOUSERSITE"] = "1"
+
+ def run_probe(copy_bits: int, assumed_align: int):
+ probe_env = env.copy()
+ probe_env["OINK_PROBE_COPY_BITS"] = str(copy_bits)
+ probe_env["OINK_PROBE_ASSUMED_ALIGN"] = str(assumed_align)
+ return subprocess.run(
+ [sys.executable, "-c", script_template],
+ env=probe_env,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.PIPE,
+ text=True,
+ timeout=120.0,
+ )
+
+ proc_256 = None
+ proc_128 = None
+ try:
+ proc_256 = run_probe(256, 32)
+ if proc_256.returncode == 0:
+ max_bits = 256
+ else:
+ proc_128 = run_probe(128, 16)
+ max_bits = 128 if proc_128.returncode == 0 else 0
+ except Exception:
+ max_bits = 0
+
+ if not _CLUSTER_DIRECT_GMEM_PROBE_WARNED and max_bits != 256:
+ _CLUSTER_DIRECT_GMEM_PROBE_WARNED = True
+ if max_bits == 128:
+ print(
+ "Oink: cluster_n>1 + direct_gmem 256b compile probe failed; "
+ "using 128b copies for the cluster ILP schedule.",
+ file=sys.stderr,
+ )
+ if proc_256 is not None and proc_256.stderr:
+ tail = "\n".join(proc_256.stderr.splitlines()[-12:])
+ print(f"Oink: probe stderr tail:\n{tail}", file=sys.stderr)
+ else:
+ rc = (
+ proc_128.returncode
+ if proc_128 is not None
+ else (proc_256.returncode if proc_256 is not None else "unknown")
+ )
+ print(
+ "Oink: cluster_n>1 + direct_gmem compile probe failed; "
+ f"falling back to staged SMEM path (returncode={rc}).",
+ file=sys.stderr,
+ )
+ failing_proc = proc_128 if proc_128 is not None else proc_256
+ if failing_proc is not None and failing_proc.stderr:
+ tail = "\n".join(failing_proc.stderr.splitlines()[-12:])
+ print(f"Oink: probe stderr tail:\n{tail}", file=sys.stderr)
+
+ _CLUSTER_DIRECT_GMEM_MAX_COPY_BITS = max_bits
+ return max_bits
+
+def _parse_version_tuple(version: str) -> Tuple[int, int, int]:
+ parts = version.split(".")
+ nums: list[int] = []
+ for part in parts[:3]:
+ match = re.match(r"^(\d+)", part)
+ nums.append(int(match.group(1)) if match is not None else 0)
+ while len(nums) < 3:
+ nums.append(0)
+ return nums[0], nums[1], nums[2]
+
+
+def _cutlass_dsl_version() -> Optional[Tuple[int, int, int]]:
+ try:
+ return _parse_version_tuple(importlib.metadata.version("nvidia-cutlass-dsl"))
+ except Exception:
+ return None
+
+
+_CUTLASS_DSL_VERSION = _cutlass_dsl_version()
+# CuTeDSL 4.3.4 tightened some kernel argument expectations (notably around
+# passing Layout/Shape/Constexpr objects into @cute.kernel functions). Keep the
+# older signature for 4.3.2, but switch to a 4.3.4-compatible signature when we
+# detect 4.3.4+ (or when version detection is unavailable).
+_KERNEL_ACCEPTS_LAYOUT_ARGS = _CUTLASS_DSL_VERSION is not None and _CUTLASS_DSL_VERSION < (4, 3, 4)
+
+if _ENABLE_CLUSTER_ILP and not _ENABLE_CLUSTER_ILP_UNSAFE:
+ # We have observed reproducible segfaults in some CuTeDSL builds when using
+ # cluster launches for this schedule. Require an explicit UNSAFE opt-in to
+ # avoid accidental crashes.
+ _ENABLE_CLUSTER_ILP = False
+ print(
+ "Oink: OINK_RMSNORM_ENABLE_CLUSTER_ILP requested but disabled by default due to "
+ "known instability; set OINK_RMSNORM_ENABLE_CLUSTER_ILP_UNSAFE=1 to force-enable.",
+ file=sys.stderr,
+ )
+
+
+def _fast_launch_enabled() -> bool:
+ return _ENABLE_FAST_LAUNCH and _FAST_LAUNCH_SUPPORTED
+
+
+def _direct_gmem_from_policy(*, default: bool) -> bool:
+ """Resolve the direct-GMEM schedule flag from the (import-time) policy string."""
+ if _DIRECT_GMEM_POLICY in {"0", "false", "no", "off"}:
+ return False
+ if _DIRECT_GMEM_POLICY in {"1", "true", "yes", "on"}:
+ return True
+ return default
+
+
+def _copy_bits_from_policy(*, default: int, can_use_256: bool) -> int:
+ """Resolve copy width (in bits) from the (import-time) policy string."""
+ if _COPY_BITS_POLICY in {"128"}:
+ return 128
+ if _COPY_BITS_POLICY in {"256"} and can_use_256:
+ return 256
+ return default
+
+
+class _StableI32Arg:
+ """A stable Int32 runtime arg (avoids per-call Int32().__c_pointers__ allocations)."""
+
+ def __init__(self, value: int):
+ self._c_value = ctypes.c_int32(int(value))
+ self._c_pointer = ctypes.cast(ctypes.pointer(self._c_value), ctypes.c_void_p)
+
+ def set(self, value: int) -> None:
+ self._c_value.value = int(value)
+
+ def __c_pointers__(self):
+ return [self._c_pointer]
+
+
+class _StableF32Arg:
+ """A stable Float32 runtime arg (avoids per-call Float32().__c_pointers__ allocations)."""
+
+ def __init__(self, value: float):
+ self._c_value = ctypes.c_float(float(value))
+ self._c_pointer = ctypes.cast(ctypes.pointer(self._c_value), ctypes.c_void_p)
+
+ def set(self, value: float) -> None:
+ self._c_value.value = float(value)
+
+ def __c_pointers__(self):
+ return [self._c_pointer]
+
+
+def _tls_fast_launch_cache() -> dict[tuple[object, ...], object]:
+ cache = getattr(_PTR_FAST_LAUNCH_TLS, "cache", None)
+ if cache is None:
+ cache = {}
+ _PTR_FAST_LAUNCH_TLS.cache = cache
+ return cache
+
+
+def _set_runtime_ptr(ptr: object, device_ptr: int) -> None:
+ # Runtime pointer objects cache a `ctypes.c_void_p` descriptor and pass
+ # its address to the compiled function. Updating `_desc.value` updates
+ # the device pointer without changing the address of the descriptor.
+ #
+ # This relies on internal CuTeDSL runtime pointer fields (`_desc`, `_pointer`,
+ # etc.). If these internals change in a future CuTeDSL upgrade, callers
+ # should catch AttributeError and fall back to the regular launch path.
+ device_ptr = int(device_ptr)
+ ptr._pointer = device_ptr # type: ignore[attr-defined]
+ if getattr(ptr, "_c_pointer", None) is None:
+ ptr.__c_pointers__() # type: ignore[attr-defined]
+ ptr._desc.value = device_ptr # type: ignore[attr-defined]
+
+
+class _PtrRmsnormFastLaunch:
+ def __init__(
+ self,
+ *,
+ compiled: object,
+ executor: object,
+ capi_func: object,
+ ptr_x: object,
+ ptr_w: Optional[object],
+ ptr_out: object,
+ arg_m: _StableI32Arg,
+ arg_n: _StableI32Arg,
+ arg_ld: _StableI32Arg,
+ arg_eps: _StableF32Arg,
+ stream: cuda.CUstream,
+ packed_args: object,
+ keepalive: tuple[object, ...],
+ ):
+ self._compiled = compiled
+ self._executor = executor
+ self._capi_func = capi_func
+ self._ptr_x = ptr_x
+ self._ptr_w = ptr_w
+ self._ptr_out = ptr_out
+ self._arg_m = arg_m
+ self._arg_n = arg_n
+ self._arg_ld = arg_ld
+ self._arg_eps = arg_eps
+ self._stream = stream
+ self._packed_args = packed_args
+ self._keepalive = keepalive
+
+ self._use_fast_launch = True
+
+ self._cuda_result = getattr(executor, "cuda_result", None)
+
+ self._last_x_ptr = -1
+ self._last_w_ptr = -1
+ self._last_out_ptr = -1
+ self._last_m = -1
+ self._last_ld = -1
+ self._last_eps = float("nan")
+
+ def launch(
+ self,
+ *,
+ x: Tensor,
+ weight: Optional[Tensor],
+ out: Tensor,
+ M: int,
+ N: int,
+ ld: int,
+ eps: float,
+ ) -> None:
+ if not _fast_launch_enabled() or not self._use_fast_launch:
+ self._fallback_launch(x=x, weight=weight, out=out, M=M, N=N, ld=ld, eps=eps)
+ return
+
+ x_ptr = x.data_ptr()
+ if x_ptr != self._last_x_ptr:
+ try:
+ _set_runtime_ptr(self._ptr_x, x_ptr)
+ self._last_x_ptr = x_ptr
+ except AttributeError:
+ self._disable_fast_launch()
+ self._fallback_launch(x=x, weight=weight, out=out, M=M, N=N, ld=ld, eps=eps)
+ return
+
+ if self._ptr_w is not None:
+ w_ptr = weight.data_ptr() # type: ignore[union-attr]
+ if w_ptr != self._last_w_ptr:
+ try:
+ _set_runtime_ptr(self._ptr_w, w_ptr)
+ self._last_w_ptr = w_ptr
+ except AttributeError:
+ self._disable_fast_launch()
+ self._fallback_launch(x=x, weight=weight, out=out, M=M, N=N, ld=ld, eps=eps)
+ return
+
+ out_ptr = out.data_ptr()
+ if out_ptr != self._last_out_ptr:
+ try:
+ _set_runtime_ptr(self._ptr_out, out_ptr)
+ self._last_out_ptr = out_ptr
+ except AttributeError:
+ self._disable_fast_launch()
+ self._fallback_launch(x=x, weight=weight, out=out, M=M, N=N, ld=ld, eps=eps)
+ return
+
+ if M != self._last_m:
+ self._arg_m.set(M)
+ self._last_m = M
+ if ld != self._last_ld:
+ self._arg_ld.set(ld)
+ self._last_ld = ld
+ if eps != self._last_eps:
+ self._arg_eps.set(eps)
+ self._last_eps = eps
+
+ # Clear the error slot before launch (mirrors JitExecutor behavior).
+ if self._cuda_result is not None:
+ self._cuda_result.value = 0
+
+ ret = self._capi_func(self._packed_args) # type: ignore[misc]
+ if ret != 0:
+ raise RuntimeError(f"CuTeDSL capi_func returned non-zero: {ret}")
+ if self._cuda_result is not None:
+ err = int(self._cuda_result.value)
+ if err != 0:
+ raise RuntimeError(f"CuTeDSL kernel launch failed (cuda_result={err})")
+
+ def _disable_fast_launch(self) -> None:
+ global _FAST_LAUNCH_SUPPORTED
+ self._use_fast_launch = False
+ _FAST_LAUNCH_SUPPORTED = False
+
+ def _fallback_launch(
+ self,
+ *,
+ x: Tensor,
+ weight: Optional[Tensor],
+ out: Tensor,
+ M: int,
+ N: int,
+ ld: int,
+ eps: float,
+ ) -> None:
+ # If the packed-args or runtime pointer mutation path stops working
+ # (e.g. due to a CuTeDSL upgrade), fall back to the regular call path.
+ dtype = TORCH2CUTE_DTYPE[x.dtype]
+ ptr_x = rt.make_ptr(dtype, x.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ ptr_out = rt.make_ptr(dtype, out.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ ptr_w = (
+ rt.make_ptr(dtype, weight.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ if weight is not None
+ else None
+ )
+ self._compiled(
+ ptr_x,
+ ptr_w,
+ None, # ptr_b
+ None, # ptr_res
+ ptr_out,
+ None, # ptr_res_out
+ None, # ptr_rstd
+ Int32(M),
+ Int32(N),
+ Int32(ld),
+ self._stream,
+ Float32(eps),
+ )
+
+
+class _PtrFusedAddRmsnormFastLaunch:
+ def __init__(
+ self,
+ *,
+ compiled: object,
+ executor: object,
+ capi_func: object,
+ ptr_x: object,
+ ptr_w: object,
+ ptr_res: object,
+ arg_m: _StableI32Arg,
+ arg_n: _StableI32Arg,
+ arg_ld_x: _StableI32Arg,
+ arg_eps: _StableF32Arg,
+ stream: cuda.CUstream,
+ assumed_align: int,
+ packed_args: object,
+ keepalive: tuple[object, ...],
+ ):
+ self._compiled = compiled
+ self._executor = executor
+ self._capi_func = capi_func
+ self._ptr_x = ptr_x
+ self._ptr_w = ptr_w
+ self._ptr_res = ptr_res
+ self._arg_m = arg_m
+ self._arg_n = arg_n
+ self._arg_ld_x = arg_ld_x
+ self._arg_eps = arg_eps
+ self._stream = stream
+ self._assumed_align = int(assumed_align)
+ self._packed_args = packed_args
+ self._keepalive = keepalive
+
+ self._use_fast_launch = True
+
+ self._cuda_result = getattr(executor, "cuda_result", None)
+
+ self._last_x_ptr = -1
+ self._last_w_ptr = -1
+ self._last_res_ptr = -1
+ self._last_m = -1
+ self._last_ld_x = -1
+ self._last_eps = float("nan")
+
+ def launch(
+ self,
+ *,
+ x: Tensor,
+ weight: Tensor,
+ residual: Tensor,
+ M: int,
+ N: int,
+ ld_x: int,
+ eps: float,
+ ) -> None:
+ if not _fast_launch_enabled() or not self._use_fast_launch:
+ self._fallback_launch(
+ x=x, weight=weight, residual=residual, M=M, N=N, ld_x=ld_x, eps=eps
+ )
+ return
+
+ x_ptr = x.data_ptr()
+ if x_ptr != self._last_x_ptr:
+ try:
+ _set_runtime_ptr(self._ptr_x, x_ptr)
+ self._last_x_ptr = x_ptr
+ except AttributeError:
+ self._disable_fast_launch()
+ self._fallback_launch(
+ x=x, weight=weight, residual=residual, M=M, N=N, ld_x=ld_x, eps=eps
+ )
+ return
+
+ w_ptr = weight.data_ptr()
+ if w_ptr != self._last_w_ptr:
+ try:
+ _set_runtime_ptr(self._ptr_w, w_ptr)
+ self._last_w_ptr = w_ptr
+ except AttributeError:
+ self._disable_fast_launch()
+ self._fallback_launch(
+ x=x, weight=weight, residual=residual, M=M, N=N, ld_x=ld_x, eps=eps
+ )
+ return
+
+ res_ptr = residual.data_ptr()
+ if res_ptr != self._last_res_ptr:
+ try:
+ _set_runtime_ptr(self._ptr_res, res_ptr)
+ self._last_res_ptr = res_ptr
+ except AttributeError:
+ self._disable_fast_launch()
+ self._fallback_launch(
+ x=x, weight=weight, residual=residual, M=M, N=N, ld_x=ld_x, eps=eps
+ )
+ return
+
+ if M != self._last_m:
+ self._arg_m.set(M)
+ self._last_m = M
+ if ld_x != self._last_ld_x:
+ self._arg_ld_x.set(ld_x)
+ self._last_ld_x = ld_x
+ if eps != self._last_eps:
+ self._arg_eps.set(eps)
+ self._last_eps = eps
+
+ if self._cuda_result is not None:
+ self._cuda_result.value = 0
+
+ ret = self._capi_func(self._packed_args) # type: ignore[misc]
+ if ret != 0:
+ raise RuntimeError(f"CuTeDSL capi_func returned non-zero: {ret}")
+ if self._cuda_result is not None:
+ err = int(self._cuda_result.value)
+ if err != 0:
+ raise RuntimeError(f"CuTeDSL kernel launch failed (cuda_result={err})")
+
+ def _disable_fast_launch(self) -> None:
+ global _FAST_LAUNCH_SUPPORTED
+ self._use_fast_launch = False
+ _FAST_LAUNCH_SUPPORTED = False
+
+ def _fallback_launch(
+ self,
+ *,
+ x: Tensor,
+ weight: Tensor,
+ residual: Tensor,
+ M: int,
+ N: int,
+ ld_x: int,
+ eps: float,
+ ) -> None:
+ dtype = TORCH2CUTE_DTYPE[x.dtype]
+ ptr_x = rt.make_ptr(
+ dtype,
+ x.data_ptr(),
+ mem_space=rt.AddressSpace.gmem,
+ assumed_align=self._assumed_align,
+ )
+ ptr_res = rt.make_ptr(
+ dtype,
+ residual.data_ptr(),
+ mem_space=rt.AddressSpace.gmem,
+ assumed_align=self._assumed_align,
+ )
+ ptr_w = rt.make_ptr(
+ dtype,
+ weight.data_ptr(),
+ mem_space=rt.AddressSpace.gmem,
+ assumed_align=self._assumed_align,
+ )
+ self._compiled(
+ ptr_x,
+ ptr_w,
+ ptr_res,
+ Int32(M),
+ Int32(N),
+ Int32(ld_x),
+ self._stream,
+ Float32(eps),
+ )
+
+
+def _get_fast_ptr_rmsnorm_launcher(
+ *,
+ compiled: object,
+ dtype: type[cutlass.Numeric],
+ N: int,
+ device_index: int,
+ stream_handle: int,
+ has_weight: bool,
+ eps: float,
+) -> Optional[_PtrRmsnormFastLaunch]:
+ if not _fast_launch_enabled():
+ return None
+ # Keyed by the compiled object identity so schedule changes (e.g. copy width,
+ # async/staged variants, etc.) never alias in the fast-launch cache.
+ key = ("ptr_fast", id(compiled), N, dtype, device_index, int(stream_handle), has_weight)
+ cache = _tls_fast_launch_cache()
+ cached = cache.get(key)
+ if cached is not None:
+ return cached # type: ignore[return-value]
+
+ # Create stable runtime args and pointer descriptors once.
+ ptr_x = rt.make_ptr(dtype, 0, mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ ptr_out = rt.make_ptr(dtype, 0, mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ ptr_w = (
+ rt.make_ptr(dtype, 0, mem_space=rt.AddressSpace.gmem, assumed_align=16) if has_weight else None
+ )
+
+ arg_m = _StableI32Arg(0)
+ arg_n = _StableI32Arg(N)
+ arg_ld = _StableI32Arg(N)
+ arg_eps = _StableF32Arg(eps)
+
+ stream = cuda.CUstream(int(stream_handle))
+
+ # Create an executor (loads the CUDA library once).
+ executor = compiled.to(device_index) # type: ignore[attr-defined]
+
+ # Use generate_execution_args once to build the packed args array, and keep
+ # any adapted args alive for the lifetime of the cache entry.
+ try:
+ exe_args, adapted_args = executor.generate_execution_args(
+ ptr_x,
+ ptr_w,
+ None, # ptr_b
+ None, # ptr_res
+ ptr_out,
+ None, # ptr_res_out
+ None, # ptr_rstd
+ arg_m,
+ arg_n,
+ arg_ld,
+ stream,
+ arg_eps,
+ )
+ packed_args = executor._get_invoke_packed_args(list(exe_args)) # type: ignore[attr-defined]
+ capi_func = compiled.capi_func # type: ignore[attr-defined]
+ except AttributeError:
+ global _FAST_LAUNCH_SUPPORTED
+ _FAST_LAUNCH_SUPPORTED = False
+ return None
+
+ keepalive: tuple[object, ...] = (
+ executor,
+ ptr_x,
+ ptr_w,
+ ptr_out,
+ arg_m,
+ arg_n,
+ arg_ld,
+ arg_eps,
+ stream,
+ *adapted_args,
+ )
+
+ launcher = _PtrRmsnormFastLaunch(
+ compiled=compiled,
+ executor=executor,
+ capi_func=capi_func,
+ ptr_x=ptr_x,
+ ptr_w=ptr_w,
+ ptr_out=ptr_out,
+ arg_m=arg_m,
+ arg_n=arg_n,
+ arg_ld=arg_ld,
+ arg_eps=arg_eps,
+ stream=stream,
+ packed_args=packed_args,
+ keepalive=keepalive,
+ )
+ cache[key] = launcher
+ return launcher
+
+
+def _get_fast_ptr_fused_add_rmsnorm_launcher(
+ *,
+ compiled: object,
+ dtype: type[cutlass.Numeric],
+ N: int,
+ device_index: int,
+ stream_handle: int,
+ copy_bits: int,
+ use_async: bool,
+ tpr: int,
+ direct_gmem: bool,
+ assumed_align: int,
+ eps: float,
+) -> Optional[_PtrFusedAddRmsnormFastLaunch]:
+ if not _fast_launch_enabled():
+ return None
+ key = (
+ "ptr_fused_add_fast",
+ id(compiled),
+ N,
+ dtype,
+ device_index,
+ int(stream_handle),
+ int(copy_bits),
+ bool(use_async),
+ int(tpr),
+ bool(direct_gmem),
+ int(assumed_align),
+ )
+ cache = _tls_fast_launch_cache()
+ cached = cache.get(key)
+ if cached is not None:
+ return cached # type: ignore[return-value]
+
+ ptr_x = rt.make_ptr(dtype, 0, mem_space=rt.AddressSpace.gmem, assumed_align=assumed_align)
+ ptr_res = rt.make_ptr(dtype, 0, mem_space=rt.AddressSpace.gmem, assumed_align=assumed_align)
+ ptr_w = rt.make_ptr(dtype, 0, mem_space=rt.AddressSpace.gmem, assumed_align=assumed_align)
+
+ arg_m = _StableI32Arg(0)
+ arg_n = _StableI32Arg(N)
+ arg_ld_x = _StableI32Arg(N)
+ arg_eps = _StableF32Arg(eps)
+
+ stream = cuda.CUstream(int(stream_handle))
+
+ executor = compiled.to(device_index) # type: ignore[attr-defined]
+
+ try:
+ exe_args, adapted_args = executor.generate_execution_args(
+ ptr_x,
+ ptr_w,
+ ptr_res,
+ arg_m,
+ arg_n,
+ arg_ld_x,
+ stream,
+ arg_eps,
+ )
+ packed_args = executor._get_invoke_packed_args(list(exe_args)) # type: ignore[attr-defined]
+ capi_func = compiled.capi_func # type: ignore[attr-defined]
+ except AttributeError:
+ global _FAST_LAUNCH_SUPPORTED
+ _FAST_LAUNCH_SUPPORTED = False
+ return None
+
+ keepalive: tuple[object, ...] = (
+ executor,
+ ptr_x,
+ ptr_w,
+ ptr_res,
+ arg_m,
+ arg_n,
+ arg_ld_x,
+ arg_eps,
+ stream,
+ *adapted_args,
+ )
+
+ launcher = _PtrFusedAddRmsnormFastLaunch(
+ compiled=compiled,
+ executor=executor,
+ capi_func=capi_func,
+ ptr_x=ptr_x,
+ ptr_w=ptr_w,
+ ptr_res=ptr_res,
+ arg_m=arg_m,
+ arg_n=arg_n,
+ arg_ld_x=arg_ld_x,
+ arg_eps=arg_eps,
+ stream=stream,
+ assumed_align=assumed_align,
+ packed_args=packed_args,
+ keepalive=keepalive,
+ )
+ cache[key] = launcher
+ return launcher
+
+
+# Local helpers for reduction, dtype mapping, and coordinate/predicate utilities.
+#
+# NOTE: Avoid `from . import ...` imports here: CuTeDSL's AST preprocessor may
+# mishandle that form (module=None in the AST). Use fully-qualified imports.
+from kernelagent_oink.blackwell import lite_quack as qutils
+from kernelagent_oink.blackwell.lite_quack import TORCH2CUTE_DTYPE, row_reduce
+
+
+# -------------------------
+# Copy helpers (allow up to 256b)
+# -------------------------
+
+@cute.jit
+def get_copy_atom_bw(
+ dtype: type[cutlass.Numeric], num_copy_elems: int, is_async: bool = False
+) -> cute.CopyAtom:
+ # cp.async (SIMT) supports up to 128b per op; use 256b for sync when possible
+ max_bits = const_expr(128 if is_async else 256)
+ num_copy_bits = const_expr(min(max_bits, num_copy_elems * dtype.width))
+ from cutlass.cute.nvgpu import cpasync
+ # Prefer GLOBAL cache policy for bulk streaming reads at large M
+ copy_op = (
+ cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL)
+ if is_async
+ else cute.nvgpu.CopyUniversalOp()
+ )
+ return cute.make_copy_atom(copy_op, dtype, num_bits_per_copy=num_copy_bits)
+
+
+@cute.jit
+def copy_tiled(
+ src: cute.Tensor,
+ dst: cute.Tensor,
+ *,
+ pred: Optional[cute.Tensor] = None,
+ num_copy_elems: int = 1,
+ is_async: bool = False,
+) -> None:
+ atom = get_copy_atom_bw(src.element_type, num_copy_elems, is_async)
+ cute.copy(atom, src, dst, pred=pred)
+
+
+# -------------------------
+# RMSNorm Kernel (SM100)
+# -------------------------
+
+
+class RMSNormSM100:
+ def __init__(
+ self,
+ N: int,
+ dtype: type[cutlass.Numeric],
+ stage: Optional[int] = None,
+ *,
+ copy_bits: int = 128,
+ use_async: bool = True,
+ direct_gmem: bool = False,
+ ):
+ self.N = N
+ self.dtype = dtype
+ # Match Quack default for RMSNorm: stage = 1 unless explicitly overridden
+ self.stage = 1 if stage is None else stage
+ self.reduction_dtype = cutlass.Float32
+ self.copy_bits = int(copy_bits)
+ self.use_async = bool(use_async)
+ self.direct_gmem = bool(direct_gmem)
+
+ def _threads_per_row(self) -> int:
+ try:
+ return self._tpr_override # type: ignore[attr-defined]
+ except Exception:
+ pass
+ # Tune mid-size buckets for large-M rows.
+ N = self.N
+ # DSv3 MLA (padded/strided) hot shape. Prefer a threads-per-row that
+ # makes the tile width exactly match N with 128b vectors (bf16/fp16),
+ # avoiding the ~33% padded work from rounding 1536 -> 2048.
+ if N == 1536 and self.dtype.width == 16:
+ return 96
+ # DSv3 default hidden size (7168). Choose a threads-per-row that matches
+ # the selected vector width to avoid padded work:
+ # - 128b copies (vec=8 for bf16/fp16): 7168/8 = 896 = 7 * 128 -> tpr=128
+ # - 256b copies (vec=16 for bf16/fp16): 7168/16 = 448 = 2 * 224 -> tpr=224
+ #
+ # The fused direct-GMEM path often uses 256b copies on 32B-aligned
+ # tensors, while the non-fused path defaults to 128b copies.
+ if N == 7168 and self.dtype.width == 16:
+ return 224 if self.copy_bits >= 256 else 128
+ # For small-N, use at least one full warp per row. The kernel
+ # implementation assumes one row per CTA; returning <32 here can
+ # produce multi-row tiles (cols_per_block > 1) which is not supported.
+ if N <= 1024:
+ return 32
+ elif N <= 4096:
+ return 128
+ elif N <= 8192:
+ # Allow an override (used by 2-rows/CTA path for N≈6k/8k)
+ try:
+ return self._tpr_override # type: ignore[attr-defined]
+ except Exception:
+ return 128
+ elif N <= 16384:
+ return 256
+ else:
+ return 256
+
+ def _cluster_n(self) -> int:
+ try:
+ return self._cluster_n_override # type: ignore[attr-defined]
+ except Exception:
+ pass
+ N = self.N
+ # Default policy
+ if N <= 8192:
+ return 1
+ if const_expr(self.dtype.width == 16):
+ if N <= 16 * 1024:
+ return 2
+ elif N <= 32 * 1024:
+ return 2
+ elif N <= 64 * 1024:
+ return 4
+ elif N <= 128 * 1024:
+ return 8
+ else:
+ return 16
+ else:
+ if N <= 32 * 1024:
+ return 1
+ elif N <= 64 * 1024:
+ return 2
+ elif N <= 128 * 1024:
+ return 4
+ elif N <= 256 * 1024:
+ return 8
+ else:
+ return 16
+
+ def _num_threads(self) -> int:
+ # Favor 128 threads up to N=16k to reduce per-row partitioning overhead.
+ # This keeps cols_per_block=1 at N=8192 (bf16), which benchmarks faster for large-M.
+ try:
+ return self._nt_override # type: ignore[attr-defined]
+ except Exception:
+ if self.N == 1536 and self.dtype.width == 16:
+ return 96
+ if self.N == 7168 and self.dtype.width == 16:
+ return 224 if self.copy_bits >= 256 else 128
+ if self.N <= 1024:
+ return 32
+ return 128 if self.N <= 16384 else 256
+
+ def _tv_layout(self, num_copy_bits: int = 256) -> Tuple[cute.Shape, cute.Layout]:
+ vecsize = num_copy_bits // self.dtype.width
+ num_threads = self._num_threads()
+ assert num_threads % cute.arch.WARP_SIZE == 0
+ tpr = self._threads_per_row()
+ cluster_n = self._cluster_n()
+ # Allow tails: compute number of vector columns with ceil
+ num_cols_vec = cute.ceil_div(self.N, vecsize)
+ num_blocks_N = cute.ceil_div(num_cols_vec, tpr * cluster_n)
+ cols_per_block = num_threads // tpr
+ tiler_mn = (cols_per_block, vecsize * num_blocks_N * tpr)
+ tv_layout = cute.make_layout(
+ ((tpr, cols_per_block), (vecsize, num_blocks_N)),
+ stride=((vecsize * cols_per_block, 1), (cols_per_block, cols_per_block * vecsize * tpr)),
+ )
+ return tiler_mn, tv_layout
+
+ def _smem_bytes(self, tiler_mn, num_warps) -> int:
+ # smem for X tile (+ residual if present) + reduction buffers + mbar(s)
+ return (
+ cute.size_in_bytes(self.dtype, cute.make_layout(tiler_mn))
+ + self.stage * num_warps * self._cluster_n() * (self.reduction_dtype.width // 8)
+ + self.stage * (cutlass.Int64.width // 8)
+ )
+
+ @cute.jit
+ def __call__(
+ self,
+ mX: cute.Tensor,
+ mW: Optional[cute.Tensor],
+ mB: Optional[cute.Tensor],
+ mRes: Optional[cute.Tensor],
+ mO: cute.Tensor,
+ mResO: Optional[cute.Tensor],
+ mRstd: Optional[cute.Tensor],
+ stream: cuda.CUstream,
+ eps: Float32 = 1e-6,
+ ):
+ # Make last dim static (N)
+ semistatic_shape = (*mX.shape[:-1], self.N)
+
+ def new_stride(t):
+ return (
+ cute.assume(t.stride[0], divby=256 // t.element_type.width),
+ t.stride[1],
+ )
+
+ mX, mRes, mO, mResO = [
+ cute.make_tensor(t.iterator, cute.make_layout(semistatic_shape, stride=new_stride(t)))
+ if const_expr(t is not None)
+ else None
+ for t in (mX, mRes, mO, mResO)
+ ]
+ assert mX.element_type == self.dtype
+ assert mO.element_type == self.dtype
+
+ copy_bits = int(self.copy_bits)
+ tiler_mn, tv_layout = self._tv_layout(num_copy_bits=copy_bits)
+ num_threads = cute.size(tv_layout, mode=[0]) if _KERNEL_ACCEPTS_LAYOUT_ARGS else self._num_threads()
+ num_warps = num_threads // cute.arch.WARP_SIZE
+ threads_per_row = tv_layout.shape[0][0] if _KERNEL_ACCEPTS_LAYOUT_ARGS else self._threads_per_row()
+ warps_per_row = max(threads_per_row // cute.arch.WARP_SIZE, 1)
+ cluster_n = self._cluster_n()
+
+ if const_expr(mW is not None):
+ mW = cute.make_tensor(
+ mW.iterator, cute.prepend(mW.layout, cute.make_layout((tiler_mn[0],), stride=(0,)))
+ )
+ if const_expr(mB is not None):
+ mB = cute.make_tensor(
+ mB.iterator, cute.prepend(mB.layout, cute.make_layout((tiler_mn[0],), stride=(0,)))
+ )
+ if const_expr(mRstd is not None):
+ mRstd = cute.make_tensor(
+ mRstd.iterator, cute.append(mRstd.layout, cute.make_layout((self.N,), stride=(0,)))
+ )
+
+ # No SMEM reload mode switch; overlap is controlled in the K-loop path
+
+ # Compute smem usage considering staged buffers.
+ #
+ # In direct-gmem mode, we skip the gmem->smem tiles entirely and only
+ # keep the reduction buffers in shared memory.
+ stage_bufs = 2 if self.stage > 1 else 1
+ tile_bytes_x = (
+ cute.size_in_bytes(self.dtype, cute.make_layout(tiler_mn)) * stage_bufs
+ if const_expr(not self.direct_gmem)
+ else 0
+ )
+ tile_bytes_res = (
+ cute.size_in_bytes(mRes.element_type, cute.make_layout(tiler_mn)) * stage_bufs
+ if const_expr(mRes is not None and not self.direct_gmem)
+ else 0
+ )
+ red_bytes = self.stage * num_warps * cluster_n * (self.reduction_dtype.width // 8)
+ # mbarriers are only allocated/used for cluster_n>1. Some CuTeDSL builds
+ # require mbarrier state to be 16B-aligned in shared memory; account for
+ # the alignment padding when computing dynamic smem bytes.
+ smem_bytes = tile_bytes_x + tile_bytes_res + red_bytes
+ if cluster_n > 1:
+ # Align up to 16B before placing the mbarrier array.
+ smem_bytes = ((smem_bytes + 15) // 16) * 16
+ smem_bytes += self.stage * (cutlass.Int64.width // 8)
+
+ kernel = (
+ self.kernel(
+ mX,
+ mW,
+ mB,
+ mRes,
+ mO,
+ mResO,
+ mRstd,
+ eps,
+ tv_layout,
+ tiler_mn,
+ const_expr(cluster_n),
+ const_expr(num_warps),
+ const_expr(warps_per_row),
+ const_expr(threads_per_row),
+ )
+ if _KERNEL_ACCEPTS_LAYOUT_ARGS
+ else self.kernel(
+ mX,
+ mW,
+ mB,
+ mRes,
+ mO,
+ mResO,
+ mRstd,
+ eps,
+ )
+ )
+ kernel.launch(
+ grid=[cute.ceil_div(mX.shape[0], tiler_mn[0]), cluster_n, 1],
+ block=[num_threads, 1, 1],
+ cluster=([1, cluster_n, 1] if cluster_n > 1 else None),
+ smem=smem_bytes,
+ stream=stream,
+ )
+
+ @cute.jit
+ def launch_from_ptrs(
+ self,
+ ptr_x: cute.Pointer,
+ ptr_w: Optional[cute.Pointer],
+ ptr_b: Optional[cute.Pointer],
+ ptr_res: Optional[cute.Pointer],
+ ptr_out: cute.Pointer,
+ ptr_res_out: Optional[cute.Pointer],
+ ptr_rstd: Optional[cute.Pointer],
+ M: Int32,
+ N_dyn: Int32,
+ ld: Int32,
+ stream: cuda.CUstream,
+ eps: Float32 = 1e-6,
+ ):
+ """Pointer-based entrypoint to reuse the existing RMSNorm schedule.
+
+ This reconstructs cute.Tensor views from raw pointers plus sizes,
+ avoiding any DLPack conversions at the Python boundary.
+ """
+ # Use a dynamic N for the leading-dimension stride so that the
+ # subsequent cute.assume(...) in __call__ sees a dynamic expression
+ # rather than a plain Python int.
+ # The compile-time N for the kernel (self.N) is still used to
+ # specialize the schedule.
+ # Assume row-major [M, N] with an arbitrary leading-dimension stride
+ # (common for padded-row / packed-attention layouts).
+ layout_mn = cute.make_layout((M, N_dyn), stride=(ld, 1))
+ layout_n = cute.make_layout((N_dyn,), stride=(1,))
+ layout_m = cute.make_layout((M,), stride=(1,))
+
+ mX = cute.make_tensor(ptr_x, layout_mn)
+ mO = cute.make_tensor(ptr_out, layout_mn)
+
+ mRes = (
+ cute.make_tensor(ptr_res, layout_mn)
+ if const_expr(ptr_res is not None)
+ else None
+ )
+ mResO = (
+ cute.make_tensor(ptr_res_out, layout_mn)
+ if const_expr(ptr_res_out is not None)
+ else None
+ )
+ mW = (
+ cute.make_tensor(ptr_w, layout_n)
+ if const_expr(ptr_w is not None)
+ else None
+ )
+ mB = (
+ cute.make_tensor(ptr_b, layout_n)
+ if const_expr(ptr_b is not None)
+ else None
+ )
+ mRstd = (
+ cute.make_tensor(ptr_rstd, layout_m)
+ if const_expr(ptr_rstd is not None)
+ else None
+ )
+
+ # Reuse the main JIT entry to launch the scheduled kernel.
+ self.__call__(mX, mW, mB, mRes, mO, mResO, mRstd, stream, eps)
+
+ @cute.jit
+ def launch_from_ptrs_fused_add_inplace(
+ self,
+ ptr_x: cute.Pointer,
+ ptr_w: cute.Pointer,
+ ptr_res: cute.Pointer,
+ M: Int32,
+ N_dyn: Int32,
+ ld_x: Int32,
+ stream: cuda.CUstream,
+ eps: Float32 = 1e-6,
+ ):
+ """Pointer-based entrypoint for vLLM-style fused_add_rms_norm semantics.
+
+ This specialized entrypoint supports:
+ - `x` / output with an arbitrary leading-dimension stride (`ld_x`), and
+ - `residual` / residual-out as a contiguous [M, N] tensor (ld_res = N).
+
+ Both `x` and `residual` are updated in-place:
+ residual <- x + residual
+ x <- RMSNorm(residual) * weight
+ """
+ layout_x = cute.make_layout((M, N_dyn), stride=(ld_x, 1))
+ layout_res = cute.make_layout((M, N_dyn), stride=(N_dyn, 1))
+ layout_n = cute.make_layout((N_dyn,), stride=(1,))
+
+ mX = cute.make_tensor(ptr_x, layout_x)
+ mO = cute.make_tensor(ptr_x, layout_x)
+ mRes = cute.make_tensor(ptr_res, layout_res)
+ mResO = cute.make_tensor(ptr_res, layout_res)
+ mW = cute.make_tensor(ptr_w, layout_n)
+
+ self.__call__(
+ mX,
+ mW,
+ None, # bias
+ mRes,
+ mO,
+ mResO,
+ None, # rstd
+ stream,
+ eps,
+ )
+
+ @cute.jit
+ def _kernel_impl(
+ self,
+ mX: cute.Tensor,
+ mW: Optional[cute.Tensor],
+ mB: Optional[cute.Tensor],
+ mRes: Optional[cute.Tensor],
+ mO: cute.Tensor,
+ mResO: Optional[cute.Tensor],
+ mRstd: Optional[cute.Tensor],
+ eps: Float32,
+ tv_layout: cute.Layout,
+ tiler_mn: cute.Shape,
+ cluster_n: cutlass.Constexpr[int],
+ num_warps: cutlass.Constexpr[int],
+ warps_per_row: cutlass.Constexpr[int],
+ threads_per_row: cutlass.Constexpr[int],
+ ):
+ tidx, _, _ = cute.arch.thread_idx()
+ bidx, _, _ = cute.arch.block_idx()
+ if const_expr(cluster_n > 1):
+ cta_rank_in_cluster = cute.arch.block_idx_in_cluster()
+ else:
+ cta_rank_in_cluster = const_expr(0)
+ n_off = cta_rank_in_cluster * tiler_mn[1]
+
+ smem = cutlass.utils.SmemAllocator()
+ # Allocate one or two SMEM buffers depending on stage depth
+ sX0 = (
+ smem.allocate_tensor(
+ mX.element_type, cute.make_ordered_layout(tiler_mn, order=(1, 0)), byte_alignment=32
+ )
+ if const_expr(not self.direct_gmem)
+ else None
+ )
+ sX1 = (
+ smem.allocate_tensor(
+ mX.element_type,
+ cute.make_ordered_layout(tiler_mn, order=(1, 0)),
+ byte_alignment=32,
+ )
+ if const_expr(self.stage > 1 and not self.direct_gmem)
+ else None
+ )
+ sRes0 = (
+ smem.allocate_tensor(
+ mRes.element_type, cute.make_ordered_layout(tiler_mn, order=(1, 0)), byte_alignment=32
+ )
+ if const_expr(mRes is not None and not self.direct_gmem)
+ else None
+ )
+ sRes1 = (
+ smem.allocate_tensor(
+ mRes.element_type,
+ cute.make_ordered_layout(tiler_mn, order=(1, 0)),
+ byte_alignment=32,
+ )
+ if const_expr(mRes is not None and self.stage > 1 and not self.direct_gmem)
+ else None
+ )
+
+ # Reduction buffers + mbar for cluster reduce (reused by row_reduce helper)
+ red_layout = cute.make_ordered_layout(
+ (num_warps // warps_per_row, (warps_per_row, cluster_n), self.stage),
+ order=(1, 0, 2),
+ )
+ reduction_buffer = smem.allocate_tensor(self.reduction_dtype, red_layout, byte_alignment=4)
+ if const_expr(cluster_n > 1):
+ # Some CuTeDSL builds appear sensitive to the shared-memory alignment of
+ # mbarrier state. `SmemAllocator.allocate_array` does not currently
+ # expose an alignment parameter, so allocate an Int64 tensor with an
+ # explicit alignment and pass its iterator as the pointer.
+ mbar_tensor = smem.allocate_tensor(
+ cutlass.Int64,
+ cute.make_layout((self.stage,), stride=(1,)),
+ byte_alignment=16,
+ )
+ mbar_ptr = mbar_tensor.iterator
+ else:
+ mbar_ptr = None
+
+ shape = mX.shape
+ idX = cute.make_identity_tensor(shape)
+ limit_k = shape[1] - n_off
+
+ # Tiled copy setup
+ num_copy_elems_X = tv_layout.shape[1][0]
+ use_async = const_expr(self.use_async and self.N >= 1024 and not self.direct_gmem)
+ copy_atom = get_copy_atom_bw(mX.element_type, num_copy_elems_X, is_async=use_async)
+ thr_copy = cute.make_tiled_copy(copy_atom, tv_layout, tiler_mn).get_slice(tidx)
+
+ # Tail predicate for the N dimension (when tile width > N). Reuse this
+ # for W/B loads so we never read past the end of those 1D tensors.
+ is_even_N_wb = const_expr(shape[1] == tiler_mn[1] * cluster_n)
+ if const_expr(not is_even_N_wb):
+ cX0 = cute.local_tile(idX, tiler_mn, (0, 0))
+ tXp_wb = qutils.predicate_k(thr_copy.partition_S(cX0), limit=limit_k)
+ else:
+ tXp_wb = None
+
+ # Weight/bias loads:
+ #
+ # - Direct-GMEM schedule: load weight/bias up front to hide latency.
+ # - Staged SMEM schedule: loading after the reduction reduces register
+ # pressure during the long-scoreboard reduction phase (better for large-M),
+ # but it measurably hurts small-M latency for the non-fused (no residual,
+ # no bias) case. For that specific case, prefetch weight up front as well.
+ tXrW = None
+ tXrB = None
+ prefetch_w_early = bool(
+ mW is not None and (self.direct_gmem or (mRes is None and mB is None))
+ )
+ if const_expr(prefetch_w_early):
+ gW = cute.local_tile(qutils.domain_offset_i64((0, n_off), mW), tiler_mn, (0, 0))
+ tXgW = thr_copy.partition_S(gW)
+ tXrW = cute.make_fragment_like(tXgW)
+ if const_expr(not is_even_N_wb):
+ tXrW.fill(0)
+ cute.copy(
+ get_copy_atom_bw(mW.element_type, num_copy_elems_X, is_async=False),
+ tXgW,
+ tXrW,
+ pred=tXp_wb,
+ )
+ if const_expr(self.direct_gmem and mB is not None):
+ gB = cute.local_tile(qutils.domain_offset_i64((0, n_off), mB), tiler_mn, (0, 0))
+ tXgB = thr_copy.partition_S(gB)
+ tXrB = cute.make_fragment_like(tXgB)
+ if const_expr(not is_even_N_wb):
+ tXrB.fill(0)
+ cute.copy(
+ get_copy_atom_bw(mB.element_type, num_copy_elems_X, is_async=False),
+ tXgB,
+ tXrB,
+ pred=tXp_wb,
+ )
+
+ # Non-persistent per-CTA execution (one tile in M)
+ self._init_cluster(tidx, mbar_ptr)
+
+ mX_i, mRes_i, mO_i, mResO_i = [
+ qutils.domain_offset_i64((bidx * tiler_mn[0], 0), t) if t is not None else None
+ for t in (mX, mRes, mO, mResO)
+ ]
+ mX_i, mRes_i, mO_i, mResO_i = [
+ qutils.domain_offset_i64((0, n_off), t) if t is not None else None
+ for t in (mX_i, mRes_i, mO_i, mResO_i)
+ ]
+ gX_i = cute.local_tile(mX_i, tiler_mn, (0, 0))
+ gO_i = cute.local_tile(mO_i, tiler_mn, (0, 0))
+ gRes_i = (
+ cute.local_tile(mRes_i, tiler_mn, (0, 0)) if const_expr(mRes is not None) else None
+ )
+ gResO_i = (
+ cute.local_tile(mResO_i, tiler_mn, (0, 0)) if const_expr(mResO is not None) else None
+ )
+ gRstd_i = (
+ cute.local_tile(mRstd, tiler_mn, (bidx, 0)) if const_expr(mRstd is not None) else None
+ )
+ cX_i = cute.local_tile(idX, tiler_mn, (bidx, 0))
+
+ # Common identity/row index partitions reused by both default and K-loop paths
+ tXcX_i = thr_copy.partition_S(cX_i)[(0, None), None, None]
+ row_i = tXcX_i[0][0]
+ tXgRstd_i = thr_copy.partition_D(gRstd_i) if const_expr(mRstd is not None) else None
+
+ # Stage-2 intra-row K-loop cp.async ping-pong (two tiles). This reduces
+ # per-thread fragment size and can improve memory-latency hiding for
+ # N=7168 at large M. It is enabled by setting `stage=2` when constructing
+ # the RMSNormSM100 op (see `_fused_add_rmsnorm_forward_ptr_inplace`).
+ if const_expr(
+ self.stage > 1 and not self.direct_gmem and use_async and cluster_n == 1 and shape[1] == 7168
+ ):
+ vecsize = tv_layout.shape[1][0]
+ tpr = threads_per_row
+ target_tile_n = const_expr(4096)
+ tile_factor = const_expr(target_tile_n // (vecsize * tpr))
+ if const_expr(tile_factor > 0):
+ tile_n = vecsize * tpr * tile_factor
+ num_tiles = cute.ceil_div(shape[1], tile_n)
+
+ tiler_mn_tile = (tiler_mn[0], tile_n)
+ sX0_tile = cute.local_tile(sX0, tiler_mn_tile, (0, 0))
+ sX1_tile = cute.local_tile(sX1, tiler_mn_tile, (0, 0))
+ sRes0_tile = (
+ cute.local_tile(sRes0, tiler_mn_tile, (0, 0))
+ if const_expr(mRes is not None)
+ else None
+ )
+ sRes1_tile = (
+ cute.local_tile(sRes1, tiler_mn_tile, (0, 0))
+ if const_expr(mRes is not None)
+ else None
+ )
+
+ tv_layout_tile = cute.make_layout(
+ ((tpr, tiler_mn[0]), (vecsize, tile_factor)),
+ stride=(
+ (vecsize * tiler_mn[0], 1),
+ (tiler_mn[0], tiler_mn[0] * vecsize * tpr),
+ ),
+ )
+ thr_copy_tile = cute.make_tiled_copy(copy_atom, tv_layout_tile, tiler_mn_tile).get_slice(
+ tidx
+ )
+
+ # Accumulate per-thread partial sums across tiles; reduce once.
+ sum_sq_thread = cute.Float32(0.0)
+
+ # Preload tile 0 into sX0/sRes0.
+ k_off0 = const_expr(0) * tile_n
+ gX_0 = cute.local_tile(
+ qutils.domain_offset_i64((0, k_off0), mX_i), tiler_mn_tile, (0, 0)
+ )
+ tXgX_0 = thr_copy_tile.partition_S(gX_0)
+ tXsX_0 = thr_copy_tile.partition_D(sX0_tile)
+ cX_0 = cute.local_tile(
+ cute.domain_offset((0, k_off0), cX_i), tiler_mn_tile, (0, 0)
+ )
+ tXc_0 = thr_copy_tile.partition_S(cX_0)
+ tXp_0 = qutils.predicate_k(tXc_0, limit=limit_k)
+
+ tXp_ping = tXp_0
+ tXp_pong = tXp_0
+
+ if row_i < shape[0]:
+ copy_tiled(tXgX_0, tXsX_0, num_copy_elems=vecsize, is_async=True, pred=tXp_0)
+ if const_expr(mRes is not None):
+ gRes_0 = cute.local_tile(
+ qutils.domain_offset_i64((0, k_off0), mRes_i),
+ tiler_mn_tile,
+ (0, 0),
+ )
+ tXgRes_0 = thr_copy_tile.partition_S(gRes_0)
+ tXsRes_0 = thr_copy_tile.partition_D(sRes0_tile)
+ copy_tiled(
+ tXgRes_0,
+ tXsRes_0,
+ num_copy_elems=vecsize,
+ is_async=True,
+ pred=tXp_0,
+ )
+ cute.arch.cp_async_commit_group()
+
+ for t in cutlass.range_constexpr(num_tiles):
+ next_t = t + 1
+ if next_t < num_tiles:
+ k_off_n = next_t * tile_n
+ gX_n = cute.local_tile(
+ qutils.domain_offset_i64((0, k_off_n), mX_i),
+ tiler_mn_tile,
+ (0, 0),
+ )
+ tXgX_n = thr_copy_tile.partition_S(gX_n)
+ cX_n = cute.local_tile(
+ cute.domain_offset((0, k_off_n), cX_i),
+ tiler_mn_tile,
+ (0, 0),
+ )
+ tXc_n = thr_copy_tile.partition_S(cX_n)
+ tXp_n = qutils.predicate_k(tXc_n, limit=limit_k)
+
+ if const_expr((t % 2) == 0):
+ tXsX_n = thr_copy_tile.partition_D(sX1_tile)
+ tXsRes_n = (
+ thr_copy_tile.partition_D(sRes1_tile) if const_expr(mRes is not None) else None
+ )
+ tXp_pong = tXp_n
+ else:
+ tXsX_n = thr_copy_tile.partition_D(sX0_tile)
+ tXsRes_n = (
+ thr_copy_tile.partition_D(sRes0_tile) if const_expr(mRes is not None) else None
+ )
+ tXp_ping = tXp_n
+
+ if row_i < shape[0]:
+ copy_tiled(
+ tXgX_n, tXsX_n, num_copy_elems=vecsize, is_async=True, pred=tXp_n
+ )
+ if const_expr(mRes is not None):
+ gRes_n = cute.local_tile(
+ qutils.domain_offset_i64((0, k_off_n), mRes_i),
+ tiler_mn_tile,
+ (0, 0),
+ )
+ tXgRes_n = thr_copy_tile.partition_S(gRes_n)
+ copy_tiled(
+ tXgRes_n,
+ tXsRes_n,
+ num_copy_elems=vecsize,
+ is_async=True,
+ pred=tXp_n,
+ )
+ cute.arch.cp_async_commit_group()
+
+ cute.arch.cp_async_wait_group(1 if next_t < num_tiles else 0)
+
+ # Current tile buffer (ping/pong).
+ if const_expr((t % 2) == 0):
+ tXsX_cur = thr_copy_tile.partition_D(sX0_tile)
+ tXsRes_cur = (
+ thr_copy_tile.partition_D(sRes0_tile) if const_expr(mRes is not None) else None
+ )
+ pred_cur = tXp_ping
+ else:
+ tXsX_cur = thr_copy_tile.partition_D(sX1_tile)
+ tXsRes_cur = (
+ thr_copy_tile.partition_D(sRes1_tile) if const_expr(mRes is not None) else None
+ )
+ pred_cur = tXp_pong
+
+ k_off = t * tile_n
+ gX_t = cute.local_tile(qutils.domain_offset_i64((0, k_off), mX_i), tiler_mn_tile, (0, 0))
+ tXgX_t = thr_copy_tile.partition_S(gX_t)
+ tXrX_t = cute.make_fragment_like(tXgX_t)
+ cute.autovec_copy(tXsX_cur, tXrX_t)
+ x_t = tXrX_t.load().to(cute.Float32)
+ if const_expr(mRes is not None):
+ gRes_t = cute.local_tile(
+ qutils.domain_offset_i64((0, k_off), mRes_i), tiler_mn_tile, (0, 0)
+ )
+ tXgRes_t = thr_copy_tile.partition_S(gRes_t)
+ tXrRes_t = cute.make_fragment_like(tXgRes_t)
+ cute.autovec_copy(tXsRes_cur, tXrRes_t)
+ x_t += tXrRes_t.load().to(cute.Float32)
+
+ if const_expr(mResO is not None):
+ gResO_t = cute.local_tile(
+ qutils.domain_offset_i64((0, k_off), mResO_i),
+ tiler_mn_tile,
+ (0, 0),
+ )
+ tXgResO_t = thr_copy_tile.partition_D(gResO_t)
+ tXrResO_t = cute.make_fragment_like(tXgResO_t)
+ tXrResO_t.store(x_t.to(tXrResO_t.element_type))
+ if row_i < shape[0]:
+ copy_tiled(
+ tXrResO_t,
+ tXgResO_t,
+ num_copy_elems=vecsize,
+ is_async=False,
+ pred=pred_cur,
+ )
+
+ sum_sq_thread = sum_sq_thread + (x_t * x_t).reduce(
+ cute.ReductionOp.ADD,
+ init_val=0.0,
+ reduction_profile=0,
+ )
+
+ sum_sq = row_reduce(
+ sum_sq_thread,
+ cute.ReductionOp.ADD,
+ threads_per_row,
+ reduction_buffer[None, None, 0],
+ mbar_ptr,
+ init_val=0.0,
+ )
+ rstd = cute.math.rsqrt(sum_sq / shape[1] + eps, fastmath=True)
+
+ if const_expr(mRstd is not None):
+ if tXcX_i[0][1] == 0 and row_i < shape[0]:
+ tXgRstd_i[0] = rstd
+
+ for t in cutlass.range_constexpr(num_tiles):
+ k_off = t * tile_n
+ cX_t = cute.local_tile(cute.domain_offset((0, k_off), cX_i), tiler_mn_tile, (0, 0))
+ tXc_t = thr_copy_tile.partition_S(cX_t)
+ tXp_t = qutils.predicate_k(tXc_t, limit=limit_k)
+
+ if const_expr((t % 2) == 0):
+ tXsX_cur = thr_copy_tile.partition_D(sX0_tile)
+ tXsRes_cur = (
+ thr_copy_tile.partition_D(sRes0_tile) if const_expr(mRes is not None) else None
+ )
+ else:
+ tXsX_cur = thr_copy_tile.partition_D(sX1_tile)
+ tXsRes_cur = (
+ thr_copy_tile.partition_D(sRes1_tile) if const_expr(mRes is not None) else None
+ )
+
+ gX_t = cute.local_tile(qutils.domain_offset_i64((0, k_off), mX_i), tiler_mn_tile, (0, 0))
+ tXgX_t = thr_copy_tile.partition_S(gX_t)
+ tXrX_t = cute.make_fragment_like(tXgX_t)
+ cute.autovec_copy(tXsX_cur, tXrX_t)
+ x_t = tXrX_t.load().to(cute.Float32)
+ if const_expr(mRes is not None):
+ gRes_t = cute.local_tile(
+ qutils.domain_offset_i64((0, k_off), mRes_i), tiler_mn_tile, (0, 0)
+ )
+ tXgRes_t = thr_copy_tile.partition_S(gRes_t)
+ tXrRes_t = cute.make_fragment_like(tXgRes_t)
+ cute.autovec_copy(tXsRes_cur, tXrRes_t)
+ x_t += tXrRes_t.load().to(cute.Float32)
+
+ y_t = x_t * rstd
+ if const_expr(mW is not None):
+ gW_t = cute.local_tile(
+ qutils.domain_offset_i64((0, k_off), mW), tiler_mn_tile, (0, 0)
+ )
+ tWgW_t = thr_copy_tile.partition_S(gW_t)
+ tWrW_t = cute.make_fragment_like(tWgW_t)
+ copy_tiled(tWgW_t, tWrW_t, num_copy_elems=vecsize, is_async=False, pred=tXp_t)
+ y_t = y_t * tWrW_t.load().to(cute.Float32)
+ if const_expr(mB is not None):
+ gB_t = cute.local_tile(
+ qutils.domain_offset_i64((0, k_off), mB), tiler_mn_tile, (0, 0)
+ )
+ tWgB_t = thr_copy_tile.partition_S(gB_t)
+ tWrB_t = cute.make_fragment_like(tWgB_t)
+ copy_tiled(tWgB_t, tWrB_t, num_copy_elems=vecsize, is_async=False, pred=tXp_t)
+ y_t = y_t + tWrB_t.load().to(cute.Float32)
+
+ gO_t = cute.local_tile(qutils.domain_offset_i64((0, k_off), mO_i), tiler_mn_tile, (0, 0))
+ tXgO_t = thr_copy_tile.partition_D(gO_t)
+ tXrO_t = cute.make_fragment_like(tXgO_t)
+ tXrO_t.store(y_t.to(tXrO_t.element_type))
+ if row_i < shape[0]:
+ copy_tiled(tXrO_t, tXgO_t, num_copy_elems=vecsize, is_async=False, pred=tXp_t)
+
+ return
+
+ # Single-stage path: one-row-per-CTA
+ tXgX_i = thr_copy.partition_S(gX_i)
+ tXgRes_i = thr_copy.partition_S(gRes_i) if const_expr(mRes is not None) else None
+ tXgO_i = thr_copy.partition_D(gO_i)
+ tXgResO_i = thr_copy.partition_D(gResO_i) if const_expr(mResO is not None) else None
+ # tXgRstd_i / tXcX_i / row_i prepared above
+ is_even_N_i = const_expr(shape[1] == tiler_mn[1] * cluster_n)
+ tXpX_i = (
+ qutils.predicate_k(thr_copy.partition_S(cX_i), limit=limit_k) if not is_even_N_i else None
+ )
+
+ tXrX = cute.make_fragment_like(tXgX_i)
+ tXrRes = cute.make_fragment_like(tXgRes_i) if const_expr(mRes is not None) else None
+ if const_expr(self.direct_gmem):
+ if const_expr(not is_even_N_i):
+ tXrX.fill(0)
+ if const_expr(tXrRes is not None):
+ tXrRes.fill(0)
+ if row_i < shape[0]:
+ cute.copy(copy_atom, tXgX_i, tXrX, pred=tXpX_i)
+ if const_expr(tXrRes is not None):
+ cute.copy(copy_atom, tXgRes_i, tXrRes, pred=tXpX_i)
+ else:
+ # If N is not a multiple of the tile width, the predicated gmem->smem
+ # copies leave out-of-bounds lanes uninitialized. Clear the SMEM tile
+ # so masked lanes read as 0 for reduction/output.
+ if const_expr(not is_even_N_i):
+ thr_copy.partition_D(sX0).fill(0)
+ if const_expr(mRes is not None):
+ thr_copy.partition_D(sRes0).fill(0)
+
+ if row_i < shape[0]:
+ cute.copy(copy_atom, tXgX_i, thr_copy.partition_D(sX0), pred=tXpX_i)
+ if const_expr(mRes is not None):
+ cute.copy(copy_atom, tXgRes_i, thr_copy.partition_D(sRes0), pred=tXpX_i)
+ if const_expr(use_async):
+ cute.arch.cp_async_commit_group()
+ cute.arch.cp_async_wait_group(0)
+
+ cute.autovec_copy(thr_copy.partition_D(sX0), tXrX)
+ if const_expr(tXrRes is not None):
+ cute.autovec_copy(thr_copy.partition_D(sRes0), tXrRes)
+ x_red = tXrX.load().to(cute.Float32)
+ if const_expr(tXrRes is not None):
+ x_red += tXrRes.load().to(cute.Float32)
+
+ if const_expr(mResO is not None):
+ tXrResO = cute.make_fragment_like(tXgResO_i)
+ tXrResO.store(x_red.to(tXrResO.element_type))
+ if row_i < shape[0]:
+ cute.copy(
+ get_copy_atom_bw(tXrResO.element_type, num_copy_elems_X, is_async=False),
+ tXrResO,
+ tXgResO_i,
+ pred=tXpX_i,
+ )
+
+ sum_sq = row_reduce(
+ x_red * x_red,
+ cute.ReductionOp.ADD,
+ threads_per_row,
+ reduction_buffer[None, None, 0],
+ mbar_ptr,
+ init_val=0.0,
+ )
+ rstd = cute.math.rsqrt(sum_sq / shape[1] + eps, fastmath=True)
+
+ if const_expr(mRstd is not None):
+ if (
+ tXcX_i[0][1] == 0
+ and row_i < shape[0]
+ and (cluster_n == 1 or cute.arch.block_idx_in_cluster() == 0)
+ ):
+ tXgRstd_i[0] = rstd
+
+ if const_expr(not self.direct_gmem and (mRes is not None or mB is not None)):
+ # Load weight/bias after the reduction so they don't inflate register
+ # pressure during the long-scoreboard reduction phase (helping occupancy
+ # when registers are the limiting factor).
+ if const_expr(mW is not None):
+ gW = cute.local_tile(qutils.domain_offset_i64((0, n_off), mW), tiler_mn, (0, 0))
+ tXgW = thr_copy.partition_S(gW)
+ tXrW = cute.make_fragment_like(tXgW)
+ if const_expr(not is_even_N_wb):
+ tXrW.fill(0)
+ cute.copy(
+ get_copy_atom_bw(mW.element_type, num_copy_elems_X, is_async=False),
+ tXgW,
+ tXrW,
+ pred=tXp_wb,
+ )
+ if const_expr(mB is not None):
+ gB = cute.local_tile(qutils.domain_offset_i64((0, n_off), mB), tiler_mn, (0, 0))
+ tXgB = thr_copy.partition_S(gB)
+ tXrB = cute.make_fragment_like(tXgB)
+ if const_expr(not is_even_N_wb):
+ tXrB.fill(0)
+ cute.copy(
+ get_copy_atom_bw(mB.element_type, num_copy_elems_X, is_async=False),
+ tXgB,
+ tXrB,
+ pred=tXp_wb,
+ )
+
+ # Reuse `x_red` (x + residual, in fp32) for the output path so we don't
+ # keep both `tXrX` and `tXrRes` fragments live across the reduction.
+ y = x_red * rstd
+ if const_expr(mW is not None):
+ y = y * tXrW.load().to(cute.Float32)
+ if const_expr(mB is not None):
+ y = y + tXrB.load().to(cute.Float32)
+
+ tXrO = cute.make_fragment_like(tXgO_i)
+ tXrO.store(y.to(tXrO.element_type))
+ if row_i < shape[0]:
+ cute.copy(
+ get_copy_atom_bw(tXrO.element_type, num_copy_elems_X, is_async=False),
+ tXrO,
+ tXgO_i,
+ pred=tXpX_i,
+ )
+
+ if _KERNEL_ACCEPTS_LAYOUT_ARGS:
+ @cute.kernel
+ def kernel(
+ self,
+ mX: cute.Tensor,
+ mW: Optional[cute.Tensor],
+ mB: Optional[cute.Tensor],
+ mRes: Optional[cute.Tensor],
+ mO: cute.Tensor,
+ mResO: Optional[cute.Tensor],
+ mRstd: Optional[cute.Tensor],
+ eps: Float32,
+ tv_layout: cute.Layout,
+ tiler_mn: cute.Shape,
+ cluster_n: cutlass.Constexpr[int],
+ num_warps: cutlass.Constexpr[int],
+ warps_per_row: cutlass.Constexpr[int],
+ threads_per_row: cutlass.Constexpr[int],
+ ):
+ self._kernel_impl(
+ mX,
+ mW,
+ mB,
+ mRes,
+ mO,
+ mResO,
+ mRstd,
+ eps,
+ tv_layout,
+ tiler_mn,
+ cluster_n,
+ num_warps,
+ warps_per_row,
+ threads_per_row,
+ )
+ else:
+ @cute.kernel
+ def kernel(
+ self,
+ mX: cute.Tensor,
+ mW: Optional[cute.Tensor],
+ mB: Optional[cute.Tensor],
+ mRes: Optional[cute.Tensor],
+ mO: cute.Tensor,
+ mResO: Optional[cute.Tensor],
+ mRstd: Optional[cute.Tensor],
+ eps: Float32,
+ ):
+ copy_bits = int(self.copy_bits)
+ tiler_mn, tv_layout = self._tv_layout(num_copy_bits=copy_bits)
+ num_threads = self._num_threads()
+ num_warps = num_threads // cute.arch.WARP_SIZE
+ threads_per_row = self._threads_per_row()
+ warps_per_row = max(threads_per_row // cute.arch.WARP_SIZE, 1)
+ cluster_n = self._cluster_n()
+ self._kernel_impl(
+ mX,
+ mW,
+ mB,
+ mRes,
+ mO,
+ mResO,
+ mRstd,
+ eps,
+ tv_layout,
+ tiler_mn,
+ const_expr(cluster_n),
+ const_expr(num_warps),
+ const_expr(warps_per_row),
+ const_expr(threads_per_row),
+ )
+
+ @cute.jit
+ def _init_cluster(self, tidx: cutlass.Int32, mbar_ptr: Optional[cute.Pointer]):
+ if const_expr(mbar_ptr is not None):
+ if tidx < self.stage:
+ cute.arch.mbarrier_init(mbar_ptr + tidx, 1)
+ cute.arch.mbarrier_init_fence()
+ cute.arch.cluster_arrive_relaxed()
+
+
+def _can_use_ptr_path(
+ x: Tensor,
+ weight: Optional[Tensor],
+ bias: Optional[Tensor],
+ residual: Optional[Tensor],
+) -> bool:
+ """Fast path precondition for the pointer-based CuTeDSL entry.
+
+ We require a row-major 2D layout where the last dimension is
+ contiguous (stride(1) == 1). The leading dimension (stride(0))
+ may be larger than N (padded-row / packed-attention layouts),
+ and is passed to the kernel as `ld`.
+ """
+ if x.stride(1) != 1:
+ return False
+ # All participating tensors are interpreted as the same element type
+ # (derived from x.dtype) in the pointer-based path. If dtypes differ,
+ # we'd read the wrong bit patterns and silently produce incorrect output.
+ if residual is not None and residual.dtype != x.dtype:
+ return False
+ if weight is not None and weight.dtype != x.dtype:
+ return False
+ if bias is not None and bias.dtype != x.dtype:
+ return False
+ # The kernel assumes `ld` satisfies a divisibility constraint used by
+ # cute.assume(..., divby=...) for vectorization.
+ elem_bits = TORCH2CUTE_DTYPE[x.dtype].width
+ divby = 256 // elem_bits
+ if (x.stride(0) % divby) != 0:
+ return False
+ # The kernel uses 128-bit vectorized copies (16B). Require at least 16B
+ # alignment on all participating tensors to avoid misaligned global loads.
+ if (x.data_ptr() % 16) != 0:
+ return False
+ if residual is not None and residual.stride(1) != 1:
+ return False
+ if residual is not None and residual.stride(0) != x.stride(0):
+ return False
+ if residual is not None and (residual.data_ptr() % 16) != 0:
+ return False
+ if weight is not None and not weight.is_contiguous():
+ return False
+ if bias is not None and not bias.is_contiguous():
+ return False
+ if weight is not None and (weight.data_ptr() % 16) != 0:
+ return False
+ if bias is not None and (bias.data_ptr() % 16) != 0:
+ return False
+ return True
+
+
+def _can_use_ptr_path_fused_add_inplace(
+ x: Tensor,
+ weight: Tensor,
+ residual: Tensor,
+) -> bool:
+ """Fast-path precondition for fused_add_rmsnorm_forward_inplace.
+
+ We allow the common vLLM layout where:
+ - `x` is strided/padded row-major (stride(1) == 1, stride(0) >= N)
+ - `residual` is contiguous row-major (stride(0) == N)
+ """
+ if x.stride(1) != 1:
+ return False
+ if residual.dtype != x.dtype:
+ return False
+ if weight.dtype != x.dtype:
+ return False
+ if residual.stride(1) != 1:
+ return False
+ if not residual.is_contiguous():
+ return False
+ if not weight.is_contiguous():
+ return False
+
+ dtype = TORCH2CUTE_DTYPE[x.dtype]
+ divby = 256 // dtype.width
+ if (x.stride(0) % divby) != 0:
+ return False
+ if (residual.stride(0) % divby) != 0:
+ return False
+
+ if (x.data_ptr() % 16) != 0:
+ return False
+ if (residual.data_ptr() % 16) != 0:
+ return False
+ if (weight.data_ptr() % 16) != 0:
+ return False
+ return True
+
+
+def _rmsnorm_forward_ptr(
+ x: Tensor,
+ weight: Optional[Tensor],
+ bias: Optional[Tensor],
+ residual: Optional[Tensor],
+ eps: float,
+ store_rstd: bool,
+) -> Tuple[Tensor, Optional[Tensor], Optional[Tensor]]:
+ """Pointer-based RMSNorm forward that bypasses DLPack entirely.
+
+ This path reconstructs cute.Tensor views from raw device pointers
+ and explicit layouts inside the JIT graph, avoiding any runtime
+ DLPack conversions while reusing the tuned RMSNormSM100 schedule.
+ """
+ assert x.is_cuda
+ assert x.dim() == 2, "Use (M, N) tensor; flatten batch/seq beforehand."
+ M, N = x.shape
+
+ # Preserve the input's 2D stride so downstream users that rely on
+ # padded-row layouts (stride0 > N) continue to see the expected layout.
+ out = torch.empty_strided(x.shape, x.stride(), device=x.device, dtype=x.dtype)
+ residual_out: Optional[Tensor] = None
+ rstd: Optional[Tensor] = None
+
+ if residual is not None:
+ residual_out = torch.empty_strided(
+ residual.shape, residual.stride(), device=residual.device, dtype=residual.dtype
+ )
+ if store_rstd:
+ rstd = torch.empty(M, device=x.device, dtype=torch.float32)
+
+ _rmsnorm_forward_ptr_into(
+ x=x,
+ weight=weight,
+ bias=bias,
+ residual=residual,
+ out=out,
+ residual_out=residual_out,
+ rstd=rstd,
+ eps=eps,
+ )
+ return out, rstd, residual_out
+
+
+def _rmsnorm_forward_ptr_into(
+ x: Tensor,
+ weight: Optional[Tensor],
+ bias: Optional[Tensor],
+ residual: Optional[Tensor],
+ out: Tensor,
+ residual_out: Optional[Tensor],
+ rstd: Optional[Tensor],
+ eps: float,
+) -> None:
+ """Internal helper that launches the pointer-based kernel into preallocated outputs.
+
+ This enables integration into frameworks like vLLM that manage their
+ own buffers and prefer in-place or out-parameter semantics.
+ """
+ assert x.is_cuda
+ assert x.dim() == 2, "Use (M, N) tensor; flatten batch/seq beforehand."
+ M, N = x.shape
+ device_index = x.get_device()
+ dtype = TORCH2CUTE_DTYPE[x.dtype]
+
+ if bias is None and residual is None and residual_out is None and rstd is None:
+ # Fast-launch path: cache packed args and update pointers/scalars in-place to
+ # avoid Python-side argument marshalling overhead that dominates small-batch cases.
+ #
+ # If fast-launch is disabled (or CuTeDSL internals changed), we fall back
+ # to calling the compiled function directly.
+ if torch.cuda.current_device() != device_index:
+ torch.cuda.set_device(device_index)
+ stream_handle = int(torch.cuda.current_stream().cuda_stream)
+ has_weight = weight is not None
+
+ stage = 1
+ compiled_key = (
+ "ptr",
+ N,
+ dtype,
+ False, # residual
+ has_weight,
+ False, # bias
+ False, # residual_out
+ False, # rstd
+ stage,
+ device_index,
+ )
+ compiled = _PTR_COMPILE_CACHE.get(compiled_key)
+ if compiled is None:
+ op = RMSNormSM100(N, dtype, stage=stage)
+ ld_val = int(x.stride(0))
+ ptr_x = rt.make_ptr(dtype, x.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ ptr_out = rt.make_ptr(dtype, out.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ ptr_w = (
+ rt.make_ptr(dtype, weight.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ if has_weight
+ else None
+ )
+ stream = cuda.CUstream(stream_handle)
+ ld = Int32(ld_val)
+ compiled = cute.compile(
+ op.launch_from_ptrs,
+ ptr_x,
+ ptr_w,
+ None, # ptr_b
+ None, # ptr_res
+ ptr_out,
+ None, # ptr_res_out
+ None, # ptr_rstd
+ Int32(M),
+ Int32(N),
+ ld,
+ stream,
+ Float32(eps),
+ )
+ _PTR_COMPILE_CACHE[compiled_key] = compiled
+
+ launcher = _get_fast_ptr_rmsnorm_launcher(
+ compiled=compiled,
+ dtype=dtype,
+ N=N,
+ device_index=device_index,
+ stream_handle=stream_handle,
+ has_weight=has_weight,
+ eps=eps,
+ )
+ ld_val = int(x.stride(0))
+ if launcher is not None:
+ launcher.launch(x=x, weight=weight, out=out, M=M, N=N, ld=ld_val, eps=eps)
+ return
+
+ ptr_x = rt.make_ptr(dtype, x.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ ptr_out = rt.make_ptr(dtype, out.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ ptr_w = (
+ rt.make_ptr(dtype, weight.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ if has_weight
+ else None
+ )
+ stream = cuda.CUstream(stream_handle)
+ ld = Int32(ld_val)
+ compiled(
+ ptr_x,
+ ptr_w,
+ None, # ptr_b
+ None, # ptr_res
+ ptr_out,
+ None, # ptr_res_out
+ None, # ptr_rstd
+ Int32(M),
+ Int32(N),
+ ld,
+ stream,
+ Float32(eps),
+ )
+ return
+
+ # Fallback: general path (supports bias/residual/rstd, but is slower to launch).
+ stage = 1
+ if torch.cuda.current_device() != device_index:
+ torch.cuda.set_device(device_index)
+ stream_handle = int(torch.cuda.current_stream().cuda_stream)
+ key = (
+ "ptr",
+ N,
+ dtype,
+ residual is not None,
+ weight is not None,
+ bias is not None,
+ residual_out is not None,
+ rstd is not None,
+ stage,
+ device_index,
+ )
+ compiled = _PTR_COMPILE_CACHE.get(key)
+ if compiled is None:
+ op = RMSNormSM100(N, dtype, stage=stage)
+ ptr_x = rt.make_ptr(dtype, x.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ ptr_out = rt.make_ptr(dtype, out.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ ptr_res = (
+ rt.make_ptr(dtype, residual.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ if residual is not None
+ else None
+ )
+ ptr_res_out = (
+ rt.make_ptr(
+ dtype, residual_out.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16
+ )
+ if residual_out is not None
+ else None
+ )
+ ptr_w = (
+ rt.make_ptr(dtype, weight.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ if weight is not None
+ else None
+ )
+ ptr_b = (
+ rt.make_ptr(dtype, bias.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ if bias is not None
+ else None
+ )
+ ptr_rstd = (
+ rt.make_ptr(
+ cutlass.Float32, rstd.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=4
+ )
+ if rstd is not None
+ else None
+ )
+ stream = cuda.CUstream(stream_handle)
+ ld = Int32(int(x.stride(0)))
+ compiled = cute.compile(
+ op.launch_from_ptrs,
+ ptr_x,
+ ptr_w,
+ ptr_b,
+ ptr_res,
+ ptr_out,
+ ptr_res_out,
+ ptr_rstd,
+ Int32(M),
+ Int32(N),
+ ld,
+ stream,
+ Float32(eps),
+ )
+ _PTR_COMPILE_CACHE[key] = compiled
+ ptr_x = rt.make_ptr(dtype, x.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ ptr_out = rt.make_ptr(dtype, out.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ ptr_res = (
+ rt.make_ptr(dtype, residual.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ if residual is not None
+ else None
+ )
+ ptr_res_out = (
+ rt.make_ptr(dtype, residual_out.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ if residual_out is not None
+ else None
+ )
+ ptr_w = (
+ rt.make_ptr(dtype, weight.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ if weight is not None
+ else None
+ )
+ ptr_b = (
+ rt.make_ptr(dtype, bias.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ if bias is not None
+ else None
+ )
+ ptr_rstd = (
+ rt.make_ptr(cutlass.Float32, rstd.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=4)
+ if rstd is not None
+ else None
+ )
+ stream = cuda.CUstream(stream_handle)
+ ld = Int32(int(x.stride(0)))
+ compiled(
+ ptr_x,
+ ptr_w,
+ ptr_b,
+ ptr_res,
+ ptr_out,
+ ptr_res_out,
+ ptr_rstd,
+ Int32(M),
+ Int32(N),
+ ld,
+ stream,
+ Float32(eps),
+ )
+
+
+def _fused_add_rmsnorm_forward_ptr_inplace(
+ x: Tensor,
+ residual: Tensor,
+ weight: Tensor,
+ eps: float,
+) -> None:
+ """Pointer-based fused_add_rmsnorm that updates `x` and `residual` in-place."""
+ assert x.is_cuda
+ assert x.dim() == 2
+ assert residual.is_cuda
+ assert residual.dim() == 2
+ assert x.shape == residual.shape
+
+ M, N = x.shape
+ device_index = x.get_device()
+ dtype = TORCH2CUTE_DTYPE[x.dtype]
+ stage = 1
+
+ if torch.cuda.current_device() != device_index:
+ torch.cuda.set_device(device_index)
+ stream_handle = int(torch.cuda.current_stream().cuda_stream)
+
+ # Latency-optimized schedule for small-M cases: avoid the gmem->smem
+ # staging path (large dynamic smem + extra barriers) and load directly
+ # from gmem into registers.
+ copy_bits = 128
+ # Use a direct-GMEM schedule (no staging SMEM tiles) for DSv3 hidden size
+ # (7168, bf16/fp16). This improves both:
+ # - small-M latency (fewer barriers + less dynamic shared memory), and
+ # - large-M bandwidth (lower overhead, better vectorization when 32B-aligned).
+ #
+ # This is a policy decision: it is tuned for DSv3's N=7168. If you want to
+ # benchmark other models/shapes, you can override it with:
+ # - OINK_RMSNORM_DIRECT_GMEM=0 (force staging/cp.async path)
+ # - OINK_RMSNORM_DIRECT_GMEM=1 (force direct-gmem path)
+ direct_gmem = _direct_gmem_from_policy(default=bool(dtype.width == 16 and N == 7168))
+ use_async = not direct_gmem
+ tpr_override: Optional[int] = None
+ nt_override: Optional[int] = None
+ cluster_n_override: Optional[int] = None
+ direct_gmem_max_copy_bits: Optional[int] = None
+
+ # Experimental stage-2 cp.async path (2-tile ping-pong) for N=7168. This is
+ # primarily about improving memory-latency hiding / reducing long-scoreboard
+ # stalls for large-M workloads.
+ if _ENABLE_STAGE2 and dtype.width == 16 and N == 7168 and M >= 4096:
+ stage = 2
+ direct_gmem = False
+ use_async = True
+
+ # Experimental ILP variant (clusters): split each row across 2 CTAs.
+ #
+ # NOTE: This is currently opt-in because some CuTeDSL builds exhibit
+ # instability with cluster launches for this specific schedule. To reduce
+ # the chance of accidental crashes, we require an additional explicit
+ # opt-in via `OINK_RMSNORM_ENABLE_CLUSTER_ILP_UNSAFE=1`.
+ if _ENABLE_CLUSTER_ILP and not _ENABLE_STAGE2:
+ if dtype.width == 16 and N == 7168 and M >= 4096:
+ cluster_n_override = 2
+ if direct_gmem:
+ # Cluster launches + direct-GMEM has exhibited reproducible compiler
+ # instability (segfaults) in some CuTeDSL builds, especially for the
+ # 256b vector path. Probe it out-of-process once so we can safely
+ # select a working copy width (or fall back to the staged SMEM path)
+ # instead of crashing the parent process.
+ max_bits = _probe_cluster_direct_gmem_max_copy_bits()
+ if max_bits == 0:
+ direct_gmem = False
+ use_async = True
+ else:
+ direct_gmem_max_copy_bits = max_bits
+
+ # Experimental per-row partitioning: use 256 threads/row for N=7168 to
+ # increase concurrency/ILP (accepts a small tail-predicate region).
+ if _ENABLE_TPR256 and cluster_n_override is None and not _ENABLE_STAGE2:
+ if dtype.width == 16 and N == 7168 and M >= 4096:
+ tpr_override = 256
+ nt_override = 256
+
+
+ can_use_256 = bool(
+ direct_gmem
+ and (
+ direct_gmem_max_copy_bits is None
+ or direct_gmem_max_copy_bits >= 256
+ )
+ and dtype.width == 16
+ and (x.data_ptr() % 32) == 0
+ and (residual.data_ptr() % 32) == 0
+ and (weight.data_ptr() % 32) == 0
+ )
+ assumed_align = 32 if can_use_256 else 16
+ if can_use_256:
+ copy_bits = 256
+
+ copy_bits = _copy_bits_from_policy(default=copy_bits, can_use_256=can_use_256)
+ if copy_bits == 128:
+ assumed_align = 16
+ elif copy_bits == 256 and can_use_256:
+ assumed_align = 32
+ else:
+ copy_bits = 128
+ assumed_align = 16
+
+ key = (
+ "ptr_fused_add_inplace",
+ N,
+ dtype,
+ stage,
+ device_index,
+ copy_bits,
+ use_async,
+ tpr_override,
+ nt_override,
+ direct_gmem,
+ cluster_n_override,
+ )
+ compiled = _PTR_COMPILE_CACHE.get(key)
+ if compiled is None:
+ op = RMSNormSM100(
+ N,
+ dtype,
+ stage=stage,
+ copy_bits=copy_bits,
+ use_async=use_async,
+ direct_gmem=direct_gmem,
+ )
+ if tpr_override is not None:
+ op._tpr_override = tpr_override # type: ignore[attr-defined]
+ if nt_override is not None:
+ op._nt_override = nt_override # type: ignore[attr-defined]
+ if cluster_n_override is not None:
+ op._cluster_n_override = cluster_n_override # type: ignore[attr-defined]
+ ptr_x = rt.make_ptr(
+ dtype, x.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=assumed_align
+ )
+ ptr_res = rt.make_ptr(
+ dtype, residual.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=assumed_align
+ )
+ ptr_w = rt.make_ptr(
+ dtype, weight.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=assumed_align
+ )
+ stream = cuda.CUstream(stream_handle)
+ ld_x = Int32(int(x.stride(0)))
+ compiled = cute.compile(
+ op.launch_from_ptrs_fused_add_inplace,
+ ptr_x,
+ ptr_w,
+ ptr_res,
+ Int32(M),
+ Int32(N),
+ ld_x,
+ stream,
+ Float32(eps),
+ )
+ _PTR_COMPILE_CACHE[key] = compiled
+ launcher = _get_fast_ptr_fused_add_rmsnorm_launcher(
+ compiled=compiled,
+ dtype=dtype,
+ N=N,
+ device_index=device_index,
+ stream_handle=stream_handle,
+ copy_bits=copy_bits,
+ use_async=use_async,
+ tpr=tpr_override or 0,
+ direct_gmem=direct_gmem,
+ assumed_align=assumed_align,
+ eps=eps,
+ )
+ if launcher is not None:
+ launcher.launch(
+ x=x,
+ weight=weight,
+ residual=residual,
+ M=M,
+ N=N,
+ ld_x=int(x.stride(0)),
+ eps=eps,
+ )
+ return
+
+ # Fast-launch is disabled/unavailable (or CuTeDSL internals changed). Fall back
+ # to calling the compiled function directly.
+ ptr_x = rt.make_ptr(dtype, x.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=assumed_align)
+ ptr_res = rt.make_ptr(
+ dtype, residual.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=assumed_align
+ )
+ ptr_w = rt.make_ptr(
+ dtype, weight.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=assumed_align
+ )
+ stream = cuda.CUstream(stream_handle)
+ ld_x = Int32(int(x.stride(0)))
+ compiled(ptr_x, ptr_w, ptr_res, Int32(M), Int32(N), ld_x, stream, Float32(eps))
+
+
+# -------------------------
+# Public API (forward + verify)
+# -------------------------
+
+
+def rmsnorm_forward(
+ x: Tensor,
+ weight: Optional[Tensor] = None,
+ bias: Optional[Tensor] = None,
+ residual: Optional[Tensor] = None,
+ eps: float = 1e-6,
+ store_rstd: bool = False,
+) -> Tuple[Tensor, Optional[Tensor], Optional[Tensor]]:
+ assert x.is_cuda
+ assert x.dim() == 2, "Use (M, N) tensor; flatten batch/seq beforehand."
+ M, N = x.shape
+ dtype = TORCH2CUTE_DTYPE[x.dtype]
+
+ # For DSv3 big-M outliers on SM100, keep using the dedicated
+ # stage-2 K-loop implementation, which is already tuned and
+ # parity-checked against the reference.
+ use_stage2_big_dsv3 = bool(
+ M >= 65536 and N in (6144, 8192) and x.dtype in (torch.float16, torch.bfloat16)
+ )
+ if use_stage2_big_dsv3:
+ try:
+ import rmsnorm_with_stage2 as rms2 # type: ignore[import-not-found]
+ except Exception:
+ rms2 = None # type: ignore[assignment]
+ if rms2 is not None:
+ y, rstd, residual_out = rms2.rmsnorm_forward_with_stage2(
+ x, weight=weight, bias=bias, residual=residual, eps=eps, store_rstd=store_rstd
+ )
+ # Preserve stride contracts for torch.compile consistency, even
+ # when using the optional stage-2 implementation.
+ if y.stride() != x.stride():
+ y_strided = torch.empty_strided(
+ x.shape, x.stride(), device=x.device, dtype=x.dtype
+ )
+ y_strided.copy_(y)
+ y = y_strided
+ if residual is not None and residual_out is not None:
+ if residual_out.stride() != residual.stride():
+ residual_out_strided = torch.empty_strided(
+ residual.shape,
+ residual.stride(),
+ device=residual.device,
+ dtype=residual.dtype,
+ )
+ residual_out_strided.copy_(residual_out)
+ residual_out = residual_out_strided
+ return y, rstd, residual_out
+
+ # Default: use the pointer-based entry whenever we can represent the
+ # inputs as a row-major [M, N] view with stride(1) == 1. For rare layouts
+ # we can't safely express without DLPack, fall back to a torch reference.
+ if _can_use_ptr_path(x, weight, bias, residual):
+ return _rmsnorm_forward_ptr(x, weight, bias, residual, eps, store_rstd)
+
+ # Safe fallback (correctness-first). This is expected to be rare in vLLM.
+ y = rmsnorm_ref(x, weight, bias, residual, eps)
+ # Preserve the input stride contract even on the fallback path so
+ # torch.compile sees a consistent output layout across all branches.
+ if y.stride() != x.stride():
+ y_strided = torch.empty_strided(x.shape, x.stride(), device=x.device, dtype=x.dtype)
+ y_strided.copy_(y)
+ y = y_strided
+ rstd = None
+ if store_rstd:
+ xf = x.float()
+ if residual is not None:
+ xf = xf + residual.float()
+ rstd = torch.rsqrt(xf.square().mean(dim=-1) + eps).to(torch.float32)
+ residual_out = None
+ if residual is not None:
+ residual_out = (x.float() + residual.float()).to(x.dtype)
+ if residual_out.stride() != residual.stride():
+ residual_out_strided = torch.empty_strided(
+ residual.shape,
+ residual.stride(),
+ device=residual.device,
+ dtype=residual.dtype,
+ )
+ residual_out_strided.copy_(residual_out)
+ residual_out = residual_out_strided
+ return y, rstd, residual_out
+
+
+def rmsnorm_ref(
+ x: Tensor,
+ w: Optional[Tensor] = None,
+ b: Optional[Tensor] = None,
+ residual: Optional[Tensor] = None,
+ eps: float = 1e-6,
+) -> Tensor:
+ xf = x.float()
+ if residual is not None:
+ xf = xf + residual.float()
+ rstd = torch.rsqrt(xf.square().mean(dim=-1, keepdim=True) + eps)
+ y = xf * rstd
+ if w is not None:
+ y = y * w.float()
+ if b is not None:
+ y = y + b.float()
+ return y.to(x.dtype)
+
+
+def fused_add_rmsnorm_forward(
+ x: Tensor,
+ residual: Tensor,
+ weight: Tensor,
+ eps: float = 1e-6,
+) -> Tuple[Tensor, Tensor]:
+ """Fused residual-add + RMSNorm for SM100 in CuteDSL.
+
+ This is a convenience wrapper around ``rmsnorm_forward`` that matches the
+ semantics of vLLM's ``fused_add_rms_norm``:
+
+ z = x + residual
+ y = RMSNorm(z, weight, eps)
+
+ It returns ``(y, z)`` where ``z`` has the same dtype/shape as the inputs.
+ """
+ assert x.is_cuda and residual.is_cuda
+ assert x.shape == residual.shape
+ assert x.dtype == residual.dtype
+
+ orig_shape = x.shape
+ N = orig_shape[-1]
+
+ x_2d = x.view(-1, N)
+ res_2d = residual.view(-1, N)
+
+ y_2d, _rstd, z_2d = rmsnorm_forward(
+ x_2d,
+ weight=weight,
+ bias=None,
+ residual=res_2d,
+ eps=eps,
+ store_rstd=False,
+ )
+
+ y = y_2d.view(orig_shape)
+ z = z_2d.view(orig_shape)
+ return y, z
+
+
+def fused_add_rmsnorm_forward_inplace(
+ x: Tensor,
+ residual: Tensor,
+ weight: Tensor,
+ eps: float = 1e-6,
+) -> Tuple[Tensor, Tensor]:
+ """In-place fused residual-add + RMSNorm matching vLLM semantics.
+
+ This variant writes:
+
+ z = x + residual (stored into ``residual``)
+ y = RMSNorm(z, w) (stored into ``x``)
+
+ i.e., it uses ``x`` as the normalized output buffer and ``residual`` as
+ the residual-out buffer, mirroring vLLM's fused_add_rms_norm kernel.
+ """
+ fused_add_rmsnorm_inplace_(x, residual, weight, eps=eps)
+ return x, residual
+
+
+def fused_add_rmsnorm_inplace_(
+ x: Tensor,
+ residual: Tensor,
+ weight: Tensor,
+ eps: float = 1e-6,
+) -> None:
+ """In-place fused residual-add + RMSNorm matching vLLM semantics.
+
+ This is the lowest-overhead Python entrypoint (returns `None`) intended
+ for performance-critical call sites like `torch.ops.oink.fused_add_rms_norm`.
+ """
+ assert x.is_cuda and residual.is_cuda
+ assert x.shape == residual.shape
+ assert x.dtype == residual.dtype
+
+ N = x.shape[-1]
+ x_2d = x if x.dim() == 2 else x.view(-1, N)
+ res_2d = residual if residual.dim() == 2 else residual.view(-1, N)
+
+ # Fast path: vLLM-compatible layout where x may be strided/padded but
+ # residual is contiguous. This updates both tensors in-place without
+ # additional allocations.
+ if _can_use_ptr_path_fused_add_inplace(x_2d, weight, res_2d):
+ _fused_add_rmsnorm_forward_ptr_inplace(x_2d, res_2d, weight, eps)
+ return None
+
+ # Fallback: allocate via the regular fused path, then copy results into
+ # the user-provided buffers so that semantics remain identical.
+ y, z = fused_add_rmsnorm_forward(x, residual, weight, eps)
+ x.copy_(y)
+ residual.copy_(z)
+ return None
+
+
+if __name__ == "__main__":
+ # Minimal ad-hoc test (functionality only). For performance comparisons, use the benchmark harness.
+ if not torch.cuda.is_available():
+ print("CUDA not available; functional test skipped.")
+ sys.exit(0)
+ M, N = 1024, 8192
+ dtype = torch.bfloat16
+ x = torch.randn(M, N, device="cuda", dtype=dtype)
+ w = torch.randn(N, device="cuda", dtype=dtype)
+ y_ref = rmsnorm_ref(x, w)
+ y, _, _ = rmsnorm_forward(x, w)
+ torch.testing.assert_close(y, y_ref, rtol=1e-3, atol=1e-3)
+ print("RMSNormSM100 correctness check passed.")
+
+# (compile cache moved to top)
From 3003c1398f700cf13095eb40f1e5bab84de013f7 Mon Sep 17 00:00:00 2001
From: Laura Wang <3700467+Laurawly@users.noreply.github.com>
Date: Tue, 6 Jan 2026 12:25:31 -0800
Subject: [PATCH 2/8] Fix oink ruff lint and add license headers
---
oink/src/kernelagent_oink/__init__.py | 14 ++++++++
.../kernelagent_oink/blackwell/__init__.py | 14 ++++++++
.../kernelagent_oink/blackwell/lite_quack.py | 20 ++++++++---
.../blackwell/oink_custom_ops.py | 16 ++++++++-
.../src/kernelagent_oink/blackwell/rmsnorm.py | 33 +++++++++++++------
5 files changed, 82 insertions(+), 15 deletions(-)
diff --git a/oink/src/kernelagent_oink/__init__.py b/oink/src/kernelagent_oink/__init__.py
index 542e59e..d9f25d0 100644
--- a/oink/src/kernelagent_oink/__init__.py
+++ b/oink/src/kernelagent_oink/__init__.py
@@ -1,3 +1,17 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
from __future__ import annotations
import logging
diff --git a/oink/src/kernelagent_oink/blackwell/__init__.py b/oink/src/kernelagent_oink/blackwell/__init__.py
index 4d21ee8..a92109a 100644
--- a/oink/src/kernelagent_oink/blackwell/__init__.py
+++ b/oink/src/kernelagent_oink/blackwell/__init__.py
@@ -1,3 +1,17 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
from __future__ import annotations
__all__ = []
diff --git a/oink/src/kernelagent_oink/blackwell/lite_quack.py b/oink/src/kernelagent_oink/blackwell/lite_quack.py
index 3c3f750..8c05b47 100644
--- a/oink/src/kernelagent_oink/blackwell/lite_quack.py
+++ b/oink/src/kernelagent_oink/blackwell/lite_quack.py
@@ -1,3 +1,17 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
"""
Lightweight local clone of the small subset of helpers that the SM100
RMSNorm CuteDSL kernels depend on.
@@ -12,9 +26,8 @@
import math
import operator
-from typing import Callable, Optional, Tuple
+from typing import Callable, Optional
-import cuda.bindings.driver as cuda # type: ignore
import torch
from torch import Tensor
@@ -23,7 +36,7 @@
from cutlass import Float32, Int32, const_expr
from cutlass.cute.runtime import from_dlpack
from cutlass.cutlass_dsl import T, dsl_user_op
-from cutlass._mlir.dialects import llvm, nvvm, vector
+from cutlass._mlir.dialects import llvm
# -------------------------
@@ -347,4 +360,3 @@ def get_sm_count(N: int, device: torch.device) -> int:
sm_count * sm_count_multiple if N <= 8192 else sm_count // 2 if N <= 16384 else sm_count * 2
)
return sm_count
-
diff --git a/oink/src/kernelagent_oink/blackwell/oink_custom_ops.py b/oink/src/kernelagent_oink/blackwell/oink_custom_ops.py
index 8225025..92423d9 100644
--- a/oink/src/kernelagent_oink/blackwell/oink_custom_ops.py
+++ b/oink/src/kernelagent_oink/blackwell/oink_custom_ops.py
@@ -1,4 +1,16 @@
-from __future__ import annotations
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
"""
Torch custom ops wrapping Oink's Blackwell RMSNorm kernels.
@@ -26,6 +38,8 @@
Mutates `x` and `residual` in-place and returns None.
"""
+from __future__ import annotations
+
import importlib
import threading
diff --git a/oink/src/kernelagent_oink/blackwell/rmsnorm.py b/oink/src/kernelagent_oink/blackwell/rmsnorm.py
index d6c2c20..a77938c 100644
--- a/oink/src/kernelagent_oink/blackwell/rmsnorm.py
+++ b/oink/src/kernelagent_oink/blackwell/rmsnorm.py
@@ -1,3 +1,17 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
"""
RMSNorm kernel for SM100 (Blackwell) in CuteDSL.
@@ -53,15 +67,15 @@
"(`cutlass`, typically provided by `nvidia-cutlass-dsl`)."
) from e
-import torch
-from torch import Tensor
+import torch # noqa: E402
+from torch import Tensor # noqa: E402
-import cuda.bindings.driver as cuda # provided by NVIDIA cuda-python
+import cuda.bindings.driver as cuda # provided by NVIDIA cuda-python # noqa: E402
-import cutlass
-import cutlass.cute as cute
-from cutlass import Float32, Int32, const_expr
-from cutlass.cute import runtime as rt
+import cutlass # noqa: E402
+import cutlass.cute as cute # noqa: E402
+from cutlass import Float32, Int32, const_expr # noqa: E402
+from cutlass.cute import runtime as rt # noqa: E402
# Simple compile cache declared early so direct execution works
_PTR_COMPILE_CACHE = {}
@@ -862,8 +876,8 @@ def _get_fast_ptr_fused_add_rmsnorm_launcher(
#
# NOTE: Avoid `from . import ...` imports here: CuTeDSL's AST preprocessor may
# mishandle that form (module=None in the AST). Use fully-qualified imports.
-from kernelagent_oink.blackwell import lite_quack as qutils
-from kernelagent_oink.blackwell.lite_quack import TORCH2CUTE_DTYPE, row_reduce
+from kernelagent_oink.blackwell import lite_quack as qutils # noqa: E402
+from kernelagent_oink.blackwell.lite_quack import TORCH2CUTE_DTYPE, row_reduce # noqa: E402
# -------------------------
@@ -2458,7 +2472,6 @@ def rmsnorm_forward(
assert x.is_cuda
assert x.dim() == 2, "Use (M, N) tensor; flatten batch/seq beforehand."
M, N = x.shape
- dtype = TORCH2CUTE_DTYPE[x.dtype]
# For DSv3 big-M outliers on SM100, keep using the dedicated
# stage-2 K-loop implementation, which is already tuned and
From 1468088ecab42c675d8e4d9e0bf465bd5f68644a Mon Sep 17 00:00:00 2001
From: Laura Wang <3700467+Laurawly@users.noreply.github.com>
Date: Tue, 6 Jan 2026 12:27:54 -0800
Subject: [PATCH 3/8] Format oink with ruff
---
oink/src/kernelagent_oink/__init__.py | 4 +-
.../kernelagent_oink/blackwell/lite_quack.py | 41 +-
.../blackwell/oink_custom_ops.py | 6 +-
.../src/kernelagent_oink/blackwell/rmsnorm.py | 472 ++++++++++++++----
4 files changed, 403 insertions(+), 120 deletions(-)
diff --git a/oink/src/kernelagent_oink/__init__.py b/oink/src/kernelagent_oink/__init__.py
index d9f25d0..bbbd7c1 100644
--- a/oink/src/kernelagent_oink/__init__.py
+++ b/oink/src/kernelagent_oink/__init__.py
@@ -94,7 +94,9 @@ def register() -> None:
# Ensure CuTeDSL sees a target arch early. If the user has already set it,
# respect their choice.
- os.environ.setdefault("CUTE_DSL_ARCH", _compute_cutedsl_arch(int(major), int(minor)))
+ os.environ.setdefault(
+ "CUTE_DSL_ARCH", _compute_cutedsl_arch(int(major), int(minor))
+ )
# Import registers the ops via torch.library.custom_op decorators.
from .blackwell import oink_custom_ops # noqa: F401
diff --git a/oink/src/kernelagent_oink/blackwell/lite_quack.py b/oink/src/kernelagent_oink/blackwell/lite_quack.py
index 8c05b47..14ae723 100644
--- a/oink/src/kernelagent_oink/blackwell/lite_quack.py
+++ b/oink/src/kernelagent_oink/blackwell/lite_quack.py
@@ -54,6 +54,7 @@
# Tensor conversion helpers
# -------------------------
+
def convert_from_dlpack(
x: Tensor,
leading_dim: int,
@@ -82,7 +83,9 @@ def convert_from_dlpack(
@dsl_user_op
-def elem_pointer(x: cute.Tensor, coord: cute.Coord, *, loc=None, ip=None) -> cute.Pointer:
+def elem_pointer(
+ x: cute.Tensor, coord: cute.Coord, *, loc=None, ip=None
+) -> cute.Pointer:
return x.iterator + cute.crd2idx(coord, x.layout, loc=loc, ip=ip)
@@ -133,7 +136,9 @@ def store_shared_remote(
).ir_value()
if const_expr(isinstance(val, float)):
val = Float32(val)
- assert isinstance(val, (Float32, Int32, cutlass.Int64)), "val must be Float32, Int32, or Int64"
+ assert isinstance(val, (Float32, Int32, cutlass.Int64)), (
+ "val must be Float32, Int32, or Int64"
+ )
suffix = {Float32: "f32", Int32: "s32", cutlass.Int64: "s64"}[type(val)]
constraint = {Float32: "f", Int32: "r", cutlass.Int64: "l"}[type(val)]
llvm.inline_asm(
@@ -155,19 +160,27 @@ def predicate_k(tAcA: cute.Tensor, limit: cutlass.Int32) -> cute.Tensor:
"""
tApA = cute.make_fragment(
cute.make_layout(
- (cute.size(tAcA, mode=[0, 1]), cute.size(tAcA, mode=[1]), cute.size(tAcA, mode=[2])),
+ (
+ cute.size(tAcA, mode=[0, 1]),
+ cute.size(tAcA, mode=[1]),
+ cute.size(tAcA, mode=[2]),
+ ),
stride=(cute.size(tAcA, mode=[2]), 0, 1),
),
cutlass.Boolean,
)
for rest_v in cutlass.range_constexpr(tApA.shape[0]):
for rest_k in cutlass.range_constexpr(tApA.shape[2]):
- tApA[rest_v, 0, rest_k] = cute.elem_less(tAcA[(0, rest_v), 0, rest_k][1], limit)
+ tApA[rest_v, 0, rest_k] = cute.elem_less(
+ tAcA[(0, rest_v), 0, rest_k][1], limit
+ )
return tApA
@dsl_user_op
-def domain_offset_i64(coord: cute.Coord, tensor: cute.Tensor, *, loc=None, ip=None) -> cute.Tensor:
+def domain_offset_i64(
+ coord: cute.Coord, tensor: cute.Tensor, *, loc=None, ip=None
+) -> cute.Tensor:
"""
Return a tensor whose iterator is offset by an Int64 byte offset
computed from `coord` and the tensor's strides.
@@ -287,7 +300,9 @@ def block_or_cluster_reduce(
"""Dispatch between block or cluster reduction depending on mbar_ptr."""
if cutlass.const_expr(mbar_ptr is None):
return block_reduce(val, op, reduction_buffer, init_val=init_val)
- return cluster_reduce(val, op, reduction_buffer, mbar_ptr, init_val=init_val, phase=phase)
+ return cluster_reduce(
+ val, op, reduction_buffer, mbar_ptr, init_val=init_val, phase=phase
+ )
@cute.jit
@@ -313,7 +328,9 @@ def row_reduce(
val = x
warp_op = {
cute.ReductionOp.ADD: operator.add,
- cute.ReductionOp.MAX: cute.arch.fmax if cutlass.const_expr(x.dtype == Float32) else max,
+ cute.ReductionOp.MAX: cute.arch.fmax
+ if cutlass.const_expr(x.dtype == Float32)
+ else max,
cute.ReductionOp.MIN: min,
cute.ReductionOp.MUL: operator.mul,
}[op]
@@ -353,10 +370,16 @@ def get_sm_count(N: int, device: torch.device) -> int:
RMSNorm kernels but lives entirely in this local module.
"""
sm_count_multiple = (
- 16 if N <= 256 else (8 if N <= 1024 else (4 if N <= 2048 else (2 if N <= 4096 else 1)))
+ 16
+ if N <= 256
+ else (8 if N <= 1024 else (4 if N <= 2048 else (2 if N <= 4096 else 1)))
)
sm_count = torch.cuda.get_device_properties(device).multi_processor_count
sm_count = (
- sm_count * sm_count_multiple if N <= 8192 else sm_count // 2 if N <= 16384 else sm_count * 2
+ sm_count * sm_count_multiple
+ if N <= 8192
+ else sm_count // 2
+ if N <= 16384
+ else sm_count * 2
)
return sm_count
diff --git a/oink/src/kernelagent_oink/blackwell/oink_custom_ops.py b/oink/src/kernelagent_oink/blackwell/oink_custom_ops.py
index 92423d9..a96a4c7 100644
--- a/oink/src/kernelagent_oink/blackwell/oink_custom_ops.py
+++ b/oink/src/kernelagent_oink/blackwell/oink_custom_ops.py
@@ -80,6 +80,7 @@ def _get_sm(device: torch.device | None = None) -> int:
# RMSNorm (functional)
#
+
@custom_op("oink::rmsnorm", mutates_args=())
def oink_rmsnorm(
x: torch.Tensor,
@@ -158,6 +159,7 @@ def oink_rmsnorm_fake(
# Fused residual-add + RMSNorm (in-place, vLLM semantics)
#
+
@custom_op("oink::fused_add_rms_norm", mutates_args=("x", "residual"))
def oink_fused_add_rms_norm(
x: torch.Tensor,
@@ -174,7 +176,9 @@ def oink_fused_add_rms_norm(
Returns:
None (mutates `x` and `residual` in-place).
"""
- assert x.is_cuda and residual.is_cuda, "oink::fused_add_rms_norm requires CUDA tensors"
+ assert x.is_cuda and residual.is_cuda, (
+ "oink::fused_add_rms_norm requires CUDA tensors"
+ )
assert x.shape == residual.shape, "x and residual must have the same shape"
assert x.dtype == residual.dtype, "x and residual must have the same dtype"
assert weight.dim() == 1, "weight must be 1D [N]"
diff --git a/oink/src/kernelagent_oink/blackwell/rmsnorm.py b/oink/src/kernelagent_oink/blackwell/rmsnorm.py
index a77938c..c7fc1b3 100644
--- a/oink/src/kernelagent_oink/blackwell/rmsnorm.py
+++ b/oink/src/kernelagent_oink/blackwell/rmsnorm.py
@@ -84,6 +84,7 @@
# pointer/scalar storage so concurrent callers don't race on in-place updates.
_PTR_FAST_LAUNCH_TLS = threading.local()
+
def _env_flag(name: str, default: bool) -> bool:
val = os.environ.get(name)
if val is None:
@@ -100,10 +101,16 @@ def _env_flag(name: str, default: bool) -> bool:
# Fused-add RMSNorm schedule knobs (read once at import time; set env vars before
# importing this module if you want to override).
-_DIRECT_GMEM_POLICY = (os.environ.get("OINK_RMSNORM_DIRECT_GMEM", "auto").strip().lower() or "auto")
-_COPY_BITS_POLICY = (os.environ.get("OINK_RMSNORM_COPY_BITS", "auto").strip().lower() or "auto")
+_DIRECT_GMEM_POLICY = (
+ os.environ.get("OINK_RMSNORM_DIRECT_GMEM", "auto").strip().lower() or "auto"
+)
+_COPY_BITS_POLICY = (
+ os.environ.get("OINK_RMSNORM_COPY_BITS", "auto").strip().lower() or "auto"
+)
_ENABLE_CLUSTER_ILP = _env_flag("OINK_RMSNORM_ENABLE_CLUSTER_ILP", default=False)
-_ENABLE_CLUSTER_ILP_UNSAFE = _env_flag("OINK_RMSNORM_ENABLE_CLUSTER_ILP_UNSAFE", default=False)
+_ENABLE_CLUSTER_ILP_UNSAFE = _env_flag(
+ "OINK_RMSNORM_ENABLE_CLUSTER_ILP_UNSAFE", default=False
+)
_ENABLE_TPR256 = _env_flag("OINK_RMSNORM_ENABLE_TPR256", default=False)
_ENABLE_STAGE2 = _env_flag("OINK_RMSNORM_ENABLE_STAGE2", default=False)
@@ -252,6 +259,7 @@ def run_probe(copy_bits: int, assumed_align: int):
_CLUSTER_DIRECT_GMEM_MAX_COPY_BITS = max_bits
return max_bits
+
def _parse_version_tuple(version: str) -> Tuple[int, int, int]:
parts = version.split(".")
nums: list[int] = []
@@ -275,7 +283,9 @@ def _cutlass_dsl_version() -> Optional[Tuple[int, int, int]]:
# passing Layout/Shape/Constexpr objects into @cute.kernel functions). Keep the
# older signature for 4.3.2, but switch to a 4.3.4-compatible signature when we
# detect 4.3.4+ (or when version detection is unavailable).
-_KERNEL_ACCEPTS_LAYOUT_ARGS = _CUTLASS_DSL_VERSION is not None and _CUTLASS_DSL_VERSION < (4, 3, 4)
+_KERNEL_ACCEPTS_LAYOUT_ARGS = (
+ _CUTLASS_DSL_VERSION is not None and _CUTLASS_DSL_VERSION < (4, 3, 4)
+)
if _ENABLE_CLUSTER_ILP and not _ENABLE_CLUSTER_ILP_UNSAFE:
# We have observed reproducible segfaults in some CuTeDSL builds when using
@@ -427,7 +437,9 @@ def launch(
self._last_x_ptr = x_ptr
except AttributeError:
self._disable_fast_launch()
- self._fallback_launch(x=x, weight=weight, out=out, M=M, N=N, ld=ld, eps=eps)
+ self._fallback_launch(
+ x=x, weight=weight, out=out, M=M, N=N, ld=ld, eps=eps
+ )
return
if self._ptr_w is not None:
@@ -438,7 +450,9 @@ def launch(
self._last_w_ptr = w_ptr
except AttributeError:
self._disable_fast_launch()
- self._fallback_launch(x=x, weight=weight, out=out, M=M, N=N, ld=ld, eps=eps)
+ self._fallback_launch(
+ x=x, weight=weight, out=out, M=M, N=N, ld=ld, eps=eps
+ )
return
out_ptr = out.data_ptr()
@@ -448,7 +462,9 @@ def launch(
self._last_out_ptr = out_ptr
except AttributeError:
self._disable_fast_launch()
- self._fallback_launch(x=x, weight=weight, out=out, M=M, N=N, ld=ld, eps=eps)
+ self._fallback_launch(
+ x=x, weight=weight, out=out, M=M, N=N, ld=ld, eps=eps
+ )
return
if M != self._last_m:
@@ -492,10 +508,19 @@ def _fallback_launch(
# If the packed-args or runtime pointer mutation path stops working
# (e.g. due to a CuTeDSL upgrade), fall back to the regular call path.
dtype = TORCH2CUTE_DTYPE[x.dtype]
- ptr_x = rt.make_ptr(dtype, x.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
- ptr_out = rt.make_ptr(dtype, out.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ ptr_x = rt.make_ptr(
+ dtype, x.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16
+ )
+ ptr_out = rt.make_ptr(
+ dtype, out.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16
+ )
ptr_w = (
- rt.make_ptr(dtype, weight.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ rt.make_ptr(
+ dtype,
+ weight.data_ptr(),
+ mem_space=rt.AddressSpace.gmem,
+ assumed_align=16,
+ )
if weight is not None
else None
)
@@ -695,7 +720,15 @@ def _get_fast_ptr_rmsnorm_launcher(
return None
# Keyed by the compiled object identity so schedule changes (e.g. copy width,
# async/staged variants, etc.) never alias in the fast-launch cache.
- key = ("ptr_fast", id(compiled), N, dtype, device_index, int(stream_handle), has_weight)
+ key = (
+ "ptr_fast",
+ id(compiled),
+ N,
+ dtype,
+ device_index,
+ int(stream_handle),
+ has_weight,
+ )
cache = _tls_fast_launch_cache()
cached = cache.get(key)
if cached is not None:
@@ -705,7 +738,9 @@ def _get_fast_ptr_rmsnorm_launcher(
ptr_x = rt.make_ptr(dtype, 0, mem_space=rt.AddressSpace.gmem, assumed_align=16)
ptr_out = rt.make_ptr(dtype, 0, mem_space=rt.AddressSpace.gmem, assumed_align=16)
ptr_w = (
- rt.make_ptr(dtype, 0, mem_space=rt.AddressSpace.gmem, assumed_align=16) if has_weight else None
+ rt.make_ptr(dtype, 0, mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ if has_weight
+ else None
)
arg_m = _StableI32Arg(0)
@@ -808,9 +843,15 @@ def _get_fast_ptr_fused_add_rmsnorm_launcher(
if cached is not None:
return cached # type: ignore[return-value]
- ptr_x = rt.make_ptr(dtype, 0, mem_space=rt.AddressSpace.gmem, assumed_align=assumed_align)
- ptr_res = rt.make_ptr(dtype, 0, mem_space=rt.AddressSpace.gmem, assumed_align=assumed_align)
- ptr_w = rt.make_ptr(dtype, 0, mem_space=rt.AddressSpace.gmem, assumed_align=assumed_align)
+ ptr_x = rt.make_ptr(
+ dtype, 0, mem_space=rt.AddressSpace.gmem, assumed_align=assumed_align
+ )
+ ptr_res = rt.make_ptr(
+ dtype, 0, mem_space=rt.AddressSpace.gmem, assumed_align=assumed_align
+ )
+ ptr_w = rt.make_ptr(
+ dtype, 0, mem_space=rt.AddressSpace.gmem, assumed_align=assumed_align
+ )
arg_m = _StableI32Arg(0)
arg_n = _StableI32Arg(N)
@@ -884,6 +925,7 @@ def _get_fast_ptr_fused_add_rmsnorm_launcher(
# Copy helpers (allow up to 256b)
# -------------------------
+
@cute.jit
def get_copy_atom_bw(
dtype: type[cutlass.Numeric], num_copy_elems: int, is_async: bool = False
@@ -892,6 +934,7 @@ def get_copy_atom_bw(
max_bits = const_expr(128 if is_async else 256)
num_copy_bits = const_expr(min(max_bits, num_copy_elems * dtype.width))
from cutlass.cute.nvgpu import cpasync
+
# Prefer GLOBAL cache policy for bulk streaming reads at large M
copy_op = (
cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL)
@@ -1037,7 +1080,10 @@ def _tv_layout(self, num_copy_bits: int = 256) -> Tuple[cute.Shape, cute.Layout]
tiler_mn = (cols_per_block, vecsize * num_blocks_N * tpr)
tv_layout = cute.make_layout(
((tpr, cols_per_block), (vecsize, num_blocks_N)),
- stride=((vecsize * cols_per_block, 1), (cols_per_block, cols_per_block * vecsize * tpr)),
+ stride=(
+ (vecsize * cols_per_block, 1),
+ (cols_per_block, cols_per_block * vecsize * tpr),
+ ),
)
return tiler_mn, tv_layout
@@ -1045,7 +1091,10 @@ def _smem_bytes(self, tiler_mn, num_warps) -> int:
# smem for X tile (+ residual if present) + reduction buffers + mbar(s)
return (
cute.size_in_bytes(self.dtype, cute.make_layout(tiler_mn))
- + self.stage * num_warps * self._cluster_n() * (self.reduction_dtype.width // 8)
+ + self.stage
+ * num_warps
+ * self._cluster_n()
+ * (self.reduction_dtype.width // 8)
+ self.stage * (cutlass.Int64.width // 8)
)
@@ -1072,7 +1121,9 @@ def new_stride(t):
)
mX, mRes, mO, mResO = [
- cute.make_tensor(t.iterator, cute.make_layout(semistatic_shape, stride=new_stride(t)))
+ cute.make_tensor(
+ t.iterator, cute.make_layout(semistatic_shape, stride=new_stride(t))
+ )
if const_expr(t is not None)
else None
for t in (mX, mRes, mO, mResO)
@@ -1082,23 +1133,34 @@ def new_stride(t):
copy_bits = int(self.copy_bits)
tiler_mn, tv_layout = self._tv_layout(num_copy_bits=copy_bits)
- num_threads = cute.size(tv_layout, mode=[0]) if _KERNEL_ACCEPTS_LAYOUT_ARGS else self._num_threads()
+ num_threads = (
+ cute.size(tv_layout, mode=[0])
+ if _KERNEL_ACCEPTS_LAYOUT_ARGS
+ else self._num_threads()
+ )
num_warps = num_threads // cute.arch.WARP_SIZE
- threads_per_row = tv_layout.shape[0][0] if _KERNEL_ACCEPTS_LAYOUT_ARGS else self._threads_per_row()
+ threads_per_row = (
+ tv_layout.shape[0][0]
+ if _KERNEL_ACCEPTS_LAYOUT_ARGS
+ else self._threads_per_row()
+ )
warps_per_row = max(threads_per_row // cute.arch.WARP_SIZE, 1)
cluster_n = self._cluster_n()
if const_expr(mW is not None):
mW = cute.make_tensor(
- mW.iterator, cute.prepend(mW.layout, cute.make_layout((tiler_mn[0],), stride=(0,)))
+ mW.iterator,
+ cute.prepend(mW.layout, cute.make_layout((tiler_mn[0],), stride=(0,))),
)
if const_expr(mB is not None):
mB = cute.make_tensor(
- mB.iterator, cute.prepend(mB.layout, cute.make_layout((tiler_mn[0],), stride=(0,)))
+ mB.iterator,
+ cute.prepend(mB.layout, cute.make_layout((tiler_mn[0],), stride=(0,))),
)
if const_expr(mRstd is not None):
mRstd = cute.make_tensor(
- mRstd.iterator, cute.append(mRstd.layout, cute.make_layout((self.N,), stride=(0,)))
+ mRstd.iterator,
+ cute.append(mRstd.layout, cute.make_layout((self.N,), stride=(0,))),
)
# No SMEM reload mode switch; overlap is controlled in the K-loop path
@@ -1114,11 +1176,14 @@ def new_stride(t):
else 0
)
tile_bytes_res = (
- cute.size_in_bytes(mRes.element_type, cute.make_layout(tiler_mn)) * stage_bufs
+ cute.size_in_bytes(mRes.element_type, cute.make_layout(tiler_mn))
+ * stage_bufs
if const_expr(mRes is not None and not self.direct_gmem)
else 0
)
- red_bytes = self.stage * num_warps * cluster_n * (self.reduction_dtype.width // 8)
+ red_bytes = (
+ self.stage * num_warps * cluster_n * (self.reduction_dtype.width // 8)
+ )
# mbarriers are only allocated/used for cluster_n>1. Some CuTeDSL builds
# require mbarrier state to be 16B-aligned in shared memory; account for
# the alignment padding when computing dynamic smem bytes.
@@ -1211,14 +1276,10 @@ def launch_from_ptrs(
else None
)
mW = (
- cute.make_tensor(ptr_w, layout_n)
- if const_expr(ptr_w is not None)
- else None
+ cute.make_tensor(ptr_w, layout_n) if const_expr(ptr_w is not None) else None
)
mB = (
- cute.make_tensor(ptr_b, layout_n)
- if const_expr(ptr_b is not None)
- else None
+ cute.make_tensor(ptr_b, layout_n) if const_expr(ptr_b is not None) else None
)
mRstd = (
cute.make_tensor(ptr_rstd, layout_m)
@@ -1303,7 +1364,9 @@ def _kernel_impl(
# Allocate one or two SMEM buffers depending on stage depth
sX0 = (
smem.allocate_tensor(
- mX.element_type, cute.make_ordered_layout(tiler_mn, order=(1, 0)), byte_alignment=32
+ mX.element_type,
+ cute.make_ordered_layout(tiler_mn, order=(1, 0)),
+ byte_alignment=32,
)
if const_expr(not self.direct_gmem)
else None
@@ -1319,7 +1382,9 @@ def _kernel_impl(
)
sRes0 = (
smem.allocate_tensor(
- mRes.element_type, cute.make_ordered_layout(tiler_mn, order=(1, 0)), byte_alignment=32
+ mRes.element_type,
+ cute.make_ordered_layout(tiler_mn, order=(1, 0)),
+ byte_alignment=32,
)
if const_expr(mRes is not None and not self.direct_gmem)
else None
@@ -1339,7 +1404,9 @@ def _kernel_impl(
(num_warps // warps_per_row, (warps_per_row, cluster_n), self.stage),
order=(1, 0, 2),
)
- reduction_buffer = smem.allocate_tensor(self.reduction_dtype, red_layout, byte_alignment=4)
+ reduction_buffer = smem.allocate_tensor(
+ self.reduction_dtype, red_layout, byte_alignment=4
+ )
if const_expr(cluster_n > 1):
# Some CuTeDSL builds appear sensitive to the shared-memory alignment of
# mbarrier state. `SmemAllocator.allocate_array` does not currently
@@ -1360,8 +1427,12 @@ def _kernel_impl(
# Tiled copy setup
num_copy_elems_X = tv_layout.shape[1][0]
- use_async = const_expr(self.use_async and self.N >= 1024 and not self.direct_gmem)
- copy_atom = get_copy_atom_bw(mX.element_type, num_copy_elems_X, is_async=use_async)
+ use_async = const_expr(
+ self.use_async and self.N >= 1024 and not self.direct_gmem
+ )
+ copy_atom = get_copy_atom_bw(
+ mX.element_type, num_copy_elems_X, is_async=use_async
+ )
thr_copy = cute.make_tiled_copy(copy_atom, tv_layout, tiler_mn).get_slice(tidx)
# Tail predicate for the N dimension (when tile width > N). Reuse this
@@ -1386,7 +1457,9 @@ def _kernel_impl(
mW is not None and (self.direct_gmem or (mRes is None and mB is None))
)
if const_expr(prefetch_w_early):
- gW = cute.local_tile(qutils.domain_offset_i64((0, n_off), mW), tiler_mn, (0, 0))
+ gW = cute.local_tile(
+ qutils.domain_offset_i64((0, n_off), mW), tiler_mn, (0, 0)
+ )
tXgW = thr_copy.partition_S(gW)
tXrW = cute.make_fragment_like(tXgW)
if const_expr(not is_even_N_wb):
@@ -1398,7 +1471,9 @@ def _kernel_impl(
pred=tXp_wb,
)
if const_expr(self.direct_gmem and mB is not None):
- gB = cute.local_tile(qutils.domain_offset_i64((0, n_off), mB), tiler_mn, (0, 0))
+ gB = cute.local_tile(
+ qutils.domain_offset_i64((0, n_off), mB), tiler_mn, (0, 0)
+ )
tXgB = thr_copy.partition_S(gB)
tXrB = cute.make_fragment_like(tXgB)
if const_expr(not is_even_N_wb):
@@ -1414,7 +1489,9 @@ def _kernel_impl(
self._init_cluster(tidx, mbar_ptr)
mX_i, mRes_i, mO_i, mResO_i = [
- qutils.domain_offset_i64((bidx * tiler_mn[0], 0), t) if t is not None else None
+ qutils.domain_offset_i64((bidx * tiler_mn[0], 0), t)
+ if t is not None
+ else None
for t in (mX, mRes, mO, mResO)
]
mX_i, mRes_i, mO_i, mResO_i = [
@@ -1424,27 +1501,39 @@ def _kernel_impl(
gX_i = cute.local_tile(mX_i, tiler_mn, (0, 0))
gO_i = cute.local_tile(mO_i, tiler_mn, (0, 0))
gRes_i = (
- cute.local_tile(mRes_i, tiler_mn, (0, 0)) if const_expr(mRes is not None) else None
+ cute.local_tile(mRes_i, tiler_mn, (0, 0))
+ if const_expr(mRes is not None)
+ else None
)
gResO_i = (
- cute.local_tile(mResO_i, tiler_mn, (0, 0)) if const_expr(mResO is not None) else None
+ cute.local_tile(mResO_i, tiler_mn, (0, 0))
+ if const_expr(mResO is not None)
+ else None
)
gRstd_i = (
- cute.local_tile(mRstd, tiler_mn, (bidx, 0)) if const_expr(mRstd is not None) else None
+ cute.local_tile(mRstd, tiler_mn, (bidx, 0))
+ if const_expr(mRstd is not None)
+ else None
)
cX_i = cute.local_tile(idX, tiler_mn, (bidx, 0))
# Common identity/row index partitions reused by both default and K-loop paths
tXcX_i = thr_copy.partition_S(cX_i)[(0, None), None, None]
row_i = tXcX_i[0][0]
- tXgRstd_i = thr_copy.partition_D(gRstd_i) if const_expr(mRstd is not None) else None
+ tXgRstd_i = (
+ thr_copy.partition_D(gRstd_i) if const_expr(mRstd is not None) else None
+ )
# Stage-2 intra-row K-loop cp.async ping-pong (two tiles). This reduces
# per-thread fragment size and can improve memory-latency hiding for
# N=7168 at large M. It is enabled by setting `stage=2` when constructing
# the RMSNormSM100 op (see `_fused_add_rmsnorm_forward_ptr_inplace`).
if const_expr(
- self.stage > 1 and not self.direct_gmem and use_async and cluster_n == 1 and shape[1] == 7168
+ self.stage > 1
+ and not self.direct_gmem
+ and use_async
+ and cluster_n == 1
+ and shape[1] == 7168
):
vecsize = tv_layout.shape[1][0]
tpr = threads_per_row
@@ -1475,9 +1564,9 @@ def _kernel_impl(
(tiler_mn[0], tiler_mn[0] * vecsize * tpr),
),
)
- thr_copy_tile = cute.make_tiled_copy(copy_atom, tv_layout_tile, tiler_mn_tile).get_slice(
- tidx
- )
+ thr_copy_tile = cute.make_tiled_copy(
+ copy_atom, tv_layout_tile, tiler_mn_tile
+ ).get_slice(tidx)
# Accumulate per-thread partial sums across tiles; reduce once.
sum_sq_thread = cute.Float32(0.0)
@@ -1499,7 +1588,13 @@ def _kernel_impl(
tXp_pong = tXp_0
if row_i < shape[0]:
- copy_tiled(tXgX_0, tXsX_0, num_copy_elems=vecsize, is_async=True, pred=tXp_0)
+ copy_tiled(
+ tXgX_0,
+ tXsX_0,
+ num_copy_elems=vecsize,
+ is_async=True,
+ pred=tXp_0,
+ )
if const_expr(mRes is not None):
gRes_0 = cute.local_tile(
qutils.domain_offset_i64((0, k_off0), mRes_i),
@@ -1538,19 +1633,27 @@ def _kernel_impl(
if const_expr((t % 2) == 0):
tXsX_n = thr_copy_tile.partition_D(sX1_tile)
tXsRes_n = (
- thr_copy_tile.partition_D(sRes1_tile) if const_expr(mRes is not None) else None
+ thr_copy_tile.partition_D(sRes1_tile)
+ if const_expr(mRes is not None)
+ else None
)
tXp_pong = tXp_n
else:
tXsX_n = thr_copy_tile.partition_D(sX0_tile)
tXsRes_n = (
- thr_copy_tile.partition_D(sRes0_tile) if const_expr(mRes is not None) else None
+ thr_copy_tile.partition_D(sRes0_tile)
+ if const_expr(mRes is not None)
+ else None
)
tXp_ping = tXp_n
if row_i < shape[0]:
copy_tiled(
- tXgX_n, tXsX_n, num_copy_elems=vecsize, is_async=True, pred=tXp_n
+ tXgX_n,
+ tXsX_n,
+ num_copy_elems=vecsize,
+ is_async=True,
+ pred=tXp_n,
)
if const_expr(mRes is not None):
gRes_n = cute.local_tile(
@@ -1574,25 +1677,35 @@ def _kernel_impl(
if const_expr((t % 2) == 0):
tXsX_cur = thr_copy_tile.partition_D(sX0_tile)
tXsRes_cur = (
- thr_copy_tile.partition_D(sRes0_tile) if const_expr(mRes is not None) else None
+ thr_copy_tile.partition_D(sRes0_tile)
+ if const_expr(mRes is not None)
+ else None
)
pred_cur = tXp_ping
else:
tXsX_cur = thr_copy_tile.partition_D(sX1_tile)
tXsRes_cur = (
- thr_copy_tile.partition_D(sRes1_tile) if const_expr(mRes is not None) else None
+ thr_copy_tile.partition_D(sRes1_tile)
+ if const_expr(mRes is not None)
+ else None
)
pred_cur = tXp_pong
k_off = t * tile_n
- gX_t = cute.local_tile(qutils.domain_offset_i64((0, k_off), mX_i), tiler_mn_tile, (0, 0))
+ gX_t = cute.local_tile(
+ qutils.domain_offset_i64((0, k_off), mX_i),
+ tiler_mn_tile,
+ (0, 0),
+ )
tXgX_t = thr_copy_tile.partition_S(gX_t)
tXrX_t = cute.make_fragment_like(tXgX_t)
cute.autovec_copy(tXsX_cur, tXrX_t)
x_t = tXrX_t.load().to(cute.Float32)
if const_expr(mRes is not None):
gRes_t = cute.local_tile(
- qutils.domain_offset_i64((0, k_off), mRes_i), tiler_mn_tile, (0, 0)
+ qutils.domain_offset_i64((0, k_off), mRes_i),
+ tiler_mn_tile,
+ (0, 0),
)
tXgRes_t = thr_copy_tile.partition_S(gRes_t)
tXrRes_t = cute.make_fragment_like(tXgRes_t)
@@ -1639,29 +1752,41 @@ def _kernel_impl(
for t in cutlass.range_constexpr(num_tiles):
k_off = t * tile_n
- cX_t = cute.local_tile(cute.domain_offset((0, k_off), cX_i), tiler_mn_tile, (0, 0))
+ cX_t = cute.local_tile(
+ cute.domain_offset((0, k_off), cX_i), tiler_mn_tile, (0, 0)
+ )
tXc_t = thr_copy_tile.partition_S(cX_t)
tXp_t = qutils.predicate_k(tXc_t, limit=limit_k)
if const_expr((t % 2) == 0):
tXsX_cur = thr_copy_tile.partition_D(sX0_tile)
tXsRes_cur = (
- thr_copy_tile.partition_D(sRes0_tile) if const_expr(mRes is not None) else None
+ thr_copy_tile.partition_D(sRes0_tile)
+ if const_expr(mRes is not None)
+ else None
)
else:
tXsX_cur = thr_copy_tile.partition_D(sX1_tile)
tXsRes_cur = (
- thr_copy_tile.partition_D(sRes1_tile) if const_expr(mRes is not None) else None
+ thr_copy_tile.partition_D(sRes1_tile)
+ if const_expr(mRes is not None)
+ else None
)
- gX_t = cute.local_tile(qutils.domain_offset_i64((0, k_off), mX_i), tiler_mn_tile, (0, 0))
+ gX_t = cute.local_tile(
+ qutils.domain_offset_i64((0, k_off), mX_i),
+ tiler_mn_tile,
+ (0, 0),
+ )
tXgX_t = thr_copy_tile.partition_S(gX_t)
tXrX_t = cute.make_fragment_like(tXgX_t)
cute.autovec_copy(tXsX_cur, tXrX_t)
x_t = tXrX_t.load().to(cute.Float32)
if const_expr(mRes is not None):
gRes_t = cute.local_tile(
- qutils.domain_offset_i64((0, k_off), mRes_i), tiler_mn_tile, (0, 0)
+ qutils.domain_offset_i64((0, k_off), mRes_i),
+ tiler_mn_tile,
+ (0, 0),
)
tXgRes_t = thr_copy_tile.partition_S(gRes_t)
tXrRes_t = cute.make_fragment_like(tXgRes_t)
@@ -1671,43 +1796,77 @@ def _kernel_impl(
y_t = x_t * rstd
if const_expr(mW is not None):
gW_t = cute.local_tile(
- qutils.domain_offset_i64((0, k_off), mW), tiler_mn_tile, (0, 0)
+ qutils.domain_offset_i64((0, k_off), mW),
+ tiler_mn_tile,
+ (0, 0),
)
tWgW_t = thr_copy_tile.partition_S(gW_t)
tWrW_t = cute.make_fragment_like(tWgW_t)
- copy_tiled(tWgW_t, tWrW_t, num_copy_elems=vecsize, is_async=False, pred=tXp_t)
+ copy_tiled(
+ tWgW_t,
+ tWrW_t,
+ num_copy_elems=vecsize,
+ is_async=False,
+ pred=tXp_t,
+ )
y_t = y_t * tWrW_t.load().to(cute.Float32)
if const_expr(mB is not None):
gB_t = cute.local_tile(
- qutils.domain_offset_i64((0, k_off), mB), tiler_mn_tile, (0, 0)
+ qutils.domain_offset_i64((0, k_off), mB),
+ tiler_mn_tile,
+ (0, 0),
)
tWgB_t = thr_copy_tile.partition_S(gB_t)
tWrB_t = cute.make_fragment_like(tWgB_t)
- copy_tiled(tWgB_t, tWrB_t, num_copy_elems=vecsize, is_async=False, pred=tXp_t)
+ copy_tiled(
+ tWgB_t,
+ tWrB_t,
+ num_copy_elems=vecsize,
+ is_async=False,
+ pred=tXp_t,
+ )
y_t = y_t + tWrB_t.load().to(cute.Float32)
- gO_t = cute.local_tile(qutils.domain_offset_i64((0, k_off), mO_i), tiler_mn_tile, (0, 0))
+ gO_t = cute.local_tile(
+ qutils.domain_offset_i64((0, k_off), mO_i),
+ tiler_mn_tile,
+ (0, 0),
+ )
tXgO_t = thr_copy_tile.partition_D(gO_t)
tXrO_t = cute.make_fragment_like(tXgO_t)
tXrO_t.store(y_t.to(tXrO_t.element_type))
if row_i < shape[0]:
- copy_tiled(tXrO_t, tXgO_t, num_copy_elems=vecsize, is_async=False, pred=tXp_t)
+ copy_tiled(
+ tXrO_t,
+ tXgO_t,
+ num_copy_elems=vecsize,
+ is_async=False,
+ pred=tXp_t,
+ )
return
# Single-stage path: one-row-per-CTA
tXgX_i = thr_copy.partition_S(gX_i)
- tXgRes_i = thr_copy.partition_S(gRes_i) if const_expr(mRes is not None) else None
+ tXgRes_i = (
+ thr_copy.partition_S(gRes_i) if const_expr(mRes is not None) else None
+ )
tXgO_i = thr_copy.partition_D(gO_i)
- tXgResO_i = thr_copy.partition_D(gResO_i) if const_expr(mResO is not None) else None
+ tXgResO_i = (
+ thr_copy.partition_D(gResO_i) if const_expr(mResO is not None) else None
+ )
# tXgRstd_i / tXcX_i / row_i prepared above
is_even_N_i = const_expr(shape[1] == tiler_mn[1] * cluster_n)
tXpX_i = (
- qutils.predicate_k(thr_copy.partition_S(cX_i), limit=limit_k) if not is_even_N_i else None
+ qutils.predicate_k(thr_copy.partition_S(cX_i), limit=limit_k)
+ if not is_even_N_i
+ else None
)
tXrX = cute.make_fragment_like(tXgX_i)
- tXrRes = cute.make_fragment_like(tXgRes_i) if const_expr(mRes is not None) else None
+ tXrRes = (
+ cute.make_fragment_like(tXgRes_i) if const_expr(mRes is not None) else None
+ )
if const_expr(self.direct_gmem):
if const_expr(not is_even_N_i):
tXrX.fill(0)
@@ -1729,7 +1888,9 @@ def _kernel_impl(
if row_i < shape[0]:
cute.copy(copy_atom, tXgX_i, thr_copy.partition_D(sX0), pred=tXpX_i)
if const_expr(mRes is not None):
- cute.copy(copy_atom, tXgRes_i, thr_copy.partition_D(sRes0), pred=tXpX_i)
+ cute.copy(
+ copy_atom, tXgRes_i, thr_copy.partition_D(sRes0), pred=tXpX_i
+ )
if const_expr(use_async):
cute.arch.cp_async_commit_group()
cute.arch.cp_async_wait_group(0)
@@ -1746,7 +1907,9 @@ def _kernel_impl(
tXrResO.store(x_red.to(tXrResO.element_type))
if row_i < shape[0]:
cute.copy(
- get_copy_atom_bw(tXrResO.element_type, num_copy_elems_X, is_async=False),
+ get_copy_atom_bw(
+ tXrResO.element_type, num_copy_elems_X, is_async=False
+ ),
tXrResO,
tXgResO_i,
pred=tXpX_i,
@@ -1775,7 +1938,9 @@ def _kernel_impl(
# pressure during the long-scoreboard reduction phase (helping occupancy
# when registers are the limiting factor).
if const_expr(mW is not None):
- gW = cute.local_tile(qutils.domain_offset_i64((0, n_off), mW), tiler_mn, (0, 0))
+ gW = cute.local_tile(
+ qutils.domain_offset_i64((0, n_off), mW), tiler_mn, (0, 0)
+ )
tXgW = thr_copy.partition_S(gW)
tXrW = cute.make_fragment_like(tXgW)
if const_expr(not is_even_N_wb):
@@ -1787,7 +1952,9 @@ def _kernel_impl(
pred=tXp_wb,
)
if const_expr(mB is not None):
- gB = cute.local_tile(qutils.domain_offset_i64((0, n_off), mB), tiler_mn, (0, 0))
+ gB = cute.local_tile(
+ qutils.domain_offset_i64((0, n_off), mB), tiler_mn, (0, 0)
+ )
tXgB = thr_copy.partition_S(gB)
tXrB = cute.make_fragment_like(tXgB)
if const_expr(not is_even_N_wb):
@@ -1818,6 +1985,7 @@ def _kernel_impl(
)
if _KERNEL_ACCEPTS_LAYOUT_ARGS:
+
@cute.kernel
def kernel(
self,
@@ -1853,6 +2021,7 @@ def kernel(
threads_per_row,
)
else:
+
@cute.kernel
def kernel(
self,
@@ -2015,7 +2184,10 @@ def _rmsnorm_forward_ptr(
if residual is not None:
residual_out = torch.empty_strided(
- residual.shape, residual.stride(), device=residual.device, dtype=residual.dtype
+ residual.shape,
+ residual.stride(),
+ device=residual.device,
+ dtype=residual.dtype,
)
if store_rstd:
rstd = torch.empty(M, device=x.device, dtype=torch.float32)
@@ -2082,10 +2254,19 @@ def _rmsnorm_forward_ptr_into(
if compiled is None:
op = RMSNormSM100(N, dtype, stage=stage)
ld_val = int(x.stride(0))
- ptr_x = rt.make_ptr(dtype, x.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
- ptr_out = rt.make_ptr(dtype, out.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ ptr_x = rt.make_ptr(
+ dtype, x.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16
+ )
+ ptr_out = rt.make_ptr(
+ dtype, out.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16
+ )
ptr_w = (
- rt.make_ptr(dtype, weight.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ rt.make_ptr(
+ dtype,
+ weight.data_ptr(),
+ mem_space=rt.AddressSpace.gmem,
+ assumed_align=16,
+ )
if has_weight
else None
)
@@ -2122,10 +2303,19 @@ def _rmsnorm_forward_ptr_into(
launcher.launch(x=x, weight=weight, out=out, M=M, N=N, ld=ld_val, eps=eps)
return
- ptr_x = rt.make_ptr(dtype, x.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
- ptr_out = rt.make_ptr(dtype, out.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ ptr_x = rt.make_ptr(
+ dtype, x.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16
+ )
+ ptr_out = rt.make_ptr(
+ dtype, out.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16
+ )
ptr_w = (
- rt.make_ptr(dtype, weight.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ rt.make_ptr(
+ dtype,
+ weight.data_ptr(),
+ mem_space=rt.AddressSpace.gmem,
+ assumed_align=16,
+ )
if has_weight
else None
)
@@ -2167,33 +2357,55 @@ def _rmsnorm_forward_ptr_into(
compiled = _PTR_COMPILE_CACHE.get(key)
if compiled is None:
op = RMSNormSM100(N, dtype, stage=stage)
- ptr_x = rt.make_ptr(dtype, x.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
- ptr_out = rt.make_ptr(dtype, out.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ ptr_x = rt.make_ptr(
+ dtype, x.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16
+ )
+ ptr_out = rt.make_ptr(
+ dtype, out.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16
+ )
ptr_res = (
- rt.make_ptr(dtype, residual.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ rt.make_ptr(
+ dtype,
+ residual.data_ptr(),
+ mem_space=rt.AddressSpace.gmem,
+ assumed_align=16,
+ )
if residual is not None
else None
)
ptr_res_out = (
rt.make_ptr(
- dtype, residual_out.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16
+ dtype,
+ residual_out.data_ptr(),
+ mem_space=rt.AddressSpace.gmem,
+ assumed_align=16,
)
if residual_out is not None
else None
)
ptr_w = (
- rt.make_ptr(dtype, weight.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ rt.make_ptr(
+ dtype,
+ weight.data_ptr(),
+ mem_space=rt.AddressSpace.gmem,
+ assumed_align=16,
+ )
if weight is not None
else None
)
ptr_b = (
- rt.make_ptr(dtype, bias.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ rt.make_ptr(
+ dtype, bias.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16
+ )
if bias is not None
else None
)
ptr_rstd = (
rt.make_ptr(
- cutlass.Float32, rstd.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=4
+ cutlass.Float32,
+ rstd.data_ptr(),
+ mem_space=rt.AddressSpace.gmem,
+ assumed_align=4,
)
if rstd is not None
else None
@@ -2216,30 +2428,50 @@ def _rmsnorm_forward_ptr_into(
Float32(eps),
)
_PTR_COMPILE_CACHE[key] = compiled
- ptr_x = rt.make_ptr(dtype, x.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
- ptr_out = rt.make_ptr(dtype, out.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ ptr_x = rt.make_ptr(
+ dtype, x.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16
+ )
+ ptr_out = rt.make_ptr(
+ dtype, out.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16
+ )
ptr_res = (
- rt.make_ptr(dtype, residual.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ rt.make_ptr(
+ dtype, residual.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16
+ )
if residual is not None
else None
)
ptr_res_out = (
- rt.make_ptr(dtype, residual_out.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ rt.make_ptr(
+ dtype,
+ residual_out.data_ptr(),
+ mem_space=rt.AddressSpace.gmem,
+ assumed_align=16,
+ )
if residual_out is not None
else None
)
ptr_w = (
- rt.make_ptr(dtype, weight.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ rt.make_ptr(
+ dtype, weight.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16
+ )
if weight is not None
else None
)
ptr_b = (
- rt.make_ptr(dtype, bias.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16)
+ rt.make_ptr(
+ dtype, bias.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=16
+ )
if bias is not None
else None
)
ptr_rstd = (
- rt.make_ptr(cutlass.Float32, rstd.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=4)
+ rt.make_ptr(
+ cutlass.Float32,
+ rstd.data_ptr(),
+ mem_space=rt.AddressSpace.gmem,
+ assumed_align=4,
+ )
if rstd is not None
else None
)
@@ -2296,7 +2528,9 @@ def _fused_add_rmsnorm_forward_ptr_inplace(
# benchmark other models/shapes, you can override it with:
# - OINK_RMSNORM_DIRECT_GMEM=0 (force staging/cp.async path)
# - OINK_RMSNORM_DIRECT_GMEM=1 (force direct-gmem path)
- direct_gmem = _direct_gmem_from_policy(default=bool(dtype.width == 16 and N == 7168))
+ direct_gmem = _direct_gmem_from_policy(
+ default=bool(dtype.width == 16 and N == 7168)
+ )
use_async = not direct_gmem
tpr_override: Optional[int] = None
nt_override: Optional[int] = None
@@ -2340,13 +2574,9 @@ def _fused_add_rmsnorm_forward_ptr_inplace(
tpr_override = 256
nt_override = 256
-
can_use_256 = bool(
direct_gmem
- and (
- direct_gmem_max_copy_bits is None
- or direct_gmem_max_copy_bits >= 256
- )
+ and (direct_gmem_max_copy_bits is None or direct_gmem_max_copy_bits >= 256)
and dtype.width == 16
and (x.data_ptr() % 32) == 0
and (residual.data_ptr() % 32) == 0
@@ -2395,13 +2625,22 @@ def _fused_add_rmsnorm_forward_ptr_inplace(
if cluster_n_override is not None:
op._cluster_n_override = cluster_n_override # type: ignore[attr-defined]
ptr_x = rt.make_ptr(
- dtype, x.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=assumed_align
+ dtype,
+ x.data_ptr(),
+ mem_space=rt.AddressSpace.gmem,
+ assumed_align=assumed_align,
)
ptr_res = rt.make_ptr(
- dtype, residual.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=assumed_align
+ dtype,
+ residual.data_ptr(),
+ mem_space=rt.AddressSpace.gmem,
+ assumed_align=assumed_align,
)
ptr_w = rt.make_ptr(
- dtype, weight.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=assumed_align
+ dtype,
+ weight.data_ptr(),
+ mem_space=rt.AddressSpace.gmem,
+ assumed_align=assumed_align,
)
stream = cuda.CUstream(stream_handle)
ld_x = Int32(int(x.stride(0)))
@@ -2444,12 +2683,20 @@ def _fused_add_rmsnorm_forward_ptr_inplace(
# Fast-launch is disabled/unavailable (or CuTeDSL internals changed). Fall back
# to calling the compiled function directly.
- ptr_x = rt.make_ptr(dtype, x.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=assumed_align)
+ ptr_x = rt.make_ptr(
+ dtype, x.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=assumed_align
+ )
ptr_res = rt.make_ptr(
- dtype, residual.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=assumed_align
+ dtype,
+ residual.data_ptr(),
+ mem_space=rt.AddressSpace.gmem,
+ assumed_align=assumed_align,
)
ptr_w = rt.make_ptr(
- dtype, weight.data_ptr(), mem_space=rt.AddressSpace.gmem, assumed_align=assumed_align
+ dtype,
+ weight.data_ptr(),
+ mem_space=rt.AddressSpace.gmem,
+ assumed_align=assumed_align,
)
stream = cuda.CUstream(stream_handle)
ld_x = Int32(int(x.stride(0)))
@@ -2486,7 +2733,12 @@ def rmsnorm_forward(
rms2 = None # type: ignore[assignment]
if rms2 is not None:
y, rstd, residual_out = rms2.rmsnorm_forward_with_stage2(
- x, weight=weight, bias=bias, residual=residual, eps=eps, store_rstd=store_rstd
+ x,
+ weight=weight,
+ bias=bias,
+ residual=residual,
+ eps=eps,
+ store_rstd=store_rstd,
)
# Preserve stride contracts for torch.compile consistency, even
# when using the optional stage-2 implementation.
@@ -2519,7 +2771,9 @@ def rmsnorm_forward(
# Preserve the input stride contract even on the fallback path so
# torch.compile sees a consistent output layout across all branches.
if y.stride() != x.stride():
- y_strided = torch.empty_strided(x.shape, x.stride(), device=x.device, dtype=x.dtype)
+ y_strided = torch.empty_strided(
+ x.shape, x.stride(), device=x.device, dtype=x.dtype
+ )
y_strided.copy_(y)
y = y_strided
rstd = None
From 4c9a826cb2b1a3d48fbffe156f10943b41ec8f80 Mon Sep 17 00:00:00 2001
From: Laura Wang <3700467+Laurawly@users.noreply.github.com>
Date: Wed, 21 Jan 2026 11:06:37 -0800
Subject: [PATCH 4/8] oink: SM100 suite refresh (strict parity + quack-style
benches)
- Switch correctness gate to PyTorch ref + record err stats\n- Tighten Softmax/LayerNorm tolerances (Quack-like)\n- Quack-style benchmark suite layout + SVG plots\n- Packaging/README polish for publishability
---
oink/README.md | 94 +-
oink/benchmarks/README.md | 146 +
oink/benchmarks/benchmark/bench_utils.py | 255 ++
.../benchmark_cross_entropy_sm100.py | 426 +++
.../benchmark_fused_add_rmsnorm_sm100.py | 296 ++
.../benchmark/benchmark_hbm_roofline_sm100.py | 226 ++
.../benchmark/benchmark_layernorm_sm100.py | 393 +++
.../benchmark/benchmark_rmsnorm_bwd_sm100.py | 434 +++
.../benchmark/benchmark_rmsnorm_sm100.py | 337 ++
.../benchmark/benchmark_softmax_sm100.py | 292 ++
.../media/sm100_bf16_oink_vs_quack.svg | 2259 +++++++++++++
.../media/sm100_bf16_oink_vs_quack_dsv3.svg | 2600 +++++++++++++++
.../sm100_bf16_oink_vs_quack_dsv3_all.svg | 2936 ++++++++++++++++
..._bf16_oink_vs_quack_dsv3_cross_entropy.svg | 1687 ++++++++++
...bf16_oink_vs_quack_dsv3_with_layernorm.svg | 2720 +++++++++++++++
...m100_bf16_oink_vs_quack_with_layernorm.svg | 2580 ++++++++++++++
.../media/sm100_fp16_oink_vs_quack.svg | 2280 +++++++++++++
.../media/sm100_fp16_oink_vs_quack_dsv3.svg | 2621 +++++++++++++++
.../sm100_fp16_oink_vs_quack_dsv3_all.svg | 2957 +++++++++++++++++
..._fp16_oink_vs_quack_dsv3_cross_entropy.svg | 1708 ++++++++++
...fp16_oink_vs_quack_dsv3_with_layernorm.svg | 2741 +++++++++++++++
...m100_fp16_oink_vs_quack_with_layernorm.svg | 2601 +++++++++++++++
.../benchmarks/readme/plot_quack_style_svg.py | 431 +++
oink/benchmarks/readme/run_sm100_suite.py | 302 ++
oink/benchmarks/readme/summarize_results.py | 205 ++
oink/pyproject.toml | 24 +-
oink/src/kernelagent_oink/__init__.py | 28 +-
.../blackwell/cross_entropy.py | 1209 +++++++
.../kernelagent_oink/blackwell/layernorm.py | 1368 ++++++++
.../kernelagent_oink/blackwell/lite_quack.py | 1001 +++++-
.../src/kernelagent_oink/blackwell/rmsnorm.py | 467 ++-
.../blackwell/rmsnorm_with_stage2.py | 805 +++++
.../src/kernelagent_oink/blackwell/softmax.py | 749 +++++
33 files changed, 39026 insertions(+), 152 deletions(-)
create mode 100644 oink/benchmarks/README.md
create mode 100644 oink/benchmarks/benchmark/bench_utils.py
create mode 100644 oink/benchmarks/benchmark/benchmark_cross_entropy_sm100.py
create mode 100644 oink/benchmarks/benchmark/benchmark_fused_add_rmsnorm_sm100.py
create mode 100644 oink/benchmarks/benchmark/benchmark_hbm_roofline_sm100.py
create mode 100644 oink/benchmarks/benchmark/benchmark_layernorm_sm100.py
create mode 100644 oink/benchmarks/benchmark/benchmark_rmsnorm_bwd_sm100.py
create mode 100644 oink/benchmarks/benchmark/benchmark_rmsnorm_sm100.py
create mode 100644 oink/benchmarks/benchmark/benchmark_softmax_sm100.py
create mode 100644 oink/benchmarks/media/sm100_bf16_oink_vs_quack.svg
create mode 100644 oink/benchmarks/media/sm100_bf16_oink_vs_quack_dsv3.svg
create mode 100644 oink/benchmarks/media/sm100_bf16_oink_vs_quack_dsv3_all.svg
create mode 100644 oink/benchmarks/media/sm100_bf16_oink_vs_quack_dsv3_cross_entropy.svg
create mode 100644 oink/benchmarks/media/sm100_bf16_oink_vs_quack_dsv3_with_layernorm.svg
create mode 100644 oink/benchmarks/media/sm100_bf16_oink_vs_quack_with_layernorm.svg
create mode 100644 oink/benchmarks/media/sm100_fp16_oink_vs_quack.svg
create mode 100644 oink/benchmarks/media/sm100_fp16_oink_vs_quack_dsv3.svg
create mode 100644 oink/benchmarks/media/sm100_fp16_oink_vs_quack_dsv3_all.svg
create mode 100644 oink/benchmarks/media/sm100_fp16_oink_vs_quack_dsv3_cross_entropy.svg
create mode 100644 oink/benchmarks/media/sm100_fp16_oink_vs_quack_dsv3_with_layernorm.svg
create mode 100644 oink/benchmarks/media/sm100_fp16_oink_vs_quack_with_layernorm.svg
create mode 100644 oink/benchmarks/readme/plot_quack_style_svg.py
create mode 100644 oink/benchmarks/readme/run_sm100_suite.py
create mode 100644 oink/benchmarks/readme/summarize_results.py
create mode 100644 oink/src/kernelagent_oink/blackwell/cross_entropy.py
create mode 100644 oink/src/kernelagent_oink/blackwell/layernorm.py
create mode 100644 oink/src/kernelagent_oink/blackwell/rmsnorm_with_stage2.py
create mode 100644 oink/src/kernelagent_oink/blackwell/softmax.py
diff --git a/oink/README.md b/oink/README.md
index 427f69f..aeb0c09 100644
--- a/oink/README.md
+++ b/oink/README.md
@@ -1,13 +1,33 @@
-# KernelAgent Oink (vLLM plugin)
+# KernelAgent-Oink
-This subproject provides an **out-of-tree vLLM plugin** that registers
-`torch.library.custom_op` entrypoints under the `oink::` namespace:
+KernelAgent-Oink is a small **CuTeDSL (CUTLASS DSL) kernel library** for
+**NVIDIA Blackwell (SM100 / GB200 / B200-class)**, bundled as a lightweight
+Python package that can be used standalone or as a **vLLM general plugin**.
-- `torch.ops.oink.rmsnorm`
-- `torch.ops.oink.fused_add_rms_norm`
+At the moment, the vLLM integration exposes the following `torch.library.custom_op`
+entrypoints under the `oink::` namespace:
-The implementation is backed by a CuTeDSL (CUTLASS) RMSNorm kernel tuned for
-**NVIDIA Blackwell (SM100)**.
+- `torch.ops.oink.rmsnorm(x, weight, eps) -> Tensor`
+- `torch.ops.oink.fused_add_rms_norm(x, residual, weight, eps) -> None` (in-place)
+
+The package also includes additional SM100 kernels used by the benchmark suite:
+LayerNorm, Softmax (fwd+bwd), and CrossEntropy (fwd+bwd).
+
+## Requirements
+
+- GPU: **SM100** for the fast CuTeDSL paths. On other GPUs, Oink falls back to
+ reference PyTorch implementations for correctness.
+- Python dependencies:
+ - `nvidia-cutlass-dsl` (CuTeDSL)
+ - `cuda-python`
+ - `torch` (provided by your environment / vLLM)
+
+Recommended env vars:
+
+```bash
+export CUTE_DSL_ARCH=sm_100a
+export PYTORCH_ALLOC_CONF=expandable_segments:True
+```
## Install (editable)
@@ -17,22 +37,23 @@ From the `KernelAgent` repo root:
pip install -e ./oink
```
-This plugin requires the CuTeDSL stack:
+For running the in-repo benchmark suite / plots:
```bash
-pip install nvidia-cutlass-dsl cuda-python
+pip install -e "./oink[bench]"
```
-## Use with vLLM
+## Usage
+
+### vLLM (general plugin)
-1. Enable the vLLM integration:
+1) Enable the plugin:
```bash
export VLLM_USE_OINK_RMSNORM=1
```
-2. Ensure vLLM keeps `rms_norm` as a custom op when using `torch.compile` /
-CUDA graphs. In Python:
+2) Ensure vLLM keeps `rms_norm` as a custom op when using `torch.compile` / CUDA graphs:
```python
from vllm import LLM
@@ -45,13 +66,44 @@ llm = LLM(
)
```
-Without `+rms_norm`, Inductor may fuse RMSNorm into larger Triton kernels and
-neither vLLM's CUDA RMSNorm nor Oink will run.
+Without `+rms_norm`, Inductor may fuse RMSNorm into larger kernels and neither
+vLLM’s CUDA RMSNorm nor Oink will run.
+
+### Direct PyTorch usage (manual op registration)
+
+For standalone use (outside vLLM), register the custom ops once:
+
+```python
+import kernelagent_oink
+import torch
+
+kernelagent_oink.register(force=True)
+
+x = torch.randn(1024, 4096, device="cuda", dtype=torch.bfloat16)
+w = torch.randn(4096, device="cuda", dtype=torch.bfloat16)
+y = torch.ops.oink.rmsnorm(x, w, 1e-6)
+```
+
+## Benchmarks
+
+The repo includes a Quack-style benchmark suite (tables + SVG plots) to compare
+Oink against Quack on SM100 and to reproduce the reported speedups.
+
+- How to run + methodology: `oink/benchmarks/README.md`
+- Pre-generated plots: `oink/benchmarks/media/`
+
+
+
+
+
+