[pd-hung]

This change does not include TileBench

Summary

• Adds the cuTile CPU interpreter.
• Connects cuTile to the shared tracing frontend.
• Adds a matmul example with visualization.
• Adds cuTile setup and usage instructions.

Test Plan

• Unit tests in tests/cutile/test_cutile.py.
• Integration tests in tests/cutile/test_cutile_tracer.py, covering kernel execution, trace save/load, and visualizer APIs.
• Tested against TileBench across all 45 operators, all passed.
• All unit, cuTile tests passed.

Related Issues

None.

Breaking Changes

None.

Checklist

• I added tests to all new functionality I added/bugs I fixed.
• I verified that a human has reviewed all code in this PR.
• I ran npm run build:frontend if the PR modified any TypeScript code.
• I made sure that my code is well documented (comments explaining strange code, docstrings for functions, website modified if new functionality added).

[pd-hung]

I’ll start resolving the merge conflicts and update the cuTile implementation to use the new tilelens.

[github-actions]

Performance Benchmark

Benchmark	main (min)	PR (min)	Change	Samples
gemm	0.072s	0.072s	+0.8%	20 / 20
gemm_oob	0.080s	0.080s	-0.7%	20 / 20
indirect_load	0.014s	0.014s	-1.4%	20 / 20
nested_loop	0.148s	0.147s	-0.8%	20 / 20
block_pointer_loop_advance	0.085s	0.083s	-2.1%	20 / 20
liger_jsd	0.100s	0.100s	-0.6%	20 / 20
flaggems_layernorm	0.261s	0.258s	-1.3%	20 / 20
swiglu	0.123s	0.120s	-1.9%	20 / 20
cross_entropy	0.688s	0.683s	-0.6%	20 / 20
fused_linear_jsd	0.152s	0.152s	-0.4%	20 / 20
Total	1.724s	1.709s	-0.8%	N/A

Iterations: 1 warmup + 20 measured
Samples are shown as main / PR; long pytest benchmarks may use fewer samples.

[Jokeren]

Please add end to end tests

[Jokeren]

cutile can be a default dependency

[Jokeren]

We could make cuTile default if you're confident

[mark14wu]

Thanks for the PR. The integration through the shared frontend adapters works well, CuTileTrace follows the same pattern as NKI, and the cuTile tests pass locally.

I reviewed this mainly by differential testing against real cuTile:

• cuda-tile 1.5.0 with [tileiras] on an RTX 4090: I launched the same kernels via ct.launch and compared result dtypes and values with the interpreter.
• I ran all 45 TileBench cuTile operators through the interpreter and compared them against both the GPU run and the torch references.

Each inline comment includes a repro. Summary:

Must fix

1. The example kernel and two test kernels do not compile on real cuTile (math.ceil in the loop bound, ct.arange called without dtype).
2. bf16 ct.sum / ct.cumsum saturate at 256 (the GPU returns 1024 for 1024 ones).
3. Tracing bf16/fp8 kernels:
   • tilelens.save() raises.
   • The visualizer silently receives empty tensors for bf16 Load/Store.
   • A bf16 matmul makes /api/data return 500.
4. ct.sum(mask) returns True instead of a count, and bool + int stays bool.
5. 0-d float64 results are narrowed to float32 (ct.sum([1e300, ...]) gives inf).

Silent divergences from GPU semantics

• Scalar typing:
  • Python scalars are hard-coded to int32/float32, so int64 constants overflow and float64 tiles lose precision.
  • ct.bid() and kernel scalar arguments are treated as weak literals.
• Result dtypes and semantics of individual ops:
  • bf16 matmul returns float32 instead of bf16.
  • NaN handling in max/min/argmax differs.
  • Float floordiv: 1.0 // 0.1 gives 9 here and 10 on the GPU.
  • bitwise_and / bitwise_or follow different promotion rules from bitwise_xor.
• NumPy/torch dtype arguments raise KeyError. This crashes 24/45 TileBench operators at commit 224ec81; all 45 pass on current TileBench main, which switched to ct.* dtypes.
• Code that should be rejected is accepted:
  • Stores implicitly narrow (f32 into f16) where cuTile raises a compile error.
  • Misspelled keyword arguments are silently swallowed.
• Coverage gaps:
  • 43 public tile ops are unpatched and fail with a misleading error.
  • @ct.function helpers and from cuda.tile import ... do not work.

Biggest lever: run cuTile's own front-end as a pre-check, the way NKITrace does with pre_trace. It runs CPU-only and takes milliseconds. It rejects most of the "accepted here, rejected on GPU" cases, and it can give the interpreter the exact result types.

Tests, CI and packaging

• +1 to the request for end-to-end tests in tests/end_to_end/.
• Validate test kernels against real cuTile rather than only against the interpreter.
• Remove pythonpath = . from pytest.ini.
• Use a version range instead of the exact pin. Making cuTile a default dependency looks cheap; see the comment on pyproject.toml.
• Keep a CI job that runs without cuTile installed.

[mark14wu]

This kernel does not compile with real cuTile. I tested cuda-tile 1.5.0 via ct.launch on a 4090:

TileTypeError: Cannot create constant from value of type builtin_function_or_method.
  examples/cutile/matmul.py, line 33, col 20-28, in matmul_kernel:
      for k in range(math.ceil(K / TILE_K)):

The interpreter runs the kernel body as plain Python, so math.ceil works there. The README points users to this example and the tracer test imports it, so please use ct.cdiv(K, TILE_K) instead. With that change the kernel compiles and produces exact results on the GPU.

[mark14wu]

dtype is a required argument of ct.arange (_stub.py: def arange(size, /, *, dtype, start=0, step=1)). Real cuTile rejects this call with:

TileTypeError: arange(): missing a required keyword-only argument: 'dtype'

L96 has the same problem. It passes here only because Builder.arange gives dtype a default (see my comment there).

[mark14wu]

In cuTile, dtype is a required keyword-only argument. Defaulting it here lets kernels that will not compile on GPU pass under the interpreter; two of the kernels in this PR's own tests are examples. Please remove the default.

[mark14wu]

Accumulating in value.dtype causes two problems:

1. bf16 saturates at 256. ml_dtypes' generic reduce adds the elements one at a time and rounds to bf16 after each step:

   • ct.sum of 1024 bf16 ones gives 256; the GPU gives 1024.
   • ct.sum of 1024 × 0.1 gives 32.0; the GPU gives 102.5.

   Any row sum of 512 or more elements (layernorm, softmax denominators) is off by 2–4x. cumsum (L726) has the same problem. fp16 and fp32 are fine because they use NumPy's pairwise sum.

2. bool. ct.sum(mask) returns True instead of a count. The GPU promotes bool to int32. cumsum and max on bool tiles behave the same way.

Suggestion: accumulate bf16/fp8 in float32 and cast back afterwards, and follow cuTile's result-dtype rules for bool.

[mark14wu]

np.float64 is a subclass of Python float, so every 0-d float64 result that reaches this line is narrowed to float32. That includes ct.sum/ct.max over all axes, ops on 0-d tiles, and 1-D @ 1-D matmul:

ct.sum(float64 [1e300, 1, 1, 1])  -> GPU: float64 1e300   | interpreter: float32 inf
s + s with s = 1e-300 (0-d load)  -> GPU: float64 2e-300  | interpreter: float32 0.0

test_mixed_float_tiles_do_not_narrow_float64 is meant to guard against this, but it only uses 1-element 1-D arrays. Checking type(value) is float instead of isinstance fixes it.

[mark14wu]

A README note is easy to miss, and non-contiguous inputs do more than look off:

• Transposed or strided inputs highlight the wrong cells, because offsets are physical but get unraveled in C order. For example, block 0 of an [::2] input highlights 0, 2, 4, 6 instead of 0–3.
• Negative-stride inputs highlight almost nothing.
• tilelens.save raises.

A runtime warning, or rejecting non-contiguous arrays in the cuTile trace path, would be safer. Related: the visualizer derives element size from the dtype name, so cuTile bool_ arrays (and int16/fp8 arrays) are treated as 4-byte and highlight the wrong cells.

[mark14wu]

On test structure: +1 to Jokeren's request. Per the repo convention, kernel-level behavior should live in tests/end_to_end/ and unit tests in tests/unit/.

More importantly, every test here checks the interpreter only against itself or NumPy. That is how the example and two of these kernels ended up not compiling on real cuTile. Suggestions:

1. Add end-to-end tests that trace real kernels, including bf16 kernels, and cover save/load plus the visualizer endpoints.
2. Add a CPU-only check that each test kernel passes cuTile's front-end (export_kernel), so the tests cannot drift from real cuTile.

Coverage also shows untested paths:

• num_tiles, equal/__eq__, minimum/builtin min, Tile.__matmul__
• most reflected operators
• the unique-index atomic_add path
• tfloat32
• the int-tile-with-float-scalar promotion branches
• most validation errors

[mark14wu]

This guards against float64 narrowing, but only with 1-element 1-D arrays. 0-d results are still narrowed (see my comment on Tile.__init__). Please add a case such as ct.sum over a float64 tile.

[mark14wu]

Thanks for the cuTile work, this is a big and useful addition. I went through the interpreter, the frontend adapter, and the tracer/visualizer changes. Inline comments below, roughly ordered by severity; most correctness items include a minimal repro.

Would like these addressed before merge

1. cuda-tile as a required dependency breaks installs on platforms without a wheel (pyproject.toml).
2. The shared STORAGE_DTYPES additions change NKI beta2 dtype validation (tilelens/utils/dtypes.py).
3. interpreted_launches accumulates records across launches (tilelens/core/frontend/cutile.py).

Correctness

• Non-deterministic out-of-bounds fill with PaddingMode.UNDETERMINED
• Tuple/list array arguments are not wrapped as Array
• Unannotated int scalars are typed by magnitude (int64) instead of int32
• atomic_cas compares numerically rather than bitwise; failed CAS lanes are traced as Stores
• Kernel-scope patching reaches into third-party modules (e.g. torch)
• numpy-integer grid dims are rejected
• flush_to_zero / RoundingMode.APPROX raise NotImplementedError

Performance

• get_offsets() is O(root array size) per traced access
• Per-op Signature.bind and per-Tile copies

Cleanup

• Duplicated padding table / dtype helpers / root-array walk
• tests/cutile/ overlaps with tests/end_to_end/test_cutile.py

Minor: the PR checklist is still unchecked, could you update it?

[mark14wu]

[blocking] cuda-tile only publishes manylinux (x86_64/aarch64) and Windows wheels and has no sdist, so as a hard dependency pip install tilelens / uv sync --extra test fails on macOS and any other platform without a wheel, including for Triton-only users.

The code already treats cuTile as optional (HAS_CUTILE guards, pytest.importorskip), so I'd move this to [project.optional-dependencies] as a cutile extra. That would also make the pip install "tilelens[cutile]" command in docs/index.html work; right now that extra doesn't exist.

[mark14wu]

[blocking] STORAGE_DTYPES is shared with the NKI beta2 simulator, where a missing key is exactly what makes _storage_dtype raise TypeError: Unsupported dtype. Adding "int64", "uint64" and "float8_e8m0fnu" here means e.g. nl.ndarray(shape, dtype="int64") now silently allocates instead of failing, so NKI kernels that wouldn't compile for Trainium pass in simulation.

Could the cuTile-only dtypes live in a cuTile-local table (e.g. STORAGE_DTYPES | {...} in the cuTile module) instead?

[mark14wu]

[blocking] Each ct.launch builds a fresh CuTileTrace but reuses the same client, and the Tracer's records list is never reset between launches.

Repro: inside with interpreted_launches(Tracer()):, two ct.launch(None, (1,), k_copy, (src, dst)) calls give [len(l.records) for l in launches] == [3, 6]. The second Launch also contains the first launch's Grid/Load/Store records, so the visualizer shows duplicated accesses.

[mark14wu]

With the default PaddingMode.UNDETERMINED, the out-of-bounds lanes of a partial tile come from np.empty, i.e. arbitrary memory. Example: loading a (4,) tile at index 1 from a 6-element array returned [4.0, 5.0, 1.4e-45, 0.0]. A partial-tile ct.sum / ct.max / ct.matmul then gives run-to-run different (possibly NaN) results, and recorded Load/Dot data isn't reproducible.

Since the value is undefined anyway, a deterministic fill (e.g. NaN for floats, a sentinel for ints) would keep traces stable and make reliance on undefined lanes visible. Alternatively raise, as load_advanced_indexing already does for this case.

[mark14wu]

Tuples and lists are passed through unchanged, so numpy arrays inside them are never wrapped as Array (and never registered by the tracer).

Repro: def k(arrs, dst): ct.load(arrs[0], (0,), (4,)) with CuTileInterpretedFunction(k).run((np.arange(4.),), dst) raises AttributeError: 'memoryview' object has no attribute 'transpose' (it ends up on ndarray.data). cuTile supports tuple/list-of-array parameters, so this probably needs to recurse into tuples/lists.

[mark14wu]

flush_to_zero=True and rounding modes other than RN/FULL (e.g. APPROX) raise NotImplementedError, so attention/softmax-style kernels that use these flags can't be interpreted or traced at all. On CPU, APPROX could fall back to the exact result, and FTZ could be emulated by flushing subnormals (or ignored with a one-time warning). Not blocking, but it would widen coverage a lot.

[mark14wu]

get_offsets() materializes byte offsets for the entire root array, and the tracer calls it on every traced Load/Store, so tracing cost is O(blocks x accesses x array size) rather than O(tile size).

Measured: a 1024x1024 tiled copy takes 0.09s untraced vs 1.26s traced. Extrapolating, 4096x4096 with 32x32 tiles is ~32k calls at ~30ms each, about 16 minutes. Computing offsets only for the accessed tile (sum(index * stride) over the access) would avoid this.

[mark14wu]

Minor perf: every ct op runs Signature.bind + apply_defaults just to validate a few kwargs, and Tile.__init__ (L281) always copies its input, including freshly computed temporaries like Tile(operation(x, y).astype(dtype)). In elementwise-heavy loops that's a few microseconds per op plus an extra allocation and copy per intermediate. Validating only the enum/rounding kwargs when they're present, and a no-copy path for internally constructed Tiles, would remove most of it.

[mark14wu]

Nit: this padding-mode-to-value table and the "nonzero padding requires a floating-point array" check duplicate Builder.load (~L684). Similarly, _is_float and _is_low_precision repeat the same four dtypes, and the while root._parent walk appears in both adapters in frontend/cutile.py (L44, L106), mirroring the tracer's _base_tensor. A module-level _PADDING_VALUES and a shared _root(array) helper would keep the two load paths from drifting apart.

[mark14wu]

AGENTS.md puts kernel-level behavior in tests/end_to_end/. This new tests/cutile/ directory overlaps with tests/end_to_end/test_cutile.py (both define copy/matmul kernels and both check the visualizer). Could these tests move into tests/end_to_end/test_cutile.py, dropping the third path from the README?