perf(qdp): remove mid-stream D2H norm validation from batch GPU-pointer encoding paths

[ryankert01]

Update (2026-06-26): Two corrections after re-checking main (full detail in the comment below).

1. The single-sample "prior work" cited below is not on main. The launch_amplitude_encode_f32_device_norm kernel exists only on an unmerged perf branch (commit 4ebe417b3); PR feat(qdp): hoist encode_from_gpu_ptr_f32 onto QuantumEncoder trait #1310 is the trait-dispatch refactor, not the device-norm kernel. The single-sample path on main still does the blocking D2H round-trip too.
2. The batch path needs neither a policy decision nor a new kernel. The batch encode kernel already reads inv_norms from device memory (amplitude.rs:660), so the fix is to defer validation: launch norm kernel → encode kernel on the same stream → single sync at the end → D2H-copy norms and validate after. This removes the mid-stream stall while keeping norm validation a hard guarantee — no strict flag, no new CUDA kernel, and no remove-vs-optional-vs-keep policy call needed.

Summary

The batch amplitude encoding paths (encode_batch_from_gpu_ptr_f32, encode_batch_from_gpu_ptr) perform a blocking CPU round-trip for norm validation between the norm kernel and the encode kernel. This serializes two GPU kernel launches unnecessarily and is the primary reason Mahout's encode-only throughput falls behind PyTorch's reference implementation.

Current pipeline (batch path)

launch_l2_norm_batch_f32   (GPU, queued)
↓
cudaStreamSynchronize      ← STALL #1
dtoh_sync_copy(norms)      ← PCIe transfer: N floats CPU-bound
CPU: check norms[i] for zero/NaN
↓
launch_amplitude_encode_batch_f32  (GPU, queued)
↓
cudaStreamSynchronize      ← STALL #2

Two syncs + one D2H copy per batch, even when all norms are valid (which is almost always).

Benchmark evidence

Running benchmark_pytorch_ref.py --qubits 16 --batches 200 --batch-size 64 (encode-only mode — both frameworks start with data already on GPU):

Framework	Throughput
PyTorch GPU	228,825 vec/s
Mahout	66,615 vec/s
Ratio	0.3x (Mahout 3.4x slower)

PyTorch's amplitude_encode uses torch.linalg.vector_norm + data / norms.clamp(min=1e-10) — everything stays on GPU, no CPU validation.

In end-to-end mode (data gen + H2D + encode), Mahout is 7.8x faster than PyTorch because the full pipeline cost dominates. The gap only appears when isolating kernel work.

Related prior work

⚠️ Correction (see Update at top): this was measured on an unmerged perf branch, not main.

The idea: a launch_amplitude_encode_f32_device_norm kernel reads inv_norm from device memory, letting the norm kernel and encode kernel chain on the same stream with a single sync at the end. On the unmerged branch (commit 4ebe417b3) this yielded 1.22x on the single-sample path (39.7 → 32.6 µs/sample at 16 qubits). It is not present on main.

The batch path has the same structure and benefits from the same idea — and, per the Update at top, it already passes inv_norms as a device pointer, so it gets there by deferring validation with no new kernel.

Related issues / PRs

• PR feat(qdp): hoist encode_from_gpu_ptr_f32 onto QuantumEncoder trait #1310 — encode_from_gpu_ptr_f32 single-sample trait dispatch
• PR [Feature] Make QDP CUDA kernel build targets configurable and future-compatible #1283 — configurable CUDA kernel build targets
• PR [Feature][QDP] F32 support for angle/basis encoders, fidelity metrics and pipeline improvements #1275 — F32 support for angle/basis encoders

[ryankert01]

Follow-up after digging into the current main tree — two corrections and a reframing that I think makes the policy questions above moot.

1. The single-sample "prior work" isn't on main

I referenced a launch_amplitude_encode_f32_device_norm kernel as already-fixed prior art. That's not in main:

• The device_norm kernel only exists on the unmerged perf-exploration branch (commit 4ebe417b3), where I measured the 1.22x.
• PR feat(qdp): hoist encode_from_gpu_ptr_f32 onto QuantumEncoder trait #1310 (which I linked) is the encode_from_gpu_ptr_f32 trait-dispatch refactor, not the device-norm optimization.
• The single-sample path on main still does the blocking round-trip too (amplitude.rs:1167), same as batch.

So this shouldn't be picked up as "we have the kernel, just port it to batch" — that building block isn't merged.

2. The batch path doesn't need a policy decision (or a new kernel)

The three options I laid out (remove / make-optional / keep) all trade the correctness guarantee for speed. But the batch encode kernel already reads inv_norms straight from device memory — launch_amplitude_encode_batch_f32(..., *inv_norms_gpu.device_ptr(), ...) (amplitude.rs:660), and the kernel itself just does v1 * inv_norm with no host dependency, assert, or trap (amplitude_encode_batch_kernel_f32 in amplitude.cu).

The only reason there's a mid-stream stall is that the D2H copy + CPU check sits between the two launches (amplitude.rs:650-656). So there's a fourth option:

Defer the validation — reorder, don't remove.

launch_l2_norm_batch_f32        (GPU, queued)
launch_amplitude_encode_batch   (GPU, queued, reads inv_norms from device)  ← chained, same stream
cudaStreamSynchronize           ← single sync at the end
dtoh_sync_copy(norms) + check   ← validate AFTER, off the critical path

This:

• removes the mid-stream cudaStreamSynchronize entirely (both kernels chain on one stream);
• keeps norm validation as a hard guarantee — same error, same encoder boundary;
• needs no new CUDA kernel for the batch path and no API/strict flag;
• costs only one wasted encode launch on the rare invalid-input error path. A zero/NaN norm just propagates inf/nan into a state vector that's discarded when we return the error — IEEE math, no CUDA trap.

If we want the deferred D2H to be trivially small, we can device-reduce the norms to a single "all-valid" flag and copy back 4 bytes instead of N floats — but that's a refinement, not required.

That answers the three questions without a vote: validation stays a guarantee (Q1), the from_gpu_ptr paths don't need to be reclassified as "trusted" (Q2), and no strict flag is needed (Q3).

3. Status / scope

• Real bottleneck, confirmed in main for batch f32/f64 (amplitude.rs:650, :417) and single-sample (:1167).
• Priority is medium, not urgent — as noted, end-to-end Mahout is already ~7.8x faster; this gap only shows when isolating kernel work.
• Proposed plan: reorder-defer the batch f32 path first (smallest change, no new kernel), then apply the same structure to batch f64 and single-sample for consistency.

I have not re-run benchmark_pytorch_ref.py against this reorder yet — the mechanism removes exactly the stall measured here, but the magnitude needs a GPU benchmark to confirm. Happy to put up the batch-f32 reorder as a draft PR for someone with a GPU to benchmark.