Test the operator, its derivatives, and its scaling separately

Use small cases that can be calculated by hand before comparing larger tensors against a higher-precision implementation. Uniform logits test the allowed-key average directly. A single dominant score tests stable normalization, while a masked dominant score tests whether exclusion happens before the maximum and denominator. Distinct sequences and heads reveal accidental broadcasting that identical random distributions can hide. Fix dropout to zero for deterministic tests, then validate dropout behavior separately.

Choose elementwise absolute-plus-relative tolerances based on dtype, accumulation, shape, and the reference’s accuracy. Relative error alone is unstable near a zero reference value, while a large absolute tolerance can hide meaningful errors in small outputs. Record maximum and aggregate errors, inspect failing coordinates, and verify finiteness. A tolerance is not permission to ignore systematic mask or head-mapping differences whose values happen to fit under one loose threshold.

Softmax is invariant to adding a common finite constant to all allowed logits in a row. Consistently permuting keys, values, and mask columns preserves a noncausal row’s output. These transformations make useful metamorphic tests, provided positional biases and causal relations are transformed consistently rather than accidentally changed. For GQA, compare a shared-KV implementation with a reference that explicitly repeats K and V according to the documented grouping, including the summed backward gradients.

Finite differences on a tiny smooth case help detect a missing scale or denominator term in backward. Use a central difference at a step that balances truncation and floating-point cancellation, and avoid crossing a discontinuous mask or routing change. A successful spot check does not prove the whole gradient tensor, so combine it with reference gradients for varied shapes, ragged boundaries, empty-row extensions, and the same logical dropout realization when replay is part of the contract.

Sweep sequence length while holding head width, feature width, batch, pattern width, dtype, and mask family fixed as appropriate. Dense attention retains quadratic leading pair work even when tiled saved state is linear. A bounded window has linear logical pairs, while a window that grows with N does not. A feature-map method’s linear claim similarly assumes a controlled feature width. Report the actual parameter values beside any fitted scaling curve.

Measure peak memory after isolating allocations and distinguish allocated from reserved pools where the tooling permits. Measure time with synchronized device timing or an equivalent documented method; asynchronous launch return time is not kernel completion time. Include kernel selection and fallback behavior, because the same high-level call can select different backends across shapes. A scaling test should explain a mechanism and identify its boundaries, not turn a few timing samples into an unsupported asymptotic proof.