Crate aria_kernel 

Scalar (+ aarch64 NEON / x86 AVX2 + optional CUDA) kernels for Aria engine.

Structs

CudaContext

Persistent cuBLAS handle + device copies of host weight buffers.

GatedDeltaStep

Bundled args for gated_delta_step (avoids clippy too_many_arguments).

Enums

ComputeBackend

Resolved backend used by Session GEMM.

ComputePref

CLI / config preference. Orthogonal to hybrid_execution.

EngineError

Unified engine error (requirements §3.6).

SimdMode

Runtime SIMD selection. Tests force SimdMode::Scalar.

Functions

attention

Causal attention for single query step against KV cache. q: [n_heads * head_dim], k_cache/v_cache: [seq, n_kv_heads * head_dim]

attention_causal

Causal attention for a batch of queries [seq_q, n_heads*head_dim]. When seq_q == seq_kv, query t attends to keys 0..=t (prefill). When seq_q == 1, equivalent to attention (decode).

attention_causal_with_scale

Prefill causal attention with an explicit softmax scale. window is the sliding-window length (None = full causal prefix).

attention_with_scale

Causal attention with an explicit softmax scale (Gemma-4 uses 1.0 after QK-norm).

cpu_simd_label

dequant_lookup_group

Codebook lookup dequant (group share): indices [k_work, n], codebook [g, kc].

fwht

In-place orthogonal FWHT on length = power of two (scale 1/sqrt(n)).

gated_delta_step

One-token Gated DeltaNet recurrence (Qwen3.5 / Bonsai linear attention).

geglu

GeGLU: gelu(gate) * up (Gemma gelu_pytorch_tanh approximates with tanh form).

gelu_pytorch_tanh

Match transformers gelu_pytorch_tanh (used by Gemma GeGLU).

hadamard_blocked_rows

Apply blocked Hadamard on rows of a row-major [rows, cols] matrix. inverse=false → per-block H@S; inverse=true → S@H.

hadamard_blocked_vec

Blocked unrotate on a length-rows vector (treat as [rows, 1]).

hdm_linear

y = W_rot @ x followed by blocked unrotate on each out_f row (HDM fused path).

kv_sliding_view

Restrict KV to the last window tokens (HF sliding-window). None is a no-op.

linear

y = x @ W^T where W is [out_features, in_features] row-major (GGUF-style).

linear_cpu

Multi-threaded linear (AVX2/FMA or NEON dots). Numerically close to linear.

matmul

C = A @ B^T style? We use row-major: out[m,n] = sum_k a[m,k] * b[k,n] with a: [m,k], b: [k,n].

matmul_blocked

Blocked matmul used as Neon / SIMD-friendly path (portable; aarch64 may specialize later).

matmul_dispatch

Dispatch scalar vs Neon (blocked) paths. Neon is available on all targets for parity tests; on aarch64 this is the production SIMD entry (intrinsics may replace the body later).

moe_topk_route

Softmax (or sigmoid) top-k MoE routing. Returns (expert_ids, normalized weights).

portable_block_signs

Portable ±1 signs matching Python portable_block_signs.

pow2_tile_sizes

Greedy largest-pow2 tiling of row count (e.g. 10 → [8, 2]).

resolve_compute

Resolve preference. Cuda never silently falls back to CPU.

rms_norm

rms_norm_gemma

Gemma-style RMSNorm: x * rrms * (1 + weight).

rope

Apply RoPE to interleaved q/k pairs for one token (head_dim even).

rope_half

HF Llama/Qwen/Gemma RoPE: rotate half of the head dims as a contiguous block.

rope_half_partial

Rotate only the first rotary_dim dims of each head (partial_rotary_factor).

rope_half_proportional

Gemma-4 global (p-RoPE): rotate the first factor * head_dim/2 pairs of rotate_half layout; remaining pairs stay identity. Frequencies use the full head_dim denominator (not the rotated subset).

short_conv_step

Causal depthwise short-conv one-token step (LFM2).

silu_vec

Elementwise SiLU (DeltaNet conv activation).

softmax

softmax_inplace

softplus

Numerically stable softplus.

swiglu