diff --git a/Project.toml b/Project.toml
index 3d75c66..b18d325 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,6 +1,6 @@
 name = "FixedEffects"
 uuid = "c8885935-8500-56a7-9867-7708b20db0eb"
-version = "3.4.1"
+version = "3.5.0"
 
 [deps]
 GroupedArrays = "6407cd72-fade-4a84-8a1e-56e431fc1533"
diff --git a/benchmarks/akm_worker_firm.jl b/benchmarks/akm_worker_firm.jl
deleted file mode 100644
index 836a27c..0000000
--- a/benchmarks/akm_worker_firm.jl
+++ /dev/null
@@ -1,209 +0,0 @@
-using FixedEffects
-using LinearAlgebra
-using Printf
-using Random
-using StatsBase
-
-# Simulated counterpart to the xhdfe AKM-style performance example:
-# log wage on seniority controls, absorbing worker, firm, and year fixed effects.
-# This package is the residualization backend, so worker-clustered SEs are out of scope here.
-# Usage: julia --project -t auto benchmarks/akm_worker_firm.jl [--seed=1234] [--method=cpu]
-#        [--double-precision=true] [--tol=1e-8] [--maxiter=Inf]
-const N_WORKERS = 50_000
-const N_FIRMS = 7_000
-const N_YEARS = 36
-const OBS_PER_WORKER = 8
-const FIRM_WINDOW = 10
-const MOVE_PROBABILITY = 0.35
-const CONTROL_NAMES = ("tenure", "tenure_sq", "experience", "experience_sq", "mover")
-
-struct AKMPanel
-    worker::Vector{Int32}
-    firm::Vector{Int32}
-    year::Vector{Int32}
-    y::Vector{Float64}
-    x::Matrix{Float64}
-end
-
-Base.length(panel::AKMPanel) = length(panel.y)
-
-function parse_options(args)
-    options = Dict{String,String}()
-
-    for arg in args
-        startswith(arg, "--") || error("Unexpected positional argument '$arg'. Use --key=value options.")
-        key_value = split(arg[3:end], "=", limit = 2)
-        length(key_value) == 2 || error("Expected --key=value, got $arg")
-        options[key_value[1]] = key_value[2]
-    end
-
-    return options
-end
-
-option(options, key, default::Int) = parse(Int, get(options, key, string(default)))
-option(options, key, default::Float64) = parse(Float64, get(options, key, string(default)))
-
-function option(options, key, default::Bool)
-    value = lowercase(get(options, key, string(default)))
-    value in ("true", "yes", "1") && return true
-    value in ("false", "no", "0") && return false
-    error("Expected boolean for --$key, got '$value'")
-end
-
-function method_option(options)
-    value = lowercase(get(options, "method", "cpu"))
-    value == "cpu" && return :cpu
-    value == "cuda" && return :CUDA
-    value == "metal" && return :Metal
-    error("Expected --method=cpu, --method=CUDA, or --method=Metal, got '$value'")
-end
-
-function maxiter_option(options)
-    value = lowercase(get(options, "maxiter", "inf"))
-    value in ("inf", "infinity") && return typemax(Int)
-    return parse(Int, value)
-end
-
-function load_backend!(method::Symbol)
-    if method == :CUDA
-        @eval using CUDA
-        CUDA.functional() || error("CUDA was requested but CUDA.functional() is false")
-    elseif method == :Metal
-        @eval using Metal
-    end
-    return nothing
-end
-
-function local_firm(rng::AbstractRNG, anchor::Int)
-    lo = max(1, anchor - FIRM_WINDOW)
-    hi = min(N_FIRMS, anchor + FIRM_WINDOW)
-    return rand(rng, lo:hi)
-end
-
-function simulate_akm_panel(; seed::Int)
-    rng = MersenneTwister(seed)
-    n = N_WORKERS * OBS_PER_WORKER
-    worker = Vector{Int32}(undef, n)
-    firm = Vector{Int32}(undef, n)
-    year = Vector{Int32}(undef, n)
-    y = Vector{Float64}(undef, n)
-    x = Matrix{Float64}(undef, n, length(CONTROL_NAMES))
-
-    worker_effect = randn(rng, N_WORKERS)
-    firm_effect = 0.6 .* randn(rng, N_FIRMS)
-    year_effect = [0.02 * (t - 1) + 0.05 * sin(2 * pi * (t - 1) / N_YEARS) for t in 1:N_YEARS]
-
-    row = 1
-    max_start_year = max(1, N_YEARS - min(OBS_PER_WORKER, N_YEARS) + 1)
-    for w in 1:N_WORKERS
-        anchor = clamp(1 + fld((w - 1) * N_FIRMS, N_WORKERS), 1, N_FIRMS)
-        current_firm = local_firm(rng, anchor)
-        start_year = rand(rng, 1:max_start_year)
-        base_experience = rand(rng, 1:20)
-        tenure = 0
-
-        for spell_t in 1:OBS_PER_WORKER
-            moved = spell_t > 1 && rand(rng) < MOVE_PROBABILITY
-            if moved
-                current_firm = local_firm(rng, anchor)
-                tenure = 0
-            elseif spell_t > 1
-                tenure += 1
-            end
-
-            calendar_year = 1 + mod(start_year + spell_t - 2, N_YEARS)
-            experience = base_experience + spell_t - 1
-            tenure_sq = tenure^2 / 100
-            experience_sq = experience^2 / 100
-            mover = moved ? 1.0 : 0.0
-
-            worker[row] = Int32(w)
-            firm[row] = Int32(current_firm)
-            year[row] = Int32(calendar_year)
-            x[row, 1] = tenure
-            x[row, 2] = tenure_sq
-            x[row, 3] = experience
-            x[row, 4] = experience_sq
-            x[row, 5] = mover
-            y[row] = 0.04 * tenure - 0.03 * tenure_sq +
-                0.015 * experience - 0.02 * experience_sq +
-                0.05 * mover +
-                worker_effect[w] + firm_effect[current_firm] +
-                year_effect[calendar_year] + 0.2 * randn(rng)
-            row += 1
-        end
-    end
-
-    order = randperm(rng, n)
-    return AKMPanel(worker[order], firm[order], year[order], y[order], x[order, :])
-end
-
-function akm_estimator_call(panel::AKMPanel;
-        method::Symbol,
-        double_precision::Bool,
-        tol::Real,
-        maxiter::Integer)
-    y = copy(panel.y)
-    x = copy(panel.x)
-    fes = [FixedEffect(panel.worker), FixedEffect(panel.firm), FixedEffect(panel.year)]
-    T = double_precision ? Float64 : Float32
-    solver = AbstractFixedEffectSolver{T}(fes, uweights(T, length(y)), Val{method})
-    variables = Vector{AbstractVector{Float64}}(undef, 1 + size(x, 2))
-    variables[1] = y
-    for j in axes(x, 2)
-        variables[j + 1] = view(x, :, j)
-    end
-    _, iterations, converged = solve_residuals!(variables, solver;
-        tol = tol,
-        maxiter = maxiter,
-        progress_bar = false)
-    beta = x \ y
-    return (beta = beta, iterations = iterations, converged = converged)
-end
-
-function print_result(result)
-    @printf("  beta:")
-    for (name, value) in zip(CONTROL_NAMES, result.beta)
-        @printf(" %s=% .4f", name, value)
-    end
-    println()
-    println("  iterations: ", join(result.iterations, ", "))
-    println("  converged:  ", join(result.converged, ", "))
-end
-
-options = parse_options(ARGS)
-seed = option(options, "seed", 1234)
-method = method_option(options)
-double_precision = option(options, "double-precision", method == :cpu)
-tol = option(options, "tol", double_precision ? 1e-8 : 1e-6)
-maxiter = maxiter_option(options)
-
-load_backend!(method)
-
-n = N_WORKERS * OBS_PER_WORKER
-println("Simulated AKM benchmark")
-println("  observations:    ", n)
-println("  worker FE:       ", N_WORKERS)
-println("  firm FE:         ", N_FIRMS)
-println("  year FE:         ", N_YEARS)
-println("  controls:        ", join(CONTROL_NAMES, ", "))
-println("  method:          ", method)
-println("  double precision:", double_precision)
-println("  tol:             ", tol)
-println("  maxiter:         ", maxiter == typemax(Int) ? "Inf" : maxiter)
-println("  note: this times FE construction, residualization of y and controls, and dense OLS; clustered SEs are not included.")
-
-panel = simulate_akm_panel(; seed = seed)
-@printf("  panel memory:    %.1f MiB\n", Base.summarysize(panel) / 2.0^20)
-
-println("\nWarmup")
-warmup = akm_estimator_call(panel; method, double_precision, tol, maxiter)
-print_result(warmup)
-
-println("\nTimed run")
-GC.gc()
-timed = @timed akm_estimator_call(panel; method, double_precision, tol, maxiter)
-@printf("  time:       %.3f s\n", timed.time)
-@printf("  allocated:  %.1f MiB\n", timed.bytes / 2.0^20)
-@printf("  gc time:    %.3f s\n", timed.gctime)
-print_result(timed.value)
diff --git a/benchmarks/benchmark.jl b/benchmarks/benchmark.jl
new file mode 100644
index 0000000..275e4a1
--- /dev/null
+++ b/benchmarks/benchmark.jl
@@ -0,0 +1,173 @@
+##############################################################################
+# Benchmarks for solve_residuals! across backends (:cpu, :CUDA, :Metal).
+# Problems are defined once; every available backend runs the same
+# specifications. Each specification residualizes 6 columns through one
+# solver, which is how FixedEffectModels calls this package (first run
+# includes compilation).
+# Usage: julia --project -t auto benchmarks/benchmark.jl
+##############################################################################
+
+using FixedEffects, Random, StatsBase
+try using CUDA catch end
+try using Metal catch end
+Random.seed!(1234)
+
+##############################################################################
+# Setup
+##############################################################################
+
+# Simple problem: N=10M, two FEs (100k × 100 groups)
+N = 10_000_000
+K = 100
+id1 = rand(1:div(N, K), N)
+id2 = rand(1:K, N)
+fes_simple = [FixedEffect(id1), FixedEffect(id2)]
+cols_simple = [rand(N) for _ in 1:6]
+
+# Hard problem: N=800k, worker-firm (40k × 5k), same construction as the
+# "difficult" setup in FixedEffectModels' benchmark/benchmark.jl
+N = 800_000
+M = 40_000
+O = 5_000
+pid = rand(1:M, N)
+fid = [rand(max(1, div(x, 8)-10):min(O, div(x, 8)+10)) for x in pid]
+fes_hard = [FixedEffect(pid), FixedEffect(fid)]
+cols_hard = [rand(N) for _ in 1:6]
+
+# Interacted fixed effects: one regressor interacted with both hard FEs
+z = rand(N)
+fes_hard_interact = [FixedEffect(pid), FixedEffect(pid; interaction = z), FixedEffect(fid), FixedEffect(fid; interaction = z)]
+
+# Three-way absorption (worker, firm, year)
+yid = rand(1:36, N)
+fes_hard_3way = [FixedEffect(pid), FixedEffect(fid), FixedEffect(yid)]
+
+# Large cardinality: N=10M, worker-firm (500k × 50k), banded construction with
+# the firm window tuned so each column takes ~20 LSMR iterations; the worker
+# coefficient tile exceeds _SORT_TILE_BYTES, so this spec runs on the sorted
+# observation layout on every backend
+N_large = 10_000_000
+M_large = 500_000
+O_large = 50_000
+pid_large = rand(1:M_large, N_large)
+fid_large = [rand(max(1, div(x, 10)-5_000):min(O_large, div(x, 10)+5_000)) for x in pid_large]
+fes_large = [FixedEffect(pid_large), FixedEffect(fid_large)]
+cols_large = [rand(N_large) for _ in 1:6]
+
+##############################################################################
+# CPU
+##############################################################################
+
+println("\n", "="^60)
+println("Backend: cpu (Float64)")
+println("="^60)
+
+feM = AbstractFixedEffectSolver{Float64}(fes_simple, uweights(length(cols_simple[1])), Val{:cpu})
+println("Simple (N=10M, 100k×100), 6 columns, first run:")
+@time solve_residuals!([copy(c) for c in cols_simple], feM; progress_bar = false)   # ~1.1 s
+println("Simple (N=10M, 100k×100), 6 columns, second run:")
+@time solve_residuals!([copy(c) for c in cols_simple], feM; progress_bar = false)   # ~0.8 s
+
+feM = AbstractFixedEffectSolver{Float64}(fes_hard, uweights(N), Val{:cpu})
+println("Hard (N=800k, 40k×5k), 6 columns, first run:")
+@time solve_residuals!([copy(c) for c in cols_hard], feM; progress_bar = false)   # ~13.6 s
+println("Hard (N=800k, 40k×5k), 6 columns, second run:")
+@time solve_residuals!([copy(c) for c in cols_hard], feM; progress_bar = false)   # ~13.8 s
+
+feM = AbstractFixedEffectSolver{Float64}(fes_hard_interact, uweights(N), Val{:cpu})
+println("Hard (N=800k, interacted), 6 columns, first run:")
+@time solve_residuals!([copy(c) for c in cols_hard], feM; progress_bar = false)   # ~22 s
+println("Hard (N=800k, interacted), 6 columns, second run:")
+@time solve_residuals!([copy(c) for c in cols_hard], feM; progress_bar = false)   # ~26 s
+
+feM = AbstractFixedEffectSolver{Float64}(fes_hard_3way, uweights(N), Val{:cpu})
+println("Hard 3-way (N=800k, 40k×5k×36), 6 columns, first run:")
+@time solve_residuals!([copy(c) for c in cols_hard], feM; progress_bar = false)   # ~20.6 s
+println("Hard 3-way (N=800k, 40k×5k×36), 6 columns, second run:")
+@time solve_residuals!([copy(c) for c in cols_hard], feM; progress_bar = false)   # ~19.9 s
+
+feM = AbstractFixedEffectSolver{Float64}(fes_large, uweights(N_large), Val{:cpu})
+println("Large (N=10M, 500k×50k, sorted layout), 6 columns, first run:")
+@time solve_residuals!([copy(c) for c in cols_large], feM; progress_bar = false)   # ~4.4 s
+println("Large (N=10M, 500k×50k, sorted layout), 6 columns, second run:")
+@time solve_residuals!([copy(c) for c in cols_large], feM; progress_bar = false)   # ~4.1 s
+
+##############################################################################
+# Metal
+##############################################################################
+
+if isdefined(Main, :Metal) && Metal.functional()
+    println("\n", "="^60)
+    println("Backend: Metal (Float32)")
+    println("="^60)
+
+    feM = AbstractFixedEffectSolver{Float32}(fes_simple, uweights(length(cols_simple[1])), Val{:Metal})
+    println("Simple (N=10M, 100k×100), 6 columns, first run:")
+    @time solve_residuals!([copy(c) for c in cols_simple], feM; progress_bar = false)   # ~9.9 s
+    println("Simple (N=10M, 100k×100), 6 columns, second run:")
+    @time solve_residuals!([copy(c) for c in cols_simple], feM; progress_bar = false)   # ~0.38 s
+
+    feM = AbstractFixedEffectSolver{Float32}(fes_hard, uweights(N), Val{:Metal})
+    println("Hard (N=800k, 40k×5k), 6 columns, first run:")
+    @time solve_residuals!([copy(c) for c in cols_hard], feM; progress_bar = false)   # ~3.5 s
+    println("Hard (N=800k, 40k×5k), 6 columns, second run:")
+    @time solve_residuals!([copy(c) for c in cols_hard], feM; progress_bar = false)   # ~0.75 s
+
+    feM = AbstractFixedEffectSolver{Float32}(fes_hard_interact, uweights(N), Val{:Metal})
+    println("Hard (N=800k, interacted), 6 columns, first run:")
+    @time solve_residuals!([copy(c) for c in cols_hard], feM; progress_bar = false)   # ~3.2 s
+    println("Hard (N=800k, interacted), 6 columns, second run:")
+    @time solve_residuals!([copy(c) for c in cols_hard], feM; progress_bar = false)   # ~1.3 s
+
+    feM = AbstractFixedEffectSolver{Float32}(fes_hard_3way, uweights(N), Val{:Metal})
+    println("Hard 3-way (N=800k, 40k×5k×36), 6 columns, first run:")
+    @time solve_residuals!([copy(c) for c in cols_hard], feM; progress_bar = false)   # ~1.6 s
+    println("Hard 3-way (N=800k, 40k×5k×36), 6 columns, second run:")
+    @time solve_residuals!([copy(c) for c in cols_hard], feM; progress_bar = false)   # ~0.95 s
+
+    feM = AbstractFixedEffectSolver{Float32}(fes_large, uweights(N_large), Val{:Metal})
+    println("Large (N=10M, 500k×50k, sorted layout), 6 columns, first run:")
+    @time solve_residuals!([copy(c) for c in cols_large], feM; progress_bar = false)   # ~2.8 s
+    println("Large (N=10M, 500k×50k, sorted layout), 6 columns, second run:")
+    @time solve_residuals!([copy(c) for c in cols_large], feM; progress_bar = false)   # ~1.0 s
+end
+
+##############################################################################
+# CUDA
+##############################################################################
+
+if isdefined(Main, :CUDA) && CUDA.functional()
+    println("\n", "="^60)
+    println("Backend: CUDA (Float32)")
+    println("="^60)
+
+    feM = AbstractFixedEffectSolver{Float32}(fes_simple, uweights(length(cols_simple[1])), Val{:CUDA})
+    println("Simple (N=10M, 100k×100), 6 columns, first run:")
+    @time solve_residuals!([copy(c) for c in cols_simple], feM; progress_bar = false)
+    println("Simple (N=10M, 100k×100), 6 columns, second run:")
+    @time solve_residuals!([copy(c) for c in cols_simple], feM; progress_bar = false)
+
+    feM = AbstractFixedEffectSolver{Float32}(fes_hard, uweights(N), Val{:CUDA})
+    println("Hard (N=800k, 40k×5k), 6 columns, first run:")
+    @time solve_residuals!([copy(c) for c in cols_hard], feM; progress_bar = false)
+    println("Hard (N=800k, 40k×5k), 6 columns, second run:")
+    @time solve_residuals!([copy(c) for c in cols_hard], feM; progress_bar = false)
+
+    feM = AbstractFixedEffectSolver{Float32}(fes_hard_interact, uweights(N), Val{:CUDA})
+    println("Hard (N=800k, interacted), 6 columns, first run:")
+    @time solve_residuals!([copy(c) for c in cols_hard], feM; progress_bar = false)
+    println("Hard (N=800k, interacted), 6 columns, second run:")
+    @time solve_residuals!([copy(c) for c in cols_hard], feM; progress_bar = false)
+
+    feM = AbstractFixedEffectSolver{Float32}(fes_hard_3way, uweights(N), Val{:CUDA})
+    println("Hard 3-way (N=800k, 40k×5k×36), 6 columns, first run:")
+    @time solve_residuals!([copy(c) for c in cols_hard], feM; progress_bar = false)
+    println("Hard 3-way (N=800k, 40k×5k×36), 6 columns, second run:")
+    @time solve_residuals!([copy(c) for c in cols_hard], feM; progress_bar = false)
+
+    feM = AbstractFixedEffectSolver{Float32}(fes_large, uweights(N_large), Val{:CUDA})
+    println("Large (N=10M, 500k×50k, sorted layout), 6 columns, first run:")
+    @time solve_residuals!([copy(c) for c in cols_large], feM; progress_bar = false)
+    println("Large (N=10M, 500k×50k, sorted layout), 6 columns, second run:")
+    @time solve_residuals!([copy(c) for c in cols_large], feM; progress_bar = false)
+end
diff --git a/benchmarks/gather_strategies.jl b/benchmarks/gather_strategies.jl
deleted file mode 100644
index 6c3f83d..0000000
--- a/benchmarks/gather_strategies.jl
+++ /dev/null
@@ -1,169 +0,0 @@
-using Random, BenchmarkTools, Base.Threads
-println("Julia ", VERSION, " — ", nthreads(), " threads")
-Random.seed!(1234)
-
-##############################################################################
-# Current serial gather (baseline)
-##############################################################################
-function gather_serial!(fecoef, refs, α, y, cache)
-    @fastmath @inbounds @simd for i in eachindex(y)
-        fecoef[refs[i]] += α * y[i] * cache[i]
-    end
-end
-
-##############################################################################
-# Approach 1: CSC-style transposed gather
-# Precompute a CSC structure: for each group k, store obs indices.
-# Each group k can be processed independently → trivially parallel, zero conflicts.
-##############################################################################
-struct CSCIndex
-    offsets::Vector{Int}
-    indices::Vector{Int}
-end
-
-function build_csc(refs::AbstractVector{<:Integer}, n::Int)
-    N = length(refs)
-    counts = zeros(Int, n)
-    @inbounds for i in 1:N
-        counts[refs[i]] += 1
-    end
-    offsets = Vector{Int}(undef, n + 1)
-    offsets[1] = 1
-    @inbounds for k in 1:n
-        offsets[k+1] = offsets[k] + counts[k]
-    end
-    indices = Vector{Int}(undef, N)
-    fill!(counts, 0)
-    @inbounds for i in 1:N
-        k = refs[i]
-        counts[k] += 1
-        indices[offsets[k] + counts[k] - 1] = i
-    end
-    return CSCIndex(offsets, indices)
-end
-
-function gather_csc_parallel!(fecoef::AbstractVector{T}, csc::CSCIndex, α, y, cache) where T
-    offsets, indices = csc.offsets, csc.indices
-    n = length(fecoef)
-    Threads.@threads for k in 1:n
-        s = zero(T)
-        @fastmath @inbounds for j in offsets[k]:(offsets[k+1]-1)
-            i = indices[j]
-            s += y[i] * cache[i]
-        end
-        @inbounds fecoef[k] += α * s
-    end
-end
-
-function gather_csc_serial!(fecoef::AbstractVector{T}, csc::CSCIndex, α, y, cache) where T
-    offsets, indices = csc.offsets, csc.indices
-    n = length(fecoef)
-    for k in 1:n
-        s = zero(T)
-        @fastmath @inbounds for j in offsets[k]:(offsets[k+1]-1)
-            i = indices[j]
-            s += y[i] * cache[i]
-        end
-        @inbounds fecoef[k] += α * s
-    end
-end
-
-##############################################################################
-# Approach 2: Per-thread accumulators with manual chunking (@spawn)
-##############################################################################
-struct PerThreadBuffers{T}
-    buffers::Vector{Vector{T}}
-end
-PerThreadBuffers{T}(n::Int, nt::Int) where T = PerThreadBuffers([zeros(T, n) for _ in 1:nt])
-
-function gather_perthread!(fecoef::AbstractVector{T}, refs, α, y, cache, ptb::PerThreadBuffers{T}) where T
-    nt = length(ptb.buffers)
-    N = length(y)
-    for buf in ptb.buffers
-        fill!(buf, zero(T))
-    end
-    chunk = cld(N, nt)
-    @sync for t in 1:nt
-        Threads.@spawn begin
-            buf = ptb.buffers[t]
-            lo = (t-1)*chunk + 1
-            hi = min(t*chunk, N)
-            @fastmath @inbounds for i in lo:hi
-                buf[refs[i]] += y[i] * cache[i]
-            end
-        end
-    end
-    @inbounds for buf in ptb.buffers
-        @simd for k in eachindex(fecoef)
-            fecoef[k] += α * buf[k]
-        end
-    end
-end
-
-##############################################################################
-# Benchmarks
-##############################################################################
-function run_bench(label, N, n_groups)
-    println("\n", "="^60)
-    println("$label: N=$N, n_groups=$n_groups (avg group size=$(N÷n_groups))")
-    println("="^60)
-
-    refs = rand(1:n_groups, N)
-    y = rand(N)
-    cache = rand(N)
-    α = 1.0
-    nt = nthreads()
-
-    csc = build_csc(refs, n_groups)
-    ptb = PerThreadBuffers{Float64}(n_groups, nt)
-
-    # Verify correctness
-    out_ref = zeros(n_groups)
-    gather_serial!(out_ref, refs, α, y, cache)
-
-    for (name, fn!) in [
-        ("CSC parallel", (out) -> gather_csc_parallel!(out, csc, α, y, cache)),
-        ("CSC serial", (out) -> gather_csc_serial!(out, csc, α, y, cache)),
-        ("Per-thread chunked", (out) -> gather_perthread!(out, refs, α, y, cache, ptb)),
-    ]
-        out_test = zeros(n_groups)
-        fn!(out_test)
-        if !isapprox(out_test, out_ref, rtol=1e-10)
-            println("  WARNING $name: INCORRECT (max diff = $(maximum(abs.(out_test .- out_ref))))")
-        end
-    end
-
-    print("  Serial (baseline):  ")
-    b0 = @benchmark gather_serial!(out, $refs, $α, $y, $cache) setup=(out=zeros($n_groups)) evals=1 samples=30
-    show(stdout, MIME("text/plain"), b0); println()
-
-    print("  CSC serial:         ")
-    b1 = @benchmark gather_csc_serial!(out, $csc, $α, $y, $cache) setup=(out=zeros($n_groups)) evals=1 samples=30
-    show(stdout, MIME("text/plain"), b1); println()
-
-    print("  CSC parallel:       ")
-    b2 = @benchmark gather_csc_parallel!(out, $csc, $α, $y, $cache) setup=(out=zeros($n_groups)) evals=1 samples=30
-    show(stdout, MIME("text/plain"), b2); println()
-
-    print("  Per-thread chunked: ")
-    b3 = @benchmark gather_perthread!(out, $refs, $α, $y, $cache, $ptb) setup=(out=zeros($n_groups)) evals=1 samples=30
-    show(stdout, MIME("text/plain"), b3); println()
-
-    t0 = median(b0).time
-    println("\n  Speedups vs serial baseline:")
-    println("    CSC serial:         $(round(t0/median(b1).time, digits=2))x")
-    println("    CSC parallel:       $(round(t0/median(b2).time, digits=2))x")
-    println("    Per-thread chunked: $(round(t0/median(b3).time, digits=2))x")
-end
-
-# Scenario 1: Few large groups (like year FE)
-run_bench("Few large groups", 10_000_000, 100)
-
-# Scenario 2: Many medium groups
-run_bench("Many medium groups", 10_000_000, 100_000)
-
-# Scenario 3: Many small groups (worker FE)
-run_bench("Many small groups (worker FE)", 800_000, 400_000)
-
-# Scenario 4: Moderate groups (firm FE)
-run_bench("Moderate groups (firm FE)", 800_000, 50_000)
diff --git a/benchmarks/solve_backends.jl b/benchmarks/solve_backends.jl
deleted file mode 100644
index bfbcb9a..0000000
--- a/benchmarks/solve_backends.jl
+++ /dev/null
@@ -1,82 +0,0 @@
-using FixedEffects, Random
-try using CUDA catch end
-try using Metal catch end
-Random.seed!(1234)
-
-##############################################################################
-# Setup
-##############################################################################
-
-# Simple problem: N=10M, two FEs (100k × 100 groups)
-N = 10_000_000
-K = 100
-id1 = rand(1:div(N, K), N)
-id2 = rand(1:K, N)
-fes_simple = [FixedEffect(id1), FixedEffect(id2)]
-x_simple = rand(N)
-
-# Hard problem: N=800k, worker-firm (400k × 50k)
-N = 800_000
-M = 400_000
-O = 50_000
-Random.seed!(1234)
-pid = rand(1:M, N)
-fid = [rand(max(1, div(x, 8)-10):min(O, div(x, 8)+10)) for x in pid]
-x_hard = rand(N)
-fes_hard = [FixedEffect(pid), FixedEffect(fid)]
-y = rand(N)
-fes_interact = [FixedEffect(pid), FixedEffect(pid; interaction = y)]
-
-##############################################################################
-# CPU
-##############################################################################
-
-println("Simple (N=10M, 100k×100), first run:") # ~3 s
-@time solve_residuals!(deepcopy(x_simple), fes_simple)
-println("Simple (N=10M, 100k×100), second run:") # ~0.5 s
-@time solve_residuals!(deepcopy(x_simple), fes_simple)
-
-println("Hard (N=800k, 400k×50k), Float32, first run:") # ~3 s
-@time solve_residuals!(deepcopy(x_hard), fes_hard; double_precision = false)
-println("Hard (N=800k, 400k×50k), Float32, second run:") # ~2.5 s
-@time solve_residuals!(deepcopy(x_hard), fes_hard; double_precision = false)
-
-println("Hard (N=800k, 400k×50k), maxiter=300, first run:") # ~2.5 s
-@time solve_residuals!(deepcopy(x_hard), fes_hard; maxiter = 300)
-println("Hard (N=800k, 400k×50k), maxiter=300, second run:") # ~2.5 s
-@time solve_residuals!(deepcopy(x_hard), fes_hard; maxiter = 300)
-
-println("Hard (N=800k, interacted), first run:") # ~3.5 s
-@time solve_residuals!(deepcopy(x_hard), fes_interact; maxiter = 300)
-println("Hard (N=800k, interacted), second run:") # ~3 s
-@time solve_residuals!(deepcopy(x_hard), fes_interact; maxiter = 300)
-
-##############################################################################
-# CUDA
-##############################################################################
-if isdefined(Main, :CUDA) && CUDA.functional()
-    println("Simple (N=10M, 100k×100), CUDA, first run:")
-    @time solve_residuals!(deepcopy(x_simple), fes_simple; method = :CUDA)
-    println("Simple (N=10M, 100k×100), CUDA, second run:")
-    @time solve_residuals!(deepcopy(x_simple), fes_simple; method = :CUDA)
-
-    println("Hard (N=800k, 400k×50k), CUDA, first run:")
-    @time solve_residuals!(deepcopy(x_hard), fes_hard; method = :CUDA)
-    println("Hard (N=800k, 400k×50k), CUDA, second run:")
-    @time solve_residuals!(deepcopy(x_hard), fes_hard; method = :CUDA)
-end
-
-##############################################################################
-# Metal
-##############################################################################
-if isdefined(Main, :Metal) && Metal.functional()
-    println("Simple (N=10M, 100k×100), Metal, first run:") # ~18 s
-    @time solve_residuals!(Float32.(deepcopy(x_simple)), fes_simple; method = :Metal, double_precision = false)
-    println("Simple (N=10M, 100k×100), Metal, second run:") # ~1.5 s
-    @time solve_residuals!(Float32.(deepcopy(x_simple)), fes_simple; method = :Metal, double_precision = false)
-
-    println("Hard (N=800k, 400k×50k), Metal, first run:") # ~3.3 s
-    @time solve_residuals!(Float32.(deepcopy(x_hard)), fes_hard; method = :Metal, double_precision = false, maxiter = 300)
-    println("Hard (N=800k, 400k×50k), Metal, second run:") # ~1.6 s
-    @time solve_residuals!(Float32.(deepcopy(x_hard)), fes_hard; method = :Metal, double_precision = false, maxiter = 300)
-end
diff --git a/benchmarks/solve_cpu.jl b/benchmarks/solve_cpu.jl
deleted file mode 100755
index 2f0f001..0000000
--- a/benchmarks/solve_cpu.jl
+++ /dev/null
@@ -1,72 +0,0 @@
-using FixedEffects, Random, Statistics
-Random.seed!(1234)
-
-# Simple problem
-N = 10_000_000
-K = 100
-id1 = rand(1:div(N, K), N)
-id2 = rand(1:K, N)
-fes = [FixedEffect(id1), FixedEffect(id2)]
-x = rand(N)
-
-@time solve_residuals!(deepcopy(x), fes)
-
-# More complicated problem (worker-firm)
-N = 800_000
-M = 400_000
-O = 50_000
-Random.seed!(1234)
-pid = rand(1:M, N)
-fid = [rand(max(1, div(x, 8)-10):min(O, div(x, 8)+10)) for x in pid]
-x = rand(N)
-fes = [FixedEffect(pid), FixedEffect(fid)]
-
-@time solve_residuals!(deepcopy(x), fes; double_precision = false)
-@time solve_residuals!(deepcopy(x), fes; maxiter = 300)
-
-# Interacted fixed effects
-y = rand(N)
-fes = [FixedEffect(pid), FixedEffect(pid; interaction = y)]
-@time solve_residuals!(deepcopy(x), fes)
-
-
-##############################################################################
-# Benchmark: simple problem (N=10M, K=100)
-##############################################################################
-println("="^60)
-println("Simple problem: N=10M, two FEs (100k and 100 groups)")
-println("="^60)
-
-N = 10_000_000
-K = 100
-id1 = rand(1:div(N, K), N)
-id2 = rand(1:K, N)
-fes = [FixedEffect(id1), FixedEffect(id2)]
-x = rand(N)
-
-# warmup
-solve_residuals!(deepcopy(x), fes)
-# benchmark
-b1 = @benchmark solve_residuals!(xc, $fes) setup=(xc=deepcopy($x)) evals=1 samples=10
-show(stdout, MIME("text/plain"), b1); println()
-println("  Median: $(round(median(b1).time/1e6, digits=1)) ms")
-
-##############################################################################
-# Benchmark: hard problem (N=800k, worker-firm)
-##############################################################################
-println("\n", "="^60)
-println("Hard problem: N=800k, worker (400k) x firm (50k)")
-println("="^60)
-
-N2 = 800_000
-M = 400_000
-O = 50_000
-pid = rand(1:M, N2)
-fid = [rand(max(1, div(x, 8)-10):min(O, div(x, 8)+10)) for x in pid]
-x2 = rand(N2)
-fes2 = [FixedEffect(pid), FixedEffect(fid)]
-
-solve_residuals!(deepcopy(x2), fes2)
-b2 = @benchmark solve_residuals!(xc, $fes2) setup=(xc=deepcopy($x2)) evals=1 samples=5
-show(stdout, MIME("text/plain"), b2); println()
-println("  Median: $(round(median(b2).time/1e6, digits=0)) ms")
diff --git a/benchmarks/worker_firm_spectrum.jl b/benchmarks/worker_firm_spectrum.jl
deleted file mode 100644
index 7ebb955..0000000
--- a/benchmarks/worker_firm_spectrum.jl
+++ /dev/null
@@ -1,85 +0,0 @@
-# Measure the leading canonical correlations sigma_k between the two FE subspaces for the
-# benchmark worker-firm data. After Jacobi scaling A'A = [[I,C],[C',I]] with sigma_k = svd(C);
-# slow LSMR convergence is driven by sigma_k -> 1. This quantifies how many slow modes exist
-# (=> how many deflation vectors k would be needed and the achievable iteration reduction).
-using Random, LinearAlgebra
-Random.seed!(1234)
-
-# --- benchmark hard scenario ---
-N = 800_000; M = 400_000; O = 50_000
-refs1 = rand(1:M, N)                                             # worker
-refs2 = [rand(max(1, div(x, 8)-10):min(O, div(x, 8)+10)) for x in refs1]  # firm
-n1 = maximum(refs1); n2 = maximum(refs2)
-
-# Jacobi scales (unweighted, no interaction): scale[g] = 1/sqrt(group count)
-function scales(refs, n)
-    c = zeros(Int, n); @inbounds for r in refs; c[r] += 1; end
-    s = zeros(n); @inbounds for g in 1:n; s[g] = c[g] > 0 ? 1 / sqrt(c[g]) : 0.0; end
-    s
-end
-s1 = scales(refs1, n1); s2 = scales(refs2, n2)
-
-# C v = gather_FE1(scatter_FE2(v)), maps firm-space (n2) -> worker-space (n1)
-function Cmul!(out, v, refs1, refs2, s1, s2)
-    fill!(out, 0.0)
-    @inbounds for i in eachindex(refs1)
-        out[refs1[i]] += s2[refs2[i]] * v[refs2[i]]
-    end
-    @inbounds for g in eachindex(out); out[g] *= s1[g]; end
-    out
-end
-
-# C' u, maps worker-space (n1) -> firm-space (n2)
-function Ctmul!(out, u, refs1, refs2, s1, s2)
-    fill!(out, 0.0)
-    @inbounds for i in eachindex(refs1)
-        out[refs2[i]] += s1[refs1[i]] * u[refs1[i]]
-    end
-    @inbounds for h in eachindex(out); out[h] *= s2[h]; end
-    out
-end
-
-# Subspace (block) iteration on C'C (acts on firm space n2=50k) for the top-k sigma^2.
-function top_sigmas(k, iters)
-    Y = randn(n2, k)
-    tmp1 = zeros(n1); ritz = zeros(k)
-    for it in 1:iters
-        # Z = (C'C) Y
-        Z = similar(Y)
-        for j in 1:k
-            Cmul!(tmp1, view(Y, :, j), refs1, refs2, s1, s2)
-            Ctmul!(view(Z, :, j), tmp1, refs1, refs2, s1, s2)
-        end
-        F = qr(Z); Q = Matrix(F.Q)
-        # Rayleigh-Ritz on Q
-        AQ = similar(Q)
-        for j in 1:k
-            Cmul!(tmp1, view(Q, :, j), refs1, refs2, s1, s2)
-            Ctmul!(view(AQ, :, j), tmp1, refs1, refs2, s1, s2)
-        end
-        H = Symmetric(Q' * AQ)
-        E = eigen(H); ev = E.values; perm = sortperm(ev, rev = true)
-        ritz = ev[perm]
-        Y = Q * E.vectors[:, perm]
-    end
-    sqrt.(clamp.(ritz, 0, Inf))
-end
-
-println("worker-firm: N=$N, n1(worker)=$n1, n2(firm)=$n2")
-k = 30
-sig = top_sigmas(k, 60)
-println("\nTop $k canonical correlations sigma_k (descending):")
-for (j, s) in enumerate(sig)
-    gap = 1 - s
-    println("  k=$(lpad(j, 2))  sigma=", round(s, digits = 6), "   1-sigma=", round(gap, sigdigits = 3))
-end
-
-# condition-number proxy and crude iteration estimate (iters ~ sqrt((1+s)/(1-s)) for sigma_max)
-s2nd = sig[2]
-println("\nsigma_1=", round(sig[1], digits = 6), " (constant mode, deflated by rank-deficiency)")
-println("sigma_2=", round(s2nd, digits = 6), "  => kappa~", round((1 + s2nd) / (1 - s2nd), digits = 1),
-        "  sqrt(kappa)~", round(sqrt((1 + s2nd) / (1 - s2nd)), digits = 1))
-nbig = count(>(0.99), sig)
-println("count(sigma>0.99) in top $k: ", nbig)
-nbig999 = count(>(0.999), sig)
-println("count(sigma>0.999) in top $k: ", nbig999)
diff --git a/ext/CUDAExt.jl b/ext/CUDAExt.jl
index 3702688..ded5a67 100644
--- a/ext/CUDAExt.jl
+++ b/ext/CUDAExt.jl
@@ -188,6 +188,82 @@ function gather_block_kernel!(coef_block, refs, qrows, y, α, k)
 	return nothing
 end
 
+##############################################################################
+##
+## 1d) Fused bidiagonalization step (single-pass LSMR iterations)
+##
+## One kernel computes u ← Σ_blocks A_b v_b + c u and accumulates ‖u‖² on the
+## fly: block-reduced in the working precision, then one Float64 atomic per
+## thread block so the cross-block sum does not drift (see the _norm2 note in
+## src/utils/lsmr.jl). This replaces one scatter launch per block plus a
+## separate norm reduction; the gathers then run through their per-block
+## strategies on the raw u. Blocks are passed as tuples of device arrays so
+## the kernel unrolls across them.
+##
+##############################################################################
+
+function FixedEffects.bidiag_forward!(u::CuVector{T}, g::FixedEffectCoefficients,
+		fem::FixedEffectLinearMapCUDA{T}, v::FixedEffectCoefficients, c::Number) where {T}
+	blocks = Tuple(fem.plan.blocks)
+	refss = map(block -> block.refs, blocks)
+	qrowss = Tuple(fem.plan.qrows)
+	vs = Tuple(v.x)
+	normacc = CUDA.zeros(Float64, 1)
+	nthreads = 256
+	nblocks = cld(length(u), nthreads)
+	@cuda threads=nthreads blocks=nblocks bidiag_scatter_kernel!(u, refss, qrowss, vs, T(c),
+		normacc, Val(nthreads))
+	fill!(g, zero(T))
+	for (coef_block, block, qrows, gather) in zip(g.x, fem.plan.blocks, fem.plan.qrows, fem.gathers)
+		_gather_block!(coef_block, block.refs, qrows, u, one(T), gather)
+	end
+	return T(sqrt(Array(normacc)[1]))
+end
+
+function bidiag_scatter_kernel!(u, refss, qrowss, vs, c, normacc, ::Val{NT}) where {NT}
+	T = eltype(u)
+	tid = Int(threadIdx().x)
+	shared = CUDA.CuStaticSharedArray(T, NT)
+	index = (Int(blockIdx().x) - 1) * NT + tid
+	stride = NT * Int(gridDim().x)
+	acc = zero(T)
+	i = index
+	@inbounds while i <= length(u)
+		ui = c * u[i] + _device_fit(refss, qrowss, vs, i)
+		u[i] = ui
+		acc += ui * ui
+		i += stride
+	end
+	@inbounds shared[tid] = acc
+	CUDA.sync_threads()
+	offset = NT ÷ 2
+	while offset > 0
+		if tid <= offset
+			@inbounds shared[tid] += shared[tid + offset]
+		end
+		CUDA.sync_threads()
+		offset ÷= 2
+	end
+	if tid == 1
+		CUDA.@atomic normacc[1] += Float64(shared[1])
+	end
+	return nothing
+end
+
+# Recursion over the block tuples, as in the CPU kernels.
+@inline _device_fit(::Tuple{}, ::Tuple{}, ::Tuple{}, i) = false
+@inline function _device_fit(refss::Tuple, qrowss::Tuple, vs::Tuple, i)
+	refs = first(refss)
+	qrows = first(qrowss)
+	vcoef = first(vs)
+	@inbounds gr = refs[i]
+	fit = zero(eltype(qrows))
+	for col in 1:size(qrows, 1)
+		@inbounds fit += vcoef[col, gr] * qrows[col, i]
+	end
+	return fit + _device_fit(Base.tail(refss), Base.tail(qrowss), Base.tail(vs), i)
+end
+
 ##############################################################################
 ##
 ## 2. FixedEffectSolverCUDA
@@ -203,6 +279,7 @@ mutable struct FixedEffectSolverCUDA{T} <: FixedEffects.AbstractFixedEffectSolve
 	v::FixedEffectCoefficients
 	h::FixedEffectCoefficients
 	hbar::FixedEffectCoefficients
+	g::FixedEffectCoefficients
 	tmp::Vector{T} # used to convert AbstractVector to Vector{T}
 end
 
@@ -214,8 +291,9 @@ function FixedEffects.AbstractFixedEffectSolver{T}(fes::Vector{<:FixedEffect}, w
 	v = FixedEffectCoefficients([CUDA.zeros(T, block_width(block), block.n) for block in m.plan.blocks])
 	h = FixedEffectCoefficients([CUDA.zeros(T, block_width(block), block.n) for block in m.plan.blocks])
 	hbar = FixedEffectCoefficients([CUDA.zeros(T, block_width(block), block.n) for block in m.plan.blocks])
+	g = FixedEffectCoefficients([CUDA.zeros(T, block_width(block), block.n) for block in m.plan.blocks])
 	tmp = zeros(T, length(weights))
-	return FixedEffectSolverCUDA{T}(m, _cu(T, weights), b, r, x, v, h, hbar, tmp)
+	return FixedEffectSolverCUDA{T}(m, _cu(T, weights), b, r, x, v, h, hbar, g, tmp)
 end
 
 function FixedEffects.update_weights!(feM::FixedEffectSolverCUDA{T}, weights::AbstractWeights) where {T}
diff --git a/ext/MetalExt.jl b/ext/MetalExt.jl
index 5e31d65..25119a6 100644
--- a/ext/MetalExt.jl
+++ b/ext/MetalExt.jl
@@ -12,6 +12,10 @@ Metal.allowscalar(false)
 ## the CPU; refs and qrows are moved to the device and consumed by fused
 ## block kernels.
 ##
+## Kernels are launched asynchronously — Metal.jl queues them in order on its
+## command queue — and the host waits only where it reads a result (the norm
+## in bidiag_forward!, the host copies in copy_internal!).
+##
 ##############################################################################
 
 ##############################################################################
@@ -30,6 +34,24 @@ function _metal_threadgroup_width()
 	return prevpow(2, width)
 end
 
+# Largest power-of-two threadgroup width the compiled pipeline allows: kernels
+# with heavy register use can cap below the device width.
+_pipeline_width(kernel) = prevpow(2, Int(kernel.pipeline.maxTotalThreadsPerThreadgroup))
+
+# Compile a kernel whose threadgroup width is also a compile-time argument
+# (the shared-memory tree size, passed as Val): if the compiled pipeline
+# allows fewer threads than requested, shrink the width and compile once more
+# (a smaller threadgroup only relaxes the pipeline limits, so this converges).
+function _sized_kernel(build, nthreads::Int)
+	kernel = build(nthreads)
+	cap = _pipeline_width(kernel)
+	if cap < nthreads
+		nthreads = cap
+		kernel = build(nthreads)
+	end
+	return kernel, nthreads
+end
+
 # Per-block plan for the adjoint gather (A'u), chosen once at construction:
 # bucketize (one threadgroup per group) for low cardinality, else atomic adds.
 struct AtomicGather end
@@ -90,9 +112,10 @@ end
 
 function _scatter_block!(y::MtlVector, refs::MtlVector, qrows::MtlMatrix,
 		coef_block::MtlMatrix, α::Number, β::Number)
-	nthreads = _metal_threadgroup_width()
-	nblocks = cld(length(y), nthreads)
-	Metal.@sync @metal threads=nthreads groups=nblocks scatter_block_kernel!(y, refs, qrows, coef_block, α, β, size(coef_block, 1))
+	kernel = @metal launch=false scatter_block_kernel!(y, refs, qrows, coef_block, α, β, size(coef_block, 1))
+	nthreads = _pipeline_width(kernel)
+	kernel(y, refs, qrows, coef_block, α, β, size(coef_block, 1);
+		threads = nthreads, groups = cld(length(y), nthreads))
 	return y
 end
 
@@ -127,17 +150,21 @@ end
 
 function _gather_block!(coef_block::MtlMatrix, refs::MtlVector, qrows::MtlMatrix,
 		y::MtlVector, α::Number, gather::BucketGather)
-	n = size(coef_block, 2)
-	nthreads = _metal_threadgroup_width()
-	Metal.@sync @metal threads=nthreads groups=n gather_block_kernel_bin!(coef_block, α, y, qrows, gather.perm, gather.offsets, Val(nthreads), size(coef_block, 1))
+	k = size(coef_block, 1)
+	kernel, nthreads = _sized_kernel(_metal_threadgroup_width()) do nt
+		@metal launch=false gather_block_kernel_bin!(coef_block, α, y, qrows, gather.perm, gather.offsets, Val(nt), k)
+	end
+	kernel(coef_block, α, y, qrows, gather.perm, gather.offsets, Val(nthreads), k;
+		threads = nthreads, groups = size(coef_block, 2))
 	return coef_block
 end
 
 function _gather_block!(coef_block::MtlMatrix, refs::MtlVector, qrows::MtlMatrix,
 		y::MtlVector, α::Number, ::AtomicGather)
-	nthreads = _metal_threadgroup_width()
-	nblocks = cld(length(y), nthreads)
-	Metal.@sync @metal threads=nthreads groups=nblocks gather_block_kernel!(coef_block, refs, α, y, qrows, size(coef_block, 1))
+	kernel = @metal launch=false gather_block_kernel!(coef_block, refs, α, y, qrows, size(coef_block, 1))
+	nthreads = _pipeline_width(kernel)
+	kernel(coef_block, refs, α, y, qrows, size(coef_block, 1);
+		threads = nthreads, groups = cld(length(y), nthreads))
 	return coef_block
 end
 
@@ -195,6 +222,87 @@ function gather_block_kernel!(coef_block, refs, α, y, qrows, k)
 	return nothing
 end
 
+##############################################################################
+##
+## 1d) Fused bidiagonalization step (single-pass LSMR iterations)
+##
+## One kernel computes u ← Σ_blocks A_b v_b + c u and accumulates ‖u‖² on the
+## fly: one partial per threadgroup, summed pairwise on the device — Metal has
+## no Float64 atomics, and the pairwise sum avoids the drift of chaining
+## Float32 atomic adds (see the _norm2 note in src/utils/lsmr.jl). The gathers
+## then run through their per-block strategies on the raw u. Blocks are passed
+## as tuples of device arrays so the kernel unrolls across them.
+##
+##############################################################################
+
+function FixedEffects.bidiag_forward!(u::MtlVector{T}, g::FixedEffectCoefficients,
+		fem::FixedEffectLinearMapMetal{T}, v::FixedEffectCoefficients, c::Number) where {T}
+	blocks = Tuple(fem.plan.blocks)
+	refss = map(block -> block.refs, blocks)
+	qrowss = Tuple(fem.plan.qrows)
+	vs = Tuple(v.x)
+	# fixed 256-thread groups, safely below any pipeline cap for this kernel
+	nthreads = 256
+	nblocks = cld(length(u), nthreads)
+	partials = MtlVector{T}(undef, nblocks)
+	@metal threads=nthreads groups=nblocks bidiag_scatter_kernel!(u, refss, qrowss, vs, T(c),
+		partials, Val(nthreads))
+	fill!(g, zero(T))
+	for (coef_block, block, qrows, gather) in zip(g.x, fem.plan.blocks, fem.plan.qrows, fem.gathers)
+		_gather_block!(coef_block, block.refs, qrows, u, one(T), gather)
+	end
+	# every kernel of this iteration is queued in order; the scalar read of the
+	# device sum below is the one host wait per iteration
+	return sqrt(sum(partials))
+end
+
+function bidiag_scatter_kernel!(u, refss, qrowss, vs, c, partials, ::Val{NT}) where {NT}
+	T = eltype(u)
+	gid = Int(threadgroup_position_in_grid().x)
+	tid = Int(thread_position_in_threadgroup().x)
+	nt = Int(threads_per_threadgroup().x)
+	shared = Metal.MtlThreadGroupArray(T, NT)
+	acc = zero(T)
+	i = thread_position_in_grid_1d()
+	if i <= length(u)
+		@inbounds begin
+			ui = c * u[i] + _device_fit(refss, qrowss, vs, i)
+			u[i] = ui
+			acc += ui * ui
+		end
+	end
+	@inbounds shared[tid] = acc
+	Metal.threadgroup_barrier(Metal.MemoryFlagThreadGroup)
+
+	offset = nt ÷ 2
+	while offset > 0
+		if tid <= offset
+			@inbounds shared[tid] += shared[tid + offset]
+		end
+		Metal.threadgroup_barrier(Metal.MemoryFlagThreadGroup)
+		offset ÷= 2
+	end
+
+	if tid == 1
+		@inbounds partials[gid] = shared[1]
+	end
+	return nothing
+end
+
+# Recursion over the block tuples, as in the CPU kernels.
+@inline _device_fit(::Tuple{}, ::Tuple{}, ::Tuple{}, i) = false
+@inline function _device_fit(refss::Tuple, qrowss::Tuple, vs::Tuple, i)
+	refs = first(refss)
+	qrows = first(qrowss)
+	vcoef = first(vs)
+	@inbounds gr = refs[i]
+	fit = zero(eltype(qrows))
+	for col in 1:size(qrows, 1)
+		@inbounds fit += vcoef[col, gr] * qrows[col, i]
+	end
+	return fit + _device_fit(Base.tail(refss), Base.tail(qrowss), Base.tail(vs), i)
+end
+
 ##############################################################################
 ##
 ## 2. FixedEffectSolverMetal
@@ -210,6 +318,7 @@ mutable struct FixedEffectSolverMetal{T} <: FixedEffects.AbstractFixedEffectSolv
 	v::FixedEffectCoefficients
 	h::FixedEffectCoefficients
 	hbar::FixedEffectCoefficients
+	g::FixedEffectCoefficients
 	tmp::Vector{T}
 end
 
@@ -223,8 +332,9 @@ function FixedEffects.AbstractFixedEffectSolver{T}(fes::Vector{<:FixedEffect}, w
 	v = FixedEffectCoefficients([Metal.zeros(T, block_width(block), block.n) for block in m.plan.blocks])
 	h = FixedEffectCoefficients([Metal.zeros(T, block_width(block), block.n) for block in m.plan.blocks])
 	hbar = FixedEffectCoefficients([Metal.zeros(T, block_width(block), block.n) for block in m.plan.blocks])
+	g = FixedEffectCoefficients([Metal.zeros(T, block_width(block), block.n) for block in m.plan.blocks])
 	tmp = zeros(T, length(weights))
-	return FixedEffectSolverMetal{T}(m, _mtl(T, weights), b, r, x, v, h, hbar, tmp)
+	return FixedEffectSolverMetal{T}(m, _mtl(T, weights), b, r, x, v, h, hbar, g, tmp)
 end
 
 
diff --git a/src/AbsorptionPlan.jl b/src/AbsorptionPlan.jl
index 87d986a..fde9099 100644
--- a/src/AbsorptionPlan.jl
+++ b/src/AbsorptionPlan.jl
@@ -58,6 +58,62 @@ function AbsorptionPlan(::Type{T}, plan::AbsorptionPlan, weights::AbstractVector
 	return AbsorptionPlan(plan.blocks, transforms, ranks, qrows)
 end
 
+## 1b') Sorted observation layout
+
+# Sorting observations by one block's groups turns that block's random
+# coefficient accesses into streams in every scatter/gather, and lets
+# its gather run race-free without per-thread buffers (row chunks can end on
+# group boundaries). Sorting pays once the block's coefficient tile outgrows
+# the caches; below this threshold it is skipped. Set to 0 to force sorting
+# (used by the tests).
+const _SORT_TILE_BYTES = Ref(1024 * 1024)
+
+# Build the plan on a sorted observation order: pick the block with the
+# largest coefficient tile (the largest random-access working set) and, when
+# its refs are unsorted and that tile is large enough to matter, permute every
+# block's refs and interactions, and the weights, by its counting-sort order.
+# Returns (plan, perm, sorted_block):
+# perm is nothing when the observations were not permuted (sorted_block may
+# still name an already-sorted block; 0 when none is sorted). Callers must
+# permute solver-side observation data (weights, right-hand sides) with perm.
+function sorted_absorption_plan(::Type{T}, fes::Vector{<:FixedEffect}, weights::AbstractVector;
+		ranktol::Union{Nothing,Real} = nothing) where {T}
+	blocks = _build_absorbed_blocks(fes)
+	perm = nothing
+	sorted_block = 0
+	# with a single block the solvers use one direct projection: no iterations
+	# to speed up, so sorting would only add the permutation passes
+	if length(blocks) > 1
+		j = argmax([block_width(block) * block.n for block in blocks])
+		if block_width(blocks[j]) * blocks[j].n * sizeof(T) > _SORT_TILE_BYTES[]
+			if issorted(blocks[j].refs)
+				sorted_block = j
+			else
+				_, _, perm = _group_permutation(blocks[j].refs, blocks[j].n)
+				blocks = [_permute_block(block, perm) for block in blocks]
+				weights = _permute_weights(weights, perm)
+				sorted_block = j
+			end
+		end
+	end
+	transforms, ranks, qrows = _build_transforms(T, blocks, weights, ranktol)
+	return AbsorptionPlan(blocks, transforms, ranks, qrows), perm, sorted_block
+end
+
+function _permute_block(block::AbsorbedBlock, perm::Vector{Int})
+	interactions = map(block.interactions) do interaction
+		if interaction isa UnitWeights
+			interaction
+		else
+			interaction[perm]
+		end
+	end
+	return AbsorbedBlock(block.refs[perm], interactions, block.n, block.input_terms)
+end
+
+_permute_weights(weights::UnitWeights, ::Vector{Int}) = weights
+_permute_weights(weights::AbstractVector, perm::Vector{Int}) = weights[perm]
+
 function _build_absorbed_blocks(fes::Vector{<:FixedEffect})
 	blocks = AbsorbedBlock[]
 	for (j, fe) in enumerate(fes)
diff --git a/src/AbstractFixedEffectSolver.jl b/src/AbstractFixedEffectSolver.jl
index b278050..e94909f 100644
--- a/src/AbstractFixedEffectSolver.jl
+++ b/src/AbstractFixedEffectSolver.jl
@@ -63,7 +63,7 @@ function solve_residuals!(r::AbstractVector{<:Real}, feM::AbstractFixedEffectSol
 	if length(feM.m.plan.blocks) == 1
 		mul!(feM.x, feM.m', feM.b, 1, 0)
 	else
-		_, ch = lsmr!(feM.x, feM.m, feM.b, feM.v, feM.h, feM.hbar; atol = tol, btol = tol, maxiter = maxiter)
+		_, ch = lsmr!(feM.x, feM.m, feM.b, feM.v, feM.h, feM.hbar, feM.g; atol = tol, btol = tol, maxiter = maxiter)
 		iter, converged = ch.mvps, ch.isconverged
 	end
 	converged || @warn "solve_residuals! did not converge within maxiter LSMR iterations; returned values may be inaccurate." iterations=iter maxiter tol
@@ -163,7 +163,7 @@ function solve_coefficients!(r::AbstractVector, feM::AbstractFixedEffectSolver{T
 		feM.b .*= sqrt.(feM.weights)
 	end
 	fill!(feM.x, zero(T))
-	_, ch = lsmr!(feM.x, feM.m, feM.b, feM.v, feM.h, feM.hbar; atol = tol, btol = tol, maxiter = maxiter)
+	_, ch = lsmr!(feM.x, feM.m, feM.b, feM.v, feM.h, feM.hbar, feM.g; atol = tol, btol = tol, maxiter = maxiter)
 	ch.isconverged || @warn "solve_coefficients! did not converge within maxiter LSMR iterations; returned values may be inaccurate." iterations=ch.mvps maxiter tol
 	recover_coefficients(feM, eltype(r)), ch.mvps, ch.isconverged
 end
diff --git a/src/CPU.jl b/src/CPU.jl
index ffda428..4df7e01 100644
--- a/src/CPU.jl
+++ b/src/CPU.jl
@@ -29,6 +29,12 @@ struct ThreadedGather{M<:AbstractMatrix}
 	buffers::Vector{M}              # one k × n accumulator per thread
 	ranges::Vector{UnitRange{Int}}  # contiguous row chunks
 end
+# For a block whose refs are sorted (see sorted_absorption_plan): row chunks
+# end on group boundaries, so threads accumulate directly into the shared
+# coefficient block — race-free without buffers, at any cardinality.
+struct SortedGather
+	ranges::Vector{UnitRange{Int}}  # row chunks aligned to group boundaries
+end
 
 mutable struct FixedEffectLinearMapCPU{T,F<:Vector{<:FixedEffect},P<:AbsorptionPlan,G<:AbstractVector} <: AbstractFixedEffectLinearMap{T}
 	fes::F
@@ -38,17 +44,23 @@ end
 
 # The struct definitions must precede this constructor (a constructor method
 # signature is evaluated at definition time, unlike ordinary function calls).
-function FixedEffectLinearMapCPU{T}(fes::Vector{<:FixedEffect},
-	weights::AbstractVector = uweights(T, length(fes[1].refs))) where {T}
-	plan = AbsorptionPlan(T, fes, weights)
+function FixedEffectLinearMapCPU{T}(fes::Vector{<:FixedEffect}, plan::AbsorptionPlan,
+		sorted_block::Integer) where {T}
 	N = length(fes[1].refs)
 	nt = nthreads()
 	ranges = _row_chunks(N, nt)
-	G = Union{SerialGather, ThreadedGather{Matrix{T}}}
-	gathers = G[_gather_strategy(T, block, N, nt, ranges) for block in plan.blocks]
+	G = Union{SerialGather, SortedGather, ThreadedGather{Matrix{T}}}
+	gathers = G[_gather_strategy(T, plan.blocks[j], N, nt, ranges, j == sorted_block)
+		for j in eachindex(plan.blocks)]
 	return FixedEffectLinearMapCPU{T,typeof(fes),typeof(plan),typeof(gathers)}(fes, plan, gathers)
 end
 
+function FixedEffectLinearMapCPU{T}(fes::Vector{<:FixedEffect},
+	weights::AbstractVector = uweights(T, length(fes[1].refs))) where {T}
+	plan = AbsorptionPlan(T, fes, weights)
+	return FixedEffectLinearMapCPU{T}(fes, plan, 0)
+end
+
 # Toggle to force the serial baseline (e.g. for benchmarking); threading is on by default.
 const _USE_THREADED_GATHER = Ref(true)
 # Threading the gather pays off only when the nt per-thread accumulators of size
@@ -57,10 +69,19 @@ const _USE_THREADED_GATHER = Ref(true)
 const _GATHER_BUFFER_BUDGET = 8 * 1024 * 1024   # bytes
 const _GATHER_MIN_ROWS = 100_000                # below this, threading overhead isn't worth it
 
-# Per block, thread the gather only if the accumulators fit in cache and N is large.
+# Per block: the designated sorted block gathers race-free over group-aligned
+# chunks; otherwise thread only if the accumulators fit in cache and N is large.
 function _gather_strategy(::Type{T}, block::AbsorbedBlock, N::Int, nt::Int,
-		ranges::Vector{UnitRange{Int}}) where {T}
+		ranges::Vector{UnitRange{Int}}, sorted::Bool) where {T}
 	k = block_width(block)
+	if sorted && nt > 1 && N >= _GATHER_MIN_ROWS
+		ranges_sorted = _group_aligned_chunks(block.refs, nt)
+		# a dominant group can swallow most rows into one chunk and serialize
+		# the pass; fall back to the buffered strategies in that case
+		if maximum(length, ranges_sorted) <= 2 * cld(N, nt)
+			return SortedGather(ranges_sorted)
+		end
+	end
 	if _USE_THREADED_GATHER[] && nt > 1 && N >= _GATHER_MIN_ROWS &&
 			nt * k * block.n * sizeof(T) <= _GATHER_BUFFER_BUDGET
 		return ThreadedGather([zeros(T, k, block.n) for _ in 1:nt], ranges)
@@ -69,6 +90,28 @@ function _gather_strategy(::Type{T}, block::AbsorbedBlock, N::Int, nt::Int,
 	end
 end
 
+# nchunks row chunks over sorted refs, each ending on a group boundary.
+function _group_aligned_chunks(refs::AbstractVector, nchunks::Int)
+	N = length(refs)
+	ranges = Vector{UnitRange{Int}}(undef, nchunks)
+	lo = 1
+	for t in 1:nchunks
+		if t == nchunks
+			hi = N
+		else
+			target = div(N * t, nchunks)
+			if target < lo
+				hi = lo - 1          # empty chunk: the previous group swallowed its share
+			else
+				hi = searchsortedlast(refs, refs[target])
+			end
+		end
+		ranges[t] = lo:hi
+		lo = hi + 1
+	end
+	return ranges
+end
+
 
 ## 1b) FixedEffectLinearMapCPU mul!
 
@@ -153,6 +196,16 @@ gather_block!(coef_block::AbstractMatrix, block::AbsorbedBlock, qrows::AbstractM
 	y::AbstractVector, α::Number, ::SerialGather) =
 	_gather_block!(coef_block, block, qrows, y, α, eachindex(y))
 
+# Sorted refs: each thread's chunk ends on a group boundary, so all threads
+# write disjoint columns of coef_block — no buffers, no merge, any cardinality.
+function gather_block!(coef_block::AbstractMatrix, block::AbsorbedBlock, qrows::AbstractMatrix,
+	y::AbstractVector, α::Number, g::SortedGather)
+	@threads for t in eachindex(g.ranges)
+		_gather_block!(coef_block, block, qrows, y, α, g.ranges[t])
+	end
+	return coef_block
+end
+
 # Threaded: each thread reduces its row chunk into a private (cache-resident) buffer,
 # then the buffers are summed into coef_block.
 function gather_block!(coef_block::AbstractMatrix, block::AbsorbedBlock, qrows::AbstractMatrix,
@@ -192,6 +245,145 @@ function _gather_block!(coef_block::AbstractMatrix, block::AbsorbedBlock,
 	return coef_block
 end
 
+##############################################################################
+##
+## 1d) Fused bidiagonalization step (single-pass LSMR iterations)
+##
+## u ← A v + c u, β = ‖u‖, g ← A'u — in as few passes over u as the gather
+## strategies allow. The mul!-based lsmr! iteration streams u about seven
+## times (one scatter per block, each reading and writing all of u, a norm
+## pass, a scale pass, one gather read per block); here the scatters of every
+## block, the norm, and — when every block has cache-resident per-thread
+## buffers — the gathers all run inside one loop over observations, so u and
+## refs/qrows are streamed once per iteration.
+##
+## Blocks are passed as tuples: the recursive helpers unroll across blocks and
+## the per-column loops stay compile-time (block_width, as in the kernels
+## above). Building the tuples is dynamic, but it happens once per call and
+## the kernels behind the function barrier specialize.
+##
+##############################################################################
+
+function bidiag_forward!(u::Vector{T}, g::FixedEffectCoefficients, fem::FixedEffectLinearMapCPU{T},
+		v::FixedEffectCoefficients, c::Number) where {T}
+	fill!(g, zero(T))
+	blocks = Tuple(fem.plan.blocks)
+	qrowss = Tuple(fem.plan.qrows)
+	vs = Tuple(v.x)
+	gs = Tuple(g.x)
+	N = length(u)
+	nt = nthreads()
+	if nt > 1 && N >= _GATHER_MIN_ROWS && all(gather -> gather isa Union{SortedGather, ThreadedGather}, fem.gathers)
+		# fully fused: each thread owns a row chunk and accumulates every
+		# block's gather — directly into g for the sorted block (its chunks
+		# end on group boundaries, so writes are disjoint), into private
+		# buffers for the others — plus a norm partial. When a sorted block
+		# is present its group-aligned chunks are used for the whole pass.
+		gathers = Tuple(fem.gathers)
+		sorted = findfirst(gather -> gather isa SortedGather, fem.gathers)
+		if sorted === nothing
+			ranges = _row_chunks(N, nt)
+		else
+			ranges = fem.gathers[sorted].ranges
+		end
+		partials = Vector{Float64}(undef, length(ranges))
+		@threads for t in eachindex(ranges)
+			targets = map(gathers, gs) do gather, coef_block
+				if gather isa SortedGather
+					coef_block
+				else
+					buf = gather.buffers[t]
+					fill!(buf, zero(T))
+					buf
+				end
+			end
+			partials[t] = _bidiag_chunk!(u, targets, blocks, qrowss, vs, T(c), ranges[t])
+		end
+		@inbounds for (coef_block, gather) in zip(g.x, fem.gathers)
+			gather isa ThreadedGather || continue
+			for buf in gather.buffers
+				@simd for idx in eachindex(coef_block)
+					coef_block[idx] += buf[idx]
+				end
+			end
+		end
+		s = sum(partials)
+	elseif nt > 1 && N >= _GATHER_MIN_ROWS
+		# fused scatters + norm in one threaded pass; the gathers of each
+		# block then run through their existing strategies on the raw u
+		ranges = _row_chunks(N, nt)
+		partials = Vector{Float64}(undef, length(ranges))
+		@threads for t in eachindex(ranges)
+			partials[t] = _bidiag_scatter_chunk!(u, blocks, qrowss, vs, T(c), ranges[t])
+		end
+		s = sum(partials)
+		for (coef_block, block, qrows, gather) in zip(g.x, fem.plan.blocks, fem.plan.qrows, fem.gathers)
+			gather_block!(coef_block, block, qrows, u, one(T), gather)
+		end
+	else
+		# serial: everything in one pass, gathers accumulated directly into g
+		s = _bidiag_chunk!(u, gs, blocks, qrowss, vs, T(c), eachindex(u))
+	end
+	return T(sqrt(s))
+end
+
+# u[i] ← c * u[i] + Σ_blocks fit_i; accumulates each block's gather into
+# `bufs` (g's own blocks when serial, one thread's private buffers when
+# threaded) and returns the Float64 sum of squares of the updated u.
+function _bidiag_chunk!(u::Vector{T}, bufs::Tuple, blocks::Tuple, qrowss::Tuple, vs::Tuple,
+		c::T, range) where {T}
+	s = 0.0
+	@inbounds for i in range
+		ui = c * u[i] + _scatter_fit(blocks, qrowss, vs, i)
+		u[i] = ui
+		s += abs2(Float64(ui))
+		_gather_accum!(bufs, blocks, qrowss, i, ui)
+	end
+	return s
+end
+
+function _bidiag_scatter_chunk!(u::Vector{T}, blocks::Tuple, qrowss::Tuple, vs::Tuple,
+		c::T, range) where {T}
+	s = 0.0
+	@inbounds for i in range
+		ui = c * u[i] + _scatter_fit(blocks, qrowss, vs, i)
+		u[i] = ui
+		s += abs2(Float64(ui))
+	end
+	return s
+end
+
+# Recursion over the block tuples: each level specializes on its block type,
+# so the inner loops over columns unroll (block_width is compile-time).
+@inline _scatter_fit(::Tuple{}, ::Tuple{}, ::Tuple{}, i) = false   # additive zero of any float type
+@inline function _scatter_fit(blocks::Tuple, qrowss::Tuple, vs::Tuple, i)
+	block = first(blocks)
+	qrows = first(qrowss)
+	vcoef = first(vs)
+	@inbounds begin
+		gr = block.refs[i]
+		fit = zero(eltype(qrows))
+		for col in 1:block_width(block)
+			fit += vcoef[col, gr] * qrows[col, i]
+		end
+	end
+	return fit + _scatter_fit(Base.tail(blocks), Base.tail(qrowss), Base.tail(vs), i)
+end
+
+@inline _gather_accum!(::Tuple{}, ::Tuple{}, ::Tuple{}, i, ui) = nothing
+@inline function _gather_accum!(bufs::Tuple, blocks::Tuple, qrowss::Tuple, i, ui)
+	buf = first(bufs)
+	block = first(blocks)
+	qrows = first(qrowss)
+	@inbounds begin
+		gr = block.refs[i]
+		for col in 1:block_width(block)
+			buf[col, gr] += ui * qrows[col, i]
+		end
+	end
+	return _gather_accum!(Base.tail(bufs), Base.tail(blocks), Base.tail(qrowss), i, ui)
+end
+
 ##############################################################################
 ##
 ## 2. FixedEffectSolverCPU
@@ -201,16 +393,22 @@ end
 mutable struct FixedEffectSolverCPU{T,M<:FixedEffectLinearMapCPU{T},C<:FixedEffectCoefficients{Matrix{T}}} <: AbstractFixedEffectSolver{T}
 	m::M
 	weights::AbstractVector
+	perm::Union{Nothing, Vector{Int}}   # observation order of the internal storage (see sorted_absorption_plan)
 	b::Vector{T}
 	r::Vector{T}
 	x::C
 	v::C
 	h::C
 	hbar::C
+	g::C
 end
 
 function AbstractFixedEffectSolver{T}(fes::Vector{<:FixedEffect}, weights::AbstractWeights, ::Type{Val{:cpu}}) where {T}
-	m = FixedEffectLinearMapCPU{T}(fes, weights)
+	plan, perm, sorted_block = sorted_absorption_plan(T, fes, weights)
+	m = FixedEffectLinearMapCPU{T}(fes, plan, sorted_block)
+	if perm !== nothing
+		weights = _permute_weights(weights, perm)
+	end
 	b = zeros(T, length(weights))
 	r = zeros(T, length(weights))
 	blocks = m.plan.blocks
@@ -218,20 +416,55 @@ function AbstractFixedEffectSolver{T}(fes::Vector{<:FixedEffect}, weights::Abstr
 	v = FixedEffectCoefficients([zeros(T, block_width(block), block.n) for block in blocks])
 	h = FixedEffectCoefficients([zeros(T, block_width(block), block.n) for block in blocks])
 	hbar = FixedEffectCoefficients([zeros(T, block_width(block), block.n) for block in blocks])
-	return FixedEffectSolverCPU(m, weights, b, r, x, v, h, hbar)
+	g = FixedEffectCoefficients([zeros(T, block_width(block), block.n) for block in blocks])
+	return FixedEffectSolverCPU(m, weights, perm, b, r, x, v, h, hbar, g)
 end
 
 function update_weights!(feM::FixedEffectSolverCPU{T}, weights::AbstractWeights) where {T}
+	if feM.perm !== nothing
+		weights = _permute_weights(weights, feM.perm)
+	end
 	feM.m.plan = AbsorptionPlan(T, feM.m.plan, weights)
 	feM.weights = weights
 	return feM
 end
 
+# Internal storage lives in the (possibly) sorted observation order; these
+# translate from and to the caller's order. The permutation loops sit behind
+# function barriers: getfield with a runtime Symbol is abstractly typed, and
+# an element-wise loop on it would dispatch dynamically on every element.
 function copy_internal!(feM::FixedEffectSolverCPU, field::Symbol, r::AbstractVector)
-	copyto!(getfield(feM, field), r)
+	dest = getfield(feM, field)
+	if feM.perm === nothing
+		copyto!(dest, r)
+	else
+		_gather_perm!(dest, r, feM.perm)
+	end
+	return dest
 end
 
 function copy_internal!(r::AbstractVector, feM::FixedEffectSolverCPU, field::Symbol)
-	copyto!(r, getfield(feM, field))
+	src = getfield(feM, field)
+	if feM.perm === nothing
+		copyto!(r, src)
+	else
+		_scatter_perm!(r, src, feM.perm)
+	end
+	return r
+end
+
+# dest[i] = r[perm[i]]
+function _gather_perm!(dest::AbstractVector, r::AbstractVector, perm::Vector{Int})
+	@inbounds for i in eachindex(dest)
+		dest[i] = r[perm[i]]
+	end
+	return dest
 end
 
+# r[perm[i]] = src[i]
+function _scatter_perm!(r::AbstractVector, src::AbstractVector, perm::Vector{Int})
+	@inbounds for i in eachindex(src)
+		r[perm[i]] = src[i]
+	end
+	return r
+end
diff --git a/src/utils/lsmr.jl b/src/utils/lsmr.jl
index a8543f7..78673b8 100644
--- a/src/utils/lsmr.jl
+++ b/src/utils/lsmr.jl
@@ -57,11 +57,18 @@ _norm2(x) = norm(x)
 
 
 ## Arguments:
-## x is initial guess for x0. Transformed in place to the solution.
-## b is also transformed in place
-## v, h, hbar are storage arrays of length size(A, 2)
-function lsmr!(x, A, b, v, h, hbar; 
-    atol::Number = 1e-6, btol::Number = 1e-6, conlim::Number = 1e8, 
+## x must be zero on entry (the solvers fill! it before calling); transformed
+## in place to the solution.
+## b is also transformed in place: it holds the UNNORMALIZED Golub-Kahan
+## vector u, i.e. u_true = u / β. Folding the 1/β into the scalar recurrences
+## lets bidiag_forward! run a whole iteration in a single pass over u: the
+## norm of the new u and the gather A'u are accumulated in the same loop that
+## builds it, which is valid because the gather is linear: A'(u/β) = (A'u)/β.
+## In exact arithmetic the iterates are identical to textbook LSMR.
+## v, h, hbar, g are storage arrays of length size(A, 2); g receives A'u
+## each iteration
+function lsmr!(x, A, b, v, h, hbar, g;
+    atol::Number = 1e-6, btol::Number = 1e-6, conlim::Number = 1e8,
     maxiter::Integer = max(size(A,1), size(A,2)), λ::Number = 0)
 
     # Sanity-checking
@@ -72,18 +79,22 @@ function lsmr!(x, A, b, v, h, hbar;
     length(v) == n || error("v has length $(length(v)) but should have length $n")
     length(h) == n || error("h has length $(length(h)) but should have length $n")
     length(hbar) == n || error("hbar has length $(length(hbar)) but should have length $n")
+    length(g) == n || error("g has length $(length(g)) but should have length $n")
     length(b) == m || error("b has length $(length(b)) but should have length $m")
-    
+
     T = Base.promote_op(/, eltype(b), eltype(A))
     Tr = real(T)
     conlim > 0 ? ctol = convert(Tr, inv(conlim)) : ctol = zero(Tr)
-    # form the first vectors u and v (satisfy  β*u = b,  α*v = A'u)    
-    u = mul!(b, A, x, -1, 1)
-    β = _norm2(u)    
-    β > 0 && rmul!(u, inv(β))    
-    mul!(v, A', u, 1, 0)
-    α = _norm2(v)
-    α > 0 && rmul!(v, inv(α))
+    # form the first vectors u and v (satisfy  β*u = b,  α*v = A'u)
+    u = b
+    β = _norm2(u)
+    α = zero(Tr)
+    if β > 0
+        mul!(v, A', u, 1, 0)
+        rmul!(v, inv(β))
+        α = _norm2(v)
+        α > 0 && rmul!(v, inv(α))
+    end
     # Initialize variables for 1st iteration.
     ζbar = α * β
     αbar = α
@@ -116,18 +127,22 @@ function lsmr!(x, A, b, v, h, hbar;
     normAr = α * β
     iter = 0
     # Exit if b = 0 or A'b = 0.
-    if normAr != 0 
+    if normAr != 0
         while iter < maxiter
             iter += 1
-            mul!(u, A, v, 1, -α)
-            β = _norm2(u)
+            # u ← A v − α u_true, with the 1/β normalization of the stored u
+            # folded into the scalar; one fused pass also returns ‖u‖ and
+            # fills g = A'u (see bidiag_forward!)
+            cu = β > 0 ? -α / β : zero(Tr)
+            β = bidiag_forward!(u, g, A, v, cu)
             if β > 0
-                rmul!(u, inv(β))
-                mul!(v, A', u, 1, -β)
+                # v ← A'(u/β) − β v = g/β − β v, then normalize
+                rmul!(v, -β)
+                axpy!(inv(β), g, v)
                 α = _norm2(v)
                 α > 0 && rmul!(v, inv(α))
             end
-        
+
             # Construct rotation Qhat_{k,2k+1}.
             αhat = sqrt(abs2(αbar) + abs2(λ))
             chat = αbar / αhat
@@ -240,3 +255,19 @@ function lsmr!(x, A, b, v, h, hbar;
     return x, ch
 end
 
+# One fused Golub-Kahan step: u ← A v + c u, then β = ‖u‖ and g ← A'u (u left
+# unnormalized — lsmr! folds the 1/β into its scalars). This fallback uses the
+# operator's mul!; backends may override it with a true single-pass
+# implementation (see src/CPU.jl).
+function bidiag_forward!(u, g, A, v, c)
+    mul!(u, A, v, 1, c)
+    β = _norm2(u)
+    if β > 0
+        # accumulate into zeroed g rather than mul!(g, A', u, 1, 0): the
+        # adjoint mul! scales g by β first, and g is uninitialized on entry
+        fill!(g, zero(eltype(g)))
+        mul!(g, A', u, 1, 1)
+    end
+    return β
+end
+
diff --git a/test/solve.jl b/test/solve.jl
index 1ce7314..ac70586 100644
--- a/test/solve.jl
+++ b/test/solve.jl
@@ -15,6 +15,12 @@ r_ols =  [-0.2015993617092453,  0.2015993617092464, -0.2015993617092463,  0.2015
 (r, iter, conv) = solve_residuals!(deepcopy(x), fes)
 @test r ≈ r_ols
 
+# a zero right-hand side converges immediately
+(rz, iterz, convz) = solve_residuals!(zeros(10), fes)
+@test convz
+@test iterz == 0
+@test all(iszero, rz)
+
 @testset "maxiter semantics" begin
 	(r0, iter0, conv0) = @test_logs (:warn, r"solve_residuals!") solve_residuals!(deepcopy(x), fes; maxiter = 0)
 	@test iter0 == 0
@@ -125,6 +131,72 @@ end
 	end
 end
 
+@testset "sorted observation layout" begin
+	n_sort = 150_000
+	id_low = mod1.(1:n_sort, 13)
+	id_high = mod1.(37 .* (1:n_sort) .+ 11, 257)
+	y_sort = sin.((1:n_sort) ./ 5) .+ cos.((1:n_sort) ./ 17)
+	weights_sort = Weights(1 .+ mod.(1:n_sort, 9) ./ 20)
+	fes_sort = [FixedEffect(id_low), FixedEffect(id_high)]
+	n_gpu = 2048
+	p1_gpu = mod1.(1:n_gpu, 32)
+	p2_gpu = mod1.((1:n_gpu) .* 7, 41)
+	x_gpu = sin.((1:n_gpu) ./ 3) .+ cos.((1:n_gpu) ./ 11)
+	weights_gpu = Weights(1 .+ mod.(1:n_gpu, 5) ./ 10)
+	fes_gpu = [FixedEffect(p1_gpu), FixedEffect(p2_gpu)]
+	atol_gpu = 1e-3
+	rtol_gpu = 1e-3
+	old_sort_tile = FixedEffects._SORT_TILE_BYTES[]
+	try
+		FixedEffects._SORT_TILE_BYTES[] = typemax(Int)
+		r_unsorted = solve_residuals!(copy(y_sort), fes_sort, weights_sort)[1]
+		coefs_unsorted = solve_coefficients!(copy(y_sort), fes_sort, weights_sort)[1]
+
+		FixedEffects._SORT_TILE_BYTES[] = 0
+		plan_sort, perm_sort, sorted_block = FixedEffects.sorted_absorption_plan(Float64, fes_sort, weights_sort)
+		@test perm_sort !== nothing
+		@test sorted_block == 2
+		@test issorted(plan_sort.blocks[sorted_block].refs)
+
+		r_sorted = solve_residuals!(copy(y_sort), fes_sort, weights_sort)[1]
+		@test r_sorted ≈ r_unsorted atol = 1e-8
+
+		coefs_sorted = solve_coefficients!(copy(y_sort), fes_sort, weights_sort)[1]
+		@test _residual_from_coefs(y_sort, fes_sort, coefs_sorted) ≈
+			_residual_from_coefs(y_sort, fes_sort, coefs_unsorted) atol = 1e-8
+
+		feM_sort = FixedEffects.AbstractFixedEffectSolver{Float64}(fes_sort, weights_sort, Val{:cpu})
+		@test feM_sort.perm !== nothing
+		new_weights_sort = Weights(1 .+ mod.(3 .* (1:n_sort), 11) ./ 15)
+		FixedEffects.update_weights!(feM_sort, new_weights_sort)
+		r_reused = solve_residuals!(copy(y_sort), feM_sort)[1]
+		r_fresh = solve_residuals!(copy(y_sort), fes_sort, new_weights_sort)[1]
+		@test r_reused ≈ r_fresh atol = 1e-8
+
+		for method in filter(!=(:cpu), method_s)
+			cpu_bucket_r = solve_residuals!(copy(x_gpu), fes_gpu, weights_gpu; double_precision = false)[1]
+			gpu_bucket_r = solve_residuals!(copy(x_gpu), fes_gpu, weights_gpu; method = method, double_precision = false)[1]
+			@test gpu_bucket_r ≈ cpu_bucket_r atol = atol_gpu rtol = rtol_gpu
+
+			id_owner_gpu = mod1.(37 .* (1:n_gpu) .+ 11, 1024)
+			fes_owner_gpu = [FixedEffect(p1_gpu), FixedEffect(id_owner_gpu)]
+			cpu_owner_r = solve_residuals!(copy(x_gpu), fes_owner_gpu, weights_gpu; double_precision = false)[1]
+			gpu_owner_r = solve_residuals!(copy(x_gpu), fes_owner_gpu, weights_gpu; method = method, double_precision = false)[1]
+			@test gpu_owner_r ≈ cpu_owner_r atol = atol_gpu rtol = rtol_gpu
+
+			# update_weights! rebuilds the device plan from the stored host plan
+			feM_gpu = FixedEffects.AbstractFixedEffectSolver{Float32}(fes_gpu, weights_gpu, Val{method})
+			new_weights_gpu = Weights(1 .+ mod.(3 .* (1:n_gpu), 11) ./ 15)
+			FixedEffects.update_weights!(feM_gpu, new_weights_gpu)
+			gpu_reused_r = solve_residuals!(copy(x_gpu), feM_gpu; tol = 1e-6)[1]
+			gpu_fresh_r = solve_residuals!(copy(x_gpu), fes_gpu, new_weights_gpu; method = method, double_precision = false)[1]
+			@test gpu_reused_r ≈ gpu_fresh_r atol = atol_gpu rtol = rtol_gpu
+		end
+	finally
+		FixedEffects._SORT_TILE_BYTES[] = old_sort_tile
+	end
+end
+
 
 fe = FixedEffect([1, 2])
 @test_throws "FixedEffects must have the same length as y" ỹ = solve_residuals!(ones(100), [fe])