Get a breakdown of memory usage by tensor. Returns: Dict mapping tensor name to size in bytes

return compute throughput in TOPS

TFLOPS with no sparse.

Get current Ditron MoE profiling state. Returns: Dictionary with profiling flags for each operation.

Walk a triton subdirectory to enumerate all subpackages dynamically.

Get the LLType of a value.

Get the underlying tensor, or None if not materialized.

scatter_idx: [M, topk]

Return True if this file resides in a git repository

Check if the type is a grid constant

Check if a Call node represents a loop modifier. Args: node: AST Call node to check ctx: Context dic

Check if the type is a scalar

Check if an object is a struct stub.

`wait_inputs` is used to build data dependency.

Args: q: (bs, seq, nheads_q, head_dim) k: (bs, seq, nheads_kv, head_dim) v: (bs, seq, nhe

Create an int64 device tensor holding data_ptr() values (simulating int32_t**).

this op assume that kv_lens has been update (kv_lens = history_kv_len + seq_len(qkv.shape[1])) inplace update new kv to key_c

Inputs: qkv: [bs, seq_len, num_q_heads + 2 * num_kv_heads, head_dim] kv_lens: [bs] q/k_rm

Full expert tables (identical on every rank via a fixed seed) + this rank's local slice. AMD layout: ``W1=[E, hidden, inter]`` (up), ``W2=[E, int

Mark a name as global/nonlocal in the current scope. This means the name refers to the outer scope, not a local definition.

Parse command output to extract performance data

Remove the active scope (e.g., when exiting a function/loop)

Prepare inputs for MoE forward pass. Args: ffn_dim: FFN intermediate dimension hidden_dim: Hidden dimension topk

Cross-rank split exchange + recv-offset + grouped-GEMM tiling. Built on-device by default (``use_device_metadata=True``); set it False to use

Print complete benchmark comparison for all configurations. Args: all_implementations: Dict mapping config keys to implementatio

Create a new nested scope (e.g., for function bodies, loops)

Read file contents and return as list of lines.

out[t] = sum_j w[t,j] * relu(x[t] @ W1[e]) @ W2[e] (model-dtype matmul, fp32 reduce).

out[t] = sum_j w[t,j] * relu(x[t] @ W1[e]) @ W2[e] (model-dtype matmul, fp32 reduce).

Global bincount of expert indices (including drop token at index num_experts).

Stable within-expert offset: out[p] = #{q < p | flat_idx[q] == flat_idx[p]}

Register a boolean operator handler.