[FEA] CuTe DSL: prepared-launch API to eliminate per-call host marshalling overhead

[zkyue]

Component: CuTe DSL

Problem

Steady-state launches of an already-compiled CuTe DSL function spend nearly all
of their host time re-deriving state that is invariant across calls. Every
compiled(*args) call re-runs ExecutionArgs.generate_execution_args
(signature rectification, adapter lookups, t.cast into freshly allocated
ctypes storage) and repopulates the packed void** array — only to end at a
~2.6 us C call into the JIT engine. Real wrappers typically pay for
argument-object construction (make_ptr, cutlass.Int32/Float32, CUstream)
on top, because they receive fresh tensors/values each step.

Measured on a plain upstream tree (trivial 4096-element f32 scale kernel,
5 arguments, one B200, host wall and CUDA-event timings agree to 0.01 us
because the launch is host-bound; 5000 calls x 3 repeats, median; numbers
vary ~±0.3 us run to run on a shared host):

path	host cost (us/call)
direct call, argument objects reused	13.5
direct call, argument objects rebuilt each call (wrapper steady state)	20.8
prepared launch (in-place argument writes + the same C call)	3.8

Isolated component medians: ~5.9 us argument-object construction (the
wrapper-side cost the rebuilt-args row adds), ~8.2 us
generate_execution_args, ~1.2 us packed-array population, ~2.6 us for the
bare C call into the engine. The cost is structural, not incidental: for a
fixed (function, argument-kinds) pair, everything except the final C call is
recomputation of per-call-invariant state.

This host time matters in real workloads, where it is directly exposed as GPU
idle at pipeline drain points (externally reported measurements, listed as
corroboration; the table above is reproducible from this tree):

• In a production Megatron-Core CuTe-DSL integration (B200), per-launch Python
  wrapper cost measured 30-220 us in steady state across compressor/top-k
  kernel wrappers; an introspection-based prototype of the mechanism proposed
  here cut it to ~3.6 us per launch (internal measurement).
• The cudnn-frontend GTC'22 Spring Cutlass Talks #427 wrapper-overhead audit found the same shape of
  problem independently: ~8.5 us of Python wrapper host time wrapping a
  5.65 us kernel — the wrapper cost exceeding the kernel itself.
• Related in this tree: [BUG] __c_pointers__() causes monotonic RSS growth due to uncached ctypes allocations in kernel launch hot path #3351 reports monotonic RSS growth from exactly these
  per-launch ctypes allocations (__c_pointers__ in the hot path) — the same
  per-call marshalling shows up both as latency and as allocator pressure.

Proposed solution

A first-class "bind once, replay cheaply" handle on compiled functions
(implemented in the linked PR):

compiled = cute.compile(fn, out_ptr, in_ptr, alpha, n, stream)
prepared = compiled.prepare(out_ptr, in_ptr, alpha, n, stream)  # validate + marshal once
for step in range(steps):
    ...
    prepared.launch(out.data_ptr(), x.data_ptr(), alpha, n, raw_stream)  # per step

• prepare(*args, **kwargs) takes the same arguments as a normal call, runs
  the normal marshalling exactly once, and snapshots the marshalled bytes into
  per-argument ctypes storage owned by the returned PreparedLaunch.
  Eligibility is validated up front and is annotation-driven where the
  declaration is informative (a declared pointer/stream/numeric slot must be
  prepared with a value of that kind): runtime pointers (make_ptr /
  pointer-address slots), numeric scalars with plain ctypes storage (Boolean,
  Int8-64, Uint8-64, Float32/64, TFloat32, plain Python bool/int/float), and
  CUDA streams. Anything else — dlpack tensors, structs, sequences, custom JIT
  arguments, scalars with bit-cast storage such as Float16 — raises
  PreparedLaunchError at prepare time with a catalog diagnostic, and the
  normal call path remains for those.
• launch(*args, **kwargs) has the same signature and return contract as
  calling the compiled function. It rewrites pointer addresses / scalar values
  / the stream handle in place through the pre-built storage and invokes the
  compiled entry point directly — no per-call object construction, no
  marshalling, no repacking. Pointer slots also accept raw integer addresses,
  stream slots accept raw integer handles or framework streams exposing an
  integer cuda_stream attribute (e.g. torch.cuda.Stream), so the hot loop
  needs zero per-call object construction.
• Launches are bitwise-identical to the direct call: the entry point reads the
  same packed bytes either way (the PR verifies per-slot packed-payload byte
  identity and compares kernel outputs bit-for-bit).
• Each instance owns its packed void** array and CUDA-result slot, so
  instances are independent and never race the executor's shared thread-local
  scratch; the model is one instance per thread (sequential cross-thread
  handoff works). The handle pins its compiled module: launching after
  jit_module.unload() raises cleanly, and a recompiled/autotuned function
  needs a new handle.
• CUDA-graph capture through a capturing stream works with standard
  frozen-at-capture semantics (also exactly what one wants inside a capture).
• compiled.to(device).prepare(...) is the multi-device variant, mirroring
  the existing to() / __call__ split.

Measured effect (same benchmark as above): 13.5-20.8 us/launch direct →
3.8 us/launch prepared (up to 5.5x on the host path); in the production
integration above, 30-220 us → ~3.6 us (order of magnitude).

Alternatives considered

• CUDA graphs: complementary, not a substitute — capture constraints
  (allocator, stream discipline), frozen arguments (scalar updates need graph
  node rewrites or extra indirection), and heavier machinery for the simple
  "same kernel, new pointers/scalars" replay case. A cheap prepared launch is
  also exactly what one wants inside a capture region.
• AOT / TVM-FFI export ([QST] Best practice of managing various functions with CuTeDSL AOT + TVM-FFI in c++ #2852): escapes Python entirely; effective but
  changes the deployment model. A prepared launch keeps the pure-Python
  workflow.
• Implicit caching inside __call__ keyed on argument identity: easy to
  silently invalidate, still pays per-call eligibility/identity checks, and
  hides the thread/lifecycle contract. An explicit handle makes ownership,
  eligibility, and threading visible.
• Downstream introspection wrappers (what our integration currently
  ships): work, but write through private executor state and break with every
  internal refactor; a first-class API removes the need for
  version-sensitive shims in every downstream project.

Additional context

Compatibility: the feature is purely additive — no existing API changes
semantics, the __call__ hot path is untouched, and ineligible argument
mixes keep working through the normal call path (prepare-time rejection is
eager and typed, so callers can fall back). The eligibility whitelist is
conservative by design; the slot mechanism extends naturally (e.g. Float16
bit-pattern writers, multi-slot dlpack descriptors) if there is interest.
Happy to adapt the surface if the team prefers a different spelling — the
mechanism is independent of the surface.

PR with implementation + tests + tutorial example: .

(Dev-loop note: testing this in-tree requires the editable install flow, fixed for split wheels in #3417.)