From 33d76621cf20a23751dc1a2dcca2809f7cc0e2cb Mon Sep 17 00:00:00 2001 From: lcy-seso Date: Mon, 3 Aug 2026 15:51:32 +0800 Subject: [PATCH 01/15] [Docs][Ops] Dtype is read at forward, never taken by an op constructor MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An op today takes a `dtype` kwarg and the caller then hands it tensors that must agree. Two sources for one fact, and the constructor's copy is the one that can be wrong. Make the tensors the only source. Consequences recorded in the slot rules: - S12/S13: `dtype` is not a kwarg, and no kernel is built in `__init__` — a kernel is dtype-specialized and there is no dtype yet. `dispatch_kernel` stays, so an unsupported arch still fails at construction. - S16: the kernel cache is keyed by `(_cache_key(*shapes), dtype)`, so a second dtype builds a second kernel instead of reusing the first. - S19: `eval_roofline` reads attributes `forward()` binds, so it is post-forward only. - The fixed-rank / arbitrary-rank split now governs when shape inference runs, not when the kernel is built. Spec only; the op layer does not conform yet. --- docs/design/ops-design-reference.md | 72 +++++++++++++++-------------- docs/design/ops-design.md | 38 +++++++-------- 2 files changed, 58 insertions(+), 52 deletions(-) diff --git a/docs/design/ops-design-reference.md b/docs/design/ops-design-reference.md index a778bd192..276b78cbc 100644 --- a/docs/design/ops-design-reference.md +++ b/docs/design/ops-design-reference.md @@ -71,7 +71,7 @@ Slot-keyed rule dictionary consumed on demand by [ops-design.md](ops-design.md) ### Slot S7: Class docstring - **Rule.** One-sentence summary; `Args:` block enumerating every `__init__` kwarg (S12) with type and short description; optional `Example:` block. -- **Derivation.** `Args` block from manifest `signature.params` + `static_dims` + `dtype`. +- **Derivation.** `Args` block from manifest `signature.params` + `static_dims`. - **Example.** ```python class ExampleCumsumFwdOp(Op): @@ -82,18 +82,17 @@ Slot-keyed rule dictionary consumed on demand by [ops-design.md](ops-design.md) Args: M: Number of rows (product of all dims except the reduction axis). N: Hidden dimension (size along the reduction axis). - dtype: Data type (float32, float16, or bfloat16). dim: Reduction dimension (default -1). kernel_map: Optional override for kernel dispatch. tune: Whether to autotune (default False). """ ``` -- **Common mistakes.** Args out of sync with `__init__`; listing tensor inputs in `Args` (they belong to `forward`). +- **Common mistakes.** Args out of sync with `__init__`; listing tensor inputs in `Args` (they belong to `forward`); documenting a `dtype` kwarg (there is none — dtype comes from the input at `forward`). ### Slot S12: `__init__` signature -- **Rule.** Keyword-only via `*`. Kwarg block order: (1) `static_dims` entries in manifest key order, no defaults; (2) `dtype`; (3) `signature.params` entries in manifest key order; (4) `kernel_map` and `tune` last. -- **Derivation.** Manifest `static_dims` + `dtype` + `signature.params`. +- **Rule.** Keyword-only via `*`. Kwarg block order: (1) `static_dims` entries in manifest key order, no defaults; (2) `signature.params` entries in manifest key order; (3) `kernel_map` and `tune` last. **`dtype` is never a kwarg** — see [Parameter design](#parameter-design). +- **Derivation.** Manifest `static_dims` + `signature.params`. - **Example.** ```python def __init__( @@ -101,32 +100,31 @@ Slot-keyed rule dictionary consumed on demand by [ops-design.md](ops-design.md) *, M: int, N: int, - dtype: torch.dtype, dim: int = -1, kernel_map: Optional[Dict[str, Kernel]] = None, tune: bool = False, ): ``` -- **Common mistakes.** Missing `*` (positional accepted); `static_dims` kwargs with defaults; params/static_dims block order inverted; kwargs not backed by a manifest source. +- **Common mistakes.** Missing `*` (positional accepted); `static_dims` kwargs with defaults; params/static_dims block order inverted; kwargs not backed by a manifest source; accepting `dtype`, `in_dtype` or `out_dtype`. ### Slot S13: `__init__` body -- **Rule.** Body sequence: (a) `self. = ` per kwarg; (b) `self.dispatch_kernel(kernel_map)`; then branch by op shape: - - **Fully-static op** (all non-static axes committed at ctor): (c-static) `self.kernel = self.kernel_map[](...)` — kernel built once at init; (d-static) optionally precompute `self._infer_output_shapes(_shape=(...))` eagerly if a caller needs the output shapes before `forward()`. The `Op` base class does not currently consume an `_output_shapes` attribute — do not introduce one unless a concrete consumer requires it. - - **Arbitrary-rank op** (at least one axis unknown until forward): (c-dyn) initialise `self._kernel_cache: Dict[Hashable, Kernel] = {}` (the cache key follows `Op._cache_key`'s `Hashable` return type — often a tuple, but overrides may return `int` or other hashables) and defer kernel construction to `forward()` keyed by `self._cache_key(*input_shapes)`; (d-dyn) defer `_infer_output_shapes` to `forward()` per unique input shape. -- **Derivation.** Each `self.*` assignment mirrors one S12 kwarg. Kernel-build positional args follow the kernel class's ctor (kernel author's API). "Fully-static" iff every `signature.inputs` shape axis is either a manifest `shape` dim name or a `static_dims` key resolvable at ctor; otherwise arbitrary-rank and the deferred branch applies. +- **Rule.** Body sequence: (a) `self. = ` per kwarg; (b) `self.dispatch_kernel(kernel_map)`; (c) initialise `self._kernel_cache: Dict[Hashable, Kernel] = {}`. **No kernel is constructed here**, for any op shape: the kernel is dtype-specialized and no dtype is known until `forward()` receives a tensor. What the ctor does resolve is the kernel *class* and the architecture check, both of which `dispatch_kernel` performs without a tensor. + - A **fully-static op** (all non-static axes committed at ctor) may still precompute `self._infer_output_shapes(_shape=(...))` eagerly if a caller needs output shapes before `forward()` — shape inference is dtype-independent. The `Op` base class does not currently consume an `_output_shapes` attribute; do not introduce one unless a concrete consumer requires it. + - An **arbitrary-rank op** (at least one axis unknown until forward) defers `_infer_output_shapes` to `forward()` per unique input shape. + - The cache key follows `Op._cache_key`'s `Hashable` return type paired with the forward-time dtype: `(self._cache_key(*input_shapes), dtype)`. +- **Derivation.** Each `self.*` assignment mirrors one S12 kwarg. Kernel-build positional args follow the kernel class's ctor (kernel author's API). "Fully-static" iff every `signature.inputs` shape axis is either a manifest `shape` dim name or a `static_dims` key resolvable at ctor; the distinction now governs only when shape inference runs, not when the kernel is built. - **Example (arbitrary-rank).** ```python self.N = N - self.dtype = dtype self.dim = dim self.tune = tune self.dispatch_kernel(kernel_map) - # M unknown at init (only N committed via static_dims); kernel - # is built lazily in forward() once M is derived. + # M unknown at init (only N committed via static_dims), and no dtype + # is known at all; the kernel is built in forward() from both. self._kernel_cache: Dict[Hashable, Kernel] = {} ``` -- **Common mistakes.** `_infer_output_shapes` called before `dispatch_kernel`; hard-coding the kernel class instead of routing through `self.kernel_map`; building the kernel in `__init__` for an arbitrary-rank op (fails when a non-static axis value is required by the kernel ctor); omitting `self._kernel_cache` initialisation for the deferred branch (first forward-time cache lookup raises `AttributeError`). +- **Common mistakes.** `_infer_output_shapes` called before `dispatch_kernel`; hard-coding the kernel class instead of routing through `self.kernel_map`; building the kernel in `__init__` (there is no dtype to build it with); storing a `self.dtype` at ctor time; omitting `self._kernel_cache` initialisation (the first forward-time cache lookup raises `AttributeError`). ### Slot S14: `default_kernel_map` property @@ -152,8 +150,8 @@ Slot-keyed rule dictionary consumed on demand by [ops-design.md](ops-design.md) ### Slot S16: `forward` body -- **Rule.** Body sequence: (a) `self._validate_dtypes(...)`; (b) validate `shape_rules` (e.g. `-x.ndim <= dim < x.ndim`) and normalise parameter-dependent axes via modulo (e.g. `dim = self.dim % x.ndim`); (c) validate each `static_dims` commitment (`x.shape[] == self.`); (d) for arbitrary-rank ops, bind `self._static_axes = frozenset({(input_index, resolved_axis)})` and look up / lazily build the kernel in `self._kernel_cache` keyed by `self._cache_key(*input_shapes)`; (e) `.contiguous()` + reshape to the kernel's expected 2D layout; (f) call the kernel; (g) restore the original shape. Fully-static ops skip the cache-lookup part of (d) since `self.kernel` was built at init. -- **Derivation.** Validation expressions come from each `static_dims` entry's `.shape[]` RHS; axis normalisation mirrors the param evaluation in `static_dims` + `shape_rules`; kernel cache key is whatever `_cache_key` projects (default: tuple of non-static-axis sizes). A kernel that pads internally returns the semantic shape, so the op does not trim. +- **Rule.** Body sequence: (a) `self._validate_dtypes(...)`; (b) validate `shape_rules` (e.g. `-x.ndim <= dim < x.ndim`) and normalise parameter-dependent axes via modulo (e.g. `dim = self.dim % x.ndim`); (c) validate each `static_dims` commitment (`x.shape[] == self.`); (d) for arbitrary-rank ops, bind `self._static_axes = frozenset({(input_index, resolved_axis)})`, then — for every op shape — look up / lazily build the kernel in `self._kernel_cache` keyed by `(self._cache_key(*input_shapes), .dtype)`; (e) `.contiguous()` + reshape to the kernel's expected 2D layout; (f) call the kernel; (g) restore the original shape. +- **Derivation.** Validation expressions come from each `static_dims` entry's `.shape[]` RHS; axis normalisation mirrors the param evaluation in `static_dims` + `shape_rules`; kernel cache key is whatever `_cache_key` projects (default: tuple of non-static-axis sizes), paired with the dtype of the dtype-defining input. A kernel that pads internally returns the semantic shape, so the op does not trim. - **Example (arbitrary-rank).** ```python self._validate_dtypes(x) @@ -172,10 +170,11 @@ Slot-keyed rule dictionary consumed on demand by [ops-design.md](ops-design.md) self.M = M # default _cache_key projects non-static axes; override for coarser # keying when kernel math permits (see Optional Hooks appendix). - key = self._cache_key(x.shape) + self.dtype = x.dtype + key = (self._cache_key(x.shape), x.dtype) if key not in self._kernel_cache: self._kernel_cache[key] = self.kernel_map["example_cumsum_fwd"]( - M, self.N, "sum", self.dtype, tune=self.tune + M, self.N, "sum", x.dtype, tune=self.tune ) kernel = self._kernel_cache[key] orig_shape = x.shape @@ -184,7 +183,7 @@ Slot-keyed rule dictionary consumed on demand by [ops-design.md](ops-design.md) y = y2.reshape(*orig_shape[:dim], *orig_shape[dim + 1 :], self.N) return y.movedim(-1, dim) ``` -- **Common mistakes.** Skipping `_validate_dtypes`; reshape before `.contiguous()`; hard-coding `x.shape[-1]` instead of the normalised `x.shape[self.dim % x.ndim]`; binding `self._static_axes` before the axis is non-negative (violates `Op._static_axes` contract); forgetting the kernel cache lookup so every forward rebuilds the kernel; trimming padded kernel output in the op instead of leaving it to the kernel; not restoring the original shape. +- **Common mistakes.** Skipping `_validate_dtypes`; keying the kernel cache on shape alone, so a second dtype silently reuses the first dtype's kernel; reshape before `.contiguous()`; hard-coding `x.shape[-1]` instead of the normalised `x.shape[self.dim % x.ndim]`; binding `self._static_axes` before the axis is non-negative (violates `Op._static_axes` contract); forgetting the kernel cache lookup so every forward rebuilds the kernel; trimming padded kernel output in the op instead of leaving it to the kernel; not restoring the original shape. ### Slot S17: `_infer_output_shapes` method body @@ -211,7 +210,7 @@ Slot-keyed rule dictionary consumed on demand by [ops-design.md](ops-design.md) ### Slot S19: `eval_roofline` method body -- **Rule.** Codegen emits a complete plain-Python body reading `self.*` attributes. Per [`roofline.md` §4.4.6](roofline.md#446-evaluator-surface-boundary) (Evaluator Surface Boundary) there is NO shared AST evaluator on L1 and NO class-level roofline expression strings (e.g. `_flops_str`, `_bytes_str`, `_roofline_vars`) that would be parsed at runtime. L1 stub raises `NotImplementedError` (FIXME staged-rollout). +- **Rule.** Codegen emits a complete plain-Python body reading `self.*` attributes, which `forward()` binds — `self.dtype` among them. `eval_roofline` is therefore defined only after at least one `forward()`. Per [`roofline.md` §4.4.6](roofline.md#446-evaluator-surface-boundary) (Evaluator Surface Boundary) there is NO shared AST evaluator on L1 and NO class-level roofline expression strings (e.g. `_flops_str`, `_bytes_str`, `_roofline_vars`) that would be parsed at runtime. L1 stub raises `NotImplementedError` (FIXME staged-rollout). - **Derivation.** Manifest `roofline.vars`, `roofline.flops`, `roofline.bytes`; see [`roofline.md` §4.4](roofline.md#44-op-codegen). - **Example.** ```python @@ -220,7 +219,7 @@ Slot-keyed rule dictionary consumed on demand by [ops-design.md](ops-design.md) bytes_ = (2 * self.M * self.N + self.N) * self.dtype.itemsize return flops, bytes_ ``` -- **Common mistakes.** Class-level roofline expression strings parsed at runtime (prohibited by §4.4.6); any `ast.parse` or shared `_safe_eval` path; returning `float` or `numpy` types (contract is `tuple[int, int]`). +- **Common mistakes.** Class-level roofline expression strings parsed at runtime (prohibited by §4.4.6); any `ast.parse` or shared `_safe_eval` path; returning `float` or `numpy` types (contract is `tuple[int, int]`); assuming `self.dtype` is set on a freshly-constructed op. ### Slot S20: Package `__init__.py` registration @@ -283,7 +282,7 @@ Per-family protocol variables, declared by L2 bases and overridden by L3 ops. | -------------- | ------------------------------------ | -------------------------------------------------------------------------------------------- | | `kernel` | `Kernel` | Kernel instance used by `forward()` | | `kernel_map` | `Optional[Dict[str, Kernel]]` | Dispatched kernels keyed by name | -| `dtype` | `Optional[torch.dtype]` | Computation dtype | +| `dtype` | `Optional[torch.dtype]` | Dtype of the most recent `forward()`; `None` before the first one | | `device` | `Optional[Union[torch.device, str]]` | Device (default `'cuda'`) | | `input_shapes` | `Optional[list[tuple]]` | Expected input tensor shapes (for introspection and non-runtime consumers) | | `_static_axes` | `frozenset[tuple[int, int]]` | Static axes as `(input_index, axis)` pairs (default `frozenset()`); consumed by `_cache_key` | @@ -292,9 +291,11 @@ Abstract interface: `default_kernel_map` (property), `forward()`. Manifest-drive ### `Kernel` base class attributes ([`tileops/kernels/kernel_base.py`](../../tileops/kernels/kernel_base.py)) +Unlike `Op`, a `Kernel` **is** constructed for one element type — it compiles a dtype-specialized program, so `dtype` is a ctor argument here. The op supplies it from the tensors at `forward()`. + | Attribute | Type | Purpose | | ------------------ | ----------------------- | ---------------------------------------------- | -| `dtype` | `Optional[torch.dtype]` | Data type | +| `dtype` | `Optional[torch.dtype]` | Element type this kernel is specialized for | | `config` | `Dict[str, Any]` | Tile configuration (block sizes, stages, etc.) | | `autotune_configs` | `Optional[list[dict]]` | Search space for autotuning | | `supported_archs` | `Optional[list[int]]` | GPU SM versions (e.g., `[80, 86, 89, 90]`) | @@ -340,21 +341,24 @@ The manifest ([`tileops/manifest/`](../../tileops/manifest/)) is the sole source ### Parameter design -Three time points: (1) manifest — constraint structure; (2) `__init__` — user commits `static_dims` values; (3) `forward` — shapes concrete, commitments validated. See [manifest.md § `static_dims`](manifest.md#static_dims). +Three time points: (1) manifest — constraint structure; (2) `__init__` — user commits `static_dims` values; (3) `forward` — shapes concrete, commitments validated, dtype read from the tensors. See [manifest.md § `static_dims`](manifest.md#static_dims). + +**Dtype belongs to time point 3, never to 2.** The tensors carry it, so requiring the caller to restate it at construction only creates a second source that can disagree with the first. Constructing an op therefore commits to shape structure and nothing about element type. -| | Fixed-rank op | Arbitrary-rank op | -| ------------------------ | ----------------------- | ------------------------------------------------------------------ | -| Manifest has `shape` | yes | no | -| `__init__` shape source | `shape` dimension names | `static_dims` | -| Undeclared dimensions | none | derived from tensor at forward time | -| Kernel construction time | init (all dims known) | init (`static_dims` known) or forward (first encounter, cached) | -| Forward cache keying | N/A (single kernel) | `_cache_key(*input_shapes)` — default non-static axes, overridable | +| | Fixed-rank op | Arbitrary-rank op | +| ------------------------ | ----------------------- | --------------------------------------------------------------------------- | +| Manifest has `shape` | yes | no | +| `__init__` shape source | `shape` dimension names | `static_dims` | +| Undeclared dimensions | none | derived from tensor at forward time | +| Kernel construction time | forward (first call) | forward (first encounter) | +| Forward cache keying | dtype | `(_cache_key(*input_shapes), dtype)` — default non-static axes, overridable | ### Calling conventions - **Fully static op:** `_infer_output_shapes` called once in `__init__`, result stored as an instance attribute. -- **Op with dynamic dims:** `_infer_output_shapes` called in `forward()` once dynamic dims resolve; kernel construction cached by `_cache_key(*input_shapes)`. -- **`_validate_dtypes`:** runs on every `forward()` call. +- **Op with dynamic dims:** `_infer_output_shapes` called in `forward()` once dynamic dims resolve. +- **Kernel construction:** always in `forward()`, cached by `(_cache_key(*input_shapes), dtype)`. +- **`_validate_dtypes`:** runs on every `forward()` call, and is the only place an op rejects a dtype. - **Non-runtime consumers** (validator, graph compiler): call `_infer_output_shapes` with concrete shape tuples without constructing tensors. Roofline consumers use interfaces in [`roofline.md`](roofline.md). ### Inheritance in family-base hierarchies diff --git a/docs/design/ops-design.md b/docs/design/ops-design.md index 14b106528..c12555b15 100644 --- a/docs/design/ops-design.md +++ b/docs/design/ops-design.md @@ -22,12 +22,16 @@ Op ← L1: thin base, shared by all ops **Do it at the first moment all required information is known, do it once, cache the result.** -| Op category | When all info is known | Behaviour | -| -------------- | ----------------------------------------------------------- | -------------------------------------------------------- | -| Fixed-rank | `__init__` (all dims provided) | `_infer_output_shapes` runs once at init. | -| Arbitrary-rank | `__init__` for `static_dims`; `forward` for everything else | Kernel built on first encounter, cached by `_cache_key`. | +| Op category | When all info is known | Behaviour | +| -------------- | ----------------------------------------------------------- | ----------------------------------------- | +| Fixed-rank | `__init__` (all dims provided) | `_infer_output_shapes` runs once at init. | +| Arbitrary-rank | `__init__` for `static_dims`; `forward` for everything else | `_infer_output_shapes` runs per shape. | -`_validate_dtypes` runs on every `forward()` call — dtype validity depends on the actual tensors passed, not just their shapes. Roofline timing and formula semantics are defined in [roofline.md](roofline.md). See [Parameter Design](ops-design-reference.md#parameter-design) for fixed-rank vs arbitrary-rank details and [Codegen Details](ops-design-reference.md#codegen) for calling conventions. +**Dtype is never a constructor parameter.** An op reads it from the input tensors in `forward()`. A caller who passes fp16 tensors gets the fp16 kernel without having said so twice, and an op can no longer be constructed in a state that disagrees with the tensors it is about to be handed. + +The kernel is dtype-specialized, so this makes kernel construction uniformly deferred to the first `forward()` — for fixed-rank and arbitrary-rank ops alike — and the kernel cache is keyed by shape *and* dtype. `dispatch_kernel()` stays in `__init__`: resolving the kernel *class* and checking the architecture needs no tensor, and keeping it there preserves fast failure on an unsupported GPU. + +`_validate_dtypes` runs on every `forward()` call — dtype validity depends on the actual tensors passed, not just their shapes. It is the only dtype gate; an op does not compare an incoming tensor against a dtype it was constructed with, because there is no such dtype. Roofline timing and formula semantics are defined in [roofline.md](roofline.md). See [Parameter Design](ops-design-reference.md#parameter-design) for fixed-rank vs arbitrary-rank details and [Codegen Details](ops-design-reference.md#codegen) for calling conventions. ## Scaffolding an Op from a Manifest Entry @@ -47,7 +51,7 @@ Provides: """ import math -from typing import Dict, Optional +from typing import Dict, Hashable, Optional import torch @@ -63,7 +67,7 @@ from ..op_base import Op ### Step 2: Class declaration + docstring + `__all__` -**Input.** Manifest entry key (= class name); `signature.inputs`, `signature.params`, `static_dims`, per-tensor `dtype` (Args block content). +**Input.** Manifest entry key (= class name); `signature.inputs`, `signature.params`, `static_dims` (Args block content). **Output.** @@ -79,7 +83,6 @@ class ExampleCumsumFwdOp(Op): Args: N: Hidden dimension (size along the reduction axis), committed at ctor via ``static_dims: N: "x.shape[dim]"``. - dtype: Data type (float32, float16, or bfloat16). dim: Reduction dimension (default -1). kernel_map: Optional override for kernel dispatch. tune: Whether to autotune (default False). @@ -92,7 +95,7 @@ class ExampleCumsumFwdOp(Op): ### Step 3: `_static_axes` + `__init__` signature and body -**Input.** `static_dims` (literal-axis → class-level `_static_axes` frozenset; param-axis → empty class-level default, bind at `forward()` after `dim % x.ndim` normalization); `signature.params`; `dtype`. +**Input.** `static_dims` (literal-axis → class-level `_static_axes` frozenset; param-axis → empty class-level default, bind at `forward()` after `dim % x.ndim` normalization); `signature.params`. **Output.** @@ -108,22 +111,20 @@ class ExampleCumsumFwdOp(Op): self, *, N: int, - dtype: torch.dtype, dim: int = -1, kernel_map: Optional[Dict[str, Kernel]] = None, tune: bool = False, ): self.N = N - self.dtype = dtype self.dim = dim self.tune = tune self.dispatch_kernel(kernel_map) # M is not a static_dim — deferred to forward() where x.ndim # is known and M is derived from the non-reduction axes. - self._kernel_cache: Dict[tuple, Kernel] = {} + self._kernel_cache: Dict[Hashable, Kernel] = {} ``` -**Validation.** Every `__init__` kwarg has a manifest source (`static_dims` or `signature.params` or `dtype`); no extras except `kernel_map` / `tune`. In particular, `M` is NOT a ctor kwarg — `ExampleCumsumFwdOp.static_dims` declares only `N`, so `M` is derived at forward time. Keyword-only via `*`, no defaults on `static_dims` entries. `_static_axes` matches the manifest axis form (literal-int axis → populated class-level frozenset; param-dependent axis → empty class-level default, bound at forward after `dim % x.ndim` normalization). +**Validation.** Every `__init__` kwarg has a manifest source (`static_dims` or `signature.params`); no extras except `kernel_map` / `tune`. `dtype` is not a kwarg — it is read from the input in `forward()`. In particular, `M` is NOT a ctor kwarg — `ExampleCumsumFwdOp.static_dims` declares only `N`, so `M` is derived at forward time. Keyword-only via `*`, no defaults on `static_dims` entries. `_static_axes` matches the manifest axis form (literal-int axis → populated class-level frozenset; param-dependent axis → empty class-level default, bound at forward after `dim % x.ndim` normalization). **Reference.** [Slot S21](ops-design-reference.md#slot-s21), [S12](ops-design-reference.md#slot-s12), [S13](ops-design-reference.md#slot-s13). @@ -155,13 +156,14 @@ class ExampleCumsumFwdOp(Op): f"got {x.shape[dim]}") # Bind _static_axes now that the concrete axis is known. self._static_axes = frozenset({(0, dim)}) - # Derive M (product of non-reduction dims) and cache kernel by (M,). + # Derive M (product of non-reduction dims); cache by shape and dtype. M = math.prod(s for i, s in enumerate(x.shape) if i != dim) self.M = M # stored for eval_roofline - key = (M,) + self.dtype = x.dtype # ditto; the op commits to no dtype before this + key = ((M,), x.dtype) if key not in self._kernel_cache: self._kernel_cache[key] = self.kernel_map["example_cumsum_fwd"]( - M, self.N, "sum", self.dtype, tune=self.tune) + M, self.N, "sum", x.dtype, tune=self.tune) kernel = self._kernel_cache[key] # Move reduction axis to last, reshape to (M, N), compute, restore. orig_shape = x.shape @@ -171,7 +173,7 @@ class ExampleCumsumFwdOp(Op): return y.movedim(-1, dim) ``` -**Validation.** `default_kernel_map` keys / values match manifest `source.kernel_map` verbatim. `forward` calls `self._validate_dtypes(...)` first (not inline dtype comparisons — that is Step 5's job). Every `static_dims` commitment is validated against the actual tensor shape at the normalized axis before the kernel is called. `_static_axes` is bound from the normalized (non-negative) axis before the kernel cache lookup. The op never trims kernel output: a kernel that pads internally returns the semantic shape. +**Validation.** `default_kernel_map` keys / values match manifest `source.kernel_map` verbatim. `forward` calls `self._validate_dtypes(...)` first (not inline dtype comparisons — that is Step 5's job). The kernel is built from `x.dtype`, and the cache key carries that dtype so a second call with a different dtype builds a second kernel rather than reusing the first. Every `static_dims` commitment is validated against the actual tensor shape at the normalized axis before the kernel is called. `_static_axes` is bound from the normalized (non-negative) axis before the kernel cache lookup. The op never trims kernel output: a kernel that pads internally returns the semantic shape. **Reference.** [Slot S14](ops-design-reference.md#slot-s14), [S15](ops-design-reference.md#slot-s15), [S16](ops-design-reference.md#slot-s16). @@ -213,7 +215,7 @@ class ExampleCumsumFwdOp(Op): return flops, bytes_ ``` -**Validation.** The body is **plain Python** reading `self.*` attributes. No class-level roofline expression strings, no `ast.parse`, no shared L1 evaluator — prohibited by [`roofline.md §4.4.6` Evaluator Surface Boundary](roofline.md#446-evaluator-surface-boundary). Return type is `tuple[int, int]`, not `float` or `numpy`. Expressions derive directly from `roofline.vars` bindings + `roofline.flops` + `roofline.bytes`; see [`roofline.md §4.4` Op Codegen](roofline.md#44-op-codegen). +**Validation.** The body is **plain Python** reading `self.*` attributes. Those attributes — `self.M` and `self.dtype` here — are bound by `forward()`, so `eval_roofline` is callable only after at least one forward; there is no ctor-time dtype to read. No class-level roofline expression strings, no `ast.parse`, no shared L1 evaluator — prohibited by [`roofline.md §4.4.6` Evaluator Surface Boundary](roofline.md#446-evaluator-surface-boundary). Return type is `tuple[int, int]`, not `float` or `numpy`. Expressions derive directly from `roofline.vars` bindings + `roofline.flops` + `roofline.bytes`; see [`roofline.md §4.4` Op Codegen](roofline.md#44-op-codegen). **Reference.** [Slot S19](ops-design-reference.md#slot-s19). From 8f17e4a941a8e4f81a9225f2499f55e0619d2ae1 Mon Sep 17 00:00:00 2001 From: lcy-seso Date: Mon, 3 Aug 2026 21:34:48 +0800 Subject: [PATCH 02/15] [Refactor][Ops] Read dtype at forward for the reduction and norm families MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 20 op constructors stop taking `dtype`; each reads it from its input in `forward()` and keys its kernel cache by shape and dtype. Also removes the `_committed_dtype` dual path (an optional ctor dtype, validated when supplied) that cumulative, softmax and the four ada/fused-add norm ops carried. `LayerNormFwdOp` gains a real kernel cache — it kept one kernel and rebuilt whenever the leading-dims product changed, which with a second dtype in play would rebuild every call. `_dtype_codegen` now emits the validator body unrolled per input instead of looping over `input_names` through `locals()`. Dynamo cannot trace `locals()`, and this body runs inside `forward()`, so an op that validates dtypes would lose `fullgraph`. The unreachable `same_as(ref)`-not-supplied branch goes with it: synthesis already rejects a ref that names no sibling input. `_SoftmaxBaseOp` keeps its inline dtype-union check rather than delegating to `_validate_dtypes`: `LogSumExpFwdOp`'s manifest entry declares only float16|bfloat16 while the kernel and its tests use float32. Narrowing the code to match is the wrong direction, and widening the manifest is a separate change. Two tests asserted that a ctor dtype disagreeing with the input raises. That coupling is gone, but the invariant underneath is not — the `dim=[]` short-circuit must still gate dtype — so they now feed a dtype outside the manifest union instead of one outside the ctor's. --- benchmarks/ops/bench_ada_layer_norm.py | 4 +- benchmarks/ops/bench_argreduce.py | 4 +- benchmarks/ops/bench_logical_reduce.py | 6 +- benchmarks/ops/bench_norm.py | 8 +- benchmarks/ops/bench_reduce.py | 16 ++-- benchmarks/ops/bench_softmax.py | 2 +- benchmarks/ops/bench_vector_norm.py | 6 +- tests/ops/test_ada_layer_norm.py | 4 +- tests/ops/test_ada_layer_norm_zero.py | 4 +- tests/ops/test_argreduce.py | 48 +++++------ tests/ops/test_cumulative.py | 30 +++---- tests/ops/test_fused_add_layer_norm.py | 6 +- tests/ops/test_fused_add_rms_norm.py | 6 +- tests/ops/test_layer_norm.py | 10 +-- tests/ops/test_logical_reduce.py | 79 +++++++++---------- tests/ops/test_normalization_alignment.py | 8 +- tests/ops/test_reduce.py | 78 +++++++++--------- .../ops/test_reduce_arithmetic_conformance.py | 6 +- tests/ops/test_reduce_boolean_conformance.py | 8 +- tests/ops/test_reduce_dim_none.py | 30 +++---- tests/ops/test_reduce_multidim.py | 62 +++++++-------- tests/ops/test_reduce_scalar_conformance.py | 14 ++-- tests/ops/test_reduce_variance_conformance.py | 12 +-- tests/ops/test_reduction_defaults.py | 56 ++++++------- tests/ops/test_reduction_scalar_input.py | 28 +++---- tests/ops/test_rms_norm.py | 6 +- tests/ops/test_softmax.py | 10 +-- tests/ops/test_vector_norm.py | 75 +++++++++--------- tests/ops/test_welford_non_aligned.py | 18 ++--- tileops/ops/_dtype_codegen.py | 71 +++++++++-------- tileops/ops/norm/ada_layer_norm.py | 12 +-- tileops/ops/norm/ada_layer_norm_zero.py | 12 +-- tileops/ops/norm/fused_add_layer_norm.py | 10 +-- tileops/ops/norm/fused_add_rms_norm.py | 10 +-- tileops/ops/norm/layer_norm.py | 30 ++++--- tileops/ops/norm/rms_norm.py | 29 ++++--- tileops/ops/reduction/argreduce.py | 8 +- tileops/ops/reduction/cumulative.py | 14 +--- tileops/ops/reduction/logical_reduce.py | 11 +-- tileops/ops/reduction/reduce.py | 38 ++++----- tileops/ops/reduction/softmax.py | 12 +-- tileops/ops/reduction/vector_norm.py | 9 +-- 42 files changed, 421 insertions(+), 489 deletions(-) diff --git a/benchmarks/ops/bench_ada_layer_norm.py b/benchmarks/ops/bench_ada_layer_norm.py index 5a5b7b6fb..2b2549cc0 100644 --- a/benchmarks/ops/bench_ada_layer_norm.py +++ b/benchmarks/ops/bench_ada_layer_norm.py @@ -29,7 +29,7 @@ def test_ada_layer_norm_bench(m: int, n: int, dtype: torch.dtype) -> None: test = AdaLayerNormWorkload(m, n, dtype) inputs = test.gen_inputs() - op = AdaLayerNormFwdOp(dtype=dtype) + op = AdaLayerNormFwdOp() bm = ManifestBenchmark(_ADA_OP_NAME, op, test) result = bm.profile(op, *inputs) BenchmarkReport.record(op, locals(), result, tag="tileops") @@ -48,7 +48,7 @@ def test_ada_layer_norm_zero_bench(m: int, n: int, dtype: torch.dtype) -> None: test = AdaLayerNormZeroWorkload(m, n, dtype) inputs = test.gen_inputs() - op = AdaLayerNormZeroFwdOp(dtype=dtype) + op = AdaLayerNormZeroFwdOp() bm = ManifestBenchmark(_ADA_ZERO_OP_NAME, op, test) result = bm.profile(op, *inputs) BenchmarkReport.record(op, locals(), result, tag="tileops") diff --git a/benchmarks/ops/bench_argreduce.py b/benchmarks/ops/bench_argreduce.py index c99ee1520..9c1834243 100644 --- a/benchmarks/ops/bench_argreduce.py +++ b/benchmarks/ops/bench_argreduce.py @@ -39,7 +39,7 @@ def test_argmax_bench(shape: tuple, dtype: torch.dtype, extra: dict) -> None: workload = ArgmaxWorkload(shape, dtype) inputs = workload.gen_inputs() - op = ArgmaxFwdOp(dtype=dtype, **extra) + op = ArgmaxFwdOp(**extra) bm = ManifestBenchmark(_ARGMAX_OP, op, workload) # FIXME(staged-rollout): ArgreduceKernel skips large-N manifest workloads # @@ -72,7 +72,7 @@ def test_argmin_bench(shape: tuple, dtype: torch.dtype, extra: dict) -> None: workload = ArgminWorkload(shape, dtype) inputs = workload.gen_inputs() - op = ArgminFwdOp(dtype=dtype, **extra) + op = ArgminFwdOp(**extra) bm = ManifestBenchmark(_ARGMIN_OP, op, workload) # FIXME(staged-rollout): ArgreduceKernel skips large-N manifest workloads # diff --git a/benchmarks/ops/bench_logical_reduce.py b/benchmarks/ops/bench_logical_reduce.py index 85f59ca8e..94efc3a8e 100644 --- a/benchmarks/ops/bench_logical_reduce.py +++ b/benchmarks/ops/bench_logical_reduce.py @@ -32,7 +32,7 @@ def test_any_bench( inputs = test.gen_inputs() op_params.setdefault("dim", -1) - op = AnyFwdOp(dtype=dtype, **op_params) + op = AnyFwdOp(**op_params) bm = ManifestBenchmark(_ANY_OP, op, test) try: result = bm.profile(op, *inputs) @@ -66,7 +66,7 @@ def test_all_bench( inputs = test.gen_inputs() op_params.setdefault("dim", -1) - op = AllFwdOp(dtype=dtype, **op_params) + op = AllFwdOp(**op_params) bm = ManifestBenchmark(_ALL_OP, op, test) try: result = bm.profile(op, *inputs) @@ -100,7 +100,7 @@ def test_count_nonzero_bench( inputs = test.gen_inputs() op_params.setdefault("dim", -1) - op = CountNonzeroFwdOp(dtype=dtype, **op_params) + op = CountNonzeroFwdOp(**op_params) bm = ManifestBenchmark(_COUNT_NONZERO_OP, op, test) try: result = bm.profile(op, *inputs) diff --git a/benchmarks/ops/bench_norm.py b/benchmarks/ops/bench_norm.py index ae013493c..92aab6788 100644 --- a/benchmarks/ops/bench_norm.py +++ b/benchmarks/ops/bench_norm.py @@ -48,7 +48,7 @@ def test_rms_norm_bench(m: int, n: int, dtype: torch.dtype, tune: bool) -> None: test = RMSNormTestBaseline(m, n, dtype) inputs = test.gen_inputs() - op = RMSNormFwdOp(normalized_shape=(n,), dtype=dtype, tune=tune) + op = RMSNormFwdOp(normalized_shape=(n,), tune=tune) bm = ManifestBenchmark(_RMS_OP_NAME, op, test) result = bm.profile(op, *inputs) BenchmarkReport.record(op, locals(), result, tag="tileops") @@ -85,7 +85,7 @@ def test_fused_add_rms_norm_bench(m: int, n: int, dtype: torch.dtype, tune: bool test = FusedAddRMSNormWorkload(m, n, dtype) inputs = test.gen_inputs() - op = FusedAddRMSNormFwdOp(dtype=dtype, tune=tune) + op = FusedAddRMSNormFwdOp(tune=tune) bm = ManifestBenchmark(_FUSED_RMS_OP_NAME, op, test) result = bm.profile(op, *inputs) BenchmarkReport.record(op, locals(), result, tag="tileops") @@ -121,7 +121,7 @@ def test_layer_norm_bench(m: int, n: int, dtype: torch.dtype, tune: bool) -> Non test = LayerNormWorkload(m, n, dtype) inputs = test.gen_inputs() - op = LayerNormFwdOp(normalized_shape=(n,), dtype=dtype, tune=tune) + op = LayerNormFwdOp(normalized_shape=(n,), tune=tune) bm = ManifestBenchmark(_LN_OP_NAME, op, test) result = bm.profile(op, *inputs) BenchmarkReport.record(op, locals(), result, tag="tileops") @@ -154,7 +154,7 @@ def test_fused_add_layer_norm_bench(m: int, n: int, dtype: torch.dtype, tune: bo test = FusedAddLayerNormWorkload(m, n, dtype) inputs = test.gen_inputs() - op = FusedAddLayerNormFwdOp(dtype=dtype, tune=tune) + op = FusedAddLayerNormFwdOp(tune=tune) bm = ManifestBenchmark(_FUSED_LN_OP_NAME, op, test) result = bm.profile(op, *inputs) BenchmarkReport.record(op, locals(), result, tag="tileops") diff --git a/benchmarks/ops/bench_reduce.py b/benchmarks/ops/bench_reduce.py index 5bf9e8742..c72167755 100644 --- a/benchmarks/ops/bench_reduce.py +++ b/benchmarks/ops/bench_reduce.py @@ -55,7 +55,7 @@ def test_sum_bench( inputs = test.gen_inputs() op_params.setdefault("dim", -1) # baseline below reduces dim=-1 - op = SumFwdOp(dtype=dtype, **op_params) + op = SumFwdOp(**op_params) bm = ManifestBenchmark(_SUM_OP, op, test) try: result = bm.profile(op, *inputs) @@ -89,7 +89,7 @@ def test_mean_bench( inputs = test.gen_inputs() op_params.setdefault("dim", -1) # baseline below mirrors the op's dim - op = MeanFwdOp(dtype=dtype, **op_params) + op = MeanFwdOp(**op_params) bm = ManifestBenchmark(_MEAN_OP, op, test) try: result = bm.profile(op, *inputs) @@ -123,7 +123,7 @@ def test_amax_bench( inputs = test.gen_inputs() op_params.setdefault("dim", -1) - op = AmaxFwdOp(dtype=dtype, **op_params) + op = AmaxFwdOp(**op_params) bm = ManifestBenchmark(_AMAX_OP, op, test) try: result = bm.profile(op, *inputs) @@ -157,7 +157,7 @@ def test_amin_bench( inputs = test.gen_inputs() op_params.setdefault("dim", -1) - op = AminFwdOp(dtype=dtype, **op_params) + op = AminFwdOp(**op_params) bm = ManifestBenchmark(_AMIN_OP, op, test) try: result = bm.profile(op, *inputs) @@ -185,7 +185,7 @@ def test_prod_bench(shape: tuple, dtype: torch.dtype) -> None: test = ProdWorkload(shape, dtype) inputs = test.gen_inputs() - op = ProdFwdOp(dtype=dtype) + op = ProdFwdOp() bm = ManifestBenchmark(_PROD_OP, op, test) try: result = bm.profile(op, *inputs) @@ -216,7 +216,7 @@ def test_std_bench( inputs = test.gen_inputs() op_params.setdefault("dim", -1) - op = StdFwdOp(dtype=dtype, correction=1, **op_params) + op = StdFwdOp(correction=1, **op_params) bm = ManifestBenchmark(_STD_OP, op, test) try: result = bm.profile(op, *inputs) @@ -250,7 +250,7 @@ def test_var_bench( inputs = test.gen_inputs() op_params.setdefault("dim", -1) - op = VarFwdOp(dtype=dtype, correction=1, **op_params) + op = VarFwdOp(correction=1, **op_params) bm = ManifestBenchmark(_VAR_OP, op, test) try: result = bm.profile(op, *inputs) @@ -284,7 +284,7 @@ def test_var_mean_bench( inputs = test.gen_inputs() op_params.setdefault("dim", -1) - op = VarMeanFwdOp(dtype=dtype, correction=1, **op_params) + op = VarMeanFwdOp(correction=1, **op_params) bm = ManifestBenchmark(_VAR_MEAN_OP, op, test) try: result = bm.profile(op, *inputs) diff --git a/benchmarks/ops/bench_softmax.py b/benchmarks/ops/bench_softmax.py index 083b1c731..a86ce92f5 100644 --- a/benchmarks/ops/bench_softmax.py +++ b/benchmarks/ops/bench_softmax.py @@ -87,7 +87,7 @@ def test_logsumexp_bench( inputs = test.gen_inputs() op_params.setdefault("dim", -1) - op = LogSumExpFwdOp(dtype=dtype, tune=True, **op_params) + op = LogSumExpFwdOp(tune=True, **op_params) bm = ManifestBenchmark(_LOGSUMEXP_OP, op, test) try: result = bm.profile(op, *inputs) diff --git a/benchmarks/ops/bench_vector_norm.py b/benchmarks/ops/bench_vector_norm.py index 93b34a511..3a3e9a2e7 100644 --- a/benchmarks/ops/bench_vector_norm.py +++ b/benchmarks/ops/bench_vector_norm.py @@ -32,7 +32,7 @@ def test_l1_norm_bench( inputs = test.gen_inputs() op_params.setdefault("dim", -1) - op = L1NormFwdOp(dtype=dtype, **op_params) + op = L1NormFwdOp(**op_params) bm = ManifestBenchmark(_L1_NORM_OP, op, test) try: result = bm.profile(op, *inputs) @@ -68,7 +68,7 @@ def test_l2_norm_bench( inputs = test.gen_inputs() op_params.setdefault("dim", -1) - op = L2NormFwdOp(dtype=dtype, **op_params) + op = L2NormFwdOp(**op_params) bm = ManifestBenchmark(_L2_NORM_OP, op, test) try: result = bm.profile(op, *inputs) @@ -104,7 +104,7 @@ def test_inf_norm_bench( inputs = test.gen_inputs() op_params.setdefault("dim", -1) - op = InfNormFwdOp(dtype=dtype, **op_params) + op = InfNormFwdOp(**op_params) bm = ManifestBenchmark(_INF_NORM_OP, op, test) try: result = bm.profile(op, *inputs) diff --git a/tests/ops/test_ada_layer_norm.py b/tests/ops/test_ada_layer_norm.py index 28ba074f9..68645ca12 100644 --- a/tests/ops/test_ada_layer_norm.py +++ b/tests/ops/test_ada_layer_norm.py @@ -65,7 +65,7 @@ def _get_tolerances(dtype: torch.dtype) -> tuple[float, float]: @AdaLayerNormFixture def test_ada_layer_norm_op(m: int, n: int, dtype: torch.dtype) -> None: test = AdaLayerNormTest(m, n, dtype) - op = AdaLayerNormFwdOp(dtype=dtype) + op = AdaLayerNormFwdOp() atol, rtol = _get_tolerances(dtype) test.check(op, *test.gen_inputs(), atol=atol, rtol=rtol) @@ -183,7 +183,7 @@ def test_ada_layer_norm_3d(batch: int, seq: int, hidden: int, dtype: torch.dtype scale = torch.randn(batch, seq, hidden, dtype=dtype, device="cuda") shift = torch.randn(batch, seq, hidden, dtype=dtype, device="cuda") - op = AdaLayerNormFwdOp(dtype=dtype) + op = AdaLayerNormFwdOp() # Reference: scale * LayerNorm(x) + shift eps = 1e-5 diff --git a/tests/ops/test_ada_layer_norm_zero.py b/tests/ops/test_ada_layer_norm_zero.py index 4bfdb4180..1382ed36f 100644 --- a/tests/ops/test_ada_layer_norm_zero.py +++ b/tests/ops/test_ada_layer_norm_zero.py @@ -66,7 +66,7 @@ def _get_tolerances(dtype: torch.dtype) -> tuple[float, float]: @AdaLayerNormZeroFixture def test_ada_layer_norm_zero_op(m: int, n: int, dtype: torch.dtype) -> None: test = AdaLayerNormZeroTest(m, n, dtype) - op = AdaLayerNormZeroFwdOp(dtype=dtype) + op = AdaLayerNormZeroFwdOp() atol, rtol = _get_tolerances(dtype) test.check(op, *test.gen_inputs(), atol=atol, rtol=rtol) @@ -130,7 +130,7 @@ def test_ada_layer_norm_zero_3d(batch: int, seq: int, hidden: int, dtype: torch. shift = torch.randn(batch, seq, hidden, dtype=dtype, device="cuda") gate = torch.randn(batch, seq, hidden, dtype=dtype, device="cuda") - op = AdaLayerNormZeroFwdOp(dtype=dtype) + op = AdaLayerNormZeroFwdOp() # Reference: gate * (scale * LayerNorm(x) + shift) eps = 1e-5 diff --git a/tests/ops/test_argreduce.py b/tests/ops/test_argreduce.py index f3421aef2..30215a20a 100644 --- a/tests/ops/test_argreduce.py +++ b/tests/ops/test_argreduce.py @@ -181,7 +181,7 @@ def test_argmax_op(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.argreduce import ArgmaxFwdOp test = ArgreduceTest(m, n, dtype, "argmax") - op = ArgmaxFwdOp(dtype=dtype, dim=-1) + op = ArgmaxFwdOp(dim=-1) test.check(op, *test.gen_inputs(), compare=_exact_compare) @@ -191,7 +191,7 @@ def test_argmax_non_contiguous(m: int, n: int, dtype: torch.dtype) -> None: x_full = torch.randn(m, n * 2, dtype=dtype, device="cuda") x = x_full[:, :n] - op = ArgmaxFwdOp(dtype=dtype, dim=-1) + op = ArgmaxFwdOp(dim=-1) ref = x.contiguous().argmax(dim=-1) y = _call(op, x) assert y.dtype == torch.int64 @@ -203,7 +203,7 @@ def test_argmax_3d(batch: int, seq: int, hidden: int, dtype: torch.dtype) -> Non from tileops.ops.reduction.argreduce import ArgmaxFwdOp x = torch.randn(batch, seq, hidden, dtype=dtype, device="cuda") - op = ArgmaxFwdOp(dtype=dtype, dim=-1) + op = ArgmaxFwdOp(dim=-1) ref = x.argmax(dim=-1) y = _call(op, x) assert y.dtype == torch.int64 @@ -215,7 +215,7 @@ def test_argmax_4d(b0: int, b1: int, b2: int, n: int, dtype: torch.dtype) -> Non from tileops.ops.reduction.argreduce import ArgmaxFwdOp x = torch.randn(b0, b1, b2, n, dtype=dtype, device="cuda") - op = ArgmaxFwdOp(dtype=dtype, dim=-1) + op = ArgmaxFwdOp(dim=-1) ref = x.argmax(dim=-1) y = _call(op, x) assert y.dtype == torch.int64 @@ -227,7 +227,7 @@ def test_argmax_1d(n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.argreduce import ArgmaxFwdOp x = torch.randn(n, dtype=dtype, device="cuda") - op = ArgmaxFwdOp(dtype=dtype, dim=-1) + op = ArgmaxFwdOp(dim=-1) ref = x.argmax(dim=-1) y = _call(op, x) assert y.dtype == torch.int64 @@ -240,7 +240,7 @@ def test_argmax_3d_dim0(batch: int, seq: int, hidden: int, dtype: torch.dtype) - from tileops.ops.reduction.argreduce import ArgmaxFwdOp x = torch.randn(batch, seq, hidden, dtype=dtype, device="cuda") - op = ArgmaxFwdOp(dtype=dtype, dim=0) + op = ArgmaxFwdOp(dim=0) ref = x.argmax(dim=0) y = _call(op, x) assert y.shape == ref.shape, f"shape mismatch: {y.shape} vs {ref.shape}" @@ -254,7 +254,7 @@ def test_argmax_3d_dim0_keepdim(batch: int, seq: int, hidden: int, dtype: torch. from tileops.ops.reduction.argreduce import ArgmaxFwdOp x = torch.randn(batch, seq, hidden, dtype=dtype, device="cuda") - op = ArgmaxFwdOp(dtype=dtype, dim=0, keepdim=True) + op = ArgmaxFwdOp(dim=0, keepdim=True) ref = x.argmax(dim=0, keepdim=True) y = _call(op, x) assert y.shape == ref.shape, f"shape mismatch: {y.shape} vs {ref.shape}" @@ -268,7 +268,7 @@ def test_argmax_4d_dim0(b0: int, b1: int, b2: int, n: int, dtype: torch.dtype) - from tileops.ops.reduction.argreduce import ArgmaxFwdOp x = torch.randn(b0, b1, b2, n, dtype=dtype, device="cuda") - op = ArgmaxFwdOp(dtype=dtype, dim=0) + op = ArgmaxFwdOp(dim=0) ref = x.argmax(dim=0) y = _call(op, x) assert y.shape == ref.shape, f"shape mismatch: {y.shape} vs {ref.shape}" @@ -282,7 +282,7 @@ def test_argmax_4d_dim0_keepdim(b0: int, b1: int, b2: int, n: int, dtype: torch. from tileops.ops.reduction.argreduce import ArgmaxFwdOp x = torch.randn(b0, b1, b2, n, dtype=dtype, device="cuda") - op = ArgmaxFwdOp(dtype=dtype, dim=0, keepdim=True) + op = ArgmaxFwdOp(dim=0, keepdim=True) ref = x.argmax(dim=0, keepdim=True) y = _call(op, x) assert y.shape == ref.shape, f"shape mismatch: {y.shape} vs {ref.shape}" @@ -296,7 +296,7 @@ def test_argmax_spec_dim(shape: tuple, dim: int, keepdim: bool, dtype: torch.dty from tileops.ops.reduction.argreduce import ArgmaxFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = ArgmaxFwdOp(dtype=dtype, dim=dim, keepdim=keepdim) + op = ArgmaxFwdOp(dim=dim, keepdim=keepdim) ref = x.argmax(dim=dim, keepdim=keepdim) y = _call(op, x) assert y.shape == ref.shape, f"shape mismatch: {y.shape} vs {ref.shape}" @@ -312,7 +312,7 @@ def test_argmin_op(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.argreduce import ArgminFwdOp test = ArgreduceTest(m, n, dtype, "argmin") - op = ArgminFwdOp(dtype=dtype, dim=-1) + op = ArgminFwdOp(dim=-1) test.check(op, *test.gen_inputs(), compare=_exact_compare) @@ -322,7 +322,7 @@ def test_argmin_non_contiguous(m: int, n: int, dtype: torch.dtype) -> None: x_full = torch.randn(m, n * 2, dtype=dtype, device="cuda") x = x_full[:, :n] - op = ArgminFwdOp(dtype=dtype, dim=-1) + op = ArgminFwdOp(dim=-1) ref = x.contiguous().argmin(dim=-1) y = _call(op, x) assert y.dtype == torch.int64 @@ -334,7 +334,7 @@ def test_argmin_3d(batch: int, seq: int, hidden: int, dtype: torch.dtype) -> Non from tileops.ops.reduction.argreduce import ArgminFwdOp x = torch.randn(batch, seq, hidden, dtype=dtype, device="cuda") - op = ArgminFwdOp(dtype=dtype, dim=-1) + op = ArgminFwdOp(dim=-1) ref = x.argmin(dim=-1) y = _call(op, x) assert y.dtype == torch.int64 @@ -346,7 +346,7 @@ def test_argmin_4d(b0: int, b1: int, b2: int, n: int, dtype: torch.dtype) -> Non from tileops.ops.reduction.argreduce import ArgminFwdOp x = torch.randn(b0, b1, b2, n, dtype=dtype, device="cuda") - op = ArgminFwdOp(dtype=dtype, dim=-1) + op = ArgminFwdOp(dim=-1) ref = x.argmin(dim=-1) y = _call(op, x) assert y.dtype == torch.int64 @@ -358,7 +358,7 @@ def test_argmin_1d(n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.argreduce import ArgminFwdOp x = torch.randn(n, dtype=dtype, device="cuda") - op = ArgminFwdOp(dtype=dtype, dim=-1) + op = ArgminFwdOp(dim=-1) ref = x.argmin(dim=-1) y = _call(op, x) assert y.dtype == torch.int64 @@ -371,7 +371,7 @@ def test_argmin_3d_dim0(batch: int, seq: int, hidden: int, dtype: torch.dtype) - from tileops.ops.reduction.argreduce import ArgminFwdOp x = torch.randn(batch, seq, hidden, dtype=dtype, device="cuda") - op = ArgminFwdOp(dtype=dtype, dim=0) + op = ArgminFwdOp(dim=0) ref = x.argmin(dim=0) y = _call(op, x) assert y.shape == ref.shape, f"shape mismatch: {y.shape} vs {ref.shape}" @@ -385,7 +385,7 @@ def test_argmin_3d_dim0_keepdim(batch: int, seq: int, hidden: int, dtype: torch. from tileops.ops.reduction.argreduce import ArgminFwdOp x = torch.randn(batch, seq, hidden, dtype=dtype, device="cuda") - op = ArgminFwdOp(dtype=dtype, dim=0, keepdim=True) + op = ArgminFwdOp(dim=0, keepdim=True) ref = x.argmin(dim=0, keepdim=True) y = _call(op, x) assert y.shape == ref.shape, f"shape mismatch: {y.shape} vs {ref.shape}" @@ -399,7 +399,7 @@ def test_argmin_4d_dim0(b0: int, b1: int, b2: int, n: int, dtype: torch.dtype) - from tileops.ops.reduction.argreduce import ArgminFwdOp x = torch.randn(b0, b1, b2, n, dtype=dtype, device="cuda") - op = ArgminFwdOp(dtype=dtype, dim=0) + op = ArgminFwdOp(dim=0) ref = x.argmin(dim=0) y = _call(op, x) assert y.shape == ref.shape, f"shape mismatch: {y.shape} vs {ref.shape}" @@ -413,7 +413,7 @@ def test_argmin_4d_dim0_keepdim(b0: int, b1: int, b2: int, n: int, dtype: torch. from tileops.ops.reduction.argreduce import ArgminFwdOp x = torch.randn(b0, b1, b2, n, dtype=dtype, device="cuda") - op = ArgminFwdOp(dtype=dtype, dim=0, keepdim=True) + op = ArgminFwdOp(dim=0, keepdim=True) ref = x.argmin(dim=0, keepdim=True) y = _call(op, x) assert y.shape == ref.shape, f"shape mismatch: {y.shape} vs {ref.shape}" @@ -427,7 +427,7 @@ def test_argmin_spec_dim(shape: tuple, dim: int, keepdim: bool, dtype: torch.dty from tileops.ops.reduction.argreduce import ArgminFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = ArgminFwdOp(dtype=dtype, dim=dim, keepdim=keepdim) + op = ArgminFwdOp(dim=dim, keepdim=keepdim) ref = x.argmin(dim=dim, keepdim=keepdim) y = _call(op, x) assert y.shape == ref.shape, f"shape mismatch: {y.shape} vs {ref.shape}" @@ -498,12 +498,12 @@ def test_argmax_dim_none(shape: tuple, dtype: torch.dtype) -> None: x = torch.randn(*shape, dtype=dtype, device="cuda") ref_flat = torch.argmax(x) - y = _call(ArgmaxFwdOp(dtype=dtype, dim=None), x) + y = _call(ArgmaxFwdOp(dim=None), x) assert y.dtype == torch.int64 assert y.shape == ref_flat.shape, f"shape mismatch: {y.shape} vs {ref_flat.shape}" assert torch.equal(y, ref_flat), f"dim=None argmax mismatch on shape={shape} dtype={dtype}" - y_keep = _call(ArgmaxFwdOp(dtype=dtype, dim=None, keepdim=True), x) + y_keep = _call(ArgmaxFwdOp(dim=None, keepdim=True), x) expected_shape = tuple(1 for _ in shape) assert y_keep.shape == expected_shape, f"keepdim shape mismatch: {y_keep.shape} vs {expected_shape}" assert torch.equal(y_keep.reshape(()), ref_flat), ( @@ -519,12 +519,12 @@ def test_argmin_dim_none(shape: tuple, dtype: torch.dtype) -> None: x = torch.randn(*shape, dtype=dtype, device="cuda") ref_flat = torch.argmin(x) - y = _call(ArgminFwdOp(dtype=dtype, dim=None), x) + y = _call(ArgminFwdOp(dim=None), x) assert y.dtype == torch.int64 assert y.shape == ref_flat.shape, f"shape mismatch: {y.shape} vs {ref_flat.shape}" assert torch.equal(y, ref_flat), f"dim=None argmin mismatch on shape={shape} dtype={dtype}" - y_keep = _call(ArgminFwdOp(dtype=dtype, dim=None, keepdim=True), x) + y_keep = _call(ArgminFwdOp(dim=None, keepdim=True), x) expected_shape = tuple(1 for _ in shape) assert y_keep.shape == expected_shape, f"keepdim shape mismatch: {y_keep.shape} vs {expected_shape}" assert torch.equal(y_keep.reshape(()), ref_flat), ( diff --git a/tests/ops/test_cumulative.py b/tests/ops/test_cumulative.py index a5b7d9010..383a68ec5 100644 --- a/tests/ops/test_cumulative.py +++ b/tests/ops/test_cumulative.py @@ -122,7 +122,7 @@ def test_cumsum_op(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.cumulative import CumsumFwdOp test = CumulativeTest((m, n), dtype, "cumsum") - op = CumsumFwdOp(dtype=dtype) + op = CumsumFwdOp() test.check(op, *test.gen_inputs(), **_tol(dtype)) @@ -132,7 +132,7 @@ def test_cumsum_non_contiguous(m: int, n: int, dtype: torch.dtype) -> None: x_full = torch.randn(m, n * 2, dtype=dtype, device="cuda") x = x_full[:, :n] - op = CumsumFwdOp(dtype=dtype) + op = CumsumFwdOp() ref = x.contiguous().float().cumsum(dim=-1).to(dtype) y = op(x) tol = _tol(dtype) @@ -144,7 +144,7 @@ def test_cumsum_3d(batch: int, seq: int, hidden: int, dtype: torch.dtype) -> Non from tileops.ops.reduction.cumulative import CumsumFwdOp x = torch.randn(batch, seq, hidden, dtype=dtype, device="cuda") - op = CumsumFwdOp(dtype=dtype) + op = CumsumFwdOp() ref = x.float().cumsum(dim=-1).to(dtype) y = op(x) tol = _tol(dtype) @@ -156,7 +156,7 @@ def test_cumsum_4d(b0: int, b1: int, b2: int, n: int, dtype: torch.dtype) -> Non from tileops.ops.reduction.cumulative import CumsumFwdOp x = torch.randn(b0, b1, b2, n, dtype=dtype, device="cuda") - op = CumsumFwdOp(dtype=dtype) + op = CumsumFwdOp() ref = x.float().cumsum(dim=-1).to(dtype) y = op(x) tol = _tol(dtype) @@ -168,7 +168,7 @@ def test_cumsum_1d(n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.cumulative import CumsumFwdOp x = torch.randn(n, dtype=dtype, device="cuda") - op = CumsumFwdOp(dtype=dtype) + op = CumsumFwdOp() ref = x.float().cumsum(dim=-1).to(dtype) y = op(x) tol = _tol(dtype) @@ -199,7 +199,7 @@ def test_cumprod_op(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.cumulative import CumprodFwdOp test = CumulativeTest((m, n), dtype, "cumprod", use_small_range=True) - op = CumprodFwdOp(dtype=dtype) + op = CumprodFwdOp() test.check(op, *test.gen_inputs(), **_cumprod_tol(dtype)) @@ -209,7 +209,7 @@ def test_cumprod_non_contiguous(m: int, n: int, dtype: torch.dtype) -> None: x_full = torch.rand(m, n * 2, dtype=dtype, device="cuda") * 0.01 + 0.99 x = x_full[:, :n] - op = CumprodFwdOp(dtype=dtype) + op = CumprodFwdOp() ref = x.contiguous().float().cumprod(dim=-1).to(dtype) y = op(x) tol = _cumprod_tol(dtype) @@ -221,7 +221,7 @@ def test_cumprod_3d(batch: int, seq: int, hidden: int, dtype: torch.dtype) -> No from tileops.ops.reduction.cumulative import CumprodFwdOp x = torch.rand(batch, seq, hidden, dtype=dtype, device="cuda") * 0.01 + 0.99 - op = CumprodFwdOp(dtype=dtype) + op = CumprodFwdOp() ref = x.float().cumprod(dim=-1).to(dtype) y = op(x) tol = _cumprod_tol(dtype) @@ -233,7 +233,7 @@ def test_cumprod_4d(b0: int, b1: int, b2: int, n: int, dtype: torch.dtype) -> No from tileops.ops.reduction.cumulative import CumprodFwdOp x = torch.rand(b0, b1, b2, n, dtype=dtype, device="cuda") * 0.01 + 0.99 - op = CumprodFwdOp(dtype=dtype) + op = CumprodFwdOp() ref = x.float().cumprod(dim=-1).to(dtype) y = op(x) tol = _cumprod_tol(dtype) @@ -245,7 +245,7 @@ def test_cumprod_1d(n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.cumulative import CumprodFwdOp x = torch.rand(n, dtype=dtype, device="cuda") * 0.01 + 0.99 - op = CumprodFwdOp(dtype=dtype) + op = CumprodFwdOp() ref = x.float().cumprod(dim=-1).to(dtype) y = op(x) tol = _cumprod_tol(dtype) @@ -269,7 +269,7 @@ def test_cumsum_dim_axis1( from tileops.ops.reduction.cumulative import CumsumFwdOp x = torch.randn(batch, hidden, seq, dtype=dtype, device="cuda") - op = CumsumFwdOp(dtype=dtype, dim=1) + op = CumsumFwdOp(dim=1) ref = x.float().cumsum(dim=1).to(dtype) y = op(x) atol = 1e-2 if dtype == torch.float16 else 1.6e-2 @@ -286,7 +286,7 @@ def test_cumprod_dim_axis1( # Values close to 1 to avoid over/underflow in cumprod over hidden dim. x = torch.rand(batch, hidden, seq, dtype=dtype, device="cuda") * 0.01 + 0.99 - op = CumprodFwdOp(dtype=dtype, dim=1) + op = CumprodFwdOp(dim=1) ref = x.float().cumprod(dim=1).to(dtype) y = op(x) tol = _cumprod_tol(dtype) @@ -311,7 +311,7 @@ def test_cumsum_backend_dispatch(M: int, N: int, dtype: torch.dtype, parallel: b from tileops.ops.reduction.cumulative import CumsumFwdOp x = torch.randn(M, N, dtype=dtype, device="cuda") - op = CumsumFwdOp(dtype=dtype, dim=-1) + op = CumsumFwdOp(dim=-1) y = op(x) ref = x.float().cumsum(dim=-1).to(dtype) @@ -333,7 +333,7 @@ def test_cumsum_parallel_scan_row_ownership(M: int, N: int) -> None: row_values = torch.arange(1, M + 1, dtype=torch.float32, device="cuda").unsqueeze(1) x = row_values.expand(-1, N).contiguous() - y = CumsumFwdOp(dtype=torch.float32, dim=-1)(x) + y = CumsumFwdOp(dim=-1)(x) # Row r holds the constant r + 1, so its cumsum is (r + 1) * [1, ..., N]. expected = row_values * torch.arange(1, N + 1, dtype=torch.float32, device="cuda") @@ -362,7 +362,7 @@ def test_cumsum_compile_fullgraph_warm_cache(M: int, N: int, dtype: torch.dtype) """ from tileops.ops.reduction.cumulative import CumsumFwdOp - op = CumsumFwdOp(dtype=dtype, dim=-1) + op = CumsumFwdOp(dim=-1) x = torch.randn(M, N, dtype=dtype, device="cuda") op(x) diff --git a/tests/ops/test_fused_add_layer_norm.py b/tests/ops/test_fused_add_layer_norm.py index 113e88073..6bdf6c454 100644 --- a/tests/ops/test_fused_add_layer_norm.py +++ b/tests/ops/test_fused_add_layer_norm.py @@ -63,7 +63,7 @@ def _get_tolerances(dtype: torch.dtype) -> tuple[float, float]: @FusedAddLayerNormFixture def test_fused_add_layer_norm_op(m: int, n: int, dtype: torch.dtype, tune: bool) -> None: test = FusedAddLayerNormTest(m, n, dtype) - op = FusedAddLayerNormFwdOp(dtype=dtype, tune=tune) + op = FusedAddLayerNormFwdOp(tune=tune) atol, rtol = _get_tolerances(dtype) test.check(op, *test.gen_inputs(), atol=atol, rtol=rtol) @@ -88,7 +88,7 @@ def test_fused_add_layer_norm_non_contiguous(m: int, n: int, dtype: torch.dtype) weight = torch.randn(n, dtype=dtype, device="cuda") bias = torch.randn(n, dtype=dtype, device="cuda") - op = FusedAddLayerNormFwdOp(M=m, N=n, dtype=dtype) + op = FusedAddLayerNormFwdOp(M=m, N=n) # Reference on contiguous copies test = FusedAddLayerNormTest(m, n, dtype) @@ -121,7 +121,7 @@ def test_fused_add_layer_norm_3d(batch: int, seq: int, hidden: int, dtype: torch bias = torch.randn(hidden, dtype=dtype, device="cuda") M = batch * seq - op = FusedAddLayerNormFwdOp(dtype=dtype) + op = FusedAddLayerNormFwdOp() test = FusedAddLayerNormTest(M, hidden, dtype) y_ref, add_ref = test.ref_program(x, residual, weight, bias) diff --git a/tests/ops/test_fused_add_rms_norm.py b/tests/ops/test_fused_add_rms_norm.py index d54b67664..7f05a5c14 100644 --- a/tests/ops/test_fused_add_rms_norm.py +++ b/tests/ops/test_fused_add_rms_norm.py @@ -51,7 +51,7 @@ def _get_tolerances(dtype: torch.dtype) -> tuple[float, float]: @FusedAddRMSNormFixture def test_fused_add_rms_norm_op(m: int, n: int, dtype: torch.dtype, tune: bool) -> None: test = FusedAddRMSNormTest(m, n, dtype) - op = FusedAddRMSNormFwdOp(dtype=dtype, tune=tune) + op = FusedAddRMSNormFwdOp(tune=tune) atol, rtol = _get_tolerances(dtype) test.check(op, *test.gen_inputs(), atol=atol, rtol=rtol) @@ -74,7 +74,7 @@ def test_fused_add_rms_norm_non_contiguous(m: int, n: int, dtype: torch.dtype) - residual = r_full[:, :n] weight = torch.randn(n, dtype=dtype, device="cuda") - op = FusedAddRMSNormFwdOp(M=m, N=n, dtype=dtype) + op = FusedAddRMSNormFwdOp(M=m, N=n) # Reference on contiguous copies test = FusedAddRMSNormTest(m, n, dtype) @@ -105,7 +105,7 @@ def test_fused_add_rms_norm_3d(batch: int, seq: int, hidden: int, dtype: torch.d weight = torch.randn(hidden, dtype=dtype, device="cuda") M = batch * seq - op = FusedAddRMSNormFwdOp(dtype=dtype) + op = FusedAddRMSNormFwdOp() test = FusedAddRMSNormTest(M, hidden, dtype) y_ref, add_ref = test.ref_program(x, residual, weight) diff --git a/tests/ops/test_layer_norm.py b/tests/ops/test_layer_norm.py index 5f82f3f33..22fac180c 100644 --- a/tests/ops/test_layer_norm.py +++ b/tests/ops/test_layer_norm.py @@ -61,7 +61,7 @@ def _get_tolerances(dtype: torch.dtype) -> tuple[float, float]: @LayerNormFixture def test_layer_norm_op(m: int, n: int, dtype: torch.dtype, tune: bool) -> None: test = LayerNormTest(m, n, dtype) - op = LayerNormFwdOp(normalized_shape=(n,), dtype=dtype) + op = LayerNormFwdOp(normalized_shape=(n,)) atol, rtol = _get_tolerances(dtype) test.check(op, *test.gen_inputs(), atol=atol, rtol=rtol) @@ -101,7 +101,7 @@ def test_layer_norm_non_contiguous(m: int, n: int, dtype: torch.dtype) -> None: weight = torch.randn(n, dtype=dtype, device="cuda") bias = torch.randn(n, dtype=dtype, device="cuda") - op = LayerNormFwdOp(normalized_shape=(n,), dtype=dtype) + op = LayerNormFwdOp(normalized_shape=(n,)) # Reference using torch.nn.functional.layer_norm x_ref = x.contiguous() @@ -133,7 +133,7 @@ def test_layer_norm_3d(batch: int, seq: int, hidden: int, dtype: torch.dtype) -> weight = torch.randn(hidden, dtype=dtype, device="cuda") bias = torch.randn(hidden, dtype=dtype, device="cuda") - op = LayerNormFwdOp(normalized_shape=(hidden,), dtype=dtype) + op = LayerNormFwdOp(normalized_shape=(hidden,)) # Reference using torch.nn.functional.layer_norm y_ref = F.layer_norm( @@ -176,7 +176,7 @@ def test_layer_norm_large_offset(m: int, n: int, dtype: torch.dtype) -> None: weight = torch.ones(n, dtype=dtype, device="cuda") bias = torch.zeros(n, dtype=dtype, device="cuda") - op = LayerNormFwdOp(normalized_shape=(n,), dtype=dtype) + op = LayerNormFwdOp(normalized_shape=(n,)) y_ref = F.layer_norm( x.float(), (n,), @@ -209,7 +209,7 @@ def test_layer_norm_rebuilds_kernel_on_m_change() -> None: n = 4096 dtype = torch.float16 - op = LayerNormFwdOp(normalized_shape=(n,), dtype=dtype) + op = LayerNormFwdOp(normalized_shape=(n,)) weight = torch.randn(n, dtype=dtype, device="cuda") bias = torch.randn(n, dtype=dtype, device="cuda") diff --git a/tests/ops/test_logical_reduce.py b/tests/ops/test_logical_reduce.py index d1815474f..b15374c99 100644 --- a/tests/ops/test_logical_reduce.py +++ b/tests/ops/test_logical_reduce.py @@ -219,7 +219,7 @@ def test_any_op(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.logical_reduce import AnyFwdOp test = LogicalReduceTest(m, n, dtype, "any") - op = AnyFwdOp(dtype=dtype, dim=-1) + op = AnyFwdOp(dim=-1) test.check(op, *test.gen_inputs(), compare=_exact_compare) @@ -229,7 +229,7 @@ def test_any_non_contiguous(m: int, n: int, dtype: torch.dtype) -> None: x_full = _make_noncontig_input(m, n, dtype) x = x_full[:, :n] - op = AnyFwdOp(dtype=dtype, dim=-1) + op = AnyFwdOp(dim=-1) ref = x.contiguous().bool().any(dim=-1) y = op(x) assert y.dtype == torch.bool @@ -241,7 +241,7 @@ def test_any_3d(batch: int, seq: int, hidden: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.logical_reduce import AnyFwdOp x = torch.randn(batch, seq, hidden, dtype=dtype, device="cuda") - op = AnyFwdOp(dtype=dtype, dim=-1) + op = AnyFwdOp(dim=-1) ref = x.bool().any(dim=-1) y = op(x) assert y.dtype == torch.bool @@ -253,7 +253,7 @@ def test_any_4d(b0: int, b1: int, b2: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.logical_reduce import AnyFwdOp x = torch.randn(b0, b1, b2, n, dtype=dtype, device="cuda") - op = AnyFwdOp(dtype=dtype, dim=-1) + op = AnyFwdOp(dim=-1) ref = x.bool().any(dim=-1) y = op(x) assert y.dtype == torch.bool @@ -265,7 +265,7 @@ def test_any_1d(n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.logical_reduce import AnyFwdOp x = _make_1d_input(n, dtype) - op = AnyFwdOp(dtype=dtype, dim=-1) + op = AnyFwdOp(dim=-1) ref = x.bool().any(dim=-1) y = op(x) assert y.dtype == torch.bool @@ -277,7 +277,7 @@ def test_any_dim(shape: tuple, dim: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.logical_reduce import AnyFwdOp x = _make_nd_input(shape, dtype) - op = AnyFwdOp(dtype=dtype, dim=dim) + op = AnyFwdOp(dim=dim) ref = x.bool().any(dim=dim) y = op(x) assert y.dtype == torch.bool @@ -290,7 +290,7 @@ def test_any_keepdim(shape: tuple, dim: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.logical_reduce import AnyFwdOp x = _make_nd_input(shape, dtype) - op = AnyFwdOp(dtype=dtype, dim=dim, keepdim=True) + op = AnyFwdOp(dim=dim, keepdim=True) ref = x.bool().any(dim=dim, keepdim=True) y = op(x) assert y.dtype == torch.bool @@ -306,7 +306,7 @@ def test_all_op(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.logical_reduce import AllFwdOp test = LogicalReduceTest(m, n, dtype, "all") - op = AllFwdOp(dtype=dtype, dim=-1) + op = AllFwdOp(dim=-1) test.check(op, *test.gen_inputs(), compare=_exact_compare) @@ -316,7 +316,7 @@ def test_all_non_contiguous(m: int, n: int, dtype: torch.dtype) -> None: x_full = _make_noncontig_input(m, n, dtype) x = x_full[:, :n] - op = AllFwdOp(dtype=dtype, dim=-1) + op = AllFwdOp(dim=-1) ref = x.contiguous().bool().all(dim=-1) y = op(x) assert y.dtype == torch.bool @@ -328,7 +328,7 @@ def test_all_3d(batch: int, seq: int, hidden: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.logical_reduce import AllFwdOp x = torch.randn(batch, seq, hidden, dtype=dtype, device="cuda") - op = AllFwdOp(dtype=dtype, dim=-1) + op = AllFwdOp(dim=-1) ref = x.bool().all(dim=-1) y = op(x) assert y.dtype == torch.bool @@ -340,7 +340,7 @@ def test_all_4d(b0: int, b1: int, b2: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.logical_reduce import AllFwdOp x = torch.randn(b0, b1, b2, n, dtype=dtype, device="cuda") - op = AllFwdOp(dtype=dtype, dim=-1) + op = AllFwdOp(dim=-1) ref = x.bool().all(dim=-1) y = op(x) assert y.dtype == torch.bool @@ -352,7 +352,7 @@ def test_all_1d(n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.logical_reduce import AllFwdOp x = _make_1d_input(n, dtype) - op = AllFwdOp(dtype=dtype, dim=-1) + op = AllFwdOp(dim=-1) ref = x.bool().all(dim=-1) y = op(x) assert y.dtype == torch.bool @@ -364,7 +364,7 @@ def test_all_dim(shape: tuple, dim: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.logical_reduce import AllFwdOp x = _make_nd_input(shape, dtype) - op = AllFwdOp(dtype=dtype, dim=dim) + op = AllFwdOp(dim=dim) ref = x.bool().all(dim=dim) y = op(x) assert y.dtype == torch.bool @@ -377,7 +377,7 @@ def test_all_keepdim(shape: tuple, dim: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.logical_reduce import AllFwdOp x = _make_nd_input(shape, dtype) - op = AllFwdOp(dtype=dtype, dim=dim, keepdim=True) + op = AllFwdOp(dim=dim, keepdim=True) ref = x.bool().all(dim=dim, keepdim=True) y = op(x) assert y.dtype == torch.bool @@ -393,7 +393,7 @@ def test_count_nonzero_op(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.logical_reduce import CountNonzeroFwdOp test = LogicalReduceTest(m, n, dtype, "count_nonzero") - op = CountNonzeroFwdOp(dtype=dtype, dim=-1) + op = CountNonzeroFwdOp(dim=-1) test.check(op, *test.gen_inputs(), compare=_exact_compare_int64) @@ -403,7 +403,7 @@ def test_count_nonzero_non_contiguous(m: int, n: int, dtype: torch.dtype) -> Non x_full = _make_noncontig_input(m, n, dtype) x = x_full[:, :n] - op = CountNonzeroFwdOp(dtype=dtype, dim=-1) + op = CountNonzeroFwdOp(dim=-1) ref = torch.count_nonzero(x.contiguous(), dim=-1).to(torch.int64) y = op(x) assert y.dtype == torch.int64 @@ -415,7 +415,7 @@ def test_count_nonzero_3d(batch: int, seq: int, hidden: int, dtype: torch.dtype) from tileops.ops.reduction.logical_reduce import CountNonzeroFwdOp x = torch.randn(batch, seq, hidden, dtype=dtype, device="cuda") - op = CountNonzeroFwdOp(dtype=dtype, dim=-1) + op = CountNonzeroFwdOp(dim=-1) ref = torch.count_nonzero(x, dim=-1).to(torch.int64) y = op(x) assert y.dtype == torch.int64 @@ -427,7 +427,7 @@ def test_count_nonzero_4d(b0: int, b1: int, b2: int, n: int, dtype: torch.dtype) from tileops.ops.reduction.logical_reduce import CountNonzeroFwdOp x = torch.randn(b0, b1, b2, n, dtype=dtype, device="cuda") - op = CountNonzeroFwdOp(dtype=dtype, dim=-1) + op = CountNonzeroFwdOp(dim=-1) ref = torch.count_nonzero(x, dim=-1).to(torch.int64) y = op(x) assert y.dtype == torch.int64 @@ -439,7 +439,7 @@ def test_count_nonzero_1d(n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.logical_reduce import CountNonzeroFwdOp x = _make_1d_input(n, dtype) - op = CountNonzeroFwdOp(dtype=dtype, dim=-1) + op = CountNonzeroFwdOp(dim=-1) ref = torch.count_nonzero(x, dim=-1).to(torch.int64) y = op(x) assert y.dtype == torch.int64 @@ -451,7 +451,7 @@ def test_count_nonzero_dim(shape: tuple, dim: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.logical_reduce import CountNonzeroFwdOp x = _make_nd_input(shape, dtype) - op = CountNonzeroFwdOp(dtype=dtype, dim=dim) + op = CountNonzeroFwdOp(dim=dim) ref = torch.count_nonzero(x, dim=dim).to(torch.int64) y = op(x) assert y.dtype == torch.int64 @@ -496,7 +496,7 @@ def test_any_smoke_float16(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.logical_reduce import AnyFwdOp test = LogicalReduceTest(m, n, dtype, "any") - op = AnyFwdOp(dtype=dtype, dim=-1) + op = AnyFwdOp(dim=-1) test.check(op, *test.gen_inputs(), compare=_exact_compare) @@ -505,7 +505,7 @@ def test_any_smoke_bfloat16(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.logical_reduce import AnyFwdOp test = LogicalReduceTest(m, n, dtype, "any") - op = AnyFwdOp(dtype=dtype, dim=-1) + op = AnyFwdOp(dim=-1) test.check(op, *test.gen_inputs(), compare=_exact_compare) @@ -514,7 +514,7 @@ def test_any_smoke_int32(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.logical_reduce import AnyFwdOp test = LogicalReduceTest(m, n, dtype, "any") - op = AnyFwdOp(dtype=dtype, dim=-1) + op = AnyFwdOp(dim=-1) test.check(op, *test.gen_inputs(), compare=_exact_compare) @@ -523,7 +523,7 @@ def test_any_smoke_int64(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.logical_reduce import AnyFwdOp test = LogicalReduceTest(m, n, dtype, "any") - op = AnyFwdOp(dtype=dtype, dim=-1) + op = AnyFwdOp(dim=-1) test.check(op, *test.gen_inputs(), compare=_exact_compare) @@ -532,7 +532,7 @@ def test_any_smoke_bool(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.logical_reduce import AnyFwdOp test = LogicalReduceTest(m, n, dtype, "any") - op = AnyFwdOp(dtype=dtype, dim=-1) + op = AnyFwdOp(dim=-1) test.check(op, *test.gen_inputs(), compare=_exact_compare) @@ -541,7 +541,7 @@ def test_all_smoke_float16(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.logical_reduce import AllFwdOp test = LogicalReduceTest(m, n, dtype, "all") - op = AllFwdOp(dtype=dtype, dim=-1) + op = AllFwdOp(dim=-1) test.check(op, *test.gen_inputs(), compare=_exact_compare) @@ -550,7 +550,7 @@ def test_all_smoke_bfloat16(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.logical_reduce import AllFwdOp test = LogicalReduceTest(m, n, dtype, "all") - op = AllFwdOp(dtype=dtype, dim=-1) + op = AllFwdOp(dim=-1) test.check(op, *test.gen_inputs(), compare=_exact_compare) @@ -559,7 +559,7 @@ def test_all_smoke_int32(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.logical_reduce import AllFwdOp test = LogicalReduceTest(m, n, dtype, "all") - op = AllFwdOp(dtype=dtype, dim=-1) + op = AllFwdOp(dim=-1) test.check(op, *test.gen_inputs(), compare=_exact_compare) @@ -568,7 +568,7 @@ def test_all_smoke_int64(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.logical_reduce import AllFwdOp test = LogicalReduceTest(m, n, dtype, "all") - op = AllFwdOp(dtype=dtype, dim=-1) + op = AllFwdOp(dim=-1) test.check(op, *test.gen_inputs(), compare=_exact_compare) @@ -577,7 +577,7 @@ def test_all_smoke_bool(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.logical_reduce import AllFwdOp test = LogicalReduceTest(m, n, dtype, "all") - op = AllFwdOp(dtype=dtype, dim=-1) + op = AllFwdOp(dim=-1) test.check(op, *test.gen_inputs(), compare=_exact_compare) @@ -586,7 +586,7 @@ def test_count_nonzero_smoke_float16(m: int, n: int, dtype: torch.dtype) -> None from tileops.ops.reduction.logical_reduce import CountNonzeroFwdOp test = LogicalReduceTest(m, n, dtype, "count_nonzero") - op = CountNonzeroFwdOp(dtype=dtype, dim=-1) + op = CountNonzeroFwdOp(dim=-1) test.check(op, *test.gen_inputs(), compare=_exact_compare_int64) @@ -595,7 +595,7 @@ def test_count_nonzero_smoke_bfloat16(m: int, n: int, dtype: torch.dtype) -> Non from tileops.ops.reduction.logical_reduce import CountNonzeroFwdOp test = LogicalReduceTest(m, n, dtype, "count_nonzero") - op = CountNonzeroFwdOp(dtype=dtype, dim=-1) + op = CountNonzeroFwdOp(dim=-1) test.check(op, *test.gen_inputs(), compare=_exact_compare_int64) @@ -604,7 +604,7 @@ def test_count_nonzero_smoke_int32(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.logical_reduce import CountNonzeroFwdOp test = LogicalReduceTest(m, n, dtype, "count_nonzero") - op = CountNonzeroFwdOp(dtype=dtype, dim=-1) + op = CountNonzeroFwdOp(dim=-1) test.check(op, *test.gen_inputs(), compare=_exact_compare_int64) @@ -613,7 +613,7 @@ def test_count_nonzero_smoke_int64(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.logical_reduce import CountNonzeroFwdOp test = LogicalReduceTest(m, n, dtype, "count_nonzero") - op = CountNonzeroFwdOp(dtype=dtype, dim=-1) + op = CountNonzeroFwdOp(dim=-1) test.check(op, *test.gen_inputs(), compare=_exact_compare_int64) @@ -622,7 +622,7 @@ def test_count_nonzero_smoke_bool(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.logical_reduce import CountNonzeroFwdOp test = LogicalReduceTest(m, n, dtype, "count_nonzero") - op = CountNonzeroFwdOp(dtype=dtype, dim=-1) + op = CountNonzeroFwdOp(dim=-1) test.check(op, *test.gen_inputs(), compare=_exact_compare_int64) @@ -646,13 +646,12 @@ def test_logical_reduce_long_sequence_tiled(op_kind: str, dtype: torch.dtype) -> } test = LogicalReduceTest(3, 33024, dtype, op_kind) op = op_map[op_kind]( - dtype=dtype, dim=-1, kernel_map={"logical_reduce": _TailBlockLogicalReduceKernel}, ) compare = _exact_compare_int64 if op_kind == "count_nonzero" else _exact_compare test.check(op, *test.gen_inputs(), compare=compare) - kernel = op._kernel_cache[(3, 33024)] + kernel = op._kernel_cache[(3, 33024, dtype)] assert kernel.config["block_m"] > test.shape[0] assert kernel.config["tile_n"] > 0 @@ -690,7 +689,7 @@ def test_logical_reduce_accepts_bool(op_name: str) -> None: import tileops.ops.reduction as mod cls = getattr(mod, op_name) - op = cls(dtype=torch.bool, dim=-1) + op = cls(dim=-1) x = torch.randint(0, 2, (_M, _N), device="cuda").bool() out = op(x) assert out.dtype == torch.bool @@ -702,7 +701,7 @@ def test_logical_reduce_accepts_bool(op_name: str) -> None: def test_count_nonzero_returns_int64() -> None: from tileops.ops.reduction.logical_reduce import CountNonzeroFwdOp - op = CountNonzeroFwdOp(dtype=torch.float16, dim=-1) + op = CountNonzeroFwdOp(dim=-1) x = torch.randn(_M, _N, dtype=torch.float16, device="cuda") out = op(x) assert out.dtype == torch.int64, ( @@ -717,7 +716,7 @@ def test_logical_reduce_returns_bool(op_name: str) -> None: import tileops.ops.reduction as mod cls = getattr(mod, op_name) - op = cls(dtype=torch.float16, dim=-1) + op = cls(dim=-1) x = torch.randn(_M, _N, dtype=torch.float16, device="cuda") out = op(x) assert out.dtype == torch.bool diff --git a/tests/ops/test_normalization_alignment.py b/tests/ops/test_normalization_alignment.py index 5de9f8796..86d47148c 100644 --- a/tests/ops/test_normalization_alignment.py +++ b/tests/ops/test_normalization_alignment.py @@ -17,7 +17,7 @@ def test_rms_norm_accepts_normalized_shape() -> None: from tileops.ops.norm.rms_norm import RMSNormFwdOp - op = RMSNormFwdOp(normalized_shape=(4096,), eps=None, dtype=torch.float16) + op = RMSNormFwdOp(normalized_shape=(4096,), eps=None) assert op.N == 4096 assert op.normalized_shape == (4096,) @@ -26,7 +26,7 @@ def test_rms_norm_accepts_normalized_shape() -> None: def test_layer_norm_accepts_normalized_shape() -> None: from tileops.ops.norm.layer_norm import LayerNormFwdOp - op = LayerNormFwdOp(normalized_shape=[4096], dtype=torch.float16) + op = LayerNormFwdOp(normalized_shape=[4096]) assert op.N == 4096 assert op.normalized_shape == (4096,) @@ -41,7 +41,7 @@ def test_rms_norm_accepts_tuple_normalized_shape_runtime() -> None: from tileops.ops.norm.rms_norm import RMSNormFwdOp - op = RMSNormFwdOp(normalized_shape=(2, 3), dtype=torch.float16) + op = RMSNormFwdOp(normalized_shape=(2, 3)) assert op.N == 6 assert op.normalized_shape == (2, 3) x = torch.randn(4, 2, 3, dtype=torch.float16, device="cuda") @@ -59,7 +59,7 @@ def test_layer_norm_accepts_tuple_normalized_shape_runtime() -> None: from tileops.ops.norm.layer_norm import LayerNormFwdOp - op = LayerNormFwdOp(normalized_shape=(2, 3), dtype=torch.float16) + op = LayerNormFwdOp(normalized_shape=(2, 3)) assert op.N == 6 assert op.normalized_shape == (2, 3) x = torch.randn(4, 2, 3, dtype=torch.float16, device="cuda") diff --git a/tests/ops/test_reduce.py b/tests/ops/test_reduce.py index 33c6a1cc0..d73d46c55 100644 --- a/tests/ops/test_reduce.py +++ b/tests/ops/test_reduce.py @@ -190,7 +190,7 @@ def test_sum_op(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.reduce import SumFwdOp test = ReduceTest(m, n, dtype, "sum") - op = SumFwdOp(dtype=dtype, dim=-1) + op = SumFwdOp(dim=-1) test.check(op, *test.gen_inputs(), **_tol(dtype)) @@ -199,7 +199,7 @@ def test_sum_tiled(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.reduce import SumFwdOp test = ReduceTest(m, n, dtype, "sum") - op = SumFwdOp(dtype=dtype, dim=-1) + op = SumFwdOp(dim=-1) test.check(op, *test.gen_inputs(), **_tol(dtype)) @@ -208,7 +208,7 @@ def test_prod_tiled(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.reduce import ProdFwdOp test = ProdTest(m, n, dtype) - op = ProdFwdOp(dtype=dtype, dim=-1) + op = ProdFwdOp(dim=-1) tol = {"atol": 5e-2, "rtol": 5e-2} if dtype != torch.float32 else {"atol": 1e-3, "rtol": 1e-3} test.check(op, *test.gen_inputs(), **tol) @@ -218,7 +218,7 @@ def test_var_tiled(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.reduce import VarFwdOp test = WelfordTest(m, n, dtype, "var", correction=1) - op = VarFwdOp(dtype=dtype, dim=-1) + op = VarFwdOp(dim=-1) test.check(op, *test.gen_inputs(), **_tol(dtype)) @@ -272,10 +272,10 @@ def test_reduce_untiled_autotune_unaligned_n() -> None: m, n, dtype = 8, 20000, torch.float16 test = ReduceTest(m, n, dtype, "sum") - op = SumFwdOp(dtype=dtype, dim=-1, tune=True) + op = SumFwdOp(dim=-1, tune=True) test.check(op, *test.gen_inputs(), **_tol(dtype)) - kernel = op._kernel_cache[(m, n)] + kernel = op._kernel_cache[(m, n, dtype)] assert not kernel._needs_tiling assert {c["block_m"] for c in kernel.autotune_configs} == {1} @@ -299,13 +299,13 @@ def test_reduce_tiled_autotune(op_kind: str) -> None: m, n, dtype = 4, 40000, torch.float16 if op_kind == "sum": test = ReduceTest(m, n, dtype, "sum") - op = SumFwdOp(dtype=dtype, dim=-1, tune=True) + op = SumFwdOp(dim=-1, tune=True) else: test = WelfordTest(m, n, dtype, "var", correction=1) - op = VarFwdOp(dtype=dtype, dim=-1, tune=True) + op = VarFwdOp(dim=-1, tune=True) test.check(op, *test.gen_inputs(), **_tol(dtype)) - kernel = op._kernel_cache[(m, n)] + kernel = op._kernel_cache[(m, n, dtype)] assert kernel._needs_tiling assert kernel.config in kernel.autotune_configs @@ -316,7 +316,7 @@ def test_sum_non_contiguous(m: int, n: int, dtype: torch.dtype) -> None: x_full = torch.randn(m, n * 2, dtype=dtype, device="cuda") x = x_full[:, :n] - op = SumFwdOp(dtype=dtype, dim=-1) + op = SumFwdOp(dim=-1) ref = x.contiguous().float().sum(dim=-1).to(dtype) y = op(x) tol = _tol(dtype) @@ -328,7 +328,7 @@ def test_sum_3d(batch: int, seq: int, hidden: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.reduce import SumFwdOp x = torch.randn(batch, seq, hidden, dtype=dtype, device="cuda") - op = SumFwdOp(dtype=dtype, dim=-1) + op = SumFwdOp(dim=-1) ref = x.float().sum(dim=-1).to(dtype) y = op(x) tol = _tol(dtype) @@ -340,7 +340,7 @@ def test_sum_4d(b0: int, b1: int, b2: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.reduce import SumFwdOp x = torch.randn(b0, b1, b2, n, dtype=dtype, device="cuda") - op = SumFwdOp(dtype=dtype, dim=-1) + op = SumFwdOp(dim=-1) ref = x.float().sum(dim=-1).to(dtype) y = op(x) tol = _tol(dtype) @@ -355,7 +355,7 @@ def test_mean_op(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.reduce import MeanFwdOp test = ReduceTest(m, n, dtype, "mean") - op = MeanFwdOp(dtype=dtype, dim=-1) + op = MeanFwdOp(dim=-1) test.check(op, *test.gen_inputs(), **_tol(dtype)) @@ -367,7 +367,7 @@ def test_amin_op(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.reduce import AminFwdOp test = ReduceTest(m, n, dtype, "amin") - op = AminFwdOp(dtype=dtype, dim=-1) + op = AminFwdOp(dim=-1) test.check(op, *test.gen_inputs(), **_tol(dtype)) @@ -379,7 +379,7 @@ def test_amax_op(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.reduce import AmaxFwdOp test = ReduceTest(m, n, dtype, "amax") - op = AmaxFwdOp(dtype=dtype, dim=-1) + op = AmaxFwdOp(dim=-1) test.check(op, *test.gen_inputs(), **_tol(dtype)) @@ -391,7 +391,7 @@ def test_prod_op(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.reduce import ProdFwdOp test = ProdTest(m, n, dtype) - op = ProdFwdOp(dtype=dtype, dim=-1) + op = ProdFwdOp(dim=-1) # Prod is more numerically sensitive tol = {"atol": 5e-2, "rtol": 5e-2} if dtype != torch.float32 else {"atol": 1e-3, "rtol": 1e-3} test.check(op, *test.gen_inputs(), **tol) @@ -405,7 +405,7 @@ def test_std_op(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.reduce import StdFwdOp test = WelfordTest(m, n, dtype, "std", correction=1) - op = StdFwdOp(dtype=dtype, dim=-1) + op = StdFwdOp(dim=-1) test.check(op, *test.gen_inputs(), **_tol(dtype)) @@ -414,7 +414,7 @@ def test_std_bessel(m: int, n: int, dtype: torch.dtype, correction: int) -> None from tileops.ops.reduction.reduce import StdFwdOp test = WelfordTest(m, n, dtype, "std", correction=correction) - op = StdFwdOp(dtype=dtype, correction=correction, dim=-1) + op = StdFwdOp(correction=correction, dim=-1) test.check(op, *test.gen_inputs(), **_tol(dtype)) @@ -426,7 +426,7 @@ def test_var_op(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.reduce import VarFwdOp test = WelfordTest(m, n, dtype, "var", correction=1) - op = VarFwdOp(dtype=dtype, dim=-1) + op = VarFwdOp(dim=-1) test.check(op, *test.gen_inputs(), **_tol(dtype)) @@ -435,7 +435,7 @@ def test_var_bessel(m: int, n: int, dtype: torch.dtype, correction: int) -> None from tileops.ops.reduction.reduce import VarFwdOp test = WelfordTest(m, n, dtype, "var", correction=correction) - op = VarFwdOp(dtype=dtype, correction=correction, dim=-1) + op = VarFwdOp(correction=correction, dim=-1) test.check(op, *test.gen_inputs(), **_tol(dtype)) @@ -447,7 +447,7 @@ def test_var_mean_op(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.reduce import VarMeanFwdOp test = WelfordTest(m, n, dtype, "var_mean", correction=1) - op = VarMeanFwdOp(dtype=dtype, dim=-1) + op = VarMeanFwdOp(dim=-1) test.check(op, *test.gen_inputs(), **_tol(dtype)) @@ -456,7 +456,7 @@ def test_var_mean_bessel(m: int, n: int, dtype: torch.dtype, correction: int) -> from tileops.ops.reduction.reduce import VarMeanFwdOp test = WelfordTest(m, n, dtype, "var_mean", correction=correction) - op = VarMeanFwdOp(dtype=dtype, correction=correction, dim=-1) + op = VarMeanFwdOp(correction=correction, dim=-1) test.check(op, *test.gen_inputs(), **_tol(dtype)) @@ -468,7 +468,7 @@ def test_var_3d(batch: int, seq: int, hidden: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.reduce import VarFwdOp x = torch.randn(batch, seq, hidden, dtype=dtype, device="cuda") - op = VarFwdOp(dtype=dtype, dim=-1) + op = VarFwdOp(dim=-1) ref = x.float().var(dim=-1, correction=1).to(dtype) y = op(x) tol = _tol(dtype) @@ -480,7 +480,7 @@ def test_std_3d(batch: int, seq: int, hidden: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.reduce import StdFwdOp x = torch.randn(batch, seq, hidden, dtype=dtype, device="cuda") - op = StdFwdOp(dtype=dtype, dim=-1) + op = StdFwdOp(dim=-1) ref = x.float().std(dim=-1, correction=1).to(dtype) y = op(x) tol = _tol(dtype) @@ -495,7 +495,7 @@ def test_sum_1d(n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.reduce import SumFwdOp x = torch.randn(n, dtype=dtype, device="cuda") - op = SumFwdOp(dtype=dtype, dim=-1) + op = SumFwdOp(dim=-1) ref = x.float().sum(dim=-1).to(dtype) y = op(x) tol = _tol(dtype) @@ -509,7 +509,7 @@ def test_var_1d(n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.reduce import VarFwdOp x = torch.randn(n, dtype=dtype, device="cuda") - op = VarFwdOp(dtype=dtype, dim=-1) + op = VarFwdOp(dim=-1) ref = x.float().var(dim=-1, correction=1).to(dtype) y = op(x) tol = _tol(dtype) @@ -527,7 +527,7 @@ def test_var_non_contiguous(m: int, n: int, dtype: torch.dtype) -> None: x_full = torch.randn(m, n * 2, dtype=dtype, device="cuda") x = x_full[:, :n] - op = VarFwdOp(dtype=dtype, dim=-1) + op = VarFwdOp(dim=-1) ref = x.contiguous().float().var(dim=-1, correction=1).to(dtype) y = op(x) tol = _tol(dtype) @@ -540,7 +540,7 @@ def test_std_non_contiguous(m: int, n: int, dtype: torch.dtype) -> None: x_full = torch.randn(m, n * 2, dtype=dtype, device="cuda") x = x_full[:, :n] - op = StdFwdOp(dtype=dtype, dim=-1) + op = StdFwdOp(dim=-1) ref = x.contiguous().float().std(dim=-1, correction=1).to(dtype) y = op(x) tol = _tol(dtype) @@ -572,7 +572,7 @@ def test_sum_spec_basic(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.reduce import SumFwdOp x = torch.randn(m, n, dtype=dtype, device="cuda") - op = SumFwdOp(dtype=dtype, dim=-1) + op = SumFwdOp(dim=-1) ref = torch.sum(x.float(), dim=-1).to(dtype) y = op(x) tol = _tol(dtype) @@ -585,7 +585,7 @@ def test_sum_spec_dim(shape: tuple, dim: int, keepdim: bool, dtype: torch.dtype) from tileops.ops.reduction.reduce import SumFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = SumFwdOp(dtype=dtype, dim=dim, keepdim=keepdim) + op = SumFwdOp(dim=dim, keepdim=keepdim) ref = torch.sum(x.float(), dim=dim, keepdim=keepdim).to(dtype) y = op(x) tol = _tol(dtype) @@ -599,7 +599,7 @@ def test_sum_spec_keepdim(shape: tuple, dim: int, keepdim: bool, dtype: torch.dt x = torch.randn(*shape, dtype=dtype, device="cuda") # Force keepdim=True regardless of fixture param to specifically test shape preservation - op = SumFwdOp(dtype=dtype, dim=dim, keepdim=True) + op = SumFwdOp(dim=dim, keepdim=True) ref = torch.sum(x.float(), dim=dim, keepdim=True).to(dtype) y = op(x) tol = _tol(dtype) @@ -613,7 +613,7 @@ def test_sum_spec_1d(n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.reduce import SumFwdOp x = torch.randn(n, dtype=dtype, device="cuda") - op = SumFwdOp(dtype=dtype, dim=-1) + op = SumFwdOp(dim=-1) ref = torch.sum(x.float(), dim=-1).to(dtype) y = op(x) tol = _tol(dtype) @@ -628,7 +628,7 @@ def test_mean_spec_dim(shape: tuple, dim: int, keepdim: bool, dtype: torch.dtype from tileops.ops.reduction.reduce import MeanFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = MeanFwdOp(dtype=dtype, dim=dim, keepdim=keepdim) + op = MeanFwdOp(dim=dim, keepdim=keepdim) ref = torch.mean(x.float(), dim=dim, keepdim=keepdim).to(dtype) y = op(x) tol = _tol(dtype) @@ -642,7 +642,7 @@ def test_amax_spec_dim(shape: tuple, dim: int, keepdim: bool, dtype: torch.dtype from tileops.ops.reduction.reduce import AmaxFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = AmaxFwdOp(dtype=dtype, dim=dim, keepdim=keepdim) + op = AmaxFwdOp(dim=dim, keepdim=keepdim) ref = torch.amax(x.float(), dim=dim, keepdim=keepdim).to(dtype) y = op(x) tol = _tol(dtype) @@ -656,7 +656,7 @@ def test_amin_spec_dim(shape: tuple, dim: int, keepdim: bool, dtype: torch.dtype from tileops.ops.reduction.reduce import AminFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = AminFwdOp(dtype=dtype, dim=dim, keepdim=keepdim) + op = AminFwdOp(dim=dim, keepdim=keepdim) ref = torch.amin(x.float(), dim=dim, keepdim=keepdim).to(dtype) y = op(x) tol = _tol(dtype) @@ -670,7 +670,7 @@ def test_prod_spec_dim(shape: tuple, dim: int, keepdim: bool, dtype: torch.dtype from tileops.ops.reduction.reduce import ProdFwdOp x = torch.rand(*shape, dtype=dtype, device="cuda") * 0.01 + 0.99 - op = ProdFwdOp(dtype=dtype, dim=dim, keepdim=keepdim) + op = ProdFwdOp(dim=dim, keepdim=keepdim) ref = torch.prod(x.float(), dim=dim, keepdim=keepdim).to(dtype) y = op(x) tol = {"atol": 5e-2, "rtol": 5e-2} if dtype != torch.float32 else {"atol": 1e-3, "rtol": 1e-3} @@ -684,7 +684,7 @@ def test_var_spec_dim(shape: tuple, dim: int, keepdim: bool, dtype: torch.dtype) from tileops.ops.reduction.reduce import VarFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = VarFwdOp(dtype=dtype, dim=dim, keepdim=keepdim) + op = VarFwdOp(dim=dim, keepdim=keepdim) ref = torch.var(x.float(), dim=dim, keepdim=keepdim, correction=1).to(dtype) y = op(x) tol = _tol(dtype) @@ -698,7 +698,7 @@ def test_std_spec_dim(shape: tuple, dim: int, keepdim: bool, dtype: torch.dtype) from tileops.ops.reduction.reduce import StdFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = StdFwdOp(dtype=dtype, dim=dim, keepdim=keepdim) + op = StdFwdOp(dim=dim, keepdim=keepdim) ref = torch.std(x.float(), dim=dim, keepdim=keepdim, correction=1).to(dtype) y = op(x) tol = _tol(dtype) @@ -712,7 +712,7 @@ def test_var_mean_spec_dim(shape: tuple, dim: int, keepdim: bool, dtype: torch.d from tileops.ops.reduction.reduce import VarMeanFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = VarMeanFwdOp(dtype=dtype, dim=dim, keepdim=keepdim) + op = VarMeanFwdOp(dim=dim, keepdim=keepdim) ref_var = torch.var(x.float(), dim=dim, keepdim=keepdim, correction=1).to(dtype) ref_mean = torch.mean(x.float(), dim=dim, keepdim=keepdim).to(dtype) var_out, mean_out = op(x) diff --git a/tests/ops/test_reduce_arithmetic_conformance.py b/tests/ops/test_reduce_arithmetic_conformance.py index 8a7a604a4..3a4f44a1a 100644 --- a/tests/ops/test_reduce_arithmetic_conformance.py +++ b/tests/ops/test_reduce_arithmetic_conformance.py @@ -85,7 +85,7 @@ def test_arithmetic_reduce_conformance( """Each (op, dim-shape, keepdim, dtype) cell must match PyTorch.""" torch.manual_seed(0) x = torch.randn(*_SHAPE, dtype=dtype, device="cuda") - op = op_cls(dtype=dtype, dim=dim, keepdim=keepdim) + op = op_cls(dim=dim, keepdim=keepdim) y = op(x) ref = _ref(torch_fn, x, dim, keepdim) assert y.shape == ref.shape, ( @@ -100,7 +100,7 @@ def test_arithmetic_reduce_conformance( def test_dim_none_keepdim_false_returns_0d(op_cls: type, torch_fn: Callable) -> None: """``dim=None, keepdim=False`` must return a 0-D tensor matching PyTorch.""" x = torch.randn(*_SHAPE, dtype=torch.float32, device="cuda") - op = op_cls(dtype=torch.float32, dim=None, keepdim=False) + op = op_cls(dim=None, keepdim=False) y = op(x) ref = _ref(torch_fn, x, None, False) assert y.ndim == 0, f"{op_cls.__name__}: expected 0-D, got shape {y.shape}" @@ -133,7 +133,7 @@ def test_arithmetic_reduce_unaligned_innermost( unaligned_shape = (4, 8, 255) dtype = torch.float16 x = torch.randn(*unaligned_shape, dtype=dtype, device="cuda") - op = op_cls(dtype=dtype, dim=dim, keepdim=False) + op = op_cls(dim=dim, keepdim=False) y = op(x) ref = _ref(torch_fn, x, dim, False) assert y.shape == ref.shape, ( diff --git a/tests/ops/test_reduce_boolean_conformance.py b/tests/ops/test_reduce_boolean_conformance.py index 2b0da36a5..aaf7ac394 100644 --- a/tests/ops/test_reduce_boolean_conformance.py +++ b/tests/ops/test_reduce_boolean_conformance.py @@ -71,7 +71,7 @@ def test_logical_reduce_conformance( zero_mask = torch.rand(_SHAPE, device="cuda") < 0.1 x = raw.masked_fill(zero_mask, 0) - op = op_cls(dtype=dtype, dim=dim, keepdim=keepdim) + op = op_cls(dim=dim, keepdim=keepdim) y = op(x) if dim is None: ref = torch_fn(x) @@ -118,7 +118,7 @@ def test_logical_reduce_unaligned_innermost( zero_mask = torch.rand(_UNALIGNED_SHAPE, device="cuda") < 0.1 x = raw.masked_fill(zero_mask, 0) - op = op_cls(dtype=dtype, dim=dim, keepdim=False) + op = op_cls(dim=dim, keepdim=False) y = op(x) ref = torch_fn(x) if dim is None else torch_fn(x, dim=dim, keepdim=False) @@ -158,7 +158,7 @@ def test_count_nonzero_conformance(dim, dtype: torch.dtype) -> None: zero_mask = torch.rand(_SHAPE, device="cuda") < 0.1 x = raw.masked_fill(zero_mask, 0) - op = CountNonzeroFwdOp(dtype=dtype, dim=dim) + op = CountNonzeroFwdOp(dim=dim) y = op(x) ref = torch.count_nonzero(x, dim=dim) @@ -190,7 +190,7 @@ def test_count_nonzero_unaligned_innermost(dim) -> None: zero_mask = torch.rand(_UNALIGNED_SHAPE, device="cuda") < 0.1 x = raw.masked_fill(zero_mask, 0) - op = CountNonzeroFwdOp(dtype=dtype, dim=dim) + op = CountNonzeroFwdOp(dim=dim) y = op(x) ref = torch.count_nonzero(x, dim=dim) diff --git a/tests/ops/test_reduce_dim_none.py b/tests/ops/test_reduce_dim_none.py index 72d6ca9be..67de4426d 100644 --- a/tests/ops/test_reduce_dim_none.py +++ b/tests/ops/test_reduce_dim_none.py @@ -86,7 +86,7 @@ def test_sum_dim_none( from tileops.ops.reduction.reduce import SumFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = SumFwdOp(dtype=dtype, dim=None, keepdim=keepdim) + op = SumFwdOp(dim=None, keepdim=keepdim) dims = _all_dims(shape) ref = torch.sum(x.float(), dim=dims, keepdim=keepdim).to(dtype) y = op(x) @@ -102,7 +102,7 @@ def test_mean_dim_none( from tileops.ops.reduction.reduce import MeanFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = MeanFwdOp(dtype=dtype, dim=None, keepdim=keepdim) + op = MeanFwdOp(dim=None, keepdim=keepdim) dims = _all_dims(shape) ref = torch.mean(x.float(), dim=dims, keepdim=keepdim).to(dtype) y = op(x) @@ -118,7 +118,7 @@ def test_amax_dim_none( from tileops.ops.reduction.reduce import AmaxFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = AmaxFwdOp(dtype=dtype, dim=None, keepdim=keepdim) + op = AmaxFwdOp(dim=None, keepdim=keepdim) dims = _all_dims(shape) ref = torch.amax(x.float(), dim=dims, keepdim=keepdim).to(dtype) y = op(x) @@ -134,7 +134,7 @@ def test_amin_dim_none( from tileops.ops.reduction.reduce import AminFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = AminFwdOp(dtype=dtype, dim=None, keepdim=keepdim) + op = AminFwdOp(dim=None, keepdim=keepdim) dims = _all_dims(shape) ref = torch.amin(x.float(), dim=dims, keepdim=keepdim).to(dtype) y = op(x) @@ -151,7 +151,7 @@ def test_prod_dim_none_rejected() -> None: from tileops.ops.reduction.reduce import ProdFwdOp with pytest.raises(TypeError, match="ProdFwdOp.dim must be int"): - ProdFwdOp(dtype=torch.float16, dim=None) + ProdFwdOp(dim=None) # Welford ops: var, std, var_mean @@ -164,7 +164,7 @@ def test_var_dim_none( from tileops.ops.reduction.reduce import VarFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = VarFwdOp(dtype=dtype, dim=None, keepdim=keepdim) + op = VarFwdOp(dim=None, keepdim=keepdim) dims = _all_dims(shape) ref = torch.var(x.float(), dim=dims, keepdim=keepdim, correction=1).to(dtype) y = op(x) @@ -180,7 +180,7 @@ def test_std_dim_none( from tileops.ops.reduction.reduce import StdFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = StdFwdOp(dtype=dtype, dim=None, keepdim=keepdim) + op = StdFwdOp(dim=None, keepdim=keepdim) dims = _all_dims(shape) ref = torch.std(x.float(), dim=dims, keepdim=keepdim, correction=1).to(dtype) y = op(x) @@ -196,7 +196,7 @@ def test_var_mean_dim_none( from tileops.ops.reduction.reduce import VarMeanFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = VarMeanFwdOp(dtype=dtype, dim=None, keepdim=keepdim) + op = VarMeanFwdOp(dim=None, keepdim=keepdim) dims = _all_dims(shape) ref_var = torch.var( x.float(), dim=dims, keepdim=keepdim, correction=1, @@ -259,7 +259,7 @@ def test_all_dim_none( from tileops.ops.reduction.logical_reduce import AllFwdOp x = _make_logical_input(shape, dtype) - op = AllFwdOp(dtype=dtype, dim=None, keepdim=keepdim) + op = AllFwdOp(dim=None, keepdim=keepdim) dims = _all_dims(shape) ref = torch.all(x.bool(), dim=dims, keepdim=keepdim) y = op(x) @@ -274,7 +274,7 @@ def test_any_dim_none( from tileops.ops.reduction.logical_reduce import AnyFwdOp x = _make_logical_input(shape, dtype) - op = AnyFwdOp(dtype=dtype, dim=None, keepdim=keepdim) + op = AnyFwdOp(dim=None, keepdim=keepdim) dims = _all_dims(shape) ref = torch.any(x.bool(), dim=dims, keepdim=keepdim) y = op(x) @@ -289,7 +289,7 @@ def test_count_nonzero_dim_none() -> None: shape = (4, 8, 256) x = torch.randn(*shape, dtype=torch.float32, device="cuda") x[x < 0] = 0.0 - op = CountNonzeroFwdOp(dtype=torch.float32, dim=None) + op = CountNonzeroFwdOp(dim=None) dims = _all_dims(shape) ref = torch.count_nonzero(x, dim=dims) y = op(x) @@ -307,7 +307,7 @@ def test_count_nonzero_dim_none_dtypes(dtype: torch.dtype) -> None: shape = (4, 8, 256) x = torch.randn(*shape, dtype=dtype, device="cuda") x[x < 0] = 0.0 - op = CountNonzeroFwdOp(dtype=dtype, dim=None) + op = CountNonzeroFwdOp(dim=None) dims = _all_dims(shape) ref = torch.count_nonzero(x, dim=dims) y = op(x) @@ -325,7 +325,7 @@ def test_l1_norm_dim_none( from tileops.ops.reduction.vector_norm import L1NormFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = L1NormFwdOp(dtype=dtype, dim=None, keepdim=keepdim) + op = L1NormFwdOp(dim=None, keepdim=keepdim) dims = _all_dims(shape) ref = torch.linalg.vector_norm( x.float(), ord=1, dim=dims, keepdim=keepdim, @@ -343,7 +343,7 @@ def test_l2_norm_dim_none( from tileops.ops.reduction.vector_norm import L2NormFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = L2NormFwdOp(dtype=dtype, dim=None, keepdim=keepdim) + op = L2NormFwdOp(dim=None, keepdim=keepdim) dims = _all_dims(shape) ref = torch.linalg.vector_norm( x.float(), ord=2, dim=dims, keepdim=keepdim, @@ -361,7 +361,7 @@ def test_inf_norm_dim_none( from tileops.ops.reduction.vector_norm import InfNormFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = InfNormFwdOp(dtype=dtype, dim=None, keepdim=keepdim) + op = InfNormFwdOp(dim=None, keepdim=keepdim) dims = _all_dims(shape) ref = torch.linalg.vector_norm( x.float(), ord=float("inf"), dim=dims, keepdim=keepdim, diff --git a/tests/ops/test_reduce_multidim.py b/tests/ops/test_reduce_multidim.py index d00ea4771..efb1a887f 100644 --- a/tests/ops/test_reduce_multidim.py +++ b/tests/ops/test_reduce_multidim.py @@ -68,7 +68,7 @@ def test_sum_multidim( from tileops.ops.reduction.reduce import SumFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = SumFwdOp(dtype=dtype, dim=dims, keepdim=keepdim) + op = SumFwdOp(dim=dims, keepdim=keepdim) ref = torch.sum(x.float(), dim=dims, keepdim=keepdim).to(dtype) y = op(x) tol = _tol(dtype) @@ -83,7 +83,7 @@ def test_mean_multidim( from tileops.ops.reduction.reduce import MeanFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = MeanFwdOp(dtype=dtype, dim=dims, keepdim=keepdim) + op = MeanFwdOp(dim=dims, keepdim=keepdim) ref = torch.mean(x.float(), dim=dims, keepdim=keepdim).to(dtype) y = op(x) tol = _tol(dtype) @@ -98,7 +98,7 @@ def test_amax_multidim( from tileops.ops.reduction.reduce import AmaxFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = AmaxFwdOp(dtype=dtype, dim=dims, keepdim=keepdim) + op = AmaxFwdOp(dim=dims, keepdim=keepdim) ref = torch.amax(x.float(), dim=dims, keepdim=keepdim).to(dtype) y = op(x) tol = _tol(dtype) @@ -114,9 +114,9 @@ def test_prod_multidim_rejected() -> None: from tileops.ops.reduction.reduce import ProdFwdOp with pytest.raises(TypeError, match="ProdFwdOp.dim must be int"): - ProdFwdOp(dtype=torch.float16, dim=[0, 1]) + ProdFwdOp(dim=[0, 1]) with pytest.raises(TypeError, match="ProdFwdOp.dim must be int"): - ProdFwdOp(dtype=torch.float16, dim=(0, 1)) + ProdFwdOp(dim=(0, 1)) @MultiDimFixture @@ -126,7 +126,7 @@ def test_amin_multidim( from tileops.ops.reduction.reduce import AminFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = AminFwdOp(dtype=dtype, dim=dims, keepdim=keepdim) + op = AminFwdOp(dim=dims, keepdim=keepdim) ref = torch.amin(x.float(), dim=dims, keepdim=keepdim).to(dtype) y = op(x) tol = _tol(dtype) @@ -144,7 +144,7 @@ def test_var_multidim( from tileops.ops.reduction.reduce import VarFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = VarFwdOp(dtype=dtype, dim=dims, keepdim=keepdim) + op = VarFwdOp(dim=dims, keepdim=keepdim) ref = torch.var(x.float(), dim=dims, keepdim=keepdim, correction=1).to(dtype) y = op(x) tol = _tol(dtype) @@ -159,7 +159,7 @@ def test_std_multidim( from tileops.ops.reduction.reduce import StdFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = StdFwdOp(dtype=dtype, dim=dims, keepdim=keepdim) + op = StdFwdOp(dim=dims, keepdim=keepdim) ref = torch.std(x.float(), dim=dims, keepdim=keepdim, correction=1).to(dtype) y = op(x) tol = _tol(dtype) @@ -174,7 +174,7 @@ def test_var_mean_multidim( from tileops.ops.reduction.reduce import VarMeanFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = VarMeanFwdOp(dtype=dtype, dim=dims, keepdim=keepdim) + op = VarMeanFwdOp(dim=dims, keepdim=keepdim) ref_var = torch.var( x.float(), dim=dims, keepdim=keepdim, correction=1, ).to(dtype) @@ -197,7 +197,7 @@ def test_logsumexp_multidim( from tileops.ops.reduction.softmax import LogSumExpFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = LogSumExpFwdOp(dtype=dtype, dim=dims, keepdim=keepdim) + op = LogSumExpFwdOp(dim=dims, keepdim=keepdim) ref = torch.logsumexp(x.float(), dim=dims, keepdim=keepdim).to(dtype) y = op(x) tol = _tol(dtype) @@ -252,7 +252,7 @@ def test_all_multidim( from tileops.ops.reduction.logical_reduce import AllFwdOp x = _make_logical_input(shape, dtype) - op = AllFwdOp(dtype=dtype, dim=dims, keepdim=keepdim) + op = AllFwdOp(dim=dims, keepdim=keepdim) ref = torch.all(x.bool(), dim=dims, keepdim=keepdim) y = op(x) assert y.shape == ref.shape, f"shape mismatch: {y.shape} vs {ref.shape}" @@ -266,7 +266,7 @@ def test_any_multidim( from tileops.ops.reduction.logical_reduce import AnyFwdOp x = _make_logical_input(shape, dtype) - op = AnyFwdOp(dtype=dtype, dim=dims, keepdim=keepdim) + op = AnyFwdOp(dim=dims, keepdim=keepdim) ref = torch.any(x.bool(), dim=dims, keepdim=keepdim) y = op(x) assert y.shape == ref.shape, f"shape mismatch: {y.shape} vs {ref.shape}" @@ -309,7 +309,7 @@ def test_count_nonzero_multidim( x = torch.randn(*shape, dtype=dtype, device="cuda") # Zero out some elements to make it interesting x[x < 0] = 0.0 - op = CountNonzeroFwdOp(dtype=dtype, dim=dims) + op = CountNonzeroFwdOp(dim=dims) ref = torch.count_nonzero(x, dim=dims) y = op(x) assert y.shape == ref.shape, f"shape mismatch: {y.shape} vs {ref.shape}" @@ -326,7 +326,7 @@ def test_l1_norm_multidim( from tileops.ops.reduction.vector_norm import L1NormFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = L1NormFwdOp(dtype=dtype, dim=dims, keepdim=keepdim) + op = L1NormFwdOp(dim=dims, keepdim=keepdim) ref = torch.linalg.vector_norm( x.float(), ord=1, dim=dims, keepdim=keepdim, ).to(dtype) @@ -343,7 +343,7 @@ def test_l2_norm_multidim( from tileops.ops.reduction.vector_norm import L2NormFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = L2NormFwdOp(dtype=dtype, dim=dims, keepdim=keepdim) + op = L2NormFwdOp(dim=dims, keepdim=keepdim) ref = torch.linalg.vector_norm( x.float(), ord=2, dim=dims, keepdim=keepdim, ).to(dtype) @@ -360,7 +360,7 @@ def test_inf_norm_multidim( from tileops.ops.reduction.vector_norm import InfNormFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = InfNormFwdOp(dtype=dtype, dim=dims, keepdim=keepdim) + op = InfNormFwdOp(dim=dims, keepdim=keepdim) ref = torch.linalg.vector_norm( x.float(), ord=float("inf"), dim=dims, keepdim=keepdim, ).to(dtype) @@ -397,8 +397,8 @@ def test_sum_empty_dim_full_reduction() -> None: from tileops.ops.reduction.reduce import SumFwdOp x = torch.randn(2, 3, 4, dtype=torch.float16, device="cuda") - op = SumFwdOp(dtype=torch.float16, dim=[], keepdim=False) - op_none = SumFwdOp(dtype=torch.float16, dim=None, keepdim=False) + op = SumFwdOp(dim=[], keepdim=False) + op_none = SumFwdOp(dim=None, keepdim=False) assert torch.allclose(op(x), op_none(x), **_tol(torch.float16)) @@ -407,8 +407,8 @@ def test_mean_empty_dim_full_reduction() -> None: from tileops.ops.reduction.reduce import MeanFwdOp x = torch.randn(2, 3, 4, dtype=torch.float16, device="cuda") - op = MeanFwdOp(dtype=torch.float16, dim=(), keepdim=True) - op_none = MeanFwdOp(dtype=torch.float16, dim=None, keepdim=True) + op = MeanFwdOp(dim=(), keepdim=True) + op_none = MeanFwdOp(dim=None, keepdim=True) assert torch.allclose(op(x), op_none(x), **_tol(torch.float16)) @@ -420,8 +420,8 @@ def test_simple_op_empty_dim_full_reduction(op_name: str) -> None: op_cls = {"amin": AminFwdOp, "amax": AmaxFwdOp, "count_nonzero": CountNonzeroFwdOp}[op_name] x = torch.randn(2, 3, 4, dtype=torch.float16, device="cuda") - y_empty = op_cls(dtype=torch.float16, dim=[])(x) - y_none = op_cls(dtype=torch.float16, dim=None)(x) + y_empty = op_cls(dim=[])(x) + y_none = op_cls(dim=None)(x) assert y_empty.shape == y_none.shape if op_name == "count_nonzero": assert (y_empty == y_none).all() @@ -436,8 +436,8 @@ def test_welford_op_empty_dim_full_reduction(op_name: str) -> None: op_cls = {"std": StdFwdOp, "var": VarFwdOp}[op_name] x = torch.randn(2, 3, 4, dtype=torch.float16, device="cuda") - y_empty = op_cls(dtype=torch.float16, dim=[], keepdim=False)(x) - y_none = op_cls(dtype=torch.float16, dim=None, keepdim=False)(x) + y_empty = op_cls(dim=[], keepdim=False)(x) + y_none = op_cls(dim=None, keepdim=False)(x) assert torch.allclose(y_empty, y_none, **_tol(torch.float16)) @@ -446,8 +446,8 @@ def test_var_mean_empty_dim_full_reduction() -> None: from tileops.ops.reduction.reduce import VarMeanFwdOp x = torch.randn(2, 3, 4, dtype=torch.float16, device="cuda") - var_e, mean_e = VarMeanFwdOp(dtype=torch.float16, dim=[], keepdim=False)(x) - var_n, mean_n = VarMeanFwdOp(dtype=torch.float16, dim=None, keepdim=False)(x) + var_e, mean_e = VarMeanFwdOp(dim=[], keepdim=False)(x) + var_n, mean_n = VarMeanFwdOp(dim=None, keepdim=False)(x) assert torch.allclose(var_e, var_n, **_tol(torch.float16)) assert torch.allclose(mean_e, mean_n, **_tol(torch.float16)) @@ -460,7 +460,7 @@ def test_prod_empty_dim_rejects() -> None: from tileops.ops.reduction.reduce import ProdFwdOp with pytest.raises(TypeError, match="ProdFwdOp.dim must be int"): - ProdFwdOp(dtype=torch.float16, dim=[], keepdim=False) + ProdFwdOp(dim=[], keepdim=False) @pytest.mark.smoke @@ -468,7 +468,7 @@ def test_logsumexp_empty_dim_rejects() -> None: from tileops.ops.reduction.softmax import LogSumExpFwdOp x = torch.randn(2, 3, 4, dtype=torch.float16, device="cuda") - op = LogSumExpFwdOp(dtype=torch.float16, dim=[], keepdim=False) + op = LogSumExpFwdOp(dim=[], keepdim=False) with pytest.raises(ValueError, match="dim=\\[\\] is not supported"): op(x) @@ -480,7 +480,7 @@ def test_all_empty_dim_is_noop() -> None: from tileops.ops.reduction.logical_reduce import AllFwdOp x = (torch.randn(2, 3, 4, device="cuda") > 0).to(torch.float16) - op = AllFwdOp(dtype=torch.float16, dim=[], keepdim=False) + op = AllFwdOp(dim=[], keepdim=False) y = op(x) assert y.shape == x.shape assert y.dtype == torch.bool @@ -493,7 +493,7 @@ def test_negative_dims_accepted() -> None: from tileops.ops.reduction.reduce import SumFwdOp x = torch.randn(4, 8, 256, dtype=torch.float16, device="cuda") - op = SumFwdOp(dtype=torch.float16, dim=[-1, 0], keepdim=False) + op = SumFwdOp(dim=[-1, 0], keepdim=False) ref = torch.sum(x.float(), dim=[0, 2], keepdim=False).to(torch.float16) y = op(x) assert y.shape == ref.shape, f"shape mismatch: {y.shape} vs {ref.shape}" @@ -506,7 +506,7 @@ def test_duplicate_dims_raises() -> None: from tileops.ops.reduction.reduce import SumFwdOp x = torch.randn(4, 8, 256, dtype=torch.float16, device="cuda") - op = SumFwdOp(dtype=torch.float16, dim=[1, 1], keepdim=False) + op = SumFwdOp(dim=[1, 1], keepdim=False) with pytest.raises(ValueError, match="Duplicate dims"): op(x) diff --git a/tests/ops/test_reduce_scalar_conformance.py b/tests/ops/test_reduce_scalar_conformance.py index 2f9e9ac5e..31c93b907 100644 --- a/tests/ops/test_reduce_scalar_conformance.py +++ b/tests/ops/test_reduce_scalar_conformance.py @@ -44,7 +44,7 @@ def test_scalar_arithmetic_reductions(dim, keepdim: bool, dtype: torch.dtype) -> ] for op_cls, torch_fn in cases: - op = op_cls(dtype=dtype, dim=dim, keepdim=keepdim) + op = op_cls(dim=dim, keepdim=keepdim) y = op(x) ref = torch_fn(x.float(), dim=dim, keepdim=keepdim).to(dtype) assert y.shape == ref.shape, ( @@ -62,7 +62,7 @@ def test_scalar_prod_reduction(dim: int, keepdim: bool, dtype: torch.dtype) -> N from tileops.ops.reduction.reduce import ProdFwdOp x = _make_scalar(dtype) - op = ProdFwdOp(dtype=dtype, dim=dim, keepdim=keepdim) + op = ProdFwdOp(dim=dim, keepdim=keepdim) y = op(x) ref = torch.prod(x.float(), dim=dim, keepdim=keepdim).to(dtype) assert y.shape == ref.shape, ( @@ -89,7 +89,7 @@ def test_scalar_welford_reductions(dim, keepdim: bool, dtype: torch.dtype) -> No ] for op_cls, ref_fn in cases: - op = op_cls(dtype=dtype, dim=dim, keepdim=keepdim) + op = op_cls(dim=dim, keepdim=keepdim) y = op(x) ref = ref_fn(x, dim=dim, keepdim=keepdim) assert y.shape == ref.shape, ( @@ -98,7 +98,7 @@ def test_scalar_welford_reductions(dim, keepdim: bool, dtype: torch.dtype) -> No ) torch.testing.assert_close(y, ref, atol=1e-4, rtol=1e-4, equal_nan=True) - op = VarMeanFwdOp(dtype=dtype, dim=dim, keepdim=keepdim) + op = VarMeanFwdOp(dim=dim, keepdim=keepdim) var_y, mean_y = op(x) ref_var, ref_mean = torch.var_mean( x.float(), dim=dim, keepdim=keepdim, correction=1, @@ -133,12 +133,12 @@ def test_invalid_dof_welford_reductions_match_pytorch( (VarFwdOp, torch.var), (StdFwdOp, torch.std), ]: - op = op_cls(dtype=torch.float32, dim=dim, keepdim=keepdim) + op = op_cls(dim=dim, keepdim=keepdim) y = op(x) ref = torch_fn(x, dim=dim, keepdim=keepdim, correction=1) torch.testing.assert_close(y, ref, atol=1e-4, rtol=1e-4, equal_nan=True) - op = VarMeanFwdOp(dtype=torch.float32, dim=dim, keepdim=keepdim) + op = VarMeanFwdOp(dim=dim, keepdim=keepdim) var_y, mean_y = op(x) ref_var, ref_mean = torch.var_mean(x, dim=dim, keepdim=keepdim, correction=1) torch.testing.assert_close(var_y, ref_var, atol=1e-4, rtol=1e-4, equal_nan=True) @@ -163,7 +163,7 @@ def test_scalar_logical_and_count_reductions( (AnyFwdOp, torch.any, torch.bool), (CountNonzeroFwdOp, torch.count_nonzero, torch.int64), ]: - op = op_cls(dtype=dtype, dim=dim) + op = op_cls(dim=dim) y = op(x) ref = torch_fn(x, dim=dim) assert y.dtype == out_dtype, f"{op_cls.__name__} scalar dtype {y.dtype}" diff --git a/tests/ops/test_reduce_variance_conformance.py b/tests/ops/test_reduce_variance_conformance.py index cf5372a95..f67a5c263 100644 --- a/tests/ops/test_reduce_variance_conformance.py +++ b/tests/ops/test_reduce_variance_conformance.py @@ -83,7 +83,7 @@ def test_var_conformance( """Each (dim-shape, correction, keepdim, dtype) cell must match torch.var.""" torch.manual_seed(0) x = torch.randn(*_SHAPE, dtype=dtype, device="cuda") - op = VarFwdOp(dtype=dtype, dim=dim, correction=correction, keepdim=keepdim) + op = VarFwdOp(dim=dim, correction=correction, keepdim=keepdim) y = op(x) ref = _ref_var(x, dim, keepdim, correction) assert y.shape == ref.shape, ( @@ -115,7 +115,7 @@ def test_std_conformance( """Each (dim-shape, correction, keepdim, dtype) cell must match torch.std.""" torch.manual_seed(0) x = torch.randn(*_SHAPE, dtype=dtype, device="cuda") - op = StdFwdOp(dtype=dtype, dim=dim, correction=correction, keepdim=keepdim) + op = StdFwdOp(dim=dim, correction=correction, keepdim=keepdim) y = op(x) ref = _ref_std(x, dim, keepdim, correction) assert y.shape == ref.shape, ( @@ -151,7 +151,7 @@ def test_var_mean_conformance( """ torch.manual_seed(0) x = torch.randn(*_SHAPE, dtype=dtype, device="cuda") - op = VarMeanFwdOp(dtype=dtype, dim=dim, correction=correction, keepdim=keepdim) + op = VarMeanFwdOp(dim=dim, correction=correction, keepdim=keepdim) out = op(x) assert isinstance(out, tuple) and len(out) == 2, ( f"VarMeanFwdOp must return a 2-tuple, got {type(out).__name__}" @@ -196,7 +196,7 @@ def test_var_unaligned_innermost(dim) -> None: torch.manual_seed(0) dtype = torch.float16 x = torch.randn(*_UNALIGNED_SHAPE, dtype=dtype, device="cuda") - op = VarFwdOp(dtype=dtype, dim=dim, correction=1, keepdim=False) + op = VarFwdOp(dim=dim, correction=1, keepdim=False) y = op(x) ref = _ref_var(x, dim, False, 1) assert y.shape == ref.shape @@ -217,7 +217,7 @@ def test_var_mean_unaligned_innermost(dim) -> None: torch.manual_seed(0) dtype = torch.float16 x = torch.randn(*_UNALIGNED_SHAPE, dtype=dtype, device="cuda") - op = VarMeanFwdOp(dtype=dtype, dim=dim, correction=1, keepdim=False) + op = VarMeanFwdOp(dim=dim, correction=1, keepdim=False) var_y, mean_y = op(x) ref_var, ref_mean = _ref_var_mean(x, dim, False, 1) assert var_y.shape == ref_var.shape @@ -240,7 +240,7 @@ def test_std_unaligned_innermost(dim) -> None: torch.manual_seed(0) dtype = torch.float16 x = torch.randn(*_UNALIGNED_SHAPE, dtype=dtype, device="cuda") - op = StdFwdOp(dtype=dtype, dim=dim, correction=1, keepdim=False) + op = StdFwdOp(dim=dim, correction=1, keepdim=False) y = op(x) ref = _ref_std(x, dim, False, 1) assert y.shape == ref.shape diff --git a/tests/ops/test_reduction_defaults.py b/tests/ops/test_reduction_defaults.py index fd2c0b716..ab61c6258 100644 --- a/tests/ops/test_reduction_defaults.py +++ b/tests/ops/test_reduction_defaults.py @@ -44,7 +44,7 @@ def test_sum_default_dim_full_reduction() -> None: from tileops.ops.reduction.reduce import SumFwdOp x = _make_float(_FLOAT_SHAPE, torch.float16) - op = SumFwdOp(dtype=torch.float16) + op = SumFwdOp() y = op(x) assert y.shape == torch.sum(x).shape @@ -54,7 +54,7 @@ def test_mean_default_dim_full_reduction() -> None: from tileops.ops.reduction.reduce import MeanFwdOp x = _make_float(_FLOAT_SHAPE, torch.float16) - op = MeanFwdOp(dtype=torch.float16) + op = MeanFwdOp() y = op(x) assert y.shape == torch.mean(x).shape @@ -64,7 +64,7 @@ def test_amax_default_dim_full_reduction() -> None: from tileops.ops.reduction.reduce import AmaxFwdOp x = _make_float(_FLOAT_SHAPE, torch.float16) - op = AmaxFwdOp(dtype=torch.float16) + op = AmaxFwdOp() y = op(x) assert y.shape == torch.amax(x).shape @@ -74,7 +74,7 @@ def test_amin_default_dim_full_reduction() -> None: from tileops.ops.reduction.reduce import AminFwdOp x = _make_float(_FLOAT_SHAPE, torch.float16) - op = AminFwdOp(dtype=torch.float16) + op = AminFwdOp() y = op(x) assert y.shape == torch.amin(x).shape @@ -84,7 +84,7 @@ def test_var_default_dim_full_reduction() -> None: from tileops.ops.reduction.reduce import VarFwdOp x = _make_float(_FLOAT_SHAPE, torch.float16) - op = VarFwdOp(dtype=torch.float16) + op = VarFwdOp() y = op(x) assert y.shape == torch.var(x).shape @@ -94,7 +94,7 @@ def test_std_default_dim_full_reduction() -> None: from tileops.ops.reduction.reduce import StdFwdOp x = _make_float(_FLOAT_SHAPE, torch.float16) - op = StdFwdOp(dtype=torch.float16) + op = StdFwdOp() y = op(x) assert y.shape == torch.std(x).shape @@ -104,7 +104,7 @@ def test_var_mean_default_dim_full_reduction() -> None: from tileops.ops.reduction.reduce import VarMeanFwdOp x = _make_float(_FLOAT_SHAPE, torch.float16) - op = VarMeanFwdOp(dtype=torch.float16) + op = VarMeanFwdOp() var_out, mean_out = op(x) ref_var, ref_mean = torch.var_mean(x) assert var_out.shape == ref_var.shape @@ -116,7 +116,7 @@ def test_all_default_dim_full_reduction() -> None: from tileops.ops.reduction.logical_reduce import AllFwdOp x = _make_logical(_LOGICAL_SHAPE, torch.float16) - op = AllFwdOp(dtype=torch.float16) + op = AllFwdOp() y = op(x) assert y.shape == torch.all(x.bool()).shape assert y.dtype == torch.bool @@ -127,7 +127,7 @@ def test_any_default_dim_full_reduction() -> None: from tileops.ops.reduction.logical_reduce import AnyFwdOp x = _make_logical(_LOGICAL_SHAPE, torch.float16) - op = AnyFwdOp(dtype=torch.float16) + op = AnyFwdOp() y = op(x) assert y.shape == torch.any(x.bool()).shape assert y.dtype == torch.bool @@ -138,7 +138,7 @@ def test_count_nonzero_default_dim_full_reduction() -> None: from tileops.ops.reduction.logical_reduce import CountNonzeroFwdOp x = _make_logical(_LOGICAL_SHAPE, torch.float16) - op = CountNonzeroFwdOp(dtype=torch.float16) + op = CountNonzeroFwdOp() y = op(x) assert y.shape == torch.count_nonzero(x).shape assert y.dtype == torch.int64 @@ -153,7 +153,7 @@ def test_prod_default_dim_last_axis() -> None: # use a narrow value range so fp16 prod is numerically stable x = torch.rand(*_FLOAT_SHAPE, dtype=torch.float16, device="cuda") * 0.01 + 0.99 - op = ProdFwdOp(dtype=torch.float16) + op = ProdFwdOp() y = op(x) assert y.shape == torch.prod(x, dim=-1).shape @@ -167,7 +167,7 @@ def test_all_empty_dim_noop(empty_dim) -> None: from tileops.ops.reduction.logical_reduce import AllFwdOp x = _make_logical(_LOGICAL_SHAPE, torch.float16) - op = AllFwdOp(dtype=torch.float16, dim=empty_dim) + op = AllFwdOp(dim=empty_dim) y = op(x) assert y.shape == x.shape assert y.dtype == torch.bool @@ -180,7 +180,7 @@ def test_any_empty_dim_noop(empty_dim) -> None: from tileops.ops.reduction.logical_reduce import AnyFwdOp x = _make_logical(_LOGICAL_SHAPE, torch.float16) - op = AnyFwdOp(dtype=torch.float16, dim=empty_dim) + op = AnyFwdOp(dim=empty_dim) y = op(x) assert y.shape == x.shape assert y.dtype == torch.bool @@ -250,7 +250,7 @@ def test_all_empty_dim_noop_rejects_cpu_tensor() -> None: from tileops.ops.reduction.logical_reduce import AllFwdOp x = (torch.randint(-1, 2, _LOGICAL_SHAPE)).to(torch.float16) # cpu - op = AllFwdOp(dtype=torch.float16, dim=[]) + op = AllFwdOp(dim=[]) with pytest.raises(ValueError, match="CUDA tensor"): op(x) @@ -260,29 +260,29 @@ def test_any_empty_dim_noop_rejects_cpu_tensor() -> None: from tileops.ops.reduction.logical_reduce import AnyFwdOp x = (torch.randint(-1, 2, _LOGICAL_SHAPE)).to(torch.float16) # cpu - op = AnyFwdOp(dtype=torch.float16, dim=[]) + op = AnyFwdOp(dim=[]) with pytest.raises(ValueError, match="CUDA tensor"): op(x) @pytest.mark.smoke -def test_all_empty_dim_noop_rejects_wrong_dtype() -> None: - """dim=[] must still validate dtype against the op's declared dtype.""" +def test_all_empty_dim_noop_rejects_undeclared_dtype() -> None: + """dim=[] must not let an input skip the manifest dtype gate.""" from tileops.ops.reduction.logical_reduce import AllFwdOp - x = _make_logical(_LOGICAL_SHAPE, torch.float32) # cuda, fp32 - op = AllFwdOp(dtype=torch.float16, dim=[]) - with pytest.raises(ValueError, match="Expected x.dtype"): + x = _make_logical(_LOGICAL_SHAPE, torch.float64) # cuda, undeclared dtype + op = AllFwdOp(dim=[]) + with pytest.raises(ValueError, match="has dtype torch.float64"): op(x) @pytest.mark.smoke -def test_any_empty_dim_noop_rejects_wrong_dtype() -> None: +def test_any_empty_dim_noop_rejects_undeclared_dtype() -> None: from tileops.ops.reduction.logical_reduce import AnyFwdOp - x = _make_logical(_LOGICAL_SHAPE, torch.float32) - op = AnyFwdOp(dtype=torch.float16, dim=[]) - with pytest.raises(ValueError, match="Expected x.dtype"): + x = _make_logical(_LOGICAL_SHAPE, torch.float64) + op = AnyFwdOp(dim=[]) + with pytest.raises(ValueError, match="has dtype torch.float64"): op(x) @@ -294,7 +294,7 @@ def test_all_empty_dim_noop_binds_roofline() -> None: from tileops.ops.reduction.logical_reduce import AllFwdOp x = _make_logical(_LOGICAL_SHAPE, torch.float16) - op = AllFwdOp(dtype=torch.float16, dim=[]) + op = AllFwdOp(dim=[]) op(x) flops, mem_bytes = op.eval_roofline() numel = x.numel() @@ -317,7 +317,7 @@ def test_any_empty_dim_noop_binds_roofline() -> None: from tileops.ops.reduction.logical_reduce import AnyFwdOp x = _make_logical(_LOGICAL_SHAPE, torch.float16) - op = AnyFwdOp(dtype=torch.float16, dim=[]) + op = AnyFwdOp(dim=[]) op(x) flops, mem_bytes = op.eval_roofline() numel = x.numel() @@ -335,7 +335,7 @@ def test_validate_dim_rejects_bool_scalar() -> None: from tileops.ops.reduction.reduce import SumFwdOp with pytest.raises(TypeError, match="dim must not be bool"): - SumFwdOp(dtype=torch.float16, dim=True) + SumFwdOp(dim=True) @pytest.mark.smoke @@ -344,4 +344,4 @@ def test_validate_dim_rejects_bool_in_list() -> None: from tileops.ops.reduction.reduce import SumFwdOp with pytest.raises(TypeError, match="must be int .not bool"): - SumFwdOp(dtype=torch.float16, dim=[True, 0]) + SumFwdOp(dim=[True, 0]) diff --git a/tests/ops/test_reduction_scalar_input.py b/tests/ops/test_reduction_scalar_input.py index 74b56bfbc..e53d9fd40 100644 --- a/tests/ops/test_reduction_scalar_input.py +++ b/tests/ops/test_reduction_scalar_input.py @@ -42,7 +42,7 @@ def test_sum_scalar_input(dim) -> None: from tileops.ops.reduction.reduce import SumFwdOp x = torch.tensor(3.5, dtype=torch.float32, device="cuda") - op = SumFwdOp(dtype=torch.float32, dim=dim) + op = SumFwdOp(dim=dim) y = op(x) ref = torch.sum(x, dim=dim) if dim is not None else torch.sum(x) assert y.shape == ref.shape @@ -55,7 +55,7 @@ def test_mean_scalar_input(dim) -> None: from tileops.ops.reduction.reduce import MeanFwdOp x = torch.tensor(2.0, dtype=torch.float32, device="cuda") - op = MeanFwdOp(dtype=torch.float32, dim=dim) + op = MeanFwdOp(dim=dim) y = op(x) ref = torch.mean(x, dim=dim) if dim is not None else torch.mean(x) assert y.shape == ref.shape @@ -68,7 +68,7 @@ def test_amax_scalar_input(dim) -> None: from tileops.ops.reduction.reduce import AmaxFwdOp x = torch.tensor(-1.5, dtype=torch.float32, device="cuda") - op = AmaxFwdOp(dtype=torch.float32, dim=dim) + op = AmaxFwdOp(dim=dim) y = op(x) ref = torch.amax(x, dim=dim) if dim is not None else torch.amax(x) assert y.shape == ref.shape @@ -81,7 +81,7 @@ def test_amin_scalar_input(dim) -> None: from tileops.ops.reduction.reduce import AminFwdOp x = torch.tensor(4.25, dtype=torch.float32, device="cuda") - op = AminFwdOp(dtype=torch.float32, dim=dim) + op = AminFwdOp(dim=dim) y = op(x) ref = torch.amin(x, dim=dim) if dim is not None else torch.amin(x) assert y.shape == ref.shape @@ -94,7 +94,7 @@ def test_prod_scalar_input(dim) -> None: from tileops.ops.reduction.reduce import ProdFwdOp x = torch.tensor(3.0, dtype=torch.float32, device="cuda") - op = ProdFwdOp(dtype=torch.float32, dim=dim) + op = ProdFwdOp(dim=dim) y = op(x) ref = torch.prod(x, dim=dim) assert y.shape == ref.shape @@ -107,7 +107,7 @@ def test_all_scalar_input(dim) -> None: from tileops.ops.reduction.logical_reduce import AllFwdOp x = torch.tensor(1.0, dtype=torch.float32, device="cuda") - op = AllFwdOp(dtype=torch.float32, dim=dim) + op = AllFwdOp(dim=dim) y = op(x) ref = torch.all(x, dim=dim) if dim is not None else torch.all(x) assert y.shape == ref.shape @@ -121,7 +121,7 @@ def test_any_scalar_input(dim) -> None: from tileops.ops.reduction.logical_reduce import AnyFwdOp x = torch.tensor(0.0, dtype=torch.float32, device="cuda") - op = AnyFwdOp(dtype=torch.float32, dim=dim) + op = AnyFwdOp(dim=dim) y = op(x) ref = torch.any(x, dim=dim) if dim is not None else torch.any(x) assert y.shape == ref.shape @@ -135,7 +135,7 @@ def test_count_nonzero_scalar_input(dim) -> None: from tileops.ops.reduction.logical_reduce import CountNonzeroFwdOp x = torch.tensor(2.5, dtype=torch.float32, device="cuda") - op = CountNonzeroFwdOp(dtype=torch.float32, dim=dim) + op = CountNonzeroFwdOp(dim=dim) y = op(x) ref = torch.count_nonzero(x, dim=dim) if dim is not None else torch.count_nonzero(x) assert y.shape == ref.shape @@ -166,7 +166,7 @@ def test_var_scalar_input(dim) -> None: from tileops.ops.reduction.reduce import VarFwdOp x = torch.tensor(1.5, dtype=torch.float32, device="cuda") - op = VarFwdOp(dtype=torch.float32, dim=dim) + op = VarFwdOp(dim=dim) expect_warn = _expect_var_warning() with warnings.catch_warnings(record=True) as op_caught: warnings.simplefilter("always") @@ -185,7 +185,7 @@ def test_std_scalar_input(dim) -> None: from tileops.ops.reduction.reduce import StdFwdOp x = torch.tensor(-0.75, dtype=torch.float32, device="cuda") - op = StdFwdOp(dtype=torch.float32, dim=dim) + op = StdFwdOp(dim=dim) expect_warn = _expect_var_warning() with warnings.catch_warnings(record=True) as op_caught: warnings.simplefilter("always") @@ -204,7 +204,7 @@ def test_var_mean_scalar_input(dim) -> None: from tileops.ops.reduction.reduce import VarMeanFwdOp x = torch.tensor(2.25, dtype=torch.float32, device="cuda") - op = VarMeanFwdOp(dtype=torch.float32, dim=dim) + op = VarMeanFwdOp(dim=dim) expect_warn = _expect_var_warning() with warnings.catch_warnings(record=True) as op_caught: warnings.simplefilter("always") @@ -238,7 +238,7 @@ def test_sum_scalar_duplicate_dim_matches_torch(dim) -> None: x = torch.tensor(1.5, dtype=torch.float32, device="cuda") with pytest.raises(RuntimeError, match="appears multiple times"): torch.sum(x, dim=list(dim)) - op = SumFwdOp(dtype=torch.float32, dim=list(dim)) + op = SumFwdOp(dim=list(dim)) with pytest.raises(RuntimeError, match="appears multiple times"): op(x) @@ -252,7 +252,7 @@ def test_var_scalar_requires_grad_preserves_grad_fn() -> None: from tileops.ops.reduction.reduce import VarFwdOp x = torch.tensor(0.5, dtype=torch.float32, device="cuda", requires_grad=True) - op = VarFwdOp(dtype=torch.float32, dim=None) + op = VarFwdOp(dim=None) with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) y = op(x) @@ -268,7 +268,7 @@ def test_var_mean_scalar_requires_grad_preserves_grad_fn() -> None: from tileops.ops.reduction.reduce import VarMeanFwdOp x = torch.tensor(1.25, dtype=torch.float32, device="cuda", requires_grad=True) - op = VarMeanFwdOp(dtype=torch.float32, dim=None) + op = VarMeanFwdOp(dim=None) with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) var_out, mean_out = op(x) diff --git a/tests/ops/test_rms_norm.py b/tests/ops/test_rms_norm.py index 7fff88221..4b7716815 100644 --- a/tests/ops/test_rms_norm.py +++ b/tests/ops/test_rms_norm.py @@ -38,7 +38,7 @@ class RMSNormFixture(FixtureBase): @RMSNormFixture def test_rms_norm_op(m: int, n: int, dtype: torch.dtype, tune: bool) -> None: test = RMSNormTest(m, n, dtype) - op = RMSNormFwdOp(normalized_shape=(n,), dtype=dtype) + op = RMSNormFwdOp(normalized_shape=(n,)) atol = 1e-2 if dtype == torch.float16 else 1.6e-2 rtol = atol test.check(op, *test.gen_inputs(), atol=atol, rtol=rtol) @@ -60,7 +60,7 @@ def test_rms_norm_non_contiguous(m: int, n: int, dtype: torch.dtype) -> None: x = x_full[:, :n] # non-contiguous slice weight = torch.randn(n, dtype=dtype, device="cuda") - op = RMSNormFwdOp(normalized_shape=(n,), dtype=dtype) + op = RMSNormFwdOp(normalized_shape=(n,)) # Reference on contiguous copy eps = 1e-6 @@ -90,7 +90,7 @@ def test_rms_norm_3d(batch: int, seq: int, hidden: int, dtype: torch.dtype) -> N x = torch.randn(batch, seq, hidden, dtype=dtype, device="cuda") weight = torch.randn(hidden, dtype=dtype, device="cuda") - op = RMSNormFwdOp(normalized_shape=(hidden,), dtype=dtype) + op = RMSNormFwdOp(normalized_shape=(hidden,)) # Reference eps = 1e-6 diff --git a/tests/ops/test_softmax.py b/tests/ops/test_softmax.py index 445a82773..6570b508c 100644 --- a/tests/ops/test_softmax.py +++ b/tests/ops/test_softmax.py @@ -304,7 +304,7 @@ def __init__(self, shape: tuple, dtype: torch.dtype, dim: int = -1): @LogSumExpFixture def test_logsumexp_op(shape: tuple, dim: int, dtype: torch.dtype, tune: bool) -> None: test = LogSumExpTest(shape, dtype, dim=dim) - op = LogSumExpFwdOp(dtype=dtype, dim=dim, tune=tune) + op = LogSumExpFwdOp(dim=dim, tune=tune) atol, rtol = _get_tolerances(dtype) test.check(op, *test.gen_inputs(), atol=atol, rtol=rtol) @@ -335,7 +335,7 @@ class LogSumExpKeepdimFixture(FixtureBase): def test_logsumexp_keepdim(shape: tuple, dim: int, dtype: torch.dtype) -> None: """Test logsumexp with keepdim=True — output retains reduced dim as size 1.""" x = torch.randn(*shape, dtype=dtype, device="cuda") - op = LogSumExpFwdOp(dtype=dtype, dim=dim, keepdim=True) + op = LogSumExpFwdOp(dim=dim, keepdim=True) y_ref = torch.logsumexp(x.float(), dim=dim, keepdim=True).to(dtype) y = op(x) @@ -405,7 +405,7 @@ def test_logsumexp_non_contiguous(shape: tuple, dtype: torch.dtype) -> None: x_full = torch.randn(m, n * 2, dtype=dtype, device="cuda") x = x_full[:, :n] - op = LogSumExpFwdOp(dtype=dtype, dim=-1) + op = LogSumExpFwdOp(dim=-1) y_ref = torch.logsumexp(x.float().contiguous(), dim=-1).to(dtype) y = op(x) @@ -468,7 +468,7 @@ class LogSumExp1DFixture(FixtureBase): def test_logsumexp_1d(n: int, dtype: torch.dtype) -> None: """Test logsumexp with 1D input -- output should be a scalar.""" x = torch.randn(n, dtype=dtype, device="cuda") - op = LogSumExpFwdOp(dtype=dtype, dim=-1) + op = LogSumExpFwdOp(dim=-1) y_ref = torch.logsumexp(x.float(), dim=-1).to(dtype) y = op(x) @@ -508,7 +508,7 @@ def test_log_softmax_rejects_multidim_before_kernel() -> None: def test_logsumexp_accepts_multidim() -> None: """LogSumExpFwdOp must accept list dim without error (multi-dim is supported).""" x = torch.randn(4, 8, device="cuda", dtype=torch.float32) - op = LogSumExpFwdOp(dtype=torch.float32, dim=[0, 1]) + op = LogSumExpFwdOp(dim=[0, 1]) y = op(x) y_ref = torch.logsumexp(x.float(), dim=[0, 1]) assert torch.allclose(y, y_ref, atol=1e-5, rtol=1e-5) diff --git a/tests/ops/test_vector_norm.py b/tests/ops/test_vector_norm.py index 4862e470f..498e23a68 100644 --- a/tests/ops/test_vector_norm.py +++ b/tests/ops/test_vector_norm.py @@ -146,7 +146,6 @@ def _make_1d_input(n: int, dtype: torch.dtype) -> torch.Tensor: def _make_op( - dtype: torch.dtype, op_kind: str, dim: int = -1, keepdim: bool = False, @@ -162,7 +161,7 @@ def _make_op( "inf": InfNormFwdOp, } cls = op_map[op_kind] - return cls(dtype=dtype, dim=dim, keepdim=keepdim, kernel_map=kernel_map, tune=tune) + return cls(dim=dim, keepdim=keepdim, kernel_map=kernel_map, tune=tune) # L1NormFwdOp tests @@ -171,7 +170,7 @@ def _make_op( @VectorNormBasicFixture def test_l1_norm_op(m: int, n: int, dtype: torch.dtype) -> None: test = VectorNormTest(m, n, dtype, "l1") - op = _make_op(dtype, "l1") + op = _make_op("l1") atol, rtol = _get_tolerances(dtype) test.check(op, *test.gen_inputs(), atol=atol, rtol=rtol) @@ -180,7 +179,7 @@ def test_l1_norm_op(m: int, n: int, dtype: torch.dtype) -> None: def test_l1_non_contiguous(m: int, n: int, dtype: torch.dtype) -> None: x_full = _make_noncontig_input(m, n, dtype) x = x_full[:, :n] - op = _make_op(dtype, "l1") + op = _make_op("l1") ref = torch.linalg.vector_norm(x.float().contiguous(), ord=1, dim=-1).to(dtype) y = op(x) atol, rtol = _get_tolerances(dtype) @@ -190,7 +189,7 @@ def test_l1_non_contiguous(m: int, n: int, dtype: torch.dtype) -> None: @VectorNorm3DFixture def test_l1_3d(batch: int, seq: int, hidden: int, dtype: torch.dtype) -> None: x = torch.randn(batch, seq, hidden, dtype=dtype, device="cuda") - op = _make_op(dtype, "l1") + op = _make_op("l1") ref = torch.linalg.vector_norm(x.float(), ord=1, dim=-1).to(dtype) y = op(x) atol, rtol = _get_tolerances(dtype) @@ -200,7 +199,7 @@ def test_l1_3d(batch: int, seq: int, hidden: int, dtype: torch.dtype) -> None: @VectorNorm4DFixture def test_l1_4d(b0: int, b1: int, b2: int, n: int, dtype: torch.dtype) -> None: x = torch.randn(b0, b1, b2, n, dtype=dtype, device="cuda") - op = _make_op(dtype, "l1") + op = _make_op("l1") ref = torch.linalg.vector_norm(x.float(), ord=1, dim=-1).to(dtype) y = op(x) atol, rtol = _get_tolerances(dtype) @@ -210,7 +209,7 @@ def test_l1_4d(b0: int, b1: int, b2: int, n: int, dtype: torch.dtype) -> None: @VectorNorm1DFixture def test_l1_1d(n: int, dtype: torch.dtype) -> None: x = _make_1d_input(n, dtype) - op = _make_op(dtype, "l1") + op = _make_op("l1") ref = torch.linalg.vector_norm(x.float(), ord=1, dim=-1).to(dtype) y = op(x) atol, rtol = _get_tolerances(dtype) @@ -223,7 +222,7 @@ def test_l1_1d(n: int, dtype: torch.dtype) -> None: @VectorNormBasicFixture def test_l2_norm_op(m: int, n: int, dtype: torch.dtype) -> None: test = VectorNormTest(m, n, dtype, "l2") - op = _make_op(dtype, "l2") + op = _make_op("l2") atol, rtol = _get_tolerances(dtype) test.check(op, *test.gen_inputs(), atol=atol, rtol=rtol) @@ -232,7 +231,7 @@ def test_l2_norm_op(m: int, n: int, dtype: torch.dtype) -> None: def test_l2_non_contiguous(m: int, n: int, dtype: torch.dtype) -> None: x_full = _make_noncontig_input(m, n, dtype) x = x_full[:, :n] - op = _make_op(dtype, "l2") + op = _make_op("l2") ref = torch.linalg.vector_norm(x.float().contiguous(), ord=2, dim=-1).to(dtype) y = op(x) atol, rtol = _get_tolerances(dtype) @@ -242,7 +241,7 @@ def test_l2_non_contiguous(m: int, n: int, dtype: torch.dtype) -> None: @VectorNorm3DFixture def test_l2_3d(batch: int, seq: int, hidden: int, dtype: torch.dtype) -> None: x = torch.randn(batch, seq, hidden, dtype=dtype, device="cuda") - op = _make_op(dtype, "l2") + op = _make_op("l2") ref = torch.linalg.vector_norm(x.float(), ord=2, dim=-1).to(dtype) y = op(x) atol, rtol = _get_tolerances(dtype) @@ -252,7 +251,7 @@ def test_l2_3d(batch: int, seq: int, hidden: int, dtype: torch.dtype) -> None: @VectorNorm4DFixture def test_l2_4d(b0: int, b1: int, b2: int, n: int, dtype: torch.dtype) -> None: x = torch.randn(b0, b1, b2, n, dtype=dtype, device="cuda") - op = _make_op(dtype, "l2") + op = _make_op("l2") ref = torch.linalg.vector_norm(x.float(), ord=2, dim=-1).to(dtype) y = op(x) atol, rtol = _get_tolerances(dtype) @@ -262,7 +261,7 @@ def test_l2_4d(b0: int, b1: int, b2: int, n: int, dtype: torch.dtype) -> None: @VectorNorm1DFixture def test_l2_1d(n: int, dtype: torch.dtype) -> None: x = _make_1d_input(n, dtype) - op = _make_op(dtype, "l2") + op = _make_op("l2") ref = torch.linalg.vector_norm(x.float(), ord=2, dim=-1).to(dtype) y = op(x) atol, rtol = _get_tolerances(dtype) @@ -275,7 +274,7 @@ def test_l2_1d(n: int, dtype: torch.dtype) -> None: @VectorNormBasicFixture def test_inf_norm_op(m: int, n: int, dtype: torch.dtype) -> None: test = VectorNormTest(m, n, dtype, "inf") - op = _make_op(dtype, "inf") + op = _make_op("inf") atol, rtol = _get_tolerances(dtype) test.check(op, *test.gen_inputs(), atol=atol, rtol=rtol) @@ -284,7 +283,7 @@ def test_inf_norm_op(m: int, n: int, dtype: torch.dtype) -> None: def test_inf_non_contiguous(m: int, n: int, dtype: torch.dtype) -> None: x_full = _make_noncontig_input(m, n, dtype) x = x_full[:, :n] - op = _make_op(dtype, "inf") + op = _make_op("inf") ref = torch.linalg.vector_norm(x.float().contiguous(), ord=float("inf"), dim=-1).to(dtype) y = op(x) atol, rtol = _get_tolerances(dtype) @@ -294,7 +293,7 @@ def test_inf_non_contiguous(m: int, n: int, dtype: torch.dtype) -> None: @VectorNorm3DFixture def test_inf_3d(batch: int, seq: int, hidden: int, dtype: torch.dtype) -> None: x = torch.randn(batch, seq, hidden, dtype=dtype, device="cuda") - op = _make_op(dtype, "inf") + op = _make_op("inf") ref = torch.linalg.vector_norm(x.float(), ord=float("inf"), dim=-1).to(dtype) y = op(x) atol, rtol = _get_tolerances(dtype) @@ -304,7 +303,7 @@ def test_inf_3d(batch: int, seq: int, hidden: int, dtype: torch.dtype) -> None: @VectorNorm4DFixture def test_inf_4d(b0: int, b1: int, b2: int, n: int, dtype: torch.dtype) -> None: x = torch.randn(b0, b1, b2, n, dtype=dtype, device="cuda") - op = _make_op(dtype, "inf") + op = _make_op("inf") ref = torch.linalg.vector_norm(x.float(), ord=float("inf"), dim=-1).to(dtype) y = op(x) atol, rtol = _get_tolerances(dtype) @@ -314,7 +313,7 @@ def test_inf_4d(b0: int, b1: int, b2: int, n: int, dtype: torch.dtype) -> None: @VectorNorm1DFixture def test_inf_1d(n: int, dtype: torch.dtype) -> None: x = _make_1d_input(n, dtype) - op = _make_op(dtype, "inf") + op = _make_op("inf") ref = torch.linalg.vector_norm(x.float(), ord=float("inf"), dim=-1).to(dtype) y = op(x) atol, rtol = _get_tolerances(dtype) @@ -348,7 +347,7 @@ def test_inf_nan_propagation(m: int, n: int, dtype: torch.dtype) -> None: x[1, -1] = float("nan") # Rows 2+ remain finite - op = _make_op(dtype, "inf") + op = _make_op("inf") ref = torch.linalg.vector_norm(x.float(), ord=float("inf"), dim=-1).to(dtype) y = op(x) @@ -383,7 +382,7 @@ class VectorNormSpecFixture(FixtureBase): def test_spec_dim0(op_kind: str, dtype: torch.dtype) -> None: """Reduce along dim=0.""" x = torch.randn(64, 512, dtype=dtype, device="cuda") - op = _make_op(dtype, op_kind, dim=0) + op = _make_op(op_kind, dim=0) ord_val = _ORD_MAP[op_kind] ref = torch.linalg.vector_norm(x.float(), ord=ord_val, dim=0).to(dtype) y = op(x) @@ -395,7 +394,7 @@ def test_spec_dim0(op_kind: str, dtype: torch.dtype) -> None: def test_spec_dim1_3d(op_kind: str, dtype: torch.dtype) -> None: """Reduce along dim=1 of a 3D tensor.""" x = torch.randn(4, 64, 512, dtype=dtype, device="cuda") - op = _make_op(dtype, op_kind, dim=1) + op = _make_op(op_kind, dim=1) ord_val = _ORD_MAP[op_kind] ref = torch.linalg.vector_norm(x.float(), ord=ord_val, dim=1).to(dtype) y = op(x) @@ -407,7 +406,7 @@ def test_spec_dim1_3d(op_kind: str, dtype: torch.dtype) -> None: def test_spec_keepdim(op_kind: str, dtype: torch.dtype) -> None: """keepdim=True preserves the reduced dimension as size 1.""" x = torch.randn(32, 512, dtype=dtype, device="cuda") - op = _make_op(dtype, op_kind, keepdim=True) + op = _make_op(op_kind, keepdim=True) ord_val = _ORD_MAP[op_kind] ref = torch.linalg.vector_norm(x.float(), ord=ord_val, dim=-1, keepdim=True).to(dtype) y = op(x) @@ -420,7 +419,7 @@ def test_spec_keepdim(op_kind: str, dtype: torch.dtype) -> None: def test_spec_dim0_keepdim(op_kind: str, dtype: torch.dtype) -> None: """dim=0 + keepdim=True.""" x = torch.randn(64, 512, dtype=dtype, device="cuda") - op = _make_op(dtype, op_kind, dim=0, keepdim=True) + op = _make_op(op_kind, dim=0, keepdim=True) ord_val = _ORD_MAP[op_kind] ref = torch.linalg.vector_norm(x.float(), ord=ord_val, dim=0, keepdim=True).to(dtype) y = op(x) @@ -459,7 +458,7 @@ class _Fixture(FixtureBase): @_DtypeSmoke_float16 def test_l1_smoke_float16(m: int, n: int, dtype: torch.dtype) -> None: test = VectorNormTest(m, n, dtype, "l1") - op = _make_op(dtype, "l1") + op = _make_op("l1") atol, rtol = _get_tolerances(dtype) test.check(op, *test.gen_inputs(), atol=atol, rtol=rtol) @@ -467,7 +466,7 @@ def test_l1_smoke_float16(m: int, n: int, dtype: torch.dtype) -> None: @_DtypeSmoke_bfloat16 def test_l1_smoke_bfloat16(m: int, n: int, dtype: torch.dtype) -> None: test = VectorNormTest(m, n, dtype, "l1") - op = _make_op(dtype, "l1") + op = _make_op("l1") atol, rtol = _get_tolerances(dtype) test.check(op, *test.gen_inputs(), atol=atol, rtol=rtol) @@ -475,7 +474,7 @@ def test_l1_smoke_bfloat16(m: int, n: int, dtype: torch.dtype) -> None: @_DtypeSmoke_float32 def test_l1_smoke_float32(m: int, n: int, dtype: torch.dtype) -> None: test = VectorNormTest(m, n, dtype, "l1") - op = _make_op(dtype, "l1") + op = _make_op("l1") atol, rtol = _get_tolerances(dtype) test.check(op, *test.gen_inputs(), atol=atol, rtol=rtol) @@ -483,7 +482,7 @@ def test_l1_smoke_float32(m: int, n: int, dtype: torch.dtype) -> None: @_DtypeSmoke_float16 def test_l2_smoke_float16(m: int, n: int, dtype: torch.dtype) -> None: test = VectorNormTest(m, n, dtype, "l2") - op = _make_op(dtype, "l2") + op = _make_op("l2") atol, rtol = _get_tolerances(dtype) test.check(op, *test.gen_inputs(), atol=atol, rtol=rtol) @@ -491,7 +490,7 @@ def test_l2_smoke_float16(m: int, n: int, dtype: torch.dtype) -> None: @_DtypeSmoke_bfloat16 def test_l2_smoke_bfloat16(m: int, n: int, dtype: torch.dtype) -> None: test = VectorNormTest(m, n, dtype, "l2") - op = _make_op(dtype, "l2") + op = _make_op("l2") atol, rtol = _get_tolerances(dtype) test.check(op, *test.gen_inputs(), atol=atol, rtol=rtol) @@ -499,7 +498,7 @@ def test_l2_smoke_bfloat16(m: int, n: int, dtype: torch.dtype) -> None: @_DtypeSmoke_float32 def test_l2_smoke_float32(m: int, n: int, dtype: torch.dtype) -> None: test = VectorNormTest(m, n, dtype, "l2") - op = _make_op(dtype, "l2") + op = _make_op("l2") atol, rtol = _get_tolerances(dtype) test.check(op, *test.gen_inputs(), atol=atol, rtol=rtol) @@ -507,7 +506,7 @@ def test_l2_smoke_float32(m: int, n: int, dtype: torch.dtype) -> None: @_DtypeSmoke_float16 def test_inf_smoke_float16(m: int, n: int, dtype: torch.dtype) -> None: test = VectorNormTest(m, n, dtype, "inf") - op = _make_op(dtype, "inf") + op = _make_op("inf") atol, rtol = _get_tolerances(dtype) test.check(op, *test.gen_inputs(), atol=atol, rtol=rtol) @@ -515,7 +514,7 @@ def test_inf_smoke_float16(m: int, n: int, dtype: torch.dtype) -> None: @_DtypeSmoke_bfloat16 def test_inf_smoke_bfloat16(m: int, n: int, dtype: torch.dtype) -> None: test = VectorNormTest(m, n, dtype, "inf") - op = _make_op(dtype, "inf") + op = _make_op("inf") atol, rtol = _get_tolerances(dtype) test.check(op, *test.gen_inputs(), atol=atol, rtol=rtol) @@ -523,7 +522,7 @@ def test_inf_smoke_bfloat16(m: int, n: int, dtype: torch.dtype) -> None: @_DtypeSmoke_float32 def test_inf_smoke_float32(m: int, n: int, dtype: torch.dtype) -> None: test = VectorNormTest(m, n, dtype, "inf") - op = _make_op(dtype, "inf") + op = _make_op("inf") atol, rtol = _get_tolerances(dtype) test.check(op, *test.gen_inputs(), atol=atol, rtol=rtol) @@ -537,7 +536,7 @@ def test_inf_smoke_float32(m: int, n: int, dtype: torch.dtype) -> None: def test_empty_dim_full_reduction_keepdim(op_kind: str, keepdim: bool) -> None: dtype = torch.float16 x = torch.randn(32, 256, dtype=dtype, device="cuda") - op = _make_op(dtype, op_kind, dim=[], keepdim=keepdim) + op = _make_op(op_kind, dim=[], keepdim=keepdim) ref = torch.linalg.vector_norm( x.float(), ord=_ORD_MAP[op_kind], dim=[], keepdim=keepdim, ).to(dtype) @@ -556,7 +555,7 @@ def test_empty_dim_full_reduction_3d_dtypes( op_kind: str, dtype: torch.dtype, ) -> None: x = torch.randn(2, 16, 128, dtype=dtype, device="cuda") - op = _make_op(dtype, op_kind, dim=[], keepdim=False) + op = _make_op(op_kind, dim=[], keepdim=False) ref = torch.linalg.vector_norm( x.float(), ord=_ORD_MAP[op_kind], dim=[], keepdim=False, ).to(dtype) @@ -572,14 +571,12 @@ def test_vector_norm_long_sequence_tiled(op_kind: str) -> None: """Exercise the N-tiled path with a tail-M block.""" dtype = torch.bfloat16 test = VectorNormTest(3, 33024, dtype, op_kind) - op = _make_op( - dtype, - op_kind, + op = _make_op(op_kind, kernel_map={"vector_norm": _TailBlockVectorNormKernel}, ) atol, rtol = _get_tolerances(dtype) test.check(op, *test.gen_inputs(), atol=atol, rtol=rtol) - kernel = op._kernel_cache[(3, 33024)] + kernel = op._kernel_cache[(3, 33024, dtype)] assert kernel.config["block_m"] > test.shape[0] assert kernel.config["tile_n"] > 0 @@ -593,11 +590,11 @@ def test_vector_norm_tiled_autotune() -> None: """ m, n, dtype = 4, 40000, torch.float16 test = VectorNormTest(m, n, dtype, "l2") - op = _make_op(dtype, "l2", tune=True) + op = _make_op("l2", tune=True) atol, rtol = _get_tolerances(dtype) test.check(op, *test.gen_inputs(), atol=atol, rtol=rtol) - kernel = op._kernel_cache[(m, n)] + kernel = op._kernel_cache[(m, n, dtype)] assert kernel._needs_tiling assert kernel.config in kernel.autotune_configs diff --git a/tests/ops/test_welford_non_aligned.py b/tests/ops/test_welford_non_aligned.py index 9506adb24..0f3b96c69 100644 --- a/tests/ops/test_welford_non_aligned.py +++ b/tests/ops/test_welford_non_aligned.py @@ -185,7 +185,7 @@ def test_var_non_aligned(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.reduce import VarFwdOp test = WelfordNonAlignedTest((m, n), dtype, "var", correction=1) - op = VarFwdOp(dim=-1, dtype=dtype) + op = VarFwdOp(dim=-1) test.check(op, *test.gen_inputs(), **_tol(dtype)) @@ -198,7 +198,7 @@ def test_std_non_aligned(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.reduce import StdFwdOp test = WelfordNonAlignedTest((m, n), dtype, "std", correction=1) - op = StdFwdOp(dim=-1, dtype=dtype) + op = StdFwdOp(dim=-1) test.check(op, *test.gen_inputs(), **_tol(dtype)) @@ -211,7 +211,7 @@ def test_var_mean_non_aligned(m: int, n: int, dtype: torch.dtype) -> None: from tileops.ops.reduction.reduce import VarMeanFwdOp test = WelfordNonAlignedTest((m, n), dtype, "var_mean", correction=1) - op = VarMeanFwdOp(dim=-1, dtype=dtype) + op = VarMeanFwdOp(dim=-1) test.check(op, *test.gen_inputs(), **_tol(dtype)) @@ -224,7 +224,7 @@ def test_var_3d_non_aligned(batch: int, seq: int, hidden: int, dtype: torch.dtyp from tileops.ops.reduction.reduce import VarFwdOp x = torch.randn(batch, seq, hidden, dtype=dtype, device="cuda") - op = VarFwdOp(dim=-1, dtype=dtype) + op = VarFwdOp(dim=-1) ref = x.float().var(dim=-1, correction=1).to(dtype) y = op(x) tol = _tol(dtype) @@ -237,7 +237,7 @@ def test_std_3d_non_aligned(batch: int, seq: int, hidden: int, dtype: torch.dtyp from tileops.ops.reduction.reduce import StdFwdOp x = torch.randn(batch, seq, hidden, dtype=dtype, device="cuda") - op = StdFwdOp(dim=-1, dtype=dtype) + op = StdFwdOp(dim=-1) ref = x.float().std(dim=-1, correction=1).to(dtype) y = op(x) tol = _tol(dtype) @@ -250,7 +250,7 @@ def test_var_mean_3d_non_aligned(batch: int, seq: int, hidden: int, dtype: torch from tileops.ops.reduction.reduce import VarMeanFwdOp x = torch.randn(batch, seq, hidden, dtype=dtype, device="cuda") - op = VarMeanFwdOp(dim=-1, dtype=dtype, correction=1) + op = VarMeanFwdOp(dim=-1, correction=1) ref_var = x.float().var(dim=-1, correction=1).to(dtype) ref_mean = x.float().mean(dim=-1).to(dtype) var_out, mean_out = op(x) @@ -274,7 +274,7 @@ def test_var_multidim_non_aligned( from tileops.ops.reduction.reduce import VarFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = VarFwdOp(dtype=dtype, dim=dims, keepdim=keepdim) + op = VarFwdOp(dim=dims, keepdim=keepdim) ref = torch.var(x.float(), dim=dims, keepdim=keepdim, correction=1).to(dtype) y = op(x) tol = _tol(dtype) @@ -292,7 +292,7 @@ def test_std_multidim_non_aligned( from tileops.ops.reduction.reduce import StdFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = StdFwdOp(dtype=dtype, dim=dims, keepdim=keepdim) + op = StdFwdOp(dim=dims, keepdim=keepdim) ref = torch.std(x.float(), dim=dims, keepdim=keepdim, correction=1).to(dtype) y = op(x) tol = _tol(dtype) @@ -310,7 +310,7 @@ def test_var_mean_multidim_non_aligned( from tileops.ops.reduction.reduce import VarMeanFwdOp x = torch.randn(*shape, dtype=dtype, device="cuda") - op = VarMeanFwdOp(dtype=dtype, dim=dims, keepdim=keepdim) + op = VarMeanFwdOp(dim=dims, keepdim=keepdim) ref_var = torch.var(x.float(), dim=dims, keepdim=keepdim, correction=1).to(dtype) ref_mean = torch.mean(x.float(), dim=dims, keepdim=keepdim).to(dtype) var_out, mean_out = op(x) diff --git a/tileops/ops/_dtype_codegen.py b/tileops/ops/_dtype_codegen.py index c9b755321..588f0c342 100644 --- a/tileops/ops/_dtype_codegen.py +++ b/tileops/ops/_dtype_codegen.py @@ -247,43 +247,48 @@ def synthesize_validate_dtypes( "op_name": op_name, } params_src = ", ".join(input_names) + # Unrolled per input, with each parameter referenced by name. A `locals()` + # lookup or a loop over `input_names` would read the same values, but + # `torch.compile` cannot trace `locals()`, and this body runs inside + # `forward()` — so an op that validates dtypes would lose `fullgraph`. src_lines = [ f"def _validate_dtypes(self, {params_src}):", f' """Synthesized from manifest signature for {op_name}."""', - " _locals = locals()", - " for _name in input_names:", - " _concrete, _refs, _dtype_str = per_input[_name]", - " _actual = _locals[_name].dtype", - " if _actual in _concrete:", - " continue", - " _matched = False", - " for _ref in _refs:", - " _ref_tensor = _locals.get(_ref)", - " if _ref_tensor is None:", - " raise ValueError(", - " f\"{op_name}: input {_name!r} declares \"", - " f\"same_as({_ref}) but {_ref!r} was not supplied\"", - " )", - " if _actual == _ref_tensor.dtype:", - " _matched = True", - " break", - " if _matched:", - " continue", - " raise ValueError(", - " f\"{op_name}: input {_name!r} has dtype {_actual}, \"", - " f\"expected {_dtype_str!r}\"", - " )", - " if combo_keys is not None:", - " _observed = tuple(_locals[_n].dtype for _n in input_names)", - " if _observed not in combo_keys:", - " _pairs = \", \".join(", - " f\"{_n}={_d}\" for _n, _d in zip(input_names, _observed)", - " )", - " raise ValueError(", - " f\"{op_name}: dtype combination ({_pairs}) is not \"", - " f\"listed in signature.dtype_combos\"", - " )", ] + for name in input_names: + concrete, refs, dtype_str = per_input[name] + closure[f"_concrete_{name}"] = frozenset(concrete) + closure[f"_dtype_str_{name}"] = dtype_str + src_lines.append(f" _actual = {name}.dtype") + src_lines.append(f" if _actual not in _concrete_{name}:") + # Every `same_as(ref)` was checked above to name a sibling input, so + # each ref is in scope as a parameter here. + if refs: + cond = " or ".join(f"_actual == {r}.dtype" for r in refs) + src_lines.append(f" if not ({cond}):") + indent = " " + else: + indent = " " + src_lines += [ + f"{indent}raise ValueError(", + f'{indent} f"{{op_name}}: input {name!r} has dtype {{_actual}}, "', + f"{indent} f\"expected {{_dtype_str_{name}!r}}\"", + f"{indent})", + ] + if combo_keys is not None: + observed = ", ".join(f"{n}.dtype" for n in input_names) + trailing = "," if len(input_names) == 1 else "" + src_lines += [ + f" _observed = ({observed}{trailing})", + " if _observed not in combo_keys:", + " _pairs = \", \".join(", + " f\"{_n}={_d}\" for _n, _d in zip(input_names, _observed)", + " )", + " raise ValueError(", + " f\"{op_name}: dtype combination ({_pairs}) is not \"", + " f\"listed in signature.dtype_combos\"", + " )", + ] exec("\n".join(src_lines), closure) _validate_dtypes = closure["_validate_dtypes"] _validate_dtypes.__name__ = "_validate_dtypes" diff --git a/tileops/ops/norm/ada_layer_norm.py b/tileops/ops/norm/ada_layer_norm.py index 29c4c219b..781fb83bc 100644 --- a/tileops/ops/norm/ada_layer_norm.py +++ b/tileops/ops/norm/ada_layer_norm.py @@ -48,17 +48,14 @@ def __init__( self, M: Optional[int] = None, N: Optional[int] = None, - dtype: Optional[torch.dtype] = None, eps: float = 1e-5, kernel_map: Optional[Dict[str, Kernel]] = None, tune: bool = False, ): self.M = M self.N = N - self.dtype = dtype self._committed_M = M self._committed_N = N - self._committed_dtype = dtype self.eps = eps self.tune = tune self.dispatch_kernel(kernel_map) @@ -112,13 +109,7 @@ def forward( raise ValueError("scale must be a CUDA tensor") if not shift.is_cuda: raise ValueError("shift must be a CUDA tensor") - expected_dtype = self._committed_dtype - if expected_dtype is not None and x.dtype != expected_dtype: - raise ValueError( - f"Expected x.dtype {expected_dtype}, got {x.dtype}" - ) - if expected_dtype is None: - expected_dtype = x.dtype + expected_dtype = x.dtype if scale.dtype != expected_dtype: raise ValueError( f"Expected scale.dtype {expected_dtype}, got {scale.dtype}" @@ -150,7 +141,6 @@ def forward( self.N = N dtype = expected_dtype assert dtype is not None - self.dtype = dtype kernel = self._get_kernel(M_actual, N, dtype, x.device.index) y = kernel(x, scale, shift) self._last_roofline_mn = (M_actual, N) diff --git a/tileops/ops/norm/ada_layer_norm_zero.py b/tileops/ops/norm/ada_layer_norm_zero.py index 3e36bc95b..5e748ae9a 100644 --- a/tileops/ops/norm/ada_layer_norm_zero.py +++ b/tileops/ops/norm/ada_layer_norm_zero.py @@ -49,17 +49,14 @@ def __init__( self, M: Optional[int] = None, N: Optional[int] = None, - dtype: Optional[torch.dtype] = None, eps: float = 1e-5, kernel_map: Optional[Dict[str, Kernel]] = None, tune: bool = False, ): self.M = M self.N = N - self.dtype = dtype self._committed_M = M self._committed_N = N - self._committed_dtype = dtype self.eps = eps self.tune = tune self.dispatch_kernel(kernel_map) @@ -120,13 +117,7 @@ def forward( raise ValueError("shift must be a CUDA tensor") if not gate.is_cuda: raise ValueError("gate must be a CUDA tensor") - expected_dtype = self._committed_dtype - if expected_dtype is not None and x.dtype != expected_dtype: - raise ValueError( - f"Expected x.dtype {expected_dtype}, got {x.dtype}" - ) - if expected_dtype is None: - expected_dtype = x.dtype + expected_dtype = x.dtype if scale.dtype != expected_dtype: raise ValueError( f"Expected scale.dtype {expected_dtype}, got {scale.dtype}" @@ -165,7 +156,6 @@ def forward( self.N = N dtype = expected_dtype assert dtype is not None - self.dtype = dtype kernel = self._get_kernel(M_actual, N, dtype, x.device.index) y = kernel(x, scale, shift, gate) self._last_roofline_mn = (M_actual, N) diff --git a/tileops/ops/norm/fused_add_layer_norm.py b/tileops/ops/norm/fused_add_layer_norm.py index bae9f433f..4f21bf13c 100644 --- a/tileops/ops/norm/fused_add_layer_norm.py +++ b/tileops/ops/norm/fused_add_layer_norm.py @@ -51,17 +51,14 @@ def __init__( self, M: Optional[int] = None, N: Optional[int] = None, - dtype: Optional[torch.dtype] = None, eps: float = 1e-5, kernel_map: Optional[Dict[str, Kernel]] = None, tune: bool = False, ): self.M = M self.N = N - self.dtype = dtype self._committed_M = M self._committed_N = N - self._committed_dtype = dtype self.eps = eps self.tune = tune self.dispatch_kernel(kernel_map) @@ -119,16 +116,14 @@ def forward( ValueError: If tensors are not on CUDA, dtypes mismatch, or shapes are incompatible with the configured dimensions. """ - expected_dtype = self._committed_dtype + expected_dtype = x.dtype for name, tensor in [("x", x), ("residual", residual), ("weight", weight), ("bias", bias)]: if not tensor.is_cuda: raise ValueError(f"{name} must be a CUDA tensor") - if expected_dtype is not None and tensor.dtype != expected_dtype: + if tensor.dtype != expected_dtype: raise ValueError( f"Expected {name}.dtype {expected_dtype}, got {tensor.dtype}" ) - if expected_dtype is None: - expected_dtype = tensor.dtype if weight.ndim != 1: raise ValueError( f"Expected weight to be 1D, got {weight.ndim}D" @@ -167,7 +162,6 @@ def forward( self.N = N dtype = expected_dtype assert dtype is not None - self.dtype = dtype kernel = self._get_kernel(M_actual, N, dtype, x.device.index) y, residual_out = kernel(x, residual, weight, bias) diff --git a/tileops/ops/norm/fused_add_rms_norm.py b/tileops/ops/norm/fused_add_rms_norm.py index d78972e38..bb109d3a3 100644 --- a/tileops/ops/norm/fused_add_rms_norm.py +++ b/tileops/ops/norm/fused_add_rms_norm.py @@ -49,17 +49,14 @@ def __init__( self, M: Optional[int] = None, N: Optional[int] = None, - dtype: Optional[torch.dtype] = None, eps: float = 1e-6, kernel_map: Optional[Dict[str, Kernel]] = None, tune: bool = False, ): self.M = M self.N = N - self.dtype = dtype self._committed_M = M self._committed_N = N - self._committed_dtype = dtype self.eps = eps self.tune = tune self.dispatch_kernel(kernel_map) @@ -115,16 +112,14 @@ def forward( ValueError: If tensors are not on CUDA, dtypes mismatch, or shapes are incompatible with the configured dimensions. """ - expected_dtype = self._committed_dtype + expected_dtype = x.dtype for name, tensor in [("x", x), ("residual", residual), ("weight", weight)]: if not tensor.is_cuda: raise ValueError(f"{name} must be a CUDA tensor") - if expected_dtype is not None and tensor.dtype != expected_dtype: + if tensor.dtype != expected_dtype: raise ValueError( f"Expected {name}.dtype {expected_dtype}, got {tensor.dtype}" ) - if expected_dtype is None: - expected_dtype = tensor.dtype if weight.ndim != 1: raise ValueError( f"Expected weight to be 1D, got {weight.ndim}D" @@ -155,7 +150,6 @@ def forward( self.N = N dtype = expected_dtype assert dtype is not None - self.dtype = dtype kernel = self._get_kernel(M_actual, N, dtype, x.device.index) y, residual_out = kernel(x, residual, weight) diff --git a/tileops/ops/norm/layer_norm.py b/tileops/ops/norm/layer_norm.py index 6fa2ff818..3028d1d23 100644 --- a/tileops/ops/norm/layer_norm.py +++ b/tileops/ops/norm/layer_norm.py @@ -53,18 +53,16 @@ def __init__( normalized_shape: Sequence[int], eps: Optional[float] = 1e-5, *, - dtype: torch.dtype, kernel_map: Optional[Dict[str, Kernel]] = None, tune: bool = False, ): self.N = normalized_shape_to_n(normalized_shape) self.normalized_shape = tuple(int(d) for d in normalized_shape) - self.dtype = dtype # Manifest declares ``eps: float | None`` with PyTorch default 1e-5. self.eps = 1e-5 if eps is None else float(eps) self.tune = tune self.dispatch_kernel(kernel_map) - self.kernel: Optional[Kernel] = None + self._kernel_cache: Dict[tuple[int, torch.dtype], Kernel] = {} self._last_m: Optional[int] = None @property @@ -72,10 +70,10 @@ def default_kernel_map(self) -> Dict[str, Kernel]: return {"layer_norm": LayerNormKernel} def eval_roofline(self) -> tuple[int, int]: - if self._last_m is None: + if self._last_m is None or self.dtype is None: raise RuntimeError( "LayerNormFwdOp.eval_roofline() requires a prior forward() " - "call to bind the leading-dims product." + "call to bind the leading-dims product and the dtype." ) elem_bytes = self.dtype.itemsize m = self._last_m @@ -109,17 +107,15 @@ def forward( raise ValueError("weight must be a CUDA tensor") if not bias.is_cuda: raise ValueError("bias must be a CUDA tensor") - if x.dtype != self.dtype: + self._validate_dtypes(x, weight, bias) + self.dtype = x.dtype + if weight.dtype != x.dtype: raise ValueError( - f"Expected x.dtype {self.dtype}, got {x.dtype}" + f"Expected weight.dtype {x.dtype}, got {weight.dtype}" ) - if weight.dtype != self.dtype: + if bias.dtype != x.dtype: raise ValueError( - f"Expected weight.dtype {self.dtype}, got {weight.dtype}" - ) - if bias.dtype != self.dtype: - raise ValueError( - f"Expected bias.dtype {self.dtype}, got {bias.dtype}" + f"Expected bias.dtype {x.dtype}, got {bias.dtype}" ) ns = self.normalized_shape @@ -143,10 +139,12 @@ def forward( weight = weight.contiguous().reshape(self.N) bias = bias.contiguous().reshape(self.N) m_actual = x.shape[0] - if self.kernel is None or m_actual != self._last_m: - self.kernel = self.kernel_map["layer_norm"]( - m_actual, self.N, self.eps, self.dtype, tune=self.tune, + key = (m_actual, x.dtype) + if key not in self._kernel_cache: + self._kernel_cache[key] = self.kernel_map["layer_norm"]( + m_actual, self.N, self.eps, x.dtype, tune=self.tune, ) + self.kernel = self._kernel_cache[key] self._last_m = m_actual y = self.kernel(x, weight, bias) diff --git a/tileops/ops/norm/rms_norm.py b/tileops/ops/norm/rms_norm.py index eb0e8d785..e3d4f6c69 100644 --- a/tileops/ops/norm/rms_norm.py +++ b/tileops/ops/norm/rms_norm.py @@ -44,7 +44,6 @@ def __init__( normalized_shape: Sequence[int], eps: Optional[float] = None, *, - dtype: torch.dtype, kernel_map: Optional[Dict[str, Kernel]] = None, tune: bool = False, ) -> None: @@ -52,29 +51,29 @@ def __init__( if len(self.normalized_shape) == 0: raise ValueError("normalized_shape must be non-empty") self.N = math.prod(self.normalized_shape) - self.dtype = dtype self.eps = _DEFAULT_EPS if eps is None else float(eps) self.tune = tune self.dispatch_kernel(kernel_map) - self._kernel_cache: Dict[int, Kernel] = {} + self._kernel_cache: Dict[tuple[int, torch.dtype], Kernel] = {} self._last_roofline_mn: Optional[Tuple[int, int]] = None @property def default_kernel_map(self) -> Dict[str, Kernel]: return {"rms_norm": RMSNormKernel} - def _get_kernel(self, m: int) -> Kernel: - if m not in self._kernel_cache: - self._kernel_cache[m] = self.kernel_map["rms_norm"]( - m, self.N, self.eps, self.dtype, tune=self.tune, + def _get_kernel(self, m: int, dtype: torch.dtype) -> Kernel: + key = (m, dtype) + if key not in self._kernel_cache: + self._kernel_cache[key] = self.kernel_map["rms_norm"]( + m, self.N, self.eps, dtype, tune=self.tune, ) - return self._kernel_cache[m] + return self._kernel_cache[key] def eval_roofline(self) -> Tuple[int, int]: - if self._last_roofline_mn is None: + if self._last_roofline_mn is None or self.dtype is None: raise RuntimeError( "RMSNormFwdOp.eval_roofline() requires a prior forward() " - "call to bind the leading-dims product." + "call to bind the leading-dims product and the dtype." ) m, n = self._last_roofline_mn elem_bytes = self.dtype.itemsize @@ -100,11 +99,11 @@ def forward(self, x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: k = len(ns) if not x.is_cuda: raise ValueError("x must be a CUDA tensor") - if x.dtype != self.dtype: - raise ValueError(f"Expected x.dtype {self.dtype}, got {x.dtype}") - if not weight.is_cuda or weight.dtype != self.dtype: + self._validate_dtypes(x, weight) + self.dtype = x.dtype + if not weight.is_cuda or weight.dtype != x.dtype: raise ValueError( - f"weight must be a CUDA tensor of dtype {self.dtype}" + f"weight must be a CUDA tensor of dtype {x.dtype}" ) if x.ndim < k or tuple(x.shape[-k:]) != ns: raise ValueError( @@ -120,6 +119,6 @@ def forward(self, x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: x_flat = x.contiguous().reshape(-1, self.N) w_flat = weight.contiguous().reshape(self.N) m = x_flat.shape[0] - y = self._get_kernel(m)(x_flat, w_flat) + y = self._get_kernel(m, x.dtype)(x_flat, w_flat) self._last_roofline_mn = (m, self.N) return y.reshape(orig_shape) diff --git a/tileops/ops/reduction/argreduce.py b/tileops/ops/reduction/argreduce.py index 8350b1782..36bc7a021 100644 --- a/tileops/ops/reduction/argreduce.py +++ b/tileops/ops/reduction/argreduce.py @@ -2,8 +2,6 @@ from typing import Dict, Optional -import torch - from tileops.kernels.kernel_base import Kernel from tileops.kernels.reduction.argreduce import ArgreduceKernel @@ -36,7 +34,6 @@ class ArgmaxFwdOp(_ReduceOpBase): def __init__( self, - dtype: torch.dtype, dim: Optional[int] = None, keepdim: bool = False, *, @@ -44,7 +41,7 @@ def __init__( tune: bool = False, ): super().__init__( - dtype=dtype, dim=dim, keepdim=keepdim, + dim=dim, keepdim=keepdim, kernel_map=kernel_map, tune=tune, ) @@ -92,7 +89,6 @@ class ArgminFwdOp(_ReduceOpBase): def __init__( self, - dtype: torch.dtype, dim: Optional[int] = None, keepdim: bool = False, *, @@ -100,7 +96,7 @@ def __init__( tune: bool = False, ): super().__init__( - dtype=dtype, dim=dim, keepdim=keepdim, + dim=dim, keepdim=keepdim, kernel_map=kernel_map, tune=tune, ) diff --git a/tileops/ops/reduction/cumulative.py b/tileops/ops/reduction/cumulative.py index c3084bdd3..a60a15570 100644 --- a/tileops/ops/reduction/cumulative.py +++ b/tileops/ops/reduction/cumulative.py @@ -20,8 +20,6 @@ class CumulativeOp(Op): op-kind dispatch string (`"sum"` or `"prod"`). Args: - dtype: Data type (float32, float16, or bfloat16). If omitted, - inferred from the first input tensor. dim: Reduction axis (default -1). Negative values are normalized at forward time (`dim % x.ndim`). kernel_map: Optional kernel override dict. @@ -37,15 +35,12 @@ class CumulativeOp(Op): def __init__( self, - dtype: Optional[torch.dtype] = None, dim: int = -1, *, kernel_map: Optional[Dict[str, Kernel]] = None, tune: bool = False, ) -> None: self.N = None - self.dtype = dtype - self._committed_dtype = dtype self.dim = dim self.tune = tune self.dispatch_kernel(kernel_map) @@ -91,10 +86,7 @@ def _get_kernel( def _validate_and_normalize_dim(self, x: torch.Tensor) -> tuple[int, int, torch.dtype]: if not x.is_cuda: raise ValueError("x must be a CUDA tensor") - if self._committed_dtype is not None and x.dtype != self._committed_dtype: - raise ValueError( - f"Expected x.dtype {self._committed_dtype}, got {x.dtype}" - ) + self._validate_dtypes(x) ndim = x.ndim if not (-ndim <= self.dim < ndim): raise ValueError( @@ -146,8 +138,6 @@ class CumsumFwdOp(CumulativeOp): SM utilization; every other shape takes the sequential scan. Args: - dtype: Optional data type (float32, float16, or bfloat16). - Preferred API infers it from ``x``. dim: Reduction axis (default -1). Negative values are normalized at forward time. kernel_map: Optional override for kernel dispatch. @@ -173,8 +163,6 @@ class CumprodFwdOp(CumulativeOp): handled inside the kernel via masked loads. Args: - dtype: Optional data type (float32, float16, or bfloat16). - Preferred API infers it from ``x``. dim: Reduction axis (default -1). Negative values are normalized at forward time. kernel_map: Optional override for kernel dispatch. diff --git a/tileops/ops/reduction/logical_reduce.py b/tileops/ops/reduction/logical_reduce.py index 62d382ba4..ba021c943 100644 --- a/tileops/ops/reduction/logical_reduce.py +++ b/tileops/ops/reduction/logical_reduce.py @@ -52,7 +52,6 @@ class AllFwdOp(_ReduceOpBase): def __init__( self, - dtype: torch.dtype, dim: Union[int, List[int], Tuple[int, ...], None] = None, keepdim: bool = False, *, @@ -62,7 +61,6 @@ def __init__( """Construct AllFwdOp. Args: - dtype: Input data type. dim: Reduction dimension (default ``None``, i.e. full reduction). Accepts ``int``, ``list[int]``, ``tuple[int, ...]``, or ``None``. @@ -71,7 +69,7 @@ def __init__( tune: Whether to autotune (default ``False``). """ super().__init__( - dtype=dtype, dim=dim, keepdim=keepdim, + dim=dim, keepdim=keepdim, kernel_map=kernel_map, tune=tune, ) @@ -126,7 +124,6 @@ class AnyFwdOp(_ReduceOpBase): def __init__( self, - dtype: torch.dtype, dim: Union[int, List[int], Tuple[int, ...], None] = None, keepdim: bool = False, *, @@ -136,7 +133,6 @@ def __init__( """Construct AnyFwdOp. Args: - dtype: Input data type. dim: Reduction dimension (default ``None``, i.e. full reduction). Accepts ``int``, ``list[int]``, ``tuple[int, ...]``, or ``None``. @@ -145,7 +141,7 @@ def __init__( tune: Whether to autotune (default ``False``). """ super().__init__( - dtype=dtype, dim=dim, keepdim=keepdim, + dim=dim, keepdim=keepdim, kernel_map=kernel_map, tune=tune, ) @@ -199,7 +195,6 @@ class CountNonzeroFwdOp(_ReduceOpBase): def __init__( self, - dtype: torch.dtype, dim: Union[int, List[int], Tuple[int, ...], None] = None, *, kernel_map: Optional[Dict[str, Kernel]] = None, @@ -207,7 +202,7 @@ def __init__( ): # count_nonzero never keeps dim (matches torch.count_nonzero) super().__init__( - dtype=dtype, dim=dim, keepdim=False, + dim=dim, keepdim=False, kernel_map=kernel_map, tune=tune, ) diff --git a/tileops/ops/reduction/reduce.py b/tileops/ops/reduction/reduce.py index 75ae10680..3ed634cd5 100644 --- a/tileops/ops/reduction/reduce.py +++ b/tileops/ops/reduction/reduce.py @@ -86,7 +86,6 @@ class _ReduceOpBase(Op): def __init__( self, - dtype: torch.dtype, dim: Union[int, List[int], Tuple[int, ...], None] = None, keepdim: bool = False, *, @@ -96,7 +95,6 @@ def __init__( """Construct a reduce op. Args: - dtype: Input data type. dim: Reduction dimension (default ``None``, i.e. full reduction). Accepts ``int``, ``list[int]``, ``tuple[int, ...]``, or ``None``. @@ -104,7 +102,6 @@ def __init__( kernel_map: Optional override for kernel dispatch. tune: Whether to autotune (default ``False``). """ - self.dtype = dtype self.dim = dim self.keepdim = keepdim self._tune = tune @@ -225,8 +222,8 @@ def _validate_input_tensor(self, x: torch.Tensor) -> None: """ if not x.is_cuda: raise ValueError("x must be a CUDA tensor") - if x.dtype != self.dtype: - raise ValueError(f"Expected x.dtype {self.dtype}, got {x.dtype}") + self._validate_dtypes(x) + self.dtype = x.dtype if x.ndim == 0: raise ValueError("Input tensor must be at least 1D") @@ -309,8 +306,8 @@ def _maybe_scalar(self, x: torch.Tensor): return None if not x.is_cuda: raise ValueError("x must be a CUDA tensor") - if x.dtype != self.dtype: - raise ValueError(f"Expected x.dtype {self.dtype}, got {x.dtype}") + self._validate_dtypes(x) + self.dtype = x.dtype self._validate_scalar_dim() self._last_roofline_mn = (1, 1) return self._scalar_forward(x) @@ -354,6 +351,11 @@ def eval_roofline(self) -> tuple[int, int]: "call to bind dynamic input shape" ) M, N = self._last_roofline_mn + if self.dtype is None: + raise RuntimeError( + f"{type(self).__name__}.eval_roofline() requires a prior forward() " + "call to bind dtype" + ) elem_bytes = self.dtype.itemsize op_kind = self._op_kind @@ -395,13 +397,15 @@ def eval_roofline(self) -> tuple[int, int]: # Kernel cache - def _get_or_create_kernel(self, M: int, N: int) -> object: - """Return a cached kernel for (M, N), creating one if needed.""" - key = (M, N) + def _get_or_create_kernel( + self, M: int, N: int, dtype: torch.dtype, + ) -> object: + """Return a cached kernel for (M, N, dtype), creating one if needed.""" + key = (M, N, dtype) if key not in self._kernel_cache: kernel_cls = self.kernel_map[self._kernel_key] self._kernel_cache[key] = kernel_cls( - M, N, self._op_kind, self.dtype, + M, N, self._op_kind, dtype, tune=self._tune, **self._build_kernel_kwargs(), ) return self._kernel_cache[key] @@ -436,7 +440,7 @@ def _prepare_input( M = prod(x.shape[:-1]) self._last_roofline_mn = (M, N) x = x.reshape(M, N) - kernel = self._get_or_create_kernel(M, N) + kernel = self._get_or_create_kernel(M, N, x.dtype) if not self._kernel_handles_padding: N_padded = align_up(N, DEFAULT_ALIGNMENT) if N_padded != N: @@ -462,7 +466,7 @@ def _prepare_input( x = x.contiguous().reshape(M, N) - kernel = self._get_or_create_kernel(M, N) + kernel = self._get_or_create_kernel(M, N, x.dtype) if not self._kernel_handles_padding: N_padded = align_up(N, DEFAULT_ALIGNMENT) @@ -563,7 +567,6 @@ class ProdFwdOp(_SimpleReduceOp): def __init__( self, - dtype: torch.dtype, dim: int = -1, keepdim: bool = False, *, @@ -573,14 +576,13 @@ def __init__( """Construct ProdFwdOp. Args: - dtype: Input data type. dim (int): reduction dimension (default ``-1``). keepdim: Whether to retain reduced dims as size 1. kernel_map: Optional override for kernel dispatch. tune: Whether to autotune (default ``False``). """ super().__init__( - dtype=dtype, dim=dim, keepdim=keepdim, + dim=dim, keepdim=keepdim, kernel_map=kernel_map, tune=tune, ) @@ -623,7 +625,6 @@ class _WelfordReduceOp(_ReduceOpBase): def __init__( self, - dtype: torch.dtype, dim: Union[int, List[int], Tuple[int, ...], None] = None, correction: int = 1, keepdim: bool = False, @@ -634,7 +635,6 @@ def __init__( """Construct a Welford-based reduce op. Args: - dtype: Input data type. dim: Reduction dimension (default ``None``, i.e. full reduction). Accepts ``int``, ``list[int]``, ``tuple[int, ...]``, or ``None``. @@ -645,7 +645,7 @@ def __init__( """ self.correction = correction super().__init__( - dtype=dtype, dim=dim, keepdim=keepdim, + dim=dim, keepdim=keepdim, kernel_map=kernel_map, tune=tune, ) diff --git a/tileops/ops/reduction/softmax.py b/tileops/ops/reduction/softmax.py index 87e7e7ed0..3b975d234 100644 --- a/tileops/ops/reduction/softmax.py +++ b/tileops/ops/reduction/softmax.py @@ -37,7 +37,6 @@ class _SoftmaxBaseOp(Op): override output reshaping if needed. Args: - dtype: Optional committed dtype for subclasses that still expose it. dim: Reduction dimension (default -1). N: Optional committed reduction dim for subclasses that still expose it. kernel_map: Optional override for kernel dispatch. @@ -57,15 +56,12 @@ class _SoftmaxBaseOp(Op): def __init__( self, - dtype: Optional[torch.dtype] = None, dim: Union[int, List[int]] = -1, N: Optional[int] = None, *, kernel_map: Optional[Dict[str, Kernel]] = None, tune: bool = False, ): - self.dtype = dtype - self._committed_dtype = dtype self.dim = dim self.N = N self.keepdim = False @@ -90,10 +86,6 @@ def _validate(self, x: torch.Tensor) -> None: "x.dtype must be float16, bfloat16, or float32, " f"got {x.dtype}" ) - if self._committed_dtype is not None and x.dtype != self._committed_dtype: - raise ValueError( - f"Expected x.dtype {self._committed_dtype}, got {x.dtype}" - ) if x.ndim == 0: raise ValueError("Input tensor must be at least 1D") self.dtype = x.dtype @@ -326,7 +318,6 @@ class LogSumExpFwdOp(_SoftmaxBaseOp): (or with size-1 if keepdim=True). Args: - dtype: Data type (float32, float16, or bfloat16). dim: Reduction dimension (default -1). keepdim: Retain reduced dimension (default False). kernel_map: Optional override for kernel dispatch. @@ -340,12 +331,11 @@ class LogSumExpFwdOp(_SoftmaxBaseOp): def __init__( self, - dtype: torch.dtype, dim: Union[int, List[int]] = -1, keepdim: bool = False, *, kernel_map: Optional[Dict[str, Kernel]] = None, tune: bool = False, ): - super().__init__(dtype=dtype, dim=dim, kernel_map=kernel_map, tune=tune) + super().__init__(dim=dim, kernel_map=kernel_map, tune=tune) self.keepdim = keepdim diff --git a/tileops/ops/reduction/vector_norm.py b/tileops/ops/reduction/vector_norm.py index d0f187416..9f94d1ea1 100644 --- a/tileops/ops/reduction/vector_norm.py +++ b/tileops/ops/reduction/vector_norm.py @@ -43,7 +43,6 @@ class L1NormFwdOp(_ReduceOpBase): def __init__( self, - dtype: torch.dtype, ord: Union[int, float] = 1, dim: Union[int, List[int], None] = None, keepdim: bool = False, @@ -58,7 +57,7 @@ def __init__( ) self.ord = ord super().__init__( - dtype=dtype, dim=dim, keepdim=keepdim, + dim=dim, keepdim=keepdim, kernel_map=kernel_map, tune=tune, ) @@ -93,7 +92,6 @@ class L2NormFwdOp(_ReduceOpBase): def __init__( self, - dtype: torch.dtype, ord: Union[int, float] = 2, dim: Union[int, List[int], None] = None, keepdim: bool = False, @@ -108,7 +106,7 @@ def __init__( ) self.ord = ord super().__init__( - dtype=dtype, dim=dim, keepdim=keepdim, + dim=dim, keepdim=keepdim, kernel_map=kernel_map, tune=tune, ) @@ -146,7 +144,6 @@ class InfNormFwdOp(_ReduceOpBase): def __init__( self, - dtype: torch.dtype, ord: Union[int, float] = inf, dim: Union[int, List[int], None] = None, keepdim: bool = False, @@ -161,7 +158,7 @@ def __init__( ) self.ord = ord super().__init__( - dtype=dtype, dim=dim, keepdim=keepdim, + dim=dim, keepdim=keepdim, kernel_map=kernel_map, tune=tune, ) From b54e1459ab2d26b40166badfedc3bebdcb23d602 Mon Sep 17 00:00:00 2001 From: lcy-seso Date: Tue, 4 Aug 2026 14:16:32 +0800 Subject: [PATCH 03/15] [Refactor][Attention] Drop is_hopper and the non-Hopper branches it gated MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The library targets H200, so `is_hopper()` was always true and every branch behind its false arm was unreachable. Removing the helper removes one axis from kernel selection entirely: what remains depends on dtype, shape and the semantic flags. Two consequences that had to be handled rather than left implicit: - The manifest named the fallbacks the ops no longer dispatch to. `source.kernel_map` now names the Wgmma/Ws kernels actually used, and `gqa_bwd_postprocess_kernel` is gone — that slot existed only on the non-Hopper path, where it held `None` on Hopper. - `test_gqa_fwd_dispatch_falls_back_off_h200` asserted the non-Hopper fallback. Its premise is gone, but the other half of what it covered is not, so it now checks that a shape outside the H200 warp-specialized contract still lands on the WGMMA kernel. The six fallback kernel classes are left in place and recorded for a separate decision: each is still exported and still passes the arch check on H200, so a caller can reach one through a `kernel_map` override. `supported_archs == [80, 89, 90]` does not identify them — several kernels carry that declaration while being the only implementation of their slot. --- benchmarks/ops/attention/bench_gqa.py | 2 - tests/ops/attention/test_gqa.py | 22 +++----- tests/ops/attention/test_gqa_decode.py | 7 --- tests/ops/attention/test_gqa_decode_paged.py | 6 --- tileops/manifest/attention.yaml | 14 +++-- tileops/ops/attention/deepseek_dsa.py | 6 --- tileops/ops/attention/deepseek_mla.py | 5 +- tileops/ops/attention/gqa.py | 57 ++++++-------------- tileops/ops/attention/mha.py | 8 +-- tileops/utils/__init__.py | 2 - tileops/utils/utils.py | 4 -- 11 files changed, 32 insertions(+), 101 deletions(-) diff --git a/benchmarks/ops/attention/bench_gqa.py b/benchmarks/ops/attention/bench_gqa.py index 150b24a8f..47e1e48f1 100644 --- a/benchmarks/ops/attention/bench_gqa.py +++ b/benchmarks/ops/attention/bench_gqa.py @@ -239,8 +239,6 @@ def _tileops_gqa_variant(op: GroupedQueryAttentionFwdOp) -> str: return "ws_noncausal" if isinstance(kernel, GQAFwdWgmmaPipelinedKernel): return "wgmma_pipelined" - if isinstance(kernel, GQAFwdKernel): - return "legacy" return kernel.__class__.__name__ diff --git a/tests/ops/attention/test_gqa.py b/tests/ops/attention/test_gqa.py index 537edc631..72927b010 100644 --- a/tests/ops/attention/test_gqa.py +++ b/tests/ops/attention/test_gqa.py @@ -8,7 +8,6 @@ from tests.test_base import FixtureBase, TestBase from tileops.kernels.attention import ( - GQAFwdKernel, GQAFwdWgmmaPipelinedKernel, GQAFwdWsPersistentCausalKernel, GQAFwdWsPersistentKernel, @@ -858,33 +857,30 @@ def test_gqa_bwd(batch: int, seq_len: int, heads: int, heads_kv: int, dim: int, @pytest.mark.smoke def test_gqa_fwd_dispatch_selects_ws_noncausal_on_h200() -> None: kernel_cls = _select_gqa_fwd_kernel_cls( - 4, 64, 4, 512, 128, False, torch.float16, hopper=True, h200=True) + 4, 64, 4, 512, 128, False, torch.float16, h200=True) assert kernel_cls is GQAFwdWsPersistentKernel @pytest.mark.smoke def test_gqa_fwd_dispatch_selects_ws_causal_on_h200() -> None: kernel_cls = _select_gqa_fwd_kernel_cls( - 4, 64, 4, 512, 128, True, torch.float16, hopper=True, h200=True) + 4, 64, 4, 512, 128, True, torch.float16, h200=True) assert kernel_cls is GQAFwdWsPersistentCausalKernel @pytest.mark.smoke def test_gqa_fwd_dispatch_falls_back_for_small_causal_shape() -> None: kernel_cls = _select_gqa_fwd_kernel_cls( - 1, 32, 8, 1024, 128, True, torch.float16, hopper=True, h200=True) + 1, 32, 8, 1024, 128, True, torch.float16, h200=True) assert kernel_cls is GQAFwdWgmmaPipelinedKernel @pytest.mark.smoke -def test_gqa_fwd_dispatch_falls_back_off_h200() -> None: - hopper_cls = _select_gqa_fwd_kernel_cls( - 4, 64, 4, 512, 128, False, torch.float16, hopper=True, h200=False) - assert hopper_cls is GQAFwdWgmmaPipelinedKernel - - non_hopper_cls = _select_gqa_fwd_kernel_cls( - 4, 64, 4, 512, 128, False, torch.float16, hopper=False, h200=False) - assert non_hopper_cls is GQAFwdKernel +def test_gqa_fwd_dispatch_without_h200_work_threshold() -> None: + """Off the H200 warp-specialized contract, dispatch takes the WGMMA kernel.""" + cls = _select_gqa_fwd_kernel_cls( + 4, 64, 4, 512, 128, False, torch.float16, h200=False) + assert cls is GQAFwdWgmmaPipelinedKernel @pytest.mark.smoke @@ -929,7 +925,6 @@ def test_gqa_prefill_dense_selector_widens_ws_capability() -> None: torch.float16, sm_scale=0.25, softcap=2.0, - hopper=True, ).__name__ == "GQAPrefillFwdWsPersistentCausalKernel" assert _select_gqa_prefill_fwd_kernel_cls( 128, @@ -937,7 +932,6 @@ def test_gqa_prefill_dense_selector_widens_ws_capability() -> None: torch.bfloat16, sm_scale=128**-0.5, softcap=0.0, - hopper=True, ).__name__ == "GQAPrefillFwdWsPersistentCausalKernel" diff --git a/tests/ops/attention/test_gqa_decode.py b/tests/ops/attention/test_gqa_decode.py index 4cc836e7e..56e26daa3 100644 --- a/tests/ops/attention/test_gqa_decode.py +++ b/tests/ops/attention/test_gqa_decode.py @@ -5,7 +5,6 @@ from tests.test_base import FixtureBase, TestBase from tileops.kernels.attention.gqa_decode import GQADecodeKernel from tileops.ops import GroupedQueryAttentionDecodeWithKVCacheFwdOp -from tileops.utils import is_hopper from workloads.attention.gqa import ( GroupedQueryAttentionDecodeWorkload, ) @@ -121,8 +120,6 @@ def test_gqa_decode_rejects_non_positive_seqlen_kv() -> None: @pytest.mark.smoke def test_gqa_decode_bs1_dispatch() -> None: """batch=1 fp16 dim-128 requests select the WS kernel; other dtypes/shapes fall back.""" - if not is_hopper(): - pytest.skip("batch=1 warp-specialized decode requires Hopper") op = GroupedQueryAttentionDecodeWithKVCacheFwdOp(1, 32, 4, 8192, 128, torch.float16) assert op._uses_bs1_fast_path() assert op.kernel.__class__.__name__ == "GQADecodeBs1Kernel" @@ -149,8 +146,6 @@ def test_gqa_decode_bs1_runtime_context_switch() -> None: Covers the crossover (1024), a balanced mid split, an aligned split needing >=3 tiles per slice, an unaligned length, and the sub-1024 non-split fallback. """ - if not is_hopper(): - pytest.skip("batch=1 warp-specialized decode requires Hopper") op = GroupedQueryAttentionDecodeWithKVCacheFwdOp(1, 32, 4, 8192, 128, torch.float16) assert op.kernel.__class__.__name__ == "GQADecodeBs1Kernel" for real, tier in ((6000, "ctx"), (3072, "ctx"), (2048, "ctx"), (1024, "ctx"), @@ -163,8 +158,6 @@ def test_gqa_decode_bs1_runtime_context_switch() -> None: @pytest.mark.smoke def test_gqa_decode_bs1_group4() -> None: """The WS kernel generalizes to a query-per-KV-head group other than 8 (here 4).""" - if not is_hopper(): - pytest.skip("batch=1 warp-specialized decode requires Hopper") op = GroupedQueryAttentionDecodeWithKVCacheFwdOp(1, 32, 8, 4096, 128, torch.float16) assert op._uses_bs1_fast_path() assert op.kernel.__class__.__name__ == "GQADecodeBs1Kernel" diff --git a/tests/ops/attention/test_gqa_decode_paged.py b/tests/ops/attention/test_gqa_decode_paged.py index b539c6131..8230f5ead 100644 --- a/tests/ops/attention/test_gqa_decode_paged.py +++ b/tests/ops/attention/test_gqa_decode_paged.py @@ -9,7 +9,6 @@ from tests.test_base import FixtureBase, TestBase from tileops.ops import GroupedQueryAttentionDecodePagedWithKVCacheFwdOp -from tileops.utils import is_hopper from workloads.attention.gqa import ( GroupedQueryAttentionDecodePagedWorkload, ) @@ -182,9 +181,6 @@ def test_gqa_decode_paged_bs1_fixed_tier_correctness( reverse_pages: bool, ) -> None: """Check both runtime tiers, including output-distinguishing page translation.""" - if not is_hopper(): - pytest.skip("batch=1 warp-specialized paged decode requires Hopper") - torch.manual_seed(0) batch, heads, heads_kv, seqlen_kv, dim, page_size = 1, 32, 4, 4096, 128, 256 dtype = torch.float16 @@ -217,8 +213,6 @@ def test_gqa_decode_paged_bs1_fixed_tier_correctness( @pytest.mark.smoke def test_gqa_decode_paged_bs1_dispatch() -> None: """Eligible Hopper requests select the paged TMA/WGMMA kernel.""" - if not is_hopper(): - pytest.skip("batch=1 warp-specialized paged decode requires Hopper") op = GroupedQueryAttentionDecodePagedWithKVCacheFwdOp( 1, 32, 4, 8192, 128, 256, torch.float16) assert op._uses_bs1_fast_path() diff --git a/tileops/manifest/attention.yaml b/tileops/manifest/attention.yaml index 294a898d5..d7c749b75 100644 --- a/tileops/manifest/attention.yaml +++ b/tileops/manifest/attention.yaml @@ -91,8 +91,7 @@ MultiHeadAttentionBwdOp: kernel: tileops/kernels/attention/gqa_bwd.py kernel_map: gqa_bwd_preprocess_kernel: FlashAttnBwdPreprocessKernel - gqa_bwd_kernel: GQABwdKernel - gqa_bwd_postprocess_kernel: FlashAttnBwdPostprocessKernel + gqa_bwd_kernel: GQABwdWgmmaPipelinedKernel op: tileops/ops/attention/mha.py test: tests/ops/attention/test_mha.py bench: benchmarks/ops/attention/bench_mha.py @@ -207,7 +206,7 @@ GroupedQueryAttentionPrefillFwdOp: gqa_prefill_fwd_kernel: GQAPrefillFwdKernel | GQAPrefillFwdWsPersistentCausalKernel gqa_prefill_square_fwd_kernel: GQAFwdWsPersistentCausalKernel gqa_prefill_varlen_fwd_kernel: GQAPrefillVarlenFwdKernel - gqa_sliding_window_varlen_fwd: GQASlidingWindowVarlenFwdKernel + gqa_sliding_window_varlen_fwd: GQASlidingWindowVarlenFwdWgmmaPipelinedKernel gqa_prefill_fp8_tensor_core_fwd_kernel: GQAFwdFP8Fa3ContractPtxAccBN224WsTmaVKernel op: tileops/ops/attention/gqa.py test: tests/ops/attention/test_gqa.py @@ -343,8 +342,7 @@ GroupedQueryAttentionBwdOp: kernel: tileops/kernels/attention/gqa_bwd.py kernel_map: gqa_bwd_preprocess_kernel: FlashAttnBwdPreprocessKernel - gqa_bwd_kernel: GQABwdKernel - gqa_bwd_postprocess_kernel: FlashAttnBwdPostprocessKernel + gqa_bwd_kernel: GQABwdWgmmaPipelinedKernel op: tileops/ops/attention/gqa.py test: tests/ops/attention/test_gqa.py bench: benchmarks/ops/attention/bench_gqa.py @@ -584,7 +582,7 @@ GroupedQueryAttentionSlidingWindowFwdOp: source: kernel: tileops/kernels/attention/gqa_sliding_window_fwd.py kernel_map: - gqa_sliding_window_fwd: GQASlidingWindowFwdKernel + gqa_sliding_window_fwd: GQASlidingWindowFwdWgmmaPipelinedKernel op: tileops/ops/attention/gqa.py test: tests/ops/attention/test_gqa_sliding_window.py bench: benchmarks/ops/attention/bench_gqa_sliding_window.py @@ -632,7 +630,7 @@ GroupedQueryAttentionSlidingWindowVarlenFwdOp: source: kernel: tileops/kernels/attention/gqa_sliding_window_varlen_fwd.py kernel_map: - gqa_sliding_window_varlen_fwd: GQASlidingWindowVarlenFwdKernel + gqa_sliding_window_varlen_fwd: GQASlidingWindowVarlenFwdWgmmaPipelinedKernel op: tileops/ops/attention/gqa.py test: tests/ops/attention/test_gqa_sliding_window_varlen.py bench: benchmarks/ops/attention/bench_gqa_sliding_window_varlen.py @@ -680,7 +678,7 @@ MultiHeadLatentAttentionDecodeWithKVCacheFwdOp: source: kernel: tileops/kernels/attention/deepseek_mla_decode.py kernel_map: - mla_decode_kernel: MLADecodeKernel + mla_decode_kernel: MLADecodeWsKernel op: tileops/ops/attention/deepseek_mla.py test: tests/ops/attention/test_deepseek_mla_decode.py bench: benchmarks/ops/attention/bench_deepseek_mla_decode.py diff --git a/tileops/ops/attention/deepseek_dsa.py b/tileops/ops/attention/deepseek_dsa.py index c423418a4..5c7b00e21 100644 --- a/tileops/ops/attention/deepseek_dsa.py +++ b/tileops/ops/attention/deepseek_dsa.py @@ -4,7 +4,6 @@ from tileops.kernels.attention import SparseMlaKernel from tileops.kernels.kernel_base import Kernel -from tileops.utils import is_hopper from ..op_base import Op @@ -107,11 +106,6 @@ def default_kernel_map(self) -> Dict[str, Kernel]: Dict[str, Kernel]: A dictionary mapping kernel names to kernel functions. The default map includes the "sparse_mla_kernel". """ - if not is_hopper(): - raise RuntimeError( - "DeepSeekSparseAttentionDecodeWithKVCacheFwdOp requires a Hopper GPU " - "(SM90) because the underlying SparseMlaKernel uses WGMMA instructions." - ) return {"sparse_mla_kernel": SparseMlaKernel} def forward(self, q: torch.Tensor, kv: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: diff --git a/tileops/ops/attention/deepseek_mla.py b/tileops/ops/attention/deepseek_mla.py index 53e631f7f..734acecfc 100644 --- a/tileops/ops/attention/deepseek_mla.py +++ b/tileops/ops/attention/deepseek_mla.py @@ -2,9 +2,8 @@ import torch -from tileops.kernels.attention import MLADecodeKernel, MLADecodeWsKernel +from tileops.kernels.attention import MLADecodeWsKernel from tileops.kernels.kernel_base import Kernel -from tileops.utils import is_hopper from ..op_base import Op @@ -39,7 +38,7 @@ def __init__(self, @property def default_kernel_map(self) -> Dict[str, Kernel]: - return {"mla_decode_kernel": MLADecodeWsKernel if is_hopper() else MLADecodeKernel} + return {"mla_decode_kernel": MLADecodeWsKernel} def forward(self, q: torch.Tensor, q_pe: torch.Tensor, k: torch.Tensor, k_pe: torch.Tensor) -> torch.Tensor: diff --git a/tileops/ops/attention/gqa.py b/tileops/ops/attention/gqa.py index c08999817..d6870c64d 100644 --- a/tileops/ops/attention/gqa.py +++ b/tileops/ops/attention/gqa.py @@ -5,16 +5,13 @@ import torch.nn.functional as F from tileops.kernels.attention import ( - FlashAttnBwdPostprocessKernel, FlashAttnBwdPreprocessKernel, - GQABwdKernel, GQABwdWgmmaPipelinedKernel, GQADecodeBs1Kernel, GQADecodeKernel, GQADecodePagedBs1Kernel, GQADecodePagedKernel, GQAFwdFP8Fa3ContractPtxAccBN224WsTmaVKernel, - GQAFwdKernel, GQAFwdWgmmaPipelinedKernel, GQAFwdWsPersistentCausalKernel, GQAFwdWsPersistentKernel, @@ -25,13 +22,11 @@ GQAPrefillPagedWithKVCacheRopeAppendKernel, GQAPrefillPagedWithKVCacheRopeFwdKernel, GQAPrefillVarlenFwdKernel, - GQASlidingWindowFwdKernel, GQASlidingWindowFwdWgmmaPipelinedKernel, - GQASlidingWindowVarlenFwdKernel, GQASlidingWindowVarlenFwdWgmmaPipelinedKernel, ) from tileops.kernels.kernel_base import Kernel -from tileops.utils import is_h200, is_hopper +from tileops.utils import is_h200 from ..op_base import Op from ..rope import _base_freqs @@ -80,7 +75,6 @@ def _supports_gqa_decode_bs1( """Return whether the common Hopper batch=1 GQA decode contract is satisfied.""" if not ( batch == 1 - and is_hopper() and dtype == torch.float16 and dim == 128 and softcap == 0.0 @@ -139,11 +133,8 @@ def _select_gqa_fwd_kernel_cls( is_causal: bool, dtype: torch.dtype, *, - hopper: bool, h200: bool, ) -> Type[Kernel]: - if not hopper: - return GQAFwdKernel if is_causal: if _supports_gqa_ws_causal(batch, heads, heads_kv, seq_len, dim, dtype, h200=h200): return GQAFwdWsPersistentCausalKernel @@ -159,11 +150,9 @@ def _select_gqa_prefill_fwd_kernel_cls( dtype: torch.dtype, sm_scale: float, softcap: float, - *, - hopper: bool, ) -> Type[Kernel]: del sm_scale, softcap - if hopper and is_causal and dim == 128 and dtype in (torch.float16, torch.bfloat16): + if is_causal and dim == 128 and dtype in (torch.float16, torch.bfloat16): return GQAPrefillFwdWsPersistentCausalKernel return GQAPrefillFwdKernel @@ -322,7 +311,6 @@ def default_kernel_map(self) -> Dict[str, Kernel]: self.dim, self.is_causal, self.dtype, - hopper=is_hopper(), sm_scale=self.sm_scale, softcap=self.softcap, ) @@ -441,13 +429,8 @@ def default_kernel_map(self) -> Dict[str, Kernel]: self.dtype, self.sm_scale, self.softcap, - hopper=is_hopper(), - ) - sliding_kernel_cls = ( - GQASlidingWindowVarlenFwdWgmmaPipelinedKernel - if is_hopper() - else GQASlidingWindowVarlenFwdKernel ) + sliding_kernel_cls = GQASlidingWindowVarlenFwdWgmmaPipelinedKernel return { "gqa_prefill_fwd_kernel": dense_kernel_cls, "gqa_prefill_square_fwd_kernel": GQAFwdWsPersistentCausalKernel, @@ -1354,10 +1337,6 @@ def __init__(self, self.dtype, tune=tune) self.kernel = self.kernel_map["gqa_bwd_kernel"]( batch, heads, heads_kv, seq_len, dim, is_causal, self.dtype, tune=tune) - if not is_hopper(): - self.post_kernel = self.kernel_map["gqa_bwd_postprocess_kernel"](batch, heads, seq_len, - dim, self.dtype, - tune=tune) @property def default_kernel_map(self) -> Dict[str, Kernel]: @@ -1365,9 +1344,7 @@ def default_kernel_map(self) -> Dict[str, Kernel]: "gqa_bwd_preprocess_kernel": FlashAttnBwdPreprocessKernel, "gqa_bwd_kernel": - GQABwdWgmmaPipelinedKernel if is_hopper() else GQABwdKernel, - "gqa_bwd_postprocess_kernel": - FlashAttnBwdPostprocessKernel if not is_hopper() else None, + GQABwdWgmmaPipelinedKernel, } def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, o: torch.Tensor, @@ -1379,7 +1356,7 @@ def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, o: torch.Te dk = torch.zeros_like(k, dtype=torch.float32) dv = torch.zeros_like(v, dtype=torch.float32) self.kernel(q, k, v, do, lse, delta, dq, dk, dv) - dq = dq.to(self.dtype) if is_hopper() else self.post_kernel(dq) + dq = dq.to(self.dtype) dk, dv = dk.to(self.dtype), dv.to(self.dtype) return dq, dk, dv @@ -1425,13 +1402,10 @@ def __init__(self, @property def default_kernel_map(self) -> Dict[str, Kernel]: - # The batch=1 warp-specialized decode kernel is Hopper-only; only expose it on - # Hopper so the op stays constructible (falling back to the architecture-agnostic - # GQADecodeKernel) on sm80/sm89. - kernel_map: Dict[str, Kernel] = {"gqa_decode_kernel": GQADecodeKernel} - if is_hopper(): - kernel_map["gqa_decode_bs1_kernel"] = GQADecodeBs1Kernel - return kernel_map + return { + "gqa_decode_kernel": GQADecodeKernel, + "gqa_decode_bs1_kernel": GQADecodeBs1Kernel, + } def _uses_bs1_fast_path(self) -> bool: """Ctor-time gate for the batch=1 warp-specialized decode kernel. @@ -1512,10 +1486,10 @@ def __init__(self, @property def default_kernel_map(self) -> Dict[str, Kernel]: - kernel_map: Dict[str, Kernel] = {"gqa_decode_paged_kernel": GQADecodePagedKernel} - if is_hopper(): - kernel_map["gqa_decode_paged_bs1_kernel"] = GQADecodePagedBs1Kernel - return kernel_map + return { + "gqa_decode_paged_kernel": GQADecodePagedKernel, + "gqa_decode_paged_bs1_kernel": GQADecodePagedBs1Kernel, + } def _uses_bs1_fast_path(self) -> bool: """Use the paged Hopper fast path only when its TMA tile stays within one page.""" @@ -1610,7 +1584,7 @@ def __init__( @property def default_kernel_map(self) -> Dict[str, Kernel]: - kernel = GQASlidingWindowFwdWgmmaPipelinedKernel if is_hopper() else GQASlidingWindowFwdKernel + kernel = GQASlidingWindowFwdWgmmaPipelinedKernel return {"gqa_sliding_window_fwd": kernel} def forward( @@ -1754,8 +1728,7 @@ def __init__( @property def default_kernel_map(self) -> Dict[str, Kernel]: - kernel = (GQASlidingWindowVarlenFwdWgmmaPipelinedKernel - if is_hopper() else GQASlidingWindowVarlenFwdKernel) + kernel = GQASlidingWindowVarlenFwdWgmmaPipelinedKernel return {"gqa_sliding_window_varlen_fwd": kernel} def forward( diff --git a/tileops/ops/attention/mha.py b/tileops/ops/attention/mha.py index c261f9d20..d5d96989f 100644 --- a/tileops/ops/attention/mha.py +++ b/tileops/ops/attention/mha.py @@ -4,16 +4,13 @@ import torch.nn.functional as F from tileops.kernels.attention import ( - FlashAttnBwdPostprocessKernel, FlashAttnBwdPreprocessKernel, - GQABwdKernel, GQABwdWgmmaPipelinedKernel, GQAFwdWsPersistentCausalKernel, MHADecodeKernel, MHADecodePagedKernel, ) from tileops.kernels.kernel_base import Kernel -from tileops.utils import is_hopper from ..op_base import Op from .gqa import ( @@ -84,7 +81,6 @@ def default_kernel_map(self) -> Dict[str, Kernel]: self.dtype, sm_scale=None, softcap=0.0, - hopper=is_hopper(), ), "gqa_prefill_square_fwd_kernel": GQAFwdWsPersistentCausalKernel, } @@ -156,9 +152,7 @@ def default_kernel_map(self) -> Dict[str, Kernel]: "gqa_bwd_preprocess_kernel": FlashAttnBwdPreprocessKernel, "gqa_bwd_kernel": - GQABwdWgmmaPipelinedKernel if is_hopper() else GQABwdKernel, - "gqa_bwd_postprocess_kernel": - FlashAttnBwdPostprocessKernel if not is_hopper() else None, + GQABwdWgmmaPipelinedKernel, } @staticmethod diff --git a/tileops/utils/__init__.py b/tileops/utils/__init__.py index b0b7fc6b5..b8af38eca 100644 --- a/tileops/utils/__init__.py +++ b/tileops/utils/__init__.py @@ -1,13 +1,11 @@ from .utils import ( get_sm_version, is_h200, - is_hopper, str2dtype, ) __all__ = [ "get_sm_version", "is_h200", - "is_hopper", "str2dtype", ] diff --git a/tileops/utils/utils.py b/tileops/utils/utils.py index ca72b49d7..37156f605 100644 --- a/tileops/utils/utils.py +++ b/tileops/utils/utils.py @@ -10,10 +10,6 @@ } -def is_hopper(): - return torch.cuda.get_device_capability() == (9, 0) - - @functools.lru_cache(maxsize=1) def is_h200(): if not torch.cuda.is_available(): From d90e3008485dea3c3ebddd759c85abe50dbecd66 Mon Sep 17 00:00:00 2001 From: lcy-seso Date: Tue, 4 Aug 2026 16:20:01 +0800 Subject: [PATCH 04/15] [Refactor][Manifest] Declare the output dtype GQA prefill was selecting undeclared MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `GroupedQueryAttentionPrefillFwdOp` takes a `dtype` constructor argument that selects the fp8 path's output element type, but no manifest param declared it — the entry's `dtype_combos` already recorded that fp8 inputs admit either fp16 or bf16 output, with nothing saying who chooses. Also drops the writer identity from the trust-model diagram and narrows `implement-op`'s blanket "do not modify manifest or design docs" to what it was protecting: the spec is not rewritten to match code that does not conform to it. Known advisory: the validator compares parameter defaults with `!=`, so a `torch.dtype` default can never match its manifest spelling. No entry could declare one before this change either. --- .claude/skills/implement-op/SKILL.md | 2 +- docs/design/manifest.md | 2 +- tileops/manifest/attention.yaml | 4 ++++ 3 files changed, 6 insertions(+), 2 deletions(-) diff --git a/.claude/skills/implement-op/SKILL.md b/.claude/skills/implement-op/SKILL.md index d4d9bdd39..27ada29cd 100644 --- a/.claude/skills/implement-op/SKILL.md +++ b/.claude/skills/implement-op/SKILL.md @@ -92,7 +92,7 @@ Record `observations` — design knowledge discovered during migration: - Edge cases - Abstraction opportunities -Do NOT modify manifest or design docs. Observations are returned to orchestrator and surfaced in PR for human review. +Do not change the manifest or a design doc to match code that does not conform — the spec is the reference, not a record of the implementation. Observations are returned to the orchestrator and surfaced in the PR. ### 6. COMMIT diff --git a/docs/design/manifest.md b/docs/design/manifest.md index 1a3874eb8..a93fd62ce 100644 --- a/docs/design/manifest.md +++ b/docs/design/manifest.md @@ -14,7 +14,7 @@ One or more YAML files per family (single file by default; large families may sh ```mermaid flowchart LR - H["Human reviewer"] -->|writes / approves| M["tileops/manifest/"] + R["Authoritative reference"] -->|specified from| M["tileops/manifest/"] M -->|reads spec from| A["Agent (codegen)"] A -->|produces| C["Op code, tests, benchmarks"] M -->|validates against| V["Validator (CI)"] diff --git a/tileops/manifest/attention.yaml b/tileops/manifest/attention.yaml index d7c749b75..6a19eda06 100644 --- a/tileops/manifest/attention.yaml +++ b/tileops/manifest/attention.yaml @@ -166,6 +166,10 @@ GroupedQueryAttentionPrefillFwdOp: window_size_right: {type: int, default: -1} backend: {type: str, default: "auto"} validate_uniform_cu_seqlens: {type: bool, default: true} + # Element type of `o`. The float16 / bfloat16 rows of dtype_combos pin it + # to q.dtype; the float8_e4m3fn rows admit either float16 or bfloat16, so + # an fp8 caller selects between them here. + dtype: {type: torch.dtype, default: float16} dtype_combos: - {q: float16, k: float16, v: float16, cu_seqlens_q: int32, cu_seqlens_kv: int32, q_scale: float32, k_scale: float32, v_scale: float32, o: float16} - {q: bfloat16, k: bfloat16, v: bfloat16, cu_seqlens_q: int32, cu_seqlens_kv: int32, q_scale: float32, k_scale: float32, v_scale: float32, o: bfloat16} From 2a3c48d0387263eac52b2a536b9e841524e153ad Mon Sep 17 00:00:00 2001 From: lcy-seso Date: Tue, 4 Aug 2026 16:37:24 +0800 Subject: [PATCH 05/15] [Refactor][MoE] Read dtype at forward for the MoE family MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit All nine MoE constructors stop taking `dtype`. The four leaves build their kernel on first forward and cache it per dtype; the five composites had nothing left to forward once the leaves stopped accepting it, so the parameter simply disappears from them. `FusedMoEExpertsNopadPersistent3WGFwdOp` anchors its shared dtype validator on `hidden_states.dtype` — the helper already required output, `w_gate_up` and `w_down` to agree with it, so the constructor's copy was the redundant one. `SharedFusedMoE` builds its shared-expert MLP kernel lazily per dtype; a `Kernel` still takes `dtype`, only the op stops deciding it. `build_activation_op` drops its `dtype` parameter: the `FusedGatedOp` it returns already accepted `None` and deferred the build to forward, so the MoE pipeline never needed to name a dtype for its activation. Call sites needed three forms handled beyond a keyword argument: a positional fourth argument to `MoeUnpermuteFwdOp` that would otherwise have bound silently to `padded_batch_sum`, and a `dict(...)` of kwargs unpacked at the call. --- .../ops/bench_moe_grouped_gemm_nopad.py | 2 +- benchmarks/ops/bench_moe_permute_nopad.py | 2 +- benchmarks/ops/bench_moe_shared_fused_moe.py | 1 - benchmarks/ops/bench_moe_unpermute.py | 2 +- tests/ops/test_fused_moe_experts.py | 51 +++++++++---------- tests/ops/test_moe_fused_moe.py | 18 +++---- tests/ops/test_moe_fused_moe_distributed.py | 6 --- tests/ops/test_moe_grouped_gemm_nopad.py | 2 +- tests/ops/test_moe_permute_nopad.py | 3 +- tests/ops/test_moe_shared_fused_moe.py | 6 --- .../test_moe_shared_fused_moe_distributed.py | 1 - tests/ops/test_moe_unpermute.py | 14 ++--- tileops/ops/moe/_activation.py | 6 +-- tileops/ops/moe/fused_moe.py | 7 --- tileops/ops/moe/routed_expert/abc.py | 4 +- .../moe/routed_expert/fused_routed_expert.py | 19 +++---- .../routed_expert/moe_grouped_gemm_nopad.py | 21 +++++--- .../moe_grouped_gemm_nopad_fused_act.py | 22 +++++--- .../ops/moe/routed_expert/permute_nopad.py | 14 +---- tileops/ops/moe/routed_expert/unpermute.py | 20 +++++--- tileops/ops/moe/shared_fused_moe.py | 32 +++++++----- 21 files changed, 120 insertions(+), 133 deletions(-) diff --git a/benchmarks/ops/bench_moe_grouped_gemm_nopad.py b/benchmarks/ops/bench_moe_grouped_gemm_nopad.py index b53f4081b..7d696d7b4 100644 --- a/benchmarks/ops/bench_moe_grouped_gemm_nopad.py +++ b/benchmarks/ops/bench_moe_grouped_gemm_nopad.py @@ -49,7 +49,7 @@ def test_moe_grouped_gemm_nopad_bench( workload = MoeGroupedGemmNopadWorkload(numel, num_experts, n, k, dtype) a, b, true_sizes, true_offsets = workload.gen_inputs() - op = MoeGroupedGemmNopadFwdOp(numel, num_experts, n, k, dtype=dtype) + op = MoeGroupedGemmNopadFwdOp(numel, num_experts, n, k) bm = ManifestBenchmark(_OP_NAME, op, workload) # Warmup: trigger JIT compilation before timed profiling. diff --git a/benchmarks/ops/bench_moe_permute_nopad.py b/benchmarks/ops/bench_moe_permute_nopad.py index 6c37a34d8..7abfdcabd 100644 --- a/benchmarks/ops/bench_moe_permute_nopad.py +++ b/benchmarks/ops/bench_moe_permute_nopad.py @@ -66,7 +66,7 @@ def test_moe_permute_nopad_bench( hidden_states, topk_ids = workload.gen_inputs() # TileOPs - op = MoePermuteNopadFwdOp(num_experts=num_experts, dtype=dtype) + op = MoePermuteNopadFwdOp(num_experts=num_experts) bm = ManifestBenchmark(_OP_NAME, op, workload) op(hidden_states, topk_ids) # warmup / JIT compile torch.cuda.synchronize() diff --git a/benchmarks/ops/bench_moe_shared_fused_moe.py b/benchmarks/ops/bench_moe_shared_fused_moe.py index d95af2da5..761235eed 100644 --- a/benchmarks/ops/bench_moe_shared_fused_moe.py +++ b/benchmarks/ops/bench_moe_shared_fused_moe.py @@ -133,7 +133,6 @@ def test_shared_fused_moe_bench( renormalize=renormalize, with_correction_bias=with_correction_bias, routed_scaling_factor=routed_scaling_factor, - dtype=dtype, shared_ffn_size=shared_ffn_size, ) op(hidden, gating, w_gate_up, w_down, correction_bias, diff --git a/benchmarks/ops/bench_moe_unpermute.py b/benchmarks/ops/bench_moe_unpermute.py index d084628b0..4e2de8447 100644 --- a/benchmarks/ops/bench_moe_unpermute.py +++ b/benchmarks/ops/bench_moe_unpermute.py @@ -63,7 +63,7 @@ def test_moe_unpermute_bench(total_tokens: int, top_k: int, hidden_size: int) -> mm2_pad, fwd_idx, topk_weights = test.gen_inputs() # TileOPs - op = MoeUnpermuteFwdOp(total_tokens, top_k, hidden_size, dtype) + op = MoeUnpermuteFwdOp(total_tokens, top_k, hidden_size) bm = ManifestBenchmark(_OP_NAME, op, test) op(mm2_pad, fwd_idx, topk_weights) # warmup / JIT compile torch.cuda.synchronize() diff --git a/tests/ops/test_fused_moe_experts.py b/tests/ops/test_fused_moe_experts.py index 6015a1349..e6b9a6c96 100644 --- a/tests/ops/test_fused_moe_experts.py +++ b/tests/ops/test_fused_moe_experts.py @@ -167,7 +167,7 @@ def test_workspace_shapes(self, moe_meta): d = moe_meta experts = FusedMoEExpertsNopadPersistent3WGFwdOp( num_tokens=d["T"], num_experts=d["E"], top_k=d["K"], - hidden_size=d["H"], ffn_size=d["F"], dtype=d["dtype"], + hidden_size=d["H"], ffn_size=d["F"], ) ws1, ws2 = experts.workspace_shapes(d["T"], d["F"], d["H"], d["K"], d["E"]) assert ws1 == (0,) and ws2 == (0,) @@ -177,7 +177,7 @@ def test_output_shape(self, moe_meta): d = moe_meta experts = FusedMoEExpertsNopadPersistent3WGFwdOp( num_tokens=d["T"], num_experts=d["E"], top_k=d["K"], - hidden_size=d["H"], ffn_size=d["F"], dtype=d["dtype"], + hidden_size=d["H"], ffn_size=d["F"], ) assert experts.output_shape(d["T"], d["H"]) == (d["T"], d["H"]) @@ -186,7 +186,7 @@ def test_make_weighted_reduce_is_noop(self, moe_meta): d = moe_meta experts = FusedMoEExpertsNopadPersistent3WGFwdOp( num_tokens=d["T"], num_experts=d["E"], top_k=d["K"], - hidden_size=d["H"], ffn_size=d["F"], dtype=d["dtype"], + hidden_size=d["H"], ffn_size=d["F"], ) assert isinstance(experts.make_weighted_reduce(), WeightedReduceNoOp) @@ -196,7 +196,7 @@ def test_forward_matches_torch_ref(self, moe_tensors): d = moe_tensors experts = FusedMoEExpertsNopadPersistent3WGFwdOp( num_tokens=d["T"], num_experts=d["E"], top_k=d["K"], - hidden_size=d["H"], ffn_size=d["F"], dtype=d["dtype"], + hidden_size=d["H"], ffn_size=d["F"], ) ref_out = _torch_ref_moe(d["hidden"], d["w1"], d["w2"], d["weights"], d["ids"]) @@ -233,7 +233,7 @@ def test_forward_fallback_path_unaligned_dims(self, caplog): ): experts = FusedMoEExpertsNopadPersistent3WGFwdOp( num_tokens=T, num_experts=E, top_k=K, - hidden_size=H, ffn_size=F_dim, dtype=dtype, + hidden_size=H, ffn_size=F_dim, ) assert any( "falling back to MoeGroupedGemmNopadKernel" in rec.message @@ -277,7 +277,7 @@ def test_forward_with_expert_map_runs(self): experts = FusedMoEExpertsNopadPersistent3WGFwdOp( num_tokens=T, num_experts=E_global, top_k=K, - hidden_size=H, ffn_size=F_dim, dtype=dtype, + hidden_size=H, ffn_size=F_dim, expert_map=expert_map, ) output = torch.empty(T, H, dtype=dtype, device="cuda") @@ -298,7 +298,7 @@ def test_forward_matches_torch_ref_activation(self, moe_tensors, activation): d = moe_tensors experts = FusedMoEExpertsNopadPersistent3WGFwdOp( num_tokens=d["T"], num_experts=d["E"], top_k=d["K"], - hidden_size=d["H"], ffn_size=d["F"], dtype=d["dtype"], + hidden_size=d["H"], ffn_size=d["F"], activation=activation, ) assert experts.activation == activation @@ -336,7 +336,7 @@ def test_use_fused_activation_parity(activation): def run(use_fused): op = FusedMoEExpertsNopadPersistent3WGFwdOp( num_tokens=T_count, num_experts=E, top_k=top_k, hidden_size=H, - ffn_size=Fdim, dtype=torch.bfloat16, activation=activation, + ffn_size=Fdim, activation=activation, use_fused_activation=use_fused) out = torch.empty(T_count, H, dtype=torch.bfloat16, device="cuda") ws = torch.empty(0, dtype=torch.bfloat16, device="cuda") @@ -364,7 +364,7 @@ def test_use_fused_activation_disabled_on_gemm_override(): from tileops.kernels.moe.moe_grouped_gemm_nopad import MoeGroupedGemmNopadKernel experts = FusedMoEExpertsNopadPersistent3WGFwdOp( num_tokens=256, num_experts=8, top_k=2, hidden_size=256, ffn_size=768, - dtype=torch.bfloat16, activation="silu_and_mul", use_fused_activation=True, + activation="silu_and_mul", use_fused_activation=True, kernel_map={"moe_grouped_gemm_kernel": MoeGroupedGemmNopadKernel}, ) assert experts.use_fused_activation is False @@ -379,7 +379,7 @@ def test_top_level_api_forwards_use_fused_activation(): from tileops.ops.moe.fused_moe import FusedMoe, FusedMoeFwdCbFwdOp, FusedMoeFwdOp from tileops.ops.moe.shared_fused_moe import SharedFusedMoE common = dict(num_tokens=256, num_experts=8, top_k=2, hidden_size=256, - ffn_size=768, dtype=torch.bfloat16) + ffn_size=768) for cls in (FusedMoe, FusedMoeFwdOp, FusedMoeFwdCbFwdOp, SharedFusedMoE): assert cls(**common, use_fused_activation=True)._experts.use_fused_activation is True assert cls(**common)._experts.use_fused_activation is False @@ -393,11 +393,10 @@ def test_use_fused_activation_rejected_with_injected_experts(): pytest.skip("Requires SM90") from tileops.ops.moe.fused_moe import FusedMoeFwdOp experts = FusedMoEExpertsNopadPersistent3WGFwdOp( - num_tokens=256, num_experts=8, top_k=2, hidden_size=256, ffn_size=768, - dtype=torch.bfloat16) + num_tokens=256, num_experts=8, top_k=2, hidden_size=256, ffn_size=768) with pytest.raises(ValueError, match="use_fused_activation"): FusedMoeFwdOp(num_tokens=256, num_experts=8, top_k=2, hidden_size=256, - ffn_size=768, dtype=torch.bfloat16, experts=experts, + ffn_size=768, experts=experts, use_fused_activation=True) @@ -406,19 +405,19 @@ class TestBuildActivationOp: @pytest.mark.smoke def test_silu_and_mul_returns_correct_type(self): from tileops.ops.elementwise import SiluAndMulFwdOp - op = build_activation_op("silu_and_mul", M=16, N=32, dtype=torch.bfloat16) + op = build_activation_op("silu_and_mul", M=16, N=32) assert isinstance(op, SiluAndMulFwdOp) @pytest.mark.smoke def test_gelu_and_mul_returns_correct_type(self): from tileops.ops.elementwise import GeluAndMulFwdOp - op = build_activation_op("gelu_and_mul", M=16, N=32, dtype=torch.bfloat16) + op = build_activation_op("gelu_and_mul", M=16, N=32) assert isinstance(op, GeluAndMulFwdOp) @pytest.mark.smoke def test_invalid_activation_raises(self): with pytest.raises(ValueError, match="activation must be one of"): - build_activation_op("unknown_act", M=16, N=32, dtype=torch.bfloat16) + build_activation_op("unknown_act", M=16, N=32) class TestFusedMoeActivationInjection: @@ -426,7 +425,7 @@ class TestFusedMoeActivationInjection: def _make_experts(self, activation="silu_and_mul"): return FusedMoEExpertsNopadPersistent3WGFwdOp( num_tokens=128, num_experts=4, top_k=2, - hidden_size=256, ffn_size=128, dtype=torch.bfloat16, + hidden_size=256, ffn_size=128, activation=activation, ) @@ -438,7 +437,7 @@ def test_injection_with_conflicting_activation_raises(self): with pytest.raises(ValueError, match="activation conflicts"): FusedMoe( num_tokens=128, num_experts=4, top_k=2, - hidden_size=256, ffn_size=128, dtype=torch.bfloat16, + hidden_size=256, ffn_size=128, experts=experts, activation="gelu_and_mul", ) @@ -449,7 +448,7 @@ def test_injection_with_matching_activation_works(self): experts = self._make_experts(activation="gelu_and_mul") moe = FusedMoe( num_tokens=128, num_experts=4, top_k=2, - hidden_size=256, ffn_size=128, dtype=torch.bfloat16, + hidden_size=256, ffn_size=128, experts=experts, activation="gelu_and_mul", ) assert moe.activation == "gelu_and_mul" @@ -461,7 +460,7 @@ def test_injection_without_activation_works(self): experts = self._make_experts() moe = FusedMoe( num_tokens=128, num_experts=4, top_k=2, - hidden_size=256, ffn_size=128, dtype=torch.bfloat16, + hidden_size=256, ffn_size=128, experts=experts, ) assert moe.activation == "silu_and_mul" @@ -472,7 +471,7 @@ def test_default_path_activation_forwarded(self): from tileops.ops.moe.fused_moe import FusedMoe moe = FusedMoe( num_tokens=128, num_experts=4, top_k=2, - hidden_size=256, ffn_size=128, dtype=torch.bfloat16, + hidden_size=256, ffn_size=128, activation="gelu_and_mul", ) assert moe.activation == "gelu_and_mul" @@ -517,7 +516,7 @@ def make_weighted_reduce(self): with pytest.raises(ValueError, match="missing the required `.activation`"): FusedMoe( num_tokens=128, num_experts=4, top_k=2, - hidden_size=256, ffn_size=128, dtype=torch.bfloat16, + hidden_size=256, ffn_size=128, experts=ExpertsWithoutActivation(), ) @@ -530,7 +529,7 @@ def test_activation_forwarded_to_routed_experts(self): from tileops.ops.moe.shared_fused_moe import SharedFusedMoE moe = SharedFusedMoE( num_tokens=128, num_experts=4, top_k=2, - hidden_size=256, ffn_size=128, dtype=torch.bfloat16, + hidden_size=256, ffn_size=128, activation="gelu_and_mul", ) assert moe.activation == "gelu_and_mul" @@ -548,7 +547,7 @@ def test_shared_expert_with_non_default_activation_raises(self): with pytest.raises(NotImplementedError, match="shared-expert path only supports"): SharedFusedMoE( num_tokens=128, num_experts=4, top_k=2, - hidden_size=256, ffn_size=128, dtype=torch.bfloat16, + hidden_size=256, ffn_size=128, shared_ffn_size=128, activation="gelu_and_mul", ) @@ -559,7 +558,7 @@ def test_shared_expert_with_default_activation_works(self): from tileops.ops.moe.shared_fused_moe import SharedFusedMoE moe = SharedFusedMoE( num_tokens=128, num_experts=4, top_k=2, - hidden_size=256, ffn_size=128, dtype=torch.bfloat16, + hidden_size=256, ffn_size=128, shared_ffn_size=128, ) assert moe.activation == "silu_and_mul" @@ -578,7 +577,7 @@ def test_fused_act_fwd_op_shape_and_values(): A = torch.randn(numel, K, dtype=torch.bfloat16, device="cuda") * 0.02 B = torch.randn(E, 2 * ffn, K, dtype=torch.bfloat16, device="cuda") * 0.02 op = MoeGroupedGemmNopad3WGFusedActFwdOp( - numel=numel, num_experts=E, ffn=ffn, k=K, dtype=torch.bfloat16, + numel=numel, num_experts=E, ffn=ffn, k=K, activation="silu_and_mul") out = op(A, B, sizes, offsets) assert out.shape == (numel, ffn) diff --git a/tests/ops/test_moe_fused_moe.py b/tests/ops/test_moe_fused_moe.py index d45865c20..d92830edc 100644 --- a/tests/ops/test_moe_fused_moe.py +++ b/tests/ops/test_moe_fused_moe.py @@ -137,7 +137,6 @@ def test_fused_moe_qwen3( num_tokens=num_tokens, num_experts=num_experts, top_k=top_k, hidden_size=hidden_size, ffn_size=ffn_size, scoring_func=scoring_func, renormalize=renormalize, - dtype=dtype, ) out_nopad = op_nopad(hidden, gating, w_gate_up, w_down) @@ -208,7 +207,7 @@ def test_fused_moe_deterministic(case): op = FusedMoe( num_tokens=nt, num_experts=ne, top_k=tk, hidden_size=hs, - ffn_size=ff, scoring_func="softmax", renormalize=False, dtype=dtype, + ffn_size=ff, scoring_func="softmax", renormalize=False, ) fk = FusedTopKOp(nt, ne, tk, "softmax", False) topk_weights, topk_ids = fk(gating) @@ -292,7 +291,6 @@ def test_fused_moe_kimi( scoring_func="sigmoid", renormalize=True, with_correction_bias=with_correction_bias, routed_scaling_factor=routed_scaling_factor, - dtype=dtype, ) out_nopad = op_nopad(hidden, gating, w_gate_up, w_down, correction_bias) @@ -344,14 +342,14 @@ def test_expert_map_local_filter() -> None: # Full output (no expert_map) op_full = FusedMoe( num_tokens=T, num_experts=E, top_k=K, - hidden_size=H, ffn_size=F, dtype=dtype, + hidden_size=H, ffn_size=F, ) out_full = op_full(hidden, gating, w_gate_up, w_down) # Rank-0 partial output (local experts 0..3) op_r0 = FusedMoe( num_tokens=T, num_experts=E, top_k=K, - hidden_size=H, ffn_size=F, dtype=dtype, + hidden_size=H, ffn_size=F, expert_map=expert_map_rank0, ) out_r0 = op_r0(hidden, gating, w_gate_up[:E // 2], w_down[:E // 2]) @@ -359,7 +357,7 @@ def test_expert_map_local_filter() -> None: # Rank-1 partial output (local experts 4..7) op_r1 = FusedMoe( num_tokens=T, num_experts=E, top_k=K, - hidden_size=H, ffn_size=F, dtype=dtype, + hidden_size=H, ffn_size=F, expert_map=expert_map_rank1, ) out_r1 = op_r1(hidden, gating, w_gate_up[E // 2:], w_down[E // 2:]) @@ -463,7 +461,6 @@ def test_fused_moe_vs_vllm( hidden_size=hidden_size, ffn_size=ffn_size, scoring_func="sigmoid", renormalize=True, with_correction_bias=True, routed_scaling_factor=routed_scaling_factor, - dtype=dtype, ) out_tileops = op(hidden, gating, w_gate_up, w_down, correction_bias) @@ -494,7 +491,7 @@ def test_fused_moe_fwd_op_identity() -> None: op = FusedMoeFwdOp( num_tokens=T, num_experts=E, top_k=K, - hidden_size=H, ffn_size=F_, dtype=dtype, + hidden_size=H, ffn_size=F_, ) out = op(hidden, gating, w_gate_up, w_down) assert out.shape == (T, H) @@ -502,7 +499,7 @@ def test_fused_moe_fwd_op_identity() -> None: ref_op = FusedMoe( num_tokens=T, num_experts=E, top_k=K, - hidden_size=H, ffn_size=F_, dtype=dtype, + hidden_size=H, ffn_size=F_, ) ref = ref_op(hidden, gating, w_gate_up, w_down) torch.testing.assert_close(out.float(), ref.float(), rtol=1e-2, atol=1e-2) @@ -527,7 +524,7 @@ def test_fused_moe_fwd_cb_op_identity() -> None: op = FusedMoeFwdCbFwdOp( num_tokens=T, num_experts=E, top_k=K, - hidden_size=H, ffn_size=F_, renormalize=True, dtype=dtype, + hidden_size=H, ffn_size=F_, renormalize=True, ) out = op(hidden, gating, correction_bias, w_gate_up, w_down) assert out.shape == (T, H) @@ -537,7 +534,6 @@ def test_fused_moe_fwd_cb_op_identity() -> None: num_tokens=T, num_experts=E, top_k=K, hidden_size=H, ffn_size=F_, scoring_func="sigmoid", renormalize=True, with_correction_bias=True, - dtype=dtype, ) ref = ref_op(hidden, gating, w_gate_up, w_down, correction_bias) torch.testing.assert_close(out.float(), ref.float(), rtol=1e-2, atol=1e-2) diff --git a/tests/ops/test_moe_fused_moe_distributed.py b/tests/ops/test_moe_fused_moe_distributed.py index bf3027373..86a449a45 100644 --- a/tests/ops/test_moe_fused_moe_distributed.py +++ b/tests/ops/test_moe_fused_moe_distributed.py @@ -137,7 +137,6 @@ def test_kimi_k2_ep_distributed(T, E_global, K, H, F, world_size): hidden_size=H, ffn_size=F, scoring_func="sigmoid", renormalize=True, with_correction_bias=True, routed_scaling_factor=2.827, - dtype=dtype, expert_map=expert_map, ) out_local = op_local(hidden, gating, w_gate_up_local, w_down_local, correction_bias) @@ -152,7 +151,6 @@ def test_kimi_k2_ep_distributed(T, E_global, K, H, F, world_size): hidden_size=H, ffn_size=F, scoring_func="sigmoid", renormalize=True, with_correction_bias=True, routed_scaling_factor=2.827, - dtype=dtype, ) out_full = op_full(hidden, gating, w_gate_up_full, w_down_full, correction_bias) @@ -213,7 +211,6 @@ def test_shared_fused_moe_ep_distributed(T, E_global, K, H, F, shared_F, world_s hidden_size=H, ffn_size=F, scoring_func="sigmoid", renormalize=True, with_correction_bias=True, routed_scaling_factor=2.827, - dtype=dtype, expert_map=expert_map, shared_ffn_size=shared_F, ) @@ -232,7 +229,6 @@ def test_shared_fused_moe_ep_distributed(T, E_global, K, H, F, shared_F, world_s hidden_size=H, ffn_size=F, scoring_func="sigmoid", renormalize=True, with_correction_bias=True, routed_scaling_factor=2.827, - dtype=dtype, shared_ffn_size=shared_F, ) shared_out_full, routed_out_full = op_full( @@ -306,7 +302,6 @@ def test_fused_moe_vs_vllm_distributed(T, E_global, K, H, F, world_size): hidden_size=H, ffn_size=F, scoring_func="sigmoid", renormalize=True, with_correction_bias=True, routed_scaling_factor=2.827, - dtype=dtype, expert_map=expert_map, ) out_tileops = op_tileops(hidden, gating, w_gate_up_local, w_down_local, correction_bias) @@ -381,7 +376,6 @@ def test_fused_moe_vs_vllm_ep_layer(T, E_global, K, H, F, world_size): hidden_size=H, ffn_size=F, scoring_func="sigmoid", renormalize=True, with_correction_bias=True, routed_scaling_factor=2.827, - dtype=dtype, expert_map=expert_map, ) out_tileops = op_tileops(hidden, gating, w_gate_up_local, w_down_local, correction_bias) diff --git a/tests/ops/test_moe_grouped_gemm_nopad.py b/tests/ops/test_moe_grouped_gemm_nopad.py index c88d7b461..7efeff1be 100644 --- a/tests/ops/test_moe_grouped_gemm_nopad.py +++ b/tests/ops/test_moe_grouped_gemm_nopad.py @@ -79,7 +79,7 @@ def test_moe_grouped_gemm_nopad_op(numel, num_experts, n, k, distribution, dtype ) a, b, true_sizes, true_offsets = test.gen_inputs() - op = MoeGroupedGemmNopadFwdOp(numel, num_experts, n, k, dtype=dtype) + op = MoeGroupedGemmNopadFwdOp(numel, num_experts, n, k) c = op(a, b, true_sizes, true_offsets) c_ref = _ref_grouped_gemm_nopad(a, b, true_sizes, true_offsets) diff --git a/tests/ops/test_moe_permute_nopad.py b/tests/ops/test_moe_permute_nopad.py index 9d2e39d36..bff53c2cd 100644 --- a/tests/ops/test_moe_permute_nopad.py +++ b/tests/ops/test_moe_permute_nopad.py @@ -114,7 +114,7 @@ class MoePermuteNopadFixture(FixtureBase): @MoePermuteNopadFixture def test_moe_permute_nopad_op(total_tokens, top_k, num_experts, hidden_size, dtype): test = MoePermuteNopadTest(total_tokens, top_k, num_experts, hidden_size, dtype) - op = MoePermuteNopadFwdOp(num_experts=num_experts, dtype=dtype) + op = MoePermuteNopadFwdOp(num_experts=num_experts) hidden_states, topk_ids = test.gen_inputs() outputs = op(hidden_states, topk_ids) @@ -135,7 +135,6 @@ def test_moe_permute_nopad_explicit_shape_mismatch_raises() -> None: top_k=2, num_experts=4, hidden_size=16, - dtype=torch.float16, ) with pytest.raises(ValueError, match="Expected total_tokens"): op(hidden_states, topk_ids) diff --git a/tests/ops/test_moe_shared_fused_moe.py b/tests/ops/test_moe_shared_fused_moe.py index 9c091883a..be8395fa3 100644 --- a/tests/ops/test_moe_shared_fused_moe.py +++ b/tests/ops/test_moe_shared_fused_moe.py @@ -33,7 +33,6 @@ def test_shared_fused_moe_basic(): num_tokens=T, num_experts=E, top_k=K, hidden_size=H, ffn_size=F, scoring_func="softmax", renormalize=False, - dtype=dtype, shared_ffn_size=F_s, ) @@ -60,7 +59,6 @@ def test_shared_fused_moe_basic(): num_tokens=T, num_experts=E, top_k=K, hidden_size=H, ffn_size=F, scoring_func="softmax", renormalize=False, - dtype=dtype, ) routed_ref = op_routed(hidden, gating, w_gate_up, w_down) torch.testing.assert_close(routed_out, routed_ref, rtol=1e-5, atol=1e-5) @@ -83,7 +81,6 @@ def test_shared_fused_moe_none(): op = SharedFusedMoE( num_tokens=T, num_experts=E, top_k=K, hidden_size=H, ffn_size=F, - dtype=dtype, ) shared_out, routed_out = op(hidden, gating, w_gate_up, w_down) @@ -132,7 +129,6 @@ def test_shared_fused_moe_tp(): num_tokens=T, num_experts=E, top_k=K, hidden_size=H, ffn_size=F, scoring_func="softmax", renormalize=False, - dtype=dtype, ) routed_ref = op_routed(hidden, gating, w_gate_up, w_down) @@ -143,7 +139,6 @@ def test_shared_fused_moe_tp(): num_tokens=T, num_experts=E, top_k=K, hidden_size=H, ffn_size=F, scoring_func="softmax", renormalize=False, - dtype=dtype, shared_ffn_size=F_s, tp_size=tp_size, tp_rank=tp_rank, ) @@ -179,7 +174,6 @@ def test_shared_fused_moe_tp_rejects_local_shards(): num_tokens=T, num_experts=E, top_k=K, hidden_size=H, ffn_size=F, scoring_func="softmax", renormalize=False, - dtype=dtype, shared_ffn_size=F_s, tp_size=tp_size, tp_rank=0, ) diff --git a/tests/ops/test_moe_shared_fused_moe_distributed.py b/tests/ops/test_moe_shared_fused_moe_distributed.py index 95257629d..8a85df7ea 100644 --- a/tests/ops/test_moe_shared_fused_moe_distributed.py +++ b/tests/ops/test_moe_shared_fused_moe_distributed.py @@ -191,7 +191,6 @@ def test_shared_expert_tp_vs_vllm(T, H, F_s, tp_size): num_tokens=T, num_experts=E, top_k=K, hidden_size=H, ffn_size=F, scoring_func="softmax", renormalize=False, - dtype=dtype, shared_ffn_size=F_s, tp_size=tp_size, tp_rank=tp_rank, ) diff --git a/tests/ops/test_moe_unpermute.py b/tests/ops/test_moe_unpermute.py index 27a8b7098..c40aeb694 100644 --- a/tests/ops/test_moe_unpermute.py +++ b/tests/ops/test_moe_unpermute.py @@ -84,7 +84,7 @@ class MoeUnpermuteFixture(FixtureBase): @MoeUnpermuteFixture def test_moe_unpermute_op(total_tokens, top_k, hidden_size, dtype): test = MoeUnpermuteTest(total_tokens, top_k, hidden_size, dtype) - op = MoeUnpermuteFwdOp(total_tokens, top_k, hidden_size, dtype) + op = MoeUnpermuteFwdOp(total_tokens, top_k, hidden_size) mm2_pad, fwd_idx, topk_weights = test.gen_inputs() output = op(mm2_pad, fwd_idx, topk_weights) @@ -105,7 +105,7 @@ def test_moe_unpermute_skewed(): fwd_idx = torch.arange(numel, dtype=torch.int32, device="cuda") % K topk_weights = torch.rand(T, K, dtype=torch.float32, device="cuda") - op = MoeUnpermuteFwdOp(T, K, H, torch.bfloat16, padded_batch_sum=numel) + op = MoeUnpermuteFwdOp(T, K, H, padded_batch_sum=numel) output = op(mm2_pad, fwd_idx, topk_weights) output_ref = _ref_moe_unpermute(mm2_pad, fwd_idx, topk_weights) @@ -132,7 +132,7 @@ def test_moe_unpermute_ep_masking(): assert (fwd_idx < 0).any(), "test must actually inject -1 slots" topk_weights = torch.rand(T, K, dtype=torch.float32, device=dev) - op = MoeUnpermuteFwdOp(T, K, H, torch.bfloat16, padded_batch_sum=numel) + op = MoeUnpermuteFwdOp(T, K, H, padded_batch_sum=numel) output = op(mm2_pad, fwd_idx, topk_weights) # Reference matching the kernel's -1 semantics: skip non-local slots. @@ -160,7 +160,7 @@ def test_moe_unpermute_out_param(): topk_weights = torch.softmax( torch.randn(T, K, dtype=torch.float32, device=dev), dim=-1) - op = MoeUnpermuteFwdOp(T, K, H, torch.bfloat16, padded_batch_sum=numel) + op = MoeUnpermuteFwdOp(T, K, H, padded_batch_sum=numel) ref = op(mm2_pad, fwd_idx, topk_weights) out = torch.empty((T, H), dtype=torch.bfloat16, device=dev) @@ -182,11 +182,11 @@ def test_moe_unpermute_scaling(): torch.randn(T, K, dtype=torch.float32, device=dev), dim=-1) scale = 2.827 - base = MoeUnpermuteFwdOp(T, K, H, torch.bfloat16, padded_batch_sum=numel) + base = MoeUnpermuteFwdOp(T, K, H, padded_batch_sum=numel) ref = base(mm2_pad, fwd_idx, topk_weights).float() * scale scaled = MoeUnpermuteFwdOp( - T, K, H, torch.bfloat16, padded_batch_sum=numel, routed_scaling_factor=scale) + T, K, H, padded_batch_sum=numel, routed_scaling_factor=scale) got = scaled(mm2_pad, fwd_idx, topk_weights).float() torch.testing.assert_close(got, ref, rtol=2e-2, atol=2e-2) @@ -203,7 +203,7 @@ def test_moe_unpermute_out_buffer_validation(): fwd_idx = torch.arange(numel, dtype=torch.int32, device=dev) topk_weights = torch.softmax( torch.randn(T, K, dtype=torch.float32, device=dev), dim=-1) - op = MoeUnpermuteFwdOp(T, K, H, torch.bfloat16, padded_batch_sum=numel) + op = MoeUnpermuteFwdOp(T, K, H, padded_batch_sum=numel) with pytest.raises(ValueError, match="contiguous"): op(mm2_pad, fwd_idx, topk_weights, diff --git a/tileops/ops/moe/_activation.py b/tileops/ops/moe/_activation.py index abde8b3c1..e6323eb0b 100644 --- a/tileops/ops/moe/_activation.py +++ b/tileops/ops/moe/_activation.py @@ -3,8 +3,6 @@ from typing import Dict, Optional -import torch - from tileops.kernels.kernel_base import Kernel from tileops.ops.elementwise import FusedGatedOp, GeluAndMulFwdOp, SiluAndMulFwdOp @@ -20,7 +18,6 @@ def build_activation_op( activation: str, M: int, N: int, - dtype: torch.dtype, kernel_map: Optional[Dict[str, Kernel]] = None, ) -> FusedGatedOp: """Construct the activation sub-Op for an MoE experts pipeline. @@ -29,7 +26,6 @@ def build_activation_op( activation: One of the keys in ``_ACTIVATION_REGISTRY``. M: Row count of the (M, 2N) gate_up tensor. N: Half column dim — the activation output width (= ffn_size). - dtype: Activation/output dtype. kernel_map: Forwarded to the inner activation op for kernel dispatch. Returns: @@ -44,5 +40,5 @@ def build_activation_op( f"activation must be one of [{allowed}], got {activation!r}" ) return _ACTIVATION_REGISTRY[activation]( - M=M, N=N, dtype=dtype, kernel_map=kernel_map, + M=M, N=N, kernel_map=kernel_map, ) diff --git a/tileops/ops/moe/fused_moe.py b/tileops/ops/moe/fused_moe.py index a06d3ae72..5b066f8c4 100644 --- a/tileops/ops/moe/fused_moe.py +++ b/tileops/ops/moe/fused_moe.py @@ -68,7 +68,6 @@ def __init__( renormalize: bool = False, with_correction_bias: bool = False, routed_scaling_factor: float = 1.0, - dtype: torch.dtype = torch.bfloat16, expert_map: Optional[torch.Tensor] = None, prepare_finalize: Optional[FusedMoEPrepareAndFinalize] = None, experts: Optional[FusedMoEExpertsModular] = None, @@ -86,7 +85,6 @@ def __init__( self.renormalize = renormalize self.with_correction_bias = with_correction_bias self.routed_scaling_factor = routed_scaling_factor - self.dtype = dtype self.expert_map = expert_map self.dispatch_kernel(kernel_map) @@ -160,7 +158,6 @@ def __init__( ffn_size=ffn_size, activation=activation, routed_scaling_factor=routed_scaling_factor, - dtype=dtype, expert_map=expert_map, kernel_map=kernel_map, use_fused_activation=use_fused_activation, @@ -229,7 +226,6 @@ def __init__( scoring_func: str = "softmax", renormalize: bool = False, routed_scaling_factor: float = 1.0, - dtype: torch.dtype = torch.bfloat16, expert_map: Optional[torch.Tensor] = None, prepare_finalize: Optional[FusedMoEPrepareAndFinalize] = None, experts: Optional[FusedMoEExpertsModular] = None, @@ -248,7 +244,6 @@ def __init__( renormalize=renormalize, with_correction_bias=False, routed_scaling_factor=routed_scaling_factor, - dtype=dtype, expert_map=expert_map, prepare_finalize=prepare_finalize, experts=experts, @@ -285,7 +280,6 @@ def __init__( scoring_func: str = "sigmoid", renormalize: bool = False, routed_scaling_factor: float = 1.0, - dtype: torch.dtype = torch.bfloat16, expert_map: Optional[torch.Tensor] = None, prepare_finalize: Optional[FusedMoEPrepareAndFinalize] = None, experts: Optional[FusedMoEExpertsModular] = None, @@ -304,7 +298,6 @@ def __init__( renormalize=renormalize, with_correction_bias=True, routed_scaling_factor=routed_scaling_factor, - dtype=dtype, expert_map=expert_map, prepare_finalize=prepare_finalize, experts=experts, diff --git a/tileops/ops/moe/routed_expert/abc.py b/tileops/ops/moe/routed_expert/abc.py index 7beb7fa9b..b0665778c 100644 --- a/tileops/ops/moe/routed_expert/abc.py +++ b/tileops/ops/moe/routed_expert/abc.py @@ -56,7 +56,9 @@ def _validate_fused_moe_experts_dtypes( """ allowed = (torch.float16, torch.bfloat16) if op_dtype not in allowed: - raise ValueError(f"op dtype must be one of {allowed}, got {op_dtype}") + raise ValueError( + f"hidden_states.dtype must be one of {allowed}, got {op_dtype}" + ) for name, t in ( ("output", output), ("hidden_states", hidden_states), diff --git a/tileops/ops/moe/routed_expert/fused_routed_expert.py b/tileops/ops/moe/routed_expert/fused_routed_expert.py index 0e44652b3..b591c48a5 100644 --- a/tileops/ops/moe/routed_expert/fused_routed_expert.py +++ b/tileops/ops/moe/routed_expert/fused_routed_expert.py @@ -110,7 +110,6 @@ def __init__( hidden_size: int, ffn_size: int, routed_scaling_factor: float = 1.0, - dtype: torch.dtype = torch.bfloat16, expert_map: Optional[Tensor] = None, gemm_kernel: Optional[type] = None, kernel_map: Optional[Dict[str, Kernel]] = None, @@ -124,7 +123,6 @@ def __init__( self.top_k = top_k self.hidden_size = hidden_size self.ffn_size = ffn_size - self.dtype = dtype numel = num_tokens * top_k num_experts_local = ( int((expert_map >= 0).sum().item()) if expert_map is not None else num_experts @@ -183,7 +181,7 @@ def __init__( self.use_fused_activation = False self._permute = MoePermuteNopadFwdOp( - num_experts=num_experts, dtype=dtype, expert_map=expert_map, + num_experts=num_experts, expert_map=expert_map, kernel_map=kernel_map, ) self.activation = activation @@ -193,27 +191,27 @@ def __init__( ) self._gemm_gate_up = MoeGroupedGemmNopad3WGFusedActFwdOp( numel=numel, num_experts=num_experts_local, - ffn=ffn_size, k=hidden_size, dtype=dtype, activation=activation, + ffn=ffn_size, k=hidden_size, activation=activation, kernel_map=kernel_map, ) self._activation_op = None else: self._gemm_gate_up = MoeGroupedGemmNopadFwdOp( numel=numel, num_experts=num_experts_local, - n=ffn_size * 2, k=hidden_size, dtype=dtype, + n=ffn_size * 2, k=hidden_size, kernel_map={"moe_grouped_gemm_kernel": kernel_cls, **(kernel_map or {})}, ) self._activation_op = build_activation_op( - activation, M=numel, N=ffn_size, dtype=dtype, kernel_map=kernel_map, + activation, M=numel, N=ffn_size, kernel_map=kernel_map, ) self._gemm_down = MoeGroupedGemmNopadFwdOp( numel=numel, num_experts=num_experts_local, - n=hidden_size, k=ffn_size, dtype=dtype, + n=hidden_size, k=ffn_size, kernel_map={"moe_grouped_gemm_kernel": kernel_cls, **(kernel_map or {})}, ) self._unpermute = MoeUnpermuteFwdOp( total_tokens=num_tokens, top_k=top_k, - hidden_size=hidden_size, dtype=dtype, padded_batch_sum=numel, + hidden_size=hidden_size, padded_batch_sum=numel, kernel_map=kernel_map, routed_scaling_factor=routed_scaling_factor, ) @@ -231,8 +229,11 @@ def _validate_dtypes( workspace1: Tensor, workspace2: Tensor, ) -> None: + # hidden_states is the dtype anchor: the helper requires output, + # w_gate_up and w_down to agree with it. + self.dtype = hidden_states.dtype _validate_fused_moe_experts_dtypes( - self.dtype, + hidden_states.dtype, output, hidden_states, w_gate_up, w_down, topk_weights, topk_ids, expert_map, workspace1, workspace2, ) diff --git a/tileops/ops/moe/routed_expert/moe_grouped_gemm_nopad.py b/tileops/ops/moe/routed_expert/moe_grouped_gemm_nopad.py index c70a6d438..95b38f907 100644 --- a/tileops/ops/moe/routed_expert/moe_grouped_gemm_nopad.py +++ b/tileops/ops/moe/routed_expert/moe_grouped_gemm_nopad.py @@ -26,13 +26,12 @@ class MoeGroupedGemmNopadFwdOp(Op): num_experts: Total number of experts E. n: Output feature dimension N (e.g. 2*ffn_size or hidden_size). k: Input feature dimension K (hidden_size or ffn_size). - dtype: Activation and weight dtype (bf16 or fp16). kernel_map: Optional kernel override dict. tune: Whether to autotune. Example: >>> op = MoeGroupedGemmNopadFwdOp(numel=16384, num_experts=256, n=4096, k=2048, - ... dtype=torch.bfloat16) + ...) >>> C = op(A, B, true_sizes, true_offsets) # [numel, N] """ @@ -42,7 +41,6 @@ def __init__( num_experts: int, n: int, k: int, - dtype: torch.dtype = torch.bfloat16, kernel_map: Optional[Dict[str, Kernel]] = None, tune: bool = False, ) -> None: @@ -50,12 +48,18 @@ def __init__( self.num_experts = num_experts self.n = n self.k = k - self.dtype = dtype + self.tune = tune self.dispatch_kernel(kernel_map) - self.kernel = self.kernel_map["moe_grouped_gemm_kernel"]( - numel, num_experts, n, k, dtype=dtype, tune=tune - ) + self._kernel_cache: Dict[torch.dtype, Kernel] = {} + + def _get_kernel(self, dtype: torch.dtype) -> Kernel: + if dtype not in self._kernel_cache: + self._kernel_cache[dtype] = self.kernel_map["moe_grouped_gemm_kernel"]( + self.numel, self.num_experts, self.n, self.k, + dtype=dtype, tune=self.tune, + ) + return self._kernel_cache[dtype] @property def default_kernel_map(self) -> Dict[str, Kernel]: @@ -79,4 +83,5 @@ def forward( Returns: C: [numel, N] GEMM output. """ - return self.kernel(a, b, true_sizes, true_offsets) + self.dtype = a.dtype + return self._get_kernel(a.dtype)(a, b, true_sizes, true_offsets) diff --git a/tileops/ops/moe/routed_expert/moe_grouped_gemm_nopad_fused_act.py b/tileops/ops/moe/routed_expert/moe_grouped_gemm_nopad_fused_act.py index ca4124200..b93d6c51b 100644 --- a/tileops/ops/moe/routed_expert/moe_grouped_gemm_nopad_fused_act.py +++ b/tileops/ops/moe/routed_expert/moe_grouped_gemm_nopad_fused_act.py @@ -24,7 +24,6 @@ class MoeGroupedGemmNopad3WGFusedActFwdOp(Op): num_experts: Number of local experts E. ffn: FFN width (output column count); B has 2*ffn rows (gate||up). k: Hidden size K. - dtype: Activation/weight dtype (bf16 or fp16). activation: "silu_and_mul" or "gelu_and_mul". kernel_map: Optional kernel override. tune: Autotune flag. @@ -32,7 +31,7 @@ class MoeGroupedGemmNopad3WGFusedActFwdOp(Op): Example: >>> op = MoeGroupedGemmNopad3WGFusedActFwdOp( ... numel=16384, num_experts=256, ffn=768, k=2048, - ... dtype=torch.bfloat16, activation="silu_and_mul") + ... activation="silu_and_mul") >>> C = op(A, B, true_sizes, true_offsets) # [numel, ffn] """ @@ -42,7 +41,6 @@ def __init__( num_experts: int, ffn: int, k: int, - dtype: torch.dtype = torch.bfloat16, activation: str = "silu_and_mul", kernel_map: Optional[Dict[str, Kernel]] = None, tune: bool = False, @@ -51,11 +49,20 @@ def __init__( self.num_experts = num_experts self.ffn = ffn self.k = k - self.dtype = dtype self.activation = activation + self.tune = tune self.dispatch_kernel(kernel_map) - self.kernel = self.kernel_map["moe_grouped_gemm_fused_act_kernel"]( - numel, num_experts, ffn, k, dtype=dtype, activation=activation, tune=tune) + self._kernel_cache: Dict[torch.dtype, Kernel] = {} + + def _get_kernel(self, dtype: torch.dtype) -> Kernel: + if dtype not in self._kernel_cache: + self._kernel_cache[dtype] = self.kernel_map[ + "moe_grouped_gemm_fused_act_kernel" + ]( + self.numel, self.num_experts, self.ffn, self.k, + dtype=dtype, activation=self.activation, tune=self.tune, + ) + return self._kernel_cache[dtype] @property def default_kernel_map(self) -> Dict[str, Kernel]: @@ -79,4 +86,5 @@ def forward( Returns: C: [numel, ffn] activated gate_up output. """ - return self.kernel(a, b, true_sizes, true_offsets) + self.dtype = a.dtype + return self._get_kernel(a.dtype)(a, b, true_sizes, true_offsets) diff --git a/tileops/ops/moe/routed_expert/permute_nopad.py b/tileops/ops/moe/routed_expert/permute_nopad.py index 0fe4cd71f..fbdfcfc3b 100644 --- a/tileops/ops/moe/routed_expert/permute_nopad.py +++ b/tileops/ops/moe/routed_expert/permute_nopad.py @@ -27,8 +27,6 @@ class MoePermuteNopadFwdOp(Op): num_experts: Total number of experts E (global count). hidden_size: Optional committed hidden dimension H. Preferred API infers it from ``hidden_states.shape[1]``. - dtype: Data type of hidden_states (bf16 or fp16). If omitted, - inferred from ``hidden_states``. expert_map: Optional [E_global] int32 tensor mapping global expert ids to local ids (-1 = not on this rank). When provided, only local token-expert pairs are counted; non-local positions get fwd_idx = -1. @@ -45,7 +43,6 @@ def __init__( top_k: Optional[int] = None, num_experts: Optional[int] = None, hidden_size: Optional[int] = None, - dtype: Optional[torch.dtype] = None, expert_map: Optional[torch.Tensor] = None, kernel_map: Optional[Dict[str, Kernel]] = None, ) -> None: @@ -53,11 +50,9 @@ def __init__( self.top_k = top_k self.num_experts = num_experts self.hidden_size = hidden_size - self.dtype = dtype self._committed_total_tokens = total_tokens self._committed_top_k = top_k self._committed_hidden_size = hidden_size - self._committed_dtype = dtype self.dispatch_kernel(kernel_map) self.expert_map = expert_map @@ -167,13 +162,6 @@ def forward( ) if self.num_experts is None: raise ValueError("num_experts must be provided at construction time") - if ( - self._committed_dtype is not None - and hidden_states.dtype != self._committed_dtype - ): - raise ValueError( - f"Expected hidden_states.dtype {self._committed_dtype}, got {hidden_states.dtype}" - ) if hidden_states.dtype not in (torch.float16, torch.bfloat16): raise ValueError( "Expected hidden_states.dtype to be torch.float16 or " @@ -183,10 +171,10 @@ def forward( raise ValueError(f"Expected topk_ids.dtype torch.int32, got {topk_ids.dtype}") dtype = hidden_states.dtype + self.dtype = dtype self.total_tokens = total_tokens self.top_k = top_k self.hidden_size = hidden_size - self.dtype = dtype self.hidden_states_shape = tuple(hidden_states.shape) self.topk_ids_shape = tuple(topk_ids.shape) kernel = self._get_kernel( diff --git a/tileops/ops/moe/routed_expert/unpermute.py b/tileops/ops/moe/routed_expert/unpermute.py index d8ffa4b66..011860e6e 100644 --- a/tileops/ops/moe/routed_expert/unpermute.py +++ b/tileops/ops/moe/routed_expert/unpermute.py @@ -40,7 +40,6 @@ def __init__( total_tokens: int, top_k: int, hidden_size: int, - dtype: torch.dtype = torch.bfloat16, padded_batch_sum: Optional[int] = None, kernel_map: Optional[Dict[str, Kernel]] = None, routed_scaling_factor: float = 1.0, @@ -48,14 +47,20 @@ def __init__( self.total_tokens = total_tokens self.top_k = top_k self.hidden_size = hidden_size - self.dtype = dtype self.padded_batch_sum = padded_batch_sum if padded_batch_sum is not None else total_tokens * top_k + self._routed_scaling_factor = routed_scaling_factor self.dispatch_kernel(kernel_map) - self.kernel = self.kernel_map["unpermute_kernel"]( - total_tokens, top_k, hidden_size, self.padded_batch_sum, - scaling=routed_scaling_factor, dtype=dtype, - ) + self._kernel_cache: Dict[torch.dtype, Kernel] = {} + + def _get_kernel(self, dtype: torch.dtype) -> Kernel: + if dtype not in self._kernel_cache: + self._kernel_cache[dtype] = self.kernel_map["unpermute_kernel"]( + self.total_tokens, self.top_k, self.hidden_size, + self.padded_batch_sum, scaling=self._routed_scaling_factor, + dtype=dtype, + ) + return self._kernel_cache[dtype] @property def default_kernel_map(self) -> Dict[str, Kernel]: @@ -80,4 +85,5 @@ def forward( Returns: output: [T, H] bf16/fp16 (``out`` if provided). """ - return self.kernel(mm2_pad, fwd_idx, topk_weights, out=out) + self.dtype = mm2_pad.dtype + return self._get_kernel(mm2_pad.dtype)(mm2_pad, fwd_idx, topk_weights, out=out) diff --git a/tileops/ops/moe/shared_fused_moe.py b/tileops/ops/moe/shared_fused_moe.py index ccc4bed96..6847b095e 100644 --- a/tileops/ops/moe/shared_fused_moe.py +++ b/tileops/ops/moe/shared_fused_moe.py @@ -80,7 +80,6 @@ def __init__( renormalize: bool = False, with_correction_bias: bool = False, routed_scaling_factor: float = 1.0, - dtype: torch.dtype = torch.bfloat16, expert_map: Optional[torch.Tensor] = None, shared_ffn_size: Optional[int] = None, tp_size: int = 1, @@ -112,7 +111,6 @@ def __init__( renormalize=renormalize, with_correction_bias=with_correction_bias, routed_scaling_factor=routed_scaling_factor, - dtype=dtype, expert_map=expert_map, activation=activation, use_fused_activation=use_fused_activation, @@ -127,21 +125,29 @@ def __init__( f"shared_ffn_size ({shared_ffn_size}) must be divisible by tp_size ({tp_size})" ) + self.num_tokens = num_tokens + self.hidden_size = hidden_size self.shared_ffn_size = shared_ffn_size self.tp_size = tp_size self.tp_rank = tp_rank # Kernel operates on the local shard size - self._shared_mlp_kernel = ( - SharedExpertMLPKernel( - num_tokens=num_tokens, - hidden_size=hidden_size, - ffn_size=shared_ffn_size // tp_size, + self._has_shared_mlp = shared_ffn_size is not None + self._shared_mlp_shard_ffn = ( + shared_ffn_size // tp_size if shared_ffn_size is not None else None + ) + self._shared_mlp_cache: Dict[torch.dtype, Kernel] = {} + + def _shared_mlp_kernel_for(self, dtype: torch.dtype) -> Kernel: + """Return the shared-expert MLP kernel for *dtype*, building on first use.""" + if dtype not in self._shared_mlp_cache: + self._shared_mlp_cache[dtype] = SharedExpertMLPKernel( + num_tokens=self.num_tokens, + hidden_size=self.hidden_size, + ffn_size=self._shared_mlp_shard_ffn, dtype=dtype, ) - if shared_ffn_size is not None - else None - ) + return self._shared_mlp_cache[dtype] @property def default_kernel_map(self) -> Dict[str, Kernel]: @@ -178,7 +184,7 @@ def forward( shared_output is a partial sum when tp_size > 1; caller must all-reduce across TP ranks. """ - if self._shared_mlp_kernel is not None: + if self._has_shared_mlp: if shared_w_gate_up is None or shared_w_down is None: raise ValueError( "shared_w_gate_up and shared_w_down must be provided " @@ -217,7 +223,9 @@ def forward( gate_up_shard = shared_w_gate_up down_shard = shared_w_down - shared_out = self._shared_mlp_kernel(hidden_states, gate_up_shard, down_shard) + shared_out = self._shared_mlp_kernel_for(hidden_states.dtype)( + hidden_states, gate_up_shard, down_shard, + ) else: shared_out = None From 9108aaa2dc2645efd90e41cb431a9284e5a80c07 Mon Sep 17 00:00:00 2001 From: lcy-seso Date: Tue, 4 Aug 2026 17:34:06 +0800 Subject: [PATCH 06/15] [Refactor][Ops] Read dtype at forward for the top-level ops MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `CBProducerOp`, `EngramGateConvFwdOp`, `EngramGateConvBwdOp`, `EngramDecodeOp` and `MeanPoolingForwardOp` build their kernel on first forward, cached per dtype. `DaCumsumFwdOp` keeps its `dtype` argument: its inputs are float32 while `dt_out` is float16/bfloat16/float32, so the output element type is a choice its inputs cannot express. The engram ops route their existing dtype check through the synthesized `_validate_dtypes`, which needs every declared input rather than just the anchor tensor. `CBProducerOp` keeps an inline check — its `_validate_dtypes` is still the L1 stub — and now compares its two inputs against each other instead of against a constructor argument. `MeanPoolingForwardOp` collects its kernel arguments through `locals()`; that dict becomes every argument except the element type, with dtype supplied at the cache lookup. Four call sites passed dtype positionally, where the constructor took it before `eps`. Removing the parameter would have bound a dtype to `eps` and carried it into the kernel — the failure surfaced as a TileLang type error about argument 3, far from the cause. --- benchmarks/ops/bench_engram.py | 6 ++-- benchmarks/ops/bench_mean_pooling.py | 1 - tests/ops/test_engram.py | 8 +++--- tests/ops/test_mamba.py | 4 +-- tests/ops/test_mean_pooling.py | 1 - tileops/ops/cb_producer.py | 22 ++++++++------ tileops/ops/engram.py | 43 +++++++++++++++++----------- tileops/ops/engram_decode.py | 22 ++++++++------ tileops/ops/mamba2_fwd.py | 1 - tileops/ops/pool.py | 16 +++++++++-- 10 files changed, 76 insertions(+), 48 deletions(-) diff --git a/benchmarks/ops/bench_engram.py b/benchmarks/ops/bench_engram.py index 27b85b2b2..5093559dd 100644 --- a/benchmarks/ops/bench_engram.py +++ b/benchmarks/ops/bench_engram.py @@ -82,7 +82,7 @@ def test_engram_gate_conv_fwd_bench(M, seq_len, d, dtype): test = EngramGateConvFwdWorkload(M, seq_len, d, dtype) inputs = test.gen_inputs() - op = EngramGateConvFwdOp(M, seq_len, d, dtype, tune=_TUNE) + op = EngramGateConvFwdOp(M, seq_len, d, tune=_TUNE) bm = ManifestBenchmark(_ENGRAM_GATE_CONV_FWD_OP, op, test) result = bm.profile(op, *inputs) BenchmarkReport.record(op, locals(), result, tag="tileops") @@ -156,7 +156,7 @@ def test_engram_gate_conv_bwd_bench(M, seq_len, d, dtype): test = EngramGateConvBwdTestBaseline(M, seq_len, d, dtype) inputs = test.gen_inputs() - op = EngramGateConvBwdOp(M, seq_len, d, dtype, tune=_TUNE) + op = EngramGateConvBwdOp(M, seq_len, d, tune=_TUNE) bm = ManifestBenchmark(_ENGRAM_GATE_CONV_BWD_OP, op, test) result = bm.profile(op, *inputs) BenchmarkReport.record(op, locals(), result, tag="tileops") @@ -242,7 +242,7 @@ def test_engram_decode_bench(batch, d_mem, d, max_conv_len, conv_kernel_size, di inputs = test.gen_inputs() op = EngramDecodeOp( - batch, d_mem, d, max_conv_len, conv_kernel_size, dilation, dtype, tune=_TUNE, + batch, d_mem, d, max_conv_len, conv_kernel_size, dilation, tune=_TUNE, ) bm = ManifestBenchmark(_ENGRAM_DECODE_OP, op, test) result = bm.profile(op, *inputs) diff --git a/benchmarks/ops/bench_mean_pooling.py b/benchmarks/ops/bench_mean_pooling.py index f6b83cb5d..38be53e77 100644 --- a/benchmarks/ops/bench_mean_pooling.py +++ b/benchmarks/ops/bench_mean_pooling.py @@ -111,7 +111,6 @@ def test_mean_pooling_bench(batch_size: int, seq_len: int, heads: int, dim: int, "chunks_per_batch": chunks_per_batch, "seq_num": seq_num, "use_offsets": use_offsets, - "dtype": dtype, "accum_dtype": accum_dtype, "tune": tune, } diff --git a/tests/ops/test_engram.py b/tests/ops/test_engram.py index 03b50c6d9..2328e2fba 100644 --- a/tests/ops/test_engram.py +++ b/tests/ops/test_engram.py @@ -71,7 +71,7 @@ class EngramGateConvFwdFixture(FixtureBase): @EngramGateConvFwdFixture def test_engram_gate_conv_fwd(M, seq_len, d, dtype, tune): test = EngramGateConvFwdTest(M, seq_len, d, dtype) - op = EngramGateConvFwdOp(M, seq_len, d, dtype, tune=tune) + op = EngramGateConvFwdOp(M, seq_len, d, tune=tune) inputs = test.gen_inputs() atol = 1e-1 if dtype == torch.float16 else 2e-1 rtol = 1e-1 @@ -150,7 +150,7 @@ class EngramGateConvBwdFixture(FixtureBase): @EngramGateConvBwdFixture def test_engram_gate_conv_bwd(M, seq_len, d, dtype, tune): test = EngramGateConvBwdTest(M, seq_len, d, dtype) - op = EngramGateConvBwdOp(M, seq_len, d, dtype, tune=tune) + op = EngramGateConvBwdOp(M, seq_len, d, tune=tune) inputs = test.gen_inputs() atol = 2e-1 if dtype == torch.float16 else 3e-1 rtol = 2e-1 @@ -247,7 +247,7 @@ class EngramDecodeFixture(FixtureBase): def test_engram_decode(batch, d_mem, d, max_conv_len, conv_kernel_size, dilation, dtype, tune): test = EngramDecodeTest(batch, d_mem, d, max_conv_len, conv_kernel_size, dilation, dtype) op = EngramDecodeOp( - batch, d_mem, d, max_conv_len, conv_kernel_size, dilation, dtype, tune=tune, + batch, d_mem, d, max_conv_len, conv_kernel_size, dilation, tune=tune, ) inputs = test.gen_inputs() atol = 5e-2 if dtype == torch.float16 else 1e-1 @@ -272,7 +272,7 @@ def test_engram_decode_multi_step(): rms_w_v = torch.ones(d, dtype=dtype, device="cuda") conv_w = torch.randn(conv_kernel_size, d, dtype=dtype, device="cuda") * 0.02 - op = EngramDecodeOp(B, d_mem, d, max_conv_len, conv_kernel_size, dilation, dtype) + op = EngramDecodeOp(B, d_mem, d, max_conv_len, conv_kernel_size, dilation) # Start with empty conv_state (like empty KV cache) conv_state = torch.zeros(B, 0, d, dtype=dtype, device="cuda") diff --git a/tests/ops/test_mamba.py b/tests/ops/test_mamba.py index 48f99ed41..b72a8d7db 100644 --- a/tests/ops/test_mamba.py +++ b/tests/ops/test_mamba.py @@ -101,7 +101,7 @@ def cb_producer_fwd_ref( pytest.param(2, 4, 64, 4, 128, torch.bfloat16, False, marks=pytest.mark.full), ]) def test_cb_producer_fwd(batch, num_chunks, chunk_len, n_groups, d_state, dtype, tune): - op = CBProducerOp(batch, num_chunks, n_groups, chunk_len, d_state, dtype=dtype, tune=tune) + op = CBProducerOp(batch, num_chunks, n_groups, chunk_len, d_state, tune=tune) seq_len = num_chunks * chunk_len C_mat = torch.randn(batch, seq_len, n_groups, d_state, dtype=dtype, device="cuda") * 0.1 B_mat = torch.randn(batch, seq_len, n_groups, d_state, dtype=dtype, device="cuda") * 0.1 @@ -123,7 +123,7 @@ def test_cb_producer_fwd_noncontiguous(): assert not C_mat.is_contiguous() assert not B_mat.is_contiguous() ref = cb_producer_fwd_ref(C_mat.contiguous(), B_mat.contiguous(), num_chunks, chunk_len, dtype) - out = CBProducerOp(batch, num_chunks, n_groups, chunk_len, d_state, dtype=dtype)(C_mat, B_mat) + out = CBProducerOp(batch, num_chunks, n_groups, chunk_len, d_state)(C_mat, B_mat) allclose_compare(out, ref, atol=1e-3, rtol=1e-3) diff --git a/tests/ops/test_mean_pooling.py b/tests/ops/test_mean_pooling.py index ba9d115e4..be33ef637 100644 --- a/tests/ops/test_mean_pooling.py +++ b/tests/ops/test_mean_pooling.py @@ -112,7 +112,6 @@ def test_mean_pooling_op(batch_size: int, seq_len: int, heads: int, dim: int, ch "chunks_per_batch": chunks_per_batch, "seq_num": seq_num, "use_offsets": use_offsets, - "dtype": dtype, "accum_dtype": accum_dtype, "tune": tune, } diff --git a/tileops/ops/cb_producer.py b/tileops/ops/cb_producer.py index 552d3e882..9771ed3e1 100644 --- a/tileops/ops/cb_producer.py +++ b/tileops/ops/cb_producer.py @@ -37,7 +37,6 @@ def __init__( n_groups: int, chunk_len: int, d_state: int, - dtype: torch.dtype = torch.float16, tune: bool = False, kernel_map: Optional[Dict[str, Kernel]] = None, ): @@ -46,13 +45,19 @@ def __init__( self.n_groups = n_groups self.chunk_len = chunk_len self.d_state = d_state - self.dtype = dtype + self.tune = tune # Use standard Op dispatch pattern self.dispatch_kernel(kernel_map) - self.kernel = self.kernel_map["cb_producer"]( - batch, num_chunks, n_groups, chunk_len, d_state, dtype, tune=tune - ) + self._kernel_cache: Dict[torch.dtype, Kernel] = {} + + def _get_kernel(self, dtype: torch.dtype) -> Kernel: + if dtype not in self._kernel_cache: + self._kernel_cache[dtype] = self.kernel_map["cb_producer"]( + self.batch, self.num_chunks, self.n_groups, self.chunk_len, + self.d_state, dtype, tune=self.tune, + ) + return self._kernel_cache[dtype] @property def default_kernel_map(self) -> Dict[str, Kernel]: @@ -76,10 +81,11 @@ def forward( """ S = self.num_chunks * self.chunk_len expected_shape = (self.batch, S, self.n_groups, self.d_state) + self.dtype = C_mat.dtype for name, t in (("C_mat", C_mat), ("B_mat", B_mat)): - if t.dtype != self.dtype: + if t.dtype != C_mat.dtype: raise ValueError( - f"{name}.dtype={t.dtype} does not match op dtype={self.dtype}" + f"{name}.dtype={t.dtype} does not match C_mat.dtype={C_mat.dtype}" ) if t.shape != torch.Size(expected_shape): raise ValueError( @@ -87,4 +93,4 @@ def forward( ) C_mat = C_mat.contiguous() B_mat = B_mat.contiguous() - return self.kernel(C_mat, B_mat) + return self._get_kernel(C_mat.dtype)(C_mat, B_mat) diff --git a/tileops/ops/engram.py b/tileops/ops/engram.py index 7a4c253a3..a8bf525c7 100644 --- a/tileops/ops/engram.py +++ b/tileops/ops/engram.py @@ -36,7 +36,6 @@ def __init__( M: int, seq_len: int, d: int, - dtype: torch.dtype, eps: float = 1e-6, tune: bool = False, kernel_map: Optional[Dict[str, Kernel]] = None, @@ -44,12 +43,17 @@ def __init__( self.M = M self.seq_len = seq_len self.d = d - self.dtype = dtype self.eps = eps + self.tune = tune self.dispatch_kernel(kernel_map) - self.kernel = self.kernel_map["engram_gate_conv_fwd"]( - M, seq_len, d, eps, dtype, tune=tune, - ) + self._kernel_cache: Dict[torch.dtype, Kernel] = {} + + def _get_kernel(self, dtype: torch.dtype) -> Kernel: + if dtype not in self._kernel_cache: + self._kernel_cache[dtype] = self.kernel_map["engram_gate_conv_fwd"]( + self.M, self.seq_len, self.d, self.eps, dtype, tune=self.tune, + ) + return self._kernel_cache[dtype] @property def default_kernel_map(self) -> Dict[str, Kernel]: @@ -84,8 +88,8 @@ def forward( """ if not H.is_cuda: raise ValueError("H must be a CUDA tensor") - if H.dtype != self.dtype: - raise ValueError(f"Expected dtype {self.dtype}, got {H.dtype}") + self._validate_dtypes(H, k, v, rms_w_h, rms_w_v, conv_w) + self.dtype = H.dtype if H.shape[-1] != self.d: raise ValueError( f"Expected hidden dim {self.d}, got {H.shape[-1]}" @@ -99,7 +103,7 @@ def forward( k = k.contiguous() v = v.contiguous() - return self.kernel(H, k, v, rms_w_h, rms_w_v, conv_w) + return self._get_kernel(H.dtype)(H, k, v, rms_w_h, rms_w_v, conv_w) class EngramGateConvBwdOp(Op): @@ -126,7 +130,6 @@ def __init__( M: int, seq_len: int, d: int, - dtype: torch.dtype, eps: float = 1e-6, tune: bool = False, kernel_map: Optional[Dict[str, Kernel]] = None, @@ -134,12 +137,17 @@ def __init__( self.M = M self.seq_len = seq_len self.d = d - self.dtype = dtype self.eps = eps + self.tune = tune self.dispatch_kernel(kernel_map) - self.kernel = self.kernel_map["engram_gate_conv_bwd"]( - M, seq_len, d, eps, dtype, tune=tune, - ) + self._kernel_cache: Dict[torch.dtype, Kernel] = {} + + def _get_kernel(self, dtype: torch.dtype) -> Kernel: + if dtype not in self._kernel_cache: + self._kernel_cache[dtype] = self.kernel_map["engram_gate_conv_bwd"]( + self.M, self.seq_len, self.d, self.eps, dtype, tune=self.tune, + ) + return self._kernel_cache[dtype] @property def default_kernel_map(self) -> Dict[str, Kernel]: @@ -186,8 +194,11 @@ def forward( """ if not dY.is_cuda: raise ValueError("dY must be a CUDA tensor") - if dY.dtype != self.dtype: - raise ValueError(f"Expected dtype {self.dtype}, got {dY.dtype}") + self._validate_dtypes( + dY, H, k, v, rms_w_h, rms_w_v, conv_w, + vhat, alpha, rrms_h, rrms_k, rrms_v, + ) + self.dtype = dY.dtype dY = dY.contiguous() H = H.contiguous() @@ -195,7 +206,7 @@ def forward( v = v.contiguous() vhat = vhat.contiguous() - return self.kernel( + return self._get_kernel(dY.dtype)( dY, H, k, v, rms_w_h, rms_w_v, conv_w, vhat, alpha, rrms_h, rrms_k, rrms_v, ) diff --git a/tileops/ops/engram_decode.py b/tileops/ops/engram_decode.py index 155ae7518..8b3569e21 100644 --- a/tileops/ops/engram_decode.py +++ b/tileops/ops/engram_decode.py @@ -38,7 +38,6 @@ def __init__( max_conv_len: int, conv_kernel_size: int, dilation: int, - dtype: torch.dtype, eps: float = 1e-6, tune: bool = False, kernel_map: Optional[Dict[str, Kernel]] = None, @@ -49,13 +48,19 @@ def __init__( self.max_conv_len = max_conv_len self.conv_kernel_size = conv_kernel_size self.dilation = dilation - self.dtype = dtype self.eps = eps + self.tune = tune self.dispatch_kernel(kernel_map) - self.kernel = self.kernel_map["engram_decode"]( - batch, d_mem, d, max_conv_len, conv_kernel_size, dilation, - eps, dtype, tune=tune, - ) + self._kernel_cache: Dict[torch.dtype, Kernel] = {} + + def _get_kernel(self, dtype: torch.dtype) -> Kernel: + if dtype not in self._kernel_cache: + self._kernel_cache[dtype] = self.kernel_map["engram_decode"]( + self.batch, self.d_mem, self.d, self.max_conv_len, + self.conv_kernel_size, self.dilation, self.eps, dtype, + tune=self.tune, + ) + return self._kernel_cache[dtype] @property def default_kernel_map(self) -> Dict[str, Kernel]: @@ -91,13 +96,12 @@ def forward( """ if not e_t.is_cuda: raise ValueError("e_t must be a CUDA tensor") - if e_t.dtype != self.dtype: - raise ValueError(f"Expected dtype {self.dtype}, got {e_t.dtype}") + self.dtype = e_t.dtype e_t = e_t.contiguous() h_t = h_t.contiguous() conv_state = conv_state.contiguous() - return self.kernel( + return self._get_kernel(e_t.dtype)( e_t, h_t, conv_state, W_K, W_V, rms_w_h, rms_w_v, conv_w, ) diff --git a/tileops/ops/mamba2_fwd.py b/tileops/ops/mamba2_fwd.py index 4f5121a3d..0b616cc0b 100644 --- a/tileops/ops/mamba2_fwd.py +++ b/tileops/ops/mamba2_fwd.py @@ -131,7 +131,6 @@ def _get_cb_producer_op( n_groups=n_groups, chunk_len=self.chunk_size, d_state=d_state, - dtype=dtype, tune=self.tune, ) return self._cb_producer_ops[key] diff --git a/tileops/ops/pool.py b/tileops/ops/pool.py index fd9d867b0..48413fe15 100644 --- a/tileops/ops/pool.py +++ b/tileops/ops/pool.py @@ -59,7 +59,6 @@ def __init__( chunks_per_batch: int, seq_num: int, use_offsets: int, - dtype: torch.dtype, accum_dtype: torch.dtype, tune: bool = False, kernel_map: Optional[Dict[str, Kernel]] = None, @@ -68,8 +67,18 @@ def __init__( for key, value in params.items(): setattr(self, key, value) + # Every kernel argument except the element type, which the first + # forward() supplies from its input. + self._kernel_params = params self.dispatch_kernel(kernel_map) - self.kernel = self.kernel_map["mean_pooling_fwd_kernel"](**params) + self._kernel_cache: Dict[torch.dtype, Kernel] = {} + + def _get_kernel(self, dtype: torch.dtype) -> Kernel: + if dtype not in self._kernel_cache: + self._kernel_cache[dtype] = self.kernel_map["mean_pooling_fwd_kernel"]( + **self._kernel_params, dtype=dtype, + ) + return self._kernel_cache[dtype] @property def default_kernel_map(self) -> Dict[str, Kernel]: @@ -81,7 +90,8 @@ def forward( offsets: torch.Tensor, indices: torch.Tensor, ) -> torch.Tensor: - return self.kernel(x, offsets, indices=indices) + self.dtype = x.dtype + return self._get_kernel(x.dtype)(x, offsets, indices=indices) def _device_index(tensor: torch.Tensor) -> int | None: From 68e0c48571a1ac1d825d5b083f3ecff07467e8bc Mon Sep 17 00:00:00 2001 From: lcy-seso Date: Tue, 4 Aug 2026 20:18:58 +0800 Subject: [PATCH 07/15] [Refactor][Attention] Read dtype at forward for the backward and decode ops `GroupedQueryAttentionBwdOp`, `MultiHeadAttentionBwdOp` and the two MHA decode ops build their kernels on first forward, cached per dtype. The GQA backward op caches the preprocess and backward kernels as one entry so a dtype switch cannot pair one dtype's preprocess kernel with another's backward kernel. `MultiHeadAttentionBwdOp` stops re-exporting `prep_kernel` / `kernel` from the GQA op it wraps: those attributes no longer exist once the kernels live in a cache, and nothing outside the op read them. Seven call sites passed dtype as the last positional argument; they are handled by position rather than by keyword, which a name-based sweep would have missed. Deferred to a follow-up commit: the four ops whose `default_kernel_map` selects a kernel class from `self.dtype` (GQA forward, GQA prefill, GQA prefill-paged, MHA forward). Those need the candidate set installed as separate map keys with the choice made at forward, which also renames manifest `source.kernel_map` entries. --- benchmarks/ops/attention/bench_gqa.py | 2 +- benchmarks/ops/attention/bench_mha.py | 2 +- benchmarks/ops/attention/bench_mha_decode.py | 2 +- .../ops/attention/bench_mha_decode_paged.py | 2 +- tests/ops/attention/test_gqa.py | 2 +- tests/ops/attention/test_mha.py | 3 +- tests/ops/attention/test_mha_decode.py | 2 +- tests/ops/attention/test_mha_decode_paged.py | 1 - tileops/ops/attention/gqa.py | 31 +++++++++---- tileops/ops/attention/mha.py | 44 ++++++++++--------- 10 files changed, 54 insertions(+), 37 deletions(-) diff --git a/benchmarks/ops/attention/bench_gqa.py b/benchmarks/ops/attention/bench_gqa.py index 47e1e48f1..bec0ed791 100644 --- a/benchmarks/ops/attention/bench_gqa.py +++ b/benchmarks/ops/attention/bench_gqa.py @@ -320,7 +320,7 @@ def test_gqa_bwd_bench( test = GroupedQueryAttentionBwdWorkload(batch, heads, heads_kv, seq_len, dim, causal, dtype) inputs = test.gen_inputs() - op = GroupedQueryAttentionBwdOp(batch, heads, heads_kv, seq_len, dim, causal, dtype, tune=tune) + op = GroupedQueryAttentionBwdOp(batch, heads, heads_kv, seq_len, dim, causal, tune=tune) bm = ManifestBenchmark(_GQA_BWD_OP, op, test) result = bm.profile(op, *inputs) BenchmarkReport.record(op, locals(), result, tag="tileops") diff --git a/benchmarks/ops/attention/bench_mha.py b/benchmarks/ops/attention/bench_mha.py index 7a6b1125a..df0c1448e 100644 --- a/benchmarks/ops/attention/bench_mha.py +++ b/benchmarks/ops/attention/bench_mha.py @@ -139,7 +139,7 @@ def test_mha_bwd_bench(batch: int, seq_len: int, heads: int, dim: int, causal: b test = MhaBwdWorkload(batch, heads, seq_len, dim, causal, dtype) inputs = test.gen_inputs() - op = MultiHeadAttentionBwdOp(batch, heads, seq_len, dim, causal, dtype, tune=tune) + op = MultiHeadAttentionBwdOp(batch, heads, seq_len, dim, causal, tune=tune) bm = ManifestBenchmark(_MHA_BWD_OP, op, test) result = bm.profile(op, *inputs) BenchmarkReport.record(op, locals(), result, tag="tileops") diff --git a/benchmarks/ops/attention/bench_mha_decode.py b/benchmarks/ops/attention/bench_mha_decode.py index 7f203f3c3..611bb9108 100644 --- a/benchmarks/ops/attention/bench_mha_decode.py +++ b/benchmarks/ops/attention/bench_mha_decode.py @@ -76,7 +76,7 @@ def test_mha_decode_bench(b: int, h: int, s_q: int, s_kv: int, d: int, dtype: to test = MhaDecodeTestBaseline(b, h, s_q, s_kv, d, dtype) inputs = test.gen_inputs() - op = MultiHeadAttentionDecodeWithKVCacheFwdOp(b, h, s_q, s_kv, d, dtype, tune=tune) + op = MultiHeadAttentionDecodeWithKVCacheFwdOp(b, h, s_q, s_kv, d, tune=tune) bm = ManifestBenchmark(_OP_NAME, op, test) result = bm.profile(op, *inputs) BenchmarkReport.record(op, locals(), result, tag="tileops") diff --git a/benchmarks/ops/attention/bench_mha_decode_paged.py b/benchmarks/ops/attention/bench_mha_decode_paged.py index 2e6afc80c..8afddb815 100644 --- a/benchmarks/ops/attention/bench_mha_decode_paged.py +++ b/benchmarks/ops/attention/bench_mha_decode_paged.py @@ -133,7 +133,7 @@ def test_mha_decode_paged_bench(batch: int, heads: int, seqlen_q: int, seqlen_kv q, k, v, real_seqlen_kv, block_table = inputs op = MultiHeadAttentionDecodePagedWithKVCacheFwdOp( - batch, heads, seqlen_q, seqlen_kv, dim, page_size, is_causal, dtype, tune=tune) + batch, heads, seqlen_q, seqlen_kv, dim, page_size, is_causal, tune=tune) bm = ManifestBenchmark(_OP_NAME, op, test) result = bm.profile(op, *inputs) BenchmarkReport.record(op, locals(), result, tag="tileops") diff --git a/tests/ops/attention/test_gqa.py b/tests/ops/attention/test_gqa.py index 72927b010..9110c5bd6 100644 --- a/tests/ops/attention/test_gqa.py +++ b/tests/ops/attention/test_gqa.py @@ -850,7 +850,7 @@ def test_gqa_prefill_varlen_rejects_unsupported_dtype() -> None: def test_gqa_bwd(batch: int, seq_len: int, heads: int, heads_kv: int, dim: int, causal: bool, dtype: torch.dtype, tune: bool) -> None: test = GroupedQueryAttentionBwdTest(batch, heads, heads_kv, seq_len, dim, causal, dtype) - op = GroupedQueryAttentionBwdOp(batch, heads, heads_kv, seq_len, dim, causal, dtype, tune=tune) + op = GroupedQueryAttentionBwdOp(batch, heads, heads_kv, seq_len, dim, causal, tune=tune) test.check(op, *test.gen_inputs(), atol=5e-3, rtol=1e-5) diff --git a/tests/ops/attention/test_mha.py b/tests/ops/attention/test_mha.py index fbd61c7a1..6361af0bf 100644 --- a/tests/ops/attention/test_mha.py +++ b/tests/ops/attention/test_mha.py @@ -159,7 +159,6 @@ def test_mha_bwd_rejects_legacy_kernel_map_keys() -> None: seq_len=128, dim=64, is_causal=False, - dtype=torch.float16, kernel_map={"mha_bwd_kernel": _FakeLegacyMhaBwdKernel}, ) @@ -168,7 +167,7 @@ def test_mha_bwd_rejects_legacy_kernel_map_keys() -> None: def test_mha_bwd(batch: int, seq_len: int, heads: int, dim: int, causal: bool, dtype: torch.dtype, tune: bool) -> None: test = MhaBwdTest(batch, heads, seq_len, dim, causal, dtype) - op = MultiHeadAttentionBwdOp(batch, heads, seq_len, dim, causal, dtype, tune=tune) + op = MultiHeadAttentionBwdOp(batch, heads, seq_len, dim, causal, tune=tune) test.check(op, *test.gen_inputs(), atol=5e-3, rtol=1e-5) diff --git a/tests/ops/attention/test_mha_decode.py b/tests/ops/attention/test_mha_decode.py index 9e371338b..cc9fe8c29 100644 --- a/tests/ops/attention/test_mha_decode.py +++ b/tests/ops/attention/test_mha_decode.py @@ -34,7 +34,7 @@ class MhaDecodeFixture(FixtureBase): def test_mha_decode(b: int, h: int, s_q: int, s_kv: int, d: int, dtype: torch.dtype, tune: bool) -> None: test = MhaDecodeTest(b, h, s_q, s_kv, d, dtype) - op = MultiHeadAttentionDecodeWithKVCacheFwdOp(b, h, s_q, s_kv, d, dtype, tune=tune) + op = MultiHeadAttentionDecodeWithKVCacheFwdOp(b, h, s_q, s_kv, d, tune=tune) test.check(op, *test.gen_inputs(), atol=2e-3, rtol=1e-5) diff --git a/tests/ops/attention/test_mha_decode_paged.py b/tests/ops/attention/test_mha_decode_paged.py index 2056a2dd1..859f36df1 100644 --- a/tests/ops/attention/test_mha_decode_paged.py +++ b/tests/ops/attention/test_mha_decode_paged.py @@ -103,7 +103,6 @@ def test_mha_decode_paged_op( dim=dim, page_size=page_size, is_causal=is_causal, - dtype=dtype, tune=tune, ) test.check(op, *test.gen_inputs(), compare=test._maxdiff_cosine_compare) diff --git a/tileops/ops/attention/gqa.py b/tileops/ops/attention/gqa.py index d6870c64d..8691cf6ff 100644 --- a/tileops/ops/attention/gqa.py +++ b/tileops/ops/attention/gqa.py @@ -1332,11 +1332,24 @@ def __init__(self, self.dtype = dtype + self.tune = tune self.dispatch_kernel(kernel_map) - self.prep_kernel = self.kernel_map["gqa_bwd_preprocess_kernel"](batch, heads, seq_len, dim, - self.dtype, tune=tune) - self.kernel = self.kernel_map["gqa_bwd_kernel"]( - batch, heads, heads_kv, seq_len, dim, is_causal, self.dtype, tune=tune) + self._kernel_cache: Dict[torch.dtype, tuple[Kernel, Kernel]] = {} + + def _get_kernels(self, dtype: torch.dtype) -> tuple[Kernel, Kernel]: + """Return (preprocess, backward) kernels for *dtype*, building once.""" + if dtype not in self._kernel_cache: + self._kernel_cache[dtype] = ( + self.kernel_map["gqa_bwd_preprocess_kernel"]( + self.batch, self.heads, self.seq_len, self.dim, dtype, + tune=self.tune, + ), + self.kernel_map["gqa_bwd_kernel"]( + self.batch, self.heads, self.heads_kv, self.seq_len, + self.dim, self.is_causal, dtype, tune=self.tune, + ), + ) + return self._kernel_cache[dtype] @property def default_kernel_map(self) -> Dict[str, Kernel]: @@ -1351,13 +1364,15 @@ def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, o: torch.Te do: torch.Tensor, lse: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: do = do.contiguous() - delta = self.prep_kernel(o, do) + self.dtype = q.dtype + prep_kernel, kernel = self._get_kernels(q.dtype) + delta = prep_kernel(o, do) dq = torch.zeros_like(q, dtype=torch.float32) dk = torch.zeros_like(k, dtype=torch.float32) dv = torch.zeros_like(v, dtype=torch.float32) - self.kernel(q, k, v, do, lse, delta, dq, dk, dv) - dq = dq.to(self.dtype) - dk, dv = dk.to(self.dtype), dv.to(self.dtype) + kernel(q, k, v, do, lse, delta, dq, dk, dv) + dq = dq.to(q.dtype) + dk, dv = dk.to(q.dtype), dv.to(q.dtype) return dq, dk, dv diff --git a/tileops/ops/attention/mha.py b/tileops/ops/attention/mha.py index d5d96989f..16327e358 100644 --- a/tileops/ops/attention/mha.py +++ b/tileops/ops/attention/mha.py @@ -117,7 +117,6 @@ def __init__(self, seq_len: int, dim: int, is_causal: bool = True, - dtype: torch.dtype = torch.float16, kernel_map: Optional[Dict[str, Kernel]] = None, tune: bool = False) -> None: self.batch = batch @@ -126,8 +125,6 @@ def __init__(self, self.dim = dim self.is_causal = is_causal - self.dtype = dtype - self.dispatch_kernel(self._gqa_kernel_map(kernel_map)) self._gqa_op = GroupedQueryAttentionBwdOp( batch=batch, @@ -136,15 +133,10 @@ def __init__(self, seq_len=seq_len, dim=dim, is_causal=is_causal, - dtype=dtype, kernel_map=self.kernel_map, tune=tune, ) self.kernel_map = self._gqa_op.kernel_map - self.prep_kernel = self._gqa_op.prep_kernel - self.kernel = self._gqa_op.kernel - if hasattr(self._gqa_op, "post_kernel"): - self.post_kernel = self._gqa_op.post_kernel @property def default_kernel_map(self) -> Dict[str, Kernel]: @@ -183,7 +175,6 @@ def __init__(self, seqlen_q: int, seqlen_kv: int, dim: int, - dtype: torch.dtype = torch.float16, kernel_map: Optional[Dict[str, Kernel]] = None, tune: bool = False) -> None: self.batch = batch @@ -192,11 +183,17 @@ def __init__(self, self.seqlen_kv = seqlen_kv self.dim = dim - self.dtype = dtype - + self.tune = tune self.dispatch_kernel(kernel_map) - self.kernel = self.kernel_map["mha_decode_kernel"]( - batch, heads, seqlen_q, seqlen_kv, dim, False, self.dtype, tune=tune) + self._kernel_cache: Dict[torch.dtype, Kernel] = {} + + def _get_kernel(self, dtype: torch.dtype) -> Kernel: + if dtype not in self._kernel_cache: + self._kernel_cache[dtype] = self.kernel_map["mha_decode_kernel"]( + self.batch, self.heads, self.seqlen_q, self.seqlen_kv, + self.dim, False, dtype, tune=self.tune, + ) + return self._kernel_cache[dtype] @property def default_kernel_map(self) -> Dict[str, Kernel]: @@ -209,7 +206,8 @@ def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Te k, pad=(0, 0, 0, 0, 0, self.seqlen_kv - real_seqlen_kv), mode='constant', value=0) v = F.pad( v, pad=(0, 0, 0, 0, 0, self.seqlen_kv - real_seqlen_kv), mode='constant', value=0) - return self.kernel(q, k, v, real_seqlen_kv) + self.dtype = q.dtype + return self._get_kernel(q.dtype)(q, k, v, real_seqlen_kv) class MultiHeadAttentionDecodePagedWithKVCacheFwdOp(Op): @@ -225,7 +223,6 @@ def __init__(self, dim: int, page_size: int, is_causal: bool = False, - dtype: torch.dtype = torch.float16, kernel_map: Optional[Dict[str, Kernel]] = None, tune: bool = False) -> None: self.batch = batch @@ -235,11 +232,17 @@ def __init__(self, self.dim = dim self.page_size = page_size self.is_causal = is_causal - self.dtype = dtype - + self.tune = tune self.dispatch_kernel(kernel_map) - self.kernel = self.kernel_map["mha_decode_paged_kernel"]( - batch, heads, seqlen_q, seqlen_kv, dim, page_size, is_causal, self.dtype, tune=tune) + self._kernel_cache: Dict[torch.dtype, Kernel] = {} + + def _get_kernel(self, dtype: torch.dtype) -> Kernel: + if dtype not in self._kernel_cache: + self._kernel_cache[dtype] = self.kernel_map["mha_decode_paged_kernel"]( + self.batch, self.heads, self.seqlen_q, self.seqlen_kv, + self.dim, self.page_size, self.is_causal, dtype, tune=self.tune, + ) + return self._kernel_cache[dtype] @property def default_kernel_map(self) -> Dict[str, Kernel]: @@ -247,4 +250,5 @@ def default_kernel_map(self) -> Dict[str, Kernel]: def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, real_seqlen_kv: torch.Tensor, block_table: torch.Tensor) -> torch.Tensor: - return self.kernel(q, k, v, real_seqlen_kv, block_table) + self.dtype = q.dtype + return self._get_kernel(q.dtype)(q, k, v, real_seqlen_kv, block_table) From 61222aa1a9b0436fb5bf463a225cc06539ef3932 Mon Sep 17 00:00:00 2001 From: lcy-seso Date: Tue, 4 Aug 2026 20:34:21 +0800 Subject: [PATCH 08/15] [Refactor][Attention] Read dtype at forward for the varlen, sliding-window and DeepSeek ops Eight more ops build their kernel on first forward, cached per dtype: the three NSA ops, the DeepSeek sparse-MLA and MLA decode ops, GQA prefill varlen, and both GQA sliding-window ops. `GroupedQueryAttentionSlidingWindowFwdOp.total_memory` computed `torch.tensor([], dtype=self.dtype).element_size()`. With no constructor dtype, `dtype=None` makes that float32, so it silently returned twice the real byte count. It now refuses to answer before the first forward rather than guessing. Three tests changed because the coupling they asserted is gone, not the invariant underneath: - an unsupported element type is still rejected, but the check moved with the dtype to first use, so the test feeds float32 tensors instead of naming float32 at construction; - `total_memory` is asserted after a forward, and before one the test now requires the refusal; - the dtype disagreement that still matters is between q, k and v, so the mismatch test makes k differ from q rather than making all three differ from a constructor argument. Call sites needed all four forms: keyword, last-positional, an explicit `params` dict inside the op, and a `**params` dict at the call. --- .../attention/bench_deepseek_dsa_decode.py | 2 +- .../attention/bench_deepseek_mla_decode.py | 2 +- benchmarks/ops/attention/bench_gqa.py | 2 +- .../ops/attention/bench_gqa_sliding_window.py | 1 - .../bench_gqa_sliding_window_varlen.py | 1 - .../ops/attention/test_deepseek_dsa_decode.py | 2 +- .../ops/attention/test_deepseek_mla_decode.py | 2 +- tests/ops/attention/test_deepseek_nsa.py | 1 - tests/ops/attention/test_deepseek_nsa_cmp.py | 2 +- tests/ops/attention/test_deepseek_nsa_topk.py | 2 +- tests/ops/attention/test_gqa.py | 30 ++-- .../ops/attention/test_gqa_sliding_window.py | 18 ++- .../test_gqa_sliding_window_varlen.py | 2 +- tileops/ops/attention/deepseek_dsa.py | 44 +++--- tileops/ops/attention/deepseek_mla.py | 17 ++- tileops/ops/attention/deepseek_nsa.py | 43 ++++-- tileops/ops/attention/gqa.py | 141 +++++++++++------- 17 files changed, 193 insertions(+), 119 deletions(-) diff --git a/benchmarks/ops/attention/bench_deepseek_dsa_decode.py b/benchmarks/ops/attention/bench_deepseek_dsa_decode.py index 21120cc2b..eefb7e919 100644 --- a/benchmarks/ops/attention/bench_deepseek_dsa_decode.py +++ b/benchmarks/ops/attention/bench_deepseek_dsa_decode.py @@ -81,7 +81,7 @@ def test_dsa_decode_bench(batch: int, heads: int, seq_len_q: int, seq_len_kv: in op = DeepSeekSparseAttentionDecodeWithKVCacheFwdOp( batch, heads, seq_len_q, seq_len_kv, dim, dim_tail, topk, stride_kv, heads_kv, - q_start_index_s, sm_scale=sm_scale, dtype=dtype, tune=tune) + q_start_index_s, sm_scale=sm_scale, tune=tune) bm = ManifestBenchmark(_OP_NAME, op, test) result = bm.profile(op, *inputs) BenchmarkReport.record(op, locals(), result, tag="tileops") diff --git a/benchmarks/ops/attention/bench_deepseek_mla_decode.py b/benchmarks/ops/attention/bench_deepseek_mla_decode.py index a21d25c8c..dd394aea1 100644 --- a/benchmarks/ops/attention/bench_deepseek_mla_decode.py +++ b/benchmarks/ops/attention/bench_deepseek_mla_decode.py @@ -72,7 +72,7 @@ def test_mla_decode_bench(batch: int, heads: int, heads_kv: int, seq_len_kv: int inputs = test.gen_inputs() op = MultiHeadLatentAttentionDecodeWithKVCacheFwdOp( - batch, heads, heads_kv, seq_len_kv, dim, dim_pe, dtype, tune=tune) + batch, heads, heads_kv, seq_len_kv, dim, dim_pe, tune=tune) bm = ManifestBenchmark(_OP_NAME, op, test) result = bm.profile(op, *inputs) BenchmarkReport.record(op, locals(), result, tag="tileops") diff --git a/benchmarks/ops/attention/bench_gqa.py b/benchmarks/ops/attention/bench_gqa.py index bec0ed791..0d3bd7a1f 100644 --- a/benchmarks/ops/attention/bench_gqa.py +++ b/benchmarks/ops/attention/bench_gqa.py @@ -455,7 +455,7 @@ def test_gqa_prefill_varlen_fwd_bench( inputs = test.gen_inputs() op = GroupedQueryAttentionPrefillVarlenFwdOp( - batch, heads, heads_kv, dim, test.max_seqlen_q, test.max_seqlen_kv, causal, dtype, tune=tune + batch, heads, heads_kv, dim, test.max_seqlen_q, test.max_seqlen_kv, causal, tune=tune ) bm = GQAPrefillVarlenFwdBenchmark(test) result = bm.profile(op, *inputs) diff --git a/benchmarks/ops/attention/bench_gqa_sliding_window.py b/benchmarks/ops/attention/bench_gqa_sliding_window.py index 48612c7b7..1f05ba51e 100644 --- a/benchmarks/ops/attention/bench_gqa_sliding_window.py +++ b/benchmarks/ops/attention/bench_gqa_sliding_window.py @@ -133,7 +133,6 @@ def test_gqa_sliding_window_fwd_bench( is_causal=is_causal, window_size_left=wl, window_size_right=wr, - dtype=dtype, tune=tune, ) bm = ManifestBenchmark(_OP_NAME, op, test) diff --git a/benchmarks/ops/attention/bench_gqa_sliding_window_varlen.py b/benchmarks/ops/attention/bench_gqa_sliding_window_varlen.py index 85db3fe76..c5b0969ee 100644 --- a/benchmarks/ops/attention/bench_gqa_sliding_window_varlen.py +++ b/benchmarks/ops/attention/bench_gqa_sliding_window_varlen.py @@ -167,7 +167,6 @@ def test_gqa_sliding_window_varlen_fwd_bench( is_causal=is_causal, window_size_left=wl, window_size_right=wr, - dtype=dtype, tune=tune, ) op.total_q = sum(seqlens_q) diff --git a/tests/ops/attention/test_deepseek_dsa_decode.py b/tests/ops/attention/test_deepseek_dsa_decode.py index ebadc428e..28536494a 100644 --- a/tests/ops/attention/test_deepseek_dsa_decode.py +++ b/tests/ops/attention/test_deepseek_dsa_decode.py @@ -80,7 +80,7 @@ def test_sparse_mla_decode(batch: int, heads: int, seq_len_q: int, seq_len_kv: i q_start_index_s, sm_scale=sm_scale, dtype=dtype) op = DeepSeekSparseAttentionDecodeWithKVCacheFwdOp( batch, heads, seq_len_q, seq_len_kv, dim, dim_tail, topk, stride_kv, heads_kv, - q_start_index_s, sm_scale=sm_scale, dtype=dtype, tune=tune) + q_start_index_s, sm_scale=sm_scale, tune=tune) test.check(op, *test.gen_inputs(), atol=3e-4, rtol=1e-5) diff --git a/tests/ops/attention/test_deepseek_mla_decode.py b/tests/ops/attention/test_deepseek_mla_decode.py index deb86cc7b..57c22b743 100644 --- a/tests/ops/attention/test_deepseek_mla_decode.py +++ b/tests/ops/attention/test_deepseek_mla_decode.py @@ -67,7 +67,7 @@ def test_mla_decode(batch: int, heads: int, heads_kv: int, seq_len_kv: int, dim: dim_pe: int, dtype: torch.dtype, tune: bool): test = MlaDecodeTest(batch, heads, heads_kv, seq_len_kv, dim, dim_pe, dtype) op = MultiHeadLatentAttentionDecodeWithKVCacheFwdOp( - batch, heads, heads_kv, seq_len_kv, dim, dim_pe, dtype, tune=tune) + batch, heads, heads_kv, seq_len_kv, dim, dim_pe, tune=tune) test.check(op, *test.gen_inputs(), atol=1e-3, rtol=1e-3) diff --git a/tests/ops/attention/test_deepseek_nsa.py b/tests/ops/attention/test_deepseek_nsa.py index 68be10905..d60082fe6 100644 --- a/tests/ops/attention/test_deepseek_nsa.py +++ b/tests/ops/attention/test_deepseek_nsa.py @@ -84,7 +84,6 @@ def test_nsa_varlen_op( "block_size": block_size, "groups": groups, "selected_blocks": selected_blocks, - "dtype": dtype, "accum_dtype": accum_dtype, "tune": tune, } diff --git a/tests/ops/attention/test_deepseek_nsa_cmp.py b/tests/ops/attention/test_deepseek_nsa_cmp.py index 6886b7221..489937430 100644 --- a/tests/ops/attention/test_deepseek_nsa_cmp.py +++ b/tests/ops/attention/test_deepseek_nsa_cmp.py @@ -105,7 +105,7 @@ def test_nsa_cmp_fwd_varlen_op( op = NSACmpFwdVarlenOp( seq_num=seq_num, c_seq_len=c_seq_len, heads=heads, dim_k=dim_k, dim_v=dim_v, group=group, - scale=scale, bc=bc, bs=bs, dtype=dtype, accum_dtype=accum_dtype, tune=tune, + scale=scale, bc=bc, bs=bs, accum_dtype=accum_dtype, tune=tune, chunk_num=test.chunk_num) test.check(op, *inputs, atol=4e-3, rtol=1e-5) diff --git a/tests/ops/attention/test_deepseek_nsa_topk.py b/tests/ops/attention/test_deepseek_nsa_topk.py index 0f1c218b7..edaecb827 100644 --- a/tests/ops/attention/test_deepseek_nsa_topk.py +++ b/tests/ops/attention/test_deepseek_nsa_topk.py @@ -207,7 +207,7 @@ def test_nsa_topk_varlen_op( inputs = test.gen_inputs() op = NSATopkVarlenOp( seq_num=seq_num, c_seq_len=c_seq_len, heads=heads, dim=dim, group=group, scale=scale, - selected_block_num=selected_block_num, bc=bc, bs=bs, dtype=dtype, + selected_block_num=selected_block_num, bc=bc, bs=bs, accum_dtype=accum_dtype, tune=tune, chunk_num=test.chunk_num) test.check_topk(op, *inputs) diff --git a/tests/ops/attention/test_gqa.py b/tests/ops/attention/test_gqa.py index 9110c5bd6..0972c4918 100644 --- a/tests/ops/attention/test_gqa.py +++ b/tests/ops/attention/test_gqa.py @@ -815,7 +815,7 @@ def test_gqa_prefill_varlen_rejects_bad_contract_inputs() -> None: [0] + torch.tensor(kv_lens).cumsum(0).tolist(), device="cuda", dtype=torch.int32) op = GroupedQueryAttentionPrefillVarlenFwdOp( - batch, heads, heads_kv, dim, max(q_lens), max(kv_lens), True, torch.float16, + batch, heads, heads_kv, dim, max(q_lens), max(kv_lens), True, validate_inputs=True) with pytest.raises(ValueError, match="Expected k shape"): op(q, k[:, :, :-1].contiguous(), v, cu_q, cu_kv) @@ -823,8 +823,7 @@ def test_gqa_prefill_varlen_rejects_bad_contract_inputs() -> None: op(q[:-1], k, v, cu_q, cu_kv) with pytest.raises(ValueError, match="max_seqlen_q"): bad_op = GroupedQueryAttentionPrefillVarlenFwdOp( - batch, heads, heads_kv, dim, max(q_lens) - 1, max(kv_lens), True, - torch.float16, validate_inputs=True) + batch, heads, heads_kv, dim, max(q_lens) - 1, max(kv_lens), True, validate_inputs=True) bad_op(q, k, v, cu_q, cu_kv) bad_cu = torch.tensor([0, 128, 96], device="cuda", dtype=torch.int32) with pytest.raises(ValueError, match="cu_seqlens_q must be non-decreasing"): @@ -833,16 +832,23 @@ def test_gqa_prefill_varlen_rejects_bad_contract_inputs() -> None: @pytest.mark.smoke def test_gqa_prefill_varlen_rejects_unsupported_dtype() -> None: + """The element type now arrives with the tensors, so the rejection does too.""" + op = GroupedQueryAttentionPrefillVarlenFwdOp( + batch=1, + heads=8, + heads_kv=2, + dim=64, + max_seqlen_q=64, + max_seqlen_kv=128, + ) + kwargs = {"dtype": torch.float32, "device": "cuda"} + q = torch.randn(64, 8, 64, **kwargs) + k = torch.randn(128, 2, 64, **kwargs) + v = torch.randn(128, 2, 64, **kwargs) + cu_q = torch.tensor([0, 64], device="cuda", dtype=torch.int32) + cu_kv = torch.tensor([0, 128], device="cuda", dtype=torch.int32) with pytest.raises(ValueError, match="Expected dtype torch.float16 or torch.bfloat16"): - GroupedQueryAttentionPrefillVarlenFwdOp( - batch=1, - heads=8, - heads_kv=2, - dim=64, - max_seqlen_q=64, - max_seqlen_kv=128, - dtype=torch.float32, - ) + op(q, k, v, cu_q, cu_kv) diff --git a/tests/ops/attention/test_gqa_sliding_window.py b/tests/ops/attention/test_gqa_sliding_window.py index 0376b17a9..8ac47992b 100644 --- a/tests/ops/attention/test_gqa_sliding_window.py +++ b/tests/ops/attention/test_gqa_sliding_window.py @@ -94,7 +94,7 @@ def test_gqa_sliding_window_fwd_op( op = GroupedQueryAttentionSlidingWindowFwdOp( batch=batch, heads=heads, heads_kv=heads_kv, seq_len=seq, dim=dim, is_causal=is_causal, window_size_left=wl, window_size_right=wr, - dtype=dtype, tune=tune) + tune=tune) test.check(op, *test.gen_inputs(), atol=1e-2, rtol=1e-2) @@ -130,9 +130,16 @@ def test_total_memory_gqa(self): op = GroupedQueryAttentionSlidingWindowFwdOp( batch=B, heads=H, heads_kv=Hkv, seq_len=S, dim=D, is_causal=True) - elem = torch.tensor([], dtype=torch.float16).element_size() + with pytest.raises(RuntimeError, match="requires a prior forward"): + _ = op.total_memory + dtype = torch.float16 + op( + torch.randn(B, S, H, D, dtype=dtype, device="cuda"), + torch.randn(B, S, Hkv, D, dtype=dtype, device="cuda"), + torch.randn(B, S, Hkv, D, dtype=dtype, device="cuda"), + ) # Q read + O write: heads each; K read + V read: heads_kv each - expected = 2 * B * S * (H + Hkv) * D * elem + expected = 2 * B * S * (H + Hkv) * D * dtype.itemsize assert op.total_memory == expected, f"got {op.total_memory}, expected {expected}" @@ -161,12 +168,13 @@ def test_invalid_window_size_right_raises(self): def float16_op(self): return GroupedQueryAttentionSlidingWindowFwdOp( batch=1, heads=4, heads_kv=2, seq_len=64, dim=64, - is_causal=True, dtype=torch.float16) + is_causal=True) @pytest.mark.smoke def test_dtype_mismatch_raises(self, float16_op): + """q/k/v must agree with each other; q is the element-type anchor.""" q = torch.randn(1, 64, 4, 64, dtype=torch.bfloat16, device="cuda") - k = torch.randn(1, 64, 2, 64, dtype=torch.bfloat16, device="cuda") + k = torch.randn(1, 64, 2, 64, dtype=torch.float16, device="cuda") v = torch.randn(1, 64, 2, 64, dtype=torch.bfloat16, device="cuda") with pytest.raises(ValueError, match="dtype"): float16_op.forward(q, k, v) diff --git a/tests/ops/attention/test_gqa_sliding_window_varlen.py b/tests/ops/attention/test_gqa_sliding_window_varlen.py index b9d3f462b..a710e9dbd 100644 --- a/tests/ops/attention/test_gqa_sliding_window_varlen.py +++ b/tests/ops/attention/test_gqa_sliding_window_varlen.py @@ -126,7 +126,7 @@ def test_gqa_sliding_window_varlen_fwd_op( op = GroupedQueryAttentionSlidingWindowVarlenFwdOp( batch=batch, heads=heads, heads_kv=heads_kv, dim=dim, is_causal=is_causal, window_size_left=wl, window_size_right=wr, - dtype=dtype, tune=tune) + tune=tune) test.check(op, *test.gen_inputs(), atol=1e-2, rtol=1e-2) diff --git a/tileops/ops/attention/deepseek_dsa.py b/tileops/ops/attention/deepseek_dsa.py index 5c7b00e21..af64f77e9 100644 --- a/tileops/ops/attention/deepseek_dsa.py +++ b/tileops/ops/attention/deepseek_dsa.py @@ -53,7 +53,6 @@ def __init__(self, q_start_index_s: int, sm_scale: Optional[float] = None, is_causal: bool = True, - dtype: torch.dtype = torch.float16, kernel_map: Optional[Dict[str, Kernel]] = None, tune: bool = False) -> None: self.batch = batch @@ -66,7 +65,6 @@ def __init__(self, self.stride_kv = stride_kv self.heads_kv = heads_kv self.sm_scale = sm_scale - self.dtype = dtype self.is_causal = is_causal if q_start_index_s != 0 and q_start_index_s <= stride_kv: @@ -79,23 +77,30 @@ def __init__(self, cp0 = q_start_index_s == 0 self.q_start_index_s = q_start_index_s + self._cp0 = cp0 + self.tune = tune self.dispatch_kernel(kernel_map) - self.kernel = self.kernel_map["sparse_mla_kernel"]( - self.batch, - self.seq_len, - self.seq_len_kv, - self.heads, - self.dim, - self.dim_tail, - self.dtype, - self.topk, - self.stride_kv, - self.q_start_index_s, - self.heads_kv, - self.sm_scale, - self.is_causal, - cp0, - tune=tune) + self._kernel_cache: Dict[torch.dtype, Kernel] = {} + + def _get_kernel(self, dtype: torch.dtype) -> Kernel: + if dtype not in self._kernel_cache: + self._kernel_cache[dtype] = self.kernel_map["sparse_mla_kernel"]( + self.batch, + self.seq_len, + self.seq_len_kv, + self.heads, + self.dim, + self.dim_tail, + dtype, + self.topk, + self.stride_kv, + self.q_start_index_s, + self.heads_kv, + self.sm_scale, + self.is_causal, + self._cp0, + tune=self.tune) + return self._kernel_cache[dtype] @property def default_kernel_map(self) -> Dict[str, Kernel]: @@ -123,4 +128,5 @@ def forward(self, q: torch.Tensor, kv: torch.Tensor, indices: torch.Tensor) -> t torch.Tensor: The result of applying the sparse attention operation on the input tensors. """ - return self.kernel(q, kv, indices) + self.dtype = q.dtype + return self._get_kernel(q.dtype)(q, kv, indices) diff --git a/tileops/ops/attention/deepseek_mla.py b/tileops/ops/attention/deepseek_mla.py index 734acecfc..8cf60514a 100644 --- a/tileops/ops/attention/deepseek_mla.py +++ b/tileops/ops/attention/deepseek_mla.py @@ -20,7 +20,6 @@ def __init__(self, seqlen_kv: int, dim: int, pe_dim: int, - dtype: torch.dtype = torch.float16, kernel_map: Optional[Dict[str, Kernel]] = None, tune: bool = False) -> None: self.batch = batch @@ -30,11 +29,16 @@ def __init__(self, self.dim = dim self.pe_dim = pe_dim - self.dtype = dtype - + self.tune = tune self.dispatch_kernel(kernel_map) - self.kernel = self.kernel_map["mla_decode_kernel"]( - batch, heads, heads_kv, seqlen_kv, dim, pe_dim, self.dtype, tune=tune) + self._kernel_cache: Dict[torch.dtype, Kernel] = {} + + def _get_kernel(self, dtype: torch.dtype) -> Kernel: + if dtype not in self._kernel_cache: + self._kernel_cache[dtype] = self.kernel_map["mla_decode_kernel"]( + self.batch, self.heads, self.heads_kv, self.seqlen_kv, + self.dim, self.pe_dim, dtype, tune=self.tune) + return self._kernel_cache[dtype] @property def default_kernel_map(self) -> Dict[str, Kernel]: @@ -42,4 +46,5 @@ def default_kernel_map(self) -> Dict[str, Kernel]: def forward(self, q: torch.Tensor, q_pe: torch.Tensor, k: torch.Tensor, k_pe: torch.Tensor) -> torch.Tensor: - return self.kernel(q, q_pe, k, k_pe) + self.dtype = q.dtype + return self._get_kernel(q.dtype)(q, q_pe, k, k_pe) diff --git a/tileops/ops/attention/deepseek_nsa.py b/tileops/ops/attention/deepseek_nsa.py index 3a933122e..d46e40fcc 100644 --- a/tileops/ops/attention/deepseek_nsa.py +++ b/tileops/ops/attention/deepseek_nsa.py @@ -32,7 +32,6 @@ def __init__( selected_block_num: int, bc: int, bs: int, - dtype: torch.dtype, accum_dtype: torch.dtype, tune: bool = False, kernel_map: Optional[Dict[str, Kernel]] = None, @@ -41,8 +40,16 @@ def __init__( for key, value in params.items(): setattr(self, key, value) + self._kernel_params = {k: v for k, v in params.items() if k != "dtype"} self.dispatch_kernel(kernel_map) - self.kernel = self.kernel_map["nsa_topk_varlen_kernel"](**params) + self._kernel_cache: Dict[torch.dtype, Kernel] = {} + + def _get_kernel(self, dtype: torch.dtype) -> Kernel: + if dtype not in self._kernel_cache: + self._kernel_cache[dtype] = self.kernel_map["nsa_topk_varlen_kernel"]( + **self._kernel_params, dtype=dtype, + ) + return self._kernel_cache[dtype] @property def default_kernel_map(self) -> Dict[str, Kernel]: @@ -51,7 +58,8 @@ def default_kernel_map(self) -> Dict[str, Kernel]: def forward(self, q: torch.Tensor, k_cmp: torch.Tensor, lse_in: torch.Tensor, offsets: torch.Tensor, chunk_offsets: torch.Tensor, token_indices: torch.Tensor) -> torch.Tensor: - return self.kernel(q, k_cmp, lse_in, offsets, chunk_offsets, token_indices) + self.dtype = q.dtype + return self._get_kernel(q.dtype)(q, k_cmp, lse_in, offsets, chunk_offsets, token_indices) class NSAFwdVarlenOp(Op): @@ -67,7 +75,6 @@ def __init__( block_size: int, groups: int, selected_blocks: int, - dtype: torch.dtype, accum_dtype: torch.dtype, tune: bool = False, kernel_map: Optional[Dict[str, Kernel]] = None, @@ -76,8 +83,16 @@ def __init__( for key, value in params.items(): setattr(self, key, value) + self._kernel_params = {k: v for k, v in params.items() if k != "dtype"} self.dispatch_kernel(kernel_map) - self.kernel = self.kernel_map["nsa_fwd_varlen_kernel"](**params) + self._kernel_cache: Dict[torch.dtype, Kernel] = {} + + def _get_kernel(self, dtype: torch.dtype) -> Kernel: + if dtype not in self._kernel_cache: + self._kernel_cache[dtype] = self.kernel_map["nsa_fwd_varlen_kernel"]( + **self._kernel_params, dtype=dtype, + ) + return self._kernel_cache[dtype] @property def default_kernel_map(self) -> Dict[str, Kernel]: @@ -86,7 +101,8 @@ def default_kernel_map(self) -> Dict[str, Kernel]: def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, block_indices: torch.Tensor, block_counts: torch.Tensor, offsets: torch.Tensor, token_indices: torch.Tensor) -> torch.Tensor: - return self.kernel(q, k, v, block_indices, block_counts, offsets, token_indices) + self.dtype = q.dtype + return self._get_kernel(q.dtype)(q, k, v, block_indices, block_counts, offsets, token_indices) class NSACmpFwdVarlenOp(Op): @@ -103,7 +119,6 @@ def __init__( scale: float, bc: int, bs: int, - dtype: torch.dtype, accum_dtype: torch.dtype, tune: bool = False, kernel_map: Optional[Dict[str, Kernel]] = None, @@ -119,15 +134,22 @@ def __init__( "scale": scale, "bc": bc, "bs": bs, - "dtype": dtype, "accum_dtype": accum_dtype, "tune": tune, } for key, value in params.items(): setattr(self, key, value) + self._kernel_params = {k: v for k, v in params.items() if k != "dtype"} self.dispatch_kernel(kernel_map) - self.kernel = self.kernel_map["nsa_cmp_fwd_varlen_kernel"](**params) + self._kernel_cache: Dict[torch.dtype, Kernel] = {} + + def _get_kernel(self, dtype: torch.dtype) -> Kernel: + if dtype not in self._kernel_cache: + self._kernel_cache[dtype] = self.kernel_map["nsa_cmp_fwd_varlen_kernel"]( + **self._kernel_params, dtype=dtype, + ) + return self._kernel_cache[dtype] @property def default_kernel_map(self) -> Dict[str, Kernel]: @@ -136,4 +158,5 @@ def default_kernel_map(self) -> Dict[str, Kernel]: def forward(self, q: torch.Tensor, k_cmp: torch.Tensor, v_cmp: torch.Tensor, offsets: torch.Tensor, chunk_offsets: torch.Tensor, token_indices: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - return self.kernel(q, k_cmp, v_cmp, offsets, chunk_offsets, token_indices) + self.dtype = q.dtype + return self._get_kernel(q.dtype)(q, k_cmp, v_cmp, offsets, chunk_offsets, token_indices) diff --git a/tileops/ops/attention/gqa.py b/tileops/ops/attention/gqa.py index 8691cf6ff..d8d4c60ec 100644 --- a/tileops/ops/attention/gqa.py +++ b/tileops/ops/attention/gqa.py @@ -767,7 +767,6 @@ def __init__( max_seqlen_q: int, max_seqlen_kv: int, is_causal: bool = True, - dtype: torch.dtype = torch.float16, sm_scale: Optional[float] = None, softcap: Optional[float] = None, validate_inputs: bool = False, @@ -781,7 +780,6 @@ def __init__( raise ValueError("max_seqlen_q must be positive") if max_seqlen_kv <= 0: raise ValueError("max_seqlen_kv must be positive") - _validate_attention_dtype(dtype) self.batch = batch self.heads = heads self.heads_kv = heads_kv @@ -789,24 +787,32 @@ def __init__( self.max_seqlen_q = max_seqlen_q self.max_seqlen_kv = max_seqlen_kv self.is_causal = is_causal - self.dtype = dtype self.sm_scale = _attention_scale(dim, sm_scale) self.softcap = _score_softcap(softcap) self.validate_inputs = validate_inputs self._roofline_kwargs = None + self.tune = tune self.dispatch_kernel(kernel_map) - self.kernel = self.kernel_map["gqa_prefill_varlen_fwd_kernel"]( - batch=batch, - heads=heads, - heads_kv=heads_kv, - dim=dim, - is_causal=is_causal, - dtype=dtype, - sm_scale=self.sm_scale, - softcap=self.softcap, - tune=tune, - ) + self._kernel_cache: Dict[torch.dtype, Kernel] = {} + + def _get_kernel(self, dtype: torch.dtype) -> Kernel: + if dtype not in self._kernel_cache: + _validate_attention_dtype(dtype) + self._kernel_cache[dtype] = self.kernel_map[ + "gqa_prefill_varlen_fwd_kernel" + ]( + batch=self.batch, + heads=self.heads, + heads_kv=self.heads_kv, + dim=self.dim, + is_causal=self.is_causal, + dtype=dtype, + sm_scale=self.sm_scale, + softcap=self.softcap, + tune=self.tune, + ) + return self._kernel_cache[dtype] @property def default_kernel_map(self) -> Dict[str, Kernel]: @@ -838,6 +844,9 @@ def _validate_forward_inputs( if not tensor.is_contiguous(): raise ValueError(f"{name} must be contiguous") + # q carries the element type; k and v must agree with it. + _validate_attention_dtype(tensors["q"].dtype) + expected_tail_shapes = { "q": (self.heads, self.dim), "k": (self.heads_kv, self.dim), @@ -849,8 +858,9 @@ def _validate_forward_inputs( raise ValueError( f"Expected {name} shape [T, {expected_tail[0]}, {expected_tail[1]}], " f"got {tuple(tensor.shape)}") - if tensor.dtype != self.dtype: - raise ValueError(f"Expected {name}.dtype {self.dtype}, got {tensor.dtype}") + if tensor.dtype != tensors["q"].dtype: + raise ValueError( + f"Expected {name}.dtype {tensors['q'].dtype}, got {tensor.dtype}") for name in ("cu_seqlens_q", "cu_seqlens_kv"): tensor = tensors[name] @@ -916,7 +926,8 @@ def forward( cu_seqlens_kv: torch.Tensor, ) -> torch.Tensor: self._validate_forward_inputs(q, k, v, cu_seqlens_q, cu_seqlens_kv) - output, _ = self.kernel( + self.dtype = q.dtype + output, _ = self._get_kernel(q.dtype)( q, k, v, cu_seqlens_q, cu_seqlens_kv, self.max_seqlen_q, self.max_seqlen_kv) self._roofline_kwargs = { "q_shape": tuple(q.shape), @@ -1320,7 +1331,6 @@ def __init__(self, seq_len: int, dim: int, is_causal: bool = True, - dtype: torch.dtype = torch.float16, kernel_map: Optional[Dict[str, Kernel]] = None, tune: bool = False) -> None: self.batch = batch @@ -1330,7 +1340,6 @@ def __init__(self, self.dim = dim self.is_causal = is_causal - self.dtype = dtype self.tune = tune self.dispatch_kernel(kernel_map) @@ -1561,7 +1570,6 @@ def __init__( is_causal: bool = True, window_size_left: int = -1, window_size_right: int = -1, - dtype: torch.dtype = torch.float16, kernel_map: Optional[Dict[str, Kernel]] = None, tune: bool = False, ) -> None: @@ -1581,21 +1589,26 @@ def __init__( self.is_causal = is_causal self.window_size_left = window_size_left self.window_size_right = window_size_right - self.dtype = dtype + self.tune = tune self.dispatch_kernel(kernel_map) - self.kernel = self.kernel_map["gqa_sliding_window_fwd"]( - batch=batch, - heads=heads, - heads_kv=heads_kv, - seq_len=seq_len, - dim=dim, - is_causal=is_causal, - window_size_left=window_size_left, - window_size_right=window_size_right, - dtype=dtype, - tune=tune, - ) + self._kernel_cache: Dict[torch.dtype, Kernel] = {} + + def _get_kernel(self, dtype: torch.dtype) -> Kernel: + if dtype not in self._kernel_cache: + self._kernel_cache[dtype] = self.kernel_map["gqa_sliding_window_fwd"]( + batch=self.batch, + heads=self.heads, + heads_kv=self.heads_kv, + seq_len=self.seq_len, + dim=self.dim, + is_causal=self.is_causal, + window_size_left=self.window_size_left, + window_size_right=self.window_size_right, + dtype=dtype, + tune=self.tune, + ) + return self._kernel_cache[dtype] @property def default_kernel_map(self) -> Dict[str, Kernel]: @@ -1622,9 +1635,9 @@ def forward( if t.device.type != 'cuda': raise ValueError( f"{name} must be on a cuda device, got {t.device}") - if t.dtype != self.dtype: + if t.dtype != q.dtype: raise ValueError( - f"{name} dtype {t.dtype} does not match op dtype {self.dtype}") + f"{name} dtype {t.dtype} does not match q dtype {q.dtype}") if not q.is_contiguous(): q = q.contiguous() if not k.is_contiguous(): @@ -1645,7 +1658,8 @@ def forward( f"v shape {v.shape} does not match expected " f"({self.batch}, {self.seq_len}, {self.heads_kv}, {self.dim})") - output, _ = self.kernel.forward(q, k, v) + self.dtype = q.dtype + output, _ = self._get_kernel(q.dtype).forward(q, k, v) return output @property @@ -1663,9 +1677,17 @@ def total_flops(self) -> int: @property def total_memory(self) -> int: - """Approximate bytes accessed: read Q/K/V, write O.""" - elem = torch.tensor([], dtype=self.dtype).element_size() - return 2 * self.batch * self.seq_len * (self.heads + self.heads_kv) * self.dim * elem + """Approximate bytes accessed: read Q/K/V, write O. + + Available after the first ``forward()``, which binds the element type + from its input; there is no element type before that. + """ + if self.dtype is None: + raise RuntimeError( + f"{type(self).__name__}.total_memory requires a prior forward() " + "call to bind the element type") + return (2 * self.batch * self.seq_len * (self.heads + self.heads_kv) + * self.dim * self.dtype.itemsize) class GroupedQueryAttentionSlidingWindowVarlenFwdOp(Op): @@ -1704,7 +1726,6 @@ def __init__( is_causal: bool = True, window_size_left: int = -1, window_size_right: int = -1, - dtype: torch.dtype = torch.float16, accum_dtype: torch.dtype = torch.float32, kernel_map: Optional[Dict[str, Kernel]] = None, tune: bool = False, @@ -1724,22 +1745,29 @@ def __init__( self.is_causal = is_causal self.window_size_left = window_size_left self.window_size_right = window_size_right - self.dtype = dtype self.accum_dtype = accum_dtype + self.tune = tune self.dispatch_kernel(kernel_map) - self.kernel = self.kernel_map["gqa_sliding_window_varlen_fwd"]( - batch=batch, - heads=heads, - heads_kv=heads_kv, - dim=dim, - is_causal=is_causal, - window_size_left=window_size_left, - window_size_right=window_size_right, - dtype=dtype, - accum_dtype=accum_dtype, - tune=tune, - ) + self._kernel_cache: Dict[torch.dtype, Kernel] = {} + + def _get_kernel(self, dtype: torch.dtype) -> Kernel: + if dtype not in self._kernel_cache: + self._kernel_cache[dtype] = self.kernel_map[ + "gqa_sliding_window_varlen_fwd" + ]( + batch=self.batch, + heads=self.heads, + heads_kv=self.heads_kv, + dim=self.dim, + is_causal=self.is_causal, + window_size_left=self.window_size_left, + window_size_right=self.window_size_right, + dtype=dtype, + accum_dtype=self.accum_dtype, + tune=self.tune, + ) + return self._kernel_cache[dtype] @property def default_kernel_map(self) -> Dict[str, Kernel]: @@ -1772,9 +1800,9 @@ def forward( if t.device.type != 'cuda': raise ValueError( f"{name} must be on a cuda device, got {t.device}") - if t.dtype != self.dtype: + if t.dtype != q.dtype: raise ValueError( - f"{name} dtype {t.dtype} does not match op dtype {self.dtype}") + f"{name} dtype {t.dtype} does not match q dtype {q.dtype}") if not t.is_contiguous(): raise ValueError(f"{name} must be contiguous") @@ -1832,7 +1860,8 @@ def forward( f"max_seqlen_q ({max_seqlen_q}) must be >= actual max Q " f"sequence length ({actual_max_q})") - output, _ = self.kernel.forward( + self.dtype = q.dtype + output, _ = self._get_kernel(q.dtype).forward( q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q) return output From 40021efd404cb11dcc7954a6859fce4c2b6e004e Mon Sep 17 00:00:00 2001 From: lcy-seso Date: Tue, 4 Aug 2026 20:44:29 +0800 Subject: [PATCH 09/15] [Refactor][Attention] Select the decode kernel at forward instead of at construction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both GQA decode ops chose their kernel *slot* from `self.dtype` inside `__init__`. The slot table already held both candidates, so only the timing was wrong: `_select_decode_kernel_key` and `_uses_bs1_fast_path` now take the element type, and the kernel is built and cached on first forward. `_uses_bs1_fast_path` stops being a property of the op. Whether a request takes the batch=1 warp-specialized kernel depends on the element type of the call, so one instance answers differently for float16 and bfloat16 — which is what the dispatch test now asserts against a single instance instead of constructing a second op it could no longer distinguish. Two paged tests moved their expectation to where the work now happens: an unsupported page layout is still rejected, but by the kernel the op builds on first use rather than by the constructor. --- benchmarks/ops/attention/bench_gqa_decode.py | 1 - .../ops/attention/bench_gqa_decode_paged.py | 2 +- tests/ops/attention/test_gqa_decode.py | 52 ++++----- tests/ops/attention/test_gqa_decode_paged.py | 45 ++++---- tileops/ops/attention/gqa.py | 101 ++++++++++-------- 5 files changed, 110 insertions(+), 91 deletions(-) diff --git a/benchmarks/ops/attention/bench_gqa_decode.py b/benchmarks/ops/attention/bench_gqa_decode.py index 4e4fcd44f..efdec9141 100644 --- a/benchmarks/ops/attention/bench_gqa_decode.py +++ b/benchmarks/ops/attention/bench_gqa_decode.py @@ -145,7 +145,6 @@ def test_gqa_decode_bench(batch: int, heads: int, heads_kv: int, seq_len_kv: int heads_kv, seq_len_kv, dim, - dtype, sm_scale=sm_scale, softcap=softcap, tune=tune, diff --git a/benchmarks/ops/attention/bench_gqa_decode_paged.py b/benchmarks/ops/attention/bench_gqa_decode_paged.py index 08f45c27e..dd17d7c93 100644 --- a/benchmarks/ops/attention/bench_gqa_decode_paged.py +++ b/benchmarks/ops/attention/bench_gqa_decode_paged.py @@ -150,7 +150,7 @@ def test_gqa_decode_paged_bench(batch: int, heads: int, heads_kv: int, seqlen_kv q, k, v, real_seqlen_kv, block_table = inputs op = GroupedQueryAttentionDecodePagedWithKVCacheFwdOp( - batch, heads, heads_kv, seqlen_kv, dim, page_size, dtype, + batch, heads, heads_kv, seqlen_kv, dim, page_size, sm_scale=sm_scale, softcap=softcap, tune=tune) bm = ManifestBenchmark(_OP_NAME, op, test) result = bm.profile(op, *inputs) diff --git a/tests/ops/attention/test_gqa_decode.py b/tests/ops/attention/test_gqa_decode.py index 56e26daa3..bb29e4566 100644 --- a/tests/ops/attention/test_gqa_decode.py +++ b/tests/ops/attention/test_gqa_decode.py @@ -39,7 +39,7 @@ class GroupedQueryAttentionDecodeFixture(FixtureBase): def test_gqa_decode(batch: int, heads: int, heads_kv: int, seq_len_kv: int, dim: int, dtype: torch.dtype, tune: bool) -> None: test = GroupedQueryAttentionDecodeTest(batch, heads, heads_kv, seq_len_kv, dim, dtype) - op = GroupedQueryAttentionDecodeWithKVCacheFwdOp(batch, heads, heads_kv, seq_len_kv, dim, dtype, tune=tune) + op = GroupedQueryAttentionDecodeWithKVCacheFwdOp(batch, heads, heads_kv, seq_len_kv, dim, tune=tune) test.check(op, *test.gen_inputs(), atol=1e-2, rtol=1e-2) @@ -67,7 +67,6 @@ def test_gqa_decode_softmax_controls(sm_scale: float | None, softcap: float | No heads_kv, seq_len_kv, dim, - dtype, sm_scale=sm_scale, softcap=softcap, ) @@ -120,23 +119,25 @@ def test_gqa_decode_rejects_non_positive_seqlen_kv() -> None: @pytest.mark.smoke def test_gqa_decode_bs1_dispatch() -> None: """batch=1 fp16 dim-128 requests select the WS kernel; other dtypes/shapes fall back.""" - op = GroupedQueryAttentionDecodeWithKVCacheFwdOp(1, 32, 4, 8192, 128, torch.float16) - assert op._uses_bs1_fast_path() - assert op.kernel.__class__.__name__ == "GQADecodeBs1Kernel" - assert op.kernel._select_tier(6000) == "ctx" - assert op.kernel._select_tier(1024) == "ctx" - assert op.kernel._select_tier(512) == "no_split" - assert op.kernel._ctx_splits_for(8192) == 32 - assert op.kernel._ctx_splits_for(2048) == 16 - assert op.kernel._ctx_splits_for(3072) == 8 - - op_bf16 = GroupedQueryAttentionDecodeWithKVCacheFwdOp(1, 32, 4, 4096, 128, torch.bfloat16) - assert not op_bf16._uses_bs1_fast_path() - assert op_bf16.kernel.__class__.__name__ == "GQADecodeKernel" - - op_batched = GroupedQueryAttentionDecodeWithKVCacheFwdOp(4, 32, 4, 4096, 128, torch.float16) - assert not op_batched._uses_bs1_fast_path() - assert op_batched.kernel.__class__.__name__ == "GQADecodeKernel" + op = GroupedQueryAttentionDecodeWithKVCacheFwdOp(1, 32, 4, 8192, 128) + assert op._uses_bs1_fast_path(torch.float16) + kernel = op._get_kernel(torch.float16) + assert kernel.__class__.__name__ == "GQADecodeBs1Kernel" + assert kernel._select_tier(6000) == "ctx" + assert kernel._select_tier(1024) == "ctx" + assert kernel._select_tier(512) == "no_split" + assert kernel._ctx_splits_for(8192) == 32 + assert kernel._ctx_splits_for(2048) == 16 + assert kernel._ctx_splits_for(3072) == 8 + + # The same instance falls back for bfloat16 — the element type is an input + # to the choice, so one op serves both paths. + assert not op._uses_bs1_fast_path(torch.bfloat16) + assert op._get_kernel(torch.bfloat16).__class__.__name__ == "GQADecodeKernel" + + op_batched = GroupedQueryAttentionDecodeWithKVCacheFwdOp(4, 32, 4, 4096, 128) + assert not op_batched._uses_bs1_fast_path(torch.float16) + assert op_batched._get_kernel(torch.float16).__class__.__name__ == "GQADecodeKernel" @pytest.mark.smoke @@ -146,11 +147,12 @@ def test_gqa_decode_bs1_runtime_context_switch() -> None: Covers the crossover (1024), a balanced mid split, an aligned split needing >=3 tiles per slice, an unaligned length, and the sub-1024 non-split fallback. """ - op = GroupedQueryAttentionDecodeWithKVCacheFwdOp(1, 32, 4, 8192, 128, torch.float16) - assert op.kernel.__class__.__name__ == "GQADecodeBs1Kernel" + op = GroupedQueryAttentionDecodeWithKVCacheFwdOp(1, 32, 4, 8192, 128) + kernel = op._get_kernel(torch.float16) + assert kernel.__class__.__name__ == "GQADecodeBs1Kernel" for real, tier in ((6000, "ctx"), (3072, "ctx"), (2048, "ctx"), (1024, "ctx"), (512, "no_split")): - assert op.kernel._select_tier(real) == tier + assert kernel._select_tier(real) == tier test = GroupedQueryAttentionDecodeTest(1, 32, 4, real, 128, torch.float16) test.check(op, *test.gen_inputs(), atol=1e-2, rtol=1e-2) @@ -158,9 +160,9 @@ def test_gqa_decode_bs1_runtime_context_switch() -> None: @pytest.mark.smoke def test_gqa_decode_bs1_group4() -> None: """The WS kernel generalizes to a query-per-KV-head group other than 8 (here 4).""" - op = GroupedQueryAttentionDecodeWithKVCacheFwdOp(1, 32, 8, 4096, 128, torch.float16) - assert op._uses_bs1_fast_path() - assert op.kernel.__class__.__name__ == "GQADecodeBs1Kernel" + op = GroupedQueryAttentionDecodeWithKVCacheFwdOp(1, 32, 8, 4096, 128) + assert op._uses_bs1_fast_path(torch.float16) + assert op._get_kernel(torch.float16).__class__.__name__ == "GQADecodeBs1Kernel" test = GroupedQueryAttentionDecodeTest(1, 32, 8, 4096, 128, torch.float16) test.check(op, *test.gen_inputs(), atol=1e-2, rtol=1e-2) diff --git a/tests/ops/attention/test_gqa_decode_paged.py b/tests/ops/attention/test_gqa_decode_paged.py index 8230f5ead..c0d9cb10d 100644 --- a/tests/ops/attention/test_gqa_decode_paged.py +++ b/tests/ops/attention/test_gqa_decode_paged.py @@ -97,7 +97,6 @@ def test_gqa_decode_paged_op( seqlen_kv=seqlen_kv, dim=dim, page_size=page_size, - dtype=dtype, tune=tune, ) test.check(op, *test.gen_inputs(), compare=test._maxdiff_cosine_compare) @@ -114,9 +113,10 @@ def test_gqa_decode_paged_non_divisible_128_page_split() -> None: block_table.copy_(torch.arange( seqlen_kv // page_size, device="cuda", dtype=torch.int32).flip(0).unsqueeze(0)) op = GroupedQueryAttentionDecodePagedWithKVCacheFwdOp( - batch, heads, heads_kv, seqlen_kv, dim, page_size, torch.float16) - assert page_size % op.kernel.config["block_N"] == 0 - assert {config["block_N"] for config in op.kernel.autotune_configs} == {64} + batch, heads, heads_kv, seqlen_kv, dim, page_size) + kernel = op._get_kernel(torch.float16) + assert page_size % kernel.config["block_N"] == 0 + assert {config["block_N"] for config in kernel.autotune_configs} == {64} test.check( op, q, k, v, real_seqlen_kv, block_table, compare=test._maxdiff_cosine_compare) @@ -127,9 +127,11 @@ def test_gqa_decode_paged_non_divisible_128_page_split() -> None: def test_gqa_decode_paged_rejects_unsupported_page_tile(page_size: int) -> None: """Reject page layouts that no supported generic N tile can cover exactly.""" seqlen_kv = page_size * 16 + op = GroupedQueryAttentionDecodePagedWithKVCacheFwdOp( + 1, 16, 4, seqlen_kv, 128, page_size) + # The page layout is rejected by the kernel, which is built on first use. with pytest.raises(ValueError, match="matches no supported block_N"): - GroupedQueryAttentionDecodePagedWithKVCacheFwdOp( - 1, 16, 4, seqlen_kv, 128, page_size, torch.float16) + op._get_kernel(torch.float16) @pytest.mark.smoke @@ -161,7 +163,6 @@ def test_gqa_decode_paged_op_softmax_controls( seqlen_kv=seqlen_kv, dim=dim, page_size=page_size, - dtype=dtype, sm_scale=sm_scale, softcap=softcap, ) @@ -188,15 +189,16 @@ def test_gqa_decode_paged_bs1_fixed_tier_correctness( batch, heads, heads_kv, seqlen_kv, dim, page_size, dtype ) op = GroupedQueryAttentionDecodePagedWithKVCacheFwdOp( - batch, heads, heads_kv, seqlen_kv, dim, page_size, dtype + batch, heads, heads_kv, seqlen_kv, dim, page_size ) q, k, v, real_seqlen_kv, block_table = test.gen_inputs() real_seqlen_kv.fill_(real_seqlen_kv_value) if reverse_pages: block_table = block_table.flip(-1).contiguous() - assert op.kernel.__class__.__name__ == "GQADecodePagedBs1Kernel" - assert op.kernel._select_tier(real_seqlen_kv_value) == ( + kernel = op._get_kernel(torch.float16) + assert kernel.__class__.__name__ == "GQADecodePagedBs1Kernel" + assert kernel._select_tier(real_seqlen_kv_value) == ( "ctx" if real_seqlen_kv_value >= 1024 else "no_split" ) test.check( @@ -214,14 +216,15 @@ def test_gqa_decode_paged_bs1_fixed_tier_correctness( def test_gqa_decode_paged_bs1_dispatch() -> None: """Eligible Hopper requests select the paged TMA/WGMMA kernel.""" op = GroupedQueryAttentionDecodePagedWithKVCacheFwdOp( - 1, 32, 4, 8192, 128, 256, torch.float16) - assert op._uses_bs1_fast_path() - assert op.kernel.__class__.__name__ == "GQADecodePagedBs1Kernel" - assert op.kernel._select_tier(1024) == "ctx" - assert op.kernel._select_tier(512) == "no_split" - assert op.kernel._ctx_splits_for(8192) == 32 - assert op.kernel._ctx_splits_for(2048) == 16 - assert op.kernel._ctx_splits_for(3072) == 8 + 1, 32, 4, 8192, 128, 256) + assert op._uses_bs1_fast_path(torch.float16) + kernel = op._get_kernel(torch.float16) + assert kernel.__class__.__name__ == "GQADecodePagedBs1Kernel" + assert kernel._select_tier(1024) == "ctx" + assert kernel._select_tier(512) == "no_split" + assert kernel._ctx_splits_for(8192) == 32 + assert kernel._ctx_splits_for(2048) == 16 + assert kernel._ctx_splits_for(3072) == 8 @pytest.mark.smoke @@ -246,9 +249,9 @@ def test_gqa_decode_paged_bs1_dispatch_fallbacks( """Unsupported shapes and features stay on the generic paged kernel.""" seqlen_kv = 8064 if page_size == 192 else 8192 op = GroupedQueryAttentionDecodePagedWithKVCacheFwdOp( - batch, 32, 4, seqlen_kv, dim, page_size, dtype, softcap=softcap) - assert not op._uses_bs1_fast_path() - assert op.kernel.__class__.__name__ == "GQADecodePagedKernel" + batch, 32, 4, seqlen_kv, dim, page_size, softcap=softcap) + assert not op._uses_bs1_fast_path(dtype) + assert op._get_kernel(dtype).__class__.__name__ == "GQADecodePagedKernel" if __name__ == "__main__": pytest.main([__file__, "-vvs"]) diff --git a/tileops/ops/attention/gqa.py b/tileops/ops/attention/gqa.py index d8d4c60ec..dc4f95dc7 100644 --- a/tileops/ops/attention/gqa.py +++ b/tileops/ops/attention/gqa.py @@ -1394,35 +1394,41 @@ def __init__(self, heads_kv: int, seqlen_kv: int, dim: int, - dtype: torch.dtype = torch.float16, sm_scale: Optional[float] = None, softcap: Optional[float] = None, kernel_map: Optional[Dict[str, Kernel]] = None, tune: bool = False) -> None: _validate_gqa_dims(heads, heads_kv, dim) - _validate_attention_dtype(dtype) self.batch = batch self.heads = heads self.heads_kv = heads_kv self.seqlen_kv = seqlen_kv self.dim = dim - self.dtype = dtype self.sm_scale = _attention_scale(dim, sm_scale) self.softcap = _score_softcap(softcap) + self.tune = tune self.dispatch_kernel(kernel_map) - self.kernel = self.kernel_map[self._select_decode_kernel_key()]( - batch, - heads, - heads_kv, - seqlen_kv, - dim, - self.dtype, - sm_scale=self.sm_scale, - softcap=self.softcap, - tune=tune, - ) + self._kernel_cache: Dict[torch.dtype, Kernel] = {} + + def _get_kernel(self, dtype: torch.dtype) -> Kernel: + if dtype not in self._kernel_cache: + _validate_attention_dtype(dtype) + self._kernel_cache[dtype] = self.kernel_map[ + self._select_decode_kernel_key(dtype) + ]( + self.batch, + self.heads, + self.heads_kv, + self.seqlen_kv, + self.dim, + dtype, + sm_scale=self.sm_scale, + softcap=self.softcap, + tune=self.tune, + ) + return self._kernel_cache[dtype] @property def default_kernel_map(self) -> Dict[str, Kernel]: @@ -1431,24 +1437,25 @@ def default_kernel_map(self) -> Dict[str, Kernel]: "gqa_decode_bs1_kernel": GQADecodeBs1Kernel, } - def _uses_bs1_fast_path(self) -> bool: - """Ctor-time gate for the batch=1 warp-specialized decode kernel. + def _uses_bs1_fast_path(self, dtype: torch.dtype) -> bool: + """Whether *dtype* routes to the batch=1 warp-specialized decode kernel. - Any batch=1 fp16 Hopper request with dim 128 and a query-per-KV-head group that fits - one wgmma tile routes to GQADecodeBs1Kernel, which then switches on the runtime KV - length in forward(). + A batch=1 fp16 request with dim 128 and a query-per-KV-head group that + fits one wgmma tile takes GQADecodeBs1Kernel, which then switches on the + runtime KV length in forward(). The element type is one of the inputs, so + this is a property of the call rather than of the op. """ return _supports_gqa_decode_bs1( self.batch, self.heads, self.heads_kv, self.dim, - self.dtype, + dtype, self.softcap, ) - def _select_decode_kernel_key(self) -> str: - if self._uses_bs1_fast_path(): + def _select_decode_kernel_key(self, dtype: torch.dtype) -> str: + if self._uses_bs1_fast_path(dtype): return "gqa_decode_bs1_kernel" return "gqa_decode_kernel" @@ -1460,7 +1467,8 @@ def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Te v = F.pad( v, pad=(0, 0, 0, 0, 0, self.seqlen_kv - real_seqlen_kv), mode='constant', value=0) - return self.kernel(q, k, v, real_seqlen_kv) + self.dtype = q.dtype + return self._get_kernel(q.dtype)(q, k, v, real_seqlen_kv) class GroupedQueryAttentionDecodePagedWithKVCacheFwdOp(Op): @@ -1475,13 +1483,11 @@ def __init__(self, seqlen_kv: int, dim: int, page_size: int, - dtype: torch.dtype = torch.float16, sm_scale: Optional[float] = None, softcap: Optional[float] = None, kernel_map: Optional[Dict[str, Kernel]] = None, tune: bool = False) -> None: _validate_gqa_dims(heads, heads_kv, dim) - _validate_attention_dtype(dtype) self.batch = batch self.heads = heads self.heads_kv = heads_kv @@ -1490,23 +1496,31 @@ def __init__(self, self.page_size = page_size if page_size <= 0: raise ValueError("page_size must be positive") - self.dtype = dtype self.sm_scale = _attention_scale(dim, sm_scale) self.softcap = _score_softcap(softcap) + self.tune = tune self.dispatch_kernel(kernel_map) - self.kernel = self.kernel_map[self._select_decode_kernel_key()]( - batch, - heads, - heads_kv, - seqlen_kv, - dim, - page_size, - self.dtype, - sm_scale=self.sm_scale, - softcap=self.softcap, - tune=tune, - ) + self._kernel_cache: Dict[torch.dtype, Kernel] = {} + + def _get_kernel(self, dtype: torch.dtype) -> Kernel: + if dtype not in self._kernel_cache: + _validate_attention_dtype(dtype) + self._kernel_cache[dtype] = self.kernel_map[ + self._select_decode_kernel_key(dtype) + ]( + self.batch, + self.heads, + self.heads_kv, + self.seqlen_kv, + self.dim, + self.page_size, + dtype, + sm_scale=self.sm_scale, + softcap=self.softcap, + tune=self.tune, + ) + return self._kernel_cache[dtype] @property def default_kernel_map(self) -> Dict[str, Kernel]: @@ -1515,25 +1529,26 @@ def default_kernel_map(self) -> Dict[str, Kernel]: "gqa_decode_paged_bs1_kernel": GQADecodePagedBs1Kernel, } - def _uses_bs1_fast_path(self) -> bool: + def _uses_bs1_fast_path(self, dtype: torch.dtype) -> bool: """Use the paged Hopper fast path only when its TMA tile stays within one page.""" return _supports_gqa_decode_bs1( self.batch, self.heads, self.heads_kv, self.dim, - self.dtype, + dtype, self.softcap, ) and GQADecodePagedBs1Kernel.block_n_for_page_size(self.page_size) is not None - def _select_decode_kernel_key(self) -> str: - if self._uses_bs1_fast_path(): + def _select_decode_kernel_key(self, dtype: torch.dtype) -> str: + if self._uses_bs1_fast_path(dtype): return "gqa_decode_paged_bs1_kernel" return "gqa_decode_paged_kernel" def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, real_seqlen_kv: torch.Tensor, block_table: torch.Tensor) -> torch.Tensor: - return self.kernel(q, k, v, real_seqlen_kv, block_table) + self.dtype = q.dtype + return self._get_kernel(q.dtype)(q, k, v, real_seqlen_kv, block_table) class GroupedQueryAttentionSlidingWindowFwdOp(Op): From 40182d534e324dd808bedbf294d8b640bb674db5 Mon Sep 17 00:00:00 2001 From: lcy-seso Date: Tue, 4 Aug 2026 21:13:42 +0800 Subject: [PATCH 10/15] [Refactor][Ops] Route every compile-dispatch instance through one registry MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `elementwise/_base.py`, `rope.py` and `dropout.py` each kept a private `WeakValueDictionary` alongside the shared one in `compile_boundary.py`, and thirteen constructors did self.dispatch_kernel(kernel_map) # sets a str key, registers self._instance_key = id(self) # overwrites it with an int _OP_REGISTRY[self._instance_key] = self `compile_boundary` documents why the key must be a string: dynamo generalizes an int custom-op argument to an unhashable `SymInt` once a second instance compiles through the same frame. Every one of the 23 wrappers was annotated `instance_key: int`, so that hazard was live rather than hypothetical. The private dictionaries are gone, the manual registrations with them, and the wrappers resolve through `get_instance`. `_IntIdentityUnaryOp`'s integer path called `_install_kernel_map` directly, which skips the registration entirely — those instances were never registered at all. It now routes through `dispatch_kernel`. Also picks up two more call sites that arrived with the rebase: a `dtype=` construction and a shape-only kernel-cache key in the autotune tests added upstream. --- tests/ops/test_logical_reduce.py | 4 +- tileops/ops/dropout.py | 11 +-- tileops/ops/elementwise/_base.py | 101 +++++++++++-------------- tileops/ops/elementwise/arithmetic.py | 4 +- tileops/ops/elementwise/clamp.py | 9 --- tileops/ops/elementwise/masked_fill.py | 6 +- tileops/ops/elementwise/nan_to_num.py | 4 +- tileops/ops/elementwise/prelu.py | 3 - tileops/ops/elementwise/where.py | 3 - tileops/ops/rope.py | 23 +++--- 10 files changed, 62 insertions(+), 106 deletions(-) diff --git a/tests/ops/test_logical_reduce.py b/tests/ops/test_logical_reduce.py index b15374c99..103e5e15e 100644 --- a/tests/ops/test_logical_reduce.py +++ b/tests/ops/test_logical_reduce.py @@ -667,10 +667,10 @@ def test_logical_reduce_tiled_autotune() -> None: m, n, dtype = 4, 40000, torch.bool test = LogicalReduceTest(m, n, dtype, "any") - op = AnyFwdOp(dtype=dtype, dim=-1, tune=True) + op = AnyFwdOp(dim=-1, tune=True) test.check(op, *test.gen_inputs(), compare=_exact_compare) - kernel = op._kernel_cache[(m, n)] + kernel = op._kernel_cache[(m, n, dtype)] assert kernel._needs_tiling assert kernel.config in kernel.autotune_configs diff --git a/tileops/ops/dropout.py b/tileops/ops/dropout.py index fb5f0758b..a8f7540cc 100644 --- a/tileops/ops/dropout.py +++ b/tileops/ops/dropout.py @@ -10,7 +10,6 @@ - training=False: identity pass-through """ -import weakref from typing import Dict, Optional import torch @@ -18,11 +17,11 @@ from tileops.kernels.dropout import DropoutKernel from tileops.kernels.kernel_base import Kernel +from .compile_boundary import get_instance from .op_base import Op __all__ = ["DropoutOp"] -_OP_REGISTRY: weakref.WeakValueDictionary = weakref.WeakValueDictionary() class DropoutOp(Op): @@ -75,8 +74,6 @@ def __init__( self._kernel_cache: Dict[tuple[int, torch.dtype, int | None], Kernel] = {} self.kernel = None - self._instance_key = id(self) - _OP_REGISTRY[self._instance_key] = self @property def default_kernel_map(self) -> Dict[str, Kernel]: @@ -128,13 +125,13 @@ def forward(self, input: torch.Tensor) -> torch.Tensor: # torch.compile registration @torch.library.custom_op("top::dropout", mutates_args=()) -def _wrapped_dropout(x: torch.Tensor, instance_key: int) -> torch.Tensor: - instance = _OP_REGISTRY[instance_key] +def _wrapped_dropout(x: torch.Tensor, instance_key: str) -> torch.Tensor: + instance = get_instance(instance_key) return instance._eager_forward(x) @_wrapped_dropout.register_fake -def _(x: torch.Tensor, instance_key: int) -> torch.Tensor: +def _(x: torch.Tensor, instance_key: str) -> torch.Tensor: return torch.empty_like(x) diff --git a/tileops/ops/elementwise/_base.py b/tileops/ops/elementwise/_base.py index 29c1c3b99..166cff757 100644 --- a/tileops/ops/elementwise/_base.py +++ b/tileops/ops/elementwise/_base.py @@ -9,7 +9,7 @@ - Concrete ops are registered via @torch.library.custom_op at package load time - Three factory functions (_register_unary_custom_op, _register_binary_custom_op, _register_fused_gated_custom_op) register every op; instances are looked up at - runtime via _OP_REGISTRY keyed by id(instance) + runtime via the shared instance registry in tileops.ops.compile_boundary Utility: - broadcast_out_shape: PyTorch broadcast output shape of two operand shapes @@ -21,7 +21,6 @@ import functools import inspect import math -import weakref from math import prod from typing import Callable, Dict, List, Optional @@ -31,13 +30,13 @@ from tileops.manifest import load_manifest from tileops.manifest.dtype_rules import promote_int_to_float_ref, same_as_ref +from ..compile_boundary import get_instance from ..op_base import Op # torch.compile registration factories (see module docstring). The registry # key is a plain int so dynamo can trace through forward() without hitting # unsupported Python side-effects. -_OP_REGISTRY: weakref.WeakValueDictionary = weakref.WeakValueDictionary() _MANIFEST_INT_SCALAR_DTYPES = ( torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64, @@ -142,12 +141,12 @@ def _register_unary_custom_op(op_cls, output_dtype_override=None): op_name = op_cls._op_name @torch.library.custom_op(f"top::elementwise_unary_{op_name}", mutates_args=()) - def _wrapped(x: torch.Tensor, instance_key: int) -> torch.Tensor: - instance = _OP_REGISTRY[instance_key] + def _wrapped(x: torch.Tensor, instance_key: str) -> torch.Tensor: + instance = get_instance(instance_key) return instance._eager_forward(x) @_wrapped.register_fake - def _(x: torch.Tensor, instance_key: int) -> torch.Tensor: + def _(x: torch.Tensor, instance_key: str) -> torch.Tensor: out_dtype = output_dtype_override if output_dtype_override is not None else x.dtype return torch.empty_like(x, dtype=out_dtype) @@ -169,8 +168,8 @@ def _register_unary_inplace_custom_op(op_cls): @torch.library.custom_op( f"top::elementwise_unary_{op_name}_inplace", mutates_args=("x",), ) - def _wrapped_inplace(x: torch.Tensor, instance_key: int) -> None: - instance = _OP_REGISTRY[instance_key] + def _wrapped_inplace(x: torch.Tensor, instance_key: str) -> None: + instance = get_instance(instance_key) result = instance._eager_forward(x) x.copy_(result.reshape(x.shape)) @@ -191,9 +190,9 @@ def _wrapped( a: torch.Tensor, b: torch.Tensor, out_shape: List[int], - instance_key: int, + instance_key: str, ) -> torch.Tensor: - instance = _OP_REGISTRY[instance_key] + instance = get_instance(instance_key) return instance._eager_forward(a, b) @_wrapped.register_fake @@ -201,7 +200,7 @@ def _( a: torch.Tensor, b: torch.Tensor, out_shape: List[int], - instance_key: int, + instance_key: str, ) -> torch.Tensor: out_dtype = torch.bool if output_bool else a.dtype return a.new_empty(out_shape, dtype=out_dtype) @@ -217,16 +216,16 @@ def _register_prelu_custom_op(op_cls): def _wrapped( x: torch.Tensor, weight: torch.Tensor, - instance_key: int, + instance_key: str, ) -> torch.Tensor: - instance = _OP_REGISTRY[instance_key] + instance = get_instance(instance_key) return instance._eager_forward(x, weight) @_wrapped.register_fake def _( x: torch.Tensor, weight: torch.Tensor, - instance_key: int, + instance_key: str, ) -> torch.Tensor: return torch.empty_like(x) @@ -247,9 +246,9 @@ def _wrapped( cond: torch.Tensor, x: torch.Tensor, y: torch.Tensor, - instance_key: int, + instance_key: str, ) -> torch.Tensor: - instance = _OP_REGISTRY[instance_key] + instance = get_instance(instance_key) return instance._eager_forward(cond, x, y) @_wrapped.register_fake @@ -257,7 +256,7 @@ def _( cond: torch.Tensor, x: torch.Tensor, y: torch.Tensor, - instance_key: int, + instance_key: str, ) -> torch.Tensor: out_shape = torch.broadcast_shapes(cond.shape, x.shape, y.shape) return x.new_empty(out_shape) @@ -282,9 +281,9 @@ def _wrapped( input: torch.Tensor, end: torch.Tensor, weight: torch.Tensor, - instance_key: int, + instance_key: str, ) -> torch.Tensor: - instance = _OP_REGISTRY[instance_key] + instance = get_instance(instance_key) return instance._eager_forward(input, end, weight) @_wrapped.register_fake @@ -292,7 +291,7 @@ def _( input: torch.Tensor, end: torch.Tensor, weight: torch.Tensor, - instance_key: int, + instance_key: str, ) -> torch.Tensor: out_shape = torch.broadcast_shapes(input.shape, end.shape, weight.shape) return input.new_empty(out_shape) @@ -313,16 +312,16 @@ def _register_masked_fill_custom_op(op_cls): def _wrapped( x: torch.Tensor, mask: torch.Tensor, - instance_key: int, + instance_key: str, ) -> torch.Tensor: - instance = _OP_REGISTRY[instance_key] + instance = get_instance(instance_key) return instance._eager_forward(x, mask) @_wrapped.register_fake def _( x: torch.Tensor, mask: torch.Tensor, - instance_key: int, + instance_key: str, ) -> torch.Tensor: out_shape = torch.broadcast_shapes(x.shape, mask.shape) return x.new_empty(out_shape) @@ -346,9 +345,9 @@ def _wrapped( input: torch.Tensor, mask: torch.Tensor, value: torch.Tensor, - instance_key: int, + instance_key: str, ) -> torch.Tensor: - instance = _OP_REGISTRY[instance_key] + instance = get_instance(instance_key) return instance._eager_forward(input, mask, value) @_wrapped.register_fake @@ -356,7 +355,7 @@ def _( input: torch.Tensor, mask: torch.Tensor, value: torch.Tensor, - instance_key: int, + instance_key: str, ) -> torch.Tensor: out_shape = torch.broadcast_shapes(input.shape, mask.shape) return input.new_empty(out_shape) @@ -384,9 +383,9 @@ def _wrapped( input: torch.Tensor, min: Optional[torch.Tensor], max: Optional[torch.Tensor], - instance_key: int, + instance_key: str, ) -> torch.Tensor: - instance = _OP_REGISTRY[instance_key] + instance = get_instance(instance_key) return instance._eager_forward(input, min, max) @_wrapped.register_fake @@ -394,7 +393,7 @@ def _( input: torch.Tensor, min: Optional[torch.Tensor], max: Optional[torch.Tensor], - instance_key: int, + instance_key: str, ) -> torch.Tensor: shapes = [input.shape] if min is not None: @@ -415,16 +414,16 @@ def _register_clamp_min_custom_op(op_cls): def _wrapped( input: torch.Tensor, min: torch.Tensor, - instance_key: int, + instance_key: str, ) -> torch.Tensor: - instance = _OP_REGISTRY[instance_key] + instance = get_instance(instance_key) return instance._eager_forward(input, min) @_wrapped.register_fake def _( input: torch.Tensor, min: torch.Tensor, - instance_key: int, + instance_key: str, ) -> torch.Tensor: out_shape = torch.broadcast_shapes(input.shape, min.shape) return input.new_empty(out_shape) @@ -440,16 +439,16 @@ def _register_clamp_max_custom_op(op_cls): def _wrapped( input: torch.Tensor, max: torch.Tensor, - instance_key: int, + instance_key: str, ) -> torch.Tensor: - instance = _OP_REGISTRY[instance_key] + instance = get_instance(instance_key) return instance._eager_forward(input, max) @_wrapped.register_fake def _( input: torch.Tensor, max: torch.Tensor, - instance_key: int, + instance_key: str, ) -> torch.Tensor: out_shape = torch.broadcast_shapes(input.shape, max.shape) return input.new_empty(out_shape) @@ -470,9 +469,9 @@ def _wrapped( x: torch.Tensor, M: int, N: int, - instance_key: int, + instance_key: str, ) -> torch.Tensor: - instance = _OP_REGISTRY[instance_key] + instance = get_instance(instance_key) return instance._eager_forward(x) @_wrapped.register_fake @@ -480,7 +479,7 @@ def _( x: torch.Tensor, M: int, N: int, - instance_key: int, + instance_key: str, ) -> torch.Tensor: return x.new_empty((M, N), dtype=x.dtype) @@ -612,9 +611,6 @@ def __init__( N_total=N_total, dtype=dtype, tune=tune, ) self.output_dtype = self._resolve_output_dtype() - # Register in global registry for torch.compile dispatch - self._instance_key = id(self) - _OP_REGISTRY[self._instance_key] = self def _build_kernel_instance( self, @@ -757,9 +753,6 @@ def __init__( self.b_numel = prod(b_shape) self.dispatch_kernel(kernel_map) self.kernel = self._build_kernel_instance(tune) - # Register in global registry for torch.compile dispatch - self._instance_key = id(self) - _OP_REGISTRY[self._instance_key] = self def _build_kernel_instance(self, tune): """Construct the kernel. Subclasses override to inject extra kwargs.""" @@ -862,9 +855,6 @@ def __init__( self._kernel_key = None if M is not None and N is not None and dtype is not None: self._ensure_kernel(M, N, dtype) - # Register in global registry for torch.compile dispatch - self._instance_key = id(self) - _OP_REGISTRY[self._instance_key] = self @property def default_kernel_map(self) -> Dict[str, Kernel]: @@ -1013,7 +1003,7 @@ class _ParametricActivationOp(_UnaryActivationMixin, UnaryOp): ``__init__`` (scalar parameter names and defaults vary per leaf): each leaf validates its scalars, populates ``self.`` for introspection, instantiates ``self.kernel`` with typed kwargs, and - registers itself with ``_OP_REGISTRY`` via the + registers itself for compile dispatch via the ``_finalize_init`` helper. ``UnaryOp.__init__`` is intentionally bypassed; ``_finalize_init`` performs the equivalent state setup. @@ -1035,7 +1025,7 @@ def _finalize_init( The leaf has already called ``self.dispatch_kernel(kernel_map)`` and instantiated its kernel directly with typed kwargs. This helper records the kernel on ``self`` and runs the - ``_OP_REGISTRY`` registration shared by every parametric leaf. + instance registration shared by every parametric leaf. """ self.N_total = N_total self.dtype = dtype @@ -1044,8 +1034,6 @@ def _finalize_init( # Surface ``output_dtype`` for ``total_memory`` accounting, as # ``UnaryOp.__init__`` does. self.output_dtype = resolve_output_dtype(type(self).__name__, dtype) - self._instance_key = id(self) - _OP_REGISTRY[self._instance_key] = self class _AlphaScaledBinaryOp(BinaryOp): @@ -1175,18 +1163,17 @@ def __init__( if dtype in type(self)._fallback_dtypes: self.N_total = N_total self.dtype = dtype - # No kernel is constructed — it is float-only. The kernel_map still - # goes through the shared install path so an override is arch-checked - # the same way as on the float path. - self._install_kernel_map(kernel_map) + # No kernel is constructed — it is float-only. Routing through + # dispatch_kernel keeps the arch check identical to the float path + # and registers the instance for the compile boundary, which a bare + # _install_kernel_map call would skip. + self.dispatch_kernel(kernel_map) self.kernel = None self.output_dtype = ( type(self)._int_output_dtype if type(self)._int_output_dtype is not None else dtype ) - self._instance_key = id(self) - _OP_REGISTRY[self._instance_key] = self return super().__init__(N_total, dtype, kernel_map=kernel_map, tune=tune) diff --git a/tileops/ops/elementwise/arithmetic.py b/tileops/ops/elementwise/arithmetic.py index fe4fc1e3b..c5b6d8f4c 100644 --- a/tileops/ops/elementwise/arithmetic.py +++ b/tileops/ops/elementwise/arithmetic.py @@ -22,7 +22,7 @@ from tileops.kernels.kernel_base import Kernel from ..op_base import Op -from ._base import _OP_REGISTRY, BinaryOp, _AlphaScaledBinaryOp +from ._base import BinaryOp, _AlphaScaledBinaryOp class AddFwdOp(_AlphaScaledBinaryOp): @@ -240,8 +240,6 @@ def __init__( self.kernel = self.kernel_map[self._op_name]( self.N_total, dtype, tune=tune, ) - self._instance_key = id(self) - _OP_REGISTRY[self._instance_key] = self @property def default_kernel_map(self) -> Dict[str, Kernel]: diff --git a/tileops/ops/elementwise/clamp.py b/tileops/ops/elementwise/clamp.py index 978574b94..0a75a4b18 100644 --- a/tileops/ops/elementwise/clamp.py +++ b/tileops/ops/elementwise/clamp.py @@ -10,7 +10,6 @@ from ..op_base import Op from ._base import ( - _OP_REGISTRY, _ClampTensorBase, _validate_scalar_param_repr, ) @@ -73,8 +72,6 @@ def __init__( has_max=self.max_shape is not None, tune=tune, ) - self._instance_key = id(self) - _OP_REGISTRY[self._instance_key] = self @property def default_kernel_map(self): @@ -167,8 +164,6 @@ def __init__( self.kernel = self.kernel_map["clamp_tensor"]( self.N_total, dtype, has_min=True, has_max=False, tune=tune, ) - self._instance_key = id(self) - _OP_REGISTRY[self._instance_key] = self @property def default_kernel_map(self): @@ -236,8 +231,6 @@ def __init__( self.kernel = self.kernel_map["clamp_tensor"]( self.N_total, dtype, has_min=False, has_max=True, tune=tune, ) - self._instance_key = id(self) - _OP_REGISTRY[self._instance_key] = self @property def default_kernel_map(self): @@ -316,8 +309,6 @@ def __init__( self.kernel = self.kernel_map["clamp"]( self.N_total, dtype, min_val=min, max_val=max, tune=tune, ) - self._instance_key = id(self) - _OP_REGISTRY[self._instance_key] = self @property def default_kernel_map(self): diff --git a/tileops/ops/elementwise/masked_fill.py b/tileops/ops/elementwise/masked_fill.py index 63ebe5ca9..e52a2c1ff 100644 --- a/tileops/ops/elementwise/masked_fill.py +++ b/tileops/ops/elementwise/masked_fill.py @@ -12,7 +12,7 @@ from tileops.kernels.kernel_base import Kernel from ..op_base import Op -from ._base import _OP_REGISTRY, _validate_scalar_param_repr +from ._base import _validate_scalar_param_repr class MaskedFillFwdOp(Op): @@ -75,8 +75,6 @@ def __init__( self._bool_storage = dtype == torch.bool kernel_dtype = torch.uint8 if self._bool_storage else dtype self.kernel = self.kernel_map["masked_fill_tensor_value"](self.N_total, kernel_dtype) - self._instance_key = id(self) - _OP_REGISTRY[self._instance_key] = self @property def default_kernel_map(self): @@ -204,8 +202,6 @@ def __init__( self.kernel = self.kernel_map["masked_fill"]( self.N_total, kernel_dtype, kernel_value, ) - self._instance_key = id(self) - _OP_REGISTRY[self._instance_key] = self @property def default_kernel_map(self): diff --git a/tileops/ops/elementwise/nan_to_num.py b/tileops/ops/elementwise/nan_to_num.py index 59a6f64cf..367b5f348 100644 --- a/tileops/ops/elementwise/nan_to_num.py +++ b/tileops/ops/elementwise/nan_to_num.py @@ -8,7 +8,7 @@ from tileops.kernels.kernel_base import Kernel from ..op_base import Op -from ._base import _OP_REGISTRY, _validate_scalar_param_repr +from ._base import _validate_scalar_param_repr class NanToNumFwdOp(Op): @@ -82,8 +82,6 @@ def __init__( self.kernel = self.kernel_map["nan_to_num"]( N_total, dtype, nan, kernel_posinf, kernel_neginf, tune=tune, ) - self._instance_key = id(self) - _OP_REGISTRY[self._instance_key] = self @property def default_kernel_map(self): diff --git a/tileops/ops/elementwise/prelu.py b/tileops/ops/elementwise/prelu.py index ab97f1ddc..c05171a98 100644 --- a/tileops/ops/elementwise/prelu.py +++ b/tileops/ops/elementwise/prelu.py @@ -9,7 +9,6 @@ from tileops.kernels.kernel_base import Kernel from ..op_base import Op -from ._base import _OP_REGISTRY class PreluFwdOp(Op): @@ -53,8 +52,6 @@ def __init__( self.inner_size = inner_size self.dispatch_kernel(kernel_map) self.kernel = self.kernel_map[self._op_name](N_total, num_channels, inner_size, dtype) - self._instance_key = id(self) - _OP_REGISTRY[self._instance_key] = self @property def default_kernel_map(self): diff --git a/tileops/ops/elementwise/where.py b/tileops/ops/elementwise/where.py index 4759af4be..e962c0fa7 100644 --- a/tileops/ops/elementwise/where.py +++ b/tileops/ops/elementwise/where.py @@ -9,7 +9,6 @@ from tileops.kernels.kernel_base import Kernel from ..op_base import Op -from ._base import _OP_REGISTRY class WhereFwdOp(Op): @@ -64,8 +63,6 @@ def __init__( self.N_total = prod(self.out_shape) if self.out_shape else 1 self.dispatch_kernel(kernel_map) self.kernel = self.kernel_map[self._op_name](self.N_total, dtype) - self._instance_key = id(self) - _OP_REGISTRY[self._instance_key] = self @property def default_kernel_map(self): diff --git a/tileops/ops/rope.py b/tileops/ops/rope.py index b95ac80be..717ba958c 100644 --- a/tileops/ops/rope.py +++ b/tileops/ops/rope.py @@ -18,12 +18,12 @@ torch.compile support: - All 5 concrete ops are registered via @torch.library.custom_op at module load time. A factory function (_register_rope_custom_op) registers every - op; instances are looked up at runtime via _OP_REGISTRY keyed by + op; instances are looked up at runtime through the shared registry in + tileops.ops.compile_boundary, keyed by id(instance). """ import math -import weakref from typing import Dict, Optional import torch @@ -38,12 +38,12 @@ RopeYarnKernel, ) +from .compile_boundary import get_instance from .op_base import Op # torch.compile registration factory: a @torch.library.custom_op + # register_fake pair per RoPE op (see module docstring). -_OP_REGISTRY: weakref.WeakValueDictionary = weakref.WeakValueDictionary() def _register_rope_custom_op(op_cls): @@ -55,12 +55,12 @@ def _register_rope_custom_op(op_cls): op_name = op_cls._op_name @torch.library.custom_op(f"top::rope_{op_name}", mutates_args=()) - def _wrapped(x: torch.Tensor, instance_key: int) -> torch.Tensor: - instance = _OP_REGISTRY[instance_key] + def _wrapped(x: torch.Tensor, instance_key: str) -> torch.Tensor: + instance = get_instance(instance_key) return instance._eager_forward(x) @_wrapped.register_fake - def _(x: torch.Tensor, instance_key: int) -> torch.Tensor: + def _(x: torch.Tensor, instance_key: str) -> torch.Tensor: return torch.empty_like(x) op_cls._wrapped = _wrapped @@ -72,12 +72,12 @@ def _register_rope_position_ids_custom_op(op_cls): @torch.library.custom_op(f"top::rope_{op_name}", mutates_args=()) def _wrapped(x: torch.Tensor, position_ids: torch.Tensor, - instance_key: int) -> torch.Tensor: - instance = _OP_REGISTRY[instance_key] + instance_key: str) -> torch.Tensor: + instance = get_instance(instance_key) return instance._eager_forward(x, position_ids) @_wrapped.register_fake - def _(x: torch.Tensor, position_ids: torch.Tensor, instance_key: int) -> torch.Tensor: + def _(x: torch.Tensor, position_ids: torch.Tensor, instance_key: str) -> torch.Tensor: return torch.empty_like(x) op_cls._wrapped = _wrapped @@ -362,9 +362,6 @@ def __init__( self._kernel_cache: Dict[tuple, Kernel] = {} self.kernel = None - # Register in global registry for torch.compile dispatch - self._instance_key = id(self) - _OP_REGISTRY[self._instance_key] = self def _get_cos_sin(self, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: """Return cached cos/sin tables, recomputing if device changed.""" @@ -555,8 +552,6 @@ def __init__( self._kernel_cache: Dict[tuple, Kernel] = {} self.kernel = None - self._instance_key = id(self) - _OP_REGISTRY[self._instance_key] = self @property def default_kernel_map(self) -> Dict[str, Kernel]: From 9680783f522ec06c4655bbadeb92a70268062d02 Mon Sep 17 00:00:00 2001 From: lcy-seso Date: Tue, 4 Aug 2026 21:42:00 +0800 Subject: [PATCH 11/15] [Refactor][Manifest] Drop the engram dtype params and compare dtype defaults as dtypes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three sequence-modeling entries still declared a `dtype` param after the engram ops stopped taking one. Every output of all three is `same_as()`, so the element type was never theirs to choose. This was the only real CI failure: `validate-manifest`, `compile-contract-gate` and `gpu-smoke` all assert the validator exits 0, so one signature divergence failed three checks. The validator also compared a `torch.dtype` default against its YAML spelling with `!=`, so `torch.float16 != "float16"` made every dtype default a mismatch — advisory locally, blocking under `--strict` in CI. No entry could declare a dtype default at all before this. Where the declared type is `torch.dtype`, the default is now resolved before comparison, with a test covering both the matching and the genuinely mismatched case. --- scripts/validate_manifest.py | 13 ++++++++++++ tests/test_validate_manifest.py | 28 +++++++++++++++++++++++++ tileops/manifest/sequence_modeling.yaml | 3 --- 3 files changed, 41 insertions(+), 3 deletions(-) diff --git a/scripts/validate_manifest.py b/scripts/validate_manifest.py index 85ad4c3c6..e87a3c51e 100755 --- a/scripts/validate_manifest.py +++ b/scripts/validate_manifest.py @@ -3176,6 +3176,19 @@ def check_c3_ctor_signature_parity( manifest_has_default = ( manifest_default is not _MISSING and manifest_default != "REQUIRED" ) + # A ``torch.dtype`` default is spelled by name in YAML ("float16"), + # so compare the resolved dtype rather than the spelling. Without this + # no entry can declare a dtype default at all. + if ( + manifest_has_default + and pattrs.get("type") == "torch.dtype" + and isinstance(manifest_default, str) + ): + import torch + + resolved = getattr(torch, manifest_default, None) + if isinstance(resolved, torch.dtype): + manifest_default = resolved compat_default = pattrs.get("compat_default", _MISSING) manifest_has_compat_default = ( compat_default is not _MISSING and not manifest_has_default diff --git a/tests/test_validate_manifest.py b/tests/test_validate_manifest.py index 46d608a41..8e56b65b7 100644 --- a/tests/test_validate_manifest.py +++ b/tests/test_validate_manifest.py @@ -12,6 +12,7 @@ from pathlib import Path import pytest +import torch pytestmark = pytest.mark.smoke @@ -2586,6 +2587,33 @@ def test_matching_defaults_pass(self, validator): "OpCompat", entry, cls ) == [] + def test_dtype_default_compared_as_dtype(self, validator): + """A ``torch.dtype`` default is spelled by name in YAML. + + ``torch.float16 != "float16"``, so comparing the spellings would make + every dtype default a mismatch and leave the manifest unable to declare + one at all. + """ + cls = _strict_op( + "OpDtypeDefault", + init=lambda self, dtype=torch.float16, kernel_map=None: None, + ) + entry = {"signature": { + "params": {"dtype": {"type": "torch.dtype", "default": "float16"}}, + }} + assert validator.check_c3_ctor_signature_parity( + "OpDtypeDefault", entry, cls, + ) == [] + + # A genuinely different dtype still fails. + entry_bad = {"signature": { + "params": {"dtype": {"type": "torch.dtype", "default": "bfloat16"}}, + }} + errs = validator.check_c3_ctor_signature_parity( + "OpDtypeDefault", entry_bad, cls, + ) + assert any("default mismatch" in e for e in errs), errs + def test_ctor_mismatches_fail(self, validator): """Case table: missing default, compat_default mismatch, kw-only.""" cases = [ diff --git a/tileops/manifest/sequence_modeling.yaml b/tileops/manifest/sequence_modeling.yaml index 3374c422f..6a0ff629d 100644 --- a/tileops/manifest/sequence_modeling.yaml +++ b/tileops/manifest/sequence_modeling.yaml @@ -25,7 +25,6 @@ EngramGateConvFwdOp: M: {type: int} seq_len: {type: int} d: {type: int} - dtype: {type: torch.dtype} eps: {type: float, default: 1.0e-6} shape_rules: - "H.shape == (M, seq_len, d)" @@ -83,7 +82,6 @@ EngramGateConvBwdOp: M: {type: int} seq_len: {type: int} d: {type: int} - dtype: {type: torch.dtype} eps: {type: float, default: 1.0e-6} shape_rules: - "dY.shape == (M, seq_len, d)" @@ -139,7 +137,6 @@ EngramDecodeOp: max_conv_len: {type: int} conv_kernel_size: {type: int} dilation: {type: int} - dtype: {type: torch.dtype} eps: {type: float, default: 1.0e-6} shape_rules: - "e_t.shape == (batch, d_mem)" From 0aa0b890ee5d57535da663aa6f487a8638201539 Mon Sep 17 00:00:00 2001 From: lcy-seso Date: Wed, 5 Aug 2026 09:47:14 +0800 Subject: [PATCH 12/15] [Refactor][Ops] Remove the last singleton dtype slots and trim the new comments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An AST sweep over the 48 migrated classes for `self.dtype`, `self.kernel` and `self.output_dtype` read on an execution path found two survivors. Neither returns a wrong answer today — both are written a few lines earlier in the same call — but both are the slot pattern that made `total_memory` silently report 4-byte elements, so they go: the varlen prefill roofline reads `q.dtype` directly, and `LayerNormFwdOp` keeps its kernel in a local instead of assigning `self.kernel` per call. `test_layer_norm_rebuilds_kernel_on_m_change` asserted through that slot; it now asserts on the cache, which is what it was really testing — and additionally that the first entry survives. Three comments introduced by this branch restated their own code or ran long; they are cut to the constraint that is not obvious from reading it. --- tests/ops/test_layer_norm.py | 7 ++++--- tileops/ops/_dtype_codegen.py | 6 ++---- tileops/ops/attention/gqa.py | 2 +- tileops/ops/elementwise/_base.py | 7 +++---- tileops/ops/norm/layer_norm.py | 4 ++-- tileops/ops/pool.py | 2 -- 6 files changed, 12 insertions(+), 16 deletions(-) diff --git a/tests/ops/test_layer_norm.py b/tests/ops/test_layer_norm.py index 22fac180c..f5cd0d639 100644 --- a/tests/ops/test_layer_norm.py +++ b/tests/ops/test_layer_norm.py @@ -215,14 +215,15 @@ def test_layer_norm_rebuilds_kernel_on_m_change() -> None: x1 = torch.randn(512, n, dtype=dtype, device="cuda") y1 = op(x1, weight, bias) - first_kernel = op.kernel + first_kernel = op._kernel_cache[(512, dtype)] assert y1.shape == x1.shape x2 = torch.randn(1024, n, dtype=dtype, device="cuda") y2 = op(x2, weight, bias) assert y2.shape == x2.shape - # Kernel should have been rebuilt for the new M. - assert op.kernel is not first_kernel + # A separate cache entry for the new M, and the first one is still there. + assert op._kernel_cache[(1024, dtype)] is not first_kernel + assert op._kernel_cache[(512, dtype)] is first_kernel y_ref = F.layer_norm( x2.float(), (n,), diff --git a/tileops/ops/_dtype_codegen.py b/tileops/ops/_dtype_codegen.py index 588f0c342..d27b577d6 100644 --- a/tileops/ops/_dtype_codegen.py +++ b/tileops/ops/_dtype_codegen.py @@ -247,10 +247,8 @@ def synthesize_validate_dtypes( "op_name": op_name, } params_src = ", ".join(input_names) - # Unrolled per input, with each parameter referenced by name. A `locals()` - # lookup or a loop over `input_names` would read the same values, but - # `torch.compile` cannot trace `locals()`, and this body runs inside - # `forward()` — so an op that validates dtypes would lose `fullgraph`. + # Unrolled per input rather than looping via `locals()`: dynamo cannot + # trace `locals()`, and this body runs inside `forward()`. src_lines = [ f"def _validate_dtypes(self, {params_src}):", f' """Synthesized from manifest signature for {op_name}."""', diff --git a/tileops/ops/attention/gqa.py b/tileops/ops/attention/gqa.py index dc4f95dc7..78aa4a030 100644 --- a/tileops/ops/attention/gqa.py +++ b/tileops/ops/attention/gqa.py @@ -660,7 +660,7 @@ def _record_roofline( "cu_seqlens_q": cu_seqlens_q, "cu_seqlens_kv": cu_seqlens_kv, "is_causal": self.is_causal, - "dtype": self.dtype, + "dtype": q.dtype, } def forward( diff --git a/tileops/ops/elementwise/_base.py b/tileops/ops/elementwise/_base.py index 166cff757..13f9552f9 100644 --- a/tileops/ops/elementwise/_base.py +++ b/tileops/ops/elementwise/_base.py @@ -1163,10 +1163,9 @@ def __init__( if dtype in type(self)._fallback_dtypes: self.N_total = N_total self.dtype = dtype - # No kernel is constructed — it is float-only. Routing through - # dispatch_kernel keeps the arch check identical to the float path - # and registers the instance for the compile boundary, which a bare - # _install_kernel_map call would skip. + # No kernel is constructed — it is float-only. dispatch_kernel, not + # _install_kernel_map: only the former registers the instance for + # the compile boundary. self.dispatch_kernel(kernel_map) self.kernel = None self.output_dtype = ( diff --git a/tileops/ops/norm/layer_norm.py b/tileops/ops/norm/layer_norm.py index 3028d1d23..b409a6fd0 100644 --- a/tileops/ops/norm/layer_norm.py +++ b/tileops/ops/norm/layer_norm.py @@ -144,9 +144,9 @@ def forward( self._kernel_cache[key] = self.kernel_map["layer_norm"]( m_actual, self.N, self.eps, x.dtype, tune=self.tune, ) - self.kernel = self._kernel_cache[key] + kernel = self._kernel_cache[key] self._last_m = m_actual - y = self.kernel(x, weight, bias) + y = kernel(x, weight, bias) return y.reshape(orig_shape) diff --git a/tileops/ops/pool.py b/tileops/ops/pool.py index 48413fe15..a4f0d50d5 100644 --- a/tileops/ops/pool.py +++ b/tileops/ops/pool.py @@ -67,8 +67,6 @@ def __init__( for key, value in params.items(): setattr(self, key, value) - # Every kernel argument except the element type, which the first - # forward() supplies from its input. self._kernel_params = params self.dispatch_kernel(kernel_map) self._kernel_cache: Dict[torch.dtype, Kernel] = {} From c3c590f3f459b9ce20b91268657e4011bd745ff2 Mon Sep 17 00:00:00 2001 From: lcy-seso Date: Wed, 5 Aug 2026 09:51:46 +0800 Subject: [PATCH 13/15] [Fix][Ops] Autotune the kernels an op keeps in a cache MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `Op.autotune` searched `dir(self)` for `Kernel`-typed attributes. Ops that build at forward time keep their kernels in a dict keyed by shape and dtype, so the search found nothing and the call silently did nothing. The hole predates this branch — the reduction family was already dict-cached — but this branch moved roughly thirty more ops off a single `self.kernel`, so it widened from one family to most of them. The search now descends through dict values and sequence items, two levels deep, which is what a backward op needs when it caches its preprocess and backward kernels as one entry. `__dict__` is skipped and kernels are deduplicated by identity: without that, a kernel bound as an attribute is tuned twice, which the existing test caught. Also corrects the `Op.kernel` contract in `ops-design-reference.md` and the class docstring, both of which still described it as the kernel `forward()` uses. --- docs/design/ops-design-reference.md | 2 +- tests/test_op_base.py | 35 ++++++++++++++++++++++++ tileops/ops/op_base.py | 41 +++++++++++++++++++++++++---- 3 files changed, 72 insertions(+), 6 deletions(-) diff --git a/docs/design/ops-design-reference.md b/docs/design/ops-design-reference.md index 276b78cbc..424f72c97 100644 --- a/docs/design/ops-design-reference.md +++ b/docs/design/ops-design-reference.md @@ -280,7 +280,7 @@ Per-family protocol variables, declared by L2 bases and overridden by L3 ops. | Attribute | Type | Purpose | | -------------- | ------------------------------------ | -------------------------------------------------------------------------------------------- | -| `kernel` | `Kernel` | Kernel instance used by `forward()` | +| `kernel` | `Kernel` | Set only by an op that holds one kernel; an op that builds per dtype uses a cache instead | | `kernel_map` | `Optional[Dict[str, Kernel]]` | Dispatched kernels keyed by name | | `dtype` | `Optional[torch.dtype]` | Dtype of the most recent `forward()`; `None` before the first one | | `device` | `Optional[Union[torch.device, str]]` | Device (default `'cuda'`) | diff --git a/tests/test_op_base.py b/tests/test_op_base.py index b4db1745d..bfc27e6fc 100644 --- a/tests/test_op_base.py +++ b/tests/test_op_base.py @@ -8,6 +8,7 @@ import warnings import pytest +import torch from tileops.kernels.kernel_base import Kernel from tileops.ops import op_base @@ -184,3 +185,37 @@ def default_kernel_map(self): FakeOp().autotune() assert sorted(tuned) == ["k1", "k2"] + + def test_autotune_reaches_kernels_held_in_a_cache(self): + """Ops that build at forward time keep kernels in a dict, not a slot.""" + tuned: list[str] = [] + + class FakeKernel(Kernel): + def __init__(self, name): + super().__init__() + self.name = name + + def forward(self): + return None + + def autotune(self, warmup=25, rep=50): + tuned.append(self.name) + + class CachingOp(Op): + def __init__(self): + self._kernel_cache = { + ((8,), torch.float16): FakeKernel("fp16"), + ((8,), torch.bfloat16): FakeKernel("bf16"), + } + # A backward op caches its kernels as one entry per dtype. + self._pair_cache = {torch.float16: (FakeKernel("pre"), FakeKernel("bwd"))} + + def forward(self, *a, **kw): + return None + + @property + def default_kernel_map(self): + return {} + + CachingOp().autotune() + assert sorted(tuned) == ["bf16", "bwd", "fp16", "pre"] diff --git a/tileops/ops/op_base.py b/tileops/ops/op_base.py index 77e2e47ea..a004943a6 100644 --- a/tileops/ops/op_base.py +++ b/tileops/ops/op_base.py @@ -12,6 +12,24 @@ # Module-level dedup for empty-static_dims warnings; keyed by Op subclass. _EMPTY_STATIC_DIMS_WARNED: set = set() +def _iter_kernels(value: object, _depth: int = 0) -> "list[Kernel]": + """Return the kernels *value* holds, descending through dicts and sequences. + + A kernel cache is a dict keyed by shape and dtype; an op that builds + several kernels together (preprocess plus backward, say) stores a tuple + per entry, so the search has to go two levels down. + """ + if isinstance(value, Kernel): + return [value] + if _depth >= 2: + return [] + if isinstance(value, dict): + return [k for v in value.values() for k in _iter_kernels(v, _depth + 1)] + if isinstance(value, (tuple, list)): + return [k for v in value for k in _iter_kernels(v, _depth + 1)] + return [] + + class Op(ABC): """Base class for TileOPs operations. @@ -30,7 +48,8 @@ class Op(ABC): >>> latency = op.profile() # Benchmark performance Attributes: - kernel: top.Kernel instance (e.g. mha_fwd_kernel) + kernel: single kernel, for ops that hold one; ops that build per + dtype keep a cache instead dtype: Data type for computation (e.g., torch.float16) device: Device for computation (e.g., 'cuda') input_shapes: Expected input tensor shapes @@ -170,11 +189,23 @@ def dispatch_kernel(self, kernel_map: Optional[dict[str, Kernel]] = None) -> Non self._instance_key = register_instance(self) def autotune(self) -> None: - """Autotune all kernels of the op""" + """Autotune every kernel the op holds. + + An op that builds its kernels at forward time keeps them in a cache + keyed by shape and dtype, so the search covers the values of dict + attributes and the items of tuple/list attributes as well as kernels + bound directly. Only kernels already built are tuned; a cache the op + has not populated yet has nothing to tune. + """ + seen: set[int] = set() for attr_name in dir(self): - attr = getattr(self, attr_name) - if isinstance(attr, Kernel): - attr.autotune() + # ``__dict__`` would re-yield every kernel bound as an attribute. + if attr_name.startswith("__"): + continue + for kernel in _iter_kernels(getattr(self, attr_name, None)): + if id(kernel) not in seen: + seen.add(id(kernel)) + kernel.autotune() @abstractmethod def forward(self, *args: object, **kwargs: object) -> Union[torch.Tensor, tuple]: From 28739b7684934f0ea347cfc9a9aa709596c271ce Mon Sep 17 00:00:00 2001 From: lcy-seso Date: Wed, 5 Aug 2026 10:02:51 +0800 Subject: [PATCH 14/15] [Fix][Ops] Validate every input, not just the anchor, before selecting a kernel MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review findings from the closing pass. Seven migrated ops chose their kernel from one tensor's dtype and never checked the others, while the manifest declares them `same_as` the anchor. An fp16 `q` with a bf16 `k` reached a kernel compiled for fp16 and reinterpreted the storage. All seven have a synthesized `_validate_dtypes`; they now call it before selecting. `FusedGatedOp.forward` built its kernel before crossing the dispatch boundary. The violation predates this branch, but nothing reached it: the constructor took a dtype and prebuilt the kernel, so the call always hit the cache. Dropping `dtype` from `build_activation_op` opened the cold path for the first time. `forward` now passes locally derived M/N and leaves construction to `_eager_forward`; a cold `fullgraph=True` call on a dtype-inferred activation op is verified. Four norm ops recorded `_last_roofline_mn` but never rebound `self.dtype`, so `eval_roofline` raised after a successful forward. The design docs claimed dtype is *never* a constructor parameter, which the seven intentional exemptions contradict; both now state the actual rule and what "determined by the inputs" means. `CBProducerOp`'s manifest still declared the removed parameter — `spec-only` kept the validator quiet. Adds the invariant the suite was missing: one instance, two dtypes, two cache entries. Every existing test builds a fresh op per dtype, so nothing covered a second dtype arriving at an instance that already had an entry. Drops two dead fragments: an NSA filter for a key that can no longer be present, and `per_input` in a closure the unrolled body stopped reading. --- docs/design/ops-design-reference.md | 2 +- docs/design/ops-design.md | 8 +- tests/ops/test_multi_dtype_instance.py | 78 +++++++++++++++++++ tileops/manifest/mamba.yaml | 1 - tileops/ops/_dtype_codegen.py | 1 - tileops/ops/attention/deepseek_dsa.py | 1 + tileops/ops/attention/deepseek_mla.py | 1 + tileops/ops/attention/deepseek_nsa.py | 6 +- tileops/ops/attention/gqa.py | 2 + tileops/ops/elementwise/_base.py | 6 +- tileops/ops/engram_decode.py | 3 + .../routed_expert/moe_grouped_gemm_nopad.py | 1 + tileops/ops/moe/routed_expert/unpermute.py | 1 + tileops/ops/norm/ada_layer_norm.py | 1 + tileops/ops/norm/ada_layer_norm_zero.py | 1 + tileops/ops/norm/fused_add_layer_norm.py | 1 + tileops/ops/norm/fused_add_rms_norm.py | 1 + 17 files changed, 104 insertions(+), 11 deletions(-) create mode 100644 tests/ops/test_multi_dtype_instance.py diff --git a/docs/design/ops-design-reference.md b/docs/design/ops-design-reference.md index 424f72c97..8952812bd 100644 --- a/docs/design/ops-design-reference.md +++ b/docs/design/ops-design-reference.md @@ -91,7 +91,7 @@ Slot-keyed rule dictionary consumed on demand by [ops-design.md](ops-design.md) ### Slot S12: `__init__` signature -- **Rule.** Keyword-only via `*`. Kwarg block order: (1) `static_dims` entries in manifest key order, no defaults; (2) `signature.params` entries in manifest key order; (3) `kernel_map` and `tune` last. **`dtype` is never a kwarg** — see [Parameter design](#parameter-design). +- **Rule.** Keyword-only via `*`. Kwarg block order: (1) `static_dims` entries in manifest key order, no defaults; (2) `signature.params` entries in manifest key order; (3) `kernel_map` and `tune` last. **`dtype` is a kwarg only when the inputs do not determine every output dtype** — see [Parameter design](#parameter-design). - **Derivation.** Manifest `static_dims` + `signature.params`. - **Example.** ```python diff --git a/docs/design/ops-design.md b/docs/design/ops-design.md index c12555b15..e811bb40e 100644 --- a/docs/design/ops-design.md +++ b/docs/design/ops-design.md @@ -27,7 +27,9 @@ Op ← L1: thin base, shared by all ops | Fixed-rank | `__init__` (all dims provided) | `_infer_output_shapes` runs once at init. | | Arbitrary-rank | `__init__` for `static_dims`; `forward` for everything else | `_infer_output_shapes` runs per shape. | -**Dtype is never a constructor parameter.** An op reads it from the input tensors in `forward()`. A caller who passes fp16 tensors gets the fp16 kernel without having said so twice, and an op can no longer be constructed in a state that disagrees with the tensors it is about to be handed. +**Dtype is not a constructor parameter when the inputs determine it.** An op reads it from the input tensors in `forward()`: a caller who passes fp16 tensors gets the fp16 kernel without having said so twice, and an op can no longer be constructed in a state that disagrees with the tensors it is about to be handed. + +An output dtype is determined by the inputs when it is `same_as(...)`, `promote_int_to_float(...)`, one concrete dtype, or a union equal to some input's. When some output dtype is an independent choice — an op that generates a tensor from parameters alone, or an fp8 path whose output may be fp16 or bf16 — the tensors are not a second source and `dtype` stays a `signature.params` entry. The kernel is dtype-specialized, so this makes kernel construction uniformly deferred to the first `forward()` — for fixed-rank and arbitrary-rank ops alike — and the kernel cache is keyed by shape *and* dtype. `dispatch_kernel()` stays in `__init__`: resolving the kernel *class* and checking the architecture needs no tensor, and keeping it there preserves fast failure on an unsupported GPU. @@ -302,8 +304,8 @@ satisfy the cold-call contract. `torch_compile_fullgraph` on an op whose compiled graph must backpropagate additionally requires registering an autograd formula for the dispatch custom op. -- Ops that pre-build their kernel at `__init__` (constructor-known shapes) - do not need the boundary; the invariant still applies to their +- An op that builds no kernel in `forward` — every kernel already in its + cache — does not need the boundary; the invariant still applies to its `forward`. ## Family-Base Refactoring diff --git a/tests/ops/test_multi_dtype_instance.py b/tests/ops/test_multi_dtype_instance.py new file mode 100644 index 000000000..a311f1e4e --- /dev/null +++ b/tests/ops/test_multi_dtype_instance.py @@ -0,0 +1,78 @@ +"""One op instance, two element types. + +Dtype used to be fixed at construction, so an instance served exactly one +element type and per-instance state could safely hold anything derived from +it. Now the tensors decide, and the family caches are hand-written per +family — this covers the invariant they all have to satisfy. + +The rest of the suite parametrizes dtype but builds a fresh op per case, so +it never exercises a second dtype through an instance that already has an +entry. +""" + +import pytest +import torch + +from tileops.ops.norm.layer_norm import LayerNormFwdOp +from tileops.ops.norm.rms_norm import RMSNormFwdOp +from tileops.ops.reduction.reduce import SumFwdOp + +_DTYPES = (torch.float16, torch.bfloat16) + + +def _assert_two_entries(op, cache_name="_kernel_cache"): + cache = getattr(op, cache_name) + assert len(cache) == 2, f"expected one entry per dtype, got {list(cache)}" + kernels = list(cache.values()) + assert kernels[0] is not kernels[1], "both dtypes reused one kernel" + + +@pytest.mark.smoke +def test_reduction_serves_two_dtypes_from_one_instance(): + op = SumFwdOp(dim=-1) + for dtype in _DTYPES: + x = torch.randn(8, 128, dtype=dtype, device="cuda") + y = op(x) + assert y.dtype == dtype + torch.testing.assert_close(y, x.sum(-1), atol=2e-2, rtol=2e-2) + _assert_two_entries(op) + + +@pytest.mark.smoke +def test_rms_norm_serves_two_dtypes_from_one_instance(): + n = 256 + op = RMSNormFwdOp(normalized_shape=(n,)) + for dtype in _DTYPES: + x = torch.randn(16, n, dtype=dtype, device="cuda") + w = torch.randn(n, dtype=dtype, device="cuda") + y = op(x, w) + assert y.dtype == dtype + _assert_two_entries(op) + + +@pytest.mark.smoke +def test_layer_norm_keys_on_both_shape_and_dtype(): + """A second dtype at the same shape must not reuse the first entry.""" + n = 256 + op = LayerNormFwdOp(normalized_shape=(n,)) + for dtype in _DTYPES: + x = torch.randn(16, n, dtype=dtype, device="cuda") + w = torch.randn(n, dtype=dtype, device="cuda") + b = torch.randn(n, dtype=dtype, device="cuda") + assert op(x, w, b).dtype == dtype + assert set(op._kernel_cache) == {(16, dt) for dt in _DTYPES} + + +@pytest.mark.smoke +def test_roofline_reports_the_most_recent_forward(): + """`self.dtype` is most-recent-forward, so bytes follow the last call.""" + op = SumFwdOp(dim=-1) + op(torch.randn(8, 128, dtype=torch.float32, device="cuda")) + _, bytes_fp32 = op.eval_roofline() + op(torch.randn(8, 128, dtype=torch.float16, device="cuda")) + _, bytes_fp16 = op.eval_roofline() + assert bytes_fp16 < bytes_fp32 + + +if __name__ == "__main__": + pytest.main([__file__, "-vvs"]) diff --git a/tileops/manifest/mamba.yaml b/tileops/manifest/mamba.yaml index 8936b705b..ba8cf5047 100644 --- a/tileops/manifest/mamba.yaml +++ b/tileops/manifest/mamba.yaml @@ -125,7 +125,6 @@ CBProducerOp: n_groups: {type: int} chunk_len: {type: int} d_state: {type: int} - dtype: {type: torch.dtype} shape_rules: - "B == batch" - "S == num_chunks * chunk_len" diff --git a/tileops/ops/_dtype_codegen.py b/tileops/ops/_dtype_codegen.py index d27b577d6..35a9b4e2d 100644 --- a/tileops/ops/_dtype_codegen.py +++ b/tileops/ops/_dtype_codegen.py @@ -240,7 +240,6 @@ def synthesize_validate_dtypes( # manifest inputs natively. A ``**kwargs`` body would need a per-call # ``Signature.bind``, which is measurable on this ``forward()`` hot path. closure: dict[str, Any] = { - "per_input": per_input, "input_names": input_names, "combo_keys": combo_keys, "ValueError": ValueError, diff --git a/tileops/ops/attention/deepseek_dsa.py b/tileops/ops/attention/deepseek_dsa.py index af64f77e9..b1160437a 100644 --- a/tileops/ops/attention/deepseek_dsa.py +++ b/tileops/ops/attention/deepseek_dsa.py @@ -128,5 +128,6 @@ def forward(self, q: torch.Tensor, kv: torch.Tensor, indices: torch.Tensor) -> t torch.Tensor: The result of applying the sparse attention operation on the input tensors. """ + self._validate_dtypes(q, kv, indices) self.dtype = q.dtype return self._get_kernel(q.dtype)(q, kv, indices) diff --git a/tileops/ops/attention/deepseek_mla.py b/tileops/ops/attention/deepseek_mla.py index 8cf60514a..8db8917e5 100644 --- a/tileops/ops/attention/deepseek_mla.py +++ b/tileops/ops/attention/deepseek_mla.py @@ -46,5 +46,6 @@ def default_kernel_map(self) -> Dict[str, Kernel]: def forward(self, q: torch.Tensor, q_pe: torch.Tensor, k: torch.Tensor, k_pe: torch.Tensor) -> torch.Tensor: + self._validate_dtypes(q, q_pe, k, k_pe) self.dtype = q.dtype return self._get_kernel(q.dtype)(q, q_pe, k, k_pe) diff --git a/tileops/ops/attention/deepseek_nsa.py b/tileops/ops/attention/deepseek_nsa.py index d46e40fcc..92d0b9506 100644 --- a/tileops/ops/attention/deepseek_nsa.py +++ b/tileops/ops/attention/deepseek_nsa.py @@ -40,7 +40,7 @@ def __init__( for key, value in params.items(): setattr(self, key, value) - self._kernel_params = {k: v for k, v in params.items() if k != "dtype"} + self._kernel_params = params self.dispatch_kernel(kernel_map) self._kernel_cache: Dict[torch.dtype, Kernel] = {} @@ -83,7 +83,7 @@ def __init__( for key, value in params.items(): setattr(self, key, value) - self._kernel_params = {k: v for k, v in params.items() if k != "dtype"} + self._kernel_params = params self.dispatch_kernel(kernel_map) self._kernel_cache: Dict[torch.dtype, Kernel] = {} @@ -140,7 +140,7 @@ def __init__( for key, value in params.items(): setattr(self, key, value) - self._kernel_params = {k: v for k, v in params.items() if k != "dtype"} + self._kernel_params = params self.dispatch_kernel(kernel_map) self._kernel_cache: Dict[torch.dtype, Kernel] = {} diff --git a/tileops/ops/attention/gqa.py b/tileops/ops/attention/gqa.py index 78aa4a030..221429866 100644 --- a/tileops/ops/attention/gqa.py +++ b/tileops/ops/attention/gqa.py @@ -1373,6 +1373,7 @@ def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, o: torch.Te do: torch.Tensor, lse: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: do = do.contiguous() + self._validate_dtypes(q, k, v, o, do, lse) self.dtype = q.dtype prep_kernel, kernel = self._get_kernels(q.dtype) delta = prep_kernel(o, do) @@ -1467,6 +1468,7 @@ def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Te v = F.pad( v, pad=(0, 0, 0, 0, 0, self.seqlen_kv - real_seqlen_kv), mode='constant', value=0) + self._validate_dtypes(q, k, v) self.dtype = q.dtype return self._get_kernel(q.dtype)(q, k, v, real_seqlen_kv) diff --git a/tileops/ops/elementwise/_base.py b/tileops/ops/elementwise/_base.py index 13f9552f9..eab6a5ee6 100644 --- a/tileops/ops/elementwise/_base.py +++ b/tileops/ops/elementwise/_base.py @@ -924,11 +924,13 @@ def _eager_forward(self, x: torch.Tensor) -> torch.Tensor: return self.kernel(x) def forward(self, x: torch.Tensor) -> torch.Tensor: + # Pass the locally derived M/N rather than building the kernel to + # populate self.M/self.N: a traced forward must not enter the TileLang + # builder, and _eager_forward builds on the far side of the boundary. M, N = self._validate_runtime_input(x) - self._ensure_kernel(M, N, x.dtype) wrapped = type(self)._wrapped if wrapped is not None: - return wrapped(x, self.M, self.N, self._instance_key) + return wrapped(x, M, N, self._instance_key) return self._eager_forward(x) diff --git a/tileops/ops/engram_decode.py b/tileops/ops/engram_decode.py index 8b3569e21..4f4a9e794 100644 --- a/tileops/ops/engram_decode.py +++ b/tileops/ops/engram_decode.py @@ -96,6 +96,9 @@ def forward( """ if not e_t.is_cuda: raise ValueError("e_t must be a CUDA tensor") + self._validate_dtypes( + e_t, h_t, conv_state, W_K, W_V, rms_w_h, rms_w_v, conv_w, + ) self.dtype = e_t.dtype e_t = e_t.contiguous() diff --git a/tileops/ops/moe/routed_expert/moe_grouped_gemm_nopad.py b/tileops/ops/moe/routed_expert/moe_grouped_gemm_nopad.py index 95b38f907..b14c5722a 100644 --- a/tileops/ops/moe/routed_expert/moe_grouped_gemm_nopad.py +++ b/tileops/ops/moe/routed_expert/moe_grouped_gemm_nopad.py @@ -83,5 +83,6 @@ def forward( Returns: C: [numel, N] GEMM output. """ + self._validate_dtypes(a, b, true_sizes, true_offsets) self.dtype = a.dtype return self._get_kernel(a.dtype)(a, b, true_sizes, true_offsets) diff --git a/tileops/ops/moe/routed_expert/unpermute.py b/tileops/ops/moe/routed_expert/unpermute.py index 011860e6e..7f13b11dc 100644 --- a/tileops/ops/moe/routed_expert/unpermute.py +++ b/tileops/ops/moe/routed_expert/unpermute.py @@ -85,5 +85,6 @@ def forward( Returns: output: [T, H] bf16/fp16 (``out`` if provided). """ + self._validate_dtypes(mm2_pad, fwd_idx, topk_weights) self.dtype = mm2_pad.dtype return self._get_kernel(mm2_pad.dtype)(mm2_pad, fwd_idx, topk_weights, out=out) diff --git a/tileops/ops/norm/ada_layer_norm.py b/tileops/ops/norm/ada_layer_norm.py index 781fb83bc..169fc0364 100644 --- a/tileops/ops/norm/ada_layer_norm.py +++ b/tileops/ops/norm/ada_layer_norm.py @@ -144,5 +144,6 @@ def forward( kernel = self._get_kernel(M_actual, N, dtype, x.device.index) y = kernel(x, scale, shift) self._last_roofline_mn = (M_actual, N) + self.dtype = expected_dtype return y.reshape(orig_shape) diff --git a/tileops/ops/norm/ada_layer_norm_zero.py b/tileops/ops/norm/ada_layer_norm_zero.py index 5e748ae9a..404928a3d 100644 --- a/tileops/ops/norm/ada_layer_norm_zero.py +++ b/tileops/ops/norm/ada_layer_norm_zero.py @@ -159,5 +159,6 @@ def forward( kernel = self._get_kernel(M_actual, N, dtype, x.device.index) y = kernel(x, scale, shift, gate) self._last_roofline_mn = (M_actual, N) + self.dtype = expected_dtype return y.reshape(orig_shape) diff --git a/tileops/ops/norm/fused_add_layer_norm.py b/tileops/ops/norm/fused_add_layer_norm.py index 4f21bf13c..c69cd6cae 100644 --- a/tileops/ops/norm/fused_add_layer_norm.py +++ b/tileops/ops/norm/fused_add_layer_norm.py @@ -166,5 +166,6 @@ def forward( kernel = self._get_kernel(M_actual, N, dtype, x.device.index) y, residual_out = kernel(x, residual, weight, bias) self._last_roofline_mn = (M_actual, N) + self.dtype = expected_dtype return y.reshape(orig_shape), residual_out.reshape(orig_shape) diff --git a/tileops/ops/norm/fused_add_rms_norm.py b/tileops/ops/norm/fused_add_rms_norm.py index bb109d3a3..e525e1d8d 100644 --- a/tileops/ops/norm/fused_add_rms_norm.py +++ b/tileops/ops/norm/fused_add_rms_norm.py @@ -154,5 +154,6 @@ def forward( kernel = self._get_kernel(M_actual, N, dtype, x.device.index) y, residual_out = kernel(x, residual, weight) self._last_roofline_mn = (M_actual, N) + self.dtype = expected_dtype return y.reshape(orig_shape), residual_out.reshape(orig_shape) From 028f168b99cbfdc11618a817111ed034a4f4d7b8 Mon Sep 17 00:00:00 2001 From: lcy-seso Date: Wed, 5 Aug 2026 10:15:25 +0800 Subject: [PATCH 15/15] [Test][Ops] Cover one instance serving two dtypes across the migrated families MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Eight cases over reduction, norm, attention decode, MoE and a top-level op. The suite parametrizes dtype everywhere but builds a fresh op per case, so nothing exercised a second dtype arriving at an instance that already had a cache entry — the invariant every hand-written family cache has to satisfy. Two cover paths nothing reached before: the decode ops select a kernel *slot* from the element type, so one instance must hold both the warp-specialized float16 kernel and the bfloat16 fallback; and a mismatch between the anchor input and the rest must be rejected rather than reinterpreted. --- tests/ops/test_multi_dtype_instance.py | 60 ++++++++++++++++++++++++++ 1 file changed, 60 insertions(+) diff --git a/tests/ops/test_multi_dtype_instance.py b/tests/ops/test_multi_dtype_instance.py index a311f1e4e..1f2802b8c 100644 --- a/tests/ops/test_multi_dtype_instance.py +++ b/tests/ops/test_multi_dtype_instance.py @@ -74,5 +74,65 @@ def test_roofline_reports_the_most_recent_forward(): assert bytes_fp16 < bytes_fp32 +@pytest.mark.smoke +def test_attention_decode_reselects_the_kernel_per_dtype(): + """The decode ops pick a kernel *slot* from the element type. + + float16 at batch=1 takes the warp-specialized kernel and bfloat16 falls + back, so one instance must hold two different kernel classes — the case + that used to need two ops with two constructor dtypes. + """ + from tileops.ops.attention.gqa import ( + GroupedQueryAttentionDecodeWithKVCacheFwdOp, + ) + + op = GroupedQueryAttentionDecodeWithKVCacheFwdOp(1, 32, 4, 8192, 128) + fp16 = op._get_kernel(torch.float16) + bf16 = op._get_kernel(torch.bfloat16) + assert fp16.__class__.__name__ == "GQADecodeBs1Kernel" + assert bf16.__class__.__name__ == "GQADecodeKernel" + _assert_two_entries(op) + + +@pytest.mark.smoke +def test_moe_unpermute_serves_two_dtypes_from_one_instance(): + from tileops.ops.moe.routed_expert.unpermute import MoeUnpermuteFwdOp + + total_tokens, top_k, hidden = 16, 2, 128 + numel = total_tokens * top_k + op = MoeUnpermuteFwdOp(total_tokens, top_k, hidden, padded_batch_sum=numel) + fwd_idx = torch.arange(numel, device="cuda", dtype=torch.int32) + for dtype in _DTYPES: + mm2_pad = torch.randn(numel, hidden, dtype=dtype, device="cuda") + weights = torch.rand(total_tokens, top_k, dtype=torch.float32, device="cuda") + assert op(mm2_pad, fwd_idx, weights).dtype == dtype + _assert_two_entries(op) + + +@pytest.mark.smoke +def test_cb_producer_serves_two_dtypes_from_one_instance(): + from tileops.ops.cb_producer import CBProducerOp + + batch, chunks, groups, chunk_len, d_state = 1, 2, 1, 64, 64 + op = CBProducerOp(batch, chunks, groups, chunk_len, d_state) + s = chunks * chunk_len + for dtype in _DTYPES: + c = torch.randn(batch, s, groups, d_state, dtype=dtype, device="cuda") + b = torch.randn(batch, s, groups, d_state, dtype=dtype, device="cuda") + assert op(c, b).dtype == dtype + _assert_two_entries(op) + + +@pytest.mark.smoke +def test_mismatched_input_dtypes_are_rejected(): + """The anchor selects the kernel; the others must agree with it.""" + n = 256 + op = RMSNormFwdOp(normalized_shape=(n,)) + x = torch.randn(16, n, dtype=torch.float16, device="cuda") + w = torch.randn(n, dtype=torch.bfloat16, device="cuda") + with pytest.raises(ValueError): + op(x, w) + + if __name__ == "__main__": pytest.main([__file__, "-vvs"])