CUDA Authoring and Support APIs

CUDA Authoring and Support APIs#

CUDA helpers#

The CUDA backend installs the Tx.cuda namespace when it is loaded:

from tvm.script import tirx as Tx

Tx.cuda.cta_sync()
leader = Tx.cuda.elect_sync()

These helpers cover operations that need multiple instructions, C/C++ expressions, descriptor packing, compiler annotations, or other behavior that is not represented as one table-driven PTX instruction. For a single supported PTX instruction, use Direct PTX Instructions.

The current Tx.cuda surface includes:

Category

Helpers

Synchronization and participation

any_sync, elect_sync, warp_sync, warpgroup_sync, cta_sync, grid_sync, cluster_sync, syncthreads_and, syncthreads_or, thread_rank, ballot_sync, __shfl_sync, __shfl_up_sync, __shfl_down_sync, __shfl_xor_sync, and __activemask

Barriers and memory ordering

mbarrier_wait, mbarrier_wait_acquire_cluster, thread_fence, atomic_add, and atomic_cas

Reductions

warp_reduce, warp_sum, warp_max, warp_min, cta_reduce, cta_sum, cta_max, cta_min, reduce_add_sync_u32, and reduce_min_sync_u32

Descriptors and addresses

wgmma.noop_barrier, wgmma.encode_matrix_descriptor, tcgen05.encode_matrix_descriptor, tcgen05.encode_instr_descriptor, tcgen05.encode_instr_descriptor_block_scaled, runtime_instr_desc, get_tmem_addr, cvta_generic_to_shared, smem_addr_from_uint64, sm100_2sm_leader_smem_addr, and mov_sreg

Loads, calls, and diagnostics

ldg, func_call, printf, trap_when_assert_failed, nano_sleep, clock64, and ffs_u32

Numeric conversion and packed math

half2float, bfloat162float, float22half2, half8tofloat8, float8tohalf8, uint_as_float, float_as_uint, make_float2, float2_x, float2_y, fmul2_rn, fadd2_rn, float22bfloat162_rn, float22bfloat162_rn_from_float2, bfloat1622float2, hmin2, hmax2, fp8x4_e4m3_from_float4, and fdividef

Instrumentation and compatibility

iket.mark, iket.range_start, iket.range_end, iket.range_push, iket.range_pop, iket.sentinel_token, iket.official_event, timer_init, timer_start, timer_end, timer_finalize, mma_store, mma_fill, mma_store_legacy, and mma_fill_legacy

NVSHMEM namespace#

The same backend installs Tx.nvshmem for my_pe, n_pes, signal_op, wait_until, quiet, fence, and barrier_all. Its getmem_nbi, putmem_nbi, and putmem_signal_nbi operations also provide .warp and .block variants.

CUDA language utilities#

tvm.backend.cuda.lang provides reusable objects for hand-written CUDA TIRx kernels, including schedulers, pipelines, roles, barriers, descriptors, and memory pools.

class tvm.backend.cuda.lang.BaseTileScheduler(*args, **kwargs)#

Base class for tile schedulers with common state and macros.

class tvm.backend.cuda.lang.ClusterPersistentScheduler2D(prefix, num_m_tiles, num_n_tiles, num_clusters, l2_group_size=8, cluster_m=1, cluster_n=1, serpentine=False)#

Persistent two-dimensional tile scheduler with group-major or serpentine traversal. init selects the first tile for a cluster, next_tile or next_tile_stride advances it, and valid reports whether the current m_idx / n_idx is in range.

class tvm.backend.cuda.lang.FlashAttentionLPTScheduler(*args, **kwargs)#

LPT scheduler with L2 swizzle for causal flash attention.

Processes high-work Q blocks (with more KV blocks to attend to) first using Longest Processing Time (LPT) scheduling. Also applies L2 cache swizzle for better cache locality across batch*head dimensions.

The LPT aspect comes from reversing m_block order: lower Q blocks have more KV blocks to process due to causal masking, so processing them first balances load.

The scheduler is only applied to non-persistent kernels.

L2 Swizzle: Groups consecutive batch*head indices together for L2 locality.

Parameters:

  • prefix (str) – Prefix for TIR variable names

  • num_batches (int) – Number of batches

  • num_heads (int) – Number of KV heads

  • num_m_blocks (int) – Number of Q blocks (M dimension tiles)

  • num_ctas (int) – Number of CTAs (should equal total_tasks for causal)

  • l2_swizzle (int) – L2 swizzle factor for cache locality

valid()#

Check if there are more tiles to process.

class tvm.backend.cuda.lang.FlashAttentionLinearScheduler(prefix, num_batches, num_heads, num_m_blocks, num_ctas)#

Linear persistent scheduler over (batch, head, m_block) tasks. init starts from a CTA index, next_tile advances by num_ctas, and valid reports whether work remains.

class tvm.backend.cuda.lang.GroupMajor3D(*args, **kwargs)#

3D grouped-row scheduler (M,N,K) with tail handling on M.

Parameters:

  • prefix (str)

  • m_tiles (int | T Expr # tiles along M (static or runtime))

  • n_tiles (int # tiles along N (static))

  • k_tiles (int # tiles along K (static))

  • group_rows (int # rows per group along M)

  • step (int = 1 # default stride for next_tile())

class tvm.backend.cuda.lang.IndexedTripleTileScheduler(*args, **kwargs)#

Scheduler that maps linear_idx to (b_idx, h_idx, q_idx) via index lists.

class tvm.backend.cuda.lang.MBarrier(*args, **kwargs)#

Mbarrier wrapper with regular mbarrier.arrive.

Parameters:

  • pool (SMEMPool) – Shared memory pool allocator.

  • depth (int) – Number of barrier slots (one per pipeline stage).

  • phase_offset (int) – XORed into the phase bit on every wait / arrive.

  • leader (Expr, optional) –

    Boolean predicate selecting the single thread that runs mbarrier.init. Defaults to T.cuda.thread_rank() == 0 – thread 0 of the enclosing CTA, which always picks exactly one thread regardless of which scope_id vars the caller declared. Override only when you want a different CTA-local thread to do the init.

    Note: the default deliberately avoids T.warp_id() / T.lane_id(). Those introduce deferred cta->warp / warp->thread ScopeIdDefs that the verifier cannot pin down unless the kernel header declares the full warp/lane chain (e.g. a single-CTA DSMEM kernel that only declares thread_id). It also avoids the synccheck false-deadlock on kernels that declare a second warp-scope id. The generated CUDA is equivalent.

remote_view(rank)#

Create a view of this barrier mapped to another CTA’s shared memory.

The returned view retains the local barrier and target CTA so arrive emits the cluster form. Its mapped buffer remains available through ptr_to for operations that consume a remote shared-memory pointer. init and wait are local-only and reject remote views.

class tvm.backend.cuda.lang.Pipeline(*args, **kwargs)#

A full/empty mbarrier pair for a software-pipelined data flow.

Pass barrier-type tags and Pipeline constructs and inits the barriers itself. Tags: "tma" (TMABar), "tcgen05" (TCGen05Bar), "mbar" (MBarrier). The barrier type and arrival count of each event stay explicit at the call site – e.g. Pipeline(pool, n, full="tma", empty="tcgen05", init_empty=NUM_CONSUMER).

Both signals are required: a Pipeline is a pair. For a one-way event (a pure “X happened” signal with no slot to recycle) use a bare barrier (TMABar/TCGen05Bar/MBarrier) directly – it has no empty side.

Parameters:

  • pool (SMEMPool) – Shared memory pool allocator.

  • stages (int) – Number of pipeline stages (barrier slots).

  • full (str) – Barrier-type tag for the full / empty signal (see above).

  • empty (str) – Barrier-type tag for the full / empty signal (see above).

  • init_full (int) – Expected arrival count for the full / empty barrier.

  • init_empty (int) – Expected arrival count for the full / empty barrier.

  • empty_phase_offset (int) – XORed into the empty barrier’s phase bit on every wait / arrive.

  • leader (Expr, optional) – Propagated to both barriers; defaults to thread 0 of the CTA.

class tvm.backend.cuda.lang.PipelineState(*args, **kwargs)#

Tracks stage and phase for a software-pipelined ring buffer.

This class does not know anything about full/empty barriers. Use it when the kernel manually waits/signals barriers, or when the stage/phase drives a ring not wrapped in a Pipeline.

Parameters:

  • depth (int) – Number of stages in the ring.

  • phase (int, optional) – Initial phase. Omit when initialization should happen later.

class tvm.backend.cuda.lang.RankAwareGroupMajorTileScheduler(*args, **kwargs)#

Group-major scheduler that applies a rank-aware remapping (remote rows first). Kept as a thin adapter because it depends on NVSHMEM rank at device-side.

class tvm.backend.cuda.lang.SMEMPool(*args, **kwargs)#

Bump allocator over a contiguous shared memory region.

Parameters:

ptr (Var or None, optional) – If omitted, an alloc_buffer([0], "uint8", scope="shared.dyn") is created automatically and commit() must be called after all allocations to emit the size annotation. If a Var is provided, the caller manages the backing buffer and commit() is a no-op.

alloc_tcgen05_mma_AB(shape, dtype='float16', swizzle_mode='auto', align=1024)#

Allocate MMA-compatible shared memory with an inferred swizzle layout.

commit(size=None)#

Emit pool size annotation into the IR.

Must be called after all alloc() / move_base_to() calls.

Parameters:

size (int, optional) – Explicit shared memory size in bytes. When None (the default), the high-water mark max_offset tracked by the allocator is used.

class tvm.backend.cuda.lang.SmemDescriptor(*args, **kwargs)#

Encoded once via init(), reused via add_16B_offset().

make_lo_uniform()#

Broadcast the lower 32 bits to all warp lanes via __shfl_sync.

class tvm.backend.cuda.lang.TCGen05Bar(*args, **kwargs)#

Barrier signaled by tcgen05 commit.

The caller is responsible for ensuring only one thread issues the commit, e.g. by wrapping the call in if T.cuda.elect_sync(): or by passing pred=T.cuda.elect_sync(). The pred form emits the predicated instruction (@p tcgen05.commit) instead of a branch and lets one elected leader predicate be shared across several commits.

class tvm.backend.cuda.lang.TMABar(*args, **kwargs)#

Barrier signaled by TMA (mbarrier.arrive.expect_tx).

When tx_count is None, falls back to a remote mbarrier.arrive (matching MBarrier.arrive defaults).

class tvm.backend.cuda.lang.TMEMPool(*args, **kwargs)#

Bump allocator over TMEM columns.

alloc_sf(shape, dtype, *, sf_per_mma, sf_reuse=1)#

Allocate a tcgen05 block-scaled SF TMEM buffer with an inferred layout.

shape last two dims are (rows, SF_K * sf_reuse) (the last dim is what gemm dispatch iterates over). When shape has 3 dims, the first is treated as a pipe-depth outer.

alloc_tcgen05_mma_A(shape, dtype='bfloat16', *, M, cta_group, ws=False, sparse=False, cols=None)#

Allocate a TMEM A operand (A-in-TMEM); layout resolved from MMA instruction params. M is the PTX tcgen05.mma instruction M (256/128/ 64), NOT per-CTA rows. See localdoc/claude_plan.txt S3a.

alloc_tcgen05_mma_D(shape, dtype='float32', *, M, cta_group, ws=False, sparse=False, group=None, cols=None)#

Allocate a TMEM D/C accumulator operand. M = PTX instruction M. group=(s,2,n) gives a flat (m,N) buffer a grouped column layout.

alloc_tcgen05_mma_C(shape, dtype='float32', *, M, cta_group, ws=False, sparse=False, group=None, cols=None)#

Allocate a TMEM D/C accumulator operand. M = PTX instruction M. group=(s,2,n) gives a flat (m,N) buffer a grouped column layout.

TMEMPool.alloc(shape, dtype='float32', *, layout=None, cols=None)#

Allocate a tensor-memory buffer from the pool. layout supplies an explicit TileLayout; without it, two-dimensional shapes use the default dense layout. cols overrides the inferred tensor-memory column count. Datapath-specific accumulators can instead use alloc_tcgen05_mma_D(...) / alloc_tcgen05_mma_C(...).

class tvm.backend.cuda.lang.WarpRole(warp_id_var, warp_id_val, regs=None, increase=False)#

A warp-level role that guards a block of code by warp_id comparison with optional register budget.

Generates:

if <warp_id_var> == <warp_id_val>:
    # if regs specified:
    T.ptx[f"setmaxnreg.{'inc' if <increase> else 'dec'}.sync.aligned.u32"](<regs>)
    <user code>

The if guard narrows the active set to the single warp; individual tile-primitive calls inside <user code> carry their own exec scope via a scope-namespace prefix (e.g. Tx.warp.copy(...)).

Parameters:

  • warp_id_var (Var) – The warp_id variable (from T.warp_id(...)).

  • warp_id_val (int) – Which warp index this role corresponds to.

  • regs (int, optional) – Register budget (passed to T.ptx.setmaxnreg). If None, no setmaxnreg is emitted.

  • increase (bool) – Direction for setmaxnreg (default False = decrease).

class tvm.backend.cuda.lang.WarpgroupRole(wg_id_var, wg_id_val, regs=None, increase=False)#

A warpgroup-level role that guards by wg_id comparison, with optional register budget.

Generates (single wg_id):

if <wg_id_var> == <wg_id_val>:
    # if regs specified:
    T.ptx[f"setmaxnreg.{'inc' if <increase> else 'dec'}.sync.aligned.u32"](<regs>)
    <user code>

Generates (range of wg_ids, e.g. wg_id_val=(0, 2)):

if 0 <= <wg_id_var> and <wg_id_var> < 2:
    T.ptx[f"setmaxnreg.{'inc' if <increase> else 'dec'}.sync.aligned.u32"](<regs>)
    <user code>

The if guard narrows the active set to the target warpgroup(s); individual tile-primitive calls inside <user code> carry their own exec scope via a scope-namespace prefix (e.g. Tx.wg.copy(...)).

Parameters:

  • wg_id_var (Var) – The warpgroup_id variable (from T.warpgroup_id(...)).

  • wg_id_val (int or tuple[int, int]) – Which warpgroup index (int) or range (start, stop) this role corresponds to.

  • regs (int, optional) – Register budget.

  • increase (bool) – Direction for setmaxnreg (default False = decrease).

IKET profiling#

NVIDIA IKET annotations and run-iket orchestration.

profile runs an explicit replayable command. run is intended for a script’s __main__ block: the parent process asks run-iket to replay the original script or python -m invocation, while the injected tracker and capture processes execute the supplied callable and then exit.

exception tvm.backend.cuda.iket.IketProfileError(message: str, *, returncode: int | None = None, command: Sequence[str] = (), output_tail: str = '', timeout: float | None = None)

An error while validating or running an official IKET profile.

class tvm.backend.cuda.iket.IketProfileResult(output_dir: Path, postprocess: str, json_traces: tuple[Path, ...], perfetto_traces: tuple[Path, ...], html_reports: tuple[Path, ...], command: tuple[str, ...] = ())

Published artifacts from a successful official IKET profile.

property trace_json: Path

Return the only JSON trace requested by this profile.

property perfetto: Path

Return the only Perfetto trace requested by this profile.

property html: Path

Return the only HTML report requested by this profile.

property trace: Any

Return the JSON trace without imposing a schema wrapper.

property launches: Any

Return trace["launches"] without reshaping it.

class tvm.backend.cuda.iket.IketProfiler(*args, **kwargs)

TIRx annotations compiled for NVIDIA’s official IKET runtime.

compile(mod, target=None, *, tir_pipeline='tirx')

Compile official IKET metadata and NativeDump placeholders.

tvm.backend.cuda.iket.profile(command: Sequence[str | PathLike[str]], *, output_dir: str | PathLike[str], profile_name: str = 'cutlass-4.6.0', postprocess: str = 'all', clobber: bool = False, cwd: str | PathLike[str] | None = None, env: Mapping[str, str | PathLike[str]] | None = None, max_ts_cnt_per_warp: int | None = None, keep: bool = False, timeout: float | None = 600.0) → IketProfileResult

Profile a replayable command with NVIDIA’s official run-iket tool.

tvm.backend.cuda.iket.run(main, *, output_dir: str | PathLike[str], profile_name: str = 'cutlass-4.6.0', postprocess: str = 'all', clobber: bool = False, cwd: str | PathLike[str] | None = None, env: Mapping[str, str | PathLike[str]] | None = None, max_ts_cnt_per_warp: int | None = None, keep: bool = False, timeout: float | None = 600.0) → IketProfileResult

Replay the active script under run-iket and execute main in its passes.

Module-level code executes once in the parent, tracker, and capture processes. main executes only in tracker and capture. Code following run executes only in the parent after a successful profile.

CUDA-specific transforms#

CUDA-specific TIRx transformations.

tvm.backend.cuda.transforms.LowerIket()

Lower frontend-only NVIDIA IKET annotations.

This pass must run after tvm.tirx.transform.SplitHostDevice and before tvm.tirx.transform.MakePackedAPI. It strips annotations for regular compilation and emits NVIDIA IKET metadata and NativeDump placeholders when the IRModule is explicitly IKET-enabled. Trace collection and postprocessing are owned by the external run-iket process.

Returns:

fpass – The result pass.

Return type:

tvm.transform.Pass

Contents