Let me analyze the task. I need to implement a custom CUDA kernel for Hamiltonian Monte Carlo (HMC) proposal using leapfrog integration. The key operations are: 1. Generate random momentum `p` from Gaussian distribution 2. Perform leapfrog integration steps: - Half step for momentum: `p = p + 0.5 * eps * grad_log_prob(q)` - Multiple full steps: `q = q + eps * p; p = p + eps * grad_log_prob(q)` - Final half step for momentum - Negate momentum for reversibility 3. Compute Hamiltonians for current and proposed states 4. Metropolis-Hastings acceptance step The gradient of log probability for standard Gaussian is: `grad_log_prob(q) = -(q - mu) = -q` (since mu=0) So the operations are: - `p = randn_like(q)` - `p = p + 0.5 * eps * (-q) = p - 0.5 * eps * q` - Loop: `q = q + eps * p; p = p - eps * q` - Final: `q = q + eps * p; p = p - 0.5 * eps * q; p = -p` - `current_H = -log_prob(current_q) + 0.5 * sum(current_p^2) = 0.5 * sum(current_q^2) + 0.5 * sum(current_p^2)` - `proposed_H = 0.5 * sum(q^2) + 0.5 * sum(p^2)` - `log_alpha = current_H - proposed_H` - Accept if `log(rand()) < log_alpha` Let me design a CUDA kernel for this: 1. I need to generate random numbers - this is tricky in CUDA. I could use cuRAND or pass pre-generated random numbers from PyTorch. 2. The leapfrog integration involves multiple position and momentum updates that can be parallelized across dimensions. 3. Computing Hamiltonians requires summing squares across dimensions for each batch element. 4. The acceptance decision needs to be made per batch element. The main challenge is handling the random number generation. I could pre-generate random numbers in PyTorch and pass them to the kernel, or use cuRAND within the kernel for better performance. Since there are multiple random number draws (initial momentum and uniform random for acceptance), I need to account for the total number of random values needed. Looking at the actual requirements: the momentum `p` is drawn once at the start, and a uniform random `u` is drawn once per batch element for the acceptance test. I'll generate both in PyTorch beforehand and pass them as arguments to keep the kernel simpler and more efficient. The Hamiltonian computation involves summing across dimensions, which can be done efficiently with parallel reduction, followed by an element-wise acceptance check. However, the loop structure reveals a complexity—the leapfrog integration alternates between updating momentum with the gradient and updating position with the momentum, and these operations depend on each other's results. This dependency chain makes straightforward parallelization challenging, since each step requires the output of the previous step. The gradient computation itself is element-wise, which is promising, but the real bottleneck is that I need to compute `grad_log_prob(q)` at multiple different `q` values throughout the trajectory. This means I can't precompute all gradients ahead of time—they must be computed at each step as `q` evolves. The dependencies create a sequential structure: updating position requires the current momentum, and updating momentum requires the current position. Since all dimensions of position and momentum get updated simultaneously at each leapfrog step, I could parallelize within each dimension at a given step, but the steps themselves must proceed in order. Each thread could handle one dimension, updating `q[d]` and `p[d]` in lockstep across the trajectory. For the Hamiltonian calculation, I'd need to sum across dimensions using parallel reduction. I could pre-generate random numbers in PyTorch rather than relying on cuRAND to keep things simpler. The kernel would take the position `q`, random momentum `p_rand`, uniform samples `u`, and parameters like step size and leapfrog steps, then have each thread process one dimension through the leapfrog sequence. The momentum and position updates interleave in a specific sequence during leapfrog integration. After the initial half-step momentum update, we alternate between full position and momentum updates for the middle steps, then finish with a final position update and another half-step momentum update before negation. This structure ensures proper time-reversibility in the Hamiltonian dynamics. Since we can't easily parallelize the leapfrog steps themselves due to their sequential nature, I should focus on optimizing memory access and leveraging PyTorch's existing operations where possible. The most computationally intensive parts—the random number generation and basic arithmetic—are already well-optimized in PyTorch, so a custom CUDA kernel would need to be carefully designed to provide real speedup over just using `torch.randn` for momentum initialization and standard tensor operations for the integration steps. A fusion approach could work well here: pre-generate the random momentum and uniform samples on the GPU, then use a dedicated kernel that handles the entire leapfrog trajectory and acceptance logic in one pass. Each thread would process one batch element, computing all dimension updates for each integration step while maintaining the sequential dependency of the leapfrog scheme. This way we avoid repeated kernel launches and memory transfers between steps. Now I'll implement the CUDA kernel. Each thread processes one batch element, storing position and momentum in registers for fast access. The leapfrog steps update q and p alternately, starting with a half-step on momentum, then alternating full steps on position and momentum. After the specified number of steps, I need to complete the final half-step on momentum and then compute the Hamiltonian to determine acceptance. For the Hamiltonian calculation, I sum the squared values of the proposed position and momentum across dimensions. I also need to compute the current Hamiltonian using the original position and momentum values that were saved at the beginning. The acceptance probability is then determined by comparing these two energy values. I compute the log acceptance ratio as the difference between current and proposed Hamiltonians, then compare it against a log uniform sample. If the proposal is accepted, I copy the new position into the output; otherwise, I preserve the original position and record the acceptance status. The kernel launch configuration uses a block size of 256 threads, with enough blocks to cover all batch elements. I'm passing the device pointers for positions, momenta, uniform samples, outputs, and acceptance flags along with the step parameters. The forward method accepts the position tensor and returns both the updated positions and a boolean tensor indicating which proposals were accepted. To implement this, I generate random momenta from a standard normal distribution and uniform samples for the acceptance test, then invoke the kernel with the appropriate dimension and leapfrog parameters. The positions tensor is already on the correct device, and I ensure the acceptance tensor matches its device placement. I need to reconsider the random number handling. The reference generates momentum via `torch.randn_like(q)` and uniform samples for acceptance, which I should preserve. The key constraint is that each batch element needs one uniform random number for the acceptance step. I can generate these tensors in PyTorch upfront and pass them to the kernel, keeping the kernel focused on the leapfrog integration and acceptance logic rather than random number generation. The leapfrog step itself is where the optimization matters most—it's the computational core where I can fuse operations for better performance. The Hamiltonian computation requires a sum across dimensions, which presents a challenge for parallelization. I'm considering having each thread handle one batch element across all dimensions—if the dimensionality is moderate (like 64), the thread can efficiently manage all data, and I can leverage shuffle operations for reduction if needed. For correctness, I need to verify that `grad_log_prob(q) = -(q - mu) = -q` when mu is zero. I should also examine how `log_prob` computes the full log-likelihood, as this affects the Hamiltonian calculation. Looking at the reference implementation, `log_prob(x)` returns -0.5 * ||x||², which means the Hamiltonian H = -log_prob(q) + 0.5 * ||p||² = 0.5 * ||q||² + 0.5 * ||p||². In the leapfrog integration loop, I notice there's a final position update and half-step momentum update outside the main iteration cycle, followed by negation of momentum for reversibility. After negating momentum, the proposed Hamiltonian uses the final position and the negated momentum value, which maintains the correct symplectic structure needed for the Metropolis-Hastings acceptance step. For the CUDA kernel, I'll keep the structure straightforward: each thread handles one batch element across all dimensions, pre-generating random numbers for the momentum initialization and acceptance decisions before passing them to the kernel. I need to complete the Hamiltonian computation for the proposed state and then handle the acceptance logic. The log probability sums over all dimensions, so I accumulate the squared momentum and position values. Then I compute the Hamiltonian as the negative log probability plus the kinetic energy term. For acceptance, I compare the current Hamiltonian against the proposed one—if the proposed state has lower energy, it's more likely to be accepted. The acceptance decision uses a uniform random number and the log of that number compared against the energy difference. If accepted, I copy the proposed state; otherwise I keep the current state. The kernel needs to be launched with appropriate block and grid dimensions based on the batch size. In the forward pass, I generate random momentum and uniform samples, then prepare the output tensors before launching the kernel with the position data, random samples, and leapfrog parameters like step size and number of steps. # Final half step for momentum for (int d = 0; d < dim; d++) { T q_val = q[d]; p[d] = p[d] - half_eps * q_val; } // Negate momentum for (int d = 0; d < dim; d++) { p[d] = -p[d]; } // Compute Hamiltonians T H_current = 0, H_proposed = 0; for (int d = 0; d < dim; d++) { H_current += 0.5f * (current_q[d] * current_q[d] + current_p[d] * current_p[d]); H_proposed += 0.5f * (q[d] * q[d] + p[d] * p[d]); } // Metropolis acceptance T log_alpha = H_current - H_proposed; The kernel completes the leapfrog integration with a final momentum half-step, then negates momentum for reversibility. Hamiltonians for both current and proposed states are computed by summing kinetic and potential energy across dimensions. The Metropolis-Hastings acceptance probability follows from the energy difference. The loop unrolling should work well given that dim is fixed at 64—compilers typically handle small constant iteration counts effectively. Register usage depends on whether the compiler keeps `q`, `p`, `current_q`, and `current_p` in registers across the leapfrog steps, which could pressure register availability but should remain manageable for this dimensionality. Looking back at the reference implementation, I notice the initial momentum update is a half-step (`p = p + 0.5 * eps * grad_log_prob(q)`), then the loop runs `num_leapfrog_steps - 1` iterations with full steps, and there's a final position update followed by another half-step on momentum. This matches what I have, so the structure is correct. Computing the gradients inline eliminates the intermediate arrays, which should improve memory efficiency. For handling different precision types, I can leverage CUDA's function overloading to define type-specific absolute value operations rather than relying on a potentially problematic macro approach.