QWIX Fixes 🪄
To preserve gradient flow during backpropagation, Qwix QAT mode does not actually cast weights to low-bit integers; instead, it does this to simulate lower-bit math while staying in 32-bit float registers.
What is fake quantization?
This famous memory-efficient mechanism avoids materializing the massive O(L2) attention matrix by tiling and computing online softmax in blocks.
What is FlashAttention (or SplashAttention)?
Developed to avoid FP16 underflow, this 16-bit floating-point format retains FP32’s 8-bit exponent range, eliminating the need for dynamic loss scaling.
What is bfloat16?
This JAX operation automatically maps a function over array axes, eliminating the need for explicit, inefficient Python loops.
What is vmap?
First introduced as a fully programmable unit in Viperfish, this core offloads embedding lookups and sparse updates from the dense processing units.
What is SparseCore?
One of the main performance benefits of weight-only quantization is reducing this major hardware bottleneck, even if the actual matrix multiplication computation is still performed in floating-point.
What is Memory Bandwidth (or Memory Footprint / I/O)?
To turn idle wait-time into peak FLOPS, custom Pallas kernels use this technique to fetch the next block of data from slow HBM into VMEM while simultaneously crunching math on the current block.
What is double buffering?
On a roofline plot, performance is plotted on the y-axis against this ratio on the x-axis.
What is Arithmetic Intensity (FLOPs per Byte)?
This core JAX transformation automatically calculates the exact derivative of a scalar-valued function, making reverse-mode backpropagation a breeze.
What is jax.grad?
Starting with Pufferfish, this feature allows two TensorCores sharing the same HBM package to work together transparently as a single logical unit.
What is Megacore?
DAILY DOUBLE!!
What are SafeTensors?
The Tokamax layer norm kernel supports standard layer norm as well as this computationally cheaper, mean-free variant that has become the standard in various models.
What is RMS Norm?
To maximize TPU systolic array efficiency, APEX is investigating this micro-scaling floating-point format as an alternative to standard INT8.
What is MXFP4 (or MX format)?
When you apply the @jax.jit decorator, JAX doesn't compile the Python code directly. Instead, it traces your function to produce this specific, strongly-typed Intermediate Representation before handing it off to XLA.
What is Jaxpr (JAX expression)?
DAILY DOUBLE!!
What are 8 and 128?
Fusing quantization with this specific mathematical transform, which relies on random sign flips, allows for efficient overlapping of the MXU and VPU.
What is Hadamard Transform?
DAILY DOUBLE!!
What is Asynchronous Remote Copy (or explicit DMAs)?
DAILY DOUBLE!!
What is the prefill phase?
To manually parallelize code across multiple hardware accelerators, developers traditionally relied on this JAX operation, though modern workflows are shifting toward jit with NamedSharding?
What is jax.pmap?
Handling tasks too complex for independent, isolated VPU lanes, this specialized subsystem is responsible for operations such as data shuffles and transposes.
What is the Cross-Lane Unit (XLU)?
To smooth out the quantization difficulty between dynamic activations and static weights, Qwix supports this algorithm which applies a mathematical "smoothing factor" across channels before clipping.
What is SmoothQuant?
You can write perfectly parallel code for all 128 VPU lanes, but if your algorithm asks two different sublanes within a lane to access different rows of data on the same clock cycle, the chip will instantly stall due to this performance-killing hardware violation.
What is a Bank Conflict?
To trade compute for a lower memory footprint, this technique drops intermediate forward-pass activations and recomputes them on the fly during the backward pass.
What is activation checkpointing (or gradient rematerialization)?
Standard Python if/else statements fail inside a jit-compiled function when the condition depends on dynamic array values; instead, you must use this specific JAX control flow primitive.
What is jax.lax.cond?
ZebraFish's structured sparsity specifically enforces this exact mathematical ratio on both the left-hand and right-hand sides to accelerate GEMMs.
What is 1:4 (or 1-in-4) sparsity?