[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.1%	20 / 20
gemm_oob	0.080s	0.081s	+0.6%	20 / 20
indirect_load	0.014s	0.014s	+0.6%	20 / 20
nested_loop	0.149s	0.149s	+0.0%	20 / 20
block_pointer_loop_advance	0.084s	0.085s	+1.0%	20 / 20
liger_jsd	0.099s	0.099s	+0.1%	20 / 20
flaggems_layernorm	0.261s	0.260s	-0.1%	20 / 20
swiglu	0.121s	0.121s	+0.3%	20 / 20
cross_entropy	0.678s	0.677s	-0.2%	20 / 20
fused_linear_jsd	0.152s	0.152s	-0.1%	20 / 20
Total	1.710s	1.711s	+0.0%	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]

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.

Added end-to-end tests in tests/end_to_end/test_cutile.py. I kept cuTile as an optional because cuda-tile only ships wheels for Linux and Windows and has no sdist, so a hard dependency breaks pip install tilelens on macOS.

[mark14wu]

I reproduced these on the PR head (cbb433c) with cuda-tile 1.8.0 (+ tileiras), torch 2.14.1+cu130 and triton 3.8.0 on an RTX 4090. In the tables, the "GPU" column is a real ct.launch on the device and "interpreter / traced" is this PR's CuTileInterpretedFunction / tilelens.trace(..., frontend="cutile"). The PR's own cuTile tests pass (tests/unit/test_cutile.py + tests/end_to_end/test_cutile.py: 69 passed).

The first six inline comments are correctness issues; the rest are suggestions and nits.

[mark14wu]

Tracing crashes on masked-off out-of-range gather/scatter indices. _cutile_indexed_adapter passes the raw indices (only masked off) to the tracer, and masked_load then indexes the offsets array with them.

@ct.kernel
def gather_neg(a, out):  # a = np.arange(4, dtype=np.float32)
    ct.store(out, (0,), ct.gather(a, ct.arange(8, dtype=ct.int32) - 6))

run	result
GPU (ct.launch)	[0 0 0 0 0 0 0 1]
CuTileInterpretedFunction	[0 0 0 0 0 0 0 1]
tilelens.trace(..., frontend="cutile")	IndexError: index -6 is out of bounds for axis 0 with size 4

ct.scatter(out, ct.arange(8, dtype=ct.int32) - 6, t) fails the same way under tracing (GPU and the untraced interpreter both give [6 7 0 0]). Replacing invalid lanes before building the record, e.g. np.where(valid, key, 0), should fix it.

[mark14wu]

propagate_nan=True is silently ignored. maximum and minimum (L953) and max (L979) accept propagate_nan through **kwargs but always use np.fmax/np.fmin, so NaNs are dropped even when the kernel asks for them to propagate.

call	GPU	interpreter / traced
ct.max(Tile([1, nan]), propagate_nan=True)	nan	1.0
ct.maximum(Tile([1, nan]), 0.0, propagate_nan=True)	[1, nan]	[1, 0]

Using np.maximum/np.minimum (and np.max) when propagate_nan=True would match the GPU.

[mark14wu]

Several Builder methods are missing keyword arguments that cuda-tile documents, so valid kernels raise TypeError in the interpreter: min (here) has no propagate_nan, neither do argmax (L985) and argmin (L1426), and astype (L858) has no rounding_mode. All of these run on GPU with cuda-tile 1.8.0:

call	GPU	interpreter / traced
ct.min(x, propagate_nan=True)	nan	TypeError: Builder.min() got an unexpected keyword argument 'propagate_nan'
ct.argmax(x, propagate_nan=True)	1	TypeError: Builder.argmax() got an unexpected keyword argument 'propagate_nan'
ct.astype(x, ct.float16, rounding_mode=ct.RoundingMode.RN)	[1.5, 2.25]	TypeError: Builder.astype() got an unexpected keyword argument 'rounding_mode'

A test that compares each Builder method's keyword parameters with the matching cuda.tile function's would keep this from drifting as new cuda-tile releases come out.

[mark14wu]

cuTile ops reached through references the static scan doesn't patch run the empty stub body. call_tile_function_from_host unwraps every tile function, including the original cuda.tile op stubs, and calls them directly. Side-effecting ops become silent no-ops.

All of these compile and give correct results on GPU:

ACTS = (ct.exp, ct.store)

class Ops:
    exp = ct.exp
    store = ct.store

def store_default(out, t, op=ct.store):
    op(out, (0,), t)

def act(x, op=ct.exp):
    return op(x)

in the kernel	GPU	interpreter / traced
ACTS[1](out, (0,), ACTS[0](t))	out written	out left untouched, no error
Ops.store(out, (0,), Ops.exp(t))	out written	out left untouched, no error
store_default(out, t)	out written	out left untouched, no error
ct.store(out, (0,), act(t))	correct	KeyError: dtype('O')

Looking func up in the originals table (L1921) and dispatching to _language_operation before unwrapping would cover these cases.

[mark14wu]

Records leak between concurrent launches. client.records is cleared here before runner.run takes cutile_execution_lock. When two threads share ct.launch inside one interpreted_launches(Tracer()) scope, thread B clears the list while A is still tracing. A can lose records, and everything A appends after the clear stays in the list and gets counted into B's launch.

Repro: two threads each ct.launch a kernel that emits 2002 records, B started 0–40 ms after A, fresh process each time. [len(l.records) for l in launches] came out [2001, 4003] in 7/7 runs; the expected result is [2002, 2002].

The tilelens.trace path doesn't reset at all: calling k = tilelens.trace("tracer", frontend="cutile")(kernel) twice gives [6, 6] records (one run alone has 3), and launches[0] is launches[1]. That part already exists in Tracer/ClientManager (a reused Triton trace behaves the same way), so it may be out of scope here. Still, it suggests the reset belongs under the execution lock at launch start, e.g. in CuTileTrace.run, or in ClientManager/Tracer for all frontends, rather than in this shim.

[mark14wu]

Suggestion: flush_to_zero=True and any rounding mode other than RN/FULL are rejected for every op. ct.exp2(x, flush_to_zero=True) and ct.truediv(x, 3.0, rounding_mode=ct.RoundingMode.APPROX) both run on GPU but raise NotImplementedError here, so kernels tuned this way (common in attention kernels) can't be traced or visualized at all. For a debugging/visualization tool, approximating with RN and no FTZ, perhaps with a one-time warning, seems more useful than refusing to run.

[mark14wu]

nit: cat checks the shapes but not the dtypes. The GPU compiler rejects ct.cat((f32_tile, i32_tile), 0) with Expected tiles to have the same dtype: float32 != int32. Here np.concatenate instead promotes the result to a float64 tile, and in my repro the kernel then failed later at ct.store with TypeError: Cannot implicitly cast to float32, which points at the wrong line. Checking that the two dtypes match here would give the same error as the compiler.

[mark14wu]

nit: For every module the kernel references (other than cuda/numpy/tilelens), this looks up every name in the kernel's co_names on that module and recursively patches whatever it finds. For example, a kernel that calls a helper probe(torch) and also uses ct.load resolves torch.load and walks torch.serialization; while the kernel runs, torch.serialization.max is the patched tile min/max wrapper (it is restored afterwards), and other threads see the replacement too. Limiting the lookup to attributes the kernel actually reads off that module would avoid touching unrelated third-party code.

[mark14wu]

perf: The tracer calls ptr.get_offsets() on the root array for every traced load/store (tracer.py L108/L128/L132). That builds an int64 offset array for the whole array, only to index out one tile's worth, so tracing cost grows with array size instead of tile size. Measured with 16 loads + 16 stores of 16×16 tiles:

array	traced	untraced
256×256	4.9 ms	0.9 ms
1024×1024	20.9 ms	0.9 ms
2048×2048	141.5 ms	1.1 ms
4096×4096	388.0 ms	1.0 ms

At 4096×4096, each access allocates a 128 MiB offsets array. Computing sum(key * stride) for only the accessed keys would make each access O(tile size).

[mark14wu]

perf nit: Tile.__init__ always makes a copy (np.array(value, copy=True)), even when a Builder op passes in a freshly computed array that nothing else references. So the result of every elementwise op, reduction and load is copied twice. An internal constructor that takes ownership of the array and only sets write=False would avoid the second copy (this one is from reading the code; I didn't measure it).