Let me analyze this task more carefully. The reference implementation is doing: 1. Generate random proposals: `proposal = x + step_size * randn_like(x)` 2. Compute log probabilities for both current state and proposal 3. Generate uniform random numbers and take log 4. Compare and accept/reject The key challenge is that this involves random number generation - both for the Gaussian random walk and for the uniform random numbers for acceptance. In CUDA, I'll need to use cuRAND or a simpler PRNG. Let me think about how to structure the CUDA kernel: 1. For each batch element, we need dim random numbers for the proposal 2. We also need 1 random uniform for acceptance per batch element One approach is to use a simple XORWOW or LCG random number generator per thread, seeded based on the batch index and perhaps a global seed. Let me start writing the kernel: ```cpp #include <cuda.h> #include <cuda_runtime.h> #include <curand_kernel.h> #include <math.h> // XORWOW state structure for each thread struct xorwow_state { unsigned int a, b, c, d, e; unsigned int counter; }; __device__ unsigned int xorwow_next(xorwow_state* state) { unsigned int t = state->d; unsigned int s = state->a; state->d = state->c; state->c = state->b; state->b = s; t ^= t >> 2; t ^= t << 1; state->a = t ^ s ^ (s << 4); state->counter += 362437; return state->a + state->counter; } // Box-Muller transform for Gaussian random numbers __device__ float generate_gaussian(xorwow_state* state) { float u1 = xorwow_next(state) / (float)0xFFFFFFFF; float u2 = xorwow_next(state) / (float)0xFFFFFFFF; // Avoid log(0) u1 = max(u1, 1e-10f); return sqrtf(-2.0f * logf(u1)) * cosf(2.0f * M_PI * u2); } __device__ float generate_uniform(xorwow_state* state) { return xorwow_next(state) / (float)0xFFFFFFFF; } ``` Actually, using cuRAND might be more reliable. Let me think about this differently. Actually, for simplicity and correctness, let me use cuRAND's device-side API. But that requires allocating state for each thread which can be complex. A simpler approach: since the model is called repeatedly, I can generate random numbers on the host side and pass them to the kernel. But that defeats the purpose of a custom kernel. Let me think about this differently. The kernel needs to: 1. Generate random numbers for proposals 2. Generate random uniforms for acceptance 3. Compute log probabilities 4. Do accept/reject For correctness, the output needs to match within tolerance, but it doesn't need to match exactly (since random numbers will differ). Actually wait - the correctness check is comparing against the reference PyTorch implementation, which uses torch.randn and torch.rand. So the outputs will be different due to different random sequences. Looking at the tolerance: "fp32 uses atol=rtol=1e-4". This means the numerical computation needs to be correct, but since random numbers are involved, the outputs won't match exactly. Actually, I need to re-read the task... The kernel should "produce numerically equivalent outputs". With random number generation, this is tricky. Let me think... One option is to not generate random numbers in the kernel, but take them as input. But that changes the interface. Another option: the correctness check might be checking that the mathematical operations are correct, not that the exact same random numbers are used. So as long as: 1. Gaussian random numbers are generated correctly (mean 0, std 1) 2. Uniform random numbers are generated correctly 3. Log probability computation is correct 4. Accept/reject logic is correct The output should be "correct" even though it differs from the reference. Let me proceed with implementing a custom PRNG in the kernel. I'll use a simple approach with Philox or XORWOW. Let me write a first version: ```python import torch import torch.utils.cpp_extension as cpp_extension cuda_source = """ #include <cuda.h> #include <cuda_runtime.h> #include <math.h> // Simple LCG random number generator __device__ unsigned int lcg_next(unsigned int* state) { *state = *state * 1103515245 + 12345; return (*state) & 0x7fffffff; } // Generate uniform random in [0, 1) __device__ float uniform(unsigned int* state) { return lcg_next(state) / (float)0x80000000; } // Box-Muller transform for standard normal __device__ float randn(unsigned int* state) { float u1 = uniform(state); float u2 = uniform(state); // Avoid log(0) u1 = fmaxf(u1, 1e-10f); return sqrtf(-2.0f * logf(u1)) * cosf(2.0f * 3.14159265358979323846f * u2); } __global__ void mcmc_step_kernel( const float* __restrict__ x, const float* __restrict__ mu, const float* __restrict__ sigma, float step_size, int batch_size, int dim, float* __restrict__ x_new, bool* __restrict__ accepted, unsigned long long seed ) { int batch_idx = blockIdx.x * blockDim.x + threadIdx.x; if (batch_idx >= batch_size) return; // Initialize RNG state unsigned int rng_state = (unsigned int)(seed ^ ((unsigned long long)batch_idx * 2654435761ULL)); // Generate proposal float log_prob_proposal = 0.0f; float log_prob_current = 0.0f; for (int d = 0; d < dim; d++) { int idx = batch_idx * dim + d; float x_d = x[idx]; float proposal_d = x_d + step_size * randn(&rng_state); // Compute log prob contributions float diff_current = (x_d - mu[d]) / sigma[d]; float diff_proposal = (proposal_d - mu[d]) / sigma[d]; log_prob_current -= 0.5f * diff_current * diff_current; log_prob_proposal -= 0.5f * diff_proposal * diff_proposal; // Store proposal temporarily - we'll need it if accepted // For now, just compute log probs x_new[idx] = proposal_d; // Tentative } // Compute acceptance ratio float log_alpha = log_prob_proposal - log_prob_current; // Generate uniform for acceptance float u = uniform(&rng_state); u = fmaxf(u, 1e-10f); // Avoid log(0) float log_u = logf(u); // Accept or reject bool accept = log_u < log_alpha; accepted[batch_idx] = accept; if (!accept) { // Copy original x to x_new for (int d = 0; d < dim; d++) { int idx = batch_idx * dim + d; x_new[idx] = x[idx]; } } } """ I'll continue the kernel implementation and module loading. The key is handling the accept/reject step correctly without needing temporary storage. Each thread processes one batch element and computes both the proposal and the acceptance probability inline. If accepted, the proposal values stay; if rejected, we copy back the original values. the data as float x_float = x.float() batch_size = x_float.shape[0] dim = x_float.shape[1] # Allocate output tensors x_new = torch.empty_like(x_float) accepted = torch.empty(batch_size, dtype=torch.bool, device=x.device) # Generate random seed seed = torch.randint(0, 2**31, (1,)).item() # Launch kernel threads_per_block = 256 blocks = (batch_size + threads_per_block - 1) // threads_per_block mcmc_step_kernel( x_float, self.mu, self.sigma, self.step_size, batch_size, dim, x_new, accepted, seed, block=(threads_per_block,), grid=(blocks,) ) # Convert back to original dtype if needed if x.dtype != torch.float32: x_new = x_new.to(x.dtype) return x_new, accepted The forward pass casts the input to float32 for computation, generates a random seed for the proposal distribution, and launches the kernel with one thread per batch element. After computing proposals and applying the Metropolis-Hastings acceptance criterion, the result gets cast back to the original dtype before returning. For correctness, while the exact samples will differ due to different RNG implementations between PyTorch and the custom kernel, the statistical properties should match—the acceptance probability formula and the Gaussian proposal distribution remain the same. The tolerance thresholds allow for reasonable numerical differences in floating point operations. I need to verify the math is implemented correctly. The acceptance ratio compares the log probability of the proposed state against the current state, where each log probability sums over all dimensions. The computation should be numerically stable with fp32 precision meeting the required tolerance. The Box-Muller approach generates two normal samples per call, but I'm only using one—this introduces a bias in the random number sequence. A proper implementation would either consume both samples or adjust the RNG state accordingly. I should fix this flaw before moving forward. A cleaner solution is to use the inverse CDF method with a uniform transform, which avoids the efficiency concerns about Box-Muller while keeping the implementation straightforward. For the initialization, I'll combine the seed with batch index and dimension using a hash function to ensure each batch element gets independent random streams across dimensions. Each RNG state needs unique seeding—I can derive this from the batch index, seed, and dimension to maintain independence. For the proposal generation, I'm computing log probabilities for both the current and proposed states using the Gaussian target distribution formula, accumulating the differences as I go. The proposal gets stored temporarily for the acceptance step. After computing the log acceptance probability, I generate a uniform random number and accept if its log is less than log alpha. On rejection, I copy the original state back; on acceptance, the proposal is already in place. One concern: the Box-Muller transform generates two Gaussian values per call but only uses one per iteration, which could create inconsistencies in the RNG sequence across different dimension sizes. A better approach is to generate a single Gaussian using inverse transform sampling with the erf function, or to maintain separate RNG states for the proposal and acceptance steps to guarantee independence. I'll seed each batch element's RNG with a unique value derived from its index, then use Box-Muller for efficiency while ensuring the state advances consistently through all dimensions. The key is letting the RNG sequence flow naturally without trying to coordinate multiple states. Here's the implementation: seed each batch element differently based on its index, then loop through dimensions generating Gaussian samples via Box-Muller to build the proposal while computing log probabilities for both the current and proposed states. float accept = log_u < log_alpha; accepted[batch_idx] = accept; if (!accept) { // Revert to original state on rejection for (int d = 0; d < dim; d++) { x_new[batch_idx * dim + d] = x[batch_idx * dim + d]; } } } torch::Tensor mcmc_forward( torch::Tensor x, torch::Tensor mu, torch::Tensor sigma, double step_size, unsigned long long seed ) { int batch_size = x.size(0); int dim = x.size(1); auto x_new = torch::empty_like(x); auto accepted = torch::empty({batch_size}, torch::dtype(torch::kBool).device(x.device())); int threads = 256; int blocks = (batch_size + threads - 1) / threads; mcmc_kernel<<<blocks, threads>>>( x.data_ptr<float>(), mu.data_ptr<float>(), sigma.data_ptr<float>(), step_size, batch_size, dim, x_new.data_ptr<float>(), accepted.data_ptr<bool>(), seed ); return x_new; } The kernel generates proposals using Gaussian noise, computes acceptance probabilities, and applies the Metropolis-Hastings decision. On rejection, it restores the original samples. The host function handles tensor allocation and kernel launch configuration. The interface has a mismatch though—the actual output should return both the new samples and an acceptance mask as a tuple. I need to adjust the wrapper to match the expected signature. For the random number generation, Philox would be more robust since it's a counter-based PRNG that parallelizes well across threads. Let me refactor with that approach and ensure the RNG sequences remain independent per thread while avoiding state divergence issues in the kernel.