Overrides the internal iteration index

Wait until all TMA descriptor previously issued are safe to be modified after tma_desc_commit_group()

L2 normalize matching FLA library implementation.

Adds a pointer offset in units of Element

Helper to add an offset in number of elements, assuming this offset is divisible by the vector size.

Increment and return an instance to self.

Return as INT4 ref (panics if not INT4).

Build a tokenized prompt of approximately target_tokens length. Uses a Gutenberg text, truncated via binary search on the tokenizer to hit th

Attention-only forward: pre-norm → attention → post-norm. Returns (hidden, residual) ready for MoE input. Used by multi-GPU EP prefil

Computes the grid shape

Hidden size of the model.

Load a BF16 expert's gate/up/down projections and quantize to INT4, INT8, or keep as BF16.

Load all GPU weights for one layer. Args: device: Device for norms, gate, shared expert, dense MLP. attn_device: Devi

Loads a fragment from memory with additional logical offset

Safe wrapper for the AVX2 INT4 matmul kernel.

Full MoE forward for a single token on one layer. Pre-quantizes the BF16 activation to INT16 once, then runs all selected experts using the integer k

Find indices of top-k largest values via partial selection sort. For small k (e.g., 10) and moderate n (e.g., 512), this is faster than full sort.

Print VRAM usage at a checkpoint for diagnostics.

Adds a pointer offset in units of Element

Returns the pointer to referenced data

Compute a single expert's output using the CPU transposed format. Dispatches to INT4 or INT8 transposed kernels based on expert.num_bits. Uses combin

Helps with visibility of barrier init operations across warps / cta / cluster Available as a separate function so as to batch inits across barriers an

Format an SSE chunk as OpenAI chat.completion.chunk JSON.

Fused add + RMSNorm (in-place on both buffers). If first_call: residual = hidden, hidden = rmsnorm(residual) Else: residual += hidden, hidden = rmsno

Generate scale permutation tables. Returns (scale_perm, scale_perm_single): - scale_perm: 64 elements, used when group_size < K - scale_perm_single:

Returns true if source is needed based on state of runtime arguments

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

Load and quantize an ungated expert (no gate_proj, just up_proj + down_proj). Returns (gate=empty, up, down) tuple compatible with ExpertWeights.

Migrate a Vec's backing memory to a specific NUMA node. Returns true if the migration was successful.

Perform a Consumer Epilogue to release all buffers

Standalone RMSNorm (non-fused).

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

Get raw CUDA pointer from tensor, or NULL.

Adds a pointer offset in units of Element

Adds a tile offset

Adds a tile offset

* Performs a non-blocking allocation of TMEM. * @param num_columns Number of columns being freed. Must be 32 <= num_columns <= 512 and power of 2.

Convert BF16 bytes to float32 numpy array.

Cache file path for Marlin format (GPU-native Marlin INT4/INT8).

Compute the number of contiguous elements needed to store a tensor with the given size

Chat completion: format messages, generate, decode.

Clears the predicate set efficiently

Returns the dim3 cluster shape.

Dequantize a slice of FP8 E4M3 bytes to BF16 u16 values, applying a per-tensor scale.

Dequantize INT4 packed weights back to f32 for verification.

Auto-detect the expert weight prefix from the weight map. Returns "model" for Qwen3/V2-Lite or "language_model.model" for Kimi K2.5.

add mask for small tiles in ELL

Clears the predicate set efficiently

Ensure we have enough pages for new_tokens more tokens.