Parallel AVX2 integer INT4 matmul — splits output rows across rayon threads. The integer kernel is ~2x faster per-element than FMA, so parallelism on

Scalar INT4 matrix-vector multiply. Matches AVX2 computation semantics: scale applied per-group, not per-element.

Safe wrapper for AVX2 FMA transposed INT4 matmul.

Vectorized INT4 matmul simulating our Rust kernel. For each row, for each group: acc += sum(signed_nibble * activation) * scale

Safe wrapper for AVX2 integer transposed INT8 matmul. Weight layout: data as i8 packed into u32 in [K, N] layout, scales [K/gs, N] as BF16. The u32 v

Get the NUMA node for a specific expert.

Extract the layer number from a tensor name like "model.layers.42.mlp.experts.0.gate_proj.weight". Returns None if no ".layers.N." pattern is found.

Parse SSE response body into (full_text, finish_reason).

Construct A Single Stage's Accumulator Shape

Pin the current thread to a specific NUMA node. This restricts the thread to CPUs on the given node.

Gets the index of the problem

Profile prefill scaling for a given model/mode. Args: use_single_call: If True, pass all tokens in one forward() call. Requir

Reset min-free tracking on all devices (e.g. after warmup, before runtime).

Set data SMEM ptr

Overrides the internal iteration index

Sets the predicate mask, overriding value stored in predicate iterator

Entry of the subtile load loop. Aux loads usually performed here Upon entry the producer acquire of the current subtile lock has completed. Upon exit

Stores a fragment to memory

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

Store a fragment to memory

Send streaming chat request. Prints tokens as they arrive. Returns (display_text, raw_text, timing) — display_text has channel markers stripp

Whether this block will perform the first iteration of this tile

Whether this block will perform the first iteration of this tile

Reset thread affinity to allow running on any node.

Wait for GPU memory to drop below 500 MB used. If memory doesn't clear within half the timeout, kills any stale krasis processes that may be

Capture a CUDA graph for the shared expert forward (M=1). Replaces ~6 kernel launches (0.64ms Python overhead) with a single graph re

Grow a text prefix until it contains enough raw tokens.

Compute PP partition — always PP=1 (all layers on primary GPU). Multi-GPU uses Expert Parallelism (EP), not Pipeline Parallelism.

Expose the WSL2 host driver binaries/libraries to subprocesses.

Extract the weights dict from an existing TransformerLayer. Reconstructs the dict format that TransformerLayer.__init__ expects, so w

Graph-compatible M=1 recurrent forward using in-place state updates. Same computation as _forward_recurrent but: - _conv_state update

Return ranking with heatmap experts first, then unranked experts for [layer_start, layer_end).

Get the PP rank index that owns a given layer.

Determine which CUDA toolkit version to install. Uses the minimum toolkit version needed for the GPU architecture. The CUDA version printed b

Get the GPU staging buffers for this layer's attention type. buf_idx: 0 or 1 (ping-pong buffer index).

Return a balanced deterministic per-layer ranking without heatmap collection.

Dispatch to INT8 or BF16 linear based on weight type.

Check if Python development headers are installed (needed for pyo3).

Probe torch in a subprocess so reinstall checks see the actual package.

Read /proc/meminfo and return dict of key -> value in KB.

Read and normalize model config.json.

Read model info from config.json for display.

vLLM-style deterministic WHT rotation signs for one attention layer.