Skip to content

[Refactor][Engram] Move the decode compile boundary to the op layer - #2280

Merged
lcy-seso merged 1 commit into
tile-ai:mainfrom
lcy-seso:refactor/engram/op-level-compile-boundary
Sep 27, 2026
Merged

lcy-seso merged 1 commit into
tile-ai:mainfrom
lcy-seso:refactor/engram/op-level-compile-boundary

Conversation

@lcy-seso

Copy link
Copy Markdown
Collaborator

Problems

  • EngramDecodeKernel registers its own torch.library custom op tileops::engram_decode with a hand-written fake, the last kernel-level custom op under src/tileops/kernels/; the compile boundary belongs to the op layer (ops-design.md § Compile Dispatch Boundary).
  • EngramDecodeFwdOp already declares compile_boundary = True and is in the compile-contract registry, so the kernel only runs inside the op's generated operator and the nested custom op is an extra dispatcher round trip.
  • The other kernel-level custom ops lost their decorators earlier but kept their bodies: 34 undecorated def _ fake functions, nine uncalled launch wrappers and fakes in DeltaNet and GLA, and comments claiming a custom op or torch.compile wrapper still exists.

Changes

  • EngramDecodeKernel.forward allocates its outputs and launches the projection and step programs directly, with the same arguments, buffers and order as the deleted custom op body.
  • Delete the undecorated fake functions, the uncalled wrappers, the imports only they used, and the stale custom-op comments across src/tileops/kernels/.
  • Surface change: the operator torch.ops.tileops.engram_decode no longer exists; nothing in the repo called it, and the op's generated operator is unchanged.
  • Eager outputs of EngramDecodeFwdOp are bit-identical (torch.equal) before and after on the three manifest rows and two extra rows (history shorter than the cache, unaligned d_mem).
  • The kernel's launched programs and configs are unchanged, so no benchmark numbers; the nightly run covers it.
  • Test node delta: 0 (no test file touched).

EngramDecodeKernel registered its own torch.library custom op and fake. EngramDecodeFwdOp already declares compile_boundary = True, so the kernel only ever runs inside the op's generated operator, where the nested custom op was an extra dispatcher round trip. The kernel now allocates its outputs and launches its two programs directly.

The other kernel-level custom ops lost their decorators earlier but kept their bodies: undecorated fake functions, uncalled launch wrappers, and comments claiming a custom op still exists. They are deleted.
@lcy-seso
lcy-seso requested review from a team and a lite review from Copilot September 27, 2026 15:48

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

@github-actions github-actions Bot added the refactor Code restructuring without behavior change label Sep 27, 2026
@lcy-seso
lcy-seso merged commit 1f6fef0 into tile-ai:main Sep 27, 2026
16 checks passed
@lcy-seso
lcy-seso deleted the refactor/engram/op-level-compile-boundary branch September 27, 2026 22:10
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

refactor Code restructuring without behavior change

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants