Inside the kernels

What fastrerandomize computes, when, and why.

FastRerandomize selects treatment assignments that balance covariates between treatment and control groups. Once outcomes are observed, it reuses the retained assignment pool for randomization tests and treatment-effect intervals. Rerandomization draws treatment assignments, measures covariate balance, and retains assignments that meet the balance criterion. At acceptance probability q = 0.001, retaining 1,000 assignments requires screening a million candidates on average, so the per-draw cost of T²(w) limits the feasible numbers of units n and covariates p.

On this page, we explain four changes that reduce repeated calculations and memory use in FastRerandomize, compare measured runtimes with hardware limits, and check numerical accuracy.

To measure computational cost at scale, we use a reference workload with 1,048,576 assignments for 4,096 units on 2,048 covariates. We retain the best 1,049. On a laptop GPU the complete R call takes 3.41 seconds.RTX 5080 Laptop GPU, float32. Median of seven warmed calls; the first call, which compiles the kernels, took 11.96 s. Full report.

When optimized fully, the compiled JAX pipeline generates, validates, and scores about 355,000 assignments per second; it uses covariates centered and scaled or whitened once per call on the host, one general matrix multiplication (GEMM) per batch, and permutations generated from counter keys.

• Before your RCT/experiment: covariates → candidate balance scores → retained assignment pool → one assignment drawn uniformly from its rows.
• After your RCT/experiment: retained pool + actual treatment assignment + observed outcomes → p-value and treatment-effect interval.

Background and Notation

Here, we have a covariate matrix X ∈ ℝ^{n × p}, which has one unit per row and one covariate per column. An assignment w ∈ {0, 1}ⁿ has entry w_j = 1 for treatment and 0 for control, for units j = 1, …, n. Every generated assignment has n₁ = Σ_j w_j treated units and n₀ = n − n₁ controls; these group sizes stay fixed across draws. The fixed treatment count also provides a validity check: every generated assignment must contain exactly n₁ ones.

A batch of B draws forms the rows of W ∈ {0, 1}^{B × n}. X_w denotes the centered and transformed covariates: standardized under diagonal scoring, or whitened under full-covariance scoring (Section 1). The product s = WX_w is a B × p matrix of treated-group covariate sums. Each row of s corresponds to one assignment and is later used to compute its T²(w) balance score, which captures how different the treatment and control groups are in their covariates.

Kernel changes

1. Scale or whiten once, in double precision

The sample covariance matrix of X is S; x̄₁ and x̄₀ are the treated and control covariate mean vectors under assignment w. Full-covariance scoring (approximate_inv = FALSE) uses Hotelling’s balance statistic:

The default, approximate_inv = TRUE, scales each covariate by its population standard deviation and ignores cross-covariances. For covariates k = 1, …, p, let X̄_k denote the column mean and v_k its population variance. With diagonal scoring,

Constant columns receive zero weight. Full-covariance scoring, by contrast, accounts for correlations: its direct quadratic form costs O(p²) per draw once the group means are known. But because X and S⁻¹ are fixed across draws, we instead form X_w = X_cL once. Here, X_c is X with each column’s mean subtracted, and L satisfies LLᵀ = S⁻¹. The quadratic form then reduces to an O(p) squared Euclidean norm of the transformed mean difference: T²(w) = (n₀n₁/n) ‖Lᵀ(x̄₁ − x̄₀)‖².

Both the diagonal and non-diagonal criteria use the same transformed-norm kernel. In each call, we prepare X_w in float64 on the host with .balance_design(), then transfer it to the device as float32. Our reference benchmark uses default diagonal scoring with the built-in threshold function, which skips covariance construction and prepares X_w in time linear in the size of X: 0.345 s here.

For either of the two criteria, centering makes the control sum the negative of the treated sum (in exact arithmetic), so the mean difference can be computed from the row vector of treated-group sums:

Writing either score as T²(w) and using 1 for the length-n all-ones vector, we retain a column-sum correction for rounding:

Computing one group sum instead of two halves the arithmetic cost. In exact arithmetic, centering makes 1ᵀX_w = 0; in float32, subtracting 1ᵀX_w/n₀ accounts for any nonzero column sum when recovering the control-group mean from s(w). WhitenedT2, six lines of R calling JAX. Centering also reduces cancellation error. In float32, subtracting nearly equal group means loses precision: with covariates near 1,000, a non-whitened per-draw kernel had a relative error of 3.4 × 10⁻³ against a float64 reference, and the whitened kernel 2.5 × 10⁻⁷.Our measurement on 1,024 draws, n = 1,024, p = 64, maximum over draws.

2. One matrix product per batch

Now, to see why this decomposition works—and why it makes batching possible—note that the control indicators are 1 − w, so the control-group sum is simply the total column sum minus the treated-group sum: (1 − w)ᵀX_w = 1ᵀX_w − s(w). Subtracting the whitened control mean from the treated mean collects s(w) into a single term:

Because X_w, n₁, and n₀ are fixed across draws, the linear projection s(w) = wᵀX_w is the only part of T²(w) that depends on the candidate assignment w. Computing s = WX_w combines all B covariate sums into a single GEMM. So, the B × p GEMM result stacks the individual row vectors s(w) in the same order as the assignments in W. A tiled GEMM reuses blocks of W and X_w in shared memory and registers, reducing repeated loads from device memory.

We quantify that data reuse through arithmetic intensity: FLOPs per byte transferred. An ideal whole-product estimate counts one read of W and X_w and one write of s. The ridge point is peak compute over peak bandwidth, about 39 FLOP/B on this GPU. In this model, bandwidth sets the throughput bound below the ridge and compute sets it above. The roofline bounds throughput by the lesser of peak FLOP/s and intensity times peak bandwidth. For F FLOPs and D bytes, the corresponding speed-of-light time floor is max(F / (peak FLOP/s), D / (peak bytes/s)). This floor provides a reference for balance-scoring time; the full pipeline also includes permutation generation and score sorting.

At the reference shape, that ideal whole-product intensity is 512 FLOP/B. The profiler reports a 256 × 128 threadblock tile. Counting input loads separately for each tile gives 0.5 × 256 × 128 / (256 + 128) = 42.7 FLOP/B, just above the 39 FLOP/B ridge.

GEMM for pairwise distances

The package’s standalone fast_distance() utility applies the same matrix-product strategy to pairwise Euclidean and Mahalanobis distances. It operates outside the assignment-screening pipeline and reduces temporary memory for distance calculations. For two datasets with n and m rows and p covariates, let a, b ∈ ℝ^p be feature vectors from the respective datasets. Expanding ‖a − b‖² = ‖a‖² + ‖b‖² − 2aᵀb avoids materializing the n × m × p difference tensor: at n = m = 1,024 and p = 256, a 4 MB inner-product matrix replaces the 1.07 GB tensor.

Choosing the batch size

For assignment generation, the automatic batch-size rule targets 2²² entries × 4 bytes = 16 MiB for the float32 assignment matrix W. By default, .resolve_batch_size() rounds 2²²/n down to a power of two, clamps B to [1,024, 65,536], and caps it at max_draws. Power-of-two rounding is a tuning heuristic; the 1,024-draw minimum can exceed the memory target at large n.

The reference benchmark explicitly sets B = 4,096 at n = 4,096, so its assignment buffer is 64 MiB (67.1 MB). The automatic rule would choose B = 1,024 for this workload.

3. Keep keys, not assignments

A pool of 1,048,576 assignments for 4,096 units would occupy 4.29 GB even as int8; only 1,049 assignments are retained. We can regenerate each full assignment w⁽ⁱ⁾ from its global draw index i and a shared root pseudorandom-number generator (PRNG) key, root_key. Its per-draw key is a pure function of (root_key, i), computed by CounterKeys. The device computes these keys as needed, without materializing the full key array. For a fixed root key, draw i always gets the same PRNG key regardless of batch size B: changing B changes throughput, but produces the exact same sequence of treatment assignments.

To turn a per-draw key into an assignment, JAX shuffles a length-n int32 payload containing n₁ ones and n₀ zeros. We use int32 because sorting int8 payloads produced incorrect treated counts on ROCm for n ≥ 256. JAX generates pseudorandom integer sort keys from the draw’s PRNG key and sorts the payload by those keys. Reproducing existing seeds requires preserving the complete mapping from (root_key, i) to an assignment, including this permutation step. Each batch generates a B × n assignment matrix in scratch memory and writes B float32 scores. Only the scores are retained across batches.

To avoid a host round trip per batch, .mc_scorers() runs all 256 batches inside one compiled program on the reference NVIDIA setup. Metal defaults to R dispatch of the compiled per-batch kernels, with a host dispatch cost for every batch. After a final argsort ranks the full score vector, the accepted assignments are regenerated from their indices.

4. Precompute candidate statistics for fiducial intervals

We select the implemented assignment uniformly from the retained pool and use that same pool for inference. After the experiment, let w_obs denote that assignment and Y ∈ ℝⁿ the observed outcomes. The package’s fiducial interval inverts randomization tests over the retained pool: it retains constant treatment effects τ for which neither one-sided test rejects at α/2, where α is the significance level.

Under this selection rule and the constant additive-effect model, enumerating every retained row with equal weight gives exact, potentially conservative randomization tests conditional on the pool. Exact inversion of the two one-sided tests has coverage of at least 1 − α conditional on that pool. A different selection rule requires inference using the actual selection probabilities.

For any vector v ∈ ℝⁿ, the difference in means under candidate assignment w is:

The observed effect is τ̂_obs = Δ_wobs(Y). Under the sharp null hypothesis of a constant treatment effect τ, the outcomes for candidate assignment w are:

Affine form in τ

Each candidate assignment fixes the weights in the difference in means, and the hypothesized effect shifts outcomes linearly, so its statistic follows a straight line; interval inversion needs only that line’s intercept and slope.

Define the candidate statistic D_τ(w) = Δ_w(Y^{(w, τ)}). For a binary assignment, it is affine in τ:

Under the constant additive-effect model, the affine identity is exact in real arithmetic. The finite endpoint search approximates the interval, and floating-point arithmetic introduces rounding error.

For the fixed-count assignments considered here, c_w = 1. At w = w_obs, b_w = c_w = 1, so the observed line has slope zero and height τ̂_obs.

On the device, CandidateMoments computes three raw sums per retained assignment in one pass. With Ȳ = (Σ_j Y_j)/n denoting the mean observed outcome, the sums give the following coefficients in exact arithmetic for the fixed treatment counts:

On the host, .fi_components() converts these sums into the intercept a_w and slope c_w − b_w once per candidate. At τ = 0, D₀(w) = a_w. The API’s two-sided p-value for zero effect compares these intercepts’ magnitudes with |τ̂_obs|.

To locate the interval endpoints, R makes 500 stochastic updates. Each update samples one retained assignment and adjusts the lower and upper τ estimates toward their tail-probability crossings. It then evaluates all candidates at 100 τ values on a grid expanded beyond these estimates, returning the range of grid values whose smaller empirical tail probability exceeds α/2. Each candidate evaluation uses one multiply-add; the search makes no further device calls.

Benchmarks and numerical checks

Change one thing at a time

The roofline models memory reuse in FLOP/byte; elapsed-time speedups also depend on cache behavior and execution overhead. To compare algorithmic changes under a fixed CPU resource budget, we benchmarked standalone Python implementations on two cores after compilation and warm-up, changing one algorithm per comparison.Two CPU cores of a cloud VM, JAX 0.12.0.1, medians of seven warmed calls.

With 256 or more covariates the balance rewrite is worth 1.5–2.2× on this CPU, close to the 2× reduction in FLOPs the cost model predicts. At these matrix sizes, X_w likely stays in CPU cache across draws, which would make arithmetic the limiting cost of the unbatched kernel. At 16 covariates both kernels are per-column reductions and the gain is small.

In fast_distance(), replacing the O(nmp) broadcast tensor with a matrix product cuts runtime by 8.2–9.6×. For the fiducial search in randomization_test(), precomputing each candidate’s intercept and slope once—instead of taking n-element inner products at every trial τ—cuts search time by 76–158×.

In a separate GPU comparison with all optimizations enabled, repeated whole-function calls on an RX 7600M XT ran 10–110× faster for Monte Carlo generation and 6–180× faster for randomization_test(), as reported in the package NEWS.

Sorting accounts for 61% of kernel time

Profiling the reference workload on the RTX 5080 Laptop GPU gives 3.7 ms per batch for balance checks (32%) and 7.1 ms for the sorting used to generate permutations (61%). All other kernels take 0.9 ms. By Amdahl’s law, eliminating the balance product would yield at most a 1.47× speedup; eliminating the sort would yield at most 2.55×.

Permutation sampling is the next optimization target. A draw needs only a random subset of n₁ units, so a selection could replace the full sort. This would change the mapping from key to assignment, so existing seeds would no longer reproduce the same assignments. Conjecture; not implemented.

Four safeguards for float32

The batch loop runs in float32 (24-bit significand, 4 bytes per value) and uses four safeguards to ensure numerical accuracy, valid assignments, and reproducibility:

• Float64 for cancellation-prone steps. Host-side centering, scaling or whitening, and the observed effect τ̂_obs use float64. Float32 rounding can change the result of the two-sided p-value comparison |a_w| ≥ |τ̂_obs| for near-ties. Using the mean observed outcome Ȳ, the near-tie tolerance is: ε = 4n·2⁻²⁴ (Σ_j |Y_j − Ȳ|) (1/n₁ + 1/n₀) Here 2⁻²⁴ is float32 unit roundoff. The factor n allows for accumulated summation error; 4 is the implementation’s extra margin. Candidates satisfying | |a_w| − |τ̂_obs| | ≤ ε are re-evaluated in float64 before this comparison.
• True FP32 products. For the balance and distance GEMMs, we pass Precision.HIGHEST to JAX so NVIDIA GPUs compute float32 dot products with 24-bit significand precision, instead of rounding inputs to TensorFloat-32 (10 fraction bits plus an implicit leading bit). This precision choice can reduce throughput relative to lower-precision modes because full FP32 products cannot use the faster TF32 Tensor Core path.
• Sort payloads. We use int32 payloads and check that every generated assignment contains exactly n₁ treated units.
• Reproduction check. Each accepted draw is regenerated from its key and re-scored. A mismatch raises an error naming the device.

In the reference run, 64 accepted scores checked against a float64 R reference had a maximum relative error of 9.07 × 10⁻⁸.

Regime, floor, and next lever

The 0.501 s floor covers balance arithmetic across 256 batches (1,048,576 draws). The measured GPU stage takes 2.951 s for generation, validation, balance scoring, and score ranking; the complete R call takes 3.410 s including host preparation and bookkeeping. Permutation sorting accounts for 61% of profiled GPU kernel time and is excluded from the balance-only floor. Per-batch kernel medians and whole-call medians were measured separately, so their sums need not match.

To reproduce the reference workload:

library(fastrerandomize)
X <- matrix(rnorm(4096 * 2048), 4096, 2048)
pool <- generate_randomizations(
  n_units = 4096, n_treated = 2048, X = X,
  randomization_type = "monte_carlo",
  max_draws = 2^20, batch_size = 4096,
  randomization_accept_prob = 0.001
)

In conclusion, by moving fixed calculations out of the draw loop, batching balance checks, regenerating assignments from keys, and reusing candidate statistics during inference, we make stringent balance criteria and treatment-effect inference practical for experiments with more units and covariates—while enforcing explicit safeguards for numerical accuracy and reproducibility.

Sources and vocabulary. Benchmarks: Speed of Light report and reproduction instructions. Roofline, arithmetic intensity, and speed-of-light reasoning follow Hwu, Kirk, and El Hajj, Programming Massively Parallel Processors, 5th ed.

References

Rebecca Goldstein, Connor T. Jerzak, Aniket Kamat, and Fucheng Warren Zhu. (2026). FastRerandomize: Fast rerandomization using accelerated computing. SoftwareX, 33, 102508. [PDF]. [BibTeX].