mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
feat(cutile): cuTile backends for RMSNorm, LayerNorm, Softmax, JSD, FusedLinearJSD (#1420)
# feat(cutile): cuTile backends for RMSNorm, LayerNorm, Softmax, JSD,
FusedLinearJSD
**Head:** `arde171:arde/liger-E-cutile` → **Base:** `main`. Final PR of
the multi-backend series; **#1419 (CuTeDSL CE) is now merged**, so this
diff is **cuTile-only** — no rebase caveat anymore.
## What this adds
Registers a **cuTile** backend family for five ops, gated strictly to
Blackwell. All impls declare **`min_cc=(10, 0)`** → only reachable on
sm_100+ (B200/B300); Hopper/Ampere transparently keep using
Triton/CuTeDSL.
| file | + |
|---|---|
| `ops/backends/_cutile/rms_norm.py` | 1221 |
| `ops/backends/_cutile/layer_norm.py` | 626 |
| `ops/backends/_cutile/softmax.py` | 448 |
| `ops/backends/_cutile/jsd.py` | 319 |
| `ops/backends/_cutile/fused_linear_jsd.py` | 64 |
| `ops/backends/_cutile/__init__.py` | 11 |
| `functional.py` | 6 (append `_cutile.*` to 5 op declarations) |
| `ops/fused_linear_jsd.py` | 19 (`_max_capability(device=None)`) |
| `testing/_op_helpers.py` | 5 (`strict_grad` flag) |
| `test/backends/test_cutile_select_mode.py` | 83 |
| `test/cutile/test_cutile_backends_parity.py` | 799 (113-case parity
matrix) |
| `test/ops/{test_rms_norm,test_layer_norm}.py` | 44 |
No CuTeDSL file touched; purely additive registration behind a Blackwell
gate.
## Default-selection is benchmark-justified (nothing defaults to a
slower kernel)
Ranks are chosen so cuTile is **only ever opt-in** — it is never
auto-selected over a faster existing path. Measured median latency
(bf16, fwd+bwd) on **B200 (sm_100)**:
| op | rank | Triton | cuTile | CuTeDSL | auto-selected |
|---|---|---|---|---|---|
| `rms_norm` | 80 | 0.262 ms | 0.473 ms | **0.190 ms** | CuTeDSL |
| `layer_norm` | 80 | 0.271 ms | 0.623 ms | **0.219 ms** | CuTeDSL |
| `softmax` | 80 | **0.151 ms** | 0.188 ms | 0.156 ms | Triton |
| `jsd` | 25 | 0.802 ms | 0.661 ms | **0.583 ms** | CuTeDSL |
| `fused_linear_jsd` | 25 | — | — | — | CuTeDSL |
- **Norms → `preference_rank=80`** (above Triton's 50): cuTile is
0.44–0.55× of Triton on B200, so it must stay opt-in.
- **JSD / FusedLinearJSD → `preference_rank=25`**: both cuTile (1.21×)
and CuTeDSL (1.38×) beat Triton, but CuTeDSL wins, so CuTeDSL stays the
default and cuTile is opt-in below it.
## Testing (deterministic `-p no:randomly`)
Full `test/backends test/ops test/cutile` + the 113-case
CuTeDSL↔cuTile↔Torch parity matrix:
- **B300 (cc 10.3):** `1149 passed, 48 skipped, 0 failed` + `113/113`
parity.
- **B200 (cc 10.0):** `1149 passed, 48 skipped, 0 failed` + `113/113`
parity.
- **H200 (cc 9.0):** green (cuTile `min_cc=(10,0)` correctly skips on
Hopper).
- ruff check + format clean. Mixed-precision affine (dW/dB in their own
dtypes), device-safe launches (all kernels use the input tensor's
device), and `strict_grad` elementwise dW/dB assertions added.
Co-authored-by: Copilot <[email protected]>
---------
Co-authored-by: Copilot <[email protected]>
This commit is contained in:
Arup De
@@ -37,6 +37,7 @@ declare_op_locations(
|
||||
"liger_kernel.ops.backends._triton.rms_norm",
|
||||
"liger_kernel.ops.backends._cutedsl.rms_norm",
|
||||
"liger_kernel.ops.backends._ascend.rms_norm",
|
||||
"liger_kernel.ops.backends._cutile.rms_norm",
|
||||
),
|
||||
)
|
||||
|
||||
@@ -45,6 +46,7 @@ declare_op_locations(
|
||||
(
|
||||
"liger_kernel.ops.backends._triton.layer_norm",
|
||||
"liger_kernel.ops.backends._cutedsl.layer_norm",
|
||||
"liger_kernel.ops.backends._cutile.layer_norm",
|
||||
),
|
||||
)
|
||||
|
||||
@@ -56,6 +58,7 @@ declare_op_locations(
|
||||
(
|
||||
"liger_kernel.ops.backends._triton.jsd",
|
||||
"liger_kernel.ops.backends._cutedsl.jsd",
|
||||
"liger_kernel.ops.backends._cutile.jsd",
|
||||
),
|
||||
)
|
||||
declare_op_locations(
|
||||
@@ -63,6 +66,7 @@ declare_op_locations(
|
||||
(
|
||||
"liger_kernel.ops.backends._triton.jsd",
|
||||
"liger_kernel.ops.backends._cutedsl.jsd",
|
||||
"liger_kernel.ops.backends._cutile.jsd",
|
||||
),
|
||||
)
|
||||
|
||||
@@ -72,6 +76,7 @@ declare_op_locations(
|
||||
(
|
||||
"liger_kernel.ops.backends._triton.softmax",
|
||||
"liger_kernel.ops.backends._cutedsl.softmax",
|
||||
"liger_kernel.ops.backends._cutile.softmax",
|
||||
),
|
||||
)
|
||||
|
||||
@@ -155,6 +160,7 @@ declare_op_locations(
|
||||
(
|
||||
"liger_kernel.ops.backends._triton.fused_linear_jsd",
|
||||
"liger_kernel.ops.backends._cutedsl.fused_linear_jsd",
|
||||
"liger_kernel.ops.backends._cutile.fused_linear_jsd",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@@ -0,0 +1,11 @@
|
||||
"""cuTile (NVIDIA cuda.tile) backend implementations for Liger-Kernel ops.
|
||||
|
||||
This package holds per-op cuTile kernel implementations registered with the
|
||||
multi-DSL dispatcher via :func:`liger_kernel.backends.register_op`. Importing
|
||||
modules in this package triggers their registration as a side effect; the
|
||||
modules themselves import ``cuda.tile``, so callers without that package will
|
||||
see :class:`ImportError`, which the dispatcher's discovery layer catches.
|
||||
|
||||
This file is intentionally empty (no ``register_vendor`` call) — vendor-level
|
||||
op swapping (NPU/XPU) happens in ``liger_kernel.ops.backends`` instead.
|
||||
"""
|
||||
@@ -0,0 +1,64 @@
|
||||
"""cuTile backend registration for ``fused_linear_jsd``."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
from liger_kernel.backends import Capability
|
||||
from liger_kernel.backends import register_op
|
||||
from liger_kernel.ops.backends._cutile.jsd import _cutile_compiler_available
|
||||
from liger_kernel.ops.fused_linear_jsd import LigerFusedLinearJSDFunction
|
||||
|
||||
_CUTILE_FLJSD_TOLERANCES = {
|
||||
torch.float16: {"atol_fwd": 5e-3, "atol_bwd": 5e-2, "rtol_fwd": 1e-3, "rtol_bwd": 1e-2},
|
||||
torch.bfloat16: {"atol_fwd": 2e-2, "atol_bwd": 1e-1, "rtol_fwd": 1e-2, "rtol_bwd": 2e-2},
|
||||
torch.float32: {"atol_fwd": 1e-5, "atol_bwd": 1e-4, "rtol_fwd": 1e-5, "rtol_bwd": 1e-4},
|
||||
}
|
||||
|
||||
|
||||
@register_op(
|
||||
"fused_linear_jsd",
|
||||
impl_name="nvidia-cutile",
|
||||
capability=Capability(
|
||||
min_cc=(10, 0),
|
||||
modules=["cuda.tile", "torch"],
|
||||
predicate=_cutile_compiler_available,
|
||||
),
|
||||
modes=("default",),
|
||||
default_mode="default",
|
||||
# Ranked above cutedsl fused_linear_jsd (rank 4) so cutedsl stays the auto-default; the
|
||||
# fused-linear path is GEMM-bound so backends are ~tied, and cuTile composes the cuTile
|
||||
# JSD primitive which is slower than cutedsl on B200 (see _cutile/jsd.py). Opt-in / fallback.
|
||||
preference_rank=25,
|
||||
tolerances=_CUTILE_FLJSD_TOLERANCES,
|
||||
notes="cuTile fused_linear_jsd for Blackwell (opt-in; cutedsl stays default). Inner JSD uses the cuTile primitive.",
|
||||
)
|
||||
def fused_linear_jsd_cutile(
|
||||
student_input: torch.Tensor,
|
||||
student_weight: torch.Tensor,
|
||||
teacher_input: torch.Tensor,
|
||||
teacher_weight: torch.Tensor,
|
||||
shift_labels: Optional[torch.Tensor] = None,
|
||||
jsd_beta: float = 0.5,
|
||||
ignore_index: int = -100,
|
||||
temperature: float = 1.0,
|
||||
*,
|
||||
mode: Optional[str] = None,
|
||||
) -> torch.Tensor:
|
||||
if mode not in (None, "default"):
|
||||
raise ValueError(f"cuTile fused_linear_jsd has only mode='default'; got mode={mode!r}.")
|
||||
return LigerFusedLinearJSDFunction.apply(
|
||||
student_input,
|
||||
student_weight,
|
||||
teacher_input,
|
||||
teacher_weight,
|
||||
shift_labels,
|
||||
jsd_beta,
|
||||
ignore_index,
|
||||
temperature,
|
||||
None,
|
||||
"nvidia-cutile",
|
||||
mode,
|
||||
)
|
||||
@@ -0,0 +1,319 @@
|
||||
"""cuTile (NVIDIA cuda.tile) backend for ``jsd``.
|
||||
|
||||
This file ports NVIDIA's cuTile JSD kernel from
|
||||
`linkedin/Liger-Kernel#1228 <https://github.com/linkedin/Liger-Kernel/pull/1228>`_
|
||||
(SPDX-License-Identifier: MIT — NVIDIA CORPORATION & AFFILIATES, 2025) and
|
||||
wraps it in PR #4's per-(op, impl) registry so it coexists with the Triton
|
||||
implementation under ``Capability``-based gating.
|
||||
|
||||
Two ops registered here:
|
||||
|
||||
1. ``jsd`` (autograd-aware) — :class:`_LigerJSDCuTileFunction` matches the
|
||||
API surface of :class:`liger_kernel.ops.jsd.LigerJSDFunction`.
|
||||
|
||||
2. ``jsd_loss_and_grad`` (non-autograd primitive) — returns per-row loss + dx
|
||||
directly, used by ``fused_linear_jsd`` so it picks up cuTile when the
|
||||
user sets ``impl="nvidia-cutile"``.
|
||||
|
||||
Capability: requires ``cuda.tile`` runtime + ``tileiras`` compiler reachable
|
||||
from the current process. We probe both at registration time so partial
|
||||
installs (runtime present, compiler missing) gate out gracefully.
|
||||
|
||||
References
|
||||
----------
|
||||
- NVIDIA TileGym Liger suite — ``src/tilegym/suites/liger/cutile/jsd.py``
|
||||
- In-repo reproduction on f2e7013f B200 sm_100 (torch 2.12, tileiras v13.2):
|
||||
**8.12× forward vs Triton** at llama_3_8b shape (BT=4096, V=128256, bf16).
|
||||
"""
|
||||
|
||||
# NOTE: same anti-PEP-563 caveat as the cuTile rms_norm / layer_norm siblings.
|
||||
# ``@ct.kernel`` introspects ``__annotations__`` to discover ``Constant``
|
||||
# parameters; ``from __future__ import annotations`` stringifies them and
|
||||
# silently breaks specialization across launches.
|
||||
|
||||
import math
|
||||
|
||||
from typing import Optional
|
||||
from typing import Tuple
|
||||
|
||||
import cuda.tile as ct
|
||||
import torch
|
||||
|
||||
from liger_kernel.backends import Capability
|
||||
from liger_kernel.backends import register_op
|
||||
from liger_kernel.ops._nvidia_shared import cutile_compiler_available as _cutile_compiler_available
|
||||
from liger_kernel.ops._nvidia_shared import next_pow2 as _next_pow2
|
||||
from liger_kernel.ops._nvidia_shared import to_local_if_dtensor as _to_local_if_dtensor
|
||||
from liger_kernel.ops.utils import ensure_contiguous
|
||||
|
||||
_ConstFloat = ct.Constant[float]
|
||||
_ConstInt = ct.Constant[int]
|
||||
_JSD_BLOCK_SIZE = 4096
|
||||
|
||||
|
||||
@ct.kernel(occupancy=ct.ByTarget(sm_100=4))
|
||||
def _jsd_kernel_ct(
|
||||
x, # (BT, V) log Q (student)
|
||||
y, # (BT, V) log P (teacher)
|
||||
loss, # (BT, V) float32 loss accumulator
|
||||
dx, # (BT, V) gradient output
|
||||
label, # (BT,) label tensor, or dummy when HAS_LABEL=0
|
||||
beta: _ConstFloat,
|
||||
inv_n_non_ignore, # runtime fp32 scalar tensor
|
||||
ignore_index: _ConstInt,
|
||||
n_cols: _ConstInt,
|
||||
BLOCK_SIZE: _ConstInt,
|
||||
HAS_LABEL: _ConstInt,
|
||||
):
|
||||
"""One row per program. Fused loss + dx; fp32 inside, dtype-cast on store."""
|
||||
row_idx = ct.bid(0)
|
||||
scale = ct.load(inv_n_non_ignore, 0, shape=())
|
||||
|
||||
if HAS_LABEL:
|
||||
lbl = ct.load(label, row_idx, shape=())
|
||||
if lbl == ignore_index:
|
||||
# Zero this row's gradient and return — no loss contribution.
|
||||
num_chunks_early = (n_cols + BLOCK_SIZE - 1) // BLOCK_SIZE
|
||||
for ci in range(num_chunks_early):
|
||||
col_indices = ct.arange(BLOCK_SIZE, dtype=ct.int32) + ci * BLOCK_SIZE
|
||||
ct.scatter(
|
||||
dx,
|
||||
(row_idx, col_indices),
|
||||
ct.full((BLOCK_SIZE,), 0.0, dtype=dx.dtype),
|
||||
check_bounds=True,
|
||||
)
|
||||
return
|
||||
|
||||
num_chunks = (n_cols + BLOCK_SIZE - 1) // BLOCK_SIZE
|
||||
for chunk_idx in range(num_chunks):
|
||||
col_indices = ct.arange(BLOCK_SIZE, dtype=ct.int32) + chunk_idx * BLOCK_SIZE
|
||||
|
||||
x_tile = ct.gather(x, (row_idx, col_indices), check_bounds=True, padding_value=-math.inf)
|
||||
y_tile = ct.gather(y, (row_idx, col_indices), check_bounds=True, padding_value=-math.inf)
|
||||
|
||||
x_f32 = ct.astype(x_tile, ct.float32)
|
||||
y_f32 = ct.astype(y_tile, ct.float32)
|
||||
|
||||
loss_tile = ct.full((BLOCK_SIZE,), 0.0, dtype=ct.float32)
|
||||
dx_tile = ct.full((BLOCK_SIZE,), 0.0, dtype=ct.float32)
|
||||
|
||||
if beta == 0.0:
|
||||
# Forward KL: P || M, M -> P
|
||||
y_max = ct.max(y_f32, 0, keepdims=True)
|
||||
y_prob = ct.exp(y_f32 - y_max) * ct.exp(y_max)
|
||||
loss_tile = y_prob * (y_f32 - x_f32)
|
||||
dx_tile = -y_prob
|
||||
elif beta == 1.0:
|
||||
# Reverse KL: Q || M, M -> Q
|
||||
x_max = ct.max(x_f32, 0, keepdims=True)
|
||||
x_prob = ct.exp(x_f32 - x_max) * ct.exp(x_max)
|
||||
loss_tile = x_prob * (x_f32 - y_f32)
|
||||
dx_tile = loss_tile + x_prob
|
||||
else:
|
||||
# Generalized JSD with mixing coefficient beta.
|
||||
x_max = ct.max(x_f32, 0, keepdims=True)
|
||||
y_max = ct.max(y_f32, 0, keepdims=True)
|
||||
max_val = ct.maximum(x_max, y_max)
|
||||
exp_max = ct.exp(max_val)
|
||||
q_prob = ct.exp(x_f32 - max_val) * exp_max
|
||||
p_prob = ct.exp(y_f32 - max_val) * exp_max
|
||||
beta_p = beta * p_prob
|
||||
one_minus_beta_q = (1.0 - beta) * q_prob
|
||||
m_prob = beta_p + one_minus_beta_q
|
||||
log_m = ct.log(m_prob)
|
||||
loss_tile = beta_p * y_f32 + one_minus_beta_q * x_f32 - m_prob * log_m
|
||||
dx_tile = one_minus_beta_q * (x_f32 - log_m)
|
||||
|
||||
loss_tile = loss_tile * scale
|
||||
dx_tile = dx_tile * scale
|
||||
|
||||
ct.scatter(loss, (row_idx, col_indices), loss_tile, check_bounds=True)
|
||||
ct.scatter(dx, (row_idx, col_indices), ct.astype(dx_tile, dx.dtype), check_bounds=True)
|
||||
|
||||
|
||||
def _launch_jsd_kernel(
|
||||
student_prob: torch.Tensor,
|
||||
teacher_prob: torch.Tensor,
|
||||
shift_labels: Optional[torch.Tensor],
|
||||
beta: float,
|
||||
ignore_index: int,
|
||||
n_non_ignore: float,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Allocate ``loss`` + ``dx`` and launch :func:`_jsd_kernel_ct`.
|
||||
|
||||
Returns ``(loss, dx)`` with shapes ``(BT, V)`` (fp32) and ``(BT, V)`` (input dtype).
|
||||
"""
|
||||
BT, V = student_prob.shape
|
||||
BLOCK_SIZE = min(_JSD_BLOCK_SIZE, _next_pow2(V))
|
||||
has_label = shift_labels is not None
|
||||
|
||||
loss = torch.zeros((BT, V), dtype=torch.float32, device=student_prob.device)
|
||||
dx = torch.empty_like(student_prob)
|
||||
|
||||
inv_n_non_ignore = torch.tensor(
|
||||
[0.0 if n_non_ignore == 0 else (1.0 / n_non_ignore)],
|
||||
dtype=torch.float32,
|
||||
device=student_prob.device,
|
||||
)
|
||||
label_arg = shift_labels if has_label else torch.empty(1, dtype=torch.int64, device=student_prob.device)
|
||||
|
||||
# Device-safe launch: dispatch selects this backend from the input tensor's
|
||||
# device, so tie the launch (and stream) to that device rather than the
|
||||
# process-current one. Mirrors ops/cutile/ops/fused_linear_cross_entropy.py.
|
||||
device = student_prob.device
|
||||
with torch.cuda.device(device):
|
||||
ct.launch(
|
||||
torch.cuda.current_stream(device),
|
||||
(BT, 1, 1),
|
||||
_jsd_kernel_ct,
|
||||
(
|
||||
student_prob,
|
||||
teacher_prob,
|
||||
loss,
|
||||
dx,
|
||||
label_arg,
|
||||
float(beta),
|
||||
inv_n_non_ignore,
|
||||
int(ignore_index),
|
||||
int(V),
|
||||
int(BLOCK_SIZE),
|
||||
int(has_label),
|
||||
),
|
||||
)
|
||||
return loss, dx
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Autograd-aware wrapper — matches LigerJSDFunction's API.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _LigerJSDCuTileFunction(torch.autograd.Function):
|
||||
"""cuTile autograd wrapper for generalized Jensen-Shannon Divergence."""
|
||||
|
||||
@staticmethod
|
||||
@ensure_contiguous
|
||||
def forward(
|
||||
ctx,
|
||||
_input: torch.Tensor,
|
||||
target: torch.Tensor,
|
||||
shift_labels: Optional[torch.Tensor],
|
||||
beta: float,
|
||||
ignore_index: int,
|
||||
) -> torch.Tensor:
|
||||
_input = _to_local_if_dtensor(_input)
|
||||
target = _to_local_if_dtensor(target)
|
||||
if shift_labels is not None:
|
||||
shift_labels = _to_local_if_dtensor(shift_labels)
|
||||
|
||||
has_label = False
|
||||
if shift_labels is not None:
|
||||
assert shift_labels.shape == (_input.shape[0],), (
|
||||
f"shift_labels must have shape (BT,). Got: {shift_labels.shape}"
|
||||
)
|
||||
shift_labels = shift_labels.contiguous()
|
||||
has_label = True
|
||||
|
||||
BT = _input.shape[0]
|
||||
if has_label:
|
||||
n_non_ignore = int((shift_labels != ignore_index).sum().item())
|
||||
else:
|
||||
n_non_ignore = BT
|
||||
|
||||
if n_non_ignore == 0:
|
||||
ctx.save_for_backward(torch.zeros_like(_input))
|
||||
return torch.tensor(0.0, device=_input.device, dtype=_input.dtype)
|
||||
|
||||
loss, dx = _launch_jsd_kernel(_input, target, shift_labels, beta, ignore_index, float(n_non_ignore))
|
||||
ctx.save_for_backward(dx)
|
||||
return torch.sum(loss).to(_input.dtype)
|
||||
|
||||
@staticmethod
|
||||
@ensure_contiguous
|
||||
def backward(ctx, grad_output: torch.Tensor):
|
||||
(dx,) = ctx.saved_tensors
|
||||
# If grad_output is the scalar 1.0 (typical when JSD is the last loss),
|
||||
# skip the elementwise multiply entirely.
|
||||
if not torch.equal(grad_output, torch.tensor(1.0, device=grad_output.device)):
|
||||
dx = grad_output * dx
|
||||
return (dx, None, None, None, None)
|
||||
|
||||
|
||||
_CUTILE_JSD_TOLERANCES = {
|
||||
torch.float16: {"atol_fwd": 5e-3, "atol_bwd": 5e-3, "rtol_fwd": 1e-3, "rtol_bwd": 1e-3},
|
||||
torch.bfloat16: {"atol_fwd": 2e-2, "atol_bwd": 2e-2, "rtol_fwd": 1e-2, "rtol_bwd": 1e-2},
|
||||
torch.float32: {"atol_fwd": 1e-5, "atol_bwd": 1e-5, "rtol_fwd": 1e-5, "rtol_bwd": 1e-5},
|
||||
}
|
||||
|
||||
|
||||
@register_op(
|
||||
"jsd",
|
||||
impl_name="nvidia-cutile",
|
||||
capability=Capability(
|
||||
min_cc=(10, 0),
|
||||
modules=["cuda.tile", "torch"],
|
||||
predicate=_cutile_compiler_available,
|
||||
),
|
||||
modes=("default",),
|
||||
default_mode="default",
|
||||
# Ranked ABOVE cutedsl JSD (rank 20 on sm_100 / 2 on sm_103) so cutedsl stays the
|
||||
# auto-default where it is faster. Measured B200 (fwd+bwd, V=128256, rel-vs-triton=0):
|
||||
# cutedsl 3.04ms < cuTile 4.09ms < Triton 6.06ms @ BT=8192 (cuTile ~1.5x vs Triton but
|
||||
# ~35% slower than cutedsl). cuTile stays below Triton (50) so it still wins as a fallback
|
||||
# when cutedsl is unavailable, and remains selectable via explicit impl="nvidia-cutile".
|
||||
preference_rank=25,
|
||||
tolerances=_CUTILE_JSD_TOLERANCES,
|
||||
notes=(
|
||||
"cuTile JSD ported from NVIDIA TileGym + PR #1228 (MIT). B200 fwd+bwd ~1.5x vs Triton "
|
||||
"but slower than cutedsl JSD, so ranked below cutedsl (opt-in / fallback above Triton)."
|
||||
),
|
||||
)
|
||||
def jsd_cutile(
|
||||
_input: torch.Tensor,
|
||||
target: torch.Tensor,
|
||||
shift_labels: Optional[torch.Tensor] = None,
|
||||
beta: float = 0.5,
|
||||
ignore_index: int = -100,
|
||||
*,
|
||||
mode: Optional[str] = None,
|
||||
) -> torch.Tensor:
|
||||
if mode not in (None, "default"):
|
||||
raise ValueError(f"cuTile jsd has only mode='default'; got {mode!r}.")
|
||||
return _LigerJSDCuTileFunction.apply(_input, target, shift_labels, beta, ignore_index)
|
||||
|
||||
|
||||
@register_op(
|
||||
"jsd_loss_and_grad",
|
||||
impl_name="nvidia-cutile",
|
||||
capability=Capability(
|
||||
min_cc=(10, 0),
|
||||
modules=["cuda.tile", "torch"],
|
||||
predicate=_cutile_compiler_available,
|
||||
),
|
||||
modes=("default",),
|
||||
default_mode="default",
|
||||
# Inner primitive the fused-linear JSD training path composes. Ranked above cutedsl's
|
||||
# jsd_loss_and_grad (rank 20 on sm_100 / 2 on sm_103) so the faster cutedsl inner kernel
|
||||
# stays the auto-default there (cutedsl > cuTile on B200; see _cutile/jsd.py notes).
|
||||
preference_rank=25,
|
||||
notes="cuTile per-chunk JSD primitive for fused_linear_jsd (opt-in; cutedsl inner stays default).",
|
||||
)
|
||||
def jsd_loss_and_grad_cutile(
|
||||
student_prob: torch.Tensor,
|
||||
teacher_prob: torch.Tensor,
|
||||
shift_labels: Optional[torch.Tensor],
|
||||
beta: float,
|
||||
ignore_index: int,
|
||||
n_non_ignore: float,
|
||||
*,
|
||||
mode: Optional[str] = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Per-chunk JSD primitive: returns ``(loss, dx)`` with shapes ``(BT, V)``.
|
||||
|
||||
Unlike the Triton sibling, ``dx`` is a **freshly allocated tensor** (not
|
||||
an in-place overwrite of ``student_prob``). Callers must rebind:
|
||||
``student_prob = dx`` if they want the old in-place semantics.
|
||||
"""
|
||||
if mode not in (None, "default"):
|
||||
raise ValueError(f"jsd_loss_and_grad_cutile: only mode='default'; got {mode!r}")
|
||||
return _launch_jsd_kernel(student_prob, teacher_prob, shift_labels, beta, ignore_index, n_non_ignore)
|
||||
@@ -0,0 +1,626 @@
|
||||
"""cuTile (NVIDIA cuda.tile) backend for ``layer_norm``.
|
||||
|
||||
This module implements LayerNorm with NVIDIA's cuTile DSL and registers it via
|
||||
``@register_op``. It exposes three forward kernel variants, selected by the
|
||||
``mode`` argument:
|
||||
|
||||
- ``"standard"`` — one row per program block; good for very wide rows.
|
||||
- ``"static_persistent"`` — NUM_SMS persistent blocks, each striding over
|
||||
rows one at a time. W and B are loaded once per block and stay register-
|
||||
resident across all rows. Best when M is much larger than NUM_SMS. Matches
|
||||
NVIDIA TileGym's reference for LayerNorm.
|
||||
- ``"multi_wave_cached"`` — single-tile-per-row with the weight and bias
|
||||
vectors loaded once per block; best for narrow rows.
|
||||
|
||||
A prior multi-row variant of ``static_persistent`` was removed because the
|
||||
row-tile / row-element index relationship in the persistent loop was wrong
|
||||
(consecutive blocks' row ranges overlapped by ``TILE_SIZE_M - 1`` rows,
|
||||
producing races and incorrect results when the runtime gate was lifted).
|
||||
|
||||
Backward is one persistent kernel that produces (dX, partial-dW, partial-dB).
|
||||
The partial dW / dB rows (one per SM) are reduced to the final ``dW`` / ``dB``
|
||||
on the host (matching the Triton kernel's ``_dW.sum(dim=0).to(W.dtype)``
|
||||
pattern).
|
||||
|
||||
Math (matches the Triton kernel and reference cuTile kernel exactly):
|
||||
|
||||
Forward (with cached Mean & RSTD for backward):
|
||||
mean = sum(X) / N
|
||||
var = sum((X - mean)^2) / N
|
||||
rstd = rsqrt(var + eps)
|
||||
Y = (X - mean) * rstd * W + B
|
||||
|
||||
Backward (per row, all reductions in fp32):
|
||||
x_hat = (X - mean) * rstd
|
||||
wdy = W * dY
|
||||
c1 = sum(x_hat * wdy) / N
|
||||
c2 = sum(wdy) / N
|
||||
dX = (wdy - x_hat * c1 - c2) * rstd
|
||||
dW += dY * x_hat (row-summed across SMs on host)
|
||||
dB += dY (row-summed across SMs on host)
|
||||
|
||||
References
|
||||
----------
|
||||
- cuTile LayerNorm reference: ``cutile-python/test/kernels/layer_norm.py`` —
|
||||
source of the ``@ct.kernel(occupancy=ct.ByTarget(sm_100=16))`` hint,
|
||||
``ct.load(..., allow_tma=False, latency=N)`` recipes, and the per-SM partial
|
||||
dW/dB pattern (the reference uses a lock-protected accumulator into a
|
||||
GROUP_SIZE_M buffer; we instead allocate one partial-row per SM and reduce
|
||||
on the host, mirroring the Triton kernel).
|
||||
- Triton reference: ``liger_kernel.ops.layer_norm`` — defines the
|
||||
computational semantics we must reproduce bit-for-bit (Mean + RSTD caches,
|
||||
the dX reduction with c1/c2 in fp32, and the dW/dB row-sum reduction on the
|
||||
host).
|
||||
- cuTile RMSNorm sibling: ``liger_kernel.ops.backends._cutile.rms_norm`` — the
|
||||
three-mode dispatch scaffolding, ``_select_mode`` heuristic, and the
|
||||
fp32-partial-dW host-reduction pattern come from this file. LayerNorm
|
||||
differs in (a) caching Mean *and* RSTD for backward, (b) accepting a real
|
||||
bias tensor (functional layer substitutes zeros when bias is None), and
|
||||
(c) computing a dB partial alongside dW.
|
||||
"""
|
||||
|
||||
# NOTE: do NOT use `from __future__ import annotations` here.
|
||||
# cuTile's @ct.kernel introspects the wrapped function's __annotations__ to
|
||||
# discover which parameters are Constants (i.e., compile-time specialized).
|
||||
# PEP-563 future annotations stringifies all annotations, which makes cuTile
|
||||
# silently miss the Constant flags, so a *single* compiled binary gets reused
|
||||
# across all callers — yielding catastrophically wrong results when N or
|
||||
# TILE_SIZE changes between launches. (Repro: shape0=(32,256) passes; shape1=
|
||||
# (128,1024) after it returns max_diff~10 because the second launch reuses
|
||||
# the binary specialized for the first launch.)
|
||||
|
||||
import math
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import cuda.tile as ct
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from liger_kernel.backends import Capability
|
||||
from liger_kernel.backends import register_op
|
||||
from liger_kernel.ops._nvidia_shared import cutile_compiler_available as _cutile_compiler_available
|
||||
from liger_kernel.ops._nvidia_shared import next_pow2 as _next_pow2
|
||||
from liger_kernel.ops._nvidia_shared import num_sms as _num_sms
|
||||
from liger_kernel.ops._nvidia_shared import to_local_if_dtensor as _to_local_if_dtensor
|
||||
|
||||
# Backward kernels above this tile width risk spilling on Blackwell (the live
|
||||
# working set is ~24 B/element × 9 vectors (x, dy, x_hat, wdy, w, dw, db, dx,
|
||||
# scratch) ≈ 216 KB at 8192 lanes, which is right at the limit of the 256 KB
|
||||
# register file). LayerNorm holds slightly more state than RMSNorm because of
|
||||
# the extra dB accumulator; we keep the same 8192 ceiling and fail-fast in
|
||||
# ``_layer_norm_backward`` if the user feeds a wider row, mirroring the
|
||||
# Triton kernel's BLOCK_SIZE assertion. The test framework's "documented
|
||||
# range limit" path will skip rather than fail when this fires.
|
||||
# Both forward and backward materialise a single ``next_pow2(N)``-wide tile per
|
||||
# row, so the supported hidden dimension is capped identically on both sides.
|
||||
_MAX_TILE = 8192
|
||||
_BWD_MAX_TILE = _MAX_TILE
|
||||
|
||||
|
||||
def _launch(device: torch.device, grid, kernel, args) -> None:
|
||||
"""Device-safe cuTile launch tied to the input tensor's CUDA device.
|
||||
|
||||
Dispatch selects this backend from the first CUDA tensor's device, so in a
|
||||
multi-GPU process the process-current device can differ. Guarding the
|
||||
device (and using that device's stream) keeps the launch on the tensor's
|
||||
GPU. Mirrors ``ops/cutile/ops/fused_linear_cross_entropy.py``.
|
||||
"""
|
||||
with torch.cuda.device(device):
|
||||
ct.launch(torch.cuda.current_stream(device), grid, kernel, args)
|
||||
|
||||
|
||||
def _select_mode(mode: Optional[str], n_rows: int, n_cols: int, device: torch.device) -> str:
|
||||
"""Pick a kernel variant when ``mode`` is ``None``.
|
||||
|
||||
Heuristic (matches the cuTile RMSNorm sibling after its multi-row removal):
|
||||
- if M > NUM_SMS * 2, use single-row ``static_persistent`` — each block
|
||||
strides over rows, amortising launch overhead and keeping W & B
|
||||
register-resident across many rows;
|
||||
- else, narrow rows (<= 4096) prefer the cached weight variant;
|
||||
- else fall through to the standard one-row-per-program kernel.
|
||||
|
||||
A prior multi-row variant of the persistent kernel was removed: the
|
||||
row-tile / row-element index relationship in the persistent loop was
|
||||
wrong (each block's row range overlapped its neighbours' by
|
||||
``TILE_SIZE_M - 1`` rows, producing races and incorrect results when the
|
||||
runtime gate was lifted). The single-row persistent kernel below matches
|
||||
NVIDIA TileGym's reference for LayerNorm.
|
||||
|
||||
Validates positive dimensions so a shape bug in the caller surfaces here
|
||||
as a clear ValueError rather than a cryptic CUDA grid error inside the
|
||||
kernel launch.
|
||||
"""
|
||||
if mode is not None:
|
||||
return mode
|
||||
if n_rows <= 0 or n_cols <= 0:
|
||||
raise ValueError(f"layer_norm_cutile: invalid shape ({n_rows}, {n_cols}); both dims must be positive.")
|
||||
sms = _num_sms(device)
|
||||
if n_rows > sms * 2:
|
||||
return "static_persistent"
|
||||
if n_cols <= 4096:
|
||||
return "multi_wave_cached"
|
||||
return "standard"
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# Forward kernels
|
||||
#
|
||||
# All three variants compute identical math; they differ only in launch
|
||||
# topology and in how aggressively they cache W/B and overlap memory loads.
|
||||
#
|
||||
# Numerics: we always upcast X, W, B to fp32 inside the kernel for the mean,
|
||||
# variance, and rstd reductions; the final Y is cast back to X's dtype on
|
||||
# store. This mirrors the Triton kernel's fp32-reductions/io-dtype-output
|
||||
# strategy. Mean and RSTD are stored in fp32 (not X's dtype) so that large-mean
|
||||
# and bf16 rows keep the precision the backward pass needs.
|
||||
#
|
||||
# Variance is computed as a MASKED centered second moment,
|
||||
# ``sum(mask * (x - mean)**2) / N``, NOT the ``E[x**2] - E[x]**2`` identity:
|
||||
# the latter catastrophically cancels in fp32 for rows with a large mean
|
||||
# (inputs ~1e4 -> var ~0/negative -> rsqrt(var+eps) -> NaN). The mask zeroes
|
||||
# the ZERO-padded lanes (index >= N) whose centered value ``-mean`` would
|
||||
# otherwise inflate the sum.
|
||||
# ===========================================================================
|
||||
@ct.kernel(occupancy=ct.ByTarget(sm_100=16))
|
||||
def _fwd_standard(
|
||||
X,
|
||||
W,
|
||||
B,
|
||||
Y,
|
||||
Mean,
|
||||
RSTD,
|
||||
N: ct.Constant[int],
|
||||
eps: ct.Constant[float],
|
||||
TILE_SIZE: ct.Constant[int],
|
||||
):
|
||||
"""One row per program block. Good when rows are wide or M is small."""
|
||||
row = ct.bid(0)
|
||||
|
||||
x = ct.load(X, index=(row, 0), shape=(1, TILE_SIZE), allow_tma=False, latency=10, padding_mode=ct.PaddingMode.ZERO)
|
||||
w = ct.load(W, index=(0,), shape=(TILE_SIZE,), allow_tma=False, latency=1, padding_mode=ct.PaddingMode.ZERO)
|
||||
b = ct.load(B, index=(0,), shape=(TILE_SIZE,), allow_tma=False, latency=1, padding_mode=ct.PaddingMode.ZERO)
|
||||
|
||||
x_f32 = ct.astype(x, np.float32)
|
||||
w_f32 = ct.astype(w, np.float32)
|
||||
b_f32 = ct.astype(b, np.float32)
|
||||
|
||||
mean = ct.sum(x_f32) / N
|
||||
# Numerically-stable variance via a masked centered second moment.
|
||||
#
|
||||
# We deliberately do NOT use the E[x**2] - E[x]**2 identity: for rows with
|
||||
# a large mean (e.g. inputs ~1e4) the two fp32 terms are nearly equal and
|
||||
# catastrophically cancel -> var ~0 or negative -> rsqrt(var+eps) -> NaN.
|
||||
#
|
||||
# padding_mode=ZERO leaves the padded lanes (index >= N) at 0, so
|
||||
# centered = 0 - mean = -mean there; masking those lanes to 0 keeps them
|
||||
# out of the variance sum. sum(x)/N and the store are unaffected by the
|
||||
# padded zeros, so mean is exact.
|
||||
cols = ct.arange(TILE_SIZE, dtype=ct.int32)
|
||||
mask_f32 = ct.astype(ct.less(cols, N), np.float32)
|
||||
centered = ct.reshape(x_f32, (TILE_SIZE,)) - mean
|
||||
centered_masked = ct.mul(centered, mask_f32)
|
||||
var = ct.sum(ct.mul(centered_masked, centered_masked)) / N
|
||||
rstd = ct.rsqrt(var + eps)
|
||||
|
||||
# Persist Mean and RSTD in fp32 so large-mean/bf16 rows keep the precision
|
||||
# the backward pass needs (backward upcasts to fp32 anyway).
|
||||
ct.store(Mean, index=(row,), tile=ct.reshape(mean, (1,)))
|
||||
ct.store(RSTD, index=(row,), tile=ct.reshape(rstd, (1,)))
|
||||
|
||||
y_f32 = ct.mul(ct.mul(centered, rstd), w_f32) + b_f32
|
||||
y = ct.astype(y_f32, x.dtype)
|
||||
ct.store(Y, index=(row, 0), tile=ct.reshape(y, (1, TILE_SIZE)), allow_tma=False, latency=3)
|
||||
|
||||
|
||||
@ct.kernel(occupancy=ct.ByTarget(sm_100=16))
|
||||
def _fwd_persistent_singlerow(
|
||||
X,
|
||||
W,
|
||||
B,
|
||||
Y,
|
||||
Mean,
|
||||
RSTD,
|
||||
n_rows: ct.Constant[int],
|
||||
N: ct.Constant[int],
|
||||
eps: ct.Constant[float],
|
||||
TILE_SIZE: ct.Constant[int],
|
||||
):
|
||||
"""NUM_SMS persistent blocks, **one row at a time** striding over rows.
|
||||
|
||||
This is the only ``static_persistent`` variant; a prior multi-row variant
|
||||
was removed (see module docstring for the index bug).
|
||||
|
||||
W and B are loaded once and cached for the lifetime of the block; X is
|
||||
reloaded per row with a higher latency hint to overlap multiple in-flight
|
||||
loads.
|
||||
"""
|
||||
pid = ct.bid(0)
|
||||
num_blocks = ct.num_blocks(0)
|
||||
|
||||
w = ct.load(W, index=(0,), shape=(TILE_SIZE,), allow_tma=False, latency=1, padding_mode=ct.PaddingMode.ZERO)
|
||||
b = ct.load(B, index=(0,), shape=(TILE_SIZE,), allow_tma=False, latency=1, padding_mode=ct.PaddingMode.ZERO)
|
||||
w_f32 = ct.astype(w, np.float32)
|
||||
b_f32 = ct.astype(b, np.float32)
|
||||
|
||||
row_idx = pid
|
||||
while row_idx < n_rows:
|
||||
x = ct.load(
|
||||
X, index=(row_idx, 0), shape=(1, TILE_SIZE), allow_tma=False, latency=10, padding_mode=ct.PaddingMode.ZERO
|
||||
)
|
||||
x_f32 = ct.astype(x, np.float32)
|
||||
|
||||
mean = ct.sum(x_f32) / N
|
||||
# Numerically-stable masked variance (see _fwd_standard for the full
|
||||
# rationale: the E[x**2]-E[x]**2 identity cancels for large-mean rows).
|
||||
cols = ct.arange(TILE_SIZE, dtype=ct.int32)
|
||||
mask_f32 = ct.astype(ct.less(cols, N), np.float32)
|
||||
centered = ct.reshape(x_f32, (TILE_SIZE,)) - mean
|
||||
centered_masked = ct.mul(centered, mask_f32)
|
||||
var = ct.sum(ct.mul(centered_masked, centered_masked)) / N
|
||||
rstd = ct.rsqrt(var + eps)
|
||||
|
||||
ct.store(Mean, index=(row_idx,), tile=ct.reshape(mean, (1,)))
|
||||
ct.store(RSTD, index=(row_idx,), tile=ct.reshape(rstd, (1,)))
|
||||
|
||||
y_f32 = ct.mul(ct.mul(centered, rstd), w_f32) + b_f32
|
||||
y = ct.astype(y_f32, x.dtype)
|
||||
ct.store(Y, index=(row_idx, 0), tile=ct.reshape(y, (1, TILE_SIZE)), allow_tma=False, latency=3)
|
||||
|
||||
row_idx = row_idx + num_blocks
|
||||
|
||||
|
||||
@ct.kernel(occupancy=ct.ByTarget(sm_100=16))
|
||||
def _fwd_cached(
|
||||
X,
|
||||
W,
|
||||
B,
|
||||
Y,
|
||||
Mean,
|
||||
RSTD,
|
||||
N: ct.Constant[int],
|
||||
eps: ct.Constant[float],
|
||||
TILE_SIZE: ct.Constant[int],
|
||||
):
|
||||
"""Single row per program with W/B cached and a latency-hinted X load.
|
||||
|
||||
Identical math to ``_fwd_standard``; the only difference is the explicit
|
||||
ordering hint (load W & B first, then issue the X load with latency=10 so
|
||||
the scheduler keeps it in flight while we do the arithmetic). On narrow
|
||||
rows this gives a ~10-15% bump over ``_fwd_standard``.
|
||||
"""
|
||||
row = ct.bid(0)
|
||||
|
||||
w = ct.load(W, index=(0,), shape=(TILE_SIZE,), allow_tma=False, latency=1, padding_mode=ct.PaddingMode.ZERO)
|
||||
b = ct.load(B, index=(0,), shape=(TILE_SIZE,), allow_tma=False, latency=1, padding_mode=ct.PaddingMode.ZERO)
|
||||
w_f32 = ct.astype(w, np.float32)
|
||||
b_f32 = ct.astype(b, np.float32)
|
||||
|
||||
x = ct.load(X, index=(row, 0), shape=(1, TILE_SIZE), allow_tma=False, latency=10, padding_mode=ct.PaddingMode.ZERO)
|
||||
x_f32 = ct.astype(x, np.float32)
|
||||
|
||||
mean = ct.sum(x_f32) / N
|
||||
# Numerically-stable masked variance (see _fwd_standard for the full
|
||||
# rationale: the E[x**2]-E[x]**2 identity cancels for large-mean rows).
|
||||
cols = ct.arange(TILE_SIZE, dtype=ct.int32)
|
||||
mask_f32 = ct.astype(ct.less(cols, N), np.float32)
|
||||
centered = ct.reshape(x_f32, (TILE_SIZE,)) - mean
|
||||
centered_masked = ct.mul(centered, mask_f32)
|
||||
var = ct.sum(ct.mul(centered_masked, centered_masked)) / N
|
||||
rstd = ct.rsqrt(var + eps)
|
||||
|
||||
ct.store(Mean, index=(row,), tile=ct.reshape(mean, (1,)))
|
||||
ct.store(RSTD, index=(row,), tile=ct.reshape(rstd, (1,)))
|
||||
|
||||
y_f32 = ct.mul(ct.mul(centered, rstd), w_f32) + b_f32
|
||||
y = ct.astype(y_f32, x.dtype)
|
||||
ct.store(Y, index=(row, 0), tile=ct.reshape(y, (1, TILE_SIZE)), allow_tma=False, latency=3)
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# Backward kernel — persistent, one block per SM, rows_per_program rows each.
|
||||
#
|
||||
# Each block walks ``rows_per_program`` consecutive rows. For each row it:
|
||||
# (1) loads X[row,:], dY[row,:], Mean[row], RSTD[row];
|
||||
# (2) computes x_hat, wdy, c1, c2 in fp32;
|
||||
# (3) writes dX[row,:] in the input dtype;
|
||||
# (4) accumulates dY * x_hat and dY into block-local fp32 vectors.
|
||||
#
|
||||
# After the row loop the block writes its accumulated dW/dB partials at row
|
||||
# `pid` of the per-SM partial buffers. The host then reduces:
|
||||
# dW = _dW.sum(dim=0).to(W.dtype)
|
||||
# dB = _dB.sum(dim=0).to(B.dtype)
|
||||
# matching the Triton kernel.
|
||||
# ===========================================================================
|
||||
@ct.kernel(occupancy=ct.ByTarget(sm_100=16))
|
||||
def _bwd(
|
||||
dY,
|
||||
X,
|
||||
W,
|
||||
Mean,
|
||||
RSTD,
|
||||
dX,
|
||||
dW_partial,
|
||||
dB_partial,
|
||||
n_rows: ct.Constant[int],
|
||||
N: ct.Constant[int],
|
||||
rows_per_program: ct.Constant[int],
|
||||
TILE_SIZE: ct.Constant[int],
|
||||
):
|
||||
"""Backward for LayerNorm with W and B (the only path the dispatcher
|
||||
exposes — torch.nn.LayerNorm always has affine params, and our
|
||||
functional.layer_norm substitutes a zeros tensor when bias is None)."""
|
||||
pid = ct.bid(0)
|
||||
row_start = pid * rows_per_program
|
||||
row_end_val = (pid + 1) * rows_per_program
|
||||
|
||||
w = ct.load(W, index=(0,), shape=(TILE_SIZE,), allow_tma=False, latency=1, padding_mode=ct.PaddingMode.ZERO)
|
||||
w_f32 = ct.astype(w, np.float32)
|
||||
|
||||
dw_accum = ct.full((TILE_SIZE,), 0.0, np.float32)
|
||||
db_accum = ct.full((TILE_SIZE,), 0.0, np.float32)
|
||||
|
||||
row_idx = row_start
|
||||
while row_idx < row_end_val and row_idx < n_rows:
|
||||
dy_row = ct.reshape(
|
||||
ct.load(
|
||||
dY,
|
||||
index=(row_idx, 0),
|
||||
shape=(1, TILE_SIZE),
|
||||
allow_tma=False,
|
||||
latency=10,
|
||||
padding_mode=ct.PaddingMode.ZERO,
|
||||
),
|
||||
(TILE_SIZE,),
|
||||
)
|
||||
x_row = ct.reshape(
|
||||
ct.load(
|
||||
X,
|
||||
index=(row_idx, 0),
|
||||
shape=(1, TILE_SIZE),
|
||||
allow_tma=False,
|
||||
latency=10,
|
||||
padding_mode=ct.PaddingMode.ZERO,
|
||||
),
|
||||
(TILE_SIZE,),
|
||||
)
|
||||
mean_val = ct.load(
|
||||
Mean, index=(row_idx,), shape=(1,), allow_tma=False, latency=1, padding_mode=ct.PaddingMode.ZERO
|
||||
)
|
||||
rstd_val = ct.load(
|
||||
RSTD, index=(row_idx,), shape=(1,), allow_tma=False, latency=1, padding_mode=ct.PaddingMode.ZERO
|
||||
)
|
||||
|
||||
x_f32 = ct.astype(x_row, np.float32)
|
||||
dy_f32 = ct.astype(dy_row, np.float32)
|
||||
mean_f32 = ct.astype(mean_val, np.float32)
|
||||
rstd_f32 = ct.astype(rstd_val, np.float32)
|
||||
|
||||
x_hat = ct.mul(x_f32 - mean_f32, rstd_f32)
|
||||
wdy = ct.mul(w_f32, dy_f32)
|
||||
|
||||
c1 = ct.sum(ct.mul(x_hat, wdy)) / N
|
||||
c2 = ct.sum(wdy) / N
|
||||
|
||||
dx_row_f32 = ct.mul(wdy - ct.mul(x_hat, c1) - c2, rstd_f32)
|
||||
ct.store(
|
||||
dX,
|
||||
index=(row_idx, 0),
|
||||
tile=ct.reshape(ct.astype(dx_row_f32, x_row.dtype), (1, TILE_SIZE)),
|
||||
allow_tma=False,
|
||||
latency=3,
|
||||
)
|
||||
|
||||
dw_accum = dw_accum + ct.mul(dy_f32, x_hat)
|
||||
db_accum = db_accum + dy_f32
|
||||
|
||||
row_idx = row_idx + 1
|
||||
|
||||
ct.store(dW_partial, index=(pid, 0), tile=ct.reshape(dw_accum, (1, TILE_SIZE)), allow_tma=False, latency=3)
|
||||
ct.store(dB_partial, index=(pid, 0), tile=ct.reshape(db_accum, (1, TILE_SIZE)), allow_tma=False, latency=3)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Forward kernel dispatch table
|
||||
# ---------------------------------------------------------------------------
|
||||
# Three modes. The ``_fwd_persistent_singlerow`` symbol name is historical —
|
||||
# the multi-row sibling was removed (see module docstring for the index bug);
|
||||
# the singlerow kernel is now the only persistent variant.
|
||||
_FWD = {
|
||||
"standard": _fwd_standard,
|
||||
"static_persistent": _fwd_persistent_singlerow,
|
||||
"multi_wave_cached": _fwd_cached,
|
||||
}
|
||||
|
||||
_VALID_MODES = (
|
||||
"standard",
|
||||
"static_persistent",
|
||||
"multi_wave_cached",
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Host-side launchers
|
||||
# ---------------------------------------------------------------------------
|
||||
def _layer_norm_forward(X, W, B, eps, mode):
|
||||
"""Launch the forward kernel; return saved tensors for backward.
|
||||
|
||||
Returns ``(Y, X_flat, Mean, RSTD, TILE_SIZE, selected_mode)`` so the
|
||||
caller can store everything for backward and report the selected mode.
|
||||
"""
|
||||
shape = X.shape
|
||||
n_cols = shape[-1]
|
||||
X_flat = X.view(-1, n_cols)
|
||||
n_rows = X_flat.shape[0]
|
||||
device = X_flat.device
|
||||
|
||||
if X_flat.shape[1] != W.shape[0]:
|
||||
raise ValueError(f"Hidden size mismatch: X has {X_flat.shape[1]}, W has {W.shape[0]}")
|
||||
if B.shape[0] != W.shape[0]:
|
||||
raise ValueError(f"Bias size mismatch: B has {B.shape[0]}, W has {W.shape[0]}")
|
||||
|
||||
if n_cols > _MAX_TILE:
|
||||
raise RuntimeError(
|
||||
f"cuTile layer_norm only supports hidden dim <= {_MAX_TILE}; got {n_cols}. "
|
||||
f"Both forward and backward materialise a single tile of this width; "
|
||||
f"use the Triton backend for wider rows."
|
||||
)
|
||||
|
||||
# Do NOT pre-quantize fp32 affine params to the activation dtype: the
|
||||
# kernels already upcast W and B to fp32 internally (see the forward and
|
||||
# backward bodies), and the autograd wrapper saves these *same* tensors for
|
||||
# backward. Rounding them here would make forward use different affine
|
||||
# values than backward differentiates, and would drop dB below the bias's
|
||||
# true precision.
|
||||
|
||||
TILE_SIZE = _next_pow2(n_cols)
|
||||
|
||||
Y = torch.empty_like(X_flat)
|
||||
# Mean and RSTD are cached for backward in fp32 (not X's dtype): for
|
||||
# large-mean or bf16 rows, storing them in a low-precision dtype loses
|
||||
# information the backward reduction needs. Backward upcasts to fp32
|
||||
# anyway, so this is the natural storage dtype.
|
||||
Mean = torch.empty(n_rows, dtype=torch.float32, device=X_flat.device)
|
||||
RSTD = torch.empty(n_rows, dtype=torch.float32, device=X_flat.device)
|
||||
|
||||
mode = _select_mode(mode, n_rows, n_cols, device)
|
||||
if mode not in _VALID_MODES:
|
||||
raise ValueError(f"cuTile layer_norm: unknown mode {mode!r}; expected one of {_VALID_MODES}")
|
||||
|
||||
kernel = _FWD[mode]
|
||||
|
||||
if mode == "static_persistent":
|
||||
grid = (_num_sms(device),)
|
||||
_launch(device, grid, kernel, (X_flat, W, B, Y, Mean, RSTD, n_rows, n_cols, eps, TILE_SIZE))
|
||||
else:
|
||||
grid = (n_rows,)
|
||||
_launch(device, grid, kernel, (X_flat, W, B, Y, Mean, RSTD, n_cols, eps, TILE_SIZE))
|
||||
|
||||
return Y.view(*shape), X_flat, Mean, RSTD, TILE_SIZE, mode
|
||||
|
||||
|
||||
def _layer_norm_backward(dY, X, W, B, Mean, RSTD, TILE_SIZE):
|
||||
"""Launch the backward kernel; return ``(dX, dW, dB)``.
|
||||
|
||||
``B`` is passed only so ``dB`` can be returned in the bias's own dtype
|
||||
(the kernel itself does not read ``B`` — ``dB`` depends only on ``dY``).
|
||||
"""
|
||||
shape = dY.shape
|
||||
n_cols = shape[-1]
|
||||
dY_flat = dY.view(-1, n_cols)
|
||||
n_rows = dY_flat.shape[0]
|
||||
device = dY_flat.device
|
||||
|
||||
if n_cols > _BWD_MAX_TILE:
|
||||
raise RuntimeError(
|
||||
f"cuTile layer_norm backward only supports hidden dim <= {_BWD_MAX_TILE}; "
|
||||
f"got {n_cols}. Use the Triton backend for wider rows."
|
||||
)
|
||||
|
||||
sms = _num_sms(device)
|
||||
rows_per_program = math.ceil(n_rows / sms)
|
||||
|
||||
dX_flat = torch.empty_like(dY_flat)
|
||||
# fp32 partials for numerical stability — one row per SM, host-reduced.
|
||||
_dW = torch.empty((sms, n_cols), dtype=torch.float32, device=W.device)
|
||||
_dB = torch.empty((sms, n_cols), dtype=torch.float32, device=W.device)
|
||||
|
||||
grid = (sms,)
|
||||
_launch(
|
||||
device,
|
||||
grid,
|
||||
_bwd,
|
||||
(dY_flat, X, W, Mean, RSTD, dX_flat, _dW, _dB, n_rows, n_cols, rows_per_program, TILE_SIZE),
|
||||
)
|
||||
|
||||
dW = _dW.sum(dim=0).to(W.dtype)
|
||||
dB = _dB.sum(dim=0).to(B.dtype) # return the bias gradient in the bias's own dtype
|
||||
return dX_flat.view(*shape), dW, dB
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# autograd.Function — same shape as the RMSNorm sibling. ``mode`` is plumbed
|
||||
# through .apply() as a positional arg (Function.apply rejects unknown
|
||||
# kwargs) and the dispatcher's mode kwarg is consumed by the registration
|
||||
# wrapper below.
|
||||
# ---------------------------------------------------------------------------
|
||||
class _LigerLayerNormCuTileFunction(torch.autograd.Function):
|
||||
"""cuTile LayerNorm. ``mode`` selects between standard / persistent /
|
||||
cached forward variants. Backward always uses the single persistent
|
||||
kernel (rows split across SMs)."""
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, X, W, B, eps, mode):
|
||||
X = _to_local_if_dtensor(X)
|
||||
|
||||
X = X.contiguous()
|
||||
W = W.contiguous()
|
||||
B = B.contiguous()
|
||||
|
||||
Y, X_flat, Mean, RSTD, TILE_SIZE, _selected_mode = _layer_norm_forward(X, W, B, eps, mode)
|
||||
ctx.TILE_SIZE = TILE_SIZE
|
||||
ctx.save_for_backward(X_flat, W, B, Mean, RSTD)
|
||||
return Y
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, dY):
|
||||
dY = _to_local_if_dtensor(dY).contiguous()
|
||||
|
||||
X, W, B, Mean, RSTD = ctx.saved_tensors
|
||||
# B is saved so dB can be returned in the bias's own dtype (and so its
|
||||
# gradient flows back through the same ctx). The backward kernel itself
|
||||
# doesn't read B — dB depends only on dY.
|
||||
dX, dW, dB = _layer_norm_backward(dY, X, W, B, Mean, RSTD, ctx.TILE_SIZE)
|
||||
# Match forward arity: (X, W, B, eps, mode)
|
||||
return dX, dW, dB, None, None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Public registration
|
||||
# ---------------------------------------------------------------------------
|
||||
@register_op(
|
||||
"layer_norm",
|
||||
impl_name="nvidia-cutile",
|
||||
capability=Capability(
|
||||
min_cc=(10, 0),
|
||||
modules=["cuda.tile", "torch"],
|
||||
predicate=_cutile_compiler_available,
|
||||
),
|
||||
modes=("standard", "static_persistent", "multi_wave_cached"),
|
||||
default_mode="static_persistent",
|
||||
# See sibling _cutile/rms_norm.py for measured-perf rationale: cuTile
|
||||
# LayerNorm also loses to both Triton and CuTeDSL across our shape sweep
|
||||
# (matches NVIDIA's own pptx — H100 fwd 0.98x, H100 bwd 0.93x, B200 bwd
|
||||
# 0.90x vs Triton). Last-fallback rank. Both forward and backward
|
||||
# materialise a single ``next_pow2(N)``-wide tile per row, so the hidden
|
||||
# dimension is capped at 8192 (``_MAX_TILE``) on both sides; wider rows
|
||||
# raise a clear error rather than silently mis-normalising.
|
||||
preference_rank=80,
|
||||
tolerances={
|
||||
torch.float16: {"atol_fwd": 5e-3, "atol_bwd": 5e-2, "rtol_fwd": 1e-3, "rtol_bwd": 1e-2},
|
||||
torch.bfloat16: {"atol_fwd": 2e-2, "atol_bwd": 1e-1, "rtol_fwd": 1e-2, "rtol_bwd": 2e-2},
|
||||
# Two-pass reduction (mean + variance) accumulates more rounding than
|
||||
# RMSNorm; atol=5e-4 absorbs N up to ~32K at unit magnitude.
|
||||
torch.float32: {"atol_fwd": 5e-4, "atol_bwd": 1e-3, "rtol_fwd": 1e-4, "rtol_bwd": 1e-3},
|
||||
},
|
||||
notes="cuTile static-persistent + multi-wave LayerNorm for Blackwell.",
|
||||
)
|
||||
def layer_norm_cutile(
|
||||
x: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor,
|
||||
eps: float,
|
||||
*,
|
||||
mode: Optional[str] = None,
|
||||
) -> torch.Tensor:
|
||||
"""Dispatch entry-point for the cuTile LayerNorm backend.
|
||||
|
||||
``bias`` is always a real tensor — :func:`liger_kernel.functional.layer_norm`
|
||||
substitutes a zeros tensor when the user passes ``None``. We add it
|
||||
unconditionally inside the kernel; we deliberately do not special-case
|
||||
"all zeros" because the cost of the add at this scale is dwarfed by the
|
||||
memory traffic for X.
|
||||
"""
|
||||
if mode is not None and mode not in _VALID_MODES:
|
||||
raise ValueError(f"cuTile layer_norm: unknown mode {mode!r}; valid modes are {_VALID_MODES}.")
|
||||
return _LigerLayerNormCuTileFunction.apply(x, weight, bias, eps, mode)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,448 @@
|
||||
"""cuTile (NVIDIA cuda.tile) backend for ``softmax``.
|
||||
|
||||
This module implements row-wise softmax with NVIDIA's cuTile DSL and registers
|
||||
it via ``@register_op``. It exposes three forward kernel variants, selected by
|
||||
the ``mode`` argument:
|
||||
|
||||
- ``"standard"`` — one row per program block, single tile per row. Loaded via
|
||||
``ct.gather`` with ``padding_value=-inf`` so the max/sum reductions ignore
|
||||
the padded lanes. Best for narrow-to-mid rows that fit in a single tile.
|
||||
- ``"static_persistent"`` — ``NUM_SMS * occupancy`` persistent blocks each
|
||||
striding over rows. Amortises launch overhead when ``M`` is much larger than
|
||||
``NUM_SMS``. Single tile per row (``TILE_SIZE = next_pow2(N)``).
|
||||
- ``"chunked"`` — a 3-pass online kernel for very wide rows (``N`` larger than
|
||||
the per-tile cap). Pass 1 finds the row max, pass 2 accumulates the
|
||||
denominator, pass 3 writes ``exp(x - m) / d``. Numerically identical to the
|
||||
multi-block Triton path.
|
||||
|
||||
Backward is a single persistent kernel: ``dx = y * (dy - sum(dy * y))`` with the
|
||||
dot-product reduction in fp32. It mirrors the Triton backward exactly. For wide
|
||||
rows the backward kernel runs the same 3-pass-free two-pass structure
|
||||
(accumulate the dot, then write dx chunk by chunk).
|
||||
|
||||
References
|
||||
----------
|
||||
- cuTile softmax reference: NVIDIA TileGym
|
||||
``src/tilegym/ops/cutile/softmax.py`` — source of the ``_softmax_kernel``
|
||||
(gather + ``padding_value=-inf``), ``_softmax_kernel_tma`` (single
|
||||
``ct.load`` per row), and ``_softmax_kernel_chunked`` (3-pass online) idioms,
|
||||
plus the ``min(NUM_SM * k, n_rows)`` persistent-grid occupancy hint.
|
||||
- Triton reference: ``liger_kernel.ops.softmax`` — defines the exact numerics
|
||||
we reproduce bit-for-bit: single-block ``y = exp(x - m) / sum(exp(x - m))``,
|
||||
multi-block online running ``(m, d)``, and the backward
|
||||
``dx = y * (dy - sum(dy * y))``.
|
||||
"""
|
||||
|
||||
# NOTE: do NOT use `from __future__ import annotations` here.
|
||||
# cuTile's @ct.kernel introspects the wrapped function's __annotations__ to
|
||||
# discover which parameters are Constants (i.e., compile-time specialized).
|
||||
# PEP-563 future annotations stringifies all annotations, which makes cuTile
|
||||
# silently miss the Constant flags, so a *single* compiled binary gets reused
|
||||
# across all callers — yielding catastrophically wrong results when N or
|
||||
# TILE_SIZE changes between launches. (Same caveat as the cuTile rms_norm /
|
||||
# jsd siblings; see their headers for the full repro.)
|
||||
|
||||
import math
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import cuda.tile as ct
|
||||
import torch
|
||||
|
||||
from liger_kernel.backends import Capability
|
||||
from liger_kernel.backends import register_op
|
||||
from liger_kernel.ops._nvidia_shared import cutile_compiler_available as _cutile_compiler_available
|
||||
from liger_kernel.ops._nvidia_shared import next_pow2 as _next_pow2
|
||||
from liger_kernel.ops._nvidia_shared import num_sms as _num_sms
|
||||
from liger_kernel.ops._nvidia_shared import to_local_if_dtensor as _to_local_if_dtensor
|
||||
from liger_kernel.ops.utils import ensure_contiguous
|
||||
|
||||
_ConstInt = ct.Constant[int]
|
||||
|
||||
# Single-tile kernels keep a full row resident in registers/smem. Above this
|
||||
# width we route to the chunked (multi-pass) kernels, mirroring the Triton
|
||||
# backend's ``n_cols <= BLOCK_SIZE`` single-block vs multi-block split. The cap
|
||||
# is a power of two so ``next_pow2(N) <= _MAX_SINGLE_TILE`` is the exact gate.
|
||||
_MAX_SINGLE_TILE = 8192
|
||||
# Chunk width used by the multi-pass (chunked) kernels. A power of two so the
|
||||
# per-chunk ``ct.arange`` tiles align with the gather/scatter bounds.
|
||||
_CHUNK_SIZE = 4096
|
||||
|
||||
|
||||
def _launch(device: torch.device, grid, kernel, args) -> None:
|
||||
"""Device-safe cuTile launch tied to the input tensor's CUDA device.
|
||||
|
||||
Dispatch selects this backend from the first CUDA tensor's device, so in a
|
||||
multi-GPU process the process-current device can differ. Guarding the
|
||||
device (and using that device's stream) keeps the launch on the tensor's
|
||||
GPU. Mirrors ``ops/cutile/ops/fused_linear_cross_entropy.py``.
|
||||
"""
|
||||
with torch.cuda.device(device):
|
||||
ct.launch(torch.cuda.current_stream(device), grid, kernel, args)
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# Forward kernels — single tile per row (``standard`` mode)
|
||||
# ===========================================================================
|
||||
@ct.kernel(occupancy=ct.ByTarget(sm_100=4))
|
||||
def _fwd_standard(
|
||||
X,
|
||||
Y,
|
||||
n_cols: _ConstInt,
|
||||
TILE_SIZE: _ConstInt,
|
||||
):
|
||||
"""One row per program block. fp32 reduction; cast back on store.
|
||||
|
||||
Loaded via ``ct.gather`` with ``padding_value=-inf`` so the padded lanes
|
||||
(when ``n_cols < TILE_SIZE``) contribute ``exp(-inf) == 0`` to the
|
||||
denominator and never win the max.
|
||||
"""
|
||||
row_idx = ct.bid(0)
|
||||
offsets = ct.arange(TILE_SIZE, dtype=ct.int32)
|
||||
|
||||
x = ct.gather(X, (row_idx, offsets), check_bounds=True, padding_value=-math.inf)
|
||||
x_f32 = ct.astype(x, ct.float32)
|
||||
|
||||
row_max = ct.max(x_f32, 0, keepdims=True)
|
||||
numerator = ct.exp(x_f32 - row_max)
|
||||
denominator = ct.sum(numerator, 0, keepdims=True)
|
||||
y = numerator / denominator
|
||||
|
||||
ct.scatter(Y, (row_idx, offsets), ct.astype(y, X.dtype), check_bounds=True)
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# Forward kernels — static persistent (``static_persistent`` mode)
|
||||
#
|
||||
# NUM_SMS*occupancy persistent blocks, each striding over rows by num_blocks.
|
||||
# Single tile per row. This is the dominant perf variant when M >> NUM_SMS
|
||||
# (the TileGym reference's default ``_softmax_kernel``).
|
||||
# ===========================================================================
|
||||
@ct.kernel(occupancy=ct.ByTarget(sm_100=4))
|
||||
def _fwd_persistent(
|
||||
X,
|
||||
Y,
|
||||
n_rows: _ConstInt,
|
||||
n_cols: _ConstInt,
|
||||
TILE_SIZE: _ConstInt,
|
||||
):
|
||||
"""Static persistent: each block processes multiple rows."""
|
||||
pid = ct.bid(0)
|
||||
num_blocks = ct.num_blocks(0)
|
||||
offsets = ct.arange(TILE_SIZE, dtype=ct.int32)
|
||||
|
||||
row_idx = pid
|
||||
while row_idx < n_rows:
|
||||
x = ct.gather(X, (row_idx, offsets), check_bounds=True, padding_value=-math.inf)
|
||||
x_f32 = ct.astype(x, ct.float32)
|
||||
|
||||
row_max = ct.max(x_f32, 0, keepdims=True)
|
||||
numerator = ct.exp(x_f32 - row_max)
|
||||
denominator = ct.sum(numerator, 0, keepdims=True)
|
||||
y = numerator / denominator
|
||||
|
||||
ct.scatter(Y, (row_idx, offsets), ct.astype(y, X.dtype), check_bounds=True)
|
||||
row_idx = row_idx + num_blocks
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# Forward kernel — chunked 3-pass (``chunked`` mode) for very wide rows
|
||||
#
|
||||
# Pass 1: running max over chunks.
|
||||
# Pass 2: accumulate denominator = sum_chunks(sum(exp(chunk - row_max))).
|
||||
# Pass 3: write exp(chunk - row_max) / denominator.
|
||||
#
|
||||
# Numerically equivalent to the Triton multi-block kernel (which folds the
|
||||
# online running (m, d) — here we split into max-then-sum, which yields the
|
||||
# same final m and the same d up to fp32 rounding).
|
||||
# ===========================================================================
|
||||
@ct.kernel(occupancy=ct.ByTarget(sm_100=4))
|
||||
def _fwd_chunked(
|
||||
X,
|
||||
Y,
|
||||
n_rows: _ConstInt,
|
||||
n_cols: _ConstInt,
|
||||
CHUNK_SIZE: _ConstInt,
|
||||
):
|
||||
"""Static persistent + 3-pass online softmax for wide rows."""
|
||||
pid = ct.bid(0)
|
||||
num_blocks = ct.num_blocks(0)
|
||||
num_chunks = (n_cols + CHUNK_SIZE - 1) // CHUNK_SIZE
|
||||
col_base = ct.arange(CHUNK_SIZE, dtype=ct.int32)
|
||||
|
||||
row_idx = pid
|
||||
while row_idx < n_rows:
|
||||
row_max = ct.full((1,), -math.inf, dtype=ct.float32)
|
||||
denominator = ct.full((1,), 0.0, dtype=ct.float32)
|
||||
|
||||
# Pass 1: row max.
|
||||
for chunk_idx in range(num_chunks):
|
||||
cols = col_base + chunk_idx * CHUNK_SIZE
|
||||
chunk = ct.gather(X, (row_idx, cols), check_bounds=True, padding_value=-math.inf)
|
||||
chunk = ct.astype(chunk, ct.float32)
|
||||
chunk_max = ct.max(chunk, 0, keepdims=True)
|
||||
row_max = ct.maximum(row_max, chunk_max)
|
||||
|
||||
# Pass 2: denominator.
|
||||
for chunk_idx in range(num_chunks):
|
||||
cols = col_base + chunk_idx * CHUNK_SIZE
|
||||
chunk = ct.gather(X, (row_idx, cols), check_bounds=True, padding_value=-math.inf)
|
||||
chunk = ct.astype(chunk, ct.float32)
|
||||
numerator = ct.exp(chunk - row_max)
|
||||
denominator = denominator + ct.sum(numerator, 0, keepdims=True)
|
||||
|
||||
# Pass 3: write softmax.
|
||||
for chunk_idx in range(num_chunks):
|
||||
cols = col_base + chunk_idx * CHUNK_SIZE
|
||||
chunk = ct.gather(X, (row_idx, cols), check_bounds=True, padding_value=-math.inf)
|
||||
chunk = ct.astype(chunk, ct.float32)
|
||||
numerator = ct.exp(chunk - row_max)
|
||||
y = numerator / denominator
|
||||
ct.scatter(Y, (row_idx, cols), ct.astype(y, X.dtype), check_bounds=True)
|
||||
|
||||
row_idx = row_idx + num_blocks
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# Backward kernels — dx = y * (dy - sum(dy * y))
|
||||
# ===========================================================================
|
||||
@ct.kernel(occupancy=ct.ByTarget(sm_100=4))
|
||||
def _bwd_persistent(
|
||||
dY,
|
||||
Y,
|
||||
dX,
|
||||
n_rows: _ConstInt,
|
||||
n_cols: _ConstInt,
|
||||
TILE_SIZE: _ConstInt,
|
||||
):
|
||||
"""Static persistent backward, single tile per row.
|
||||
|
||||
The dot-product ``sum(dy * y)`` is reduced in fp32 (matching Triton's
|
||||
``tl.float32`` accumulator). Padded lanes contribute ``0 * 0 == 0`` to the
|
||||
dot because both dY and Y gather with ``padding_value=0``.
|
||||
"""
|
||||
pid = ct.bid(0)
|
||||
num_blocks = ct.num_blocks(0)
|
||||
offsets = ct.arange(TILE_SIZE, dtype=ct.int32)
|
||||
|
||||
row_idx = pid
|
||||
while row_idx < n_rows:
|
||||
dy = ct.gather(dY, (row_idx, offsets), check_bounds=True, padding_value=0.0)
|
||||
y = ct.gather(Y, (row_idx, offsets), check_bounds=True, padding_value=0.0)
|
||||
|
||||
dy_f32 = ct.astype(dy, ct.float32)
|
||||
y_f32 = ct.astype(y, ct.float32)
|
||||
|
||||
dot = ct.sum(dy_f32 * y_f32, 0, keepdims=True)
|
||||
dx = y_f32 * (dy_f32 - dot)
|
||||
|
||||
ct.scatter(dX, (row_idx, offsets), ct.astype(dx, dX.dtype), check_bounds=True)
|
||||
row_idx = row_idx + num_blocks
|
||||
|
||||
|
||||
@ct.kernel(occupancy=ct.ByTarget(sm_100=4))
|
||||
def _bwd_chunked(
|
||||
dY,
|
||||
Y,
|
||||
dX,
|
||||
n_rows: _ConstInt,
|
||||
n_cols: _ConstInt,
|
||||
CHUNK_SIZE: _ConstInt,
|
||||
):
|
||||
"""Two-pass backward for wide rows.
|
||||
|
||||
Pass 1: accumulate ``dot = sum(dy * y)`` over chunks (fp32).
|
||||
Pass 2: write ``dx = y * (dy - dot)`` chunk by chunk.
|
||||
|
||||
Matches the Triton ``_softmax_multi_block_backward_kernel`` exactly.
|
||||
"""
|
||||
pid = ct.bid(0)
|
||||
num_blocks = ct.num_blocks(0)
|
||||
num_chunks = (n_cols + CHUNK_SIZE - 1) // CHUNK_SIZE
|
||||
col_base = ct.arange(CHUNK_SIZE, dtype=ct.int32)
|
||||
|
||||
row_idx = pid
|
||||
while row_idx < n_rows:
|
||||
dot = ct.full((1,), 0.0, dtype=ct.float32)
|
||||
|
||||
# Pass 1: accumulate the dot product.
|
||||
for chunk_idx in range(num_chunks):
|
||||
cols = col_base + chunk_idx * CHUNK_SIZE
|
||||
dy = ct.gather(dY, (row_idx, cols), check_bounds=True, padding_value=0.0)
|
||||
y = ct.gather(Y, (row_idx, cols), check_bounds=True, padding_value=0.0)
|
||||
dot = dot + ct.sum(ct.astype(dy, ct.float32) * ct.astype(y, ct.float32), 0, keepdims=True)
|
||||
|
||||
# Pass 2: write dx.
|
||||
for chunk_idx in range(num_chunks):
|
||||
cols = col_base + chunk_idx * CHUNK_SIZE
|
||||
dy = ct.gather(dY, (row_idx, cols), check_bounds=True, padding_value=0.0)
|
||||
y = ct.gather(Y, (row_idx, cols), check_bounds=True, padding_value=0.0)
|
||||
dx = ct.astype(y, ct.float32) * (ct.astype(dy, ct.float32) - dot)
|
||||
ct.scatter(dX, (row_idx, cols), ct.astype(dx, dX.dtype), check_bounds=True)
|
||||
|
||||
row_idx = row_idx + num_blocks
|
||||
|
||||
|
||||
_VALID_MODES = ("standard", "static_persistent", "chunked")
|
||||
|
||||
|
||||
def _select_mode(mode: Optional[str], n_rows: int, n_cols: int, device: torch.device) -> str:
|
||||
"""Pick a forward kernel variant when ``mode`` is ``None``.
|
||||
|
||||
Heuristic (mirrors the cuTile rms_norm sibling + TileGym reference):
|
||||
- Wide rows (``next_pow2(N) > _MAX_SINGLE_TILE``) must use ``chunked``
|
||||
— a single tile would exceed the per-row register/smem budget.
|
||||
- Otherwise, when ``M`` is much larger than ``NUM_SMS`` use
|
||||
``static_persistent`` to amortise launch overhead.
|
||||
- Otherwise fall back to the simple one-row-per-program ``standard``.
|
||||
|
||||
Validates positive dims so caller shape bugs surface here as a clean
|
||||
ValueError rather than a cryptic CUDA grid error.
|
||||
"""
|
||||
if mode is not None:
|
||||
# Explicit single-tile modes cannot handle a row wider than the
|
||||
# per-row tile budget; reject up-front rather than compiling/launching
|
||||
# an oversized tile that silently corrupts results. Only ``chunked``
|
||||
# streams arbitrarily wide rows.
|
||||
if mode in ("standard", "static_persistent") and _next_pow2(n_cols) > _MAX_SINGLE_TILE:
|
||||
raise ValueError(
|
||||
f"cuTile softmax mode={mode!r} only supports rows with "
|
||||
f"next_pow2(N) <= {_MAX_SINGLE_TILE}; got N={n_cols} "
|
||||
f"(next_pow2={_next_pow2(n_cols)}). Use mode='chunked' or auto "
|
||||
f"selection (mode=None) for wider rows."
|
||||
)
|
||||
return mode
|
||||
|
||||
if n_rows <= 0 or n_cols <= 0:
|
||||
raise ValueError(f"softmax_cutile: invalid shape ({n_rows}, {n_cols}); both dims must be positive.")
|
||||
|
||||
if _next_pow2(n_cols) > _MAX_SINGLE_TILE:
|
||||
return "chunked"
|
||||
if n_rows > _num_sms(device) * 2:
|
||||
return "static_persistent"
|
||||
return "standard"
|
||||
|
||||
|
||||
def _persistent_grid(n_rows: int, device: torch.device) -> tuple:
|
||||
"""Saturate at most ``NUM_SMS * 4`` blocks but never more than ``n_rows``.
|
||||
|
||||
Matches TileGym's ``num_programs = min(NUM_SM * 4, n_rows)`` occupancy hint.
|
||||
"""
|
||||
return (min(_num_sms(device) * 4, max(n_rows, 1)), 1, 1)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Host-side launchers
|
||||
# ---------------------------------------------------------------------------
|
||||
def _softmax_forward(x: torch.Tensor, mode: Optional[str]):
|
||||
"""Launch the forward kernel; return ``(Y, selected_mode)``."""
|
||||
*batch, n_cols = x.shape
|
||||
x2d = x.contiguous().view(-1, n_cols)
|
||||
n_rows = x2d.shape[0]
|
||||
|
||||
selected = _select_mode(mode, n_rows, n_cols, x2d.device)
|
||||
if selected not in _VALID_MODES:
|
||||
raise ValueError(f"cuTile softmax: unknown mode {selected!r}; expected one of {_VALID_MODES}")
|
||||
|
||||
y2d = torch.empty_like(x2d)
|
||||
device = x2d.device
|
||||
|
||||
if selected == "standard":
|
||||
tile_size = _next_pow2(n_cols)
|
||||
_launch(device, (n_rows, 1, 1), _fwd_standard, (x2d, y2d, n_cols, tile_size))
|
||||
elif selected == "static_persistent":
|
||||
tile_size = _next_pow2(n_cols)
|
||||
_launch(device, _persistent_grid(n_rows, device), _fwd_persistent, (x2d, y2d, n_rows, n_cols, tile_size))
|
||||
else: # chunked
|
||||
chunk = min(_CHUNK_SIZE, _next_pow2(n_cols))
|
||||
_launch(device, _persistent_grid(n_rows, device), _fwd_chunked, (x2d, y2d, n_rows, n_cols, chunk))
|
||||
|
||||
return y2d.view(*batch, n_cols), selected
|
||||
|
||||
|
||||
def _softmax_backward(dy: torch.Tensor, y: torch.Tensor, selected_mode: str):
|
||||
"""Launch the backward kernel; return ``dX``."""
|
||||
*batch, n_cols = dy.shape
|
||||
dy2d = dy.contiguous().view(-1, n_cols)
|
||||
y2d = y.contiguous().view(-1, n_cols)
|
||||
n_rows = dy2d.shape[0]
|
||||
dx2d = torch.empty_like(dy2d)
|
||||
device = dy2d.device
|
||||
|
||||
if selected_mode == "chunked":
|
||||
chunk = min(_CHUNK_SIZE, _next_pow2(n_cols))
|
||||
_launch(device, _persistent_grid(n_rows, device), _bwd_chunked, (dy2d, y2d, dx2d, n_rows, n_cols, chunk))
|
||||
else:
|
||||
tile_size = _next_pow2(n_cols)
|
||||
_launch(device, _persistent_grid(n_rows, device), _bwd_persistent, (dy2d, y2d, dx2d, n_rows, n_cols, tile_size))
|
||||
|
||||
return dx2d.view(*batch, n_cols)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# autograd.Function — mode is threaded through ctx so backward picks the
|
||||
# matching kernel layout (single-tile vs chunked).
|
||||
# ---------------------------------------------------------------------------
|
||||
class _LigerSoftmaxCuTileFunction(torch.autograd.Function):
|
||||
"""cuTile softmax. ``mode`` flows in as a positional arg (resolved in the
|
||||
registration wrapper) so this class stays ``.apply()``-compatible.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
@ensure_contiguous
|
||||
def forward(ctx, input_: torch.Tensor, mode: Optional[str]):
|
||||
input_ = _to_local_if_dtensor(input_)
|
||||
y, selected_mode = _softmax_forward(input_, mode)
|
||||
ctx.save_for_backward(y)
|
||||
ctx.selected_mode = selected_mode
|
||||
return y
|
||||
|
||||
@staticmethod
|
||||
@ensure_contiguous
|
||||
def backward(ctx, grad_output):
|
||||
(y,) = ctx.saved_tensors
|
||||
grad_output = _to_local_if_dtensor(grad_output)
|
||||
dx = _softmax_backward(grad_output, y, ctx.selected_mode)
|
||||
# forward arity: (input_, mode)
|
||||
return dx, None
|
||||
|
||||
|
||||
_CUTILE_SOFTMAX_TOLERANCES = {
|
||||
torch.float16: {"atol_fwd": 5e-3, "atol_bwd": 5e-2, "rtol_fwd": 1e-3, "rtol_bwd": 1e-2},
|
||||
torch.bfloat16: {"atol_fwd": 2e-2, "atol_bwd": 1e-1, "rtol_fwd": 1e-2, "rtol_bwd": 2e-2},
|
||||
torch.float32: {"atol_fwd": 1e-5, "atol_bwd": 1e-4, "rtol_fwd": 1e-5, "rtol_bwd": 1e-4},
|
||||
}
|
||||
|
||||
|
||||
@register_op(
|
||||
"softmax",
|
||||
impl_name="nvidia-cutile",
|
||||
capability=Capability(
|
||||
min_cc=(10, 0),
|
||||
modules=["cuda.tile", "torch"],
|
||||
predicate=_cutile_compiler_available,
|
||||
),
|
||||
modes=("standard", "static_persistent", "chunked"),
|
||||
default_mode="static_persistent",
|
||||
# Keep cuTile as the last fallback. CuTe DSL is faster on measured
|
||||
# production shapes; cuTile remains useful for very wide rows where the
|
||||
# single-tile DSL paths are range-capped.
|
||||
preference_rank=80,
|
||||
tolerances=_CUTILE_SOFTMAX_TOLERANCES,
|
||||
notes="cuTile static-persistent + chunked softmax for Blackwell (sm_100).",
|
||||
)
|
||||
def softmax_cutile(
|
||||
x: torch.Tensor,
|
||||
*,
|
||||
mode: Optional[str] = None,
|
||||
) -> torch.Tensor:
|
||||
"""Dispatch entry-point for the cuTile softmax backend.
|
||||
|
||||
Resolves ``mode`` (auto if ``None``), then forwards to
|
||||
:class:`_LigerSoftmaxCuTileFunction`. The wrapper exists because the
|
||||
dispatcher passes ``mode`` as a kwarg and ``Function.apply`` does not
|
||||
accept unknown kwargs.
|
||||
"""
|
||||
if mode is not None and mode not in _VALID_MODES:
|
||||
raise ValueError(f"cuTile softmax: unknown mode {mode!r}; valid modes are {_VALID_MODES}.")
|
||||
return _LigerSoftmaxCuTileFunction.apply(x, mode)
|
||||
@@ -25,14 +25,21 @@ _TORCH_VERSION = Version(torch.__version__.split("+")[0])
|
||||
_ADDMM_SUPPORTS_OUT_DTYPE = _TORCH_VERSION >= Version("2.8.0")
|
||||
|
||||
|
||||
def _max_capability() -> int:
|
||||
def _max_capability(device=None) -> int:
|
||||
# Runtime dispatch-plane key (e.g. 100 = sm_100/B200, 103 = sm_103/B300).
|
||||
# Called per wrapper invocation (never memoized) so mock-cc suites key
|
||||
# separately per call, same idiom as capability.is_satisfied.
|
||||
#
|
||||
# ``device`` selects WHICH GPU's capability is read. It must be threaded
|
||||
# from the tensor whose chunk budget we are sizing (``student_input``) so a
|
||||
# heterogeneous process does not size a cuda:1 input by cuda:0's plane.
|
||||
# ``None`` preserves the legacy current-device behavior; a non-CUDA device
|
||||
# (e.g. a CPU ``student_input`` while CUDA happens to be visible) falls back
|
||||
# to the tight pre-Blackwell budget via the exception guard below.
|
||||
if torch.cuda.is_available():
|
||||
try:
|
||||
maj, minor = torch.cuda.get_device_capability()
|
||||
except (IndexError, RuntimeError, AssertionError): # pragma: no cover
|
||||
maj, minor = torch.cuda.get_device_capability(device)
|
||||
except (IndexError, RuntimeError, AssertionError, ValueError): # pragma: no cover
|
||||
return 0
|
||||
return 10 * maj + minor
|
||||
return 0
|
||||
@@ -98,7 +105,11 @@ def fused_linear_jsd_forward(
|
||||
# constant, sm_100 the B200 constant; every pre-Blackwell device (Hopper,
|
||||
# Ampere, CPU/no-CUDA where _max_capability()==0) keeps the original tight
|
||||
# budget (constant 1 == ceil(V / H)), so non-Blackwell memory is unchanged.
|
||||
_cc = _max_capability()
|
||||
# Size the chunk budget by the plane of the GPU that actually holds the
|
||||
# fused-linear input, NOT the process-current device — otherwise a cuda:1
|
||||
# (B200) input could be sized by a cuda:0 (B300) plane and get ~2x the
|
||||
# intended transient-memory budget in a heterogeneous process.
|
||||
_cc = _max_capability(student_input.device)
|
||||
chunk_mem_const = _CHUNK_MEM_CONST_B300 if _cc > 100 else (_CHUNK_MEM_CONST_B200 if _cc == 100 else 1)
|
||||
inc_factor = triton.cdiv(V, chunk_mem_const * H)
|
||||
chunk_size = triton.next_power_of_2(triton.cdiv(BT, inc_factor)) # ceil(BT / inc_factor)
|
||||
|
||||
@@ -344,6 +344,7 @@ def assert_op_correctness(
|
||||
seed: int = 0,
|
||||
device: Optional[torch.device] = None,
|
||||
extra: Optional[dict] = None,
|
||||
strict_grad: bool = False,
|
||||
) -> None:
|
||||
"""Run ``dispatch(op_name, ..., backend=backend)`` and verify forward AND
|
||||
backward against the registered PyTorch reference impl.
|
||||
@@ -466,6 +467,8 @@ def assert_op_correctness(
|
||||
|
||||
dw_msg = _close_message("dw", weight.grad, weight_ref.grad, atol=atol_bwd, rtol=rtol_bwd)
|
||||
if dw_msg is not None:
|
||||
if strict_grad:
|
||||
raise AssertionError(f"[{ctx}] {dw_msg}")
|
||||
ok, sum_msg = sum_based_dw_close(weight.grad, weight_ref.grad, rtol=rtol_bwd, atol=atol_bwd)
|
||||
if not ok:
|
||||
raise AssertionError(f"[{ctx}] {dw_msg}\nElement-wise dw failed; sum-based fallback also failed: {sum_msg}")
|
||||
@@ -477,6 +480,8 @@ def assert_op_correctness(
|
||||
if op_name == "layer_norm" and bias is not None and bias.requires_grad:
|
||||
db_msg = _close_message("db", bias.grad, bias_ref.grad, atol=atol_bwd, rtol=rtol_bwd)
|
||||
if db_msg is not None:
|
||||
if strict_grad:
|
||||
raise AssertionError(f"[{ctx}] {db_msg}")
|
||||
ok, sum_msg = sum_based_dw_close(bias.grad, bias_ref.grad, rtol=rtol_bwd, atol=atol_bwd)
|
||||
if not ok:
|
||||
raise AssertionError(
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
"""cuTile ``_select_mode`` shape/width validation.
|
||||
|
||||
These are pure host-side selector tests (no kernel launch), so they run on any
|
||||
box where the ``_cutile`` module *imports* — CUDA is not required. Only a
|
||||
genuine import failure (e.g. the ``cuda.tile`` package is absent) is skipped;
|
||||
real regressions in the selector surface as failures rather than skips.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
|
||||
def _import(module, name):
|
||||
try:
|
||||
mod = __import__(
|
||||
f"liger_kernel.ops.backends._cutile.{module}",
|
||||
fromlist=[name],
|
||||
)
|
||||
except ImportError as e:
|
||||
pytest.skip(f"cuTile backend unimportable on this host: {e}")
|
||||
return getattr(mod, name)
|
||||
|
||||
|
||||
# A device descriptor is required by the new signature but is only *used* on the
|
||||
# auto path (mode=None) for the SM-count heuristic. The zero/negative-dim and
|
||||
# explicit-mode cases below return/raise before touching it, so a descriptor for
|
||||
# a possibly-absent GPU is fine.
|
||||
_DEV = torch.device("cuda:0")
|
||||
|
||||
|
||||
class TestRmsNormSelectMode:
|
||||
@pytest.mark.parametrize("n_rows,n_cols", [(0, 256), (256, 0), (-1, 256), (256, -1), (0, 0)])
|
||||
def test_zero_or_negative_dims_raise(self, n_rows, n_cols):
|
||||
_select_mode = _import("rms_norm", "_select_mode")
|
||||
with pytest.raises(ValueError) as ei:
|
||||
_select_mode(None, n_rows, n_cols, _DEV)
|
||||
assert "invalid shape" in str(ei.value)
|
||||
|
||||
def test_explicit_mode_returned_as_is(self):
|
||||
_select_mode = _import("rms_norm", "_select_mode")
|
||||
# Explicit mode is returned verbatim (caller owns dim validation) and is
|
||||
# a plain string — the obsolete (mode, tile_m) tuple contract is gone.
|
||||
assert _select_mode("standard", 0, 0, _DEV) == "standard"
|
||||
|
||||
|
||||
class TestLayerNormSelectMode:
|
||||
@pytest.mark.parametrize("n_rows,n_cols", [(0, 256), (256, 0), (0, 0)])
|
||||
def test_zero_or_negative_dims_raise(self, n_rows, n_cols):
|
||||
_select_mode = _import("layer_norm", "_select_mode")
|
||||
with pytest.raises(ValueError) as ei:
|
||||
_select_mode(None, n_rows, n_cols, _DEV)
|
||||
assert "invalid shape" in str(ei.value)
|
||||
|
||||
def test_explicit_mode_returned_as_is(self):
|
||||
_select_mode = _import("layer_norm", "_select_mode")
|
||||
assert _select_mode("static_persistent", 0, 0, _DEV) == "static_persistent"
|
||||
|
||||
|
||||
class TestSoftmaxSelectMode:
|
||||
@pytest.mark.parametrize("n_rows,n_cols", [(0, 256), (256, 0), (0, 0)])
|
||||
def test_zero_or_negative_dims_raise(self, n_rows, n_cols):
|
||||
_select_mode = _import("softmax", "_select_mode")
|
||||
with pytest.raises(ValueError) as ei:
|
||||
_select_mode(None, n_rows, n_cols, _DEV)
|
||||
assert "invalid shape" in str(ei.value)
|
||||
|
||||
def test_explicit_single_tile_mode_rejects_wide_row(self):
|
||||
_select_mode = _import("softmax", "_select_mode")
|
||||
# next_pow2(9000) = 16384 > 8192 → the single-tile modes must reject it
|
||||
# instead of compiling/launching an oversized tile.
|
||||
for mode in ("standard", "static_persistent"):
|
||||
with pytest.raises(ValueError) as ei:
|
||||
_select_mode(mode, 8, 9000, _DEV)
|
||||
assert "next_pow2" in str(ei.value)
|
||||
|
||||
def test_explicit_chunked_mode_accepts_wide_row(self):
|
||||
_select_mode = _import("softmax", "_select_mode")
|
||||
# chunked streams arbitrarily wide rows, so it is not rejected.
|
||||
assert _select_mode("chunked", 8, 9000, _DEV) == "chunked"
|
||||
|
||||
def test_explicit_single_tile_mode_accepts_supported_width(self):
|
||||
_select_mode = _import("softmax", "_select_mode")
|
||||
assert _select_mode("standard", 8, 8192, _DEV) == "standard"
|
||||
@@ -0,0 +1,799 @@
|
||||
"""Pinned ``nvidia-cutile`` parity + correctness tests for the five cuTile
|
||||
backends (RMSNorm, LayerNorm, Softmax, JSD, FusedLinearJSD).
|
||||
|
||||
These are the tests requested by the PR-E code review. They exercise the fixes
|
||||
landed in ``src/liger_kernel/ops/backends/_cutile/``:
|
||||
|
||||
* forward AND backward parity vs the Triton reference in fp16 / bf16 / fp32,
|
||||
using each op's *registered* tolerances;
|
||||
* mixed-precision affine — bf16/fp16 activations with **fp32** affine params —
|
||||
checking the output, ``dX``, ``dW`` (and LayerNorm ``dB``) against an fp32
|
||||
ground-truth reference, and that ``dW`` / ``dB`` come back in the parameter
|
||||
dtype (the pre-fix backend silently quantized them);
|
||||
* forward-selector boundaries for the three norm/softmax variant selectors
|
||||
(``N=1``, non-powers-of-two, both sides of each cutoff, 8192, 8193), and that
|
||||
explicit invalid softmax modes raise clearly;
|
||||
* JSD ``beta`` in ``{0, 0.5, 1, 0.1}``, no-labels / scattered-ignore /
|
||||
all-ignore, and non-unit upstream gradient;
|
||||
* non-contiguous inputs for every public op.
|
||||
|
||||
The whole module skips cleanly when ``cuda.tile`` is not importable (so CPU /
|
||||
Hopper CI stays green). Individual tests skip when the cuTile impl is not
|
||||
*registered* for the op on this device (i.e. compute capability < 10.0), but a
|
||||
real assertion failure is never turned into a skip.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
# Skip the entire module cleanly when the cuTile compiler stack is absent.
|
||||
pytest.importorskip("cuda.tile", reason="cuda.tile (cuTile) not importable on this host")
|
||||
|
||||
import liger_kernel.functional as F # noqa: E402 (importing also registers the op discovery map)
|
||||
|
||||
from liger_kernel.backends.dispatch import available_impls # noqa: E402
|
||||
from liger_kernel.backends.registry import get_discovery_failures # noqa: E402
|
||||
from liger_kernel.backends.registry import get_registered # noqa: E402
|
||||
from liger_kernel.ops._nvidia_shared import cutile_compiler_available # noqa: E402
|
||||
from liger_kernel.testing import assert_op_correctness # noqa: E402
|
||||
|
||||
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
|
||||
|
||||
_DTYPES = [torch.float16, torch.bfloat16, torch.float32]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Impl discovery helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
def _impl_matching(op: str, needle: str) -> Optional[str]:
|
||||
for name in available_impls(op):
|
||||
if needle in name:
|
||||
return name
|
||||
return None
|
||||
|
||||
|
||||
def _cutile_supported() -> tuple[bool, str]:
|
||||
"""Is this host *genuinely* capable of running the cuTile impls?
|
||||
|
||||
``cuda.tile`` importability is already guaranteed by the module-level
|
||||
``importorskip``. Here we check the remaining hard requirements the
|
||||
``Capability`` gate enforces: a CUDA device, compute capability >= (10, 0),
|
||||
and a reachable ``tileiras`` compiler. When all three hold, absence of the
|
||||
registered impl is a real import/registration regression — NOT an
|
||||
"unsupported machine" — so callers must fail instead of skip.
|
||||
"""
|
||||
if not torch.cuda.is_available():
|
||||
return False, "CUDA device not available"
|
||||
cc = torch.cuda.get_device_capability()
|
||||
if cc < (10, 0):
|
||||
return False, f"device sm_{cc[0]}{cc[1]} < required sm_100 (cuTile needs compute capability >= 10.0)"
|
||||
if not cutile_compiler_available():
|
||||
return False, "cuTile (tileiras) compiler not reachable"
|
||||
return True, ""
|
||||
|
||||
|
||||
def _cutile_impl(op: str) -> str:
|
||||
impl = _impl_matching(op, "cutile")
|
||||
if impl is not None:
|
||||
return impl
|
||||
|
||||
supported, reason = _cutile_supported()
|
||||
if not supported:
|
||||
# Genuinely unsupported host (CPU / cc < 10.0 / missing compiler): skip.
|
||||
pytest.skip(f"nvidia-cutile not available for {op!r}: {reason}")
|
||||
|
||||
# Compatible host (CUDA + sm>=100 + cuTile compiler) but the impl is still
|
||||
# absent -> a real backend import/registration regression that discovery
|
||||
# swallowed. Surface it and FAIL rather than mask it as a skip.
|
||||
failures = get_discovery_failures()
|
||||
detail = (
|
||||
"; ".join(f"{path} -> {etype.__name__}: {msg}" for path, (etype, msg) in failures.items())
|
||||
if failures
|
||||
else "none recorded"
|
||||
)
|
||||
registered = get_registered(op, "nvidia-cutile")
|
||||
if registered is None:
|
||||
pytest.fail(
|
||||
f"nvidia-cutile impl for {op!r} did NOT register on a cuTile-capable host "
|
||||
f"(CUDA + sm>=100 + tileiras present). This is an import/registration regression, "
|
||||
f"not an unsupported machine. Recorded discovery failures: {detail}"
|
||||
)
|
||||
# Registered but filtered out by the capability gate on a host we consider
|
||||
# compatible -> a capability-evaluation mismatch worth surfacing.
|
||||
pytest.fail(
|
||||
f"nvidia-cutile impl for {op!r} is registered but reported unavailable on a cuTile-capable host. "
|
||||
f"Capability reason: {registered.capability.reason_if_unsatisfied()!r}. "
|
||||
f"Recorded discovery failures: {detail}"
|
||||
)
|
||||
|
||||
|
||||
def _triton_impl(op: str) -> str:
|
||||
impl = _impl_matching(op, "triton")
|
||||
if impl is None:
|
||||
pytest.skip(f"triton reference not registered for {op!r}")
|
||||
return impl
|
||||
|
||||
|
||||
def _tols(op: str, impl: str, dtype: torch.dtype) -> dict:
|
||||
registered = get_registered(op, impl)
|
||||
entry = dict((registered.tolerances if registered is not None else {}).get(dtype, {}))
|
||||
defaults = {
|
||||
torch.float16: {"atol_fwd": 1e-2, "rtol_fwd": 1e-2, "atol_bwd": 5e-2, "rtol_bwd": 5e-2},
|
||||
torch.bfloat16: {"atol_fwd": 2e-2, "rtol_fwd": 2e-2, "atol_bwd": 1e-1, "rtol_bwd": 5e-2},
|
||||
torch.float32: {"atol_fwd": 1e-5, "rtol_fwd": 1e-5, "atol_bwd": 1e-4, "rtol_bwd": 1e-4},
|
||||
}
|
||||
for k, v in defaults[dtype].items():
|
||||
entry.setdefault(k, v)
|
||||
return entry
|
||||
|
||||
|
||||
def _assert_close(actual, expected, atol, rtol, label):
|
||||
torch.testing.assert_close(
|
||||
actual.to(torch.float32),
|
||||
expected.to(torch.float32),
|
||||
atol=atol,
|
||||
rtol=rtol,
|
||||
msg=lambda m: f"[{label}] {m}",
|
||||
)
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# Ground-truth references (identical semantics to the Triton kernels)
|
||||
# ===========================================================================
|
||||
def _ref_rms_norm(x, weight, eps, offset, casting_mode):
|
||||
in_dtype = x.dtype
|
||||
if casting_mode == "llama":
|
||||
xf = x.to(torch.float32)
|
||||
rstd = torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps)
|
||||
return (xf * rstd).to(in_dtype) * (offset + weight)
|
||||
if casting_mode == "gemma":
|
||||
xf = x.to(torch.float32)
|
||||
rstd = torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps)
|
||||
return ((xf * rstd) * (offset + weight.to(torch.float32))).to(in_dtype)
|
||||
# none — cuTile keeps input-dtype math when W matches the activation but
|
||||
# promotes to fp32 when W is a wider dtype (its backward reduces dW as
|
||||
# ``x.float() * rstd.float()``). Mirror that fp32-faithful normalization so
|
||||
# the elementwise dW/dX comparison reflects the gradient the kernel really
|
||||
# computes (the mixed-precision test always uses a wider fp32 weight). The
|
||||
# kernel forms ``mean(x*x)`` from input-dtype products (``ct.mul(x, x)``)
|
||||
# then normalizes/reduces in fp32, so square in the input dtype first.
|
||||
mean_sq = x.pow(2).to(torch.float32).mean(-1, keepdim=True)
|
||||
rstd = torch.rsqrt(mean_sq + eps)
|
||||
return ((x.to(torch.float32) * rstd) * (offset + weight)).to(in_dtype)
|
||||
|
||||
|
||||
def _ref_layer_norm(x, weight, bias, eps):
|
||||
in_dtype = x.dtype
|
||||
xf = x.to(torch.float32)
|
||||
mean = xf.mean(-1, keepdim=True)
|
||||
var = xf.var(-1, keepdim=True, unbiased=False)
|
||||
xhat = (xf - mean) * torch.rsqrt(var + eps)
|
||||
return (xhat * weight.to(torch.float32) + bias.to(torch.float32)).to(in_dtype)
|
||||
|
||||
|
||||
def _ref_jsd(log_q, log_p, shift_labels, beta, ignore_index):
|
||||
lq = log_q.float()
|
||||
lp = log_p.float()
|
||||
if beta == 0.0:
|
||||
loss = torch.exp(lp) * (lp - lq)
|
||||
elif beta == 1.0:
|
||||
loss = torch.exp(lq) * (lq - lp)
|
||||
else:
|
||||
m = beta * torch.exp(lp) + (1.0 - beta) * torch.exp(lq)
|
||||
log_m = torch.log(m)
|
||||
loss = beta * torch.exp(lp) * lp + (1.0 - beta) * torch.exp(lq) * lq - m * log_m
|
||||
if shift_labels is not None:
|
||||
mask = (shift_labels != ignore_index).unsqueeze(-1).float()
|
||||
loss = loss * mask
|
||||
n = float(mask.sum().item())
|
||||
if n == 0:
|
||||
return torch.zeros((), dtype=log_q.dtype, device=log_q.device)
|
||||
else:
|
||||
n = float(log_q.shape[0])
|
||||
return (loss.sum() / n).to(log_q.dtype)
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# 1. Forward + backward parity vs the Triton reference (fp16 / bf16 / fp32)
|
||||
# ===========================================================================
|
||||
_NORM_PARITY_SHAPES = [(256, 1024), (512, 2048)]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", _DTYPES)
|
||||
@pytest.mark.parametrize("shape", _NORM_PARITY_SHAPES)
|
||||
@pytest.mark.parametrize("casting_mode", ["llama", "gemma", "none"])
|
||||
def test_rms_norm_fwd_bwd_parity(dtype, shape, casting_mode):
|
||||
"""Forward + backward parity for the cuTile RMSNorm vs the fp32 PyTorch
|
||||
reference (identical semantics to the Triton kernel), using the op's
|
||||
registered tolerances. Compared against the reference rather than the Triton
|
||||
backend directly so a Triton bug can't mask a cuTile bug (and vice-versa) —
|
||||
the convention used by ``test/ops/test_jsd.py``.
|
||||
"""
|
||||
ci = _cutile_impl("rms_norm")
|
||||
assert_op_correctness("rms_norm", ci, shape, dtype, casting_mode=casting_mode, strict_grad=True)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", _DTYPES)
|
||||
@pytest.mark.parametrize("shape", _NORM_PARITY_SHAPES)
|
||||
@pytest.mark.parametrize("with_bias", [True, False])
|
||||
def test_layer_norm_fwd_bwd_parity(dtype, shape, with_bias):
|
||||
"""Forward + backward parity (output, dX, dW, dB) for cuTile LayerNorm vs the
|
||||
fp32 PyTorch reference, using the op's registered tolerances."""
|
||||
ci = _cutile_impl("layer_norm")
|
||||
assert_op_correctness("layer_norm", ci, shape, dtype, extra={"include_bias": with_bias}, strict_grad=True)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", _DTYPES)
|
||||
@pytest.mark.parametrize("shape", [(256, 1024), (512, 4096)])
|
||||
def test_softmax_fwd_bwd_parity(dtype, shape):
|
||||
"""Forward + backward parity for cuTile softmax vs ``torch.softmax`` (fp32)."""
|
||||
ci = _cutile_impl("softmax")
|
||||
M, N = shape
|
||||
tols = _tols("softmax", ci, dtype)
|
||||
g = torch.Generator(device="cpu").manual_seed(0)
|
||||
x_cpu = torch.randn(M, N, generator=g)
|
||||
dy_cpu = torch.randn(M, N, generator=g)
|
||||
|
||||
x = x_cpu.to("cuda", dtype).detach().requires_grad_(True)
|
||||
xr = x_cpu.to("cuda", dtype).detach().requires_grad_(True)
|
||||
y = F.softmax(x, impl=ci)
|
||||
yr = torch.softmax(xr.float(), dim=-1).to(dtype)
|
||||
dy = dy_cpu.to("cuda", dtype)
|
||||
y.backward(dy)
|
||||
yr.backward(dy.clone())
|
||||
|
||||
_assert_close(y, yr, tols["atol_fwd"], tols["rtol_fwd"], f"softmax/{dtype} fwd")
|
||||
_assert_close(x.grad, xr.grad, tols["atol_bwd"], tols["rtol_bwd"], f"softmax/{dtype} dx")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", _DTYPES)
|
||||
@pytest.mark.parametrize("beta", [0.0, 0.5, 1.0])
|
||||
def test_jsd_fwd_bwd_parity(dtype, beta):
|
||||
"""Forward + backward parity for cuTile JSD vs the fp32 PyTorch reference."""
|
||||
ci = _cutile_impl("jsd")
|
||||
BT, V = 64, 512
|
||||
tols = _tols("jsd", ci, dtype)
|
||||
lq0, lp = _jsd_inputs(BT, V, dtype, seed=BT * 100 + V)
|
||||
|
||||
lq = lq0.clone().requires_grad_(True)
|
||||
out = F.jsd(lq, lp, None, beta=beta, ignore_index=-100, impl=ci)
|
||||
out.backward()
|
||||
|
||||
lqr = lq0.clone().requires_grad_(True)
|
||||
ref = _ref_jsd(lqr, lp, None, beta, -100)
|
||||
ref.backward()
|
||||
_assert_close(out, ref, tols["atol_fwd"], tols["rtol_fwd"], f"jsd/{dtype}/beta={beta} fwd")
|
||||
_assert_close(lq.grad, lqr.grad, tols["atol_bwd"], tols["rtol_bwd"], f"jsd/{dtype}/beta={beta} dx")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", _DTYPES)
|
||||
@pytest.mark.parametrize("beta", [0.0, 0.5, 1.0])
|
||||
def test_fused_linear_jsd_parity_vs_triton(dtype, beta):
|
||||
"""Forward + backward parity for cuTile fused_linear_jsd vs the Triton
|
||||
backend. The two share the outer chunked cuBLAS implementation and differ
|
||||
only in the inner cuTile JSD primitive, so they agree tightly; comparing
|
||||
them directly pins the inner primitive."""
|
||||
ci = _cutile_impl("fused_linear_jsd")
|
||||
ti = _triton_impl("fused_linear_jsd")
|
||||
BT, H, V = 64, 128, 512
|
||||
tols = _tols("fused_linear_jsd", ci, dtype)
|
||||
g = torch.Generator(device="cpu").manual_seed(0)
|
||||
si_cpu = torch.randn(BT, H, generator=g)
|
||||
sw_cpu = torch.randn(V, H, generator=g)
|
||||
ti_cpu = torch.randn(BT, H, generator=g)
|
||||
tw_cpu = torch.randn(V, H, generator=g)
|
||||
|
||||
def run(impl):
|
||||
si = si_cpu.to("cuda", dtype).detach().requires_grad_(True)
|
||||
sw = sw_cpu.to("cuda", dtype).detach().requires_grad_(True)
|
||||
t_in = ti_cpu.to("cuda", dtype).detach()
|
||||
tw = tw_cpu.to("cuda", dtype).detach()
|
||||
out = F.fused_linear_jsd(si, sw, t_in, tw, None, jsd_beta=beta, temperature=1.0, impl=impl)
|
||||
out.backward()
|
||||
return out, si.grad, sw.grad
|
||||
|
||||
oc, dsic, dswc = run(ci)
|
||||
ot, dsit, dswt = run(ti)
|
||||
_assert_close(oc, ot, tols["atol_fwd"], tols["rtol_fwd"], f"fljsd/{dtype}/beta={beta} fwd")
|
||||
_assert_close(dsic, dsit, tols["atol_bwd"], tols["rtol_bwd"], f"fljsd/{dtype}/beta={beta} d_student_input")
|
||||
_assert_close(dswc, dswt, tols["atol_bwd"], tols["rtol_bwd"], f"fljsd/{dtype}/beta={beta} d_student_weight")
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# 2. Mixed-precision affine: bf16/fp16 activations + fp32 affine params
|
||||
# ===========================================================================
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
@pytest.mark.parametrize("casting_mode", ["llama", "gemma", "none"])
|
||||
def test_rms_norm_mixed_precision_affine(dtype, casting_mode):
|
||||
"""fp32 weight must NOT be silently quantized to the activation dtype, and
|
||||
backward must differentiate exactly what forward computed. Compared against
|
||||
an fp32 ground-truth reference (identical to Triton's promotion semantics).
|
||||
"""
|
||||
ci = _cutile_impl("rms_norm")
|
||||
M, N = 256, 1024
|
||||
tols = _tols("rms_norm", ci, dtype)
|
||||
g = torch.Generator(device="cpu").manual_seed(1)
|
||||
x_cpu = torch.randn(M, N, generator=g)
|
||||
w_cpu = torch.randn(N, generator=g)
|
||||
|
||||
def build():
|
||||
x = x_cpu.to("cuda", dtype).detach().requires_grad_(True)
|
||||
w = w_cpu.to("cuda", torch.float32).detach().requires_grad_(True) # fp32 affine
|
||||
return x, w
|
||||
|
||||
x, w = build()
|
||||
y = F.rms_norm(x, w, 1e-6, casting_mode=casting_mode, impl=ci)
|
||||
y.backward(torch.ones_like(y))
|
||||
|
||||
xr, wr = build()
|
||||
yr = _ref_rms_norm(xr, wr, 1e-6, 0.0, casting_mode)
|
||||
yr.backward(torch.ones_like(yr))
|
||||
|
||||
# dW and (here) x-grad dtype must be preserved.
|
||||
assert w.grad.dtype == torch.float32, f"dW dtype {w.grad.dtype} != fp32 (weight was fp32)"
|
||||
assert x.grad.dtype == dtype, f"dX dtype {x.grad.dtype} != {dtype}"
|
||||
|
||||
_assert_close(y, yr, tols["atol_fwd"], tols["rtol_fwd"], f"rms_norm-mixed/{casting_mode}/{dtype} fwd")
|
||||
_assert_close(x.grad, xr.grad, tols["atol_bwd"], tols["rtol_bwd"], f"rms_norm-mixed/{casting_mode}/{dtype} dx")
|
||||
_assert_close(w.grad, wr.grad, tols["atol_bwd"], tols["rtol_bwd"], f"rms_norm-mixed/{casting_mode}/{dtype} dw")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
def test_layer_norm_mixed_precision_affine(dtype):
|
||||
ci = _cutile_impl("layer_norm")
|
||||
M, N = 256, 1024
|
||||
tols = _tols("layer_norm", ci, dtype)
|
||||
g = torch.Generator(device="cpu").manual_seed(1)
|
||||
x_cpu = torch.randn(M, N, generator=g)
|
||||
w_cpu = torch.randn(N, generator=g)
|
||||
b_cpu = torch.randn(N, generator=g)
|
||||
|
||||
def build():
|
||||
x = x_cpu.to("cuda", dtype).detach().requires_grad_(True)
|
||||
w = w_cpu.to("cuda", torch.float32).detach().requires_grad_(True) # fp32 affine
|
||||
b = b_cpu.to("cuda", torch.float32).detach().requires_grad_(True) # fp32 affine
|
||||
return x, w, b
|
||||
|
||||
x, w, b = build()
|
||||
y = F.layer_norm(x, w, b, 1e-6, impl=ci)
|
||||
y.backward(torch.ones_like(y))
|
||||
|
||||
xr, wr, br = build()
|
||||
yr = _ref_layer_norm(xr, wr, br, 1e-6)
|
||||
yr.backward(torch.ones_like(yr))
|
||||
|
||||
assert w.grad.dtype == torch.float32, f"dW dtype {w.grad.dtype} != fp32"
|
||||
assert b.grad.dtype == torch.float32, f"dB dtype {b.grad.dtype} != fp32 (bias was fp32)"
|
||||
assert x.grad.dtype == dtype
|
||||
|
||||
_assert_close(y, yr, tols["atol_fwd"], tols["rtol_fwd"], f"layer_norm-mixed/{dtype} fwd")
|
||||
_assert_close(x.grad, xr.grad, tols["atol_bwd"], tols["rtol_bwd"], f"layer_norm-mixed/{dtype} dx")
|
||||
_assert_close(w.grad, wr.grad, tols["atol_bwd"], tols["rtol_bwd"], f"layer_norm-mixed/{dtype} dw")
|
||||
_assert_close(b.grad, br.grad, tols["atol_bwd"], tols["rtol_bwd"], f"layer_norm-mixed/{dtype} db")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
@pytest.mark.parametrize("w_dtype,b_dtype", [(torch.float32, torch.bfloat16), (torch.bfloat16, torch.float32)])
|
||||
def test_layer_norm_affine_dtype_mismatch(dtype, w_dtype, b_dtype):
|
||||
"""W and B with DIFFERENT dtypes: dW must come back in W.dtype and dB in
|
||||
B.dtype *independently* (the pre-fix bug cast dB to W.dtype), and each must
|
||||
be elementwise-correct vs an fp32 ground-truth reference. A same-dtype test
|
||||
can never catch dB being reduced/cast to W.dtype instead of B.dtype."""
|
||||
assert w_dtype != b_dtype
|
||||
ci = _cutile_impl("layer_norm")
|
||||
M, N = 256, 1024
|
||||
tols = _tols("layer_norm", ci, dtype)
|
||||
g = torch.Generator(device="cpu").manual_seed(1)
|
||||
x_cpu = torch.randn(M, N, generator=g)
|
||||
w_cpu = torch.randn(N, generator=g)
|
||||
b_cpu = torch.randn(N, generator=g)
|
||||
|
||||
def build():
|
||||
x = x_cpu.to("cuda", dtype).detach().requires_grad_(True)
|
||||
w = w_cpu.to("cuda", w_dtype).detach().requires_grad_(True)
|
||||
b = b_cpu.to("cuda", b_dtype).detach().requires_grad_(True)
|
||||
return x, w, b
|
||||
|
||||
x, w, b = build()
|
||||
y = F.layer_norm(x, w, b, 1e-6, impl=ci)
|
||||
y.backward(torch.ones_like(y))
|
||||
|
||||
xr, wr, br = build()
|
||||
yr = _ref_layer_norm(xr, wr, br, 1e-6)
|
||||
yr.backward(torch.ones_like(yr))
|
||||
|
||||
# Each param gradient must independently preserve its OWN dtype.
|
||||
assert w.grad.dtype == w_dtype, f"dW dtype {w.grad.dtype} != W.dtype {w_dtype}"
|
||||
assert b.grad.dtype == b_dtype, f"dB dtype {b.grad.dtype} != B.dtype {b_dtype}"
|
||||
assert x.grad.dtype == dtype
|
||||
|
||||
_assert_close(y, yr, tols["atol_fwd"], tols["rtol_fwd"], f"ln-wbmismatch/{dtype}/W{w_dtype}/B{b_dtype} fwd")
|
||||
_assert_close(x.grad, xr.grad, tols["atol_bwd"], tols["rtol_bwd"], f"ln-wbmismatch/{dtype} dx")
|
||||
_assert_close(w.grad, wr.grad, tols["atol_bwd"], tols["rtol_bwd"], f"ln-wbmismatch/{dtype} dw")
|
||||
_assert_close(b.grad, br.grad, tols["atol_bwd"], tols["rtol_bwd"], f"ln-wbmismatch/{dtype} db")
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# 3. Forward-selector boundary cases + explicit-mode rejection
|
||||
# ===========================================================================
|
||||
# Both sides of the RMSNorm/LayerNorm width cutoff (4096) and of the persistent
|
||||
# row cutoff (M vs NUM_SMS), plus N=1, non-powers-of-two, and the 8192 cap.
|
||||
_NORM_BOUNDARY = [
|
||||
(1, 1), # single row, single col
|
||||
(8, 3), # non-pow2 narrow, small M (standard/cached path)
|
||||
(4, 4095), # just below the 4096 cached cutoff
|
||||
(4, 4096), # exactly the cutoff
|
||||
(4, 4097), # just above the cutoff (standard path)
|
||||
(2048, 256), # large M -> static_persistent
|
||||
(4, 8192), # widest supported single tile
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", _NORM_BOUNDARY)
|
||||
def test_rms_norm_selector_boundaries(shape):
|
||||
ci = _cutile_impl("rms_norm")
|
||||
M, N = shape
|
||||
dtype = torch.float32
|
||||
g = torch.Generator(device="cpu").manual_seed(2)
|
||||
x = torch.randn(M, N, generator=g).to("cuda", dtype).detach().requires_grad_(True)
|
||||
w = torch.randn(N, generator=g).to("cuda", dtype).detach().requires_grad_(True)
|
||||
y = F.rms_norm(x, w, 1e-6, casting_mode="llama", impl=ci)
|
||||
yr = _ref_rms_norm(x.detach(), w.detach(), 1e-6, 0.0, "llama")
|
||||
_assert_close(y, yr, 1e-4, 1e-4, f"rms_norm-boundary fwd {shape}")
|
||||
y.backward(torch.ones_like(y))
|
||||
assert torch.isfinite(x.grad).all() and torch.isfinite(w.grad).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", _NORM_BOUNDARY)
|
||||
def test_layer_norm_selector_boundaries(shape):
|
||||
ci = _cutile_impl("layer_norm")
|
||||
M, N = shape
|
||||
dtype = torch.float32
|
||||
g = torch.Generator(device="cpu").manual_seed(2)
|
||||
x = torch.randn(M, N, generator=g).to("cuda", dtype).detach().requires_grad_(True)
|
||||
w = torch.randn(N, generator=g).to("cuda", dtype).detach().requires_grad_(True)
|
||||
b = torch.randn(N, generator=g).to("cuda", dtype).detach().requires_grad_(True)
|
||||
y = F.layer_norm(x, w, b, 1e-6, impl=ci)
|
||||
yr = _ref_layer_norm(x.detach(), w.detach(), b.detach(), 1e-6)
|
||||
_assert_close(y, yr, 1e-3, 1e-3, f"layer_norm-boundary fwd {shape}")
|
||||
y.backward(torch.ones_like(y))
|
||||
assert torch.isfinite(x.grad).all() and torch.isfinite(w.grad).all() and torch.isfinite(b.grad).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("op,fn", [("rms_norm", F.rms_norm), ("layer_norm", F.layer_norm)])
|
||||
def test_norm_hidden_dim_over_cap_raises(op, fn):
|
||||
"""N > 8192 must raise a clear, consistent error in forward (matching the
|
||||
backward cap) rather than silently mis-normalising."""
|
||||
ci = _cutile_impl(op)
|
||||
x = torch.randn(4, 8193, device="cuda", dtype=torch.float32, requires_grad=True)
|
||||
w = torch.randn(8193, device="cuda", dtype=torch.float32, requires_grad=True)
|
||||
with pytest.raises(RuntimeError, match="only supports hidden dim"):
|
||||
if op == "rms_norm":
|
||||
fn(x, w, 1e-6, impl=ci)
|
||||
else:
|
||||
b = torch.randn(8193, device="cuda", dtype=torch.float32)
|
||||
fn(x, w, b, 1e-6, impl=ci)
|
||||
|
||||
|
||||
_SOFTMAX_BOUNDARY = [
|
||||
(1, 1),
|
||||
(8, 3), # non-pow2 narrow
|
||||
(4, 8191), # just below single-tile cap
|
||||
(4, 8192), # exactly at the cap
|
||||
(4, 8193), # just above -> auto picks chunked
|
||||
(2048, 256), # large M -> static_persistent
|
||||
(4, 18432), # very wide -> chunked
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", _SOFTMAX_BOUNDARY)
|
||||
def test_softmax_selector_boundaries(shape):
|
||||
ci = _cutile_impl("softmax")
|
||||
M, N = shape
|
||||
dtype = torch.float32
|
||||
g = torch.Generator(device="cpu").manual_seed(3)
|
||||
x = torch.randn(M, N, generator=g).to("cuda", dtype).detach().requires_grad_(True)
|
||||
y = F.softmax(x, impl=ci) # auto mode
|
||||
yr = torch.softmax(x.detach().float(), dim=-1)
|
||||
_assert_close(y, yr, 1e-4, 1e-4, f"softmax-boundary fwd {shape}")
|
||||
y.backward(torch.randn_like(y))
|
||||
assert torch.isfinite(x.grad).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["standard", "static_persistent"])
|
||||
def test_softmax_explicit_singletile_mode_rejects_wide(mode):
|
||||
"""Explicit single-tile modes must reject a row wider than the tile budget
|
||||
(next_pow2(N) > 8192) instead of compiling/launching an oversized tile."""
|
||||
ci = _cutile_impl("softmax")
|
||||
x = torch.randn(8, 9000, device="cuda", dtype=torch.float32, requires_grad=True)
|
||||
with pytest.raises(ValueError, match="next_pow2"):
|
||||
F.softmax(x, impl=ci, mode=mode)
|
||||
|
||||
|
||||
def test_softmax_invalid_mode_raises():
|
||||
ci = _cutile_impl("softmax")
|
||||
x = torch.randn(8, 128, device="cuda", dtype=torch.float32)
|
||||
with pytest.raises(ValueError, match="unknown mode"):
|
||||
F.softmax(x, impl=ci, mode="does_not_exist")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("op", ["rms_norm", "layer_norm"])
|
||||
def test_norm_invalid_mode_raises(op):
|
||||
ci = _cutile_impl(op)
|
||||
x = torch.randn(8, 128, device="cuda", dtype=torch.float32)
|
||||
w = torch.randn(128, device="cuda", dtype=torch.float32)
|
||||
with pytest.raises(ValueError, match="unknown mode"):
|
||||
if op == "rms_norm":
|
||||
F.rms_norm(x, w, 1e-6, impl=ci, mode="does_not_exist")
|
||||
else:
|
||||
b = torch.randn(128, device="cuda", dtype=torch.float32)
|
||||
F.layer_norm(x, w, b, 1e-6, impl=ci, mode="does_not_exist")
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# 4. JSD beta / ignore-index / non-unit upstream gradient coverage
|
||||
# ===========================================================================
|
||||
def _jsd_inputs(BT, V, dtype, seed=7):
|
||||
torch.manual_seed(seed)
|
||||
raw_q = torch.randn(BT, V, device="cuda", dtype=dtype)
|
||||
raw_p = torch.randn(BT, V, device="cuda", dtype=dtype)
|
||||
return raw_q.log_softmax(-1).detach(), raw_p.log_softmax(-1).detach()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("beta", [0.0, 0.1, 0.5, 1.0])
|
||||
def test_jsd_beta_values(beta):
|
||||
ci = _cutile_impl("jsd")
|
||||
dtype = torch.float32
|
||||
tols = _tols("jsd", ci, dtype)
|
||||
lq0, lp = _jsd_inputs(64, 512, dtype)
|
||||
lq = lq0.clone().requires_grad_(True)
|
||||
out = F.jsd(lq, lp, None, beta=beta, ignore_index=-100, impl=ci)
|
||||
out.backward()
|
||||
|
||||
lqr = lq0.clone().requires_grad_(True)
|
||||
ref = _ref_jsd(lqr, lp, None, beta, -100)
|
||||
ref.backward()
|
||||
_assert_close(out, ref, tols["atol_fwd"], tols["rtol_fwd"], f"jsd-beta={beta} fwd")
|
||||
_assert_close(lq.grad, lqr.grad, tols["atol_bwd"], tols["rtol_bwd"], f"jsd-beta={beta} dx")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"label_kind",
|
||||
["none", "scattered_ignore", "all_ignore"],
|
||||
)
|
||||
def test_jsd_ignore_patterns(label_kind):
|
||||
ci = _cutile_impl("jsd")
|
||||
dtype = torch.float32
|
||||
BT, V = 32, 256
|
||||
tols = _tols("jsd", ci, dtype)
|
||||
lq0, lp = _jsd_inputs(BT, V, dtype, seed=11)
|
||||
ignore_index = -100
|
||||
|
||||
if label_kind == "none":
|
||||
labels = None
|
||||
elif label_kind == "scattered_ignore":
|
||||
labels = torch.arange(BT, device="cuda", dtype=torch.int64) % V
|
||||
labels[::3] = ignore_index # every third row ignored
|
||||
else: # all_ignore
|
||||
labels = torch.full((BT,), ignore_index, device="cuda", dtype=torch.int64)
|
||||
|
||||
lq = lq0.clone().requires_grad_(True)
|
||||
out = F.jsd(lq, lp, labels, beta=0.5, ignore_index=ignore_index, impl=ci)
|
||||
out.backward()
|
||||
|
||||
if label_kind == "all_ignore":
|
||||
# Every row ignored -> zero loss and zero gradient (the reference is a
|
||||
# disconnected zero scalar, so there is nothing to differentiate).
|
||||
assert float(out.item()) == 0.0
|
||||
assert torch.count_nonzero(lq.grad) == 0
|
||||
return
|
||||
|
||||
lqr = lq0.clone().requires_grad_(True)
|
||||
ref = _ref_jsd(lqr, lp, labels, 0.5, ignore_index)
|
||||
ref.backward()
|
||||
_assert_close(out, ref, tols["atol_fwd"], tols["rtol_fwd"], f"jsd-{label_kind} fwd")
|
||||
_assert_close(lq.grad, lqr.grad, tols["atol_bwd"], tols["rtol_bwd"], f"jsd-{label_kind} dx")
|
||||
|
||||
|
||||
def test_jsd_non_unit_upstream_grad():
|
||||
ci = _cutile_impl("jsd")
|
||||
dtype = torch.float32
|
||||
tols = _tols("jsd", ci, dtype)
|
||||
lq0, lp = _jsd_inputs(48, 256, dtype, seed=13)
|
||||
scale = 2.75
|
||||
|
||||
lq = lq0.clone().requires_grad_(True)
|
||||
out = F.jsd(lq, lp, None, beta=0.5, ignore_index=-100, impl=ci)
|
||||
out.backward(torch.tensor(scale, device="cuda"))
|
||||
|
||||
lqr = lq0.clone().requires_grad_(True)
|
||||
ref = _ref_jsd(lqr, lp, None, 0.5, -100)
|
||||
ref.backward(torch.tensor(scale, device="cuda"))
|
||||
_assert_close(lq.grad, lqr.grad, tols["atol_bwd"], tols["rtol_bwd"], "jsd-nonunit-grad dx")
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# 5. Non-contiguous inputs for every public op
|
||||
# ===========================================================================
|
||||
def test_rms_norm_non_contiguous():
|
||||
ci = _cutile_impl("rms_norm")
|
||||
dtype = torch.float32
|
||||
g = torch.Generator(device="cpu").manual_seed(5)
|
||||
# Transpose to make the row dimension non-contiguous.
|
||||
base = torch.randn(1024, 256, generator=g).to("cuda", dtype)
|
||||
x = base.t().detach().requires_grad_(True) # (256, 1024), non-contiguous
|
||||
assert not x.is_contiguous()
|
||||
w = torch.randn(1024, generator=g).to("cuda", dtype).detach().requires_grad_(True)
|
||||
y = F.rms_norm(x, w, 1e-6, casting_mode="llama", impl=ci)
|
||||
yr = _ref_rms_norm(x.detach().contiguous(), w.detach(), 1e-6, 0.0, "llama")
|
||||
_assert_close(y, yr, 1e-4, 1e-4, "rms_norm-noncontig fwd")
|
||||
y.backward(torch.ones_like(y))
|
||||
assert torch.isfinite(x.grad).all() and torch.isfinite(w.grad).all()
|
||||
|
||||
|
||||
def test_layer_norm_non_contiguous():
|
||||
ci = _cutile_impl("layer_norm")
|
||||
dtype = torch.float32
|
||||
g = torch.Generator(device="cpu").manual_seed(5)
|
||||
base = torch.randn(1024, 256, generator=g).to("cuda", dtype)
|
||||
x = base.t().detach().requires_grad_(True)
|
||||
assert not x.is_contiguous()
|
||||
w = torch.randn(1024, generator=g).to("cuda", dtype).detach().requires_grad_(True)
|
||||
b = torch.randn(1024, generator=g).to("cuda", dtype).detach().requires_grad_(True)
|
||||
y = F.layer_norm(x, w, b, 1e-6, impl=ci)
|
||||
yr = _ref_layer_norm(x.detach().contiguous(), w.detach(), b.detach(), 1e-6)
|
||||
_assert_close(y, yr, 1e-3, 1e-3, "layer_norm-noncontig fwd")
|
||||
y.backward(torch.ones_like(y))
|
||||
assert torch.isfinite(x.grad).all()
|
||||
|
||||
|
||||
def test_softmax_non_contiguous():
|
||||
ci = _cutile_impl("softmax")
|
||||
dtype = torch.float32
|
||||
g = torch.Generator(device="cpu").manual_seed(5)
|
||||
base = torch.randn(1024, 256, generator=g).to("cuda", dtype)
|
||||
x = base.t().detach().requires_grad_(True) # (256, 1024)
|
||||
assert not x.is_contiguous()
|
||||
y = F.softmax(x, impl=ci)
|
||||
yr = torch.softmax(x.detach().float(), dim=-1)
|
||||
_assert_close(y, yr, 1e-4, 1e-4, "softmax-noncontig fwd")
|
||||
y.backward(torch.randn_like(y))
|
||||
assert torch.isfinite(x.grad).all()
|
||||
|
||||
|
||||
def test_jsd_non_contiguous():
|
||||
ci = _cutile_impl("jsd")
|
||||
dtype = torch.float32
|
||||
tols = _tols("jsd", ci, dtype)
|
||||
torch.manual_seed(17)
|
||||
BT, V = 256, 64
|
||||
# Build VALID log-probs first (row-wise log_softmax over the V last dim),
|
||||
# then take a strided row-slice to obtain a genuinely NON-contiguous VIEW
|
||||
# whose rows are still valid log-probs (each retained row is untouched).
|
||||
# The previous version re-ran log_softmax after transposing, which
|
||||
# materialised a fresh contiguous tensor and never exercised a
|
||||
# non-contiguous JSD input at all.
|
||||
lq_full = torch.randn(2 * BT, V, device="cuda", dtype=dtype).log_softmax(-1)
|
||||
lp_full = torch.randn(2 * BT, V, device="cuda", dtype=dtype).log_softmax(-1)
|
||||
lq = lq_full[::2].detach().requires_grad_(True) # (BT, V) rows stride 2V -> non-contig
|
||||
lp = lp_full[::2].detach()
|
||||
# Assert the exact layout that reaches the kernel, immediately before dispatch.
|
||||
assert not lq.is_contiguous() and not lp.is_contiguous()
|
||||
assert lq.shape == (BT, V)
|
||||
|
||||
out = F.jsd(lq, lp, None, beta=0.5, ignore_index=-100, impl=ci)
|
||||
lqr = lq.detach().clone().requires_grad_(True) # contiguous copy of the same values
|
||||
ref = _ref_jsd(lqr, lp, None, 0.5, -100)
|
||||
out.backward()
|
||||
ref.backward()
|
||||
_assert_close(out, ref, tols["atol_fwd"], tols["rtol_fwd"], "jsd-noncontig fwd")
|
||||
_assert_close(lq.grad, lqr.grad, tols["atol_bwd"], tols["rtol_bwd"], "jsd-noncontig dx")
|
||||
|
||||
|
||||
def test_fused_linear_jsd_non_contiguous():
|
||||
ci = _cutile_impl("fused_linear_jsd")
|
||||
ti = _triton_impl("fused_linear_jsd")
|
||||
dtype = torch.float32
|
||||
tols = _tols("fused_linear_jsd", ci, dtype)
|
||||
BT, H, V = 64, 128, 256
|
||||
g = torch.Generator(device="cpu").manual_seed(19)
|
||||
si_cpu = torch.randn(H, BT, generator=g) # transposed layout
|
||||
sw_cpu = torch.randn(V, H, generator=g)
|
||||
ti_cpu = torch.randn(BT, H, generator=g)
|
||||
tw_cpu = torch.randn(V, H, generator=g)
|
||||
|
||||
def run(impl):
|
||||
si = si_cpu.to("cuda", dtype).t().detach().requires_grad_(True) # (BT, H) non-contig
|
||||
# Assert the exact layout that reaches the fused kernel, before dispatch.
|
||||
assert not si.is_contiguous() and si.shape == (BT, H)
|
||||
sw = sw_cpu.to("cuda", dtype).detach().requires_grad_(True)
|
||||
t_in = ti_cpu.to("cuda", dtype).detach()
|
||||
tw = tw_cpu.to("cuda", dtype).detach()
|
||||
out = F.fused_linear_jsd(si, sw, t_in, tw, None, jsd_beta=0.5, impl=impl)
|
||||
out.backward()
|
||||
return out, si.grad, sw.grad
|
||||
|
||||
oc, dsic, dswc = run(ci)
|
||||
ot, dsit, dswt = run(ti)
|
||||
_assert_close(oc, ot, tols["atol_fwd"], tols["rtol_fwd"], "fljsd-noncontig fwd")
|
||||
_assert_close(dsic, dsit, tols["atol_bwd"], tols["rtol_bwd"], "fljsd-noncontig d_student_input")
|
||||
_assert_close(dswc, dswt, tols["atol_bwd"], tols["rtol_bwd"], "fljsd-noncontig d_student_weight")
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# 6. Fused-linear JSD semantics vs an INDEPENDENT torch reference
|
||||
# ===========================================================================
|
||||
# The parity-vs-Triton tests above only pin the inner cuTile primitive against
|
||||
# Triton (both share the outer chunked cuBLAS loop, so a shared bug would hide).
|
||||
# These cases build the whole fused_linear_jsd forward+backward from plain torch
|
||||
# ops in fp32 and differentiate through it, then compare the cuTile op against
|
||||
# that ground truth across: ignored labels, non-unit upstream gradient,
|
||||
# non-default temperature, asymmetric beta, and a multi-outer-chunk shape.
|
||||
def _ref_fused_linear_jsd(si_cpu, sw_cpu, t_in_cpu, tw_cpu, labels, beta, ignore_index, temperature, upstream):
|
||||
"""fp32 autograd ground truth for fused_linear_jsd -> (loss, d_si, d_sw)."""
|
||||
si = si_cpu.to("cuda", torch.float32).detach().requires_grad_(True)
|
||||
sw = sw_cpu.to("cuda", torch.float32).detach().requires_grad_(True)
|
||||
t_in = t_in_cpu.to("cuda", torch.float32).detach()
|
||||
tw = tw_cpu.to("cuda", torch.float32).detach()
|
||||
student_logits = (si @ sw.t()) / temperature
|
||||
teacher_logits = (t_in @ tw.t()) / temperature
|
||||
log_q = student_logits.log_softmax(-1)
|
||||
log_p = teacher_logits.log_softmax(-1)
|
||||
loss = _ref_jsd(log_q, log_p, labels, beta, ignore_index)
|
||||
loss.backward(upstream)
|
||||
return loss.detach(), si.grad, sw.grad
|
||||
|
||||
|
||||
def _run_cutile_fused(si_cpu, sw_cpu, t_in_cpu, tw_cpu, labels, beta, ignore_index, temperature, upstream, impl):
|
||||
si = si_cpu.to("cuda", torch.float32).detach().requires_grad_(True)
|
||||
sw = sw_cpu.to("cuda", torch.float32).detach().requires_grad_(True)
|
||||
t_in = t_in_cpu.to("cuda", torch.float32).detach()
|
||||
tw = tw_cpu.to("cuda", torch.float32).detach()
|
||||
out = F.fused_linear_jsd(
|
||||
si, sw, t_in, tw, labels, jsd_beta=beta, ignore_index=ignore_index, temperature=temperature, impl=impl
|
||||
)
|
||||
out.backward(upstream)
|
||||
return out.detach(), si.grad, sw.grad
|
||||
|
||||
|
||||
# (name, BT, H, V, beta, temperature, upstream_scale, label_kind)
|
||||
# multi_chunk uses V >> chunk_mem_const*H so the outer loop runs >1 chunk on
|
||||
# BOTH B200 (const 8) and B300 (const 16): V=4096 > 16*64.
|
||||
_FLJSD_SEMANTIC_CASES = [
|
||||
("ignored_labels", 32, 128, 512, 0.5, 1.0, 1.0, "scattered_ignore"),
|
||||
("non_unit_grad", 32, 128, 512, 0.5, 1.0, 2.75, "none"),
|
||||
("non_default_temp", 32, 128, 512, 0.5, 2.0, 1.0, "none"),
|
||||
("asymmetric_beta_low", 32, 128, 512, 0.1, 1.0, 1.0, "none"),
|
||||
("asymmetric_beta_high", 32, 128, 512, 0.9, 1.0, 1.0, "none"),
|
||||
("multi_chunk", 256, 64, 4096, 0.5, 1.0, 1.0, "none"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("case", _FLJSD_SEMANTIC_CASES, ids=[c[0] for c in _FLJSD_SEMANTIC_CASES])
|
||||
def test_fused_linear_jsd_semantics_vs_torch(case):
|
||||
name, BT, H, V, beta, temperature, scale, label_kind = case
|
||||
ci = _cutile_impl("fused_linear_jsd")
|
||||
dtype = torch.float32
|
||||
tols = _tols("fused_linear_jsd", ci, dtype)
|
||||
g = torch.Generator(device="cpu").manual_seed(hash(name) & 0xFFFF)
|
||||
si_cpu = torch.randn(BT, H, generator=g)
|
||||
sw_cpu = torch.randn(V, H, generator=g)
|
||||
ti_cpu = torch.randn(BT, H, generator=g)
|
||||
tw_cpu = torch.randn(V, H, generator=g)
|
||||
|
||||
if label_kind == "scattered_ignore":
|
||||
labels = torch.arange(BT, device="cuda", dtype=torch.int64) % V
|
||||
labels[::3] = -100
|
||||
else:
|
||||
labels = None
|
||||
upstream = torch.tensor(scale, device="cuda", dtype=torch.float32)
|
||||
|
||||
oc, dsic, dswc = _run_cutile_fused(si_cpu, sw_cpu, ti_cpu, tw_cpu, labels, beta, -100, temperature, upstream, ci)
|
||||
orf, dsir, dswr = _ref_fused_linear_jsd(si_cpu, sw_cpu, ti_cpu, tw_cpu, labels, beta, -100, temperature, upstream)
|
||||
_assert_close(oc, orf, tols["atol_fwd"], tols["rtol_fwd"], f"fljsd-sem/{name} fwd")
|
||||
_assert_close(dsic, dsir, tols["atol_bwd"], tols["rtol_bwd"], f"fljsd-sem/{name} d_student_input")
|
||||
_assert_close(dswc, dswr, tols["atol_bwd"], tols["rtol_bwd"], f"fljsd-sem/{name} d_student_weight")
|
||||
@@ -57,13 +57,20 @@ def test_layer_norm_correctness(backend, shape, dtype, with_bias):
|
||||
"""
|
||||
if backend == "__none__":
|
||||
pytest.skip("No layer_norm backends registered in this environment")
|
||||
assert_op_correctness(
|
||||
"layer_norm",
|
||||
backend,
|
||||
shape,
|
||||
dtype,
|
||||
extra={"include_bias": with_bias},
|
||||
)
|
||||
try:
|
||||
assert_op_correctness(
|
||||
"layer_norm",
|
||||
backend,
|
||||
shape,
|
||||
dtype,
|
||||
extra={"include_bias": with_bias},
|
||||
)
|
||||
except RuntimeError as e:
|
||||
# See test_rms_norm.py: honour a backend's documented hidden-dim range
|
||||
# limit (cuTile caps forward and backward at 8192) as a clean skip.
|
||||
if "only supports hidden dim" in str(e) or "out of range" in str(e).lower():
|
||||
pytest.skip(f"backend {backend} documents range limit: {e}")
|
||||
raise
|
||||
|
||||
|
||||
def test_layer_norm_explicit_backend_unavailable():
|
||||
|
||||
@@ -59,13 +59,22 @@ def test_rms_norm_correctness(backend, shape, dtype, casting_mode):
|
||||
"""
|
||||
if backend == "__none__":
|
||||
pytest.skip("No rms_norm backends registered in this environment")
|
||||
assert_op_correctness(
|
||||
"rms_norm",
|
||||
backend,
|
||||
shape,
|
||||
dtype,
|
||||
casting_mode=casting_mode,
|
||||
)
|
||||
try:
|
||||
assert_op_correctness(
|
||||
"rms_norm",
|
||||
backend,
|
||||
shape,
|
||||
dtype,
|
||||
casting_mode=casting_mode,
|
||||
)
|
||||
except RuntimeError as e:
|
||||
# Backends may document a hidden-dim range and reject wider rows in
|
||||
# forward *and* backward (e.g. cuTile caps at 8192). assert_op_correctness
|
||||
# only skips on the backward path; catch the forward-side rejection here
|
||||
# so the documented limit reads as a clean skip, not a failure.
|
||||
if "only supports hidden dim" in str(e) or "out of range" in str(e).lower():
|
||||
pytest.skip(f"backend {backend} documents range limit: {e}")
|
||||
raise
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", _REGISTERED_BACKENDS or ["__none__"])
|
||||
|
||||
Reference in New Issue
Block a user