[CuTeDSL] Add prepared-launch API for low-overhead kernel replay

[zkyue]

Fixes #3418 ([FEA] CuTe DSL: prepared-launch API to eliminate
per-call host marshalling overhead).

What

A first-class "bind once, replay cheaply" handle on compiled CuTe DSL
functions:

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, handle)  # in-place + direct call

Steady-state launches of an already-compiled function currently re-marshal
every argument on every call (ExecutionArgs.generate_execution_args plus
packed void** regeneration; wrappers typically also rebuild the argument
objects themselves each step) before a ~2.6 us C call. On a B200 with a
trivial 5-argument kernel that is 13.5-20.8 us/launch of host time;
prepared.launch is 3.8 us/launch (table below).

Design

• JitCompiledFunction.prepare(*args, **kwargs) -> PreparedLaunch validates
  eligibility once, runs the normal marshalling exactly once, and snapshots
  the marshalled bytes into private per-slot ctypes storage owned by the
  returned handle. It delegates to JitExecutor.prepare through the extracted
  _ensure_default_executor() (same lock/weakref-proxy pattern as before), so
  compiled.to(device).prepare(...) is the multi-device variant — the API
  mirrors the existing to() / __call__ split exactly.
  CudaDialectJitCompiledFunction inherits prepare(); cuda_jit_executor.py
  needs zero changes.
• PreparedLaunch.launch(*args, **kwargs) -> int | None has the same
  signature and return contract as calling the compiled function: keyword
  arguments and omitted defaulted parameters rectify through the existing
  get_rectified_args (raising the shared binding diagnostics); full
  positional calls are the fast path. It rewrites per-argument slots in place
  (no allocation, no repacking, no ExecutionArgs regeneration) and invokes
  the compiled entry point directly, checking a private CUDA-result slot.
  Pointer slots accept runtime Pointer objects or raw integer addresses;
  scalar slots accept Python scalars or Numeric instances (written with
  plain ctypes assignment semantics — a value the storage rejects, such as a
  float for an integer slot, raises a typed error instead of applying the
  implicit DSL cast a direct call would); stream slots accept CUDA stream
  objects, raw integer handles, or framework streams exposing an integer
  cuda_stream attribute (torch.cuda.Stream).
• Eligibility (validated at prepare, not at launch) is annotation-driven
  where the declaration is informative — a declared pointer, stream, or
  numeric slot must be prepared with a value of that kind — and value-driven
  for unannotated slots: runtime pointers (make_ptr / pointer-address
  slots; the marshalling contract is verified against the pointer's _desc
  storage, so contract lookalikes are rejected), numeric scalars with plain
  ctypes storage (Boolean, Int8-64, Uint8-64, Float32/64, TFloat32, plain
  Python bool/int/float), CUDA streams. Anything else (dlpack tensors,
  structs, sequences, custom JIT arguments, bit-cast scalar storage such as
  Float16, and TVM-FFI compiled functions, which dispatch through their own
  entry point) raises PreparedLaunchError — a DSLUserCodeError carried by
  six new LAUNCH_* entries in the DiagId catalog, so errors render like
  every other DSL diagnostic. The normal call path remains for rejected
  mixes.
• Bitwise identity: the initial slot bytes are the slow path's bytes
  (memmoved from the storages generate_execution_args just produced), and
  launch-time writes produce the same bytes the slow path would marshal for
  the same values of the prepared kinds. Verified per-slot (packed payload
  byte comparison) and end-to-end (bitwise output comparison in the tests).
• Ownership and lifecycle: each instance owns its packed void** array and
  CUDA-result slot (the executor's are shared/thread-local), so instances
  never race each other or the normal call path. The slow path's scalar
  ctypes storage is ephemeral, so the handle snapshots into cells it owns.
  Launching after jit_module.unload() raises cleanly; a recompiled or
  autotuned function needs a new handle. One instance per thread is the
  documented model (instances share no mutable state; sequential cross-thread
  handoff works). CUDA-graph capture through a capturing stream works with
  standard frozen-at-capture semantics.

What this PR does NOT change

• No changes to __call__'s hot path; the only change to existing code is
  extracting _ensure_default_executor() from run_compiled_program
  (verbatim, same locking).
• No changes to cuda_jit_executor.py.
• No behavior change for non-users; the feature is purely additive.

Performance

B200 (sm_100), trivial 4096-element f32 scale kernel (~2 us device time,
host-bound), 300 warmup + 5000 timed calls x 3 repeats, median; host wall =
perf_counter around the loop, cuda-event = events on the stream around the
loop (agreement to 0.01 us confirms the host is the bottleneck). Numbers vary
~±0.3 us run to run on a shared host:

path	host wall (us/call)	cuda-event (us/call)
direct call, reused arg objects	13.48	13.49
direct call, rebuilt arg objects (wrapper steady state)	20.76	20.76
prepared.launch, raw int addresses/handle	3.77	3.78
prepared.launch, rebuilt pointer objects	6.40	6.40

Slow-path host breakdown (isolated medians): ~5.9 us argument-object
construction (wrapper-side, applies to the rebuilt-args row), ~8.2 us
generate_execution_args, ~1.2 us packed-array population, ~2.6 us bare
entry-point call. In a production Megatron-Core integration the same
mechanism cut 30-220 us/launch wrapper steady state to ~3.6 us (externally
measured corroboration; see the issue).

Tests

test/python/CuTeDSL/test_prepared_launch.py (unittest style, matches the
existing suite; device behavior is exercised on a real sm_100 kernel, and
eligibility rejections also use host-only signatures):

• bitwise output parity with the direct call (incl. prepare before any call);
• pointer/scalar/stream mutation between launches, raw ints, kwargs,
  defaulted parameters omitted at prepare and launch, framework streams at
  launch; Float64/Int64/Boolean width coverage with bitwise checks;
• fast-path contract: launches never re-enter generate_execution_args
  (monkeypatch guard) and reuse the packed array in place;
• storage stability: 1000-launch read-modify-write accumulate under periodic
  gc + allocator churn;
• eligibility rejections (dlpack, sequence, Float16, multi-slot custom
  __c_pointers__, pointer-contract lookalikes, wrong-kind values for
  declared pointer/stream slots, framework stream at prepare, extra tail
  args, missing args, TVM-FFI-style non-JitExecutor executors) and that the
  direct call keeps working for rejected mixes;
• misuse: launch arg-count mismatches (shared binding diagnostics), bad
  pointer/scalar/stream/boolean values (naming the argument), launch after
  jit_module.unload();
• threading: concurrent per-thread instances + sequential cross-thread
  handoff;
• to(device).prepare(...);
• CUDA-graph capture + replay and frozen-at-capture semantics both directions;
• an in-test micro-benchmark asserting prepared is not slower than direct
  (direction-only, load-tolerant).

28 tests; full test/python/CuTeDSL/ suite: 34 passed (includes the
pre-existing test_struct_in_if.py), zero regressions.

A tutorial example (examples/python/CuTeDSL/dsl_tutorials/prepared_launch.py)
follows the by-passing-dlpack argument style of call_bypass_dlpack.py.

Notes for reviewers

• The eligibility whitelist is deliberately conservative for v1; the slot
  mechanism extends naturally (Float16/BFloat16 bit-pattern writers,
  multi-slot dlpack descriptors) if there is interest.
• Launch-time value checking is intentionally light — ctypes assignment
  semantics with typed errors on rejected values (out-of-range integers
  truncate to the declared width, like the existing fast path; wrong-kind
  scalars, non-stream stream values, and non-bool Boolean values raise
  PreparedLaunchError). Full _full_arg_check validation still runs at
  prepare via the normal marshalling; kind-level validation is done once at
  prepare, annotation-driven.
• Related issues: [BUG] __c_pointers__() causes monotonic RSS growth due to uncached ctypes allocations in kernel launch hot path #3351 (RSS growth from per-launch __c_pointers__
  allocations — this PR's replay path performs none), [QST] Best practice of managing various functions with CuTeDSL AOT + TVM-FFI in c++ #2852 (AOT/TVM-FFI
  export — complementary deployment model).

(To run the tests in-tree via the editable flow on current wheels you'll also want #3417.)