Stage 1+2 of #162 (CUDA Graph capture of training step). The AdamW
kernels previously took the step counter as a host scalar arg, which
gets baked into kernel args at CUDA Graph capture time — replays would
freeze the counter and produce wrong bias-correction values.
Both AdamW variants now read the step from a device pointer, advanced
by a tiny 1-thread `increment_counter` kernel that goes inside the
captured region. Each replay correctly increments and observes the
new step value.
Kernel changes:
adamw_step.cu:
- adamw_step: int step → const int* step_ptr
- adamw_increment_counter: new, +=1 on step_ptr[0]
mamba2_alpha_kernel.cu:
- mamba2_alpha_adamw_step_devscale: int t → const int* step_ptr
- mamba2_alpha_increment_step_counter: new
Rust changes:
trainer/optim.rs (AdamW):
- host `step: i32` → device `step_count_d: CudaSlice<i32>`
- step(): launch increment kernel BEFORE adamw kernel; both read
device counter via pointer arg.
- step_count(): test-only accessor, mapped-pinned readback (sync).
mamba2_block.rs (Mamba2AdamW):
- kept host `step_count: i32` for legacy paths (`step`,
`step_from_buffers`) which aren't capture-compatible anyway
(host grad-norm dtoh, host scalar grad_scale).
- added device `step_count_d: CudaSlice<i32>` for the production
gpu_clip path; advances via `kernel_increment_step` kernel
inside the captured region.
- adamw_apply_devscale: `t: i32` → `step_d: &CudaSlice<i32>`.
Validation:
- 4 adamw_invariants tests pass (step_count_increments specifically
exercises the device counter).
- 10 mamba2_block lib tests pass (training_loop_decreases_loss
exercises legacy host-counter path).
- Synthetic overfit smoke: initial=0.25 → final=0.0006 (matches
pre-refactor trajectory bit-for-bit-equivalent).
Stage 3+4 (capture brackets + first-call-capture / subsequent-replay
state machine in step_batched) follows in the next commit.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
ab94ce2a49