Epilogue implementation for non-TMA store version. :param epi_tidx: Thread index :type epi_tidx: cutlass.Int32 :para

Returns the biased exponent

Execute the given function for each supplemental read tensor.

Returns a pointer

Return if the address in on the diagonal

Given the `obj`, recursively go through it to extract all contained C pointers

Get the current k_index, k_tile_count, and local split_kv value for the MLA kernel. :param split_kv: Split_kv value :type split_kv: c

Get the current k_index, k_tile_count, and local split_kv value for the MLA kernel. :param split_kv: Split_kv value :type split_kv: c

Get the registered JIT argument adapter for the given argument.

Get A extents. fprop: A extents array contains [N,D,H,W,C]. Turn that into ((W,H,D,N), (C)) dgrad: A extents array contains [N,Z,P,Q,K]. Turn that int

Get the SM (compute capability) version of a CUDA device.

Simple heuristics to select the optimal SMEM layout atom based on the majorness, the data type, and the major mode size. :param major_mode: T

Calculate the number of unmasked trips for the current block. This represents the number of trips that don't require special

Initializes the workspace

Return the underlying MLIR vector value.

Check if a field is a constexpr field.

Loads a fragment from memory at the location pointed to by the iterator.

Perform a Producer Epilogue to prevent early exit of blocks in a Cluster

Loads a fragment from memory

Loads a fragment from memory with additional logical offset

Loads a fragment from memory

Loads a fragment from memory with additional logical offset

Loads a fragment from memory

Loads a fragment from memory

Loads a fragment

Normalize a field specifier to its IR logical field name. Accepted inputs: - Enum value present in admissible_fields (must expose _to_i

Convert cumulative end offsets to per-group sizes.

Returns whether the provided operation class supports the provided data type combination :param op_class: operation class to conside

Pack values to an alloca that lays out in the order of the values. Parameters ---------- current_block : ir.Block

The full procedural name indicates architecture, extended name, tile size, and layout.

Returns a string name for debugging

Raise an error and return -1. Instead of concatenating parts at compile time, we define each part as a global string and call a helpe

Obtains raw bits

Remove edge src -> dst

Create a return statement.

Runs the kernel

Executes one test

Executes one test

Executes one test

Notify the iterator which k-group it is currently pointing to. This does not advance the iterator. Rather, it overrides its internal tracking with co

Sigmoid activation function: f(x) = 1 / (1 + exp(-x))

Compute element-wise sine of the input tensor. :param a: Input tensor (in radians) :type a: Union[TensorSSA, Numeric] :param fastmath: En

Softmax for one k-tile. Updates the related pipeline states and returns the computed results. :param common_params: The common parameters

Compute softmax on attention scores from QK matrix multiplication. This method handles the softmax computation for either the first or second

Stores a fragment to memory

Store

Store

Store a fragment to memory

Stores a fragment to memory at the location pointed to by the iterator

Store a fragment to memory

Store a fragment to memory

Stores a fragment

/////////////////////////////////////////////////////////////////////////////////////////////// Strided dgrad helper functions

Return the output tensor (concept: cutlass_cppgen.backend.evt.ir.tensor)

Returns a list of valid tile descriptions for the operations :returns: list of valid tile descriptions for the operations :r

UMMA arrive for MMA_2x1SM + TMA_LOAD_MULTICAST combination

Constructs the arguments structure given the configuration and arguments

Generate a yield operation. It it used to return values from a loop, if-else, or while region.