GPU engine dst-output routing: tensors written by kernels return zero via .Data() — freezes PatchTST training on GB10

[dndungu]

Summary

Any *TensorNumeric[float32] passed as the destination to a GPU engine op (MatMul, Add, Sub, Mul, etc.) reads as all-zero via .Data() immediately after the op returns, on NVIDIA GB10 (DGX Spark, aarch64, unified memory). This breaks PatchTST training in github.com/zerfoo/zerfoo — loss freezes at a byte-identical deterministic value (0.268357 on a specific test config) because backward-pass gradients never reach the CPU-visible side for the optimizer to read.

The exact same training loop converges correctly on the CPU engine (86% loss reduction in 10 epochs on 15 samples via TestPatchTST_TrainWindowed_EngineConvergence in timeseries/patchtst_test.go), isolating the fault to ztensor's GPU engine implementation.

Minimal reproducer (zerfoo side)

cd github.com/zerfoo/zerfoo
git checkout main
go test ./timeseries/ -run TestPatchTST_TrainWindowed_EngineConvergence -count=1 -v
# PASSES: loss[0]=5.169868 -> loss[9]=0.702752  (86% reduction)

The same test, pointed at a GPU engine constructed via compute.NewGPUEngine[float32](ops), would fail with static loss. The test currently uses newTestEngine() which returns a CPU engine. Running the same binary on DGX via scripts/bench-spark.sh -samples 1000 -channels 5 -epochs 3 -cleanup shows:
```
epoch 1: loss=0.268357 ok
epoch 2: loss=0.268357 ok
epoch 3: loss=0.268357 ok
convergence: FAILED
```

Empirical diagnostic (seven probes on DGX, branch debug/diag-grad-zero-source-v2)

First-batch gradient/parameter state logged at key points in trainWindowedGPU:

```
P0 post-zero grads.patchEmbW = [0 0 0 0] (expected)
P7 post-encoderBackward dX = [0 0 0 0] BROKEN
P8 post-posEmb CPU accumulation = [0 0 0 0] (consumes dX via .Data())
P1 pre-MatMul fc.dPEW = [0 0 0 0] (already zero before use)
P2 post-Add grads.patchEmbW = [0 0 0 0] (Add of zeros = zero)
P3 pre-gradclip = [0 0 0 0]
P4 adamw-read gradTs[0] = [0 0 0 0]
P5 adamw-post-update paramTs[0] = [0.01112, ...] near-init (AdamW ~= 0 on zero grads)
P6 batch 1 forward entry = same as P5 (AdamW writeback + forward read work fine)
```

Every tensor that was written by a GPU engine op returns zero via `.Data()`. The CPU/GPU disconnect is at the engine-op boundary.

What's NOT the bug (ruled out by multi-day investigation in zerfoo)

1. Wrapper aliasing: `grads.X[i]` and `gradTs[i]` ARE the same `*TensorNumeric` wrapper (sentinel verified in zerfoo PR #369).
2. Stream synchronization: adding `engine.Sync()` (via anonymous interface type-assert on `*GPUEngine[T]`) before `.Data()` reads did NOT change the frozen loss. Not a D2H race.
3. Once-per-training caches: rebuilding `paramTs`/`gradTs` per batch (zerfoo PR #365) was a no-op — the wrappers were already aliased with the live struct fields.
4. Storage-identity sentinel: the check passes because wrappers are genuinely the same instance — the disconnect is deeper, inside the engine op's output buffer routing.
5. Kernel timing: ztensor stream Sync barriers before the CPU-side posEmb accumulation AND before grad-clip had zero effect.

Storage kind flip (context)

Storage-kind probe across all 37 grad tensors on first-batch backward (zerfoo branch `debug/v3-storage-kind-probe`, scratch `.claude/scratch/v3-storage-kind-result.md`):

• Setup: all 37 tensors are `*tensor.CPUStorage[float32]` (allocated via `tensor.New[float32]` which hardcodes `NewCPUStorage`).
• Post-first-backward: 36 of 37 flip to `*tensor.GPUStorage[float32]`. The flip happens via `makeGPUResult` calling `dst[0].SetStorage(gs)` at `ztensor/compute/gpu_kernels.go:121-132`.
• Lone holdout: `posEmb` grad (idx=2) stays CPUStorage because its update is a pure CPU loop at `patchtst_gpu_train.go:1012-1019` (`dPosData[j] += dXData[i]`) that never hits `makeGPUResult`.

After the flip, `GPUStorage.Slice()` at `ztensor/tensor/gpu_storage.go:215-250` performs `make([]T, s.length)` + D2H memcpy on every `.Data()` call. The D2H returns all zeros.

Hypotheses for the real bug (need ztensor expertise to pick)

1. (α) Kernel writes to a different buffer than the one installed by SetStorage. `makeGPUResult` allocates a fresh device buffer, installs it on `dst` via `SetStorage`, but the underlying CUDA kernel launch uses a different pointer (the old storage's pointer, a scratch buffer, or a pre-allocated op output buffer).
2. (β) Kernel writes correctly, but SetStorage happens AFTER the kernel, installing a fresh zero-init buffer that clobbers the kernel's output.
3. (γ) D2H memcpy reads from a buffer that was never written. `GPUStorage.Slice()`'s memcpy source pointer diverges from the kernel's write target. Possibly a double-allocation where `makeGPUResult` returns one buffer but the CUDA call went to another.
4. (δ) In-place aliasing (dst == src0) broken. `engine.Add(ctx, a, b, a)` handled incorrectly — kernel reads src0 after dst has been swapped to a fresh buffer, so reads empty.

Suggested trace points inside ztensor:

• At entry to `makeGPUResult`: log the caller-supplied `dst` pointer and its existing Storage device pointer.
• Just before the CUDA kernel launch: log the device pointer the kernel is writing to.
• Just after the kernel launch but before `SetStorage`: log the same.
• After `SetStorage`: log `dst[0].GetStorage()` device pointer.
• Compare: kernel-write-target vs post-SetStorage device pointer.

Environment

• Host: DGX Spark (GB10, NVIDIA), aarch64, unified memory
• CUDA: current ztensor CUDA runtime binding
• Go: 1.25 (via ztensor module)
• zerfoo branch: main @ current HEAD (all docs in `docs/devlog.md` entry dated 2026-04-08 titled "FINAL")

Cross-references

• zerfoo issue #368 (PR #365 non-functional bisect marker)
• zerfoo issue #364 (paramTs staleness — false concern, resolved)
• zerfoo issue #367 (posEmb idx=2 storage flip anomaly — not an anomaly, it's the clean CPU control case)
• zerfoo PR #369 (Storage-identity sentinel — fixes false-positive panic but doesn't unblock training)
• Full investigation history: https://github.com/zerfoo/zerfoo/blob/main/docs/devlog.md (entries dated 2026-04-08)

What we need from ztensor

1. Trace/fix the `makeGPUResult` output-routing bug per the hypotheses above.
2. Consider making `.Data()` on a GPU tensor implicitly sync the stream before the D2H memcpy (systemic safety net; relevant even if the primary bug is buffer routing).
3. Add a minimal test at the ztensor level: run `engine.Add(a, b, c)` on a GPU engine with known inputs, read `c.Data()`, assert equals expected. If this test FAILS, the bug is reproducible without zerfoo at all.

Priority

Blocks PatchTST training on DGX. CPU training still works on DGX as a workaround (unverified at full 5K x 10ch x 3ep scale but expected to converge given local CPU test passes).

[dndungu]

Update 2026-04-09: Not reproducible at ztensor primitive level on DGX GB10

Added TestGPUEngine_PatchTSTBackward_DstRoundTrip in compute/gpu_dst_roundtrip_test.go (branch fix/issue-79-matmul-accumulate-repro, commit 3e538e6). It ports the exact op sequence from zerfoo/timeseries/patchtst_gpu_train.go:1022-1031:

Transpose(patches -> patchesT)
Zero(dPEW)                              # pre-allocated CPU-wrapper tensor
MatMul(patchesT, dX, dPEW)              # fresh write via GPU kernel
Add(gradW, dPEW, gradW)                 # in-place accumulate into CPU-wrapper
gradW.Data()                            # AdamW-side read

gradW is pre-seeded with nonzero values so a silent-zero readback is detectable. Ran on DGX GB10 via Spark pod ztensor-issue79-repro-1775759440 (manifest at docs/bench/manifests/issue-79-repro.yaml).

Results on DGX GB10

TestGPUEngine_Add_DstRoundTrip_OutOfPlace        PASS
TestGPUEngine_Add_DstRoundTrip_InPlace           PASS
TestGPUEngine_Add_DstRoundTrip_RepeatedInPlace   PASS
TestGPUEngine_Add_DstRoundTrip_NoExplicitSync    PASS
TestGPUEngine_PatchTSTBackward_DstRoundTrip      PASS

What this rules out

None of the four hypotheses (alpha / beta / gamma / delta) from the issue body is triggered by the patch-embedding backward sequence on small shapes. The makeGPUResult / SetStorage / GPUStorage.Slice() path is NOT silently zeroing dst tensors, at least not on its own and not at these shapes.

What it does NOT rule out

1. Shape regime. My test uses totalRows=4, patchLen=3, dModel=2. Real PatchTST runs with thousands of rows and dModel in the hundreds. Arena pressure, alignment, or allocator fragmentation may only manifest at production sizes.
2. Multi-op batch state. The full trainWindowedGPU runs dozens of GPU ops plus encoderBackward per batch. The bug may require that longer chain, not just the patch-embedding slice.
3. CPU-loop interleave. The dPosData[j] += dXData[sOff+j] CPU loop at patchtst_gpu_train.go:1012-1019 touches dX.Data() mid-backward, forcing a D2H sync in the middle of GPU work. Not covered here.
4. zerfoo-side wrapper state. The gradTs rebuild path in zerfoo could alias storage differently than a fresh tensor.New wrapper in a standalone test.

Recommended next steps

• Expand the ztensor test to realistic PatchTST shapes (e.g. totalRows in the thousands, dModel=64+, multi-channel) and rerun.
• OR instrument trainWindowedGPU itself to log device pointers before/after each op around the grad accumulation, rather than trying to lift primitives into ztensor tests.
• This issue should probably be relabeled as a zerfoo integration bug unless the larger-shape ztensor test above reproduces it.

Infra notes captured in docs/devlog.md (2026-04-09 entry). Diagnostic test file and manifest merged to main in e04540e.

[dndungu]

Update 2026-04-09 (2): realistic shapes also PASS on DGX GB10

Expanded the test to match cmd/bench_train/main.go defaults and ran on DGX GB10:

TestGPUEngine_PatchTSTBackward_RealisticShapes  PASS  totalRows=1600 dModel=64 iters=20
TestGPUEngine_PatchTSTBackward_LargerBatch      PASS  totalRows=3200 dModel=64 iters=20

Configuration: PatchLength=8, DModel=64, BatchSize=64, Channels=5, NumPatches=5 -> totalRows=1600. Each test runs 20 Transpose -> Zero -> MatMul -> in-place Add cycles against a pre-seeded gradW tensor. After every iteration: gradW.Data() is read, checked for all-zero, and a sample of positions is compared to the analytical seed + n*contrib value within tolerance. No failures across 40 total accumulations.

Pod: ztensor-issue79-repro-1775761950. Branch fix/issue-79-matmul-accumulate-repro @ c7739cf.

Conclusion

The bug is not triggered by the patch-embedding backward op chain even at production shapes. Three remaining suspects now narrow significantly:

1. encoderBackward's full op chain (my test skips it).
2. CPU-loop dPosData += dXData interleave at patchtst_gpu_train.go:1012-1019 forcing a mid-backward D2H on the same stream.
3. zerfoo-side gradTs wrapper/arena state that doesn't match fresh tensor.New wrappers.

I recommend pivoting: instrument trainWindowedGPU directly (log device pointers before/after each op on the first batch) rather than continuing to lift primitives into ztensor tests. This issue may belong in zerfoo, not ztensor.

[dndungu]

Resolved on the zerfoo side via zerfoo/zerfoo#371 (merged as 1ebdb787).

Root cause: ztensor/compute/gpu_engine_memory.go:614 GPUEngine.Reshape takes a variadic dst parameter (shared contract with CPUEngine.Reshape) but its zero-copy GPUStorage fast-path at line 644 ignores dst entirely — it returns a fresh tensor aliasing the source storage and never touches the caller-provided dst buffer.

PatchTST's GPU backward at timeseries/patchtst_gpu_train.go:999 discarded the Reshape return value and relied on fc.dX (the dst) being populated. Since GPU Reshape silently left fc.dX as its pre-allocated zeros, encoderBackward saw an all-zero input and every downstream gradient was zero → frozen loss at 0.268357 across every epoch on DGX GB10.

Smoking gun from Wave 7 gpuProbe instrumentation (first batch, samples=500×5×1 on DGX):

post-dFlat-matmul      ptr=0x...628c000 l2=0.0833596 allZero=false
post-dFlat-reshape-dX  ptr=0x...f244000 l2=0         allZero=true
encBwd:entry:dX        ptr=0x...f244000 l2=0         allZero=true

Fix (zerfoo side): capture the Reshape return value and feed it into encoderBackward. Validated on DGX (samples=5000×10×3):

epoch 1: loss=0.027676
epoch 2: loss=0.019176
epoch 3: loss=0.018603
convergence: OK (32.8% reduction)

Follow-up for ztensor (this repo): GPUEngine.Reshape should honor the dst argument per the compute.Engine contract — either copy source data into dst's storage, or SetStorage dst to alias the view, so the next caller that discards the return value doesn't re-hit this silent-zero trap. Filing a separate issue.

Closing this investigation issue — the convergence regression is fixed on main.